[論文レビュー] SE(3) Equivariant Graph Neural Networks with Complete Local Frames

Weitao Du, He Zhang+5 more|arXiv (Cornell University)|Oct 26, 2021

Machine Learning in BioinformaticsBiochemistry, Genetics and Molecular Biology被引用数 22

ひとこと要約

本稿では、外積を用いて構築される完全な局所フレームを用いることで、参照なしで忠実な幾何テンソル表現を可能にする、計算的に効率的なSE(3)同変グラフニューラルネットワークClofNetを提案する。高次元の幾何的量を正規直交フレームに射影することにより、多数体ダイナミクスのモデリングおよび分子コンformation生成の両面で最先端の性能を達成し、優れたサンプル効率と低い同変性誤差を示す。

ABSTRACT

Group equivariance (e.g. SE(3) equivariance) is a critical physical symmetry in science, from classical and quantum physics to computational biology. It enables robust and accurate prediction under arbitrary reference transformations. In light of this, great efforts have been put on encoding this symmetry into deep neural networks, which has been shown to improve the generalization performance and data efficiency for downstream tasks. Constructing an equivariant neural network generally brings high computational costs to ensure expressiveness. Therefore, how to better trade-off the expressiveness and computational efficiency plays a core role in the design of the equivariant deep learning models. In this paper, we propose a framework to construct SE(3) equivariant graph neural networks that can approximate the geometric quantities efficiently. Inspired by differential geometry and physics, we introduce equivariant local complete frames to graph neural networks, such that tensor information at given orders can be projected onto the frames. The local frame is constructed to form an orthonormal basis that avoids direction degeneration and ensure completeness. Since the frames are built only by cross product operations, our method is computationally efficient. We evaluate our method on two tasks: Newton mechanics modeling and equilibrium molecule conformation generation. Extensive experimental results demonstrate that our model achieves the best or competitive performance in two types of datasets.

研究の動機と目的

• 3次元多数体系のSE(3)同変ディープラーニングにおいて表現力と計算効率のバランスをとる課題に対処すること。
• 従来の同変GNNが不完全な径方向フレームに依存することで生じる方向の退化を克服すること。
• 参照なしで幾何的量（例：ベクトル、ねじり力）の損失なしの表現を同変的に実現すること。
• 任意の順序のテンソルを扱えるフレームワークを構築し、高次元の持ち上げや複雑なテンソル演算を用いずに群同変性を維持すること。
• 物理的モデリングおよび分子コンformation生成のタスクにおいて、優れた一般化性能とデータ効率を示すこと。

提案手法

• 粒子のベクトル同士の外積を用いて3次元空間における正規直交基底を形成する、完全な局所フレームを導入し、完全性を確保するとともに方向の退化を回避する。
• スカラライゼーションブロックを介して幾何的テンソル（スカラー、ベクトル、高次元）をこれらの局所フレームに射影し、メッセージパッシング用のスカラー係数に変換する。
• スカラー係数のメッセージパッシングに標準的なグラフニューラルネットワークを用い、フレームの同変性のおかげでSE(3)同変性を保持する。
• Clebsch-Gordan分解や球面調和関数に依存する手法と比較して、計算コストが低い外積演算のみでフレームを構築する。
• 外積の再帰的適用と正規直交性の維持により、高次元テンソル入力を扱えるフレーム構築法を拡張する。
• サンプリングにアニールドランジュヴィンダイナミクスを用いるスコアベースの生成モデリング設定にフレームワークを統合し、分子コンformation生成を実現する。

実験結果

リサーチクエスチョン

• RQ1 高次元の持ち上げやテンソル分解の高コストを回避できる、計算的に効率的なSE(3)同変GNNを設計できるか？
• RQ2 完全な局所フレームの使用は、径方向のみまたは不完全なフレーム法と比較して、幾何的表現の表現力と正確性を向上させるか？
• RQ3 本手法は、分布外や低データレジームの物理系において、どの程度一般化できるか？
• RQ4 ダイナミクスモデリングおよび分子コンformation生成の両タスクで、最先端の性能を達成しつつも強い同変性を維持できるか？
• RQ5 外積によるフレーム構築は、代替の同変関数空間表現と比較して、効率性と安定性に優れているか？

主な発見

• ClofNetはPOSデータセットで最高の性能を示し、内挿MSEが0.138、外挿MSEが2.502と、Radial Field（7.665）やEGNN（6.449）を大きく上回った。
• 同変性誤差が非常に低く（$\Delta_{EQ} = 4.13 \cdot 10^{-6}$）あり、SE(3)同変性の数値的保持が強く示された。
• GEOM-QM9およびGEOM-Drugsデータセットでは、視覚化された最良一致構造から、ClofNetが低RMSDで高品質な分子コンformationを生成した。
• 訓練データのたった0.4%のみを用いても、競争力ある結果を達成しており、優れたサンプル効率を示した。
• アブレーションスタディにより、ねじり力のような複雑な幾何的量をモデル化する上で、完全な局所フレームの使用が不可欠であることが確認された。
• フレーム構築に外積演算のみを用いることで、高価なテンソル積分解を回避し、フレームワークの計算的効率が裏付けられた。