Where Should RL Post-Training Compute Go? Model Size, Search, Learning, and Feedback

Resumo Técnico: Para Onde Deve Ir o Computo de Pós-Treinamento de RL?

Declaração do Problema
O pós-treinamento por Aprendizado por Reforço (RL) é cada vez mais utilizado para adaptar modelos de fundação para raciocínio, planejamento e pipelines de aprendizado de robótica baseados em feedback. No entanto, as análises de escalonamento atuais tipicamente resumem o custo desses processos usando uma única métrica agregada: o total de FLOPs. Este artigo argumenta que tal reporte é subespecificado. Sob um orçamento fixo de pós-treinamento, a alocação de computo é um problema de alocação de recursos restrito que envolve quatro dimensões acopladas:

1. Tamanho do Modelo: A capacidade da política sendo treinada.
2. Busca/Rollout: A geração de trajetórias candidatas.
3. Aprendizado/Atualizações: A otimização da política baseada em gradiente.
4. Feedback/Avaliação de Recompensa: O custo de inferência de modelos de recompensa ou verificadores.

A questão central abordada é: Sob um orçamento fixo de FLOPs, como o computo deve ser distribuído entre esses componentes para maximizar o desempenho? Os autores postulam que políticas maiores consomem mais computo por token, comprando, portanto, menos rollouts e passos de atualização sob o mesmo orçamento. Além disso, a introdução de mecanismos de feedback aprendidos (ex: Modelos de Recompensa de Processo ou PRMs) introduz um novo custo de computo que compete com os orçamentos de rollout e atualização, um trade-off obscurecido pelo reporte de FLOPs totais.

Metodologia
O estudo introduz um framework de contabilidade de FLOPs especificamente para o GRPO (Group Relative Policy Optimization) de pós-treinamento. O computo total ($C_{total}$) é decomposto em três termos distintos:
$$C_{total} = C_{search} + C_{learning} + C_{reward}$$
Onde:

• $C_{search}$: Geração de rollout autorregressiva.
• $C_{learning}$: Otimização da política (contabilizando os parâmetros do adaptador LoRA e os custos do backbone congelado).
• $C_{reward}$: Inferência de recompensa ou verificador.

A variável de alocação primária é a fração de atualização ($\rho$), definida como a razão entre o computo de aprendizado e o computo total. Os autores conduzem experimentos usando políticas Qwen2.5 adaptadas com LoRA no dataset de raciocínio matemático Polaris-53K. O design experimental inclui:

• Estágio A: Variar tamanhos de modelo (1.5B, 3B, 7B) para estudar a alocação busca-aprendizado.
• Estágio B & D: Comparar diferentes sistemas de recompensa (esparsa, estruturada, densa, proxy-PRM e PRM real) para avaliar como o custo de feedback desloca as fronteiras de alocação.
• Avaliação: As métricas incluem recompensa nativa de treinamento, precisão pass@1 downstream no GSM8K e MATH-500, e uma pontuação de qualidade de processo de um "juiz comum".

Para lidar com o alto custo da busca exaustiva, os autores introduzem o RACE (Reward-Aware Compute Allocation), um protocolo de piloto-grade diagnóstico. O RACE executa pequenas grades piloto IsoFLOP para estimar as respostas locais de alocação de recompensa e identificar se uma configuração pertence a um regime de "busca intensa", "atualização intensa" ou "recompensa intensa". Os autores declaram explicitamente que o RACE é uma ferramenta diagnóstica para identificação de regimes, não uma garantia de melhoria em dados não vistos (held-out).

Principais Contribuições

1. Framework de Contabilidade de FLOPs: Uma decomposição formal do computo de pós-treinamento de RL em componentes de rollout, atualização e avaliação de recompensa, permitindo a comparação de execuções que são mecanisticamente diferentes, apesar de possuírem idênticos FLOPs totais.
2. Acoplamento de Tamanho de Modelo e Alocação: Evidência empírica mostrando que a escolha do modelo e a alocação de treinamento são inseparáveis. Sob um orçamento fixo, políticas maiores compram inerentemente menos atualizações e rollouts, alterando a fronteira de alocação ótima.
3. Dependência de Recompensa e Alvo: Demonstração de que a "melhor" alocação depende do sistema de recompensa (baseado em regras vs. PRM) e do alvo de avaliação (recompensa nativa vs. precisão downstream vs. qualidade de processo). Por exemplo, o feedback baseado em PRM desloca a fração de atualização ótima ($\rho$) para baixo, porque uma parte do orçamento é consumida pela inferência do modelo de recompensa.
4. Protocolo RACE: Um método diagnóstico para identificar regimes de alocação via pequenas grades piloto, destinado a guiar execuções de validação em vez de substituí-las.

Resultos

• Fronteiras de Alocação Condicionais: A alocação ótima não é um ponto fixo, mas uma fronteira condicional. Configurações de busca intensa subutilizam a capacidade de aprendizado, enquanto configurações de atualização intensa tornam-se limitadas pela exploração.
• Impacto do Sistema de Recompensa: Recompensas baseadas em regras (esparsa, estrutída, densa) tendem a favorecer regimes de alta atualização ($\rho \approx 0.72$). Em contraste, o feedback do tipo PRM (proxy e real) desloca a fronteira de recompensa nativa para frações de atualização mais baixas ($\rho \approx 0.44–0.45$), porque o orçamento fixo deve agora cobrir a inferência do modelo de recompensa.
• Desacordo de Alvos: Diferentes alvos de avaliação selecionam diferentes alocações ótimas.
  • Recompensa Nativa: Favorece a alocação que maximiza o sinal de feedback específico usado durante o treinamento.
  • Pass@1 Downstream: Frequentemente favorece regimes de alta atualização (ex: a configuração $a180$) através de diversos sistemas de recompensa.
  • Juiz Comum: Frequentemente atinge o pico em regimes de atualização média (ex: $a100$) para recompensas densas e PRM.
  • Este desacordo indica que uma receita que maximiza uma métrica pode não maximizar outra.
• Desempenho do RACE: O RACE identifica com sucesso o regime correto (atualização alta vs. baixa) em configurações de validação com alta precisão (taxa de 1.0 de mesmo regime), mas não garante uma melhoria global na recompensa de dados não vistos (held-out) em relação aos baselines.

Significância e Alegações
O artigo alega modestamente que o pós-treinamento de RL é um problema de alocação de recursos restrito, em vez de uma simples lei de escala (scaling law). Os autores argumentam que reportar apenas os FLOPs totais esconde os trade-offs mecanísticos entre capacidade do modelo, geração de experiência e avaliação de feedback.

A significância reside na recomendação para a comunidade reportar computo consciente da alocação: os FLOPs totais junto com o detalhamento específico de rollout, atualização e avaliação de recompensa, além da fração de atualização resultante $\rho$. Essa transparência é crucial para sistemas de aprendizado de robótica, onde melhorar um módulo de raciocínio ou de feedback via pós-treinamento de RL envolve decidir se deve investir recursos limitados em modelos maiores, busca mais extensa ou mecanismos de feedback mais fortes. O artigo conclui que não existe uma "melhor" receita universal; a alocação ótima depende conjuntamente do tamanho do modelo, do design da recompensa e do alvo de avaliação específico que está sendo otimvelizado.