# Equivariant Graph Attention Networks
for Molecular Property Prediction

Tuan Le  Frank Noé  Djork-Arné Clevert

###### Abstract

Learning and reasoning about 3D molecular structures with varying size is an
emerging and important challenge in machine learning and especially in drug
discovery. Equivariant Graph Neural Networks (GNNs) can simultaneously
leverage the geometric and relational detail of the problem domain and are
known to learn expressive representations through the propagation of
information between nodes leveraging higher-order representations to
faithfully express the geometry of the data, such as directionality in their
intermediate layers. In this work, we propose an equivariant GNN that operates
with Cartesian coordinates to incorporate directionality and we implement a
novel attention mechanism, acting as a content and spatial dependent filter
when propagating information between nodes. We demonstrate the efficacy of our
architecture on predicting quantum mechanical properties of small molecules
and its benefit on problems that concern macromolecular structures such as
protein complexes.

Machine Learning, ICML

##  1 Introduction

Predicting molecular properties is of central importance to applications such
as drug discovery or protein design. In silico (computational) methods with
fast and precise predictions can significantly accelerate the overall process
of finding better molecular candidates in a faster and cheaper way. Learning
on $3$D environments of molecular structures is a rapidly growing area of
machine learning with promising applications but also domain-specific
challenges. While Deep Learning (DL) has replaced hand-crafted features to a
large extent, many advances are crucially determined through inductive biases
in deep neural networks. Developed neural models should maintain an efficient
and accurate representation of structures with even up to thousand of atoms
and correctly reason about their $3$D geometry independent of orientation and
position. A powerful method to restrict a neural network to the functions of
interest, such as a molecular property, is to exploit the symmetry of the data
by constraining equivariance with respect to transformations from a certain
symmetry group (Battaglia et al., 2018; Bronstein et al., 2021).

Graph Neural Networks (GNNs) have been applied on a widespread of molecular
structures, such as in the prediction of quantum chemistry properties of small
molecules (Schütt et al., 2017; Gilmer et al., 2017) but also on
macromolecular structures like proteins (Fout et al., 2017; Ingraham et al.,
2019) due to the natural representation of structures as graphs, with atoms as
nodes and edges drawn based on bonding or spatial proximity. These networks
generally encode the $3$D geometry in terms of rotationally invariant
representations, such as pairwise distances when modelling local interactions
which leads to a loss of directional information, while the addition of
angular information into network architecture has shown to be beneficial in
representing the local geometry (Klicpera et al., 2020).

Neural models that preserve equivariance when working on point clouds in $3$D
space have been proposed (Thomas et al., 2018; Anderson et al., 2019; Fuchs et
al., 2020; Batzner et al., 2021) which can be described as Tensorfield
Networks. These physics-inspired models leverage higher-order tensor
representations and require additional calculations, to construct the basis
for the transformations of their learnable kernels, which can be expensive to
compute.
While these models enable the interaction between different-order
representations, (often referred to as type-$l$ representation), many data
types are often restricted to scalar values (type-$0$ e.g., temperature or
energy) and $3$D vectors (type-$1$ e.g., velocity or forces). Another choice
of using more information with the (limited) data at hand and build data-
efficient models on point clouds111An example of more information preservation
is when considering relative positions between points in $3$D space, where the
information of orientation is maintained, as opposed when only the (scalar;
invariant) distances between points are considered. through equivariant
functions is to operate directly on Cartesian coordinates (Satorras et al.,
2021; Schütt et al., 2021; Jing et al., 2021b, a) and explicitly define the
(equivariant) transformations which is conceptually simpler and does not
require the basis calculations as in Tensorfield Networks.

In this work, we introduce Equivariant Graph Attention Networks (EQGAT)
operating on point clouds that can scale up to hundreds of atoms when
considering larger systems, such as proteins or protein-ligand complexes but
also achieves state-of-art results on the prediction of quantum mechanical
properties of small molecules. Our model architecture implements a novel
attention-mechanism which is invariant to global rotations and translations of
inputs and includes spatial- but also content related information which serves
as powerful edge embedding when propagating information in the Message Passing
Neural Networks (MPNNs) (Gilmer et al., 2017) framework.

##  2 Methods and Related Work

#### Preliminaries.

In this work we consider vector spaces for representing a point cloud graph
$G=(V,P)$ to a feature space. The point cloud has the vertex set
$V=\\{v_{1},\dots,v_{N}\\}$ and positional set
$P=\\{\vec{p}_{1},\dots\,\vec{p}_{N}\\}$, where $\vec{p}_{i}\in\mathbb{R}^{3}$
states the Cartesian coordinate of node $i$. Let
$\mathbf{P}\in\mathbb{R}^{N\times 3}$ denote the coordinate matrix of the
point-cloud $G$ with an arbitrary ordering $\pi(\cdot)$ along the first axis.
A graph adjacency matrix $\mathbf{A}\in\mathbb{R}^{N\times N}$ can be
constructed by defining a distance cutoff $c$ and set $a_{ij}=1$ if
$d_{ij}=||\vec{p}_{i}-\vec{p}_{j}||_{2}<c$ and $0$ else.
The (symmetric) Euclidean distance matrix
$\mathbf{D}\in\mathbb{R}_{+}^{N\times N}$ can be obtained using the canonical
inner product:

$(\mathbf{D})_{i,j}=\sqrt{\vec{p}_{i}^{\top}\vec{p}_{i}+\vec{p}_{j}^{\top}\vec{p}_{j}-2\vec{p}_{i}^{\top}\vec{p}_{j}}.$
(1)

We aim to develop a Graph Neural Network (GNN) model that transforms feature
embeddings invariantly or equivariantly to arbitrary rotations in 3D space.
The advantage of including equivariant features into the model as a structural
prior lies in the fact that some task-related predicted quantities such as
forces $\vec{F}$ in Molecular Dynamics, or the dipole moment $\vec{\mu}$ in
the multipole expansion of the electron density (Schütt et al., 2021) are
geometric quantities and transform as order-1 tensor when a rotation is
performed on the analyzed rigid body. Even though some properties of interest
might in fact not be geometric tensors, i.e., are scalar representations,
preserving the information of the geometry within the network architecture,
has shown to be beneficial for their prediction (Miller et al., 2020; Schütt
et al., 2021).

Formally, in our work we aim to model scalar-valued features
$\mathbf{s}\in\mathbb{R}^{F_{s}}$ and vector-valued features
$\vec{\mathbf{v}}\in\mathbb{R}^{3\times F_{v}}$ in our architecture separately
under the constraint that scalar (type-$0$) features transform invariantly
under rotations and vector (type-$1$) features transform equivariantly under
rotations.

###  2.1 Invariance and Equivariance

Let $X\subset\mathbb{R}^{F_{s}}\times\mathbb{R}^{3\times F_{v}}$ be the
considered vector space. Our feature representation can be expressed as tuple
of scalar- and vector components $\mathbf{x}=(\mathbf{s},\vec{\mathbf{v}})$,
where $\mathbf{s}$ is the invariant and $\vec{\mathbf{v}}$ is the equivariant
representation that transforms accordingly when a group action is applied on
them. The symmetry considered in this paper is the $\textsc{SO}(3)$ group of
orthogonal matrices with determinant $1$ in $\mathbb{R}^{3\times 3}$, i.e.,

$\displaystyle\textsc{SO}(3)=\\{\mathbf{R}\in\mathbb{R}^{3\times 3}:$
$\displaystyle\mathbf{R}^{\top}\mathbf{R}=\mathbf{R}\mathbf{R}^{\top}=\mathbf{I}_{3},~{}~{}\text{det}(\mathbf{R})=1\\}$
(2)

Let $T_{g}:X\xrightarrow{}{}X$ be a set of transformations on $X$ for the
group $g\in\textsc{SO}(3)$. The action of $g=\mathbf{R}\in\textsc{SO}(3)$ on
the vector space $X$ is defined as:

$T_{g}(\mathbf{x})=(\mathbf{s},\mathbf{R}\vec{\mathbf{v}})~{}.$ (3)

The group action in Eq. (3) reveals that scalar component $\mathbf{s}$ is not
affected by the rotation, as this is not a geometric quantity in
$\mathbb{R}^{3}$ and therefore invariant222To be precise, the group action
applied on the scalar embedding is the trivial representation, which is
represented as the identity map through the $F_{s}$-dimensional unit matrix..
The vector feature $\vec{\mathbf{v}}$ however, transforms as an order-$1$
tensor and is therefore rotated.

We say a function $f:X\xrightarrow{}{}Y$ is equivariant to $g$ if there exists
and equivalent transformation on its output space $S_{g}:Y\xrightarrow[]{}Y$
such that following holds:

$f(T_{g}(\mathbf{x}))=S_{g}(f(\mathbf{x}))~{}.$ (4)

We note that a composition of equivariant functions satisfying Eq. (4) is
equivariant again. In our work, $f$ is a message passing (MP) layer (Gilmer et
al., 2017) that updates node feature representation
$\\{\mathbf{x}_{i}\\}_{i=1}^{N}$ of the point cloud in an iterative manner
such that

$\mathbf{x}^{(l+1)}=(\mathbf{s}^{(l+1)},\vec{\mathbf{v}}^{(l+1)})=f((\mathbf{s}^{(l)},\vec{\mathbf{v}}^{(l)}))=f(\mathbf{x}^{(l)})~{}.$
(5)

Combining the group action in Eq. (3) with the equivariance property of the
function $f$ in (4) results into

$\displaystyle f(T_{g}(\mathbf{x}^{(l)}))$
$\displaystyle=S_{g}(f(\mathbf{x}^{(l)}))$
$\displaystyle=S_{g}(\mathbf{x}^{(l+1)})=(\mathbf{s}^{(l+1)},\mathbf{R}\vec{\mathbf{v}}^{(l+1)})~{},$
(6)

which shows that the group action $g$ commutes with the function $f$.

#### Tensor / Outer Product.

The tensor product333Often also referred to as Outer or Kronecker product.
$\otimes$ between two vectors $\mathbf{a}\in\mathbb{R}^{n}$ and
$\mathbf{b}\in\mathbb{R}^{k}$ is computed as

$\mathbf{a}\otimes\mathbf{b}=(\mathbf{a}\mathbf{b}^{\top})\in\mathbb{R}^{n\times
k}~{},$ (7)

and returns a matrix given two vectors. The tensor/outer product will be a
useful operation to construct equivariant features by combining type-$1$ and
type-$0$ representations.

###  2.2 Message Passing Neural Networks (MPNNs)

MPNNs (Gilmer et al., 2017) construct complex representations of vertices
within their local neighbourhood through an iterative exchange of messages
followed by updates of vertex features. Since MPNNs utilize shared trainable
layers among nodes, permutation equivariance is preserved. As mentioned in
Section 2, edges between vertices of the point cloud are specified by their
relative position $\vec{p}_{ij}=\vec{p}_{j}-\vec{p}_{i}$ of vertices $i,j$
within a local neighbourhood through a distance cutoff $c>0$ , i.e.,
$\mathcal{N}(i)=\\{j:||\vec{p}_{ij}||_{2}=d_{ij}<c\\}$. The (Euclidean)
distance function in (1) is an invariant function as the canonical inner
product between two vectors $\vec{p}_{i}$ and $\vec{p}_{j}$ after the group
action $\mathbf{R}\in\textsc{SO}(3)$ does not change, i.e.,
$(\mathbf{R}\vec{p}_{i})^{\top}(\mathbf{R}\vec{p}_{j})=(\vec{p}_{j}^{\top}\mathbf{R}^{\top})(\mathbf{R}\vec{p}_{j})=\vec{p}_{i}^{\top}\mathbf{I}_{3}\vec{p}_{j}=\vec{p_{i}}^{\top}\vec{p}_{j}$.
Hence, the inner product, or a composition of it, like the norm, is a natural
way to obtain $\textsc{SO}(3)$ invariant features from equivariant features.
The usage of the tensor/outer product and inner product to obtain equivariant
and invariant features, respectively, will be explained in the next Section.

Standard MPNNs implement a message and update function for feature
representation as

$\displaystyle\mathbf{m}_{i}^{(l+1)}=\sum_{j\in\mathcal{N}(i)}f_{m}^{(l)}(\mathbf{x}_{i}^{(l)}\mathbf{x}_{j}^{(l)},\vec{p}_{ji})=\sum_{j\in\mathcal{N}(i)}\mathbf{m}_{ji}^{(l)}~{},$
(8)
$\displaystyle\mathbf{x}_{i}^{(l+1)}=f_{u}^{(l)}(\mathbf{x}_{i}^{(l)},\mathbf{m}_{i}^{(l+1)})~{}.$
(9)

In our paper, we aim to implement the message ($f_{m}$) and update ($f_{u}$)
function to be equivariant, i.e., to satisfy the property in Equation (2.1).

Figure 1:  (a) Overview of the EQGAT encoder architecture with initial scalar
features obtained through an embedding layer. (b) Update layer to enable
information flow between scalar- and vector features. (c) Message function
that creates vector features through the tensor product of relative positions
with filtered scalar features combined with existing filtered vector features.
Scalar messages are obtained via filtered value tensors. (d) Feature Attention
calculation that inputs SO(3) invariant representations and outputs a content
as well as spatial based filter.

###  2.3 Related Work

Neural networks that specifically achieve E($3$) or SE($3$) equivariance have
been proposed in Tensorfield Networks (TFN) (Thomas et al., 2018) and its
variants in the covariant Cormorant (Anderson et al., 2019), NequIP (Batzner
et al., 2021) and SE(3)-Transformer (Fuchs et al., 2020). With TFNs,
equivariance is achieved through the usage of equivariant function spaces such
as Spherical Harmonics combined with Clebsch-Gordan coefficients, while others
resort to lifting the spatial space to higher-dimensional spaces such as Lie
group spaces (Finzi et al., 2020). Since no restriction on the order of
tensors is imposed on these methods, sufficient expressive power of these
models is guaranteed, but at a cost of excessive computational calculations
with increased time and memory. To circumvent the expensive computational
cost, another line of research proposed to directly implement equivariant
operations in the original (Cartesian) space, providing and efficient approach
to preserve equivariance without performing complex space transformations as
introduced in the E($n$)-GNN (Satorras et al., 2021), GVP (Jing et al., 2021b,
a) and PaiNN (Schütt et al., 2021) architectures.

Our proposed model implements equivariant operations in the original Cartesian
space and includes a continuous filter through the self-attention coefficients
which serve as spatial- and content based edge embedding in the message
propagation, as opposed to the PaiNN model where the filter solely depends on
the distance. The E($n$)-GNN architecture does not learn higher-dimensional
type-$1$ vector features, but only updates given type-$1$ features through a
weighted linear combination of such, where the (learnable) scalar weights are
obtained from invariant embeddings. The GVP model which was initially designed
to work on macromolecular structures includes a complex message functions of
concatenated node- and edge features composed with a series of GVP-blocks that
enables information exchange between type-$0$ and type-$1$ features, with a
potential disadvantage of discontinuities through non-smooth components for
distances close to the cutoff.
A concurrent work by (Thölke & Fabritiis, 2022) proposed the Equivariant-
Transformer (ET), which is similar to ours, and obtains strong results on QM9
(Ramakrishnan et al., 2014), MD17 (Chmiela et al., 2017) and ANI-1 (Smith et
al., 2017) datasets. Within their model, to initialize node embeddings, they
implement an additional neighbouring embedding $n_{i}$ for target node $i$,
which can thought of a continuous filter convolution as proposed in the SchNet
architecture (Schütt et al., 2017). This neighbouring embedding $n_{i}$ is
then combined with a self-embedding, generally resembling a graph
convolutional layer, which our model does not incorporate. Our model differs
in a distinguished message function where we implement a feature-attention
embedding that aims to filter scalar-channels from neighbouring nodes.
Furthermore, our message function incorporates an intermediate embedding that
is content- as well as spatial dependent to modulate the vector-channels,
while their filter solely depends on spatial information via interatomic
distances. As we benchmark our model on large molecular structures, we assign
different channel sizes, $F_{s},F_{v}$ for scalar and vector features,
respectively, to be able to train deep models on learning problems that
include proteins.

##  3 Equivariant Graph Attention Networks

The traditional Transformer model as introduced by Vaswani et al. (Vaswani et
al., 2017) has revolutionized the field of Natural Language Processing (NLP)
(Devlin et al., 2019; Dai et al., 2019) and image analysis (Ramachandran et
al., 2019; Zhao et al., 2020; Dosovitskiy et al., 2021). The self-attention
module which lies in the core of Transformers, initially designed to operate
on sequence of tokens in NLP, consists in essence of two components: input-
dependent attention-weights between any two elements of the set, and an
embedding for each set-element, called value. Since graphs, are naturally
represented as sets with underlying structure through the (sparse)
connectivity between vertices, the implementation of self-attention, following
the MPNN paradigm, was introduced in Graph Attention Networks (GATs)
(Veličković et al., 2018). Transformer-like GNNs (that work on a fully-
connected 2D graph) were recently introduced in (Kreuzer et al., 2021; Ying et
al., 2021) and have found success in several graph-learning benchmarks due to
their incorporation of (learnable) relative positional encodings into the
attention function.
Our proposed Equivariant Graph Attention Networks (EQGAT) operates in $3$D
space and implements the message passing for each target node $i$ on its local
neighbourhood $\mathcal{N}(i)$ as defined in Section 2 to avoid the quadratic
complexity of the vanilla self-attention when one target node would interact
with all other nodes in the point cloud. We emphasize that the integration of
local neighbourhoods manifests as a powerful inductive bias and in a bio-
chemistry context, coincides with the assumption that a large part of energy
variations can be attributed to local interactions, although the influence and
importance of non-local effects in machine-learned force-fields has been
recently analyzed in (Unke et al., 2021).
In the following Subsections, we will introduce the function components of our
EQGAT network displayed in Figure 1a and highlight its invariance and
equivariance properties.

###  3.1 Input Embedding

We initially embed atoms of small molecules or proteins based on their nuclear
charge $Z$. As nuclear charges are discrete and bounded, we utilize a
trainable embedding look-up table commonly used in NLP to map the $i-$th atom
charge $z_{i}$ to a feature representation
$\mathbf{s}_{i}^{(0)}\in\mathbb{R}^{F_{s}}$ as

$\mathbf{s}_{i}^{(0)}=\text{embed}(z_{i}),$ (10)

which provides a starting (invariant) scalar representation of the node prior
to the message passing.
As in most cases, no directional information for atoms are available, we
initialize the vector features as zero tensor
$\vec{\mathbf{v}}_{i}^{(0)}=\vec{\mathbf{0}}\in\mathbb{R}^{3\times F_{v}}$.

###  3.2 Distance Encoding

Interatomic distances $d_{ji}=||\vec{p}_{ji}||_{2}\in\mathbb{R}_{+}$ are
embedded using the Bessel radial basis function (RBF) as introduced by
Klicpera et al. into a representation
$\mathbf{e}^{\text{RBF}}(d_{ji})\in\mathbb{R}^{K}$ that serves as distance
encoding into the attention mechanism. The Bessel RBF function reads

$\displaystyle{e}^{\text{RBF}}_{k}(d_{ji})=\sqrt{\frac{2}{c}}\frac{\sin{(\frac{k\pi}{c}}d_{ji})}{d_{ji}},$

where $c$ is the distance cutoff and $k=1,\dots,K$. In similar fashion to the
continuous filter convolution in the SchNet architecture (Schütt et al.,
2017), the deterministic radial basis encoding is further transformed using a
trainable linear layer to obtain an edge embedding
$\tilde{\mathbf{e}}_{ji}:=\tilde{\mathbf{e}}(d_{ji})\in\mathbb{R}^{F_{s}}$,

$\displaystyle\tilde{\mathbf{e}}_{ji}=\mathbf{W}_{e}\mathbf{e}^{\text{RBF}}(d_{ji})+\mathbf{b}_{e}~{}.$

As the final edge embedding $\mathbf{e}_{ji}$ ought to transition smoothly to
$\mathbf{0}$ to avoid discontinuities when $d_{ji}\xrightarrow[]{}c$ , we
apply a cosine-cutoff function $\kappa:\mathbb{R}_{+}\mapsto[0,1]$ as
introduced in (Behler, 2011)

$\displaystyle\kappa(d_{ji})=\frac{1}{2}\left(\cos{(\frac{\pi
d_{ji}}{c})}+1\right)\cdot\mathbbm{1}[d_{ji}\leq c]~{},$

which leads to the final edge embedding:

$\displaystyle\mathbf{e}_{ji}$ $\displaystyle:=\mathbf{e}(d_{ji})$ (11)
$\displaystyle=\kappa(d_{ji})\cdot\tilde{\mathbf{e}}_{ji}=\kappa(d_{ji})\cdot(\mathbf{W}_{e}\mathbf{e}^{\text{RBF}}(d_{ji})+\mathbf{b}_{e})~{}.$

The edge-embedding $\mathbf{e}_{ji}\in\mathbb{R}^{F_{s}}$ in Equation (11) can
be interpreted as a convolutional filter that solely depends on an invariant
representation, i.e., the distance between the nodes $j$ and $i$.

###  3.3 Feature Attention

Our attention mechanism depends on invariant scalar representations and
includes the distance encoding in Eq. (11) to incorporate spatial
connectivity. A novelty of our approach is that the attention coefficient
between two vertices $j$ and $i$ is in fact obtained per feature-channel
instead for the entire embedding. Formally, given the invariant scalar
embeddings $\mathbf{s}_{i},\mathbf{s}_{j}\in\mathbb{R}^{F_{s}}$ we use two
linear layers to obtain query- and key embeddings as

$\displaystyle\mathbf{q}_{i}$
$\displaystyle=\mathbf{W}_{q}\mathbf{s}_{i}+\mathbf{b}_{q}~{}\in\mathbb{R}^{F_{s}}~{},$
$\displaystyle\mathbf{k}_{j}$
$\displaystyle=\mathbf{W}_{k}\mathbf{s}_{j}+\mathbf{b}_{k}~{}\in\mathbb{R}^{F_{s}}~{}.$
(12)

We proceed to compute a content and spatial dependent anisotropic embedding
through the elementwise product as

$\tilde{\mathbf{a}}_{ji}=\mathbf{q}_{i}\odot\mathbf{k}_{j}\odot\mathbf{e}_{ji}~{}.$
(13)

The embedding vector $\tilde{\mathbf{a}}_{ji}\in\mathbb{R}^{F_{s}}$ contains
both, semantic information through the query and key representations as well
as spatial information via the relative positional encoding expressed through
the edge embedding.
Next, we utilize another weight matrix
$\mathbf{W}_{a}\in\mathbb{R}^{(F_{s}+2F_{v})\times F_{s}}$ to obtain an
embedding vector to compute the feature-attention and filter components for
the equivariant features,

$\mathbf{a}_{ji}=\mathbf{W}_{a}\tilde{\mathbf{a}}_{ji}=[\tilde{\bm{\alpha}}_{ji},\bm{\beta}_{ji},\bm{\gamma}_{ji}]\in\mathbb{R}^{F_{s}+2F_{v}}~{}.$
(14)

The feature-attention is calculated using the $\bm{\tilde{\alpha}}_{ji}$
tensor as follows

$\bm{\alpha}_{ji}=\frac{\sigma(\bm{\tilde{\alpha}}_{ji})}{\sum_{j^{\prime}\in\mathcal{N}(i)}\sigma(\bm{\tilde{\alpha}}_{j^{\prime}i})}~{}\in(0,1)^{F_{s}},$
(15)

where $\sigma(\cdot)$ is the sigmoidal activation function and applied
componentwise for each channel.

###  3.4 Value Transforms

Another constituent of the self-attention mechanism is the pointwise transform
of neighbouring (source) nodes which will be multiplied with the output of the
feature-attention mechanism described in Eq. (14). Value tensors for scalar-
and vector features are computed using weight matrices
$\mathbf{W}_{sv}\in\mathbb{R}^{F_{s}\times(F_{s}+2F_{v})}$, and
$\mathbf{W}_{vv}\in\mathbb{R}^{F_{v}\times F_{v}}$.

$\displaystyle\mathbf{v}_{s,j}=\mathbf{W}_{sv}\mathbf{s}_{j}+\mathbf{b}_{sv}~{},$
$\displaystyle\vec{\mathbf{v}}_{v,j}=\vec{\mathbf{v}}_{j}\mathbf{W}_{vv},$
(16)

where the product on $\vec{\mathbf{v}}_{j}\in\mathbb{R}^{3\times F_{v}}$ is
applied on the feature axis. By such application of the linear transformation
in Eq. (3.4), the equivariance property in Eq. (3) is maintained.
The rationale behind the additional $2F_{v}$ output channels for the scalar
embedding $\mathbf{v}_{s,j}$ is associated to the construction of equivariant
features through the tensor product with invariant features as will be
explained later.
We split the value tensor of the scalar part $\mathbf{v}_{s,j}$ into three
tensors $[\mathbf{v}_{s_{0},j},\mathbf{y}_{v_{0},j},\mathbf{y}_{v_{1},j}]$.

#### Self-Attention as Continuous Filter Convolution.

The invariant feature-attention coefficients in Eq. (15) depend on content
related as well as spatial based information and act as a filter when
elementwise multiplied with the pointwise transforms of scalar-value tensors
$\mathbf{v}_{s_{0},j}$ to obtain the scalar message-embedding $\forall
j\in\mathcal{N}(i)$:

$\mathbf{m}_{s,ji}=\bm{\alpha}_{ji}\odot\mathbf{v}_{s_{0},j}~{}.$ (17)

Furthermore, we compute two additional elementwise products to obtain filtered
message tensors:

$\displaystyle\mathbf{m}_{v_{0},ji}=\bm{\beta}_{ji}\odot\mathbf{y}_{v_{0},j}~{},$
$\displaystyle\mathbf{m}_{v_{1},ji}=\bm{\gamma}_{ji}\odot\mathbf{y}_{v_{1},j}~{}.$
(18)

We highlight that the message embeddings in Equations 17 and 3.4 are
calculated using $\textsc{SO}(3)$ invariant representations through the
initial query-key computation in (3.3) and only the channels of the
$\alpha$-coefficients are bounded in the unit interval to obtain ”modulated”
scalar representations when multiplied with the $\mathbf{v}_{s_{0},j}$ tensor.

Our model differs from the SE(3)-Transformer proposed by Fuchs et al. as we do
not rely on Spherical Harmonics calculation and Clebsch-Gordan decomposition
to build equivariant functions, but we explicitly design functions that are
equivariant and operate on 3D-Cartesian coordinates for faster and more
memory-efficient calculations as described in the next paragraph.
Additionally, our proposed model aims to decouple the information flow between
scalar- and vector representations within the attention-mechanism by just
using scalar/invariant representations $\mathbf{a}_{ji}$ in Equation (14) to
modulate scalar and vector embeddings, while the (later applied) update
function is used to enable an information flow between vector representations
to scalar embeddings.
An important fact is that the attention coefficient(s) which serve as a filter
are required to be SO($3$) invariant. Such requirement originates from the
idea that the attention coefficient is multiplied with a type-1 feature444When
considering the vector features. which itself transforms, when a group action
(in our case a Rotation), is applied. For the case that one would want to
include type-1 features into the attention calculation, as proposed by (Fuchs
et al., 2020) through additional type-1 related query and key representations,
a contraction along the spatial axis, such as an inner product between those
pairs is required to obtain SO(3) invariance. We tried such construction in
our initial development of the attention function, but found worse performance
in validation and additional computational complexity.

#### Building Equivariant Features.

In case that no initial vector features such as velocity, or forces are
available, equivariant representations in the point cloud can be constructed
as a function of (relative) position $\vec{p}_{ij}=\vec{p}_{j}-\vec{p}_{i}$.
Relative position vectors in 3D space can have unbounded norms, so a common
practice is to obtain relative positions with unit norm
$\vec{p}_{e,ij}=\frac{\vec{p}_{ij}}{||\vec{p}_{ij}||_{2}}$ which describe
points on the $2-$dimensional unit sphere
$S^{2}=\\{\vec{r}\in\mathbb{R}^{3}:||\vec{r}||_{2}=1\\}$.
Equivariant interactions between node $j$ and $i$ are modelled through the
tensor-product (cf. Eq. (7)) of the equivariant representation
$\vec{p}_{e,ji}\in\mathbb{R}^{3}$ with invariant representation
$\mathbf{m}_{v_{0},ji}\in\mathbb{R}^{F_{v}}$.
Formally, equivariant features are obtained as

$\vec{\mathbf{y}}_{0,ji}=\vec{p}_{e,ji}\otimes\mathbf{m}_{v_{0},ji}=\vec{p}_{e,ji}\mathbf{m}_{v_{0},ji}^{\top}~{}~{}\in\mathbb{R}^{3\times
F_{v}},$ (19)

and (hidden) equivariant type-$1$ embeddings are filtered using the
(elementwise) scalar multiplication

$\vec{\mathbf{y}}_{1,ji}=(\mathbf{1}\otimes\mathbf{m}_{v_{1},ji})\odot\vec{\mathbf{a}}_{v,ji}=(\mathbf{1}\mathbf{m}_{v_{1},ji}^{\top})\odot\vec{\mathbf{a}}_{v,ji}~{},$
(20)

where $\mathbf{1}\in\mathbb{R}^{3}$ with $1$’s in the components and
$\vec{\mathbf{a}}_{v,ji}$ is obtained through the cross product $\times$ as:

$\displaystyle\vec{\mathbf{a}}_{v,ji}=\vec{\mathbf{v}}_{j}\times\vec{\mathbf{v}}_{i}+\vec{\mathbf{v}}_{v,j}~{}.$

The rationale for including the cross product of the two transformed vector
features is to enable interaction between type-1 features and is inspired from
TFNs (Thomas et al., 2018) which models the tensor product between two type-1
vectors resulting in a rank-2 tensor. Such rank-2 tensor includes
antisymmetric elements presented in the cross product. We combine the two
equivariant representations into a final aggregated equivariant message
embedding

$\displaystyle\vec{\mathbf{m}}_{v,i}=\sum_{j\in\mathcal{N}(i)}(\vec{\mathbf{y}}_{0,ji}+\vec{\mathbf{y}}_{1,ji})~{}.$
(21)

The aggregated invariant message embedding is obtained in the same fashion,
using filtered messages (see. Eq. (17)) between target node $i$ and its
neighbours

$\mathbf{m}_{s,i}=\sum_{j\in\mathcal{N}(i)}\mathbf{m}_{s,ji}~{}.$ (22)

We follow our vector space terminology and write

$\displaystyle\mathbf{m}_{i}=(\mathbf{m}_{s,i},\vec{\mathbf{m}}_{v,i})\in\mathbb{R}^{F_{s}}\times\mathbb{R}^{3\times
F_{v}}~{},$

as the final aggregated message embedding (cf. Eq. (8)) which by design also
satisfies the equivariance property in Eq. (2.1). We prove this claim in the
Appendix A.
At this point, it is worth mentioning that the tensor product applied in
Equations 19 and 20 reduces to a scalar multiplication in case the (invariant)
embedding(s) $\mathbf{m}_{v,ji}$ are one-dimensional as in the E($n$)-GNN
architecture (Satorras et al., 2021). Furthermore, notice that by such
construction, higher-dimensional type-$1$ embeddings can also be constructed
in $\mathbb{R}^{n}$ instead of $\mathbb{R}^{3}$.

###  3.5 Update

We update the hidden embedding for each target node $i$ by adding the
aggregated message $\mathbf{m}_{i}$ with the previous hidden state

$\mathbf{\tilde{x}}_{i}\xleftarrow{}\mathbf{x}_{i}+\mathbf{m}_{i}~{}.$ (23)

The state $\mathbf{\tilde{x}}_{i}$ in Equation (23) includes complex
transformation performed in the $\mathbf{m}_{i}$ embedding via the attention
mechanism, while no self-interaction transformation on
$\mathbf{x}_{i}\in\mathbb{R}^{F_{s}}\times\mathbb{R}^{3\times F_{v}}$ has been
applied. We propose to first update the $i-$th’s hidden state in a residual-
connection manner and then apply a pointwise update-layer similar as in the
PaiNN architecture (Schütt et al., 2021) with the use of gated equivariant
non-linearities (Weiler et al., 2018) to combine invariant and equivariant
information in the $\tilde{\mathbf{x}}$ representation to update its state
(cf. Eq. (9)) into
$\mathbf{x}_{i}\xleftarrow[]{}f_{u}(\tilde{\mathbf{x}}_{i})$ as illustrated in
Figure 1b.

##  4 Experiments and Results

Table 1: QM$9$ Mean Absolute Error (MAE) on the test set.
Models displayed with * use different (random) train/val and test splits. For
our EQGAT, we report averaged MAEs over 3 runs.

Tasks | $\alpha$ | $\Delta\epsilon$ | $\epsilon_{\textsc{HOMO}}$ | $\epsilon_{\textsc{LUMO}}$ | $\mu$ | $C_{\nu}$
---|---|---|---|---|---|---
Units | bohr3 | meV | meV | meV | $D$ | $\frac{\text{cal}}{\text{mol K}}$
NMP | .092 | 69 | 43 | 38 | .030 | .040
SchNet∗ | .235 | 63 | 41 | 34 | .033 | .033
Cormorant | .085 | 61 | 34 | 38 | .038 | .026
LieConv | .084 | 49 | 30 | 25 | .032 | .038
DimeNet++${}^{\text{*}}$ | .044 | 33 | 25 | 20 | .030 | .023
SE(3) Tr. | .142 | 53 | 35 | 33 | .051 | .054
E($n$)-GNN | .071 | 48 | 29 | 25 | .029 | .031
ET∗ | .010 | 38 | 21 | 18 | .002 | .026
PaiNN${}^{\text{*}}$ | .045 | 46 | 28 | 20 | .012 | .024
EQGAT | .063 | 44 | 26 | 22 | .014 | .027
EQGAT∗ | .053 | 32 | 20 | 16 | .011 | .024

We test the effectiveness of our proposed EQGAT model on four publicly
available molecular benchmark datasets which pose significant challenges for
the development of efficient and accurate prediction models in small molecules
drug discovery but also protein design on different data scales.

###  4.1 QM9

The QM9 dataset (Ramakrishnan et al., 2014) is a chemical property regression
benchmark and consists of $134$k small molecules with up to $29$ atoms per
molecule. Molecules are represented as point clouds with each atom having a
$3$D position and a five dimensional one-hot encoding that describes the atom
type (H, C, N, O, F), and additional features that can be derived from the 2D
topological graph, such as bond-types. In our experiments, we only use the
atom positions and atom types/charges as input features for our EQGAT model as
described in the Methods Section 3.1. To compare our method with the
literature, we import the training, validation- and test-splits from (Anderson
et al., 2019) which consists of $100$k, $18$k and $13$k compounds,
respectively. We adopt the hyperparameters from (Satorras et al., 2021) and
implement a $7$-layer EQGAT-encoder with $F_{s}=128$ scalar and $F_{v}=32$
vector channels as well as $K=20$ radial basis functions to encode interatomic
distances. For the downstream networks we only utilize scalar embeddings from
the last hidden layer. For a detailed architecture of the entire network with
approximately $1.1$M trainable parameters, we refer to Section B of the
Appendix.
We optimized and report the Mean Absolute Error (MAE) between prediction and
ground truths on 6 targets and compare to NMP (Gilmer et al., 2017), SchNet
(Schütt et al., 2017), Cormorant (Anderson et al., 2019), LieConv (Finzi et
al., 2020), DimeNet++ (Klicpera et al., 2020), SE(3)-Transformer (Fuchs et
al., 2020), E$(n)$-GNN (Satorras et al., 2021), ET (Thölke & Fabritiis, 2022)
and PaiNN (Schütt et al., 2021).

Table 2: Benchmark results on three ATOM3D tasks.
We report the results for the Baseline models from (Townshend et al., 2021).
For the E($n$)-GNN, PaiNN and our EQGAT model, we report averaged metrics over
3 runs.

Tasks | PSR $(\uparrow)$ | RSR $(\uparrow)$ | LBA $(\downarrow)$
---|---|---|---
Metric | Global $R_{S}$ | Mean $R_{S}$ | Global $R_{S}$ | Mean $R_{S}$ | RMSE
CNN | 0.431 | 0.789 | 0.264 | 0.372 | 1.415
GNN | 0.515 | 0.755 | 0.234 | 0.512 | 1.570
GVP-GNN | 0.511 | 0.845 | 0.211 | 0.330 | 1.594
E($n$)-GNN | 0.466 | 0.789 | - | - | 1.558
PaiNN | 0.485 | 0.808 | - | - | 1.548
EQGAT | 0.576 | 0.849 | 0.322 | 0.365 | 1.489

Results of our proposed model are reported in Table 1. Our EQGAT architecture
obtains very strong performance on the 6 predicted targets compared to recent
state-of-the-art models while our model similar to the E($n$)-GNN, ET, and
PaiNN model implements equivariant functions in the original Cartesian space
without the need of higher-order tensor representations. We believe that the
implementation of output-specific layers for the prediction of certain targets
such as the magnitude of the dipole moment $\mu$ or the electronic spatial
extent $\langle R^{2}\rangle$ (not listed in Table 1) as proposed by (Schütt
et al., 2021) using the GNN encoder’s scalar and vector features
$\mathbf{x}^{(L)}=(\mathbf{s}^{(L)},\vec{\mathbf{v}}^{(L)})$ in combination
with the global geometry via Cartesian coordinates
$\\{\vec{p}_{i}\\}_{i=1}^{N}$ might lead to better performance on those
targets. As we adopt the overall architecture for all 6 targets, we did not
implement output-specific (decoder) layers for each target. In comparison to
the concurrent ET architecture (Thölke & Fabritiis, 2022), our model obtains
similar results on the 6 targets while being parameter-lighter ($1.1$M vs.
$6.9$M).

We refer the reader to the Appendix B for the architectural details and
performance of our models on all $12$ QM$9$ targets including energy
predictions.

###  4.2 ATOM3D

The ATOM3D benchmark (Townshend et al., 2021) is a collection of eight tasks
and datasets for learning on atomic-level 3D molecular structures of different
kinds, i.e., proteins, RNAs, small molecules and complexes. Since proteins
perform specific biological functions essential for all living organisms and
hence, play a key role when investigating the most fundamental questions in
the life sciences, we focus our experiments on the learning problems often
encountered in structural biology with different difficulties due to data
scarcity and varying structural sizes. We use provided training, validation
and test splits from ATOM3D and refer the interested reader to the original
work (Townshend et al., 2021) for more details. For all benchmarks, we compare
against the Baseline CNN and GNN models provided by the authors from ATOM3D,
the GVP-GNN reported in (Jing et al., 2021a) and PaiNN (Schütt et al., 2021)
as well as the E($n$)-GNN (Satorras et al., 2021) architectures using our own
implementations.
All of our implemented models utilize a 5-layer GNN encoder with $128$
(scalar) channels as hidden invariant embedding. For our EQGAT model we apply
$16$ vector channels, while the PaiNN architecture models $128$ vector
channels and the E($n$)-GNN does not include learnable vector-feature but just
updates the positional coordinates. We refer the reader to the Appendix C for
a detailed description concerning the implementation of the networks on the
tasks.

The Protein and RNA Structure Ranking tasks (PSR / RSR) in ATOM3D are both
regression tasks with the objective to predict the quality score in terms of
Global Distance Test (GDT_TS) or Root-Mean-Square Deviation (RMSD) for
generated Protein and RNA models wrt. to its experimentally determined ground-
truth structure. Being able to reliably rank a biopolymer structure requires a
model to accurately learn the atomic environments such that discrepancies
between a ground truths state an its corrupted version can be distinguished.
We evaluated our model on the biopolymer ranking and obtained the best results
on the current benchmark, reported in Table 2 in terms of Spearman rank
correlation. Our proposed model performs particularly well on the PSR task
outperforming the GVP-GNN (Jing et al., 2021a). We noticed that the RSR
benchmark was particularly difficult to validate as only a few dozen
experimentally determined RNA structures are existent to date, and the
structural models generated in the ATOM3D framework are labeled with the RMSD
to its native structure, which is known to be sensitive to outlier regions,
for exampling by inadequate modelling of loop regions (Zemla, 2003), while the
GDT_TS metric might be a better suited target to predict a ranking for
generated RNA structures as in the PSR benchmark. In our experiments, we tried
to obtain comparable results for the PaiNN and E($n$)-GNN architecture on the
biopolymer benchmarks but found it difficult to make our implementation train
on these two benchmarks after careful comparison to the authors original
source code.

Another challenging and important task for drug discovery projects is
estimating the binding strength (affinity) of a candidate drug atomistic’s
interaction with a target protein. We use the ligand binding affinity (LBA)
dataset and found that among the GNN architectures, our proposed model obtains
the best results, while also being computationally cheap and fast to train.
The best performing model in the LBA-task is a 3D CNN model which works on the
joint protein-ligand representation using voxel space and enforcing
equivariance through data augmentation. The inferior performance of all
equivariant GNNs might be caused by the need of larger filters to better
capture the locality, where 3D CNNs have an advantage when using voxel
representations. Furthermore, as all GNN models jointly represent ligand- and
protein as one graph by connecting vertices through a distance cutoff of
$4.5\textup{\AA}$, we believe that such union leads to an information loss of
distinguishing the atom identity from the ligand and protein. A promising
direction to investigate is to incorporate a ligand and protein GNN encoder
seperately and merge the two embeddings prior the binding affinity prediction,
similar to Graph Matching Networks (Li et al., 2019).

Computational Efficiency. We assess the computational efficiency of the
proposed equivariant Graph Attention Network and compare it against our
implemented PaiNN and E($n$)-GNN architectures on the PSR, RSR and LBA
benchmarks from ATOM3D. As these datasets consist of graphs with up to
hundreds of atoms, computationally- and memory efficient models are preferred
such that batches of graphs can be stored on GPU memory and trained fast. We
measure the inference time of a random batch comprising 10 macromolecular
structures on an NVIDIA V$100$ GPU and observe the fastest execution time for
our proposed model on the LBA dataset with $23.67$ms against the PaiNN and
E($n$)-GNN with $30.75$ms and $25.76$ms, respectively.
Most notably on the PSR and RSR datasets which consist of large biopolymer
structures, the time difference is more obvious for EQGAT with $42.79$ms
against PaiNN ($67.53$ms) and E($n$)-GNN ($56.66$ms) on the PSR dataset, while
the inference on the RNA benchmark takes most time with $54.80$ms for EQGAT
against $90.38$ms (PaiNN) and $76.45$ms for the E($n$)-GNN architecture.
The inferior execution time for the PaiNN architecture is most likely
attributed to the fact that the number of scalar and vector channels is set to
the same value. As the number of scalar channels is set to $128$ for all
Atom3D models, the PaiNN architecture implements vector features of size
$[\textsc{\\#Nodes},3,128]$ which becomes very memory-intensive for large
graphs and requires a higher amount of FLOPS in its convolutional layers.
Despite the inferior performance of the PaiNN architecture on the ATOM3D
benchmark, it is worth highlighting that the PaiNN model was initially
designed to predict potential energy surfaces of small molecules with high
fidelity, while its architectural components most likely require further re-
implementation to be usable on larger complexes.

##  5 Conclusion

In this work, we introduce a novel attention-based graph neural network for
the prediction of molecular properties of systems with varying size. Our
proposed equivariant Graph Attention Network makes use of rotationally
equivariant features in their intermediate layers to faithfully represent the
geometry of the data, while being computationally efficient, as all
equivariant functions are directly implement in the original Cartesian space.
We demonstrate the potential of EQGAT by succesfully applying it to a wide
range of molecular tasks, such as the prediction of quantum mechanical
properties of small molecules, but also on learning problems that involve
large systems such as proteins, or protein-ligand complexes.

## Code and Data Availability

The code and data will be made available upon official publication.

## References

  * Anderson et al. (2019) Anderson, B., Hy, T.-S., and Kondor, R.  Cormorant: Covariant molecular neural networks, 2019.
  * 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., Gulcehre, C., Song, F., Ballard, A., Gilmer, J., Dahl, G., Vaswani, A., Allen, K., Nash, C., Langston, V., Dyer, C., Heess, N., Wierstra, D., Kohli, P., Botvinick, M., Vinyals, O., Li, Y., and Pascanu, R.  Relational inductive biases, deep learning, and graph networks, 2018.
  * Batzner et al. (2021) Batzner, S., Musaelian, A., Sun, L., Geiger, M., Mailoa, J. P., Kornbluth, M., Molinari, N., Smidt, T. E., and Kozinsky, B.  E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials, 2021.
  * Behler (2011) Behler, J.  Atom-centered symmetry functions for constructing high-dimensional neural network potentials.  _The Journal of Chemical Physics_ , 134(7):074106, 2011.  doi: 10.1063/1.3553717.
  * Bronstein et al. (2021) Bronstein, M. M., Bruna, J., Cohen, T., and Veličković, P.  Geometric deep learning: Grids, groups, graphs, geodesics, and gauges, 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.  doi: 10.1126/sciadv.1603015.  URL https://www.science.org/doi/abs/10.1126/sciadv.1603015.
  * Dai et al. (2019) Dai, Z., Yang, Z., Yang, Y., Carbonell, J., Le, Q. V., and Salakhutdinov, R.  Transformer-xl: Attentive language models beyond a fixed-length context, 2019.
  * Devlin et al. (2019) Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K.  Bert: Pre-training of deep bidirectional transformers for language understanding, 2019.
  * Dosovitskiy et al. (2021) Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N.  An image is worth 16x16 words: Transformers for image recognition at scale, 2021.
  * Finzi et al. (2020) Finzi, M., Stanton, S., Izmailov, P., and Wilson, A. G.  Generalizing convolutional neural networks for equivariance to lie groups on arbitrary continuous data, 2020.
  * Fout et al. (2017) Fout, A., Byrd, J., Shariat, B., and Ben-Hur, A.  Protein interface prediction using graph convolutional networks.  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.
  * Fuchs et al. (2020) Fuchs, F. B., Worrall, D. E., Fischer, V., and Welling, M.  Se(3)-transformers: 3d roto-translation equivariant attention networks, 2020.
  * Gilmer et al. (2017) Gilmer, J., Schoenholz, S. S., Riley, P. F., Vinyals, O., and Dahl, G. E.  Neural message passing for quantum chemistry, 2017.
  * Ingraham et al. (2019) Ingraham, J., Garg, V., Barzilay, R., and Jaakkola, T.  Generative models for graph-based protein design.  In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buc, F., Fox, E., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems_ , volume 32. Curran Associates, Inc., 2019.
  * Jing et al. (2021a) Jing, B., Eismann, S., Soni, P. N., and Dror, R. O.  Equivariant graph neural networks for 3d macromolecular structure, 2021a.
  * Jing et al. (2021b) Jing, B., Eismann, S., Suriana, P., Townshend, R. J. L., and Dror, R.  Learning from protein structure with geometric vector perceptrons, 2021b.
  * Klicpera et al. (2020) Klicpera, J., Groß, J., and Günnemann, S.  Directional message passing for molecular graphs, 2020.
  * Kreuzer et al. (2021) Kreuzer, D., Beaini, D., Hamilton, W. L., Létourneau, V., and Tossou, P.  Rethinking graph transformers with spectral attention, 2021.
  * Li et al. (2019) Li, Y., Gu, C., Dullien, T., Vinyals, O., and Kohli, P.  Graph matching networks for learning the similarity of graph structured objects, 2019.
  * Miller et al. (2020) Miller, B. K., Geiger, M., Smidt, T. E., and Noé, F.  Relevance of rotationally equivariant convolutions for predicting molecular properties, 2020.
  * Ramachandran et al. (2019) Ramachandran, P., Parmar, N., Vaswani, A., Bello, I., Levskaya, A., and Shlens, J.  Stand-alone self-attention in vision models, 2019.
  * Ramakrishnan et al. (2014) Ramakrishnan, R., Dral, P., Rupp, M., and von Lilienfeld, A.  Quantum chemistry structures and properties of 134 kilo molecules.  _Scientific Data_ , 1, 08 2014.  doi: 10.1038/sdata.2014.22.
  * Satorras et al. (2021) Satorras, V. G., Hoogeboom, E., and Welling, M.  E(n) equivariant graph neural networks, 2021.
  * Schütt et al. (2017) Schütt, K. T., Kindermans, P.-J., Sauceda, H. E., Chmiela, S., Tkatchenko, A., and Müller, K.-R.  Schnet: A continuous-filter convolutional neural network for modeling quantum interactions, 2017.
  * Schütt et al. (2021) Schütt, K. T., Unke, O. T., and Gastegger, M.  Equivariant message passing for the prediction of tensorial properties and molecular spectra, 2021.
  * Smith et al. (2017) Smith, J., Isayev, O., and Roitberg, A.  Ani-1, a data set of 20 million calculated off-equilibrium conformations for organic molecules.  _Scientific Data_ , 4:170193, 12 2017.  doi: 10.1038/sdata.2017.193.
  * Thölke & Fabritiis (2022) Thölke, P. and Fabritiis, G. D.  Equivariant transformers for neural network based molecular potentials.  In _International Conference on Learning Representations_ , 2022.  URL https://openreview.net/forum?id=zNHzqZ9wrRB.
  * Thomas et al. (2018) Thomas, N., Smidt, T., Kearnes, S., Yang, L., Li, L., Kohlhoff, K., and Riley, P.  Tensor field networks: Rotation- and translation-equivariant neural networks for 3d point clouds, 2018.
  * Townshend et al. (2021) Townshend, R. J. L., Vögele, M., Suriana, P., Derry, A., Powers, A., Laloudakis, Y., Balachandar, S., Jing, B., Anderson, B., Eismann, S., Kondor, R., Altman, R. B., and Dror, R. O.  Atom3d: Tasks on molecules in three dimensions, 2021.
  * Unke et al. (2021) Unke, O. T., Chmiela, S., Gastegger, M., Schütt, K. T., Sauceda, H. E., and Müller, K.-R.  Spookynet: Learning force fields with electronic degrees of freedom and nonlocal effects.  _Nature Communications_ , 12(1), Dec 2021.  ISSN 2041-1723.  doi: 10.1038/s41467-021-27504-0.  URL http://dx.doi.org/10.1038/s41467-021-27504-0.
  * Vaswani et al. (2017) Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I.  Attention is all you need, 2017.
  * Veličković et al. (2018) Veličković, P., Cucurull, G., Casanova, A., Romero, A., Liò, P., and Bengio, Y.  Graph attention networks, 2018.
  * Weiler et al. (2018) Weiler, M., Hamprecht, F. A., and Storath, M.  Learning steerable filters for rotation equivariant cnns, 2018.
  * 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 bad for graph representation?, 2021.
  * Zemla (2003) Zemla, A.  Lga – a method for finding 3d similarities in protein structures.  _Nucleic acids research_ , 31:3370–4, 08 2003.  doi: 10.1093/nar/gkg571.
  * Zhao et al. (2020) Zhao, H., Jia, J., and Koltun, V.  Exploring self-attention for image recognition, 2020.

## Encoder Details and Hyperparameters

All EQGAT models in this paper were trained on a single Nvidia Tesla V$100$
GPU.

Table 3: Description of architectural parameters on the QM9 and ATOM3D benchmarks. Parameter | QM9 | LBA | PSR | RSR
---|---|---|---|---
Learning rate (lr.) | $5\cdot 10^{-4}$ | $10^{-4}$ | $10^{-4}$ | $10^{-4}$
Maximum epochs | $300$ | $20$ | $20$ | $20$
Early stopping patience | $20$ | $10$ | $10$ | $10$
Lr. patience | $5$ | $4$ | $4$ | $4$
Lr. decay factor | $0.75$ | $0.5$ | $0.5$ | $0.5$
Batch size | $128$ | $16$ | $16$ | $16$
Num. layers | $7$ | $5$ | $5$ | $5$
Num. RBFs | $20$ | $16$ | $16$ | $16$
Cutoff [Å] | $5.0$ | $4.5$ | $4.5$ | $4.5$
Scalar channels $F_{s}$ | $128$ | $128$ | $128$ | $128$
Vector channels $F_{v}$ | $32$ | $16$ | $16$ | $16$
Num. parameters |  $1.1$M |  $741$k |  $795$k |  $795$k

Our model implements a Layer Normalization tailored for scalar and vector
features as proposed by (Jing et al., 2021b) and is applied at the beginning
of every EQGAT convolutional layer.

##  Appendix A Proof Equivariance

We prove the rotation equivariance for the two tensors
$\\{\vec{\mathbf{y}}_{0,ji},\vec{\mathbf{y}}_{1,ji}\\}$ in Equations (19) and
(20) by analyzing the properties of the tensor product $\otimes$. To recap,
the computation is proceed as follows

$\displaystyle\vec{\mathbf{y}}_{0,ji}=\vec{p}_{e,ji}\otimes\mathbf{m}_{v_{0},ji}=\vec{p}_{e,ji}\mathbf{m}_{v_{0},ji}^{\top}~{}~{}\in\mathbb{R}^{3\times
F_{v}},$

as well as

$\displaystyle\vec{\mathbf{y}}_{1,ji}=(\mathbf{1}\otimes\mathbf{m}_{v_{1},ji})\odot\vec{\mathbf{a}}_{v,ji}=(\mathbf{1}\mathbf{m}_{v_{1},ji}^{\top})\odot\vec{\mathbf{a}}_{v,ji}~{},$

where
$\vec{\mathbf{a}}_{v,ji}=\vec{\mathbf{v}}_{j}\times\vec{\mathbf{v}}_{i}+\vec{\mathbf{v}}_{v,j}$
is obtained via the cross-product between two vector features.
Since geometric tensors considered in this work are of type-$1$ only, the
group action represented as rotation matrix $\mathbf{R}\in\textsc{SO}(3)$ only
acts on $\vec{p}$ as well as $\vec{\mathbf{v}}$, where we excluded indices for
brewity and better reading. If the point cloud is rotated, as defined in Eq.
(3), (relative) position as well vector features change to

$\displaystyle\vec{p}\xrightarrow{G}\mathbf{R}\vec{p}~{},$
$\displaystyle\vec{\mathbf{v}}\xrightarrow{G}\mathbf{R}\vec{\mathbf{v}}~{},$
(24)

while the cross product between two vector features
$\vec{\mathbf{v}}_{0},\vec{\mathbf{v}}_{1}$ is invariant to rotation,
resulting to the property

$\displaystyle(\mathbf{R}\vec{\mathbf{v}}_{0}\times\mathbf{R}\vec{\mathbf{v}}_{1})=\mathbf{R}(\vec{\mathbf{v}}_{0}\times\vec{\mathbf{v}}_{1})~{}.$

The message tensor $\mathbf{m}$ is an $\textsc{SO}(3)$ invariant embedding and
is not affected by any rotation, but its group representation acting on
$\mathbf{m}$ is defined as the identity map.
With definition of the tensor product $\otimes$ in Eq. (7), we can conclude
that

$(\mathbf{R}\vec{p})\otimes\mathbf{m}=(\mathbf{R}\vec{p})\mathbf{m}^{\top}=\mathbf{R}(\vec{p}\mathbf{m}^{\top})=\mathbf{R}(\vec{p}\otimes\mathbf{m})=\mathbf{R}\vec{\mathbf{y}}_{0}~{},$
(25)

which proves rotation equivariance for the first equation.
The second equation changes to the following expression when a rotation is
performed:

$\displaystyle(\mathbf{1}\otimes\mathbf{m})\odot(\mathbf{R}\vec{\mathbf{v}}_{0}\times\mathbf{R}\vec{\mathbf{v}}_{1}+\mathbf{R}\vec{\mathbf{v}})$
$\displaystyle=(\mathbf{1}\mathbf{m}^{\top})\odot\mathbf{R}(\vec{\mathbf{v}}_{0}\times\vec{\mathbf{v}}_{1})+\mathbf{R}\vec{\mathbf{v}}=(\mathbf{1}\mathbf{m}^{\top})\odot\mathbf{R}\vec{\mathbf{a}}$
(26)
$\displaystyle=\mathbf{R}\left(\mathbf{1}\mathbf{m}^{\top}\odot\vec{\mathbf{a}}\right)=\mathbf{R}\left(\mathbf{1}\otimes\mathbf{m}\odot\vec{\mathbf{a}}\right)=\mathbf{R}\vec{\mathbf{y}}_{1}~{},$
(27)

where we use the property that the elementwise product $\odot$ between
$\mathbf{1}\mathbf{m}^{\top}\in\mathbb{R}^{3\times F_{v}}$ and
$\mathbf{R}\vec{\mathbf{a}}\in\mathbb{R}^{3\times F_{v}}$ is applied
elementwise, allowing to factorization of the rotation matrix $\mathbf{R}$.
As the summation of $\mathbf{R}\vec{\mathbf{y}}_{0}$ with
$\mathbf{R}\vec{\mathbf{y}}_{1}$ is an equivariant function, we conclude that
$\mathbf{R}(\vec{\mathbf{y}}_{0}+\vec{\mathbf{y}}_{1})$ is equivariant.

##  Appendix B QM9

Table 4: QM$9$ Mean Absolute Error (MAE) on the test set. Models displayed with * use different (random) train/val and test splits. For our EQGAT, we report averaged MAEs over 3 runs. Tasks | $\alpha$ | $\Delta\epsilon$ | $\epsilon_{\textsc{HOMO}}$ | $\epsilon_{\textsc{LUMO}}$ | $\mu$ | $C_{\nu}$ | $G$ | $H$ | $\langle R^{2}\rangle$ | $U$ | $U_{0}$ | ZPVE
---|---|---|---|---|---|---|---|---|---|---|---|---
Units | bohr3 | meV | meV | meV | $D$ | $\frac{\text{cal}}{\text{mol K}}$ | meV | meV | bohr2 | meV | meV | meV
NMP | .092 | 69 | 43 | 38 | .030 | .040 | 19 | 17 | .180 | 20 | 20 | 1.50
SchNet∗ | .235 | 63 | 41 | 34 | .033 | .033 | 14 | 14 | .073 | 19 | 14 | 1.70
Cormorant | .085 | 61 | 34 | 38 | .038 | .026 | 20 | 21 | .961 | 21 | 22 | 2.03
LieConv | .084 | 49 | 30 | 25 | .032 | .038 | 22 | 24 | .800 | 19 | 19 | 2.28
DimeNet++${}^{\text{*}}$ | .044 | 33 | 25 | 20 | .030 | .023 | 8 | 7 | .331 | 6 | 6 | 1.21
SE(3) Tr. | .142 | 53 | 35 | 33 | .051 | .054 | - | - | - | - | - | -
E($n$)-GNN | .071 | 48 | 29 | 25 | .029 | .031 | 12 | 12 | .106 | 12 | 11 | 1.55
ET∗ | .010 | 38 | 21 | 18 | .002 | .026 | 7.64 | 6.48 | .015 | 6.30 | 6.24 | 2.12
PaiNN${}^{\text{*}}$ | .045 | 46 | 28 | 20 | .012 | .024 | 7.35 | 5.98 | .066 | 5.83 | 5.85 | 1.28
EQGAT | .063 | 44 | 26 | 22 | .014 | .027 | 12 | 13 | .257 | 13 | 13 | 1.50
EQGAT∗ | .053 | 32 | 20 | 16 | .011 | .024 | 23 | 24 | .382 | 25 | 25 | 2.00

For the QM$9$ benchmark (Ramakrishnan et al., 2014), we used the pre-processed
dataset partitions from (Anderson et al., 2019; Satorras et al., 2021).
Following (Satorras et al., 2021), we normalized all properties by subtracting
the mean and dividing by the Mean Absolute Deviation. For all targets, we only
retrieve the scalar-embedding $\mathbf{s}^{(L)}_{i}\in\mathbb{R}^{128}$ from
the last GNN layer and pool all (scalar)-node embeddings through summation to
obtain the graph embedding

$\mathbf{g}=\sum_{i}\mathbf{s}_{i}^{(L)}~{}.$ (28)

The graph embedding $\mathbf{g}$ is consequently used to predict the target
property using a 2-layer MLP following the steps as
$\mathbf{g}\xrightarrow{}\textsc{Linear(128,
128)}\xrightarrow{}\textsc{SiLU}\xrightarrow{}\textsc{Linear(128,
1)}\xrightarrow{}\textsc{Property}$.
We further note that our proposed EQGATs produce results on par with the best
performing methods on the non-energy variables, however perform inferior to
the state of the art on the energy variables $(G,H,U,U_{0})$. We hypothesize
that these targets could benefit from more involved problem-tailored output-
layers and additional optimization strategies, such as those compared against
with.
As the comparison of our architecture to recent methods is not straightforward
due to the usage of varying data splits, we also performed the experiments on
random splits from the QM9 dataset, by creating splits with the same size as
in the initial setting of (Anderson et al., 2019). We noticed performance
discrepancy between the EQGAT and EQGAT* model on the energy predictions, when
the later model was trained on 3 random splits of the entire QM9 dataset. As
we mostly refer and compare to the (Anderson et al., 2019) datasplits
retrieved from the Github repository of the authors from (Satorras et al.,
2021), we observe similar perfomance between EQGAT and E$(n)$-GNN, while ours
performs better on the (norm of the dipole moment) $\mu$ target, which is in
fact a vector-quantity.
For better comparison, we hope that in the future pre-processed sets for the
training-, validation- and test splits are used.

##  Appendix C ATOM3D

We use the provided datasplits from the ATOM3D benchmark (Townshend et al.,
2021) and their open-source python package to process the training, validation
and test splits. The LBA dataset consists on average of $382$ atoms, while the
PSR and RSR datasets contains larger graphs with an average size of $1384$ and
$2161$. Statistics are taken from (Jing et al., 2021a).
For the LBA task, we extract the scalar and vector embeddings,
$\mathbf{x}^{(L)}=(\mathbf{s}^{(L)},\vec{\mathbf{v}}^{(L)})$ from the GNN
encoder’s last layer and implement a 2-layer MLP node decoder using Gated
Equivariant layers as proposed in (Schütt et al., 2021) that outputs scalar
and vector features of dimensionality 1, i.e., ${\tilde{s}}_{i}\in\mathbb{R}$
and $\vec{\tilde{v}}_{i}\in\mathbb{R}^{3}$ for all nodes $i=1,\dots,N.$ For
the final prediction $\hat{y}$ we calculate

$\hat{y}=\sum_{i=1}^{N}||\tilde{s}_{i}\otimes\vec{p}_{i}+\vec{\tilde{v}}_{i}||_{2}^{2}~{},$
(29)

assuming the center of mass being $\vec{p}=0$.

The decoder network for the PSR and RSR tasks are implemented using a single
layer of Gated Equivariant Layer on the GNN encoder’s last layer output
$\mathbf{x}_{i}^{(L)}$ returning $\tilde{\mathbf{x}}_{i}^{(L)}$ which is then
required for mean-pooling, where only the output scalar feature
$\mathbf{s}=\frac{1}{N}\sum_{i=1}^{N}\mathbf{\tilde{s}}_{i}$ is used as input
for a 2-layer MLP as follows $\mathbf{s}\xrightarrow{}\textsc{Linear(128,
64)}\xrightarrow{}\textsc{SiLU}\xrightarrow{}\textsc{Linear(64,
1)}\xrightarrow{}\textsc{Property}$.
As described in the E($n$)-GNN (Satorras et al., 2021) model trained on the
QM$9$ dataset, we implement and train the invariant E($n$)-GNN model without
the positional coordinates update equation on all three ATOM3D tasks. To
enable stable training for the E($n$)-GNN model, we further transform raw
distances between nodes $j$ and $i$ using a reciprocal function
$e(d_{ji})=\frac{1}{1+d_{ji}}$ as opposed to the authors initial
implementation using the squared euclidean distance.
We noticed that including layer normalization between consecutive
convolutional layers for the E$(n)$-GNN and PaiNN architectures enabled stable
training on the PSR and LBA benchmark datasets. This is mostly related to the
fact that macromolecules consists of larger neighbourhoods when a radial
cutoff is applied. Nonetheless, we still perform a sum-aggregation for the
message embedding of target node $i$, such that the magnitude of neighboring
embeddings is recognized by the network.