Resumo Técnico: Otimização de Rollout Adaptativa entre Épocas para Pós-Treinamento de RL (CERO)

1. Formulação do Problema

O pós-treinamento de Grandes Modelos de Linguagem (LLMs) via Aprendizado por Reforço (RL), particularmente utilizando algoritmos como o GRPO, depende tipicamente da amostragem de um número fixo de rollouts por prompt. Esta estratégia de alocação uniforme ignora a heterogeneidade nos sinais de treinamento fornecidos por diferentes prompts. Alguns prompts são triviais (gerando sucesso ou falha consistentes com baixa variância), enquanto outros são difíceis demais ou já dominados, oferecendo retornos marginais decrescentes. Consequentemente, uma parte significativa do orçamento global de rollout é desperdiçada em prompts com valor de aprendizado limitado.

O artigo formula isso como um problema de alocação de recursos online entre épocas (cross-epoch) sob um orçamento global fixo $B$. Diferente de métodos anteriores que realizam realocação dentro do lote (within-batch) ou filtragem pós-geração, este trabalho aborda a alocação de rollouts tanto entre prompts quanto entre épocas de treinamento. O objetivo é maximizar a utilidade cumulativa das amostras coletadas, onde a utilidade está sujeita a retornos decrescentes ao nível do prompt. O problema é matematicamente definido como a maximização da soma das funções de utilidade específicas de cada prompt $U_i(\sum_k N_{i,k})$ sujeita à restrição de orçamento global $\sum_{i,k} N_{i,k} \le B$, onde $N_{i,k}$ é o número de rollouts atribuídos ao prompt $i$ na época $k$.

2. Metodologia: CERO

Os autores propõem o CERO (Cross-Epoch Adaptive Rollout Optimization), um algoritmo primal-dual bayesiano online projetado para alocar rollouts adaptativamente com base na estimativa de informatividade do prompt.

2.1 Informatividade e Design de Utilidade

• Modelagem de Posterior Beta: Para cada prompt $x_i$, o CERO mantém uma distribuição posterior Beta sobre sua probabilidade de sucesso latente $p_i$, atualizada com base no histórico de rollouts $H_i^k$.
• Pontuação de Informatividade ($q_{k,i}$): Em vez de usar uma estimativa pontual da probabilidade de sucesso, o CERO estima a variância de Bernoulli esperada: $q_{k,i} = E[p_i(1-p_i) | H_i^k]$. Esta pontuação é maximizada quando $p_i \approx 0.5$ (dificuldade intermediária) e desaparece para prompts fáceis ($p_i \approx 1$) ou difíceis ($p_i \approx 0$), identificando efetivamente os prompts que oferecem os sinais de aprendizado mais incertos e, portanto, mais informativos.
• Função de Utilidade: Uma função de utilidade côncava e de saturação é definida como $U_i(n) = 1 - \exp(-\eta q_i n)$, onde $n$ é o número cumulativo de rollouts. Isso captura os retornos decrescentes: os primeiros rollouts em um prompt informativo são altamente valiosos, mas o valor marginal diminui à medida que mais amostras são coletadas.

2.2 Otimização via Reformulação Fenchel-Dual

O acoplamento temporal da função de utilidade (que depende das alocações cumulativas) torna a otimização direta online difícil. O CERO aborda isso através de uma reformulação Fenchel-dual:

• O problema é transformado em um objetivo dual envolvendo variáveis duais ao nível do prompt $\theta_i$ e um multiplicador de orçamento global $\mu$.
• Esta reformulação decompõe o problema temporalmente acoplado em subproblemas de alocação por época, coordenados pelas variáveis duais.

2.3 Algoritmo Primal-Dual Online

O CERO opera como uma camada de inserção (drop-in) sobre pipelines de RL existentes (ex: GRPO):

1. Decisão de Alocação: Em cada época, para cada prompt, o algoritmo resolve um subproblema linear. Um prompt recebe o máximo de rollouts permitidos ($N_{max}$) se o seu valor marginal estimado (variável dual $\theta_i$) exceder o preço do orçamento global atual ($\mu$); caso contrário, recebe zero.
2. Atualização de Feedback: Após gerar respostas e observar recompensas, a posterior Beta para o prompt é atualizada, e a pontuação de informatividade $q_{k,i}$ é renovada.
3. Atualizações de Variáveis Duais:
   • Ao nível do prompt ($\theta$): Atualizada via gradiente descendente online projetado para equilibrar os rollouts alocados contra a utilidade implicada pela estimativa de informatividade atual.
   • Ao nível do orçamento ($\mu$): Atualizada com base na discrepância entre o consumo real de rollouts e a taxa de orçamento alvo. Se o consumo exceder a taxa, $\mu$ aumenta, tornando a alocação futura mais seletiva; se estiver abaixo, $\mu$ diminui.

3. Principais Contribuições

1. Formulação do Problema: O artigo é o primeiro a formular a alocação de rollout de pós-treinamento de LLM como um problema de otimização online de orçamento fixo entre épocas com retornos decrescentes específicos de cada prompt, acoplando explicitamente as decisões ao longo do horizonte de treinamento via um orçamento global compartilhado.
2. Design do Algoritmo: Introduz o CERO, um algoritmo primal-dual bayesiano que utiliza estimativas de posterior do valor de rollout do prompt para guiar o orçamento. Utiliza a variância de Bernoulli posterior esperada como um substituto (proxy) para a capacidade de aprendizado e emprega uma reformulação Fenchel-dual para lidar com a não-separabilidade temporal.
3. Validação Teórica e Empírica:
   • Teoria: Sob utilidades de prompt fixas, os autores provam um limite de regret de $O(\sqrt{K})$ contra um benchmark de alocação offline, onde $K$ é o número de épocas.
   • Experimentos: Experimentos extensos em tarefas de raciocínio matemático (AIME, AMC) em quatro LLMs de código aberto (variando de 1.5B a 7B parâmetros) demonstram que o CERO supera consistentemente o GRPO vanilla.

4. Resultados Experimentais

Os autores avaliaram o CERO no dataset DAPO-Math-17K utilizando DeepSeek-R1-Distill-1.5B, Qwen3-4B (Base e Instruct) e Qwen2.5-Math-7B.

• Desempenho: O CERO alcançou melhorias consistentes sobre o GRPO vanilla em todos os modelos e benchmarks. Os ganhos médios variaram de +1,53 a +6,04 pontos, dependendo do modelo.
• Análise de Eficiência: A "Razão de Prompt Eficaz" (a fração de prompts em um lote contendo respostas bem-sucedidas e falhas, que geram gradientes não nulos) foi significativamente maior para o CERO do que para o GRPO.
  • No baseline GRPO, a razão de prompt eficaz caiu drasticamente conforme o treinamento progredia, indicando que o modelo havia dominado muitos prompts, levando a grupos de rollout degenerados.
  • O CERO manteve uma razão de prompt eficaz mais alta ao desviar adaptativamente o orçamento de prompts dominados para aqueles com maior incerteza.
• Sobrecarga (Overhead): O mecanismo de atualização de alocação introduz uma sobrecarga computacional desprezível (menos de um segundo por atualização).

5. Significância e Alegações

O artigo afirma que o orçamento adaptativo de rollout é uma direção fundamentada e eficaz para melhorar a eficiência de amostragem no pós-treinamento de RL. Ao se afastar da alocação uniforme, o CERO garante que o orçamento de geração fixo seja gasto em prompts que fornecem os sinais de treinamento mais informativos.

Os autores posicionam o CERO como uma camada de coleta de dados adaptativa de "inserção direta" (drop-in) que modifica quais prompts recebem rollouts e quantos, sem alterar a regra de gradiente de política subjacente. O trabalho destaca que a ineficiência do pós-treinamento de RL atual provém não apenas da regra de atualização da política, mas da estratégia de coleta de dados suboptimal.

Limitações observadas pelos autores:

• A garantia teórica de regret assume utilidades de prompt fixas, enquanto a implementação utiliza estimativas de posterior dinâmicas por época. Assim, a análise não captura totalmente a não-estacionariedade induzida pelas atualizações da política e renovações de posterior.
• A validação empírica é atualmente limitada a tarefas de raciocínio matemático; a validação mais ampla em outros domínios e modelos de recompensa é sugerida como trabalho futuro.