Resumen Técnico: Muown - Control de Normalización por Fila para la Optimización Muon

Declaración del Problema

Muon ha surgido como un optimizador competitivo para el pre-entrenamiento de modelos de lenguaje grande (LLM), superando a menudo a AdamW al aprovechar la estructura matricial de las capas lineales de los Transformers. Sin embargo, observaciones recientes indican que el comportamiento de Muon a gran escala es sensible a la descomposición de pesos (weight decay) y propenso a inestabilidades numéricas, como desbordamientos en bfloat16 y logits de atención que explotan.

Un problema crítico identificado es la deriva ascendente consistente de la norma espectral ($\|W_t\|_{S_\infty}$) de las matrices de pesos durante el entrenamiento cuando se utiliza Muon sin descomposición de pesos desacoplada. Esta deriva viola la condición espectral requerida para un aprendizaje de características estable ($\|W\|_{S_\infty} = \Theta(\sqrt{m/n})$). Aunque la descomposición de pesos desacoplada puede mitigar esto reduciendo uniformemente la matriz de pesos, es una intervención no dirigida que escala todo el espectro independientemente de qué direcciones específicas impulsan el crecimiento, lo que podría ralentizar la convergencia en las fases tempranas del entrenamiento. El artículo busca identificar el impulsor empírico específico de esta deriva de la norma espectral y proponer un mecanismo de control más dirigido.

Metodología: Muown

Los autores introducen Muown, un reemplazo directo para Muon que aborda la deriva de la norma espectral mediante una estrategia de reparametrización en lugar de restricciones agregadas.

1. Descomposición de la Norma Espectral

La idea central es una descomposición de la norma espectral en dos factores (Proposición 1):
$$\|W_t\|_{S_\infty}^2 = \|g_t\|_\infty^2 \cdot \lambda_{\max}(P_t C_t P_t)$$
Donde:

• $g_t = \|W_t\|_{row}$ es el vector de magnitudes de fila (normas euclidianas de las filas).
• $D_t$ es la matriz normalizada por filas ($W_t = \text{Diag}(g_t)D_t$).
• $C_t = D_t D_t^\top$ es la matriz de Gram de las filas de norma unitaria.
• $P_t$ es una matriz diagonal de magnitudes de fila normalizadas.
• $\lambda_{\max}(P_t C_t P_t)$ representa el factor de coherencia de filas.

El análisis empírico en modelos de 500M revela que el factor de magnitud de fila ($\|g_t\|_\infty$) es el impulsor principal de la deriva de la norma espectral, moviéndose al unísono con $\|W_t\|_{S_\infty}$. En contraste, el factor de coherencia de filas se mantiene bien comportado y acotado durante todo el entrenamiento.

2. Reparametrización y Optimización

Basándose en este diagnóstico, Muown reparametriza la matriz de pesos $W$ como:
$$W(g, R) = \text{Diag}\left(\frac{g}{\|R\|_{row}}\right) R$$
Aquí, $g \in \mathbb{R}^m$ representa el vector explícito de magnitudes de fila, y $R \in \mathbb{R}^{m \times n}$ representa el componente direccional.

• Componente Direccional ($R$): Muon se aplica sin cambios a $R$, tratándolo bajo la geometría de la norma espectral (descenso espectral más pronunciado).
• Componente de Magnitud ($g$): Las magnitudes de fila $g$ se promueven a variables explícitas del optimizador. Se actualizan bajo la geometría $\ell_\infty$, que es la norma dual natural para la norma del operador de matrices diagonales.
• Regla de Actualización: La actualización de $g$ utiliza un optimizador Adam (actuando como SignSGD suavizado bajo la geometría $\ell_\infty$) para manejar el ruido estocástico, mientras que la actualización de $R$ sigue el procedimiento estándar de Muon con ortogonalización Newton-Schulz.

Este enfoque se implementa implícitamente dentro del estado del optimizador, dejando el paso hacia adelante (forward pass) de la arquitectura del modelo sin cambios.

Contribuciones Clave

1. Diagnóstico de la Deriva Espectral: El artículo descompone la norma espectral para identificar la magnitud máxima de fila ($\|g_t\|_\infty$) como el impulsor empírico del crecimiento de la norma espectral bajo Muon, distinguiéndola del factor de coherencia de filas estable.
2. Algoritmo Muown: Un nuevo optimizador que trata las magnitudes de fila como variables entrenables explícitas actualizadas bajo la geometría $\ell_\infty$, preservando al mismo tiempo las actualizaciones direccionales de Muon. Esto evita la reducción indiscriminada de la descomposición de pesos y los priores de ortogonalidad de las reparametrizaciones espectrales.
3. Garantías de Convergencia: Los autores demuestran que Muown alcanza tasas de convergencia óptimas no convexas en regímenes deterministas y estocásticos. El análisis se realiza en una norma dual mixta ($\ell_1$ para magnitudes, Schatten-1 para direcciones) alineada con las geometrías de actualización. El límite teórico sugiere una reducción en el coeficiente de ruido estocástico efectivo en comparación con Muon.
4. Validación Empírica: Experimentos extensos en FineWeb-Edu a través de tamaños de modelo desde 124M hasta 2.7B parámetros demuestran mejoras consistentes.

Resultados Experimentales

• Mejoras en la Perplejidad: Muown logra consistentemente una perplejidad menor que Muon, SOAP, AdamW y Lion en todas las escalas de modelo probadas (124M a 2.7B) y tasas de aprendizaje.
• Robustez a Hiperparámetros:
  • Tasa de Aprendizaje: Muown amplía la meseta de tasas de aprendizaje casi óptimas a través de diferentes anchos de modelo, reduciendo la sensibilidad al ajuste de la tasa de aprendizaje por escala.
  • Descomposición de Pesos: Muown alivia la necesidad de ajustar la descomposición de pesos; funciona bien sin descomposición de pesos ($\lambda=0$), mientras que Muon requiere un ajuste cuidadoso (óptimo $\lambda=0.1$) para evitar la divergencia o un rendimiento subóptimo.
• Estabilidad: Muown previene la deriva sistemática de la norma espectral observada en Muon. En regímenes de alta tasa de aprendizaje (por ejemplo, $4 \times 10^{-3}$), Muon a menudo diverge, mientras que Muown permanece estable y cerca de lo óptimo.
• Propiedades del Gradiente: El análisis muestra que Muown mantiene un rango efectivo normalizado más alto de los gradientes de minibatch y exhibe un coeficiente de ruido estocástico más bajo durante todo el entrenamiento en comparación con Muon.
• Sobrecarga Computacional: El método introduce una sobrecarga de tiempo por paso negligible. Cuando el optimizador se fragmenta (sharded) entre rangos (entorno distribuido), los tiempos de paso para Muown y Muon se vuelven indistinguibles. La sobrecarga de memoria es mínima (lineal en el número de filas), añadiendo solo unos pocos megabytes por capa.

Significado y Afirmaciones

El artículo afirma que la deriva de la norma espectral observada en Muon no es una propiedad irreducible del optimizador, sino una consecuencia del crecimiento no controlado de la magnitud de fila. Al aislar este impulsor y promoverlo a una variable explícita del optimizador, Muown logra:

• Mejora en la Convergencia: Mejor perplejidad y estabilidad sin la necesidad de una descomposición de pesos agresiva.
• Solidez Teórica: Tasas de convergencia óptimas en un marco consciente de la geometría que tiene en cuenta explícitamente la naturaleza mixta de las actualizaciones.
• Eficiencia Práctica: Un reemplazo directo que iguala el tiempo de ejecución de Muon en entornos distribuidos mientras ofrece una superior robustez ante las elecciones de hiperparámetros.

Los autores posicionan a Muown como una solución que cierra la brecha entre los requisitos teóricos de estabilidad del aprendizaje de características y las necesidades prácticas del entrenamiento de LLM a gran escala, ofreciendo una alternativa más dirigida a la descomposición de pesos desacoplada.