Trajectory-Refined Distillation

TL;DR 针对 on-policy 蒸馏的 prefix failure 问题，提出轨迹级修正蒸馏(TRD)，在教师指导下修正学生 rollouts 后再蒸馏，缓解梯度碎片并提升推理正确率与覆盖度。

问题

On-policy distillation (OPD) 已成为大型语言模型后训练的核心范式，通过沿学生自身生成轨迹（rollout）施加逐 token 教师监督来提升推理质量。然而，现有 OPD 方法普遍存在一个结构性缺陷——prefix failure：当学生生成的前缀偏离教师分布时，逐 token 的 KL 监督会退化为双峰混合分布，导致梯度碎片化，使模型无法聚焦单一有效的修正路径。基于 token 级损失截断或重加权的常见干预手段未能根治此问题，因为它们仍忽略轨迹的整体连贯性。

这一失败模式之所以关键，是因为它直接削弱了 OPD 在复杂推理任务中的效用。数学推导、代码生成等场景要求模型在中间步骤犯错后仍能自我修正，但碎片化梯度会阻碍模型从完整正确序列中学习，并限制其对替代有效推导路径的探索。业界正积极寻求既能保留 on-policy 探索优势，又能提供连贯校正信号的方法。

可以类比：在自动驾驶行为克隆中，若逐帧纠正司机错误而不考虑轨迹上下文，得到的修正动作可能互相矛盾；更稳健的做法是重规划一段完整轨迹。

核心洞察

• Prefix failure 揭示了 token-level 蒸馏的内在缺陷：密集的逐 token 监督导致教师混合分布双峰化与梯度碎片化，使 token 级损失截断或重加权等局部干预失效。这一观察将改进方向从损失函数微调转向轨迹级结构化修正，为蒸馏优化提供了更高层次的视角。
• TRD 通过在蒸馏前由教师修正学生 rollout 的前缀，从源头消除错误传播，同时引入教师引导的替代推理路径，即便原始解答正确也能扩展搜索空间。该方法将教师信号从分布匹配升级为生成过程的结构性引导，显著提升了 on-policy 学习的探索效率与覆盖范围。

方法

输入

学生模型在训练问题上的采样轨迹（rollout）及对应的教师模型（或自蒸馏中的特权学生）。

关键模块：轨迹级修正

1. 前缀失败检测：在标准的在线策略蒸馏（OPD）中，教师对学生的逐 token 分布进行监督。但研究发现，当学生生成的某个前缀（prefix）存在瑕疵时，教师在该前缀下的条件分布会呈现出双峰混合（bimodal mixture），导致后续 token 的梯度方向碎片化，形成“前缀失败”。
2. 轨迹级重写：TRD 不直接调整 token 级损失权重，而是对学生的整个 rollout 进行修正。利用教师指导（如通过额外的 refinement prompt），为每条轨迹生成一条修正后的版本，同时确保修正轨迹仍在学生策略的支持内（on-policy support），避免偏离自身分布。
3. 蒸馏执行：在修正后的轨迹上执行标准的 token 级知识蒸馏（如正向 KL 散度或反向 KL 散度），因为此时前缀已被修复，教师分布不再双峰，梯度稳定一致。

输出

得到一个在修正轨迹上蒸馏的学生模型，有效提升单次尝试准确率并拓宽推理覆盖范围。

与同类方法的差异

与仅依靠 token 级损失截断（truncation）或重加权（reweight）的方法不同，TRD 从根源上解决梯度碎片化：通过轨迹级修正，将 OPD 重新对齐到教师单一分布的指导下，避免了碎片化的教师混合分布对训练造成的扭曲。

实验

实验设计

实验覆盖 On-Policy Distillation (OPD) 和 On-Policy Self-Distillation (OPSD) 两种范式，在多种规模的基础模型上评测。作者构造了数学与代码推理的训练轨迹，对比 TRD 与多种 token-level 基线（如 loss truncation、reweighting），并分析 rollout 过程中的 prefix failure 现象。评估指标包括单次准确率（single-attempt accuracy）和多次采样的 Pass@k，以同时衡量精度和推理覆盖度。

关键发现

TRD 在所有规模和基准上一致优于先前方法，prefix failure 是 OPD 性能瓶颈的直接原因：token-level 密集监督会形成 bimodal teacher mixture 和碎片化梯度，而 TRD 通过在蒸馏前修正学生 rollout 中的错误前缀，从源头缓解了这一问题。即使原始 rollout 已经正确，TRD 也能借助 teacher 引导探索替代推导路径，从而拓宽推理覆盖。Trajectory 分析表明，TRD 的修正动作单调提升验证器准确率，并且缩短了解题路径长度。在困难问题上，Pass@k 的提升尤为明显，说明 TRD 有效扩展了探索前沿。

与基线对比解读

Token-level 干预（截断或重加权）无法解决 prefix failure，因为它们仍然作用于被 bimodal mixture 污染的损失信号。TRD 的 trajectory-level 修正直接修改 rollout 数据分布，使得后续蒸馏的梯度更加连贯。OPSD 设置下，参数共享的 teacher 同样能从这种修正中受益，表明方法具有通用性。TRD 对过滤操作敏感——简单丢弃错误轨迹反而会损害 Pass@16，说明保留失败样本并修正比直接剔除更能提升探索多样性。

行业影响

落地场景

TRD 直接作用于大语言模型的后训练蒸馏阶段，尤其适用于需复杂多步推理的下游应用。典型场景包括：

• 代码助手：用强模型作为教师，蒸馏出轻量代码生成模型，提升补全和建议的准确率，同时减少延迟。
• 数学 / 逻辑解题：在数学辅导或自动推理系统中，蒸馏模型能生成更完整的解题路径，覆盖更多有效解法。
• 对话式 AI：在需要多轮推理的客服或知识助手场景，小模型通过 TRD 获得接近大模型的推理能力。

商业价值

TRD 主要带来推理成本下降与输出质量提升的双重收益：

• 降本：蒸馏出的小模型可本地部署，避免高频调用大体量教师 API，显著降低单次推理成本；同时小模型推理延迟更低，支持更高并发。
• 体验提升：通过纠正前缀，模型首次生成正确答案的比例（pass@1）和答案多样性（pass@k）均有提升，减少用户反复重试，直接改善交互体验。
• 覆盖更广：尤其在困难问题上的扩展能力，让蒸馏模型具备处理长尾复杂输入的能力，扩大产品可服务范围。

与现有工作流的集成

TRD 作为蒸馏前的轨迹级预处理步骤，可无缝嵌入现有训练流水线：

• 训练侧：在生成学生 rollout 后，用教师模型对输出序列进行前缀纠正（仅修改导致梯度碎片的起始片段），再执行标准蒸馏损失。与现有 token 级损失截断 / 重加权方法兼容。
• 推理侧：产出的蒸馏模型为标准自回归模型，无需特殊推理引擎，可直接用 vLLM、TGI 等框架部署。
• 可叠加：TRD 同样适用于自蒸馏场景（OPSD），可在无外部教师时仅用学生自身的高质量分支进行折优。

具体落地用例

1. 电商智能客服：用 GPT-4o 等大模型作为教师，蒸馏出 7B 级别的客服模型。当用户询问诸如“我买的两件衣服怎么组合优惠最多”等需要推理的问题时，TRD 蒸馏模型能给出步骤清晰、正确的答案，避免因前缀错误导致中途偏离。本地部署的 7B 模型推理成本仅为 API 调用的几十分之一，且响应时间可控在 200ms 内。

2. 代码辅助平台：在 IDE 插件或代码托管平台中，使用 TRD 从强代码模型（如 DeepSeek-Coder-V2）蒸馏出 1B–3B 参数的轻量模型，用于实时代码补全和 bug 修复建议。当原始补全可能因前缀歧义而错误时，TRD 纠正前缀使得后续生成更符合正确语法和逻辑，减少用户手动修正的次数，提升开发效率。

集成难度低，只需在现有蒸馏框架中增加一次教师引导的轨迹修正步骤，即可获得显著收益。

局限

• **计算开销与延迟增加**：TRD 在蒸馏前需对学生每次 rollout 执行轨迹级提炼，依赖额外教师前向传播。与标准 on-policy 蒸馏相比，这显著增加了训练计算负担和时间成本。对于大规模生产环境或需要高频迭代的应用，该开销可能成为瓶颈。论文未给出详细的吞吐与延迟对比分析，也未讨论如何通过缓存或少次提炼降低开销。
• **对教师模型质量的依赖**：TRD 高度依赖教师模型提供准确的纠正信号。若教师本身在特定领域有偏差或能力不足，提炼步骤可能引入噪声甚至误导学生，尤其当师生能力差距大时，教师的纠正可能偏离学生应遵循的合理探索路径，反而损害泛化性。论文仅在强教师设定下验证，未分析教师退化情况。
• **任务与模型泛化性未充分验证**：实验集中于数学推理和代码生成任务，基模型主要为 Qwen3-8B 等中等规模模型。TRD 假设学生生成具有清晰前缀结构的推理链，对于开放域对话、摘要等非推理式任务，轨迹级纠正的收益和适用性未知。此外，方法在更大规模模型、不同架构上的效果有待检验。