为什么 Muon 优于 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；当曲率异质性足够强时​​，在相同步数后也会产生较低的局部二次损失。