Resumen Técnico: Aprendizaje por Refuerzo a partir de Retroalimentación Rica con DAgger Distribucional

Formulación del Problema

El artículo aborda las limitaciones del paradigma dominante en el entrenamiento de modelos de razonamiento: el Aprendizaje por Refuerzo con Recompensas Verificables (RLVR). En RLVR, los modelos generan múltiples soluciones candidatas y a cada una se le asigna una única recompensa binaria (correcto/incorrecto) basada en la respuesta final. Este enfoque sufre de dos fallos críticos:

1. Asignación de Crédito Deficiente: La recompensa terminal se transmite uniformemente a todos los tokens en una respuesta, fallando al distinguir qué pasos de razonamiento específicos contribuyeron al éxito o al fracaso.
2. Inaplicabilidad a Dominios No Verificables: El RLVR depende de la verificación automática de las respuestas finales, lo que lo hace inadecuado para tareas donde las respuestas de verdad absoluta (ground-truth) no están disponibles o son difíciles de comprobar.

Aunque muchos entornos ofrecen "retroalimentación rica" (por ejemplo, trazas de ejecución, registros de errores, correcciones de expertos o autoevaluaciones), los métodos actuales que intentan aprovechar esto mediante la autodestilación on-policy (por ejemplo, SDPO, OPSD) enfrentan obstáculos teóricos y prácticos. Los enfoques existentes suelen minimizar las $f$-divergencias (como la KL inversa o Jensen-Shannon) entre una política estudiante y un profesor condicionado por la retroalimentación. Los autores argumentan que estos métodos carecen de garantías para la mejora monotónica de la política y no logran realizar una asignación de crédito a nivel de secuencia, lo que a menudo conduce a una convergencia subóptima.

Metodología: DistIL

Los autores proponen DistIL (Aprendizaje de Imitación Distribucional), una variante distribucional del algoritmo clásico DAgger adaptada para el aprendizaje por refuerzo con retroalimentación rica.

Objetivo Central

DistIL optimiza una pérdida de entropía cruzada hacia adelante (forward cross-entropy) entre la política estudiante ($\pi_\theta$) y la política del profesor condicionada por la retroalimentación ($\pi_T$) en los estados visitados por el estudiante:
$$\mathcal{L}_{\text{DistIL}}(\theta) := \mathbb{E}_{x \sim \rho, y \sim \pi_\theta(\cdot|x)} \left[ \sum_{t=1}^H H_\times(\pi_T(\cdot | s_t), \pi_\theta(\cdot | s_t)) \right]$$
donde $H_\times(p, q) = -\mathbb{E}_{y \sim p}[\log q(y)]$ es la entropía cruzada hacia adelante. Crucialmente, el profesor $\pi_T$ está condicionado por la retroalimentación rica $f$ (por ejemplo, registros de ejecución o verdad absoluta) pero se mantiene fijo durante la actualización del estudiante.

Gradiente y Asignación de Crédito

A diferencia de los métodos de autodestilación previos que aproximan los gradientes utilizando únicamente discrepancias locales por token (efectivamente congelando la distribución de estados), DistIL utiliza gradientes completos. El gradiente se descompone en dos términos:

1. Imitación Local: Un término ponderado por el profesor que actualiza al estudiante para que coincida con la distribución del profesor en el estado actual.
2. Asignación de Crédito Consciente del Futuro: Un término que propaga la discrepancia de la entropía cruzada observada en pasos de tiempo futuros hacia las decisiones anteriores. Esto asegura que las elecciones de tokens tempranos que influyen en la trayectoria hacia un estado con alta discrepancia entre profesor y estudiante reciban el crédito o la penalización apropiados.

El algoritmo admite profesores de caja negra (modelos externos o expertos humanos), ya que no requiere acceso a las probabilidades logarítmicas del profesor, solo muestras o la capacidad de estimar el objetivo.

Contribuciones Clave

1. Limitaciones Teóricas de los Métodos Existentes

El artículo proporciona pruebas formales que demuestran que los objetivos de autodestilación on-policy existentes basados en $f$-divergencias (incluyendo la KL inversa y Jensen-Shannon) sufren de dos problemas fundamentales:

• Sin Garantía de Mejora Monotónica: Incluso si el profesor tiene una recompensa esperada estrictamente mayor que el estudiante, minimizar una $f$-divergencia no garantiza que un paso de gradiente natural aumentará la recompensa esperada del estudiante. La dirección de la actualización puede aumentar inadvertidamente la probabilidad de acciones subóptimas debido al peso de las razones logarítmicas.
• Convergencia Subóptima vía Aproximación Local: Al ignorar el término de "crédito futuro" (el impacto de las acciones actuales en las distribuciones de estados futuros), las aproximaciones de gradiente locales pueden converger a políticas estrictamente subóptimas. Los autores demuestran que, en MDPs simples de dos pasos, las actualizaciones locales pueden fallar al asignar crédito a las decisiones tempranas que son cruciales para alcanzar estados de alta recompensa.

2. Garantías Teóricas de DistIL

Se demuestra que DistIL posee propiedades que alinean la optimización con la maximización de la recompensa:

• Mejora Monotónica de la Política: Bajo supuestos estándar de realización local, el objetivo de entropía cruzada hacia adelante garantiza que moverse hacia un mejor profesor aumenta la recompensa esperada del estudiante.
• Límites de Arrepentimiento (Regret Bounds): El algoritmo admite un límite de arrepentimiento sublineal, con tasas de convergencia que dependen de los coeficientes de concentrabilidad y la varianza del profesor.
• Optimización de Pass@N: Se demuestra que el objetivo maximiza un límite inferior de la verosimilitud del éxito ponderada por el profesor. Esto proporciona una explicación fundamentada de por qué DistIL mejora las métricas Pass@N (la probabilidad de generar al menos una respuesta correcta en $N$ muestras) para todo $N$.

3. Rendimiento Empírico

DistIL fue evaluado en tres dominios:

• Razonamiento Científico: En SciKnowEval (Química, Física, Biología, Ciencia de Materiales), DistIL superó a SDPO, GRPO (on-policy y off-policy) y logró puntuaciones más altas de Best@16 y Maj@16. Notablemente, DistIL mostró un entrenamiento estable, mientras que SDPO exhibió oscilaciones significativas y colapso de rendimiento en ciertos dominios.
• Programación (Coding): En LiveCodeBench, utilizando trazas de ejecución como retroalimentación, DistIL superó a SDPO y GRPO tanto en exactitud exacta como en precisión parcial.
• Razonamiento Matemático Difícil: En problemas desafiantes (AIME, HMMT) donde el modelo base tiene una tasa de éxito de cero, los métodos RLVR (GRPO) fallaron debido a señales de ventaja nulas. DistIL aprovechó con éxito las soluciones de verdad absoluta para mejorar el rendimiento, superando a las líneas base de autodestilación (SDPO, OPSD) y SFT.

Significado y Reivindicaciones

El artículo afirma que DistIL cierra la brecha entre el aprendizaje por imitación y el aprendizaje por refuerzo al proporcionar un marco teóricamente fundamentado para aprender de la retroalimentación rica. Su importancia radica en:

1. Corregir el Objetivo: Pasar de la minimización de la $f$-divergencia (que puede estar desalineada con la recompensa) a la entropía cruzada hacia adelante, la cual garantiza una mejora monotónica cuando el profesor es mejor.
2. Permitir el Aprendizaje a Nivel de Secuencia: Al utilizar gradientes completos, DistIL resuelve el problema de la asignación de crédito inherente al razonamiento de largo horizonte, asegurando que las decisiones tempranas sean evaluadas según sus consecuencias posteriores.
3. Robustez Práctica: El método funciona con profesores de caja negra y diversos tipos de retroalimentación (desde la corrección binaria hasta los registros de ejecución), lo que lo hace aplicable a una gama más amplia de tareas de razonamiento que el RLVR estándar.

Los autores concluyen que la combinación de un objetivo alineado con la recompensa y una asignación de crédito consciente del futuro es esencial para destilar eficazmente la retroalimentación rica en modelos de razonamiento, lo que conduce a ganancias empíricas consistentes sobre las líneas base del estado del arte.