[論文レビュー] Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs

Saro Passaro, C. Lawrence Zitnick|arXiv (Cornell University)|Feb 7, 2023

Advanced Neural Network Applications被引用数 27

ひとこと要約

本論文は、主要軸をエッジの方向に合わせることで SO(3) 等変畳み込みを効率的な SO(2) 演算に変換し、テンソル積の計算量を O(L^6) から O(L^3) に削減する方法を提示し、Equivariant Spherical Channel Network (eSCN) を導入して OC-20/OC-22 データセットで最先端の結果を達成する。

ABSTRACT

Graph neural networks that model 3D data, such as point clouds or atoms, are typically desired to be $SO(3)$ equivariant, i.e., equivariant to 3D rotations. Unfortunately equivariant convolutions, which are a fundamental operation for equivariant networks, increase significantly in computational complexity as higher-order tensors are used. In this paper, we address this issue by reducing the $SO(3)$ convolutions or tensor products to mathematically equivalent convolutions in $SO(2)$ . This is accomplished by aligning the node embeddings' primary axis with the edge vectors, which sparsifies the tensor product and reduces the computational complexity from $O(L^6)$ to $O(L^3)$, where $L$ is the degree of the representation. We demonstrate the potential implications of this improvement by proposing the Equivariant Spherical Channel Network (eSCN), a graph neural network utilizing our novel approach to equivariant convolutions, which achieves state-of-the-art results on the large-scale OC-20 and OC-22 datasets.

研究の動機と目的

• SO(3) 同変グラフニューラルネットワークの計算効率を改善するために、テンソル積を SO(2) 畳み込みへ削減する。
• 膨大なコストを伴うことなく、実際に高次数の不可約表現を利用可能にする。
• 大規模な原子データセット（OC-20/OC-22）で最先端の性能を示す。
• GNNにおけるメッセージ伝搬の SO(3) と SO(2) の定式化間の原理的な結びつきを提供する。

提案手法

• 不可約表現の主軸をエッジの方向に合わせてテンソル積を疎化する。
• 回転された SO(3) 畳み込みは円調和関数上の等価な SO(2) 畳み込みとして計算できることを示す。
• 完全な Clebsch–Gordan 係数の計算を回避し複雑さを O(L^3) に削減する、計算的に効率的な定式化を導出する。
• SO(2) 定式化を利用したメッセージ伝搬のための Equivariant Spherical Channel Network (eSCN) を提案する。
• GPU に優しいリシェイピングを用いて、m チャンネルと次数にわたる畳み込みを parallel化する。

実験結果

リサーチクエスチョン

• RQ1 SO(3) 等変畳み込みを、同変性を損なうことなく効率的な SO(2) 畳み込みを用いて計算できるか。
• RQ2 軸整列により、3D データの GNN において高次数の不可約表現を実用的に使用できるか。
• RQ3 大規模なオープン触媒データセットにおける eSCN の性能は、最先端の等変 GNN と比較してどうか。

主な発見

• 軸揃え SO(2) 定式化を用いると、テンソル積の計算は O(L^6) から O(L^3) に削減される。
• 主軸をエッジ方向に合わせると Y^(l) 項が疎となり、計算が高速化され、Clebsch–Gordan 係数の必要性がなくなる。
• eSCN は OC-20 および OC-22 で最先端の結果を達成し、力予測の改善が顕著（OC-20で 9%、OC-22で 21%）、リラックス構造予測での改善は 15%。”
• 不可約表現を適切に回転させると、SO(3) の等変演算は一般化された SO(2) 畳み込みの集合として見ることができる。