Resumen Técnico: MONA – Optimizador de Muón con Aceleración de Nesterov

1. Planteamiento del Problema

El entrenamiento de modelos de lenguaje grandes (LLM) exige cada vez más optimizadores con una eficiencia de muestra superior. Aunque AdamW ha sido el estándar durante más de una década, han surgido alternativas recientes como Muon. Muon mejora los métodos de primer orden tratando las matrices de pesos como unidades geométricas únicas y aplicando ortogonalización de matrices (mediante la iteración Newton-Schulz) al búfer de momento. Este enfoque se alinea con el descenso de máxima pendiente bajo la norma espectral, ofreciendo aproximadamente un $2\times$ de eficiencia computacional sobre AdamW en el entrenamiento de Mezclas de Expertos (MoE) a gran escala.

Sin embargo, Muon comparte una limitación fundamental con todos los métodos de gradiente de primer orden: carece de mecanismos explícitos para explorar el paisaje de pérdida más allá de la información local del gradiente. En consecuencia, Muon puede quedar atrapado en mínimos locales agudos, particularmente en escenarios de entrenamiento con lotes grandes. Los mínimos agudos se asocian con una generalización deficiente, mientras que los mínimos más planos tienden a rendir un mejor desempeño en tareas posteriores. Aunque el optimizador ALTO introdujo un término de aceleración basado en el promedio móvil exponencial (EMA) de las diferencias de gradiente para escapar de los mínimos agudos, su diseño como un adaptador general para Adam y LAMB presenta desafíos de integración con el proceso de ortogonalización altamente no lineal de Muon.

2. Metodología: MONA

Los autores proponen MONA (Optimizador de Muón con Aceleración de Nesterov), que integra sin fisuras un término de aceleración consciente de la curvatura directamente en el marco de ortogonalización de Muon.

Algoritmo Central

MONA modifica la tubería estándar de Muon introduciendo un término de aceleración antes de la acumulación de momento y la ortogonalización. El proceso de actualización para un parámetro matricial $W$ en la iteración $k$ procede de la siguiente manera:

1. Cálculo de la Diferencia de Gradiente: Calcule la diferencia entre el gradiente estocástico actual $G_k$ y el gradiente anterior $G_{k-1}$:
   $$D_k = G_k - G_{k-1}$$
2. Actualización del Búfer de Aceleración: Mantenga un promedio móvil exponencial de estas diferencias:
   $$A_k = \beta_a A_{k-1} + (1 - \beta_a) D_k$$
   Aquí, $\beta_a$ controla la memoria de la aceleración.
3. Gradiente Consciente de la Curvatura: Construya un gradiente aumentado $\tilde{G}_k$ sumando el término de aceleración escalado al gradiente crudo:
   $$\tilde{G}_k = G_k + \alpha A_k$$
   El coeficiente $\alpha$ escala la aceleración. La motivación teórica sugiere que $G_k - G_{k-1}$ aproxima $H_k(\theta_k - \theta_{k-1})$, donde $H_k$ es el Hessiano. Por lo tanto, el término apunta lejos de los mínimos agudos (donde los cambios de gradiente son grandes) y permanece pequeño en regiones planas.
4. Ortogonalización de Muon: El gradiente aumentado se alimenta en los pasos estándar de momento y ortogonalización de Muon:
   $$M_k = \mu M_{k-1} + \tilde{G}_k$$
   $$O_k = \text{Newton-Schulz}(M_k)$$
   $$W_{k+1} = W_k - \eta (\gamma O_k + \lambda W_k)$$

Fundamento del Diseño

Aplicar el término de aceleración antes de la ortogonalización es crítico. Si se aplicara después, interrumpiría la estructura ortogonal de la que depende Muon para la regularización de la norma espectral. Al transformar primero el gradiente crudo en una dirección consciente de la curvatura, el búfer de momento captura estas direcciones enriquecidas, las cuales luego se procesan a través de la tubería estándar de Muon.

Para parámetros de valor vectorial (por ejemplo, incrustaciones, sesgos), MONA utiliza AdamW por defecto, adherido a las convenciones existentes de Muon.

3. Análisis Teórico

El artículo proporciona un análisis de convergencia para MONA bajo supuestos estándar (suavidad L, gradientes no sesgados con varianza acotada y alineación direccional esperada).

• Tasa de Convergencia: Los autores demuestran que MONA alcanza una tasa de convergencia de $O(K^{-1/2})$ en entornos no convexos, igualando las tasas asintóticas de SGD y AdamW. El término de aceleración afecta las constantes en el límite, pero no la tasa asintótica.
• Escape de Mínimos Agudos: El análisis formaliza cómo el término de aceleración resiste selectivamente los mínimos agudos. Cerca de un punto estacionario, la aceleración esperada $E[A_k]$ es proporcional a la suma ponderada por el Hessiano de las direcciones pasadas. Para direcciones con valores propios grandes (mínimos agudos), la magnitud de la aceleración es grande, empujando al optimizador hacia afuera. Para valores propios pequeños (direcciones planas), la aceleración es pequeña, permitiendo una convergencia estable. Este mecanismo preserva la regularización de la norma espectral de Muon mientras añade conciencia de la curvatura.

4. Resultados Empíricos

MONA fue evaluado en tres escalas de modelos de preentrenamiento MoE (1B, 6B y 68B parámetros) y una etapa de ajuste fino supervisado (SFT).

Rendimiento de Preentrenamiento

• Modelos: MOE-1B-A0d2B (400B tokens), MOE-6B-A0d5B (1.2T tokens) y MOE-68B-A3B (700B tokens).
• Métricas: Se monitoreó la pérdida de validación en los dominios de código, razonamiento matemático, inglés general y texto académico chino.
• Hallazgos: MONA logró consistentemente una pérdida de validación inferior tanto a Muon como a AdamW en las tres escalas de modelos y en todos los dominios. La brecha de rendimiento fue particularmente evidente en el modelo más grande de 68B, donde MONA demostró una convergencia más estable.

Benchmarks de Tareas Posteriores (MOE-68B-A3B)

El modelo de 68B entrenado con MONA fue evaluado en una suite completa de benchmarks:

• Capacidad General: MONA obtuvo la puntuación promedio más alta (0.4557) en benchmarks que incluyen MMLU, MMLU-Pro, CMMLU, CEVAL y BBH, superando a Muon (0.4478) y AdamW (0.4382).
• Razonamiento Matemático y Generación de Código: MONA mostró un rendimiento superior en GSM8K, GPQA y varias pruebas de generación de código (HumanEval+, MBPP+, BigCodeBench, LiveCodeBench). Notablemente, MONA logró ganancias significativas en tareas que requieren representaciones de código transferibles.
• Ajuste Fino Supervisado (SFT): Cuando el checkpoint preentrenado con MONA se sometió a SFT en datos de código, mantuvo una pérdida de entrenamiento más baja durante todo el proceso en comparación con la línea base preentrenada con Muon. Las evaluaciones posteriores al SFT en la suite BigCode confirmaron que los modelos inicializados con MONA obtuvieron puntuaciones más altas en 6 de las 7 tareas en comparación con los modelos inicializados con Muon.

Eficiencia y Despliegue

Para abordar la sobrecarga de memoria de almacenar el gradiente anterior ($G_{k-1}$) y el búfer de aceleración ($A_k$), los autores introdujeron MONA-Lite:

• Cuantización: Almacenar estados auxiliares en bfloat16 (BF16) en lugar de float32.
• Cálculo en Flujo: Calcular las diferencias de gradiente in situ para eliminar la necesidad de un búfer dedicado $G_{k-1}$.
• Resultado: Estas estrategias reducen la sobrecarga de memoria adicional en aproximadamente un 75% sin sacrificar la calidad del entrenamiento. El análisis de la sobrecarga computacional mostró que MONA es solo ~1% más lento a nivel del paso del optimizador, sin diferencia práctica en el tiempo de iteración de extremo a extremo en comparación con Muon.

5. Significado y Afirmaciones

El artículo afirma que MONA cierra exitosamente la brecha entre la eficiencia geométrica de Muon y la necesidad de aceleración consciente de la curvatura. Su significado principal radica en:

1. Integración Algorítmica: Demostrar que los términos de aceleración pueden integrarse en marcos de ortogonalización matricial sin interrumpir sus propiedades geométricas.
2. Escalabilidad: Validar el método en una amplia gama de tamaños de modelo (de 1B a 68B parámetros) y conteos de tokens (hasta 1 billón de tokens), mostrando mejoras consistentes sobre AdamW y Muon.
3. Generalización: Proporcionar evidencia empírica de que escapar de los mínimos agudos durante el preentrenamiento conduce a una mejor generalización en tareas posteriores, particularmente en razonamiento complejo y generación de código.
4. Practicidad: Ofrecer una variante eficiente en memoria (MONA-Lite) que hace viable el optimizador para entornos con recursos limitados.

Los autores concluyen que MONA representa un optimizador de última generación para el entrenamiento de LLM a gran escala, logrando una convergencia superior y un mejor rendimiento en tareas posteriores, manteniendo al mismo tiempo la eficiencia computacional del marco de Muon.