PaperlayerView exact revision on arXiv

Source-linked AI summary

Convergence of Practical Muon with Finite Newton-Schulz Iterations and Nesterov Momentum

Hanyng Peng, Hui Wang, Yue Yu

arXiv:2609.39595v1cs.LG

TL;DR

Practical Muon combines layer-wise finite Newton–Schulz updates and Nesterov momentum, but its convergence analysis must retain these features on a coupled objective. The paper analyzes that iteration with joint gradient-dependent smoothness and bounded layer-wise variance. It proves an O(T^-1/4) expected average Frobenius gradient-norm bound under coupled schedules, without explicit dimension or rank factors when the block count and problem constants are fixed.

  • Problem

    Existing Muon analyses treat different ingredients, leaving the combined layer-wise finite-polynomial update with Nesterov momentum to be analyzed directly on a coupled objective.

    Source: §1 · pp. 1–2

  • Method

    A joint descent-and-tracking analysis preserves each block’s finite Newton–Schulz map and Nesterov recursion under gradient-dependent smoothness and bounded layer-wise conditional variance.

    Source: §1 · p. 2, §3 · p. 5, §5 · p. 12

  • Results

    O(T^-1/4) expected average Frobenius gradient-norm bound, with no explicit matrix-dimension or rank factors for fixed block count and problem constants.

    Source: §1 · p. 3, §5 · p. 12

  • Takeaways & Limitations

    Finite orthogonalization provides sufficient alignment and update-energy control for convergence without a positive lower bound on nonzero output singular values.

    Source: §1 · p. 3, §5 · p. 12

    The theorem establishes dimension independence but does not claim a uniformly small Newton–Schulz constant.

    Source: §4.2 · p. 11

Abstract

from arXiv · show

Practical Muon maintains momentum and performs a small, fixed number of Newton--Schulz iterations separately for each parameter matrix, often with a Nesterov correction. We analyze these layer-wise finite-step updates jointly on a coupled nonconvex objective, rather than replacing them by exact polar factors or one global orthogonalization. Under gradient-dependent $(\mathcal L_0,\mathcal L_1,q)$-smoothness and conditionally unbiased stochastic gradients with bounded layer-wise variance, we establish an $\mathcal O(T^{-1/4})$ bound on the expected average Frobenius gradient norm. The analysis retains the Nesterov recursion and requires neither bounded stochastic gradients, symmetric noise, nor a uniform positive lower bound on the nonzero output singular values. Its constants contain no explicit matrix-dimension or rank factors when the number of blocks and problem constants are fixed. The proof follows a descent inequality and a decomposition of the momentum tracking error into initialization, noise, and drift. For the original five-step quintic, we verify the required scalar-map bounds analytically; the result also allows step-dependent coefficients satisfying the same bounds. A complementary nuclear-norm result quantifies rank dependence under a stronger spectral condition. The vanishing rate uses coupled learning-rate and momentum schedules, including the standard single-coefficient Nesterov rule.

1 Introduction

Practical Muon combines layer-wise momentum, finite Newton–Schulz orthogonalization, and often a Nesterov correction on a coupled objective. The paper analyzes these features together rather than substituting exact polar factors or one global orthogonalization.

1 Introduction

  • The coupled objective permits different block dimensions and noise levels; all gradient blocks are evaluated at one iterate, and changing one block can affect gradients in others.The analysis uses joint smoothness and does not assume a separable loss or independent noise across layers.Source: §1 · p. 1
  • The contribution is a direct convergence analysis of layer-wise Frobenius-normalized finite polynomial updates with Nesterov momentum under joint gradient-dependent smoothness and bounded layer-wise conditional variance.The paper distinguishes this combination from prior analyses of individual ingredients and notes that bounded variance does not cover the heavier-tailed regimes of Choudhury et al. (2026).Source: §1 · p. 2
  • For a fixed finite iteration count, the scalar map can approach zero on small singular values, so the analysis uses bounded alignment ratios rather than a positive output-singular-value floor.This supports descent and update-energy bounds without a rank factor while retaining the finite map.Source: §1 · p. 3
  • The analysis yields an O(T^-1/4) bound on the expected average Frobenius gradient norm under coupled learning-rate and momentum schedules, including equal-coefficient Nesterov.For fixed block count and problem constants, the bound has no explicit matrix-dimension or rank factors.Source: §1 · p. 3
  • The proof tracks initialization bias, accumulated stochastic noise, and gradient drift, and verifies the scalar-map conditions for the original five-step quintic.The finite-step result needs neither bounded stochastic gradients nor symmetric noise; a complementary nuclear-norm bound has rank dependence under a stronger spectral condition.Source: §1 · p. 3

2 Preliminary

The preliminary setup defines practical Muon as a layer-wise matrix update with momentum, a Nesterov correction, and potentially inexact orthogonalization. Its assumptions control joint objective changes and layer-wise conditional variance without imposing bounded stochastic gradients or independent noise across layers.

2 Preliminary

  • The momentum setting includes ordinary momentum and the single-coefficient Nesterov correction used in the reference implementation.The orthogonalization routine may be inexact, with its finite Newton–Schulz specialization specified later.Source: §2 · p. 3
  • Each matrix block is normalized using its own momentum norm, while all gradient blocks are evaluated at the same iterate before the joint update.The update is not cyclic block-coordinate descent, and cross-block derivatives are not set to zero.Source: §2 · p. 3
  • The iteration retains layer-wise matrix updates, a finite Newton–Schulz map, and a Nesterov correction rather than analyzing an exact-polar or globally orthogonalized surrogate.The mathematical iteration is studied in exact arithmetic, without decoupled weight decay or a separate optimizer for other parameter groups.Source: §2 · p. 3
  • Finite polynomial orthogonalization maps normalized singular values through ϕJ, where J is the number of polynomial steps; a positive floor on nonzero output singular values is an additional condition, not an automatic consequence.The output values are assumed nonnegative in the stated spectral setup.Source: §2 · p. 4
  • The objective assumptions include global gradient-dependent smoothness for simultaneous block changes and lower boundedness on the parameter space.When q = 0, the smoothness power is interpreted as 1, including at a zero gradient.Source: §2 · p. 4
  • The stochastic oracle is conditionally unbiased with layer-wise variance bounds, but the analysis does not require bounded gradient magnitudes, symmetric noise, sub-Gaussian tails, or independence across layers.The sampling index is independent of the past, which includes the current iterate.Source: §2 · pp. 4–5

3 Main Results

The results use a shared descent-and-tracking argument to analyze layer-wise Muon on a coupled objective. Finite-map control yields a dimension-independent Frobenius stationarity rate without a positive output singular-value floor, while a complementary nuclear-norm result retains rank dependence under a stronger spectral condition.

3.1 A common descent and tracking argument

  • The common proof telescopes a descent inequality and bounds Nesterov tracking error through initialization, stochastic noise, and gradient drift under joint updates.Layer-specific noise levels are retained, and the argument controls drift for the joint move on the coupled objective.Source: §3.1 · p. 5, §3 · p. 5
  • The alternative output-singular-value route gives a nuclear-norm stationarity bound with explicit rank dependence and requires a positive output singular-value floor.The finite-map estimates instead lead to the Frobenius result without that spectral condition.Source: §3.1 · p. 5
  • Finite-map control uses the ratio of output to input singular values, avoiding a positive output singular-value floor and rank bounds for momentum or tracking error.The Frobenius stationarity result therefore does not require a rank conversion.Source: §3.1 · p. 7, §3.2 · p. 8
  • For the original five-step quintic, certified analytic scalar-map bounds are available; the stated J = 5 bounds include ¯ρ5 < 13.371.The bounds are certified, not claimed to be sharp.Source: §3.1 · p. 8

3.2 Dimension-independent convergence of practical Muon

  • Theorem 3.4 and Corollary 3.5 give a dimension-independent O(T^-1/4) Frobenius stationarity rate under coupled learning-rate and momentum schedules, including the standard Nesterov choice.The bound has no explicit matrix-dimension or rank factor when the block count and problem constants are fixed; the constants need not remain fixed across architectures.Source: §3.2 · p. 9

4 Discussion

The discussion frames Muon’s schedules as coupled theoretical choices, not architecture-agnostic rules, and distinguishes rank-sensitive nuclear-norm guarantees from rank-free Frobenius guarantees. Dimension independence and implementation generality hold only within stated bounds on problem constants and output scaling.

4.1 Momentum schedules

  • The vanishing rate relies on a scheduled first-momentum gap; fixing it leaves a stochastic term that does not vanish with the horizon.The Nesterov correction does not change the leading O(T^-1/4) rate, though its coefficients trade off initialization, drift, and stochastic terms.Source: §4.1 · pp. 10–11
  • Corollary 3.2’s rank-sensitive schedule balances noise, accumulated rank, and horizon; Corollary 3.5’s finite-step Frobenius schedule has no explicit rank factor.The rank-sensitive stepsize is not a stand-alone law under arbitrary fixed momentum.Source: §4.1 · p. 10

4.2 Dimension independence

  • With fixed block count and problem constants, Corollary 3.5 has no explicit dimension or rank factor, but its constants need not remain fixed across architectures.The finite-map constants depend on the polynomial map and need not be small.Source: §4.2 · p. 11

4.3 Implementation scope

  • The proof covers inexact orthogonalization with controlled alignment and magnitude, but does not assert a finite-precision error guarantee.Its alternative finite-map route avoids a uniform positive output-singular-value floor by establishing Frobenius rather than nuclear-norm stationarity.Source: §4.3 · p. 12
  • Fixed output multipliers preserve dimension independence only when their maximum and the reciprocal of their minimum stay bounded independently of matrix size.Unbounded size-dependent multipliers cannot be absorbed into a dimension-independent constant.Source: §4.3 · p. 12

5 Conclusion

The conclusion establishes convergence for layer-wise Muon with finite Newton–Schulz updates and Nesterov momentum, without requiring a positive lower bound on nonzero output singular values. It also identifies a complementary rank-dependent result under stronger spectral control.

5 Conclusion

  • O(T^-1/4) expected average Frobenius gradient norm is obtained with coupled learning-rate and momentum schedules, without explicit dimension or rank factors when block count and problem constants are fixed.The result includes the standard equal-coefficient Nesterov rule and analytically verified scalar-map conditions for the original five-step quintic.Source: §5 · p. 12
  • The momentum tracking error is separated into initialization, stochastic-noise, and gradient-drift terms.The shared descent-and-tracking argument supports the convergence bound.Source: §5 · p. 12
  • Finite orthogonalization provides sufficient alignment and update-energy control without a positive lower bound on nonzero output singular values.A complementary nuclear-norm result quantifies rank dependence under stronger spectral control.Source: §5 · p. 12

A Useful Lemmas

  • The threshold in s corresponding to the smallest feasible x is within a factor of 2 of the optimal threshold.The lemma characterizes the threshold through the smallest positive solution of its stated inequality.Source: A Useful Lemmas · p. 14

B Proof of Theorem 3.1

The proof combines a descent inequality with a decomposition of Nesterov momentum tracking error into initialization, noise, and drift. Joint control of these terms yields the stated convergence bound while accounting for simultaneous coupled block updates.

B Proof of Theorem 3.1

  • The descent step uses the inexact orthogonalization direction and telescopes the resulting expected inequality over iterations.The proof then rearranges terms using the stated smoothness assumption.Source: B Proof of Theorem 3.1 · pp. 16–17
  • The Nesterov tracking discrepancy is unrolled and decomposed into initialization, stochastic-noise, and gradient-drift contributions.The proof bounds the initialization term by a geometric series.Source: B Proof of Theorem 3.1 · p. 18
  • Conditional unbiasedness removes noise cross terms, while the layer-wise variance bound controls the remaining stochastic-gradient contribution.The noise estimate uses Jensen’s inequality and the martingale-difference property.Source: B Proof of Theorem 3.1 · pp. 18–19
  • Gradient drift is bounded across layers using joint smoothness and the update-energy bound, rather than a single-layer bound on a displacement in all layers.The argument uses Cauchy–Schwarz across layers and controls the geometric sum.Source: B Proof of Theorem 3.1 · p. 19
  • Telescoping the descent inequality and absorbing the update terms gives the claimed convergence bound, with the absorption coefficient accounting for simultaneous coupled block updates.The proof applies joint smoothness to simultaneous changes rather than applying a one-layer displacement bound across all layers.Source: B Proof of Theorem 3.1 · pp. 19–20

C Proof of Corollary 3.2

The corollary proof balances schedule-dependent terms to obtain the vanishing-rate guarantee. Its schedule conditions ensure the Nesterov momentum projection does not change the final rate.

C Proof of Corollary 3.2

  • For a > 1/2, 1(1 − β1)2a = O(1 − β1); since β2 ≤ 1, the upper envelope of the three leading terms in Eq. (44) is minimized.Source: C Proof of Corollary 3.2 · p. 20
  • The minimized upper envelope has a leading T^-1/4 term after balancing the schedule parameters.The displayed bound includes problem constants and the initial objective gap.Source: C Proof of Corollary 3.2 · pp. 20–21
  • The proof chooses the horizon using the stated threshold condition, applying Lemma A.3 or a direct check when qL1 = 0.It then reformulates the bound for substitution into the convergence argument.Source: C Proof of Corollary 3.2 · p. 21

D.1 Proof of Proposition 3.3

The finite Newton–Schulz map is analyzed through its action on normalized singular values and scalar-map bounds. For the fixed quintic, the proof verifies those bounds analytically and notes that step-dependent coefficients require their own checks.

D.1 Proof of Proposition 3.3

  • Each polynomial step preserves the singular-vector bases, and the resulting scalar map yields bounds used to control the finite orthogonalization direction.The proof derives an angle bound by maximizing an expression over A > 0.Source: §D.1 · pp. 21–22

D.2 Analytic constants for the fixed quintic

  • For the original quintic, rational brackets for its extrema establish the interval bounds, and iteration yields ϕJ(s) ≥ (17/25)s for s in (0, 1].The proof also bounds all iterates in [0, 5/4].Source: §D.2 · p. 22
  • The fixed-coefficient quintic’s scalar constants do not bound a matrix’s smallest nonzero singular value; step-dependent schedules require positivity and reachable intervals to be checked separately.The required estimates also hold when Nt = 0, with Ot = 0.Source: §D.2 · p. 22, §D.1 · p. 22

D.3 Proof of Theorem 3.4

The proof combines descent and telescoping with a three-part expansion of each block’s Nesterov tracking error. It bounds initialization, stochastic noise, and gradient drift before absorbing the remaining gradient term.

D.3 Proof of Theorem 3.4

  • The proof telescopes a descent inequality, expands Nesterov tracking error, bounds its three components, and absorbs the remaining gradient term.The orthogonalization routine enters through alignment and update-energy bounds.Source: §D.3 · p. 23
  • Initialization error is bounded by a geometric series, completing the three-part tracking-error control used in the descent argument.The proof then substitutes the bounds and absorbs the remaining term to obtain the result.Source: §D.3 · pp. 23–24
  • Conditional unbiasedness removes temporal cross terms within each block, so its variance bound and a geometric series control the noise term without requiring cross-layer independence.The stated estimate is expressed through the layer-specific variance level.Source: §D.3 · pp. 23–24
  • The drift bound applies Cauchy–Schwarz across blocks before joint smoothness, controlling changes in true gradients under simultaneous updates.The argument uses the iterate difference and bounds the resulting finite sums.Source: §D.3 · p. 24

D.4 Proof of Corollary 3.5

The schedules balance the leading terms in the convergence bound, with learning rate and momentum gap decaying at specified powers of T. The argument also addresses zero noise or drift constants and bounds the initial gradient under the stated optimization assumptions.

D.4 Proof of Corollary 3.5

  • γ ≍ T^-3/4 and δ ≍ T^-1/2 balance the leading terms in the convergence bound.The corollary allows any positive γ0 and δ0 with these exponents, including verification when D = 0 or SF = 0.Source: §D.4 · p. 24
  • For a > 1/2, all exponents other than the leading 1/4 are strictly larger, yielding the stated Frobenius stationarity claim under the corollary’s threshold condition.The proof uses ∥∇F(Xt)∥F ≤ Ht.Source: §D.4 · pp. 24–25
  • The initial gradient bound follows from the stated optimization assumptions and is not a claim about a network-width limit.The proof applies a smoothness-based inequality at X1 and uses the objective lower bound.Source: §D.4 · p. 25

D.5 Scope of the angular refinement

The angular refinement is conditional: it adds spectral-direction smoothness and uses a sufficient test to give a deterministic certificate. The discussion also distinguishes limits of the test from limitations of the finite-map update bound.

D.5 Scope of the angular refinement

  • Failure of the sufficient test does not establish negative alignment, and a small gradient in one block does not imply stationarity of the full parameter tuple.The test only indicates that its sufficient lower bound is unavailable.Source: §D.5 · p. 25
  • Spectral-direction smoothness is an additional hypothesis, not implied by Assumption 2.3 with a dimension-independent constant.The refinement defines hγ from that additional condition and uses it in a sufficient test involving the block momentum and tracking error.Source: §D.5 · p. 25
  • If the tracking error stays uniformly below ē on a sample path, the test yields either a bound at its first failure or a bound on the minimum gradient when it holds throughout.The certificate depends on a pathwise maximum error, including initialization transients; expected error bounds cannot replace it.Source: §D.5 · p. 25
  • For the fixed quintic, any uniform Frobenius update bound needs magnitude constant at least uJ, even though the optimization bound has no explicit dimension factor.The argument uses equal normalized singular values at rank r to show the universal magnitude constant cannot generally be removed.Source: §D.5 · pp. 25–26
Loading 2609.39595v1…