Verstärkungsbewusste Wissensdistillation für das Schlussfolgern mit großen Sprachmodellen
Reinforcement-aware Knowledge Distillation for LLM Reasoning
February 26, 2026
Autoren: Zhaoyang Zhang, Shuli Jiang, Yantao Shen, Yuting Zhang, Dhananjay Ram, Shuo Yang, Zhuowen Tu, Wei Xia, Stefano Soatto
cs.AI
Zusammenfassung
Verstärkendes Lernen (RL) nach dem Training hat kürzlich bedeutende Fortschritte bei Large Language Models (LLMs) mit langen Denkketten erzielt, doch die hohen Inferenzkosten solcher Modelle motivieren die Distillation in kleinere Studentenmodelle. Die meisten bestehenden Wissensdistillationsmethoden (KD) sind für überwachtes Feinjustieren (SFT) konzipiert und stützen sich auf feste Lehrer-Traces oder auf Kullback-Leibler (KL)-Divergenz basierende Regularisierung zwischen Lehrer und Student. In Kombination mit RL leiden diese Ansätze oft unter Verteilungsinkongruenz und Zielkonflikt: Die Lehrer-Aufsicht stimmt möglicherweise nicht mit der sich entwickelnden Rollout-Verteilung des Studenten überein, und der KL-Regularisierer kann mit der Belohnungsmaximierung konkurrieren und erfordert eine sorgfältige Loss-Balance. Um diese Probleme zu adressieren, schlagen wir RL-bewusste Distillation (RLAD) vor, die selektive Imitation während des RL durchführt – sie lenkt den Studenten nur dann in Richtung des Lehrers, wenn dies die aktuelle Policy-Aktualisierung verbessert. Unsere Kernkomponente, Trust Region Ratio Distillation (TRRD), ersetzt den KL-Regularisierer zwischen Lehrer und Student durch ein PPO/GRPO-artiges Likelihood-Ratio-Ziel, das an einer Mischung aus Lehrer- und alter Policy verankert ist. Dies ermöglicht vorteilsbewusste, durch eine Trust Region begrenzte Distillation auf Studenten-Rollouts und balanciert Exploration, Exploitation und Imitation natürlich aus. Über diverse Logik- und Mathematik-Benchmarks hinweg übertrifft RLAD konsistent Offline-Distillation, Standard-GRPO und auf KL basierende On-Policy-Lehrer-Studenten-Wissensdistillation.
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.