Resumo Técnico: Otimização de Política Pass@K (PKPO)

Definição do Problema

Algoritmos de Aprendizado por Reforço (RL) aplicados a Grandes Modelos de Linguagem (LLMs) frequentemente amostram múltiplas tentativas de solução ($n > 1$) para um único problema, mas otimizam com base no desempenho de amostras individuais (pass@1). Essa abordagem prioriza a utilidade de uma única amostra de alta recompensa em detrimento da utilidade coletiva do conjunto, efetivamente subutilizando a capacidade de amostragem. Consequentemente, métodos padrão de gradiente de política podem falhar em explorar o espaço de soluções suficientemente, particularmente em tarefas mais difíceis onde a probabilidade de uma única amostra correta é baixa. Embora existam métodos de busca em tempo de inferência, o acoplamento entre a inferência e as atualizações do modelo sugere que a otimização ingênua da recompensa esperada de uma única amostra é subótima para tarefas onde a correção pode ser verificada, mas as soluções não são conhecidas a priori.

Metodologia

Os autores propõem a Otimização de Política Pass-at-k (PKPO), um framework que otimiza diretamente a métrica pass@k (ou seu análogo contínuo, maxg@k), definida como a esperança do máximo de recompensa sobre $k$ amostras independentes extraídas da política.

Contribuições Teóricas Principais

A principal contribuição é a derivação de novos estimadores de baixa variância e não enviesados tanto para o objetivo pass@k quanto para o seu gradiente. Esses estimadores funcionam através da média de estimadores simples sobre todas as $\binom{n}{k}$ subconjuntos de tamanho $k$ de um lote de $n$ amostras.

1. Recompensas Binárias (Seção 2):

   • Para rótulos de correção binários ($f \in \{0, 1\}$), os autores derivam um estimador não enviesado $\rho(n, c, k)$ para o pass@k, onde $c$ é a contagem de amostras corretas.
   • Eles provam que um estimador não enviesado existe apenas se $n \ge k$.
   • Um estimador de gradiente correspondente não enviesado é derivado, atribuindo pesos às amostras com base em sua correção e na probabilidade combinatória de sua inclusão em um subconjunto bem-sucedido.

2. Recompensas Contínuas (Seção 3):

   • O framework generaliza para recompensas contínuas $g(x)$ via o objetivo maxg@k: $E[\max_{i=1}^k g(x_i)]$.
   • Ao assumir que as recompensas estão ordenadas ($g_1 \le g_2 \le \dots \le g_n$), os autores derivam um estimador não enviesado $\rho^{(g)}(n, c, k)$ que pondera as recompensas ordenadas usando coeficientes binomiais.
   • Um estimador de gradiente não enviesado é construído, onde o peso $s_i$ para uma amostra é uma combinação linear de recompensas, computada eficientemente em tempo $O(k + n \log n)$.

3. Redução de Variância (Seção 4):

   • Para reduzir a variância das estimativas de gradiente sem introduzir viés, os autores propõem baselines de Leave-One-Out (LOO).
   • Eles introduzem uma baseline específica $s^{(loo-1)}_i$ que subtrai o maxg@k de subconjuntos de tamanho $k-1$ (excluindo a amostra atual) do maxg@k de subconjuntos de tamanho $k$. Isso generaliza trabalhos anteriores (especificamente [TZSM25]), que eram restritos ao caso $n=k$.
   • Esta formulação permite uma otimização robusta mesmo quando $n > k$, um cenário comum no treinamento de LLMs onde grandes lotes são amostrados.

Implementação

O método é projetado como um "substituto direto" (drop-in replacement) para transformações de recompensa padrão em qualquer algoritmo de gradiente de política (ex: PPO). Dado um lote de $n$ amostras e suas recompensas brutas, as recompensas são transformadas usando os estimadores derivados (especificamente a função sloo_minus_one no código fornecido) antes de serem enviadas para a etapa de atualização da política.

Resultados Principais

Os autores validam o PKPO através de experimentos de brinquedo (toy experiments) e benchmarks do mundo real usando GEMMA2 (parâmetros 2B e 9B) e LLAMA3.1 (8B parâmetros).

1. Experimentos de Brinquedo:

   • Em um experimento de política Gaussiana 1D, o PKPO consegue deslocar a média da política ótima para ser mais "tolerante ao risco" para um $k$ mais alto, permitindo que o modelo explore regiões com menor probabilidade, mas maior potencial de recompensa.
   • A análise de variância confirma que o estimador sloo_minus_one proposto reduz significativamente a variância do gradiente em comparação com a partição ingênua ou estimadores sem baselines.

2. Benchmarks do Mundo Real (MATH, Coding, ARC-AGI-1):

   • Otimização Seletiva: Definir o alvo de otimização $k_{opt}$ igual ao métrica de avaliação $k_{eval}$ produz o melhor desempenho. Por exemplo, otimizar para pass@8 melhora significativamente os escores de avaliação de pass@8 em comparação com a otimização para pass@1.
   • Resolvendo Problemas Mais Difíceis: No desafiador benchmark ARC-AGI-1, a otimização convencional de pass@1 estagna, falhando em melhorar a taxa de resolução cumulativa. Em contraste, o PKPO com $k_{opt} > 1$ "desbloqueia" o aprendizado, alcançando taxas de resolução e pontuações pass@k mais altas para todos os $k$.
   • Estratégia de Recozimento (Annealing): Os autores demonstram que o recozimento de $k$ durante o treinamento (começando com um $k_{opt}$ alto para incentivar a exploração, depois reduzindo para $k_{opt}=1$ para consolidar) alcança um forte desempenho de pass@1 junto com ganhos significativos de pass@k, evitando o trade-off típico onde a melhoria do pass@k degrada o pass@1.
   • Comparação: O PKPO supera o caso especial de $n=k$ (como em [TZSM25]) e baselines de regularização de entropia, particularmente em tarefas difíceis onde a exploração é crítica.

Significância e Alegações

O artigo afirma que o PKPO fornece um método robusto, eficiente e teoricamente fundamentado para otimizar a utilidade coletiva de um conjunto de amostras, em vez de apenas a média ou a melhor amostra individual.

• Desbloqueio da Exploração: Os autores argumentam que, ao priorizar a utilidade conjunta (a probabilidade de pelo menos uma de $k$ amostras ter sucesso) sobre a utilidade da amostra individual, o método melhora a exploração. Isso é crucial para resolver problemas mais difíceis onde o "sinal" (uma solução correta) é esparso no espaço de busca.
• Flexibilidade: Ao contrário de métodos anteriores que eram restritos a $k=n$ (otimizando para o melhor de todo o lote) ou $k=1$, o PKPO permite a otimização robusta para qualquer $k$ arbitrário $\le n$.
• Impacto Prático: O método permite que praticantes ajustem o trade-off entre exploração e explotação via parâmetro $k_{opt}$, permitindo que modelos resolvam tarefas complexas de raciocínio e codificação que de outra forma seriam inacessíveis à otimização padrão de pass@1.

O trabalho posiciona o PKPO como uma generalização e redução de variância de métodos existentes de gradiente de política de estatística de ordem, oferecendo um caminho prático para escalar o treinamento de RL para LLMs em tarefas que exigem exploração profunda.