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

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

arXiv: 2607.02499v1

论文信息

标题: 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

arXiv ID: 2607.02499v1

PDF 链接: 下载 PDF

3 分钟速览

  • 研究问题：这篇论文探索机器学习原子间势（MLIP）训练中长期以来被忽视的优化器选择问题。社区默认使用 Adam/AdamW，但最近在大型语言模型中兴起的矩阵结构优化器能否为 MLIP 训练带来加速和精度提升？

  • 核心方法：将 Muon、SOAP 和 SOAP-Muon 三种矩阵结构优化器集成到 nequip 框架，在液体水（NequIP 模型）和固体电解质 CDP（Allegro 模型）两个系统上系统对比 AdamW，并考察不同力标签比例下的表现。

  • 关键结果：SOAP 相比 AdamW 在两个系统上分别将训练收敛时间缩短 4.9 倍（CDP）和 5.8 倍（水）（见图 2）。SOAP-Muon 仅用 50% 力标签即达到 AdamW 全量力监督的精度水平（见图 1）。

  • 主要局限：仅在两个系统和两种等变架构上验证，尚未在通用势函数预训练或微调场景测试。Muon 在水的实验中持续弱于 AdamW，稳定性和泛化性需进一步研究。

  • 适合读者：从事计算材料、计算化学中 MLIP 训练的研究者和工程师；对优化算法在科学机器学习中应用感兴趣的量化研究员；神经网络训练优化的开发者。

论文背景和研究动机

机器学习原子间势已经成为化学和材料科学中原子尺度模拟的标志性工具。从水的热力学、非晶硅的性质到 HIV 衣壳蛋白的全原子模拟，MLIP 使得原本需要昂贵量子力学计算的模拟变得可行。

在追求更高精度和更强泛化能力的进程中，MLIP 社区的投资集中在两个方向：一是新架构——从早期的描述子方法发展到等变图神经网络和注意力机制；二是新数据集——大规模量子化学数据库覆盖有机分子、催化体系和无机材料。然而，第三个改进轴——决定学习动态的训练技术——却极少受到关注。几乎所有的 MLIP 训练流程默认使用 Adam 或其变体 AdamW，即使在训练通用势函数时亦是如此。

与此同时，大语言模型的规模扩展重新点燃了对超越 Adam 的优化器的兴趣。这些新方法的核心思想是利用神经网络权重张量的矩阵结构进行预条件处理，而不是像 Adam 那样仅依赖逐元素的对角缩放。代表性方法包括：基于显式 Kronecker 分解预条件的 Shampoo、SOAP、Kron；基于半正交化的 Muon、Scion、Gluon；以及两者的混合方案 SOAP-Muon 等。这些方法在大语言模型训练中展现出更快的收敛速度，但它们在 MLIP 训练中的潜力仍未被系统探索。

本文的工作恰好填补这一空白。作者选择 Muon、SOAP 和 SOAP-Muon 作为代表性矩阵结构优化器，在物理上有重要意义的两个体系上对 AdamW 进行系统基准测试，并特别关注稀疏力监督场景——这与更高层次量子力学理论（如耦合簇方法、扩散蒙特卡洛）中力标签昂贵甚至无法获取的现实需求密切相关。

核心方法和技术细节

MLIP 训练的数学框架

MLIP 需要拟合一个将原子位置 {rj,Zj}\{\mathbf{r}_j, Z_j\} 映射到总能量 EθE_\theta 的神经网络。能量分解为局域原子贡献之和以保证尺寸外延性：

Eθ({rj,Zj})=∑iεi,θ({rj,Zj}j∈Ni)E_\theta(\{\mathbf{r}_j, Z_j\}) = \sum_i \varepsilon_{i,\theta}(\{\mathbf{r}_j, Z_j\}_{j\in\mathcal{N}_i})

力则通过自动微分从能量对原子位置的负梯度获得：

Fi,θ=−∂Eθ∂ri\mathbf{F}_{i,\theta} = -\frac{\partial E_\theta}{\partial \mathbf{r}_i}

这一设计天然保证能量守恒。训练损失函数是能量和力误差的加权组合（式 3）。在稀疏力监督场景中，力损失仅在部分构型上计算，极端情况 λF=0\lambda_F = 0 对应纯能量训练。

从 Adam 到矩阵结构优化器

优化器的核心任务是对梯度进行预条件处理，以应对损失函数在不同方向上的各向异性。对于一个向量参数 ww 及其梯度 gg，预条件更新的一般形式为 w←w−ηP−1gw \leftarrow w - \eta P^{-1} g。

Adam 使用最简单的对角预条件 P=diag⁡(v^t)1/2P = \operatorname{diag}(\hat{v}_t)^{1/2}，其中 v^t\hat{v}_t 是逐元素梯度平方的滑动平均。这种设计在计算上极其高效，但完全忽略了参数矩阵的结构信息。

当参数是矩阵 W∈Rm×nW \in \mathbb{R}^{m\times n} 时，一个自然的推广是假设预条件矩阵具有 Kronecker 乘积结构。论文详细描述了这一思路的演化脉络：

  • Shampoo（2018）：定义左预条件子 L∈Rm×mL \in \mathbb{R}^{m\times m} 和右预条件子 R∈Rn×nR \in \mathbb{R}^{n\times n} 分别为梯度矩阵行协方差和列协方差的累积。更新形式为 W←W−ηL−1/4GR−1/4W \leftarrow W - \eta L^{-1/4} G R^{-1/4}。

  • Muon（2024）：基于 Bernstein 和 Newhouse 的洞察——当去除累积步骤时，Shampoo 退化为将梯度矩阵投影到最近的半正交矩阵。Muon 先施加 Nesterov 式动量，再用 Newton-Schulz 迭代实现高效正交化（见算法 1）。

  • SOAP（2025）：核心发现是 Shampoo 等价于在预条件子特征空间中运行 Adafactor。SOAP 将梯度转换到这个特征空间执行 AdamW 更新，再投影回原参数空间。SOAP 成为本文最稳健的优化器选择。

  • SOAP-Muon（2025）：在 SOAP 的基础上叠加 Muon 风格的正交化步骤，形成混合方案（见算法 2）。

技术实现的复杂性

将矩阵结构优化器集成到 nequip 框架并非易事。论文详细描述了如何应对 e3nn 的特殊参数存储方式——其等变算子的权重以扁平化一维向量存储，但这些权重在前向传播中实际上作为结构化矩阵或张量块起作用。作者对 e3nn.o3.Linear 的权重按照指令恢复为二维矩阵块施加 Muon 更新，而对 FullyConnectedTensorProduct 的三维张量块则暂时保持原始形式交给 Adam 处理（见附录 A.2 及表 2）。

创新点和贡献

1. 首次系统验证矩阵结构优化器在 MLIP 训练中的优势

在此之前，矩阵结构优化器主要在大语言模型领域获得验证。Liu 等人（2026b）曾在 MLIP 微调中讨论优化器选择，但仅限于对角预条件方法。Koker 等人（2025）展示了 Muon 在经济型基础势函数训练中的初步结果，但缺乏系统性消融实验。本文首次在两个物性体系、两种等变架构、全量到稀疏力监督的完整谱系上进行系统对比。

2. 发现优化器选择是 MLIP 训练中被忽视的关键设计维度

论文明确提出，训练技术——特别是优化器选择——应当与架构和数据集并列，成为 MLIP 的 “一等公民” 设计轴。在力标签稀疏的场景下，这一发现的经济价值尤为突出：SOAP-Muon 仅使用 50% 力标签即可达到 AdamW 全监督的精度水平（见图 1 和表 5），这意味着当力标签昂贵时，可以通过优化器选择大幅降低数据成本。

3. 揭示 Muon 的不稳定性并提供实用建议

Muon 在水体系上的持续劣于 AdamW 是一个重要发现。这一结果暗示正交化步骤本身可能是性能退化的主要来源（见第 5 节讨论），与物理信息神经网络中 Muon 的稳定性问题相呼应。论文基于完整实验数据推荐 SOAP 作为 “不进一步调参即可替代 AdamW 的最佳默认选择”。

实验结果分析

全监督：精度提升与加速并存

在全额能量和力标签下，矩阵结构优化器在两个体系上均显著优于 AdamW（见表 1、图 1）：

  • CDP/Allegro：SOAP 能量 MAE 降低 9%，SOAP-Muon 力 MAE 降低 14%；
  • 水/NequIP：SOAP-Muon 能量 MAE 降低 24%，SOAP 力 MAE 降低 19%。

更引人注目的是时间效率。尽管矩阵结构方法单步计算开销更大，但由于更好的预条件加速了收敛，所需训练轮数减少的幅度远大于单步开销增加。SOAP 达到 AdamW 最优力精度的时间在 CDP 上加速 4.9 倍，水中加速 5.8 倍（图 2）。

稀疏力监督：优势进一步扩大

随着力标签比例降低，矩阵结构优化器的优势愈发明显：

  • 纯能量训练（力标签占比 0%）：SOAP-Muon 在 CDP 上将能量和力 MAE 分别降低 47% 和 60%；SOAP 在水中将能量和力误差分别降低 26% 和 23%。这证明改进的优化有助于学到的势能面捕获物理上有意义的梯度结构。

  • 极端稀疏场景（力标签仅 5%）：在 CDP 上，SOAP-Muon 保持稳定，力 MAE 仍优于 AdamW 在 100% 力标签下的表现。更重要的是，AdamW 在 5% 力标签下的模型在分子动力学模拟中出现灾难性不稳定性，轨迹在数步内发散；而 SOAP-Muon 模型保持物理真实性，计算的质子扩散活化能 Ea=0.42  eVE_a = 0.42\;\mathrm{eV} 与实验范围 0.39–0.43  eV0.39\text{--}0.43\;\mathrm{eV} 吻合（见附录图 7）。

体系特异性与 Muon 的退化

Muon 在 CDP 上表现尚可，但在水的所有力标签比例设置下均不如 AdamW。论文分析认为正交化步骤是主要退化来源，而 SOAP 的适应性预条件部分缓解了这一效应。SOAP-Muon 也存在敏感性——在水上需要额外调参（动量系数和奇异值指数 ρ\rho，从默认的 ρ=0\rho=0 调为 ρ=0.5\rho=0.5 需要完整 SVD 计算）（见附录 B 中 SOAP-Muon 稳定性部分）。

实践建议

基于论文发现，对 MLIP 训练实践者提出以下建议：

1. 将 SOAP 作为 AdamW 的替代方案优先尝试

在所有测试设置中，SOAP 表现出最稳健的优势，无需大量额外调参。对于从零开始训练 MLIP 的工作，尤其是当追求的是最佳最终精度而非最短训练时间时，SOAP 是当前最安全的默认选择。

2. 在力标签昂贵的计算层级使用 SOAP 或 SOAP-Muon

当参考数据来自耦合簇、扩散蒙特卡洛等力计算代价高昂的方法时，矩阵结构优化器的优势最为突出。实验表明，SOAP-Muon 用 50% 力标签甚至可匹配 AdamW 全量监督的精度。这提供了以更少资源获得同等质量模型的实际路径。

3. 关注优化器配置与模型架构的交互

论文发现水/NequIP 对 SOAP-Muon 的超参数（动量系数 β1\beta_1、β2\beta_2 和奇异值指数 ρ\rho）比 CDP/Allegro 更敏感。在部署到新体系时，建议参照论文做法进行轻量级学习率和权重衰减搜索，对 SOAP-Muon 额外搜索动量组合 {(0.9,0.95),(0.95,0.95),(0.95,0.98)}\{ (0.9, 0.95), (0.95, 0.95), (0.95, 0.98) \}。

4. 避免在无调参情况下使用 Muon 作为唯一优化器

由于 Muon 在水体系上的系统性退化，不建议将其作为直接替代 AdamW 的方案。如果确实需要使用 Muon 以获得其平方根时间复杂度的计算效率，应与自适应预条件方法结合，或至少在新体系上进行充分验证。

5. 在评估 MLIP 质量时超越数值误差指标

论文展示了优化器选择会影响分子动力学模拟中的物理忠实性：AdamW 在 5% 力标签下的模型尽管数值误差尚可接受，但在 MD 模拟中出现结构崩溃。对实际部署而言，MLIP 的最终评判应是物理观测量的复现能力，而非仅看能量和力 MAE。