[論文レビュー] GemNet-OC: Developing Graph Neural Networks for Large and Diverse Molecular Simulation Datasets

Johannes Gasteiger, Muhammed Shuaibi+5 more|arXiv (Cornell University)|Apr 6, 2022

Machine Learning in Materials Science被引用数 54

ひとこと要約

GemNet-OCはOC20で最先端の結果を達成しつつ、訓練は約10x速く、論文はデータセットの規模、多様性、ドメインシフトがGNN設計に与える影響を分析し、効率的な開発の代表的なサブセットとしてOC-2Mを導入している。

ABSTRACT

Recent years have seen the advent of molecular simulation datasets that are orders of magnitude larger and more diverse. These new datasets differ substantially in four aspects of complexity: 1. Chemical diversity (number of different elements), 2. system size (number of atoms per sample), 3. dataset size (number of data samples), and 4. domain shift (similarity of the training and test set). Despite these large differences, benchmarks on small and narrow datasets remain the predominant method of demonstrating progress in graph neural networks (GNNs) for molecular simulation, likely due to cheaper training compute requirements. This raises the question -- does GNN progress on small and narrow datasets translate to these more complex datasets? This work investigates this question by first developing the GemNet-OC model based on the large Open Catalyst 2020 (OC20) dataset. GemNet-OC outperforms the previous state-of-the-art on OC20 by 16% while reducing training time by a factor of 10. We then compare the impact of 18 model components and hyperparameter choices on performance in multiple datasets. We find that the resulting model would be drastically different depending on the dataset used for making model choices. To isolate the source of this discrepancy we study six subsets of the OC20 dataset that individually test each of the above-mentioned four dataset aspects. We find that results on the OC-2M subset correlate well with the full OC20 dataset while being substantially cheaper to train on. Our findings challenge the common practice of developing GNNs solely on small datasets, but highlight ways of achieving fast development cycles and generalizable results via moderately-sized, representative datasets such as OC-2M and efficient models such as GemNet-OC. Our code and pretrained model weights are open-sourced.

研究の動機と目的

• 小規模データセットでのGNNの改善が、大規模で多様な分子データへ転移するかを調査する

提案手法

• OC20をベースとした多層レベルの相互作用階層とエッジ/原子埋め込みを用いたGemNet-OCベースラインを開発する
• グラフ構築を、固定距離のカットオフから固定数の最近傍近傍に置き換える
• 放射状・角度項の計算コストを削減するため、基底関数を簡略化・最適化する
• 長距離の原子レベル経路を伴う、原子-原子、エッジ-原子、エッジ-エッジ相互作用を含む階層的相互作用スキームを導入する
• 相互作用ブロックからの埋め込みを出力し結合して、最終的なエネルギー/力の予測を改善する
• OC20とOC-2Mで再現性のある開発を可能にするオープンソースコードと事前学習済みウェイトを提供する

実験結果

リサーチクエスチョン

• RQ1 4つのデータセット特性（化学的多様性、系のサイズ、データセットサイズ、ドメインシフト）は、GNN設計の決定にどのような影響を与えるか？
• RQ2 OC20で訓練されたモデルが、訓練コストを削減しつつ最先端の性能を達成できるか？
• RQ3 OC-2Mは、より高速な開発のためにOC20のトレンドと相関する信頼できるサブセットか？
• RQ4 モデルの構成要素の効果は、小規模と大規模データセット、およびOC20のサブセット間で異なるか？
• RQ5 大規模な分子データセットに対して、迅速でスケーラブルな学習を最も実現するアーキテクチャと訓練戦略は何か？

主な発見

• GemNet-OCはOC20タスクで最先端の成果を達成し、従来の大規模モデルより訓練が約10x速い
• GemNet-OCはOC-2MおよびOC20の訓練データで従来モデルを上回り、はるかに少ない訓練データを使用する
• 小規模データと大規模で多様なデータセットにおけるモデル構成要素のパフォーマンスには顕著な差異がある
• OC-2Mは完全なOC20の結果と良く相関し、より速く、代表的なモデル開発を可能にする
• 長距離の原子レベル経路と最適化された基底関数を備えた相互作用階層は、大規模システムの効率的な取り扱いを可能にする