[論文レビュー] Unified 2D and 3D Pre-Training of Molecular Representations

J. Zhu, Yingce Xia+6 more|arXiv (Cornell University)|Jul 14, 2022

Machine Learning in Materials Science被引用数 5

ひとこと要約

本論文は、グラフニューラルネットワークを用いて原子座標と原子間距離を統合的に符号化することで、分子表現学習のための統合的2次元（2D）および3次元（3D）事前学習フレームワークを提案する。モデルは11の分子性質予測タスクのうち10つで最先端性能を達成し、3Dコンformation生成において顕著な改善を示し、2Dのみのタスクでは平均8.3%の精度向上を達成し、コンformation生成ベンチマークでも優れた結果を示した。

ABSTRACT

Molecular representation learning has attracted much attention recently. A molecule can be viewed as a 2D graph with nodes/atoms connected by edges/bonds, and can also be represented by a 3D conformation with 3-dimensional coordinates of all atoms. We note that most previous work handles 2D and 3D information separately, while jointly leveraging these two sources may foster a more informative representation. In this work, we explore this appealing idea and propose a new representation learning method based on a unified 2D and 3D pre-training. Atom coordinates and interatomic distances are encoded and then fused with atomic representations through graph neural networks. The model is pre-trained on three tasks: reconstruction of masked atoms and coordinates, 3D conformation generation conditioned on 2D graph, and 2D graph generation conditioned on 3D conformation. We evaluate our method on 11 downstream molecular property prediction tasks: 7 with 2D information only and 4 with both 2D and 3D information. Our method achieves state-of-the-art results on 10 tasks, and the average improvement on 2D-only tasks is 8.3%. Our method also achieves significant improvement on two 3D conformation generation tasks.

研究の動機と目的

• 既存手法が2D分子グラフと3Dコンformationを別々に扱うという制限を解決すること。これは、両者が相補的であるにもかかわらずである。
• 2Dおよび3D分子表現を1つのモデルですべて統合的に学習する統合的ディープラーニングフレームワークの開発。
• 複数の自己教師的タスクで事前学習することで、分子性質予測および3Dコンformation生成の性能を向上させること。
• モデルの頑健性と一般化性能を向上させるために、置換不変性および回転平行移動不変性を学習目的に組み込むこと。

提案手法

• 原子座標と原子間距離を、最初に全結合ネットワークで符号化し、その後原子表現とグラフニューラルネットワークを介して統合する。
• 3つの事前学習タスクを採用する：マスキングされた原子および座標の再構築、2Dグラフからの3Dコンformation生成、3Dコンformationからの2Dグラフ生成。
• 対称的なサブ構造を扱うために、置換不変損失を用い、原子の入れ替えが学習安定性に影響しないようにする。
• 3Dコンformation生成に回転平行移動不変損失を適用することで、空間変換下での等長性を保証する。
• PCQM4Mv2データセットを用いて事前学習を行い、分子の2Dグラフと3Dコンformationの両方を提供する。
• 下流タスク（分子性質予測およびコンformation生成）において、トランスファー学習を用いて微調整する。

実験結果

リサーチクエスチョン

• RQ1 2Dおよび3D分子表現の統合的事前学習は、下流の分子性質予測性能を向上させることができるか？
• RQ2 損失関数に置換不変性を組み込むことで、対称的な分子サブ構造におけるモデルの頑健性と性能はどのように向上するか？
• RQ3 2Dおよび3Dの統合的事前学習は、2Dまたは3Dの監視のみを用いる手法と比較して、3Dコンformation生成をどの程度向上させるか？
• RQ4 各事前学習目的（再構築、生成）が最終的なモデル性能に果たす寄与度は何か？
• RQ5 サイクル整合性とマスキングされた原子再構築を活用することで、3D構造を持たない分子に対しても統合モデルは一般化可能か？

主な発見

• 提案手法は、11の分子性質予測タスクのうち10つで最先端性能を達成し、2Dのみのタスクでは平均8.3%の向上を示した。
• GEOM-QM9およびGEOM-Drugsデータセットにおいて、すべてのベースラインを上回る性能を示し、コンフォーメーションカバレッジ（COV）とマッチングスコア（MAT）が最高を記録した。
• アブレーションスタディにより、3つの事前学習目的を併用することで最良の性能が得られ、単一目的学習と比較してMATスコアが最大15%向上した。
• 置換不変損失の導入により、対称的サブ構造における性能が顕著に向上し、アブレーション実験では予測誤差が平均20%削減された。
• 初期座標をランダムに設定し、マスキング再構築とサイクル整合性を活用することで、3D構造を持たない分子に対しても良好な一般化性能を示した。
• 統合的事前学習フレームワークにより、3Dコンformation生成に対する強力なゼロショット転移が可能となり、モデルの頑健性と一般化能力が裏付けられた。