[論文レビュー] Molecular Geometry Pretraining with SE(3)-Invariant Denoising Distance Matching

Shengchao Liu, Hongyu Guo+1 more|arXiv (Cornell University)|Jun 27, 2022

Machine Learning in Materials ScienceMaterials Science被引用数 17

ひとこと要約

本稿では、座標のノイズ除去を通じてポテンシャルエネルギー面をモデル化することにより、SE(3)不変のノイズ除去距離マッチングを活用して、ロバストな分子表現を学習する新しい3次元分子幾何学事前学習フレームワークであるGeoSSL-DDMを提案する。座標のノイズ除去をスコアマッチングを用いてペアワイズ距離のノイズ除去に変換することで、22の下流幾何分子予測タスクで最先端の性能を達成し、既存の2次元および3次元事前学習ベースラインを上回る。

ABSTRACT

Molecular representation pretraining is critical in various applications for drug and material discovery due to the limited number of labeled molecules, and most existing work focuses on pretraining on 2D molecular graphs. However, the power of pretraining on 3D geometric structures has been less explored. This is owing to the difficulty of finding a sufficient proxy task that can empower the pretraining to effectively extract essential features from the geometric structures. Motivated by the dynamic nature of 3D molecules, where the continuous motion of a molecule in the 3D Euclidean space forms a smooth potential energy surface, we propose GeoSSL, a 3D coordinate denoising pretraining framework to model such an energy landscape. Further by leveraging an SE(3)-invariant score matching method, we propose GeoSSL-DDM in which the coordinate denoising proxy task is effectively boiled down to denoising the pairwise atomic distances in a molecule. Our comprehensive experiments confirm the effectiveness and robustness of our proposed method.

研究の動機と目的

• 3次元分子幾何学データのための効果的な事前学習手法の不足に取り組むこと。これは、正確な分子性質予測に不可欠である。
• ポテンシャルエネルギー面（PES）における分子コンフォーマーの動的かつノイズの多い性質を捉える自己教師付き代理タスクを開発すること。
• 直接的な3次元座標のノイズ除去の課題を乗り越えるために、SE(3)不変スコアマッチングを用いて、タスクを原子間距離のノイズ除去に変換すること。
• 3次元空間における回転および並進変動に対して不変である、ロバストな分子表現学習を可能にすること。
• 多様な3次元分子予測ベンチマークにおいて、提案手法の優位性を実証的に検証すること。

提案手法

• クリーンな分子幾何学とそのノイズの加えられた摂動との間の相互情報量を最大化する一般化された幾何的自己教師付き表現学習フレームワークであるGeoSSLを提案する。
• 座標のノイズ除去を、回転および並進不変性を保証するためのペアワイズ原子距離のノイズ除去に再定式化するSE(3)不変スコアマッチング手法であるGeoSSL-DDMを導入する。
• NCSNにインspiredされたマルチレイヤーのノイズ除去スコアマッチング目的関数を採用し、複数のノイズ除去段階にわたり段階的ノイズ除去を可能にする。
• 真のデータ分布を近似し、クリーンな幾何学とノイズの加えられた幾何学との間の相互情報量を最大化するために、エネルギーに基づくモデリング（EBM）を用いる。
• グラフニューラルネットワーク（GNN）を用いて分子構造を符号化し、ノイズ除去された距離を予測する。損失関数はスコアマッチングを介して最適化される。
• 最適化を現実可能にするために、相互情報量の下界を適用し、クリーンな幾何学をノイズの加えられた入力から再構築するようにモデルを効果的に訓練する。

実験結果

リサーチクエスチョン

• RQ1 3次元分子幾何学に基づく自己教師付き事前学習フレームワークは、既存の2次元グラフベース手法を上回り、下流の分子予測タスクで優れた性能を発揮できるか？
• RQ2 3次元座標の直接的ノイズ除去よりも、ペアワイズ原子距離のノイズ除去は、3次元分子表現学習においてより効果的かつ現実的な代理タスクであるか？
• RQ3 ノイズ除去レイヤーの数が、3次元分子事前学習におけるモデル性能とロバスト性に与える影響は何か？
• RQ4 SE(3)不変スコアマッチングは、分子の滑らかなポテンシャルエネルギー面を効果的に捉えつつ、回転および並進変動に対して不変であるか？
• RQ5 提案手法は、力予測やエネルギー回帰などの多様な分子データセットおよび下流タスクに一般化可能か？

主な発見

• GeoSSL-DDMは22の下流幾何分子予測タスクで最先端の性能を達成し、2次元グラフベースおよび他の3次元手法を含む9つの事前学習ベースラインを上回った。
• MD17データセットでは、力予測において平均絶対誤差が1.637 ± 0.01にまで低下し、比較されたすべての手法の中で最低であった。
• 特定の力予測タスクにおいて、教師ありベースラインと比較して平均絶対誤差を最大20％まで低減し、自己教師付き事前学習による優れた一般化性能を示した。
• 実験の結果、ノイズ除去レイヤー数（L > 1）を増やすことでさらなる性能向上が得られ、単一レイヤー手法に比べてマルチステージノイズ除去の利点が確認された。
• GeoSSL-DDMは、評価されたすべてのタスクで並列研究[76]を上回り、特に深層ノイズ除去ネットワークを用いる場合に顕著な優位性を示した。これは、マルチレイヤーのスコアマッチングの利点を示している。
• アブレーションスタディの結果、座標の直接的ノイズ除去よりもペアワイズ距離のノイズ除去がより効果的であり、SE(3)不変性が性能とロバスト性を顕著に向上させた。