Reinforcement-aware Knowledge Distillation for LLM Reasoning

Resumo Técnico: Destilação de Conhecimento com Sensibilidade à RL para Raciocínio de LLM (RLAD)

1. Declaração do Problema

Grandes Modelos de Linguagem (LLMs) alcançaram capacidades significativas de raciocínio por meio de pós-treinamento com Aprendizado por Reforço (RL), particularmente em tarefas de cadeia de pensamento (CoT) longa. No entanto, os altos custos de inferência desses modelos grandes tornam necessária a destilação para modelos estudantes menores. Os métodos existentes de Destilação de Conhecimento (KD) enfrentam limitações críticas quando aplicados ao pós-treinamento de RL:

• Desajuste de Distribuição: Métodos padrão de KD (offline, off-policy ou on-policy) frequentemente dependem de traços fixos do professor ou de regularização de divergência KL estática. Em RL, a política do estudante evolui continuamente, causando o deslocamento da distribuição de rollout. Alvos de destilação estáticos tornam-se subótimos ou instáveis à medida que o estudante se afasta da distribuição inicial do professor.
• Interferência de Objetivos: Métodos que integram KD no loop de RL (por exemplo, KDRL) tipicamente adicionam um termo de divergência KL entre o estudante e o professor como uma penalidade de regularização. Isso cria um problema de otimização multiobjetivo onde o objetivo de maximização de recompensa compete com o objetivo de imitação. Essa competição exige um ajuste cuidadoso de hiperparâmetros (balanceamento de perda) e pode levar a dinâmicas de treinamento instáveis, pois a supervisão do professor pode não estar alinhada com as estimativas de vantagem atuais do estudante.

2. Metodologia: Destilação com Sensibilidade à RL (RLAD)

Os autores propõem o RLAD, um framework projetado para realizar a imitação seletiva durante o pós-treinamento de RL. O princípio central é que o estudante deve ser guiado pelo professor apenas quando tal orientação melhora a atualização da política atual sob o objetivo de RL.

Componente Chave: Razão de Região de Confiança (TRRD)
Em vez de usar um regularizador KL independente, o RLAD introduz o TRRD, que substitui a razão de amostragem de importância padrão no Group Relative Policy Optimization (GRPO) por uma mistura professor-estudante.

• Mecanismo: O TRRD define uma nova razão de verossimilhança $r^{TRRD}_{i,t}$ para o $t$-ésimo token da $i$-ésima resposta:
  $$r^{TRRD}_{i,t}(\theta_S) = \left( \frac{\pi_{\theta_S}(y_t | x, y_{<t})}{\pi_{\theta_S,old}(y_t | x, y_{<t})} \right)^\alpha \left( \frac{\pi_{\theta_S}(y_t | x, y_{<t})}{\pi_{\theta_T}(y_t | x, y_{<t})} \right)^{1-\alpha}$$
  Aqui, $\pi_{\theta_S}$ é a política atual do estudante, $\pi_{\theta_S,old}$ é a política anterior do estudante e $\pi_{\theta_T}$ é a política do professor. $\alpha \in [0, 1]$ é um coeficiente de mistura.
• Interpretação de Região de Confiança: Esta razão pode ser vista como uma atualização estilo PPO onde a região de confiança é ancorada em uma mistura do antigo estudante e do professor ($r_{\pi_{mix}}$), em vez de apenas o antigo estudante.
• Seleção Sensível à Vantagem: A atualização é ponderada pela vantagem relativa ao grupo $\hat{A}_{i,t}$.
  • Vantagem Positiva: Se a vantagem for positiva, o estudante é encorajado a aumentar a probabilidade do token. Se o professor também atribuir alta probabilidade, a âncora da mistura é maior, permitindo que o estudante aumente sua probabilidade antes de atingir o limite de clipping (permitindo atualizações maiores). Se o professor atribuir baixa probabilidade, a âncora é menor, restringindo a magnitude da atualização.
  • Vantagem Negativa: Se a vantagem for negativa, o estudante é encorajado a diminuir a probabilidade. Se o professor atribuir alta probabilidade, o limite de clipping impede que o estudante suprima agressivamente o token, preservando comportamentos alinhados ao professor.
  • Vantagem Próxima de Zero: A influência do professor torna-se negligenciável, evitando interferência quando o sinal de recompensa não é informativo.

Este design equilibra implicitamente exploração, explotação e imitação ao garantir que a orientação do professor seja modulada pelo sinal de recompensa (vantagem) e contida dentro de uma região de confiança.

3. Principais Contribuições

1. Framework RLAD: Um novo framework de destilação para pós-treinamento de RL que integra a orientação do professor via Trust Region Ratio Distillation (TRRD), substituindo a regularização KL estática por um objetivo de razão de verossimilhança clipado e ponderado pela vantagem.
2. Análise Teórica: Uma análise demonstrando que o TRRD permite a imitação seletiva condicionada à recompensa. Ele preserva atualizações limitadas pela região de confiança e equilibra naturalmente o trade-off entre seguir o professor e otimizar as recompensas sem exigir um balanceamento de perda delicado.
3. Validação Empírica: Experimentos extensos em tarefas de raciocínio lógico (K&K Logistics) e raciocínio matemático de contexto longo (AIME, AMC, MATH) mostrando melhorias consistentes sobre os baselines.

4. Resultados Experimentais

Os autores avaliaram o RLAD em tarefas de raciocínio lógico e raciocínio matemático complexo usando modelos estudantes variando de 0.6B a 8B de parâmetros e professores de até 32B.

• Raciocínio Lógico (K&K Logistics):

  • Usando um professor Qwen3-8B, o RLAD melhorou significativamente a precisão média para Qwen3-0.6B (de 0.76 para 0.94 em 8K de contexto) e Qwen3-1.7B (de 0.95 para 0.99 em 8K de contexto) comparado ao GRPO padrão.
  • Os ganhos foram mais pronunciados nos subconjuntos mais difíceis (PPL7 e PPL8), onde o RLAD manteve uma margem clara sobre o KDRL.
  • A dinâmica de treinamento mostrou convergência mais rápida e recompensas finais mais altas em comparação tanto com o GRPO quanto com o KDRL.

• Raciocínio Matemático:

  • Modelos Base (estilo R1-Zero): No Qwen3-8B-Base, o RLAD melhorou o desempenho médio em +5.5 pontos sobre o GRPO em tarefas de contexto longo (30K).
  • Modelos Pós-treinados: O RLAD superou a destilação explícita de KL (KDRL) e a destilação offline.
  • Especificidades de Métrica: O RLAD mostrou maiores ganhos no Pass@1 (Mean@32) comparado ao KDRL, que tendeu a mostrar maiores ganhos relativos no Pass@32. Os autores interpretam isso como o RLAD impulsionando com sucesso a melhoria da política alinhada à recompensa (explotação), enquanto os ganhos do KDRL foram mais impulsionados pela pura imitação de soluções diversas do professor.
  • Estabilidade: O RLAD exibiu dinâmicas de treinamento mais estáveis com menos oscilações comparado ao KDRL, que sofreu com problemas de acoplamento de objetivos.

• Robustez e Eficiência:

  • Professor Fraco: Quando o professor era mais fraco que o estudante, o RLAD permaneceu próximo ao desempenho do GRPO, enquanto o KDRL degradou-se mais significativamente.
  • Hiperparâmetros: O RL_AD foi considerado amplamente insensível ao coeficiente de mistura $\alpha$ (fixado em 0.5), ao contrário do KDRL que é sensível ao peso de KL $\lambda$.
  • Eficiência: Integrar o RLAD adicionou apenas cerca de 12% de latência de treinamento comparado ao GRPO padrão, já que o professor é usado apenas para o cálculo da log-probabilidade, não para a geração de trajetórias.

5. Significância e Alegações

O artigo afirma que o RLAD aborda o desalinhamento fundamental entre a supervisão estática do professor e as atualizações dinâmicas de política de RL. Ao incorporar a orientação do professor no mecanismo de atualização da região de confiança, o RLAD evita o "ajuste frágil" exigido pela regularização KL multiobjetivo.

Os autores asseguram que o RLAD permite que o estudante se mova em direção a comportamentos de alta vantagem dentro de uma distribuição âncora enquanto explora oportunisticamente a orientação do professor. Isso resulta em:

• Desempenho Superior: Superação consistente do GRPO, KDRL e destilação offline, particularmente em benchmarks desafiadores.
• Estabilidade: Dinâmicas de treinamento mais estáveis e melhor convergência.
• Otimização Real de RL: Diferente de métodos que dependem fortemente de imitação (que podem aumentar o Pass@32 mas não o Pass@1), o RLAD prioriza o comportamento alinhado à recompensa, levando a ganhos mais fortes nas respostas corretas mais prováveis do modelo (Pass@1).

O artigo conclui que, embora a abordagem dependa de modelos professores de pesos abertos e bem calibrados para acesso aos logits, ela representa um passo significativo para tornar a destilação compatível com a natureza dinâmica do pós-treinamento de RL.