Muon+: Towards Better Muon via One Additional Normalization Step

Die Arbeit stellt Muon+ vor, einen effizienten Optimierer, der durch eine zusätzliche Normalisierungsschritt nach der Orthogonalisierung die Leistung des ursprünglichen Muon-Optimierers beim Vortraining von Sprachmodellen verschiedener Größen konsistent verbessert.

Das Wesentliche

MUON+: Der „Glättungs-Schritt" für bessere KI-Trainings

Stell dir vor, du möchtest einen riesigen, komplexen Berg aus Daten (ein großes Sprachmodell wie GPT oder LLaMA) besteigen. Um diesen Berg zu erklimmen, brauchst du einen sehr klugen Führer, der dir sagt, in welche Richtung du treten sollst. In der Welt der künstlichen Intelligenz nennen wir diesen Führer einen Optimierer.

Bisher war ein Führer namens Muon sehr beliebt. Er war besonders gut darin, den Pfad zu glätten, indem er sicherstellte, dass die Schritte des KI-Modells nicht in die falsche Richtung „verkrampfen" (ein technisches Problem namens Rank Collapse). Er nutzte eine Art mathematisches Ballett, um die Bewegungen der KI perfekt aufeinander abzustimmen.

Aber die Forscher von UCSB haben sich gefragt: „Was wäre, wenn wir dem Führer noch eine kleine, zusätzliche Hilfe geben?"

Das Ergebnis ist MUON+.

Die einfache Idee: Der „Trocken-Tuch"-Effekt

Stell dir vor, du hast gerade einen Tanzschritt ausgeführt (das ist der Schritt des Muon-Optimierers). Du bist zwar in der richtigen Richtung, aber vielleicht bist du ein bisschen schief gelandet oder hast zu viel Schwung in die falsche Richtung mitgenommen.

Der alte Muon hat gesagt: „Gut gemacht, weiter so!"
Der neue MUON+ sagt: „Moment mal, bevor du den nächsten Schritt machst, wisch dir kurz das Gesicht ab und richte dich auf."

Diese „Wisch-Bewegung" ist der zusätzliche Normalisierungsschritt.

• Die Metapher: Stell dir vor, du malst ein Bild. Der Muon-Optimierer sorgt dafür, dass deine Pinselstriche orthogonal (im rechten Winkel) zueinander stehen, damit das Bild nicht verzerrt wird. MUON+ fügt einen Schritt hinzu, bei dem du nach jedem Strich prüfst: „Ist meine Handkraft überall gleich stark? Nicht zu fest hier, nicht zu schwach dort?" Und dann passt du die Kraft sofort an.

Was bringt das?

Die Forscher haben dieses kleine Extra bei verschiedenen KI-Modellen getestet – von kleinen Modellen (wie ein kleiner Roboter) bis hin zu riesigen Modellen (wie ein Super-Genie).

1. Bessere Ergebnisse: In fast jedem Test war MUON+ schneller und genauer als das alte Muon. Die KI lernte schneller und machte weniger Fehler (gemessen an der „Perplexität", was man sich wie eine Art „Verwirrtheits-Maß" vorstellen kann: Je niedriger, desto weniger verwirrt ist die KI).
2. Robustheit: Selbst wenn die Lernrate (wie schnell die KI lernt) nicht perfekt eingestellt war, schaffte es MUON+ immer noch gut. Das alte Muon wäre hier oft ins Straucheln gekommen.
3. Skalierbarkeit: Es funktioniert nicht nur bei kleinen Modellen, sondern auch, wenn man die KI extrem lange trainiert (mit Milliarden von Texten). Das ist wie ein Marathonläufer, der nicht nur am Start schnell ist, sondern auch nach 40 Kilometern noch frisch bleibt.

Warum ist das wichtig?

Das Training von großen KI-Modellen kostet heute eine unvorstellbare Menge an Strom und Rechenzeit (Millionen von Dollar). Wenn man einen Optimierer wie MUON+ verwendet, der nur eine winzige mathematische Änderung macht, aber dafür das Training effizienter macht, spart das:

• Geld: Weniger Rechenzeit.
• Zeit: Man kommt schneller zu einem besseren Modell.
• Energie: Weniger CO2-Ausstoß durch weniger Server-Laufzeit.

Fazit

Die Botschaft der Forscher ist einfach: Manchmal muss man nicht das ganze Auto neu bauen, um schneller zu fahren. Manchmal reicht es, einen kleinen, cleveren Schritt hinzuzufügen – wie das Abwischen des Gesichts nach dem Tanz –, um sicherzustellen, dass der nächste Schritt perfekt sitzt.

MUON+ ist dieser kleine, aber mächtige Schritt, der KI-Modelle stabiler, schneller und effizienter macht.

Technische Zusammenfassung

1. Problemstellung

Das Training großer Sprachmodelle (LLMs) wie GPT, LLaMA oder Gemini ist aufgrund der extremen Größe der Modelle und der Trainingsdatensätze mit enormen Rechenkosten verbunden. Obwohl Optimierer wie Adam und AdamW dominieren, gibt es einen starken Bedarf an effizienteren Alternativen, um die Vor-Trainingskosten zu senken.

Der Muon-Optimierer hat sich als vielversprechende Lösung etabliert. Sein Kernkonzept besteht darin, die Gradientenmatrix durch Orthogonalisierung (mittels Newton-Schulz-Iterationen) zu transformieren, um den „Rank Collapse" (den Verlust an Rang der Gradientenmatrix) zu verhindern. Dies hat zu einer weit verbreiteten Adoption in führenden Modellen (z. B. Kimi, GLM) geführt.

Das spezifische Problem: Trotz der Erfolge von Muon besteht die Frage, ob die Leistung durch die Orthogonalisierung allein maximiert wird oder ob weitere, einfache Modifikationen die Stabilität und Effizienz des Trainings weiter steigern können. Bisherige Arbeiten haben bereits Normalisierungsschritte in Varianten wie NorMuon oder Maon eingeführt, doch deren spezifischer Beitrag im Vergleich zu komplexeren Anpassungen (wie zweiten Momenten oder Mannigfaltigkeits-Formulierungen) war noch nicht vollständig isoliert und verstanden.

2. Methodik: MUON+

Die Autoren schlagen MUON+ vor, eine einfache, aber effektive Erweiterung des Muon-Optimierers. Die Methode fügt einen zusätzlichen Normalisierungsschritt direkt nach der Orthogonalisierung ein.

• Der Update-Prozess:

  1. Momentum-Aktualisierung: Wie bei Muon wird der Momentum-Term $M_t$ aktualisiert.
  2. Orthogonalisierung: Die Matrix $M_t$ wird orthogonalisiert ($O_t = \text{Ortho}(M_t)$), typischerweise approximiert durch Newton-Schulz-Iterationen.
  3. Neuer Schritt (MUON+): Auf die orthogonalisierte Matrix wird ein Normalisierungsoperator $\text{Norm}(d)$ angewendet.
  4. Gewichtsupdate: Die Parameter werden aktualisiert basierend auf dem normalisierten, orthogonalen Gradienten.

• Normalisierungsrichtungen:
  Die Autoren untersuchen verschiedene Richtungen für die $\ell_2$-Normalisierung:

  • Spaltenweise (Column-wise): Normalisierung der Spaltenvektoren.
  • Reihenweise (Row-wise): Normalisierung der Zeilenvektoren.
  • Kombiniert: Sequenzielle Anwendung (z. B. erst Spalten, dann Reihen).

Der Algorithmus bleibt dabei strukturell sehr ähnlich zu Muon, fügt jedoch diese eine Normalisierungsoperation hinzu, um die Skalierung der Updates zu stabilisieren.

3. Wichtige Beiträge

• Einführung von MUON+: Ein neuer Optimierer, der Muon durch einen einzigen zusätzlichen Normalisierungsschritt verbessert.
• Umfassende Evaluation: Die Leistung wurde über ein breites Spektrum an Modellarchitekturen (GPT- und LLaMA-Stil) und Größen (von 60M bis 1B Parameter) getestet.
• Skalierbarkeitstests: Evaluation sowohl im „compute-optimalen" Regime (Token-zu-Parameter-Verhältnis $\approx 20$) als auch im „Overtraining"-Regime mit industriellen Verhältnissen ($\approx 200$).
• Ablationsstudien: Systematische Analyse der Einflussfaktoren:
  • Isolierung des Normalisierungseffekts von anderen Komponenten (wie zweiten Momenten bei NorMuon).
  • Vergleich verschiedener Normalisierungsrichtungen.
  • Robustheit gegenüber verschiedenen Approximationsmethoden für die Orthogonalisierung (SVD-Näherungen).

4. Ergebnisse

Die experimentellen Ergebnisse zeigen konsistent und signifikante Verbesserungen von MUON+ gegenüber dem Basis-Muon-Optimierer:

• Perplexity-Verbesserung: MUON+ erzielt auf allen getesteten Modellen (GPT-Small bis Large, LLaMA-60M bis 1B) niedrigere Validierungs-Perplexities.
  • Bei GPT-Modellen wurden Verbesserungen von bis zu 2.02 Perplexity-Punkten (bei GPT-Small) erreicht.
  • Bei LLaMA-Modellen waren die Gewinne ebenfalls konsistent (z. B. -0.50 bei LLaMA-60M).
• Robustheit bei Overtraining: Selbst bei einem extrem hohen Token-zu-Parameter-Verhältnis von 200 (72 Milliarden Tokens für GPT-Base und LLaMA-350M) behielt MUON+ seine Überlegenheit bei und zeigte stabilere Optimierungsverläufe ohne Degradation in späteren Trainingsphasen.
• Hyperparameter-Stabilität: MUON+ ist weniger empfindlich gegenüber der Wahl der Lernrate. Bei suboptimalen (zu großen) Lernraten zeigt MUON+ eine deutlich geringere Leistungsverschlechterung als Muon.
• Ablationserkenntnisse:
  • Der Normalisierungsschritt allein ist der Haupttreiber für die Leistungssteigerung. Zusätzliche Komponenten wie die Anpassung des zweiten Moments (wie in NorMuon) brachten in den getesteten Szenarien keine weiteren signifikanten Vorteile.
  • Die Reihenfolge der Normalisierung (Spalten vs. Reihen) spielt eine Rolle; die Kombinationen (z. B. col_row oder row_col) erzielten die besten Ergebnisse.
  • Die Methode ist unabhängig von der spezifischen SVD-Approximation (z. B. Jordan, You, PolarExpress) robust.

5. Bedeutung und Fazit

Die Arbeit demonstriert, dass die Leistung von Optimierern für große Sprachmodelle durch eine strukturelle Normalisierung der orthogonalen Updates signifikant gesteigert werden kann, ohne die Komplexität des Algorithmus drastisch zu erhöhen.

• Praktische Relevanz: Da MUON+ einfach zu implementieren ist (einzeilige Änderung im Update-Schritt) und in bestehenden Pipelines (wie denen von Kimi oder GLM) leicht integrierbar ist, bietet es einen sofortigen Weg zur Verbesserung des Trainings von Foundation-Modellen.
• Theoretischer Einsicht: Die Ergebnisse legen nahe, dass die Stabilität des Trainings weniger von komplexen Anpassungen des zweiten Moments abhängt, sondern primär von der geometrischen Strukturierung (Normalisierung) der orthogonalen Updates.
• Zukunftsausblick: MUON+ etabliert einen neuen Standard für das Vor-Training großer Modelle, insbesondere in Szenarien mit begrenzten Rechenressourcen oder extrem langen Trainingshorizonten, und unterstreicht die Wichtigkeit von Normalisierungsschritten in der Matrix-Optimierung.

Zusammenfassend bietet MUON+ einen klaren, reproduzierbaren und effektiven Weg, um die Effizienz und Endqualität von großen Sprachmodellen zu verbessern, indem er ein einfaches Prinzip (Normalisierung nach Orthogonalisierung) konsequent anwendet.