🤖 AI

Trajectory-Refined Distillation

本文介绍了轨迹精炼蒸馏（Trajectory-Refined Distillation, TRD），这是一种新颖的在策略（on-policy）蒸馏方法，通过在教师指导下对学生模型的展开进行轨迹级修正，缓解了结构性的“前缀失效”（prefix failure）问题，从而提升了在各种基准测试和模型规模下的推理准确性与覆盖度。

原作者： Li Jiang, Haoran Xu, Yichuan Ding, Amy Zhang

发布于 2026-06-09

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Li Jiang, Haoran Xu, Yichuan Ding, Amy Zhang

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在教一名学生（AI）如何解决一个复杂的数学问题。你有一位非常出色的老师（一个更高级的 AI），他知道正确答案。

在被称为**在策略蒸馏（On-Policy Distillation, OPD）**的标准方法中，流程如下：

学生尝试独立解决问题。
当学生写下步骤时，老师会观察并说道：“对于你刚刚写的这个特定词汇，下一个词出现的概率是这样的。”
学生试图调整自己的大脑，以匹配老师给出的这些概率。

问题：“前缀失败”（Prefix Failure，即“走错路”）

该论文指出，这种标准方法存在一个主要缺陷，称为前设失败（Prefix Failure）。

想象一下，学生正在沿着一条路径解题。在第 3 步时，他们转错了弯。他们现在进入了一条死胡同。

老师的困境： 老师看到学生迷路了。老师知道正确的路径应该以“等等，实际上……”开头，然后转向完全不同的方向。
混乱： 但老师被迫必须关注学生当前所在的位置（死胡同）。因此，老师的建议变成了一种混乱的混合体：“留在这个死胡同里，因为你已经在这里了”以及“立即转身”。
结果： 学生得到了一个混乱的信号。他们试图向老师学习，但因为学生卡在了错误的路径上，老师的建议变得支离破碎。学生学到了如何“擅长迷路”，而不是学习如何回到正轨。论文称之为“双峰混合（bimodal mixture）”——老师试图同时教授两个相反的东西，导致学生感到困惑。

现有的解决方案试图通过仅仅微调特定词汇的“评分”（损失函数）来修复这个问题，但它们并没有解决学生仍然卡在错误路径上的本质问题。

解决方案：轨迹精炼蒸馏（Trajectory-Refined Distillation, TRD）

作者提出了一种新的方法——轨迹精炼蒸馏（TRD）。与其仅仅对学生的每一个混乱词汇进行评分，不如改变工作流程：

草稿： 学生仍然对问题进行第一次尝试（“原始展开/raw rollout”）。
重写： 在评分之前，老师会拿到那个混乱的草稿并对其进行重写。老师修正了错误的转向，纠正了逻辑，并基于学生的起始点生成一个清晰、完美的解决方案版本。
教学： 学生随后从这个清晰、经过修正的版本中学习，而不是从混乱的原稿中学习。

类比：

旧方法： 你写了一个故事，但在中间出现了一个情节漏洞。编辑读完后说：“在你第三段的‘猫’这个词处，你应该写‘狗’。”但你仍然卡在第三段，试图弄清楚为什么你最初写的是“猫”。这让人很困惑。
TRD 方法： 你写了一个带有情节漏洞的故事。你的编辑拿走你的草稿，修复了那个情节漏洞，并重写了整个中间部分，使故事流畅自然。然后，你通过学习编辑的重写版本来学习如何写出更好的故事。

为什么这种方法更好

解决了根本原因： 通过在学生尝试学习之前就修正“错误的转向”，老师的建议不再是混乱的。它变成了一条清晰、单一的路径。
探索新路径： 即使学生第一次尝试时就答对了，老师也可能会展示一种不同的、更短或更聪明的解题方式。这教会了学生具备更高的灵活性。
效率： 论文发现，“重写”后的路径通常比学生原本啰嗦的尝试要短得多且更清晰，这实际上加快了训练速度。

结果

研究人员在困难的数学竞赛（如 AIME 和 HMMT）和编程挑战上测试了该方法。他们发现：

新方法（TRD）始终优于旧方法。
它在帮助 AI 解决那些通常会让 AI 陷入困境的最难问题方面表现尤为出色。
无论老师是一个独立的、更大的 AI，还是同一个 AI 利用“特权信息（privileged information）”扮演其自身的“聪明”老师，该方法都有效。

简而言之，论文认为，要有效地教导 AI，你不应该仅仅在它犯错的过程中逐词地评判它的错误；相反，你应该让老师先修正整条路径，然后再让学生从这段修正后的旅程中学习。

技术摘要：轨迹精炼蒸馏 (Trajectory-Refined Distillation, TRD)

1. 问题陈述：在策略蒸馏中的前缀失效 (Prefix Failure)

本文识别了大型语言模型 (LLM) 在在策略蒸馏 (On-Policy Distillation, OPD) 和在策略自我蒸馏 (On-Policy Self-Distillation, OPSD) 中存在的一个根本性的结构性局限，称之为前缀失效 (prefix failure)。

在标准的 OPD/OPSD 中，学生模型生成一个展开过程（轨迹） $y_o$ ，教师模型基于该特定轨迹提供密集的逐 token 监督。作者认为，当学生模型的展开过程采取了错误的推理路径（即“失败的前缀”）时，教师的监督信号在结构上会受到损害：

双峰教师混合 (Bimodal Teacher Mixture)： 当前缀错误时，教师分布会变成两种模式的混合：一种是继续当前的失败前缀（为了保持序列一致性），另一种是转向正确的解法。
碎片化梯度 (Fragmented Gradients)： 标准的逐 token 损失函数（KL 散度）是在冻结的失败轨迹上评估监督的。因此，梯度更新是碎片化的。教师不断地在错误的上下文（失败的前缀）条件下推荐修正 token（例如，“Wait”或“Actually”），而不是基于正在展开的修正路径进行条件化。
Token 级修复的无效性： 现有的补救措施（如损失裁剪、Top-K 截断、重加权）都在 token 损失层面进行操作。它们试图选择“信息丰富的”token 或抑制“噪声”token，但未能解决根源问题：失败的前缀本身在结构上仍然保持完整。传递给学生的监督对与理想的修正路径相背离，导致模型无法有效地学习如何从错误中恢复。

2. 方法论：轨迹精炼蒸馏 (TRD)

为了缓解前缀失效，作者提出了轨迹精炼蒸馏 (TRD)，该方法将干预层级从 token 层面提升到了轨迹层面。

核心机制

TRD 不直接使用原始的学生展开 $y_o$ 进行蒸馏，而是在计算损失之前引入了一个精炼步骤：

采样 (Sampling)： 学生模型 $\pi_\theta$ 生成一个原始的在策略展开 $y_o \sim \pi_\theta(\cdot | x)$ 。
精炼 (Refinement)： 教师模型 $\pi_T$ $π_{T}$ （或在 OPSD 中，由自身作为教师，并以如标准答案 $y^*$ $y^{*}$ 等特权信息为条件）生成一个精炼轨迹 $y_r$ $y_{r}$ 。
- 在 OPD 中： $y_r \sim \pi_T(\cdot | x, y_o)$ 。
- 在 OPSD 中： $y_r \sim \text{sg}[\pi_\theta(\cdot | x, y_o, y^*)]$ 。
- 至关重要的是，通过提示教师根据参考解法来修正学生的展开过程，从而有效地“重写”错误的路径部分，同时保留学生在有效部分的推理模式。
蒸馏 (Distillation)： 使用精炼后的轨迹 $y_r$ 作为标准分布匹配损失（针对全词表的 Forward KL）的监督目标。

理论依据

TRD 通过确保监督上下文沿着修正路径而非失败的前缀增长，解决了“碎片化梯度”问题。

理想梯度恢复： 通过对 $y_r$ 进行监督，每个 token 的上下文都包含了修正后的前缀，这使得教师能够提供连贯的多步修正信号 ( $g_{ideal}$ )，而不是重复且上下文不匹配的信号 ( $g_{frag}$ )。
探索 (Exploration)： 即使在 $y_o$ 正确的情况下，TRD 也会让学生接触到教师建议的其他有效推导路径，从而扩大了学生学习到的正确推理路径的分布。

算法

该过程总结于算法 1：

采样 $y_o \sim \pi_\theta(\cdot | x)$ 。
生成 $y_r \sim \pi_T(\cdot | x, y_o)$ 。
使用 $\pi_\theta$ 与教师在 $y_r$ 上的分布之间的 KL 散度来更新 $\theta$ 。

3. 核心贡献

识别前缀失效： 本文将“前缀失效”形式化为在策略蒸馏性能下降的一个结构性原因，解释了为什么 token 级的干预无法从推理错误中恢复。
轨迹级修正： TRD 被提出为首个显式针对恢复失败前缀的轨迹级优化设计，同时在不产生高昂计算开销的情况下尊重了在策略支持集 (on-policy support) 的约束。
双重机制改进： TRD 在两个维度上提升了性能：
- 错误修正： 它通过在蒸馏前修正有问题的路径来缓解前缀失效。
- 探索： 它通过向学生展示替代的有效推导过程，扩大了推理覆盖范围，即使原始展开是正确的。
适用于自我蒸馏： 该方法可以自然地扩展到 OPSD，其中学生利用特权信息作为自身的教师，从而无需额外的教师模型。

4. 实验结果

作者在五个竞赛数学基准测试（AIME24/25, HMMT25, BeyondAIME, AMOBench）和代码生成基准测试（HumanEval+, MBPP+, LiveCodeBench）上评估了 TRD，使用的是 Qwen3 模型（1.7B, 4B, 8B）。

主要发现

卓越性能： 在 OPD 和 OPSD 设置下，TRD 在各项指标上均一致优于先前的基准模型（包括使用/不使用裁剪/截断的 Forward/Reverse KL）。
在困难任务上显著增益： 在最难的基准测试中，提升最为明显。
- AMOBench： 在 OPSD 设置下，相比基础模型，TRD 使 Qwen3-8B 的 Pass@16 实现了 ~50% 的相对提升。
- AIME24/25： TRD 显示出持续的增益，特别是对于较小的模型（例如，Qwen3-1.7B 在 AIME24 上提升了 +4.6%）。
稳定性： 不同于某些在特定基准测试上表现下降的基线（例如 Reverse KL 在 AIME24 上性能下降），TRD 从未低于基础模型的性能水平。
效率： 在 OPSD 设置中，精炼后的轨迹 $y_r$ 比原始展开 $y_o$ 短得多（中值长度压缩了约 9 倍），尽管增加了采样步骤，但仍减少了约 60% 的训练实际耗时。
代码生成： TRD 在 HumanEval+ 和 MBPP+ 上达到或超过了基准水平，但在 LiveCodeBench 上的增益有限，这表明教师的精炼能力可能受限于任务难度。

5. 重要性与主张

本文声称 TRD 代表了 LLM 后训练策略的一次必要转变。作者认为：

Token 级干预是不充分的： 前缀失效的结构性特征要求进行轨迹级的修正。
保留了在策略支持集： TRD 在实现这些修正的同时，并未违反在策略约束（不同于使用固定的专家策略），使其能够兼容标准的计算高效型学习范式。
自我提升是可行的： TRD 在 OPSD 设置中的成功表明，模型可以利用特权信息（地面真值/Ground Truth）有效地精炼自身的推理路径，而无需依赖另一个更大的教师模型，这为轻量化自我提升提供了一条路径。

这项工作表明，未来的蒸馏方法必须考虑推理轨迹的结构完整性，而不仅仅是单个 token 的统计对齐。