[論文レビュー] A Self-Attention Ansatz for Ab-initio Quantum Chemistry

Ingrid von Glehn, James S. Spencer+1 more|arXiv (Cornell University)|Nov 24, 2022

Machine Learning in Materials Science被引用数 30

ひとこと要約

Psiformerはab-initio量子化学の自己注意ベースの波動関数 Ansatz で、特に大きな分子に対してFermiNetおよびFermiNet+SchNetと比べて基底状態エネルギーを改善します。

ABSTRACT

We present a novel neural network architecture using self-attention, the Wavefunction Transformer (Psiformer), which can be used as an approximation (or Ansatz) for solving the many-electron Schrödinger equation, the fundamental equation for quantum chemistry and material science. This equation can be solved from first principles, requiring no external training data. In recent years, deep neural networks like the FermiNet and PauliNet have been used to significantly improve the accuracy of these first-principle calculations, but they lack an attention-like mechanism for gating interactions between electrons. Here we show that the Psiformer can be used as a drop-in replacement for these other neural networks, often dramatically improving the accuracy of the calculations. On larger molecules especially, the ground state energy can be improved by dozens of kcal/mol, a qualitative leap over previous methods. This demonstrates that self-attention networks can learn complex quantum mechanical correlations between electrons, and are a promising route to reaching unprecedented accuracy in chemical calculations on larger systems.

研究の動機と目的

• 外部データを用いずに、多電子シュレディンガー方程式を解くための自己注意ベースのニューラルネット Ansatz を開発する。
• 既存のニューラルネット Ansatz と比較して分子系の基底状態エネルギー計算の精度を向上させる。
• 注意機構ベースのアーキテクチャがより大きな分子へ拡張可能であることを示す。
• Jastrow因子とエンベロープ関数を介して cusp 条件などの物理的制約を組み込む。

提案手法

• Psi_theta を exp(J_theta) と行列式の和の積として定義する。J_theta は Jastrow に似た前因子であり、Phi_theta^k は自己注意ネットワークを用いて学習する。
• スピンを入力に連結した電子-核相互作用の入力特徴をモデル化するため、単一ストリームのマルチヘッド自己注意アーキテクチャを使用する。
• 大きな分離幾何を扱うための入力リスケーリングを用い、指数的エンベロープによって境界条件を課す。残差接続と tanh 非線形性を適用。
• レイリー商最適化とメトロポリス-ヘイスティングスサンプリングを用いた変分量子モンテカルロで学習する。KFAC（Kronecker-factored Approximate Curvature）で最適化する。
• ベンゼンを含む小規模および大規模分子ベンチマークで Psiformer を FermiNet および FermiNet+SchNet と比較し、精度向上を評価する。

実験結果

リサーチクエスチョン

• RQ1 自己注意ベースのニューラルネットが ab-initio 量子化学における多電子シュレディンガー方程式に対して競争力のあるまたは優れた Ansatz を提供できるか。
• RQ2 Psiformer は既存のニューラルネット Ansatz と比べて基底状態エネルギーの精度を改善するか、特により大きな分子系で。
• RQ3 スピンの組み込み、Jastrowによる cusp 処理、自己注意型の電子-核結合を取り入れることが、学習と安定性をシステムサイズ全体でどう影響するか。
• RQ4 自己注意アーキテクチャは、従来の手法と比較して、十数個から数十個の電子を有する分子へどの程度スケールできるか。

主な発見

• Psiformer は FermiNet および FermiNet+SchNet よりも著しくエネルギーを改善し、特に大規模な系で顕著。
• 小さな分子では、巨大な Psiformer がオゾンのような一部の系で CCSD(T)/CBS に約 1 kcal/mol の内側まで達する。
• 42–74電子分子（ベンゼンからCCl4）に対して、Psiformer は FermiNet を上回り、いくつかの公表済み DMC エネルギーをも上回る。ベンゼンと CCl4 で顕著な改善。
• Psiformer のエネルギーは系のサイズとともに改善し、巨大分子で従来ベースラインに対して最大約 ~75 kcal/mol (100 mHa) の改善を示す。
• LayerNorm や他の学習選択が最適化に影響を与え、Psiformer は一般に構成間で優れた精度を維持する。