PaperlayerView exact revision on arXiv

Source-linked AI summary

Proposing Better Rollouts for On-Policy Distillation

Sirui Chen, Yunzhe Qi, Haobo Xu, Jiaru Zou, Mengting Ai, Zihao Li, Jingrui He

arXiv:2610.36477v1cs.LGcs.AI

TL;DR

On-policy distillation can require teachers to supervise student-sampled trajectories that are difficult for the teacher to reach, creating a mismatch across rounds. MPOPD derives a student-anchored, teacher-guided rollout distribution and approximates it with a proposal policy. Across 16 reasoning benchmarks it consistently outperforms OPD and recent variants, with higher average accuracy and shorter generation on Qwen3-4B-Base math tasks.

  • Problem

    Because OPD samples from the student, the teacher must supervise low-reachability student trajectories, where its token-level guidance can be noisy or misleading.

    Source: §1 · p. 1

  • Method

    MPOPD derives a geometric mixture of student and teacher policies from relative teacher support under a local trust-region constraint, then approximates it with a proposal policy.

    Source: §1 · p. 2

  • Results

    Across 16 math, knowledge, and medical reasoning benchmarks MPOPD consistently outperforms OPD and recent variants; on Qwen3-4B-Base math tasks, average accuracy improves by +6.91 over OPD while average generation length falls by 41%.

    Source: §4.2 · p. 7, §6 · p. 9

  • Takeaways & Limitations

    Diagnostics indicate that gains come from shifting rollouts toward more teacher-reachable trajectories, making rollout policy a critical design choice for OPD.

    Source: §6 · p. 9

    The empirical study focuses on the Qwen3 family, with main evaluations at the 1.7B and 4B scales across math, knowledge QA, and medical reasoning.

    Source: C Discussion · p. 23

Abstract

from arXiv · show

On-policy distillation (OPD) improves language-model post-training by letting the student learn from its own sampled trajectories under teacher supervision, aligning the training process with inference-time behavior. However, the rollout policy in OPD is usually treated as fixed: each round samples directly from the current student. We show that this choice can create a compounding teacher--student mismatch, because it requires the teacher to provide effective supervision on whichever trajectories the student samples, including ones less reachable by the teacher itself. On such trajectories, teacher supervision can become noisy or even misleading, leading to undesirable student patterns such as repetition. Moreover, since the updated student in turn defines the rollout policy of subsequent rounds, these patterns persist and potentially amplify over training. This failure mode suggests that improving OPD requires treating the rollout policy itself as an optimization target, rather than only modifying the distillation objective. We propose Mixture Proposal On-Policy Distillation (MPOPD), which derives a student-anchored, teacher-guided rollout distribution from a local trust-region objective based on relative teacher support. The resulting target is a geometric mixture of the student and teacher policies, and we introduce a proposal policy that approximates this mixture, enabling rollout generation in a single forward pass. Experiments on 16 math, knowledge, and medical reasoning benchmarks show consistent gains over OPD and recent variants. On math reasoning, our method improves average accuracy by 6.91 over OPD on Qwen3-4B-Base while reducing average generation length by 41%, suggesting that rollout policy design provides an effective way to improve OPD.

1 Introduction

On-policy distillation can compound teacher–student mismatch when student rollouts enter regions where the teacher provides unreliable supervision. MPOPD instead optimizes the rollout policy, and the paper reports consistent benchmark gains over OPD and recent variants.

1 Introduction

  • Student trajectories can enter low teacher-reachability regions, where supervision becomes noisy or misleading and degenerate patterns can persist across distillation rounds.The paper identifies repetition and entropy collapse as related failure modes reported in prior studies.Source: §1 · pp. 1–2
  • Under OPD, the 10th-percentile similarity between teacher completions drops by 0.24 from its peak, while the fraction of prompts with a lower PRM score reaches about 24% by training’s end.The paper reports that its method keeps both measures stable or near zero.Source: Figure 1 · p. 2, §1 · p. 2
  • MPOPD optimizes rollouts using relative teacher support within a local trust region, yielding a geometric mixture of student and teacher policies that a proposal policy approximates for single-pass generation.The method shifts probability toward teacher-supported trajectories while remaining close to the current student.Source: §1 · p. 2
  • Across 16 math, knowledge, and medical reasoning benchmarks, MPOPD consistently outperforms OPD and recent variants; on Qwen3-4B-Base, average accuracy improves by +6.91 over OPD while generation length falls by 41%.The reported math results also show an average accuracy gain of +1.46 on Qwen3-1.7B-Base.Source: §1 · p. 2

2 Setup: The Rollout Policy in On-Policy Distillation

Standard OPD uses the current student as its rollout policy, aligning samples with student behavior but exposing the teacher to low-reachability trajectories. The section frames rollout design as a local correction that increases teacher reachability without replacing the student.

2 Setup: The Rollout Policy in On-Policy Distillation

  • Standard OPD sets the rollout policy to the current student, so training uses trajectories sampled from the student’s own behavior.At each round, the teacher supervises those sampled responses through a token-level objective.Source: §2 · p. 3
  • The proposed rollout policy should upweight tokens more reachable under the teacher while deviating from the student only locally.This treats the rollout policy as a design choice and a correction to student sampling, not a replacement.Source: §2 · p. 3
  • A multiplicative reweighting keeps the rollout anchored to the student while redistributing probability toward trajectories better aligned with teacher behavior.The section previews an exponential tilt and a learnable proposal policy for efficient rollout generation.Source: §2 · p. 3

3 Mixture Proposal for On-Policy Distillation

MPOPD reweights student rollouts according to relative teacher support while constraining their divergence from the student. A learnable proposal approximates the resulting tilted target so rollout generation does not require synchronized teacher and student decoding.

3.1 Reweighting by Relative Teacher Support

  • Relative teacher support is the log probability ratio between teacher and student for a candidate token, prioritizing tokens the teacher supports more strongly than the student does.The paper motivates this signal as favoring effective teacher supervision on tokens the student has not yet fully learned.Source: §3.1 · p. 3
  • Maximizing expected relative teacher support under a local KL constraint yields a student-anchored geometric mixture of student and teacher policies.The tilt strength η interpolates between the student at η = 0 and the teacher at η = 1.Source: §3.1 · p. 4
  • The benefit of tilting depends on the spread of relative-support scores under the student and the trust-region radius, with diminishing returns as local deviations grow.When score spread is small, the student already agrees with the teacher up to a constant offset and tilting has little to redistribute.Source: §3.1 · p. 4

3.2 Learning a Proposal Rollout Policy

  • Directly sampling the tilted target would require synchronized teacher and student forward passes at each decoding step, so MPOPD trains a proposal policy to approximate it.Rollouts are sampled from the proposal, while the student update remains defined against the student and teacher.Source: §3.2 · p. 4
  • Training alternates between student updates from proposal-sampled trajectories and proposal updates toward the next tilted target; rollout generation then uses the proposal alone.The student update corrects for the proposal–student gap with clipped importance weights.Source: §3.2 · p. 5

4 Experiments

Experiments evaluate MPOPD on math, knowledge QA, and medical reasoning, comparing its performance and training costs with OPD and other variants. Results include higher accuracy, shorter generations, and diagnostic evidence about rollout quality and training efficiency.

4.1 Experimental Setup

  • The evaluation covers two Qwen3-Base student scales, eight math benchmarks, three knowledge QA benchmarks, and five medical reasoning benchmarks.The math experiments train on DAPO-Math-17K; the additional domains use ARPO-10k and M23K, respectively.Source: §4.1 · pp. 5–6

4.2 Main Results

  • On Qwen3-4B-Base, MPOPD achieves the best accuracy on all eight math benchmarks, improving over the strongest baseline by +4.31 on average.On Qwen3-1.7B-Base, it is best on 7 of 8 benchmarks, with an average gain of +0.82.Source: §4.2 · p. 6, Table 1 · p. 6
  • On Qwen3-4B-Base, MPOPD reduces average generation length by 41% relative to OPD while improving average accuracy by +6.91.The length gap mainly comes from baselines producing degenerate responses that fail to terminate; Figure 3 illustrates OPD repetition and EOPD spurious behavior at the 16.4k-token limit.Source: §4.2 · p. 7, Figure 3 · p. 7
  • On Qwen3-4B-Base, MPOPD achieves the best accuracy on all knowledge QA and medical reasoning tasks, raising strongest-baseline averages from 65.34 to 66.54 and 58.03 to 59.17, respectively.Compared with OPD, average generation length drops from 10.91k to 2.65k on knowledge QA and from 8.00k to 0.88k on medical reasoning.Source: §4.2 · pp. 7–8, Table 2 · p. 7

4.3 Training Efficiency

  • MPOPD incurs about 13% extra wall-clock time on average, while faster rollout generation partially offsets the added proposal-policy training cost.Against OPD with two student updates per round, MPOPD reaches 35.37% average accuracy in 131.12s per round, versus 28.59% in 150.01s; its rollout generation takes 94.68s versus about 130s.Source: §4.3 · p. 8, Figure 4 · p. 8, Table 3 · p. 8

4.4 Analysis

  • MPOPD earns higher training reward while reducing responses that reach the maximum generation length by 45 points, and its generated trajectories receive +0.11 teacher log-probability.The student–teacher entropy gap shrinks by 0.04 and the fraction of loss from low-confidence teacher positions drops by 17 points. Accuracy varies within about 0.76 across tilt strengths and peaks at η = 0.20.Source: §4.4 · pp. 8–9, Figure 5 · p. 8, Table 4 · p. 9

5 Related Work

Prior work spans dense teacher-guided distillation and reward-based reasoning post-training. OPD research has addressed its objectives and rollout heuristics, while this paper studies the rollout policy itself.

Reasoning in Large Language Models

  • Reinforcement learning with verifiable rewards uses sparse outcome-level feedback, whereas distillation supplies dense token-level teacher signals [42] [59] [64] [13] [45], [41] [25].Distillation is also used with reasoning corpora [12] [61] [2] and to transfer reasoning across scales and modalities [4] [37].Source: §5 · p. 9

On-Policy Distillation

  • OPD, which lets students generate trajectories for token-level teacher supervision to address off-policy train–test mismatch, has been explored through cross-family, multi-teacher [9] [50], black-box [58], and self-distillation [15] [62] [20] approaches.Source: §5 · p. 9
  • Reported OPD failure modes include repetition, overlong generations, and entropy collapse [8] [24] [26] [31].Prior remedies mainly alter the supervision target [1], [57] [24], [21] [26], or use heuristic teacher–student interleaving for rollouts.Source: §5 · p. 9

6 Conclusion

The conclusion presents MPOPD as a rollout-policy optimization that targets compounding teacher–student mismatch and consistently outperforms OPD and recent variants across 16 reasoning benchmarks. The appendix describes training configurations across math, knowledge QA, and medical reasoning.

6 Conclusion

  • MPOPD uses a local trust-region objective under relative teacher support to derive a geometric student–teacher mixture, approximated by a learnable proposal policy for single-forward-pass rollouts.Across 16 math, knowledge, and medical reasoning benchmarks, it consistently outperforms OPD and recent variants; diagnostics associate the gains with rollouts shifted toward more teacher-reachable trajectories.Source: §6 · p. 9

A.1.1 Training

  • Math and knowledge QA use batch size 256, while medical reasoning uses batch size 128; the configuration summary covers training across domains.Medical reasoning trains on a 5K subset; math and knowledge QA use each dataset’s full split.Source: §A.1.1 · p. 17, Table 5 · p. 17

A.1.2 Evaluation

Evaluation uses separate prompting and decoding protocols for math, knowledge QA, and medical reasoning. The implementation appendix describes teacher queries, student updates, proposal updates, and baseline parity.

A.1.2 Evaluation

  • Math evaluation uses step-by-step boxed-answer prompts, temperature 0.8, top-p 0.9, and up to 16,384 tokens, with 32 completions per problem except eight on OlympiadBench.Accuracy and average generated length are reported.Source: §A.1.2 · p. 17
  • Knowledge QA uses zero-shot MCQA for GPQA-Diamond and ARC-Challenge, while MMLU-Pro uses subject-specific few-shot validation exemplars.ARC-Challenge and MMLU-Pro use 500 randomly sampled test questions each; GPQA-Diamond uses its full test set, with eight completions per question.Source: §A.1.2 · p. 17
  • Medical reasoning uses a step-by-step MCQA prompt and deterministic decoding at temperature 0.0, top-p 1.0, and one completion per question.The maximum generation length is 8,192 tokens; accuracy and average generated length are reported.Source: §A.1.2 · pp. 17–18

A.2 Implementation Details

  • Training rollouts come from the proposal policy, while the teacher supplies token-level reference probabilities on top-16 support for distillation and target construction.Each round applies one student gradient step and one proposal inner step; baselines share the rollout backend, optimizer, chat template, data, and evaluation infrastructure, varying only method-specific objective terms and coefficients.Source: §A.2 · p. 18

B.1 A Closer Look at Accuracy and Generation Length

The appendix tests whether MPOPD’s math results depend on decoding budget, termination, proposal approximation, model scale, and teacher support size. Its analyses report gains at matched budgets, shorter post-answer generations, close proposal tracking, and improvements with a larger teacher and student.

B.1 A Closer Look at Accuracy and Generation Length

  • At the training-matched 4K budget, MPOPD gains 5.00 and 2.37 accuracy points over OPD and VETO, respectively, and outperforms both at 8K and 16K.The improvement therefore cannot be attributed solely to evaluation beyond the 4K training length.Source: §B.1 · p. 19, Table 6 · p. 19
  • 35.37% vs. 31.06% overall and 61.25% vs. 55.13% among terminated responses favor MPOPD over VETO, despite similar termination rates of 55.09% and 53.27%.OPD and MPOPD have similar answer presence and first-answer positions; OPD’s terminated-response accuracy is based on only 0.60% of its responses and is unstable.Source: §B.1 · pp. 19–20, Table 6 · p. 19
  • Only 0.01% of MPOPD responses terminate without a parseable answer; tokens after the first correct answer fall from 12.38K for OPD to 0.34K for MPOPD.At least 94.6% of the correct-response length reduction occurs after the first correct answer on every benchmark; 95.78% of correct OPD responses repeat the already scored boxed answer.Source: Table 6 · p. 19, §B.1 · p. 20

B.2 Proposal Approximation and Effective KL Radius

  • At the final step, full-vocabulary approximation KLs are below 0.01; positive rank correlations indicate proposal shifts generally track target shifts, not equal displacement magnitudes.The comparison is measured on teacher top-16 support and across prefixes; the exact geometric-mixture and trust-region interpretation applies to the ideal target.Source: §B.2 · pp. 20–21, Table 7 · p. 20
  • Most prefixes receive only a small shift under the global η, while a smaller right tail receives larger shifts, despite the induced prefix-dependent effective KL radius.Table 7(b) reports median, 70th, and 80th percentiles of KL(q∗∥π) on the full vocabulary at four training steps.Source: Table 7 · p. 20, §B.2 · p. 21

B.3 Scaling to a Larger Teacher and Student

  • 3.18 and 2.48 points: MPOPD improves average accuracy over OPD for Qwen3-4B-Base and Qwen3-8B-Base, respectively, with the larger teacher.MPOPD also exceeds VETO by 2.17 and 1.36 points, respectively.Source: §B.3 · p. 21, Figure 6 · p. 21

B.4 Impact of the Teacher Support Size

  • 1.23 lower average accuracy: reducing teacher support from k = 16 to k = 4 lowers performance, while k = 16, 32, and 128 span only 0.33 points.The scores for k = 16, 32, and 128 are 35.37, 35.33, and 35.66, respectively; gains are not monotonic across support sizes.Source: §B.4 · p. 22, Table 8 · p. 22

C Discussion

The discussion explains why the learned proposal generates from the tilted target rather than merely reweighting student samples, and outlines possible extensions and the study’s scope. The appendix derives the geometric-mixture form of the trust-region solution, while the broader-impact section discusses potential benefits and deployment risks.

C Discussion

  • Reweighting sampled student trajectories cannot recover low-probability prefixes that were never visited; the learned proposal changes rollout generation toward the tilted target.Clipped importance weights instead keep student updates aligned with the OPD objective on proposal-generated rollouts.Source: C Discussion · p. 23
  • The study uses a fixed tilt strength and reports Qwen3-family evaluations, leaving adaptive schedules, lighter proposal parameterizations, cross-family pairs, larger scales, and additional modalities for future work.A lightweight LoRA proposal is suggested to reduce training-time overhead at larger student scales.Source: C Discussion · p. 23

D Declaration of LLM Usage and Broader Impact

  • The paper frames MPOPD as potentially improving distillation budget efficiency and access to reasoning capabilities, while noting risks from over-reliance, incorrect outputs, and misuse in high-stakes settings.The authors recommend verification, monitoring, and domain-specific safeguards before deployment.Source: D Declaration of LLM Usage and Broader… · p. 24

E.1 Proof of Proposition 3.1

  • The geometric-mixture solution follows from a local trust-region problem: exponential tilting by the teacher–student score yields a normalized mixture of student and teacher conditionals.When the KL constraint is active, its multiplier sets the tilt strength.Source: §E.1 · p. 25

E.2 Proof of Corollary 3.1

The proof analyzes the local trust-region solution as the KL radius shrinks. Its expansion links the small-tilt behavior to the score variance, with a constant score yielding no first-order improvement.

E.2 Proof of Corollary 3.1

  • For a sufficiently small trust-region radius, the corollary gives a local improvement expansion for the tilted distribution, whose first-order behavior is governed by score variance.When the score is constant on the student’s support, the first-order improvement is zero.Source: §E.2 · p. 26

F Optimization Analysis under Structured Teacher Target Noise

The analysis models teacher-target noise as conditionally unbiased, with its second moment controlled by teacher reachability risk. It derives a proposal-induced one-step descent bound and identifies conditions under which tilted rollouts improve that bound.

F Optimization Analysis under Structured Teacher Target Noise

  • The analysis assumes teacher-target perturbations are conditionally unbiased and that their conditional second moment is controlled by reachability risk.Higher risk is assigned to prefixes that are less reachable under the teacher, for example via a nondecreasing function of teacher prefix negative log-likelihood.Source: F Optimization Analysis under… · p. 27
  • The one-step result bounds descent for the clean objective induced by each rollout occupancy, combining a clean-gradient signal with a penalty tied to reachability-risk exposure.The comparison is between rollout-induced objectives, not a guarantee on a fixed evaluation objective.Source: F Optimization Analysis under… · pp. 27–28
  • Tilting has a no-smaller proposal-induced lower bound when it raises the clean descent signal and/or reduces reachability-risk exposure without an excessive second-moment penalty.The geometric tilt does not guarantee lower risk; the local diagnostic is that greater relative teacher support tends to lead to lower-risk next prefixes.Source: F Optimization Analysis under… · p. 28

G Geometric and Arithmetic Interpolation

The appendix distinguishes geometric and arithmetic interpolation by the variational objectives they encode: the geometric path fits a student-anchored reverse-KL trust region, while the arithmetic path suits coverage-oriented objectives. Case studies illustrate qualitative differences in answer quality and output length.

G Geometric and Arithmetic Interpolation

  • Arithmetic interpolation provides bounded importance ratios useful for student-coverage settings, whereas the paper’s student-anchored reverse-KL formulation yields geometric interpolation.Neither path is an unconditional winner: reducing reachability-risk exposure with the geometric path depends on a covariance condition.Source: G Geometric and Arithmetic Interpolation · pp. 29–30
  • The geometric path is the reverse-KL variational solution: it starts from the student conditional and multiplicatively reweights tokens by relative teacher support.The arithmetic path instead combines student and teacher conditionals additively and corresponds to a different proposal objective.Source: G Geometric and Arithmetic Interpolation · pp. 29–30

H More Case Studies

  • In the Minerva Math examples, MPOPD preserves a parseable final answer after correct reasoning and maintains a quark-count invariant that the four shown baselines lose.Source: Figure 7 · p. 31, Figure 8 · p. 32
  • In the OlympiadBench examples, MPOPD uses a short symbolic path and preserves the domain constraint; shown baselines fail to produce parsed answers, and one exhausts a 16.4k-token budget after answering.The examples are OlympiadBench ids 2548 and 2783.Source: Figure 9 · p. 33, Figure 10 · p. 34

Related papers

Loading 2610.36477v1…