Learning Reasoning Reward Models from Expert Demonstration via Inverse Reinforcement Learning

Título: Aprendendo Modelos de Recompensa de Raciocínio a partir de Demonstrações de Especialistas via Aprendizado por Reforço Inverso

1. Problema e Motivação

As abordagens atuais para melhorar o raciocínio em Grandes Modelos de Linguagem (LLMs) enfrentam duas limitações principais:

• Ajuste Fino Supervisionado (SFT): Baseia-se na imitação de traços de especialistas. Embora eficaz, é fundamentalmente imitativo e não aprende um objetivo explícito em nível de passo para avaliar trajetórias de raciocínio alternativas, especialmente quando o modelo sai do domínio das demonstrações (off-demonstration states).
• Aprendizado por Reforço (RL) com Recompensas de Resultado: Assume a existência de um verificador bem especificado (ground-truth) que fornece recompensas apenas no final da sequência. Isso ignora a qualidade dos passos intermediários e não oferece feedback processual.

O artigo propõe preencher essa lacuna aprendendo recompensas de raciocínio densas diretamente a partir de demonstrações de especialistas, sem depender de verificadores externos ou funções de recompensa projetadas manualmente.

2. Metodologia: AIRL para LLMs

Os autores adaptam o Aprendizado por Reforço Inverso Adversarial (AIRL) para o contexto de geração de linguagem. O objetivo é inferir uma função de recompensa $r_\phi$ que capture a qualidade do processo de raciocínio, permitindo otimizar a política do modelo.

Arquitetura do Modelo

• Discriminador ($D_\phi$): Um modelo classificador token-a-token que distingue entre traços de raciocínio de especialistas e traços gerados pela política atual.
• Política ($\pi_\theta$): O LLM gerador que é otimizado para maximizar a recompensa aprendida.
• Definição de Positivos e Negativos:
  • Para evitar que o discriminador aprenda apenas "estilo" ou "proveniência" (source bias), os rótulos são definidos com base na concordância da resposta final.
  • Positivos ($D_{pos}$): Traços de especialistas + traços da política que produzem a mesma resposta correta que o especialista.
  • Negativos ($D_{neg}$): Traços da política com resposta incorreta + traços de especialistas "corrompidos" sinteticamente (ex: troca de números, operadores invertidos) para aumentar a robustez.

Granularidade da Recompensa

O método introduz uma máscara binária $m_t$ para controlar a granularidade da supervisão, permitindo comparar diferentes níveis de densidade:

1. Esparsa (Sparse): Recompensa apenas no token final ($t=T$).
2. Intervalar (Interval): Recompensa a cada $k$ tokens.
3. Densa (Dense): Recompensa em todos os tokens ($t=1 \dots T$).
4. Passo a Passo (Step-wise): Recompensa em delimitadores de sentença/lógica.

Para granularidades esparsas ou intervalares, a recompensa é "preenchida" (backfilling) para todos os tokens subsequentes até o próximo ponto supervisionado.

Treinamento e Estabilidade

• Otimização Alternada: O discriminador e a política são atualizados alternadamente.
• GRPO (Group Relative Policy Optimization): A política é atualizada usando GRPO, adaptado para recompensas densas. As recompensas são normalizadas dentro de um grupo de amostras para reduzir a variância.
• Estabilização: Para evitar o colapso de modo e a instabilidade comum no treinamento adversarial, os autores utilizam:
  • Replay Buffers: Acumulam traços positivos e negativos de iterações anteriores.
  • Múltiplos Passos de Discriminador: O discriminador é atualizado várias vezes ($N_{disc}$) por atualização da política.
  • Clipping de Recompensa: Limita os valores da recompensa para evitar explosão numérica.

3. Contribuições Principais

1. Adaptação do AIRL para LLMs: Primeira aplicação de IRL adversarial para aprender recompensas de raciocínio em múltiplas granularidades diretamente de demonstrações, sem verificadores externos.
2. Reutilização da Recompensa: Demonstra que uma única recompensa aprendida pode ser usada tanto para treinamento da política quanto para reordenamento (reranking) em tempo de inferência.
3. Diagnóstico de Erros: A recompensa densa permite localizar o passo exato onde uma trajetória de raciocínio diverge do caminho correto, fornecendo feedback interpretável.

4. Resultados Experimentais

Os experimentos foram conduzidos em três benchmarks: GSM8K (matemática), MEDREASON (raciocínio médico) e MMLU-PRO (ciência geral), utilizando diversos backbones (Qwen2.5, Llama3, Qwen3).

D1: Sinal de Treinamento (Supervisão)

• As recompensas aprendidas via AIRL superaram o SFT em 8 de 9 configurações (combinações de modelo e dataset).
• A granularidade Esparsa mostrou-se a mais estável e robusta, enquanto a Densa foi mais sensível à instabilidade de treinamento e colapso de modo, embora tenha potencial para maior precisão.
• Em cenários difíceis (ex: Qwen2.5-7B no MEDREASON), a recompensa densa falhou inicialmente, mas a granularidade intervalar recuperou o desempenho, superando o SFT.

D2: Assistência em Tempo de Inferência (Reranking)

• O modelo de recompensa aprendida atuou como um reordenador eficaz sob um orçamento fixo de amostragem (Best-of-16).
• Ganhos significativos foram observados: até +17.4 pontos percentuais no GSM8K (Qwen2.5-7B com recompensa densa) e +10.7 pontos no MMLU-PRO.
• A combinação de recompensa aprendida com votação majoritária (majority voting) mostrou-se ainda mais eficaz em orçamentos maiores.

D3: Localização de Erros e Transferência

• Transferência: As recompensas aprenderam generalizaram bem entre diferentes tarefas e backbones (ex: um modelo treinado em GSM8K melhorou o reranking em MMLU-PRO).
• Localização de Erros: A análise token-a-token mostrou que a recompensa permanece positiva em passos corretos e cai abruptamente no primeiro passo errôneo, permitindo diagnosticar onde o raciocínio falhou (ex: em problemas de matemática, a recompensa cai exatamente na operação aritmética incorreta).

5. Significado e Conclusão

Este trabalho demonstra que é possível extrair um sinal de raciocínio intermediário reutilizável apenas a partir de demonstrações, sem necessidade de verificadores programáticos complexos.

• Ponte entre Imitação e Otimização: O método une a simplicidade da imitação (SFT) com a capacidade de otimização baseada em recompensas do RL.
• Diagnóstico Interpretável: A capacidade de identificar o momento exato da falha no raciocínio abre caminho para correções direcionadas e intervenções em tempo de teste.
• Desafios: O trade-off entre a densidade da recompensa e a estabilidade da otimização adversarial permanece um ponto de atenção, sugerindo que recompensas mais esparsas podem ser mais robustas para treinamento, enquanto as densas são superiores para diagnóstico e reranking.

Em suma, o framework proposto oferece uma via escalável para melhorar a capacidade de raciocínio de LLMs, transformando demonstrações brutas em sinais de supervisão processual ricos e reutilizáveis.