cs.CLOct 13, 2026

Titan-Llama: Retrofitting Pre-trained LLMs with External Neural Memory Modules

Authors: Sarah Pan, Ali Behrouz

Organizations: Massachusetts Institute of Technology · Cornell University

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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 GitHub

Figures & tables

Figure 1: Titan-Llama architecture at a glance. Left: The inner loss generates gradients at test-time w.r.t. MLP weights, teaching it to memorize historical attention key-value bindings. Right: The query associated with a local hidden state is used to retrieve a value from the NMM. This value, approximating the historical attention output, is then used as a residual and injected into the lambda-gated output stream.
Figure 2: Long-context language modeling and decode latency. Left. Change in per-token negative log-likelihood relative to full attention, measured on Pile. Titan-Llama with 1024-token segments beats full attention from 16K through 128K. Right. Median single-token decode latency on an A100 at batch size one, after prefill and warmup (lower is better). Titan-Llama’s CUDA-graph-enabled decode is faster than the measured full-attention, Mamba-2, Gated DeltaNet, and In-Place TTT implementations across the reported lengths ( Dao and Gu, 2024 ; Yang et al., 2025 ) .
Figure 3: The effects of scaling chunk size on downstream performance and efficiency. Left. Scaling chunk size improves performance on long-context reasoning tasks like QuALITY and LongHealth. Middle. RULER trends are mixed. Longer chunks help at intermediate lengths but do not consistently improve associative recall. Detailed RULER results can be found in Table 5 in the Appendix. Right. Larger chunks increase prefill latency as memory bandwidth saturates.
 S-NIAH-1 (pass-key)S-NIAH-2 (number)S-NIAH-3 (UUID)Relative decode time ↓\downarrow
Model8K16K32K8K16K32K8K16K32KMean
Full attention100.0100.0100.099.0100.0100.064.064.083.01.00
Mamba-299.049.026.099.043.028.077.036.024.02.70
Gated DeltaNet100.050.026.0100.043.027.091.045.023.02.59
LaCT98.043.022.0100.043.027.074.034.023.045.8†45.8^{\dagger}
TTT-E2E100.046.024.099.043.028.077.044.024.0115†115^{\dagger}
Table 1: RULER S-NIAH accuracy (%) across context lengths. Titan-Llama matches full attention on 32K pass-key retrieval and achieves the strongest non-full-attention UUID scores at 8K–32K. Relative decode time ( ↓\downarrow ) averages next-token latency divided by full-attention latency at 8K, 16K, and 32K on A100. †\dagger marks timings of publicly available code that may not utilize caching.
ModelSingle needleMulti-keyMulti-query
 4K8K16K32K4K8K16K32K4K8K16K32K
Full attention100.099.3100.099.382.074.063.357.399.594.094.083.5
Segmented attn.10.70.00.00.00.00.00.00.00.00.00.00.0
Titan-Llama97.396.088.078.791.388.066.739.396.590.570.534.5
ModelMulti-valueFWE/CWEHotpotQA
 4k8k16k32k4k8k16k32k4k8k16k32k
Table 2: RULER accuracy by task family and context length. Titan-Llama recovers much of the accuracy lost to attention segmentation, although the gap to full attention widens on some tasks at longer contexts. Multi-key and FWE/CWE results exceed the full attention baseline.
Figure 4: Increasing the proportion of Llama 3.2 1B layers retrofitted with NMMs increases token decode throughput, decreases prefill latency, and also decreases memory footprint during token decoding. Peak memory footprint during prefill for partial Titan-Llama hybrids exceeds full-attention Llama while retrofitting all layers leads to lower peak memory footprint. Benchmarking was performed on a single 80GB NVIDIA A100 GPU.
Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Inference cost when later attention layers are restricted to local windows. The plot shows a 512-token window. Restricting more layers limits the growth of the KV cache and makes decode throughput less sensitive to prefill length.
ModelQuALITYLongHealth
Full attention (Llama-3.2-1B-Instruct)43.34±1.08\mathbf{43.34}\pm 1.0854.25±2.49\mathbf{54.25}\pm 2.49
Local-visibility variants
Segmented attention, no NMM26.815.8
Segmented attention, with NMM35.232.3
Appendix
Table 3: Multiple-choice accuracy when the later attention layers are also restricted to local windows. Neural memory improves both tasks relative to segmented attention alone, but does not recover the full-attention baseline.
ModelSingle needleMulti-keyMulti-queryMulti-value
 4k8k16k32k4k8k16k32k4k8k16k32k4k8k16k32k
Segmented attn.0.03.31.30.70.31.30.00.02.01.80.80.31.03.82.00.0
Titan-Llama1.04.74.01.71.32.30.71.03.52.81.01.31.03.81.81.0
ModelFWE/CWEHotpotQASQuADVT
 4k8k16k32k4k8k16k32k4k8k16k32k4k8k16k32k
Segmented attn.22.710.14.22.717.09.01.02.031.41.00.00.00.40.00.00.0
Appendix
Table 4: RULER accuracy by task family and context length when later attention layers are restricted to local windows. Results use cached autoregressive decoding. Accuracy remains low on most tasks, indicating that the NMM does not recover the information available through the later full-attention layers in the main configuration.
Chunk sizeSingle needleMulti-keyMulti-queryMulti-value
 4k8k16k32k4k8k16k32k4k8k16k32k4k8k16k32k
51297.396.088.078.791.388.066.739.396.590.570.534.588.081.063.053.0
102495.392.791.078.781.379.363.036.392.885.051.242.287.068.849.038.2
204890.392.791.751.766.374.067.032.391.388.881.219.287.579.282.225.0
Chunk sizeFWE/CWEHotpotQASQuADVT
 4k8k16k32k4k8k16k32k4k8k16k32k4k8k16k32k
Appendix
Table 5: RULER accuracy by task family, context length, and chunk size. Bold marks the best result among the three chunk sizes for each task and length.
More in this topic

Long-Context Language Model Inference

Long-context language model inference makes processing and generating over lengthy prompts faster and more memory-efficient.

See all 221 papers

Explore similar work

CardsList
May 12, 2026

Training-Inference Consistent Segmented Execution for Long-Context LLMs

Transformer-based large language models face severe scalability challenges in long-context generation due to the computational and memory costs of full-context attention. Under practical computation and memory constraints, many inference-efficient long-context methods improve efficiency by adopting bounded-context or segment-level execution only during inference, while continuing to train models under full-context attention, resulting in a mismatch between training and inference execution and state-transition semantics. Based on this insight, we propose a training-inference consistent segment-level generation framework, in which training and inference follow the same segment-level forward execution semantics. During training, consistency with inference is enforced by restricting gradient propagation to KV states carried over from the immediately preceding segment, while permitting head-specific access to past KV states during the forward pass without involving them in gradient propagation. Across long-context benchmarks, our approach achieves performance comparable to full-context attention, while achieving competitive latency-memory trade-offs against strong inference-efficient baselines, and substantially improving scalability at very long context lengths (e.g., approximately 6x lower peak prefill memory at 128K compared to full-context attention with FlashAttention).

Xianpeng Shang, Jiang Li, Zehua Duo +2

College of Computer Science, Inner Mongolia University, Hohhot 010021, China · National & Local Joint Engineering Research Center of Intelligent Information Processing Technology for Mongolian, Hohhot 010021, China · Inner Mongolia Key Laboratory of Multilingual Artificial Intelligence Technology, Hohhot 010021, China +1

Self-AttentionLong-Context Language ModelingLong-Context InferenceEfficient Language Model Inference
AskPDF
Mar 18, 2026

Learning When to Attend: Conditional Memory Access for Long-Context LLMs

Language models struggle to generalize beyond pretraining context lengths, limiting long-horizon reasoning and retrieval. Continued pretraining on long-context data can help but is expensive due to the quadratic scaling of Attention. We observe that most tokens do not require (Global) Attention over the entire sequence and can rely on local context. Based on this, we propose L2A (Learning To Attend), a layer that enables conditional (token-wise) long-range memory access by deciding when to invoke global attention. We evaluate L2A on Qwen 2.5 and Qwen 3 models, extending their effective context length from 32K to 128K tokens. L2A matches the performance of standard long-context training to within 3% while skipping Global Attention for ∼\sim80% of tokens, outperforming prior baselines. We also design custom Triton kernels to efficiently implement this token-wise conditional Attention on GPUs, achieving up to ∼\sim2×\times improvements in training throughput and time-to-first-token over FlashAttention. Moreover, L2A enables post-training pruning of highly sparse Global Attention layers, reducing KV cache memory by up to 50% with negligible performance loss. Our code is released under Apache 2.0 at https://github.com/awslabs/hybrid-model-factory/tree/main/examples/research/L2A.

Sakshi Choudhary, Aditya Chattopadhyay, Luca Zancato +4

Department of Electrical and Computer Engineering, Purdue University, USA · AWS Agentic AI, USA

Self-AttentionLong-Context Language ModelingLLM Inference AccelerationSparse AttentionAttention MechanismsLong-Context Reasoning in Language Models
AskPDF
Aug 18, 2026

MoNe: Modular Neural Memory for Efficient Long Context Inference

We present MoNe, a lightweight modular neural memory that attaches to any frozen pretrained Transformer to enable long-context inference without retraining. MoNe reads context in fixed-size segments via test-time learning of fast-weight neural memory networks with layer-localized gradient updates; at inference, the memory generates keys and values from the query tokens alone, with no context tokens re-read. This two-phase design decouples inference cost from context length, achieving O(N)O(N) preprocessing and O(1)O(1) query cost with peak GPU memory that does not grow with NN. At 128K tokens, MoNe reduces both compute and peak GPU memory by approximately 80% compared to ICL with only 6.4% parameter overhead. MoNe generalizes to context lengths far beyond the backbone's native window, achieving strong performance on needle-in-a-haystack and word extraction benchmarks from RULER, where ICL degrades sharply.

Wonguk Cho, Kyubyung Chae, Tribhuvanesh Orekondy +6

Qualcomm AI Research

Long-Context Language Model InferenceMemory-Augmented Language ModelsLLM Inference AccelerationTest-Time Training
AskPDF