Resumen Técnico: Cómo Comienzas es Cómo Razonas: Impulsando la Exploración en RLVR mediante Priors Ajustados por Prefijos

1. Declaración del Problema

El Aprendizaje por Refuerzo con Recompensas Verificables (RLVR) ha surgido como un paradigma poderoso para mejorar las capacidades de razonamiento de los Modelos de Lenguaje Grande (LLM) en tareas como la resolución de problemas matemáticos y la demostración de teoremas. Sin embargo, los autores identifican una limitación crítica en los enfoques actuales de RLVR: el colapso de la entropía.

Aunque el RLVR a menudo mejora la precisión de una sola ejecución (Pass@1), frecuentemente falla en expandir la cobertura de las trayectorias de razonamiento exitosas. Esto se manifiesta como una reducción en la diversidad de salida, donde el modelo converge a un conjunto estrecho de rutas de razonamiento. Las técnicas de exploración existentes, como la regularización de entropía o las recompensas intrínsecas, a menudo tratan la distribución del modelo base como fija e intentan modificar el objetivo de RL o el procedimiento de muestreo. El artículo argumenta que estos métodos son insuficientes porque el RL está fuertemente moldeado por la distribución previa del modelo base. Además, las técnicas de exploración pasiva a menudo degradan la calidad de la generación, resultando en ejecuciones ruidosas que no contribuyen a un razonamiento efectivo.

El desafío central es fomentar una exploración efectiva—descubrir comportamientos de razonamiento diversos, relevantes para la tarea y de alta calidad—sin colapsar la política en trayectorias redundantes o de baja calidad, y sin el costo computacional de actualizar los parámetros completos del modelo.

2. Metodología: Marco IMAX

Los autores proponen Maximización de Información con Exploración Aumentada (IMAX), un marco ligero que desplaza el objetivo de optimización desde los parámetros del LLM hacia el ajuste de prompts suaves. En lugar de modificar el proceso de RL sobre un prior fijo, IMAX aprende un conjunto de prefijos suaves que remodelan activamente el prior del modelo sobre las trayectorias de razonamiento.

2.1 Razonamiento Condicionado por Prefijos

El mecanismo central implica un conjunto de $C$ prefijos suaves (tokens virtuales continuos) $\mathcal{E} = \{e_1, \dots, e_C\}$. Para un prompt de entrada dado $x$, un prefijo $e_z$ se antepone a la secuencia de incrustaciones de entrada. El LLM base permanece congelado, y la política $\pi_E(y|x, z)$ se condiciona tanto al prompt como a la identidad del prefijo seleccionado $z$.

• Inicialización: Los prefijos se inicializan a partir de un conjunto de prompts de texto semánticamente significativos (por ejemplo, "Introducir variables explícitamente", "Particionar en casos") generados por un LLM externo. Esto asegura que los prefijos comiencen con estrategias de razonamiento distintas.
• Control: Cada prefijo actúa como una "perilla de control", induciendo una distribución de ejecución distinta y guiando al modelo hacia diferentes regiones de su espacio de razonamiento (por ejemplo, soluciones algebraicas vs. diagramáticas vs. similares a código).

2.2 Recompensa de Maximización de Información (InfoMax)

Para evitar que el conjunto aprendido de prefijos colapse en controles redundantes (donde diferentes prefijos inducen trayectorias idénticas), los autores introducen una recompensa InfoMax. El objetivo es maximizar la información mutua (MI) entre la identidad del prefijo $Z$ y la trayectoria generada $Y$, condicionado al prompt $X$:
$$I(Z; Y | X) = H(Y | X) - H(Y | X, Z)$$
Intuitivamente, esto fomenta que diferentes prefijos produzcan comportamientos de razonamiento válidos y distinguibles para la misma pregunta.

2.3 Límite Inferior Variacional y Entrenamiento

Calcular directamente la información mutua es intratable. Los autores derivan un límite inferior variacional utilizando una posterior variacional $q_\phi(z | x, y)$, que es un clasificador ligero entrenado para predecir qué prefijo generó una respuesta específica $y$ dado el prompt $x$.

• Forma Inversa: Utilizan la descomposición inversa de la MI: $I(Z; Y | X) = H(Z | X) - H(Z | X, Y)$.
• Formulación de la Recompensa: La recompensa InfoMax se define como $\log q_\phi(z | x, y)$. El objetivo total para el conjunto de prefijos se convierte en:
  $$\max_E \mathbb{E} [R(x, y) + \beta \log q_\phi(z | x, y)]$$
  donde $R(x, y)$ es la recompensa verificable estándar (correctitud) y $\beta$ es un hiperparámetro.

2.4 Optimización Alternada

El procedimiento de entrenamiento sigue un esquema de optimización alternada (Inferencia Variacional de Ascenso de Coordenadas):

1. Actualizar Posterior ($q_\phi$): Con el conjunto de prefijos fijo, actualizar el clasificador $q_\phi$ para predecir la identidad del prefijo a partir del par prompt-respuesta (generado sin el prefijo suave para evitar la fuga de información).
2. Actualizar Conjunto de Prefijos ($E$): Con $q_\phi$ fijo, actualizar los prefijos suaves utilizando un algoritmo de Optimización de Política Relativa por Grupos (GRPO) aumentado con la recompensa InfoMax.
3. Muestreo Estratificado: Durante el entrenamiento, cada prompt se empareja con identidades de prefijo distintas para asegurar una cobertura uniforme del espacio de prefijos dentro de cada grupo GRPO.

3. Contribuciones Clave

1. Paradigma Ligero de RLVR: El artículo introduce un método para mejorar el razonamiento optimizando un conjunto de prompts suaves en lugar de los parámetros completos del modelo. Esto remodela la distribución de generación condicional de un modelo congelado sin la sobrecarga computacional del ajuste fino completo.
2. Recompensa InfoMax para la Diversidad: Los autores proponen una recompensa InfoMax a nivel de respuesta que fomenta explícitamente que diferentes prefijos induzcan trayectorias de razonamiento distinguibles y relevantes para la tarea, evitando el colapso del conjunto de prefijos en controles redundantes.
3. Validación Empírica: Extensos experimentos en puntos de referencia de razonamiento matemático (GSM8K, MATH-500, Minerva, AMC) a través de tres escalas de modelo (Qwen2.5-1.5B, Qwen3-4B, Qwen3-8B) demuestran que IMAX mejora tanto la calidad del razonamiento como la diversidad.

4. Resultados Experimentales

Los autores evaluaron IMAX frente a varias líneas base, incluido GRPO de política completa, inicialización con prompts de texto y otros métodos de RLVR orientados a la diversidad (DIVER, MERCI).

• Mejoras de Rendimiento: IMAX superó consistentemente a las líneas base estándar de RLVR ajustadas por prefijos.
  • En Qwen3-4B, IMAX mejoró Pass@4 en hasta 11.60% y Avg@4 en 10.57% en comparación con las líneas base estándar en puntos de referencia específicos.
  • En Qwen3-8B, DAPO+IMAX mostró mejoras significativas sobre DAPO-prefix en GSM8K, MATH-500 y Minerva.
• Diversidad y Cobertura:
  • Análisis t-SNE: Las visualizaciones de las incrustaciones de respuesta mostraron que los prefijos inducidos por IMAX resultaron en grupos más separados en comparación con otros métodos, indicando una mayor diversidad intra-grupo.
  • Componentes de Razonamiento: Los mapas de calor revelaron que diferentes prefijos indujeron con éxito patrones de razonamiento distintos (por ejemplo, el Prefijo 1 favoreciendo la manipulación algebraica, el Prefijo 2 favoreciendo la partición de casos y la verificación), alineándose con su inicialización semántica.
  • Rendimiento del Oráculo: El rendimiento combinado de múltiples prefijos (Oráculo) superó al mejor prefijo individual, confirmando que los prefijos cubren regiones exitosas complementarias del espacio de problemas.
• Eficiencia: IMAX requiere un costo de entrenamiento significativamente menor (3.0x–6.0x menos horas de GPU) en comparación con GRPO de parámetros completos, ya que solo actualiza las incrustaciones ligeras de prefijos y la cabeza del clasificador.

5. Significado y Afirmaciones

El artículo afirma que cómo un modelo comienza un proceso de razonamiento (el prior) dicta fundamentalmente cómo razona. Al tratar los prompts suaves como variables de control reutilizables, IMAX ofrece una solución escalable al problema de exploración en RLVR.

• Agnóstico del Algoritmo: El marco está diseñado para integrarse perfectamente en pipelines de RLVR existentes (por ejemplo, GRPO, DAPO).
• Diversidad Estructurada: A diferencia de los métodos que simplemente aumentan la aleatoriedad, IMAX fomenta una diversidad de razonamiento estructurada, asegurando que la exploración conduzca a rutas semánticamente distintas y de alta calidad.
• Limitaciones Modestas: Los autores reconocen que el método está diseñado principalmente para tareas de razonamiento, mostrando ganancias limitadas en tareas generales de seguimiento de instrucciones. Además, la efectividad depende de la calidad de los prompts de texto iniciales, lo que puede requerir diseño manual o reajuste entre dominios. Los prefijos aprendidos también introducen una sobrecarga persistente, aunque pequeña, durante la inferencia.

En resumen, IMAX demuestra que optimizar el "comienzo" del proceso de generación mediante prefijos suaves es una estrategia altamente efectiva para impulsar la exploración y mejorar las capacidades de razonamiento en LLMs, ofreciendo una alternativa computacionalmente eficiente al aprendizaje por refuerzo de política completa.