DOCSEARCHbyAlgolia
Submit Docs

titans-pytorch

repository·main·Indexed 24 days ago

https://github.com/lucidrains/titans-pytorch

An unofficial PyTorch implementation of the Titans paper, focusing on neural memory modules and architectures that learn to memorize at test time. It provides the NeuralMemory module, MemoryAsContextTransformer, and various memory model architectures including MemoryMLP, MemoryAttention, FactorizedMemoryMLP, MemorySwiGluMLP, and GatedResidualMemoryMLP. The library also implements ImplicitMLPAttention and NestedAttention for advanced memory-augmented networks.

Tokens
4.6K
Snippets
8
Records
28
Agent score
83%

What's inside titans-pytorch

  3. How NeuralMemory stores and retrieves information

    main

    The NeuralMemory module operates via two primary internal methods:

    1. store_memories (The Update Step)

    This method takes a sequence and updates the memory model's weights. It calculates the 'surprise' (the gradient of the loss between the current memory model's prediction and the target value) and applies this surprise as a weight update. It supports:

    • Adaptive Learning Rates: Per-token learning rates.
    • Momentum: Using associative scans to accumulate updates.
    • Weight Decay: Using associative scans to implement learned forgetting.

    2. retrieve_memories (The Query Step)

    This method uses the current weights of the memory model to fetch information. It transforms the input sequence into queries and uses the memory_model (via functional_call) to produce values based on the current weights. This is mathematically similar to a fast-weight memory or linear attention mechanism.

  4. Use the MemoryAsContextTransformer

    main

    The MemoryAsContextTransformer provides a transformer architecture using the MAC configuration, integrating memory as context.

    Parameters:

    • num_tokens: Number of tokens.
    • dim: Dimensionality.
    • depth: Number of layers.
    • segment_len: Local attention window size.
    • num_persist_mem_tokens: Number of persistent memory tokens.
    • num_longterm_mem_tokens: Number of long-term memory tokens.

    Methods:

    • forward(token_ids, return_loss=False): Returns the output (or loss if return_loss=True).
    • sample(token_ids, length): Samples new tokens based on the provided input sequence.
    import torch
from titans_pytorch import MemoryAsContextTransformer

transformer = MemoryAsContextTransformer(
    num_tokens = 256,
    dim = 256,
    depth = 2,
    segment_len = 128,              # local attention window size
    num_persist_mem_tokens = 4,
    num_longterm_mem_tokens = 16,
)

token_ids = torch.randint(0, 256, (1, 1023))

loss = transformer(token_ids, return_loss = True) # (1, 1023, 256)
loss.backward()

# after much training

sampled = transformer.sample(token_ids[:, :4], 512)

  5. Use the NeuralMemory module

    main

    The NeuralMemory module implements the neural memory component from the Titans paper. It takes a sequence and returns retrieved information along with the updated memory state.

    Parameters:

    • dim: The dimensionality of the input/output.
    • chunk_size: The size of the memory chunks. Setting this to a smaller value can improve performance on shorter sequences but will increase memory usage.
    import torch
from titans_pytorch import NeuralMemory

mem = NeuralMemory(
    dim = 384,
    chunk_size = 64 # set to smaller chunk size for better perf on smaller sequence lengths (but more memory usage)
).cuda()

seq = torch.randn(2, 1024, 384).cuda()
retrieved, mem_state = mem(seq)

assert seq.shape == retrieved.shape

  6. Configure NeuralMemory initialization parameters

    main

    When initializing NeuralMemory, several advanced configuration options are available to control the memory dynamics:

    ParameterDescription
    dimFeature dimension of the input/output
    chunk_sizeSize of chunks used for storing to memory model weights
    headsNumber of memory heads
    modelThe nn.Module used as the memory model (defaults to MemoryMLP)
    store_memory_loss_fnLoss function used to calculate the 'surprise' for weight updates
    adaptive_step_transformFunction to transform the learned adaptive step size
    momentumBoolean, whether to use momentum in updates
    momentum_orderThe order of momentum to use
    num_kv_per_tokenNumber of key/value pairs each token can emit
    spectral_norm_surprisesWhether to apply Newton-Schulz spectral norming to updates
    gated_transitionWhether to use a learned gate for transitioning from initial to updated weights
    use_accelerated_scanWhether to use an accelerated associative scan implementation
    store_with_lookahead_valueWhether to use values from the next timestep for gradients
    default_model_kwargsDictionary of arguments passed to the MemoryMLP model

  7. Implement incremental decoding with NestedAttention cache

    main

    To use NestedAttention for autoregressive generation, set return_kv_cache=True during the initial pass to obtain the cache. In subsequent steps, pass the tokens (typically just the last token) and the retrieved cache back into the forward method.

    nested_attn = NestedAttention(512)
tokens = torch.randn(1, 1024, 512)

# Initial pass (e.g., prompt processing)
out1, cache = nested_attn(tokens, return_kv_cache=True)

# Incremental step (e.g., generating the next token)
out2, cache = nested_attn(tokens[:, -1:], cache=cache, return_kv_cache=True)
    nested_attn = NestedAttention(512)

tokens = torch.randn(1, 1024, 512)

out1, cache = nested_attn(tokens, return_kv_cache = True)
out2, cache = nested_attn(tokens[:, -1:], cache = cache, return_kv_cache = True)

assert out1.shape == tokens.shape
assert out2.shape == (1, 1, 512)

  8. Initialize ImplicitMLPAttention

    main

    The ImplicitMLPAttention class implements a neural memory attention mechanism where each key-value pair forms an implicit weight (memory) of dimensions (dim_key, dim_values). Chaining these pairs creates an implicit MLP structure inspired by TTT (Test-Time Training) and Titans architectures.

    Parameters

    • dim: The input dimension of the tokens.
    • mlp_hiddens: A tuple[int, ...] defining the dimensions of the implicit MLP. The first and last elements represent the input and output dimensions of the MLP chain, while the intermediate elements represent the hidden layers. It must have at least 2 elements.
    • activation: The activation function used between implicit layers (default: nn.SiLU()).
    • heads: The number of attention heads (default: 8).
    • talking_heads: If True, applies a Conv2d layer with a Dirac initialization to allow communication between heads (default: True).
    • prenorm: If True, applies RMSNorm to the input tokens before processing (default: True).
    • keys_rmsnorm: If True, applies RMSNorm to the keys (default: True).

  9. Use MemoryAsContextTransformer for long-context modeling

    main

    The MemoryAsContextTransformer is the primary class for implementing a transformer architecture that integrates long-term memory tokens and neural memory. It supports segmented attention, axial positional embeddings, and hyper-connections to manage multiple residual streams.

    Key features include:

    • Long-term memory tokens: Learned parameters interspersed into the sequence.
    • Neural Memory: An optional NeuralMemory module that can be integrated into the layers.
    • Flex Attention: Support for PyTorch's flex_attention for efficient masked attention.
    • Sliding Window Attention: Optional local attention mechanism.
    • Value Residuals: Ability to pass residuals into the attention mechanism.

    To use it, initialize the class with your desired dimensions and memory configurations, then call .forward() for training/inference or .sample() for autoregressive generation.

  10. Use NestedAttention in forward pass

    main

    The forward method processes input tokens and optionally manages a KV cache for autoregressive generation.

    Arguments:

    • tokens: Input tensor of shape (batch, seq_len, dim).
    • cache: An optional cache tuple containing ((cache_keys, cache_values), (cache_nested_keys, cache_nested_values)) to support incremental decoding.
    • return_kv_cache: If True, the method returns the updated KV cache along with the output.

    Returns:

    • If return_kv_cache=False: Returns the output tensor of shape (batch, seq_len, dim).
    • If return_kv_cache=True: Returns a tuple (output, cache) where cache is the updated KV cache.

  11. Initialize NestedAttention

    main

    The NestedAttention class implements a hierarchical attention mechanism. It projects input tokens into queries, keys, and values, and performs nested attention passes.

    Parameters:

    • dim: The input dimension of the tokens.
    • dim_head: The dimension of each attention head (default: 64).
    • heads: The number of attention heads (default: 8).
    • prenorm: If True, applies nn.RMSNorm to the input tokens before processing (default: True).
    • keys_rmsnorm: Whether to use RMSNorm on keys (default: True).

  12. Use GEGLU and FeedForward layers

    main

    The module provides a GEGLU activation and a FeedForward block commonly used in modern transformers.

    • GEGLU: Implements Gated Linear Unit with SiLU activation. It splits the input into two parts, applies SiLU to one, and multiplies it by the other.
    • FeedForward: A standard feedforward block consisting of RMSNorm, a linear expansion layer, GEGLU, and a linear contraction layer.

    Arguments for FeedForward:

    • dim: Input/output dimension.
    • mult: Multiplier for the inner dimension (default is 4).