Organizations: Department of Computer Science, University of California, Los Angeles, CA 90095, USA · School of Computing and Data Science, the University of Hong Kong
The Muon optimizer applies spectral orthogonalization to matrix-valued updates and has shown strong performance in large-scale neural network training, yet the mechanisms of this transformation in feature learning remain poorly understood. In this work, we investigate this question through a tractable factual-recall model, where a fact maps each subject-relation pair to an answer, and a linear transformer learns the subject- and relation-dependent information required to recover this mapping. The transformer is optimized with gradient flow (GF), spectral GF, or Sign GF, which are continuous-time limits of gradient descent, Muon, and Adam, respectively. Prior studies (Nichani et al., 2025) have shown that when the number of subjects exceeds the number of relations, GF learns relation-dependent information before subject-dependent information, producing a feature-separation phase during training. We characterize this separation with the learning times when the subject- and relation-dependent components of the prediction reach a target accuracy. With S subjects and R relations, GF has a learning-time ratio of Θ(S/R), whereas Spectral GF reduces this ratio to Θ(1). In addition, for fixed S and R, the subject- and relation-dependent errors decay as 1/(TlogT) in training time T under GF, but as exp(−poly(T)) under spectral GF. Finally, we show that GF and spectral GF are equivariant under orthogonal transformations of the token embeddings, whereas Sign GF is not: Different orthonormal embeddings can potentially produce no feature separation, a large feature-separation phase, or even a reversed learning order. These results provide a mechanistic view of how spectral orthogonalization can fundamentally reshape feature-learning dynamics.
Figures & tables
Figure 1: Illustration of feature separation and softmax saturation. (a): Subject vs. relation components of model prediction accuracy. The top left portion of the plot for GF stands for feature separation, i.e., high accuracy of relation component but low accuracy of subject component; plot of Spectral GF is closer to the diagonal, indicating weaker feature separation; different embeddings in Sign GF can cause strong feature separation, weak separation, or reversed learning order (bottom right). (b, c): Dynamics of subject and relation component accuracies in GF and Spectral GF. For (S,R)=(64,16) , the learning-time ratio of GF is approximately 2 (consistent with theoretical prediction Θ(S/R) ); the ratio is close to 1 in Spectral GF. Additionally, GF slows down when entering the high-accuracy regime due to softmax saturation, while Spectral GF does not.
Table 1: Comparison of learning times of GF, Spectral GF, and Sign GF dynamics in the factual recall task with S subjects and R relations, assuming S≥R≥4 . The triplet after Sign GF stands for the embeddings corresponding to the subject, relation, and answer tokens: “O” and “H” stand for one-hot and Hadamard embeddings, respectively (See Section 5.1 for definition). We use TSδ and TRδ to denote the learning times when the subject- and relation-dependent components of the prediction reach the target accuracy of 1−δ under assumptions about the model initialization, and the ratio TSδ/TRδ measures the duration of the feature-separation phase.
ϕS
ϕR
ϕA
α˙S
α˙R
β˙S
β˙R
TSδ
TRδ
TSδ/TRδ
O
O
O
2R
2S
1
1
Θ(1)
Θ(1)
Θ(1)
O
O
H
S
R
1
1
Θ((R/S)1/4)
Θ((S/R)1/4)
Θ(R/S)
O
H
O
2R
SR
1
1/R
Θ(1)
Θ(1)
Θ(1)
O
H
H
S
1
1
1/R
Θ((R/S)1/4)
Θ((SR)1/4)
Θ(1/S)
H
O
O
SR
2S
1/S
1
Θ(1)
Θ(1)
Θ(1)
H
O
H
1
R
1/S
1
Θ((SR)1/4)
Θ((S/R)1/4)
Θ(R)
Table 2: Summary of dynamics and learning times for the eight possible combinations of embeddings. “O” and “H” stand for the one-hot and the Hadamard embeddings, respectively. The Θ(⋅) in the learning times hides logarithmic factors of δ−1 as well as S and R .
Figure 2: Scaling of learning-time ratio with S .
Figure 3: Scaling of learning-time ratio with R .
Figure 4: Scaling of learning times with δ .
Figure 5: Scaling of learning-time ratio with S .
Figure 6: Scaling of learning-time ratio with R .
Figure 7: Scaling of learning times with δ .
Figure 8: SGD, AdamW, and Muon in realistic factual recall task learned by Pythia 35M. The solid line stands for the overall error rate, the dashed line stands for the error probability of the prediction having the wrong subject entity, and the dotted line stands for the error probability of the prediction having the wrong relation entity.
Neural network training dynamics studies how parameters and representations change during optimization and how those changes shape learning, stability, and generalization.
Muon has recently emerged as a strong optimizer for large-scale deep learning, where it reshapes gradient updates through approximate orthogonalization and has been reported to outperform Adam and AdamW in large language model training. Its empirical success has motivated a growing body of theoretical work that interprets Muon as steepest descent under the spectral norm. Yet it remains unclear which of Muon's advantages stem from its update rule itself and which are artifacts of the scale, architecture, and data of modern deep networks. In this work, we isolate the optimizer from these confounding factors by studying Muon on a simple, well-understood, and spectrally structured problem: low-rank matrix factorization. Through a controlled comparison against carefully tuned adaptive baselines, we find that Muon does not consistently outperform AdamW in this setting and that several previously reported advantages are sensitive to hyperparameter choices. Our results provide a more nuanced picture of when spectrum-aware orthogonalization is beneficial and argue for evaluating modern optimizers on controlled problems in addition to end-to-end benchmarks.
Ali Parviz, Gal Mishne, Alex Cloninger
†. Halicioğlu Data Science Institute, UC San Diego. · ∗. Department of Mathematics, UC San Diego.
Muon orthogonalizes the momentum buffer before each update, replacing its singular values with ones via Newton-Schulz iterations. This simple change lets Muon tolerate far larger learning rates and converge faster than other optimizers, but why? We show that the mechanism is spectral flattening, and develop two results around it. First, we prove that Muon's maximal stable step size scales with the average singular value of the gradient rather than the largest, which bottlenecks standard gradient descent. Second, we recast Muon as a preconditioned gradient method and show, under a Kronecker-factored curvature model, that it improves the effective convergence factor, with the improvement controlled by the spectrum of the gradient covariance. Extensive experiments validate both results: Muon remains stable at learning rates that cause SGD to diverge within the first few iterations, and reaches accuracy milestones several epochs earlier even at identical step sizes. Taken together, our results offer a principled, geometric explanation for Muon's empirical success.
Hanoi University of Science and Technology Hanoi, Vietnam · Indiana University 107 S. Indiana Ave, Bloomington, IN 47405, USA · Monash University Clayton, VIC 3800, Australia
The Muon optimizer orthogonalizes each update matrix, setting all of its singular values to one, and has proven highly effective for training large language models. It remains unclear, however, whether this success comes from amplifying small singular directions that gradient descent neglects or from suppressing large, degenerate directions that disrupt training. We introduce Spectrally Targeted Muon, which orthogonalizes only the singular values above or below a threshold τ, so that varying τ interpolates between normalized SGD and Muon. It isolates the relevant singular subspaces with projections computed by Newton-Schulz iteration on a shifted Gram matrix, so no SVD is needed. We evaluate these variants on the CIFAR-10 and NanoGPT speedruns, tracking the effective rank of gradient, update, and weight matrices and a new metric, the alignment of updates with the tangent space of the weight matrix's isospectral manifold. We find three things. First, the small singular values of the momentum are not noise. Orthogonalizing everything except the few largest singular values of each matrix nearly matches Muon while touching only a small fraction of the momentum, whereas orthogonalizing only the top falls well short even though it holds almost all of it. On language models every momentum singular value is far below one, so targeted orthogonalization can only amplify, and Muon wins by making directions that are too small to train on at their raw scale trainable. Second, shrinking the largest singular values is what keeps the parameter spectrum flat; this is cheap, and it is not what drives the loss. Third, AdamW differs from Muon mainly in how slowly it builds structure, which explains its slower start and why warmup helps AdamW but only hurts Muon.
Vishrut Goyal, Rohan Ramkumar
Department of Computer Science, Stanford University · Department of Electrical Engineering, Stanford University.