Resumen Técnico: rePIRL: Aprender PRM con RL Inverso para el Razonamiento de LLM

Declaración del Problema

En el razonamiento de Modelos de Lenguaje de Gran Escala (LLM), los Modelos de Recompensa de Proceso (PRM) se utilizan para proporcionar recompensas a nivel de token, con el objetivo de mejorar la eficiencia del entrenamiento, reducir la varianza y prevenir el "reward hacking" (manipulación de la recompensa). Sin embargo, los métodos existentes para aprender PRMs efectivos enfrentan limitaciones significativas:

• Métodos Offline: Los enfoques supervisados (p. ej., Uesato et al., 2022; Lightman et al., 2023) y los trabajos de tipo DPO requieren acceso a anotaciones a nivel de token o etiquetas de preferencia de trayectorias expertas, lo cual es costoso de obtener. Los métodos basados en Búsqueda de Árbol de Monte Carlo (MCTS) (p. ej., Wang et al., 2024) requieren ejecutar una política experta para generar trayectorias, lo cual es computacionalmente costoso e inviable si la política experta no está disponible.
• Métodos Online: Enfoques como PRIME (Cui et al., 2025a) dependen de suposiciones fuertes respecto a las funciones de recompensa y al modelo de política bajo entrenamiento. Los métodos de recompensa intrínseca (p. ej., Zhao et al., 2025b) utilizan la entropía o la confianza del modelo como recompensas, pero sufren de "colapso de entropía", donde el rendimiento se degrada en las etapas finales del entrenamiento debido a un comportamiento excesivo de búsqueda de modo (mode-seeking).

El desafío central es aprender un PRM efectivo con mínimas suposiciones: específicamente, sin requerir anotaciones de recompensa a nivel de token, etiquetas de preferencia, acceso a la política experta o suposiciones fuertes sobre la forma de la función de recompensa.

Metodología: Marco de Trabajo rePIRL

Los autores proponen rePIRL, un marco inspirado en el Aprendizaje por Refuerzo Inverso (IRL) que aprende PRMs únicamente a partir de trayectorias expertas. El método modela el razonamiento de LLM de múltiples pasos como un Proceso de Decisión de Markov (MDP) a nivel de token y emplea un proceso de aprendizaje dual que actualiza la política y el PRM de manera intercambiable.

1. Fundamento Teórico

• Construcción del Objetivo: Siguiendo el marco clásico de IRL, los autores definen una variable oculta $o_t \in \{0, 1\}$ que indica si una acción es muestreada de la política experta. Modelan la probabilidad de que una acción sea muestreada por el experto utilizando una función basada en energía: $p(o_t=1|s_t, a_t) \propto \exp(r_\phi(s_t, a_t))$.
• Maximización de la Verosimilitud: La función de recompensa $r_\phi$ se aprende maximizando la verosimilitud de las trayectorias expertas. Esto conduce a un objetivo que involucra una función de partición $z(\phi)$, la cual es intratable para los LLM debido a los vastos espacios de estados y acciones.
• Tractabilidad mediante Muestreo: Para resolver la intratabilidad de la función de partición, los autores aplican un truco de gradiente y muestreo de importancia (importance sampling). En lugar de muestrear de la política óptima (que es desconocida), muestrean trayectorias $\bar{\tau}$ de la política actual $\pi_\theta$. El gradiente se aproxima utilizando pesos de importancia $w = \exp(r_\phi(\bar{\tau})) / \pi_\theta(\bar{\tau})$.
• Optimización de la Política: Para prevenir el colapso del modelo (políticas deterministas), la política se actualiza utilizando un objetivo de RL de Máxima Entropía (Haarnoja et al., 2018), que equilibra la maximización de la recompensa con la regularización de la entropía. Esto se resuelve utilizando algoritmos como RLOO, PPO o GRPO.

2. Algoritmo de Aprendizaje

El bucle de entrenamiento alterna entre la actualización del PRM y la política:

1. Rollout: Generar rollouts de la política y obtener demostraciones expertas.
2. Cálculo de Recompensa: Calcular recompensas de resultado (outcome rewards) y recompensas de proceso promedio.
3. Actualización del PRM: Actualizar el modelo de recompensa utilizando una función de pérdida que contrasta las trayectorias expertas contra los rollouts de la política, ponderados por puntuaciones de muestreo de importancia. Crucialmente, los autores utilizan recompensas de proceso promedio en lugar de recompensas acumulativas para mitigar el reward hacking.
4. Actualización de la Política: Actualizar el modelo de política utilizando las recompensas de proceso aprendidas combinadas con recompensas de resultado (si están disponibles), empleando objetivos de máxima entropía.

3. Unificación de Métodos Existentes

Una contribución teórica clave es demostrar que rePIRL puede unificar varios métodos de aprendizaje de PRM existentes (DPO, DQO, Math-Shepherd, PRIME) bajo un solo marco mediante la introducción de suposiciones adicionales específicas. Esto posiciona a rePIRL como el método más general, que requiere las menores suposiciones.

Principales Contribuciones

1. Marco Novedoso: Introducción de rePIRL, un marco inspirado en IRL para aprender PRMs de trayectorias expertas sin anotaciones a nivel de token, etiquetas de preferencia o acceso a la política experta.
2. Teoría Unificada: Un análisis teórico que muestra que rePIRL unifica diversos métodos de aprendizaje de PRM offline y online de vanguardia (incluyendo DPO, DQO, PRIME y enfoques basados en MCTS) bajo un único marco de aprendizaje, aclarando sus suposiciones subyacentes.
3. Algoritmo Escalable: Desarrollo de un algoritmo de aprendizaje escalable diseñado para LLMs, utilizando muestreo de importancia para manejar funciones de partición intratables y RL de máxima entropía para asegurar la estabilidad del entrenamiento.
4. Validación Empírica: Evaluación exhaustiva en benchmarks de matemáticas y programación demostrando un rendimiento superior sobre las líneas base.

Resultados Experimentales

Los autores evaluaron rePIRL utilizando Qwen2.5-3B-Instruct y Qwen3-4B-Base en siete benchmarks estándar (AIME-2024, AMC, Math-500, MinervaMath, OlympiadBench, Leetcode, LiveCodeBench).

• Rendimiento: rePIRL alcanzó resultados de vanguardia (SOTA) entre los modelos de escala 3B y 4B. Por ejemplo, en Qwen3-4B-Base, alcanzó un 41.6% de precisión promedio en benchmarks de matemáticas y un 39.8% en benchmarks de codificación, superando bases sólidas como PRIME, MCTS y RLOO vanilla.
• Comparación: Mostró mejoras absolutas de 3.2% (matemáticas) y 4.3% (codificación) sobre RLOO vanilla con Qwen3-4B-Base.
• Aplicaciones:
  • Entrenamiento en Tiempo de Prueba (Test-time Training): Cuando se entrena únicamente con señales de PRM sin recompensas de resultado, rePIRL igualó el rendimiento de los modelos entrenados con recompensas de resultado verificables.
  • Escalamiento en Tiempo de Prueba (Test-time Scaling): El uso del PRM como verificador para seleccionar mejores rollouts produjo mejoras de rendimiento monotónicas, superando al voto mayoritario (majority voting).
  • Problemas Difíciles: El PRM proporcionó señales tempranas informativas para problemas difíciles donde las recompensas de resultado eran dispersas, lo que llevó a una convergencia más rápida en comparación con el entrenamiento basado solo en resultados.
• Eficiencia: rePIRL introdujo solo un exceso de tiempo de entrenamiento de aproximadamente 1.4x en comparación con el RL basado en resultados (p. ej., RLOO), mientras evitaba la necesidad de redes de valor separadas.

Significado y Reivindicaciones

El artículo afirma ser el primer trabajo en generalizar el Aprendizaje por Refuerzo Inverso para el razonamiento de LLM y el primero en unificar los métodos de aprendizaje de PRM existentes en un solo marco.

• Mínimas Suposiciones: Los autores enfatizan que rePIRL opera bajo las suposiciones mínimas necesarias, lo que lo hace más generalizable que los métodos que requieren anotaciones a nivel de token o acceso específico a la política experta.
• Estabilidad: El marco aborda el problema del colapso de entropía común en los métodos de recompensa intrínseca mediante el uso de RL de máxima entropía, asegurando una dinámica de entrenamiento estable.
• Utilidad: El PRM entrenado se presenta como un módulo reutilizable e independiente que beneficia las aplicaciones posteriores, tales como el entrenamiento en tiempo de prueba y el escalamiento, particularmente en escenarios donde las recompensas de resultado no están disponibles o son dispersas.

Los autores concluyen que, si bien rePIRL añade una perspectiva teórica unificada y utilidad práctica, el trabajo futuro podría explorar la combinación con métodos de recompensa intrínseca, comparaciones sistemáticas de heurísticas de RL online y el escalamiento hacia tareas agénticas y modelos más grandes.