[論文レビュー] Generalizing Neural Wave Functions

Nicholas Gao, Stephan Günnemann|arXiv (Cornell University)|Feb 8, 2023

Machine Learning in Materials ScienceMaterials Science被引用数 3

ひとこと要約

本稿では、Globeと呼ばれるグラフニューラルネットワーク（GNN）を用いた軌道埋め込み手法を紹介する。Globeは空間的メッセージパッシングを用いて局所的な電子構造を学習することで、さまざまな分子にわたるニューラル波動関数の一般化を可能にする。さらに、Moonと呼ばれるサイズ整合性を満たしたニューラル波動関数を提案し、従来手法に比べ最大4.5倍速い収束を達成するとともに、系のサイズに応じてエネルギーが加法的にスケーリングされる正確なスケーリングを維持する。これにより、一度のモデルで異なる分子のシュレーディンガー方程式を同時に解くことが可能になる。

ABSTRACT

Recent neural network-based wave functions have achieved state-of-the-art accuracies in modeling ab-initio ground-state potential energy surface. However, these networks can only solve different spatial arrangements of the same set of atoms. To overcome this limitation, we present Graph-learned orbital embeddings (Globe), a neural network-based reparametrization method that can adapt neural wave functions to different molecules. Globe learns representations of local electronic structures that generalize across molecules via spatial message passing by connecting molecular orbitals to covalent bonds. Further, we propose a size-consistent wave function Ansatz, the Molecular orbital network (Moon), tailored to jointly solve Schrödinger equations of different molecules. In our experiments, we find Moon converging in 4.5 times fewer steps to similar accuracy as previous methods or to lower energies given the same time. Further, our analysis shows that Moon's energy estimate scales additively with increased system sizes, unlike previous work where we observe divergence. In both computational chemistry and machine learning, we are the first to demonstrate that a single wave function can solve the Schrödinger equation of molecules with different atoms jointly.

研究の動機と目的

• 既存のニューラル波動関数には、同じ分子の異なる幾何構造には一般化可能だが、異なる分子組成には一般化できないという制限があることの克服。
• 異なる原子数や軌道数を有する複数の異なる分子を同時にシュレーディンガー方程式を解ける1つのニューラル波動関数の実現。
• サイズ整合性をニューラル波動関数に実装し、系のサイズに比例してエネルギーが加法的にスケーリングされることを保証することで、量子化学的物理的正確性を満たすこと。
• 空間的メッセージパッシングと原子核のアンカーポイントを活用した波動関数アーキテクチャの設計により、変分モンテカルロ訓練における収束速度の向上。

提案手法

• Globeは、分子軌道を3次元空間に埋め込むためにGNNにおける空間的メッセージパッシングを用い、異なる分子間で局所的な電子構造を学習可能にする。
• 分子軌道を共有結合に接続することで、異なる数の軌道を有する分子間での一般化が可能になるように、軌道埋め込みを学習する。
• 波動関数アーキテクチャであるMoonは、局所的な電子-原子核相互作用に注目し、原子核をメッセージパッシングのアンカーポイントとして用いることでサイズ整合性を実現する。
• 波動関数は、Globeが再パrameter化した軌道重みを用いたデターミナントの和として構築され、多様な分子における一括学習が可能になる。
• 本手法は二段階のニューラルネットワークを用いる：一つは軌道重みの再パラメータライゼーション（Globe）を担当し、もう一つは最終的な波動関数の計算（Moon）を担当する。
• 学習はエネルギー最小化を目的とした変分モンテカルロ法で実施され、Globeの再パラメータライゼーションは1ステップあたり一度計算され、サンプリングやエネルギー評価と独立して処理される。

Figure 1: Globe steps. 1) Atom embeddings are obtained by message passing between atoms. 2) Orbital locations and embeddings are determined by core orbitals and molecular bonds. 3) Orbital embeddings are updated via a unidirectional message passing. For clarity, we omitted messages from each atom to

実験結果

リサーチクエスチョン

• RQ1 同じ分子の異なる幾何構造だけでなく、異なる原子組成を有する複数の異なる分子に対して、1つのニューラル波動関数を同時に学習可能か？
• RQ2 微分可能でスケーラブルなニューラル波動関数アーキテクチャにおいて、サイズ整合性をどのように実装できるか？
• RQ3 GNNにおける空間的メッセージパッシングは、異なる数の軌道を有する分子間で軌道埋め込みを一般化可能か？
• RQ4 個別学習と比較して、複数の分子を同時に学習させることで収束が早まるか？
• RQ5 異なる分子に一般化されたニューラル波動関数のエネルギー推定値は、系のサイズに加法的にスケーリングされるか？

主な発見

• Moonは、従来手法に比べて4.5倍少ない学習ステップで収束し、エネルギー推定値は同等または低くなることを確認し、最適化収束の高速化を実証した。
• Moonのエネルギー推定値は、系のサイズに加法的にスケーリングされるが、従来手法では発散的スケーリングを示すのに対し、サイズ整合性が確認された。
• Globeは、追加の自己無撞着（SCF）計算を必要とせず、より大きな分子系への転送が可能になり、直接的に大きな鎖状系での評価が可能になった。
• 水素鎖系において、GlobeとMoonを組み合わせた場合、系のサイズが増大しても原子あたりのエネルギーが一定を保ち、強固な拡張性と物理的整合性を示した。
• エチレンやシクロブタジエンといったより大きな分子に対しても、事前学習なしにGlobeとMoonは収束に成功したが、他の手法は事前学習なしでは収束しなかった。
• Moonは、60万パラメータ少ないにもかかわらず、FermiNetと同等の精度を達成し、PsiFormerと同等の性能を示した。これは、高いパラメータ効率性を示している。

Figure 2: Illustration of Moon. 1) We initialize electron embeddings by aggregating their local neighborhood of electrons. 2) Nuclei aggregate electron embeddings via message passing. 3) Nuclei embeddings are iteratively updated. 4) Nuclei embeddings are structurally diffused towards the electrons v