轨迹精炼蒸馏(TRD)
阅读原文· arxiv.org在同策略蒸馏(OPD)中,密集的逐token教师监督会导致前缀失败(双峰教师混合与碎片化梯度),逐token损失截断或重加权无法解决。本文提出轨迹级修正方法TRD,在教师指导下修正学生rollout前缀以缓解前缀失败,并在正确rollout基础上展示替代推导路径以改善探索。TRD可推广至同策略自蒸馏(OPSD)。在多项基准和多尺度模型上,TRD一致优于先前基线,提升单次准确率并拓宽推理覆盖范围。
On-policy distillation (OPD) has become a central post-training tool for large language models (LLMs), providing dense per-token teacher supervision along the student's own rollouts. In this work, we identify a common structural cause underlying OPD, which we call prefix failure. Under prefix failure, dense per-token supervision induces a bimodal teacher mixture and fragmented gradients that token-level loss truncation or reweighting fail to address. This observation motivates us to move beyond token-level loss interventions toward trajectory-level output corrections. We thus propose Trajectory-Refined Distillation (TRD), a trajectory-level correction method that revises the student's rollout under the teacher guidance while within on-policy support. By correcting problematic prefixes before distillation, TRD mitigates prefix failure at its source. Moreover, TRD improves the exploration by exposing the student to alternative valid derivations under teacher guidance, even when the original rolls are already correct. TRD can also be applied to on-policy self-distillation (OPSD), a parameter-sharing variant that uses the student model conditioned on privileged informations as the teacher. Across a wide range of benchmarks and base models at multiple scales, TRD consistently outperforms prior baselines, improving single-attempt accuracy and broadening reasoning coverage. Code is available at https://github.com/louieworth/trd