ReCo: Reweighting GRPO Against Distributional Concentration

Problema

O Group Relative Policy Optimization (GRPO) emergiu como um método padrão de aprendizado por reforço para o pós-treinamento de modelos de linguagem, particularmente em domínios com recompensas verificáveis, como matemática e codificação. Embora o GRPO melhore a eficiência de amostragem para pequenos orçamentos de amostragem (baixo $k$ em Pass@$k$), observações recentes indicam que ele falha em superar, e frequentemente apresenta desempenho inferior ao modelo base, quando o orçamento de amostragem é grande (alto $k$). Isso sugere uma redução na cobertura de caminhos de raciocínio. O artigo identifica que o GRPO concentra a massa de probabilidade em respostas que o modelo base já gera com alta probabilidade, reduzindo assim a diversidade das trajetórias de resolução de problemas e limitando a capacidade do modelo de descobrir soluções corretas que residem em regiões de menor probabilidade da distribuição base.

Metodologia

Os autores propõem o ReCo (Reweighting GRPO Against Distributional Concentration), um método projetado para mitigar essa concentração modificando a atualização do GRPO em dois níveis:

1. Redirecionamento de Peso ao Nível da Resposta (Response-Level Reweighting):
   No GRPO padrão, respostas de alta probabilidade são amostradas com mais frequência dentro de um grupo de rollout. Como a vantagem do grupo é calculada pela média entre as respostas amostradas, essas trajetórias frequentes acumulam uma massa de gradiente agregada maior simplesmente devido à ocorrência repetida. O ReCo aborda isso normalizando a contribuição do gradiente de cada resposta pela sua contagem de ocorrência esperada dentro do grupo.

   • O peso é definido como o inverso da contagem esperada: $w_i = \frac{1}{G \cdot \bar{\pi}_{\theta_{old}}(o_i | q)}$, onde $\bar{\pi}$ é a verossimilhança normalizada pelo comprimento.
   • Isso evita que respostas frequentes dominem o gradiente do nível do grupo, garantindo que a atualização não reforce desproporcionalmente trajetórias que já são prováveis sob a política antiga.

2. Redirecionamento de Peso ao Nível do Token (Token-Level Reweighting):
   O GRPO padrão utiliza uma razão de importância ($\pi_\theta / \pi_{\theta_{old}}$) para escalar os gradientes ao nível do token. Isso cria um ciclo de feedback onde tokens que se tornam mais prováveis sob a política atual recebem atualizações maiores, reforçando ainda mais pontos de decisão saturados e reduzindo caminhos alternativos.

   • O ReCo substitui a razão de importância padrão por uma razão baseada em variância: $r_{var} = \frac{p_\theta(1 - p_\theta)}{p_{old}(1 - p_{old})}$.
   • O termo $p(1-p)$ representa a variância de Bernoulli. Esta razão atribui escalas de atualização maiores para pontos de decisão não saturados (onde $p$ não está próximo de 0 ou 1, significando que escolhas de tokens alternativas permanecem plausíveis) e reduz o peso de escolhas saturadas (onde $p \approx 1$).
   • Este mecanismo incentiva atualizações em pontos de decisão onde a política ainda não colapsou para uma única escolha, preservando assim a diversidade do raciocínio.

Principais Contribuições

• Análise da Concentração do GRPO: O artigo rastreia a redução no Pass@$k$ para grandes valores de $k$ até dois mecanismos específicos na atualização do GRPO: a dominância de respostas frequentes no nível do grupo e a amplificação de tokens saturados no nível do token.
• Algoritmo ReCo: A introdução de um esquema de redirecionamento de peso em dois níveis que normaliza as contribuições das respostas pela ocorrência esperada e substitui a razão de importância do token por uma métrica baseada em variância.
• Validação Empírica: Avaliação extensiva em diferentes famílias de modelos (Qwen2.5-Math, Llama-3.1) e escalas (1.5B, 7B, 8B) em cinco benchmarks de raciocínio matemático (AIME 2024/2025, AMC 2023, OlympiadBench, MATH 500).

Resultados

• Desempenho Pass@$k$: O ReCo supera consistentemente o GRPO no Pass@$k$ para valores grandes de $k$ (ex: Pass@64), indicando uma melhor cobertura de trajetórias de raciocínio corretas. Para valores pequenos de $k$, o ReCo permanece comparável ao GRPO, mantendo a eficiência de amostragem.
• Recuperação do Modelo Base: No backbone Llama-3.1-8B-Instruct, onde o GRPO padrão apresentou queda de desempenho no Pass@$k$ em relação ao modelo base, o ReCo reduz significativamente essa degradação e até supera o modelo base em certos benchmarks (AIME 2024/2025).
• Métricas de Diversidade: A análise da dinâmica de treinamento mostra que, enquanto o GRPO perde diversidade rapidamente (diminuindo o Distinct-2, aumentando o Self-BLEU), o ReCo mantém a diversidade de rollout próxima à do modelo base durante todo o treinamento.
• Variância de Token: O ReCo mantém uma variância de Bernoulli ao nível do token substancialmente maior nos segmentos iniciais e médios das respostas em comparação ao GRPO, indicando que os caminhos de raciocínio permanecem não saturados por mais tempo.
• Generalização: O método também melhora o Pass@$k$ em benchmarks de raciocínio de código (HumanEval, HumanEval+, LiveCodeBench), sugerindo que o problema de concentração não se limita ao raciocínio matemático.
• Compatibilidade: O ReCo pode ser combinado com outras variantes do GRPO (ex: DAPO, Regra 80/20) para gerar melhorias adicionais.

Significância

O artigo afirma que o ReCo aborda os mecanismos internos do GRPO que levam à concentração distributiva sem exigir objetivos de diversidade externos ou regularizadores. Ao redirecionar o peso da atualização com base na probabilidade de ocorrência e na saturação dos pontos de decisão, o ReCo preserva um conjunto mais amplo de caminhos de raciocínio corretos. Isso permite que o modelo retenha os ganhos de eficiência de amostragem do aprendizado por reforço enquanto expande sua cobertura do espaço de soluções, mitigando efetivamente o "colapso de diversidade" observado no treinamento padrão do GRPO. O trabalho sugere que as limitações do RLVR em expandir a capacidade de raciocínio devem-se, em parte, à dinâmica de otimização da própria atualização da política, que pode ser corrigida através de um redirecionamento de peso direcionado.