Pass the Baton: When AI Teachers Learn When to Step In — Relay-OPD Explained
*Full English rendering of a Chinese forum post analyzing the paper "Pass the Baton: Trajectory-Relayed On-Policy Distillation" (arXiv:2607.26057). Metaphors and explanatory style preserved.*
> *"Education is not the filling of a pail, but the lighting of a fire. But when the fire strays off course, you need a hand extended at just the right moment."* — adapted from W.B. Yeats
The relay race that shouldn't be lost
Imagine a 4×100m relay. Runner three fumbles the baton handoff — it's crooked — but instead of stopping to regrip, he sprints on. Every later leg is built on that unstable foundation, and the team loses. Not for lack of speed, but because an early error was never corrected, so all subsequent effort pointed the wrong way.
This is a metaphor for Prefix Failure in knowledge distillation. When a small student model learns from a large teacher via on-policy distillation (OPD), the student generates its own reasoning trajectories. Once it goes off track mid-solution, every subsequent token is conditioned on the flawed prefix — producing "misleading continuations" that waste compute and yield unreliable supervision. The paper describes this as teacher-student continuation asymmetry on failed prefixes: on a bad prefix, the teacher's next-token distribution tends to *correct*, while the student's tends to *continue* the error, driving KL divergence sharply upward.
Background: from distillation to OPD
Knowledge distillation (Hinton et al., 2015) is an apprenticeship: a large teacher (e.g., GPT-4, DeepSeek-R1) transfers its skill to a small student (0.6B–1.7B) not by copying parameters, but by letting the student observe *how the teacher thinks*.
On-policy distillation (OPD) goes further: instead of imitating the teacher's finished answers, the student writes its own attempt and receives per-token feedback — like a teacher grading the student's own essay rather than having them trace a model answer.
But OPD inherits the crooked-baton problem: once the student errs early, the teacher must supervise a chain of reasoning built on a mistake. Prior fixes — truncation, weighted losses on post-error tokens, and early stopping (FastOPD) — are band-aids; none fixes the root cause.
Relay-OPD: Detect, Relay, Resume
Core idea: rather than preventing errors, let the teacher briefly take the baton when an error occurs, run a correct segment, then hand it back.
1. Detection — a free signal
The teacher-student continuation asymmetry is itself the trigger: the KL divergence between teacher and student next-token distributions is already computed as the OPD loss, and it is *recycled* as a failure signal. When it exceeds a threshold, the student has likely strayed. No manual labels, no extra compute, and it is naturally adaptive (easy tasks → few interventions; hard tasks → more).2. Relay — the teacher's correction leg
At the trigger point, the student's prefix is frozen (errors included), the teacher continues generating from it for a limited number of tokens — the teacher leg — producing a corrected continuation.3. Resume — handing back
After the leg, the student resumes generation from the corrected position, now facing a sound foundation instead of a misleading one. The relay budget limits how long the teacher runs, preventing the student from collapsing into a teacher clone and concentrating help at early positions, where most prefix failures occur.The math
Standard OPD minimizes, over prefix tokens \(y_{<t}\):Relay-OPD instead builds a relay trajectory:
and minimizes the KL on this trajectory, preserving OPD's on-policy nature while repairing prefix failure.
Budget Thermostat
A Budget Thermostat adaptively adjusts the trigger threshold over training: lower early (teacher intervenes often), raised later (student gains independence) — maintaining a target average relay frequency. Like a bike-riding coach: hands on the seat in week one, a loose grip in week three, watching from a bench by month two.Experiments
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.
Results:
- On the 1.7B student: best or second-best on all 8 math benchmarks; +5.73% average over standard OPD; +1.49% over the strongest baseline (FastOPD).
- Training trajectory length reduced by over 50% — faster and cheaper training.
- On AMOBench, problems unsolvable by baselines reached Pass@16 of 0.39 — Relay-OPD *expands the range of solvable problems*, not just accuracy.
- Scaffolding (Jerome Bruner): temporary support that is gradually withdrawn — exactly what the Budget Thermostat implements algorithmically.
- Control theory: standard OPD is almost purely feedforward, so small errors accumulate. Relay-OPD inserts a feedback loop that breaks the error-accumulation chain — like an autonomous car combining map navigation with live lidar.
- Art analogy: the teacher doesn't repaint the canvas for the student; at the spot where perspective is wrong, the teacher paints a few corrective strokes and lets the student continue.
- 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.
Ablations: always-relay degrades performance (smart triggering is necessary); adaptive relay length beats fixed length; removing the thermostat causes instability and late-training regression; forward KL works better than reverse KL within the relay framework.
Why it works: scaffolding and feedback control
Implications and outlook
Key takeaways: (1) errors themselves aren't the problem — continuation built on them is; (2) teacher/student roles should be dynamic, not fixed; (3) *timing* of intervention matters more than its intensity.
Potential applications beyond math: code generation (preventing cascading compile errors), long-form text (catching contradictory character setups early), multi-step decision-making (a takeover mechanism at critical decision points), and scientific hypothesis vetting.
Limitations: binary relay triggers (no partial takeover), a fixed teacher, correction at only one point per trajectory, and reliance on KL divergence as the difficulty signal. Future directions include gradient- or uncertainty-based triggers, multi-teacher committees, and multi-relay trajectories.
Closing thought
In the ancient Olympic torch relay, the flame must never break — but it also mustn't be passed blindly. Relay-OPD shows that AI training can be an elegant relay: not the teacher always running ahead, nor the student abandoned on the track, but an experienced hand steadying the baton at the crucial instant — then letting go.