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

# Titan-Llama: Retrofitting Pre-trained LLMs with External Neural Memory Modules
Source: [https://arxiv.org/html/2610.19210](https://arxiv.org/html/2610.19210)
Ali BehrouzAffiliation:Cornell UniversityEmail:[ab2947@cornell\.edu](mailto:)

###### 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](https://github.com/sarahpannn/titans-pytorch/tree/main)\.

## 1Introduction

Transformers are efficient to train at scale because attention permits parallel computation across sequence positions\. At inference time, however, retaining keys and values for an increasingly long context cause the KV cache to grow\. Linear attention and state\-space models offer more efficient recurrent inference, but adopting a different architecture generally requires training a new model rather than reusing a pre\-trained attention backbone\([Katharopoulos et al\., 2020](https://arxiv.org/html/2610.19210#bib.bib35);[Gu and Dao, 2024](https://arxiv.org/html/2610.19210#bib.bib37)\)\. Scaling large language models requires advances in both representational capacity and inference efficiency, which naturally motivates retrofitting the architecture of transformer\-based models to better suit them for inference\.

The growing KV cache raises the question of how should a pre\-trained transformer best retain information from distant tokens without necessarily storing every historical key and value\. Titans\([Behrouz et al\., 2024](https://arxiv.org/html/2610.19210#bib.bib18)\)represents distant information in an RNN dubbed the neural memory module \(NMM\), effectively parameterizing historical tokens as fast weight updates to a fixed\-size NMM\. This motivates our design in using full attention for local windows and representing older tokens in a fixed\-size MLP\. It is also worth noting that work in KV cache compaction similarly argues that exacting inference tasks do not necessarily require full fidelity on the KV cache\([Zhang et al\., 2023](https://arxiv.org/html/2610.19210#bib.bib20);[Zweiger et al\., 2026](https://arxiv.org/html/2610.19210#bib.bib21)\)\.

Our work combines these two ideas and develops a method to post\-train lightweight, external modules to improve inference efficiency\. We partition attention into local and historical contributions, preserve the local contribution, and train a neural memory module \(NMM\) to provide information from historical context\. At test time, a self\-supervised key\-value reconstruction loss updates the NMM as new chunks arrive\. This objective supplies a self\-supervised learning signal without labels, although reconstructing key\-value pairs does not guarantee exact reproduction of historical attention\. Notably, we keep the pre\-trained backbone frozen and are able to recover significant performance while training external weights amounting to only 1% of the total parameter footprint\.

Two main considerations guide our design:

1. 1\.How can we post\-train a pre\-trained LLM to use Titans as a means to achieve fixed\-size context compaction?
2. 2\.How can we make the resulting model efficient in practice, accounting for the cost of updating memory during prefill as well as storing and reading it during decode?

Historical windowsNMM∇θNMMℒmem\\nabla\_\{\\theta\_\{\\mathrm\{NMM\}\}\}\\mathcal\{L\}\_\{\\mathrm\{mem\}\}TTT inner loop\(update during prefill\)kkvvCurrent windowGQA Attentionqqkkvvλ\\lambda1−λ1\-\\lambdaLinear\+σ\+\\,\\sigmahth\_\{t\}qqO^hist\\hat\{O\}\_\{\\mathrm\{hist\}\}OlocalO\_\{\\mathrm\{local\}\}\+\+OutputO=λ​O^hist\+\(1−λ\)​OlocalO=\\lambda\\,\\hat\{O\}\_\{\\mathrm\{hist\}\}\+\(1\-\\lambda\)\\,O\_\{\\mathrm\{local\}\}

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\.
## 2Related Work

##### Hierarchical Architectures\.

Prior work has found success using hierarchical systems to model long dependencies, or more precisely, different mechanisms to model local and distant relationships\([Beltagy et al\., 2020](https://arxiv.org/html/2610.19210#bib.bib4);[Qiu et al\., 2020](https://arxiv.org/html/2610.19210#bib.bib5);[Kitaev et al\., 2020](https://arxiv.org/html/2610.19210#bib.bib6);[Bulatov et al\., 2022](https://arxiv.org/html/2610.19210#bib.bib10)\)\. Many of these approaches, however, do not take advantage of caching and must recompute global attention scores with every new token generated\([Patel et al\., 2024](https://arxiv.org/html/2610.19210#bib.bib8);[Ainslie et al\., 2020](https://arxiv.org/html/2610.19210#bib.bib9)\)\. Instead, we try to design a method that aligns with the abstraction of prefill and decode phases where historical context can be stored in the NMM via test\-time training at prefill time\.

Other hierarchical approaches apply downsampling operations to the sequence itself, whereas our time steps follow the chronological arrival of the context\.\([Hwang et al\., 2025](https://arxiv.org/html/2610.19210#bib.bib1);[Nawrot et al\., 2022](https://arxiv.org/html/2610.19210#bib.bib2);[Slagle, 2024](https://arxiv.org/html/2610.19210#bib.bib3)\)\.

##### Retrofitting Pre\-trained Models\.

Modifications to the attention architecture to leverage transformers’ training parallelism to fit inference\-time objectives have also been studied in prior literature\([Bick et al\., 2025](https://arxiv.org/html/2610.19210#bib.bib11);[Wang et al\., 2025a](https://arxiv.org/html/2610.19210#bib.bib12)\)\. These approaches differ in how much of the original model they preserve and how they train the replacement mechanism\. For instance, some approaches use distillation objectives to fully convert architectures from scratch while others fine\-tune modular components\([Eyuboglu et al\., 2025](https://arxiv.org/html/2610.19210#bib.bib7);[Kojima et al\., 2025](https://arxiv.org/html/2610.19210#bib.bib22)\)\. Our work aims to fine\-tune a small, external set of weights that, through training on next token prediction, can learn how to best modulate a frozen backbone\.

##### Test\-time Training\.

Test\-time training \(TTT\) methods apply gradient updates to a subset of weights at inference time, creating a meta\-objective for the slow weights to be able to “learn to learn”\([Sun et al\., 2020](https://arxiv.org/html/2610.19210#bib.bib13);[Sun et al\., 2025](https://arxiv.org/html/2610.19210#bib.bib14);[Finn et al\., 2017](https://arxiv.org/html/2610.19210#bib.bib15)\)\. Recent studies utilize TTT approaches to extend the effective context length or to use it to greater effect\([Bansal et al\., 2025](https://arxiv.org/html/2610.19210#bib.bib17)\)\. However, some methods keep the underlying architecture intact, with gradients applying to the MLP layers, decreasing throughput\([Feng et al\., 2026](https://arxiv.org/html/2610.19210#bib.bib16)\)\. The hierarchical approach we use replaces attention for longer\-range dependencies with RNNs, which actually improves throughput and average memory footprint for long\-context tasks

The TTT works that provide us the most inspiration and are most similar to ours are Titans\([Behrouz et al\., 2024](https://arxiv.org/html/2610.19210#bib.bib18)\)conceptually and LaCT\([Zhang et al\., 2025](https://arxiv.org/html/2610.19210#bib.bib19)\)methodologically, but we focus on the post\-training dynamics of retrofitting LLMs rather than pre\-training or video\. Importantly, we also use a self\-supervised inner objective rather than final logit prediction, as other TTT approaches do\([Tandon et al\., 2025](https://arxiv.org/html/2610.19210#bib.bib29);[Kuratov et al\., 2026](https://arxiv.org/html/2610.19210#bib.bib28);[Behrouz et al\., 2025](https://arxiv.org/html/2610.19210#bib.bib36)\)\.

Titans\.To motivate the attention KV binding used in Titans from first principles, consider the vanilla attention operation, which applies a softmax operation to𝐪⊤​𝐤\\mathbf\{q\}^\{\\top\}\\mathbf\{k\}before applying it to the𝐯\\mathbf\{v\}matrix\([Sun et al\., 2025](https://arxiv.org/html/2610.19210#bib.bib14)\)\. In abstract, we can view this as a two\-layer MLP where the softmax serves as an activation function and maps keys to values\([Han et al\., 2026](https://arxiv.org/html/2610.19210#bib.bib33);[Wen et al\., 2026](https://arxiv.org/html/2610.19210#bib.bib34)\)\. In linear attention, by replacing the softmax with a homomorphic function, we collapse the MLP into a single linear map, thereby reducing representational capabilities\([Katharopoulos et al\., 2020](https://arxiv.org/html/2610.19210#bib.bib35)\)\. In the Titans view, efficiency can be achieved by encoding the fast\-weight MLP explicitly in neural architecture and simulating fast weights via test\-time training \(TTT\)\([Behrouz et al\., 2024](https://arxiv.org/html/2610.19210#bib.bib18);[Behrouz et al\., 2025](https://arxiv.org/html/2610.19210#bib.bib36)\)\.

## 3Method

### 3\.1Approximating Historical Attention as a Residual

As shown in Figure[1](https://arxiv.org/html/2610.19210#S1.F1), our Titan\-Llama layers encode distant information in a fixed\-size RNN state and local information with attention, loosely resembling the distinction between lossy long\-term memories and vivid recent ones\. We now show that contributions to attention from local and historical contexts are easy to partition\. This allows us to approximate full attention by predicting the historical output and the mixing weight\.

For a query𝐪\\mathbf\{q\}at positionttand segment lengthkk, letbt=k⁡⌊t/k⌋b\_\{t\}=k\\lfloor t/k\\rfloordenote the start of the current segment\. We partition the available keys and values into historical positionsℋ=\{i:i<bt\}\\mathcal\{H\}=\\\{i:i<b\_\{t\}\\\}and local positionsℒ=\{i:bt≤i≤t\}\\mathcal\{L\}=\\\{i:b\_\{t\}\\leq i\\leq t\\\}\. With scoressi=𝐪⊤​𝐤i/ds\_\{i\}=\\mathbf\{q\}^\{\\top\}\\mathbf\{k\}\_\{i\}/\\sqrt\{d\}, define

Zhist=∑i∈ℋesi,Zlocal=∑i∈ℒesi,Z\_\{\\mathrm\{hist\}\}=\\sum\_\{i\\in\\mathcal\{H\}\}e^\{s\_\{i\}\},\\qquad Z\_\{\\mathrm\{local\}\}=\\sum\_\{i\\in\\mathcal\{L\}\}e^\{s\_\{i\}\},and the attention outputs normalized within each partition:

Ohist=∑i∈ℋesi​𝐯iZhist,Olocal=∑i∈ℒesi​𝐯iZlocal\.O\_\{\\mathrm\{hist\}\}=\\frac\{\\sum\_\{i\\in\\mathcal\{H\}\}e^\{s\_\{i\}\}\\mathbf\{v\}\_\{i\}\}\{Z\_\{\\mathrm\{hist\}\}\},\\qquad O\_\{\\mathrm\{local\}\}=\\frac\{\\sum\_\{i\\in\\mathcal\{L\}\}e^\{s\_\{i\}\}\\mathbf\{v\}\_\{i\}\}\{Z\_\{\\mathrm\{local\}\}\}\.
As in the blockwise formulation of exact attention by[Dao et al\. \(2022\)](https://arxiv.org/html/2610.19210#bib.bib27), we can partition the softmax numerator and denominator into historical and local contributions, withλ\\lambdarepresenting the fraction of attention mass assigned to history:

Oattn=∑i∈ℋesi​𝐯i\+∑j∈ℒesj​𝐯j∑i∈ℋesi\+∑j∈ℒesj,λ=ZhistZhist\+Zlocal\.O\_\{\\text\{attn\}\}=\\frac\{\\sum\_\{i\\in\\mathcal\{H\}\}e^\{s\_\{i\}\}\\mathbf\{v\}\_\{i\}\+\\sum\_\{j\\in\\mathcal\{L\}\}e^\{s\_\{j\}\}\\mathbf\{v\}\_\{j\}\}\{\\sum\_\{i\\in\\mathcal\{H\}\}e^\{s\_\{i\}\}\+\\sum\_\{j\\in\\mathcal\{L\}\}e^\{s\_\{j\}\}\},\\qquad\\lambda=\\frac\{Z\_\{\\mathrm\{hist\}\}\}\{Z\_\{\\mathrm\{hist\}\}\+Z\_\{\\mathrm\{local\}\}\}\.
![Refer to caption](https://arxiv.org/html/2610.19210v1/ppl_latency_e2e_comparison.png)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](https://arxiv.org/html/2610.19210#bib.bib38);[Yang et al\., 2025](https://arxiv.org/html/2610.19210#bib.bib39)\)\.Finally, we can construct the full attention output as a linear interpolation between the historical and local outputs with parameterλ\\lambda:

O=λ​softmax⁡\(𝐪⊤​𝐊histd\)​𝐕hist⏟historical attention\+\(1−λ\)​softmax⁡\(𝐪⊤​𝐊locald\)​𝐕local⏟local attentionO=\\lambda\\underbrace\{\\operatorname\{softmax\}\\left\(\\frac\{\\mathbf\{q\}^\{\\top\}\\mathbf\{K\}\_\{\\text\{hist\}\}\}\{\\sqrt\{d\}\}\\right\)\\mathbf\{V\}\_\{\\text\{hist\}\}\}\_\{\\text\{historical attention\}\}\+\(1\-\\lambda\)\\underbrace\{\\operatorname\{softmax\}\\left\(\\frac\{\\mathbf\{q\}^\{\\top\}\\mathbf\{K\}\_\{\\text\{local\}\}\}\{\\sqrt\{d\}\}\\right\)\\mathbf\{V\}\_\{\\text\{local\}\}\}\_\{\\text\{local attention\}\}\(1\)
O=λ​Ohist\+\(1−λ\)​Olocal\.\\boxed\{O=\\lambda O\_\{\\text\{hist\}\}\+\(1\-\\lambda\)O\_\{\\text\{local\}\}\.\}
Our approach replaces the historical termOhistO\_\{\\text\{hist\}\}with a readO^hist\\hat\{O\}\_\{\\text\{hist\}\}from recurrent memory and predicts its mixing weight\. In practice, we predict bothλ\\lambdaandOhistO\_\{\\text\{hist\}\}and approximate full attention outputs with the expression:λ^​O^hist\+\(1−λ^\)​Olocal\\hat\{\\lambda\}\\hat\{O\}\_\{\\text\{hist\}\}\+\(1\-\\hat\{\\lambda\}\)O\_\{\\text\{local\}\}\. Importantly, this identity does not imply that historical attention can be compressed without loss\.

### 3\.2Reading and Writing at Test\-time

Following the derivation in Section[3\.1](https://arxiv.org/html/2610.19210#S3.SS1), we seek to approximate historical attention, orO^hist\\hat\{O\}\_\{\\text\{hist\}\}, with the external neural memory module \(NMM\) as used in Titans\([Behrouz et al\., 2024](https://arxiv.org/html/2610.19210#bib.bib18)\)\. We make the design choice to blend the retrieved output with a learned gate as a residual added to the local attention output, prior to passing the result to the remaining layer computation\. This is similar to other implementations that add the value read from the RNN to the residual stream after attention\([Feng et al\., 2026](https://arxiv.org/html/2610.19210#bib.bib16)\)and use it to replace attention entirely in certain layers\([Tandon et al\., 2025](https://arxiv.org/html/2610.19210#bib.bib29)\)\.

We inject the RNN read as a residual into the attention output because this is a natural implementation of equation[1](https://arxiv.org/html/2610.19210#S3.E1)and aligns with the inner\-loss objective described below\. This also allows us to modulate an existing LLM backbone without changing any of the parameters already set\.

Read\.As shown in Figure[1](https://arxiv.org/html/2610.19210#S1.F1), we use the NMM to memorize the attention key–value mappings from historical attention windows and provide a residual to the current local\-attention output, with the two contributions balanced by a per\-channelλ\\lambdavector\. More specifically, for each query𝐪\\mathbf\{q\}, the NMM retrieves a historical attention residual𝐎^hist=ℳ⁡\(𝐪\)\\hat\{\\mathbf\{O\}\}\_\{\\text\{hist\}\}=\\mathcal\{M\}\(\\mathbf\{q\}\), while a learned gate is computed per channel as:

λ^=σ⁡\(𝐖λ⊤​RMSNorm⁡\(𝐪\)\+𝐛λ\)\.\\mathbf\{\\hat\{\\lambda\}\}=\\sigma\\left\(\\mathbf\{W\}\_\{\\lambda\}^\{\\top\}\\operatorname\{RMSNorm\}\(\\mathbf\{q\}\)\+\\mathbf\{b\}\_\{\\lambda\}\\right\)\.We combine this local attention output with the NMM retrieval as𝐎^full attention=λ^​𝐎^hist\+\(1−λ^\)​𝐎local\\hat\{\\mathbf\{O\}\}\_\{\\text\{full attention\}\}=\\hat\{\\lambda\}\\hat\{\\mathbf\{O\}\}\_\{\\text\{hist\}\}\+\(1\-\\hat\{\\lambda\}\)\\mathbf\{O\}\_\{\\text\{local\}\}\. We show this expression approximates𝐎full attention\\mathbf\{O\}\_\{\\text\{full attention\}\}in Section[3\.1](https://arxiv.org/html/2610.19210#S3.SS1)\. For tokens in a chunk, the NMM read uses the state accumulated from preceding chunks\. During inference, once a full chunk has been processed, its keys and values are used to compute the loss that updates the NMM\.

Write\.To update the NMM at inference time, we minimize a squared\-error key–value reconstruction loss over each multi\-token chunkcc, binding a map between𝐊\\mathbf\{K\}and𝐕\\mathbf\{V\}:

ℒc=‖ℳ⁡\(𝐊c\)−𝐕c‖2\.\\mathcal\{L\}\_\{c\}=\\left\\\|\\mathcal\{M\}\(\\mathbf\{K\}\_\{c\}\)\-\\mathbf\{V\}\_\{\\text\{c\}\}\\right\\\|^\{2\}\.
This loss gives the memory module a training signal at every chunk without requiring labels or a teacher attention output\. It encourages the module to retain associations present in the historical keys and values, which future queries may use for retrieval\. A future query, however, need not coincide with a stored key, and historical attention can combine values from several positions\. Key–value reconstruction is therefore a surrogate for useful memory, rather than an exact objective for reproducing historical attention\.

Optimization is done via gradient descent on the NMM once per chunk using learned momentum, decay, and step\-size parameters:

ℳc=\(1−αc\)​ℳc−1\+𝐒c,𝐒c=ηc​𝐒c−1−θc​∇𝐌ℒc​\(ℳc−1\)\.\\mathcal\{M\}\_\{c\}=\(1\-\\alpha\_\{c\}\)\\mathcal\{M\}\_\{c\-1\}\+\\mathbf\{S\}\_\{c\},\\qquad\\mathbf\{S\}\_\{c\}=\\eta\_\{c\}\\mathbf\{S\}\_\{c\-1\}\-\\theta\_\{c\}\\nabla\_\{\\mathbf\{M\}\}\\mathcal\{L\}\_\{c\}\(\\mathcal\{M\}\_\{c\-1\}\)\.
In practice, we do not perform this update on a per\-token scale, and instead use fixed\-size token batches to take advantage of modern hardware\. For even more efficiency, we use an offline variant of this algorithm by batching the gradient update over multiple chunks\.

### 3\.3Architecture

We make several architectural choices to adapt pretrained LLMs for efficient inference\. First, blockwise attention allows us to decouple long\-term relationships within the context\. This mask complicates the interpretation of equation[1](https://arxiv.org/html/2610.19210#S3.E1)because earlier layers also use segmented attention, so the historical keys and values differ from those produced by the full\-attention model\. Blockwise attention also provides a natural threshold for NMM hidden\-state updates and inference\-time gains from sparse computation\.

As shown in equation[1](https://arxiv.org/html/2610.19210#S3.E1), attention outputs are the mass\-weighted sum of local and historical contributions, which we aim to simulate via a predictedλ\\lambda\. This is nontrivial as our use ofλ\\lambdaessentially substitutes the usage of a softmax normalization as a means of token reweighting and mixing\. We choose to utilize a linear layer on top of the hidden input stream to predictλ^\\hat\{\\lambda\}as it is a query\-dependent quantity, and we use it to interpolate betweenOlocalO\_\{\\text\{local\}\}andO^hist\\hat\{O\}\_\{\\text\{hist\}\}\.

![Refer to caption](https://arxiv.org/html/2610.19210v1/long_context_quality_prefill_summary.png)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](https://arxiv.org/html/2610.19210#A3.T5)in the Appendix\.Right\.Larger chunks increase prefill latency as memory bandwidth saturates\.

## 4Experiments

We convert the first four layers of Llama\-3\.2\-1B\-Instruct\([Grattafiori et al\., 2024](https://arxiv.org/html/2610.19210#bib.bib30)\)to blockwise attention and add a neural memory module for each KV head\. The remaining layers retain full attention\. This setting tests whether memory can recover information no longer accessible through attention in the converted layers, without requiring every layer to rely on the new mechanism\.

Each three\-layer NMM learns historical key\-value bindings using a self\-supervised test\-time objective\. During post\-training, we optimize the external modules for language modeling while keeping the pre\-trained backbone frozen\. We conduct a series of experiments to measure how much long\-text performance the NMM recovers, which forms of recall it supports, and how our architectural design affects inference costs\.

### 4\.1Post\-Training for Long Context Reasoning

We post\-train Titan\-Llama on long sequences using a standard next\-token prediction loss, while the NMM continues to update its fast weights through the self\-supervised key\-value objective in Section[3\.2](https://arxiv.org/html/2610.19210#S3.SS2)\. The language modeling loss trains the added components to make those memory reads useful to the frozen backbone\. We initially considered distilling the backbone’s historical attention output into the NMM, but found this difficult to scale and ineffective in practice\. Under our final objective,O^hist\\hat\{O\}\_\{\\text\{hist\}\}need not reproduce that output exactly and instead learns a contribution that helps predict future tokens under segmented attention\.

We show that training Titan\-Llama on next\-token prediction, the details of which can be found in Appendix[A](https://arxiv.org/html/2610.19210#A1), improves not just next\-token prediction but also recall and reasoning over long\-term dependencies\. We evaluate long\-context question answering on QuALITY\([Pang et al\., 2022](https://arxiv.org/html/2610.19210#bib.bib23)\)and LongHealth\([Adams et al\., 2024](https://arxiv.org/html/2610.19210#bib.bib24)\)\. Results are shown in Figure[3](https://arxiv.org/html/2610.19210#S3.F3)and compared against a segmented attention baseline to show the effect induced by adding the NMM\.

Figure[2](https://arxiv.org/html/2610.19210#S3.F2)shows how language\-modeling loss and decode latency change with context length\. Loss delta is the difference in per\-token negative log\-likelihood from full\-attention Llama, so values below zero indicate better performance\. With 1024\-token segments, Titan\-Llama stays below zero from 16K through 128K tokens, while the measured Mamba\-2, Gated DeltaNet, and LaCT baselines incur larger loss penalties at longer contexts\. Because the model is training on sequences of at most 16K tokens, the measurements beyond 16K also test how its language\-modeling behavior extends beyond the training length\.

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\.†\\daggermarks timings of publicly available code that may not utilize caching\.
### 4\.2RULER Evaluation

Aggregate language\-modeling loss does not tell us which kinds of historical information are best encoded by the NMM\. The RULER benchmark measures different types of recall over long sequences, including retrieval of individual values, multiple associations, and word\-frequency information\([Hsieh et al\., 2024](https://arxiv.org/html/2610.19210#bib.bib31)\)\. The RULER results in Table[2](https://arxiv.org/html/2610.19210#S4.T2)compare full attention, segmented attention without memory, and Titan\-Llama\. The NMMs are able to effectively enrich representations enough to recover the majority of backbone performance, particularly for single\-needle and multi\-key retrieval at lengths up to 16K\.

Unexpectedly, Titan\-Llama also exceeds the full\-attention baseline on multi\-key retrieval at 4K–16K and on frequent/common word extraction at 8K–32K\. Multi\-value retrieval remains below full attention for every length, suggesting greater difficulty returning several distinct values from the same context\. Performance declines at 32K, beyond the 16K training length, suggesting limited length generalization in this setting\. This reveals that the NMM can perform semantic retrieval, accumulation, and recall of individual values, but is less proficient at recovering several values at once\.

When compared to other efficient architectures in Table[1](https://arxiv.org/html/2610.19210#S4.T1), Titan\-Llama matches the full attention baseline on pass\-key retrieval and exceeds full attention performance on UUID recall at 8K and 16K\. On 32K pass\-key retrieval, it also retains full accuracy where the measured efficient\-model baselines decline\. These comparisons show that compacting history can preserve strong retrieval on some tasks, even though the broader RULER results show that it does not preserve every form of recall equally well\. Especially compared to RNN\-based models with limited full\-attention layers such as Mamba\-2 and Gated DeltaNet, Titan\-Llama’s performance does not degrade as severely with increasing context lengths\.

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\.
### 4\.3Inference Gains

While various studies have applied TTT to modulate existing architecture in transformer, full attention oftentimes serves as a lower bound on performance and a ceiling on inference efficiency\.\([Feng et al\., 2026](https://arxiv.org/html/2610.19210#bib.bib16)\)Because our use of TTT adapts a cheaper replacement to long\-context relationships, we are able to harness greater inference efficiency\. The decode comparison in Figure[2](https://arxiv.org/html/2610.19210#S3.F2)shows this in effect as Titan\-Llama has lower median single\-token latency than the measured full\-attention and efficient\-model baselines across the reported context lengths\.

Figure[4](https://arxiv.org/html/2610.19210#S4.F4)separately measures how these costs change as we convert more layers\. At longer prompt lengths, Titan\-Llama achieves lower prefill latency than the full\-attention backbone, and the advantage grows with the number of converted layers, the details of which can be found in Appendix[B](https://arxiv.org/html/2610.19210#A2)\. In fact, the prefill process for Titan\-Llama is more than 20% faster than full\-attention at long prefill lengths with only five converted layers\. The steady\-state decode memory usage for Titan\-Llama is also 24% lower than full attention with the KV cache, though peak memory usage during prefill is more expensive for partial hybrids\. This reflects our intuition that our approach should result in inference economics gains, with one practical limitation that thepeakmemory footprint of Titan\-Llama exceeds that of full attention in shorter sequence lengths and in hybrid variants\.

![Refer to caption](https://arxiv.org/html/2610.19210v1/prefill_decode_sweep_layer_counts.jpg)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\.
### 4\.4Scaling chunk size

The chunk size in Titan\-Llama corresponds to the number of tokens in a local window as well as the effective token batch size for the inner objective loss\. A similar hyperparameter affects not just inference dynamics but capabilities in other studies as well[Feng et al\. \(2026\)](https://arxiv.org/html/2610.19210#bib.bib16)\. Changing it affects both what information remains available through exact local attention and how often the NMM is updated\. The RULER and long\-context reasoning results seen previously utilized 512\-token segments\. Figure[3](https://arxiv.org/html/2610.19210#S3.F3)varies the chunk size from 512 to 2048 tokens\. Larger chunks improve accuracy on QuALITY and LongHealth, but their effects on RULER are mixed: they help some task–length combinations while reducing accuracy on others, including several needle retrieval settings\.

We also study the effects of scaling context length on RULER performance and show in Appendix Table[5](https://arxiv.org/html/2610.19210#A3.T5)that shorter context lengths lead to better performance on NIAH\-style tasks while longer chunk sizes lend better semantic retrieval\.

Larger chunks also increase measured prefill latency\. Chunk size thus sets a task\-dependent balance among local attention, the frequency of memory updates, and inference cost, rather than providing a uniform accuracy gain\.

## 5Discussion

The idea of using test\-time training on an MLP to improve language model capabilities, whether on an existing part of LLM architecture or on external modules, has been well explored in a substantial body of work\.\([Behrouz et al\., 2024](https://arxiv.org/html/2610.19210#bib.bib18);[Feng et al\., 2026](https://arxiv.org/html/2610.19210#bib.bib16);[Tandon et al\., 2025](https://arxiv.org/html/2610.19210#bib.bib29);[Sun et al\., 2025](https://arxiv.org/html/2610.19210#bib.bib14);[Zhang et al\., 2025](https://arxiv.org/html/2610.19210#bib.bib19)\)Compared to these studies, however, our main contribution seeks to take advantage of the theoretical understanding of the updated set of weights as a memory mechanism and use it tosubstitutethe long\-range attention calculations of a pre\-trained backbone\.

To that end, our results corroborate the current line of research which utilizes TTT to improve language modeling, while demonstrating practical inference\-efficiency benefits\. Similar to[Zhang et al\. \(2025\)](https://arxiv.org/html/2610.19210#bib.bib19);[Behrouz et al\. \(2024\)](https://arxiv.org/html/2610.19210#bib.bib18), we show that TTT using KV binding effectively learns to do language modeling\. And like[Feng et al\. \(2026\)](https://arxiv.org/html/2610.19210#bib.bib16), we show that TTT can extend the context window for certain tasks\.

Our work also allows for open questions such as scaling NMM size or changing the NMM architecture\. Further, understanding the mechanisms behind memory retrieval is a natural extension of this work, and exploring the use of inner\-loop objectives other than KV binding could lead to improvements on this method\.

#### Acknowledgments

We would like to extend our appreciation to the MIT Office of Research Computing and Data for providing the computational resources to run our experiments\. We would also like to extend thanks to Sarah Dufays for her help in early versions of this project\.

## References

- Adamset al\.\(2024\)L\. Adams, F\. Busch, T\. Han, J\. Excoffier, M\. Ortala, A\. Löser, H\. JWL\. Aerts, J\. N\. Kather, D\. Truhn, and K\. BressemLongHealth: a question answering benchmark with long clinical documents\.External Links:2401\.14490,[Link](https://arxiv.org/abs/2401.14490)Cited by:[§4\.1](https://arxiv.org/html/2610.19210#S4.SS1.p2.1)\.
- Ainslieet al\.\(2020\)J\. Ainslie, S\. Ontanon, C\. Alberti, V\. Cvicek, Z\. Fisher, P\. Pham, A\. Ravula, S\. Sanghai, Q\. Wang, and L\. YangETC: encoding long and structured inputs in transformers\.External Links:2004\.08483,[Link](https://arxiv.org/abs/2004.08483)Cited by:[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px1.p1.1)\.
- Baiet al\.\(2024\)Y\. Bai, X\. Lv, J\. Zhang, Y\. He, J\. Qi, L\. Hou, J\. Tang, Y\. Dong, and J\. LiLongAlign: a recipe for long context alignment of large language models\.External Links:2401\.18058,[Link](https://arxiv.org/abs/2401.18058)Cited by:[§A\.1](https://arxiv.org/html/2610.19210#A1.SS1.p1.1)\.
- Bansalet al\.\(2025\)R\. Bansal, A\. Zhang, R\. Tiwari, L\. Madaan, S\. S\. Duvvuri, D\. Khatri, D\. Brandfonbrener, D\. Alvarez\-Melis, P\. Bhargava, M\. S\. Kale, and S\. JelassiLet’s \(not\) just put things in context: test\-time training for long\-context llms\.External Links:2512\.13898,[Link](https://arxiv.org/abs/2512.13898)Cited by:[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px3.p1.1)\.
- Behrouzet al\.\(2025\)A\. Behrouz, Z\. Li, P\. Kacham, M\. Daliri, Y\. Deng, P\. Zhong, M\. Razaviyayn, and V\. MirrokniATLAS: learning to optimally memorize the context at test time\.External Links:2505\.23735,[Link](https://arxiv.org/abs/2505.23735)Cited by:[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px3.p2.1),[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px3.p3.1)\.
- Behrouzet al\.\(2024\)A\. Behrouz, P\. Zhong, and V\. MirrokniTitans: learning to memorize at test time\.External Links:2501\.00663,[Link](https://arxiv.org/abs/2501.00663)Cited by:[§1](https://arxiv.org/html/2610.19210#S1.p2.1),[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px3.p2.1),[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px3.p3.1),[§3\.2](https://arxiv.org/html/2610.19210#S3.SS2.p1.1),[§5](https://arxiv.org/html/2610.19210#S5.p1.1),[§5](https://arxiv.org/html/2610.19210#S5.p2.1)\.
- Beltagyet al\.\(2020\)I\. Beltagy, M\. E\. Peters, and A\. CohanLongformer: the long\-document transformer\.External Links:2004\.05150,[Link](https://arxiv.org/abs/2004.05150)Cited by:[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px1.p1.1)\.
- Bicket al\.\(2025\)A\. Bick, K\. Y\. Li, E\. P\. Xing, J\. Z\. Kolter, and A\. GuTransformers to ssms: distilling quadratic knowledge to subquadratic models\.External Links:2408\.10189,[Link](https://arxiv.org/abs/2408.10189)Cited by:[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px2.p1.1)\.
- Bulatovet al\.\(2022\)A\. Bulatov, Y\. Kuratov, and M\. S\. BurtsevRecurrent memory transformer\.External Links:2207\.06881,[Link](https://arxiv.org/abs/2207.06881)Cited by:[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px1.p1.1)\.
- Daoet al\.\(2022\)T\. Dao, D\. Y\. Fu, S\. Ermon, A\. Rudra, and C\. RéFlashAttention: fast and memory\-efficient exact attention with io\-awareness\.External Links:2205\.14135,[Link](https://arxiv.org/abs/2205.14135)Cited by:[§3\.1](https://arxiv.org/html/2610.19210#S3.SS1.p3.1)\.
- Dao and Gu \(2024\)T\. Dao and A\. GuTransformers are ssms: generalized models and efficient algorithms through structured state space duality\.External Links:2405\.21060,[Link](https://arxiv.org/abs/2405.21060)Cited by:[Figure 2](https://arxiv.org/html/2610.19210#S3.F2)\.
- Eyubogluet al\.\(2025\)S\. Eyuboglu, R\. Ehrlich, S\. Arora, N\. Guha, D\. Zinsley, E\. Liu, W\. Tennien, A\. Rudra, J\. Zou, A\. Mirhoseini, and C\. ReCartridges: lightweight and general\-purpose long context representations via self\-study\.External Links:2506\.06266,[Link](https://arxiv.org/abs/2506.06266)Cited by:[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px2.p1.1)\.
- Fenget al\.\(2026\)G\. Feng, S\. Luo, K\. Hua, G\. Zhang, D\. He, W\. Huang, and T\. CaiIn\-place test\-time training\.External Links:2604\.06169,[Link](https://arxiv.org/abs/2604.06169)Cited by:[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px3.p1.1),[§3\.2](https://arxiv.org/html/2610.19210#S3.SS2.p1.1),[§4\.3](https://arxiv.org/html/2610.19210#S4.SS3.p1.1),[§4\.4](https://arxiv.org/html/2610.19210#S4.SS4.p1.1),[§5](https://arxiv.org/html/2610.19210#S5.p1.1),[§5](https://arxiv.org/html/2610.19210#S5.p2.1)\.
- Finnet al\.\(2017\)C\. Finn, P\. Abbeel, and S\. LevineModel\-agnostic meta\-learning for fast adaptation of deep networks\.External Links:1703\.03400,[Link](https://arxiv.org/abs/1703.03400)Cited by:[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px3.p1.1)\.
- Grattafioriet al\.\(2024\)A\. Grattafiori, A\. Dubey, A\. Jauhri, A\. Pandey, A\. Kadian, A\. Al\-Dahle, A\. Letman, A\. Mathur, A\. Schelten, A\. Vaughan, A\. Yang, A\. Fan, A\. Goyal, A\. Hartshorn, A\. Yang, A\. Mitra, A\. Sravankumar, A\. Korenev, A\. Hinsvark, A\. Rao, A\. Zhang, A\. Rodriguez, A\. Gregerson, A\. Spataru, B\. Roziere, B\. Biron, B\. Tang, B\. Chern, C\. Caucheteux, C\. Nayak, C\. Bi, C\. Marra, C\. McConnell, C\. Keller, C\. Touret, C\. Wu, C\. Wong, C\. C\. Ferrer, C\. Nikolaidis, D\. Allonsius, D\. Song, D\. Pintz, D\. Livshits, D\. Wyatt, D\. Esiobu, D\. Choudhary, D\. Mahajan, D\. Garcia\-Olano, D\. Perino, D\. Hupkes, E\. Lakomkin, E\. AlBadawy, E\. Lobanova, E\. Dinan, E\. M\. Smith, F\. Radenovic, F\. Guzmán, F\. Zhang, G\. Synnaeve, G\. Lee, G\. L\. Anderson, G\. Thattai, G\. Nail, G\. Mialon, G\. Pang, G\. Cucurell, H\. Nguyen, H\. Korevaar, H\. Xu, H\. Touvron, I\. Zarov, I\. A\. Ibarra, I\. Kloumann, I\. Misra, I\. Evtimov, J\. Zhang, J\. Copet, J\. Lee, J\. Geffert, J\. Vranes, J\. Park, J\. Mahadeokar, J\. Shah, J\. van der Linde, J\. Billock, J\. Hong, J\. Lee, J\. Fu, J\. Chi, J\. Huang, J\. Liu, J\. Wang, J\. Yu, J\. Bitton, J\. Spisak, J\. Park, J\. Rocca, J\. Johnstun, J\. Saxe, J\. Jia, K\. V\. Alwala, K\. Prasad, K\. Upasani, K\. Plawiak, K\. Li, K\. Heafield, K\. Stone, K\. El\-Arini, K\. Iyer, K\. Malik, K\. Chiu, K\. Bhalla, K\. Lakhotia, L\. Rantala\-Yeary, L\. van der Maaten, L\. Chen, L\. Tan, L\. Jenkins, L\. Martin, L\. Madaan, L\. Malo, L\. Blecher, L\. Landzaat, L\. de Oliveira, M\. Muzzi, M\. Pasupuleti, M\. Singh, M\. Paluri, M\. Kardas, M\. Tsimpoukelli, M\. Oldham, M\. Rita, M\. Pavlova, M\. Kambadur, M\. Lewis, M\. Si, M\. K\. Singh, M\. Hassan, N\. Goyal, N\. Torabi, N\. Bashlykov, N\. Bogoychev, N\. Chatterji, N\. Zhang, O\. Duchenne, O\. Çelebi, P\. Alrassy, P\. Zhang, P\. Li, P\. Vasic, P\. Weng, P\. Bhargava, P\. Dubal, P\. Krishnan, P\. S\. Koura, P\. Xu, Q\. He, Q\. Dong, R\. Srinivasan, R\. Ganapathy, R\. Calderer, R\. S\. Cabral, R\. Stojnic, R\. Raileanu, R\. Maheswari, R\. Girdhar, R\. Patel, R\. Sauvestre, R\. Polidoro, R\. Sumbaly, R\. Taylor, R\. Silva, R\. Hou, R\. Wang, S\. Hosseini, S\. Chennabasappa, S\. Singh, S\. Bell, S\. S\. Kim, S\. Edunov, S\. Nie, S\. Narang, S\. Raparthy, S\. Shen, S\. Wan, S\. Bhosale, S\. Zhang, S\. Vandenhende, S\. Batra, S\. Whitman, S\. Sootla, S\. Collot, S\. Gururangan, S\. Borodinsky, T\. Herman, T\. Fowler, T\. Sheasha, T\. Georgiou, T\. Scialom, T\. Speckbacher, T\. Mihaylov, T\. Xiao, U\. Karn, V\. Goswami, V\. Gupta, V\. Ramanathan, V\. Kerkez, V\. Gonguet, V\. Do, V\. Vogeti, V\. Albiero, V\. Petrovic, W\. Chu, W\. Xiong, W\. Fu, W\. Meers, X\. Martinet, X\. Wang, X\. Wang, X\. E\. Tan, X\. Xia, X\. Xie, X\. Jia, X\. Wang, Y\. Goldschlag, Y\. Gaur, Y\. Babaei, Y\. Wen, Y\. Song, Y\. Zhang, Y\. Li, Y\. Mao, Z\. D\. Coudert, Z\. Yan, Z\. Chen, Z\. Papakipos, A\. Singh, A\. Srivastava, A\. Jain, A\. Kelsey, A\. Shajnfeld, A\. Gangidi, A\. Victoria, A\. Goldstand, A\. Menon, A\. Sharma, A\. Boesenberg, A\. Baevski, A\. Feinstein, A\. Kallet, A\. Sangani, A\. Teo, A\. Yunus, A\. Lupu, A\. Alvarado, A\. Caples, A\. Gu, A\. Ho, A\. Poulton, A\. Ryan, A\. Ramchandani, A\. Dong, A\. Franco, A\. Goyal, A\. Saraf, A\. Chowdhury, A\. Gabriel, A\. Bharambe, A\. Eisenman, A\. Yazdan, B\. James, B\. Maurer, B\. Leonhardi, B\. Huang, B\. Loyd, B\. D\. Paola, B\. Paranjape, B\. Liu, B\. Wu, B\. Ni, B\. Hancock, B\. Wasti, B\. Spence, B\. Stojkovic, B\. Gamido, B\. Montalvo, C\. Parker, C\. Burton, C\. Mejia, C\. Liu, C\. Wang, C\. Kim, C\. Zhou, C\. Hu, C\. Chu, C\. Cai, C\. Tindal, C\. Feichtenhofer, C\. Gao, D\. Civin, D\. Beaty, D\. Kreymer, D\. Li, D\. Adkins, D\. Xu, D\. Testuggine, D\. David, D\. Parikh, D\. Liskovich, D\. Foss, D\. Wang, D\. Le, D\. Holland, E\. Dowling, E\. Jamil, E\. Montgomery, E\. Presani, E\. Hahn, E\. Wood, E\. Le, E\. Brinkman, E\. Arcaute, E\. Dunbar, E\. Smothers, F\. Sun, F\. Kreuk, F\. Tian, F\. Kokkinos, F\. Ozgenel, F\. Caggioni, F\. Kanayet, F\. Seide, G\. M\. Florez, G\. Schwarz, G\. Badeer, G\. Swee, G\. Halpern, G\. Herman, G\. Sizov, Guangyi, Zhang, G\. Lakshminarayanan, H\. Inan, H\. Shojanazeri, H\. Zou, H\. Wang, H\. Zha, H\. Habeeb, H\. Rudolph, H\. Suk, H\. Aspegren, H\. Goldman, H\. Zhan, I\. Damlaj, I\. Molybog, I\. Tufanov, I\. Leontiadis, I\. Veliche, I\. Gat, J\. Weissman, J\. Geboski, J\. Kohli, J\. Lam, J\. Asher, J\. Gaya, J\. Marcus, J\. Tang, J\. Chan, J\. Zhen, J\. Reizenstein, J\. Teboul, J\. Zhong, J\. Jin, J\. Yang, J\. Cummings, J\. Carvill, J\. Shepard, J\. McPhie, J\. Torres, J\. Ginsburg, J\. Wang, K\. Wu, K\. H\. U, K\. Saxena, K\. Khandelwal, K\. Zand, K\. Matosich, K\. Veeraraghavan, K\. Michelena, K\. Li, K\. Jagadeesh, K\. Huang, K\. Chawla, K\. Huang, L\. Chen, L\. Garg, L\. A, L\. Silva, L\. Bell, L\. Zhang, L\. Guo, L\. Yu, L\. Moshkovich, L\. Wehrstedt, M\. Khabsa, M\. Avalani, M\. Bhatt, M\. Mankus, M\. Hasson, M\. Lennie, M\. Reso, M\. Groshev, M\. Naumov, M\. Lathi, M\. Keneally, M\. Liu, M\. L\. Seltzer, M\. Valko, M\. Restrepo, M\. Patel, M\. Vyatskov, M\. Samvelyan, M\. Clark, M\. Macey, M\. Wang, M\. J\. Hermoso, M\. Metanat, M\. Rastegari, M\. Bansal, N\. Santhanam, N\. Parks, N\. White, N\. Bawa, N\. Singhal, N\. Egebo, N\. Usunier, N\. Mehta, N\. P\. Laptev, N\. Dong, N\. Cheng, O\. Chernoguz, O\. Hart, O\. Salpekar, O\. Kalinli, P\. Kent, P\. Parekh, P\. Saab, P\. Balaji, P\. Rittner, P\. Bontrager, P\. Roux, P\. Dollar, P\. Zvyagina, P\. Ratanchandani, P\. Yuvraj, Q\. Liang, R\. Alao, R\. Rodriguez, R\. Ayub, R\. Murthy, R\. Nayani, R\. Mitra, R\. Parthasarathy, R\. Li, R\. Hogan, R\. Battey, R\. Wang, R\. Howes, R\. Rinott, S\. Mehta, S\. Siby, S\. J\. Bondu, S\. Datta, S\. Chugh, S\. Hunt, S\. Dhillon, S\. Sidorov, S\. Pan, S\. Mahajan, S\. Verma, S\. Yamamoto, S\. Ramaswamy, S\. Lindsay, S\. Lindsay, S\. Feng, S\. Lin, S\. C\. Zha, S\. Patil, S\. Shankar, S\. Zhang, S\. Zhang, S\. Wang, S\. Agarwal, S\. Sajuyigbe, S\. Chintala, S\. Max, S\. Chen, S\. Kehoe, S\. Satterfield, S\. Govindaprasad, S\. Gupta, S\. Deng, S\. Cho, S\. Virk, S\. Subramanian, S\. Choudhury, S\. Goldman, T\. Remez, T\. Glaser, T\. Best, T\. Koehler, T\. Robinson, T\. Li, T\. Zhang, T\. Matthews, T\. Chou, T\. Shaked, V\. Vontimitta, V\. Ajayi, V\. Montanez, V\. Mohan, V\. S\. Kumar, V\. Mangla, V\. Ionescu, V\. Poenaru, V\. T\. Mihailescu, V\. Ivanov, W\. Li, W\. Wang, W\. Jiang, W\. Bouaziz, W\. Constable, X\. Tang, X\. Wu, X\. Wang, X\. Wu, X\. Gao, Y\. Kleinman, Y\. Chen, Y\. Hu, Y\. Jia, Y\. Qi, Y\. Li, Y\. Zhang, Y\. Zhang, Y\. Adi, Y\. Nam, Yu, Wang, Y\. Zhao, Y\. Hao, Y\. Qian, Y\. Li, Y\. He, Z\. Rait, Z\. DeVito, Z\. Rosnbrick, Z\. Wen, Z\. Yang, Z\. Zhao, and Z\. MaThe llama 3 herd of models\.External Links:2407\.21783,[Link](https://arxiv.org/abs/2407.21783)Cited by:[§4](https://arxiv.org/html/2610.19210#S4.p1.1)\.
- Gu and Dao \(2024\)A\. Gu and T\. DaoMamba: linear\-time sequence modeling with selective state spaces\.External Links:2312\.00752,[Link](https://arxiv.org/abs/2312.00752)Cited by:[§1](https://arxiv.org/html/2610.19210#S1.p1.1)\.
- Hanet al\.\(2026\)D\. Han, Y\. Li, T\. Li, Z\. Cao, Z\. Wang, J\. Song, Y\. Cheng, B\. Zheng, and G\. HuangViT3\{\}^\{3\}: unlocking test\-time training in vision\.External Links:2512\.01643,[Link](https://arxiv.org/abs/2512.01643)Cited by:[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px3.p3.1)\.
- Hsiehet al\.\(2024\)C\. Hsieh, S\. Sun, S\. Kriman, S\. Acharya, D\. Rekesh, F\. Jia, Y\. Zhang, and B\. GinsburgRULER: what’s the real context size of your long\-context language models?\.External Links:2404\.06654,[Link](https://arxiv.org/abs/2404.06654)Cited by:[§4\.2](https://arxiv.org/html/2610.19210#S4.SS2.p1.1)\.
- Hwanget al\.\(2025\)S\. Hwang, B\. Wang, and A\. GuDynamic chunking for end\-to\-end hierarchical sequence modeling\.External Links:2507\.07955,[Link](https://arxiv.org/abs/2507.07955)Cited by:[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px1.p2.1)\.
- Katharopouloset al\.\(2020\)A\. Katharopoulos, A\. Vyas, N\. Pappas, and F\. FleuretTransformers are rnns: fast autoregressive transformers with linear attention\.External Links:2006\.16236,[Link](https://arxiv.org/abs/2006.16236)Cited by:[§1](https://arxiv.org/html/2610.19210#S1.p1.1),[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px3.p3.1)\.
- Kitaevet al\.\(2020\)N\. Kitaev, Ł\. Kaiser, and A\. LevskayaReformer: the efficient transformer\.External Links:2001\.04451,[Link](https://arxiv.org/abs/2001.04451)Cited by:[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px1.p1.1)\.
- Kojimaet al\.\(2025\)Y\. Kojima, J\. Xu, X\. Zou, and X\. WangLoRA\-ttt: low\-rank test\-time training for vision\-language models\.External Links:2502\.02069,[Link](https://arxiv.org/abs/2502.02069)Cited by:[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px2.p1.1)\.
- Kuratovet al\.\(2026\)Y\. Kuratov, M\. Kairov, A\. Bulatov, I\. Rodkin, and M\. BurtsevGradMem: learning to write context into memory with test\-time gradient descent\.External Links:2603\.13875,[Link](https://arxiv.org/abs/2603.13875)Cited by:[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px3.p2.1)\.
- Nawrotet al\.\(2022\)P\. Nawrot, S\. Tworkowski, M\. Tyrolski, Ł\. Kaiser, Y\. Wu, C\. Szegedy, and H\. MichalewskiHierarchical transformers are more efficient language models\.External Links:2110\.13711,[Link](https://arxiv.org/abs/2110.13711)Cited by:[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px1.p2.1)\.
- Panget al\.\(2022\)R\. Y\. Pang, A\. Parrish, N\. Joshi, N\. Nangia, J\. Phang, A\. Chen, V\. Padmakumar, J\. Ma, J\. Thompson, H\. He, and S\. R\. BowmanQuALITY: question answering with long input texts, yes\!\.External Links:2112\.08608,[Link](https://arxiv.org/abs/2112.08608)Cited by:[§A\.1](https://arxiv.org/html/2610.19210#A1.SS1.p1.1),[§4\.1](https://arxiv.org/html/2610.19210#S4.SS1.p2.1)\.
- Patelet al\.\(2024\)P\. Patel, E\. Choukse, C\. Zhang, A\. Shah, Í\. Goiri, S\. Maleki, and R\. BianchiniSplitwise: efficient generative llm inference using phase splitting\.External Links:2311\.18677,[Link](https://arxiv.org/abs/2311.18677)Cited by:[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px1.p1.1)\.
- Penedoet al\.\(2024\)G\. Penedo, H\. Kydlíček, L\. B\. allal, A\. Lozhkov, M\. Mitchell, C\. Raffel, L\. V\. Werra, and T\. WolfThe fineweb datasets: decanting the web for the finest text data at scale\.External Links:2406\.17557,[Link](https://arxiv.org/abs/2406.17557)Cited by:[§A\.1](https://arxiv.org/html/2610.19210#A1.SS1.p1.1)\.
- Qiuet al\.\(2020\)J\. Qiu, H\. Ma, O\. Levy, S\. W\. Yih, S\. Wang, and J\. TangBlockwise self\-attention for long document understanding\.External Links:1911\.02972,[Link](https://arxiv.org/abs/1911.02972)Cited by:[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px1.p1.1)\.
- Slagle \(2024\)K\. SlagleSpaceByte: towards deleting tokenization from large language modeling\.External Links:2404\.14408,[Link](https://arxiv.org/abs/2404.14408)Cited by:[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px1.p2.1)\.
- Sunet al\.\(2025\)Y\. Sun, X\. Li, K\. Dalal, J\. Xu, A\. Vikram, G\. Zhang, Y\. Dubois, X\. Chen, X\. Wang, S\. Koyejo, T\. Hashimoto, and C\. GuestrinLearning to \(learn at test time\): rnns with expressive hidden states\.External Links:2407\.04620,[Link](https://arxiv.org/abs/2407.04620)Cited by:[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px3.p1.1),[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px3.p3.1),[§5](https://arxiv.org/html/2610.19210#S5.p1.1)\.
- Sunet al\.\(2020\)Y\. Sun, X\. Wang, Z\. Liu, J\. Miller, A\. A\. Efros, and M\. HardtTest\-time training with self\-supervision for generalization under distribution shifts\.External Links:1909\.13231,[Link](https://arxiv.org/abs/1909.13231)Cited by:[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px3.p1.1)\.
- Tandonet al\.\(2025\)A\. Tandon, K\. Dalal, X\. Li, D\. Koceja, M\. Rød, S\. Buchanan, X\. Wang, J\. Leskovec, S\. Koyejo, T\. Hashimoto, C\. Guestrin, J\. McCaleb, Y\. Choi, and Y\. SunEnd\-to\-end test\-time training for long context\.External Links:2512\.23675,[Link](https://arxiv.org/abs/2512.23675)Cited by:[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px3.p2.1),[§3\.2](https://arxiv.org/html/2610.19210#S3.SS2.p1.1),[§5](https://arxiv.org/html/2610.19210#S5.p1.1)\.
- Wanget al\.\(2025a\)J\. Wang, D\. Paliotta, A\. May, A\. M\. Rush, and T\. DaoThe mamba in the llama: distilling and accelerating hybrid models\.External Links:2408\.15237,[Link](https://arxiv.org/abs/2408.15237)Cited by:[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px2.p1.1)\.
- Wanget al\.\(2025b\)S\. Wang, G\. Zhang, L\. L\. Zhang, N\. Shang, F\. Yang, D\. Chen, and M\. YangLoongRL: reinforcement learning for advanced reasoning over long contexts\.External Links:2510\.19363,[Link](https://arxiv.org/abs/2510.19363)Cited by:[§A\.1](https://arxiv.org/html/2610.19210#A1.SS1.p1.1)\.
- Wenet al\.\(2026\)Q\. Wen, Z\. Huang, X\. Meng, W\. He, and C\. LiMixture\-of\-top\-k attention: efficient attention via scalable fast weights\.External Links:2602\.01219,[Link](https://arxiv.org/abs/2602.01219)Cited by:[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px3.p3.1)\.
- Yanget al\.\(2025\)S\. Yang, J\. Kautz, and A\. HatamizadehGated delta networks: improving mamba2 with delta rule\.External Links:2412\.06464,[Link](https://arxiv.org/abs/2412.06464)Cited by:[Figure 2](https://arxiv.org/html/2610.19210#S3.F2)\.
- Zhanget al\.\(2025\)T\. Zhang, S\. Bi, Y\. Hong, K\. Zhang, F\. Luan, S\. Yang, K\. Sunkavalli, W\. T\. Freeman, and H\. TanTest\-time training done right\.External Links:2505\.23884,[Link](https://arxiv.org/abs/2505.23884)Cited by:[§2](https://arxiv.org/html/2610.19210#S2.SS0.SSS0.Px3.p2.1),[§5](https://arxiv.org/html/2610.19210#S5.p1.1),[§5](https://arxiv.org/html/2610.19210#S5.p2.1)\.
- Zhanget al\.\(2023\)Z\. Zhang, Y\. Sheng, T\. Zhou, T\. Chen, L\. Zheng, R\. Cai, Z\. Song, Y\. Tian, C\. Ré, C\. Barrett, Z\. Wang, and B\. ChenH2\{\}\_\{2\}o: heavy\-hitter oracle for efficient generative inference of large language models\.External Links:2306\.14048,[Link](https://arxiv.org/abs/2306.14048)Cited by:[§1](https://arxiv.org/html/2610.19210#S1.p2.1)\.
- Zweigeret al\.\(2026\)A\. Zweiger, X\. Fu, H\. Guo, and Y\. KimFast kv compaction via attention matching\.External Links:2602\.16284,[Link](https://arxiv.org/abs/2602.16284)Cited by:[§1](https://arxiv.org/html/2610.19210#S1.p2.1)\.

## Appendix ATraining Details

### A\.1Datasets

We train on one billion tokens of packed sequences from FineWeb, with a maximum sequence length of 16K tokens\([Penedo et al\., 2024](https://arxiv.org/html/2610.19210#bib.bib32)\)\. We then perform supervised fine\-tuning on LongAlign\([Bai et al\., 2024](https://arxiv.org/html/2610.19210#bib.bib25)\)and LoongRL\([Wang et al\., 2025b](https://arxiv.org/html/2610.19210#bib.bib26)\)\. A final fine\-tuning stage uses the QuALITY training set to adapt the model to multiple\-choice answering\([Pang et al\., 2022](https://arxiv.org/html/2610.19210#bib.bib23)\)\.

### A\.2Hyperparameters

The first four Llama layers each contain one neural memory module \(NMM\) per KV head\. Each NMM is a three\-layer network with a hidden dimension of 64 and GELU activations\. The added components contain 17\.23M trainable parameters in total\.

We use a batch size of 64 during the initial training stage, corresponding to approximately one million tokens per batch at the maximum sequence length\. Supervised fine\-tuning uses a batch size of 128 with variable answer lengths\. The inner\-loss batch size equals the chunk size\. We evaluate chunk sizes of 512, 1024, and 2048 tokens in Section[4\.4](https://arxiv.org/html/2610.19210#S4.SS4)\. Compute constraints prevented a broad search over these architectural and training choices\.

![Refer to caption](https://arxiv.org/html/2610.19210v1/lalrl_prefix_depth_prefill_decode_sweep_fixed.jpg)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\.

## Appendix BInference Implementation

Autoregressive decoding performs many small operations for each generated token\. At batch size one, the cost of launching those operations can be substantial\. For the measurements in Figure[4](https://arxiv.org/html/2610.19210#S4.F4), we compile the Titan\-Llama layers in CUDA Graph “reduce\-overhead” mode and replay their single\-token decode computations\. Compilation and graph capture take place during warmup, outside the timed decode steps\.

The captured computation uses a preallocated KV ring buffer for each attention segment\. Cache positions and memory\-store indices are updated in place on the GPU, and attention operands retain fixed shapes\. We capture separate paths for ordinary tokens and for tokens that complete a memory\-update chunk\. The remaining full\-attention layers are compiled without graph capture so that they retain their optimized attention implementation\.

CUDA Graphs reduce launch overhead but still execute the underlying kernels\. A fused implementation of the memory read, update, and output blend could reduce intermediate memory traffic and launch overhead further\. Such an implementation would need to preserve the ordering of reads and updates at chunk boundaries\.

## Appendix CAdditional Experiments

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\.### C\.1Restricting the Visibility of Later Attention Layers

In the main configuration, the layers after the Titan\-Llama layers retain full attention\. They can therefore access earlier tokens directly, even though the converted layers use segmented attention and neural memory\. To test how much the model relies on this remaining full attention, we also restrict the later layers to local windows\. In this variant, neural memory is the only route by which information from a completed segment can affect later segments\.

The restriction reduces inference cost, as shown in Figure[5](https://arxiv.org/html/2610.19210#A1.F5), but it also causes a large loss in accuracy\. On QuALITY, adding the NMM raises accuracy from 26\.8 to 35\.2; on LongHealth, it raises accuracy from 15\.8 to 32\.3 \(Table[3](https://arxiv.org/html/2610.19210#A3.T3)\)\. Both results remain below the full\-attention baseline\. The RULER results in Table[4](https://arxiv.org/html/2610.19210#A3.T4)show the same limitation more sharply: the NMM provides occasional gains over segmented attention, but neither configuration reliably retrieves information across segments\.

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\.
### C\.2Chunk Size and RULER Performance

Chunk size controls both the span of local attention and the number of tokens used for each memory update\. Table[5](https://arxiv.org/html/2610.19210#A3.T5)reports RULER results for 512\-, 1024\-, and 2048\-token chunks\. The effect depends on the task and context length\. The 512\- token setting is strongest on most multi\-key and multi\-value retrieval settings, while larger chunks improve some individual results, including 16K SQuAD and multi\-value accuracy\. These results do not support a single best chunk size across the RULER tasks\.

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