EBPO: Empirical Bayes Shrinkage for Stabilizing Group-Relative Policy Optimization

Resumo Técnico: EBPO – Otimização de Política com Encolhimento Bayesiano Empírico

1. O Problema

O Reinforcement Learning with Verifiable Rewards (RLVR) tem sido fundamental para melhorar as capacidades de raciocínio em Modelos de Linguagem de Grande Escala (LLMs), especialmente em tarefas como matemática e geração de código. A abordagem dominante atual é o Group Relative Policy Optimization (GRPO), que elimina a necessidade de uma rede de valor (critic) ao normalizar recompensas dentro de um grupo de respostas geradas para um mesmo prompt.

No entanto, o GRPO enfrenta dois desafios críticos de estabilidade:

1. Alta Variância do Estimador: Em cenários com restrições computacionais, onde o tamanho do grupo ($G$) é pequeno, a média e o desvio padrão calculados apenas localmente (dentro do grupo) sofrem de alta variância, resultando em gradientes ruidosos.
2. Gradientes Desvanecidos (Vanishing Gradients) em Regimes Saturados: Se um modelo falha em todas as tentativas de um prompt difícil (todas as recompensas são zero), a vantagem relativa dentro do grupo torna-se zero. Isso resulta em um gradiente nulo, desperdiçando o passo de treinamento. Da mesma forma, se todas as respostas forem corretas, não há sinal relativo para melhoria.

Abordagens anteriores tentaram mitigar isso filtrando exemplos problemáticos, o que leva ao desperdício de dados. O GRPO padrão frequentemente exige grupos grandes para suprimir o ruído, aumentando drasticamente o custo computacional.

2. Metodologia: EBPO

O artigo propõe o EBPO (Empirical Bayes Policy Optimization), um novo quadro de trabalho que regulariza as bases (baselines) locais baseadas em grupos, "emprestando força" das estatísticas globais acumuladas da política.

Conceito Central:
Em vez de estimar a base isoladamente para cada grupo, o EBPO utiliza um estimador de encolhimento (shrinkage estimator). Ele equilibra dinamicamente a estatística local do grupo com uma priors global atualizada via algoritmo online.

Mecanismos Principais:

• Estimativa Bayesiana Empírica: Assume-se que a probabilidade latente de sucesso de um prompt é retirada de uma distribuição global caracterizada pelo desempenho histórico da política.
• Bases Inteligentes (Smart Baselines): A base para um prompt $q$ é uma média ponderada entre a média local do grupo ($\mu_{group}$) e a média global da política ($\mu_{glob}$):
  $$V^{EB}_q = (1 - S_q)\mu_{group} + S_q\mu_{glob}$$
  Onde $S_q$ é o fator de encolhimento, determinado pela razão entre a variância dentro do grupo e a variância entre grupos.
• Atualização Online: As estatísticas globais ($\mu_{glob}$, variâncias) são atualizadas dinamicamente usando o Algoritmo de Welford, garantindo escalabilidade e estabilidade sem a necessidade de armazenar grandes buffers de dados.
• Amostragem Agrupada (Clustered Sampling): Para lidar com a heterogeneidade dos dados (ex: misturar aritmética simples com cálculo complexo), o EBPO organiza o fluxo de treinamento em clusters coesos (por tópico ou dificuldade). Isso permite que o estimador online se adapte a distribuições locais consistentes, reduzindo o erro de estimativa da prior.

3. Contribuições Teóricas

Os autores fornecem garantias teóricas rigorosas para o EBPO em comparação com o GRPO:

• Redução de Erro Quadrático Médio (MSE): O EBPO garante um MSE estritamente menor na estimativa da dificuldade do prompt em comparação com a média amostral não enviesada do GRPO.
• Gradientes Não-Desvanecidos: Em cenários de falha saturada (todas as recompensas zero), o EBPO produz um gradiente não nulo e informativo. O modelo recebe uma penalidade adaptativa baseada na dificuldade global da tarefa, evitando que o passo de treinamento seja desperdiçado.
• Conservação de Entropia: O EBPO regulariza a covariância entre a probabilidade da política e a vantagem, limitando o decaimento da entropia. Isso previne o colapso prematuro da política e encoraja uma exploração mais robusta.

4. Resultados Experimentais

O EBPO foi avaliado em diversos benchmarks de raciocínio matemático (AIME-2024/2025, OlympiadBench, MATH-500, AMC23) usando modelos de diferentes escalas (LLaMA3-8B, Qwen3-8B, Qwen3-14B).

Principais Achados:

• Desempenho Superior: O EBPO superou consistentemente o GRPO e outras variantes (como DAPO e Dr. GRPO). No modelo Qwen3-8B, o EBPO alcançou uma melhoria média de mais de 5% no Pass@1 em comparação ao GRPO.
• Eficiência de Amostra (Small Group Sizes): O EBPO demonstrou ser altamente eficiente em cenários com recursos limitados. Com um tamanho de grupo pequeno ($G=8$), o EBPO superou o GRPO em 11,28% em média. Isso permite treinar com grupos menores sem sacrificar a estabilidade.
• Estabilidade de Treinamento:
  • O EBPO manteve normas de gradiente robustas mesmo quando o GRPO apresentava sinais desvanecidos.
  • A divergência KL por passo (medida de mudança na política) foi estritamente limitada no EBPO, prevenindo colapsos catastróficos observados no GRPO em estágios tardios.
  • A entropia da política foi mantida mais alta, indicando melhor exploração.
• Aprendizado de Currículo: A combinação do EBPO com amostragem baseada em dificuldade (do fácil ao difícil) resultou em ganhos adicionais significativos, especialmente em benchmarks de elite como AIME e OlympiadBench.

5. Significância e Impacto

O EBPO representa um avanço significativo na otimização de políticas para LLMs com recompensas verificáveis:

1. Solução para Instabilidade: Resolve o problema fundamental da variância em grupos pequenos e do colapso de gradiente em falhas totais, permitindo treinamento mais estável e eficiente.
2. Redução de Custos: Ao permitir o uso de grupos menores ($G$) sem perda de desempenho, reduz drasticamente o custo computacional de treinamento de RLVR.
3. Generalização: A abordagem de "emprestar força" das estatísticas globais via inferência Bayesiana Empírica é um novo paradigma que pode ser aplicado além do raciocínio matemático, beneficiando qualquer tarefa de RL onde os sinais de recompensa são esparsos ou ruidosos.

Em suma, o EBPO demonstra que ancorar o aprendizado por reforço em raciocínios semânticos compartilhados (via estatísticas globais e agrupamento) oferece um caminho teoricamente robusto e computacionalmente eficiente para alcançar desempenho de nível de elite em raciocínio matemático.