Resumo Técnico: rePIRL: Aprender PRM com Aprendizado por Reforço Inverso para Raciocínio de LLM

Declaração do Problema

No raciocínio de Grandes Modelos de Linguagem (LLMs), os Modelos de Recompensa de Processo (PRMs) são utilizados para fornecer recompensas intermediárias ao nível de token, visando melhorar a eficiência do treinamento, reduzir a variância e prevenir o "reward hacking" (subversão de recompensa). No entanto, os métodos existentes para aprender PRMs eficazes enfrentam limitações significativas:

• Métodos Offline: Abordagens supervisionadas (ex: Uesato et al., 2022; Lightman et al., 2023) e trabalhos do tipo DPO exigem acesso a anotações ao nível de token ou rótulos de preferência de trajetórias de especialistas, que são caros de obter. Métodos baseados em Busca em Árvore Monte Carlo (MCTS) (ex: Wang et al., 2024) exigem a execução de uma política especialista para gerar trajetórias, o que é computacionalmente caro e inviável se a política especialista não estiver disponível.
• Métodos Online: Abordagens como o PRIME (Cui et al., 2025a) dependem de suposições fortes em relação às funções de recompensa e ao modelo de política sob treinamento. Métodos de recompensa intrínseca (ex: Zhao et al., 2025b) utilizam a entropia ou confiança do modelo como recompensas, mas sofrem de "colapso de entropia", onde o desempenho degrada nos estágios finais do treinamento devido a um comportamento excessivo de busca de modo (mode-seeking).

O desafio central é aprender um PRM eficaz com o mínimo de suposições: especificamente, sem exigir anotações de recompensa ao nível de token, rótulos de preferência, acesso à política especialista ou suposições fortes sobre a forma da função de recompensa.

Metodologia: Estrutura rePIRL

Os autores propõem o rePIRL, uma estrutura inspirada em Aprendizado por Reforço Inverso (IRL) que aprende PRMs apenas a partir de trajetórias de especialistas. O método modela o raciocínio de LLM de múltiplos passos como um Processo de Decisão de Markov (MDP) ao nível de token e emprega um processo de aprendizado duplo que atualiza a política e o PRM de forma intercambiável.

1. Fundamentação Teórica

• Construção do Objetivo: Seguindo a estrutura clássica de IRL, os autores definem uma variável oculta $o_t \in \{0, 1\}$ indicando se uma ação é amostrada da política especialista. Eles modelam a probabilidade de uma ação ser amostrada do especialista usando uma função baseada em energia: $p(o_t=1|s_t, a_t) \propto \exp(r_\phi(s_t, a_t))$.
• Maximização de Verossimilhança: A função de recompensa $r_\phi$ é aprendida maximizando a verossimilhança das trajetórias de especialistas. Isso leva a um objetivo que envolve uma função de partição $z(\phi)$, que é intratável para LLMs devido aos vastos espaços de estado e ação.
• Tratabilidade via Amostragem: Para resolver a intratabilidade da função de partição, os autores aplicam um truque de gradiente e amostragem de importância. Em vez de amostrar da política ótima (que é desconhecida), eles amostram trajetórias $\bar{\tau}$ da política atual $\pi_\theta$. O gradiente é aproximado usando pesos de importância $w = \exp(r_\phi(\bar{\tau})) / \pi_\theta(\bar{\tau})$.
• Otimização da Política: Para evitar o colapso do modelo (políticas determinísticas), a política é atualizada usando um objetivo de RL de Entropia Máxima (Haarnoja et al., 2018), que equilibra a maximização da recompensa com a regularização de entropia. Isso é resolvido usando algoritmos como RLOO, PPO ou GRPO.

2. Algoritmo de Aprendizado

O ciclo de treinamento alterna entre a atualização do PRM e da política:

1. Rollout: Gerar rollouts da política e buscar demonstrações de especialistas.
2. Cálculo de Recompensa: Calcular recompensas de resultado (outcome rewards) e recompensas de processo médias.
3. Atualização do PRM: Atualizar o modelo de recompensa usando uma função de perda que contrasta trajetórias de especialistas contra rollouts da política, ponderados por escores de amostragem de importância. Crucialmente, os autores utilizam recompensas de processo médias em vez de recompensas cumulativas para mitigar o reward hacking.
4. Atualização da Política: Atualizar o modelo de política usando as recompensas de processo aprendidas combinadas com recompensas de resultado (se disponíveis), empregando objetivos de entropia máxima.

3. Unificação de Métodos Existentes

Uma contribuição teórica fundamental é demonstrar que o rePIRL pode unificar vários métodos existentes de aprendizado de PRM (DPO, DQO, Math-Shepherd, PRIME) sob uma única estrutura ao introduzir suposições adicionais específicas. Isso posiciona o rePIRL como o método mais geral, exigindo o menor número de suposições.

Principais Contribuições

1. Estrutura Inovadora: Introdução do rePIRL, uma estrutura inspirada em IRL para aprender PRMs a partir de trajetórias de especialistas sem necessidade de anotações ao nível de token, rótulos de preferência ou acesso à política especialista.
2. Teoria Unificada: Uma análise teórica mostrando que o rePIRL unifica diversos métodos de aprendizado de PRM offline e online de última geração (incluindo DPO, DQO, PRIME e abordagens baseadas em MCTS), esclarecendo suas suposições subjacentes.
3. Algoritmo Escalável: Desenvolvimento de um algoritmo de aprendizado escalável adaptado para LLMs, utilizando amostragem de importância para lidar com funções de partição intratáveis e RL de entropia máxima para garantir a estabilidade do treinamento.
4. Validação Empírica: Avaliação abrangente em benchmarks de matemática e codificação demonstrando desempenho superior em relação aos baselines.

Resultados Experimentais

Os autores avaliaram o rePIRL utilizando Qwen2.5-3B-Instruct e Qwen3-4B-Base em sete benchmarks padrão (AIME-2024, AMC, Math-500, MinervaMath, OlympiadBench, Leetcode, LiveCodeBench).

• Desempenho: O rePIRL alcançou resultados de estado da arte entre modelos de escala 3B e 4B. Por exemplo, no Qwen3-4B-Base, atingiu 41,6% de acurácia média em benchmarks de matemática e 39,8% em benchmarks de codificação, superando baselines fortes como PRIME, MCTS e o RLOO vanilla.
• Comparação: Apresentou melhorias absolutas de 3,2% (matemática) e 4,3% (codificação) sobre o RLOO vanilla com o Qwen3-4B-Base.
• Aplicações:
  • Treinamento em Tempo de Teste (Test-time Training): Quando treinado exclusivamente com sinais de PRM sem recompensas de resultado, o rePIRL igualou o desempenho de modelos treinados com recompensas de resultado verificáveis.
  • Escalabilidade em Tempo de Teste (Test-time Scaling): O uso do PRM como um verificador para selecionar melhores rollouts resultou em melhorias monotônicas de desempenho, superando a votação majoritária.
  • Problemas Difíceis: O PRM forneceu sinais precoces informativos para problemas difíceis onde as recompensas de resultado eram esparsas, levando a uma convergência mais rápida em comparação com o treinamento baseado apenas em resultados.
• Eficiência: O rePIRL introduziu apenas um overhead de tempo de treinamento de aproximadamente 1,4x em comparação ao RL baseado em resultados (ex: RLOO), enquanto evita a necessidade de redes de valor separadas.

Significância e Alegações

O artigo afirma ser o primeiro trabalho a generalizar o Aprendizado por Reforço Inverso para o raciocínio de LLM e o primeiro a unificar os métodos de aprendizado de PRM existentes em uma única estrutura.

• Mínimas Suposições: Os autores enfatizam que o rePIRL opera sob as suposições mínimas necessárias, tornando-o mais generalizável do que métodos que exigem anotações ao nível de token ou acesso específico à política especialista.
• Estabilidade: A estrutura aborda o problema do colapso de entropia comum em métodos de recompensa intrínseca ao utilizar RL de entropia máxima, garantindo uma dinâmica de treinamento estável.
• Utilidade: O PRM treinado é apresentado como um módulo reutilizável e independente que beneficia aplicações downstream, como treinamento em tempo de teste e escalabilidade, particularmente em cenários onde recompensas de resultado não estão disponíveis ou são esparsas.

Os autores concluem que, embora o rePIRL adicione uma perspectiva teórica unificada e utilidade prática, trabalhos futuros poderiam explorar a combinação com métodos de recompensa intrínseca, comparações sistemáticas de heurísticas de RL online e a escalabilidade para tarefas agênticas e modelos maiores.