arXiv:2608.01672cs.CLcs.AI2026-08

让模型在推理时学会记住未来有用的信息,提升长文本理解能力。

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

  • 用教师-学生结构动态分配记忆容量,指导模型记住关键上下文。
  • 在长文本建模任务中优于DeltaNet等方法,提升推理阶段表现。
  • 仅需轻量级改造,即可让预训练模型获得持续学习能力。

有效的长上下文建模不仅在于保留更多历史信息,更在于保存未来可能有用的那部分。测试时训练(TTT)是一种在线更新参数的长上下文建模方法,但现有方法仅优化重建或在线适应目标,未考虑保留信息的未来价值。本文提出测试时上下文蒸馏(TTCD),通过自监督目标为未来使用分配有限记忆容量。具体地,利用长窗口教师监督短窗口学生,隐藏状态差异提供密集自监督信号,引导模型记忆对后续词预测至关重要的上下文信息。我们聚焦一种就地变体——就地TTCD(IP-TTCD),使用现有MLP参数作为快速权重。在长上下文语言建模任务上,IP-TTCD在从零训练时持续优于DeltaNet、Gated DeltaNet、滑动窗口注意力和传统TTT。此外,IP-TTCD使预训练Transformer模型在推理阶段通过持续预训练调整参数,仅通过轻量级架构改进即获得长上下文能力。结果表明TTCD是迈向架构级持续学习的重要一步。

原文摘要 · Abstract (English)

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.

测试时训练长上下文持续学习自监督

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。