Technical Summary: Spectral Flattening Is All Muon Needs

Problem Statement

Training modern neural networks involves optimizing over highly non-convex landscapes characterized by flat regions, sharp directions, and saddle points. The choice of optimizer is critical for stability and efficiency. While coordinate-wise adaptive methods like Adam and AdamW have improved robustness, they do not fully address the geometric challenges posed by matrix-valued parameters. Recently, Muon has emerged as a high-performing alternative that orthogonalizes the momentum buffer before updates, allowing for larger learning rates and faster convergence. However, the theoretical mechanisms explaining why Muon outperforms standard Gradient Descent (SGD) in terms of maximal stable step sizes and convergence rates remain under-explored. Specifically, existing theory has not quantitatively explained how orthogonalization alters the learning rate bottleneck or the effective condition number of the optimization problem.

Methodology

The paper provides a theoretical analysis of Muon by framing its update rule as a form of spectral flattening. The authors develop two primary theoretical results and validate them experimentally:

1. Theoretical Framework:

   • Muon Definition: Muon updates weights $W$ by computing the gradient $G_t$, then applying Newton–Schulz iterations to approximate the polar factor $O_t \approx (G_t G_t^\top)^{-1/2} G_t = UV^\top$. This operation preserves the singular vectors of the gradient while replacing singular values with ones.
   • Hessian Approximation: The analysis utilizes a Kronecker-factored curvature model (K-FAC), approximating the Hessian as $H \approx X^\top X \otimes G_t G_t^\top$, where $X$ represents input activations and $G_t$ the gradient.
   • Learning Rate Analysis: The authors derive exact one-step descent thresholds for both SGD and Muon. They compare the maximal stable learning rate ($\eta_{max}$) for both optimizers, showing that SGD is bottlenecked by the largest singular value of the gradient ($\sigma_{max}$), whereas Muon's stability is governed by the average singular value ($\frac{1}{m}\sum \sigma_i$).
   • Convergence Analysis: Muon is recast as a preconditioned gradient method with preconditioner $P = I_n \otimes (G_t G_t^\top)^{-1/2}$. Under Polyak–Łojasiewicz (PL) conditions and relative smoothness, the authors analyze the convergence factor. They prove that the effective condition ratio for Muon ($\tilde{\alpha}/\tilde{\beta}$) is strictly better than that of SGD ($\alpha/\beta$) when the gradient covariance is anisotropic, specifically improving by a factor related to the spectral spread of the gradient.

2. Experimental Design:

   • Setup: Experiments were conducted on CIFAR-10 using CifarNet. To isolate the optimizer's effect on matrix-valued parameters, a dual-optimizer strategy was employed: convolutional kernels were updated by the target optimizer (Muon or SGD), while biases and other parameters used a fixed reference optimizer.
   • Stability Tests: Learning rate sweeps were performed without Batch Normalization (BN) to observe intrinsic stability.
   • Convergence Tests: Both optimizers were run at identical learning rates (and separately at their respective best rates) to measure convergence speed and per-step convergence ratios ($r_t$).
   • Normalization Principle: A new normalization scheme, FrobNorm (normalizing by the Frobenius norm), was proposed and tested to verify the theoretical claim that reducing the spectral scale of inputs ($\lambda_{max}(X^\top X)$) enables higher learning rates.

Key Contributions

1. Maximal Learning Rate Derivation: The paper proves that Muon's maximal stable step size scales with the average singular value of the gradient rather than the largest.

   • For SGD: $\eta_{max} \propto \frac{2}{\lambda_{max}(H)} \approx \frac{2}{\lambda_{max}(X^\top X) \sigma_{max}(G)}$.
   • For Muon: $\eta_{max} \propto \frac{2}{\lambda_{max}(H)} \frac{\sum \sigma_i}{m} \approx \frac{2}{\lambda_{max}(X^\top X) \sigma_{max}(G)} \frac{\sum \sigma_i}{m \sigma_{max}(G)}$.
   • This demonstrates that spectral flattening prevents a single dominant singular direction from dictating the step-size constraint, allowing Muon to tolerate significantly larger learning rates.

2. Convergence Rate Acceleration: The authors show that Muon acts as a preconditioner that improves the effective condition ratio.

   • Under K-FAC, the convergence factor for Muon is $1 - \frac{\tilde{\alpha}}{\tilde{\beta}}$, where $\frac{\tilde{\alpha}}{\tilde{\beta}} = \frac{\alpha}{\beta} \sqrt{\frac{\lambda_{min}(G_t G_t^\top)}{\lambda_{max}(G_t G_t^\top)}}$.
   • Since $\frac{\lambda_{min}}{\lambda_{max}} < 1$, the term $\sqrt{\dots}$ is less than 1, but the derivation shows $\frac{\tilde{\alpha}}{\tilde{\beta}} > \frac{\alpha}{\beta}$ in the context of the preconditioned system, leading to a strictly faster linear convergence rate when the gradient covariance is ill-conditioned.

3. New Normalization Principle: The analysis reveals that the maximal stable learning rate is inversely proportional to $\lambda_{max}(X^\top X)$. The authors propose that normalization layers should aim to reduce this spectral scale. They validate this with FrobNorm, showing it enables training at high learning rates (e.g., 0.25) where standard training fails, particularly for SGD.

Experimental Results

• Stability: In the absence of Batch Normalization, SGD diverged within the first few iterations at learning rates where Muon remained stable and continued to decrease loss. Muon demonstrated controlled growth in parameter and gradient norms, whereas SGD exhibited catastrophic explosion.
• Convergence Speed: At identical learning rates (e.g., $\eta=0.05$), Muon reached accuracy milestones (70%–91%) several epochs earlier than SGD. The empirical convergence ratio $r_t$ was consistently lower for Muon, confirming the theoretical prediction of a faster linear rate.
• Normalization Impact: The application of FrobNorm allowed both optimizers to train stably at high learning rates (0.25 for CifarNet, 0.2 for a GPT-2-style Transformer), validating the hypothesis that controlling the spectral scale of layer inputs is a general principle for enabling larger step sizes.

Significance and Claims

The paper claims to offer a principled, geometric explanation for Muon's empirical success. It posits that Muon's advantage is not merely a heuristic but a direct consequence of spectral flattening, which:

1. Decouples the learning rate constraint from the largest singular value, allowing for larger step sizes.
2. Acts as a one-sided preconditioner that improves the effective condition number of the optimization landscape, leading to faster convergence.

The authors acknowledge limitations, noting that their theoretical analysis assumes deterministic full-batch settings and exact polar factors, whereas practical implementations use stochastic gradients, finite Newton–Schulz iterations, and momentum. They state that extending these results to stochastic settings and larger-scale architectures remains an important direction for future work. The paper concludes that these results help transform Muon's empirical performance into a more precise theoretical understanding, potentially guiding the design of future optimizers and normalization strategies.