Reformulate LLM Reinforcement Learning for Efficient Training under Black-box Discrepancy

Resumen Técnico: Reformulación del Aprendizaje por Refuerzo de LLM para un Entrenamiento Eficiente bajo Discrepancia de Caja Negra

1. Planteamiento del Problema

El Aprendizaje por Refuerzo (RL) se ha convertido en un paradigma crítico de post-entrenamiento para mejorar las capacidades de razonamiento de los Grandes Modelos de Lenguaje (LLMs). Sin embargo, el post-entrenamiento mediante RL sufre frecuentemente de un rendimiento subóptimo impredecible o colapsos totales del entrenamiento. Investigaciones recientes atribuyen estos fallos a una omnipresente discrepancia entre entrenamiento e inferencia (o desajuste).

Esta discrepancia surge porque las demandas industriales obligan a menudo a los profesionales a utilizar motores y formatos de precisión dispares (por ejemplo, diferentes kernels, estrategias de paralelización o precisiones numéricas como BF16 frente a FP8) para los mismos pesos del modelo durante el entrenamiento frente a la inferencia. Esto crea una desviación oculta entre las distribuciones de probabilidad de acción de la política de entrenamiento ($\pi_{Train}$) y la política de inferencia ($\pi_{Inf}$).

Las estrategias de mitigación existentes presentan limitaciones significativas:

• Alineación de Precisión Uniforme: El simple hecho de reducir la precisión del entrenamiento para que coincida con la de la inferencia (por ejemplo, a FP16) reduce el rango de valores representables y exacerba los riesgos de explosión de gradientes.
• Intervenciones Heurísticas: Métodos como el enmascaramiento de gradientes de tokens con alta discrepancia simplemente suprimen el impacto de los datos ruidosos. No abordan la causa raíz y pasan por alto el objetivo final: optimizar la política para el despliegue de inferencia, no solo para el proceso de entrenamiento.
• Alineación de Infraestructura: Intentar alinear los operadores de bajo nivel y los stacks de hardware es laborioso y a menudo inviable debido a la miríada de casos particulares.

El desafío central es permitir que la política de entrenamiento perciba dinámicamente y anule activamente su discrepancia con la política de inferencia sin sacrificar el objetivo primario de maximización de la recompensa.

2. Metodología: MDP con Restricción de Discrepancia (DCMDP)

Los autores proponen un cambio de paradigma de los Procesos de Decisión de Markov (MDP) estándar a un MDP con Restricción de Discrepancia (DCMDP). Este marco trata la discrepancia entre entrenamiento e inferencia como un problema de optimización de caja negra, permitiendo que el sistema gestione configuraciones heterogéneas de entrenamiento e inferencia de forma fluida.

Componentes Clave:

• Métrica de Discrepancia:
  Los autores evalúan empíricamente tres señales de discrepancia a nivel de token: Diferencia de Probabilidad ($\delta_{diff}$), Ratio de Probabilidad ($\delta_{ratio}$) y Diferencia de Log-Probabilidad ($\delta_{log}$). Identifican que la diferencia absoluta de probabilidad ($\delta_{diff} = |p_{Train} - p_{Inf}|$) es la penalización más efectiva. Ofrece una regularización "plástica" que corrige la deriva relevante para el despliegue sin restringir excesivamente la exploración impulsada por la recompensa, a diferencia de las penalizaciones basadas en ratios, que pueden ser demasiado conservadoras para tokens de baja probabilidad.

• La Región de Tolerancia de Discrepancia:
  Un hallazgo empírico crítico es la existencia de una región de tolerancia de la discrepancia. Penalizar agresivamente las discrepancias menores y despreciables al inicio del entrenamiento suprime la exploración de la política y reduce la eficiencia del aprendizaje. Por el contrario, reducir la discrepancia excesiva fuera de esta región mejora la consistencia de la optimización.

  • Mecanismo: La formulación DCMDP solo aplica penalizaciones cuando la discrepancia excede un umbral predefinido $c$. Dentro de esta región de tolerancia, la política tiene permitido explorar libremente para maximizar recompensas.

• Optimización de Doble Objetivo mediante Relajación de Lagrange:
  El problema se formula como la maximización de la recompensa esperada sujeta a una restricción sobre la discrepancia promedio a nivel de token:
  $$\max_{\theta} J(\theta) = \mathbb{E}_{\tau \sim \pi_{\theta}} \left[ \sum r(o_t) \right]$$
  $$\text{sujeto a } \frac{1}{|\tau|} \sum |\pi_{Train}(o_t) - \pi_{Inf}(o_t)| \leq c$$

Para resolver esto, los autores emplean la Relajación de Lagrange. Un multiplicador aprendible $\lambda$ ajusta dinámicamente el peso de la penalización de la discrepancia:
* Si la discrepancia viola el límite de tolerancia ($>c$), $\lambda$ aumenta, fortaleciendo la penalización para devolver la política a la zona segura.
* Si la discrepancia está dentro del límite, $\lambda$ se atenúa hacia cero, permitiendo que la política se concentre en la maximización de la recompensa.

El objetivo final (Dual-GRPO) inyecta esta penalización adaptativa directamente en la ventaja normalizada por grupo:
$$\hat{A}^{DCMDP}_{i,t} = \hat{A}_{i,t} - \lambda_k \cdot \delta_{i,t}$$

3. Contribuciones Clave

1. Compensación Autónoma de la Discrepancia: El artículo demuestra que la incorporación de una penalización robusta de caja negra (diferencia de probabilidad absoluta a nivel de token) dota a la política de entrenamiento con la capacidad de erradicar autónomamente el riesgo de colapso del entrenamiento inducido por la discrepancia.
2. Cambio de Paradigma a DCMDP: Los autores reformulan el RL de LLM como un MDP con Restricción de Discrepancia. Este enfoque equilibra óptimamente los objetivos duales de maximizar la recompensa de la tarea y minimizar la discrepancia entre entrenamiento e inferencia, yendo más allá de la simple escalarización de objetivos.
3. Paradigma de Entrenamiento Heterogéneo: El método permite una configuración de entrenamiento donde los modelos se optimizan en entornos de alta fidelidad y uso intensivo de recursos (por ejemplo, precisión BF16) mientras se alinean explícitamente para despliegues de inferencia de bajo costo y recursos limitados (por ejemplo, FP8). Esto convierte la heterogeneidad de la infraestructura de una desventaja en un controlable knob de optimización.

4. Resultados Experimentales

El método fue evaluado en modelos Qwen3-8B (denso) y Qwen3-30B-A3B (Mixture-of-Experts) a través de seis benchmarks de razonamiento matemático (AIME24/25, AMC23, MATH-500, Minerva, OlympiadBench).

• Entrenamiento/Inferencia BF16: DC-GRPO superó consistentemente al baseline de GRPO vanilla en todos los benchmarks para ambos modelos, logrando puntuaciones promedio más altas (por ejemplo, Qwen3-8B mejoró de 67.3 a 71.0). Esto confirma que imponer restricciones de discrepancia mejora, en lugar de obstaculizar, la optimización de la recompensa primaria.
• Entrenamiento Heterogéneo (BF16) + Despliegue (FP8): En este escenario, el GRPO vanilla sufrió una degradación significativa de rendimiento o colapso debido al aumento del gap entre entrenamiento e inferencia. DC-GRPO mitigó esto con éxito, preservando y a menudo superando el rendimiento del baseline de solo BF16.
• Estudios de Ablación: El método demostró ser robusto a las variaciones en el multiplicador de Lagrange inicial ($\lambda_0$) y la tasa de aprendizaje dual ($\eta_\lambda$). Se encontró que el presupuesto de tolerancia ($c$) es crítico; establecerlo demasiado pequeño causó una sobre-regularización, mientras que establecerlo demasiado grande causó que el método degenerara en un GRPO vanilla.

5. Significado y Reivindicaciones

El artículo afirma que DCMDP proporciona un marco algorítmico unificado para mitigar la inestabilidad entre entrenamiento e inferencia sin requerir optimizaciones exhaustivas y personalizadas de los stacks de hardware o software subyacentes. Al tratar la discrepancia como una señal de retroalimentación de caja negra, el método permite:

• Entrenamiento Estable: Previniendo los colapsos abruptos comunes en el post-entrenamiento de RL.
• Exploración Eficiente: Permitiendo que la política explore libremente dentro de una región de discrepancia segura.
• Eficiencia de Recursos: Permitiendo el uso de configuraciones de entrenamiento de alta fidelidad para simular y adaptarse a entornos de despliegue restringidos, facilitando así despliegues de modelos en el mundo real más eficientes.

Los autores reconocen limitaciones, señalando que su alcance empírico está actualmente limitado a las precisiones BF16 y FP8 y familias de modelos específicas, requiriéndose trabajos futuros para extender el método a otras combinaciones de backend y arquitecturas de modelos.