Muon Outperforms Adam in Tail-End Associative Memory Learning
Authors: Shuche Wang
Fengzhuo Zhang
Jiaxiang Li
Cunxiao Du
Chao Du
Tianyu Pang
Zhuoran Yang
Mingyi Hong
Vincent Y. F. Tan
Venue: ICLR 2026 Poster (2026)
Keywords: Transformers
Muon
Optimization
Primary area: optimization

Abstract: 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 'head' classes are extremely frequent, while a vast number of 'tail' classes are individually rare. 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.

Published as a conference paper at ICLR 2026
MUON OUTPERFORMS ADAM IN TAIL-END ASSOCIATIVE MEMORY LEARNING
Shuche Wang
Fengzhuo Zhang!"!
Jiaxiang Li Cunxiao Du Chao Du Mingyi Hong² Vincent Y. F. Tan'
Tianyu Pang* 3
Zhuoran Yang'
National University of Singapore
University of Minnesota
Sea AI Lab Yale University
shuche, wanggu. nus edu fzzhang@u. nus edu basonljx3 60 1m@@mmmll... com cnsdunmégmail com (duchao, t tianyuaang
shuoran.yangäyan yang@yale..dd mhong@umn. edu vtanênus vtan@nus.edu.sg edu.
ABSTRACT
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 as- sociative memory parameters of LLMs, namely the Value and Output (VO) atten- tion weights and Feed-Forward Networks (FFNs), are the primary contributors to Muon's 's superiority. Motivated by this associative memory view, then explain Muon's superiority on real-world corpora, which are intrinsically heavy-tailed: a few 'head" classes extremely frequent, while vast number of 'tail' classes individually rare. The superiority is explained through two key properties: (i) its update rule consistently yields 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 an- alyzing a one-layer associative memory model under class-imbalee imbalanced data. We prove that Muon consistently achieves balanced learning across classes regard- less of feature embeddings, whereas Adam can induce large disparities in learning errors depending on embedding properties. In summary, our empirical observa- tions and theoretical analyses reveal Muon's '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.
INTRODUCTION
The effectiveness of Adam (Kingma & Ba, 2015) across diverse training scenarios has made it one of the most widely used optimizers for neural networks, serving as a cornerstone the tremendous successes of Large Language Models (LLMs).. Building this foundation, Muon (Jordan et al., 2024) has emerged as a matrix-paramete optimizer designed to surpass Adam. Empirical studies show that Muon is nearly 2 times faster than Adam across wide range of model sizes and architectures (Liu et al., 2025; Jordan et al., 2024). Its key innovation is to replace the raw gradient with the sum of its normalized orthogonal factors, which can be interpreted performing steepest descent with respect to the spectral norm (Bernstein & Newhouse, 2024).
However, despite its empirical success, a rigorous understanding of why and how Muon outperforms Adam in transformers remains incomplete. In particular, the steepest gradient descent interpretation does not clarify why optimization with respect to the spectral norm, as in Muon, should outperform optimization with respect to the infinity norm (for vectors), as in Adam. Consequently, convergence analyses of Muon derived from this interpretation fail to account for its observed superiority over Adam (Li & Hong, 2025; Shen et al., 2025).
"Equal contribution.
Project Lead. Fengzhuo Zhang conducted this work during his internship Sea AI Lab.

---

Published as a conference paper at ICLR 2026
This paper takes the first step toward understanding the mechanisms underlying Muon's superiority over Adam in training LLMs. Specifically, we ask the following two questions:
1. Which transformer components benefit most from Muon's 's matrix-nomm-bart based optimization compared to Adam?
2. What structural features of the transformer allow Muon to optimize these com- ponents more effectively?
To answer the first question, apply Muon to different transformer components. Our experiments consistently show that the more rapid convergence of the validation loss using the Muon optimizer compared to Adam is primarily due to the former's focus the value-output (VO) matrices of the attention mechanism and the Feed-Forward Networks (FFN) blocks. This leads to our first key insight: VO and FFN blocks, which serve as the primary associative memory stores in the model (Geva et al., 2020; Bietti et al., 2023), are the main beneficiaries of Muon's 's optimization strategy.
Building on this, we address the second question linking Muon's update mechanism to the learning dynamics of associative memory. Prior work suggests that the behavior of these memory compo- nents can be modeled as a sum of outer products representing stored facts (Meng et al., 2022a). Since Muon's 's update assigns equal update magnitudes to each outer product of the gradient corre- sponding to orthogonal singular directions, we hypothesize that optimizes associative memories more effectively than Adam because: (i) Muon's spectral normalization procedure balances the rates of learning of these outer products. (ii) Thus, when training on heavy-tailed data (i.e., where a few 'head' classes are extremely frequent, while vast number of 'tail' classes are individually rare), Muon reduces the dominance of frequent (head) facts and enables more effective learning from infrequent (tail) facts compared to Adam.
We validate these hypotheses through a combination of empirical analysis and theoretical modeling. Empirically, we conduct two experiments. First, measure the singular value spectra of the weight matrices and show that Muon consistently yields more isotropic representations than Adam, indi- cating that its normalization prevents spectral energy from concentrating in dominant components. Second, we evaluate the performance of both optimizers on a knowledge-intensive. heavy-tailed task to demonstrate the practical benefit of Muon's more balanced updates: while both optimizers perform well on head classes (frequent in training data), Muon outperforms Adam on tail classes (rare training data), leading to more stable and uniform convergence.
Theoretically, we focus on a one-layer linear associative memory model to rigorously explain these empirical findings. Under class imbalance in the training data, mimicking heavy-tailed distribution, we show that Muon maintains balanced learning across classes, regardless of the feature embeddings.
contrast, we prove that Adam's performance is unstable and strongly dependent on the embedding structure, which can lead to large disparities in learning error across classes. By closely examining the parameter updates, we find that the singular spectrum of weight matrices trained by Muon is nearly isotropic, whereas Adam's 's is uneven.
Summarizing the empirical and thcoretical findings, we identify a clear mechanism underlying Muon's 'ss superiority: The Muon update rule is aligned with the outer-product structure of linear associative memories, enabling more balanced and effective learning of tail classes in heavy- tailed distributions compared with Adam.
2 PRELIMINARIES
Muon (Jordan et al., 2024) is an optimizer tailored for matrix parameters that replaces the raw (or momentum) gradient with the sum of its normalized orthogonal factors, producing scale-invariant, norm-controlled update direction. For weight matrix W E R at step let G2 VwL(W, denote its gradient. Muon maintains a momentum accumulator of gradients as B, = μB,,1 1 + G, with Bo O, and μ [0. 1). At each step, Muon computes the Singular Value Decomposition (SVD) of B, as B, = V. with U, E mxrs V E Raxr,, where ri = rank (B,), and form the nearest (eemi)-ortoogee
matrix O U.V.". Then Muon updates the parameter as W4+ = W, Th O,. In practice one can approximate O using fixed number (e.g., 5) of Newton- Schulz iterations applied to (B, B,)-1/2 11/2, which avoids the full SVD while preserving the seale normalization effect. Detailed introduction of Muon is in the related works section (Appendix C).
2

---

Published as a conference paper at ICLR 2026
Transformers serve as the backbone of LLMs, predicting the probability of the next token given sequence of N tokens. A sequence of N tokens is embedded into matrix X E [R The first layer takes X (9) as the input, and each subsequent layer takes the previous layer's output as its input. Every layer l E [L] processes its input through two sequential components: an attention module and a FFN module. The attention module computes
H
H = x((11)
V.A with wil (((-)) sm(x((-1), T With w
K.A q. (()x)) (2.1)
where sm(-) is the column-wise softmax operator, H is the number of attention heads, W w(e Q.h: K.A E Redexd capture token relationships, and W E Rd.xd w E Rdddd apply linear transformations. The ferd-forward module then updates the representation as
x0 H (E(H)(), wo w H( + will o (w
(2.2)
where (-) is the element-wise activation function, and ( parameters. gated variant replaces the standard form with
E xd willl E Rdxdy are learnable
in
E fgae (H,), W w-(f) out gate = will (a(W! H()) o (wll) H)))
gate
where ④ is the Hadamard product, and w gate E Rd xd is an additional mapping, After L layers, the final hidden state of the last token, x(4), is projected by the language model head Enead E I2 xd to produce logits E (L), which has a vocabulary of size of K.
Associative memory refers to architectures that store and retrieve patterns based on learned asso- ciations between inputs and outputs. Recent research has examined linear associative memory in LLMs. Specifically, consider triplet (s, r, o). where is the subject, the relation, and the object (e.g., 8 Thhe United Nations headquarters". r ="is located in", =NNew York City"). A linear associative memory W maps key vector e, encoding rs) value vector e, encoding o, such that 6,- = We, holds for all possible (s, o). Under the orthogonality of embeddings e, and co. W can be expressed W = Σ; eo, eg,, where the summation taken over the indexes of facts. These facts naturally emerge in the token association in the pretraining data, e.g., the coappearance of ""paceX"" and "Elon Musk", and are learned by LLMs in the form of associative memories. Prior work has investigated associative memory in both attention and FFN modules. In the attention module, Bietti et al. (2023) showed that the parameter Wo can serve as a linear associative mem- ory when Wy is fixed. Since Wo and Wv play symmetric roles, we also treat Wy part of the associative memory parameters. It is therefore natural to consider VO jointly: several works (Lin et al., 2024; Wang et al., 2025) have shown that the value and output matrices play similar roles and can be analyzed together in practice, even in multi-query attention (MQA) and grouped-query attention (GQA) settings. In FFN, works knowledge editing (Geva al., 2020; Dai et al., 2021; Meng et al., 2022a;b) have identified the module as functioning as an associative memory, which can be well approximated by linear associative memory models. In fact, they demonstrate that we can manually update the knowledge in Large Language Models (LLM)s using least squares on the FFN parameters (Meng et al., 2022a;b; Fang et al., 2024). Thus, throughout this paper, we refer to Wo. Wv, and FFN in LLMs as the associative memory parameters.
3
MAIN RESULTS
3.1 ASSOCLATIVE MEMORIES ARE MAIN BENEFICIARIES OF MUON
In this section, we identify the transformer components that benefit most from Muon by measuring validation loss on the Wib dataset using 160M NanoGPT model. We adopt two-stage protocol. First, in the "Independent Blocks" setting, we apply Muon to a single block at a time while keeping all other blocks on Adam, covering the attention projections Wo, WK. Wv, Wo and the feed-forward matrices Wmn Wout Second, in the "Combined Configurations" setting, we apply Muon to the most impactful subsets identified in the first stage examine whether partial application can recover the performance gains of full Muon. As introduced in Section 2, we evaluate both gated and non-gated FFN variants of NanOGPT. The experimental details are in Appendix F.
3

---

Published as a conference paper at ICLR 2026
MacsiAll
ManiAlIAT AT.. FFN)
Maaiig Am Atn, FTN)
Xaxo/VD Arm, PwN.eA/b/OK ArD)
Masssjv Amniacctaaa Art, 1130
Maan/VU Ar.. Wi.ltAdaing Aln. Wya?
6.8
Maaaj Alt Aus. FFN0
6.-
Mare/Vn Are. W.AAAoo
Validatioon Los
MacatO Anal ManntK
FFN
Validetio Lown
MauuVV Maan/80 And
5.5
Maust (4 Mkke(T AIs, Weed
MaootO Ars. FPN Alm. Km))
Macai W. MlAdartA Ats. W.
5.0
4.0
4.0
3.5
3.5
2000
4000
6000
10000
2100
4000
6000
30000
Training Steps
Training Steps
(a) Independent blocks: Val loss over training
(b) Combined configurations: Val loss over training
4.2
4.0
3.031
40-
Vilidaatoon Ea0s
1,806
Validaation Laws
3.8
-1354
3.771
3,307
3.50
1,47%
3,67
0.6
1.565
0,6
3.583
1/465
3,404
3.4
34-
32
Meaice Amap
AlmerPal
32
Unww aae mi
Usw del Ril
Alsurio All Road
Aade mw
PAme me
Mome
Mangly
(c) Independent blocks: Val loss at step 10,000
(d) Combined configurations: Val loss at step 10,000
Figure 1: Validation loss comparison on the 160M NanoGPT model with non-gated FFN under different Muon/Adam assignments. Panels (a) and (b) show the validation loss over training steps for the Independent Blocks and Combined Configurations settings, respectively. Panels (c) and (d) report the corresponding validation loss at step 10,000 for each mode, summarizing the final performance of the Independent Blocks and Combined Configurations.
Figure 1 presents our results. We first examine the independent-block experiments for attention. From Figures ((a) and (c)), the VO weights Wy, Wo (Muon on VO / Adam on QK and FFN) show substantially larger gains under Muon than the QK weights Wq. Wx (Muon on QK/Adan on VO and FFN). Notably, applying Muon to only Wv or only Wo already yields much larger gains than applying it to QK. For the FFN, we find that W W gate and W... all benefit from Muon, with Wod. yielding stronger improvements than W.n- As we show in Appendix G.2 these trends persist even after controlling for parameter count.
After identifying the importance of each module, the combined configurations aim to quantify their contributions to the full Muon. Guided by the independent-block findings, we first observe that VO+FFN already closely tracks-and in our runs nearly recovers-the full-Muon trajectory in Fig- ure 1(b). This indicates that applying Muon QK contributes little its overall performance. Importantly, this effect is not due to the logit explosion reported by Team et al. (2025) in large Mixture of Experts (MoE) models; logit values for our setting do not explode, as reported in Ap- pendix G. The small remaining gap between full Muon and VO+FFN may stem from the fact that VO+FFN adopts the same learning rate as full Muon without further tuning.
To isolate the contributions of Wo and Wy within VO+FFN, we perform ablations starting from the VO+FFN setting: we keep Muon on FFN and on only one of Wo or Wr, reverting the other to Adam (i.e., V+FFN and O+FFN). Both ablations degrade performance, with the V+FFN variant dropping more indicating that Wo is more influential than Wy. Overall, applying Muon to VO+FFN is critical for recovering full-Muon performance. The same qualitative patterns hold for the gated FFN variant reported in Appendix G.3, and are further confirmed on a larger 0.7B model in Appendix G.4, demonstrating the robustness of our findings.
4

---

Published as a conference paper at ICLR 2026
Observation Muon is most effective when applied VO and FFN; in particular, applying Muon to only VO+FFN almost recovers the full-Muon trajectory.
Remark 3.1. We emphasize that this observation is not trivial consequence of parameter counting; although QK and VO are equal in size, VO proves substantially more influential.
As introduced Section 2, prior works discover that the common role of VO and FFN is that they both serve as the associative memories for transformers, which store facts and knowledge. Furthermore, Bietti et al. (2023) and Meng et al. (2022a) show that the linear associative memories well approximate them. Specifically, for a set of facts represented by key-value pairs {(eg, Co:)
the memory matrix W can be constructed as sum of outer products, i... W Ei c,, egp, where the summation is taken over the index i of K facts. To make this more concrete, consider toy example with two orthogonal facts in R2:
Fact I: ("the capital of France") Es, [1, [1,0]T, ("Paris") co.
Fact 2: ("the capital of Italy") C2z = [0, ("Rome") coz = [0, 1]T.
The resulting memory matrix is W Co, tog e. 12.2 which correctly stores these facts since We,, = Co, for i = 1.2.
Learning linear associative memories is particularly well-suited to Muon's update mechanism. Con- cretely, the gradient G E Rdddd of the loss with respect to the linear associative memory weight W can be expressed as sum of outer products via SVD as G = USV Ei= s,u,"," Muon computes its update (without momentum) by normalizing away the singular values, forming the orthogonal factor O = UVT = Ei= u.u.T Following the toy example, consider training the memory parameter W with & loss, i.e, ei e, - We,, + c2lleo2 We,, where C₁ > represent the importance frequency of each fact the current training batch. The corresponding gradient is G = G1 Co, e + 6oz = diag(ci,c)) factor becomes O = UV 122 Cor C,, + Cog szy which is simply the sum of the constituent facts' outer products. Crucially, the update Ó assigns equal weight to both Fact I and Fact 2, regard- less of their original coefficients e₁ and C2 in the gradient. This illustrates how Muon normalizes the updates across orthogonal facts, allowing it to learn both frequent (large C1 ) and infrequent (small (g) facts uniformly. Comparing this with the linear associative memory Σ Po. we see that Muon updates all "orthogonal" facts at the same rate. Later, we will see that the singular values S of the gradient G' encode the frequencies of knowledge in the training data under cross-entropy loss in Sections 3.3 and 4. By normalizing away S to form its update, Muon can therefore learn both frequent and infrequent facts more uniformly than wradient-magniuudebane based optimizers like Adam.
62). Consequently, Muon's 's normalized update
We verify this insight from two perspectives. First, from the view of weight spectra, the weight ma- trices learned with Muon exhibit more isotropic singular-value spectrum than those learned with Adam, indicating that knowledge, regardless of its frequency, is represented with comparable mag- nitude. Second, at the level of overall knowledge acquisition, Muon yields more balanced learning across entities and frequencies (head and tail) than Adam. We examine these two consequences in the following sections.
3.2 MUON CONSISTENTLY LEARNS MORE ISOTROPIC WEIGHTS THAN ADAM
To validate that Muon can shape the weight matrices more evenly across directions, we conducted a spectral analysis of them. For weight matrix with non-zero singular values o (Ø1, .2,...,07 ox), we define the normalized singular energy distribution nent q; is qi of? EØ= o3. This distribution represents the fraction of energy captured by each corresponding singular vector. Based on this, we introduce several metrics to characterize the isotropy of the spectrum: normalized SVD entropy defined as Huem
(41, ....... 9n), where each compo-
1 Σ log 9 exp (- Ei=1 qi log q:). Top-k energy fraction defined as (a) = -' of Ey o3. and eigenvalue quantile ratio defined as (/((oP))(Q((o?") Detailed explanations of these metrics are in Appendix F.2. Intuitively, more isotropic weights correspond to larger values of normalized SVD entropy and effective rank, and smaller Top-k energy fraction and eigenvalue quantile ratio.
log
effective rank defined as eRank(o)
5

---

Published as a conference paper at ICLR 2026
ebjny
attory L7)
Tepitt:12
Te
Manm Ith
MardAJAts
270
3BC
2000
2P8
0001
7000
181
330
306
1601
200
62 106
Treaing Says
18"
075035 Rude
200
226
3001
1200e
10
2000
20:0
1000
2003
3000
282
390
32.0
1000
2001
120
(a) VO(Non- gated FFN)
(b) (NonNNoggateed FFN)
Figure 2: Spectral Dynamics of Transformer Weight Matrices During Training. Each panel reports four metrics characterizing singular value distributions: SVD entropy, Top10E, cRank, and Q75/Q25 ratio. The four subplots correspond to different weight matrix groups: (a) VO and (b) Wout-
The spectral analysis in Figure 2, focusing on the key associative memory components from Ob- servation 1, shows that Muon systematically reshapes the learned weight matrices relative to Adam. The results, averaged over 10 random seeds, demonstrate that: (i) Muon produces much more isotropic singular spectrum than Adam from the start of training, whereas Adam's 's isotropy fluctu- ates significantly over the course of optimization. (ii) The isotropy of Muon stable across random initializations, indicated by the negligible error bars in Figure 2, while Adam is more sensitive to initialization. These findings suggest that Muon consistently promotes richer and more diverse features in the model's most critical memory components, conclusion we summarize below. The results for the gated FFN architecture and other weights are in Appendix G.3 and G.5, respectively.
Muon consistently yields more isotropic weight matrices with broadly distributed spectral energy than Adam, both throughout training and across random initializations, thereby sup- porting richer feature representations.
Observation 2:
Empirically, we also find that Muon learns more isotropic QK weights than Adam. However, as discussed in Section 3.1, QK weights are not part of the linear associative memory mechanism and are therefore not expected to benefit from the isotropic property of the weight matrices.
Our results differ fundamentally from the spectral analysis in Liu et al. (2025) for three reasons. First, we decompose the parameters according to associative memories, whereas Liu et al. (2025) aggregates them, obscuring the essential components driving Muon's behavior. Second, we investigate the instability of Adam under random initialization (i.c., random seeds), which further establish theoretically in Section 4. Finally, our analysis focuses on dense architectures, while Liu et al. (2025) centers Mfxture-of-Expertt (MoE) models.
3.3
MUON ACQUIRES KNOWLEDGE MORE EVENLY COMPARED TO ADAM
Our previous findings indicate that the Muon optimizer is particularly important for the associative memory components of the model, where learns more isotropic weights. To examine the overall effects of learning associative memories, we turn to a knowledge-intensinn question question-answering answering (QA) task. The task is based on a synthetic QA dataset containing biographical information (e-g.. name, birthday, and company) for over 200,000 individuals (Allen-Zhu & Li, 2024). To capture the heavy- tailed nature of real-world knowledge, we control the frequency each individual's appearance in the training set that it follows a power-law distribution (Figure 3(a)), thereby inducing varying levels of difficulty in learning knowledge about different individuals. A 160M NanoGPT model is trained answer questions about this biographical information. The performance is evaluated via the First Token Accuracy (FTA) on the answers, following Allen-ZIe -Zuu & Li (2024). Further details on the dataset are provided in Appendix F.3. We include SGD baseline for Adam and Muon.

---

Published as a conference paper at ICLR 2026
Glaa Groep
1.0
baee
+ Sooples 167
Accunic (F (FTA)
Cod
3.6
Lecug
treep
0.g
0.0
2305
Skeps
(a) Sample/class
Muon
(c) Adam
1.0
Acuuasy (FTA)
Accuaac (PTA)
0.8
6.5
3.6
0.
02
68
2000
$700
<020
5000
106re
2010
4266
5000
3000
10000
2300
4000
0000
8900
1000
Traiting Sieps
(d) SGD+MMmmntum
(c) Muon VO.FFN/O.FF/)
(f) Muoon(QK/AAdmm(VO.FF))
Figure Performance comparison of different optimizers for transformers with non-gated FFN on heavy-tailed knowledge task. (a) Sample distribution per class, following power law. (b-d) Performance of Muon, Adam, and SGD+Momentun, (e) Muon applied to VO and FFN, with Adam on QK. (f) Muon applied to QK, with Adam on VO and FFN.
The results in Figure 3 lead to unequivocal conclusion about the efficacy of different optimizers under data imbalance. high-frequency (head) classes, all optimizers perform well, with Muon, Adam, and even SGD+Momentum rapidly reaching near-perfect accuracy (Figure 3(b-d)). Consis- tent with prior work on heavy-tailed distributions (Kunstner et al., 2024). Adam maintains clear advantage over SGD, which struggles with tail classes. Our key finding, however, is that Muon sub- stantially outperforms Adam on low-frequency (tail) data, achieving faster and more uniform con- vergence across all frequencies. Moreover, the consistently tighter error bars for Muon-especiial relative to Adam-refledt lower variance and a more stable learning process.
Furthermore, the hybrid configurations in Figure 3(e-f) clarify where Muon matters most. Applying Muon to VO+FFN (with QK on Adam) yields strong gains on rare classes and markedly reduces the head-tail gap, whereas applying Muon only QK (with VO+FFN on Adam) yields only limited improvement. This mirrors Observation 1: VO+FFN is the most effective target set, as it concen- trates the model's associative memory. Results for the gated FFN, which show the same pattern, are provided in Appendix G.7. Additional experiments in Appendix G.8 vary the degree of fact imbal- ance, and show that the average FTA gap between Muon and Adam shrinks as the data distribution becomes more uniform. Together with the Wikitext103 results in Appendix G.9, which exhibit the same qualitative behavior on a standard language modeling benchmark, these findings further sup- port the view that Muon's 's advantage is tightly linked heavy-tailed imbalance. We summarize these findings as Observation 3.
Observation 3:
In heave- chowledge-intensiee tasks, Muon matches Adam's strong perfor- mance the head classes while substantially improving learning on tail classes, narrowing the head-tail gap and accelerating convergence.
addition to the knowledge acquisition task, whose success primarily depends on learning associative-memndt parameters (VO and FFN), also evaluate in-context linear regression task in Appendix 1.10, which primarily depends on learning the QK parameters. In contrast to the above observation, Muon achieves performance the tail class similar that of Adam in this task. This consistent with Ohservation 1, which indicates that the QK parameters are not the main source of Muon's superiority.

---

Published as a conference paper at ICLR 2026
3-3
142
3+2
103
104
104
Sig30 Capped
364
San Caplee
egattt Draaaan
1d4
364
Digple
Shars Eiagtl
Xde
Pepolatco
Prpciatee
(a) Average Angles Between E,/E, (b) One-step Optimization Results (c) Multi-step Optimization Results
Figure 4: (a) Average angles between E, or F. in FFN at layers 5, 10, 15, 20, 25 of Llama3- 8b-instruct. (b) Results of one-step GD, SignGD, and Muon with both coupled and decoupled embeddings. For GD, the outcomes under the two embedding types coincide. We use different step sizes to obtain different levels of population loss and (W) Results of multi-step GD, SignGD, and Muon with both coupled and decoupled embeddings. We use different numbers of optimization steps to obtain different levels of population loss and AWW).
4 CASE STUDY OF ONE-LAYER MODELS
We now analyze three optimizers complement the preceding empirical observations. We first introduce an abstraction that captures their key dynamics and then present both empirical and theoretical results. As shown in Eqns. (2.1) and (2.2). structural property of associative memory parameters is that their output is added directly to the hidden states, which are subsequently processed by the language model head. Motivated by this property, our abstraction retains the associative memory and language model head, while replacing all preceding modules with given feature embeddings.
Muon, and Gradient Descent (GD) (as a baseline)
Consider triplets {(s; Ti, o;) K (introduced in Section 2), where subject-relation pairs (s,, r:) and objects o; are embedded into the columns of matrices E BRRx and E R.,.K respectively. A linear associative memory W Rd,xd, predicts the object for query Fgg with probabilities (Ex) sm(E) Sn(EWWE) R.K. The objective to minimize the population cross-entropy loss L(W)= = E Pk log[sw(( (F&)) where pI: is the frequency or probability of the k-th triplet. We analyze three optimizers initialized W, = 0, all simplified by disabling momentum for clarity. (i) (ii) Adam: Following prior work (Kunstner et al., 2024; Bernstein & Newhouse, 2024), set B3 32 reducing to SignGD: W W₁ W sign( (V W C(W.)) (iii) Muon: The update is We+1 We - nU.VV" where U,E,V,T is the SVD of V L(W)). This simplified form, U.V.T. the projection of the gradient onto the nearest orthogonal matrix. We then state the assumptions for our results.
GD: W = W. V
Assumption 4.1. The embeddings E and are orthonormal, i.e., E E IK.K.
The unit-norm requirement rules feature-level imbalance, which would otherwise couple with the imbalance induced by and complicate the analysis. Our techniques can be directly applied even without this unit-norm requirement. The orthogonality assumption is intuitively plausible, as different concepts are independent and do not influence one another. We empirically verify this on Llama3-Sb-instrud (Dubey al., 2024). Following Fang al. (2024), extract E; and E, in FFN across layers for 3, 000 knowledge items Counterfact (Meng et al., 2022a) and compute average angles between them (see Appendix F.4 for details). As shown in Figure 4(a), these angles are near 90", confirming approximate orthogonality. For K independent concepts, orthogonality requires d,, d, 2 K. For simplicity, set d, d, K in what follows.
Assumption 4.2. The first L triplets share the same probability and together contribute a total mass of a, i.e., Pk a/L L for is E [L]. The remaining triplets also share the same probability and together contribute a total mass of a, i.e., P = (1 a)/(K - L) for
L.
This assumption states that the data imbalance is between two classes among the K triplets Defining = K, the ratio a/3 quantifies the degree of balance: if a > 3, the first L triplets appear more frequently during learning, and vice versa. This simplified two-class setting is sufficient to capture
8

---

Published as a conference paper at ICLR 2026
the primary differences between optimizers; the multi-class case follows directly from proof by extending the SVD calculation.
Throughout Section 4 we will also refer to the imbalance ratio, defined as the ratio between the minse(
minimal and maximal frequencies of triplets, i.e., r = ma E (0. 1]. Under Assumption 4.2
with parameters a and 3 = L/K, this reduces to r( = min a(1-1) ([1-a)' a(1-8) J We keep the
2(1-a)
two-mass (o.,3) parametrization because allows us to write the gradient and its SVD in closed form while capturing the same dependence on class imbalance as using directly; the multiclass case follows from the same SVD calculation.
4.1 EXPERIMENTAL RESULTS
Under Assumptions 4.1 and 4.2, we evaluate GD, SignGD, and Muon for a = 0.8, 3 = 0.2, considering two embeddings for E and E: (i) support-decoupled: the supports (indices of non- zero entries) of different E; or F, are disjoint; (ii) support-coupled: supports may overlap. We study two optimization protocols, initializing Wo = xd,xa. (i) one-step: take a single update with a scaled step size to obtain range of L(W) values; (ii) multi-step: run multiple updates to reduce L(W), varying the number steps. Experimental details are in Appendix F.5. To quantify learning imbalance across K knowledge items, examine the relationship between population loss C(W) and maximal probability gap A (W) == jax,,j{(K [fw (E,))- [E where [fw (E,)): denotes the probability assigned to the correct item i. A larger AWW indicates greater imbalance.
Across both optimization-step protocols and embeddings (Figures 4(b), 4(c)), we observe that (i) For all optimizers, A(W) first increases and then decreases as L(W) decreases. Early in training, when correct probabilities are near 0, imbalance is pronounced; later, when all items are well learned (e.g,, probabilities 2 0.9), imbalance diminishes. (ii) For both embedding regimes, GD and Muon behave consistently: GD exhibits substantial imbalance, whereas Muon remains much more balanced across items. (iii) SignGD also demonstrates unstable behavior; its imbalance resembles GD in the and Muon in the decoupled embedding case.
coupled embedding
Because one-step and multi-step experiments align qualitatively, we first analyze the one-step set- ting for clarity. This simplification is common in theoretical studies of neural network dynamics (Ba et al.. 2022: Dandi et al., 2023), and our techniques extend directly-albei with more algebra-to the multi-step case. As demonstration, Theorem 4.4 provides a multi-step analysis of Muon.
4.2 THEORETICAL RESULTS
For each optimizer, we choose a step size 1/ so that some class already attains correct-class proba- bility at least after one update, and then we report the smallest correct-class probability across classes at the same 17. Equation 4.1 formalizes this procedure.
Eóp = inf min [Sw, (Ex)la n0 KE(K)
max [fw, (En)]a 2 e, W,, -Wo 17 (W)) (4.1) KE(K)
where opt E {GD, SignGD, Muon} and (Wo) denotes the parameter update of optimizer opt at Wo: and W,, denotes the parameter obtained after one step of optimizer opt with step size 1) starting from Wo. i.e., W,, Wo 17 Go (Wo). Specifically, GGD (Wo) V ((W)). GSigaG (Wo) sign(Vw ((W)), and Guuon MamnnWW (Wo) Ugnorm(Eo) Vo where Vo is the SVD of V (WW)- Note that 2xpt E [0, - E] and ((W) are related as A(W) = 1 - Copt 2 O. When oor ~ 1-c, opt achieves balanced learning across facts; in contrast, when espt O, imbalanced learning ensues.
Theorem 4.3. Let r := min Pt | Pk (under Assumption 4.2, r = r(a, Assumptions 4.1 and 4.2 hold, with fixed a, such that * B, and K goes infinity, obtain the following results for one-step GD, Muon, and Adam.
GD: For any E and E satistifying Assumption 4.1, we have
Φep = ( where r(a, B) maxç ps:
min& Pk min
aa(1- - 3) 8(1 [1))
<1
[[1 - aa) a(1 3)

---

Published as a conference paper at ICLR 2026
Muon: For any E and E satistifying Assumption 4.1, we have
1 ( REIKKT)) PMuon 1 11- -1-e(1+0) log K )) and GMuon (Wo) -EE t0
where JK.K K E RR is the matrix with all elements equal 1. The big-O notation for matrices means that for A = OB), each entry satisfies As O(Bi)) for all i,
Adam: There exist E and E satisfying Assumption 4. such that PSignGD 2 1 1-c. There also exist E' and E' satisfying Assumption 4.1 such that
Omin (iignGD ( (Wo)
SignGD = O(E K K-03),, and mmax SignGD (Wo)) 3 25%,
where max and min are the largest and smallest singular values, respectively.
Interpretation Theorem 4.3. These theoretical results align with Observations 2 and 3, and Figures 4(b) and 4(c): Muon maintains balanced learning with near-isotropic updates, GD is highly sensitive to data imbalance, and Adam varies widely across embeddings.
the one-step update, when the maximum correct-class probability across items is at least 1-c, the item with the minimum correct-class probability satisfies: (i) Muon: 2 1 11-el e(1 + O( Jug K )). which indicates learning is essentially balanced across items with a near-isotropic update (singular values nearly equal); (ii) GD: O (((rr() (r(a.))-1), which is strongly controlled by data imbalance via r(a, (8) (balanced when = 1, severe imbalance when < 1); (iii) Adam: embedding dependent; it can match Muon with disjoint supports (e.g., E = E IK,K), achieving 1 but drop Occ-0.7 K K-0.3) with overlap; its update may exhibit pronounced spectral decay (amin/ / & 25%), unlike the near- uniform singular values of Muon. A detailed discussion of Theorem 4.3 is provided in Appendix E.
In the following, we extend our techniques of one-step analysis to the multi-step analysis of Muon. Parallel Eqn. (4.1), we define the infimum correct-class probability for the multi-step optimizer as lop inf; [adf:[minee]]][ww (Ex) ) max ( (Ex)]a 2 1-6, where W. = W 1h: (W) Here, we assume that the learning rates (17 he21 are determined by a fixed schedule prior to optimization. Although the quantity implicitly depends on this schedule, we omit it from the notation for lign for brevity. We emphasize that different schedules may affect the value of / that attains the infimum in Copt but they do not influence the balance behavior that we present.
Theorem 4.4. If Assumptions 4.1 and 4.2 hold, then multi-step Muon achieves
1 11+0 log
log K )) (08KK) and GMuon (W,) -EEE -ET ( for any 12 O. K
Munnn 2 1 21-
The proof is provided in Appendix I. We note that the multi-step analysis of Muon shares similar characteristics as the one-step version in Theorem 4.3.
CONCLUSION
Our work takes the first step toward unveiling why and how Muon outperforms Adam. Through ablations of Muon's effect different Transformer components and by relating these results to the balanced learning of associative memories, we conclude that the Muon update rule is aligned with the outer-prode nnter-product structure of linear associative memories, enabling more balanced and effective learning of tail classes in heavy-tailed distributions. Intuitively, this property Muon may extend beyond outer products to higher-order tensor products, an exciting direction for future work.
ACKNOWLEDGMENTS
S. Wang and V. Y. F. Tan are supported by a Singapore Ministry of Education (MOE) AcRF Tier 2 grant under grant number A-8004062-00-000 M. Hong and Li are supported partly by NSF grant SCCS-2450988. NSF CIF-2414372, NSF PPCN-2311007
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Ruslan Abdulkadirov, Pavel Lyakhov, and Nikolay Nagornov. Survey of optimization algorithms in modern neural networks. Mathematics, 11(11))2466, 2023.
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 3.3, knowledge capacity scaling laws. arXiv preprint arXiv:2404.0544 2024.
Orly Alter, Patrick O Brown, and David Botstein. Singular value decomposition for genome-wide expression data processing and modeling. Proceedings of the National Academy of Sciences, 97
(18):10101-10106, 2000.
Kang An. Yuxing Liu, Rui Pan, Yi Ren, Shiqian Ma. Donald Goldfarb, and Tong Zhang. Asgo: Adaptive structured gradient optimization. arXiv preprint arXiv:2503.2076,, 2025.
Jimmy Ba. Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High- dimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems, 35:37932-37946, 2022.
Jeremy Bernstein and Laker Newhouse. Old optimizer, new norm: anthology. arXiv preprint arXiv:2409.20325, 2024.
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformer: A memory viewpoint. Advances in Neural Information Processing Systems, 36: 1560-1588, 2023.
Xiangyi Chen, Sijia Liu, Ruoyu Sun, and Mingyi Hong. On the convergence of a class of adam-type algorithms for non-convex optimization. International Conference on Learning Representa- tions, 2019. URL https:://// wpwpeeeeeeeeeeeeeeeeeeeeeeeeeeee wwwerreeeeeeeeeeeeeeeeeeeeeeeeee wwt./e////oo foorum (d=H1==x333 9t..
Damai Dai, Dong, Yaru Hao, Zhifang Sui, Baobao Chang, and Furu Wei. Knowledge neurons in pretrained transformers. arXiv preprint vrXiv:2104.08600 2021.
Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. How two-layer neural networks learn, one (giant) step at a time. arXiv preprint arXiv:2305. 18270, 2023.
Alexandre Défossez, Léon Bottou, Francis Bach, and Nicolas Usunier. A simple convergence proof of adam and adagrad. arXiv preprint #xXiv:2003.02399 2020.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang. Angela Fan, et al. The llama 3 herd of models. arXiv e-prints, pp. arXiv-2407, 2024.
Junfeng Fang, Houcheng Jiang, Kun Wang, Yunshan Ma, Shi Jie, Xiang Wang, Xiangnan He, and Tat-Seng Chua. Alphaedit: Null-space constrained knowledge editing for language models. arXiv preprint arXiv:2410. 02355, 2024.
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in neural information processing systems, 35:30583-30598 2022.
Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. arXiv preprint arXiv:20/2. 20..:333...4993,, 2020.
Ekaterina Grishina, Matvey Smirnov, and Maxim Rakhuba. Accelerating newton-schulz iteration for orthogonalization via chebyshev-type polynomials. arXiv preprint arXiv:2506.10933 2025.
John Hopfield. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the national academy of sciences, 7954))5554-2558 1982.
Keller Jordan, Yuchen Jin, Vlado Boza, You Jiacheng, Franz Cecista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. URL https://elllrjjrrdnn. github. io/posts/muon, 6, 2024.

---

Published as a conference paper at ICLR 2026
Aishwarya Kamath, Johan Ferret, Shreya Pathak, Nino Vieillard, Ramona Merhej. Sarah Perrin, Tatiana Matejovicova, Alexandre Ramé, Morgane Rivière, et al. Gemma 3 technical report. arXiv preprint drXiv:2503.19786 2025.
Gyu Yeol Kim and Min-hwan Oh. Convergence of muon with newton-schulz. arXiv preprint arXiv:2601. 191,, 2026.
Diederik Kingma and Jimmy Lei Ba. Adam: A method for stochastic gradient descent. In ICLR: international conference on learning representations, pp. 1-15, 2015.
Teuvo Kohonen. Correlation matrix memories. IEEE transactions on computers, 100(4):353-359,
2009.
Dmitry Kovalev. Understanding gradient orthogonalization for deep learning via non-euclidean trust-region optimization. arXiv preprint 0rXiv:2503.12644 2025.
Frederik Kunstner, Alan Milligan, Robin Yadav, Mark Schmidt, and Alberto Bietti. Heavy-tailed class imbalance and why adam outperforms gradient descent on language models. Advances in Neural Information Processing Systems, 37:30106-300 -01448, 2024.
Tim Kiz-Kit Lau, Long, and Weijie Su. Polargrad: A class of matrix-gradien optimizers from 2025.
unifying preconditioning perspective. arXiv preprint
Omer Levy, Minjoon Seo, Eunsol Choi, and Luke Zettlemoyer. Zero-shot relation extraction via reading comprehension. arXiv preprint arXiv:/706.044 04/ 15, 2017.
Haochuan Li, Alexander Rakhlin, and Ali Jadbabaie. Convergence of adam under relaxed assump- tions. Advances in Neural Information Prucessing Systems, 36:52166-52196, 2023.
Jiaxiang Li and Mingyi Hong. A note on the convergence of muon and further. arXiv e-prints, pp. arXiv-2502, 2025.
Chi-Heng Lin, Shangqian Gao, James Seale Smith, Abhishek Patel, Shikhar Tuli, Yilin Shen, Hongxia Jin, and Yen-Chang Hsu. Modegpt: Modular decomposition for large language model compression. arXiv preprint wwXiv:2408.09632. 2024.
Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, et al. Muon is scalable for Ilm training. arXiv preprint a5Xiv:2502. 16982, 2025.
Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt. Advances neural information processing systems, 35:17359-17372, 2022a.
Kevin Meng, Arnab Sen Sharma, Alex Andonian, Yonatan Belinkov, and David Bau. Mass-editing memory in a transformer. arXiv preprint arXiv:2210.07229, 2022b.
Eshaan Nichani, Jason D Lee, and Alberto Bietti. Understanding factual recall in transformers via associative memories. arXiv preprint arXiv:24/2.065538 2024.
Antonio Orvieto, Samuel Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pas- canu, and Soham De. Resurrecting recurrent neural networks for long sequences. In International Conference on Machine Learning, pp. 16670-26699 PMIR, 2023.
Yan Pan and Yuanzhi Li. Toward understanding why adam converges faster than sgd for transform- ers. arXiv preprini arXiv:2306.00204 2023.
Thomas Pethick, Wanyun Xie, Kimon Antonakopoulos, Zhenyu Zhu, Antonio Silveti-Falls, and Volkan Cevher. Training deep learning models with norm-constrained lmos. arXiv preprint arXiv:2502.07529, 2025.
Olivier Roy and Martin Vetterli. The effective rank: A measure of effeetive dimensionality. In 2007 15th European signal processing conference, pp. 606-610. IEEE, 2007.
Naoki Sato, Hiroki Naganuma, and Hideaki liduka. Analysis of muon's 's convergence and critical batch size. arXiv preprint arXiv:2507.01558, 2025.
12

---

Published as a conference paper at ICLR 2026
Ishaan Shah, Anthony M Polloreno, Karl Stratos, Philip Monk, Adarsh Chaluvaraju, Andrew Hojel, Andrew Ma, Anil Thomas, Ashish Tanwer, Darsh Shah, et al. Practical efficiency of muon for pretraining, arXiv preprint Stti::25000022222 2025.
Wei Shen, Ruichuan Huang, Minhui Huang, Cong Shen, and Jiawei Zhang. On the convergence analysis of muon. arXiv preprinl arXiv:2505.23777, 2025.
Chongjie Si, Debing Zhang, and Wei Shen. Adamuon: Adaptive muon optimizer. arXiv preprint vrXiv:2507.11005 2025.
Kimi Team, Yifan Bai, Yiping Bao, Guanduo Chen, Jiahao Chen, Ningxin Chen, Ruijue Chen, Yanru Chen, Yuankun Chen, Yutian Chen, et al. Kimi k2: Open agentic intelligence. arXiv preprint M.XXi::5577.255::4 2025.
Bhavya Vasudeva, Puneesh Deora, and Christos Thrampoulidis. On generalization of spectral gra- dient descent: case study on imbalanced data. ffigh-dimensional Learning Dynamics 2025, 2025.
Jinbo Wang. Mingze Wang. Zhanpeng Zhou, Junchi Yan, Lei Wu, et al. The sharpness dis- parity principle transformers for accelerating language model pre-training. arXiv preprint arXiv:2502. 19002, 2025.
Kaiyue Wen, David Hall, Tengyu Ma, and Percy Liang. Fantastic pretraining optimizers and where to find them. arXiv preprint arXiv:2509.0200 02046,, 2025.
David J Willssaw, O Peter Buneman, and Hugh Christopher Longuet-Higgins. Non-holographic associative memory. Nature, 222(5197):960 962, 1969.
Greg Yang. James Simon, and Jeremy Bernstein. A spectral condition for feature learning. arXiv preprint wrXiv:2310.178111 2023.
Yushun Zhang, Congliang Chen, Naichen Shi, Ruoyu Sun, and Zhi-Quan Luo. Adam converge without any modification on update rules. Advances in neural information processing systems, 35:28386 21-28386-28399,0 2022.
Yushun Zhang, Congliang Chen, Tian Ding, Ziniu Li, Ruoyu Sun, and Zhiquan Luo. Why trans- formers need adam: A hessian perspective. Advances in neural information processing systems, 37:131786- 131233, 2024a.
Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Diederik P Kingma, Yinyu Ye, Zhi- thi-Quan Luo, and Ruoyu Sun. Adam-mini: Use fewer learning rates gain more. arXiv preprint ar Xii:224006167933 2024b.
Zeyu Zhang, Akide Liu, Ian Reid, Richard Hartley, Bohan Zhuang, and Hao Tang. Motion mamba: Efficient and long sequence motion generation. In European Conference on Computer Vision, pp- 265-282. Springer, 2024c.
Rosie Zhao, Depen Morwani, David Brandfonbrener, Nikhil Vyas, and Sham Kakade. Deconstruct- ing what makes a good optimizer for language models. arXiv preprint vrXiv22407.07972 2024.
Dongruo Zhou, Jinghui Chen, Yuan Cao, Ziyan Yang, and Quanquan Gu. On the convergence of adaptive gradient methods for nonconvex optimization. arXiv preprint arXiv: .060000.0566,... 2018.
Fangyu Zou, Li Shen, Zequn Jie, Weizhong Zhang, and Wei Liu. sufficient condition for conver- gences of adam and rmsprop. In Proceedings of the IEEE/CVF Conference on computer vision and pattern recognition, pp. 11127-11135, 2019.
13

---

Published as a conference paper at ICLR 2026
A USE OF LARGE LANGUAGE MODELS (LLMS)
Large language models (LLMS) were used solely to aid and polish the writing of this paper. The authors generated all research ideas, methods, analyses, and results independently. LLM assistance was limited to improving clarity, grammar, and readability of the manuscript text. No content was fabricated introduced by the LLM beyond these language refinements.
B NOTATIONS
Let [M] for the set f1...., J]. For matrix X R X2 is its i-th column and X. 1 is its last column. K is the K x K identity matrix, J is all-ones vector and JK, is the all-ones matrix. denotes the element-wise product.
C RELATED WORKS
Adam, proposed by Kingma & Ba (2015), was designed to make GD adaptive to the complex optimization landscape of neural networks. Existing works analyze Adam from two primary per- spectives: online optimization and feature learning. The online convex optimization view focuses on Adam's properties when optimizing convex or non-convex loss functions. From this perspective, Chen al. (2019) and Zhou et al. (2018) derive non-convex convergence results for Adam, and series of subsequent works continuously relaxed the required assumptions for Adam's 's convergence while tightening its convergence rate. For instance, Zou al. (2019) proposes a set of easy-to-verine sufficient conditions for Adam's 'ss update rules to guarantee convergence. Défossez et al. (2020) de- rives the tightest dependency on the heavy ball momentum parameters. More recently, Zhang et al. (2022) demonstrates that Adam can converge without modification of its procedures, and Li et al. (2023) relaxes the smoothness assumption by employing an adaptive Lipschitz constant for gradi- ents. The feature learning view, on the other hand, highlights the relationship between deep learning characteristics and Adam, focusing more on how Adam's 's mechanisms influence the properties of learned features within deep networks. For example, Pan & Li (2023) examines the sharpness of GD and Adam and relates Adam's 's superiority to its low sharpness. Kunstner et al. (2024) finds that Adam is better at learning heavy-tailed distributions than GD. Furthermore, Zhang et al. (2024a) shows that Adam is adaptive to heterogeneous Hessian structures, thus optimizing faster than GD. In a spirit similar to our work, recent studies have also used ablation experiments to deconstruct Adam's 's effectiveness. For instance, Zhao et al. (2024) and Zhang et al. (2024b) conduct detailed ablations on Adam's hyperparameters and components, identifying that its benefits are particularly pronounced for the first and last embedding layers of language models. While these works focus on understanding Adam's 's existing components, our study applies a similar ablation methodology to understand the impact of different optimizer, Muon, on the internal modules of a Transformer. More literature on Adam included the survey by Abdulkadirov et al. (2023).
Muon, proposed by Jordan et al. (2024), applies spectral normalization of the gradient to update parameters. At a high level, Muon can be understood as steepest descent with respect to the matrix operator norm (Bernstein & Newhouse, 2024). Alternatively, it can be viewed as maximizing the feature update subject to a parameter update constraint (Yang et al., 2023). Experiments show that Muon consistently outperforms Adam across diverse model sizes and architectures, including dense transformers and Mixture-of--Exerts (Liu et al., 2025; Jordan et al., 2024). Building on this, Si al. (2025) introduces an adaptive variant of Muon. To explain its advantages, Lau al. (2025) introduces unifying preconditioning framework, distinguishing optimizers that address curvature anisotropy (like Adam) from those that address gradient anisotropy (like Muon), and proposes a generalized optimizer class named PolarGrad. Sato et al. (2025) and Shah et al. (2025) examine the critical batch size of Muon, while other works analyze its convergence in convex and non- convex settings (Li & Hong, 2025; An et al., 2025; Kovalev, 2025; Pethick et al., 2025; Shen et al., 2im & Oh (2026) derives the convergence bound of Muon, including the influence of NS steps. Furthermore, Grishina et al. (2025) proposes accelerating these NS steps via Chebyshev- optimized coefficients. Concurrently, Vasudeva et al. (2025) studies Muon on shallow ViTs for computer vision, grounding their results for gradient descent and Muon in linear regression. In contrast, we investigate Muon in the context of LLMs, focusing on its effects on associative memory in next-token prediction. Recent works have also investigated the scalability of the Muon optimizer.
14

---

Published as a conference paper at ICLR 2026
For instance, Wen al. (2025) reports that the benefits of Muon diminish with scale (dropping from 1.4x gain at O.1B to 1.1x at 1.2B), whereas Liu et al. (2025) observes that Muon maintains à Rs 2x FLOP-efficiency advantage over Adam
on 32B models.
Associative Memories have a long history in neural network design and knowledge storage (Hop- field, 1982; Kohonen, 2009; Willshaw et al., 1969). They have inspired architectures capable of retaining long histories, including RNNs (Orvieto et al., 2023) and Mamba (Zhang et al., 2024c). With the success of transformers, recent work has examined them through the lens of associative memories. Geva et al. (2020) and Dai et al. (2021) show that feed-forward modules store knowledge in Wour while Bietti al. (2023) demonstrates that the attention output matrix Wo also encodes associations of knowledge. Building on these findings, series of works edit knowledge directly by modifying these weights (Meng et al., 2022b; Fang et al., 2024). Beyond empirical results, theo- retical analyses have further clarified how transformers leverage associative memories: Bietti et al. (2023) conducts dynamic analysis of memory formation, while Nichani al. (2024) constructs explicit associative memory mechanisms in both attention and feed-forward modules.
D STEEPEST DESCENT VIEW UNDERSTANDING MUON AND ADAM
Bernstein & Newhouse (2024) showed that many popular deep learning optimizers can be under- stood through the unifying framework of steepest descent, once their exponential moving averages (EMAs) are disabled. This perspective shifts the focus from heuristic or second-order motivations to a more fundamental, geometric view: the choice of an optimizer is equivalent to choosing a specific nomm measure the "size" of the weight update.
The Steepest Descent Framework. The core idea to find a weight update, Aw, that minimizes a local quadratic approximation of the loss function. This is formulated as the following optimization problem:
A Aw' argmin
g Aw +
2 |Awl?
Aw
where g is the gradient of the loss, À >> is "sharpness" parameter that controls the step size, and || || is chosen norm.
The solution to this problem can be expressed as:
Aw -1 dd
where the step size | and the update direction d = arg max = tt. Here, - denotes the dual norm of) ||- (defined as ||y||- sup y T x). The key insight is that different choices of the norm lead different update directions d, recovering the update rules of well-known optimizers.
Muon Steepest Descent under Spectral Norm. The update rule of the Muon optimizer is derived by applying the steepest descent framework to weight matrices equipped with the spectral norm, denoted in the paper the || -lle operator norm (defined as its largest singular value, |
l2-t,z = mmax (A) = s Ax|22) For gradient matrix G. the problem is to find the
update A W that solves:
AW argmin AW
(G, + AAW P
-2
The solution to this problem is directly determined by the Singular Value Decomposition (SVD) of the gradient, G = UEVT. The resulting update direction, which maximizes alignment with the gradient under the spectral norm constraint, is shown to be UV The corresponding dual norm of the gradient, ||G - which seales the step size, is found to be tr(E), the sum of the singular values. Combining these components yields the final steepest descent update rule:
tr( tr( A
AW
UV
This demonstrates that Muon's 's core operation is principled descent step where the singular vectors of the gradient determine the direction, and the sum of its singular values scales the step size.
15

---

Published as a conference paper at ICLR 2026
Adam as Steepest Descent under &oo Norm. Adam can be understood steepest descent the flattened parameter vector w when the space is equipped with the vector infinitty norm (foo (defined as the maximum absolute value of its elements, ||x max; (zt)). For gradient vector g, the optimization problem is to find the update Aw that solves:
Aw"
argmin
g
Aw Aw
The update direction that maximizes alignment with the gradient g under the infinity norm constraint the sign of the gradient, sign(g). The corresponding dual norm of the gradient, ||gll which scales the step size, is the & norm, ||gl] (the sum of the absolute values of its elements, ||x||1 Σ, [z.). Combining these components yields the final steepest descent update rule:
Aw |g sign(g).
This reveals that Adam's 's fundamental operation corresponds to descent step where each parameter moves with the same magnitude, determined only by its gradient's 's sign.
E DETAILED DISCUSSION OF THE THEOREM 4.3
The proof of Theorem 4.3 is provided in Appendix H. We now explain the results for the three optimizers separately. For GD, the quantity r(a, (3) S measures the imbalance of the data dis- tribution: r(a, 3) = 1 corresponds to perfectly balanced data, while r(a, 8) < 1 indicates severe imbalance. The results show that if one set of (s, o) triplets is learned with the correct-class prob- ability [Fw)(F)]] , of at least 1 - e, then there exists another triplet whose correct-class probability is 7(c K)(). Thus, GD is highly sensitive to data imbalance: as the training distribu- tion becomes more imbalanced, the dispersion of correct-class, probabilities across items increases, i.e. the maximal probability gap A(W) grows and [[j) [/w (B&)Ja decreases. This mirrors the message Figure 4(b), 4(c), and Figure 3(d) Section 3.3.
In contrast, Muon learns in a balanced fashion, unaffected by data imbalance for any embeddings E and E. Our results show that when the best-learned triplet achieves correct-class probability of at least 1 - c, the worst-learned triplet has a comparable correct-class probability at least - e(1 + Olog K/K))- This justifies Observation 3. Furthermore, consistent with Observation 2, Muon's update Guoon rule allocates equal strength to all update directions; equivalently, the singular values of G Guoo (Wo) are nearly identical.
Our analysis shows that Adam's performance is unstable with respect to the embeddings and E, reflected by the large error bars in Observations 2 and 3. Adam's element-wing normalization disrupts the inherent matrix structure of the gradient. When embeddings of different triplets have disjoint supports (e.g., = E IR,x), Adam can optimize parameters in a balanced manner. However, when embeddings overlap. the sign operator in Adam can introduce imbalance. In par- ticular, the worst-optimized triplet may then have correct-class probability O K -033). These exponents (0.3, 0.7) are intrinsic to Adam's update under certain embeddings and are independent of a or B. Moreover, the Adam update GSignGD (Wo) exhibits pronounced spectral decay- for ex- ample, its smallest singular value can be less than 25% of the largest-unlike unlike the nearly uniform singular values of Muon. This spectral decay explains the poor isotropy reported in Observation 2.
F EXPERIMENTAL DETAILS
F.1 EXPERIMENTAL DETAILS OF TRAINING ON FINEWEB
When training 160M models on Fine Web, we disable weight decaying and Nesterov acceleration for both Adam and Muon. Thus, we only compare their performance along. To set the learning rate,
conduct a grid search on 1 x 10 5 x 10 20 2 x 10 2 x 10 10-2 2.5 x 10 10-3 3 2 x 10-3 3 x 10 3 10-3
10 10-4,2 2,2 10 When conducting the "Independent Blocks and "Combined Configuration experiments Section 3.1, we just fix the learning rate of Muon. We set B = 0.8, 3g = 0.95 for Adam and set 3 0.95 for Muon. When training 0.7B models Fine conduct grid
16

---

Published as a conference paper at ICLR 2026
search of learning rate 2 x 10-3, 10 3, 10 x 10 We set 31 0.9. 32 = 0.95 for
Adam and set 3 = 0.95 for Muon. We do not adopt group query attention in the structure; thus, the parameter sizes of Wo. WR. Wy, and Wo are the same. We conduct experiments on 8 A100 with 80 GB memory.
F.2 ISOTROPICITY METRICS EXPLANATIONS
Normalized SVD Entropy. This metric, adapted from Alter al. (2000), quantifies the uniformity of the singular energy distribution. A higher entropy value indicates more isotropic matrix where energy distributed evenly across many directions. It is defined as the Shannon entropy of the distribution q, normalized by the maximum possible entropy: Horm
log x Di=1 q log 9s-
Effective Rank. The effective rank (Roy & Vetterli, 2007) provides continuous measure of the number of significant singular dimensions used by the matrix. is calculated as the exponentiation of the unnormalized Shannon entropy, which corresponds the perplexity of the energy distribution: eRank(o) = exp (- - log 4:).
Top-k Energy Fraction. This metric measures the concentration of energy within the Top- principal singular components. Assuming the singular values are sorted in descending order
(01 2 02 2.... ... on), it is the cumulative sum of the first & energy fractions: TopE, (o) Σy
Eigenvalue Quantile Ratio. To measure the spread of the singular energy distribution while being robust to extreme outliers, we compute the ratio of the 75th percentile (Qs) to the 25th percentile (Q1) of the eigenvalues {o?)== (r) = Qa((a]))
Qi((a;1)
F.3 DATASET DETAILS FOR THE HEAVY-TAN TAIL KNOWLEDGE TASK
Following Allen-Zhu & Li (2024), the foundation of our knowledge-intensiie task is of question-answering (QA) pairs derived from synthetically generated biographies. Each biography constructed from combination of seven key attributes: name, birthdate, birthplace, educational institution, major, employer, and workplace. The attribute values are sampled from predefined lists, creating diverse set of entities. Specifically, we use approximately 400 first names, 1000 surnames, 300 educational institutions, 100 majors, and 300 employers. Each synthetic individual is assigned unique combination of these attributes, forming a distinct biographical profile. For example, generated biography might look like this:
Ashton Hilda Older has birthday that falls February 01, 2063. Miami, FL is the birthplace of he. He is an alumnus of Saddleback College. He has General Literature education. He works closely with BlockFi. For professional growth, he chose to relocate to Jersey City.
This text is generated by combining the structured attributes (name, date, location, etc.) with set of sentence templates.
A predefined set of QA templates is then used to generate the final training data. These templates contain placeholders corresponding to the biographical attributes. By formatting these templates with the information from each synthetic biography, generate collection of concrete QA pairs for each entity. For example, for the entity "Ashton Hilda Older", we can generate the following six QA pairs:
17

---

Published as a conference paper at ICLR 2026
1. What is the birth date of Ashton Hilda Older?
4. What major did Ashton Hilda Older study?
Answer: February 01, 2063.
Answer: General Literature.
2. What is the birth city of Ashton Hilda Older?
5. Which company did Ashton Hilda Older work for?
Answer: Miami, FL.
Answer: BlockFi.
3. Which university did Ashton Hilda Older study?
6. Where did Ashton Hilda Older work? Answer: Jersey City.
Answer: Saddleback College.
To evaluate the optimizers on knowledge-intensive task with data imbalance, we constructed synthetic dataset where the number of guestion-answering (QA) samples per class follows a power- law distribution. This is designed simulate real-world scenarios where few entities (the "head")) are highly represented, while most entities (the "tail") are rare.
The generation process is controlled by an integer parameter, m. The classes are organized into m + 1 groups, indexed from g 0 to
Group y contains Ng classes, where No and N, 29 for g >
Each class within group is allocated specific number of "selections," Sy 2m-9.
For each selection, we generate ng unique QA pairs by formatting templates with bio- graphical information corresponding to that class.
Thus, the total number of QA samples for any given class in group is Sg x Uga- This structure ensures that the single class group 0 has the most samples, while the numerous classes in group have the fewest.
In our experiment, we set the parameters to m 15 and Age 6. This results in a dataset with total of 25 32, 768 classes. The number of samples per class ranges from 196, 608 for the head class (group 0) down to just 6 for each of the 16, 384 tail classes (group 15). The final distribution visualized Figure 3(a) in the main text.
To evaluate the model's performance on this pure memory task, we measure the First Token Accu- racy (FTA) on the answers. This metric assesses the model's ability to correctly recall information by checking if the first generated token of the answer matches the ground truth. Furthermore, to un- derstand how optimizers handle data imbalance, we analyze the FTA across different data frequency groups, from high- frequerry (head) low-frequency (tail) data.
F.4 EXPERIMENTAL DETAILS ABOUT ANGLES BETWEEN ASSOCIATIVE MEMORIES EMREDDINGS
Following Fang et al. (2024), analyze the associative memories in the FFN modules. To obtain E, use the activations within the feed-forward modules, and for E,, we take the corresponding module outputs. We evaluate knowledge items from two widely used datasets: Counterfact (Meng al., 2022a) and ZsRE (Levy al., 2017). Results on Counterfact are shown in Figure 4(a), while results on ZsRE are provided in Figure 18 in Appendix G.I
F.5 EXPERIMENTAL DETAILS OF ONE-LAYER MODELS
We set the hyperparameters as K 999, a = 0.8, 3 = 0.2. For the support-decoupled setting, we set E and as identity matrices. For the support-couples setting, we set and É according to the construction presented in the proof of Theorem 4.3 in Appendix H.
18

---

Published as a conference paper at ICLR 2026
G ADDITIONAL EXPERIMENTAL RESULTS
G.1 MAXLOGIT PER LAYER ON THE 160M NANOGPT MODEL VIA MUON OPTIMIZER
In this subsection, we present the MaxLogit values for each layer of the 160M NanoGPT model trained using the Muon Optimizer. Following Gemma 3 (Kamath et al., 2025), we introduce RM- SNorm to the attention mechanism. The attention mechanism in our model is defined as follows:
O Softmax(@K") Q = RMSNorm QQ) K RMSNorm(K) K)
where RMSNorm is defined as RMSNorm(a) VILL
with d being the dimension of x.
MaxLogit is defined as:
Smax max o: 1.3
k;
representing the maximum value in the attention scores before softmax normalization.
The MaxLogit values for each layer are summarized in Table 1.
Table I: MaxLogit values per layer on the 160M NanoGPT model via Muon Optimizer.
Layer
3
10
11
12
MaxLogit 8,396 6.880
6.009
7.676 6.349
5.890
7.688
6.314
6,205
5.613
6.033
6.371
Recent reports Team al. (2025) have shown a potential "MaxLogit explosion" phenomenon, where Smax grows steadily (often near-linearly) during training, leading to overly peaked attention, gra- dient spikes, and degraded optimizer comparisons. We included this measurement to rule out the possibility that Muon's 's comparatively smaller impact on the QK blocks (relative to VO/FFN) is sim- ply due to suppressing such an instability. In our 160M setting, with RMSNorm applied to both Q and K (following Gemma 3), the per-layer MaxLogit values remain moderate and show no runaway growth. Thus, for this model size and normalization scheme, differences in Muon's 's effectiveness across components cannot be attributed to avoiding a MaxLogit explosion in attention.
G.2 CONTROLLING FOR PARAMETER COUNT IN COMPONENT-WIS ABLATIONS
A potential confounding factor in our ablation studies (Section 3.1) is that different model compo- nents contain different numbers of parameters. One might argue that applying Muon to a larger component naturally yields greater gains simply because more parameters are being optimized dif- ferently. To disentangle the effect of component type from the effect of parameter count, we measure the performance gain per parameter.
We measure the validation-loss improvement at 10,000 steps when applying Muon to a single com- ponent (QK, VO, Win Wan) relative full-Adam baseline. This gain is then normalized by the number of parameters in that specific component. For the 160M model, the parameter counts satisfy [Wvl = Wol = (Wo |W&l and Wi, = W ourl 4 (Wol.
Figure 5(a) reports the validation loss at 10,000 steps, normalized by the number of parameters in each component. This result shows that the normalized gain for VO Attn is approximately times greater than that for QK Attn, even though both components have the same number of parameters. The gains for Wi and our are also substantially higher (over 3x) than for QK Attn. Although Wi. and W have twice as many parameters as QK, their normalized gains are far more than half the gain of QK.
To provide even more direct evidence, we designed the second experiment where the number of parameters optimized by Muon is held exactly equal across different components. We achieve this by comparing three configurations: Muon applies to QK matrices in all layers; but Muon applies to Wm and Wont matrices in only the huddbreee layers. In this setup, the total number of parameters of QK, Win and W optimized by Muon are identical. The results in Figure 5(b) show that even when optimizing an identical number of parameters, the gain from applying Muon to Win or W Wi.. is much more than the gain from applying it to the QK blocks.
19

---

Published as a conference paper at ICLR 2026
7.
6.0
StaetV Ane Plaz. Mys WAkaarCKK ARe. Pic)
Validation Los
Vol Volistation Lave
MaaxVVo Ana Mis. Myaakhhaassk
0.3
Muue W Midenn ataa Wea W..)
NivwiWW ai Nareall Arc, Wa. Wass)
4.1
2040
4000
6000
10033
10000
2100
4000
6000
10000
Training Steps
Training Steps
(a) Independent Blocks with Gated FFN
(b) Combined Configuration with Gated FFN
Figure 6: Validation loss comparison on the 160M NanoGPT model with gated FFN under different Muon/Adam assignments. Panels (a) and (b) show the validation loss over training steps for the Independent Blocks and Combined Configurations settings, respectively.
6.14
Normaliod Gain vs Adam
A.I1
Cain ve Adam
Acaeet
2x187/
Maun/Weal
(a) Val. loss gain per parameter.
(b) Val. loss gain under equal parameter budget.
Figure 5: Component-wise validation-loss gain of Muon over Adam at 10,000 steps.
These results demonstrate that Muon's effectiivnnes is not simply about the quantity of parameters, but highly specific to the function the parameters. The associative memory components (VO, Wi. Wo.) derive much larger benefit per parameter, reinforcing central claim that Muon excels at optimizing these specific parts of the Transformer architecture.
G.3 ADDITIONAL RESULTS FOR GATED FFN ON FINEWER EB
To verify that our findings in Section 3.1 are not specific to the non-gated FFN architecture, we repeat the same "Independent Blocks" and "Combined Configurations" experiments on the 160M NanoGPT model with gated FFN. The results are presented Figure 6.
The conclusions are almost identical to those from the non-gated setting (Figure 1). Specifically, in both the independent and combined settings, applying Muon to VO+FFN yields the most significant validation loss reduction, closely tracking the performance of full Muon. In contrast, applying Muon only the QK blocks provides minimal benefit over the Adam baseline. This confirms that our finding-that that the associative memory components (VO and FFN) the primary beneficiaries of Muon-is robust to variations in the Transformer architecture, holding for both gated and non non-gated FFNs.
Furthermore, we analyze the spectral dynamics of the weight matrices for the gated FFN model, with results for the VO and Wa. matrices shown in Figure 7. The trends consistent with Ob- servation 2 from the main text: for both matrices, Muon leads to significantly higher SVD entropy and effective rank (eRank) compared to Adam. This indicates that Muon encourages the learning
20

---

Published as a conference paper at ICLR 2026
1pp
Tepit1:12
Te
SND
200
3BC
2000
0001
7000
101
330
306
1601
200
62 100
12"
5001
075035 RusH
400
226
3091
12000
2000
20:0
1000
2003
3000
202
330
30.0
1500
2001
120
PX0
12
(a) VO(Gated FFN)
(b) Watee FFN)
Figure 7: Spectral Dynamics of Transformer Weight Matrices During Training, Each panel reports four metries characterizing singular value distributions: SVD entropy, Top10E, eRank, and Q75/Q25 ratio. The four subplots correspond to different weight matrix groups: (a) VO and (b) Wou-
of more distributed, higher-dimensional representations in the associative memory components, finding that holds true for the gated FFN architecture as well.
G.4 SCALING TO THE 0.7B NANOGPT MODEL
To evaluate the scalability of our findings, we extend our experiments from the 160M model to a larger 0.7B parameter model. This section presents the results of this scaled-up analysis, examining whether the advantages of Muon observed in the smaller model persist at a larger scale.
633
AlAAdan
MaveOOK AmalAdan Al. FFN)
MassCO Ant FFN)
60-
3.5
MicerD Ane BNNAMRoKKY Ano)
Manil Aln, EEN MA Ats))
Mo Ans OK Arm)
5.0
MazetO Ans PPN AAmIR Anto
Munel MSS
Mice Arn, DYNUACmm: Ams)
Valtaana
Valiariim
35
39-
2000
6008
8003
10000
2080
4000
6000
5600
10000
Trining Stepe
Texining Steps
(a) Non-gated FFN
(b) Gated FFN
Figure 8: Validation loss comparison on the 0.7B NanoGPT model. (a) Combined configuration with non-gated feed-forward networks.(b) Combined configuration with gated feed-forward networks.
21

---

Published as a conference paper at ICLR 2026
VVD Entropy (1)
Wic Entropy
Wost Entropy (1)
0.9
0.98
11.05
(1)
0.8
Entooge
0.7
0.94
0.6
0.92
0.4
Maneall
0.90
0.0
4000
Trining Stepe
Steque W <Rann (1)
Stege
VO- oRank
Way <Rann (1)
5.00
1100
1000
1000
802-
<Rask
@Rosk
000
600
700
300
600
2000
.000
6000 Trening Sleps
1000
10000
2000
4r8m
(000
1000
10000
2000
4000
6000
8000
1003)
Training Stepee
Training Stepe
Figure 9: Spectral Dynamics of Weight Matrices During Training on the 0.7B NanoGPT model.
Enropyy (t)
(I
W Eartogy
1.44
0.000
0.12
5V0 Sy0
2:75
2-7
02-
0:5
A1
OXE
9800
10890
33.0
4000
69R
5000
10000
0000
20089
6000 Trulilng
180
V/O claak (1)
Wr,
Weste <Raak
<Rark
1100
0228-111
100
1850
-00
1850
150
26.4
-090
2.31
(5.9)
20.0
-2:0
5.01
1000)
25.0
4.53
00.00
2.59
3008
2000
4.03
00.00
1881
18.99
Figure 10: Spectral Dynamics of Weight Matrices During Training on the 0.7B NanoGPT model with the Gated FFN.
Figure shows the validation loss curves for various optimizer configurations. Consistent with our findings on the 160M model, applying Muon to all components achieves the lowest validation loss, outperforming Adam bascline. The hybrid experiments further reinforce our earlier conclusions: applying Muon to only the VO and FFN components yields performance nearly identical to that of the full Muon optimizer, whereas applying only to the QK components offers little advantage over Adam.
The spectral dynamics, shown in Figures 9 and 10, also align with Observation 2. For the VO, Win Wgate gate (in model with Gated FFN) and Wan matrices, Muon leads to higher SVD entropy and eRank compared to Adam, indicating that encourages the learning more distributed, higher- dimensional representations. Overall, these results demonstrate that the benefits of Muon and the underlying mechanisms scale to larger models.
G.5 ADDITIONAL RESULTS ABOUT SPECTRAL DYNAMICS OF TRANSFORMER WEIGHT MATRICES DURING TRAINING
To complement the main-text analysis (Fig. 2), we also evaluate spectral dynamics during training for the 160M NT model with both non-gated and gated feed-forward networks (Fig. 11).
22

---

Published as a conference paper at ICLR 2026
The analysis includes Win for both configurations, well as the gate matrix W gate for the gated version. The conclusions are consistent across all three matrices and mirror the non gated setting: with Muon, SVD entropy and eRank increase, while Top-k energy and the Q75/25 ratio decrease, consistent with Observation 2 in the main text.
#0Oeppp1
HD
beniy lave
Dalay lime
(a) Win (Non-Gated FFN)
(b) Win (Gated FFN)
(c) gate
Spectral Dynamics of FFN Weight Matrices During Training on the 160M NanoGPT model. Each panel reports four metrics characterizing singular value distributions: SVD entropy, TopI0E, eRank, and Q75/Q25 ratio. The subplots correspond to different weight matrices: (a) Wm. (non-gated), (b) W (gated), and (c) W game (gated).
Figure 11:
G.6 DETAILED EXPERIMENT RESULTS ABOUT HEAVY-TAI IMBALANCE KNOWLEDGE TASK
To complement the qualitative trends shown in Section 3.3 (Fig. 3), we report the exact First Token Accuracy (FTA) for selected tail groups at three training checkpoints (2k, 5k, 10k steps). We focus on groups g = 11, 13, 15, which represent increasingly rare (mid-tail, tail, extreme tail) frequency bands in the power-law distribution (recall that larger g implies fewer samples per class). The tables contrast full Muon, Adam, only VO&FFN or only to QK). The numbers highlight: (i) Muon's rapid convergence on rare groups (already strong by 2k, near-saturated by 5k), (ii) Adam's (aaa's persistent head-tail gap, and (iii) the dominant contribution applying Muon to VO&FFN for tail generalization (the VO&FFN hybrid closely tracks full Muon, whereas the QK-only hybrid lags). These quantitative results substantiate Observation that Muon delivers more balanced learning.
and two hybrid configurations (Muon applied
Table 2: Heavy-tail knowledge task: Group performance by optimizer (2,000 steps)
Optimizer
Group
Muon
Adam
SGD+Mom.
Muon(VO, FFN)
Muon(QK)
11
0.854 0.029
0.312 0.043
0.156 0.037
0.814 0.022
0.472 0.041
13
0.386 = 0.029
0.146 = 0.015
0.120 = 0.012
0,256 + 0.030
0.154 0.032
15
0.140 = 0.027
0.090 = 0.031
0.082 = 0.013
0.114 0.023
0.086 + 0.037
Table 3: Heavy-tail knowledge task: Group performance by optimizer (5,000 steps)
Optimizer
Group
Muon
Adam
SGD+Mom.
Muon(VO. FFN)
Muon(QK)
11
0.996 = 0.006
0.936 = 0.039
0.314 = 0.021
0,992 0.005
0.970 0.007
13
0.964 0.023
0.298 0.074
0.148 0.013
0,934 0.015
0.354 0.032
15
0.320 = 0.028
0.110 # 0.027
0.084 + 0.011
0.254 0.026
0.118 + 0.019
23

---

Published as a conference paper at ICLR 2026
Table 4: Heavy-tail knowledge task: Group performance by optimizer (10,000 steps)
Optimizer
Group
Muon
Adam
SGD+Mom.
Muon(VO, FFN)
Muon(QK)
11
1.000 + 0.000
1.000 = 0.000
0,422 + 0.023
1,000 + 0.000
1.000 + 0.000
13
1.000
0.000
0.890 0.042
0.294 + 0.013
0.998 0.002
0.940 0.034
15
0.976 0.006
0.264 0.048
0.126 0.021
0,954 + 0.021
0.286 0.039
G.7 ADDITIONAL. EXPERIMENT RESULTS ABOUT HEAVY-TAIL IMBALANCE KNOWLEDGE TASK WITH GATED FEED-FORWARD NETWORKS
This subsection complements the main heavy-tail results Section 3.3 by studying the gated feed- forward networks (Gated FFN) variant. We follow the same presentation order as the main text: first an overview figure (sample distribution and learning curves under different optimizers), then tables reporting the exact First- Token Accuracy (FTA) for tail groups 9 {11, 13. 15) three training checkpoints (2k, 5k, 10k steps). The findings mirror the non-gated setting: (i) full Muon consistently outperforms Adam and SGD+Momentum on rare classes and reaches high accuracy earlier; (ii) the VO&FFN-hybrid (Muon applied to VO and FFN while Adam is used for QK) closely tracks full Muon, indicating that VO&FFN are the primary levers for tail generalization; (iii) the QK- only hybrid offers limited gains. Overall, the gated FFN does not change the qualitative conclusions about where Muon helps most. See Fig. 12 and Tables 5-7 for details.
Doog
Deap?
.0 Sanglees
Drgg
Doee
Aucumsy (FTA'
0.8
Godd
3.6
0.
02
16
10¹
10
10
2020
4106
6600
9000
1000
3300
4000
6803
9303
1000
Index (sumaa))
Thiting Staps
Trining Stips
(a) Sample/class
(b) Muon
(c) Adam
Ascumcy (TTA)
0.8
Axcumny (FTA)
2.8
3.6
0.2
02
2003
4303
3000
10000
2020
4166
660
3000
1000
2300
4000
6003
1000
Traming
Stepe
Thining Stps
Thining
(d) SGD+Momentum
(e) Muon(VO, FFN)/Adam(QK)
(f) Muun((K//Adamm(VO.FF))
Figure 12: Performance comparison of different optimizers on heavy-tailed knowledge task with gated feed-forward networks. (a) The distribution of samples per class follows power law. (b-d) Performance of Muon, Adam, and SGD+Momentum optimizers. (e) Muon (VO, FFN)/Adam (QK). (f) Muon (QK)/Adam (VO, FFN).
Table 5: Heavy-tail knowledge task with the Gated FFN: Group performance by optimizer (2,000 steps)
Optimizer
Group
Muon
Adam
SGD+Mom.
Muon(VO, FFN)
Muon(QK)
11
0.896 = 0.009
0.214 + 0.063
0.146 + 0.018
0.892 0.021
0.330 0.042
13
0.478 0.034
0.116 0.030
0.110 0.007
0.458 0.037
0.140 0.019
0.178 0.018
0.086 + 0.013
0.074 + 0.009
0.166+0.011
0.090 0.020
24

---

Published as a conference paper at ICLR 2026
Table 6: Heavy-tail knowledge task with the Gated FFN: Group performance by optimizer (5,000 steps)
Optimizer
Group
Muon
Adam
SGD+Mom.
Muon(VO, FFN)
Muon(QK)
0.998 0.002
0.928 0.024
0,252 = 0.016
0,990 =
0.960 0.032
13
0.990 = 0.010
0.216 = 0.052
0.156 = 0.024
0.968 0.028
0.290 + 0.046
15
0.510 # 0.039
0.092 # 0.015
0.080 + 0.016
0.468 0.016
0.098 + 0.013
Table Heavy-tail knowledge task with the Gated FFN: Group performance by optimizer (10,000 steps)
Optimizer
Group
Muon
Adam
SGD+Mom.
Muon(VO, FFN)
Muon(QK)
1.000 0.000
0.998 0.002
0.322 0.011
1,000 0.000
1.000 0.000
1.000 = 0.000
0.948 0.027
0.304 = 0.017
1.000 0.000
0.946 + 0.026
15
0.994 = 0.006
0.244 = 0.085
0.148 + 0.015
0.990 0.010
0.274 # 0.042
G.8 IMPACT OF DATA IMBALANCE LEVEL
To further investigate how the degree of data imbalance affects the performance gap between Muon and Adam, we conduct an ablation study on the heavy-tail knowledge task with varying levels of class imbalance. We compare three settings:
High Imbalance (base 2.0): This is the default setting used our main experiments (Section 3), where the number of samples per class follows the power-law construction in Section F.3 with base 2.0.
Medium Imbalance (base 1.2): A less skewed version of the same construction, where the base is reduced to 1.2 so that the head-tail ratio is smaller.
Uniform: A balanced setting where each group contains the same number of classes and each class is assigned the same number of QA samples.
The results presented in Figure 13. From left to right, the panels correspond the high- imbalance, medium-immalance, and uniform settings, each plotting the average First Token Ac- curacy (FTA) over all groups for Adam and Muon. As the data distribution becomes more uniform, the performance gap between Muon and Adam steadily shrinks, and in the uniform case the two optimizers behave very similarly, indicating that Muon's 's advantage is most pronounced in highly imbalanced, heavy-tailed regimes.
High-Imballnnee
Meddian-Imbbllanee
Uniform
1.0
3.0
1.0
0.0
0.5
0.8
Overal PTX
Overall FTA
0.4
2.4
0.4
Training
Stepa
Figure 13: Comparison Muon and Adam under different levels class imbalance on the heavy-tail knowledge task. From left to right, the panels correspond the High Imbalance (base 2.0), Medium Imbalance (base 1.2), and Uniform settings.
The datasets mixed with different levels of tailedesss exhibit two properties: (1) Figure 14(a) shows that the single-token distributions are not exactly the same, i.e., facts and tokens cannot be
25

---

Published as a conference paper at ICLR 2026
perfectly decoupled; and (2) the token distribution in the uniform mixture still follows Zipf's law (Figure 14(b)). Thus, we conclude that, under different levels of heavy- seavy-tailedness in the pretraining data, the benefit of Muon over Adam varies when the token distribution remains close Zipf's law: the more uniform the mixture, the smaller the gain of Muon over Adam.
10a
Power Mixing
10a
Uniferm Mixing
Frequency (log scale)
Uniform Miding
Froquency (log scale)
10a
10a
101
104
10a
10a
10a
10a
10a
10
0
500
1000 1500 2000 2500 3000 3500
1000 1500 2000 2500 3000 3500 Token Rank (sorted by Uniform Mixing freq.)
500
Token Rank (sorted by Power-law freq.)
(a) Freq. Comparison (Sorted by Power-law Rank)
(b) Uniform Mixing (Sorted by Uniform Rank)
Figure 14: frequecyyy profiles for the synthetic heavy-tail QA task. (a) Compares the token frequencies of the original Power-law mixing (red curve) and fact-balanced Uniform mixing (blue dots), using the token rank from the power-law mixing. (b) Shows the token frequency profile for the Uniform mixing data, re-sorted by its own token frequencies.
G.9 ADDITIONAL RESULTS ON WIKITEXTI0S
To verify that our observations Wib and the synthetic heavy-tail knowledge task transfer to a more standard language modeling benchmark, we additionally train 160M NanoGPT models on the Wikitext103 dataset. We keep the model architecture and most hyperparameters identical the Fin setup and only retune the learning rate for each optimizer with a small grid search.
Figure 15 provides overview of this setting. Panel (a) shows the empirical token frequency distribution of Wikitext103, which exhibits a clear heavy-tail pattern: a small number of tokens appear very frequently, while many tokens rare. In the plot, the vocabulary is partitioned into frequency-based groups, each containing approximately 10% of the tokens (from most frequent to rarest), make head and tail behavior more comparable. Panels (b) and (c) report the training loss curves for Adam and Muon, respectively. Consistent with our main results, Muon converges faster and reaches a lower training loss than Adam.
Figure 16 further highlights the difference between the two optimizers by plotting their training losses on the same axes. Looking from the head group to the tail group, the performance gap between Muon and Adam steadily widens: while the two optimizers behave similarly on high-frequency (head) tokens, Muon remains much stronger on mid- and low-frequency (tail) tokens. In addition, the error bars for Adam grow substantially toward the tail, indicating unstable generalization on rare tokens, whereas Muon stays consistently stable across all groups.
Figure 16 also reports the two hybrid configurations. The Muon(VO, FFN) variant, which applies Muon only to the value//utput and -redwrd blocks while keeping Adam on QK, almost overlaps with the full Muon curve, showing that most of the improvement comes from these components. In contrast, the Muon (QK)-only variant is very close to the Adam baseline, suggesting that using Muon solely the QK blocks brings limited benefit.
26

---

Published as a conference paper at ICLR 2026
5
Trahs
100
1800
Tranzng Maags
Truaing Skig
Token frequency
Muon
Adam
Figure 15: Performance comparison of different optimizers Wikitext103, (a) Token frequency distribution in the Wikitext103 training corpus, showing pronounced heavy-tail structure. (b) Training loss curve for Muon. (c) Training loss curve for Adam.
Adre
Moreno 2010
Hood//a FWO
Nasdfo F350
Hunalo 23N
MucuPo.1EM
Triin Lote
2003 600E 0303 5000 10000 Transing Stago Group
2850 60E0 1693 18000
2333 4206 6836 9300 19030
2003 4000 0060 1303 10000
2000 4950 0EEO 3890 18303
Truaaay Staga
"raaning Sugg
Trusing Siupa
Truining Stipe
Group
Group
Group
Group
Mion
ieki 210
Trais Line
Traee
Trantng Stage
Traasn Stoge
Training Skige
Truzang Stogs
Treisig Stiga
Figure 16: Training loss comparison on Wikitext103 across head and tail token groups under dif- ferent optimizer configurations. The curves correspond Adam, Muon, and two hybrid variants that apply Muon only to VO&FFN or only to QK. In these figures, the results of Muon(VO, FFN) coincide with those of Muon, while the results of Muon(QK) coincide with those of Adam.
G.10
ADDITIONAL RESULTS ON LINEAR REGRESSION
Adam: Thil
Adam: Head Tail
Muon: Mid
4.0
Mid
3.5
3.0
Average Validation
2.5
2.0
1.5
1.0
0.5
2000
4000
6000
8000
10000
Training Steps
Figure 17: Validation loss on linear regression across head and tail groups under different Adam and Muon.
27

---

Published as a conference paper at ICLR 2026
To further demonstrate that Muon is ineffective optimizing the QK parameters in the attention module, we consider an in-context linear regression task (Garg et al., 2022), which heavily relies on the QK parameters. In this task, the model is prompted with number of demonstrations (I4, Vi Vi); with y/: - = w and query Zg. where x; E Rd for E [K] and g. E Rdd The model is expected to output 24 w. Intuitively, the QK parameters capture the correlations between the demonstrations and the query and use them estimate x.g w. Following (Garg al., 2022), we train the model with lg loss. To test the efficacy the optimizers under a heavy-tailed task distribution, we partition into groups supported
mutually orthogonal subspaces, which appear in the training data with different frequencies. We perform a grid search over learning rates for Adam and Muon and report the results in Figure 17.
Figure 17 shows that Adam and Muon achieve similar performance across different groups. In particular, both optimizers effectively learn the head class but barely improve the tail class. This behavior is in sharp contrast to the results in Section 3.3, where Muon substantially outperforms Adam on tail classes. Hence, the linear regression experiment further supports our claim that the main benefit of Muon does not come from optimizing the QK parameters.
G.11 ADDITIONAL RESULTS ABOUT ANGLES BETWEEN ASSOCIATIVE MEMORIES EMBEDDINGS
Aagles Embeddings (Degrees)
Layer
Figure 18: Average angles between e, or en for items in ZsRE at layers 5, 10, 15, 20, 25 of Llama3- 8b-instruct.
H PROOF OF THEOREM 4.3
We separately derive the results for GD, Muon, and Adam in the following proof. For all of them, we define
Mep inf {1r 0 max [Ew)] Sc. where W Wo 77 Gupu Ga (H.1)
kE[K]
The quantity Vep represents the minimal step size for at least one triplet to be learned with error probability less than From the definition, we have that
Papt
min [S-n Goje (Fi)) kE[K]
Step 1: Caleulations GD.
We define the score of k'tth object for the k-th subject-relation pair with the parameter W as
s(k', k, W) = cxp( ET. W Ex)
Ek exp( E.T, WEA)
At Wo = Od.,d,, we have that
1
s(K' K. Wo)
for all k, k E [K].
K
28

---

Published as a conference paper at ICLR 2026
Proposition J.1 shows that the gradient is
1-a
WWW)(W)) = ELLEIL +
E EL+I:K E.tIKEL+1.K EL+1-K
E.K.LEIL LK
K --L
BJK.K-LELtaA -K-LK
1-a
(H.2)
From the gradient, it is easy to see that the first L triplets (s. o) share the same learning behavior, and the last K L triplets also share the same behavior. Thus, we calculate the results for k = and h = L + 1. The calculation for k = depends on evaluating its score function, which takes [ the form _ V w ((W E1, for k" E [l,.... K). Based on the gradient in (H.2) and the orthonormality of the embeddings, it evaluates a/ L for the case K"
and for all k" *
This leads to numerator in the softmax score of exp(1 L)L,, while the denominator sum consists A similar calculation for L+ shows
of one term exp(1) a/ L) and K terms of exp(0) that the argument of the exponent for the correct object, 17 ET.+1l 1
Vw [(W)) EL+1. evaluates to 7 (1 - (K -.). By defining 71 = a/KK) and 72 = (1 aa/((1(1(11 - 3) based on the problem setup (L = 3,),, we have that
exp(v12) exp(m)) K K [F-Nwe(ELtt)llnn exp(112)
exp(mm)
[(Ei) =
) + K - -1'
where 7/1 and 722 are defined as
1 a 72
71
3K
(1 3)K
Then we derive that To calculate the desired quantity, we define the quantity r(a, 3) to evaluate the balance of data as
nAD max(7, 72
log (cc11
1)(K -
(H.3)
= min { (a 3) 3(1 a(1-a)] B1 a) aal- a((13)) 33
r(a, B) =
Some basic calculations show that
min KE[K]
Con =
(H.4)
+ (1- (1 (K - 1)
When < 1, with the fact that
1 = =1 0 as r + 0, we have that
min
Gan (E+) x Rr().
kE[K]
Thus, the proof for the convergence of GD has been established.
Step 2: Calculations of Muon.
For Muon, we first calculate the SVD of the gradient. fact, we can write the gradient in Eqn. (H.2)
} ET
-WWC(W)) E. diag
I T,
1 1-
I
In
I T. K K- K-L
1- --
K
= T.
The SVD calculation of X T can be directly derived from Proposition J.3. Thus, the SVD of the gradient (WW((W) (F. U)E(F V)T. The update quantity G (Wo)
29

---

Published as a conference paper at ICLR 2026
Unnorm V of Muon is
GMuunl (Wo)
E1.L RL R E... EL-IKK - L,K L.K-I. 1
1
(KK ((--))E:tA))
L)E1.r. LE1+1xIK-I))
[K(a((K - 1)3 +1- a))
[K (K-14) L)a ET...
L(1 a)
K kI E.TAIx
F.LL F.L + FLTI.K F.T 1
1
(1 - -133 ELLLJE.LEI
+
K
3
1
( (B111 -a) -1 -.
B(1 a) L. a(1 -(L-3))BHKK } (H.5)
where = (a - 3)3 + (1 - 23333 the matrices RT.II-1 and -L- are defined in Proposition J.3, and the second equality results from the following facts
RL.L-1 = II, - I.I.I.
1
RK-LKK-LIRR--LLE ---K--LL -K-LK-L --L---L--KKLK-- K L-L -
Although the gradient composed of heterogeneous components from ËLL, Eiz and FLLIKK FLTIKK we can bound the convergence rate of [f-nGrum (F*)] for any k: an upper (resp- lower) bound is obtained by increasing (resp. decreasing) the coefficient of E, while decreasing (resp. increasing) that of Fgg- F.T for k' & k. In fact, Eqn. (H.5) implies that there exists constant C' > 0 such that the dynamics of the fastest- and slowest-learning triplets are bounded by those along the following two update directions.
2C'
C E.JKKET
-G (Wo) 1+
(BLLET + BLAKKEETIKK K
K
2C
C
-Mon (Wo) 1
-.J.KE.. (ÉI.LET. + ELLL.K K
K
Concretely, the rate of score increase for the correct object of the k-th triplet, which is given by the term - (W)) Ex fastest-learning triplet is lower-bounded by the corresponding rate derived from - (Wo). while the rate for the slowest-learning triplet is upper- upper-bounded by that from - -G (Wo). Thus, we only need to focus Muonn (Wo) and Muoe (Wo) to calculate the desired quantity. Following the similar procedures for GD to derive Eqn. (H.4), we have that for any 1) such that ma [Sw, (Fx))s k 2
in the exponent of the softmax score, is bounded. The rate for the
Muon
1 - (where W,, Wo 17 GMaon (Wo)), the following holds
1 - min [fw, (E+) c +1 (1) 1)r(k)-1' KE[K]
(H.6)
where r(K) = (K- 2C)/(K + +C). We further have that
(1 (k), ( ()
40
E -logkK log(K -11) -10((K1)))
(1 t) exp
log 1-c
KC2C
4C
-og((K-1)) - -log(K - -1) +0
(log K)²
c)
1
log 1-e
K + +0
K2
log K
)) +
(H.7)
30

---

Published as a conference paper at ICLR 2026
where the first equality results from the basic calculations, the second equality results from that exp(x) = 1 + x Oxx²) when I - O. Combining Eqn. (H.6) and (H.7), we have that
PMunn 21- =1-e(1
log K K )
Thus, prove the desired results for Muon.
Step 3:
Calculations of Adam.
The proof of the results for Adam is conducted under two cases. We will construct different embed- dings E' and E' in these two cases. In the first case, E E IK,K. With such embedding and sufficiently large K, we have that
GGiignD (W) (ign( L(W))) = 21 JK.K. Under such setting, all triplets share the same dynamic. Thus, have that
PSignGD
In the second case, we set E and E as block wise diagonal matrices. Here we set the block size as 3, ic., requiring that K mod O. Such a requirement can be satisfied infinitely often when K - Then the sufficient and necessary condition of Assumption 4.1 is that each x 3 block contains an orthonormal basis. To achieve this, define the following matrix.
[os a b cos sin sin cos cos sin c sin a cos R(a, b)) = sin cos b cos + cos sin c sin cos bsin c + cos a cos sin b sin
cos sin 6 sin ain cos b
sin b c
is obvious that R(a, b, e) R(a, b, = 13,3. Then we set E and E as E KK/3.K/3 @ R(3.638, 2.949, 5.218), 1=/K//.K/3 & R(1.715, 0.876, 3.098),
where @ is the Kronecker product. With these specifications and sufficiently large K, the Adam update matrix is
Gig (Wo)
K//3KK/3 A+ A @ &B,
where and B are specified as
0
0
A
2
0
2
-
2
22
These show that the diagonal block of -Gsigaci (Wo) is
1
A B
1
Since the k-th and (k + 3)th triplets share the same learning dynamics for all & E [K - 3], we focus on the learning dynamics of & = 1, 2, 3. We have that
R(3.638, 2.949,5.2189 (A + B) R(1.715, 0.876,
1.46552253
1.0132908
-0.111795637
0.0732561
1.00709257
-1.26935805
0.0544114
0.89611102
1.54147329
R(3.638, 5.249)
B R(1.715, 0.876,
0.192881 119288146
-1.24460331
-1.4058011
0. 2217512115
-1.2977753
-1.46585978
0. 12780259 -0.82466989
-0.93147899.
From the last columns of these matrices, following the similar procedures for GD to derive Eqn. (H.3), we have that
GD 5 1.541 + 0.930
log - 1)(K - 1)] 2.471
log [(6-1 - 1)(K - 1)]
31

---

Published as a conference paper at ICLR 2026
Then, from the first columns of these matrices, we have that
min [ Gogun (Ex)] 2 c+ (K- KE [K]
where
1.466+-0.20
1.668 2.471
2.471
Thus, we have that
ESgaati -(E-TK-1) - 00
Then we calculate the singular values of (Wo) We define the eigen vectors k, as Using the orthogonal invariance of singular values, -G -Gsignn (Wo) shares the singular values with the following matrix
U. i.e., ÜT IK/3,K/SUU = diag(K/3, 0...............
,0). ,
( & 13.3)( (GiignD((W)))) @ 13,3)
K//, & A+ (U 1K/3KK/3U) @ B diag(A- - KB33, ....,, A).
Thus, the singular values of A are also the singular values of (Wo). We have that
min SignGD (Wo)
miin (A)
o m (A) S 25%.
max SignGD (Wo))
Thus, we conclude the proof of Theorem 4.3.
I PROOF OF THEOREM 4.4
The proof of Theorem 4.4 takes two steps. In the first step, we derive the share form of W, along the whole optimization trajectory. In the second step, build the desired results the basis of step 1. Throughout the proof, we will write WMMoon as W, for the ease of presentation.
Step 1: Derive the shared forms of W, and GMuon-
We will derive the forms of W, along the optimization trajectory via the induction method. We first state hypothesis and then prove it.
Hypothesis 1. For any optimization step index E [T], the parameters W, can be expressed
W. = ÊX,E,
X, As + Co.
where A₂ and C2 are
ct Ju 21 JK-LL
22. Jr,KL
A diag(a II,b IK-L), G
d JK-LKKLL
where at, b,, 11
21 _t 4
12
22 E R are real numbers such that (1) at = b 2 0, and (2) cy =
K) for i,j [2].
When t = 0, it is obvious to verify that Wo Od,,d, satisfying this hypothesis with a, b, cf c.- c? 22 = O. Then we assume that this hypothesis holds for will prove that it holds for + Since W++1 W, 7/+1 innorm(,,) we need to show that - n.nom VE, satisfies the hypothesis. We define the score of k'-th object for the k-th subject-relation pair with the parameter W as
t), and we
exp(ET. W..) ΣK exp( ETT, WE,)
s(k', ks W)
According to the symmetry of We, we have that
s(k,k, k, Wt) s(1, 1, W.) for all L.
skk,,, k, W.) =(K, K, W.) for all k L.
32

---

Published as a conference paper at ICLR 2026
sK, Wi) s(2,1, for all k 2 L, *k.
ki Wa) = =,(K,1 W:) for all & L, k > L.
K, Wa) =((K K, W)) for all k, K L, k' k..
sK,, W.) = =s K, W.) for all k > L, k' S L.
Thus, Proposition J.1 shows that the gradient of W; is
-Vw-(W,)) E(T, B.)
where r2 and B, are defined as
r, diag
(1 (1+ s(,11 s(1, W.
1-a
K- + s K, W,) s(K, K, W.))
(,, JLLL
1-a + s11, K, Wi) JL,KLL
B, = a sK, Wi) JKLL
- 1, K, Wi) -KKLLK-L
Thus, Proposition J.2 shows that
MMon (Wi) El [dia(Ix)++ +
On Ju. Cz JK-L. L.I.
C12 Jr,K-I
ETT 1)
On JK JK-L,K- 1,KKII
where
Cn = V1.VV.1 U1
U U1 U1,2V22
C12 =
BK
VB(1 -3)
Cn = U2, V1,1 + 22 V1.
U2 V2 + U22 V2,2
Oz
VB - K)
(1 3)K
where U. V. E g2x2 are the orthonormal matrices defined in Proposition J.2. Since W4+1 We Mc-1 G (We), it obvious that ac+1 bl+1. The orthonormality of U and V implies that |0r VVal < 1. Thus, we have
V..VV.
V1,VV. - -0 -
3K
K
This further implics that c:41 (((111/K))))) The proofs for other of the proof.
are similar. This completes
Step 2: Establish the convergence results.
We note that this analysis is very similar to the proof of Muon in Theorem 4.3. Concretely, for Wt, the coefficients a, by, cf c.c,,.. 21 c²2 from multiple-step optimization share the same property with those of the one-step results. It means that there exists constant C >> such that the dynamics of the fastest- and slowest-learning triplets are bounded by those along the following two update directions in only one step.
20' K
(EL.L + ELLLK E -. E.
1+
Muun
-
1
2C
(E + BLANK B + K
C. EJK,KE..
Muon
K
The remaining analysis is then exactly the same as that of Theorem 4.3. Thus, we conclude the proof of Theorem 4.4.
33

---

Published as a conference paper at ICLR 2026
J SUPPORTING PROPOSITIONS
Proposition J.1. We define the score of k/-th object for the k-th subject-relation pair with the pa- rameter W
explEx, W Ex) s(k', k. W)
ΣAT exp(E E W
When the parameter W is trained with loss
C(W) Σ
pk log [fw (Ex)
k=1
the gradient of W is
K
Vw-(WW = Σ pk [1 [1- [1- W)
E
Σ K'
s(K
W) Ex Eg
Proof of Proposition J.I. The proof just follows from the basic calculus. Thus, we omit them here.
Proposition J.2. Let X = A kxk The matrix A diag(a IL, b l---)) is diagonal matrix whose first L diagonal elements are and the last K L. elements are b with a, b >> 0. The matrix C is block-wise constant matrix defined as
C11 JLLL C21 JK-L,L
C12 JL,KLL
C
022 222JK-LL. -K-KKLLK-]
Then X = UVVT, Here M, V,, U are defined as follows, All them can be decomposed into three blocks, each corresponding subspace. The first subspace
x $
S = 0K-
x ====0, = O, and R
The dimension of this space L 1. The singular value X corresponding to this subspace is a. The block of columns in both U and V that forms an orthonormal basis for this subspace is given by
RI,I-1 0 OK-L
where the columns of the matrix R E R((((-1 form an orthonormal basis for the subspace {a E R 1 = 0}. The second subspace is
S2 - [0 y TIK-==O 0, and E R
The dimension of this space is K L 1. The singular value of X corresponding this subspace The block of columns both U and V that forms an orthonormal basis for this subspace is given by
OL,K-L-1
RK-L-KII-I
where the columns of the matrix Rx- K,K- ---1 II-1 B(K--LX(X------ form orthonormal basis for the subspace (y Ry TTx = 0}. The remaining 2-dimensional subspace is induced by
2 x 2 matrix defined as
[a g
M
didag(51,8)))
7
8
where the elements M are defined as
Leai
B == -
L)c12,
7 VL(K L) c.
5= b+ (K L)ezz-
34

---

Published as a conference paper at ICLR 2026
The singular values $1, 82
T TT 4A
a= + + +2 A = (ad Byy)
$1.2
The singular values of X this subspace are 51 and 82. The corresponding right singular vectors (u)) and left singular vectors (u+), which form columns of V and U respectively, are given by:
0x = Vi + V2,,2, u, e1 + U2 for i 1,2,
where the vectors é and are defined as
0r
=
0K-L
LK - K-14
In summary, the SVD of X
Σ
diag(a -1-1, 1K K-L-1: 81, 82):
RI,L-1
0L.K L-1
V
OK- OK-LL--1
RK-LK- ---1
U1 U2
RI,I-1
U
OK OI-L.II-11
RK-LK-IA
U1, U2
Proof of Proposition J.2. We first prove the results for S₁. For any vector e in S1, it is direct to verify that
x X
a²
Ox OK-L
0x
Thus, the singular value of X corresponding to the subspace spanned by the vector [e --]] and the corresponding columns of V form an orthonormal basis for S.. For the U calculation, we have that
_ is
x
0K-L = a
0K-L
Thus, the corresponding left singular vectors (columns of U) are identical to the right singular vectors for this subspace. A similar calculation can be done for Sz. The remaining vectors are orthogonal both S₁ and S2 and thus take the form of
t P1c1 P262,
14 P3c1 P402 for 1, with P1. Ps. R.
By solving the equation X Xu; Auy, we can show that the corresponding singular values and coefficients P1.P2, Ps, P4 coincide with those in the SVD of M, as can be verified by simple calcu- lations. Thus, we conclude the proof of Proposition J.2.
Proposition J.3. Let (a I[. b. Ik-L]T= R and X diag(r) K- TT RRK where a, b 0. Then the SVD of X UEVT is that
22 (K - L) 12
Σ diag
a
1.
K
R1,--1
OL, --1 Rk L, b-1
V
0K-LL-1
, 11,,
R1AL-1
OL,K I--1 RK-L,K-L-
U
OK-LLL11
U1, U2
Here, the columns of the matrix RL E (L-(1)) form orthonormal basis for the subspace of vectors in R orthogonal Tt. Similarly, the columns of Rx RK-LKK L,K --11 ((K-))((-L-I) form
35

---

Published as a conference paper at ICLR 2026
an orthonormal basis for the subspace of vectors in RK-1. orthogonal to 1 These correspond to the subspaces S1 and S2 defined as:
S - 1
x OK-L. (
1 x I = O, and æ E ER.
S2 =
1 [0 [00
y M = O, and y E IR K-
1
The vectors U1, U2, 11 U2 are
av K
I1
bVL
[ 1)
O2
= Vo2(KK L) + DLL
VL
0K-L
VK L
[ 1
Or.
0/
Va²(K L) + B2L
+
[K-L]
1
(K (K- L)
I
L
O2
VKL(K-1
0K-L
[K-L
U2
VK 1K.
Proof of Proposition J.3. This proposition is direct corollary of Proposition J.2. The matrix X diag(#) K I x x. T is an instance of the general form A + C from Proposition J.2.
The diagonal part is A = diag(æ) diag(a Iz, b --L)). The off- off-diagonal part is C K -IK
T.. We can write C' block form:
C
1
bI k- -
1 aJLL bJLKK-L K aJK-L, bJK-LK-L
K
[K-L
This corresponds to setting the block-wiss wise constants Proposition J.2 to:
C11 aa/K, 612 bb/K, C21 -a/K, 022 b/ K.
Substituting these into the formulas for a. 7, y, from Proposition J.2 gives:
a a T aa (/K)/KK--- = a(K L)/K
V VL(K - L)(-b/K)
V L(K aa)
8 b (K (L)(-bKK) bL/K
These coefficients define the 2 matrix M from Proposition J.2 for this specific case. We now analyze this matrix M. A key observation is that its determinant is zero:
(M)(M) = =ō - - a(K L)bL PL K
(LI
(-b)(-a)=0
K
K2
Since the determinant zero, one its singular values must zero. The other singular value, $1 be calculated from the squared Frobenius norm (sum of squares of elements), which is also the sum of squared singular values (sf + $3);
2
K 2K
1)2
L(K
L)D2
L(K
L)a²
121,2
0 =a= =a²
+ =a²+
+- +
+
+
K2
K²
K2
K2
a² K
L.) + 1T.
K
This confirms the singular values stated in the proposition. The singular vectors U1, U2, 01, U2 be derived by performing the SVD on this specific 2 x matrix M.
K HEAVY-TAILEDNESS OF GRADIENT OF LLMS
In this section, we discuss how our insight about the gradient in the one-layer model generalizes to the multi-layer model. In the following analysis, we focus the FFN modules in the model, and
36

---

Published as a conference paper at ICLR 2026
the attention module be similarly analyzed. The illustration of this multi-layer model of FFN modules is shown the following Figure 19.
(ww(EW)
LM Head
y sm(E x)
Later Layers y h(W ((W)(F))
Residual
Current FFN
Residual
= Wout 9 (Ex), g(E,) = a (x)
Ek
Previous Layers
Figure 19: The illustration of the multi-layer model of FFN modules.
We abstract the feature extraction role of all the previous layers and the in-projection of the current FFN as function as follows.
V W omn (Win>) (I)
where W ou,, Wi E Rdxd are weight matrices, E IRd is the output of all the previous layers, and g: Rdd + Rd abstracts the role of feature learned in Win Abstracting all the later layers function h R. Rd, the function h may also take all previous tokens as inputs, which we omit from the notation for brevity. The whole model written as
fw (Ex) = sm E Wo((E)+ ((WO())(()(E))))
where É dR is the parameter of the language model head, K is the alphabet size, and E& is the hidden state of the last token in the training context that precedes the k-th token the alphabet, at the layer where associative memory is present. Without loss of generality, assume that the next token the k-th token in E. Then the loss function on the pretraining data is
K
L(W)
Σ fw (Ex))) k=1
where Pk is the frequency of k-th token. We that this simplification of what happens in the pretraining, where the frequencies of token associations instead of the single token matter. However, such simplification does not influence our main message. In the heavy-tailed dataset, egg. WikiText1033, py decays Pk a 1-1 for k E [K]. Then the gradient of W is
K
ph V (osl/w((x)lle VwC(W) = Σ
1=
K - (EW)s E.o(Ex)T}
Vw log[sw (Ex)la- (la.d + (Wg(EE)))
37

---

Published as a conference paper at ICLR 2026
where Ia Rddd is the identity matrix and Jn is the Jacobian of the function h. Two structural properties of the gradient are worth highlighting. First, Vw L(W) is heavy-tailed, since is a weighted sum of per-token gradients with geometrically decaying weights Pa = a- for & E [K]. Second, the gradient decomposes as a sum of outer products Ex g E&) T. Our theoretical analysis in Section 4 focuses on the simplified setting h = 0 and g(x) = For general h and g. the Jacobian Jn acts as preconditioner on the gradient, and each outer product is formed between the feature in
the language-mode head E and the transformed feature g in each layer. Thus, our intuition extends to this more general multi-layer setting.
38


Reviews:
[{"review_id": "aOUTky0JC7", "reviewer": "Reviewer_E1d8", "summary": "The submission investigates the reason for the improved performance of Muon over Adam on transformers. It establishes empirically that the weights of a transformer architecture that benefit most from Muon-style updates are the non-key/query parameters of the attention modules and the parameters of the feedforward networks. To explain this behaviour, the authors propose an explanation based on different frequencies of \"fact\" representations in an associative memory module, and show theoretically that Muon-style updates lead to more uniform convergence across different frequencies.", "questions": "**Why does Muon help more on the VO/FFN than on the QK weights?**  \nThis is the main distinction found in the first set of experiments. \nThe theory tries to explain why the VO and FFN weights are helped by Muon by drawing on the connection to associative memories, \nbut I have some reservations about this explanation.\n\nThe QK weights are also part of the attention mechanism and form an outer-product based on the representation of the tokens. Why would the associative-memory explanation not apply to them as well? The FFN weights do not seem to follow this outer-product structure, so why is the associate-memory explanation relevant for them?\n\nIt might be that the associative-memory explanation is most relevant for the VO weights and do not explain the benefits on FFN weights or the reduced benefits on QK weights.\nBut if that is the case, this should be acknowledged more explicitly in the submission.\n\nBeyond this theoretical explanations, it would be helpful for future work if the authors could look at the experiments and data they have collected to see if there is some empirical difference between the QK weights and the VO/FFN weights that could explain the different behavior, for example a different spread in the spectrum of the singular values of the inputs of each layer? This is not strictly necessary but even a speculative explanation (based on some empirical observation) might be helpful for future work,\n\n**Why the focus on Gated vs. Ungated FFN?**  \nIs there a-priori a reason to believe that Muon would behave differently? \n\n**How is the step-size set in Figure 1?**  \nIt is not clear to me whether the results presented in Figure 1 use the same step-size for both Adam and Muon, \nor use a tuned step-size for each optimizer. Could the authors clarify this point?\n\nThe later would seem more appropriate as it is not clear that the same step-size is optimal for both optimizers, especially when used at the same time for different weights.\nAlthough it would then require a comparison against using Adam for both subset of weights with a different step-size for each subset to account for the possibility that the improvement comes from using different step-sizes.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "To my knowledge, this is the submission that attempts to provide empirical evidence for the source of the benefit of Muon updates over Adam on language models. The identification of specific weights that benefit most from Muon and the corresponding explanations help us understand the difficulties of optimization in transformer-based architecture, and hopefully design better methods. The experiment on the synthetic dataset with imbalanced facts is well-designed an provides good evidence for the claim that Muon helps most on associative memories. \n\nI am generally positive about the submission. My initial score for the review is low due to the issues listed below. Those should be addressable with some revisions to the presentation and an additional experiment. I am expecting to increase my score if that can be added.", "weaknesses": "My main issue is with the presentation of the results, specifically in Figure 1, 2, 4 and §4.1.Those are dense and require a lot of care on the reader's part to understand what is being shown. This complexity does not appear necessary. A pass on the issues are detailed below could improve the clarity of the presentation significantly.\n\nOn the evidence, the main argument of the submission, that Muon helps with imbalance in \"facts\" stored in an associative memory. It appears difficult to distinguish between this explanation and the more general imbalance in token frequencies in text data for the success of Adam. See below of a possible additional experiment that would strengthen the argument of the submission.\n\nThe submission would be much stronger with those changes, and I would be happy to increase my score if they are addressed.\n\n## Distinguishing between frequency imbalance in \"facts\" vs. general imbalance in text\n\nThe argument presented by the submission is that the benefit of Muon over Adam stems from an imbalance in the frequencies of different \"facts\" stored in an associative memory.\nThe controlled experiment with a synthetic dataset in §4.1 is a very nice illustration of this point. However, we could also expect this imbalance to stem simply from the distribution of tokens in text data, which follow Zipf's law. Those have been shown to lead to a better performance of Adam over SGD (Kunstner et al., 2024) primarily due to the update in the first and last embedding layers (Zhao et al., 2024; Zhang et al., 2024). One hypothesis could be that Muon helps on additional layers for the same issue of imbalanced tokens, not specifically facts in associate memories. Adam cannot correct these imbalances beyond the embedding layers because they are no longer axis-aligned, while Muon could. \n\nThe main difference I am drawing between those two hypotheses is the whether the main difficulty is due to an imbalance in token frequencies or an imbalance in the frequencies of \"facts\" to be stored in associate memories. As facts are also represented as tokens, there is some overlap between those ideas, but there is a way to disentangle these two explanations. If the arguments of the submission is correct, that Muon target imbalances in \"facts\" and associate memories beyond mere imbalance in token frequencies, I would expect the performance gap between Muon and Adam to decrease as the fact imbalance is reduced. The synthetic dataset prepared by the authors could be modified to exhibit various levels of imbalance, varying the imbalance continuously between the approximate power-law distribution in Figure 3 to a uniform distribution. The imbalance in token frequencies should still follow a power-law (assuming the different templates are rich enough to approximate natural text, which should be checked), but the imbalance in \"facts\" would be controlled. If the performance gap between Muon and Adam shrinks as the \"fact\" imbalance becomes more uniform, this would give a lot more strength to the proposed fact-based explanation.\n\n\n\n## Presentation of the results\n\nThe specific recommendations below are only one possible way to improve the presentation, meant to help communicate the difficulties I encountered in understanding the results. The authors know their work best and might have other ideas on how to improve the presentation that are more appropriate. \n\n\n**Figure 1**:  \nThere are too many lines to be able to easily tell which combination of architecture/optimizer is performing best.\nThis detailed plot could be useful to readers wanting to see the behaviour over time, \nbut the main point of Figure 1 would be better illustrated by a simple bar plot of the final performance, \nusing more legible distinctions (colors/makers) or having the description of the setup directly above the bar. \n\nThe figure also combines gated and non-gated variants of the architecture, \nbut it is not clear why this distinction is important for the point being made.\nThe plot could be made much simpler by focusing on one of the two variants, \nand pointing the reader to \"the same happens on the gated version, see Appendix X\".\n\n**Figure 2**  \nsuffers from a similar issue.\nThe plot shows 4 metrics measuring the spread of the eigenvalues, \nsome which should be small while other should be large, \nfor 2 weights (VO, $W_{\\mathrm{out}}$) on 2 architectures (gated, non-gated).\nThis makes it hard to extract the main point of the figure at a glance.\nA simpler variant of the figure could focus on one architecture and one metric,\nrelegating the others to the appendix with the comment that a similar behaviour is observed,\nor use this space to show the full distribution of the spectrum over time.\n\nAn interesting addition to Figure 2 could be to show the QK weights. The Muon update is apparently less helpful for those weights; does the spectrum give a hint as to why?\n\n**Figure 4 (b, c)**:  \nI had a hard time understanding those figures, and what those metrics tell us about the update. \nI am not sure which is the best way to improve the presentation, but here are things I found confusing: \n\nThe figure 4b is implicitly parameterized by a step-size which is being varied. This can only be understood from the main text (not the caption), which should be fixed. Even with this information, it is not easy to see the \"step-size\" on the plot. This might be because the lines \"start\" at the bottom right ($\\eta = 0$ is at $(10^{.8}, 0)$) and go to the top left as $\\eta$ increases, going against the standard left-to-right/top-to-bottom reading direction. Flipping the x and y axes might help. Another option could be to use a plot with 2 y-axes (or two subplots on top of each other) that would show the $\\Delta$ vs. step-size and the corresponding loss vs. step-size. The same treatment could be applied to Figure 4c (with number of updates as the x-axis). Currently, 4b and 4c have a different experimental setting (varying step-size vs. varying time) but it is not clear from the figures/caption that they differ in this way.\n\n**Theoretical results in §4.2**  \nAssumption 4.2 seems overly complex for its uses in Theorem 4.3. The parameters $\\alpha$ and $\\beta$ appear to impact the results only through the ratio $r(\\alpha,\\beta)$. Since it happens to be $\\frac{\\min_i p_i}{\\max_i p_i}$, wouldn't it be easier to just define the degree of imbalance as $r = \\min_i p_i / \\max_i p_i$ directly, and avoid the extra notation of $\\alpha, \\beta$? If not, writing the overall loss after Assumption 4.2 as a function of $\\alpha, \\beta$ might be helpful.\n\nThe only place where I see their use is in the proof for Adam, but the range of behavior observed under the $(\\alpha,\\beta)$ assumption would also be observed on a model with imbalance ratio $r$ since it's strictly less restrictive. If I'm missing something and the $\\alpha, \\beta$ are important for the proof, they should still be mentioned, but an explanation of what $r(\\alpha,\\beta)$ is should be given around Assumption 4.2 to help the reader.\n\nCould some of the notation be simplified by assuming that the frequencies are sorted? For example $\\min_{k\\in[K]}$ and $\\max_{k\\in[K]}$ in Eq. 4.1 could be replaced by $p_K$ and $p_1$ if we assume the probabilities are decreasing.\n\nThe description of the stopping condition would also benefit from a slightly more wordy explanation. The start of §4.2 takes a while to parse. Something like \"For each optimizer, we take one step with step-sizes $\\eta$, selected such that the accuracy of the most frequent class reaches $1-\\epsilon$. We then look at the accuracy achieved on the least frequent class.\"\n\n\n## Additional references \n\nThe submission should cite and discuss the following works which have done similar ablation studies on Adam to justify its performance on language models and identify that they help most with the first and last embedding layers. This does not detract from the contributions of the submission, but those works should be acknowledge.$\n\n- Deconstructing What Makes a Good Optimizer for Language Models\n  Rosie Zhao, Depen Morwani, David Brandfonbrener, Nikhil Vyas, Sham Kakade\n  https://arxiv.org/abs/2407.07972\n- Adam-mini: Use Fewer Learning Rates To Gain More\n  Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Diederik P. Kingma, Yinyu Ye, Zhi-Quan Luo, Ruoyu Sun\n  https://arxiv.org/abs/2406.16793\n\n\n## Minor/Typos\n\n- \"intrinsically heavy-tailed: a few classes (tail classes) appear far less frequently than others\"\n  Is this a typo? Presumably if only a few classes are rare, their impact on the overall loss is small and Muon would not help much.\n- It is not clear where Theorem 4.3 ends. Could the paragraph breaks \"For GD, ... For Muon, ... For Adam, ...\" be placed in an itemize block to help that?\n- \"highly sensitive\" as a description of Adam's behavior in Figure 2 seems exaggerated. \"More sensitive\" would be appropriate.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LSAZsoZet1", "reviewer": "Reviewer_W4iz", "summary": "The paper studies what transformer components benefit most from Muon as compared to Adam. The authors find that the main benefit comes from the VO and FFN blocks, which correspond to the associative memory stores. The authors use a heavy-tailed knowledge task to compare Adam and Muon and provide theoretical insights on one-layer models.", "questions": "- Do you have any intuition on why benefits from Muon seem to diminish with scale (/if this pertains to associative memory)?\n- It seems like the Newton-Schulz application could be an interesting knob to turn here -- have you tried varying the Newton Schulz coefficients/iterations to vary how much of the spectrum Muon is capturing and test how this affects associative memory?\n- What LR schedule did you use for Adam and Muon respectively?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The imbalanced knowledge experiment is a nice set up and the ablations over which layers most benefit from Muon are interesting\n- Provides an interesting connection between optimizer performance and prior work on associative memory \n- The writing and research motivation are clear", "weaknesses": "- Experiments are limited to a synthetic heavy-tailed knowledge task. It would be interesting to see in addition a task that relies heavily on QK to support the claim that Muon is not as important to apply to QK.\n- Language modeling results are limited to 160M parameter transformers. It would be interesting to see the main ablation (Muon only on VO+FFN) at a larger scale and performance measured on downstream tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gaEnrFarVx", "reviewer": "Reviewer_Yhie", "summary": "The authors mainly did an abalation study of Muon. They showed two major and novel findings: 1. applying muon only on VO+FFN can recover the full muon training. 2. Muon narrows the head-tail performance gap when learning highly imbalance datasets. They also provided a preliminary theory to explain their findings.", "questions": "N/A", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The ablation study of this work is sound to me. Muon only works for VO and FFN. This finding is somewhat interesting and might inspire following works to promote Muon's ability over QK.\n2. The heavy-tail findings is also interesting, which promotes the appliablity of Muon in practice.", "weaknesses": "**The novelty and contributions of this paper could be further improved.**\n- Three observations are provided. Among these, the second observation is already presented in Liu et al. (2025). Though the differences between them are noted by the authors, such as Liu et al. (2025) conduced experiments on MoE models whereas this paper focues on Dense Models, these reasons cannot convince me about the novelty. \n- The first observations are good, yet I expect more analysis. A central question is \"Why Muon is more effective on VO+FNN than QK\", yet this paper didn't address this question. They attribute this phenomenon to the associative memory property of VO and FFN, but the causal relationship between these is not clear. Is the heavy-tail task ability of Muon relate to the effectiveness of Muon on VO+FNN?\n- In addition, via the third observation, the authors further highlighted the connection between Muon and the associative memory. However, I also doubt that if the third observation can fully supports that Muon acquires knowledge more evenly than AdamW. Muon indeed prompts the performance on heavy-tail tasks, but it might restrict to knowledge acquring tasks. I expect more experiments across various tasks.\n\n**The toy theory seems to be intuitive**\n- The theory seems to be intuitive. Muon produces updates evenly spread over all directions compared to GD and SignGD. Then Muon must give even/uniform weight matrices, and thereby in this simple associative memory modeling, mitigating the unbalance learning problem.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "k3g3FOPFeQ", "reviewer": "Reviewer_sugx", "summary": "The paper studies why Muon often outperforms Adam in LLM-pretraining. The central claim is that Muon’s advantage concentrates in associative-memory components, specifically the Value/Output (VO) attention matrices and FFNs. Muon's update rule yields a more isotropic singular spectrum, which in turn improves learning on tail classes in heavy-tailed distributions. The authors support this with \n1. component-wise optimizer ablations in Transformers,showing that VO+FFN accounts for most of the validation-loss gains for Muon.\n2. measurements of more isotropic singular-value spectra under Muon than under Adam. \n3. a knowledge-intensive QA task reflecting heavy-tailed data where Muon shows larger gains on tail examples.\n4. a theoretical analysis of a one-layer linear associative memory under class imbalance showing Muon’s superiority.", "questions": "1. **Capacity control** If you equalize parameter counts (e.g., only optimize $W_{in}$ parameter in some depth by Adam), do Muon-over-Adam gaps persist in QK vs. FFN?\n2. **Multiple kinds of heavy-tailness** Can you connect the heavy-tailness of data and knowledge with the heavy-tailness of parameter and update in the context of LLM?\n3. **Logic gap on associative-memory mechanism** The paper does not yet provide strong causal evidence that Muon surpasses Adam *specifically* because of associative-memory effects. The superiority on the knowledge-intensive QA task might not be attributable to associative memory per se. VO+FFN may contribute more ('may' is due to weakness 2 and 4), but the mechanistic link to associative memory remains suggestive rather than conclusive.\n\nOverall, the paper discusses a vital problem in the optimizer community. However, currentlymy tentative recommendation is borderline Weak Accept/Accept. I’m open to revising my score if these concerns are addressed.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. **Focused component ablations.** The blockwise ablation is informative and directly isolates VO/FFN as primary beneficiaries, supporting the associative-memory lens.\n\n2. **Heavy-tail perspective with tail-class gains**. The paper emphasizes tail performance and links it to spectral isotropy, providing a plausible causal chain: Muon → more balanced singular spectrum → better tail learning.\n\n3. **Theory that matches the story.** The one-layer associative-memory model with class imbalance theoretically explains why Muon outperforms Adam under heavy-tailed class distributions, through the model is simple.\n\n4. **Muon wins over Adam, and community does not know why!**\n5. **LLM settings are diverse.** Covers both gated and non-gated 160M/700M LLM.", "weaknesses": "1. **Multiple kinds of heavy-tailness** The paper presents servral kinds of heavy-tailness: data, knowledge, parameter and update. There are evidence in the paper to connect the heavy-tailness of data with the heavy-tailness of knowledge, there are also evidence to connect the heavy-tailness of parameter with the heavy-tailness of update. However, it remains unclear how to connect the heavy-tailness of data and knowledge with the heavy-tailness of parameter and update in the LLM.\n2. **$W_V$ as associative memory** Prior work rarely frames $W_V$ as an associative-memory parameter. Although the paper argues $W_V$ and $W_O$ play symmetric roles, that symmetry breaks under MQA/GQA. Treating VO jointly is understandable as a number of literatures[2][3] say that VO(Value-Output) are similar and can be considered together in practice even under the context of MQA/GQA, but a restatement is needed. \n3. **Observation 2 seems to be known in literature** The figure 4 in [3] already points out Observation 2 that Muon consistently yields more isotropic weight matrices with broadly distributed spectral energy than Adam.\n4. **FFN over QK may be the result of parameter counting** FFN over QK may be the result of parameter counting, not because of associative memory.\n5. **Reproducibility** The `model` are missing from the released code, making the work not fully reproducible.\n\n[1] Lin, C. H., Gao, S., Smith, J. S., Patel, A., Tuli, S., Shen, Y., ... & Hsu, Y. C. MoDeGPT: Modular Decomposition for Large Language Model Compression. In The Thirteenth International Conference on Learning Representations. \\\n[2] Wang, J., Wang, M., Zhou, Z., Yan, J., & Wu, L. The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Training. In Forty-second International Conference on Machine Learning. \\\n[3] Liu, J., Su, J., Yao, X., Jiang, Z., Lai, G., Du, Y., ... & Yang, Z. (2025). Muon is scalable for LLM training. arXiv preprint arXiv:2502.16982.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)