The Spectral Dynamics and Noise Geometry of Muon

Resumen Técnico: La Dinámica Espectral y la Geometría del Ruido de Muon

Planteamiento del Problema
El artículo investiga el sesgo de regularización implícita inducido por el optimizador Muon, el cual actualiza los pesos utilizando el factor polar de la matriz de gradiente en lugar del gradiente mismo. Dado que un gradiente $G = U\Sigma V^\top$, Muon actualiza mediante $UV^\top$, descartando efectivamente las magnitudes relativas de los valores singulares mientras preserva las direcciones singulares activas. Aunque Muon ha demostrado éxito empírico en el preentrenamiento de transformadores, su mecanismo de selección teórico sigue siendo incierto. Se abordan dos hipótesis contrapuestas:

1. ¿Es Muon simplemente una forma de normalización del gradiente matricial?
2. ¿Implica la geometría de la norma espectral un sesgo hacia soluciones de bajo rango (minimización de la norma nuclear)?

Los autores argumentan que ambas son incompletas. En su lugar, proponen que Muon induce un sesgo de espectro de actualización plano (flat-update-spectrum bias), asignando una amplitud igual a todas las direcciones singulares activas. El artículo busca caracterizar este sesgo teórica y empíricamente, específicamente en el contexto de la regresión subdeterminada donde el optimizador debe seleccionar un interpolante finito de un manifold afín de soluciones de pérdida cero.

Metodología
El estudio emplea una combinación de análisis de sistemas dinámicos de tiempo continuo, cálculo variacional y experimentos empíricos controlados.

1. Marco Teórico:

   • Flujo Polar Proyectado: Los autores analizan el flujo de tiempo continuo $\dot{W} = -\eta P_\perp P(W)$, donde $P(W)$ es el factor polar de $W$ y $P_\perp$ es el proyector ortogonal sobre el espacio tangente del manifold de interpolación $\mathcal{M} = \{W : XW = Y\}$.
   • Suposición de Marco Compartido (SF): Se introduce una condición estructural donde las matrices de gradiente y de peso comparten marcos singulares ($U_G = U_W, V_G = V_W$) para simplificar el análisis. Se demuestran teoremas sin esta suposición para el flujo auto-polar, mientras que la suposición SF cierra la brecha hacia el flujo real de Muon impulsado por el gradiente.
   • Funcionales: El análisis utiliza un funcional espectral por pares $R_{pw}(W) = -\sum_{i \neq j} \log(\sigma_i + \sigma_j)$ y su versión invariante de escala $\tilde{R}_{pw}$ para rastrear la forma espectral.
   • Geometría del Ruido: Se realiza una derivación del tamaño de lote crítico utilizando un cruce de señal-ruido linealizado de un solo paso, analizando la derivada de Fréchet del mapa polar.

2. Validación Empírica:

   • Detección de Matrices (Matrix Sensing): Experimentos controlados en tareas de regresión lineal subdeterminada comparan Muon contra minimizadores de norma nuclear (resueltos vía CVXPY) y descenso de gradiente.
   • Preentrenamiento de Transformers: Preentrenamiento a pequeña escala de NanoGPT (124M) y un control ViT emparejado sobre CIFAR-10 para observar la dinámica espectral (rango estable, entropía espectral) en regímenes de aprendizaje profundo.
   • Ablaciones: Pruebas sobre el acoplamiento de la decaimiento de pesos (acoplado vs. desacoplado) y el desalineamiento del marco compartido para verificar la robustez.

Contribuciones Clave

1. Teorema de Sesgo Espectral de un Solo Paso:
   Bajo suposiciones de alineación explícitas, los autores demuestran que la actualización polar es la única elección de maximización de entropía de un solo paso entre actualizaciones acotadas que siguen las direcciones singulares del gradiente pero ignoran el espectro actual del peso. Establecen un criterio de aplanamiento espectral: el espectro normalizado se aplana si y solo si la geometría de la medición preserva preferencialmente las direcciones correspondientes a los mayores valores singulares (cuantificado por un criterio de signo que involucra $\alpha_i q_i$).

2. Dinámica Exacta de los Valores Singulares:
   El artículo deriva la dinámica exacta para el flujo polar proyectado, mostrando que los valores singulares decaen a tasas $\dot{\sigma}_i = -\eta \alpha_i$, donde $\alpha_i$ depende de la alineación entre los vectores singulares izquierdos del peso y el espacio nulo de la matriz de medición. Esto conduce a una caracterización variacional donde, a una norma de Frobenius fija, el único minimizador del funcional espectral es un espectro plano (todos los valores singulares no nulos son iguales).

3. Separación de Explicaciones Existentes:

   • Normalización: Los experimentos muestran que el descenso de gradiente con norma emparejada no reproduce el comportamiento de Muon.
   • Norma Nuclear: Los resultados teóricos y empíricos demuestran que los puntos límite de Muon no coinciden con los minimizadores de norma nuclear. Mientras que la minimización de la norma nuclear favorece la concentración espectral (bajo rango), Muon favorece el aplanamiento espectral (alto rango estable).

4. Utilidad Dependiente del Régimen:
   El artículo establece que el sesgo de espectro plano de Muon no es universalmente superior. En el preentrenamiento pequeño de NanoGPT, Muon preserva el rango estable y mejora la pérdida de validación en relación con AdamW. Sin embargo, en un control ViT pequeño emparejado, el ranking se invierte, sugiriendo que el sesgo es beneficioso solo cuando muchas direcciones espectrales deben permanecer activas.

Resultos

• Teóricos: El Teorema 1 prueba que el flujo auto-polar proyectado impulsa el sistema hacia un espectro plano si la geometría de la medición satisface un criterio de signo específico. El Teorema 2 acota la desviación cuando la suposición de Marco Compartido solo se cumple aproximadamente. El Teorema 4 proporciona una fórmula para el tamaño de lote crítico $B_{crit}$ en el SGD de mapa polar, vinculándolo a la sensibilidad al ruido del mapa polar $S(\mu) = \sum_{i \neq j} (\sigma_i + \sigma_j)^{-2}$.
• Empíricos (Detección de Matrices): Muon converge a soluciones con normas nucleares significativamente más altas (1.29x–2.02x) y mayor entropía espectral en comparación con el interpolante de mínima norma nuclear, falsificando la hipótesis de que Muon minimiza la norma nuclear.
• Empíricos (Transformers): En NanoGPT, Muon mantiene una mayor entropía espectral por capa y un mayor rango estable en comparación con AdamW, que tiende a colapsar el rango en tareas de lenguaje. Por el contrario, en el control ViT, AdamW supera a Muon, apoyando la conclusión de "dependencia del régimen".
• Sensibilidad al Ruido: La fórmula del tamaño de lote crítico derivada muestra que $B_{crit}$ está gobernado por la sensibilidad del mapa polar al ruido, la cual se infla por valores singulares pequeños o casi degenerados. Los espectros planos a un radio de Frobenius fijo minimizan esta sensibilidad.

Significancia y Reivindicaciones
El artículo afirma proporcionar la primera caracterización estructural del sesgo implícito de Muon como un mecanismo de espectro de actualización plano en lugar de un efecto de bajo rango o de simple normalización.

• Alcance Modesto: Los autores declaran explícitamente que los resultados son dependientes del régimen. Muon no es un optimizador universal; su utilidad surge en regímenes espectrales donde mantener la actividad en muchas direcciones es beneficioso.
• Límites Teóricos: Los teoremas son locales y estructurales, dependiendo de suposiciones como brechas de valores singulares simples y geometrías de medición específicas. El artículo no pretende explicar el éxito de Muon en todos los escenarios de entrenamiento de LLM a gran escala, señalando que la reducción de las capas de los transformadores al modelo de regresión subdeterminada estudiado sigue siendo una "observación condicional" que requiere suposiciones no verificadas sobre los rangos de activación.
• Perspectiva Práctica: El trabajo ofrece una herramienta de diagnóstico para profesionales: si una aplicación requiere preservar el rango estable y una actividad espectral amplia, el sesgo de espectro plano de Muon es ventajoso; si el problema es inherentemente de bajo rango o bien condicionado, el sesgo puede ser innecesario o perjudicial.

El artículo concluye identificando problemas abiertos, incluyendo la caracterización de las geometrías de medición específicas que garantizan el criterio de aplanamiento globalmente y la extensión del análisis de sensibilidad al ruido a casos rectangulares y de rango deficiente.