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)

Working with the memory

The memory is state that lives on the model and changes during forward.

y = model(x)                     # reads from memory and writes to it
y = model(x, is_training=False)  # reads only; the memory stays as it is

model.reset_memory()             # forget everything, e.g. between documents
model.to_device("cuda")          # move the model together with its memory

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.

Differences from the paper

This is an alpha-stage reimplementation, and it departs from the paper in ways that matter if you plan to reproduce its results:

• One memory is shared by the whole batch and persists across forward calls, so outputs depend on batch composition and call order. Use a batch size of 1 and call reset_memory() between sequences for clean behaviour.
• reset_memory() draws new random memory weights; it does not return to a learned starting point, so results after a reset are not reproducible without fixing the random seed.
• The memory's key and value projections and its learning-rate, momentum and forgetting gates are not trained; they keep their initial values.
• In MAG and MAL, persistent tokens drop out of the attention window for positions beyond window_size.
• MAC queries the memory with the average of a chunk, giving one recalled token per chunk.

These are tracked in the roadmap.

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