[論文レビュー] Efficient and Accurate Physics-aware Multiplex Graph Neural Networks for 3D Small Molecules and Macromolecule Complexes

Shuo Zhang, Yang Liu+1 more|arXiv (Cornell University)|Jun 6, 2022

Machine Learning in Materials Science被引用数 6

ひとこと要約

本稿では、局所的および非局所的分子相互作用を別々のメッセージパッシング方式を用いてモデル化する物理知覚型マルチプレックスグラフニューラルネットワーク、PaxNetを提案する。これにより、3D小分子およびマクロ分子複合体におけるスカラーおよびベクトル的性質の予測が、効率的かつ高精度に行えるようになる。PaxNetは、最先端のベースラインと比較して、小分子では予測誤差を15%削減し、メモリ使用量を73%削減する。マクロ分子複合体では、推論時間を85%短縮し、メモリ使用量を33%削減する。

ABSTRACT

Recent advances in applying Graph Neural Networks (GNNs) to molecular science have showcased the power of learning three-dimensional (3D) structure representations with GNNs. However, most existing GNNs suffer from the limitations of insufficient modeling of diverse interactions, computational expensive operations, and ignorance of vectorial values. Here, we tackle these limitations by proposing a novel GNN model, Physics-aware Multiplex Graph Neural Network (PaxNet), to efficiently and accurately learn the representations of 3D molecules for both small organic compounds and macromolecule complexes. PaxNet separates the modeling of local and non-local interactions inspired by molecular mechanics, and reduces the expensive angle-related computations. Besides scalar properties, PaxNet can also predict vectorial properties by learning an associated vector for each atom. To evaluate the performance of PaxNet, we compare it with state-of-the-art baselines in two tasks. On small molecule dataset for predicting quantum chemical properties, PaxNet reduces the prediction error by 15% and uses 73% less memory than the best baseline. On macromolecule dataset for predicting protein-ligand binding affinities, PaxNet outperforms the best baseline while reducing the memory consumption by 33% and the inference time by 85%. Thus, PaxNet provides a universal, robust and accurate method for large-scale machine learning of molecules. Our code is available at https://github.com/zetayue/Physics-aware-Multiplex-GNN.

研究の動機と目的

• 既存のGNNが多様な分子相互作用、特に3D構造における相互作用をモデル化する点で抱える限界を解消すること。
• 局所的および非局所的相互作用のモデル化を分離することで、計算コストを低減し、高価な角度関連計算を回避すること。
• スカラー的性質に加えて、双極子モーメントや力といったベクトル的性質の正確な予測を可能にすること。
• リソース制約下でのタンパク質-リガンド系など、マクロ分子複合体の効率性とスケーラビリティを向上させること。
• 大規模な分子表現学習のための統一的で、強固かつ高精度なフレームワークの開発

提案手法

• PaxNetは、各「プレックス」が異なる種類の相互作用（例：局所的共有結合、非局所的非共有結合）を表すマルチプレックスグラフ構造を採用し、相互作用タイプごとに別々のメッセージパッシングを可能にする。
• 局所的（1-hopおよび2-hop）および非局所的（グローバル）メッセージパッシングを分離することで、より正確かつ効率的に相互作用をモデル化する物理的インスピレーションに基づく設計を導入する。
• 原子の位置から導出される幾何学的ベクトルとメッセージ重みを組み合わせ、双極子モーメントなどのベクトル的性質を直接予測する。
• 異なる相互作用タイプからの寄与を動的に重み付けるために、メッセージパッシング方式全体に学習可能なアテンション機構を適用する。
• 角計算を物理的インスパイラションに基づく近似に置き換えることで、O(Nk²)の計算複雑度を回避する修正されたメッセージパッシング方式を採用する。
• 各原子にベクトルを関連させ、メッセージパッシングによる更新を通じて、スカラーおよびベクトル的分子性質のエンドツーエンド学習を可能にする。

実験結果

リサーチクエスチョン

• RQ1 局所的および非局所的相互作用のための別々のメッセージパッシングを備えたマルチプレックスGNNフレームワークは、3D分子性質の予測精度を向上させることができるか？
• RQ2 相互作用モデル化を分離することで、すべての相互作用を均一に扱う従来のGNNと比較して、計算コストはどのように低減されるか？
• RQ3 双極子モーメントなどのベクトル的性質を、スカラーのみを予測するモデルよりも高精度に予測できるか、その程度はどの程度か？
• RQ4 PaxNetは、大規模なマクロ分子複合体において、大幅にメモリ使用量と推論時間を削減しながらも、高いパフォーマンスを維持できるか？
• RQ5 アテンション、プレックス間通信、およびハイパーホップメッセージパッシングといった個々のコンポーネントが、全体のモデルパフォーマンスに果たす貢献度はどの程度か？

主な発見

• QM9データセットでは、PaxNetが双極子モーメント予測の平均絶対誤差（MAE）を0.0108 Dまで低下させ、最良のベースライン（0.021 D）を50%以上上回った。
• QM9では、PaxNetが最良のベースラインと比較して15%低い予測誤差と73%少ないメモリ使用量を達成した。
• PDBbindデータセットでは、PaxNetが最良のベースラインと比較して33%少ないメモリ消費量と85%短い推論時間を記録しながら、優れたパフォーマンスを達成した。
• アブレーションスタディでは、アテンションモジュールを削除すると、スカラー的ターゲットでMAEが最大4%上昇し、ベクトル的ターゲットで3%上昇した。これにより、その重要性が示された。
• プレックス間通信やメッセージパッシング方式を削除すると、スカラー的ターゲットで最大28%、ベクトル的ターゲットで最大102%のパフォーマンス低下が生じ、統合的設計の必要性が確認された。
• 局所的メッセージパッシングに2ホップの幾何学的情報（θ）を活用することで、U₀では64%、μでは111%のパフォーマンス向上が達成され、より豊かな幾何学的符号化の価値が浮き彫りになった。