Denoise First, Orthogonalize Later: Understanding Momentum in Muon via Spectral Filtering

Resumen Técnico: "Primero eliminar el ruido, después ortogonalizar: Entendiendo el momento en Muon mediante el filtrado espectral"

Planteamiento del problema
El optimizador Muon ha demostrado recientemente un sólido rendimiento empírico en el entrenamiento de modelos de lenguaje extensos (LLMs) al ortogonalizar aproximadamente una matriz de momento mediante una iteración de Newton–Schulz. Sin embargo, el papel teórico del momento dentro de Muon sigue sin estar claro. Los análisis existentes o bien eliminan el momento para estudiar las actualizaciones espectrales de forma aislada, o bien lo mantienen sin explicar sus beneficios empíricos. Además, la literatura teórica sobre Muon suele sugerir que un momento pesado puede degradar las tasas de convergencia o no ofrecer ningún beneficio, lo que contradice las observaciones prácticas donde el momento es esencial. Asimismo, el orden específico de las operaciones —aplicar el momento antes de la ortogonalización (Pre-polar) frente a ortogonalizar antes del momento (Post-polar)— carece de una explicación rigurosa, a pesar de que el método Pre-polar supera significativamente al Post-polar en la práctica.

Metodología
Los autores proponen un modelo de señal más perturbación de gradiente estructurado para analizar Muon. Asumen que el gradiente $G_t$ en la iteración $t$ se descompone en un componente de señal coherente y de bajo rango $G^{\text{sig}}_t$ y una perturbación de varianza acotada y media cero (BVMZOS) $\Xi_t$.

La contribución teórica central es el análisis del búfer de momento $M_t = \beta M_{t-1} + (1-\beta)G_t$ como un filtro espectral. Los autores investigan tres variantes de canalización para la actualización de pesos $\Delta W_t$:

1. Pre-polar (Muon estándar): $\Delta W_t \propto O(M_t)$, donde el momento se aplica primero, seguido del factor polar $O(\cdot)$.
2. Post-polar (Orthogonal-SGDM): $\Delta W_t \propto \tilde{M}_t$, donde el factor polar se aplica primero al gradiente bruto, y luego se acumula el momento.
3. Solo-polar: $\Delta W_t \propto O(G_t)$, un modelo base sin momento.

El análisis se basa en el tamaño de ventana efectivo $T = 1/(1-\beta)$ y utiliza herramientas de la teoría de matrices aleatorias, incluyendo el teorema $\sin \Theta$ de Wedin para la perturbación de subespacios singulares, para cuantificar cómo el momento afecta a los valores singulares y los subespacios de la matriz de gradiente.

Contribuciones clave

1. El momento como filtro espectral (Teorema 1 y Corolario 1):
   Los autores demuestran que, bajo su modelo de gradiente, el momento actúa como un filtro de paso bajo que suprime la perturbación (ruido) mientras preserva la señal coherente. Específicamente, a medida que el tamaño de la ventana efectiva $T$ aumenta, los valores singulares de la cola (dominados por el ruido) decaen a una tasa de $O(T^{-1/4})$, mientras que los valores singulares principales permanecen alejados de cero. Esto crea un brecha espectral ampliada entre los componentes de señal y ruido. En consecuencia, los subespacios singulares de los primeros $r$ elementos del búfer de momento $M_t$ se alinean de manera más fiable con los subespacios de la señal real que los subespacios del gradiente bruto.

2. No conmutatividad y orden de la canalización (Teorema 2):
   El artículo establece que el momento y la ortogonalización no conmutan.

   • Dominancia Pre-polar: Cuando el momento se aplica primero, este elimina el ruido del flujo de gradientes, permitiendo que el factor polar subsiguiente opere sobre una matriz con una estructura de subespacio singular estabilizada. Los autores demuestran que el método Pre-polar logra una mayor alineación de la señal esperada que tanto el método Post-polar como los modelos base de Solo-polar.
   • Fallo de las alternativas: Para el método Post-polar, el hecho de ortogonalizar el gradiente ruidoso primero "borra" la brecha de amplitud entre las direcciones de señal y ruido. El factor polar reemplaza todos los valores singulares por unos, haciendo imposible distinguir la señal del ruido en los pasos de momento posteriores. Teóricamente, las canalizaciones Post-polar y Solo-polar fallan en recuperar eficazmente la dirección de la señal, especialmente en regímenes de baja relación señal-ruido (SNR).

3. Separación cuantitativa en baja SNR (Teorema 3):
   En un modelo gaussiano de pico de rango 1 con baja SNR, los autores muestran que la alineación de la señal del canal de Solo-polar desaparece a medida que la dimensión de la matriz $m \to \infty$. En contraste, la canalización Pre-polar continúa recuperando la señal objetivo siempre que la ventana de momento $T$ sea suficientemente grande.

Resultados
Los hallazgos teóricos se validan mediante experimentos extensos:

• Visualización de la brecha espectral: Utilizando sondas estacionarias en NanoGPT, los autores demuestran que el aumento del coeficiente de momento $\beta$ ensancha la brecha entre los valores singulares principales y de la cola del búfer de momento, lo cual es consistente con el Teorema 1.
• Fiabilidad del subespacio: Las mediciones empíricas de los errores de alineación de subespacios muestran que la desviación entre los subespacios singulares del búfer de momento y los subespacios de la señal real disminuye a medida que $T$ aumenta, validando el Corolario 1.
• Comparación de canalizaciones: A través de modelos sintéticos, CIFAR-10 y el entrenamiento de NanoGPT/LLaMA 350M, la canalización Pre-polar supera consistentemente a la Post-polar y a la de Solo-polar en términos de alineación de señal y pérdida de validación. Notablemente, el rendimiento de la Post-polar permanece plano independientemente de $\beta$, confirmando la predicción teórica de que ortogonalizar antes de promediar no puede recuperar las direcciones de la señal perdidas por el ruido.
• Entrenamiento de extremo a extremo: En tareas de preentrenamiento de LLM (NanoGPT y LLaMA 350M), la canalización Pre-polar (Muon estándar) logra una pérdida de validación menor en comparación con las otras variantes, lo que respalda la afirmación de que el momento es un componente esencial para la recuperación de la señal, y no simplemente un truco de aceleración.

Significancia
Este artículo proporciona la primera justificación teórica de por qué Muon requiere momento y por qué el orden de las operaciones (eliminar el ruido antes de ortogonalizar) es crítico. Replantea el momento en los optimizadores basados en matrices no solo como un mecanismo de aceleración, sino como un filtro espectral que estabiliza los subespacios singulares frente al ruido. El principio rector "primero eliminar el ruido, luego ortogonalizar" ofrece una regla de diseño para futuros optimizadores conscientes de la matriz que empleen pasos espectrales no lineales, sugiriendo que separar la señal del ruido debe preceder a cualquier normalización o ortogonalización espectral. Los autores señalan que, si bien su teoría explica los beneficios del momento en Muon, también resalta las limitaciones en los límites teóricos actuales que no logran predecir estos beneficios, señalando una brecha entre la teoría de optimización estándar y el comportamiento práctico de los optimizadores espectrales.