Geometric and Physical Quantities improve E(3) Equivariant Message Passing
Authors: Johannes Brandstetter
Rob Hesselink
Elise van der Pol
Erik J Bekkers
Max Welling
Venue: ICLR 2022 Spotlight (2022)
Keywords: equivariant graph neural networks
steerable message passing
non-linear convolutions
molecular modeling
covariant information

Abstract: Including covariant information, such as position, force, velocity or spin is important in many tasks in computational physics and chemistry. We introduce Steerable E($3$) Equivariant Graph Neural Networks (SEGNNs) that generalise equivariant graph networks, such that node and edge attributes are not restricted to invariant scalars, but can contain covariant information, such as vectors or tensors. Our model, composed of steerable MLPs, is able to incorporate geometric and physical information in both the message and update functions.
Through the definition of steerable node attributes, the MLPs provide a new class of activation functions for general use with steerable feature fields. We discuss ours and related work through the lens of equivariant non-linear convolutions, which further allows us to pin-point the successful components of SEGNNs: non-linear message aggregation improves upon classic linear (steerable) point convolutions; steerable messages improve upon recent equivariant graph networks that send invariant messages. We demonstrate the effectiveness of our method on several tasks in computational physics and chemistry and provide extensive ablation studies.

Published as a conference paper at ICLR 2022
GEOMETRIC AND PHYSICAL QUANTITIES IMPROVE E(3) EQUIVARIANT MESSAGE PASSING
Johannes Brandstetter" University of Amsterdam Johannes Kepler University Linz brandstetter@ee jku.
Rob Hesselink*
University Amsterdam r.d. www.shesssssssses.
Elise van der Pol UvA-Bosch DeltaLab University of Amsterdam e. e. vanderpoleuva.....
Erik J Bekkers
Max Welling
University of Amsterdam e. www.kekseuuaa..... nl
UvA-Bosch DeltaLab University of Amsterdam m. we wellingeuva
ABSTRACT
Including covariant information, such position, force, velocity spin is impor- tant in many tasks in computational physies and chemistry. We introduce Steerable E(3) Equivariant Graph Neural Networks (SEGNNs) that generalise equivariant graph networks, such that node and edge attributes are not restricted to invariant scalars, but can contain covariant information, such as vectors or tensors. This model, composed of steerable MLPs, is able to incorporate geometric and physical information both the message and update functions. Through the definition steerable node attributes, the MLPs provide new class of activation functions for general use with steerable feature fields. We discuss ours and related work through the lens of equivariant non-linear convolutions. which further allows pin-point the successful components SEGNNs: non-linear message aggregation improves upon classic linear (steerable) point convolutions; steerable messages improve upon recent equivariant graph networks that send invariant messages. We demonstrate the effectiveness of our method on several tasks in computational physics and chemistry and provide extensive ablation studies.
INTRODUCTION
The success of Convolutional Neural Networks (CNNs) (LeCun et al., 1998 2015 Schmidhuber 2015) Krizhevsky et al. 2012) is key factor for the rise of deep learning, attributed to their capability of exploiting translation symmetries, hereby introducing a strong inductive bias. Recent work has shown that designing CNNs to exploit additional symmetries via group convolutions has even further
increased their performance (Cohen & Welling 2016 2017, Worrall et 2017 Cohen al. [ Weiler et 2018 Bekkers et al.) 2018 Bekkers) 2019 Weiler & Cesa 2019). Graph neural networks (GNNs) and CNNs are closely related to each other via their aggregation of local information. More precisely, CNNs can be formulated as message passing layers Gilmer et al. 2017) based on a sum aggregation of messages that are obtained by relative position-dependent linear transformations of neighbouring node features. The power of message passing layers is, however, that node features are transformed and propagated a highly non-linear manner. Equivariant GNNs have been proposed before either PointConv-type (Wu al. 2019)
2018, Kondor& & Trivedi] 2018)
Kristof et al. 2017) implementations of steerable (Thomas et al.) 2018 Anderson et al.) 2019 Fuchs
al.] 2020) regular group convolutions (Finzi et 2020 The most important component in these methods are the convolution layers. Although powerful, such layers only (pseudd) linearly transform the graphs and lineaiiy is only obtained via point-wise activations.
Methods such SE(3) transfforossss Fuuhss et al. 2020) and Cormorant Anderson et al. 2019 include an input-dependent attention component that augments the convolutions.

---

Published as a conference paper at ICLR 2022
In this paper, we propose non- non-linear E(3) equivariant message passing layers using the same principles that underlie steerable group convolutions, and view them as non-linear group convolutions. Central to our method is the use of steerable vectors and their equivariant transformations to represent and process node features; we present the underlying mathematics of both in Sec. 2 and illustrate in Fig. 1 on molecular graph. As consequence, information at nodes and edges can now be rotationally invariant (scalar) or covariant (vector, tensor). In steerable message passing frameworks, the Clebsch-Gordan (CG) tensor product is used steer the update and message functions by geometric information such as relative orientation (pose). Through notion of steerable node attributes we provide new class of equivariant activation functions for general use with steerable feature fields (Weiler et al. [0188] Thomas et al. 2018). Node attributes can include information such as node velocity, force, or atomic spin. Currently, especially molecular modelling, most datasets are build up merely of atomic number and position information. In this paper, we demonstrate the potential of enriching node attributes with more geometric and physical quantities. We demonstrate the effectiveness of SEGNNs by setting new state of the art n-body toy datasets, in which our method leverages the abundance of geometric and physical quantities available. We further test our model the molecular datasets QM9 and OC20. Although here only (relative) positional information is available as geometric quantity, SEGNNs achieve state of the art on the IS2RE dataset of OC20, and competitive performance QM9, For all experiments we provide extensive ablation studies.
The main contributions of this paper are: (i) A generalisation of equivariant GNNs such that node and edge attributes are not restricted scalars. (ii) new class of equivariant activation functions for stecrable vector fields, based on the introduction of steerable node attributes and steerable multi-layer perceptrons, which permit the injection of geometric and physical quantities into node updates. (iii) A unifying view on various equivariant GNNs through the definition of non-linear convolutions. (iv) Extensive experimental ablation studies that shows the benefit of steerable over non-steerable (invariant) message passing, and the benefit of non-linear over linear convolutions.
Figure 1: Commutation diagram for an equivariant operator ø applied to a 3D molecular graph with steerable node features (visualised as spherical functions); As the molecule rotates, so do the node features. The use of steerable vectors allows neural networks to exploit, embed, or learn geometric cues such as force and velocity vectors.
2 GENERALISED E(3) EQUIVARIANT STEERABLE MESSAGE PASSING
Message passing networks. Consider a graph 9 = (V, E). with nodes v; E V and edges eij E E, with feature vectors f, E Re attached to each node, and edge attributes a, E Re attached each edge. Graph neural networks (GNNs) (Scarselli et al.] 2009 Kipf & Welling 2017 Defferrard et al. 2016 Battaglia et al. 2018) are designed to learn from graph-structured data and are by construction permutation equivarant with respect to the input. A specific type of GNNs are message passing networks (Gilmer et al.) 2017), where layer updates node features via the following steps:
compute message m.j from node by ey:
maj Φm (fe, (..f),aa))
(I)
aggregate messages and update node features o;:
f =s f, Σ ma
(2)
JEN(()
where M(1) represents the set of neighbours of node v. and on and o, are commonly parameterised by multilayer perceptrons (MIPs).
2

---

Published as a conference paper at ICLR 2022
Equivariant message passing networks. Our objective is to build graph neural networks that are robust to rotations, reflections, translations and permutations. This is a desirable property since some prediction tasks, such molecular energy prediction, require E(3) invariance, whereas others, like force prediction, require equivariance. From a technical point of view, equivariance of a function ø to certain transformations means that for any transformation parameter g and all inputs x have Ty [[(=)] = [x) where T' and T3g denote transformations on the input and output domain of o, respectively. Equivariant operators äpplied to atomic graphs allow us preserve the geometric structure of the system as well as enriching it with increasingly abstract directional information. We build E(3) equivariant GNNs by constraining the functions om and Φy of Eqs. 2 to be equivariant, which in return guarantees equivariance of the entire network. In the following, we introduce the core components behind our method; full mathematical details and background can be found App.
Steerable features. In this work, we achieve equivariant graph neural networks by working with steerable feature vectors, which we denote with a tilde, e.g. vector h is steerable. Steerability of a vector means that for a certain transformation group with transformation parameters g, the vector transforms via matrix-vector multiplication D(g).. For example, a Euclidean vector in R is steerable for rotations g R SO(3) by multiplying the vector with rotation matrix, thus ggg) = R. We are however not restricted to only work with 3D vectors; via the construction of steerable vector spaces,
generalise the notion of 3D rotations arbitrarily large vectors.
Central to our approach is the use of Wigner-D matrices (g) These are (21 + I x 21 + 1)- dimensional matrix representations that act on dimenssioal vector spaces. These vector spaces that are transformed by (-th degree Wigner-D matrices will be referred to type-l steerable vector spaces and denoted with V/- We note that we can combine two independent steerable vector spaces Vi, and Vi, of type 4 and 22 by the direct sum, denoted by V Vi, @ Viz- Such combined vector space then transforms by the direct sum of Wigner D-matrices, i.e., via D D (g) + D (g). which is a block-diagonal matrix with the Wigner-D matrices along the diagonal. We denote the direct sum of type-l vector spaces up degree 1 = L by V, == Vo @ V11 $ + Vt. and n copies of the the same vector space with nV := V + V .... + V. Regular MLPs are based on transformations
n times
between di d-dimensional type-0 vector spaces i.e., R = dVo. and are special case of our steerable MLPs that act on steerable vector spaces of arbitrary type.
Steerable MLPs. Like regular MLPs, steerable MLPs are constructed by interleaving linear map- pings mmntrix-vecttr multiplications) with non nin-lineasi Now however, the linear maps transform between steerable vector spaces layer have the same functional form as regular MIPs, although, in our case, the linear transformation matrices W;, a. defined below, are conditioned on geometric information (e.g. relative atom positions) which is encoded in the steerable vector å. Both vectors å and h are steerable vectors. In this work we will however use the vector the information flow of h through the network. In order to guarantee that W maps between steerable vector spaces, the matrices are defined via the Clebsch- Gordan tensor product. By construction the resultant MLPs are equivariant for every transformation parameter g via
to layer i via h Wi Steerable MLPs thus
have geometric and structural information encoded and the steer
()) = (L.
provided that the steerable vectors á that condition the MLPs are also obtained equivariantly.
Spherical harmonic embedding of vectors. We can convert any vector x E E into a type-l vector through the evaluation of spherical harmonics yi()) S2 + R at x
For any x E TR3
T
ā (() = (!)
(4)
is type-/ steerable vector. The spherical harmonic functions ( are functions on the sphere S and we visualise them as such in Figure? 2 geometric and physical information into steerable MLPs
We will use spherical harmonic embeddings to include
Mapping between steerable vector spaces. The Clebsch- Gornha (CG) tensor product &ey Ví, Vi, + Vi is a bilinear operator that combines two 0(3) steerable input vectors of types I and
II order be 0(3) not just SO(3) equiva-auv,,,,,nt,, we include reflections. See App. A for detail.
3

---

Published as a conference paper at ICLR 2022
h0
(x)
C-)
Σ (C
149 (())) M
x2(()) Σ En
14
Figure 2: Left: representation of an 0(3) steerable vector h Vz Vo @ Vi V2, a spherical harmonic embedding of vector x, e.g. relative orientation, velocity or force. Right: for each subspace the embedding with the basis functions yI is shown (a). The transformation of h via D (g) acts on each subspace of Vi separately (b).
l2 and returns another steerable vector of type 1. Let h E Vi = 2R2++ denote a steerable vector of type and bm its components with m -, -1 1 1....,I.
I. The CG tensor product is given by
4
( m
(1)) (ta) Pi
(5)
in which w is learnable parameter that scales the product, and Cil.m) Gordan coefficients that ensure that the resulting vector is type-I steerable. The CG tensor product is a sparse tensor product, as generally many coefficients are zero. Most notably, C(l.m)
(l..m))aa.m) are the Clebsch-
(ts,m,)((2,mm) whenever < 141 12l or / > 1 + l2. While Eq. (5) only describes the product between steerable vectors of single type, c.g. h(.)) E Vi, and h (1x) E VV2, it is directly extendable mixed type steerable vectors that may have multiple chamnels/muliipiccttesss within a type. In this case, every input to output sub- vector pair gets its index in similar way the weights in a standard linear layer are indexed with input-output indices. We then denote the CG product with W with boldfaced W to indicate that it is parametrised by a collection of weights. In order to stay close to standard notation used with MLPs, we treat the CG product with a fixed vector à in one of its inputs as steerable linear layer conditioned å, denoted with
Wah h @ey à or
Wa(d)h = h cy W(d) à
where the latter indicates that the CG weights depend on some quantity d, e.g. relative distances.
Steerable activation functions. The common recipe for deep neural networks is alternate linear layers with element-wise non-linear activation functions. In the steerable setting, careful consideration is required to ensure that the activation functions are equivariant; currently available classes of activations include Fourier-based (Cohen et al.) 20188), Norm-altering (Thomas et al., 2018). or gated non-linearities (Weiler et al.) 2018). We use gated -oneariassss in all our architectures and shortly discuss their working principle in Sec.CCinn the appendix. The resulting steerable MI.Ps themselves in turn provide a new class of steerable activation functions, that is the first its kind in directly leveraging local geometric cues. Namely, through steerable node attributes å, either derived from the physical setup (forces, velocities) or from predictions (similar to gating). The MLPs can be applied node-wise and be generally used in steerable feature fields as non-linear activations.
2.1 STEERABLE E(3) EQUIVARIANT GRAPH NEURAL NETWORKS
We extend the message passing equations ( 2 and define a message passing layer that updates steerable node features f, E Vz at node v via the following steps:
mij- Φm ff, f, |||
x; xill2,
P. Φs f Σ
mijs à;
(8)
JEN(()
4

---

Published as a conference paper at ICLR 2022
Here, |x, - x2112 is the squared relative distance between two nodes 24 and 1/y, 9m and Φ, are O(3) steerable MLPs, äty E VE and å, E Vz are stecrable edge and node attributes. If additional attributes exist, such as pair-wise distance ||x; x;Il, they can either be concatenated to the attributes that condition our steerable MIPs, or as is more commonly done (Sec. (3 add as inputs to on and Øy- We do the latter, and stack all inputs to single steerable vector by which the steerable MLP layer is given: h = a oWaaaWa a. hi)), with hς = fr. f, |xx xsll²) E V + Vj @ Vo. where Vj is the user specified steerable vector space of node representations. The message network Φm is steered via edge attributes ā.j. and the node update network dy is similarly steered via node attributes ä,.
Injecting geometric and physical quantities. In order to make SEGNNs more expressive, we include geometric and physical information in the edge and node updates. For that purpose, the edge attributes are obtained via the spherical harmonic embedding (Eq. (4)) of relative positions, in most cases, but possibly also relative force or relative momentum. The node attributes could e.g. be the average edge embedding of relative positions over neighbours of a node, i.e., å W Ejev(o) aij. and could additionally include node force, spin or velocities, as we do in the N-body experiment. The use of steerable node attributes in the steerable MLPs that define Øy allows us to not just integrate geometric cues into the message functions, but also leverage in the node updates. We observe that the more geometric and physical quantities are injected the better SEGNNs perform.
MESSAGE PASSING AS CONVOLUTION, RELATED WORK
Recent literature shows a trend towards building architectures that improve performance by means of maximally preserving equivariance through groups convolutions (either in regular censor-product form, see [B.. Convolutions, however, are "just" linear operators and non-linearities are only introduced through point-wise activation functions. This is in contrast to architectures that are built without explicit use of group convolutions, but instead rely on the highly non-linear framework of message passing. In the following we show that many related works are connected through a notion of non-linear convolution, term that we coin based on the following. Any linear operator which equivariant is a group convolution (Kondor & Trivedi] 2018 Bekkers 2019) and their discrete implementations can be written in message passing form. We then call any non- lon-linear operator that is equivariant, and which can be written in simple message passing form, a non-linear convolution. This framing allows to place related work in a unifying context and to identify two important aspects of successful architectures: (i) equivariant layers improve upon invariant ones and (ii) non- non-linear layers improve upon linear ones. Both come together in SEGNNs via steerable non- non-linear convolutions.
Point convolutions as equivariant linear message passing, Consider a feature map f IRd + R A convolution layer (defined via cross-correlation)) with point-wise non-linearity is given by
((() -00 (xx)=o(A ( -[xxx=))((xxddx xx)((x)dxx
with W Rd + 1 a convolution kernel that provides for every relative position a matrix that linearly transforms features from RR R Point convolutions, generally referred to as PointConvs (Wu et al. 2019). and SchNet (Kristof et al., 2017) implement Eq. (9) on point clouds. For sparse input feature map consisting of ceration-feature pairs (x,, f,), the point convolution is given by f - E)eN(() W(x -x,)f, which describes linear message passing layer of Eqs. 2
in which the messages are m.j W(x, - xx and the message update f = E, myj- In the above convolutions, the transformation matrices W are conditioned on relative positions Xy- x,, which is typically done in one of the following three approaches. (i) Classically, feature maps are processed on dense regular grids with shared neighbourhoods and thus the transformations W.j can be stored for finite set of relative positions Xj x in a single tensor. This method however does not generalise to non-uniform grids such as point clouds. Continuous kernel methods parametrise the transformations either by (ii) expanding W into a continuous basis or (iii) directly parametrising them with MLPs via W (xx, x;) MIP((x, x:). Steerable kernel method are of type (ii) and rely on a steerable
*See [Lagg & Weiler 2020 for general theory for G-steerable kernel constraints for compact groups.
5

---

Published as a conference paper at ICLR 2022
basis, such 3D spherical harmonics Yn via
W(x; x,)
Σ Σ
( -x;1-xx+111)0 x (x- x:)
(10)
with basis coefficients Wh typically depending on pair-wise distances. We next show how such kernels connect to steerable group convolutions and provide a detailed background in App. B
Steerable (group) convolutions. In context of the steerable framework of Sec. 2 linear feature
transformations W are equivalent to steerable linear transformations (Eq. ( conditioned on the scalar "1". L.e, with h, E Rd in the usual and h E dVo the steerable setting, messages are obtained by
mij- W(x; x;)h;
m:j W x,)
(11) f
When the transformations parametrised in steerable basis 10) we can make the identification"
m ((,, x,) h;)
m.j-
Wa,, (1xx; x:1l)b;
(12)
in which à, X, xx) are spherical harmonic embeddings (Eq. (4) of x, ---x,, X,, and the weights that parametrise the CG tensor product depend on pair-wise distances. We can thus perform convolutions either in the regular or in the steerable setting, where the latter has the benefit of allowing us to directly derive what the convolution result would be if the kernel were to be rotated via
my W(R (x, - -x,)
my = D(R) a., (Ixx, - x,I)h,,
(13)
Steerable vectors obtained via convolutions with steerable kernels thus generate signals on 0(3) via the Wigner-D D matrices. This relation is in fact established by the inverse Fourier transform on O(3)
(ef. App. A). by which we can treat steerable feature vectors f; at each location as a function on the group 0(3). Steerable convolutions thus produce feature maps on the full group H(3) that for every possible translation/positinn this mechanism of transforming convolution kernels via the group action (via W(g xx,) = WR (xx x)) that underlies group convolutions (Cohen & Welling 2016). Message passing via explicit kemel rotations (l.h.s. of (13) corresponds to regular group convolations, and via steerable transformations (r.h.s. of (13) to steerable group convolutions.
and rotation R provide a feature response f(x, R). It is precisely
The equivariant steerable methods (Thomas et al. 2018 Anderson et al. 2019: Miller et al., 2020
Fuchs et al. 2020) that we compare against in our experiments can all be written in convolutionn form
Σ JEN(()
Wa, (1x; - x; xiIl);,
or
F
Wa. (fr, f,, ||x xx1))f,
(14)
JEN()
where, in the latter case, the linear transformations additionally depend on an input dependent at- tention mechanism as in (Anderson et al. [2019] Fuchs et al. (020), and can be seen as steerable PointConv version of attentive group convolutions (Romero et al.) 2020). these attention-based cases, convolutions are augmented with input dependent weights a via Wa., (f,, f,, ||x; - x:
(ff,, ) a.y (l|x; x=l)). This makes the convolution non-linear, however, the transformation of input features still happens linearly and thus describes what one may call a pseudo-linear transforma- tion. Finally, the recently proposed LieConv (Finzi et al. 2020) and NequlP (Batzner et al. 1 also fall the convolutional message passing class. LieCony is a PointConv-type variation of regular group convolutions on Lie groups (Bekkers) 2019). NeuqIP follows the approach of Tensor Field Networks (Thomas et al.) 2018). and weighs interactions using an MLP with radial basis functions as input. These functions are obtained solution of the Schrödinger equation. Finally, steerable methods fall into a more general class of coordinate independent convolutions (Weiler et al. 2021 which allow to ensure equivariance locally, even when global symmetries can not be defined.
Equivariant message passing as non-linear convolution. EGNNs (Satorras et al., 2021) are equivariant transformations in E(n) and outperform most aforementioned steerable methods. This is somewhat surprising as it sends invariant messages, which are obtained via MLPs of the form
mij MLP(f, fy, |xx x,1 = (.. (,))))
(15)
*Exact correspondence is obtained by sum reduction over the steerable vector components (App. B

---

Published as a conference paper at ICLR 2022
where h, = f,,,,,,,,,, -x,1 These messages resemble the convolutional messages of point convolutions due to their dependency on relative positions. There are, however, two important differences: (i) the messages are non-linear transformations of the neighbouring feature values f via an ML. and (ii) the messages are only conditioned on the distance between point pairs, and are therefore F(n) invariant. As such, regard EGNN layers as non-linear convolutions with isotropic message functions (the non-linear counterpart of rotationally invariant kernels). In our work, we lift the isotropy constraint and generalise non-linear steerable convolutions via messages of the form
oj = MLPa., a., (f, 1x, xx(1²) = ( (m) mug (... (a( (1) Ay h,))))
(16)
with h; = (f., f x x2 E Vy @ Vy @ Vo. The MLP is then conditioned on attribute å.j which could e.g. be a spherical harmonic embedding of x, -x,. This allows for the creation of messages more general than those found in convolution, while carrying covariant geometrical information.
Related equivariant message passing methods. A different but also fully message passing based approach can be found in Geometric Vector Perceptrons (GVP) (Jing et al. 2020). Vector Neurons (Deng et al. 202T), ) and PaiNN (Schütt et al. 202T). Compared to SEGNN Ns which treat equivariant information as fully steerable features, these architectures update realar-values attributes using the norm of vector-valued attributes, and therefore with 0(3) invariant information. These methods restrict the flow of information between attributes of different types, whereas the Clebsch- Gordan tensor product in SEGNNs allows for interaction between spherical harmonics of all orders throughout the network. Methods such as Dimenet++ (Klicpera et al. (2020). SphereNet (Liu et al.) 2021), and GemNet (Klicpera et al., 2021 incorporate relative orientation in a second-order message passing scheme that considers angles between the central point and neighhours-of -f-uus--o--nnnoouuss. In contrast, our method directly leverages angular information in a first-order message passing scheme using stcerable vectors. We attribute the success of such methods to the fact that they equivariantly process point clouds of local orientations (x,, ray) E RS x S², defined by relative positions between atoms, by sending messages between edges (local orientations) rather than nodes. As such, they can be thought of as non- non-linear regular group convolutions on the homogeneous space of positions and orientations (R33 5²) with isotropic (zonal) message functions, where the symmetry constraint is induced by the quotient RR x S = SE(3)/SO(2) (Bekkers) 2019 Thm. 1).
4 EXPERIMENTS
Implementation details. The implementation of SEGNN's O(3) steerable MI.Ps is based on the e3nn library (Geiger et al., 2021a). We either define the steerable vector spaces as V = nV (N-body, QM9 experiments), I.e., copies of steerable vector spaces up to order Imaz: or by dividing an n-dimensional vector space V into L approximately equally large type-l sub-vector spaces (OC20 experiments) as is done in Finzi et al. () Furthermore, for a fair comparison between experiments with different maa. we choose n such that the total number of weights in the CG products corresponds to that of regular (type-O) linear layer. Further implementation details are in App. C
SEGNN architectures and ablations. We consider several variations of SEGNNs. On all tasks we have at least one fully steerable (l > 0, > 0) SEGNN tuned for the specific task hand. We perform ablation experiments to investigate two main principles that sets SEGNNs apart from the literature. A1 The case of non-steerable vs steerable EGNNs obtained by applying the same SEGNN network with different specifications of maximal spherical harmonic order in the feature vectors (l)) and the attributes (l,). EGNN (Satorras et al.) (202)) arises as a special case with I la O. These models will be labelled SEGNN. A2 In this ablation, we use steerable equivariant point cony methods (Thomas et al., 2018) with messages as in Eq. [4] and regular gated non- non-linearities as activation/update function, labelled SF and compare it to the same network but with messages obtained non-linear manner via 2-layer steerable MI.Ps as in Eq. (16), labelled as SEnon linear-
N-body system. The charged N-body particle system experiment (Kipf et al. 2018) consists of 5 particles that carry positive or negative charge, having initial position and velocity in 3- dimensional space. The task is to estimate all particle positions after 1.000 timesteps. We build upon the experimental setting introduced in (Satorras et al., 2021). Steerable architectures are designed such that the parameter budget at I, input the relative position to the center of the system and the velocity as vectors of type | = 1 with odd parity. We further input the norm of the velocity as scalar, which altogether results in an input
and Ta = I matches that of the EGNN implementation. We
7

---

Published as a conference paper at ICLR 2022
Table Mcan Squared Error (MSE) in the N-body system experiment, and forward time in seconds for a batch size of 100 samples running on GeForce RTX 2080 Ti GPU. Results except for EGNN and SEGNN are taken from (Satorras et al.
2021 and verified. Runtimes are re- measured.
Method
MSE
Time [s]
SE(3)-Tr (Fuchs et al. 2020
.0244
.0742
TFN
Thomas
2018
.0155
.0182
NMP
Gilmer et 2017
.0107
.0017
Radial Freld (Köhler el
2019
.0104
.0019
EGNN (Satorras et al.
2021)
.0070 .00022
.0029
Skinear (1y 2, 2)
.01161 .00021
.0640
SEaweiinr (11 1,
.0660 .00019
.0310
SEGNNG (1/ 1,16
.0566 .00025
.0250
SEGNNo. (1, 1,14
.00431.0000 .00015
.0260
vector h E Vo @ Vi @ V1. The output is embedded as difference vector to the initial position (odd parity), i.e. E V1. In doing so, we keep E(3) equivariance for vector valued inputs and outputs. The edge attributes are obtained via the spherical harmonic embedding of xj Eq. 4. Messages additionally have the product of charges and absolute distance included. SEGNN architectures are compared to steerable equivariant (SEnmear) and steerable non-linear point conv methods (SEmon-lincar) Results and ablation studies are shown in Tab. [ A full ablation is outlined in App. C Stecrable architectures obtain the best results for 1 1 and 4a 1, and don't 't benefit from higher orders of l, and la- We consider a first SEGNN architecture where the node attributes the averaged edge embeddings, i.e. mean only geometric information is used. The second SEGNN architecture has the spherical harmonics embedding of the velocity added to the node attributes. It can thus leverage geometric (orientation) and physical (velocity) information, and is consequently labelled SEGNNG#P. Including physical information in addition to geometric information in the node updates considerably boosts SEGNN performance. We further test SEGNNs on a gravitational 100-body system (Sec. in the appendix).
as described in
relative orientation, labelled SEGNNG since
QM9. The QM9 dataset (Ramakrishnan et al., 2014 Ruddigkeit et al. 2012 2 consists of small
molecules up to 29 atoms. where each atom is described with 3D position coordinates and one-hot mode embedding of its atomic type (H, C, N, O, F). The aim is to regress various chemical properties for each of the molecules, optimising on the mean absolute error (MAE) between predictions and ground truth. We use the dataset partitions from (onderonnooon et al. 2019). Table[]shows SEGNN results on the QM9 dataset. In Table show that by steering with the relative orientation between atoms we observe that for higher (maximum) orders of steerable feature vectors, the performance increases, especially when small cutoff radius is chosen. While previous methods use relatively large cutoff' radii of 4.5-1 11 A, we use a cutoff radius of 2A. Doing so results in a sharp reduction of the number of messages per layer, as shown in App. C Tables 3 together show that SEGNNs outperform an architecturally comparable baseline EGNN, whilst stripping away attention modules from it and reducing graph connectivity from fully connected note that runtime is still limited by the relatively expensive calculation of the Clebsch-Gorde tensor products. We further note that SEGNNS produce results on par with the best performing methods on the non-energy variables, however lag behind state of the art on the energy variables (G, H, U, Uo). We conjecture that such targets could benefit from more involved (e.g. including attention or boighbour-neighbow interactions) or problem- problem-tailored architectures, such as those compared against.
only 2 distant atoms. It is however apt to
OC20. The Open Catalyst Project OC20 dataset (Zitnick et al. 2020 Chanussot et al., 202 I),
consists of molecular adsorptions onto surfaces. We focus on the Initial Structure to Relaxed Energy (IS2RE) task, which takes as input an initial structure and targets the prediction of the energy in the final, relaxed state. The IS2RE training set consists of over 450,000 catalyst adsorbate combinations with 70 atoms on average. Optimisation is done for MAE between the predicted and ground truth energy. Additionally, performance is measured in the percentage of structures in which the predicted energy is within 0.02 eV threshold (EwT). The four test splits contain in-distribution (ID) catalysts and adsorbates, out-of-domain adsorbates (OOD Ads), diitri--ddtiiiiiiinnn catalysts (OOD Cat), and adsorbates and catalysts (OOD Both). Table 44 shows SEGNN results on the OC20 dataset and comparisons with existing methods. We compare to models which have obtained results by training on the IS2RE training set such as SphereNet (Liu et al. (022)) and DimeNet++ (Klicpera
8

---

Published as a conference paper at ICLR 2022
al. 2019 2020). The best SEGNN performance is seen for f = and L. (see App. C A full ablation study, comparing performance and runtime for different orders ly and la is found in App. C
Table 2: Performance comparison on the QM9 dataset. Numbers are reported for Mean Absolute Error (MAB) between model predictions and ground truth.
Task Units
A mcV
CHомO meV
éLUMO meV
G
H
R²
U
Un mcV
ZPVE
bohr
D
cal/mol K
mcV
mcV
bohr
mcV
mcV
NMP
.092
69
43
38
.030
.440
19
17
.180
20
20
1.50
SchNet
.355
63
34
.033
003
14
14
.073
19
14
1.70
Cormorant
085
61
20
21
22
LINet
088
68
35
043
.011
14
14
.354
14
13
1.56
LieConv
.84
49
25
.03
.038
22
24
.800
19
19
2.28
TEN
.22
58
40
38
064
.011
SE(3)-Tr.
142
53
35
33
.051
054
DimeNet+++
.043
32
19
.029
7
6
1.21
SphereNet *
.046
32
18
.06
.021
.292
7
1.12
PaiNN
.045
45
20
.012
.024
7
.066
5
5
1.28
EGNN
.071
48
29
25
.029
.031
12
12
.066
12
12
1.55
SEGNN (Ours)
.060
42
24
21
.023
.031
15
16
.660
13
15
1.62
these methods use different train/val/test partitions.
Table QM9 ablation study to compare SEGNN performances for different (maximum) orders of steerable feature vectors (ly) and attributes (l,). Models with only trivial features akin the invariant EGNN. The method with a fully connected graph uses soft edge estimation (Satorras et al. 2021b. Forward time is measured for a batch of 128 samples running on GeForce RTX 3090 GPU.
Task
As
EhоMO meV
ELUMO meV
μ
C2
Units
Cutoff radius
bohr
meV
D)
cal/mool
Time [s]
(SIEGNN 0,
.091
53
34
28
.042
.043
0.016
(SJEGNN 12
GÅ
.105
57
36
31
.047
.047
0.014
SEGNN dy 1.44
GÅ
.080
49
29
25
.032
.034
0.058
(SIEGNN iy - 0,14
2A
.240
98
60
.340
.077
0.014
SEGNN (1y 1.10
2Å
.074
48
27
25
.031
.035
0.046
SEGNN (1y 2.14
2
.060
42
24
21
.023
.031
0.096
Table 4: Comparison on the OC20 IS2RE task in terms of Mean Absolute Error (MAE) between model predictions and ground truth energy and predictions within 0.02 eV of the ground truth (EwT)). Numbers are taken from the OC20 leaderboard. SEGNNs outperform all competitors.
Energy MAE [eV] 1
EwT +
Model
OOD Ads
OODCat
OOD Both
ID
OOD Ads
OOD
OOD Both
Median baseline
1.7499
1.7993
L.7090
1.6636
0.71%
0.72%
0.89%
0.74%
CGCNN
0.6149
0.9155
0.6219
0.8511
3.40%
1.93%
3.10%
2.00%
SchNet
0.7342
0,7037
2.96%
2.21%
EdgeUpdatcNet
0.5839
0.7252
0.6016
0.6862
3.48%
2.35%
3.30%
2.57%
EnergyNet
0.6366
0.717
0.6387
0.6626
3.30%
2.20%
3.07%
2.34%
DimeNet++
0.5620
0.7252
0.5756
0.6613
4.25%
2.07%
4.10%
2,41%
SphereNet
0.5630
0.7030
0,5710
0.6380
4.47%
2.29%
4.09%
2.41%
SEGNN (Ours)
0.5327
0.6921
0.5369
0.6790
5.37%
2.46%
4.91%
2.63%
5 CONCLUSION
We have introduced SEGNNs which generalise equivariant graph networks, such that node and edge information is not restricted to be invariant (scalar), but can also be vector- or tensor-valued. SEGNNs are the first networks which allow the steering of node updates by leveraging geometric and physical cues, introducing a new class of equivariant activation functions. We demonstrate the potential of SEGNNs by applying it to wide range of different tasks. Extensive ablation studies have further shown the benefit of steerable over non-steerable (invariant) message passing, and the benefit of non-linear linear convolutions. On the OC20 ISRE taks, SEGNNs outperform all competitors.

---

Published as a conference paper at ICLR 2022
6 REPRODUCIBILITY STATEMENT
We have included error bars, reproducibility checks and ablation studies wherever we found necessary and appropriate. For example, for the N-body experiment we have reproduced the results Satorras et al. (2021), we have ablated different (maximum) orders of steerable feature vectors, and we have stated mean and standard deviation of the results which we have obtained by running the same experiments eight times with different initial seeds. For the OC20 experiments, we have stated official numbers from the Open Catalyst Project challenge in the paper. Thus, our comparisons to other methods are obtained from their respective best entries to the challenge. We have further ablated runtimes, different cutoff radii and different (maximum) orders of stecarable feature vectors, Sec. in the appendix. We have done this by running the experiments eight times with different initial random seeds to be able to report mean and standard deviation of the results. For the reproducibility of the QM9 experiments, we have uploaded our code and included a command which reproduces results of the a variable.
We have described our architecture Sec. 2.1 and provided further implementation details in Appendix Section C We have not introduced new mathematical results. However, have used concepts from different mathematical fields and therefore included a self-contained mathematical exposition in our appendix. We have further included proof of equivariance and properties of Wigner D matrices and spherical harmonics in the appendix. In Sec. 3 we have introduced a unifying view on various equivariant graph neural networks through the definition of non-linear convolutions, which allows us to draw comparisons with many other methods: we verify our findings in the N- N-body experiments of Sec.
ETHICAL STATEMENT
The societal impact of SEGNNs is difficult to predict. However, SEGNNs are well suited for physical and chemical modeling and therefore potential shortcuts for computationally expensive simulations. And if used such, SEGNNs might potentially be directly or indirectly related to reducing the carbon footprint. On the downside, relying on simulations always requires rigorous cross-checks and monitoring, especially when simulations or simulated quantities are learned.
ACKNOWLEDGMENTS
Johannes Brandstetter thanks the Institute of Advanced Research in Artificial Intelligence (IARAD) and the Federal State Upper Austria for the support. This work is part the research programme VENI (grant number 17290), financed by the Dutch Research Council (NWO). The authors thank Markus Holzleitner for helpful comments on this work.
REFERENCES
Sanchez-Gonzaler Alvaro, Godwin Jonathan, Pfaff Tobias, Ying Rex, Leskovee Jure, and Battaglia Peter. Learning to simulate complex physics with graph networks. In Proceedings of the 37th International Conference on Machine Learning, volume 19, pp. 8459-8468, 2020.
Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Peter W. Battaglia, Jessica B. Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, Caglar Gulcehre, Francis Song, Andrew Ballard, Justin Gilmer, George Dahl, Ashish Vaswani, Kelsey Allen, Charles Nash, Victoria Langston, Chris Dyer, Nicolas Heess, Daan Wierstra, Pushmeet Kohli, Matt Botvinick, Oriol Vinyals, Yujia Li, and Razvan Pascanu. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv: 006.0/00.00266 2018.
10

---

Published as a conference paper at ICLR 2022
Simon Batzner, Tess E Smidt, Lixin Sun, Jonathan P Mailoa, Mordechai Kornbluth, Nicola Molinari, and Boris Kozinsky. Se (q)equvarr graph neural networks for data-efficient and accurate interatomic potentials. arXiv preprint arXtv:2101.03164, 2021.
Erik J Bekkers. B-spline enns on lie groups. In International Conference on Learning Representations, 2019.
Erik Bekkers, Maxime W Lafarge, Mitko Veta, Koen AJ Eppenhof, Josien PW Pluim, and Remco Duits. Roto-translation covariant convolutional networks for medical image analysis. In Interna- fional conference on medical image computing and computer-assisted intervention, pp. 440-448. Springer, 2018.
Lukas Biewald. Experiment tracking with weights and biases, 2020. URL https:://ww. wando com/ Software available from wandb.com.
Lowik Chanussot, Abhishek Das, Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Morgane Riviere, Kevin Tran, Javier Heras-Domingo, Caleb Ho, Weihua Hu, Aini Palizhati, Anuroop Sriram, Brandon Wood, Junwoong Yoon, Devi Parikh, C. Lawrence Zitnick, and Zachary Ulissi. Open catalyst 2020 (oc20) dataset and community challenges. ACS Catalysis, 600))605-66772 2020, doi: 0c04525
URL btttos://dot org/ 10. www.:1122aaccccccco
Lowik Chanussot, Abhishek Das, Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Morgane Riviere, Kevin Tran, Javier Heras- Deras-Dommnngo, Caleb Ho. Weihua Hu, Aini Palizhati, Anuroop Sriram, Brandon Wood, Junwoong Yoon, Devi Parikh, Lawrence Zitnick, and Zachary Ulissi. The open catalyst 2020 (oc20) dataset and community challenges, 2021.
Taco Cohen and Max Welling. Group equivariant convolutional networks. In International conference machine learning, pp. 2990-2999, PMLR, 2016.
Taco Cohen and Max Welling. Steerable cnns. In International Conference on Learning Represen- lations (ICLR), 2017.
Taco S. Cohen, Mario Geiger, Jonas Koehler, and Max Welling. Spherical enns. In International Conference on Learning Representations (ICLR), 2018
Michaël Defferrard, Xavier Bresson, and Pierre Vandergheynst. Convolutional neural networks graphs with fast localized spectral filtering. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett (eds.), Advances Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016.
Walter Dehnen and Justin Read. N-body simulations of gravitational dynamics. The European Physical Journal Plus, 126(5):1-28. 2011.
Congyue Deng, Or Litany, Yueqi Duan, Adrien Poulenard, Andrea Tagliasacchi, and Leonidas Guibas. Vector neurons: A general framework for -equvvrrua networks. arXiv preprint 10t.:2104.12299, 2021.
Matthias Fey and Jan E. Lenssen. Fast graph representation learning with PyTorch Geometric. In ICLR Workshop on Representation Learning on Graphs and Manifolds, 2019.
Mare Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolutional neural networks for equivariance to lie groups on arbitrary continuous data. International Conference Machine Learning, pp. 3165-3176. PMLR, 2020.
Mare Finzi, Max Welling, and Andrew Gordon Wilson. A practical method for constructing equivari- ant multilayer perceptrons for arbitrary matrix groups. arXiv preprint 104:::21040094:: 2021.
William T Freeman, Edward H Adelson, et al. The design and use of steerable filters. IEEE
Transactions Pattern analysis and machine intelligence, 13(9):891-906, 1991.
Fabian Fuchs, Daniel E Worrall, Volker Fischer, and Max Welling. Se transforrrr:::: 3d
iito-translation equivariant attention networks. arXiv preprint arXiv:2006. 10503, 2020.
11

---

Published as a conference paper at ICLR 2022
Mario Geiger, Tess Smidt, Alby M., Benjamin Kurt Miller, Wouter Boomsma, Bradley Dice, Kos- tiantyn Lapchevskyi, Maurice Weiler, Michal Tyszkiewicz, Simon Batzner, Jes Frellsen, Nuri Jung, Sophia Sanborn, Josh Rackers, and Michael Bailey. e3nn/e3nn: 2021-04-21, April 2021a. URL
ht.tps://doi.. .org .tt.//10.8820 5281 /ww/dzeoo....m 4708275
Mario Geiger, Tess Smidt, Alby M., Benjamin Kurt Miller, Wouter Boomsma, Bradley Dice, Kos- tiantyn Lapchevskyi, Maurice Weiler, Michat Tyszkiewicz, Simon Batzner, Jes Frellssen, Nuri Jung, Sophia Sanborn, Josh Rackers, and Michael Bailey. e3nn/e3nn: 2021-05-10, May 2021b. URL
httpss//doi owtp:///ooooog//..... wtt.://////..2/22eeeee/1
5281 www.dd/ 2 4745784
Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry, In International Conference on Machine Learning, pp.
1263-1272. PMLR, 2017.
Yacov. Hel-Or and Patrick Teo. Canonical decomposition of steerable functions. In Proceedings CVPR IEEE Computer Society Conference on Computer Vision and Pattern Recognition, pp. 809-816, 1996.
Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael JL Townshend, and Ron Dror. Learning from protein structure with geometric vector perceptrons. arXiv preprint arXiv:2009.01411 2020.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint 1412,6980, 2014.
Thomas Kipf, Ethan Fetaya, Kuan-Chich Wang, Max Welling, and Richard Zemel. Neural relational inference for interacting systems. In International Conference on Machine Learning, pp. 2688- 2697. PMLR, 2018.
Thomas N. Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In International Conference on Learning Representations (ICLR), 2017.
Johannes Klicpera, Janek Groß, and Stephan Günnemann. Directional message passing for molecular graphs. International Conference on Learning Representations, 2019.
Johannes Klicpera, Shankari Giri, Johannes T Margraf, and Stephan Günnemann. Fast and uncertainty-awar directional message passing for non-equilibrium molecules. arXiv preprint arXiv:2011.14115 2020.
Johannes Klicpera, Florian Becker, and Stephan Günnemann. Gemnet: Universal directional graph neural networks for molecules. arXiv preprint arXiv:2106.0880 08903, 2021.
Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: sampling configurations for multi-body systems with symmetric energies. arXiv preprint arXiv:1910.00753. 2019.
Risi Kondor and Shubhendu Trivedi. On the generalization of equivariance and convolution in neural networks the action of compact groups. In International Conference on Machine Learning, pp. 2747-2755. PMLR, 2018
Risi Kondor, Zhen Lin, and Shubhendu Trivedi. Clebsch Clebsch-gordan nets: fully fourier space spherical convolutional neural network. Advances in Neural Information Processing Systems, 11:10117- 10126, 2018.
Schütt Kristof, Kindermans Pieter-Jan, Sauceda Huziel, Chmiela Stefan, Tkatchenko Alexandre, and Klaus-Robert Müller. Schnet: continuous-fitter convolutional neural network for modeling quantum interactions. In Proceedings of the 31st International Conference on Neural Information Processing Systems, pp. 992-1002,2011 2017.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey Hinton. volutional neural networks. Advances in neural information processing systems. 25:1097-1105, 2012.
Imagenet classification with deep con-
Leon Lang and Maurice Weiler. wigner-eckan theorem for group equivariant convolution kerelss. In International Conference on Learning Representations, 2020.
12

---

Published as a conference paper at ICLR 2022
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient- learning applied to document recognition. Proceedings of the IEEE, 86(11):2278-25 2324, 1998.
Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. nature, 521(7553):436 2015.
Yi Liu, Limei Wang, Meng Liu, Xuan Zhang, Bora Oztekin, and Shuiwang Ji. Spherical message passing for 3d graph networks. arXiv preprint arXiv:2102.05013, 2021.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint
2017.
Andreas Mayr, Sebastian Lehner, Arno Mayrhofer, Christoph Kloss, Sepp Hochreiter, and Johannes Brandstetter. Boundary graph neural networks for 3d simulations. arXiv preprini arXiv:2106.1122: 2021.
Benjamin Kurt Miller, Mario Geiger, Tess Smidt, and Frank Noé. Relevance of rotationally equivariant convolutions for predicting molecular properties. arXiv preprint arXiv:2008. 08461, 2020.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. arXiv preprint arXiv:1912.01703 2019.
Prajit Ramachandran, Barret Zoph, and Quoe V. Le. Searching for activation functions. arXiv preprint arXiv:1710.05941, 2017.
Raghunathan Ramakrishnan, Pavlo Dral, Matthias Rupp, and Anatole von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific Data, 1, 08 2014.
David Romero, Erik Bekkers, Jakub Tomczak, and Mark Hoogendoorn. Attentive group equivariant convolutional networks. In Proceedings of the 37th International Conference Machine Learning, Proceedings of Machine Learning Research, pp. 8188-8199 2020.
Lars Ruddigkeit, Ruud Van Deursen, Lorenz C Blum, and Jean-Louis Reymond. Enumeration of 166 billion organic small molecules in the chemical universe database gdb-17. Journal of chemical information and modeling, 52(1 182444-88887 2012.
Jun J. Sakurai and Jim Napolitano. Modern Quantum Mechanics. Cambridge University Press, edition, 2017.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling, (n) equivariant graph neural networks. arXiv preprint arXiv:2102.09844, 2021.
Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The graph neural network model. JEEE Transactions on Neural Networks, 20(1):61-80, 2009.
Jürgen Schmidhuber, Deep learning in neural networks: An overview. Neural networks, 61:85-117, 2015.
Kristof T Schütt, Oliver T Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. arXiv preprint arXiv:2102.03150, 2021.
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang. Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds. arXiv preprint 1802.088000022199 2018.
Pfaff Tobias, Fortunato Meire, Gonzaezz Alvaro, and Battaglia Peter W. Learning mesh- based simulation with graph networks, 2020.
Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. Instance normalization: The missing ingredient for fast stylization. arXiv preprint arXiv: 1607.0802, 2016.
Maurice Weiler and Gabriele Cesa. General e (q)eeuriar stecrable cnns. Advances in Neural Information Processing Systems, 1434444444345. 2019.
13

---

Published as a conference paper at ICLR 2022
Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma, and Taco S Cohen. 3d steerable cnns: Learning rotationally equivariant features in volumetric data. Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
Maurice Weiler, Patrick Forré, Erik Verlinde, and Max Welling. Coordinate independent convolutional networks-isometry and gauge equivariant convolutions on riemannian manifolds. arXiv preprint arXiv:2106.06020, 2021.
Marysia Winkels and Taco Cohen. 3d g-enns for pulmonary nodule detection. In International Conference Medi- cal Imaging with Deep Learning (MIDL), 2018.
Daniel Worrall and Gabriel Brostow. Cubenet: Equivariance to 3d rotation and translation. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 567-584, 2018.
Daniel E Worrall, Stephan J Garbin, Daniyar Turmukhambetov, and Gabriel J Brostow. Harmonic networks: Deep translation and rotation equivariance. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 5028-5037, 2017.
Wenxuan Wu, Zhongang Qi, and Li Fuxin. Pointconv: Deep convolutional networks on 3d point clouds. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pp. 9621-9630, 2019.
Tian Xie and Jeffrey C Grossman. Crystal graph convolutional neural networks for accurate and interpretable prediction of material properties. Physical review letters, 120(1 (0014))143001) 2018.
C. Lawrence Zitnick, Lowik Chanussot, Abhishek Das, Siddharth Goyal, Javier Heras- MeraasDommingo, Caleb Ho, Weihua Hu, Thibaut Lavril, Aini Palizhati, Morgane Riviere, Muhammed Shuaibi, Anuroop Sriram, Kevin Tran, Brandon Wood, Junwoong Yoon, Devi Parikh, and Zachary Ulissi, introduction electrocatalyst design using machine learning for renewable energy storage, 2020.
14

---

Published as a conference paper at ICLR 2022
A MATHEMATICAL BACKGROUND
This appendix provides the mathematical background and intuition for steerable MLPs. We remark that the reader may appreciate several related works, such as, (Thomas al.] 2018 Anderson et al., 2019 Fuchs et al. 2020), as excellent alternative resource to get acquainted with the group/representation theory used in this paper. In this appendix we introduce the theory from our own perspective which is tuned towards the idea of steerable MLPs and our viewpoint on group convolutions. It provides complementary intuition to the aforementioned resources. The main concepts explained in this appendix are:
1. Group definition and examples of groups (Section A.I). The entire framework builds upon notions from group theory and as such formal definition is in order. In this paper, model transformations such as translation, rotation and reflection as groups.
2. Invariance, equivariance and representations (Section[A.2 A function is said to be invariant to a transformation if its output is unaffected by a transformation of the input. A function is said to be equivariant if its output transforms predictably under a transformation of the input. In order to make the definition precise, we need a definition of representations; representation formalises the notion of transformations applied to vectors in the context of group theory.
3. Steerable vectors, Wigner-D matrices and irreducible representations (Section A.3) Whereas regular MLPs work with feature vectors whose elements are scalars, our steerable MLPs work with feature vectors consisting of steerable feature vectors. Steerable feature vectors are vectors that transform via so-called Wigner-D matrices, which are representations of the orthogonal group O(3). Wigner-D matrices are the smallest possible group representa- tions and can be used define any representation (or conversely, any representation can be reduced to tensor product of Wigner-D matrices via change of basis). As such, the Wigner-D matrices are irreducible representations.
4. Spherical harmonics (Section A.4). Spherical harmonies are a class of functions on the sphere 2 and can be thought of as Fourier basis on the sphere. We show that spherical harmonics are steered by the Wigner-D matrices and interpret steerable vectors as functions on S², which justifies the glyph visualisations used in this paper. Moreover, spherical harmonics allow the embedding of three-dimensional displacement vectors into arbitrarily large steerable vectors.
5. Beebsch-Gordan tensor product and steerable MLPs (Section A.5) In a regular MLP one maps between input and output vector spaces linearly via matrix vector multiplication and applies non-linearities afterwards. In steerable MLPs one maps between steerable input and steerable output vector spaces via the Beebcch-Gordan tensor product. Akin to the learnable weight matrix in regular MI.ss, the learnable flebsch-Gordan tensor product is the workhorse of our steerable MI.Ps
After these concepts are introduced will in Section B revisit the convolution operator in the light of the steerable, group theoretical viewpoint that we take in this paper. In particular, we show that steerable group convolutions are equivalent linear group convolutions with convolution kernels expressed in a spherical harmonic basis. With this in mind, we argue that our approach via message passing can be thought of building neural networks via non-linear group convolutions.
A.1 GROUP DEFINITION AND THE GROUPS E(3) AND 0(3)
Group definition. A group is an algebraic structure that consists of a set G and binary operator " the group product, that satisfies the following axioms: Closure: for all h, g E G we have h g E G; Identity: there exists an identity element e E G; Inverse: for each g E G there exists inverse element g E G such that g .g..---- -g=g g = e; and Associativity: for each g, h, i E G we have (g h)) i i= g. ())
The Euclidean group E(3). In this work, we are interested in the group of three-dimensional translations, rotations, and reflections which is denoted with E(3), the 3D Euclidean group. Such
Fach of these works presents unique view points that greatly influenced the writing of this appendix.
15

---

Published as a conference paper at ICLR 2022
transformations are parametrised by pairs of translation vectors R and orthogonal transformation matrices R E (3).. The E(3) group product and inverse are defined by
g g = RRx' + x, xRR), g - R x x, R R-
(x',)) with y R), Rx) E B(3). One can readily see that with these definitions all four group axioms are satisfied, and that it therefore defines group. The group product can be seen as description for how two E(3) transformations parametrised by g and g' applied one after another are described by single transformation parametrised by g g'. The transformations themselves act on the vector space 3D positions via the group action, which we also denote with ; via
g y Ry + x;
where g (x, R) B(3) and y E R
The orthogonal group 0(3) and special orthogonal group SO(3). The group E(3) = R (3)
a semi-direct product (denoted with x) of the group of translations RR with the group of orthogonal transformations 0(3). This means that we can conveniently decompose tran-trrnsoormmaion in 0(3)-transformation (rotation and/or reflection) followed by a translation. In this work we will mainly focus on dealing with O(3) transformations as translations are trivially dealt with. When representing the group elements of O(3) with matrices R. as we have done before, the group product and inverse are simply given by the matrix- prrricctprruttttec and inverse Le., with g R. g R' = the group product and inverse are defined by
9. g = R.
R
The group acts 3 by matrix-vector multiplication, i.e., 9 y = Ry. The group elements of 0(3) are square matrices with determinant -1 or 1. Their action on IR3 defines a reflection and/or rotation.
The special orthogonal group SO(3) has the same group product and inverse, but excludes reflections. The group thus consists of matrices with determinant 1.
The sphere 2 is homogeneous space of SO(3). The sphere is not a group as we cannot define a group product on S that satisfies the group axioms. It can be convenient treat it as homogeneous space the groups 0(3) or SO(3). A space X' is called a homogeneous space of a group G for any two points x, y E X' there exists group element g E G such that y 9
The sphere S a homogeneous space of the rotation group SO(3) since any point on the sphere can be reached via the rotation of some reference vector. Consider for example XYX parametrisation of SO(3) rotations in which three rotations are applied one after another via
Ro.2, Ro.n, R.s.n, Ryy.,
(A.1)
with nr and ny denoting unit vectors along the and y axis, and R.,, n. denotes rotation of degrees around axis n.. We can model points on the sphere in a similar way via Euler angles via
Do.8 == R.,
(A.2)
So, with two rotation angles, any point on the sphere can be reached. In the above we set y = 0 in the parametrisation of the rotation matrix (an element from SO(3)) that rotates the reference vector ny, but should be clear that with any y the same point na,y is reached. This means that there are many group elements in SO(3) that all map n. the same point on the sphere.
A.2 INVARIANCE, EQUIVARIANCE AND REPRESENTATIONS
Group representations. We previously defined the group product, which tells us how elements of a group interact. We also showed that the groups E(3) and O(3) can transform the three dimensional vector space R3 via the group action. We usually think of the groups E(3) and O(3) groups that describe transformations on R but these groups are not restricted to transformations on R and can generally act arbitrary vector spaces via representations. A representation is invertible linear
16

---

Published as a conference paper at ICLR 2022
transformation p(g) V V parametrised by group elements g G that acts vector space V, and which follows the group structure (it is group homomorphism) via
p(g)p(h)) p(g h)v.
with v V.
A representation can act on infinite-dimensionne vector spaces such as functions. E.g., the so-called left-regular representations Lg of E(3) on functions f R + R on IR is given by
ie.. it transforms the function by letting g-1 act on the domain from the left. Here we used the notation Lg[f] to indicate that Cg transforms the function f first, which creates a new function (x)), which is then sampled at x.
When representations transform finite dimensional vectors V = RR,, they are d dimensional matrices. In this work, we denote such matrix representations with boldface D. A familiar example of a matrix representation of O(3) on TR3 are the matrices g = R E (()) themselves, i.e., D(g)x x Rx.
Finally, any two representations, say D(g) and D' D'(g), are equivalent if they relate via similarity transform via
D'(g) ( QD
i.e., such representations describe one and the same thing but in a different basis, and the change of basis is carried out by Q. Now that representations have been introduced we can formally define equivariance.
Invariance and equivariance. Equivariance is property of an operator ó o: y that maps between input and output vector spaces X' and y. Given a group G and its representations pK and p which transform vectors in X' and Y' respectively, an operator : X + y is said to be equivariant if satisfies the following constraint
/ (o)l((=)) ( (g)[z)) for all g E G, E
(A.3)
Thus, with equivariant map, the output transforms predictably with transformations the input. One might say that no information gets lost when the input is transformed, merely re-structured. One way to interpret Eq. (A.3) is therefore that the operators px (g) : X X' and (g) y describe the same transformation, but in different spaces.
Invariance is special case of equivariance in which /² = I for all g E G. L.e,, an operator Φ X' + y is said to be invariant if it satisfies the following constraint
(<x) = <p (g)[z)) for all E G,,EE
(A.4)
Thus, with an invariant operator, the output of & is unaffected by transformations applied to the input.
A.3 STEERABLE VECTORS, WIGNER- D MATRICES AND IRREDUCIBLE REPRESENTATIONS
One strategy to build equivariant MLPs is define input and output spaces of the MLPs and define how these spaces transform under the action of a group. This then sets an equivariance constraint on the operator that maps between these spaces. By only working with such equivariant operators we can guarantee that the entire learning framework is equivariant.
In our work, the proposed graph neural networks are translation equivariant by construction as any form of spatial information only enters the pipeline in the form of relative positions between nodes (x, - x.). Then, any remaining operations are designed to O(3) equivariant such that, together with the given translation equivariance, the complete framework is fully E(3) equivariant. Since translations are trivially dealt with, we focus on SO(3) and 0(3) and show how build equivariant MLPs through the use of the Clebsch-Gordan tensor product.
Wigner-D matrices are irreducible representations. For SO(3) there exists a collection of repre- sentations, indexed with their order 2 0, which act on vector spaces of dimension 2/ + 1. These representations are called Wigner-D matrices and we denote them with ( (g). The use of Wigner-I>
17

---

Published as a conference paper at ICLR 2022
matrices is motivated by the fact that any matrix representation Dgg) of SO(3) that acts some vector space V can "reduced" to an equivalent block diagonal matrix representation with Wigner-D matrices along the diagonal:
(
D(g) (D())) @ D ( Q-1
((l))g)
Q
(A.5)
with Q the change of basis that makes them equivalent. The individual Wigner-D matrices themselves cannot be reduced and are hence irreducible representations of SO(3). Thus, since the block diagonal representations are equivalent to we may as well work with them instead. This convenient since each block, i.e., each Wigner-D matrix D(t.,, only factorise V = Vi, ① Vi2 @ ...., which motivates the use of steerable vector spaces and their direct sums as presented in Sec.
sub-space Vi, of V. As such we can
The Wigner-D matrices are the irreducible representations of SO(3), but we can easily adapt these representations to be suitable for O(3) by including the group of reflections as direct product. We will still refer to these representations as Wigner-D matrices in the entirety of this work, opting to avoid the distinction in favour of clarity of exposition. We further remark that explicit forms of the Wigner-D matrices can e.g. be found books such as in Sakurai & Napolitano (017) and their numerical implementations code libraries such as the e3nn library (Geiger al. 2021a).
Steerable vector spaces. The (21 + 1) dimensional vector space on which a Wigner-D matrix of order acts will be called a rype steerable vector space and is denoted with V/. E.g.. type-3 vector V3 transformed by g E 0(3) via h + (g(g) h. We remark that this definition is equivalent to the definition of steerable functions commonly used in computer vision (Freeman et al. 1991 Hel-Or Teo (996) via the viewpoint that steerable vectors can regarded as the basis coefficients of function expanded in spherical harmonic basis. We clicit this viewpoint in Sec. A.4 and B
At this point we are already familiar with type-O and type-l steerable vector spaces. Namely, type-0 vectors & Vo R are scalars, which are invariant to transformations g E 0(3), i.e. Do(g)h h. Type- features are vectors h E IR3 which transform directly via the matrix representation of the group, i.e., D2 (g)h- Rh.
A.4 SPHERICAL HARMONICS
Spherical harmonics. Related the Wigner-D matrices and their steerable vector spaces the spherical harmonics Spherical harmonics are a class of functions on the sphere s², akin to the familiar circular harmonics that are best known the 1D Fourier basis. As with a Fourier basis, spherical harmonics form an orthonormal basis for functions on $2, In this work we use the real-valued spherical harmonies and denote them with y $2 R.
Spherical Harmonics are Wigner-D functions. One can also think of spherical harmonics as functions v.() SSO3) + E SO(3) that are invariant sub-group of rotations via
yr (P...) = (!) (R.,a,, Dz) =: (1()) (Ro,9.7)
in which used the parametrisation for S and O(3) given in A.2 and A.1 respectively. Then, by definition, (D) invariant with respect to rotation angle 7, i.... Vye(o,27)
Y (Ra,,,v) =
y-(1) (R.,,,,0. This viewpoint of regarding the spherical harmonies as y-invariant functions on 0(3) helps us to draw the connection to the Wigner-D functions (()) that make up the 2/ + 27 + elements the Wigner-D matrices. Namely, the n = column Wigner-D functions are also
y i-iivariant and, in fact, correspond (up to normalisation factor) to the spherical harmonics via
(
(na.8) = /21
(
m0 (Ra,8.7)
(A.6)
'Solutions to Laplace's ' equation called harmonics. Solutions of Laplace's equation the sphere therefore called spherical harmonies.
18

---

Published as a conference paper at ICLR 2022
The mapping from vectors into spherical harmonic coefficients is equivariant. It then directly follows that vectors of spherical harmonics are steerable by the Wigner-D matrices of the same degree.
Let (n) (Y (m).... y (n))² " be the embedding of direction vector E S in spherical harmonics. Then this vector embedding equivariant as it satisfies
VREOOO) VnES=
4 (R'nn)
(R'n) = D (R')a
(A.7)
Using the S and 0(3) parametrization of A.2 and A.1 this is derived as follows
D (R'R R.))
D (Ra
nh) (R'na,,))) =
(R')
11 (R'na
(R'Ra.
(0) (Ro
10
Y (na..)
= -D'(RR)
= D (R')a (na.a)
(Da.o)
Steerable vectors represent steerable functions on S2 Just like the 1D Fourier basis forms complete orthonormal basis for ID functions, the spherical harmonics y!( form an orthonormal basis for (5"),,, the space of square integrable functions the sphere. Any function on the sphere can thus be represented by a steerable vector á E Vo @ V/2 @ when it is expressed a spherical harmonic basis via
((n)
y m m (n).
(A.8)
120 m=-1
Since spherical harmonies form orthonormal basis, the coefficient vector directly be obtained by taking --inner products of the function I with the spherical harmonics, i.e.,
(S,Y )L)(s)) / f(n) (n)dn.
(A.9)
Equation (A.9) is sometimes referred to the Fourier transform on S², and Eq. (A.8) the inverse spherical Fourier transform. Thus, one can identify steerable vectors a with functions on the sphere S via the spherical Fourier transform.
In this paper, we visualise such functions on S via glyph-visualisation which are obtained as surface plots by taking S2 and scaling it by [n)] sign (f(n) f(n):
[nm > 0 | <0
n [n(())]
| n S²
}
where each point on this surface color-coded with the function value f(n). The visualisations are thus color-coded spheres that are stretched in each direction n via (n)]]
Steerable vectors also represent steerable functions on 0(3). In order to draw a connection between group equivariant message passing and group convolutions, as we did in Sec. of the main paper, it important to understand that steerable vectors also represent functions on the group SO(3) via an
transform. The collection of (21 + 1)² Wigner-D functions ((() form an orthonormal basis for (S0(3))). This orthonormal basis allows for Fourier transform that maps between the function space IL. (SO(3)) and steerable vector space V; the forward and inverse Fourier transform on SO(3) are respectively given by
1 f(g)
(
mn (g)dg.
(A.10)
J0(3)
f(g)
Σ
Σ
mn mnn (g)
(
(A.1I)
120 m= n==!
19

---

Published as a conference paper at ICLR 2022
with dg the Haar measure of the group. Noteworthy, the forward Fourier transform generates matrix of Fourier coefficients, rather than a vector in spherical case. The coefficent matrix is steerable by left-multiplication with the Wigner-D matrices of the same type D
A.5 CLEBSCH-GORDAE PRODUCT AND STEERABLE MLPs
In regular MLP one maps between input and output vector spaces linearly via matrix- multiplication and applies non-linearitin afterwards. In steerable MLPs, one maps between steerable input and output vector spaces via the Gordch tensor product and applies non-linearities afterwards. Akin to the learnable weight matrix in regular MLPs, the learnable Glebsch-Gordan tensor product is the main workhorse our steerable MLPs.
Clebsch-Gordas tensor product. The Clebsch-Gorder (CG) tensor product allows map between stecrable input and output spaces. While there is much to be said about tensors and tensor products in general, multiplication between all components of two input vectors. E.g., with two vectors h1 E IRds and h2 E Rdd, the tensor product is given by
here intend to foeus intuition. In general, tensor product involves the
h1h2 h2h1
h h2
hg 2 h h₂h
hgh2
which we can flatten into di dz-dimeensiona vector via an operation which we denote with vec(h, @ h2). In our steerable setting we would like to work exclusively with steerable vectors and as such we would like for any two steerable vectors h2 E Vi, and h2 E Vi,, that the tensor product's output is again steerable with 0(3) representation D(g) such that the following equivariance constraint is satisfied:
(hg h2) (D (b)) ( (g)hn) @ (D (1)) (2)
(A.12)
Via the identity voc(AXB (BT @ A) vec(X),
show that the output is indeed steerable:
vec (g)))
vec (D (g)h b ((a)) (()TT(g)
D (g) @ (g) voe h hinng
The resulting vector vee(hi & hig) is thus steered by representation D(g) D (g) & (g) Since any matrix representation of 0(3) can be reduced to a direct sum of Wigner-D matrices (see (A.5)), the resulting vector be organised via change of basis into parts that individually transform
via Wigner-D matrices of different type. Lee. h vec(h; @ h2) E V Vo @ V1 ...., with Vi the steerable sub-vector spaces of type
With the CG tensor product directly obtain the vector components for the steerable sectors of type as follows. Let (1)) E Vi = 2 denote steerable vector type and (( components with -l, -1 1,....l.
I. Then the m-th component of the type I sub-vector of the output of the tensor product between two stecrable vectors type l and la is given by
(1())) Beg (0)) w
(.m) (frmi)llamaa
h (1) (1) Pi Pa
(A.13)
in which is learnable parameter that scales the product and ((l.m)
(11mm))amm are the Clebsch- Gordan coefficients. The CG tensor product is a sparse tensor product, as generally many of the (l.m)
(1rm1 (la.m) components are zero. Most notably, ((l.m)
(11.m))(ssmm = 0 whenever I < |/1 12l or 1> l2. E.g., type-O and type-1 feature vector cannot create type-2 feature vector. Well known examples of the CG product are the scalar product (11 = 0, ly = 1, = 1), which takes as input scalar and type-1 vector to generate type-1 vector, the dot product (1, 1,12 1, 0) and cross product (1) 1 = 1, -01=11 1).
Examples of instances of the CG product
20

---

Published as a conference paper at ICLR 2022
The product of two scalars is CG product which takes input two type-0 "vectors" to generate another scalar 141 = 0, l2 = 0, = 0).
The scalar product, which takes as input scalar and a type-1 vector to generate type-1 vector (/ 0, 12 1,7 1).
The dot product, which takes two type-1 vectors as input to produce a scalar (41 = 1,12 1,7 0).
The cross product, which takes two type-1 vectors input to produce another type-1 I vector (41 = 1, = 1 1).
In a standard linear layer, an input vector is transformed via multiplication with the weight matrix, where weights and elements of the input vector are simply multiplied. For forl: > 0 elements of the Clebsch-Gorts tensor product of Eq. (A.13) more than one mathematical operation needed for the connection with one weight, as can easily be verified if hm and hmgg are thought of as e.g. type-1 I vectors. This makes calculation the Clebsch-Gordan tensor product slow for higher order irreps, as can be observed in all experiments.
Steerable MLPs. The CG tensor product thus allows us to map two steerable input vectors to a new output vector and can furthermore be extended to define tensor product between steerable vectors of mixed types. The CG tensor product can be parametrised by weights where the product scaled with some weight for each triplet of types (l,, 12, 1) for which the CG coefficients non-zerd We indicate such CG tensor products with @ While in principle the CG tensor product takes two steerable vectors as input, in this work we mainly use with one of its input vectors "fixed". The CG tensor product can then be regarded as linear layer that maps between two steerable vector spaces and we denote this Wah == h @ W å, with à the steerable vector that is considered to be fixed. With this viewpoint we can design MLPs the same way as we are used to with regular linear layers, and establish clear analogies with convolution layers (Sec. of the main paper).
While in general the CG tensor product between two steerable vectors of type I and l2 can contain steerable vectors up to degree 1 + 12. one typically "band-limits" the output vector by only considering the steerable vectors up to some degree Lmaz- In our experiments we band-limit the hidden feature representations to maaz ly.
Finally, the amount of interaction between the steerable sub-vectors in the hidden representations of degree Vi=í, = Vo + V1 + Vi, is determined by the maximum order la of the steerable vector à on which the stecrable linear layer is conditioned. Namely, the CG tensor product only produces type steerable sub-vectors for 1 lal 2 I S f + lal- For example, it is not sensible to embed positions steerable vectors up to order la 5 when the hidden representations degree I 2; the lowest steerable vector type that can be produced with the CG product of a 5 sub-vector of with any sub-vector of the hidden representation vector will be / = 3 and since the hidden representations are aand-limited to maximum type of ly 2 higher order vectors will be ignored.
of
B STEERABLE GROUP CONVOLUTIONS
Steerable functions. Through the equivalence of stecrable vectors and spherical functions via the spherical Fourier transform. it is clear that our definition of steerable vectors coincides with the more common definition of steerable functions, as commonly used computer vision (Freeman al., 1991). A function if is called steerable (Hel-Or & Teo 1996) under transformation group G if any transformation E G on f can be written as linear combination of a fixed, finite set of n basis functions {
fg = Σ od(9)s x= a))
(B.1)
in which C, is the left regular representation of the group G that performs the transformation on (see section A.2 In case of functions in spherical harmonic basis, which we denote with
In terms of the e3nn library (Geiger al. 2021al one then says path exists between these 3 types.
21

---

Published as a conference paper at ICLR 2022
( La Ei: -m
am y (n), it directly follows from Eq. A.7 that they steerable via (gfa(n) (D(g)a
(B.2)
In terms of the steerable function definition in Eq. B.I.. this means that a(g) D(g)a a and Φ = -Yu)) y xo ,"..)", i.. the set of basis functions flattened into a vector.
translation-rotation template machting motivation for (steerable) group convolutions. The notion of steerability becomes particularly clear when viewed in the computer vision context (Freeman al. 1991). where one may be interested in detecting visual features under arbitrary rotations. Let us consider the case of 3D cross-correlations of kernel k R + f: IR R:
RR with an input feature map
('(x - (k f)() - / 1(x' (x)-xxx(x)xxx
(B.3)
We think of such kernel k as describing a particular visual pattern. In many applications, we want detect such patterns under arbitrary rotations. For example, in 3D medical image data there is no preferred orientation and features (such as blood vessels, lesions, ...) can appear under any angle, and the same holds for particular atomic patterns in molecules. So ideally, one wants to apply the convolution kernel under all such transformations and obtain feature maps via
fxx. R)
k(R (x'
(B.4)
By repeating the convolutions with rotated kernels we are able to detect the presence of a certain feature at all possible angles. In a group convolution context (Cohen Welling 2016), the above is what one usually calls lifting group convolution (Bekkers 2019) (feature maps are lifted from R to the group SB(3)).
Group convolutions. In regular group convolutional neural networks one continues to work with such higher dimensional feature maps in which the kernels are also functions on the group. The lifting and subsequent group convolutions then all have the same form and are defined via the group action on R and group product respectively via
f'(g) - kgg xx)((x)xx''
(B.5)
(((g) f [SE(S) k(g gg))(('''''
(B.6)
where g E SE(3), dg the Haar measure on the group and where in B.5 and (.6) respectively denote the group action on 3R3 and group product of SE(3) (cf. Section A.I. Note that Eq. BB.) exactly the same as Eq. (B.4) but in different notation.
The lifting group convolution thus generates a function on the joint space of positions Rd3 and rotations SO(3). In numerical implementations, this space needs to be discretised, i.e., for a particular finite grid of rotations we want to store the results of the convolutions for each rotation R. This approach then requires that the convolution kernel is continuous and can be sampled under all transformations. Hence, such kernels can be expanded in a continuous basis such as spherical harmonics (Weiler et alll,,2000 2018), B B-splines (Bekkers) 2019) they can be parametrised via MLPs (Finzi et al. 2020). Alternatielyy, the kernels are only transformed via a sub-group of transformations in E(3) that leaves the grid on which the kernel is defined intact, in (Worrall & Brostow 2018: Winkels & Cohen 2018). An advantage of regular group convolution methods is that normal point-wise activation functions can directly be applied to the feature maps; down-side is that these methods are only equivariant to the sub-group on which they are discretised. When expressing the convolution kernel in terms of spherical harmonics. however, there no need for such a discretisation at all and one can obtain the response of the convolution at any rotation R after convolving with the basis functions. This works as follows.
Steerable template matching. Suppose 3D convolution kernel that is expanded in a spherical harmonic basis up to degree follows
k(x)= ke(])>|) (x)
ΣV
Σ
(B.7)
|
m=-1
22

---

Published as a conference paper at ICLR 2022
(0) (1) 000 ...)² the vector of basis coefficients that can depend on ||x||_ The coef- with é (co -0
ficients can e.g. be parametrised with an MLP that takes as input ||*|| and returns the coefficient vector. We note that such coefficients are then O(3) invariant, i.c., VREO3) : c(I] Furthermore, we labelled the vector with a " to indicate it is a steerable vector as it represents the coefficients relative spherical harmonic basis. then follows (from Eq. (B.2h) that the kernel steerable via
c(||x|)
k: <R- = (R)(((( (x),
i.e., via transformation of the coefficient vector è by 0(3) representation D(R).
This steerability property. together with linearity of the convolutions and basis expansion, implies that with such steerable convolution kernels we can obtain their convolutional response at any rotation directly from convolutions with the basis functions. Instead of first expanding the kernel in the basis by taking a weighted sum of basis functions with their corresponding coefficients, and only then doing the convolutions, we can change this order and first do the convolution with the basis functions and sum afterwards. In doing so we create vector of responses f(x) (Jo (x), (x), (x).... of which the elements are given by
(x) = ((C,) (() { (x)
(B.8)
(1xx|) fxx(x')
Then the result with the convolution kernel (Eq. (B.7) is obtained simply by sum over the vector components which we denote with SumReduce as follows
1,
L Σ m=-
(x) = SumReducee (f( =
Σ
((()
(x)
1,m
If one were to be interested in the rotated filter response (x( via the matrix representation D(R) and only then do the reduction. Lee., once the convolutions of (B.8) are done the lifting group convolution result is directly obtained via
first rotate the steerable vector
fx( R) SumReduce (DR)
(B.9)
1,
convolution with kernel expressed in a spherical harmonics basis thus generates a signal on O(3) at each point x.
When only considering the m component for each order / in the kernel parametrisation, the kernels have axial symmetry in them due to which the convolution results in point-wise spherical signal. This in fact the situation that we consider we do not index the weights with m. This choice leads to computational efficiency. It is however not necessary to constrain the kernels as such and continue with stcerable group convolutions on the full space R33 SO(3), this would require working with CG products with weights indexed by both m and m' (Weiler et al. 2018). Since this work we choose to with steerable vector spaces V mo Vo @ m, Vi ... with arbitrary multiplicities my for each type, we implement what is called generalised steerable convolutions, of which the SO(3) convolutions are special case with my (21 1) Kondor et al. 2018).
our weighted CG products (Eq. 55) where
Steerable group convolutions. When working with stecrable convolution kernels one does not have to work with a grid on 0(3). There are reasons avoid working with a grid on 0(3) as one cannot numerically obtain exact equivariance if the chosen grid is not a sub-group of O(3). When limiting to a discrete subgroups one can guarantee exact equivariance to the sub-group, but ideally one obtains equivariance to the entire group 0(3). Steerable methods provide a way to build neural networks entirely independent any sampling on O(3), since, as we have seen, steerable convolutions directly result in functions the entire group 0(3) via stecrable vectors. That is, each location have steerable vector f(x) which represents a function on 0(3) via the inverse Fourier transform given in Eq. (T. In fact, the sum reduction in Eq. B.9 corresponds to a discrete inverse Fourier transform.
Finally, let us revisit steerable group convolutions one more time but now in the steerable vector notation used throughout this paper. The steerable convolution, as described in (B.8) is operator
23

---

Published as a conference paper at ICLR 2022
that maps between steerable vector fields of different types. In the above example the input feature map was one that only provided a single scalar value per location x, i.c. a type-O steerable vector, steerable vector field containing steerable vectors up to type L. The transition
and the output from type-O vector field to a type-1. vector field happened via tensor products with steerable vectors of spherical harmonics, and this tensor product convolution of a single type-11 steerable input field with a convolution kernel that is expanded in
parametrised by the coefficients è. Suppose
Txx|| The kernel can spherical harmonic basis of only type l/ via k(x) -mm-ro (0 ([[[
then be represented as steerable vector field in itself as kxx) (x[[x (x)),, with á the type-12 spherical harmonic embedding given in Eq. 5 of the main paper. Such a steerable convolution maps
from input feature maps f R3 - Vi, using convolution kernel k R + V1, to an output 1R3 V via
/ f(x) ([[x a(x' - x)dx'.
f(()
Here we assumed steerable feature vector fields of a single type, but in general such convolutions can map between vector fields of mixed type analogous to the standard convolutional case in CNNs where mappings occur between multi-channel type-0 vector fields.
Steerable and regular group convolutions. In conclusion, both steerable and regular group con- volutions produce feature maps on the entire group E(3) and they are equivalent when the regular convolution kernel is expanded in a steerable basis. In regular group convolutions, the response is stored on particular grid which is e.g. the Cartesian product of regular 3D grid with a par- ticular discretisation of O(3). In regular group convolutions, we directly index the responses with (x, R.) E B(3) + f(x, R). In steerable convolutions the filter responses - f(x) are stored at each spatial grid point x in steerable vectors f(x) from which functions on the full group O(3) can be obtained via inverse Fourier transform As such one recovers the regular representation via (x, R) F- [F( (X)... Regular group convolutions can however not be perfectly equivariant to all transformations in O(3) due to discretisaton artifacts, or they are only equivariant to a discrete sub-group of O(3). Steerable group convolutions on the other hand exactly equivariant to all transformations in 0(3) via steerable representations of O(3).
"This Fourier transform in fact enables working with classic point-wise activation functions that could be applied point-wiss to each location in spherical 0(3) signal as in (Cohen al. 2018). is important to note that one cannot readily apply activation functions such as ReLU directly to the steerable coefficients, but one could sample the O(3) signal on grid via the inverse Fourier transform. apply the activation function, and transform back into the steerable hasis. This is the approach taken in (Cohen al, 2018). In this paper work entirely in the steerable domain (0(3) Fourier space) and work with gated non-linearities (Weiler al. (018).
24

---

Published as a conference paper at ICLR 2022
C EXPERIMENTS
C.I PSEUDOCODE OF SEGNN AND ABLATED ARCHITECTURES
In this subsection, we provide pseudocode for our implementation SEGNN and steerable ablations. Pseudocode for steerable equivariant point conv layers such as (Thomas et al. 2018) with messages as Eq. (14) and regular gated non-lined con-lineaiities as activation/update function, labelled SEnmeur is outlined in Alg. Pseudocode for the same network but with messages obtained in non-linear manner via 2-layer steerable MLPs as Eq. (16) is outlined in Alg. 2 This layer is labelled as SEnon-innea-
Finally, pseudocode for SEGNN implementation is provided in Alg. 3 SEGNN layers allow for the
inclusion of geometric and physical quantities via node attributes ä;- These new node update layers can be regarded as new steerable activation function and allow functionality which is not possible in layers such as SEiineur or S linear. A more detailed explanation of the Clebsch-Gordan tensor and
product (CGTensorProduct), the spherical harmonic embedding (SphericalllammoniiEmmble the application of gated non-linearities can be found in the e enn (Geiger et al. 2021a) library.
Injecting geometric and physical quantities. In the following algorithms the relative position vector Xij between node f; and node fj is given as an example for a geometric quantity used to steer the message layers. To get more concrete, we look at Alg. 3 (SEGNN) and discuss edge attributes and node attributes. Steerable attributes first of all have to be embedded via spherical harmonics. Steerable edge attribute å.j in most cases consist of relative positions. Instead of or additional to relative positions also relative forces relative velocities could be used. And finally, also node specific quantities such as force can be added together to serve as edge attributes. Steerable node attributes are in most cases real physical quantities. In Alg. they are sketched with v! and vř. and might comprise velocity, acceleration, spin, or force. However, also adding the spherical harmonics embedding of relative positions at node f, results in steerable node attribute. Lastly, it is note that if multiple steerable attributes exist they can be either added or concatenated. However, the latter results in significantly more weights in the CG tensor product and does not result in better performance in any of the conducted experiments.
Gated non-rincasssies. The gate activation is direct sum of two sets of irreps. The first set the set of scalar irreps passed through activation functions. The second set is the set of higher order irreps, which is multiplied by an additional set of scalar irreps that introduced solely for the activation layer and passed through activation functions. For example, in Alg. [ the scalar irreps gij the additional scalar irreps, introduced to "gate" the non-scalar irreps of m.j.
Algorithm 1 Code sketch of an SEjincar message passing layer. The SEjinear layer updates steerable node features f/ < f. The additional scalar irreps gi are used to "gate" the non-zero order irreps as introduced in Weiler et al. (2018). A fully documented code for this layer will be applicable in our repo.
Require: f, Xij
> Steerable nodes f,, relative position vector Xij between node and node f,
function 733.TESoo.PRODDccTnnp
input2)
output +- CCTensorProducc(nnput,, inntiitttttiiiiin)))
Apply CG tensor product following Eq. ( D Add bias to zero order irreps
output + output bias
return output
end function
a; + SphericllfammonidEEmmcdddngngng
Spherical harmonic embedding of x;j (Bq. Concatenate input for messages between f., f, & Calculate messages mij and scalar irreps gij
h. + f f, + ||x,,||
mij @ gij f- O3_TENSOR_PRODUCTIO ā:y)
m; + g:
Em.
m.y >E',
Aggregate messages m, and scalar irreps gij
f. + Gate(m,, Swish(g,))
P Transform nodes via gated non-linearisies
25

---

Published as a conference paper at ICLR 2022
Algorithm/ Code sketch of an SEgon-lincar message passing layer. The SEacn-lincar layer updates steerable node features f/ +- f. The additional scalar irreps gi are used to "gate" the non-zen order irreps introduced in Weiler et al. (0188) A fully documented code for this layer will be applicable in our repo
Require:
f, Xij
Steerable nodes f,, relative position vector Xij between node and node fj
function O3_TENSOR PETENSSORPRODDUCC(nnnnt... inpur2)
output +- CGTeesooProoduc((npttl,, inputt22) input2)
Apply CG tensor product following Eq. (6) D> Add bias to zero order irreps
output + output bias
return output
end function
function 03_ 23_TENSOR TT__TSSSRR,PRDUCTT PRODUCT_SWISH
input2)
output @ gi + TTTENSOR_PRDUUCTII input2)
Output plus scalar irreps g: Transform output via gated non-linearities
oute +- Gate(output, Swish(g.))
return output
end function
ā. + Sphericlllamonicineeddiggsss
Spherical harmonic embedding of Xij (Eq. Concatenate input for messages between f,, f, First non-linear message layer b Second linear message layer
h. f, @ f, + ||x,,|2
mij f- 03 TENSOR PRODUCT SWISH GATE(h. āzg))
mij @ Bij < OS_TENSOR,PRODUCTTO
å:z)
m gi + Σ m.j >
Aggregate messages m.j and scalar irreps gij
f; + Gate(m; Swish(g:)
D Transform nodes via gated non-linearities
Algorithm 3 Code sketch an SEGNN message passing layer. The SEGNN layer updates steerable node features f/ < f, The additional scalar irreps gi are used to "gate" the non-zero order irreps as introduced in Weiler et al. (2018). A fully documented code for this layer will be applicable in our repo
Require: f, Xij. v vē Steerable nodes f,, relative position vector X,y between node and node f,. geometric or physical quantities v! vf such as velocity, acceleration, spin, or force.
function 03_ TE.TENSOR Ww_.TENSOR_RRRRD PRODUCT(inuul,, input2)
output +- CCTensoPPooduccinppul, in
Apply CG tensor product following Eq. (6) > Add bias to zero order irreps
output < output bias
return output
end function
function 03 TENSOR PRODUCT SWISH GATE(inputl, input2)
output @ g: +- S.STEESOO_.RRDUCTinput input2)
Output plus scalar irreps g, Transform output via gated non-linearities
outputgated f- Gate(output, Swish(g.))
return output
end function
ā + SpterceaHHammoiidEnmcceiegagg
Spherical harmonic embedding of x.y (Eq.
v
> Spherical harmonic embedding of v! (Eq.
V.
SpperiiallfammoninEnneeeeeggggg
Spherical harmonic embedding of v? (Eq.
Σ ā + + +
>> Node attributes
bi + f. @ f, @ |x
Concatenate input for messages between f,, f
mij 03_ WWW.SORRPRDDUCT.SWISSS.
å.z)
First non-linear message layer Second non-linear message layer
m;j + 03_ "i3_TENSOR_PPee TERSORR.RRDDCT..SSSS.......
āg)
m
Σ m.
Aggregate messages my
03. TENSOR PRODUCT SWISH GATE(É, SWISH.GATBCC m, ā.)
First non-linear node update layer Second linear node update layer
f.+ O3_ TENSOR PRODUCT(f, ā,)
26

---

Published as a conference paper at ICLR 2022
C.2 EXPERIMENTAL DETAILS
Compared methods. We compared against methods discussed in Sec. 3 of the main paper. namely
NMP (Gilmer al.) 2017). SchNet (Kristof al. 2017), Cormorant (Anderson al. 2019).
LINet (Miller et al.) 2020). LieConv (Finzi et al.) 2020), TEN (Thomas et al. 2018). SE(3)- )
transformer (Fuchs et al. 2020), and EGNN (Satorras al., 2021). We additionally compare
to CGCNN (Xie & Grossman 20188) PaiNN (Schitt et al. 2021 Dimenet++ (Klicpera et al. 2020)
and SphereNet (Liu et al. 2021).
Implementation details for N-body dataset. SEGNN architectures, point conv methods (SEunaar). and stecarable non non-linear point conv methods (SEnon-linca) consist of three parts (sequentially applicd):
1. Embedding network: Input + {CGa, + SwishGate CG,,). J. wherein CG2, denotes steerable linear layer conditioned on node attributes à; and which are applied per node.
2. Four message passing layers as described in Sec. 2. for SEGNN architectures and 7 convolution layers for SHiinear and minnnnlinnar ablations.
3. A prediction network: {CG2. - Swish -> Linear - MeanPool -> Linear Swish
Linear), in which CG2, denotes a steerable linear layer conditioned on node attributes à; and which maps vector of 64 scalar (type-0) features. The remaining layers are regular linear layers, in Satorras et al. (2021). CG2, and Linear applied per node.
All steerable architectures are designed such that the parameter budget l and 1 matches that of the tested EGNN implementation. We optimise models using the Adam optimiser (Kingma & Ba 2014) with learning rate Ie-4, weight decay 1e-8 and batch size 100 for 10000 epochs and minimise the mean absolute error.
The SwishGate refers to gated non-linearity (Weiler et al. 2018) when >> 0, and a swish activation (Ramachandran et al. 2017) otherwise. A message network 6m consists of {CGa;; SwishGate 4 CGa,, + SwishGate], while the update network o consists of {CGa, + SwishGate - CGs, InstanceNorm), with an instance normalisation (Ulyanov et al. 2016) implementation that is compatible with steerable vectors, based on the BatchNorm implementation of the e3nn library (Geiger et al., (2021a). In the default setting, however, the instance normalisation is turned off. We use skip-connections the message passing layers, adapting Eq. (8) to yield f/ + do (f,, E, m.y. ā.).
The convolution layers which are used instead of the message passing layers for SEinnear and Snnon-lincar ablations consist of:
SEnnear {CGa.s EjEN(O) - SwishGate + InstanceNorm! }
SEnon-lincari {CGa., + SwishGate + CGa., + E)EN(()) SwishGate -> InstanceNorm]
27

---

Published as a conference paper at ICLR 2022
Table C.1: A full ablation study on the performance and runtime different (maximum) orders of steerable features (ly) and attributes (la). The ablation study includes steerable linear convolutions, steerable non-linear convolutions, as well as SEGNN approaches. SE linear+++ +++ uses the structural information (velocity) in the two embedding and read-out layers. Performance is measure using the Mean Squared Error (MSE).
Method
MSE
Time [s]
Skiinear y 0, la
0)
SEiineur (1
.0352
.003
SEiineur (
.0130
.031
SEiincu f
.0121
.035
SF lincar f 2, la 2)
.0166
.065
SEinner ( 2, 3)
.0121
.075
nonn-linea ( 0, l,
.0382
.044
SE nonon-innar (1 f O, la
.0388
.004
SE non--line f = 1, La 1)
.0061
.030
SEgon-liness
(1 f 1, la 2)
.0060
.034
SE non-lireu
( = 2, 2, 2)
.0061
.665
SEmon-lincar
( 2, 3)
.0600
.074
SE mon--line
+ (ly 0, O)
.0392
.004
S
+ (1 0, la 1)
.0390
.004
nEnon-
++ (1y = L, 1)
.0057
.031
SE linear
++ (1 1, 2)
.0057
.036
SE linear ++ (1g 2, 2)
.0088
.667
(15 = 2, 3)
.0062
.078
SEGNNG ( 0. = 0)
.0099
.004
SEGNNG (1
0., la 1)
.0098
.015
SEGNNG (1
1, la = 1)
.0056
.025
SEGNNG (1
1, la = 2)
.0058
.031
SEGNNG (1
la 2)
.0060
SEGNNG (1
2, la 3)
.068
GEO (l D,
SEGNNG: (
la
SEGNN (
.0043
.026
SEGNNG+P (l
.0044
.032
SEGNNG,P (
.0041
.063
SEGNN
( 2, .041
.071
Table C.2: SEGNN and EGNN performance comparison the N-body system experiment. Perfor- mance is compared for different number of available training samples, and measured using the Mean Squared Error (MSB). SEGNNs significantly more data efficient.
Method
Training samples
MSF
SEGNN Cly 11,1, 1)
1000
.0068+ .00023
SEGNN 0 =1,1_
3000
.0043+
.00015
SEGNN G 11.1g
10000
.0037:: .00014
EGNN
1000
.00941 .00035
EGNN
3000
.00704: .00022
EGNN
10000
.0088: .00018
28

---

Published as a conference paper at ICLR 2022
Gravitational N-body dataset. The gravitational 100-body particle system experiment is self- designed toy experiment, which is conceptually similar to the charged N-body experiment described in the main paper. However, extend the number of particles from 5 to 100, use gravitational interaction, no boundary conditions. Particle positions are initialised from a unit Gaussian, particle velocities are initialised with norm equal one and random direction and particle mass is set one. Trajectories are generated with by integrating gravitational acceleration in natural units using the leapfrog integrator with dt 0.001 for 5000 steps. Force smoothing of 0.1 is applied (Dehnen & Read 2011).
We predict position or force at step encountered in physical and astrophysical simulations (Alvaro et 20200 Tobias et al.) [2200] Mayr
given the state
3. Such problems are often
al. 202T), where usually the force (or acceleration) is used to Euler-update the positions. Thus,
reliable force prediction is of great interest. An exemplary trajectory is shown in Fig. C.1 For the
following experiments, we use 10.000 simulation trajectories in the training set, and 2.000 trajectories
for the validation and test set, respectively. The implemented architectures are 4 layer graph neural networks with roughly equal parameter budget, and we train each network for 250 epochs. The
SEGNN layers exactly the same as described above for the charged N-body dataset, the EGNN and MPNN layers are the same as described in Satorras et al. (2021).
Figure C.I: Trajectory between f and of 100 particles under gravitational interaction.
Marker shows final position at end of simulation, with opacity decreasing over time.
Table C shows results for different numbers of considered interactions (neighbours) for SEGNNs, as well as for message passing networks (MPNNs) and EGNNs. For SEGNN implementation, we input the relative position to the center of the system and the velocity vectors of type = 1 with odd parity. We further input the norm of the velocity as scalar, which altogether results in an input vector E Vo e V1 @ V1. The output is embedded as difference vector to the initial position (odd parity) the directly predicted force after 1,000 timesteps, i.e. ö V1, In doing so, we keep E(3) equivariance for vector valued inputs and outputs. The edge attributes are obtained via the spherical harmonic embedding of x, X; as described in Eq. The node attributes comprise the sum of the edge attributes plus the spherical harmonic embedding of the velocity. The MPNNs take the position and the velocity as input. Due to the vector- valued inputs and outputs, equivariance not preserved. For EGNNs, equivariance is only preserved for the prediction of future positions, for the force prediction is not.
MPNNs and SEGNN(1 0,4, 0) are conceptually very similar. Interestingly, MPNNs, which use Py PyTorch's 's Linear operation, are approximately times faster than our SEGNN implementation, which uses Einsum operations. In general, the position and force prediction for the 100-body problem is intrinsically very difficult due to the wide range of different dynamics, where the main error contributions arise from few outlier trajectories. However, aside from unavoidable errors,
29

---

Published as a conference paper at ICLR 2022
SEGNNs generalise much better than default MPNNS EGNNs do. The results strongly suggest the applicability of SEGNNs to large physical (simulation) datasets.
Table C.3: Mean Squared Error (MSE) for positional (pos) and force prediction in the gravitational
N-body system experiment, and forward time in seconds for batch size of 20 samples running NVIDIA GeForce RTX 3090 GPU.
neighbours
20 neighbours
50 neighbours
Method
pos
force
Time[s]
pos
force
Time[s]
pos
force
Timefs]
MPNN
.297
.299
.0012
.277
.273
.0044
.262
.268
.0029
EGNN
unstable
.0244
.256
unstable
.0255
.239
unstable
.0047
SEGNN( 0,
.296
.0855
.266
.276
.0088
.251
SEGNN( 1,14
.273
.0008
.277
.244
.0212
.212
.0416
30

---

Published as a conference paper at ICLR 2022
Implementation details for QM9. Our QM9 hyperparameterss strongly resemble those found in (Saoorras et al. 2021), whose architecture we took as a starting point for our own. The network consists of three parts (sequentially applied):
Embedding network: Input - {CGs, + SwishGate - CGa, J. wherein CGa; denotes steerable linear layer conditioned node attributes à and which are applied per node.
2. Seven message passing layers as described in Sec. 2.1
3. A prediction network: {CG,, R. > Swish - Linear >> MeanPool > Linear - Swish
Linear), in which CG2, denotes steerable linear layer conditioned on node attributes a; and which maps vector of 128 scalar (type-0) features. The remaining layers are regular linear layers, as in Satorras et al. 2021 CG2, and Linear are applied per node.
Message network, update network and skip-connection are used as deseribed for the N-body task. In contrast the other tasks, the hidden irreps correspond N copies for every order, selected to correspond to 128-dimensional linear layer in terms of number of parameters.
We optimise models using the Adam optimiser (Kingma & Ba 2014) with learning 5e-4, weight decay 1e-8 and batch size 128 for 1000 epochs and minimise the mean absolute error. The learning rate is reduced by factor of ten at 80% and 90% of training. Targets are normalised by subtracting the mean and dividing by the mean absolute deviation. We found that pre-processing the molecular graphs for a specific cutoff radius on QM9 was beneficial.
function of cutoff radius for the train partition of QM9. shows that a network with cutoff radius of 2Á sends 6x fewer messages on average in every layer compared to cutoff radius of 5Á, with far smaller variance.
Figure C.2 shows the number of messages
400
350
300
250
Mean number
200
100
Transition from disconnected to connected graphs
2
5
8
Radius
Figure C.2: Mean number of messages as function of cutoff radius for the QM9 train partition. Standard deviation is indicated by the shaded region. The dashed black line indicates the transition from disconnected to connected graphs and may be thought of as minimum cutoff radius.
Implementation details for the Open Catalyst Project OC20 dataset. Our OCP hyperparam- eters resemble those from the N-body task, but with an overall larger network architecture. The network consists of three parts (sequentially applied):
Embedding network: Input {CG,, + SwishGate CG,, J. wherein CG,, denotes steerable linear layer conditioned on node attributes à; and which are applied per node.
2. Seven message passing layers as described in Sec. 2.1
3. A prediction network: (CGg, - Swish - Linear MeanPool > Linear Swish
Linear), in which CG2, denotes steerable linear layer conditioned on node attributes à; and which maps vector of 256 scalar (type-0) features. The remaining layers are regular linear layers, in (Saoorras et al.
2021
CGO and Linear are applied per node.
31

---

Published as a conference paper at ICLR 2022
Message network, update network and skip- coppeenneetions used deseribed for the QM9 tasks. The imput is a -i6ennnnnsioall one-hot representation of atom types, all hidden vectors are 256- dimensional or their steerable "weight balanced" equivalent (ef Sec. 4), where we divide into equally large type-l sub-vector spaces.
We optimise models using the A optimiser (Loshchilov & Hutter 2017) with learning rate le-4, weight decay le-8 and batch size for 20 epochs and minimise the mean absolute error. The learning rate reduced by a factor of 50% and 75% of training.
Ablation study for the Open Catalyst Project OC20 dataset, speed comparison. Table C. shows SEGNN ablation of the performance when using different (maximum) orders of steerable features y and attributes la- ablation for different cutoff radii of 5Á and 3 as well as runtime comparisons, All results are obtained on the validation set. The case l, = and Iy = corresponds to our EGNN implementation. The SEiineur corresponds to implementation outlined in Alg. and can be seen as instance of e.g. TFNs TTomas et al. (2018). All methods have been optimised using roughly the same parameter budget SEGNNs.
Table C.4: Comparison the validation sets of the OC20 IS2RE tasks. Model selection is done on those datasets which are conceptually the nearest. We compare the performance for the usage of different (maximum) orders of steerable features (ly) and attributes (l,) ) as well as the performance for cutoff radii of 5Á and 3Á on the OC20 dataset. The case L, and Is corresponds to our EGNN implementation. The SEnnca corresponds to our implementation as outlined in Alg. and can be seen as instance of e.g. TFNs (Thomas al. (2018). Runtime numbers are reported for one forward pass and are measured for batch of size 8 on Nvidia Tesla V100 GPU. Higher orders of y and in general do not improve results. The Dimenet++ runtime is given as reference.
Energy MAE [oV] L
EwT
Time
Model
0OD Ads
Cat
OOD
OOD Ads
OOD
Both
EGNN
0.5497
0.6851
0.5519
0.6102
4.99%
2.50%
4.71%
2,88%
0.025
EGNNo, 0. 3A
0.5931
0.7411
0.6182
0.6564
4.87%
2.48%
4.34%
2.63%
0.015
SEincelh
1,5
0.5841
0.7655
0.6358
0.7001
4.32%
2.51%
4.55%
2.66%
0.028
SEmer
1,3
0.6020
0.7991
0.6631
0.7104
4.27%
2.29%
4.12%
2.56%
0.022
SEGNN 1.1. 1,84)
0.5341
5.32%
2.80%
3.09%
SEGNN,, 11.6 1,300
0.5523
0.6761
0.5609
0.6127
5.03%
2.59%
4.41%
2.70%
0,040
SEGNNI, -1.1 2.550
0.5337
0.6419
0.5389
0,5764
5,41%
2.78%
4.89%
3.09%
0.112
SEGNNI,
0.5497
0.6662
0,5587
0.6060
4.98%
2.51%
4.40%
2.60%
0.051
SEGNNA, = 2.1. 2.33
0.5369
0.6473
0.5335
0.5822
5.22%
2.82%
4.8%%
2,98%
0.234
Dimenet++
0.062
D LICENSES
This work would not have been possible without the existence of free software. As we have built extensively on existing assets, we mention them and their licenses here.
Our codebase uses Python under the PSF License Agreement, using NumPy under the BSD 3-Clause "New" or "Revised" License. Our models were implemented PyTorch (Paszke et al. 2019) (BSD license) using CUDA (proprietary license). We used PyTorch extensions such PyTorch Geometric (ey & Lenssen 2019) (MIT license) and e3nn (Geiger et al. 2021b) (MIT license). For tracking runs used Weights & Biases (Biewald 2020
(MIT license).
The QM9 (Ramakrishnan et al. 2014: Ruddigkeit et al.) 2012) dataset was used under CCO 1.0 Universal (CCO 1.0) Public Domain Dedication and OC20 CCanussot et al. (020)) was used under Creative Commons Attribution 4.0 License. The QM9 dataloaders were adapted from (Anderson
al. 2019) under Educational and Not-for-Profit Research License.
32


Reviews:
[{"review_id": "u9iJDAO9U2-", "reviewer": "Reviewer_vbV1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I have listed the contributions of this submission in the summary above. I will mainly focus on the questions below:\n\n1.\tThis paper is not easy to follow. I understand that this is not owing the writing (the presentation is good anyway) but due to the abundant preliminaries to understand the geometry of steerable functions. The authors are suggested to separate the introductions of the preliminaries, related work, and the proposed methodology, instead of introducing them as a whole (section 2 and 3). \n\n2.\tThe authors claim that the proposed SEGNN (Eq. (16)) is a general form of various previous steerable models (Eq. (14-15)). In my understanding, the main difference of SEGNN from previous steerable models (Thomas et al., 2018 and Fuchs et al., 2020) is that the hidden feature h in Eq. (16) is now a concatenation of the node steerable features and the squared relative distance. The idea of making the MLP (i.e. W) conditional on the attribute a_{ij} (e.g. be a spherical harmonic embedding of xj-xi) has been explored before. So, it seems the contribution is minor. Do I miss something? Perhaps, the authors should specify the form of each previous model under the same denotations and pipeline, from which we can easily check the difference between different models.\n\n3.\t In terms of the nonlinearity, the authors claim that “Although powerful, such layers only (pseudo1) linearly transform the graphs and non-linearity is only obtained via point-wise activations.” However, in my opinion, Eq. (16) first computes the linear steerable convolution and then performs the nonlinear activation \\sigma which is also point-wise. So, why we have the claim that the convolution is nonlinear? By the way, what is the form of \\sigma? Does it directly follow the previous papers?\n\n4.\tI do not check the very details of the experimental part. Just justified from the performance, it seems on QM9, the performance between SEGNN and EGNN is close. Given that more complexity is involved in SEGNN for the computation of Spherical harmonic and Clebsch-Gordan products, the benefit of involving steerable features is not well verified experimentally. Besides, on OC20, there is no comparison with EGNN, why? Regarding the ablation studies, the difference between SEGNN-linear and -nonlinear is unclear. And it is suggested that the authors should compare their SEGNN with previous models under a more complete and fair setting. Since they are claimed as a specific case of SEGNN, the authors can keep other parts unchanged and check which component works actually. \n", "summary_of_the_paper": "This paper introduces a kind of Steerable E(3) Equivariant Graph Neural Networks (SEGNNs), exhibiting these merits: 1) Compared to EGNN (Satorras et al., 2021) and other concurrent works, SEGNN does not restrict the message passing to be scalars but to be steerable features, and it also permits the injection of geometric and physical quantities into node updates. 2) Compared to previous steerable architectures (such as the works by Thomas et al., 2018 and Fuchs et al., 2020), SEGNN employs more general steerable convolutions, built upon the message passing network (Gilmer et al., 2017). The experiments are conducted on several benchmarks: N-body, QM9 and OC20.", "summary_of_the_review": "I generally think the proposed general steerable GNN is meaningful. Yet, given certain unclarities related to the novelty and experimental significance, I initially suggest that it is just marginally abve the acceptance bar. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "stqZEaFUDJ", "reviewer": "Reviewer_k14d", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "### Strengths\n- There are several existing equivariant graph neural networks including TFN, SE(3)-transformers, Schnet, and E(n)-Equivariant GNNs (EGNN).  I was thus a bit surprised to see there does not exist a full-power equivariant GNN in the popular message passing flavor (MPNNs).  By “full power” I mean able to handle a range of equivariant feature types and learn over a large space of interwiners between them.  Schnet and EGNNs are MPNNs which are technically equivariant but primarily handle scalar features (invariants) and to a limited extent vector features (standard rep) but without representing the full space of equivariant linear mappings.  The current work is thus novel and extends to a very reasonable combination of successful aspects since both MPNNs and equivariant networks with higher order features types have both proved successful in the past.\n- In particular the current method uses anisotropic convolutions which is one of the primary advantages of uses non-invariant equivariant feature types. \n- The method achieves competitive results with several baselines on an N-body problem and two molecular properties tasks.\n- One of the more interesting aspects of the method is that convolution is taken not necessarily over the usual spatial base space.  The graphs here are embedded in R^3, however, the features may be considered as living over a higher dimensional state space consisting of position and velocity.  In physics, it is often more reasonable to compute over this higher dimensional space, and the author find that doing so improves their model.  \n\n### Weaknesses\n- The paper suffers from some lack of precision and formalism.  The authors describe their method as “non-linear” convolution.  With respect to which definition of convolution?  Section 3, in particular is very difficult to read.  The first part appears to try to compare regular group convolutions and steerable group convolutions but without defining either or stating a precise theorem. \n- Section 3 is also claimed as a contribution in the form of “unifying view.”  As it stands, I would not consider that to be the case.    \n- The authors claim, for example, that their method provides a new class of activation functions for equivariant NNs, but so far as I can tell, the non-linearity is comparable to TFN, this claim merely stems from considering a NN as a non-linearity. This does not seem like a useful extension of the terminology and obfuscates the contribution. \n\n### Questions\n- Why are TFN and EGNN not used as baselines for OC20? \n- DimeNet, SphereNet, PaiNN are all shown with better numbers for QM9, but an asterisk notes they use a different set.  It is hard for me to know whether it is reasonable for me to use these numbers to compare.  Is the code or data not available to make a fair comparison? \n- On page 7, it is noted that l_a = l_f = 1 usually worked best.  Isn’t this case already considered by SchNet and EGNN? ", "summary_of_the_paper": "This paper proposes SEGNNs, a message passing neural network which is equivariant to 3D rotations.  The message passing neural network is a type of graph neural network consisting of edge networks and node networks.  Equivariance is achieved by enforcing SO(3) equivariance on the node and edge networks using what the authors call Steerable MLPs.  Steerable MLPs use spherical harmonics and Clebsch-Gordon coefficients similar to Tensor Field Networks but without aggregation.  SEGNNs may be considered non-linear convolutional networks with equivariant anisotropic messages.  The base space for the convolution may be taken to be not just spatial, but a state space consisting of, e.g. position and velocity.  The method achieves competitive results with several baselines on an N-body problem and two molecular properties tasks.\n", "summary_of_the_review": "Although parts of the paper were difficult to understand and some of the contributions are questionable, the core contribution, making equivariant GNNs in the message passing flavor is a useful step and together with considering higher-dimensional state spaces seems to result in improved performance.  Thus I tend towards accept, but I would prefer revisions. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "fwzHYn5rT8K", "reviewer": "Reviewer_zF5H", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper is well-motivated and well-written. I have few concerns below for the authors to address.\n\nMinor concerns:\n1. How to understand that using the spherical harmonic embedding of relative positions is including physical information? Is that a pure geometric information?\n2. How to include relative force or relative momentum in the molecular tasks, e.g., the QM9 and OC20 experiments used in the paper.\n3. I assume that EGNN should be an important baseline of this paper, but seems that EGNN has not been reported on OC20.\n4. I'm very curious that why equivarience is important (or say, how to prove it) for the molecular tasks, e.g., QM9 and OC20, which induces a series of equivarient model architecture development. As far as I know, some methods which aren't equivarient perform well on similar tasks, e.g., Very Deep Graph Neural Networks Via Noise Regularisation, Godwin et al.\n5. The performance of the proposed method is not very competitive. For example, UNiTE largely outperform baseline methods (SphereNet, PaiNN, etc) on QM9, but the performance gain of proposed SEGNN is marginal. Similarily, GemNet (Klicpera et al., NeurIPS 2021) and NoisyNode (Godwin et al.) perform well on OC20 IS2RE.\n6. Considering that the recent GemNet is also a powerful equivarient GNN, would the authors make a more comprehensive discussion to explain the pros/cons between these two methods?\n7. Equivarient model is a hot topic today in the community, and they could be developed by different techniques, e.g., integral (group equivarient CNN, spherical CNN,...), tensor product (TFN, SE(3)-Tr, ...), or equivarient coordinates (EGNN, AF2, ...). Would the authors explain more about how should we think about these equivarient models, how to justify whether an equivarient model is better than others, and what's the future directions of equivarient models?\n8. I'm not very familiar with the mathmatics in this paper, so I could not guarantee all the theories are correct in the paper. I'm open to other reviewers opinions.\n", "summary_of_the_paper": "This study presents a improved version of EGNN (E(n) Equivariant GNN, Victor et al. 2021), where the node features and edge attributes are replaced by steerable vectors. In addition, the MLPs in message passing should also be replaced by steerable MLPs. The advantage of steerable vectors on the node and edge is that attributes are not restricted to scalars but also can be covariant information (e.g., vectors, tensors). The experimental results on N-body, QM9 and OC20 demonstrate the effectiveness of proposed method.", "summary_of_the_review": "Overall, I enjoy reading this paper, but I have minor concerns. I'm definitely willing to raise my score if my concerns are well addressed.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "EIJvpwOheB", "reviewer": "Reviewer_aA3f", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The starting point for this paper is the observation that regular MPNNs (of which CNNs can be seen to be a special case) get their power from the fact that node features (which are usually invariant scalars) are transformed and propagated in a non-linear manner. Work on equivariant GNNs  is more general, but they still linearly transform the graph and the non-linearity is only obtained via point-wise activations. The authors propose to generalize this formalism and propose E(3) equivariant graph networks, viewing the message passing phase itself as performing non-linear convolution. Further, covariant information is incorporated into the nodes and edges, and they also transform equivariantly. A major inspiration is a recent work by Satorras et al. Indeed, the paper can be seen building on it directly while adding in what they call \"non-linear convolutions\" and augmenting the nodes and edges to include covariant information. \n\nTo go into more detail about some of the above points: the usual message passing formalism is defined through equations (1) and (2). phi_m is the message function and phi_f is the aggregation/update function. phi_m and phi_f are required to be equivariant, specifically to the group E(3). Furthermore, physical and geometric information is incorporated in the node and edge attributes as steerable vectors which transform covariantly as the message passing propagates. For this, the authors define steerable MLPs on page 3 (same functional form as MLPs, but the weight matrices are conditioned on geometric information using the CG products, and thus support covariance). The exact form is encapsulated in equation (6). Section 2.1 is essentially the same scheme as in Satorras et al. but with the steerable MLP added in. Section 3 goes over the related work on viewing message passing as convolution. The usual convolution layers are discussed first (eq. 6), followed by their usual treatment in point clouds aka \"point convolution.\" The important point is that the matrices W (the convolution kernel) are conditioned on relative positions. The exact form depends on the type of input data. The setup used in the paper corresponds to a steerable convolution given in equations (12) and (13) (the relative distances are embedded in the spherical harmonic space). Because of the CG products that W crucially relies on, the authors interpret their method as performing a non-linear convolution (the approach of Satorras can be seen as corresponding to the special case when the kernel has an isotropic form). \n\nQuite simply, I see the paper as proposing an extension of the approach of Satorras. The message passing scheme is the same. The relative distances are embedded into the spherical harmonic space, the node and edge attributes are steerable, and the message passing ensures these are transformed covariantly. In this sense, I feel like the paper has limited novelty, or at least overstates its contributions (this happens frequently throughout the paper).  \n\nIn the beginning one gets the impression that something new is going on with \"non-linear convolutions\", but it just amounts to taking CG products, an idea that has been used in other contexts quite successfully 2-3 years ago. It is also said that incorporating steerability into GNNs is a new contribution -- this is only true with respect to the work of Satorras which just passes invariant messages, which was indeed pleasantly surprising at the time in terms of performance and simplicity. For instance the paper \"Covariant Compositional Networks for Learning Graphs.\", Risi Kondor, Hy Truong Son, Horace Pan, Brandon Anderson, Shubhendu Trivedi, has considered this problem in detail. The case for steerability is mentioned in the abstract, and the intro. Besides, the \"type 0\" and \"type k\" message passing types are also discussed in detail later. The difference is that node and edge attributes are not required to be covariant with respect to some Euclidean group. However, the main idea there, although it just deals with the permutation group, is about steerability. One can replace the defining representation there with a representation for a group of interest, along with attaching steerable vectors on nodes and edges, and their approach will also work with the same import as reported in this paper. Papers such as \"Vector Neurons: A General Framework for SO(3)-Equivariant Networks\" Congyue Deng, Or Litany, Yueqi Duan, Adrien Poulenard, Andrea Tagliasacchi, Leonidas Guibas, are not cited and should be discussed properly, along with some related papers such as Jing et al. (2020). These are more related to this paper than the author let on. \nTaking tensor products can also become prohibitive quite quickly, unless I missed it (and I am sorry if I did), I don't see a discussion about this in the main paper (nor about what L is considered and what effect it has on the runtime versus accuracy curve). \n\nThe experimental results reported in the paper are quite encouraging, and for me the hallmark of the paper. \n\nMinor comment:\nPage 6, below equation (15): \"There are, however, two import differences\" -> \"There are, however, two important differences\"", "summary_of_the_paper": "The paper extends the recently proposed E(n) GNN (Satorras et al.) such that node and edge attributes can contain covariant information. This allows the easy incorporation of geometric and physical information in the MPNN framework in a steerable manner. Excellent results are reported across a set of benchmarks.", "summary_of_the_review": "I agonized over the rating to give this paper. On one hand, it reports a clear and unambiguous benefit of incorporating covariant information in the node and edge attributes. It is also well-motivated and well written and is definitely valuable. However, the technical contribution of the paper at this point (in how the equivariant neural network literature is progressing) is relatively minor. The paper also seems to overclaim its contributions, one case being in the sense of how it relates to steerable graph networks (although this happens because they are positioning themselves in contrast to Satorras et al., Thomas et al., etc). While I am a big fan of simple methods that just work well, I feel conflicted about this paper. I am currently rating the paper a 6 (which is mostly due to their experimental results), I would be willing to adjust my rating if the authors can present arguments showing that my concerns are unfounded or not rational. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "8O8P5KQEC9C", "reviewer": "Reviewer_Mqkb", "summary": "", "questions": "", "limitations": "", "rating": 10, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "In my opinion, this is an outstanding paper that synthesizes an exciting area of recent work into a relatively authoritative general framework. I found the exposition to be extremely clear, especially given the technical nature of the subject matter. I also really liked the related works section that did a great job of positioning this paper in the context of the literature. Finally, the experimental results and ablation studies are thorough and seem to show a significant advantage for equivariant networks.\n\nTo summarize what I thought were the strengths of the paper:\n\n1. The exposition was very clear throughout, but especially in the description of the model, the introduction of steerable vector spaces, the affine transformations, and related work. The supplementary material was also very thorough and contained a lot of really interesting and useful information. I would almost encourage the authors to publish a longer form version of this paper with everything in the main text.\n\n2. I really liked the identification of the direct sum of steerable vector spaces (including multiplicity) with a type system that makes it straight-forward to add new physical quantities.\n\n3. The formulation seems very general; while the authors apply it to MPNNs as defined in Gilmer et al. it seems like you could plug in the equivariant MLPs into any graph-network architecture.\n\n4. The included code looks relatively clean; do the authors plan to open source it?\n\n5. The empirical results seem very strong and three pretty different datasets. The Open Catalyst results are particularly nice to see, since the dataset seems much harder / larger than QM9 or the N-body system.\n\n6. The ablations seem pretty thorough and it’s nice to see that equivariance seems to improve results of the l=0 (invariant) baseline.\n\nWhile I did really like this paper, I do think there are a few places where it could be a bit more clear.\n\n1. Although most of the design decisions seemed quite reasonable, I thought the paper could do a better job describing how the conditioning vectors, a, were chosen. Are they always the node or edge features? For the N-body system, were the conditioning vectors the velocity vector or the displacement vector to the center of the system? Did you try any other choices?\n\n2. I would include a description of the gated nonlinearities in the SI of the paper (if not also in the main text). I would also be a little bit more straightforward about the fact that you use gated non-linearities. In particular, in the text you currently write [1] but in the text you only ever use gated non-linearities. I would lead with the choice you made and then mention that one could use any equivariant MLP.\n\n3. One of the things I liked about this paper was how straightforward the proposed equivariant MLPs are. The submitted code uses e3nn, which (while very nice) is a pretty complex codebase. Can the authors discuss this choice vs writing things from scratch using Clebsch-Gordan coefficients as described in the paper?\n\n4. Do the authors ever use the fact that the weights can be made dependent on a parameter? I couldn't find any place in the paper where this was used (perhaps I missed it), but if it's not used then I'm not sure it adds much to the discussion.\n\n[1] _The common recipe for deep neural networks is to alternate linear layers with element-wise non-linear activation functions. In the steerable setting, careful consideration is required to ensure that the activation functions are equivariant; currently available classes of activations include Fourier-based (Cohen et al., 2018), norm-altering (Thomas et al., 2018), or gated non-linearities (Weiler et al., 2018). All these can be used in alternation with (6). The resulting steerable MLPs themselves in turn provide a new class of steerable activation functions that is able to directly leverage local geometric cues. Namely, through steerable node attributes a, either derived from the physical setup (forces, velocities) or from predictions (similar to gating), the MLPs can be applied node-wise and generally used in steerable feature fields as non-linear activations._\n", "summary_of_the_paper": "Motivated by a wide array of applications across the physical sciences, there has recently been a flurry of interest in building neural networks that are equivariant to the symmetries of the Euclidean group in three-dimensions. In my opinion, this paper presents a particularly elegant and general framework for building equivariant networks that generalizes a number of recent results. There are two essential ingredients: 1) steerable linear transformations and 2) steerable nonlinearities. Both pieces are defined over direct sums of steerable vector spaces and, with these two ingredients, it is easy to promote any message-passing neural network to an equivariant one. By describing the problem in terms of direct sums of steerable vector spaces, there is sort of a nice “type checking” nature when wanting to include new geometric quantities.", "summary_of_the_review": "A very strong paper that provides an elegant framework for building E(3) equivariant graph networks. The proposal is very general and well-motivated; it is also well-situated within the literature. The exposition is generally high quality, and the experiments are rigorous and compelling. The supplementary material gives thorough background and the code seems pretty clean. A few comments / questions about some details of the paper remain outstanding.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "1kVen_62dP2", "reviewer": "Reviewer_SAGC", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is well written. The main strengths of this paper are the strong empirical results on recently introduced benchmark (OC20) for learning on small molecules. Authors put down a good effort to jot connection between various convolutional and message passing algorithms in the main paper. Although I enjoyed reading the manuscript I think it would benefit from a clearer presentation of the novel technical contributions. For example, some form of comparison using pseudo code will help readers.\n\nThe overall contribution and novelty is minimal. As pointed earlier, SEGNNs is an extension of EGNNs (Satorras et al., 2021) wherein the linear layer and nonlinear activation within MLPs are replaced by their steerable version. The steerable linear transformations was earlier proposed in (Anderson, Fuchs, Thomas) while steerable activation is extended from (Wieler).\n\nOther drawback and feedback: \n\n1. For comparison on QM9, SEGNN performs better than EGNN on first six properties. However, SEGNNs performance on last six properties (non-energy ones) is much worse than previous state-of-the-art. Author points that more involved architecture might perform better. Given the level of complexity in SEGNN, I feel SEGNN is quite involved. How about increasing order or cutoff radius ?\n2. From an ablation study on cutoff radius and order provided in Table C.2, I noticed that result improves considerably on increasing the cutoff radius. Can you also provide ablation results for higher cutoff radius (say at 7, 10, 15 A) ? Will lower order and increased cutoff radii give similar result ?\n3. I suggest that for Table 2 and elsewhere, please present all baselines and results with fixed number of decimal points. This makes it easier to compare across models. Also, for each of the properties, please highlight the best model in bold. \n4. The model claims scalability to large molecules due to the use of fewer MP steps. But the application of same on large molecule datasets is missing. OC20 contains maximum of 70 atoms which is not large enough. Some datasets such as GEOM-Drugs contains max. of 180 atoms. Moreover, given that the model is applicable to any physical system, I would like to see the performance on fluid simulation (A).\n5. What is the actual difference between SE(non-linear) and SEGNN_G ?\n6. The results on N-body system in Table 1 suggest that in comparison to EGNN, SEGNN (order 1) significantly reduce error but at the cost of ~10x slower runtime.\n7. What is odd parity ?\n\nA. Learning to Simulate Complex Physics with Graph Networks, ICML 2020.", "summary_of_the_paper": "This work extends the concept of steerable kernel to nonlinear transformations in MLP. The steerable MLPs are then incorporated into recently introduced equivariant message passing layers thus forming Steerable E(3) Equivariant GNNs (SEGNNs). At the outset SEGNNs are generalisation of EGNNs (Satorras et al., 2021) with non-isotropic message functions utilising relative pose information instead of mere distances. When applied to molecular property prediction tasks, it achieves mixed results on QM9 data and state-of-the-art results on IS2RE task of OC20 dataset.", "summary_of_the_review": "Drawing connection with steerable network, this work attempts to rectify a missing link in prior art for equivariant network. Although, the technical contribution is minimal, the paper would have stand out if the empirical results were strong. I suggest that authors improve upon their results and further demonstrate wide applicability.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This work combines steerable MLPs with equivariant message passing layers to form Steerable E(3) Equivariant GNNs (SEGNNs). It extends previous work such as Schnet and EGNNs, by allowing equivariant tensor messages (in contrast to scalar or vector messages). The paper also provides a unifying view of related work which is a nice overview for the ML community. It is overall well written, but would benefit from further revision to improve readability in some parts (in particular section 3, cf. reviews).
It shows strong empirical results on the IS2RE task of the OC20 dataset and mixed results on the QM9 dataset.
Decision: Accept (Spotlight)