Distillazione della Consapevolezza del Rinforzo per il Ragionamento nei LLM
Reinforcement-aware Knowledge Distillation for LLM Reasoning
February 26, 2026
Autori: Zhaoyang Zhang, Shuli Jiang, Yantao Shen, Yuting Zhang, Dhananjay Ram, Shuo Yang, Zhuowen Tu, Wei Xia, Stefano Soatto
cs.AI
Abstract
L'addestramento post-rinforzo (RL) ha recentemente prodotto notevoli miglioramenti nei grandi modelli linguistici (LLM) per il ragionamento a catena di pensiero complessa, ma l'elevato costo inferenziale di tali modelli motiva la distillazione in studenti più piccoli. La maggior parte dei metodi di distillazione della conoscenza (KD) esistenti sono progettati per l'addestramento supervisionato (SFT), basandosi su tracce fisse del insegnante o su una regolarizzazione basata sulla divergenza di Kullback-Leibler (KL) insegnante-studente. Se combinati con l'RL, questi approcci spesso soffrono di disallineamento distributivo e interferenza degli obiettivi: la supervisione del insegnante potrebbe non allinearsi con la distribuzione in evoluzione dello studente, e il regolarizzatore KL può competere con la massimizzazione della ricompensa e richiedere un bilanciamento accurato delle perdite. Per affrontare questi problemi, proponiamo la distillazione consapevole dell'RL (RLAD), che esegue un'imitazione selettiva durante l'RL – guidando lo studente verso il insegnante solo quando ciò migliora l'aggiornamento della politica corrente. Il nostro componente principale, la Distillazione del Rapporto nella Regione di Fiducia (TRRD), sostituisce il regolarizzatore KL insegnante-studente con un obiettivo a rapporto di verosimiglianza in stile PPO/GRPO ancorato a una miscela insegnante–politica precedente, producendo una distillazione vantaggio-consapevole e delimitata dalla regione di fiducia sui rollout dello studente e bilanciando naturalmente esplorazione, sfruttamento e imitazione. In vari benchmark di ragionamento logico e matematico, RLAD supera costantemente la distillazione offline, il GRPO standard e la distillazione della conoscenza insegnante-studente on-policy basata su 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.