DyJR: Preserving Diversity in Reinforcement Learning with Verifiable Rewards via Dynamic Jensen-Shannon Replay

Each language version is independently generated for its own context, not a direct translation.

¡Hola! Imagina que estás entrenando a un genio matemático (una Inteligencia Artificial) para que resuelva problemas muy difíciles. El método tradicional es como darle un examen, ver si acertó, y luego tirar el examen a la basura inmediatamente para pasar al siguiente. Esto es muy ineficiente: el genio olvida lo que aprendió en los intentos anteriores y a veces se vuelve "obcecado", pensando que solo hay una forma correcta de resolver todo, perdiendo su creatividad.

El paper que me has pasado presenta una solución brillante llamada DyJR. Vamos a explicarlo con una analogía sencilla: El "Gimnasio de la Memoria".

1. El Problema: El Entrenador que Olvida y el Atleta Obsesivo

En el entrenamiento actual (llamado GRPO), el entrenador (la IA) hace un ejercicio, ve si funciona, y luego borra la memoria de ese intento para hacer el siguiente.

El riesgo: Si el entrenador solo se fija en la respuesta perfecta de hoy, puede volverse un "robot". Solo aprenderá un camino para resolver problemas y dejará de explorar otras formas creativas. Es como un atleta que solo practica una sola carrera y olvida cómo correr en diferentes terrenos.
El error de los métodos anteriores: Otros intentaron guardar todos los exámenes pasados en una biblioteca gigante para repasarlos. Pero esto tiene dos problemas:
1. Cuesta mucho dinero y energía (memoria de computadora) guardar todo.
2. Confunde al atleta: Si le haces repasar exámenes de hace un año, cuando su nivel era muy bajo, le enseñas cosas que ya no le sirven o lo confundes.

2. La Solución: DyJR (El Entrenador Sabio)

DyJR cambia las reglas del juego con dos ideas geniales:

A. La "Caja de Recuerdos Dinámica" (El Buffer)

Imagina que en lugar de guardar todos los exámenes de la vida del atleta, tienes una caja mágica que solo guarda los últimos intentos.

La regla de oro: Solo guardamos los intentos que fueron recientes y exitosos.
El truco: Al principio del entrenamiento, cuando el atleta está aprendiendo y probando muchas cosas locas (alta diversidad), la caja se hace grande para guardar esos momentos creativos. A medida que el atleta se vuelve experto y se estabiliza, la caja se hace pequeña, guardando solo lo más reciente.
Por qué funciona: Es como tener un entrenador que te dice: "Recuerda cómo pensabas hace un momento, cuando estabas explorando, no lo que pensabas hace un año". Esto mantiene la frescura y evita que el atleta se vuelva rígido.

B. El "Espejo de Diversidad" (La Regularización JS)

Aquí está la parte más mágica. En lugar de decirle al atleta: "¡Haz exactamente lo mismo que hiciste en el pasado!" (lo cual lo vuelve un robot), DyJR le dice: "Mira el espejo de tus intentos pasados exitosos y asegúrate de que tu forma de pensar hoy sea similar en variedad a la de ayer".

La analogía: Imagina que el atleta tiene un "espejo" que refleja todas las formas creativas en las que resolvió problemas ayer.
El objetivo: No se trata de copiar la respuesta exacta, sino de no perder la capacidad de pensar de muchas formas. Si el atleta empieza a pensar que solo hay una solución (el Rank-1, o la opción más obvia), el espejo le da un "empujón" suave para que vuelva a considerar otras opciones (Rank-2, Rank-3).
Resultado: El atleta sigue siendo muy inteligente, pero no se vuelve un robot aburrido. Sigue explorando caminos nuevos.

3. ¿Qué logran con esto?

En pruebas reales (resolviendo problemas de matemáticas y creando código SQL):

Son más rápidos: No necesitan guardar montañas de datos, ahorrando memoria de computadora.
Son más inteligentes: Resuelven problemas más difíciles porque no se quedan atascados en un solo método.
Son más creativos: Mantienen la capacidad de encontrar múltiples soluciones, no solo la primera que se les ocurre.

En resumen

DyJR es como un entrenador que sabe que la diversidad es más importante que la perfección inmediata. En lugar de llenar la cabeza del estudiante con todos sus exámenes pasados (lo cual es pesado y confuso), le enseña a recordar cómo pensaba cuando estaba lleno de ideas frescas, asegurándose de que nunca pierda su capacidad de explorar y crear.

Es un cambio de mentalidad: No se trata de repetir lo que ya sabes, sino de preservar tu capacidad de descubrir cosas nuevas.

Each language version is independently generated for its own context, not a direct translation.

Aquí tienes un resumen técnico detallado del artículo "DyJR: Preserving Diversity in Reinforcement Learning with Verifiable Rewards via Dynamic Jensen-Shannon Replay" en español.

1. Planteamiento del Problema

El aprendizaje por refuerzo (RL) ha demostrado ser fundamental para mejorar las capacidades de razonamiento de los Modelos de Lenguaje Grande (LLM), especialmente mediante algoritmos on-policy como GRPO (Group Relative Policy Optimization) que utilizan recompensas verificables (RLVR). Sin embargo, estos métodos adolecen de dos problemas críticos:

Ineficiencia de Muestra: Al ser on-policy, GRPO descarta los datos de los rollouts (trayectorias generadas) después de una sola actualización, desperdiciando recursos computacionales costosos.
Colapso de Modos (Mode Collapse) y Sobreajuste: Los métodos existentes de Experience Replay (reutilización de datos históricos) intentan corregir la ineficiencia tratando las trayectorias pasadas como nuevas instancias positivas para actualizaciones directas del gradiente. Esto provoca que el modelo se sobreajuste a soluciones específicas, pierda su capacidad de exploración y colapse en un único modo de respuesta (dependencia excesiva en tokens de rango 1), reduciendo la diversidad y la robustez.

La premisa central de los autores es que, en tareas de razonamiento, el valor de los datos históricos no radica principalmente en reforzar la precisión, sino en sustentar la diversidad de las estrategias de solución.

2. Metodología: DyJR (Dynamic Jensen-Shannon Replay)

Los autores proponen DyJR, un marco de regularización que redefine el papel del replay buffer. En lugar de usar datos históricos para actualizaciones directas de política, se utilizan como una restricción de distribución para mantener la diversidad.

A. Construcción de Datos Dinámica (Buffer Adaptativo)

En lugar de almacenar masivamente todas las trayectorias, DyJR implementa una estrategia no uniforme basada en la proximidad temporal:

Buffer de Edad Máxima (Max Age - M): El buffer retiene estrictamente muestras perfectas (recompensa = 1) generadas dentro de los últimos $M$ pasos. Esto asegura que la distribución de referencia evolucione junto con el modelo, evitando el desajuste de distribución (distribution shift) de datos muy antiguos.
Selección de Datos Sensible al Tiempo: Se utiliza un protocolo FIFO (First-In-First-Out) y una estrategia de admisión basada en la confianza (confidence-stratified).
- Prioriza trayectorias de alta confianza (múltiples respuestas correctas en un grupo).
- Calentamiento (Warm-up): Durante las primeras etapas del entrenamiento (cuando la entropía cae rápidamente), el tamaño del buffer se expande temporalmente para capturar patrones de alta entropía y evitar el colapso prematuro.

Ajuste de la Utilización de Datos (Regularización JS)

DyJR evita actualizar la política directamente con los datos del buffer. En su lugar, introduce un término de regularización basado en la Divergencia de Jensen-Shannon (JS):

Distribución de Referencia ( $Q_B$ ): Se construye a partir de la mezcla de políticas históricas almacenadas en el buffer.
Función de Pérdida: Se minimiza la divergencia JS entre la política actual ( $\pi_\theta$ $π_{θ}$ ) y la distribución de referencia del buffer ( $Q_B$ $Q_{B}$ ).
- A diferencia de la Divergencia KL Directa (que tiende a cubrir todos los modos y puede suavizar en exceso), la JS es simétrica y acotada, proporcionando un equilibrio más robusto.
- Esto actúa como un "ancla" flexible que impide que el modelo se desvíe demasiado de los caminos exitosos diversos sin forzar un ajuste agresivo.

La función de pérdida total combina el objetivo de GRPO (en línea) con la regularización JS:
$\mathcal{L}_{total}(\theta) = \mathcal{L}_{GRPO}(\theta) + \alpha_{JS} \cdot \mathcal{L}_{JS}(\theta)$

3. Contribuciones Clave

Redefinición del Paradigma de Replay: Cambian el objetivo de la optimización de precisión (usar datos pasados para ganar más) a la regularización de diversidad (usar datos pasados para evitar el colapso).
Estrategia de Construcción de Datos Dinámica: Introducen un mecanismo de buffer que se expande y contrae según la fase de entrenamiento y utiliza una ventana de tiempo estricta (FIFO) para mantener la relevancia de los datos, reduciendo drásticamente la huella de memoria.
Regularización JS Efectiva: Demuestran que la divergencia JS es superior a la KL Directa o Inversa en este contexto, ya que maneja mejor la naturaleza no estacionaria y heterogénea de los datos de replay en RLVR.

4. Resultados Experimentales

Los autores evaluaron DyJR en tareas de razonamiento matemático (Qwen3-4B) y generación de SQL (Llama-3.1-8B).

Rendimiento en Matemáticas:
- DyJR superó significativamente a GRPO y a otros métodos de replay (RLEP, Ex-GRPO, DPH-RL).
- Logró un 34.1% de precisión promedio en 6 benchmarks matemáticos, una mejora de 4.3 puntos sobre GRPO (29.8%).
- La mejora fue consistente tanto en problemas fáciles (AMC23, +7.4%) como en extremadamente difíciles (HMMT25, +2.9%).
Rendimiento en Text-to-SQL:
- En los conjuntos de datos BIRD y Spider, DyJR superó a GRPO en Pass@1 (+3.3% en BIRD, +5.0% en Spider) y en Pass@16, demostrando generalización cruzada.
Eficiencia:
- A diferencia de métodos como RLEP que requieren almacenar múltiples soluciones completas (alto costo de memoria), DyJR requiere almacenar solo una fracción de los datos (aprox. 2k pares vs 28k), manteniendo una sobrecarga de memoria de GPU insignificante (<1GB).
- La complejidad temporal es marginal, manteniendo un rendimiento de entrenamiento comparable al GRPO original.

5. Análisis y Significado

Dinámica de Entropía y Diversidad: El análisis de la evolución de los tokens de rango-k muestra que GRPO sufre un colapso rápido (probabilidad de Rank-1 > 90% desde el inicio). DyJR, por el contrario, mantiene una entropía más saludable y redistribuye la probabilidad hacia Rank-2 y Rank-3, lo que indica una capacidad de exploración sostenida.
Importancia de la Proximidad Temporal: Los experimentos de ablación revelan que un tamaño de ventana temporal pequeño ( $M=8$ ) es óptimo. Ventanas más grandes ( $M=64$ ) degradan el rendimiento, confirmando que los datos muy antiguos son perjudiciales debido a la divergencia de la política.
Significado: DyJR demuestra que la clave para escalar el RL en razonamiento no es simplemente acumular más datos, sino gestionar inteligentemente la diversidad de los datos históricos. Al priorizar la diversidad sobre la precisión pura en la fase de replay, se logra un equilibrio superior entre exploración y explotación, permitiendo que los modelos descubran soluciones óptimas sin caer en óptimos locales prematuros.

En conclusión, DyJR ofrece una solución eficiente y escalable para el problema del colapso de modos en el RL de LLMs, estableciendo un nuevo estándar para el uso de replay buffers en tareas de razonamiento complejo.