Are complicated loss functions necessary for teaching LLMs to reason?

Resumo Técnico: Simplificação de Funções de Perda para Raciocínio em LLMs

1. O Problema

Avanços recentes em Grandes Modelos de Linguagem (LLMs) destacaram a importância das técnicas de post-training (pós-treinamento) para melhorar capacidades de raciocínio e matemática. O GRPO (Group Relative Policy Optimization) emergiu como um método promissor, eliminando a necessidade de um modelo de valor (value model) ao estimar vantagens baseadas em recompensas relativas dentro de um grupo de respostas.

No entanto, a função de perda do GRPO é complexa, combinando três componentes principais:

1. Estimativa de vantagem relativa ao grupo.
2. Restrições estilo PPO (Proximal Policy Optimization), como clipping da razão de políticas.
3. Regularização KL (Kullback-Leibler).

A questão central levantada pelos autores é: todos esses componentes são realmente necessários? A complexidade atual pode ser uma "superengenharia" (overengineering), e métodos mais simples poderiam alcançar resultados iguais ou superiores com maior transparência e eficiência.

2. Metodologia

Os autores realizaram uma análise sistemática e ablativa do GRPO, isolando e removendo componentes individuais para avaliar seu impacto no aprendizado e no raciocínio matemático.

Abordagem Experimental:

• Datasets: Treinamento realizado no conjunto de dados GSM8K (matemática de nível escolar).
• Benchmarks: Avaliação em 9 benchmarks diversos, incluindo matemática em inglês (GSM8K, MATH, OlympiadBench), matemática em chinês (CMATH, CN-Middle-School) e STEM (MMLU-STEM, Gaokao2024).
• Modelos: Testes realizados em modelos de pequeno e médio porte: Qwen2.5 (0.5B e 1.5B parâmetros) e Llama3.2 (1B parâmetro).
• Configuração: Uso de LoRA (rank 128), geração de grupos de 8 completions por prompt, e recompensas baseadas em formato e correção.

Variações Analisadas:
Os autores compararam o GRPO original contra três variantes simplificadas e outros métodos:

1. GRPO-Pos (Apenas Vantagens Positivas): Remove o feedback negativo, treinando apenas nas ações que superam a linha de base (baseline).
2. RGRA (REINFORCE with Group Relative Advantage): A proposta principal. Remove as restrições estilo PPO (clipping e razão de políticas), mantendo apenas a estimativa de vantagem relativa ao grupo e a regularização KL. É essencialmente uma variante do algoritmo REINFORCE.
3. REINFORCE Direto: Remove a vantagem relativa ao grupo, treinando diretamente na recompensa bruta.
4. RAFT (Rejection Sampling Fine-Tuning): Método de reamostragem simples baseado em cross-entropy.

3. Contribuições Chave e Descobertas

O estudo identificou três conclusões fundamentais sobre a arquitetura de aprendizado por reforço para LLMs:

1. O Feedback Negativo é Indispensável:

   • Métodos que ignoram o feedback negativo (como GRPO-Pos e RAFT) levam a colapso do treinamento (training collapse), especialmente em modelos menores (0.5B).
   • O modelo tende a explorar respostas triviais e curtas para maximizar a recompensa sem aprender raciocínio real (fenômeno de reward hacking).
   • Conclusão: O aprendizado requer penalização de ações ruins, não apenas reforço de boas ações.

2. A Estimativa de Vantagem é Crucial:

   • O uso de recompensas diretas (REINFORCE puro) sem estimativa de vantagem desestabiliza o aprendizado, mesmo em modelos maiores (1.5B).
   • A normalização das recompensas dentro de um grupo (vantagem relativa) é essencial para a estabilidade.

3. Restrições Estilo PPO Não São Necessárias:

   • A remoção do clipping e da razão de políticas (componentes do PPO) não prejudica a estabilidade ou o desempenho.
   • Pelo contrário, a simplificação resulta em um método mais transparente e, em muitos casos, mais eficiente.

4. Resultados

A variante proposta, RGRA, demonstrou desempenho superior ou equivalente ao GRPO original em quase todos os cenários:

• Desempenho em Benchmarks:
  • O RGRA superou o GRPO em 17 de 27 comparações individuais nos benchmarks de matemática em inglês.
  • No benchmark de matemática em chinês, o RGRA alcançou as melhores médias de precisão (ex: 55.1% no Qwen2.5-0.5B vs 51.4% do GRPO).
  • Em benchmarks STEM, o RGRA também mostrou as maiores melhorias nos modelos Qwen.
• Estabilidade de Treinamento:
  • Enquanto GRPO-Pos e RAFT sofreram colapso (redução drástica no comprimento da resposta e recompensa), o RGRA manteve dinâmicas de treinamento estáveis, com trajetórias de recompensa comparáveis ao GRPO.
• Emergência de Comportamento de Raciocínio:
  • Modelos treinados com RGRA e GRPO desenvolveram trilhas de raciocínio explícitas (passos intermediários, reavaliação de respostas).
  • Modelos com RAFT ou GRPO-Pos falharam em gerar raciocínio, entregando apenas respostas finais diretas, muitas vezes incorretas.

5. Significado e Conclusão

O artigo desafia a suposição de que funções de perda complexas e inspiradas no PPO são obrigatórias para o alinhamento e raciocínio em LLMs.

• Simplicidade e Eficiência: O RGRA prova que é possível obter (e até superar) o estado da arte em raciocínio matemático removendo a complexidade do clipping do PPO, mantendo apenas a estimativa de vantagem relativa ao grupo.
• Transparência: Ao eliminar componentes desnecessários, o RGRA oferece uma abordagem mais transparente para entender como o aprendizado por reforço molda o raciocínio.
• Direção Futura: O trabalho sugere que futuras pesquisas em RL para LLMs podem focar em otimizações mais simples e robustas, evitando a "superengenharia" de algoritmos complexos que não trazem benefícios adicionais de desempenho.

Em suma, o estudo demonstra que feedback negativo e estimativa de vantagem são os pilares essenciais, enquanto as restrições de estabilidade do PPO podem ser dispensadas em favor de uma abordagem baseada em REINFORCE mais limpa e eficaz.