[論文レビュー] Energy-Motivated Equivariant Pretraining for 3D Molecular Graphs

Rui Jiao, Jiaqi Han+3 more|arXiv (Cornell University)|Jul 18, 2022

Machine Learning in Materials Science被引用数 5

ひとこと要約

本稿では、E(3)-不変メッセージパッシングと物理的インスピレーションを受ける自己教師付き事前学習タスク——リーマン・ガウスノイズ除去による等価力予測とノイズスケール予測——を活用する3D分子グラフ向けの自己教師付き事前学習フレームワーク、3D-EMGPを提案する。この手法は、MD17およびQM9ベンチマークで最先端の性能を達成し、3D座標の増強を施した2D事前学習ベースラインでさえも上回っている。

ABSTRACT

Pretraining molecular representation models without labels is fundamental to various applications. Conventional methods mainly process 2D molecular graphs and focus solely on 2D tasks, making their pretrained models incapable of characterizing 3D geometry and thus defective for downstream 3D tasks. In this work, we tackle 3D molecular pretraining in a complete and novel sense. In particular, we first propose to adopt an equivariant energy-based model as the backbone for pretraining, which enjoys the merits of fulfilling the symmetry of 3D space. Then we develop a node-level pretraining loss for force prediction, where we further exploit the Riemann-Gaussian distribution to ensure the loss to be E(3)-invariant, enabling more robustness. Moreover, a graph-level noise scale prediction task is also leveraged to further promote the eventual performance. We evaluate our model pretrained from a large-scale 3D dataset GEOM-QM9 on two challenging 3D benchmarks: MD17 and QM9. Experimental results demonstrate the efficacy of our method against current state-of-the-art pretraining approaches, and verify the validity of our design for each proposed component.

研究の動機と目的

• 分子動力学やエネルギー予測などのタスクに不可欠な3D幾何的情報を捉えられていない既存の2D分子事前学習手法の限界を是正すること。
• E(3)-不変モデルと3Dに適した事前学習目的を用いて、3D分子幾何の完全な活用を図る自己教師付き事前学習フレームワークの開発。
• 3D空間変換に対してロバストで一般化性の高い性能を向上させるために、事前学習損失における対称性不変性（E(3)-不変性）を保証すること。
• 複数の3Dバックボーンとベンチマークを対象としたアブレーションおよび一般化性の検証を通じて、提案された構成要素の有効性を評価すること。

提案手法

• 3D回転、平行移動、反転に対して不変である原子レベルの力を予測するE(3)-不変メッセージパッシングバックボーンを採用。これは、等価GNN（例：EGNN）に基づく。
• 力予測に基づくノードレベルの事前学習タスクを導入。E(3)-不変性を維持するため、リーマン・ガウス分布を用いたノイズ除去損失に再定式化。
• ノイズスケールの予測を対象としたグラフレベルの代理タスクを設計。これにより、ロバストネスと一般化性が向上する。
• 2つの事前学習目的を統合し、3D-EMGPと呼ぶ単一のフレームワークを構築。大規模なGEOM-QM9データセット（3Dコンformationを含む）を用いて学習。
• 微分可能なエネルギーに基づくモデリングアプローチを採用：ノード表現のプールからエネルギーを予測し、力をエネルギーの勾配として計算することで、グローバルな整合性を向上。
• 微調整により、エネルギー可視化およびエネルギー分布評価用のエネルギー投影ヘッドを追加し、下流の3Dタスクに事前学習モデルを適用。

実験結果

リサーチクエスチョン

• RQ1 3D分子グラフ向けの自己教師付き事前学習フレームワークは、既存の2D事前学習手法を3Dベンチマークタスクで上回ることができるか？
• RQ2 リーマン・ガウスノイズ除去を介してE(3)-不変性を強制することで、モデルのロバストネスと性能にどのような影響を与えるか？
• RQ3 ノードレベル（力予測）とグラフレベル（ノイズスケール予測）の事前学習タスクを組み合わせることで、最終的な性能にどの程度寄与するか？
• RQ4 エネルギーに基づくモデリング戦略（予測エネルギーの勾配としての力の計算）は、直接的な力予測よりも優れた性能を発揮するか？
• RQ5 提案された事前学習フレームワークは、異なる3D GNNバックボーンに一般化可能か？

主な発見

• 3D-EMGPはMD17力予測ベンチマークでテストMAE 0.0876を達成し、2番目に良いベースライン（0.0968）を顕著に上回り、3D座標の増強を施した2D事前学習手法ですら上回っている。
• アブレーションスタディでは、リーマン・ガウスノイズ除去スキームを削除するとMAEが0.0912に上昇し、E(3)-不変性とロバストネスを維持する上でその重要性が示された。
• グラフレベルのノイズスケール予測タスクは有意に寄与しており、削除するとMAEが0.0876から0.0905に上昇し、一般化性向上の価値が裏付けられた。
• エネルギーに基づく力予測戦略（エネルギーの勾配としての力計算）は、直接的な力予測（0.0914）よりも低いMAE（0.0876）を達成し、グローバルな幾何的パターンをよりよく捉えていることが確認された。
• 異なるバックボーンにわたる一般化性は良好で、SchNetおよびTorchMD-ETに3D-EMGPを適用した際も一貫して性能向上が観察された。
• エネルギー分布の可視化では、3D-EMGPはコンformationの周囲で滑らかで準安定なエネルギー表面を生成している一方、2D事前学習手法は粗いエネルギー分布を示しており、物理的エネルギー分布との整合性が優れていることが示された。