Destilações com Trajetória Refinada

Resumo

A destilação on-policy (OPD) tornou-se uma ferramenta central de pós-treinamento para modelos de linguagem de grande escala (LLMs), fornecendo supervisão densa do professor por token ao longo das próprias trajetórias geradas pelo estudante. Neste trabalho, identificamos uma causa estrutural comum subjacente à OPD, que denominamos falha de prefixo. Sob a falha de prefixo, a supervisão densa por token induz uma mistura bimodal do professor e gradientes fragmentados que a truncagem ou reponderação de perdas no nível do token não conseguem resolver. Essa observação nos motiva a avançar além das intervenções de perda no nível do token em direção a correções de saída no nível da trajetória. Propomos, assim, a Destilação Refinada por Trajetória (TRD), um método de correção no nível da trajetória que revisa a rolagem do estudante sob a orientação do professor, mantendo-se dentro do suporte on-policy. Ao corrigir prefixos problemáticos antes da destilação, a TRD mitiga a falha de prefixo em sua origem. Além disso, a TRD melhora a exploração ao expor o estudante a derivações alternativas válidas sob a orientação do professor, mesmo quando as rolagens originais já estão corretas. A TRD também pode ser aplicada à autodestilação on-policy (OPSD), uma variante com compartilhamento de parâmetros que utiliza o modelo estudante condicionado a informações privilegiadas como professor. Em uma ampla gama de referências e modelos base em múltiplas escalas, a TRD supera consistentemente as linhas de base anteriores, melhorando a precisão de tentativa única e ampliando a cobertura de raciocínio. O código está disponível em 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