让小模型更高效学习大模型的推理能力,解决强化学习训练中的知识迁移难题。
Reinforcement-aware Knowledge Distillation for LLM Reasoning
- 基于强化学习动态选择性模仿,只在有益时引导学生
- 新方法使小模型在多个数学逻辑任务上超越现有技术
- 适合追求高效推理的小模型部署场景
强化学习后训练显著提升了大语言模型的长链推理能力,但高推理成本促使将其压缩为小型学生模型。现有知识蒸馏方法多针对监督微调设计,依赖固定教师轨迹或教师-学生KL散度正则化。结合强化学习时,这些方法常因分布不匹配和目标冲突而表现不佳:教师监督与学生滚动分布不一致,且KL正则项会与奖励最大化竞争,需精细调整损失权重。为此,我们提出强化学习感知蒸馏(RLAD),在强化学习过程中进行选择性模仿——仅当有助于当前策略更新时才引导学生向教师学习。核心组件信任区域比率蒸馏(TRRD)以教师-旧策略混合为基础,采用类似PPO/GRPO的似然比目标,实现优势感知、信任区域约束的蒸馏,自然平衡探索、利用与模仿。在多种逻辑推理与数学基准测试中,RLAD持续优于离线蒸馏、标准GRPO及基于KL的在线教师-学生蒸馏方法。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。