[论文解读] Long-Short-Range Message-Passing: A Physics-Informed Framework to Capture Non-Local Interaction for Scalable Molecular Dynamics Simulation

Yunyang Li, Yusong Wang+8 more|arXiv (Cornell University)|Apr 26, 2023

Machine Learning in Materials Science被引用 6

一句话总结

该论文提出了一种物理信息图神经网络框架——长-短程消息传递（LSR-MP），通过基于分割的消息传递机制显式建模长程相互作用，从而增强等变图神经网络。该方法在MD22和Chignolin数据集上实现了高达40%的误差降低，同时保持了计算效率和等变性。

ABSTRACT

Computational simulation of chemical and biological systems using ab initio molecular dynamics has been a challenge over decades. Researchers have attempted to address the problem with machine learning and fragmentation-based methods. However, the two approaches fail to give a satisfactory description of long-range and many-body interactions, respectively. Inspired by fragmentation-based methods, we propose the Long-Short-Range Message-Passing (LSR-MP) framework as a generalization of the existing equivariant graph neural networks (EGNNs) with the intent to incorporate long-range interactions efficiently and effectively. We apply the LSR-MP framework to the recently proposed ViSNet and demonstrate the state-of-the-art results with up to 40% MAE reduction for molecules in MD22 and Chignolin datasets. Consistent improvements to various EGNNs will also be discussed to illustrate the general applicability and robustness of our LSR-MP framework. The code for our experiments and trained model weights could be found at https://github.com/liyy2/LSR-MP.

研究动机与目标

• 解决现有等变图神经网络在捕捉大分子系统中长程和多体相互作用方面的局限性。
• 在不牺牲计算可扩展性或旋转等变性的情况下，高效地将长程相互作用整合到图神经网络势能中。
• 在多种EGNN架构和分子数据集上推广该框架，证明其广泛适用性。
• 通过受分割启发的消息传递方案，引入非局部相互作用，提升从头算分子动力学模拟的精度。
• 通过解耦短程与长程消息传递，保持数据效率并避免信息过度压缩。

提出的方法

• LSR-MP框架引入了双层消息传递架构：短程模块用于截断半径内的局部多体相互作用，长程模块用于片段之间的相互作用。
• 采用基于分割的方法建模长程相互作用，将分子划分为子系统，并显式计算子系统间的相互作用。
• 长程消息传递被实现为片段与原子之间的二部图网络，通过向量和标量嵌入变换保持旋转等变性。
• 通过确保所有消息传递操作与旋转变换可交换，框架保持了等变性，其证明基于向量和标量嵌入的旋转等变更新规则。
• 该方法使用加权损失函数，结合能量和力预测的均方误差，通过超参数控制能量与力精度之间的权衡。
• 其具体实现为ViSNet-LSRM，即一种集成长程消息传递机制的ViSNet变体，通过向量-标量交互机制保持等变性。

实验结果

研究问题

• RQ1 消息传递框架能否在保持旋转等变性的同时，有效捕捉分子系统中的长程相互作用？
• RQ2 基于分割的长程消息传递如何提升大型分子系统中图神经网络势能的准确性？
• RQ3 LSR-MP框架在不同EGNN主干网络和分子数据集上的泛化能力如何？
• RQ4 与标准EGNN相比，集成长程相互作用是否能降低能量和力预测的误差？
• RQ5 该框架在建模非局部相互作用时，能否同时保持计算效率和数据效率？

主要发现

• 与基线模型相比，LSR-MP框架在MD22和Chignolin数据集上的能量和力预测误差最高降低40%。
• ViSNet-LSRM作为LSR-MP的实现，以更少的参数量和更短的训练时间达到最先进性能，优于以往模型。
• 该框架在多种EGNN主干网络（包括PaiNN、ET和Equiformer）上均表现出性能提升，证明了其广泛的兼容性和鲁棒性。
• 通过数学上证明的向量和标量嵌入在旋转下的变换规则，该方法在理论上保持了旋转等变性。
• 长程消息传递机制有效捕捉了长程电荷和范德华力等非局部相互作用，而这些是标准局域消息传递所忽略的。
• 消融研究证实，长程模块对性能提升有显著贡献，尤其在具有强非局部相互作用的系统中。