Technische Zusammenfassung: LionMuon

Problemstellung

Das Training von Large Language Models (LLMs) ist ein massives Optimierungsproblem, bei dem die pro-Schritt-Rechenkosten (FLOPs) und die Konvergenzgeschwindigkeit die primären Determinanten der endgültigen Rechenkostenrechnung sind. Aktuelle Optimierer stehen vor einem Zielkonflikt zwischen der Qualität der Aktualisierung und der Recheneffizienz:

• Vorzeichenbasierte Methoden (z. B. Lion, Signum): Diese erzeugen kostengünstige, koordinatenweise Vorzeichen-Aktualisierungen ($W_{t+1} = W_t - \eta \cdot \text{sign}(\hat{G}_t)$). Sie sind rechnerisch günstig und speichereffizient, liefern jedoch oft schwächere Aktualisierungsrichtungen im Vergleich zu spektralen Methoden.
• Spektrale Methoden (z. B. Muon): Diese berechnen eine Matrix-Vorzeichen-Aktualisierung via Newton-Schulz-Iterationen ($W_{t+1} = W_t - \eta \cdot \text{msign}(\hat{G}_t)$). Obwohl sie stärkere spektrale Richtungen liefern, die häufig zu einer schnelleren Konvergenz führen, verursachen sie aufgrund mehrerer Matrixmultiplikationen und erhöhter Kommunikationskosten in verteilten Umgebungen (erforderlich für vollständiges Matrix-All-Gather) signifikant höhere Kosten pro Schritt.

Die zentrale Frage lautet, ob es möglich ist, die niedrigen Kosten vorzeichenbasierter Aktualisierungen mit den hochwertigen Richtungen spektraler Aktualisierungen zu kombinieren, um einen überlegenen Kosten-Qualitäts-Kompromiss zu erreichen.

Methodik: LionMuon

Die Autoren schlagen LionMuon vor, einen Optimierer, der in einem festen Intervall $P$ zwischen Lion-artigen Vorzeichen-Aktualisierungen und Muon-artigen spektralen Aktualisierungen wechselt.

Kernalgorithmus

1. Wechselnder Zeitplan: Der Algorithmus führt jeden $P$-ten Iterationsschritt einen Muon-Schritt (Newton-Schulz-spektrale Aktualisierung) aus und für die verbleibenden $P-1$ Iterationen Lion-artige Vorzeichen-Aktualisierungen.
2. Gemeinsamer Momentum-Puffer: Beide Aktualisierungstypen teilen sich einen einzigen Dual-Exponential Moving Average (EMA) Momentum-Puffer ($M_t$). Die Richtung $\hat{G}_t$ wird über Lions Interpolation berechnet: $\hat{G}_t = \beta_1 M_{t-1} + (1-\beta_1)G_t$.
3. Speichereffizienz: LionMuon verwaltet nur einen Momentum-Puffer, was der Speichernutzung von Lion/Signum entspricht und genau die Hälfte der von AdamW beträgt.
4. Sonderfälle:
   • $P=1$: Reines Muon.
   • $P=\infty$: Reines Lion (oder Signum, falls $\beta_1=\beta_2$).
   • $P \in (1, \infty)$: Die vorgeschlagene Wechselmethode.
   • SignMuon: Eine Variante, die einen Single-EMA-Momentum ($\beta_1 = \beta_2$) mit demselben wechselnden Zeitplan verwendet.

Theoretische Analyse

Die Autoren liefern Konvergenzschranken für LionMuon unter schwerfälligen (heavy-tailed) Rauschen, einem Zustand, der beim LLM-Training häufig vorkommt.

• Interpolierte Konstanten: Die Konvergenzschranke hängt von periodenmittelten Glättungskonstanten ($\bar{L}$) und Rauschniveaus ($\bar{\rho}$) ab, die zwischen den Konstanten von reinem Muon (Spektralnorm) und reinem Lion (Unendlich-Norm) interpolieren.
• Optimale Periode: Die Theorie legt nahe, dass eine optimale Periode $P^*$ existiert, die einen Kompromissfaktor minimiert, der die Verhältnisse der Glättungskonstanten ($L_\infty / L_2$) und der Rauschniveaus ($\rho_1 / \rho_{nuc}$) umfasst.
• Komplexität: Die hergeleiteten Komplexitätsschranken zeigen, dass LionMuon unter spezifischen Konfigurationen, insbesondere wenn der Kompromissfaktor $\phi(P) < 1$ ist, die gleiche Genauigkeit wie reines Muon oder Lion mit weniger Operationen erreichen kann.

Hauptbeiträge

1. SignMuon: Ein alternierender Optimierer, der jeden $P$-ten Iterationsschritt einen Muon-Schritt und ansonsten einen Signum-artigen Vorzeichen-Schritt ausführt. Die Autoren zeigen, dass selbst diese Single-EMA-Variante den Trade-off zwischen Verlust und FLOPs bei $P=2$ gegenüber reinem Muon verbessert.
2. LionMuon: Der Hauptbeitrag, der das Momentum von Signum durch die Dual-EMA-Regel von Lion ersetzt. Er bewahrt die Kosteneinsparungen von SignMuon und verbessert gleichzeitig die Konvergenz weiter. Er erfordert nur einen zusätzlichen ganzzahligen Hyperparameter ($P$) und ist ein nahezu direktes Ersatzmodul für Lion.
3. Theoretische Schranken: Scharfe Komplexitätsschranken unter schwerfälligen Rauschen, die die Interpolation zwischen Muon- und Lion-Regimes explizit modellieren und die rechenoptimale Periode $P$ vorhersagen.
4. Empirische Validierung: Umfassende Experimente über mehrere Datensätze (FineWeb, SlimPajama, WikiText-103) und Architekturen (GPT-2, LLaMA) in den Skalierungen von 124M, 355M und 720M Parametern.

Experimentelle Ergebnisse

• Skalierung 124M: LionMuon (speziell mit $P=2$) dominiert Muon, Lion, Signum und AdamW in Bezug auf die Pareto-Optimalität auf allen getesteten Datensätzen und Architekturen. Es erreicht einen niedrigeren Validierungsverlust mit etwa 5,5 % weniger gesamten Trainings-FLOPs im Vergleich zu reinem Muon.
• Skalierung (355M & 720M): Der Vorteil der Konvergenzqualität bleibt bei größeren Skalierungen erhalten. Bei 355M übertreffen LionMuon und SignMuon ($P=2$) reines Muon. Bei 720M schlägt SignMuon ($P=2$) immer noch reines Muon, obwohl sich der Dual-EMA-Vorteil von LionMuon gegenüber SignMuon in diesem spezifischen Budget für untertrainierte Modelle nicht reproduzieren ließ.
• Verteiltes Training: Die Methode bietet strukturelle Kommunikationseinsparungen. Da Lion-Schritte elementweise und lokal zu den Shards sind, halbiert die Einstellung von $P=2$ die Häufigkeit der teuren All-Gather-Operationen, die für den Newton-Schulz-Schritt von Muon erforderlich sind, unabhängig von der Batch-Größe.

Bedeutung und Behauptungen

Die Arbeit behauptet, dass LionMuon einen strukturell überlegenen Ansatz für das LLM-Training bietet, indem er die Häufigkeit teurer spektraler Aktualisierungen von der Iterationszahl entkoppelt.

• Effizienz: Er behält die Wirksamkeit von Muon-Schritten bei, senkt jedoch die durchschnittlichen Kosten pro Iteration und entspricht der Speichereffizienz vorzeichenbasierter Methoden.
• Robustheit: Die Methode ist über verschiedene Modellgrößen und Datensätze hinweg robust und definiert konsistent die Pareto-Grenze für Verlust-vs-FLOPs.
• Theoretische Übereinstimmung: Die empirischen Ergebnisse stimmen mit den theoretischen Vorhersagen bezüglich der optimalen Periode $P$ überein und deuten darauf hin, dass mittlere Werte (z. B. $P=2$ oder $P=5$) oft den besten Ausgleich zwischen der teuren spektralen Richtung und der günstigen Vorzeichenrichtung bieten.

Die Autoren schließen, dass LionMuon eine praktische, direkte Erweiterung für das moderne LLM-Training ist, die den aktuellen Stand der Technik (Muon) sowohl in Bezug auf die Konvergenzqualität als auch auf die Recheneffizienz strikt verbessert. Als zukünftige Arbeiten werden eine adaptive Auswahl von $P$ und die Skalierung über 720M hinaus vorgeschlagen.