To Use or not to Use Muon: How Simplicity Bias in Optimizers Matters

Diese Arbeit untersucht die potenziellen Nachteile des Optimierers Muon, insbesondere den Verlust eines durch SGD bewahrten Einfachheitsbias, der zu einer stärkeren Neigung zu Spurious Features und einer geringeren Fähigkeit, gemeinsame Strukturen zu erkennen, führen kann.

Das Wesentliche

🚀 Das Dilemma: Der Rennwagen oder der Entdecker?

Stell dir vor, du möchtest ein neues Haus bauen (ein KI-Modell trainieren). Dafür brauchst du einen Bauleiter (einen Optimierer), der entscheidet, wie die Steine gelegt werden, um das Haus schnell und stabil zu errichten.

Lange Zeit war SGD (Stochastischer Gradientenabstieg) der Standard-Bauleiter. Er ist etwas langsamer, aber sehr methodisch. Er baut das Haus Schritt für Schritt: Erst das Fundament, dann die Wände, dann das Dach. Er lernt die einfachen Dinge zuerst und arbeitet sich dann zu den komplexen Details vor. Man nennt das einen „Einfachheits-Bias" (Simplicity Bias). Er ignoriert unnötigen Schnickschnack und sucht nach der einfachsten, robustesten Lösung.

Dann kam Muon auf den Plan. Muon ist wie ein superschneller Rennwagen-Bauleiter. Er ist extrem schnell und kann das Haus in der Hälfte der Zeit fertigstellen. Viele haben ihn sofort übernommen, weil er so effizient wirkt.

Aber die Forscher dieser Arbeit (Dragutinović & Ranganath) fragen sich:
Ist Muon wirklich besser? Oder ist er nur schnell, aber baut ein Haus, das bei Sturm zusammenfällt?

🔍 Die Entdeckung: Was Muon anders macht

Die Forscher haben herausgefunden, dass Muon einen entscheidenden Unterschied hat: Er lernt alles auf einmal, statt Schritt für Schritt.

• SGD (Der langsame Entdecker): Er lernt zuerst die wichtigsten Muster (z. B. „Das ist eine Katze"). Erst wenn er das sicher kann, lernt er die Details (z. B. „Die Katze hat einen braunen Fleck"). Wenn er auf ein Hindernis (ein mathematisches Problem) stößt, bleibt er kurz stehen, überlegt und geht dann weiter. Dieser „Stopp" ist wichtig, weil er sicherstellt, dass das Modell wirklich die Struktur der Welt versteht und nicht nur auswendig lernt.
• Muon (Der schnelle Alleskönner): Er nimmt sich vor, alles gleichzeitig zu lernen. Er ignoriert die Hindernisse und schießt einfach durch. Das macht ihn unglaublich schnell. Aber weil er nicht Schritt für Schritt vorgeht, verliert er den „Einfachheits-Bias". Er lernt nicht nur die wichtigen Regeln, sondern auch den unnötigen Schnickschnack und die Zufälligkeiten in den Daten.

🧩 Die zwei Experimente: Wo Muon scheitert

Die Forscher haben zwei Szenarien getestet, um zu zeigen, warum Muons Geschwindigkeit manchmal ein Nachteil ist.

1. Das Puzzle-Rätsel (Geteilte Strukturen)

Stell dir vor, du musst einem Kind beibringen, Zahlen zu erkennen.

• Der Trick: Die Zahlen kommen in verschiedenen Farben vor (Rot, Blau, Grün). Die eigentliche Regel ist: „Die Form der Zahl ist wichtig, die Farbe nicht."
• SGD: Lernt zuerst die Form (die wichtige Regel). Das Kind versteht: „Eine 1 ist eine 1, egal ob rot oder blau." Es kann die Regel auf neue Farben übertragen.
• Muon: Lernt Form und Farbe gleichzeitig. Das Kind merkt sich: „Die rote 1 ist eine 1, die blaue 1 ist eine 1." Es hat die Regel nicht verstanden, sondern nur die spezifischen Beispiele auswendig gelernt. Wenn du ihm eine gelbe 1 zeigst, ist es verwirrt.
• Ergebnis: Muon war schneller im Training, aber das Kind konnte die Regel nicht auf neue Situationen übertragen. SGD war langsamer, aber das Kind war schlauer.

2. Der falsche Hinweis (Spurious Features)

Stell dir vor, du trainierst eine KI, um Autos zu erkennen. In den Trainingsfotos stehen die Autos immer auf einer grünen Wiese.

• Der Fehler: Die KI könnte denken: „Grüner Hintergrund = Auto". Das ist ein falscher Hinweis (ein „spurious feature").
• SGD: Lernt zuerst die wichtigen Teile (Räder, Scheinwerfer). Erst wenn es das sicher kann, merkt es vielleicht, dass der grüne Hintergrund auch da ist. Wenn du die KI später auf einer grauen Straße testest (ohne Gras), erkennt sie das Auto trotzdem, weil sie die wahren Merkmale gelernt hat.
• Muon: Lernt Räder und grünen Hintergrund gleichzeitig. Da der grüne Hintergrund sehr auffällig ist, lernt die KI ihn sofort mit. Sie verlässt sich zu stark auf den falschen Hinweis. Auf der grauen Straße versagt sie komplett.
• Ergebnis: Muon ist schneller, aber anfälliger für Täuschungen. SGD ist robuster.

💡 Die große Lehre: Geschwindigkeit ist nicht alles

Die Botschaft der Arbeit ist wie folgt:

Wenn du einen neuen Optimierer (wie Muon) entwickelst, darfst du nicht nur auf die Geschwindigkeit schauen. Du musst auch fragen: „Wie denkt dieser Optimierer?"

• Muon ist wie ein Student, der alles auf einmal auswendig lernt, um die Prüfung schnell zu bestehen. Er ist schnell, aber wenn die Fragen anders formuliert sind, scheitert er.
• SGD ist wie ein Student, der erst die Grundprinzipien versteht und dann die Details lernt. Er braucht länger, aber er kann das Gelernte auf neue Probleme anwenden.

Fazit:
Muon ist ein tolles Werkzeug, wenn es nur um rohe Geschwindigkeit geht. Aber wenn wir KI-Modelle bauen, die wirklich verstehen sollen, wie die Welt funktioniert (und nicht nur Daten auswendig lernen), dann ist die langsame, schrittweise Methode von SGD manchmal besser. Wir müssen aufpassen, dass wir beim Streben nach Geschwindigkeit nicht die Intelligenz unseres Modells verlieren.

Kurz gesagt: Ein schneller Optimierer ist gut, aber ein weiser Optimierer ist besser. Und manchmal ist „langsam und stetig" der Weg zum Sieg.

Technische Zusammenfassung

Problemstellung

Der Adam-Optimierer war lange Zeit der Standard für das Training tiefer neuronaler Netze. Kürzlich hat der Muon-Optimierer (MomentUm Orthogonalized by Newton-Schulz) aufgrund seiner überlegenen Trainingsgeschwindigkeit an Popularität gewonnen und wird zunehmend als neuer Standard für das Training von Sprachmodellen (z. B. NanoChat) diskutiert.

Das zentrale Problem dieses Papers ist jedoch, dass die induktiven Verzerrungen (Biases) von Muon kaum verstanden sind. Während Muon die Verlustlandschaft schneller durchläuft als traditionelle Methoden wie SGD (Stochastic Gradient Descent), ist unklar, welche Pfade es nimmt und welche Konsequenzen dies für die finale Lösung hat. Die Autoren argumentieren, dass die reine Fokussierung auf die Trainingsgeschwindigkeit (Greedy-Benefits) die potenziellen Nachteile ignoriert, die durch den Verlust eines natürlichen Simplicity Bias (Einfachheitsverzerrung) entstehen.

Methodik

Die Autoren untersuchen die Auswirkungen von Muon durch eine Kombination aus theoretischer Analyse und empirischen Experimenten:

1. Theoretische Analyse (Deep Linear Networks):

   • Um die Dynamik analytisch handhabbar zu machen, analysieren die Autoren tiefe lineare Netze (2-Schichten).
   • Sie führen einen vereinfachten Optimierer namens Spectral Gradient Descent (Spectral GD) ein. Dieser übernimmt die Kernidee von Muon (Orthogonalisierung der Update-Schritte mittels SVD), verzichtet aber auf Momentum und approximierende Verfahren (Newton-Schulz), um eine exakte mathematische Analyse zu ermöglichen.
   • Der Vergleich erfolgt zwischen klassischem Gradient Descent (GD) und Spectral GD.

2. Empirische Experimente:

   • Die Autoren testen die theoretischen Vorhersagen in zwei spezifischen Szenarien, in denen der Verlust des Simplicity Bias kritisch sein könnte:
     • Shared Representations (Routing-Aufgabe): Ein Multi-Modal-Setup, bei dem das Modell eine gemeinsame zugrundeliegende Struktur über verschiedene Eingabedomänen hinweg lernen muss, um Out-of-Distribution (OOD) zu generalisieren.
     • Spurious Features (MNIST): Ein Klassifikationsproblem mit MNIST-Daten, bei dem ein irrelevanter, aber stark korrelierter Pixel (spurious feature) als Täuschungsmöglichkeit dient.

Theoretische Erkenntnisse & Key Contributions

• Verlust des Simplicity Bias:

  • GD (Gradient Descent): Folgt einer „Saddle-to-Saddle"-Dynamik. Es lernt die Lösung sequenziell: Zuerst werden die dominanten Singulärwerte (hohe Varianz/Komponenten) gelernt, danach folgen die schwächeren. Dies führt zu einer schrittweisen Erhöhung des Rangs der Lösung und wirkt als implizite Regularisierung (Simplicity Bias).
  • Spectral GD (und Muon): Lernt alle Singulärwerte gleichzeitig und mit gleicher Geschwindigkeit. Es umgeht die Sattelpunkte, die GD durchqueren muss, was die Geschwindigkeit erklärt. Allerdings fehlt dabei die sequenzielle Komplexitätssteigerung.

• Konsequenzen:

  • Der Verlust des Simplicity Bias bedeutet, dass Muon/Spectral GD keine implizite „Lehrplan"-Struktur (Curriculum) hat, die zuerst einfache, dominante Muster lernt. Stattdessen versucht es, alle Muster (sowohl signifikante als auch Rauschen) gleichzeitig zu memorieren.

Ergebnisse

1. Versagen beim Erlernen gemeinsamer Repräsentationen (Routing-Aufgabe):

   • In der Routing-Aufgabe muss das Modell eine gemeinsame Abbildung lernen, die über verschiedene Eingabedomänen hinweg gilt.
   • SGD: Lernt die gemeinsame, niedrigrangige Struktur (Rank = 4) und generalisiert erfolgreich auf nicht gesehene Eingabe-Ausgabe-Paare.
   • Spectral GD (Muon): Lernt zwar die Trainingsdaten perfekt, aber durch Memorization (Auswendiglernen) statt durch Strukturverständnis. Die Lösung hat einen deutlich höheren effektiven Rang und generalisiert schlecht auf OOD-Daten.

2. Empfindlichkeit gegenüber spurious Features:

   • In der MNIST-Aufgabe mit einem spurious Pixel (der stark mit der Klasse korreliert, aber irrelevant ist):
   • SGD: Lernt zuerst die dominanten Merkmale (die tatsächlichen Ziffern) und ignoriert das spurious Feature zunächst länger. Dies führt zu einer besseren Peak-Genauigkeit auf sauberen Validierungsdaten, wenn Early Stopping verwendet wird.
   • Muon: Lernt sowohl die Ziffern als auch das spurious Feature gleichzeitig. Es beginnt früher, sich auf das spurious Feature zu verlassen, was zu einer schnelleren Verschlechterung der Leistung auf Daten ohne spurious Features führt.
   • Nuance: Wenn das spurious Feature jedoch so dominant wird, dass es die eigentliche Aufgabe überdeckt, kann Muon sogar besser abschneiden, da es nicht so stark auf die sequenzielle Reihenfolge angewiesen ist.

Signifikanz und Fazit

Das Paper liefert eine wichtige Warnung an die Community: Trainingsgeschwindigkeit ist nicht das einzige Kriterium für die Qualität eines Optimierers.

• Induktive Verzerrungen sind entscheidend: Die Wahl des Optimierers bestimmt nicht nur die Konvergenzgeschwindigkeit, sondern auch die Art der gefundenen Lösung. Muon entfernt den Simplicity Bias, der bei SGD natürlich vorhanden ist und oft für bessere Generalisierung und das Finden robusterer Strukturen sorgt.
• Trade-off: Der Geschwindigkeitsvorteil von Muon (durch das Umgehen von Sattelpunkten und paralleles Lernen) kommt auf Kosten der Fähigkeit, gemeinsame zugrundeliegende Strukturen zu entdecken und sich gegen spurious Korrelationen zu schützen.
• Zukunftsausblick: Die Autoren fordern, dass bei der Entwicklung neuer Optimierer die induktiven Verzerrungen systematisch analysiert werden müssen. Ein idealer zukünftiger Optimierer könnte die Geschwindigkeit von Muon mit dem strukturellen Bias von SGD kombinieren, um sowohl schnell als auch robust zu sein.

Zusammenfassend zeigt das Paper, dass Muon in Szenarien, in denen das Verständnis der zugrundeliegenden Datenstruktur wichtiger ist als das reine Auswendiglernen von Mustern (z. B. bei mathematischen Aufgaben oder komplexen Generalisierungsaufgaben), schlechter abschneiden kann als etablierte Methoden wie SGD.