Learning What to Remember: Test-Time Training via Context Distillation
Zixuan Wang, Xingyu Dang, Rui-Jie Zhu, Zixin Wen, Hengyu Fu, Wenhao Chai, Jason D. Lee
cs.CL, cs.AI, cs.LG
Submitted: 2026-08-03
Code: https://github.com/dangxingyu/ttcd
License: http://creativecommons.org/licenses/by-nc-sa/4.0/
The gist: Effective long-context modeling is not merely about retaining more of the past, but about preserving the information that may prove relevant later.
Terminology
Abstract
Effective long-context modeling is not merely about retaining more of the past, but about preserving the information that may prove relevant later. Test-time training (TTT) is an appealing approach that performs online parameter updates for long-context modeling, yet existing TTT methods only optimize either reconstruction or online adaptation objectives without considering the future utility of retained information. In this work, we propose Test-Time Context Distillation (TTCD), a TTT framework that introduces a self-supervised objective for allocating limited memory capacity for future use. Specifically, TTCD uses a long-window teacher to supervise the fast weights of a short-window student, where the hidden-state discrepancy between them offers a dense, self-supervised signal guiding the model to memorize the contextual information crucial for future token predictions. We focus on an in-place variant: In-Place TTCD (IP-TTCD), which uses the existing MLP parameters as the fast weights. Experiments on long-context language modeling tasks show IP-TTCD consistently outperforms DeltaNet, Gated DeltaNet, sliding-window attention, and TTT when pre-trained from scratch. Furthermore, IP-TTCD allows pre-trained transformer models to adapt their parameters during inference through continual pre-training, gaining long-context capabilities with only a lightweight architectural augmentation. Our results position TTCD as a step toward architectural continual learning.
Sources
- SmolLM2: When Smol Goes Big -- Data-Centric Training of a Small Language Model
- An Empirical Evaluation of Generic Convolutional and Recurrent Networks for Sequence Modeling
- Titans: Learning to Memorize at Test Time
- Lessons from the Trenches on Reproducible Evaluation of Language Models
- KV-Distill: Nearly Lossless Learnable Context Compression for LLMs
- Empirical Evaluation of Gated Recurrent Neural Networks on Sequence Modeling
- Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge
- Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality
- Cartridges: Lightweight and general-purpose long context representations via self-study
- Hungry Hungry Hippos: Towards Language Modeling with State Space Models
- The Pile: An 800GB Dataset of Diverse Text for Language Modeling
- Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs
- Mamba: Linear-Time Sequence Modeling with Selective State Spaces
- Efficiently Modeling Long Sequences with Structured State Spaces
- Log-Linear Attention
- Measuring Massive Multitask Language Understanding
- RULER: What's the Real Context Size of Your Long-Context Language Models?
- MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies
- Characterizing Prompt Compression Methods for Long Context Inference
- Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention
Related papers
- Exploring Solution Divergence and Its Effect on Large Language Model Problem Solving
- Ishigaki-IDS-Bench: A Benchmark for Generating Information Delivery Specification from BIM Information Requirements
- Subliminal Steering: Stronger Encoding of Hidden Signals
- MedStruct-S: A Benchmark for Key Discovery, Key-Conditioned QA and Semi-Structured Extraction from OCR Clinical Reports
- The End of Transformers? On Challenging Attention and the Rise of Sub-Quadratic Architectures
- Untangling the Mechanisms of Misleading Context in Medical Question Answering