Distillation raffinée par trajectoire

Résumé

La distillation sur politique (on-policy distillation, OPD) est devenue un outil central de post-entraînement pour les grands modèles de langage (LLMs), offrant une supervision dense par token de l’enseignant le long des propres rollouts de l’étudiant. Dans ce travail, nous identifions une cause structurelle commune sous-jacente à l’OPD, que nous appelons défaillance du préfixe. En cas de défaillance du préfixe, la supervision dense par token induit un mélange bimodal de l’enseignant et des gradients fragmentés que la troncature ou la repondération de perte au niveau des tokens ne parviennent pas à résoudre. Cette observation nous incite à dépasser les interventions sur la perte au niveau des tokens pour aller vers des corrections de sortie au niveau de la trajectoire. Nous proposons ainsi la distillation affinée par trajectoire (Trajectory-Refined Distillation, TRD), une méthode de correction au niveau de la trajectoire qui révise le rollout de l’étudiant sous la guidance de l’enseignant tout en restant dans le support on-policy. En corrigeant les préfixes problématiques avant la distillation, TRD atténue la défaillance du préfixe à sa source. De plus, TRD améliore l’exploration en exposant l’étudiant à des dérivations valides alternatives sous la guidance de l’enseignant, même lorsque les rollouts originaux sont déjà corrects. TRD peut également être appliqué à l’auto-distillation sur politique (on-policy self-distillation, OPSD), une variante avec partage des paramètres qui utilise le modèle étudiant conditionné sur des informations privilégiées comme enseignant. À travers un large éventail de benchmarks et de modèles de base à plusieurs échelles, TRD surpasse systématiquement les références antérieures, améliorant la précision en un seul essai et élargissant la couverture de raisonnement. Le code est disponible à l’adresse https://github.com/louieworth/trd.

English

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