Source-linked AI summary
Muon Learns Facts Better: Understanding the Role of Spectral Orthogonalization
Xuheng Li, Qiwei Di, Yuan Cao, Quanquan Gu
TL;DR
Despite Muon’s empirical success, how spectral orthogonalization changes feature learning remains poorly understood. The paper analyzes factual recall under GF, Spectral GF, and Sign GF, finding that Spectral GF balances subject- and relation-feature learning while Sign GF’s dynamics depend on token embeddings.
Problem
Despite Muon’s empirical success, the mechanism of spectral orthogonalization in feature learning remains incompletely understood, while factual recall allows subject- and relation-dependent learning to be characterized separately.
Source: §1 · p. 2Method
The paper studies GF, Spectral GF, and Sign GF on factual recall, reducing the dynamics to an invariant low-dimensional manifold where predictions factor into subject- and relation-dependent components.
Source: §1 · p. 3Results
Spectral GF balances subject- and relation-feature learning and reduces their learning-time ratio, while Sign GF can exhibit varied learning dynamics under different orthonormal embeddings.
Source: §1 · p. 3Takeaways & Limitations
In this factual-recall analysis, spectral orthogonalization weakens feature separation and avoids the softmax-saturation-related δ−1 dependence in learning times.
Source: §1 · p. 3, §7 · p. 19The spectral GF analysis depends more directly on symmetric positive initialization; the effects of general initialization remain for future work.
Source: §4.3 · p. 12
Abstract
from arXiv · showhide
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{\Theta}(\sqrt{S/R})$, whereas Spectral GF reduces this ratio to $\widetilde{\Theta}(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.
1 Introduction
The paper studies how spectral orthogonalization changes feature learning in factual recall, where a prediction depends on both subject and relation information. It finds that Spectral GF balances their learning more than GF, while Sign GF can vary with token embeddings.
1 Introduction
- The factual-recall task maps subject-relation pairs to answers and separates the learning of subject- and relation-dependent information.Earlier work found that GF learns relation-dependent features before subject-dependent features when subjects substantially outnumber relations.Source: §1 · p. 2
- The study compares GF, Spectral GF, and Sign GF, continuous-time limits of gradient descent, Muon, and Adam, respectively.Muon applies spectral orthogonalization to momentum for matrix-valued parameters.Source: §1 · pp. 1, 3
- Spectral GF reduces the subject-to-relation learning-time ratio to eΘ(1), compared with GF’s eΘ(√(S/R)), and avoids the high-accuracy slowdown attributed to softmax saturation.The paper also reports that spectral normalization reduces GF’s δ^-1 dependence in learning times to a logarithmic scale.Source: §1 · p. 3
- Sign GF is not equivariant to orthogonal embedding transformations, so different orthonormal embeddings can produce strong or weak feature separation or reversed learning order.GF and Spectral GF are equivariant under those transformations.Source: §1 · p. 3
- The analysis identifies a low-dimensional invariant manifold on which predictions factor into subject- and relation-dependent components.This reduction enables analysis of high-dimensional attention-weight dynamics through a small set of scalar parameters.Source: §1 · p. 3
2 Related Work
Prior work connects spectral optimization to margins, convergence, and more balanced learning, while factual-recall research studies how transformers acquire facts. This paper focuses on how optimizer geometry shapes the acquisition of complementary subject and relation information.
2.1 Understanding the Muon Optimizer
- Prior analyses connect spectral optimization to operator-norm margins, optimizer geometry, convergence guarantees, and empirical behavior; this work complements them by studying finite-time feature acquisition.The cited margin analyses concern asymptotic implicit bias, unlike this paper’s learning times to a prescribed accuracy.Source: §2.1 · p. 4
- Related Muon studies report balanced learning across classes, principal components, low-frequency components, or spectral modes, while this work examines complementary features within one prediction.In factual recall, pairs and answers are sampled uniformly, yet each subject occurs less often than each relation when S > R.Source: §2.1 · pp. 4–5
- Muon’s balanced learning can benefit factual recall, but does not imply that removing simplicity bias universally improves generalization.Other studied settings report impaired shared-representation discovery or greater reliance on spurious features; factual recall instead requires both subject and relation information for the answer.Source: §2.1 · p. 5
2.2 Factual Recall and Associative Memory
- Unlike associative-memory studies of mappings stored in one matrix, this work analyzes attention weights that combine subject and relation information from a sequence.Its closest antecedent found relation-before-subject learning under GF; this paper compares GF, Spectral GF, and Sign GF and studies Sign GF’s embedding dependence.Source: §2.2 · pp. 5–6
3 Problem Setting
The section sets up a factual-recall task and linear-attention model to study how continuous-time optimizer dynamics learn subject- and relation-dependent information. It also establishes how orthogonal embedding changes affect those dynamics.
3.1 Data Model
- Facts bijectively map subject-relation pairs to answers, with uniform sampling isolating the effect of different subject and relation set sizes.Inputs place the subject and relation at distinct positions among noise tokens, followed by an EOS token; the next-token label is the mapped answer.Source: §3.1 · p. 6
3.2 Linear Attention
- A single-head linear attention model uses fixed orthonormal token embeddings and trains merged attention weights to predict the next token by minimizing cross-entropy.The trainable parameters merge the O and V weights and the K and Q weights.Source: §3.2 · pp. 6–7
3.3 Optimizers and Continuous-Time Dynamics
- GF follows negative loss gradients, Sign GF uses elementwise gradient signs, and Spectral GF models Muon through spectral orthogonalization.Sign GF is presented as an idealized continuous-time model of Adam; Spectral GF studies the ideal setting where spectral orthogonalization is exact.Source: §3.3 · p. 7
3.4 Orthogonal Equivariance and Embedding Dependence
- GF and Spectral GF dynamics correspond bijectively across orthogonal embedding transformations, whereas Sign GF is not generally equivariant under those transformations.The stated distinction follows from spectral orthogonalization's transformation property, which the elementwise sign operator does not generally have.Source: §3.4 · pp. 7–8
4 GF versus Spectral GF
The section reduces factual-recall training to a symmetry-based manifold and compares GF with Spectral GF there. Spectral orthogonalization balances subject- and relation-dependent learning, while GF slows at high accuracy due to softmax saturation.
4.1 Reduced Parameterization and Learning Times
- Symmetry motivates a low-dimensional manifold M, where predictions factor into subject- and relation-dependent components and the reduced dynamics exactly describe training.The dynamics preserve M when training begins on it.Source: §4 · p. 8, §4.1 · pp. 8–9
- The analysis assumes symmetric, positive initialization; matching initial magnitudes isolates learning-time differences due to the data and optimization dynamics.The positivity assumption avoids sign changes in the reduced dynamics and simplifies learning-time expressions.Source: §4.1 · p. 9
4.3 Spectral GF Dynamics
- Spectral orthogonalization reweights subject- and relation-dependent gradient factors and weakens softmax-saturation slowdown; the β parameter for the smaller margin can increase almost linearly.This makes learning-time dependence on the target error rate milder.Source: §4.3 · p. 12
4.2 GF Dynamics
- In GF, relation parameters receive larger updates, while the subject-dependent component slows at high accuracy because of softmax saturation.The GF analysis treats subject and relation parameters as separate systems.Source: §4.2 · pp. 10–11, §4.3 · p. 12
4.3 Spectral GF Dynamics
- Under Spectral GF, the relation learning-time transition is asymmetric and depends on whether target relation accuracy arrives before the subject component is well learned; subject learning time has no phase transition.The asymmetry follows from coupled βS and βR dynamics.Source: §4.3 · p. 13
- Spectral GF’s learning-time ratio is eΘ(1), with no polynomial dependence on S/R, unlike GF.The ratio is considered both for constant target error and as the target error tends to zero.Source: §4.3 · p. 13
5 Sign GF and Embedding Dependence
The section constructs orthonormal token embeddings from one-hot and Hadamard choices, then analyzes how those choices affect Sign GF. Across the eight combinations, Sign GF’s learning-time ratio can vary substantially, including a reversed learning order.
5.1 Orthonormal Embedding Constructions
- Each of the subject, relation, and answer token families can use one-hot or Hadamard embeddings, yielding eight combinations.The construction lifts the embeddings into orthogonal subspaces; the Hadamard construction assumes S and R are powers of two for convenience.Source: §5.1 · p. 13, Table 2 · p. 13, §5.2 · p. 14
5.2 Sign GF Dynamics
- For all eight embedding combinations, the manifold M is invariant under Sign GF, and Proposition 5.1 gives its parameter dynamics and learning times.The proposition’s learning-time result assumes sufficiently small initial α and β parameters.Source: §5.2 · p. 14
- For Sign GD, embedding choice affects the learning-time ratio: with certain embeddings it is eΘ(1), matching Spectral GF, while it can reach eΘ(A) or fall below 1, letting the subject reach target accuracy earlier.Assumption 4.2 serves to avoid sign changes in the manifold parameters but is not essential.Source: §5.2 · p. 14
- For the embeddings of interest, learning times depend only logarithmically on δ^-1; whether element-wise sign similarly weakens softmax saturation for general embeddings remains conjectural.The general-embedding question is left for future work.Source: §5.2 · p. 14
6 Experiments
Experiments test the predicted learning-time behavior in ODE simulations, single-layer attention training, and a realistic factual-recall task. Muon learns the subject-dependent component earlier in the realistic task, while the other experiments largely follow the predicted scaling patterns.
6 Experiments
- The experiments progress from ODE simulations on invariant manifold M to single-layer attention training and a small language model trained on synthetic biographies.The later stages introduce stochastic gradients, momentum, and trajectories outside M.Source: §6 · p. 14
6.1 (Normalized) Gradient Flow Dynamics
- The ODE experiments simulate the theoretical setting directly and focus on GF and Spectral GF dynamics.The learning-time ratio is plotted against S and R, and learning times are also examined as δ varies.Source: §6.1 · p. 14, §6.1.1 · p. 15, §6.1.2 · p. 16
6.1.1 Learning-Time Ratio Scaling with S and R
- In ODE simulations, the GF learning-time ratio versus S has a log-log slope near 0.5, while the Spectral GF ratio stays close to 1 across S and R.The experiments vary S and R from 4 to 1,024 and set δ = 0.1.Source: §6.1.1 · p. 15, Figure 2 · p. 15, Figure 3 · p. 15
6.1.2 Learning-Time Scaling with δ
- For fixed (S, R) = (64, 16), GF learning times scale with δ−1 for small δ and exhibit a phase transition for large δ; Spectral GF learning times depend logarithmically on δ−1.Source: §6.1.2 · p. 16
6.2 Single-Layer Attention Training
- With stochastic gradients, momentum, and off-manifold iterates, single-layer attention training reproduces the predicted learning-time-ratio scaling; SGD shows a δ phase transition, while other optimizers scale logarithmically with δ^-1.The setup uses the original factual-recall model and a cosine learning-rate schedule.Source: §6.2 · p. 16
6.3 Realistic Factual Recall in Transformers
- In Pythia’s 8,192-person factual-recall task, Muon’s subject error decays earlier than SGD’s and AdamW’s, although all three reach similar accuracies after 30,000 steps.Relation-component errors are much smaller than subject-component errors, leaving overall error close to subject error.Source: §6.3 · p. 18, Figure 8 · p. 19
7 Conclusion
The paper analyzes GF, Spectral GF, and Sign GF on an invariant manifold where predictions factor into subject and relation components. Spectral GF balances their learning rates, while Sign GF depends on the token embeddings.
7 Conclusion
- Spectral GF learns subject- and relation-dependent components at a more balanced rate than GF, weakening GF’s phase separation.Its spectral normalization also avoids the δ^-1 dependence in learning times caused by softmax saturation.Source: §7 · p. 19
- Sign GF can display diverse learning dynamics because its behavior depends on the token embeddings.The analysis includes a simple factual-recall data model and studies continuous-time dynamics with full gradient access.Source: §7 · p. 19
- The analysis is restricted to single-head linear attention with merged K-Q and O-V weights on an invariant manifold.The paper identifies generalizing these assumptions as future work.Source: §7 · p. 19
A Discussion about Learning Rate Schedule
A learning-rate schedule can be incorporated by reparameterizing time. In this time variable, scheduled dynamics become unit-learning-rate dynamics.
A Discussion about Learning Rate Schedule
- The time reparameterization expresses scheduled dynamics through a transformed time variable and transformed parameters.The supplied expressions define the transformed parameters as Θ composed with the inverse time map.Source: A Discussion about Learning Rate… · p. 20
- Under time reparameterization, dynamics with an arbitrary learning-rate schedule become unit-learning-rate dynamics.The schedule is incorporated through a change of time variable.Source: A Discussion about Learning Rate… · pp. 19–20
- The subject-relation learning-time ratio can be viewed as the ratio of cumulative learning rates needed to reach the target threshold.The passage describes this ratio as more intrinsic to the optimization dynamics than either learning time alone.Source: A Discussion about Learning Rate… · p. 20
B Proof of Structural Results
The proof establishes that orthogonal transformations preserve logits and relate the transformed gradients to the original gradients. It uses these properties to show that the transformed parameters follow GF and Spectral GF dynamics.
B Proof of Structural Results
- Under the orthogonal transformation, logits are invariant because the transformation matrix satisfies U⊤U = I.The proof uses the resulting equality of model predictions under the parameter transformation.Source: §B.1 · p. 20
- The transformed gradients for WOV and WKQ are obtained by conjugating the original gradients with Q and Q⊤.The proof states these gradient relationships for both parameter matrices.Source: §B.1 · pp. 20–21
- The transformed parameters follow GF and Spectral GF dynamics under their respective transformations.The proof explicitly concludes this for GF and then for Spectral GF.Source: §B.1 · p. 21
B.2 Proof of Proposition 4.1 and Proposition 4.3
The appendix derives gradient and model-output expressions, studies spectral normalization and element-wise signs, and reduces a margin derivative to a scalar expression for integration.
B.2 Proof of Proposition 4.1 and Proposition 4.3
- The derivation treats the noise contribution as independent of the subject tokens, relation tokens, and next token, then combines noise, subject, and relation terms in the expected-gradient calculation.Source: §B.2 · pp. 24–26
B.3 Proof of Proposition 4.7
- The proof of Proposition 4.7 reduces spectral orthogonalization to normalizing a matrix by ∥u∥.The passage introduces a lemma for applying spectral orthogonalization to the relevant matrix.Source: §B.3 · p. 26
B.4 Proof of Proposition 5.1
- The proof of Proposition 5.1 calculates element-wise gradient signs for different combinations of subject, relation, and answer embeddings.It combines lemmas describing how those signs vary across embedding choices.Source: §B.4 · pp. 27–28
C.1 Proof of Theorem 4.4
- The proof of Theorem 4.4 expresses the derivative of qS using qS and a constant, then separates variables and integrates.The proof treats the relation-side margin equations similarly.Source: §C.1 · p. 29
C.2 Proof of Theorem 4.8
This appendix proves learning-time results for Spectral GF and Sign GF and supplies supporting embedding and dynamics lemmas. The proof combines bounds and examines how embedding choices enter the Sign GF analysis.
C.2 Proof of Theorem 4.8
- Lemma C.2 characterizes Spectral GF learning times for subject and relation components to reach 0.5 accuracy, and gives corresponding manifold parameters.The proof combines matching upper and lower bounds for Tδ and discusses when its two approximations apply.Source: §C.2 · pp. 31, 34, 35
C.3 Proof of Learning Times for Sign GF
- Theorem C.3 treats a general four-component ODE and states learning-time bounds for its subject and relation components.The supplied theorem and conclusion passages do not include the bound expressions.Source: §C.3 · p. 36
D.1 Proof of Lemma B.1
- The proof of Lemma B.1 constructs a singular value decomposition using subject and relation embedding matrices made orthogonal by Q_S and Q_R.The proof notes that the original V_S and V_R matrices do not have the same properties as the orthonormal embedding basis.Source: §D.1 · pp. 37–38
D.2 Proof of Lemma B.2 and Lemma B.3
- Lemma B.2's proof evaluates the element-wise sign for four combinations of subject and answer embeddings: one-hot or Hadamard.The four cases are enumerated by which embedding type is used for each.Source: §D.2 · pp. 38–40
D.3 Proof of Lemma C.1
- The proof of Lemma C.1 distinguishes S=R from S>R when analyzing the relation-component dynamics.For S>R, it considers the relation-component dynamics from initialization.Source: §D.3 · p. 40
D.4 Proof of Lemma C.2
This appendix derives the relation-component learning time and manifold parameters, alongside auxiliary lemmas about orthogonal matrices and embedding properties.
D.4 Proof of Lemma C.2
- The proof of Lemma C.2 uses the initialization assumptions on α0 and β0 and derives the relation-dependent learning time and manifold parameters.The supplied passages do not include the resulting expressions.Source: §D.4 · p. 42
E Auxiliary Lemmas
- The auxiliary lemmas state properties of the subject- and relation-dependent embeddings, including an orthogonality relation between them.Lemma E.2 lists properties of the embeddings defined earlier; its supplied proof invokes their orthonormality.Source: E Auxiliary Lemmas · pp. 42–43