# Masked Attention is All You Need for Graphs

David Buterez  Jon Paul Janet  Dino Oglic  Pietro Lio

###### Abstract

Graph neural networks (GNNs) and variations of the message passing algorithm
are the predominant means for learning on graphs, largely due to their
flexibility, speed, and satisfactory performance. The design of powerful and
general purpose GNNs, however, requires significant research efforts and often
relies on handcrafted, carefully-chosen message passing operators. Motivated
by this, we propose a remarkably simple alternative for learning on graphs
that relies exclusively on attention. Graphs are represented as node or edge
_sets_ and their connectivity is enforced by masking the attention weight
matrix, effectively creating custom attention patterns for each graph. Despite
its simplicity, masked attention for graphs (MAG) has state-of-the-art
performance on long-range tasks and outperforms strong message passing
baselines and much more involved attention-based methods on over 55 node and
graph-level tasks. We also show significantly better transfer learning
capabilities compared to GNNs and comparable or better time and memory
scaling. MAG has sub-linear memory scaling in the number of nodes or edges,
enabling learning on dense graphs and future-proofing the approach.

Machine Learning, ICML

##  1 Introduction

The field of geometric deep learning (GDL) seeks to describe, understand, and
even “unify” deep learning strategies for data structures such as sets, grids,
and graphs by leveraging the fundamental concepts of symmetry and invariance.
A remarkably successful application of GDL is learning on graphs, abstractions
that represent relationships between items of a set, and which naturally
describe real-world phenomena such as social, biological, or transportation
networks, as well as objects like molecules. In particular, deep learning for
molecules has the potential to accelerate and even revolutionise fields such
as drug discovery and materials science, being one of the main factors behind
the accelerated development of GDL, as well as one of the earliest (Duvenaud
et al., 2015; Kearnes et al., 2016; Gilmer et al., 2017). General purpose
learning on graphs is typically specified as an instance of _message passing_
, an iterative algorithm where one must define a message function which
aggregates information from a given node’s neighbourhood, as well as a node
(possibly also edge) update function to incorporate the encoded messages.

Despite the overall success and wide adoption of GNNs, several fundamental
problems have been highlighted over time. Firstly, although the message
passing framework is highly customisable through user-defined, learnable
message and node update functions, the design of novel layers is a difficult
research problem, where improvements take years to achieve and often rely on
hand-crafted operators. This is particularly the case for general purpose GNNs
that do not exploit additional input modalities such as atomic coordinates.
For example, principal neighbourhood aggregation (PNA) (Corso et al., 2020) is
regarded as one of the most powerful message passing layers, but it is built
using a collection of hand-picked neighbourhood aggregation functions, it
requires a dataset degree histogram which must be pre-computed prior to
learning, and further uses hand-picked degree scalers. Another example is
given by Graph Attention Networks (GAT) (Veličković et al., 2018), one of the
most popular graph layers and one of the earliest efforts to combine attention
with GNNs. It has been shown, afterwards, that the original formulation of GAT
has limited expressive power, and that a simple reordering of the operations
can improve performance (Brody et al., 2022).

Secondly, the nature of message passing imposes certain limitations which have
shaped most of the GNN literature. One of the most prominent examples is the
readout function used to combine node-level representations into a single
graph-level representation, and which is required to be permutation invariant
with respect to the node order. Thus, the default choice for most GNNs remains
a simple, non-learnable function such as sum, mean, or max, despite the
potential limited expressivity. Recently, it has been shown that breaking this
constraint can lead to improved performance, and might be acceptable in
certain scenarios such as learning on molecules, where the inputs can be
presented in a canonical order (Buterez et al., 2022). Furthermore, expressive
readout functions that are based on Set Transformers (Lee et al., 2019) have
been shown to consistently improve performance of GNNs regardless of the
underlying message passing algorithm (Buterez et al., 2022, 2023b).

Thirdly, the majority of GNN architectures are plagued by the well-known
oversmoothing and oversquashing problems that limit the depth of GNNs, as well
as difficulties in modelling long-range relationships, all of which are
consequences of aggregating information from exponentially larger
neighbourhoods (Alon & Yahav, 2021). The proposed solutions typically take the
form of message regularisation schemes (Godwin et al., 2022; Zhao & Akoglu,
2020; Cai et al., 2021). However, there is generally not a consensus on the
right architectural choices for deep GNNs. Separately, standard GNNs have also
shown limitations in terms of transfer learning and strategies such as pre-
training and fine-tuning, as opposed to other families of neural networks such
as language models (Hu et al., 2020). For certain types of data and tasks,
non-standard GNNs that leverage attention-based readouts are currently the
only way to effectively perform transfer learning (Buterez et al., 2024,
2023a).

Perhaps due to the progress that message passing neural networks have enabled
on a wide range of tasks, alternative paradigms for learning on graphs are
relatively underdeveloped. The attention mechanism (Vaswani et al., 2017) is
one of the main sources of innovation within graph learning, either by
directly incorporating attention within message passing (Veličković et al.,
2018; Brody et al., 2022), by formulating graph learning as a language
processing task (Ying et al., 2021; Kim et al., 2022), or by combining
standard GNN layers with an attention mechanism (Rampasek et al., 2022;
Shirzad et al., 2023; Buterez et al., 2022, 2023b).

Here, we propose a novel graph learning framework characterised by its
simplicity and lack of message passing layers, instead being based on a
classical attention mechanism with masking to create custom attention
patterns. Compared to some of the previous works that have formulated graph
learning as a language modelling task, we take a different approach and
consider graph learning as a learning task on _sets_ , where the graph
connectivity (i.e., adjacency matrix) is enforced by masking the pairwise
attention weight matrix and allowing only values that correspond to graph
connections. We term this architecture Masked Attention for Graphs (MAG) and
demonstrate that it can be customised to propagate information across nodes
(MAGN) or edges (MAGE). It is general purpose, in the sense that it only
relies on the graph structure and possibly node and edge features, and it is
not restricted to a particular domain such as chemistry. The simplicity of the
architecture is further demonstrated by the fact that MAG does not use
positional or similar encodings, it does not encode graph structures as tokens
or other language (sequence) specific concepts, and it does not require any
pre-computations.

Figure 1: Overview of the Masked Attention for Graphs (MAG) architecture. An
important aspect is the choice of processing the node feature set (MAGN) or
the edge feature set (MAGE), and the appropriate masking algorithm.
Conventional self attention blocks (S) can be alternated with masked variants
(M) as desired (here, a choice of SMSM is depicted). The decoder is only
required for graph-level tasks.

Despite its simplicity, MAG generally outperforms strong message passing
baselines and much more involved attention-based algorithms. Our empirical
evaluation covers long-range molecular benchmarks and over 55 tasks from
different domains such as quantum mechanics, molecular docking, physical
chemistry, biophysics, bioinformatics, computer vision, social networks,
functional call graphs, and synthetic graphs. This emphasises the fact that
the carefully-selected and hand-crafted nature of most message passing
algorithms can be easily superseded by attention itself, without the need to
explicitly define any operator. Beyond benchmarking, we also explore a recent
research direction that showed the transformative benefits of transfer
learning in drug discovery and quantum mechanics (Buterez et al., 2024). Here,
we leverage a newly-published, refined version of the qm9 dataset at a higher
level of theory (Fediai et al., 2023), and show that MAG is a viable and well
performing transfer learning strategy, while GNNs are limited.

MAG is arguably one of the most straightforward ways to apply attention on
graphs. Thanks to modern implementations of exact attention, MAG scales sub-
linearly in terms of memory with the number of nodes (MAGN) or edges (MAGE),
as it largely relies on (masked) self-attention. Although current libraries
are not optimised for masked attention and many optimisations are possible
(see Appendix A), both the training time and the memory consumption are
competitive or even better than plain message passing.

##  2 Related Work

Graph neural networks with adaptive readouts – A recent research trend
consists of following standard GNN layers with an attention-based pooling
(readout) function, for example using the Set Transformer (Buterez et al.,
2022, 2024, 2023b), or standard Transformers (Jain et al., 2021). The change
to a more expressive readout function has provided consistent uplifts in most
supervised learning tasks, and has enabled easy and effective transfer
learning for molecules through a pre-training and fine-tuning workflow.

Transformers for graphs – One of the most popular approaches is the direct
application of standard Transformers designed for language modelling to
graphs. Graphormer (Ying et al., 2021) achieves this through an involved and
computation-heavy suite of pre-processing steps, involving a centrality,
spatial, and edge encoding. Graphormer models the mean readout function from
GNNs through a special “virtual” node that is added to the graph and connected
to every other node. Another approach within this paradigm is the Tokenized
Graph Transformer (TokenGT) (Kim et al., 2022), which treats all the nodes and
edges as independent tokens. To adapt sequence learning to the graph domain,
TokenGT encodes the graph information using node identifiers derived from
orthogonal random features or Laplacian eigenvectors, and learnable type
identifiers for nodes and edges. Such models rely on embedding layers and are
thus limited to integer node or edge features. Moreover, the high number of
sequence-based architectures (e.g. Big Bird (Zaheer et al., 2020), Performer
(Choromanski et al., 2021), etc.) represents a large unexplored territory for
graphs. Spectral Attention Networks (Kreuzer et al., 2021) use a
computationally-expensive positional encoding based on the Laplacian and two
complementary attention mechanisms over nodes. A different philosophy is that
of a message passing and Transformer hybrid. Such an approach is taken by the
GraphGPS framework (Rampášek et al., 2022), which alternates message passing
layers with Transformer layers. Like previous works, GraphGPS puts a large
emphasis on different types of encodings, proposing and analysing positional
and structural encodings, further divided into local, global, and relative
encodings. Exphormer (Shirzad et al., 2023) is an evolution of GraphGPS which
adds virtual global nodes and sparse attention based on expander graphs. While
effective, such frameworks do still rely on message passing and are thus not
purely attention-based solutions. Limitations include dependence on
approximations (Performer for both, expander graphs for Exphormer), decreased
performance when encodings (GraphGPS) or special nodes (Exphormer) are
removed, and scalability (GraphGPS). Notably, Exphormer is the first work to
consider custom attention patterns for graphs in the form of node-level
neighbourhoods.

##  3 Preliminaries

Graphs – A graph is a tuple $\mathcal{G}=(\mathcal{V},\mathcal{E})$ where
$\mathcal{V}$ represents the set of nodes,
$\mathcal{E}\subseteq\mathcal{V}\times\mathcal{V}$ is the set of edges, and
$N_{n}=|\mathcal{V}|$, $N_{e}=|\mathcal{E}|$. Nodes are associated with
feature vectors $\mathbf{x}_{u}$ of dimension $d_{n}$ for all nodes
$u\in\mathcal{V}$, and $d_{e}$-dimensional edge features $\mathbf{e}_{uv}$ for
all edges $e\in\mathcal{E}$. The node features are collected as rows in a
matrix $\mathbf{X}\in\mathbb{R}^{N_{n}\times d_{n}}$, and similarly for edge
features into $\mathbf{E}\in\mathbb{R}^{N_{e}\times d_{e}}$. The graph
connectivity information can be represented as an adjacency matrix
$\mathbf{A}$, where $\mathbf{A}_{uv}=1$ if $(u,v)\in\mathcal{E}$ and
$\mathbf{A}_{uv}=0$ otherwise, although an edge list representation is often
more practical. GNNs also define the neighbourhood of a node $u\in\mathcal{V}$
as $\mathcal{N}_{u}=\\{v\mid(u,v)\in\mathcal{E}\lor(v,u)\in\mathcal{E}\\}$.
Many GNNs can be described as a form of message passing, which takes the
general form of
$m_{u}^{t+1}=\sum_{v\in\mathcal{N}_{u}}{\texttt{message}_{t}\left(\mathbf{x}_{u}^{t},\mathbf{x}_{v}^{t},\mathbf{e}_{uv}\right)}$
and $x_{u}^{t+1}=\texttt{update}_{t}\left(x_{u}^{t},m_{u}^{t+1}\right)$
(Gilmer et al., 2017), where message and update are the message and node
update functions, respectively, and $t$ is the current time step. As described
above, message passing is highly customisable, for example by novel
definitions of the message and update functions, their inputs, or even adding
new update steps such as the graph-level state embedding and update proposed
in Chen et al. (2021) for multi-fidelity applications. For graph-level
prediction tasks, a readout or pooling function $\bigoplus$ must be used,
aggregating all the learnt node representations:
$\mathbf{x}_{\mathcal{G}}=\bigoplus_{u\in\mathcal{V}}\left(\mathbf{x}_{u}\right)$.

Set Transformer – The Set Transformer is an encoder-decoder attention-based
architecture for learning on sets. It leverages the scaled dot product and
multihead attention

1 from torch_geometric import unbatch_edge_index

2 function node_mask(b_ei, b_map, B, M)

3 # b_ei is the batched edge index

4 # b_map maps nodes to graphs

5 # B, M batch, respectively mask size

6 mask $\leftarrow$ torch.full(size=(B, M, M), fill=False)

7 graph_idx $\leftarrow$ b_map.index_select(0, b_ei[0, :])

8 eis $\leftarrow$ unbatch_edge_index(b_ei, b_map)

9 ei $\leftarrow$ torch.cat(eis, dim=1)

10mask[graph_idx, ei[0, :], ei[1, :]] $\leftarrow$ True

11return $\scriptstyle\mathtt{\sim}$mask

Algorithm 1 Node masking algorithm, where the inputs follow PyTorch Geometric
conventions.

mechanisms proposed in Vaswani et al. (2017) to define multihead attention
blocks (MABs), self attention blocks (SABs), and pooling by multihead
attention (PMA) blocks. We do not recapitulate the definition of multihead
attention and assume instead that a function
$\text{Multihead}\left(\mathbf{Q},\mathbf{K},\mathbf{V}\right)$ is available
(for query, key, and value matrices, respectively). The original Set
Transformer is given by
$\text{ST}(\mathbf{X})=\text{Decoder}(\text{Encoder}(\mathbf{X}))$ with:

$\displaystyle\text{MAB}(\mathbf{X},\mathbf{Y})$
$\displaystyle=\mathbf{H}+\text{Linear}_{\phi}(\mathbf{H})$ (1)
$\displaystyle\text{and }\mathbf{H}$
$\displaystyle=\mathbf{X}+\text{Multihead}(\mathbf{X},\mathbf{Y},\mathbf{Y}),$
(2) $\displaystyle\text{SAB}(\mathbf{X})$
$\displaystyle=\text{MAB}\left(\mathbf{X},\mathbf{X}\right),$ (3)
$\displaystyle\text{Encoder}(\mathbf{X})$
$\displaystyle=\text{SAB}^{n}\left(\mathbf{X}\right),$ (4)
$\displaystyle\text{PMA}_{k}(\mathbf{Z})$
$\displaystyle=\text{MAB}\left(\mathbf{S}_{k},\text{Linear}_{\phi}(\mathbf{Z})\right),$
(5) $\displaystyle\text{Decoder}(\mathbf{Z})$
$\displaystyle=\text{Linear}_{\phi}\left(\text{SAB}^{n}\left(\text{PMA}_{k}(\mathbf{Z})\right)\right).$
(6)

and where $\text{Linear}_{\phi}$ is a linear layer followed by an activation
function $\phi$, $\text{SAB}^{n}(\cdot)$ represents $n$ subsequent
applications of a SAB, and $\mathbf{S}_{k}$ is a tensor of $k$ learnable seed
vectors that are randomly initialised ($\text{PMA}_{k}$ outputs $k$ vectors)
(Buterez et al., 2023b). Notably, the original Set Transformer is designed
exclusively for learning on sets, uses layer normalisation (LN) in a post-LN
fashion (Xiong et al., 2020a), and does not use any form of positional or
structural encodings.

Efficient attention – Recently, Flash attention has enabled exact attention
with linear memory scaling and faster training and inference (up to x10)
thanks to more efficiently utilising the architecture of modern graphics
processing units (GPUs) (Dao et al., 2022; Dao, 2023). For self-attention, an
exact attention implementation has also been developed that scales with the
square root of the sequence length in terms of memory, and has comparable run
time with standard implementations (Rabe & Staats, 2021).

##  4 Masked attention for graphs

We formulate graph learning as a learning problem on _sets_. The main learning
mechanism consists of applying attention directly to the node or edge feature
matrix by means of SABs. We propose _masking_ the pairwise attention weight
matrix according to the node or edge adjacency information in order to
incorporate the graph structure. We call this approach Masked Attention for
Graphs (MAG).

Inputs – MAG supports two main ways of information propagation: (1) on nodes
(MAGN), more specifically on the node feature matrix $\mathbf{X}$, or (2) on
edges (MAGE), i.e. on the edge feature matrix $\mathbf{E}$.

[!ht]  Edge masking algorithm, where the inputs follow PyTorch Geometric
conventions. T is the transpose. Helper functions are explained in Appendix B.

from mag import consecutive, first_unique_index

2

function edge_adjacency(b_ei)

3 E $\leftarrow$ b_ei.size(1)

4 source_nodes $\leftarrow$ b_ei[0]

5 $\hphantom{\text{{source\\_nodes}}}\mathllap{\text{{target\\_nodes}}}$
$\leftarrow$ b_ei[1]

6

87 # unsqueeze and expand  exp_src $\leftarrow$ source_nodes.unsq(1).exp((-1,
E))

9 $\hphantom{\text{{exp\\_src}}}\mathllap{\text{{exp\\_trg}}}$ $\leftarrow$
target_nodes.unsq(1).exp((-1, E))

10

11 src_adj $\leftarrow$ exp_src == T(exp_src)

12 $\hphantom{\text{{src\\_adj}}}\mathllap{\text{{trg\\_adj}}}$ $\leftarrow$
exp_trg == T(exp_trg)

13 $\hphantom{\text{{src\\_adj}}}\mathllap{\text{{cross}}}$ $\leftarrow$
(exp_src == T(exp_trg)) logical_or

14 (exp_trg == T(exp_src))

15

16 return (src_adj logical_or trg_adj logical_or cross)

17function edge_mask(b_ei, b_map, B, M)

18 mask $\leftarrow$ torch.full(size=(B, M, M), fill=False)

19 edge_to_graph $\leftarrow$ b_map.index_select(0, b_ei[0, :])

20

21 edge_adj $\leftarrow$ edge_adjacency(b_ei)

22 ei_to_original $\leftarrow$ consecutive(

23 first_unique_index(edge_to_graph), b_ei.size(1))

24

25 edges $\leftarrow$ edge_adj.nonzero()

26 graph_idx $\leftarrow$ edge_to_graph.idx_select(0, edges[:, 0])

27 coord_1 $\leftarrow$ ei_to_original.idx_select(0, edges[:, 0])

28 coord_2 $\leftarrow$ ei_to_original.idx_select(0, edges[:, 1])

29

30 mask[graph_idx, coord_1, coord_2] $\leftarrow$ True

31 return $\scriptstyle\mathtt{\sim}$mask

Masking – We extend the Set Transformer with masked equivalents of the MAB and
SAB (MSAB). Practically, this means supplying the new blocks with a mask
tensor of shape $B\times N_{d}\times N_{d}$ for MAGN and $B\times N_{e}\times
N_{e}$ for MAGE, where $B$ is the batch size. In a naive implementation,
masking simply means replacing the targeted values of the scaled $QK^{T}$
product with negative infinity (or a very large negative value for stability)
before softmax (Appendix F for the equation). The correct mask for each batch
is different and must thus be computed on the fly. For MAGN, the mask allows
only adjacent nodes. In other words, the $N_{d}\times N_{d}$ portion of the
mask corresponds directly to the node adjacency matrix (Algorithm 1). For
MAGE, the mask operates on the set of edges and must allow only edges that
share a common node. While this computation is not as trivial as the MAGN
mask, both kinds of masks can be efficiently computed using exclusively tensor
operations (Section 4).

Architecture – At a high level, a complete MAG model takes the form of an
encoder where MSAB and SAB blocks are alternated as desired, with a PMA-based
decoder (Figure 1). Compared to the Set Transformer, we have also adapted MAG
to use a pre-LN architecture with layer or batch normalisation and optionally
include multi-layer perceptrons (MLPs) after multihead attention. For graph-
level tasks, the PMA module of the Set Transformer acts as an equivalent to
the readout function in GNNs, but fully based on attention, while for node-
level tasks PMA is not required.

Table 1: Test set mean absolute error (MAE) or average precision (AP) for two long-range molecular benchmarks. All the results except for MAGE are extracted from (Tönshoff et al., 2023). The number of layers for pept-struct, respectively pept-func is given as $(\cdot/\cdot)$. Dataset |  GCN (6/6) |  GIN (10/8) |  GraphGPS (8/6) |  Exphormer (4/8) |  MAGE (3/4)
---|---|---|---|---|---
pept-struct (MAE $\downarrow$) |  $0.246$ $\pm$ | $0.0007$ |  $0.2473$ $\pm$ | $0.0017$ |  $0.2509$ $\pm$ | $0.0014$ |  $0.2481$ $\pm$ | $0.0007$ |  $0.245329301$ $\pm$ | $0.000297453$
pept-func (AP $\uparrow$) |  $0.686$ $\pm$ | $0.005$ |  $0.6621$ $\pm$ | $0.0067$ |  $0.6534$ $\pm$ | $0.0091$ |  $0.6527$ $\pm$ | $0.0043$ |  $0.686267567$ $\pm$ | $0.004352919$
Table 2: Test set Matthews correlation coefficient (MCC) for 3 node-level classification tasks, presented as mean $\pm$ standard deviation from 5 different runs. The highest mean values are highlighted in bold. Dataset | GCN | GAT | GATv2 | GIN | PNA | MAGN
---|---|---|---|---|---|---
ppi $\uparrow$ |  $0.465077143907547$ $\pm$ | $0.022862258475938636$ |  $0.8206786870956421$ $\pm$ | $0.027530049860917918$ |  $0.7236320707947015$ $\pm$ | $0.3881943479190861$ |  $0.35471993088722226$ $\pm$ | $0.022346513820448957$ |  $0.8287733674049378$ $\pm$ | $0.005196909272198615$ |  $0.9885392665863038$ $\pm$ | $0.0006107937592923426$
CiteSeer $\uparrow$ |  $0.30353574156761165$ $\pm$ | $0.03971441229040849$ |  $0.19309382140636439$ $\pm$ | $0.0329168272454626$ |  $0.20521373748779292$ $\pm$ | $0.045510821851640064$ |  $0.22781004607677455$ $\pm$ | $0.025327781492425105$ |  $0.013885515648871641$ $\pm$ | $0.01845756823953702$ |  $0.538603150844574$ $\pm$ | $0.022423541774820147$
Cora $\uparrow$ |  $0.4821257650852203$ $\pm$ | $0.06321092918897829$ |  $0.37335475683212277$ $\pm$ | $0.012007111729171077$ |  $0.43192169070243835$ $\pm$ | $0.04347835004836811$ |  $0.391590827703476$ $\pm$ | $0.01897646414361149$ |  $0.04265761561691756$ $\pm$ | $0.04681603852353973$ |  $0.696277630329132$ $\pm$ | $0.01700565865669409$

##  5 Experiments

Our evaluation encompasses: (1) an extensive suite of over 55 benchmarks from
various domains, where the goal is to compare MAG with representative message
passing networks (GCN, GAT, GATv2, GIN, and PNA) and Transformers for graphs
(Graphormer, TokenGT). Due to space limitations, we cannot present the results
for all methods in the main part of the paper, so we select the most
representative ones based on the task, with the rest presented in the
Appendix. Graphormer and TokenGT are applicable only to a subset of tasks due
to limitations discussed in Related Work. They also fail for many datasets due
to very high memory requirements (CPU and/or GPU). We cover long-range tasks
(Section 5.1), node-level tasks (Section 5.2) and graph-level tasks (Section
5.3). (2) A transfer learning performance evaluation (Section 5.4) based on a
new, refined variant of qm9 and following the recent framework proposed by
Buterez et al. (2024). (3) The time and memory characteristics for all
discussed methods (Section 5.5).

###  5.1 Long-range tasks

Graph learning with Transformers has traditionally been evaluated on long-
range graph benchmarks (LRGB) (Dwivedi et al., 2022). However, it was recently
shown that simple GNNs outperform most attention-based methods (Tönshoff et
al., 2023). Nonetheless, we evaluated MAGE on two long-range molecular tasks
(Table 1) and conclude that it outperforms GraphGPS, Exphormer, GCN, and GIN.
Despite using half the number of layers as other methods or less, MAGE matches
the 2nd model on the pept-struct leaderboard, and is within the top 5 for
pept-func (as of January 2024). Remarkably, MAG is the only top method that is
exclusively based on attention (i.e. no message passing or hybrid), does not
use any positional or similar encoding, and is general purpose (i.e. not
specifically built for molecules). We also did not use hyperparameter
optimisation or sophisticated schedulers (Appendix C).

###  5.2 Node-level tasks

Generally, node-level tasks take the form of citation networks of different
sizes, and are not as varied as graph-level problems. Nonetheless, they
represent an interesting case for MAG as PMA is not needed. Here, we selected
3 representative datasets: ppi, CiteSeer, and Cora. In particular, MAGN is the
most natural choice as it works over node representations. Our results
indicate that MAGN is the best performing method by a large margin (Table 2).
Graphormer and TokenGT are not available for node-level classification, and
they would not work due to the large graph sizes.

###  5.3 Graph-level tasks

Graph-level tasks are generally more varied, as they often originate from
different domains and require a readout function for GNNs and a PMA module for
MAG. Here, we do not focus on the differences between readouts and instead
choose a reasonably strong baseline (4-layer GNNs and mean readout; Appendix C
for the experimental setup). A more granular evaluation of readouts has been
done by (Buterez et al., 2022), concluding that the difference between
readouts is small for most tasks.

qm9 is a quantum mechanics dataset consisting of 133,885 small organic
molecules and 19 regression targets given by quantum properties (Ramakrishnan
et al., 2014). We report results for all 19 targets in Table 3, with GCN and
TokenGT separately in Appendix D due to limited space. We observe that for 15
out of 19 properties, MAGE is the best performing method. For the dipole
moment ($\mu$), PNA is better by a slight margin, while for HOMO, LUMO, and
the HOMO-LUMO gap ($\Delta\epsilon$), Graphormer is stronger. We expect that
attention-based methods would be the best suited for intensive and localised
properties like the HOMO and LUMO, which agrees with recent literature on this
topic (Buterez et al., 2023b). The fact that MAGE is not as competitive on
these tasks could be explained by MAGE converging quicker than other methods
and getting stuck in local minima, especially for HOMO/LUMO which take a long
time to converge for most methods. We report the results of altering the
number and order of attention blocks in Appendix I.

Table 3: Test set root mean squared error (RMSE, standard for quantum mechanics) qm9, and $R^{2}$ for the rest of tasks, presented as mean $\pm$ standard deviation from 5 different runs. The lowest (qm9) and highest (rest) mean values are highlighted in bold. Only one of Graphormer/TokenGT was chosen for spacing reasons, based on competitiveness and lack of out-of-memory errors (oom). Any results not displayed here (e.g. GCN, TokenGT) are available in Appendix D. | Property | GAT | GATv2 | GIN | PNA | Graphormer | MAGE
---|---|---|---|---|---|---|---
qm9 |  $\mu$ $\downarrow$ |  $0.6117233026$ $\pm$ | $0.00322791030417739$ |  $0.6054740547999999$ $\pm$ | $0.006568307253910333$ |  $0.6148490484$ $\pm$ | $0.004785032811027966$ |  $0.573500891$ $\pm$ | $0.004342105898209744$ |  $0.6342938602$ $\pm$ | $0.004869978039706592$ |  $0.6139972304$ $\pm$ | $0.015654514385105636$
$\alpha$ $\downarrow$ |  $2.663973035418298$ $\pm$ | $0.25244283100145354$ |  $1.863891228055802$ $\pm$ | $0.2795333984842851$ |  $1.1815054031344825$ $\pm$ | $0.10219442305891348$ |  $1.001273514296955$ $\pm$ | $0.01599294577680106$ |  $0.568715371596984$ $\pm$ | $0.16393874177743337$ |  $0.47676611551534853$ $\pm$ | $0.012005002030815186$
$\epsilon_{\text{HOMO}}$ $\downarrow$ |  $0.1239681316$ $\pm$ | $0.003314684464323611$ |  $0.11666108420000001$ $\pm$ | $0.0012842508330091484$ |  $0.1195951676$ $\pm$ | $0.0012943660592088296$ |  $0.108021064$ $\pm$ | $0.001231428119351671$ |  $0.104775851$ $\pm$ | $0.0036275429013649858$ |  $0.11114805679999999$ $\pm$ | $0.003119979091718997$
$\epsilon_{\text{LUMO}}$ $\downarrow$ |  $0.1395206008$ $\pm$ | $0.0036951568187456045$ |  $0.1281089538$ $\pm$ | $0.0031740911498961692$ |  $0.13335046779999998$ $\pm$ | $0.00186180246289051$ |  $0.1147641956$ $\pm$ | $0.0011389572800677408$ |  $0.1111764658$ $\pm$ | $0.0025760152606515573$ |  $0.12704481080000002$ $\pm$ | $0.0009926947853880807$
$\Delta\epsilon$ $\downarrow$ |  $0.18422291732115537$ $\pm$ | $0.0024710793089515893$ |  $0.17435110946604515$ $\pm$ | $0.005431627014889936$ |  $0.1851481248459784$ $\pm$ | $0.001993287398384209$ |  $0.15751472430741725$ $\pm$ | $0.0012991034958098758$ |  $0.15884127268237536$ $\pm$ | $0.006698267921781763$ |  $0.18655848180555562$ $\pm$ | $0.0013864649479679875$
$\langle R^{2}\rangle$ $\downarrow$ |  $53.027619422$ $\pm$ | $4.062079126887478$ |  $47.812721974$ $\pm$ | $3.502585904736325$ |  $36.136113162$ $\pm$ | $0.1764717036367732$ |  $35.167091398$ $\pm$ | $0.348914751219973$ |  $31.648359963999997$ $\pm$ | $1.7944812189277686$ |  $28.569369802$ $\pm$ | $0.2632541893990392$
ZPVE $\downarrow$ |  $0.24064087731445105$ $\pm$ | $0.012656330427125584$ |  $0.1375877996547808$ $\pm$ | $0.020768341432428502$ |  $0.07721947728101804$ $\pm$ | $0.0063411183757174664$ |  $0.060956764236741824$ $\pm$ | $0.0025961967619237106$ |  $0.06072110088150688$ $\pm$ | $0.018500042354629327$ |  $0.02578744427174226$ $\pm$ | $0.000867405321348238$
$U_{0}$ $\downarrow$ |  $609.2545884699109$ $\pm$ | $68.91246798482135$ |  $329.5560245536768$ $\pm$ | $30.078194009649607$ |  $143.71066676082495$ $\pm$ | $15.267857172681275$ |  $100.48961222525227$ $\pm$ | $5.430613325154432$ |  $31.06019904915944$ $\pm$ | $11.447185151284497$ |  $9.934471503301662$ $\pm$ | $2.897252802290114$
$U$ $\downarrow$ |  $579.7456709272891$ $\pm$ | $78.18870316631192$ |  $319.36901774083105$ $\pm$ | $11.626977447369674$ |  $160.75278502468748$ $\pm$ | $46.209004269069034$ |  $100.40824367152953$ $\pm$ | $4.267309465079941$ |  $28.295482919184415$ $\pm$ | $18.807401031149467$ |  $10.047206778242886$ $\pm$ | $3.183884353263835$
$H$ $\downarrow$ |  $558.8919696371956$ $\pm$ | $71.12768538135153$ |  $310.9509549285775$ $\pm$ | $20.79414183324842$ |  $166.09629835960604$ $\pm$ | $41.685033236591075$ |  $102.84723246890955$ $\pm$ | $4.131815007261605$ |  $32.17624357934126$ $\pm$ | $11.431784703033006$ |  $11.054349856643729$ $\pm$ | $6.366421675276654$
$G$ $\downarrow$ |  $580.0280637853255$ $\pm$ | $68.4077800897621$ |  $320.5260489867811$ $\pm$ | $36.1375269073437$ |  $166.65868119200604$ $\pm$ | $54.18739740144541$ |  $102.42570293923484$ $\pm$ | $3.3612975719955993$ |  $19.89765339556925$ $\pm$ | $11.131743596330399$ |  $10.687896060994543$ $\pm$ | $6.337137186845336$
$c_{\text{V}}$ $\downarrow$ |  $0.9020331324184475$ $\pm$ | $0.15856427749274968$ |  $0.6355779562722873$ $\pm$ | $0.0224647004101471$ |  $0.5322265537909983$ $\pm$ | $0.039091974899190945$ |  $0.3890311747141559$ $\pm$ | $0.008857172159100635$ |  $0.465136218323983$ $\pm$ | $0.082881583328882$ |  $0.17387085680363815$ $\pm$ | $0.001166652341224601$
$U_{0}^{\text{ATOM}}$ $\downarrow$ |  $3.960773652210654$ $\pm$ | $0.6088762804966179$ |  $1.7677899727201005$ $\pm$ | $0.21581125328331746$ |  $1.1907888182589328$ $\pm$ | $0.05851557004396988$ |  $0.9575399360092511$ $\pm$ | $0.033472195956417525$ |  $0.4493103691513176$ $\pm$ | $0.049956492659528444$ |  $0.24094167894120844$ $\pm$ | $0.005769893746834744$
$U^{\text{ATOM}}$ $\downarrow$ |  $4.525032283007901$ $\pm$ | $0.32566518657167676$ |  $1.918839575022136$ $\pm$ | $0.2537237951528616$ |  $1.1875472189218579$ $\pm$ | $0.05025222857522079$ |  $0.9684797332668202$ $\pm$ | $0.021979652129970773$ |  $0.4717307664336431$ $\pm$ | $0.0939828634468958$ |  $0.24315737562659173$ $\pm$ | $0.005146016048884781$
$H^{\text{ATOM}}$ $\downarrow$ |  $3.7799190025250615$ $\pm$ | $0.6981536661091328$ |  $1.8059939092647184$ $\pm$ | $0.2382935220717042$ |  $1.1502412167773588$ $\pm$ | $0.03490654193975345$ |  $0.9639097900384808$ $\pm$ | $0.02524761193614273$ |  $0.4208364564507344$ $\pm$ | $0.034523773738333936$ |  $0.24543222960369357$ $\pm$ | $0.003695878189629133$
$G^{\text{ATOM}}$ $\downarrow$ |  $3.5005126755614064$ $\pm$ | $0.6695755314457321$ |  $1.7300406929708285$ $\pm$ | $0.10545412078920889$ |  $1.0700908459341352$ $\pm$ | $0.03941343434861267$ |  $0.8546739891402237$ $\pm$ | $0.024348148668088133$ |  $0.3764059317504705$ $\pm$ | $0.026607713474123114$ |  $0.22466719771921126$ $\pm$ | $0.015174473269980511$
|  $A$ $\downarrow$ |  $13.4502139476$ $\pm$ | $4.20356070133326$ |  $5.6968419636$ $\pm$ | $5.789613294398423$ |  $3.8274557144$ $\pm$ | $4.078319314816777$ |  $27.988005322799996$ $\pm$ | $38.022971371441514$ |  $1.6090884928000002$ $\pm$ | $0.1892611641565722$ |  $0.7464662176$ $\pm$ | $0.10619528617484261$
|  $B$ $\downarrow$ |  $0.23087088160000002$ $\pm$ | $0.014304232859457052$ |  $0.24962173979999996$ $\pm$ | $0.03255811973612031$ |  $0.22591736440000001$ $\pm$ | $0.017653905093035747$ |  $0.25791816160000003$ $\pm$ | $0.018555011493651515$ |  $0.11685811799999998$ $\pm$ | $0.02604917624377956$ |  $0.0791156356$ $\pm$ | $0.011033295424966503$
|  $C$ $\downarrow$ |  $0.1779067728$ $\pm$ | $0.03206959033621605$ |  $0.21522022979999997$ $\pm$ | $0.030866224016824595$ |  $0.22185758240000003$ $\pm$ | $0.06596110496453271$ |  $0.2060523158$ $\pm$ | $0.018838481328648383$ |  $0.11650425639999999$ $\pm$ | $0.02116471174434527$ |  $0.0497399854$ $\pm$ | $0.011808119227201711$
MolNet |  FreeSolv $\uparrow$ |  $0.9531048546041083$ $\pm$ | $0.009524258414881455$ |  $0.9385841718355937$ $\pm$ | $0.02727459582698729$ |  $0.721769474503605$ $\pm$ | $0.4336237303339043$ |  $0.38568995605894585$ $\pm$ | $0.5099657486714516$ |  $0.9174857803705214$ $\pm$ | $0.006988592017555867$ |  $0.9580144172261275$ $\pm$ | $0.002335062092319676$
lipo $\uparrow$ |  $0.7809456175983587$ $\pm$ | $0.012884532753263647$ |  $0.7814330774573854$ $\pm$ | $0.008761639847206584$ |  $0.7813629873259241$ $\pm$ | $0.008632346125741845$ |  $0.8005897509865727$ $\pm$ | $0.012876418831898969$ | oom |  $0.7054820809251442$ $\pm$ | $0.013484371316680255$
esol $\uparrow$ |  $0.8634920600038448$ $\pm$ | $0.010856213298923612$ |  $0.8526111608874576$ $\pm$ | $0.010152367331517802$ |  $0.8920266832005354$ $\pm$ | $0.008692499186842405$ |  $0.8835049558878794$ $\pm$ | $0.007677406300999873$ |  $0.905310880959911$ $\pm$ | $0.012549082656384585$ |  $0.9289722655120787$ $\pm$ | $0.007363159304838348$
|  | GAT | GATv2 | GIN | PNA | TokenGT | MAGE
dockstring |  esr2 $\uparrow$ |  $0.5736072996030306$ $\pm$ | $0.010357615384481638$ |  $0.5772809421982601$ $\pm$ | $0.006012868454111018$ |  $0.5946087705160163$ $\pm$ | $0.007412942207378858$ |  $0.6123148656735278$ $\pm$ | $0.0035790334805497546$ |  $0.48082527753542$ $\pm$ | $0.009288835200440722$ |  $0.6273723557040729$ $\pm$ | $0.002807812201563067$
f2 $\uparrow$ |  $0.7939308303587385$ $\pm$ | $0.019461587265551934$ |  $0.8284409051621033$ $\pm$ | $0.009716288977250648$ |  $0.8458393492584367$ $\pm$ | $0.0019396860536288645$ |  $0.8540260735639794$ $\pm$ | $0.004303890024604707$ |  $0.769889003336125$ $\pm$ | $0.0034224298844486947$ |  $0.875940831226665$ $\pm$ | $0.002093778354916085$
kit $\uparrow$ |  $0.7992795762836284$ $\pm$ | $0.0033017845953203944$ |  $0.8030903210151739$ $\pm$ | $0.0031166835899172876$ |  $0.8045744869110258$ $\pm$ | $0.005453082900418009$ |  $0.8181661599353148$ $\pm$ | $0.0019733347167171615$ |  $0.659688433717898$ $\pm$ | $0.023577659720417712$ |  $0.8021298321711784$ $\pm$ | $0.0025591732121592967$
parp1 $\uparrow$ |  $0.7880449467163745$ $\pm$ | $0.03895671346748012$ |  $0.8618604660632396$ $\pm$ | $0.009405607476330203$ |  $0.8758152645711214$ $\pm$ | $0.002240762354650523$ |  $0.8828739418040751$ $\pm$ | $0.0020405184378235893$ |  $0.7911301171608788$ $\pm$ | $0.010487693287877693$ |  $0.9097779089664982$ $\pm$ | $0.001520146385327911$
pgr $\uparrow$ |  $0.496140731201996$ $\pm$ | $0.02414531721731501$ |  $0.49665918194535197$ $\pm$ | $0.031790358166305356$ |  $0.5611367832006834$ $\pm$ | $0.015914503889233444$ |  $0.5008834417399128$ $\pm$ | $0.0607417435049365$ |  $0.5040454252623364$ $\pm$ | $0.0028443808644646523$ |  $0.6789017459172216$ $\pm$ | $0.002390602986508577$
Table 4: Test set Matthews correlation coefficient (MCC) for 3 graph-level classification tasks from MoleculeNet, presented as mean $\pm$ standard deviation from 5 different runs. The highest mean values are highlighted in bold. GCN and TokenGT are available in Appendix D. Dataset | GAT | GATv2 | GIN | PNA | Graphormer | MAGE
---|---|---|---|---|---|---
bbbp $\uparrow$ |  $0.710335087776184$ $\pm$ | $0.04215453286433914$ |  $0.7178290486335754$ $\pm$ | $0.04346931022648294$ |  $0.712753438949585$ $\pm$ | $0.01521694420520734$ |  $0.720511519908905$ $\pm$ | $0.026699789297735827$ | oom |  $0.7616158366203308$ $\pm$ | $0.05489501868097872$
bace $\uparrow$ |  $0.5934321761131287$ $\pm$ | $0.024780634954809772$ |  $0.5952930212020874$ $\pm$ | $0.010638616273522814$ |  $0.5640460014343261$ $\pm$ | $0.05268601881929813$ |  $0.44834697172045707$ $\pm$ | $0.2522901530983836$ |  $0.09102762341499324$ $\pm$ | $0.06863972263940434$ |  $0.6456059455871582$ $\pm$ | $0.022338364641736597$
hiv $\uparrow$ |  $0.3827834963798523$ $\pm$ | $0.02379494629458785$ |  $0.3776569187641144$ $\pm$ | $0.06699457269644934$ |  $0.3464722990989685$ $\pm$ | $0.04870448044807584$ |  $0.3776120781898498$ $\pm$ | $0.05225855561999408$ | oom |  $0.4266544163227081$ $\pm$ | $0.023232756631788004$

dockstring (García-Ortegón et al., 2022) is a recent drug discovery data
collection consisting of molecular docking scores for 260,155 small molecules
and 5 high-quality targets from different protein families that were selected
as a regression benchmark, with different levels of difficulty: parp1 (enzyme,
easy), f2 (protease, easy to medium), kit (kinase, medium), esr2 (nuclear
receptor, hard), and pgr (nuclear receptor, hard). The tasks are expected to
be challenging, as the docking score depends on the 3D structure of the
ligand–target complex. Furthermore, a cluster split into train (221,274
molecules) and test (38,881 molecules) sets is provided, ensuring a more
difficult and meaningful benchmark. We have selected a random subset of 19,993
molecules from the original train set as a validation set (the test set is not
modified). We report results for the 5 targets in Table 3 and observe that
MAGE is the strongest method for 4 of the tasks. Remarkably, MAGE matches or
outperforms the strongest methods in the original manuscript (Attentive FP, a
GNN based on attention, Xiong et al., 2020b) despite using 20,000 less
training molecules. Moreover, while the other methods are competitive for the
4 easy or medium difficulty tasks, MAGE outperforms the others by a large
margin for the most difficult target (pgr).

Table 5: Test set Matthews correlation coefficient (MCC) for graph-level classification tasks from various domains, presented as mean $\pm$ standard deviation from 5 runs. All models use MAGE, except the last 6 (MAGN). The highest mean values are highlighted in bold. | Dataset | GCN | GAT | GATv2 | GIN | PNA | MAG
---|---|---|---|---|---|---|---
|  MalNetTiny $\uparrow$ |  $0.846909726$ $\pm$ | $0.008525725$ |  $0.872072268$ $\pm$ | $0.012074483$ |  $0.875785911$ $\pm$ | $0.011278045$ |  $0.900038433$ $\pm$ | $0.00788677$ |  $0.88843044$ $\pm$ | $0.009356888$ |  $0.914639081$ $\pm$ | $0.010521251$
cv |  mnist $\uparrow$ |  $0.8547129271999999$ $\pm$ | $0.0020944243482163066$ |  $0.9705775261999999$ $\pm$ | $0.002569258811755584$ |  $0.9725290776$ $\pm$ | $0.0027878071272239443$ |  $0.8830966234$ $\pm$ | $0.007055951813629732$ |  $0.9759323238000001$ $\pm$ | $0.003663654381721627$ |  $0.9698804498$ $\pm$ | $0.003570527120474143$
cifar10 $\uparrow$ |  $0.45687335119999994$ $\pm$ | $0.0031158203098276783$ |  $0.6261574628$ $\pm$ | $0.007951211539509962$ |  $0.6300276398$ $\pm$ | $0.005195035094275953$ |  $0.4413229824$ $\pm$ | $0.01876954052486123$ |  $0.6500860212$ $\pm$ | $0.01639014671601997$ |  $0.6284652112$ $\pm$ | $0.010776121687335799$
BioInf. |  enzymes $\uparrow$ |  $0.6137649178$ $\pm$ | $0.028953430107558742$ |  $0.697439802$ $\pm$ | $0.0425950784319399$ |  $0.7155687928$ $\pm$ | $0.0274472639075072$ |  $0.6667206884$ $\pm$ | $0.04484606091874496$ |  $0.7269003274$ $\pm$ | $0.023700107643505487$ |  $0.7111327886581421$ $\pm$ | $0.03195198919384762$
proteins $\uparrow$ |  $0.20661196340000001$ $\pm$ | $0.10037970927235781$ |  $0.41280053260000005$ $\pm$ | $0.040848791883428645$ |  $0.4670306088$ $\pm$ | $0.04973596073654339$ |  $0.482513249$ $\pm$ | $0.07575550117555234$ |  $0.43682518579999996$ $\pm$ | $0.10321250850444982$ |  $0.5457976044$ $\pm$ | $0.06429192768767972$
|  dd $\uparrow$ |  $0.488106388$ $\pm$ | $0.016628777627086536$ |  $0.457036662$ $\pm$ | $0.05396712135324707$ |  $0.42737657440000004$ $\pm$ | $0.043783968668912966$ |  $0.3896193088$ $\pm$ | $0.0659628090371931$ |  $0.402002281$ $\pm$ | $0.1049344235010979$ |  $0.569248$ $\pm$ | $0.039797$
synthetic |  synthetic $\uparrow$ |  $1.0$ $\pm$ | $0.0$ |  $1.0$ $\pm$ | $0.0$ |  $1.0$ $\pm$ | $0.0$ |  $0.986970532$ $\pm$ | $0.029134776158658206$ |  $1.0$ $\pm$ | $0.0$ |  $1.0$ $\pm$ | $0.0$
synthetic n. $\uparrow$ |  $-0.010243167800000001$ $\pm$ | $0.05222915363540241$ |  $0.56139161$ $\pm$ | $0.11801431281725723$ |  $0.794536948$ $\pm$ | $0.11268526842132642$ |  $0.4688645662$ $\pm$ | $0.140754049353754$ |  $0.9100602031999999$ $\pm$ | $0.0339484696133309$ |  $0.947882128$ $\pm$ | $0.029134776158658206$
synthie $\uparrow$ |  $0.8686718224$ $\pm$ | $0.03712171558885936$ |  $0.2639232814$ $\pm$ | $0.06646334093797508$ |  $0.3044165074$ $\pm$ | $0.06890456724332734$ |  $0.5961084725999999$ $\pm$ | $0.06737753673731017$ |  $0.8362539290000001$ $\pm$ | $0.08356585650644709$ |  $0.9157447099685669$ $\pm$ | $0.05408323631022407$
triangles $\uparrow$ |  $0.169029191$ $\pm$ | $0.012335114421629058$ |  $0.15319949400000002$ $\pm$ | $0.02246838983612614$ |  $0.13770303579999998$ $\pm$ | $0.0242579583217896$ |  $0.14511427859999998$ $\pm$ | $0.012863146383130604$ |  $0.066064243$ $\pm$ | $0.015174058327148294$ |  $0.22003355940000002$ $\pm$ | $0.05375818377725348$
colors-3 $\uparrow$ |  $0.3049089848$ $\pm$ | $0.01315279837938683$ |  $0.21642020639999998$ $\pm$ | $0.0170262764692737$ |  $0.2513988822$ $\pm$ | $0.007482215046828316$ |  $0.289394337$ $\pm$ | $0.01744351438102221$ |  $0.382303095$ $\pm$ | $0.011124402837864485$ |  $0.7511434912$ $\pm$ | $0.00851483235893665$
social |  imdb-binary $\uparrow$ |  $0.5944314598$ $\pm$ | $0.04023574610864177$ |  $0.5248554110000001$ $\pm$ | $0.045355310307425754$ |  $0.5125346838$ $\pm$ | $0.049321961947283825$ |  $0.540641123$ $\pm$ | $0.05117334349724364$ |  $0.5015860200000001$ $\pm$ | $0.04656045750645811$ |  $0.6249738216$ $\pm$ | $0.06255087528952709$
imdb-multi $\uparrow$ |  $0.1756723674$ $\pm$ | $0.025226220311375915$ |  $0.19813333760000001$ $\pm$ | $0.015259434558779875$ |  $0.1724162996$ $\pm$ | $0.014678983004899749$ |  $0.1901872516$ $\pm$ | $0.018567491956091756$ |  $0.1981180162$ $\pm$ | $0.018781910798294752$ |  $0.21063815959999999$ $\pm$ | $0.03175813395155078$
reddit-binary $\uparrow$ |  $0.536093545$ $\pm$ | $0.023082116507776085$ |  $0.534048474$ $\pm$ | $0.0503704088624253$ |  $0.34961970119999997$ $\pm$ | $0.20214788829145072$ |  $0.49128948440000003$ $\pm$ | $0.03962042740511072$ |  $0.534781283$ $\pm$ | $0.06688766596191131$ |  $0.6063191891999999$ $\pm$ | $0.041312515221784$
reddit-m-5k $\uparrow$ |  $0.3436703028$ $\pm$ | $0.013897738526435775$ |  $0.31345064040000004$ $\pm$ | $0.008257567667880926$ |  $0.29100785240000004$ $\pm$ | $0.02465313904292598$ |  $0.3094648836$ $\pm$ | $0.01098135709150021$ |  $0.34007665519999997$ $\pm$ | $0.018874048069614558$ |  $0.349989283$ $\pm$ | $0.01432880054554965$
reddit-m-12k $\uparrow$ |  $0.3409200252$ $\pm$ | $0.006357959602463455$ |  $0.3068252622$ $\pm$ | $0.014207963154336117$ |  $0.2616211712$ $\pm$ | $0.015011991575280593$ |  $0.3298817932$ $\pm$ | $0.004727645815888881$ |  $0.3527527452$ $\pm$ | $0.011675878673556775$ |  $0.368852693$ $\pm$ | $0.028065668481956213$
twitch egos $\uparrow$ |  $0.374348432$ $\pm$ | $0.0025100871441150205$ |  $0.38107041119999996$ $\pm$ | $0.003241664993182388$ |  $0.384127599$ $\pm$ | $0.002325862404661776$ |  $0.3812508702$ $\pm$ | $0.003689127238540532$ |  $0.3891746398$ $\pm$ | $0.0034478945694306332$ |  $0.39302604198455804$ $\pm$ | $0.001559818067394581$
reddit thr. $\uparrow$ |  $0.5628166316$ $\pm$ | $0.002138440050109315$ |  $0.565382278$ $\pm$ | $0.0011997262988240453$ |  $0.5655566573999999$ $\pm$ | $0.0014248534528123677$ |  $0.5668736818$ $\pm$ | $0.0014438245472725333$ |  $0.566987765$ $\pm$ | $0.0010541922459217413$ |  $0.5673637152$ $\pm$ | $0.001057896435324195$
github star. $\uparrow$ |  $0.2742180704$ $\pm$ | $0.011319049212392648$ |  $0.3019396722$ $\pm$ | $0.006512505704881434$ |  $0.2050568126$ $\pm$ | $0.11128871464281592$ |  $0.28256318560000004$ $\pm$ | $0.025061307658719514$ |  $0.3550015092$ $\pm$ | $0.02449082649516103$ |  $0.31061212420463563$ $\pm$ | $0.018235945976027877$

We further extend our evaluation with a collection of datasets that covers
multiple domains such as bioinformatics, physical chemistry, computer vision,
social networks, functional call graphs, and synthetic graphs (Tables 3, 4 and
5). We observe that MAGE and MAGN are competitive with the other methods and
in most cases better.

###  5.4 Transfer learning

We follow the recipe recently outlined for drug discovery and quantum
mechanics by Buterez et al. (2024). Here, we leverage a recently-published
refined version of the qm9 HOMO and LUMO energies (Fediai et al., 2023), which
provides alternative DFT calculations based on the correlation-consistent
basis set aug-cc-DZVP and the PBE functional, as well as calculations at the
(more accurate) eigen-value-self-consistent GW level of theory. The transfer
learning setup consists of randomly selected training, validation, and test
sets of 25K, 5K, and respectively 10K molecules with GW calculations (from the
total 133,885). As outlined by Buterez et al. (2024), transfer learning can be
performed in transductive or inductive setups. In the transductive case, test
set DFT-level measurements are used for pre-training, while in the inductive
setting they are not. Here, we perform transfer learning by pre-training a
model on the DFT target for a fixed number of epochs and then fine-tuning it
on the subset of 25K GW calculations. In the transductive case, pre-training
occurs on the full set of 133K DFT calculations, while in the inductive case
the DFT test set values are removed (note that the evaluation is done on the
test set GW values). The results (Table 6) indicate that MAGE improves thanks
to transfer learning from DFT by 45% (HOMO) and 53% (LUMO) in the challenging
inductive case, with 10 to 20-fold improvements for the transductive case,
while GNNs improve only by a modest amount.

Table 6: Transfer learning performance (RMSE) on qm9 for HOMO and LUMO, presented as mean $\pm$ standard deviation from 5 different runs on test sets. The lowest mean values are highlighted in bold. Task | Strategy | GCN | GAT | GATv2 | GIN | PNA | MAGE
---|---|---|---|---|---|---|---
HOMO $\downarrow$ | GW |  $0.23012437217652812$ $\pm$ | $0.0036644964069379926$ |  $0.20651051831717374$ $\pm$ | $0.0016182177197634894$ |  $0.20952276285062502$ $\pm$ | $0.0018187021002665873$ |  $0.2075313080347611$ $\pm$ | $0.0014948022604596744$ |  $0.18640467009916634$ $\pm$ | $0.001368160082676368$ |  $0.1365047604336419$ $\pm$ | $0.0011043899933781897$
Trans. |  $0.21299290219588657$ $\pm$ | $0.0009576881906116281$ |  $0.16066327355610976$ $\pm$ | $0.000310737250238902$ |  $0.14859657993629538$ $\pm$ | $0.00012061929507060947$ |  $0.15349459886780573$ $\pm$ | $0.00041719234129667006$ |  $0.1413485030779497$ $\pm$ | $0.00008595460828917416$ |  $0.0070027358056375$ $\pm$ | $0.00006149382247068463$
Ind. |  $0.21080607846682559$ $\pm$ | $0.0006148872602508036$ |  $0.1783826508929181$ $\pm$ | $0.00035664180333540887$ |  $0.17846199018556672$ $\pm$ | $0.0003231695488317062$ |  $0.185292948351458$ $\pm$ | $0.00016744940107373682$ |  $0.1680676730378752$ $\pm$ | $0.00012432225351275252$ |  $0.09431803641576318$ $\pm$ | $0.00004462601399396683$
LUMO $\downarrow$ | GW |  $0.1959630929092371$ $\pm$ | $0.001449565786176188$ |  $0.19135032379881994$ $\pm$ | $0.0015673427568889085$ |  $0.19496575231336818$ $\pm$ | $0.002323548304504633$ |  $0.18930690164914396$ $\pm$ | $0.0011337092693161914$ |  $0.18313335941652445$ $\pm$ | $0.0027778823016168575$ |  $0.12108962371898213$ $\pm$ | $0.0009430255084660352$
Trans. |  $0.20770791846990672$ $\pm$ | $0.000793553971470199$ |  $0.16424414608473253$ $\pm$ | $0.000287053259700656$ |  $0.1652769871736944$ $\pm$ | $0.0002741973349826362$ |  $0.16683475696390956$ $\pm$ | $0.00025076375870373936$ |  $0.15923357020923748$ $\pm$ | $0.0002439859152211106$ |  $0.0102717764677073$ $\pm$ | $0.00031589308194151453$
Ind. |  $0.19803494437200878$ $\pm$ | $0.00048666208185973916$ |  $0.171031077463694$ $\pm$ | $0.0004820261332570306$ |  $0.17305384996190526$ $\pm$ | $0.00018461070653856605$ |  $0.17187033320690528$ $\pm$ | $0.0005931161197957346$ |  $0.16525529430581326$ $\pm$ | $0.000628422160929613$ |  $0.07899072113644254$ $\pm$ | $0.00019256848906605754$
Table 7: Average training time per epoch (s) and used memory (GB) for all methods, presented as mean $\pm$ std. from 5 epochs and with the number of parameters (#). For qm9 (103,542 train items), the maximum number of nodes/edges per graph is 29 (MAGN), respectively 56 (MAGE). DD (942 train items) has 5,748 maximum nodes and 28,534 maximum edges. All algorithms use bfloat16 mixed training. (*) Parameters for masked blocks are counted as for normal self-attention, despite a very large number of connections being dropped. Method |  # | qm9 | dd
---|---|---|---
|  | Time (s) | Memory (GB) | Time (s) | Memory (GB)
GCN | 158K |  $13.741277$ $\pm$ | $0.474285$ |  $0.09605$ $\pm$ | $0.000938$ |  $2.196680$ $\pm$ | $0.107517$ |  $0.322542$ $\pm$ | $0.024475$
GAT | 10.1M |  $25.603818$ $\pm$ | $0.350959$ |  $1.142109$ $\pm$ | $0.073678$ |  $4.050491$ $\pm$ | $0.077201$ |  $3.614651$ $\pm$ | $0.550235$
GATv2 | 20.1M |  $28.817296$ $\pm$ | $0.65316$ |  $1.357277$ $\pm$ | $0.085771$ |  $5.309763$ $\pm$ | $0.035361$ |  $5.514671$ $\pm$ | $0.669779$
GIN | 433K |  $14.953487$ $\pm$ | $1.242019$ |  $0.119118$ $\pm$ | $0.002297$ |  $2.129896$ $\pm$ | $0.032087$ |  $0.282696$ $\pm$ | $0.058784$
PNA | 6.9M |  $66.61936$ $\pm$ | $1.326496$ |  $2.529585$ $\pm$ | $0.196117$ |  $12.053648$ $\pm$ | $0.045797$ |  $6.773801$ $\pm$ | $0.966106$
Graphormer | 22.3M |  $186.597434$ $\pm$ | $4.498434$ |  $1.952449$ $\pm$ | $0$ | oom
TokenGT | 8.1M |  $30.763488$ $\pm$ | $0.257696$ |  $1.214251$ $\pm$ | $0$ | oom
MAGN (naive) | 8.3M* |  $16.457141$ $\pm$ | $0.708097$ |  $0.314573$ $\pm$ | $0$ | oom
MAGN (mem-efficient) |  $15.379194$ $\pm$ | $0.703972$ |  $0.268435$ $\pm$ | $0$ |  $13.532658$ $\pm$ | $0.032929$ |  $1.193279$ $\pm$ | $0.000000$
MAGE (naive) | 8.6M* |  $26.11023$ $\pm$ | $0.306541$ |  $0.612368$ $\pm$ | $0.004689$ | oom
MAGE (mem-efficient) |  $22.580969$ $\pm$ | $0.486256$ |  $0.442499$ $\pm$ | $0$ |  $348.221254$ $\pm$ | $1.321884$ |  $21.899719$ $\pm$ | $1.194906$

###  5.5 Time and memory utilisation

In MAG, the most computation-intensive component is the encoder. All encoder
blocks perform (masked) self attention, which can be efficiently implemented
with $\mathcal{O}(\sqrt{N})$ memory complexity ($N_{n}$, $N_{e}$ for MAGN,
respectively MAGE). The time complexity is as for standard attention. The
decoder PMA encodes cross-attention between the full set outputted by the
encoder and a set of learnable $k$ vectors, benefitting from Flash attention
(linear memory and time scaling). A decoder with a single PMA block and $k=1$,
as used here, is even more efficient. The mask tensor requires $B\times N^{2}$
memory (Appendix A for a solution). However, it does not require gradients and
MAGE runs with up to $N_{e}\approx 30,000$ edges on a consumer GPU with 24GB.
Indeed, we report competitive time and memory utilisation (Table 7).

##  6 Discussion

We presented an end-to-end approach that leverages attention in a novel way
for learning on graphs and demonstrated its effectiveness relative to message
passing and more involved attention-based methods. Our approach is end-to-end
in the sense that it replaces both message passing and readout/pooling
functions with attention mechanisms. The former is facilitated by modularly
stacking self and masked attention blocks (node or edge-based masking). In the
future, node- and edge-masked blocks might be interspersed. Both node and edge
masked attention mechanisms can be implemented using a few lines of code
(Algorithms 1 and 4). Despite its simplicity, MAG consistently outperforms
other message passing baselines and more involved Transformer-based methods
while supporting transfer learning through pre-training and fine-tuning, which
does not work well with classical GNNs, and scales favourably in terms of time
and memory. Still, several software limitations impede a more efficient MAG
implementation (see Appendix A). Compared to Transformers, we remarked several
distinguishing features. It performs as well as presented here without any
sophisticated learning rate scheduler (warm-up, cosine annealing, etc.) or
optimiser. Interestingly, for about half the datasets MAG performed better
with batch normalisation instead of layer normalisation (more in Appendix G).
Also, MAG does not use any form of positional encoding, which is against
current trends for GNNs and Transformers. Future research might investigate
ways to combine node and edge processing, devise efficient attention
implementations for graphs, study the importance of positional encodings,
integrate successful concepts from other domains such as sparse expert models
(Fedus et al., 2022), or study multi-modality.

## 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.
  * Brody et al. (2022) Brody, S., Alon, U., and Yahav, E.  How attentive are graph attention networks?  In _International Conference on Learning Representations_ , 2022.
  * Buterez et al. (2022) Buterez, D., Janet, J. P., Kiddle, S. J., Oglic, D., and Liò, P.  Graph neural networks with adaptive readouts.  In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), _Advances in Neural Information Processing Systems_ , 2022.
  * Buterez et al. (2023a) Buterez, D., Janet, J. P., Kiddle, S. J., and Liò, P.  Mf-pcba: Multifidelity high-throughput screening benchmarks for drug discovery and machine learning.  _Journal of Chemical Information and Modeling_ , 63(9):2667–2678, 2023a.  doi: 10.1021/acs.jcim.2c01569.  PMID: 37058588.
  * Buterez et al. (2023b) Buterez, D., Janet, J. P., Kiddle, S. J., Oglic, D., and Liò, P.  Modelling local and general quantum mechanical properties with attention-based pooling.  _Communications Chemistry_ , 6(1):262, Nov 2023b.  ISSN 2399-3669.  doi: 10.1038/s42004-023-01045-7.
  * Buterez et al. (2024) Buterez, D., Janet, J. P., Kiddle, S., Oglic, D., and Liò, P.  Transfer learning with graph neural networks for improved molecular property prediction in the multi-fidelity setting.  _ChemRxiv_ , 2024.  doi: 10.26434/chemrxiv-2022-dsbm5-v3.
  * Cai et al. (2021) Cai, T., Luo, S., Xu, K., He, D., Liu, T.-Y., and Wang, L.  Graphnorm: A principled approach to accelerating graph neural network training.  In _2021 International Conference on Machine Learning_ , May 2021.
  * Chen et al. (2021) Chen, C., Zuo, Y., Ye, W., Li, X., and Ong, S. P.  Learning properties of ordered and disordered materials from multi-fidelity data.  _Nature Computational Science_ , 1(1):46–53, 01 2021.  ISSN 2662-8457.
  * Choromanski et al. (2021) Choromanski, K. M., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., Hawkins, P., Davis, J. Q., Mohiuddin, A., Kaiser, L., Belanger, D. B., Colwell, L. J., and Weller, A.  Rethinking attention with performers.  In _International Conference on Learning Representations_ , 2021.
  * Corso et al. (2020) Corso, G., Cavalleri, L., Beaini, D., Liò, P., and Velickovic, P.  Principal neighbourhood aggregation for graph nets.  In _Proceedings of the 34th International Conference on Neural Information Processing Systems_ , NIPS’20, Red Hook, NY, USA, 2020. Curran Associates Inc.  ISBN 9781713829546.
  * Dao (2023) Dao, T.  Flashattention-2: Faster attention with better parallelism and work partitioning, 2023.
  * Dao et al. (2022) Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Ré, C.  FlashAttention: Fast and memory-efficient exact attention with IO-awareness.  In _Advances in Neural Information Processing Systems_ , 2022.
  * Duvenaud et al. (2015) Duvenaud, D., Maclaurin, D., Aguilera-Iparraguirre, J., Gómez-Bombarelli, R., Hirzel, T., Aspuru-Guzik, A., and Adams, R. P.  Convolutional networks on graphs for learning molecular fingerprints.  In _Proceedings of the 28th International Conference on Neural Information Processing Systems - Volume 2_ , NIPS’15, pp. 2224–2232, Cambridge, MA, USA, 2015. MIT Press.
  * Dwivedi et al. (2022) Dwivedi, V. P., Rampášek, L., Galkin, M., Parviz, A., Wolf, G., Luu, A. T., and Beaini, D.  Long range graph benchmark.  In _Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track_ , 2022.
  * Fediai et al. (2023) Fediai, A., Reiser, P., Peña, J. E. O., Friederich, P., and Wenzel, W.  Accurate gw frontier orbital energies of 134 kilo molecules.  _Scientific Data_ , 10(1):581, Sep 2023.  ISSN 2052-4463.  doi: 10.1038/s41597-023-02486-4.
  * Fedus et al. (2022) Fedus, W., Dean, J., and Zoph, B.  A review of sparse expert models in deep learning, 2022.
  * García-Ortegón et al. (2022) García-Ortegón, M., Simm, G. N. C., Tripp, A. J., Hernández-Lobato, J. M., Bender, A., and Bacallado, S.  Dockstring: Easy molecular docking yields better benchmarks for ligand design.  _Journal of Chemical Information and Modeling_ , 62(15):3486–3502, 2022.  doi: 10.1021/acs.jcim.1c01334.  PMID: 35849793.
  * 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 _Proceedings of the 34th International Conference on Machine Learning - Volume 70_ , ICML’17, pp. 1263–1272. JMLR.org, 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.
  * Hu et al. (2020) Hu, W., Liu, B., Gomes, J., Zitnik, M., Liang, P., Pande, V., and Leskovec, J.  Strategies for pre-training graph neural networks.  In _International Conference on Learning Representations_ , 2020.
  * Jain et al. (2021) Jain, P., Wu, Z., Wright, M. A., Mirhoseini, A., Gonzalez, J. E., and Stoica, I.  Representing long-range context for graph neural networks with global attention.  In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), _Advances in Neural Information Processing Systems_ , 2021.
  * Kearnes et al. (2016) Kearnes, S., McCloskey, K., Berndl, M., Pande, V., and Riley, P.  Molecular graph convolutions: moving beyond fingerprints.  _Journal of Computer-Aided Molecular Design_ , 30(8):595–608, Aug 2016.  ISSN 1573-4951.  doi: 10.1007/s10822-016-9938-8.
  * Kim et al. (2022) Kim, J., Nguyen, D. T., Min, S., Cho, S., Lee, M., Lee, H., and Hong, S.  Pure transformers are powerful graph learners.  In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), _Advances in Neural Information Processing Systems_ , 2022.
  * Kreuzer et al. (2021) Kreuzer, D., Beaini, D., Hamilton, W. L., Létourneau, V., and Tossou, P.  Rethinking graph transformers with spectral attention.  In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), _Advances in Neural Information Processing Systems_ , 2021.
  * Lee et al. (2019) Lee, J., Lee, Y., Kim, J., Kosiorek, A., Choi, S., and Teh, Y. W.  Set transformer: A framework for attention-based permutation-invariant neural networks.  In _Proceedings of the 36th International Conference on Machine Learning_ , pp. 3744–3753, 2019.
  * Lefaudeux et al. (2022) Lefaudeux, B., Massa, F., Liskovich, D., Xiong, W., Caggiano, V., Naren, S., Xu, M., Hu, J., Tintore, M., Zhang, S., Labatut, P., and Haziza, D.  xformers: A modular and hackable transformer modelling library.  https://github.com/facebookresearch/xformers, 2022.
  * Loshchilov & Hutter (2019) Loshchilov, I. and Hutter, F.  Decoupled weight decay regularization.  In _International Conference on Learning Representations_ , 2019.
  * Paszke et al. (2019) Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S.  Pytorch: An imperative style, high-performance deep learning library.  In _Advances in Neural Information Processing Systems 32_ , pp. 8024–8035. Curran Associates, Inc., 2019.
  * Rabe & Staats (2021) Rabe, M. N. and Staats, C.  Self-attention does not need o(n${}^{\mbox{2}}$) memory.  _CoRR_ , abs/2112.05682, 2021.
  * 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):140022, Aug 2014.  ISSN 2052-4463.  doi: 10.1038/sdata.2014.22.
  * Rampasek et al. (2022) Rampasek, L., Galkin, M., Dwivedi, V. P., Luu, A. T., Wolf, G., and Beaini, D.  Recipe for a general, powerful, scalable graph transformer.  In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), _Advances in Neural Information Processing Systems_ , 2022.
  * 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.  _Advances in Neural Information Processing Systems_ , 35, 2022.
  * Shazeer (2020) Shazeer, N.  GLU variants improve transformer.  _CoRR_ , abs/2002.05202, 2020.
  * Shirzad et al. (2023) Shirzad, H., Velingker, A., Venkatachalam, B., Sutherland, D. J., and Sinop, A. K.  Exphormer: Scaling graph transformers with expander graphs, 2023.
  * Tönshoff et al. (2023) Tönshoff, J., Ritzert, M., Rosenbluth, E., and Grohe, M.  Where did the gap go? reassessing the long-range graph benchmark.  In _The Second Learning on Graphs Conference_ , 2023.
  * Vaswani et al. (2017) Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L. u., and Polosukhin, I.  Attention is all you need.  In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems_ , volume 30. Curran Associates, Inc., 2017.
  * Veličković et al. (2018) Veličković, P., Cucurull, G., Casanova, A., Romero, A., Liò, P., and Bengio, Y.  Graph attention networks.  In _International Conference on Learning Representations_ , 2018.
  * Wolf et al. (2020) Wolf, T., Debut, L., Sanh, V., Chaumond, J., Delangue, C., Moi, A., Cistac, P., Rault, T., Louf, R., Funtowicz, M., Davison, J., Shleifer, S., von Platen, P., Ma, C., Jernite, Y., Plu, J., Xu, C., Le Scao, T., Gugger, S., Drame, M., Lhoest, Q., and Rush, A.  Transformers: State-of-the-art natural language processing.  In Liu, Q. and Schlangen, D. (eds.), _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations_ , pp. 38–45, Online, October 2020. Association for Computational Linguistics.  doi: 10.18653/v1/2020.emnlp-demos.6.
  * Xiong et al. (2020a) Xiong, R., Yang, Y., He, D., Zheng, K., Zheng, S., Xing, C., Zhang, H., Lan, Y., Wang, L., and Liu, T.-Y.  On layer normalization in the transformer architecture.  In _Proceedings of the 37th International Conference on Machine Learning_ , ICML’20. JMLR.org, 2020a.
  * Xiong et al. (2020b) Xiong, Z., Wang, D., Liu, X., Zhong, F., Wan, X., Li, X., Li, Z., Luo, X., Chen, K., Jiang, H., and Zheng, M.  Pushing the boundaries of molecular representation for drug discovery with the graph attention mechanism.  _Journal of Medicinal Chemistry_ , 63(16):8749–8760, 2020b.  doi: 10.1021/acs.jmedchem.9b00959.  PMID: 31408336.
  * Ying et al. (2021) 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 _Thirty-Fifth Conference on Neural Information Processing Systems_ , 2021.
  * Zaheer et al. (2020) Zaheer, M., Guruganesh, G., Dubey, K. A., Ainslie, J., Alberti, C., Ontanon, S., Pham, P., Ravula, A., Wang, Q., Yang, L., and Ahmed, A.  Big bird: Transformers for longer sequences.  In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), _Advances in Neural Information Processing Systems_ , volume 33, pp. 17283–17297. Curran Associates, Inc., 2020.
  * Zhao & Akoglu (2020) Zhao, L. and Akoglu, L.  Pairnorm: Tackling oversmoothing in gnns.  In _International Conference on Learning Representations_ , 2020.

##  Appendix A Limitations

In terms of limitations, we highlight that the available libraries are not
optimised for masking or custom attention patterns. This is most evident for
very dense graphs (tens of thousands of edges or more). Memory efficient and
Flash attention are available natively in PyTorch (Paszke et al., 2019)
starting from version 2, as well as in the xFormers library (Lefaudeux et al.,
2022). More specifically, we have tested at least 5 different implementations
of MAG: (1) leveraging the MultiheadAttention module from PyTorch, (2)
leveraging the MultiHeadDispatch module from xFormers, (3) a manual
implementation of multihead attention, relying on PyTorch’s
scaled_dot_product_attention function, (4) a manual implementation of
multihead attention, relying on xFormers’ memory_efficient_attention, and (5)
a naive implementation. Options (1) - (4) can all make use of efficient and
fast implementations. However, we have observed performance differences
between the 4 implementations, as well as compared to a naive implementation.
This behaviour is likely due to the different low-level kernel
implementations. Moreover, Flash attention does not support custom attention
masks as there is little interest for such functionality from a language
modelling perspective.

Although the masks can be computed efficiently during training, all frameworks
require the last two dimensions of the input mask tensor to be of shape
$(N_{d},N_{d})$ for nodes or $(N_{e},N_{e})$ for edges, effectively squaring
the number of nodes or edges. However, the mask tensors are very sparse and a
sparse tensor alternative could greatly reduce the memory consumption for
large and dense graphs. Such an option exists for the PyTorch native
attention, but it is currently broken.

Another possible optimisation would be to use nested (ragged) tensors to
represent graphs, since padding is currently necessary to ensure identical
dimensions for attention. A prototype nested tensor attention is available in
PyTorch; however, not all the required operations are supported and converting
between normal and nested tensors is slow.

For all implementations, it is required that the mask tensor is repeated by
the number of attention heads (e.g. 8 or 16). However, a notable bottleneck is
encountered for the MultiheadAttention and MultiHeadDispatch variants
described above, which require that the repeats happen in the batch dimension,
i.e. requiring 3D mask tensors of shape $(B\times H,N,N)$, where $H$ is the
number of heads. The other two efficient implementations require a 4D mask
instead, i.e. $(B,H,N,N)$, where one can use PyTorch’s expand function instead
of repeat. The expand alternative does not use any additional memory, while
repeat requires $\times H$ memory. Note that it is not possible to reshape the
4D tensor created using expand without using additional memory.

##  Appendix B Helper functions

consecutive is a helper function that generates consecutive numbers starting
from 0, with a length specified in its tensor argument as the difference
between adjacent elements, and a second integer argument used for the last
length computation, e.g. consecutive($[1,4,6]$, $10$) $=[0,1,2,0,1,0,1,2,3]$,
and first_unique_index finds the first occurrence of each unique element in
the tensor (sorted), e.g. first_unique_index($[3,2,3,4,2]$) $=[1,0,3]$.

##  Appendix C Experimental setup

We follow a simple and universal experimental protocol to ensure that it is
possible to compare the results of different methods and to evaluate a large
number of datasets with high throughput. We chose a number of reasonable
hyperparameters and settings for all methods, regardless of their nature (GNN
or attention-based). This includes the AdamW optimiser (Loshchilov & Hutter,
2019), learning rate ($0.0001$), batch size ($128$), 32-bit training (without
mixed precision), early stopping with a patience of 30 epochs (100 for the
very small datasets such as freesolv), and gradient clipping (set to the
default value of $0.5$). Furthermore, we used a simple learning rate scheduler
that halved the learning rate if no improvement was encountered for 15 epochs
(half the early stopping patience).

For GNNs, we used 4 graph layers for all algorithms (GCN, GIN, GAT, GATv2,
PNA) and the mean readout function, a node dimension of $64$, and a hidden
dimension for the graph layers of $256$. All layers use batch normalisation.
Settings specific to some algorithms were also given by reasonable defaults,
such as $8$ attention heads for GAT(v2), and $5$ towers for PNA. In certain
instances such as small datasets with dense graphs, the defaults selected
above can lead to out-of-memory errors, even for GNNs, in particular the more
computationally-intensive algorithms such as GAT(v2) or PNA. In such cases, we
lower the settings that were the most likely cause, such as the hidden
dimension or batch size, to the next (lower) power of two. In some cases,
mixed precision training (using the bfloat16 type) was performed as an
alternative if Ampere-class GPUs were available.

For Graphormer and TokenGT, we leverage the huggingface (Wolf et al., 2020)
implementation. Our selection of reasonable defaults include 3 layers, a
hidden and embedding size of $512$, and 8 attention heads, in addition to the
defaults mentioned earlier such as learning rate and batch size. Graphormer in
particular can be difficult to train, and in such cases we reduce the
complexity of the model using the same strategies as above. However, note that
many out-of-memory errors for this family of models are not due to GPU memory,
but RAM; we attempted to use up to 256GB, but conceded if it did not work.

For MAG, the same suite of “general” defaults such as the learning rate and
batch size apply. We also generally follow the same configuration for all
datasets. However, since MAG is a new architecture with many unknowns, we
generally evaluate a small number of variations for each dataset and select
the best one according to the validation metrics. The variations typically
involve choosing MAGN or MAGE, batch or layer normalisation, the number (3 or
4) and order of self-attention and masked self attention blocks (e.g. SMM,
MSMM, etc.), the hidden size ($256$ or $512$) and the number of attention
heads ($8$ or $16$). We generally prefer MAGE in all situations and only
consider MAGN for datasets where MAGE would take a very long time to run or
requires heavy modifications to the default parameters. This is because MAGE
can naturally incorporate edge features and almost always performs better. We
have also found SwiGLU to often be better than plain MLPs (Shazeer, 2020).

Graphormer (and TokenGT to a lower extent) have very expensive pre-processing
steps which require up to hundreds of GBs of storage space to cache
intermediary results. The alternative would be to not use caching; however,
this means that everything must be stored in memory, resulting in almost
immediate crashing. As a further complication, for large datasets suck as
dockstring, Graphormer would run for a few epochs but spontaneously crash,
most likely due to high memory utilisation during training. Combined with the
fact that one epoch took several hours, we have included Graphormer results
only for a minority of datasets.

### Hyperparameter optimisation

Other than the basic filtering described above for MAG, we did not use any
techniques for tuning. In particular, we did not perform hyperparameter
optimisation and have not tuned aspects of the networks such as the optimiser,
learning rate, batch size, dropout, hidden dimensions, etc. We acknowledge
that we are not using optimal parameters for the majority of models. However,
this is also true for MAG, and tuning every model presented here would
dramatically increase the time and resource utilisation (as well as the
financial costs associated with it), defeating the purpose of presenting a
simple yet effective alternative to GNNs.

### Evaluation

Generally, for a self-contained evaluation we split all datasets using a
random 80%, 10%, 10% split for train, validation, and test. The same data
splits are used for the different evaluated algorithms. Some datasets, such as
mnist or dockstring are provided with existing train, test, and optionally
validation splits. If such splits are available through PyTorch Geometric or
from the authors (such as dockstring), we use them and we do not perform our
own random splits. For all datasets and models, we provide results from 5
different runs (seeds).

##  Appendix D Additional results

The missing GCN, TokenGT, and Graphormer results for some of the datasets
presented in the main text are presented below in Tables 8 and 9. These
complete the results from Tables 3 and 4.

Table 8: Test set root mean squared error (RMSE) for qm9 and $\text{R}^{2}$ for the others, for GCN and TokenGT, presented as mean $\pm$ standard deviation from 5 different runs. | Property | GCN | TokenGT
---|---|---|---
qm9 | $\mu$ |  $0.6745743364$ $\pm$ | $0.006933984290360895$ |  $1.0003418700000002$ $\pm$ | $0.004548326300329208$
$\alpha$ |  $3.5084559487999996$ $\pm$ | $0.18430228740797702$ |  $2.1439044564$ $\pm$ | $0.06966659397253458$
$\epsilon_{\text{HOMO}}$ |  $0.14451352680000001$ $\pm$ | $0.004143409926233732$ |  $0.26493893479999997$ $\pm$ | $0.010394688435768408$
$\epsilon_{\text{LUMO}}$ |  $0.1536381066$ $\pm$ | $0.0008669119370972958$ |  $0.4227532568$ $\pm$ | $0.013949440303808642$
$\Delta\epsilon$ |  $0.2132146896$ $\pm$ | $0.0014809216234380235$ |  $0.556362002$ $\pm$ | $0.008453743508148207$
$\langle R^{2}\rangle$ |  $67.737992444$ $\pm$ | $5.718781024859155$ |  $177.69592664$ $\pm$ | $6.394972919890789$
ZPVE |  $0.23126787059999998$ $\pm$ | $0.0010708327919441063$ |  $0.1382764998$ $\pm$ | $0.0031473814200245445$
$U_{0}$ |  $573.31943614$ $\pm$ | $47.41046911878416$ |  $228.487603984$ $\pm$ | $198.09857140407738$
$U$ |  $594.54783412$ $\pm$ | $57.80518261016987$ |  $228.48638613600002$ $\pm$ | $198.09751016260756$
$H$ |  $575.0972620800001$ $\pm$ | $48.66365465597905$ |  $228.48637986600002$ $\pm$ | $198.09751432518846$
$G$ |  $594.07371094$ $\pm$ | $56.999665735383935$ |  $228.49044466$ $\pm$ | $198.10102804379832$
$c_{\text{V}}$ |  $1.339669154$ $\pm$ | $0.2606712003389698$ |  $1.051119233$ $\pm$ | $0.025202749244745203$
$U_{0}^{\text{ATOM}}$ |  $3.989557961$ $\pm$ | $0.03867132439966954$ |  $2.0237579842$ $\pm$ | $0.3017541728303727$
$U^{\text{ATOM}}$ |  $4.1940800798$ $\pm$ | $0.3257170165795751$ |  $2.2179237248$ $\pm$ | $0.186405232032096$
$H^{\text{ATOM}}$ |  $4.0534259776$ $\pm$ | $0.02201026675754947$ |  $2.1613375456$ $\pm$ | $0.09732061561443657$
$G^{\text{ATOM}}$ |  $3.7046056548000004$ $\pm$ | $0.016787459006093523$ |  $1.9144232420000002$ $\pm$ | $0.14262474064714065$
| $A$ |  $1.1679455317999998$ $\pm$ | $0.10861048195386805$ |  $0.008707393599999998$ $\pm$ | $0.014000227948560203$
| $B$ |  $0.28494551220000003$ $\pm$ | $0.01779822904993714$ |  $0.4010663332$ $\pm$ | $5.801902274185263\text{\times}{10}^{-07}$
| $C$ |  $0.2361907992$ $\pm$ | $0.018418575005509703$ |  $0.349094864$ $\pm$ | $2.3628055356701814\text{\times}{10}^{-07}$
MolNet | FreeSolv |  $0.3394724596720835$ $\pm$ | $0.5111978898407118$ |  $0.8565087372182878$ $\pm$ | $0.015122426595436543$
lipo |  $0.7120830620519311$ $\pm$ | $0.009073406642521932$ |  oom
esol |  $0.863572673045691$ $\pm$ | $0.009223431543601559$ |  $0.7776586885826735$ $\pm$ | $0.00772350191729414$
|  | GCN | Graphormer
dockstring | esr2 |  $0.5291013622893503$ $\pm$ | $0.009627268565126702$ |  oom
f2 |  $0.7782939224110621$ $\pm$ | $0.0031419784756594763$ |  oom
kit |  $0.758943373474925$ $\pm$ | $0.0028655571948525142$ |  oom
| parp1 |  $0.8122799639236348$ $\pm$ | $0.0034660345423232525$ |  oom
| pgr |  $0.3648950275940546$ $\pm$ | $0.009090128585157106$ |  oom
Table 9: Test set Matthews correlation coefficient (MCC) for GCN and TokenGT for 3 graph-level classification tasks from MoleculeNet, presented as mean $\pm$ standard deviation, over 5 runs. Dataset | GCN | TokenGT
---|---|---
bace |  $0.3462121233344078$ $\pm$ | $0.31426194734962465$ |  $0.18967240005731578$ $\pm$ | $0.08990886312828858$
bbbp |  $0.6757150411605835$ $\pm$ | $0.02480650400007105$ |  $0.3121381223201751$ $\pm$ | $0.0865077820933659$
hiv |  $0.3923294901847839$ $\pm$ | $0.031157433525451424$ |  oom

##  Appendix E Dataset statistics

We present a summary of all the used datasets, along with their size and the
maximum number of nodes and edges encountered in a graph in the dataset (Table
10). The last two are important as they determine the shape of the mask and of
the inputs for the attention blocks. Technically, we require that the maximum
number of nodes/edges is determined per batch and the tensors to be padded
accordingly. This per-batch maximum is lower than the dataset maximum for most
batches. However, certain operations such as layer norm., if performed over
the last two dimensions, require a constant value. To enable this, we use the
dataset maximum.

Table 10: Summary of used datasets, their size, and the max. number of nodes (N) and edges (E) seen in a graph in the dataset. | Dataset |  Size |  N |  E
---|---|---|---|---
lrgb | pept-struct | $15\,535$ | $444$ | $928$
pept-func | $15\,535$ | $444$ | $928$
node | ppi | $24$ | $3\,480$ | $106\,754$
cora | $1$ | $2\,708$ | $10\,556$
CiteSeer | $1$ | $3\,327$ | $9\,104$
| qm9 | $133\,885$ | $29$ | $56$
MoleculeNet | FreeSolv | $642$ | $44$ | $92$
lipo | $4\,200$ | $216$ | $438$
esol | $1\,128$ | $119$ | $252$
bbbp | $2\,039$ | $269$ | $562$
bace | $1\,513$ | $184$ | $376$
hiv | $41\,127$ | $438$ | $882$
| dockstring | $260\,060$ | $164$ | $342$
| MalNetTiny | $5\,000$ | $4\,994$ | $20\,096$
CV | mnist | $70\,000$ | $75$ | $600$
cifar10 | $60\,000$ | $150$ | $1\,200$
BioInfo | enzymes | $600$ | $126$ | $298$
proteins | $1\,113$ | $620$ | $2\,098$
dd | $1\,178$ | $5\,748$ | $28\,534$
synthetic | synthetic | $300$ | $100$ | $392$
synthetic new | $300$ | $100$ | $396$
synthie | $400$ | $100$ | $424$
triangles | $45\,000$ | $100$ | $396$
colors-3 | $10\,500$ | $200$ | $794$
social | imdb-binary | $1\,000$ | $136$ | $2\,498$
imdb-multi | $1\,500$ | $89$ | $2\,934$
reddit-binary | $2\,000$ | $3\,782$ | $8\,142$
reddit-m-5k | $4\,999$ | $3\,648$ | $9\,566$
reddit-m-12k | $11\,929$ | $3\,782$ | $10\,342$
twitch egos | $127\,094$ | $52$ | $1\,572$
reddit thr. | $203\,088$ | $97$ | $370$
github star. | $12\,725$ | $957$ | $9\,336$

##  Appendix F Masked attention equation

Masked attention can be thought of as a custom attention pattern, which for
graphs was described succinctly by (Shirzad et al., 2023). An adaptation for
masking would be:

$\displaystyle\text{MaskedSelfAttention}(\mathbf{X},\mathbf{M})=$ (7)
$\displaystyle\sum_{j=1}^{h}\mathbf{W}^{j}_{O}\left(\mathbf{W}^{j}_{V}\mathbf{X_{M}}\right)\sigma\left(\left(\mathbf{W}^{j}_{K}\mathbf{X_{M}}\right)^{T}(\mathbf{W}^{j}_{Q}\mathbf{X_{M}})\right)$

where $h$ is the number of attention heads, $\mathbf{X}$ is the input features
matrix, $\mathbf{M}$ is the custom attention pattern/mask, which restricts the
attention to a subset of elements of $\mathbf{X}$, denoted by
$\mathbf{X_{M}}$,
$\mathbf{W}^{j}_{Q},\mathbf{W}^{j}_{K},\mathbf{W}^{j}_{V},\mathbf{W}^{j}_{O}$
are weight matrices corresponding to queries, keys, values, and outputs,
respectively, and $\sigma$ is the softmax function.

##  Appendix G Layer vs batch normalisation

Contrary to standard Transformers and current trends, we have found that
simply replacing the layer normalisation (LN) operation within MAG with batch
normalisation (BN) can dramatically improve performance on a number of
datasets. LN was our default initial choice and we have found that it works
better for the standard qm9 dataset and its properties (Table 3). For the
other datasets and tasks (excluding the 19 qm9 properties but including the
qm9 GW tasks from Table 6 and the pept-struct and pept-fn datasets), although
we have not exhaustively tested LN vs BN models, we have generally observed
that LN is preferable for 20 tasks and BN for 17. More specifically, BN was
preferable for the dockstring properties parp1, esr2, and pgr, as well as for
the datasets: cifar10, colors-3, dd, esol, lipo, pept-struct, pept-fn,
proteins, reddit threads, reddit-multi-12k, synthie, triangles, enzymes, and
twitch egos. Apart from the fact that most molecular datasets tend to do
better with LN, there is no obvious indication of which normalisation
technique might be preferable for certain dataset types.

##  Appendix H Experimental platform

Representative versions of the software used as part of this paper include
Python 3.11.6, PyTorch version 2.1.1 with CUDA 11.8, PyTorch Geometric 2.4.0,
PyTorch Lightning 1.9.5, huggingface transformers version 4.35.2, and xFormers
version 0.0.23. We have also tested our code with CUDA $\geq$ 12.0. It is
worth noting that attention masking and efficient implementations of attention
are early features that are advancing quickly. This means that their behaviour
might change unexpectedly and there might be bugs. For example, PyTorch 2.1.1
recently fixed a bug that concerned non-contiguous custom attention masks in
the scaled_dot_product_attention function.

In terms of hardware, the GPUs used include an NVIDIA RTX 3090 with 24GB VRAM,
NVIDIA V100 with 16GB or 32GB of VRAM, and NVIDIA A100 with 40GB of VRAM.
Recent, efficient implementations of attention are optimised for the newest
GPU architectures, generally starting from Ampere (RTX 3090 and A100).
However, while slower, it is possible to run memory efficient attention on
V100 GPUs.

##  Appendix I Multiple configurations of MAGE

We summarise the results of running MAGE on the qm9 property $\alpha$ (alpha)
in Table 11. Compared to the main results of Table 3, we used an early
stopping patience of 10 epochs and trained for a maximum of 150 epochs.

Table 11: Test set RMSE for MAGE configurations for $\alpha$ (qm9), presented as mean $\pm$ standard deviation from 5 different runs. Configuration | MAGE
---|---
MMMSP |  $0.49865540843895706$ $\pm$ | $0.011639542191863255$
MMSSP |  $0.5012473834837001$ $\pm$ | $0.007484177491151632$
MMSP |  $0.5026090505779304$ $\pm$ | $0.010119298682555934$
MMSMP |  $0.5055836816731379$ $\pm$ | $0.003292502880082702$
MSMMP |  $0.5076063884402546$ $\pm$ | $0.00981395687917195$
MMMMP |  $0.5086624586894055$ $\pm$ | $0.006953097409000435$
MSSMP |  $0.5087148590628632$ $\pm$ | $0.00413042663723419$
SMSSP |  $0.5126171173014373$ $\pm$ | $0.012028356145963483$
SMMSP |  $0.5152218244035589$ $\pm$ | $0.007596666506242643$
MSSSP |  $0.5167563402602465$ $\pm$ | $0.014023907440300458$
MSMP |  $0.5175664162150926$ $\pm$ | $0.008183732580774403$
MSMSP |  $0.5189215309964502$ $\pm$ | $0.010260387834085966$
SMSMP |  $0.5195859059845295$ $\pm$ | $0.01030811996935556$
SMMMP |  $0.5216753620453077$ $\pm$ | $0.009659879372934633$
MMMP |  $0.522169198580239$ $\pm$ | $0.00323572613620821$
SSMMP |  $0.5248669555724181$ $\pm$ | $0.011680701297498354$
MSP |  $0.5254668333742372$ $\pm$ | $0.007992422030158665$
MSSP |  $0.526543497457852$ $\pm$ | $0.015102443906296788$
SMMP |  $0.5272624355370156$ $\pm$ | $0.017174448642144444$
SSMSP |  $0.5284746230927171$ $\pm$ | $0.008967395618901962$
SMSP |  $0.5295456320600895$ $\pm$ | $0.005173455059068523$
SSSMP |  $0.5393291069351986$ $\pm$ | $0.013608827969422012$
MMP |  $0.5452990625013332$ $\pm$ | $0.017425677907731026$
SSMP |  $0.5459180905618931$ $\pm$ | $0.009814865157605771$
MP |  $0.5612864877788828$ $\pm$ | $0.007385477748245008$
SMP |  $0.5669066123903705$ $\pm$ | $0.010843513748946546$
SSSSP |  $0.5903769337719205$ $\pm$ | $0.013925916460663948$
SSSP |  $0.6036681408298253$ $\pm$ | $0.012567360341743383$
SSP |  $0.6418977373618453$ $\pm$ | $0.012684036561076101$
SP |  $0.6648514434313151$ $\pm$ | $0.026470828140827057$

The configurations are given as a sequence of masked attention (M) or self-
attention (S) blocks, followed by a pooling by multihead attention block (P).
While performance is slightly lower than Table 3 due to the training changes,
we notice that models that use exclusively self-attention blocks and no masked
blocks are worse than models that predominantly rely on masked attention, as
expected. It is also possible to use self-attention blocks after pooling by
multihead attention. However, that would greatly increase the number of
possible configurations and we do not evaluate this option here.