On the Position Bias of On-Policy Distillation

On-Policy Distillation (OPD) is a popular method for improving reinforcement learning in large language models by using a teacher model to provide dense, token-level supervision. However, standard OPD treats all tokens in a generated sequence as equally important. This paper investigates whether this uniform weighting is optimal, given that student-generated rollouts often drift away from the teacher's distribution as the sequence length increases.

The authors identify a "position bias" in OPD, where early tokens in a student's response are significantly more valuable for learning than later ones. They formalize this as a constrained optimization problem, where the student's policy is updated within a local trust region. By solving this, they derive an optimal policy that reweights tokens based on the likelihood ratio between the teacher and the student. To implement this, they introduce Importance-Weighted On-Policy Distillation (IW-OPD), which uses a normalized, accumulated discrepancy between teacher and student probabilities to dynamically downweight later tokens in a sequence.

• Position Bias: Empirical evidence shows that training on only the first 30% of tokens in a rollout performs as well as or better than training on the full sequence, while training on the last 30% yields minimal learning.
• Performance Gains: IW-OPD consistently outperforms standard OPD in both convergence speed and final accuracy. On the AIME-2025 benchmark, IW-OPD achieved performance improvements of up to 6.9 points at early training stages.
• Scalability: The advantage of IW-OPD increases as the teacher-student mismatch grows, making it particularly effective for distilling knowledge from stronger teachers into smaller student models.