Research map / Architectures and efficiency / Attention and transformers

Learning What to Remember: Test-Time Training via Context Distillation

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: State space and recurrent sequence models · 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 Test-Time Context Distillation (TTCD), where a long-window teacher supervises fast weights of a short-window student, and an in-place variant (IP-TTCD) that updates MLP down-projections, improving long-context modeling from scratch and via continual pre-training.

Problem

Test-time training objectives (reconstruction or next-token prediction) do not explicitly prioritize retaining past information that will be useful for future predictions under fixed-capacity fast-weight memory.

Approach

Use the hidden-state discrepancy between a long-window self-teacher and a short-window student (both sliding-window views of the same model) as a dense self-supervised online update target for fast weights, implemented in-place on the MLP down-projection with a chunk-wise, off-policy update and a learned low-rank correction.

Key results

• At 760M from scratch, IP-TTCD reaches 21.96 RULER NIAH average at 32K versus 9.29 for IP-TTT, and 39.90 common-sense average versus 39.24 for IP-TTT.
• On LLaMA-3.1-8B continual pre-training, IP-TTCD improves RULER average to 73.75 vs 73.39 for IP-TTT and 72.85 for CPT (64K context).
• On SmolLM2-360M continual pre-training under full attention, IP-TTCD achieves 33.55 at 32K vs 29.32 for IP-TTT; on SmolLM2-1.7B it reaches 19.17 at 64K vs 14.44 for IP-TTT.

Surprising findings

• IP-TTT with input-embedding target outperforms its recommended current-layer hidden-state target, which the authors found unstable.
• Under SWA inference, the CPT, IP-TTT and IP-TTCD variants perform within about one RULER point of each other on SmolLM2.

Limitations

• Naive two-kernel implementation slows prefill up to 1.3x at small teacher windows; training 1.1-1.2x slower than Transformer baselines.
• From-scratch pre-training limited to 760M parameters; scaling to larger models from scratch remains open.
• Theoretical analysis is a single-layer, simplified one-step argument under stated assumptions.

Experimental setup

Datasets: Long-Data-Collections Books3 ProLong 64K book subset RULER BABILong MAD

Benchmarks: RULER BABILong MAD HellaSwag ARC-Easy ARC-Challenge MMLU sliding-window perplexity on Books

Metrics: perplexity accuracy average accuracy

Baselines: DeltaNet Gated DeltaNet sliding-window attention IP-TTT CPT Base (Yarn-8x) LaCT

Builds on: In-Place TTT DeltaNet Gated DeltaNet context distillation LaCT

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)
• Test-Time Training Provably Improves Transformers as In-context Learners (ICML 2025 · similarity 0.90)
• Rethinking Expressivity and Efficiency in Test-Time Training (NeurIPS 2026 · similarity 0.89)
• Titans: Learning to Memorize at Test Time (NeurIPS 2025 · similarity 0.88)
• The Surprising Effectiveness of Test-Time Training for Few-Shot Learning (ICML 2025 · similarity 0.88)