Resumen Técnico: SAGE – Moldeando Anclajes para la Exploración Guiada en RLVR de LLMs

1. Declaración del Problema

Los avances recientes en Aprendizaje por Refuerzo con Recompensas Verificables (RLVR) han demostrado la capacidad de mejorar la precisión pass@1 de los Modelos de Lenguaje Grande (LLMs) en tareas de razonamiento. Sin embargo, ha surgido una limitación crítica: mientras que RLVR a menudo aumenta las tasas de éxito de una sola muestra, frecuentemente falla en generar ganancias comparables en pass@k (la probabilidad de encontrar una respuesta correcta entre $k$ muestras).

Análisis previos sugieren que esta discrepancia surge porque el RLVR estándar induce colapso de modos (o "retracción del soporte"). En lugar de descubrir nuevas estrategias de razonamiento, la política colapsa hacia un subconjunto estrecho de modos de razonamiento preexistentes y de alta densidad presentes en el modelo de referencia base. Este comportamiento se atribuye a las restricciones estructurales del término de regularización estándar de Kullback-Leibler inverso (reverse-KL). Aunque el reverse-KL proporciona la estabilidad de entrenamiento necesaria al anclar la política a una distribución de referencia ($\pi_{ref}$), suprime inherentemente las trayectorias de baja densidad, impidiendo efectivamente que el modelo explore caminos de razonamiento alternativos y válidos que se encuentran fuera del soporte inmediato del modelo de referencia.

Los intentos de resolver esto eliminando el término KL o reemplazándolo con forward-KL han demostrado ser subóptimos. Eliminar el término KL a menudo conduce a la manipulación de recompensas (reward hacking) e inestabilidad, mientras que forward-KL, aunque expande el soporte, frecuentemente asigna masa de probabilidad a regiones irrelevantes para la recompensa, fallando en mejorar el pass@1. El desafío central es retener los beneficios estabilizadores del reverse-KL mientras se habilita una expansión controlada del soporte empírico.

2. Metodología: Marco SAGE

Los autores proponen SAGE (Moldeando Anclajes para la Exploración Guiada), un marco que reinterpreta la regularización reverse-KL no como una restricción estática que debe debilitarse, sino como un ancla dinámica cuya distribución objetivo puede remodelarse para guiar la exploración.

Mecanismo Central

El RLVR reverse-KL estándar optimiza:
$$J(\theta) = \mathbb{E}_{x,y \sim \pi_\theta} [r(x, y) - \beta D_{KL}(\pi_\theta \parallel \pi_{ref})]$$
Esto conduce a una solución estacionaria donde la política es una inclinación exponencial de la referencia: $\pi_\theta \propto \pi_{ref} \exp(r/\beta)$. Esta actualización multiplicativa amplifica las regiones de alta densidad de $\pi_{ref}$ mientras suprime las de baja densidad.

SAGE modifica esto introduciendo una función guía $q(x, y) > 0$ que remodela la propia distribución de anclaje. El nuevo objetivo se convierte en:
$$J'(\theta) = \mathbb{E}_{x,y \sim \pi_\theta} [r(x, y) - \beta \tilde{D}_{KL}(\pi_\theta \parallel q \cdot \pi_{ref})]$$
Aquí, $\tilde{D}_{KL}$ es una divergencia pseudo-KL definida para funciones no normalizadas. La solución estacionaria se convierte en:
$$\pi_\theta(y | x) \propto q(x, y) \pi_{ref}(y | x) \exp\left(\frac{1}{\beta} r(x, y)\right)$$
Al diseñar $q(x, y)$ para que sea grande en modos de razonamiento subexplorados pero potencialmente válidos, SAGE amplifica su masa de probabilidad sin romper la estabilidad proporcionada por el ancla $\pi_{ref}$.

Instanciación mediante Señales Intrínsecas

Para hacer $q(x, y)$ práctica sin requerir información de recompensa explícita durante la construcción de la guía, los autores proponen el uso de señales intrínsecas derivadas del propio proceso de razonamiento:

1. Exploración Aleatoria: Una línea base que utiliza perturbaciones estocásticas.
2. Exploración a Nivel de Token: Utiliza la sorpresa ($-\log p$) para sesgar hacia elecciones léxicas raras. Los autores notan que esto captura desviaciones a nivel superficial pero puede no reflejar diferencias estructurales en el razonamiento.
3. Exploración a Nivel de Rama (SAGE-Rama): La variante principal propuesta. Utiliza la entropía ($H_t$) como señal para la ramificación del razonamiento. La función guía se define como:
   $$q_{Branch}(y_t | x, y_{<t}) = 1 + \gamma (H_t - \tau)_+$$
   donde $H_t$ es la entropía a nivel de token, $\tau$ es un umbral de activación y $\gamma$ controla la magnitud. Esto amplifica selectivamente la exploración en puntos de decisión de alta incertidumbre (nodos de ramificación) mientras suprime el ruido en pasos simbólicos deterministas.

Garantías Teóricas

El artículo proporciona un análisis teórico (Teorema 4.3) que demuestra que, a diferencia del RLVR estándar que está limitado por el soporte empírico del modelo de referencia, SAGE puede teóricamente expandir el soporte empírico. Específicamente, si la función guía $q$ es suficientemente grande para una solución válida en términos de recompensa $y^*$ que actualmente está fuera del soporte empírico de la política de referencia, SAGE puede asignarle una probabilidad no despreciable, recuperando efectivamente modos de razonamiento "ocultos".

3. Contribuciones Clave

1. Análisis de las Limitaciones de RLVR: Los autores identifican que el ancla reverse-KL, aunque estabilizadora, es la causa estructural principal del colapso de modos en RLVR. Argumentan que simplemente eliminar o reemplazar el término KL es insuficiente, ya que interrumpe el equilibrio entre eficiencia y cobertura.
2. Marco SAGE: Proponen un método principiado para habilitar una expansión controlada del soporte empírico remodelando la distribución ancla KL mediante una función guía $q(x, y)$, en lugar de eliminar el ancla.
3. Integración de Señales Intrínsecas: Demuestran cómo señales intrínsecas ligeras (específicamente entropía y sorpresa) pueden utilizarse para construir funciones guía efectivas que no requieren pasadas forward adicionales ni datos externos.
4. Validación Empírica: Experimentos extensivos muestran que las variantes de SAGE, particularmente la exploración a nivel de rama, mejoran consistentemente tanto el pass@1 como el pass@k en diversos benchmarks de razonamiento matemático.

4. Resultados Experimentales

Los autores evaluaron SAGE en los conjuntos de datos AIME 2024-2025, AMC23 y MATH-500 utilizando el modelo Qwen2.5-Math-7B-Base y el algoritmo GRPO.

• Ganancias de Rendimiento: La variante GRPO + Rama logró el mejor rendimiento promedio, mejorando el pass@1 de 0.266 (GRPO vanilla) a 0.284 y el pass@256 de 0.714 a 0.729. Superó a las líneas base en los tres conjuntos de datos.
• Comparación con Alternativas:
  • GRPO sin KL: Mostró recompensas de entrenamiento más altas pero un rendimiento significativamente peor en tiempo de prueba (pass@1 y pass@256), confirmando el riesgo de manipulación de recompensas sin el ancla KL.
  • Forward-KL: Mejoró ligeramente el pass@256 pero no logró igualar las ganancias de pass@1 de SAGE, indicando que asigna masa a regiones fuera del objetivo.
  • Otros Algoritmos de RL: SAGE se integró exitosamente con BNPO y DAPO, resolviendo consistentemente el compromiso donde estos métodos típicamente mejoraban el pass@1 a costa del pass@k.
• Análisis de Patrones de Razonamiento: El análisis cualitativo reveló que, mientras que el RLVR estándar reforzaba patrones dominantes como "configuración de restricciones", SAGE aumentó significativamente la frecuencia de la prueba por contradicción (de 0.8% en GRPO a 3.3% en GRPO+Rama), un modo de razonamiento raro pero válido.
• Robustez: En evaluaciones fuera de distribución (benchmark de Caballeros y Caballeros Mentirosos), los modelos entrenados con SAGE mantuvieron el rendimiento bajo cambios de distribución, mientras que los modelos GRPO estándar exhibieron una degradación significativa, sugiriendo que SAGE mitiga el colapso de modos.
• Generalización: El método también fue efectivo en modelos de razonamiento largo (DeepSeek-R1-Distill-Qwen-7B) y modelos más pequeños (Qwen2.5-1.5B/3B), demostrando robustez a través de diferentes escalas de modelos y paradigmas de razonamiento.

5. Significado y Afirmaciones

El artículo afirma que SAGE ofrece un enfoque práctico y principiado para mitigar el afilamiento de modos en el RLVR reverse-KL. Al tratar la regularización KL como un instrumento controlable en lugar de una restricción rígida, SAGE permite a los LLMs descubrir y utilizar diversas trayectorias de razonamiento que son válidas pero subrepresentadas en el modelo base.

Los autores enfatizan que SAGE es modular y ortogonal a las técnicas existentes (como el modelado de ventajas o el muestreo dinámico), ya que modifica la propia distribución de anclaje en lugar de la estructura de optimización. Posicionan a SAGE como un marco general para la expansión estable y compatible con KL del soporte empírico.

Limitaciones y Alcance:
Los autores reconocen modestamente que SAGE opera dentro del soporte de la política de referencia fija. Redistribuye la masa de probabilidad entre trayectorias que ya tienen una probabilidad no nula bajo $\pi_{ref}$; no puede generar patrones de razonamiento completamente nuevos que sean imposibles de producir para el modelo base. Notan que, aunque la destilación de modelos maestros más fuertes puede expandir la frontera de razonamiento, SAGE la complementa permitiendo una exploración estable dentro de esa frontera. La instanciación actual depende de señales intrínsecas ligeras, y los autores sugieren que aprender funciones guía más expresivas a partir de modelos más fuertes es una dirección prometedora para trabajos futuros.