Versterkingsbewuste Kennisdistillatie voor Redeneervaardigheden van Grote Taalmodellen
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
Samenvatting
Versterkend leren (VL) na de training heeft recentelijk grote vooruitgang geboekt bij grote taalmodellen (GTM's) met lange redeneerketens, maar de hoge inferentiekosten van dergelijke modellen stimuleren de distillatie naar kleinere studentmodellen. De meeste bestaande kennisdistillatie (KD) methoden zijn ontworpen voor supervised fine-tuning (SFT) en zijn afhankelijk van vaste leraarsporen of Kullback-Leibler (KL) divergentie-gebaseerde regularisatie tussen leraar en student. In combinatie met VL leiden deze benaderingen vaak tot distributiemismatch en doelstellingsinterferentie: de supervisie van de leraar komt mogelijk niet overeen met de evoluerende rollout-distributie van de student, en de KL-regularisator kan concurreren met beloningsmaximalisatie en vereist een zorgvuldige balancering van de verliesfunctie. Om deze problemen aan te pakken, stellen wij Reinforcement Learning-Aware Distillation (RLAD) voor, dat selectieve imitatie uitvoert tijdens VL – waarbij de student alleen naar de leraar wordt geleid wanneer dit de huidige policy-update verbetert. Onze kerncomponent, Trust Region Ratio Distillation (TRRD), vervangt de KL-regularisator tussen leraar en student door een PPO/GRPO-achtig likelihood-ratio doelwit verankerd aan een mengsel van de leraar en het oude beleid, wat resulteert in voordeel-bewuste, trust-region-gebonden distillatie op student-rollouts en een natuurlijke balans biedt tussen exploratie, exploitatie en imitatie. Over diverse logische redeneer- en rekenbenchmarks heen presteert RLAD consistent beter dan offline distillatie, standaard GRPO en KL-gebaseerde on-policy kennisdistillatie tussen leraar en student.
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.