Resumo Técnico: REFT – Exploração de Trajetórias com Diversificação do Primeiro Token

Declaração do Problema

A Aprendizagem por Reforço com Recompensas Verificáveis (RLVR) emergiu como um método primário para aprimorar as capacidades de raciocínio de modelos de linguagem grandes (LLMs) sem depender de trajetórias de raciocínio rotuladas. Em métodos de RLVR agrupados, como GRPO e DAPO, uma política gera um conjunto finito de trajetórias (trilhas de raciocínio) para um único prompt, que são pontuadas por um verificador automático. O sinal de aprendizado é derivado das recompensas relativas dentro desse grupo.

Um gargalo crítico nesse paradigma é a diversidade das trajetórias. Se todas as trajetórias em um grupo seguirem caminhos de raciocínio semelhantes ou receberem recompensas idênticas, o grupo fornece pouco sinal contrastivo para atualizações da política. Abordagens existentes tentam ampliar a exploração aumentando a temperatura de amostragem, ramificando-se em pivôs de raciocínio de alta entropia ou selecionando respostas finais diversas. No entanto, esses métodos frequentemente assumem que os tokens iniciais do prefixo são andaimes fixos ou que a exploração valiosa reside apenas em etapas de raciocínio substantivas de alta incerteza.

Os autores identificam uma ineficiência estrutural: o primeiro token gerado após o marcador de raciocínio (<thought>) é tipicamente um iniciador de discurso de baixo conteúdo (por exemplo, "Primeiro," "Vamos," "Para"). Políticas de amostragem padrão exibem uma distribuição fortemente acentuada sobre esses tokens, enviesando pesadamente o grupo de trajetórias para um único iniciador. Crucialmente, os autores observam que, embora a probabilidade desses primeiros tokens seja altamente concentrada, a correção das trajetórias resultantes é quase plana entre os 20 principais candidatos. Isso cria uma oportunidade de "baixa carga, alta alavanca": o primeiro token tem carga semântica mínima, mas exerce uma alta alavanca distribucional, roteando toda a continuação subsequente para diferentes regiões semânticas.

Metodologia: REFT

O artigo introduz o REFT (Exploração de Trajetórias com Diversificação do Primeiro Token), uma intervenção mínima no pipeline de RLVR que diversifica o primeiro token sem alterar a arquitetura do modelo, o verificador, a função de recompensa ou o objetivo de RL.

Mecanismo Central

O REFT substitui a amostragem estocástica padrão do primeiro token por uma estratégia de amostragem estratificada:

1. Seleção de Candidatos: Para um prompt dado $x$, a distribuição da política na posição do primeiro token é avaliada. Os $N$ tokens semânticos válidos superiores (excluindo espaços em branco e tokens especiais) são identificados como o conjunto de candidatos $F_N(x)$.
2. Amostragem Uniforme: De $F_N(x)$, $K$ primeiros tokens distintos são amostrados uniformemente sem reposição para formar um conjunto $S_K(x)$.
3. Alocação de Orçamento: O tamanho total do grupo de trajetórias $G$ é dividido igualmente entre os $K$ tokens selecionados. Para cada primeiro token selecionado $f \in S_K(x)$, o modelo gera $G/K$ continhações condicionadas a $f$, utilizando o decodificador de base inalterado para o restante da sequência.
4. Treinamento: O grupo diversificado resultante é processado pelo algoritmo padrão de RLVR (por exemplo, GRPO ou DAPO) sem correções de ponderação de importância para a discrepância de amostragem do primeiro token.

Escolhas de Design Chave

• Baixa Carga, Alta Alavanca: O método visa uma posição onde o modelo está "artificialmente certo" (alta probabilidade para um iniciador específico), mas onde a correção é estável entre alternativas.
• Sem Ruído Global: Ao contrário do aumento da temperatura, que achata distribuições em toda a trajetória (potencialmente prejudicando etapas de raciocínio críticas posteriores), o REFT injeta diversidade apenas no primeiro token, preservando a integridade da cadeia de raciocínio subsequente.
• Substituição Direta: O REFT não requer alterações no objetivo de RL subjacente, na estimativa de vantagem ou na lógica do verificador.

Contribuições Principais

1. Insight Diagnóstico: Os autores demonstram que a distribuição do primeiro token em modelos de raciocínio é fortemente concentrada, mas desacoplada da correção. Eles mostram que primeiros tokens de baixa probabilidade dentro dos 20 principais podem levar a trajetórias corretas com a mesma frequência que o token mais provável, desde que a continuação seja amostrada normalmente.
2. Efeito de Roteamento: Experimentos confirmam que diferentes primeiros tokens roteiam o modelo para regiões distintas de continuação semântica. A amostragem uniforme dos 20 primeiros tokens superiores produz diversidade semântica significativamente maior na continuação restante (após remover o primeiro token) em comparação com a amostragem padrão.
3. Algoritmo REFT: Um método simples e computacionalmente eficiente para diversificar grupos de trajetórias, realocando o orçamento fixo de trajetórias entre iniciadores de primeiro token plausíveis.
4. Mitigação do Excesso de Crédito: O artigo mostra que as atualizações padrão do GRPO tendem a afiar a preferência pelo primeiro token de posição 1 ao longo das etapas de treinamento, mesmo quando o verificador não apoia fortemente essa preferência. O REFT mitiga isso garantindo a cobertura de múltiplos iniciadores plausíveis, impedindo que a política atribua crédito excessivo a uma preferência discursiva específica.

Resultados Experimentais

Os autores avaliaram o REFT em quatro modelos base (Qwen2.5-0.5B, 3B, 7B e Llama3.2-3B) e dois objetivos de RLVR (GRPO e DAPO) em três conjuntos de dados (GSM8K, BigMath-Easy, BigMath).

• Ganhos de Desempenho: O REFT melhora consistentemente o desempenho agregado nas métricas Pass@1, Pass@8 e Pass@64.
  • No GSM8K, o GRPO e o DAPO aumentados com REFT superam suas linhas de base.
  • Na "Média Matemática" (média de MATH-500, AIME 2024/2025, AMC 2023), os ganhos são particularmente pronunciados, sugerindo que o REFT é eficaz quando uma exploração mais ampla é crítica.
  • Melhorias são observadas mesmo na menor escala (0,5 bilhão de parâmetros), indicando que o método não depende de grande capacidade do modelo.
• Análise de Diversidade:
  • Diversidade de Treinamento: O REFT produz os grupos de trajetórias mais semanticamente diversos durante o treinamento.
  • Concentração de Inferência: Paradoxalmente, modelos treinados com REFT produzem saídas mais concentradas no momento da inferência, sugerindo que o método ajuda a política a aprender a convergir para padrões bem-sucedidos de forma mais eficaz após explorar caminhos diversos.
  • Decomposição de Variância Zero: O REFT reduz a fração de grupos "todos errados" (grupos de variância zero onde nenhuma trajetória correta é encontrada) sem aumentar a saturação "todos corretos", indicando uma exploração mais útil.
• Eficiência: O REFT introduz sobrecarga de tempo de execução negligenciável. Em alguns casos, reduz ligeiramente as horas de GPU devido a comprimentos médios de conclusão mais curtos resultantes de maior precisão.

Significado e Alegações

O artigo afirma que o REFT oferece uma "simples inversão da intuição usual de exploração". Em vez de buscar diversidade em bifurcações de raciocínio de alta entropia e incertas, a diversidade útil pode ser recuperada de escolhas de prefixo de baixa carga que são fortemente enviesadas, mas fracamente ligadas à correção.

Os autores enfatizam que o REFT não é um novo algoritmo de RL, mas uma substituição direcionada para a primeira decisão do amostrador de trajetórias. Seu significado reside em:

• Eficiência: Amplia o sinal de treinamento sem aumentar o orçamento de trajetórias ou o custo computacional.
• Robustez: Melhora tanto a precisão de amostra única (Pass@1) quanto a probabilidade de encontrar uma solução correta dentro de um orçamento finito (Pass@8/64), abordando o trade-off frequentemente observado no RLVR onde a melhoria do Pass@1 estreita a distribuição acessível.
• Generalizabilidade: O método funciona em diferentes escalas de modelo e objetivos de RLVR, sugerindo que o fenômeno de "roteamento do primeiro token" é uma característica fundamental dos modelos de raciocínio atuais que pode ser explorada para um melhor treinamento.

O artigo conclui que o primeiro token após o marcador de raciocínio é um ponto de intervenção inesperadamente barato e eficaz para a exploração em RLVR, permitindo que os modelos explorem regiões de continuação mais amplas enquanto mantêm a estabilidade da trilha de raciocínio.