[论文解读] Geometric and Physical Quantities Improve E(3) Equivariant Message Passing

J. Brandstetter, Rob Hesselink+3 more|arXiv (Cornell University)|Oct 6, 2021

Advanced Graph Neural Networks被引用 41

一句话总结

该论文提出 SEGNNs，一种通用的 E(3) 等变图神经网络框架，使用可导向向量空间和 Clebsch-Gordan 乘积在消息与更新步骤中注入几何和物理量，在若干物理/化学基准上实现最先进或具有竞争力的结果。

ABSTRACT

Including covariant information, such as position, force, velocity or spin is important in many tasks in computational physics and chemistry. We introduce Steerable E(3) Equivariant Graph Neural Networks (SEGNNs) that generalise equivariant graph networks, such that node and edge attributes are not restricted to invariant scalars, but can contain covariant information, such as vectors or tensors. This model, composed of steerable MLPs, is able to incorporate geometric and physical information in both the message and update functions. Through the definition of steerable node attributes, the MLPs provide a new class of activation functions for general use with steerable feature fields. We discuss ours and related work through the lens of equivariant non-linear convolutions, which further allows us to pin-point the successful components of SEGNNs: non-linear message aggregation improves upon classic linear (steerable) point convolutions; steerable messages improve upon recent equivariant graph networks that send invariant messages. We demonstrate the effectiveness of our method on several tasks in computational physics and chemistry and provide extensive ablation studies.

研究动机与目标

• 将等变 GNN 泛化以允许节点和边的属性不仅限于标量，而包含协变信息（向量、张量）。
• 引入可导向的 MLP（steerable MLPs）和可导向激活函数（steerable activation functions），将几何/物理量注入消息传递。
• 提供对等变 GNN 的非线性卷积统一视角，并识别提升性能的关键组成部分。
• 展示有效性并对 N-body、QM9 和 OC20 数据集进行广泛的消融分析。

提出的方法

• 使用 steerable 向量空间 V_L = V0 ⊕ V1 ⊕ ... ⊕ VL 以及 steerable MLPs，其线性映射通过 Clebsch-Gordan 乘积受几何信息条件约束。
• 将节点/边属性表示为可导向向量（包括相对位置的球面调和嵌入）。
• 定义非线性可导向消息传递，消息依赖于可导向边属性和平方距离，更新由节点属性引导。
• 结合球面调和嵌入，将向量转换为 type-l 可导向向量用于 CG 张量乘积。
• 提出通过可导向的节点属性实现的可导向激活函数，使协变信息的非线性、等变处理成为可能。
• 将 SEGNN 与非线性卷积联系起来，并与线性等变/卷积基线进行比较。

实验结果

研究问题

• RQ1 在 GNN 中，节点和边的属性能否扩展为除标量之外的协变量（向量、张量），并保持 E(3) 等变性？
• RQ2 非线性、可导向的消息传递层是否比线性、不变/可导向的消息传递在物理/化学任务上具有更高的性能？
• RQ3 几何和物理量（如速度、力）如何影响 SEGNN 在分子和 N-body 问题上的性能？
• RQ4 在该框架下，构成有效 E(3) 等变 GNN 的关键组成部分是什么？

主要发现

• SEGNN 在 l_f = 1 且 l_a = 1 的设置下，在各基准测试中达到最先进或具竞争力的结果，优于若干基线。
• 同时考虑几何和物理量（G+P）显著提升性能，相较仅几何（G）或不变方法。
• 非线性可导向卷积在比线性可导向消息传递和不变消息传递时具有优势。
• 对相对位置使用球面调和嵌入使基于 CG 的消息变换能够有效实现可导向处理。
• SEGNN 在 N-body、QM9 和 OC20 IS2RE 上表现强劲，同时保持旋转/等变性的一致性。