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: This paper proposes a 3D pretraining method for molecular property prediction. As 3D information is infeasible to compute at the scale required by real-world applications, this paper reasons about the geometry of molecules given only their 2D molecular graphs.   During pretraining with molecules whose 3D information is known beforehand, this paper uses a 2D GNN to encode these molecules. Then, it maximizes the mutual information between 3D summary vectors and the encoded representations for injecting 3D information into the representations. During fine-tuning, the model can take 2D molecules as input. The pretraining phase ensures that the representations during fine-tuning contain latent 3D information.   The paper pre-trains on three datasets, QM9, GEOM-Drugs, and QMugs, and tests on both quantum mechanical properties (QM9 and GEOM-Drugs) and non-quantum properties (10 datasets, e.g. HIV and BACE). The paper claims that significant improvements are obtained for quantum mechanical properties. Also, the method does not suffer from the negative transfer.