Muon$^2$: Boosting Muon via Adaptive Second-Moment Preconditioning
For practitioners training large foundation models, Muon$^2$ offers a more efficient optimizer that reduces computational overhead while improving performance.
Muon$^2$ improves the Muon optimizer by adding adaptive second-moment preconditioning before orthogonalization, reducing Newton-Schulz iterations by 40% while consistently outperforming Muon and its variants in GPT and LLaMA pre-training from 60M to 1.3B parameters.
Muon has emerged as a promising optimizer for large-scale foundation model pre-training by exploiting the matrix structure of neural network updates through iterative orthogonalization. However, its practical efficiency is limited by the need for multiple Newton--Schulz (NS) iterations per optimization step, which introduces non-trivial computation and communication overhead. We propose Muon$^2$, an extension of Muon that applies Adam-style adaptive second-moment preconditioning before orthogonalization. Our key insight is that the core challenge of polar approximation in Muon lies in the ill-conditioned momentum matrix, of which the spectrum is substantially improved by Muon$^2$, leading to faster convergence toward a practically sufficient orthogonalization. We further characterize the practical orthogonalization quality via directional alignment, under which Muon$^2$ demonstrates dramatic improvement over Muon at each polar step. Across GPT and LLaMA pre-training experiments from 60M to 1.3B parameters, Muon$^2$ consistently outperforms Muon and recent Muon variants while reducing NS iterations by 40\%. We further introduce Muon$^2$-F, a memory-efficient factorized variant that preserves most of the gains of Muon$^2$ with negligible memory overhead.