Trajektorienverfeinerte Destillation

Zusammenfassung

On-Policy-Destillation (OPD) hat sich zu einem zentralen Nachbearbeitungswerkzeug für große Sprachmodelle (LLMs) entwickelt, indem sie eine dichte Token-für-Token-Lehrerüberwachung entlang der studenteneigenen Rollouts bereitstellt. In dieser Arbeit identifizieren wir eine gemeinsame strukturelle Ursache für OPD, die wir als Präfixfehler bezeichnen. Bei Präfixfehlern induziert die dichte Token-für-Token-Überwachung eine bimodale Lehrermischung und fragmentierte Gradienten, die durch Token-Level-Verlustkürzung oder -Neugewichtung nicht behoben werden können. Diese Beobachtung motiviert uns, über Token-Level-Verlustinterventionen hinaus zu Korrekturen auf Trajektorienebene zu gehen. Daher schlagen wir Trajektorien-verfeinerte Destillation (TRD) vor, eine Korrekturmethode auf Trajektorienebene, die den Rollout des Studenten unter der Lehrerführung innerhalb des On-Policy-Unterstützungsbereichs revidiert. Durch die Korrektur problematischer Präfixe vor der Destillation mildert TRD den Präfixfehler an seiner Quelle. Darüber hinaus verbessert TRD die Exploration, indem es den Studenten unter Lehrerführung alternativen gültigen Ableitungen aussetzt, selbst wenn die ursprünglichen Rollouts bereits korrekt sind. TRD kann auch auf On-Policy-Selbstdestillation (OPSD) angewendet werden, eine Variante mit gemeinsamen Parametern, die das Studentenmodell, konditioniert auf privilegierte Informationen, als Lehrer verwendet. Über eine breite Palette von Benchmarks und Basismodellen auf mehreren Skalen hinweg übertrifft TRD durchgängig frühere Baselines, verbessert die Einzelversuchsgenauigkeit und erweitert die Abdeckung des logischen Denkens. Der Code ist verfügbar unter 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