A transformer memory module that updates itself during inference

What it does

titans-pytorch is an unofficial reimplementation of the Titans paper, which bolts a learnable neural memory module onto standard transformers. The gist is test-time training: instead of freezing weights during inference, the model updates its memory to absorb long contexts on the fly. The repo provides a standalone NeuralMemory layer and a full MemoryAsContextTransformer that weaves memory tokens into a local-attention backbone. The author also plans to tinker with deeper memory architectures beyond the paper’s simple 1–4 layer MLP if early experiments continue to look promising.

The interesting bit

The MAC (Memory as Context) configuration mixes persistent memory tokens, long-term memory tokens, and a sliding local attention window, letting the model decide what to remember without ballooning the KV cache. A tunable chunk_size parameter hints that memory updates happen in chunks, offering a dial to trade memory pressure for sequence-level performance. The README is admirably candid about the whole thing being unofficial and exploratory.

Key highlights

• Standalone NeuralMemory module plus a full MemoryAsContextTransformer ready for token-level training and sampling.
• Configurable segment_len for local attention, plus separate slots for persistent and long-term memory tokens.
• Ships with a train_mac.py experiment script and a linked Colab for quick smoke tests.
• Author is already eyeing architectural variants beyond the paper’s shallow MLP memory design.

Caveats

• The README offers no training curves, benchmarks, or throughput numbers, so you are buying the architecture on spec.
• The author explicitly hedges that deeper memory explorations will only continue if the basic version works well, which it may not.

Verdict

Worth a look if you are researching long-context transformers or need a clean, hackable PyTorch sketch of test-time training. Move along if you require a production-hardened memory layer with proven performance data.