Траекторно-уточненная дистилляция

Аннотация

Дистилляция на политике (OPD) стала основным инструментом пост-тренировки для больших языковых моделей (LLM), обеспечивая плотное пословное (по-токеновое) обучение от учителя на собственных траекториях студента. В данной работе мы выявляем общую структурную причину, лежащую в основе OPD, которую мы называем отказом префикса. При отказе префикса плотное пословное обучение порождает бимодальную смесь учителя и фрагментированные градиенты, которые не удается устранить с помощью усечения или перевзвешивания потерь на уровне токенов. Это наблюдение побуждает нас выйти за рамки вмешательств в функцию потерь на уровне токенов и перейти к коррекции выходов на уровне траекторий. Таким образом, мы предлагаем дистилляцию с уточнением траекторий (TRD) — метод коррекции на уровне траекторий, который исправляет развертку студента под руководством учителя, оставаясь в рамках поддержки на политике. Исправляя проблемные префиксы до дистилляции, TRD устраняет отказ префикса в его источнике. Кроме того, TRD улучшает исследование, предоставляя студенту альтернативные корректные выводы под руководством учителя, даже когда исходные развертки уже верны. TRD также может применяться к само-дистилляции на политике (OPSD) — варианту с разделением параметров, где в качестве учителя используется модель студента, дополненная привилегированной информацией. На широком спектре бенчмарков и базовых моделей различных масштабов TRD последовательно превосходит предыдущие базовые подходы, повышая точность за одну попытку и расширяя покрытие рассуждений. Код доступен по адресу 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