Resumen Técnico: Límite de Convergencia y Tamaño de Lote Crítico del Optimizador Muon

Planteamiento del Problema

Aunque el optimizador Muon ha demostrado un fuerte rendimiento empírico como un potencial sucesor de los optimizadores estándar como AdamW, sus fundamentos teóricos siguen estando subdesarrollados. La literatura existente carece de pruebas de convergencia formales para Muon, particularmente en entornos prácticos que involucran el impulso de Nesterov y el decaimiento de pesos (weight decay). Este artículo tiene como objetivo cerrar estas brechas proporcionando un análisis de convergencia riguroso y derivando un límite inferior para el tamaño de lote crítico de Muon.

Metodología

Los autores llevan a cabo un análisis teórico de Muon a través de cuatro configuraciones prácticas:

1. Sin impulso de Nesterov y sin decaimiento de pesos.
2. Con impulso de Nesterov y sin decaimiento de pesos.
3. Sin impulso de Nesterov y con decaimiento de pesos.
4. Con impulso de Nesterov y con decaimiento de pesos.

Marco Teórico:

• Análisis de Convergencia: Los autores derivan límites superiores para la norma nuclear promedio esperada del gradiente, $\frac{1}{T}\sum_{t=0}^{T-1} \mathbb{E}[\|\nabla f(W_t)\|_*]$. Utilizan supuestos estándar de suavidad $L$ y varianza de gradiente acotada. Una innovación clave es el análisis de la dirección de búsqueda de Muon, la cual se obtiene mediante la ortogonalización de una matriz de impulso $C_t$. Los autores contabilizan el rango de la matriz de error de impulso ($r_1$) y la dirección de búsqueda ($r_2$), en lugar de la dimensión completa del parámetro.
• Aproximación de Newton-Schulz: Dado que Muon aproxima prácticamente la ortogonalización mediante la iteración de Newton-Schulz en lugar de una SVD exacta, los autores analizan el impacto de este error de aproximación ($\delta(k)$) en la convergencia, demostrando que introduce solo un término negligible.
• Derivación del Tamaño de Lote Crítico (CBS): Los autores modelan el número de pasos de entrenamiento $T$ necesarios para alcanzar una precisión objetivo $\epsilon$ como una función del tamaño de lote $b$. Al minimizar la complejidad SFO (Stochastic First-Order), derivan un límite inferior para el tamaño de lote crítico $b^*$.

Validación Experimental:
Los autores validan las predicciones teóricas mediante una estrategia de dos niveles:

1. Experimentos de Control de Full-Muon: Una tarea sintética de Profesor-Estudiante donde Muon actualiza todos los parámetros, coincidiendo estrictamente con los supuestos teóricos.
2. Experimentos Híbridos de Escala Práctica: Entrenamiento de ResNet-18/VGG-16 en conjuntos de datos CIFAR y Llama3.1 (320M) en el corpus C4. En estos entornos, Muon se aplica a los parámetros con forma de matriz mientras que AdamW se encarga del resto, reflejando un despliegue estándar.

Contribuciones Clave

1. Pruebas de Convergencia para Cuatro Variantes

El artículo establece garantías de convergencia para las cuatro variantes de Muon. Los límites superiores derivados para la norma nuclear promedio del gradiente son del orden de:
$$O\left(\frac{1}{\eta T} + \sqrt{\frac{(1-\beta)r_1}{b}} + \bar{r}\eta\right)$$
(con ligeras variaciones para los términos de Nesterov y decaimiento de pesos).

• Significancia: A diferencia de los optimizadores basados en vectores donde la varianza escala con el número total de parámetros, el límite de Muon escala con el rango $r_1$ de la matriz de error de impulso. La evidencia empírica sugiere que $r_1 \ll n$ en regímenes sobreparametrizados, lo que implica un límite teórico más ajustado.
• Norma Nuclear: El análisis utiliza la norma nuclear, que es un criterio más estricto que la norma de Frobenius utilizada en muchos análisis basados en vectores, sugiriendo que Muon respeta la estructura intrínseca de matriz de los parámetros de la red neuronal.

2. El Papel del Decaimiento de Pesos y la Estabilidad

Los autores demuestran que la incorporación del decaimiento de pesos asegura la acotación casi segura tanto de las normas de los parámetros ($\|W_t\|_F$) como de las normas de los gradientes ($\|\nabla f(W_t)\|_F$) sin depender de la restricción de gradientes acotados.

• Condición de Estabilidad: Derivan una condición necesaria para la estabilidad: la tasa de aprendizaje $\eta$ y el coeficiente de decaimiento de pesos $\lambda$ deben satisfacer $\eta \leq 1/\lambda$. Los resultados experimentales confirman que el entrenamiento se vuelve inestable cuando $\eta$ supera este umbral.
• Regularización Implícita: Aunque el decaimiento de pesos no mejora el orden de la tasa de convergencia asintótica, proporciona regularización implícita al acotar las normas, lo que explica su rendimiento empírico superior.

3. Límite Inferior del Tamaño de Lote Crítico (CBS)

El artículo deriva un límite inferior para el tamaño de lote crítico $b^*$, mostrando su dependencia de los hiperparámetros:
$$b^*_{\text{Muon}} > \frac{9(1-\beta)(1+\sqrt{2\beta})^2 r_1 \sigma^2}{(1-\lambda)^2 \epsilon^2}$$

• Dependencia de Hiperparámetros: La fórmula revela que:
  • Impulso ($\beta$): A medida que $\beta \to 1$, el término $(1-\beta)$ desaparece, sugiriendo que el tamaño de lote crítico disminuye.
  • Decaimiento de Pesos ($\lambda$): El factor $1/(1-\lambda)^2$ indica que aumentar el decaimiento de pesos aumenta el tamaño de lote crítico.
  • Rango ($r_1$): El límite escala con el rango del error de impulso, vinculando el CBS con la estructura de bajo rango de los gradientes.
• Escalamiento Cualitativo: Los autores enfatizan que la fórmula proporciona una ley de escalamiento cualitativa en lugar de una predicción absoluta, ya que depende de constantes específicas del problema ($\sigma^2, r_1, \epsilon$). Sin embargo, la ratio entre los diferentes CBS de las variantes de Muon depende solo de $\beta$ y $\lambda$.

Resultados

Hallazgos Teóricos:

• La tasa de convergencia coincide con los métodos estándar de tipo SGD no convexo ($O(T^{-1/4})$ bajo esquemas específicos de hiperparámetros) pero con un término de varianza más ajustado debido a la dependencia del rango.
• El error de aproximación de Newton-Schulz tiene un efecto limitado en la convergencia, decayendo exponencialmente con el número de iteraciones.

Hallazgos Experimentales:

• Estabilidad: Experimentos en ResNet-18/CIFAR-10 validan la condición de estabilidad $\eta \leq 1/\lambda$. Las normas de los gradientes y la pérdida de entrenamiento se minimizan cerca de este límite.
• Velocidad de Convergencia: Muon con impulso de Nesterov y decaimiento de pesos logra consistentemente la disminución más rápida de la pérdida y la norma del gradiente en comparación con AdamW, Momentum SGD y otras variantes de Muon.
• Validación del Tamaño de Lote Crítico (CBS):
  • Efectos de los Hiperparámetros: Los experimentos confirman que añadir impulso de Nesterov desplaza el CBS hacia valores más grandes, y aumentar el decaimiento de pesos incrementa aún más el CBS. Por ejemplo, aumentar $\lambda$ de 0 a 0.5 resultó en un desplazamiento del CBS de $\approx 1024$ a $\geq 4096$, coincidiendo con el escalamiento predicho de $1/(1-\lambda)^2$ (ratio de 4.0).
  • Escalamiento del Rango: Variar el ancho del modelo mostró que el rango efectivo $r_1$ escala con el ancho, y el producto $r_1 \cdot \sigma^2$ permanece relativamente constante, explicando por qué los valores empíricos de CBS parecían planos a través de los anchos a pesar de la dependencia teórica del rango.
  • Eficiencia: Muon generalmente requiere menos pasos y una menor complejidad SFO que AdamW, particularmente en tamaños de lote pequeños a moderados.

Significancia y Reivindicaciones

El artículo afirma proporcionar el primer análisis de convergencia riguroso de Muon que incorpora simultáneamente el impulso de Nesterov y el decaimiento de pesos, abordando un vacío en la comprensión teórica de este optimizador emergente.

• Perspectiva Teórica: Demuestra que el decaimiento de pesos no es meramente una herramienta de regularización, sino un mecanismo que asegura la acotación y la estabilidad, eliminando la necesidad de supuestos de gradientes acotados.
• Guía Práctica: El límite inferior del CBS derivado ofrece una guía accionable para los profesionales. Clarifica cómo el ajuste de $\beta$ y $\lambda$ afecta cualitativamente al tamaño de lote óptimo, sugiriendo que las configuraciones de alto impulso pueden entrenarse eficientemente con lotes más pequeños, mientras que un mayor decaimiento de pesos requiere lotes más grandes para mantener la eficiencia.
• Validación: El trabajo cierra la brecha entre la teoría y la práctica validando estas predicciones a través de cargas de trabajo de visión y modelado de lenguaje, confirmando que las leyes de escalamiento teórico se mantienen incluso en entornos de optimización híbridos donde Muon se aplica solo a parámetros con forma de matriz.

Los autores señalan modestamente que su análisis de matriz única no modela explícitamente la heterogeneidad por capa, pero argumentan que la alta correlación de las escalas de ruido de gradiente entre capas respalda la generalizabilidad de sus hallazgos. También reconocen que, aunque Muon supera a Shampoo en sus experimentos, derivar límites análogos para Shampoo sigue siendo un problema abierto.