Distillation de connaissances sensible au renforcement pour le raisonnement des grands modèles de langage
Reinforcement-aware Knowledge Distillation for LLM Reasoning
February 26, 2026
Auteurs: Zhaoyang Zhang, Shuli Jiang, Yantao Shen, Yuting Zhang, Dhananjay Ram, Shuo Yang, Zhuowen Tu, Wei Xia, Stefano Soatto
cs.AI
Résumé
L'apprentissage par renforcement (RL) en post-formation a récemment permis des progrès majeurs dans les grands modèles de langage (LLM) pour le raisonnement à longue chaîne de pensée, mais le coût d'inférence élevé de ces modèles motive leur distillation vers des modèles étudiants plus petits. La plupart des méthodes de distillation de connaissances (KD) existantes sont conçues pour l'affinage supervisé (SFT), reposant sur des traces fixes de l'enseignant ou sur une régularisation basée sur la divergence de Kullback-Leibler (KL) entre l'enseignant et l'étudiant. Lorsqu'elles sont combinées avec le RL, ces approches souffrent souvent d'un décalage de distribution et d'interférences d'objectifs : la supervision de l'enseignant peut ne pas s'aligner sur la distribution d'exploration évolutive de l'étudiant, et le régulariseur KL peut entrer en concurrence avec la maximisation de la récompense et nécessiter un équilibrage minutieux des pertes. Pour résoudre ces problèmes, nous proposons la distillation sensible au RL (RLAD), qui effectue une imitation sélective pendant le RL – guidant l'étudiant vers l'enseignant uniquement lorsque cela améliore la mise à jour de la politique actuelle. Notre composant central, la Distillation par Ratio de Région de Confiance (TRRD), remplace le régulariseur KL enseignant-étudiant par un objectif de rapport de vraisemblance de type PPO/GRPO ancré à un mélange enseignant-ancienne-politique, produisant une distillation avantageuse et bornée par une région de confiance sur les explorations de l'étudiant, et équilibrant naturellement l'exploration, l'exploitation et l'imitation. Sur divers benchmarks de raisonnement logique et mathématique, RLAD surpasse systématiquement la distillation hors ligne, le GRPO standard et la distillation de connaissances enseignant-étudiant en ligne basée sur la KL.
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.