Resumen Técnico: Muon aprende características más robustas y transferibles que Adam

Declaración del Problema

Si bien el optimizador Muon (Momentum Ortogonalizado por Newton-Schulz) ha demostrado recientemente una eficiencia de vanguardia en el preentrenamiento de Modelos de Lenguaje de Gran Escala (LLMs) y clasificadores de visión —entrenando a menudo casi el doble de rápido que Adam y SGD—, la naturaleza específica de sus ventajas en el aprendizaje de características sigue sin estar clara. La literatura existente ha explicado mayoritariamente el éxito de Muon a través de lentes de optimización teórica, tales como el descenso de gradiente más pronunciado bajo normas espectrales, la dinámica de memoria asociativa o el manejo de datos de cola pesada. Sin embargo, ha faltado una investigación sistemática sobre si Muon aprende mejores características que Adam y SGD, y cómo se manifiestan estas características en las representaciones del modelo. Este artículo aborda la pregunta: ¿Aprende Muon características más robustas y transferibles que Adam y SGD, y qué propiedades caracterizan esta ventaja?

Metodología

Los autores emplean un enfoque dual que combina una extensa evaluación empírica a través de dominios de visión y lenguaje con un análisis teórico en un problema de clasificación estilizado.

Evaluación Empírica

El estudio evalúa la calidad de las características mediante dos métricas primarias: robustez y transferibilidad.

• Robustez: Los modelos preentrenados (ResNet-18, ViT-S, GPT-2, GPT-2 Medium) son probados con entradas corruptas. Los modelos de visión se evalúan en ImageNet-C (15 tipos de corrupción), mientras que los modelos de lenguaje se evalúan en un benchmark recién construido, FineWeb10B-C, que aplica ataques de errores tipográficos a nivel de carácter (inserciones, eliminaciones, sustituciones, intercambios).
• Transferibilidad: Las características preentrenadas se adaptan a tareas posteriores. Para visión, se entrenan clasificadores lineales sobre backbones congelados en cuatro conjuntos de datos (EuroSAT, Food-101, Flowers-102, Stanford Cars). Para lenguaje, los modelos completos se ajustan (fine-tuning) en conjuntos de datos de seguimiento de instrucciones (Alpaca, Dolly, WizardLM).
• Análisis de Estados Ocultos: Para comprender el mecanismo detrás de estas mejoras de rendimiento, los autores analizan los estados ocultos por capa utilizando:
  • Márgenes de Logit: Sondas lineales entrenadas (marco de trabajo tuned-lens) decodifican los estados ocultos para calcular el margen entre la clase/token correcto y el distractor más fuerte.
  • Diversidad Espectral: Se computan el rango efectivo y la fracción de energía Top-k de las matrices de estados ocultos para medir la diversidad y la distribución de la información a través de direcciones ortogonales.
• Protocolo Experimental: Todas las comparaciones siguen un protocolo de presupuesto igualado donde la arquitectura, los datos y la duración del entrenamiento se fijan, mientras que los hiperparámetros (tasa de aprendizaje, scheduler) se ajustan independientemente para cada optimizador para aislar las diferencias de representación del esfuerzo de ingeniería.

Análisis Teórico

Los autores proporcionan soporte teórico utilizando un modelo de clasificación lineal de una sola capa con características de múltiples componentes. El modelo de datos asume que las características consisten en un componente compartido a nivel de bloque (clase gruesa) y un componente específico de la clase (clase de grano fino). El análisis compara la dinámica de entrenamiento de el Descenso de Gradiente (GD), Adam y Muon bajo tamaños de paso infinitesimales y momentum cero, centrándose en el estado donde cada optimizador alcanza por primera vez un umbral de pérdida específico.

Contribuciones Clave y Resultados

1. Evidencia Empírica de una Calidad de Característica Superior

El artículo establece que Muon aprende consistentemente características que son más robustas y transferibles que las aprendidas por Adam y SGD:

• Robustez: Los modelos preentrenados con Muon alcanzan la precisión media más alta en ImageNet-C y la menor perplejidad en FineWeb10B-C a través de todos los tipos de corrupción y arquitecturas. Notablemente, Muon supera a Adam en transformadores de visión y modelos de lenguaje, y supera a SGD en modelos de lenguaje.
• Transferibilidad: Las características preentrenadas con Muon se transfieren más eficazmente a tareas posteriores. En visión, Muon iguala o supera a Adam y SGD en todos los benchmarks posteriores. En lenguaje, los modelos preentrenados con Muon logran una menor perplejidad tras el ajuste de instrucciones en comparación con los modelos preentrenados con Adam.

2. Mecanismos de Estados Ocultos

Las ventajas empíricas están directamente vinculadas a propiedades específicas de las representaciones aprendidas:

• Mayores Márgenes de Logit: El sondeo por capas revela que Muon produce márgenes de logit significativamente mayores en la mayoría de las capas en comparación con Adam y SGD. Esto indica una separación más amplia entre las predicciones correctas y los distractores, lo que correlaciona directamente con la robustez observada.
• Mayor Rango Efectivo: El análisis espectral muestra que Muon alcanza un rango efectivo más alto y una fracción de energía Top-10 más baja a través de las capas. Esto sugiere que las representaciones de Muon distribuyen la información de manera más uniforme a través de direcciones ortogonales, creando un subespacio de características más rico que facilita una mejor adaptación a las tareas posteriores.

3. Garantías Teóricas

En el entorno de características de múltiples componentes estilizado, los autores demuestran que la normalización espectral de Muon elimina los desequilibrios de valores singulares en la matriz de gradiente.

• Ventaja de Margen: Muon alcanza un margen de clasificación estrictamente mayor que Adam y GD en cualquier pérdida de entrenamiento igualada (fuera de una ventana de pérdida terminal exponencialmente pequeña). Esto se debe a que la normalización espectral de Muon equilibra la contribución de las direcciones de características ortogonales, mientras que la normalización por coordenadas de Adam y el flujo de gradiente bruto de GD conducen a una dependencia excesiva de componentes específicos.
• Ventaja de Rango Efectivo: Muon alcanza demostrablemente un rango efectivo de características estrictamente superior al de Adam y GD en pérdida igualada, lo que respalda el hallazgo empírico de representaciones más diversas.

Significado y Reivindicaciones

El artículo afirma "desmitificar" las ventajas de Muon al desplazar el enfoque de la eficiencia de optimización hacia la calidad de la característica. Los autores argumentan que el éxito de Muon no es simplemente un resultado de una convergencia más rápida, sino que proviene de su capacidad para aprender representaciones que son inherentemente más robustas a las perturbaciones de entrada y más adaptables a nuevas tareas.

La importancia de este trabajo radica en:

1. Conectar la Optimización con la Representación: Establece un vínculo concreto entre la regla de actualización consciente de la matriz de Muon y las propiedades geométricas de las características aprendidas (márgenes de logit y rango efectivo).
2. Validación Teórica: Proporciona la primera prueba teórica de que Muon induce calidades de características sistemáticamente diferentes y superiores (mayores márgenes, mayor rango efectivo) en comparación con Adam y GD bajo condiciones de pérdida igualada.
3. Implicaciones Prácticas: Los hallazgos sug로 que para aplicaciones que requieren robustez ante cambios en la distribución o una fuerte transferibilidad a tareas posteriores, Muon es una opción superior a los optimizadores adaptativos estándar como Adam, incluso cuando los presupuestos de entrenamiento son iguales.

Los autores señalan modestamente que su análisis actual se centra en LLMs y clasificadores de visión, dejando el desempeño de Muon en otras familias de modelos (por ejemplo, modelos de difusión) como una dirección para trabajos futuros.