[論文レビュー] Highly Accurate Quantum Chemical Property Prediction with Uni-Mol+

Shuqi Lu, Zhifeng Gao+3 more|arXiv (Cornell University)|Mar 16, 2023

Machine Learning in Materials Science被引用数 6

ひとこと要約

Uni-Mol+ は、1D/2D 入力を用いて低コストの3次元分子コンformationを生成し、2トラックのTransformerモデルを用いてそれらを DFT で最適化された平衡コンformationへ段階的に改善することで、量子化学的 (QC) 特性を予測する深層学習フレームワークを提案する。PCQM4MV2 および OC20 ベンチマークにおいて最先端の性能を達成し、OC20 テストセットでは 0.3745 eV の MAE と 11.29% の EwT を達成し、従来手法を著しく上回っている。

ABSTRACT

Recent developments in deep learning have made remarkable progress in speeding up the prediction of quantum chemical (QC) properties by removing the need for expensive electronic structure calculations like density functional theory. However, previous methods learned from 1D SMILES sequences or 2D molecular graphs failed to achieve high accuracy as QC properties primarily depend on the 3D equilibrium conformations optimized by electronic structure methods, far different from the sequence-type and graph-type data. In this paper, we propose a novel approach called Uni-Mol+ to tackle this challenge. Uni-Mol+ first generates a raw 3D molecule conformation from inexpensive methods such as RDKit. Then, the raw conformation is iteratively updated to its target DFT equilibrium conformation using neural networks, and the learned conformation will be used to predict the QC properties. To effectively learn this update process towards the equilibrium conformation, we introduce a two-track Transformer model backbone and train it with the QC property prediction task. We also design a novel approach to guide the model's training process. Our extensive benchmarking results demonstrate that the proposed Uni-Mol+ significantly improves the accuracy of QC property prediction in various datasets. We have made the code and model publicly available at \url{https://github.com/dptech-corp/Uni-Mol}.

研究の動機と目的

• 1D/2D の分子表現と、正確なDFT計算に必要な3次元平衡コンformationとの間のギャップを解消すること。
• SMILES や2次元グラフから直接 QC 特性を予測する従来の深層学習モデルの限界を克服し、精度に不可欠なコンformation依存性を捉えること。
• ニューラルネットワークを用いて、RDKit などから得られる原始的な3次元コンformationを DFT で最適化された平衡コンformationへ段階的に改善する手法を開発すること。
• トレーニング時に擬似軌道サンプリング戦略を導入することで、分布シフトの影響を軽減し、分布外設定における一般化性とロバスト性を向上させること。
• DFT計算に依存するのを減らし、大スケールでの分子特性スクリーニングに適したスケーラブルでエンドツーエンドのフレームワークを確立すること。

提案手法

• RDKit などの低コストなツールを用いて、SMILES や2次元分子グラフから初期の3次元分子コンformationを生成する。
• 原子表現用とペア表現用の2トラックのTransformerバックボーンを採用し、両者に3次元空間的およびグラフの位置符号化を追加して強化する。
• 原始的なコンformationを DFT で最適化された平衡コンformationへ向かって、学習されたニューラルネットワークを用いて3次元座標を段階的に更新する。
• コンformationの改善プロセスを監視するために、RDKit で生成されたコンformation と DFT で最適化されたコンformation の間の擬似線形軌道に沿って中間コンformation をサンプリングする、革新的なトレーニング戦略を採用する。
• ペア表現の改善に、繰り返しの外積、三角形更新、全結合ネットワークを組み合わせたマルチヘッドアテンション機構を統合する。
• 大規模な、原始的コンformation と DFT コンformation のペairedデータセットを活用し、L1損失によるQC特性予測と、ノイズ除去目的のコンformationの改善の両方を用いて、エンドツーエンドでモデルをトレーニングする。

実験結果

リサーチクエスチョン

• RQ1 3次元分子コンformationを DFT で最適化された構造へと学習的に改善することで、深層学習モデルが量子化学的特性予測において最先端の精度を達成できるか？
• RQ2 1D/2D 入力からの直接的モデリングと比較して、コンformationの改善プロセスを段階的な更新の系列としてモデル化することで、予測精度がどのように向上するか？
• RQ3 トレーニング時に擬似軌道と中間コンformation のサンプリングを用いることで、一般化性能が向上し、分布シフトが軽減される程度はどの程度か？
• RQ4 強化されたペア表現と空間的符号化を備えた2トラックのTransformerアーキテクチャは、3次元分子幾何学タスクにおいて、標準的なGNN やTransformerを上回る性能を発揮するか？
• RQ5 提案手法は、OC20ベンチマークに含まれるような分布外の分子データセットに対しても、効果的に一般化できるか？

主な発見

• Uni-Mol+ は、OC20 IS2RE テストセットで 0.3745 eV の平均絶対誤差 (MAE) を達成し、以前の最先端手法（DRFormer）の 0.3865 eV と比較して顕著に優れている。
• PCQM4MV2 ベンチマークでは、すべてのカテゴリで最小の MAE を達成しており、インドメイン（ID）では 0.3795 eV、分布外吸着（OOD Ads.）では 0.4526 eV、両方分布外（OOD Both）では 0.4021 eV であった。
• OC20 テストセットでは 11.29% の EwT（エネルギー加重トップ1正答率）を達成し、以前の最先端（DRFormer 9.18%）を2ポイント以上上回っている。
• アブレーションスタディの結果、2トラックのTransformerバックボーンと擬似軌道サンプリング戦略の両方が性能向上に不可欠であることが確認され、いずれかのコンponentを除去すると顕著な精度低下が生じた。
• モデルは分布外データに対しても強くロバストであり、OC20 および PCQM4MV2 ベンチマークのすべての OOD カテゴリで最高の EwT スコアを達成した。
• 本手法はインドメイン設定だけでなく、挑戦的な分布外設定においても最先端の結果を達成しており、強力な一般化能力を示している。