On-Policy Distillation with Best-of-N Teacher Rollout Selection

• URL: http://arxiv.org/abs/2605.09725v1
• Date: Sun, 10 May 2026 19:49:00 GMT
• Title: On-Policy Distillation with Best-of-N Teacher Rollout Selection
• Abstract summary: We propose BRTS, a Best-of-N Rollout Teacher Selection framework for on-policy distillation.<n>BRTS augments standard student-context OPD with a teacher-context supervision branch constructed from the teacher trajectory.<n>BRTS improves over standard OPD on challenging reasoning benchmarks, with the largest gains on harder datasets.
• Score: 54.91780727674628
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: On-policy distillation (OPD), which supervises a student on its own sampled trajectories, has emerged as a data-efficient post-training method for improving reasoning while avoiding the reward dependence of reinforcement learning and the catastrophic forgetting often observed in standard supervised fine-tuning. However, standard OPD typically computes teacher supervision under noisy student-generated contexts and often relies on a single stochastic teacher rollout per prompt. As a result, the supervision signal can be high-variance: the sampled teacher trajectory can be incorrect, uninformative, or poorly matched to the student's current reasoning behavior. To address this limitation, we propose BRTS, a Best-of-N Rollout Teacher Selection framework for on-policy distillation. BRTS augments standard student-context OPD with a teacher-context supervision branch constructed from the curated teacher trajectory. Rather than distilling from the first sampled teacher rollout, BRTS samples a small pool of teacher trajectories and selects the auxiliary trajectory using a simple priority rule: correctness first, student alignment second. When multiple correct teacher trajectories are available, BRTS chooses the one most aligned with the student's current behavior; when unconditioned teacher samples fail on harder prompts, it invokes a ground-truth-conditioned recovery step to elicit a natural derivation. The selected trajectory is then used to provide reliable teacher-context supervision inside the OPD loop, augmented with an auxiliary loss on the teacher trajectory. Experiments on AIME 2024, AIME 2025, and AMC 2023 show that BRTS improves over standard OPD on challenging reasoning benchmarks, with the largest gains on harder datasets. Our code is available at https://github.com/BWGZK-keke/BRTS.

Related papers

• When Teacher Guidance Misleads: Reward-Aligned On-Policy Distillation [21.86522736658343]
  We propose Reward-Aligned On-Policy Distillation (RA-OPD) to mitigate misaligned teacher guidance.<n> RA-OPD selects more reliable trajectories to improve student model performance without requiring additional computational cost.<n>We evaluate RA-OPD on math and code benchmarks using models from the Qwen3 family and the DeepSeek-R1 family.
  arXiv  Detail & Related papers  (2026-08-28T06:05:50Z)
• ReOrder-OPD:Reliability-Aware Prompt Ordering for On-Policy Distillation [25.69241346799495]
  We define prompt-level teacher continuation reliability $R$ as the teacher's probability of reaching a correct answer from a student prefix.<n>We show that high-$R$ prompts yield larger OPD gains and that descending-$R$ training outperforms random and ascending orders.
  arXiv  Detail & Related papers  (2026-08-11T13:27:56Z)
• Pass the Baton: Trajectory-Relayed On-Policy Distillation [23.50689651536186]
  On-policy distillation grounds token-level supervision in the student's own trajectory, yet suffers from prefix failure.<n>We identify a teacher-student continuation asymmetry on failed prefixes, and convert it into a label-free handoff trigger in Relay On-Policy Distillation (Relay-OPD)<n>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.
  arXiv  Detail & Related papers  (2026-07-28T17:59:46Z)
• H$^2$SD: Hybrid Hindsight Self-Distillation [61.71131241473028]
  Reinforcement learning with verifiable rewards (RLVR) provides reliable outcome supervision for language model reasoning.<n>Existing self-distillation methods add a privileged teacher but typically assign it a fixed role.<n>We introduce Hybrid Hindsight Self-Distillation ($mathrmH2mathrmSD$), which jointly adapts teacher context and update strategy to trajectory correctness.
  arXiv  Detail & Related papers  (2026-07-21T10:47:27Z)
• Reward-Gated On-Policy Distillation [58.187458236302575]
  On-policy distillation is a powerful way to transfer reasoning ability from a strong teacher to a smaller student.<n>We introduce RG-OPD: Reward-Gated On-Policy Distillation that uses verifier feedback to decide when teacher logits should be trusted.<n>Across reasoning and coding benchmarks, RG-OPD produces stronger distilled students, outperforming both vanilla reverse-KL distillation and the recent TSD-KD baseline.
  arXiv  Detail & Related papers  (2026-07-04T21:51:22Z)
• Your Teacher Can't Help You Here: Combating Supervision Fidelity Decay in On-Policy Distillation [81.10000755917712]
  On-policy distillation transfers reasoning capabilities by training a student model on its own generated trajectories using token-level feedback from a teacher.<n>As student-generated prefixes lengthen, the teacher's next-token distribution becomes less confident and less discriminative.<n>To mitigate SFD, we introduce textbfLookahead Group Reward (ours).
  arXiv  Detail & Related papers  (2026-05-29T04:39:20Z)
• Prefix Teach, Suffix Fade: Local Teachability Collapse in Strong-to-Weak On-Policy Distillation [49.117085054884676]
  On-policy distillation trains a student model on its own rollouts using dense feedback from a stronger teacher.<n>We operationalize this principle through a trajectory-specific release rule.<n> Experimental results across strong-to-weak distillation tasks indicate that this release rule consistently outperforms standard full-trajectory OPD.
  arXiv  Detail & Related papers  (2026-05-13T15:05:30Z)
• Prune-OPD: Efficient and Reliable On-Policy Distillation for Long-Horizon Reasoning [66.52232008796294]
  Prune-OPD dynamically aligns training budgets with supervision quality.<n>It reduces training time by 37.6%--68.0% while preserving, and often improving, performance on challenging benchmarks.
  arXiv  Detail & Related papers  (2026-05-08T14:38:53Z)
• Preference-Based Self-Distillation: Beyond KL Matching via Reward Regularization [18.027254451537342]
  Existing self-distillation methods largely reduce learning to KL matching toward the context-augmented teacher model.<n>We propose textbfPreference-textbfBased textbfSelf-textbfDistillation (textbfPBSD), which revisits on-policy self-distillation through a reward-regularized perspective.
  arXiv  Detail & Related papers  (2026-05-06T15:31:50Z)
• TCOD: Exploring Temporal Curriculum in On-Policy Distillation for Multi-turn Autonomous Agents [55.27396165691312]
  We identify a key limitation of vanilla OPD in multi-turn agent settings, which we term Trajectory-Level KL Instability.<n>We propose TCOD, a framework that controls the trajectory depth exposed to the student and progressively expands it with a curriculum schedule.<n> Experimental results across four student-teacher pairs show that TCOD mitigates KL escalation and enhances KL stability throughout training, improving agent performance by up to 18 points over vanilla OPD.
  arXiv  Detail & Related papers  (2026-04-27T03:38:27Z)
• Learning beyond Teacher: Generalized On-Policy Distillation with Reward Extrapolation [57.524909883706556]
  On-policy distillation (OPD) has demonstrated strong empirical gains in improving student performance.<n>This work introduces a flexible reference model and a reward scaling factor that controls the relative weight of the reward term against the KL regularization.<n>In particular, in the setting where we merge the knowledge from different domain experts, obtained by applying domain-specific RL to the same student model, ExOPD enables the student to even surpass the teacher's performance boundary.
  arXiv  Detail & Related papers  (2026-02-12T16:14:29Z)
• Which Reasoning Trajectories Teach Students to Reason Better? A Simple Metric of Informative Alignment [82.00769536768509]
  Rank-Surprisal Ratio is a simple metric that captures both alignment and informativeness to assess the suitability of a reasoning trajectory.<n>We demonstrate its practical utility in both trajectory selection and teacher selection.
  arXiv  Detail & Related papers  (2026-01-20T18:58:10Z)
• Distilling Reasoning into Student LLMs: Local Naturalness for Selecting Teacher Data [18.97748910748554]
  We introduce Local Naturalness, which measures the student's log-probabilities over short, sequential reasoning steps conditioned only on a small local window.<n>When mixing answers from many teachers, Local Naturalness boosts a 32B student's accuracy on math benchmarks by 9.4pp over global selection.<n>These results highlight the power of localized data quality evaluation and data mixing for more effective reasoning distillation.
  arXiv  Detail & Related papers  (2025-10-05T01:15:32Z)
• From Correction to Mastery: Reinforced Distillation of Large Language Model Agents [13.982204994247718]
  Large Language Model agents excel at solving complex tasks through iterative reasoning and tool use.<n>Existing distillation approaches train smaller students to imitate full teacher trajectories.<n>We propose SCoRe, a student-centered framework in which the student generates training trajectories and the teacher corrects only the earliest error.
  arXiv  Detail & Related papers  (2025-09-12T15:34:07Z)
• Merge-of-Thought Distillation [23.53356244978525]
  Merge-of-Thought Distillation (MoT) is a lightweight framework that alternates between teacher-specific supervised fine-tuning branches and weight-space merging the resulting student variants.<n>On competition math benchmarks, applying MoT to a Qwen3-14B student surpasses strong models including Deepseek-R1, Qwen3-32B, and OpenAI-O1.<n>MoT consistently outperforms the best single-teacher distillation, improves general reasoning beyond mathematics, and shows robustness to distribution-shifted and peer-level teachers.
  arXiv  Detail & Related papers  (2025-09-10T17:46:57Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.