Delving into Muon and Beyond: Deep Analysis and Extensions

Resumen Técnico: Profundizando en Muon y Más Allá: Análisis Detallado y Extensiones

Planteamiento del Problema
El optimizador Muon ha ganado atención recientemente por su rendimiento empírico en el entrenamiento de modelos de lenguaje de gran escala, utilizando actualizaciones ortogonalizadas a nivel de matriz. Sin embargo, los mecanismos subyacentes de Muon y su relación con optimizadores adaptativos establecidos como Adam permanecen insuficientemente comprendidos. La literatura existente a menudo confunde los efectos de Muon con técnicas auxiliares (por ejemplo, QK-Norm, QK-Clip) o carece de una comparación sistemática contra líneas de base estándar. Además, no está claro si la ortogonalización espectral de Muon representa una estrategia de optimización universalmente superior o una forma específica de estabilización, y si puede ser mejorada o generalizada.

Metodología
Los autores proponen un marco espectral unificado para analizar Muon y sus variantes. Definen a Muon como el extremo $p=0$ de una familia de transformaciones espectrales de la forma $U\Sigma^p V^\top$, donde $U\Sigma V^\top$ es la Descomposición en Valores Singulares (SVD) de la matriz de actualización. El estudio investiga exponentes intermedios $p \in \{1, 1/2, 1/4, 0\}$:

• $p=1$: Identidad (actualizaciones estándar).
• $p=1/2, 1/4$: Compresión espectral parcial.
• $p=0$: Ortogonalización completa (Muon).

Estas transformaciones se aplican a dos tipos de entrada distintos:

1. Momento del primer instante ($O_{mom}$): Equivalente a las entradas de momentum SGD.
2. Actualizaciones normalizadas por RMS ($O_{rms}$): Equivalente a las entradas de Adam, donde los gradientes se normalizan por el segundo momento.

Para permitir el cálculo eficiente de potencias fraccionarias ($p=1/2, 1/4$) sin una SVD explícita, los autores desarrollan una iteración de Newton-Schulz acoplada. Este algoritmo calcula simultáneamente la raíz cuadrada de la matriz y su inversa utilizando únicamente multiplicaciones de matrices, asegurando la estabilidad numérica y la escalabilidad.

La evaluación experimental se lleva a cabo en un modelo GPT-2 de 124 millones de parámetros entrenado en OpenWebText. El estudio emplea una configuración altamente controlada:

• Tasas de Aprendizaje Desacopladas: Se utilizan tasas de aprendizaje separadas para los parámetros de matriz y de vector para aislar el efecto del optimizador.
• Sin Técnicas Auxiliares: El decaimiento de peso (weight decay), QK-Norm y QK-Clip están desactivados para asegurar comparaciones justas.
• Ajuste Sistemático: Se realiza una búsqueda de tasa de aprendizaje en dos etapas (barrido logarítmico grueso seguido de un refinamiento de grano fino) para las ocho variantes del optimizador (4 exponentes espectrales $\times$ 2 tipos de entrada).

Contribuciones Clave

1. Marco Espectral Unificado: El artículo introduce un marco formal $U\Sigma^p V^\top$ que posiciona a Muon como un caso específico ($p=0$) y propone variantes intermedias ($p=1/2, 1/4$) para interpolar entre las actualizaciones estándar y la ortogonalización completa.
2. Algoritmo Eficiente: Se desarrolla una iteración de Newton-Schulz acoplada para calcular eficientemente las actualizaciones espectrales fraccionarias ($p=1/2, 1/4$), evitando el costo computacional de una SVD explícita.
3. Comparación Empírica Rigurosa: El estudio proporciona una comparación controlada entre Muon, Adam y sus variantes espectrales, aislando los efectos espectrales de factores de confusión como el decaimiento de peso o las técnicas de normalización auxiliar.

Resultos Clave
Los experimentos arrojan tres hallazgos principales:

1. Estabilización de las Actualizaciones del Primer Instante: En entradas de primer instante (mSGD), Muon ($p=0$) proporciona una estabilización significativa, permitiendo tasas de aprendizaje mucho más altas y mayor robustez en comparación con mSGD estándar. Sin embargo, la compresión parcial ($p=1/4$) a menudo logra un mejor rendimiento final una vez establecida la estabilidad.
2. Ganancias Limitadas en Entradas Normalizadas por RMS: Cuando se aplica a entradas normalizadas por RMS (estilo Adam), la compresión espectral ofrece rendimientos decrecientes. La ortogonalización completa (AdamZ/Muon aplicado a $O_{rms}$) no supera a Adam estándar y puede ser inferior a la compresión parcial. Esto sugiere que la normalización RMS elemento a elemento ya mitiga gran parte de la anisotropía que la compresión espectral busca abordar.
3. Superioridad de los Métodos de Segundo Momento: A través de todos los exponentes espectrales, los optimizadores que utilizan la normalización del segundo momento (basada en RMS) superan consistentemente a sus contrapartes de primer momento (solo momentum) tanto en estabilidad como en la pérdida final.

Significancia y Reivindicaciones
El artículo concluye que Muon debe entenderse mejor como una forma efectiva de normalización espectral en lugar de un método de optimización universalmente superior.

• Estabilización vs. Rendimiento: La contribución principal de Muon es estabilizar las actualizaciones del primer instante al descartar las magnitudes de los valores singulares, lo que evita que las direcciones con grandes valores singulares dominen. Sin embargo, este "aplanamiento" es un mecanismo rudimentario; elimina el filtrado basado en la magnitud, lo que potencialmente amplifica las direcciones de gradiente ruidosas (valores singulares pequeños) y resta pasos a los modos dominantes informativos.
• Utilidad Contextual: Si bien la compresión espectral ofrece beneficios reales de estabilización, especialmente cuando la entrada de actualización no está ya normalizada, la ortogonalización completa ($p=0$) no es un reemplazo para los optimizadores modernos de segundo momento normalizados por RMS como Adam. Los autores argumentan que las ventajas reportadas de Muon en trabajos previos pueden atribuirse a la combinación específica de técnicas utilizadas (por ejemplo, QK-Norm) más que a la ortogonalización en sí misma.

El trabajo enfatiza que, aunque las transformaciones espectrales son herramientas poderosas para controlar la anisotropía, no superan intrínsecamente los beneficios del escalado de segundo momento adaptativo en regímenes de entrenamiento a gran escala.