[論文レビュー] 3D Infomax improves GNNs for Molecular Property Prediction

H. Stärk, Dominique Beaini+5 more|arXiv (Cornell University)|Oct 8, 2021

Machine Learning in Materials Science参考文献 53被引用数 85

ひとこと要約

本論文は 3D Infomax を提案する。自己教師付き事前学習法で、2D GNN に3D分子ジオメトリについて暗黙裡に推論させ、下流の量子特性予測を向上させ、データセット間の一般化を可能にする。

ABSTRACT

Molecular property prediction is one of the fastest-growing applications of deep learning with critical real-world impacts. Including 3D molecular structure as input to learned models improves their performance for many molecular tasks. However, this information is infeasible to compute at the scale required by several real-world applications. We propose pre-training a model to reason about the geometry of molecules given only their 2D molecular graphs. Using methods from self-supervised learning, we maximize the mutual information between 3D summary vectors and the representations of a Graph Neural Network (GNN) such that they contain latent 3D information. During fine-tuning on molecules with unknown geometry, the GNN still generates implicit 3D information and can use it to improve downstream tasks. We show that 3D pre-training provides significant improvements for a wide range of properties, such as a 22% average MAE reduction on eight quantum mechanical properties. Moreover, the learned representations can be effectively transferred between datasets in different molecular spaces.

研究の動機と目的

• 明示的な3D構造が利用できない大規模な分子特性予測を動機付ける。
• 2Dグラフから潜在的な3D情報を2Dグラフからエンコードするような事前学習スキームを開発する。
• 学習された3D対応表現が下流の量子特性予測を改善することを示す。
• 3D Infomax表現が異なる分子空間やデータセットに跨って一般化することを示す。
• 事前学習中に複数のコンフォーマーを使用する利点を探り、性能をさらに向上させる。

提案手法

• 2D分子グラフ上で2D GNN f^aを事前学習し、コンフォーマーから導出される3D表現 z^b との相互情報を最大化して z^a を生成する。
• 同一分子に対して z^a と z^b を整列させ、非対応ペアを排除するNT-Xent型のコントラスト損失を用いる。
• 分子のすべてのコンフォーマーを正例として扱い、コントラスト損失(L^multi3D)を計算することで、複数のコンフォーマー(c)を組み込む。
• 3Dネットワークは3D座標を点群として処理し、距離ベースのMPNNと正弦波特徴写像 γ(d_uv) によって SE(3)-不変の z^b に対し、対となる距離をエンコードする。
• 下流タスクに対して2D入力のみを用いて2Dネットワークをファインチューニングし、暗黙的に学習された3D情報を活用して予測を向上させる。

実験結果

リサーチクエスチョン

• RQ1 2D GNNは、3Dコンフォーマーを用いた対比的な事前学習を通じて、2D分子グラフから潜在的な3D情報をエンコードすることを学べるのか？
• RQ2 3D Infomax事前学習モデルは、2Dのみの事前学習ベースラインと比較して、QM9 などの量子特性予測を改善するのか？
• RQ3 学習された表現は、サイズや組成が異なる分子空間間で転用可能か？
• RQ4 事前学習で複数のコンフォーマーを利用することで、下流の特性に追加の利得が得られるか？

主な発見

• 3D Infomaxは量子特性予測を大幅に改善し、QM9プロパティで平均MAEを約22%削減するなどの効果を達成した。
• 事前学習済み表現は、異なる分子空間に対して一般化し、小分子から薬物様分子への転移やその逆でも利益を示す。
• 事前学習における複数のコンフォーマーの利用は追加の性能向上をもたらすが、3つ以上のコンフォーマーでの追加利得は次第に小さくなる。
• 3D Infomaxは、GraphCLや距離予測前学習といった従来のSSLベースラインを、評価データセット全体で上回る。
• 本手法は負の転移を示さず、従来のいくつかの事前学習アプローチとは異なる。