[1909.08423] Deep neural network solution of the electronic Schrodinger equation [paper-reading]

[yoheikikuta]

論文リンク

https://arxiv.org/abs/1909.08423

公開日（yyyy/mm/dd）

2019/09/16

概要

量子化学で多体相互作用を計算する手法として PauliNet というモデルを提案。
Slater-Jastrow-Backflow 型の波動関数を試行関数として準備するところに人間の知識を入れて学習を効率的にしつつ、これまでは決めでパラメタを手で与えていた部分を DNN で学習することで高い表現力を確保することを可能にした。
物理的要請（例えば電子の相対位置のみが重要）をうまく取り込むために graph convolutional neural network の一種である SchNet を改変して用いている。
学習は教師なしの変分量子モンテカル法で、変分原理に従いエネルギーを最小化するようなパラメタを学習するというものになっている。
電子数 $ N $ に対して $ O(N^4) $ くらいで計算可能で、先行研究の $ O(N^7) $ くらいの手法と同等程度の性能を発揮できた。
同時期に DeepMind が提案した FermiNet は知識の encoding が少なくてより大量に学習するモデルで、精度は少し劣るが学習効率は 2 桁くらい良いと主張している。

著者実装：https://github.com/deepqmc/deepqmc

[yoheikikuta]

PauliNet という名前がパンチありすぎるな〜と思って読んでみることにする。
ちなみに同時期に DeepMind が FermiNet https://deepmind.com/blog/article/FermiNet というモデルを提案している。
量子化学計算の文脈ではないけど DiracNets https://arxiv.org/abs/1706.00388 というのもあったな。
みんな名前でやりたい放題やんけ。機械学習界隈の論文は目立つために名前をキャッチーなものにする傾向が強いよね。

名前がどうこうは置いておいて、中身を読んでいく。

[yoheikikuta]

intro の書き出しがパワーを感じる。

A solution of the time-independent electronic Schrödinger equation of a given atomic system provides, 
in principle, full access to its chemical properties. 

それはそう。
まさに原理的には可能というやつ。

電子数が数十個程度の分子に対しては配置間相互作用法や結合クラスター法、多次元量子モンテカルロ法、などで全電子エネルギーの有効数字が6桁の精度で解けているとのこと。2012 の段階で。手法の詳細は一旦おいといて、これらの手法は典型的には電子数を $ N $ としたときに $ N^7 $ やそれ以上で計算コストが大きくなるため電子数が増えると使えなかった。
密度汎関数法は $ N^3 $ で計算できるし精度も結構良いのでかなり使えるが、高精度の計算をしたいという時には精度に限界がある。

計算量を抑えつつ精度を上げたい！というのがモチベーション。

[yoheikikuta]

計算コストの $ N^7 $ とかはちょっと直感的にはよく分からない（というか具体的な手法もよく知らない）のだが、ほとんどの量子化学の手法では電子の波動関数をスレーター行列式の線形結合で表現している、ということから電子数が増えると計算量がめちゃ増えることは想像がつく。
厳密には無限個の分子軌道と無限個のスレーター行列式が必要になり、それを有限個に近似したとしてもスピン軌道の組み合わせは組み合わせ爆発的に増えるので大変。

ちなみにこれが配置間相互作用法というもので、波動関数を様々なスピン軌道の組のスレーター行列式の線形結合で表現して変分法でその係数とエネルギー固有状態を求めるというものらしい。Hartree Fock 法が全波動関数を一組のスピン軌道から成るスレーター行列式で表現している平均場近似であったことを思い出すと、これは電子相関を取り入れるための自然な拡張であると思われる。

スレーター行列式は量子化学の手法において重要で、手法によって異なる役割を持っているとのこと。
以下の 2nd quantization に属するものはスレーター行列式を基底として問題が解かれるもの（上で説明した配置間相互作用法）で、これは多体相互作用をより正確に取り入れるためにより多くの項（すなわちより多くのスレーター行列式）が必要になるため横軸の number of determinants が多くなっている。
1st quantization に属するものは実空間で問題を解き（電子の入れ替えに対する反対称性は演算子によるものではなく、波動関数で $ r_i \leftrightarrow r_j $ 入れ替えで符号が変わるということで担保）、電子運動の相関の大部分を明示的に取り扱う。物理的知識を encode した何かしらのパラメタでパラメトライズされた波動関数を ansatz を準備し、そこで determinant が使われたりするが、その必要な数は 2nd quantization と比べると少ないというものになるとのこと。正直あまりよく分かってないが、ansatz としてうまいものを準備したりそのパラメタをうまく決めたりする方法があったりすることで効率的になるということなのかな。

[yoheikikuta]

上の図で出てくる QMC は quantum monte-carlo の略。
ただし量子モンテカルロといっても色々あるので、ここではいわゆる変分量子モンテカルロのみを対象にする（他には経路積分モンテカルロとかもある）。

変分量子モンテカルロは量子状態に対してあるパラメタでパラメトライズされた特定の形を考え、変分法によってハミルトニアンの基底状態を求めるという手法。
日本語で読める資料だと https://whyitsso.net/physics/quantum_mechanics/quantum-monte-carlo.html が簡潔にまとまっていて良かった。今回は深入りせずこのページに書かれていることくらいを押さえて論文を読んでいくことにする。

[yoheikikuta]

量子状態に対してあるパラメタでパラメトライズされた特定の形 のところで neural networks でパラメトライズしようというのがこの論文のアイデアをナイーブに表現したものになる。
物理的な知識を encode して、neural networks の表現力を活かして学習がうまくできれば良い精度で解けるだろうというお話。
変分法を使った教師無し学習（エネルギー最小化を目指して収束するまで）になるので、既存の学習データを必要としない。学習のテクニックとかも発展してきているので、話だけ聞くとやってやれないことはない気はしてくる。

ちなみにもっと大胆にいわゆる第一原理的手法（ボトムアップで非相対論的シュレディンガー方程式を解いて性質を理解する）で精度に限界がない形で neural networks で強く相関がある波動関数を求めるという先行研究もあったようだが、波動関数に対して ansatz を用いておらず（つまり物理的な知識を encode せずに学習だけでいい感じに学ばせようとして）ハートリーフォック近似の精度にも達してないと紹介されている。まあある程度は知識を外から与えないと流石に厳しい気がする。ってかその場合反対称性とかどうなるんだ？まあこの論文とは直接関係ないのでいったんスルー。

これまでの説明を踏まえて、この論文の概要を三行で書くと以下のようになる。

• これまで ad-hoc に入れられていた Jastrow factor と backflow transformation を neural network でパラメトライズして最適な形に学習できるようにした（こいつらがなんなのかはあとでみる）
• ベースとして多参照ハートリーフォック法の結果が使われるようにして、そこから補正として電子相関が学習されるようにすることで変分法の初期値として良いものを担保した
• 結果として N^4 くらいのスケーリングで高い精度を達成

以降ではこの具体的な内容の理解に努める。

---

ちなみに FermiNet では基本的な反対称性以外の ansatz を用いずより大量のパラメタを使って学習し、計算量は PauliNet よりだいぶ大きい（計算コストが 2 桁くらい上と書いてある）が精度は高いものを達成している、とのこと。

[yoheikikuta]

手法について見ていく。

まず、変分原理に基づいて全電子のエネルギーを最小化するような波動関数を教師なし学習で見つけるというのが基本戦略。
物理であればエネルギーについてはよく知っている（核子のエネルギーとかポテンシャルとか）のでこれは「正しい」ものを選ぶことができる。括弧付きで書いたのは典型的な機械学習の手法とかだとどのような形にすべきかの指導原理がない場合もあるが、物理であればそれがあるという意味。
あとは波動関数を人間の物理学の知識を encode しつつパラメトライズ（下記では θ）したものとして準備し、学習によってエネルギーを最小にするような θ を見つける、という手順である。

（エルミート共役がなかったりして正しい表記じゃないんだけど...）
お察し（？）の通り、この期待値の積分を真面目にやろうとすると各電子の位置を考えなければならないので電子の数が増えると立ち行かなくなる。そこで変分量子モンテカルロを使う。雑にいうとハミルトニアンの左隣に $ \psi \psi^{-1} $ を入れて local energy $ E_{loc} [\psi] (r) = \psi^{-1} H \psi $ を定義した時に、この local energy を確率分布 $ |\psi (r)|^2 $ にしたがってサンプリングして計算することで期待値を計算する式に読み替え、このサンプリングをよしなな回数で打ち切って期待値を近似するという手法である。

local energy という言葉は初めて聞いたけど、これは座標に依存するという意味だろう。
あとは波動関数を $ \psi_θ $ に置き換えて、total energy を loss function として weigthed adam optimizer として最小エネルギーと（その配位を表現する θ を求める）というのが形式的な流れ。
（微分計算がコスト高くて不安定なので local energy を clip するとかもあるみたいだけど、自分の興味はそこの詳細まではないのでスキップ）

波動関数 $ \psi_θ $ としてどのような形のものを準備するかが腕の見せ所になる。
この論文では以下のような形を仮定する。

これには以下の情報が含まれる。

• チルダなしの φ ((1) 式の小文字のファイ) が一つの電子の分子軌道を表す。このチルダなしの φ ((1) 式の小文字のファイ) は初期値としてハートリーフォック法で求めたものを PauliNet の input にする。
• 電子のポテンシャルは原子核の位置で cusp (微分が不連続になる尖った点) を持つという物理学の知識があるので、その影響を取り込む必要がある。電子の cusp 条件は γ に取り込む。これは先行研究で簡単化して取り込む手法があり、γ に関してはパラメタ依存性がなく与えられる。原子核の cusp 条件は φ ((1) 式の小文字のファイ) に取り込む。これは DNN として SchNet という graph convolutional neural network を改修して使うことで満たすようにする（このモデルについては後述）。
• J_θ は Jastrow factor というもので、これはもともと波動関数を (電子のスピン up の slater 行列式) x (電子のスピン down の slater 行列式) x exp(Σ u(r_{ij})) のように factorrize するのが有効と知られていたものを踏襲している。先行研究ではここはある決まった形で導入していたが、これを parametrize して学習しようということで J_θ (r) という形にしている。このままだと元々の相対座標で決まるという情報が含まれていないので、これも SchNet のアーキテクチャで担保するようにする。
• チルダありの φ ((1) 式の小文字のファイ) を分解して出てくる f が backflow と呼ばれるものである。もともと N 個の電子がある時、それに対応する N 個の分子軌道を選び、一つのスレーター行列式を選ぶというのが標準的なやり方だった。これだと可能な分子軌道の組み合わせのうちごく少数の組み合わせしか取れないので、それを解決するために色々な分子軌道の組み合わせを考えてスレーター行列式をたくさん作ってその線形結合を取るというのが配置間相互作用法だった。これを効率的にするために、まず N 個の電子に対して Neural Network でそれぞれの相互作用を考慮した N 個の出力を出す。この N 個の出力を仮想的な電子（これは他の電子との相関をいい感じに考慮してくれてる）として、それらに対して N 個の分子軌道を選び一つのスレーター行列式を構成する、という手法らしい。これにより少ないスレーター行列式でも表現力が高まるというもの。ちなみに backflow 自体は元々ある発明で、そこに Neural Network を利用するというのがこの論文での発明っぽい。
• とはいえスレーター行列をたくさん準備する方が精度があがるので、それは p の添字である程度の数準備するというものになっている。

だいぶ複雑な感じだが、ここまで述べたことを図示したものが以下。
一番上の X,Y が input っぽい形で書いてあるけど trainable になってるのがちょっと読みづらい感じもするが、これは潜在表現に対応するもので、学習によっていい感じに backflow とか Jastrow を構築できるできるような潜在表現となる（SchNet で学習する）。

[yoheikikuta]

SchNet https://arxiv.org/abs/1706.08566 は graph convolutional neural network で、従来はノード間の相互作用をグラフ上の距離のみでモデリングしていたところを、学習によってノード間の相互作用をよしなに考慮して特徴ベクトルを更新することでノード距離によらずに相互作用を表現するようなモデルらしい。全然ちゃんと読んでない。

今の話で言えば、そもそも graph convolution neural network でノード間の距離というのが重要なので、これが粒子の相対位置が重要というのにマッチしているので使っている。上で述べたような SchNet の特徴がどれくらい効いてるかはよくわからんけど、相互作用をいい感じに取り入れてくれそうというくらいで採用してるのかな。
一応こういう感じに modify して PauliNet では使っているというのが以下の式で出てくる。Y, x がそれぞれ核子と電子の特徴ベクトルで z が message である。他にも色々記号が出てきてちゃんと理解できてないが、実空間の相対座標に基づいて電子と核子の message が update されてそれをもとに電子の特徴ベクトルが update されるのを見て取ることができる。核子の特徴ベクトルは fixed geometry として変分量子モンテカルロで一度最適化してるので更新はしなくてよい。

[yoheikikuta]

実験結果は以下。
先行研究と比べるとかなり良い精度であることが分かる。correlation energy とは exact energy から HF で計算される energy を引いたもので、それがどれくらい再現できてるかをパーセンテージで表している。
論文ではあまり結果の図が載っていないが、YouTube にアップされてる著者の talk https://www.youtube.com/watch?v=_bdvpmleAgw ではもっと色々な図が出てくるのと backflow の説明とかもあるので分かりやすい。この論文だけだとどれくらい PauliNet が効率的かというのが結果だけみるとあまりはっきりしないけど、動画では $ O(N^7) $ の手法と比べて $ O(N^4) $ の PauliNet がほぼ同程度の性能を出していることが示されている。

[yoheikikuta]

あまりちゃんと読みきれなかったが、概略はこんなもんかということで終了としておく。
物理の知識を入れ込みつつ DNN も取り入れることで効率的で精度が高い手法を生み出したというのはなかなか面白い。一方で比較的少ない知識 encode で大量に学習することでこれよりも高い精度を発揮したという FermiNet が出てきたのも時代って感じがする。

諸々の前提知識がなくてちゃんと理解するのはまだまだできてないが、今後どういう風に発展していくかはたまにチェックしていきたい。