---
title: Muon-based Gradient Optimizer (MuSGD)
url: https://www.emergentmind.com/topics/muon-based-gradient-optimizer-musgd
type: topic
---

# Muon-based Gradient Optimizer (MuSGD)

The Muon-based Gradient Optimizer (commonly termed MuSGD) is a matrix-aware stochastic optimization method characterized by spectral-norm orthogonalization of momentum or gradient updates. Originally developed for efficient and stable pre-training of large language models, MuSGD proceeds via momentum accumulation and orthogonalization using polar decomposition or Newton–Schulz iterations, followed by a geometrically scaled parameter update. The approach has led to substantial performance gains over AdamW and standard SGD in both theoretical convergence and empirical efficiency, especially on large-scale transformer architectures. The Muon+ variant introduces an additional per-row or per-column normalization after the orthogonalization step, yielding further robustness and improvements in validation perplexity across GPT and LLaMA model families [2602.21545].

## 1. Mathematical Formulation and Algorithmic Workflow

A single MuSGD step at iteration $t$ for a matrix parameter $W_t \in \mathbb{R}^{m \times n}$ proceeds as follows:

\[
\begin{aligned}
M_t &= \mu\, M_{t-1} + (1-\mu)\, G_t, \\
O_t &= \operatorname{Ortho}(M_t), \quad \operatorname{Ortho}(M) = M (M^\top M)^{-1/2} \approx U V^\top, \\
W_t &= W_{t-1} - \eta \sqrt{m/n}\, O_t,
\end{aligned}
\]
where $G_t$ is the stochastic gradient, $\mu$ the momentum coefficient, and $\eta$ the learning rate. The orthogonalization ensures that $O_t^T O_t \approx I$, enforcing a unit spectral norm and equivalently projecting the update onto the Stiefel manifold [2602.21545, 2509.24406].

In practical implementations, the polar factor is computed efficiently using Newton–Schulz iterations of the form:

\[
\begin{aligned}
X_0 &= \tfrac{1}{\alpha} M, \quad \alpha = \| M \|_2, \\
X_{k+1} &= \tfrac{3}{2} X_k - \tfrac{1}{2} X_k (X_k^\top X_k),
\end{aligned}
\]
with $K$ typically fixed at 5 for stability and throughput. This procedure approximates $M (M^\top M)^{-1/2}$ with controllable error in singular values [2602.21545].

In Muon+, an additional normalization pass is applied:

\[
O_t = \operatorname{Norm}_{(d)}(U_t), \quad \text{(with $d$ specifying "row", "col", etc.)}
\]
where, for instance, column-wise,
\[
\operatorname{Norm}_{(\text{col})}(X) = X D_{\text{col}}^{-1}, \quad D_{\text{col}} = \operatorname{diag}\left( \left(\sum_{i} x_{i1}^2+\varepsilon\right)^{1/2}, \dots \right)
\]
This stabilizes the scale of updates and controls the sensitivity to $\eta$ [2602.21545].

## 2. Spectral Regularization and Theoretical Principles

Muon’s key innovation is the use of spectral-norm constraints at each update. This can be formalized in the Lion-$\mathcal{K}$ mirror descent family, where the nuclear norm $K(X) = \| X \|_*$ defines the regularization and $\operatorname{msign}(X)$ (matrix sign function) acts as a subgradient preconditioner. The Muon update is then equivalent to solving:

\[
\min_X f(X) \ \text{subject to}\ \|X\|_2 \le \frac{1}{\lambda}
\]
for weight-decay parameter $\lambda$. The optimizer enforces $\|W\|_2 \le 1/\lambda$ throughout training, leading to spectral regularization and improved generalization. These updates are strictly dual to enforcing spectral constraints at each iteration via Fenchel conjugacy and KKT stationarity [2506.15054].

On non-square matrices—or with block-structured neural weights—this spectral flattening (enforcing all singular values to unity or near unity) yields highly controlled update directions. The explicit orthogonalization acts as a geometry-aware normalization, enhancing both convergence and step-size robustness [2505.23737].

## 3. Convergence Rates and Variance-Reduction Extensions

Standard stochastic MuSGD achieves a nonconvex convergence rate of $O(T^{-1/4})$ on the expected gradient norm. Under additional smoothness and PL conditions, the optimizer enjoys $O(T^{-1/2})$ or $O(T^{-1/3})$ (with variance reduction) rates, matching lower bounds for this problem class:

- **Option EMA (standard Muon)**: $O(T^{-1/4})$ on ergodic gradient norm [2509.15816].
- **Muon-VR2 (variance reduction)**: $O(T^{-1/3})$ with two-batch correction and properly coupled step/momentum schedules [2509.15816, 2512.16598].

Variance-reduced momentum (MVR) techniques integrated in the Gluon-MVR-2 framework yield the optimal nonconvex rate of $O(K^{-1/3})$. These refinements involve inner-outer double buffering and strong per-layer relative smoothness assumptions, amplifying stability at large batch sizes [2512.16598].

Table: Convergence Rates

| Method                | Nonconvex Rate      | Reference         |
|-----------------------|--------------------|-------------------|
| SGD, AdamW            | $O(T^{-1/2})$      | Standard Theory   |
| MuSGD Standard        | $O(T^{-1/4})$      | [2509.15816]      |
| MuSGD Variance-Reduced| $O(T^{-1/3})$      | [2512.16598]      |

## 4. Empirical Performance and Robustness

Extensive pretraining benchmarks confirm that Muon and Muon+ consistently outperform AdamW across architectures and training regimes:

- **GPT-style models (124M–774M)**: Muon+ improves validation perplexity by up to 2 points over Muon (e.g., 29.66 $\to$ 27.64, GPT-Small) [2602.21545].
- **LLaMA-style models (60M–1B)**: Similar robust improvements (e.g., LLaMA-1B, 10.68 $\to$ 10.31 PPL).
- **Compute-optimal (T2P $\approx$ 20) and over-training (T2P $\approx$ 200)**: Perplexity improvements persist with extended tokens.

| Model         | Muon PPL | Muon+ PPL | ΔPPL  |
|---------------|----------|-----------|-------|
| GPT-Small     | 29.66    | 27.64     | -2.02 |
| GPT-Base      | 21.70    | 19.98     | -1.72 |
| GPT-Large     | 17.82    | 16.91     | -0.91 |
| LLaMA-60M     | 25.75    | 25.25     | -0.50 |
| LLaMA-1B      | 10.68    | 10.31     | -0.37 |

Overhead from the Muon+ normalization step is negligible ($O(mn)$) relative to the Newton–Schulz orthogonalization kernel ($O(K m n^2)$). Muon+ requires no additional hyperparameter tuning relative to base Muon; the effective learning-rate window often widens due to increased scale-invariance and normalized update geometry [2602.21545]. Integration is straightforward: the normalization stage is a direct one-liner after the polar factor.

In the compute/epoch Pareto regime, Muon reaches target loss with half the training of AdamW while maintaining or improving perplexity. Additional robustness gains are reported for large-batch training, reduced grokking latency, and alleviation of spectral collapse in deep vision transformers (ViTs) [2509.24406, 2605.24770, 2504.16041].

## 5. Comparative Analysis and Theoretical Insights

Muon uniformly outperforms SGD and AdamW in settings where the Hessian or gradient covariance is low-rank, block-diagonal, or highly anisotropic—regimes typical in modern transformers and wide MLPs. The spectral-norm constraint enables learning rates up to the scale of the *average* singular value of the gradient (not the largest), a mechanism termed "spectral flattening" [2605.13079]. This greatly increases both the maximal stable step and the convergence rate under Kronecker-factored or K-FAC curvature models.

Muon’s update direction coincides with steepest descent under spectral norm and implements the natural gradient on the Stiefel manifold for square matrices [2509.24406]. In broader terms, the Muon step falls under non-Euclidean mirror descent, interpreted as a linear minimization oracle (LMO) constrained to a spectral-norm ball, with decoupled weight decay enforcing explicit operator-norm regularization throughout training [2506.15054, 2512.16598].

## 6. Practical Integration and Recommendations

- **Per-layer application**: Muon and Muon+ are applied to all non-embedding, non-norm matrix parameters; AdamW is recommended for scalar and embedding tokens as well as (optionally) LayerNorm layers.
- **Hyperparameters**: Default learning rates for Muon+ are robust; a typical range is $1 \times 10^{-3}$ to $5 \times 10^{-4}$, with momentum $\mu=0.9$ and 5 Newton–Schulz steps per iteration.
- **Scaling considerations**: The normalization factor $\sqrt{m/n}$ is included to match the update scale across rectangular matrices; weight decay and gradient clipping are unchanged.
- **Implementation**: Replace the orthogonal update step in preexisting Muon code with a normalized variant as specified. Adjust the normalization axis (“col”, “row”, or both) empirically; column-then-row is recommended for best robustness.
- **System overhead**: Newton–Schulz iterations dominate per-layer time; the added normalization has minimal overhead—empirically under 5% of optimizer compute [2602.21545].

## 7. Significance and Extensions

Muon+ establishes a new standard for matrix-aware optimization in deep learning, yielding pervasive improvements in convergence, loss, and stability. Its one-line normalization enhancement generalizes readily to new architectures, requires no retuning, and is robust across scales. Theoretical grounding within non-Euclidean mirror descent, spectral regularization, and blockwise adaptive trust-regions supports Muon+’s empirical success. Extensive pretraining and fine-tuning trials on both language models and vision transformers corroborate its efficiency and generality [2602.21545, 2509.24406, 2605.24770].

The Muon framework continues to expand, with recent advances such as curvature-aware extensions (e.g., Mousse), mixed Muon–SGD hybrids, and fully schedule-free variants. These lines of work aim to harness Muon’s spectral geometry while further refining its adaptivity, scaling behavior, and generalization guarantees.

Source: https://www.emergentmind.com/topics/muon-based-gradient-optimizer-musgd