How Much Orthogonalization Does Muon Need?

Resumen Técnico: ¿Cuánta ortogonalización necesita Muon?

Planteamiento del problema
El optimizador Muon mejora el entrenamiento de redes neuronales reemplazando las actualizaciones de momentum mal condicionadas por actualizaciones aproximadamente semi-ortogonales. Mientras que el álgebra numérica lineal estándar busca una descomposición polar de alta precisión (mapeando todos los valores singulares a 1), el contexto de Muon sugiere que la precisión exacta puede no ser necesaria. La cuestión central abordada por este trabajo es: ¿Cuánta precisión de ortogonalización requiere realmente Muon para preservar la calidad del entrenamiento? Específicamente, el artículo investiga si el costo computacional de las iteraciones de polinomios de alto grado (utilizadas en métodos existentes como Muon-Jordan y Polar Express) está justificado, o si una aproximación más barata y de menor grado puede lograr resultados de entrenamiento comparables.

Metodología
Los autores proponen un esquema de Newton–Schulz cúbico relajado (denominado cubic5) diseñado específicamente para la banda de valores singulares de baja precisión (bfloat16) de Muon.

1. Objetivos Relajados: A diferencia de la descomposición polar clásica que apunta a la precisión de máquina, este método apunta a una banda relajada de $[0.7, 1.3]$ para los valores singulares. Reconoce que en bfloat16, los valores singulares por debajo de un límite inferior efectivo ($l_0 = 7 \times 10^{-3}$) no son prácticamente significativos.
2. Derivación de Coeficientes Adaptativos: Los autores derivan un polinomio cúbico adaptativo $f_t(x) = a_t x + b_t x^3$ en cada iteración. Los coeficientes se calculan basándose en el límite inferior de valor singular peor de caso actual ($l_t$) y un pico superior fijo ($u=1.3$). Esto asegura que el polinomio mapee el intervalo protegido $[l_t, r_t]$ hacia la banda objetivo mientras equilibra los extremos para maximizar el nuevo límite inferior.
3. Eficiencia Computacional:
   • Cubic5: Un paso cúbico único requiere dos multiplicaciones de matrices dominantes (formar la matriz de Gram y multiplicarla de vuelta). Cinco pasos requieren 10 multiplicaciones de matrices dominantes.
   • Línea base (Muon-Jordan/Polar Express): Un paso quíntico requiere tres multiplicaciones de matrices dominantes. Cinco pasos requieren 15 multiplicaciones de matrices dominantes.
   • El método utiliza el "truco de la transpuesta" para matrices altas para minimizar las operaciones en la matriz de Gram más pequeña.

Contribuciones Clave

• Derivación de un Esquema Cúbico Relajado: El artículo deriva un esquema de Newton–Schulz cúbico de cinco pasos con coeficientes adaptados al límite inferior efectivo de bfloat16, logrando la banda espectral objetivo en cinco iteraciones.
• Caracterización de la Frontera Costo-Calidad: El trabajo establece que la calidad del entrenamiento no está gobernada monótonamente por la precisión de la descomposición polar. Demuestra que una aproximación cúbica de menor costo puede alcanzar la paridad de rendimiento con métodos quínticos de mayor costo e incluso con la SVD exacta en regímenes específicos.
• Validación Empírica a Escala: Los autores validan el método a través de:
  • Diagnósticos Sintéticos: Confirmando que la transferencia del polinomio escalar se mantiene en la aritmética de matrices de baja precisión.
  • Ablaciones de NanoGPT: Comparando cubic5, Polar Express truncado, Muon-Jordan truncado y SVD explícita en FP32 sobre GPT-2 Small.
  • Entrenamiento a Gran Escala: Entrenamiento de modelos híbridos MoE/Mamba (Nemotron-3-Nano) que van desde 1B hasta 4B de parámetros utilizando Megatron-LM.

Resultados

• Benchmarks Sintéticos: El método cubic5 mapea con éxito los valores singulares en el rango protegido $[7 \times 10^{-3}, 1]$ hacia la banda objetivo sin explosión numérica. Los microbenchmarks en GPUs NVIDIA H200 muestran que cubic5 reduce el costo de la subrutina de ortogonalización en relación con Polar Express, con aceleraciones que oscilan entre 1.16x y 1.52x dependiendo de la forma de la matriz (rectangular vs. cuadrada).
• Ablaciones de GPT-2 Small:
  • Un cubic5 de cinco pasos logra una pérdida de validación final de 3.2811.
  • Un Polar Express de cinco pasos logra 3.2798.
  • Un Muon-Jordan de cinco pasos logra 3.2822.
  • La SVD explícita en FP32 logra 3.2809.
  • Observación: La brecha entre las variantes más fuertes es menor a 0.003 en la pérdida de validación. Notablemente, la SVD exacta no supera a los métodos de aproximación polinómica, y añadir un quinto paso a Muon-Jordan degradó ligeramente el rendimiento en comparación con cuatro pasos.
• Escalamiento de Híbrido MoE/Mamba (1B–4B):
  • En modelos de 1B y 2B, cubic5 superó ligeramente a Muon-Jordan (menor pérdida de validación por $\approx 10^{-3}$).
  • En el modelo de 4B, cubic5 fue ligeramente superior ($\approx 10^{-3}$) pero las trayectorias de entrenamiento fueron casi indistinguibles tras la fase inicial.
  • La diferencia relativa en la pérdida de entrenamiento decayó rápidamente, fluctuando alrededor de cero durante el resto del presupuesto de entrenamiento.

Significancia y Reclamaciones
El artículo reclama modestamente que cubic5 es una variante práctica y de bajo costo para la ortogonalización de Muon, en lugar de una regla de actualización universalmente superior.

• Optimizador vs. Precisión de Descomposición: Los resultados apoyan la hipótesis de que Muon se beneficia de un efecto de "remodelación espectral" en lugar de una estricta precisión de descomposición polar. El mapa de valores singulares exacto que optimiza el entrenamiento depende de la fase de entrenamiento y la tasa de aprendizaje, no solo de la convergencia matemática a 1.
• Utilidad Práctica: cubic5 ofrece un punto fundamentado en la frontera costo-calidad. Reduce el conteo de multiplicaciones de matrices dominantes en aproximadamente un tercio (10 vs. 15) manteniendo la paridad de calidad de entrenamiento empírica con la actualización quíntica estándar de Muon-Jordan en los regímenes probados.
• Limitaciones: Los autores establecen explícitamente que estos resultados establecen la paridad de calidad de entrenamiento, no aceleraciones de tiempo de ejecución (throughput) de extremo a extremo. La aceleración real en tiempo de reloj depende de los kernels de hardware, el batching y la fracción de tiempo que la etapa de ortogonalización ocupa dentro del bucle completo del optimizador. El artículo no afirma que cubic5 sea un mejor algoritmo de descomposición polar que Polar Express en un sentido matemático, sino que es un transformador espectral más eficiente para el contexto específico del entrenamiento de Muon.