GPS++: Reviving the Art of Message Passing for Molecular Property Prediction
Authors: Dominic Masters
Josef Dean
Kerstin Klaeser
Zhiyi Li
Samuel Maddrell-Mander
Adam Sanders
Hatem Helal
Deniz Beker
Andrew W Fitzgibbon
Shenyang Huang
Ladislav Rampášek
Dominique Beaini
Venue: Accepted by TMLR ()

Abstract: We present GPS++, a hybrid Message Passing Neural Network / Graph Transformer model for molecular property prediction. Our model integrates a well-tuned local message passing component and biased global attention with other key ideas from prior literature to achieve state-of-the-art results on large-scale molecular dataset PCQM4Mv2. Through a thorough ablation study we highlight the impact of individual components and find that nearly all of the model’s performance can be maintained without any use of global self-attention, showing that message passing is still a competitive approach for 3D molecular property prediction despite the recent dominance of graph transformers. We also find that our approach is significantly more accurate than prior art when 3D positional information is not available.

Published in Transactions on Machine Learning Research (07/2023)
GPS++: Reviving the Art of Message Passing for Molecular Property Prediction
Dominic Masters Adam Sanders', Hatem Helal', Deniz Beker', Andrew Fitzgibbon', Shenyang Huang?. 3, Ladislav Rampášek 5 Dominique Beain]
Josef De Kerstin Klaser Zhiyi Li', Sam Maddrell-Mander"
*Graphcore
Valence
Mila
Québec AT Institute
McGill University
Université de Montréal
Reviewed OpenReview:
Abstract
We present GPS++, a hybrid Message Passing Neural Network / Graph Transformer model for molecular property prediction. Our model integrates a well-tuned local message passing component and biased global attention with other key idcas from prior literature to achieve state-of-the-arl results on large-scale molecular dalaset. PCQM4Mv2. Through thorough ablation study highlight the impact of individual components and find that nearly all of the model's performance can be maintained without any use of global self-attention, showing that message passing is still a competitive approach for 3D molecular property prediction despite the recent dominance of graph transformers. We also find that our approach is significantly more accurate than prior art when 3D positional information is not available.
Table Comparison of model size and accuracy on large-scale molecular property prediction dataset PCQM4Mv2.
Model
Validation MAE (meV) +
Model
# Params
Type
GIN-virtua] [uu et al.
2021
6.7M
MPNN
108.3
GPS [Rampašek et al. 2022
19.4M
Ilyrrid
85.8
GEM-2 LLu
2022a
32.1M
Transformer
79.3
Global- Vini-
Wang el al.
2022b
78.5M
Transformer
78.4
Transformer-M M [uo et al.
2022
69.0M
Transformer
77.2
GPS++ [MPNN only]
40.0M
MPNN
77.2
GPS++
44.3M
Hybrid
76.6
1 Introduction
Among many scientific areas, deep learning having a transformative impact on molecular property prediction tasks for biological and chemical applications (Keith et al., 2021 Reiser al. 2022). Tn particular, the aim is to replace or augment expensive and/or time-consuming experiments and first-principles methods with more efficient machine lcarning models. While there is long history of machine learning methods in this field, two particular approaches have been dominant as of late for processing graph-structured molecular data: message passing neural networks (MPNNs) iteratively build graph representations of molecules by sending information explicitly along edges defined by bonds (Gilmer et al. 2017 Battaglia et al. 2018); while graph transformers treat the nodes (atoms) as all-to-all connected and employ global self-attention approaches (Vaswani et al.) 2017), optionally building in local inductive biases through augmented attention (Ying et al. [021a] Luo
et al. 2022). While transformers have been extremely successful in other domains, the quadratic complexity
"Equal Contribution

---

Published in Transactions on Machine Learning Research (07 (07/202) 2023)
Atom/Bond
Calegorical Features
iife cariit
Corrupt
GPS++ module
Noisy Nodes/Edges Ricconstruction Loss
Global
Structural
Edge
Meesage Paszing
MLP
Encode pe feture MuP
Node
MLP
Pool
MLP
L-1 Lose
Blased Sett- Attantion
Attontion Bias
Ground Truth
Posibions
HOMO-LUMO Gap
(hai
Mask?
SE/3) Equivariant
Posstional Danaising Loes
Noise Prediction
x16
Noise
Figure 1: Augmented General Powerful Scalable (GPS++) Graph Transformer overview. (left) GPS+ takes on the input a featurised molecular graph. Chemical, positional, and structural node (atom) and edge (bond) features are detailed in section 4.2 Geometric 3D information is provided only optionally during training, e.g. inference on PCQM4Mv2 test set. needs to be done without. their explicit knowledge. (middle) A stack of GPS++ modules that combine custom local MPNN and global biased self-attention mechanism (akin to Transformer-M) to learn expressive representations, see section 4.1 (right) Global graph pooling with final prediction head, and auxiliary denoising tasks, see section 5.2
with respect to the number of nodes is a significant obstacle for larger molocules. This makes MPNNs an attractive approach due to their linear scaling with graph size, however, issues like oversmoothing (Li et al., 2018), oversquashing, and underreaching (Alon & Vahav 2021 have been found to limit their effectiveness.
In this work we focus on the task of predicting the HOMO-LUMO energy gap, important quantum chemistry property being the minimum energy needed to excite an electron in the molecular structure. This property is typically calculated using Density Functional Theory (DFT) (Khnn & Sham 1965), the de facto method used for accurately predicting quantum phenomena across a range of molecular systems. Unfortunately, traditional DFT can be extremely computationally expensive, prohibiting the efficient exploration of chemical space (Dobson 2004,, with some approaches taking more than 8 hours per molecule per CPU [Axelrod & Gomez-Bombarelij 2022). Within this context the motivation for replacing it with fast and accurate machine learning models is clear. While this task does aim to accelerate the development of alternatives to DFT it also serves as proxy for other molecular property prediction tasks. Therefore, it can potentially benefit a range of scientific applications in fields like computational chemistry, material sciences and drug discovery.
The PubChemQC project (aakata Shimazaki 2017 is one of the largest widely available DFT databases, and from it derived the PCQM4Mv2 dataset, released as part of the Open Graph Benchmark Large Scale Challenge (OGB-LSC) (Hu et al. 2021), which has served as a popular best for development and benchmarking of novel graph neural networks (GNNs). The original OGB-LSC 2021 competition motivated wide range of solutions using both MPNNs and transformers. However, following the success of the winning method Graphormer (Ying et al. 2021a), subsequent work has resulted in a large mumber of graph transformer methods for this task with comparatively little attention given to the message passing methods that have been successful in other graph-structure tasks.
In this work we build on the work of Rampášok et al. (2022) that advocates for a hybrid approach, including both message passing and transformer components in their General, Powerful, Scalable (GPS) framework. Specifically, we build GPS++ which combines a large and expressive gressage-passingg module with biased self-attention layer to maximise the benefit of local inductive biases while still allowing for effective global communication. Furthermore, by integrating a grouped input masking method (Luo et al. 2022) to exploit available 3D positional information and carefully crafting a range of diverse input features we achieve the best-reported result on the PCQM4Mv2 validation data split of 76.6 meV mean absolute error (MAE).
Next, we perform an extensive ablation study to understand the impact of different model components on the performance. Surprisingly, we find that even without global self-attention mechanism (seen in graph
2

---

Published in Transactions on Machine Learning Research (07/2023)
transformer architectures), extremely competitive performance can be achieved. Therefore, argue that MPNNs remain viable in the molecular property prediction task and hope the results presented in this work can spark a renewed interest in this area. We also observe that when solely focusing on 2D molecular data (without 3D conformer coordinates), our proposed GPS++ model significantly outperforms other such works.
Our contributions can be summarised follows:
We show that our hybrid MPNN MONNNTTransee /Transformer model, GPS++ is a parameter-effcient and effective approach to molecular property prediction, achieving state-oftthe-ar MAE scores for PCQM4Mv2 even when compared to parametrically larger models.
We find that even without self-attention, our model matches the performance of prior state-of theffftth-artt transformers, highlighting that well-optimised MPNNs are still highly effective in this domain.
We investigate how different model components affect the model performance, in particular highlight- ing the impact of the improved chemical fealure choice, 3D positional features, structural encodings and architectural components.
Reproducibility: Source code to reproduce our results can be found at: https///ggthhu..omm//rpphoorrr ogb- lsc- pcqm4mv2
2 Related Work
Before the advent of deep learning for graph representation learning (Bronstein et al. 2021 molecular representation in chemoinformatics rested on feature engineering of descriptors or fingerprints Rogers &
Hahn 2010 Keith et al. 2021). After learnable fingerprints (Duvenaud et al. 2015 the general framework of GNNs, often based on MPNNs, has been rapidly gaining adoption (Reiser et al. 2022). Not all approaches
follow the MPNN model of restricting compute to the sparse connectivity graph of the molecule, for example
the continuous-fltee convolutions SchNet (Schütt et al. 2018) and the transformer-bast Grover (Rong
et al.
2020
employ dense all-to-all compute patterns.
GNN method development has also boen facilitated by the availability of well motivated (suites of) benchmark- ing datasels such as QM9 (Ramakrishnan et al. 2014), MoleculeNet (Wu et al. 2018), Open Graph Benchmark
(OGB) (Hu et al. 2020), OGB-LSC PCQM4Mv2 (Hu et al.) 2021). Therapeutics Data Commons (Huang ct al. 2022) for molccular property prediction or MD17 (Chmicla et al., 2017p, ANI-1 (Smith et al. 2017). [Tran el al. 2022) for structural conformers or molecular force field predictions.
Open Calalyst.
Depending on the nature of the task, the geometric 3D structure of input molecules needs to be taken into account. This imposes an important constraint on the models, that need to be roto-translation invariant in case of global molecular property prediction and roto-translation equivariant in case of conformer force field prediction, in addition to input permutation invariance and equivariance, respectively. For the latter case, specialised geometric GNN models have been proposed, such as SchNet Schütt et al. (2018), DimeNet(++ Gasteiger et al. (2020b a), PaiNN SShütt et al. (2021), NequlP Batzner et al. (2022) or a transformer- TorchMD-Net Thölke De Fabritiis (202bb).
In this work we are primarily motivated by the PCQM4Mv2 dataset that is a part of the large-scale graph ML challenge (Hu et al.) 2021), which contains uniquely large numbers of graphs, and has the particular characteristic that the molecular 3D structure is only provided for the training portion of the dataset, but not at test time. This has motivated methods specifically equipped to handle such a scenario: Noisy
Nodes denoising autoencoding (Godwin el. al.) 2022) and pretraining (Zaidi et al.) 2022). GEM-2 (Liu et al.) 2022a), ViSNet (Wang et al. 2022b). Methods based on global self-attention (Vaswani et al., 2017) became particularly dominant after Graphormer (Ying et al. 2021a) won OGB-LSC 2021, which spurred development of transformer- based methods: SAN (Kreuzer et al. 2021), EGT (Hussain al. 2022), GPS (Rampášek
et al. 2022
TokenGT (Kim et al. 2022
Transformer-M [Luo al.] 2022).
3

---

Published in Transactions on Machine Learning Research (07 2022)
3 Preliminaries
Throughout the paper we use the following notation. Bold lowercase letters v are (row) vectors, bold uppercase letters M are matrices, with individual clements denoted by non-bold letters i.. 2k or Mpg- Blackboard bold lowercase letters v are categorical (integer-valued) vectors. In general, we denote by [vx] [e the vertical concatenation (stacking) of vectors Vi- Vertical concatenation is also denoted by semicolon, | i.e. [VL ; ... vJ] = [3 1=1 [v, for J]]. Horizontal concatenation, which typically means concatenation along the feature dimension, is denoted by vertical bar, i.e. [v1 v2].
A molecule is represented as graph G = (V,E) for nodes V and edges E. In this reprosentation, each node E V is an alom in the molecule and each edge (u, v) E & is a chemical bond between two atoms. The number of atoms in the molecule denoted by N = [V] and the number of edges M = -E|.
Each node and edge is associated with a list of categorical features x, Patom and euv Proud, respectively, for sstmm atom features and Dbond bond features. A further set of 3D atom positions R [r] : ... EN E Nx3, extracted from original DFT calculations, is provided for training data, but crucially not for validation and test data.
Our algorithm operates on edge, node, and global features. Node features in layer f are denoted by [xf ...; xW] Edge features en E Redae are concatenated into the edge feature matrix E = =Ceu for (u, v) E E]. Global features are defined per layer as g E dgllobal
x E Roooe,, and are concatenated into the N x dnode matrix X
= X²=
We also define an atlention bias matrix B E RNN,, computed from the input graph topology and 3D atom positions, described later in section 4.2
4
GPS++
Our GPS++ model closely follows the GPS framework set out by Rampášek et al. (2022). This work presents flexible model structure for building hybrid MPNN/Transtonn Transformer models for graph-structured input data. We build a specific implementation of GPS that focuses maximising the benefit of the inductive biases of the graph structure and 3D positional information. We do this by building a large and expressive MPNN component and biasing our attention component with structural and positional information. We also allow global information to be propagated through two mechanisms, namely the global attention and by using a global feature in the MPNN.
As displayed in Figure 1 the main GPS++ block (section 4.1
combines the benefits of both message passing and attention layers by rumning them in parallel before combining them with simple summation and MLP; this layer is repeated 16 times. This main trunk of processing is preceded by an Encoder function responsible for encoding the input information into the latent space (section 4.2 and is followed by simple Decoder function (section 5.2
Feature engineering is also used improve the representation of the atoms/bonds, provide rich positional and structural features that increase expressivity, and to bias the attention weights with a distance embedding.
4.1
GPS++ Block
The GPS++ block is defined as follows for layers é (see section 4.2 for the definitions of X E g., B):
E g +
I+1 GPS++ xx E g. B)
(1)
computed as
Y B
g .+1 MPNN (x E .'¹) Z' BiasedAttn (x', (X*B)
(2)
(3)
Vi:
4+1 FFN (yf +2)
(4)
4

---

Published in Transactions on Machine Learning Research (07/202) /203)
X
Figh1
MLP stobal
MLPPEN
concat
MLP podie
Biasedicta
MPNN
ML. Podge
CORGAE
X
E
g
X
gather, (X) xx, (u,-) E R-
gather, (X) (.. for (1,0) E] MMxia
scatter, ESsnnce Nu for VI Nooi
scatter, (X) Cuoee for V Nxd.
sum(-)
The main GPS++ processing block (left) is composed of local message passing MPNN module and biased global attention BiasedAttn module. (right) A diagram of the used MPNN block. The gather, scatter and sum operations highlight changes in tensor shapes and are defined in equations to
Figure 2:
Our MPNN module is a variation on the neural message passing module with edge and global features Gilmer el al.
2017 Battaglia et al. 2018 Bronstein el al. 2021 We choose this form to maximise the expressivity of the model (Veličković 2023) with the expectation that overfitting will be less of an issue with PCQM4Mv2, compared to other molecular datasets, due to its size. The essential components (excluding dropout and layer norm) of the MPNN module are defined as follows (see Figure for graphical representation, and A Appendix A.1 for the exact formulation):
Y E++1
E+T
MPNN (x+ E g
(5)
computed as
v),
c MLPedge ([xi x Cuo
g'
(6)
Vi
x MLPnode
x
Σ
Tus
Σ
Σ
g
(7)
(M.))E
(EDEE
(U.
receiver mesanges
adjacent
g' MML globbal
g
Σ
Σ
Eue
(8)
JEV
(u.wiee
Wi y: x + x
(9)
Y(u v):
d+1 ctu
Eye
(10)
g c+1 g¹ + +g"
(11)
The three networks MLP,, for (node, edge, global} each have two layers with GELU activation functions and an expanded intermediate hidden dimension of 4d,,
This message passing block is principally the most similar to Battaglia et al. 20018 However, we draw the reader's attention to a few areas that differ from common approaches. Firstly, we choose decouple the three latent sizes, setting duode 256, dedge 128 and dglobal = 64, because we found that increasing dedge and dgtobal does not improve task performance. Secondly (and relatedly) we aggregate the adjacent node
5

---

Published in Transactions on Machine Learning Research (07/2023)
representations well the more complex edge feature messages in Equation 7 this is similar to running simple Graph Convolutional Network (GCN) in parallel to the more complex message calculation. We do this to allow the higber-dimennsonal node features Xy to bypass the lower-dimensional messages e preventing potential compression. Finally, we also aggregate not only the sender messages from edges directed towards the node, but also the receiver messages compuled on the edges directed away from the node (labelled Equation 7 By concatenating these two terms rather than summing them we maintain directionality data in each edge, encouraging message information to flow bidirectionally but not requiring the MLPedge to learn directional invariance.
Our BiasedAttn module follows the form of the self-attention layer of Luo et al. (2022) where standard self-attention block (Vaswani et al., 2017) is biased by a structural prior derived from the input data. In our work, the bias B made up two components, a shortest path distance (SPD) embedding and a 3D distance bias derived from the molecular conformations as described section
The FFN module takes a similar form to MLPe though with additional dropout terms (see Appendix A.1 for full details).
4.2 Input Feature Engineering
As described in section the dataset samples include the graph structure G, set of categorical features for the atoms and bonds Xj, Buv and the 3D node positions r,- It has been shown that there are many benefits to augmenting the imput data with additional structural, positional, and chemical information (Rampášck et al. 2022 Wang et al. 2022a, Dwivedi et. al., 2022). Therefore, we combine several feature sources when computing the input to the first GPS++ layer. There are four feature tensors to initialise; node state, edge
state, whole graph state and attention biases.
X Xatom
LapVec
LLppVal
X RRW
X Cent x
(12)
X Dense
all)
Xddooo
(13)
E Dense bbond
E
R M xdedxe
(14)
g" Embed Aalobal (0)
E Regtthal
(15)
= BSPD + 3D
E NNN
(16)
Here the node features X are built from categorical atom features, graph Laplacian positional encod-
ings [Kreuzer et al. 2021 Dwivedi & Bresson 2020 random walk structural encodings (Dwivedi et al. 2022
local graph centrality encodings (Ying el al.J 2021a Shi el al., 2022) and a 3D centrality encoding (Luo el al.,
2022). The edge features E are derived from categorical bond features and bond lengths, and the attention bias uses SPD and 3D distances. The global features are initialised with a learned constant embedding in latent space. These input features are further described in Appendix A.2
Chemical Features The categorical features x, encapsulate the known chemical properties of the atoms and bonds, for example, the atomic number, the bond type or the number of attached hydrogens (which are not explicitly modelled nodes), as well as graphical properties like node degree or whether an atom or bond is within a ring. The set that is used is not determined by the dataset and augmenting or modifying the chemical input features is a common strategy for improving results.
By default, the PCQM4Mv2 dataset uses a set of atom and 3 bond features (described Table A.1 There is, however, a wide range of chemical fcatures that can be extracted from the periodic table or using tools like RDKit Ying et al. (2021b) have shown that extracting additional atom level properties can be beneficial when trying to predict the HOMO-LUMO energy gap, defining total of 28 atom and 5 bond features. We explore the impact of number of additional node and edge features and sweep wide range of possible combinations.
In particular, we expand on the set defined by Ying et al. (022b)) with three additional atom features derived from the periodic table, the atom group (column), period (row) and clement type (often shown by colour). These are intended generalise to unseen atom types better than single integer atomic number feature, and we found that them to be particularly beneficial (shown later in ablation Table 5)
6

---

Published in Transactions on Machine Learning Research (07/2023)
We also found that in many cases removing features was beneficial, for example, we found that generally our models performed better when excluding information about chiral tag and replacing it with chiral centers. We further observe that our best feature combinations all consist of only 8 node features, where the majority of the imput features stay consistent between the sets. We show the three best feature sets found in Table A.1 and use Set for all experiments unless otherwise stated (i.e., during ablations).
Comparison to GPS The GPS+ model is a specific instantiation of the GPS framework, and therefore has several similarities with the reference GPS model implemented by Rampášek et al. (2022) ) The high level architecture of each layer remains the same, using an MPNN module and self-attention module in parallel before combining their node updates with an MLP. However, the design of these two modules is quite different in our implementation. GPS++ uses a bespoke MPNN instead of the GatedGCN (Bresson & Laurent, 2018) employed by GPS. Our MPNN is most similar to Battaglia et al. (2018) but with some key differences described in section 4.1 such as GCN-like adjacent node aggregation and reverse message passing. For the telf-antention module, where GPS uses unbiased attention like Vaswani et al. (017). GPS++ implements learned bias terms BSPD and B3D to incorporate the graph structure. Considering the input features of each model, GPS++ and GPS share the use of the Graph Laplacian and Random walk lastuunallstrucuuaall encodings X Lap LLapVal and XRW whilst. GPS++ additionally uses node centrality encoding X XCent and the spatial node and edge features X3D and E3D. Moreover, the atomic features Xatom in GPS++ are refinement of the defaults used by GPS, including the addition of the novel atomic Group, Period and Element Type features. Finally, while GPS trains by minimising the error on only the HOMO-LUMO prediction task, GPS++ utilises the following extra regularisation techniques in order to effectively scale parameters on finite data: a node/edge feature denoising task (Godwin et al. 2022 a 3D spatial denoising task (Luo et al. 2022). stochastic depth (Huang el al. 2016) and pervasive element-wise dropout. The additional training lasks are described further in section 5
5
Experimental Setup
5.1
Dataset
The PCQM4Mv2 dataset (uu et al. 2021 consists of 3.7M molecules defined by their SMILES strings. Each molecule has on average 14 atoms and 15 chemical bonds. However, as the bonds are undirected in nature and graph neural networks act on directed edges, two bidirectional edges are used to represent each chemical bond.
The 3.7M molecules are separated into standardised sets, namely into training (90%), validation (2%), test-dev (4%) and test-challenge (4%) sets using scaffold split where the HOMO-LUMO gap targets are only publicly available for the training and validation splits. PCQM4Mv2 also provides conformation for each molecule in the training split, i.e., a position in 3D space for each atom such that each molecular graph is in a relaxed low-energy state. Crucially, the validation and test sets have no such 3D information provided.
5.2 Model Training
Training Configuration Our model training setup uses the Adam optimiser (Kingina & Ba 2015 with a gradient clipping value of 5, peak learning rate of 4e-4 and the model is trained for a total of 450 epochs. We used learning rate warmup period of 10 epochs followed by a lincar decay schedule.
Decoder and Loss The final model prediction is formed by global sum-pooling of all node representations and then passing through 2-layer MLP. The regression loss is the mean absolute error (L1 loss) between scalar prediction and the ground truth HOMO-LUMO gap value.
Noisy Nodes/Edges Noisy nodes Godwin el al. 2022: Zaidi et al. 2022 has previously been shown to be beneficial for molecular GNNs including on the PCQM4M dataset. The method adds noise to the input data then trics to reconstruct the uncorrupted data in an auxiliary task. Its benefits are expected to be two-fold: it adds regularisation by inducing some noise on the input, but also combats oversmoothing by forcing the node-level information to remain discriminative throughout the model. This has been shown to be particularly
7

---

Published in Transactions on Machine Learning Research (07/2023)
beneficial when training deep GNNss Godwin al. 2022). We follow the method of Godwin et al. (2022) that applies noise to the categorical node features by randomly choosing a different category with probability Pooruppt: and we further extend this to the categorical edge features. A simple categorical cross-entropy loss is then used to reconstruct the uncorrupted features at the output. We set Pcorrupt = 0.01 and weight the cross-entropp losses such that they have a ratio 1:1. 1:1.2:1.2 for losses HOMO- Www.-LLMMNNNOOOoooooooo htwp:/oooossoodeesoooseeggeeee The impact of Noisy Eddes can be seen in Table of the ablation study.
Grouped Input Masking The 3D positional features R. are only defined for the training data. We must. therefore make use of them in training without requiring them for validation/tess /test. We found that the method proposed by Luo et al. (2022) achieved the most favourable results we adopt variation hereon referred to as grouped inpul masking.
Atom distances calculated from 3D positions are embedded into vector space R via K = 128 Gaussian kernel functions. We then process these distance embeddings in three ways to produce attention biases (B"))
node features (X3D) and edge features (E$D) (exact formulations can be found Appendix A.2)
This method stochastically masks out any features derived from the 3D positional features R to build robustness to their absence. Specifically, this is done by defining two input sets to be masked:
Spatial (x3D ESD B3D), X Topological (BSPD)
and three potential masking groups: 1. Mask Sppttial, 2. Mask X Topologgiall and 3. No masking. These masking groups are then sampled randomly throughout training with ratio 1:3:1. If 3D positions are not defined, egg. in validation/test, masking group 1 is always used.
3D Denoising Alongside grouped imput masking. Luo et al. 2022) also add noise to the 3D atom positions R during training before computing any derivative features (e.g., 3D Bias, 3D Centrality), and predict the atom-wise noise as an auxiliary self-supervised training task. We closely follow their implementation for GPS++ adding an SF(())-quivariin self-attention layer a noise prediction head, which encourages the model to preserve 3D information until the final layer well further develop spatial reasoning. See Appendix section A.2 for full details.
6 Results
Model Accuracy In Table 2 we compare the single model performance of GPS++ with results from the literature, in particular, pay close attention to the Transformer-M (Luo et al.] 2022) model as it has the best results for a transformer, but also because we use their input masking and denoising method for incorporating 3D features; this allows us to gain some tangible insights into the value of hybrid/MPNN MPNNN approaches vs. transformers.
Comparing the best GPS++ result with prior work, we set new state of the art MAE score on the PCQM4Mv2 validation set, outperforming not only the parametrically comparable Transformer-M Medium but also the much larger Transformer-M Large. While it would be tempting to assume that the Multi-Headed Self Attention (MHSA) component of our model is the most significant due to the prevalence of transformers among the other top results, we find that we can maintain competitive accuracy with no attention at all, falling back to a pure MPNN, passing messages only between bonded atoms (nodes). Furthermore, we show that this MPNN model is more accurate in the absence of 3D features, even beating the hybrid GPS++ in this setting. We believe this shows a strong reliance of the transformer's global attention mechanism on 3D positional information to learn atom relationships and, in contrast, the power of encoding spatial priors into the model using message passing along molecular bonds in the MPNN. Further, by halving the number of GPS++ layers from 16 to 8 to reach parametric parity, Table 2 shows that the MPNN-only GPS++ model compares favourably to the original GPS implementation. We investigate the impact of other components of
the model more in section 7
Model Throughput Molecules in PCQM4Mv2 have an average of 14.3 nodes (atoms) per graph, and our full GPS++ model trains at 17,500 graphs per second on 16 TPUs. This means each epoch completes in
8

---

Published in Transactions on Machine Learning Research (07/2023)
Table 2: Comparison of single model performance on PCQM4Mv2 dataset. Note Transformer-M and
GioNal-ViSNet concurrent work.
Model
# Param.
Model Type
Valid MAE (meV)
without 3D Positional In Information
GCN-virtual (Hu et al.
2021
4.9M
MPNN
115.3
GIN-virtual Hu et
2021
6.7M
MPNN
108.3
GRPE (Park et al. 2022
46.2M
Transformer
89.0
Transformer-Bi Medium
No 3D Positions]
(Luo al. 2022
47.1M
Transformer
87.8
EGT (Hussain al. 2022)
89.3M
Transformer
86.9
Graphormer (Shi et al.
2022
48.3M
Transformer
86.4
GPS [Rampásck ct. al.
2022
19.4M
Hybrid
85.8
GPS++ 18 Layers, No 3D Positins
22.4M
Hybrid
83.6
GPS++ No 3D Positions]
44.2M
Hybrid
82.6
GPS++ Layers, MPNN only, No 3D Positions]
20.3M
MPNN
82.2
GPS++ MPNN only, No 3D Positions
40.0M
MPNN
81.8
with 3D Positional Information
GEM-2 (Liu et al.
2022a
32. 1M
Transformer
79.3
GPS++ [8 Layers, MPNN only
22.5M
MPNN
79.3
Transformer-M [Medium]
(Lno et al. 2022
47.
Transformer
78.7
GPS++ [8 Layers]
22.7M
Hybrid
78.6
Global-ViSNet [Wang et al.
2022b
78.5M
Transformer
78.4
Transformer-M Large] Luo et al 2022
69. 0M
Transformer
77.2
GPS++ [MPNN only]
30.3M
MPNN
77.2
GPS++
44.5M
Hybrid
76.6
minutes and a full 450 epoch training run takes less than 24 hours. The MPNN-only version of GPS++ (i.e. disabling the MHSA module) has almost double the throughput at 33,000 graphs per second, partially due to the reduced compute costs but also as a result of increasing the micro batch size thanks to reduced memory pressure. Figure 3 compares the models for range training epochs, and shows the MPNN-only version can oulperform prior results from GEM-2 and Global-ViSNe with only 5-6 hours of training, and match the prior SOTA Transformer-M M after 13 hours. High throughput allows for rapid iteration and extensive ablation studies (sce section
7 Full details of the hardware used and acceleration methods employed can be found in
Appendix
MPNN Hytrid
79.5
Valdaton MGE (meV)
GBM-2
78.5
GobaliiNNe
Pensformentt
77.0
Taining Time (Hours)
Figure 3: Training time versus validation error for the full hybrid GPS++ model versus the MPNN-only model. Comparable training times for GEM-2, Global-ViSNet and Transformer-M were unavailable at the time of writing.
Fine-tuning Accuracy In order to Lest. GPS++ in the presence of 3D positional data during test time and to investigate the generalisability the model, we fine-tune GPS++ on different tasks from the the quantum chemistry benchmark QM9 (Ruddigkeit et al.] 2012 Ramakrishnan et al., 2014).. The QM9 datasct comprises 134k small organic molecules in equilibrium state which contain up heavy atoms (C, O, N, F) and, unlike in the PCQ dataset, 3D positional information available at test time. Though the benchmark provides 12

---

Published in Transactions on Machine Learning Research (07/2023)
labels for each molecule, we found that 4 very closely related energy labels Uo, U, H and G were particularly sensitive to hyperparameters, random seeds and label normalisation choices, so we have omitted these labels to reduce the time and compute costs of the generalisation study. QM9 does not provide a standardised dataset split, we therefore follow several previous works (Luo et al. 2022 Thölke & De Fabritiis 2022a) and randomly select 10,000 molecules for validation and 10,831 for testing; all remaining molecules are used during training- Due to the small size of this dataset and the high variance in test accuracy observed for different train/tes besttteetttr splits, we fine-tune 3 baseline GPS+ checkpoints on 5 random split seeds for each task and report both the mean and standard deviation across all 15 runs. We disable grouped imput masking during fine-tuning, to maximise the benefit of QM9's 's always-present 3D features. For each label we performed a hyperparameter sweep, settling on learning rates in the range [0.0004, 0.001]. epochs in the range [1000, 2000], 3D denoising loss weighting in the range (0.00001, 0.75], and in some cases removing some dropout operations. In particular, we found that the HOMO, LUMO and Gap labels benefited from a much higher 3D denoising loss weight than protraining (0.75 vs 0.1). By contrast, for the label R² we found the best results by sotting the 3D denoising weight to 0.00001 and disabled dropout in the node MLP as well as stochastic depth, suggesting the signal-to-noise ratio in this label is much lower.
The results are shown in Table 3 Despite differences in the types of molccules, the DFT configurations used to generate the equilibrium states and the availability of 3D positions at test time, GPS++ generalises very well to this new dataset for tasks like HOMO, LUMO and Gap which are closely related to the PCQM4Mw2 pretraining task. The performance of GPS++ on less closely related labels generally falls behind the state of the art, the weakest being ZPVE and R², though in each case the results are well within the distribution of reasonable results from prior work. This bias of the model to particular kind of task appears to be the norm throughout the QM9 results; each model that sets the state of the art in one label also exhibits particularly weak performance in one more other labels. Beyond its current strengths, it may be possible to improve GPS++'s weaker scores by implementing task-specific prediction heads. For example, Thölke & De Fabritiis (022aa) add a rotationally-oquivrrian output layer just when predicting the Electronic Spatial Extent (R"), as this is a highly spatial task.
Table 3: QM9 finc-tuning results. Test MAE and STD reported over 5 test split seeds and 3 pre-trained GPS++ checkpoints (15 runs total). The first and second best results for each label are highlighted.
Homo meV
Lumo meV
Gap meV
Cv cal/mol K
ZPVE meV
R²
Model
Pretrained
D
3
Schnet Schütl
2018
41
34
0.033
0.033
1.70
0.073
0.24
MGCN Lu
20
42
57
0.038
0.056
1.12
0.11
0.030
LieConv 1 Finzi
2020
30
25
49
0.038
0.032
2.28
0.80
0.081
SB(3)- Transtorrrt Fichs ot
2020
35
33
53
0.054
0.051
0.14
DimeNett+ Gassteggr el al.
24.6
19.5
32.6
0.023
0.030
1.21
0.33
0.044
GEM FFng et al 2021
33.8
27.7
52.1
0.035
0.034
1.73
0.089
0.081
N Schitti of all 20
27.6
20.4
45.7
0.024
0.012
1.28
0.068
0.045
LieTF Hutchinson el all 2021
33
27
0.035
0.041
2.10
0.45
0.082
TorchMD-Net (1 Thhkke &r De Fabritiis
2022b
20.3
17.5
36. 1
0.026
0.011
1.84
0.033
0.059
EGNN (atorma
2021
25
0.031
0.029
1.55
0.11
Sphere (Liu et all 20222
22.8
18.9
31.1
0.022
0.026
1.12
0.29
0.016
SEGNN [Bandstetter al.
2022
24
21
0.031
0.023
1.62
0.66
0.060
EQGAT Le ct aL. 2022
20
16
32
0.024
0.011
2.00
0.38
0.053
3D Informx Stark et m.
2022
V
29.8
25.7
48.8
0.033
0.034
1.67
0.12
0.075
3D-MGP (Jaa et aL. 2022
<
21.3
18.2
37.1
0.026
0.020
1.38
0.092
0.057
NoisyNode,
Codwin et all
2022
20.4
18.6
28.6
0.025
0.025
1.16
0.70
0.052
Trans Transforme--X M Lao el al.
2022
17.5
16.2
27.4
0.022
0.037
1.18
0.075
0.041
NN
Zaidi et al 2022
14.9
14.7
22.0
0.020
0.016
1.02
0.44
0.040
GPS++
16.1
14.6
26.2
0.028
0.024
1.85
0.40
0.062
0.001
0.001
7 Ablation Study
The top-performing GPS++ model was attained empirically, resulting in a complex combination of input features, architectural choices and loss functions. In this section we assess the contribution of each feature to final task performance on PCQM4Mv2, and find that much of the performance can be retained by a simplified design. All results listed in the ablation study are obtained by training the final GPS++ model from scratch
10

---

Published in Transactions on Machine Learning Research (07/2023)
for 200 epochs with one more features disabled, averaging MAE over 5 runs. Average batch size is kept constant (926 nodes per batch) for all runs to keep them directly comparable.
Node and Edge Input Features The PCQM4Mv2 dataset provides 9 chemical node and 3 edge features for each sample, and following prior work GPS++ incorporates additional features and preprocessing strategies. Table isolates the contribution of cach of these additions (excluding 3D features).
Table 4: Ablation of node and edge features.
A MAE (meV)
Removed Feature
Valid
Train
None (Baseline)
77.2
53.6
Random Walk Structural Encoding
+ 1.2
+
Shortest Path Distance Bias
0.8
Laplacian Positional Encoding
+ 0.3
+
Local Centrality Encoding
+ 0.1
Bond Lengths
+ 0.1
Noisy 3D Positions
1.6
3.6
Noisy Nodes
0.6
0.2
Noisy Edges
0.1
0.1
In line with previous work [Rampášek et al. 2022), the Random Walk Structural Encoding is the most impactful individual input feature, degrading validation performance by 1.2 mcV upon removal, whilst the Graph Laplacian Positional Encodings (both eigenvectors and eigenvalues) have minimal benefit. The Shortest Path Distance Bias the self-attention module makes a meaningful contribution, though the memory cost. of storing all-to-all node distances is non-trivial. Other features, despite being beneficial when they were added at intermediate steps in the development of the model, are shown to have become redundant in the final GPS++ composition: the Graphormer-style Local Centrality Encoding (i.e. embedding the degree of the node) can be derived from the Random Walk features, and the use of Bond Lengths as an edge feature in the MPNN appears to be unnecessary in the presence of more comprehensive 3D fealures elsewhere in the network. Table 4 also shows that whilst the regularisation effect of the noisy nodes loss is beneficial, the contribution of noisy edges negligible. This may be due the simplicity (and hence case of denoising) of the limited edge features in PCQM4M+2.
Table 5: Choices of chemical features.
Atom Group,
MAE (meV)
Feature Set
Period & Type
Set Size
Valid
Train
Set 1 (Baseline)
77.2
53.6
Set 2
14
+ 0.2
0.0
Set. 3
+ 0.2
0.3
Original
Original
+ 0.5
0.7
Set 1
11
0.6
0.4
Ying21
36
+ 4.1
3.1
Ying21
33
+ 4.4
3.0
Table compares the possible choices of chemical input features described section 4.2 Original the set of features included in PCQM4Mv2; Ying21 refers to all 28 node features and edge features in the PCQ superset defined by Ying et al. (2021b;; Set 1-3 are performant subsets of 11 node features and 3 edge features selected from all of the above, defined fully in Table A. and obtained via combinatorial scarch; Atom Group, Period & Type refers to our contribution of 3 additional atom features relating to an alom's position in the periodic table, intended to be more generalised than the alomic number. The results clearly show that training on all 36 available input features causes significant overfitting and degrades task performance by introducing noise. Morcover, the new Atom Group, Period Type features consistently provide meaningfully richer atom representations than the atomic number feature which present in every feature set.. We believe future work similar datasets should consider incorporating this simple addition.
11

---

Published in Transactions on Machine Learning Research (07/2023)
3D Input Features and Self-Attentin The 3D node position information given for the training set in PCQM4Mv2 is provided input to the final model in three forms: the Bond Length edge features in the MPNN, the all-to-all 3D Bias map added to the self-attention module, and the node-wise 3D Centrality Encoding added in the node encoder. The Bond Lengths are shown to have minimal impact in Table ! so Table G explores the 3D Bias and Centrality as well as how they interact with the MHSA module. In rows 1-4 we see that including at least one of the pair critical to the final model's performance, but neither makes the other redundant. The 3D Bias is higher fidelity than the 3D Centrality (i.e, an all-to-all distance map versus node-wise distance sum) and is the stronger of the pair, but the effectiveness of the simpler 3D Centrality should be noted since it does not require the use of custom biased self-attention layers. Additionally, we found that training stability of the MHSA module suffered in the absence of either 3D feature, requiring us to halve the learning rale to 2e-4 for convergence unless the MHSA module was disabled.
Table 6: Ablation of self-attention and 3D features.
Feature
A MAE (meV)
#
MHSA
3D Bias
3D Centrality
Valid
Train
1
V
77.2
53.6
2
V
+ 0.6
2.3
3
V
+ 1.1
3.6
4
+ 5.1
10.0
5
+ 4.3
16.4
6
+ 0.9
1.5
trained with half learning rate for stability
Interestingly, when the 3D Bias term is not used, Table G rows 4-5 show that it is preferable to remove the MHSA. This may imply that the Shortest Path Distance Bias term (the other bias in the MHSA, which attempts to provide a non-3D global graph positional encoding) is insufficient to modulate attention between atoms, and that instead spatial relationships are key for this quantum chemistry task. For example, when a molecule is folded over and two aloms may be very close in 3D space but far apart in the graph topology.
Model Architecture GPS++ uses a large MPNN module comprising 70% of the model paramcters and using extensively tuned hyper-parameters. Table ablates the components of the model network archilecture, and shows that the MPNN module is unsurprisingly the single most important component for task performance. Outside of the MPNN, the MHSA (with accompanying attention biases) and FFN contribute approximately equally to the final performance, though they have opposite impacts training MAE. Within the MPNN, the largest task impact arises from removing the edge features and edge MLP from all layers (eur and MLP in Equation 3) falling back on simple GCN-style message passing that concatenates neighbouring node features when computing messages. Table 7 also shows that the use of Adjacent Node Feature Aggregation (i.e., the sum of adjacent node features x in Equation 7 allows large node features of size 256 to bypass compression into edge messages of size 128, affording a small but meaningful MAE improvement. It likely due to this feature that we have not found a benefit to increasing the edge latent size to match the node size. Both the use of Global Features and Sender Message Aggregation within the MPNN shown to be of comparable importance to the MHSA and FFN modules. Global Features, denoted as g² in Equation 6 and 7. allow the MPNN some degrec of global rcasoning. whilst Sender Mcssage Aggregation (the sum of e in Equation 7 uses outgoing messages in addition to the usual incoming messages to update a node's features. There is some overlap in purpose of the Global Features and the MHSA module, so the table also shows the impact of removing both at once. The resulting performance degradation is significantly greater than the sum of their individual impacts, which may imply each feature is able partially compensate for the loss of the other in the individual ablations.
12

---

Published in Transactions on Machine Learning Research (07/2023)
Table 7: Ablation of network architecture features.
MAE (meV)
Removed Feature
Valid
Train
Params (A)
None (Baseline)
77.2
53.6
44M ( 0%)
MPNN
7.7
10.9
14M (-70%)
Edge Features
4.2
+ 3.7
33M
(-26%)
Global Features
1.2
0.0
40M -8%]
Sender Message Aggregation
0.7
0.4
38M
(-14%)
Adjacent Node Aggregation
0.4
1.3
36M
(-19%)
MIISA
0.9
40M
(-10%)
FFN
1.0
1.3
36M
(999)
MHSA and Global Features
37M
(-18%)
8
Discussion
In this work, we define GPS++, hybrid MPNN/ OPPNNTTrasformic optimised for the PCQM4Mv2 molccular property prediction task (Hu al. (022).. Our model builds on previous work (Rampášek al. 2022 Luo
et al. 2022 Godwin et al. 2022) with a particular focus building powerful and expressive message passing component comprising 70% of model parameters. We showed that our GPS++ model has state-of-the-art performance on this uniquely large-scale dataset, and that despite recent trends towards using only graph transformers in molecular property prediction, GPS++ retains almost all its performance as a pure MPNN with the global attention module ablated. Finally, we consider the case where 3D positions are not available and show that our GPS++ models are significantly better than transformers, where significant drop-off is observed across all prior work. We also found that MPNN-only model performs better than hybrid under these circumstances; this indicates strong reliance between effective attention and the availability of 3D positional information.
Whilst these results are interesting for problems with small molecules, testing much larger molecules, for example, peptides, proteins, or RNAs, which can have hundreds or thousands of atoms, is a tougher challenge. Under these conditions, the linear complexity of MPNNs with graph size suggests some computational benefits over global attention, however, it is still be determined if the downsides of issues like underreaching outweigh these benefits. Nevertheless, domain and hope that it inspires revival of message passing and hybrid models for molecular property prediction when dealing with large dalasets.
believe that results highlight the strength of MPNNs in this
References
Uri Alon and Eran Yahav. On the bottleneck of graph neural networks and its practical implications. International Conference on Learning Representations, 2021.
Simon Axelrod and Rafael Gomez-Bombarell Geom, energy mneggy-nnottttdd molecular conformations for property prediction and molecular generation. Scientific Data, 9(1): 14, 2022.
Jimmy Lei Ba, Jamic Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXin:1607.064450 2016.
Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXie:1806. 01261, 2018.
Simon Batzner, Albort Musaclian, Lixin Sun, Mario Geiger, Jonathan Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E Smidt, and Boris Kozinsky. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature communnccatinns, 13(1):1-11, 11, 2022.
Jenna Bilbrey, Kristina M. Herman, Henry Sprucill, Soritis Xanthcas, Payel Das, Manuel Lopez Roldan, Mike Kraus, Hatem Helal, and Sutanay Choudhury. Reducing down(stream) time: Pretraining molecular gnns using heterogeneous ai accelerators. Xiv preprint 2211: 0.44998 2022.
13

---

Published in Transactions on Machine Learning Research (07/2023)
Johannes Brandstetter, Rob Hesselink, Elise van der Pol, Erik J Bekkers, and Max Welling. Geometric and physical quantities improve e(3) equivariant message passing. In International Conference on Learning Representations, 2022. URL
Xavier Bresson and Thomas Laurent. Residual gated graph convnets, 2018.
Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learning: Grids, groups, graphs geodesies, and gauges. arXii::104.13478, 2021.
Stefan Chmiela, Alexandre Tkatchenko, Huziel E Sauceda, Igor Poltavsky, Kristof T Schütt, and
Müller. Machine learning of accurate energy-conserving molecular force fields. Science advances, 3(5): e1603015, 2017.
Christopher M. Dobson. Chemical space and biology. Nature, 432((019):824 828, December 2004. ISSN 1476-4687. 10i (000100///auu
David K Duvenaud, Dougal Maclaurin, Jorge Iparraguirre, Rafael Bombarell, Timothy Hirzel, Alán Aspuru- Guzik, and Ryan P Adams. Convolutional networks on graphs for learning molccular fingerprints. Advances in neural information processing systems, 28, 2015.
Vijay Prakash Dwivedi and Xavier Bresson. A generalization of transformer networks to graphs. arXhv:2012.0999) 2020.
Vijay Prakash Dwivedi, Anh Tuan Luu, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. Graph neural networks with learnable structural and positional representations. In International Conference on Learning Representations, 2022.
Xiaomin Fang, Lihang Liu, Jieqiong Lei, Donglong He, Shanzhuo Zhang, Jingbo Zhou, Fan Wang, Hua Wu, and Haifeng Wang. Chemrl-gem: Geometry enhanced molecular representation learning for property prediction. arXiv preprint (XX::2106.06130 2021.
Matthias Fey and Jan Eric Lenssen. Fast graph representation learning with PyTorch Geometric. In ICLR Workshop on Representation Learning on Graphs and Manifolds, 2019.
Mare Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolutional neural networks for equivariance to lie groups arbitrary continuous data. In International Conference Machine Learning, pp. 2165-3176. PMLR, 2020.
Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. Se (3)-transffrrmer:: 3d roto-translation equivariant attention networks. Advances in Neural Information Processing Systems, 33:1970-195 19911, 2020.
Johannes Gasteiger, Shankari Giri, Johannes T. Margraf, and Stephan Günnemann. Fast and uncertainty- aware directional message passing for non-equilibrium molecules. In Machine Learning for Molecules Workshop, NewrlPS, 2020a.
Johannes Gasteiger, Janek Groß, and Stephan Günneemann. Directional message passing for molecular graphs. In International Conference on Learning Representations (ICLR), 2020b.
Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George Dahl. Neural message passing for quantum chemistry. In International conference on machine learning, pp. 2017.
PMLR,
Jonathan Godwin, Michael Schaarschmidt, Alexander L Gaunt, Alvaro Sanchez-Gonzalez Yulia Rubanova, Petar Veličković, James Kirkpatrick, and Peter Battaglia. Simple GNN regularisation for 3D molecular property prediction and beyond. In International Conference on Learning Representations, 2022.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (GELUs). arXiv:1606.08// 084160088111 2016.
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. 34th Conference on Neural Information Processing Systems, 2020.
14

---

Published in Transactions on Machine Learning Research (07/2023)
Weihua Hu, Matthias Fey, Hongyu Ren, Maho Nakata, Yuxiao Dong, and Jure Leskovee. OGB-LSC: A large-scale challenge for machine learning on graphs. In 35th Conference on Newral Information Processing Systems: Datasets and Benchmarks Track, 2021.
Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian Q. Weinberger. Deep networks with stochastic depth. In Bastian Leibe, Jiri Matas, Nicu Sebe, and Max Welling (eds.), Compuler Vision ECCV 2016, pp. 646 661, Cham, 2016. Springer International Publishing, ISBN 978-3-319-46493-00
Kexin Huang, Tianfan Fu, Wenhao Gao, Yue Zhao, Yusuf Roohani, Jure Leskovec, Connor W Coley, Cao Xiao, Jimeng Sun, and Marinka Zitnik. Artificial intelligence foundation for therapeutic science. Nature Chemical Biology, 18(10):1033-1086, 2022.
Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Mia Xu Chen, Dehao Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V. Le, Yonghui Wu, and Zhifeng Chen. Gpipe: Efficient. training of giant neural networks using pipeline parallelism. arXic preprint orXi::1811.068 2018.
Md Shamim Hussain, Mohammed J Zaki, and Dharmashankar Subramanian. Global self-attention replacement for graph convolution. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp.
2022.
Michael J Hutchinson, Charline Le Lan, Sheheryar Zaidi, Emilien Dupont, Yee Whye Teh, and Hyunjik Kim. Lietransformer: Equivariant self-attention for lie groups. International Conference on Machine Learning, pp. 4533-4543. PMLR, 2021.
Rui Jiao, Jiaqi Han, Wenbing Huang, Yu Rong, and Yang Liu. 3d equivariant molecular graph pretraining. arXiv preprint arXiv:2207.0889, 2022.
John A. Keith, Valentin Vassilev- Galindo, Bingqing Cheng, Stefan Chmiela, Michael Gastegger, Klaus- Robert Müller, and Alexandre Tkatchenko. Combining machine lcarning and computational chem- istry for predictive insights into chemical systems. Chemical Reviews, 121 9872, 2021. 10100.001//asscceeeee www.schemrrv..c000 PMTD: 34232033.
Jinwoo Kim, Tien Dat Nguyen, Sconwoo Min, Sungjun Cho, Moontae Lee, Honglak Lee, and Seunghoon Hong. Pure transformers powerful graph learners. In NewrIPS, 2022.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR (Poster), 2015. 6980
URL http:://arxiv..og/abbb//////////
Walter Kohn and Lu Jeu Sham. Self-consistent equations including exchange and correlation effects. Phys. Rev., 140:A1133 A1138, Nov 1965.
Mario Michael Krell, Matej Kosec, Sergio P. Perez, and Andrew Fitzgibbon. Efficient sequence packing without cross-contamination:: Accelerating large language models without impacting performance. arXiv preprint arXiv:2107.02027, 2021.
Devin Kreuzer, Dominique Beaini, William L. Hamilton, Vincent Létourneau, and Prudencio Tossou. Rethinking graph transformers with spectral attention. In Advances in Neural Information Processing Systems, 2021.
Tuan Le, Frank Noé, and Djork-Arné Clevert. Equivariant graph attention networks for molecular property prediction. arXiv preprint arXim:2202.09820 2022.
Qimai Li, Zhichao Han, and Xiao-Ming Wu. Deeper insights into graph convolutional networks for semi- supervised learning. In AAAI, pp. 1538-3545, 2018.
Lihang Liu, Donglong He, Xiaomin Fang, Shanzhuo Zhang, Fan Wang, Jingzhou Hc, and Hua Wu. GEM-2: Next generation molecular property prediction network with many-body and full-range interaction modeling. arxiv preprint ar arXi::2208.0540 2022a.
15

---

Published in Transactions on Machine Learning Research (07/2023)
Yi Liu, Limei Wang, Meng Liu, Yuchao Lin, Xuan Zhang, Bora Oztekin, and Shuiwang Ji. Spherical message passing for 3d molecular graphs. In International Conference on Learning Regresentations (ICLR), 2022b.
Chengqiang Lu, Qi Liu, Chao Wang, Zhenya Huang, Peize Lin, and Lixin He. Molecular property prediction: A multilevel quantum interactions modcling perspective. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 1052-1060, 2019.
Shengjüe Luo, Tianlang Chen, Yixian Xu, Shuxin Zheng, Tie-Yan Liu, Liwei Wang, and Di He. One transformer can understand both 2D & 3D molecular data. orXiv:2210.011 765, 2022.
Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, and Hao Wu. Mixed precision training. arXiv preprinl ar Xrx 777000000744 2017.
Maho Nakata and Tomomi Shimazaki. Pubchemqe project: large-scale first-principles electronic structure database for data-driven chemistry. Journal of chemical information and modeling, 37(66)1330-13300 2017.
Wonpyo Park, Woonggi Chang, Donggeon Lee, Juntac Kim, and Scung won Hwang. GRPE: Relative positional encoding for graph transformer. arXiw:22201.12787, 2022.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp. and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1-7, 2014.
Ladislav Rampášck, Mikhail Galkin, Vijay Prakash Dwivedi, Anh Than Lun, Guy Wolf, and Dominique Beaini. Recipe for General, Powerful, Scalable Graph Transformer. 2xXi:2205.12444 2022.
RDKit. RDKit: Open-source cheminformatics. http://ww.rdktt
[Online].
Patrick Reiser, Marlen Neubert, André Eberhard, Luca Torresi, Chen Zhou, Chen Shao, Houssam Metni, Clint van Hoesel, Henrik Schopmans, Timo Sommer, et al. Graph neural networks for materials science and chemistry. Communications Materials, 3(1):93, 2022.
David Rogers and Mathew Hahn. Extendedl-comnectiiiy fingerprints. Journal of chemical information and modeling, 50(5):742-755 2010.
Yu Rong, Yatao Bian, Tingyang Xu, Wei yang Xie, Ying Wei, Wen bing Huang, and Junzhou Huang. Grover: Self-supervised message passing transformer on large-scale molecular data. ArXiv, abs/2007.02$55 002835, 2020.
Lars Ruddigkeit, Rund Van Deursen, Lorenz C Blum, and Jean-Louis Reymond. Enumcration of 166 billion organic small molecules in the chemical universe database gdb-17. Journal of chemical information and modeling, 52(11))2864 -8755, 2012.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling, E (n) equivariant graph neural networks. International conference on machine learning, pp. 0223-932.. PMLR, 2021.
Kristof Schüitt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. Tn International Conference Machine Learning, pp. 9377-9388 PMLR, 2021.
Kristof T Schütt, Huziel E Sauceda, P-J Kindermans, Alexandre Tkatchenko, and K-R Müller. Schnet-a deep learning architecture for molecules and materials. The Journal of Chemical Physics, 148(24):241722, 2018.
Yu Shi, Shuxin Zheng, Guolin Ke, Yilei Shen, Jiacheng You, Jiyan He, Shengjie Luo, Chang Liu, Di He, and Tie Yan Liu. Benchmarking graphormer on large-scale molecular modeling datasets. arXiv:2203.04: 048100, 2022.
Justin S Smith, Olexandr Isayev, and Adrian E Roitberg. Ani-1, data set of 20 million calculated off-equilibrium conformations for organic molecules. Scientifie data, 4(1):1 8, 2017.
16

---

Published in Transactions on Machine Learning Research (07/2023)
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Lya Sutskever, and Ruslan Salakhutdinov. Dropout: simple way to prevent neural networks from overfitting. The journal of machine learning research, 15(1): 1929 1958, 2014.
Hannes Stärk, Dominique Beaini, Gabriele Corso, Prudencio Tossou, Christian Dallago, Stephan Günnemann, and Pietro Liò. 3d infomax improves gnns for molecular property prediction. In International Conference on Machine Learning, pp. 20502220502. PMLR, 2022.
Philipp Thölke and Gianni De Fabritiis. Equivariant transformers for neural network based molecular potentials. International Conference on Learning Representations, 2022a.
Philipp Thölke and Gianni De Fabritiis. Torchmd-net: Equivariant transformers for neural network based molecular potentials. arXiv preprint arXiv:220. 0, 2022b.
Richard Tran, Janice Lan, Muhammed Shuaibi, Siddharth Goyal, Brandon M Wood, Abhishek Das, Javier Hcras-Domingo, Adeesh Kolluru, Ammar Rizvi, Nima Shoghi, et al. The open catalyst 2022 (oc22) datasot. and challenges for oxide electrocatalysis. XrXiv:2266.068 08917, 2022.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Newral Information Processing Systems, 30,
2017.
Petar Veličković. Everything is connected: Graph neural networks. Curr Opin Struet Biol, 79, 2023.
Haorui Wang, Haoteng Yin, Muhan Zhang, and Pan Li. Equivariant and stable positional encoding for powerful graph neural networks. In International Conference on Learning Representations, 2022a.
Yusong Wang, Shaoning Li, Tong Wang, Zun Wang, Xinheng He, Bin Shao, and Tie-Yan Liu. How to better introduce geometric information in equivariant message passing? https://github: htw.g/ttt.u...oo.bbbb.oe.. visnet/ Global- wttaaa/iiiiieelooommmteeeee//eee Tectssstteeeeiitteeeeeeeehhcchee Repotttt 2022b. Accessed: 2022-11-16.
Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesso, Ancesh S Pappu, Karl Leswing, and Vijay Pande. Moleculenet: a benchmark for molecular machine learning. Chemical science, 9 (2):513-588 530, 2018.
Chengxuan Ying, Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di Hc, Yanming Shen, and Tic- Yan Liu. Do translormers really perform badly for graph representation? Tn Advances in Neurul Information Processing Systems, 2021a.
Chengxuan Ying, Mingqi Yang, Shuxin Zheng, Guolin Ke, Shengjie Luo, Tianle Cai, Chenglin Wu, Yuxin Wang, Yanming Shen, and Di He. First place solution of KDD Cup 2021 & OGB large-scale challenge graph prediction track. arXiv:2106. 00002206608879, 2021b.
Sheheryar Zaidi, Michael Schaarschmidt, James Martens, Hyunjik Kim, Yee Whye Teh, Alvaro Sanchez- Gonzalez, Peter Battaglia, Razvan Pascanu, and Jonathan Godwin. Pre-training via denoising for molecular property prediction. Xie preprinl 25i:22266000333 2022.
A Detailed Model Description
A.1
GPS++ Block Complete
The GPS++ block is defined as follows for layers l (see section A.2 for the definitions of X X
.)).
x
E
f-1
g
s+1 GPS++ (x' E g. B)
(17)
computed as
Y!
g
e+1 MPNN (x' E g)
(18)
Z' BiasedAttn (X',B)
(19)
Vi:
(+1
FFN (yf + +2
(20)
17

---

Published in Transactions on Machine Learning Research (07 (07/2023)
The MPNN module A simplified version of the MPNN module is presented in section 4.1 is follows:
The full description
Y E g MPNN (x', E²,gg)
(21)
computed as
v(u,)):
cyu [x x
e
g
(22)
v(u
e = Dropouto.0oss
(MLPedgo (ctv)
(23)
Vi: x = MLP MLPpade
x
e (u,i)EE
Σ
Σ (iF)E
Σ (u.S)EE
g
(24)
sender messages
receiver messsgggs
adjacent nodes
g' MLP global
Σ euv (L,
(25)
Vi: y LayerNor (x.) + +x{)
(26)
V(u, v):
e ue
e eu
(27)
Dropouto. (g') tg¹
(28)
where Dropout, Srivastava et al. 2014 masks by zero each element with probability p and LayerNorm follows the normalisation procedure of Ba et al. (0016). The three networks MLP,, for E {node, edge, global] each have two layers and are defined by:
y MLP,,
(29)
computed as
X GELU(De))e(x))
E RRd, Rd
(30)
y Dense(LayerNorm(X)
(31)
where GELU is from Hendrrcks Gimpel (2016
The BiasedAttn module Our BiasedAttn module follows the form of the self-attention layer in Luo et al. (2022) where a standard self-attention block (Vaswani et al. 2017] is biased by structural prior derived from the input graph. In our work the bias B is made up of two components, a Shortest Path Distance embedding and a 3D Distance Bias derived from the molecular conformations as described in section A.2 Single-head biased attention defined by:
Z BiasedAttn(X,B) B)
(32)
computed as
A X (XWx) +B B
E R N
(33)
Vdnude
Á Dropout (Softmax (A)) (XWv)
R Xddoode
(34)
Z GraphDropout $0.3 (A) +X
E Nxdode
(35)
for learnable weight matrices Wo. K. Wy E Ru ddaade and output in R. Nxdnode, though in practice we use 32 attention heads which are mixed before GraphDropout using affine projection Wp R deddd Xdnade
18

---

Published in Transactions on Machine Learning Research (07/2023)
The FFN module Finally, the feed-forward network module takes the form:
y FFN
(36)
computed as
X =
Dropout,
E ER ggddnods
(37)
y GraphDropout (Deesss)))) +x E Rdddde
(38)
Unless otherwise stated, the dropout probability 0. GraphDropout, also known as Stochastic Depth or LayerDrop (Huang et al. 2016), masks whole graphs together rather than individual nodes or features, relying on skip connections to propagate activatitions.
A.2 Input Feature Engineering
As described in section the dataset samples include the graph structure 9, a set of categorical features for the atoms and bonds X,, Guv, and the 3D node positions It has been shown that there many benefits to augmenting the input data with additional structural, positional, and chemical information (Rampášck
et al. 2022 Wang et al. 2022a Dwivedi et al., 2022). Therefore, we combine several feature sources when computing the input to the first GPS++ layer. There are four feature tensors to initialise; node state, edge
state, whole graph state and attention biases.
X all
X atto
X LapVec
X LLppVal
X RW
X Cent
X 3D
X Dense X
R Nxdnoda
(39)
E Dense bond
E
MM Xdedge
(40)
0 Embed Latoall (0)
E Rdiaba
(41)
B SPD 3D
NxN
(42)
The various components of each these equations are defined over the remainder of this section. The encoding of these fcatures also makes recurring use of the following two generic functions. Firstly, two-layer MLPencoder that projects features to fixed-size latent space:
y = MLP Peecoddr (z), where E E
(43)
computed as
T ReLUiiDense(Layerroom((3)))
E 1R
(44)
y Dropouto. (Bepottt_)))
R
(45)
Secondly, function Embeda(j) E Rd which selects the jth row from an implicit learnable weight matrix.
Chemical Features The chemical features used in this work are shown in Table A.1 To embed the categorical chemical features from the dataset X,, into a continuous vector space, we learn simple embedding vector for each category, sum the embeddings for all categories, and then process it with an MLP to produce atom and Ebond iee.
Vi
X atom
MLP node
Embeda4(j)
(46)
jEv
Vi Dropout
18
gatoo
Rududd
(47)
V(u, v)
bund
MLP edge
Embed6 (i)
(48)
jEaus band euv
V(u v)
exe
band
Dropout 18
E dddeog
(49)
Here MLPnode and MLPedge refer the functions by the same names used in Eq. 23 and yet parameterised independently.
24 in the MPNN module,
19

---

Published in Transactions on Machine Learning Research (07/2023)
Table A.1: Chemical input feature selection for PCQM4Mv2.
Feature Set
Node features
Original
Set
Set 2
Set 3
Ying21
Atomic number
Group
Period
Element type
Chiral tag
Degree
Formal charge
# Hydrogens
# Radical electrons
Hybridisation
Is aromatic
Is in ring
Is chiral center Explicit valence Implicit valence
Total valence
Total degree
Default valence
# Outer elecrons
Van der Waals radius
Covalent radius
# Bonds radius N 2
Gasteiger charge
Is donor
Is acceptor
Edge features
Bond type
Bond stereo
Is conjugated
Is in ring
Bond direction
20

---

Published in Transactions on Machine Learning Research (07/2023)
Graph Laplacian Positional Encodings (Kreuzer et al. 2021 Dwivedi & Bresson 2020
Given graph
with adjacency matrix A and degree matrix D, the eigendeccopoositoon of the graph Laplacian L is formulated into a global positional encoding as follows.
Vi x LappVec MLP MLPenoodrr ( ( [i. 2.. (..,....p)))
ER
where L = D = AUAU
(50)
A'
Vi: x; LapVal MLP MLPencoder
E =R332
||A'|
where A diag(A) [2... kLap
(51)
To produce fixed shape inputs despite variable numbers of eigenvalues eigenvectors per graph, we truncate / pad to the lowest 7 eigenvalues, excluding the first trivial eigenvalue A1 = O. We also randomise the cigenvector sign every epoch which is otherwise arbitrarily defined.
Random Walk Structural Encoding (Dwivedi et al. 2022) This feature captures the probability that random graph walk starting at node i will finish back at node i, and is computed using powers of the transition matrix P. This feature captures information about the local structures in the neighbourhood around each node, with the degree of locality controlled by the number of steps. For this submission random walks from 1 up to k.t = 16 steps were computed to form the feature vector.
Vi: x RW
(P) )i (",,,,,
(PARm
where P = D D-1A A
(52)
: RW MLP Pencoder (x)W)) E RR32
-MLPnnoddr
(53)
The graph centrality encoding is intended to allow the network to gauge the importance of node based on its connectivity, by embedding the degree (number of incident edges) of each node into a learnable feature vector.
Local Graph Centrality Encoding Ying et al. 202 Shi et al. 2022
Vi:
Cent Embed64 (Da) E 6
(54)
Shortest Path Distance Attention Bias Graphormer (Ying et al. 2021a Shi et al. 2022 showed that graph topology information can be incorporated into a node transformer by adding learnable biases to the self-attention matrix depending on the distance between node pairs. During data preprocessing the SPD map A E NNXN computed where Aij is the number of edges the shortest continuous path from node node j. During training each integer distance is embedded as a scalar attention bias term create the SPD attention bias map BSPD N
Vi,j:
B = Embed, (Ax) E R
(55)
Single-headed attention assumed throughout this report for simplified notation, however, upon extension to multi-headed attention, one bias learned per distance per head.
Embedding 3D Distances Using the 3D positional information provided by the dataset comes with a number of inherent difficulties. Firstly, the task is invariant to molecular rotations and translations, however, the 3D positions themselves are not. Secondly, the 3D conformer positions are only provided for the training data, not the validation or test data. To deal with these two issues and take advantage of the 3D positions provided we follow the approach of Luo et al. (2022).
To ensure rotational and translational invariance
only the distances between atoms, not the positions directly. To embed the scalar distances into vector space 1 we first apply K = 128 Gaussian kernel functions, where the tth function is defined as
[r
r
Vi,j:
(56)
ig
10
1
(wnt) (27)2)
Vi
exp
ER
(57)
2
la*
21

---

Published in Transactions on Machine Learning Research (07/2023)
with learnable parameters μ² and o.. The K elements concatenated into vector Day. We then process these distance embeddings in three ways to produce attention biases, node features and edge features.
3D Distance Attention Bias The 3D attention bias map B E NxN allows the model to modulate the information flowing between two node representations during self-attention based on the spatial distance between them, and are calculated as per Luo et al. (2022
Vi
B kj
MLP (b))) E R
(58)
Upon extension to multi-headed allention with 32 heads MLPhiasaD instead projects to R32.
Bond Length Encoding Whilst B makes inter-node distance information available to the self-attention
module in a dense all-to-all manner as a matrix of simple scalar biases, we also make this information available the MPNN module in sparse but high-dimensional manner as edge features ESD [e for (u, E] calculated
v(u,u)
3D euv
MLP encoder (1) RR32
(59)
Global 3D Centrality Encoding The 3D node centrality features X3D = [x [x?D: by summing the embedded 3D distances from node i to all other nodes. Since the sum commutes this feature cannot be used determine the distance to specific node, so serves centrality encoding rather than a positional encoding.
3D are computed
Vi:
3D = x
W-3D Σ Vis E
(60)
JEV
Here W3D E Kx32 is a linear projection to the same latent size as the other encoded features.
3D Denoising We also closely follow the approach of Luo al. 2022 implement an auxiliary self- supervised 3D denoising task during training. Before the 3D distance inputs are embedded (Equation 56 the atom positions are substituted by R+ oe, where R scaling factor set to 0.2, and e R R-x3 are Gaussian noise vectors. GPS++ computes ésk: the predicted value of Esk,
Σ A JEV
Eis
W ER
(61)
where
Á (' (X
+B B
(62)
Vaudoo
A Dropout (Softmax (A)) NXN
(63)
Here X is the node output of the final GPS++ layer, A is the k-th element of directional vector Ft Ty from atom j to atom i, and w Q w W E Rdoode xduade, WDD E Rd x1 are learnable weight matrices. V/ A cosine similarity loss is computed between and 6, and we set the loss weight ratio of Homo-Lumo MAF vs 3D Denoising to 10:
B Hardware and Acceleration
B.1
Hardware
We train our models using a BOW-POD16 which contains 16 IPU processors, delivering total of 5.6 petaFT.OPS of float16 compute and 14.4 GB of in-processor SRAM which is accessible at aggregate
22

---

Published in Transactions on Machine Learning Research (07/2023)
bandwidth of over petabyte per second. This compute and memory is then distributed evenly over 1472 tiles per processor. This architecture has two key attributes that enable high performance on GNN and other Al workloads (Bilbrey et al. 2022): memory is kept as close to the compute as possible (i.e., using on-chip SRAM rather than off-chip DRAM) which maximises bandwidth for a nominal power budget; and compute is split up inlo many small independent arithmetic units meaning that any available parallelism can be extremely well utilised. In particular this enables very high performance for sparse communication ops, like gather and scatter, and achieves high FLOP utilisation even with complex configurations of smaller matrix multiplications. Both of these cases particularly prevalent in MPNN structures like those found in GPS++
To exploit the architectural benefits of the IPU and maximise utilisation, understanding the program structure ahead of time key. This means all programs must be compiled end-to-end, opening up a range of opportunities for optimisation but also adding the constraint that tensor shapes must be known and fixed at. compile time.
B.2 Batching and Packing
To enable fixed lensor sizes with variable sized graphs it is common to pad the graphs to the max node and edge size in the dataset. This, however, can lead lots of compute being wasted on padding operations, particularly in cases where there are large variations in the graph sizes. To combat this it is common pack number of graphs into fixed size shape to minimise the amount of padding required, this is an abstraction that is common in graph soflware frameworks like PyTorch Geometric (Fey & Lenssen 2019) and has been shown to achieve as much as 2x throughput improvement for variable length sequence models (Krell et al.) 2021). Packing graphs into one single large pack, however, has a couple of significant downsides: the memory and compute complexity of all-to-all attention layers is (n²)) in the pack size not the individual graph sizes, and allowing arbitrary communication between all nodes in the pack forces the compiler to choose sub-optimal parallelisation schemes for the gather/scatter operations.
To strike balance between these two extremes we employ two tiered hierarchical batching scheme that packs graphs into fixed size but then batches multiple packs to form the micro-batch. We define the maximum pack size to be 60 nodes, 120 edges and 8 graphs then use a simple streaming packing method where graphs are added to the pack until either the total nodes, edges or graphs exceeds the maximum size. This achieves 87% packing efficiency of the nodes and edges with on average 3.6 graphs per pack, though we believe that this could be increasod by employing a more complex packing strategy (Krell et al., 2021p. We then form micro-batches of 8 packs which are pipelined (Huang el al., 2018) over IPUs accunulating over 8 micro- micro-batches batches and replicated 4 times to form global batch size of 921 graphs distributed over 16 IPUs. For the MPNN-only model, the micro-batch size was increased to 15, forming global batches of 1737 graphs
B.3 Numerical Precision
To maximisc compute throughput and maximisc memory efficiency it is now common practice to use lower precision numerical formats in deep learning (Micikevicius et al., 2017). On IPUs using float16 increases the peak FTOP rate by 4x compared to float32 but also makes more effective usage of the high bandwidth on-chip SRAM. For this reason we use float16 for nearly all compute but also use float16 for the majority 2 of the weights, this is made possible, without loss of accuracy, by enabling the hardware-level stochastic rounding of values. While we do use a loss scaling (Micikevicius et al., 2017) value of 1024 we find that our results are robust to a wide range of choices. We also use float16 for the first-order moment in Adam but keep the second-order moment in float32 due to the large dynamic range requirements of the sum-of-squares.
IA few operations like the sum squares the variance caleulations are up-cast to float32 by the compiler. 2 small number of weights are kept float32 for simplicity of the code rather than numerical stability.
23


Reviews:
[{"review_id": "YANWaG6cRw", "reviewer": "Reviewer_Souw", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: Although only study a very specific application in graph learning, the paper clearly defines its scope and provides enough context for the problem setup and discussion of prior works. The proposed architecture combines the graph transformers and the message-passing networks and achieves state-of-the-art results in the OGB-LSC task. The paper provides extensive ablation empirical studies on the properties of the node/edge/3D features as well as the effectiveness of the massage-passing neural module. The architecture of the proposed method is well presented and the structure of the paper is easy to follow.\n\nWeaknesses: I have two main concerns for this work: novelty and generalizability. The proposed architecture largely resembles that of GPS, with limited modifications. Also, it is not clear the motivations behind the design of the new GPS++ block. The paper only uses a single dataset in all experiments. It is not clear whether the proposed method only achieves better performance in this specific case or it's general enough to be used in other molecular graph learning tasks.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "In this paper, the authors propose a novel graph-learning architecture designed specifically for molecular property prediction tasks. The proposed model, namely GPS++, consists of a combination of architectures/modules that exists in prior works.\nThe paper primarily focuses on the molecules HOMO-LUMO energy gap prediction and shows the proposed method achieves state-of-the-art results on the OGB-LSC dataset PCQM4Mv2. In addition to the main results, the paper also includes extensive ablation studies to showcase the contribution of each neural module and input features to the final results.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* Include empirical studies on more datasets/tasks to examine the general applicability of the GPS++ module.\n* Add more discussions on the motivations and reasoning behind the design of the new neural architecture.", "broader_impact_concerns": "I don't have ethical concerns about this work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "EZDz3ez9el", "reviewer": "Reviewer_T4FP", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "#### Strengths ####\n\n(1) This work provides exhaustive experiments to demonstrate a lot number of different model architecture designs and feature engineering. The extensive experimental results could be useful for the community for future research.\n\n(2) The writing and the organization of the paper is super clear and easy to follow.\n\n(3) The good empirical performance of the dedicatedly designed model, including feature engineering.\n\n#### Weaknesses ####\n\n(1) One of the biggest weaknesses is the novelty of the method. Since GPS++ closely follows the general design of previous models, the architecture itself is not new. Since TMLR does not require a lot of novelty, I will not criticize this aspect in terms of making an acceptance decision.\n\n(2) The experiments are only on a single dataset and task, although the dataset is super large. Readers may wonder if the exhaustive model designs and featurization work for other molecular property prediction tasks. So I highly recommend adding at least one more experiment to make the paper stronger.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work proposes a deep model for graph learning, particularly predicting the HOMO-LUMO gap of molecules. It closely follows the design of a previous work GPS, which combines the operation of message passing and transformer. The main difference between this GPS++ work, compared to GPS, lies in a well-tuned MPNN model, feature engineering, and dedicatedly designed training objective.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "(1) It would be better to evaluate the GPS++ method on more tasks. ", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "OUuSP0zslS", "reviewer": "Reviewer_9ysq", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\nThe main strength of this paper is the detailed empirical evaluation and the extensive ablation studies of the different design choices of the network architecture. The authors utilize the generality of the GPS framework to combine various established feature engineering methods and network components, and in doing so achieve a state-of-the-art performance on the selected task. Furthermore, the detailed ablation studies can potentially be used to guide the design decisions of researchers or engineers when encountered with a similar task in the future.\nAnother strength of the work is its clearness and readability, with the authors clearly stating the reason and justification for each design decision, and having a well defined notation, making the equations easy to follow.\nWeaknesses:\nThe two main weaknesses of this work are its narrow focus and the lack of novelty. \nMost of the other state-of-the-art GNN models included in the paper as a point of comparison have been evaluated on a number of graph-based tasks and datasets, showing their generality and potential for wide range application. On the other hand the GPS++ was only evaluated on the HOMO-LUMO gap prediction task of the PCQM4Mv2 dataset. Due to the large number and specificity of the components used in GPS++,  it is unclear whether its good performance generalizes to other graph-based tasks or whether its architecture is fine-tuned to provide exceptional results only for the selected task.\nIn addition to the narrow focus, the work also fails to present any novel theoretical contributions to the field, although to the credit of the authors they also do not claim any such contribution. The model is a specific instantiation of the GPS framework, which is clearly signaled by its name GPS++, and largely uses already established features and network components that are introduced in a lot of the competing models. As mentioned previously, the contribution mainly comes in the way that these components are combined together and the extensive ablation studies. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors present a hybrid graph neural network model called the GPS++, utilizing elements of both MPNNs and transformer architectures, focused on the HOMO-LUMO gap prediction task on the  PCQM4Mv2 dataset. The paper is a mainly empirical contribution, including and evaluating various features and network blocks within the GPS framework, resulting in a very good performance on the selected regression task, outperforming other state-of-the-art GNN models. The paper also includes extensive ablation studies, clearly showcasing the contribution of the individual features/network components to the overall performance, leading to the interesting finding that most of the performance can be attributed to the MPNN component as opposed to the attention component of the architecture, with the MPNN alone showing better performance on the 2d structure task.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Overall, I think the paper is good as it is, the authors set a specific objective and achieve it, giving the paper a narrow but clear structure. While it would certainly be beneficial to include results for other similar datasets such as QM9, PDBBind, etc, which would significantly strengthen the work, I think the work as it is just about satisfies the acceptance criteria of TMLR and I would recommend it for acceptance without requiring any major changes in the content. \nHowever, for the purposes of reproducibility, I believe it is also important for the authors to provide access to the code used to produce the results, and unless I missed it i did not see the code included as a part of the submission. Not having the code available is potentially a reason to not accept the paper, since reproducibility is even more important given the empirical nature of this work and lack of theoretical novelty.\n\nIn term of minor changes, some more details on the training procedure would be appreciated, specifically about how the auxiliary task of 3d positions denoising is incorporated into the training process. The authors give ratios of the losses for the noisy nodes/edges task during training, but no such details are given for the 3d denoising task for training.", "broader_impact_concerns": "I have no concerns about the ethical implications of this work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "EmQQDOwQdK", "reviewer": "Reviewer_uvAP", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "This paper studied how to optimize the input features to improve the performance of GPS. The proposed method achieved state-of-the-art performance with an interesting hybrid architecture that combines transformers and message-passing neural networks, which is a type of graph neural networks.\nThe paper reads well, and the discussion in the paper is straightforward. The authors introduced the augmented general powerful scalable graph transformer in Fig 1., the new layer GPS++ layer, and feature sets.\nAn extensive ablation study supports the value of each component in the proposed method.\nHowever, it is not clear which component is newly proposed in this paper. The authors did not explicitly compare the proposed method with the previous method/baseline GPS regarding neural network architectures.\nIn addition, in Table 2, the main experimental results raise a couple of concerns. First, in the experimental setting without 3D positional information, GPS has a twice smaller number of parameters than GPS++. This is a huge difference. Also, the proposed Hybrid method GPS++ GPS++ [No 3D Position] with more parameters underperforms GPS++ [MPNNonly, No 3D Position].\nAlthough the proposed method achieved considerable performance gain in the setting with limited information, with 3D positional information, the performance gain is marginal, and GPS is not even compared.\nIn sum, the proposed extension has limited novelty, and performance gain is questionable.\n\n**Strengths**:\n\n1. The proposed method achieved strong performance and good efficiency by extensive optimization for architectural components, features, and structural encodings.\n2. Extensive ablation study is provided and each component in the final model is well justified by empirical performance gain.\n\n**Weaknesses**:\n\n1. Although this paper admitted that the proposed method is a specific implementation of GPS, no clear comparison between GPS and GPS++ layers. So, it is hard to evaluate the contributions of this paper unless the differences are explicitly discussed. \n2. This paper is incremental and the method section mainly introduces the implementation details including feature engineering. It may be crucial for improving performance but beyond the specific benchmarks the feature engineering has limited academic merit. \n3. ~~In the setting with 3D positional information, the performance gain is marginal. Also, in Table 2, compared to GPS, GPS++ seems to improve the performance but the model size is doubled. So, it may not be a fair comparison. The authors need to provide well-tuned GPS with a comparable number of parameters~~. [addressed by authors' feedback:GPS++ (8 layers)]\n4. Table 2, GPS is missing in with 3D Positional information.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper extends GPS. The proposed method GPS++ is a hybrid model that combines message-passing neural networks and graph transformer models. The authors tackled molecular property prediction. The proposed method achieved state-of-the-art performance. Lastly, the authors provide extensive ablation studies to investigate the effects of various features, structural encodings, and architectural components.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. ~~Overall the paper reads well. However, as mentioned above, the main contribution/difference between GPS and the proposed model needs to be discussed more clearly.~~ [addressed by author feedback, Section 4. in the revised version]\n2. ~~Given graph laplacian, eig vec, eig vals can be used for positional embeddings or features. However, it is not clear how to compute them for heterogeneous graphs with various edge/node types. Did you convert the heterogenous graphs into a homogeneous graph and compute $X^\\text{LapVec}$ and $X^\\text{LapVal}$?~~ [addressed by author feedback] treated as a homogeneous graph\n3. ~~In Eq (7), why is $\\sum_{(i,v) \\in E}x_v^l$ included? Often, message passing along reversed edges improve the performance.~~ [addressed by author feedback]\n4. ~~The feature set Ying21 is missing in Table A.1~~. [addressed by author feedback]", "broader_impact_concerns": "I have no ethical concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Both reviewers and authors engaged in a fruitful discussion that resulted in substantial improvements of the paper. 

All reviewers like the paper especially after the revisions but are also concerned about the lack of changes over the original GPS method. In the end the reviewers were convinced that the changes merit publication.
Decision: Accept (poster)