PaperlayerView exact revision on arXiv

Source-linked AI summary

Pass the Baton: Trajectory-Relayed On-Policy Distillation

Haolei Xu, Xiaowen Xu, Haiwen Hong, Zixuan Ni, Hongxing Li, Yiwen Qiu, Weiming Lu, Yongliang Shen

arXiv:2607.26057v1cs.CLcs.AI

TL;DR

On-policy distillation can fail when students commit early to wrong reasoning directions, producing unreliable, wasteful continuations. Relay-OPD detects teacher–student divergence and briefly relays control to the teacher, outperforming OPD and FastOPD across eight mathematical benchmarks while shortening training trajectories.

  • Problem

    On-policy distillation suffers prefix failure: early wrong reasoning directions produce long, misdirected continuations with unreliable supervision and substantial compute waste.

  • Method

    Relay-OPD detects label-free teacher–student continuation divergences and inserts budgeted, brief teacher legs into student trajectories before returning control to the student.

  • Results

    +5.73% over standard OPD and +1.49% over FastOPD on average for the 1.7B student across eight mathematical reasoning benchmarks, with over 50% shorter training trajectories.

  • Takeaways & Limitations

    Early, local teacher intervention at detected failure points improves on-policy distillation while keeping relay trajectories close to the student policy.

  • Takeaways & Limitations

    Evaluation is limited to mathematical reasoning with Qwen3 teacher–student pairs, and the method presumes teachers redirect failed prefixes more reliably than students.

Abstract

from arXiv · show

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%.

1 INTRODUCTION

Relay-OPD addresses prefix failure by detecting teacher–student continuation asymmetry and inserting brief, early teacher interventions into student-generated trajectories. Across eight mathematical reasoning benchmarks, it improves over OPD and FastOPD while reducing training trajectory length.

  • Prefix failure causes long misdirected continuations, unreliable supervision, and substantial compute waste after the student commits early to a wrong reasoning direction.
  • On failed prefixes, teachers tend to re-examine and redirect, while students continue in the wrong direction; this observable asymmetry defines a label-free handoff trigger.
  • 0.35% teacher-generated tokens increase accuracy from 27.73 to 34.96 (+7.23%), showing that correction can be early and highly local.
  • Relay-OPD interleaves student generation with brief teacher legs at detected triggers, then returns control to the student through a budgeted relay rollout.
  • +5.73% over standard OPD and +1.49% over FastOPD are achieved for the 1.7B student across eight math reasoning benchmarks, while training trajectories shrink by over 50%.

2 RELATED WORK

On-policy distillation supervises tokens on trajectories sampled from the student, grounding transfer in the student’s visited states rather than teacher-generated data. However, this grounding also imports student failures, making supervision unreliable when prefixes drift from teacher-supported reasoning states, especially in long-chain reasoning.

  • On-policy distillation: On-policy distillation has the teacher score tokens on student-sampled prefixes, grounding supervision in the student’s own state distribution.Compared with offline training on teacher-generated data, it yields faster and stronger transfer and is standard in post-training pipelines.
  • Limitations: Student-trajectory grounding imports failures, shrinking distillation gains when teacher and student reasoning patterns are incompatible.The passage attributes this limitation to incompatibility between teacher and student reasoning patterns.
  • Limitations: Early directional deviations in long-chain reasoning compound autoregressively into extended misdirected continuations, making subsequent supervision increasingly unreliable.Reliability declines as student prefixes drift from teacher-supported states.

3 METHOD

Relay-OPD detects reasoning-direction failures online, briefly hands control to the teacher, and resumes student generation on the resulting relay trajectory. Its state-driven construction, trajectory-level objective, and single-engine speculative decoding provide corrected context while preserving exact sampling and limiting intervention.

  • Relay trajectory construction: Relay trajectories alternate student and teacher legs, with a budget limiting takeovers and focusing intervention early in generation.Teacher-leg extension saturates: from L = 3 to L = 6, accuracy plateaus around 41–44 while teacher token ratio rises from 17.52% to 28.52%.
  • Handoff trigger: Relay-OPD detects handoff triggers when the teacher favors a reflection token while the student’s top-K support contains none, indicating divergent reasoning directions.The trigger requires no external verifier or process labels; K controls its sensitivity.
  • Optimization objective: 43.55 accuracy at L = 6 with a 28.52% teacher token ratio remains below 60.55 for the independently generating teacher, motivating optimization on observed relay trajectories.The relay objective is therefore designed around supervision from the generated relay process rather than ideal teacher-only trajectories.
  • Optimization objective: The objective optimizes actually generated tokens across the entire relay trajectory, so teacher legs correct subsequent student context and directly contribute training tokens.When L = 0, the teacher leg contains only the replacement token, yielding minimal prefix correction.
  • Efficient implementation: Relay-OPD uses a single state-switched decoding engine that alternates student and teacher targets during generation and terminates at end-of-sequence or length limits.The engine directly emits the teacher’s trigger token at each teacher-leg start, while later teacher-leg positions use speculative decoding.
  • Efficient implementation: Every verified position is sampled from its target policy, so the single-engine implementation exactly reproduces the two-model relay process while batching teacher-leg draft verification.Teacher legs are distributionally identical to direct teacher continuation from the extended prefix, and student legs to ordinary student sampling.

4 EXPERIMENTS

Across eight mathematical reasoning benchmarks, Relay-OPD delivers consistently strong gains over standard OPD and trajectory-intervention baselines for both student sizes. Its improvements are associated with shorter training trajectories, adaptive teacher intervention, and ablations supporting brief, selective relays with the relay-token objective.

  • Overall performance: Relay-OPD achieves the best or second-best results across all eight benchmarks for both student models.Experiments use Qwen3-4B-Instruct-2507 as teacher, Qwen3-0.6B/1.7B-Non-Thinking as students, and DAPO-Math-17K training data.
  • Overall performance: 46.96 average accuracy on 1.7B exceeds standard OPD by +5.73% and FastOPD by +1.49%.AIME 2025 and AIME 2026 improve over OPD by +7.29% and +7.19%, respectively; the 0.6B student improves over OPD by +3.01% and FastOPD by +0.62%.
  • Training token efficiency: 35 training steps and 2,296 average rollout tokens let Relay-OPD reach its 1.7B optimum before OPD at step 55 and FastOPD at step 45.The 2,296-token rollout length is a 50.7% reduction from OPD’s 4,658 and is shorter than FastOPD’s 2,709.
  • Training dynamics: Teacher token ratio falls from roughly 13% to 2%–3% after about 20 steps, while relay-budget exhaustion decreases from roughly 75%–85% to 50%–60%.Relay-OPD also maintains consistently higher policy entropy than OPD and FastOPD during the first 60 steps.
  • Role of the teacher leg: 46.25 average accuracy with a teacher leg surpasses Trigger-stop’s 43.48 by +2.77%, supporting corrected context and local reasoning demonstrations beyond truncation.The comparison sets the maximum number of takeovers to M = 1; Trigger-stop generates no teacher tokens, whereas Relay-OPD uses L = 3.

5 CONCLUSION · A TRAINING CONFIGURATION · A.1 HANDOFF TRIGGER WORD LIST

Relay-OPD addresses prefix failure by detecting online teacher–student reasoning-direction divergences and enabling local teacher intervention at failure points. Its limited relay budget focuses intervention on critical early positions while limiting departure from the student policy, and the handoff criterion uses a specified reflection-token set.

  • 5 CONCLUSION: Relay-OPD detects teacher–student reasoning-direction divergences online to address prefix failure in on-policy distillation.The method lets the teacher intervene locally at detected failure points.
  • 5 CONCLUSION: The teacher intervenes locally at detected reasoning failure points, creating a targeted handoff mechanism.Intervention is performed at detected failure points rather than throughout the trajectory.
  • A TRAINING CONFIGURATION: A limited relay budget concentrates intervention on critical early positions while limiting departure from the student policy.This budget is part of the training configuration described for Relay-OPD.
  • 5 CONCLUSION: Across two Qwen3 student models and eight mathematical reasoning benchmarks, Relay-OPD outperforms standard OPD by +5.73%.The reported comparison covers both student-model sizes and all eight benchmarks.
  • A.1 HANDOFF TRIGGER WORD LIST: The handoff criterion uses a reflection token set R consisting of base words together with their case and leading-space variants.The listed set is part of the handoff-trigger implementation.
  • A.1 HANDOFF TRIGGER WORD LIST: The handoff-trigger word list includes Wait, But, Hmm, Actually, Hold, However, Yet, Oh, Alternatively, No, Ah, Oops, and Well.These are the base words in the listed reflection token set.

A.2 TRAINING HYPERPARAMETERS

This section lists the complete training hyperparameters for the online distillation methods and GRPO, with the online-distillation settings presented in Table 3 and the GRPO settings in Table 4.

  • A.2 TRAINING HYPERPARAMETERS: The complete GRPO training hyperparameters are provided in Table 4.The passage states that Table 4 lists the GRPO hyperparameters.
  • A.2 TRAINING HYPERPARAMETERS: The complete training hyperparameters for online distillation methods are provided in Table 3.Table 3 is titled “Training hyperparameters for online distillation methods.”

B PROMPTS · C ALGORITHM AND BASELINES · C.1 RELAY-OPD ALGORITHM

The section specifies the prompts used for training and inference, the TRD rewrite prompt, and the Relay-OPD relay-rollout construction procedure. Relay-rollout construction samples student drafts, detects handoffs from teacher and student token distributions, and permits an S→T transition under a relay budget.

  • B.1 TRAINING AND INFERENCE PROMPT: Training and inference use a template that asks the model to reason step by step and place the final answer within \boxed{}.The template supplies the problem as user content and begins the assistant response with <think>.
  • B.2 TRD REWRITE PROMPT: The TRD rewrite prompt asks the model to rewrite its mathematical solution.It is adopted from the original TRD paper by Jiang et al. (2026).
  • B.2 TRD REWRITE PROMPT: The rewrite instructions preserve the original solution’s structure and reasoning path while fixing computational or logical errors.They also require retaining correct intermediate steps and meaningful work, then outputting only the rewritten solution.
  • C.1 RELAY-OPD ALGORITHM: Relay-OPD constructs relay trajectories from prompt x, student policy π¯θ, teacher policy πT, a reflection-token set R, handoff top-K, and relay budget (M, L).The procedure initializes an empty trajectory, takeover count j ← 0, and state s ← S.
  • C.1 RELAY-OPD ALGORITHM: At each step, the student samples a draft token, while the teacher supplies an argmax token and the student supplies its top-K token set.The prefix is h ← (x, z), and the rollout continues until EOS or maximum length.
  • C.1 RELAY-OPD ALGORITHM: The handoff criterion is ϕ(h) = 1[aT ∈ R] · 1[KS(h) ∩ R = ∅], triggering an S→T transition when s = S, ϕ(h) = 1, and j < M.Thus, the teacher argmax must be a reflection token while the student top-K set contains no reflection token.

C.2 METHOD COMPARISON OVERVIEW … D SUPPLEMENTARY RESULTS

The paper compares methods by policy type, trajectory source, training loss, and method-specific settings. TRD, SKD, and FastOPD differ mainly in how they construct or truncate training trajectories and provide supervision.

  • C.2 METHOD COMPARISON OVERVIEW: Table 5 organizes compared methods by policy type, trajectory source, training loss, and method-specific settings.The comparison uniformly uses top-k=128 for FKL-based papers; SKD uses top-k=5 instead of its original acceptance top-k=25.
  • C.3 TRD: TRD creates training data offline by rewriting a student-generated trajectory with a teacher-conditioned prompt.The student first generates yo from x, then the teacher produces rewritten trajectory yr using x and yo.
  • C.3 TRD: TRD applies top-k FKL with teacher probabilities renormalized over a support set of k = 128.The objective is defined on the rewritten trajectory, using distinct teacher and student contexts.
  • C.3 TRD: TRD conditions teacher supervision on the student’s initial trajectory yo, which the student cannot access during training or inference.For rewritten-trajectory token prediction, the student context uses x and yr,<t, while the teacher context additionally includes yo.
  • C.4 SKD: SKD constructs mixed teacher–student trajectories through speculative decoding, using the student as draft model and teacher as target model.At each prefix, the teacher uses a top-K support set with K = 5; rejected draft tokens receive teacher-sampled replacements.
  • C.4 SKD: SKD marks replacement positions as teacher-owned and optimizes top-k FKL on the mixed trajectory with shared teacher–student prefixes.The shared prefix is ht = (x, z<t).
  • C.5 FASTOPD: FastOPD truncates student rollouts at a fixed length B before applying the standard OPD objective.It replaces N and y with NB and y(B), changing only the rollout length participating in training, not the token-level learning signal.

D.1 COMPLETE SENSITIVITY RESULTS · D.2 FASTOPD TRUNCATION LENGTH SWEEP · E CASE STUDIES

The sensitivity analysis reports complete per-benchmark relay-budget and handoff top-K sweeps, while the FastOPD truncation sweep identifies 4,096 tokens as the best average-accuracy configuration for both students.

  • D.1 COMPLETE SENSITIVITY RESULTS: D.1 COMPLETE SENSITIVITY RESULTS reports complete per-benchmark results for the relay-budget and handoff top-K sweeps.These sweeps are described as the complete results corresponding to §4.5.
  • D.1 COMPLETE SENSITIVITY RESULTS: The sensitivity table varies one hyperparameter at a time using Qwen3-4B-Instruct-2507 as teacher and Qwen3-1.7B-Non-Thinking as student.The table specifies the teacher, student, and one-at-a-time hyperparameter variation.
  • D.1 COMPLETE SENSITIVITY RESULTS: Standard OPD is represented by K = |V| in the sensitivity comparison.The table defines this setting explicitly as the Standard OPD reference.
  • D.1 COMPLETE SENSITIVITY RESULTS: Bold values denote the best result in each sensitivity-table column.The table caption states the convention used to identify column-wise best results.
  • D.2 FASTOPD TRUNCATION LENGTH SWEEP: D.2 FASTOPD TRUNCATION LENGTH SWEEP reports complete results across fixed truncation lengths.The fixed-length comparisons are provided in Table 7.
  • D.2 FASTOPD TRUNCATION LENGTH SWEEP: 4,096 tokens yields the highest average accuracy for both student models.This configuration is consequently used for the results reported in Table 1.
  • D.2 FASTOPD TRUNCATION LENGTH SWEEP: Table 7 marks the best checkpoint's training step with subscripts in Avg and defines Train Len as average rollout response length through that checkpoint.These conventions describe how average accuracy and training length are presented in the truncation comparison.

E.1 RELAY-OPD TAKEOVER CASE · E.2 TRD REWRITTEN TEXT · E.3 SKD REPETITION PATTERNS

The examples show Relay-OPD correcting a failed arithmetic continuation through teacher takeover, while TRD exposes rewritten-trajectory artifacts and SKD can perpetuate repetitive answer blocks. Together, these cases illustrate intervention benefits and failure modes in trajectory-based distillation.

  • E.1 RELAY-OPD TAKEOVER CASE: With only 5 RM remaining after 15 apples, the student incorrectly concludes that buying both a mango and papaya is affordable, producing 17 fruits.The ground-truth maximum is 15 fruits.
  • E.1 RELAY-OPD TAKEOVER CASE: 74.4084%: Faced with the failed prefix, the teacher selects “But,” while the student favors “So” or “Now” to extend its conclusion.This distributional asymmetry triggers the teacher takeover.
  • E.1 RELAY-OPD TAKEOVER CASE: The teacher leg identifies that the required 9 RM for one mango and one papaya exceeds the 5 RM remaining, redirecting the trajectory away from the invalid 17-fruit answer.The intervention explicitly compares the remaining budget with the combined minimum cost.
  • E.2 TRD REWRITTEN TEXT: 18.96%: Trajectories contain at least one rewriting artifact such as references to the original solution, revision descriptions, or reviewer-style openings.These expressions were matched against teacher-generated rewritten trajectories and indicate that the artifacts are not isolated.
  • E.2 TRD REWRITTEN TEXT: Teacher-generated rewritten text can acknowledge an incorrect parity assumption and uncertainty before settling on a periodic answer such as 674.The fragment describes revising after manual simulation appeared inconsistent.
  • E.3 SKD REPETITION PATTERNS: SKD accepts a student token whenever it lies within the teacher’s top-K support, allowing repetitive continuations to persist even when the teacher would terminate or redirect.Generic distributional agreement does not explicitly detect that the established generation pattern has become erroneous.
  • E.3 SKD REPETITION PATTERNS: 73.7035%: The teacher’s top-1 token is <|im_end|>, but the student assigns 78.5176% to opening another paragraph; with K = 5, that token remains in teacher support at 16.4455%.The student draft is therefore accepted, and recurring distributions can produce an infinite loop of final-answer blocks.

F LIMITATIONS

Relay-OPD is evaluated only on mathematical reasoning with Qwen3 teacher–student pairs, while its application beyond this setting remains future work. Its transfer may require adjusting the reflection-token set for other model families and assumes a teacher that redirects failed prefixes.

  • Evaluation scope: Evaluation covers mathematical reasoning with Qwen3 teacher–student pairs in the strong-to-weak setting targeted by prior on-policy distillation work.The passage does not report evaluation on other tasks or model families.
  • Future applications: Applying the task-agnostic relay mechanism to code generation or agentic tool use is left for future work.These domains are identified as examples beyond the reported evaluation scope.
  • Model-family transfer: Switching to a different model family may require adjusting the reflection-token set described in Appendix A.1.The limitation concerns transfer across model families rather than the relay mechanism’s stated task agnosticism.
  • Teacher assumption: Relay-OPD presumes a teacher whose continuations redirect failed prefixes more readily than the student.The supplied passage introduces this teacher-behavior assumption but is truncated before providing further qualification.
Loading 2607.26057v1…