The Active Ingredient in Muon's Grokking Muon 泛化的活性成分

TL;DR

• The speedup of the Muon optimizer in reaching grokking thresholds is primarily driven by momentum orthogonalization via Newton-Schulz iteration, not spectral-norm constraints.
• Orthogonalizing optimizers achieve generalization at approximately 3x lower spectral norms and settle into lower-norm solutions compared to non-orthogonalized counterparts.
• Reducing Newton-Schulz iterations from five to one accelerates initial convergence but significantly increases solution fragility and susceptibility to transient collapse, especially at higher learning rates.
• Spectral scaling can be removed without measurable performance cost, simplifying the optimizer implementation while maintaining robustness.

Why It Matters

This analysis clarifies the mechanistic drivers behind the superior performance of the Muon optimizer, helping practitioners understand that orthogonalization is the critical component for efficient grokking rather than spectral constraints. By identifying the trade-offs between convergence speed and solution stability based on iteration counts, it provides actionable guidance for tuning hyperparameters in optimization-heavy tasks. Furthermore, the finding that spectral scaling is unnecessary allows for simplified optimizer implementations without sacrificing performance.

Technical Details

• Ablation Study: The authors decomposed Muon's components, revealing that an "orthogonalize-only" variant matches full Muon performance, whereas a "spectral-only" variant performs no better than AdamW and is unreliable across various learning rates.
• Mechanistic Analysis: Orthogonalizing optimizers reach generalization at roughly three times lower spectral norms. When controlling for embedding movement, these optimizers converge to lower-norm solutions rather than merely perturbing embeddings less.
• Iteration Count Trade-off: Using only one Newton-Schulz iteration speeds up threshold crossing but creates fragile solutions prone to transient collapse. The canonical five iterations provide a robust balance, ensuring stability across different learning rates.
• Methodological Rigor: The study employs multi-seed and multi-learning-rate sweeps to stress-test results. It introduces a stability-aware metric distinguishing between "first-crossing" and "sustained-grok" times to avoid misleading claims about optimizer speed.
• Code Release: Full training and analysis code has been released to support reproducibility of the findings.

Industry Insight

Practitioners using or experimenting with the Muon optimizer should prioritize implementing effective momentum orthogonalization over complex spectral norm constraints, as the latter contributes minimally to the observed speedups. When optimizing for speed, reducing Newton-Schulz iterations may offer gains but requires careful monitoring of solution stability, particularly at higher learning rates where fragility increases. Finally, removing spectral scaling can simplify the optimizer architecture and reduce computational overhead without impacting final model performance, streamlining deployment pipelines.

TL;DR

• Muon优化器在模运算任务上比AdamW更快实现“顿悟”（Grokking），其核心加速机制在于动量的正交化而非谱范数约束。
• 正交化优化器能在约3倍更低的谱范数下达到泛化，且倾向于收敛到更低范数的解，而非仅仅减少嵌入扰动。
• 将Newton-Schulz迭代次数从5次减至1次虽能加速初期收敛，但会导致解的脆弱性和瞬态崩溃风险，5次迭代是兼顾速度与稳定性的最佳选择。
• 研究证实可以移除谱缩放（spectral scaling）而不影响性能，并提出了区分“首次跨越”与“持续顿悟”时间的稳定性评估方法。

为什么值得看

这篇文章深入剖析了近期热门的Muon优化器为何能加速深度学习中的“顿悟”现象，通过消融实验明确了正交化机制的关键作用，为优化器设计提供了理论依据。对于AI从业者而言，理解这一机制有助于在实际训练中调整超参数（如迭代次数），在收敛速度与模型稳定性之间取得最佳平衡。

技术解析

• 消融实验结论：通过多种子和多学习率扫描，研究证明Muon的加速优势主要来源于正交化（Newton-Schulz迭代）。仅使用正交化的变体性能与完整Muon相当，而仅使用谱范数约束的变体则与AdamW无异且不可靠。
• 机制分析：正交化优化器在达到泛化阈值时，其谱范数约为非正交化方法的1/3。在控制嵌入移动量的前提下，正交化方法收敛到更低范数的解空间，表明其找到了更优的几何结构。
• 迭代次数权衡：Newton-Schulz迭代次数直接影响稳定性。1次迭代速度快但解脆弱，易随学习率增加发生瞬态崩溃；5次迭代提供了对 learning rate 的鲁棒性，是推荐配置。
• 方法论改进：指出传统的“更快”指标可能具有误导性，建议同时报告首次达到泛化阈值的时间（first-crossing time）和持续保持泛化的时间（sustained-grok times），以全面评估优化器的稳定性。

行业启示

• 优化器选型策略：在使用Muon等高级优化器时，不应盲目追求极致速度而牺牲稳定性。保留完整的正交化迭代步骤（如5次Newton-Schulz迭代）对于确保模型在复杂任务中的鲁棒性至关重要。
• 泛化能力评估：行业应重新审视“顿悟”相关的评估标准，从单纯关注收敛速度转向关注解的稳定性和泛化持久性，避免被短期的性能提升所误导。
• 简化配置可能性：研究表明谱缩放并非必需，这为简化优化器实现、降低计算开销提供了依据，开发者可尝试移除该组件以优化训练效率。