Muon Learns Facts Better: Understanding the Role of Spectral Orthogonalization
cs.LG, math.OC, stat.ML
Submitted: 2026-10-02
Updated: 2026-10-02
Comments: 47 pages, 8 figures
License: http://creativecommons.org/licenses/by/4.0/
The gist: 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
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 Θ(sqrt 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/(T T) in training time T under GF, but as (-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.
Sources
- Old Optimizer, New Norm: An Anthology
- Learning Associative Memories with Gradient Descent
- On the Implicit Bias of Adam
- To Use or not to Use Muon: How Simplicity Bias in Optimizers Matters
- Implicit Bias of Spectral Descent and Muon on Multiclass Separable Data
- Risk and parameter convergence of logistic regression
- SOAP, Muon, and Beyond: Pushing LLM Pretraining Scales
- Sharp Capacity Scaling of Spectral Optimizers in Learning Associative Memory
- Understanding Gradient Orthogonalization for Deep Learning via Non-Euclidean Trust-Region Optimization
- Muon in Associative Memory Learning: Training Dynamics and Scaling Laws
- Muon is Scalable for LLM Training
- Gradient Descent Maximizes the Margin of Homogeneous Neural Networks
- Muon learns balanced solutions in matrix factorization without slow saddle-to-saddle dynamics
- Muon Learns More Robust and Transferable Features than Adam
- On the Convergence Analysis of Muon
- Kimi K2: Open Agentic Intelligence
- Flavors of Margin: Implicit Bias of Steepest Descent in Homogeneous Neural Networks
- Muon Outperforms Adam in Tail-End Associative Memory Learning
- Implicit Bias of AdamW: $\ell_\infty$ Norm Constrained Optimization
- DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligence
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks