Delving into Muon and Beyond: Deep Analysis and Extensions

Technische Zusammenfassung: Einblick in Muon und darüber hinaus: Tiefe Analyse und Erweiterungen

Problemstellung
Der Muon-Optimizer hat kürzlich durch seine empirische Leistungsfähigkeit beim Training großer Sprachmodelle unter Verwendung von orthogonalisierten Updates auf Matrizebene Aufmerksamkeit erregt. Die zugrunde liegenden Mechanismen von Muon und dessen Beziehung zu etablierten adaptiven Optimierern wie Adam sind jedoch noch unzureichend verstanden. Die bestehende Literatur vermischt die Effekte von Muon oft mit Hilfstechniken (z. B. QK-Norm, QK-Clip) oder lässt einen systematischen Vergleich gegen Standard-Baselines vermissen. Zudem bleibt unklar, ob Muons spektrale Orthogonalisierung eine universell überlegene Optimierungsstrategie oder eine spezifische Form der Stabilisierung darstellt und ob sie verbessert oder generalisiert werden kann.

Methodik
Die Autoren schlagen ein einheitliches spektrales Framework vor, um Muon und seine Varianten zu analysieren. Sie definieren Muon als den $p=0$-Endpunkt einer Familie von spektralen Transformationen der Form $U\Sigma^p V^\top$, wobei $U\Sigma V^\top$ die Singulärwertzerlegung (SVD) der Update-Matrix ist. Die Studie untersucht Zwischen-Exponenten $p \in \{1, 1/2, 1/4, 0\}$:

• $p=1$: Identität (Standard-Updates).
• $p=1/2, 1/4$: Partielle spektrale Kompression.
• $p=0$: Volle Orthogonalisierung (Muon).

Diese Transformationen werden auf zwei unterschiedliche Input-Typen angewendet:

1. First-Moment Momentum ($O_{mom}$): Äquivalent zu Momentum-SGD-Inputs.
2. RMS-normalisierte Updates ($O_{rms}$): Äquivalent zu Adam-Inputs, bei denen Gradienten durch das zweite Moment normalisiert werden.

Um eine effiziente Berechnung fraktionaler Potenzen ($p=1/2, 1/4$) ohne explizite SVD zu ermöglichen, entwickeln die Autoren eine gekoppelte Newton-Schulz-Iteration. Dieser Algorithmus berechnet gleichzeitig die Quadratwurzel der Matrix und deren Inversen unter Verwendung reiner Matrixmultiplikationen, was numerische Stabilität und Skalierbarkeit gewährleistet.

Die experimentelle Evaluierung erfolgt auf einem GPT-2-Modell mit 124 Mio. Parametern, das auf OpenWebText trainiert wurde. Die Studie nutzt ein streng kontrolliertes Setup:

• Entkoppelte Lernraten: Separate Lernraten werden für Matrix- und Vektorparameter verwendet, um den Effekt des Optimierers zu isolieren.
• Keine Hilfstechniken: Weight Decay, QK-Norm und QK-Clip sind deaktiviert, um einen fairen Vergleich zu gewährleisten.
• Systematisches Tuning: Eine zweistufige Lernratensuche (grobe logarithmische Durchmusterung gefolgt von einer feingranularen Verfeinerung) wird für alle acht Optimizer-Varianten (4 spektrale Exponenten $\times$ 2 Input-Typen) durchgeführt.

Wesentliche Beiträge

1. Einheitliches spektrales Framework: Das Paper führt ein formales Framework $U\Sigma^p V^\top$ ein, das Muon als einen spezifischen Fall ($p=0$) positioniert und Zwischenvarianten ($p=1/2, 1/4$) vorschlägt, um zwischen Standard-Updates und voller Orthogonalisierung zu interpolieren.
2. Effizienter Algorithmus: Eine gekoppelte Newton-Schulz-Iteration wird entwickelt, um fraktionale spektrale Updates ($p=1/2, 1/4$) effizient zu berechnen, wodurch der Rechenaufwand einer expliziten SVD vermieden wird.
3. Rigoroser empirischer Vergleich: Die Studie liefert einen kontrollierten Vergleich zwischen Muon, Adam und deren spektralen Varianten, wobei spektrale Effekte von konfundierenden Faktoren wie Weight Decay oder Hilfs-Normalisierung isoliert werden.

Zentrale Ergebnisse
Die Experimente liefern drei primäre Erkenntnisse:

1. Stabilisierung von First-Moment-Updates: Bei First-Moment-Inputs (mSGD) bietet Muon ($p=0$) eine signifikante Stabilisierung, was wesentlich höhere Lernraten und eine größere Robustheit im Vergleich zu Standard-mSGD ermöglicht. Eine partielle Kompression ($p=1/4$) erzielt jedoch oft eine bessere Endleistung, sobald die Stabilität etabliert ist.
2. Begrenzte Gewinne bei RMS-normalisierten Inputs: Bei Anwendung auf RMS-normalisierte Inputs (Adam-Stil) führt spektrale Kompression zu abnehmenden Grenzerträgen. Volle Orthogonalisierung (AdamZ/Muon angewandt auf $O_{rms}$) übertrifft Standard-Adam nicht und kann schlechter als partielle Kompression sein. Dies deutet darauf hin, dass die elementweise RMS-Normalisierung bereits einen Großteil der Anisotropie mildert, die die spektrale Kompression adressiert.
3. Überlegenheit von Second-Moment-Methoden: Über alle spektralen Exponenten hinweg schneiden Optimierer, die eine Second-Moment-Normalisierung (RMS-basiert) nutzen, konsistent besser ab als ihre First-Moment-Gegenstücke (nur Momentum) – sowohl in Bezug auf die Stabilität als auch auf den finalen Loss.

Bedeutung und Behauptungen
Das Paper kommt zu dem Schluss, dass Muon am besten als eine effektive Form der spektralen Normalisierung und nicht als universell überlegene Optimierungsmethode verstanden werden sollte.

• Stabilisierung vs. Leistung: Muons primärer Beitrag liegt in der Stabilisierung von First-Moment-Updates, indem die Beträge der Singulärwerte verworfen werden, was verhindert, dass Richtungen mit großen Singulärwerten dominieren. Dieses „Abflachen“ ist jedoch ein grober Mechanismus; es entfernt die größenbasierte Filterung, was potenziell dazu führt, dass verrauschte Gradientenrichtungen (kleine Singulärwerte) verstärkt werden und informative dominante Modi unterschätzt werden.
• Kontextuelle Nützlichkeit: Während spektrale Kompression echte Stabilisierungsvorteile bietet, insbesondere wenn der Input-Update noch nicht normalisiert ist, ist die volle Orthogonalisierung ($p=0$) kein Ersatz für moderne RMS-normalisierte Second-Moment-Optimierer wie Adam. Die Autoren argumentieren, dass die berichteten Vorteile von Muon in vorangegangenen Arbeiten eher auf die spezifische Kombination von Techniken (z. B. QK-Norm) zurückzuführen sind als auf die Orthogonalisierung selbst.

Die Arbeit betont, dass spektrale Transformationen zwar mächtige Werkzeuge zur Kontrolle der Anisotropie sind, aber die Vorteile der adaptiven Second-Moment-Skalierung in groß angelegten Trainingsregimen nicht inhärent übertreffen.