Pass the Baton: Trajectory-Relayed On-Policy Distillation Explained

Forum topic · 小凯 · 2026-07-29

Summary

This post is a detailed Chinese-language walkthrough of the paper "Pass the Baton: Trajectory-Relayed On-Policy Distillation" (Relay-OPD). It explains the prefix-failure problem in on-policy knowledge distillation: when a small student model makes an error early in its reasoning trajectory, all subsequent generation is built on a wrong prefix, forcing the teacher to supervise misleading continuations. Relay-OPD solves this with a relay mechanism: teacher-student continuation asymmetry (measured by KL divergence, already computed as the training loss) triggers intervention; the teacher takes over generation for a limited correction segment; then the student resumes from the corrected point. A Budget Thermostat dynamically adjusts the trigger threshold across training, so teacher intervention fades as the student improves—an algorithmic version of educational scaffolding. On math reasoning benchmarks (AIME 2024/2025, HMMT 2025, BeyondAIME, AMOBench) with Qwen3-4B-Instruct as teacher and Qwen3-0.6B/1.7B as students, Relay-OPD improves average accuracy by 5.73% over standard OPD and 1.49% over FastOPD on the 1.7B student, while cutting training trajectory length by more than 50%. Ablations confirm the value of smart triggering, adaptive relay length, and the budget thermostat.

This forum post is an in-depth Chinese explainer of "Pass the Baton: Trajectory-Relayed On-Policy Distillation" (arXiv:2607.26057), written in a Feynman-style narrative built around a relay-race metaphor: a runner fumbles the baton handoff but keeps sprinting with a crooked grip, and every later leg of the race is built on that early mistake. The paper addresses exactly this failure mode in knowledge distillation.

Key points

Background: Knowledge distillation (Hinton et al., 2015) transfers capabilities from a large teacher model to a small student model. On-Policy Distillation (OPD) improves on offline imitation by having the student generate its own reasoning trajectories and receive teacher supervision on them—like a teacher grading a student's own essay rather than having the student copy a model answer.
The problem — Prefix Failure: If the student errs early (e.g., adding fractions by summing denominators), every subsequent token continues from that wrong prefix. The teacher is then forced to supervise "misleading continuations." The authors identify a teacher–student continuation asymmetry on failed prefixes: the teacher tends to *correct* direction, while the student tends to *persist* along the wrong path, driving their distributions apart (high KL divergence).

The Relay-OPD mechanism

Relay-OPD runs as a three-step relay: Detect → Relay → Resume.

1. Detection (free signal): The KL divergence between teacher and student next-token distributions—already computed as the OPD loss—doubles as an error trigger. No extra labeling or computation is needed; the signal is naturally adaptive (easier tasks trigger less). 2. Relay (teacher leg): When KL divergence exceeds a threshold, the student's prefix is frozen, and the teacher generates a limited correction segment from the erroneous prefix. 3. Resume (student leg): The student continues generating from the teacher's corrected point. The relay length is capped by a relay budget, concentrating scarce teacher tokens at early, error-prone positions.

The relay trajectory is:

\tilde{y} = [y_{student}[1:t_{trigger}], \ y_{teacher}[t_{trigger}+1:t_{trigger}+L], \ y_{student}[t_{trigger}+L+1:T]]

and training minimizes

L_{Relay\text{-}OPD} = \sum_{t} D_{KL}(\pi_{\theta}(\cdot | x, \tilde{y}_{<t}) \| \pi_{teacher}(\cdot | x, \tilde{y}_{<t}))

preserving OPD's on-policy character.

Budget Thermostat: Too little teacher help leaves errors uncorrected; too much makes the student a teacher clone. The thermostat monitors average relay frequency and dynamically adjusts the trigger threshold—intervening often early in training and progressively withdrawing support, mirroring Jerome Bruner's scaffolding pedagogy.

Experimental results

Setup: Teacher = Qwen3-4B-Instruct; students = Qwen3-0.6B and Qwen3-1.7B. Benchmarks: AIME 2024/2025, HMMT 2025, BeyondAIME, AMOBench, plus HumanEval+/MBPP+/LiveCodeBench.
1.7B student: best or second-best on all 8 math benchmarks; +5.73% average over standard OPD, +1.49% over FastOPD.
0.6B student: consistent improvements.
Efficiency: training trajectory length reduced by more than 50%.
Capability expansion: on AMOBench, Pass@16 rises from 0 to 0.39 on problems baselines cannot solve—Relay-OPD extends the range of solvable problems, not just accuracy.

Ablations: always-on relaying hurts (over-reliance); adaptive relay length beats fixed length; removing the budget thermostat destabilizes training; forward KL works better than reverse KL within the relay framework.

Interpretations and outlook

The post frames Relay-OPD through three lenses: educational scaffolding (progressively removing support), control theory (inserting a feedback loop to break feed-forward error accumulation), and artistic practice (a teacher correcting a few strokes of a student's draft rather than repainting it).

Suggested applications beyond math: code generation (stopping cascading syntax errors), long-form text (fixing early consistency violations), multi-step decision systems (human-takeover logic at critical points), and scientific hypothesis generation.

Limitations noted: binary (all-or-nothing) triggering, a fixed teacher, single correction point per trajectory, and reliance on KL divergence as the difficulty signal. Future directions include gradual/partial takeover, teacher ensembles or evolving teachers, multi-relay trajectories, and alternative triggers (gradient-, uncertainty-, or task-based).

References cited in the post

Xu, H., Xu, X., & Hong, H. (2026). Pass the Baton: Trajectory-Relayed On-Policy Distillation. arXiv:2607.26057.
Hinton, G., Vinyals, O., & Dean, J. (2015). Distilling the knowledge in a neural network. arXiv:1503.02531.
Gu, Y., et al. (2024). On-policy distillation for language models. arXiv:2404.08527.
Song, M., & Zheng, M. (2026). A survey of on-policy distillation for large language models. arXiv:2604.00626.
Yang, A., et al. (2025). Qwen3 technical report. arXiv:2505.09388.
Bruner, J. S. (1978). The role of dialogue in language acquisition.