Uniform-Correct Policy Optimization: Breaking RLVR's Indifference to Diversity

1. Declaração do Problema

Aprendizado por Reforço com Recompensas Verificáveis (RLVR) melhorou significativamente a precisão de uma única tentativa (Pass@1) de Modelos de Linguagem de Grande Escala (LLMs) em tarefas de raciocínio, como matemática. No entanto, um efeito colateral crítico foi identificado: Colapso de Diversidade.

• O Fenômeno: Embora o RLVR aumente a probabilidade de gerar uma resposta correta, ele frequentemente reduz drasticamente o Pass@K (a probabilidade de que pelo menos uma das $K$ respostas amostradas seja correta). Isso indica que o espaço de soluções do modelo contrai-se, concentrando a massa de probabilidade em um subconjunto estreito de soluções corretas, enquanto suprime outros caminhos de raciocínio válidos.
• A Causa Raiz: Os autores identificam uma limitação estrutural nos objetivos padrão de RLVR (como GRPO - Group Relative Policy Optimization). Esses objetivos são indiferentes à forma como a massa de probabilidade é distribuída entre as soluções corretas.
  • Matematicamente, qualquer política que coloque 100% de sua massa no conjunto de saídas corretas $Y^+$ rende a mesma recompensa máxima, independentemente de essa massa estar concentrada em uma única solução (colapsada) ou distribuída uniformemente entre todas as soluções válidas (diversa).
  • Combinado com a dinâmica de amostragem on-policy, essa indiferença cria um ciclo de retroalimentação autorreforçador: soluções amostradas ligeiramente mais frequentemente recebem atualizações de gradiente maiores, aumentando ainda mais sua probabilidade e frequência de amostragem, levando a um colapso de modo irreversível.

2. Metodologia

A. Análise Teórica

Os autores formalizam o mecanismo de colapso e definem a estrutura da política ótima:

1. Mecanismo de Colapso: Eles provam que, sob o GRPO, o viés de amostragem leva a contagens de amostras desiguais. As atualizações de gradiente escalonam com essas contagens, amplificando a probabilidade de soluções dominantes e suprimindo as minoritárias. Uma vez que uma solução cai abaixo de um limiar de amostragem ($\pi(y) \ll 1/K$), ela não recebe sinal de gradiente e não pode se recuperar.
2. Definição da Política Ótima: Os autores propõem que a política ótima deve ser a Política Uniformemente-Correta (UCP). Esta política satisfaz dois critérios:
   • Robustez: Maximiza a massa de probabilidade retida no pior caso se um subconjunto de soluções corretas for desabilitado (por exemplo, devido a uma mudança de distribuição).
   • Optimalidade Regularizada por Entropia: Maximiza exclusivamente o objetivo de equilibrar correção com diversidade (entropia) dentro do conjunto correto.
   • Definição: A UCP coloca probabilidade zero em saídas incorretas e distribui a probabilidade uniformemente entre todas as soluções corretas ($\pi(y|x) = 1/|Y^+|$ para $y \in Y^+$).

B. Solução Proposta: Uniform-Correct Policy Optimization (UCPO)

Para alcançar a Política Uniformemente-Correta, os autores propõem o UCPO, uma modificação ao objetivo do GRPO.

• Função Objetivo: O UCPO amplia a maximização padrão de correção com uma penalidade de uniformidade condicional.
  $$\mathcal{L}_{UCPO} = \underbrace{\log Z_\theta(x)}_{\text{Massa de Correção}} - \tau \underbrace{KL(u(\cdot|x) \parallel q_\theta(\cdot|x))}_{\text{Penalidade de Uniformidade}}$$

  • $Z_\theta(x)$: Massa total de probabilidade nas saídas corretas.
  • $q_\theta(\cdot|x)$: A distribuição condicional da política sobre as saídas corretas.
  • $u(\cdot|x)$: A distribuição uniforme sobre o conjunto de saídas corretas.
  • $\tau$: Um hiperparâmetro que controla a força da penalidade de uniformidade.

• Mecanismo de Gradiente:

  • O GRPO padrão pondera os gradientes pela frequência de amostragem atual (favorecendo modos dominantes).
  • O UCPO repondera o sinal de gradiente para contrabalançar esse viés. Ele interpola entre a distribuição condicional atual e o alvo uniforme.
  • Ponderação Inversa: O estimador de gradiente atribui pesos maiores a soluções corretas sub-representadas (aquelas com baixa probabilidade atual) e pesos menores às dominantes. Isso redistribui efetivamente o "orçamento de vantagem" para amplificar respostas corretas minoritárias sem sacrificar a massa total de correção.

3. Contribuições Principais

1. Identificação da Indiferença: O artigo prova formalmente que os objetivos padrão de RLVR são indiferentes à diversidade intra-conjunto-correto, permitindo que a dinâmica estocástica de treinamento leve a um colapso de modo irreversível.
2. Caracterização Teórica: Estabelece a Política Uniformemente-Correta como a solução ótima única sob critérios de robustez e regularização por entropia.
3. Design de Algoritmo (UCPO): Deriva o UCPO, uma modificação simples, porém eficaz, ao GRPO que impõe a Política Uniformemente-Correta como seu ótimo único.
4. Redistribuição de Gradiente: Introduz um mecanismo para redistribuir sinais de gradiente em direção a respostas corretas sub-representadas, quebrando o ciclo de colapso autorreforçador.

4. Resultados Experimentais

Os autores avaliaram o UCPO em três modelos (1,5B, 7B parâmetros) e cinco benchmarks de raciocínio matemático (AIME24/25, AMC, MATH, OlympiadBench).

• Melhoria Pass@K: O UCPO melhora consistentemente o Pass@K (especificamente Pass@64) mantendo um Pass@1 competitivo.
  • AIME24 (modelo 7B): Alcançou uma melhoria absoluta de +10% no Pass@64 em comparação ao GRPO padrão.
  • Melhoria Média: Em todos os modelos, o UCPO melhorou o Pass@64 médio em +3,5 a +4,8 pontos sobre o GRPO.
• Métricas de Diversidade:
  • O UCPO alcançou até 45% de diversidade maior ao nível de equação dentro do conjunto correto em comparação ao GRPO.
  • Diferentemente da regularização global por entropia (que frequentemente espalha massa para respostas incorretas), o UCPO mantém alta massa de correção ($Z_\theta \approx 1$) enquanto diversifica o conjunto correto.
• Comparação com Baselines: O UCPO superou outros métodos de preservação de diversidade, como Regularização por Entropia (Ent-Reg), controle de Covariância KL e Focal GRPO. A regularização global por entropia mostrou-se frágil e frequentemente degradou o Pass@1 ou falhou em melhorar a diversidade dentro do conjunto correto.
• Eficiência: O UCPO introduz sobrecarga computacional negligenciável (apenas reponderação escalar de vantagens) e não requer rollouts adicionais ou modelos de referência.

5. Significado

Este trabalho aborda uma falha fundamental nos paradigmas atuais de treinamento RLVR para modelos de raciocínio. Ao mudar o foco de apenas "encontrar uma resposta correta" para "manter uma distribuição diversa de todas as respostas corretas", o UCPO:

• Previne a Contração do Espaço de Soluções: Garante que os modelos não "esqueçam" caminhos de raciocínio válidos durante o ajuste fino.
• Aumenta a Robustez: Uma política que cobre múltiplas soluções corretas é mais robusta a mudanças de distribuição ou ruído do verificador.
• Fornece um Alvo Principiado: Vai além de bônus de entropia heurísticos para um objetivo fundamentado teoricamente que visa explicitamente a estrutura da política ótima.

Em resumo, o UCPO oferece uma solução escalável e teoricamente sólida para o problema do colapso de diversidade no RLVR, permitindo que LLMs alcancem maior cobertura de múltiplas amostras sem sacrificar a precisão de uma única tentativa.