On-policy distillation (OPD) provides teacher supervision on student-generated trajectories, but a single trajectory realizes only one continuation at each decision, limiting supervision coverage under a finite sampling budget. We propose BranchOPD (Branching On-Policy Distillation), which selects informative decision points using student uncertainty and teacher-student disagreement, resamples student-policy continuations from shared prefixes, and applies token-level teacher supervision to both main trajectories and branches. BranchOPD operates at token-level decisions in mathematics and within-turn decisions in agent tasks. Shared-prefix reuse reduces student-generated token costs by approximately 19% on ALFWorld and 31% on WebShop relative to seven independent trajectories. We also introduce protocol-boundary repair (PBR), which uses targeted supervision under protocol constraints to mitigate repetition and wasted generation. BranchOPD outperforms Vanilla OPD and Multi-Rollout OPD (n=7) across all evaluated benchmarks, with gains over Vanilla OPD of 18.55 percentage points in mean mathematics accuracy over 16 sampled answers (avg@16) and 13.8 points in mean ALFWorld ID/OOD success. Mechanistic analyses and selection ablations show that branching at uncertain decisions expands subsequent behavior coverage and outperforms random selection and turn-initial branching.
Figures & tables
Figure 1: Resampling from one prefix on a mathematics problem and a logged ALFWorld trajectory. Scores mark the fork; q is the student probability of each continuation’s first token.
Figure 2: Overview of BranchOPD. Fork points are selected on the main trajectory, suffixes are resampled from the student, and the teacher scores every generated token; shared prefixes are counted once. The bottom examples show concrete PBR constraints used in Section 5.5 .
Figure 3: ALFWorld generation cost and success with PBR, for three (left) and seven (right) trajectory slots. OPD, Independent, and Reweighting denote Vanilla OPD, Multi-Rollout OPD, and Entropy-Weighted OPD.
Figure 4: Training dynamics with PBR: mathematics response length and truncation rate (fraction of responses stopped at the length limit), and ALFWorld turn-level protocol violation (fraction of turns with a violation). Repetition penalties delay budget exhaustion but leave generation near the cap; PBR reduces both truncation and violations.
Appendix figures & tables10 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 5: Trajectory-level protocol violation (PVP): the fraction of the 192 main and branch trajectories that contain a violation. All methods shown use PBR; BranchOPD-F denotes forced-token expansion (Table 8 ).
Figure 6: Probability mass available for alternative decisions. (a) “Mass kept by the main path” is the probability of the sampled main token, πθold(a0∣s) . (b) “Mass discarded” and MK denote the alternative-token mass 1−πθold(a0∣s) , i.e., the mass left unrealized by the main draw. (c) All supervised positions compared with selected forks. ALFWorld and WebShop use training dumps; mathematics and Search-QA use reconstructed summaries. In the plot, q denotes πθold and turn-open-RS denotes Turn-initial branching.
Figure 7: Empirical CDF of alternative mass over supervised positions (ALFWorld step 200, WebShop step 210), accumulated over 20 equal-width bins. BranchOPD-F denotes forced-token expansion (Table 8 ), and “uniform” denotes uniform weights on its branches.
Figure 8: Behavior coverage of resampled branches over training. “New state” in the plot denotes new-action coverage (NSR). High disagreement and NSR mean that branches reach behavior not present in the main trajectory.
Figure 9: Rescue and harm rates during training, each conditioned on the outcome of the main trajectory. “Mean net” is R↑−R↓ averaged over training steps. “Spoil” means that all branches fail although the main trajectory succeeds; the main trajectory itself stays in the training batch.
Figure 10: Fork location and distance to the first changed action. (a) Fork position relative to the full trajectory and (b) relative to the selected turn, where 0 and 1 mark the start and end. (c) Number of generated student tokens from the fork to the first changed action.
On-policy distillation (OPD) combines student-generated rollouts with dense token-level supervision from a teacher, but providing such supervision for every rollout requires substantial teacher computation. We introduce Success-Referenced On-Policy Distillation (SR-OPD), which reduces this cost by selecting which prompts and rollouts receive teacher supervision. When the student produces both successful and failed rollouts for the same prompt, a successful rollout can serve as a natural reference for selecting failed rollouts. SR-OPD therefore focuses on such prompts and prioritizes failed rollouts whose hidden-state trajectories show sustained divergence from a successful reference, while accounting for estimated teacher-input cost. Across three teacher-student pairs and six mathematical reasoning benchmarks, SR-OPD uses only 3.46-5.02% of the teacher-input tokens required by Vanilla OPD in the one-pass setting while maintaining comparable reasoning performance. Under a controlled setting matched to 5% of Vanilla OPD's teacher-input budget, further experiments support both key design choices: focusing supervision on prompts with both successful and failed rollouts, and using successful rollouts to guide failure selection. These results indicate that a student's own successful behavior can serve as a useful reference for allocating teacher supervision under a fixed teacher-input budget.
Xiang Chen, Futao Su, Kong Wang +2
Tongji University · Xiaohongshu · Shanghai Jiao Tong University +1
On-policy distillation (OPD) grounds token-level supervision in the student's own trajectory, yet suffers from prefix failure: once the student commits to a wrong reasoning direction, all subsequent generation builds on this deviation, producing misdirected continuations that elicit unreliable supervision and waste compute. We identify a teacher-student continuation asymmetry on failed prefixes, where the teacher tends to redirect while the student continues along the original direction, and convert it into a label-free handoff trigger in Relay On-Policy Distillation (Relay-OPD). During training, Relay-OPD constructs relay trajectories by letting the teacher briefly take over at detected trigger points to produce a teacher leg, after which the student resumes and is optimized on the resulting trajectory. A limited relay budget concentrates intervention on critical early positions while limiting departure from the student policy. With a Qwen3-4B-Instruct-2507 teacher and Qwen3-0.6B/1.7B-Non-Thinking students on eight mathematical reasoning benchmarks, Relay-OPD achieves the best or second-best results on every benchmark, outperforming standard OPD by +5.73% and the strongest baseline FastOPD by +1.49% on average for 1.7B, with consistent gains at 0.6B. Training trajectory length is reduced by over 50%.
Haolei Xu, Xiaowen Xu, Haiwen Hong +5
1Zhejiang University · 2Yuvion Team, Alibaba Group
On-policy self-distillation (OPSD) provides dense teacher targets, but evaluates them only along student-sampled rollouts. When the privileged teacher favors an alternative action at a visited prefix, OPSD can provide a target for the branch decision but cannot supervise the successor contexts induced by that action unless the student samples it. This creates a training-time data-collection bottleneck and suggests a different role for teacher-student disagreement: proposing a trajectory branch rather than identifying a sufficient local repair. Our diagnostic framework using controlled token interventions reveals that a teacher-preferred token at peak disagreement can improve student continuation success, while its local corrective value is limited. Motivated by this finding, we introduce a simple branch-regenerate-distill algorithm, Trajectory-Intervention Self-Distillation (TISD). TISD forces a teacher-selected branch action, returns suffix generation to the student, and distills the full trajectory under the privileged-context-conditioned teacher. Across the coding models, TISD improves average Avg@4 over SDPO by 1.2 percentage points. Across the science domains, it improves average Avg@128 by 0.8 points under an equal-step budget and by 0.3 points under an equal-time budget. These results support teacher-guided branching as a way to expose useful successor contexts for self-distillation.