Muon übertrifft Adam beim Lernen von assoziativen Speichern im Tail-End-Bereich

papers.abstract

Der Muon-Optimierer ist durchweg schneller als Adam beim Training von Large Language Models (LLMs), doch der Mechanismus, der seinem Erfolg zugrunde liegt, bleibt unklar. Diese Arbeit beleuchtet diesen Mechanismus durch die Linse des assoziativen Gedächtnisses. Durch das Ablieren der Transformer-Komponenten, die von Muon optimiert werden, zeigen wir, dass die Parameter des assoziativen Gedächtnisses von LLMs, nämlich die Value- und Output-Aufmerksamkeitsgewichte (VO) und die Feed-Forward Networks (FFNs), die Hauptbeiträge zur Überlegenheit von Muon sind. Motiviert durch diese Sichtweise des assoziativen Gedächtnisses erklären wir dann die Überlegenheit von Muon auf realen Korpora, die intrinsisch schwergewichtig sind: einige Klassen (Tail-Klassen) treten weitaus seltener auf als andere. Die Überlegenheit wird durch zwei Schlüsseleigenschaften erklärt: (i) seine Aktualisierungsregel erzeugt durchweg ein isotroperes Spektrum der Singulärwerte als Adam; und infolgedessen (ii) optimiert es auf schwergewichtigen Daten die Tail-Klassen effektiver als Adam. Über empirische Belege hinaus bestätigen wir diese Erkenntnisse theoretisch, indem wir ein einlagiges Modell des assoziativen Gedächtnisses unter klassenungleichen Daten analysieren. Wir beweisen, dass Muon unabhängig von Feature-Embeddings durchweg ein ausgewogenes Lernen über die Klassen hinweg erreicht, während Adam je nach den Eigenschaften der Embeddings große Disparitäten in den Lernfehlern induzieren kann. Zusammenfassend offenbaren unsere empirischen Beobachtungen und theoretischen Analysen den Kernvorteil von Muon: seine Aktualisierungsregel stimmt mit der äußeren Produktstruktur linearer assoziativer Gedächtnisse überein, was ein ausgewogeneres und effektiveres Lernen von Tail-Klassen in schwergewichtigen Verteilungen ermöglicht als Adam.

English

The Muon optimizer is consistently faster than Adam in training Large Language Models (LLMs), yet the mechanism underlying its success remains unclear. This paper demystifies this mechanism through the lens of associative memory. By ablating the transformer components optimized by Muon, we reveal that the associative memory parameters of LLMs, namely the Value and Output (VO) attention weights and Feed-Forward Networks (FFNs), are the primary contributors to Muon's superiority. Motivated by this associative memory view, we then explain Muon's superiority on real-world corpora, which are intrinsically heavy-tailed: a few classes (tail classes) appear far less frequently than others. The superiority is explained through two key properties: (i) its update rule consistently yields a more isotropic singular spectrum than Adam; and as a result, (ii) on heavy-tailed data, it optimizes tail classes more effectively than Adam. Beyond empirical evidence, we theoretically confirm these findings by analyzing a one-layer associative memory model under class-imbalanced data. We prove that Muon consistently achieves balanced learning across classes regardless of feature embeddings, whereas Adam can induce large disparities in learning errors depending on embedding properties. In summary, our empirical observations and theoretical analyses reveal Muon's core advantage: its update rule aligns with the outer-product structure of linear associative memories, enabling more balanced and effective learning of tail classes in heavy-tailed distributions than Adam.

Muon übertrifft Adam beim Lernen von assoziativen Speichern im Tail-End-Bereich

Muon Outperforms Adam in Tail-End Associative Memory Learning

papers.abstract

Support