3D Pre-training improves GNNs for Molecular Property Prediction

Molecular property prediction is one of the fastest-growing applications of deep learning with critical real-world impacts . Including 3D molecular structure as input to learned models improves their performance for many molecular tasks . However , this information is infeasible to compute at the scale required by several real-world applications . We propose pre-training a model to reason about the geometry of molecules given only their 2D molecular graphs . Using methods from self-supervised learning , we maximize the mutual information between 3D summary vectors and the representations of a Graph Neural Network ( GNN ) such that they contain latent 3D information . During fine-tuning on molecules with unknown geometry , the GNN still produces implicit 3D information and can use it to improve downstream tasks . We show that 3D pre-training provides significant improvements for a wide range of properties , such as a 22 % average MAE reduction on eight quantum mechanical properties . Moreover , the learned representations can be effectively transferred between datasets in different molecular spaces . 1 INTRODUCTION . The understanding of molecular and quantum chemistry is a rapidly growing area for deep learning with models having direct real-world impacts in quantum chemistry ( Dral , 2020 ) , protein structure prediction ( Jumper et al. , 2021 ) , materials science ( Schmidt et al. , 2019 ) , and drug discovery ( Stokes et al. , 2020 ) . In particular , for the task of molecular property prediction , GNNs have had great success ( Yang et al. , 2019 ) . GNNs operate on the molecular graph by updating each atom ’ s representation based on the atoms connected to it via covalent bonds . However , these models reason poorly about other important interatomic forces that depend on the atoms ’ relative positions in space . Previous works showed that using the atoms ’ 3D coordinates in space improves the accuracy of molecular property prediction ( Schütt et al. , 2017 ; Klicpera et al. , 2020b ; Liu et al. , 2021 ; Klicpera et al. , 2021 ) . However , using classical molecular dynamics simulations to explicitly compute a molecule ’ s geometry before predicting its properties is computationally intractable for many real-world applications . Even recent Machine Learning ( ML ) methods for conformation generation ( Xu et al. , 2021b ; Shi et al. , 2021 ; Ganea et al. , 2021 ) are still too slow for large-scale applications . Our Solution : 3D Infomax We pre-train a GNN to encode implicit 3D information in its latent vectors using publicly available molecular structures . A GNN is pre-trained by maximizing the mutual information ( MI ) between its embedding of a 2D molecular graph and a representation capturing the 3D information that is produced by a separate network . This way , the GNN learns to produce latent 3D information using only the information given by the 2D molecular graphs . After pre-training , the weights can be transferred and fine-tuned on molecular datasets where no 3D information is available . For those molecules , the GNN is still able to produce implicit 3D information that can be used to inform property predictions . Several other self-supervised learning ( SSL ) methods that do not use 3D information have been proposed and evaluated to pre-train GNNs and obtain better property predictions after fine-tuning ( Hu et al. , 2020b ; You et al. , 2020 ; Xu et al. , 2021a ) . These often rely on augmentations ( such as removing atoms ) that significantly alter the molecules while assuming that their properties do not change . Meanwhile , 3D Infomax pre-training teaches the model to reason about how atoms interact in space , which is a principled and generalizable form of information . We analyze our method ’ s performance by pre-training with multiple 3D datasets before evaluating on ten quantum mechanical properties and ten datasets with biological , pharmacological , or chemical properties . 3D Infomax improves property predictions by large margins and the learned representations are highly generalizable : significant improvements are obtained even when the molecular space of the pre-training dataset is vastly different ( e.g. , in size ) from the kinds of molecules in the downstream tasks . While conventional pre-training methods sometimes suffer from negative transfer ( Pan & Yang , 2010 ) , i.e. , a decrease in performance , this is not observed for 3D Infomax . Our main contributions are : . • A 3D pre-training method that enables GNNs to reason about the geometry of molecules given only their 2D molecular graphs , which improves property predictions . • Experiments showing that our learned representations are meaningful for various quantum mechanical , chemical , biological , or pharmacological tasks , without negative transfer . • Empirical evidence that the embeddings generalize across different molecular spaces . • An approach to leverage information from multiple molecular conformers that further im- proves downstream property predictions and an evaluation to what extent this is possible . 2 BACKGROUND . Molecular Conformers We pre-train models to learn 3D information and transfer it to downstream molecular property prediction tasks . This 3D information is not given by a single set of coordinates . For a given molecular graph there are multiple conformers , i.e. , probable arrangements of the atoms , which can lead to different chemical properties . It is usual to consider only the conformers of lower energy since they have a higher probability of naturally occurring . Several tools exist to compute conformers ranging from methods based on classical force fields to slower but more accurate molecular dynamics simulations . Methods such as RDKit ’ s ETKDG algorithm ( Landrum , 2016 ) are fast but in our experiments we find that their less accurate 3D information does not necessarily improve predictions . The popular metadynamics method CREST ( Grimme , 2019 ) offers a good tradeoff between speed and accuracy but still requires about 6 hours per druglike molecule per CPU-core ( Axelrod & Gomez-Bombarelli , 2020 ) . This highlights the need to capture 3D information without explicitly computing structures , especially for drug-discovery screening datasets comprising of millions or billions of molecules ( Gorgulla et al. , 2020 ) . Symmetries of Molecules A molecule ’ s conformation does not change if all the atom coordinates are jointly translated or rotated around a point , i.e. , molecules are symmetric with respect to these two types of transformations which is also called SE ( 3 ) symmetry . Note that some molecules ( called chiral ) are not invariant to reflections : their properties depend on their chirality . Deep learning architectures that capture these symmetries are usually more sample efficient , and they generalize to all symmetric inputs the architecture has been designed for ( Bronstein et al. , 2021 ) . In our method , the produced representations of the 3D structure respect these symmetries of molecules . Scaffold Split The preferred way to evaluate molecular models is to use a scaffold split when generating the train-test sets , such that the molecules from the test set do not share a scaffold with those of the training set . This helps to avoid overestimating the generalization power ( Yang et al. , 2019 ) since ML models tend to memorize and overfit these structures . We use the common Bemis-Murcko scaffold ( Bemis & Murcko , 1996 ) ( see Figure 5 in Appendix A ) . Molecules that share identical scaffolds are put into the same set , i.e. , each scaffold goes into either the train , validation or test set . Graph Neural Networks We make use of GNNs to predict molecular properties given a molecular graph . Many GNNs can be described in the framework of Message Passing Neural Networks ( MPNNs ) ( Gilmer et al. , 2017 ) , such as the PNA model ( Corso et al. , 2020 ) which we employ . The aim of MPNNs is to learn a representation of a graph G = ( V , E ) with vertices V connected by edges E . They do so by iteratively applying message-passing layers and then combining all vertex representations in a readout function . A message-passing layer first creates messages for each edge based on the vertices it connects , then each vertex representation is updated by aggregating the messages of all connected edges and combining them with the previous layer ’ s representation . The messages are usually created by multi-layer perceptrons ( MLPs ) and are aggregated via permutation invariant functions such as taking their mean , max , or sum . After the message-passing layers , another permutation invariant function is used as readout to obtain a final graph level representation . 3 RELATED WORK . Molecular property predictions While ours is the first work on pre-training GNNs for molecular property prediction using 3D information , it heavily draws from the fields of SSL and ML for molecules . An important milestone for the latter was Gilmer et al . ( 2017 ) introducing MPNNs after which GNNs became popular for quantum chemistry ( Brockschmidt , 2020 ; Tang et al. , 2020 ; Withnall et al. , 2020 ) , drug discovery ( Li et al. , 2017 ; Stokes et al. , 2020 ; Torng & Altman , 2019 ) , and molecular property prediction in general ( Coley et al. , 2019 ; Hy et al. , 2018 ; Unke & Meuwly , 2019 ) . The field is well established with easily accessible molecular datasets driving progress ( Wu et al. , 2017 ; Hu et al. , 2020a ) and rigorous evaluations of MPNNs for property prediction ( Yang et al. , 2019 ) showing the effectiveness of the approach . While these GNNs have had great successes by operating on the 2D graph , many tasks on molecules can be improved by additionally using 3D information . A simple approach is to use bond lengths as edge features ( Chen et al. , 2020a ) , but methods that capture more molecular geometry improve on this such as SchNet ( Schütt et al. , 2017 ) . Similarly , DimeNet ( Klicpera et al. , 2020b ; a ) proposed extracting more 3D information via bond angles , which further improved quantum mechanical property prediction . Spherical Message Passing ( SMP ) ( Liu et al. , 2021 ) included another angular quantity , and GemNet ( Klicpera et al. , 2021 ) developed an approach to also capture torsion angles , such that all relative atom positions are uniquely defined . Equivariant Graph Neural Networks ( EGNN ) ( Satorras et al. , 2021 ) achieved the same by operating on all pairwise atom distances . Self-Supervised Learning attempts to find supervision signals in unlabelled data to learn meaningful representations . In particular , contrastive learning ( van den Oord et al. , 2018 ; Gutmann & Hyvärinen , 2010 ; Belghazi et al. , 2018 ; Hjelm et al. , 2019 ) is a popular class of methods . These learn representations by comparing the embeddings of similar and dissimilar inputs and have achieved impressive results in computer vision ( Chen et al. , 2020b ; Caron et al. , 2020 ) . Learning from unlabeled data also is a critical challenge in molecular chemistry since datasets are relatively small due to experimental costs . Several works have explored contrastive learning variants in the context of molecular graphs for non-quantum molecular properties ( Hu et al. , 2020b ; Wang et al. , 2021 ; You et al. , 2020 ; 2021 ; Xu et al. , 2021a ) . The improvements these methods provided in molecular property prediction are still limited and often fail to generalize .

Summary: The authors present 3D Infomax, a graph neural network (GNN) pre-training solution that leverages 3D information to generate better learned embeddings and improve performance on down-stream prediction tasks where 3D information would be useful but not easily obtainable.  The approach is useful for a range of downstream tasks involving molecules, including ones that are quantum mechnical, biological, and pharmacological in nature.  They also demonstrate that the use of multiple 3D conformations (thereby encoding the inherent flexibility of molecules) further improves performance.