titans-unofficial

An unofficial PyTorch implementation of
Titans: Learning to Memorize at Test Time (Behrouz, Zhong and Mirrokni).

Titans pairs attention, which acts as short-term memory, with a small neural network that serves as long-term memory. That network keeps learning while the model reads: it is updated by gradient descent on each token, with momentum and a forgetting term, at inference time as well as during training.

pip install titans-unofficial

import torch
from titans import TitansMAC

model = TitansMAC(d_model=128, n_layers=2, n_heads=4)

x = torch.randn(2, 300, 128)  # (batch, sequence, d_model)
y = model(x)                  # same shape; the memory learns as it reads

The models work on embeddings, not token IDs, so they slot in where a transformer encoder stack would go. Add your own embedding layer and output head; the examples show how.

Variants

The paper describes three ways to connect the memory to attention. All three share the same constructor arguments and the same forward(x, is_training=True).

from titans import TitansMAC, TitansMAG, TitansMAL

mac = TitansMAC(d_model=128, n_layers=2, n_heads=4, chunk_size=128)
mag = TitansMAG(d_model=128, n_layers=2, n_heads=4, window_size=128)
mal = TitansMAL(d_model=128, n_layers=2, n_heads=4, window_size=128)

All three variants are causal: the output at a position depends only on that position and the ones before it, so they can be trained as language models.

Working with the memory

Every sequence in a batch gets its own memory. It starts from weights the model has learned, is updated token by token as the sequence is read, and is thrown away when forward returns. Calls are therefore independent: the same input always gives the same output, and the model's parameters never change at inference time.

y = model(x)                     # each sequence reads from and writes to its own memory
y = model(x, is_training=False)  # reads only: the memory stays at its learned starting point

is_training controls the memory only and is independent of model.train() and model.eval(), which still switch dropout as usual. Leave it on at inference time if you want the test-time learning the paper is about.

To feed one long input in several calls, keep the memory between them:

model = TitansMAC(d_model=128, n_layers=2, n_heads=4, persist_memory=True)

for piece in pieces_of_a_long_document:  # each of shape (batch, tokens, d_model)
    y = model(piece)                     # remembers the earlier pieces

model.reset_memory()                     # before the next document or a new batch size

The memory updates are differentiable, so training with an ordinary optimizer also trains the memory's starting weights, its key, value and query projections, and the gates that set its learning rate, momentum and forgetting.

Differences from the paper

This is an alpha-stage reimplementation. It follows the paper's equations for the memory (Section 3.1) and its three architectures (Section 4), with these departures:

• The gradient the memory learns from is clipped to a norm of 1 per sequence. Without it the updates diverge on inputs of ordinary scale. Pass max_grad_norm=None to NeuralMemoryModule to turn it off.
• The forgetting gate is shifted so that an untrained memory keeps about 98% of its weights per token.
• The memory is written one token at a time. The chunked, parallel training of Section 3.2 is not implemented, so long sequences are slow.
• MAC keeps the model causal by letting a token see what the memory recalled for itself and for earlier tokens only.
• MAL shares one set of memory parameters between its layers, each with its own memory state, and places the memory after attention in each block.
• The depthwise-separable convolutions and SiLU activations of Section 4.4 are not used.

No results from the paper have been reproduced with this code.

Examples

The example scripts are not part of the installed package. Run them from a clone:

git clone https://github.com/itsmesherry/titans-unofficial.git
cd titans-unofficial
pip install -e ".[examples]"
cd examples
python simple_example.py

Citation

This repository is not affiliated with the authors. If you use it, cite the paper:

@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}
}

Project

Changelog · Contributing · Roadmap · MIT License