mlsanityData to 14 Oct 2026

Towards Training Billion Parameter Graph Neural Networks for Atomic Simulations

arXiv 2203.09697 ↗First posted 18 Mar 2022cs.LG physics.comp-ph stat.MLcited by 25

Not on the map. Why?

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.

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

23 of 37 references.

References · 23

  1. 2014Adam: A Method for Stochastic Optimizationcited by 37,240
  2. 2020Language Models are Few-Shot Learnerscited by 32,419
  3. 2017Neural Message Passing for Quantum Chemistrycited by 3,487
  4. 2019Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelismcited by 1,868
  5. 2018Graph Neural Networks: A Review of Methods and Applicationscited by 1,576
  6. 2019Energy and Policy Considerations for Deep Learning in NLPcited by 1,254
  7. 2021Scaling Vision Transformerscited by 999
  8. 2017SchNet: A continuous-filter convolutional neural network for modeling quantum interactionscited by 571
  9. 2020Directional Message Passing for Molecular Graphscited by 561
  10. 2019Quantifying the Carbon Emissions of Machine Learningcited by 521
  11. 2021GemNet: Universal Directional Graph Neural Networks for Moleculescited by 343
  12. 2020Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Moleculescited by 255
11 more
  1. 2020The Open Catalyst 2020 (OC20) Dataset and Community Challengescited by 246
  2. 2021Spherical Message Passing for 3D Graph Networkscited by 168
  3. 2019AliGraph: A Comprehensive Graph Neural Network Platformcited by 153
  4. 2020DistDGL: Distributed Graph Neural Network Training for Billion-Scale Graphscited by 139
  5. 2020OrbNet: Deep Learning for Quantum Chemistry Using Symmetry-Adapted Atomic-Orbital Featurescited by 64
  6. 2021ForceNet: A Graph Neural Network for Large-Scale Quantum Calculationscited by 58
  7. 2021Rotation Invariant Graph Neural Networks using Spin Convolutionscited by 52
  8. 2020Reducing Communication in Graph Neural Network Trainingcited by 50
  9. 2020An Introduction to Electrocatalyst Design using Machine Learning for Renewable Energy Storagecited by 45
  10. 2020Accelerating Graph Sampling for Graph Machine Learning using GPUscited by 29
  11. 2018Neural Message Passing with Edge Updates for Predicting Properties of Molecules and Materialscited by 21

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2022: 7; 2023: 6; 2024: 2; 2025: 7; 2026: 3.

25 papers.

Sort:

By area:

  1. 2022Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphscited by 264
  2. 2023EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representationscited by 204
  3. 2022Pre-training via Denoising for Molecular Property Predictioncited by 124
  4. 2025UMA: A Family of Universal Models for Atomscited by 105
  5. 2022The Open Catalyst 2022 (OC22) Dataset and Challenges for Oxide Electrocatalystscited by 95
  6. 2022GemNet-OC: Developing Graph Neural Networks for Large and Diverse Molecular Simulation Datasetscited by 61
  7. 2022Polyhistor: Parameter-Efficient Multi-Task Adaptation for Dense Vision Taskscited by 52
  8. 2023From Molecules to Materials: Pre-training Large Generalizable Models for Atomic Property Predictioncited by 39
  9. 2023LinGCN: Structural Linearized Graph Convolutional Network for Homomorphically Encrypted Inferencecited by 27
  10. 2022Towards Sparsification of Graph Neural Networkscited by 18
  11. 2023The Open DAC 2023 Dataset and Challenges for Sorbent Discovery in Direct Air Capturecited by 18
  12. 2022The Open MatSci ML Toolkit: A Flexible Framework for Machine Learning in Materials Sciencecited by 14
13 more
  1. 2023Accelerating Molecular Graph Neural Networks via Knowledge Distillationcited by 11
  2. 2025Efficient Distributed Optimization under Heavy-Tailed Noisecited by 7
  3. 2025The Open DAC 2025 Dataset for Sorbent Discovery in Direct Air Capturecited by 7
  4. 2025DistMLIP: A Distributed Inference Platform for Machine Learning Interatomic Potentialscited by 5
  5. 2023DR-Label: Improving GNN Models for Catalysis Systems by Label Deconstruction and Reconstructioncited by 4
  6. 2025Data-driven particle dynamics: Structure-preserving coarse-graining for emergent behavior in non-equilibrium systemscited by 3
  7. 2026Enhancing Molecular Property Predictions by Learning from Bond Modelling and Interactionscited by 3
  8. 2026Speculative Sampling For Faster Molecular Dynamicscited by 1
  9. 2024Lightweight Geometric Deep Learning for Molecular Modelling in Catalyst Discoverycited by 0
  10. 2024Distribution Learning for Molecular Regressioncited by 0
  11. 2025Distributed Equivariant Graph Neural Networks for Large-Scale Electronic Structure Predictioncited by 0
  12. 2025Layer-to-Layer Knowledge Mixing in Graph Neural Network for Chemical Property Predictioncited by 0
  13. 2026Breaking the Training Barrier of Billion-Parameter Universal Machine Learning Interatomic Potentialscited by 0