# Extreme Acceleration of Graph Neural Network-based Prediction Models for Quantum Chemistry

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 71 references that are in the map.
- First posted on arXiv 25 Nov 2022.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [How Powerful are Graph Neural Networks?](https://mlsanity.com/paper/1810.00826/) (2018)
- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials](https://mlsanity.com/paper/2101.03164/) (2021)
- [In-Datacenter Performance Analysis of a Tensor Processing Unit](https://mlsanity.com/paper/1704.04760/) (2017)
- [GraphSAINT: Graph Sampling Based Inductive Learning Method](https://mlsanity.com/paper/1907.04931/) (2019)
- [Directional Message Passing for Molecular Graphs](https://mlsanity.com/paper/2003.03123/) (2020)
- [Cluster-GCN: An Efficient Algorithm for Training Deep and Large Graph Convolutional Networks](https://mlsanity.com/paper/1905.07953/) (2019)
- [The Open Catalyst 2020 (OC20) Dataset and Community Challenges](https://mlsanity.com/paper/2010.09990/) (2020)
- [Deep Graph Library: A Graph-Centric, Highly-Performant Package for Graph Neural Networks](https://mlsanity.com/paper/1909.01315/) (2019)
- [Molecular Transformer - A Model for Uncertainty-Calibrated Chemical Reaction Prediction](https://mlsanity.com/paper/1811.02633/) (2018)
- [DistDGL: Distributed Graph Neural Network Training for Billion-Scale Graphs](https://mlsanity.com/paper/2010.05337/) (2020)
- [DGL-KE: Training Knowledge Graph Embeddings at Scale](https://mlsanity.com/paper/2004.08532/) (2020)
- [GNNAdvisor: An Adaptive and Efficient Runtime System for GNN Acceleration on GPUs](https://mlsanity.com/paper/2006.06608/) (2020)
- [Reducing Communication in Graph Neural Network Training](https://mlsanity.com/paper/2005.03300/) (2020)
- [Dissecting the Graphcore IPU Architecture via Microbenchmarking](https://mlsanity.com/paper/1912.03413/) (2019)
- [Molecular Mechanics-Driven Graph Neural Network with Multiplex Graph for Molecular Structures](https://mlsanity.com/paper/2011.07457/) (2020)
- [PipeGCN: Efficient Full-Graph Training of Graph Convolutional Networks with Pipelined Feature Communication](https://mlsanity.com/paper/2203.10428/) (2022)
- [Accelerating Training and Inference of Graph Neural Networks with Fast Sampling and Pipelining](https://mlsanity.com/paper/2110.08450/) (2021)
- [Hierarchical modeling of molecular energies using a deep neural network](https://mlsanity.com/paper/1710.00017/) (2017)
- [Distributed Hybrid CPU and GPU training for Graph Neural Networks on Billion-Scale Graphs](https://mlsanity.com/paper/2112.15345/) (2021)
- [Large-scale graph representation learning with very deep GNNs and self-supervision](https://mlsanity.com/paper/2107.09422/) (2021)
- [HydroNet: Benchmark Tasks for Preserving Intermolecular Interactions and Structural Motifs in Predictive and Generative Models for Molecular Data](https://mlsanity.com/paper/2012.00131/) (2020)
- [Operation-Level Performance Benchmarking of Graph Neural Networks for Scientific Applications](https://mlsanity.com/paper/2207.09955/) (2022)

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

- [Towards Foundational Models for Molecular Learning on Large-Scale Multi-Task Datasets](https://mlsanity.com/paper/2310.04292/) (2023)
- [BitGNN: Unleashing the Performance Potential of Binary Graph Neural Networks on GPUs](https://mlsanity.com/paper/2305.02522/) (2023)
- [PopSparse: Accelerated block sparse matrix multiplication on IPU](https://mlsanity.com/paper/2303.16999/) (2023)

## Abstract (arXiv)

Molecular property calculations are the bedrock of chemical physics. High-fidelity textitab initio modeling techniques for computing the molecular properties can be prohibitively expensive, and motivate the development of machine-learning models that make the same predictions more efficiently. Training graph neural networks over large molecular databases introduces unique computational challenges such as the need to process millions of small graphs with variable size and support communication patterns that are distinct from learning over large graphs such as social networks. This paper demonstrates a novel hardware-software co-design approach to scale up the training of graph neural networks for molecular property prediction. We introduce an algorithm to coalesce the batches of molecular graphs into fixed size packs to eliminate redundant computation and memory associated with alternative padding techniques and improve throughput via minimizing communication. We demonstrate the effectiveness of our co-design approach by providing an implementation of a well-established molecular property prediction model on the Graphcore Intelligence Processing Units (IPU). We evaluate the training performance on multiple molecular graph databases with varying degrees of graph counts, sizes and sparsity. We demonstrate that such a co-design approach can reduce the training time of such molecular property prediction models from days to less than two hours, opening new possibilities for AI-driven scientific discovery.