[論文レビュー] Ab-Initio Potential Energy Surfaces by Pairing GNNs with Neural Wave Functions

Nicholas Gao, Stephan Günnemann|arXiv (Cornell University)|Oct 11, 2021

Machine Learning in Materials Science参考文献 47被引用数 5

ひとこと要約

本稿では、変分モンテカルロ法（VMC）を用いて複数の分子幾何配置に対してシュレーディンガー方程式を同時に解くための新規アーキテクチャ、PESNetを提案する。グラフニューラルネットワーク（GNN）を用いて1回の学習プロセスで異なる配置に対する波動関数パラメータを予測することで、FermiNetなどの最先端手法と比較して最大40倍高速に学習を達成しながら、精度は同等またはそれを上回る。

ABSTRACT

Solving the Schrödinger equation is key to many quantum mechanical properties. However, an analytical solution is only tractable for single-electron systems. Recently, neural networks succeeded at modeling wave functions of many-electron systems. Together with the variational Monte-Carlo (VMC) framework, this led to solutions on par with the best known classical methods. Still, these neural methods require tremendous amounts of computational resources as one has to train a separate model for each molecular geometry. In this work, we combine a Graph Neural Network (GNN) with a neural wave function to simultaneously solve the Schrödinger equation for multiple geometries via VMC. This enables us to model continuous subsets of the potential energy surface with a single training pass. Compared to existing state-of-the-art networks, our Potential Energy Surface Network PESNet speeds up training for multiple geometries by up to 40 times while matching or surpassing their accuracy. This may open the path to accurate and orders of magnitude cheaper quantum mechanical calculations.

研究の動機と目的

• 変分モンテカルロ（VMC）シミュレーションにおいて、各分子幾何配置ごとに別々に学習する神経波動関数の高い計算コストを解消すること。
• 再学習を必要としない1つのニューラルネットワークを用いてポテンシャルエネルギー表面（PES）を連続的にモデル化すること。
• 物理的対称性（並進、回転、反転）を波動関数のパrameter化に統合し、一般化性と不変性を向上させること。
• 既存の神経波動関数モデルと同等またはそれを上回る精度を達成しながら、学習時間を著しく短縮するスケーラブルで効率的なフレームワークを開発すること。

提案手法

• 分子の核の幾何配置に基づいて、神経波動関数（WFModel）のパラメータを予測するためのメタグラフニューラルネットワーク（MetaGNN）を用いる。
• MetaGNNを、各入力幾何配置に対して変分エネルギーを最小化する波動関数を生成するように学習させ、VMCによって最適化する。
• GNNの入力に等変性座標系を設計することで、物理的対称性（並進、回転、反転）への不変性を強制する。
• 電子配置の確率的サンプリングを用いて波動関数を最適化する変分モンテカルロフレームワークを活用する。
• 複数の分子幾何配置にわたる同時学習を可能にするために、PESNetアーキテクチャを一括でエンドツーエンドに学習する。
• GNNの出力によって条件づけられる共通の波動関数アーキテクチャ（例：FermiNetにインspired）を用い、連続するPES領域における一般化を可能にする。

実験結果

リサーチクエスチョン

• RQ1 1つのニューラルネットワークアーキテクチャが、複数の分子幾何配置を同時に正確な電子波動関数を予測できるか？
• RQ2 GNNで条件づけられた波動関数モデルは、最先端の波動関数を各幾何配置ごとに別々に学習する手法と比較して、どの程度学習時間を短縮できるか？
• RQ3 GNN入力に物理的対称性を統合することで、予測されたポテンシャルエネルギー表面の一般化性と不変性はどの程度向上するか？
• RQ4 再学習を伴わずに、連続するPES領域内の未観測の幾何配置にも一般化できるか？
• RQ5 多配置ポテンシャルエネルギー表面をモデル化する際、学習効率と予測精度のトレードオフはどのようなものか？

主な発見

• PESNetは、H₄⁺、H₄、H₁₀、N₂、シクロブタジエンを含む複数の系において、FermiNetと比較して最大47倍の学習時間短縮を達成した。
• H₄⁺では、A100 GPUを用いて20 GPU時間で学習が完了したのに対し、FermiNetは127時間であった。精度は同等またはそれを上回った。
• H₄では24時間 vs. 118時間、N₂では89時間 vs. 4196時間と、PESNetはPauliNetおよびDeepErwinの両方と比較して、精度と学習効率の両面で優れた性能を示した。
• 1回の学習プロセスでポテンシャルエネルギー表面の連続的部分を捉えることができ、配置数に対して線形スケーラビリティを実現した。
• PESNetは、シクロブタジエンのような複雑な分子を含む全テスト系で最先端の精度を達成し、計算コストを著しく削減した。
• 等変性座標系による物理的対称性の統合は、予測性能に影響を与えることなく、モデルの安定性と一般化性を向上させた。