Resumen Técnico: Más Allá del Colapso de Modos: Coincidencia de Distribuciones para el Razonamiento Diverso

1. Declaración del Problema

El artículo identifica el colapso de modos como una limitación fundamental en los métodos de aprendizaje por refuerzo (RL) en política, particularmente aquellos utilizados para entrenar Modelos de Razonamiento a Gran Escala (LRMs) mediante Aprendizaje por Refuerzo con Recompensas Verificables (RLVR). Aunque métodos como la Optimización de Políticas Relativa por Grupos (GRPO) han logrado éxitos en la resolución de problemas matemáticos y la generación de código, muestran una tendencia a concentrar la masa de probabilidad en una única trayectoria de alta recompensa una vez descubierta. Este comportamiento suprime estrategias de solución alternativas, detiene la exploración y conduce a un razonamiento frágil que no generaliza.

Los autores argumentan que este fenómeno no es meramente un artefacto de implementación, sino una consecuencia directa del objetivo de optimización: el RL en política estándar minimiza implícitamente la divergencia Kullback-Leibler (KL) inversa ($D_{KL}(\pi_\theta \| p^*)$). La KL inversa es inherentemente "buscadora de modos", amplificando la primera trayectoria exitosa encontrada en lugar de mantener una distribución sobre múltiples soluciones de alta recompensa diversas. En contraste, la divergencia KL directa ($D_{KL}(p^* \| \pi_\theta)$) exhibe un comportamiento "cubridor de modos", fomentando que la política coincida con el soporte completo de la distribución objetivo, pero suele ser intratable porque requiere muestreo de la distribución global ponderada por recompensas.

2. Metodología: DMPO

Para abordar el colapso de modos sin requerir muestreo global intratable, los autores proponen la Optimización de Políticas de Coincidencia de Distribuciones (DMPO).

Idea Central

DMPO aproxima la minimización de la KL directa a nivel de grupo en lugar de globalmente. En lugar de muestrear de una distribución objetivo global, el algoritmo construye una distribución objetivo sobre un grupo específico de trayectorias muestreadas para un prompt dado.

Formulación Algorítmica

1. Distribución Objetivo a Nivel de Grupo: Para un grupo de $G$ trayectorias $\mathcal{O} = \{o_1, \dots, o_G\}$ muestreadas de la política actual $\pi_{\theta_{old}}$, se construye una distribución objetivo $p(o_i | \mathcal{O})$ proporcional al exponencial de sus recompensas (una distribución de Boltzmann):
   $$p(o_i | \mathcal{O}) = \frac{\exp(r(o_i)/\alpha)}{\sum_{j=1}^G \exp(r(o_j)/\alpha)}$$
   Esto es computacionalmente tratable ya que solo requiere la función de partición a nivel de grupo.

2. Distribución de Política a Nivel de Grupo: Para alinear la política con este objetivo, el método define una distribución de política $q_\theta(o_i | \mathcal{O})$ basada en log-verosimilitudes normalizadas por longitud para tener en cuenta el sesgo de la longitud de la trayectoria:
   $$q_\theta(o_i | \mathcal{O}) = \frac{\exp(\varphi(o_i))}{\sum_{j=1}^G \exp(\varphi(o_j))}$$
   donde $\varphi(o_i)$ es la probabilidad logarítmica promedio por token.

3. Pérdida de Coincidencia de Distribuciones: La política se alinea con el objetivo minimizando la divergencia entre $q_\theta$ y $p$. Los autores utilizan el Error Cuadrático Medio (MSE) en lugar de la divergencia KL para estabilidad numérica, evitando la explosión de gradientes cuando la probabilidad de la política se aproxima a cero:
   $$\mathcal{L}_{DM}(\theta) = \frac{1}{G} \sum_{i=1}^G (p(o_i | \mathcal{O}) - q_\theta(o_i | \mathcal{O}))^2$$

4. Objetivo Unificado: El objetivo final combina el término estándar de maximización de recompensas de GRPO con la regularización de coincidencia de distribuciones:
   $$\mathcal{L}_{DMPO}(\theta) = \mathcal{L}_{GRPO}(\theta) + \lambda \mathcal{L}_{DM}(\theta)$$
   Aquí, $\mathcal{L}_{GRPO}$ impulsa la explotación (maximización de recompensas), mientras que $\mathcal{L}_{DM}$ impone la exploración al evitar que la política colapse a un único modo dentro de la región de alta recompensa.

3. Contribuciones Clave

A. Algoritmo DMPO

El artículo presenta DMPO como una solución principada y práctica al colapso de modos. Al aproximar la KL directa a nivel de grupo, logra un comportamiento cubridor de modos sin requerir una función de partición aprendida ni muestreo global. El método es compatible con los marcos existentes de GRPO, requiriendo únicamente la adición de un término de regularización.

B. MM-NP-Bench

Los autores introducen MM-NP-Bench, un benchmark multimodal diseñado para evaluar rigurosamente la exploración en modelos de razonamiento.

• Alcance: Extiende el NP-Bench basado en texto para incluir representaciones visuales de 10 tareas de optimización combinatoria NP-difíciles (por ejemplo, TSP, Coloreado de Grafos, Cobertura de Vértices).
• Evaluación con Doble Métrica: Para hacer observable el colapso de modos, el benchmark utiliza dos métricas:
  • Tasa de Éxito (SR): Mide la satisfacción de restricciones (factibilidad).
  • Ratio de Calidad (QR): Mide la calidad de optimización relativa a una línea base heurística.
  • Significado: El colapso de modos se manifiesta como una alta SR pero baja QR (encontrar soluciones válidas pero subóptimas).
• Infraestructura: Incluye generadores paramétricos, verificadores basados en reglas y solucionadores heurísticos para proporcionar líneas base casi óptimas.

C. Validación Empírica

El artículo proporciona experimentos extensos que demuestran que DMPO supera a líneas base sólidas (GRPO, GSPO, GPG, FlowRL, etc.) en múltiples dominios.

4. Resultados

Rendimiento en Tareas de Optimización

• MM-NP-Bench (Multimodal): DMPO logró un Ratio de Calidad del 43.1%, una mejora relativa del 12% sobre GRPO (38.4%). También logró la Tasa de Éxito más alta (61.9% frente a 55.7%).
• NP-Bench (Solo texto): DMPO logró un Ratio de Calidad del 43.9%, una mejora relativa del 9% sobre GRPO (40.1%).
• Ablación: El método mostró robustez ante elecciones de hiperparámetros ($\lambda \in [0.3, 0.7]$ y $\alpha \in [0.8, 1.2]$).

Generalización y Transferencia

• Razonamiento Matemático: El entrenamiento en tareas de optimización con DMPO mejoró el rendimiento en benchmarks de razonamiento matemático (AIME, AMC, etc.) en un 2.0% en comparación con GRPO.
• Tareas Fuera de Dominio: DMPO demostró una mejora del 2.3% en benchmarks de razonamiento visual fuera de dominio (LogicVista, MathVista, etc.), lo que sugiere que el entrenamiento que preserva la diversidad mejora las capacidades de razonamiento general.

Dinámica de Entrenamiento

El análisis de los registros de entrenamiento reveló que, mientras GRPO sufre de "colapso de longitud" (degeneración en respuestas cortas y triviales) y mesetas prematuras de recompensa, DMPO sostiene una mejora continua y preserva cadenas de razonamiento robustas (~400 tokens). DMPO mantiene una entropía más baja que GRPO, indicando "diversidad estructurada" (ejecución confiable de diferentes estrategias lógicas) en lugar de confusión estocástica.

5. Significado y Afirmaciones

El artículo afirma establecer la coincidencia de distribuciones como un enfoque práctico y principado para prevenir el colapso de modos en el RL en política. Los autores argumentan que:

1. Identificación de la Causa Raíz: El colapso de modos en RLVR es fundamentalmente causado por la naturaleza buscadora de modos de la minimización de la KL inversa.
2. Solución Factible: La minimización de la KL directa, tradicionalmente considerada intratable, puede aproximarse efectivamente a nivel de grupo para preservar la diversidad de soluciones.
3. Beneficio General: Preservar la diversidad durante el entrenamiento en tareas específicas de optimización se transfiere a capacidades de razonamiento general mejoradas, sugiriendo que la exploración es un componente crítico del razonamiento robusto en Modelos de Razonamiento a Gran Escala.

El trabajo posiciona a DMPO no meramente como un regularizador de diversidad, sino como un mecanismo que permite el descubrimiento de soluciones de mayor calidad que los objetivos buscadores de modos pasan por alto, mejorando así tanto el rendimiento de optimización como la robustez del razonamiento a través de modalidades.