Learning Reasoning Reward Models from Expert Demonstration via Inverse Reinforcement Learning

1. El Problema

Los enfoques actuales para mejorar el razonamiento en LLMs se basan principalmente en dos estrategias, ambas con limitaciones fundamentales:

• Ajuste Fino Supervisado (SFT): Entrena el modelo para imitar las trazas de expertos. Es fundamentalmente imitativo; optimiza la reproducción de trayectorias observadas pero no aprende un objetivo explícito por paso para evaluar o corregir razonamientos que se desvían de la demostración.
• Aprendizaje por Refuerzo (RL) con recompensas de resultado: Utiliza recompensas binarias al final de la secuencia (ej. ¿es correcta la respuesta?). Esto asume la existencia de un verificador bien especificado y no proporciona retroalimentación sobre dónde falló el proceso de razonamiento intermedio.

Existe una brecha: se necesitan métodos que aprendan recompensas de proceso (nivel de paso) directamente de demostraciones expertas, sin depender de verificadores externos costosos, para guiar tanto el entrenamiento como la inferencia.

2. Metodología: Aprendizaje por Refuerzo Inverso Adversarial (AIRL)

Los autores proponen un marco de Aprendizaje por Refuerzo Inverso Adversarial (AIRL) adaptado a LLMs para inferir una función de recompensa de razonamiento densa a partir de demostraciones expertas.

Componentes Clave:

• Formulación del Problema: Se trata como un juego minimax entre un discriminador (modelo de recompensa $D_\phi$) y una política de razonamiento ($\pi_\theta$). El objetivo es aprender una recompensa que asigne puntuaciones altas a las trazas expertas y bajas a las generadas por la política, mientras la política intenta maximizar la recompensa aprendida.
• Definición de Positivos y Negativos (Sin Verificadores Externos):
  • A diferencia de métodos anteriores que separan por fuente (experto vs. modelo), este método utiliza la coincidencia de la respuesta final.
  • Positivos ($D_{pos}$): Trazas expertas + Trazas de la política que coinciden con la respuesta experta.
  • Negativos ($D_{neg}$): Trazas de la política con respuestas incorrectas + Trazas expertas "corrompidas" sintéticamente (para robustez).
  • Esto fuerza al discriminador a aprender la calidad del razonamiento y no solo el estilo o la fuente.
• Granularidad de la Recompensa: El marco evalúa tres niveles de resolución de asignación de crédito mediante una máscara binaria $m_t$:
  1. Esparsa (Sparse): Recompensa solo al final ($t=T$).
  2. Intervalo (Interval): Recompensa cada $k$ tokens.
  3. Densa (Dense): Recompensa en cada token.
  • Nota: Para las máscaras esparsas o de intervalo, se utiliza un mecanismo de "backfilling" (relleno hacia atrás) para propagar la recompensa a los tokens intermedios no supervisados.
• Entrenamiento de la Política: Se utiliza GRPO (Group Relative Policy Optimization) para actualizar la política $\pi_\theta$ basándose en las recompensas implícitas derivadas del discriminador. Se implementan mecanismos de estabilización como buffers de repetición (replay buffers) y múltiples pasos de actualización del discriminador por paso de política para evitar el colapso de modos.

3. Contribuciones Principales

1. Adaptación de AIRL para LLMs: Se adapta el aprendizaje por refuerzo inverso adversarial para aprender recompensas de razonamiento a múltiples granularidades directamente desde demostraciones, sin necesidad de funciones de recompensa programadas manualmente.
2. Reutilización Dual: Se demuestra que una única recompensa aprendida sirve para dos propósitos complementarios:
   • Como señal de optimización para el entrenamiento de la política (superando al SFT).
   • Como clasificador/reordenador en tiempo de inferencia (reranking) bajo un presupuesto de muestreo fijo.
3. Localización de Errores: El modelo de recompensa densa permite diagnosticar en qué paso exacto una trayectoria de razonamiento se desvía de la ruta correcta, facilitando la identificación de modos de fallo.

4. Resultados Experimentales

Los experimentos se realizaron en tres benchmarks de razonamiento: GSM8K (matemáticas), MEDREASON (razonamiento clínico) y MMLU-PRO (ciencia general). Se probaron varios backbones (Qwen2.5, Llama3, Qwen3).

• Como Señal de Entrenamiento (D1):
  • Las recompensas aprendidas por AIRL superaron al SFT en 8 de 9 configuraciones (par backbone-dataset).
  • La granularidad Esparsa fue la más estable y robusta.
  • La granularidad Densa mostró resultados mixtos: aunque potente, fue más propensa a inestabilidad y colapso de modos en ciertos backbones (ej. Qwen2.5-7B en MEDREASON), aunque en otros casos (Qwen3-4B) logró rendimientos muy altos.
• Como Asistente en Inferencia (D2 - Reranking):
  • El uso de la recompensa para reordenar 16 muestras (Best-of-16) mejoró significativamente el rendimiento.
  • Se logró una mejora de hasta 17.4 puntos porcentuales en GSM8K (con recompensa densa en Qwen2.5-7B).
  • En general, el reordenamiento guiado por recompensa superó consistentemente a la selección aleatoria y a baselines basados en probabilidad logarítmica.
• Generalización y Localización (D3):
  • Transferencia: Las recompensas aprendidas en una tarea o backbone funcionaron bien al transferirse a otros (ej. un modelo entrenado en MEDREASON mejoró el rendimiento en GSM8K), sugiriendo que aprende señales de razonamiento reutilizables.
  • Diagnóstico: Las recompensas densas mostraron capacidad para mantener puntuaciones positivas en pasos intermedios correctos y caer abruptamente en el primer paso erróneo, permitiendo identificar el punto exacto de fallo en la cadena de pensamiento.

5. Significado y Conclusión

Este trabajo cierra la brecha entre la imitación pura (SFT) y la optimización basada en recompensas externas. Al aprender una función de recompensa de razonamiento intermedio directamente de las demostraciones, el método ofrece:

• Eficiencia: Elimina la necesidad de construir verificadores de reglas complejos para cada tarea.
• Interpretabilidad: Proporciona una señal de retroalimentación paso a paso que permite diagnosticar errores de razonamiento.
• Versatilidad: Un único modelo de recompensa sirve para entrenar modelos más inteligentes y para mejorar la selección de respuestas en tiempo de ejecución.

El estudio concluye que, aunque existe una compensación entre la densidad de la recompensa y la estabilidad de la optimización, el marco AIRL permite recuperar señales de razonamiento intermedio reutilizables, mejorando tanto la capacidad de generación como la verificación de LLMs en tareas complejas.