Towards Training Billion Parameter Graph Neural Networks for Atomic Simulations
This work addresses a key bottleneck in catalyst discovery for combating climate change by enabling larger, more accurate GNNs, though it is incremental as it builds on existing models like DimeNet++ and GemNet.
The paper tackles the challenge of scaling memory-intensive Graph Neural Networks (GNNs) for atomic simulations by introducing Graph Parallelism, a method to distribute input graphs across multiple GPUs, enabling training of models with billions of parameters and achieving relative improvements of 15% on force MAE and 21% on AFbT metrics on the OC20 dataset.
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.