---
title: 'Muon-NSGD: Matrix-Aware Gradient Optimizer'
url: https://www.emergentmind.com/topics/muon-nsgd-optimizer
type: topic
---

# Muon-NSGD: Matrix-Aware Gradient Optimizer

Muon-NSGD (Natural Spectral Gradient Descent) is a matrix-aware, non-Euclidean gradient-based optimizer designed for efficient, stable training of neural networks—particularly transformer architectures—by leveraging geometric properties of matrix-structured parameters. The method replaces traditional Euclidean updates with steps adapted to the spectral structure of matrix weights, projecting gradient-based momenta onto the Stiefel manifold via the matrix-sign function, and employs efficient Newton–Schulz-based approximations to avoid explicit singular value decompositions. Muon-NSGD has demonstrated strong empirical performance in large-batch regimes and when combined with contemporary architectural components such as Mixture-of-Experts (MoE) and Multi-Head Latent Attention (MLA), and is supported by rigorous theoretical guarantees of stability, implicit spectral regularization, and convergence under modern nonconvex settings [2509.24406].

## 1. Fundamental Algorithmic Structure

Muon-NSGD constructs its update by applying a momentum-smoothed gradient transformation and projecting the resulting matrix onto the Stiefel manifold, thus ensuring the step direction is orthogonal and has spectral norm 1. Mathematically, for a matrix parameter $W_t$ at layer $\ell$, the iteration proceeds:

- **Momentum update:**
  \[
  M_t = \beta M_{t-1} + (1-\beta) G_t
  \]
  where $G_t = \nabla_{W_\ell} \mathcal L(W_t)$ is the stochastic gradient, $\beta$ is the momentum parameter.

- **Matrix-sign projection (orthogonalization):**
  \[
  U_t = \mathrm{msign}(M_t), \qquad \mathrm{msign}(M) = U_{:,1:r} V_{:,1:r}^\top = M(M^\top M)^{-1/2}
  \]
  where $M = U \Sigma V^\top$ is the SVD, and $r$ is the rank. This operation enforces that all singular values of $U_t$ are equal to 1, guaranteeing $\|U_t\|_2 = 1$.

- **Parameter update (with weight decay):**
  \[
  W_{t+1} = W_t - \eta_t (U_t + \lambda W_t)
  \]
  with learning rate $\eta_t$ and weight decay $\lambda$.

This process performs a natural gradient step on the Stiefel manifold and acts as a steepest descent step under the spectral norm [2509.24406, 2502.02900].

## 2. Efficient Computation: Newton–Schulz Orthogonalization

Direct computation of the matrix sign via full SVD is computationally prohibitive. Instead, Muon-NSGD employs a truncated Newton–Schulz iteration to approximate $M_t(M_t^\top M_t)^{-1/2}$. The iteration with coefficients $(a, b, c)$ (optimized to control the contraction of singular values) is

\[
X_0 = M_t / \|M_t\|_F
\]
\[
X_{k} = a X_{k-1} + b X_{k-1}(X_{k-1}^\top X_{k-1}) + c X_{k-1}(X_{k-1}^\top X_{k-1})^2
\]
for a small fixed $K$ (typically $K=5$) and coefficients $(3.4445, -4.7750, 2.0315)$. The result is re-scaled (e.g., by $s = \rho \sqrt{n}$ with $\rho=0.2$) to match the norm statistics of AdamW's step, mitigating variational instability during early training [2509.24406, 2601.19156].

This approach achieves near-optimal orthogonalization and removes the $\sqrt{\text{rank}}$ penalty characteristic of vector-based momentum methods, while avoiding the overhead of SVD, with doubly-exponential contraction of the approximation error as the number and degree of Newton–Schulz steps increase. In practice, $q=2,\, \kappa=2$ yields negligible deviation from exact polar iteration, while maintaining high efficiency [2601.19156].

## 3. Theoretical Guarantees and Spectral Regularization

Muon-NSGD imposes implicit spectral regularization through its update, stabilizing the optimization trajectory by flattening the gradient spectrum. The projection step solves

\[
U_t = \arg\min_{U : \sigma_i(U) = 1} \|U - M_t\|_F^2
\]
which ensures that no single direction in the parameter space is allowed to dominate, directly countering gradient explosion. The optimizer is additionally equivalent to computing the direction of steepest descent under the spectral norm constraint:

\[
U_t = \arg\max_{\|U\|_2 = 1} \operatorname{Tr}(G_t^\top U)
\]
which leads to robust control over the effective step size and direction.

For square matrices, Muon-NSGD's step coincides with a Riemannian (natural) gradient on the Stiefel manifold under the canonical metric, and $\mathrm{msign}(G)$ acts as a geodesic retraction of the ambient gradient projected onto the tangent space [2509.24406].

Muon-NSGD obtains $O(1/\sqrt{T})$ convergence in the stochastic nonconvex setting with adaptive learning rate schedules and remains provably stable for large-batch regimes. Notably, the Newton–Schulz approximation maintains the same convergence rate as the SVD-polar ideal, up to a constant factor vanishing doubly-exponentially fast in the number of steps and degree of polynomial approximation [2509.24406, 2601.19156, 2502.02900, 2507.01598].

## 4. Relationship to the Spectral-Compression Family and Non-Euclidean Descent

Muon-NSGD is the $p=0$ endpoint of a broader spectral-transformation family:
\[
\Psi_p(O_t) = U \Sigma^p V^\top
\]
with $p \in [0,1]$, interpolating between standard updates ($p=1$, identity), partial compression ($p=\frac12, \frac14$), and full flattening ($p=0$, Muon). The $p=0$ transform outputs the “polar factor” $UV^\top$, which enforces unit singular values, thus flattening the spectrum completely.

Empirical findings indicate that Muon-NSGD (p=0, momentum input) dramatically stabilizes standard momentum SGD, widening its learning-rate stability range. However, the advantage relative to adaptive optimizers such as Adam is context-dependent: under second-moment (RMS) normalization, as in Adam, full spectral flattening may degrade performance, while partial spectral compression ($p\approx\frac12$) often achieves optimal trade-offs [2602.04669].

Within the non-Euclidean gradient descent and LMO (Linear Minimization Oracle) frameworks, Muon-NSGD’s update corresponds to steepest-descent under the spectral norm, and can be interpreted as constrained or regularized descent on a product norm that aggregates spectral and (for non-matrix blocks) adaptive norms (e.g., $\ell_\infty$, AdaGrad- or Adam-type norms) [2510.09827, 2510.00643].

## 5. Empirical Performance and Practical Implementation

Muon-NSGD achieves high data efficiency and compute optimality, especially with modern architectural advances. Empirically, Muon reaches target loss with $48$–$52\%$ of the training computation of AdamW at matched or improved final perplexity [2509.24406]. When combined with MLA and MoE, configurations such as MLA+MoE+Muon achieve up to $68\%$ memory reduction and $3.2\times$ inference speedup alongside $8$–$12\%$ perplexity improvements. In Transformer kernels, Muon scales efficiently from $30$M to $200$M parameter decoders.

In the context of “grokking” experiments, Muon-NSGD led to a statistically significant acceleration of generalization onset compared to AdamW, reducing the mean “grokking” epoch from $153.09$ to $102.89$ across modular arithmetic and related tasks (t = 5.0175, p = 6.33e-08) [2504.16041].

The optimizer is sensitive to hyperparameter selection in its vanilla (constrained) forms—particularly learning rates—yet variants such as MuonMax (a regularized version mixing spectral and adaptive norms in an $\ell_2$ product) and Momo-regularized Muon substantially widen the stable regions and facilitate practical deployment [2510.09827].

Summary of recommended hyperparameters for canonical Muon-NSGD (see [2509.24406]):
- Learning-rate decay: $\eta_t = \eta_0 / \sqrt{t}$ or cosine schedule.
- Momentum: $\beta = 0.9$.
- Newton–Schulz steps: $K=5$, $(a,b,c) = (3.4445, -4.7750, 2.0315)$.
- RMS scaling: $\rho=0.2$; match RMS of AdamW to prevent early-stage instability.
- Weight decay: $\lambda \in [0.05,0.1]$.
- Gradient clipping: global norm $\leq 1.0$.
- Precision: mixed bfloat16 with float32 accumulators.

## 6. Variants, Distributed Implementations, and Extensions

Several Muon-NSGD variants address optimizer robustness, distributed computation, and computational bottlenecks:
- **MuonMax:** Replaces the strict spectral-ball update with a regularized step sensitive to aggregated spectral- and adaptive-norm magnitudes, exhibiting much greater learning-rate robustness while preserving strong generalization [2510.09827].
- **Momo-regularized Muon:** Incorporates model-based momentum and error-clipping using a “truncated model” framework for effective self-tuning of step sizes, leading to hyperparameter-insensitivity.
- **EF21-Muon:** Provides the first convergence-guaranteed, communication-efficient distributed framework for non-Euclidean LMO-based optimizers, leveraging bidirectional error-feedback with compression. EF21-Muon can recover Muon-NSGD as a special case in the uncompressed limit and retains optimal convergence rates even under non-Euclidean smoothness structures and $(L^0, L^1)$-layerwise regularity [2510.00643].

Momentum Variance Reduction (MVR) techniques further accelerate Muon-NSGD, reducing complexity from $O(K^{-1/4})$ to $O(K^{-1/3})$ in the non-convex regime. When integrated in a block-wise fashion (Gluon-MVR), these approaches achieve improvements in iteration complexity and time-to-accuracy, as confirmed by controlled LLM pretraining experiments [2512.16598].

## 7. Theoretical and Practical Trade-offs

Muon-NSGD occupies a unique Pareto frontier in the optimizer landscape: it achieves spectral regularization without requiring expensive per-iteration SVD, scales efficiently with large parameter matrices, and delivers substantial stabilization benefits for sequence model training in regimes where standard SGD and Adam-type optimizers struggle with gradient anisotropy or explosion.

However, full spectral flattening can “crude-ify” the update: it discards spectrum magnitude information, which may under-step in “stiff” directions or over-step in noisy directions. In extensive benchmarking, partial spectral compression (e.g., $p=1/2$) is sometimes favored over full flattening ($p=0$) in terms of final validation loss and stable learning-rate regime, especially when second-moment normalization is already present (e.g., Adam, AdamS). A plausible implication is that spectral normalization principally addresses stabilization, while performance optimization may require milder spectrum flattening or hybrid norm designs [2602.04669].

## References

- "Muon: Training and Trade-offs with Latent Attention and MoE" [2509.24406]
- "A Note on the Convergence of Muon" [2502.02900]
- "Delving into Muon and Beyond: Deep Analysis and Extensions" [2602.04669]
- "Convergence of Muon with Newton-Schulz" [2601.19156]
- "An Exploration of Non-Euclidean Gradient Descent: Muon and its Many Variants" [2510.09827]
- "Muon in Associative Memory Learning: Training Dynamics and Scaling Laws" [2602.05725]
- "Analysis of Muon's Convergence and Critical Batch Size" [2507.01598]
- "Muon is Provably Faster with Momentum Variance Reduction" [2512.16598]
- "Muon Optimizer Accelerates Grokking" [2504.16041]
- "Error Feedback for Muon and Friends" [2510.00643]

Source: https://www.emergentmind.com/topics/muon-nsgd-optimizer