Muon learns balanced solutions in matrix factorization without slow saddle-to-saddle dynamics Mark Rhee∗ UC Berkeley
Jamie Simon Imbue, UC Berkeley
Dhruva Karkada UC Berkeley
arXiv:2606.30509v1 [cs.LG] 29 Jun 2026
Abstract Matrix factorization (i.e., problems of the form minP ,Q ∥M ∗ − P Q⊤ ∥2F ) is a minimal learning problem that exhibits both nonlinear parameter dynamics and representation learning. In this setting, we study how parameter trajectories under the Muon optimizer differ from those of gradient descent. We identify three main dynamical differences: 1) Muon avoids the slow saddleto-saddle dynamics from small initialization. Muon instead learns all the top modes of M ∗ at the same rate, with the smaller modes converging first. 2) Muon remains stable even when the learning rate exceeds the critical threshold set by the local loss sharpness. This frees the learning rate from the condition number of the problem, enabling rapid convergence via exponential learning rate annealing. 3) Once the are aligned with each other and the target, Muon √ weights p flow conserves the matrix quantity P ⊤ P − Q⊤ Q, while gradient flow is known to conserve the matrix P ⊤ P − Q⊤ Q. Despite having distinct conserved quantities, both optimizers find the so-called balanced solution from vanishing initialization. When training from small random initialization, the weights spontaneously align early in training. We derive the alignment rates in simple settings and show that they predict the empirical alignment rates in general. Finally, we exploit structural properties of Muon to construct a learning rate schedule that achieves near-perfect alignment in only two optimization steps.
1
Introduction
Muon is a first-order optimizer which has recently found success in accelerating training in large models. Its implementation details vary; in this paper, we study the momentum-free update rule \ ∆W = −η · ∇ W L,
(Muon)
b denotes the where W is a weight matrix, η is the learning rate, ∇W L is the gradient, and G b = U V ⊤. orthogonalization of G: if its singular value decomposition is G = USV ⊤ , then G Therefore, all Muon updates are (semi-)orthogonal matrices scaled by η. In contrast, the regular gradient descent update, ∆W = −η · ∇W L, directly inherits the spectrum of the gradient. In this paper, we study the consequences of this difference for matrix factorization problems, L(P , Q) =
1 2 M ∗ − P Q⊤ F , 2
where M ∗ is a target matrix and P and Q are trainable weights. ∗ Correspondence to [email protected]. Code is publicly available at https://github.com/dkarkada/muon-mfac.
This work was completed during a research fellowship program at Imbue.
(1)
singular values of P
Muon
singular values of P
Gradient Descent
optimization time t
optimization time t
Multiplicative dynamics → slow saddle- Additive dynamics → uniform growth, to-saddle learning. Each of the weight’s no intermediate saddles. After alignment, singular values grows at a rate proportional each singular value increases by η per step, to its current magnitude. Singular directions so all singular directions are learned at the slowly escape the origin one at a time. Larger same rate with no intermediate saddles. Thus, modes converge before smaller ones. smaller modes converge before larger ones.
initial loss
Final loss
102
10 8
critical LR
Final loss
108
10−4
10−10
10−2
10−1
learning rate
10 −4
10 −10
100
initial loss
10 2
10 −2
10 −1
learning rate
10 0
Maximum stable learning rate. Optimiza- Always stable. The update magnitude is tion diverges if the learning rate exceeds the fixed and cannot diverge. With large learning critical threshold 2/λmax (∇2 L). The rate of rates, optimization proceeds quickly to the convergence is thus limited by the condition solution manifold and oscillates there until number of the problem. the learning rate is annealed. 2.0
2.0
1.5
1.5
¢ 2 conserved
1.0
¢ 1 conserved
1.0
0.5
0.5 0.0
¢1 ¢2
® ¢(t); ¢(t ¤ ) F
® ¢(t); ¢(t ¤ ) F
¢1 ¢2
t¤
0.0
optimization time t
Flow lines are hyperbolic. The quantity ∆2 := P ⊤ P −Q⊤ Q is conserved under gradient flow. Trajectories in parameter space have hyperbolic geometry.
t¤
optimization time t
Flow lines are affine. Afterp alignment, the √ quantity ∆1 := P ⊤ P − Q⊤ Q is constant under Muon flow; the flow equations are Hamiltonian with H(P , Q) = Tr(∆1 ). Optimization trajectories are parallel rays.
Figure 1: Muon vs. gradient descent in matrix factorization. See Appendix A for experimental details.
2
Matrix factorization problems are interesting because the model P Q⊤ exhibits nonlinear parameter dynamics, learns nontrivial representations of the data (e.g., in word embedding algorithms or PCA), and has large-width inductive biases that illuminate those of neural networks. In addition, matrix factorization finds use in applications such as recommender systems and LoRA fine-tuning (Srebro 2004; Hu et al. 2022). We show that, starting from small initial weights, the Muon learning dynamics unfold as follows: (1) At early time, the right singular vectors of P and Q align with each other, and the singular vectors of the model P Q⊤ align with those of the target. We derive the alignment rates in simple settings in Section 4 and find a close match with general-case empirics. (2) As alignment completes, the singular values of each weight matrix grow uniformly. Each singular value of the model saturates when it reaches its target, oscillating ad infinitum within a neighborhood of the target. Full convergence can only be achieved by annealing the learning rate; geometric annealing suffices. We derive these facts in Section 3. We emphasize that the post-alignment behaviors in particular are qualitatively different from gradient descent, which exhibits stepwise learning from small initialization and whose convergence is sensitive to the learning rate (see Figure 1). In Section 5, we design a learning rate schedule that achieves alignment in only two optimization steps, in the case where the model has enough parameters to completely learn the target. The schedule consists of three stages: a small initial step that aligns P with Q, followed by a large step that simultaneously aligns the model with the target and drives it to the solution manifold, and finally an annealing stage during which the model converges to a solution.
1.1
Setup and notation
We use capital boldface to denote matrices and lowercase boldface for vectors. Parenthesized subscripts denote tensor elements (e.g., A(ij) is a scalar). Parenthesized superscripts denote optimization iterates (e.g., P (t) is P after t steps); we leave off the superscript when it is unimportant or clear from context. For a matrix A, ∥A∥F denotes the Frobenius norm, equal to the square root of the sum of squared singular values (equivalently, of squared entries); ∥A∥∗ denotes the nuclear norm, the sum of the singular values. 2
The learning problem is L(P , Q) = 12 M ∗ − P Q⊤ . M ∗ ∈ Rn×n is a fixed square target matrix.1 P ∈ Rn×d and Q ∈ Rn×d are both trainable weights whose discrete dynamics are driven by the Muon iteration rule with learning rate η: d [ ∆P = −η · ∇ P L = −η · RQ
(2)
\ ⊤ [ ∆Q = −η · ∇ Q L = −η · R P ,
(3)
where we define the (time-dependent) residual as R := M ∗ − P Q⊤ . We use the wide hat to denote 1We focus on square M ∗ for simplicity; our analysis extends straightforwardly to non-square targets.
3
the matrix sign function, i.e., polar orthogonalization:2 b := msign(G) = U V ⊤ G
where G = U SV ⊤ . (0)
(4) 2
α The initial conditions are sampled from the isotropic Gaussian: P(ij) ∼ N (0, max(n,d) ) and likewise for Q. The initialization variance depends on whether the model is over-parameterized (d ≥ n) or under-parameterized (d < n), because we want the relative size of the update to scale only with the initialization scale α and the learning rate η, and not with the model dimensions:
∥∆P ∥2 ∥∆Q∥2 η 2 min(n, d) η2 ≈ ∼ = (α2 / max(n, d)) · nd α2 ∥P (0) ∥2 ∥Q(0) ∥2
(5)
using standard concentration arguments.3 This convention is consistent with prior literature using Muon at scale (Liu et al. 2025); however, this argument privileges the Frobenius norm, and is therefore inconsistent with µP , where the native norm is the spectral norm (Qiu et al. 2026; Yang et al. 2023). (We are not constrained by µP scaling since we are not taking the d/n → ∞ limit.) To understand the dynamics of singular vector alignment in this flow-like regime, we define the following SVDs: ⊤ M ∗ = U ∗S∗V ∗⊤ P Q⊤ = Umod Smod Vmod (6) and
P = UP SP VP⊤
Q = UQ SQ VQ⊤ .
(7)
Intuitively, the alignment between the weights is quantified by the overlap4 between VP and VQ , and the model-target alignment is quantified by the overlap between Umod and U ∗ , and between Vmod and V ∗ . If these singular vectors are fully aligned, and if SP ≈ SQ (which is the case when the initialization is small, see Section 3), then the optimization trajectory will converge near the set of balanced solutions {P , Q : P Q⊤ = topd (M ∗ ), P ⊤ P = Q⊤ Q}. Without loss of generality, we will always assume that M ∗ is diagonal with nonnegative entries. We do this by simply rotating the whole problem into the SVD basis of M ∗ . Concretely, define the rotated weight matrices P̃ := U ∗ ⊤ P and Q̃ := V ∗ ⊤ Q. Then, the loss becomes L̃ = ∥S ∗ − P̃ Q̃⊤ ∥2F and the update rule becomes d̃ [ ∆P̃ = −η · ∇ P̃ L̃ = η · RQ̃ \ [ ⊤ ∆Q̃ = −η · ∇ Q̃ L̃ = η · R̃ P̃ , where R̃ = S ∗ − P̃ Q̃⊤ is the rotated residual. Since the equations are the same in the new variables, and since the initialization distributions of P̃ and Q̃ are the same as those of P and Q, we can simply drop the tildes to work directly in the diagonal basis of M ∗ . 2 If G is a vector, then matrix sign reduces to unit-normalization, justifying the hat notation. We do not use hats to denote statistical estimates (since our problem, as framed, does not involve statistical estimation). 3We will later find that ∥P (t) ∥ grows linearly with t, so one can warm up the learning rate linearly and still ensure ∥∆P ∥/∥P ∥ ≪ 1 throughout the growth phase, i.e., that the trajectory is well-described by a continuous flow. 4 See Section 4 for precise definitions of alignment and overlap.
4
2
Related Work
Muon in matrix factorization. Closest to our work, Ma et al. (2026) and Kang et al. (2026) analyze Muon in the matrix factorization setting. Both identify a uniform growth phase in which every singular value increases at the same rate, giving linear convergence whose iteration complexity is independent of the target’s condition number. Both works rigorously prove convergence rates using techniques that avoid directly characterizing the alignment phase. We instead pursue a mechanistic understanding of alignment and growth via scaling arguments. Our analysis, while less rigorous, is intuitive and correctly predicts the empirical alignment dynamics in general settings. Learning dynamics of gradient descent in matrix factorization. Matrix factorization is a standard problem setting to study deep learning phenomena. From small initialization, gradient descent on a matrix factorization objective has a pronounced low-rank bias, greedily recovering the target one effective rank at a time via sigmoidal parameter trajectories (Saxe et al. 2014; Gunasekar et al. 2017; Arora et al. 2019; Gidel et al. 2019; Z. Li et al. 2021). The same objective is also a minimal model of representation learning: Karkada et al. (2025) show that a matrix factorization objective correctly predicts the learning dynamics of the word2vec word embedding algorithm. A recurring theme is that training from small initialization separates into an early alignment phase and a later growth phase. Atanasov et al. (2022) term this silent alignment: the singular vectors of the model align with those of the target while the loss is still essentially flat, on a timescale governed by the ratio between the target scale and the initialization scale, after which the singular values grow. Asymptotically, gradient descent provably aligns the factors both with each other and with the target (Ji et al. 2019; Radhakrishnan et al. 2020). The growth phase itself proceeds as saddle-to-saddle dynamics: trajectories linger near a sequence of low-rank saddles and switch on one mode at a time, producing the characteristic staircase loss curve (Jacot et al. 2021). These dynamics are organized by a conservation law, namely that the difference of squared factor norms is invariant under gradient flow, so that small initialization keeps the two factors balanced throughout training (S. S. Du et al. 2018; Marcotte et al. 2023). Our analysis revisits each of these three phenomena—alignment, growth, and conservation—under Muon. When Muon outperforms gradient descent. Muon is an optimizer that flattens the spectrum of matrix-valued updates (Jordan et al. 2024). It has driven strong empirical results (Liu et al. 2025). Bernstein et al. (2024) cast this update as steepest descent under the spectral norm. A consistent picture across works that compare Muon with gradient descent is that Muon helps precisely when the problem carries heterogeneous spectral structure—a low-rank, ill-conditioned, or imbalanced component that GD learns slowly. This mechanism has been made precise for class-imbalanced classification, where spectral descent and Muon learn all principal components at equal rates and thereby recover the tail classes that GD and Adam learn slowly (Vasudeva et al. 2025; Wang et al. 2025; Fan et al. 2025; B. Li et al. 2026), and for deep networks more broadly, where low-stable-rank activations induce an ill-conditioned Euclidean landscape to which spectral updates are naturally adapted (Davis et al. 2025; Braun et al. 2026). Ill-conditioned matrix factorization is the same phenomenon in a pure optimization setting, with the spread of target singular values playing the role of the imbalance ratio.
5
3
Learning dynamics after weight alignment
We will first assume that the weights are already fully aligned, i.e., that the model is aligned with the target (P Q⊤ commutes with M ∗ ) and the weights are aligned with each other (P ⊤ P commutes with Q⊤ Q). Then the residual R = M ∗ − P Q⊤ is diagonal in the target’s SVD basis, and the learning problem decouples into n independent scalar problems (Ma et al. 2026): 1 ∥R∥2 2 n 1X ∗ (m − pµ qµ )2 = 2 µ=1 µ
(8)
L=
(9)
where m∗µ , pµ , qµ are the µ-th largest singular values of M ∗ , P , Q respectively. Without loss of generality, we can directly optimize pµ and qµ , due to the model’s internal rotation symmetry (i.e., neither the task nor the learning rule prefers a particular basis in the model’s hidden latent space). Each of these scalar problems minimizes the loss L(p, q) = 12 (m∗ − pq)2 . The Muon update is ∆p = η · rq b = η · sign(r) ∆q = η · rp b = η · sign(r)
(10) (11)
where r := m∗ − pq is the residual, and we assumed nonzero singular values, i.e., p, q > 0. Already, we see some striking differences from gradient descent (see Figure 2 for a visualization): • The magnitude of the update is fixed. This implies that each singular value flows at the same uniform rate, with no dependence on its current value or the target value. This sharply contrasts gradient descent, where singular values grow exponentially, separating from each other in the small initialization regime, leading to stepwise learning. • For fixed η, the iterates can neither diverge nor converge. The “effective sharpness” is zero everywhere except on the measure-zero set where the gradient is ill-defined, so all learning rates are stable. On the other hand, the trajectory does not slow down near the minimizer, causing the iterates to oscillate across the solution manifold due to the changing sign(r). • When training from small initialization, r > 0 before reaching the solution manifold. In the flow limit, ṗ = q̇ = 1, and p − q is a conserved quantity. Another way to understand the last observation is to import the following fact from classical mechanics: a flow (ṗ, q̇) is Hamiltonian if it is generated by some scalar H(p, q) via q̇ = ∂H/∂p and ṗ = −∂H/∂q. The r > 0 Muon flow is generated by the Hamiltonian H(p, q) = p−q. Hamiltonian flows have the property that H is an invariant of the flow, and is therefore a conserved quantity. The orbits of this Hamiltonian are parallel in parameter space; this contrasts gradient flow, whose flow lines are hyperbolas (since gradient flow conserves the quantity p2 − q 2 , through a different, non-Hamiltonian mechanism). p √ It follows that the matrix generalization of this conserved quantity is ∆1 := P ⊤ P − Q⊤ Q. Likewise, the flow-generating Hamiltonian is H(P , Q) = Tr ∆1 = ∥P ∥∗ − ∥Q∥∗ , where ∥ · ∥∗ is the nuclear norm. If the weights are aligned and ∆1 ≈ 0, then Muon will converge to a near-balanced solution. 6
Gradient flow
1.5
1.5
1.0
1.0
mini
mize rs
0.5
mini
mize rs
0.5
0.0
0.0 p
1
2
0 "Effective" loss surface along q = p
0 Loss surface along q = p
Muon flow
2.0
q
q
2.0
init
converge
0
p
1
2
p
1
2
init
oscillate
0
p
1
2
Figure 2: In GD, flow lines are hyperbolic with slow saddle points; in Muon, they are parallel and uniform. (Top.) We show optimization flow trajectories in parameter space for scalar factorization, L(p, q) = 12 (1 − pq)2 , comparing gradient descent (left) and Muon (right). Arrow lengths indicate update magnitude, and the solution manifold pq = 1 is shown in blue. (Bottom.) We show a slice of the (effective) loss landscape along the line p = q. The “effective loss” Leff is defined such that the GD trajectory on Leff gives the same trajectory as Muon on L. In GD, small initialization is cursed by slow convergence due to the saddle point at the origin. Muon escapes this curse since its flow rate is uniform arbitrarily close to the origin. However, unlike GD, Muon with constant learning rate oscillates about the solution manifold without converging, so it is necessary to anneal the learning rate.
We note that the logical relation between the conservation law and alignment is different from GD. In GD, the conserved quantity holds at all times and is a consequence of the learning rule and symmetry of the problem, and can thus be used to argue that the weights must quickly align. In Muon, however, we assume alignment as a precondition and use it to derive a conservation law; the conservation law does not hold before alignment is achieved. Clearly, the dynamics of alignment are important to understand; we study them in detail in Section 4.
3.1
Oscillations near the solution manifold
With constant learning rate, the Muon iterates oscillate about the solution manifold without converging. To converge, then, it is necessary to anneal the learning rate. In particular, geometrically annealing the learning rate by a factor of 1/2 at each step gives linear convergence. We give a direct proof here in the case of aligned weights near the solution manifold, and refer to Ma et al. (2026) for a more general analysis.
7
The conservation law ensures that the parameters stay balanced. Assume as the inductive √ √ hypothesis that at time t, |p(t) − m∗ | < η (t−1) . Without loss of generality, assume p(t) − m∗ < 0 (the √ (t−1) p(t) − m∗ > 0 case holds by a symmetric argument). Then p(t+1) = p(t) +η (t) where η (t) = η 2 . The inductive step maintains the inductive hypothesis as follows: p(t+1) −
√
m∗ = p(t) −
√
m∗ +
< −η (t−1) + = The convergence rate of p(t) →
3.2
√
η (t−1) 2
η (t−1) 2
η (t−1) = η (t) . 2
m∗ is thus upper bounded by the geometric decay of η (t) .
Oscillation traps near the origin
The solution manifold is not the only place where the iterates can oscillate. If the initial weights are very misaligned and the learning rate is too large, the iterates can oscillate across the origin for a long time. One can understand this from the flow map in Figure 2 by considering the initialization (p0 , q0 ) = (−α, α) with η = α; the iterates oscillate indefinitely between the second and fourth quadrants. In both the scalar and matrix cases, these oscillation traps can be avoided by going deeper into the flow regime (i.e., decreasing η/α). In the matrix case, this trap can also be avoided by heavy overparameterization, which makes catastrophic misalignment much less likely.
4
Why we can expect model subspaces to align quickly
In matrix factorization problems, the most important aspects of learning are contained in the alignment dynamics. It is during alignment that the model learns to directly encode the most important target directions within the singular directions of the weights themselves. This nontrivial behavior is key to learning task-relevant features; over-parameterized models in the neural tangent kernel regime do not align, despite ultimately solving the matrix factorization problem with zero loss. Therefore, it is important to understand the mechanism of rapid spontaneous weight alignment when training from small random initialization. We begin by formally defining alignment, a scalar quantity associated to a pair of positive semieVe ⊤ ) that quantifies the extent to definite matrices (with eigendecompositions V SV ⊤ and Ve S which they share a jointly-ordered eigenbasis. For example, when both spectra are non-degenerate, the two matrices are fully aligned iff V ⊤ Ve is diagonal with entries ±1. In this case, we can define the scalar alignment metric anondegen ∈ [0, 1]: anondegen :=
n n 1X 1X eµ⊤ ) = eµ ). Tr (vµ vµ⊤ )(e vµ v cos2 (vµ , v n µ=1 n µ=1
(12)
Thus, the trace of the product of projectors is simply the squared cosine similarity, a natural and intuitive measure of overlap. 8
1 ¡ ain
10 −1
10 −2
® = 10 ¡4
r ¡4 (theo
® = 10 ¡5
10 −2
10 −1
10 0
Over-parameterized (d=n = 4:0) (1 ¡ aleft=right ¼ 0 throughout)
t
1 ¡ aleft=right
:2 ¡1
t
left right
10 −1
10 0
Exactly-parameterized (d=n = 1:0) (1 ¡ aleft=right ¼ 0 throughout)
1 ¡ ain
10 0
Under-parameterized (d=n = 0:25) (1 ¡ ain ¼ 0 throughout)
10 −2
y)
® = 10 ¡3 10
−3
10 −6
10 −4 10 −5 t=tconverge
10 −3
10 −3
10 −6
10 −4 10 −5 t=tconverge
10 −3
10 −3
10 −6
10 −4 10 −5 t=tconverge
10 −3
Figure 3: Alignment rates vary with aspect ratio d/n and occurs earlier at smaller initialization scale. We show the internal alignment metric ain and the left and right target alignment metrics aleft and aright for matrix factorization problem (n = 128, d ∈ [32, 128, 512]). The target has planted low-rank-signal-in-noise structure: M ∗ = G + λU V ⊤ , where G is random Gaussian, λ is the planted signal strength, and U , V ∈ Rn×32 are semi-orthogonal. See Appendix A for details.
However, we must generalize this notion for matrices with (nearly) degenerate subspaces. This is because the singular vectors spanning such subspaces are unstable under perturbation. To ensure that our alignment metric is not sensitive to such instabilities, we define n′ h i 1 1 X e e a(V , V , S, S) := ′ Tr (Vµ Vµ⊤ )(Veµ Veµ⊤ ) n µ=1 dim(µ)
(13)
n′
2 1 X 1 = ′ Vµ⊤ Veµ , n µ=1 dim(µ) F
(14)
where µ enumerates the n′ nearly-degenerate subspaces of the PSD matrices, and Vµ ∈ Rn×dim(µ) is the submatrix of V containing the µth nearly-degenerate subspace. It remains to choose how to demarcate these subspaces jointly between the two matrices, especially when their spectra differ. In our experiments, we use a convention in which indices i and j are grouped in a degenerate subspace if either si /sj ≈ 1 or s̃i /s̃j ≈ 1 (see Appendix A). With a quantitative alignment metric in hand, we may study the following alignment quantities: Internal alignment.
:= a(VP , VQ , SP , SQ )
(15)
aleft := a(Umod , U ∗ , Smod , S ∗ ) aright := a(Vmod , V ∗ , Smod , S ∗ ).
(16)
ain
Left target alignment. Right target alignment.
(17)
In both our theory and experiments, we will be interested in the misalignment quantity 1 − a (which, in some sense, represents the average sin2 (θ) between subspaces). Misalignment typically starts at some O(1) quantity and rapidly decays to zero over the course of optimization. Our goal is to characterize this decay rate. 9
This analysis requires working in cases; the dynamics of internal alignment are qualitatively distinct from those of target alignment. In particular: • In the under-parameterized regime (d < n), internal alignment comes essentially for free, while target alignment decays via power iteration like −1
1 − aleft ≈ 1 − aright ∼ (t/αeff )−2(1−s
)
,
(18)
where t is training time, αeff := max(η, α) is the effective5 initialization scale, and s ≥ 1 is the relative spectral gap in the target. We derive this rate theoretically in the n = 2, d = 1 problem L(p) = 12 ∥diag([s, 1]) − pp⊤ ∥2 . • In the over-parameterized regime (d > n), target alignment comes essentially for free, while internal misalignment decays like 1 − ain ∼ (t/αeff )−4 .
(19)
We derive this rate theoretically in the problem L(p, q) = 12 (m∗ − p⊤ q)2 . • When d = n, we empirically find that both types of alignment happen quickly. In sum, target misalignment decays polynomially with exponent at most −2, with rates slowing as the spectral gap in the target gets smaller. Internal misalignment decays even more quickly: polynomially with exponent −4. In all cases, alignment occurs earlier with smaller effective initialization scale. These results are summarized in Figure 3.
4.1
A simple model of under-parameterized alignment
To study model-target alignment, we analyze the “outer product factorization” problem: minimize L(p) = 12 ∥M ∗ − pp⊤ ∥2 , where M ∗ = diag(s, 1) with s ≥ 1. Since both factors share the same parameters (tied vector weights), internal alignment is trivially satisfied; the target misalignment is ε := 1 − aright = 1 − (p̂⊤ ê)2 , where ê := (1, 0)⊤ is the dominant eigenvector of M ∗ . From small initialization, ∥p∥2 ≪ 1 for much of optimization, and the residual is dominated by the target. Then the Muon flow is approximately ṗ ≈
M ∗p . ∥M ∗ p∥
(20)
Absorbing the normalization into a rescaled clock via dt = ∥M ∗ p∥ dτ gives dp/dτ = M ∗ p. With symmetric initialization p(0) = [α, α], this integrates to sτ p e p(τ ) = α , ∥M ∗ p(τ )∥ = α s2 e2sτ + e2τ . (21) τ e Rτ √ Returning to physical time via t(τ ) = α 0 s2 e2sτ ′ + e2τ ′ dτ ′ , for s > 1 the first term dominates: t(τ ) ∼ α esτ
=⇒
τ (t) ∼ 1s ln(t/α).
(22)
5 If the step size η is larger than the true initialization scale α, then the effective initialization scale is η, since the first step will take the model to that scale.
10
Substituting back, p(0) (t) ∼ t,
p(1) (t) ∼ α(t/α)1/s .
(23)
1
(24)
The target misalignment is then ε(t) =
p2(1) p2(0) + p2(1)
=
1 + p(0) /p(1)
2 ∼
1 . 1 + (t/α)2(1−s−1 )
Discretizing the flow with step size η, we get two qualitative regimes. When η/α ≪ 1 we are near the flow regime, and the above analysis holds. However, when η/α ≫ 1, the first step increases the parameter norms from α to ≈ η while barely changing the alignment. Therefore, the effective initialization scale is αeff := max(η, α), and the misalignment decays asymptotically like 1 − aright ∼
t max(η, α)
−2(1−s−1 ) .
(25)
At worst, as the relative spectral gap closes (s → 1+ ), the exponent vanishes and alignment stalls. At best, ε ∼ (t/αeff )−2 when s → ∞.
4.2
A simple model of over-parameterized alignment
To study internal alignment, we analyze the “inner product factorization” problem: minimize L(p, q) = 12 (m∗ − p⊤ q)2 . Since the target m∗ is scalar, the model p⊤ q is always aligned to it; the misalignment between the vectors p and q is ε := 1 − ain = 1 − (p̂⊤ q̂)2 = sin2 (p, q). With small initialization, the early-time residual r := m∗ − p⊤ q is positive. The Muon flow equations are b b ṗ = q, q̇ = p. (26) These are Hamilton’s equations for H := ∥p∥ − ∥q∥. Assume balanced initialization ∥p(0)∥ = ∥q(0)∥ = α; then √ H ≡ 0. We may define the radial variable x := ∥p∥ = ∥q∥, whose evolution is ẋ = pb⊤ qb = 1 − ε. The Hamiltonian is invariant under (p, q) 7→ (Rp, Rq), so Noether’s theorem gives the conserved angular momentum L2 := ∥p∥2 ∥q∥2 − (p⊤ q)2 = x4 ε = const. (27) √ √ Differentiating x4 ε = L2 and using ẋ = 1 − ε gives ε̇ = −4ε 1 − ε/x. Eliminating x using the angular momentum conservation law gives √ 5/4 4 1 − ε √ ε̇ = −ε . (28) L Separating variables yields an elliptic integral, but we can extract the asymptotic alignment behavior by taking ε → 0. This gives ε(t) = 1 − ain (t) ∼ 11
t √ L
−4 .
(29)
2 4 4 At O(1) for rotation-invariant init in d ≥ 2), so √ initialization L = α ε0 ≈ α (since ε0 = L ∼ α and the flow predicts 1 − ain ∼ (t/α)−4 . Discretizing the flow with step size η as in the under-parameterized setting, we get
1 − ain ∼
t max(η, α)
−4 .
(30)
Though this rate always holds at large t, the absolute misalignment can drop dramatically (almost to zero) in the first two steps, if η ≫ α. We discuss this trick in the next section.
5
The spiked learning rate schedule
Using the theoretical insight we gained from Section 3 and Section 4, we construct a learning rate schedule that achieves near-perfect alignment in two steps, in the exactly- and over-parameterized regimes. We do not claim that such a schedule is advisable in general settings; instead, it is a proof-of-concept that theoretical insight about learning dynamics can inform practice. We assume d ≥ n and small (semi-)orthogonal initial weights.6 Then, at initialization, it holds that P P ⊤ = QQ⊤ = α2 In and M ∗ ≫ P Q⊤ . The initial updates will be: 1 ∗ \ [ ∇ P L ≈ −M Q = − Q α 1 ∗ \ [ ∇Q L ≈ −M P = − P . α
(31) (32)
The diagonal M ∗ term in the gradient will simply rescale the already orthogonal rows of P and Q. However, the orthogonalization step of Muon will then immediately rescale the rows again to make them unit length. Thus, at initialization, the updates are approximately: η (0) (0) Q α η (0) (0) Q(1) ≈ Q(0) + P . α
P (1) ≈ P (0) +
(33) (34)
If we set the initial learning rate η (0) equal to α, then P and Q align perfectly in exactly one step: P (1) ≈ P (0) + Q(0) ≈ Q(1) . The second update will approximately be: ∗ P (1) P (2) ≈ P (1) + η (1) M\
(35) (36)
(2)
≈Q
where we used the fact that M ∗ ≫ P (1) (Q(1) )⊤ (as long as the initial learning rate is small) ∗ P (1) and that P (1) ≈ Q(1) . If the learning rate η (1) is sufficiently large, the update matrix M\ 6 In the heavily overparameterized setting, d ≫ n, the following assumptions hold approximately with high probability under random Gaussian initialization, and the rest of the argument carries over as well.
12
Loss
1.00
2.0
0.75
1.5
0.6 ´(t)
0.50
1.0
0.25 0.00
Learning Rate
L(t)
s^i (t)
Singular Values of PQ >
0.2
0.5 0
4
8 iteration
12
0.0
0.4
0
4
8 iteration
12
0.0
0
4
8 iteration
12
Figure 4: The spiked learning rate schedule in the exactly-parameterized setting (n = d = 25) with orthogonal initialization at scale α = 10−4 . (Left.) Singular values of P Q⊤ over 16 iterations. (Middle.) Loss L(P , Q). (Right.) Learning rate schedule. The p first step (η (0) = α = 10−4 ) achieves (1) near-perfect internal alignment. The second step (η = s∗max /2 ≈ 0.707) simultaneously achieves near-perfect external alignment and grows all singular values toward s∗max /2. Geometrically annealing by 1/2 thereafter yields rapid convergence. ∗ P (1) ≈ Q(2) . As long as n ≤ d, this will overwhelm P (1) and Q(1) such that P (2) ≈ η (1) M\ (2) (2) forces the rows of P (which are exactly the rows of Q ) to be mutually orthogonal such that P (2) (Q(2) )⊤ ≈ (η (1) )2 In , thus achieving near-perfect alignment with the target. See Figure 4 for empirical validation of our assumptions and the efficacy of our resulting schedule.
We are free to set the initial learning rate η (0) and the learning rate spike η (1) to any values that satisfy η (0) = α and η (1) ≫ η (0) . After the two steps that achieve alignment, the objective decouples into n independent scalar problems; geometrically annealing the learning rate guarantees linear convergence as shown in Section 3.1. Limitations. Our analysis is restricted to simple matrix factorization problems. We focus on the unmodified spectral gradient descent update rule, i.e., ignoring the effects of momentum, weight decay, and Newton-Schulz iteration. Moving beyond these simplifications is an interesting and important area for future work. Acknowledgements. We thank Joey Turnbull and Jianhao Ma for useful discussions. This work was funded by Imbue under the Feature Lab (FLAB) initiative. Author contributions. JS and DK conceived the idea in the car while crossing the Bay Bridge and oversaw the project’s subsequent development. MR and DK jointly performed the theoretical analysis, coded the numerical experiments, and wrote the manuscript with input and guidance from JS.
13
References Arora, Sanjeev et al. (2019). “Implicit regularization in deep matrix factorization”. In: Advances in Neural Information Processing Systems (NeurIPS). arXiv:1905.13655 (cit. on p. 5). Atanasov, Alexander, Blake Bordelon, and Cengiz Pehlevan (2022). “Neural networks as kernel learners: The silent alignment effect”. In: International Conference on Learning Representations (ICLR). arXiv:2111.00034 (cit. on p. 5). Bernstein, Jeremy and Laker Newhouse (2024). “Old optimizer, new norm: An anthology”. In: arXiv preprint arXiv:2409.20325 (cit. on p. 5). Braun, Guillaume et al. (2026). “Spectral gradient descent mitigates anisotropy-driven misalignment: A case study in phase retrieval”. In: arXiv preprint arXiv:2601.22652 (cit. on p. 5). Davis, Damek and Dmitriy Drusvyatskiy (2025). “When do spectral gradient updates help in deep learning?” In: arXiv preprint arXiv:2512.04299 (cit. on p. 5). Du, Simon S., Wei Hu, and Jason D. Lee (2018). “Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced”. In: Advances in Neural Information Processing Systems (NeurIPS). arXiv:1806.00900 (cit. on p. 5). Fan, Chen, Mark Schmidt, and Christos Thrampoulidis (2025). “Implicit bias of spectral descent and Muon on multiclass separable data”. In: Advances in Neural Information Processing Systems (NeurIPS). arXiv:2502.04664 (cit. on p. 5). Gidel, Gauthier, Francis Bach, and Simon Lacoste-Julien (2019). “Implicit regularization of discrete gradient dynamics in linear neural networks”. In: Advances in Neural Information Processing Systems (NeurIPS). arXiv:1904.13262 (cit. on p. 5). Gunasekar, Suriya et al. (2017). “Implicit regularization in matrix factorization”. In: Advances in Neural Information Processing Systems (NeurIPS). arXiv:1705.09280 (cit. on p. 5). Hu, Edward J. et al. (2022). “LoRA: Low-rank adaptation of large language models”. In: International Conference on Learning Representations (ICLR). arXiv:2106.09685 (cit. on p. 3). Jacot, Arthur et al. (2021). “Saddle-to-saddle dynamics in deep linear networks: Small initialization training, symmetry, and sparsity”. In: arXiv preprint arXiv:2106.15933 (cit. on p. 5). Ji, Ziwei and Matus Telgarsky (2019). “Gradient descent aligns the layers of deep linear networks”. In: International Conference on Learning Representations (ICLR). arXiv:1810.02032 (cit. on p. 5). Jordan, Keller et al. (2024). Muon: An optimizer for hidden layers in neural networks. https : //kellerjordan.github.io/posts/muon/ (cit. on p. 5). Kang, Changmin et al. (2026). “Uniform spectral growth and convergence of Muon in LoRA-style matrix factorization”. In: arXiv preprint arXiv:2602.06385 (cit. on p. 5). Karkada, Dhruva et al. (2025). “Closed-form training dynamics reveal learned features and linear structure in word2vec-like models”. In: Advances in Neural Information Processing Systems (NeurIPS). arXiv:2502.09863 (cit. on p. 5). Li, Binghui et al. (2026). “Muon in associative memory learning: Training dynamics and scaling laws”. In: arXiv preprint arXiv:2602.05725 (cit. on p. 5). Li, Zhiyuan, Yuping Luo, and Kaifeng Lyu (2021). “Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning”. In: International Conference on Learning Representations (ICLR). arXiv:2012.09839 (cit. on p. 5). Liu, Jingyuan et al. (2025). “Muon is scalable for LLM training”. In: arXiv preprint arXiv:2502.16982 (cit. on pp. 4, 5).
14
Ma, Jianhao et al. (2026). “Preconditioning benefits of spectral orthogonalization in Muon”. In: arXiv preprint arXiv:2601.13474 (cit. on pp. 5, 6, 7). Marcotte, Sibylle, Rémi Gribonval, and Gabriel Peyré (2023). “Abide by the law and follow the flow: Conservation laws for gradient flows”. In: Advances in Neural Information Processing Systems (NeurIPS). arXiv:2307.00144 (cit. on p. 5). Qiu, Shikai et al. (2026). “Hyperparameter transfer enables consistent gains of matrix-preconditioned optimizers across scales”. In: Advances in Neural Information Processing Systems 38, pp. 130867– 130911 (cit. on p. 4). Radhakrishnan, Adityanarayanan et al. (2020). “On alignment in deep linear neural networks”. In: arXiv preprint arXiv:2003.06340 (cit. on p. 5). Saxe, Andrew M., James L. McClelland, and Surya Ganguli (2014). “Exact solutions to the nonlinear dynamics of learning in deep linear neural networks”. In: International Conference on Learning Representations (ICLR). arXiv:1312.6120 (cit. on p. 5). Srebro, Nathan (2004). “Learning with matrix factorizations”. PhD thesis. MIT (cit. on p. 3). Vasudeva, Bhavya et al. (2025). “How Muon’s spectral design benefits generalization: A study on imbalanced data”. In: arXiv preprint arXiv:2510.22980 (cit. on p. 5). Wang, Shuche et al. (2025). “Muon outperforms Adam in tail-end associative memory learning”. In: arXiv preprint arXiv:2509.26030 (cit. on p. 5). Yang, Greg, James B Simon, and Jeremy Bernstein (2023). “A spectral condition for feature learning”. In: arXiv preprint arXiv:2310.17813 (cit. on p. 4).
15
A
Experimental details
All experiments solve the matrix factorization problem L(P , Q) = 12 ∥M ∗ − P Q⊤ ∥2 with two trainable factors P ∈ Rn×d and Q ∈ Rn×d , optimized either by plain (momentum-free, no weight decay) gradient descent or by the momentum-free Muon update of Section 2. The following conventions apply: the target M ∗ is taken diagonal (we always work in its SVD basis, without loss of generality), the factors are initialized i.i.d. Gaussian with P(ij) , Q(ij) ∼ N (0, α2 / max(n, d)) at initialization scale α, and η denotes the learning rate. Our Muon implementation computes the b = U V ⊤ from a full SVD G = U SV ⊤ of the gradient, update direction as the exact polar factor G rather than the Newton–Schulz approximation used in practice; this removes orthogonalization error as a confound. Code is publicly available at https://github.com/dkarkada/muon-mfac.
A.1
Summary of dynamical differences (Figure 1)
This figure comprises three independent experiments (one per row, GD in the left column and Muon in the right), all run on the same target. The target is the diagonal n × n matrix with n = 32 and singular values sµ = 10 µ−1 for µ ∈ [1, . . . , n]. The models are all exactly-parameterized, d = n. Singular-value trajectories (panels a, b). We train from initialization scale α = 5 × 10−3 for 1000 steps, with η = 5 × 10−3 for Muon and η = 1.5 × 10−2 for GD. Each panel plots all 32 singular values of the factor P as a function of optimization time, colored from the largest mode to the smallest. The axes are unlabeled because the panels are meant to be read qualitatively. The time axis is scaled linearly and the singular value axis ranges linearly from [0, 3.7]. Final loss versus learning rate (panels c, d). Here we sweep the learning rate to probe stability. Both axes are scaled logarithmically. We use a larger initialization scale α = 0.1 and sweep 200 learning rates spaced logarithmically over [0.05, 20] × ηc with ηc = 0.14, a reference rate marking GD’s critical (divergence) threshold. Each Muon run is 500 steps and each GD run is 2000 steps. The plotted “final loss” is the average of the last two iterates’ losses (to suppress the terminal oscillation of Muon) normalized by the initial loss. The red dashed line marks the GD ηc and the gray dashed line marks the initial loss. Conserved quantities (panels e, f). We train from α = 5 × 10−3 for 1200 steps, with η = 5 × 10−4√for Muonp and η = 2.5 × 10−3 for GD, and track the two candidate conserved matrices ∆1 := P ⊤ P − Q⊤ Q (Muon) and ∆2 := P ⊤ P − Q⊤ Q (GD). To test conservation we fix a reference iterate t∗ = 80 (chosen after alignment, during the growth phase) and plot, for each ∆ ∈ {∆1 , ∆2 }, the overlap ⟨∆(t), ∆(t∗ )⟩F /∥∆(t∗ )∥2F . This overlap equals 1 exactly when ∆(t) = ∆(t∗ ), so a curve pinned at 1 for t > t∗ indicates that quantity is conserved. Under GD (panel e) it is ∆2 that stays fixed; under Muon (panel f) it is ∆1 .
16
A.2
Alignment rates (Figure 3)
This experiment measures how quickly the weights align, as a function of the aspect ratio d/n and the initialization scale α, and compares the measured rates to the power laws derived in Section 4. All runs use Muon (no GD comparison here) in float64 for numerical precision, and each configuration is averaged over 3 random seeds. Target. We fix n = 128 and build a planted-signal-in-noise target: form M = diag(λ 1r ) + √ G/ n with planted rank r = 32, planted signal strength λ = 4, and G i.i.d. standard Gaussian; the target singular values are then 10× the singular values of M , floored at 2 to keep the noise bulk well-conditioned. Equivalently M ∗ = G′ + λ U V ⊤ with U , V ∈ Rn×32 semi-orthogonal, as stated in the caption. As elsewhere we rotate into the target’s SVD basis, so M ∗ is diagonal. The relative spectral gap s between the planted and bulk singular values sets the under-parameterized alignment exponent 2(1 − s−1 ); we compute its effective value by taking the mean relative gap between signal modes and the bulk edge. Sweep. We vary the aspect ratio over d ∈ {32, 128, 512}, i.e. d/n ∈ {0.25, 1, 4} (under-, exactly-, and over-parameterized), and the initialization scale over α ∈ {10−5 , 10−4 , 10−3 }, with learning rate η = 0.2 α held small relative to α. Metrics and axes. We report the misalignments 1 − ain , 1 − aleft , 1 − aright using the degeneracyaware alignment metric a(·) of Section 4. The degenerate-subspace grouping merges adjacent modes i, j unless both their relative singular-value gap exceeds a tolerance (here 20%, with a small absolute floor) in both the model and the reference spectra; this is the convention referenced in Section 4. The horizontal axis is rescaled optimization time t/tconverge , where tconverge is the step at which uniform growth would carry the top mode to its target. Panels. The three columns correspond to the three aspect ratios. In the under-parameterized column (d/n = 0.25) internal alignment is essentially free (1 − ain ≈ 0), so we plot target misalignment 1 − aleft (dotted) and 1 − aright (dashed), overlaid with the theoretical power law −1 t−2(1−s ) . In the exactly- and over-parameterized columns (d/n = 1, 4) target alignment is essentially free, so we plot internal misalignment 1 − ain ; the over-parameterized column is overlaid with the predicted t−4 law. Within each panel the three colors correspond to the three initialization scales α, illustrating that smaller α (equivalently smaller αeff = max(η, α)) aligns earlier in rescaled time.
17
A.3
Spiked learning rate schedule (Figure 4)
This experiment demonstrates the spiked learning rate schedule of Section 5 in the exactlyparameterized setting and verifies that alignment and convergence can be achieved in a small number of steps. Setup. We take n = d = 25 and a diagonal target M ∗ ∈ Rn×n with spectrum s∗µ = (4 + µ)−1 for µ ∈ [1, . . . , n], normalized so that s∗1 = 1. Factors P , Q ∈ Rn×d are initialized orthogonally: we draw W1 , W2 ∼ N (0, 1)n×d , compute their QR factorizations Wk = Qk Rk , and set P (0) = α Q1 diag(sign(diag(R1 ))) (and analogously for Q(0) ), where the sign correction makes the initialization unique. The initialization scale is α = 10−4 . Schedule. We run T = 16 steps of the momentum-free Muon update (exact polar factor via full SVD) with the following three-phase learning rate schedule: 1. Alignment step (t = 0): η (0) = α = 10−4 . As argued in Section 5, setting η (0) = α achieves near-perfect internal alignment in a single step. p p 2. Spike (t = 1): η (1) = s∗max /2 = 1/2 ≈ 0.707. This large step simultaneously achieves near-perfect external (target) alignment and uniformly grows the singular values of P Q⊤ toward s∗max /2. 3. Geometric anneal (t ≥ 2): η (t) = η (1) · 2−(t−1) . As shown in Section 3.1, geometric annealing guarantees linear convergence once the iterates are close to the solution manifold. Diagnostics. We track three quantities at each step: (i) the singular values of P (t) (Q(t) )⊤ , (ii) the loss L(P (t) , Q(t) ), and (iii) the learning rate η (t) , plotted against iteration on a linear axis (panels a, b, c respectively).
18