Anchor: Augmented Cross-Layer Recurrent State for Long-Context Retrieval in Test-Time Training

Vimal William; Yuanzhe Li; Danella Zhao; Huanrui Yang

cs.LGNov 3, 2026

Anchor: Augmented Cross-Layer Recurrent State for Long-Context Retrieval in Test-Time Training

Authors: Vimal William, Yuanzhe Li, Danella Zhao, Huanrui Yang

Organizations: School of Electrical, Computer & Software Engineering, University of Arizona, Tucson, AZ, USA.

Abstract

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 ( $\ell\in\{0,1,28,\dots,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\times$ .

Figure 2: Layer-wise contribution of each branch to long-context retrieval. $\Delta\mathrm{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 ( $\ell\in\{2,\dots,27\}$ ) but dominates at the boundary layers ( $\ell\in\{0,1,28,\dots,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 $\downarrow$	State (MB) $\downarrow$
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 (

\ell\in\{0,1,28,29,30,31\}

) raises out-of-domain perplexity on WikiText by

48\times

, while disabling it in the interior layers (

\ell\in\{2,\dots,27\}

) leaves perplexity unchanged and shrinks the persistent state by

5.3\times

. SWA remains active in all layers.

Figure 3: Improved Attention Transfer in Stage-1 Linearization. Held-out attention-transfer loss $\mathcal{L}_{\mathrm{MSE}}$ (left) and perplexity (right) during Stage 1. (a) Training on $\mathcal{L}_{\mathrm{MSE}}$ alone matches the teacher’s attention outputs closely, reducing the loss by over $25\times$ , 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 $\mathcal{L}_{\mathrm{CE}}$ 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 ( ,

\ell\in\{2,\dots,27\}

) are drawn separately from anchor pairs (

\blacktriangle

\ell\in\{0,1,28,\dots,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\times

(

201.0\rightarrow 12.6

MB at 32k context). Each reader accesses the shared state through its own learned read map

\chi(\cdot)

. 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

l_{1},\dots,L

from a single shared state.

		State	Tput	Commonsense							Recall
Model	Config	(MB) $\downarrow$	(k/s) $\uparrow$	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$	$\mathbf{74.02}$	$91.81$	$81.03$	$53.82$	$\mathbf{75.55}$
	SWA	—	$\mathbf{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 =

8\text{K}

and B =

19

Model	Config.	E2E Latency Over Seq. Length (ms)
Model	Config.	2K	4K	8K	16K
LLaMa	Baseline	$29164$	$51049$	$95387$	OOM
	Anchor -F	$38028$	$40719$	$48471$	$56516$
	Anchor -I	$38770$	$38333$	$45219$	$53775$
	Anchor	$\mathbf{19050}$	$\mathbf{20854}$	$\mathbf{22816}$	$\mathbf{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\times

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\times$ 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 $\times$ while using a 10.2 $\times$ 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	$\mathbf{100.0}$	$\mathbf{67.9}$	$100.0$	$99.9$	$100.0$	$\mathbf{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

\chi_{\ell}(\cdot)

. 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$	$\mathbf{85.2}$
Anchor -Hybrid	$\mathbf{25}$	$\mathbf{98.3}$	$\mathbf{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{\times}10^{-4}$	$1{\times}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

Panel of Figure 11 — 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 $\theta$ 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

Memory-augmented language models combine persistent or temporary memory with neural computation to retain, retrieve, and use information beyond the active context.

See all 300 papers

MARCH: Scaling Recurrent Memory with Content-Routed State Anchors
Aug 12, 2026Ming Zhang, Kaisen Yang, Shu Yu +6Long-Context Retrieval In-Context Retrieval
Shanghai AI Laboratory · Tsinghua University · Fudan University
Latent Recurrent Transformer: Architecture Exploration, Training Strategies, and Scaling Behavior
May 26, 2026Zeyi Huang, Xuehai He, LiLiang Ren +8Language Model Pretraining Transformer Inference
Microsoft · University of Washington · University of Wisconsin-Madison
Titan-Llama: Retrofitting Pre-trained LLMs with External Neural Memory Modules
Oct 13, 2026Sarah Pan, Ali BehrouzLong-Context Language Model Inference Memory-Augmented Language Models
Massachusetts Institute of Technology · Cornell University

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

MARCH: Scaling Recurrent Memory with Content-Routed State Anchors

Ming Zhang, Kaisen Yang, Shu Yu, Ermo Hua+5

Shanghai AI Laboratory · Tsinghua University · Fudan University

Long-Context Retrieval In-Context Retrieval Memory-Augmented Language Models Persistent Memory for Language Models Long-Context Modeling State Space Models

Open paper AskPDF arXiv

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.

Abstract

Memory-Augmented Language Models

Explore similar work

MARCH: Scaling Recurrent Memory with Content-Routed State Anchors

Latent Recurrent Transformer: Architecture Exploration, Training Strategies, and Scaling Behavior

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