Muon优化器在小模型上实现高效训练,结合新架构可大幅降内存提速度。
Muon: Training and Trade-offs with Latent Attention and MoE
- 基于流形自然梯度的理论设计,避免梯度爆炸,收敛更快。
- 大批次训练下数据效率提升,比AdamW少用48-52%计算量。
- 与多头隐空间注意力和专家混合模型联动,内存减68%,推理快3.2倍。
我们对仅使用小型到中型解码器(30M - 200M参数)的Muon优化器进行了全面的理论与实证研究,重点分析其数学基础、收敛特性及与现代架构优化的协同效应。在已有工作表明Muon具备可扩展性的基础上,我们提供了严格的理论分析:(i) 在标准假设下的收敛速率;(ii) 防止梯度爆炸的谱正则化性质;(iii) 与斯蒂费尔流形上的自然梯度下降的联系;(iv) 在谱范数下等价于最陡下降法。关键发现是,Muon在计算-时间权衡中拓展了帕累托前沿,维持大批次下的优异数据效率,这一结论在多个模型规模上得到验证。实验显示,Muon在达到目标损失时仅需AdamW 48-52%的训练计算量,同时保持或提升最终困惑度,与大规模结果一致。当与多头隐空间注意力(MLA)和专家混合(MoE)结合时,实现乘法级效率增益:MLA+MoE+Muon实现68%内存减少和3.2×推理加速,同时困惑度提升8-12%。我们提供了15个架构与优化器组件的详细流程、100多次训练的稳定性分析,以及实践指南,包括由~\cite{su2024muonblog}优化的Newton-Schulz系数(3.4445, -4.7750, 2.0315)。理论分析与全面实验确立了Muon作为比AdamW更稳健、更原则化的替代方案,尤其在结合现代效率技术与大批次训练时表现卓越。
原文摘要 · Abstract (English)
We present a comprehensive theoretical and empirical study of the Muon optimizer for training transformers only with a small to medium decoder (30M - 200M parameters), with an emphasis on its mathematical foundations, convergence properties and synergistic interactions with modern architectural optimizations. Building on recent work showing Muon's scalability, we provide rigorous theoretical analysis including: (i)showing the convergence rate under standard assumptions, (ii) spectral regularization properties that prevent gradient explosion, (iii) connection to natural gradient descent on the Stiefel manifold, and (iv) equivalence to steepest gradient descent under the spectral norm. Crucially, we demonstrate that Muon expands the Pareto frontier in the compute-time trade-off by maintaining superior data efficiency at large batch sizes, a key finding of~\cite{essentialai2025muon} that we validate across our model scales. Empirically, Muon reaches the target loss with 48-52\% of the training calculated by AdamW while maintaining or improving the final perplexity, consistent with larger-scale results. When combined with Multi-Head Latent Attention (MLA) and Mixture-of-Experts (MoE), we observe multiplicative efficiency gains: MLA+MoE+Muon achieves 68\% memory reduction and 3.2$\times$ inference speedup, while improving perplexity by 8-12\%. We provide detailed procedures on 15 architectural and optimizer components, stability analyzes across 100+ training runs, and practical implementation guidelines including Newton-Schulz coefficients $(3.4445, -4.7750, 2.0315)$ optimized by~\cite{su2024muonblog}. Our theoretical analysis and comprehensive experiments establish Muon as a principled, robust alternative to AdamW that particularly excels when combined with modern efficiency techniques and large-batch training regimes.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。