PaperlayerView exact revision on arXiv

Source-linked AI summary

Muon with Finite Newton-Schulz: The Smoothing Benefit in Nonsmooth Nonconvex Optimization

Mingyi Li, Taira Tsuchiya

arXiv:2608.26288v1cs.LGmath.OCstat.ML

TL;DR

Prior analyses largely replace Muon’s finite Newton–Schulz iteration with the exact polar map or treat finite depth as error, leaving its role in nonsmooth optimization unresolved. The paper analyzes the actual finite-depth, momentum-based update through online-to-nonconvex conversion and shows that its smoothing supports convergence, with logarithmic depth in accuracy and stated sample-complexity guarantees.

  • Problem

    Existing theory either replaces finite Newton–Schulz with the exact polar update or treats it as approximation error, leaving whether the actual finite-depth update benefits nonsmooth nonconvex optimization unresolved.

  • Method

    The paper views momentum Muon with finite Newton–Schulz as an online learner with a smoothed spectral potential and converts discounted regret into stationarity guarantees.

  • Results

    q = O(log(1/ε)) Newton–Schulz steps per round find a (ρ, ε)-stationary point in expectation within O(ρ−1ε−3 + ε−2) stochastic gradient evaluations for Lipschitz nonsmooth nonconvex objectives.

  • Takeaways & Limitations

    Finite Newton–Schulz can act as a smoothing mechanism that enables convergence for nonsmooth objectives, whereas exact-polar Muon may fail to converge.

  • Takeaways & Limitations

    The guarantees differ from deployed Muon through fixed pre-run normalization, classical rather than empirically tuned coefficients, and random intermediate gradient queries.

Abstract

from arXiv · show

Muon has emerged as a strong optimizer for the matrix-valued parameters in large language model pretraining, approximately orthogonalizing its momentum with a few Newton-Schulz iterations. Existing theory either replaces this iteration with the exact polar factor it approximates, or treats its finite depth as an approximation error, and thus the iteration Muon actually runs can only hurt the guarantees. We show that finite Newton-Schulz can instead be beneficial for nonsmooth nonconvex optimization. To this end, we analyze Muon through the online-to-nonconvex conversion, which views the update rule as an online learner and converts its regret bound into a stationarity guarantee. The finite Newton-Schulz iteration smooths the discontinuous polar map into a Lipschitz map of the singular values, and Muon with finite Newton-Schulz can be regarded as an online learner with a smoothed spectral potential. This smoothing is exactly what the conversion needs: we prove that a Newton-Schulz depth growing only logarithmically in the target accuracy suffices for convergence to stationary points in nonsmooth nonconvex optimization, whereas Muon with the exact-polar update may fail to converge. The resulting sample complexity bounds match the best-known guarantees for nonsmooth nonconvex optimization and are optimal for smooth nonconvex optimization up to problem-dependent factors. The argument extends beyond Newton-Schulz to general spectral maps with the same smoothing property.

1 Introduction

Muon uses momentum orthogonalized by finite Newton–Schulz iterations, but prior theory either replaces those iterations with the exact polar factor or treats them as error. This paper instead shows that finite depth smooths the update and enables nonsmooth stationarity guarantees.

  • Muon and prior theory: Muon maintains gradient momentum and approximately orthogonalizes it with a few Newton–Schulz matrix-polynomial iterations.The operation requires only matrix–matrix multiplications, supporting Muon’s practical use in large-scale LLM pretraining.
  • Muon and prior theory: Existing analyses either replace finite Newton–Schulz with the exact polar factor or treat finite depth as an approximation error.A separate line handles nonsmooth objectives through online learning but analyzes an update different from the one Muon performs.
  • Main contribution: Finite Newton–Schulz replaces the discontinuous polar map with a Lipschitz singular-value map, enabling the online-to-nonconvex conversion to provide stationarity guarantees.The paper’s answer is affirmative: finite Newton–Schulz can benefit nonsmooth nonconvex optimization rather than merely introduce error.
  • Main guarantee: q = O(log(1/ε)) Newton–Schulz steps per round yield a (ρ, ε)-stationary point in expectation within O(ρ−1ε−3 + ε−2) stochastic gradient evaluations.This informal theorem applies to Lipschitz objectives that may be neither smooth nor convex, with an unbiased stochastic gradient oracle and bounded noise.
  • Technical mechanism: The analysis bounds discounted regret through a smoothed spectral potential whose penalty decreases and stability increases geometrically with Newton–Schulz depth.Balancing these terms gives sublinear discounted regret and yields the nonsmooth convergence result.
  • Technical mechanism: The framework extends beyond Newton–Schulz to general spectral maps with the same smoothing property and covers smooth relaxations of the polar transformation.Finite Newton–Schulz itself induces an FTRL regularizer without requiring a separate relaxation.

2 Optimization setting and the Muon update

The paper formalizes matrix optimization, relaxed stationarity, and a Muon learner that combines momentum with finite Newton–Schulz orthogonalization. Its analyzed procedure uses bounded operator-norm gradients, fixed normalization, random interpolation queries, and a randomized averaged output.

  • Optimization setting: The optimization problem uses matrix parameters and a differentiable objective that is bounded below, without assuming smoothness.For locally Lipschitz nonsmooth objectives, the paper describes a differentiable surrogate obtained by uniform-ball smoothing.
  • Optimization setting: (ρ, ε)-stationarity is defined through finitely supported distributions centered at W whose expected operator-norm displacement is at most ρ.This relaxed criterion is the target used for nonsmooth objectives.
  • Conversion protocol: The discounted O2NC procedure selects an update direction, queries the stochastic gradient at a random point between consecutive iterates, and returns a randomized exponentially weighted average iterate.This protocol retains momentum and finite Newton–Schulz but differs from practical Muon in several respects, including normalization.
  • Muon update: Muon maintains momentum M, assumes ∥Gs∥op ≤ Gop, normalizes M by Gop, applies q Newton–Schulz steps, and outputs an operator-norm-bounded action.The action is scaled by radius D for use by the online learner.
  • Muon update: The analyzed Newton–Schulz map focuses on the degree-five transformation commonly used in Muon and applies q-fold composition to normalized singular values.The scalar map is defined on [0, 1] and extended as 1 for inputs at least 1.
  • Conversion protocol: The learner’s online action determines the optimizer’s update direction through the discounted online-to-nonconvex conversion.The paper explicitly notes that this learner is not identical to practical Muon implementations.

3 Discounted online-to-nonconvex conversion

Discounted online-to-nonconvex conversion uses an online learner to choose update directions and translates discounted regret into relaxed stationarity. In this framework, exact-polar Muon corresponds to follow-the-leader, while finite Newton–Schulz can provide the stability needed for convergence.

  • Conversion mechanism: Discounted O2NC converts an online learner’s regret bound into a guarantee for (ρ, ε)-stationarity.The framework uses the same discount factor β as Muon’s momentum update.
  • Conversion mechanism: The conversion queries the stochastic gradient at a uniformly random point on the segment between consecutive iterates.This choice makes the expected inner product with the update equal the objective change without requiring smoothness.
  • Exact-polar update: Exact-polar Muon is equivalent to follow-the-leader over the operator-norm ball for the discounted loss sequence.Its discounted O2NC analysis therefore depends on the regret behavior of follow-the-leader.
  • Finite-depth smoothing: For normalized singular values, increasing Newton–Schulz depth decreases the polar-approximation gap but increases the map’s Lipschitz constant.This is the penalty–stability tradeoff used in the finite-depth analysis.
  • Exact-polar update: Because follow-the-leader need not have sublinear regret, discounted O2NC alone gives no general nonsmooth stationarity guarantee for exact-polar Muon.This aligns with counterexamples where exact-polar Muon with current-iterate subgradients fails to converge.

4 Discounted regret for a matrix online learner with a general spectral map

The paper analyzes a momentum-based matrix learner using a general spectral map and derives discounted-regret bounds through a smoothed spectral potential. The bounds separate approximation, stability, and potential-variation effects, supporting later nonsmooth and smooth optimization guarantees.

  • Online learner: The online-to-nonconvex conversion reduces stationarity analysis to bounding discounted regret for an online learner over the operator-norm ball.The learner maintains momentum and applies a general spectral map to it.
  • Spectral-map assumptions: The spectral map h is continuous, nondecreasing, bounded in [0, 1], and constrained by Assumption 4.1.Its approximation error measures the gap from the exact polar map, while smaller error requires a larger Lipschitz constant.
  • Smoothed potential: The update is gradient-based prediction with a time-varying smoothed spectral potential evaluated at the momentum.For h = hq and Ct = Gop, it recovers the finite Newton–Schulz update under the stated operator-norm bound; replacing Hh with polar recovers exact-polar Muon.
  • Regret guarantees: Theorem 4.2 provides a coarser regret bound sufficient for nonsmooth nonconvex guarantees and a Bregman-divergence bound used for smooth-objective guarantees.The coarser bound has the standard penalty–stability structure, while the finer bound supports optimal O(1/ε^2) deterministic dependence for smooth objectives.
  • Regret decomposition: The discounted-regret decomposition contains a penalty term, a stability term, and variation of the potentials across rounds.The penalty measures the gap between the base and smoothed potentials, stability accumulates Bregman divergences, and potential variation compares successive potentials.
  • Regret decomposition: The penalty is controlled by approximation error, stability by Lip(h), and the potential-variation term is nonpositive when Cs+1 ≥ βCs.The latter follows from convexity under the scale condition.

5 Regret and stationarity guarantees for Muon with finite Newton–Schulz

Finite Newton–Schulz yields discounted-regret and stationarity guarantees by balancing polar-approximation error against update stability. With logarithmic depth, Muon reaches stationary points for nonsmooth objectives, while smooth-objective rates match established bounds.

  • 5.1 Discounted regret for Muon with finite Newton–Schulz: Muon’s finite Newton–Schulz map is analyzed through online-to-nonconvex conversion as an online learner with a smoothed spectral potential.The conversion translates discounted regret into stationarity guarantees.
  • 5.1 Discounted regret for Muon with finite Newton–Schulz: Lip(hq) = Aq, while the approximation error satisfies 1/(2Aq) ≤ ∆(hq) ≤ 1/(Aq + 1) ≤ 1.Balancing these depth-dependent terms controls the regret bound.
  • 5.1 Discounted regret for Muon with finite Newton–Schulz: As q increases, polar-approximation error decreases but the Lipschitz constant grows, creating a penalty–stability tradeoff absent from exact-polar analysis.The exact-polar limit has vanishing approximation error but an uncontrolled stability term.
  • 5.3 Sample complexity in smooth nonconvex optimization: For smooth objectives, the deterministic ε^-2 and stochastic ε^-4 dependencies match known lower bounds, while leading terms coincide with exact-polar Muon.The comparison is under the stated smoothness and bounded-variance assumptions.
  • 5.4 FTRL interpretation: Finite Newton–Schulz corresponds to FTRL with a depth-controlled spectral regularizer that approaches follow-the-leader as q →∞.At q = 0, the regularizer is Euclidean quadratic on the operator-norm ball.

6 Conclusion and future work

The paper interprets finite Newton–Schulz as smoothing rather than approximation error, enabling nonsmooth stationarity guarantees through a logarithmic-depth penalty–stability balance. The analysis matches exact-polar leading terms for smooth objectives but does not yet cover several deployed Muon variants.

  • 6 Conclusion and future work: Finite Newton–Schulz smooths the polar map, and logarithmic depth balances approximation error with stability to obtain nonsmooth convergence.Exact-polar Muon may fail to converge in this setting.
  • 6 Conclusion and future work: The analysis assumes a fixed known normalization, classical Taylor-truncation coefficients, and gradient queries at random intermediate points rather than practical Muon’s deployed choices.Extending the theory to data-dependent normalization, tuned coefficients, and current-iterate gradients remains future work.

A Comparison with prior convergence guarantees

The appendix compares gradient-oracle guarantees for nonsmooth and smooth nonconvex optimization, including how each method treats polar computation. It also records supporting singular-value lemmas used in the analysis.

  • Nonsmooth objectives: Table 1 compares gradient-oracle complexity under the (ρ, ε)-stationarity criterion for nonsmooth nonconvex objectives.Its notation sets r = m, defines ν, and abbreviates Newton–Schulz as NS.
  • Smooth objectives: Table 2 compares gradient-oracle complexity for ε-stationary points in smooth nonconvex optimization.The comparison specializes to operator-norm smoothness and nuclear-norm stationarity, with deterministic and stochastic columns.
  • Nonsmooth objectives: The leading ρ−1ε−3 term matches the bounds for Pion and Leon.Pion’s stated bound assumes exact expected polar factors, while its Monte Carlo implementation uses k = T samples per round.
  • Nonsmooth objectives: Exact-polar Muon lacks a general nonsmooth stationarity guarantee from the discounted O2NC framework alone.The appendix notes that some parameters for Pion and Leon are not directly comparable.
  • Smooth objectives: The work’s leading smooth-objective terms coincide with prior exact-polar Muon guarantees.Compared with Kim and Oh, its bounds omit χq and improve the middle term from r2σ2/ε2 to rσ2/ε2, although the algorithms differ slightly.
  • Supporting lemmas: The appendix collects singular-value lemmas covering convexity, gradients, Fenchel conjugates, and Lipschitz bounds for spectral functions.These lemmas support the paper’s spectral-potential analysis.

C Proof of Lemma 3.1

The proof models oracle noise as a matrix-valued martingale difference sequence and uses conditional moment bounds to control the stochastic terms. Norm comparison and Jensen’s inequality complete the argument.

  • Proof of Lemma 3.1: The proof defines oracle noise Et := Gt − ∇L(f Wt) and places it within a filtration generated by initialization and round-by-round randomness.Wt and Xt are measurable with respect to the preceding-round history.
  • Proof of Lemma 3.1: Conditional unbiasedness and variance control imply E[Et | Ft−1] = 0 and E[∥Et∥2_F] ≤ σ2.The tower property transfers the conditional assumptions to the filtration used in the proof.
  • Proof of Lemma 3.1: The noise sequence is therefore a matrix-valued martingale difference sequence.This structure yields conditional orthogonality used in the subsequent bound.
  • Proof of Lemma 3.1: The proof applies ∥·∥∗ ≤ √r ∥·∥F and Jensen’s inequality to finish the estimate.The preceding probabilistic bounds are combined to complete the lemma.

D.1 Newton–Schulz iterations at a general Taylor order

The appendix generalizes the finite Newton–Schulz analysis from the degree-five polynomial to every Taylor order κ ≥ 1. The same penalty–stability tradeoff persists, with constants determined by Aκ.

  • General Taylor order: For Taylor order κ ≥ 1, the truncated Newton–Schulz scalar iteration is xj+1 = fκ(xj), where fκ has degree 2κ + 1.The associated map hκ,q is defined by iterating fκ and extending it by 1 for inputs at least 1.
  • General Taylor order: The analysis recovers the Section 5 map at κ = 2, yielding h2,q = hq.Thus the degree-five case is a specialization of the general construction.
  • General Taylor order: The penalty term decreases at rate Aκ^−q for every Taylor order κ.This establishes the same penalty–stability tradeoff used in the main analysis.
  • General Taylor order: The Section 5 argument extends to Newton–Schulz polynomials obtained by truncating the Taylor expansion at any finite order.The generalization replaces A with Aκ in the relevant bounds.
  • General Taylor order: Lip(hκ,q) = Aκ^q, with the map remaining in [0, 1] on that interval.The Lipschitz factor follows from the slope bound at the origin and composition across q iterations.

D.3 Proof of Theorem 5.3

The proof of Theorem 5.3 combines the finite-depth Newton–Schulz regret analysis with stochastic-gradient and smoothing bounds. It selects accuracy-dependent parameters so the resulting terms meet the stationarity target.

  • Theorem 5.3: When ε > BNS, the definition of BNS implies every point W satisfies the required stationarity bound.The proof uses BNS ≥ (1 + 2A)√r Γ > 3√r Γ and handles this case directly.
  • Theorem 5.3: For ε ≤ BNS, the proof sets 1 − β = (ε/3BNS)2 and verifies the depth condition from Theorem 5.1.The chosen depth is shown to equal the depth required by the earlier theorem.
  • Theorem 5.3: The proof controls stochastic-gradient terms using conditional unbiasedness and variance at most σ2.The query-point gradient gs and momentum Ms−1 are measurable with respect to the round history.
  • Theorem 5.3: The Bregman-divergence decomposition separates the update into a conditional-mean component and a zero-mean oracle-noise component.The cross term has zero conditional expectation because E[Ms − c Ms | Qs] = 0.
  • Theorem 5.3: The smoothing map contributes through ∆(hq) ≤ A−q and the Lipschitz-controlled spectral potential.The proof uses the bound ∥∇Ψ(M)∥op ≤ 1 because hq takes values in [0, 1].
  • Theorem 5.3: Choosing ρ := ε/(12Lop) and the corresponding β bounds the six error terms, while controlling the 1/T terms completes the theorem.The proof also selects q = ⌈logA(12rGop/ε)⌉ and uses the resulting bound on Aq.

E.1 General spectral map h (proof of Theorem 5.7)

Finite Newton–Schulz steps induce a singular-value map that can be represented through a smoothed spectral potential and an FTRL update. As depth increases, this update approaches exact-polar Muon.

  • Spectral potential: The smoothed potential’s Fenchel conjugate is obtained by summing its scalar conjugate over the singular values.The scalar conjugate scales as Cϕ*_h for a potential scaled by C.
  • Conjugate bound: For a∈[0,1], the scalar conjugate satisfies ϕ*_h(a)≤∆(h), while for a>1 it is infinite.The bound follows from the scalar conjugate’s construction and the extension of the regularizer.
  • Spectral map: q Newton–Schulz steps replace each singular value x_i with h_q(x_i)=f^q(x_i), preserving singular values in [0,1].The resulting iteration is the singular-value map H_hq(Y).
  • FTRL representation: The q-step Newton–Schulz update for momentum is the FTRL update for discounted linear losses.This connects the implemented matrix iteration to the online-learning analysis.
  • Exact-polar limit: As q→∞, ∆(h_q)→0, so the FTRL update approaches the exact-polar action −D polar(M_{t−1}).The associated endpoint values decrease toward zero with depth.

F Guarantees for other spectral maps

The stationarity analysis extends from finite Newton–Schulz to other spectral maps satisfying the same smoothing assumptions. Their guarantees expose an unavoidable approximation–stability tradeoff.

  • Generalization: General spectral maps satisfying Assumption 4.1 inherit stationarity guarantees for nonsmooth nonconvex objectives.The bound balances approximation error ∆(h) against the map’s Lipschitz constant Lip(h).
  • Approximation–stability tradeoff: A smaller approximation error ∆(h) requires a larger Lipschitz constant Lip(h), and this tradeoff is unavoidable under Assumption 4.1.The result applies to smooth relaxations of the polar map as well as finite Newton–Schulz.
  • Smooth polar relaxation: For the smooth polar relaxation e_h,a, guarantees analogous to Theorems 5.1 and 5.3 are controlled by smoothing parameter a rather than iteration depth q.The relaxation satisfies the required continuity, monotonicity, and range properties.
  • Intrinsic lower bound: Every admissible map obeys ∆(h)≥1/(2 Lip(h)), preventing both approximation error and Lipschitz constant from being simultaneously small.This lower bound makes the tradeoff intrinsic rather than specific to Newton–Schulz.

G Numerical experiments

The synthetic experiment tests whether Newton–Schulz depth produces the predicted penalty–stability tradeoff. Intermediate depths achieve fast, stable descent, while exact-polar and very shallow or deep updates each exhibit a drawback.

  • Numerical experiments: At learning rate 0.05, Muon (SVD) fluctuates around high loss, while q=10 is also unstable.The exact-polar update and the deepest tested finite iteration both fail to provide stable descent at the larger learning rate.
  • Numerical experiments: q=0 is stable but slow, whereas q=2 and q=5 descend quickly and remain stable.The intermediate depths avoid both the instability of larger depth and the slow progress of the shallowest update.
  • Penalty–stability tradeoff: The observed depth dependence qualitatively matches the tradeoff between better polar approximation and a larger Lipschitz constant.The experiment was designed to examine this behavior rather than benchmark individual optimization performance.
  • Baselines: Pion and Leon are stable, but Pion is slower in wall-clock time because each step computes several perturbed polar factors.The comparison uses the same problem instance and learning rates for all methods.
Loading 2608.26288v1…