Research map / Architectures and efficiency / Attention and transformers

Rethinking Expressivity and Efficiency in Test-Time Training

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: Sparse attention for long-context LLMs · 0 citations (OpenAlex)

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 E²-TTT, a closed-form chunk-parallel reformulation of token-wise TTT with per-token momentum and decay that exactly recovers chunk-end states under chunk-start gradient evaluation, and trains 340M/1.3B hybrid (SWA + TTT) models on FineWeb-Edu.

Problem

Token-wise TTT is expressive but sequential and hardware-inefficient, while chunk-wise TTT (LaCT) is fast but collapses per-token momentum and decay, losing expressivity and length extrapolation.

Approach

Derives two scalar kernels (momentum kernel and weight kernel) from unrolling the token-wise momentum-plus-decay recurrence, so chunk-end fast weights and momentum are computed exactly in parallel from chunk-start gradients via log-space cumulative sums and two weighted gradient aggregations.

Key results

• 1.3B E2-TTT_MLP reaches 15.3 LAMBADA perplexity and 54.5% zero-shot average vs 53.8% for the strongest baseline (HQLT).
• On S-NIAH-1 at 16K (8x the 2K training context), E2-TTT_SwiGLU retains 93.6% accuracy while LaCT falls to 3.0% and HQLT to about 25%.
• On LongBench (14 tasks), E2-TTT_SwiGLU averages 14.1 vs 12.1 for HQLT, 10.3 for Mamba2, and 7.7 for LaCT.

Surprising findings

• A chunk-averaged variant with per-token factors collapsed to chunk scalars matches the full model in-domain (15.6 vs 15.8 LAMBADA ppl) yet collapses on long-context retrieval, suggesting the retrieval gains come from the per-token kernel.
• The TTT branch alone without sliding-window attention scores 0.0% on S-NIAH-1 at all lengths, because chunk-start output cannot resolve within-chunk dependencies.

Limitations

• Exact equivalence holds only under chunk-start gradient evaluation, an approximation relative to fully sequential token-wise TTT.
• Zero-shot commonsense accuracy gains over baselines are modest (within about two points); in-domain gains are mainly in perplexity.
• Retrieval performance on FDA and SWDE still trails full-attention Transformer++.

Experimental setup

Datasets: FineWeb-Edu LLaVA-Video-178K

Benchmarks: WikiText perplexity LAMBADA PIQA HellaSwag WinoGrande ARC-e ARC-c SWDE SQuAD FDA S-NIAH-1 S-NIAH-2

Metrics: perplexity accuracy average score per-token loss

Baselines: Transformer++ DeltaNet Mamba2 HQLT LaCT LaCT-Matched Qwen3VL (SFT)

Builds on: test-time training LaCT HQLT Gated DeltaNet Mamba2

Closest papers by meaning

• In-Place Test-Time Training (ICLR 2026 · similarity 0.90)
• Let's (not) just put things in Context: Test-time Training for Long-context LLMs (ICLR 2026 · similarity 0.90)
• Learning What to Remember: Test-Time Training via Context Distillation (NeurIPS 2026 · similarity 0.89)
• ATTS: Asynchronous Test-Time Scaling via Conformal Prediction (ICLR 2026 · similarity 0.89)
• Compute Aligned Training: Optimizing for Test Time Inference (NeurIPS 2026 · similarity 0.89)
• The Surprising Effectiveness of Test-Time Training for Few-Shot Learning (ICML 2025 · similarity 0.88)