Resumo Técnico: SAGE – Moldando Âncoras para Exploração Guiada em RLVR de LLMs

1. Declaração do Problema

Avanços recentes em Aprendizado por Reforço com Recompensas Verificáveis (RLVR) demonstraram a capacidade de melhorar a precisão pass@1 de Modelos de Linguagem de Grande Escala (LLMs) em tarefas de raciocínio. No entanto, uma limitação crítica emergiu: enquanto o RLVR frequentemente aumenta as taxas de sucesso de amostras únicas, ele frequentemente falha em produzir ganhos comparáveis no pass@k (a probabilidade de encontrar uma resposta correta entre $k$ amostras).

Análises anteriores sugerem que essa discrepância surge porque o RLVR padrão induz colapso de modos (ou "encolhimento de suporte"). Em vez de descobrir novas estratégias de raciocínio, a política colapsa em um subconjunto estreito de modos de raciocínio pré-existentes e de alta densidade presentes no modelo de referência base. Esse comportamento é atribuído às restrições estruturais do termo de regularização Kullback-Leibler reverso (reverse-KL) padrão. Embora o reverse-KL forneça estabilidade de treinamento necessária ao ancorar a política a uma distribuição de referência ($\pi_{ref}$), ele inerentemente suprime trajetórias de baixa densidade, impedindo efetivamente que o modelo explore caminhos de raciocínio alternativos e válidos que estão fora do suporte imediato do modelo de referência.

Tentativas de resolver isso removendo o termo KL ou substituindo-o por forward-KL provaram-se subótimas. Remover o termo KL frequentemente leva a hacking de recompensa e instabilidade, enquanto o forward-KL, embora expanda o suporte, frequentemente aloca massa de probabilidade para regiões irrelevantes para a recompensa, falhando em melhorar o pass@1. O desafio central é reter os benefícios estabilizadores do reverse-KL enquanto permite expansão empírica controlada do suporte.

2. Metodologia: Framework SAGE

Os autores propõem o SAGE (Moldando Âncoras para Exploração Guiada), um framework que reinterpreta a regularização reverse-KL não como uma restrição estática a ser enfraquecida, mas como uma âncora dinâmica cuja distribuição-alvo pode ser remodelada para guiar a exploração.

Mecanismo Central

O RLVR reverse-KL padrão otimiza:
$$J(\theta) = \mathbb{E}_{x,y \sim \pi_\theta} [r(x, y) - \beta D_{KL}(\pi_\theta \parallel \pi_{ref})]$$
Isso leva a uma solução estacionária onde a política é um inclinação exponencial da referência: $\pi_\theta \propto \pi_{ref} \exp(r/\beta)$. Essa atualização multiplicativa amplifica regiões de alta densidade de $\pi_{ref}$ enquanto suprime aquelas de baixa densidade.

O SAGE modifica isso introduzindo uma função guia $q(x, y) > 0$ que remodela a própria distribuição da âncora. O novo objetivo torna-se:
$$J'(\theta) = \mathbb{E}_{x,y \sim \pi_\theta} [r(x, y) - \beta \tilde{D}_{KL}(\pi_\theta \parallel q \cdot \pi_{ref})]$$
Aqui, $\tilde{D}_{KL}$ é uma divergência pseudo-KL definida para funções não normalizadas. A solução estacionária torna-se:
$$\pi_\theta(y | x) \propto q(x, y) \pi_{ref}(y | x) \exp\left(\frac{1}{\beta} r(x, y)\right)$$
Ao projetar $q(x, y)$ para ser grande para modos de raciocínio subexplorados mas potencialmente válidos, o SAGE amplifica sua massa de probabilidade sem quebrar a estabilidade fornecida pela âncora $\pi_{ref}$.

Instanciação via Sinais Intrínsecos

Para tornar $q(x, y)$ prático sem exigir informações de recompensa explícitas durante a construção do guia, os autores propõem o uso de sinais intrínsecos derivados do próprio processo de raciocínio:

1. Exploração Aleatória: Uma linha de base usando perturbações estocásticas.
2. Exploração Nível de Token: Usa surpresa ($-\log p$) para enviesar escolhas léxicas raras. Os autores observam que isso captura desvios de superfície, mas pode não refletir diferenças estruturais de raciocínio.
3. Exploração Nível de Ramo (SAGE-Branch): A variante principal proposta. Usa entropia ($H_t$) como sinal para ramificação de raciocínio. A função guia é definida como:
   $$q_{Branch}(y_t | x, y_{<t}) = 1 + \gamma (H_t - \tau)_+$$
   onde $H_t$ é a entropia nível de token, $\tau$ é um limiar de ativação e $\gamma$ controla a magnitude. Isso amplifica seletivamente a exploração em pontos de decisão de alta incerteza (nós de ramificação) enquanto suprime ruído em passos simbólicos determinísticos.

Garantias Teóricas

O artigo fornece uma análise teórica (Teorema 4.3) demonstrando que, ao contrário do RLVR padrão, que é limitado pelo suporte empírico do modelo de referência, o SAGE pode teoricamente expandir o suporte empírico. Especificamente, se a função guia $q$ for suficientemente grande para uma solução válida em termos de recompensa $y^*$ que está atualmente fora do suporte empírico da política de referência, o SAGE pode atribuir a ela uma probabilidade não desprezível, efetivamente recuperando modos de raciocínio "ocultos".

3. Contribuições Principais

1. Análise das Limitações do RLVR: Os autores identificam que a âncora reverse-KL, embora estabilizadora, é a causa estrutural primária do colapso de modos no RLVR. Eles argumentam que simplesmente remover ou substituir o termo KL é insuficiente, pois interrompe o trade-off entre eficiência e cobertura.
2. Framework SAGE: Eles propõem um método principled para permitir expansão controlada do suporte empírico remodelando a distribuição da âncora KL via uma função guia $q(x, y)$, em vez de remover a âncora.
3. Integração de Sinais Intrínsecos: Eles demonstram como sinais leves e intrínsecos (especificamente entropia e surpresa) podem ser usados para construir funções guia eficazes que não exigem passagens adiantadas adicionais ou dados externos.
4. Validação Empírica: Experimentos extensivos mostram que variantes do SAGE, particularmente a exploração nível de ramo, melhoram consistentemente tanto o pass@1 quanto o pass@k em diversos benchmarks de raciocínio matemático.

4. Resultados Experimentais

Os autores avaliaram o SAGE nos conjuntos de dados AIME 2024-2025, AMC23 e MATH-500 usando o modelo Qwen2.5-Math-7B-Base e o algoritmo GRPO.

• Ganhos de Desempenho: A variante GRPO + Branch alcançou o desempenho médio mais forte, melhorando o pass@1 de 0,266 (GRPO puro) para 0,284 e o pass@256 de 0,714 para 0,729. Ela superou as linhas de base em todos os três conjuntos de dados.
• Comparação com Alternativas:
  • GRPO sem KL: Mostrou recompensas de treinamento mais altas, mas desempenho significativamente pior no tempo de teste (pass@1 e pass@256), confirmando o risco de hacking de recompensa sem a âncora KL.
  • Forward-KL: Melhorou ligeiramente o pass@256, mas não conseguiu igualar os ganhos de pass@1 do SAGE, indicando que aloca massa para regiões fora do alvo.
  • Outros Algoritmos de RL: O SAGE foi integrado com sucesso ao BNPO e ao DAPO, resolvendo consistentemente o trade-off onde esses métodos tipicamente melhoravam o pass@1 à custa do pass@k.
• Análise de Padrões de Raciocínio: Análise qualitativa revelou que, enquanto o RLVR padrão reforçava padrões dominantes como "configuração de restrições", o SAGE aumentou significativamente a frequência de prova por contradição (de 0,8% no GRPO para 3,3% no GRPO+Branch), um modo de raciocínio raro mas válido.
• Robustez: Em avaliações fora da distribuição (benchmark Knights and Knaves), modelos treinados com SAGE mantiveram o desempenho sob mudança de distribuição, enquanto modelos GRPO padrão exibiram degradação significativa, sugerindo que o SAGE mitiga o colapso de modos.
• Generalização: O método também foi eficaz em modelos de raciocínio longo (DeepSeek-R1-Distill-Qwen-7B) e modelos menores (Qwen2.5-1.5B/3B), demonstrando robustez em diferentes escalas de modelo e paradigmas de raciocínio.

5. Significado e Alegações

O artigo afirma que o SAGE oferece uma abordagem prática e principled para mitigar o afiamento de modos no RLVR reverse-KL. Ao tratar a regularização KL como um instrumento controlável em vez de uma restrição rígida, o SAGE permite que LLMs descubram e utilizem trajetórias de raciocínio diversas que são válidas, mas sub-representadas no modelo base.

Os autores enfatizam que o SAGE é modular e ortogonal a técnicas existentes (como moldagem de vantagem ou amostragem dinâmica), pois modifica a própria distribuição da âncora em vez da estrutura de otimização. Eles posicionam o SAGE como um framework geral para expansão empírica estável e compatível com KL.

Limitações e Escopo:
Os autores modestamente reconhecem que o SAGE opera dentro do suporte da política de referência fixa. Ele redistribui massa de probabilidade entre trajetórias que já têm probabilidade não nula sob $\pi_{ref}$; ele não pode gerar padrões de raciocínio inteiramente novos que sejam impossíveis para o modelo base produzir. Eles observam que, embora a destilação de modelos professores mais fortes possa expandir a fronteira de raciocínio, o SAGE complementa isso permitindo exploração estável dentro dessa fronteira. A instanciação atual depende de sinais intrínsecos leves, e os autores sugerem que aprender funções guia mais expressivas a partir de modelos mais fortes é uma direção promissora para trabalhos futuros.