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
发布日期: 2026-08-03
💡 一句话要点
提出测试时间上下文蒸馏方法以优化长上下文建模
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 长上下文建模 测试时间训练 上下文蒸馏 自监督学习 语言模型
📋 核心要点
- 现有的测试时间训练方法未能充分考虑保留信息的未来效用，导致长上下文建模效果不佳。
- 本文提出的TTCD框架通过自监督目标优化信息保留，利用长窗口教师指导短窗口学生的快速权重更新。
- 实验结果显示，IP-TTCD在长上下文语言建模任务中显著优于多种基线方法，展现出良好的适应性和性能提升。
📝 摘要（中文）
有效的长上下文建模不仅仅是保留更多的过去信息，而是要保留未来可能相关的信息。测试时间训练（TTT）是一种在线参数更新的有效方法，但现有的TTT方法仅优化重建或在线适应目标，而未考虑保留信息的未来效用。本文提出了测试时间上下文蒸馏（TTCD）框架，引入自监督目标以分配有限的内存容量用于未来使用。TTCD使用长窗口教师来监督短窗口学生的快速权重，二者之间的隐藏状态差异提供了密集的自监督信号，指导模型记忆对未来标记预测至关重要的上下文信息。我们专注于一种就地变体：就地TTCD（IP-TTCD），它使用现有的MLP参数作为快速权重。实验表明，IP-TTCD在长上下文语言建模任务中始终优于DeltaNet、Gated DeltaNet、滑动窗口注意力和TTT。
🔬 方法详解
问题定义：本文旨在解决现有测试时间训练方法在长上下文建模中未能有效考虑未来信息保留的问题。现有方法往往只关注重建或在线适应目标，导致信息利用不充分。
核心思路：TTCD框架通过引入自监督目标，优化信息的保留与利用。具体而言，利用长窗口教师的指导来监督短窗口学生的快速权重更新，从而提高模型对未来上下文的记忆能力。
技术框架：TTCD的整体架构包括长窗口教师和短窗口学生两个主要模块。教师模型负责提供稳定的上下文信息，而学生模型则快速适应当前输入，通过两者之间的隐藏状态差异进行自监督学习。
关键创新：TTCD的核心创新在于引入了自监督信号来优化信息保留的未来效用，这与现有方法的单一目标优化形成了显著区别。通过这种设计，模型能够更有效地记忆对未来预测重要的上下文信息。
关键设计：在IP-TTCD中，使用现有的MLP参数作为快速权重，确保了模型在推理过程中的高效性。同时，设计了特定的损失函数来量化教师与学生之间的隐藏状态差异，以指导模型的学习过程。
🖼️ 关键图片
📊 实验亮点
实验结果表明，IP-TTCD在长上下文语言建模任务中相较于DeltaNet、Gated DeltaNet和滑动窗口注意力等基线方法，表现出显著的性能提升，尤其是在从零开始预训练的情况下，展现出更强的适应性和效果。
🎯 应用场景
该研究的潜在应用领域包括自然语言处理、对话系统和长文本生成等场景。通过优化长上下文建模，TTCD能够提升模型在处理复杂语言任务时的表现，具有重要的实际价值和未来影响。
📄 摘要（原文）
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.