Fast and Effective On-Policy Distillation from Reasoning Prefixes
Paper · 2026 · Dongxu Zhang, Zhichao Yang, Sepehr Janghorbani et al.
Summary
On-policy distillation (OPD) trains a student LLM on its own sampled outputs, using token-level feedback from a teacher. This avoids the exposure bias of SFT on teacher-written data. However, OPD becomes very costly when reasoning traces run to tens of thousands of tokens. The authors ran two diagnostic studies. First, during full OPD, the reverse-KL loss is concentrated in the early part of each trajectory. Second, an untrained student solves noticeably more AIME-24 problems when it continues from a short teacher-written prefix.
Based on these findings, they propose prefix OPD. Student generation stops after a token budget, and only that truncated segment is distilled. An optional linear schedule gradually lengthens the budget during training. At inference, the length cap is removed.
They distilled Qwen3-8B into Qwen3-1.7B-Base and Qwen3-8B-Base using OpenThoughts3 questions. Prefix OPD performed about as well as full OPD on math and out-of-domain reasoning benchmarks while using 2×–40× fewer training FLOPs. This makes on-policy distillation practical for long chain-of-thought models.
Key points
- Setup: The teacher is Qwen3-8B, and the students are Qwen3-1.7B-Base and Qwen3-8B-Base. Training questions come from OpenThoughts3. The SeqKD baseline uses QwQ-32B responses from OpenThoughts3.
- Objective and stability fix: Training uses a sampled reverse-KL estimate with importance sampling. The
<think>token is forced at the end of the prompt, because learning special tokens through sampled KL produced unreadable output. - Prefix scheduling: The prefix starts at 1 token and grows by ΔL = 256 tokens after each gradient step. At test time, generation is allowed up to 16,384 tokens.
- Teacher-prefix probe: Prefixes of 2048–8192 teacher tokens raised untrained students' AIME-24 accuracy. Most of the gain came from the first few thousand tokens.
- Learning-rate study (Qwen3-8B-Base, prefix 1024): AIME-24 rose from 5.2% to about 25–26%. MMLU-Pro rose from 26.5% to as high as 55.0%, and GPQA from 17.6% to 27.9%. Lowering the learning rate did not improve out-of-domain generalization.
Fields, topics and concepts
📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 works
🧭 Topics 2Research subjects and languages the work is about.Reasoning · 1,426 worksEfficient models & compression · 242 works
Cite
Cite as Dongxu Zhang, Zhichao Yang, Sepehr Janghorbani, Jun Han, Andrew Ressler II, Qian Qian, Gregory D Lyng, Sanjit Singh Batra and Robert E. Tillman. 2026. Fast and Effective On-Policy Distillation from Reasoning Prefixes. Findings of the Association for Computational Linguistics: ACL 2026.
@inproceedings{zhang2026fast,
title = {{Fast and Effective On-Policy Distillation from Reasoning Prefixes}},
author = {Dongxu Zhang, Zhichao Yang, Sepehr Janghorbani, Jun Han, Andrew Ressler II, Qian Qian, Gregory D Lyng, Sanjit Singh Batra and Robert E. Tillman},
booktitle = {Findings of the Association for Computational Linguistics: ACL 2026},
year = {2026},
doi = {10.18653/v1/2026.findings-acl.1276},
url = {https://doi.org/10.18653/v1/2026.findings-acl.1276}
}References
- 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. Cristian Buciluǎ, Rich Carua… find
- Tianzhe Chu, Yuexiang Zhai, Jihan Yang, Shengbang Tong, Saining Xie, Dale Schuurmans, Quoc V Le, Sergey Levine, and Yi Ma. 2025. SFT memorizes, RL generalizes: A comparative study of foundation model post-training. In Forty-second International Conference on Machine Learning.
- Mengnan Du, Fengxiang He, Na Zou, Dacheng Tao, and Xia Hu. 2023. Shortcut learning of large language models in natural language understanding. Communications of the ACM, 67(1):110–120. find
- Robert Geirhos, Jörn-Henrik Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix A Wichmann. 2020. Shortcut learning in deep neural networks. Nature Machine Intelligence, 2(11):665–673. find
- Zorik Gekhman, Gal Yona, Roee Aharoni, Matan Eyal, Amir Feder, Roi Reichart, and Jonathan Herzig. 2024. Does fine-tuning LLMs on new knowledge encourage hallucinations? In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, Miami, Florida, USA. Association for Com… find
- Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang. 2024. MiniLLM: Knowledge distillation of large language models. In The Twelfth International Conference on Learning Representations. find
- Etash Kumar Guha, Ryan Marten, Sedrick Keh, Negin Raoof, Georgios Smyrnis, Hritik Bansal, Marianna Nezhurina, Jean Mercat, Trung Vu, Zayne Rea Sprague, Ashima Suvarna, Benjamin Feuer, Leon Liangyu Chen, Zaid Khan, Eric Frankel, Sachin Grover, Caroline Choi, Niklas Muennighoff, Shiye Su, and 32 other… — OpenThoughts: Data Recipes for Reasoning Models
- Horace He and Thinking Machines Lab. 2025. Defeating nondeterminism in llm inference. Thinking Machines Lab: Connectionism. Https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/. find
- Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. 2021. Measuring mathematical problem solving with the math dataset. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2).
- Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. 2015. Distilling the knowledge in a neural network. Preprint, arXiv:1503.02531.
- Haiduo Huang, Jiangcheng Song, Yadong Zhang, and Pengju Ren. 2025. Selectkd: Selective token-weighted knowledge distillation for llms. arXiv preprint arXiv:2510.24021. — [arXiv:2510.24021](https://arxiv.org/abs/2510.24021) find
- Ke Ji, Jiahao Xu, Tian Liang, Qiuzhi Liu, Zhiwei He, Xiaoyuan Liu, Xingyu Chen, Junying Chen, Benyou Wang, Zhaopeng Tu, Haitao Mi, and Dong Yu. 2025. find
- The first few tokens are all you need: An efficient and effective unsupervised prefix fine-tuning method for reasoning models. In The Thirty-ninth Annual Conference on Neural Information Processing Systems. find
- Adam Tauman Kalai, Ofir Nachum, Santosh S Vempala, and Edwin Zhang. 2025. Why language models hallucinate. arXiv preprint arXiv:2509.04664. — [arXiv:2509.04664](https://arxiv.org/abs/2509.04664) find
- Damjan Kalajdzievski. 2024. Scaling laws for forgetting when fine-tuning large language models. arXiv preprint arXiv:2401.05605. — [arXiv:2401.05605](https://arxiv.org/abs/2401.05605) find
- Yoon Kim and Alexander M Rush. 2016. Sequence-level knowledge distillation. In Proceedings of the 2016 conference on empirical methods in natural language processing, pages 1317–1327.
- Jongwoo Ko, Tianyi Chen, Sungnyun Kim, Tianyu Ding, Luming Liang, Ilya Zharkov, and Se-Young Yun. 2025. DistiLLM-2: A contrastive approach boosts the distillation of LLMs. In Forty-second International Conference on Machine Learning.
- Jongwoo Ko, Sungnyun Kim, Tianyi Chen, and Se-Young Yun. 2024. DistiLLM: Towards streamlined distillation for large language models. In Forty-first International Conference on Machine Learning.
- Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. 2023. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th symposium on operating systems principles, pages 611–626.
- Nathan Lambert, Jacob Morrison, Valentina Pyatkin, Shengyi Huang, Hamish Ivison, Faeze Brahman, Lester James V. Miranda, Alisa Liu, Nouha Dziri, Shane Lyu, Yuling Gu, Saumya Malik, Victoria Graf, Jena D. Hwang, Jiangjiang Yang, Ronan Le Bras, Oyvind Tafjord, Chris Wilhelm, Luca Soldaini, and 4 other… — Tulu 3: Pushing Frontiers in Open Language Model Post-Training
- Alexander Lin, Jeremy Wohlwend, Howard Chen, and Tao Lei. 2020. Autoregressive knowledge distillation through imitation learning. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 6121–6133, Online. Association for Computational Linguistics. find
- Jiacheng Lin, Zhongruo Wang, Kun Qian, Tian Wang, Arvind Srinivasan, Hansi Zeng, Ruochen Jiao, Xie Zhou, Jiri Gesi, Dakuo Wang, Yufan Guo, Kai Zhong, Weiqi Zhang, sujay sanghavi, Changyou Chen, Hyokun Yun, and Lihong Li. 2026. SFT doesn’t always hurt general capabilities: Revisiting domain-specific… find
- Zhenghao Lin, Zhibin Gou, Yeyun Gong, Xiao Liu, yelong shen, Ruochen Xu, Chen Lin, Yujiu Yang, Jian Jiao, Nan Duan, and Weizhu Chen. 2024. Not all tokens are what you need for pretraining. In The Thirty-eighth Annual Conference on Neural Information Processing Systems. find
- Kevin Lu and Thinking Machines Lab. 2025. On-policy distillation. Thinking Machines Lab: Connectionism. Https://thinkingmachines.ai/blog/on-policy-distillation. find
- Yun Luo, Zhen Yang, Fandong Meng, Yafu Li, Jie Zhou, and Yue Zhang. 2025. An empirical study of catastrophic forgetting in large language models during continual fine-tuning. IEEE Transactions on Audio, Speech and Language Processing. find
- Hamed Mahdavi, Alireza Hashemi, Majid Daliri, Pegah Mohammadipour, Alireza Farhadi, Samira Malek, Yekta Yazdanifard, Amir Khasahmadi, and Vasant G Honavar. 2025. Brains vs. bytes: Evaluating LLM proficiency in olympiad mathematics. In Second Conference on Language Modeling. find
- Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, and 1 others. 2022. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730–27744.
- Aldo Pareja, Nikhil Shivakumar Nayak, Hao Wang, Krishnateja Killamsetty, Shivchander Sudalairaj, Wenlong Zhao, Seungwook Han, Abhishek Bhandwaldar, Guangxuan Xu, Kai Xu, Ligong Han, Luke Inglis, and Akash Srivastava. 2025. Unveiling the secret recipe: A guide for supervised fine-tuning small LLMs. I… find
- Team Qwen. 2024. Qwq: Reflect deeply on the boundaries of the unknown, november 2024. URL https://qwenlm.github. io/blog/qwq-32b-preview. ↗
- Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, and 1 others. 2019. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9. find
- David Rein, Betty Li Hou, Asa Cooper Stickland, Jackson Petty, Richard Yuanzhe Pang, Julien Dirani, Julian Michael, and Samuel R Bowman. 2024. Gpqa: A graduate-level google-proof q&a benchmark. In First Conference on Language Modeling.
- Gavin A Rummery and Mahesan Niranjan. 1994. Online Q-learning using connectionist systems, volume 37. University of Cambridge, Department of Engineering Cambridge, UK. find
- Guangming Sheng, Chi Zhang, Zilingfeng Ye, Xibin Wu, Wang Zhang, Ru Zhang, Yanghua Peng, Haibin Lin, and Chuan Wu. 2025. Hybridflow: A flexible and efficient rlhf framework. In Proceedings of the Twentieth European Conference on Computer Systems, EuroSys ’25, page 1279–1297, New York, NY, USA. Assoc…
- Shenzhi Wang, Le Yu, Chang Gao, Chujie Zheng, Shixuan Liu, Rui Lu, Kai Dang, Xiong-Hui Chen, Jianxin Yang, Zhenru Zhang, Yuqiong Liu, An Yang, Andrew Zhao, Yang Yue, Shiji Song, Bowen Yu, Gao Huang, and Junyang Lin. 2025. Beyond the 80/20 rule: High-entropy minority tokens drive effective reinforcem…
- Yubo Wang, Xueguang Ma, Ge Zhang, Yuansheng Ni, Abhranil Chandra, Shiguang Guo, Weiming Ren, Aaran Arulraj, Xuan He, Ziyan Jiang, and 1 others. 2024. Mmlu-pro: A more robust and challenging multi-task language understanding benchmark. Advances in Neural Information Processing Systems, 37:95266–95290…
- Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M. Dai, and Quoc V Le. 2022. Finetuned language models are zero-shot learners. In International Conference on Learning Representations.
- Ronald J Williams. 1992. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning, 8(3):229–256. find
- Xurong Xie, Zhucun Xue, Jiafu Wu, Jian Li, Yabiao Wang, Xiaobin Hu, Yong Liu, and Jiangning Zhang. 2026. Llm-oriented token-adaptive knowledge distillation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 40, pages 34070– 34078.
- An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, and 1 others. 2025. Qwen3 technical report. arXiv preprint arXiv:2505.09388.
Updated
This page in Markdown → https://md.research.elamur.ai/free/en/papers/zhang-2026-fast-effective-policy-distillation-reasoning/