Trajectverfijnde destillatie

Samenvatting

On-policy distillatie (OPD) is een centraal post-traininginstrument geworden voor grote taalmodellen (LLM's), dat dichte per-token begeleiding door de leraar biedt tijdens de eigen rollouts van de student. In dit werk identificeren we een veelvoorkomende structurele oorzaak die ten grondslag ligt aan OPD, die we prefixfalen noemen. Bij prefixfalen induceert dichte per-token begeleiding een bimodale mix van leraren en gefragmenteerde gradiënten die token-niveau verliesafkapping of herweging niet kan aanpakken. Deze observatie motiveert ons om verder te gaan dan interventies op token-niveau en over te stappen naar correcties op trajectniveau. Daarom stellen we Traject-Verfijnde Distillatie (TRD) voor, een correctiemethode op trajectniveau die de rollout van de student herziet onder begeleiding van de leraar, terwijl deze binnen de on-policy ondersteuning blijft. Door problematische prefixen te corrigeren vóór distillatie, vermindert TRD prefixfalen bij de bron. Bovendien verbetert TRD de exploratie door de student bloot te stellen aan alternatieve geldige afleidingen onder begeleiding van de leraar, zelfs wanneer de oorspronkelijke rollouts al correct zijn. TRD kan ook worden toegepast op on-policy zelfdistillatie (OPSD), een parameter-delende variant die het studentmodel, geconditioneerd op bevoorrechte informatie, als leraar gebruikt. Over een breed scala aan benchmarks en basismodellen op meerdere schalen presteert TRD consequent beter dan eerdere basislijnen, verbetert het de nauwkeurigheid bij een enkele poging en verbreedt het de redeneerdekking. Code is beschikbaar op 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