Learning What to Remember: Test-Time Training via Context Distillation
2026-08-03 • Computation and Language
Computation and LanguageArtificial IntelligenceMachine Learning
AI summaryⓘ
The authors study how to help language models remember and use important information from earlier in a long text. They propose a method called Test-Time Context Distillation (TTCD), where a large "teacher" model guides a smaller "student" model to focus on remembering useful past details. Their in-place version, IP-TTCD, updates existing model parts during reading to better predict future words. Experiments show IP-TTCD works better than previous methods, and it can improve pre-trained models with a small change during inference. This approach helps models learn to handle longer contexts more effectively over time.
long-context modelingtest-time trainingself-supervised learningcontext distillationtransformer modelshidden statesparameter updateslanguage modelingcontinual learningMLP (multi-layer perceptron)
Authors
Zixuan Wang, Xingyu Dang, Rui-Jie Zhu, Zixin Wen, Hengyu Fu, Wenhao Chai, Jason D. Lee
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 \textbf{T}est-\textbf{T}ime \textbf{C}ontext \textbf{D}istillation (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.