[論文レビュー] Rotation Invariant Graph Neural Networks using Spin Convolutions

Muhammed Shuaibi, Adeesh Kolluru+5 more|arXiv (Cornell University)|Jun 17, 2021

Machine Learning in Materials Science参考文献 34被引用数 51

ひとこと要約

SpinConv は、各エッジの局所フレームと残りの回転自由度に対するスピン畳み込みを用いて回転不変なグラフニューラルネットワークを原子構造に適用し、 energy と force の推定において OC20 で最先端の結果を達成します。

ABSTRACT

Progress towards the energy breakthroughs needed to combat climate change can be significantly accelerated through the efficient simulation of atomic systems. Simulation techniques based on first principles, such as Density Functional Theory (DFT), are limited in their practical use due to their high computational expense. Machine learning approaches have the potential to approximate DFT in a computationally efficient manner, which could dramatically increase the impact of computational simulations on real-world problems. Approximating DFT poses several challenges. These include accurately modeling the subtle changes in the relative positions and angles between atoms, and enforcing constraints such as rotation invariance or energy conservation. We introduce a novel approach to modeling angular information between sets of neighboring atoms in a graph neural network. Rotation invariance is achieved for the network's edge messages through the use of a per-edge local coordinate frame and a novel spin convolution over the remaining degree of freedom. Two model variants are proposed for the applications of structure relaxation and molecular dynamics. State-of-the-art results are demonstrated on the large-scale Open Catalyst 2020 dataset. Comparisons are also performed on the MD17 and QM9 datasets.

研究の動機と目的

• 回転不変 GNN による原子力とエネルギーの効率的で DF Tote-like な予測を動機づける。
• 角度情報を捉えつつ回転不変性を失わないように、エッジごとの局所座標系を開発する。
• 回転自由度の残りに対するスピン畳み込みを提案し、回転不変なエッジメッセージを実現する。
• エネルギー保全と直接力推定のバランスを取るための energy-centric および force-centric SpinConv 変種を提示する。
• OC20 での最先端性能と MD17 および QM9 での競合的結果を示す。

提案手法

• 原子構造を原子をノード、カットオフ内の隣接エッジを持つグラフとして表現する。
• スピン畳み込みを含む共有更新を用いて、距離エンコーディングと原子番号埋め込みを含むエッジメッセージ h_st を反復的に計算する。
• ソース s とターゲット t によって定義される局所フレームを用い、s の周りの球面上に隣接を射影し、st 軸周りの roll 自由度に対してスピン畳み込みを適用する。
• 訓練目的として、エネルギー中心経路（力は -∂E/∂x からの力の導出）または力中心経路（力を直接予測）を採用する。
• スピン畳み込みのためにグリッドベースまたは球面調和関数ベースの球面表現を採用し、回転不変性を保証するためにプーリングを行う。
• メッセージ更新に原子番号情報を融合する埋め込みブロック（mixture-of-experts 風）を組み込む。
• Robustness のため訓練時に全体構造を回転させるオプションを用意し、力ブロックの近似的な回転同値性を促進する。

実験結果

リサーチクエスチョン

• RQ1 spin-convolution を用いた回転不変 GNN が原子系のエネルギーと力場を正確に学習できるか。
• RQ2 energy-centric SpinConv と force-centric SpinConv のエネルギー保存と予測精度の比較。
• RQ3 representation の選択（グリッドベース vs 球面調和関数）とモデルサイズが OC20、MD17、QM9 の性能に与える影響。
• RQ4 複数回の乱れ回転を用いた回転頑健性の導入が力の予測をどれだけ改善するか。
• RQ5 SpinConv は OC20 の in-domain vs out-of-domain（吸着体/触媒）分割においてどう機能するか。

主な発見

• 力中心の訓練を用いた SpinConv は OC20 S2EF タスクで、試験済みモデルの中で最も低いエネルギー MAE および力 MAE を達成した。
• Grid-based なスピン畳み込みは訓練が速く、OC20 に対して球面調和関数と競合する精度を得る；両表現は類似の精度を示す。
• energy-centric SpinConv はエネルギーの精度を向上させるが、全体の S2EF 指標と IS2RS の結果では力中心 variante に劣る。
• OC20 IS2RE では SpinConv が従来アプローチを上回り、特に緩和エネルギー予測の直接回帰設定で優れている。
• OC20 IS2RS では SpinConv が最良の ADwT（分割ごとに平均 53.62%）を達成する一方、DimeNet++-large が最良の AFbT（21.82%）を獲得しており、SpinConv は推論の速さでも優位を保つ（約263.2h vs 814.6h）。
• MD17 および QM9 では SpinConv が最先端手法に対して競合的または優れた性能を示し、特に力中心の定式化と適切な表現サイズを活用した場合に顕著。