[論文レビュー] Supervised Pretraining for Molecular Force Fields and Properties Prediction

Xiang Gao, Weihao Gao+4 more|arXiv (Cornell University)|Nov 23, 2022

Computational Drug Discovery Methods被引用数 5

ひとこと要約

本論文では、3次元幾何学的構造と原子電荷を入力とし、分子エネルギーを教師ラベルとする8600万個の分子を用いて、グラフニューラルネットワークの教師あり事前学習を提案する。この手法は、転移学習を通じて、分子性質および力場予測タスクにおける下流性能を顕著に向上させる。一方、線形プロービングにより、構造関連タスクに明示的な教師信号を用いなくても、モデルが豊富な構造的表現を学習していることが示された。

ABSTRACT

Machine learning approaches have become popular for molecular modeling tasks, including molecular force fields and properties prediction. Traditional supervised learning methods suffer from scarcity of labeled data for particular tasks, motivating the use of large-scale dataset for other relevant tasks. We propose to pretrain neural networks on a dataset of 86 millions of molecules with atom charges and 3D geometries as inputs and molecular energies as labels. Experiments show that, compared to training from scratch, fine-tuning the pretrained model can significantly improve the performance for seven molecular property prediction tasks and two force field tasks. We also demonstrate that the learned representations from the pretrained model contain adequate information about molecular structures, by showing that linear probing of the representations can predict many molecular information including atom types, interatomic distances, class of molecular scaffolds, and existence of molecular fragments. Our results show that supervised pretraining is a promising research direction in molecular modeling

研究の動機と目的

• 事前学習のための大規模ラベル付きデータセットを活用することで、分子機械学習におけるデータ不足問題に対処する。
• 事前学習フェーズからの転移学習を通じて、分子性質および力場予測タスクにおける一般化性能を向上させる。
• 分子エネルギーに対する教師あり事前学習が、さまざまな下流タスクに有用な構造的情報を暗黙的に学習できるかどうかを検証する。
• 物理的に根拠のある正確なラベル（分子エネルギー）を用いた事前学習が、ランダムまたは自己教師あり事前学習よりも優れた下流性能をもたらすことを示す。

提案手法

• 3次元座標と原子電荷を入力とし、分子エネルギーを教師ラベルとする8600万個の分子のデータセット上で、メッセージパッシング型グラフニューラルネットワーク（EGNN）を事前学習する。
• 標準的なEGNNアーキテクチャを用い、アテンションベースのメッセージパッシングにより分子グラフを符号化し、最終的なリードアウト層でエネルギーを予測する。
• 事前学習中に、エネルギーの原子座標に関する勾配のL2ノルムを最小化する力の正則化損失を適用し、滑らかな力の予測を促進する。
• 限られたラベル付きデータを用いて、事前学習済みモデルを下流の分子性質および力場予測タスクで微調整する。
• 学習済み表現の能力を評価するために、線形プロービングを実施し、原子種別、原子間距離、分子スケルトン、機能的群の予測を行う。
• すべての線形プロービングタスクに負の対数尤度損失を用い、学習済み表現の質を評価する。

実験結果

リサーチクエスチョン

• RQ1 分子エネルギーに対する教師あり事前学習は、分子性質および力場予測タスクにおける下流性能を向上させることができるか？
• RQ2 物理的に根拠のある大規模ラベル（分子エネルギー）を用いた事前学習は、自己教師ありまたはタスクに無関係な事前学習よりも一般化性能が優れているか？
• RQ3 エネルギー事前学習から得られる表現が、原子種別、原子間距離、機能的群といった分子構造的情報のどの程度を符号化しているか？
• RQ4 事前学習済みモデルは、未観測の分子スケルトンに対してゼロショット予測を一般化できるか？

主な発見

• 事前学習済みモデルの微調整は、7つの分子性質予測タスクおよび2つの力場予測タスクにおいて、初期化から学習を開始する場合と比較して、性能を顕著に向上させる。
• 線形プロービングにより、事前学習済みモデルの表現が、ランダム重みを持つモデルと比較して、原子種別、原子間距離、分子スケルトン、および85種類の機能的群を著しく低い誤差で予測できることを示した。
• 深い層を通じて原子種別情報が保持されていることから、効果的なメッセージパッシングと表現学習が行われていることが示された。
• 力の正則化技術により、力場予測タスクにおける一般化性能が向上し、勾配が正確に必要となるタスクに適したモデルとなった。
• 事前学習済みモデルは、いくつかの下流タスクで新たなSOTA（最先端）性能を達成しており、未観測の分子スケルトンに対するゼロショット一般化テストでも成功を収めた。
• 結果から、自己教師あり事前学習の代替として、分子エネルギーに対する教師あり事前学習が有望であることが示された。これは、より優れた下流性能をもたらすとともに、構造的情報を暗黙的に捉え込んでいる。