Elamur Research⌕
☰
home / papers / 2026 / Reinforcement-aware Knowledge Distillation for LLM Reasoning

Reinforcement-aware Knowledge Distillation for LLM Reasoning

Paper · 2026 · Zhaoyang Zhang, Shuli Jiang, Yantao Shen et al.

Summary

Long chain-of-thought reasoning models trained with RL are expensive to run, so practitioners distill them into smaller students. Most distillation recipes were built for SFT. They either imitate fixed teacher traces or add a teacher–student KL penalty. Inside an RL loop, both approaches cause problems. The student's rollouts drift away from what the teacher would produce, and the KL term competes with reward maximization, so it needs careful weighting.

The authors (AWS Agentic AI, ICML 2026) propose RLAD, which only pulls the student toward the teacher when that helps the current policy update. Its core, Trust Region Ratio Distillation (TRRD), drops the separate KL loss. Instead it uses a clipped PPO/GRPO-style likelihood ratio whose anchor mixes the teacher with the old student policy. The result is distillation on student rollouts that is weighted by advantage and bounded by a trust region.

On K&K Logistics and on long-context math benchmarks, RLAD beats GRPO, KDRL and offline distillation. The largest gains appear on the hardest problems, and training is more stable. This makes RLAD a practical option for building compact reasoning models.

Key points

  • On K&K Logistics with a Qwen3-8B teacher, RLAD raises average accuracy as follows:
    • Qwen3-0.6B: 0.76 → 0.94 at 8K context, and 0.70 → 0.90 at 2K.
    • Qwen3-1.7B: 0.95 → 0.99 at 8K, and 0.86 → 0.93 at 2K.
    • The clearest margins over KDRL are on the hardest subsets, PPL7 and PPL8.
  • On long-context math with Qwen3-8B/Qwen3-32B teachers, RLAD beats standard GRPO on the average across five test sets and nine metrics:
    • +2.5 for Qwen3-1.7B-Base and +5.5 for Qwen3-8B-Base at 30K context.
    • +6.6 and +2.6 for the two post-trained models under an 8K response budget.
    • Gains are strongest on AIME24/25 under Pass@32 and small on MATH500.
  • KDRL's improvements look driven by imitation: it gains more on Pass@32 than on Pass@1, and its validation curves oscillate. RLAD gives steadier training and larger Pass@1 gains. It is also less sensitive to its mixture coefficient α than KDRL is to its KL weight λ.
  • Cost was measured with a Qwen3-8B-Base student and a Qwen3-32B teacher on 32 H200 GPUs. RLAD adds only about 12% training latency, because the teacher only scores student rollouts and does not generate new ones. No teacher is needed at inference time.
  • The main limitation is that RLAD needs access to teacher logits. It therefore suits open-weight teachers, and extending it to closed-source teachers is left for future work.

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

 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 3Research subjects and languages the work is about.Language models & pretraining · 4,145 worksReasoning · 1,948 worksEfficient models & compression · 328 works

Cite

Cite as Zhaoyang Zhang, Shuli Jiang, Yantao Shen, Yuting Zhang, Dhananjay Ram, Shuo Yang, Zhuowen Tu, Wei Xia and Stefano Soatto. 2026. Reinforcement-aware Knowledge Distillation for LLM Reasoning. arXiv:2602.22495.

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

@misc{zhang2026reinforcement,
  title = {{Reinforcement-aware Knowledge Distillation for LLM Reasoning}},
  author = {Zhaoyang Zhang, Shuli Jiang, Yantao Shen, Yuting Zhang, Dhananjay Ram, Shuo Yang, Zhuowen Tu, Wei Xia and Stefano Soatto},
  year = {2026},
  doi = {10.48550/arxiv.2602.22495},
  eprint = {2602.22495},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2602.22495}
}

References

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

  1. Agarwal, R., Vieillard, N., Zhou, Y., Stanczyk, P., Garea, S. R., Geist, M., and Bachem, O. On-policy distillation of language models: Learning from self-generated mistakes, 2024a. In The Twelfth International Conference on Learning Representations, 2024.
  2. Hinton, G. E., Vinyals, O., and Dean, J. Distilling the knowledge in a neural network. ArXiv, abs/1503.02531, 2015. URL https://api.semanticscholar.org/CorpusID:7200347. ↗
  3. Agarwal, R., Vieillard, N., Zhou, Y., Stanczyk, P., Garea, S. R., Geist, M., and Bachem, O. On-policy distillation of language models: Learning from self-generated mistakes. In The Twelfth International Conference on Learning Representations, 2024b.
  4. Hsieh, C.-Y., Li, C.-L., kuan Yeh, C., Nakhost, H., Fujii, Y., Ratner, A., Krishna, R., Lee, C.-Y., and Pfister, T. Distilling step-by-step! outperforming larger language models with less training data and smaller model sizes. In Findings of the Association for Computational Linguistics: ACL 2023, p…
  5. AI-MO. Aimo validation aime dataset, 2024a. https://huggingface.co/datasets/AI-MO/aimo-validation-aime, 2024. ↗
  6. AI-MO. Aimo validation amc dataset, 2024b. https://huggingface.co/datasets/AI-MO/aimo-validation-amc, 2024. ↗
  7. Hu, J., Zhang, Y., Han, Q., Jiang, D., Zhang, X., and Shum, H.-Y. Open-reasoner-zero: An open source approach to scaling up reinforcement learning on the base model, 2025, 2025.
  8. Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., Hesse, C., and Schulman, J. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168, 2021.
  9. Hui, B., Yang, J., Cui, Z., Yang, J., Liu, D., Zhang, L., Liu, T., Zhang, J., Yu, B., Dang, K., et al. Qwen2. 5-coder technical report. arXiv preprint arXiv:2409.12186, 2024.
  10. Ko, J., Park, S., Jeong, M., Hong, S., Ahn, E., Chang, D.-S., and Yun, S.-Y. Revisiting intermediate layer distillation for compressing language models: An overfitting perspective. In Vlachos, A. and Augenstein, I. (eds.), Findings of the Association for Computational Linguistics: EACL 2023, pp. 158… find
  11. Gu, Y., Dong, L., Wei, F., and Huang, M. Minillm: Knowledge distillation of large language models. In The Twelfth International Conference on Learning Representations, 2024. find
  12. Gu, Y., Zhou, H., Meng, F., Zhou, J., and Huang, M. Mini-PLM: Knowledge distillation for pre-training language models. In The Thirteenth International Conference on Learning Representations, 2025. find
  13. Guo, D., Yang, D., Zhang, H., Song, J., Zhang, R., Xu, R., Zhu, Q., Ma, S., Wang, P., Bi, X., et al. Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint arXiv:2501.12948, 2025.
  14. Ko, J., Kim, S., Chen, T., and Yun, S.-Y. DistiLLM: Towards streamlined distillation for large language models. In Forty-first International Conference on Machine Learning, 2024. URL https://openreview.net/ forum?id=lsHZNNoC7r. ↗
  15. He, J., Liu, J., Liu, C. Y., Yan, R., Wang, C., Cheng, P., Zhang, X., Zhang, F., Xu, J., Shen, W., Li, S., Zeng, L., Wei, T., Cheng, C., An, B., Liu, Y., and Zhou, Y. Skywork open reasoner series, 2025, 2025. Notion Blog. find
  16. Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J., Zhang, H., and Stoica, I. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Principles, pp. 611–626, 2023.
  17. Hendrycks, D., Burns, C., Kadavath, S., Arora, A., Basart, S., Tang, E., Song, D., and Steinhardt, J. Measuring mathematical problem solving with the MATH dataset. Li, M., Zhou, F., and Song, X. BiLD: Bi-directional logits difference loss for large language model distillation. In Rambow, O., Wanner,… find
  18. Lin, A., Wohlwend, J., Chen, H., and Lei, T. Autoregressive knowledge distillation through imitation learning. In Webber, B., Cohn, T., He, Y., and Liu, Y. (eds.), Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 6121– 6133, Online, November 2020. A… find
  19. Liu, J., Zhang, C., Guo, J., Zhang, Y., Que, H., Deng, K., Liu, J., Zhang, G., Wu, Y., Liu, C., et al. DDK: Distilling domain knowledge for efficient large language models. Advances in Neural Information Processing Systems, 37: 98297–98319, 2024.
  20. Pan, J., Shen, W., Huang, S., Zhou, Q., and Zhang, Y. Pre-dpo: Improving data utilization in direct preference optimization using a guiding reference model, 2025, 2025. find
  21. Rafailov, R., Sharma, A., Mitchell, E., Manning, C. D., Ermon, S., and Finn, C. Direct preference optimization: Your language model is secretly a reward model. Advances in Neural Information Processing Systems, 36, 2023. rawsh. American mathematics competitions, 2024. https://huggingface.co/datasets… ↗
  22. Schulman, J., Levine, S., Abbeel, P., Jordan, M., and Moritz, P. Trust region policy optimization. In International conference on machine learning, pp. 1889–1897. find
  23. Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. Proximal policy optimization algorithms, 2017. arXiv preprint arXiv:1707.06347, 2017. find
  24. Shao, Z., Wang, P., Zhu, Q., Xu, R., Song, J., Bi, X., Zhang, H., Zhang, M., Li, Y., Wu, Y., and et al. Deepseekmath: Pushing the limits of mathematical reasoning in open language models, 2024. arXiv preprint arXiv:2402.03300, 2024.
  25. Liu, Z., Chen, C., Li, W., Qi, P., Pang, T., Du, C., Lee, W. S., and Lin, M. Understanding r1-zero-like training: A critical perspective, 2025. arXiv preprint arXiv:2503.20783, 2025.
  26. Sheng, G., Zhang, C., Ye, Z., Wu, X., Zhang, W., Zhang, R., Peng, Y., Lin, H., and Wu, C. Hybridflow: A flexible and efficient rlhf framework, 2024. arXiv preprint arXiv: 2409.19256, 2024.
  27. Luo, M., Tan, S., Wong, J., Shi, X., Tang, W. Y., Roongta, M., Cai, C., Luo, J., Li, L. E., Popa, R. A., and Stoica, I. Deepscaler: Surpassing o1-preview with a 1.5b model by scaling rl, 2025, 2025. Notion Blog. find
  28. Team, K., Bai, Y., Bao, Y., Chen, G., Chen, J., Chen, N., Chen, R., Chen, Y., Chen, Y., Chen, Y., Chen, Z., Cui, J., Ding, H., Dong, M., Du, A., Du, C., Du, D., Du, Y., Fan, Y., Feng, Y., Fu, K., Gao, B., Gao, H., Gao, P., Gao, T., Gu, X., Guan, L., Guo, H., Guo, J., Hu, H., Hao, X., He, T., He, W.,… find
  29. Meta. Llama 3.2: Revolutionizing edge AI and vision with open, customizable models. https://ai.meta.com/blog/ llama-3-2-connect-2024-vision-edge-\\ mobile-devices/, 2024. ↗
  30. Muennighoff, N., Yang, Z., Shi, W., Li, X. L., Fei-Fei, L., Hajishirzi, H., Zettlemoyer, L., Liang, P., Candès, E., and Hashimoto, T. s1: Simple test-time scaling, 2025.
  31. OpenCompass. Aime 2025 dataset, 2025. https://huggingface.co/datasets/opencompass/AIME2025, 2025. ↗
  32. H., Zhang, D., Zhang, H., Zhang, W., Zhang, X., Zhang, Y., Zhang, Y., Zhang, Y., Zhang, Y., Zhang, Y., Zhang, Y., Zhang, Z., Zhao, H., Zhao, Y., Zheng, H., Zheng, S., Zhou, J., Zhou, X., Zhou, Z., Zhu, Z., Zhuang, W., and Zu, X. Kimi k2: Open agentic intelligence, 2025a. URL https://arxiv.org/abs/25… ↗
  33. J., Zhou, J., Zhou, J., Lin, J., Dang, K., Bao, K., Yang, K., Yu, L., Deng, L., Li, M., Xue, M., Li, M., Zhang, P., Wang, P., Zhu, Q., Men, R., Gao, R., Liu, S., Luo, S., Li, T., Tang, T., Yin, W., Ren, X., Wang, X., Zhang, X., Ren, X., Fan, Y., Su, Y., Zhang, Y., Zhang, Y., Wan, Y., Liu, Y., Wang,… — Qwen3 Technical Report
  34. Team, K., Du, A., Gao, B., Xing, B., Jiang, C., Chen, C., Li, C., Xiao, C., Du, C., Liao, C., and et al. Kimi k1.5: Scaling reinforcement learning with llms, 2025b. arXiv preprint arXiv:2501.12599, 2025.
  35. Ye, Y., Huang, Z., Xiao, Y., Chern, E., Xia, S., and Liu, P. Limo: Less is more for reasoning, 2025.
  36. Tian, X., Zhao, S., Wang, H., Chen, S., Peng, Y., Ji, Y., Zhao, H., and Li, X. Deepdistill: Enhancing llm reasoning capabilities via large-scale difficulty-graded data training, 2025. find
  37. Yu, Q., Zhang, Z., Zhu, R., Yuan, Y., Zuo, X., Yue, Y., Fan, T., Liu, G., Liu, L., Liu, X., and et al. Dapo: An open-source llm reinforcement learning system at scale, 2025. arXiv preprint arXiv:2503.14476, 2025.
  38. Wang, W., Wei, F., Dong, L., Bao, H., Yang, N., and Zhou, M. Minilm: Deep self-attention distillation for task-agnostic compression of pre-trained transformers. Advances in neural information processing systems, 33: 5776–5788, 2020. find
  39. Yue, Y., Chen, Z., Lu, R., Zhao, A., Wang, Z., Song, S., and Huang, G. Does reinforcement learning really incentivize reasoning capacity in llms beyond the base model?, 2025a. arXiv preprint arXiv:2504.13837, 2025.
  40. Wen, L., Cai, Y., Xiao, F., He, X., An, Q., Duan, Z., Du, Y., Liu, J., Tang, L., Lv, X., Zou, H., Deng, Y., Jia, S., and Zhang, X. Light-r1: Curriculum sft, dpo and rl for long cot from scratch and beyond, 2025, 2025.
  41. Yue, Y., Yuan, Y., Yu, Q., Zuo, X., Zhu, R., Xu, W., Chen, J., Wang, C., Fan, T., Du, Z., Wei, X., Yu, X., Liu, G., Liu, J., Liu, L., Lin, H., Lin, Z., Ma, B., Zhang, C., Zhang, M., Zhang, W., Zhu, H., Zhang, R., Liu, X., Wang, M., Wu, Y., and Yan, L. Vapo: Efficient and reliable reinforcement learn… — VAPO: Efficient and Reliable Reinforcement Learning for Advanced Reasoning Tasks
  42. Xie, C., Huang, Y., Zhang, C., Yu, D., Chen, X., Lin, B. Y., Li, B., Ghazi, B., and Kumar, R. On memorization of large language models in logical reasoning, 2025a. URL https://arxiv.org/abs/2410.23123. ↗
  43. Zeng, W., Huang, Y., Liu, Q., Liu, W., He, K., Ma, Z., and He, J. Simplerl-zoo: Investigating and taming zero reinforcement learning for open base models in the wild, 2025, 2025.
  44. Xie, T., Gao, Z., Ren, Q., Luo, H., Hong, Y., Dai, B., Zhou, J., Qiu, K., Wu, Z., and Luo, C. Logic-rl: Unleashing llm reasoning with rule-based reinforcement learning, 2025, 2025b.
  45. Zhang, J. and Zuo, C. Grpo-lead: A difficulty-aware reinforcement learning approach for concise mathematical reasoning in language models, 2025, 2025.
  46. Xu, H., Zhu, Q., Deng, H., Li, J., Hou, L., Wang, Y., Shang, L., Xu, R., and Mi, F. Kdrl: Post-training reasoning llms via unified knowledge distillation and reinforcement learning, 2025. URL https://arxiv.org/abs/ 2506.02208. ↗
  47. Zhang, S., Zhang, X., Sun, Z., Chen, Y., and Xu, J. Dual-space knowledge distillation for large language models. In Al-Onaizan, Y., Bansal, M., and Chen, Y.-N. (eds.), Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pp. 18164–18181, Miami, Florida, USA, Novemb…
  48. Xu, W., Han, R., Wang, Z., Le, L. T., Madeka, D., Li, L., Wang, W. Y., Agarwal, R., Lee, C.-Y., and Pfister, T. Speculative knowledge distillation: Bridging the teacher-student gap through interleaved sampling. arXiv preprint arXiv:2410.11325, 2024.
  49. Yan, J., Li, Y., Hu, Z., Wang, Z., Cui, G., Qu, X., Cheng, Y., and Zhang, Y. Learning to reason under off-policy guidance, 2025, 2025.
  50. Yang, A., Li, A., Yang, B., Zhang, B., Hui, B., Zheng, B., Yu, B., Gao, C., Huang, C., Lv, C., Zheng, C., Liu, D., Zhou, F., Huang, F., Hu, F., Ge, H., Wei, H., Lin, H., Tang, J., Yang, J., Tu, J., Zhang, J., Yang, J., Yang, Metric GRPO KDRL KDRL KDRL RLAD RLAD RLAD RLAD RLAD −3 −2 −1 -λ = 10 λ = 10… find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/zhang-2026-reinforcement-aware-knowledge-distillation-llm/