[论文解读] Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs

Yi-Lun Liao, Tess Smidt|arXiv (Cornell University)|Jun 23, 2022

Machine Learning in Materials Science被引用 63

一句话总结

Equiformer 是一个 SE(3)/E(3)-等变图变换器，它整合不可约表示和一个等变图注意力机制，以在 3D 原子级图上预测量子性质，具有竞争力的准确性和更高的效率。

ABSTRACT

Despite their widespread success in various domains, Transformer networks have yet to perform well across datasets in the domain of 3D atomistic graphs such as molecules even when 3D-related inductive biases like translational invariance and rotational equivariance are considered. In this paper, we demonstrate that Transformers can generalize well to 3D atomistic graphs and present Equiformer, a graph neural network leveraging the strength of Transformer architectures and incorporating SE(3)/E(3)-equivariant features based on irreducible representations (irreps). First, we propose a simple and effective architecture by only replacing original operations in Transformers with their equivariant counterparts and including tensor products. Using equivariant operations enables encoding equivariant information in channels of irreps features without complicating graph structures. With minimal modifications to Transformers, this architecture has already achieved strong empirical results. Second, we propose a novel attention mechanism called equivariant graph attention, which improves upon typical attention in Transformers through replacing dot product attention with multi-layer perceptron attention and including non-linear message passing. With these two innovations, Equiformer achieves competitive results to previous models on QM9, MD17 and OC20 datasets.

研究动机与目标

• 通过融入 3D 几何归纳偏置，激发并使 Transformer 能泛化到 3D 原子级图。
• 通过用等变对应操作替换 Transformer 操作并加入张量积来开发 Equiformer。
• 引入一个等变图注意力机制，以结合内容信息和几何信息。
• 在 QM9、MD17 和 OC20 上进行评估，以证明与最先进模型相比的竞争性能。

提出的方法

• 用 SE(3)/E(3)-等变运算替换 Transformer 组件，包括在不可约表示特征上的张量积。
• 使用深度分解张量积以有效耦合类型-L向量并控制参数增长。
• 采用新颖的等变图注意力，在标量（type-0）特征上使用 MLPL 注意力，在高阶特征上进行非线性信息传递。
• 使用原子和边-度嵌入对输入进行嵌入，并结合径向基和径向函数以实现基于距离的权重条件化。
• 对非标量不可约表示应用门控型非线性激活以保持等变性，并采用深度分解方法来控制计算复杂性。

实验结果

研究问题

• RQ1 当配备 SE(3)/E(3)-等变特征时，Transformer 能否有效地泛化到 3D 原子级图？
• RQ2 所提出的等变图注意力在 3D 图上是否比标准点积注意力更有效？
• RQ3 对 Transformer 进行最小修改、采用等变运算，是否能在 QM9、MD17 和 OC20 上匹配或超过之前的 SE(3)/E(3)-等变架构？
• RQ4 在 Equiformer 中使用更高阶的不可约表示（Lmax）时，训练时间和可扩展性有哪些权衡？

主要发现

• Lmax=2-3 的 Equiformer 在 QM9、MD17 和 OC20 数据集上取得与先前模型相比具有竞争力的结果。
• 在 QM9 上，Equiformer 在 12 个回归任务中的总体结果优于包括 NequIP 和 TorchMD-NET 在内的若干基线。
• 在 MD17 上，具有更高 Lmax 的 Equiformer 在多种分子上的能量和力的 MAE 有所提升，相较于 TorchMD-NET 和 NequIP 等替代方案。
• 在 OC20 上，当使用 IS2RE 数据训练时，Equiformer 在 IS2RE 性能上优于最先进的 SEGNN 和 Graphormer，并且相比于先前的模型，训练时间可降低 2.3x 到 15.5x。
• Equiformer 的等变图注意力（基于 MLP 的注意力，带有非线性信息传递）比点积注意力更具表达力，并提升等变 Transformer 的性能。
• 通常提高 Lmax 会改善 OC20 和 MD17 的结果，但在某些情况下可能导致过拟合。