Technisches Fazit: Muown – Zeilennorm-Steuerung für die Muon-Optimierung

Problemstellung

Muon hat sich als ein leistungsfähiger Optimierer für das Vortraining großer Sprachmodelle (LLM) etabliert, der AdamW häufig durch Ausnutzung der Matrixstruktur linearer Transformer-Schichten übertrifft. Neuere Beobachtungen deuten jedoch darauf hin, dass das Verhalten von Muon im großen Maßstab empfindlich auf Gewichtsabnahme (Weight Decay) reagiert und anfällig für numerische Instabilitäten ist, wie etwa Überläufe im bfloat16-Format und explodierende Attention-Logits.

Ein kritisches identifiziertes Problem ist der konsistente Aufwärtstrend der Spektralnorm ($\|W_t\|_{S_\infty}$) von Gewichtsmatrizen während des Trainings, wenn Muon ohne entkoppelte Gewichtsabnahme verwendet wird. Diese Drift verletzt die spektrale Bedingung, die für stabiles Feature-Learning erforderlich ist ($\|W\|_{S_\infty} = \Theta(\sqrt{m/n})$). Während entkoppelte Gewichtsabnahme dies durch eine gleichmäßige Verkleinerung der Gewichtsmatrix mildern kann, stellt dies einen nicht zielgerichteten Eingriff dar, der das gesamte Spektrum skaliert, unabhängig davon, welche spezifischen Richtungen das Wachstum antreiben. Dies kann die Konvergenz in frühen Trainingsphasen potenziell verlangsamen. Die Arbeit zielt darauf ab, den spezifischen empirischen Treiber dieser Spektralnorm-Drift zu identifizieren und einen gezielteren Steuerungsmechanismus vorzuschlagen.

Methodik: Muown

Die Autoren stellen Muown vor, einen direkten Ersatz für Muon, der die Spektralnorm-Drift durch eine Umparametrisierungsstrategie anstelle von aggregierten Einschränkungen adressiert.

1. Spektralnorm-Zerlegung

Der Kern der Erkenntnis ist eine Zerlegung der Spektralnorm in zwei Faktoren (Proposition 1):
$$\|W_t\|_{S_\infty}^2 = \|g_t\|_\infty^2 \cdot \lambda_{\max}(P_t C_t P_t)$$
Wobei:

• $g_t = \|W_t\|_{row}$ der Vektor der Zeilenbeträge (Euklidische Normen der Zeilen) ist.
• $D_t$ die zeilennormierte Matrix ist ($W_t = \text{Diag}(g_t)D_t$).
• $C_t = D_t D_t^\top$ die Gram-Matrix der Einheitszeilen ist.
• $P_t$ eine Diagonalmatrix der normalisierten Zeilenbeträge ist.
• $\lambda_{\max}(P_t C_t P_t)$ den Zeilen-Kohärenz-Faktor darstellt.

Eine empirische Analyse an 500-M-Modellen zeigt, dass der Zeilenbetrags-Faktor ($\|g_t\|_\infty$) der primäre Treiber der Spektralnorm-Drift ist und sich synchron mit $\|W_t\|_{S_\infty}$ bewegt. Im Gegensatz dazu bleibt der Zeilen-Kohärenz-Faktor während des gesamten Trainings wohlverhalten und beschränkt.

2. Umparametrisierung und Optimierung

Basierend auf dieser Diagnose parametrisiert Muown die Gewichtsmatrix $W$ wie folgt um:
$$W(g, R) = \text{Diag}\left(\frac{g}{\|R\|_{row}}\right) R$$
Hierbei repräsentiert $g \in \mathbb{R}^m$ den expliziten Zeilenbetrags-Vektor und $R \in \mathbb{R}^{m \times n}$ die Richtungskomponente.

• Richtungskomponente ($R$): Muon wird unverändert auf $R$ angewendet und behandelt diese unter der Spektralnorm-Geometrie (steilster spektraler Abstieg).
• Betragskomponente ($g$): Die Zeilenbeträge $g$ werden zu expliziten Optimierer-Variablen befördert. Sie werden unter der $\ell_\infty$-Geometrie aktualisiert, was die natürliche Dualnorm für die Operatornorm von Diagonalmatrizen darstellt.
• Aktualisierungsregel: Die Aktualisierung für $g$ nutzt einen Adam-Optimierer (der als geglätteter SignSGD unter $\ell_\infty$-Geometrie wirkt), um stochastisches Rauschen zu handhaben, während die Aktualisierung für $R$ dem Standard-Muon-Verfahren mit Newton-Schulz-Orthogonalisierung folgt.

Dieser Ansatz wird implizit innerhalb des Optimierer-Zustands implementiert, wobei der Vorwärtsdurchlauf der Modellarchitektur unverändert bleibt.

Hauptbeiträge

1. Diagnose der Spektral-Drift: Die Arbeit zerlegt die Spektralnorm, um die maximale Zeilenbetragsgröße ($\|g_t\|_\infty$) als empirischen Treiber des Spektralnorm-Wachstums unter Muon zu identifizieren und sie vom stabilen Zeilen-Kohärenz-Faktor zu unterscheiden.
2. Muown-Algorithmus: Ein neuer Optimierer, der Zeilenbeträge als explizite trainierbare Variablen behandelt, die unter $\ell_\infty$-Geometrie aktualisiert werden, während die Richtungsaktualisierungen von Muon erhalten bleiben. Dies vermeidet die unselektive Verkleinerung durch Gewichtsabnahme und die Orthogonalitäts-Apriori von spektralen Umparametrisierungen.
3. Konvergenzgarantien: Die Autoren beweisen, dass Muown optimale Konvergenzraten in nicht-konvexen Szenarien sowohl im deterministischen als auch im stochastischen Regime erreicht. Die Analyse erfolgt in einer gemischten Dualnorm ($\ell_1$ für Beträge, Schatten-1 für Richtungen), die mit den Aktualisierungsgeometrien übereinstimmt. Die theoretische Schranke deutet auf eine Reduktion des effektiven stochastischen Rausch-Koeffizienten im Vergleich zu Muon hin.
4. Empirische Validierung: Umfassende Experimente auf FineWeb-Edu über Modellgrößen von 124M bis 2,7B Parametern zeigen konsistente Verbesserungen.

Experimentelle Ergebnisse

• Perplexitäts-Verbesserungen: Muown erzielt konsistent eine niedrigere Perplexität als Muon, SOAP, AdamW und Lion über alle getesteten Modellskalen (124M bis 2,7B) und Lernraten hinweg.
• Robustheit gegenüber Hyperparametern:
  • Lernrate: Muown erweitert das Plateau nahezu optimaler Lernraten über verschiedene Modellbreiten hinweg und reduziert die Empfindlichkeit gegenüber einer skalenabhängigen Lernraten-Tuning.
  • Gewichtsabnahme: Muown mildert die Notwendigkeit des Tunings der Gewichtsabnahme; es funktioniert gut ohne Gewichtsabnahme ($\lambda=0$), wohingegen Muon eine sorgfältige Abstimmung (optimales $\lambda=0,1$) erfordert, um Divergenz oder suboptimale Leistung zu verhindern.
• Stabilität: Muown verhindert die systematische Spektralnorm-Drift, die bei Muon beobachtet wird. In Regimen mit hohen Lernraten (z. B. $4 \times 10^{-3}$) divergiert Muon häufig, während Muown stabil und nahezu optimal bleibt.
• Gradienteneigenschaften: Analysen zeigen, dass Muown einen höheren normalisierten effektiven Rang von Minibatch-Gradienten beibehält und während des gesamten Trainings einen niedrigeren stochastischen Rausch-Koeffizienten aufweist als Muon.
• Rechen-Overhead: Die Methode führt zu einem vernachlässigbaren Overhead pro Schritt. Wenn der Optimierer über Ranks hinweg geteilt wird (verteilte Umgebung), werden die Schrittzzeiten für Muown und Muon ununterscheidbar. Der Speicher-Overhead ist minimal (linear in der Anzahl der Zeilen) und fügt nur wenige Megabyte pro Schicht hinzu.

Bedeutung und Behauptungen

Die Arbeit behauptet, dass die bei Muon beobachtete Spektralnorm-Drift keine irreduzible Eigenschaft des Optimierers ist, sondern eine Folge unkontrollierten Zeilenbetrags-Wachstums. Durch die Isolierung dieses Treibers und seine Beförderung zu einer expliziten Optimierer-Variable erreicht Muown:

• Verbesserte Konvergenz: Bessere Perplexität und Stabilität ohne die Notwendigkeit aggressiver Gewichtsabnahme.
• Theoretische Fundierung: Optimale Konvergenzraten in einem geometrie-bewussten Rahmen, der die gemischte Natur der Aktualisierungen explizit berücksichtigt.
• Praktische Effizienz: Ein direkter Ersatz, der in verteilten Umgebungen die Laufzeit von Muon erreicht und gleichzeitig eine überlegene Robustheit gegenüber Hyperparameter-Wahl bietet.

Die Autoren positionieren Muown als eine Lösung, die die Lücke zwischen den theoretischen Stabilitätsanforderungen des Feature-Learnings und den praktischen Bedürfnissen des großskaligen LLM-Trainings schließt und eine gezieltere Alternative zur entkoppelten Gewichtsabnahme bietet.