[論文レビュー] GemNet: Universal Directional Graph Neural Networks for Molecules

Johannes Gasteiger, Florian Becker+1 more|arXiv (Cornell University)|Jun 2, 2021

Computational Drug Discovery Methods被引用数 127

ひとこと要約

GemNetは、回転不変/協変な分子予測の球面表現の普遍性を確立し、方向性エッジ埋め込みを用いた幾何学的な2ホップメッセージパッシングを導入し、複数のMDデータセットで力の予測で最先端を達成します。

ABSTRACT

Effectively predicting molecular interactions has the potential to accelerate molecular dynamics by multiple orders of magnitude and thus revolutionize chemical simulations. Graph neural networks (GNNs) have recently shown great successes for this task, overtaking classical methods based on fixed molecular kernels. However, they still appear very limited from a theoretical perspective, since regular GNNs cannot distinguish certain types of graphs. In this work we close this gap between theory and practice. We show that GNNs with spherical representations are indeed universal approximators for predictions that are invariant to translation, and equivariant to permutation and rotation. We then discretize such GNNs via directed edge embeddings and two-hop message passing, and incorporate multiple structural improvements to arrive at the geometric message passing neural network (GemNet). We demonstrate the benefits of the proposed changes in multiple ablation studies. GemNet outperforms previous models on the COLL, MD17, and OC20 datasets by 34%, 41%, and 20%, respectively, and performs especially well on the most challenging molecules. Our implementation is available online.

研究の動機と目的

• 分子グラフに対する通常のGNNの表現力のギャップを解消し、回転不変および協変な予測を可能にする。
• 回転不変および協変なタスクに対する球面表現の普遍性を証明する。
• 幾何情報を備えた実用的な、方向性エッジベースの2ホップメッセージパッシングフレームワークを開発し、精度とサンプル効率を向上させる。

提案手法

• SO(3)不変および置換対称性を持つ関数に対する球面表現の普遍性を証明する。
• 球面表現を有向エッジ埋め込みに離散化し、エッジ間の2ホップメッセージパッシングをモデル化する。
• 球面フーリエ-ベessel基底を用いた、原子間の方向と角度（φとθを含む）を用いた幾何的メッセージパッシングを導入する。
• 対称的なメッセージパッシングと効率的な二項演算層を組み込み、安定性と計算効率を向上させる。
• 標準化の代わりに事前に決定されたスケーリング因子で活性化分散を安定化させる。
• GemNetアーキテクチャを、2ホップ幾何メッセージパッシング、1ホップ幾何メッセージパッシング、原子自己相互作用という3つの相互作用形で開発する。

実験結果

リサーチクエスチョン

• RQ1 分子GNNにおいて球面表現は回転不変および協変な予測に対して普遍性を達成できるか？
• RQ2 方向性（エッジベース）の表現をどのように離散化して、実用的で表現力のある分子GNNを作るか？
• RQ3 エネルギーと力の予測精度を最適化するために、どのようなアーキテクチャ革新（対称的メッセージ、バイリニア層、分散安定化）が最も効果的か？
• RQ4 GemNetはCOLL、MD17、OC20の力とエネルギー予測で最先端モデルとどう比べられるか？

主な発見

• GemNetは、COLL、MD17、OC20のデータセットにおいて前モデルを著しく上回り、エラーを大幅に低減する。
• 方向性エッジ埋め込みを用いた2ホップ幾何メッセージパッシングは、特に非平面分子の難易度が高いケースで性能を大幅に向上させる。
• 対称的なメッセージパッシング、バイリニア層、および活性化分散の安定化は、アブレーション研究で各要素の利益が確認され、精度に有意に寄与する。
• 直接的な力の予測はトレーニングと推論の速度を高める可能性があるが、データセットとタスクによってMAEにトレードオフが生じる。
• GemNet-TとGemNet-Qの varianteは計算コストと精度の間で異なるトレードオフを提供し、より複雑な組み合わせデータセットでGemNet-Qが大きな利得を提供する。