Source-linked AI summary
A Note on the Convergence of Muon
Jiaxiang Li, Mingyi Hong
TL;DR
The note studies convergence of the Muon optimizer and a closely related spectral steepest-descent scheme for stochastic optimization. It derives convergence results for both versions and identifies a batch-free convergence limitation for the generic scheme.
Problem
The note examines the convergence of Muon, a new optimizer for pretraining large models, and its closely related spectral optimizer.
Method
The analysis characterizes the spectral-norm update through singular-value decomposition and develops a Frobenius-norm descent analysis for the optimizer updates.
Results
The note establishes convergence bounds for Muon and the related generic scheme under specified step-size, momentum, and batch-size conditions.
Takeaways & Limitations
Muon's convergence can be analyzed through specialized descent inequalities, while the generic scheme requires batch growth and does not achieve batch-free convergence under the presented analysis.
Takeaways & Limitations
The generic scheme does not achieve batch-free convergence in this analysis because a key bound uses convexity of the quadratic and norm functions.
Abstract
from arXiv · showhide
In this note, we inspect the convergence of a new optimizer for pretraining LLMs, namely the Muon optimizer. Such an optimizer is closely related to a specialized steepest descent method where the update direction is the minimizer of the quadratic approximation of the objective function under spectral norm. We provide the convergence analysis on both versions of the optimizer and discuss its implications.
1 Introduction
The note introduces Muon as an optimizer for large-model training and relates it to a spectral optimizer whose update uses the SVD of the gradient-related matrix.
- Muon is introduced as a new optimizer motivated by its success in training large models.
- The optimization variable X is an m×n matrix, with m ≥ n assumed without loss of generality.
- Muon’s O step can be written as O_t = U V^⊤ when B_t = UΣV^⊤ is its singular value decomposition.
- The note also studies a closely related spectral optimizer derived from a quadratic approximation under the matrix spectral norm.
- Under the spectral norm, the update is Δ_t = −η_t∥B_t∥_* U V^⊤, using the SVD B_t = UΣV^⊤.
2 Convergence of Muon
The convergence analysis establishes results for Muon under Frobenius-norm smoothness and bounded stochastic-gradient variance, with different batch and momentum choices yielding distinct guarantees.
- Muon update: Muon’s update (2) is characterized as a heavy-ball method, unlike the Nesterov-type method in (1.2).
- Assumptions: The analysis assumes Lipschitz continuity of the gradient in Frobenius norm and an unbiased stochastic-gradient estimator with variance bounded by σ^2.
- Proof strategy: The proof develops a Frobenius-norm descent lemma and bounds the update through SVD-based quantities and gradient-estimation errors.
- Convergence of Muon: For update (2), Theorem 2.1 provides convergence guarantees under parameter choices involving β, α, R, T, L, and batch size B.
- Parameter choices: With B = 1, the analysis uses a parameter choice for 1 − β that depends on T, while arbitrary constant β requires B = T.
3 Convergence of the generic scheme in Bernstein and Newhouse (2024)
The generic scheme is analyzed under Lipschitz-gradient and unbiased bounded-variance assumptions, yielding a convergence result for suitable step size and momentum parameters. The analysis also identifies limitations for batch-free convergence and arbitrary norms.
- The proof analyzes the heavy-ball and mini-batch scheme by taking conditional expectations, summing inequalities, and selecting γ and η to control the resulting bound.The argument uses independence of the samples ξ_t,i and a suitable positive constant γ.
- The analysis assumes a Lipschitz gradient and an unbiased stochastic gradient estimator with variance bounded by σ^2.
- Theorem 3.1 establishes convergence for update (3.1) when η and β satisfy the stated step-size and momentum conditions.The result applies with β chosen as any constant in (0, 1).
- The resulting bound depends on the initial objective gap, η, L, β, and γ through the displayed inequality.
- Batch-free convergence is not achieved because the proof relies on convexity of quadratic and norm functions, while the sharper cited bound is known only for vector norms.
- Replacing the norm with an arbitrary norm leaves an unavoidable constant variance term even with mini-batch updates.