Parameter-efficient linearization turns a pretrained autoregressive Transformer into an efficient model by hybridizing sliding-window attention (SWA) with a constant-size recurrent memory. But does every layer's recurrent state contribute to long-context retrieval? Investigating an SWA and Test-Time Training (TTT) hybrid, we find that retrieval is sparse across depth, concentrated in a few anchor layers, and that the remaining states hold similar information expressed through layer-specific transformations. We exploit both observations in Anchor: TTT memory placed only at the anchor layers, and a cross-layer recurrent state written once and read by many layers, each through its own learned read map that realizes a layer-specific stretch and rotation, along with demonstrated generalization to other hybrid attentions. On Llama-3.1-8B and Mistral-7B, Anchor delivers 2.3-3.1x higher recall and 6.7-8.9 points higher MMLU accuracy than LoLCATs with a 10x smaller recurrent state, retains 91-95% of the teacher's MMLU accuracy, and cuts end-to-end latency by up to 4.3x relative to softmax attention.
Figures & tables
Figure 1: Recurrent states across layers are redundant, so one shared state suffices. Left: Debiased CKA between the TTT states of every pair of layers in Llama-3.1-8B, measured on WikiText for targets beyond the local window. The anchor layers ( ℓ∈{0,1,28,…,31} , dashed boxes) are strongly aligned with one another (mean CKA 0.92 ), whereas the interior layers are not ( 0.49 ). Right: Number of recurrent states per block as the ratio of hybrid to softmax-attention layers grows. A per-layer design keeps one state per hybrid layer; Anchor keeps a single shared state, reducing the count by 3 – 6× .
Figure 2: Layer-wise contribution of each branch to long-context retrieval. ΔCE (nats) on AR-32k when the TTT (top) or SWA (bottom) branch is removed at a single layer, broken down by the distance between a query and its target. Warm cells indicate degradation, that is, a larger increase in loss when the branch is removed. The TTT branch contributes almost nothing across the interior layers ( ℓ∈{2,…,27} ) but dominates at the boundary layers ( ℓ∈{0,1,28,…,31} ), whereas SWA contributes broadly at every depth. The color scale is symmetric-log with a linear region below 0.01 , and both panels share one scale.
Configuration
PPL ↓
State (MB) ↓
Llama-3.1-8B (baseline)
7.36
–
Non-unified state
10.19
201.3
− Anchor TTT
487.56
163.6
− Interior TTT
10.01
37.7
Table 1: Ablation of the test-time-training (TTT) module on long-context performance. Disabling the TTT update in the anchor layers ( ℓ∈{0,1,28,29,30,31} ) raises out-of-domain perplexity on WikiText by 48× , while disabling it in the interior layers ( ℓ∈{2,…,27} ) leaves perplexity unchanged and shrinks the persistent state by 5.3× . SWA remains active in all layers.
Figure 3: Improved Attention Transfer in Stage-1 Linearization. Held-out attention-transfer loss LMSE (left) and perplexity (right) during Stage 1. (a) Training on LMSE alone matches the teacher’s attention outputs closely, reducing the loss by over 25× , yet perplexity rises from 5.3K to 236K: the student reproduces each layer’s output while losing the ability to model language. (b) Adding next-token supervision LCE resolves this, lowering both losses together and reaching a perplexity of 34, which provides a stable initialization for the TTT parameters.
Figure 4: Representational similarity of TTT states on (a) PG-19 and (b) WikiText. Each point is one layer pair, plotted as debiased CKA against mean cosine similarity on cross-window targets, that is, those whose source lies beyond the local window. Interior pairs ( , ℓ∈{2,…,27} ) are drawn separately from anchor pairs ( ▲ , ℓ∈{0,1,28,…,31} ). Anchor pairs are more strongly aligned than interior pairs on both corpora ( 0.97 versus 0.82 on PG-19, 0.92 versus 0.49 on WikiText). High CKA with near-zero cosine indicates closely related information encoded in different bases.
Figure 5: Anchor architecture and recurrent-state coverage. Left: Anchor-F maintains a private recurrent state in every layer, while Anchor keeps one shared state per anchor group and removes the interior states entirely, reducing recurrent state by 16× ( 201.0→12.6 MB at 32k context). Each reader accesses the shared state through its own learned read map χ(⋅) . Right: coverage over query (rows) and key (columns) positions. Green marks the local sliding window, red the span served by the recurrent state. Anchor-F requires one such state per layer; Anchor serves layers l1,…,L from a single shared state.
State
Tput
Commonsense
Recall
Model
Config
(MB) ↓
(k/s) ↑
LMB.
PIQA
Hella.
Wino.
ARC-e
ARC-c
Avg.
SWDE
FDA
SQuAD
Avg.
LLaMA
Baseline
—
1.73
75.28
81.12
79.04
73.64
81.52
53.50
74.02
91.81
81.03
53.82
75.55
SWA
—
8.32
75.00
80.47
78.45
73.09
81.10
52.73
73.47
27.72
12.43
48.66
29.60
LoLCATs
129.0
3.18
69.09
81.61
78.56
73.95
82.87
55.20
73.55
18.99
1.72
32.57
17.76
Anchor -F
201.0
3.41
74.21
80.36
76.63
72.69
80.56
49.40
72.31
50.50
18.60
60.46
43.19
Anchor -I
176.0
3.66
73.04
79.16
77.37
72.69
79.59
47.44
71.55
48.51
19.24
54.02
40.59
Table 2: Zero-shot downstream reasoning and recall performance. Accuracy of the fine-tuned Anchor variants on six common-sense reasoning benchmarks (LAMBADA, PIQA, HellaSwag, WinoGrande, ARC-e, ARC-c) and three recall-intensive benchmarks (SWDE, FDA, SQuAD), compared against the softmax-attention baseline. Anchor -F gives every sliding window its own TTT module; Anchor -I shares one TTT module across the initial and final layers while the interior layers keep separate modules; Anchor shares one module across the anchor layers and removes the interior modules entirely. (k/s) - K Token/sec ; measured at N = 8K and B = 19 .
Model
Config.
E2E Latency Over Seq. Length (ms)
2K
4K
8K
16K
LLaMa
Baseline
29164
51049
95387
OOM
Anchor -F
38028
40719
48471
56516
Anchor -I
38770
38333
45219
53775
Anchor
19050
20854
22816
27915
Mistral
Baseline
28872
53100
100300
OOM
Table 3: Study on End-to-end latency (ms) of Anchor variants compared with the softmax-attention baselines (Llama and Mistral) across sequence lengths with the batch-size (B) of 19 . Anchor reduces end-to-end latency by up to 4.3× relative to the baseline at 8K, and still runs at 16K, where both baselines exhaust memory ( OOM ) : Out-Of-Memory.
Figure 6: Study on Performance–efficiency trade-off. Anchor evaluated on reasoning (MMLU) and recall (SWDE, FDA, SQuAD) benchmarks, together with end-to-end throughput and recurrent state size. (a) Reasoning accuracy of the Anchor series against LoLCATs; targeted placement of the cross-layer recurrent state improves reasoning. (b) Throughput (tokens/s) against reasoning accuracy; Anchor attains 4.3× higher throughput than the baseline at comparable accuracy. (c) Recall accuracy against recurrent state size; Anchor improves recall over LoLCATs by 2.3–3.1 × while using a 10.2 × smaller state.
Model
IC & Noisy
Fuzzy
Mem.
Sel. Copy
Comp.
Avg.
Transformer
98.7
55.7
100.0
99.9
100.0
90.9
3:1 -Hybrid
96.8
37.7
99.9
86.4
100.0
84.1
Anchor -F
100.0
58.0
100.0
99.9
100.0
91.6
Anchor -I
100.0
67.9
100.0
99.9
100.0
93.5
Table 4: Mechanistic architecture design (MAD) synthetic tasks. Accuracy (%) of each architecture trained from scratch on the MAD suite. Anchor -I attains the highest average and the largest gain on fuzzy recall.
Figure 7: Hybrid architecture of Anchor and full self-attention. Anchor and full self-attention layers are interleaved at a 3:1 ratio, so that softmax attention accounts for 25% of all layers. Within each group, the three Anchor layers retain their own sliding-window attention for local context but share a single recurrent state for long-range context: the leading layer writes the state, and each layer reads it through its own learned read map χℓ(⋅) . The subsequent full self-attention layer attends over the entire sequence, providing exact retrieval at regular depth intervals, while the recurrent state per group remains constant in sequence length.
Model
Attn. (%)
4K
8K
16K
LLaMA-3-8B (teacher)
100
92.3
90.5
85.7
Lizard (hybrid)
50
92.5
91.2
85.2
Anchor -Hybrid
25
98.3
97.3
82.9
Table 5: Evaluation on Long-Context Retrieval (RULER). We evaluate Anchor -hybrid on RULER at sequence lengths from 4 K to 16 K tokens and compare it with the LLaMA-3-8B teacher and Lizard, a prior hybrid long-context linearization method.
Figure 8: The learned reader rotation is concentrated in a few planes. Rotation angles of Q ’s planes, ranked (log scale), at L31; median over 32 heads, whiskers show the IQR. Shaded bands span from the learned profile to a random rotation with equal total squared angle (dark: learned above null; light: below). The learned map exceeds the null only in the leading planes (up to plane 5 for Llama-3.1-8B and plane 3 for Mistral-7B), and its top 5 planes carry 82% and 92% of the squared rotation versus 26% for the null.This holds in all 128 reader-layer heads of both backbones.
Figure 9: Cross-layer similarity generalizes across linearization methods. Debiased CKA against mean cosine similarity for every pair of layers, on the recurrent-state readouts of our TTT hybrid (Anchor-F), LoLCATs, and Liger on Mistral-7B over identical PG-19 and WikiText inputs. In all three models, layers share information while their cosine similarity remains small, indicating that this information is encoded in different bases.
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 10: Needle-in-a-haystack retrieval. Accuracy across context lengths (1K–8K) and needle depths for Llama-3.1-8B, a LoLCATs hybrid ( 18.75% SoftMax Attention), and the Anchor 3:1 hybrid. Anchor reaches a mean accuracy of 0.95 (baseline 1.00), whereas the LoLCATs hybrid ( 18.75% SoftMax Attention) (0.33) retrieves reliably only at 1K or when the needle lies near the end of the context.
Stage 1
Stage 2
Motivational Study
Dataset
emozilla/pg19
Train / Val docs
200 / 20
200 / 20
Learning Rate (LR)
2×10−4
1×10−4
Max steps
400
200
Eval every
25
25
Appendix
Table 6: Training Configuration of Anchor -F Architecture. we trained the architecture separately for motivational study and main study with different dataset to understand to architecture behavior in different corpus and to maintain the fairness with other baselines. Context Length = 8192 and Window Size = 512 .
Figure 11: Geometry of the learned read maps. Each column is one reader layer, and the rows show the readout as it passes through the two stages of the read map Q : a stretch followed by a rotation. Points are readout vectors projected onto the layer’s dominant rotation plane and coloured by position in the sequence; arcs trace the rotation step, which preserves the radius within the plane. Panel titles give the rotation angle θ in that plane and the share of readout energy it carries. Both backbones follow the same pattern, with the rotation concentrated in a few planes (Figure 8). Head 0 of 32; axes are scaled per layer.
Memory-augmented language models combine persistent or temporary memory with neural computation to retain, retrieve, and use information beyond the active context.
Transformers owe much of their strong long-context retrieval capability to a token-level memory that grows with context length. This flexibility, however, incurs a quadratic computation complexity during training and a key--value cache that grows linearly during autoregressive inference. Recurrent alternatives offer efficient decoding by compressing the entire history into a fixed-size state, but often underperform on recall-intensive tasks since earlier associations usually get overwritten by subsequent updates, and only the most recent contextual information is retained. In this paper, we introduce Memory-Anchor Routing across Context History (MARCH), a network architecture that effectively scales state-space models beyond a fixed-size dimension, while maintaining computational efficiency over long-sequences. MARCH periodically caches cumulative recurrent-state checkpoints as state anchors and associates each anchor with a compact, content-conditioned anchor key. This lets MARCH maintain a memory bank, which can grow as context length increases, providing a controllable trade-off between historical resolution and memory cost. At each token, MARCH produces an anchor query to attend all causally available state anchors, and the output is calculated as an attention-style aggregation over all historical anchors along the current state. We show that after standard pretraining, MARCH consistently outperforms multiple linear attention variants across commonsense reasoning, LongBench, and in-context retrieval. These results demonstrate that content-routed state caching substantially strengthens recurrent long-range memory while preserving its native computation path.
Ming Zhang, Kaisen Yang, Shu Yu +6
Shanghai AI Laboratory · Tsinghua University · Fudan University
We study Latent Recurrent Transformer (LRT), a lightweight augmentation of autoregressive transformers that reuses a high-level source-layer hidden state from the previous token as recurrent memory for the next token. Because this state is already computed during ordinary decoding, LRT introduces a cross-token, cross-layer latent pathway while preserving the standard attention mechanism, KV-cache interface, and one model forward per generated token. To pretrain this recurrence without sequentially unrolling the full sequence, we introduce interleaved parallel training: one full-sequence initialization forward constructs a shared buffer, followed by sequential refinement of disjoint position subsets with parallel computation within each subset. This provides every token with recurrent-memory-aware supervision at approximately 2x ideal token compute. Across 1.3B- and 2.1B-parameter nanochat-style backbones and a wide range of training budgets, LRT improves both BPB and CORE under matched effective compute. Additionally, LRT outperforms two-forward PonderLM-2 and matches a three-loop Transformer in BPB, while retaining one-forward-per-token decoding with 9% latency overhead over the standard Transformer.
Zeyi Huang, Xuehai He, LiLiang Ren +8
Microsoft · University of Washington · University of Wisconsin-Madison
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
Sarah Pan, Ali Behrouz
Massachusetts Institute of Technology · Cornell University