Resumo Técnico: Domando Tokens Extremos com GRPO Consciente de Covariância

Declaração do Problema

O Otimização de Política Relativa por Grupos (GRPO) tornou-se uma abordagem padrão para aprimorar as capacidades de raciocínio de Modelos de Linguagem de Grande Escala (LLMs), particularmente em tarefas matemáticas e de codificação. No entanto, o GRPO enfrenta uma limitação crítica ao equilibrar o trade-off entre exploração (busca por novas estratégias de raciocínio) e exploração (refinamento de soluções boas conhecidas).

O artigo identifica que esse desequilíbrio está fundamentalmente ligado às dinâmicas de entropia da política durante o treinamento. Baseando-se em insights teóricos de Cui et al. (2025a), os autores observam que as mudanças na entropia da política são governadas pela covariância entre os log-probabilidades dos tokens e suas estimativas de vantagem correspondentes. A análise empírica revela que uma pequena fração de tokens com valores de covariância extremos domina desproporcionalmente as atualizações da política. Esses outliers causam instabilidade na entropia, levando a uma superconfiança em soluções subótimas (exploração excessiva) ou a dinâmicas de treinamento instáveis (exploração excessiva), resultando, em última análise, em desempenho downstream subótimo.

Metodologia

A solução proposta é o GRPO Consciente de Covariância (CW-GRPO), um método de otimização livre de hiperparâmetros que repondera dinamicamente os sinais de vantagem com base na covariância em nível de token.

1. Fundamentação Teórica

O método baseia-se na observação de que a mudança na entropia da política ($\Delta H$) sob atualizações de gradiente de política natural é aproximadamente proporcional à covariância negativa entre os log-probabilidades dos tokens e as vantagens:
$$\Delta H \approx -\eta \cdot \text{Cov}_t(\log \pi_\theta(o_t|q, o_{<t}), \hat{A}_i)$$
Valores extremos nessa covariância afastam a entropia de um estado estável e equilibrado.

2. Reponderação de Vantagem com Kernel Gaussiano

Para mitigar o efeito desestabilizador de tokens extremos, os autores introduzem um mecanismo de reponderação:

• Cálculo de Covariância: Para cada token $t$ na resposta $i$, o método calcula a covariância entre o log-probabilidade centrado e a vantagem centrada:
  $$c_{i,t} = (\log \pi_\theta(o_{i,t}|\dots) - \overline{\log \pi}) \cdot (\hat{A}_i - \overline{\hat{A}})$$
• Ponderação Gaussiana: Um kernel gaussiano é aplicado para suprimir tokens com covariâncias de alta magnitude. A largura de banda é definida como o desvio padrão empírico ($\sigma$) da distribuição de covariância:
  $$w_{i,t} = \exp\left(-\frac{c_{i,t}^2}{2\sigma^2}\right)$$
  Esta função filtra suavemente tokens "outliers" que, de outra forma, causariam colapso ou explosão da entropia, preservando ao mesmo tempo os sinais de aprendizado de tokens com covariâncias moderadas.
• Normalização: Os pesos são normalizados para manter a magnitude esperada da perda em todo o lote.
• Objetivo Modificado: A perda em nível de token é atualizada multiplicando a vantagem $\hat{A}_i$ pelo fator de reponderação $\tilde{w}_{i,t}$ antes de aplicar o objetivo padrão do GRPO.

Contribuições Principais

1. Identificação de Tokens Extremos: O artigo demonstra empiricamente que uma pequena fração de tokens com valores de covariância extremos impulsiona desproporcionalmente as atualizações da política e a instabilidade da entropia no GRPO padrão.
2. Otimização Livre de Hiperparâmetros: O método proposto CW-GRPO reduz automaticamente o peso das atualizações extremas usando um kernel gaussiano, sem exigir hiperparâmetros adicionais ajustáveis.
3. Dinâmicas de Entropia Estabilizadas: Ao suprimir atualizações de covariância extremas, o método mantém a entropia da política em um nível razoável, equilibrando efetivamente a exploração e a exploração ao longo do treinamento.

Resultados Experimentais

Os autores avaliaram o CW-GRPO em duas escalas de modelo (1,5B e 7B) usando o conjunto de dados de treinamento Open-RS e cinco benchmarks de raciocínio matemático (AIME24, MATH-500, AMC23, Minerva, OlympiadBench).

• Ganhos de Desempenho:
  • Modelo 1,5B (DeepSeek-R1-Distill-Qwen): O CW-GRPO alcançou uma melhoria média de +6,4% sobre o modelo base e +2,7% sobre o GRPO padrão. Ganhos notáveis foram observados em AMC23 e Minerva.
  • Modelo 7B (Qwen2.5-Math): O CW-GRPO alcançou uma melhoria média de +4,3% sobre o modelo base e +2,4% sobre o GRPO padrão, com ganhos significativos em AMC23 e OlympiadBench.
• Estabilidade da Entropia:
  • O GRPO padrão exibiu instabilidade significativa de entropia, com desempenho degradando à medida que o treinamento avançava (por exemplo, as pontuações no MATH-500 caíram de 85,0% para 79,8% à medida que a entropia flutuava).
  • O CW-GRPO manteve desempenho consistente em todos os checkpoints de treinamento (86,2%–87,0% no MATH-500), validando que o controle de valores de covariância extremos correlaciona-se com desempenho de raciocínio estável e superior.

Significado e Alegações

O artigo afirma que sua abordagem fornece um mecanismo principiado para equilibrar exploração e exploração no GRPO, sem a necessidade de ajuste manual de hiperparâmetros. Ao abordar diretamente a causa raiz da instabilidade da entropia — atualizações de covariância em nível de token extremas — o método permite uma otimização de política mais robusta.

Os autores posicionam este trabalho como um passo em direção a um aprendizado por reforço mais estável para raciocínio, demonstrando que moderar atualizações extremas leva a melhorias consistentes downstream em diferentes arquiteturas de modelos e benchmarks matemáticos. Eles reconhecem limitações, observando que os experimentos foram conduzidos em modelos de até 7B parâmetros e focaram principalmente no raciocínio matemático, sugerindo que validação adicional em modelos de maior escala e tarefas diversas é necessária para confirmar a ampla aplicabilidade.