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

by   Hatem Helal, et al.

Molecular property calculations are the bedrock of chemical physics. High-fidelity ab 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.


page 3

page 4

page 9

page 17


Gated Graph Recursive Neural Networks for Molecular Property Prediction

Molecule property prediction is a fundamental problem for computer-aided...

On Graph Neural Network Ensembles for Large-Scale Molecular Property Prediction

In order to advance large-scale graph machine learning, the Open Graph B...

Multitask Learning On Graph Neural Networks Applied To Molecular Property Predictions

Prediction of molecular properties, including physico-chemical propertie...

Retention Time Prediction for Chromatographic Enantioseparation by Quantile Geometry-enhanced Graph Neural Network

A new research framework is proposed to incorporate machine learning tec...

ProGReST: Prototypical Graph Regression Soft Trees for Molecular Property Prediction

In this work, we propose the novel Prototypical Graph Regression Self-ex...

Chemical Property Prediction Under Experimental Biases

The ability to predict the chemical properties of compounds is crucial i...

Improving Molecular Graph Neural Network Explainability with Orthonormalization and Induced Sparsity

Rationalizing which parts of a molecule drive the predictions of a molec...

Please sign up or login with your details

Forgot password? Click here to reset