Resumo Técnico: Sobre Estimativas de Vantagem para Gradientes de Política Max@K

Declaração do Problema

O aprendizado por reforço com recompensas verificáveis (RLVR) é um paradigma dominante para o pós-treinamento de grandes modelos de linguagem (LLMs) em tarefas de raciocínio. No entanto, esses cenários normalmente dependem de recompensas esparsas baseadas em resultados (por exemplo, um sinal binário de correto/incorreto), o que torna a exploração difícil em espaços combinatórios amplos. Embora métodos como regularização de entropia e bônus baseados em contagem existam, eles introduzem objetivos auxiliares que exigem um ajuste cuidadoso.

Uma abordagem complementar é otimizar diretamente objetivos de tempo de inferência, como pass@K (a probabilidade de pelo menos uma das $K$ conclusões amostradas estar correta) e sua generalização contínua, max@K (o valor esperado do máximo de recompensa entre $K$ amostras). Diversos trabalhos recentes propuseram estimadores de gradiente de política (PG) não enviesados para esses objetivos. No entanto, as relações entre esses estimadores permanecem obscuras. Especificamente, estimadores existentes frequentemente utilizam sinais, baselines e normalizações diferentes. Um problema crítico identificado neste trabalho é que os principais estimadores (por exemplo, baseados em Melhoria Esperada - Expected Improvement) são não enviesados como estimadores de PG, mas produzem vantagens não centralizadas. Uma vantagem não centralizada (com média não nula) corresponde a um baseline subestimado, o que pode direcionar a otimização para ações subótimas e aumentar a variância do gradiente.

Metodologia

O artigo propõe o MaxPO (Max@K Policy Optimization), um método que introduz um baseline de Deixe-Dois-Fora (L2O - Leave-Two-Out) para corrigir a estimativa de vantagem para objetivos de max@k.

1. Análise de Estimadores Existentes

Os autores revisitam o estimador de Melhoria Esperada (EI) usado em trabalhos anteriores (ex: Walder e Karkhanis [54]). Eles demonstram que, embora o estimador baseado em EI seja um estimador de PG não enviesado, o sinal de vantagem resultante $s(a)$ é estritamente não negativo (devido à operação ReLU na definição de EI). Consequentemente, $E[s(a)] > 0$, o que significa que a vantagem não é centrada. Essa falta de centralização deixa margem para redução de variância.

2. O Baseline Deixe-Dois-Fora (L2O)

Para alcançar uma vantagem exatamente centrada enquanto preserva a não tendência (unbiasedness) de PG, os autores introduzem o baseline L2O.

• O Desafio: Um baseline padrão de Deixe-Um-Fora (L1O) falha no cenário de max@K porque a recompensa de uma amostra $i$ pode influenciar implicitamente o cálculo da vantagem de outras amostras $j$ (pois $i$ atua como um comparador potencial na operação de max). Essa dependência viola a condição de independência necessária para a estimativa de PG não enviesada.
• A Solução: O baseline L2O $b_{-i}^{L2O}$ é computado pela média dos escores de Melhoria Esperada de todas as outras amostras $j$, onde cada escore é computado sobre um subconjunto de comparadores que exclui tanto a amostra alvo $i$ quanto a própria amostra $j$.
• Propriedades:
  • Independência: Ao excluir $i$ do conjunto de comparadores para todo $j$, o baseline $b_{-i}^{L2O}$ é independente da recompensa $r_i$, garantindo que o estimador de PG permaneça não enviesado.
  • Centralização: Os autores provam que a vantagem média realizada no lote é exatamente zero para cada lote amostrado: $\frac{1}{B}\sum_{i=1}^B (s_i - b_{-i}^{L2O}) = 0$.
  • Eficiência: Embora uma computação ingênua seria $O(B^3)$, os autores derivam uma implementação vetorizada usando formulações baseadas em ReLU que computa o baseline L2O em $O(B^2)$, tornando-o viável para o treinamento de LLMs acelerado por GPU.

3. Vantagem Canônica de Lote Finito

O artigo deriva uma vantagem canônica de lote finito para max@K, definida como $u_i - v_i$:

• $u_i$: O valor esperado condicional de max@K de um grupo contendo a resposta $i$.
• $v_i$: O valor esperado de max@K de "deixe-um-fora" computado de grupos excluindo $i$.
  Os autores mostram que o estimador MaxPO (EI + L2O) recupera exatamente este sinal canônico ($s_i - b_{-i}^{L2O} = u_i - v_i$). Isso fornece uma visão unificada de estimadores existentes, classificando-os como "não centralizados" (ex: apenas EI/PKPO) e "(normalizados) canônicos" (ex: MaxPO, e certas variantes analíticas de pass@K).

4. Integração no Treinamento de LLM

O MaxPO é integrado em algoritmos de RL baseados em grupos (como o GRPO). Para cada prompt, um grupo de $G$ saídas é amostrado. O baseline L2O é computado dentro de cada grupo para gerar vantagens de nível de sequência, que são então transmitidas para as perdas de nível de token.

Principais Contribuições

1. Identificação de Vantagens Não Centralizadas: O artigo identifica que os principais estimadores de PG para max@K (especificamente os baseados em EI) são não enviesados como estimadores de gradiente, mas falham em produzir vantagens centralizadas, levando a propriedades de variância subótimas.
2. Baseline Deixe-Dois-Fora (L2O): Um design de baseline inovador que preserva a não tendência de PG no cenário não linear de max@K, ao mesmo tempo que garante que a vantagem realizada do lote seja de média zero.
3. Unificação Canônica: A derivação da vantagem canônica de lote finito ($u_i - v_i$) e a prova de que o MaxPO recupera este sinal, oferecendo um framework unificado para comparar e categorizar estimadores de pass@K/max@K existentes.
4. Implementação Eficiente: Um algoritmo vetorizado $O(B^2)$ para computar o baseline L2O, permitindo a aplicação prática no pós-treinamento de LLMs.

Resultados Experimentais

Os autores validam o método através de experimentos controlados de bandit/labirinto e tarefas de raciocínio de LLM.

• Redução de Variância: Em ambientes de bandit, o baseline L2O reduz significativamente a variância do gradiente em comparação com os estimadores de EI brutos, com o efeito crescendo conforme o tamanho do espaço de ações aumenta. No treinamento de LLM (Llama-3.2-3B-Instruct), o método reduziu o proxy de variância do gradiente em 77,4% em relação ao PKPO (o baseline de apenas EI).
• Ganhos de Desempenho: Em cinco benchmarks de raciocínio matemático (AIME24, AIME25, AMC23, MATH500, Minerva), o MaxPO superou baselines fortes, incluindo PKPO, GRPO e Entropy-Adv.
  • No Qwen2.5-Math-7B, o MaxPO melhorou o pass@256 médio da tarefa em 5,2% sobre o PKPO.
  • No Llama-3.2-3B-Instruct, ele melhorou o pass@256 em 2,4% sobre o PKPO.
• Escalabilidade de Computação de Inferência: Os ganhos de desempenho do MaxPO foram mais pronunciados em níveis mais altos de computação de inferência (maior $k$ em pass@k), sugerindo que o método ajuda a política a utilizar melhor as amostras adicionais de tempo de teste.

Significância e Alegações

O artigo afirma que o MaxPO fornece uma fundamentação principista para projetar e comparar estimadores de gradiente de política para objetivos de tempo de inferência como max@K. Ao identificar a forma de vantagem canônica e corrigir o viés de estimadores existentes, o trabalho oferece um método para melhorar a eficiência e estabilidade de modelos de raciocínio de pós-treinamento.

Os autores observam modestamente que sua análise foca em estimativa on-policy de lote finito, sob amostragem i.i.d. Eles reconhecem limitações em relação ao treinamento off-policy, amostragem adaptativa e cenários com gerações correlacionadas como direções para trabalhos futuros. Além disso, destacam que, embora o método melhore a otimização, a melhoria da otimização de objetivos de raciocínio de amostragem repetida carrega impactos mais amplos, potencialmente fortalecendo modelos que poderiam ser usados de forma indevida para assistência enganosa ou resolução automatizada de problemas.