Delving into Muon and Beyond: Deep Analysis and Extensions

Technical Summary: Delving into Muon and Beyond: Deep Analysis and Extensions

Problem Statement
The Muon optimizer has recently gained attention for its empirical performance in training large language models, utilizing matrix-level orthogonalized updates. However, the underlying mechanisms of Muon and its relationship to established adaptive optimizers like Adam remain insufficiently understood. Existing literature often conflates Muon's effects with auxiliary techniques (e.g., QK-Norm, QK-Clip) or lacks a systematic comparison against standard baselines. Furthermore, it is unclear whether Muon's spectral orthogonalization represents a universally superior optimization strategy or a specific form of stabilization, and whether it can be improved or generalized.

Methodology
The authors propose a unified spectral framework to analyze Muon and its variants. They define Muon as the $p=0$ endpoint of a family of spectral transformations of the form $U\Sigma^p V^\top$, where $U\Sigma V^\top$ is the Singular Value Decomposition (SVD) of the update matrix. The study investigates intermediate exponents $p \in \{1, 1/2, 1/4, 0\}$:

• $p=1$: Identity (standard updates).
• $p=1/2, 1/4$: Partial spectral compression.
• $p=0$: Full orthogonalization (Muon).

These transformations are applied to two distinct input types:

1. First-moment momentum ($O_{mom}$): Equivalent to momentum SGD inputs.
2. RMS-normalized updates ($O_{rms}$): Equivalent to Adam inputs, where gradients are normalized by the second moment.

To enable efficient computation of fractional powers ($p=1/2, 1/4$) without explicit SVD, the authors develop a coupled Newton-Schulz iteration. This algorithm simultaneously computes the matrix square root and its inverse using only matrix multiplications, ensuring numerical stability and scalability.

The experimental evaluation is conducted on a 124M parameter GPT-2 model trained on OpenWebText. The study employs a highly controlled setup:

• Decoupled Learning Rates: Separate learning rates are used for matrix and vector parameters to isolate the optimizer's effect.
• No Auxiliary Techniques: Weight decay, QK-Norm, and QK-Clip are disabled to ensure fair comparisons.
• Systematic Tuning: A two-stage learning rate search (coarse logarithmic sweep followed by fine-grained refinement) is performed for all eight optimizer variants (4 spectral exponents $\times$ 2 input types).

Key Contributions

1. Unified Spectral Framework: The paper introduces a formal framework $U\Sigma^p V^\top$ that positions Muon as a specific case ($p=0$) and proposes intermediate variants ($p=1/2, 1/4$) to interpolate between standard updates and full orthogonalization.
2. Efficient Algorithm: A coupled Newton-Schulz iteration is developed to compute fractional spectral updates ($p=1/2, 1/4$) efficiently, avoiding the computational cost of explicit SVD.
3. Rigorous Empirical Comparison: The study provides a controlled comparison between Muon, Adam, and their spectral variants, isolating spectral effects from confounding factors like weight decay or auxiliary normalization techniques.

Key Results
The experiments yield three primary findings:

1. Stabilization of First-Moment Updates: On first-moment inputs (mSGD), Muon ($p=0$) provides significant stabilization, allowing for much higher learning rates and greater robustness compared to standard mSGD. However, partial compression ($p=1/4$) often achieves better final performance once stability is established.
2. Limited Gains on RMS-Normalized Inputs: When applied to RMS-normalized inputs (Adam-style), spectral compression yields diminishing returns. Full orthogonalization (AdamZ/Muon applied to $O_{rms}$) does not outperform standard Adam and can be inferior to partial compression. This suggests that element-wise RMS normalization already mitigates much of the anisotropy that spectral compression targets.
3. Superiority of Second-Moment Methods: Across all spectral exponents, optimizers utilizing second-moment normalization (RMS-based) consistently outperform their first-moment (momentum-only) counterparts in terms of both stability and final loss.

Significance and Claims
The paper concludes that Muon is best understood as an effective form of spectral normalization rather than a universally superior optimization method.

• Stabilization vs. Performance: Muon's primary contribution is stabilizing first-moment updates by discarding singular value magnitudes, which prevents directions with large singular values from dominating. However, this "flattening" is a crude mechanism; it removes magnitude-based filtering, potentially amplifying noisy gradient directions (small singular values) and under-stepping on informative dominant modes.
• Contextual Utility: While spectral compression offers genuine stabilization benefits, especially when the input update is not already normalized, full orthogonalization ($p=0$) is not a replacement for modern RMS-normalized second-moment optimizers like Adam. The authors argue that the reported advantages of Muon in prior work may be attributed to the specific combination of techniques used (e.g., QK-Norm) rather than the orthogonalization itself.

The work emphasizes that while spectral transformations are powerful tools for controlling anisotropy, they do not inherently surpass the benefits of adaptive second-moment scaling in large-scale training regimes.