The Spectral Dynamics and Noise Geometry of Muon

Technical Summary: The Spectral Dynamics and Noise Geometry of Muon

Problem Statement
The paper investigates the implicit regularization bias induced by the Muon optimizer, which updates weights using the polar factor of the gradient matrix rather than the gradient itself. Given a gradient $G = U\Sigma V^\top$, Muon updates via $UV^\top$, effectively discarding the relative magnitudes of singular values while preserving the active singular directions. While Muon has demonstrated empirical success in transformer pretraining, its theoretical selection mechanism remains unclear. Two competing hypotheses are addressed:

1. Is Muon merely a form of matrix-gradient normalization?
2. Does the spectral-norm geometry imply a bias toward low-rank solutions (nuclear-norm minimization)?

The authors argue both are incomplete. Instead, they propose that Muon induces a flat-update-spectrum bias, allocating equal amplitude across all active singular directions. The paper seeks to characterize this bias theoretically and empirically, specifically in the context of underdetermined regression where the optimizer must select a finite interpolant from an affine manifold of zero-loss solutions.

Methodology
The study employs a combination of continuous-time dynamical systems analysis, variational calculus, and controlled empirical experiments.

1. Theoretical Framework:

   • Projected Polar Flow: The authors analyze the continuous-time flow $\dot{W} = -\eta P_\perp P(W)$, where $P(W)$ is the polar factor of $W$ and $P_\perp$ is the orthogonal projector onto the tangent space of the interpolation manifold $\mathcal{M} = \{W : XW = Y\}$.
   • Shared-Frame (SF) Assumption: A structural condition where the gradient and weight matrices share singular frames ($U_G = U_W, V_G = V_W$) is introduced to simplify analysis. Theorems are proven without this assumption for the self-polar flow, while the SF assumption bridges the gap to the actual gradient-driven Muon flow.
   • Functionals: The analysis utilizes a pairwise spectral functional $R_{pw}(W) = -\sum_{i \neq j} \log(\sigma_i + \sigma_j)$ and its scale-invariant version $\tilde{R}_{pw}$ to track spectral shape.
   • Noise Geometry: A critical batch size derivation is performed using a one-step linearized signal-to-noise crossover, analyzing the Fréchet derivative of the polar map.

2. Empirical Validation:

   • Matrix Sensing: Controlled experiments on underdetermined linear regression tasks compare Muon against nuclear-norm minimizers (solved via CVXPY) and gradient descent.
   • Transformer Pretraining: Small-scale pretraining of NanoGPT (124M) and a matched ViT control on CIFAR-10 to observe spectral dynamics (stable rank, spectral entropy) in deep learning regimes.
   • Ablations: Tests on weight decay coupling (coupled vs. decoupled) and shared-frame misalignment to verify robustness.

Key Contributions

1. One-Step Spectral-Bias Theorem:
   Under explicit alignment assumptions, the authors prove that the polar update is the unique one-step entropy-maximizing choice among bounded updates that follow gradient singular directions but ignore the current weight spectrum. They establish a spectral flattening criterion: the normalized spectrum flattens if and only if the measurement geometry preferentially preserves the directions corresponding to the largest singular values (quantified by a sign criterion involving $\alpha_i q_i$).

2. Exact Singular-Value Dynamics:
   The paper derives exact dynamics for the projected polar flow, showing that singular values decay at rates $\dot{\sigma}_i = -\eta \alpha_i$, where $\alpha_i$ depends on the alignment between the weight's left singular vectors and the null space of the measurement matrix. This leads to a variational characterization where, at a fixed Frobenius norm, the unique minimizer of the spectral functional is a flat spectrum (all non-zero singular values equal).

3. Separation from Existing Explanations:

   • Normalization: Experiments show that norm-matched gradient descent does not reproduce Muon's behavior.
   • Nuclear Norm: Theoretical and empirical results demonstrate that Muon's limit points do not match nuclear-norm minimizers. While nuclear-norm minimization favors spectral concentration (low-rank), Muon favors spectral flattening (high stable rank).

4. Regime-Dependent Utility:
   The paper establishes that Muon's flat-spectrum bias is not universally superior. In small NanoGPT pretraining, Muon preserves stable rank and improves validation loss relative to AdamW. However, in a matched small-ViT control, the ranking reverses, suggesting the bias is beneficial only when many spectral directions need to remain active.

Results

• Theoretical: Theorem 1 proves that the projected self-polar flow drives the system toward a flat spectrum if the measurement geometry satisfies a specific sign criterion. Theorem 2 bounds the deviation when the Shared-Frame assumption is only approximately met. Theorem 4 provides a formula for the critical batch size $B_{crit}$ in polar-map SGD, linking it to the polar-map's noise sensitivity $S(\mu) = \sum_{i \neq j} (\sigma_i + \sigma_j)^{-2}$.
• Empirical (Matrix Sensing): Muon converges to solutions with significantly higher nuclear norms (1.29x–2.02x) and higher spectral entropy compared to the minimum nuclear-norm interpolant, falsifying the hypothesis that Muon minimizes nuclear norm.
• Empirical (Transformers): In NanoGPT, Muon maintains higher per-layer spectral entropy and stable rank compared to AdamW, which tends to collapse rank in language tasks. Conversely, in the ViT control, AdamW outperforms Muon, supporting the "regime-dependent" conclusion.
• Noise Sensitivity: The derived critical batch size formula shows that $B_{crit}$ is governed by the polar-map's sensitivity to noise, which is inflated by small or near-degenerate singular values. Flat spectra at fixed Frobenius radius minimize this sensitivity.

Significance and Claims
The paper claims to provide the first structural characterization of Muon's implicit bias as a flat-update-spectrum mechanism rather than a low-rank or simple normalization effect.

• Modest Scope: The authors explicitly state that the results are regime-dependent. Muon is not a universal optimizer; its utility arises in spectral regimes where maintaining activity across many directions is beneficial.
• Theoretical Limits: The theorems are local and structural, relying on assumptions like simple singular value gaps and specific measurement geometries. The paper does not claim to explain Muon's success in all large-scale LLM training scenarios, noting that the reduction of transformer layers to the studied underdetermined regression model remains a "conditional observation" requiring unverified assumptions about activation ranks.
• Practical Insight: The work offers practitioners a diagnostic tool: if an application requires preserving stable rank and broad spectral activity, Muon's flat-spectrum bias is advantageous; if the problem is inherently low-rank or well-conditioned, the bias may be unnecessary or detrimental.

The paper concludes by identifying open problems, including characterizing the specific measurement geometries that guarantee the flattening criterion globally and extending the noise sensitivity analysis to rectangular and rank-deficient cases.