Distilação de Conhecimento Consciente do Reforço para Raciocínio em LLMs

Resumo

O pós-treinamento por aprendizagem por reforço (RL) tem recentemente impulsionado ganhos significativos em modelos de linguagem grandes (LLMs) com raciocínio de cadeia longa de pensamento (chain-of-thought), mas o alto custo de inferência desses modelos motiva a sua destilação em estudantes menores. A maioria dos métodos existentes de destilação de conhecimento (KD) são projetados para ajuste fino supervisionado (SFT), dependendo de traços fixos do professor ou de regularização baseada na divergência de Kullback-Leibler (KL) entre professor e estudante. Quando combinadas com RL, essas abordagens frequentemente sofrem com desajuste de distribuição e interferência de objetivo: a supervisão do professor pode não estar alinhada com a distribuição evolutiva de *rollout* do estudante, e o regularizador de KL pode competir com a maximização de recompensa e exigir um balanceamento cuidadoso da função de perda. Para resolver essas questões, propomos a destilação consciente do RL (RLAD), que realiza imitação seletiva durante o RL — guiando o estudante em direção ao professor apenas quando isso melhora a atualização de política atual. Nosso componente central, a Destilação por Razão da Região de Confiança (TRRD), substitui o regularizador de KL professor-estudante por um objetivo de razão de verossimilhança no estilo PPO/GRPO ancorado a uma mistura professor-política-antiga, produzindo uma destilação consciente da vantagem e limitada pela região de confiança nos *rollouts* do estudante, e balanceando naturalmente a exploração, a exploração (*exploitation*) e a imitação. Em diversos benchmarks de raciocínio lógico e matemática, o RLAD supera consistentemente a destilação offline, o GRPO padrão e a destilação de conhecimento professor-estudante *on-policy* baseada em 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.

Distilação de Conhecimento Consciente do Reforço para Raciocínio em LLMs

Reinforcement-aware Knowledge Distillation for LLM Reasoning

Resumo

Support