大規模パラメータを持つグラフニューラルネットワークを原子シミュレーションに向けた訓練へ

最近のグラフニューラルネットワーク（GNN）における進歩は、原子シミュレーションのモデリングに活用され、気候変動対策に必要なエネルギー革命への重要な一歩である触媒発見を革新する可能性を持っています。しかし、このタスクで最も効果的と証明されたGNNは、三重や四重の原子間相互作用などのグラフ内の高次相互作用をモデル化するため、メモリ負荷が高く、これらのモデルをスケーリングすることが困難となっています。本論文では、Graph Parallelismという手法を導入します。これは複数のGPUに入力グラフを分散させる方法であり、数百億乃至数十億パラメータを持つ非常に大きなGNNの学習を可能にします。我々はこの手法を実証的に評価するために、最近提案されたDimeNet++およびGemNetモデルのパラメータ数を1桁以上増加させることでスケーリングを行いました。大規模なOpen Catalyst 2020 (OC20) データセットにおいて、これらのグラフ並列化されたモデルは以下の相対的な改善をもたらしました：1) S2EFタスクにおける力の平均絶対誤差（MAE）指標で15%の改善と 2) IS2RSタスクにおけるAFbT指標で21%の改善です。これにより新たな最先端の結果が樹立されました。