Elamur Research⌕
☰
home / papers / 2026 / Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes

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

Paper · 2026 · Yuqian Fu, Haohuan Huang, Kaiwen Jiang et al.

Summary

On-policy distillation (OPD) trains a student LLM on its own rollouts while a stronger teacher scores each token. Most pipelines implement it by comparing teacher and student log-probabilities only on the token the student actually sampled. The authors examine why this signal becomes unstable on long generations.

Theoretically, they show that the token-level estimator is biased compared with sequence-level reverse-KL. However, its worst-case variance grows much more slowly with length. A toy two-task experiment confirms that coupling updates to future rewards inflates gradient variance and causes policy drift.

Empirically, with Qwen2.5-7B-Instruct as student and OpenThinker3-7B as teacher, they find three problems:

  • the reward is skewed toward penalties,
  • teacher scores become misleading once the student's prefix drifts,
  • tokenizer and special-token mismatches distort the comparison.

Their fix replaces the single-token signal with a renormalized, truncated reverse-KL over the teacher's top-K tokens. It is combined with top-p rollout sampling and special-token masking. On math reasoning and mixed ALFWorld-plus-math training, the recipe yields steadier optimization and about a 19.8% gain over sampled-token OPD. This makes it a cheap, practical upgrade for post-training pipelines.

Key points

  • Worst-case variance bound: token-level OPD scales as O(T²), while the sequence-level estimator scales as O(T⁴). A discounted γ-estimator interpolates between the two, and larger γ produces higher gradient variance.
  • Single-task math ablation (AIME24 avg@32):
    • sampled-token OPD: 20.4
    • sampled-token OPD + top-p: 21.6
    • teacher top-K alone: 17.7
    • teacher top-K + top-p: 23.6
    • for reference, the student scores 10.0 and the teacher 63.3
  • Multi-task training (ALFWorld + MATH500, AIME24, AIME25, Minerva, OlympiadBench):
    • the average reasoning score rises from 34.8 (sampled-token OPD) to 41.7 (ours without mask)
    • the masked variant reaches 97.7 on ALFWorld
  • Renormalization within the truncated support is essential: without it, training collapses quickly. Very small K values or unconstrained rollouts (top_p=1.0) also destabilize learning.
  • Building the support from the teacher's top-K works best in multi-task training. Student top-K and teacher top-K plus the sampled token are roughly equal in the single-task setting but degrade sharply on math in the multi-task setting.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🏢Artificial Intelligence in MedicineToronto · Canada3 works · 170 citations

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 worksArts and Humanities · 17,959 worksLanguage and Linguistics · 17,789 works

🧭 Topics 1Research subjects and languages the work is about.Efficient models & compression · 328 works

Cite

Cite as Yuqian Fu, Haohuan Huang, Kaiwen Jiang, Jiacai Liu, Zhuo Jiang, Yuanheng Zhu and Dongbin Zhao. 2026. Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes. arXiv:2603.25562.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2603.25562 · arXiv 2603.25562

@misc{fu2026revisiting,
  title = {{Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes}},
  author = {Yuqian Fu, Haohuan Huang, Kaiwen Jiang, Jiacai Liu, Zhuo Jiang, Yuanheng Zhu and Dongbin Zhao},
  year = {2026},
  doi = {10.48550/arxiv.2603.25562},
  eprint = {2603.25562},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2603.25562}
}

References

47 references, 27 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos Garea, Matthieu Geist, and Olivier Bachem. On-policy distillation of language models: Learning from self-generated mistakes. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.… ↗
  2. Nicolas Boizard, Kevin El Haddad, CELINE HUDELOT, and Pierre Colombo. Towards cross-tokenizer distillation: the universal logit distillation loss for LLMs. Transactions on Machine Learning Research, 2025. ISSN 2835-8856. URL https://openreview.net/forum?id=bwRxXiGO9A. ↗
  3. Howard Chen, Noam Razin, Karthik Narasimhan, and Danqi Chen. Retaining by doing: The role of on-policy data in mitigating forgetting. arXiv preprint arXiv:2510.18874, 2025. find
  4. DeepSeek-AI. Deepseek-v4: Towards highly efficient million-token context intelligence, 2026.
  5. Lang Feng, Zhenghai Xue, Tingcong Liu, and Bo An. Group-in-group policy optimization for LLM agent training. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. URL https://openreview.net/forum?id=QXEhBMNrCW. ↗
  6. Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang. MiniLLM: Knowledge distillation of large language models. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.net/forum?id=5h0qf7IBZZ. ↗
  7. Etash Guha, Ryan Marten, Sedrick Keh, Negin Raoof, Georgios Smyrnis, Hritik Bansal, Marianna Nezhurina, Jean Mercat, Trung Vu, Zayne Sprague, et al. OpenThoughts: Data recipes for reasoning models. arXiv preprint arXiv:2506.04178, 2025.
  8. Chaoqun He, Renjie Luo, Yuzhuo Bai, Shengding Hu, Zhen Thai, Junhao Shen, Jinyi Hu, Xu Han, Yujie Huang, Yuxiang Zhang, et al. OlympiadBench: A challenging benchmark for promoting AGI with olympiad-level bilingual multimodal scientific problems. In Proceedings of the 62nd Annual Meeting of the Assoc…
  9. Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the MATH dataset. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021. URL ht…
  10. Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531, 2015.
  11. Jonas Hübotter, Frederike Lübeck, Lejs Behric, Anton Baumann, Marco Bagatella, Daniel Marta, Ido Hakimi, Idan Shenfeld, Thomas Kleine Buening, Carlos Guestrin, et al. Reinforcement learning via self-distillation. arXiv preprint arXiv:2601.20802, 2026.
  12. Woogyeol Jin, Taywon Min, Yongjin Yang, Swanand Ravindra Kadhe, Yi Zhou, Dennis Wei, Nathalie Baracaldo, and Kimin Lee. Entropy-aware on-policy distillation of language models. In The 1st Workshop on Scaling Post-training for LLMs, 2026a. URL https://openreview.net/forum?id=WSRQ37tzk1. ↗
  13. Woogyeol Jin, Taywon Min, Yongjin Yang, Swanand Ravindra Kadhe, Yi Zhou, Dennis Wei, Nathalie Baracaldo, and Kimin Lee. Entropy-aware on-policy distillation of language models. arXiv preprint arXiv:2603.07079, 2026b.
  14. Yoon Kim and Alexander M Rush. Sequence-level knowledge distillation. In Proceedings of the 2016 conference on empirical methods in natural language processing, pp. 1317–1327, 2016.
  15. Jongwoo Ko, Sara Abdali, Young Jin Kim, Tianyi Chen, and Pashmina Cameron. Scaling reasoning efficiently via relaxed on-policy distillation. arXiv preprint arXiv:2603.11137, 2026.
  16. Aitor Lewkowycz, Anders Andreassen, David Dohan, Ethan Dyer, Henryk Michalewski, Vinay Ramasesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman-Solo, et al. Solving quantitative reasoning problems with language models. Advances in neural information processing systems, 35:3843–3857, 2022.
  17. Yaxuan Li, Yuxin Zuo, Bingxiang He, Jinqian Zhang, Chaojun Xiao, Cheng Qian, Tianyu Yu, Huan-ang Gao, Wenkai Yang, Zhiyuan Liu, et al. Rethinking on-policy distillation of large language models: Phenomenology, mechanism, and recipe. arXiv preprint arXiv:2604.13016, 2026.
  18. Jiacai Liu, Yingru Li, Yuqian Fu, Jiawei Wang, Qian Liu, and Zhuo Jiang. When speed kills stability: Demystifying RL collapse from the training-inference mismatch, September 2025. URL https://richardli.xyz/rl-collapse. ↗
  19. Kevin Lu and Thinking Machines Lab. On-policy distillation. Thinking Machines Lab: Connectionism, 2025. doi: 10.64434/tml.20251026. https://thinkingmachines.ai/blog/on-policy-distillation. — [doi:10.64434/tml.20251026](https://doi.org/10.64434/tml.20251026) ↗
  20. Benjamin Minixhofer, Ivan Vulić, and Edoardo Ponti. Universal cross-tokenizer distillation via approximate likelihood matching. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. URL https://openreview.net/forum?id=DxKP2E0xK2. ↗
  21. Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730–27744, 2022.
  22. Carlos Miguel Patiño, Kashif Rasul, Quentin Gallouédec, Ben Burtenshaw, Sergio Paniego, Vaibhav Srivastav, Thibaud Frere, Ed Beeching, Lewis Tunstall, Leandro von Werra, and Thomas Wolf. Unlocking on-policy distillation for any model family, 2025. find
  23. Qwen, An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, Huan Lin, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Yang, Jiaxi Yang, Jingren Zhou, Junyang Lin, Kai Dang, Keming Lu, Keqin Bao, Kexin Yang, Le Yu, Mei Li, Mingfeng Xu… find
  24. Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01108, 2019.
  25. John Schulman. Approximating KL divergence. John Schulman’s Homepage, 5, 2020. find
  26. Guangming Sheng, Chi Zhang, Zilingfeng Ye, Xibin Wu, Wang Zhang, Ru Zhang, Yanghua Peng, Haibin Lin, and Chuan Wu. HybridFlow: A flexible and efficient RLHF framework. In Proceedings of the Twentieth European Conference on Computer Systems, EuroSys ’25, pp. 1279–1297. ACM, March 2025. doi: 10.1145/3…
  27. Mohit Shridhar, Xingdi Yuan, Marc-Alexandre Cote, Yonatan Bisk, Adam Trischler, and Matthew Hausknecht. ALFWorld: Aligning text and embodied environments for interactive learning. In International Conference on Learning Representations, 2021. URL https://openreview.net/forum?id=0IOX0YcCdTn. ↗
  28. Haoqing Wang, Xiang Long, Ziheng Li, Yilong Xu, Tingguang Li, and Yehui Tang. To mix or to merge: Toward multi-domain reinforcement learning for large language models. arXiv preprint arXiv:2602.12566, 2026a. — [arXiv:2602.12566](https://arxiv.org/abs/2602.12566) find
  29. Yinjie Wang, Xuyang Chen, Xiaolong Jin, Mengdi Wang, and Ling Yang. Openclaw-RL: Train any agent simply by talking. arXiv preprint arXiv:2603.10165, 2026b. — [arXiv:2603.10165](https://arxiv.org/abs/2603.10165) find
  30. Bangjun Xiao, Bingquan Xia, Bo Yang, Bofei Gao, Bowen Shen, Chen Zhang, Chenhong He, Chiheng Lou, Fuli Luo, Gang Wang, et al. MiMo-V2-Flash technical report. arXiv preprint arXiv:2601.02780, 2026.
  31. An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint arXiv:2505.09388, 2025.
  32. Chenxu Yang, Chuanyu Qin, Qingyi Si, Minghui Chen, Naibin Gu, Dingyu Yao, Zheng Lin, Weiping Wang, Jiaqi Wang, and Nan Duan. Self-distilled RLVR. arXiv preprint arXiv:2604.03128, 2026a.
  33. Wenkai Yang, Weijie Liu, Ruobing Xie, Kai Yang, Saiyong Yang, and Yankai Lin. Learning beyond teacher: Generalized on-policy distillation with reward extrapolation. arXiv preprint arXiv:2602.12125, 2026b. — [arXiv:2602.12125](https://arxiv.org/abs/2602.12125) find
  34. Feng Yao, Liyuan Liu, Dinghuai Zhang, Chengyu Dong, Jingbo Shang, and Jianfeng Gao. Your efficient RL framework secretly brings you off-policy RL training, August 2025. URL https://fengyao.notion.site/off-policy-rl. ↗
  35. Shunyu Yao, Howard Chen, John Yang, and Karthik Narasimhan. WebShop: Towards scalable real-world web interaction with grounded language agents. Advances in Neural Information Processing Systems, 35:20744–20757, 2022.
  36. Chenlu Ye, Xuanchang Zhang, Yifan Hao, Zhou Yu, Ziji Zhang, Abhinav Gullapalli, Hao Chen, Jing Huang, and Tong Zhang. Adaptive layerwise perturbation: Unifying off-policy corrections for LLM RL. arXiv preprint arXiv:2603.19470, 2026a. — [arXiv:2603.19470](https://arxiv.org/abs/2603.19470) find
  37. Tianzhu Ye, Li Dong, Zewen Chi, Xun Wu, Shaohan Huang, and Furu Wei. Black-box on-policy distillation of large language models. arXiv preprint arXiv:2511.10643, 2025. — [arXiv:2511.10643](https://arxiv.org/abs/2511.10643) find
  38. Tianzhu Ye, Li Dong, Xun Wu, Shaohan Huang, and Furu Wei. On-policy context distillation for language models. arXiv preprint arXiv:2602.12275, 2026b.
  39. Qiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan, Xiaochen Zuo, Yu Yue, Weinan Dai, Tiantian Fan, Gaohong Liu, Juncai Liu, Lingjun Liu, Xin Liu, Haibin Lin, Zhiqi Lin, Bole Ma, Guangming Sheng, Yuxuan Tong, Chi Zhang, Mofan Zhang, Ru Zhang, Wang Zhang, Hang Zhu, Jinhua Zhu, Jiaze Chen, Jiangjie Chen,… — DAPO: An Open-Source LLM Reinforcement Learning System at Scale
  40. Aohan Zeng, Xin Lv, Zhenyu Hou, Zhengxiao Du, Qinkai Zheng, Bin Chen, Da Yin, Chendi Ge, Chengxing Xie, Cunxiang Wang, et al. GLM-5: from vibe coding to agentic engineering. arXiv preprint arXiv:2602.15763, 2026.
  41. Juzheng Zhang, Abhimanyu Hans, John Kirchenbauer, Micah Goldblum, Ashwinee Panda, and Tom Goldstein. Learning from mixed rollouts: Logit fusion as a bridge between imitation and exploration. Notion Blog, 2026. URL https://juzhengz.notion.site/logit-fusion. ↗
  42. Lunjun Zhang and Jimmy Ba. EMA policy gradient: Taming reinforcement learning for LLMs with EMA anchor and Top-k KL. arXiv preprint arXiv:2602.04417, 2026. find
  43. Yifan Zhang and Team Math-AI. American invitational mathematics examination (AIME) 2024. find
  44. Wei Zhao, Zhe Li, Yige Li, Ye Zhang, and Junfeng Sun, 2024. find
  45. Yifan Zhang and Team Math-AI. American invitational mathematics examination (AIME) 2025. find
  46. Wei Zhao, Zhe Li, Yige Li, Ye Zhang, and Junfeng Sun, 2025. find
  47. Siyan Zhao, Zhihui Xie, Mengchen Liu, Jing Huang, Guan Pang, Feiyu Chen, and Aditya Grover. Self-distilled reasoner: On-policy self-distillation for large language models. arXiv preprint arXiv:2601.18734, 2026.

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/fu-2026-revisiting-policy-distillation-empirical-failure/