PaperlayerView exact revision on arXiv

Source-linked AI summary

Convergence of Muon with Newton-Schulz

Gyu Yeol Kim, Min-hwan Oh

arXiv:2601.19156v1stat.MLcs.LGmath.OC

TL;DR

Practical MUON uses a few Newton–Schulz steps to orthogonalize momentum, while prior theory largely analyzed an exact SVD-polar substitute. This paper analyzes the practical method and proves matching stationary-point convergence up to a rapidly vanishing factor, while showing sharper rank dependence than SGD with momentum. The results explain why low-degree, few-step Newton–Schulz can retain exact-polar behavior with lower per-iteration cost.

  • Problem

    Prior MUON analyses typically replace practical finite-step Newton–Schulz orthogonalization with an exact SVD polar step, leaving practical nonconvex convergence and approximation effects unresolved.

  • Method

    The paper analyzes MUON with finite Newton–Schulz momentum orthogonalization using a one-step descent inequality that explicitly incorporates polar approximation error εq.

  • Results

    MUON with Newton–Schulz converges to a stationary point at the SVD-polar iteration rate up to a constant factor, while weakening rank dependence relative to SGD with momentum.

  • Takeaways & Limitations

    The approximation factor approaches 1 doubly exponentially with q and improves with polynomial degree κ, supporting a few-step, SVD-free, GPU-friendly implementation.

  • Takeaways & Limitations

    The convergence analysis assumes a continuously differentiable L-smooth nonconvex objective and bounded mini-batch variance under an operator–nuclear geometry.

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper closes the gap between practical MUON, which uses finite Newton–Schulz orthogonalization, and prior analyses based on exact SVD polar steps. It proves comparable convergence while explaining the effects of Newton–Schulz accuracy, polynomial degree, and matrix structure.

  • Motivation: MUON respects matrix structure by orthogonalizing momentum matrices instead of applying vector-based updates to vectorized parameters.This targets search directions aligned with matrix geometry in linear layers and attention projections.
  • Open questions: Prior MUON theory typically replaces practical Newton–Schulz orthogonalization with an exact SVD-based polar step, leaving finite-step convergence unresolved.The paper asks whether the SVD-free implementation admits nonconvex guarantees and how its rates compare with the idealization.
  • Comparisons: The paper also compares MUON with SGD using momentum to study whether matrix momentum orthogonalization avoids rank-dependent losses.This comparison addresses the accuracy–compute trade-off and the geometric mechanisms behind any convergence gap.
  • Main results: The analysis proves that practical MUON converges to a stationary point at the SVD-polar rate up to a constant depending on the Newton–Schulz approximation error εq.The result applies under standard smoothness assumptions and measures stationarity using the gradient’s nuclear norm.
  • Main results: The approximation error εq shrinks doubly exponentially with the number q of Newton–Schulz steps and improves as the polynomial degree κ increases.Thus, only a few steps can make the practical convergence rate arbitrarily close to the exact-polar rate.

2 RELATED WORK

Related work distinguishes MUON’s SVD-free momentum orthogonalization from both earlier SVD-based methods and curvature-preconditioning optimizers. The paper’s novelty is to analyze finite Newton–Schulz accuracy, rank dependence, and wall-clock implications directly.

  • Muon and momentum orthogonalization: Earlier Orthogonal-SGDM orthogonalizes gradients by SVD before applying momentum, whereas MUON applies momentum first and uses Newton–Schulz instead.The MUON design uses matrix multiplications and is therefore faster to compute than SVD-based orthogonalization.
  • Second-order preconditioners vs. MUON: MUON is not a second-order preconditioner: it orthogonalizes momentum without estimating or inverting curvature.Curvature preconditioning and momentum projection target different mechanisms and may be complementary.
  • Practical efficiency of MUON: Practical MUON reports motivate theory that is SVD-free and GPU-aligned, matching the optimizer’s efficient matrix-multiplication implementation.The analysis is designed around the Newton–Schulz step used in practice and can extend to other MUON variants.
  • Convergence analysis of MUON: Prior analyses either assume exact SVD polar steps, use geometries that obscure rank benefits, or omit Newton–Schulz approximation accuracy.This paper instead analyzes how the step count q and polynomial degree κ affect MUON’s convergence rate.
  • Key distinctions: The paper reports nonconvex stationary-point convergence, a rapidly vanishing Newton–Schulz constant factor, and sharper rank dependence than SGD with momentum.These results connect finite-step theory with MUON’s practical wall-clock behavior.

3 PRELIMINARIES

The paper formulates nonconvex matrix optimization under operator-norm smoothness and bounded stochastic variance, then specifies MUON's Newton–Schulz-based orthogonalization and its error measures.

  • Problem setting: The objective is nonconvex, bounded below, and L-Lipschitz smooth under the operator norm, with the nuclear norm as its dual.These assumptions define the matrix-optimization setting used for convergence analysis.
  • Problem setting: Mini-batch stochastic gradients are unbiased, with variance bounded by σ^2/B for batch size B.The variance assumption is used as a standard stochastic-optimization condition.
  • Problem setting: An ε-stationary point is defined by E[∥∇f(W)∥∗] ≤ ε, a condition stronger than stationarity measured in the Frobenius norm.The nuclear norm is the convergence metric for matrix-valued parameters.
  • MUON algorithm: MUON updates matrix parameters using momentum that is scaled and orthogonalized through q Newton–Schulz polynomial steps before the parameter update.The practical method uses matrix multiplications instead of exact SVD orthogonalization.
  • Newton–Schulz orthogonalization: Newton–Schulz recursively applies a degree-κ polynomial to XX⊤ so the scaled momentum becomes nearly orthogonal while preserving the unit spectral ball.The number of steps controls orthogonality accuracy, and the update is inversion-free and SVD-free.
  • Newton–Schulz orthogonalization: The analysis measures both the output's orthogonality residual and its operator-norm polar approximation error relative to the exact SVD polar factor.These quantities quantify how closely finite Newton–Schulz iterations approximate exact orthogonalization.

4 MAIN RESULTS

The paper proves nonconvex convergence for practical MUON using finite Newton–Schulz orthogonalization, matching exact SVD-polar iteration complexity up to a factor that rapidly approaches one. It also compares MUON with SVD and SGD with momentum under a common nuclear-norm stationarity metric, showing benefits from approximation efficiency and rank dependence.

  • MUON with Newton–Schulz: Theorem 1 gives the first nonconvex convergence guarantee for MUON with a finite number of Newton–Schulz steps, establishing convergence to an ϵ-stationary point.The result analyzes the practically used SVD-free optimizer rather than replacing Newton–Schulz with an exact polar step.
  • Comparison with exact polar: MUON with Newton–Schulz matches the exact SVD-polar iteration complexity up to a multiplicative factor χ_q controlled by the polar-approximation error ε_q.The exact-polar rate is recovered by setting ε_q = 0.
  • Approximation error: χ_q approaches 1 doubly exponentially in q, with faster convergence for larger Newton–Schulz polynomial degree κ.The residual satisfies a multi-step contraction governed by the exponent (κ+1)^q, which controls the approximation error and hence χ_q.
  • Wall-clock implications: A few Newton–Schulz steps make the iteration-complexity gap to exact SVD negligible while retaining substantially lower per-iteration cost through matrix multiplications.This provides the theoretical explanation for the practical wall-clock advantage of SVD-free MUON.
  • Comparison with SGD with momentum: Under the same nuclear-norm stationarity metric, MUON removes the deterministic √r loss relative to SGD with momentum and sharpens rank dependence in stochastic terms.The polar update is aligned with the leading singular structure of the gradient, unlike Frobenius-aligned vector updates.

5 NUMERICAL EXPERIMENTS

Experiments evaluate Newton–Schulz step count and polynomial degree across models and datasets. Increasing the step count makes MUON’s optimization behavior approach SVD-based MUON, while lower-cost Newton–Schulz updates improve wall-clock performance and larger polynomial degrees trade computation for faster loss reduction.

  • Newton–Schulz step ablation: On CIFAR-10 with CIFARNET, MUON with q = 1 outperforms SGD with momentum, while q ∈ {2, 3} nearly coincides with SVD-based MUON in train and test loss.The experiment uses 50 epochs, batch size B = 512, and a degree-2 Newton–Schulz polynomial.
  • Newton–Schulz step ablation: MUON with q = 2 or 3 reaches a given test loss substantially faster in wall-clock time than the SVD variant.The difference reflects the lower per-iteration cost of Newton–Schulz updates.
  • Polynomial degree ablation: Increasing polynomial degree κ ∈ {1, ..., 5} at fixed q = 3 improves loss reduction at a fixed epoch but lengthens each optimization step.The observed accuracy–time trade-off matches the predicted residual contraction δ_{j+1} ≤ δ_j^(κ+1).
  • Rank dependence: Across ranks r ∈ {16, 32, 64, 128, 216}, SGD-M’s epoch-averaged nuclear gradient norm grows with rank, whereas MUON variants remain nearly flat.The SGD-M slope is approximately 0.3, consistent with MUON removing the deterministic √r penalty and weakening stochastic rank dependence.

6 CONCLUSION

The paper proves that practical MUON with finite Newton–Schulz orthogonalization reaches the SVD-polar iteration rate up to a factor that rapidly vanishes toward one, while weakening rank dependence relative to SGD with momentum. These results connect MUON’s practical SVD-free speed with its theoretical convergence behavior.

  • Conclusion: Finite Newton–Schulz steps recover the SVD-polar iteration complexity up to χ_q, which shrinks doubly exponentially with q and improves with κ.The resulting method remains SVD-free and GPU-friendly.
  • Conclusion: Compared with SGD with momentum under the same metric, MUON weakens rank dependence and helps close the theory–practice gap.The conclusion attributes MUON’s practical performance to the combination of near-ideal iteration complexity and cheaper matrix-multiplication updates.

USE OF LARGE LANGUAGE MODELS

The paper states that LLMs were used only to improve exposition, while the supplied technical passages develop MUON, its Newton–Schulz orthogonalization, and convergence analysis.

  • LLMs were used solely for writing clarity, grammar, and style, not for research ideation, algorithm design, theoretical analysis, or experiments.
  • Newton–Schulz preserves the unit spectral ball and retains the polar factor across its iterations.
  • The algorithm uses momentum, scaling, q Newton–Schulz steps, and a degree-κ polynomial before updating the weights.
  • Under standard smoothness assumptions, MUON with Newton–Schulz converges to a nuclear-norm stationary point at the idealized SVD-polar rate up to a polar-approximation constant.

C.2 THEOREM 3 (SGD WITH MOMENTUM)

This section analyzes SGD with momentum through scaled momentum recursions, descent arguments, and stochastic-gradient control. The supplied passages emphasize the smoothness and variance assumptions used to obtain a stationary-point complexity.

  • SGD with momentum is represented using scaled stochastic and true full-batch momentum recursions.
  • The proof controls updates using smoothness, operator-to-Frobenius norm bounds, and a learning-rate condition that removes the final descent term.
  • The analysis separates momentum error into deterministic and stochastic terms before taking expectations and averaging over iterations.
  • The nuclear norm is bounded by √r times the Frobenius norm, introducing the rank-dependent comparison factor used in the SGD analysis.
  • The section concludes that an ϵ-nuclear-norm stationary point can be obtained with a stated complexity.

D NEWTON–SCHULZ LEMMAS: PROOFS

The Newton–Schulz lemmas establish that the iterates preserve the relevant matrix structure, retain the exact polar factor, and reduce orthogonality residuals rapidly with polynomial degree.

  • Polar factor invariance: Newton–Schulz updates preserve the polar factor of the momentum matrix at every iteration.
  • Support invariance: The iterates preserve the momentum matrix’s range, so the analysis does not require a full-rank assumption.
  • Spectral update: The update acts spectrally by mapping each eigenvalue λ to λ[pκ(λ)]^2 while keeping the iterates inside the unit spectral ball.
  • Residual decay: For a degree-κ polynomial, the orthogonality residual decays with order κ+1 at each Newton–Schulz step.
  • Residual decay: With sufficiently small initial residual ρ, q steps yield doubly exponential bounds such as ρ^(2^q) or ρ^(3^q), depending on the stated setting.

E WALL-CLOCK VIA COMPUTATIONAL COMPLEXITY.

Newton–Schulz replaces SVD-based orthogonalization with GPU-efficient matrix multiplications. For practical low step counts and polynomial degrees, this produces substantial per-iteration wall-clock savings while retaining near-exact orthogonalization.

  • GEMM throughput is much higher than SVD throughput in practice because GEMMs achieve greater hardware utilization.
  • Each Newton–Schulz step forms XXᵀ and applies a degree-κ polynomial using matrix multiplications, whereas SVD incurs an additional cubic term.
  • For q ∈ {2, 3} and κ ∈ {1, 2}, the algebraic factor is O(0.3∼1), making speedup primarily depend on the GEMM/SVD efficiency ratio.
  • A GEMM/SVD efficiency ratio of 4∼10 typically yields a multi-× per-iteration speedup for Newton–Schulz over SVD.
  • Newton–Schulz scales linearly in q, uses only GPU-efficient GEMMs, and can achieve near-exact orthogonalization for small q and modest κ.

F.1 EXPERIMENTAL SETTING

The experiments compare MUON using finite Newton–Schulz orthogonalization, exact SVD-based polar orthogonalization, and SGD with momentum under matched CIFAR-10 training conditions. Evaluation tracks loss over epochs and wall-clock time while varying Newton–Schulz steps.

  • Task and model: Experiments use CIFAR-10 with 50k training and 10k test examples, channel-wise normalization, and a compact approximately 2M-parameter CIFARNET.The model contains convolutional groups and a linear head, with specified initialization and normalization choices.
  • Evaluation: Performance is measured by cross-entropy train and test loss plus cumulative wall-clock time, aggregated as mean ± 1 standard deviation over 5 runs.Plots provide both epoch-aligned and time-aligned views.
  • Training setup: Training runs for 50 epochs with batch size B = 512 on GPU, using independent seeds and a CPU fallback with B = 256.The default schedule includes a 5% linear warm-up followed by cosine decay to zero.
  • Optimizers: The comparison includes MUON with q ∈ {0, 1, 2, 3} Newton–Schulz steps, exact-polar MUON using SVD, and SGD with momentum.The q = 0 variant is a normalization-only ablation without orthogonalization.
  • Orthogonalization: All optimizers share the same schedules and auxiliary updates, while MUON orthogonalization uses a fixed degree-2 Newton–Schulz polynomial and Frobenius normalization.The polynomial update is applied to normalized momentum matrices, with a transpose trick for tall matrices.

G ADDITIONAL NUMERICAL EXPERIMENTS

Additional experiments span image, language, and scale settings, while targeted ablations examine polynomial degree, rank dependence, and batch size. They report accuracy–time trade-offs and support MUON’s weaker rank dependence than SGD with momentum.

  • Additional tasks: Additional evaluations cover MLP on MNIST, ResNet-18 on CIFAR-100, WideResNet-28-10 on Tiny-ImageNet, and NanoGPT and GPT-2 models on FineWeb.The compared optimizers are SGD with momentum, MUON with 1–3 Newton–Schulz steps, and exact-SVD MUON.
  • Degree κ ablation: The degree-κ residual recursion contracts as δ_{j+1} ≤ (δ_j)^(κ+1), linking higher polynomial degree to faster orthogonality-residual decay.Coefficients are generated from the Taylor series of λ^(-1/2) at 1.
  • Degree κ ablation: At fixed q = 3, increasing Newton–Schulz polynomial degree κ ∈ {1, …, 5} improves loss at a fixed epoch but increases time per step.The experiment isolates polynomial degree while holding the step count and other training components fixed.
  • Rank dependence: Across ranks r ∈ {16, 32, 64, 128, 216}, SGD-M has a positive log–log slope of about 0.3, whereas both MUON variants are nearly flat.After division by √r, MUON slopes are about −0.4 for SVD and −0.61 for Newton–Schulz, while SGD-M is about −0.21.
  • Batch size: For batch sizes B ∈ {64, 128, 256, 512, 1024}, B = 1024 achieves the best time-to-accuracy in the reported runs.Larger batches reduce gradient noise but also reduce orthogonalization frequency per epoch.

H.4 DEGREE-2 NS POLYNOMIAL VS. AD-HOC DEGREE-2 NS POLYNOMIAL

The section compares a degree-2 Newton–Schulz polynomial with an ad-hoc polynomial using train/test loss and wall-clock plots, while showing that the ad-hoc polynomial violates the monotonicity assumptions required for the guarantees.

  • Figure 11 compares the NS polynomial and ad-hoc polynomial using train/test loss versus epoch and wall-clock.
  • The ad-hoc polynomial p_ad-hoc(λ) = 3.4445 − 4.7750 λ + 2.0315 λ^2 does not remain within [0, 1] on [0, 1].Its range is [0.701, 3.4445].
  • τ_ad-hoc is not monotone non-decreasing on [0, 1], because its derivative is negative at λ = 1.The reported endpoint value is τ′_ad-hoc(1) = −0.5087.
  • The function increases and then decreases, with a local maximum at approximately λ ≈ 0.308.
  • Because the monotonicity premise of Lemma 2 fails for p_ad-hoc, the guarantees apply to p_κ but not to the ad-hoc quadratic, which remains an empirical heuristic.The condition τ(1) = p(1)^2 = 0.7012 < 1 nevertheless holds.
Loading 2601.19156v1…