Titans: Learning to Memorize at Test Time

Table of Contents

Paper References

Features

Core Features

Backend-Specific Features

Test Coverage

Architecture Overview

Memory Perspective

Architecture Variants

Quick Comparison

When to Use Each Variant

MAC: Memory as Context (Section 4.1)

h_t = M*_{t-1}(q_t) # Eq. 21: Retrieve from memory
S̃^(t) = [persistent] || h_t || x # Eq. 22: Concatenate
y_t = Attn(S̃^(t)) # Eq. 23: Segmented attention
M_t = M_{t-1}(y_t) # Eq. 24: Update memory
o_t = y_t ⊗ M*_t(y_t) # Eq. 25: Output gating

MAG: Memory as Gate (Section 4.2)

x̃ = [persistent] || x # Eq. 26: Add persistent tokens
y = SW-Attn*(x̃) # Eq. 27: Sliding window attention
o = y ⊗ M(x̃) # Eq. 28: Element-wise gating

MAL: Memory as Layer (Section 4.3)

x̃ = [persistent] || x # Eq. 29: Add persistent tokens
y = M(x̃) # Eq. 30: Memory layer
o = SW-Attn(y) # Eq. 31: Attention on memory output

Neural Long-term Memory

Core Equations (Section 3.1)

l(M; x_t) = ||M(k_t) - v_t||2

M_t = (1 - α_t) · M_{t-1} + S_t

S_t = η_t · S_{t-1} - θ_t · ∇l(M_{t-1}; x_t)
 \_________/ \____________________/
 Past Surprise Momentary Surprise

Key Innovations

Installation

Basic Installation (PyTorch)

git clone https://github.com/yourusername/Google-Titans-replication.git
cd Google-Titans-replication
uv sync

With Training Dependencies

uv sync --extra train

With All Extras (Development)

uv sync --all-extras

MLX Requirements

# MLX is included in default dependencies
uv sync

Quick Start

PyTorch Quick Start

import torch
from titans import TitansConfig, TitansMAC, TitansMAG, TitansMAL
# Configuration
config = TitansConfig(
 dim=512,
 num_heads=8,
 num_layers=6,
 vocab_size=32000,
 chunk_size=512, # For MAC
 window_size=512, # For MAG/MAL
 num_persistent_tokens=16,
 num_memory_layers=2, # Deep memory
)
# Create model
model = TitansMAC(config) # or TitansMAG, TitansMAL
# Forward pass
input_ids = torch.randint(0, config.vocab_size, (2, 1024))
logits, states = model(input_ids)
# Continue with states for next segment
input_ids_next = torch.randint(0, config.vocab_size, (2, 512))
logits_next, states = model(input_ids_next, states=states)

MLX Quick Start

import mlx.core as mx
from titans_mlx import TitansConfig, TitansMAC, TitansMAG, TitansMAL
# Configuration (same as PyTorch)
config = TitansConfig(
 dim=512,
 num_heads=8,
 num_layers=6,
 vocab_size=32000,
 chunk_size=512,
 window_size=512,
 num_persistent_tokens=16,
 num_memory_layers=2,
)
# Create model
model = TitansMAC(config)
mx.eval(model.parameters()) # Evaluate parameters
# Forward pass
input_ids = mx.random.randint(0, config.vocab_size, (2, 1024))
logits, states = model(input_ids)
mx.eval(logits) # Force evaluation
# Continue with states
input_ids_next = mx.random.randint(0, config.vocab_size, (2, 512))
logits_next, states = model(input_ids_next, states=states)

Standalone Neural Memory

# PyTorch
from titans import TitansConfig, NeuralLongTermMemory
import torch
config = TitansConfig(dim=512, num_memory_layers=2)
memory = NeuralLongTermMemory(config)
x = torch.randn(2, 100, 512)
output, state = memory(x)
output2, state2 = memory(x, state=state) # Continue with state
# MLX
from titans_mlx import TitansConfig, NeuralLongTermMemory
import mlx.core as mx
config = TitansConfig(dim=512, num_memory_layers=2)
memory = NeuralLongTermMemory(config)
mx.eval(memory.parameters())
x = mx.random.normal((2, 100, 512))
output, state = memory(x)
mx.eval(output)

Pretraining

PyTorch Pretraining

Option 1: HuggingFace Streaming (Simple, No Setup)

# Train with FineWeb-Edu streaming
uv run python scripts/pretrain.py --model mac \
 --dataset HuggingFaceFW/fineweb-edu \
 --dataset-subset sample-10BT \
 --tokenizer NousResearch/Llama-2-7b-hf \
 --dim 512 --num-layers 12 \
 --mixed-precision bf16
# Full training (340M params)
uv run python scripts/pretrain.py --model mac \
 --dataset HuggingFaceFW/fineweb-edu \
 --dataset-subset sample-10BT \
 --tokenizer NousResearch/Llama-2-7b-hf \
 --dim 1024 --num-layers 24 --num-heads 16 \
 --batch-size 8 --gradient-accumulation-steps 32 \
 --lr 4e-4 --mixed-precision bf16 --wandb

Option 2: Pre-tokenized Local Dataset (Fastest)

# Step 1: Pre-tokenize (one time)
uv run python scripts/pretokenize.py \
 --dataset HuggingFaceFW/fineweb-edu \
 --subset sample-10BT \
 --tokenizer NousResearch/Llama-2-7b-hf \
 --output data/fineweb-tokenized \
 --seq-len 4096 \
 --num-proc 8
# Step 2: Train with pre-tokenized data
uv run python scripts/pretrain.py --model mac \
 --local-dataset data/fineweb-tokenized \
 --tokenizer NousResearch/Llama-2-7b-hf \
 --dim 512 --num-layers 12 \
 --mixed-precision bf16

Other Options

# Demo with synthetic data (quick test)
uv run python scripts/pretrain.py --model mac --dim 256 --epochs 10
# Train with local text file
uv run python scripts/pretrain.py --model mag \
 --data path/to/corpus.txt \
 --tokenizer gpt2
# Resume from checkpoint
uv run python scripts/pretrain.py --model mac \
 --resume checkpoints/latest.pt

PyTorch Training Options

CUDA Optimizations

Distributed Training (Multi-GPU)

# Multi-GPU with DDP (auto-detects GPUs)
uv run accelerate launch scripts/pretrain_distributed.py \
 --model mac --dim 512 \
 --local-dataset data/fineweb-tokenized
# Multi-GPU with custom config
uv run accelerate launch --config_file configs/fsdp_config.yaml \
 scripts/pretrain_distributed.py --model mac --dim 1024

MLX Pretraining

# Demo with synthetic data
uv run python scripts/pretrain_mlx.py --model mac --dim 256 --epochs 10
# Train with FineWeb-Edu
uv run python scripts/pretrain_mlx.py --model mac \
 --dataset HuggingFaceFW/fineweb-edu \
 --tokenizer meta-llama/Llama-2-7b-hf \
 --dim 512 --num-layers 12
# Full training
uv run python scripts/pretrain_mlx.py --model mac \
 --dataset HuggingFaceFW/fineweb-edu \
 --tokenizer meta-llama/Llama-2-7b-hf \
 --dim 1024 --num-layers 24 --num-heads 16 \
 --batch-size 4 --gradient-accumulation-steps 32 \
 --dtype float16 --wandb
# Train with local text
uv run python scripts/pretrain_mlx.py --model mag \
 --data path/to/corpus.txt
# Resume from checkpoint
uv run python scripts/pretrain_mlx.py --model mac \
 --resume checkpoints_mlx/latest.safetensors

MLX Training Options

Inference

PyTorch Inference

# Generate text
uv run python scripts/inference.py \
 --checkpoint checkpoints/best_model.pt \
 --prompt "Once upon a time" \
 --max-tokens 100
# Interactive mode
uv run python scripts/inference.py \
 --checkpoint checkpoints/best_model.pt \
 --interactive
# With sampling parameters
uv run python scripts/inference.py \
 --checkpoint checkpoints/best_model.pt \
 --prompt "The meaning of life is" \
 --temperature 0.8 \
 --top-p 0.9 \
 --max-tokens 200
# With quantization
uv run python scripts/inference.py \
 --checkpoint checkpoints/best_model.pt \
 --prompt "Hello" \
 --quantize int8

PyTorch Inference Options

MLX Inference

# Generate text
uv run python scripts/inference_mlx.py \
 --checkpoint checkpoints_mlx/best_model.safetensors \
 --prompt "Once upon a time" \
 --max-tokens 100
# Interactive mode
uv run python scripts/inference_mlx.py \
 --checkpoint checkpoints_mlx/best_model.safetensors \
 --interactive
# With quantization and benchmark
uv run python scripts/inference_mlx.py \
 --checkpoint checkpoints_mlx/best_model.safetensors \
 --prompt "Hello" \
 --quantize 8 \
 --benchmark

MLX Inference Options

Benchmarks

Model Quality (from Paper Table 1 & 5)

Inference Speed (Apple M4 Pro)

Numerical Parity

Configuration Reference

TitansConfig Parameters

API Reference

PyTorch API

from titans import (
 # Configuration
 TitansConfig,
 # Models
 TitansMAC,
 TitansMAG,
 TitansMAL,
 TitansLMM,
 # Components
 NeuralLongTermMemory,
 MemoryState,
 SlidingWindowAttention,
 SegmentedAttention,
 PersistentMemory,
)
# Model forward signature
logits, states = model(input_ids, states=None)
# input_ids: (batch, seq_len) - Token IDs
# states: Optional list of MemoryState
# Returns: logits (batch, seq_len, vocab_size), new states
# Memory forward signature
output, state = memory(x, state=None, return_state=True)
# x: (batch, seq_len, dim)
# state: Optional MemoryState
# Returns: output (batch, seq_len, dim), new state

MLX API

from titans_mlx import (
 # Configuration
 TitansConfig,
 # Models
 TitansMAC,
 TitansMAG,
 TitansMAL,
 TitansLMM,
 # Components
 NeuralLongTermMemory,
 MemoryState,
 SlidingWindowAttention,
 SegmentedAttention,
 PersistentMemory,
 # Optimizations
 compile_model, # Note: Limited support
 compile_function,
 get_device_info,
 # Metal Kernels (benchmarking only)
 metal_silu_gate,
 metal_memory_update,
 metal_rope,
)

MLX Optimizations

Gradient Computation

# Efficient gradient via matmul (avoids huge intermediate tensors)
# Instead of: expand_dims + outer product + sum
# We use: reshape + matmul
delta_flat = delta.reshape(batch_seq, -1) # (B*S, D_out)
act_flat = act.reshape(batch_seq, -1) # (B*S, D_in)
grad_w = delta_flat.T @ act_flat # (D_out, D_in)

Why Not mx.compile?

Metal Kernels

Recommended Practices

import mlx.core as mx
from titans_mlx import TitansConfig, TitansMAC
# 1. Use float16 for training (default)
# Saves memory, marginal speed difference on Apple Silicon
# 2. Evaluate parameters after creation
model = TitansMAC(config)
mx.eval(model.parameters())
# 3. Evaluate outputs when needed
logits, states = model(input_ids)
mx.eval(logits) # Force computation
# 4. Use larger batches to amortize overhead
# batch_size=4 or higher recommended
# 5. Disable convolution if dimensions mismatch
config = TitansConfig(..., use_conv=False)

Troubleshooting

Common Issues

PyTorch

# Reduce batch size or use gradient accumulation
--batch-size 2 --gradient-accumulation-steps 64

# Use bf16 instead of fp16, or reduce learning rate
--mixed-precision bf16 --lr 2e-4

MLX

# Disable convolution
config = TitansConfig(..., use_conv=False)

# Normal - MLX compiles on first call
# Subsequent iterations will be faster

# Force garbage collection
import gc
gc.collect()
mx.metal.clear_cache() # If available

Numerical Differences

Development

Project Structure

titans-pytorch/
├── src/
│ ├── titans/ # PyTorch implementation
│ │ ├── __init__.py
│ │ ├── config.py # TitansConfig
│ │ ├── memory.py # Neural Long-term Memory
│ │ ├── attention.py # Attention modules
│ │ ├── persistent.py # Persistent Memory
│ │ ├── models.py # MAC, MAG, MAL, LMM
│ │ └── triton_kernels.py # Triton optimizations
│ │
│ └── titans_mlx/ # MLX implementation
│ ├── __init__.py
│ ├── config.py
│ ├── memory.py
│ ├── attention.py
│ ├── persistent.py
│ ├── models.py
│ ├── optimizations.py # MLX optimizations
│ └── metal_kernels.py # Metal kernels
│
├── scripts/
│ ├── pretrain.py # PyTorch training (optimized)
│ ├── pretrain_distributed.py # Multi-GPU training (Accelerate)
│ ├── pretrain_mlx.py # MLX training
│ ├── pretokenize.py # Dataset pre-tokenization
│ ├── inference.py # PyTorch inference
│ └── inference_mlx.py # MLX inference
│
├── tests/
│ ├── test_memory.py
│ ├── test_attention.py
│ ├── test_models.py
│ ├── test_persistent.py
│ └── test_numerical_parity.py # MLX vs PyTorch
│
├── examples/
│ ├── basic_usage.py
│ └── long_sequence.py
│
└── pyproject.toml

Running Tests

# All tests
uv run pytest tests/ -v
# Specific test file
uv run pytest tests/test_numerical_parity.py -v
# With coverage
uv run pytest tests/ --cov=titans --cov=titans_mlx --cov-report=term-missing

Linting

uv run ruff check src/ tests/ scripts/
uv run ruff format src/ tests/ scripts/

Citation

@article{behrouz2024titans,
 title={Titans: Learning to Memorize at Test Time},
 author={Behrouz, Ali and Zhong, Peilin and Mirrokni, Vahab},
 journal={arXiv preprint arXiv:2501.00663},
 year={2024}
}
@article{dinepi2025titans,
 title={Titans Revisited: A Lightweight Reimplementation and Critical Analysis of a Test-Time Memory Model},
 author={Di Nepi, Gavriel and Siciliano, Federico and Silvestri, Fabrizio},
 journal={arXiv preprint arXiv:2510.09551},
 year={2025}
}

License