3D P RE - TRAINING IMPROVES GNN S FOR M OLECULAR P ROPERTY P REDICTION
Anonymous authors
Paper under double-blind review
A BSTRACT
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 ﬁne-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 signiﬁcant 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.
I NTRODUCTION
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 (Sch ¨ utt 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 ﬁne-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 ﬁne-tuning (Hu et al., 2020b; You et al., 2020; Xu et al., 2021a).
These often rely on augmentations (such as removing atoms) that signiﬁcantly 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: signiﬁcant 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 improves downstream property predictions, and an evaluation to what extent this is possible.
B ACKGROUND
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 ﬁelds to slower but more accurate molecular dynamics simulations.
Methods such as RDKit’s ETKDG algorithm (Landrum, 2016) are fast but in our experiments we ﬁnd 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 reﬂections: their properties depend on their chirality.
Deep learning architectures that capture these symmetries are usually more sample efﬁcient, 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 overﬁt 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 ﬁrst 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 ﬁnal graph level representation.
R ELATED W ORK
Molecular property predictions While ours is the ﬁrst work on pre-training GNNs for molecular property prediction using 3D information, it heavily draws from the ﬁelds 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 ﬁeld 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 (Sch¨utt 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 deﬁned.
Equivariant Graph Neural Networks (EGNN) (Satorras et al., 2021) achieved the same by operating on all pairwise atom distances.
Self-Supervised Learning attempts to ﬁnd supervision signals in unlabelled data to learn meaningful representations.
In particular, contrastive learning (van den Oord et al., 2018; Gutmann &
Hyv¨arinen, 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.
M ETHOD
Our goal is a method that produces latent 3D information from a 2D molecular graph and use it for more accurate molecular property predictions.
For this purpose, we propose to pre-train a GNN to produce latent representations of a molecule’s 3D geometry.
Through this pre-training, the network learns the connection between the 2D information of molecular graphs and their 3D conformers.
After pre-training, we transfer the weights and ﬁne-tune them on property prediction tasks as visualized in Figure 2.
The GNN’s produced 3D information can be used to improve predictions.
The following presents our 3D Infomax pre-training before explaining the baselines we compare with.
4. 3D I NFOMAX P RE -
[Figure]
3D Infomax uses two different models, as visualized in Figure 2.
Firstly, the model that should be pre-trained which we call 2D network f a since its inputs are 2D molecular graphs G =
( V , E ) with atoms V and bonds E from which it produces a representation f a ( G ) =
z a ∈ R d z .
Secondly, the 3D network f b which encodes the atoms’ 3D coordinates R = { r v } v ∈V in a 3D representation f b ( R ) =
z b ∈ R d z .
Our pre-training can also be understood from a knowledge distillation perspective where the student 2D network learns from the teacher 3D network to produce 3D information.
Contrastive Framework To teach the 2D network f a to produce 3D information from the 2D graph inputs we maximize the mutual information between the latent 2D representations z a and 3D representations z b .
Intuitively, we wish to maximize the agreement between z a and z b if they are derived from the same molecule.
For this purpose, we use contrastive learning (visualized in
Figure 3).
We consider a batch of N molecular graphs {G i } i ∈{ 1 ...N } with their atom coordinates { R i } i ∈{ 1 ...N } from which the networks produce multiple representations z ai and z bi .
The ﬁrst objective of contrastive learning is to maximize the representations’ similarity if they are a positive pair, meaning that they come from the same molecule (same index i ).
The second objective is to enforce dissimilarity between negative pairs z ai and z bk where i (cid:54)
=
k , i.e., the 2D and 3D representations in the batch should be dissimilar if they come from different molecules.
These objectives are captured in the popular NTXent loss (Chen et al., 2020b) and we use a similar loss to jointly optimize our models:
[Equation]
where sim ( z a , z b ) =
z a · z b / ( (cid:107) z a (cid:107)(cid:107) z b (cid:107) ) is the cosine similarity and τ is a temperature parameter which can be seen as weight for the most similar negative pair.
While different combinations of contrastive losses and SSL are possible to learn a joint embedding space between 2D and 3D representations, we found the above loss to perform best.
Other methods (Barlow Twins (Zbontar et al., 2021), BYOL (Grill et al., 2020), VICReg (Bardes et al., 2021)) are explored in Appendix C.4.
Multiple Conformers For most molecules, there are multiple low-energy stable conformers.
Instead of only using the most probable conformer (with the lowest energy), we found that leveraging structural information from multiple conformers provides signiﬁcant beneﬁts.
To achieve this, we now consider the c highest probability conformers { R ji } j ∈{ 1 ...c } of the i-th molecule.
If there are fewer than c conformers for a molecule, the lowest energy conformer is repeated.
Our choice for the following approach is justiﬁed by its good trade-off between simplicity and performance in the comparisons with other possible methods in Appendix C.2.
For every molecule the 3D network now takes all conformers as input and produces their latent 3D representations { z bi,j } j ∈{ 1 ...c } .
The objective is to maximize the similarity between z ai and all conformer representations z bi,j that stem from the same molecule (see Figure 3).
As such we modify our loss to sum over the similarities of all conformers to obtain the ﬁnal loss:
[Equation]
3D Network
As the 2D network is an arbitrary GNN that should be pre-trained for which we choose
Principal Neighborhood Aggregation (PNA) (Corso et al., 2020), the last missing part of our method is the 3D network.
Its purpose is to encode as much 3D information as possible into the vector z b and it does not have access to the 2D information used by the 2D network such as atom or bond features.
Otherwise, the mutual information could be increased by both networks encoding this information instead of the desired 3D information.
Therefore, the 3D network only takes the atom coordinates as input.
Our concrete architecture encodes the 3D information given by the pairwise
Euclidean distances of all atoms.
This representation uniquely deﬁnes all relative atom positions and is invariant to translation and rotation, as desired, but also to reﬂection, which is a disadvantage since molecular properties can change under reﬂections (chiral molecules).
Using all pairwise distances also means that the method’s complexity is quadratic in the number of atoms but this is feasible for drug-like molecules.
The pairwise distances d uv between atoms u and v are ﬁrst mapped to a higher dimensional space using high-frequency functions.
This motivation for this is to enable deep networks to better ﬁt data with high-frequency variation (Rahaman et al., 2019; Tancik et al., 2020).
This scenario is present in our case with small differences in bond lengths.
These bond distances and their small variations might be the most important ones assuming that close-by atoms have the most relevant interactions.
As such, we use the following mapping γ : R (cid:55)→ R 2 F +1 with the number of frequencies F set to 4 :
[Equation]
The further components can be seen as an MPNN (Gilmer et al., 2017) operating on the fully connected graph of a molecule with the encoded distances as edge features and a constant learned vector as node features.
We use these initial node representations instead of atom features (as in the 2D network) such that the mutual information cannot be increased by both networks only encoding the atom feature information.
Instead, 3D information has to be captured to solve the objective.
The message passing layers iteratively encode the 3D information into the node features, which are pooled to produce the 3D representation z b .
The differences to standard MPNNs are detailed in
Appendix A. Instead of the presented architecture, a 3D GNN such as SMP (Liu et al., 2021) operating on learned node embeddings could also be used.
We justify choosing our architecture with experiments in Appendix C.1 that also show the γ mapping’s effectiveness by ablating it.
4. P RE - TRAINING B ASELINES
Distance Predictor A simpler method to 3D pre-train a GNN instead of 3D Infomax is by directly predicting all atom distances.
To predict the distance between node v and u , we concatenate their representations h u , h v ∈ R d h that were produced by the GNN and feed them to an MLP that produces a single scalar U : R 2 d h
(cid:55)→ R .
The distance prediction dist uv is then given by
[Equation]
where (cid:107)
Conformer Generation As a second 3D pre-training alternative to 3D Infomax, we pre-train a GNN by generating molecular conformers via the state-of-the-art (SOTA) method GeoMol (Ganea et al., 2021).
To predict correct 3D conformers, the GNN has to encode 3D information in its hidden layers which can potentially be transferred and used to inform downstream property predictions.
GraphCL We compare against the conventional augmentation-based pre-training method Graph
Contrastive Learning (GraphCL) (You et al., 2020) with the settings of JOAO (You et al., 2021) since it outperformed other SSL approaches for multiple molecular tasks.
It uses a common selfsupervised objective in which the model has to learn to produce representations that are invariant to augmentations.
We use randomly dropping nodes with a ratio of 0 .
2 on both branches of the SSL setup since JOAO found this combination of augmentations to work particularly well for molecules.
E XPERIMENTS
Data and Setup For pre-training, we use three datasets of molecules with 3D conformer information: QM9 (Ramakrishnan et al., 2014) which contains 134 k small molecules (18 atoms on average) with a single conformer, GEOM-Drugs (Axelrod & Gomez-Bombarelli, 2020) with 304 k molecules and QMugs (Isert et al., 2021) with 665 k. GEOM-Drugs and QMug, both consist of larger drug-like molecules (44.4 and 30.6 atoms on average) with multiple conformers.
For ﬁne-tuning we consider quantum properties on the one side and biological, chemical, and pharmacological properties on the other side.
We predict ten quantum properties of QM9 and GEOM-Drugs (a disjoint half of the datasets if the other half was used for pre-training), for which we employ a random split.
For non-quantum properties, we use ten Open Graph Benchmark (OGB)
(Hu et al., 2020a)
datasets with their standard scaffold splits.
We use OGB’s atom and bond featurization for all datasets.
Details for all used data are in Appendix B.1.
We choose PNA (Corso et al., 2020) as the GNN to pre-train due to its simplicity and SOTA performance for molecular tasks.
The reported conﬁdence intervals are one standard deviation calculated from six random weight initializations, unless stated otherwise.
All baselines we compare with use the same GNN as our 3D Infomax method and all experimental settings are detailed in Appendix B. Code to 3D pre-train a GNN or to reproduce results is available at https: //anonymous.4open.science
/r/3141 .
5. Q UANTUM M ECHANICAL P ROPERTIES
We use 3D Infomax to pre-train three different instances of PNA (1) on 50 k molecules from QM using a single conformer, (2) on 140 k of GEOM-Drugs with 5 conformers and, (3) on 620 k of
QMugs using 3 conformers.
For comparison, we use two different pre-training methods.
These are
GraphCL (You et al., 2020) as described in Section 4.2 and pre-training by predicting the Gibbs free enery of GEOM-Drugs’ pre-training subset (labeled PropPred ).
All pre-training methods use a batch size of 500 .
After pre-training, the models are ﬁne-tuned on 50 k molecules from QM9 (in Table 1) or 140 k from
GEOM-Drugs (in Table 2) that have no overlap with the molecules from the pre-training data.
On the same molecules, we also train PNA with random weight initialization (labeled Rand Init ) to compare how much the downstream performance is improved by the different pre-training methods.
Furthermore, we train and test the 3D GNN SMP (Liu et al., 2021) on the same molecules with 3D coordinates generated by RDKit’s ETKDG algorithm (Landrum, 2016) which can be done in a fast manner (labeled RDKit SMP ).
Lastly, we evaluate SMP using the accurate ground truth 3D conformers of QM9 which were computed with time-consuming simulations that would be infeasible for many real-world applications.
These structures are not available to the other methods.
[Table] [Table]
Table 1 shows that 3D Infomax pre-training leads to large improvements over the randomly initialized baseline and over GraphCL with all three pre-training datasets.
After 3D pre-training on one half of QM9, the average of the MAE decreases by 22%.
Comparing 3D Infomax on GEOM-Drugs with GraphCL shows that even though the latter is pre-trained on two times as many molecules from the same dataset, 3D pre-training is always better by a large margin.
Pre-training with the disjoint half of QM9 performs best since it shares the molecular space of the test set.
Nevertheless, the learned representations also generalize well: pre-training on GEOMDrugs and QMugs leads to improvements of 19% and 18% respectively, even though QM9 contains much smaller molecules with on average 18 atoms compared to the 44.4 atoms for the drug-like molecules of GEOM-Drugs.
[Table]
While 3D Infomax yields large improvements, the
MAE is still substantially higher than that of SMP which uses the 3D information explicitly.
One reason for this is likely that QM9’s properties are conformerspeciﬁc.
There might be a maximum accuracy that can be achieved if only the molecule is known and not for which conformer the property should be predicted.
Nevertheless, this performance gap suggests that there is still room for improvement.
Table 2 further conﬁrms that 3D Infomax substantially improves quantum property predictions and generalizes out-of-distribution.
Our method outperforms GraphCL, even though GraphCL also sees the ﬁne-tuning molecules during pre-training.
Moreover, we observe strong generalization when pre-training with QM9 and ﬁne-tuning on GEOMDrugs.
In this case, the pre-training data only contains the elements C, H, N, O, and F while the target data contains eleven additional elements that are unseen during pre-training.
Such consistent and out-of-distribution improvements can be explained by the type of information captured with 3D Infomax.
Learning to reason about molecular geometry and its impact does not depend on the data’s molecular space and therefore it is not necessary to have a high similarity between the molecules during pre-training and ﬁne-tuning.
Another advantage of 3D Infomax is its comparably fast convergence.
Pre-training on 620 k molecules of QMugs with 3 conformers takes 12 hours, compared to 71 hours for GraphCL on 280 k molecules of GEOM-Drugs.
5. P REDICTIVE 3D PRE - TRAINING
[Table]
3D pre
-training by directly predicting 3D quantities is simpler than 3D Infomax and would be preferable in case of similar gains.
Therefore, we compare with the baselines in Section 4.2 using the same 140 k molecules of GEOM-Drugs for all 3D pre-training methods.
Dist-pred refers to predicting all atom distances of the highest probability conformer and Conf-gen means pre-training by predicting up to 10 conformers.
Table 3 shows that 3D Infomax pre-training is always superior to the predictive baselines, and is the only method to not suffer from negative transfer (Pan & Yang, 2010)
[Table]
5. N ON -Q UANTUM P ROPERTIES
[Table]
In the previous sections
[Table]
we found that 3D Infomax yields large improvements for predicting quantum properties.
For non-quantum properties, there is less empirical evidence that 3D information improves prediction accuracy.
Nevertheless, for tasks such as binding prediction in the bace dataset, we would expect it to be helpful
[Table]
and we compare different methods pre-trained with GEOM-Drugs.
In Table
[Table]
we ﬁnd that 3D Infomax improved performance for 4 out of 10 OGB datasets.
In contrast to the results for quantum mechanical property predictions (Section 5.1)
[Table]
it is not always superior to GraphCL.
However, 3D Infomax never decreases performance which can be valuable in practice and make the method worth employing for non-quantum properties as well.
When investigating for which tasks 3D Infomax is useful, we see that abstract tasks such as predicting clinical test outcomes ( clintox ) beneﬁt less from 3D Infomax.
The most signiﬁcant improvements are rather possible for tasks like predicting solubility and lipophilicity in esol and lipo .
These are more directly related to molecular mechanics and a molecule’s intrinsic quantum properties (e.g., the dipole moment/polarity is important for predicting lipophilicity).
These tasks do not depend on how a molecule will interact with others to result in, e.g., different effects on patients.
Additionally, for datasets like bace with its binding prediction task where 3D information should be valuable, the improvements are only modest.
This could suggest that our method does not capture all of the 3D information that is relevant for predicting protein binding and there is still room for improvement.
Another explanation is that a molecule’s geometry is less helpful for bace since the protein and the binding pocket the molecule has to ﬁt into are not known.
5. N UMBER OF C ONFORMERS AND PRE - TRAINING M OLECULES
Figure 4’s left plot highlights the beneﬁt of using more than a single conformer.
However the marginal gain reduces as higher energy conformers are added and beyond a certain point ( three conformers) the reduced focus on the most likely conformers worsens the downstream performance.
This is in line with the observation that, on average, three conformers are enough to cover 70% of the cumulative Boltzmann weight for GEOM-Drugs.
Additionally, experiments in Appendix
C.2 show that using multiple conformers is essential when pre-training with QMugs: the MAE for
QM9’s homo property is 82 .
57 with a single conformer while it improves to 70 .
77 when using three.
In the right plot of Figure 4 we can observe performance improving as the size of the pre-training dataset increases.
However, the returns are diminishing and we cannot claim that even larger pretraining datasets are likely to drastically improve performance.
C ONCLUSION
We presented a pre-training strategy, 3D Infomax, that teaches a GNN to produce latent 3D and quantum information from 2D molecular graphs.
This can later be used during ﬁne-tuning to improve molecular property predictions while retaining the inference speed of a standard GNN operating on 2D molecular graphs.
We found consistently large improvements ( ∼ 22%) for quantum properties, overshadowing the gains possible with conventional SSL methods.
The embedded 3D knowledge can be transferred across highly different types of molecules (e.g., from molecules with an average of 18 atoms to drug-like molecules with 44.4 atoms), since the representations capture a principled form of information that is known to be useful for molecular tasks.
Similarly, we demonstrated that learned embeddings are transferable across different physical, biological, and pharmaceutical tasks.
Lastly, we observed that using multiple molecular conformers during pre-training provides valuable additional information to further improve downstream property predictions.
E THICS S TATEMENT
We see the most important application and the largest possible positive societal impact of our method in the domain of drug discovery.
Our 3D Infomax pre-training method improves molecular property predictions with implicit 3D information while being fast at inference.
This means that our approach can be used when predicting the properties of vast amounts of molecules to identify potential drug candidates and in large-scale virtual screening.
Applications to materials science could also help in developing new materials for energy storage or energy generation (e.g., solar panels), which are crucial for transitioning to clean energy and preventing climate change.
We believe the greatest possible risks to be side effects such as practitioners blindly trusting molecular property predictions of learned models.
Research on explainability of GNN’s predictions would be crucial for solving this problem.
Additional possible risks come from malicious applications of molecular property predictions, such as developing chemical or biological weapons.
R EPRODUCIBILITY S TATEMENT
All code to reproduce our results is available at https://anonymous.4open.science/r/ 3141 .
This repository also provides an explanation for running the experiments.
Moreover, we detail our hyperparameter search spaces and ﬁnal parameter settings in Appendix B.3 for all our used architectures and baselines.
Furthermore, Appendix B.1 provides the exact datasets we use and their availability with guides for downloading them in our repository.
Additionally, there are additional clarifying explanations for the exact architecture of the 3D network in
Appendix A. Lastly, we state the seeds which we use for the different runs in our repository as well as Appendix B and completely specify our used hard- and software in Appendix B.5.
A F URTHER E XPLANATIONS
3D Network Details
The l-th layer of the 3D network takes two sets as input.
First, n 2 − n edge representations { d luv ∈ R d
d | u, v ∈ V ∧ u (cid:54) = v } (the edges of a complete graph without selfloops).
In the ﬁrst layer they are given by the encoded distances fed through an initial feed-forward network U init :
R 2 F +1 (cid:55)→ R d d which projects them to the hidden dimension of the edges d 0 uv =
U init ( γ ( d uv )) .
The second input is a set of n atom representations { h l 1 , . . .
h ln } with dimensionality R d h .
In the ﬁrst layer, the atom representations are all set to the same learned vector that is initialized with a standard normal.
With (cid:107) meaning concatenation, every layer updates the edge and atom representations and iteratively encodes 3D information into them as follows:
[Equation]
The layer is parameterized by three MLPs where the ﬁrst one updates the edges U edge : R 2 d h + d d (cid:55)→
R d d .
The second one updates the atom representations U h : R d h + d d (cid:55)→ R d h .
The third one
U softedge : R d d (cid:55)→ R is followed by the logistic sigmoid function to create a value betweenand 1 that can be seen as a soft edge weight telling us how probable an edge is for each message m uv as it is done by Satorras et al.
To produce the ﬁnal 3D representation z b , all atom representations are aggregated by concatenating their mean, their maximum, and their standard deviation and feeding this through a ﬁnal feedforward network U : R 3 d h (cid:55)→ R d z .
B E XPERIMENTAL D ETAILS B. D ATA D ETAILS
We use three datasets containing 3D information for pre-training with diversity in molecule size, and the number of molecules as can be seen in Table 5.
The pre-training datasets are:
QM9 1 (Ramakrishnan et al., 2014) contains 134k stable small organic molecules of 5 elements (CHONF).
Every molecule has the 3D coordinates of one low-energy conformer and is annotated with 12 quantum mechanical properties as regression targets.
The molecules are considered very small, with at most 9 heavy atoms.
GEOM-Drugs 2 (Axelrod & Gomez-Bombarelli, 2020) consists of 304k realistically-sized biologically and pharmacologically relevant molecules of 16 elements, annotated with multiple 3D conformers, the ensemble Gibbs free energy, and the ensemble energy as regression targets.
For the average molecule, 70% of the Boltzmann weight is captured by just three conformers as can be seen in Figure 6a where we also provide a histogram for the number of molecules that have a certain amount of conformers in Figure 6b.
The conformers are generated using CREST (Grimme, 2019).
QMugs 3 (Isert et al., 2021) has 665k drug-like molecules with three diverse conformers each and multiple conformer speciﬁc quantum mechanical properties as regression tasks.
The conformers are generated using CREST (Grimme, 2019).
For ﬁne-tuning, we use a variety of datasets that cover a wide range of domains and applications.
The molecular properties are relevant for quantum mechanics, physical chemistry, biophysics, and physiology such that we can obtain a good estimate of how valuable our 3D pre-training is for each domain.
For quantum mechanical properties, which are often speciﬁc to a conformer, it is clear that 3D information is important and there has been a lot of evidence that learned methods highly beneﬁt from its use (Klicpera et al., 2020b;a;
Liu et al., 2021; Sch¨utt et al., 2017).
For these properties, the interest is in how much our method can leverage this information and transfer it to molecules where no 3D geometry is available.
Meanwhile, for biological or physiological properties such as blood-brain barrier penetration, it is not as clear if improvements from 3D information are to be expected.
As such, this question needs to be answered next to how much of the beneﬁts 3D pre-training recovers.
For this purpose, we use the following molecular graph datasets which are mainly taken from MoleculeNet (Wu et al., 2017) and we use the scaffold splits 4 with an 80/10/10 split ratio provided by OGB Hu et al.
(2020a).
The ﬁne-tuning datasets are:
QM9 and GEOM-Drugs: On these 3D datasets we also ﬁne-tune and evaluate the quantum mechanical properties of one half of the datasets with a random split.
This is done after either pre-training on another 3D dataset (generalization), or after pre-training on the other half of the same dataset (in distribution).
ESOL: 1128 common organic small molecules with water solubility data (log solubility in mols per liter).
Lipo: Experimental data for the octanol/water distribution coefﬁcient of 4200 molecules.
FreeSolv:
The hydration free energy of 642 molecules in water.
HIV: 41k molecules with binary labels for HIV virus replication inhibition.
BACE:
Binary labels of binding results for inhibitors of human β -secretase 1 formolecules.
BBBP: 2039 molecules with binary labels for blood-brain barrier penetration. 8.
Tox21: 7831 molecules with binary labels of their toxic for 12 different targets.
ToxCast: 8576 molecules with binary labels of toxicity experiment outcomes withtargets.
SIDER: 1427 approved drugs with 27 different side effect groups and the task is to predict whether the drug is in the side effect group.
ClinTox: 1477 drugs with two binary annotations where the ﬁrst is to predict toxicity in clinical trials and the second is the FDA approval status.
The reason why muv and pcba are the only datasets from the OGB benchmark suite which we omit is their larger size.
[Table]
B. U NITS AND M EANING OF Q UANTUM P ROPERTIES
GEOM-Drugs dataset numbers the unit kcal/mol
Gibbs refers to ensemble Gibbs free energy
[Table]
and (cid:104) E (cid:105) to ensemble energy
[Table]
[Figure]
B. P ARAMETER D ETAILS
The hyperparameters for SMP are taken from the ofﬁcial repository 5 where Liu et al.
(2021) provide their code and we predict gap even though it could be calculated as | homo − lumo | .
The parameter search space and ﬁnal parameters for the PNA architecture are speciﬁed in Table 7 and those of the 3D network in Table 8.
Pre-training: We use Adam with a start learning rate of 8 × 10 − 5 and a batch size of 500.
The learning rate schedule during pre-training starts with 700 optimization steps of linear warmup followed by the schedule given by the ReduceLROnPlateau scheduler by PyTorch 6 with reduction parameter 0.6, patience 25, and a cooldown of 20.
Fine-tuning quantum mechanical properties: We use Adam with a start learning rate of 7 × 10 − 5 , weight decay 1 × 10 − 11 and a batch size of 128.
For the learning rate schedule, we ﬁrst perform warmup as follows.
We consider three different sets of learnable parameters: (1) the batch norm parameters, (2) all newly initialized parameters that were not transferred, and (3) all parameters.
For these sets, we increase the learning rate in this order from 0 to the start learning rate with linear interpolation.
For parameter group one we warm up for 700 steps, 700 steps for group 2, andsteps for group 3.
After that we use the schedule given by the ReduceLROnPlateau with reduction parameter 0.5, patience 25, and a cooldown of 20.
Fine-tuning non-quantum properties: We use Adam with a start learning rate of 1 × 10 − 3 and a batch size of 32.
The learning rate schedule is the same as for the quantum mechanical properties.
The experiment on the non-quantum properties has different hyperparameters for PNA since the smaller datasets are easily overﬁtted on with the large architecture we use for the quantum mechanical properties.
A smaller PNA yields better performance with the random initialization baseline.
Therefore the PNA in these experiments has a hidden dimension of 50 and 3 message passing layers as propagation depth.
Apart from that, it is the same as the PNA described in Table 7.
[Table] [Table]
[Table] [Table] [Table]
[Table] [Table]
momentum [0.
[Table]
B. C ONFIDENCE I NTERVAL D ETAILS
All the speciﬁed conﬁdence intervals in our work are standard deviations calculated from different weight initializations using the seeds [1 , 2 , 3 , 4 , 5 , 6] or [1 , 2 , 3 , 4] .
The following tables provide additional conﬁdence intervals for the results in the main text.
[Table]
[Table] [Table]
I MPLEMENTATION
Code to 3D pre-train a GNN or to reproduce results is available at https://anonymous.
4open.science/r/3141 .
All experiments were implemented in PyTorch (Paszke et al., 2017) using the deep learning libraries for processing graphs Pytorch Geometric (Fey & Lenssen, 2019) and Deep Graph Library (Wang et al., 2019).
The code we use for SMP (Liu et al., 2021) is under the GNU General Public License v3.0 and we use their implementation after discussing it with the ﬁrst author of the paper and under the consideration that their project welcomed our contributions to their library.
The experiments were conducted on two different machines while the same system was always used in direct comparisons.
The ﬁrst machine has an AMD Ryzen 1700 CPU @ 3.70Ghz, 16GB of RAM, and an Nvidia GTX 1060 GPU with 6GB vRAM.
The second system contains two Intel Xeon Gold 6248 CPUs @ 2.50GHz each with 20/40 cores, 400GB of RAM, and four Quadro RTX 8000 GPUs with 46GB vRAM of which only a single one was used for each experiment.
All mentions and of training time refer to the second system.
[Table]
C A DDITIONAL R ESULTS C. D IFFERENT 3D N ETWORKS AND A BLATION
In this section, we justify our the design of our 3D network which we call Net3D in this comparison.
In Table 12 we compare Net3D with different alternative 3D networks which are the 3D GNNs SMP and EGNN operating on learned node embeddings similar to Net3D .
Additionally, we ablate the use of our γ function that maps the pairwise distances to a higher dimensional space since it would be an unnecessary complication if it provides no beneﬁt.
We call it Net3D w/o γ if Net3D directly operates on the pairwise distances.
In Table 12 we can observe that Net3D yields the best downstream performance and that using our γ function is a valuable component of it.
Possibly EGNN would beneﬁt similarly from this encoding.
SMP’s downstream performance is the worst which could be expected since the 3D input representation which it uses does not uniquely deﬁne all the relative positions in a molecule.
We note that SMP is able to distinguish chiral molecules unlike the other 3D networks but this advantage cannot be evaluated with our experiments on quantum mechanical properties.
Chirality only becomes relevant when considering the interactions between molecules and in these situations,
SMP might be able to leverage its advantage such that our evaluation could be criticized as not entirely fair.
Additionally, SMP has much lower memory requirements since it does not suffer from the quadratic complexity of EGNN and Net3D in the molecule size.
Nevertheless, Net3D performs the best and for drug-like molecules, the quadratic complexity is not problematic.
C. D IFFERENT M ETHODS FOR M ULTIPLE C ONFORMERS
We test three main approaches and variations of them for incorporating the 3D information of multiple conformers to justify our choice in the main text.
The most straightforward one is conformer sampling .
We use one of the single conformer setups but when sampling the batch, we additionally sample j ∈ { 1 . . .
c i } and use the single conformer R ji .
The probability of sampling a conformer is either distributed uniformly (so 1 /c
i is the probability for each j ) or given by the Boltzmann weight of each conformer.
multi3D is the approach from the main text where we include multiple conformers as additional positive pairs in contrastive learning.
For each molecule ( G i , { R ji } j ∈{ 1 ...c i } )
we choose the c lowest energy conformers to have a ﬁxed number of them.
If there are fewer than c conformers for a molecule ( c i < c ) then the lowest energy conformer is repeated.
For every molecule the 3D network now takes all c conformers { R ji } j ∈{ 1 ...c } as input and produces their latent 3D representations { z bi,j } j ∈{ 1 ...c } which we can see as additional positive samples.
In our contrastive setting, we, therefore, want the similarity between z a i and all conformer representations that come from the same molecule z bi,j to be high.
As such we modify our loss to obtain:
[Equation]
One concern with this formulation is the following.
Let us consider a single molecule.
The objective of high similarity between the many 3D representations and the single 2D representation might be easier to solve through encoding the same 2D information in the 3D representations instead of capturing the 3D information of all conformers in the single 2D representation.
The 2D network would therefore not learn to produce 3D information from its 2D inputs because the mutual information could be maximized through encoded 2D information.
To address this problem, multi3D+2D is our third approach.
The 2D network is now modiﬁed to produce c many latent 2D representations
f a ( G i ) =
{ z ai,j } j ∈{ 1 ...c } which are compared to all 3D representations of the same molecule in a similarity function sim .
We simply use this similarity in the loss instead of the cosine similarity.
Intuitively, the 2D network now has to produce an embedding for each 3D conformer.
One way to deﬁne such a similarity between two same-sized sets of vectors is to use the sum of all pairwise cosine similarities (for brevity we drop the subscript and only write { z ai,j } to mean the set of all representations corresponding to the i-th molecule):
[Equation]
More principled would be to ﬁnd the optimal transport matching with the highest cosine similarity, such that one 2D representation always corresponds to one 3D representation.
However, this approach was not computationally feasible with the batch sizes we use in contrastive learning.
We instead opt for an upper bound on the maximum similarity matching.
For every 3D representation, we choose the 2D representation that has the highest similarity.
This way one 2D representation could be associated with multiple 3D embeddings and we no longer have a mass preserving matching:
[Equation]
Beyond these similarity measures, we explore additional ones based on the inverse of different distance functions and asymmetric metrics such as the maximum mean discrepancy (Gretton et al., 2012) or the KL- and JS-Divergence when interpreting the conformer representations as samples from a normal distribution.
Results We evaluate which of the different approaches best leverage the additional conformer’s information to justify our choice for multi3D in the main text.
Another hypothesis we wish to test is that for smaller molecules such as those in QM9, the ability to make predictions informed by multiple conformers is not as important as for larger drug-like molecules.
The reasoning is that a single conformer takes most of the Boltzmann weight for QM9’s molecules due to the fewer degrees of freedom.
We test conformer sampling , multi3D , and multi3D+2D when pre-training on either QMugs or one half of GEOM-Drugs and ﬁne-tuning on QM9 or the other half of GEOM-Drugs.
In QMugs we have three diverse conformers available for each molecule which are all used, while for GEOM-Drugs different numbers of conformers are available of which we use the ﬁve with the highest Boltzmann weight i.e. lowest energy.
If there are fewer than ﬁve we duplicate the lowest energy conformer (see Section 4.1 for details).
We recall that for the multi3D+2D loss sets with as many elements as conformers are produced by the 2D and 3D networks.
Both the discussed sim all and sim max are used as similarity measures between those sets.
For the conformer sampling strategies of using a uniform weighting or sampling conformers according to their Boltzmann weight, we do not evaluate the latter on QMugs since we do not have it available with exactly three conformers per molecule.
In Table 13 we can observe that there are large improvements possible when using multiple conformers.
After pre-training on QMugs, the MAE when predicting the homo property decreases from 82 .
57 to 70 .
77 and from . 1966 to .
1831 for predicting the Gibbs free energy for GEOM-Drugs.
Notably, these improvements are much larger than when pre-training with GEOM-Drugs.
This is likely because the GEOM-Drugs dataset contains the lowest energy conformers and we always use the most probable one with the highest Boltzmann weight when pre-training with a single conformer.
Meanwhile, the QMugs dataset contains three diverse conformers per molecule and not the ones with
[Table]
the highest Boltzmann weight.
Pre-training with the lowest energy conformer from GEOM-Drugs already captures most of the relevant information and using more is not as beneﬁcial.
However, for
QMugs, using the information of all three diverse conformers is crucial.
Similar to the small improvements over the random initialization baseline with GEOM-Drugs, the different methods for using multiple conformers mostly perform the same when pre-training with
GEOM-Drugs.
When pre-training with QMugs instead, the MAEs are overall slightly worse and we ﬁnd multi3D to perform the best.
Note that this is with the slight caveat that the epoch at which pre-training is stopped for all methods was chosen based on where multi3D had the lowest MAE.
Due to these results, we consider multi3D and conformer sampling with uniform weighting as our best methods since multi3D performs slightly better with pre-training on QMugs but conformer sampling is simpler and especially uses much less memory.
For multi3D , all the conformers need to be processed in parallel, and training with more than 5 conformers and a batch size of 500 would not be possible on a 48GB vRAM GPU.
The hypothesis that the downstream performance on the smaller molecules of QM9 would beneﬁt less from using multiple conformers than the molecules of GEOM-Drugs clearly does not hold.
Surprisingly, the improvements on the small molecules of QM9 are larger.
C. D IFFERENT L OSSES
We compare the different losses to estimate and maximize the mutual information.
For this purpose we pre-train PNA on 50 000 molecules from QM9 and another instance on 140 000 molecules of
GEOM-Drugs, both with a single conformer.
We do so with the Donsker-Varadhan (Hjelm et al., 2019) estimator, the Jensen-Shannon estimator (Hjelm et al., 2019), noise contrastive estimation of mutual information (InfoNCE), and our loss.
For our loss, we search over seven temperature parameters τ ∈ [0 .
7] and choose τ = 0 .
In Table 14 we see that 3D pre-training on GEOM-Drugs or QM9 can yield signiﬁcant improvements for predicting quantum mechanical properties, especially when using InfoNCE and our loss as objectives.
These two objectives perform better than the Donsker-Varadhan and Jensen-Shannon estimator in every case and the Jensen-Shannon objective is superior to the Donsker-Varadhan estimator which seems to yield no signiﬁcant improvements over random initialization.
The superiority of the Jensen-Shannon loss over the Donsker-Varadhan alternative is in line with the ﬁndings of
Hjelm et al.
(2019) in their different setting on images.
While our loss seems to perform better than InfoNCE in three settings, this might be due to the additional investment in searching through temperature parameters for our loss.
[Table]
C. SSL M ETHODS
Here we compare our 3D Infomax pre-training against three additional SSL methods.
These are
Barlow Twins (Zbontar et al., 2021), multi-modal Bootstrap your own latent (BYOL) (Grill et al., 2020)
[Table]
and Variance-Invariance-Covariance Regularization (VICReg) (Bardes et al., 2021).
We pretrain these methods on one half of QM9.
For a fair comparison, we search through 8 different hyperparameter settings based on the downstream performance on the QM9 homo property.
After these method-speciﬁc hyperparameters were selected, we tuned every method with a random search over the same search space.
For 3D Infomax, we vary the temperature of our loss τ .
When using BYOL we try different decay rates γ for the exponential moving average weight copying.
Here, we include γ = 0 making our setup similar to a multi-modal version of SimSiam (Chen & He, 2020).
For Barlow Twins, the hyperparameter is λ weighting the redundancy loss.
Lastly, for VICReg we vary µ and ν , the parameters for the variance and the covariance regularization:
3D Infomax with our loss: τ ∈ [0 . 02 , 0 .
7] where τ = 0 .performed the best.
Multi-modal BYOL: γ ∈ [0 , 0 . 0005 , 0 .
07] where γ = 0 .performed the best.
Barlow-Twins: λ ∈ [0 . 002 , 0 . 0039 , 0 . 005 , 0 .
where γ = 0 .
0039 performed the best. 4.
VICReg: λ = 1 ; µ ∈ [1 , 0 . 5] ; ν ∈ [0 . 02 , 0 . 04 , 0 . 1 , 0 . 3] where λ = 1 , µ = 1 , ν = 0 .performed the best
[Table]
The results in Table 15 showcase that 3D Infomax clearly is the superior method in our setting.
It decreases the MAE from 82.10 ± 0.33 to 68 .
96 ± 0 .
32 while the other methods either lead to no improvement or to the much smaller drop to 79.16 ± 0.58 for BYOL.
This is not due to collapse to a constant solution since we can observe a high variance between the representations in a batch for all methods.
Furthermore, with the ﬁnal parameter settings, all methods were able to achieve a low value for their loss during pre-training, both on the training and validation data and there are no optimization issues.
Intuitively, the results can be explained by 3D Infomax being the only method that optimizes a lower bound on the mutual information, which potentially makes it especially ﬁt for our setting.
The other approaches have no direct relation to mutual information and instead rely on maximizing a notion of similarity with tricks to prevent collapse.
While this might work for conventional SSL, we see no success in our scenario where the rigorous guarantee on maximizing the mutual information seems valuable.
Another reason for the poor performance of BYOL and especially Barlow Twins and VICReg might be that they rely on having symmetric networks to generate the compared representations.
In our scenario, we have very little similarity between the architectures with our 2D and 3D networks operating on different modalities.
This hypothesis would ﬁt in line with the ﬁndings of Bardes et al.
(2021) and Zbontar et al.
(2021) where introducing asymmetries between the networks hurt performance.
C. P RE - TRAINING A 3D GNN
We try to use our 3D Infomax setup to pre-train a 3D GNN.
For this purpose, we employ SMP (Liu et al., 2021) as 3D network during pre-training with half of the QM9 dataset.
We then transfer it’s weights to and ﬁne tune them using the accurate 3D conformers of QM9 to predict the dataset’s properties.
We compare this with SMP trained on the same molecules with randomly initialized weights.
The only architectural difference between the networks is that the pre-trained GNN does not use atom features for the reasons explained in the 3D Network paragraph in Section 4.1.
[Table]
In Table 16, we ﬁnd that pre-training improves the 3D GNN’s performance.
This may be due to the covalent bonding structure and other 2D information edge information that is available during pre-training and which SMP usually cannot use since it employs a radius graph.
This is the case even though the pre-trained SMP does not have access to the atom features.
Pre-training 3D GNNs might be an intersting future direction to attempt beating the state-of-the-art methods for predicting quantum properties with accurate 3D information.
C. C OMBINING P RE -T RAINING M ETHODS
Here we simply use GraphCL’s node drop augmentation for the 2D graph and the 3D information (removing all pairwise distances for a removed atom) with a drop ratio of 0.2 during our 3D pretraining process.
[Table]