Relay On-Policy Distillation: Passing the Baton When Students Go Astray

Paper: 2607.26057 Authors: Haolei Xu, Xiaowen Xu, Haiwen Hong, Zixuan Ni, Hongxing Li, Yiwen Qiu, Weiming Lu, Yongliang Shen Categories: cs.CL, cs.AI

The Gap

On-policy distillation (OPD) is the go-to method for transferring reasoning ability from a large teacher to a small student: let the student generate its own trajectory, then supervise it with teacher logits at every token. The “on-policy” part matters — the student learns from its own explorations, not the teacher’s. But this very property creates a nasty failure mode the field has largely papered over.

Prefix failure. Once the student commits to a wrong reasoning direction in the first few tokens, every subsequent token is generated conditioned on that mistake. The teacher’s supervision is now anchored to a trajectory the teacher itself would never have produced. The result: misdirected continuations, unreliable gradients, and a lot of wasted compute. FastOPD (Hong et al., 2025) tried to fix this by skipping teacher logit computation for “easy” tokens, but it treats the symptom (compute waste) without addressing the root cause (wrong prefixes poisoning the entire trajectory).

The authors identify a subtle but exploitable pattern: teacher-student continuation asymmetry. When facing a failed prefix, the teacher tends to redirect — it would produce tokens that pivot away from the mistake. The student, by contrast, just keeps barreling along the wrong path. This asymmetry is detectable without extra labels, and that’s the key insight.

On-policy distillation
         |
         v
Student commits to wrong prefix
         |
         v
All subsequent tokens conditioned on mistake
         |
         v
Teacher would redirect, student continues wrong
(continuation asymmetry)
         |
         v
Detect asymmetry as trigger signal
(no additional labels needed)
         |
         v
Relay-OPD: teacher takes over at trigger point
produces "teacher leg" trajectory
         |
         v
Student resumes on corrected path
optimized on relay trajectory
         |
         v
+5.73% over OPD, +1.49% over FastOPD
50%+ reduction in training trajectory length

The Increment

One sentence: Before this paper, on-policy distillation blindly let the student generate entire trajectories even when early mistakes made the rest worthless; after this paper, the teacher can step in at critical moments to reroute the trajectory, turning wasted compute into useful supervision.

Core Mechanism

Relay-OPD has three moving parts: a trigger detector, a relay mechanism, and a relay budget. The training loop starts the same way as standard OPD — the student generates tokens autoregressively. But at each token position, a lightweight check compares the student’s next-token distribution against what the teacher would do. When the divergence crosses a threshold (the asymmetry signal), a trigger fires.

At the trigger point, the teacher takes over and generates a short “relay leg” — a continuation that reflects the teacher’s preferred reasoning direction. This leg is typically just a few tokens. Once the relay leg is complete, the baton is passed back to the student, which now resumes generating from this corrected vantage point. The full trajectory (student prefix + teacher leg + student continuation) becomes the supervision signal. Crucially, the student is optimized on the entire relay trajectory, learning from the teacher’s redirection as if it were its own.

The relay budget is the final piece. Not every trigger leads to an intervention. The budget caps the total number of relay tokens per training example, forcing the system to concentrate its interventions on the most critical early positions — exactly where prefix failure does the most damage. This keeps the student’s policy from drifting too far from its own distribution (you still want on-policy learning) while fixing the worst mistakes.

Student generates tokens
         |
         v
[trigger detector]
   |            |
   v            v
no trigger     trigger fired
   |            |
   v            v
continue       teacher generates relay leg
student        (short continuation)
generation        |
                  v
               student resumes from
               corrected context
                  |
                  v
               optimize student on
               full relay trajectory
                  |
                  v
               decrement relay budget

Here’s the structural metaphor that makes this click. Think of a relay race with a coach on the sidelines. The student is the first runner on the track. Most of the time, the runner does fine — keep running. But the coach (teacher) is watching for moments when the runner starts veering toward the wrong lane. When that happens, the coach doesn’t yell instructions (token-level KL doesn’t help when the prefix is already poisoned). Instead, the coach calls a baton pass — an experienced anchor runner (the teacher) takes over for a leg, gets the baton back on course, then hands it back to the original runner. The relay budget is the rule that says “you only get three substitutions per race” — you can’t just have the anchor runner do the whole thing, because the first runner still needs to learn to run. The relay leg is short because the anchor runner only needs to correct the trajectory, not finish the race. And the whole resulting path — including the anchor’s leg — becomes the training example for the first runner to study afterward. The genius is that the coach doesn’t need a separate signal to know when to call the pass: the asymmetry between what the anchor runner would do and what the first runner is doing *is the signal.

Key Concepts

• Prefix failure: Imagine you’re writing an essay and your first sentence sets up the wrong frame — say, you start arguing the opposite side. Every paragraph after that is technically well-written but arguing the wrong thing. In OPD, the student model generates tokens left-to-right. If the first few tokens commit to a flawed reasoning path (say, applying the wrong formula), every subsequent token is generated conditioned on that mistake. The teacher’s supervision at position 100 is trying to correct the student’s output, but the student’s context at position 100 already contains 99 tokens of wrong reasoning. The supervision becomes unreliable because the teacher is being asked to score a trajectory it would never have produced. The fix isn’t better supervision — it’s preventing the trajectory from going off the rails in the first place.

• Teacher-student continuation asymmetry: This is the paper’s key empirical observation. Take a context where the student has just gone wrong. Now ask: what would the teacher generate next? And what would the student generate next? The teacher tends to produce tokens that redirect — acknowledging the mistake and pivoting. The student tends to produce tokens that continue the mistake. This divergence is measurable and reliable. It’s not about which model is “smarter” — it’s about the fact that the teacher has seen enough correct reasoning paths that it knows when to course-correct, while the student hasn’t. The paper converts this observation into a trigger: when the asymmetry is large, intervene.

• Relay budget: A hard cap on how many tokens the teacher is allowed to generate per training example. This is the constraint that keeps Relay-OPD actually on-policy. Without a budget, you’d just have the teacher generate everything — that’s standard off-policy distillation, and it loses the benefits of on-policy learning. The budget forces the system to be strategic: spend your intervention tokens where they matter most (early positions where prefix failure occurs) and let the student handle the rest. The paper shows that even a small budget (concentrated on critical early positions) is enough to get most of the gains.

Framework Shift

Before (standard OPD):              After (Relay-OPD):

Student: [A B C D E F G ...]        Student: [A B]
            |                              |
            v                              v
      all tokens generated          C is wrong -> trigger
      by student                         |
            |                            v
            v                       Teacher: [C'] (relay leg)
      teacher supervises                   |
      every token                          v
      (including wasted ones        Student: [D E F G ...]
       on wrong prefix)                   |
            |                              v
            v                       teacher supervises
      full-length trajectories       relay trajectory
      lots of wasted compute         shorter, more useful

From blind generation to targeted intervention, the core shift is treating on-policy distillation as a collaborative relay rather than a solo performance.

Expert Assessment

Problem choice: This is a real gap. Prefix failure has been acknowledged informally in the distillation community but nobody attacked it head-on with a clean mechanism. The observation that teacher and student diverge on failed prefixes is genuinely useful and probably underlies failure modes in other on-policy methods too. It sits at the intersection of efficient training and knowledge transfer — both hot areas — so the timing is good.

Method maturity: Clever insight, not brute force. The relay mechanism is simple enough to implement (trigger detection + teacher takeover + budget) and doesn’t require architectural changes or extra learned components. The trigger based on distributional asymmetry is elegant — it repurposes information that’s already being computed. One concern: the threshold for triggering might be sensitive and the paper could be more forthcoming about how much tuning it requires. There may be simpler heuristics (e.g., entropy spikes) that achieve 80% of the benefit with 20% of the complexity, but the authors don’t explore this.

Experimental integrity: Solid but not exhaustive. Eight mathematical reasoning benchmarks with two student sizes (0.6B and 1.7B) using Qwen3 models — this is a fair and representative setup for the claimed contribution. The gains are consistent across benchmarks and model sizes, which is reassuring. The +5.73% over OPD and +1.49% over FastOPD are meaningful. The 50% trajectory length reduction is a strong efficiency story. Red flags: only mathematical reasoning is tested — no general language tasks, no code generation, no commonsense reasoning. The teacher (4B) is not that much larger than the student (1.7B), which raises the question of whether the method helps more with larger teacher-student gaps. Ablations on relay budget and trigger threshold are present but could be deeper.

Writing quality: Generally clean and well-structured. The title “Pass the Baton” is memorable and accurate. The core contribution is communicated clearly within the first two pages. Where they cut corners: the related work section is thin — they cite FastOPD and standard OPD but don’t engage with the broader landscape of distillation interventions (e.g., speculative decoding parallels, curriculum learning). The discussion section is brief and doesn’t address failure modes or when relay intervention might hurt. A stronger discussion of *why the teacher’s relay leg produces better supervision (beyond “it redirects”) would elevate the theoretical contribution.

Verdict: weak accept — Clean method on a real problem with solid but not exhaustive experiments; the core insight is transferable and the efficiency gains are compelling, but the narrow evaluation domain and thin theoretical analysis keep it from a strong accept.

Takeaways

Three things worth stealing:

1. The asymmetry-as-signal pattern: You don’t need extra labels or a separate classifier to detect when distillation is going wrong. Just compare what the teacher and student would generate from the same context — large divergence means trouble. This is applicable to any teacher-student setup where both models are available during training.

2. Intervention budgeting: If you’re doing any kind of corrective intervention during training (not just distillation — could be reward shaping, data augmentation, etc.), cap the intervention with a budget and spend it on the highest-leverage positions. The paper shows you don’t need to intervene everywhere to get most of the benefit.

3. The relay framing itself: Instead of thinking of distillation as “teacher supervises student end-to-end,” think of it as a collaborative relay where different agents can contribute legs of the trajectory. This framing opens up possibilities beyond just teacher-student — multiple specialist models could contribute relay legs for different parts of the reasoning process.