Fast and flexible long-range models for atomistic machine learning

本文介绍了一个快速、灵活且模块化的框架（在 PyTorch 和 JAX 中实现），该框架将成熟的长程相互作用算法（如埃瓦尔德求和与粒子网格埃瓦尔德方法）整合到原子机器学习之中，从而实现了物理长程力与局部模型的无缝结合，以克服在描述静电作用及其他长程效应方面的局限性。

要点

想象一下，你正在试图预测一大群人在巨型体育场内将如何移动和互动。在原子世界里，科学家们利用“机器学习”（AI）来做这件事。通常，这些 AI 模型就像戴着马眼罩的人：它们只关注那些紧挨着或站在正旁边的邻居。这对于短程相互作用（比如人群中的握手或碰撞）来说效果很好。

然而，原子之间也存在“长程”关系。这就像体育场里的扩音器：即使你离得很远，你仍然能听到音乐（或感受到静电）。在物理学中，这被称为静电学。传统的 AI 模型往往忽略这一点，因为计算体育场里每一个原子如何影响其他每一个原子，计算成本太高了。

这篇论文介绍了一套新的工具包（适用于 PyTorch 和 JAX 的库），它就像是为这些 AI 模型配备的超高效音响系统。它让 AI 能够“听到”远处的原子，而不会被缓慢、繁重的计算所拖累。

以下是他们解决方案的分解，使用了简单的类比：

1. 问题：“马眼罩”与“整个体育场”

大多数原子 AI 模型依赖于“局域性”规则：“我只关心手臂范围内能触及的原子。”

• 问题所在：这对于离子晶体（盐）或水等物质是行不通的，因为电力会延伸到整个系统。忽略“远处的观众”会导致对材料行为的预测错误。
• 旧有的修补方法：以前的尝试就像试图手动向体育场里的每一个人逐一喊话。虽然准确，但速度极慢且难以设置。

2. 解决方案：“网格”与“拆分”

作者构建了一个框架，将物理学中三种经典的快速方法引入现代 AI 世界。他们称之为范围分离。

将两个原子之间的相互作用想象成一次对话：

• 耳语（短程）：这是原子靠近时发生的情况。它复杂且具体。AI 通过观察直接邻居（即“耳语”）来处理这部分。
• 广播（长程）：这是平滑、衰减缓慢且能延伸到远处的电力。新方法不使用网格（或网）来捕捉“广播”，而不是计算每一个连接。

类比：
想象你试图计算房间里的温度。

• 旧方法：你测量空气中每一个点的温度，然后取平均值。（太慢了）。
• 新方法（PME/P3M）：你在墙上放置一个传感器网格（网格）。你使用一种快速的数学技巧（傅里叶变换）来计算网格上的“平滑”热流，然后只需检查人们（原子）站立的具体位置。这要快得多，并且即使房间变得巨大，其扩展性也很好。

3. “净化”的描述符（“外部”视角）

该论文的一项巧妙创新是他们称之为**外部势特征（EPFs）**的东西。

• 问题：如果你试图描述原子上的“长程”力，信号通常会被其直接邻居的“短程”噪声所淹没。这就像站在电钻旁边试图听远处的警笛声。
• 修补方法：作者创建了一个“过滤器”，在数学上静音了直接邻居。他们只让 AI“听”到某个圆圈之外的原子。
• 结果：这为 AI 提供了长程环境的“干净”信号，然后它可以将其与处理“电钻”（短程）噪声的单独模型相结合。这使得整个系统更准确且更易于训练。

4. 为什么它很灵活（“乐高”方法）

作者没有建造一台僵硬的机器，而是建造了一套乐高积木。

• 模块化：你可以将这些长程计算器连接到任何现有的 AI 模型上。
• 可微分：由于他们使用流行工具（PyTorch 和 JAX）构建了它，AI 可以自动找出如何调整自己的设置（例如电荷应该有多强），以便从数据中学习。这就像一辆在行驶中就能调整自己引擎的汽车。
• 快速：他们在多达 260,000 个原子的系统上进行了测试。他们的方法速度足够快，可以运行以前对机器学习来说太慢的模拟。

5. 他们实际做了什么（基准测试）

这篇论文并没有声称治愈了某种疾病或发现了新材料。相反，他们通过以下方式证明了他们的工具有效：

• 速度测试：显示他们的代码在大型系统上的运行速度与行业标准物理软件（LAMMPS）一样快（甚至更快）。
• 准确性测试：显示当他们模拟水或盐晶体时，结果与已知的物理学完全吻合。
• 学习测试：显示 AI 可以仅通过查看数据来“学习”原子的正确电荷，而无需事先被告知答案。

总结

简而言之，这篇论文提供了一个快速、灵活且模块化的工具包，让 AI 模型能够“看到”原子之间的长距离电力。通过将问题拆分为“近距离”和“远距离”部分，并使用智能网格系统来计算远距离部分，他们使机器学习能够以高精度和速度处理复杂材料（如盐和水），而以前要高效地做到这一点是非常困难的。

技术摘要

技术摘要：用于原子机器学习的快速且灵活的长程模型

1. 问题陈述

大多数原子机器学习（ML）模型依赖于局域性假设，将系统能量分解为短程、以原子为中心的贡献之和。虽然这种方法效率很高，但它无法准确描述由长程相互作用主导的物理现象，最显著的是静电作用和色散力。这些相互作用对于离子材料、极性系统、层状材料和分子晶体至关重要，影响着介电常数、声子谱和结构稳定性等性质。

现有尝试引入长程效应的方法通常存在两个主要局限性：

1. 实施障碍：长程相互作用的高效算法（如 Ewald 求和、粒子网格 Ewald）传统上是在经典分子动力学（MD）代码中实现的，但难以集成到现代可微分的机器学习框架中。
2. 描述符污染：许多包含长程项的机器学习方法仍然依赖于混合了短程和长程信息的描述符。由于原子处的势能在数值上主要由近邻原子主导，“长程”信号往往受到短程贡献的污染，使得单独隔离和学习非局域效应变得困难。

2. 方法论

作者提出了一个框架及参考实现（用于 PyTorch 的 torch-pme 和用于 JAX 的 jax-pme），将成熟的长程算法集成到原子机器学习中。核心方法论包括：

A. 范围分离与算法

该框架实施了一种范围分离策略，将成对势 $v(r)$ 分解为短程（$v_{SR}$）和长程（$v_{LR}$）分量：
$$v(r) = v_{SR}(r) + v_{LR}(r)$$
短程部分通过截断半径 $r_{cut}$ 内的近邻列表直接求和计算。长程部分则采用以下方法处理：

• Ewald 求和：适用于小到中等规模系统，利用实空间和倒易空间求和。
• 粒子网格方法（PME, P3M, SPME）：适用于大规模系统，这些方法将粒子电荷插值到网格上，执行快速傅里叶变换（FFT）以计算倒易空间贡献，并实现 $O(N \log N)$ 的标度。
• 泛化：该实现支持任意反幂律势 $v(r) \propto 1/r^p$（例如，$p=1$ 对应库仑势，$p=6$ 对应色散力），利用广义不完全 Gamma 函数进行范围分离。

B. 模块化与可微分架构

该库采用模块化结构设计，包含：

• 势函数类：计算 $v(r)$、$v_{SR}(r)$、$v_{LR}(r)$ 以及傅里叶变换 $\hat{v}_{LR}(k)$。
• 网格插值器：将粒子位置和伪电荷转换为密度网格，并将场插值回粒子位置。
• K 空间滤波器：执行傅里叶域卷积。
• 计算器：组合这些模块以评估势能和力。
  至关重要的是，所有组件均在自动微分框架（PyTorch/JAX）中实现，使得长程模型能够与局部机器学习方案无缝结合，并通过梯度下降优化参数（如原子电荷、相互作用指数）。

C. 外部势特征（EPFs）

为了解决短程污染问题，作者引入了外部势特征（EPFs）。与对所有近邻求和的标准势函数不同，EPFs 使用平滑过渡函数 $f_{trans}(r)$ 明确排除截断半径 $r_{cut}$ 内原子的贡献。这产生了仅包含长程信息的“纯化”描述符，使其适合与单独的短程机器学习模型结合使用。

D. 自动超参数调整

该框架包含内置功能，可自动调整数值参数（网格间距、实空间截断、展宽参数 $\sigma$），以在最小化计算时间的同时满足目标力精度 $\epsilon_{target}$。

3. 主要贡献

1. 参考实现：发布 torch-pme 和 jax-pme，为原子机器学习提供高效、可微分的 Ewald、PME 和 P3M 算法实现。
2. 纯化描述符：形式化并实现外部势特征（EPFs），以将长程贡献从短程噪声中隔离出来。
3. 灵活性：支持任意晶胞（包括三斜晶系）、任意幂律指数（$p > 0$），以及直接从数据中学习相互作用参数（电荷、指数）的能力。
4. 集成：模块化设计允许这些物理长程计算器作为构建复杂等变机器学习架构（如长程等变或 LODE 特征）的基石。

4. 结果与基准测试

该论文通过多个基准测试验证了该框架：

• 精度：实现针对各种晶体结构（如 NaCl、CsCl）达到了目标力精度（相对误差低至 $10^{-9}$），使用了 Ewald 和基于网格的方法。自动调整过程成功收敛至这些目标。
• 计算成本：
  • 对于小系统（$N < 1000$），Ewald 实现具有竞争力，但由于初始化开销，略慢于 LAMMPS。
  • 对于较大系统（$N > 10^4$），基于网格的（PME/P3M）实现展示了预期的 $O(N \log N)$ 标度，在 $N=10^4$ 时比 $O(N^2)$ 的 Ewald 方法快约 5 倍。
  • 在速度和精度方面，这些实现与 LAMMPS 的 P3M 实现具有竞争力。
• 分子动力学（MD）：使用 torch-pme PME 实现进行的 2 纳秒 NpT 刚性 SPC/E 水模拟，产生的径向分布函数和等温压缩率值与纯 LAMMPS 模拟一致，验证了其作为经验力场引擎的用途。
• 学习能力：
  • 该框架成功学习了 NaCl 结构的正确原子电荷，并恢复了相互作用势的正确函数形式（$1/r$）。
  • 在用于有机分子的“第三代”神经网络势中，一个结合短程 SOAP 神经网络与长程库仑项（使用 EPFs）的模型，尽管仅使用单个描述符，但其精度与之前使用更复杂 LODE 描述符的工作相当。

5. 意义与主张

作者声称，这项工作提供了一个快速、灵活且模块化的框架，弥合了经典长程静电算法与现代原子机器学习之间的差距。

• 可及性：通过在流行的机器学习库（PyTorch/JAX）中提供这些算法，该工作消除了在机器学习模型中实现高效长程相互作用的障碍。
• 模块化：短程和长程分量的分离允许构建“范围分离”模型，其中物理相互作用可作为构建更复杂架构（包括预测张量性质或电子密度的架构）的基石。
• 可扩展性：粒子网格方法的使用确保了这些长程模型可扩展至大规模系统（$N \sim 10^5$），克服了机器学习工作流中二次标度 Ewald 求和的局限性。
• 纯化：EPFs 的引入提供了一种原则性的方法来构建真正长程的描述符，避免了标准基于势的描述符中普遍存在的短程信息冗余问题。

该论文 concludes 指出，这些库旨在鼓励开发更标准化、高效且可扩展的长程机器学习模型，超越简单的点电荷近似，转向更通用且物理信息更丰富的架构。