Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

SR-TTT: Surprisal-Aware Residual Test-Time Training

Swamynathan V P

TL;DR

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.

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: https://github.com/swamynathanvp/Surprisal-Aware-Residual-Test-Time-Training.

SR-TTT: Surprisal-Aware Residual Test-Time Training

TL;DR

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.

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: https://github.com/swamynathanvp/Surprisal-Aware-Residual-Test-Time-Training.
Paper Structure (12 sections, 1 equation, 4 figures)

This paper contains 12 sections, 1 equation, 4 figures.

Table of Contents

  1. Introduction
  2. Related Work
  3. Methodology
  4. The Surprisal Filter
  5. Residual Cache & Alpha Fusion
  6. Experiments
  7. Model Configuration
  8. The Two-Stage Curriculum
  9. Results
  10. RoPE Extrapolation Failure
  11. Limitations
  12. Conclusion and Future Work

Figures (4)

  • Figure 1: Training loss comparison showing the Two-Stage Curriculum for SR-TTT. Graphs are stacked to display full trajectory dynamics.
  • Figure 2: Overall loss trajectory comparison between the Pure TTT Baseline and Two-Stage SR-TTT methodologies.
  • Figure 3: Needle-in-a-Haystack Exact Match performance at sequence length 2048. SR-TTT demonstrates massive improvements at mid-sequence depths.
  • Figure 4: Cache utilization statistics demonstrating saturation and Layer 1--3 $\alpha$ gates opening to $\approx 10\%$.