超越 Adam：SOAP 与 Muon 实现更快速、标签高效型机器学习原子间势训练

论文信息

标题: Beyond Adam: SOAP and Muon for Faster, Label-Efficient Training of Machine Learning Interatomic Potentials

作者: Gil Harari, Yoel Zimmermann, Ola Tangen Kulseng, et al.

发布日期: 2026-07-02

PDF 链接: 下载 PDF

研究背景与核心动机

机器学习原子间势函数（MLIP）已成为化学与材料科学中原子级模拟的核心工具，它们通过学习第一性原理计算产生的高精度能量和力，大幅加速了从水热力学到固态电解质等系统的研究。社区在模型架构和新数据集上投入巨大，典型成果包括等变图神经网络 NequIP、局部等变模型 Allegro 以及大规模量子化学数据库。然而，训练过程的优化器选择长期以来被严重忽视，绝大多数工作直接套用 Adam 或 AdamW，这一 “默认选项” 几乎从未被质疑。

与此同时，大语言模型的规模化训练重新点燃了对矩阵结构优化器的兴趣。这类方法通过利用权重张量中的矩阵结构，突破了 Adam 逐元素对角线缩放的局限。其中有两大代表方向：一是以 Shampoo、SOAP 为代表的显式 Kronecker 预处理方法；二是以 Muon 为代表的正交化矩阵更新方法，以及融合二者的 SOAP-Muon。尽管它们在 LLM 中已展现出更快的收敛速度和更高的最终精度，但在 MLIP 训练中的潜力几乎无人涉足。本文正是为了填补这一空白，系统性将 Muon、SOAP 和 SOAP-Muon 引入 MLIP 训练，并在液体水和固态电解质 CsH₂PO₄ (CDP) 两个物理体系上与 AdamW 进行严格基准测试，尤其关注在力标签稀疏场景下的表现。

核心方法：矩阵结构优化器的工作原理

MLIP 训练与优化视角

MLIP 将原子结构和元素种类映射为体系总能量 $E_\theta$ ，通常写为原子能量之和以保证尺寸一致性，并可自动求导得到原子力 $\mathbf{F}_{i,\theta} = -\partial E_\theta/\partial \mathbf{r}_i$ 。标准训练目标由能量和力的均方误差加权组合构成。当力标签昂贵或缺失时，仅用能量项训练（ $\lambda_F=0$ ）就成为一种重要的实际场景，因为耦合簇、扩散蒙特卡罗等方法计算力代价极高，而即使力标签可获得，仅使用部分力的稀疏训练也能降低内存和时间开销。

优化器的核心是预条件处理。参数更新的通用形式为 $w \leftarrow w - \eta P^{-1} g$ ，其中 $P$ 是对损失局部曲率的近似。Adam 将 $P$ 视为元素级平方梯度的对角矩阵，这忽略了权重矩阵的行列结构。矩阵结构优化器则假设对于矩阵参数 $W$ ，其梯度 $G$ 的预处理可以分解为左右矩阵的作用： $W \leftarrow W - \eta L^{-p} \, G \, R^{-p}$ ，这来源于对全 Hessian 的 Kronecker 近似 $P \approx L^T \otimes R$ 。

Muon：半正交更新

Muon 吸收了这一思想的一个极端形式：它发现如果去掉 Shampoo 中的累积步骤，预条件更新就退化为对动量化梯度矩阵的半正交化——即在 Frobenius 范数下寻找与梯度矩阵最接近的半正交矩阵。实现时，Muon 先对矩阵参数做 Nesterov 加速的动量更新，然后通过 Newton–Schulz 迭代高效地计算出正交化的更新方向。此过程中的关键是将中间层的权重矩阵和非矩阵参数（如偏置、嵌入）分离，后者仍由 Adam 更新。对于 NequIP 中等变算子内部以平坦向量形式存储的结构化权重，需要先恢复其二维或多维块状形式，更新后再压平，这使得集成并不平凡。

SOAP：Shampoo 特征空间中的自适应优化

SOAP 意识到 Shampoo 本身等价于在左、右预条件子的特征空间内运行 Adafactor。因此 SOAP 的思路是：维护梯度外积的指数移动平均 $L$ 和 $R$ ，每若干步更新它们的特征向量 $Q_L, Q_R$ ，将梯度投影到该特征基上： $\tilde{G} = Q_L^T G Q_R$ ，然后在此空间内执行标准的 AdamW 更新（计算一阶矩、二阶矩和方差归一化），再将更新方向反投影回原空间： $U = Q_L \tilde{U} Q_R^T$ 。这种设计将 Adam 自适应性与二阶结构信息无缝结合，实践中比 Muon 更加鲁棒，且保留了似然于 Adam 的计算图。

SOAP-Muon：双重融合

SOAP-Muon 在 SOAP 的反投影方向 $U$ 后额外引入一个可选的 Muon 风格正交化步骤：对 $U$ 执行 Newton–Schulz 或 SVD 变换（带可调节的奇异值幂次 $\rho$ ），并可选择做 RMS 归一化。这使得更新同时受益于特征空间自适应性与矩阵正交化，但在某些任务上可能出现不稳定性，需要调整动量参数和 $\rho$ 值。

创新点与主要贡献

系统对比与基准：首次将三类矩阵结构优化器整合到同一 MLIP 框架（nequip），并对液体水和固态 CDP 两个性质差异显著的体系进行全面评估。消除了以往仅关注架构或数据集而忽视优化器的盲点。
稀疏力监督场景下的显著优势：实验表明，SOAP-Muon 仅用 50% 的力标签即可达到 AdamW 在全部力标签下的精度；在能量仅训练和极端稀疏（5% 力标签）时，SOAP 家族较 AdamW 可将能量误差降低近一半，力误差降低 60%，展现出极强的数据效率，这直接回应了更高精度数据昂贵时标的需求。
物理真实性的验证：不仅报告误差数值，还通过分子动力学模拟考察径向分布函数、均方位移以及质子扩散活化能等实验可观测量，证明优化器带来的精度提升最终转化为更可靠的物理性质预测，而 AdamW 在某些稀疏标签下会完全失效。
实用建议与稳定性分析：给出了不同体系下各优化器的最佳超参数、参数分组策略以及 SOAP-Muon 需注意的稳定性调节方法，为从业者提供了可操作的替代方案。

实验结果深度剖析

全力监督下的准确性与加速

在 CDP (Allegro) 上，SOAP 取得最佳能量 MAE（降低 9%），SOAP-Muon 取得最佳力 MAE（降低 14%）；Muon 也有提升，但弱于前两者。在水 (NequIP) 上，SOAP-Muon 降低能量误差达 24%，SOAP 降低力误差 19%，而 Muon 反而不如 AdamW，尤其是在能量预测上。壁钟时间加速更为惊人：SOAP 在 CDP 上达到 AdamW 最优力的速度快 4.9 倍，水上快 5.8 倍，即便矩阵操作的每步额外开销也被有效降低的训练次数所补偿。

稀疏力标签下的表现

能量仅训练时，SOAP-Muon 在 CDP 上能量误差仅为 AdamW 的 53%，力误差仅为 40%；水上 SOAP 也分别优化 26% 和 23%。随着力标签比例逐渐降低至 5%，SOAP 和 SOAP-Muon 依然稳定，而 AdamW 精度急剧恶化。特别地，在 5% 力标签下，CDP 的 SOAP-Muon 模型不仅保持稳定，其质子扩散活化能（0.42 eV）与实验值（0.39–0.43 eV）高度吻合；而 AdamW 则轨迹发散，无法正常模拟。水上各优化器在稀疏时仍能维持合理的 RDF，但 SOAP 在 5% 和 10% 力标签下力误差仍明显优于 AdamW。

不同优化器的行为差异

SOAP 在两个体系上均表现最稳健，可看作可直接替换 AdamW 的首选。SOAP-Muon 潜力更大但超参数敏感，尤其在 NequIP 水上需要精细调节动量系数和 $\rho=0.5$ 避免震荡。Muon 的退化表明纯正交化在某些权重矩阵的谱特性下可能引入不稳定性，而 SOAP 的预条件有效缓解了这一问题。这一发现与物理信息神经网络中对 Muon 进行谱引导修正的必要性相呼应。

实践应用建议与未来方向

对于 MLIP 从业者，建议将 SOAP 作为默认替代 AdamW 的优化器，以在几乎不增加调参负担的情况下获得更快收敛和更高精度，尤其在可获取的力标签有限时（如耦合簇或多体微扰计算）。若算力允许进一步调参，SOAP-Muon 可在某些体系上获得额外增益，但须警惕其需要更小学习率和特殊动量配置的风险。集成 SOAP 时，务必像作者那样正确重塑等变网络中的扁平参数块，并分离嵌入和读出层用 Adam 更新。

未来研究可扩展到：

大规模基础模型：当前结果局限于小规模单体系训练，需验证 SOAP 在像 UMA、SevenNet 等通用基础势能的预训练和微调中是否依然优势明显，这对减少全元素力标签需求量意义重大。
架构泛化：在注意力机制 MLIP 和其他等变框架上测试，研究 SOAP 与不同归纳偏置的相互作用。
稳定性理论：深入理解 Muon 在部分权重上失效的谱原因，并提出自动化调整策略（如自适应 $\rho$ ）。
与二阶优化结合：探索在 SOAP 特征空间中融合 Hessian 信息（如 K-FAC），进一步提升条件数。

总结与展望

本文以扎实的实验揭示了优化器选择在 MLIP 训练中的关键作用：矩阵结构优化器 SOAP 和 SOAP-Muon 在收敛速度、最终精度和稀疏标记效率上全面超越沿用多年的 AdamW，且这种优势最终转化为准确的物理模拟。SOAP 的鲁棒性使其成为即插即用的改进方案，而 SOAP-Muon 则为极端数据效率场景提供了武器。这项工作将训练技巧提升为与架构和数据并重的设计维度，为机器学习势能迈向更大规模、更高准确性和更通用的未来开辟了一条新路径。