[論文レビュー] ForceNet: A Graph Neural Network for Large-Scale Quantum Calculations

Weihua Hu, Muhammed Shuaibi+6 more|arXiv (Cornell University)|Mar 2, 2021

Machine Learning in Materials Science参考文献 53被引用数 24

ひとこと要約

ForceNet は、OC20 での原子力学的フォース予測において物理ベースの GNN を上回る、スケーラブルな force-centric グラフニューラルネットワークであり、トレーニングと推論の大幅なスピードアップを提供します。ForceNet-large は、DimeNet++ に匹敵するかそれ以上を、はるかに少ない計算量で実現します。

ABSTRACT

With massive amounts of atomic simulation data available, there is a huge opportunity to develop fast and accurate machine learning models to approximate expensive physics-based calculations. The key quantity to estimate is atomic forces, where the state-of-the-art Graph Neural Networks (GNNs) explicitly enforce basic physical constraints such as rotation-covariance. However, to strictly satisfy the physical constraints, existing models have to make tradeoffs between computational efficiency and model expressiveness. Here we explore an alternative approach. By not imposing explicit physical constraints, we can flexibly design expressive models while maintaining their computational efficiency. Physical constraints are implicitly imposed by training the models using physics-based data augmentation. To evaluate the approach, we carefully design a scalable and expressive GNN model, ForceNet, and apply it to OC20 (Chanussot et al., 2020), an unprecedentedly-large dataset of quantum physics calculations. Our proposed ForceNet is able to predict atomic forces more accurately than state-of-the-art physics-based GNNs while being faster both in training and inference. Overall, our promising and counter-intuitive results open up an exciting avenue for future research.

研究の動機と目的

• 大規模データを用いて原子力の高価な量子物理計算を高速かつ正確に近似することを動機づける。
• アーキテクチャに明示的な物理制約に依存しない、柔軟で表現力豊かな force-centric GNN の開発。
• 回転関連の物理特性を暗黙的に強制するために、物理ベースのデータ拡張を活用する。
• OC20 を対象としたスケーラビリティと効率性の検討、非平衡構造を含む200Mを超える構造データセット。

提案手法

• 原子をノードとし、エッジが条件付きフィルター畳み込みによって3D相互作用を捉えるエンコーダ–デコーダ型 GNN アーキテクチャを用いて ForceNet を開発する。
• 完全な3Dの相対配置と回転共変の方向情報をエンコードするエッジ特徴 e_st を使用する。
• 生のエッジ特徴をより豊かな表現に変換する基底関数 B を導入（球関数の variantes を含む）し、エッジおよびノードのMLPに表現力豊かな非線形性 Swish を適用する。
• トレーニング時に回転データ拡張を適用して、明示的なアーキテクチャ回転制約を課すことなく回転共変予測を促す。
• 最終ノード埋め込みからデコーダMLPを介して力の予測を直接構築し、エネルギー保存制約を明示的に課さないことでスケーラビリティを得る。
• OC20 S2F タスクで SchNet、DimeNet++、DimeNet++-large、GNS と比較し、トレーニング/推論予算を慎重に制御する。
• アイデンティティ、線形+活性化、ガウス、正弦、球関数などの基底関数の範囲と、Swish が最良の活性化関数であることを含む活性化を評価し、性能への影響を検討する。

実験結果

リサーチクエスチョン

• RQ1 明示的な回転制約を持たない force-centric GNN が、大規模かつ非平衡なデータセットで物理ベースの GNN を上回る最先端の予測を達成できるか。
• RQ2 アーキテクチャの選択（基底関数・活性化・エッジ条件付け）とデータ拡張が力の予測精度と回転共変性にどう影響するか。
• RQ3 ForceNet を OC20 で拡張する際のトレーニング/推論の効率性のトレードオフは既存モデルとどう違うか。
• RQ4 回転拡張は回転非不変モデルにどの程度回転共変性を付与できるか。

主な発見

• ForceNet は比較可能な計算量の下で OC20 の原子ごとの力予測において物理ベースの GNN を上回り、トレーニングと推論がより高速。
• ForceNet-large は DimeNet++-large と同等またはそれ以上の力予測を達成しつつ、計算量を大幅に低減（トレーニング約8.2x–18.9x低減、推論約2.4x–3.0x高速、比較に依存）する。
• ForceNet–large は平均力 MAE で DimeNet++ を上回り、トレーニングと推論を大幅に高速化した。初期結果は、SchNet に対して約6.3x の計算削減で類似の性能を達成、DimeNet++ に対しては推論で最大約6.5x の高速化。
• 回転データ拡張と大規模データは、回転下での不安定性を低減し、ForceNet の予測を回転共変挙動に近づけるのに有効である。
• 表現力のあるエッジレベル計算と球基底＋Swish 活性化の組み合わせは ForceNet の性能にとって不可欠であり、これらを欠くと精度が劣化する。