Pre-training via Denoising for Molecular Property Prediction
Abstract
Many important problems involving molecular property prediction from 3D structures have limited data, posing a generalization challenge for neural networks.
In this paper, we describe a pre-training technique based on denoising that achieves a new state-of-the-art in molecular property prediction by utilizing large datasets of 3D molecular structures at equilibrium to learn meaningful representations for downstream tasks.
Relying on the well-known link between denoising autoencoders and score-matching, we show that the denoising objective corresponds to learning a molecular force field – arising from approximating the Boltzmann distribution with a mixture of Gaussians – directly from equilibrium structures.
Our experiments demonstrate that using this pre-training objective significantly improves performance on multiple benchmarks, achieving a new state-of-the-art on the majority of targets in the widely used QM9 dataset.
Our analysis then provides practical insights into the effects of different factors – dataset sizes, model size and architecture, and the choice of upstream/downstream datasets – on pre-training.
Introduction
The success of the best performing neural networks in vision and natural language processing (NLP) relies on pre-training the models on large datasets to learn meaningful features for downstream tasks [15, 60, 16, 11, 19].
For molecular property prediction from 3D structures (a point cloud of atomic nuclei in R 3 ), the problem of how to similarly learn such representations remains open.
For example, none of the best models on the widely used QM9 benchmark use any form of pre-training [ e.g. 33, 39, 56, 65], in stark contrast with vision and NLP.
Effective methods for pre-training could have a significant impact on fields such as drug discovery and material science.
In this work we focus on the problem of how large datasets of 3D molecular structures can be utilized to improve performance on downstream molecular property prediction tasks that also rely on 3D structures as input.
We address the question: how can one exploit large datasets like PCQM4Mv2,that contain over 3 million structures, to improve performance on datasets such as DES15K that are orders of magnitude smaller?
Our answer is a form of self-supervised pre-training that generates useful representations for downstream prediction tasks, leading to state-of-the-art (SOTA) results.
[Figure]
t r e a m
D o w n s t r e a m
[Figure]
Inspired by recent advances in noise regularization for graph neural networks (GNNs) [23], our pre-training objective is based on denoising in the space of structures (and is hence self-supervised).
Unlike existing pre-training methods, which largely focus on 2D graphs, our approach targets the setting where the downstream task involves 3D point clouds defining the molecular structure.
Relying on the well-known connection between denoising and score-matching [69, 61, 25], we show that the denoising objective is equivalent to learning a particular force field, adding a new interpretation of denoising in the context of molecules and shedding light on how it aids representation learning.
The contributions of our work are summarized as follows:
• We investigate a simple and effective method for pre-training via denoising in the space of 3D structures with the aim of improving downstream molecular property prediction from such 3D structures.
Our denoising objective is shown to be related to learning a specific force field.
• Our experiments demonstrate that pre-training via denoising significantly improves performance on multiple challenging datasets that vary in size, nature of task, and molecular composition.
This establishes that denoising over structures successfully transfers to molecular property prediction, setting, in particular, a new state-of-the-art on 10 out of 12 targets in the widely used QM9 dataset.
Figure 1 illustrates performance on one of the targets in QM9.
• We make improvements to a common GNN, in particular showing how to apply Tailored Activation Transformation (TAT) [79] to Graph Network Simulators (GNS) [51], which is complementary to pre-training and further boosts performance.
• We analyze the benefits of pre-training by gaining insights into the effects of dataset size, model size and architecture, and the relationship between the upstream and downstream datasets.
Related Work
Related work is discussed in detail in Appendix A and summarized in this section.
First, various works have formulated methods for pre-training using graph data [40, 27, 74, 32, 37], rather than 3D point clouds of atom nuclei as in this paper.
To our knowledge, this work is the first to pre-train models for downstream tasks involving molecular property prediction from 3D structures.
Second, denoising has been applied to learning representations in various other contexts [70, 71, 50, 51, 46, 23] and for score-based generative modelling [69, 61, 62, 25, 26, 75, 57].
We make a formal connection between the two in a molecular context in Section 3.2.1.
Lastly, the design of equivariant neural network architectures has been the dominant approach for improving molecular property prediction models recently [ e.g. 33, 39, 56, 65, 29, 53, 1].
Our approach is architecture-agnostic and shown to improve both a GNN and Transformer-based model.
Methodology 3.1 Problem Setup
Molecular property prediction consists of predicting scalar quantities given the structure of one or more molecules as input.
Each data example is a labelled set specified as follows: we are provided with a set of atoms S
= { ( a 1 , p 1 ) , . . .
, ( a | S | , p | S | ) } , where a i ∈ { 1 , . .
118 } and p i ∈ R 3 are the atomic number and 3D position respectively of atom i in the molecule, alongside a label y ∈ R .
We assume that the model, which takes S as input, is any architecture consisting of a backbone, which first processes S to build a latent representation of it, followed by a vertex-level or graph-level “decoder”, that returns per-vertex predictions or a single prediction for the input respectively.
3.2 Pre-training via Denoising
Given a dataset of molecular structures, we pre-train the network by denoising the structures, which operates as follows.
Let D structures =
{ S 1 , . . .
, S n } denote the upstream dataset of equilibrium structures, and let GNN θ denote a graph neural network with parameters θ which takes S ∈ D structures as input and returns per-vertex predictions GNN θ ( S ) =
(ˆ ϵ 1 , . . .
, ˆ ϵ | S | ) .
The precise parameterization of the models we consider in this work is described in Section 3.3 and
Appendix B.
Starting with an input molecule S ∈ D structures , we perturb it by adding i.i.d.
Gaussian noise to its atomic positions p i .
That is, we create a noisy version of the molecule:
[Equation]
˜
[Equation]
ϵ
[Equation]
The noise scale σ is a tuneable hyperparameter (an interpretation of which is given in Section 3.2.1).
We train the model as a denoising autoencoder by minimizing the following loss with respect to θ :
[Equation]
[Table]
The distribution p ( ˜ S, S ) corresponds to sampling a structure S from D structures and adding noise to it according to Equation (1).
Note that the model predicts the noise, not the original coordinates.
Next, we motivate denoising as our pre-training objective for molecular modelling.
3.2.1 Denoising as Learning a Force Field
Datasets in quantum chemistry are typically generated by minimizing expensive-to-compute interatomic forces with methods such as density functional theory (DFT) [44].
We speculate that learning this force field would give rise to useful representations for downstream tasks, since molecular properties vary with forces and energy.
Therefore, a reasonable pre-training objective would be one that involves learning the force field.
Unfortunately, this force field is either unknown or expensive to evaluate, and hence it cannot be used directly for pre-training.
An alternative is to approximate the data-generating force field with one that can be cheaply evaluated and use it to learn good representations – an approach we outline in this section.
Using the well-known link between denoising autoencoders and score-matching [69, 61, 62], we can show that the denoising objective in
Equation (2) is equivalent to learning a particular force field directly from equilibrium structures with some desirable properties.
For clarity, in this subsection we condition on and suppress the atom types and molecule size in our notation, specifying a molecular structure by its coordinates x ∈ R 3 N (with
N as the size of the molecule).
From the perspective of statistical physics, a structure x can be treated as a random quantity sampled from the Boltzmann distribution p physical ( x ) ∝ exp( − E ( x )) , where E ( x ) is the (potential) energy of x .
According to p physical , low energy structures have a high probability of occurring.
Moreover, the per-atom forces are given by ∇ x log p physical ( x ) =
−∇ x E ( x ) , which is referred to as the force field.
Our goal is to learn this force field.
However both the energy function E and distribution p physical are unknown, and we only have access to a set of equilibrium structures x 1 , . . .
, x n that locally minimize the energy E .
Since x 1 , . . .
, x n are then local maxima of the distribution p physical , our main approximation is to replace p physical with a mixture of Gaussians centered at the data:
[Equation]
where we define q σ (˜ x | x i )
=
N (˜ x ; x i , σ 2 I 3 N ) .
This approximation captures the fact that p physical will have local maxima at the equilibrium structures, vary smoothly with x and is computationally convenient.
Learning the force field corresponding to q σ (˜ x ) now yields a score-matching objective:
[Equation]
As shown by Vincent [69], and recently applied to generative modelling [61, 62, 25, 57, 75], this objective is equivalent to the denoising objective.
Specifically, defining q 0 ( x ) = 1 n (cid:80) ni
=1 δ ( x = x i ) to be the empirical distribution and q σ (˜ x , x )
=
q σ (˜ x | x ) q 0 ( x ) , the objective in Equation (3) is equivalent to:
[Equation]
We notice that the RHS corresponds to the earlier denoising loss in Equation (2) (up to a constant factor of 1 /σ applied to GNN θ that can be absorbed into the network).
To summarize, denoising equilibrium structures corresponds to learning the force field that arises from approximating the distribution p physical with a mixture of Gaussians.
Note that we can interpret the noise scale σ as being related to the sharpness of p physical or E around the local maxima x i .
We also remark that the equivalence between Equation (3) and the LHS of Equation (4) does not require q σ ( ˜ x
| x i ) to be a
Gaussian distribution [69], and other choices will lead to different denoising objectives, which we leave as future work.
See Appendix C for technical caveats and details.
3.2.2 Noisy Nodes: Denoising as an Auxiliary Loss
Recently, Godwin et al.
[23] also applied denoising as an auxiliary loss to molecular property prediction, achieving significant improvements on a variety of molecular datasets.
In particular, their approach, called Noisy Nodes, consisted of augmenting the usual optimization objective for predicting y with an auxiliary denoising loss.
They suggested two explanations for why Noisy
Nodes improves performance.
First, the presence of a vertex-level loss discourages oversmoothing [14, 12] of vertex/edge features after multiple message-passing layers – a common problem plaguing
GNNs – because successful denoising requires diversity amongst vertex features in order to match the diversity in the noise targets
ϵ i .
Second, they argued that denoising can aid representation learning by encouraging the network to learn aspects of the input distribution.
Since Noisy Nodes incorporates denoising only as an auxiliary task, the representation learning benefits of denoising are limited to the downstream dataset on which it is used as an auxiliary task.
Our approach is to apply denoising as a pre-training objective on another large (unlabelled) dataset of structures to learn higher-quality representations, which results in better performance.
3.3 GNS and GNS-TAT
The main two models we consider in this work are Graph Net Simulator (GNS) [51], which is a type of GNN, and a better-performing variant we contribute called GNS-TAT.
GNS-TAT makes use of a recently published network transformation method called Tailored Activation Transforms (TAT) [79], which has been shown to prevent certain degenerate behaviors at initialization in deep MLPs/convnets that are reminiscent of oversmoothing in GNNs (and are also associated with training difficulties).
While GNS is not by default compatible with the assumptions of TAT, we propose a novel GNN initialization scheme called “Edge-Delta” that makes it compatible by initializing to zero the weights that carry “messages” from vertices to edges.
This marks the first application of TAT to any applied problem in the literature.
See Appendix B for details.
Experiments
The goal of our experimental evaluation in this section is to answer the following questions.
First, does pre-training a neural network via denoising improve performance on the downstream task compared to training from a random initialization?
Second, how does the benefit of pre-training depend on the relationship between the upstream and downstream datasets?
Our evaluation involves four realistic and challenging molecular datasets, which vary in size, compound compositions (organic or inorganic) and labelling methodology (DFT- or CCSD(T)-generated), as described below.
[Figure]
P e r c e n t a g e o f e l e m e n t s c o v e r e d b y P
C Q M
M v
4.1 Datasets and Training Setup
Datasets.
First, the main dataset we use for pre-training is PCQM4Mv2 [43], which contains 3. million organic molecules, specified by their 3D structures at equilibrium calculated using DFT. {+
 The molecules+} in PCQM4Mv2 only contain one label, however the labels are not used as denoising only requires the structures.
The large scale and diversity of PCQM4Mv2 makes it well-suited for pretraining via denoising.
Second, as a dataset for fine-tuning, we use QM9 [47], which contains around 130,000 small organic molecules and is widely used as a molecular property prediction benchmark [33, 21, 53, 20, 29, 56, 65, 23].
Each molecule is specified by its structure alongside 12 associated molecular property labels.
Third, Open Catalyst 2020 (OC20) [13] is a recent large benchmark of interacting surfaces and adsorbates relevant to catalyst discovery.
OC20 contains various tasks, such as predicting the relaxed state energy from an initial high-energy structure (IS2RE).
We explore different combinations of upstream and downstream tasks as described in Section 4.3.
Lastly, DES15K [18] is a small dataset we use for fine-tuning, which contains around 15,000 dimer geometries ( i.e. molecule pairs) with non-covalent molecular interactions.
Each pair is labelled with its interaction energy computed using the gold-standard CCSD(T) method [3].
CCSD(T) is usually both more expensive and accurate than DFT, which is used for all aforementioned datasets.
See Appendix F for further details and a discussion about the choice of using DFT-generated structures for pre-training.
Figure 2 (right) shows what percentage of elements appearing in each of QM9, OC20 and DES15K also appear in PCQM4Mv2.
Whereas QM9 is fully covered by PCQM4Mv2, we observe that
DES15K has less element overlap with PCQM4Mv2 and less than < 30% of elements in OC20 are contained in PCQM4Mv2.
This is owing to the fact that surface molecules in OC20 are inorganic lattices, none of which appear in PCQM4Mv2.
This suggests that we can expect least transfer from PCQM4Mv2 to OC20.
We also compare PCQM4Mv2 and QM9 in terms of the molecular compositions, i.e. the number of atoms of each element, that appear in each.
Due to presence of isomers, both datasets contain multiple molecules with the same composition.
For each molecular composition in QM9, Figure 2 (left) shows its frequency in both QM9 and PCQM4Mv2.
We observe that most molecular compositions in QM9 also appear in PCQM4Mv2.
We also remark that since pre-training is self-supervised using only unlabelled structures, test set contamination is not possible – in fact, PCQM4Mv2 does not have most of the labels in QM9.
Training setup.
GNS/GNS-TAT were implemented in JAX [8] using Haiku and Jraph [24, 22].
All experiments were averaged over 3 seeds.
Detailed hyperparameter and hardware settings can be found in Appendices G and H.
4.2 Results on QM
We evaluate two variants of our model on QM9 in Table 1, GNS-TAT with Noisy Nodes trained from a random initialization versus pre-trained parameters.
Pre-training is done on PCQM4Mv2 via
[Table] denoising. For best performance on QM9, we found that using atom type masking and prediction during pre-training additionally helped [27].
We fine-tune a separate model for each of thetargets, as usually done on QM9, using a single pre-trained model.
This is repeated for three seeds (including pre-training).
Following customary practice, hyperparameters, including the noise scale for denoising during pre-training and fine-tuning, are tuned on the HOMO target and then kept fixed for all other targets.
We first observe that GNS-TAT with Noisy Nodes performs competitively with other models and significantly improves upon GNS with Noisy Nodes, revealing the benefit of the TAT modifications.
Utilizing pre-training then further improves performance across all targets, achieving a new state-of-the-art compared to prior work for 10 out of 12 targets.
Interestingly, for the electronic spatial extent target (cid:10) R 2 (cid:11) , we found GNS-TAT to perform worse than other models, which may be due to the optimal noise scale being different from that of other targets.
4.3 Results on OC
Next, we consider the Open Catalyst 2020 benchmark focusing on the downstream task of predicting the relaxed energy from the initial structure (IS2RE).
We compared GNS with Noisy Nodes trained from scratch versus using pre-trained parameters.
We experimented with two options for pre-training: (1) pre-training via denoising on PCQM4Mv2, and (2) pre-training via denoising on OC20 itself.
For the latter, we follow Godwin et al.’s
[2022] approach of letting the denoising target be the relaxed structure, while the perturbed input is a random interpolation between the initial and relaxed structures with added Gaussian noise – this corresponds to the IS2RS task with additional noise.
As shown in Figure 3 (left), pre-training on PCQM4Mv2 offers no benefit for validation performance on IS2RE, however pre-training on OC20 leads to considerably faster convergence but the same final performance.
The lack of transfer from PCQM4Mv2 to OC20 is likely due to the difference in nature of the two datasets and the small element overlap as discussed in Section 4.1 and Figure(right).
On the other hand, faster convergence from using parameters pre-trained on OC20 suggests that denoising learned meaningful features.
Unsurprisingly, the final performance is unchanged since the upstream and downstream datasets are the same in this case, so pre-training with denoising is identical to the auxiliary task of applying Noisy Nodes.
The performance achieved is also competitive with other models in the literature as shown in Table 7.
4.4 Results on DES15K
In our experiments so far, all downstream tasks were based on DFT-generated datasets.
While DFT calculations are more expensive than using neural networks, they are relatively cheap compared to even higher quality methods such as CCSD(T) [3].
In this section, we evaluate how useful pre-training on DFT-generated structures from PCQM4Mv2 is when fine-tuning on the recent DES15K dataset which contains higher quality CCSD(T)-generated interaction energies.
Moreover, unlike QM9, inputs from DES15K are systems of two interacting molecules and the dataset contains only around 15, examples, rendering it more challenging.
We compare the test performance on DES15K achieved by
GNS-TAT with Noisy Nodes when trained from scratch versus using pre-trained parameters from
[Figure]
PCQM4Mv2.
As a baseline, we also include pre-training on PCQM4Mv2 using 2D-based AttrMask [27] by masking and predicting atomic numbers.
Figure 3 (right) shows that using Noisy Nodes significantly improves performance compared to training from scratch, with a further improvement resulting from using pre-training via denoising.
AttrMask underperforms denoising since it likely does not fully exploit the 3D structural information.
Importantly, this shows that pre-training by denoising structures obtained through relatively cheap methods such as DFT can even be beneficial when fine-tuning on more expensive and smaller downstream datasets.
See Appendix I.2 for similar results on another architecture.
Analysis 5.1 Pre-training a Different Architecture
To explore whether pre-training is beneficial beyond GNS/GNS-TAT, we applied pre-training via denoising to the TorchMD-NET architecture [65].
TorchMD-NET is a transformer-based architecture whose layers maintain per-atom scalar features x i ∈ R F and vector features v i ∈ R 3 × F , where
F is the feature dimension, that are updated in each layer using a self-attention mechanism.
We implemented denoising by using gated equivariant blocks [72, 56] applied to the processed scalar and vector features.
The resulting vector features are then used as the noise prediction.
[Table]
In Table 2, we evaluate the effect of adding
Noisy Nodes and pre-training to the architecture on the HOMO and LUMO targets in QM9.
Pre-training yields a boost in performance, allowing the model to achieve
SOTA results.
Note that the results shown for TorchMD-NET are from our runs using Thölke and De Fabritiis’s [2022] opensource code, which led to slightly worse results than their published ones (our pre-training results still outperform their published results).
Our code for experiments on TorchMD-NET is available open-source.
5.2 Varying Dataset Sizes
We also investigate how downstream test performance on the HOMO target in QM9 varies as a function of the number of upstream and downstream training examples.
First, we compare the performance of GNS-TAT with Noisy Nodes either trained from scratch or using pre-trained parameters for different numbers of training examples from QM9; we also include the performance of just GNS-TAT.
As shown in Figure 4 (left), pre-training improves the downstream performance for all dataset sizes.
The difference in test MAE also grows as the downstream training data reduces.
[Figure] Second, we assess the effect of varying the amount of pre-training data while fixing the downstream dataset size for both GNS and GNS-TAT as shown in Figure 4 (middle).
For both models, we find that downstream performance generally improves as upstream data increases, with saturating performance for GNS-TAT.
More upstream data can yield better quality representations
[Figure]
5.3 Varying Model Size
We study the benefit of pre-training as models are scaled up on large downstream datasets.
Recall that the S2EF dataset in OC20 contains around 130 million DFT evaluations for catalytic systems, providing three orders of magnitude more training data than QM9.
We compare the performance of four GNS models with sizes ranging from 10 million to 1.2 billion parameters scaled up by increasing the hidden layer sizes in the MLPs.
Each is pre-trained via denoising using the trajectories provided for the IS2RE/IS2RS tasks as described in Section 4.3.
We also compare this to amillion parameter variant of GNS trained from scratch.
As shown in Figure 4 (right), the pre-trained models continue to benefit from larger model sizes.
We also observe that pre-training is beneficial, as the model trained from scratch underperforms in comparison: the 130 million parameters model trained from scratch is outperformed by a pre-trained model of less than half the size.
5.4 Additional Experimental Results
Additional experiments can be found in Appendix I.
Freezing pre-trained parameters. Appendix I.1 demonstrates that the features learned during pre-training are significantly more discriminative than random features, by considering the effect of freezing the parameters of the model backbone during fine-tuning.
Denoising improves force prediction.
In Appendix I.4, we empirically verify that pre-training via denoising also improves atomic force prediction, as suggested by our analysis in Section 3.2.1.
We apply TorchMD-NET to MD17 (aspirin) and GNS to OC20 for force prediction tasks, finding improved performance as a result of pre-training in both cases.
Conclusion
We investigated pre-training neural networks by denoising in the space of 3D molecular structures.
We showed that denoising in this context is equivalent to learning a force field, motivating its ability to learn useful representations and shedding light on successful applications of denoising in other works [23].
This technique enabled us to utilize existing large datasets of 3D structures for improving performance on various downstream molecular property prediction tasks, setting a new SOTA in some cases such as QM9.
More broadly, this bridges the gap between the utility of pre-training in vision/NLP and molecular property prediction from structures.
We hope that this approach will be particularly impactful for applications of deep learning to scientific problems.
Limitations of our work and avenues for future work are discussed in Appendix D.
A Related Work
Pre-training of GNNs.
Various recent works have formulated methods for pre-training using graph data [40, 27, 74, 32], rather than 3D point clouds of atom nuclei as in this paper.
Approaches based on contrastive methods rely on learning representations by contrasting different views of the input graph [63, 68, 78, 37] or bootstrapping [64].
Autoregressive or reconstruction-based approaches, such as ours, learn representations by requiring the model to predict aspects of the input graph [27, 28, 49, 36].
Most methods in the current literature are not designed to handle 3D structural information, focusing instead on 2D graphs.
The closest work to ours is GraphMVP [37], where 3D structure is treated as one view of a 2D molecule for the purpose of upstream contrastive learning.
Their work focuses on downstream tasks that only involve 2D information, while our aim is to improve downstream models for molecular property prediction from 3D structures.
After the release of this pre-print, similar ideas have been studied by Jiao et al.
[30] and Liu et al.
Denoising, representation learning and score-matching.
Noise has long been known to improve generalization in machine learning [59, 7].
Denoising autoencoders have been used to effectively learn representations by mapping corrupted inputs to original inputs [70, 71].
Specific to GNNs [5, 54, 10], randomizing input graph features has been shown to improve performance [27, 52].
Applications to physical simulation also involve corrupting the state with Gaussian noise [50, 51, 46].
Our work builds on Noisy Nodes [23], which incorporates denoising as an auxiliary task to improve performance, indicating the effectiveness of denoising for molecular property prediction ( cf.
Section 3.2.2).
Denoising is also closely connected to score-matching [69], which has become popular for generative modelling [61, 62, 25, 26, 75, 57].
We also rely on this connection to show that denoising structures corresponds to learning a force field.
Equivariant neural networks for 3D molecular property prediction.
Recently, the dominant approach for improving models for molecular property prediction from 3D structures has been through the design of architectures that incorporate roto-translational inductive biases into the model, such that the outputs are invariant to translating and rotating the input atomic positions.
A simple way to achieve this is to use roto-translation invariant features as inputs, such as inter-atomic distances [55, 67], angles [34, 33, 58, 39], or the principal axes of inertia [23].
There is also broad literature on equivariant neural networks, whose intermediate activations transform accordingly with rototranslations of inputs thereby naturally preserving inter-atomic distance and orientation information.
Such models can be broadly categorized into those that are specifically designed for molecular property prediction [65, 56, 6, 1, 42] and general-purpose architectures [53, 20, 29, 66, 35, 9].
Our pre-training technique is architecture-agnostic, and we show that it can be applied to enhance performance in both a GNN-based architecture [51] and a Transformer-based one [65].
We conjecture that similar improvements will hold for other models.
B Architectural Details B.1 Standard GNS
As our base model architecture we chose a Graph Net Simulator (GNS) [51], which consists of an
E NCODER which constructs a graph representation from the input S , a P ROCESSOR of repeated message passing blocks that update the latent graph representation, and a D ECODER which produces predictions.
Our implementation follows Godwin et al.’s
[2022] modifications to enable molecular and graph-level property predictions which has been shown to achieve strong results across different molecular prediction tasks without relying on problem-specific features.
In the E NCODER , we represent the set of atoms S
= { ( a 1 , p 1 ) , ( a 2 , p 2 ) , . . .
, ( a | S | , p | S | ) } as a directed graph G =
( V, E ) where V = { v 1 , v 2 , . .
. , v | S | } and E = { e i,j } i,j are the sets of “featurized” vertices and edges, respectively.
Edges e i,j ∈ E are constructed whenever the distance between the i -th and j -th atoms is less than the connectivity radius R cut , in which case we connect v i and v j with a directed edge e i,j from i to j that is a featurization of the displacement vector p j − p i .
Meanwhile, for the i -th atom, v i is given by a learnable vector embedding of the atomic number a i .
The P ROCESSOR consists of L message-passing steps that produce intermediate graphs G 1 , . . .
, G L (with the same connectivity structure as the initial one).
Each of these steps computes the sum of a shortcut connection from the previous graph, and the application of an Interaction Network [4].
Interaction Networks first update each edge feature by applying an “edge update function” to a combination of the existing feature and the features of the two connected vertices.
They then update each vertex feature by applying a “vertex update function” to a combination of the existing feature and the (new) edge features of incoming edges.
In GNS, edge update functions are 3 hidden layer fullyconnected MLPs, using a “shifted softplus” (ssp ( x ) = log(0 .
5 e x + 0 . 5) ) activation function, applied to the concatenation of the relevant edge and vertex features, followed by a layer normalization layer.
Vertex update functions are similar, but are applied to the concatenation of the relevant vertex feature and sum over relevant edge features.
In our implementation of GNS we applied the same P ROCESSOR in sequence three times (with shared parameters), with the output of each being decoded to produce a prediction and corresponding loss value.
The loss for the whole model is then given by the average of these.
(Test-time predictions are meanwhile computed using only the output of the final P ROCESSOR .)
The D ECODER is responsible for computing graph-level and vertex-level predictions from the output of each P ROCESSOR .
Vertex-level predictions, such as noise as described in Section 3.2, are decoded using an MLP applied to each vertex feature.
Graph-level predictions ( e.g. energies) are produced by applying an MLP to each vertex feature, aggregating the result over vertices (via a sum), and then applying another MLP to the result.
B.2 GNS with Tailored Activation Transformation (GNS-TAT)
Tailored Activation Transformation (TAT) [79] is a method for initializing and transforming neural networks to make them easier to train, and is based on a similar method called Deep Kernel Shaping (DKS)
TAT controls the propagation of “q values”, which are initialization-time approximations to dimension-normalized squared norms of the network’s layer-wise activation vectors, and “c values”, which are cosine similarities between such vectors (for different inputs).
In other words, q values approximate ∥
z ( x ) ∥ 2 / dim( z ( x )) , where z ( x ) denotes a layer’s output as a function of the network’s input x , and c values approximate z ( x ) ⊤ z ( x ′ ) / ( ∥ z ( x ) ∥∥
z ( x ′ ) ∥ ) , where x ′ is another possible network input.
In standard deep networks, c values will converge to a constant value c ∞ ∈ [0 , 1] , so that “geometric information” is lost, which leads to training difficulties [41].
DKS/TAT prevents this convergence through a combination of careful weight initialization, and transformations to the network’s activation functions and sum/average layers.
Oversmoothing [14, 12, 48, 81, 76, 80, 17] is a phenomenon observed in GNN architectures where vertex/edge features all converge to approximately the same value with depth, and is associated with training difficulties.
It is reminiscent of how, when c ∞ = 1 , feature vectors will converge with depth to a constant input-independent vector in standard deep networks.
It therefore seems plausible that applying TAT to GNNs may help with the oversmoothing problem and thus improve training performance. [Figure] [Figure]
⨁ … …
[Figure] [Figure] Unfortunately, the GNS architecture violates two key assumptions of TAT.
Firstly, the sums over edge features (performed in the vertex update functions) violate the assumption that all sum operations must be between the outputs of linear layers with independently sampled initial weights.
Secondly,
GNS networks have multiple inputs for which information needs to be independently preserved and propagated to the output, while DKS/TAT assumes a single input (or multiple inputs whose representations evolve independently in the network).
To address these issues we introduce a new initialization scheme called “Edge-Delta”, which initializes to zero the weights that multiply incoming vertex features in the edge update functions (and treats these weights as absent for the purpose of computing the initial weight variance).
This approach is inspired by the use of the “Delta initialization" [2, 73] for convolutional networks in DKS/TAT, which initializes filter weights of the non-central locations to zero, thus allowing geometric information, in the form of c values, to propagate independently for each location in the feature map.
When using the Edge-Delta initialization, edge features propagate independently of each other (and of vertex features), through what is essentially a standard deep residual network (with edge update functions acting as the residual branches), which we will refer to as the “edge network”.
Given the use of Edge-Delta we can then apply TAT to GNS as follows 6 .
First, we replace GNS’s activation functions with TAT’s transformed Leaky-ReLU activation functions (or “Tailored ReLUs”), which we compute with TAT’s η parameter set to 0 . 8 , and its “subnetwork maximizing function” defined on the edge network 7 .
We also replace each sum involving shortcut connections with weighted sums, whose weights are 0 .
9 and √ 1 − 0 .
9 2 for the shortcut and non-shortcut branches respectively.
We retain the use of layer normalization layers in the edge/vertex update functions, but move them to before the first fully-connected layer, as this seems to give the best performance.
As required by
TAT, we use a standard Gaussian fan-in initialization for the weights, and a zero initialization for the biases, with Edge-Delta used only for the first linear layer of the edge update functions.
Finally, we replace the sum used to aggregate vertex features in the D ECODER with an average.
See Figuresand 6 for an illustration of these changes.
We experimented with an analogous “Vertex-Delta” initialization, which initializes to zero weights in the vertex update functions that multiply summed edge features, but found that Edge-Delta gave the best results.
This might be because the edge features, which encode distances between vertices (and are best preserved with the Edge-Delta approach), are generally much more informative than the vertex features in molecular property prediction tasks.
We also ran informal ablation studies, and found that each of our changes to the original GNS model contributed to improved results, with the use of Edge-Delta and weighted shortcut sums being especially important.
C Denoising as Learning a Force Field
We specify a molecular structure as x = ( x (1) , . . .
, x ( N ) )
∈ R 3 N , where x ( i ) ∈ R 3 is the coordinate of atom i .
Let E ( x ) denote the total (potential) energy of x , such that −∇ x E ( x ) are the forces on the atoms.
As discussed in Section 3.2.1, learning the force field, i.e. the mapping x (cid:55)→ −∇ x E ( x ) , is a reasonable pre-training objective.
Furthermore, learning the force field can be viewed as scorematching if we define the distribution p physical ( x ) ∝ exp( − E ( x )) and observe that the score of p physical is the force field: ∇ x log p physical ( x ) =
−∇ x E ( x ) .
However, a technical caveat is that p physical is an improper probability density, because it cannot be normalized due to the translation invariance of E .
Writing the translation of a structure as x + t : =
( x (1) + t , . . .
, x ( N ) + t ) where t ∈ R 3 is a constant vector, we have E ( x + t ) =
E ( x ) .
This implies that the normalizing constant (cid:82) R 3 N p physical ( x ) d x diverges to infinity.
To remedy this, we can restrict ourselves to the (3 N − 3)
-dimensional subspace V := { x ∈ R 3 N | (cid:80) i
x
( i ) = 0 } ⊆
R 3 N consisting of the mean-centered structures, over which p physical can be defined as a normalizable distribution.
Proceeding similarly as Section 3.2.1, let x 1 , . . .
, x n ∈ V be a set of mean-centered equilibrium structures.
For any ˜ x ∈ V , we now approximate
[Equation]
where the Gaussian distributions q σ (˜ x | x i ) are defined on V as:
[Equation]
For convenience, we have expressed structures as vectors in the ambient space R 3 N , however they are restricted to lie in the smaller space V .
Note that the normalizing constant accounts for the fact that
V is (3 N − 3) -dimensional.
As before, we define q 0 ( x ) = 1 n (cid:80) ni
=1 δ ( x = x i ) to be the empirical distribution and q σ (˜ x , x ) =
q σ (˜ x | x ) q 0 ( x ) .
The score-matching objective is given by:
[Equation]
where the expectation is now over V .
As shown by Vincent [69], minimizing the objective above is equivalent to the minimizing the following objective:
[Equation]
This is recognized as a denoising objective, because ∇ ˜ x log q σ (˜ x | x ) =
( x − ˜ x ) /σ 2 .
A practical implication of this analysis is that the noise ( x − ˜ x ) /σ 2 ∈ V should be mean-centered, which is intuitive since it is impossible to predict a translational component in the noise.
We include a proof of the equivalence between Equations (5) and (6) for completeness:
Proposition 1 (Vincent [69]) .
The minimization objectives J 1 ( θ ) and J 2 ( θ ) are equivalent.
Proof.
We first observe:
[Equation]
where C 1 , C 2 are constants independent of θ .
Therefore, it suffices to show that the middle terms on the RHS are equal.
Since expectations over q σ (˜ x ) and q σ (˜ x , x ) are restricted to V ⊆ R 3 N , we apply a change of basis to write them as integrals against the (3 N − 3) -dimensional
Lebesgue measure.
Pick an orthonormal basis { v 1 , . .
. , v 3 N − 3 } ⊆ R 3 N for V and let P V =
[ v 1 , . . .
, v 3 N − 3 ] ∈ R 3 N × 3( N − 1) be the projection matrix, so
z = P ⊺ V x expresses a mean-centered structure x in terms of the coordinates of the chosen basis for V .
Noting that P V has orthonormal columns and that it yields a bijection between V and R 3 N − 3 , we calculate:
[Equation]
D Limitations & Future Work
We have shown that pre-training can significantly improve performance for various tasks.
One additional advantage of pre-trained models is that they can be shared in the community, allowing practitioners to fine-tune models on their datasets.
However, unlike vision and NLP, molecular networks vary widely and the community has not yet settled on a “standard” architecture, making pre-trained weights less reusable.
Moreover, the success of pre-training inevitably depends on the relationship between the upstream and downstream datasets.
In the context of molecular property prediction, understanding what aspects of the upstream data distribution must match the downstream data distribution for transfer is an important direction for future work.
More generally, pre-training models on large datasets incurs a computational cost.
However, our results show that pre-training for 3D molecular prediction does not require the same scale as large NLP and vision models.
We discuss considerations on the use of compute and broader impact in Appendix E.
E Broader Impact
Who may benefit from this work?
Molecular property prediction works towards a range of applications in materials design, chemistry, and drug discovery.
Wider use of pre-trained models may accelerate progress in a similar manner to how pre-trained language or image models have enabled practitioners to avoid training on large datasets from scratch.
Pre-training via denoising is simple to implement and can be immediately adopted to improve performance on a wide range of molecular property prediction tasks.
As research converges on more standardized architectures, we expect shared pre-trained weights will become more common across the community.
Potential negative impact and ethical considerations.
Pre-training models on large structure datasets incurs additional computational cost when compared to training a potentially smaller model with less capacity from scratch.
Environmental mitigation should be taken into account when pretraining large models [45].
However, the computational cost of pre-training can and should be offset by sharing pre-trained embeddings when possible.
Moreover, in our ablations of upstream dataset sizes for GNS-TAT, we observed that training on a subset of PCQM4Mv2 was sufficient for strong downstream performance.
In future work, we plan to investigate how smaller subsets with sufficient diversity can be used to minimize computational requirements, e.g. by requiring fewer gradient steps
. F Datasets
PCQM4Mv2.
The main dataset we use for pre-training is PCQM4Mv2 [43] (license: CC BY 4.0), which contains 3,378,606 organic molecules, specified by their 3D structures at equilibrium (atom types and coordinates) calculated using DFT.
Molecules in PCQM4Mv2 have around 30 atoms on average and vary in terms of their composition, with the dataset containing 22 unique elements in total.
The molecules in PCQM4Mv2 only contain one label, unlike e.g. QM9, which containslabels per molecule, however we do not use these labels as denoising only requires the structures.
QM9.
QM9 is a dataset [47] (license: CCBY 4.0) with approximately 130,000 small organic molecules containing up to nine heavy C, N, O, F atoms, specified by their structures.
Each molecule has 12 different labels corresponding to different molecular properties, such as highest occupied molecular orbital (HOMO) energy and internal energy, which we use for fine-tuning.
OC20.
Open Catalyst 2020 [13] (OC20, license: CC Attribution 4.0) is a recent large benchmark containing trajectories of interacting surfaces and adsorbates that are relevant to catalyst discovery and optimization.
This dataset contains three tasks: predicting the relaxed state energy from the initial structure (IS2RE), predicting the relaxed structure from the initial structure (IS2RS) and predicting the energy and forces given the structure at any point in the trajectory (S2EF).
For IS2RE and IS2RS, there are 460,000 training examples, where each data point is a trajectory of a surface-adsorbate molecule pair starting with a high-energy initial structure that is relaxed towards a low-energy, equilibrium structure.
For S2EF, there are 113 million examples of (non-equilibrium) structures with their associated energies and per-atom forces.
DES15K. DES15K [18] (license: CC0 1.0) is a small dataset containing around 15,000 interacting molecule pairs, specifically dimer geometries with non-covalent molecular interactions.
Each pair is labelled with the associated interaction energy computed using the coupled-cluster method with single, double, and perturbative triple excitations (CCSD(T)) [3], which is widely regarded as the gold-standard method in electronic structure theory.
Usage of DFT-generated structures for pre-training.
The structures in PCQM4Mv2 are obtained using DFT calculations, which in principle could have also been used to generate labels for molecular properties in DFT-generated downstream datasets, such as QM9.
However, there are multiple reasons why denoising remains a desirable pre-training objective in such settings.
First, although there is a computational cost for generating datasets such as PCQM4Mv2, it is now openly available and part of our aim is to understand how to leverage such datasets for tasks on other datasets (analogous to how ImageNet is expensive to build, but once it is available, it is important to understand how it can improve downstream performance on other datasets).
Second, even if PCQM4Mv2 contained all the labels in QM9, pre-training via denoising structures allows one to pre-train a single, label-agnostic model which can be individually fine-tuned on any of the targets in QM9.
This is substantially cheaper than per-target pre-training, and the resulting pre-trained model is also re-useable for differing needs and downstream tasks.
In other settings where the downstream dataset is generated using more expensive methods than DFT, pre-training via denoising DFT-relaxed structures can also be helpful and has a clear benefit, as shown by our experiment on the CCSD(T)-generated dataset DES15K where denoising on PCQM4Mv improves performance for a downstream task involving a “higher” level of theory such as CCSD(T).
Generally, we emphasize that the methodology of denoising structures can be applied to any dataset of structures (regardless of whether they are computed using DFT or not).
We hope that denoising will be useful in the future for learning representations by pre-training on structures obtained through other methods such as experimental data (in which case labeling may be expensive) and databases generated by other models such as AlphaFold (where only structures are available) [31].
G Experiment Setup and Compute Resources
Below, we list details on our experiment setup and hardware resources used.
GNS & GNS-TAT.
GNS-TAT training for QM9, PCQM4Mv2 and DES15K was done on a cluster of 16 TPU v3 devices and evaluation on a single V100 device.
GNS training for OC20 was done on 8 TPU v4 devices, with the exception of the 1.2 billion parameters variant of the model, which was trained on 64 TPU v4 devices.
Pre-training on PCQM4Mv2 was executed for 3 · 10 5 gradient updates (approximately 1.5 days of training).
Fine-tuning experiments were run until convergence for
QM9 ( 10 6 gradient updates taking approximately 2 days) and DES15K ( 10 5 gradient updates taking approximately 4 hours) and stopped after 5 · 10 5 gradient updates on OC20 (2.5 days) to minimize hardware use (the larger models keep benefiting from additional gradient updates).
TorchMD-NET.
We implemented denoising for TorchMD-NET on top of Thölke and De Fabritiis’s [2022] open-source code.
8 Models were trained on QM9 using data parallelism over two NVIDIA
RTX 2080Ti GPUs.
Pre-training on PCQM4Mv2 was done using three GPUs to accommodate the larger molecules while keeping the batch size approximately the same as QM9.
All hyperparameters except the learning rate schedule were kept fixed at the defaults.
Pre-training took roughly 24 hours, whereas fine-tuning took around 16 hours.
Hyperparameter optimization.
We note that effective pre-training via denoising requires sweeping noise values, as well as loss co-efficients for denoising and atom type recovery.
For GNS/GNS-TAT, we relied on the hyperparameters published by Godwin et al.
[23] but determined new noise values for pre-training and fine-tuning by tuning over a grid of approximately 5 values each on PCQM4Mv and QM9 (for the HOMO target).
We used the same values for DES15K without modification.
We also ran a similar number of experiments to determine cosine cycle parameters for learning rates.
H Hyperparameters
We report the main hyperparameters used for GNS and GNS-TAT below.
[Table]
[Table]
[Table]
[Table]
r
[Table]
= 0 , σ =
[Table]
[Table]
[Table]
[Table]
[Table]
[Table] [Table] [Table] Atom type loss coefficient 4.
I Additional Experimental Results I.1 Freezing Pre-trained Parameters
[Figure]
ones
We perform an experiment to assess how useful the features learned by pre-training are if they are not fine-tuned for the downstream task but kept fixed instead.
Specifically, on the
HOMO target in QM9, we freeze the backbone of the model and fine-tune only the decoder ( cf.
Appendix B).
To evaluate this, we compare it to using random parameters from initialization for the model’s backbone, which allows us to isolate how useful the pre-trained features are.
As described in Appendix B, the decoder is a simple module involving no message-passing.
Figure 7 shows that only training the decoder while keeping the pre-trained parameters fixed results in test MAE of 40 meV, which is worse than fine-tuning the entire model but substantially better than the performance of >100 meV in test MAE resulting from training the decoder when the remaining parameters are randomly initialized.
This suggests that the features learned by denoising are more discriminative for downstream prediction than random features.
We note that training only the decoder is also substantially faster than training the entire network – one batch on a single V100 GPU takes 15ms, which is 50 × faster than one batch using 16 TPUs for the full network.
[Table]
I.2 Performance of TorchMD-NET on DES15K
In addition to the experiments involving GNS-TAT on DES15K in Section 4.
[Table]
we also consider the performance of TorchMD-NET on DES15K with and without pre-training on PCQM4Mv2.
As shown in Table
[Table]
TorchMD-NET outperforms GNS-TAT when trained from scratch
[Table]
and pre-training then yields a further boost in performance as with GNS-TAT.
I.3 Comparison of OC20 IS2RE Pre-training Performance with Other Architectures
Table 7 compares the performance of our models with various other architectures proposed in prior work.
GNS yields SOTA performance on each of the four validation sets for the direct IS2RE task.
As discussed in Section 4.3 and shown in Figure 3 (left), the model pre-trained on OC20 itself achieves
SOTA performance with faster convergence than all other GNS variants.
Note that since we pre-train on OC20 itself, the pre-trained model performs equally well at convergence as the model trained from scratch with noisy nodes ( cf.
Section 4.3).
I.4 Pre-training via Denoising Improves Force Prediction
As shown in Section 3.2.1, denoising structures corresponds to learning an approximate force field directly from equilibrium structures.
We explore whether pre-training via denoising would therefore also improve models trained to predict atomic forces.
First, we compare the performance of TorchMDNET for force prediction on the MD17 (aspirin) dataset with and without pre-training on PCQM4Mv2.
Table 8 shows that pre-training improves force prediction.
[Table]
[Table] [Table] [Table] Second, the OC20 dataset contains a force prediction task (S2EF), where a model is trained to predict point-wise energy and forces (each point being a single DFT evaluation during a relaxation trajectory) from a given 3D structure.
Tables 9 andshow the performance of GNS when trained from scratch vs. pre-trained via denoising on equilibrium structures in OC20, as described in Section 4.3.
We show two metrics for measuring force prediction performance on each of the four validation datasets: mean absolute error (lower is better) and cosine similarity (higher is better).
We observe that the pre-trained model improves upon the model trained from scratch for both metrics and all four validation datasets, with improvements up to 15%.
[Table]
[Table]