[論文レビュー] Multi-task learning for electronic structure to predict and explore molecular potential energy surfaces

Zhuoran Qiao, Feizhi Ding+6 more|arXiv (Cornell University)|Nov 5, 2020

Machine Learning in Materials Science参考文献 17被引用数 9

ひとこと要約

本論文では、低コストなタイトバインディング計算から得られる対称性適合原子軌道（SAAO）特徴量を活用して、DFTレベルの精度で分子エネルギーと力を予測するが、計算コストを1000分の1に抑えたマルチタスク・グラフニューラルネットワーク、OrbNetを提案する。解析的勾配と補助電子構造性質のマルチタスク学習を組み合わせることで、QM9およびコンフォーマント最適化タスクにおいて最先端の性能を達成するとともに、化学的空間全体にわたる汎用性を維持する。

ABSTRACT

We refine the OrbNet model to accurately predict energy, forces, and other response properties for molecules using a graph neural-network architecture based on features from low-cost approximated quantum operators in the symmetry-adapted atomic orbital basis. The model is end-to-end differentiable due to the derivation of analytic gradients for all electronic structure terms, and is shown to be transferable across chemical space due to the use of domain-specific features. The learning efficiency is improved by incorporating physically motivated constraints on the electronic structure through multi-task learning. The model outperforms existing methods on energy prediction tasks for the QM9 dataset and for molecular geometry optimizations on conformer datasets, at a computational cost that is thousand-fold or more reduced compared to conventional quantum-chemistry calculations (such as density functional theory) that offer similar accuracy.

研究の動機と目的

• DFTレベルの精度で分子ポテンシャルエネルギー面を予測できる機械学習モデルを開発すること。
• ドメイン特化型で物理的に根拠のある特徴量を用いることで、化学的空間および非平衡状態の幾何構造への一般化性と転送性を向上させること。
• グラフニューラルネットワークフレームワーク内での解析的勾配導出により、力と応答性質のエンドツーエンド微分可能化を実現すること。
• 密度行列の射影や双極子モーメントなどの補助電子構造性質に対するマルチタスク学習を組み込むことで、学習効率とモデルの忠実度を向上させること。
• QM9やコンフォーマント最適化を含むベンチマークデータセット上でモデルの性能を実証し、優れた精度とスケーラビリティを示すこと。

提案手法

• 本モデルは、対称性適合原子軌道（SAAO）基底における低コストなGFN-xTB計算から得られるノードおよびエッジ特徴量を用いたグラフニューラルネットワークアーキテクチャを採用し、フォック行列、密度行列、ハミルトニアン行列、重なり行列、距離行列を含む。
• 密度行列の射影、双極子モーメント、線形応答性質などの補助的ターゲットを予測するためのマルチタスク学習を可能にするために、原子固有のおよび分子全体のグローバルなアテンション特徴量を導入する。
• 全エネルギーは、タイトバインディングエネルギーとニューラルネットワークで予測された補正項の和として計算され、力は全エネルギー式の解析的勾配により得られる。
• すべての電子構造項に対して解析的勾配が導出され、数値微分を用いずとも力と応答性質の微分可能予測が可能になる。
• エネルギーおよび補助的ターゲット予測のための残差ニューラルネットワークデコーダを採用し、GradNormによる適応的損失重み付けによりマルチタスクの目的をバランスさせる。
• SAAO行列要素およびグローバルな対称性適合記述子を特徴量として用い、物理的整合性と転送性を保証するように、分子グラフ上でエンドツーエンドで学習を行う。

実験結果

リサーチクエスチョン

• RQ1 低コストな電子構造特徴量に基づいて学習したマルチタスクディープラーニングモデルは、分子エネルギーと力の予測においてDFTレベルの精度に達することができるか？
• RQ2 密度行列の射影や双極子モーメントなどの補助電子構造ターゲットを組み込むことで、学習効率とモデルの一般化性能はどのように向上するか？
• RQ3 本モデルは、未学習の分子や非平衡状態の幾何構造、特に大規模系に対してもどの程度一般化できるか？
• RQ4 本手法は、従来のDFT計算と比較してどの程度計算コストの優位性を示すか？
• RQ5 グラフニューラルネットワークフレームワーク内で解析的勾配を効果的に導出し、数値近似なしに正確な力の予測を実現できるか？

主な発見

• QM9データセットにおいて、エネルギー予測の平均絶対誤差が1.34 kcal/molという最先端の性能を達成し、既存手法を上回った。
• コンフォーマントデータセットにおける分子幾何最適化では、最適化された結合長の平均二乗偏差が0.046 ÅとDFTに近い精度を達成した。
• 従来のDFTと比較して計算コストが1000倍以上低減され、大規模なポテンシャルエネルギー面の迅速な探索が可能になった。
• 密度行列の射影などの補助的ターゲットを含むマルチタスク学習により、単一タスクベースラインと比較して誤差が15%低減され、一般化性能が向上した。
• 解析的勾配フレームワークにより、平均絶対誤差が0.024 kcal/mol/Åという高精度な原子間力の予測が可能となり、DFTと同等の性能を示した。
• 本モデルは強力な転送性を示し、学習データセットに含まれる分子よりもはるかに大きな分子、特に複雑な有機および無機系においてもエネルギーと力を正確に予測できた。