Transformer models retrofitted with memory modules for faster long context
Titan-Llama: Retrofitting Pre-trained LLMs with External Neural Memory Modules
Computation and LanguageArtificial Intelligence
Summary
Language models can struggle to remember and efficiently use long pieces of text because their usual method requires a lot of computing power. The authors introduce Titan-Llama, a way to improve these models after they are already trained by adding special memory parts that keep important information for longer. This new approach makes the models more accurate on reading long documents and faster to use when working with medium-length contexts. This work helps adapt existing language models to work better with long texts without retraining from scratch.
What this means in practice
- •For language model engineers: Update existing attention-based models with hierarchical memory to improve accuracy and speed on tasks requiring long text understanding without retraining.
- •For ai infrastructure teams: Reduce computational costs and memory requirements during long-context inference by retrofitting pre-trained transformers with compressed recurrent memory modules.
Authors
Sarah Pan, Ali Behrouz
Abstract
Transformers are efficient to pre-train at scale because attention permits parallel computation across sequence positions, but long-context inference requires increasingly large KV caches. Efficient recurrent alternatives reduce the asymptotic complexity of inference, but their sequential dependence limits token-wise parallelism during training, making them harder to scale. Further, the current paradigm of extending long context tasks creates competing notions of retaining information precisely and keeping long-context inference efficient. To address this gap, we introduce Titan-Llama, a post-training method for retrofitting attention-based language models with a hierarchical memory architecture that updates parameters at test-time. Titan-Llama applies a blockwise attention mask and preserves full-fidelity attention over recent tokens while compressing historical context into recurrent, fixed-size neural memory modules (NMMs). Our formulation allows pre-trained attention layers to be replaced without training a new language model from scratch. On Llama-3.2-1B-Instruct, Titan-Llama raises accuracy over segmented attention windows from 23.87\% to 40.51\% on QuALITY and from 14.50\% to 40.75\% on LongHealth. Inference benchmarks show more than 20\% faster prefill, 37\% faster decode at medium context lengths, and 24\% lower steady-state decode memory usage in the evaluated configurations. Together, these results demonstrate that incorporating memory architecture mechanisms that utilize test-time training effectively modulates long-context behavior of pre-trained models, offering a practical route to adapting pre-trained transformers for efficient long-context inference. Our code is released on Github \href{https://github.com/sarahpannn/titans-pytorch/tree/main}{GitHub}