[論文レビュー] Scalable variational Monte Carlo with graph neural ansatz

Li Yang, Wen-Jun Hu+1 more|arXiv (Cornell University)|Nov 25, 2020

Model Reduction and Neural Networks参考文献 42被引用数 11

ひとこと要約

本論文は、グラフ畳み込みネットワーク（GCNs）に基づくグラフニューラルアンサンブル（GNA）を導入し、幾何構造に依存しない普遍的な変分モンテカルロ（VMC）シミュレーションを可能にする。パラメータ共有と128 TPUコアを用いたスケーラブルな分散学習を活用することで、2次元のカグラ、三角格子、およびランダムに接続された格子（最大432サイト）において、正確な基底状態エネルギーの予測が可能となり、小規模な事前学習済み系からの転移学習によって大規模系への高速収束が実現された。

ABSTRACT

Deep neural networks have been shown as a potentially powerful ansatz in variational Monte Carlo for solving quantum many-body problems. We propose two improvements in this direction. The first is graph neural ansatz (GNA), which is a variational wavefunction universal to arbitrary geometry. GNA results in accurate ground-state energies on 2D Kagome lattices, triangular lattices, and randomly connected graphs. Secondly, we design a distributed workflow on multiple accelerators to scale up the computation. We compute Kagome lattices with sizes up to 432 sites on 128 TPU cores. The parameter sharing nature of the GNA also leads to transferability across different system sizes and geometries.

研究の動機と目的

• 任意の格子幾何構造（周期的でない、不規則な構造を含む）を扱える普遍的な量子多体系の変分アンサンブルの開発を目的とする。
• 前例のニューラルネットワークアンサンブル（CNN や RBM）が格子の対称性やグリッド表現に制限を受ける問題を克服することを目的とする。
• 分散アクセラレータ（例：TPU）を用いたスケーラブルなVMCシミュレーションを実現し、メモリおよび計算上のボトルネックを克服することを目的とする。
• 訓練済みGNAモデルの異なる系サイズおよび幾何構造間での転送性を示し、学習時間の短縮と収束性の向上を実現することを目的とする。
• 強い相関を持つ量子系を統一的かつスケーラブルにシミュレートするフレームワークを提供することを目的とし、複素数波動関数やフェルミ系への応用拡張も含む。

提案手法

• 各格子サイトをノードとして扱い、隣接グラフ上のメッセージパッシングにより更新される学習可能な埋め込みを有する、グラフ畳み込みネットワーク（GCN）を用いたグラフニューラルアンサンブル（GNA）を提案。これにより、多体波動関数がエンコードされる。
• 残差接続、ReLU非線形性、層正則化を備えたGCNアーキテクチャを採用。ノード特徴量は正規化された隣接行列と、層間で共有される重み行列を用いて更新される。
• 最終的なノード埋め込みを集約し、1層の全結合層を介して波動関数のスカラー対数振幅にマッピングすることで、変分エネルギーの微分可能最適化が可能となる。
• 複数のアクセラレータ（例：TPU）を用いた重要度サンプリング勾配最適化（ISGO）の分散実装を採用。これにより、大規模系におけるミニバッチサンプリングとパラメータ更新が可能となる。
• 格子構築時に隣接行列の反対側のエッジを接続することで周期的境界条件を強制し、トーラス幾何構造のシミュレーションを可能にする。
• 小規模系（例：12サイト）でGNAを事前学習し、大規模系（例：432サイト）で微調整することで転移学習を実施。幾何構造を越えたパラメータ共有を活用する。

実験結果

リサーチクエスチョン

• RQ1 GNNベースのアンサンブルは、構造的変更なしにカグラ、三角格子、およびランダムに接続されたグラフなど、多様な格子幾何構造に一般化可能か？
• RQ2 GNAにおけるパラメータ共有は、小規模系から大規模系への転移学習をどの程度促進し、学習収束性と安定性を向上させるか？
• RQ3 TPUなどの分散アクセラレータを用いることで、GNAを備えたVMCフレームワークは、大規模系（例：432サイト）にスケーリング可能か？単一デバイスのメモリ制限を克服できるか？
• RQ4 異なる相互作用強度（t/V）と格子種別において、GNA-VMC手法の基底状態エネルギー推定精度は、正確な対角化結果と比較してどの程度高いか？
• RQ5 異なる幾何構造（例：三角格子やランダムグラフ）で事前学習したモデルでも、目的のカグラ格子への収束が加速するか？

主な発見

• GNAは2次元カグラ格子および三角格子において、基底状態エネルギーを正確な対角化結果と比較して0.1%以内の精度で予測した。これはt/V比が大きくても成立する。
• 432サイトの12×12×3カグラ格子において、128 TPUコアを用いて基底状態エネルギーを正確に計算した。これは単一アクセラレータでは以前は不可能であったシステムサイズである。
• 12サイトの小規模系で事前学習したGNAは、収束までの学習イテレーションをわずか5回にまで短縮した。一方、初期化がランダムなモデルは40回のイテレーション経過しても収束しなかった。
• GNAは転送性を示した。三角格子やランダムに接続されたグラフで事前学習したモデルも、目的のカグラ格子への学習を顕著に加速させた。これは、幾何構造を越えた強力な一般化能力を示している。
• スケーラブルなISGO実装により、システムを分割し複数アクセラレータ間で勾配を集約することで、効率的な分散学習が可能となり、大規模VMCが現実可能となった。
• 非正則および不規則なグラフを含む多様な幾何構造においても、高い精度を維持した。これにより、GNAが変分アンサンブルとしての普遍性を有していることが確認された。