Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

On the Position Bias of On-Policy Distillation

About

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.

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

Related benchmarks

TaskDatasetResultRank
Mathematical ReasoningHMMT 2025--
241
Mathematical ReasoningAIME 25
Avg@320.497
50
Code GenerationMBPP+
Accuracy (%)70.9
38
Mathematical ReasoningHMMT25
Avg@320.287
34
Code GenerationHE+
Accuracy78.7
32
Mathematical Problem SolvingAIME 24
Mean@32 Accuracy58.7
16
Mathematical Problem SolvingAIME25
Mean@32 Accuracy46.7
16
Mathematical Problem SolvingHMMT25
Mean@32 Accuracy25
16
Mathematical ReasoningAIME 24
Mean@32 Accuracy57.5
16
Multi-task Model EvaluationAIME24, AIME25, HMMT25, HE+, MBPP+ Average
Average Accuracy55.3
16
Showing 10 of 11 rows

Other info

Follow for update