Convergence of Muon with Newton-Schulz

AI summary

Overview

• Research area: Nonconvex stochastic optimization theory for matrix-structured parameters, specifically the Muon optimizer and its Newton–Schulz orthogonalization step.
• Technical level: Advanced. The paper is a theoretical convergence analysis built on Schatten norms (nuclear, spectral, Frobenius), Lipschitz smoothness in operator–nuclear geometry, and polynomial convergence arguments, supported by a numerical study.
• Scope: The paper proves that Muon exactly as proposed and used in practice — momentum orthogonalization approximated by a finite number of Newton–Schulz steps rather than an exact SVD polar factor — converges to a stationary point, and it quantifies how quickly the Newton–Schulz approximation error vanishes.

What This Paper Is About

Muon is an optimizer for matrix-shaped parameters that orthogonalizes the momentum matrix before each update. In practice this orthogonalization is not computed by an exact singular value decomposition (SVD), which is accurate but expensive, but is approximated by a small fixed number of Newton–Schulz steps using only matrix multiplications. Existing convergence analyses typically analyze an idealized variant that replaces the Newton–Schulz step with an exact SVD-based polar factor, leaving open whether the actual SVD-free implementation inherits any nonconvex convergence guarantee. The paper closes this gap by analyzing Muon with a finite number of Newton–Schulz steps, bounding the resulting approximation penalty, and comparing its rank dependence against SGD with momentum under the same stationarity metric.

Key Contributions

1. First convergence result for Muon with Newton–Schulz (Theorem 1). The paper presents what it states is the first nonconvex convergence guarantee for Muon with a finite number of Newton–Schulz steps, as originally proposed and practically used. It notes that even for convex optimization, convergence of Muon with Newton–Schulz had not been shown previously. The key distinction from prior analyses is that Newton–Schulz is not replaced by the exact polar computed via SVD.
2. Analysis of the polar approximation error and wall-clock convergence (Theorems 2 and 4). The polar approximation error due to replacing SVD with Newton–Schulz is proven to decay doubly exponentially with the number of Newton–Schulz steps q and to decay with the polynomial degree κ.
3. Sharper rank dependence in Muon with Newton–Schulz (Theorem 3, Table 1). The paper shows that Muon with Newton–Schulz sharpens the convergence rate by a factor of the square root of the rank of the momentum matrix compared with the vector-based counterpart, SGD with momentum.
4. A convergence bound for SGD with momentum under the same metric (Theorem 3). Stated as new to the authors' knowledge and possibly of independent interest, providing a clean reference point for vector-based updates under nuclear-norm stationarity.

Main Findings

• Muon with Newton–Schulz converges at the SVD idealization's rate up to a constant. Under Lipschitz smoothness (Assumption 1) and bounded variance (Assumption 2), with stepsize η = sqrt((1−β)D/(TL)) and momentum β = 1 − min{sqrt(LDB)/(σ sqrt(rT)), 1} where r = min{m, n}, Theorem 1 gives (1/T) Σ E[||∇f(W_{t−1})||_*] ≤ χ_q · O(sqrt(LD/T) + σr/sqrt(BT) + (rσ²LD/(BT))^{1/4}). The corresponding iteration complexity to reach ε-stationarity is T = O(max{χ_q²LD/ε², χ_q²r²σ²/(Bε²), χ_q⁴rσ²LD/(Bε⁴)}).
• The Newton–Schulz penalty χ_q → 1 doubly exponentially in q. Theorem 2 gives δ_{t,q} ≤ δ_{t,0}^{(κ+1)^q}, hence ε_q ≤ 1 − sqrt(1 − δ_0^{(κ+1)^q}) ≤ δ_0^{(κ+1)^q} and χ_q = 1/(1 − ε_q) ≤ 1/sqrt(1 − δ_0^{(κ+1)^q}), where δ_0 = sup_t δ_{t,0} < 1. The bound improves with larger polynomial degree κ.
• Per-step residual contraction drives the decay. Lemma 3 establishes φ(u) ≤ u^{κ+1} on [0, 1], giving δ_{t,j+1} ≤ δ_{t,j}^{κ+1} and δ_{t,q} ≤ δ_{t,0}^{(κ+1)^q}. Lemma 1 links the orthogonality residual to the polar approximation error: δ_{t,q} = 1 − λ_min^+ and ε_{t,q} = 1 − sqrt(1 − δ_{t,q}).
• Muon removes a square-root-of-rank penalty relative to SGD with momentum. Theorem 3 gives the SGD-with-momentum rate O(sqrt(rLD/T) + (r²σ²LD/(BT))^{1/4}), with iteration complexity O(max{rLD/ε², r²σ²LD/(Bε⁴)}). This carries a sqrt(r) factor in the deterministic term that the Muon rate does not.
• The idealized SVD variant is a special case. Theorem 4 sets ε_q = 0 in Theorem 1, yielding O(sqrt(LD/T) + σr/sqrt(BT) + (rσ²LD/(BT))^{1/4}), with complexity O(max{LD/ε², r²σ²/(Bε²), rσ²LD/(Bε⁴)}). This is described as an "oracle" baseline that clarifies the gap Newton–Schulz must close.
• Iteration count and wall-clock time diverge in favor of Newton–Schulz. Because a small number of Newton–Schulz steps already yields essentially the SVD iteration rate while using only matrix multiplications, the paper argues the practical implementation attains substantially faster wall-clock convergence.
• Numerical ablation on q matches the theory. On CIFAR-10 (50k/10k) with a CifarNet CNN of approximately 2.2M parameters, batch size B = 512, κ = 2, 50 epochs, averaged over five random seeds, Muon with q = 1 already outperforms SGD-M, and q ∈ {2, 3} nearly coincides with SVD-based Muon in both train and test loss. Muon with q = 2 or 3 reaches a given test loss substantially faster in wall-clock time than the SVD variant.
• Ablation on polynomial degree κ shows an accuracy–time trade-off. With q = 3 fixed and κ ∈ {1, …, 5}, increasing κ improves optimization (faster loss drop at a fixed epoch) but lengthens each step, mirroring the residual contraction δ_{j+1} ≤ δ_j^{κ+1} while computation scales with polynomial evaluations.
• Rank-dependence experiment. Varying the monitored layer's effective rank r ∈ {16, 32, 64, 128, 216} and plotting epoch-averaged ||∇f(W)||_* on a log–log scale, SGD-M shows a positive slope of approximately 0.3, whereas Muon and its variants are nearly flat, in line with Theorems 3 and 4.
• Additional experiments are reported at other scales. The paper reports appendix experiments with an MLP of 0.5M parameters on MNIST; ResNet-18 (11.2M) on CIFAR-100; WideResNet-28-10 (36.6M) on Tiny-ImageNet; NanoGPT (124M, Transformer) on FineWeb; and a GPT-2-based model (1.3B, Transformer) on FineWeb.

Methodology in Plain English

The authors keep the original Muon algorithm intact rather than substituting a convenient idealization. They define a scaled momentum N_t = (1 − β)M_t so the exponential moving average follows the original update rule, then apply a descent lemma to the parameter update W_t ← W_{t−1} − ηO_t. This produces an inner product between the gradient and the Newton–Schulz output, which they split twice: first separating the momentum mismatch term, and then separating the exact polar factor P_t = Polar(N_t) = Polar(M_t) from the Newton–Schulz output O_t. The second split is what lets them isolate the Newton–Schulz error as the term ⟨N_t, O_t − P_t⟩_F and control it through the polar approximation error ε_q. Averaging the resulting one-step descent inequality over T iterations and choosing η and β as specified yields the convergence rate.

To bound ε_q, they connect it to an orthogonality residual δ that measures how far the Newton–Schulz iterate is from having orthonormal columns. They then show that one Newton–Schulz step maps that residual through the function φ(u) = 1 − (1−u)[p_κ(1−u)]², and prove φ(u) ≤ u^{κ+1}. Applying this q times gives the doubly exponential bound. For comparison, they derive the same style of guarantee for SGD with momentum and for an idealized Muon that uses an exact SVD polar step, all under identical smoothness and bounded-variance assumptions and all measured by the nuclear norm of the gradient so the rates are directly comparable.

Why This Matters

• Impact on research: The work narrows the practice–theory gap for Muon by analyzing the actual SVD-free algorithm instead of an exact-polar idealization, and shows that an SVD-free method can be analyzed rigorously in nonconvex settings. It also supplies a convergence guarantee for SGD with momentum under nuclear-norm stationarity, and the authors state their analysis can be adapted to other variants of Muon. They explicitly position Muon as not a second-order method — it neither estimates nor inverts curvature — and note that it can be complementary to curvature-preconditioning methods such as Shampoo and SOAP rather than directly comparable.
• Real-world applications (as scoped by the paper's own experiments and framing):
  • Large language model pretraining, where the paper reports experiments with NanoGPT (124M) and a GPT-2-based model (1.3B) on FineWeb, and cites large-scale training reports for Muon.
  • Training convolutional image models such as ResNet-18 (11.2M) on CIFAR-100, WideResNet-28-10 (36.6M) on Tiny-ImageNet, and CifarNet (approximately 2.2M) on CIFAR-10.
  • Settings where SVD is a per-iteration bottleneck and matrix multiplications are preferable, including memory-aware and communication-aware distributed training, which the related work section cites as motivation for SVD-free, GPU-aligned theory.
  • Practical selection of the Newton–Schulz step count and polynomial degree, since the theory quantifies the accuracy–compute trade-off between q, κ, and convergence.
• Industry relevance: The results provide a theoretical justification for using only a few low-degree Newton–Schulz steps, which is exactly what practitioners do, and explain why this matches exact-polar behavior at much faster wall-clock time. Because Newton–Schulz uses only matrix multiplications, the guidance is aligned with hardware-efficient large-scale training.

Future Directions

• Extending the analysis to Muon variants. The paper states its Newton–Schulz analysis can be adapted to other variants, including communication- and memory-aware versions cited in related work, but does not carry those extensions out here.
• A fuller theory of the degree–accuracy trade-off. The κ effect enters the bound as δ^{κ+1} per step, but the paper only studies the κ ∈ {1, …, 5} trade-off empirically in an appendix ablation with q = 3 fixed.
• Convex-case guarantees. The paper notes that even for convex optimization, convergence of Muon with Newton–Schulz had not been shown previously; establishing that case remains open.
• Composition with second-order preconditioners. The paper observes that Muon and curvature-based methods such as Shampoo and SOAP target different mechanisms (projection/normalization versus curvature preconditioning) and may be complementary, leaving the combined or comparative theory unexplored.
• Beyond the stated assumptions. The guarantees rest on operator–nuclear Lipschitz smoothness (Assumption 1) and bounded variance with mini-batch variance at most σ²/B (Assumption 2); whether the results hold under weaker or alternative geometric assumptions is not addressed.

Target Audience

• Optimization and machine learning theorists interested in nonconvex convergence rates for matrix-parameter updates, Schatten-norm geometry, and the analysis of iterative polynomial approximations.
• Researchers working on optimizer design, particularly those studying Muon, Orthogonal-SGDM, Shampoo, SOAP, Adam, and SGD with momentum.
• Practitioners training large models — large language models, transformers, and convolutional networks — who choose the number of Newton–Schulz steps and the polynomial degree and want principled guidance on the accuracy–compute trade-off.
• Engineers focused on training efficiency, for whom the wall-clock argument (SVD-free, matrix-multiplication-only orthogonalization) is the most directly actionable takeaway.

Authors’ abstract

We analyze Muon as originally proposed and used in practice -- using the momentum orthogonalization with a few Newton-Schulz steps. The prior theoretical results replace this key step in Muon with an exact SVD-based polar factor. We prove that Muon with Newton-Schulz converges to a stationary point at the same rate as the SVD-polar idealization, up to a constant factor for a given number $q$ of Newton-Schulz steps. We further analyze this constant factor and prove that it converges to 1 doubly exponentially in $q$ and improves with the degree of the polynomial used in Newton-Schulz for approximating the orthogonalization direction. We also prove that Muon removes the typical square-root-of-rank loss compared to its vector-based counterpart, SGD with momentum. Our results explain why Muon with a few low-degree Newton-Schulz steps matches exact-polar (SVD) behavior at a much faster wall-clock time and explain how much momentum matrix orthogonalization via Newton-Schulz benefits over the vector-based optimizer. Overall, our theory justifies the practical Newton-Schulz design of Muon, narrowing its practice-theory gap.