Uniform-Correct Policy Optimization: Breaking RLVR's Indifference to Diversity

1. Planteamiento del Problema

El Aprendizaje por Refuerzo con Recompensas Verificables (RLVR) ha mejorado significativamente la precisión de un solo intento (Pass@1) de los Modelos de Lenguaje Grande (LLMs) en tareas de razonamiento como las matemáticas. Sin embargo, se ha identificado un efecto secundario crítico: Colapso de la Diversidad.

• El Fenómeno: Aunque RLVR aumenta la probabilidad de generar una respuesta correcta, a menudo reduce drásticamente el Pass@K (la probabilidad de que al menos una de las $K$ respuestas muestreadas sea correcta). Esto indica que el espacio de soluciones del modelo se contrae, concentrando la masa de probabilidad en un subconjunto estrecho de soluciones correctas mientras suprime otros caminos de razonamiento válidos.
• La Causa Raíz: Los autores identifican una limitación estructural en los objetivos estándar de RLVR (como GRPO - Group Relative Policy Optimization). Estos objetivos son indiferentes a cómo se distribuye la masa de probabilidad entre las soluciones correctas.
  • Matemáticamente, cualquier política que coloque el 100% de su masa en el conjunto de salidas correctas $Y^+$ produce la misma recompensa máxima, independientemente de si esa masa está concentrada en una sola solución (colapsada) o distribuida uniformemente entre todas las soluciones válidas (diversa).
  • Combinado con la dinámica de muestreo en política, esta indiferencia crea un bucle de retroalimentación de auto-refuerzo: las soluciones muestreadas ligeramente más a menudo reciben actualizaciones de gradiente más grandes, aumentando aún más su probabilidad y frecuencia de muestreo, lo que lleva a un colapso de modos irreversible.

2. Metodología

A. Análisis Teórico

Los autores formalizan el mecanismo de colapso y definen la estructura de la política óptima:

1. Mecanismo de Colapso: Demuestran que bajo GRPO, el sesgo de muestreo conduce a conteos de muestras desiguales. Las actualizaciones de gradiente escalan con estos conteos, amplificando la probabilidad de las soluciones dominantes y suprimiendo las minoritarias. Una vez que una solución cae por debajo de un umbral de muestreo ($\pi(y) \ll 1/K$), no recibe señal de gradiente y no puede recuperarse.
2. Definición de la Política Óptima: Los autores proponen que la política óptima debe ser la Política Uniforme-Correcta (UCP). Esta política satisface dos criterios:
   • Robustez: Maximiza la masa de probabilidad retenida en el peor caso si un subconjunto de soluciones correctas se deshabilita (por ejemplo, debido a un cambio de distribución).
   • Optimalidad Regularizada por Entropía: Maximiza de manera única el objetivo de equilibrar la corrección con la diversidad (entropía) dentro del conjunto correcto.
   • Definición: La UCP coloca probabilidad cero en salidas incorrectas y distribuye la probabilidad uniformemente entre todas las soluciones correctas ($\pi(y|x) = 1/|Y^+|$ para $y \in Y^+$).

B. Solución Propuesta: Uniform-Correct Policy Optimization (UCPO)

Para lograr la Política Uniforme-Correcta, los autores proponen UCPO, una modificación al objetivo de GRPO.

• Función Objetivo: UCPO augmenta la maximización estándar de corrección con una penalización de uniformidad condicional.
  $$\mathcal{L}_{UCPO} = \underbrace{\log Z_\theta(x)}_{\text{Masa de Corrección}} - \tau \underbrace{KL(u(\cdot|x) \parallel q_\theta(\cdot|x))}_{\text{Penalización de Uniformidad}}$$

  • $Z_\theta(x)$: Masa de probabilidad total en salidas correctas.
  • $q_\theta(\cdot|x)$: La distribución condicional de la política sobre salidas correctas.
  • $u(\cdot|x)$: La distribución uniforme sobre el conjunto de salidas correctas.
  • $\tau$: Un hiperparámetro que controla la fuerza de la penalización de uniformidad.

• Mecanismo de Gradiente:

  • GRPO estándar pondera los gradientes por la frecuencia de muestreo actual (favoreciendo modos dominantes).
  • UCPO repondera la señal de gradiente para contrarrestar este sesgo. Interpola entre la distribución condicional actual y el objetivo uniforme.
  • Ponderación Inversa: El estimador de gradiente asigna pesos más altos a las soluciones correctas subrepresentadas (aquellas con baja probabilidad actual) y pesos más bajos a las dominantes. Esto redistribuye efectivamente el "presupuesto de ventaja" para amplificar las respuestas correctas minoritarias sin sacrificar la masa total de corrección.

3. Contribuciones Clave

1. Identificación de la Indiferencia: El artículo demuestra formalmente que los objetivos estándar de RLVR son indiferentes a la diversidad intra-conjunto-correcto, permitiendo que la dinámica estocástica de entrenamiento impulse un colapso de modos irreversible.
2. Caracterización Teórica: Establece la Política Uniforme-Correcta como la única solución óptima bajo criterios de robustez y regularización por entropía.
3. Diseño de Algoritmo (UCPO): Deriva UCPO, una modificación simple pero efectiva a GRPO que impone la Política Uniforme-Correcta como su óptimo único.
4. Redistribución de Gradientes: Introduce un mecanismo para redistribuir las señales de gradiente hacia respuestas correctas subrepresentadas, rompiendo el ciclo de colapso de auto-refuerzo.

4. Resultados Experimentales

Los autores evaluaron UCPO en tres modelos (1.5B, 7B parámetros) y cinco benchmarks de razonamiento matemático (AIME24/25, AMC, MATH, OlympiadBench).

• Mejora en Pass@K: UCPO mejora consistentemente el Pass@K (específicamente Pass@64) mientras mantiene un Pass@1 competitivo.
  • AIME24 (modelo 7B): Logró una mejora absoluta del +10% en Pass@64 en comparación con GRPO estándar.
  • Mejora Promedio: En todos los modelos, UCPO mejoró el Pass@64 promedio en +3.5 a +4.8 puntos sobre GRPO.
• Métricas de Diversidad:
  • UCPO logró hasta un 45% más de diversidad a nivel de ecuación dentro del conjunto correcto en comparación con GRPO.
  • A diferencia de la regularización global de entropía (que a menudo dispersa masa hacia respuestas incorrectas), UCPO mantiene una alta masa de corrección ($Z_\theta \approx 1$) mientras diversifica el conjunto correcto.
• Comparación con Líneas Base: UCPO superó a otros métodos de preservación de diversidad como Regularización de Entropía (Ent-Reg), control de Covarianza-KL y Focal GRPO. Se demostró que la regularización global de entropía es frágil y a menudo degrada el Pass@1 o no logra mejorar la diversidad dentro del conjunto correcto.
• Eficiencia: UCPO introduce una sobrecarga computacional insignificante (solo reponderación escalar de ventajas) y no requiere rodajes adicionales ni modelos de referencia.

5. Significado

Este trabajo aborda una falla fundamental en los paradigmas actuales de entrenamiento RLVR para modelos de razonamiento. Al cambiar el enfoque de simplemente "encontrar una respuesta correcta" a "mantener una distribución diversa de todas las respuestas correctas", UCPO:

• Previene la Contracción del Espacio de Soluciones: Asegura que los modelos no "olviden" caminos de razonamiento válidos durante el ajuste fino.
• Mejora la Robustez: Una política que cubre múltiples soluciones correctas es más robusta a cambios de distribución o ruido en el verificador.
• Proporciona un Objetivo Principiado: Va más allá de bonos de entropía heurísticos hacia un objetivo teóricamente fundamentado que apunta explícitamente a la estructura de la política óptima.

En resumen, UCPO ofrece una solución escalable y teóricamente sólida al problema del colapso de diversidad en RLVR, permitiendo que los LLMs logren una mayor cobertura de múltiples muestras sin sacrificar la precisión de un solo intento.