text:
\section{Mathematical background}
\label{app:math}
This appendix provides the mathematical background and intuition for steerable MLPs. We remark that the reader may appreciate several related works, such as, \citep{thomas2018tensor,anderson2019covariant,fuchs2020se}, as excellent alternative resources\footnote{Each of these works presents unique view points that greatly influenced the writing of this appendix.} 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:
\begin{enumerate}
    \item \textit{Group definition} and \textit{examples of groups} (Section \ref{sec:groupdefinition}). The entire framework builds upon notions from group theory and as such a formal definition is in order. In this paper, we model transformations such as translation, rotation and reflection as groups.
    \item \textit{Invariance}, \textit{equivariance} and \textit{representations} (Section \ref{sec:equivariance}). 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; a representation formalises the notion of transformations applied to vectors in the context of group theory.
    \item \textit{Steerable vectors}, \textit{Wigner-D matrices} and \textit{irreducible representations} (Section \ref{sec:steerable}). 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 representations and can be used to define any representation (or conversely, any representation can be reduced to a tensor product of Wigner-D matrices via a change of basis). As such, the Wigner-D matrices are irreducible representations.
    \item \textit{Spherical harmonics} (Section \ref{sec:sphericalharmonics}). Spherical harmonics are a class of functions on the sphere $S^2$ and can be thought of as a 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^2$, 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.
    \item \emph{Clebsch-Gordan tensor product} and \textit{steerable MLPs} (Section \ref{sec:clebschgordan}). 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 Clebsch-Gordan tensor product. Akin to the learnable weight matrix in regular MLPs, the learnable Glebsch-Gordan tensor product is the workhorse of our steerable MLPs.
\end{enumerate}

After these concepts are introduced we will in Section \ref{sec:groupconvs} 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 to 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 as building neural networks via non-linear group convolutions.

\subsection{Group definition and the groups E($3$) and O($3$)}
\label{sec:groupdefinition}
\paragraph{Group definition.}
A group is an algebraic structure that consists of a set $G$ and a binary operator $\cdot$, the group product, that satisfies the following axioms: \textit{Closure}: for all $h,g \in G$ we have $h\cdot g \in G$; \textit{Identity}: there exists an identity element $e \in G$; \textit{Inverse}: for each $g \in G$ there exists an inverse element $g^{-1} \in G$ such that $g^{-1} \cdot g = g \cdot g^{-1} = e$; and \textit{Associativity}: for each $g,h,i \in G$ we have $ (g \cdot h) \cdot i = g \cdot (h \cdot i)$.

\paragraph{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 transformations are parametrised by pairs of translation vectors $\mathbf{x}\in \mathbb{R}^3$ and orthogonal transformation matrices $\mathbf{R} \in \text{O}(3)$. The E(3) group product and inverse are defined by
\begin{align*}
g \cdot g' &:= (\mathbf{R} \mathbf{x}' + \mathbf{x}, \mathbf{R} \mathbf{R}') \ , \\
g^{-1} &:= (\mathbf{R}^{-1} \mathbf{x}, \mathbf{R}^{-1}) \ ,
\end{align*}
with $g=(\mathbf{x},\mathbf{R}), g'= (\mathbf{x}',\mathbf{R}) \in \text{E}(3)$. One can readily see that with these definitions all four group axioms are satisfied, and that it therefore defines a group. The group product can be seen as a description for how two E($3$) transformations parametrised by $g$ and $g'$ applied one after another are described by single transformation parametrised by $g \cdot g'$. The transformations themselves act on the vector space of 3D positions via the group action, which we also denote with $\cdot$, via
\begin{equation*}
g \cdot \mathbf{y}:= \mathbf{R} \mathbf{y} + \mathbf{x},
\end{equation*}
where $g = (\mathbf{x},\mathbf{R}) \in \text{E}(3)$ and $\mathbf{y} \in \mathbb{R}^3$.

\paragraph{The orthogonal group O($3$) and special orthogonal group SO($3$).}
The group E$(3) = \mathbb{R}^3 \rtimes \text{O}(3)$ is a semi-direct product (denoted with $\rtimes$) of the group of translations $\mathbb{R}^3$ with the group of orthogonal transformations O($3$). This means that we can conveniently decompose E($3$)-transformations in an O($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 $\mathbf{R}$, as we have done before, the group product and inverse are simply given by the matrix-product and matrix-inverse. I.e., with
 $g = \mathbf{R}, g'=\mathbf{R}' \in G = \text{O}(3)$ the group product and inverse are defined by
\begin{align*}
g \cdot g' &:= \mathbf{R} \mathbf{R}' \ , \\
g^{-1} &:= \mathbf{R}^{-1} \ .
\end{align*}
The group acts on $\mathbb{R}^3$ by matrix-vector multiplication, i.e., $g \cdot \mathbf{y} := \mathbf{R} \mathbf{y}$. The group elements of O(3) are square matrices with determinant $-1$ or $1$. Their action on $\mathbb{R}^3$ 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$.

\paragraph{The sphere $S^2$ is a homogeneous space of SO(3).}
The sphere is not a group as we cannot define a group product on $S^2$ that satisfies the group axioms. It can be convenient to treat it as a homogeneous space of the groups O(3) or SO(3). A space $\mathcal{X}$ is called a homogeneous space of a group $G$ if for any two points $x,y \in \mathcal{X}$ there exists a group element $g \in G$ such that $y = g \cdot x$.

The sphere $S^2$ is 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 an XYX parametrisation of SO(3) rotations in which three rotations are applied one after another via
\begin{equation}
\label{eq:O3parametrization}
\mathbf{R}_{\alpha,\beta,\gamma} = \mathbf{R}_{\alpha,\mathbf{n}_x} \mathbf{R}_{\beta,\mathbf{n}_y} \mathbf{R}_{\gamma,\mathbf{n}_x} \ ,
\end{equation}
with $\mathbf{n}_x$ and $\mathbf{n}_y$ denoting unit vectors along the $x$ and $y$ axis, and $\mathbf{R}_{\alpha,\mathbf{n}_x}$ denotes a rotation of $\alpha$ degrees around axis $\mathbf{n}_x$. We can model points on the sphere in a similar way via Euler angles via
\begin{equation}
\label{eq:S2parametrization}
\mathbf{n}_{\alpha,\beta} := \mathbf{R}_{\alpha,\beta,0}\mathbf{n}_x \ .
\end{equation}
So, with two rotation angles, any point on the sphere can be reached. In the above we set $\gamma=0$ in the parametrisation of the rotation matrix (an element from SO($3$)) that rotates the reference vector $\mathbf{n}_x$, but it should be clear that with any $\gamma$ the same point $\mathbf{n}_{\alpha,\gamma}$ is reached. This means that there are many group elements in SO($3$) that all map $\mathbf{n}_x$ to the same point on the sphere.

\subsection{Invariance, equivariance and representations}
\label{sec:equivariance}
\paragraph{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 $\mathbb{R}^3$ via the group action. We usually think of the groups E($3$) and O($3$) as groups that describe transformations on $\mathbb{R}^3$, but these groups are not restricted to transformations on $\mathbb{R}^3$ and can generally act on arbitrary vector spaces via representations. A \textit{representation} is an invertible linear transformation $\rho(g):{V} \rightarrow {V}$ parametrised by group elements $g \in G$ that acts on some vector space ${V}$, and which follows the group structure (it is a group homomorphism) via
\begin{equation*}
\rho(g)\rho(h)v = \rho(g \cdot h)v \ ,
\end{equation*}
with $v \in V$.

A representation can act on \textit{infinite-dimensional vector spaces} such as functions. E.g., the so-called \textit{left-regular representations} $\mathcal{L}_g$ of E($3$) on functions $f:\mathbb{R}^3\rightarrow \mathbb{R}$ on $\mathbb{R}^3$ is given by
\begin{equation*}
\mathcal{L}_g[f](\mathbf{x}) = f(g^{-1} \mathbf{x}) \ ,
\end{equation*}
i.e., it transforms the function by letting $g^{-1}$ act on the domain from the left. Here we used the notation $\mathcal{L}_g[f]$ to indicate that $\mathcal{L}_g$ transforms the function $f$ first, which creates a new function $\mathcal{L}(g)[f](\mathbf{x})$, which is then sampled at $\mathbf{x}$.

When representations transform \textit{finite dimensional vectors} $v \in V = \mathbb{R}^d$, they are $d \times d$ dimensional matrices. In this work, we denote such \textit{matrix representations} with boldface $\mathbf{D}$. A familiar example of a matrix representation of O($3$) on $\mathbb{R}^3$ are the matrices $g = \mathbf{R} \in \text{O}(3)$ themselves, i.e., $\mathbf{D}(g) \mathbf{x} = \mathbf{R} \mathbf{x}$.

Finally, any two representations, say $\mathbf{D}(g)$ and $\mathbf{D}'(g)$, are \textit{equivalent} if they relate via a similarity transform via
\begin{equation*}
\mathbf{D}'(g) = \mathbf{Q}^{-1} \mathbf{D}(g) \mathbf{Q} \ ,
\end{equation*}
i.e., such representations describe one and the same thing but in a different basis, and the change of basis is carried out by $\mathbf{Q}$. Now that representations have been introduced we can formally define equivariance.

\paragraph{Invariance and equivariance.}
Equivariance is a property of an operator $\phi: \mathcal{X} \rightarrow \mathcal{Y}$ that maps between input and output vector spaces $\mathcal{X}$ and $\mathcal{Y}$. Given a group $G$ and its representations $\rho^{\mathcal{X}}$ and $\rho^{\mathcal{Y}}$ which transform vectors in $\mathcal{X}$ and $\mathcal{Y}$ respectively, an \textit{operator $\phi: \mathcal{X} \rightarrow \mathcal{Y}$ is said to be equivariant if it satisfies the following constraint}
\begin{align}
\label{eq:equivariance_app}
    \rho^{\mathcal{Y}}(g)[\phi(x)]=\phi(\rho^{\mathcal{X}}(g)[x]) \ , \ \text{for all } g\in G, x \in \mathcal{X} \ .
\end{align}
Thus, with an equivariant map, the output transforms predictably with transformations on the input. One might say that no information gets lost when the input is transformed, merely re-structured. One way to interpret Eq.~\eqref{eq:equivariance_app} is therefore that the operators $\rho^{\mathcal{X}}(g):\mathcal{X} \rightarrow \mathcal{X}$ and $\rho^{\mathcal{Y}}(g):\mathcal{Y} \rightarrow \mathcal{Y}$ describe the same transformation, but in different spaces.

Invariance is a special case of equivariance in which $\rho^{\mathcal{Y}} = \mathcal{I}^{\mathcal{Y}}$ for all $g \in G$. I.e.,  \textit{an operator $\phi: \mathcal{X} \rightarrow \mathcal{Y}$ is said to be invariant if it satisfies the following constraint}
\begin{align}
\label{eq:invariance_app}
    \phi(x)=\phi(\rho^{\mathcal{X}}(g)[x]) \ , \ \text{for all } g\in G, x \in \mathcal{X} \ .
\end{align}
Thus, with an invariant operator, the output of $\phi$ is unaffected by transformations applied to the input.
\iffalse
\rob{This means that, for E(3), all information regarding orientation and position is lost by applying an invariant transformation.}
\erik{Correct! Hence one wants to work with equivariant operators as then such geometric information can still be exploited in subsequent operations}
\fi

\subsection{Steerable vectors, Wigner-D matrices and irreducible representations}
\label{sec:steerable}
One strategy to build equivariant MLPs is to 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 ($\mathbf{x}_j - \mathbf{x}_i$). Then, any remaining operations are designed to be 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 O($3$) and show how to build equivariant MLPs through the use of the Clebsch-Gordan tensor product.

\paragraph{Wigner-D matrices are irreducible representations.}
For SO($3$) there exists a collection of representations, indexed with their order $l \geq 0$, which act on vector spaces of dimension $2l + 1$. These representations are called Wigner-D matrices and we denote them with $\mathbf{D}^{(l)}(g)$. The use of Wigner-D matrices is motivated by the fact that any matrix representation $\mathbf{D}(g)$ of SO($3$) that acts on some vector space $V$ can be ``reduced'' to an equivalent block diagonal matrix representation with Wigner-D matrices along the diagonal:
\begin{equation}
\label{eq:blockdiagonal}
\mathbf{D}(g)  = \mathbf{Q}^{-1} (\mathbf{D}^{(l_1)}(g) \oplus \mathbf{D}^{(l_2)}(g) \oplus \dots) \mathbf{Q}= \mathbf{Q}^{-1}
\begin{pmatrix}
\mathbf{D}^{(l_1)}(g) & & \\
& \mathbf{D}^{(l_2)}(g) & \\
& & \ddots
\end{pmatrix}
\mathbf{Q} \ ,
\end{equation}
with $\mathbf{Q}$ the change of basis that makes them equivalent. The individual Wigner-D matrices themselves cannot be reduced and are hence \textit{irreducible representations} of SO(3). Thus, since the block diagonal representations are equivalent to $\mathbf{D}$ we may as well work with them instead. This is convenient since each block, i.e., each Wigner-D matrix $\mathbf{D}^{(l_i)}$, only acts on a sub-space $V_{l_1}$ of $V$. As such we can factorise $V=V_{l_1} \oplus V_{l_2} \oplus \dots$, which motivates the use of steerable vector spaces and their direct sums as presented in
Sec.~\ref{sec:method}.

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 a 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 \citep{sakurai2017} and their numerical implementations in code libraries such as the \texttt{e3nn} library \citep{geiger2021}.

\paragraph{Steerable vector spaces.}
The $(2 l + 1)$-dimensional vector space on which a Wigner-D matrix of order $l$ acts will be called \textit{a type $l$ steerable vector space} and is denoted with $V_l$. E.g., a type-3 vector $\mathbf{h} \in V_3$ is transformed by $g \in \text{O}(3)$ via $\mathbf{h} \mapsto \mathbf{D}_3(g) \mathbf{h}$. We remark that this definition is equivalent to the definition of \textit{steerable functions} commonly used in computer vision~\citep{freeman1991design,helor1996steerable} via the viewpoint that steerable vectors can be regarded as the basis coefficients of a function expanded in a spherical harmonic basis. We elicit this viewpoint in Sec.~\ref{sec:sphericalharmonics} and \ref{sec:groupconvs}.

At this point we are already familiar with type-0 and type-1 steerable vector spaces. Namely, type-0 vectors $h \in V_0 = \mathbb{R}$ are scalars, which are invariant to transformations $g \in \text{O}(3)$, i.e., $\mathbf{D}_0(g) h = h$. Type-1 features are vectors $\mathbf{h} \in \mathbb{R}^3$ which transform directly via the matrix representation of the group, i.e., $\mathbf{D}_1(g) \mathbf{h} = \mathbf{R} \mathbf{h}$.

\subsection{Spherical harmonics}
\label{sec:sphericalharmonics}

\paragraph{Spherical harmonics.}
Related to the Wigner-D matrices and their steerable vector spaces are the spherical harmonics\footnote{Solutions to Laplace's equation are called harmonics. Solutions of Laplace's equation on the sphere are therefore called spherical harmonics.}. Spherical harmonics are a class of functions on the sphere $S^2$, akin to the familiar circular harmonics that are best known as the 1D Fourier basis. As with a Fourier basis, spherical harmonics form an orthonormal basis for functions on $S^2$. In this work we use the real-valued spherical harmonics and denote them with $Y^{(l)}_m:S^2 \rightarrow \mathbb{R}$.

\paragraph{Spherical Harmonics are Wigner-D functions.}
One can also think of spherical harmonics as functions $\tilde{Y}^{(l)}_m:\text{SO}(3) \rightarrow \mathbb{R}$ on SO$(3)$ that are invariant to a sub-group of rotations via
\begin{equation*}
Y^{(l)}_m(\mathbf{n}_{\alpha,\beta}) = Y^{(l)}_m(\mathbf{R}_{\alpha,\beta,\gamma}\mathbf{n}_x)=:\tilde{Y}_m^{(l)}(\mathbf{R}_{\alpha,\beta,\gamma}) \ ,
\end{equation*}
in which we used the parametrisation for $S^2$ and O($3$) given in (\ref{eq:S2parametrization}) and (\ref{eq:O3parametrization}) respectively. Then, by definition, $\tilde{Y}_m^{(l)}$ is invariant with respect to rotation angle $\gamma$, i.e.., $\forall_{\gamma \in [0,2 \pi)}: \;\; \tilde{Y}^{(l)}_m(\mathbf{R}_{\alpha,\beta,\gamma}) = \tilde{Y}^{(l)}_m(\mathbf{R}_{\alpha,\beta,0})$.
This viewpoint of regarding the spherical harmonics as $\gamma$-invariant functions on O($3$) helps us to draw the connection to the Wigner-D functions $D_{mn}^{(l)}$ that make up the $2l+1\times2l+1$ elements of the Wigner-D matrices. Namely, the $n=0$ column of Wigner-D functions are also $\gamma$-invariant and, in fact, correspond (up to a normalisation factor) to the spherical harmonics via
\begin{equation}
Y^{(l)}_m(\mathbf{n}_{\alpha,\beta}) = \frac{1}{\sqrt{2l+1}} D^{(l)}_{m0}(\mathbf{R}_{\alpha,\beta,\gamma}) \ .
\end{equation}

\paragraph{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 $\mathbf{a}^{(l)}(\mathbf{n}) := (Y^{(l)}_{-l}(\mathbf{n}), \dots, Y^{(l)}_{l}(\mathbf{n}))^T$ be the embedding of a direction vector $\mathbf{n} \in S^2$ in spherical harmonics. Then this vector embedding is equivariant as it satisfies
\begin{equation}
\label{eq:shequivariance}
\forall_{\mathbf{R}' \in \text{SO(3)}} \;\; \forall_{\mathbf{n} \in S^2}: \;\;\;\; \mathbf{a}^{(l)}(\mathbf{R}'\mathbf{n}) = \mathbf{D}^{(l)}(\mathbf{R}') \mathbf{a}^{(l)}({\mathbf{n}}) \ .
\end{equation}
Using the $S^2$ and O($3$) parametrization of (\ref{eq:S2parametrization}) and (\ref{eq:O3parametrization}) this is derived as follows
\begin{multline*}
\mathbf{a}^{(l)}(\mathbf{R}'\mathbf{n}_{\alpha,\beta}) =
\begin{pmatrix}
Y_{-l}^{(l)}(\mathbf{R}'\mathbf{n}_{\alpha,\beta}) \\
\vdots \\
Y_{l}^{(l)}(\mathbf{R}'\mathbf{n}_{\alpha,\beta}) \\
\end{pmatrix} =
\begin{pmatrix}
D_{-l0}^{(l)}(\mathbf{R}'\mathbf{R}_{\alpha,\beta,\gamma}) \\
\vdots \\
D_{l0}^{(l)}(\mathbf{R}'\mathbf{R}_{\alpha,\beta,\gamma}) \\
\end{pmatrix} =
\mathbf{D}^{(l)}(\mathbf{R}')
\begin{pmatrix}
D_{-l0}^{(l)}(\mathbf{R}_{\alpha,\beta,\gamma}) \\
\vdots \\
D_{l0}^{(l)}(\mathbf{R}_{\alpha,\beta,\gamma}) \\
\end{pmatrix} \\
\mathbf{D}^{(l)}(\mathbf{R}')
\begin{pmatrix}
Y_{-l}^{(l)}(\mathbf{R}'\mathbf{n}_{\alpha,\beta}) \\
\vdots \\
Y_{l}^{(l)}(\mathbf{R}'\mathbf{n}_{\alpha,\beta}) \\
\end{pmatrix} =
\mathbf{D}^{(l)}(\mathbf{R}')\mathbf{a}^{(l)}(\mathbf{n}_{\alpha,\beta}) \ .
\end{multline*}

\paragraph{Steerable vectors represent steerable functions on $S^2$.} Just like the 1D Fourier basis forms a complete orthonormal basis for 1D functions, the spherical harmonics $Y^{(l)}_m$ form an orthonormal basis for $\mathbb{L}_2(S^2)$, the space of square integrable functions on the sphere. Any function on the sphere can thus be represented by a steerable vector $\mathbf{a} \in V_0 \oplus V_1 \oplus \dots$ when it is expressed in a spherical harmonic basis via
\begin{equation}
\label{eq:inversefourier}
f(\mathbf{n}) = \sum_{l\geq 0} \sum_{m=-l}^l a_m^{(l)} Y_m^{(l)}(\mathbf{n}) \ .
\end{equation}
Since spherical harmonics form an orthonormal basis, the coefficient vector $\mathbf{a}$ can directly be obtained by taking $\mathbb{L}_2(S^2)$-inner products of the function $f$ with the spherical harmonics, i.e. ,
\begin{equation}
\label{eq:fourier}
a_m^{(l)} = ( f, Y_m^{(l)})_{\mathbb{L}_2(S^2)} =\int_{S^2} f(\mathbf{n}) Y_m^{(l)}(\mathbf{n}) {\rm d}\mathbf{n} \ .
\end{equation}
Equation (\ref{eq:fourier}) is sometimes referred to as the Fourier transform on $S^2$, and Eq.~\eqref{eq:inversefourier} as the inverse spherical Fourier transform. Thus, one can identify steerable vectors $\mathbf{a}$ with functions on the sphere $S^2$ via the spherical Fourier transform.

In this paper we visualize such functions on $S^2$ via glyph-visualizations which are obtained as surface plots
\begin{equation*}
\left\{ \; \mathbf{n} \, |f(\mathbf{n})| \;\;\; \left| \;\;\; \mathbf{n} \in S^2 \right. \; \right\} \;\;\; \Longleftrightarrow \;\;\;
\vcenter{\hbox{\includegraphics[width=4cm]{Figures/single_harmonic.png}}}
\end{equation*}
and each point on this surface is color-coded with the function value $f(\mathbf{n})$. The visualisations are thus color-coded spheres that are stretched in each direction $\mathbf{n}$ via $\lVert f(\mathbf{n}) \rVert$.

\paragraph{Steerable vectors also represent steerable functions on O($3$).}
In order to draw a connection between group equivariant message passing and group convolutions, as we did in Sec.~\ref{sec:convolution} of the main paper, it is important to understand that steerable vectors also represent functions on the group SO($3$) via an SO($3$)-Fourier transform. The collection of $(2l+1)^2$ Wigner-D functions $D_{mn}^{(l)}$ form an orthonormal basis for $\mathbb{L}_2(\text{SO}(3))$.

This orthonormal basis allows for a Fourier transform that maps between the function space $\mathbb{L}_2(\text{SO}(3))$ and steerable vector space $V$; the forward and inverse Fourier transform on SO($3$) are respectively given by
\begin{align}
a_{mn}^{(l)} &=\int_{\text{O}(3)} f(g) D_{mn}^{(l)}(g) {\rm d}g \ , \label{eq:O3Fourier}\\
f(g) &= \sum_{l \geq 0} \sum_{m=-l}^l\sum_{n=-l}^l a_{mn}^{(l)} D_{mn}^{(l)}(g) \ , \label{eq:O3FourierInv}
\end{align}
with ${\rm d}g$ the Haar measure of the group. Noteworthy, the forward Fourier transform generates a 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 $\mathbf{D}^{(l)}$.

\subsection{Clebsch-Gordan product and steerable MLPs}
\label{sec:clebschgordan}
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 output vector spaces via the Clebsch-Gordan 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 of our steerable MLPs.

\paragraph{Clebsch-Gordan tensor product.}
The Clebsch-Gordan (CG) tensor product allows us to map between steerale input and output spaces. While there is much to be said about tensors and tensor products in general, we here intend to focus on intuition. In general, a tensor product involves the multiplication between all components of two input vectors. E.g., with two vectors $\mathbf{h}_1 \in \mathbb{R}^{d_1}$ and $\mathbf{h}_2 \in \mathbb{R}^{d_2}$, the tensor product is given by
\begin{equation*}
\mathbf{h}_1 \otimes \mathbf{h}_2 = \mathbf{h}_1\mathbf{h}_2^T =
\begin{pmatrix}
h_1 h_1 & h_1 h_2 & \dots \\
h_2 h_1 & h_2 h_2 & \dots \\
\vdots & \vdots & \ddots
\end{pmatrix} \ ,
\end{equation*}
which we can flatten into a $d_1d_2$-dimensional vector via an operation which we denote with $\operatorname{vec}(\mathbf{h}_1 \otimes \mathbf{h}_2)$.
In our steerable setting we would like to work exclusively with steerable vectors and as such we would like for any two steerable vectors $\tilde{\mathbf{h}}_1\in V_{l_1}$ and $\tilde{\mathbf{h}}_2\in V_{l_2}$, that the tensor product's output is again steerable with a O($3$) representation $\mathbf{D}(g)$ such that the following equivariance constraint is satisfied:
\begin{equation}
    \mathbf{D}(g) (\tilde{\mathbf{h}}_1 \otimes \tilde{\mathbf{h}}_2) = (\mathbf{D}^{(l_1)}(g) \tilde{\mathbf{h}}_1) \otimes (\mathbf{D}^{(l_2)}(g) \tilde{\mathbf{h}}_2) \ .
\end{equation}
Via the identity $\operatorname{vec}(\mathbf{A}\mathbf{X}\mathbf{B}) = (\mathbf{B}^T \otimes \mathbf{A}) \operatorname{vec}(\mathbf{X})$, we can show that the output is indeed steerable:
\begin{multline*}
\operatorname{vec}\left(\; (\mathbf{D}^{(l_1)}(g)\tilde{\mathbf{h}}_1) (\mathbf{D}^{(l_2)}(g)\tilde{\mathbf{h}}_2)^T \;\right)
=
\operatorname{vec}\left(\; \mathbf{D}^{(l_1)}(g)\tilde{\mathbf{h}}_1 \tilde{\mathbf{h}}_2^T \mathbf{D}^{(l_2)T}(g) \;\right) \\
=
\left( \mathbf{D}^{(l_2)}(g) \otimes \mathbf{D}^{(l_1)}(g) \right) \operatorname{vec}\left(\; \tilde{\mathbf{h}}_1 \tilde{\mathbf{h}}_2^T  \;\right) \ .
\end{multline*}
The resulting vector $\operatorname{vec}(\tilde{\mathbf{h}_1} \otimes \tilde{\mathbf{h}}_2)$ is thus steered by a representation $\mathbf{D}(g) = \mathbf{D}^{(l_2)}(g) \otimes \mathbf{D}^{(l_1)}(g)$. Since any matrix representation of O($3$) can be reduced to a direct sum of Wigner-D matrices (see (\ref{eq:blockdiagonal})), the resulting vector can be organised via a change of basis into parts that individually transform via Wigner-D matrices of different type. I.e. $\tilde{\mathbf{h}} = \operatorname{vec}(\tilde{\mathbf{h}}_1 \otimes \tilde{\mathbf{h}}_2) \in V = V_0 \oplus V_1 \oplus \dots$, with $V_l$ the steerable sub-vector spaces of type $l$.

With the CG tensor product we directly obtain the vector components for the steerable sub-vectors of type $l$ as follows. Let $\tilde{\mathbf{h}}^{(l)} \in V_l = \mathbb{R}^{2l+1}$ denote a steerable vector of type $l$ and $h^{(l)}_m$ its components with $m=-l,-l+1,\dots,l$.
Then the $m$-th component of the type $l$ sub-vector of the output of the tensor product between two steerable vectors of type $l_1$ and $l_2$ is given by
\begin{align}
(\tilde{\mathbf{h}}^{(l_1)} \otimes_{cg} \tilde{\mathbf{h}}^{(l_2)})^{(l)}_{m} = \sum_{m_1=-l_1}^{l_1} \sum_{m_2=-l_2}^{l_2} C^{(l,m)}_{(l_1,m_1)(l_2,m_2)} h^{(l_1)}_{m_1} h^{(l_2)}_{m_2} \ ,
\label{eq:CG_productappendix}
\end{align}
in which $C^{(l,m)}_{(l_1,m_1)(l_2,m_2)}$ are the Clebsch-Gordan coefficients.
\iffalse
\rob{Can we give some extra intuition here like: The Clebsch-Gordan coefficients carry out a change of basis that ensures that the resulting vector is again steerable, as seen previously in eqn A.26.} \erik{Small addition: The Clebsch-Gordan coefficients carry out a change of basis such that the resulting vector is again steerable relative to our standard basis of Wigner-D matrices, as seen previously in Eq.\eqref{eq:blockdiagonal}}
\fi
The CG tensor product is a sparse tensor product, as generally many of the $C^{(l,m)}_{(l_1,m_1)(l_2,m_2)}$ components are zero. Most notably, $C^{(l,m)}_{(l_1,m_1)(l_2,m_2)}=0$ whenever $ l < |l_1 - l_2|$ or $l > l_1 + l_2$. E.g., a type-0 and a type-1 feature vector cannot create a type-2 feature vector. Well known examples of the CG product are the scalar product ($l_1=0, l_2=1, l=1$), which takes as input a scalar and a type-1 vector to generate a type-1 vector, the dot product ($l_1=1, l_2=1, l=0$) and cross product ($l_1=1, l_2=1, l=1$).

\paragraph{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 a tensor product between steerable vectors of mixed types. The CG tensor product can be parametrised by weights where the product is scaled with some weight $w$ for each triplet of types $(l_1, l_2, l)$ for which the CG coefficients are non-zero\footnote{In terms of the \texttt{e3nn} library \citep{geiger2021} one then says a path exists between these 3 types.}.  We indicate such CG tensor products with $\otimes^{\mathbf{W}}_{cg}$. While in principle the CG tensor product takes two steerable vectors as input, in this work we mainly use it with one of its input vectors ``fixed''. The CG tensor product can then be regarded as a linear layer that maps between two steerable vector spaces and we denote this $\mathbf{W}_{\tilde{\mathbf{a}}} \tilde{\mathbf{h}} := \tilde{\mathbf{h}} \otimes_{cg}^\mathbf{W} \tilde{\mathbf{a}}$, with $\tilde{\mathbf{a}}$ the steerable vector that is considered to be fixed. With this viewpoint we can design MLPs in the same way as we are used to with regular linear layers, and establish clear analogies with convolution layers
(Sec.~\ref{sec:convolution} of the main paper).

While in general the CG tensor product between two steerable vectors of type $l_1$ and $l_2$ can contain steerable vectors up to degree $l=l_1 + l_2$, one typically ``band-limits'' the output vector by only considering the steerable vectors up to some degree $l_{max}$. In our experiments we band-limit the hidden feature representations to $l_{max}=l_f$.

Finally, the amount of interaction between the steerable sub-vectors in the hidden representations of degree $\tilde{\mathbf{h}} \in V_{L=l_f} = V_0 \oplus V_1 \oplus \dots V_{l_f}$ is determined by the maximum order $l_a$ of the steerable vector $\tilde{\mathbf{a}}$ on which the steerable linear layer is conditioned. Namely, the CG tensor product only produces type $l$ steerable sub-vectors for $|l_f - l_a| \geq l \leq |l_f + l_a|$. For example, it is not sensible to embed positions as steerable vectors up to order $l_a=5$ when the hidden representations are of degree $l_f=2$; the lowest steerable vector type that can be produced with the CG product of a $l=5$ sub-vector of $\tilde{\mathbf{a}}$ with any sub-vector of the hidden representation vector will be $l=3$ and since the hidden representations are band-limited to a maximum type of $l_f=2$ higher order vectors will be ignored.
\iffalse
\johannes{from here on it is a bit heard to read if you don't know what is coming. We might now need this\\ } \erik{I think we can then remove this then since we are not really experimenting this way, but we could add it later on once we have obtained more insight in what optimal configurations are.} Conversely, the band-limit $l_f$ of the hidden representations can be arbitrarily high as with any $l_a>0$ interaction takes place between sub-vectors of adjacent type. E.g. with $l_f=5$ and $l_a=1$ it takes five steerable linear layers $\mathbf{W}_{\tilde{\mathbf{a}}}$ before information in type-$0$ vectors of the hidden representation is propagated to type-$l_f$ sub-vectors, whereas this happens in a single layer when $l_a=5$. \rob{This why the experiments in section \ref{sec:experiments} use embeddings with $l_f$ and $l_a$ closely together.}
\fi

\section{Steerable group convolutions}
\label{sec:groupconvs}

\paragraph{Steerable functions.}
Through the equivalence of steerable vectors and spherical functions via the spherical Fourier transform, it is clear that our definition of \textit{steerable vectors} coincides with the more common definition of \textit{steerable functions}, as commonly used in computer vision \citep{freeman1991design}. A function $f$ is called steerable~\citep{helor1996steerable} under a transformation group $G$ if any transformation
$g \in G$ on $f$ can be written as a linear combination of a fixed, finite set of $n$ basis functions $\{ \phi_i\}$:
\begin{align}
\label{eq:steerablefunction}
    \mathcal{L}_g f = \sum_{i=1}^n \alpha_i(g) \phi_i = \boldsymbol{\alpha}^T(g) \mathbf{\Phi} \ ,
\end{align}
in which $\mathcal{L}_g$ is the left regular representation of the group $G$ that performs the transformation on $f$ (see section \ref{sec:equivariance}). In case of functions in a spherical harmonic basis, which we denote with $f_{\mathbf{a}}:\sum_{l\geq 0} \sum_{m=-l}^l a_m^{(l)} Y_m^l(\mathbf{n})$, it directly follows from Eq.~(\ref{eq:shequivariance}) that they are steerable via
\begin{equation}
\label{eq:shsteerable}
\mathcal{L}_g f_{\mathbf{a}}(\mathbf{n})
= f_{\mathbf{D}(g)\mathbf{a}}(\mathbf{n}) \ .
\end{equation}
In terms of the steerable function definition in Eq.~(\ref{eq:steerablefunction}), this means that $\boldsymbol{\alpha}(g) = \mathbf{D}(g) \mathbf{a}$ and $\boldsymbol{\Phi} = (Y_0^{(0)}, Y_{-1}^{(1)}, Y_0^{(1)}, \dots)^T$, i.e. the set of basis functions flattened into a vector.

\paragraph{A translation-rotation template machting motivation for (steerable) group convolutions.}
The notion of steerability becomes particularly clear when viewed in the computer vision context \citep{freeman1991design}, where one may be interested in detecting visual features under arbitrary rotations. Let us consider the case of 3D cross-correlations of a kernel ${k}:\mathbb{R}^3 \rightarrow \mathbb{R}$ with an input feature map ${f}:\mathbb{R}^3 \rightarrow \mathbb{R}$:
\begin{equation}
\label{eq:normalconvolution}
{f}'(\mathbf{x}) = ({k} \star {f})(\mathbf{x}) = \int_{\mathbb{R}^3} {k}(\mathbf{x}' - \mathbf{x}) {f}(\mathbf{x}') {\rm d}\mathbf{x}' \ .
\end{equation}
We think of such a kernel ${k}$ as describing a particular visual pattern. In many applications, we want to 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
\begin{equation}
\label{eq:liftingconv}
{f}'(\mathbf{x}, \mathbf{R}) = \int_{\mathbb{R}^3} {k}(\mathbf{R}^{-1}(\mathbf{x}' - \mathbf{x})) {f}(\mathbf{x}') {\rm d}\mathbf{x}' \ .
\end{equation}
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 \citep{cohen2016group}, the above is what one usually calls a lifting group convolution~\citep{bekkers2019b} (feature maps are lifted from $\mathbb{R}^3$ to the group SE$(3)$).

\paragraph{Group convolutions.}
In \textit{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 $\mathbb{R}^3$ and group product respectively via
\begin{align}
{f}'(g) &= \int_{\mathbb{R}^3} {k}(g^{-1} \cdot \mathbf{x}') {f}(\mathbf{x}') {\rm d}\mathbf{x}' \ , \label{eq:liftingconvapp} \\
{f}'(g) &= \int_{SE(3)} {k}(g^{-1} \cdot g') {f}(g') {\rm d}g' \label{eq:groupconvapp} \ ,
\end{align}
where $g \in \text{SE}(3)$, ${\rm d}g$ the Haar measure on the group and where $\cdot$ in \eqref{eq:liftingconvapp} and \eqref{eq:groupconvapp} respectively denote the group action on $\mathbb{R}^3$ and group product of $\text{SE}(3)$ (cf. Section \ref{sec:groupdefinition}). Note that Eq.~\eqref{eq:liftingconvapp} is exactly the same as  Eq.~\eqref{eq:liftingconv} but in different notation.

The lifting group convolution thus generates a function on the joint space of positions $\mathbb{R}^3$ 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 $\mathbf{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 \citep{weiler2018}, B-splines \citep{bekkers2019b} or they can be parametrised via MLPs \citep{finzi2020generalizing}. Alternatively, 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, as in \citep{worrall2018cubenet,winkels20183d}. An advantage of regular group convolution methods is that normal point-wise activation functions can directly be applied to the feature maps; a 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 is no need for such a discretisation at all and one can obtain the response of the convolution at any rotation $\mathbf{R}$ after convolving with the basis functions. This works as follows.

\paragraph{Steerable template matching.}
Suppose a 3D convolution kernel that is expanded in a spherical harmonic basis up to degree $L$ as follows
\begin{equation}
\label{eq:shkernelapp}
k(\mathbf{x}) = k_{\tilde{\mathbf{c}}(\lVert \mathbf{x} \rVert)}(\mathbf{x}) := \sum_l^L \sum_{m=-l}^l c^{(l)}_m(\lVert \mathbf{x} \rVert) Y^{(l)}_m\left(\tfrac{\mathbf{x}}{\lVert \mathbf{x} \rVert}\right) \ ,
\end{equation}
with $\tilde{\mathbf{c}} = (c_0^{(0)},c_{-1}^{(1)},c_0^{(1)} \dots)^T$ the vector of basis coefficients that can depend on $\lVert \mathbf{x} \rVert$. The coefficients can e.g. be parametrised with an MLP that takes as input $\lVert \mathbf{x} \rVert$ and returns the coefficient vector. We note that such coefficients are then O($3$) invariant, i.e., $\forall_{\mathbf{R} \in \text{O(3)}}:\,\tilde{\mathbf{c}}(\lVert\mathbf{R} \mathbf{x}\rVert)=\tilde{\mathbf{c}}(\lVert\mathbf{x}\rVert)$. Furthermore, we labelled the vector with a ``$\;\tilde{\,}\;$'' to indicate it is a steerable vector as it represents the coefficients relative to a spherical harmonic basis. It then follows (from Eq.~\eqref{eq:shsteerable}) that the kernel is steerable via
\begin{equation*}
k(\mathbf{R}^{-1} \mathbf{x}) = k_{\mathbf{D}(\mathbf{R})\tilde{\mathbf{c}}(\lVert \mathbf{x} \rVert)}(\mathbf{x}) \ ,
\end{equation*}
i.e., via a transformation of the coefficient vector $\tilde{\mathbf{c}}$ by it O(3) representation $\mathbf{D}(\mathbf{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 a vector of responses $\tilde{\mathbf{f}}(\mathbf{x}) = (f^{(0)}_0(\mathbf{x}), f^{(1)}_{-1}(\mathbf{x}), f^{(1)}_0(\mathbf{x}) , \dots )^T$ of which the elements are given by
\begin{align}
\label{eq:steerableconv}
f^{(l)}_m(\mathbf{x}) &= ((c^{l}_m Y^{(l)}_m) \star f^{in})(\mathbf{x}) \\
&= \int_{\mathbb{R}^3}  c^{(l)}_m(\lVert \mathbf{x} \rVert) Y^{(l)}_m(\mathbf{x}' - \mathbf{x}) f(\mathbf{x}) {\rm d}\mathbf{x} \ . \nonumber
\end{align}

Then the result with the convolution kernel (Eq.~(\ref{eq:shkernelapp})) is obtained simply by a sum over the vector components which we denote with $\underset{l,m}{\operatorname{SumReduce}}$ as follows
\begin{equation*}
f'(\mathbf{x}) = \underset{l,m}{\operatorname{SumReduce}}(\tilde{\mathbf{f}}'(\mathbf{x})) := \sum_{l}^L \sum_{m=-l}^l f'^{(l)}_m(\mathbf{x}) \ .
\end{equation*}
If one were to be interested in the rotated filter response at $\mathbf{x}$ one can first rotate the steerable vector $\tilde{\mathbf{f}}'(\mathbf{x})$ via the matrix representation $\mathbf{D}(\mathbf{R})$ and only then do the reduction. I.e., once the convolutions of \eqref{eq:steerableconv} are done the lifting group convolution result is directly obtained via
\begin{equation}
\label{eq:sumreduce}
f'(\mathbf{x}, \mathbf{R}) = \underset{l,m}{\operatorname{SumReduce}}(\mathbf{D}(\mathbf{R})\tilde{\mathbf{f}}'(\mathbf{x})) \ .
\end{equation}
A convolution with a kernel expressed in a spherical harmonics basis thus generates a signal on O(3) at each point $\mathbf{x}$.

When only considering the $m=0$ component for each order $l$ in the kernel parametrisation, the kernels have an axial symmetry in them due to which the convolution results in a point-wise spherical signal. This is in fact the situation that we consider in our weighted CG products (Eq.~\ref{eq:CG_product}) where 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 steerable group convolutions on the full space $\mathbb{R}^3 \rtimes SO(3)$, this would require working with CG products with weights indexed by both $m$ and $m'$ \citep{weiler2018}. Since in this work we choose to with steerable vector spaces $V=m_0 V_0 \oplus m_1 V_1 \oplus \dots$ with arbitrary multiplicities $m_l$ for each type, we implement what is called generalised steerable convolutions, of which the SO(3) convolutions are a special case with $m_l=(2l+1)$ \citep{kondor2018clebsch}.

\paragraph{Steerable group convolutions.}
When working with steerable convolution kernels one does not have to work with a grid on O(3). There are reasons to avoid working with a grid on O($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 O($3$). Steerable methods provide a way to build neural networks entirely independent of any sampling on O($3$), since, as we have seen, steerable convolutions directly result in functions on the entire group O($3$) via steerable vectors. That is, \emph{at each location $\mathbf{x}$ we have a steerable vector $\tilde{\mathbf{f}}(\mathbf{x})$ which represents a function on O($3$) via the inverse Fourier transform given in Eq.~\eqref{eq:O3FourierInv}}. In fact, the sum reduction in Eq.~\eqref{eq:sumreduce} 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 \eqref{eq:steerableconv}, is an operator 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 $\mathbf{x}$, i.e. a type-0 steerable vector, and the output was a steerable vector field containing steerable vectors up to type $L$. The transition from type-0 vector field to a type-$L$ vector field happened via tensor products with steerable vectors of spherical harmonics, and this tensor product was parametrised by the coefficients $\tilde{c}$. Suppose a convolution of a single type-$l_1$ steerable input field with a convolution kernel that is expanded in a spherical harmonic basis of only type $l_2$ via $k(\mathbf{x})=\sum_{m=-l_2}^{l_2} w(\lVert \mathbf{x}\rVert) Y^{(l)}_m\left(\tfrac{\mathbf{x}}{\lVert \mathbf{x} \rVert}\right)$. The kernel can then be represented as a steerable vector field in itself as $\tilde{\mathbf{k}}(\mathbf{x}) = w(\lVert \mathbf{x}\rVert) \tilde{\mathbf{a}}(\mathbf{x})$, with $\tilde{\mathbf{a}}$ the type-$l_2$
spherical harmonic embedding given in Eq.~5 of the main paper.
Such a steerable convolution maps from input feature maps $\tilde{\mathbf{f}}:\mathbb{R}^3 \rightarrow V_{l_1}$ using a convolution kernel $\tilde{\mathbf{k}}:\mathbb{R}^3 \rightarrow V_{l_2}$ to an output $\tilde{\mathbf{f}}':\mathbb{R}^3 \rightarrow V_{l}$ via
\begin{equation*}
\tilde{\mathbf{f}}'(\mathbf{x}) = \int_{\mathbb{R}^3}   \tilde{\mathbf{f}}(\mathbf{x}) \otimes_{cg}^{w(\lVert \mathbf{x} \rVert)}\tilde{\mathbf{a}}(\mathbf{x}'-\mathbf{x}) {\rm d}\mathbf{x}' \ .
\end{equation*}
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.

\paragraph{Steerable and regular group convolutions.}
In conclusion, both steerable and regular group convolutions 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 a particular grid which is e.g. the Cartesian product of a regular 3D grid with a particular discretisation of O($3$). In regular group convolutions, we directly index the responses with $(\mathbf{x},\mathbf{R}) \in \text{E}(3) \mapsto \mathbf{f}(\mathbf{x},\mathbf{R})$. In steerable convolutions the filter responses $\mathbf{x} \rightarrow \tilde{\mathbf{f}}(\mathbf{x})$ are stored at each spatial grid point $\mathbf{x}$ in steerable vectors $\tilde{\mathbf{f}}(\mathbf{x})$ from which functions on the full group O(3) can be obtained via an inverse Fourier transform\footnote{This Fourier transform in fact enables working with classic point-wise activation functions that could be applied point-wise to each location in a spherical or O(3) signal as in \citep{cohen2018spherical}. It 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 a grid via the inverse Fourier transform, apply the activation function, and transform back into the steerable basis. This is the approach taken in \citep{cohen2018spherical}. In this paper we work entirely in the steerable domain (O($3$) Fourier space) and work with gated non-linearities \citep{weiler2018}.}. As such one recovers the regular representation via $(\mathbf{x},\mathbf{R}) \mapsto \mathcal{F}^{-1}[\tilde{\mathbf{f}}(\mathbf{x})](\mathbf{R})$. 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 are exactly equivariant to all transformations in O($3$) via steerable representations of O($3$).

\iffalse
\section{Non-linear steerable group convolutions}
\erik{to write: I'm going to focus now on the main paper}

\section{Stuff I'm going to reorganise:}
\subsection{Equivariant message passing}
\subsection*{\color{red}-SOME THINGS MIGHT GO TO THE APPENDIX-}
In order to make the message passing framework equivariant to translations, rotations, and reflections we will replace the MLPs that define the learnable functions $\phi_m$ and $\phi_f$ by \textit{equivariant MLPs}, that are steered by directional information such as relative positions between node pairs. The notion of equivariance will soon be formalised, but for now we can think of equivariant MLPs as just regular MLPs that are constrained to only represent functions that behave predictably under certain transformations. For example, a convolutional neural network is an MLP that is constrained to only represent translation equivariant functions, since the linear layer between feature maps is constrained to be a convolution operator. This notion of equivariance can be extended to transformations beyond just translations. In our work, we focus on MLPs that are constrained to be equivariant to rotations, translations and reflections of the data. Since our approach builds upon core concepts from group theory we first provide some formal definitions in order to explain and specify our method.

\paragraph{Equivariance}.
Given a group $G$ and a set of transformations $T_g: \mathcal{X} \rightarrow \mathcal{X}$ on $\mathcal{X}$,
a function $\phi: \mathcal{X} \rightarrow \mathcal{Y}$ is called equivariant with respect to $G$ if there exists a transformation $S_g: \mathcal{Y}\rightarrow\mathcal{Y}$ on the output space $\mathcal{Y}$ such that
\begin{align}
    S_g [\phi(x)]=\phi(T_g[x]) \ , \ \text{for all } g\in G, x \in \mathcal{X} \ .\label{eq:equivariance}
\end{align}

One way to interpret equation~\eqref{eq:equivariance} is therefore that the operators $S_g$ and $T_g$ describe the same transformation but in different spaces. Invariant functions are a special case of equivariant functions, where $S_g = \mathcal{I}_{\mathcal{Y}}$ for all $g \in G$. Equivariance provides us with a principled way of encoding symmetries into neural networks and improves generalisation. For example, the output of a segmentation task is equivariant to translations, as translating the image yields a similarly translated segmentation mask.

\paragraph{Group representations.}
In the above we defined the group product, which defines how to transform elements in the group, but also showed how the groups E($3$) and O($3$) act on a three dimensional vector space. We usually think of the groups E($3$) and O($3$) as groups that describe transformations on $\mathbb{R}^3$, but these groups are not restricted to transformations on $\mathbb{R}^3$ be can generally act on arbitrary vector spaces via representations. A representation is an invertible linear transformation $\rho(g):{V} \rightarrow {V}$ parametrised by group elements $g \in G$ that acts on some vector space ${V}$, and which follows the group structure (it is a group homomorphism) via
$$
\rho(g)\rho(h)v = \rho(g \cdot h)v \ ,
$$
with $v \in V$. When the representations transform finite dimensional vectors $v \in V = \mathbb{R}^d$, the representations are $\mathbb{R}^d \times \mathbb{R}^d$ dimensional matrices. For example, the representation of O($3$) on $\mathbb{R}^3$ are directly the group elements themselves, i.e., $\rho(\mathbf{R}) \mathbf{x} = \mathbf{R} \mathbf{x}$.

\paragraph{Wigner-D matrices.}
For O($3$) there exists a collection of so-called irreducible representations, indexed with their degree $l \geq 0$, which act on vector spaces of dimension $2l + 1$. These representations are called Wigner-D matrices and we denote them with $\mathbf{D}^{(l)}(g)$.

\paragraph{O($3$) steerable vector spaces.}
We call a function $f$ steerable~\citep{freeman1991design,helor1996steerable} under a transformation group $G$ if any transformation
$g \in G$ on $f$ can be written as a linear combination of a fixed, finite set of $n$ basis functions $\{ \phi_i\}$:
\begin{align}
    g f = \sum_{i=1}^n \alpha_i \phi_i = \boldsymbol{\alpha}^T \mathbf{\Phi} \ .
\end{align}
The functions $\alpha_i$ are the steering functions of $f$ and are associated with the basis $\{ \phi_i\}$.
For O($3$) steerable vector spaces, vectors are represented as functions on the sphere in the basis of $Y_m^{(l)}$, which are the \textit{spherical harmonics}, $Y_m^{(l)}: \mathcal{S}^2 \rightarrow \mathcal{C}$
Spherical harmonics are square-integrable, complex-valued functions on the sphere $\mathcal{S}^2$, $m=-l, -l+1,\ldots,l$.
Wigner-D matrices $\mathbf{D}^{(l)}(g)$ are the steering functions.
In our case, we call a vector space to be steerable of type $l$ if its elements transform via the $l$-th degree Wigner-D matrices $\mathbf{D}^{(l)}$, and we denote the space with $V_l$. E.g., we say a type-3 vector $\mathbf{h} \in V_3$ can be transformed by $g \in \text{O}(3)$ via $\mathbf{h} \mapsto \mathbf{D}^{(3)}(g)\mathbf{h}$. Explicit forms of the Wigner-D matrices can be found in~\citet{sakurai2017}. Notably, type-0 vectors $h \in V_0 = \mathbb{R}$ are just scalars which are invariant to transformations $g \in \text{O}(3)$, i.e., $\mathbf{D}_0(g) h = h$. Type-1 features are vectors $\mathbf{h} \in \mathbb{R}^3$ which transform directly via the matrix representation of the group, i.e., $\mathbf{D}_1(g) \mathbf{h} = \mathbf{R} \mathbf{h}$.

\paragraph{Equivariance, Steerable MLPs.}
It is our aim to construct multi-layer perceptrons that are equivariant to $\text{O}(3)$ transformations. We will do so by constructing functions $\phi: V_k \rightarrow V_{l}$ which are equivariant to transformations in $\text{O}(3)$ via the Wigner-D matrices. I.e., for every input vector $\mathbf{h} \in V_k$ and every $g \in \text{O}(3)$ we want the following to hold
\begin{equation}
\phi(\mathbf{D}^{(k)}(g) \mathbf{h}) = \mathbf{D}^{(l)} \phi( \mathbf{h} ) \ .
\end{equation}
We will call an equivariant multi-layer perceptron with steerable input and output spaces a \emph{steerable MLP}.
Note that when the function $\phi: V_k \rightarrow V_0$ maps  a type-k input vector to a scalar (type-0 vector), the function is invariant
since
\begin{equation}
\phi(\mathbf{D}^{(k)} \mathbf{h}) = \phi(\mathbf{h}) \ .
\end{equation}

\paragraph{Mapping between O(3) steerable vector spaces of different types: The Clebsch-Gordan product.}
The Clebsch-Gordan (CG) tensor product $\otimes_{cg}: V_{l_1} \times V_{l_2} \rightarrow V_{l}$ is a bilinear operator that combines two O($3$) steerable input vectors of types $l_1$ and $l_2$ and returns another steerable vector of type $l$. Let $\mathbf{h}^{(l)} \in V_l = \mathbb{R}^{2l+1}$ denote a steerable vector of type $l$, and $h^{(l)}_m$ its components with $m=-l,-l+1,\dots,l$. The basis functions of $h^{(l)}_m$ are the spherical harmonics $Y^{(l)}_m$.
The CG tensor product is given by
\begin{align}
(\mathbf{h}^{(l_1)} \otimes_{cg}^{w} \mathbf{h}^{(l_2)})^{(l)}_{m} = \sum_{m_1=-l_1}^{l_1} \sum_{m_2=-l_2}^{l_2} {\color{red} w \;} C^{(l,m)}_{(l_1,m_1)(l_2,m_2)} h^{(l_1)}_{m_1} h^{(l_2)}_{m_2} \ ,
\label{Aeq:CG_product}
\end{align}
\iffalse
$$
(\mathbf{h}^{(l_1)} \otimes_{cg}^{\mathbf{w}} \mathbf{h}^{(l_2)})^{(l)}_{m} = \sum_{m_1=-l_1}^{l_1} \sum_{m_2=-l_2}^{l_2} {\color{red}w_{m_1m_2}^{(l)}}(\lVert \mathbf{x}_j - \mathbf{x}_i \rVert) C^{(l,m)}_{(l_1,m_1)(l_2,m_2)} h^{(l_1)}_{m_1} h^{(l_2)}_{m_2},
$$
$$
(\mathbf{h}^{(l_1)} \otimes_{cg}^{{\rm MLP}(\lVert \mathbf{x}_j - \mathbf{x}_i \rVert)} \mathbf{h}^{(l_2)})^{(l)}_{m} = \sum_{m_1=-l_1}^{l_1} \sum_{m_2=-l_2}^{l_2} {\color{red}{\rm MLP}(\lVert \mathbf{x}_j - \mathbf{x}_i \rVert)} C^{(l,m)}_{(l_1,m_1)(l_2,m_2)} h^{(l_1)}_{m_1} h^{(l_2)}_{m_2},
$$
\fi
\textcolor{red}{in which $w$ is a learnable parameter that scales the product}, and
$C^{(l,m)}_{(l_1,m_1)(l_2,m_2)}$ are the Clebsch-Gordan coefficients, which ensure hat the resulting vector is steerable of type $l$. The CG product describes a sparse tensor product, as generally many of the $C^{(l,m)}_{(l_1,m_1)(l_2,m_2)}$ components are zero. Most notably, $C^{(l,m)}_{(l_1,m_1)(l_2,m_2)}=0$ for any $m > |l_1 + l_2|$ and $m < |l_1 - l_2|$. As such, one can for example not generate a type-2 vector out of a type-0 (a scalar) and a type-1 vector, which is however possible for two type-1 vectors.

{\color{red}Example of non-zero CG coefficients for cross product?}

Examples of instances of the CG product are:
\begin{itemize}
    \item The product of two scalars is a CG product which takes as input two type-0 ``vectors'' to generate another scalar ($l_1 = 0, l_2 = 0, l = 0$).
    \item The scalar product, which takes as input a scalar and a type-1 vector to generate a type-1 vector ($l_1=0, l_2=1, l=1$).
    \item The dot product, which takes two type-1 vectors as input to produce a scalar ($l_1=1, l_2=1, l=0$).
    \item The cross product, which takes two type-1 vectors as input to produce another type-1 vector ($l_1=1, l_2=1, l=1$).
    \item \textcolor{red}{The outer product, which takes two type-1 vectors as input to produce a type-2 vector ($l_1=1, l_2=1, l=2$).}
\end{itemize}
{\color{red} CG product between mixed steerable vector spaces}.

\paragraph{Steerable MLPs.}
In order to keep notation as close as possible to standard notation used in MLPs, we refer to the CG product which is used with a fixed vector $\mathbf{a}$ in one of its inputs as a \emph{steerable linear layer conditioned on $\mathbf{a}$}. We denote this as follows
\begin{align}
\mathbf{W}_{\mathbf{a}} \mathbf{h} := \mathbf{h} \otimes_{cg}^\mathbf{w} \mathbf{a} \ .
\label{Aeq:steerable_layer}
\end{align}
We can now define steerable MLPs by stacking layers and applying non-linearities via
$$
\mathbf{h}^{i} = \sigma(\mathbf{W}_{\mathbf{a}}^{i}(\mathbf{h}^{i-1})) \ ,
$$
in which $\sigma$ is an equivariant activation function. We propose gated non-linearities~\citep{weiler2018} for $l\geq0$ terms of the CG product, and swish activation function~\citep{ramachandran2017searching} for $l=0$ terms.

\subsection{A note on spherical harmonics}

\paragraph{Visualizing $O(3)$ steerable vectors}
{\color{red} Erik: insert background on representing functions with spherical harmonics, which justifies our visualizations}

\paragraph{Steerable group convolutions}
{\color{red} Erik: insert background on steerable group convolutions. In essence we are doing non-linear steerable group convolutions}
\fi

\section{Experiments}\label{sec:appendix_experiments}

\paragraph{Compared methods.}
We compared against methods discussed in Sec.~3 of the main paper, namely
NMP~\citep{gilmer2017neural}, SchNet~\citep{schutt2017schnet},
Cormorant~\citep{anderson2019covariant}, L1Net~\citep{miller2020relevance}, LieConv~\citep{finzi2020generalizing}, TFN~\citep{thomas2018tensor}, SE(3)-transformer~\citep{fuchs2020se}, and EGNN~\citep{satorras2021n}. We additionally compare to CGCNN~\citep{xie2018crystal},
PaiNN~\citep{schutt2021equivariant},
NequIP~\citep{batzner2021se},
Dimenet++~\citep{klicpera2020dimenet_plusplus} and SphereNet~\citep{liu2021spherical}.

\paragraph{Implementation details for N-body dataset.}
SEGNN architectures, point conv methods (SE$_{\text{linear}}$), and steearable non-linear point conv methods (SE$_{\text{non-linear}}$)
consist of three parts (sequentially applied):
\begin{enumerate}
    \item Embedding network: Input $\rightarrow$ \{$\text{CG}_{\tilde{\mathbf{a}}_i}$ $\rightarrow$ SwishGate $\rightarrow$ $\text{CG}_{\tilde{\mathbf{a}}_i}$\}, wherein $\text{CG}_{\tilde{\mathbf{a}}_i}$ denotes a steerable linear layer conditioned on node attributes $\tilde{\mathbf{a}}_i$ and which are applied per node.
    \item Four message passing layers as described in Sec.~\ref{sec:segnns} for SEGNN architectures and 7 convolution layers for SE$_{\text{linear}}$ and SE$_{\text{non-linear}}$ ablations.
    \item A prediction network: \{$\text{CG}^0_{\tilde{\mathbf{a}}_i}$ $\rightarrow$ Swish $\rightarrow$ Linear $\rightarrow$ MeanPool $\rightarrow$ Linear $\rightarrow$ Swish $\rightarrow$ Linear\}, in which $\text{CG}^0_{\tilde{\mathbf{a}}_i}$ denotes a steerable linear layer conditioned on node attributes $\tilde{\mathbf{a}}_i$ and which maps to a vector of 64 scalar (type-0) features. The remaining layers are regular linear layers, as in \citet{satorras2021n}. $\text{CG}^0_{\tilde{\mathbf{a}}_i}$ and Linear are applied per node.

    All steerable architectures are designed such that the parameter budget at  $l_f=1$ and $l_a=1$ matches that of the tested EGNN implementation.
    We optimise models using the Adam optimiser \citep{kingma2014adam} with learning rate 1e-4, weight decay 1e-8 and batch size 100 for 10000 epochs and minimise the mean absolute error. .
\end{enumerate}
The SwishGate refers to a gated non-linearity \citep{weiler2018} when $l > 0$, and a swish activation \citep{ramachandran2017searching} otherwise. A message network $\phi_m$ consists of \{$\text{CG}_{\tilde{\mathbf{a}}_{ij}}$ $\rightarrow$ SwishGate $\rightarrow$ $\text{CG}_{\tilde{\mathbf{a}}_{ij}}$ $\rightarrow$ SwishGate\}, while the update network $\phi_f$ consists of \{$\text{CG}_{\tilde{\mathbf{a}}_{i}}$ $\rightarrow$ SwishGate $\rightarrow$ $\text{CG}_{\tilde{\mathbf{a}}_{i}}$  $\rightarrow$ InstanceNorm\}, with an instance normalisation \citep{ulyanov2016instance} implementation that is compatible with steerable vectors, based on the BatchNorm implementation of the \texttt{e3nn} library \citep{geiger2021}. We use skip-connections in the message passing layers,
adapting Eq.~(\ref{eq:segnn3})
to yield $\tilde{\mathbf{f}}'_i = \tilde{\mathbf{f}}_i + \phi_f(\tilde{\mathbf{f}}_i, \sum_j \tilde{\mathbf{m}}_{ij}, \tilde{\mathbf{a}}_i )$.

The convolution layers which are used instead of the message passing layers for SE$_{\text{linear}}$ and SE$_{\text{non-linear}}$ ablations consist of:
\begin{itemize}
    \item  SE$_{\text{linear}}$: \{$\text{CG}_{\tilde{\mathbf{a}}_{ij}}$
    $\rightarrow$ $\sum_{j \in \mathcal{N}(i)}$ $\rightarrow$ SwishGate\ $\rightarrow$ InstanceNorm \}
     \item  SE$_{\text{non-linear}}$: \{$\text{CG}_{\tilde{\mathbf{a}}_{ij}}$
    $\rightarrow$ SwishGate $\rightarrow$ $\text{CG}_{\tilde{\mathbf{a}}_{ij}}$
    $\rightarrow$ $\sum_{j \in \mathcal{N}(i)}$ $\rightarrow$ SwishGate\ $\rightarrow$ InstanceNorm\}

\end{itemize}

\begin{table}[h]
\centering
\tiny
\caption{A full ablation study on the performance and runtime of different  (maximum) orders of steerable features ($l_f$) and attributes ($l_a$). The ablation study includes steerable linear convolutions, steerable non-linear convolutions, as well as SEGNN approaches.  SE$_{\text{non-linear}}$++ uses the structural information (velocity) in the two embedding and read-out layers. }
\begin{adjustbox}{width=0.49\textwidth}
  \begin{tabular}{l c c}
  \toprule
    Method & MSE & Time [s] \\
    \midrule
    SE$_{\text{linear}}$ ($l_f=0, l_a=0$) & .0344 & .0032 \\
    SE$_{\text{linear}}$ ($l_f=0, l_a=1$) & .0352 & .0036 \\
    SE$_{\text{linear}}$ ($l_f=1, l_a=1$) & .0130 & .031 \\
    SE$_{\text{linear}}$ ($l_f=1, l_a=2$) & .0121 & .035 \\
    SE$_{\text{linear}}$ ($l_f=2, l_a=2$) & .0116 & .065 \\
    SE$_{\text{linear}}$ ($l_f=2, l_a=3$) & .0121 & .075 \\
    \midrule
    SE$_{\text{non-linear}}$ ($l_f=0, l_a=0$) & .0382 & .0038 \\
    SE$_{\text{non-linear}}$ ($l_f=0, l_a=1$) & .0388 & .0040 \\
    SE$_{\text{non-linear}}$ ($l_f=1, l_a=1$) & .0061 & .030 \\
    SE$_{\text{non-linear}}$ ($l_f=1, l_a=2$) & .0060 & .034 \\
    SE$_{\text{non-linear}}$ ($l_f=2, l_a=2$) & .0061 & .065 \\
    SE$_{\text{non-linear}}$ ($l_f=2, l_a=3$) & .0060 & .074 \\
    \midrule
    SE$_{\text{non-linear}}$++ ($l_f=0, l_a=0$) & .0392 & .0038 \\
    SE$_{\text{non-linear}}$++ ($l_f=0, l_a=1$) & .0390 & .0041 \\
    SE$_{\text{non-linear}}$++ ($l_f=1, l_a=1$) & .0057 & .031 \\
    SE$_{\text{non-linear}}$++ ($l_f=1, l_a=2$) & .0057 & .036 \\
    SE$_{\text{non-linear}}$++ ($l_f=2, l_a=2$) & .0058 & .067 \\
    SE$_{\text{non-linear}}$++ ($l_f=2, l_a=3$) & .0062 & .078 \\
    \midrule
    SEGNN$_\text{G}$ ($l_f=0, l_a=0$) & .0099 & .0036 \\
    SEGNN$_\text{G}$ ($l_f=0, l_a=1$) & .0098 & .0050 \\
    SEGNN$_\text{G}$ ($l_f=1, l_a=1$) & .0056 & .025 \\
    SEGNN$_\text{G}$ ($l_f=1, l_a=2$) & .0058 & .031 \\
    SEGNN$_\text{G}$ ($l_f=2, l_a=2$) & .0060 & .061 \\
    SEGNN$_\text{G}$ ($l_f=2, l_a=3$) & .0062 & .068 \\
    \midrule
    SEGNN$_\text{G+P}$ ($l_f=0, l_a=0$) & .0102 & .0037 \\
    SEGNN$_\text{G+P}$ ($l_f=0, l_a=1$) & .0096 & .0052 \\
    SEGNN$_\text{G+P}$ ($l_f=1, l_a=1$) & .0043 & .026 \\
    SEGNN$_\text{G+P}$ ($l_f=1, l_a=2$) & .0044 & .032 \\
    SEGNN$_\text{G+P}$ ($l_f=2, l_a=2$) & .0041 & .063 \\
    SEGNN$_\text{G+P}$ ($l_f=2, l_a=3$) & .0041 & .071 \\
    \bottomrule
  \end{tabular}
\end{adjustbox}
\label{tab:app_nbody_ablation}
\end{table}

\paragraph{Implementation details for QM9.}
Our QM9 hyperparameters strongly resemble those found in \citep{satorras2021n}, whose architecture we took as a starting point for our own. The network consists of three parts (sequentially applied):
\begin{enumerate}
    \item Embedding network: Input $\rightarrow$ \{$\text{CG}_{\tilde{\mathbf{a}}_i}$ $\rightarrow$ SwishGate $\rightarrow$ $\text{CG}_{\tilde{\mathbf{a}}_i}$\}, wherein $\text{CG}_{\tilde{\mathbf{a}}_i}$ denotes a steerable linear layer conditioned on node attributes $\tilde{\mathbf{a}}_i$ and which are applied per node.
    \item Seven message passing layers as described in Sec.~\ref{sec:segnns}.
    \item A prediction network: \{$\text{CG}^0_{\tilde{\mathbf{a}}_i}$ $\rightarrow$ Swish $\rightarrow$ Linear $\rightarrow$ MeanPool $\rightarrow$ Linear $\rightarrow$ Swish $\rightarrow$ Linear\}, in which $\text{CG}^0_{\tilde{\mathbf{a}}_i}$ denotes a steerable linear layer conditioned on node attributes $\tilde{\mathbf{a}}_i$ and which maps to a vector of 128 scalar (type-0) features. The remaining layers are regular linear layers, as in \citet{satorras2021n}. $\text{CG}^0_{\tilde{\mathbf{a}}_i}$ and Linear are applied per node.
\end{enumerate}

Message network, update network and skip-connections are used as described for the N-body task. In contrast to the other tasks, the hidden irreps correspond to $N$ copies for every order, selected to correspond to a 128-dimensional linear layer in terms of number of parameters.

We optimise models using the Adam optimiser \citep{kingma2014adam} with learning rate 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 a 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.

Figure~\ref{fig:radiusmessages} shows the number of messages as a function of cutoff radius for the train partition of QM9. It shows that a network with cutoff radius of 2\r{A} sends $\sim 6\times$ fewer messages on average in every layer compared to a cutoff radius of 5\r{A}, with far smaller variance.

\begin{figure}
    \centering
    \includegraphics[width=0.8\textwidth]{Figures/radius.png}
    \caption{Mean number of messages as a 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 a minimum cutoff radius.}
    \label{fig:radiusmessages}
\end{figure}

\paragraph{Implementation details for the Open Catalyst Project OC20 dataset.}

Our OCP hyperparameters resemble those from the N-body task, but with an overall larger network architecture. The network consists of three parts (sequentially applied):
\begin{enumerate}
    \item Embedding network: Input $\rightarrow$ \{$\text{CG}_{\tilde{\mathbf{a}}_i}$ $\rightarrow$ SwishGate $\rightarrow$ $\text{CG}_{\tilde{\mathbf{a}}_i}$\}, wherein $\text{CG}_{\tilde{\mathbf{a}}_i}$ denotes a steerable linear layer conditioned on node attributes $\tilde{\mathbf{a}}_i$ and which are applied per node.
    \item Seven message passing layers as described in Sec.~\ref{sec:segnns}.
    \item A prediction network: \{$\text{CG}^0_{\tilde{\mathbf{a}}_i}$ $\rightarrow$ Swish $\rightarrow$ Linear $\rightarrow$ MeanPool $\rightarrow$ Linear $\rightarrow$ Swish $\rightarrow$ Linear\}, in which $\text{CG}^0_{\tilde{\mathbf{a}}_i}$ denotes a steerable linear layer conditioned on node attributes $\tilde{\mathbf{a}}_i$ and which maps to a vector of 256 scalar (type-0) features. The remaining layers are regular linear layers, as in \citep{satorras2021n}. $\text{CG}^0_{\tilde{\mathbf{a}}_i}$ and Linear are applied per node.
\end{enumerate}
Message network, update network and skip-connections are used as described for the QM9 tasks.
The input is a 256-dimensional one-hot representation of atom types, all hidden vectors are 256-dimensional or their steerable ``weight balanced'' equivalent
(cf Sec.~\ref{sec:experiments}),
where we divide into equally large type-$l$ sub-vector spaces.

We optimise models using the AdamW optimiser~\citep{loshchilov2017decoupled} with learning rate 1e-4, weight decay 1e-8 and batch size 8 for 20 epochs and minimise the mean absolute error. The learning rate is reduced by a factor of ten at 50$\%$ and 75$\%$ of training.

\paragraph{Ablation study for the Open Catalyst Project OC20 dataset, speed comparison.}
Table~\ref{tab:app_nbody_ablation} shows SEGNN ablation of the performance when using different (maximum) orders of steerable features $l_f$ and attributes $l_a$, ablation for different cutoff radii of 5\r{A} and 3\r{A}, as well as runtime comparisons. All results are obtained on the validation set.

\begin{table*}[!htb]
\centering
\caption{Comparison on the validation sets of the OC20 IS2RE tasks. Compared is the performance for the usage of different (maximum) orders of steerable features ($l_f$) and attributes ($l_a$), as well as the performance for cutoff radii of 5\r{A} and 3\r{A} on the OC20 dataset. Runtime numbers are reported for one forward pass and are measured for a batch of size 8 on a Nvidia Tesla V100 GPU. Higher orders of $l_f$ and $l_a$ in general do not improve results. The Dimenet++ runtime is given as reference.}
\begin{adjustbox}{width=1\textwidth}
\begin{tabular}{l c c c c | c c c c | c }
\toprule
\\[-0.4em]
& \multicolumn{4}{c |}{Energy MAE [eV] $\downarrow$} & \multicolumn{4}{c}{EwT $\uparrow$} & Time [s]  \\
\cline{2-9}
\\[-0.4em]
    Model  & ID & OOD Ads & OOD Cat & OOD Both & ID & OOD Ads & OOD Cat & OOD Both & - \\
\midrule
SEGNN{\scriptsize($l_f = 0, l_a = 0$, 5\r{A})}  &  0.5497 & 0.6851 & 0.5519 & 0.6102 & 4.99$\%$ & 2.50$\%$ & 4.71$\%$ & 2.88${\%}$ & 0.025 \\
SEGNN{\scriptsize($l_f = 0, l_a = 0$, 3\r{A})}  &  0.5931 & 0.7411 & 0.6182 & 0.6564 & 4.87$\%$ & 2.48$\%$ & 4.34$\%$ & 2.63${\%}$ & 0.015 \\
\textbf{SEGNN}{\scriptsize($\mathbf{l_f = 1, l_a = 1}$, \textbf{5\r{A}})}  &  0.5310 & 0.6432 & 0.5341 & 0.5777 & 5.32$\%$ & 2.80$\%$ & 4.89$\%$ & 3.09${\%}$ & 0.080 \\
SEGNN{\scriptsize($l_f = 1, l_a = 1$), 3\r{A})}  &  0.5523 & 0.6761 & 0.5609 & 0.6127 & 5.03$\%$ & 2.59$\%$ & 4.41$\%$ & 2.70${\%}$ & 0.040 \\
SEGNN{\scriptsize($l_f = 1, l_a = 2$), 5\r{A})}  &  0.5337 & 0.6419 & 0.5389 & 0.5764 & 5.41$\%$ & 2.78$\%$ & 4.89$\%$ & 3.09${\%}$ & 0.112 \\
SEGNN{\scriptsize($l_f = 1, l_a = 2$), 3\r{A})}  &  0.5497 & 0.6662 & 0.5587 & 0.6060 & 4.98$\%$ & 2.51$\%$ & 4.40$\%$ & 2.60${\%}$ & 0.051 \\
SEGNN{\scriptsize($l_f = 2, l_a = 2$), 3\r{A})}  &  0.5369 & 0.6473 & 0.5335 & 0.5822 & 5.22$\%$ & 2.82$\%$ & 4.87$\%$ & 2.98${\%}$ & 0.234 \\
Dimenet++  &  - & - & - & - & - & - & - & - &  0.062 \\
\bottomrule
\end{tabular}
\end{adjustbox}
\label{tab:ocp_ablation}
\end{table*}

\section{Introduction}
The success of Convolutional Neural Networks (CNNs)~\citep{lecun1998gradient, lecun2015deep, schmidhuber2015deep, krizhevsky2012imagenet} is a 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 ~\citep{cohen2016group, cohen2017steerable, worrall2017harmonic, cohen2018spherical, kondor2018generalization, weiler2018, bekkers2018roto,bekkers2019b}.
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 \citep{gilmer2017neural} based on a sum aggregation of messages that are obtained by relative position-dependent \textit{linear} transformations of neighbouring node features. The power of message passing layers is, however, that node features are transformed and propagated in a highly \textit{non-linear} manner.
Equivariant GNNs have been proposed before as either PointConv-type \citep{wu2019pointconv,schutt2017schnet} implementations of steerable \citep{thomas2018tensor, anderson2019covariant, fuchs2020se} or regular group convolutions \citep{finzi2020generalizing}. The most important component in these methods are the convolution layers. Although powerful, such layers only (pseudo\footnote{Methods such as SE(3)-transformers \citep{fuchs2020se} and Cormorant \citep{anderson2019covariant} include an input-dependent attention component that augments the convolutions.}) linearly transform the graphs and non-linearity is only obtained via point-wise activations.

In this paper, we propose non-linear E(3) \textit{equivariant message passing} layers using the same principles that underlie steerable group convolutions, and view them as \textit{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.~\ref{sec:method} and illustrate it in Fig.~\ref{fig:commutative} on a molecular graph. As a 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 to steer the update and message functions by geometric information such as relative
orientation (pose). Through a notion of steerable node attributes we provide a new class of equivariant activation functions for general use with steerable feature fields \citep{weiler2018,thomas2018tensor}. Node attributes can include information such as node velocity, force, or atomic spin. Currently, especially in 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 a new state of the art on the n-body toy dataset, in which our method leverages the abundance of geometric and physical quantities available.
We further test our model on the molecular datasets QM9 and OC20. Although here only (relative) positional information is available as geometric quantites, our SEGNNs further achieve state of the art on the IS2RE dataset of OC20, and achieve competitive performance on QM9.
For all experiments we provide extensive ablation studies.
The main contributions of this paper are:
\begin{itemize}
    \item Generalisation of equivariant graph networks such that node and edge attributes are not restricted to be scalars.
    \item A new class of equivariant activation functions for steerable 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.
    \item A unifying view on various equivariant graph neural networks through the definition of non-linear convolutions.
    \item 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.
\end{itemize}
\vspace*{-0.3cm}
\begin{figure}[h]
    \centering
    \includegraphics[width=0.75\textwidth]{Figures/Equivariance.png}
    \caption{Commutation diagram for equivariant operator $\phi$ 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.
    }
    \label{fig:commutative}
\end{figure}

\section{Generalised E($3$) equivariant steerable message passing}
\label{sec:method}
\paragraph{Message passing networks.}
Consider a graph $\mathcal{G} =(\mathcal{V},\mathcal{E})$, with nodes $v_i \in \mathcal{V}$ and edges $e_{ij} \in \mathcal{E}$, with feature vectors $\mathbf{f}_i \in \mathbb{R}^{c_n}$ attached to each node,
and edge attributes $\mathbf{a}_{ij} \in \mathbb{R}^{c_e}$ attached to each edge.
Graph neural networks (GNNs)~\citep{scarselli2008, kipf2017semisupervised, defferrard2016, battaglia2018}
are designed to learn from graph-structured data and are by construction permutation
equivariant with respect to the input. A specific type of GNNs are message passing networks~\citep{gilmer2017neural}, where a layer updates node features via the following steps:
\begin{align}
\text{compute message $\textbf{m}_{ij}$ from node $v_j$ to $v_i$:}\hspace{10mm} & \mathbf{m}_{ij} = \phi_m\left(\mathbf{f}_i, \mathbf{f}_j,\mathbf{a}_{ij}\right) \ , \label{eq:message1} \\
\text{aggregate messages and update node features $v_i$:}\hspace{10mm} & \mathbf{f}'_i = \phi_f\left(\mathbf{f}_i, \sum_{j \in \mathcal{N}(i)} \mathbf{m}_{ij} \right) \ ,  \label{eq:message3}
\end{align}

where $\mathcal{N} (i)$ represents the set of neighbours of node $v_i$, and $\phi_m$ and $\phi_f$ are commonly parameterised by multilayer perceptrons (MLPs).
\iffalse
In step 1, a learnable function $\phi_m: \mathbb{R}^{c_l} \times \mathbb{R}^{c_l} \times \mathbb{R}^{a_e} \rightarrow \mathbb{R}^m$ computes an $m$-dimensional message $\mathbf{m}_{ij} \in \mathbb{R}^m$ that encodes for the information that is passed on from a neighbouring node $j$ to the to-be-updated node $i$. The message is constructed using the feature vectors $\mathbf{f}^l_i$ and $\mathbf{f}^l_j$ at node $i$ and $j$ respectively, as well as the edge attribute vector $\mathbf{a}_{ij}$.
In step 2, the messages coming in from all neighbours are aggregated into a single message $\mathbf{m}_i$.
$\mathcal{N} (i)$ represents the set of neighbours of node $v_i$.
In step 3, the new node features (at the next layer $l+1$) are computed using the current feature vector (at layer $l$), the aggregated incoming messages and the vertex attribute vector $\mathbf{a}_{i}$, via another learnable function $\phi_f: \mathbf{R}^{c_l} \times \mathbf{R}^m \times \mathbb{R}^{a_v}\rightarrow \mathbb{R}^{c_{l+1}}$.
\fi

\textbf{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 as molecular energy prediction, require E(3) invariance, whereas others, like force prediction, require equivariance.
From a technical point of view, equivariance of a function $\phi$ to certain transformations means that for any transformation parameter $g$ and all inputs $x$ we have
$
    T'_g [\phi(x)]=\phi(T_g[x]) \ ,
    \label{eq:equivariance_main}
$
where $T_g$ and $T'_g$ denote transformations on the input and output domain of $\phi$, respectively.
Equivariant operators applied to atomic graphs allow us to 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 $\phi_m$ and $\phi_f$ of Eqs.~(\ref{eq:message1}-\ref{eq:message3}) to be equivariant, which in return guarantees equivariance of the entire neural network. The following paragraphs introduce the core components behind our method; full mathematical details and background can be found in App.~\ref{app:math}.

\iffalse
\textbf{Equivariance. } From a technical point of view, equivariance of a function $\phi$ to certain transformations means that for any transformation parameter $g$ and all inputs $x$ we have
\begin{equation}
    T'_g [\phi(x)]=\phi(T_g[x]) \ ,
    \label{eq:equivariance_main}
\end{equation}
where $T_g$ and $T'_g$ denote transformations on the input and output domain of $\phi$, respectively. In less formal terms, Eq.~(\ref{eq:equivariance_main}) tells us that equivariant operators handle all transformations equally well and that no information is lost when the input is transformed, since we can always recover the untransformed scenario via the inverse transformation $T'_{g^{-1}}$ on the output.
Our main group of interest is the 3D Euclidean group E$(3) = \mathbb{R}^3 \rtimes \text{O}(3)$, which is a semi-direct product (denoted with $\rtimes$) of the group of translations $\mathbb{R}^3$ and the group of orthogonal transformations O($3$).
\fi

\textbf{Steerable features. } In this work, we achieve equivariant graph neural networks by working with \emph{steerable feature vectors}, which we denote with a tilde, e.g. a vector $\tilde{\mathbf{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 $\mathbf{D}(g) \tilde{\mathbf{h}}$. For example, a Euclidean vector in $\mathbb{R}^3$ is steerable for rotations $g= \mathbf{R} \in \text{SO}(3)$ by multiplying the vector with a rotation matrix, thus $\mathbf{D}(g) = \mathbf{R}$. We are however not restricted to only work with 3D vectors; via the construction of steerable vector spaces, we can generalise the notion of 3D rotations to arbitrarily large vectors.
\newline
Central to our approach is the use of \textit{Wigner-D matrices} $\mathbf{D}^{(l)}(g)$ \footnote{In order to be O(3)—not just SO(3)—equivariant, we include reflections. See App.~\ref{app:math} for more detail.}. These are $(2l+1 \times 2l+1)$-dimensional matrix representations that act on $(2l+1)$-dimensional vector spaces. These vector spaces that are transformed by $l$-th degree  Wigner-D matrices will be referred to as \textit{type-$l$ steerable vector spaces} and denoted with $V_l$.
We note that we can combine two independent steerable vector spaces $V_{l_1}$ and $V_{l_2}$ of type $l_1$ and $l_2$ by the direct sum, denoted by $V = V_{l_1} \oplus V_{l_2}$. Such a combined vector space then transforms by the direct sum of Wigner D-matrices, i.e., via $\mathbf{D}(g) = \mathbf{D}^{(l_1)}(g) \oplus \mathbf{D}^{(l_2)}(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 to degree $l=L$ by $V_L:=V_0 \oplus V_1 \oplus \dots \oplus V_L$, and $n$ copies of the the same vector space with $nV:=\underbrace{V \oplus V \oplus \ldots \oplus V}_{n \ \text{times}}$. Regular MLPs are based on transformations between $d$-dimensional type-0 vector spaces i.e., $\mathbb{R}^d = dV_0$, and are a special case of our steerable MLPs that act on steerable vector spaces of arbitrary type.

\textbf{Steerable MLPs. } Like regular MLPs, steerable MLPs are constructed by interleaving linear mappings (matrix-vector multiplications) with non-linearities. Now however, the linear maps transform between steerable vector spaces at layer $i-1$ to layer $i$ via
$
\tilde{\mathbf{h}}^{i} = \mathbf{W}_{\tilde{\mathbf{a}}}^{i}\tilde{\mathbf{h}}^{i-1} \ .
$
Steerable MLPs thus have the same functional form as regular MLPs, although, in our case, the linear transformation matrices $\mathbf{W}_{\tilde{\mathbf{a}}}^i$, defined below, are conditioned on geometric information (e.g. relative atom positions) which is encoded in the steerable vector $\tilde{\mathbf{a}}$.
Both vectors $\tilde{\mathbf{a}}$ and $\tilde{\mathbf{h}}$ are steerable vectors. In this work we will however use the vector $\tilde{\mathbf{a}}$ to have geometric and structural information encoded and the steer the information flow of $\tilde{\mathbf{h}}$ through the network.
In order to guarantee that $\mathbf{W}_{\tilde{\mathbf{a}}}^i$ 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
\begin{equation}
{\rm \widetilde{MLP}}( \mathbf{D}(g) \tilde{\mathbf{h}}_0) = \mathbf{D}'(g) {\rm \widetilde{MLP}}(\tilde{\mathbf{h}}_0) \ ,
\end{equation}
provided that the steerable vectors $\tilde{\mathbf{a}}$ that condition the MLPs are also obtained equivariantly.

\textbf{Spherical harmonic embedding of vectors. }
We can convert any vector $\mathbf{x} \in \mathbb{R}^3$ into a type-$l$ vector through the evaluation of \textit{spherical harmonics} $Y^{(l)}_m:S^2 \rightarrow \mathbb{R}$ at $\tfrac{\mathbf{x}}{\lVert \mathbf{x} \rVert}$. For any $\mathbf{x} \in \mathbb{R}^3$
\begin{equation}
\label{eq:embedding}
    \tilde{\mathbf{a}}^{(l)} = \left(Y^{(l)}_m\left(\tfrac{\mathbf{x}}{\lVert \mathbf{x} \rVert}\right)\right)_{m=-l,-l+1,\dots,l}^T
\end{equation}
is a type-$l$ steerable vector. The spherical harmonic functions $Y_m^{(l)}$ are functions on the sphere $S^2$ and we visualise them as such in Figure \ref{fig:spherical_harmonics}..
We will use spherical harmonic embeddings to include geometric and physical information into steerable MLPs.

\begin{figure}%
    \centering
    \subfloat{{\includegraphics[width=0.48\textwidth]{Figures/embedding.png} }}%
    \quad
    \subfloat{{\includegraphics[width=0.48\textwidth]{Figures/embedding_rotated.png} }}%
    \caption{Left: representation of an O($3$) steerable vector $\tilde{\mathbf{h}} \in V_L = V_0 \oplus V_1 \oplus V_2$,
    a spherical harmonic embedding of vector $\mathbf{x}$, e.g. relative orientation, velocity or force.
    Right: for each subspace the embedding with the basis functions $Y_m^{(l)}$ is shown (a). The transformation of $\tilde{\mathbf{h}}$ via $\mathbf{D}^{(l)}(g)$ acts on each subspace of $V_l$ separately (b).}%
    \label{fig:spherical_harmonics}%
\end{figure}

\textbf{Mapping between steerable vector spaces. }
The Clebsch-Gordan (CG) tensor product $\otimes_{cg}: V_{l_1} \times V_{l_2} \rightarrow V_{l}$ is a bilinear operator that combines two O($3$) steerable input vectors of types $l_1$ and $l_2$ and returns another steerable vector of type $l$.
Let $\tilde{\mathbf{h}}^{(l)} \in V_l = \mathbb{R}^{2l+1}$ denote a steerable vector of type $l$ and $h^{(l)}_m$ its components with $m=-l,-l+1,\dots,l$.
The CG tensor product is given by
\begin{equation}
(\tilde{\mathbf{h}}^{(l_1)} \otimes_{cg}^{w} \tilde{\mathbf{h}}^{(l_2)})^{(l)}_{m} = w \sum_{m_1=-l_1}^{l_1} \sum_{m_2=-l_2}^{l_2}  \, C^{(l,m)}_{(l_1,m_1)(l_2,m_2)} h^{(l_1)}_{m_1} h^{(l_2)}_{m_2} \ ,
\label{eq:CG_product}
\end{equation}

in which $w$ is a learnable parameter that scales the product, and
$C^{(l,m)}_{(l_1,m_1)(l_2,m_2)}$ are the Clebsch-Gordan coefficients that ensure that the resulting vector is type-$l$ steerable. The CG tensor product is a sparse tensor product, as generally many coefficients are zero. Most notably, $C^{(l,m)}_{(l_1,m_1)(l_2,m_2)}=0$ whenever $ l < |l_1 - l_2|$ or $l > l_1 + l_2$.
While Eq.~(\ref{eq:CG_product}) only describes the product between steerable vectors of a single type, e.g. $\tilde{\mathbf{h}}^{(l_1)} \in V_{l_1}$ and $\tilde{\mathbf{h}}^{(l_2)} \in V_{l_2}$, it is directly extendable to mixed type steerable vectors that may have multiple channels/multiplicities within a type. In this case, every input to output sub-vector pair gets its own index in a similar way as the weights in a standard linear layer are indexed with input-output indices. We then denote the CG product with $\otimes_{cg}^{\mathbf{W}}$ with boldfaced $\mathbf{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 $\tilde{\mathbf{a}}$ in one of its inputs as a \emph{steerable linear layer conditioned on $\tilde{\mathbf{a}}$}, denoted with
\begin{align}
\mathbf{W}_{\tilde{\mathbf{a}}} \tilde{\mathbf{h}} := \tilde{\mathbf{h}} \otimes_{cg}^\mathbf{W} \tilde{\mathbf{a}} \ ,
\;\;\;\;\; \text{or} \;\;\;\;\; \mathbf{W}_{\tilde{\mathbf{a}}}(d) \tilde{\mathbf{h}} := \tilde{\mathbf{h}} \otimes_{cg}^{\mathbf{W}(d)} \tilde{\mathbf{a}} \ ,
\label{eq:steerable_layer}
\end{align}
where the latter indicates that the CG weights depend on some quantity $d$, e.g. relative distances.

\textbf{Steerable activation functions. }
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 \citep{cohen2018spherical}, norm-altering \citep{thomas2018tensor}, or gated non-linearities \citep{weiler2018}.
All these can be used in alternation with (\ref{eq:steerable_layer}). 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 $\tilde{\mathbf{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.

\subsection{Steerable E($3$) Equivariant Graph Neural Networks}
\label{sec:segnns}
We extend the message passing equations \eqref{eq:message1}-\eqref{eq:message3} and define a message passing layer that updates steerable node features $\tilde{\mathbf{f}}_i \in V_L$ at node $v_i$ via the following steps:
\begin{align}
\tilde{\mathbf{m}}_{ij} &= \phi_m\left(\tilde{\mathbf{f}}_i, \tilde{\mathbf{f}}_j, \lVert \mathbf{x}_j - \mathbf{x}_i \rVert^2, \tilde{\mathbf{a}}_{ij}\right) \label{eq:segnn1} \ , \\
\tilde{\mathbf{f}}'_i &= \phi_f\left(\tilde{\mathbf{f}}_i, \sum_{j \in \mathcal{N}(i)} \tilde{\mathbf{m}}_{ij}  , \tilde{\mathbf{a}}_i \right) \ . \label{eq:segnn3}
\end{align}
Here, $\lVert \mathbf{x}_j - \mathbf{x}_i \rVert^2$ is the squared relative distance between two nodes $v_i$ and $v_j$,
$\phi_m$ and $\phi_f$ are O($3$) steerable MLPs,
$\tilde{\mathbf{a}}_{ij} \in V_L$ and $\tilde{\mathbf{a}}_i \in V_L$ are steerable edge and node attributes.
If additional attributes exist, such as pair-wise distance $\lVert \mathbf{x}_j - \mathbf{x}_i \rVert$, they can either be concatenated to the attributes that condition our steerable MLPs, or as is more commonly (Sec.~\ref{sec:convolution}) done add them as inputs to $\phi_m$ and $\phi_f$. We do the latter, and stack all inputs to a single steerable vector by which the steerable MLP layer is given:
$\tilde{\mathbf{h}}^{1} = \sigma ( \mathbf{W}_{\tilde{\mathbf{a}}_{ij}} \tilde{\mathbf{h}}_i^0 ) \ ,$
with $\tilde{\mathbf{h}}_i^0 = \left(\tilde{\mathbf{f}}_i, \tilde{\mathbf{f}}_j, \lVert \mathbf{x}_j - \mathbf{x}_i \rVert^2 \right) \in V_f \oplus V_f \oplus V_0$, where $V_f$ is the user specified steerable vector space of node representations. The message network $\phi_m$ is steered via edge attributes $\tilde{\mathbf{a}}_{ij}$, and the node update network $\phi_f$ is similarly steered via node attributes
$\tilde{\mathbf{a}}_{i}$.

\textbf{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.~(\ref{eq:embedding})) 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.,
$\tilde{\mathbf{a}}_i = \frac{1}{\lvert \mathcal{N}(i)  \rvert} \sum_{j \in \mathcal{N}(i)} \tilde{\mathbf{a}}_{ij}$, 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 $\phi_f$ allows us to not just integrate geometric cues into the message functions, but also leverage it in the node updates.
We observe that the more geometric and physical quantities are injected to better SEGNNs perform.

\section{Message Passing as Convolution, Related Work} \label{sec:convolution}
In this section, we revisit related work from a convolutional perspective. We build this section on our formulation of steerable layers as introduced in Eq.~(\ref{eq:steerable_layer}). Such a general framework does not only help to understand and better categorise the existing architectures in literature, but also allows us to clearly motivate the novelty of our approach.
The theoretical findings are tested and confirmed in Sec.~\ref{sec:experiments}.

\textbf{Convolutions and point convolutions. }
Consider a feature map $\mathbf{f}:\mathbb{R}^d\rightarrow \mathbb{R}^{c_l}$ defined on an Euclidean input space.
A convolution layer (defined via cross-correlation) with a point-wise non-linearity $\sigma$ is then given by
\begin{align}
\label{eq:continuous_convolution}
\mathbf{f}'(\mathbf{x}) = \sigma \bigg( \int_{\mathbb{R}^d} \mathbf{W}(\mathbf{x}' - \mathbf{x}) \mathbf{f}(\mathbf{x}'){\rm d}\mathbf{x}' \bigg) \ ,
\end{align}
with $\mathbf{W}:\mathbb{R}^d \rightarrow \mathbb{R}^{c_{l+1} \times c_{l}}$ being a convolution kernel that provides for every relative position a matrix that transforms features from $\mathbb{R}^{c_l}$ to $\mathbb{R}^{c_{l+1}}$.
\textit{Point convolutions}, generally referred to as PointConvs~\citep{wu2019pointconv}, and SchNet~\citep{schutt2017schnet} implement Eq.~(\ref{eq:continuous_convolution}) on point clouds. For a sparse input feature map consisting of location-feature pairs $(\mathbf{x}_i, \mathbf{f}_i)$, the point convolution is given by
$
    \mathbf{f}'_i = \sum_{j \in \mathcal{N}(i)} \mathbf{W}(\mathbf{x}_j - \mathbf{x}_i) \mathbf{f}_i \ ,
$
which describes a message passing layer of Eqs.~(\ref{eq:message1})-(\ref{eq:message3}) in which the messages are $\mathbf{m}_j = \mathbf{W}(\mathbf{x}_j -\mathbf{x}_i) \mathbf{f}_j$ and the message update $\mathbf{f}'_i = \sum_j \mathbf{m}_{ij}$.

In the above convolutions, the linear transformation matrices $\mathbf{W}$ are conditioned on relative positions $\mathbf{x}_j - \mathbf{x}_i$, which is typically done in one of the following three approaches. (i) Classically, data on a dense discrete grid with shared neighbourhoods is used and transformations $\mathbf{W}_{ij}$ for relative positions $\mathbf{x}_j - \mathbf{x}_i$ are stored. This method however does not generalise to non-uniform grids such as point clouds.
(ii) Continuous kernel methods parametrise the transformations either by expanding it into a continuous basis that can be sampled at arbitrary locations, i.e.,
$\mathbf{W}(\mathbf{x}_j - \mathbf{x}_i)=\sum_b w_b B_b(\mathbf{x}_j - \mathbf{x}_i)$, where $w_b$ are trainable parameters and $B_b$ the basis functions, or by parametrising the transformations as MLPs via $\mathbf{W}(\mathbf{x}_j - \mathbf{x}_i) = \text{MLP}(\mathbf{x}_j - \mathbf{x}_i)$.
(iii) Steerable kernel methods expand the linear transformations in a steerable basis, such as 3D spherical harmonics, i.e.,
\begin{equation}
\label{eq:shkernel}
\mathbf{W}(\mathbf{x}_j - \mathbf{x}_i) = \sum_l \sum_{m=-l}^l \mathbf{W}_{m}^{(l)}(\lVert \mathbf{x}_j - \mathbf{x}_i\rVert) Y^{(l)}_m(\mathbf{x}_j - \mathbf{x}_i) \ ,
\end{equation}
where basis coefficients typically depend on the distance between point pairs. In what follows we establish a relation between the second and third type of kernel parametrisations.

\textbf{Steerable (group) convolutions. }
In our steerable setting, linear feature transformations $\mathbf{W}$ are equivalent to steerable linear transformations conditioned on the scalar ``1''. I.e, with $\mathbf{h}_i \in \mathbb{R}^d$ and $\tilde{\mathbf{h}} \in dV_0$ the messages are given by
\begin{align}
    \mathbf{m}_{ij} = \mathbf{W}(\mathbf{x}_j - \mathbf{x}_i) \mathbf{h}_i \;\;\;\; \Leftrightarrow \;\;\;\; \tilde{\mathbf{m}}_{ij} = \mathbf{W}_1(\mathbf{x}_j - \mathbf{x}_i) \tilde{\mathbf{h}}_i \ .
\end{align}
When the transformations are parametrised in a steerable basis (\ref{eq:shkernel}) we can make the identification\footnote{Exact correspondence is obtained by a sum reduction over the steerable vector components (App.~\ref{sec:groupconvs}).}
\begin{equation}
\label{eq:steerableconv}
    \mathbf{m}_{ij} = \mathbf{W}(\mathbf{x}_j - \mathbf{x}_i) \mathbf{h}_i \;\;\;\; \Leftrightarrow \;\;\;\;
    \tilde{\mathbf{m}}_{ij} = \mathbf{W}_{\tilde{\mathbf{a}}_{ij}}(\lVert \mathbf{x}_j - \mathbf{x}_i\rVert) \tilde{\mathbf{h}}_i \ ,
\end{equation}
in which $\tilde{\mathbf{a}}_{ij}$ are the spherical harmonic embeddings (Eq.~\eqref{eq:embedding}) of $\mathbf{x}_j - \mathbf{x}_i$, and the weights that parametrise the CG tensor product depend on the distance between point pairs. Thus, convolutions with kernels expanded in a spherical harmonic basis can either be done as usual via regular, or via steerable convolutions. The advantage of steerable convolutions is that we can directly derive what the result would be if the kernel were to be rotated by making use of the Wigner-D matrices via
\begin{equation}
\label{eq:rotatedconv}
    \mathbf{m}_{ij} = \mathbf{W}(\mathbf{R}^{-1}( \mathbf{x}_j - \mathbf{x}_i)) \mathbf{h}_i \;\;\;\; \Leftrightarrow \;\;\;\;
    \tilde{\mathbf{m}}_{ij} = \mathbf{D}(\mathbf{R}) \mathbf{W}_{\tilde{\mathbf{a}}_{ij}}(\lVert \mathbf{x}_j - \mathbf{x}_i\rVert) \tilde{\mathbf{h}}_i \ .
\end{equation}
In fact, through the identification of steerable vectors with functions on O(3) via the inverse Fourier transform of,
we can treat the steerable feature vectors $\tilde{\mathbf{f}}_i$ at each location $\mathbf{x}_i$ as functions on the group O(3). As such, steerable convolutions produce feature maps on the full group E(3) that for every possible translation/position $\mathbf{x}$ and rotation $\mathbf{R}$ provide a feature response $\mathbf{f}(\mathbf{x}, \mathbf{R})$. It is precisely this mechanism of transforming convolution kernels via the group action (via $\mathbf{W}(g^{-1} \cdot \mathbf{x}_j) = \mathbf{W}(\mathbf{R}^{-1}(\mathbf{x}_j - \mathbf{x})$) that underlies group convolutions \citep{cohen2016group}. As a result, message passing via explicit kernel rotations (l.h.s. of (\ref{eq:rotatedconv})) corresponds to \textit{regular group convolutions}, and via steerable transformations (r.h.s. of (\ref{eq:rotatedconv})) corresponds to \textit{steerable group convolutions}.

\iffalse
\erik{somewhere mention this: CGCNNs~\citep{xie2018crystal} aggregate updates without explicit dependencies on $\mathbf{x}_j - \mathbf{x}_i$:
\begin{align}
    \mathbf{m}_i = \sum_{j \in \mathcal{N}(i)} \mathbf{W}_1 \mathbf{h}_i \ ,
\end{align}
where $\mathbf{h}_i=(\mathbf{f}_i, \mathbf{f}_j, \mathbf{u}_{i,j})$ and $\mathbf{u}_{i,j}$ is an edge feature vector. }
\fi

The equivariant steerable methods \citep{thomas2018tensor,anderson2019covariant,miller2020relevance,fuchs2020se} that we compare against in our experiments can all be written in convolution form
\begin{align}
    \label{eq:steerable_convolution}
    \tilde{\mathbf{f}}'_i = \sum_{j \in \mathcal{N}(i)}\mathbf{W}_{\tilde{\mathbf{a}}_{ij}}(\lVert \mathbf{x}_j - \mathbf{x}_i\rVert) \tilde{\mathbf{f}}_i \ , \;\;\;\;\;\;\; \text{or} \;\;\;\;\;\;\;
    \tilde{\mathbf{f}}'_i = \sum_{j \in \mathcal{N}(i)}\mathbf{W}_{\tilde{\mathbf{a}}_{ij}}(\tilde{\mathbf{f}}_i, \tilde{\mathbf{f}}_j, \lVert \mathbf{x}_j - \mathbf{x}_i\rVert) \tilde{\mathbf{f}}_i \ ,
\end{align}
where, in the latter case, the linear transformations additionally depend on an input dependent attention mechanism as in \citep{anderson2019covariant,fuchs2020se}, and can be seen as a steerable PointConv version of attentive group convolutions \citep{romero2020}. In these attention-based cases, convolutions are augmented with input dependent weights $\alpha$ via $\mathbf{W}_{\tilde{\mathbf{a}}_{ij}}(\tilde{\mathbf{f}}_i, \tilde{\mathbf{f}}_j, \lVert \mathbf{x}_j - \mathbf{x}_i\rVert) = \alpha(\tilde{\mathbf{f}}_i, \tilde{\mathbf{f}}_j) \mathbf{W}_{\tilde{\mathbf{a}}_{ij}}(\lVert \mathbf{x}_j - \mathbf{x}_i \rVert)$. 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 transformation.
\iffalse
To finish our exposition, we use the notation of Eq.(\ref{eq:steerable_layer}) and rewrite Eq.(\ref{eq:steerable_convolution}) as:
\begin{align}
    \tilde{\mathbf{f}}'_i = \sum_{j \in \mathcal{N}(i)}\tilde{\mathbf{f}}_i \otimes_{cg}^{\mathbf{W}(\lVert \mathbf{x}_j - \mathbf{x}_i\rVert)} \tilde{\mathbf{a}}_{ij} \ , \;\;\;\;\;\;\; \text{and} \;\;\;\;\;\;\;
    \tilde{\mathbf{f}}'_i = \sum_{j \in \mathcal{N}(i)}\tilde{\mathbf{f}}_i \otimes_{cg}^{\mathbf{W}(\tilde{\mathbf{f}}_i, \tilde{\mathbf{f}}_j, \lVert \mathbf{x}_j - \mathbf{x}_i\rVert )} \tilde{\mathbf{a}}_{ij} \ .
\end{align}
\fi
Finally, the recently proposed LieConv~\citep{finzi2020generalizing} and NequIP~\citep{batzner2021se}
also fall in the convolutional message passing class.
LieConv is a PointConv-type variation of \emph{regular group convolutions} on Lie groups~\citep{bekkers2019b}.
NeuqIP follows the approach of Tensor Field Networks~\citep{thomas2018tensor}, and weighs interactions using an MLP with radial basis functions as input. These functions are obtained as solution of the Schrödinger equation. Thus, learnable filters consisting of spherical harmonics  convolve nearby points based on the distance to the central point.

\textbf{Equivariant message passing as non-linear convolution. }
EGNNs \citep{satorras2021n} are equivariant to transformations in E($n$) and outperform most aforementioned steerable methods. This is somewhat surprising as it sends \textit{invariant} messages, which are obtained via MLPs of the form
\begin{equation}
\label{eq:egnnmessages}
\mathbf{m}_{ij} = \text{MLP}(\mathbf{f}_i, \mathbf{f}_j, \lVert \mathbf{x}_j - \mathbf{x}_i \rVert^2) = \sigma(\mathbf{W}^{(k)} (\ldots(\sigma(\mathbf{W}^{(1)} \mathbf{h}_i)))) \ ,
\end{equation}
where $\mathbf{h}_i=(\mathbf{f}_i, \mathbf{f}_j, \lVert \mathbf{x}_j - \mathbf{x}_i \rVert^2)$. These messages resemble the convolutional messages of point convolutions due to their dependency on relative positions. There are, however, two import differences: (i) the messages are non-linear transformations of the neighbouring feature values $\mathbf{f}_j$ via an MLP and (ii) the messages are only conditioned on the distance between point pairs, and are therefore E($n$) invariant. As such, we regard EGNN layers as \textit{non-linear convolutions} with \textit{isotropic message} functions (the non-linear counterpart of a rotationally invariant kernel).
In our work, we lift the isotropy constraint and generalise to non-linear steerable convolutions via messages of the form
\begin{align}
\label{eq:segnnmessage}
    \tilde{\mathbf{m}}_{ij} = \widetilde{\text{MLP}}_{\tilde{\mathbf{a}}_{ij}}(\tilde{\mathbf{f}}_i, \tilde{\mathbf{f}}_j, \lVert \mathbf{x}_j - \mathbf{x}_i \rVert^2) =  \sigma(\mathbf{W}^{(n)}_{\tilde{\mathbf{a}}_{ij}} (\ldots(\sigma(\mathbf{W}^{(1)}_{\tilde{\mathbf{a}}_{ij}} \tilde{\mathbf{h}}_i)))) \ ,
\end{align}
with $\tilde{\mathbf{h}}_i=(\tilde{\mathbf{f}}_i, \tilde{\mathbf{f}}_j, \lVert \mathbf{x}_j - \mathbf{x}_i \rVert^2) \in V_f \oplus V_f \oplus V_0$. The MLP is then conditioned on attribute  $\tilde{\mathbf{a}}_{ij}$, which could e.g. be a spherical harmonic embedding of $\mathbf{x}_j - \mathbf{x}_i$. This allows for the creation of messages more general than those found in convolution, while carrying covariant geometrical information.

\textbf{Related equivariant message passing methods. }
A different but also fully message passing based approach can be found in Geometric Vector
Perceptrons (GVP)~\citep{jing2020learning} and PaiNN~\citep{schutt2021equivariant}.
Compared to SEGNNs which treat equivariant information as fully steerable features, these GVP and PaiNN update scalar-valued attributes using the norm of vector-valued attributes, and therefore with O($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. Another approach for incorporating relative orientation is to utilise invariant angles. Methods such as Dimenet++ \citep{klicpera2020dimenet_plusplus} and SphereNet \citep{liu2021spherical} do so by relying on a \textit{second-order} message passing scheme that considers angles between the central point and neighbours-of-neighbours. These angles are not affected by rigid body motions and can therefore be used without making the model overly sensitive to O$(3)$ transformations. In contrast, our method can directly leverage angular information in a \textit{first-order} message passing scheme through the use of steerable vectors.

\section{Experiments}
\label{sec:experiments}

\textbf{Implementation details. }
The implementation of SEGNN's O(3) steerable MLPs is based on the \texttt{e3nn} library \citep{geiger2021}. We either define the steerable vector spaces as $V=nV_{L=l_{max}}$ (N-body, QM9 experiments), i.e., $n$ copies of steerable vector spaces up to order $l_{max}$, 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 \citep{finzi2021practical}. Furthermore, for a fair comparison between experiments with different $l_{max}$, we choose $n$ such that the total number of weights in the CG products corresponds to that of a regular (type-0) linear layer.
Further implementation details are in App.~\ref{sec:appendix_experiments}.

\textbf{SEGNN architectures and ablations. }
We consider several variations of SEGNNs.
On all tasks we have at least one fully steerable ($l_f>0,l_a>0$) SEGNN tuned for the specific task at hand.
We perform ablation experiments to investigate two main principles that sets SEGNNs apart from the literature.
\textbf{A1} The case of non-steerable vs steerable EGNNs is obtained by applying the same SEGNN network with different specifications of maximal spherical harmonic order in the feature vectors ($l_f$) and the attributes ($l_a$). EGNN \citep{satorras2021n} arises as a special case with $l_f=l_a=0$. These models will be labelled SEGNN.
\textbf{A2} In this ablation, we use steerable equivariant point conv methods ~\citep{thomas2018tensor}
with messages as in Eq.~(\ref{eq:steerable_convolution}) and regular gated non-linearities as activation/update function, labelled SE$_\text{linear}$, and compare it to the same network but with messages obtained in a non-linear manner via 2-layer steerable MLP as in Eq.~(\ref{eq:segnnmessage}), labelled as SE$_\text{non-linear}$.

\textbf{N-body system. }
The charged n-body particle system experiment is a toy experiment used in
\citep{kipf2018neural, fuchs2020se, satorras2021n}.
It consists of 5 particles that carry a positive or negative charge,
and have an initial position and an initial velocity in a 3-dimensional space. The task is to estimate the positions of the five particles after 1.000 timesteps. We build upon the code and the experimental setting introduced in~\citep{satorras2021n} (3.000 particles, 10.000 epochs, 4 layer networks for the SEGNN and more layers for those ablation where we don't fill the parameter budget).
Steerable architectures are designed such that the parameter budget at  $l_f=1$ and $l_a=1$ matches that of the tested EGNN implementation. More details are provided in App.~\ref{sec:appendix_experiments}. For the implementation, we input the relative position to the center of the system and the velocity as vectors of type $l=1$ with odd parity. We further input the norm of the velocity as scalar, which altogether results in an input vector $\tilde{\mathbf{h}} \in V_0 \oplus V_1 \oplus V_1$. The output is embedded as difference vector to the initial position (odd parity), i.e. $\tilde{\mathbf{o}} \in V_1$.
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 $\mathbf{x}_j - \mathbf{x}_i$ as described in Eq.~\ref{eq:embedding}. Messages additionally have the product of charges and absolute distance included.

SEGNN architectures are compared to steerable equivariant point conv methods (SE$_{\text{linear}}$) and steerable non-linear point conv methods (SE$_{\text{non-linear}}$).
For all cases, we use two embedding and two read-out layers.
Results and ablation studies are shown in Tab.~\ref{tab:n-body}. A full ablation on the performance and runtime of different (maximum) orders of steerable feature vectors ($l_f$) and attributes ($l_a$) is outlined in App.~\ref{sec:appendix_experiments}.
For most steerable architectures, the best results are obtained for $l_f=1$ and $l_a=1$, and architectures don't seem to benefit from higher orders of $l_f$ and $l_a$.
We consider two SEGNN architectures: one where the node attributes are the averaged edge embeddings, i.e. mean over relative orientation. This architecture is labelled SEGNN$_{\text{G}}$ since 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 the whole geometric (relative orientation) and physical (velocity) information,
and is consequently labelled SEGNN$_{\text{G+P}}$.
Including physical information in addition to geometric information in the node updates considerably boosts SEGNN performance,
which demonstrates the potential of the new class of SEGNN node activation functions.

\begin{table}[]
    \centering
    \caption{Mean Squared Error (MSE) in the N-body system experiment, and forward time in seconds for a
batch size of 100 samples running on a GeForce RTX 2080 Ti GPU. Results except for EGNN and SEGNN are taken from~\citep{satorras2021n} and verified. Runtimes are re-measured.
}
    \footnotesize
    \begin{adjustbox}{width=0.55\textwidth}
    \begin{tabular}{l l l}
  \toprule
    Method & MSE & Time [s] \\
    \midrule
    Linear & .0819 & .0001 \\
    SE(3)-Tr.~\citep{fuchs2020se} & .0244 & .0742 \\
    TFN~\citep{thomas2018tensor} & .0155 & .0182 \\
    NMP~\citep{gilmer2017neural} & .0107 & .0017 \\
    Radial Field~\citep{kohler2019equivariant} & .0104 & .0019 \\
    EGNN~\citep{satorras2021n} & .0070$\pm$ .00022  & .0029 \\
    \midrule
    SE$_{\text{linear}}$ ($l_f=2, l_a=2$) & .0116$\pm$ .00021 & .064 \\
    SE$_{\text{non-linear}}$ ($l_f=1, l_a=1$) & .0060$\pm$ .00019 & .031 \\
    SEGNN$_\text{G}$ ($l_f=1, l_a=1$) & .0056$\pm$ .00025 & .025 \\
    SEGNN$_\text{G+P}$ ($l_f=1, l_a=1$) & .0043$\pm$ .00015 & .026 \\
    \bottomrule \\
  \end{tabular}
  \end{adjustbox}
    \label{tab:n-body}
\end{table}

\paragraph{QM9.}
The QM9 dataset~\citep{ramakrishnan2014QM9, ruddigkeit2012enumeration} consists of small molecules up to 29 atoms,
where each atom is described with 3D position coordinates and a five dimensional 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.
We optimise and report the mean
absolute error (MAE) between predictions and ground truth,
using the dataset partitions and dataloaders from ~\citet{anderson2019covariant}, which splits the dataset into 100K molecules for training, 18K for validation and
13K for testing.
In Table~\ref{tab:qm9_ablation} we 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 a small cutoff radius is chosen.
While previous methods use relatively large cutoff radii of 4.5-11\r{A}, we use a cutoff radius of 2\r{A}. Doing so results in a sharp reduction of the number of messages per layer, as shown in App.~\ref{sec:appendix_experiments}.
Tables \ref{tab:qm9} and \ref{tab:qm9_ablation} together show that SEGNNs outperform an architecturally comparable baseline EGNN \citep{satorras2021n}, whilst stripping away attention modules from it and reducing graph connectivity from fully connected to only 2\r{A} distant atoms.
It is however apt to note that runtime is still limited by the relatively expensive calculation of the Clebsch-Gordan 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$, $U_0$). We conjecture that such targets could benefit from more involved (e.g. including attention or neighbour-neighbour interactions) or problem-tailored architectures, such as those compared against.

\begin{table*}[h]
  \centering
  \caption{Performance comparison on the QM9 dataset. Numbers are reported for Mean Absolute Error (MAE) between model predictions and ground truth.}
  \begin{adjustbox}{width=1.\textwidth}
  \begin{tabular}{l c c c c c c c c c c c c c c c c c c}
  \toprule
    Task & $\alpha$ & $\Delta \varepsilon$ & $\varepsilon_{\mathrm{HOMO}}$ & $\varepsilon_{\mathrm{LUMO}}$ & $\mu$ & $C_{\nu}$ & $G$ & $H$ & $R^2$ & $U$ & $U_0$ & ZPVE \\
    Units & bohr$^3$ & meV & meV & meV & D & cal/mol K & meV & meV & bohr$^3$ & meV & meV & meV \\
    \midrule
    NMP & .092 & 69 & 43 & 38 & .030 & .040 & 19 & 17 & .180 & 20 & 20 & 1.500 \\
    SchNet * & .235 & 63 & 41 & 34 & .033 & .033 & 14 & 14 & .073 & 19 & 14 & 1.700 \\
    Cormorant & .085 & 61 & 34 & 38 & .038 & .026 & 20 & 21 & .961 & 21 & 22 & 2.027  \\
    L1Net & .088 & 68 & 46 & 35 & .043 & .031 & 14 & 14 & .354 & 14 & 13 &  1.561\\
    LieConv & .084 & 49 & 30 & {25} & .032 & .038 & 22 & 24 & .800 & 19 & 19 & 2.280 \\
    TFN & .223 & 58 & 40 & 38 & .064 & .101 & - & - & - & - & - & - \\
    SE(3)-Tr. & .142 & 53 & 35 & 33 & .051 & .054 & - & - & - & - & - & -  \\
    DimeNet++ * & .0435 & 32.6 & 24.6 & 19.5 & .0297 & .0230 & 7.56 & 6.53 & .331 & 6.28 & 6.32 & 1.21 \\
    SphereNet * & .0465 & 32.3 & 23.6 & 18.9 & .0269 & .0215 & 8.0 & 6.40 & .292 & 7.33 & 6.26 & 1.12\\
    PaiNN * & .045 & 45.7 & 27.6 & 20.4 & .012 & .024 & 7.35 & 5.98 & .066 & 5.83 & 5.85 & 1.28\\
    EGNN & {.071}  & {48} & {29} & {25} & .029 & .031 & 12 & 12 & .106 & 12 & 12 & 1.55  \\
    \midrule
    SEGNN (Ours) & .060 & 42  & 24 & 21 & .023 & .031 & 15 & 16 & .66 & 13 & 15 & 1.62 \\
    \bottomrule
    & \multicolumn{12}{l}{* these methods use different train/val/test partitions.}
  \end{tabular}
  \end{adjustbox}
  \label{tab:qm9}
\end{table*}

\begin{table*}[h]
  \centering
  \caption{
  QM9 ablation study to compare SEGNN performances for different (maximum) orders of steerable feature vectors ($l_f$) and attributes ($l_a$). Models with only trivial features are akin to the invariant EGNN. The method with a fully connected graph uses soft edge estimation \citep{satorras2021n}. Forward time is measured for a batch of 128 samples running on a GeForce RTX 3090 GPU.}
  \begin{adjustbox}{width=0.9\textwidth}
  \footnotesize
  \begin{tabular}{l c c c c c c c c c c c c c c}
  \toprule
    Task & & $\alpha$ & $\Delta \varepsilon$ & $\varepsilon_{\mathrm{HOMO}}$ & $\varepsilon_{\mathrm{LUMO}}$ & $\mu$ & $C_{\nu}$ & \\
    Units & Cutoff radius &  bohr$^3$ & meV & meV & meV & D & cal/mol K & Time [s] \\
    \midrule
    (S)EGNN {\tiny($l_f = 0, l_a = 0$)} & - &  .091  & 53 & 34 & 28 & .042   & .043 & 0.016   \\
    (S)EGNN {\tiny($l_f = 0, l_a = 0$)} & $2$\r{A} &  .24  & 98 & 60 & 60 & .34   &  .077 & 0.014  \\
    SEGNN {\tiny($l_f = 1, l_a = 2$)}   & $2$\r{A} & .074 & 48 & 27 & 25 & .031  &  .035 & 0.048  \\
    SEGNN {\tiny($l_f = 2, l_a = 3$)}   & $2$\r{A} & .060 & 42 &  24     & 21 & .023 & .031 & 0.097  \\
    \bottomrule
  \end{tabular}
  \end{adjustbox}
  \label{tab:qm9_ablation}
\end{table*}

\paragraph{OC20.}
The Open Catalyst Project OC20 dataset~\citep{zitnick2020introduction, chanussot2021open}, consists of molecular adsorptions onto surfaces where
82 different adsorbates are considered, consisting of oxygen, hydrogen, carbon and nitrogen atoms.
We focus on the Initial Structure to Relaxed Energy
(IS2RE) task, which takes as input an initial structure
and has the target to predict the energy in the final, relaxed state.
The IS2RE training set consists of over 450,000 catalyst adsorbate combinations where graphs comprise 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 a $0.02~$eV threshold (EwT).
The four test splits contain in-distribution (ID) catalysts and adsorbates, out-of-domain adsorbates (OOD Ads), out-of-distribution catalysts (OOD Cat), and out-of-distribution adsorbates and catalysts (OOD Both).
Table~\ref{tab:ocp} 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.
This includes methods like
SphereNet \citep{liu2021spherical} and
DimeNet++~\citep{klicpera2019dimenet, klicpera2020dimenet_plusplus}.
The best SEGNN performance is seen for $l_f=1$ and $l_a=1$ (see App.~\ref{sec:appendix_experiments}).
A full ablation study, comparing performance and runtime for different orders of $l_f$ and $l_a$ can be found in App.~\ref{sec:appendix_experiments}).

\begin{table*}[!htb]
\centering
\caption{Comparison on the OC20 IS2RE task in terms of Mean Absolute Error (MAE) between model predictions and ground truth energy and \% of predictions within $\epsilon=0.02~$eV of the ground truth (EwT). Numbers are taken from the OC20 leaderboard. SEGNNs outperform all competitors.}
\begin{adjustbox}{width=1.\textwidth}
\scriptsize
\begin{tabular}{l c c c c | c c c c }
\toprule
\\[-0.4em]
& \multicolumn{4}{c |}{Energy MAE [eV] $\downarrow$} & \multicolumn{4}{c}{EwT $\uparrow$}  \\
\cline{2-9}
\\[-0.4em]
    Model  & ID & OOD Ads & OODCat & OOD Both & ID & OOD Ads & OOD Cat & OOD Both \\
\midrule
Median baseline &  1.7499 & 1.8793 & 1.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.6387 & 0.7342 & 0.6616 & 0.7037 & 2.96$\%$ & 2.33$\%$ & 2.94$\%$ & 2.21$\%$ \\
EdgeUpdateNet & 0.5839 & 0.7252 & 0.6016 & 0.6862 & 3.48$\%$ & 2.35$\%$ & 3.3$\%$ & 2.57$\%$ \\
EnergyNet & 0.6366 & 0.717 & 0.6387 & 0.6626 & 3.3$\%$ & 2.2$\%$ & 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 & \textbf{0.6380} & 4.47$\%$ & 2.29$\%$ & 4.09$\%$ & 2.41$\%$ \\
\midrule
SEGNN (Ours) &  \textbf{0.5327} & \textbf{0.6921} & \textbf{0.5369} & 0.6790 & \textbf{5.37}$\boldsymbol{\%}$ & \textbf{2.46}$\boldsymbol{\%}$ & \textbf{4.91}$\boldsymbol{\%}$ & \textbf{2.63}$\boldsymbol{\%}$ \\
\bottomrule
\end{tabular}
\end{adjustbox}
\label{tab:ocp}
\end{table*}

\section{Conclusion}\label{sec:limitations}
We have introduced Steerable E($3$) Equivariant Graph Neural Networks (SEGNNs) which generalise equivariant graph neural networks, such that information at nodes and edges is not restricted to be invariant (scalar), but can also be covariant (vector, tensor).
The key ingredient of SEGNNs is a new class of equivariant activation functions for steerable vector fields, based on the introduction of steerable node attributes and steerable MLPs.
Including geometric and physical information in the node updates is a unique feature that considerably boosts model performance; we demonstrate the potential of the new class of steerable node activation functions
and consider it a promising new direction in computational physics and chemistry.
Extensive ablation studies have further shown the benefit of steerable over non-steerable (invariant) message passing, and the benefit of non-linear over linear convolutions.
Furthermore, on the OC20 ISRE taks, SEGNNs outperform all competitors.

\section*{Acknowledgements}

Johannes Brandstetter thanks the Institute of Advanced Research in Artificial Intelligence (IARAI) 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.

redpajama_set_name: RedPajamaArXiv