How Much Orthogonalization Does Muon Need?

Technische Zusammenfassung: Wie viel Orthogonalisierung benötigt Muon?

Problemstellung
Der Muon-Optimizer verbessert das Training neuronaler Netze, indem er schlecht konditionierte Momentum-Updates durch annähernd semi-orthogonale Updates ersetzt. Während die klassische numerische Lineare Algebra auf eine hochpräzise Polardekomposition abzielt (Abbildung aller Singulärwerte auf 1), legt der Kontext von Muon nahe, dass exakte Genauigkeit nicht unbedingt erforderlich ist. Die zentrale Frage, die in dieser Arbeit behandelt wird, lautet: Wie viel Genauigkeit bei der Orthogonalisierung benötigt Muon tatsächlich, um die Trainingsqualität zu erhalten? Speziell untersucht die Arbeit, ob die Rechenkosten hochgradiger Polynom-Iterationen (verwendet in bestehenden Methoden wie Muon-Jordan und Polar Express) gerechtfertigt sind oder ob eine günstigere, niedergradige Approximation vergleichbare Trainingsergebnisse erzielen kann.

Methodik
Die Autoren schlagen einen relaxierten kubischen Newton–Schulz-Zeitplan (bezeichnet als cubic5) vor, der speziell für das Niedrigpräzisions-Band (bfloat16) von Muon entwickelt wurde.

1. Relaxierte Zielwerte: Im Gegensatz zur klassischen Polardekomposition, die auf Maschinengenauigkeit abzielt, zielt diese Methode auf ein relaxiertes Band von $[0,7, 1,3]$ für die Singulärwerte ab. Sie berücksichtigt, dass in bfloat16 Singulärwerte unterhalb einer effektiven Untergrenze ($l_0 = 7 \times 10^{-3}$) praktisch nicht mehr sinnvoll sind.
2. Ableitung adaptiver Koeffizienten: Die Autoren leiten in jeder Iteration ein adaptives kubisches Polynom $f_t(x) = a_t x + b_t x^3$ ab. Die Koeffizienten werden basierend auf der aktuellen Worst-Case-Untergrenze der Singulärwerte ($l_t$) und einem festen oberen Peak ($u=1,3$) berechnet. Dies stellt sicher, dass das Polynom das geschützte Intervall $[l_t, r_t]$ in das Zielband abbildet und dabei die Endpunkte ausbalanciert, um die neue Untergrenze zu maximieren.
3. Recheneffizienz:
   • Cubic5: Ein einziger kubischer Schritt erfordert zwei dominante Matrizenmultiplikationen (Bildung der Gram-Matrix und deren Rückmultiplikation). Fünf Schritte erfordern 10 dominante Matrizenmultiplikationen.
   • Baseline (Muon-Jordan/Polar Express): Ein quintischer Schritt erfordert drei dominante Matrizenmultiplikationen. Fünf Schritte erfordern 15 dominante Matrizenmultiplikationen.
   • Die Methode nutzt den „Transpose-Trick“ für hohe Matrizen, um die Operationen auf der kleineren Gram-Matrix zu minimieren.

Wesentliche Beiträge

• Ableitung eines relaxierten kubischen Zeitplans: Die Arbeit leitet einen fünfstufigen kubischen Newton–Schulz-Zeitplan mit Koeffizienten ab, die an die effektive bfloat16-Untergrenze angepasst sind, und erreicht das Ziel-Spektralband in fünf Iterationen.
• Charakterisierung der Kosten-Qualitäts-Frontier: Die Arbeit stellt fest, dass die Trainingsqualität nicht monoton von der Genauigkeit der Polardekomposition abhängt. Sie zeigt, dass eine kostengünstigere kubische Approximation die Performance-Parität mit höherwertigen quintischen Methoden und sogar mit exakter SVD in spezifischen Regimen erreichen kann.
• Empirische Validierung über verschiedene Skalen hinweg: Die Autoren validieren die Methode durch:
  • Synthetische Diagnostik: Bestätigung, dass der skalare Polynomtransfer in der Niedrigpräzisions-Arithmetik der Matrizen Bestand hat.
  • NanoGPT Ablationen: Vergleich von cubic5, trunkiertem Polar Express, trunkiertem Muon-Jordan und expliziter FP32 SVD auf GPT-2 Small.
  • Großskaliges Training: Training von hybriden MoE/Mamba-Modellen (Nemotron-3-Nano) mit einer Größe von 1B bis 4B Parametern unter Verwendung von Megatron-LM.

Ergebnisse

• Synthetische Benchmarks: Die cubic5-Methode bildet Singulärwerte im geschützten Bereich $[7 \times 10^{-3}, 1]$ erfolgreich in das Zielband ab, ohne numerische Instabilitäten (Blow-up) zu erzeugen. Microbenchmarks auf NVIDIA H200 GPUs zeigen, dass cubic5 die Kosten der Orthogonalisierungs-Subroutine im Vergleich zu Polar Express reduziert, mit Beschleunigungen zwischen 1,16x und 1,52x, abhängig von der Matrizenform (rechteckig vs. quadratisch).
• GPT-2 Small Ablationen:
  • Fünf-Schritt cubic5 erreicht einen finalen Validierungsverlust von 3,2811.
  • Fünf-Schritt Polar Express erreicht 3,2798.
  • Fünf-Schritt Muon-Jordan erreicht 3,2822.
  • Explizite FP32 SVD erreicht 3,2809.
  • Beobachtung: Die Differenz zwischen den stärksten Varianten liegt bei weniger als 0,003 im Validierungsverlust. Bemerkenswerterweise übertrifft die exakte SVD die approximativen Polynom-Methoden nicht, und das Hinzufügen eines fünften Schritts zu Muon-Jordan verschlechterte die Performance im Vergleich zu vier Schritten leicht.
• Hybride MoE/Mamba Skalierung (1B–4B):
  • Bei 1B und 2B Modellen übertraf cubic5 Muon-Jordan leicht (niedrigerer Validierungsverlust um $\approx 10^{-3}$).
  • Beim 4B Modell war cubic5 leicht höher ($\approx 10^{-3}$), aber die Trainingsverläufe waren nach der Initialphase nahezu ununterscheidbar.
  • Der relative Unterschied im Trainingsverlust nahm schnell ab und fluktuierte für den Rest des Trainingsbudgets um Null.

Bedeutung und Ansprüche
Die Arbeit behauptet bescheiden, dass cubic5 eine praktische, kostengünstige Variante für die Muon-Orthogonalisierung ist, und nicht eine universell überlegene Update-Regel.

• Optimizer vs. Dekompositionsgenauigkeit: Die Ergebnisse stützen die Hypothese, dass Muon von einem „spektralen Umgestaltungs-Effekt“ profitiert und nicht von einer strikten Genauigkeit der Polardekomposition. Die exakte Singulärwert-Abbildung, die das Training optimiert, hängt von der Trainingsphase und der Lernrate ab, nicht nur von der mathematischen Konvergenz gegen 1.
• Praktischer Nutzen: cubic5 bietet einen fundierten Punkt auf der Kosten-Qualitäts-Frontier. Es reduziert die Anzahl der dominanten Matrizenmultiplikationen um etwa ein Drittel (10 statt 15) und behält dabei die empirische Trainingsqualität der Standard-Muon-Jordan-Quintic-Updates in den getesteten Regimen bei.
• Limitierungen: Die Autoren stellen explizit klar, dass diese Ergebnisse eine Qualitätsparität beim Training etablieren, nicht eine Steigerung des End-to-End-Durchsatzes. Die tatsächliche Wandzeit-Beschleunigung hängt von Hardware-Kerneln, Batching und dem Anteil ab, den der Orthogonalisierungsschritt innerhalb des vollständigen Optimizer-Loops einnimmt. Die Arbeit behauptet nicht, dass cubic5 ein mathematisch besserer Polardekompositions-Algorithmus als Polar Express ist, sondern dass es eine effizientere spektrale Transformation für den spezifischen Kontext des Muon-Trainings darstellt.