Learning What to Remember: Test-Time Training via Context Distillation¶
会议: NeurIPS 2026
arXiv: 2608.01672
代码: https://github.com/dangxingyu/ttcd
领域: 模型压缩
关键词: 测试时训练, 上下文蒸馏, 快速权重记忆, 长上下文语言建模, 在位自适应
一句话总结¶
针对传统测试时训练（TTT）与循环注意力架构在有限容量下缺乏对历史信息“未来效用”建模的瓶颈，本文提出测试时上下文蒸馏（TTCD），以长短窗口两路视图的表征残差作为自监督信号更新 MLP 快速权重，以极低开销实现了对远端关键信息的定向记忆与高效检索。
研究背景与动机¶
随着大语言模型在超长链条推理、长文档理解以及具身智能轨迹规划中的广泛应用，长上下文建模已成为扩展模型能力的核心基石。标准 Transformer 的全注意力机制虽然能够无损保留并精确检索完整的历史上下文，但其计算与显存开销随序列长度呈二次方爆炸增长，成为制约大规模落地的算力壁垒。人类在阅读长篇代码库或大部头专著时，并不需要对历史字句进行无差别的机械复写，而是会主动提取并压缩具有长远预测价值的关键事实，同时主动遗忘无关噪声。受此启发，近来涌现的线性循环网络（如 DeltaNet、Gated DeltaNet）与测试时训练（Test-Time Training, TTT）试图将历史上下文动态压缩至固定容量的循环隐状态或动态演化的序列级“快速权重”（Fast Weights）中，以实现恒定的推理计算开销。
然而，在有限容量状态的强约束下，模型面临的核心难题在于：每一个时间步究竟应该选择性地写入、持久化保留或主动遗忘哪些信息。现有 TTT 的目标函数设计暴露出明显的结构性短板： 基于自编码或特征重构的目标函数（如 LaCT、Titans）本质上要求隐状态均匀地记忆整个上下文历史，无法区分高价值的预测线索与琐碎的局部噪声；而直接采用下一个 Token 预测（NTP）的目标函数又极易陷入近因偏差（Recency Bias），使得快速权重的梯度更新主要被局部前驱词的显式共现所主导，难以将预测增益有效归因到遥远的背景语境中。
本文的切入视角是引入上下文蒸馏（Context Distillation）范式，将其改造成一个具备前瞻性的在线自监督记忆准则。通过在共享骨干网络下构建具备远端视野的长窗口“自教师”（Teacher）与仅见局部近邻的短窗口“自学生”（Student），两者在当前 Token 处投影隐状态的残差便天然量化了“远端上下文对局部预测所带来的纯增益”。核心 idea：将长短滑动窗口的隐状态差异作为测试时自监督蒸馏损失，以在线方式将远端关键信息蒸馏至 MLP 固有投影矩阵中作为动态快速权重，使有限容量模型在测试时能依据信息的未来效用主动学会“该记什么”。
方法详解¶
整体框架¶
测试时上下文蒸馏（TTCD）统一了上下文压缩与测试时在线自适应。给定输入序列分块，在每一块计算中，共享同一套慢权重的骨干网络同时运行两路因果滑动窗口：长窗口自教师（窗口大小为 \(w_T\)）观测包含丰富远端前缀的上下文，而短窗口自学生（窗口大小为 \(w_S < w_T\)）仅观测最近邻局部上下文。在进入 MLP 模块后，两者隐状态表征的差距被转化为表征级自监督蒸馏损失，通过梯度反传更新当前模块的快权重。
在实际实现上，本文提出在位变体 IP-TTCD（In-Place TTCD），复用标准 Transformer 中 SwiGLU MLP 现有的下投影矩阵（\(W_{\text{down}}\)）作为承载动态记忆的快权重。为了消除跨块顺序递归带来的串行瓶颈，IP-TTCD 采用离策略（Off-policy）近似在基准权重上并行计算两路激活，结合轻量一维因果卷积与特征归一化，通过因果前缀扫描完成块级在线更新。推理阶段，通过两路滑动窗口注意力融合算子（Dual-Window Fusion Kernel）实现键值对加载与内积复用，几乎不引入额外推理延迟。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入序列切块<br/>Chunk C_i (大小 C)"] --> B["双窗口因果注意力并行表征<br/>长窗口教师 w_T + 短窗口学生 w_S"]
B --> C["双路激活特征提取与局部混合<br/>短因果卷积 conv_T / conv_S"]
C --> D["隐状态残差与键向量构建<br/>残差 R_i 与学生归一化键 K_i"]
D --> E["在位快权重更新与因果前缀累加<br/>低秩修正 ΔW_down = β_i R_i K_i^T"]
E --> F["输出表征与下游推理生成<br/>远端知识注入当前局部预测"]关键设计¶
1. 隐状态上下文蒸馏损失：以表征差异作为前瞻性记忆监督信号
传统 TTT 采用自重构目标会导致容量被无关背景均分，而 NTP 目标又被局部强信号占据。TTCD 针对这一痛点，通过设计长短窗口两路计算，精确提炼出远端背景对当前预测的增量贡献。在第 \(i\) 个数据块 \(\mathcal{C}_i\) 内，对于每个 Token 位置 \(t\)，分别获取前缀处于长窗口 \(w_T\) 与短窗口 \(w_S\) 下经过 MLP 激活前的隐状态 \(h_t^{(T)}\) 和 \(h_t^{(S)}\)，进而得到对应的 MLP 输出 \(y_t^{(T)}\) 与 \(y_t^{(S)}\)。蒸馏损失直接对齐二者表征，并截断教师端梯度： $$ \mathcal{L}_i^{\mathrm{CD}} = \frac{1}{|\mathcal{C}i|} \sum) \big|_2^2 $$ 该自监督损失迫使学生去拟合“如果拥有更远视野本应具备的表征”，从而在没有任务真实标签的情况下，自发以紧凑的特征梯度指引模型将具有未来预测潜力的远端语义注入记忆系统。}_i} \frac{1}{2} \big| y_t^{(S)} - \mathsf{stopgrad}(y_t^{(T)
2. 在位快速权重与离策略并行更新：消除架构冗余并打破递归串行壁垒
为了实现低成本的即插即用部署，IP-TTCD 避免引入独立的外部记忆模块，而是将 SwiGLU MLP 中的下投影矩阵 \(W_{\text{down}}\) 视作序列内演化的快权重。标准顺序递归要求块 \(i\) 的前向计算依赖块 \(i-1\) 更新后的权重，严重阻碍并行化。IP-TTCD 提出离策略（Off-policy）近似机制：在所有块内，教师和学生的前向中间激活一律基于初始慢权重 \(W_{\text{down}}^{(0)}\) 提取，而后将各块计算出的低秩校正量通过因果关联进行累加。这种解耦使得各块的注意力与 MLP 激活能够完全并行前向，将时序更新退化为极轻量的参数矩阵前缀和扫描，完美适配现代 GPU 的长上下文吞吐需求。
3. 轻量特征变换与低秩修正算子：提升记忆写入的稳定性与可检索性
直接对原始高维激活求外积容易导致写入记忆的噪声发散并出现检索漂移。IP-TTCD 在更新规则中引入了结构化增强设计：首先对教师和学生的中间激活 \(Z_i^{(T)}\) 和 \(Z_i^{(S)}\) 分别施加轻量深度短因果卷积 \(\text{conv}_T\) 和 \(\text{conv}_S\)，引入局部多 Token 混合表征；随后引入可学习线性变换 \(W_{\text{proj}}\) 和学生激活的 \(\ell_2\) 范数归一化，分别构建“值残差矩阵”\(R_i\) 和“查询键矩阵”\(K_i\)： $$ R_i = W_{\text{down}}^{(0)} W_{\text{proj}} \big(\bar{Z}_i^{(T)} - \bar{Z}_i^{(S)}\big), \qquad K_i = \ell_2\text{-norm}\big(\bar{Z}_i^{(S)}\big) $$ 块级的在位参数低秩更新由此定义为 \(\Delta W_{\text{down}}^{(i)} = \beta_i R_i K_i^\top\)。其中 \(R_i\) 精确指定了“长短视差蕴含的增量信息”（写入什么内容），而规范化的 \(K_i\) 则充当寻址键（指定未来何时以何种条件被重新检索激活），形成了兼具数学可解释性与数值稳定性的动态联想记忆机制。
4. 理论保障：上下文蒸馏在序列内部的跨时间步解算与重放
在理论层面，论文严格证明了上下文蒸馏能够实现上下文学习（ICL）解向后续相关查询的精确迁移。设上下文片段 \((x_{c_1}, \ldots, x_{c_2})\) 位于远端，前序查询点 \(q_1\) 处由于具备长窗口教师支持而产生非琐碎的预测信号增益：\(y_{q_1}^{(T)\top} (y_{q_1}^{(T)} - y_{q_1}^{(S)}) \ge c_{\text{gap}} > 0\)。对于远在后文出现且自身注意力无法覆盖支持上下文的相关查询点 \(q_2\)，只要其学生端激活特征与 \(q_1\) 保持正向对齐（\(z_{q_1}^{(S)\top} z_{q_2}^{(S)} \ge c_{\text{align}} > 0\)），前序块更新所诱导的学生表征漂移满足： $$ y_{q_1}^{(T)\top} \Delta y_{q_2}^{(S)} \ge \eta \cdot c_{\text{gap}} \cdot c_{\text{align}} $$ 该结论在理论上保证了：即使原始事实已完全滑出局部注意力窗口，前序蒸馏写入快权重的残差增益依然能在后续查询点被精准回放激活。
损失函数 / 训练策略¶
模型整体训练损失采用标准下一个词预测交叉熵（对最终生成的 Token 序列）与中间各 TTT 层的上下文蒸馏自监督损失两阶段/多任务协同。在从头预训练阶段，优化器采用专门针对矩阵参数优化的 Muon 配合 WSD（Warmup-Stable-Decay）学习率调度，教师卷积初始化为恒等或零，学生卷积随机初始化；而在持续预训练阶段（Continual Pre-training, CPT），则使用 AdamW 优化器，保持学生卷积零初始化以保证初始模型表征无扰动，并在 64K 长文本语料（如 ProLong 密集长文本集）上进行 10B Token 的微调扩容。
实验关键数据¶
主实验¶
论文在从头预训练与持续预训练两大设定下展开验证。主实验对比了 760M 参数规模下各主流长序列建模架构在 32K 长度预训练后的长上下文大海捞针检索能力（RULER NIAH 平均准确率）、常识推理能力以及长程推理（BABILong 平均准确率）。
| 模型架构 (760M) | RULER NIAH (4K) | RULER NIAH (8K) | RULER NIAH (16K) | RULER NIAH (32K) | 常识推理 Avg (Zero-shot) | BABILong (32K) |
|---|---|---|---|---|---|---|
| DeltaNet | 41.67% | 38.25% | 24.83% | 19.33% | 37.38% | 1.94% |
| Gated DeltaNet | 32.21% | 29.88% | 22.29% | 18.62% | 39.40% | 2.84% |
| 滑动窗口注意力 (SWA) | 44.88% | 22.58% | 10.79% | 4.21% | 39.08% | 4.40% |
| In-Place TTT (IP-TTT) | 59.83% | 33.88% | 23.92% | 9.29% | 39.24% | 4.46% |
| IP-TTCD (本文) | 65.79% | 46.92% | 41.75% | 21.96% | 39.90% | 6.73% |
在持续预训练设置下，将 IP-TTCD 插件式集成到开源主流基座模型中，在 ProLong 64K 数据上微调 10B Token 后，在 RULER 综合长上下文基准（全部 13 个子任务）上测评：
| 基座模型与持续预训练配置 | RULER (4K) | RULER (8K) | RULER (16K) | RULER (32K) | RULER (64K) | 全长度综合平均 (Avg) |
|---|---|---|---|---|---|---|
| SmolLM2-360M (Base 原始) | 58.13% | 39.75% | 0.00% | 0.00% | 0.00% | 19.58% |
| SmolLM2-360M + YaRN-8x 外推 | 35.59% | 7.21% | 6.28% | 6.18% | 4.23% | 11.90% |
| SmolLM2-360M + 全注意力 CPT | 50.53% | 42.74% | 39.92% | 27.24% | 10.09% | 34.10% |
| SmolLM2-360M + IP-TTT CPT | 50.20% | 42.43% | 40.47% | 29.32% | 10.35% | 34.55% |
| SmolLM2-360M + IP-TTCD CPT (本文) | 49.79% | 42.13% | 40.59% | 33.55% | 9.79% | 35.17% |
| SmolLM2-1.7B + 全注意力 CPT | 63.74% | 57.35% | 54.74% | 44.18% | 13.38% | 46.68% |
| SmolLM2-1.7B + IP-TTT CPT | 67.41% | 60.13% | 54.89% | 41.53% | 14.44% | 47.68% |
| SmolLM2-1.7B + IP-TTCD CPT (本文) | 67.61% | 59.58% | 55.50% | 42.96% | 19.17% | 48.96% |
| LLaMA-3.1-8B + 全注意力 CPT | 93.87% | 92.58% | 87.32% | 65.19% | 25.30% | 72.85% |
| LLaMA-3.1-8B + IP-TTCD CPT (本文) | 93.72% | 92.32% | 88.59% | 66.48% | 27.65% | 73.75% |
消融实验¶
在合成序列记忆评测基准 MAD（包含上下文召回、含噪召回、模糊召回、记忆压缩与选择性复制）以及 340M 语言模型架构组件消融中，全面评估各关键设计模块的有效性：
| 评估基准与测试配置 (MAD Benchmark) | In-context Recall | Noisy Recall | Fuzzy Recall | Compress | Memorize | Selective Copy | Overall 综合得分 |
|---|---|---|---|---|---|---|---|
| DeltaNet (循环基线) | 99.9% | 86.7% | 25.1% | 42.3% | 55.5% | 100.0% | 67.8% |
| Gated DeltaNet (门控基线) | 99.9% | 99.7% | 18.2% | 47.3% | 86.9% | 99.9% | 74.4% |
| Large Chunk TTT (LaCT) | 100.0% | 100.0% | 18.6% | 48.3% | 80.7% | 99.8% | 73.8% |
| In-Place TTT (IP-TTT) | 98.1% | 97.1% | 52.4% | 43.4% | 82.9% | 96.6% | 78.4% |
| IP-TTCD (完整模型) | 99.3% | 98.8% | 56.2% | 50.0% | 88.9% | 99.2% | 82.1% |
| 340M 架构设计组件消融变量 | 滑动窗口困惑度 (PPL @ 32K) | 影响说明 |
|---|---|---|
| 完整设计 (Short Conv + Learnable Proj \(W_{\text{proj}}\)) | 最佳 (曲线最低，持续下降至 32K) | 协同利用局部特征平滑与最优空间对齐，泛化效果最佳 |
| 仅短卷积 (w/o \(W_{\text{proj}}\)) | 明显劣于完整模型 (PPL 上升约 0.8) | 缺乏线性投影时跨空间表征残差匹配难度加大 |
| 仅线性投影 (w/o Short Conv) | 劣于完整模型 (PPL 上升约 1.2) | 丧失局部时序邻域混合，对噪声敏感 |
| 变换块大小 Chunk Size = 512 | 基准表现良好 | 细粒度更新但吞吐受限 |
| 变换块大小 Chunk Size = 2048 | 表现更优且硬件利用率最高 | 大块更新平滑梯度噪声，契合张量计算范式 |
关键发现¶
- 32K 极端检索优势翻倍：在从头预训练的 760M 模型上，当序列扩展至 32K 时，由于 SWA 彻底丧失远端信息（仅剩 4.21%），IP-TTT 因目标受限跌落至 9.29%，而 IP-TTCD 仍维持 21.96% 的检出率，性能领先 IP-TTT 达到 2.36 倍，证实了上下文蒸馏对长程信息抗衰减的独到价值。
- 攻克模糊与记忆压缩瓶颈：在 MAD 合成基准中，前期诸如 Gated DeltaNet 与 LaCT 虽在精确召回和复制任务上触顶（~100%），但在最具挑战性的模糊召回（Fuzzy Recall）和信息压缩（Compress）任务上表现疲软（低于 50%）。IP-TTCD 在这两项上分别大幅突破至 56.2% 和 50.0%，带动综合总分从 78.4% 跃升至 82.1%。
- 大分块计算收益更高：消融实验揭示，增大 TTT 的分块粒度（从 512 增至 2048）非但没有破坏记忆的时效性，反而凭借更充沛的跨 Token 上下文聚合降低了瞬时梯度抖动，同时显著释放了 GPU 张量核心的吞吐红利。
亮点与洞察¶
- 反直觉的记忆监督源：传统模型往往以标签数据或自重构为目标，本文跳出定势，巧妙地将“同一模型不同上下文视野的输出差值”重塑为无监督记忆分配信号，为自监督持续学习提供了优雅的降维视角。
- 在位双赢的架构轻量化：摒弃了外挂复杂循环细胞或专用神经网络的重型方案，直接原位更新 MLP 的下投影矩阵，在几乎零新增常驻参数的前提下赋予了静态 Transformer 动态终身记忆机制。
- 软硬件协同的双窗口融合优化：针对自注意力与学生注意力的重叠区间，利用 Triton 开发融合单次扫描算子与 LSE 组合技巧，将双路窗口计算开销压制在基准的 1%~6% 误差范围内，打破了算法创新必伴随巨大延迟的宿命。
局限与展望¶
- 作者承认的局限：朴素实现的双路滑动窗口计算在小教师窗口时会带来高达 1.3 倍的推理 Prefill 耗时开销；尽管经过定制 Triton 算子优化将开销削减，但训练耗时相较标准 Transformer 仍有 1.1~1.2 倍的固定额外开销；此外，从头预训练受限于算力仅验证至 760M 规模。
- 实验与设定层面的局限：在持续预训练设置中，虽然全注意力评测有提升，但在滑动窗口推断（SWA）场景下各基线差异收敛在 1 个点以内，说明快权重的记忆提取仍高度依赖注意力机制的局部提示能力。
- 后续演进方向：未来可探索向超长多轮对话与自主 Agent 历史轨迹存储的迁移；进一步结合低精度量化（如 FP8/FP4）的快权重就地累加，进一步削减大模型部署时的显存带宽开销。
相关工作与启发¶
- vs IP-TTT (In-Place Test-Time Training)：IP-TTT 确立了利用 MLP 下投影作为快权重的有效范式，但其依赖输入嵌入对齐或下一个词预测（NTP）作为测试时训练目标，易产生近因偏置且更新不稳定。本文引入长短窗口上下文蒸馏，提供了更具全局前瞻性、且具备明确数学收敛证明的记忆写入目标。
- vs DeltaNet / Gated DeltaNet：Delta 规则系列依赖矩阵循环状态进行线性外积写入，受限于固定的显式状态更新法则，难以复用成熟的预训练 Transformer 权重。IP-TTCD 完美保留了标准 Transformer 预训练权重体系，既能从头训练也能以持续学习形式低成本接入已有大模型。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ [将长短窗口上下文视差转化为测试时自监督记忆更新信号，概念精巧且洞察深刻]
- 实验充分度: ⭐⭐⭐⭐⭐ [覆盖合成基准、从头预训练多尺度评测、持续预训练拓展至 8B 模型及定制算子效能实测]
- 写作质量: ⭐⭐⭐⭐⭐ [逻辑链条严密，理论推导详尽完备，图表绘制精当]
- 价值: ⭐⭐⭐⭐⭐ [为 Transformer 克服长上下文二次方计算复杂度与记忆分配难题指引了切实可行的演进路径]