LLMs Can Learn to Reason Via Off-Policy RL

Este trabajo propone OAPL, un nuevo algoritmo de aprendizaje por refuerzo fuera de política que aprovecha la desincronización inherente en el entrenamiento distribuido de LLMs para lograr un rendimiento superior en matemáticas y codificación con una generación de datos significativamente menor y una mayor escalabilidad en el tiempo de prueba.

Lo esencial

¡Hola! Imagina que quieres enseñar a un robot muy inteligente (una Inteligencia Artificial) a resolver problemas matemáticos difíciles o a escribir código de computadora. Para hacerlo, usamos un método llamado Aprendizaje por Refuerzo, que es como darle premios al robot cuando acierta y "regañarlo" suavemente cuando falla.

El problema es que, hasta ahora, la mayoría de estos robots aprendían de una manera muy estricta y lenta, llamada "On-Policy" (en política). Vamos a usar una analogía para entenderlo mejor.

🏃‍♂️ La vieja forma: El corredor y el entrenador que se pelean

Imagina un entrenador (el cerebro que aprende) y un corredor (el robot que genera respuestas).

En el método antiguo (como el que usa la famosa IA DeepSeek-R1), el entrenador y el corredor deben estar siempre sincronizados.

1. El corredor hace una carrera (genera una respuesta).
2. El entrenador la mira inmediatamente, le da consejos y actualiza sus instrucciones.
3. Inmediatamente, el corredor debe usar esas nuevas instrucciones para la siguiente carrera.

El problema: En la vida real, esto es imposible de mantener perfecto. A veces, el corredor está en una computadora muy rápida y el entrenador en otra más lenta. O a veces, el entrenador está pensando en una estrategia nueva mientras el corredor ya está usando la vieja.
Esto crea un "desfase". El entrenador está dando consejos basados en un corredor que ya no existe (o que es diferente). Es como si un entrenador de fútbol le gritara tácticas a un jugador que ya cambió de equipo hace 10 minutos. El resultado es que el entrenamiento se vuelve inestable, el robot se confunde y a veces olvida lo que sabía (colapso de la entropía).

Para arreglar esto, los científicos anteriores intentaban "engañar" al sistema usando matemáticas complejas (llamadas muestreo por importancia) para decir: "Oye, aunque el corredor era un poco diferente, vamos a tratar sus respuestas como si fueran del entrenador". Esto es como intentar corregir un mapa antiguo con una lupa muy pequeña: funciona, pero es lento y propenso a errores.

🚀 La nueva forma: OAPL (El método de "Aprovechar el retraso")

Los autores de este paper dicen: "¡Esperen! ¿Por qué luchamos contra el desfase? ¡Aceptémoslo!".

Presentan un nuevo algoritmo llamado OAPL. Imagina que en lugar de obligar al entrenador y al corredor a estar sincronizados al segundo, les decimos:

"Corredor, tú sigue generando respuestas con tus instrucciones actuales. Entrenador, tú toma esas respuestas, úsalas para aprender, y no te preocupes si el corredor ya cambió un poco de estrategia mientras tú pensabas".

La analogía del "Entrenador con Memoria":
OAPL funciona como un entrenador muy sabio que sabe que su equipo cambia.

1. El entrenador no intenta corregir el pasado.
2. En su lugar, usa una fórmula matemática especial (una "regresión cuadrática") que le permite aprender de las respuestas viejas sin confundirse.
3. Solo se actualiza al entrenador y al corredor cada cierto tiempo (cada 50 o 100 pasos), en lugar de hacerlo en cada paso.

Esto es como si el entrenador dijera: "Tengo un montón de videos de partidos pasados (datos viejos). Voy a estudiarlos todos juntos para mejorar mi estrategia general, en lugar de intentar corregir cada jugada en tiempo real".

🌟 ¿Por qué es genial esto? (Los beneficios)

1. Es más rápido y eficiente: El paper dice que OAPL logra los mismos resultados que los métodos antiguos usando 3 veces menos datos. Es como aprender a tocar la guitarra en 3 meses en lugar de en 9.
2. Es más estable: Los robots que usan OAPL no se "vuelven locos" ni olvidan lo que saben. Mantienen una buena variedad de respuestas (no colapsan), lo que les permite encontrar soluciones creativas.
3. Funciona con un desfase enorme: Pueden entrenar con un retraso de más de 400 pasos sin problemas. Imagina que el entrenador está 400 días atrás en el tiempo y el corredor está en el presente, y aun así, ¡el entrenador sigue enseñando perfectamente!
4. Mejor rendimiento en pruebas difíciles: En matemáticas (como olimpiadas de matemáticas) y en programación, OAPL superó a los modelos anteriores, obteniendo mejores resultados incluso cuando se le pedía generar muchas opciones para elegir la mejor (Pass@k).

🎯 En resumen

La idea central del paper es que no necesitamos que el robot y el entrenador estén sincronizados al milímetro para aprender bien.

• Antes: Intentábamos forzar la sincronización perfecta o usar parches matemáticos complejos para corregir los errores.
• Ahora (OAPL): Aceptamos que el robot y el entrenador van a estar un poco desfasados, y diseñamos un método de aprendizaje que es tan robusto que ese desfase ni siquiera importa.

Es como pasar de intentar caminar en una cuerda floja (donde un pequeño error te hace caer) a caminar en un barco grande (donde las olas y los movimientos no te hacen perder el equilibrio). ¡Y el barco llega más rápido a la meta!

Resumen técnico

Resumen Técnico: LLMs Can Learn to Reason Via Off-Policy RL

1. El Problema: La Falacia de la "Política en Línea" en el Entrenamiento de RL

El entrenamiento de Grandes Modelos de Lenguaje (LLMs) mediante Aprendizaje por Refuerzo (RL) para potenciar capacidades de razonamiento (como en DeepSeek-R1) depende tradicionalmente de algoritmos en línea (on-policy) como PPO o GRPO. Estos algoritmos asumen que los datos de entrenamiento se generan utilizando la misma política que se está optimizando.

Sin embargo, en la infraestructura moderna de RL distribuido, esta suposición se rompe por diseño debido a dos factores principales:

1. Diferencias de Implementación: El motor de entrenamiento (ej. HuggingFace) y el motor de inferencia (ej. vLLM) pueden producir log-probabilidades diferentes para la misma secuencia, incluso con los mismos pesos, debido a diferencias en los kernels o implementaciones.
2. Latencia Asíncrona: En pipelines asíncronos, el motor de inferencia a menudo utiliza una versión antigua de los pesos del entrenador (retraso de política o policy lag).

Esto convierte los datos en fuera de línea (off-policy) por naturaleza. Las soluciones anteriores intentan corregir esto mediante:

• Muestreo de Importancia (IS): Añadir pesos para corregir la discrepancia, lo que introduce alta varianza en la función de pérdida.
• Alineación de Motores: Modificar el motor de inferencia para que coincida exactamente con el entrenador, lo que ralentiza la generación y no cierra completamente la brecha en entornos asíncronos.

2. Metodología: OAPL (Optimal Advantage-based Policy Optimization with Lagged Inference policy)

Los autores proponen OAPL, un algoritmo de RL fuera de línea que abraza la naturaleza off-policy en lugar de combatirla. En lugar de forzar que los datos parezcan on-policy, OAPL trata la discrepancia entre el entrenador ($\pi$) y el motor de inferencia ($\pi_{vllm}$) como un problema de RL regularizado por KL.

Componentes Clave del Algoritmo:

• Objetivo de Optimización: Se formula un objetivo que maximiza la recompensa mientras minimiza la divergencia KL respecto a la política de inferencia actual ($\pi_{vllm}$), no a una referencia fija antigua.
  $$\max_{\pi} \mathbb{E}[r(x, y)] - \beta \text{KL}(\pi || \pi_{vllm})$$
• Solución de Cierre: Utilizando la solución analítica de este problema, se deriva una función de pérdida basada en la ventaja óptima ($A^*$). La pérdida es una regresión cuadrática simple:
  $$\min_{\pi} \sum \left( \beta \ln \frac{\pi(y|x)}{\pi_{vllm}(y|x)} - (r(x, y) - \hat{V}^*(x)) \right)^2$$
  Donde $\hat{V}^*(x)$ es una estimación del valor óptimo calculada a partir de un grupo de rollos (rollouts) generados por $\pi_{vllm}$.
• Flujo Asíncrono:
  1. El motor de inferencia ($\pi_{vllm}$) genera datos de forma asíncrona y los almacena en un búfer.
  2. El entrenador ($\pi$) actualiza los pesos minimizando la pérdida de regresión cuadrática usando esos datos.
  3. Sincronización Infrecuente: Los pesos de $\pi$ y $\pi_{vllm}$ se sincronizan solo cada $L$ iteraciones (ej. cada 50 o 100 pasos). Entre sincronizaciones, el algoritmo opera completamente off-policy sin necesidad de ratios de importancia ni recorte (clipping).

3. Contribuciones Clave

• Nueva Perspectiva Teórica: Demuestran que la condición on-policy no es necesaria para el post-entrenamiento de LLMs. El uso de datos off-policy puede ser más eficiente y estable si se formula correctamente.
• Algoritmo OAPL: Un método simple, escalable y fácil de implementar que elimina la necesidad de:
  • Muestreo de importancia (IS) y sus ratios.
  • Operadores de recorte (clipping) complejos.
  • Modificaciones en el motor de inferencia.
• Robustez al Retraso: Funciona eficazmente con retrasos de política de más de 400 pasos de gradiente (100 veces más off-policy que los enfoques anteriores), permitiendo una arquitectura de entrenamiento altamente asíncrona.
• Mejora en la Escalabilidad de Prueba (Test-Time Scaling): A diferencia de métodos que colapsan la entropía, OAPL mantiene una entropía saludable, mejorando las métricas Pass@k para valores altos de $k$ (hasta 256).

4. Resultados Experimentales

Los autores evaluaron OAPL en matemáticas competitivas y generación de código:

• Razonamiento Matemático (AIME, HMMT, BRUMO):

  • OAPL superó consistentemente a la línea base GRPO (con IS) en todas las métricas Pass@k.
  • Mostró una convergencia más estable y una mayor precisión final.
  • Evitó el colapso de entropía observado en GRPO, lo que permitió un mejor rendimiento en Pass@5 y Pass@10.
  • Mantuvo la estabilidad incluso con un intervalo de sincronización de $L=100$ pasos.

• Generación de Código (LiveCodeBench):

  • OAPL igualó o superó el rendimiento del modelo DeepCoder (entrenado con GRPO y heurísticas complejas).
  • Eficiencia de Muestras: Logró este rendimiento utilizando 3 veces menos generaciones durante el entrenamiento (~200k muestras vs ~650k de DeepCoder).
  • Demostró que el entrenamiento off-policy extremo (sin sincronización durante 1 época completa, ~400 pasos) es viable y efectivo.

5. Significado e Impacto

Este trabajo representa un cambio de paradigma en el post-entrenamiento de LLMs:

1. Eficiencia Computacional: Al eliminar la necesidad de sincronización frecuente y el cálculo de ratios de importancia, OAPL reduce drásticamente la sobrecarga computacional y permite un uso más eficiente de los recursos de hardware (GPUs).
2. Escalabilidad: Facilita el entrenamiento a gran escala en arquitecturas distribuidas donde la latencia es inevitable, eliminando el cuello de botella de la sincronización on-policy.
3. Calidad del Modelo: Al mantener una mayor entropía y evitar el colapso de la distribución, los modelos entrenados con OAPL muestran una mejor capacidad de razonamiento bajo múltiples intentos (Pass@k alto), lo cual es crucial para aplicaciones de razonamiento complejo.

En conclusión, el paper argumenta que abrazar la naturaleza off-policy de los sistemas de RL modernos, en lugar de intentar corregirla, conduce a algoritmos más simples, estables y eficientes para el desarrollo de LLMs con capacidades de razonamiento.