PaperlayerView exact revision on arXiv

Source-linked AI summary

On the Position Bias of On-Policy Distillation

Yan Xie, Sijie Zhu, Tiansheng Wen, Bo Chen, Yifei Wang

arXiv:2606.22600v3cs.LGcs.AI

TL;DR

OPD’s uniform token weighting overlooks the declining supervision quality of later prefixes as student rollouts drift from the teacher distribution. The paper derives IW-OPD from constrained optimization, using accumulated prefix discrepancy to reweight tokens, and reports faster convergence and better final performance than standard OPD across same-size and cross-scale settings. The main scope caveat is that larger-scale validation remains future work.

  • Problem

    Uniformly averaging OPD’s token-level KL terms overlooks position-dependent supervision quality as later student prefixes drift from the teacher distribution.

  • Method

    IW-OPD derives token importance from a constrained local projection and uses cumulative prefix discrepancy to reallocate update weight toward teacher-compatible prefixes.

  • Results

    IW-OPD consistently improves OPD’s sample efficiency and final performance across evaluated teacher–student pairs and benchmarks, including a 1.9-point math-average gain in a larger distillation setting.

  • Takeaways & Limitations

    Effective on-policy supervision should account for trajectory context, not only token-level disagreement, while standard OPD remains the dense supervision floor.

  • Takeaways & Limitations

    Validation at larger student scales remains future work, and the cumulative prefix discrepancy is only a conservative compatibility proxy rather than an exact density-ratio correction.

Abstract

from arXiv · show

On-Policy Distillation (OPD) improves the learning efficiency of standard reinforcement learning through dense, token-level supervision from teachers. In the standard KL objective of OPD, token-level losses are uniformly averaged, implying equal weights for all tokens. However, we discover that not all tokens are created equal: as student rollouts grow longer, they deviate further from the teacher's distribution, leading to degraded supervision quality at later positions. As a result, OPD using only the first 30% of tokens can perform comparably to using all tokens, whereas OPD using only the last 30% of tokens barely learns anything. In this work, we provide a principled understanding of this issue through the lens of constrained optimization. Based on these insights, we derive Importance-Weighted On-Policy Distillation (IW-OPD), in which the weight assigned to each token depends on the accumulated discrepancy between the student's and teacher's distributions, naturally upweighting earlier tokens and downweighting later ones with larger deviations. We show that IW-OPD converges significantly faster than OPD, with better learning efficiency, and achieves better final performance than standard OPD in both same-size and cross-scale settings, improving performance up to 6.9 points on AIME-2025.

1 Introduction

OPD improves learning efficiency with dense token-level teacher supervision, but uniformly weighting tokens overlooks the declining quality of supervision along weak-student rollouts. The paper explains this position bias through constrained optimization and introduces IW-OPD to reallocate supervision toward teacher-compatible prefixes.

  • OPD replaces sparse trajectory-level rewards with dense token-level supervision from a stronger teacher on student-generated prefixes.This substantially improves learning efficiency over sparse trajectory-level rewards in LLM post-training.
  • Later student-rollout prefixes drift farther from the teacher distribution, making early tokens more valuable and later tokens less reliable for supervision.The standard objective uniformly aggregates token-level KL terms despite this position-dependent supervision quality.
  • The paper frames OPD as a finite-budget allocation problem and derives optimal weights from the teacher-to-student likelihood ratio.The constrained local projection allocates limited update budget toward prefixes where teacher supervision remains compatible with the student trajectory.
  • IW-OPD consistently improves OPD, with stronger teachers becoming more sample-efficient and smaller students obtaining larger gains.The paper reports faster convergence and stronger final performance across evaluated settings.
  • IW-OPD uses token-level importance estimated from teacher–student discrepancy to correct OPD’s position bias.Its contributions include identifying the phenomenon, deriving the objective, and evaluating its effects across teacher–student settings.

2 Preliminaries

This section contrasts sparse on-policy reinforcement learning with OPD’s dense token-level teacher supervision. OPD samples responses from the student and applies token-local updates using teacher–student distribution gaps.

  • On-Policy RL: On-policy RL samples trajectories from the current policy and optimizes a trajectory-level reward.The reward may come from a reward model or verifier, while trajectory-level advantages assign the same credit to every token.
  • On-Policy Distillation: OPD replaces sparse trajectory-level rewards with dense token-level supervision from a teacher model.This provides supervision at the prefixes visited by student-generated responses.
  • On-Policy Distillation: OPD samples responses from the student and decomposes its sequence-level objective into token-local semi-gradient updates.The sampled prefixes are treated as fixed during the practical update.
  • On-Policy Distillation: The token-level update uses an advantage derived from the teacher–student distribution gap.This contrasts with trajectory-level RL, where the same trajectory reward advantage is assigned across tokens.

3 Position Bias in On-Policy Distillation

OPD exhibits a position bias: early-token supervision is more useful than later-token supervision because student trajectories increasingly diverge from teacher-preferred paths. A finite-budget constrained-optimization view explains this pattern through likelihood-ratio reweighting.

  • 3.1 The Position Bias Phenomenon in OPD: Prefix-30 supervision performs comparably to standard OPD, whereas suffix-30 supervision produces substantially lower rewards throughout training.The controlled comparison fixes the supervision budget and varies only the supervised response segment.
  • 3.1 The Position Bias Phenomenon in OPD: OPD’s mean token-level KL decreases but plateaus at a non-zero residual, indicating that teacher–student divergence persists after training.The figure also compares token-level reverse KL before and after OPD training.
  • 3.1 The Position Bias Phenomenon in OPD: The teacher–student divergence decreases by only 20% after OPD converges, while student performance saturates.The paper interprets this as OPD optimizing the student distribution within a small local region.
  • 3.2 Understanding Position Bias from a Finite-Budget Allocation Perspective: OPD can therefore be modeled as constrained optimization in which the student remains within a local KL-defined update region.This finite update budget motivates viewing supervision as an allocation problem.
  • 3.2 Understanding Position Bias from a Finite-Budget Allocation Perspective: The constrained optimum reweights the student policy by the teacher-to-student likelihood ratio.The normalizing factor and exponent are determined by the local-update constraint.
  • 3.2 Understanding Position Bias from a Finite-Budget Allocation Perspective: Prefix likelihood ratios tend to shrink at later positions when early student decisions move reasoning away from the teacher-preferred region.This explains why later rollout tokens receive less useful supervision under standard OPD.

4 Importance-Weighted On-Policy Distillation

IW-OPD reallocates OPD’s finite update budget using prefix-level importance weights derived from teacher–student discrepancy. The method stabilizes these weights and improves learning efficiency across student–teacher scales.

  • Importance-weighted objective: IW-OPD reweights OPD token losses using teacher–student likelihood-ratio information, emphasizing prefixes that provide meaningful learning signals.The objective is derived by projecting toward the teacher under a constrained local-update budget and then changing measure back to student-sampled rollouts.
  • Importance-weighted objective: The detached normalized prefix weight acts multiplicatively on the OPD policy-gradient signal, reallocating finite update budget across positions.This operational form is explicitly identified as the correction for OPD’s position bias.
  • Stable token-level importance estimates: Longer prefixes make raw likelihood-ratio weights vary dramatically, introducing severe training instability through exponentially accumulated probability products.The method therefore uses stabilization strategies rather than directly applying the raw ratio.
  • Stable token-level importance estimates: Small α values only roughly match standard OPD, while α = 1 and α = 0.1 substantially degrade training because α controls weight concentration and flatness.The reported stable settings are α = 0.01 and α = 0.001.
  • Stable token-level importance estimates: Log scaling, sign correction, within-sample normalization, and interpolation with OPD stabilize the importance signal and reduce cancellation or extreme-scale effects.The unsigned variant produces better results than the original signed version in the reported comparison, and γ = 0.5 is a practical default.

5 Experiments

Experiments across teacher scales, student sizes, benchmarks, and ablations show that IW-OPD improves OPD’s efficiency and final performance by adaptively emphasizing teacher-compatible prefixes.

  • Main results: IW-OPD improves final average performance over standard OPD across evaluated teacher–student pairs and benchmarks.With a Qwen3-235B-A22B teacher and Qwen3-30B-A3B student, it improves the math average by 1.9 points.
  • Sample efficiency: At step 10, IW-OPD improves the average score from 52.0 to 55.5 and AIME25 from 42.4 to 49.3 in the Qwen3-30B-A3B → Qwen3-4B setting.The step-10 checkpoint already matches the final OPD checkpoint in average performance.
  • Teacher scaling: IW-OPD makes stronger teachers more sample-efficient despite larger teacher–student trajectory mismatch.For the Qwen3-4B student, IW-OPD10 with the 30B-A3B teacher reaches 55.5 average and 49.3 AIME25, versus 55.1 and 46.7 with the 4B teacher.
  • Student scaling: With a Qwen3-4B teacher, IW-OPD’s final improvement over OPD grows from +1.0 points for 4B students to +1.9 points for 0.6B students.Relative gains increase from +1.8% to +9.5% as student size decreases.
  • Stronger-teacher scaling: With a Qwen3-235B-A22B teacher and Qwen3-30B-A3B student, IW-OPD gains 1.3 points on AIME24, 2.2 points on AIME25, and 2.1 points on HMMT25.The gains show usefulness beyond weak-student regimes.
  • Ablations: The cumulative-share ablation gives +5.6, outperforming fixed-ratio, hand-designed, and linear-decay prefix schedules.The result indicates that trajectory adaptivity matters beyond simply preferring earlier positions.
  • Ablations: Unsigned discrepancy gives +5.6 versus +2.6 for signed accumulation because signed terms can cancel across model disagreements.The absolute statistic treats each disagreement as evidence of departure from the shared prefix region.
  • Ablations: The practical surrogate should blend with standard OPD: it preserves dense supervision while allocating extra budget to compatible prefixes.Using the ideal likelihood-ratio weight alone collapses performance, whereas the blended surrogate is more robust.

6 Related Work

Related work situates IW-OPD among on-policy distillation, selective supervision, credit assignment, and reweighted policy-update methods. Its distinction is allocating dense teacher supervision by prefix compatibility.

  • On-policy and token-selective distillation: Classical KD trains on teacher-generated data, whereas OPD supervises student-sampled rollouts with dense teacher feedback.Prior OPD methods include GKD, MiniLLM, and extensions studying on-policy teacher supervision.
  • On-policy and token-selective distillation: Selective distillation methods choose where to supervise using sequence-level curricula or token-level weights based on frequency, difficulty, teacher confidence, and related criteria.IW-OPD instead uses prefix compatibility as its allocation principle.
  • Credit assignment and reweighted policy updates: RLVR assigns sparse outcome rewards over long traces, while process supervision and process rewards provide denser step-level feedback.Recent token-level analyses identify high-entropy, critical, and reasoning tokens as disproportionate learning drivers.

7 Discussion and Conclusion

The paper frames OPD’s position bias as a finite-budget local-projection problem and concludes that supervision should account for trajectory context. Its derivation uses a constrained KL projection whose solution interpolates between student and teacher policies.

  • Discussion: OPD supervision is often more useful near the beginning of student rollouts than near the end, creating a position bias.The teacher–student distribution gap is not known in advance, so compatible trajectory regions cannot be predefined.
  • Conclusion: IW-OPD reallocates gradient budget toward teacher-compatible prefixes using stable cumulative prefix discrepancy while retaining standard OPD as a dense supervision floor.Experiments show improved sample efficiency and final performance across same-family, cross-scale, and stronger-teacher settings.
  • Constrained projection: The constrained projection assumes distributions have common support where the relevant KL divergences are finite.The appendix denotes the local projected distribution by q⋆ and the causal prefix weight by rθ.
  • Constrained projection: In the local-update regime, the trust-region constraint is active and the constrained problem has a unique optimum because the KL objective is strictly convex over a convex KL ball.If the teacher is feasible, the constraint can become inactive and recover the teacher policy.
  • Constrained projection: The derivation forms a Lagrangian with a trust-region multiplier λ and normalization multiplier µ, then sets the functional derivative with respect to q(y) to zero.The resulting stationarity equations yield the projected policy’s closed form.
  • Constrained projection: The stationarity condition gives log q(y) as a weighted combination of log πT(y) and log πθ(y), with coefficients determined by λ.Writing α = 1/(1+λ) makes the interpolation coefficient explicit.
  • Constrained projection: The interpolation coefficient α is induced by the trust-region radius ρ, increasing with ρ when teacher and student policies differ.For 0 < ρ < DKL(πT ∥ πθ), α lies between 0 and 1.

A.2 Derivation of the Importance-Weighted OPD Objective

IW-OPD derives a causal token-level weighting from constrained local projection toward the teacher. The resulting normalized prefix ratio is detached and multiplies the standard OPD token update, with interpolation controlled by γ.

  • Under a fixed local-update budget, the projected-KL problem yields a likelihood-ratio weighting of student trajectories.The derivation treats OPD as allocating limited update capacity toward teacher-compatible prefixes.
  • The trajectory-level ratio must be replaced by a causal prefix ratio because token-local OPD updates depend only on the sampled prefix h_t.Future suffix tokens should not determine the coefficient for position t.
  • The token-level IW-OPD surrogate multiplies each standard log-ratio term by a normalized prefix ratio.The normalizer is position-specific and the prefix ratio is detached during the local update.
  • Stop-gradient prevents prefix-ratio and normalizer derivatives from introducing sequence-level credit-assignment terms.Detachment preserves the single-step token-local semi-gradient used by OPD.
  • IW-OPD is implemented in the same PPO pipeline as OPD, using student-sampled responses and directly supplied token-level advantages.The experiments use matched data order and seed sets, while γ=0.5 is the practical default.

C.4 Evaluation Protocol

The evaluation uses standardized math and code benchmarks with fixed sampling and validation procedures. Reported comparisons preserve matched seeds, training protocols, and model–teacher settings.

  • Math evaluation samples 32 responses per problem and scores extracted boxed answers with symbolic equivalence checking.Sampling uses temperature 1.0, top-p 1.0, maximum length 16384, and seed matching across methods.
  • Code evaluation uses EvalPlus on HumanEval+ and MBPP+ with greedy single-sample pass-rate measurement.
  • Converged checkpoints are selected with the same validation protocol within each model–teacher setting before held-out math and code evaluation.Checkpoints with subscript 10 denote training step 10.

D Combination experiments with other reward design methods

IW-OPD can be combined with alternative reward designs because it reallocates an existing OPD signal rather than depending on a particular reward formulation. The IW-ExOPD combination improves over ExOPD for both evaluated student sizes.

  • IW-OPD changes how an already-computed OPD signal is allocated across positions, making it orthogonal to reward design.The paper applies prefix-dependent importance weighting to ExOPD and calls the combination IW-ExOPD.
  • IW-ExOPD improves over ExOPD for both Qwen3-4B and Qwen3-1.7B students distilled from Qwen3-30B-A3B.

E Limitations

The analysis has a stated scope boundary: experiments are conducted at the 4B student scale, and validation at larger scales remains future work.

  • Experiments are conducted at the 4B student scale, so validation at larger student scales remains future work.The normalized prefix weight also retains mild residual nonuniformity at convergence and is described as a conservative compatibility proxy.
Loading 2606.22600v3…