Resumo Técnico: Raciocínio Aprimorado de LLMs por Otimização de Funções de Recompensa com Aprendizado por Reforço Orientado a Busca

Declaração do Problema

O raciocínio matemático serve como um marco crítico para Modelos de Linguagem Grandes (LLMs), exigindo não apenas a geração fluente de texto, mas consistência intermediária, transporte de quantidades e respostas finais precisas. Embora o Aprendizado por Reforço (RL) tenha surgido como um mecanismo padrão de pós-treinamento para alinhar LLMs a comportamentos desejados, sua eficácia é altamente sensível ao design da função de recompensa. Abordagens atuais frequentemente dependem de especificação manual e orientada por especialistas da função de recompensa, o que é propenso a erros, requer tentativa e erro significativo e pode inadvertidamente induzir "hacking de recompensa" ou comportamentos frágeis. Além disso, a literatura existente trata amplamente a função de recompensa como infraestrutura fixa (por exemplo, regras criadas manualmente ou modelos de recompensa treinados) em vez de um objeto de otimização em si. Este artigo aborda a questão de se a especificação de recompensa usada durante o pós-treinamento por RL pode ser sistematicamente pesquisada e melhorada, particularmente em um cenário onde o modelo base é mantido fixo e a função de recompensa é a variável de design primária.

Metodologia

Os autores propõem um framework orientado a busca que trata a especificação da função de recompensa como um problema de otimização iterativa. O pipeline opera em cinco rodadas, gerando e avaliando funções de recompensa candidatas usando um modelo base fixo e um LLM de fronteira como gerador.

Componentes Principais

1. Modelo Base e Configuração de Treinamento:

   • Modelo: Llama-3.2-3B-Instruct da Meta, ajustado finamente usando Adaptação de Baixo Rank (LoRA) com rank de 32.
   • Algoritmo: Otimização de Política Relativa em Grupo (GRPO), que substitui a rede crítica separada no PPO pela normalização relativa ao grupo das amostras geradas, reduzindo a sobrecarga de memória e melhorando o raciocínio matemático.
   • Conjunto de Dados: GSM8K (7.473 exemplos de treinamento, 1.319 exemplos de teste).
   • Formato de Saída: Saída estruturada com tags <thinking> e <solution> para separar o raciocínio da resposta final.

2. Recompensas Base Fixas:
   Três recompensas base protegidas são incluídas em toda execução de treinamento e não podem ser substituídas por candidatos gerados:

   • match_format_exactly: Recompensa o pareamento correto das tags <thinking> e <solution>.
   • match_format_approximately: Fornece sinais graduais para presença/ausência de tags.
   • check_answer: Avalia a resposta numérica extraída contra a verdade fundamental com pontuação em níveis (correspondência exata, normalização, verificações de razão).

3. Loop de Geração Orientado a Busca:

   • Gerador: Um LLM de fronteira (Kimi K2) gera 10 funções de recompensa candidatas por rodada (50 no total ao longo de 5 rodadas).
   • Prompting: O gerador recebe 20 exemplos aleatórios do GSM8K, uma especificação de assinatura e um resumo classificado do desempenho de rodadas anteriores para orientar a geração e evitar duplicação semântica.
   • Validação: Candidatos passam por uma validação em sandbox de quatro etapas (verificação de importação AST, execução bloqueada, namespace restrito e chamada de sonda) para garantir segurança e executabilidade.
   • Triagem: Candidatos válidos são treinados por 500 passos GRPO no conjunto de treinamento do GSM8K juntamente com as três recompensas base.
   • Classificação: Modelos são avaliados no conjunto de teste do GSM8K usando a pontuação F1 (média harmônica de precisão e recall) como métrica principal de classificação. O resumo classificado é alimentado de volta na próxima rodada de geração.

4. Construção de Ensemble:
   Após a busca, as recompensas de melhor desempenho são combinadas em ensembles. Sete configurações foram testadas, variando em tamanho (3 a 10 recompensas descobertas) e estratégia de seleção (melhor classificada vs. baseada em diversidade entre rodadas). A recompensa do ensemble é a soma elemento a elemento das três recompensas base e das recompensas descobertas selecionadas.

Contribuições Principais

1. Formulação da Busca de Recompensa: O artigo formula a especificação de recompensa para raciocínio de LLMs como um loop iterativo gerar–validar–treinar–classificar sobre código Python executável, tratando a própria recompensa como um objeto de busca.
2. Integração GRPO: O framework aplica este loop de busca especificamente sob GRPO, um otimizador de política computacionalmente eficiente e relevante para raciocínio, demonstrando que a descoberta de recompensa é viável dentro das restrições de ajuste fino eficiente em parâmetros.
3. Análise de Ensemble: O estudo examina se a combinação de múltiplas recompensas descobertas produz sinais mais fortes do que recompensas individuais, encontrando que ensembles de 5 ou mais recompensas superam significativamente recompensas únicas ou configurações apenas com base.
4. Rigor Empírico: O artigo fornece uma análise abrangente incluindo intervalos de confiança bootstrap, testes de McNemar pareados com correção de Bonferroni e auditorias específicas para hacking de recompensa (por exemplo, preenchimento de linhas) e conformidade de formato.

Resultados

• Trajetória de Busca: Ao longo de cinco rodadas, a pontuação F1 média das recompensas geradas aumentou de 0,596 (Rodada 1) para 0,632 (Rodada 5). A recompensa individual superior (thinking_steps_count) alcançou um F1 de 0,787.
• Desempenho do Ensemble: A melhor configuração de ensemble (ens_diverse_rounds, combinando 3 recompensas base + 5 recompensas descobertas diversas) alcançou um F1 de 0,795 e uma precisão de 0,660 (IC 95% [0,635, 0,686]).
• Significância Estatística: Testes de McNemar pareados com correção de Bonferroni ($\alpha \approx 0,0024$) mostraram que todos os ensembles com 5 ou mais recompensas descobertas foram estatisticamente indistinguíveis entre si, mas significativamente melhores que a linha de base apenas com base (F1 = 0,609) e ensembles menores.
• Experimentos de Controle:
  • Linha de Base Aleatória: Um ensemble de 5 recompensas sorteado aleatoriamente colapsou para um F1 de 0,047, demonstrando que o ganho é impulsionado pelo loop de feedback classificado, e não meramente pelo sinal aditivo de ter mais recompensas.
  • Robustez da Semente: O re-treinamento do melhor ensemble com três sementes diferentes resultou em pontuações F1 de $0,785 \pm 0,013$, confirmando que a ordenação qualitativa dos resultados é robusta à variância de treinamento.
• Auditoria de Hacking de Recompensa: A análise da recompensa estrutural superior (thinking_steps_count) mostrou que respostas corretas naturalmente continham mais linhas de raciocínio e maior densidade aritmética do que as incorretas, sugerindo que o modelo não explorou a recompensa via preenchimento de linhas vazias.
• Comparação com Linha de Base: O modelo Llama-3.2-3B treinado com GRPO superou cinco modelos prontos para uso na faixa de 2B–4B parâmetros (incluindo variantes Qwen e Phi) tanto em precisão estrita quanto em F1.

Significado e Alegações

O artigo alega que a especificação de recompensa no pós-treinamento por RL pode ser tratada como um objeto pesquisável de forma confiável, automatizando um processo tipicamente dependente de intuição humana.

• Acessibilidade: Todo o pipeline roda em uma única GPU de nível de consumidor (NVIDIA RTX 5090) em aproximadamente 40 horas, tornando a descoberta automatizada de recompensas acessível a orçamentos acadêmicos e de pequena indústria.
• Auditabilidade: Diferentemente de modelos de recompensa de caixa preta, as recompensas descobertas são código Python legível que os praticantes podem auditar, modificar e combinar.
• Generalizabilidade: Embora avaliado no GSM8K, os autores postulam que a "receita" se transfere para outros domínios de raciocínio (por exemplo, geração de código, resolução de problemas científicos) e oferece uma alternativa prática ao treinamento caro de modelos de recompensa.
• Modéstia: Os autores afirmam explicitamente que a contribuição é metodológica, e não uma alegação de um novo estado da arte no GSM8K em si. A evidência primária é interna ao pipeline: a melhoria iterativa das recompensas, o colapso dos controles aleatórios e a robustez da ordenação do ensemble entre sementes. O artigo não alega que as recompensas descobertas são universalmente ótimas, mas sim que o processo de busca identifica efetivamente sinais estruturais de alto desempenho para um modelo base fixo.