What Shared Prefixes Hide: Trajectory Dropout for On-Policy Distillation
cs.AI
Submitted: 2026-09-27
Updated: 2026-09-29
License: http://creativecommons.org/licenses/by/4.0/
The gist: On-policy distillation (OPD) trains a student model on its own trajectories using dense token-level feedback from a stronger teacher model.
Terminology
Abstract
On-policy distillation (OPD) trains a student model on its own trajectories using dense token-level feedback from a stronger teacher model. Since each update is conditioned on the reasoning prefix already generated by the student, the prefix also shapes how effectively teacher feedback is converted into learning. We find that shared prefixes can lead to weak token-level updates, a phenomenon we call Prefix-Induced Supervision Attenuation (PISA). This attenuation arises in two common cases. (i) High student confidence can weaken corrective gradients even when the teacher disagrees. (ii) Tokens that rely on earlier reasoning can receive learning signals as weak as those for simple local continuations. To solve this problem, we propose Trajectory Dropout, a simple training-time intervention that exposes these weakened signals. The student first performs a standard full-context rollout to generate a complete trajectory. During training, we randomly drop a certain proportion of the student's reasoning trajectory, while the teacher continues to observe the complete trajectory for token-level supervision. This intervention strengthens corrections for overconfident predictions and introduces additional supervision at prefix-sensitive positions. Trajectory Dropout consistently improves average performance across teacher--student model pairs of different scales and six mathematical reasoning benchmarks, while also yielding gains on two out-of-domain benchmarks. It can also be flexibly integrated into existing OPD variants with negligible computational overhead, further improving their performance. These results demonstrate that Trajectory Dropout provides a simple mechanism for strengthening token-level supervision across model scales and OPD objectives.
Sources
- On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes
- Demystifying Training-Time Augmentation for Data-Constrained Language Model Pretraining
- Trajectory-Refined Distillation
- Bridging Reasoning Trajectories in On-Policy Distillation via Near-Future Guidance
- Filter, Then Reweight: Rethinking Optimization Granularity in On-Policy Distillation
- ReNIO: Reweighting Negative Trajectory Importance for LLM On-Policy Distillation
- Same Evidence, Different Answers: Canonical-Context On-Policy Distillation for Multi-Turn Language Models
- Prefix Teach, Suffix Fade: Local Teachability Collapse in Strong-to-Weak On-Policy Distillation
- Your Teacher Can't Help You Here: Combating Supervision Fidelity Decay in On-Policy Distillation
- GPQA: A Graduate-Level Google-Proof Q&A Benchmark
- Not All Disagreement Is Learnable: Token Teachability in On-Policy Distillation
- Qwen3 Technical Report
- On-Policy Context Distillation for Language Models
- Dense Supervision, Sparse Updates: On the Sparsity and Geometry of On-Policy Distillation
- Self-Distilled Reasoner: On-Policy Self-Distillation for Large Language Models
Related papers
- MAVEN-T: Reinforced Heterogeneous Distillation for Real-Time Multi-Agent Trajectory Prediction
- Model Discovery Agent: LLM-assisted Bayesian experiment design for data-efficient discovery of mechanistic world models
- Agents' Last Exam
- The Clinician's Veto: Navigating Trust, Liability, and Uncertainty in Autonomous AI Prescribing
- MindHelper: Closed-Loop Embodied Mental-State Reasoning for Precision Intervention
- Incumbent Advantage: Brand Bias and Cognitive Manipulation Dynamics in LLM Recommendation Systems