🤖 machine learning

Convergence Bound and Critical Batch Size of Muon Optimizer

This paper provides theoretical convergence proofs for the Muon optimizer across various settings, demonstrates that weight decay ensures bounded parameter and gradient norms without requiring bounded-gradient assumptions, and derives a hyperparameter-dependent lower bound on the critical batch size that governs its training efficiency.

Original authors: Naoki Sato, Hiroki Naganuma, Hideaki Iiduka

Published 2026-06-09

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Naoki Sato, Hiroki Naganuma, Hideaki Iiduka

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a giant, complex robot (a neural network) to recognize cats in photos or write stories. To do this, you need a "coach" (an optimizer) that tells the robot how to adjust its internal settings to get better. For a long time, the standard coach has been AdamW, a very reliable trainer. But recently, a new coach named Muon has arrived, and it's showing incredible results in practice.

This paper is like a detective report that tries to explain why Muon is so good and gives us a rulebook for how to use it most efficiently. Here is the breakdown in simple terms:

1. The Core Idea: Seeing the Shape, Not Just the List

Most coaches treat the robot's brain like a giant, messy list of numbers. They look at the errors and say, "Move this number up, move that number down."

Muon is different. It looks at the robot's brain as a collection of shapes (matrices).

The Analogy: Imagine you are trying to straighten a crumpled piece of paper.
- Old Coaches (AdamW): They try to smooth it out by pushing individual points on the paper randomly. It works, but it's a bit messy.
- Muon: It looks at the whole sheet of paper. It realizes the paper has a specific "fold" or structure. Instead of just pushing, it performs a special maneuver (called orthogonalization) to straighten the paper perfectly along its natural lines. It finds the "cleanest" direction to move, ignoring how loud or noisy the error signal is.

2. The "Stability" Secret: The Brake Pedal

The paper discovered a crucial rule for using Muon with a feature called Weight Decay (which acts like a brake to stop the robot from getting too wild).

The Discovery: There is a strict relationship between how hard you push the robot (the Learning Rate) and how hard you press the brake (Weight Decay).
The Rule: If you push too hard without enough brake, the robot goes out of control. The paper proves mathematically that the "push" must never exceed the "brake" capacity. Specifically, the Learning Rate must be smaller than 1 / Weight Decay.
The Result: When you follow this rule, the robot's settings stay within safe, predictable limits. It doesn't explode or go crazy. This is a huge advantage because it means you don't have to assume the errors are small; the math guarantees the robot stays stable on its own.

3. The "Sweet Spot" Batch Size (The Critical Batch Size)

When training these robots, you can show them one photo at a time, or a whole pile of photos at once (a Batch).

The Problem: If you show too few photos, the robot learns slowly. If you show too many, the computer gets busy processing the pile, but the robot doesn't learn much faster. There is a "Goldilocks" point where you get the most learning for the least amount of computer work. This is called the Critical Batch Size.
Muon's Secret Sauce: The paper derived a formula for this "Goldilocks" point for Muon.
- The Momentum Factor: Muon uses a "momentum" setting (like a heavy flywheel that keeps the robot moving in the same direction). The paper found that if you turn up the momentum (make the flywheel heavier), the "Goldilocks" batch size gets smaller. You can use smaller batches and still be efficient.
- The Brake Factor: If you use a stronger brake (Weight Decay), the "Goldilocks" batch size gets larger. You need to show the robot more photos at once to get the best efficiency.

4. Why This Matters (The "Rank" Advantage)

The paper explains that Muon is efficient because it understands that the robot's brain isn't actually as complicated as it looks.

The Analogy: Imagine a 100x100 grid of lights (10,000 lights). Most of the time, only a few patterns of lights are actually turning on. The rest are just noise.
Muon's Edge: Muon realizes that the "real" information is low-rank (simple patterns). It ignores the noise. Because of this, it doesn't need a massive batch size to figure out the right direction. It can learn effectively with smaller batches compared to other methods, saving time and energy.

Summary of Findings

It Works: Muon is mathematically proven to converge (learn successfully) in four different settings (with/without momentum, with/without weight decay).
It's Stable: Using weight decay with the right learning rate guarantees the robot won't go crazy, even without assuming the errors are small.
It's Efficient: The paper gives you a formula to find the perfect batch size.
- High Momentum = You can use smaller batches.
- High Weight Decay = You should use larger batches.
Real-World Check: The authors tested this on image recognition (cats/dogs) and language models (writing text). The experiments confirmed that the math matches reality: Muon is faster and more efficient than the old standard, AdamW, especially when you tune the batch size according to their new rules.

In short, this paper takes a "black box" optimizer that was working well in practice, opens it up, and explains the mechanics, proving it's safe to use and telling you exactly how to tune it for maximum speed.

Technical Summary: Convergence Bound and Critical Batch Size of Muon Optimizer

Problem Statement

While the Muon optimizer has demonstrated strong empirical performance as a potential successor to standard optimizers like AdamW, its theoretical underpinnings remain underdeveloped. Existing literature lacks formal convergence proofs for Muon, particularly in practical settings involving Nesterov momentum and weight decay. Furthermore, the relationship between Muon's hyperparameters and its critical batch size (the batch size that minimizes the stochastic first-order oracle (SFO) complexity of training) is not theoretically characterized. This paper aims to bridge these gaps by providing rigorous convergence analysis and deriving a lower bound for Muon's critical batch size.

Methodology

The authors conduct a theoretical analysis of Muon across four practical configurations:

Without Nesterov momentum and without weight decay.
With Nesterov momentum and without weight decay.
Without Nesterov momentum and with weight decay.
With Nesterov momentum and with weight decay.

Theoretical Framework:

Convergence Analysis: The authors derive upper bounds on the average expected nuclear norm of the gradient, $\frac{1}{T}\sum_{t=0}^{T-1} \mathbb{E}[\|\nabla f(W_t)\|_*]$ . They utilize standard assumptions of $L$ -smoothness and bounded gradient variance. A key innovation is the analysis of Muon's search direction, which is obtained by orthogonalizing a momentum matrix $C_t$ . The authors account for the rank of the momentum error matrix ( $r_1$ ) and the search direction ( $r_2$ ), rather than the full parameter dimension.
Newton-Schulz Approximation: Since Muon practically approximates the orthogonalization via the Newton-Schulz iteration rather than exact SVD, the authors analyze the impact of this approximation error ( $\delta(k)$ ) on convergence, showing it introduces only a negligible term.
Critical Batch Size (CBS) Derivation: The authors model the number of training steps $T$ required to reach a target precision $\epsilon$ as a function of batch size $b$ . By minimizing the SFO complexity ( $T \times b$ ), they derive a lower bound for the critical batch size $b^*$ .

Experimental Validation:
The theoretical predictions are validated through a two-level strategy:

Controlled Full-Muon Experiments: A synthetic Teacher-Student task where Muon updates all parameters, strictly matching theoretical assumptions.
Practical-Scale Hybrid Experiments: Training ResNet-18/VGG-16 on CIFAR datasets and Llama3.1 (320M) on the C4 corpus. In these settings, Muon is applied to matrix-shaped parameters while AdamW handles the rest, reflecting standard deployment.

Key Contributions

1. Convergence Proofs for Four Variants

The paper establishes convergence guarantees for all four Muon variants. The derived upper bounds for the average expected gradient norm are of the order:
$O\left(\frac{1}{\eta T} + \sqrt{\frac{(1-\beta)r_1}{b}} + \bar{r}\eta\right)$
(with slight variations for Nesterov and weight decay terms).

Significance: Unlike vector-based optimizers where variance scales with the full parameter count, Muon's bound scales with the rank $r_1$ of the momentum error matrix. Empirical evidence suggests $r_1 \ll n$ in over-parameterized regimes, implying a tighter theoretical bound.
Nuclear Norm: The analysis uses the nuclear norm, which is a stricter criterion than the Frobenius norm used in many vector-based analyses, suggesting Muon respects the intrinsic matrix structure of neural network parameters.

2. Role of Weight Decay and Stability

The authors prove that incorporating weight decay ensures the almost-sure boundedness of both parameter norms ( $\|W_t\|_F$ ) and gradient norms ( $\|\nabla f(W_t)\|_F$ ) without relying on the restrictive assumption of bounded gradients.

Stability Condition: They derive a necessary condition for stability: the learning rate $\eta$ and weight decay coefficient $\lambda$ must satisfy $\eta \leq 1/\lambda$ . Experimental results confirm that training becomes unstable when $\eta$ exceeds this threshold.
Implicit Regularization: While weight decay does not improve the asymptotic convergence rate order, it provides implicit regularization by bounding norms, which explains its superior empirical performance.

3. Critical Batch Size (CBS) Lower Bound

The paper derives a lower bound for the critical batch size $b^*$ , showing its dependence on hyperparameters:
$b^*_{\text{Muon}} > \frac{9(1-\beta)(1+\sqrt{2\beta})^2 r_1 \sigma^2}{(1-\lambda)^2 \epsilon^2}$

Hyperparameter Dependence: The formula reveals that:
- Momentum ( $\beta$ ): As $\beta \to 1$ , the term $(1-\beta)$ vanishes, suggesting the critical batch size decreases.
- Weight Decay ( $\lambda$ ): The factor $1/(1-\lambda)^2$ indicates that increasing weight decay increases the critical batch size.
- Rank ( $r_1$ ): The bound scales with the rank of the momentum error, linking CBS to the low-rank structure of gradients.
Qualitative Scaling: The authors emphasize that the formula provides a qualitative scaling law rather than an absolute prediction, as it depends on problem-specific constants ( $\sigma^2, r_1, \epsilon$ ). However, the ratio of CBS between different Muon variants depends only on $\beta$ and $\lambda$ .

Results

Theoretical Findings:

The convergence rate matches standard non-convex SGD-type methods ( $O(T^{-1/4})$ under specific hyperparameter schedules) but with a tighter variance term due to rank dependence.
The Newton-Schulz approximation error has a limited effect on convergence, decaying exponentially with the number of iterations.

Experimental Findings:

Stability: Experiments on ResNet-18/CIFAR-10 validate the stability condition $\eta \leq 1/\lambda$ . Gradient norms and training loss are minimized near this boundary.
Convergence Speed: Muon with Nesterov momentum and weight decay consistently achieves the fastest decrease in loss and gradient norm compared to AdamW, Momentum SGD, and other Muon variants.
Critical Batch Size Validation:
- Hyperparameter Effects: Experiments confirm that adding Nesterov momentum shifts the CBS to larger values, and increasing weight decay further increases the CBS. For instance, increasing $\lambda$ from 0 to 0.5 resulted in a CBS shift from $\approx 1024$ to $\geq 4096$ , matching the predicted $1/(1-\lambda)^2$ scaling (ratio of 4.0).
- Rank Scaling: Varying the model width showed that the effective rank $r_1$ scales with width, and the product $r_1 \cdot \sigma^2$ remains relatively constant, explaining why empirical CBS values appeared flat across widths despite theoretical rank dependence.
- Efficiency: Muon generally requires fewer steps and lower SFO complexity than AdamW, particularly at small to moderate batch sizes.

Significance and Claims

The paper claims to provide the first rigorous convergence analysis of Muon that simultaneously incorporates Nesterov momentum and weight decay, addressing a gap in the theoretical understanding of this emerging optimizer.

Theoretical Insight: It demonstrates that weight decay is not merely a regularization tool but a mechanism that ensures boundedness and stability, removing the need for bounded-gradient assumptions.
Practical Guidance: The derived CBS lower bound offers actionable guidance for practitioners. It clarifies how tuning $\beta$ and $\lambda$ qualitatively affects the optimal batch size, suggesting that high-momentum configurations can be trained efficiently with smaller batches, while higher weight decay requires larger batches to maintain efficiency.
Validation: The work bridges the gap between theory and practice by validating these predictions across vision and language modeling workloads, confirming that the theoretical scaling laws hold even in hybrid optimization settings where Muon is applied only to matrix-shaped parameters.

The authors modestly note that their single-matrix analysis does not explicitly model layer-wise heterogeneity, but they argue that the high correlation of gradient noise scales across layers supports the generalizability of their findings. They also acknowledge that while Muon outperforms Shampoo in their experiments, deriving analogous bounds for Shampoo remains an open problem.