How Much Orthogonalization Does Muon Need?

Technical Summary: How Much Orthogonalization Does Muon Need?

Problem Statement
The Muon optimizer improves neural network training by replacing ill-conditioned momentum updates with approximately semi-orthogonal updates. While standard numerical linear algebra aims for high-accuracy polar decomposition (mapping all singular values to 1), the Muon context suggests that exact accuracy may be unnecessary. The central question addressed by this work is: How much orthogonalization accuracy does Muon actually require to preserve training quality? Specifically, the paper investigates whether the computational cost of high-degree polynomial iterations (used in existing methods like Muon-Jordan and Polar Express) is justified, or if a cheaper, lower-degree approximation can achieve comparable training results.

Methodology
The authors propose a relaxed cubic Newton–Schulz schedule (termed cubic5) designed specifically for Muon's low-precision (bfloat16) singular value band.

1. Relaxed Targets: Unlike classical polar decomposition which targets machine precision, this method targets a relaxed band of $[0.7, 1.3]$ for singular values. It acknowledges that in bfloat16, singular values below an effective lower bound ($l_0 = 7 \times 10^{-3}$) are not practically meaningful.
2. Adaptive Coefficient Derivation: The authors derive an adaptive cubic polynomial $f_t(x) = a_t x + b_t x^3$ at each iteration. The coefficients are calculated based on the current worst-case lower singular value bound ($l_t$) and a fixed upper peak ($u=1.3$). This ensures the polynomial maps the protected interval $[l_t, r_t]$ into the target band while balancing the endpoints to maximize the new lower bound.
3. Computational Efficiency:
   • Cubic5: A single cubic step requires two dominant matrix multiplications (forming the Gram matrix and multiplying it back). Five steps require 10 dominant matrix multiplications.
   • Baseline (Muon-Jordan/Polar Express): A quintic step requires three dominant matrix multiplications. Five steps require 15 dominant matrix multiplications.
   • The method utilizes the "transpose trick" for tall matrices to minimize operations on the smaller Gram matrix.

Key Contributions

• Derivation of a Relaxed Cubic Schedule: The paper derives a five-step cubic Newton–Schulz schedule with coefficients adapted to the bfloat16 effective lower bound, achieving the target spectral band in five iterations.
• Cost-Quality Frontier Characterization: The work establishes that training quality is not monotonically governed by polar-decomposition accuracy. It demonstrates that a lower-cost cubic approximation can reach performance parity with higher-cost quintic methods and even exact SVD in specific regimes.
• Empirical Validation Across Scales: The authors validate the method through:
  • Synthetic Diagnostics: Confirming that the scalar polynomial transfer holds in low-precision matrix arithmetic.
  • NanoGPT Ablations: Comparing cubic5, truncated Polar Express, truncated Muon-Jordan, and explicit FP32 SVD on GPT-2 Small.
  • Large-Scale Training: Training hybrid MoE/Mamba models (Nemotron-3-Nano) ranging from 1B to 4B parameters using Megatron-LM.

Results

• Synthetic Benchmarks: The cubic5 method successfully maps singular values in the protected range $[7 \times 10^{-3}, 1]$ into the target band without numerical blow-up. Microbenchmarks on NVIDIA H200 GPUs show cubic5 reduces the orthogonalization subroutine cost relative to Polar Express, with speedups ranging from 1.16x to 1.52x depending on matrix shape (rectangular vs. square).
• GPT-2 Small Ablations:
  • Five-step cubic5 achieves a final validation loss of 3.2811.
  • Five-step Polar Express achieves 3.2798.
  • Five-step Muon-Jordan achieves 3.2822.
  • Explicit FP32 SVD achieves 3.2809.
  • Observation: The gap between the strongest variants is less than 0.003 in validation loss. Notably, the exact SVD does not outperform the approximate polynomial methods, and adding a fifth step to Muon-Jordan slightly degraded performance compared to four steps.
• Hybrid MoE/Mamba Scaling (1B–4B):
  • On 1B and 2B models, cubic5 slightly outperformed Muon-Jordan (lower validation loss by $\approx 10^{-3}$).
  • On the 4B model, cubic5 was slightly higher ($\approx 10^{-3}$) but the training trajectories were nearly indistinguishable after the initial phase.
  • The relative difference in training loss decayed rapidly, fluctuating around zero for the remainder of the training budget.

Significance and Claims
The paper modestly claims that cubic5 is a practical, low-cost variant for Muon orthogonalization, rather than a universally superior update rule.

• Optimizer vs. Decomposition Accuracy: The results support the hypothesis that Muon benefits from a "spectral reshaping" effect rather than strict polar decomposition accuracy. The exact singular value map that optimizes training depends on the training phase and learning rate, not just mathematical convergence to 1.
• Practical Utility: cubic5 offers a principled point on the cost-quality frontier. It reduces the dominant matrix multiplication count by roughly one-third (10 vs. 15) while maintaining empirical training quality parity with the standard Muon-Jordan quintic update in the tested regimes.
• Limitations: The authors explicitly state that these results establish training-quality parity, not end-to-end throughput speedups. The actual wall-clock speedup depends on hardware kernels, batching, and the fraction of time the orthogonalization step occupies within the full optimizer loop. The paper does not claim cubic5 is a better polar decomposition algorithm than Polar Express in a mathematical sense, but rather a more efficient spectral transform for the specific context of Muon training.