Resumo Técnico: Gaste Suas Rollouts Onde Conta

Declaração do Problema

No pós-treinamento de Modelos de Linguagem de Grande Escala (LLMs) via Aprendizado por Reforço (RL), especificamente no cenário online e on-policy, o gargalo computacional dominante não é o passo de otimização, mas a geração de rollouts (amostragem). Métodos de otimização de política baseados em grupos, como a Otimização de Política Relativa de Grupo (GRPO), calculam vantagens comparando múltiplas rollouts geradas para o mesmo prompt. No entanto, as abordagens existentes alocam indiscriminadamente o orçamento de rollouts em todos os prompts de um lote, independentemente de seu potencial de aprendizado.

Isso leva a uma ineficiência significativa. Prompts com distribuições de recompensa colapsadas (onde todas as rollouts produzem recompensas idênticas, seja todas corretas ou todas incorretas) produzem vantagem zero e não contribuem com nenhum sinal de gradiente, ainda que consumam recursos de amostragem caros. Por outro lado, prompts com alta variância de recompensa fornecem os sinais de aprendizado mais fortes. Os métodos atuais falham em distinguir entre esses regimes online, desperdiçando computação em dados não informativos. Embora estratégias de superamostragem como o DAPO melhorem a exploração, elas escalonam os custos de amostragem linearmente sem filtrar prompts de baixo sinal.

Metodologia: Pilot-Commit

Os autores propõem o Pilot-Commit (PC), um framework de alocação de rollouts consciente do orçamento, projetado para desacoplar a avaliação de prompts da exploração. O método opera sob a descoberta teórica de que atualizações baseadas em grupos são mais eficazes em regimes de alta variância de recompensa, onde a magnitude do gradiente da política é maximizada.

O framework consiste em duas etapas por passo de treinamento:

1. Etapa Pilot: Uma fração do orçamento total de rollouts é usada para amostrar um pequeno número de rollouts ($n_{pilot}$) para cada prompt no lote de amostragem. O sistema estima a taxa de sucesso empírica (e, portanto, a variância de recompensa) para cada prompt sob a política atual.

   • Prompts são filtrados com base em limiares assimétricos. Prompts considerados "muito fáceis" (taxa de sucesso $> p_{upper}$) são removidos da amostragem futura, pois seu sinal de gradiente diminui. Prompts considerados "muito difíceis" (taxa de sucesso $< p_{lower}$) são adiados para o próximo epoch em vez de serem descartados, reconhecendo que podem se aproximar do regime de alta variância à medida que a política melhora.
   • Prompts que caem dentro da faixa de variância alvo (tipicamente próximo de $p \approx 0.5$) são selecionados para a etapa de commit.

2. Etapa Commit: O orçamento de rollouts restante ($n_{commit}$) é alocado exclusivamente aos prompts que passaram pelo filtro pilot. Esses prompts recebem a maior parte dos recursos de amostragem para gerar os dados de treinamento.

3. Treinamento: A política é atualizada usando a união das rollouts pilot e commit para os prompts selecionados.

Otimizações Chave:

• Vinculação (Binding): Para reduzir a sobrecarga de inferência, a etapa pilot do passo $t-1$ é vinculada à etapa commit do passo $t$, permitindo passagens de inferência sobrepostas e introduzindo apenas um atraso off-policy mínimo.
• Buffer de Replay: Prompts excedentes que passam pelo filtro pilot são armazenados em um buffer de replay para preencher o lote de treinamento se a etapa pilot atual produzir poucos sobreviventes, e prompts envelhecidos são removidos para evitar obsolescência.
• Estratégia de Remoção (Eviction): Prompts que são consistentemente resolvidos (todas as rollouts pilot corretas) são removidos permanentemente do pool de amostragem para evitar orçamento desperdiçado.

Contribuições Chave

• Insight Teórico: O artigo fornece uma prova analítica (Proposição 3.1) demonstrando que a magnitude do gradiente de grupo do GRPO é maximizada quando os resultados de recompensa dentro de um grupo são maximamente incertos (ou seja, alta variância). Isso estabelece a variância de recompensa como um proxy direto para a informatividade do prompt.
• Design do Framework: O Pilot-Commit introduz uma estratégia de alocação de duas etapas inovadora que realoca dinamicamente um orçamento fixo de rollouts com base em estimativas online e em tempo real da informatividade do prompt, evitando a necessidade de modelos auxiliares ou estatísticas históricas obsoletas.
• Implementação do Sistema: O método é implementado como uma nova receita no framework de RL verl, apresentando otimizações práticas como vinculação off-policy e buffers de replay para garantir escalabilidade.

Resultados Experimentais

Os autores avaliaram o Pilot-Commit em escalas de modelos variando de 1,5B a 14B parâmetros em benchmarks de raciocínio matemático (incluindo AIME, AMC, MATH 500 e DeepMath).

• Eficiência de Rollouts: O Pilot-Commit consistentemente atinge a precisão de linha de base alvo com significativamente menos rollouts cumulativas em comparação com GRPO e DAPO.
  • Em comparação com o GRPO, o PC atinge a precisão alvo 1,5× a 1,9× mais rápido em termos de rollouts cumulativas.
  • Em comparação com o DAPO, o PC é 2,3× a 4,0× mais rápido em rollouts cumulativas.
• Tempo de Relógio (Wall-Clock Time): Mesmo em ambientes altamente paralelizados (64 GPUs H200), a redução no custo de amostragem se traduz em acelerações práticas, com o PC sendo 1,21× mais rápido que o GRPO e 1,83× mais rápido que o DAPO no tempo de relógio por passo.
• Qualidade de Aprendizado: O PC mantém uma variância média de recompensa mais alta durante todo o treinamento em comparação com as linhas de base, indicando que as atualizações estão concentradas em prompts com sinais de gradiente mais fortes. Consequentemente, o PC frequentemente atinge uma precisão de pico mais alta que as linhas de base dentro do mesmo número de passos de treinamento.
• Robustez de Hiperparâmetros: Estudos de ablação mostram que o método é robusto a variações nas proporções de divisão pilot/commit e limiares de filtragem, desde que os limiares não sejam definidos em extremos que ou desperdiçam orçamento em prompts sem esperança ou descartam aqueles borderline informativos.

Significado e Alegações

O artigo afirma que a alocação de rollouts é um problema de "primeira classe" no pós-treinamento de RL. Ao deslocar o foco de "quantas rollouts gerar" para "em quais prompts gastá-las", o Pilot-Commit aborda a ineficiência fundamental da amostragem uniforme.

Os autores enfatizam que sua abordagem não altera o objetivo de treinamento nem impõe uma ordem fixa de currículo. Em vez disso, adapta-se online à política em evolução, garantindo que o orçamento computacional seja gasto onde o sinal de aprendizado é mais forte. Os resultados sugerem que, para o pós-treinamento de RL baseado em grupos, uma alocação inteligente de amostragem pode gerar ganhos substanciais de eficiência sem sacrificar, e frequentemente melhorando, o desempenho final do modelo. O trabalho é apresentado como uma melhoria complementar às otimizações de throughput em nível de sistema, focando na eficácia dos dados sendo amostrados, e não apenas na velocidade de geração.