Research map / Architectures and efficiency / Attention and transformers
Learning to (Learn at Test Time): RNNs with Expressive Hidden States
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Test-Time Training (TTT) layers whose hidden state is a learned model updated by a self-supervised gradient step, instantiated as TTT-Linear and TTT-MLP, and shows they match or beat Mamba and Transformers on long-context language modeling up to 32k.
Problem
Existing RNN layers compress context into a fixed-size hidden state whose limited expressive power prevents them from exploiting long context, while self-attention has quadratic cost.
Approach
Make the RNN hidden state a machine learning model f with weights W, and the update rule one (mini-batch) gradient step on a self-supervised reconstruction loss, so the layer performs training at test time; the self-supervised task's views (theta_K, theta_V, theta_Q) are learned in the outer loop. A dual form enables efficient matmul-based computation.
Key results
- At 8k context on the Pile, TTT-Linear and TTT-MLP achieve better perplexity than Mamba across model sizes (125M-1.3B).
- On the Books 32k setting with matched FLOPs as Mamba 1.4B, TTT-Linear (M) and TTT-MLP (M) outperform Mamba, and TTT-MLP (T) slightly beats Mamba at 32k; Mamba perplexity plateaus after 16k while TTT keeps decreasing.
- Ablation from linear attention (15.23 ppl) to TTT-Linear (11.09 ppl) on 125M models; mini-batch TTT gives the largest gain (-1.70), and the dual form gives over 5x training speedup on TPUs.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| improve the linear attention implementation (ablation) | Linear attention · Pile | perplexity: 15.91 -> 15.23 | Helped |
| implement the linear-attention-equivalent layer as TTT (ablation) | Linear attention improved · Pile | perplexity: 15.23 -> 15.23 | No clear effect |
| add learnable W0 (ablation) | TTT equivalence · Pile | perplexity: 15.23 -> 15.27 | Hurt |
| add LayerNorm and a residual connection in f (ablation) | TTT with learnable W0 · Pile | perplexity: 15.27 -> 14.05 | Helped |
| use mini-batch TTT with b=16 instead of b=T=2048 (ablation) | TTT with LN and residual in f · Pile | perplexity: 14.05 -> 12.35 | Helped |
| add learnable eta (ablation) | TTT with mini-batch TTT · Pile | perplexity: 12.35 -> 11.99 | Helped |
| use the Mamba backbone (ablation) | TTT with learnable eta · Pile | perplexity: 11.99 -> 11.09 | Helped |
| use TTT-Linear or TTT-MLP instead of Mamba | Mamba · Pile, 8k context | perplexity | Helped |
| use TTT-Linear or TTT-MLP instead of Mamba | Mamba · Books, 32k context | perplexity | Helped |
| use TTT-MLP instead of TTT-Linear at matched FLOPs | TTT-Linear · Books, short versus long context | perplexity | Mixed |
| use the Transformer backbone instead of the Mamba backbone (ablation) | TTT layers · Pile and Books | perplexity | Hurt |
| use TTT-MLP instead of TTT-Linear (ablation) | TTT layers with the Transformer backbone · Pile and Books | perplexity | Helped |
Surprising findings
- TTT-MLP is effective in FLOPs but its wall-clock time grows much more relative to FLOPs, so it is slower in practice than TTT-Linear.
- Trained Transformers show no clean linear scaling-law fit on these settings, so the authors connect points instead of fitting scaling laws.
Limitations
- TTT-MLP has high wall-clock cost due to memory I/O despite good FLOP efficiency.
- Experiments limited to 125M-1.3B parameters and up to 32k context due to academic compute.
- Systems training kernel not written; only inference kernels for GPU, training on TPUs.
Experimental setup
Datasets: The Pile Books3
Benchmarks: The Pile (2k, 8k context perplexity) Books3 (1k-32k context perplexity)
Metrics: perplexity FLOPs latency (ms)
Baselines: Transformer++ Mamba linear attention DeltaNet TF finetune
Builds on: Mamba linear attention DeltaNet dynamic evaluation fast weight programmers
Closest papers by meaning
- Explaining Modern Gated-Linear RNNs via a Unified Implicit Attention Formulation (ICLR 2025 · similarity 0.89)
- M$^2$RNN: Non-Linear RNNs with Matrix-Valued States for Scalable Language Modeling (NeurIPS 2026 · similarity 0.86)
- Rethinking Expressivity and Efficiency in Test-Time Training (NeurIPS 2026 · similarity 0.86)
- RAT: Bridging RNN Efficiency and Attention Accuracy via Chunk-based Sequence Modeling (NeurIPS 2025 · similarity 0.86)
- Learning What to Remember: Test-Time Training via Context Distillation (NeurIPS 2026 · similarity 0.86)
- Measuring In-Context Computation Complexity via Hidden State Prediction (ICML 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.