[论文解读] Rotation Invariant Graph Neural Networks using Spin Convolutions

Muhammed Shuaibi, Adeesh Kolluru+5 more|arXiv (Cornell University)|Jun 17, 2021

Machine Learning in Materials Science参考文献 34被引用 51

一句话总结

SpinConv 引入一个对旋转不变的图神经网络，用于原子结构，通过对边使用局部坐标系和对剩余旋转自由度的自旋卷积，在能量和力估计上实现 OC20 的最新结果。

ABSTRACT

Progress towards the energy breakthroughs needed to combat climate change can be significantly accelerated through the efficient simulation of atomic systems. Simulation techniques based on first principles, such as Density Functional Theory (DFT), are limited in their practical use due to their high computational expense. Machine learning approaches have the potential to approximate DFT in a computationally efficient manner, which could dramatically increase the impact of computational simulations on real-world problems. Approximating DFT poses several challenges. These include accurately modeling the subtle changes in the relative positions and angles between atoms, and enforcing constraints such as rotation invariance or energy conservation. We introduce a novel approach to modeling angular information between sets of neighboring atoms in a graph neural network. Rotation invariance is achieved for the network's edge messages through the use of a per-edge local coordinate frame and a novel spin convolution over the remaining degree of freedom. Two model variants are proposed for the applications of structure relaxation and molecular dynamics. State-of-the-art results are demonstrated on the large-scale Open Catalyst 2020 dataset. Comparisons are also performed on the MD17 and QM9 datasets.

研究动机与目标

• 以旋转不变的 GNNs 为手段，推动对原子力和能量的高效、接近密度泛函理论（DFT）预测。
• 开发一个每条边的局部坐标系，使角度信息得以捕获，同时不损失旋转不变性。
• 提出对剩余旋转自由度进行自旋卷积，以实现旋转不变量的边信息传递。
• 给出以能量为中心与以力为中心的 SpinConv 变体，以在能量守恒与直接力估计之间取得平衡。
• 在 OC20 上展示最先进的性能，在 MD17 和 QM9 上展现有竞争力的结果。

提出的方法

• 将原子结构表示为图结构，原子作为节点，邻居边在截断距离内。
• 用一个共享更新迭代计算边消息 h_st，其中包含自旋卷积、距离编码和原子序数嵌入。
• 使用由源节点 s 和目标节点 t 定义的局部坐标系，将邻居投影到以 s 为中心的球面上，然后在 st 轴周围对旋转自由度进行自旋卷积。
• 在训练目标中，采用能量为中心的路径（通过 -∂E/∂x 得到力）或力为中心的路径（直接预测力）。
• 采用基于网格的或球面调和函数的球面表示来执行自旋卷积，随后进行池化以确保旋转不变性。
• 加入一个嵌入块（混合专家风格）以将原子序数信息融入消息更新。
• 在训练期间可选地对整个结构进行旋转，以提高鲁棒性并在力块中促进近似的旋转等变性。

实验结果

研究问题

• RQ1 基于自旋卷积的旋转不变 GNN 是否能够准确学习原子系统的能量和力场？
• RQ2 能量为中心与力为中心的 SpinConv 变体在能量守恒和预测精度方面有何比较？
• RQ3 表示选择（网格基 vs 球面调和函数）和模型规模对 OC20、MD17、QM9 表现的影响？
• RQ4 通过多次随机旋转引入旋转鲁棒性在多大程度上提升力预测？
• RQ5 SpinConv 在 OC20 的域内与域外（吸附体/催化剂）分割上的表现如何？

主要发现

• 在 OC20 的 S2EF 任务中，采用力为中心训练的 SpinConv 获得了最低的能量 MAE 和力 MAE。
• 网格基自旋卷积训练更快，在 OC20 上与球面调和函数相比具有竞争力的精度；两种表示均获得相似的准确度。
• 以能量为中心的 SpinConv 提高了能量精度，但在整体 S2EF 指标和 IS2RS 结果中不及以力为中心的变体。
• 在 OC20 的 IS2RE 上，SpinConv 超越了先前方法，特别是在放松能量预测的直接回归设置中。
• 在 OC20 的 IS2RS 上，SpinConv 获得最好 的 ADwT（在各分割上的平均为 53.62%），但 DimeNet++-large 获得最佳 AFbT（21.82%），SpinConv 仍提供更快的推理时间（约 263.2h 对 814.6h）。
• 在 MD17 和 QM9 中，SpinConv 显示出与最先进方法竞争甚至优越的表现，特别是在利用力为中心的公式和合适的表示尺寸时。