The Active Ingredient in Muon's Grokking
The Active Ingredient in Muon’s Grokking
Muon 优化器中实现“顿悟”(Grokking)的核心要素
Abstract: The Muon optimizer reaches the grokking threshold on modular arithmetic faster than AdamW. Prior work attributes this to “spectral-norm constraints plus orthogonalized momentum” but does not isolate which mechanism matters. To better understand Muon’s behavior, we run multi-seed and multi-learning-rate sweeps to decompose and stress-test the effect.
摘要: Muon 优化器在模运算任务上达到“顿悟”(grokking)阈值的速度比 AdamW 更快。先前的研究将其归因于“谱范数约束加正交化动量”,但并未明确究竟哪种机制起到了关键作用。为了更好地理解 Muon 的行为,我们进行了多随机种子和多学习率的扫描实验,以分解并压力测试这些效应。
First, an ablation shows the speedup comes from orthogonalization (the Newton-Schulz iteration): orthogonalize-only matches full Muon, whereas spectral-only is no faster than AdamW and is unreliable, and this verdict holds across learning rates.
首先,消融实验表明,加速效果源于正交化(Newton-Schulz 迭代):仅使用正交化的效果与完整版 Muon 相当,而仅使用谱约束的效果并不比 AdamW 快且不可靠,这一结论在不同学习率下均成立。
Second, a mechanistic analysis finds that orthogonalizing optimizers reach generalization at roughly 3x lower spectral norm and, controlling for how much the embedding actually moves, settle into a lower-norm solution rather than simply perturbing the embedding less.
其次,机制分析发现,使用正交化的优化器在谱范数降低约 3 倍时即可达到泛化,并且在控制嵌入层实际移动幅度的情况下,它们倾向于收敛到范数更低的解,而不仅仅是减少了对嵌入层的扰动。
Third, reducing the Newton-Schulz iteration count from five to one accelerates reaching the threshold but makes the grokked solution fragile, prone to transient collapse, with fragility that grows with learning rate; a single iteration is fast and stable only at small learning rate, while the canonical five iterations are the learning-rate-robust choice. We also show spectral scaling can be dropped at no measured cost.
第三,将 Newton-Schulz 迭代次数从 5 次减少到 1 次虽然能加速达到阈值,但会使“顿悟”后的解变得脆弱,容易出现瞬时崩溃,且这种脆弱性随学习率的增加而增强;单次迭代仅在小学习率下既快速又稳定,而标准的 5 次迭代则是对学习率更具鲁棒性的选择。我们还证明,在不产生可测性能损失的情况下,可以舍弃谱缩放(spectral scaling)。
A methodological thread runs throughout: under a stability-aware metric, “faster” claims about grokking optimizers can invert, so we report both first-crossing and sustained-grok times. To support reproducibility, we release our full training and analysis code at this https URL.
贯穿全文的一个方法论核心是:在考虑稳定性的指标下,关于优化器“更快”实现顿悟的结论可能会发生反转,因此我们同时报告了“首次跨越阈值时间”和“持续顿悟时间”。为了支持可复现性,我们已在指定链接发布了完整的训练和分析代码。