[論文レビュー] A new perspective on building efficient and expressive 3D equivariant graph neural networks

Weitao Du, Yuanqi Du+6 more|arXiv (Cornell University)|Apr 7, 2023

Machine Learning in Materials Science被引用数 7

ひとこと要約

本稿は、3次元同相性の階層を介した局所的からグローバルな幾何的表現力の分析を通じて、効率的で表現力のある3次元等長グラフニューラルネットワーク（GNN）を設計するための新しい理論的枠組みを提示する。LSE（局所的サブ構造符号化）とFTE（フレーム遷移符号化）を統合したLEFTNetを提案し、分子性質予測タスクにおいて最先端の性能を達成し、スカラー性質では平均で20%、ベクトル性質では16%の向上を達成した。

ABSTRACT

Geometric deep learning enables the encoding of physical symmetries in modeling 3D objects. Despite rapid progress in encoding 3D symmetries into Graph Neural Networks (GNNs), a comprehensive evaluation of the expressiveness of these networks through a local-to-global analysis lacks today. In this paper, we propose a local hierarchy of 3D isomorphism to evaluate the expressive power of equivariant GNNs and investigate the process of representing global geometric information from local patches. Our work leads to two crucial modules for designing expressive and efficient geometric GNNs; namely local substructure encoding (LSE) and frame transition encoding (FTE). To demonstrate the applicability of our theory, we propose LEFTNet which effectively implements these modules and achieves state-of-the-art performance on both scalar-valued and vector-valued molecular property prediction tasks. We further point out the design space for future developments of equivariant graph neural networks. Our codes are available at \url{https://github.com/yuanqidu/LeftNet}.

研究の動機と目的

• 局所的からグローバルな幾何的同相性の分析を通じて、3次元等長GNNの表現力の評価のための理論的枠組みを確立すること。
• 表現力と計算効率の両立を可能にする、LSEおよびFTEという主要なアーキテクチャ的要素を同定すること。
• 不変な幾何的メッセージと等長の集約を組み合わせることで、グローバルな3次元構造を効果的に捉えることができることを示すこと。
• 将来の3次元幾何的深層学習モデルのためのモジュラー設計のレシピを提供すること。

提案手法

• GNNが幾何的構造を捉える表現力を分析するため、3次元同相性の階層（木構造、三角形、部分グラフ）を提案する。
• 微分幾何学からのスカラー化技術を用いて、局所的な3次元幾何的パターンを不変に実現するメカニズムとして、局所的サブ構造符号化（LSE）を導入する。
• メッセージ伝達中の局所フレームの変化をモデル化するため、フレーム遷移符号化（FTE）を提案し、グローバルな幾何的表現を可能にする。
• SE(3)対称性を保持するため、等長テンソル更新とLSE・FTEを統合したモジュラーなアーキテクチャであるLEFTNetを設計する。
• 高次体積オーダーのポテンシャルモデルにインspiredされた、テンソル積または和を用いたフレームベースのメッセージ集約メカニズムを採用する。
• フレームバンドルにおけるスカラー化を適用することで、局所的幾何的特徴の不変符号化を可能にし、LSEにとって不可欠な要因となる。

実験結果

リサーチクエスチョン

• RQ1 不変スカラーは、3次元幾何的パターンを符号化する際に、どの程度の表現力を有するのか？
• RQ2 グローバルな3次元幾何を捉えるために、メッセージ伝達における等長性は本当に必要なのか？
• RQ3 どのようにして、対称性を保ったまま、局所的幾何的情報を体系的に集約し、グローバルな3次元構造を表現できるのか？
• RQ4 3次元等長GNNにおいて、高い表現力と計算効率を両立させるために、どのようなアーキテクチャ的要素が必要なのか？

主な発見

• LEFTNetはQM9データセットにおいて最先端の性能を達成し、12の分子性質予測タスクの平均で、前回のSOTA結果を20%向上させた。
• rMD17データセットでは、力重み付き損失が100の条件下で、8つの分子のうち7つですべてのベースラインを上回り、平均で16%の向上を達成した。
• QM9データセットの両方のデータ分割において、12の性質のうち10つで2位または最良の性能を記録し、堅牢性と一般化性能を示した。
• アブレーションスタディの結果、LSEおよびFTEの両方が性能向上に不可欠であり、それぞれが最終的な結果に顕著な寄与をしていることが確認された。
• このフレームワークは分子に限らず、接空間上のフレーム定義を用いることで、多様体上の幾何的メッシュグラフへも拡張可能である。