Elamur Research⌕
☰
home / papers / 2026 / On-Policy Context Distillation for Language Models

On-Policy Context Distillation for Language Models

Paper · 2026 · Tianzhu Ye, Li Dong, Xun Wu et al.

Summary

Prompts can supply a language model with instructions or lessons learned, but that knowledge vanishes when the context is cleared. Classic context distillation tries to bake such context into the weights. It trains on a fixed set of teacher outputs with forward KL, which causes a train/inference mismatch (exposure bias) and spreads probability too broadly. The authors from Microsoft Research propose On-Policy Context Distillation (OPCD). A student without the context generates its own responses. At every token, it is pushed toward a teacher that sees the context, by minimising reverse KL, approximated over the student's top-k tokens. They test two uses. The first is experiential knowledge distillation: the model writes reusable lessons from its own unlabeled solution traces, accumulates them, and then absorbs them. The second is system prompt distillation, for medical QA and safety classification. On math, text-based games and these domain tasks, OPCD beats off-policy baselines and forgets less on out-of-distribution tasks. It also lets smaller models learn from a larger teacher's experience. This makes it relevant for cheaply turning accumulated prompt knowledge into persistent model behaviour.

Key points

  • On math validation, Qwen3-8B scores 75.1 with no experience and 70.5 with raw traces in context. Extracted knowledge raises it to 77.4, and adding OPCD reaches 79.7.
  • In safety system prompt distillation (student Qwen2.5-3B-Instruct, frozen teacher Qwen2.5-7B-Instruct), OPCD raises in-distribution accuracy. It also keeps medical (OOD) accuracy near the initial student, about 4 points above off-policy context distillation.
  • On Frozen Lake, OPCD scores roughly 2% higher on IF-Eval than the context distillation baseline.
  • A frozen-teacher setup beats self-distillation (shared, simultaneously updated weights) and trains more stably:
    • Sokoban with Qwen3-4B-Instruct-2507: 53.9 vs 18.8.
    • Medical prompt with Qwen2.5-3B-Instruct: 56.8 vs 50.0.
  • Smaller Qwen3 students improve via OPCD from a frozen Qwen3-8B teacher. By contrast, putting the teacher's knowledge directly into their context hurts them. Code is available at aka.ms/opcd-code.

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

🏢Microsoft Research (United Kingdom)Cambridge · United Kingdom15 works · 3,537 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 2Research subjects and languages the work is about.Language models & pretraining · 4,145 worksEfficient models & compression · 328 works

Cite

Cite as Tianzhu Ye, Li Dong, Xun Wu, Shaohan Huang and Furu Wei. 2026. On-Policy Context Distillation for Language Models. arXiv:2602.12275.

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

@misc{ye2026policy,
  title = {{On-Policy Context Distillation for Language Models}},
  author = {Tianzhu Ye, Li Dong, Xun Wu, Shaohan Huang and Furu Wei},
  year = {2026},
  doi = {10.48550/arxiv.2602.12275},
  eprint = {2602.12275},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2602.12275}
}

References

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

  1. Amanda Askell, Yuntao Bai, Anna Chen, Dawn Drain, Deep Ganguli, T. J. Henighan, Andy Jones, Nicholas Joseph, Benjamin Mann, Nova Dassarma, Nelson Elhage, Zac Hatfield-Dodds, Danny Hernandez, John Kernion, Kamal Ndousse, Catherine Olsson, Dario Amodei, Tom B. Brown, Jack Clark, Sam McCandlish, Chris… find
  2. 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.
  3. Francesco Barbieri, Jose Camacho-Collados, Luis Espinosa Anke, and Leonardo Neves. Tweeteval: Unified benchmark and comparative evaluation for tweet classification. In Findings of the association for computational linguistics: EMNLP 2020, pages 1644–1650, 2020.
  4. Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, et al. Language models are few-shot learners. In Proceedings of NeurIPS, 2020.
  5. Yumin Choi, Jinheon Baek, and Sung Ju Hwang. System prompt optimization with meta-learning. arXiv preprint arXiv:2505.09666, 2025. Model Experience Type Accuracy Qwen3-8B w/o Experience 75.1 Qwen3-8B Qwen3-8B Raw Trace Knowledge + OPCD 70.5 77.4 79.7 Table 6: Using raw response traces from previous… — [arXiv:2505.09666](https://arxiv.org/abs/2505.09666) find
  6. Bowen Cao, Deng Cai, and Wai Lam. Infiniteicl: Breaking the limit of context window size via long short-term memory transformation. In Findings of the Association for Computational Linguistics: ACL 2025, pages 11402–11415, 2025. find
  7. 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
  8. Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Jingyuan Ma, Rui Li, Heming Xia, Jingjing Xu, Zhiyong Wu, Baobao Chang, Xu Sun, Lei Li, and Zhifang Sui. A survey on in-context learning. In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen, editors, Proceedings of the 2024 Conference on Empirical Methods…
  9. Leon Guertler, Bobby Cheng, Simon Yu, Bo Liu, Leshem Choshen, and Cheston Tan. Textarena. arXiv preprint arXiv:2504.11442, 2025. — [arXiv:2504.11442](https://arxiv.org/abs/2504.11442) find
  10. Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783, 2024.
  11. Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang. MiniLLM: On-policy distillation of large language models. In The Twelfth International Conference on Learning Representations, 2024.
  12. Jonas Hübotter, Frederike Lübeck, Lejs Behric, Anton Baumann, Marco Bagatella, Daniel Marta, Ido Hakimi, Idan Shenfeld, Thomas Kleine Buening, Carlos Guestrin, and Andreas Krause. Reinforcement learning via self-distillation, 2026.
  13. Kevin Lu and Thinking Machines Lab. On-policy distillation. Thinking Machines Lab: Connectionism, 2025. https://thinkingmachines.ai/blog/on-policy-distillation. ↗
  14. Ioannis Mollas, Zoe Chrysopoulou, Stamatis Karlos, and Grigorios Tsoumakas. Ethos: a multi-label hate speech detection dataset. Complex & Intelligent Systems, 8(6):4663–4678, 2022. find
  15. Ankit Pal, Logesh Kumar Umapathi, and Malaikannan Sankarasubbu. Medmcqa: A large-scale multi-subject multi-choice dataset for medical domain question answering. In Conference on health, inference, and learning, pages 248–260. PMLR, 2022. — MedMCQA : A Large-scale Multi-Subject Multi-Choice Dataset for Medical domain Question Answering
  16. Emiliano Penaloza, Dheeraj Vattikonda, Nicolas Gontier, Alexandre Lacoste, Laurent Charlin, and Massimo Caccia. Privileged information distillation for language models, 2026. find
  17. Paul Röttger, Bertie Vidgen, Dong Nguyen, Zeerak Talat, Helen Margetts, and Janet Pierrehumbert. Hatecheck: Functional tests for hate speech detection models. In Proceedings of the 59th annual meeting of the association for computational linguistics and the 11th international joint conference on nat…
  18. Idan Shenfeld, Mehul Damani, Jonas Hübotter, and Pulkit Agrawal. Self-distillation enables continual learning, 2026.
  19. Charlie Snell, Dan Klein, and Ruiqi Zhong. Learning by distilling context, 2022.
  20. Idan Shenfeld, Jyothish Pari, and Pulkit Agrawal. Rl’s razor: Why online reinforcement learning forgets less. arXiv preprint arXiv:2509.04259, 2025. — [arXiv:2509.04259](https://arxiv.org/abs/2509.04259) find
  21. Sai Wang, Yu Wu, and Zhongwen Xu. Cogito, ergo ludo: An agent that learns to play by reasoning and planning. arXiv preprint arXiv:2509.25052, 2025. — [arXiv:2509.25052](https://arxiv.org/abs/2509.25052) find
  22. Tianzhu Ye, Li Dong, Zewen Chi, Xun Wu, Shaohan Huang, and Furu Wei. Black-box on-policy distillation of large language models, 2026. find
  23. 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.
  24. An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Li Chengyuan, Dayiheng Liu, Fei Huang, Haoran Wei, et al. Qwen2.5 Technical Report. arXiv preprint arXiv:2412.15115, 2025. find
  25. Qiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan, Xiaochen Zuo, Yu Yue, Weinan Dai, Tiantian Fan, Gaohong Liu, Lingjun Liu, et al. Dapo: An open-source llm reinforcement learning system at scale. arXiv preprint arXiv:2503.14476, 2025.
  26. Jeffrey Zhou, Tianjian Lu, Swaroop Mishra, Siddhartha Brahma, Sujoy Basu, Yi Luan, Denny Zhou, and Le Hou. Instruction-following evaluation for large language models. arXiv preprint arXiv:2311.07911, 2023.
  27. Eric Zelikman, Yuhuai Wu, Jesse Mu, and Noah Goodman. STar: Bootstrapping reasoning with reasoning. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
  28. 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, 2026.

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/ye-2026-policy-context-distillation-language-models/