SPS: Steering Probability Squeezing for Better Exploration in Reinforcement Learning for Large Language Models

Resumen Técnico: SPS (Steering Probability Squeezing)

1. El Problema: Limitaciones de la Exploración en RL para LLMs

El aprendizaje por refuerzo (RL) se ha convertido en un paradigma fundamental para entrenar modelos de lenguaje grandes (LLMs) orientados al razonamiento, utilizando señales de recompensa basadas en reglas (como la corrección de respuestas matemáticas). Sin embargo, el artículo identifica una limitación crítica en los métodos actuales de RL (como GRPO):

• Enfoque en Pass@1 vs. Pass@k: El RL tiende a mejorar drásticamente la tasa de éxito de una sola muestra (Pass@1), pero ofrece una exploración limitada de trayectorias de razonamiento diversas, lo que resulta en un rendimiento deficiente en escenarios de múltiples muestras (Pass@k).
• El Efecto de "Apretado" (Squeezing Effect): Los autores descubren que el RL introduce un sesgo sistemático en la redistribución de la masa de probabilidad. Cuando se aplican actualizaciones de gradiente negativo a respuestas de baja probabilidad, la masa de probabilidad eliminada no se redistribuye equitativamente hacia alternativas prometedoras. En su lugar, es absorbida desproporcionadamente por la respuesta dominante (la más probable o "greedy").
• Consecuencia: Esto provoca una concentración excesiva de la distribución de salida en un subconjunto estrecho de trayectorias de alta recompensa, reduciendo la entropía y limitando la capacidad del modelo para descubrir nuevas y mejores rutas de razonamiento.

2. Metodología: SPS (Steering Probability Squeezing)

Para mitigar este efecto de apretado y fomentar una exploración genuina, los autores proponen SPS, un nuevo paradigma de entrenamiento que intercala el RL convencional con Aprendizaje por Refuerzo Inverso (IRL).

• Concepto Central: En lugar de permitir que la masa de probabilidad se concentre ciegamente en las trayectorias dominantes, SPS "redirige" este apretado hacia regiones subexploradas que podrían contener soluciones correctas.
• Flujo de Entrenamiento Iterativo:
  1. Fase de RL (On-Policy): Se realiza un entrenamiento estándar (ej. GRPO) para explorar el conjunto de datos y generar rollouts (muestras de respuesta).
  2. Fase de IRL (Inverse RL): Se utilizan los rollouts generados en la fase anterior como "demostraciones" (sin supervisión externa ni recompensas adicionales). Se emplea un objetivo de KL Forward para remodelar la distribución de la política.
     • El objetivo es que la política actual se alinee con la distribución empírica de los rollouts, redistribuyendo la masa de probabilidad desde las respuestas excesivamente dominantes hacia trayectorias menos probables pero potencialmente valiosas.
  3. Estrategia Iterativa: Este ciclo RL-IRL se repite, permitiendo una redistribución progresiva de la masa de probabilidad y evitando la convergencia prematura a comportamientos "greedy".
• Énfasis en Trayectorias de Baja Probabilidad (L2TE): Para potenciar la exploración, SPS incluye una estrategia que muestrea preferentemente trayectorias con baja probabilidad inicial durante la fase de IRL, amplificando la señal de aprendizaje de soluciones subexploradas.

3. Contribuciones Clave

1. Análisis de la Dinámica de Aprendizaje: Identificación y formalización teórica del "efecto de apretado" (squeezing effect) como una limitación intrínseca del RL basado en recompensas binarias, que restringe la exploración y establece un límite superior empírico para el Pass@k.
2. Propuesta de SPS: Desarrollo de un marco de entrenamiento que utiliza IRL para remodelar explícitamente la distribución de trayectorias inducida, facilitando una exploración más amplia sin introducir supervisión externa.
3. Validación Empírica: Demostración de que alternar entre RL e IRL es un camino efectivo para extender la capacidad de exploración de modelos de razonamiento, superando a los métodos de RL puros.

4. Resultados Experimentales

Los autores evaluaron SPS en cinco benchmarks de matemáticas de nivel olímpico (AIME, BRUMO, HMMT) utilizando modelos de la serie Qwen2.5-Math (1.5B y 7B parámetros).

• Mejora en Pass@k: SPS logró mejoras consistentes y sustanciales en Pass@128 en comparación con baselines como GRPO, DAPO y GSPO.
  • Ejemplo destacado: En el modelo Qwen2.5-Math-1.5B, SPS alcanzó un puntaje de 63.33 en Pass@128 en el benchmark BRUMO, una mejora de +10.00 puntos respecto al GRPO estándar.
• Eficiencia en Datos Limitados: SPS demostró ser particularmente efectivo en configuraciones con pocos datos (ej. 3k muestras), donde los modelos pequeños suelen sufrir más de sobreajuste y colapso de entropía.
• Análisis de Diversidad: Las métricas de similitud de trayectorias mostraron que SPS genera respuestas más diversas que el RL puro, confirmando que la mejora no es solo un ajuste de probabilidad, sino una verdadera expansión de las rutas de razonamiento.
• Costo Computacional: La fase de IRL añade un sobrecosto computacional marginal (aprox. 3% del tiempo total de iteración), ya que no requiere entrenamiento de modelos de recompensa externos ni inferencia costosa adicional.

5. Significado e Impacto

• Reevaluación del RL en LLMs: El trabajo desafía la noción de que el RL es suficiente por sí solo para la exploración en razonamiento complejo, señalando que sin mecanismos de control de la distribución, el RL tiende a colapsar la diversidad.
• Nueva Dirección de Investigación: Introduce el uso de IRL no para imitar expertos externos, sino como una herramienta interna para "desapretar" (steer) la distribución de probabilidad generada por el propio modelo, preservando la masa de probabilidad en trayectorias diversas.
• Escalabilidad: Al demostrar mejoras significativas en modelos pequeños (1.5B) y grandes (7B) con un costo computacional bajo, SPS ofrece una vía práctica y escalable para mejorar la capacidad de razonamiento de los LLMs en tareas de alta complejidad.

En conclusión, SPS representa un avance metodológico al integrar la redistribución de probabilidad mediante IRL dentro del ciclo de RL, resolviendo el problema fundamental de la falta de exploración y permitiendo a los modelos descubrir un espectro más amplio de soluciones correctas.