cs.LGFeb 26, 2026

SR-TTT Does Not Learn Retrieval: A Correction and Mechanistic Post-Mortem of Surprisal-Aware Residual Test-Time Training

Authors: Swamynathan V P

Organizations: Department of Computer Science and Engineering Amrita School of Computing Amrita Vishwa Vidyapeetham Chennai, Tamil Nadu, India

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Test-Time Training (TTT) language models replace the KV-cache with fast weights updated during inference, achieving O(1) memory but suffering catastrophic failure on exact-recall tasks. Version 1 of this work proposed SR-TTT, which routes high-surprisal tokens to a sparse exact-attention Residual Cache, and reported large Needle-in-a-Haystack gains. We show those gains were evaluation artifacts: the loss and metric read logits at the answer positions rather than one position earlier, training both models to copy an answer already visible in their input (a model trained on retrieval-impossible data reaches 100% accuracy under the flawed metric); additionally, the cache attended non-causally over future tokens, including the answer itself. We release a corrected implementation with startup causality self-tests, then ask whether the SR-TTT hypothesis survives correction. It does not, and the failure decomposes into two independent, separately measured bottlenecks. Storage: surprisal gating is systematically position-biased - the TTT reconstruction loss requires burn-in before a needle becomes relatively surprising, so early-context needles are stored at near-zero rates (0-1% containment at depth 0.1) exactly where long-context memory matters most. Addressing: with storage solved by an oracle and with new trainable read-time projections, per-slot attention supervision raises addressing mass 2.5x (0.06 -> 0.15) yet token accuracy is statistically unchanged, and retrieval extracts only approx. 0.06 nats of the 2.30-nat needle; position-free content addressing cannot resolve ordered slots whose contents are near-interchangeable. Exact match remains 0% in all 2,250 paired trials across all corrected conditions. We retract the claims of v1 and offer the corrected codebase, diagnostic protocol, and negative results as a cautionary reference for surprise-gated memory architectures.

Figures & tables

Figure 1: The storage bottleneck. (a) Natural (unforced) needle flagging during Stage-2 training collapses as the distance curriculum moves the needle earlier in context. (b) At evaluation, cache containment of the needle is near zero at shallow depths and peaks mid-sequence: reconstruction-loss surprise requires burn-in before a needle is relatively surprising.
Figure 2: The addressing–accuracy dissociation under oracle storage (containment =100%=100\% ). (a) Per-slot oracle supervision raises addressing mass 2.5×2.5\times over the unsupervised variant. (b) Token accuracy is nonetheless statistically indistinguishable between variants and an order of magnitude short of exact-match viability.
Figure 3: Information budget under oracle storage at the training length: retrieval recovers ≈\approx 0.06 nats/token of the 2.303-nat digit needle.
ConditionModelExact matchToken acc. (1024/2048/4096)Pooled McNemar vs. A
Natural storageA (pure TTT)0%3.0% / 9.7% / 4.3%—
Natural storageB-full0%6.7% / 10.7% / 8.2%p=1.0p=1.0 (0 discordant)
Natural storageB-noOracle0%4.6% / 10.2% / 7.1%p=1.0p=1.0 (0 discordant)
Oracle storageB-full0%12.0% / 17.4% / 11.3%—
Oracle storageB-noOracle0%12.9% / 16.1% / 12.2%—
Table 1: Corrected paired evaluation (50 trials/cell, pooled over depths; generation exact match; identical seeded samples). “Oracle storage” force-inserts the needle at evaluation (upper bound). Chance token accuracy is 10%.
More in this topic

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 291 papers

Explore similar work

CardsList
Oct 12, 2026

Error Correction in Trained TTT Layers

Test-time-training (TTT) layers update fast-weight memory while processing an input sequence. A gradient-based update, however, does not by itself establish how the memory revises an existing association. We study the structure and empirical role of prediction-derived updates in released TTT implementations. For dot-product objectives, we characterize the absence of a prediction-error residual and establish conditions under which inner-loop ascent and descent are equivalent after reparameterization. For the LayerNorm-based squared-error objective used by TTT-Linear and TTT-MLP, we derive the inner gradient and identify a prediction-derived component that cancels at the released initialization in the zero-epsilon limit. We then measure the full normalized prediction's contribution on three trained checkpoints. Removing this contribution while retaining the LayerNorm Jacobian gives a mean gradient norm ratio of 0.25 to 0.32 across the models. On 100,303 tokens from War and Peace, the same intervention with local per-token norm matching raises perplexity by 2.8-2.9% on TTT-Linear-760M and 1.3B, whereas freezing all fast-weight updates raises it by factors of 4.2-4.5. Separate synthetic tests illustrate the dependence of plain delta-rule overwriting on step size. The results distinguish a prediction-derived update from its behavioral value and show why implementation analysis and live interventions are both needed to assess test-time memory.

Yuhe Hu

Duke University

Test-Time TrainingMemory-Augmented Language ModelsNeural Network Training Dynamics
AskPDF
Aug 21, 2026

Rethinking Expressivity and Efficiency in Test-Time Training

Test-Time Training (TTT) enables long-context processing via continuous weight updates during inference, but current methods struggle to balance the expressivity of per-token update dynamics with the hardware efficiency of chunk-wise approximations. We propose E2^2-TTT (Expressive and Efficient TTT) to bridge this gap. Under the standard approximation of taking gradients at the chunk-start weights, we derive a closed-form state transition that exactly reproduces the chunk-end fast-weight and momentum states of the per-token recurrence. This enables fully parallelized chunk-level training while preserving the temporal structure of the update rule that prior chunk-wise methods discard. We validate E2^2-TTT by training models up to 1.3B parameters from scratch. It performs on par with previous TTT and hybrid attention baselines in language modeling while outperforming them on in-context retrieval. Its advantage is most pronounced in length extrapolation: on the standard ``Needle in a Haystack'' passkey test, it retains over 90% accuracy at 8×8\times the training context length. Meanwhile, E2^2-TTT can match the training throughput of efficient chunk-wise methods, demonstrating that it effectively reconciles expressivity with efficiency. The code is available at https://github.com/zeyun-zhong/E2-TTT.

Zeyun Zhong, Joya Chen, Manuel Martin +3

Karlsruhe Institute of Technology (KIT) · Fraunhofer IOSB · National University of Singapore +2

Test-Time TrainingLong-Context Language ModelingEfficient Language Model Training
AskPDF
Feb 24, 2026

Test-Time Training with KV Binding Is Secretly Linear Attention

Test-time training (TTT) with KV binding as sequence modeling layer is commonly interpreted as a form of online meta-learning that memorizes a key-value mapping at test time. However, our analysis reveals multiple phenomena that contradict this memorization-based interpretation. Motivated by these findings, we revisit the formulation of TTT and show that a broad class of TTT architectures can be expressed as a form of learned linear attention operator. Beyond explaining previously puzzling model behaviors, this perspective yields multiple practical benefits: it enables principled architectural simplifications, admits fully parallel formulations that preserve performance while improving efficiency, and provides a systematic reduction of diverse TTT variants to a standard linear attention form. Overall, our results reframe TTT not as test-time memorization, but as learned linear attention with enhanced representational capacity. Project page: https://research.nvidia.com/labs/sil/projects/tttla/.

Junchen Liu, Sven Elflein, Or Litany +2

NVIDIA, Toronto, Ontario, Canada · University of Toronto, Toronto, Ontario, Canada · Vector Institute, Toronto, Ontario, Canada +1

Test-Time TrainingLinear AttentionSequence Modeling
AskPDF