ミューオンは末尾連想記憶学習においてAdamを上回る性能を示す

要旨

Muonオプティマイザは、大規模言語モデル（LLMs）の訓練においてAdamよりも一貫して高速であるが、その成功の背後にあるメカニズムは未だ不明である。本論文は、連想記憶の観点からこのメカニズムを解明する。Muonによって最適化されるトランスフォーマーコンポーネントを除去することで、LLMsの連想記憶パラメータ、すなわちValueおよびOutput（VO）アテンション重みとフィードフォワードネットワーク（FFNs）が、Muonの優位性の主要な要因であることを明らかにする。この連想記憶の視点に基づき、Muonの現実世界のコーパスにおける優位性を説明する。これらのコーパスは本質的に裾が重い分布を持ち、少数のクラス（裾クラス）が他のクラスよりもはるかに少ない頻度で出現する。Muonの優位性は、以下の二つの主要な特性によって説明される：(i) その更新ルールは、Adamよりも一貫して等方的な特異値スペクトルを生成する；(ii) その結果、裾が重いデータにおいて、裾クラスをAdamよりも効果的に最適化する。実験的証拠に加えて、クラス不均衡データ下での一層連想記憶モデルの分析を通じて、これらの知見を理論的に確認する。Muonは、特徴埋め込みに関わらずクラス間で均衡の取れた学習を一貫して達成するのに対し、Adamは埋め込み特性に依存して学習誤差に大きな差を生じさせる可能性があることを証明する。要約すると、我々の実験的観察と理論的分析は、Muonの核心的な利点を明らかにする：その更新ルールは線形連想記憶の外積構造と整合し、裾が重い分布における裾クラスのより均衡の取れた効果的な学習を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.

ミューオンは末尾連想記憶学習においてAdamを上回る性能を示す

Muon Outperforms Adam in Tail-End Associative Memory Learning

要旨

Support