Muon Learns Facts Better: Understanding the Role of Spectral Orthogonalization

Xuheng Li; Qiwei Di; Yuan Cao; Quanquan Gu

cs.LGOct 2, 2026

Muon Learns Facts Better: Understanding the Role of Spectral Orthogonalization

Authors: Xuheng Li, Qiwei Di, Yuan Cao, Quanquan Gu

Organizations: Department of Computer Science, University of California, Los Angeles, CA 90095, USA · School of Computing and Data Science, the University of Hong Kong

Abstract

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 $\widetildeΘ(\sqrt{S/R})$ , whereas Spectral GF reduces this ratio to $\widetildeΘ(1)$ . In addition, for fixed $S$ and $R$ , the subject- and relation-dependent errors decay as $1/(T\log T)$ in training time $T$ under GF, but as $\exp(-\mathrm{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

Panel of Figure 1 — 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 $\widetilde{\Theta}(\sqrt{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.

Optimization dynamics	GF	Spectral GF	Sign GF (O, O, O)	Sign GF (H, O, H)	Sign GF (O, H, H)
$T_{S}^{\delta}$	$\widetilde{\Theta}(S\sqrt{R}\cdot{\color[rgb]{1,0,0}\delta^{-1}})$	$\widetilde{\Theta}((SR)^{1/4})$	$\widetilde{\Theta}(1)$	$\widetilde{\Theta}((SR)^{1/4})$	$\widetilde{\Theta}((R/S)^{1/4})$
$T_{R}^{\delta}$	$\widetilde{\Theta}(R\sqrt{S}\cdot{\color[rgb]{1,0,0}\delta^{-1}})$	$\widetilde{\Theta}((SR)^{1/4})$	$\widetilde{\Theta}(1)$	$\widetilde{\Theta}((S/R)^{1/4})$	$\widetilde{\Theta}((SR)^{1/4})$
$T_{S}^{\delta}/T_{R}^{\delta}$	$\widetilde{\Theta}(\sqrt{S/R})$	$\widetilde{\Theta}(1)$	$\widetilde{\Theta}(1)$	$\widetilde{\Theta}(\sqrt{R})$	$\widetilde{\Theta}(1/\sqrt{S})$

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\geq R\geq 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

T_{S}^{\delta}

and

T_{R}^{\delta}

to denote the learning times when the subject- and relation-dependent components of the prediction reach the target accuracy of

1-\delta

under assumptions about the model initialization, and the ratio

T_{S}^{\delta}/T_{R}^{\delta}

measures the duration of the feature-separation phase.

$\bm{\phi}_{S}$	$\bm{\phi}_{R}$	$\bm{\phi}_{A}$	$\dot{\alpha}_{S}$	$\dot{\alpha}_{R}$	$\dot{\beta}_{S}$	$\dot{\beta}_{R}$	$T_{S}^{\delta}$	$T_{R}^{\delta}$	$T_{S}^{\delta}/T_{R}^{\delta}$
O	O	O	$2\sqrt{R}$	$2\sqrt{S}$	$1$	$1$	$\widetilde{\Theta}(1)$	$\widetilde{\Theta}(1)$	$\widetilde{\Theta}(1)$
O	O	H	$\sqrt{S}$	$\sqrt{R}$	$1$	$1$	$\widetilde{\Theta}((R/S)^{1/4})$	$\widetilde{\Theta}((S/R)^{1/4})$	$\widetilde{\Theta}(\sqrt{R/S})$
O	H	O	$2\sqrt{R}$	$\sqrt{SR}$	$1$	$1/\sqrt{R}$	$\widetilde{\Theta}(1)$	$\widetilde{\Theta}(1)$	$\widetilde{\Theta}(1)$
O	H	H	$\sqrt{S}$	$1$	$1$	$1/\sqrt{R}$	$\widetilde{\Theta}((R/S)^{1/4})$	$\widetilde{\Theta}((SR)^{1/4})$	$\widetilde{\Theta}(1/\sqrt{S})$
H	O	O	$\sqrt{SR}$	$2\sqrt{S}$	$1/\sqrt{S}$	$1$	$\widetilde{\Theta}(1)$	$\widetilde{\Theta}(1)$	$\widetilde{\Theta}(1)$
H	O	H	$1$	$\sqrt{R}$	$1/\sqrt{S}$	$1$	$\widetilde{\Theta}((SR)^{1/4})$	$\widetilde{\Theta}((S/R)^{1/4})$	$\widetilde{\Theta}(\sqrt{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

\widetilde{\Theta}(\cdot)

in the learning times hides logarithmic factors of

\delta^{-1}

as well as

S

and

R

Panel of Figure 2 — Figure 2: Scaling of learning-time ratio with $S$ .

Panel of Figure 3 — Figure 3: Scaling of learning-time ratio with $R$ .

Panel of Figure 4 — Figure 4: Scaling of learning times with $\delta$ .

Panel of Figure 5 — Figure 5: Scaling of learning-time ratio with $S$ .

Panel of Figure 6 — Figure 6: Scaling of learning-time ratio with $R$ .

Panel of Figure 7 — Figure 7: Scaling of learning times with $\delta$ .

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

Neural network training dynamics studies how parameters and representations change during optimization and how those changes shape learning, stability, and generalization.

See all 507 papers

Reassessing Muon for Matrix Factorization
Jul 14, 2026Ali Parviz, Gal Mishne, Alex CloningerDeep Learning Optimization Muon Optimizer
†. Halicioğlu Data Science Institute, UC San Diego. · ∗. Department of Mathematics, UC San Diego.
Spectral Flattening Is All Muon Needs: How Orthogonalization Controls Learning Rate and Convergence
May 13, 2026Tien-Phat Nguyen, Truong Nguyen, Minh-Phuc Truong +3Momentum Methods Gradient Descent
Hanoi University of Science and Technology Hanoi, Vietnam · Indiana University 107 S. Indiana Ave, Bloomington, IN 47405, USA · Monash University Clayton, VIC 3800, Australia
Spectrally Targeted Muon
Oct 7, 2026Vishrut Goyal, Rohan RamkumarSpectral Optimization Deep Learning Optimization
Department of Computer Science, Stanford University · Department of Electrical Engineering, Stanford University.

Jul 14, 2026J/K move · Enter open · S save

Reassessing Muon for Matrix Factorization

Ali Parviz, Gal Mishne, Alex Cloninger

†. Halicioğlu Data Science Institute, UC San Diego. · ∗. Department of Mathematics, UC San Diego.

Deep Learning Optimization Muon Optimizer Adaptive Gradient Methods

Open paper AskPDF arXiv

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.

Abstract

Neural Network Training Dynamics

Explore similar work

Reassessing Muon for Matrix Factorization

Spectral Flattening Is All Muon Needs: How Orthogonalization Controls Learning Rate and Convergence

Spectrally Targeted Muon