Muon：隐式注意力与专家混合模型的训练与权衡
机器学习
2025-09-30 v1
摘要
我们提出了一项关于Muon优化器训练Transformer的全面理论与实证研究，仅使用小型至中型解码器（30M-200M参数），重点探讨其数学基础、收敛特性以及与现代架构优化的协同交互。基于近期展示Muon可扩展性的工作，我们提供了严格的理论分析，包括：(i) 在标准假设下证明收敛速率，(ii) 防止梯度爆炸的谱正则化特性，(iii) 与Stiefel流形上自然梯度下降的联系，以及(iv) 在谱范数下等价于最陡梯度下降。关键的是，我们证明Muon在计算-时间权衡中扩展了帕累托前沿，通过在大批量训练下保持优越的数据效率，这是\cite{essentialai2025muon}中的关键发现，我们在我们模型规模上进行了验证。实证上，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的一个原则性、鲁棒的替代方案，在与现代效率技术和大批量训练策略结合时尤其表现出色。
引用
@article{arxiv.2509.24406,
title = {Muon: Training and Trade-offs with Latent Attention and MoE},
author = {Sushant Mehta and Raj Dandekar and Rajat Dandekar and Sreedath Panat},
journal= {arXiv preprint arXiv:2509.24406},
year = {2025}
}