cs.LGFeb 26, 2026

Reinforcement-aware Knowledge Distillation for LLM Reasoning

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

Organizations: AWS Agentic AI

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.

Figures & tables

Explore similar work

CardsList
  1. Distilled Reinforcement Learning for LLM Post-training

    Jul 19, 2026Chen Wang, Zhaochun Li, Jionghao Bai +4RL for Language Model ReasoningCross-Architecture Knowledge Distillation

  2. The Weakest Link: Distilling LLM Reasoning with Worst-Case Constrained Reinforcement Learning

    Sep 29, 2026Matthieu Zimmer, Xiaotong Ji, Tu Nguyen +1Language Model DistillationCoT Distillation

  3. Decoupling KL and Trajectories: A Unified Perspective for SFT, DAgger, Offline RL, and OPD in LLM Distillation

    May 16, 2026Anhao Zhao, Haoran Xin, Yingqi Fan +3Reasoning DistillationOn-Policy Distillation