Resumen Técnico: Entrenamiento Conjunto de Predicción Multi-Token en Aprendizaje por Refuerzo mediante Calibración Óptima de Coeficientes

Declaración del Problema

El Aprendizaje por Refuerzo a partir de Recompensas Verificables (RLVR) se ha convertido en el paradigma estándar para mejorar las capacidades de razonamiento de los Modelos de Lenguaje Grandes (LLM). Simultáneamente, la Predicción Multi-Token (MTP), un objetivo auxiliar entrenado durante el preentrenamiento para predecir múltiples tokens futuros, ha demostrado beneficios significativos en la eficiencia de muestras y como modelo borrador para la decodificación especulativa.

Intuitivamente, la optimización conjunta de MTP y los objetivos de RL durante el post-entrenamiento debería permitir que los modelos se beneficien de la supervisión multi-token mientras adquieren capacidades de razonamiento. Sin embargo, la evidencia empírica muestra que el entrenamiento conjunto MTP-RL a menudo conduce a una degradación significativa del rendimiento en el modelo principal. En consecuencia, las prácticas actuales (por ejemplo, en marcos como veRL y slime) suelen desconectar los gradientes de MTP del modelo principal durante el entrenamiento de RL. Esta desconexión obliga a los parámetros de MTP a adaptarse únicamente mediante un ajuste fino post-hoc independiente, sacrificando las ganancias potenciales de la supervisión multi-token simultánea. El problema central abordado es: ¿Por qué falla el entrenamiento conjunto MTP-RL y se puede diseñar una estrategia fundamentada para habilitar un entrenamiento conjunto efectivo?

Metodología

Análisis Teórico

Los autores abordan el problema desde una perspectiva de optimización, analizando el efecto por paso de MTP sobre el objetivo de RL $J(\theta)$. Asumiendo una función objetivo $L$-suave, derivan una cota inferior para la mejora por paso. Esta descomposición revela que el efecto de MTP está gobernado por dos términos en competencia:

1. Término de Correlación de Primer Orden: Determinado por el acuerdo direccional (producto interno) entre el gradiente de RL ($g_{RL}$) y el gradiente de MTP ($g_{MTP}$).
2. Penalización de Perturbación de Segundo Orden: Un término no positivo que refleja la perturbación introducida por la magnitud del gradiente de MTP.

La mejora del rendimiento ocurre solo si el término de correlación supera al término de penalización.

Análisis de Regímenes Existentes

Utilizando esta descomposición, el artículo unifica y explica tres regímenes de entrenamiento existentes:

• Desconexión (Detach): $g_{MTP} = 0$ con respecto al modelo principal. $\Delta_{MTP} = 0$. Esto es seguro pero renuncia a las ganancias potenciales.
• Pérdida de Entropía Cruzada (CE): El término de correlación está cerca de cero en expectativa porque RL pondera más las muestras de alta recompensa mientras que CE trata todas las muestras por igual. En consecuencia, la penalización negativa de segundo orden domina, causando degradación.
• Pérdida de Política: Utiliza la misma formulación de objetivo de RL para MTP. Inicialmente, los gradientes están alineados y el término de correlación domina, generando ganancias. Sin embargo, a medida que avanza el entrenamiento y el modelo entra en regiones planas del paisaje de pérdida, la correlación decae mientras que la norma del gradiente (y por tanto la penalización) persiste. Esto conduce a una curva de rendimiento de "subida-luego-bajada" donde un coeficiente fijo no puede rastrear la transición de fase.

Solución Propuesta: Calibración Óptima de Coeficientes (OCC)

Para abordar la transición de fase, los autores proponen la Calibración Óptima de Coeficientes (OCC), un esquema adaptativo que rastrea el coeficiente óptimo de MTP ($\lambda$) en línea.

1. Óptimo en Forma Cerrada: El análisis produce una expresión en forma cerrada para el $\lambda^*$ óptimo en cada paso, que es proporcional a la relación entre la correlación de gradientes ($c$) y la norma al cuadrado del gradiente ($v^2$).
2. Proxy de Log-Probabilidad: Calcular gradientes de modelo completo para $c$ y $v^2$ es prohibitivamente costoso en sistemas distribuidos a gran escala. OCC utiliza un proxy de log-probabilidad derivado del cambio en la log-probabilidad bajo actualizaciones de pequeño paso ($\delta \approx \log \pi_\theta - \log \pi_{old}$).
   • Proxy de Alineación ($\hat{c}_t$): Estima el producto interno.
   • Proxy de Varianza ($\hat{v}^2_t$): Estima la norma al cuadrado.
3. Actualización Adaptativa: El coeficiente se actualiza como $\lambda_t = \lambda_+ \cdot \frac{\hat{c}_t}{\hat{v}^2_t + \epsilon}$. Esto permite que OCC ajuste dinámicamente el peso de MTP, manteniendo ganancias positivas durante todo el entrenamiento sin calcular gradientes completos.

Contribuciones Clave

1. Unificación Teórica: El artículo proporciona un análisis teórico que unifica los regímenes de Desconexión, Entropía Cruzada y Pérdida de Política, explicando su éxito o fracaso basándose en la interacción entre la correlación de primer orden y la perturbación de segundo orden.
2. Algoritmo OCC: Propone la Calibración Óptima de Coeficientes, un procedimiento adaptativo que calibra el coeficiente de MTP en línea utilizando un proxy de log-probabilidad. Este método requiere una sobrecarga computacional despreciable y evita el cálculo de gradientes de modelo completo.
3. Validación Empírica: Extensos experimentos en múltiples modelos (MiMo-7B, GLM-4.5-Air) y algoritmos (DAPO, GSPO) en benchmarks de razonamiento matemático de nivel competitivo demuestran que OCC coincide o supera consistentemente la línea base de Desconexión.

Resultados Experimentales

Los autores evaluaron OCC en seis benchmarks: AIME24, AIME25, AMC, MATH, Minerva y OlympiadBench.

• Rendimiento: En todos los benchmarks y escalas de modelo (7B denso y 106B MoE), OCC superó consistentemente a la línea base de Desconexión y a la Pérdida de Política con coeficiente fijo.
  • En MiMo-7B con DAPO, OCC logró una precisión promedio del 61.7%, superando a Desconexión (58.9%) y a Pérdida de Política (57.4%).
  • En GLM-4.5-Air, OCC logró una precisión promedio del 67.6% frente al 65.9% de Desconexión.
• Dinámica de Entrenamiento: Mientras que la Pérdida de Política mostró una tendencia de subida-luego-bajada (degradándose en etapas posteriores), OCC mantuvo mejoras estables durante todo el entrenamiento. La Pérdida CE rindió consistentemente por debajo de Desconexión.
• Eficiencia: OCC introdujo una sobrecarga despreciable. El tiempo de entrenamiento por paso en tiempo real para OCC (8.07s) fue comparable a Desconexión (8.11s), mientras que la retropropagación de modelo completo requirió 45.06s (una ralentización de 5.6×).
• Ablación: El análisis de sensibilidad mostró que el factor de escala predefinido $\lambda_+$ es robusto en un rango de valores, con $\lambda_+ = 1.0$ produciendo el mejor rendimiento promedio.

Significado y Afirmaciones

El artículo afirma que el fracaso del entrenamiento conjunto MTP-RL no es inherente, sino que surge de una discrepancia entre coeficientes estáticos y la deriva dinámica de la alineación de gradientes durante el entrenamiento. Al proporcionar una descomposición teórica del paisaje de optimización y proponer OCC, los autores demuestran que MTP puede ser reintroducido de forma segura en el post-entrenamiento de RL.

El significado radica en permitir que los modelos adquieran simultáneamente capacidades de razonamiento mediante RL y aprovechen representaciones multi-token sin sacrificar el rendimiento o incurrir en costos computacionales prohibitivos. El trabajo sugiere que la calibración adaptativa es esencial para el entrenamiento conjunto en paisajes de aprendizaje profundo no convexos, ofreciendo una alternativa fundamentada al estándar industrial actual de desconexión de gradientes.

Limitaciones

Los autores señalan que su análisis teórico se basa en suposiciones de $L$-suavidad que pueden no cumplirse exactamente en paisajes no convexos, proporcionando orientación cualitativa en lugar de cuantitativa. El proxy de log-probabilidad es una aproximación de primer orden que puede perder fidelidad bajo tasas de aprendizaje grandes. Además, los experimentos se centran en el razonamiento matemático con recompensas verificables; la generalización a configuraciones como RLHF con modelos de recompensa aprendidos o generación abierta queda por validar.