Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Reinforcement-aware Knowledge Distillation for LLM Reasoning

Zhaoyang Zhang, Shuli Jiang, Yantao Shen, Yuting Zhang, Dhananjay Ram, Shuo Yang, Zhuowen Tu, Wei Xia, Stefano Soatto

TL;DR

RL-aware distillation (RLAD), which performs selective imitation during RL -- guiding the student toward the teacher only when it improves the current policy update, consistently outperforms offline distillation, standard GRPO, and KL-based on-policy teacher-student knowledge distillation.

Abstract

Reinforcement learning (RL) post-training has recently driven major gains in long chain-of-thought reasoning large language models (LLMs), but the high inference cost of such models motivates distillation into smaller students. Most existing knowledge distillation (KD) methods are designed for supervised fine-tuning (SFT), relying on fixed teacher traces or teacher-student Kullback-Leibler (KL) divergence-based regularization. When combined with RL, these approaches often suffer from distribution mismatch and objective interference: teacher supervision may not align with the student's evolving rollout distribution, and the KL regularizer can compete with reward maximization and require careful loss balancing. To address these issues, we propose RL-aware distillation (RLAD), which performs selective imitation during RL -- guiding the student toward the teacher only when it improves the current policy update. Our core component, Trust Region Ratio Distillation (TRRD), replaces the teacher-student KL regularizer with a PPO/GRPO-style likelihood-ratio objective anchored to a teacher--old-policy mixture, yielding advantage-aware, trust-region-bounded distillation on student rollouts and naturally balancing exploration, exploitation, and imitation. Across diverse logic reasoning and math benchmarks, RLAD consistently outperforms offline distillation, standard GRPO, and KL-based on-policy teacher-student knowledge distillation.

Reinforcement-aware Knowledge Distillation for LLM Reasoning

TL;DR

RL-aware distillation (RLAD), which performs selective imitation during RL -- guiding the student toward the teacher only when it improves the current policy update, consistently outperforms offline distillation, standard GRPO, and KL-based on-policy teacher-student knowledge distillation.

Abstract

Reinforcement learning (RL) post-training has recently driven major gains in long chain-of-thought reasoning large language models (LLMs), but the high inference cost of such models motivates distillation into smaller students. Most existing knowledge distillation (KD) methods are designed for supervised fine-tuning (SFT), relying on fixed teacher traces or teacher-student Kullback-Leibler (KL) divergence-based regularization. When combined with RL, these approaches often suffer from distribution mismatch and objective interference: teacher supervision may not align with the student's evolving rollout distribution, and the KL regularizer can compete with reward maximization and require careful loss balancing. To address these issues, we propose RL-aware distillation (RLAD), which performs selective imitation during RL -- guiding the student toward the teacher only when it improves the current policy update. Our core component, Trust Region Ratio Distillation (TRRD), replaces the teacher-student KL regularizer with a PPO/GRPO-style likelihood-ratio objective anchored to a teacher--old-policy mixture, yielding advantage-aware, trust-region-bounded distillation on student rollouts and naturally balancing exploration, exploitation, and imitation. Across diverse logic reasoning and math benchmarks, RLAD consistently outperforms offline distillation, standard GRPO, and KL-based on-policy teacher-student knowledge distillation.
Paper Structure (17 sections, 14 equations, 2 figures, 5 tables)

This paper contains 17 sections, 14 equations, 2 figures, 5 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Background and Notations
  4. RL-Aware Distillation (RLAD)
  5. Trust Region Ratio Distillation (TRRD)
  6. Experiments
  7. Logical Reasoning Task
  8. Experimental Setup
  9. Main Results
  10. Complex Math Reasoning Task
  11. Experimental Setup
  12. Main Results
  13. Efficiency Analysis
  14. Conclusion
  15. Ablation on the choice of $\alpha$
  16. ...and 2 more sections

Figures (2)

  • Figure 1: Comparisons of training dynamics RL training on logistics reasoning in terms of both training and validation average rewards for Qwen3-0.6B abd Qwen3-1.7B model.
  • Figure 2: Comparison of validation accuracy on AIME-24 dataset dynamics during RL training on the math reasoning task for Qwen3-8B-Base, reported in terms of Mean@32 (as the mean Pass@1 over 32 runs) and Pass@32.