Destilación Refinada por Trayectorias

Resumen

La destilación en política (OPD, por sus siglas en inglés) se ha convertido en una herramienta central para el entrenamiento posterior de modelos de lenguaje de gran escala (LLMs), proporcionando una supervisión densa por token por parte del profesor a lo largo de las propias trayectorias generadas por el estudiante. En este trabajo, identificamos una causa estructural común subyacente a la OPD, a la que denominamos fallo de prefijo. Bajo el fallo de prefijo, la supervisión densa por token induce una mezcla bimodal del profesor y gradientes fragmentados que la truncación o reponderación de pérdidas a nivel de token no logran abordar. Esta observación nos motiva a ir más allá de las intervenciones sobre la pérdida a nivel de token, hacia correcciones de salida a nivel de trayectoria. Por lo tanto, proponemos la Destilación Refinada por Trayectoria (TRD, por sus siglas en inglés), un método de corrección a nivel de trayectoria que revisa la trayectoria generada por el estudiante bajo la guía del profesor, manteniéndose dentro del soporte en política. Al corregir los prefijos problemáticos antes de la destilación, TRD mitiga el fallo de prefijo en su origen. Además, TRD mejora la exploración al exponer al estudiante a derivaciones válidas alternativas bajo la guía del profesor, incluso cuando las trayectorias originales ya son correctas. TRD también se puede aplicar a la autodestilación en política (OPSD), una variante de parámetros compartidos que utiliza el modelo estudiante condicionado a información privilegiada como profesor. En una amplia gama de puntos de referencia y modelos base a múltiples escalas, TRD supera consistentemente a las líneas de base anteriores, mejorando la precisión en intentos únicos y ampliando la cobertura de razonamiento. El código está disponible en 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