S IMPLE GNN N OISE R EGULARISATION FOR 3D M OLECULAR P ROPERTY P REDICTION AND B EYOND
Anonymous authors
Paper under double-blind review
A BSTRACT
Graph Neural Networks (GNNs) have been proven effective across a wide range of molecular property prediction and structured learning problems.
However, their efﬁciency is known to be hindered by practical challenges such as oversmoothing.
We introduce “Noisy Nodes”, a very simple technique for improved training of
GNNs, in which we corrupt the input graph with noise and add a noise correcting node-level loss.
Adding noise helps overﬁtting and the noise correction loss helps ameliorate oversmoothing by encouraging diverse node latents.
Our regulariser applies well-studied methods in simple, straightforward ways which allows even generic architectures not designed for quantum chemistry to achieve state of the art results.
We also demonstrate the effectiveness of Noisy Nodes with non-spatial architectures on Open Graph Benchmark (OGB) datasets.
Our results suggest
Noisy Nodes can serve as a complementary building block in the GNN toolkit for 3D molecular property prediction and beyond.
I NTRODUCTION
While Graph Neural Networks have demonstrated success in a wide variety of tasks (Zhou et al., 2020a; Wu et al., 2020; Bapst et al., 2020; Schütt et al., 2017; Klicpera et al., 2020a), it has been proposed that in practice “oversmoothing” (Chen et al., 2019) limits their ability to beneﬁt from overparametrization.
Here we propose a simple noise regulariser, Noisy Nodes, and demonstrate how it overcomes these challenges across a range of datasets and architectures, achieving top results on
OC20 IS2RS & IS2RE direct, QM9 and OGBG-PCQM4Mv1.
Graph Neural Networks (GNNs) are a family of neural networks that operate on graph structured data by iteratively passing learned messages over the graph’s structure (Scarselli et al., 2009;
Bronstein et al., 2017; Gilmer et al., 2017; Battaglia et al., 2018; Shlomi et al., 2021).
Oversmoothing is a phenomenon where a GNN’s latent node representations become increasing indistinguishable over successive steps of message passing (Chen et al., 2019).
Once these representations are oversmoothed, further message-passing cannot improve expressive capacity, and so performance does not improve.
We speculate that GNN architectures applied on the benchmarks we study here suffer from these effects and could have improved accuracy if ameliorated.
Our “Noisy Nodes” method is an elegant, i.e. extremely simple, technique for regularising GNNs and associated training procedures.
During training, our noise regularisation approach corrupts the input graph’s attributes with noise, and adds a per-node noise correction term.
We posit that our Noisy
Nodes approach is effective because the model is rewarded for maintaining and reﬁning distinct node representations through message passing to the ﬁnal output, which causes it to resist oversmoothing.
Like denoising autoencoders, it may also encourage the model to explicitly learn the manifold on which the uncorrupted input graph’s features lie, analogous to a form of representation learning.
When applied to 3D molecular prediction tasks, it encourages the model to distinguish between low and high energy states.
We ﬁnd that applying Noisy Nodes reduces oversmoothing for shallower networks, and allows us to see improvements with added depth for some architectures, even on tasks for which depth was assumed to be unhelpful.
This study’s approach is to investigate the combination of Noisy Nodes with generic, popular baseline
GNN architectures.
For 3D Molecular prediction we use a standard architecture working on 3D points developed for particle ﬂuid simulations, the Graph Net Simulator (GNS) (Sanchez-Gonzalez* et al., 2020), which has also been used for molecular property prediction (Hu et al., 2021b).
Without using Noisy Nodes the GNS is not a competitive model, but using Noisy Nodes allows the GNS to achieve top performance on three 3D molecular property prediction tasks: the OC20 IS2RE direct task by 43% over previous work, 12% on OC20 IS2RS direct, and top results on 3 out of 12 of the QM9 tasks.
For non-spatial GNN benchmarks we test a MPNN (Gilmer et al., 2017) on OGBG-MOLPCBA and OGBG-PCQM4M (Hu et al., 2021a) and again see signiﬁcant improvements.
Finally, we applied Noisy Nodes to a GCN (Kipf & Welling, 2016), arguably the most popular and simple GNN, trained on OGBN-Arxiv and see similar results.
These results suggest Noisy Nodes can serve as a complementary GNN building block for 3D molecular property prediction and beyond.
R ELATED W ORK
Oversmoothing.
Recent work has aimed to understand why it is challenging to realise the beneﬁts of training deeper GNNs (Wu et al., 2020).
A key contribution has been the analysis of “oversmoothing” which describes how all node features become almost identical in GCNs after a few layers.
Since ﬁrst being noted in (Li et al., 2018) oversmoothing has been studied extensively and regularisation techniques have been suggested to overcome it
(Chen et al., 2019; Cai & Wang, 2020;
Rong et al., 2019;
Zhou et al., 2020b; Yang et al., 2020; Do et al., 2021).
A recent paper, (Li et al., 2021), ﬁnds, as in previous work, (Li et al., 2019; 2020), the optimal depth for some datasets they evaluate on to be far lower (5 for OGBN-Arxiv from the Open Graph Benchmark (Hu et al., 2020a), for example) than the 1000 layers possible.
Denoising Models.
Training neural networks with noise has a long history (Sietsma & Dow, 1991;
Bishop, 1995).
Of particular relevance are Denoising Autoencoders (Vincent et al., 2008) in which an autoencoder is trained to map corrupted inputs ˜ x to uncorrupted inputs x .
Denoising Autoencoders have found particular success as a form of pre-training for representation learning (Vincent et al., 2010).
More recently, in research applying GNNs to simulation (Sanchez-Gonzalez et al., 2018;
Sanchez-Gonzalez* et al., 2020; Pfaff et al., 2020)
Gaussian noise is added during training to input positions of a ground truth simulator to mimic the distribution of errors of the learned simulator.
Pre-training methods (Devlin et al., 2019; You et al., 2020; Thakoor et al., 2021) are another similar approach; most similarly to our method Hu et al.
(2020b) apply a reconstruction loss to graphs with masked nodes to generate graph embeddings for use in downstream tasks.
FLAG (Kong et al., 2020) adds adversarial noise during training to input node features as a form of data augmentation for GNNs that demonstrates improved performance for many tasks.
It does not add an additional auxiliary loss, which we ﬁnd is essential for addressing oversmoothing.
In other related GNN work, (Sato et al., 2021) use random input features to improve generalisation of graph neaural networks.
Finally, we take inspiration from (Vincent et al., 2008; 2010; Vincent, 2011; Song & Ermon, 2019) which use the observation that noised data lies off the data manifold for representation learning and generative modelling.
Machine Learning for 3D Molecular Property Prediction.
One application of GNNs is to speed up quantum chemistry calculations which operate on 3D positions of a molecule (Duvenaud et al., 2015; Gilmer et al., 2017; Schütt et al., 2017;
Hu et al., 2021b).
In these graphs atoms are nodes and edges are constructed between close 3D neighbours.
Common goals are the prediction of molecular properties (Ramakrishnan et al., 2014), forces (Chmiela et al., 2017), energies (Chanussot* et al., 2020) and charges (Unke & Meuwly, 2019).
These datasets have spurred the development of GNNs that embed 3D physical symmetry inductive biases such as the rotation equivariance of forces.
Such inductive biases typically improve performance and sample complexity.
A common approach to embed physical symmetries is to design a network that predicts a rotation and translation invariant energy (Schütt et al., 2017; Klicpera et al., 2020a;
Liu et al., 2021).
The input features of such models include distances (Schütt et al., 2017), angles (Klicpera et al., 2020b;a) or torsions and higher order terms (Liu et al., 2021).
An alternative approach to embedding symmetries is to design a rotation equivariant neural network that use equivariant representations (Thomas et al., 2018;
Köhler et al., 2019; Kondor et al., 2018; Fuchs et al., 2020; Batzner et al., 2021; Anderson et al., 2019;
Satorras et al., 2021).
Machine Learning for Bond and Atom Molecular Graphs .
Predicting properties from molecular graphs without 3D points, such as graphs of bonds and atoms, is studied separately and often used to benchmark generic graph property prediction models such as GCNs (Hu et al., 2020a) or GATs (Veliˇckovi´c et al., 2018).
Models developed for 3D molecular property prediction cannot be applied to bond and atom graphs.
Common datasets that contain such data are OGBG-MOLPCBA and OGBG-MOLHIV.
P RELIMINARIES : G RAPH P REDICTION P ROBLEM
Let G =
( V, E, g ) be an input graph.
The nodes are V = { v 1 , . .
. , v | V | } , where v i ∈ R d v .
The directed, attributed edges are E =
e 1 , . .
, e | E | } : each edge includes a sender node index, receiver node index, and edge attribute, e k
= ( s k , r k , e k ) , respectively, where s k , r k ∈ { 1 , . .
, | V |} and e k ∈ R d e .
The graph-level property is g ∈ R d g .
The goal is to predict a target graph, G (cid:48) , with the same structure as G , but different node, edge, and/or graph-level attributes.
We denote ˆ G (cid:48) as a model’s prediction of G (cid:48) .
Some error metric deﬁnes quality of ˆ G (cid:48) with respect to the target G (cid:48) , Error ( ˆ G (cid:48) , G (cid:48) ) , which the training loss terms are deﬁned to optimize.
In this paper the phrase “message passing steps” is synonymous with “GNN layers”.
N OISY N ODES
One way to incentivise diverse node representations is to have diverse node-level targets; in order to do well on such a task, the GNN must have diverse latents at the penultimate layer of the network.
However, many problems are not node level prediction tasks, or do not contain sufﬁciently diverse node targets.
Noisy Nodes tackles this problem by adding a diverse noise correction target, modifying the original graph prediction problem deﬁnition in several ways.
It introduces a graph corrupted by noise, ˜
G = ( ˜ V , ˜ E, ˜ g ) , where ˜ v i ∈ ˜
V is constructed by adding noise, σ i , to the input nodes, ˜ v i
= v i
+ σ i .
The edges, ˜ E , and graph-level attribute, ˜ g , can either be uncorrupted by noise (i.e., ˜ E
= E , ˜ g
= g ), calculated from the noisy nodes (for example in a nearest neighbors graph), or corrupted independent of the nodes—these are minor choices that can be informed by the speciﬁc problem setting.
Our method requires a noise correction target to prevent oversmoothing by enforcing diversity in the last layers of the GNN, which for most problems can be achieved with an additional auxiliary denoising autoencoder loss.
For example, where the Error is deﬁned with respect to graph-level predictions (e.g., predict the minimum energy value of some molecular system), a second output head can be added to the GNN architecture which requires denoising the inputs as targets. This output head is weighted by a constant weighting coefﬁcient λ ∈ R .
Because the v (cid:48) i ∈ V (cid:48) is not speciﬁed by the goal, we can set v (cid:48)
i = v i and train the model to predict the uncorrupted input nodes.
The added noise prevents overﬁtting, and the denoising loss ensures a diverse node level target.
Such a recipe also encourages the model to take advantage of message passing.
For example, consider a group of three atoms corrupted by noise, leading to two atoms being very close together (and so with a very high interatomic force).
To denoise the interatomic distances correctly the GNN must use message passing to triangulate between the three atoms.
In Figure 2 we illustrate the impact of Noisy Nodes on oversmoothing by plotting the Mean Absolute
Distance (MAD) (Chen et al., 2020) of the residual updates of each layer of a GNN (a GNS) trained on the QM9 (Ramakrishnan et al., 2014) dataset.
MAD is a measure of the diversity of graph node features, often used to quantify oversmoothing, the higher the number the more diverse the node features, the lower the number the less diverse.
In this plot we can see that for Noisy Nodes the node updates remain diverse for all of the layers, whereas without Noisy Nodes diversity is lost afterlayers. Further analysis of MAD across seeds andd with unsorted layers can be seen in Appendix Figures 7 and 6.
The Graph Manifold Learning Perspective.
By using an implicit mapping from corrupted data to clean data, the Noisy Nodes objective encourages the model to learn the manifold on which the clean data lies— we speculate that the GNN learns to go from low probability graphs to high probability graphs.
In the autoencoder case the GNN learns the manifold of the input data.
When node targets are provided, the GNN learns the manifold of the target data (e.g. the manifold of atoms at equilibrium).
We speculate that such a manifold may include commonly repeated substructures that are useful for downstream prediction tasks.
A similar motivation can be found in (Vincent et al., 2010; Song &
Ermon, 2019)
The Energy Perspective for Molecular Property Prediction.
Local, random distortions of the geometry of a molecule at a local energy minimum are almost certainly higher energy conﬁgurations.
As such, a task that maps from a noised molecule to a local energy minimum is learning a mapping from high energy to low energy.
Data such as QM9 contains molecules at local minima.
Some problems have input data that is already high energy, and targets that are at equilibrium.
For these datasets we can generate new high energy states by adding noise to the inputs but keeping the equilibrium target the same, Figure 6 demonstrates this approach.
To preserve translation invariance we use displacements between input and target ∆ , the corrected target after noise is ∆ − σ .
3D M OLECULAR P ROPERTY P REDICTION E XPERIMENTS AND R ESULTS
In this section we evaluate how a popular, simple model, the GNS (Sanchez-Gonzalez* et al., 2020) performs on 3D molecular prediction tasks when combined with Noisy Nodes.
The GNS was originally developed for particle ﬂuid simulations, but has recently been adapted for molecular property prediction (Hu et al., 2021b).
We ﬁnd that Without Noisy Nodes the GNS architecture is not competitive, but by using Noisy Nodes we see improved performance comparable to the use of specialised architectures.
We made minor changes to the GNS architecture. We featurize the distance input features using radial basis functions. We group layer weights, similar to tied weights used in Li et al.
(2021) for reduced parameter counts; for a group size of n the ﬁrst n layer weights are repeated, i.e. the ﬁrst layer with a group size of 10 has the same weights as the 11 th , 21 st , 31 st layers and so on.
n contiguous blocks of layers are considered a single group.
Finally we ﬁnd that decoding the intermediate latents and adding a loss after each group aids training stability. The decoder is shared across groups.
We tested this architecture on three challenging molecular property prediction benchmarks:
OC20 (Chanussot* et al., 2020) IS2RS & IS2RE, and QM9 (Ramakrishnan et al., 2014).
These benchmarks are detailed below, but as general distinctions, OC20 tasks use graphs 2-20x larger than
QM9.
While QM9 always requires graph-level prediction, one of OC20’s two tasks (IS2RS) requires node-level predictions while the other (IS2RE) requires graph-level predictions.
All training details may be found in the Appendix.
performance and realises no beneﬁt from depth. B) After adding a position node loss, performance improves as depth increases.
C)
As we add Noisy Nodes and parameters the model achieves SOTA, even with 3 layers, and stops overﬁtting.
D)
Adding Noisy Nodes allows a model with even fully shared weights to achieve SOTA.
5. O PEN C ATALYST
Dataset.
The OC20 dataset (Chanussot* et al., 2020) (CC Attribution 4.0) describes the interaction of a small molecule (the adsorbate) and a large slab (the catalyst), with total systems consisting of 20-200 atoms simulated until equilibrium is reached.
We focus on two tasks; the Initial Structure to Resulting Energy (IS2RE) task which takes the initial structure of the simulation and predicts the ﬁnal energy, and the Initial Structure to Resulting Structure (IS2RS) which takes the initial structure and predicts the relaxed structure.
Note that we train the more common “direct” prediction task that map directly from initial positions to target in a single forward pass, and compare against other models trained for direct prediction.
We do not train or compare against “relaxation” models, which have reported fewer results and use far more data.
However, it should be noted that recent promising models such as GemNet (Klicpera et al., 2021) have only been evaluated on relaxation tasks, but we would expect them to perform well on direct prediction. Models are evaluated on 4 held out test sets.
Four canonical validation datasets are also provided.
Test sets are evaluated on a remote server hosted by the dataset authors with a very limited number of submissions per team.
Noisy Nodes in this case consists of a random jump between the initial position and relaxed position.
During training we ﬁrst sample uniformly from a point in the relaxation trajectory or interpolate uniformly between the initial and ﬁnal positions. , and then add I.I.D Gaussian noise with mean zero and σ = 0 .
The Noisy Node target is the relaxed structure.
We ﬁrst convert to fractional coordinates (i.e. use the periodic unit cell as the basis) which render the predictions of our model invariant to rotations, and append the following rotation and translation invariant vector ( αβ T , βγ T , αγ T ,
| α | , | β | , | γ | ) ∈ R 6 to the edge features where α, β, γ are vectors of the unit cell.
This additional vector provides rotation invariant angular and extent information to the GNN.
[Table]
[Table] [Table] [Table]
[Table]
IS2RE Results.
In Figure 3 we show how using Noisy Nodes allows the GNS to achieve state of the art performance.
Figure 3 A shows that without any auxiliary node target, an IS2RE GNS achieves poor performance even with increased depth.
The fact that increased depth does not result in improvement supports the hypothesis that GNS suffers from oversmoothing.
As we add a node level position target in B) we see better performance, and improvement as depth increases, validating our hypothesis that node level targets are key to addressing oversmoothing.
In C) we add noisy nodes and parameters, and see that the increased diversity of the node level predictions leads to very signiﬁcant improvements and SOTA, even for a shallow 3 layer network.
D) demonstrates this effect is not just due to increased parameters - SOTA can still be achieve with shared layer weights .
In Table 1 we conduct an ablation on our hyperparameters, and again demonstrate the improved performance of using Noisy Nodes.
Results were averaged over 3 seeds and standard errors on the best obtained checkpoint show little sensitivity to initialisation.
We conducted an ablation on ID comparing sampling from a relaxation trajectory and interpolating between initial & ﬁnal positions, which found that interpolation improved our score from 0.48 to 0.45.
Our best hyperparameter setting was 100 layers which achieved a 43.3% relative performance improvement against SOTA results (Table 2).
Due to limited permitted test submissions, results presented here were from one test upload of our best performing validation seed.
IS2RS Results.
In Table 3 we perform the same ablation studies on the IS2RS validation set using the Average Distance within Threshold (ADwT) metric and observe the same relative effects.
In
Table 4 we see that GNS + Noisy Nodes is signiﬁcantly better than the only other reported IS2RS direct result, ForceNet, itself a GNS variant.
Similarly, to IS2RE, we tried an ablation using random interpolated positions between input and target instead of sampling from mid trajectory, and noted an improvement from 54.3% to 54.5% on OOD both validation.
We note a drop between the validation and test sets, which we speculate is due to distribution shift.
We also notice a consistent pattern that having energy predictions as an auxiliary loss inhibits both relaxation and direct IS2RE predictions, we speculate that adding more capacity to the positional decoder head may help ameliorate such effects.
[Table] [Table]
Dataset.
The QM9 benchmark (Ramakrishnan et al., 2014) contains 134k molecules in equilibrium with up to 9 heavy C, O, N and F atoms, targeting 12 associated chemical properties (License: CCBY 4.0).
We use 114k molecules for training, 10k for validation and 10k for test.
All results are on the test set.
We subtract a ﬁxed per atom energy from the target values computed from linear regression to reduce variance.
We perform training in eV units for energetic targets
[Table]
and evaluate using MAE.
We summarise the results across the targets using mean standardised MAE (std. MAE) in which
MAEs are normalised by their standard deviation, and mean standardised logMAE.
Std.
MAE is dominated by targets with high relative error such as ∆ (cid:15) , whereas logMAE is sensitive to outliers such as (cid:10) R 2 (cid:11) .
As is standard for this dataset, a model is trained separately for each target.
For this dataset we add I.I.D Gaussian noise with mean zero and σ = 0 .
02 to the input atom positions.
A denoising autoencoder loss is used.
Results In Table 6 we can see that adding Noisy Nodes signiﬁcantly improves results by 23.1% relative for GNS, making it competitive with specialised architectures.
To understand the effect of adding a denoising loss, we tried just adding noise and found no where near the same improvement (Table 6).
A GNS-10 + Noisy Nodes with 30 layers achieves top results on 3 of the 12 targets and comparable performance on the remainder (Table 6).
On the std.
MAE aggregate metric GNS + Noisy Nodes performs better than all other reported results, showing that Noisy Nodes can make even a generic model competitive with models hand-crafted for molecular property prediction.
The same trend is repeated for an rotation invariant version of this network that uses the principle axes of inertia ordered by eigenvalue as the co-ordinate frame (Table 5).
(cid:10) R 2 (cid:11) , the electronic spatial extent, is an outlier for GNS + Noisy Nodes.
Interestingly, we found that without noise GNS-10 + Noisy Nodes achieves 0.33 for this target.
We speculate that this target is particularly sensitive to noise, and the best noise value for this target would be signiﬁcantly lower than for the dataset as a whole.
[Table]
GNS
[Table]
Noise But No Node Target
[Table]
Nodes
[Table]
[Table]
[Table]
Graphormer (Ying et al., 2021)
[Table]
N ON -S PATIAL T ASKS
The previous experiments use the 3D geometries of atoms, and models that operate on 3D points.
However, the recipe adding a denoising auxiliary loss can be applied to other graphs with different types of features.
In this section we apply Noisy Nodes to additional datasets with no 3D points, using different GNNs, and show analagous effects to the 3D case.
Details of the hyperparameters, models and training details can be found in the appendix.
6. OGBG-PCQM4M
This dataset from the OGB benchmarks consists of molecular graphs which consist of bonds and atom types, and no 3D or 2D coordinates.
To adapt Noisy Nodes to this setting, we randomly ﬂip node and edge features at a rate of 5% and add a reconstruction loss.
We evaluate Noisy Nodes using an MPNN + Virtual Node (Gilmer et al., 2017).
The test set is not currently available for this dataset.
In Table 7 we see that for this task Noisy Nodes enables a 50 layer MPNN to reach state of the art results.
Before adding Noisy Nodes, adding capacity beyond 16 layers did not improve results.
6. OGBG-MOLPCBA
The OGBG-MOLPCBA dataset contains molecular graphs with no 3D points, with the goal of classifying 128 biological activities.
On the OGBG-MOLPCBA dataset we again use an MPNN +
Virtual Node and random ﬂipping noise.
In Figure 4 we see that adding Noisy Nodes improves the performance of the base model, accentuated for deeper networks.
Our 16 layer MPNN improved from 27.6% ± 0.004 to 28.1% ± 0.002 Mean Average Precision (“Mean AP”).
Figure 5 demonstrates how Noisy Nodes improves performance during training.
Of the reported results, our MPNN is most similar to GCN 1 + Virtual Node and GIN + Virtual Node (Xu et al., 2018) which report results of 24.2% ± 0.003 and 27.03% ± 0.003 respectively. Further ablations on noise values, DropNode (Do ﬂipping of input categories improves the performance of MPNNs, and the effect is accentuated with depth. leads to a consistent improvement.
et al., 2021) and DropEdge (Rong et al., 2019) can be found in the appendix, which demonstrate that noisy nodes can be combined with different methods addressing oversmoothing.
6. OGBN-ARXIV
The above results use models with explicit edge updates, and are reported for graph prediction.
To test the effectiveness with Noisy Nodes with GCNs, arguably the simplest and most popular GNN, we use OGBN-ARXIV, a citation network with the goal of predicting the arxiv category of each paper.
Adding Noisy Nodes, with noise as input dropout of 0.1, to 4 layer GCN with residual connections improves from 72.39% ± 0.002 accuracy to 72.52% ± 0.003 accuracy.
A baseline 4 layer GCN on this dataset reports 71.71% ± 0.002.
The SOTA for this dataset is 74.31% (Sun & Wu, 2020).
6. L IMITATIONS
We have not demonstrated the effectiveness of Noisy Nodes in small data regimes, which may be important for learning from experimental data.
The representation learning perspective requires access to a local minimum conﬁguration, which is not the case for all quantum modeling datasets.
We have also not demonstrated the combination of Noisy Nodes with more sophisticated 3D molecular property prediction models such as DimeNet++(Klicpera et al., 2020a), such models may require an alternative reconstruction loss to position change, such as pairwise interatomic distances.
We leave this to future work.
Noisy Nodes requires careful selection of the form of noise, and a balance between the auxiliary and primary losses.
This can require hyper parameter tuning, and models can be sensitive to the choice of these parameters.
Noisy Nodes has a particular effect for deep GNNs, but depth is not always an advantage.
There are situations, for example molecular dynamics, which place a premium on very fast inference time.
However even at 3 layers (a comparable depth to alternative architectures)
the
GNS architecture achieves state of the art validation OC20 IS2RE predictions (Figure 3).
Finally, returns diminish as depth increases indicating depth is not the only answer (Table 1).
C ONCLUSIONS
In this work we present Noisy Nodes, a novel regularisation technique for GNNs with particular focus on 3D molecular property prediction.
Noisy nodes helps address common challenges around oversmoothed node representations, shows beneﬁts for GNNs of all depths, but in particular improves performance for deeper GNNs.
We demonstrate results on challenging 3D molecular property prediction tasks, and some generic GNN benchmark datasets.
We believe these results demonstrate
Noisy Nodes could be a useful building block for GNNs for molecular property prediction and beyond.
R EPRODUCIBILITY STATEMENT
Code for reproducing OGB-PCQM4M results using Noisy Nodes is available on github, and was prepared as part of a leaderboard submission.
We withhold the link here to preserve the double blind review process but will provide the link for camera ready.
We provide detailed hyper parameter settings for all our experiments in the appendix, in addition to formulae for computing the encoder and decoder stages of the GNS.
E THICS STATEMENT
Who may beneﬁt from this work?
Molecular property prediction with GNNs is a fast-growing area with applications across domains such as drug design, catalyst discovery, synthetic biology, and chemical engineering.
Noisy Nodes could aid models applied to these domains.
We also demonstrate on OC20 that our direct state prediction approach is nearly as accurate as learned relaxed approaches at a small fraction of the computational cost, which may support material design which requires many predictions.
Finally, Noisy Nodes could be adapted and applied to many areas in which GNNs are used—for example, knowledge base completion, physical simulation or trafﬁc prediction.
Potential negative impact and reﬂection.
Noisy Nodes sees improved performance from depth, but the training of very deep GNNs could contribute to global warming.
Care should be taken when utilising depth, and we note that Noisy Nodes settings can be calibrated at shallow depth.
A A PPENDIX
The following sections include details on training setup, hyper-parameters, input processing, as well as additional experimental results.
A. A DDITIONAL M ETRICS FOR O PEN C ATALYST IS2RS T EST S ET
Relaxation approaches to IS2RS minimise forces with respect to positions, with the expectation that forces at the minimum are close to zero.
One metric of such a model’s success is to evaluate the forces at the converged structure using ground truth Density Functional Theory calculations and see how close they are to zero.
Two metrics are provided by OC20 (Chanussot* et al., 2020) on the
IS2RS test set: Force below Threshold (FbT), which is the percentage of structures that have forces below 0.05 eV/Angstrom, and Average Force below Threshold (AFbT) which is FbT calculated at multiple thresholds.
The OC20 project computes test DFT calculations on the evaluation server and presents a summary result for all IS2RS position predictions.
Such calculations take 10-12 hours and they are not available for the validation set.
Thus, we are not able to analyse the results in Tables 8 and 9 in any further detail.
Before application to catalyst screening further work may be needed for direct approaches to ensure forces do not explode from atoms being too close together.
A. M ORE DETAILS ON GNS ADAPTATIONS FOR MOLECULAR PROPERTY PREDICTION . Encoder.
The node features are a learned embedding lookup of the atom type, and in the case of OC20 two additional binary features representing whether the atom is part of the adsorbate or catalyst and whether the atom remains ﬁxed during the quantum chemistry simulation.
The edge features, e k are the distances | d | featurised using c Radial Bessel basis functions, ˜ e RBF,c = (cid:113) [Equation] , and the edge vector displacements, d , normalised by the edge distance:
[Equation]
Decoder
The decoder consists of two parts, a graph-level decoder which predicts a single output for the input graph, and a node-level decoder which predicts individual outputs for each node.
The graph-level decoder implements the following equation:
[Equation]
Where a Proc i are node latents from the Processor, a Enc i are node latents from the Encoder, W Enc and
W Proc are linear layers,
b Enc and b
Proc are biases, and | V | is the number of nodes.
The node-level decoder is simply an MLP applied to each a Proc i which predicts a ∆ i .
A. M ORE DETAILS ON MPNN FOR OGBG-PCQM4M AND OGBG-MOLPCBA
Our MPNN follows the blueprint of Gilmer et al.
We use (cid:126)h ( t ) v to denote the latent vector of node v at message passing step t , and (cid:126)m ( t ) uv to be the computed message vector for the edge between nodes u and v at message passing step t .
We deﬁne the update functions as: [Equation] [Equation] Where the message function ψ t +1 and the update function φ t +1 are MLPs.
We use a “Virtual Node” which is connected to all other nodes to enable long range communication.
Out readout function is an MLP.
No spatial features are used.
A. E XPERIMENT SETUP FOR 3D M OLECULAR M ODELING
Open Catalyst.
All training experiments were ran on a cluster of TPU devices.
For the Open Catalyst experiments, each individual run (i.e. a single random seed) utilised 8 TPU devices on 2 hosts (4 per host) for training, and 4 V100 GPU devices for evaluation (1 per dataset).
Each Open Catalyst experiment was ran until convergence for up to 200 hours.
Our best result, the large 100 layer model requires 7 days of training using the above setting.
Each conﬁguration was run at least 3 times in this hardware conﬁguration, including all ablation settings.
We further note that making effective use of our regulariser requires sweeping noise values.
These sweeps are dataset dependent and can be carried out using few message passing steps.
QM9.
Experiments were also run on TPU devices.
Each seed was run using 8 TPU devices on a single host for training, and 2 V100 GPU devices for evaluation.
QM9 targets were trained between 12-24 hours per experiment.
Following Klicpera et al.
(2020b)
we deﬁne std.
MAE as :
[Equation]
and logMAE as:
[Equation]
[Table]
A. O VER S MOOTHING A NALYSIS FOR GNS
In addition to Figure we repeat the analysis with a mean MAD over 3 seeds 7.
Furthermore we remove the sorting layer by MAD value and ﬁnd the trend holds.
A. N OISE A BLATIONS FOR OGBG-MOLPCBA
We conduct a noise ablation on the random ﬂipping noise for OGBG-MOLPCBA with an 8 layer
MPNN + Virtual Node and ﬁnd that our model is not very sensitive to the noise value (Table 10) but degrades from 0.1.
A. D ROP E DGE & D ROP N ODE A BLATIONS FOR OGBG-MOLPCBA
We conduct an ablation with our 16 layer MPNN using DropEdge at a rate of 0.1 as an alternative approach to improving oversmoothing and ﬁnd it does not improve performance for ogbg-molpcba (Table 11), similarly we ﬁnd DropNode (Table 12) does not improve performance.
In addition, we ﬁnd that these two methods can’t be combined well together, reaching a performance of 27.0% ±
[Table]
[Table]
We also measure the MAD of the node latents for each layer and ﬁnd the indeed Noisy Nodes is more effective at addressing oversmoothing in Figure 8.
A. T RAINING C URVES FOR OC20 N OISY N ODES A BLATIONS D EMONSTRATING O VERFITTING
Figure
A.
P SEUDOCODE FOR 3D M OLECULAR P REDICTION T RAINING S TEP
G =
( V, E, g ) //
Input graph ˜ G
= G // Initialize noisy graph λ //
Noisy Nodes Weight if not_provided( V (cid:48) )
then
V (cid:48) ← V endif predict_differences then ∆ = { v (cid:48) i − v i | i ∈ 1 , . . . , | V |} endforeach i ∈ 1 , . . . , | V | do σ i = sample_node_noise(shape_of( v i ));
˜ v i
=
v i
+ σ i ; if predict_differences then ˜ ∆ i =
∆ i − σ i ; end endfor ˜ E = recompute_edges( ˜ V ); ˆ G (cid:48) = GNN( ˜ G ); if predict_differences then V (cid:48)
=
˜∆ i ; end
Loss = λ
NoisyNodesLoss( ˆ G (cid:48) , V (cid:48) )
+ PrimaryLoss( ˆ G (cid:48) , V (cid:48) ));
Loss
.minimise()
A. T RAINING D ETAILS
Our code base is implemented in JAX using Haiku and Jraph for GNNs, and Optax for training (Bradbury et al., 2018; Babuschkin et al., 2020; Godwin* et al., 2020; Hennigan et al., 2020).
Model selection used early stopping.
All results reported as an average of 10 random seeds.
OGBG-PCQM4M & OGBG-MOLPCBA were trained with 16 TPUs and evaluated with a single V100 GPU.
OGBN-Arxiv was trained and evalated with a single TPU
3D Molecular Prediction
We minimise the mean squared error loss on mean and standard deviation normalised targets and use the Adam (Kingma & Ba, 2015) optimiser with warmup and cosine decay.
For OC20 IS2RE energy prediction we subtract a learned reference energy, computed using an MLP with atom types as input.
For the GNS model the node and edge latents as well as MLP hidden layers were sized 512, withlayers per MLP and using shifted softplus activations throughout.
OC20 & QM9 Models were trained on 8 TPU devices and evaluated on a single V100 GPUs.
We provide the full set of hyper-parameters and computational resources used separately for each dataset in the Appendix.
All noise levels were determined by sweeping a small range of values ( ≈ 10 ) informed by the noised feature covariance.
Non Spatial Tasks A. H YPER - PARAMETERS
Open Catalyst.
We list the hyper-parameters used to train the default Open Catalyst experiment.
If not speciﬁed otherwise (e.g. in ablations of these parameters), experiments were ran with this conﬁguration.
[Table] Dynamic batch sizes refers to constructing batches by specifying maximum node, edge and graph counts (as opposed to only graph counts) to better balance computational load.
Batches are constructed until one of the limits is reached.
Parameter updates were smoothed using an EMA for the current training step with the current decay value computed through decay = min ( decay, (1 .
+ step ) /
(10 . 0 + step ) .
As discussed in the evaluation, best results on Open Catalyst were obtained by utilising a 100 layer network with group size 10.
QM9 Table 14 lists QM9 hyper-parameters which primarily reﬂect the smaller dataset and geometries with fewer long range interactions.
For U 0 , U , H and G we use a slightly larger number of graphs per batch - 16 - and a smaller position loss co-efﬁcient of 0.01.
OGBG-PCQM4M Table 15 provides the hyper parameters for OGBG-PCQM4M.
OGBG-MOLPCBA Table 16 provides the hyper parameters for the OGBG-MOLPCBA experiments.
OGBN-ARXIV Table 17 provides the hyper parameters for the OGBN-Arxiv
[Table] [Table]
[Table] [Table]