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

2026-08-03

AI Summary

This work addresses the inefficiency of existing test-time training (TTT) methods in long-context modeling, which often neglect the utility of retained information for future predictions. The authors propose a Test-Time Context Distillation (TTCD) framework that leverages a long-window teacher model to supervise the fast weights of a short-window student model. By exploiting discrepancies in hidden states, TTCD generates a novel self-supervised signal oriented toward future predictive utility, dynamically guiding the model to prioritize memorizing information critical for subsequent predictions. This approach introduces, for the first time, a future-utility-based self-supervised objective combined with in-situ MLP parameter updates, yielding a lightweight continual learning architecture. Experiments demonstrate that IP-TTCD, trained from scratch, outperforms DeltaNet, Gated DeltaNet, sliding-window attention, and conventional TTT methods on long-context language modeling tasks, and further enables pretrained Transformers to effectively acquire long-context capabilities through inference-time continual pretraining.

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.

Problem

• long-context modeling
• test-time training
• memory allocation
• future relevance
• context distillation

Innovation

• Test-Time Training
• Context Distillation
• Long-Context Modeling
• Fast Weights
• Continual Learning