Technical Summary: On the Position Bias of On-Policy Distillation

Problem Statement

On-Policy Distillation (OPD) enhances the learning efficiency of Reinforcement Learning (RL) for Large Language Models (LLMs) by providing dense, token-level supervision from a teacher model to a student model generating its own rollouts. However, standard OPD objectives uniformly average token-level losses (typically via KL divergence), implicitly assuming equal value for all tokens in a sequence.

The authors identify a critical Position Bias in this approach: as student rollouts extend, the generated trajectory increasingly deviates from the teacher's distribution. Consequently, the quality of teacher supervision degrades at later positions. Empirical observations reveal that:

1. Asymmetry in Supervision Value: Training using only the first 30% of tokens (prefix) yields performance comparable to or better than using all tokens, whereas training on the last 30% (suffix) provides negligible learning benefits.
2. Persistent Divergence: Even after convergence, the KL divergence between the teacher and student remains significant, suggesting that OPD only optimizes the student within a small local region of the policy space.
3. Inefficiency: Uniformly weighting all tokens wastes the limited gradient update budget on later tokens where the student has already drifted into Out-of-Distribution (OOD) regions relative to the teacher, rendering the supervision signal noisy or unhelpful.

Methodology: Importance-Weighted On-Policy Distillation (IW-OPD)

To address this bias, the authors propose Importance-Weighted On-Policy Distillation (IW-OPD), a method that dynamically adjusts token weights based on the accumulated discrepancy between the student and teacher distributions.

Theoretical Foundation

The authors frame OPD as a constrained optimization problem (a local projection) where the student policy is updated within a trust region defined by a KL divergence budget $\rho$.

• Optimal Policy Derivation: Solving the constrained minimization of $D_{KL}(q \| \pi_T)$ subject to $D_{KL}(q \| \pi_\theta) \leq \rho$ yields a closed-form optimal policy $q^\star$.
• Likelihood Ratio Reweighting: The solution $q^\star$ reweights the student policy $\pi_\theta$ by the likelihood ratio $(\pi_T / \pi_\theta)^\alpha$. This implies that tokens where the teacher's probability is significantly higher than the student's (high compatibility) should receive higher weight, while tokens where the student has drifted (low ratio) should be downweighted.

Practical Implementation

Directly using the raw likelihood ratio is numerically unstable due to the exponential accumulation of probabilities over long sequences. IW-OPD introduces a stable, unsigned cumulative prefix discrepancy as a proxy for the importance weight:

1. Accumulated Discrepancy: Instead of the raw ratio, the method sums the absolute token-level advantages ($|A^{OPD}_k|$) from the start of the sequence up to the current token $t$. This captures the total "drift" from the teacher's preferred path.
2. Normalization: The accumulated discrepancy is standardized within each sample to a range $[0, 1]$, ensuring stable scaling.
3. Interpolation: To prevent the complete removal of supervision for drifted tokens, the final weight is an interpolation between the standard OPD signal (weight 1) and the importance weight derived from the discrepancy.
   • Formula: $\tilde{r}^{IW-OPD}_t = 1 + \gamma \cdot (1 - \frac{\sum_{k<t} |A^{OPD}_k|}{\sum_{k<T} |A^{OPD}_k|})$.
   • Here, $\gamma$ controls the contribution of the importance weighting.

This approach naturally upweights early tokens (where the prefix is compatible with the teacher) and downweights later tokens (where the accumulated drift is large), effectively reallocating the gradient budget to the most informative parts of the trajectory.

Key Contributions

1. Identification of Position Bias: The paper empirically demonstrates and theoretically explains that OPD suffers from a position bias where early tokens provide significantly more valuable supervision than later tokens due to distributional drift.
2. Constrained Optimization Perspective: The authors provide a principled understanding of OPD as a constrained local projection, deriving that the optimal update strategy inherently requires ratio-based reweighting.
3. IW-OPD Algorithm: They propose IW-OPD, a practical, stable algorithm that estimates token importance using unsigned cumulative prefix discrepancy. It requires no additional teacher evaluations beyond standard OPD.
4. Empirical Validation: Extensive experiments show that IW-OPD converges faster and achieves higher final performance than standard OPD across various model scales and teacher-student configurations.

Experimental Results

The authors evaluated IW-OPD using Qwen3 models (0.6B, 1.7B, 4B, 30B, and 235B) as teachers and students on mathematical reasoning (AIME-2024, AIME-2025, HMMT-2025) and code generation (HumanEval+, MBPP+) benchmarks.

• Performance Gains: IW-OPD consistently outperforms standard OPD.
  • In the Qwen3-1.7B student setting with a Qwen3-30B teacher, IW-OPD improved AIME-2025 accuracy by +6.9 points at step 10 and +1.7 points at convergence.
  • Across all student scales (0.6B to 30B) and teacher settings, IW-OPD improved the average performance score.
• Sample Efficiency: IW-OPD converges significantly faster. In several cases, the IW-OPD checkpoint at step 10 matched or exceeded the final performance of standard OPD.
• Scaling Effects:
  • Stronger Teachers: IW-OPD makes stronger teachers more sample-efficient. While standard OPD struggles initially with larger teacher-student gaps (requiring more steps to benefit from the stronger teacher), IW-OPD immediately leverages the stronger teacher's signal.
  • Smaller Students: The relative gains of IW-OPD increase as the student model becomes smaller (e.g., +9.5% relative gain for the 0.6B student vs. +1.8% for the 4B student), indicating it is particularly effective in cross-scale distillation where drift is more severe.

Significance and Claims

The paper claims that effective on-policy supervision must account for trajectory context, not just token-level disagreement. By recognizing that teacher supervision is not uniformly reliable along a student-generated trajectory, IW-OPD reallocates the finite gradient budget to "teacher-compatible" prefixes.

The authors conclude that this simple modification—downweighting tokens where the student has drifted from the teacher's distribution—significantly improves both the sample efficiency and the final performance of on-policy distillation. This insight suggests that future on-policy methods should consider the cumulative history of the rollout (the prefix) when assigning credit or supervision weights, rather than treating all tokens in a trajectory as independent and equally valuable.