[論文レビュー] 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を導入します。これは、 steerableベクトル空間とClebsch-Gordan積を用いて幾何学的・物理量をメッセージと更新ステップの両方に注入する一般的なE(3)同変性GNNフレームワークで、いくつかの物理/化学ベンチマークにおいて最先端または競合的な結果を達成します。

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.

研究の動機と目的

• ノード属性とエッジ属性をスカラー以外（ベクトル、テンソル）を含む共変情報として一般化し、E(3)同変性を保つ。
• メッセージ伝播に幾何学的・物理的量を注入するためのsteerable MLPとsteerable活性化関数を導入する。
• 非線形の畳み込みとしての同変性GNNの統一的な見方を提供し、性能を向上させる重要な要素を特定する。
• N体、QM9、OC20データセットに対する有効性を実証し、広範なアブレーションを行う。

提案手法

• Steerable vector spaces V_L = V0 ⊕ V1 ⊕ ... ⊕ VL と、幾何情報をClebsch-Gordan積を介して条件づけた線形写像を持つsteerable MLPを用いる。
• ノード/エッジ属性をSteerableベクトルとして表現する（相対位置の球面調和関数埋め込みを含む）。
• Steerable edge属性と二乗距離に依存するメッセージと、ノード属性により更新を導く非線形のsteerableメッセージ伝播を定義する。
• 球面調和関数の埋め込みを取り入れて、ベクトルをCGテンソル積の型-l steeredベクトルへ変換する。
• 幾何学的・物理的量を有効に処理する非線形で同変な処理を可能にするsteerableノード属性によるsteerable活性化関数を提案する。
• SEGNNsを非線形畳み込みに関連づけ、線形の同変/畳み込みベースラインと比較する。

実験結果

リサーチクエスチョン

• RQ1 GNNにおけるノードおよびエッジ属性をスカラー以外の共変量（ベクトル、テンソル）として拡張し、E(3)同変性を保つことは可能か。
• RQ2 非線形でsteerableなメッセージ伝播層は、物理/化学タスクにおいて線形・不変/steerableなメッセージ伝播より性能を向上させるか。
• RQ3 幾何学的・物理的量（速度・力など）は、分子・N体問題におけるSEGNNsの性能にどのような影響を与えるか。
• RQ4 このフレームワークにおける効果的なE(3)同変性GNNsの本質的要素は何か。

主な発見

• l_f = 1 および l_a = 1 のSEGNNsは、ベンチマーク全体で最先端または競合的な結果を達成し、いくつかのベースラインを上回る。
• 幾何的量と物理的量の両方を含むG+Pの組み合わせは、幾何学のみ（G）または不変手法より大幅に性能を向上させる。
• 非線形のsteerable畳み込みは、線形のsteerableメッセージ伝播および不変メッセージ伝播よりも有利な点を提供する。
• 相対位置の球面調和関数埋め込みを用いると、CGベースのメッセージ変換が効果的に機能する。
• SEGNNsはN-bodyタスク、QM9、OC20 IS2REで強力な結果を示し、回転・同変性の整合性を維持する。