[論文レビュー] Towards Training Billion Parameter Graph Neural Networks for Atomic Simulations

Anuroop Sriram, Abhishek Das+3 more|arXiv (Cornell University)|Mar 18, 2022

Machine Learning in Materials Science被引用数 9

ひとこと要約

本稿では、複数のGPUに渡って入力グラフを分散することで、10億パラメータ級のグラフニューラルネットワーク（GNN）を原子シミュレーション用にトレーニングする手法「グラフ並列化」を提案する。これにより、原子の三重項や四重項といった高次相互作用をモデル化できる大規模なGNNのスケーラブルなトレーニングが可能になる。Open Catalyst 2020（OC20）データセットにおいて、先行モデルと比較して力のMAEを15%低減、AFbTを21%向上させ、最先端の結果を達成した。

ABSTRACT

Recent progress in Graph Neural Networks (GNNs) for modeling atomic simulations has the potential to revolutionize catalyst discovery, which is a key step in making progress towards the energy breakthroughs needed to combat climate change. However, the GNNs that have proven most effective for this task are memory intensive as they model higher-order interactions in the graphs such as those between triplets or quadruplets of atoms, making it challenging to scale these models. In this paper, we introduce Graph Parallelism, a method to distribute input graphs across multiple GPUs, enabling us to train very large GNNs with hundreds of millions or billions of parameters. We empirically evaluate our method by scaling up the number of parameters of the recently proposed DimeNet++ and GemNet models by over an order of magnitude. On the large-scale Open Catalyst 2020 (OC20) dataset, these graph-parallelized models lead to relative improvements of 1) 15% on the force MAE metric for the S2EF task and 2) 21% on the AFbT metric for the IS2RS task, establishing new state-of-the-art results.

研究の動機と目的

• 原子シミュレーション用に10億パラメータ級のGNNをトレーニングする課題に取り組むこと、特に原子の三重項や四重項といった高次相互作用をモデル化する場合を対象とする。
• 数百万個の中小規模のグラフを含むデータセットに対して、メモリと計算リソースの制限が顕著な従来のGNNトレーニング手法の限界を乗り越えること。
• 大規模なOpen Catalyst 2020（OC20）データセット上で、DimeNet++ や GemNet-T といった最先端のGNNをスケーラブルにトレーニング可能にする。
• グラフ並列化によるモデルサイズの拡大が、原子の力と緩和構造の予測精度に顕著な向上をもたらすことを示すこと。
• 拡張グラフネットワーク（EGN）フレームワーク内において、等長性やエネルギー保存を考慮したモデルを含む、多数のメッセージパッシングGNNアーキテクチャに一般化可能なフレームワークを提供すること。

提案手法

• 入力グラフを複数のGPUに分割することで、1台のデバイスに収まらない大規模なGNNのトレーニングを可能にする、グラフ並列化を提案する。
• グラフネットワーク（GN）フレームワークを拡張グラフネットワーク（EGN）フレームワークに拡張し、原子の三重項や四重項といった高次相互作用を明示的にモデル化する。
• 高次相互作用（例：結合角）の計算をGPU間で分散させ、メモリと計算の両方を効率的にデバイス間で分配する。
• 各GPUが入力グラフのサブセットを処理するデータ並列に類似した戦略を採用しつつ、高次項においても完全なグラフ接続性を維持する。
• GPipeなどの既存のモデル並列化技術と統合することで、非常に大規模なGNNの幅方向および深さ方向のスケーリングを可能にする。
• グラフ分割とモデル並列化を組み合わせたハイブリッドトレーニング戦略を採用し、GNN内のパrameter数と相互作用の複雑さの両方をスケーリングする。

実験結果

リサーチクエスチョン

• RQ1 高次相互作用に起因する高いメモリと計算負荷を抱えるにもかかわらず、10億パラメータ級のGNNを原子シミュレーション用に効果的にトレーニングできるか？
• RQ2 数百万個の小規模な分子グラフを含む分子データセットにおいて、グラフ並列化は従来の分散GNNトレーニング手法と比較して、スケーラビリティとパフォーマンスの点で優れているか？
• RQ3 グラフ並列化によりトレーニングされたより大きなGNNは、力のMAEやAFbTといった重要な原子シミュレーションベンチマークで、予測精度をどの程度向上させるか？
• RQ4 グラフ並列化は、エネルギー保存や回転等長性をモデル化するGNNアーキテクチャに対しても一般化可能か？
• RQ5 このような大規模モデルのトレーニングによる環境的影響は何か。また、計算リソースの戦略的選定によって、それらの影響をどのように軽減できるか？

主な発見

• グラフ並列化により、1億パラメータを超えるDimeNet++およびGemNet-Tモデルのトレーニングが可能となり、従来モデルから10倍以上もパラメータ数を拡大した。
• グラフ並列化されたGemNet-XLモデルは、S2EFタスクにおいて、従来の最先端モデルと比較して力のMAEを15%相対的に改善した。
• 同じモデルは、IS2RSタスクにおけるAFbT指標でも21%相対的に向上させ、OC20データセットで新たな最先端を樹立した。
• 本手法は一般化可能であり、EGNフレームワーク内において、等長性やエネルギー保存を考慮したモデルを含む多数のメッセージパッシングGNNアーキテクチャに適用可能である。
• Tesla V100 32GB搭載のクラウドGPU上でGemNet-XLをトレーニングした結果、推定されるCO2排出量は3490〜8052 kg CO2 eq.に達し、ロサンゼルスからニューヨークへの往復フライト16回分に相当する。
• 著者らは、環境影響を低減するため、炭素オフセットの取り組みを積極的に行う計算リソースを選択することを推奨しており、本研究のトレーニングは100%オフセット済みであると述べている。