Resumo Técnico: Reutilize seus FLOPs: Escalando RL em Problemas Difíceis através do Condicionamento em Prefixos Muito Off-Policy

1. Declaração do Problema

O aprendizado por reforço (RL) tornou-se o método de fato para aumentar as capacidades de raciocínio de Grandes Modelos de Linguagem (LLMs), particularmente em matemática e codificação. No entanto, os métodos de RL on-policy padrão enfrentam um gargalo crítico quando aplicados a problemas difíceis onde a taxa de acerto do modelo base é próxima de zero (ex: pass@2k ≈ 0).

Neste regime, a política de RL raramente amostra um traço de raciocínio correto. Consequentemente:

• Sinais de Aprendizado Vanishing (Desaparecimento): As estimativas do gradiente da política tornam-se zero ou próximas de zero porque todas as trajetórias amostradas recebem recompensa zero.
• Desperdício de Computação: O sistema despende quantidades massivas de FLOPs em amostragem sem receber qualquer sinal de aprendizado, causando a estagnação do treinamento e o platô das recompensas.
• Ineficiência das Soluções Existentes:
  • Supervised Fine-Tuning (SFT) em traços off-policy: Embora isso possa impulsionar o modelo, frequentemente leva ao colapso da entropia e à memorização, dificultando a exploração subsequente.
  • RL Off-Policy Direto: Usar traços off-policy diretamente via ponderação de importância sofre de estimativas de gradiente de alta variância e instabilidade.

O desafio central é como reutilizar a computação "desperdiçada" de execções de inferência ou rodadas de RL anteriores (que podem conter traços corretos raros) para guiar a política de RL online sem introduzir as instabilidades associadas ao direcionamento direto off-policy.

2. Metodologia: PrefixRL

Os autores propõem o PrefixRL, um framework que condiciona o processo de RL on-policy em prefixos de traços off-policy bem-sucedidos, em vez de usar esses traços como alvos de supervisão direta.

Mecanismo Central

1. Coleta de Dados: Um conjunto de dados de traços off-policy ($D_{off}$) é coletado, tipicamente via rejeição de amostragem em larga escala no modelo base até que um traço correto seja encontrado para cada problema.
2. Aumentação de Problemas: Para cada problema original $x$, o sistema extrai um prefixo (os primeiros $h$ tokens) de um traço off-policy correto $y_x$. Este prefixo é anexado ao prompt do problema original para criar um problema prefixado: $\text{concat}(x, y_{x, 1:h})$.
3. Objetivo de Treinamento: O modelo é treinado usando RL on-policy em uma mistura de:
   • Problemas originais (sem prefixo).
   • Problemas prefixados.
   • Crucialmente: Os gradientes são mascarados nos tokens do prefixo off-policy. O modelo é atualizado apenas nos tokens que ele gera após o prefixo.

Garantias Teóricas

O artigo fornece provas teóricas em relação ao objetivo do PrefixRL:

• Consistência: Se os prefixos off-policy forem corretos e realizáveis por uma política dentro da classe do modelo, qualquer maximizador do objetivo PrefixRL também é um maximizador do objetivo de RL padrão. O método não enviesa a solução ótima.
• Eficiência de Amostra: O PrefixRL reduz o gap de subotimalidade em comparação com o RL padrão. Ao colocar a política em estados de maior recompensa (estados onde uma solução correta já está parcialmente formada), ele reduz a carga de exploração. Limites teóricos sugerem uma redução no número de amostras on-policy necessárias, escalando favoravelmente com o comprimento do contexto.

Fenômeno Chave: Back-Generalization (Retro-Generalização)

Empiricamente, os autores descobrem a back-generalization: treinar apenas em problemas prefixados aumenta significativamente o desempenho nos problemas originais, sem prefixo, embora o modelo nunca tenha sido explicitamente treinado nas versões sem prefixo durante essa fase.

• Mecanismo: O modelo não simplesmente imita o prefixo. Em vez disso, ele aprende estratégias que se generalizam para o estado sem prefixo.
• Descoberta de Estratégia: O modelo pode "desaprender" estratégias subótimas presentes no prefixo e descobrir novas estratégias melhores. Ele altera a distribuição do próximo token nos estados não vistos, efetivamente realizando um "backtracking" a partir dos estados de alta recompensa fornecidos pelo prefixo para encontrar caminhos melhores no espaço sem prefixo.
• Relacionabilidade: A back-generalization é mais eficaz quando o prefixo (ou o problema em contexto) é semanticamente relacionado ao problema alvo.

3. Principais Contribuições

1. Framework PrefixRL: Um método inovador para reciclar FLOPs off-policy condicionando em prefixos corretos, evitando as instabilidades do RL off-policy direto ou o colapso de entropia induzido por SFT.
2. Análise Teórica: Prova de que o PrefixRL é consistente com os objetivos de RL padrão e oferece ganhos prováveis de eficiência de amostra ao reduzir a carga de exploração.
3. Descoberta da Back-Generalization: Identificação de um fenômeno onde o condicionamento em prefixos pode levar a um desempenho superior em tarefas sem prefixo, muitas vezes superando a generalização simples entre tarefas relacionadas.
4. Loop de Auto-Melhoria: Demonstração de um pipeline prático onde traços off-policy são gerados via rejeição de amostragem do próprio modelo base, criando um ciclo de auto-melhoria sem modelos professores externos.

4. Resultados Experimentais

Os autores avaliaram o PrefixRL em benchmarks de raciocínio matemático difíceis (DAPO, OMNI-MATH, AIME '25, HMMT '25) usando modelos Llama-3.1-8B e Qwen-3-4B.

• Eficiência de Computação: O PrefixRL alcança a mesma recompensa de treinamento 2× mais rápido que a base mais forte (SFT em dados off-policy seguido de RL), mesmo após contabilizar o custo computacional da rejeição de amostragem inicial.
• Ganhos de Acurácia:
  • Em problemas de treinamento difíceis, o PrefixRL aumenta a acurácia final em >45% (mais de 3× relativo) em comparação com a base.
  • Em benchmarks não vistos como AIME '25, o PrefixRL melhora o pass@1 em 12% sobre a base mais forte em uma comparação de computação equivalente.
• Robustez à Origem: O método permanece eficaz mesmo quando os prefixos off-policy são originários de uma família de modelos diferente (ex: usando prefixos Qwen para treinar uma política Llama), demonstrando flexibilidade em cenários práticos.
• Dinâmica de Treinamento:
  • Estabilidade: Ao contrário do SFT+RL, o PrefixRL preserva a entropia a nível de token, evitando o colapso da exploração.
  • Qualidade do Gradiente: O PrefixRL exibe normas de gradiente mais altas e menor variância em comparação com o RL off-policy padrão (que sofre de pesos de importância de alta variância).
  • Eficiência: O modelo gera respostas mais curtas em média, indicando que atinge passos decisivos mais cedo e reduz o "vagueio improdutivo".

5. Significância e Alegações

O artigo afirma que o PrefixRL muda fundamentalmente a abordagem de escalar RL em problemas difíceis. Em vez de ver os dados off-policy como um alvo para imitação (o que leva à instabilidade ou colapso) ou descartá-los como ruído, o PrefixRL os trata como um mecanismo de condicionamento para guiar a exploração.

A significância reside em:

• Reuso de FLOPs: Transforma a computação anteriormente desperdiçada (tentativas falhas que eventualmente renderam um traço correto via rejeção de amostragem) em um sinal direto para aprendizado.
• Solução do Regime de "Problemas Difíceis": Permite o treinamento de RL em problemas onde os rollouts on-policy quase nunca têm sucesso, um regime onde o RL padrão estagna.
• Back-Generalization: A descoberta de que o condicionamento em prefixos pode levar a um desempenho superior em tarefas sem prefixo sugere um mecanismo poderoso para alavancar a computação prévia que vai além da simples generalização, potencialmente remodelando como futuros algoritmos de RL incorporam dados históricos.

Os autores concluem que aproveitar a back-generalization é um ingrediente crucial para impulsionar a computação prévia e permitir a auto-melhoria contínua em LLMs de raciocínio.