[论文解读] Learning Atomic Multipoles: Prediction of the Electrostatic Potential with Equivariant Graph Neural Networks

Moritz Thürlemann, Lennard Böselt+1 more|arXiv (Cornell University)|Oct 11, 2021

Machine Learning in Materials Science参考文献 83被引用 47

一句话总结

该论文提出了一种等变图神经网络（GNN），可从分子几何结构中预测至四极矩阶的原子多极矩，从而在无需昂贵的量子力学计算的情况下实现精确的电势能面（ESP）预测。该模型通过设计实现旋转对称性，其ESP预测精度接近DFT参考值，同时相比QM方法将计算成本降低了数个数量级。

ABSTRACT

The accurate description of electrostatic interactions remains a challenging problem for fitted potential-energy functions. The commonly used fixed partial-charge approximation fails to reproduce the electrostatic potential at short range due to its insensitivity to conformational changes and anisotropic effects. At the same time, possibly more accurate machine-learned (ML) potentials struggle with the long-range behaviour due to their inherent locality ansatz. Employing a multipole expansion offers in principle an exact treatment of the electrostatic potential such that the long-range and short-range electrostatic interaction can be treated simultaneously with high accuracy. However, such an expansion requires the calculation of the electron density using computationally expensive quantum-mechanical (QM) methods. Here, we introduce an equivariant graph neural network (GNN) to address this issue. The proposed model predicts atomic multipoles up to the quadrupole, circumventing the need of expensive QM computations. By using an equivariant architecture, the model enforces the correct symmetry by design without relying on local reference frames. The GNN reproduces the electrostatic potential of various systems with high fidelity. Possible use cases for such an approach include the separate treatment of long-range interactions in ML potentials, the analysis of electrostatic potential surfaces, and for the static multipoles in polarizable force fields.

研究动机与目标

• 为解决经典力场中固定电荷近似在捕捉各向异性和构象依赖的电势效应方面的不准确性。
• 通过显式处理高阶多极矩，克服局部机器学习势能的长程精度限制。
• 用可微分、对称性保持的机器学习模型替代昂贵的多极矩量子力学计算。
• 开发一种可预测原子多极矩（单极至四极矩）且对旋转具有等变性的模型，避免依赖局部参考系。
• 实现多极矩在极化力场、ESP表面分析及机器学习势能中长程校正的实际应用。

提出的方法

• 采用基于Satorras等人方法的E(n)-等变GNN架构，以确保多极矩预测中的旋转不变性。
• 构建图表示，其中节点为原子，边编码截断半径内的原子间向量与距离。
• 使用基于原子间向量的向量积与张量积作为特征，通过等变消息传递进行更新。
• 通过将消息传递后的表示与可学习读出层结合，输出原子多极矩（q, μ, θ）。
• 采用基于预测多极矩与参考MBIS推导多极矩、以及ESP值之间差异的损失函数，进行端到端训练。
• 基于邻近原子向量的稀疏局部原子环境表示，实现对未见构象与分子的泛化能力。

实验结果

研究问题

• RQ1 机器学习模型能否在保持旋转对称性的同时，高精度预测至四极矩阶的原子多极矩？
• RQ2 由机器学习预测的多极矩所生成的ESP是否与DFT参考多极矩的精度相当？
• RQ3 该模型能否在不依赖局部参考系的情况下，泛化至未见的分子构象与多样的化学环境？
• RQ4 在ESP保真度与计算成本方面，该模型相较于固定电荷近似与QM参考计算的表现如何？
• RQ5 该模型能否用于改进机器学习势能或极化力场中的长程电势效应？

主要发现

• 与DFT参考值相比，该模型预测的原子多极矩（至四极矩）平均绝对误差为：单极矩0.013 e，偶极矩0.022 D，四极矩0.018 e·bohr²。
• 由多极矩预测的电势能面（ESP）在范德华表面上与DFT参考ESP的平均绝对误差为0.0082 a.u.。
• 即使使用机器学习预测的多极矩而非QM推导的多极矩，引入预测的偶极矩与四极矩仍使ESP误差降低50%以上。
• 该模型的推理速度为每原子0.026–0.075 µs（边特征计算），全预测为3.26–8.21 µs/原子，支持分子模拟中的实时应用。
• 等变架构确保了正确的旋转行为，无需数据增强或局部参考系，提升了泛化能力与物理一致性。
• 该模型在未见构象与分子上泛化良好，在1,013,949个构象的测试集中表现一致，覆盖广泛的化学空间。