Muon-NSGD: Natural Spectral Gradient Descent

• Muon-NSGD is a matrix-structured optimizer that applies spectral normalization to uniformly balance singular value directions for efficient convergence.
• It leverages SVD-based polar updates and Newton–Schulz iterations to achieve implicit spectral regularization and iteration complexity independent of the condition number.
• Practical implementations in LoRA and large-scale generative modeling demonstrate enhanced generalization, robust learning dynamics, and improved optimization stability.

Muon-NSGD (Natural Spectral Gradient Descent) is a matrix-structured optimization paradigm that performs steepest descent with respect to the spectral (operator) norm. By orthogonalizing the weight update direction via the matrix sign or polar factor, Muon-NSGD uniformly balances progress across singular-value directions and provides implicit spectral regularization. This optimizer underpins state-of-the-art methods in large-scale generative modeling, scientific machine learning, matrix factorization, LLM adaptation, and ill-posed inverse problems. Its core mechanism, theoretical convergence, and distinctive spectral behavior have been rigorously analyzed in recent literature, particularly in the context of low-rank adaptation (LoRA) and structured architectures.

1. Algorithmic Structure and Mathematical Foundations

Muon-NSGD updates matrices by stepwise projection onto the set of orthonormal (or unit spectral-norm) directions. At each iteration, for a weight matrix $W \in \mathbb{R}^{m \times n}$ with gradient $G$:

• SVD-based polar update:

$$G = U \Sigma V^\top, \quad \text{then} \quad P(G) = U V^\top,$$

so the update is

$W_{t+1} = W_t - \eta P(G).$

• In the momentum and practical variants,

$$M_t = \beta M_{t-1} + (1-\beta) G_t, \quad U_t = M_t (M_t^\top M_t)^{-1/2}.$$

The main weight step is

$W_{t+1} = W_t - \eta (U_t + \lambda W_t).$

Momentum and weight decay are handled analogously to AdamW (Mehta et al., 29 Sep 2025).

• The Newton–Schulz iteration,

$$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,$$

with coefficients (3.4445, -4.7750, 2.0315), efficiently approximates the inverse square root for spectral normalization (Mehta et al., 29 Sep 2025, Kim et al., 27 Jan 2026).

This approach coincides with the Riemannian (natural) gradient for the Stiefel manifold when $W$ is square, and more generally, it is the mirror descent under the nuclear norm, with spectral-norm constraint induced by decoupled weight decay (Chen et al., 18 Jun 2025).

2. Spectral Dynamics in Low-Rank Matrix Factorization and LoRA

In the LoRA-style parameterization, the target is optimizing a product $W = AB$ for $A \in \mathbb{R}^{m\times r}$, $B \in \mathbb{R}^{r\times n}$, where $r \ll \min\{m, n\}$. The Muon/SpecGD update on low-rank factors applies spectral orthogonalization to each gradient: $$\widetilde{G}_{A,t} = (G_{A,t} G_{A,t}^\top)^{-1/2} G_{A,t}, \quad \widetilde{G}_{B,t} = (G_{B,t} G_{B,t}^\top)^{-1/2} G_{B,t}$$ followed by

$$A_{t+1} = A_t - \eta_t \widetilde{G}_{A,t}, \quad B_{t+1} = B_t - \eta_t \widetilde{G}_{B,t}.$$

A central finding is the "equal-rate" spectral growth phenomenon: all singular values of $AB$ grow in near-perfect synchrony, in stark contrast to the largest-first, spectrum-fanning dynamics of standard gradient descent or AdamW. This is formalized through continuous-time analysis (SpecGF), where squared root singular values evolve according to

$$\frac{d}{dt} \sqrt{d_i(t)} = 1 + O(\varepsilon^{1/4})$$

for all active modes, and convergence to global minima is guaranteed under mild initialization conditions and with ℓ₂ regularization (Kang et al., 6 Feb 2026).

3. Theoretical Properties and Convergence Guarantees

Muon-NSGD satisfies the following key theoretical results:

• Convergence Rates: Mirror descent arguments and spectral-norm smoothness yield $O(1/\sqrt{T})$ nonconvex convergence rates for standard settings. In convex or PL conditions, linear or $O(1/T)$ rates are achieved. The use of Newton–Schulz approximation introduces a constant factor, which decays doubly exponentially in the number of iterations and polynomial degree (Kim et al., 27 Jan 2026).
• Independence from Condition Number: In matrix factorization and related spectral settings, Muon-NSGD achieves iteration complexity independent of the condition number $\kappa$, unlike GD or coordinatewise optimizers that scale with $\kappa$ or $\sqrt{\kappa}$ (Ma et al., 20 Jan 2026).
• Spectral Regularization: The unit spectral norm of each update prevents gradient explosion and enforces low-rank isotropy, acting as an implicit preconditioner. The iterates remain confined to a spectral-norm ball determined by the weight decay parameter (Chen et al., 18 Jun 2025, Li et al., 5 Feb 2025).
• Strict-saddle Avoidance: Owing to the real-analytic nature of the dynamics and the topological properties of matrix factorization landscapes, almost all initializations avoid strict saddles and converge to global minima if iterates remain bounded (Kang et al., 6 Feb 2026).

4. Practical Implementation and Algorithmic Variants

Implementation leverages low-rank SVD or Newton–Schulz approximations to enforce feasible computational cost per iteration:

• For each matrix parameter, accumulate Polyak or Nesterov momentum.
• Spectral normalization is achieved via SVD (for small-to-medium matrices) or 3–5 Newton–Schulz steps (large-scale settings).
• Learning rates, momentum, and weight decay follow standard schedules, but RMS-matching of update magnitude is recommended to match AdamW scales.
• Spectral-norm projection can be imposed post-update to enforce hard constraints (Mehta et al., 29 Sep 2025, Chen et al., 18 Jun 2025).

Pseudocode in canonical form:

M_t = beta * M_{t-1} + (1 - beta) * G_t X_0 = M_t / ||M_t||_F for k in range(K): X_k = a * X_{k-1} + b * X_{k-1} @ (X_{k-1}.T @ X_{k-1}) + c * X_{k-1} @ (X_{k-1}.T @ X_{k-1}) @ (X_{k-1}.T @ X_{k-1}) U_t = X_K W_{t+1} = W_t - eta * (U_t + lambda * W_t)

with (a, b, c) tuned for numerical stability and efficiency (Mehta et al., 29 Sep 2025, Kim et al., 27 Jan 2026).

For LoRA or matrix factorization, update orthonormalized gradients for both low-rank factors separately, preserving the same spectral dynamics (Kang et al., 6 Feb 2026).

5. Spectral Behavior, Applications, and Empirical Results

Muon-NSGD exhibits unique spectral signatures:

• Uniform Spectral Growth: In LoRA, all singular values of the parameter product grow uniformly. Muon thus preserves effective rank and mode diversity, contrasting with spectrum collapse under AdamW (Kang et al., 6 Feb 2026).
• Isotropic Learning: In associative memory, phase retrieval, and imbalanced settings, Muon-NSGD equalizes update amplitudes across modes or class frequencies, yielding exponential speedup over GD in tail-dominated or highly non-uniform spectra (Li et al., 5 Feb 2026, Braun et al., 30 Jan 2026).
• Improved Scaling Laws: In LLM pretraining, Muon achieves steeper learning curves ($\gamma \approx 0.80$ in budget-scaling exponents) versus SGD ($\gamma \approx 0.45$) (Li et al., 5 Feb 2026).
• Enhanced Generalization and Stability: Enforced spectral constraints and curvature-aware steps accelerate transitions from memorization to generalization, as observed in grokking benchmarks where Muon reduces the mean epoch of generalization by ~50 epochs versus AdamW (Tveit et al., 22 Apr 2025).

Empirical comparisons consistently demonstrate that Muon-NSGD outpaces AdamW and standard SGD in effective rank preservation, learning efficiency, and robustness to pathological spectra or initialization (Kang et al., 6 Feb 2026, Li et al., 5 Feb 2026, Braun et al., 30 Jan 2026).

6. Connections to Natural Gradient, Geometric Optimization, and Extensions

Muon-NSGD implements the steepest descent update under the spectral norm, closely related to mirror descent with nuclear norm (trace class) geometry (Chen et al., 18 Jun 2025). The update

$W_{t+1} = W_t - \eta (W_t W_t^\top)^{-1/2} G_t$

corresponds to a natural gradient step on the Stiefel manifold, and the matrix sign operator is the (sub)gradient map of the nuclear norm. This spectral "whitening" achieves geometric conditioning similar to block-wise Fisher preconditioning, but with exact singular vector decomposition.

• In physics-informed and scientific ML tasks, Muon-NSGD has been extended with mode-wise step adaptation via RSAV (relaxed scalar auxiliary variables), leading to provable energy dissipation, positivity, and linear convergence under PL conditions. These variants surpass AdamW and vanilla Muon in PINNs, DeepONets, and stiff PDE training (Lu et al., 18 Feb 2026).
• The optimizer design admits generalization: block-wise or low-rank spectral normalization, curvature integration, and hybrid variants, all retaining the regularization and preconditioning benefits intrinsic to Muon-NSGD's spectral geometry (Ma et al., 20 Jan 2026, Mehta et al., 29 Sep 2025, Kang et al., 6 Feb 2026).

7. Comparative and Implementation Considerations

• Computational Overhead: Full SVD scales as $O(m n \min\{m, n\})$, but Newton–Schulz and randomized SVD approaches yield speedups of 4–10× per step for large matrices (Kim et al., 27 Jan 2026, Mehta et al., 29 Sep 2025).
• Memory: For Transformer layers (1–4K dims), Muon is competitive with Adam in space, requiring storage for momenta and possibly SVD factors.
• Hyperparameter Tuning: Learning rates and momentum largely parallel AdamW defaults, though spectral scaling and decay can be matched via RMS normalization. Parameter settings $K=3$–5 (NS steps), momentum $\beta=0.9$, and polynomial degree $\kappa=2$ are effective (Mehta et al., 29 Sep 2025, Kim et al., 27 Jan 2026).
• Suitability: Muon-NSGD is particularly effective for dense layers in LLMs, spectral learning tasks, and situations with ill-conditioned Hessians or long-tailed data distributions.

Key primary sources for Muon-NSGD theory, dynamics, and applications include (Kang et al., 6 Feb 2026, Ma et al., 20 Jan 2026, Mehta et al., 29 Sep 2025, Kim et al., 27 Jan 2026, Chen et al., 18 Jun 2025, Li et al., 5 Feb 2026), and (Lu et al., 18 Feb 2026). These works provide comprehensive mathematical derivations, algorithmic designs, and large-scale empirical validations across machine learning and scientific domains.