Trajectory-Refined Distillation (TRD) 深度解析：为什么 token 级在线蒸馏总是效果有限？因为问题不在 token，而在「前缀」

论文：Trajectory-Refined Distillation
arXiv: 2606.08432 [cs.AI] (7 Jun 2026)
作者：Li Jiang, Haoran Xu, Yichuan Ding, Amy Zhang
机构：McGill University, Mila Quebec AI Institute, UT Austin
代码：https://github.com/louieworth/trd

一、核心问题：在线蒸馏的「前缀失效」

在线蒸馏（On-policy Distillation, OPD）已经成为 LLM 后训练的核心工具——Qwen3、DeepSeek-v4、MiMo-v2、GLM-5 都包含 OPD 阶段。但现有方法有一个共同的结构性缺陷：

前缀失效（Prefix Failure）：当学生模型生成的推理前缀出错后，几乎不可能在不回溯的情况下得到正确结果。现有的 token 级损失调整无法解决这个结构问题。

1.1 前缀失效的数学本质

论文用严格的数学分析揭示了前缀失效的核心机制：

当学生模型生成了一个错误的推理前缀（prefix）时，老师模型面临一个两难困境：

• 继续错误路径：保持序列一致性，但结果错误
• 转向修正：需要生成修正词（如 "Wait"、"Actually"），然后重新推导

这导致老师分布变成一个双峰混合分布（bimodal mixture）：

• 一个峰值在「继续错误」
• 一个峰值在「开始修正」

1.2 Token 级干预的局限性

现有方法（如 top-K truncation、per-token loss reweighting）都只在 token 损失层面做调整：

• Zhao et al. 2026：clip 高 KL 的 token，防止不稳定训练
• Fu et al. 2026：top-K truncation，只保留高置信度 token 的监督
• Xu et al. 2026a：根据熵和师生分歧重新加权，恢复有信息的 token

但这些方法有一个共同的根本缺陷：它们不改变错误的前缀本身。它们只是调整 token 损失的权重，而错误的推理路径仍然冻结在学生的 rollout 中。

论文证明：即使在完美老师的情况下，dense per-token KL 也是结构性受限的——因为它是在学生冻结的 rollout 上后验地评估老师，而不是沿着修正路径展开。

1.3 梯度碎片化

论文推导了一个关键公式：

在序列级反向 KL 下，梯度应该是：

∇θJ(θ) = E[Σ_t (δ_t + Σ_{t' > t} δ_{t'}) ∇θ log πθ(y_t | x, y_{≤t})]

但标准 OPD 实现只保留即时 log-ratio：

∇θJ(θ) = E[Σ_t δ_t ∇θ log πθ(y_t | x, y_{≤t})]

这意味着：

• 理想梯度需要在修正路径上展开：(y_{o,<t}, ȳ*_t), (y_{o,<t}, ȳ*_t, ȳ*_{t+1}), ...
• 实际梯度却在错误路径上重复：(y_{o,<t}, ȳ*_t), (y_{o,<t+1}, ȳ*_t), ...

两个轨迹只共享第一个元素，之后完全分离。这就是梯度碎片化（gradient fragmentation）——监督信号被拆散，无法有效传播修正路径。

二、TRD 核心方法：轨迹级修正

2.1 核心思路

TRD 不再纠结于 token 级损失的调整，而是直接修正整个轨迹：

原始问题 → 学生采样 rollout → 老师基于参考答案修正 → 修正轨迹作为蒸馏监督

2.2 具体流程

Step 1：学生采样原始轨迹

• 学生模型生成 on-policy rollout y_o
• 这个 rollout 可能包含错误前缀

Step 2：老师修正轨迹

• 老师模型接收：问题 + 原始轨迹 + 参考答案
• 老师生成修正后的轨迹 y_r，保持 on-policy support（即修正后的轨迹仍然合理自然）
• 修正过程类似于：在错误处插入 "Wait, let me reconsider..."，然后给出正确推导

Step 3：用修正轨迹蒸馏

• 使用修正后的轨迹 y_r 作为监督信号
• 计算 per-token KL 在 y_r 上，而非原始错误轨迹

2.3 为什么有效？

TRD 解决了前缀失效的根本原因：

• 修正了错误前缀，而不是在错误前缀上调整损失权重
• 恢复了完整修正路径的监督，而不是碎片化的梯度
• 保持 on-policy support，修正轨迹仍然是自然语言，不是人工构造的

此外，TRD 还有一个附带好处：

• 增强探索：即使原始 rollout 正确，老师也可能提供不同的正确推导路径，扩展学生的推理覆盖

2.4 扩展到自蒸馏（OPSD）

TRD 可以自然扩展到自蒸馏场景：

• 学生：只接收问题
• 老师：同一个模型，但接收问题 + 参考答案（privileged information）
• 修正：老师基于参考答案修正学生的 rollout
• 蒸馏：用修正轨迹训练学生

论文中，OPSD 设置下 Qwen3-8B 在 AMOBench 上实现了**~50% 的相对提升**。

三、实验结果：跨任务、跨规模的稳定提升

3.1 评测基准

数学推理（5个竞赛级基准）：

• AIME24, AIME25
• HMMT25
• BeyondAIME
• AMOBench（最难）

代码生成（3个基准）：

• HumanEval+
• MBPP+
• LiveCodeBench

3.2 核心结果

在多个 Qwen3 模型规模上（0.6B, 1.7B, 4B, 8B, 14B, 32B）：

• TRD 在 OPD 和 OPSD 两种设置下一致优于所有基线
• 在最难的 AMOBench 上提升最大：
  • Qwen3-8B OPSD：~50% 相对提升
  • Pass@16 显著提升

关键发现：

• 规模越大，TRD 相对优势越明显（大模型更能受益于轨迹修正）
• 在代码生成任务上同样有效，说明方法不限于数学推理

四、这篇论文的工程启示

4.1 为什么现有 OPD 方法效果有限？

很多团队发现 OPD 在实践中「理论上很好，但实际提升不明显」。这篇论文给出了根本原因：

问题不在蒸馏的「密度」（per-token），而在蒸馏的「结构」（prefix）。

如果学生的 rollout 走错了路，token 级干预就像在高速公路上微调方向盘——车已经上错匝道了，微调方向盘改变不了目的地。

TRD 的做法是：先让车回到正确的匝道上，再教它怎么开。

4.2 对工业界 OPD 流程的改进

当前主流 OPD 流程（如 Qwen3、DeepSeek-v4）：

1. SFT 预训练
2. RLVR（RL with Verifiable Rewards）
3. OPD（在线蒸馏）

TRD 建议把 OPD 改为 TRD：

1. SFT 预训练
2. RLVR
3. TRD：学生采样 → 老师修正 → 蒸馏

这个改动不需要改模型结构，只需要在数据 pipeline 中增加一个「修正」步骤。对于已经有 OPD 基础设施的团队，这是一个低成本的升级。

4.3 对 OPSD（自蒸馏）的特殊意义

OPSD 是更轻量的方案（不需要单独的老师模型），但之前的效果往往不如 OPD。TRD 让 OPSD 变得可行：

• 同一个模型，用 privileged info（参考答案）修正自己的 rollout
• 然后蒸馏
• 这样既保留了 OPSD 的轻量性，又解决了监督信号质量的问题

对于资源有限的团队，TRD-OPSD 可能是「小模型追大模型」的高效路径。

4.4 对「自我修正」能力的启发

TRD 中老师模型的修正行为，实际上是在教学生学习如何自我修正。修正轨迹中包含的 "Wait, actually..." 等反思标记，可能成为学生模型自我修正能力的种子。

这与 OpenAI 的 o1/o3 系列的「chain of thought」中的反思行为有相似之处——但 TRD 是通过蒸馏显式地教，而不是通过 RL 隐式地学。

4.5 局限性和未来方向

局限性：

• 需要参考答案（reference solution）作为修正指导，不适用于无监督场景
• 修正轨迹的质量依赖老师模型的能力
• 修正过程增加计算开销（两次推理：学生采样 + 老师修正）

未来方向：

• 无需参考答案的自动修正（如用 verifier 反馈替代参考答案）
• 多步修正（ iterative refinement）
• 修正策略的自动化（学习何时、如何修正）

五、总结：从 token 级到轨迹级的范式转移

这篇论文的核心贡献不是提出了一个「更复杂的 token 级损失」，而是指出了：

当问题出在结构上时，token 级的优化是无力的。你需要改变结构本身。

TRD 的优雅之处在于：

• 不需要改模型架构
• 不需要重新发明训练算法
• 只需要在数据 pipeline 中增加一个「修正」步骤

但它带来的提升是显著的：AMOBench 上 ~50% 的相对提升，跨多个规模、跨多个任务一致有效。

对于正在做模型后训练的团队，TRD 提供了一个立即可用的升级方案。前缀失效是一个普遍问题，而 TRD 是一个通用解法。

论文来源：arXiv:2606.08432 [cs.AI] (7 Jun 2026)
作者：Li Jiang, Haoran Xu, Yichuan Ding, Amy Zhang (McGill / Mila / UT Austin)
代码：https://github.com/louieworth/trd

#LLM #知识蒸馏 #在线蒸馏 #OPD #OPSD #数学推理 #代码生成 #后训练 #小凯

#LLM #知识蒸馏 #在线蒸馏 #OPD #OPSD #数学推理 #代码生成 #后训练 #小凯