轨迹精炼蒸馏：通过轨迹级修正解决策略内蒸馏的前缀失效问题 Trajectory-Refined Distillation
TRD通过轨迹级修正解决OPD前缀失效问题
前置知识
On-Policy Distillation (OPD)
策略内蒸馏是大语言模型后训练的核心技术，它不在固定语料上训练，而是在学生模型自己生成的轨迹上计算教师的逐token监督信号。给定输入x，学生采样rollout，然后在学生访问的每个前缀上评估教师分布，通过逐token的KL散度损失进行训练。这种方法保留了策略采样的性质和密集的token级学习信号。
OPD是现代LLM后训练包括Qwen3、DeepSeek-v4等的核心组件，理解其工作原理和失效模式是读懂本文的基础
On-Policy Self-Distillation (OPSD)
策略内自蒸馏是OPD的特殊形式，教师和学生来自同一个模型但在不同上下文条件下实例化，从而消除了单独教师模型的需求。给定问题-解答对，教师策略接收特权信息如参考答案或推理轨迹，教师和学生共享相同参数。标准OPSD损失在教师分支上使用stop-gradient操作。
OPSD是轻量级替代方案，也是本文重点验证的场景之一，理解其与OPD的区别有助于理解TRD的适用范围
Prefix Failure
前缀失效是本文形式化的结构性问题，指学生的rollout走上了错误的推理路径，该前缀的任何延续都无法在不回溯或反思的情况下到达正确解。令事件表示前缀失效事件即前缀包含与参考解矛盾且不能延伸到参考解的推理错误。当该事件成立时，教师变成双峰混合分布：一个峰继续失败的前缀，另一个峰转向参考解，导致密集监督变成噪声甚至对抗信号。
前缀失效是本文识别的核心问题，是所有token级干预方法无法解决的结构性瓶颈，理解它才能理解TRD的必要性
Forward KL vs Reverse KL
前向KL散度是模式覆盖的，由教师分布加权；反向KL散度是模式寻找的，由学生分布加权。在前缀失效下，前向KL被修正起始区域主导，其模式覆盖性质强迫学生进入该OOD模式，可能导致模式崩塌；反向KL被错误延续区域主导，因为修正信号影响有限，更新集中在失败轨迹而非恢复token上。
KL方向的选择与前缀失效的相互作用是token级方法失效的关键机制，理解这一点才能明白为什么TRD采用前向KL
Fragmented Gradient
碎片化梯度指在学生冻结的rollout上评估逐tokenKL导致的结构问题。理想教师会沿着修正路径自回归展开，产生监督对。但密集KL下，上下文由冻结的学生轨迹决定而非展开的修正，教师提供的监督对形成碎片化序列，上下文沿着错误延续增长而目标停在修正起始token，导致梯度在完全不同的上下文-令牌对上评估。
碎片化梯度是前缀失效下token级方法的根本局限性，即使完美教师也无法恢复修正路径，这解释了为什么需要轨迹级方法
研究动机
现有的策略内蒸馏方法在实践中经常表现不佳，exhibiting多种失效模式，包括模式崩塌、轨迹膨胀、监督信号消失甚至误导学生。最近的行业发布如Qwen3、DeepSeek-v4、MiMo-v2和GLM-5都包含OPD阶段，但研究表明vanilla OPD和OPSD配方往往达不到预期。例如，在OPSD训练过程中，通过一阶段验证器结果分割的正确和错误rollout的逐tokenKL都保持接近零，表明教师和学生仍然对齐，密集KL在前缀失效饱和后不再提供信号。教师和学生困惑度差距在整个训练过程中保持接近零，特权条件提供的增量监督几乎消失。此外，教师在6到8千分之一的每个位置质量上分配给16个认知起始token，匹配碎片化梯度预测的修正起始token重复签名。这些观察表明，即使有理想教师，逐tokenKL在结构上也受到限制，因为它是沿着学生rollout的后验逐token目标。
本文的目标是本文的具体目标是解决OPD和OPSD中的前缀失效问题。作者认识到直接优化前缀失效是棘手的：前缀是否失效只有在完整轨迹验证后才会揭示，而定位失效索引需要搜索词汇表规模的k次方个长度为k的延续。因此，作者将目标放宽到轨迹级代理，在学生策略的支持约束下最大化数据集上的预期验证器通过率。注意这是轨迹级分布支持约束，优化目标镜像标准RLVR目标，但在OPD更新中不直接优化此目标，而是定义上游轨迹构建任务，在标准的参数上OPD优化之前，构建在当前学生支持内获得更高验证器通过率的轨迹。
与已有工作不同的是，本文的独特切入角度是从轨迹层面而非token层面解决问题。现有的缓解措施通过损失设计操作，将冒犯的前缀保持不变，包括裁剪高KL项、教师top-K截断、重新加权损失等。这些措施虽然在相反方向控制token级主导——裁剪抑制过度主导的前向KL项，截断或重新加权放大欠加权的反向KL教师模式——但都保持失败前缀不变。作者认识到这些失效与前缀失效机制一致，即使没有明确这样表述。在前向KL下，教师加权的修正或OOD模式可能已经主导损失；在反向KL下，修正模式被学生加权目标欠加权。所有这些干预在访问的对集内重新加权项，而不是将梯度移动到修正路径对集上。因此，本文提出在轨迹层面优化，这是第一个明确针对优化目标同时尊重策略内约束而不会产生禁止的额外计算开销的轨迹级优化设计。
核心方法
轨迹精炼蒸馏TRD的整体思路是在轨迹层面进行修正，而非在token层面进行损失干预。直觉上，当学生的原始轨迹出现前缀失效时，与其在token损失上做各种裁剪和重新加权，不如让教师直接帮助学生修正整个轨迹。技术路线上，TRD首先从当前策略采样一个原始的策略内rollout，然后请求教师构建精炼轨迹，方式为从教师分布采样。在OPSD中，相同的backbone通过额外条件化参考解实现此教师查询。关键地，在原始rollout上条件化将精炼轨迹锚定到学生策略已经展示的推理模式，即策略支持内，而教师重写错误部分以直接缓解前缀失效。精炼轨迹随后用于后续OPD更新的监督。沿着精炼轨迹的逐tokenKL减少对双峰教师混合的暴露，并恢复理想梯度，因为监督上下文沿着精炼轨迹而非原始rollout增长。TRD还通过标准OPD提升学生的探索性：在正确的原始rollout上，标准OPD提供很少新信号，它倾向于仅仅强化学生已经产生的高概率解决方案路径；相反，精炼轨迹从教师分布采样，揭示相同答案的替代有效推导，即教师建议但很少从学生策略单独采样的路径，从而扩展学生受监督的正确推理轨迹集合。
TRD的核心创新在于从token级损失干预转向轨迹级输出修正。与现有方法在损失函数层面操作包括裁剪、截断、重新加权不同，TRD在蒸馏之前修正学生轨迹本身。具体来说，给定问题-解答对，首先采样原始策略内rollout，然后提示教师模型产生该rollout的精炼版本，由参考解引导。精炼轨迹然后用作后续蒸馏的监督。TRD的关键insight是：在原始rollout上条件化将精炼轨迹锚定到学生策略已经展示的推理模式即策略支持内，而教师重写错误部分以直接缓解前缀失效。这与两个极端选择形成对比：设置分布等于学生策略退回到标准OPD，策略约束满足但未能缓解前缀失效；设置分布等于专家策略即产生参考解的专家策略，精确达到验证通过但一般违反策略内支持约束，因此落在可行集之外并破坏OPD依赖的策略内特征。TRD通过轨迹层面优化优化目标同时尊重策略内约束而不会产生禁止的额外计算开销，这是首次明确针对该目标的轨迹级优化设计。
方法步骤详情
TRD方法的完整步骤如下：首先，给定来自训练数据集的提示，学生模型采样一个自回归rollout，这是原始的学生生成轨迹。第二步，在OPD设置中，教师模型基于原始rollout和提示生成精炼轨迹；在OPSD设置中，由于教师和学生共享相同backbone，教师策略接收特权信息包括参考解，生成精炼轨迹，其中stop-gradient操作符应用于教师分支。第三步，使用精炼轨迹计算蒸馏损失，具体采用前向KL散度与全词汇表匹配。第四步，通过梯度下降更新学生参数。这个过程在训练过程中迭代进行，每批次都先采样原始rollout，再用教师精炼，最后在精炼轨迹上计算损失更新参数。整个算法在保持策略内支持的同时，通过轨迹级修正直接缓解前缀失效，为蒸馏提供更有效的监督信号。
技术新颖性
TRD的技术新颖性体现在多个方面：首先，这是首个明确针对轨迹级代理优化目标同时尊重策略内约束的方法。其次，TRD揭示了前缀失效的结构机制，包括双峰教师混合和碎片化梯度，为token级OPD方法的失效提供了理论基础。第三，TRD在轨迹层面而非token层面解决问题，避免了loss-level干预的结构性局限。第四，TRD可以自然扩展到自蒸馏设置OPSD，显示其通用性。第五，TRD在精炼轨迹上使用前向KL与全词汇表匹配，提供模式覆盖监督并稳定梯度估计，与反向KL相比更好地利用了精炼轨迹。第六，TRD不仅缓解失败前缀，还扩展学生的探索性，即使在正确的原始rollout上也能提供替代有效推导。最后，实验验证表明TRD缓解了原始rollout上观察到的问题：逐tokenKL恢复，教师和学生困惑度差距打开，教师的认知起始质量减少到原始rollout的3倍以下，表明教师信号真正在训练过程中传递给学生。
实验结果
本文在OPD和OPSD设置下评估了TRD与四个密集KL基准，跨越数学和代码基准，组织评估围绕三个问题：TRD在OPD和OPSD设置下的性能如何，探索和利用权衡如何；在相同学生规模下，哪种精炼信号对TRD更有效；精炼如何改变训练轨迹和测试时rollout行为。在OPD设置下使用Qwen3-8B作为教师，平均16次采样结果显示TRD在两个学生规模上都改进了基线模型的利用，每个块的八个基准中有七个最佳或并列最佳。增益对于较小的Qwen3-1.7B学生最大，例如AIME24上提升4.6个百分点。Qwen3-4B-Instruct-2507块更具诊断性：几乎所有在原始rollout上训练的OPD变体都无法匹配基线模型，相比之下，在精炼轨迹上训练保留了较强学生的基线能力并将它们转化为广泛增益。Pass@16结果显示增益集中在较难的数学基准上，TRD在两个规模上都给出最佳AMOBench结果，对Qwen3-1.7B改进基线5.1个百分点，对Qwen3-4B-Instruct-2507改进12.8个百分点。在OPSD设置下，平均16次采样结果显示TRD在两个规模和每个基准上都是最佳的，从不低于基线。Pass@16结果显示TRD的轨迹级精炼与token级干预最清晰分离，在Qwen3-8B上，TRD在AMOBench和HMMT25上分别提升50%相对增益和15%。轨迹分析显示验证器通过率从原始rollout上的65.8%提高到精炼轨迹上的81.4%，长度分布压缩约9倍即中位数7.7千字符降到0.88千字符向参考解规模即参考解中位数约0.49千字符，训练墙钟时间减少约60%抵消了额外精炼轨迹采样成本。在AMOBench评估rollout上，TRD略短于平均长度即18.9千字符降到18.5千字符，覆盖增益在k等于1时 modest 但随采样预算扩大，在k等于128时达到53.8%对比46.7%。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| AIME24数学推理 | Avg@16 | 49.4% (OPD, 1.7B) | Base 44.8% | +4.6个百分点 |
| AIME24数学推理 | Pass@16 | 80.0% (OPD, 1.7B) | Base 76.7% | +3.3个百分点 |
| AMOBench数学推理 | Pass@16 | 35.9% (OPD, 1.7B) | Base 12.8% | +12.8个百分点 |
| AMOBench数学推理 | Pass@16 | 61.5% (OPSD, 8B) | Base 41.0% | +20.5个百分点 |
| HMMT25数学推理 | Avg@16 | 44.5% (OPSD, 8B) | Base 41.5% | +3.0个百分点 |
| HumanEval+代码生成 | Avg@16 | 63.2% (OPD, 1.7B) | Base 62.3% | +0.9个百分点 |
| LiveCodeBench代码生成 | Avg@16 | 32.9% (OPD, 1.7B) | Base 32.3% | +0.6个百分点 |
局限与改进
作者承认的局限性包括：第一，TRD需要一个额外采样预算来构建精炼轨迹。这个开销部分被更短精炼轨迹上更快的KL训练抵消；在Qwen3-8B上，总墙钟时间几乎匹配密集KL基准。第二，TRD依赖于教师以缓解前缀失效同时保持精炼轨迹接近学生策略内分布的方式引导精炼的能力。这个限制在更强的教师下不太严重，因为更强的教师原则上可以以promise优化优化目标。基于我自己的观察，TRD的局限性还包括：首先，在代码任务如LiveCodeBench上，所有方法包括TRD都未能匹配基线模型，表明当前教师可能无法在这些较难代码任务上提供有效精炼。其次，TRD的增益在较容易的基准如AIME24和AIME25上饱和，边际改进有限。第三，TRD需要额外的教师查询，增加了计算成本，虽然部分被训练效率提升抵消，但在资源受限场景下仍有考虑。最后，TRD的有效性可能依赖于教师与学生的能力差距，如果教师能力不足，精炼质量可能受限。
独立分析的弱点
TRD的独立分析弱点包括：首先，TRD的有效性在不同任务类型上表现不一致。在数学推理任务特别是竞争数学如AMOBench上表现优异，但在代码生成任务如LiveCodeBench上增益有限甚至没有改进。改进方向可以是针对不同任务类型设计差异化的精炼策略，或者在代码任务上集成代码特定的验证器和反馈机制。第二，TRD依赖于教师模型的精炼能力，如果教师本身能力不足或与学生能力差距过大，精炼效果可能受限。改进方向可以包括设计教师质量评估机制，自适应调整精炼强度，或者使用多个教师模型集成提高精炼鲁棒性。第三，TRD的轨迹精炼过程可能引入计算开销，虽然论文声称墙钟时间几乎匹配，但在实际大规模部署中仍需考虑。改进方向可以是优化精炼效率，例如缓存常见错误模式的精炼结果，或者设计轻量级精炼策略。第四，TRD的前向KL选择可能在某些场景下不如反向KL，特别是在模式崩塌风险高的场景。改进方向可以是自适应KL方向选择，或者设计混合策略结合两者的优势。第五，TRD目前主要在单轮蒸馏中验证，多轮精炼的累积效应和策略偏移问题尚未充分探索。改进方向可以是设计渐进式精炼策略，或者引入偏移检测和校正机制。
未来方向
作者提出的未来工作方向和基于成果可延伸的探索包括：首先，可以探索更强的教师模型对TRD性能的影响。虽然论文使用Qwen3-8B作为教师，但更大的教师或不同架构的教师可能提供更好的精炼质量，特别是在代码任务上。其次，可以研究多轮精炼策略，即迭代地应用TRD，观察是否能够累积改进或产生策略偏移。第三，可以探索TRD与其他蒸馏技术的结合，例如与强化学习与可验证奖励或监督微调的混合策略，综合利用不同方法的优势。第四，可以研究TRD在不同规模模型上的缩放行为，从更小的基础模型到更大的模型，观察TRD的收益是否随着模型规模变化。第五，可以探索TRD在其他任务类型上的应用，例如代码生成、逻辑推理、常识推理等，验证TRD的通用性。第六，可以研究TRD的精炼质量度量，设计自动化评估精炼轨迹质量的方法，这对于实际部署和调试TRD系统很重要。最后，可以探索TRD在多语言和跨任务场景中的应用，观察精炼策略是否能够泛化到不同语言和任务域。
复现评估
TRD的复现评估：开源方面，论文提供代码链接https://github.com/louieworth/trd，表明代码是公开的。数据方面，数学训练使用DeepScaleR数学语料库，约4万问题和解答；代码训练使用TACO算法代码生成语料库，约2.5万训练问题，具有参考解和测试用例。评估基准包括五个数学基准AIME24、AIME25、HMMT25、BeyondAIME、AMOBench和三个代码基准HumanEval+、MBPP+、LiveCodeBench v6。算力方面，论文使用Qwen3模型家族，在OPD设置中使用Qwen3-8B作为教师，学生为Qwen3-1.7B和Qwen3-4B-Instruct-2507；在OPSD设置中，教师和学生共享相同backbone，使用Qwen3-4B-Instruct-2507和Qwen3-8B。训练细节在附录C中提供，包括响应长度设置即数学38912，代码16384、采样预算即K等于16和验证器配置。难度方面，复现需要访问模型权重、训练数据和评估基准，可能需要显著计算资源。论文报告了训练墙钟时间，在Qwen3-8B上总墙钟时间几乎匹配密集KL基准，这有助于估计计算需求。然而，某些设置如OPSD的特权条件化需要仔细实现，可能需要额外调试。总体而言，鉴于代码开源、数据集公开和详细实验设置，复现应该是可行的，但需要相当的机器学习经验和计算资源。
论文图表
图2展示了前缀失效下的机制分析。左图和中图分别展示了前向KL和反向KL在前缀失效下的分布表现，可以看到前向KL被修正起始区域主导，反向KL被错误延续区域主导。右图展示了教师在前缀失效下提供的监督信号，显示了双峰混合分布的结构。这三个子图共同说明了token级KL在前缀失效下的失效机制。
Figure 2对理解论文的理论基础至关重要，因为它直观地展示了前缀失效的结构机制，包括双峰教师混合和不同KL方向的行为差异。这个图帮助读者理解为什么token级干预方法会失效，以及为什么需要轨迹级方法。它是论文第4节Prefix Failure in Token-Level On-Policy Distillation的核心视觉说明，为TRD的方法论提供了理论支撑。
图3展示了标准OPSD在前缀失效下的三个经验观察。左图显示按正确和错误rollout分割的逐tokenKL，两者在原始rollout上都保持接近零；中图显示教师和学生困惑度差距，差距在整个训练过程中保持接近零；右图显示教师的认知起始质量，教师在6到8千分之一的每个位置质量上分配给16个认知起始token。这三个测量共同验证了前缀失效机制，并包含了TRD方法精炼轨迹的参考曲线，显示TRD缓解了这些问题。
Figure 3对理解论文的实验验证部分非常重要，因为它提供了前缀失效机制的实证证据。三个子图分别从监督退化、困惑度差距缩窄和认知起始质量增益三个角度验证了理论分析，同时展示TRD如何缓解这些问题。这个图是论文第4.3节Experimental Validation of Prefix Failure的核心，帮助读者理解前缀失效不仅是理论分析，而且在实际训练中确实发生。