# GPS++: Reviving the Art of Message Passing for Molecular Property Prediction

Dominic Masters  Josef Dean  Kerstin Klaser  Zhiyi Li  Sam Maddrell-Mander
Adam Sanders  Hatem Helal  Deniz Beker  Andrew Fitzgibbon  Shenyang Huang
Ladislav Rampášek  Dominique Beaini

###### 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, contrary to expectations set by recent
trends, find that nearly all of the model’s performance can be maintained
without any use of global self-attention. We also show that our approach is
significantly more accurate than prior art when 3D positional information is
not available.

Molecular property prediction, Graph learning, Hybrid MPNN/Transformer, OGB-
LSC PCQM4Mv2

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 §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 a custom
local MPNN and a global biased self-attention mechanism (akin to
Transformer-M) to learn expressive representations, see §4.1. (right) Global
graph pooling with final prediction head, and auxiliary denoising tasks, see
§5.2.

##  1 Introduction

Among many scientific areas, deep learning is having a transformative impact
on molecular property prediction tasks for biological and chemical
applications (Keith et al., 2021; Reiser et al., 2022). In particular, the aim
is to replace or augment expensive and/or time-consuming experiments and
first-principles methods with more efficient machine learning models. While
there is a 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.,
2021a; Luo et al., 2022). While transformers have been extremely successful in
other domains, the quadratic complexity with respect to the number of nodes is
a significant obstacle for larger molecules. 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 &
Yahav, 2021) have been found to limit their effectiveness.

Table 1: Comparison of ensembled model accuracy on OGB-LSC PCQM4Mv2 dataset with under four-hour inference time constraint. Strong single-model performance and efficient implementation on Graphcore IPU hardware enables a large ensemble to achieve the state-of-the-art accuracy. Model/Method (Ensemble Size) | Test-Challenge
---|---
MAE (meV) $\downarrow$
lightgray!20Transformer-M, GNN and ResNet |
lightgray!20 (35 models) Darabi et al. (2022) | 72.3
Transformer-M-ViSNet and Pretrained-3D-ViSNet |
(22 models) Wang et al. (2022c) | 72.3
lightgray!20GPS++ (112 models) | 71.9

In this work we focus on the task of predicting the HOMO-LUMO energy gap, an
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) (Kohn & 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 more than 8 hours per molecule per CPU (Axelrod & Gomez-
Bombarelli, 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 a 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 (Nakata & Shimazaki, 2017) is one of the largest widely
available DFT databases, and from it is derived the PCQM4Mv2 dataset, released
as a part of the Open Graph Benchmark Large Scale Challenge (OGB-LSC) (Hu et
al., 2021), which has served as a popular testbed for development and
benchmarking of novel graph neural networks (GNNs). The original OGB-LSC 2021
competition motivated a 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 number of graph
transformer methods for this task with comparatively little attention given to
the message passing methods that have been successful in other graph-
structured tasks.

In this work we build on the work of Rampášek 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 message-passing module with a
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, carefully crafting a range of diverse
input features, and building a large ensemble we achieve the best-reported
result on the PCQM4Mv2 test-challenge data split of 71.9 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 a global self-attention mechanism (seen in graph transformer
architectures), competitive performance can be achieved. Therefore, we argue
that MPNNs remain competitive 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 as follows:

  * •

We show that our hybrid MPNN/Transformer model, GPS++, is a parameter-
efficient and effective approach to molecular property prediction, achieving
lower MAE for PCQM4Mv2 compared to models of similar size and comparable to
the best-reported error for larger models.

  * •

We find that even without self-attention, our model is competitive with prior
state-of-the-art 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 highlighting the impact of the improved chemical feature choice, 3D
positional features and architectural components.

  * •

We show that a large diverse ensemble of GPS++ models achieves a test MAE
result on PCQM4Mv2 that is the state-of-the-art at time of submission (Table
1).

Reproducibility: Source code to reproduce our results can be found at:
https://github.com/graphcore/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).

GNN method development has also been facilitated by the availability of well
motivated (suites of) benchmarking datasets 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 et
al., 2022) for molecular property prediction or MD17 (Chmiela et al., 2017),
ANI-1 (Smith et al., 2017), Open Catalyst (Tran et al., 2022) for structural
conformers or molecular force field predictions.

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), 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
et al., 2022) and pretraining (Zaidi et al., 2022), GEM-2 (Liu et al., 2022),
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 et al., 2022), GPS (Rampášek
et al., 2022), TokenGT (Kim et al., 2022), or Transformer-M (Luo et al.,
2022), that have superseded message-passing GNNs.

##  3 Preliminaries

Throughout the paper we use the following notation. Bold lowercase letters
$\mathbf{v}$ are (row) vectors, bold uppercase letters $\mathbf{M}$ are
matrices, with individual elements denoted by non-bold letters i.e. $v_{k}$ or
$M_{pq}$. Blackboard bold lowercase letters $\mathbbm{v}$ are categorical
(integer-valued) vectors. In general, we denote by
$\bigl{[}\mathbf{v}_{k}\bigr{]}_{k\in K}$ the vertical concatenation
(stacking) of vectors $\mathbf{v}_{k}$. Vertical concatenation is also denoted
by a semicolon, i.e.
$\bigl{[}\mathbf{v}_{1}~{};~{}...~{};~{}\mathbf{v}_{J}\bigr{]}=\bigl{[}\mathbf{v}_{j}\bigr{]}_{j=1}^{J}=\bigl{[}\mathbf{v}_{j}~{}\mathsf{for}~{}j\in\\{1..J\\}\bigr{]}$.
Horizontal concatenation, which typically means concatenation along the
feature dimension, is denoted by a vertical bar, i.e.
$\bigl{[}\mathbf{v}_{1}\mid\mathbf{v}_{2}\bigr{]}$.

A molecule is represented as a graph $\mathcal{G}=(\mathcal{V},\mathcal{E})$
for nodes $\mathcal{V}$ and edges $\mathcal{E}$. In this representation, each
node $i\in\mathcal{V}$ is an atom in the molecule and each edge
$(u,v)\in\mathcal{E}$ is a chemical bond between two atoms. The number of
atoms in the molecule is denoted by $N=|\mathcal{V}|$ and the number of edges
is $M=|\mathcal{E}|$.

Each node and edge is associated with a list of categorical features
$\mathbbm{x}_{i}\in\mathbb{Z}^{D_{\text{atom}}}$ and
$\mathbbm{e}_{uv}\in\mathbb{Z}^{D_{\text{bond}}}$, respectively, for
$D_{\text{atom}}$ atom features and $D_{\text{bond}}$ bond features. A further
set of 3D atom positions
$\mathbf{R}=\bigl{[}\mathbf{r}_{1}~{};~{}...~{};~{}\mathbf{r}_{N}\bigr{]}\in\mathbb{R}^{N\times
3}$, 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 $\ell$ are denoted by
$\mathbf{x}_{i}^{\ell}\in\mathbb{R}^{{d_{\text{node}}}}$, and are concatenated
into the $N\times{d_{\text{node}}}$ matrix
$\mathbf{X}^{\ell}=\bigl{[}\mathbf{x}^{\ell}_{1}~{};~{}...~{};~{}\mathbf{x}_{N}^{\ell}\bigr{]}$.
Edge features $\mathbf{e}_{uv}^{\ell}\in\mathbb{R}^{{d_{\text{edge}}}}$ are
concatenated into the edge feature matrix
$\mathbf{E}^{\ell}=\bigl{[}\mathbf{e}^{\ell}_{uv}~{}\mathsf{for}~{}(u,v)\in\mathcal{E}\bigr{]}$.
Global features are defined per layer as
$\mathbf{g}^{\ell}\in\mathbb{R}^{{d_{\text{global}}}}$.

We also define an attention bias matrix $\mathbf{B}\in\mathbb{R}^{N\times N}$,
computed from the input graph topology and 3D atom positions, described later
in §4.2.

##  4 GPS++

Our GPS++ model closely follows the GPS framework set out by Rampášek et al.
(2022). This work presents a flexible model structure for building hybrid
MPNN/Transformer models for graph-structured input data. We build a specific
implementation of GPS that focuses on 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.

The main GPS++ block (§4.1) combines the benefits of both message passing and
attention layers by running them in parallel before combining them with a
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 (§4.2), and is followed by a simple
Decoder function (§5.2).

Feature engineering is also used to improve the representation of the
atoms/bonds, to 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 $\ell>0$ (see §4.2 for the
definitions of $\mathbf{X}^{0},\mathbf{E}^{0},\mathbf{g}^{0},\mathbf{B}$):

$\displaystyle\mathbf{X}^{\ell+1},\mathbf{E}^{\ell+1},\mathbf{g}^{\ell+1}$
$\displaystyle=\texttt{GPS++}\left(\mathbf{X}^{\ell},\mathbf{E}^{\ell},\mathbf{g}^{\ell},\mathbf{B}\right)$
(1) computed as $\displaystyle\mathbf{Y}^{\ell},\ \mathbf{E}^{\ell+1},\
\mathbf{g}^{\ell+1}$
$\displaystyle=\texttt{MPNN}\left(\mathbf{X}^{\ell},\mathbf{E}^{\ell},\mathbf{g}^{\ell}\right),$
(2) $\displaystyle\mathbf{Z}^{\ell}$
$\displaystyle=\texttt{BiasedAttn}\left(\mathbf{X}^{\ell},\mathbf{B}\right),$
(3) $\displaystyle\forall i:~{}~{}~{}\mathbf{x}_{i}^{\ell+1}$
$\displaystyle=\texttt{FFN}\left(\mathbf{y}_{i}^{\ell}+\mathbf{z}_{i}^{\ell}\right)$
(4)

Our MPNN module is a variation on the neural message passing module with edge
and global features (Gilmer et al., 2017; Battaglia et al., 2018; Bronstein et
al., 2021). We choose this form to maximise the expressivity of the model 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 2 for a graphical representation, and Appendix A.1 for
the exact formulation):

$\displaystyle\mathbf{Y}^{\ell},\ \mathbf{E}^{\ell+1},\
\mathbf{g}^{\ell+1}=\texttt{MPNN}\left(\mathbf{X}^{\ell},\mathbf{E}^{\ell},\mathbf{g}^{\ell}\right),$
(5) computed as $\displaystyle\forall(u,v):~{}~{}~{}$
$\displaystyle\bar{\mathbf{e}}_{uv}^{\ell}=\texttt{MLP}_{\textrm{edge}}\left(\bigl{[}\mathbf{x}_{u}^{\ell}\mid\mathbf{x}_{v}^{\ell}\mid\mathbf{e}_{uv}^{\ell}\mid\mathbf{g}^{\ell}\bigr{]}\right)$
(6) $\displaystyle\forall i:~{}~{}~{}$
$\displaystyle{\mathbf{c}}_{i}^{\ell}=\Biggl{[}\mathbf{x}_{i}^{\ell}~{}\Bigg{|}\underbrace{\sum_{(u,i)\in\mathcal{E}}\bar{\mathbf{e}}_{ui}^{\ell}}_{\begin{subarray}{c}\text{sender}\\\
\text{messages}\end{subarray}}~{}\Bigg{|}\underbrace{\sum_{(i,v)\in\mathcal{E}}\bar{\mathbf{e}}_{iv}^{\ell}}_{\begin{subarray}{c}\text{receiver}\\\
\text{messages}\end{subarray}}~{}\Bigg{|}\underbrace{\sum_{(u,i)\in\mathcal{E}}\mathbf{x}_{u}^{\ell}}_{\begin{subarray}{c}\text{adjacent}\\\
\text{nodes}\end{subarray}}~{}\Bigg{|}~{}\mathbf{g}^{\ell}\Biggr{]}$ (7)
$\displaystyle\forall i:~{}~{}~{}$
$\displaystyle\bar{\mathbf{x}}_{i}^{\ell}=\texttt{MLP}_{\textrm{node}}(\mathbf{c}_{i}^{\ell})$
(8)
$\displaystyle\bar{\mathbf{g}}^{\ell}=\texttt{MLP}_{\textrm{global}}\left(\Biggl{[}{\mathbf{g}^{\ell}~{}\Bigg{|}\sum_{j\in\mathcal{V}}\bar{\mathbf{x}}_{j}^{\ell}~{}\Bigg{|}\sum_{(u,v)\in\mathcal{E}}\bar{\mathbf{e}}_{uv}^{\ell}}\Biggr{]}\right)$
(9) $\displaystyle\forall i:~{}~{}~{}$
$\displaystyle\mathbf{y}_{i}^{\ell}=\bar{\mathbf{x}}_{i}^{\ell}+\mathbf{x}_{i}^{\ell}$
(10) $\displaystyle\forall(u,v):~{}~{}~{}$
$\displaystyle\mathbf{e}_{uv}^{\ell+1}=\bar{\mathbf{e}}_{uv}^{\ell}+\mathbf{e}_{uv}^{\ell}$
(11)
$\displaystyle\mathbf{g}^{\ell+1}=\bar{\mathbf{g}}^{\ell}+\mathbf{g}^{\ell}.$
(12)

The three networks $\texttt{MLP}_{\eta}$ for
$\eta\in\\{\textrm{node},\textrm{edge},\textrm{global}\\}$ each have two
layers with GELU activation functions and an expanded intermediate hidden
dimension of $4d_{\eta}$.

This message passing block is principally the most similar to Battaglia et al.
(2018). However we draw the reader’s attention to a few areas that differ from
common approaches. First we aggregate over the adjacent node representations
as well as the more complex edge feature messages in Equation 7, this is
similar to running a simple graph convolutional network (GCN) in parallel to
the more complex message calculation. We also aggregate over both sender and
receiver edges but critically concatenate rather than sum the results to
maintain some degree of directionality. Finally, we choose to decouple the
three latent sizes, setting ${d_{\text{node}}}=256$, ${d_{\text{edge}}}=128$
and $d_{\textrm{global}}=64$. While the primary reason is that, empirically,
we find increasing ${d_{\text{edge}}}$ and $d_{\textrm{global}}$ does not
improve performance, this choice also leads to greater efficiency.

Figure 2: The main GPS++ processing block (left) is composed of a local
message passing MPNN module and a 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 as above.

Our BiasedAttn module follows the form of the self-attention layer of Luo et
al. (2022) where a standard self-attention block (Vaswani et al., 2017) is
biased by a structural prior derived from the input data. In our work, the
bias $\mathbf{B}$ is made up of two components, a shortest path distance (SPD)
embedding and a 3D distance bias derived from the molecular conformations as
described in §4.2. The FFN module takes a similar form to
$\texttt{MLP}_{\textrm{node}}$ though with additional dropout terms (see
Appendix A.1 for full details).

###  4.2 Input Feature Engineering

As described in §3, the dataset samples include the graph structure
$\mathcal{G}$, a set of categorical features for the atoms and bonds
$\mathbbm{x}_{i}$, $\mathbbm{e}_{uv}$, and the 3D node positions
$\mathbf{r}_{i}$. It has been shown that there are many benefits to augmenting
the input data with additional structural, positional, and chemical
information (Rampášek 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.

$\displaystyle\mathbf{X}^{all}$
$\displaystyle=\bigl{[}\mathbf{X}^{\text{atom}}\mid\mathbf{X}^{\text{LapVec}}\mid\mathbf{X}^{\text{LapVal}}\mid\mathbf{X}^{\text{RW}}\mid\mathbf{X}^{\text{Cent}}\mid\mathbf{X}^{\text{3D}}\bigr{]}$
$\displaystyle\mathbf{X}^{0}$
$\displaystyle=\texttt{Dense}(\mathbf{X}^{all})\qquad\qquad\in\mathbb{R}^{N\times{d_{\text{node}}}}$
(13) $\displaystyle\mathbf{E}^{0}$
$\displaystyle=\texttt{Dense}(\bigl{[}\mathbf{E}^{\text{bond}}\mid\mathbf{E}^{\text{3D}}\bigr{]})\
\ \in\mathbb{R}^{M\times{d_{\text{edge}}}}$ (14) $\displaystyle\mathbf{g}^{0}$
$\displaystyle=\texttt{Embed}_{{d_{\text{global}}}}(0)\qquad\quad\ \ \
\in\mathbb{R}^{{d_{\text{global}}}}$ (15) $\displaystyle\mathbf{B}$
$\displaystyle=\mathbf{B}^{\mathrm{SPD}}+\mathbf{B}^{\text{3D}}\qquad\qquad\in\mathbb{R}^{N\times
N}$ (16)

Here the node features $\mathbf{X}^{0}$ are built from categorical atom
features, graph Laplacian positional encodings (Kreuzer et al., 2021; Dwivedi
& Bresson, 2020), random walk structural encodings (Dwivedi et al., 2022),
local graph centrality encodings (Ying et al., 2021a; Shi et al., 2022) and a
3D centrality encoding (Luo et al., 2022). The edge features $\mathbf{E}^{0}$
are derived from categorical bond features and bond lengths, and the attention
bias uses SPD and 3D distances. These input features are further described in
Appendix A.2.

#### Chemical Features

The categorical features $\mathbbm{x}$, $\mathbbm{e}$ 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 as 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 9 atom and 3 bond features
(described in Table A.1). There is, however, a wide range of chemical features
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 a
total of 28 atom and 5 bond features. We explore the impact of a number of
additional node and edge features and sweep a wide range of possible
combinations.

In particular, we expand on the set defined by Ying et al. (2021b) with three
additional atom features derived from the periodic table, the atom group
(column), period (row) and element type (often shown by colour). We found that
these three additional features were particularly beneficial.

We also found that in many cases _removing_ features was beneficial, for
example, we found that all 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 input features stay consistent between the
sets. We show the three best feature sets found in Table A.1 and use _Set 1_
for all experiments unless otherwise stated (e.g., during ablations and
ensembling).

##  5 Experimental Setup

###  5.1 Dataset

The PCQM4Mv2 dataset (Hu 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 by OGB (Hu et al.,
2021), namely into training (90%), validation (2%), test-dev (4%) and test-
challenge (4%) sets using a scaffold split where the HOMO-LUMO gap targets are
only publicly available for the training and validation splits.

###  5.2 Model Training

#### Training Configuration

Our model training setup uses the Adam optimiser (Kingma & Ba, 2015) with a
gradient clipping value of 5, a peak learning rate of 4e-4 and the model is
trained for a total of 450 epochs. We used a learning rate warmup period of 10
epochs followed by a linear decay schedule.

#### Decoder and Loss

The final model prediction is formed by global sum-pooling of all node
representations and then passing it through a 2-layer MLP. The regression loss
is the mean absolute error (L1 loss) between a scalar prediction and the
ground truth HOMO-LUMO gap value.

#### Noisy Nodes/Edges

Noisy nodes (Godwin et 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 tries 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 beneficial when
training deep GNNs (Godwin et 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 $p_{\text{corrupt}}$, 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 $p_{\text{corrupt}}=0.01$ and weight the cross-entropy losses
such that they have a ratio 1:1.2:1.2 for losses HOMO-
LUMO:NoisyNodes:NoisyEdges.

Table 2: Comparison of single model performance on PCQM4Mv2 dataset. Note
Transformer-M and Global-ViSNet are concurrent work.

1lightgray!20  Model # Param. Model Type Valid. MAE (meV) $\downarrow$ without
3D Positional Information GCN-virtual (Hu et al., 2021) 4.9M MPNN 115.3 GIN-
virtual (Hu et al., 2021) 6.7M MPNN 108.3 GRPE (Park et al., 2022) 46.2M
Transformer 89.0 Transformer-M (Medium, No 3D Positions) (Luo et al., 2022)
47.1M Transformer 87.8 EGT (Hussain et al., 2022) 89.3M Transformer 86.9
Graphormer (Shi et al., 2022) 48.3M Transformer 86.4 GPS (Rampášek et al.,
2022) 19.4M Hybrid 85.8 lightblue GPS++ (No 3D Positions) 44.3M Hybrid 82.0
lightblue GPS++ (No MHSA, No 3D Positions) 40.0M MPNN 81.0 with 3D Positional
Information Transformer-M (Medium, No Denoising) (Luo et al., 2022) 47.1M
Transformer 81.1 GEM-2 (Liu et al., 2022) 32.1M Transformer 79.3 Transformer-M
(Medium) (Luo et al., 2022) 47.1M Transformer 78.7 lightblue GPS++ (No MHSA)
40.0M MPNN 78.6 Global-ViSNet (Wang et al., 2022b) 78.5M Transformer 78.4
lightblue GPS++ 44.3M Hybrid 77.8 Transformer-M (Large) (Luo et al., 2022)
69.0M Transformer 77.2

#### Grouped Input Masking

The 3D positional features $\mathbf{R}$ are only defined for the training
data. We must therefore make use of them in training without requiring them
for validation/test. We found that the method proposed by Luo et al. (2022)
achieved the most favourable results so we adopt a variation hereon referred
to as _grouped input masking_.

Atom distances calculated from 3D positions are embedded into vector space
$\mathbb{R}^{K}$ via $K=128$ Gaussian kernel functions. We then process these
distance embeddings in three ways to produce attention biases
($\mathbf{B}^{\text{3D}}$), node features ($\mathbf{X}^{\text{3D}}$) and edge
features ($\mathbf{E}^{\text{3D}}$) (exact formulations can be found in
Appendix A.2).

This method stochastically masks out any features derived from the 3D
positional features $\mathbf{R}$ to build robustness to their absence.
Specifically, this is done by defining two input sets to be masked:

$\displaystyle\mathcal{X}^{\text{Spatial}}=\\{\mathbf{X}^{\text{3D}},\mathbf{E}^{\text{3D}},\mathbf{B}^{\text{3D}}\\},~{}~{}~{}~{}\mathcal{X}^{\text{Topological}}=\\{\mathbf{B}^{\mathrm{SPD}}\\}$

and three potential masking groups: 1. Mask $\mathcal{X}^{\text{Spatial}}$, 2.
Mask $\mathcal{X}^{\text{Topological}}$, and 3. No masking. These masking
groups are then sampled randomly throughout training with ratio 1:3:1. If 3D
positions are not defined, e.g. in validation/test, masking group 1 is always
used.

#### Training Time

Our GPS++ model trains at 17,500 graphs per second on 16 IPUs which means each
epoch completes in 3 minutes and a full 450 epoch training run takes 24 hours.
Full details of the hardware used and acceleration methods employed can be
found in Appendix B.

##  6 Results

#### Single Model Performance

In Table 2, we compare the single model performance of GPS++ with results from
the literature, in particular, we 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 method for incorporating 3D features;
this allows us to gain some tangible insights into the value of hybrid/MPNN
approaches vs. transformer.

Comparing the best GPS++ result with prior work, we see improvement over all
results apart from the largest Transformer-M that is only narrowly better,
however, we see that GPS++ overtakes the Transformer-M with comparable
parameters. Another interesting feature of the Transformer-M is that it
implements a 3D denoising loss, not present in GPS++ due to Transformer-M
being released towards the end of GPS++ development, which achieves a
significant performance boost. Without this feature, in a regime more similar
to ours, we see a significant performance reduction in Transformer-M relative
to GPS++.

While it would be tempting to assume that the attention element of our model
is the most significant due to the prevalence of transformers among the other
top results, we find that we can maintain a large proportion of the 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 global attention
mechanisms on 3D positional information and, in contrast, the power of using
message passing along molecular bonds. We further investigate the impact of
different elements of the model in §7.

#### Ensembled Model Performance

To test our GPS++ model on the official PCQM4Mv2 test set we build a large
ensemble to maximise the model accuracy. For this we focus on both the size of
ensemble as well as the diversity which are both well-known factors for
achieving good ensemble performance (Zhou et al., 2002; Lakshminarayanan et
al., 2017). Furthermore, as we are aiming to infer the model on the test-
challenge set we also include the validation set in training. The strategy
taken to build the ensemble is outlined in Appendix C. In the end we build an
ensemble of 112 models using 7 different configurations of GPS++ which
achieves a final test-challenge MAE of 71.9 meV that beats all prior work on
this dataset (see Table 1). Inference was completed in 1 hour 32 minutes using
1 IPU and AMD EPYC 7742 64-Core CPU, including all the feature processing and
program compilation, which is well under the 4 hour constraint typically
applied to this task.

##  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 for 200 epochs with
one or more features disabled, averaging MAE over 5 runs.

#### 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 3 isolates the contribution of each of these
additions (excluding 3D features).

Table 3: Ablation of node and edge features.

1lightgray!20  Removed Feature $\Delta$ MAE (meV) Valid Train white None
(Baseline) 78.4 52.3 Random Walk Structural Enc \+ 1.3 \+ 1.3 Laplacian
Positional Enc \- 0.1 \- 0.3 Local Centrality Enc 0.0 \+ 0.6 Bond Lengths 0.0
\- 0.4 Noisy Edges \+ 0.1 \- 0.2 Noisy Nodes \+ 0.5 \+ 0.3

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.3meV upon removal, whilst the Graph Laplacian Positional
Encodings (both eigenvectors and eigenvalues) are found to have no benefit.
Other features, despite being beneficial when they were added at intermediate
steps in the development of the model, are shown to have become similarly
redundant in the final GPS++ composition: the Graphormer-style Local
Centrality Encoding (i.e. embedding the degree of the node) carries no
information that the Random Walks do not, and the use of Bond Lengths as an
edge feature in the MPNN appears to be unnecessary in the presence of more
comprehensive 3D features elsewhere in the network. Though both features are
equally redundant for task performance, the "2D" graph Centrality Encoding is
contributing much more to overfitting than the "3D" Bond Length, a pattern
that is evident elsewhere in the ablations. Table 3 also shows that whilst the
regularisation effect of the noisy nodes loss is beneficial, the contribution
of noisy edges is negligible. This may be due to the simplicity (and hence
ease of denoising) of the limited edge features in PCQM4Mv2.

Table 4 compares the possible choices of chemical input features described in
§4.2: Original is the set of features included in PCQM4Mv2; Ying21 refers to
all 28 node features and 5 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.1 and
obtained via a combinatorial search; Atom Group, Period & Type refers to our
contribution of 3 additional atom features relating to an atom’s position in
the periodic table, intended to be more generalised than the atomic number.
The results clearly show that training on all 36 available input features
causes significant overfitting and degrades task performance by introducing
noise. Moreover, the original PCQM4Mv2 features are not an optimal subset; the
significant difference between Original and Set 1 is largely due to the
removal of the original Chiral Tag feature, which has been replaced by a
Chiral Center boolean. Lastly, the table shows that Atom Group, Period & Type
provide meaningfully richer information than the atomic number feature which
is present in every feature set, but are not immune to exacerbating
overfitting in cases where too many features are used.

Table 4: Choices of chemical features.

1lightgray!20  Atom Group, $\Delta$ MAE (meV) Feature Set Period & Type Set
Size Valid Train white Set 1 (Baseline) $\checkmark$ 14 78.4 52.3 Set 2
$\checkmark$ 14 \- 0.1 \+ 0.0 Set 3 $\checkmark$ 14 \- 0.1 \- 0.2 Set 1
$\cdot$ 11 \+ 0.5 \- 0.4 Original $\checkmark$ 15 \+ 1.2 \+ 0.1 Original
$\cdot$ 12 \+ 1.7 \+ 0.5 Ying21 $\checkmark$ 36 \+ 5.4 \- 5.2 Ying21 $\cdot$
33 \+ 5.0 \- 4.4

#### 3D Input Features and MHSA

The 3D node position information given for the training set in PCQM4Mv2 is
provided as input to the final model in two forms: the all-to-all 3D Bias map
added to the self-attention module, and the node-wise 3D Centrality Encoding
added in the encoder. The first four rows of Table 7 show that they both
improve task performance independently and in combination, but to different
degrees. Whilst the 3D Bias is meaningfully beneficial to the valid MAE, the
3D Centrality has a greater impact and is also responsible for a significant
reduction in overfitting visible in the training MAE, suggesting improved
generalisation capabilities. Additionally, we found that training stability of
the MHSA module suffered in the absence of either 3D feature, requiring us to
halve the learning rate to 2e-4 for convergence unless the MHSA module was
disabled.

Table 5: Ablation of self-attention and 3D features.

1lightgray!20  Feature $\Delta$ MAE (meV) white# MHSA SPD Bias 3D Bias 3D
Centrality Valid Train 1 $\checkmark$ $\checkmark$ $\checkmark$ $\checkmark$
78.4 52.3 2 $\checkmark$ $\checkmark$ $\checkmark$ $\cdot$ \+ 2.0 -8.4 $*$ 3
$\checkmark$ $\checkmark$ $\cdot$ $\checkmark$ \+ 0.9 -2.5 $*$ 4 $\checkmark$
$\checkmark$ $\cdot$ $\cdot$ \+ 3.6 -8.9 $*$ 5 $\checkmark$ $\cdot$ $\cdot$
$\cdot$ \+ 3.5 -8.2 $*$ 6 $\cdot$ $\cdot$ $\cdot$ $\cdot$ \+ 2.5 -9.9 7
$\checkmark$ $\cdot$ $\cdot$ $\checkmark$ \+ 0.9 -1.8 $*$ 8 $\cdot$ $\cdot$
$\cdot$ $\checkmark$ \+ 0.9 -0.5 9 $\checkmark$ $\cdot$ $\checkmark$
$\checkmark$ \+ 0.1 +0.3 white $*$ trained with half learning rate for
stability

Interestingly, in the absence of both 3D features, Table 7 rows 4-6 show that
it is preferable to remove the MHSA. Even with the 3D Centrality feature added
to nodes in rows 7-8, a "standard" unbiased MHSA is not beneficial to task
performance and increases overfitting. This implies that the Shortest Path
Distance Bias term, which attempts to provide a non-3D global graph positional
encoding to the MHSA, is insufficient to modulate attention between atoms, and
that instead spatial relationships are key for this quantum chemistry task.
Indeed, in rows 4 vs 5, 3 vs 7 and 1 vs 9 the SPD Bias is shown to increase
overfitting whilst having a negligible valid MAE impact. This result stands
out against previous work where the SPD Bias was an effective component of
state-of-the-art graph transformers on PCQM4Mv2; one explanation may be that
in the absence of an MPNN module graph transformers rely on the SPD Bias to
identify neighbourhoods and recreate some of the functionality of MPNNs within
the MHSA layer, as described in section 3.3 of Ying et al. (2021a).

#### Model Architecture

Table 6: Ablation of network architecture features.

1lightgray!20  $\Delta$ MAE (meV) Removed Feature Valid Train $\Delta$ Param.
white None (Baseline) 78.4 52.3 44.3M MPNN \+ 13.3 \+ 21.1 -69.8% Edge
Features \+ 4.4 \+ 3.7 -26.3% Global Features \+ 1.0 \+ 0.4 \- 8.0% Sender
Message Aggregation \+ 0.7 \- 0.4 -14.2% Adjacent Node Aggregation \+ 0.4 \+
2.1 -18.9% MHSA \+ 0.9 \- 0.5 \- 9.5% FFN \+ 0.8 \+ 0.9 -19.0% MHSA and Global
Features \+ 3.5 \+ 0.2 -17.5%

GPS++ uses a large MPNN module comprising 70% of the model parameters and
using extensively tuned hyper-parameters. Table 6 ablates the components of
the model network architecture, 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 on training MAE. Within the MPNN, the largest task impact
arises from removing the edge features and edge MLP from all layers
($\mathbf{e}_{uv}$ and $\texttt{MLP}_{\textrm{edge}}$ in Equation 6), falling
back on simple GCN-style message passing that concatenates neighbouring node
features when computing messages. Table 6 also shows that the use of Adjacent
Node Feature Aggregation (i.e., the sum of adjacent node features
$\mathbf{x}_{u}^{\ell}$ 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 is 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
are shown to be of comparable importance to the MHSA and FFN modules. Global
Features, denoted as $\mathbf{g}^{\ell}$ in Equation 6 and 7, allow the MPNN
some degree of global reasoning, whilst Sender Message Aggregation (the sum of
$\bar{\mathbf{e}}_{iv}^{\ell}$ 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 to partially
compensate for the loss of the other in the individual ablations.

##  8 Discussion

In this work, we define GPS++, a hybrid MPNN/ Transformer, optimised for the
PCQM4Mv2 molecular property prediction task (Hu et al., 2021). Our model
builds on previous work (Rampášek et al., 2022; Luo et al., 2022; Godwin et
al., 2022) with a particular focus on building a powerful and expressive
message passing component. Despite strong trends to use transformers for
molecular property prediction, we showed that our GPS++ model is very
competitive with prior work achieving the best results when comparing similar
sized models, only marginally worse than the best larger transformer model. To
our surprise, we also found that without any global attention mechanism at
all, GPS++ (now a pure MPNN) retains almost all of its performance. Finally,
we consider the case where 3D positions are not available and show that our
GPS++ models are significantly better than Transformers, where a significant
drop-off is observed across all prior work. We also found that our MPNN-only
model performs better than our hybrid under these circumstances, this
indicates a strong reliance between effective attention and the availability
of 3D positional information.

Finally, we built a large diverse ensemble of GPS++ models and achieved an MAE
of 71.9 meV on the PCQM4Mv2 test-challenge set with four-hour inference time
budget, the state-of-the-art at time of submission.

Whilst these results are interesting for problems with small molecules,
testing on 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 to be
determined if the downsides of issues like underreaching outweigh these
benefits. Nevertheless, we believe that our results highlight the strength of
MPNNs in this domain and hope that it inspires a revival of message passing
and hybrid models for molecular property prediction when dealing with large
datasets.

## References

  * Alon & Yahav (2021) Alon, U. and Yahav, E.  On the bottleneck of graph neural networks and its practical implications.  In _International Conference on Learning Representations_ , 2021.
  * Axelrod & Gomez-Bombarelli (2022) Axelrod, S. and Gomez-Bombarelli, R.  Geom, energy-annotated molecular conformations for property prediction and molecular generation.  _Scientific Data_ , 9(1):1–14, 2022.
  * Ba et al. (2016) Ba, J. L., Kiros, J. R., and Hinton, G. E.  Layer normalization.  _arXiv:1607.06450_ , 2016.
  * Battaglia et al. (2018) Battaglia, P. W., Hamrick, J. B., Bapst, V., Sanchez-Gonzalez, A., Zambaldi, V., Malinowski, M., Tacchetti, A., Raposo, D., Santoro, A., Faulkner, R., et al.  Relational inductive biases, deep learning, and graph networks.  _arXiv:1806.01261_ , 2018.
  * Bilbrey et al. (2022) Bilbrey, J. A., Herman, K. M., Sprueill, H., Xantheas, S. S., Das, P., Roldan, M. L., Kraus, M., Helal, H., and Choudhury, S.  Reducing down(stream)time: Pretraining molecular gnns using heterogeneous ai accelerators.  _arXiv preprint arXiv:2211.04598_ , 2022.
  * Bronstein et al. (2021) Bronstein, M. M., Bruna, J., Cohen, T., and Veličković, P.  Geometric deep learning: Grids, groups, graphs, geodesics, and gauges.  _arXiv:2104.13478_ , 2021.
  * Chmiela et al. (2017) Chmiela, S., Tkatchenko, A., Sauceda, H. E., Poltavsky, I., Schütt, K. T., and Müller, K.-R.  Machine learning of accurate energy-conserving molecular force fields.  _Science advances_ , 3(5):e1603015, 2017.
  * Darabi et al. (2022) Darabi, S., Fazeli, S., Liu, J., Milesi, A., Morkisz, P., Puget, J.-F., and Titericz, G.  Heterogeneous ensemble of models for molecular property prediction.  https://https://ogb.stanford.edu/paper/neurips2022/pcqm4mv2_NVIDIA-PCQM4Mv2.pdf, 2022\.  Accessed: 2023-01-24.
  * Dobson (2004) Dobson, C. M.  Chemical space and biology.  _Nature_ , 432(7019):824–828, December 2004.  ISSN 1476-4687.  doi:10.1038/nature03192.
  * Duvenaud et al. (2015) Duvenaud, D. K., Maclaurin, D., Iparraguirre, J., Bombarell, R., Hirzel, T., Aspuru-Guzik, A., and Adams, R. P.  Convolutional networks on graphs for learning molecular fingerprints.  _Advances in neural information processing systems_ , 28, 2015.
  * Dwivedi & Bresson (2020) Dwivedi, V. P. and Bresson, X.  A generalization of transformer networks to graphs.  _arXiv:2012.09699_ , 2020.
  * Dwivedi et al. (2022) Dwivedi, V. P., Luu, A. T., Laurent, T., Bengio, Y., and Bresson, X.  Graph neural networks with learnable structural and positional representations.  In _International Conference on Learning Representations_ , 2022.
  * Fey & Lenssen (2019) Fey, M. and Lenssen, J. E.  Fast graph representation learning with PyTorch Geometric.  In _ICLR Workshop on Representation Learning on Graphs and Manifolds_ , 2019.
  * Gilmer et al. (2017) Gilmer, J., Schoenholz, S. S., Riley, P. F., Vinyals, O., and Dahl, G. E.  Neural message passing for quantum chemistry.  In _International conference on machine learning_ , pp. 1263–1272. PMLR, 2017.
  * Godwin et al. (2022) Godwin, J., Schaarschmidt, M., Gaunt, A. L., Sanchez-Gonzalez, A., Rubanova, Y., Veličković, P., Kirkpatrick, J., and Battaglia, P.  Simple GNN regularisation for 3D molecular property prediction and beyond.  In _International Conference on Learning Representations_ , 2022.
  * Hendrycks & Gimpel (2016) Hendrycks, D. and Gimpel, K.  Gaussian error linear units (GELUs).  _arXiv:1606.08415_ , 2016.
  * Hu et al. (2020) Hu, W., Fey, M., Zitnik, M., Dong, Y., Ren, H., Liu, B., Catasta, M., and Leskovec, J.  Open Graph Benchmark: Datasets for Machine Learning on Graphs.  _34th Conference on Neural Information Processing Systems_ , 2020\.
  * Hu et al. (2021) Hu, W., Fey, M., Ren, H., Nakata, M., Dong, Y., and Leskovec, J.  OGB-LSC: A large-scale challenge for machine learning on graphs.  In _35th Conference on Neural Information Processing Systems: Datasets and Benchmarks Track_ , 2021.
  * Huang et al. (2016) Huang, G., Sun, Y., Liu, Z., Sedra, D., and Weinberger, K. Q.  Deep networks with stochastic depth.  In Leibe, B., Matas, J., Sebe, N., and Welling, M. (eds.), _Computer Vision – ECCV 2016_ , pp. 646–661, Cham, 2016. Springer International Publishing.  ISBN 978-3-319-46493-0.
  * Huang et al. (2022) Huang, K., Fu, T., Gao, W., Zhao, Y., Roohani, Y., Leskovec, J., Coley, C. W., Xiao, C., Sun, J., and Zitnik, M.  Artificial intelligence foundation for therapeutic science.  _Nature Chemical Biology_ , 18(10):1033–1036, 2022.
  * Huang et al. (2018) Huang, Y., Cheng, Y., Bapna, A., Firat, O., Chen, M. X., Chen, D., Lee, H., Ngiam, J., Le, Q. V., Wu, Y., and Chen, Z.  Gpipe: Efficient training of giant neural networks using pipeline parallelism.  _arXiv preprint arXiv:1811.06965_ , 2018.
  * Hussain et al. (2022) Hussain, M. S., Zaki, M. J., and Subramanian, D.  Global self-attention as a replacement for graph convolution.  In _Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining_ , pp. 655–665, 2022.
  * Keith et al. (2021) Keith, J. A., Vassilev-Galindo, V., Cheng, B., Chmiela, S., Gastegger, M., Müller, K.-R., and Tkatchenko, A.  Combining machine learning and computational chemistry for predictive insights into chemical systems.  _Chemical Reviews_ , 121(16):9816–9872, 2021\.  doi:10.1021/acs.chemrev.1c00107.  PMID: 34232033.
  * Kim et al. (2022) Kim, J., Nguyen, T. D., Min, S., Cho, S., Lee, M., Lee, H., and Hong, S.  Pure transformers are powerful graph learners.  In _NeurIPS_ , 2022.
  * Kingma & Ba (2015) Kingma, D. P. and Ba, J.  Adam: A method for stochastic optimization.  In _ICLR (Poster)_ , 2015.  URL http://arxiv.org/abs/1412.6980.
  * Kohn & Sham (1965) Kohn, W. and Sham, L. J.  Self-consistent equations including exchange and correlation effects.  _Phys. Rev._ , 140:A1133–A1138, Nov 1965.
  * Krell et al. (2021) Krell, M. M., Kosec, M., Perez, S. P., and Fitzgibbon, A.  Efficient sequence packing without cross-contamination: Accelerating large language models without impacting performance.  _arXiv preprint arXiv:2107.02027_ , 2021.
  * Kreuzer et al. (2021) Kreuzer, D., Beaini, D., Hamilton, W. L., Létourneau, V., and Tossou, P.  Rethinking graph transformers with spectral attention.  In _Advances in Neural Information Processing Systems_ , 2021.
  * Lakshminarayanan et al. (2017) Lakshminarayanan, B., Pritzel, A., and Blundell, C.  Simple and scalable predictive uncertainty estimation using deep ensembles.  _Advances in neural information processing systems_ , 30, 2017.
  * Li et al. (2018) Li, Q., Han, Z., and Wu, X.  Deeper insights into graph convolutional networks for semi-supervised learning.  In _AAAI_ , pp. 3538–3545, 2018.
  * Liu et al. (2022) Liu, L., He, D., Fang, X., Zhang, S., Wang, F., He, J., and Wu, H.  GEM-2: Next generation molecular property prediction network with many-body and full-range interaction modeling.  _arXiv preprint arXiv:2208.05863_ , 2022.
  * Luo et al. (2022) Luo, S., Chen, T., Xu, Y., Zheng, S., Liu, T.-Y., Wang, L., and He, D.  One transformer can understand both 2D & 3D molecular data.  _arXiv:2210.01765_ , 2022.
  * Micikevicius et al. (2017) Micikevicius, P., Narang, S., Alben, J., Diamos, G., Elsen, E., Garcia, D., Ginsburg, B., Houston, M., Kuchaiev, O., Venkatesh, G., and Wu, H.  Mixed precision training.  _arXiv preprint arXiv:1710.03740_ , 2017.
  * Nakata & Shimazaki (2017) Nakata, M. and Shimazaki, T.  Pubchemqc project: a large-scale first-principles electronic structure database for data-driven chemistry.  _Journal of chemical information and modeling_ , 57(6):1300–1308, 2017.
  * Park et al. (2022) Park, W., Chang, W., Lee, D., Kim, J., and won Hwang, S.  GRPE: Relative positional encoding for graph transformer.  _arXiv:22201.12787_ , 2022.
  * Ramakrishnan et al. (2014) Ramakrishnan, R., Dral, P. O., Rupp, M., and Von Lilienfeld, O. A.  Quantum chemistry structures and properties of 134 kilo molecules.  _Scientific data_ , 1(1):1–7, 2014.
  * Rampášek et al. (2022) Rampášek, L., Galkin, M., Dwivedi, V. P., Luu, A. T., Wolf, G., and Beaini, D.  Recipe for a General, Powerful, Scalable Graph Transformer.  _arXiv:2205.12454_ , 2022.
  * (38) RDKit.  RDKit: Open-source cheminformatics.  http://www.rdkit.org.  [Online].
  * Reiser et al. (2022) Reiser, P., Neubert, M., Eberhard, A., Torresi, L., Zhou, C., Shao, C., Metni, H., van Hoesel, C., Schopmans, H., Sommer, T., et al.  Graph neural networks for materials science and chemistry.  _Communications Materials_ , 3(1):93, 2022.
  * Rogers & Hahn (2010) Rogers, D. and Hahn, M.  Extended-connectivity fingerprints.  _Journal of chemical information and modeling_ , 50(5):742–754, 2010.
  * Shi et al. (2022) Shi, Y., Zheng, S., Ke, G., Shen, Y., You, J., He, J., Luo, S., Liu, C., He, D., and Liu, T.-Y.  Benchmarking graphormer on large-scale molecular modeling datasets.  _arXiv:2203.04810_ , 2022.
  * Smith et al. (2017) Smith, J. S., Isayev, O., and Roitberg, A. E.  Ani-1, a data set of 20 million calculated off-equilibrium conformations for organic molecules.  _Scientific data_ , 4(1):1–8, 2017.
  * Srivastava et al. (2014) Srivastava, N., Hinton, G., Krizhevsky, A., Sutskever, I., and Salakhutdinov, R.  Dropout: a simple way to prevent neural networks from overfitting.  _The journal of machine learning research_ , 15(1):1929–1958, 2014.
  * Tran et al. (2022) Tran, R., Lan, J., Shuaibi, M., Goyal, S., Wood, B. M., Das, A., Heras-Domingo, J., Kolluru, A., Rizvi, A., Shoghi, N., et al.  The open catalyst 2022 (oc22) dataset and challenges for oxide electrocatalysis.  _arXiv:2206.08917_ , 2022.
  * Vaswani et al. (2017) Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I.  Attention is all you need.  _Advances in Neural Information Processing Systems_ , 30, 2017.
  * Wang et al. (2022a) Wang, H., Yin, H., Zhang, M., and Li, P.  Equivariant and stable positional encoding for more powerful graph neural networks.  In _International Conference on Learning Representations_ , 2022a.
  * Wang et al. (2022b) Wang, Y., Li, S., Wang, T., Wang, Z., He, X., Shao, B., and Liu, T.-Y.  How to better introduce geometric information in equivariant message passing?  https://github.com/ogb-visnet/Global-ViSNet/blob/master/ViSNet_Tech_Report.pdf, 2022b.  Accessed: 2022-11-16.
  * Wang et al. (2022c) Wang, Y., Li, S., Wang, T., Wang, Z., He, X., Shao, B., and Liu, T.-Y.  An ensemble of VisNet, Transformer-M, and pretraining models for molecular property prediction in OGB Large-Scale Challenge @ NeurIPS 2022, 2022c.
  * Wu et al. (2018) Wu, Z., Ramsundar, B., Feinberg, E. N., Gomes, J., Geniesse, C., Pappu, A. S., Leswing, K., and Pande, V.  Moleculenet: a benchmark for molecular machine learning.  _Chemical science_ , 9(2):513–530, 2018.
  * Ying et al. (2021a) Ying, C., Cai, T., Luo, S., Zheng, S., Ke, G., He, D., Shen, Y., and Liu, T.-Y.  Do transformers really perform badly for graph representation?  In _Advances in Neural Information Processing Systems_ , 2021a.
  * Ying et al. (2021b) Ying, C., Yang, M., Zheng, S., Ke, G., Luo, S., Cai, T., Wu, C., Wang, Y., Shen, Y., and He, D.  First place solution of KDD Cup 2021 & OGB large-scale challenge graph prediction track.  _arXiv:2106.08279_ , 2021b.
  * Zaidi et al. (2022) Zaidi, S., Schaarschmidt, M., Martens, J., Kim, H., Teh, Y. W., Sanchez-Gonzalez, A., Battaglia, P., Pascanu, R., and Godwin, J.  Pre-training via denoising for molecular property prediction.  _arXiv preprint arXiv:2206.00133_ , 2022.
  * Zhou et al. (2002) Zhou, Z.-H., Wu, J., and Tang, W.  Ensembling neural networks: many could be better than all.  _Artificial intelligence_ , 137(1-2):239–263, 2002.

##  Appendix A Detailed Model Description

###  A.1 GPS++ Block Complete

The GPS++ block is defined as follows for layers $\ell>0$ (see §A.2 for the
definitions of $\mathbf{X}^{0},\mathbf{E}^{0},\mathbf{g}^{0}$).

$\displaystyle\mathbf{X}^{\ell+1},\mathbf{E}^{\ell+1},\mathbf{g}^{\ell+1}$
$\displaystyle=\texttt{GPS++}\left(\mathbf{X}^{\ell},\mathbf{E}^{\ell},\mathbf{g}^{\ell},\mathbf{B}\right)$
(17) computed as $\displaystyle\mathbf{Y}^{\ell},\ \mathbf{E}^{\ell+1},\
\mathbf{g}^{\ell+1}$
$\displaystyle=\texttt{MPNN}\left(\mathbf{X}^{\ell},\mathbf{E}^{\ell},\mathbf{g}^{\ell}\right),$
(18) $\displaystyle\mathbf{Z}^{\ell}$
$\displaystyle=\texttt{BiasedAttn}\left(\mathbf{X}^{\ell},\mathbf{B}\right),$
(19) $\displaystyle\forall i:~{}~{}~{}\mathbf{x}_{i}^{\ell+1}$
$\displaystyle=\texttt{FFN}\left(\mathbf{y}_{i}^{\ell}+\mathbf{z}_{i}^{\ell}\right)$
(20)

#### The MPNN module

A simplified version of the MPNN module is presented in §4.1. The full
description is as follows:

$\displaystyle\mathbf{Y}^{\ell},\ \mathbf{E}^{\ell+1},\
\mathbf{g}^{\ell+1}=\texttt{MPNN}\left(\mathbf{X}^{\ell},\mathbf{E}^{\ell},\mathbf{g}^{\ell}\right),$
(21) computed as $\displaystyle\forall(u,v):~{}~{}~{}$
$\displaystyle{\mathbf{c}}_{uv}^{\ell}=\bigl{[}\mathbf{x}_{u}^{\ell}\mid\mathbf{x}_{v}^{\ell}\mid\mathbf{e}_{uv}^{\ell}\mid\mathbf{g}^{\ell}\bigr{]}$
(22) $\displaystyle\forall(u,v):~{}~{}~{}$
$\displaystyle\bar{\mathbf{e}}_{uv}^{\ell}=\texttt{Dropout}_{0.0035}\left(\texttt{MLP}_{\textrm{edge}}\left({\mathbf{c}}_{uv}^{\ell}\right)\right)$
(23) $\displaystyle\forall i:~{}~{}~{}$
$\displaystyle{\mathbf{c}}_{i}^{\ell}=\Biggl{[}\mathbf{x}_{i}^{\ell}~{}\Bigg{|}\underbrace{\sum_{(u,i)\in\mathcal{E}}\bar{\mathbf{e}}_{ui}^{\ell}}_{\begin{subarray}{c}\text{sender}\\\
\text{messages}\end{subarray}}~{}\Bigg{|}\underbrace{\sum_{(i,v)\in\mathcal{E}}\bar{\mathbf{e}}_{iv}^{\ell}}_{\begin{subarray}{c}\text{receiver}\\\
\text{messages}\end{subarray}}~{}\Bigg{|}\underbrace{\sum_{(u,i)\in\mathcal{E}}\mathbf{x}_{u}^{\ell}}_{\begin{subarray}{c}\text{adjacent}\\\
\text{nodes}\end{subarray}}~{}\Bigg{|}~{}\mathbf{g}^{\ell}\Biggr{]}$ (24)
$\displaystyle\forall i:~{}~{}~{}$
$\displaystyle\bar{\mathbf{x}}_{i}^{\ell}=\texttt{MLP}_{\textrm{node}}(\mathbf{c}_{i}^{\ell})$
(25)
$\displaystyle\bar{\mathbf{g}}^{\ell}=\texttt{MLP}_{\textrm{global}}\left(\Biggl{[}{\mathbf{g}^{\ell}~{}\Bigg{|}\sum_{j\in\mathcal{V}}\bar{\mathbf{x}}_{j}^{\ell}~{}\Bigg{|}\sum_{(u,v)\in\mathcal{E}}\bar{\mathbf{e}}_{uv}^{\ell}}\Biggr{]}\right)$
(26) $\displaystyle\forall i:~{}~{}~{}$
$\displaystyle\mathbf{y}_{i}^{\ell}=\texttt{LayerNorm}(\texttt{Dropout}_{0.3}(\bar{\mathbf{x}}_{i}^{\ell})+\mathbf{x}_{i}^{\ell})$
(27) $\displaystyle\forall(u,v):~{}~{}~{}$
$\displaystyle\mathbf{e}_{uv}^{\ell+1}=\bar{\mathbf{e}}_{uv}^{\ell}+\mathbf{e}_{uv}^{\ell}$
(28)
$\displaystyle\mathbf{g}^{\ell+1}=\texttt{Dropout}_{0.35}(\bar{\mathbf{g}}^{\ell})+\mathbf{g}^{\ell},$
(29)

where $\texttt{Dropout}_{p}$ (Srivastava et al., 2014) masks by zero each
element with probability $p$ and LayerNorm follows the normalisation procedure
of Ba et al. (2016). The three networks $\texttt{MLP}_{\eta}$ for
$\eta\in\\{\textrm{node},\textrm{edge},\textrm{global}\\}$ each have two
layers and are defined by:

$\displaystyle\mathbf{y}$ $\displaystyle=\texttt{MLP}_{\eta}(\mathbf{x})$ (30)
$\displaystyle\textrm{computed as}\ \ \ \bar{\mathbf{x}}$
$\displaystyle=\texttt{GELU}(\texttt{Dense}(\mathbf{x}))$
$\displaystyle\in\mathbb{R}^{4d_{\eta}}$ (31) $\displaystyle\mathbf{y}$
$\displaystyle=\texttt{Dense}(\texttt{LayerNorm}(\bar{\mathbf{x}}))$
$\displaystyle\in\mathbb{R}^{d_{\eta}}\ $ (32)

where GELU is from Hendrycks & 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 a 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
§A.2. Single-head biased attention is defined by:

$\displaystyle\mathbf{Z}$
$\displaystyle=\texttt{BiasedAttn}(\mathbf{X},\mathbf{B})$ (33) computed as
$\displaystyle\mathbf{A}$
$\displaystyle=\frac{\left(\mathbf{X}\mathbf{W}_{Q}\right)\left(\mathbf{X}\mathbf{W}_{K}\right)^{\top}}{\sqrt{{d_{\text{node}}}}}+\mathbf{B}$
$\displaystyle\in\mathbb{R}^{N\times N}\ \ \ $ (34)
$\displaystyle\bar{\mathbf{A}}$
$\displaystyle=\textrm{{Dropout}}_{0.3}\left(\texttt{Softmax}\left(\mathbf{A}\right)\right)\left(\mathbf{X}\mathbf{W}_{V}\right)$
$\displaystyle\in\mathbb{R}^{N\times{d_{\text{node}}}}$ (35)
$\displaystyle\mathbf{Z}$
$\displaystyle=\texttt{GraphDropout}_{\frac{\ell}{L}0.3}\left(\bar{\mathbf{A}}\right)+\mathbf{X}$
$\displaystyle\in\mathbb{R}^{N\times{d_{\text{node}}}}$ (36)

for learnable weight matrices
$\mathbf{W}_{Q},\mathbf{W}_{K},\mathbf{W}_{V}\in\mathbb{R}^{{d_{\text{node}}}\times{d_{\text{node}}}}$
and output in $\mathbb{R}^{N\times{d_{\text{node}}}}$, though in practice we
use 32 attention heads which are mixed before GraphDropout using an affine
projection
$\mathbf{W}_{P}\in\mathbb{R}^{{d_{\text{node}}}\times{d_{\text{node}}}}$.

#### The FFN module

Finally, the feed-forward network module takes the form:

$\displaystyle\mathbf{y}$ $\displaystyle=\texttt{FFN}(\mathbf{x})$ (37)
computed as $\displaystyle\bar{\mathbf{x}}$
$\displaystyle=\texttt{Dropout}_{p}(\texttt{GELU}(\texttt{Dense}(\mathbf{x}))\qquad\quad\
\ \in\mathbb{R}^{4{d_{\text{node}}}}$ (38) $\displaystyle\mathbf{y}$
$\displaystyle=\texttt{GraphDropout}_{\frac{\ell}{L}0.3}(\texttt{Dense}(\bar{\mathbf{x}}))+\mathbf{x}\
\in\mathbb{R}^{{d_{\text{node}}}}$ (39)

Unless otherwise stated, the dropout probability $p=0$, however, we experiment
with other values when ensembling multiple model variants. 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 §3, the dataset samples include the graph structure
$\mathcal{G}$, a set of categorical features for the atoms and bonds
$\mathbbm{x}_{i}$, $\mathbbm{e}_{uv}$, and the 3D node positions
$\mathbf{r}_{i}$. It has been shown that there are many benefits to augmenting
the input data with additional structural, positional, and chemical
information (Rampášek 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.

$\displaystyle\mathbf{X}^{all}=\bigl{[}\mathbf{X}^{\text{atom}}\mid\mathbf{X}^{\text{LapVec}}\mid\mathbf{X}^{\text{LapVal}}\mid\mathbf{X}^{\text{RW}}\mid\mathbf{X}^{\text{Cent}}\mid\mathbf{X}^{\text{3D}}\bigr{]}$
$\displaystyle\mathbf{X}^{0}$
$\displaystyle=\texttt{Dense}(\mathbf{X}^{all})\qquad\qquad\in\mathbb{R}^{N\times{d_{\text{node}}}}$
(40) $\displaystyle\mathbf{E}^{0}$
$\displaystyle=\texttt{Dense}(\bigl{[}\mathbf{E}^{\text{bond}}\mid\mathbf{E}^{\text{3D}}\bigr{]})\
\ \in\mathbb{R}^{M\times{d_{\text{edge}}}}$ (41) $\displaystyle\mathbf{g}^{0}$
$\displaystyle=\texttt{Embed}_{{d_{\text{global}}}}(0)\qquad\quad\ \ \
\in\mathbb{R}^{{d_{\text{global}}}}$ (42) $\displaystyle\mathbf{B}$
$\displaystyle=\mathbf{B}^{\mathrm{SPD}}+\mathbf{B}^{\text{3D}}\qquad\qquad\in\mathbb{R}^{N\times
N}$ (43)

The various components of each of these equations are defined over the
remainder of this section. The encoding of these features also makes recurring
use of the following two generic functions. Firstly, a two-layer
$\texttt{MLP}_{\mathrm{encoder}}$ that projects features to a fixed-size
latent space:

$\displaystyle y$
$\displaystyle=\texttt{MLP}_{\textrm{encoder}}(x),\quad\text{where}\,x\in\mathbb{R}^{h}$
(44) computed as $\displaystyle\bar{x}$
$\displaystyle=\texttt{ReLU}(\texttt{Dense}(\texttt{LayerNorm}(x)))$
$\displaystyle\in\mathbb{R}^{2h}$ (45) $\displaystyle y$
$\displaystyle=\texttt{Dropout}_{0.18}(\texttt{Dense}(\texttt{LayerNorm}(\bar{x})))$
$\displaystyle\in\mathbb{R}^{32}$ (46)

Secondly, a function $\texttt{Embed}_{d}(j)\in\mathbb{R}^{d}$ which selects
the $j^{\mathrm{th}}$ 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 $\mathbbm{x}_{i}$,
$\mathbbm{e}_{uv}$ into a continuous vector space, we learn a simple embedding
vector for each category, sum the embeddings for all categories, and then
process it with an MLP to produce $\mathbf{X}^{\text{atom}}$ and
$\mathbf{E}^{\text{bond}}$, i.e.

$\displaystyle\forall i:~{}~{}~{}\bar{\mathbf{x}}_{i}^{\text{atom}}$
$\displaystyle=\texttt{MLP}_{\textrm{node}}\left(\sum\nolimits_{j\in\mathbbm{x}_{i}}\texttt{Embed}_{64}(j)\right)$
(47) $\displaystyle\forall i:~{}~{}~{}\mathbf{x}_{i}^{\text{atom}}$
$\displaystyle=\texttt{Dropout}_{0.18}(\bar{\mathbf{x}}_{i}^{\text{atom}})\in\mathbb{R}^{{d_{\text{node}}}}$
(48) $\displaystyle\forall(u,v):~{}~{}~{}\bar{\mathbf{e}}_{uv}^{\text{bond}}$
$\displaystyle=\texttt{MLP}_{\textrm{edge}}\left(\sum\nolimits_{j\in\mathbbm{e}_{uv}}\texttt{Embed}_{64}(j)\right)$
(49) $\displaystyle\forall(u,v):~{}~{}~{}\mathbf{e}_{uv}^{\text{bond}}$
$\displaystyle=\texttt{Dropout}_{0.18}(\bar{\mathbf{e}}_{uv}^{\text{bond}})\in\mathbb{R}^{{d_{\text{edge}}}}$
(50)

Here $\texttt{MLP}_{\textrm{node}}$ and $\texttt{MLP}_{\textrm{edge}}$ refer
to the functions by the same names used in Eq. 23 and 25 in the MPNN module,
yet parameterised independently.

Table A.1: Chemical input feature selection for PCQM4Mv2.

1lightgray!20  Feature Set Node features Original Set 1 Set 2 Set 3 Atomic
number $\checkmark$ $\checkmark$ $\checkmark$ $\checkmark$ Group $\cdot$
$\checkmark$ $\checkmark$ $\checkmark$ Period $\cdot$ $\checkmark$
$\checkmark$ $\checkmark$ Element type $\cdot$ $\checkmark$ $\checkmark$
$\checkmark$ Chiral tag $\checkmark$ $\cdot$ $\cdot$ $\cdot$ Degree
$\checkmark$ $\checkmark$ $\checkmark$ $\checkmark$ Formal charge $\checkmark$
$\checkmark$ $\cdot$ $\checkmark$ # Hydrogens $\checkmark$ $\checkmark$
$\checkmark$ $\checkmark$ # Radical electrons $\checkmark$ $\checkmark$
$\checkmark$ $\checkmark$ Hybridisation $\checkmark$ $\cdot$ $\checkmark$
$\checkmark$ Is aromatic $\checkmark$ $\checkmark$ $\checkmark$ $\cdot$ Is in
ring $\checkmark$ $\checkmark$ $\checkmark$ $\checkmark$ Is chiral center
$\cdot$ $\checkmark$ $\checkmark$ $\checkmark$ 0white Edge features Bond type
$\checkmark$ $\checkmark$ $\cdot$ $\cdot$ Bond stereo $\checkmark$
$\checkmark$ $\checkmark$ $\checkmark$ Is conjugated $\checkmark$ $\cdot$
$\checkmark$ $\checkmark$ Is in ring $\cdot$ $\checkmark$ $\checkmark$
$\checkmark$

#### Graph Laplacian Positional Encodings (Kreuzer et al., 2021; Dwivedi &
Bresson, 2020)

Given a graph with adjacency matrix $\mathbf{A}$ and degree matrix
$\mathbf{D}$, the eigendecomposition of the graph Laplacian $\mathbf{L}$ is
formulated into a global positional encoding as follows.

$\displaystyle\forall i:~{}~{}~{}\mathbf{x}_{i}^{\mathrm{LapVec}}$
$\displaystyle=\texttt{MLP}_{\textrm{encoder}}(\mathbf{U}\left[i,\,2\dots
k^{\mathrm{Lap}}\right])\in\mathbb{R}^{32},$ $\displaystyle\mathrm{where}\
\mathbf{L}=\mathbf{D}-\mathbf{A}=\mathbf{U}^{\top}\mathbf{\Lambda}\mathbf{U}$
(51) $\displaystyle\forall i:~{}~{}~{}\mathbf{x}_{i}^{\mathrm{LapVal}}$
$\displaystyle=\texttt{MLP}_{\textrm{encoder}}\left(\frac{\mathbf{\Lambda}^{\prime}}{||\mathbf{\Lambda}^{\prime}||}\right)\in\mathbb{R}^{32},$
$\displaystyle\mathrm{where}\
\mathbf{\Lambda}^{\prime}=\mathrm{diag}(\mathbf{\Lambda})\left[2\dots
k^{\mathrm{Lap}}\right]$ (52)

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 $\Lambda_{11}=0$. We also randomise the
eigenvector sign every epoch which is otherwise arbitrarily defined.

#### Random Walk Structural Encoding (Dwivedi et al., 2022)

This feature captures the probability that a random graph walk starting at
node $i$ will finish back at node $i$, and is computed using powers of the
transition matrix $\mathbf{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^{\mathrm{RW}}=16$ steps were computed to form the feature
vector.

$\displaystyle\forall i:~{}~{}~{}\bar{\mathbf{x}}_{i}^{\mathrm{RW}}$
$\displaystyle=\left[(\mathbf{P}^{1})_{ii},\,(\mathbf{P}^{2})_{ii},\,\cdots,\,(\mathbf{P}^{k^{\mathrm{RW}}})_{ii}\right]$
$\displaystyle\mathrm{where}\ \mathbf{P}=\mathbf{D}^{-1}\mathbf{A}$ (53)
$\displaystyle\forall i:~{}~{}~{}\mathbf{x}_{i}^{\mathrm{RW}}$
$\displaystyle=\texttt{MLP}_{\textrm{encoder}}\left(\bar{\mathbf{x}}_{i}^{\mathrm{RW}}\right)\in\mathbb{R}^{32},$
(54)

#### Local Graph Centrality Encoding (Ying et al., 2021a; Shi et al., 2022)

The graph centrality encoding is intended to allow the network to gauge the
importance of a node based on its connectivity, by embedding the degree
(number of incident edges) of each node into a learnable feature vector.

$\displaystyle\forall i:~{}~{}~{}\mathbf{x}_{i}^{\text{Cent}}$
$\displaystyle=\texttt{Embed}_{64}\left(D_{ii}\right)\in\mathbb{R}^{64}$ (55)

#### 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 $\Delta\in\mathbb{N}^{N\times N}$
is computed where $\Delta_{ij}$ is the number of edges in the shortest
continuous path from node $i$ to node $j$. During training each integer
distance is embedded as a scalar attention bias term to create the SPD
attention bias map $\mathbf{B}^{\mathrm{SPD}}\in\mathbb{R}^{N\times N}$.

$\displaystyle\forall
i,j:~{}~{}~{}B_{ij}^{\mathrm{SPD}}=\texttt{Embed}_{1}\left(\Delta_{ij}\right)\in\mathbb{R}$
(56)

Single-headed attention is assumed throughout this report for simplified
notation, however, upon extension to multi-headed attention, one bias is
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 we use only the distances
between atoms, not the positions directly. To embed the scalar distances into
vector space $\mathbb{R}^{K}$ we first apply $K=128$ Gaussian kernel
functions, where the $k^{\text{th}}$ function is defined as

$\displaystyle\forall i,j:~{}~{}~{}\bar{\psi}_{ij}^{k}$
$\displaystyle=\frac{||\mathbf{r}_{i}-\mathbf{r}_{j}||-\mu^{k}}{|\sigma^{k}|}$
(57) $\displaystyle\forall i,j:~{}~{}~{}\psi_{ij}^{k}$
$\displaystyle=-\frac{1}{\sqrt{2\pi}\left|\sigma^{k}\right|}\exp\left(-\frac{1}{2}\left(\bar{\psi}_{ij}^{k}\right)^{2}\right)\in\mathbb{R}$
(58)

with learnable parameters $\mu^{k}$ and $\sigma^{k}$. The $K$ elements are
concatenated into vector $\boldsymbol{\psi}_{ij}$. 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 $\mathbf{B}^{\text{3D}}\in\mathbb{R}^{N\times N}$
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)

$\displaystyle\forall i,j:~{}~{}~{}\mathbf{B}_{ij}^{\text{3D}}$
$\displaystyle=\texttt{MLP}_{\textrm{bias3D}}(\boldsymbol{\psi}_{ij})\in\mathbb{R}$
(59)

Upon extension to multi-headed attention with 32 heads
$\texttt{MLP}_{\mathrm{bias3D}}$ instead projects to $\mathbb{R}^{32}$.

#### Bond Length Encoding

Whilst $\mathbf{B}^{\text{3D}}$ 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 to the
MPNN module in a sparse but high-dimensional manner as edge features
$\mathbf{E}^{\text{3D}}=\bigl{[}\mathbf{e}^{\text{3D}}_{uv}~{}\mathsf{for}~{}(u,v)\in\mathcal{E}\bigr{]}$
calculated as

$\displaystyle\forall(u,v):~{}~{}~{}\mathbf{e}_{uv}^{\text{3D}}$
$\displaystyle=\texttt{MLP}_{\textrm{encoder}}(\boldsymbol{\psi}_{uv})\in\mathbb{R}^{32}$
(60)

#### Global 3D Centrality Encoding

The 3D node centrality features
$\mathbf{X}^{\text{3D}}=\left[\mathbf{x}_{1}^{\text{3D}};\,\dots;\,\mathbf{x}_{N}^{\text{3D}}\right]$
are computed by summing the embedded 3D distances from node $i$ to all other
nodes. Since the sum commutes this feature cannot be used to determine the
distance to a specific node, so serves as a centrality encoding rather than a
positional encoding.

$\displaystyle\forall i:~{}~{}~{}\mathbf{x}_{i}^{\text{3D}}$
$\displaystyle=W^{\text{3D}}\sum_{j\in\mathcal{V}}\boldsymbol{\psi}_{ij}\in\mathbb{R}^{32}$
(61)

Here $W^{\text{3D}}\in\mathbb{R}^{K\times 32}$ is a linear projection to the
same latent size as the other encoded features.

Table A.2: Ensembled model performance on PCQM4Mv2 dataset. Models in the proxy set are trained on the train+half_valid data split whereas those in the full set are trained on all available data. | Proxy Set | Main Set |
---|---|---|---
|  | Valid MAE |  | Ensembling
Case | # Models | Avg. | Ensembled | # Models | Weight
1: Baseline | 10 | 0.0755 | 0.0725 | 35 | 1
2: No Atomic Number | 4 | 0.0761 | 0.0734 | 16 | 0.5
3: FNN Dropout = 0.412 | 8 | 0.0759 | 0.0729 | 14 | 1
4: FNN Dropout = 0.412; No Atomic Number | 5 | 0.0761 | 0.0736 | 7 | 0.5
5: Feature Set 2† | 4 | 0.0755 | 0.0731 | 15 | 1
6: Feature Set 3† | 4 | 0.0754 | 0.0731 | 14 | 1
7: Masking Weights = [1,2,2] | 4 | 0.0754 | 0.0730 | 15 | 1
All | 39 | 0.0756 | 0.0722 | 112 |
† As defined in Table A.1.

##  Appendix B Hardware and Acceleration

###  B.1 Hardware

We train our models using a BOW-POD16 which contains 16 IPU processors,
delivering a total of 5.6 petaFLOPS of float16 compute and 14.4 GB of in-
processor SRAM which is accessible at an aggregate bandwidth of over a
petabyte per second. This compute and memory is then distributed evenly over
1472 tiles per chip. This architecture has two key attributes that enable high
performance on GNN and other AI 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 into 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 are 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 is 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 tensor 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 to 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 to _pack_ a number of graphs into a fixed size shape to minimise the
amount of padding required, this is an abstraction that is common in graph
software 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 $\mathcal{O}(n^{2})$ 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 a balance between these two extremes we employ a two tiered
hierarchical batching scheme that packs graphs into a 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 increased by employing a more complex packing
strategy (Krell et al., 2021). We then form micro-batches of 8 packs which are
pipelined (Huang et al., 2018) over 4 IPUs accumulating over 8 micro-batches
and replicated 4 times to form a global batch size of 921 graphs distributed
over 16 IPUs.

###  B.3 Numerical Precision

To maximise compute throughput and maximise 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 FLOP
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 all111A few
operations like the sum of squares the variance calculations are up-cast to
float32 by the compiler. compute but also use float16 for the majority222A
small number of weights are kept in float32 for simplicity of the code rather
than numerical stability. 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.

##  Appendix C Ensembling

###  C.1 Dataset Splits

While the original training set already contains 98% of the labelled data and
all of the data with 3D positional information, we still aim to train on all
available data for our final model. This, however, comes with many pitfalls
due to inability to calculate a reliable measure of performance. Therefore, to
understand the value of this additional data we consider an intermediate split
configuration where we randomly sample half of the validation set (not
resampled per run) to train on, holding the other half out for validation. As
a result we have three dataset splits to consider: original, train+valid and
train+half_valid.

###  C.2 Ensembling

Ensembling models has long been used to improve generalisation of machine
learning models and has become an indispensable tool for practitioners
entering machine learning competitions. Here we outline our ensembling
strategy that aims to achieve confidence while training without a validation
set, but also allows on-the-fly tuning and weighting of model ensembles. The
main idea is to build two comparative model sets to ensemble, a _proxy_ set
and a _main_ set. The proxy set is designed to be qualitatively similar to the
main set but maintain a clean held out validation set by training on
train+half_valid, whereas the main set is trained on train+valid . We aim for
the proxy set to match the main set in all aspects apart from the training
data and the number of models; to be able to focus computational resources on
the main set we aim to build the proxy set to be approximately 25% of its
size.

Diversity of models in an ensemble is a well known way to boost model
performance (Zhou et al., 2002; Lakshminarayanan et al., 2017). In this work
we propose six adjustments to the hyperparameters to form seven different
model configurations to ensemble. In choosing these configurations we aim to
build diversity in three main areas: input feature description (cases 2, 4, 5
and 6), regularisation strength (cases 2 and 3) and reliance on 3D features
(case 7). As described in in the previous paragraph we build a _proxy_ set of
models to help guide our ensembling strategy as well as a _main_ set trained
on all available data. The details of our seven configurations and their
evaluation are shown in Table A.2.

Analysing the proxy set we found that reducing the weighting for the two worst
performing models gave a small improvement to our final ensembled MAE, we
therefore apply this is our main final ensemble too.

The final 112 model ensemble achieves an MAE of 71.9 meV on the test challenge
set. This was completed in 1 hour 32 minutes using 1 IPU and AMD EPYC 7742
64-Core CPU including all the feature processing, program compilation and
model inference.