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

Sarah Pan; Ali Behrouz

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

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

Panel of Figure 1 — 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.

Model	8K	16K	32K	8K	16K	32K	8K	16K	32K	Mean
	S-NIAH-1 (pass-key)			S-NIAH-2 (number)			S-NIAH-3 (UUID)			Relative decode time $\downarrow$
Full attention	100.0	100.0	100.0	99.0	100.0	100.0	64.0	64.0	83.0	1.00
Mamba-2	99.0	49.0	26.0	99.0	43.0	28.0	77.0	36.0	24.0	2.70
Gated DeltaNet	100.0	50.0	26.0	100.0	43.0	27.0	91.0	45.0	23.0	2.59
LaCT	98.0	43.0	22.0	100.0	43.0	27.0	74.0	34.0	23.0	$45.8^{\dagger}$
TTT-E2E	100.0	46.0	24.0	99.0	43.0	28.0	77.0	44.0	24.0	$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.

Model	Single needle				Multi-key				Multi-query
	4K	8K	16K	32K	4K	8K	16K	32K	4K	8K	16K	32K
Full attention	100.0	99.3	100.0	99.3	82.0	74.0	63.3	57.3	99.5	94.0	94.0	83.5
Segmented attn.	10.7	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0
Titan-Llama	97.3	96.0	88.0	78.7	91.3	88.0	66.7	39.3	96.5	90.5	70.5	34.5
Model	Multi-value				FWE/CWE				HotpotQA
	4k	8k	16k	32k	4k	8k	16k	32k	4k	8k	16k	32k

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.

Local-visibility variants
Model	QuALITY	LongHealth
Full attention (Llama-3.2-1B-Instruct)	$\mathbf{43.34}\pm 1.08$	$\mathbf{54.25}\pm 2.49$
Segmented attention, no NMM	26.8	15.8
Segmented attention, with NMM	35.2	32.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.

Model	Single needle				Multi-key				Multi-query				Multi-value
	4k	8k	16k	32k	4k	8k	16k	32k	4k	8k	16k	32k	4k	8k	16k	32k
Segmented attn.	0.0	3.3	1.3	0.7	0.3	1.3	0.0	0.0	2.0	1.8	0.8	0.3	1.0	3.8	2.0	0.0
Titan-Llama	1.0	4.7	4.0	1.7	1.3	2.3	0.7	1.0	3.5	2.8	1.0	1.3	1.0	3.8	1.8	1.0
Model	FWE/CWE				HotpotQA				SQuAD				VT
	4k	8k	16k	32k	4k	8k	16k	32k	4k	8k	16k	32k	4k	8k	16k	32k
Segmented attn.	22.7	10.1	4.2	2.7	17.0	9.0	1.0	2.0	31.4	1.0	0.0	0.0	0.4	0.0	0.0	0.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 size	Single needle				Multi-key				Multi-query				Multi-value
	4k	8k	16k	32k	4k	8k	16k	32k	4k	8k	16k	32k	4k	8k	16k	32k
512	97.3	96.0	88.0	78.7	91.3	88.0	66.7	39.3	96.5	90.5	70.5	34.5	88.0	81.0	63.0	53.0
1024	95.3	92.7	91.0	78.7	81.3	79.3	63.0	36.3	92.8	85.0	51.2	42.2	87.0	68.8	49.0	38.2
2048	90.3	92.7	91.7	51.7	66.3	74.0	67.0	32.3	91.3	88.8	81.2	19.2	87.5	79.2	82.2	25.0
Chunk size	FWE/CWE				HotpotQA				SQuAD				VT
	4k	8k	16k	32k	4k	8k	16k	32k	4k	8k	16k	32k	4k	8k	16k	32k

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.

Long-Context Language Model Inference

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

See all 224 papers

Training-Inference Consistent Segmented Execution for Long-Context LLMs
May 12, 2026Xianpeng Shang, Jiang Li, Zehua Duo +2Self-Attention Long-Context Language Modeling
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
Learning When to Attend: Conditional Memory Access for Long-Context LLMs
Mar 18, 2026Sakshi Choudhary, Aditya Chattopadhyay, Luca Zancato +4Self-Attention Long-Context Language Modeling
Department of Electrical and Computer Engineering, Purdue University, USA · AWS Agentic AI, USA
MoNe: Modular Neural Memory for Efficient Long Context Inference
Aug 18, 2026Wonguk Cho, Kyubyung Chae, Tribhuvanesh Orekondy +6Long-Context Language Model Inference Memory-Augmented Language Models
Qualcomm AI Research

May 12, 2026J/K move · Enter open · S save

Training-Inference Consistent Segmented Execution for Long-Context LLMs

Xianpeng Shang, Jiang Li, Zehua Duo, Qianyi Cai+1

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-Attention Long-Context Language Modeling Long-Context Inference Efficient Language Model Inference

Open paper AskPDF arXiv

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).

Abstract

Long-Context Language Model Inference

Explore similar work

Training-Inference Consistent Segmented Execution for Long-Context LLMs

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

MoNe: Modular Neural Memory for Efficient Long Context Inference