Recurrent Memory Transformer PyTorch

For training on extremely long sequences, use the RecurrentMemoryTransformerWrapper. This wrapper supports memory_replay_backprop, a memory-efficient training technique from the Memformer paper.

Workflow:

1. Wrap your RecurrentMemoryTransformer instance with RecurrentMemoryTransformerWrapper.
2. Pass a long sequence (e.g., shape (batch, total_seq_len)) to the wrapper.
3. Set memory_replay_backprop = True in the call to enable efficient backpropagation through the memory segments.

import torch
from recurrent_memory_transformer_pytorch import (
    RecurrentMemoryTransformer,
    RecurrentMemoryTransformerWrapper
)

model = RecurrentMemoryTransformer(
    num_tokens = 256,
    num_memory_tokens = 128,
    dim = 512,
    depth = 6,
    seq_len = 1024,
    use_flash_attn = True,
    causal = True
)

# Wrap the model for long sequence training
model = RecurrentMemoryTransformerWrapper(model).cuda()

# Example long sequence
seq = torch.randint(0, 256, (4, 65536)).cuda()

# Train with memory replay backpropagation
loss = model(seq, memory_replay_backprop = True)

The RecurrentMemoryTransformer allows for processing sequences in segments while passing memory tokens to future segments.

Key Arguments:

• num_tokens: Vocabulary size.
• num_memory_tokens: Number of memory tokens (determines the information bottleneck).
• dim: Model dimension.
• depth: Number of transformer layers.
• causal: Whether to use autoregressive masking.
• seq_len: Sequence length of a single segment.
• use_flash_attn: Whether to use Flash Attention.

Returns:

• logits: The output logits for the current segment.
• mem: The updated memory tokens to be passed to the next segment.
• _: An additional return value (often None in basic mode).

import torch
from recurrent_memory_transformer_pytorch import RecurrentMemoryTransformer

model = RecurrentMemoryTransformer(
    num_tokens = 20000,
    num_memory_tokens = 128,
    dim = 512,
    depth = 6,
    causal = True,
    dim_head = 64,
    heads = 8,
    seq_len = 1024,
    use_flash_attn = True
)

x = torch.randint(0, 256, (1, 1024))

# Initial forward pass
logits1, mem1, _ = model(x)        # (1, 1024, 20000), (1, 128, 512), None

# Subsequent pass using memory from previous segment
logits2, mem2, _ = model(x, mem1)  # (1, 1024, 20000), (1, 128, 512), None

To use XL memories (which provide additional context from the immediate past), set use_xl_memories = True and specify xl_mem_len. This can help prevent RMT memories from over-memorizing the immediate preceding text.

New Arguments:

• use_xl_memories: Boolean to enable XL memory.
• xl_mem_len: Length of the XL memory (can be shorter than seq_len).

Returns:

• logits: Output logits.
• mem: Updated RMT memory tokens.
• xl_mem: A list containing the XL memories (e.g., [(batch, 1, dim, dim)]).

import torch
from recurrent_memory_transformer_pytorch import RecurrentMemoryTransformer

model = RecurrentMemoryTransformer(
    num_tokens = 20000,
    num_memory_tokens = 128,
    dim = 512,
    depth = 6,
    causal = True,
    dim_head = 64,
    heads = 8,
    seq_len = 1024,
    use_flash_attn = True,
    use_xl_memories = True,
    xl_mem_len = 512
)

x = torch.randint(0, 256, (1, 1024))

# Initial pass with XL memories
logits1, mem1, xl_mem1 = model(x)

# Subsequent pass passing both RMT memory and XL memories
logits2, mem2, xl_mem2 = model(x, mem1, xl_memories = xl_mem1)