[论文解读] MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields

Ilyes Batatia, Dávid Péter Kovács+3 more|arXiv (Cornell University)|Jun 15, 2022

Machine Learning in Materials Science被引用 269

一句话总结

MACE 引入高阶（多体）等变消息传递到原子间势，能够在两层网络上实现最先进的准确性并在 rMD17、3BPA 和 AcAc 的基准测试上实现快速训练。

ABSTRACT

Creating fast and accurate force fields is a long-standing challenge in computational chemistry and materials science. Recently, several equivariant message passing neural networks (MPNNs) have been shown to outperform models built using other approaches in terms of accuracy. However, most MPNNs suffer from high computational cost and poor scalability. We propose that these limitations arise because MPNNs only pass two-body messages leading to a direct relationship between the number of layers and the expressivity of the network. In this work, we introduce MACE, a new equivariant MPNN model that uses higher body order messages. In particular, we show that using four-body messages reduces the required number of message passing iterations to just two, resulting in a fast and highly parallelizable model, reaching or exceeding state-of-the-art accuracy on the rMD17, 3BPA, and AcAc benchmark tasks. We also demonstrate that using higher order messages leads to an improved steepness of the learning curves.

研究动机与目标

• 为化学与材料科学提供快速且准确的力场建模的动机。
• 通过引入更高维度的信息来克服两体 MPNN 的局限性。
• 开发可扩展、并行化的架构，具有高效张量运算。
• 在具有挑战性的基准测试（rMD17、3BPA、AcAc）上展示最先进的准确性。
• 研究身体序（body order）和等变性如何影响学习曲线与泛化。

提出的方法

• 引入 MACE，通过包含多达 nu-体相互作用的分层多体消息构造。
• 将消息表示为在邻居集合上的求和，随身体序的增加而变化（方程7）。
• 通过张量积和 Clebsch–Gordan 耦合构造更高阶的特征以实现等变性（方程8–10）。
• 使用带残差连接的线性更新来传播消息（方程12）。
• 读取特征的不可变部分以得到原子位能，作为层级贡献之和（方程13）。
• 实现高效的循环张量收缩，使高阶消息传递在计算上可行。

实验结果

研究问题

• RQ1 高阶（多体）信息在一个等变 MPNN 中是否能在不牺牲准确性的前提下减少所需的消息传递层数？
• RQ2 增加身体序和等变性如何影响力场预测的学习曲线和数据效率？
• RQ3 更高阶、等变的消息在域外分子构型上的外推和泛化是否更优？
• RQ4 与现有最先进的等变 MPNN 相比，MACE 在标准基准上的速度和可扩展性表现如何？

主要发现

• 更高阶的消息将所需层数降至两层，同时在基准任务上达到或超过最先进的准确性。
• 具有足够身体序的不可变消息可以接近等变模型的性能，而更高阶的等变消息可获得最佳结果。
• MACE 的 L=2（两层，2L+M 等变消息）在 3BPA 外推测试中优于竞争模型，并在基准测试中展现出更优的能量与力预测。
• 相比先前的顶级模型，MACE 在训练和评估速度上具有显著优势，同时保持或提升准确性。
• 学习曲线表明引入更高的身体序信息会改变并加 steep 数据效率斜率，证实了学习动力学的改进。