Muon Outperforms Adam in Tail-End Associative Memory Learning
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.
68dd3d8cd80e1286a364fdb3
Shuche Wang
68dd3d8cd80e1286a364fdb4
Fengzhuo Zhang
68dd3d8cd80e1286a364fdb5
Jiaxiang Li
68dd3d8cd80e1286a364fdb6
Cunxiao Du
68dd3d8cd80e1286a364fdb7
Chao Du
68dd3d8cd80e1286a364fdb8
Tianyu Pang
68dd3d8cd80e1286a364fdb9
Zhuoran Yang
claimed_verified
2025-10-20T12:14:28.576Z
646d769cda8e99940b71928e
/avatars/acee495a23362aa39b3d3e75c9afd967.svg
Zhuoran Yang
user
zhuoranyang
68dd3d8cd80e1286a364fdba
Mingyi Hong
68dd3d8cd80e1286a364fdbb
Vincent Y. F. Tan
Muon Outperforms Adam in Tail-End Associative Memory Learning
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 optimizer outperforms Adam in training LLMs by effectively optimizing associative memory parameters and balancing learning across classes in heavy-tailed data.
Muon optimizer Adam Large Language Models (LLMs) associative memory transformer components Value and Output (VO) attention weights Feed-Forward Networks (FFNs) heavy-tailed data singular spectrum tail classes class-imbalanced data balanced learning outer-product structure linear associative memories