# Towards Training Billion Parameter Graph Neural Networks for Atomic Simulations

arXiv: https://arxiv.org/abs/2203.09697 · HTML: https://mlsanity.com/paper/2203.09697/

## Key facts

- Cited by 25 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 37 references that are in the map.
- First posted on arXiv 18 Mar 2022.

## Builds on (most cited first, 23 of 23)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Graph Neural Networks: A Review of Methods and Applications](https://mlsanity.com/paper/1812.08434/) (2018)
- [Energy and Policy Considerations for Deep Learning in NLP](https://mlsanity.com/paper/1906.02243/) (2019)
- [Scaling Vision Transformers](https://mlsanity.com/paper/2106.04560/) (2021)
- [SchNet: A continuous-filter convolutional neural network for modeling quantum interactions](https://mlsanity.com/paper/1706.08566/) (2017)
- [Directional Message Passing for Molecular Graphs](https://mlsanity.com/paper/2003.03123/) (2020)
- [Quantifying the Carbon Emissions of Machine Learning](https://mlsanity.com/paper/1910.09700/) (2019)
- [GemNet: Universal Directional Graph Neural Networks for Molecules](https://mlsanity.com/paper/2106.08903/) (2021)
- [Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Molecules](https://mlsanity.com/paper/2011.14115/) (2020)
- [The Open Catalyst 2020 (OC20) Dataset and Community Challenges](https://mlsanity.com/paper/2010.09990/) (2020)
- [Spherical Message Passing for 3D Graph Networks](https://mlsanity.com/paper/2102.05013/) (2021)
- [AliGraph: A Comprehensive Graph Neural Network Platform](https://mlsanity.com/paper/1902.08730/) (2019)
- [DistDGL: Distributed Graph Neural Network Training for Billion-Scale Graphs](https://mlsanity.com/paper/2010.05337/) (2020)
- [OrbNet: Deep Learning for Quantum Chemistry Using Symmetry-Adapted Atomic-Orbital Features](https://mlsanity.com/paper/2007.08026/) (2020)
- [ForceNet: A Graph Neural Network for Large-Scale Quantum Calculations](https://mlsanity.com/paper/2103.01436/) (2021)
- [Rotation Invariant Graph Neural Networks using Spin Convolutions](https://mlsanity.com/paper/2106.09575/) (2021)
- [Reducing Communication in Graph Neural Network Training](https://mlsanity.com/paper/2005.03300/) (2020)
- [An Introduction to Electrocatalyst Design using Machine Learning for Renewable Energy Storage](https://mlsanity.com/paper/2010.09435/) (2020)
- [Accelerating Graph Sampling for Graph Machine Learning using GPUs](https://mlsanity.com/paper/2009.06693/) (2020)
- [Neural Message Passing with Edge Updates for Predicting Properties of Molecules and Materials](https://mlsanity.com/paper/1806.03146/) (2018)

## Built on by (most cited first, 25 of 25)

- [Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs](https://mlsanity.com/paper/2206.11990/) (2022)
- [EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations](https://mlsanity.com/paper/2306.12059/) (2023)
- [Pre-training via Denoising for Molecular Property Prediction](https://mlsanity.com/paper/2206.00133/) (2022)
- [UMA: A Family of Universal Models for Atoms](https://mlsanity.com/paper/2506.23971/) (2025)
- [The Open Catalyst 2022 (OC22) Dataset and Challenges for Oxide Electrocatalysts](https://mlsanity.com/paper/2206.08917/) (2022)
- [GemNet-OC: Developing Graph Neural Networks for Large and Diverse Molecular Simulation Datasets](https://mlsanity.com/paper/2204.02782/) (2022)
- [Polyhistor: Parameter-Efficient Multi-Task Adaptation for Dense Vision Tasks](https://mlsanity.com/paper/2210.03265/) (2022)
- [From Molecules to Materials: Pre-training Large Generalizable Models for Atomic Property Prediction](https://mlsanity.com/paper/2310.16802/) (2023)
- [LinGCN: Structural Linearized Graph Convolutional Network for Homomorphically Encrypted Inference](https://mlsanity.com/paper/2309.14331/) (2023)
- [Towards Sparsification of Graph Neural Networks](https://mlsanity.com/paper/2209.04766/) (2022)
- [The Open DAC 2023 Dataset and Challenges for Sorbent Discovery in Direct Air Capture](https://mlsanity.com/paper/2311.00341/) (2023)
- [The Open MatSci ML Toolkit: A Flexible Framework for Machine Learning in Materials Science](https://mlsanity.com/paper/2210.17484/) (2022)
- [Accelerating Molecular Graph Neural Networks via Knowledge Distillation](https://mlsanity.com/paper/2306.14818/) (2023)
- [Efficient Distributed Optimization under Heavy-Tailed Noise](https://mlsanity.com/paper/2502.04164/) (2025)
- [The Open DAC 2025 Dataset for Sorbent Discovery in Direct Air Capture](https://mlsanity.com/paper/2508.03162/) (2025)
- [DistMLIP: A Distributed Inference Platform for Machine Learning Interatomic Potentials](https://mlsanity.com/paper/2506.02023/) (2025)
- [DR-Label: Improving GNN Models for Catalysis Systems by Label Deconstruction and Reconstruction](https://mlsanity.com/paper/2303.02875/) (2023)
- [Data-driven particle dynamics: Structure-preserving coarse-graining for emergent behavior in non-equilibrium systems](https://mlsanity.com/paper/2508.12569/) (2025)
- [Enhancing Molecular Property Predictions by Learning from Bond Modelling and Interactions](https://mlsanity.com/paper/2603.00568/) (2026)
- [Speculative Sampling For Faster Molecular Dynamics](https://mlsanity.com/paper/2606.02455/) (2026)
- [Lightweight Geometric Deep Learning for Molecular Modelling in Catalyst Discovery](https://mlsanity.com/paper/2404.10003/) (2024)
- [Distribution Learning for Molecular Regression](https://mlsanity.com/paper/2407.20475/) (2024)
- [Distributed Equivariant Graph Neural Networks for Large-Scale Electronic Structure Prediction](https://mlsanity.com/paper/2507.03840/) (2025)
- [Layer-to-Layer Knowledge Mixing in Graph Neural Network for Chemical Property Prediction](https://mlsanity.com/paper/2510.20236/) (2025)
- [Breaking the Training Barrier of Billion-Parameter Universal Machine Learning Interatomic Potentials](https://mlsanity.com/paper/2604.15821/) (2026)

## Abstract (arXiv)

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.