Muon+: Towards Better Muon via One Additional Normalization Step

Este trabajo presenta Muon+, una mejora del optimizador Muon que añade un paso de normalización tras la ortogonalización y que demuestra consistentemente un mejor rendimiento en el entrenamiento de modelos de lenguaje de diversos tamaños y arquitecturas.

Lo esencial

¡Claro que sí! Imagina que entrenar una Inteligencia Artificial gigante (como las que escriben textos o crean imágenes) es como entrenar a un equipo olímpico de atletas para que corran la maratón perfecta.

Aquí te explico el papel "MUON+" usando una analogía sencilla:

1. El Problema: El Equipo se Desorganiza

Para entrenar a estos modelos, usamos un "entrenador" (un algoritmo matemático llamado Muon). Este entrenador tiene una tarea difícil: corregir los pasos de miles de atletas al mismo tiempo.

El problema es que, a veces, los atletas se mueven de forma caótica. Algunos corren muy rápido, otros muy lento, y sus movimientos se "apilan" unos sobre otros, perdiendo dirección. En matemáticas, a esto le llaman "colapso de rango". El entrenador original (Muon) intentaba arreglar esto haciendo que los atletas se organizaran en filas y columnas perfectas (una técnica llamada ortogonalización). Funcionaba bien, pero no era perfecto.

2. La Solución: El "Ajuste Final" (MUON+)

Los autores de este paper (Ruijie Zhang y su equipo) se dieron cuenta de algo genial. Decidieron: "Oye, después de que el entrenador organiza a los atletas en filas perfectas, ¿por qué no les damos un pequeño empujón extra para asegurarnos de que todos tengan la misma fuerza?".

Ese "pequeño empujón" es el MUON+. Es un paso extra muy simple: normalización.

La analogía de la orquesta:
Imagina que el entrenador (Muon) ha logrado que todos los músicos toquen la misma nota al mismo tiempo (ortogonalización). ¡Genial! Pero, ¿qué pasa si el violinista toca a volumen 10 y el tamborilero a volumen 1? La música suena mal.

• Muon: Alinea a los músicos para que toquen juntos.
• MUON+: Alinea a los músicos Y ajusta los volúmenes para que todos suenen equilibrados.

3. ¿Qué lograron?

El equipo probó esta idea con modelos de diferentes tamaños (desde pequeños como un pájaro hasta gigantes como un elefante) y en dos escenarios:

• Entrenamiento normal: Como correr una maratón estándar.
• Entrenamiento extremo (Overtraining): Como correr una maratón de ultra-distancia (muchísimos datos).

El resultado fue increíble:

• Más rápido: Los modelos aprendían mejor y más rápido.
• Más estable: El entrenamiento no se "descontrolaba" tanto.
• Mejor calidad: Al final, el modelo entendía el lenguaje humano de forma más precisa (menos errores, como un estudiante que saca mejores notas).

4. ¿Por qué es importante?

Antes, para entrenar a estos gigantes, se necesitaban superordenadores y mucho dinero. Muon ya ayudaba a ahorrar costos, pero MUON+ es como un "hack" simple que mejora el rendimiento sin necesidad de hardware nuevo.

Es como si descubrieras que, simplemente ordenando mejor tu mochila antes de salir a correr, puedes correr más rápido y sin cansarte tanto, sin necesidad de comprar zapatillas nuevas.

En resumen:

El paper dice: "El método Muon ya era bueno organizando el caos, pero si le añadimos un paso extra para equilibrar la fuerza de cada movimiento (normalización), el resultado es un modelo más inteligente, más rápido y más robusto, sin complicar demasiado las cosas".

¡Es una mejora elegante y poderosa que podría hacer que las futuras Inteligencias Artificiales sean más eficientes y accesibles para todos!

Resumen técnico

Aquí tienes un resumen técnico detallado del artículo "MUON+: Towards Better Muon via One Additional Normalization Step" en español.

Resumen Técnico: MUON+

1. El Problema

El entrenamiento de modelos fundacionales de lenguaje (LLMs) a gran escala enfrenta costos computacionales prohibitivos a medida que aumentan el tamaño de los parámetros y los conjuntos de datos. Aunque optimizadores como Adam y AdamW son dominantes, la comunidad busca alternativas más eficientes. Recientemente, el optimizador Muon ha demostrado un rendimiento prometedor al ortogonalizar la matriz de momento mediante iteraciones de Newton-Schulz, lo que ayuda a contrarrestar el colapso de rango de los gradientes.

Sin embargo, a pesar de su éxito, los autores identifican que el optimizador Muon original puede beneficiarse de una mejora estructural simple. El objetivo es mejorar la estabilidad de la optimización y la calidad final del modelo en regímenes de entrenamiento computacionalmente óptimos y en escenarios de "sobre-entrenamiento" (long-horizon), donde la relación tokens-parámetros (T2P) es muy alta.

2. Metodología: MUON+

Los autores proponen MUON+, una mejora directa y sencilla sobre el optimizador Muon. La innovación central consiste en añadir un paso de normalización adicional inmediatamente después del paso de ortogonalización.

• Mecanismo de Actualización:
  El optimizador Muon estándar actualiza los pesos $W_t$ basándose en una matriz ortogonal $O_t$ derivada del momento $M_t$:
  $$O_t = \text{Ortho}(M_t)$$
  $$W_t = W_{t-1} - \eta \cdot \sqrt{m/n} \cdot O_t$$

  En MUON+, se introduce un operador de normalización $\text{Norm}^{(d)}$ sobre la matriz ortogonalizada antes de la actualización:
  $$O_t = \text{Norm}^{(d)}(\text{Ortho}(M_t))$$
  $$W_t = W_{t-1} - \eta \cdot \sqrt{m/n} \cdot O_t$$

• Direcciones de Normalización:
  El estudio evalúa varias direcciones para la normalización $\ell_2$:

  • Columna (col): Normalizar cada columna de la matriz.
  • Fila (row): Normalizar cada fila de la matriz.
  • Combinadas: Aplicar secuencialmente normalización por columna y fila (o viceversa), denotadas como col_row y row_col.

• Implementación:
  El algoritmo mantiene la aproximación de la descomposición en valores singulares (SVD) mediante iteraciones de Newton-Schulz (usualmente 5 iteraciones) para calcular la parte ortogonal, pero añade una capa de normalización que controla la magnitud de las actualizaciones, actuando como un regularizador estructural.

3. Contribuciones Clave

1. Propuesta de MUON+: Introducir un paso de normalización post-ortogonalización que es simple de implementar pero altamente efectivo.
2. Evaluación Exhaustiva: Validación en una amplia gama de escalas de modelos (desde 60M hasta 1B de parámetros) y arquitecturas (GPT y LLaMA).
3. Análisis de Regímenes de Entrenamiento:
   • Óptimo Computacional: Entrenamiento con una relación T2P de $\approx 20$.
   • Sobre-entrenamiento (Overtraining): Extensión de la relación T2P a niveles industriales de $\approx 200$ (entrenando con 72 mil millones de tokens).
4. Estudios de Ablación: Identificación de que la normalización es el factor principal de mejora, superando a otras modificaciones complejas como la adaptación de segundo momento o formulaciones de variedades (manifolds).

4. Resultados Experimentales

Los experimentos se realizaron utilizando el conjunto de datos FineWeb en GPUs H100/A100 con precisión mixta (bfloat16).

• Rendimiento en Modelos GPT (130M - 774M):

  • MUON+ superó consistentemente a Muon en todas las escalas.
  • GPT-Small (124M): Reducción de la perplejidad de validación de 29.66 a 27.64 (-2.02).
  • GPT-Base (362M): Reducción de 21.70 a 19.98 (-1.72).
  • GPT-Large (774M): Reducción de 17.82 a 16.91 (-0.91).

• Rendimiento en Modelos LLaMA (60M - 1B):

  • MUON+ también superó a AdamW y Muon en todas las configuraciones.
  • LLaMA-1B: Logró una perplejidad de 10.31 frente a 10.68 de Muon y 14.38 de AdamW.
  • Las direcciones de normalización col_row y row_col mostraron el mejor rendimiento, siendo la normalización por fila (row) ligeramente superior a la de columna (col).

• Escalabilidad y Sobre-entrenamiento:

  • En experimentos de sobre-entrenamiento (72B tokens, T2P $\approx 200$), MUON+ mantuvo una ventaja consistente sobre Muon (ej. GPT-Base: 15.84 vs 16.97).
  • Las curvas de pérdida mostraron que la brecha de rendimiento se mantiene estable durante todo el entrenamiento, indicando que MUON+ no degrada en etapas avanzadas de optimización.

• Robustez ante la Tasa de Aprendizaje:

  • MUON+ demostró ser menos sensible a la selección de la tasa de aprendizaje, manteniendo un rendimiento estable en un rango más amplio de hiperparámetros en comparación con Muon.

• Independencia del Método de Ortogonalización:

  • La mejora se mantuvo consistente independientemente del método utilizado para aproximar la parte ortogonal (You, Jordan, o PolarExpress), lo que sugiere que la ganancia proviene de la normalización y no de la aproximación SVD específica.

5. Significado e Impacto

• Simplicidad y Eficacia: El trabajo demuestra que una modificación mínima (una sola línea de normalización) puede generar ganancias sustanciales en el entrenamiento de LLMs, desafiando la noción de que se necesitan algoritmos complejos para mejorar el rendimiento.
• Interpretación Teórica: Los autores sugieren que el motor principal de la mejora no es la adaptación de segundo momento (como en Adam) ni la complejidad de las variedades, sino la normalización estructural de las actualizaciones ortogonales. Esto estabiliza la magnitud de los gradientes y previene inestabilidades numéricas.
• Aplicabilidad Industrial: Al funcionar bien en relaciones T2P extremadamente altas ($\approx 200$), MUON+ es una candidata viable para el entrenamiento de modelos de próxima generación donde la eficiencia computacional y la estabilidad a largo plazo son críticas.
• Reproducibilidad: El código está disponible públicamente, facilitando la adopción inmediata por parte de la comunidad de investigación e industria.

En conclusión, MUON+ establece un nuevo estándar simple pero potente para la optimización de LLMs, demostrando que la normalización post-ortogonalización es un componente esencial para la estabilidad y el rendimiento en el entrenamiento de modelos a gran escala.