Skip to content
Preprint

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

Aug 2026 · 0 citations · 55 references
Computer Science

TL;DR

This work proposes TTCD, a TTT framework that introduces a self-supervised objective for allocating limited memory capacity for future use, and focuses on an in-place variant: In-Place TTCD, which uses the existing MLP parameters as the fast weights.

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.

View source

Similar papers

Preprint Jul 2026

Self-Guided Test-Time Training for Long-Context LLMs

A simple method, Self-Guided TTT (S-TTT), which improves accuracy for both Qwen3-4B-Thinking-2507 and Llama-3.1-8B-Instruct, achieving up to a 15% relative improvement.

Xinyu Zhu, Zhenqin Xu, Xiaohan Wei et al. · 0 citations
Preprint Aug 2026

Rethinking Expressivity and Efficiency in Test-Time Training

Test-Time Training (TTT) enables long-context processing via continuous weight updates during inference, but current methods struggle to balance the expressivity of per-token update dynamics with the hardware efficiency of chunk-wise approximations. We propose E$^2$-TTT (Expressive and Efficient TTT) to bridge this gap. Under the standard approximation of taking gradients at the chunk-start weights, we derive a closed-form state transition that exactly reproduces the chunk-end fast-weight and momentum states of the per-token recurrence. This enables fully parallelized chunk-level training while preserving the temporal structure of the update rule that prior chunk-wise methods discard. We validate E$^2$-TTT by training models up to 1.3B parameters from scratch. It performs on par with previous TTT and hybrid attention baselines in language modeling while outperforming them on in-context retrieval. Its advantage is most pronounced in length extrapolation: on the standard ``Needle in a Haystack''passkey test, it retains over 90% accuracy at $8\times$ the training context length. Meanwhile, E$^2$-TTT can match the training throughput of efficient chunk-wise methods, demonstrating that it effectively reconciles expressivity with efficiency. The code is available at https://github.com/zeyun-zhong/E2-TTT.

Zeyun Zhong, Joya Chen, Manuel Martín et al. · 1 citation
Preprint Aug 2026

Where to Look Matters: On-Policy Self-Distillation for Long-Video Understanding

Vision-language models (VLMs) have made substantial progress in long-video understanding, with standard backbone models typically answering questions from frames sampled across the full video. However, as videos become longer, the full-video context inevitably contains more question-irrelevant temporal content, which can distract the model from the evidence needed to answer a specific question. We empirically find that focusing the visual input on short annotated clue intervals containing question-relevant evidence consistently improves prediction accuracy across model scales compared with using the corresponding full videos, while requiring fewer input frames. Based on this finding, we introduce Clue-OPSD, a clue-privileged on-policy self-distillation framework for long-video understanding. During training, a full-video student learns from a self-teacher conditioned on the corresponding clue interval by aligning their next-token distributions along student-generated trajectories. Clue-OPSD thus uses clue intervals as privileged supervision without relying on ground-truth answer labels, while requiring no clue annotations or additional modules at inference time. Extensive experiments across multiple long-video understanding benchmarks and Qwen3.5 model scales demonstrate consistent improvements over the corresponding backbone models and strong performance against supervised post-training baselines.

Kaishen Wang, Dong-Di Zhao, Yijun Liang et al. · 0 citations
Preprint Aug 2026

Anti-Shortcut Distillation via Temporal Negative Knowledge Transfer

Knowledge distillation (KD) trains a compact student by attracting it towards a converged teacher. It is silent about which directions the teacher itself learned to suppress: repulsive and bias-aware objectives exist, but none exploits the teacher's own trajectory to identify what the student should avoid. We observe that the missing signal is already encoded in the teacher's optimization trajectory: features that an early-stage teacher emphasizes but that a converged teacher attenuates are precisely the shortcut directions worth pushing the student away from. We instantiate this observation as \textbf{A}nti-\textbf{S}hortcut \textbf{D}istillation (ASD), a push--pull KD framework that treats the converged teacher $\Tfinal$ as a positive semantic anchor and an early-checkpoint teacher $\Tearly$ as a temporal negative reference. ASD couples two losses: a temporal contrastive loss ($\Ltc$) that places the early-teacher feature as a same-sample negative against in-batch and memory-bank final-teacher features in an InfoNCE objective; and a shortcut suppression loss ($\Lss$) that penalizes student projection onto the top eigenvectors of $\E[\Dh\Dh^{\top}]$, the uncentered second-moment matrix of early-to-final feature displacements. Across 13 teacher--student pairs on CIFAR-100, ImageNet-100, and TinyImageNet, ASD attains the highest clean top-1 accuracy on more than 10 pairs and outperforms standard KD on 12. On CIFAR-100-C corruption robustness, ASD obtains the lowest mean Corruption Error ($86.1$\,mCE) on the most challenging cross-architecture pair (WRN-40-2$\to$ShuffleNet-V2). Mechanistic diagnostics confirm the intended geometry: the ASD student is systematically anti-aligned with the shortcut direction, while its projection onto the robust subspace is substantially larger ($0.45$ vs.\ $0.12$).

Syed Muhammad Raza, Omer Tariq, J. Son · 0 citations
#artificial intelligence Review Aug 2026

When to Review: Spaced Repetition for Continual Pre-Training of Language Models

Spaced Repetition Training (SRT) is introduced, a continual learning framework inspired by cognitive science, which schedules sample-rehearsal using the SuperMemo-2 (SM-2) algorithm, and preserves broad benchmark performance that naive continual pre-training and uniform replay substantially degrade.

Alankar Atreya, Devesh Batra, Yoages Kumar Mantri et al. · 0 citations
Preprint Jul 2026

Train Smarter, Not Longer: Memorization-Guided Data Reuse for Efficient LLM Training

This work proposes "Memorization-guided Data Reuse", a training paradigm that adaptively determines when and how data should be reused, enabling principled decisions on the number of training epochs and the scheduling of data replays.

Jingwei Zuo, Cong Zeng, Ilyas Chahed et al. · 0 citations