aiqbench analysis
This research identifies "prefix failure" as a structural bottleneck in on-policy distillation where token-level interventions prove insufficient. The proposed Trajectory-Refined Distillation (TRD) addresses this by implementing trajectory-level output corrections that revise student rollouts before distillation, thereby improving exploration and reasoning coverage. The method demonstrates consistent performance gains across diverse benchmarks and scales, offering a robust upgrade to existing post-training pipelines.
Trajectory-Refined Distillation
arXiv:2606.08432v1 Announce Type: new Abstract: 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
Abstract: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 this https URL
Artificial Intelligence (cs.AI)
arXiv:2606.08432 [cs.AI]
(or arXiv:2606.08432v1 [cs.AI] for this version)
https://doi.org/10.48550/arXiv.2606.08432 Focus to learn more arXiv-issued DOI via DataCite (pending registration)
arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.
Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.
Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.
Full text under the publisher's open licence · 286 words
Related coverage
Frontier post-training recipe review with Finbarr Timbers
Interconnects · Jun 16, 2026
Reflect-VLA: Internalizing Reflective Guidance via Dual-Path On-Policy Self-Distillation for Autonomous Driving
arXiv · Oct 12, 2026
MeSD: Multi-Evidence Self-Distillation for VideoLLM
arXiv · Oct 7, 2026
Distill Locally, Schedule Globally: Flow Maps for Few-Step Text-to-Speech
arXiv · Sep 30, 2026