[論文レビュー] Better, Faster Fermionic Neural Networks

James S. Spencer, David Pfau+2 more|arXiv (Cornell University)|Nov 13, 2020

Machine Learning in Materials Science参考文献 15被引用数 7

ひとこと要約

本稿では、FermiNetのネットワーク幅と行列式数を増加させることで、アルゴンを含む第二周期元素において化学的精度に到達するよう改善された。JAXへの移行とアーキテクチャの簡素化により、学習時間は10倍短縮され、ボシクロブタンやシクロブタジエンの自己転移反応といった困難な系においても最先端の結果を達成した。

ABSTRACT

The Fermionic Neural Network (FermiNet) is a recently-developed neural network architecture that can be used as a wavefunction Ansatz for many-electron systems, and has already demonstrated high accuracy on small systems. Here we present several improvements to the FermiNet that allow us to set new records for speed and accuracy on challenging systems. We find that increasing the size of the network is sufficient to reach chemical accuracy on atoms as large as argon. Through a combination of implementing FermiNet in JAX and simplifying several parts of the network, we are able to reduce the number of GPU hours needed to train the FermiNet on large systems by an order of magnitude. This enables us to run the FermiNet on the challenging transition of bicyclobutane to butadiene and compare against the PauliNet on the automerization of cyclobutadiene, and we achieve results near the state of the art for both.

研究の動機と目的

• 第二周期の元素、特にアルゴンにおいて、従来のFermiNetバージョンが達成できなかった化学的精度にまで精度を向上させること。
• 大規模な多電子系における学習時間と計算コストを顕著に削減すること。
• アーキテクチャの改善と最適化された学習プロトコルにより、収束性とエネルギー精度を向上させること。
• 多電子状態系としてのシクロブタジエン自己転移反応やボシクロブタンの遷移状態において、FermiNetをPauliNetと比較すること。
• ニューラルネットワークベースの量子化学波動関数アンザッツにおける、精度と学習速度のトレードオフを評価すること。

提案手法

• FermiNetアーキテクチャの1電子ストリームの幅と行列式の数を増加させ、表現力と精度を向上させた。
• FermiNetの実装をTensorFlowからJAXに移行させ、高速な計算とより優れた自動微分を可能にした。
• 完全な共分散エナベロープを削除し、等方的重み初期化を採用することでネットワークを簡素化し、複雑さを低減した。
• 1回の重み更新ごとのMCMCサンプリングステップを10から50に延長し、平衡化を改善し、学習中のノイズを低減した。
• 収束を保証するため、20万回から30万～40万回の重み更新まで学習を延長した。
• エネルギー期待値を最小化するために、確率的勾配降下法とKFAC最適化を用いた変分モンテカルロ法を採用した。

実験結果

リサーチクエスチョン

• RQ1 FermiNetアーキテクチャのサイズを単に拡大することで、アルゴンのような第二周期元素において化学的精度に到達できるか？
• RQ2 アーキテクチャの簡素化とJAXへの移行は、精度を損なわずに学習時間をどの程度短縮できるか？
• RQ3 ボシクロブタンからブタジエンへの反応やシクロブタジエンの自己転移反応といった多電子状態系において、FermiNetはPauliNetに比べてどの程度の性能を示すか？
• RQ4 挑戦的な有機反応における基底状態と遷移状態の相対エネルギーを計算する際、FermiNetの相対的精度はどの程度か？
• RQ5 FermiNetはなぜPauliNetよりも学習が遅いのか？このギャップは、精度を損なわずに縮小可能か？

主な発見

• 行列式数と1電子ストリームの幅を増加させたことで、硫黄を除くすべての第二周期元素において化学的精度（1 kcal/mol以内）を達成した。
• アルゴンに対してエネルギー値 -527.5405(1) $E_h$ を達成し、正確な値と化学的精度で一致した。
• 学習時間はTensorFlow版の11,520 GPU時間からJAX版（簡素化アーキテクチャ）の1,104 GPU時間にまで短縮され、10倍の改善が達成された。
• ボシクロブタンからブタジエンへの遷移において、FermiNetは最高の利用可能手法と同等の化学的精度での相対エネルギー差を達成した。
• シクロブタジエンの自己転移反応において、FermiNetはPauliNetよりもエネルギーを約70 m$E_h$低く計算し、相対エネルギーは10.3 ± 0.1 kcal/molを達成し、MR-CC結果と同等の水準であった。
• 1イテレーションあたりの計算コストがPauliNetの約5倍高いにもかかわらず、多電子状態系においてFermiNetは優れた絶対エネルギー精度を示した。