PaperlayerView exact revision on arXiv

Source-linked AI summary

Convergence guarantees for Muon: New parameter regimes and generalizations

Arthur C. B. de Oliveira, Dhruv D. Jatkar, Guilherme S. Vicinansa, Eduardo D. Sontag

arXiv:2609.30546v1math.NAcs.AIcs.NAmath.OC

TL;DR

Classical full-sequence convergence guarantees for Muon are challenging, and existing analyses do not provide them. The paper interprets regularized Muon as preconditioned heavy-ball, extends the normalization structure to Muesterov, and proves convergence results; its experiments illustrate the role of learning rate and regularization in convergence behavior.

  • Problem

    Classical full-sequence convergence guarantees for Muon are lacking, while richer optimization methods can require maintaining and manipulating additional matrix statistics.

    Source: §1 · p. 2

  • Method

    The paper treats Newton–Schulz regularization as a bounded preconditioner for heavy-ball analysis and applies the same normalization structure to extrapolated gradients in Muesterov.

    Source: §1 · p. 2, §2.1 · p. 3, §3.1 · p. 6

  • Results

    With suitable hyperparameters, regularized Muon’s gradient norms converge to zero, and function values converge linearly under a global Polyak–Łojasiewicz condition; Muesterov has analogous guarantees.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Synthetic and transformer-training experiments suggest that learning rate and Newton–Schulz regularization jointly control the transition between oscillatory sign-like dynamics and stable convergence.

    Source: §1 · p. 2, §6 · p. 10

    The general convergence theorem imposes a momentum restriction that the authors say is conservative for Muon specifically and may be tightened using Newton–Schulz structure.

    Source: §3.1 · p. 6

Abstract

from arXiv · show

In this paper, we establish the first asymptotic convergence guarantees for the Muon algorithm through a more accurate proxy for the Newton-Schultz iteration than the typical matrix sign function. We prove that, for appropriate choices of hyperparameters, the iterates satisfy $\lim_{k\to\infty}\|\nabla f(x_k)\|=0$, and, under a global Polyak-\L{}ojasiewicz condition, that the sequence of function values converges linearly. The key insight is that the regularization, implicit in Muon's Newton-Schulz implementation, induces a bounded preconditioner, exposing Muon as a \emph{preconditioned Polyak heavy-ball} method and enabling a classical Lyapunov analysis. This observation naturally motivates applying the same preconditioning structure to the Nesterov gradient evaluation. We formalize this idea by introducing \emph{Muesterov}, a Nesterov-based variant of Muon, and prove that it enjoys the same convergence guarantees, extending the theoretical framework beyond the heavy-ball setting. Numerical experiments on a scalar cross-entropy problem corroborate the theory and illuminate the joint role of the learning rate and the Newton-Schulz regularizer in controlling convergence. Preliminary numerical simulations training the nanoGPT dataset provide intuition regarding the relevance of the observations in this paper to practical applications.

1 Introduction

Richer matrix-based optimization methods can improve convergence but often require additional matrix statistics, while classical full-sequence convergence guarantees for Muon remain challenging. The paper uses regularization to analyze Muon as preconditioned heavy-ball, extends the approach to Muesterov, and studies how learning rate and regularization relate to convergence.

1 Introduction

  • Muon approximates the matrix sign with a few Newton–Schulz iterations, an efficient route to gradient-matrix normalization that can be implemented in low precision and parallelized.The paper situates this approach among richer gradient and curvature methods that typically maintain and manipulate additional matrix statistics.Source: §1 · p. 2
  • Existing Muon analyses do not provide classical full-sequence convergence guarantees, in part because the exact matrix-sign preconditioner is not uniformly bounded.For rank-r matrices, the sign operator has Frobenius norm √r, obstructing direct use of standard Lyapunov arguments.Source: §1 · p. 2
  • The paper establishes asymptotic convergence guarantees for regularized Muon by interpreting its implicit regularization as a bounded preconditioner for Polyak heavy-ball.With suitable hyperparameters, gradients converge to zero; under a global Polyak–Łojasiewicz condition, function values converge linearly.Source: §1 · p. 2
  • The paper introduces Muesterov by applying Muon’s normalization structure to extrapolated gradients, and proves analogous convergence guarantees.The variant is motivated by the cited advantages of Nesterov over Polyak heavy-ball.Source: §1 · p. 2
  • Synthetic and transformer-training experiments examine how learning rate and Newton–Schulz regularization jointly relate to oscillatory sign-like dynamics and stable convergence.The fixed-batch nanoGPT experiments are described as preliminary illustrations of the qualitative behavior, not benchmark evaluations.Source: §1 · p. 2

2 Problem Setup and Preliminary Analysis

The analysis considers smooth optimization problems, including a global Polyak–Łojasiewicz setting, and relates practical Muon to idealized sign-based and regularized soft-sign formulations. The formulations coincide under stated conditions, but regularization can produce vanishing updates near critical points and restore asymptotic convergence.

2 Problem Setup and Preliminary Analysis

  • The setup considers differentiable matrix-valued objectives, with smoothness and global Polyak–Łojasiewicz assumptions used in the convergence analysis.The stated PL setting also assumes a proper objective with a finite attained minimum.Source: §2.1 · p. 3, §2 · p. 3

2.1 Equivalent Muon formulations

  • Canonical Muon applies Newton–Schulz iterations to the gradient matrix normalized by its Frobenius norm plus a regularizer.The analysis introduces ideal matrix-sign and regularized soft-sign formulations as proxies for the practical algorithm.Source: §2.1 · p. 3
  • Under the stated initialization and scaling, canonical Muon and ideal Muon generate identical parameter trajectories, with momentum differing by a rescaling.This identifies ideal Muon with a preconditioned Polyak heavy-ball iteration.Source: §2.1 · pp. 3–4
  • Unlike ideal Muon’s constant-magnitude sign updates, soft-sign normalization permits vanishing updates near critical points, enabling asymptotic convergence analysis.The soft-sign map applies scalar soft sign to singular values and is well-defined for all input matrices.Source: §2.1 · p. 4

2.2 Why so many Muons?

  • When regularization is weak, the three formulations behave nearly identically and fail to converge asymptotically; with sufficient regularization, canonical and soft Muon converge while ideal Muon orbits the minimizer.The paper aligns canonical and soft Muon’s local behavior near the origin by matching their scalar first-order expansions, with δ ≈ 484.876 ε in the simulations.Source: Figure 1 · p. 4, §2.2 · p. 5

3 Muon and the preconditioned Polyak heavy-ball

With a uniformly bounded preconditioner, preconditioned heavy-ball methods converge to a stationary point for suitable learning rates and momentum; under a global Polyak–Łojasiewicz condition, function values converge linearly. Soft Muon obtains the needed bounds through trajectory boundedness under stated assumptions.

3 Muon and the preconditioned Polyak heavy-ball

  • With a uniformly bounded preconditioner and suitable learning-rate and momentum choices, preconditioned heavy-ball iterates converge to a first-order stationary point.Under a global Polyak–Łojasiewicz condition, function values also converge linearly.Source: §3 · p. 5

3.1 Why soft Muon works

  • Soft Muon’s regularization gives the preconditioner a uniform lower bound, but a uniform upper bound requires bounded iterates.The preconditioner can grow with the momentum variable, so establishing trajectory boundedness is the key step for applying the general convergence result.Source: §3.1 · pp. 5–6
  • For any initialization, an unconditional bound on the update prevents arbitrarily large one-step excursions and enables a proof that the Soft Muon trajectory is bounded.The bounded-trajectory theorem assumes a proper, L-smooth objective with a unique global minimizer and a global Polyak–Łojasiewicz inequality.Source: §3.1 · p. 6
  • For suitable learning rates and momentum below the stated threshold, Soft Muon has bounded iterates and momentum and its gradient norm tends to zero.The result combines trajectory boundedness with the general preconditioned heavy-ball convergence theorem.Source: §3.1 · p. 6
  • The theorem’s momentum threshold is below the empirically typical value around 0.95, but the paper says this does not imply larger momentum necessarily breaks convergence for Muon.The authors identify exploiting Newton–Schulz structure to tighten the constraint as future work.Source: §3.1 · p. 6
  • The same normalization-based preconditioning perspective motivates investigating whether convergence guarantees extend to Nesterov-type acceleration.Source: §3.1 · p. 6

4 A Nesterov-based Muon-type algorithm

The paper extends preconditioned heavy-ball convergence results to Nesterov methods when the preconditioner is uniformly bounded. It then introduces Muesterov, which applies Muon’s normalization structure to extrapolated gradients and obtains analogous convergence guarantees.

4 A Nesterov-based Muon-type algorithm

  • With a uniformly bounded preconditioner and suitable learning-rate and momentum choices, preconditioned Nesterov methods have gradients at extrapolated points converge to zero.Under a global Polyak–Łojasiewicz condition, function values converge linearly.Source: §4 · p. 7

4.1 The Muesterov algorithm

  • Muesterov replaces Soft Muon’s heavy-ball momentum with a Nesterov extrapolation step while preserving its normalization-based structure.Its gradient is evaluated at the extrapolated point, and the update uses an analogous preconditioner.Source: §4.1 · p. 7
  • Under a global Polyak–Łojasiewicz inequality and suitable learning-rate and momentum choices, Muesterov’s iterates and extrapolated points remain bounded.The bounded trajectories ensure the induced preconditioner remains uniformly bounded, allowing the general Nesterov convergence result to apply.Source: §4.1 · p. 7
  • For suitable parameter choices, Muesterov’s gradient evaluated at extrapolated points converges to zero.The paper presents this as an extension of Muon’s convergence framework to a Nesterov-based setting.Source: §4.1 · p. 7

5 Simulations

Synthetic simulations examine how the learning rate and normalization regularizer shape optimizer dynamics; a fixed-batch nanoGPT experiment tests whether the qualitative interaction is visible in transformer training. Across both settings, stronger regularization can smooth oscillatory behavior, while the nanoGPT results are diagnostic rather than a standard training benchmark.

5.1 Building intuition through a synthetic example

  • The scalar cross-entropy problem isolates normalization and regularization effects in a transparent testbed for Muon, Soft Muon, and Muesterov.The simulations vary ε and α to examine their effects on training cost.Source: §5.1 · pp. 7–8, Figure 2 · p. 8
  • In the scalar cross-entropy simulation, increasing ε stabilizes Soft Muon and Muesterov, while larger δ produces a similar qualitative transition for Canonical Muon.For small regularization, all methods show persistent oscillations; excessively large regularization eventually slows convergence.Source: §5.1 · p. 8
  • Muesterov reaches lower loss faster than Canonical and Soft Muon under identical hyperparameters in the simple 1D simulation.The experiments suggest that learning rate and normalization regularizer must be tuned jointly.Source: §5.1 · p. 8

5.2 Training nanoGPT

  • In fixed-batch nanoGPT training, increasing δ from 0.1 to 1 substantially smooths the trajectory and reduces repeated large-amplitude peak–trough events.The remaining transient is followed by near-monotone descent toward a fixed point.Source: Figure 3 · p. 9, §5.2 · p. 9
  • The nanoGPT experiment deliberately reuses one batch and raises the script-level learning rate, so its loss curve is a fixed-batch diagnostic rather than a sampled training curve.The observed regularization interaction is qualitative evidence in a controlled setting, not a standard large-scale benchmark.Source: §5.2 · pp. 8–9

6 Conclusion

The paper interprets Newton–Schulz regularization as bounded preconditioning, yielding asymptotic convergence guarantees for Muon through a classical Lyapunov analysis. It also introduces Muesterov, while experiments suggest that learning rate and regularization jointly control oscillatory versus stable dynamics.

6 Conclusion

  • The implicit Newton–Schulz regularization acts as bounded preconditioning, enabling a classical Lyapunov analysis and the first asymptotic convergence guarantees for Muon.The guarantees hold under suitable hyperparameter choices.Source: §6 · p. 10
  • Muesterov preserves Muon’s normalization structure in a Nesterov-based variant and enjoys analogous convergence guarantees.Source: §6 · p. 10
  • Experiments suggest that learning rate and Newton–Schulz regularization jointly control the transition between oscillatory sign-like dynamics and stable convergence.Understanding these mechanisms in practical large-scale training remains future work.Source: §6 · p. 10

A.1 Proof of Proposition 1

The proof establishes an algebraic equivalence between the sign-based update sequences after rescaling the momentum variable. Positive scaling preserves the sign operator, so the parameter updates agree.

A.1 Proof of Proposition 1

  • With N(Bk) = sign(Bk), the two sequences have identical parameter iterates and momentum variables related by the stated rescaling.Source: §A.1 · p. 11
  • The proof uses sign(Hk) = sign(Bk) under the positive scaling between the momentum variables.Source: §A.1 · p. 11
  • Because both sequences share initialization and update law, the scaled momentum relation yields the claimed trajectory equivalence.Source: §A.1 · p. 11

A.2 Proof of Proposition 2

The proof establishes equivalence between the soft-sign update sequences when the regularizer is scaled with the momentum variable. The resulting soft-sign evaluations and parameter updates coincide.

A.2 Proof of Proposition 2

  • With N(Bk) = S(Bk, ε), the sequences have identical parameter iterates and momentum variables related by the stated rescaling.Source: §A.2 · p. 11
  • Scaling both the momentum matrix and regularizer by 1 − β preserves the soft-sign evaluation: S(Hk, ε(1−β)) = S(Bk, ε).Source: §A.2 · p. 11
  • Shared initialization and update laws then give the claimed equivalence of parameter trajectories.Source: §A.2 · p. 11

A.3 Proof of Theorem 1

The proof establishes a Lyapunov decrease for the preconditioned heavy-ball iteration when the preconditioner is uniformly bounded and the momentum and step size are sufficiently small.

A.3 Proof of Theorem 1

  • For uniformly bounded P_k, the theorem provides a positive step-size threshold α⋆(β) for each momentum β in [0, β⋆).The statement assumes p_min I ⪯ P_k ⪯ p_max I for all k.Source: §A.3 · p. 12
  • The proof uses Φ_k = Φ(x_k, e_k) as a Lyapunov function and seeks strict decrease for sufficiently small α and β.The lag error is e_k = H_k − ∇f(x_k).Source: §A.3 · p. 12
  • The Lyapunov difference is nonpositive when its quadratic form in ∥G_k∥ and ∥e_k∥ is positive semidefinite.The stated conditions are C_G > 0, C_e > 0, and 4C_GC_e > C_c^2.Source: §A.3 · p. 12
  • The proof selects α⋆(β) as the smaller positive root of Q, ensuring the required coefficient and discriminant conditions hold for α in (0, α⋆(β)).The root construction uses positivity of Q on that interval and verifies the coefficient conditions there.Source: §A.3 · p. 13

A.4 Proof of Corollary 1

Under the global Polyak–Łojasiewicz condition and a uniformly bounded preconditioner, the proof obtains a geometric decrease of its nonnegative Lyapunov function, yielding linear convergence of function values.

A.4 Proof of Corollary 1

  • With a global Polyak–Łojasiewicz condition and uniformly bounded P_k, the proof establishes linear convergence of the function values for admissible β and α.The parameters use the thresholds from Theorem 1.Source: §A.4 · p. 13
  • The proof bounds Φ_k by a fixed factor ρ in (0, 1) at each iteration, then iterates this contraction.It first obtains Φ_{k+1} ≤ (1 − δ)Φ_k for δ > 0.Source: §A.4 · p. 14

A.5 Proof of Theorem 2

The appendix establishes boundedness and convergence results for both heavy-ball and Nesterov variants under stated smoothness, preconditioner, and Polyak–Łojasiewicz assumptions. For each variant, sufficiently small step sizes and an admissible momentum range support stationarity, with linear function-value convergence under the global Polyak–Łojasiewicz condition.

A.6 Proof of Corollary 2

  • For algorithm (8), sufficiently small α with β < 1/2 yields bounded iterates and lim_{k→∞}∥∇f(x_k)∥ = 0.The stated threshold is β⋆ = 1/2.Source: §A.6 · p. 15

A.7 Proof of Theorem 3

  • For the preconditioned Nesterov iteration, uniformly bounded P_k permits an admissible momentum range and positive step-size threshold.The theorem states the result for β in [0, β⋆) and α in (0, α⋆(β)).Source: §A.7 · p. 15

A.8 Proof of Corollary 3

  • Under a global Polyak–Łojasiewicz inequality with µ > 0, for every β ∈ [0, β⋆) and α ∈ (0, α⋆(β)), the nonnegative Lyapunov function contracts geometrically: Φk+1 ≤ ρΦk for some ρ ∈ (0, 1).Source: §A.8 · p. 18

A.9 Proof of Theorem 4

  • For preconditioned Nesterov algorithm (16), a global Polyak–Łojasiewicz condition and sufficiently small α ensure bounded sequences {x_k} and {y_k}.The result requires β in the stated interval [0, β∗).Source: §A.9 · p. 19

A.10 Proof of Corollary 4

  • For algorithm (16), sufficiently small α and admissible β yield bounded iterates and lim_{k→∞}∥∇f(y_k)∥ = 0.The corollary combines the boundedness result with the Nesterov convergence theorem.Source: §A.10 · p. 20

Related papers

Loading 2609.30546v1…