Towards Training Billion Parameter Graph Neural Networks for Atomic Simulations

by   Anuroop Sriram, et al.

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 task and 2) 21 state-of-the-art results.


page 1

page 2

page 3

page 4


Weisfeiler and Leman Go Neural: Higher-order Graph Neural Networks

In recent years, graph neural networks (GNNs) have emerged as a powerful...

ForceNet: A Graph Neural Network for Large-Scale Quantum Calculations

With massive amounts of atomic simulation data available, there is a hug...

Rotation Invariant Graph Neural Networks using Spin Convolutions

Progress towards the energy breakthroughs needed to combat climate chang...

Catalyst Property Prediction with CatBERTa: Unveiling Feature Exploration Strategies through Large Language Models

Efficient catalyst screening necessitates predictive models for adsorpti...

How Do Graph Networks Generalize to Large and Diverse Molecular Systems?

The predominant method of demonstrating progress of atomic graph neural ...

Scientific Computing Algorithms to Learn Enhanced Scalable Surrogates for Mesh Physics

Data-driven modeling approaches can produce fast surrogates to study lar...

Web Image Context Extraction with Graph Neural Networks and Sentence Embeddings on the DOM tree

Web Image Context Extraction (WICE) consists in obtaining the textual in...

Please sign up or login with your details

Forgot password? Click here to reset