Resumo Técnico: SARA (Alocação de Rollout Adaptativa Sequencial)

Definição do Problema

O Aprendizado por Reforço com Recompensas Verificáveis (RLVR) é atualmente limitado pelo custo da geração de rollouts. Em estimadores baseados em grupos, como o Group Relative Policy Optimization (GRPO), a contribuição de um prompt para o gradiente da política depende da variância das recompensas dentro de seu grupo amostrado. Se um grupo estiver "saturado" (todas as respostas estão corretas ou todas estão incorretas), a variância da recompensa é zero, resultando em uma vantagem normalizada evanescente e sem sinal de aprendizado.

Métodos existentes para mitigar esse desperdício enfrentam um trade-off:

1. Avaliar-então-filtrar (ex: Dynamic Sampling/DS): Estes métodos sobreamostram um grande conjunto de candidatos, geram grupos completos para todos e descartam os saturados. Embora isso garanta um lote limpo de grupos eficazes, incorre em um custo de rollout massivo (frequentemente 4$\times$ ou mais do que a amostragem uniforme) porque se paga pela geração completa de prompts que serão eventualmente descartados.
2. Prever-então-selecionar: Estes métodos estimam a dificuldade do prompt antes da amostragem para priorizar prompts promissores. Embora evitem rollouts extras, dependem de previsões que podem ser frágeis quando a política muda rapidamente, levando a lotes poluídos se as previsões forem imprecisas.

Ambas as abordagens decidem ao nível do prompt antes de observar a dinâmica interna do grupo. No entanto, o artigo observa que a eficácia de um grupo é frequentemente decidida precocemente dentro da própria sequência de seus rollouts. Gastar o orçamento de um grupo completo em um prompt que já revelou que será saturado é computacionalmente dispendioso.

Metodologia: SARA

Os autores propõem o SARA (Sequential Adaptive Rollout Allocation), que reformula a coleta de rollout por etapa como um problema de alocação sequencial com restrição de orçamento (parada ótima). Em vez de gerar um número fixo de rollouts ($k$) para cada prompt, o SARA sonda os prompts em rodadas em lote, atualizando crenças e tomando decisões com base nos resultados observados.

Mecanismos Principais

1. Modelagem Bayesiana: Para cada prompt $q$, o SARA mantém uma distribuição posterior Beta sobre sua taxa de sucesso latente $\gamma_q$. Inicialmente, utiliza-se uma priori uniforme. Após observar $n$ rollouts com $s$ sucessos, a posterior é atualizada para $\text{Beta}(\alpha_0 + s, \beta_0 + n - s)$.
2. Preditor de Eficácia de Forma Fechada: O SARA calcula a probabilidade preditiva posterior ($p_{\text{eff}}$) de que um grupo de tamanho $k$ seja "eficaz" (resultados mistos) dado o prefixo atual.
   • Se o prefixo já for misto ($1 \le s \le n-1$), $p_{\text{eff}} = 1$.
   • Se o prefixo for todo-falha ou todo-sucesso, $p_{\text{eff}}$ é calculado analiticamente usando a função Beta. Para uma priori uniforme e um prefixo de todos-falhas, isso simplifica para $p_{\text{eff}}(n, 0) = \frac{k-n}{k+1}$.
3. Regra de Parada de Dois Limiares: Com base em $p_{\text{eff}}$, o SARA aplica uma regra de decisão sequencial reminiscente do Teste de Razão de Verossimilhança Sequencial (SPRT) de Wald:
   • COMMIT (Comprometer): Se o grupo for misto (eficaz), ele é adicionado ao lote de treinamento imediatamente.
   • ABANDON (Abandonar): Se $p_{\text{eff}}$ cair abaixo de um limiar inferior $\tau_{\text{low}}$, o prompt é considerado provavelmente saturado. O orçamento restante para este prompt é liberado.
   • CONTINUE (Continuar): Caso contrário, o prompt recebe mais um rollout.
4. Realocação de Orçamento: O orçamento liberado dos prompts abandonados é imediatamente realocado para novos prompts do pool. Isso garante que um orçamento total fixo gere mais grupos eficazes do que a alocação uniforme.

Propriedades Algorítmicas

• Ortogonalidade: O SARA opera na etapa de coleta de rollout, tornando-se compatível com qualquer estratégia de seleção de prompt (ex: pode ser composto com o Dynamic Sampling).
• Sem Rollouts Extras: Ao contrário de métodos preditivos que exigem chamadas de modelos auxiliares para estimar a dificuldade, o SARA utiliza apenas os rollouts que o otimizador geraria de qualquer maneira.
• Sincronização: O algoritmo roda em lotes síncronos por rodada para manter o throughput de inferência, tipicamente exigindo apenas 2 a 4 rodadas de sincronização por etapa.

Principais Contribuições

1. Reenquadramento do Problema: Os autores identificam a "decidibilidade precoce" da eficácia do grupo e recategorizam a coleta de rollout como um problema de alocação sequencial, distinto da seleção de prompt.
2. Algoritmo SARA: Eles derivam um preditor Beta-Binomial de forma fechada e uma regra de parada de dois limiares, criando um alocador livre de predição-e-rollout que se integra aos pipelines de GRPO existentes.
3. Garantias Teóricas:
   • Confiabilidade do Abandono: A probabilidade de abandonar incorretamente um grupo eficaz é limitada pelo limiar $\tau_{\text{low}}$.
   • Economia de Rollouts: O número esperado de rollouts gastos por prompt é estritamente menor que o $k$ fixo usado no Dynamic Sampling, com economias que aumentam conforme o tamanho do grupo $k$ cresce.
   • Dominância de Rendimento: Com um orçamento fixo, o SARA garante um número maior ou igual de grupos eficazes comparado à alocação uniforme.
   • Elo com o Gradiente: Maximizar o rendimento de grupos eficazes maximiza diretamente um limite inferior da norma do quadrado do gradiente esperado do GRPO.
4. Validação Empírica: Experimentos extensos em tarefas de raciocínio matemático e planejamento usando modelos de 1.5B e 3B.

Resultados Experimentais

Avaliado em uma única GPU com modelos R1-Distill-Qwen-1.5B e Qwen2.5-3B em datasets como MATH, AIME24 e Countdown:

• Eficiência vs. Dynamic Sampling (DS): O SARA iguala a precisão do Dynamic Sampling (que usa um oráculo para filtrar grupos saturados) enquanto utiliza 22% menos rollouts.
• Composição com Seleção Preditiva: Combinar o SARA com o Dynamic Sampling (SLA + DPS) produz a melhor precisão, superando ligeiramente o oráculo do DS, enquanto utiliza 67% menos rollouts que o DS.
• Economia de Tokens: Como os traços de "todos-falha" abandonados tendem a ser os mais longos, as economias de tokens são ainda mais pronunciadas do que as de rollouts.
• Robustez: Ao contrário da seleção preditiva, que degrada conforme a política muda, o SARA mantém uma fração de lote quase 100% eficaz durante todo o treinamento ao confiar na verificação in-sample.
• Compatibilidade: O SARA melhora o desempenho de vários algoritmos de RL (PPO, GRPO, RLOO, Reinforce++) ao substituir a coleta uniforme de rollouts.

Significância e Alegações

O artigo alega que o SARA oferece uma solução de "melhor dos dois mundos" ao eliminar a necessidade de sobreamostragem cara (como o DS) e evitar a fragilidade das previsões pré-amostragem. Ao aproveitar a evidência estatística presente dentro do próprio grupo de rollout, o SARA alcança alta eficiência de treinamento sem chamadas de modelos auxiliares.

Os autores posicionam o SARA como uma alavanca fundamental de eficiência para RLVR, particularmente à medida que os tamanhos de grupo aumentam para redução de variância. Eles observam que, embora o método assuma recompensas verificáveis binárias e rollouts i.i.d. dentro de um grupo, a lógica central de alocação sequencial é ortogonal à seleção de prompt e métodos de controle de comprimento, permitindo futuras extensões para recompensas contínuas e rollouts estruturados em árvore. O trabalho demonstra que economias significativas de computação em pós-treinamento de LLMs de raciocínio são alcançáveis através de estratégias de parada ótima, e não apenas por uma melhor curadoria de prompts.