PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Sarah Pan, Ali Behrouz

arXiv:2610.19210v1cs.CLcs.AI

TL;DR

Long-context inference makes transformer KV caches grow, while efficient recurrent alternatives can be harder to scale through training. Titan-Llama retrofits pre-trained attention models with test-time neural memory, reporting faster prefill and lower steady-state decode memory usage.

  • Problem

    Long-context transformer inference requires increasingly large KV caches, while recurrent alternatives reduce inference complexity but limit token-wise parallelism during training.

    Source: §1 · p. 1

  • Method

    Titan-Llama post-trains external neural memory modules to supply historical attention contributions while preserving local attention and injecting memory reads as residuals into a frozen backbone.

    Source: §1 · p. 1, §3.2 · p. 4, §3.3 · p. 5

  • Results

    More than 20% faster prefill than full attention at long prefill lengths with five converted layers, and 24% lower steady-state decode memory usage than full attention with the KV cache.

    Source: §4.3 · p. 7

  • Takeaways & Limitations

    The results demonstrate practical inference-efficiency benefits and that test-time training with key–value binding can learn language modeling and extend context for certain tasks.

    Source: §5 · p. 8

    Open questions include scaling NMM size, changing its architecture, understanding memory retrieval, and testing inner-loop objectives beyond KV binding.

    Source: §5 · p. 8

Abstract

from arXiv · show

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}

1 INTRODUCTION

Titan-Llama adapts pre-trained transformers for long-context inference by preserving local attention and using external neural memory for historical context. It targets both fixed-size context compaction and the practical costs of updating and reading memory.

1 INTRODUCTION

  • Transformers parallelize training but accumulate growing KV caches at inference, while recurrent alternatives can require training a new model instead of reusing a pre-trained attention backbone.The paper motivates retaining distant information without storing every historical key and value, using full attention for local windows and fixed-size memory for older tokens.Source: §1 · p. 1
  • Titan-Llama post-trains external modules to replace historical attention contributions while preserving local attention, using test-time gradients to store key–value bindings and retrieve historical residuals.The retrieved value is injected into the output stream; the pre-trained backbone stays frozen, with external weights amounting to 1% of the total parameter footprint.Source: §1 · pp. 1–2, Figure 1 · p. 2
  • At test time, a self-supervised key–value reconstruction loss updates the NMM as chunks arrive, without labels, though reconstruction does not guarantee exact historical attention.The authors keep the pre-trained backbone frozen while training external weights amounting to 1% of the total parameter footprint.Source: §1 · p. 2
  • The design asks how to post-train an LLM for fixed-size context compaction and how to account for memory-update costs during prefill and storage and reading costs during decode.These are the two stated design considerations.Source: §1 · p. 2

2 RELATED WORK

Related work addresses long dependencies through hierarchical architectures, model retrofits, and test-time training. Titan-Llama is positioned as a post-training approach that uses a self-supervised inner objective to substitute memory for long-range attention.

Hierarchical Architectures

  • Hierarchical models separate local and distant relationships, but some recompute global attention scores at each generated token; this work stores historical context in an NMM during prefill.Other hierarchical approaches downsample the sequence, whereas this method follows chronological context arrival.Source: §2 · p. 2

Retrofitting Pre-trained Models

  • Retrofit approaches vary in how much of the original model they preserve and how they train replacements; Titan-Llama fine-tunes a small external weight set to modulate a frozen backbone.The cited alternatives include distillation for full architecture conversion and fine-tuning modular components.Source: §2 · pp. 2–3

Test-time Training

  • Test-time training updates weights during inference, and some methods retain the architecture while applying gradients to MLP layers, decreasing throughput; this work replaces long-range attention with RNNs.The paper describes the replacement as improving throughput and average memory footprint for long-context tasks.Source: §2 · p. 3
  • Titan-Llama is conceptually inspired by Titans and methodologically by LaCT, but focuses on retrofitting LLMs after pre-training and uses a self-supervised inner objective rather than final-logit prediction.The passage contrasts its post-training focus with pre-training or video work.Source: §2 · p. 3
  • The Titans view treats attention as key-to-value mapping and simulates fast weights through test-time training in an explicit MLP, while linear attention collapses that mapping into a single linear map.The passage describes the latter as reducing representational capabilities.Source: §2 · p. 3

3 METHOD

Titan-Llama approximates historical attention with recurrent memory while retaining local attention, then updates that memory at test time. Its architecture uses blockwise attention and a query-dependent gate to combine local and retrieved historical outputs.

3.1 APPROXIMATING HISTORICAL ATTENTION AS A RESIDUAL

  • Historical and local attention contributions are partitioned, then recombined by predicting the historical output and its mixing weight.The historical term is replaced by a recurrent-memory read; the paper cautions that this approximation does not imply lossless compression of historical attention.Source: §3.1 · pp. 3–4
  • Titan-Llama with 1024-token segments beats full attention in Pile per-token negative log-likelihood from 16K through 128K, and its CUDA-graph decode is faster than measured baselines across reported lengths.Figure 2 reports median single-token decode latency on an A100 at batch size one after prefill and warmup.Source: Figure 2 · p. 4

3.2 READING AND WRITING AT TEST-TIME

  • The NMM retrieves a historical attention residual that is blended with local attention through a learned per-channel gate.The retrieved output is combined with local attention as ˆλ ˆOhist + (1 − ˆλ)Olocal.Source: §3.2 · pp. 4–5
  • A squared-error key–value reconstruction loss updates the NMM on each chunk without labels, but serves only as a surrogate for reproducing historical attention.Updates use gradient descent with learned momentum, decay, and step-size parameters; fixed-size token batches and offline multi-chunk batching are used for efficiency.Source: §3.2 · p. 5

3.3 ARCHITECTURE

  • Blockwise attention separates long-term relationships and provides a threshold for NMM hidden-state updates and sparse-computation gains.The paper notes that earlier segmented-attention layers make historical keys and values differ from those of the full-attention model.Source: §3.3 · p. 5
  • A linear layer predicts query-dependent ˆλ to interpolate between local attention and the retrieved historical output.The gate is intended to approximate the attention-mass weighting of local and historical contributions.Source: §3.3 · p. 5
  • Larger chunks improve QuALITY and LongHealth performance, yield mixed RULER trends, and increase prefill latency; Titan-Llama matches full attention on 32K pass-key retrieval.Table 1 also reports the strongest non-full-attention UUID scores at 8K–32K.Source: Figure 3 · p. 5, Table 1 · p. 6

4 EXPERIMENTS

The experiments test whether post-training external neural memory modules can recover long-context performance while keeping a frozen backbone and improving inference efficiency. Results show strong but task-dependent recall and inference gains, with chunk-size and peak-memory tradeoffs.

4 EXPERIMENTS

  • The experiments convert the first four Llama-3.2-1B-Instruct layers to blockwise attention with an NMM per KV head, leaving the remaining layers at full attention.The setup tests memory recovery in converted layers without requiring every layer to use the new mechanism; external modules are optimized while the pretrained backbone stays frozen.Source: §4 · p. 6

4.1 POST-TRAINING FOR LONG CONTEXT REASONING

  • With 1024-token segments, Titan-Llama’s language-modeling loss delta stays below zero from 16K through 128K tokens, while measured Mamba-2, Gated DeltaNet, and LaCT incur larger penalties at longer contexts.Training sequences are at most 16K tokens, so measurements beyond that length test extension beyond the training length.Source: §4.1 · p. 6

4.2 RULER EVALUATION

  • Table 2 reports recovery of much of the accuracy lost to segmented attention, especially for single-needle and multi-key retrieval up to 16K, while multi-value retrieval remains weaker.Titan-Llama exceeds full attention on some multi-key and frequent-word extraction settings, but performance declines at 32K, beyond the 16K training length.Source: Table 2 · p. 7, §4.2 · p. 7
  • Titan-Llama matches full attention on pass-key retrieval, exceeds it on UUID recall at 8K and 16K, and retains full accuracy on 32K pass-key retrieval where measured efficient-model baselines decline.The broader RULER results show that recall performance varies by task type.Source: §4.2 · p. 7

4.3 INFERENCE GAINS

  • At long prefill lengths, Titan-Llama is more than 20% faster than full attention with five converted layers and uses 24% less steady-state decode memory, though partial hybrids can raise peak prefill memory.Titan-Llama also has lower median single-token latency than measured full-attention and efficient-model baselines across reported context lengths.Source: §4.3 · pp. 7–8, Figure 4 · p. 8

4.4 SCALING CHUNK SIZE

  • Larger chunks improve QuALITY and LongHealth accuracy but have mixed effects on RULER, helping some task-length combinations while reducing accuracy on others, including needle retrieval.Longer chunks also increase prefill latency; chunk size balances local attention, memory-update frequency, and inference cost.Source: §4.4 · p. 8

5 DISCUSSION

The work uses test-time training to treat updated weights as a memory mechanism for replacing long-range attention calculations in a pre-trained backbone. Its results support this approach while leaving NMM scaling and architecture as open questions.

5 DISCUSSION

  • The main contribution is using updated weights as memory to substitute for long-range attention calculations in a pre-trained backbone.The paper distinguishes this focus from prior work exploring test-time training on MLPs to improve language-model capabilities.Source: §5 · p. 8
  • The results corroborate test-time training for language modeling and report practical inference-efficiency benefits.The authors also report that TTT using KV binding learns language modeling and extends context windows for certain tasks.Source: §5 · p. 8
  • Scaling NMM size, changing its architecture, understanding memory retrieval, and testing other inner-loop objectives remain open questions.The authors identify KV binding as the current inner-loop objective.Source: §5 · p. 8

A.1 DATASETS

Training uses packed FineWeb sequences, followed by supervised fine-tuning on LongAlign and LoongRL, then QuALITY training data for multiple-choice adaptation.

A.1 DATASETS

  • Training uses one billion packed FineWeb tokens with a maximum sequence length of 16K, followed by supervised fine-tuning on LongAlign and LoongRL and QuALITY multiple-choice adaptation.The final fine-tuning stage uses the QuALITY training set.Source: §A.1 · p. 13

A.2 HYPERPARAMETERS

The appendix reports the NMM configuration and training batch setup, and describes an inference-cost plot for a 512-token local window.

A.2 HYPERPARAMETERS

  • The first four Llama layers use one three-layer, 64-hidden-dimension GELU NMM per KV head, adding 17.23M trainable parameters.These added components make up the reported trainable parameter total.Source: §A.2 · p. 13
  • Initial training uses batch size 64, while supervised fine-tuning uses batch size 128; the inner-loss batch size equals the chunk size.The evaluated chunk sizes are 512, 1024, and 2048 tokens, and compute constraints prevented a broad search over these choices.Source: §A.2 · p. 13
  • With a 512-token window, restricting more later attention layers limits KV-cache growth and makes decode throughput less sensitive to prefill length.The figure presents inference cost when later attention layers are restricted to local windows.Source: Figure 5 · p. 13

B INFERENCE IMPLEMENTATION

The inference implementation uses CUDA Graph replay for single-token decoding and a preallocated KV ring buffer with in-place index updates. Table 3 reports that NMM improves accuracy over segmented attention alone, while CUDA Graphs do not eliminate kernel execution.

B INFERENCE IMPLEMENTATION

  • For batch-size-one decoding, CUDA Graph reduce-overhead mode replays captured single-token computations, with compilation and capture excluded from timed steps.The implementation uses separate captured paths for ordinary tokens and tokens completing a memory-update chunk.Source: B INFERENCE IMPLEMENTATION · p. 13
  • A preallocated KV ring buffer keeps attention operands at fixed shapes while cache positions and memory-store indices update in place on the GPU.The implementation captures separate paths for ordinary tokens and chunk-completing tokens.Source: B INFERENCE IMPLEMENTATION · p. 13
  • Table 3 reports that neural memory improves multiple-choice accuracy on both tasks over segmented attention alone, but does not recover the full-attention baseline.The table covers settings where later attention layers are also restricted to local windows.Source: Table 3 · p. 14
  • CUDA Graphs reduce launch overhead but still execute underlying kernels; fusing memory reads, updates, and output blending could further reduce traffic and overhead if chunk-boundary ordering is preserved.The proposed fused implementation must preserve the ordering of reads and updates at chunk boundaries.Source: B INFERENCE IMPLEMENTATION · p. 14

C.1 RESTRICTING THE VISIBILITY OF LATER ATTENTION LAYERS

Restricting later attention layers to local windows lowers inference cost but leaves neural memory as the only cross-segment route and reduces accuracy. Chunk-size results vary by task and context length, with no single setting best across RULER tasks.

C.1 RESTRICTING THE VISIBILITY OF LATER ATTENTION LAYERS

  • Later layers retain full attention in the main configuration, but restricting them to local windows makes neural memory the only route for information from completed segments.Source: §C.1 · p. 14
  • On QuALITY, adding the NMM raises accuracy from 26.8 to 35.2; on LongHealth, it raises accuracy from 15.8 to 32.3, though both remain below full attention.Source: §C.1 · p. 14
  • On most RULER tasks, accuracy remains low, and neither restricted configuration reliably retrieves information across segments.Source: §C.1 · p. 14, Table 4 · p. 14

C.2 CHUNK SIZE AND RULER PERFORMANCE

  • Chunk size sets both local-attention span and tokens per memory update; 512 tokens is strongest on most multi-key and multi-value retrieval settings, but larger chunks improve some individual results.The results do not support one best chunk size across RULER tasks.Source: §C.2 · p. 14, Table 5 · p. 15

Related papers

Loading 2610.19210v1…