arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2606.04662cs.LGcs.AI
收藏

为什么 Muon 优于 Adam:曲率视角

Why Muon Outperforms Adam: A Curvature Perspective

National University of Singapore(新加坡国立大学);Yale University(耶鲁大学);University of Minnesota(明尼苏达大学)

Shuche Wang, Fengzhuo Zhang, Jiaxiang Li, Dirk Bergemann, Zhuoran Yang

AI总结:

从曲率视角出发,通过泰勒展开和曲率分解,发现 Muon 因更低的归一化方向锐度(NDS)而比 Adam 实现更大的一步损失下降,数据不平衡和层内曲率是其主要优势来源。

摘要:

Muon 在大语言模型训练中相比 Adam 将训练效率提升约两倍,但这一优势的局部几何来源尚不清楚。我们的工作首次从曲率视角尝试揭开 Muon 优于 Adam 的原因。首先,我们对训练损失曲面应用二阶泰勒近似,表明在匹配验证损失下,Muon 比 Adam 实现更大的一步损失下降。两种优化器的一阶增益相当,但 Muon 始终承受更小的二阶曲率惩罚。其次,我们将该曲率惩罚分解为更新范数的平方和归一化方向锐度(NDS)。我们发现 Muon 和 Adam 的更新范数相当,因此 Muon 更小的曲率惩罚源于更低的 NDS,而非更新尺度。第三,我们研究训练数据和模型结构如何塑造 Muon 的 NDS 优势。使用具有受控不平衡的 Zipf-概率上下文无关文法(PCFG)数据,我们表明数据不平衡放大了 Muon 相对于 Adam 的 NDS 优势。进一步的层内/跨层分解表明,在训练的中后期,Muon 更低的 NDS 主要由更小的层内曲率维持。除了经验证据,我们还分析了具有异质曲率和梯度对齐于高曲率模式的风格化二次问题。我们证明 Muon 通过平衡曲率组间的更新能量,实现了比 GD 更低的平均 NDS;当曲率异质性足够强时,在相同步数后这也产生更低的局部二次损失。

英文摘要:

Muon improves training efficiency over Adam in large language-model training by about two times, but the local geometric source of this advantage remains unclear. Our work takes a first step toward demystifying Muon's superiority over Adam from a curvature perspective. First, we apply a second-order Taylor approximation to the training landscape and show that Muon achieves a larger one-step loss decrease than Adam at matched validation loss. The two optimizers have comparable first-order gains, but Muon consistently incurs a smaller second-order curvature penalty. Second, we decompose this curvature penalty into the squared update norm and Normalized Directional Sharpness (NDS). We find that Muon and Adam have comparable update norms, so Muon's smaller curvature penalty is driven by lower NDS, not update scale. Third, we study how training data and model structure shape Muon's NDS advantage. Using Zipf-Probabilistic Context-Free Grammar (PCFG) data with controlled imbalance, we show that data imbalance amplifies Muon's NDS advantage over Adam. A within-/cross-layer decomposition further shows that, in the middle and late stages of training, Muon's lower NDS is mainly sustained by smaller within-layer curvature. Beyond empirical evidence, we analyze stylized quadratic problems with heterogeneous curvature and gradient alignment toward high-curvature modes. We prove that Muon attains a smaller average NDS than GD by balancing update energy across curvature groups; when curvature heterogeneity is sufficiently strong, this also yields lower local quadratic loss after the same number of steps.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑