Resumen Técnico: Dropout-GRPO para el Razonamiento Latente Continuo

1. Planteamiento del Problema

El artículo aborda una incompatibilidad estructural fundamental entre la Optimización de Política Relativa de Grupo (GRPO) y los modelos de razonamiento latente continuo (específicamente la arquitectura COCONUT).

• El Requisito de GRPO: GRPO depende de la diversidad de $K$ rollouts dentro de un grupo para calcular una ventaja de media de grupo no nula ($A^{(k)} = r^{(k)} - \mu_r$). Si todos los rollouts producen trayectorias idénticas, la ventaja colapsa a cero, deteniendo la optimización.
• El Desafío del Razonamiento Latente: A diferencia de los modelos estándar de Cadena de Pensamiento (CoT) que derivan la diversidad mediante el muestreo estocástico de tokens, los modelos de razonamiento latente continuo (como COCONUT) operan mediante recurrencia de estados ocultos determinista. Dado un prompt y parámetros fijos, múltiples rollouts producen trayectorias latentes e identidades de respuesta idénticas.
• La Consecuencia: Aplicar GRPO estándar a estos modelos resulta en una varianza cero ($\sigma_r = 0$), lo que provoca que el gradiente de la política se desvanezca y se impida la generación de cualquier señal de aprendizaje.

2. Metodología: Dropout-GRPO

Los autores proponen Dropout-GRPO, un método que inyecta la estocasticidad necesaria en el proceso de razonamiento latente a través de un dropout estructurado, manteniendo al mismo tiempo las garantías teóricas de la inferencia variacional.

Mecanismo Central: Dropout Variacional de Máscara Compartida

En lugar del muestreo estándar a nivel de token, el método introduce estocasticidad a nivel de parámetro:

1. Generación de la Máscara: Para cada rollout $k$, se extrae una única máscara de Bernoulli $\xi^{(k)}$.
2. Consistencia: Esta máscara se mantiene constante a través de todos los $T$ pasos de recurrencia latente dentro de ese rollout específico.
3. Perturbación: La máscara induce una perturbación estructurada $\tilde{\theta}(\xi)$ a los pesos de la red. La recurrencia latente se convierte en $h_t = f_{\tilde{\theta}(\xi)}(h_{t-1})$.
4. Replay de la Máscara: Crucialmente, la semilla del RNG para la máscara se almacena durante la fase de rollout. Durante la actualización de la política, la misma máscara se regenera y se aplica. Esto asegura que la trayectoria evaluada para la recompensa sea bit-idéntica a la trayectoria utilizada para el cálculo del gradiente cuando $\theta = \theta_{old}$.

Interpretación Teórica

• Promedio de Modelo Bayesiano: Siguiendo a Gal y Ghahramani, la máscara compartida trata cada rollout como una única muestra posterior de una distribución estructurada sobre los parámetros. La política marginal $\bar{\pi}_\theta$ se interpreta como un promedio de modelo bayesiano.
• Insesgadez del Gradiente: Los autores demuestran que el gradiente sustituto utilizando la ventaja de media única ($A^{(k)} = r^{(k)} - \mu_r$) es un estimador insesgado del gradiente de la recompensa esperada de la política marginal, escalado por $(1 - 1/K)$.
• Reducción de Varianza: Al reutilizar la misma máscara (Números Aleatorios Comunes) durante la actualización que la utilizada durante el rollout, el método reduce la varianza en comparación con el muestreo de una nueva máscara en el momento de la actualización.

Refinamientos de Implementación

Para estabilizar el entrenamiento, se introducen cuatro refinamientos específicos:

1. KL $k_3$ Huberizada: Reemplaza el truncamiento estándar en el estimador de la divergencia KL con una extensión de tipo Huber para evitar que los gradientes se desvanezcan para tokens con grandes log-ratios, evitando "tokens escapados" que derivan sin regularización.
2. Recocido de Región de Confianza (Trust-Region Annealing): El coeficiente KL $\beta$ se recoce proporcionalmente con la tasa de aprendizaje $\eta_t$ para evitar que la política quede anclada en un equilibrio donde la tracción de la KL cancele el gradiente de la política.
3. Filtro de Precisión a Nivel de Grupo: Los grupos con una precisión empírica $\mu_r$ fuera de una ventana configurada (por ejemplo, $[0.20, 0.95]$) son descartados. Esto actúa como un currículo emergente, enfocando el cómputo en prompts con varianza no trivial.
4. Salto Coordinado por DDP: Asegura la consistencia del entrenamiento distribuido al descartar grupos.

3. Contribuciones Clave

1. Nuevo Paradigma para el RL Latente: El artículo identifica la naturaleza determinista de la recurrencia latente como la causa del fallo de GRPO y propone el dropout de máscara compartida como una solución mínima y teóricamente limpia. Es el primer método que aplica con éxito el RL relativo de grupo a modelos de razonamiento latente continuo.
2. Justificación Teórica:
   • Demuestra la definición de la Jacobiana latente $\partial h_T / \partial \theta$ bajo la regularidad estándar de Transformer.
   • Demuestra la insesgadez del gradiente sustituto para la política marginal.
   • Demuestra que el replay de la máscara proporciona una reducción de varianza mediante números aleatorios comunes.
3. Validación Empírica: Aplica con éxito GRPO al modelo COCONUT en el benchmark GSM8K, un régimen donde el GRPO determinista falla completamente.
4. Implementación de Referencia: Libera el código que realiza el replay de la máscara bit-idéntica a lo largo de toda la recurrencia latente.

4. Resultados Experimentales

Los experimentos se realizaron utilizando Qwen2.5-1.5B como modelo base con una profundidad latente de $T=6$.

• Mejora de Rendimiento: Dropout-GRPO mejoró la precisión GSM8K pass@1 de un baseline de COCONUT del 27.29% al 29.01%.
• Recuperación de Capacidades: El método revirtió una regresión de rendimiento de 1.33 puntos en el conjunto de datos SVAMP causada por el entrenamiento de currículo, recuperando el rendimiento al 44.00% (igualando la varianza del modelo base).
• Estudios de Ablación:
  • GRPO Determinista: Falló completamente, produciendo una trayectoria de entrenamiento plana idéntica a la inicialización de SFT debido a la ventaja cero.
  • REINFORCE + Dropout (EMA): Falló debido a la alta varianza; el gradiente de un solo rollout era demasiado ruidoso para estabilizar la inicialización de SFT.
  • Dropout-GRPO: La única configuración que produjo una señal de aprendizaje viable, confirmando que las bases relativas al grupo son esenciales para la reducción de la varianza en este entorno.

5. Significado y Reivindicaciones

El artículo posiciona este trabajo como un enfoque práctico y teóricamente fundamentado para el post-entrenamiento de modelos de lenguaje de gran escala (LLM) con razonamiento latente.

• Factibilidad: Demuestra que el post-entrenamiento mediante RL es factible para modelos que razonan directamente en estados ocultos, una capacidad previamente bloqueada por la falta de diversidad en los rollouts.
• Mecanismo: Replantea el dropout no como un regularizador (que a menudo se desactiva en los LLM modernos), sino como una fuente de estocasticidad exógena estructurada necesaria para la optimización variacional de la política.
• Limitaciones: Los autores son modestos respecto a la magnitud de la mejora (una ganancia de 1.72 puntos en GSM8K), atribuyéndolo al pequeño tamaño de la columna vertebral (1.5B) y la escasez de conjuntos de datos de razonamiento de pasos concisos adecuados. Observan que la política desplegada (inferencia sin dropout) difiere de la política marginal optimizada, lo que introduce un sesgo dependiente de las constantes de Lipschitz.

En resumen, Dropout-GRPO cierra la brecha entre el aprendizaje por refuerzo relativo de grupo y el razonamiento latente continuo, permitiendo la optimización de modelos que anteriormente no podían utilizar GRPO debido a sus dinámicas internas deterministas.