How Much Orthogonalization Does Muon Need?

2026-05-29

AI Summary

This work investigates the trade-off between orthogonality precision and training efficiency in the Muon optimizer for neural networks. The authors propose a low-overhead polar decomposition scheme based on a five-step cubic Newton–Schulz iteration (cubic5), which substantially reduces the number of matrix multiplications compared to conventional approaches. Experiments on GPT-2 Small and billion- to tens-of-billions-parameter hybrid MoE/Mamba models demonstrate that cubic5 achieves final losses comparable to those obtained with high-precision singular value decomposition (SVD), despite its significantly lower computational cost. These findings reveal a non-monotonic relationship between training quality and polar decomposition accuracy, underscoring the effectiveness and potential of low-cost orthogonalization strategies in large-scale neural network optimization.

Abstract

Muon optimizers improve neural-network training by replacing ill-conditioned momentum updates with approximately semi-orthogonal updates. This motivates a practical question: how much orthogonalization does Muon actually require? We study this question using a relaxed cubic Newton--Schulz schedule derived directly for Muon's low precision singular value band. The resulting five-step cubic construction uses ten dominant matrix multiplications, compared with fifteen for five quintic Newton--Schulz iterations. The cubic schedule is not intended as a more accurate polar solver; instead, it is a principled low-cost variant that lets us probe the relation between polar accuracy, spectral shaping, and training quality. Across synthetic diagnostics, NanoGPT ablations, and training experiments on hybrid MoE/Mamba models, we find that training quality is not governed monotonically by polar-decomposition accuracy: truncated Polar Express, Muon-Jordan, cubic Newton--Schulz, and an explicit FP32 SVD polar factor can reach nearly indistinguishable final loss on GPT-2 Small, and cubic5 matches the Muon-Jordan quintic update within about $10^{-3}$ validation loss on hybrid MoE/Mamba models with one billion to four billion parameters. These results support cubic5 as a practical low-cost Muon orthogonalization variant, with empirical evidence of training-quality parity in the settings tested.

Problem

• orthogonalization
• Muon optimizer
• polar decomposition
• training quality
• neural network optimization

Innovation

• Muon optimizer
• orthogonalization
• Newton-Schulz iteration
• polar decomposition
• low-cost approximation