Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes

Yuqian Fu; Haohuan Huang; Kaiwen Jiang; Jiacai Liu; Zhuo Jiang; Yuanheng Zhu; Dongbin Zhao

cs.LGMar 26, 2026

Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes

Authors: Yuqian Fu, Haohuan Huang, Kaiwen Jiang, Jiacai Liu, Zhuo Jiang, Yuanheng Zhu, Dongbin Zhao

Organizations: State Key Laboratory of Multimodal Artificial Intelligence Systems, CASIA · School of Artificial Intelligence, UCAS · Fudan University · Independent Researcher

Abstract

On-policy distillation (OPD) is increasingly used in LLM post-training because it can leverage a teacher model to provide dense supervision on student rollouts. The standard implementation, however, usually reduces distribution matching to a sampled-token log-ratio, which can make the learning signal fragile on long rollouts whose prefixes drift away from the teacher's typical support. We revisit this formulation from both theoretical and implementation perspectives. Theoretically, token-level OPD is biased relative to sequence-level reverse-KL minimization, but admits a substantially tighter worst-case variance bound; a controlled synthetic study further shows that stronger future-reward coupling increases gradient variance and destabilizes training. Empirically, we identify three failure modes of sampled-token OPD: imbalanced token-level supervision, unreliable teacher guidance on student-generated prefixes, and tokenizer or special-token mismatch. These findings motivate teacher top-K local support matching, a truncated reverse-KL objective that compares teacher and student distributions over a teacher-supported token set at each prefix, together with top-p rollout sampling and special-token masking. Across single-task reasoning and multi-task benchmarks spanning agentic and reasoning settings, this objective improves optimization stability and yields a +19.8% performance gain over standard sampled-token OPD baselines, providing a practical recipe for more stable on-policy distillation.

Figures & tables

Panel of Figure 1 — Figure 1: Effect of increasing $\gamma$ in the toy experiment. Larger $\gamma$ yields a higher and more persistent variance regime and, in the sequence-level limit, drifting policies in state space.

Figure 2: Scatter of token probabilities (teacher vs. student) at the sampled-token OPD first iteration. The sampled-token signal is heavily skewed toward penalizing the current token.

Figure 3: Example of unreliable teacher guidance. The student falls into a repetition loop, yet the teacher remains locally aligned with the student on repeated tokens, indicating that sampled-token OPD may fail to penalize this behavior.

Figure 4: Distribution of teacher–student log-probability gaps across token positions. Several later position buckets exhibit wider lower tails and more extreme values, indicating a noisier teacher signal on long-horizon student-generated rollouts.

Figure 5: Token-level comparison can penalize semantically correct outputs due to tokenizer mismatch between the teacher and student.

Method	Math500	AIME24	AIME25	Minerva	OlympiadBench	Avg.
Qwen2.5-7B-It	68.2	13.3	0.0	26.5	32.9	28.2
OpenThinker3-7B	92.2	53.3	40.0	39.0	55.6	56.0
Sampled-token OPD	80.0	10.0	16.7	32.4	43.1	36.4
Sampled-token OPD w/ mask	81.4	26.7	16.7	34.2	44.7	40.7
Ours w/o mask	80.4	23.3	26.7	34.2	43.9	41.7
Ours w/ mask	82.0	23.3	23.3	34.9	43.9	41.5

Table 1: Results on single-task math reasoning.

Method	Agentic	Reasoning
Method	ALFWorld	MATH500	AIME24	AIME25	Minerva	OlympiadBench	Avg.
Qwen2.5-7B-It	21.9	68.2	13.3	0.0	26.5	32.9	28.2
GiGPO-Qwen2.5-7B-It-Alfworld	95.3	–	–	–	–	–	–
OpenThinker3-7B	–	92.2	53.3	40.0	39.0	55.6	56.0
Sampled-token OPD	90.6	74.8	13.3	13.3	32.1	40.5	34.8
Sampled-token OPD w/ mask	93.8	76.0	20.0	13.3	33.5	40.4	36.6

Table 2: Results on batch-alternating multi-task training over ALFWorld and math reasoning.

Method	AIME24 avg@32
Qwen2.5-7B-Instruct	10.0
OpenThinker3-7B	63.3
Sampled-token OPD	20.4
Sampled-token OPD + top- $p$	21.6
Teacher top- $K$	17.7
Teacher top- $K$ + top- $p$	23.6

Table 3: Component ablation under single-task math training.

Panel of Figure 6 — Figure 6: Ablations of the main design choices. Renormalization is required for stability, very small support sets hurt learning, and unconstrained rollout sampling degrades optimization.

Table 4: Alternative top-

K

support variants under single-task and multi-task settings.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 11 — Figure A1: Toy experiment with random seed 42: gradient variance and state visitation.

Panel of Figure 12 — Figure A2: Toy experiment with random seed 43: gradient variance and state visitation.

Panel of Figure 13 — Figure A3: Toy experiment with random seed 2026: gradient variance and state visitation.

Category	Item	Math ( $\lx@sectionsign$ 4.2 , $\lx@sectionsign$ 4.3 )	ALFWorld ( $\lx@sectionsign$ 4.3 )	WebShop ( $\lx@sectionsign$ G.1 )
Model	Student	Qwen2.5-7B-Instruct	Qwen2.5-7B-Instruct	Qwen2.5-1.5B-Instruct
Model	Teacher	OpenThinker3-7B	GiGPO-Qwen2.5-7B-It-ALFWorld	GiGPO-Qwen2.5-1.5B-It-WebShop
Environment	Max prompt length	2048	2048	2048
	Max response length	16384	512	512
	Max turns	–	30	15
Rollout	Group size	8	8	8

Appendix

Table A1: Training and evaluation settings. * Rollout top-

p

is used in local support matching, but not in the sampled-token OPD baseline. † Training steps denote the total number of training updates; in the alternating multi-task setting, 400 total steps correspond to 200 math updates and 200 ALFWorld updates.

Method	Task score	Success rate
Qwen2.5-1.5B-It	12.7	2.3
GiGPO-Qwen2.5-1.5B-Webshop ( Feng et al., 2025 )	81.9	66.4
Sampled-token OPD	73.0	50.0
Ours	75.1	57.8

Appendix

Table A2: Results on single-task WebShop with Qwen2.5-1.5B-Instruct as the student. Local support matching improves over sampled-token OPD on both task score and success rate, reducing the gap to the teacher. We do not include a masking variant here because the teacher is obtained by RL training from the same base model, making tokenizer-mismatch effects less central in this setting.

Figure A4: Single-task training curves for math reasoning. Local support matching improves both training reward and evaluation performance over the course of training.

Figure A5: Multi-task learning curves for ALFWorld and math reasoning. The main gains appear on the math side while agentic performance remains strong.

Panel of Figure 18 — Figure A6: Optimization and alignment diagnostics. Relative to sampled-token OPD, local support matching yields smaller gradient norms, fewer clipping-boundary hits, shorter responses, and teacher–student log-probability gaps closer to zero.

Table A3: Ablation of five top-

K

variants under single-task and multi-task settings.

Panel of Figure 20 — Figure A7: Even after the student has effectively reached an answer, the teacher can still assign high conditional probability to low-information continuations.

Panel of Figure 21 — Figure A8: Loops, overlong chain-of-thought traces, and punctuation-heavy continuations can remain locally rewardable.

Figure A9: The teacher still assigns high probability to several tokens after the student drifts into malformed non-English output.

Language Model Distillation

Language model distillation transfers a larger model’s capabilities to smaller or specialized models using outputs, probabilities, or internal representations.

See all 489 papers

Unbiased Top- $k$ Estimation for On-Policy Distillation
Sep 28, 2026Linjian Meng, Siyuan Gan, YuHan Li +5Language Model Distillation On-Policy Distillation
Shanghai Artificial Intelligence Laboratory, Shanghai, China · State Key Laboratory of Novel Software Technology, Nanjing University, Nanjing, China
Demystifying On-Policy Distillation: Roles, Pathologies, and Regulations
Jul 15, 2026Rui Wang, Hongru Wang, Yi Chen +4RL for Language Model Reasoning Language Model Distillation
†The Chinese University of Hong Kong · ‡Tencent AI Lab
The Many Faces of On-Policy Distillation: Pitfalls, Mechanisms, and Fixes
May 11, 2026Siqi Zhu, Xuyan Ye, Hongyu Lu +2On-Policy Self-Distillation On-Policy Distillation
1UIUC · 2Renmin University of China · †Work done during an internship at UIUC. +1

Sep 28, 2026J/K move · Enter open · S save

Unbiased Top- $k$ Estimation for On-Policy Distillation

Linjian Meng, Siyuan Gan, YuHan Li, Xiran Wang+4

Shanghai Artificial Intelligence Laboratory, Shanghai, China · State Key Laboratory of Novel Software Technology, Nanjing University, Nanjing, China

Language Model Distillation On-Policy Distillation

Open paper AskPDF arXiv

On-policy distillation (OPD) is becoming an important component of large language model (LLM) post-training for transferring the reasoning capability of a strong teacher LLM to a weaker student LLM. OPD trains the student by minimizing the reverse KL divergence between the teacher and the student via rollouts generated by the student's policy. However, estimating the gradient of the reverse KL divergence in OPD remains a challenge. Using only the sampled token from the student-generated rollout is computationally cheap but provides limited distributional supervision, which will degrade accuracy. In addition, using the full vocabulary provides complete distributional supervision but is computationally expensive. Therefore, recent works propose Top-

k

OPD (TK-OPD) that use selected top-

k

tokens, which provides richer distributional supervision than sampled-token estimation at substantially lower computational cost than full-vocabulary estimation. Unfortunately, using only the selected top-

k

tokens induces bias, leading to accuracy degradation, as the probability mass outside the selected top-

k

tokens is discarded. To address the bias of TK-OPD, we propose Tail-Corrected Top-

k

On-Policy Distillation (TT-OPD). It preserves the advantages of TK-OPD, including rich distributional supervision and low computational cost, while providing an unbiased estimator of the gradient of the reverse KL divergence. The key insight of TT-OPD is to use not only the selected top-

k

tokens, but also the sampled token from the student-generated rollout, thereby recovering the discarded probability mass in expectation, avoiding the bias. Experimental results demonstrate that TT-OPD significantly outperforms other tested OPD variants.

Abstract

Language Model Distillation

Explore similar work

Unbiased Top-kkk Estimation for On-Policy Distillation

Demystifying On-Policy Distillation: Roles, Pathologies, and Regulations

The Many Faces of On-Policy Distillation: Pitfalls, Mechanisms, and Fixes

Unbiased Top- $k$ Estimation for On-Policy Distillation