How Much Orthogonalization Does Muon Need?

Résumé technique : Quel degré d'orthogonalisation le Muon nécessite-t-il ?

Énoncé du problème
L'optimiseur Muon améliore l'entraînement des réseaux de neurones en remplaçant les mises à jour de momentum mal conditionnées par des mises à jour approximativement semi-orthogonales. Alors que l'algèbre linéaire numérique standard vise une décomposition polaire de haute précision (mappant toutes les valeurs singulières à 1), le contexte de Muon suggère qu'une précision exacte peut être superflue. La question centrale abordée par ce travail est la suivante : Quel degré de précision d'orthogonalisation le Muon nécessite-t-il réellement pour préserver la qualité de l'entraînement ? Plus précisément, l'étude examine si le coût computationnel des itérations polynomiales de degré élevé (utilisées dans les méthodes existantes comme Muon-Jordan et Polar Express) est justifié, ou si une approximation plus économique, de degré inférieur, peut obtenir des résultats d'entraînement comparables.

Méthodologie
Les auteurs proposent un schéma de Newton-Schulz cubique relaxé (nommé cubic5) conçu spécifiquement pour la bande de valeurs singulières en basse précision (bfloat16) de Muon.

1. Cibles relaxées : Contrairement à la décomposition polaire classique qui cible la précision machine, cette méthode cible une bande relaxée de $[0,7, 1,3]$ pour les valeurs singulières. Elle reconnaît qu'en bfloat16, les valeurs singulières inférieures à une borne inférieure effective ($l_0 = 7 \times 10^{-3}$) ne sont pas pratiquement significatives.
2. Dérivation de coefficients adaptatifs : Les auteurs dérivent un polynôme cubique adaptatif $f_t(x) = a_t x + b_t x^3$ à chaque itération. Les coefficients sont calculés sur la base de la borne inférieure la plus faible actuelle ($l_t$) et d'un sommet supérieur fixe ($u=1,3$). Cela garantit que le polynôme projette l'intervalle protégé $[l_t, r_t]$ dans la bande cible tout en équilibrant les extrémités pour maximiser la nouvelle borne inférieure.
3. Efficacité computationnelle :
   • Cubic5 : Une étape cubique unique nécessite deux multiplications matricielles dominantes (formation de la matrice de Gram et multiplication inverse). Cinq étapes nécessitent 10 multiplications matricielles dominantes.
   • Référence (Muon-Jordan/Polar Express) : Une étape quintique nécessite trois multiplications matricielles dominantes. Cinq étapes nécessitent 15 multiplications matricielles dominantes.
   • La méthode utilise l'astuce de la transposée ("transpose trick") pour les matrices hautes afin de minimiser les opérations sur la matrice de Gram plus petite.

Principales contributions

• Dérivation d'un schéma cubique relaxé : Le papier dérive un schéma de Newton-Schulz cubique en cinq étapes dont les coefficients sont adaptés à la borne inférieure effective du bfloat16, atteignant la bande spectrale cible en cinq itérations.
• Caractérisation de la frontière Coût-Qualité : Ce travail établit que la qualité de l'entraînement n'est pas gouvernée de manière monotone par la précision de la décomposition polaire. Il démontre qu'une approximation cubique moins coûteuse peut atteindre la parité de performance avec les méthodes quintiques plus coûteuses et même avec la SVD exacte dans des régimes spécifiques.
• Validation empirique à travers différentes échelles : Les auteurs valident la méthode via :
  • Diagnostics synthétiques : Confirmation que le transfert polynomial scalaire tient dans l'arithmétique matricielle de basse précision.
  • Ablations NanoGPT : Comparaison de cubic5, Polar Express tronqué, Muon-Jordan tronqué et SVD explicite en FP32 sur GPT-2 Small.
  • Entraînement à grande échelle : Entraînement de modèles hybrides MoE/Mamba (Nemotron-3-Nano) allant de 1B à 4B de paramètres avec Megatron-LM.

Résultats

• Benchmarks synthétiques : La méthode cubic5 parvient à mapper les valeurs singulières dans la plage protégée $[7 \times 10^{-3}, 1]$ dans la bande cible sans explosion numérique. Des microbenchmarks sur des GPU NVIDIA H200 montrent que cubic5 réduit le coût de la sous-routine d'orthogonalisation par rapport à Polar Express, avec des accélérations allant de 1,16x à 1,52x selon la forme de la matrice (rectangulaire vs carrée).
• Ablations GPT-2 Small :
  • Un cubic5 en cinq étapes atteint une perte de validation finale de 3,2811.
  • Un Polar Express en cinq étapes atteint 3,2798.
  • Un Muon-Jordan en cinq étapes atteint 3,2822.
  • Une SVD explicite en FP32 atteint 3,2809.
  • Observation : L'écart entre les variantes les plus fortes est inférieur à 0,003 en perte de validation. Notamment, la SVD exacte ne surpasse pas les méthodes polynomiales approximatives, et l'ajout d'une cinquième étape à Muon-Jordan dégrade légèrement la performance par rapport à quatre étapes.
• Mise à l'échelle Hybride MoE/Mamba (1B–4B) :
  • Sur les modèles 1B et 2B, cubic5 surpasse légèrement Muon-Jordan (perte de validation plus faible d'environ $10^{-3}$).
  • Sur le modèle 4B, cubic5 est légèrement plus élevé ($\approx 10^{-3}$), mais les trajectoires d'entraînement sont presque indiscernables après la phase initiale.
  • La différence relative de perte d'entraînement a décroît rapidement, fluctuant autour de zéro pour le reste du budget d'entraînement.

Signification et affirmations
Le papier affirme modestement que cubic5 est une variante pratique et à bas coût pour l'orthogonalisation de Muon, plutôt qu'une règle de mise à jour universellement supérieure.

• Optimiseur vs Précision de la décomposition : Les résultats soutiennent l'hypothèse selon laquelle Muon bénéficie d'un effet de "reformatage spectral" plutôt que d'une stricte précision de décomposition polaire. La carte exacte des valeurs singulières qui optimise l'entraînement dépend de la phase d'entraînement et du taux d'apprentissage, et non seulement de la convergence mathématique vers 1.
• Utilité pratique : cubic5 offre un point de contrôle principled sur la frontière coût-qualité. Il réduit le nombre de multiplications matricielles dominantes d'environ un tiers (10 contre 15) tout en maintenant une parité de qualité d'entraînement empirique avec la mise à jour quintique standard de Muon-Jordan dans les régimes testés.
• Limites : Les auteurs précisent explicitement que ces résultats établissent une parité de qualité d'entraînement, et non des accélérations de débit (throughput) de bout en bout. L'accélération réelle en temps réel dépend des kernels matériels, du batching et de la fraction de temps que occupe l'étape d'orthogonalisation dans la boucle complète de l'optimiseur. Le papier ne prétend pas que cubic5 est un meilleur algorithme de décomposition polaire que Polar Express d'un point de vue mathématique, mais plutôt une transformation spectrale plus efficace pour le contexte spécifique de l'entraînement de Muon.