Paper deep dive
SR-TTT: Surprisal-Aware Residual Test-Time Training
Swamynathan V P
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 88%
Last extracted: 7/20/2026, 10:29:05 AM
Summary
The paper introduces SR-TTT (Surprisal-Aware Residual Test-Time Training), a hybrid architecture that augments Test-Time Training (TTT) language models with a loss-gated sparse memory mechanism to solve the 'Lost in the Middle' problem. By routing only highly surprising, incompressible tokens to a Residual Cache while maintaining O(1) memory for background context, SR-TTT significantly improves exact-recall performance on Needle-in-a-Haystack tasks compared to pure TTT baselines.
Entities (12)
Relation Signals (12)
SR-TTT → isa → Hybrid Architecture
confidence 95% · SR-TTT provides a robust proof-of-concept that O(1) memory Test-Time Training architectures can achieve exact-recall... by utilizing a loss-gated Residual Cache.
Residual Cache → isusedby → SR-TTT
confidence 92% · By dynamically routing only incompressible, highly surprising tokens to a traditional exact-attention Residual Cache, SR-TTT preserves O(1) memory
SR-TTT → usescomponent → Residual Cache
confidence 92% · SR-TTT maintains the infinite-context compression of standard TTT while incorporating a parallel memory track consisting of a Surprisal Filter, a Residual Cache, and a Fusion Gate.
SR-TTT → outperforms → Pure TTT
confidence 90% · Our evaluation demonstrates that SR-TTT substantially mitigates the exact-recall failures of infinite-context models... while the Pure TTT baseline explicitly forgets the needle
SR-TTT → solvesproblem → Needle-in-a-Haystack
confidence 90% · We evaluated SR-TTT using an 8-character alphanumeric Needle-in-a-Haystack protocol... SR-TTT demonstrates massive improvements at mid-sequence depths.
Test-Time Training → suffersfrom → Lost in the Middle
confidence 90% · Despite these efficiency gains, O(1) memory models suffer from a severe “Lost in the Middle” problem.
SR-TTT → usescomponent → Surprisal Filter
confidence 90% · SR-TTT maintains the infinite-context compression of standard TTT while incorporating a parallel memory track consisting of a Surprisal Filter...
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Test-Time Training (TTT) language models achieve theoretically infinite context windows with an O(1) memory footprint by replacing the standard exact-attention KV-cache with hidden state ``fast weights'' W_fast updated via self-supervised learning during inference. However, pure TTT architectures suffer catastrophic failures on exact-recall tasks (e.g., Needle-in-a-Haystack). Because the fast weights aggressively compress the context into an information bottleneck, highly surprising or unique tokens are rapidly overwritten and forgotten by subsequent token gradient updates. We introduce SR-TTT (Surprisal-Aware Residual Test-Time Training), which resolves this recall failure by augmenting the TTT backbone with a loss-gated sparse memory mechanism. By dynamically routing only incompressible, highly surprising tokens to a traditional exact-attention Residual Cache, SR-TTT preserves O(1) memory for low-entropy background context while utilizing exact attention exclusively for critical needles. Our complete implementation, training scripts, and pre-trained weights are open-source and available at: this https URL.
Tags
Links
- Source: https://arxiv.org/abs/2603.06642v1
- Canonical: https://arxiv.org/abs/2603.06642v1
Trouble viewing inline? Open PDF directly →
Full Text
16,445 characters extracted from source content.
Expand or collapse full text
SR-T: SURPRISAL-AWARE RESIDUAL TEST-TIME TRAINING ∗ Swamynathan V P Department of Computer Science and Engineering Amrita School of Computing Amrita Vishwa Vidyapeetham Chennai, Tamil Nadu, India vpswamynathan0311@gmail.com ABSTRACT Test-Time Training (T) language models achieve theoretically infinite context windows with anO(1)memory footprint by replacing the standard exact-attention KV-cache with hidden state “fast weights” (W fast ) updated via self-supervised learning during inference. However, pure T architectures suffer catastrophic failures on exact-recall tasks (e.g., Needle-in-a-Haystack). Be- cause the fast weights aggressively compress the context into an information bottleneck, highly surprising or unique tokens are rapidly overwritten and forgotten by subsequent token gradi- ent updates. We introduce SR-T (Surprisal-Aware Residual Test-Time Training), which re- solves this recall failure by augmenting the T backbone with a loss-gated sparse memory mechanism. By dynamically routing only incompressible, highly surprising tokens to a tradi- tional exact-attention Residual Cache, SR-T preservesO(1)memory for low-entropy back- ground context while utilizing exact attention exclusively for critical needles. Our complete im- plementation, training scripts, and pre-trained weights are open-source and available at:https: //github.com/swamynathanvp/Surprisal-Aware-Residual-Test-Time-Training. Keywords: Test-Time Training, Large Language Models, Long-Context Memory, Exact Recall, Sparse Attention, Hybrid Architecture, Curriculum Learning 1 Introduction The sequence length of large language models is traditionally bottlenecked by theO(N 2 )computational andO(N) memory complexities of the Transformer’s KV-cache [1]. Recent Test-Time Training (T) models propose an alternative: replacing the explicitly stored cache with a set of fast weights that summarize the sequence history. These weights are updated during the forward pass using a self-supervised objective, yielding anO(1)memory footprint [3]. Despite these efficiency gains,O(1)memory models suffer from a severe “Lost in the Middle” problem [6]. Standard T representations compress information uniformly, forcing a trade-off between retaining recent token dependencies and preserving distinct, rare tokens (such as specific names, IDs, or exact alphanumeric strings) from early in the sequence. In long-context Needle-in-a-Haystack evaluations [7], these rare needles are overwritten by the continuous stream of updates from background tokens. This recall-versus-compression tension is not unique to T. Prior work on linear attention and state-space models has established that fixed-size recurrent states face a fundamental recall bottleneck once the sequence length exceeds the state capacity [11,12]. Hybrid architectures that combine compressed states with selective exact attention have emerged as a promising mitigation strategy [11, 13, 10]. To address this, we propose SR-T. We hypothesize that exact recall in infinite-context models can be rescued without sacrificing theO(1)asymptotic benefits by pairing the T backbone with a highly selective, sparse Resid- ual Cache. Unlike prior hybrid approaches that use fixed sliding windows or attention-score heuristics, SR-T uses the T inner-loop reconstruction loss itself as a principled signal for identifying incompressible tokens. ∗ Citation: Swamynathan V P. SR-T: Surprisal-Aware Residual Test-Time Training. 2026. arXiv:2603.06642v1 [cs.LG] 26 Feb 2026 SR-T: Surprisal-Aware Residual Test-Time Training Code and experimental data can be reproduced via our public repository:https://github.com/swamynathanvp/ Surprisal-Aware-Residual-Test-Time-Training. 2 Related Work Test-Time Training (T). The concept of adapting model parameters at inference time was introduced by Sun et al. [2] for distribution-shift robustness in vision. The application of T to completely replace the KV-cache as a sequence modeling primitive was formalized in T-E2E by Sun et al. [3]. Their architecture demonstrated that training a linear model over hidden states during inference could achieve strong perplexity over infinite contexts. However, the exact recall problem remains a known limitation of aggressively compressed recurrent states. Concurrently, T Done Right [4] proposed large-chunk T with Muon optimization, combining window attention with T layers for improved efficiency, though without explicitly addressing the recall failure mode through surprisal-based routing. Gated Linear Recurrent Models. Recent work on Gated Differentiable Working Memory (GDWM) [5] and modern linear RNNs such as Mamba [8] and RWKV [9] have attempted to solve the context capacity problem using data- dependent gating and chunk-wise linear attention. The Titans architecture [10] introduced a neural memory module with explicit memorization and forgetting mechanisms for long-context modeling. Unlike GDWMs and Titans, which learn to gate the update to the recurrent state, SR-T specifically isolates anomalous tokens and physically routes them to an external sparse memory buffer, bypassing the recurrent bottleneck entirely for tokens that cannot be safely compressed. Hybrid Compressed-Exact Architectures. The recall limitations of sub-quadratic models have motivated several hybrid approaches. BASED [11] combines linear attention with sliding window attention, traversing the Pareto frontier of the recall-memory tradeoff by varying the window size and feature dimension. LESS [13] learns a low-rank residual between sparse KV-cache outputs and full attention, synthesizing recurrence with selective exact computation. TRIM- KV [14] learns per-token retention scores for memory-bounded cache eviction. MInference [15] leverages dynamic sparse attention patterns for long-context acceleration. Our approach differs from these methods in two key respects: (i) we use the T reconstruction loss as a principled, self-supervised routing signal rather than attention-score heuristics or learned gates, and (i) the Residual Cache operates as a parallel memory track alongside the T fast weights rather than as a modification to the KV-cache of a standard Transformer. KV-Cache Compression. Orthogonal to architectural changes, significant work has focused on compressing the KV-cache itself through quantization [16], eviction policies [17], and token merging [18]. While these methods reduce memory for standard Transformers, they do not address the fundamental compression bottleneck of recurrent or T-based architectures. 3 Methodology SR-T maintains the infinite-context compression of standard T while incorporating a parallel memory track consisting of a Surprisal Filter, a Residual Cache, and a Fusion Gate. 3.1 The Surprisal Filter During the T inner-loop forward pass, we calculate the per-token reconstruction lossL t =∥z t − v t ∥ 2 . Tokens are flagged as “surprising” (incompressible) if they satisfy a dual-track condition: 1. The per-token loss exceeds an Exponential Moving Average (EMA) smoothed percentile threshold. 2. The mean loss of the local chunk containing the token also exceeds a proportional threshold. Letτ EMA be the EMA-smoothed95 th percentile of recent reconstruction losses. A token at positiontis routed to the cache if L t > τ EMA and the surrounding chunk loss L chunk > 0.8· τ EMA . 3.2 Residual Cache & Alpha Fusion Flagged tokens (their post-RoPE Keys and Values) are parked in a fixed-capacity Residual Cache with a priority-based eviction policy. To retrieve this information, a multi-head attention module queries the cache using the current T query projections. The output is fused back into the main T stream via a learned gate vector α: Output = T(x) + α· CacheAttention(x)(1) 2 SR-T: Surprisal-Aware Residual Test-Time Training Early experimentation revealed that standard Sigmoid gating suffered from dying gradients. To resolve this, we utilize a direct clamp parameterization for the gates:α = clamp(θ gate , 0,α max ), ensuring stable gradient flow during the critical cache-integration phase. 4 Experiments We evaluated SR-T using an 8-character alphanumeric Needle-in-a-Haystack protocol embedded within the TinySto- ries dataset [19]. TinyStories provides a structured, low-entropy background, allowing the Surprisal Filter to clearly isolate high-entropy needles. 4.1 Model Configuration Both the pure T baseline and the SR-T models were configured with15.8M parameters (4 Layers,d model = 256) and trained for 10,000 steps at a sequence length of 2048. 4.2 The Two-Stage Curriculum Training SR-T directly from scratch end-to-end fails due to “Cold Start Noise.” In the early stages of training, the T backbone produces uncalibrated representations. Consequently, the network minimizes its loss by forcing theα gates to 0.0, effectively shutting off the cache and acting as a pure T model. We circumvented this using a Two-Stage Curriculum: • Stage 1 (Steps 1–7,000): Base T training with the Residual Cache disabled. • Stage 2 (Steps 7,001–10,000): The T backbone parameters are frozen, and the cache is enabled. By freezing the backbone, the network is forced to route gradients exclusively through theαmodule to minimize the remaining cross-entropy loss, successfully pulling the gates open. 5 Results Our evaluation demonstrates that SR-T substantially mitigates the exact-recall failures of infinite-context models at the evaluated training length of 2048 tokens. As shown in Figure 3, while the Pure T baseline explicitly forgets the needle when it falls out of the recent sliding window, the SR-T Surprisal Filter successfully routes the needle to the cache, resulting in massive retrieval improvements: • Depth 0.50: +23% Exact Match Improvement (from 10% to 33%) • Depth 0.75: +20% Exact Match Improvement (from 17% to 37%) Furthermore, Figure 4 confirms that the Two-Stage Curriculum successfully forced theαgates open to approximately 10% at the deeper semantic layers, validating the selective routing hypothesis. 5.1 RoPE Extrapolation Failure At the 4096 context length, both architectures suffered a complete collapse (0% exact match). This is an anticipated architectural limitation: the models were explicitly trained at a maximum sequence length of 2048. Standard Rotary Position Embeddings (RoPE) [20] suffer catastrophic zero-shot extrapolation failures when evaluated at unseen positional frequencies significantly beyond their training length. 6 Limitations We identify three primary limitations of this work: First, all experiments were conducted at a small scale (15.8M parameters, 2048 context length). While this is sufficient to validate the core mechanism, it remains an open question whether the Surprisal Filter thresholds and Two-Stage Curriculum transfer effectively to billion-parameter models with longer training contexts. Scaling behavior is a critical direction for future investigation. 3 SR-T: Surprisal-Aware Residual Test-Time Training (a) Model A (Pure T Baseline) Training. (b) Model B (Two-Stage SR-T) Training. Notice the Stage 2 cache warmup at step 7,000. Figure 1: Training loss comparison showing the Two-Stage Curriculum for SR-T. Graphs are stacked to display full trajectory dynamics. Figure 2: Overall loss trajectory comparison between the Pure T Baseline and Two-Stage SR-T methodologies. 4 SR-T: Surprisal-Aware Residual Test-Time Training Figure 3: Needle-in-a-Haystack Exact Match performance at sequence length 2048. SR-T demonstrates massive improvements at mid-sequence depths. Figure 4: Cache utilization statistics demonstrating saturation and Layer 1–3 α gates opening to≈ 10%. 5 SR-T: Surprisal-Aware Residual Test-Time Training Second, the RoPE extrapolation wall (Section 5.1) prevents zero-shot evaluation at context lengths beyond the training length. This is not a limitation of the SR-T mechanism itself, but it restricts the current evaluation to the trained sequence length. Third, the priority-based eviction mechanism in the Residual Cache functioned well at 2048-token sequences, but under extreme contextual volume the fixed-capacity cache converges to standard FIFO behavior, potentially evicting important earlier needles when the cache saturates. A learned eviction policy would be more robust. 7 Conclusion and Future Work SR-T provides a robust proof-of-concept thatO(1)memory Test-Time Training architectures can achieve exact-recall without compromising their compression benefits by utilizing a loss-gated Residual Cache. The Two-Stage Curriculum proved essential in overcoming the cold-start noise that ordinarily prevents dynamic memory integration. Future work will address the limitations identified in Section 6. First, we plan to implement YaRN [21] or Dynamic NTK interpolation to resolve the RoPE scaling wall, enabling zero-shot extrapolation to 4096+ lengths without explicitly increasing the baseline sequence training length. Second, we aim to replace the heuristic eviction policy with a fully learned scoring mechanism for cache eviction, drawing on ideas from TRIM-KV [14] and attention-based importance scoring. Third, scaling experiments to larger model sizes and longer training contexts will be necessary to validate the generality of the surprisal-based routing mechanism. References [1]Vaswani, A., et al. (2017). Attention Is All You Need. Advances in Neural Information Processing Systems (NeurIPS). [2]Sun, Y., Wang, X., Liu, Z., Miller, J., Efros, A., & Hardt, M. (2020). Test-Time Training with Self-Supervision for Generalization under Distribution Shifts. International Conference on Machine Learning (ICML). [3]Sun, Y., et al. (2024). Learning to (Learn at Test Time): RNNs with Expressive Hidden States. International Conference on Machine Learning (ICML). [4] Zhang, T., et al. (2025). Test-Time Training Done Right. arXiv preprint arXiv:2505.23884. [5]Anonymous Authors. (2026). Gated Differentiable Working Memory (GDWM) for Long-Context Language Modeling. arXiv preprint arXiv:2601.12906. [6]Liu, N. F., et al. (2024). Lost in the Middle: How Language Models Use Long Contexts. Transactions of the Association for Computational Linguistics (TACL). [7] Kamradt, G. (2023). Needle In A Haystack – Pressure Testing LLMs. GitHub Repository. [8] Gu, A., & Dao, T. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv preprint arXiv:2312.00752. [9] Peng, B., et al. (2023). RWKV: Reinventing RNNs for the Transformer Era. Findings of EMNLP. [10] Behrouz, A., et al. (2025). Titans: Learning to Memorize at Test Time. arXiv preprint arXiv:2501.00663. [11]Arora, S., et al. (2024). Simple Linear Attention Language Models Balance the Recall-Throughput Tradeoff . International Conference on Machine Learning (ICML). [12] Arora, S., et al. (2024). Zoology: Measuring and Improving Recall in Efficient Language Models. International Conference on Learning Representations (ICLR). [13]Dong, H., et al. (2024). Get More with LESS: Synthesizing Recurrence with KV Cache Compression for Efficient LLM Inference. International Conference on Machine Learning (ICML). [14] Anonymous Authors. (2025). Cache What Lasts: Token Retention for Memory-Bounded KV Cache in LLMs. arXiv preprint arXiv:2512.03324. [15] Jiang, H., et al. (2024). MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention. Advances in Neural Information Processing Systems (NeurIPS). [16]Hooper, C., et al. (2024). KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantiza- tion. arXiv preprint arXiv:2401.18079. [17]Zhang, Z., et al. (2024). H 2 O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models. Advances in Neural Information Processing Systems (NeurIPS). 6 SR-T: Surprisal-Aware Residual Test-Time Training [18]Nawrot, P., et al. (2024). Dynamic Memory Compression: Retrofitting LLMs for Accelerated Inference. Interna- tional Conference on Machine Learning (ICML). [19]Eldan, R., & Li, Y. (2023). TinyStories: How Small Can Language Models Be and Still Speak Coherent English?. arXiv preprint arXiv:2305.07759. [20]Su, J., et al. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv preprint arXiv:2104.09864. [21]Peng, B., et al. (2023). YaRN: Efficient Context Window Extension of Large Language Models. arXiv preprint arXiv:2309.00071. 7