3D Pre-training improves GNNs for Molecular Property Prediction
Authors: Hannes Stärk
Dominique Beaini
Gabriele Corso
Prudencio Tossou
Christian Dallago
Stephan Günnemann
Pietro Lio
Venue: ICLR 2022 Submitted (2022)
Keywords: graph neural networks
deep learning
self-supervised learning
molecular representation

Abstract: 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 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 generates 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.

Under review as a conference paper at ICLR 2022
3D PRE-TRAINING IMPROVES GNNS FOR MOLECULAR PROPERTY PREDICTION
Anonymous authors Paper under double-blind review
ABSTRACT
Molecular property prediction is one of the fastest-growing applications deep learning with critical real-world impacts. Including 3D molecular structure input to learned models improves their performance for many molecular tasks. However, this information is infeasible compute at the scale required by sev- eral 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 sulf-vuue learning, we maximize the mutual information between 3D summary vectors and the representations that they contain latent 3D information. During fine-tuning on molecules with un- known geometry, the GNN still produces implicit 3D information and can use improve downstream tasks. We show that 3D pre- pre-training provides significant im- provements for a wide range of properties, such as a 22% average MAE reduction eight quantum mechanical properties. Moreover, the learned representations can be effectively transferred between datasets in different molecular spaces.
Graph Neural Network (GNN) such
INTRODUCTION
The understanding of molecular and quantum chemistry is rapidly growing area for deep learning with models having direct real-world impacts quantum chemistry (Dral 2020), protein structure prediction (Jumper al. 2021). materials science (Schmidt et 2019), and drug discovery (Stokes) jet al., 2020). In particular, for the task of molecular property prediction, GNNs have had great success Yang 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ütt et al. 2017 Klicpera et al. 2020b Liu et al. 2021 Klicpera et al. 2021
However, using classical molecular dynamics simulations explicitly compute a molecule's geom- etry 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. 202T are still too slow for large-scale applications.
Our Solution: 3D Infomax We pre-train GNN to encode implicit 3D information in its latent vec- tors using publicly available molecular structures. A GNN is pre- pre-trained by maximizing the mutual information (MI) between its embedding of 2D molecular graph and representation capturing the 3D information that is produced by 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 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 inform property predictions.
Several other self 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. 20206b] 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

---

Under review as a conference paper at ICLR 2022
Setting
Approaches
Result
without information which properties have to be predicted.
Standard Approach: molocular graph
X
property
the only input
GNN
Predictions are but liss
2D
and based
accurate since 5D information cannot leve
Explicit 30 Approach:
Employ classic [1] mathodk compute coordinates and use them as Inout to 3D Graph Noural Network.
3D GNN
property
Accurate predictions methods for generating coordinates are too slaw for real-world applications.
Pre-train 2D Het with the molecules for 3D information avatlable to generate impl 3D intormation in latnnt
1.
maximize MI
2. property
During fine mene-taning the Net stil generates latent Information uses inform property predictions.
2D
Prediction ons are more accurate than meihhus that
3D
Transfer weights 20 Net and fine-ture tu-ettuee for predicting molecular properties.
2D Net
3D Net
2D Net
Molecules with 3D information that used for pro training.
Inferwnce and only
2D
2D
uses a single forward pass the chosen 20 Net.
3D
Figure 1: The considered problem setting and the motivation for our 3D Infomax pre-training.
change. Mcanwhile, 3D Infomax pre- pre-training teaches the model to reason about how atoms interact in space, which is principled and generalizable form of information.
We analyze our method's performance by pre-training with multiple 3D datasets before evaluating on quantum mechanical properties and ten datasets with biological, pharmacological, or chem- ical properties. 3D Infomax improves property predictions by large margins and the learned repre- sentations 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 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 given molecular graph there are multiple conformers, i.e., probable arrangements of the atoms, which can lead different chemical properties. usual 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 algo- rithm (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 hours per drug- like molecule per CPU-core (Axelrod & Gomez- Bombarelli, 2020). This highlights the need to cap- ture 3D information without explicitly computing structures, especially for discovery screening datasets comprising of millions or billions of molecules (Gorgulla al. 2020
2

---

Under review as a conference paper at ICLR 2022
Symmetries of Molecules A molecule's conformation does not change if all the atom coordinates are jointly translated or rotated around a point, i.c., 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. [02I). 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 gen- erating the train-test sets, such that the molecules from the test set do not share scaffold with those of the training set. This helps to avoid overestimating the generalization power (Yang et al.) 2019) since MI. models tend to memorize and overfit these structures. We use the common Remis-Mureko 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 molecu- lar 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 representation of a graph G = (V.E) with vertices V connected by edges E. They do by iteratively applying prassage-passing layers and then combining all ver- tex representations in a readout function. A mnsssage-passingg 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 s representation. The messages are usually created by multi-layer perceptrons (MLPs) and are aggregated via permuta- tion invariant functions such as taking their mean, max, or sum. After the message-passing layers, another permutation invariant function 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 molecu- lar 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 af-
ter 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., Hy et al. 2018] & Meuwly ly,
2019).
The field is well established with casily accessible molecular datasets driving progress (Wu
ct al.
2017:
Hu et al. 2020a) and rigorous evaluations of MPNNs for property prediction (Yang
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ütt al. 2017). Similarly, DimeNet (Klicpera et al. 2020b a) proposed extracting more 3D information via bond angles, which further improved quantum mechanical prop- erty prediction. Spherical Message Passing (SMP) (Liu al. 2021 included another angular quan- tity, and GemNet (Kliçpera et al. 202I) developed an approach to also capture torsion angles, such that all relative atom positions are uniquely defined. Equivariant Graph Neural Networks (EGNN) (Saoorra et al. 2021) achieved the same by operating on all pairwise atom distances.
Self-Supervised Learning attempts to find supervision signals in unlabelled data learn mean- ingful representations. In particular, contrastive learning (van den Oord et al.) 2018 Gutmann & Hyvärinen] 2010 (Belghaii et al., 2018 [Hjelm et al. (2019) a popular class of methods. These learn representations by comparing the embeddings of similar and dissimilar inputs and have achieved im- pressive 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 ct al. 202T] 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.
3

---

Under review as a conference paper at ICLR 2022
4
METHOD
Our goal is method that produces latent 3D information from a 2D molecular graph and use it for more accurate molecular property predictions. For this purpose, propose pre-train GNN to produce latent representations of a molecule's 3D geometry. Through this pre- 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 fine-tune them on property prediction tasks as visu- The GNN's produced 3D information can used to improve predictions. The following presents our 3D Infomax pre-training before explaining the baselines we compare with.
alized in Figure 2
4.1
3D INFOMAX PRE-TRAINING
1.
maximize MI
2. property
3D Net
fe 2D Net
fe 2D Net
3D
G 2D
G
2D
Figure 2: We first pre-train a 2D network f by maximizing the mutual information (MI) between its representation za of a molecular graph G and 3D representation 2" produced from the molecules' conformers R', In step 2, the weights of f" are transferred and fine-tuned to predict properties.
3D Infomax uses two different models, as visualized in Figure Firstly, the model that should be pre-trained which we call 2D network fe since its inputs are 2D molecular graphs G (V, E) with atoms V and bonds from which produces representation (g) = E R. Secondly, the 3D network f which encodes the atoms' 3D coordinates R = (rr in a 3D representation (*(
z² E Rdd Our pre-training can also be understood from 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 fa to produce 3D information from the 2D graph inputs, maximize the mutual information between the latent 2D representations z" and 3D representations z². Intuitively, we wish to maximize the agreement between z" and z² if they derived from the same molecule. For this purpose, we contrastive learning (visualized in Figure We consider a batch of N molecular graphs {Gi)ie(L...N) with their atom coordinates (R. (E...E from which the networks produce multiple representations 2 and 2/-
The first objective of contrastive learning is maximize the representations' similarity if they are positive pair, meaning that they come from the same molecule (same index i). The second objective to enforce dissimilarity between negative pairs z and = where i of 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 al. 2020b) and we use a similar loss to jointly optimize our models:
asim(2f..am 2///T
C
log
where sim(22",2) 2² = which can seen as weight for the most similar negative pair. While different combinations of contrastive losses and SSL possible to learn a joint embedding space between 2D and 3D repre- sentations, we found the above loss to perform best. Other methods (Barlow Twins (Zbontar et al.) 2021), BYOL (Grill et al. 2020), VICReg (Bardes al. 2021) explored in Appendix C.4
211[[1== I|]==* is the cosine similarity and is temperature parameter
4

---

Under review as a conference paper at ICLR 2022
z
221
22
f 2D Net
3D Net
f 2D Net
3D Net
1 10
2
2 0 3 Aa
3
200
2
single conformer
two conformers
Figure 3: The single conformer example shows batch of three molecular graphs as input to the 2D network with the corresponding three conformers as input the 3D model. During 3D pre- training, the contrastive loss L enforces high similarity between latent representations that come from the same molecule (green arrows) while encouraging dissimilarity otherwise (red arrows). This is depicted for the first molecule, but the same is calculated for the second and third. The final loss is the average. The multiple conformer example on the right shows two conformers per molecule c = 2, and the loss is adjusted to treat all of them as positive pairs if they come from the same molecule and as negative pairs otherwise. Our loss CmulttDD achieves this.
Multiple Conformers For most molecules, there are multiple low-energy stable conformers. In- stead of only using the most probable conformer (with the lowest energy), we found that leveraging structural information from multiple conformers provides significant benefits. To achieve this, we now consider the highest probability conformers { { [yer.......) of the i-th molecule. If there are fewer than conformers for molecule, the lowest energy conformer is repeated. Our choice for the following approach is justified 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 3DD representations ) The objective is maximize the similarity between 2 and all conformer representations z², that stem from the same molecule (see Figure 3 As such, we modify
our loss to over the similarities of all conformers obtain the final loss.
N
Σj sim( #L.//7
Cma
log
N
j=1 (
3D Network As the 2D network an arbitrary GNN that should be pre-trained for which we choose Principal Neighborhood Aggregation (PNA) (Corso 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 2 and does not have access the 2D information used by the 2D network such atom or bond features. Otherwise, the mutual information could be increased by both networks encoding this information instead of the desired 3D information.. Our concrete architecture encodes the 3D information given by the pairwise Fuclidean distances of all atoms. This representation uniquely defines all relative atom positions and is invariant to translation and rotation, as desired, but also to reflection, which is disadvantage since molecular properties can change under reflections (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 duv between atoms u and v are first mapped to a higher dimensional space using feghhfreuuency functions. This motivation for this is enable deep networks better fit data with figh-freeuency 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
5

---

Under review as a conference paper at ICLR 2022
might be the most important ones assuming that close-by atoms have the most relevant interactions. As such, we use the following mapping Y RR + R2F+ with the number of frequencies F set to 4:
((dwv) = (dues sin(dwr ) cos(duv
sin(dwu/2 /2f
cos /2f)))/2P1)))))))))
(3)
The further components can be seen as an MPNN Giimmr et al. 2017) operating on the fully con- nected graph of molecule with the encoded distances edge features and 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². The differences to standard MPNNs are detailed in Appendix A Instead of the presented architecture, a 3D GNN such as SMP (Liu al.) (02T) op- erating on learned node embeddings could also be used. We justify choosing our architecture with experiments in Appendix C.T that also show the y mapping's 's effectiveness by ablating it.
4.2 PRE-TRAINING BASELINES
A simpler method 3D pre-train a GNN instead of 3D Infomax is by directly predicting all atom distances. To predict the distance between node e and u, we concatenate their representations h... h Rdd that were produced by the GNN and feed them to an MLP that produces single scalar U Rddd + RR. The distance prediction distu is then given by
Distance Predictor
distuv softplus(U h..) + Uh. h)))
(4)
where || denotes concatenation and softplus(x)) log(1 + e*). The node representations are con- catenated in both orders and fed to the MLP to ensure that the function is symmetric. The final loss to pre-train f is the mean squared error between the predicted and true distances.
Conformer Generation As 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. 1 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-base 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 common self- supervised objective in which the model has to learn to produce representations that are invariant to augmentations. We use randomly dropping nodes with ratio of 0.2 both branches of the SSL setup since JOAO found this combination of augmentations to work particularly well for molecules.
5 EXPERIMENTS
Data and Setup For pre- pre-training. we use three datasets of molecules with 3D conformer informa- tion: QM9 (Ramakrishnan et al. 2014) which contains 134k small molecules (18 atoms on average) 2020) with 304k molecules
with a single conformer, GEOM DrugsD (Axelrod & KGomez-Bombann and QMugs (Isert al.) 2021 with 665k. GEOM-Drugs and QMug, both consist of larger drug-like like molecules (44.4 and 30.6 atoms on average) with multiple conformers. For fine-tuning we consider quantum properties
the one side and biological, chemical, and pharmacological properties on the other side. We predict ten quantum properties of QM9 and GEOM-Drugh (a disjoint half of the datasets if the other half was used for pre-trainings for which we employ 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.I
We choose PNA (Corso et al. 2020) as the GNN to pre-train due to its simplicity and SOTA per- formance for molecular tasks. The reported confidence intervals are one standard deviation cal- culated from six random weight initializations, unless stated otherwise. All baselines we com- pare with the same GNN tailed in Appendix B Code to 3D treit a GNN or to reproduce results is available at https:: 7/anonymous. 4open. aci ence/
3D Infomax method and all experimental settings are de-

---

Under review as a conference paper at ICLR 2022
5.1 QUANTUM MECHANICAL PROPERTIES
We use 3D Infomax to pre-train three different instances of PNA (1) on 50k molecules from QM9 using single conformer, (2) 140k of GEOM-Drugs with 5 conformers and, (3) on 620k of QMugs using conformers. For comparison, we use two different pre-training methods. These are GraphCL (You al. 2020 described in Section 4.2 and pre-training by predicting the Gibbs free enery of GEOM- GEOM-Drugs pre-training subset (labeled PropPred). All pre-training methods use
batch size of 500.
After pre-training, the models are fine-tuned on 50k molecules from QM9 (in Table or 140k from GEOM-Drugs (in Table 2 that have no overlap with the molecules from the pre-traming data. On the same molecules, we also train PNA with random weight initialization (labeled Rand Init) to compare how much the downstream performance 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 fast manner (labeled RDKit SMP). Using conformers generated by the SOTA learned method GeoMol (Ganea et al.] 2021 always performed worse (Appendix C.6). 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 I: Mean Absolute Error (MAE) for QM9's properties. 3D Infomax is tested with three different pre- tre-training datasets and GraphCL uses a two times larger subset GEOMMDrugs. True 3D SMP is 3D GNN using ground truth 3D coordinates (hidden from other methods). Details on confidence intervals are in Appendix B Colors indicate improvement (lower MAE) or worse performance compared the randomly initialized (Rand Init) model.
Pre-training baselines PropPred
Our 3D Infomax Drugs
RDKit SMP
True 3D
Target
Rand Init
GraphCL
QM9
QMugs
SMP
μ 0.4133=00003
0.3937
0.3975
0.3507
0.3512
0.3668
0.4344
0,0726
0.3972=00004
0.3295
0.3732
0.3268
0.2959
0.2807
0.3020
0.1542
homo
82.10
79.57
93.11
68.96
70.78
70.77
82.51
56.19
lumo
85.72:1162
80.81
99.84
69.51
71.38
78.10
80.36
43.58
gap
123.08 3.99
120.08
131.99
101.71
102.59
103.85
114.24
85.10
r2
22. 20
21.84
29.21
17.39
18.96
18.00
22.63
1.51
ZPVE
15.08
12.39
11.17
7.966
9.677
12,06
5.18
2.69
0. 0.1000=00004
0.1422
0.1795
0.1306
0.1409
0.1208
0.1419
0.0498
Table shows that 3D Infomax pre-training leads to large improvements over the randomly initial-
ized baseline and over GraphCL with all three pre-training datasets. After 3D pre-training one
half of QM9, the average decrease in MAE is 22%. Comparing 3D Infomax on GEOM-Drugs with GraphCL. shows that even though the latter is pre-trained two times many molecules from the same dataset, 3D pre-training is always better by 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 gener- alize well: pre-training on GEOM-Drugs and QMugs leads to improvements of 19% and 18% respectively, even though QM9 contains much smaller molecules with on average atoms compared to the 44.4 atoms for the drug-like molecules of GEOM-Drugs.
Table 2: The MAE for GEOM-Dru DrugsD properties. 3D Infomax compared with GraphCL and no pre-training.
Method
Gibbs
(E)
Rand Init
.2035
.1026
GraphCL
.1941
.0995
3D Infomax QM9
.152
.0968
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 conformer-specte specific. There might be maximum accuracy that can be achieved if only the molecule is known and not for which conformer the property should be predicted. Nevertheless, this perfor- mance gap suggests that there is still room for improvement.
3D Infomax Drugs
.1811
.0952
3D Infomax QMugs
.1835
.0965
7

---

Under review as a conference paper at ICLR 2022
Table 2 further confirms that 3D Infomax substantially improves quantum property predictions and generalizes out-of-distrbuution. Our method outperforms GraphCL, even though GraphCL also sees the fine-tuning molecules during pre-training. Moreover, observe strong generalization when pre-training with QM9 and fine-tuning on GEOM-Drugs. 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 high similarity between the molecules during pre-training and fine-tuning.
Another advantage of 3D Infomax is its comparably fast convergence. Pre-training on 620k molecules of QMugs with 3 conformers takes 12 hours, compared to 71 hours for GraphCI on 280k molecules of GEOM-Drugs.
5.2 PREDICTIVE 3D PRE-TRAINING
Table 3:
Comparison of 3D Infomax against predictive 3D pre-training baselines. Shown is the MAE for predicting QM9's properties. Colors indicate improvement (lower MAE) or worse perfor- mance compared to the randomly initialized (Rand Init) model.
Method
μ
a
homo
lumo
gap
r2
ZPVE
Cp
Rand Init
0.4148
0.3348
82.10
87.74
120.94
22.14
15.08
0.1670
Dist-pred
0.4626
0.3570
80.58
84.93
116.21
29.23
25.91
0.1587
Conf-gen
0.3940
0.4219
79.75
79.16
110.72
20.86
21.10
0.1555
3D Infomax
0.3512
0.2959
70.78
71.38
102.59
18.96
9.677
0.1409
3D pre-training by directly predicting 3D quantities is simpler than 3D Infomax and would be prefer- able in case of similar gains. Therefore, we compare with the baselines in Section Jection[..2ussinnn 4.2jusing the same 140k molecules of GEOM-Drugs for all 3D pre-traie 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 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).
5.3 NON-QUANTUX PROPERTIES
Table 4: Comparison of 3D pre-training baselines and GraphCL against 3D Infomax on various OGB datasets. Shown is either the Root Mean Squared Error (RMSE) (lower is better) or the area under the ROC-curve (ROC-AUC) (higher is better). Colors indicate improvement, worse perfor- mance, or no significant change compared to the randomly initialized (Rand Init) model.
Dataset
Metric
Rand Init
Dist-pred
Conf-gen
GraphCL
3D Infomax
esol
RMSE +
0.947
0.986:10.025
0.867 10.045
0.959
0.804
lipo
RMSE L
0.739+0.000
0.718:10.00
0.757 20.035
0.714
0.695
freesolv
RMSE J.
2.233
2.486
2.428
1..24
2.
bace
ROC-AUC f
78.13 1.30
76.51 1.99
80.02-11.58
77.18
79.42-1194 94
bbbp
ROC-AUC
68.27=1.98
66,06#:
16+2.22.24
71.06=2.00
69.10++17
tox21
ROC-AUC
73.88=0.64
73.87+0.43
75.24+1.00
78.92=0.61
74,46++0.4
toxcast
ROC-AUC
63.622..48
61.58.10.88
64.74.11.20
64.95
64.41
clintox
ROC-AUC
58.98
55.77 1:588
64.27-15.22
51.07 1.5.52
19.43
sider ROC-AUC
55.95=3.27
57.13++1.89
56.3444420
57.32=5.00
53.37+3.34
hiv
ROC-AUC f
77.06=3.16
75.66++1.26
76.57+1.39
76.06
76.08+1133
In the previous sections, we found that 3D Infomax yields large improvements for predicting quan- tum 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, would expect it be helpful, and we compare different methods pre- pre-trained with FEOM-Drugh
In Table we find that 3D Infomax improved performance for 4 out of 10 OGB datasets. In contrast to the results for quantum mechanical property predictions (Section 5. is not always superior to
8

---

Under review as a conference paper at ICLR 2022
GraphCL and Conf-gen. However, 3D Infomax never decreases performance which can be valuable in practice and make the method worth employing for non non-quantum properties as well.
When investigating for which tasks 3D Infomax is useful, we see that abstract tasks such pre- dicting clinical test outcomes (clintox) benefit less. The most significant improvements are rather possible for tasks like predicting solubility and lipophilicity in esol and lipo. These are more di- rectly related to molecular mechanics and a molecule's intrinsic properties (e.g., the dipole mo- ment/polarity is important for predicting lipophilicity). They do not depend on how molecule will interact with others to result in, e.g., different effects on patients. For such tasks, Conf-get often leads to significant improvements, providing further evidence for the value of 3D pre-training.
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 molecule's geometry is less helpful for bace since the geometry of the protein and the binding pocket the molecule has to fit into are not known.
5.4
NUMBER OF CONFORMERS AND PRE-TRAINING MOLECULES
0.185
QM9
GEOM-Drugs
0.184
71.4
MAE homo
71.2
0.183
MAE Gibbs
MAE homo
70.8
0.162
70.6
00.88
70.4
00.800
35k
280k
number of conformers
Pre-train dataset size
Figure 4:
The left plot shows the MAE for the QM9's homo and GEOM-Drugs' Gibbs property when varying the number of GEOM-Drugs' conformers used during pre-training, The right plot depicts the MAE when using different numbers of molecules of GEOM-Drugs during pre-training.
Figure ss left plot highlights the benefit of using more than a single conformer. However, the marginal gain reduces higher energy conformers added and beyond certain point (around three conformers), the reduced focus on the most likely conformers worsens the downstream perfor- mance. This is in line with the observation that, average, three conformers enough to cover 70% of the cumulative Boltzmann weight for GEOM- Drugs. Additionally, experiments in Appendix C.2 show that using multiple conformers 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 77 when using three.
In the right plot of Figure 4 we can observe the performance improving as the size of the pre-training dataset increases. However, the returns are diminishing, and we cannot claim that even larger pre- training datasets are likely to drastically improve performance.
6 CONCLUSION
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 fine-tuning im- prove molecular property predictions while retaining the inference speed of a standard GNN op- erating 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 principled form of information that is known to be useful for molecular tasks. Similarly, we demon- strated that learned embeddings are transferable across different physical, biological, and pharma- ceutical tasks. Lastly, we observed that using multiple molecular conformers during pre-training provides valuable additional information to further improve downstream property predictions.

---

Under review as a conference paper at ICLR 2022
7 ETHICS STATEMENT
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 approach can be used when predicting the properties of vast amounts of molecules to identify potential drug materials science could also help
candidates and in large-scale virtual screening. Applications 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 molecu- lar property predictions of learned models. Research on explainability of GNN's predictions would be erucial for solving this problem. Additional possible risks come from malicious applications of molecular property predictions, such as developing chemical or biological weapons.
8 REPRODUCIBILITY STATEMENT
All code to reproduce results is available at
4open. swrereenscienneerrrrre
3141] This repository also provides an explanation for running the experiments. Moreover, we detail our hyperparameter search spaces and final parameter settings in Appendix B.3 for all our used architectures and baselines. Furthermore, Appendix [..] 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 and completely specify our used hard- and software in Appendix B.5
REFERENCES
Simon Axelrod and Rafael for property prediction and molecular generation. arXiv preprint arXiv:2006.05531, 2020.
Geom: Energy anrotttdd molecular conformations
Adrien Bardes, Jean Ponce, and Yann LeCun. Vicreg: Variance-ivvaaance-ovvvviiceee heg.:/mmmeeoooocaaaaaaeeeerrriia 2021.
for self-supervised learning. CoRR,
Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeswar, Sherjil Ozair, Yoshua Bengio, R. Devon Hjelm, and Aaron C. Courville. Mutual information neural estimation. In Proceedings of the 35th International Conference Machine Learning, ICML 2018, Stockholm, Sweden, July 10-15, 2018, volume 80 Proceedings of Machine Learning Research, pp. 530-539,000 2018.
Guy W. Bemis and Mark A. Murcko. The properties of known drugs. 1. molecular frameworks. Journal of Medicinal Chemistry, 39(15):8887-289, 1996.
Mare Brockschmidt. Gnn-film: Graph neural networks with feature-wise linear modulation. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 1144-1152. PMLR, 2020.
Michael M. Bronstein, Joan Bruna, Taco Cohen, and Petar Velickovic. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. CORR, abs/2104.1 13478. 2021.
Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tiee Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference Neural Information Pro- cessing Systems 2020, NewrIPS 2020, December 6-12, 2020, virtual, 2020.
Pengfei Chen, Weiwen Liu, Chang- Yu Hsich, Guangyong Chen, and Pheng Ann Heng. Utilizing edge features in graph neural networks via variational information maximization, 2020a.
10

---

Under review as a conference paper at ICLR 2022
Ting Chen, Simon Korblithh, Mohammad Norouzi, and Geoffrey E. Hinton. A simple framework for contrastive learning of visual representations. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 1597-1607. PMLR, 2020b.
Xinlei Chen and Kaiming He. Exploring simple siamese representation learning.
CoRR,
0bs/2011.10566. 2020.
Connor W. Coley, Wengong Jin, Luke Rogers, Timothy Jamison, Tommi S. Jaakkola, William H. Green, Regina Barzilay, and Klavs F. Jensen. A craph-convoluttonall neural network model for the prediction of chemical reactivity. Chem. Sci., 10:370-377 2019.
Gabriele Corso, Luca Cavalleri, Dominique Beaini, Pietro Liò, and Petar Veličković. Principal neighbourhood aggregation for graph nets. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 13260- 13271. Curran Associates, Inc., 2020.
Pavlo O. Dral. Quantum chemistry in the age of machine learning. The Journal of Physical Chem- istry Letters, 116))2336-2377, 2020.
Matthias Fey and Jan E. Lenssen. Fast graph representation learning with PyTorch Geometric. In ICLR Workshop on Representation Learning on Graphs and Manifolds, 2019.
Octavian- Eugen Ganea, Lagnajit Pattanaik, Connor W. Coley, Regina Barzilay, Klavs Jensen, William H. Green, and Tommi S. Jaakkola. Geomol: Torsional geometric generation of molecular 3d conformer ensembles, 2021.
Justin Gilmer, Samuel Schoenholz, Patrick Riley, Oriol Vinyals, and George E. Dahl. Neural message passing for quantum chemistry. In Yee Whye Teh (ed.), Proceedings of the 34th Interna- tional Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pp. 1263-1272. PMLR, 2017.
Christoph Gorgulla, Andras Boeszoermenyi, Zi-Fu Wang, Patrick D. Fischer, Paul W. Coote, Kr- ishna M. Padmanabha Das, Yehor S. Malets, Dmytro S. Radchenko, Yurii Moroz, David A. Scott, Konstantin Fackeldey, Moritz Hoffmann, Iryna lavniuk, Gerhard Wagner, and Haribabu Arthanari. An open-source drug discovery platform enables ultra-large virtual screens. Nature, 58007805):663-6688 668, 2020.
Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. Journal of Machine Learning Research, 13(25):723-7733 2012.
Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre H. Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Ávila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, Bilal Piot, Koray Kavukcuoglu, Rémi Munos, and Michal Valko. Bootstrap your own latent A new approach to self-supervised learning. Hugo Larochelle, Mare' Aurelio Ranzato, Raia Had- sell, Maria-Florina Balcan, and Hsuan-Tiees Lin (eds.), Advances in Neural Information Process- ing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
Stefan Grimme. Exploration of chemical compound, conformer, and reaction space with meta- dynamics sinulations based on tight-binding quantum chemical iiriiaations..orra of Chemical Theory and Computation, 15(5):2847-285 2019.
Michael Gutmann and Aapo Hyvärinen. Noise-contrastive estimation: new estimation principle for unnormalized statistical models. In Yee Whye Teh and D. Mike Titterington (eds.), Proceed- ings of the Thirteenth International Conference on Artificial Intelligence and Statistics, AISTATS 2010, Chia Laguna Resori, Sardinia, Itlly, May 13-15, 2010, volume 9 of JMLR Proceedings, pp- 297-304. JMLR.org, 2010.
R. Devon Hjelm, Alex Fedorov, Samuel Marriiidnr,,,,onn Karan Grewal, Philip Bachman, Adam Trischler, and Yoshua Bengio. Learning deep representations by mutual information estimation and maximization. 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.

---

Under review as a conference paper at ICLR 2022
Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. In Hugo Larochelle. Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020a.
Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay S. Pande, and Jure Leskovec. Strategies for pre-training graph neural networks. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenRe- view.net, 2020b.
Truong Son Hy. Shubhendu Trivedi, Horace Pan, Brandon M. Anderson, and Risi Kondor. Pre- dicting molecular properties with covariant compositional networks. The Journal of Chemical Physics, 148(24))241745, 2018.
Clemens Isert, Kenneth Atz, José Jiménez-Luna, and Gisbert Schneider, Qmugs: Quantum mechan- ical properties of drug-like molecules, 2021.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, Alex Bridgland, Clemens Meyer, Simon A. A. Kohl, Andrew Ballard, Andrew Cowie, Bernardino Romera- Paredes, Stanislav Nikolov, Rishub Jain, Jonas Adler, Trevor Back, Stig Petersen, David Reiman, Ellen Clancy, Michal Zielinski, Martin Steinegger, Michalina Pacholska, Tamas Berghammer, Se- bastian Bodenstein, David Silver, Oriol Vinyals, Andrew W. Senior, Koray Kavukcuoglu, Push- meet Kohli, and Demis Hassabis. Highly accurate protein structure prediction with alphafold. Nature, 596(7873) 2021.
Johannes Klicpera, Shankari Giri, Johannes T. Margraf, and Stephan Günnemann. Fast and uncertainty-aware directional message passing for non-equilibrium molecules. In NeurlPS-W, 2020a.
Johannes Klicpera, Janek GroB, and Stephan Günnemann. Directional message passing for molec- ular graphs. In International Conference on Learning Representations (ICLR), 2020b.
Johannes Klicpera, Florian Becker, and Stephan Günnemann. Gemnet: Universal directional graph neural networks for molecules. NeurIPS, 2021.
Greg Landrum. Rdkit: Open-source cheminformatics software, 2016.
Junying Li, Deng Cai, and Xiaofei He. Learning graph-level representation for drug discovery. CoRR, abs// 2017.
Yi Liu. Limei Wang, Meng Liu, Xuan Zhang, Bora Oztekin, and Shuiwang Ji. Spherical message passing for 3d graph networks. CORR, 102.05013 2021.
Sinno Jialin Pan and Qiang Yang. A survey transfer learning. IEEE Trans. Knowl, Data Eng., 22 (105:1345-1359, 2010.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch, 2017.
Nasim Rahaman, Aristide Baratin, Devansh Arpit, Felix Draxler, Min Lin, Fred Hamprecht, Yoshua Bengio, and Aaron Courville. On the spectral bias of neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 5301-5310. PMLR, 6 2019.
Raghunathan Ramakrishnan, Pavlo O. Dral, Matthias Rupp, and O. Anatole von Lilienfeld. Quan- tum chemistry structures and properties of 134 kilo molecules. Scientific Data, 1(1):140022, 2014.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural net- works. CoRR, abs42102.09844, 2021.
12

---

Under review as a conference paper at ICLR 2022
Jonathan Schmidt, Mário R. G. Marques, Silvana Botti, and Miguel A. L. Marques. Recent ad- vances and applications of machine learning in solid-state materials science. npj Computational Materials, 5(1):83, 8 2019.
Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmiela, Alexandre Tkatchenko, and Klaus-Robert Müller. Schnet: A continuous- filtte convolutional neural net- work for modeling quantum interactions. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, V. N. Vishwanathan, and Roman Garnett (eds.). Advances in Neural Information Processing Systems 30: Anmual Conference on Neural Information Process- ing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 991-1001, 2017.
Chence Shi, Shitong Luo, Minkai Xu, and Jian Tang. Learning gradient fields for molecular confor- mation generation. In Marina Meila and Tong Zhang (eds.). Proceedings of the 38th International Conference on Machine Learning. ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 9558-9566 PMLR, 2021.
Jonathan M. Stokes, Kevin Yang, Kyle Swanson, Wengong Jin, Andres Cubillos-Ruiz Nina M. Donghia, Craig R. MacNair, Shawn French, Lindsey A. Carfrae, Zohar Bloom-Ackermann, Vic- toria M. Tran, Anush Chiappino-Pepe. Ahmed H. Badran, Ian W. Andrews, Emma Chory, George M. Church, Eric D. Brown, Tommi S. Jaakkola, Regina Barzilay, and James J. Collins. A deep learning approach to antibiotic discovery. Cell, 180(4):688-7022 -022222.113 2020.
Matthew Tancik, Pratul P. Srinivasan, Ben Mildenhall, Sara Fridovich-Kan Keil, Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan Barron, and Ren Ng. Fourier features let net- works learn high frequency functions in low dimensional domains. NeurlPS, 2020.
Bowen Tang, Skyler T. Kramer, Meijuan Fang, Yingkun Qiu, Zhen Wu, and Dong Xu. A self- attention based message passing neural network for predicting molecular lipophilicity and aque- ous solubility. J. Cheminformatics, 12(1):15, 2020.
Wen Torng and Russ Altman. Graph convolutional neural networks for predicting drug-tare ttugetarget interactions. J. Chem. Inf. Model., 59(10):4131 11144411114449, 2019.
Oliver T. Unke and Markus Meuwly, Physnet: A neural network for predicting energies, forces, dipole moments, and partial charges. Journal of Chemical Theory and Compulation, 15(6):3678- 3693, 2019.
Aäron den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding. CoRR, abs/1 807. .03748, 2018.
Minjie Wang, Da Zheng, Zihao Ye, Quan Gan, Mufei Li, Xiang Song, Jinjing Zhou, Chao Ma, Lingfan Yu, Yu Gai, Tianjun Xiao, Tong He, George Karypis. Jinyang Li. and Zheng Zhang. Deep graph library: graph-centrie, highly-performant package for graph neural networks. arXiv preprint arXiv: 1rXXvv9909.00315: 2019.
Yuyang Wang, Jianren Wang, Zhonglin Cao, and Amir Barati Farimani. MolcIr: Molecular con- trastive learning of representations via graph neural networks. CoRR, abs/2102. 2021.
Michael Withnall, Edvard Lindelöf, Ola Engkvist, and Hongming Chen. Building attention and edge message passing neural networks for bioactivity and physical-chemical property prediction. J. Cheminformatics, 12(1):1, 2020.
Zhenqin Wu, Bharath Ramsundar, Evan N. Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh S. Pappu, Karl Leswing. and Vijay S. Pande. Moleculenet: A benchmark for molecular machine learning. CoRR, hts//7033.0064. 2017.
Minghao Xu, Hang Wang, Bingbing Ni, Hongyu Guo, and Jian Tang. Self-supervised graph-level representation learning with local and global structure. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Learning. ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 11548- 11558. PMLR, 2021a.
13

---

Under review as a conference paper at ICLR 2022
Minkai Xu, Shitong Luo, Yoshua Bengio, Jian Peng, and Jian Tang. Learning neural generative dynamics for molecular conformation generation. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021b.
Kevin Yang, Kyle Swanson, Wengong Jin, Connor Coley, Philipp Eiden, Hua Gao, Angel Guzman- Perez, Timothy Hopper, Brian Kelley, Miriam Mathea, Andrew Palmer, Volker Settels, Tommi Jaakkola, Klavs Jensen, and Regina Barzilay. Analyzing learned molecular representations for property prediction. Journal of Chemical Information and Modeling, 59(8):3370-3388 2019.
Yuning You, Tianlong Chen, Yongduo Sui, Ting Chen, Zhangyang Wang, and Yang Shen. Graph contrastive learning with augmentations. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Had- sell, Maria-Florina Balcan, and Hsuan-Tiee Lin (eds.), Advances in Neural Information Process- ing Systems 33: Anmual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
Yuning You, Tianlong Chen, Yang Shen, and Zhangyang Wang. Graph contrastive learning auto- mated. In Marina Meila and Tong Zhang (eds.), Proceedings the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 12121-12132. PMLR, 2021.
Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny. Barlow twins: Self-supervised learning via redundancy reduction. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning. ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 12310-12320. PMLR, 2021.
A FURTHER EXPLANATIONS
3D Network Details The 1-th layer of the 3D network takes sets as input. First, n² edge representations {d E Rdda u, E V Aus + v} (the edges of a complete graph without self- loops). In the first layer they are given by the encoded distances fed through an initial feed-forward network Unnst 12 + IRda which projects them to the hidden dimension of the edges d Uinse( (ydwe)) The second input a set of n atom representations (h}.... N2 with dimensionality Rd,, In the first layer, the atom representations are all set to the same learned vector that is initialized with a standard normal. With meaning concatenation, every layer updates the edge and atom representations and iteratively encodes 3D information into them as follows:
mav Uedge (Th,, || h || d(,..l)
(5)
d =d d' + mau
(6)
1411 Unh
Σ Muv @UUsoftedgee (Mwet
The layer is parameterized by three MLPs where the first one updates the edges Uedge 2R, + Rdd, The second one updates the atom representations Un, Rdh-da + Rdd The third one Usaftedge : RRda + IR. is followed by the logistic sigmoid function to create a value between and 1 that can be seen as soft edge weight telling us how probable an edge is for each message may as is done by /Satorras et al. (2021).
To produce the final 3D representation 2b, all atom representations are aggregated by concatenat- ing their mean, their maximum, and their standard deviation and feeding this through final feed- forward network U 3dd Rd=
B EXPERIMENTAL DETAILS
B. DATA DETAILS
We use three datasets containing 3D information for pre-training with diversity in molecule size and the number of molecules, can be seen in Table 5 The pre-training datasets are:
14

---

Under review as a conference paper at ICLR 2022
Molecules
Scaffold
HO
Figure 5: Depiction of two groups of molecules where all the molecules in the top have the same Bemis-Murcko scaffold (Bemis & Murcko, 1996) which different from the scaffold to which the two molecules in the bottom row belong. We easily obtain the scaffold of molecule using RDKit (Landrum 2016).
QM Ramaakishhan et al., 2014) contains 134k stable small organic molecules of 5 ele- ments (CHONF). Every molecule has the 3ID 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.
2. GEOM-Drugs? Axeerodd & Gomez-Bombareli 2020) consists of 304k realistically-sized biologically and pharmacologically relevant molecules of 16 elements, annotated with mul- tiple 3D conformers, the ensemble Gibbs free energy. and the ensemble energy as regres- sion targets. For the average molecule. 70% of the Boltzmann weight is captured by just three conformers as can be seen Figure [a] where we also provide a histogram for the number of molecules that have a certain amount of conformers in Figure 6b) The conform- ers are generated using CREST (Grimme 2019).
3. QMug (Isert et al. 2021) has 665k drug-like molecules with three diverse conformers each and multiple conformer specific quantum mechanical properties regression tasks. The conformers are generated using CREST (Grimme 2019).
For fine-tuning, we use 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 3D information is important and there has been a lot of evidence that learned methods highly benefit from its (Klicpera al. 2020bb Liu et al.) 2021 Schütt et al.) (017). 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.
often specific to conformer, it is clear that
Meanwhile, for biological or physiological properties such as blood-brain barrier penetration, it is clear if improvements from 3D information are be expected. As such, this question needs to be answered next to how much of the benefits 3D pre-training recovers. For this purpose, we use the following molecular graph datasets, which are mainly taken from MoleculeNet (Wu 2017) and we use the scaffold splits" with an 80/10/10 split ratio provided by OGB Hu et al. (2020a) The fine-tuning datasets are:
1. QM9 and GEOM-Drugs: On these 3D datasets we also tne-tune and evaluate the quan- tum mechanical properties of one half of the datasets with 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).
15

---

Under review as a conference paper at ICLR 2022
2. ESOL: 1128 common organic small molecules with water solubility data (log solubility in
mols per liter).
3. Lipo: Experimental data for the octanol/water distribution coefficient of 4200 molecules.
4. FreeSolv: The hydration free energy of 642 molecules water.
5. HTV; 41 molecules with binary labels for HIV virus replication inhibition.
6. BACE: Binary labels of binding results for inhibitors of human ß-secretase 1 for 1512
molecules.
7. 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.
9. ToxCast: 8576 molecules with binary labels of toxicity experiment outcomes with 617 targets.
10. 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.
11. ClinTox: 1477 drugs with two binary annotations where the first is to predict toxicity in clinical trials and the second the FDA approval status.
The reason why muy and pcba are the only datasets from the OGB benchmark suite which we omit their larger size.
Table 5: Statistics the used datasets. In the upper section are datasets with 3D information, which we use for pre-training, and the datasets in the bottom section do not contain additional 3D annotations.
Dataset
#Molecules
Avg. #Atoms
Avg. #Bonds
split
QM9
130831
18.0
18.6
random
GEOM-Drugs
304 293
44.4
46.4
random
QMugs
665 911
30.6
33.4
random
esol
1128
13.3
13.7
scaffold
1ipo
4200
27.0
29.5
scaffold
freesolv
642
8.7
8.4
scaffold
bace
1512
34.1
36.9
scaffold
bbbp
2039
24.1
26.0
scaffold
hiv
41 127
25.5
27.5
scaffold
tox21
7831
18.6
19.3
scaffold
toxcast
8576
18.8
19.3
scaffold
clintox
1477
26.2
27.9
scaffold
sider
1427
33.6
35.4
scaffold
B.2 UNITS AND MEANING OF QUANTUM PROPERTIES
For the GEOM-Drugs dataset, all reported numbers have the unit kcal/mol, Gibbs refers to the ensemble Gibbs free energy, and (E') the ensemble energy.
Table Units and description quantum mechanical properties of the QM9 dataset.
Property Unit
Description
Debye Bohr
Dipole moment
a homo
Isotropic polarizability
meV
Energy of Highest occupied molecular orbital (HOMO)
lumo
meV
Energy of Lowest occupied molecular orbital (LUMO)
gap
meV
Gap, difference between LUMO and HOMO
r2
2
Electronic spatial extent
ZPVE
meV
Zero point vibrational energy
Cu
molk
Heat capacity at 298.15 K
16

---

Under review as a conference paper at ICLR 2022
05
0V2/000 nuinbor nO00022 cc000mm
00
175
25
Coirt
100
76
0.0
0.2
0.4
00
0.0
1.0 0
15
20
amound cactumd variance
n. p.oonformrss
(a)
(b)
Figure 6: The average number of conformers necessary to cover a certain amount of Boltzmann weight in GEOM-Drugs. For a given amount of cumulative Boltzmann weight on the horizontal axis, the vertical axis shows the average number of conformers necessary pass that threshold. b) Histogram of how many molecules there are in GEOM-Drugh with a certain amount of conformers. The histogram is created for 1000 molecules of GFOM-Drugs.
B.3 PARAMETER DETAILS
The hyperparameters for SMP are taken from the official repositoryS where Liu al. (2021) provide their code, and we predict gap even though it could be calculated as homo - lumo The parameter search space and final parameters for the PNA architecture are specified in Table 7 and those of the 3D network in Table
Pre-training: We use Adam with a start learning rate of x 10-5 and batch size of 500. The learn- ing rate schedule during pre-training starts with 700 optimization steps of linear warmup followed by the schedule given by the ReduceLROn LROnPlateau scheduler by PyToret with reduction parameter 0.6, patience 25, and cooldown of 20.
Fine-tuning quantum mechanical properties: We use Adam with start learning rate of 7 10-5 weight decay 10-11 and batch size of 128. For the learning rate schedule, we first 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 to the start learning rate with linear interpolation. For parameter group one, we warm up for 700 steps, 700 steps for group 2, and 350 steps for group 3. After that we use the schedule given by the ReduceLLROnPiateen with reduction parameter 0.5, patience 25, and cooldown of 20.
Fine-tuning non-quantum properties: We use Adam with a start learning rate of 1 x 10-3 and a batch size 32. The learning rate schedule 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 overfitted with the large architecture use for the quantum mechan- ical properties. A smaller PNA yields better performance with the random initialization baseline. Therefore the PNA in these experiments has 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
htt.:p/s/.oorrr....ooooo//ccates Wecducellkonp1aaaeeu.hta
wrariiim..r__sheeee 1r_schedulerr
17

---

Under review as a conference paper at ICLR 2022
Table 7: Search space for the 2D network PNA through which searched to obtain strong baseline performance on the energy of the highest occupied molecular orbital (homo) property of the QM9 dataset. The parameters top to bottom. After this was completed for all parameters, we performed a second round of tuning for subset of them. The final parameters are marked in bold.
tuned in the order in which they listed in this table from
Parameter
Search Space
propagation depth hidden dimension
[4, 5, .71
[40. 50, 75, 90, 100, 150, 200 .300]
message MLP layers
[1, 2, 3]
update MLP layers
[1, 2, 3]
[mean, max, min, std, sum]. [mean, max, min], [mean, max, sum], [mean, max, min, std], [max, sum], [sum] [identity], fidentity, amplification, attenuation]
aggregators
scalers
readout aggregators dropout
[mean]. [sum], [mean, max, sum]. [mean, max, min, sum]
[0, 0.05, 0.1, 0.2]
batchnorm after MLPs
True/False True/False
batchnorm in MLPs
readout MLP layers
[1, 2, 3]
batchnorm momentum
[0.1, 0.9, 0.93]
Table 8: Search space for the 3D network Net3D through which we searched to obtain strong baseline performance the homo property of the QM9 dataset and we considered the size of the network where parameters leading to less memory use are preferred. The parameters were tuned in the order in which they are listed in this table from top to bottom. After this was completed for all parameters, we performed a second round of tuning for a subset of them. The final parameters are marked in bold.
Parameter
Search Space
propagation depth hidden dimension F used in R ~> 2R2F+1 message MLP layers update MLP layers readout aggregators dropout batchnorm after MLPs readout MIPP layers
[1, 3, 4, 5] [10, 20, 40, 60, 80, 100]
[0. 3, 4, 8, 10, 50]
[1, 2, 3]
[1, 2, 3]
[mean], [sum], [mean, max, min], [mean, max, min, sum]
[0, 0.05, 0.1, 0.2, 0.5]
True/False
[1, 3]
batchnorm momentum
[0.1, 0.9, 0.93]
18

---

All standard deviations are Under review as a conference paper at ICLR 2022
B.4 CONFIDENCE INTERVAL DETAILS
All the specified confidence 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, 3, 4]. The following tables provide additional confidence intervals for the results the main text.
Table 9: Additional confidence intervals of our method in Table [ calculated from 4 seeds except for the homo property where are used.
Our 3D Infomax Drugs
Target
Rand Init
Q.99
QMugs
0,4133:10.003
0.3507:10.005
0.010
0.3668:1.004
0.3972:+0.014
0.3268=00006
0.0099=0.009
0.2807
homo
82.10::0.33
68.96=0.32
70.78=0.82
70.77
lumo
85.72++1.62
6..51+00554
71.38
08.19990.69
gap
123.08
101.71++.03
102.59+3.27
103.85++192
r2
22.14:+0.21
17.39+00.44
18.96=0.69
18.00++0.40
ZPVE
15.08:12.83
7.966111.77
9.677
12.06=
Cp
0.1670:10.004
0.1306:10.009
0.1409
0.1008:0.000
Table 10: Additional confidence intervals for Table 2
Method
Gibbs
(E)
Rand Init
.2035+
0.0011
.1026=t 0.0017
GraphCL
.1941
.0995
3D Infomax QM9
.1852
.0968
3D Infomax Drugs
.1811
.0952
3D Infomax QMugs
.1835
.0965
Table 11: Additional confidence intervals for Table
Method
homo
lumo
gap
r2
ZPVE
Co
Rand Init
+0.003 +0.010
+0.014
+0.33
11.62
+3.98
+0.21
+2.83
#0.004
3D Infomax
#0.009
#0.82
#0.74
#3.27
#0.69
11.9
#0.01
B.5 IMPLEMENTATION
Code to 3D pre-train a GNN or to reproduce results is available at ltttps://anonnntt All experiments were implemented PyTorch (Paszke et 2017)
4cpen. wcwence/r//ee
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., (021) is under the GNU General Public License v3.0 and we use their implementation after discussing it with the
first 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 first 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 single one was used for each experiment. All mentions and of training time refer to the second system.
19

---

Under review as a conference paper at ICLR 2022
Table 12: Comparison of 3D networks. The MAE of the homo property pre-training and fine-tuning on different halves of QM9. Net3D w/o refers to dropping the distance encoding of Net3D. Net3D achieves the best MAE.
Method
CM9 MAE
Rand Init
82.10
SMP
72.37
EGNN
70.46
Net3D w/o 7
70.34
Net3D 68.96-10.32
C ADDITIONAL RESULTS
C.I DIFFERENT 3D NETWORKS AND ABLATION
In this section, we justify the design of our 3D network, which we call Ne13D in this comparison. In Table 12 we compare Ner3D with different alternative 3D networks, which are the 3D GNNs SMP and EGNN operating on learned node embeddings similar to Ner3D. Additionally, we ablate the use of our y function that maps the pairwise distances to higher dimensional space since would be an unnecessary complication if it provides no benefit. We call it Ner3D 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 y function is a valuable component of it. Possibly EGNN would benefit 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 define 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.2 DIFFERENT METHODS FOR MULTIPLE CONFORMERS
We test three main approaches and variations of them for incorporating the 3D information of mul- tiple 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 E {1...cc) )) and use the single conformer Rd The probability of sampling a conformer is either distributed uniformly (so 1/c; is the probability for each j) or given by the Boltzmann weight of each conformer.
multi3D is the approach from the main text where include multiple conformers as additional positive pairs in contrastive learning. For each molecule (G,, (R) )j(()...))) we choose the low- est energy conformers to have a fixed number of them. If there are fewer than e conformers for molecule (c: <), then the lowest energy conformer is repeated. For every molecule the 3D network now takes all e conformers {R] (E....) as input and produces their latent 3D represen- tations {=f,) (eE,))e(L...)) which we can see as additional positive samples. In our contrastive setting, we, therefore, want the similarity between 2 and all conformer representations that come from the same molecule 2t, to be high. As such, we modify our loss to obtain:
N
Σj=1 simead (2f s/T
log
(8)
N
(8) 2/, j=1
20

---

Under review as a conference paper at ICLR 2022
One concern with this formulation is the following. Let consider 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 cap- turing 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 modified to produce many latent 2D representations f²(G;) = {zf, (e((.... which are compared to all 3D representations of the same molecule 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 define such 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 {2ī } to mean the set of all representations corresponding to the i-th molecule):
simau (1 (1 (2
simcos (2 2(,1)
More principled would be to find the optimal transport matching with the highest cosine similar- ity, such that one 2D representation always corresponds to one 3D representation. However, this approach was not computationally feasible with the batch sizes use in contrastive learning. We instead opt for an upper bound on the maximum similarity matching. For every 3D representation, 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 match- ing:
simmaz (1 (+
max simcoo(Ef. x(.k). JE((I..))
(10)
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 the KL- and JS-Divergence when interpreting the conformer representations samples from normal distribution.
Results We evaluate which of the different approaches best leverage the additional conformer's information to justify our choice for multi3l) 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 important as for larger drug-like molecules. The reasoning is that 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 fine-tuning QM9 the other half of GHOM-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 five with the highest Boltzmann weight i.c. lowest energy. If there are fewer than five we duplicate the lowest energy conformer (see Section 4.1 for details). We recall that for the mulri3D+2D loss sets with as many elements as conformers are produced by the 2D and 3D networks. Both the discussed simall and simmar are used as similarity measures between those sets. For the conformer sampling strategies of using a uniform weighting 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 con- formers. 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 always use the most probable one with the highest Boltzmann weight when pre-training with a single con- former. Meanwhile, the QMugs dataset contains three diverse conformers per molecule and not
21

---

Under review as a conference paper at ICLR 2022
Comparison of strategies for using multiple conformers. The middle double-column shows the results for pre-training on one half of GEOM-Drugs and the right double-column cor- responds to pre-training on QMugs, and the second row indicates what dataset was used for fine- tuning. The Random Init row shows the performance when training from scratch without any pre- training. For QM9, the reported number is the MAE of the homo property, and for GEOM-Drugs it is the MAE when predicting the ensemble Gibbs free energy. There are large improvements from using multiple conformers, but the differences between the methods are small.
Table 13:
GROM-Drugs pre-training
QMugs pre-training
Loss/Estimator
QM9
GEOM-Drugs
QM9
GEOM-Drugs
Rand Init
82.10
2.05
02.30
.0335..0011
single conformer
71.66
.1844
82.57
.1966
uniform sampling
70.66
.1823
72.94
.1874
boltzmann sampling
70.93
.1846
multi3D
70.78
.1811
70.77
.1831
stmall
71.11
.1849
72.40
.1936
simmaz
70.81
.1896
71.15
.1840
the ones with 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 beneficial. 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--rrugs When pre-training with QMugs instead, the MAES are overall slightly worse, and find mulni3D perform the best. Note that this is with the slight caveat that the epoch which pre-training is stopped for all methods was chosen based on where multiDD had the lowest MAE.
Due to these results, consider multiDD, and conformer sampling with uniform weighting as our best methods since multi3D performs slightly better with pre-training on QMugs but conformer sampling 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 48GB vRAM GPU.
The hypothesis that the downstream performance on the smaller molecules of QM9 would benefit 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.3 DIFFERENT LOSSES
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 040 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 J E [0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.7] and choose r 0.1.
In Table 14 we see that 3D pre-training on GEOM-Drugs or QM9 can yield significant improve- ments for predicting quantum mechanical properties, especially when using InfoNCE and our loss as objectives. These two objectives perform better than the Donsker-Varae Varadhan, and Jensen-Shannon estimator in every case and the Jensen-Shane Shann-oh objective superior to the Donsker-Varadhan es- timator, which seems to yield significant improvements over random initialization. The superi- ority of the Jensen- Shannon loss over the Donsker- Varadhan alternative in line with the findings of Hjelm et al.) 2019)) in their different setting images. While our loss seems perform better than InfoNCE in three settings, this might be due to the additional investment in searching through temperature parameters for our loss.
22

---

Under review as a conference paper at ICLR 2022
Table 14: Comparison of mutual information estimators for 3D Infomax. The middle double-column shows the results for pre-training on one half of QM9, and the right double-colume corresponds to pre-training on one half of GEOM-Drugs, and the second row indicates what dataset was used for fine-tuning. The Rand Init row shows the performance when training from scratch without any pre-training. For QM9 the reported number is the MAE of the homo, and for GEOM-Drugs it the MAE when predicting the ensemble Gibbs free energy.
OM9 pre-training GEOM-Drugs
GEOM-Drugs pre-training
Loss/Estimator
QM9
QM9
GEOM-Drugs
Rand Init
82.10
.2035+.0011
82.10
.2035
Donsker- Varadhan
82,49
.2152
85.46
.2033
Sensen-Shannon
80.71
.2078
81.61
.2047
InfoNCE
75.81
.1938
79.31
.1894
our loss
68.96
.1945
71.66
.1844
C.4 SSL METHODS
Here compare our 3D Infomax pre-training against three additional SSL methods. These Barlow Twins (Zbontar et al. 202T), multi-modal Bootstrap your own latent (B YOL) (Grill et al.] Covariance Regularization (VICReg) (Bardes al. 2021). We pre-
2020), and Variance-Invariance eevariance-Covariance train these methods on one half QM9. For a fair comparison, we search through different hyperparameter settings based on the downstream performance on the QM9 homo property. After these method-specific 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 r. When using BYOL we try different decay rates y for the exponential moving average weight copying. Here, we include y 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 v, the parameters for the variance and the covariance regularization:
1. 3D Infomax with our loss: T E [0.02, 0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.7] where т = 0.01 performed the best.
2. Multi-modal BYOL: [0, 0.0005, 0.001, 0.005, 0.01, 0.03, 0.05, 0.07] where y 0.005 performed the best.
3. Barlow-Twins: 0.0039 performed the best.
E j0.002, 0.0039, 0,0,0000000000,,00,,000,0,,,00,, 0.02] where ?
4. VICReg: À 1; E [1,0.5]: v E [0.02, 0.04, 0.3] where = 1, / = 1, = 0.04
performed the best
Table 15: Comparison of latent space SSL. methods. The numbers show the MAE when predicting QM9's homo property after pre-training one half of QM9 with the given method and fine-tuning on the other half of QM9. The Rand Init column shows the MAE without pre-training and with random weight initialization. 3D Infomax is our best latent space SSL method.
Random Init
3D Infomax
BYOL
Barlow Twins
VICReg
QM9 MAE
82.10
68.96 0.32
79. 19.16
82.38
85.15
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 collapse to a constant solution since we can observe high variance between the representations in a batch for all methods. Furthermore, with the final parameter settings, all methods were able to achieve low value for their loss during pre-training both on the training and validation data and there are no optimization issues.
23

---

Under review as a conference paper at ICLR 2022
Intuitively, the results can be explained by 3D Infomax being the only method that optimizes lower bound on the mutual information, which potentially makes it especially fit for our setting. The other approaches have direct relation to mutual information and instead rely maximizing notion of similarity with tricks to prevent collapse. While this might work for conventional SSL,, we see no success in scenario where the rigorous guarantee 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 different modalities. This hypothesis would fit in line with the findings of Bardes et al. (202T) and /hoonarr al. (202T) where introducing asymmetries between the networks hurt performance.
C.5 PRE-TRAINING A 3D GNN
We try to use our 3D Infomax setup to pre-train a 3D GNN. For this purpose, we employ SMP < jtt al. 202T 3D network during pre-training with half of the QM9 dataset. We then transfer it's weights and fine tune them using the accurate 3D conformers of the other half of QM9's molecules 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 16: MAE for predicting QM9's molecular properties. SMP is tested with random weight initialization and with the weights obtained from using it as 3D network our 3D Infomax pre- training setup.
Target
SMP Rand Init
SMP pre-trained
0.0726
0.0801
0.1542
0.1276
homo
56.19
44.50
lumo
43.58
37.48
gap
85.10
70.45
r2
1.51
1.12
ZPVE
2.69
2.43
Gy
0.0498
0.0421
In Table 16 we find that pre-training improves the 3D GNN's performance. This may be due to the covalent bonding structure and other 2D edge information that is available during pre-training and which SMP usually cannot use since it employs radius graph. This is the case the pre-trained SMP does not have access to the atom features. Pre-training 3D GNNs might be an interesting future direction attempt beating the offte-of-the-art methods for predicting quantum properties with accurate 3D information.
though
C.6 CHEAP NEURAL CONFORMERS AS 3D GNN INPUT
In Section 5.1
used RDKit's ETKDG algorithm Landrum 2016 to generate inaccurate but
cheap and fast conformers and employed them as inputs to the 3D GNN SMP (Liu et al. [02T) Here, we attempt the same with conformers generated by the SOTA deep learning method for con- formation generation which is GeoMol (Ganea et al. 2021). For this purpose, we train GeoMol with 50k molecules of QM9 and use it to generate the conformations for the rest of QM9, SMP is then trained on 50k different molecules to predict their properties, either using RDKit's conformers or those of GeoMol. This enables a fair comparison with 3D Infomax, which uses the same molecules for pre-training that were used to train GeoMol. When visually inspecting some of the conformers generated by GeoMol, we found that they were sometimes of poor quality for molecules with rings and contained outliers with conformations that seem particularly unrealistic.
24

---

Under review as a conference paper at ICLR 2022
Table 17: MAE for QM9's properties. 3D Infomax tested with three different pre-training datasets and compared with the 3D GNN SMP using explicit 3D coordinates. The conformers are generated using either the classical method RDKit ETKDG or the learned method GeoMol. Colors indicate improvement (lower MAE) or worse performance compared to the randomly initialized (Rand Init) model.
Our 3D Infomax
RDKit SMP
GeoMol
Target
Rand Init
QM9
Drugs
QMugs
SMP
0.4133:40.00
0.3507
0.3512
0.3668
0.4344
0.6046
0.3972:+0.014
0.3268
0.2959
0.2807
0.3020
0.8435
homo lumo
82.10.10.33
68.96
70.78
70.77
82.51
195.0
85.72.1162
69.51
71.38
78.10
80.36
201.4
gap
123.0811399
101.71
102.59
103.85
114.24
284.1
12
12.14
17.39
18.96
18.00
22.63
65.84
ZPVE
15.08+2.83
7.966
9.677
12.06
5.18
17.40
Cy
0.1670
0.1306
0.1409
0.1208
0.1419
0,5467
Table 17 shows that SMP performs poorly with the conformers generated by GeoMol and using those generated by RDKit is always superior. This is the even though the average accuracy of GeoMol's conformers is comparable to that RDKit ETKDG's conformers when GeoMol is trained on all of QM9 (Ganea al. [02I). We hypothesize that the high MABs with GeoMol's conformers occur since they contained some particularly unrealistic outlier conformations, and SMP is not able handle those well.
C.7 COMBINING
METHODS
Here 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 pre- training process.
Table 18: Comparison of performance when combining 3D pre-training with conventional pre- training by randomly dropping nodes on the 2D or 3D side (labeled 3D Infomax + for various biophysical property OGB datasets. GraphCL is another pre-trained bascline. Shown is either the RMSE indicated by where lower values are better or the area under the curve of the Receiver Operator Characteristic (ROC-AUC) indicated by f where higher values are better. Colors indicate improvement, worse performance, or no significant change compared to the randomly initialized (Rand Init) model. 3D Infomax either on par with random initialization better. There is no negative transfer as there is with GraphCL.
dataset
Rand Init
GraphCL
3D Infomax
3D Infomax
esoll lipol
0.947
0.95940.047
0.894
0.918+40037
0.739
0.7141..011
0.695=40.01
0.710++0007
freesolv.
20.261
3.
20.27
40.323
bacet
18.13
77.18
79.42.11.4
79.28
bbbpt
68.27.11.99
71.06
69.10.1107
68.64
tox211
13.88
78.92
74.46
13.63
toxcastt
63.6240048
64.95=0.31
64.414-088
63.95400.8
clintoxt
58.9845.40
51,07=5.52
59,43+4.21
83.5943.64
sider
55.95++3.27
57.32
53.3743.34
58,43
hiy?
77.06+43116
76.06-1.06
76.08+11.3
75.38
25


Reviews:
[{"review_id": "cvrK8xplPb8", "reviewer": "Reviewer_7nK8", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\n1. The method proposed in this paper is novel and sound. Incorporating 3D information for molecular property prediction should improve its performance as demonstrated by this and previous work.\n\n2. The experiments are well designed. The authors choose to pre-train on three datasets (QM9, GEOM-Drugs, and QMugs) and evaluate for both quantum mechanical properties and non-quantum properties. Yet some setting descriptions and results are missing (see weaknesses). \n\n3. The paper is clearly written and easy to follow (with some exceptions shown below). \n\n\nWeaknesses:\n\n1. Some experimental settings are unclear. For example, for QM9 and other datasets, what are the input features for your method and the baselines? There is a bunch of work using QM9 for evaluation. If they follow the same setting, the authors may consider listing their results as well for a comparison. Otherwise, it is hard to tell whether this paper has followed the standard way for evaluation on QM9 and whether the evaluation is meaningful and fair. \n\n2. Some evaluation results are missing. For example, in Table 2, the results for Rand Int and GraphCL are shown, yet PropPred is not shown. Is there a specific reason for this? Similarly, in Table 4, the results of Dist-pred and Conf-gen are missing. \n\n3. What does PropPred mean in the paper? The paper states that \"we pre-train baselines by either predicting the properties of GEOM-Drugs’ pre-training subset (labeled Prop Pred).\" Which baselines are referring to? What do the properties of GEOM-Drugs correspond to here?\n\n4. Standard deviation is missing in Tables 1, 2, and 3.\n\n5. The settings for baselines are unclear.  Which GNN is used for the baselines like GraphCL? How many layers are used? \n\n6. The paper claims that conformation methods are slow and shows a previous method that takes 6 hours/molecule for conformation. I am not sure whether this claim is correct. Are neural conformers (e.g. https://arxiv.org/abs/1904.00314) as slow as these methods? In my opinion, the inference of neural models shall be relatively faster.\n\nMinor point:\n\nI think the method has some similarities with knowledge distillation. The model takes a 3D GNN as 'teacher', while the 2D GNN learns from the 3D GNN to obtain latent 3D information. Clearly, the methods proposed in this paper have differences from knowledge-distillation-based methods. The paper might discuss these differences and compare with a knowledge distillation baseline if possible. \n\nTypo:\n\nIncomplete sentence: Including 3D molecular structure as input to learned models their performance for many molecular tasks.", "summary_of_the_paper": "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. \n\nDuring 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. \n\nThe 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. ", "summary_of_the_review": "The method proposed in this paper is sound. However, the evaluation has some flaws as discussed in the main review section.\n\nI will consider raising my scores if these questions are properly solved during the discussion period.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "cthkpaZXSh", "reviewer": "Reviewer_SAmw", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I found the introduction and background good to read and easy to understand. The paper is supported by great images to explain what was done. It also admits the weakness of the method with biological systems. The non-QM property part is not as clearly written and I found it harder to understand which dataset belongs to which statement. \n-\tPage 4 line 5: I believe the statement that it learns quantum mechanical (QM) interactions is not backed up. The 3D structures of molecules are partially determined by steric hindrance and not purely by QM interactions. E.g. when the model is pretrained on a smaller set of elements and used for elements it has not seen this shows it learned the structure and not QM. \n-\tPage 3 line 3: What is SE(3) symmetry? \n-\tTable 1: Is there a reason why two properties (cv and alpha) are better predicated with the pretraining of QMugs than with the same dataset (QM9)? \n\nSome very minor points:\n-\tFor the review is would be with advantage to add a column with line numbers. Further, exchanging numbers for the citations instead of name and year, this improved the fluent reading experience. \n-\tPage 1 figure 1: MI could be mentioned in the Figure comment or before on page 1 paragraph 4 (our solution) as this abbreviation was not used yet. \n-\tAbstract, Page 2 line 4, page 6: Inconsistency with the number of quantum mechanical properties. 8 are given in the tables but often 10 are mentioned. \n-\tPage 2 last line: Year is missing in the citation. \n-\tPage 4 last line: The word hardest is confusing here. \n-\tPage 7 Table 1: The targets are never mentioned or explained. Only the units are given in the SI. Table 6, here also the names of the variables would be great. \n- Page 8 after table 4, last sentence of the paragraph: Incomplete and I do not understand it.  \n-\tPage 10ff: The citations are not normed, e.g. PMID \n", "summary_of_the_paper": "The paper uses the 3D structures of molecules to pre-train graph neural networks (GNN) representation of 2D molecules. This improves the molecular property prediction of some quantum mechanical (from 8 tested 8 were improved) and non-quantum mechanical properties (from 10 tested 4 were improved). ", "summary_of_the_review": "For molecules the 3D information is essential and this paper shows an affordable method to include this for improved property prediction. There are two main advantages, firstly it can be used for a dataset where 3D information is not available and secondly it is cheaper than implementing 3D information in the representation. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "FjA3z85O5VK", "reviewer": "Reviewer_Awh3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\nThis paper starts with an interesting direction, to use 3D geometry to augment the 2D representation. This is an interesting and well-motivated problem.\n\n-----\n\nWeaknesses:\n\n(1) The empirical result is the most important point, but it has the following problems.\n1. Motivation contradiction. The empirical result does not match the motivation. In the abstract, it says that `fine-tuning on molecules with unknown geometry, …`; in the intro, it says `After pre-training, … fine-tuned on molecular datasets where no 3D information is available.` However, in Sec 5.1, it is tested on `50k molecules from QM9 or 140k from GEOM-Drugs`, where both include 3D information. Besides, as we can tell from the last column of Table 1, 3D geometric GNN is overwhelmingly better, so it is not quite clear why authors want to test on this dataset. Similar concerns on Table 2-3.\n2. The empirical results are confusing. If we ignore the motivation contradiction issue, then according to Method section, this work pre-trains one 2D GNN and one 3D GNN, and downstream tasks include both 2D and 3D information. However, in Table1-3, the authors only check the 2D fine-tuning. If the main focus is on performance improvement, why not focus on 3D downstream tasks? If the performance is about 2D downstream tasks, then why consider these tasks with 3D geometry available? Same concerns on Table 2-3.\n3. Performance not good enough and not enough SSL baselines. As pointed out by authors in Intro and Related Work, there are many SSL works along this research line, like AttrMasking, ContextPred, JOAO. GraphCL is more recent than the previous two, but JOAO is more recent. Besides, GraphCL does not show an overwhelming best performance on these tasks, and instead, it shows better performance on different types of graph data. So I’m just wondering why authors didn’t try this? And Table 4 result shows that 3D InfoMax has very limited improvement, or sometimes has even worse performance. This fails to support the usefulness of 3D Infomax.\n4. Confusing results. Table 1 result is confusing. It doesn’t specify which datasets are used for the pre-training baselines.\n\n(2) Some descriptions are not clear and terminologies should be specified more clearly.\n1. GNN can be adopted on both 2D and 3D graphs, which is not clearly specified in the paper. For example: the abstract says `GNN still generates implicit 3D information …`, and the GNN here refers to the 2D GNN.\n2. In Intro, `We pre-train a GNN to encode implicit 3D information …`. This expression is not clear. A better way would be: This work is training 2 GNN (one 2D GNN and one 3D GNN) at the pre-training stage; with MI (SSL), the 2D representation can encode implicit 3D information. Then only 2D GNN is adopted for fine-tuning because downstream tasks have no geometry info.\n3. In Intro, `This way, the GNN learns to generate latent 3D information using only the information given by the 2D molecular graphs.` The GNN here also refers to the 2D GNN.\n4. The authors mention `2D GNN generate 3D information` at multiple places, like in Intro and Method. But the main solution is to adopt InfoMax, which is a contrastive SSL method; the word `generate` typically implies generative SSL [1]. Authors might better tune these words carefully or explicitly specify them in the paper.\n\nWithout specifying these terminologies explicitly and clearly (e.g., at the beginning of the paper), it is hard for readers to follow the logic of this work.\n\n(3) Some claims about the related work in this paper are too bold.\nIn Intro, `It is not clear why these learned representations should be informative and generalize well.` But all the papers, AttrMasking, ContextPred, GraphLoG, GraphCL, JOAO have claimed the intuitions of the proposed methods, some are supported with ablation studies.\nI acknowledge that 3D info is helpful for 2D, but the above claims on existing works are not correct.\n\n(4) Some text words in Fig 1, Fig 2, Fig 3 seem to be hand-written? If so, authors may as well update them in a more formal way.\n\n(5) The 3D Net / 3D GNN representation is wrong.\nIn Sec 4.1, the equation on 3D Net / 3D GNN is wrong. The representation should be f^b(V, R), i.e., the input to 3D GNN includes both the atom types and atom positions. This holds for SchNet, DimeNet, EGNN, SphereNet.\n\n(6) The objective function is wrong.\nThe NTXent loss doesn't seem correct to me.\nThe authors can verify this from the definition of Cross-Entropy, and here I put a more straightforward explanation.\nIn eq(1) in [2], the numerator is the positive pair; the denominator includes 1 positive pair and 2(N-1) negative paris, i.e., 2N-1 pairs in total.\nIn eq(1) of this paper, however, the numerator is the positive pair, and the denominator only includes N-1 negative pairs.\nThus, this is definitely not NTXent (similar for the following objective functions).\n\nI haven’t found the codes attached in SI, so I just played around with it on my own implementation, and sometimes (on some datasets) it can lead to negative loss. I’m not sure which one is used in this work, and maybe it can still have some positive effects on downstream tasks, but definitely not NTXent loss.\n\n-----\nMinor points:\n\n(1) Sec 4.2 can be moved to Sec 5.\n\n-----\n[1]  Liu, Xiao, et al. \"Self-supervised learning: Generative or contrastive.\" IEEE Transactions on Knowledge and Data Engineering (2021).\n\n[2] Chen, Ting, et al. \"A simple framework for contrastive learning of visual representations.\" International conference on machine learning. PMLR, 2020.\n", "summary_of_the_paper": "This paper starts with an interesting direction, to use 3D geometry to augment the 2D representation. From the technical point of view, it adopts SSL to maximize the representations of 2 views, so as to enable the 2D GNN to encode 3D geometry information, which can be beneficial for downstream tasks.\n", "summary_of_the_review": "This paper is an empirical work, so the empirical results are the most important. However, as listed above, there are some key issues with the empirical performance (point 1). In addition, there are also some other concerns: the writing issues (point 2&3), visualization issues (point 4), and equation issues (point 5-6).", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "23933zHJnuM", "reviewer": "Reviewer_2Vfa", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The major strength of this work is the extensive empirical evidence supporting the claims that the proposed self-supervised learning scheme (mutual information between 2D GNN and 3D GNN learned embeddings) significantly improves performance on downstream tasks as compared to other pre-training strategies.  In addition, the language is simple and clear, and the figures provide easily understandable graphical depictions of the method.  \n\nThe major weakness, which is addressed by the authors, is that these 2D models, even when pre-trained with 3D information, still significantly underperform methods that can leverage 3D information in the downstream task.  This implies that in settings where accuracy is key, one should still opt for the computational expense of generating the 3D conformers directly.  While the authors argue that the gold standard methods of generating 3D conformers can be prohibitively expensive, I would like to see what the performance of 3D GNN operating on conformers such as those that might be generated by a fast method like GeoMol.  This would allow for better evaluation of the tradeoff between accuracy and speed provided by explicit 3D conformer modeling.", "summary_of_the_paper": "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.", "summary_of_the_review": "Overall, I think this is a strong paper due to the compelling results and clarity of communication.  I currently recommend weak accept as I would like to see the experiment I mentioned above (using lower quality, but quick-to-generate 3D conformers) to better understand the tradeoffs involved.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper aims to use pre-training to bridge the gap in performance between 2D GNN and 3D GNN. Specifically, during pretraining, it trains both 2D GNNs and 3D GNNs on data equipped with 3D geometry to maximize the mutual information between the 2D GNN representation with the 3D GNN representation. The proposed approach is interesting and novel and the paper presents some promising results showing that the pre-training does provide some benefits for downstream tasks where 3D geometry information is not available in comparison to several other baseline pretraining methods. While the reviewers agree that property prediction without only 2D graph is a practically important setting for high throughput screening, there are concerns about whether the current set of results paint a clear picture on the benefits and superiority of the proposed methods to alternatives (e.g., vs conf-gen) even after the revision. This is not due to lacking of results, but more of a presentation issue where results are not organized and discussed clearly to provide a coherent story.  We do see clear and strong potential for this paper but it needs a careful rewrite/re-organization to tease out the key messages and how the experiments support them.
Decision: Reject