[论文解读] Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics

Albert Musaelian, Simon Batzner+5 more|arXiv (Cornell University)|Apr 11, 2022

Machine Learning in Materials Science被引用 21

一句话总结

Allegro 是一个严格局部的 E(3)-等变神经网络原子间势，在小分子基准上达到最先进的准确性，同时实现可扩展的并行仿真并对分布外数据具有强鲁棒性。

ABSTRACT

A simultaneously accurate and computationally efficient parametrization of the energy and atomic forces of molecules and materials is a long-standing goal in the natural sciences. In pursuit of this goal, neural message passing has lead to a paradigm shift by describing many-body correlations of atoms through iteratively passing messages along an atomistic graph. This propagation of information, however, makes parallel computation difficult and limits the length scales that can be studied. Strictly local descriptor-based methods, on the other hand, can scale to large systems but do not currently match the high accuracy observed with message passing approaches. This work introduces Allegro, a strictly local equivariant deep learning interatomic potential that simultaneously exhibits excellent accuracy and scalability of parallel computation. Allegro learns many-body functions of atomic coordinates using a series of tensor products of learned equivariant representations, but without relying on message passing. Allegro obtains improvements over state-of-the-art methods on the QM9 and revised MD-17 data sets. A single tensor product layer is shown to outperform existing deep message passing neural networks and transformers on the QM9 benchmark. Furthermore, Allegro displays remarkable generalization to out-of-distribution data. Molecular dynamics simulations based on Allegro recover structural and kinetic properties of an amorphous phosphate electrolyte in excellent agreement with first principles calculations. Finally, we demonstrate the parallel scaling of Allegro with a dynamics simulation of 100 million atoms.

研究动机与目标

• 开发一个严格局部的等变原子间势，达到或超过信息传递方法的准确性。
• 展示对大系统的可扩展性和并行性，同时不牺牲保真度。
• 展示对分布外和更高温度数据的强传递性。
• 通过在复杂电解质的 MD 模拟中再现结构和动力学性质来验证物理真实感。

提出的方法

• 引入 Allegro，这是一个任意深度的 E(3)-等变网络，具有标量（不变）和张量（等变）潜在空间。
• 将总能量分解为按物种和按中心对的能量，以确保尺寸可扩展的局部能量贡献。
• 使用张量积层通过嵌入环境（邻居球面谐波的加权和）传播信息，并通过密度技巧降低计算成本。
• 嵌入两体标量特征，然后在各层中迭代更新不变和等变潜在状态，随后通过 MLP 预测对能量。
• 对内部特征和目标进行归一化，以提高训练稳定性以及在不同系统规模和组成下的尺寸可扩展性。

实验结果

研究问题

• RQ1 严格局部的等变神经网络是否能够达到与信息传递原子间势相媲美甚至更优的最先进精度？
• RQ2 本地化的 Allegro 模型对分布外配置和更高温度的泛化能力有多好？
• RQ3 Allegro 在不损失精度与能量守恒的前提下，是否可扩展到非常大规模的系统？
• RQ4 Allegro 能否在复杂电解质的 MD 模拟中真实再现结构和动力学性质？

主要发现

分子	能量 (meV)	力 (meV/Å)
Aspirin	6.2	2.3
Azobenzene	2.8	2.3
Benzene	0.3	0.2
Ethanol	0.9	2.1
Malonaldehyde	1.5	3.6
Naphthalene	1.2	0.9
Paracetamol	2.9	4.9
Salicylic acid	1.8	2.9
Toluene	1.6	1.8
Uracil	0.6	1.8

• Allegro 在修订后的 MD-17 基准测试上实现了力分量 MAE 的最先进水平，并在若干小分子目标上显示出具有竞争力或优越的性能。
• 在 QM9 基准上，Allegro 的单层张量积层就能超越现有的深度 MPNN 和变换器。
• Allegro 显示出对分布外数据的强传递性，优于其他局部 MLIP。
• 使用 Allegro 的分子动力学在 Li3PO4 的结构和动力学性质方面与第一性原理计算高度一致。
• 该模型在并行上可扩展到模拟多达 1 0000 万原子的系统，展示了实际的大规模适用性。