Resumen Técnico: Por qué Muon supera a Adam: Una perspectiva de la curvatura

Declaración del problema

Muon ha emergido como un optimizador altamente eficiente para el preentrenamiento de Modelos de Lenguaje de Gran Escala (LLM), demostrando una convergencia de entrenamiento aproximadamente 2$\times$ más rápida en comparación con el optimizador estándar Adam en diversas escalas de modelos. Mientras que trabajos recientes han atribuido esta ventaja a mecanismos como la memoria asociativa y la naturaleza de cola larga de los datos, la fuente geométrica local de la superioridad de Muon sigue sin estar clara. Este artículo aborda dos preguntas primarias:

1. ¿Qué propiedad específica del paisaje de optimización respalda la ventaja de Muon sobre Adam?
2. ¿Cómo influyen los factores de preentrenamiento, específicamente la distribución de los datos de entrenamiento y la estructura del modelo, en esta propiedad?

Metodología

Los autores emplean una combinación de análisis empírico sobre trayectorias de entrenamiento de LLM y análisis teórico sobre problemas de optimización cuadrática estilizados.

Análisis Empírico:

• Descomposición de Taylor de segundo orden: Los autores analizan la disminución de la pérdida en un paso ($\Delta D$) descomponiéndola en una ganancia de primer orden (alineación del gradiente) y una penalización de curvatura de segundo orden utilizando una expansión de Taylor: $\Delta D \approx \langle G, Z \rangle - \frac{1}{2}\langle Z, H[Z] \rangle$.
• Definición de Métrica: Para aislar el efecto de la dirección de la actualización del escalar de la actualización, introducen la Nitidez Direccional Normalizada (NDS), definida como $SF(W; Z) = \langle Z, H[Z] \rangle / \|Z\|_F^2$. Esto permite que la penalización de curvatura se factorice en la norma de la actualización y la NDS.
• Configuración Experimental: Los experimentos se realizan en un modelo NanoGPT de 124 millones de parámetros entrenado con el conjunto de datos FineWeb. Las comparaciones entre Muon y Adam se realizan en niveles de pérdida de validación emparejados para asegurar la equidad.
• Variaciones de Datos y Estructura:
  • Desequilibrio de Datos: Se genera datos sintéticos utilizando una Gramática Libre de Contexto Probabilística de Zipf (PCFG) con niveles de desequilibrio controlados (exponente de Zipf $s$) para probar el efecto de las distribuciones de cola pesada.
  • Descomposición de Capas: La NDS se descompone en contribuciones dentro de la capa (bloques de Hessiano diagonales) y entre capas (interacciones fuera de la diagonal) para comprender cómo diferentes componentes del modelo contribren a la ventaja de la curvatura.

Análisis Teórico:

• Modelo Cuadrático Estilizado: Los autores construyen un problema de optimización cuadrática que refleja las características del paisaje de los LLM, basado en cuatro supuestos verificados empíricamente:
  1. Bajo Rango de Kronecker: El Hessiano admite una aproximación de Kronecker de bajo rango.
  2. Diagonalización Simultánea: Los factores de Kronecker comparten bases ortogonales comunes.
  3. Heterogeneidad de Curvatura: El espectro del Hessiano exhibe una fuerte heterogeneidad (distribución de valores propios de cola larga).
  4. Alineación de Gradiente: Los gradientes se alinean fuertemente con las direcciones de alta curvatura.
• Comparación: Las pruebas teóricas comparan Muon (actualizaciones espectralmente normalizadas) contra el Descenso de Gradiente (GD). Se señala que Adam se comporta de manera similar a GD en este entorno estilizado específico, lo que permite a los autores centrarse en el contraste entre la normalización espectral y las actualizaciones basadas en gradiente.

Contribuciones Clave y Resultados

1. La Fuente de la Ventaja de Muon: Reducción de la Penalización de Curvatura
El análisis revela que Muon logra una mayor disminución de la pérdida en un paso que Adam principalmente porque incurre en una penalización de curvatura de segundo orden significativamente menor, no porque logre una mayor ganancia de gradiente de primer orden.

• Observación 1: En pérdidas de validación emparejadas, Muon y Adam muestran disminuciones de primer orden comparables, pero la penalización de curvatura de Muon es sustancialmente menor.

2. El Mecanismo: Menor Nitidez Direccional Normalizada (NDS)
La reducción de la penalización de curvatura es impulsada por la dirección de la actualización de Muon y no por su tamaño de paso.

• Observación 2: Muon y Adam tienen normas de actualización ($\|Z\|_F$) comparables. Sin embargo, Muon exhibe una NDS notablemente inferior (aproximadamente 1.76$\times$ menor que Adam en promedio). Esto indica que la normalización espectral de Muon dirige las actualizaciones hacia direcciones con menor curvatura local.

3. El Rol del Desequilibrio de Datos
La ventaja de NDS de Muon se amplifica por el desequilibrio de los datos.

• Observación 3: Utilizando datos de Zipf-PCFG, los autores muestran que a medida que el conjunto de datos se vuelve más desequilibrado (exponente de Zipf $s$ más alto), la NDS para ambos optimizadores aumenta, pero la brecha entre Adam y Muon se ensancha significativamente (aumentando 1.8$\times$ de $s=0$ a $s=1$). Muon es más robusto ante el aumento de la nitidez inducida por los datos de cola pesada.

4. El Rol de la Estructura del Modelo
La ventaja de la curvatura cambia a lo largo del entrenamiento.

• Observación 4: Si bien tanto los componentes de NDS dentro de la capa como entre capas son menores para Muon, el componente dentro de la capa se vuelve cada vez más dominante en la ventaja de Muon durante las etapas media y tardía del entrenamiento. La contribución entre capas de Muon cae rápidamente, mientras que Adam mantiene un equilibrio estable entre la nitidez dentro y entre capas.

5. Justificación Teórica
El artículo demuestra que en problemas cuadráticos con curvatura heterogénea y alineación de gradiente hacia modos de alta curvatura:

• Muon obtiene un NDS promedio más pequeño que GD porque su normalización espectral equilibra la energía de la actualización a través de los grupos de curvatura, mientras que GD concentra la energía en los modos de alta curvatura.
• Cuando la heterogeneidad de la curvatura es suficientemente fuerte, Muon logra una pérdida cuadrática local menor que GD después del mismo número de pasos.

Significado y Reivindicaciones

El artículo afirma dar el "primer paso" hacia la desmitificación de la superioridad de Muon desde una perspectiva de la curvatura local. La conclusión central es que la eficiencia de Muon proviene de su capacidad para realizar actualizaciones espectralmente normalizadas que inherentemente poseen una menor nitidez direccional, minimizando así el costo de segundo orden de los pasos de optimización.

Los autores posicionan explícitamente este trabajo como complementario a las explicaciones existentes "centradas en la memoria" (por ejemplo, Wang et al., 2025b), ofreciendo un mecanismo geomético distinto: en la misma iteración, Muon paga una penalización de nitidez direccional menor que Adam. El trabajo es modesto en alcance, centrándose en LLMs causales e identificando un mecanismo local en lugar de reclamar garantías de convergencia global o aplicabilidad a todas las clases de modelos (como los modelos de difusión), las cuales se dejan para trabajos futuros.