Computer Science > Machine Learning
arXiv:2607.20512 (cs)
Submitted on 6 Jul 2026
Authors: Yufeng Wang
Abstract
The Muon optimizer reaches the grokking threshold on modular arithmetic tasks significantly faster than AdamW. Prior work attributes this advantage to "spectral-norm constraints combined with orthogonalized momentum," but does not isolate which specific mechanism drives the improvement. To better understand Muon's behavior, we conduct multi-seed and multi-learning-rate sweeps to decompose and stress-test the effect.
First, ablation studies reveal that the speedup stems from orthogonalization via the Newton-Schulz iteration: an orthogonalize-only variant matches full Muon performance, whereas spectral-only constraints are no faster than AdamW and prove unreliable—a result that holds across a range of learning rates. Second, mechanistic analysis finds that orthogonalizing optimizers achieve generalization at roughly three times lower spectral norm. Controlling for actual embedding movement, these optimizers settle into a lower-norm solution rather than simply perturbing the embedding less. Third, reducing the Newton-Schulz iteration count from five to one accelerates reaching the grokking threshold but makes the grokked solution fragile, prone to transient collapse—with fragility increasing alongside learning rate. A single iteration proves fast and stable only at small learning rates, while the canonical five iterations offer learning-rate-robust stability. We also show that spectral scaling can be dropped with no measurable performance cost.
A methodological thread runs throughout: under a stability-aware metric, claims of "faster" grokking can invert. We therefore report both first-crossing and sustained-grok times. To support reproducibility, we release our full training and analysis code at this GitHub repository.
Subjects
- Machine Learning (cs.LG) (Primary)
- Artificial Intelligence (cs.AI)
Cite as
- arXiv: 2607.20512 [cs.LG]
- Version: v1
- DOI: 10.48550/arXiv.2607.20512 (pending registration)
Submission History
From: Yufeng Wang [view email]
- v1: Mon, 6 Jul 2026 22:35:57 UTC (34 KB)
via ArXiv LG