Muon optimizer study reveals orthogonalization is key to faster grokking

A new arXiv paper by Yufeng Wang systematically dissects the Muon optimizer to determine why it reaches the grokking threshold on modular arithmetic faster than AdamW. Through multi-seed, multi-learning-rate sweeps, the author ablates two candidate mechanisms: spectral-norm constraints and orthogonalized momentum (via Newton-Schulz iteration). The results are clear: orthogonalization alone replicates full Muon's speedup, whereas spectral-only performs no better than AdamW and is unreliable across learning rates. Mechanistically, orthogonalizing optimizers achieve generalization at roughly 3x lower spectral norm, settling into a lower-norm solution rather than simply perturbing embeddings less.

The study also uncovers a critical robustness trade-off. Reducing the Newton-Schulz iteration count from the canonical five to a single iteration accelerates reaching the grokking threshold but makes the solution fragile — prone to transient collapse, with fragility increasing at higher learning rates. Single-iteration is fast and stable only at small learning rates, while five iterations provide learning-rate-robust performance. Additionally, the paper finds that spectral scaling can be dropped with no measured cost. A methodological note: the author advocates reporting both first-crossing and sustained-grok times, as speed comparisons can invert under stability-aware metrics. Full code is released for reproducibility.