Fast and Accurate Predictions from 3D Molecular Structures

Nature is fond of geometry, and different molecular structures show unique geometric traits.  However, training neural networks to predict properties of such geometric structures comes with a unique set of challenges.   Works presented here detail numerous contributions toward solving those distinct problems, including: 1) introducing one of the largest machine learning (ML) benchmark datasets for chemistry; 2) presenting graph-theoretic approaches for both development and interpretation of neural network models; 3) providing reference implementations via message-passing frameworks, such as PyTorch Geometric; and 4) exploring co-design of novel software-hardware architectures for rapid training of these models.

These neural networks require approaches that are cognizant of the connectivity structure as well as the 3D coordinates of atoms.  This is addressed by a family of graph neural networks known as "3D-GNNs".  

As molecular structures grow larger, the number of possible input states for 3D-GNNs increases exponentially.  This presents a challenge for generalizability of the models and requires innovation to answer questions, such as can a model trained over a region of the chemical space be applied to another?  How do we push the limits of today's ML frameworks and hardware architectures to train these models in hours?

We are a team of researchers from national laboratories, academia and industry interested in answering these questions.  Our team hails from Pacific Northwest National Laboratory (USA), Argonne National Laboratory (USA), Graphcore (UK), IBM Research (USA) and the University of Washington (USA).  Check out our recent work!