# Steerable Transformers

Soumyabrata Kundu
Department of Statistics
University of Chicago
Chicago, IL 60637
<EMAIL_ADDRESS>
&Risi Kondor
Department of Computer Science
Department of Statistics
Computational and Applied Mathematics
University of Chicago
Chicago, IL 60637
<EMAIL_ADDRESS>

###### Abstract

In this work we introduce Steerable Transformers, an extension of the Vision
Transformer mechanism that maintains equivariance to the special Euclidean
group $\mathrm{SE}(d)$. We propose an equivariant attention mechanism that
operates on features extracted by steerable convolutions. Operating in Fourier
space, our network utilizes Fourier space non-linearities. Our experiments in
both two and three dimensions show that adding a steerable transformer encoder
layer to a steerable convolution network enhances performance.

##  1 Introduction

Transformers have emerged as the preferred architecture for natural language
processing tasks, with recent powerful models like Chat-GPT employing this
framework. Their relatively straightforward design, coupled with their
remarkable success, has led to their widespread adoption across various
domains, including image classification [11], object detection [5], and graph-
based problems [12]. The self-attention mechanism [2] employed in transformer
architectures has proven to be crucial for capturing relationships between
different parts of input sequences. Dosovitskiy et al. [11] introduced
transformers as an alternative to traditional convolutional architectures for
image classification tasks. Unlike convolutional neural networks (CNNs), which
focus on local neighborhoods, transformers excel at capturing relations across
different parts of the input, including in images.

Equivariant neural network architectures have gained significant popularity in
recent years due to their inherent ability to comprehend the underlying
symmetries of problems, making them highly effective tools for real-world
scenarios characterized by such symmetries. For instance, achieving
equivariance to the permutation group $S_{n}$ is crucial in graph-based
problems, where the structure of the graph remains invariant under
permutations. Similarly, in image-related tasks, it is desirable to have
equivariance to rigid body transformations such as rotations and translations.
Recent studies have demonstrated the remarkable efficacy of these equivariant
architectures in addressing such challenges without relying on brute-force
techniques like data augmentation [40, 20].

In image classification tasks, steerable convolutions have proven to be a
powerful tool for achieving equivariance to the group $\mathrm{SE}(d)$, which
encompasses rotations and translations in $d$ dimensions [40, 10, 37].
Steerable convolution networks primarily operate in Fourier space, leveraging
representations of the underlying group to encode rotation and translation
equivariant features from the data [9, 8, 37]. Nevertheless, these
architectures, rooted in convolutions, retain the ability to learn local
neighborhoods. Conversely, transformer architectures excel in learning
relationships between different regions of an image. Our contribution lies in
integrating these two concepts—vision transformers and steerable
convolutions—yielding steerable transformers capable of effectively capturing
both local and global patterns in images.

In this paper we introduce a novel steerable transformer architecture for
image recognition. The architecture draws inspiration mainly from the works of
Dosovitskiy et al. [11] and Fuchs et al. [13]. Dosovitskiy et al. [11]
pioneered the use of transformers in vision tasks by segmenting images into
sequences. They also introduced a hybrid approach, incorporating learned
features from CNN layers instead of directly inputting images into
transformers. On the other hand, Fuchs et al. [13] introduced an attention-
based architecture for processing 3D point cloud data, integrating the self-
attention mechanism into Tensor Field Networks [32]. Similar to steerable CNNs
[9] and Tensor Field Networks [32], steerable transformers operate in Fourier
space, leveraging representations of the underlying symmetry group to learn
equivariant features. The self-attention mechanism in steerable transformers
is designed akin to that in [13]. Similar to [33], we use this kind of
architecture in addition with steerable convolutions[40, 37]. Other
transformer based equivariant architecture like [42, 27] base their design on
[24], where they have completely replaced the convolution architecture for an
attention based mechanism. Both Dosovitskiy et al. [11] and Fuchs et al. [13]
employ learnable positional encodings. While the seminal work by Vaswani et
al. [33] utilizes predefined functions for positional encoding, our approach
combines both strategies by multiplying a fixed function with a learnable
scalar. Unlike [42, 27], who employ invariant positional encoding, we use
equivariant positional encoding.

##  2 Related Work

The original idea of attention was introduced by Bahdanau et al. [2], for
sequence modeling. Vaswani et al. [33] introduced the transformer encoder-
decoder architecture for translation tasks in natural language processing.
Dosovitskiy et al. [11] proposed the transformer architecture for vision
tasks. However these transformer architectures are not linear in their inputs
and hence break the equivariance structure of traditional CNN. Since then
there has been a flurry of papers trying to make these structures equivariant.
While [11] advocated for a hybrid architecture where the features are
extracted from traditional CNN. Ramachandran et al. [24] introduced a stand
alone attention based architecture which completely replaces traditional CNNs.

Shaw et al. [29] proposed learnable relative positional encoding instead of
absolute positional encoding in the original work [33], to encode the relative
positions of input sequences. Romero et al. [27] showed that vision
transformers [11] with relative positional encoding is translation
equivariant. They extended the idea of [24] to incorporate rotation and
translation equivariance in the 2D regime, by lifting the attention mechanism
to be defined on the group $\mathrm{SE}(2)$. For this they proposed a
positional encoding operator and use regular representations of the group
instead of irreducible representations. Xu et al. [42] introduces a novel
positional encoding in the method proposed in [27] to enhance its performance.
Romero et al. [26] incorporated attention in convolution itself by defining
attention score on the equivariance groups.

In the 3D regime there have been efforts to use attention in vision task,
however to the best of our knowledge, only on point cloud data and not
volumetric data. Fuchs et al. [13] proposed an $\mathrm{SE}(3)$ equivariant
architecture for point cloud data by adding self-attention mechanism to the
rotation, translation and permutation equivariant Tensor Field network [32].
Chen at al. [7] achieved equivariance in attention mechanisms by mapping the
data onto $\mathrm{SE}(3)$ equivariant and invariant features. These maps are
based on the work of Villar et al. [34]. Hutchinson et al.
citephutchinson2021lietransformer introduced an attention mechanism that is
equivariant to the action of Lie groups using regular representations of the
group and built a transformer architecture based on that. Chatzipantazis et
al. [6] introduced $\mathrm{SE}(3)$ equivariant attention for shape
reconstruction tasks.

##  3 Background

In this section we introduce the relevant background required to design the
steerable transformer architecture.

###  3.1 The Multihead Self-Attention Mechanism

The attention mechanism, initially introduced by Bahdanau et al. [2] for
sequence-to-sequence tasks, involves three main components: query vectors
$q_{i}$, key vectors $k_{i}$, and value vectors $v_{i}$ for each input in a
sequence. These vectors are utilized to compute attention weights, denoted
$\alpha_{ij}$, which determine the relevance of each input vector in producing
the output. The score function $s(q_{i},k_{j})$ calculates the compatibility
between query and key vectors. The softmax of these scores yields the
attention weights $\alpha_{ij}$. Finally, we get the output by linearly
combining $\alpha_{ij}$ with the value vectors $v_{i}$:

$\textsc{Attn}(q_{i},\\{k_{i}\\},\\{v_{i}\\})=\sum_{j=1}^{N}\alpha_{ij}v_{j},\quad\quad\quad\quad\alpha_{ij}=\frac{\exp(s(q_{i},k_{j}))}{\sum_{j=1}^{N}\exp(s(q_{i},k_{j}))}.$

The score function $s(q_{i},k_{j})$ in the scaled dot-product attention
mechanism [33], for a fixed $i$, is calculated as:

$s(q_{i},k_{j})=\frac{q_{i}^{T}k_{j}}{\sqrt{d_{K}}}.$

Vaswani et al. [33] proposed a Multihead version of this attention mechanism
which involves stacking multiple attention mechanisms, concatenating their
outputs and applying a linear transformation to get the final result:

$\textsc{Multi-
Attn}(q_{i},\\{k_{i}\\},\\{v_{i}\\})=\text{Concat}(\textsc{Attn}^{1},\dots,\textsc{Attn}^{h})W_{O},$

where $W_{O}\in\mathbb{R}^{hd_{V}\times d_{\text{model}}}$. Here
$d_{\text{model}}$ is the dimension if the output which is same as the
dimension of the input vectors. This dimension is fixed through the
Transformer encoder structure. Typically, the query, key, and value vectors
are obtained through learnable linear embedding followed by adding _positional
encoding_.

#### Positional Encoding:

Positional encoding in transformers is a technique employed to incorporate
positional information about words or tokens into the input embeddings of a
sequence. In vision tasks, these encodings inform the network about the
original locations of different portions of the image. In the seminal work by
Vaswani et al. [33], they utilized sine and cosine functions of varying
frequencies for positional encoding. Relative positional encoding, introduced
by Shaw et al. [29], represents an extension or variation of traditional
positional encoding. Its goal is to use learnable position encoding that
capture the relative distances between input sequences. As we explain below,
relative positional encoding is crucial for achieving equivariance.

###  3.2 Vision Transformers

Dosovitskiy et al. [11] extended the use of the transformer architecture to
vision tasks. In their approach, they divided the image into equal-sized
portions and fed them into the encoder part of the transformer architecture.
They utilize the features obtained from the encoder for classification.
Assuming a single input channel, the vision transformer layer can be
summarized as follows:

$\displaystyle
z_{0}=[\textbf{x}_{1}E,\textbf{x}_{2}E,\dots,\textbf{x}_{N}E]+E_{\text{pos}}$
$\displaystyle E\in\mathbb{R}^{p\times d_{\text{model}}},\ \
E_{\text{pos}}\in\mathbb{R}^{N\times d_{\text{model}}}$ $\displaystyle
z_{\ell}^{\prime}=\textsc{Multi-Attn}(\textsc{LN}(z_{\ell-1}))+z_{\ell-1}$
$\displaystyle\ell=0,1,\dots L$ $\displaystyle
z_{\ell}^{\prime}=\textsc{MLP}(\textsc{LN}(z_{\ell-1}))+z^{\prime}_{\ell}$
$\displaystyle\ell=0,1,\dots L.$

Here, $\mathbf{x}_{i}\in\mathbb{R}^{p}$ represents a portion of the image
which serves as the input, $E$ is the learnable linear embedding matrix, and
$E_{\text{pos}}$ is the positional encoding. LN denotes a normalization layer,
and MLP stands for a multilayer perceptron. $d_{\text{model}}$ is the
dimension of the embedding fed into the transformer, which remains fixed
throughout the entire encoder layer. Typically, the MLP layer comprises two
hidden layers separated by a non-linearity [33, 11].

###  3.3 Equivariance

Equivariance refers to a property where a mapping between vector spaces
preserves the action of a group $G$. For each $g\in G$, a linear map
$T_{g}:\mathcal{V}\to\mathcal{V}$ is defined on a vector space $\mathcal{V}$.
Given actions $T_{g}$ and $T^{\prime}_{g}$ on vector spaces $\mathcal{V}$ and
$\mathcal{W}$ respectively, a map $\phi:\mathcal{V}\to\mathcal{W}$ is
considered equivariant if:

$\phi(T_{g}(\mathbf{v}))=T^{\prime}_{g}(\phi(\mathbf{v}))\quad\quad\quad\quad\mathbf{v}\in\mathcal{V}.$

Equivariance naturally emerges as a constraint in various real-world
scenarios. For instance, in Graph Neural Networks, equivariance to the
permutation group $S_{n}$ is fundamental [31, 35]. Similarly, in vision tasks,
equivariance to rotation and reflection symmetries, i.e., the action of the
Euclidean group, is desirable. There is now an extensive literature on
equivariant neural networks, in particular, focusing on cases where $\phi$ is
linear [40, 10, 37, 9, 8, 36, 16, 1].

###  3.4 Steerable Neural Networks

Steerable Networks are engineered to exhibit equivariance to the special
Euclidean group $\mathrm{SE}(d)$ in $d$ dimensions. Consider an
$d$-dimensional image with a single input channel, represented by the function
$f^{\textrm{in}}:\mathbb{R}^{d}\to\mathbb{R}$, which has compact support.
These networks employ steerable convolution layers, which can be expressed as
convolutions on the group $\mathrm{SE}(d)$ [37, 36]:

$\displaystyle
f^{\textrm{out}}(\textbf{x},R)=\int_{\mathrm{SO}(d)}\int_{\mathbb{R}}f^{\textrm{in}}(\textbf{x}+R\textbf{y})\hskip
1.00006ptw(\textbf{y})\,d\textbf{y}\,d\mu(R^{\prime}),$ (1) $\displaystyle
f^{\textrm{out}}(\textbf{x},R)=\int_{\mathrm{SO}(d)}\int_{\mathbb{R}}f^{\textrm{in}}(\textbf{x}+R\textbf{y},RR^{\prime})\hskip
1.00006ptw(\textbf{y},R^{\prime})\,d\textbf{y}\,d\mu(R^{\prime}).$ (2)

Here $\mu$ denotes the Haar measure on $\mathrm{SO}(d)$. Equation (1)
represents the formula for the first layer, while (2) represents subsequent
layers. The presence of two distinct formulas is due to the fact that the
input to the first layer is a vector field on $\mathbb{R}^{d}$, whereas the
input to subsequent layers is a vector field on $\mathrm{SE}(d)$. These linear
mappings possess a key property: whenever the input map $f^{\textrm{in}}$
undergoes a transformation by an element in $\mathrm{SE}(d)$, the resulting
output $f^{\textrm{out}}$ also undergoes the same transformation. This
property means that if under the action of the group $f^{\textrm{in}}$ changes
to

${f^{\textrm{in}}}^{\prime}(\textbf{x})=(t,R)\cdot
f^{\textrm{in}}(\textbf{x})=f^{\textrm{in}}(R\textbf{x}+t)$ (3)

then $f^{\textrm{out}}$ changes to

${f^{\textrm{out}}}^{\prime}(\textbf{x},R^{\prime})=(t,R)\cdot
f^{\textrm{out}}(\textbf{x},R^{\prime})=f^{\textrm{out}}(R\textbf{x}+t,RR^{\prime}).$
(4)

For compact groups like $\mathrm{SO}(d)$, where discretization is not
straightforward, it is advantageous to compute the rotation component of the
output feature map of (2) in Fourier space [40]. Leveraging the convolution
theorem, equation (2) in Fourier space becomes a product of the Fourier
transforms of the input and the weights. The Fourier transform of a function
$f$ on a compact group can be computed by integrating it against its
irreducible representations or _irreps_ (see Appendix for definition):

$\widehat{f}(\rho)=\int_{\mathrm{SO}(d)}f(R)\,\rho(R)\,d\mu(R).$

These Fourier space vector fields possess the property that under input
transformations like (3), the output transforms to
$\rho(R)\widehat{f^{\textrm{out}}}(R\mathbf{x}+t,\rho)$ [9, 37].

In the 2D case, the irreps of $\mathrm{SO}(2)$ are just complex exponential
indexed by integers, $\rho_{k}(\theta)=e^{\iota k\theta}$. Therefore, the
$k$’th frequency component of the output feature map is a function
$\widehat{f^{\textrm{out}}_{k}}:\mathbb{R}^{2}\to\mathbb{C}$. For the 3D case,
the irreps of $\mathrm{SO}(3)$ are indexed by integers $\ell\geq 0$,
represented by $\rho_{\ell}(R)=D^{\ell}(R)$, known as the Wigner matrices
[38], which are unitary matrices of dimension $2\ell+1$. Correspondingly, the
$\ell^{\text{th}}$ Fourier component of the output feature map is a function
$\widehat{f^{\textrm{out}}_{\ell}}:\mathbb{R}^{3}\to\mathbb{C}^{2\ell+1}$.

#### Non-linearities:

In steerable neural networks, utilizing non-linearities like ReLU in Fourier
space can disrupt equivariance. One approach to mitigate this issue is to
transition back to real space, apply the desired non-linearity there, and then
return to Fourier space. This method, for instance, is employed in [10].
However, this back-and-forth transformation between real and Fourier space can
be computationally expensive and error-prone, especially on groups like
$\text{SO}(3)$ where the grid is non-uniform.

As an alternative, Fourier space non-linearities are preferred in steerable
neural networks. These non-linearities act directly on Fourier space while
preserving equivariance. Several such non-linearities can be found in the
literature. For instance, Worall et al. [40] apply non-linearity to the norm
of the Fourier vector, as the norm remains invariant under rotations. Other
works such as [16, 1] utilize the Clebsch–Gordan non-linearity, which involves
a tensor product of Fourier vectors followed by a Clebsch-Gordan
decomposition. Weiler et al. [37] introduce another steerable convolution
filter to serve as a non-linearity. These methods enable effective non-linear
transformations in Fourier space while maintaining equivariance.

##  4 Method

Figure 1: The schematic illustrates the steerable attention mechanism (left)
and a steerable transformer encoder layer (right); c.f. Figure 1 in [11]). The
steerable attention mechanism is shown for a single head ($h=1$) and one query
dimension $(d_{K}=1)$. The stacks of feature maps represent different Fourier
components, and the _H-Nonlinearity_ denotes the harmonic non-linearity
described in [40].

The methodology behind steerable transformers builds on the foundation
established by vision transformers [11], which rely on linear embeddings of
fixed-size image patches along with positional encoding for these patches. A
hybrid architecture discussed in the same work integrates features maps
produced by CNNs as input to the transformer encoder. In the case of steerable
transformers, we adopt this hybrid approach, with feature maps generated by
steerable convolution layers serving as input to the transformer encoder.

In Section 3.4, we discussed that the output of steerable networks is a
compactly supported function on a continuous domain. However, practical neural
networks operate on finite resolution rasterized images. Consequently, in
real-world implementations, the output features for a particular irrep are
represented as a collection of complex vectors on a grid. For simplicity, let
us assume that there are $N$ such grid points. For $C$ channels, the input
sequence to the transformer is described by:

$f^{\textrm{in}}(\mathbf{x}_{i},\rho)\in\mathbb{C}^{d_{\rho}\times C},\quad
i=1,2,\dots,N,\quad k\in\mathbb{Z}.$ (5)

In two dimensions, $d_{\rho}$ is always $1$, while in three dimensions, the
$\ell^{\text{th}}$ component has $d_{\rho}=2\ell+1$.

###  4.1 Positional Encoding

Before we delve into the steerable version of self-attention, let us first
explore a steerable approach to positional encoding. In the seminal work by
Vaswani et al. [33], sine and cosine functions of varying frequencies were
utilized to encode positions. In contrast, Dosovitskiy et al. [11] employed
learnable positional embeddings for vision transformers. In our approach, we
adopt a hybrid method, by having both learnable and fixed components in
positional encoding.

Relative positional encoding, introduced by Shaw et al. [29], considers
positional encoding as learned functions of the relative distances between
input sequences. Romero et al. [27] demonstrated that self-attention
mechanisms with relative positional encoding exhibit translation equivariance.
Therefore, the positional embeddings we seek are functions
$P_{ij}^{\rho}=P(\mathbf{x}_{i}-\mathbf{x}_{j},\rho)$, inherently invariant to
translations.

To introduce the notion of steerability, we aim for $P_{ij}^{\rho}$ to
transform as $\rho(R)P_{ij}^{\rho}$ whenever the grid locations
$\mathbf{x}_{i}$ undergo a rotation $R$. Essentially we will be using the
Fourier basis to enforce the equivariance property (3), (4). In two
dimensions, this boils down to:

$P(\mathbf{x},k)=\phi(r,k)e^{-ik\theta},\quad\mathbf{x}=[r\cos\theta,r\sin\theta]^{T},\quad
k\in\mathbb{Z}.$ (6)

Here, $\phi(r,k)$ modulates the encoding strength for nearby and distant
points. In the 3D case, we have

$P(\mathbf{x},\ell)=\phi(r,l)Y^{\ell}(\theta,\phi),\quad\mathbf{x}=[r\sin\theta\cos\phi,r\sin\theta\sin\phi,r\cos\theta]^{T},\quad\ell\in\mathbb{Z}_{\geq
0},$ (7)

where $Y^{\ell}$ denotes the spherical harmonics [4]. The steerability
property of positional encoding implies $P_{ii}^{\rho}=0$ for any $\rho$
except for the constant representation, i.e., $k=0$ and $\ell=0$ in two and
three dimensions, respectively (as $P_{ii}^{\rho}=\rho(R)P_{ii}^{\rho}$ for
any $R$). This condition can be enforced by setting $\phi$ to zero when $r=0$.
For our experiments, we employed
$\phi(r,\rho)=w_{\rho}e^{-r^{2}}\mathbbm{1}_{r>0}$, where $w_{\rho}$ is a
learnable scalar. This kind of exponential modulation of the radial component
is also seen in [40, 37], and the resulting positional encoding assigns higher
weight to neighboring points, and the weights decrease as points move further
apart. In our actual implementation, we have different learnable scalars not
only for different Fourier components but also for different heads and query
dimensions.

###  4.2 Steerable Self-Attention Mechanism

The query, key, and value embeddings are obtained by multiplying the input
$f^{\textrm{in}}$ in (5) with learnable weight matrices. Additionally,
positional encoding is incorporated into the key embeddings [29].

$\displaystyle\mathbf{q}_{i}^{\rho}=f^{\textrm{in}}(\mathbf{x}_{i},\rho)\mathbf{W}_{Q}^{\rho}$
$\displaystyle\mathbf{W}_{Q}^{\rho}\in\mathbb{C}^{d_{\text{model}}\times
d_{K}}$
$\displaystyle\mathbf{k}_{ij}^{\rho}=f^{\textrm{in}}(\mathbf{x}_{i},\rho)\mathbf{W}_{K}^{\rho}+P(\mathbf{x}_{i}-\mathbf{x}_{j},\rho)$
$\displaystyle\mathbf{W}_{K}^{\rho}\in\mathbb{C}^{d_{\text{model}}\times
d_{K}}$
$\displaystyle\mathbf{v}_{i}^{\rho}=f^{\textrm{in}}(\mathbf{x}_{i},\rho)\mathbf{W}_{V}^{\rho}$
$\displaystyle\mathbf{W}_{V}^{\rho}\in\mathbb{C}^{d_{\text{model}}\times
d_{V}}$

Next, the score function for each Fourier component is calculated by taking
the dot product between the query and key embeddings. The individual scores
from each Fourier component can be mixed together using scalar weights. The
attention scores are obtained by applying the softmax function to the raw
scores. Since we are dealing with complex numbers, we use the absolute value
of these scores:

$\displaystyle
s(\mathbf{q}_{j}^{\rho},\mathbf{k}_{ij}^{\rho})=\frac{\text{vec}(\mathbf{q}_{j}^{\rho})^{\dagger}\text{vec}(\mathbf{k}_{ij}^{\rho})}{\sqrt{d_{K}}}\qquad
s_{ij}^{\rho}=\sum_{\rho^{\prime}}w^{(1)}_{\rho\rho^{\prime}}s_{ij}^{\rho^{\prime}}\qquad\alpha_{ij}^{\rho}=\frac{\exp(|s_{ij}^{\rho}|)}{\sum_{j^{\prime}=1}^{N}|s_{ij^{\prime}}^{\rho}|}$

Here, the _vec_ operation flattens the matrix into a vector. Since the vectors
are complex-valued, the dot product involves the conjugate transpose, denoted
by $\cdot^{\dagger}$, instead of a regular transpose. Finally, the output for
each position and Fourier component is computed as a weighted sum of the value
vectors using the attention scores, and then combined again using scalar
weights:

$f^{\textrm{out}}(\mathbf{x}_{i},\rho)=\sum_{\rho^{\prime}}\sum_{j=1}^{N}w^{(2)}_{\rho\rho^{\prime}}\alpha_{ij}^{\rho^{\prime}}\mathbf{v}_{j}^{\rho}.$

The scalar weights
$w^{(1)}_{\rho\rho^{\prime}},w^{(2)}_{\rho\rho^{\prime}}\in\mathbb{C}$ can be
fixed, learnable, or a combination of both, offering a variety of models to
work with. This self-attention mechanism maintains equivariance, as
demonstrated by the following proposition.

######  Proposition 1.

If the input to a steerable transformer self-attention layer transforms as
$f^{\textrm{in}}(\mathbf{x}_{i},\rho)\mapsto\rho(R)\,f^{\textrm{in}}(R\mathbf{x}_{i}+\mathbf{t},\rho)$,
then the output also transforms as
$f^{\textrm{out}}(\mathbf{x}_{i},\rho)\mapsto\rho(R)\,f^{\textrm{out}}(R\mathbf{x}_{i}+\mathbf{t},\rho)$.

Despite the softmax operation’s lack of equivariance, in this scenario, the
dot product structure of the score function maintains rotational invariance,
thereby preserving equivariance. For multihead attention, this mechanism is
repeated independently for $h$ heads. Finally, the outputs are concatenated,
and each Fourier component is scaled by another matrix
$\mathbf{W}_{O}^{\rho}\in\mathbb{C}^{hd_{V}\times d_{\text{model}}}$.
Consistent with [33], we choose $d_{V}=d_{K}=d_{\text{model}}/h$.

###  4.3 Position-wise Feed Forward Layer

In a transformer architecture, the self-attention layer is typically followed
by a multilayer perceptron (MLP) layer. This MLP layer commonly consists of
two linear layers separated by a non-linearity [33]. The linear layers
themselves do not affect equivariance as they are applied to the input from
the right:

$f^{\textrm{out}}(\mathbf{x}_{i},\rho)=\sigma\left(f^{\textrm{in}}(\mathbf{x}_{i},\rho)W_{1}\right)W_{2}\quad\quad\text{where}\quad\quad
W_{1},W_{2}^{T}\in\mathbb{C}^{d_{\text{model}}\times d_{\text{hidden}}}.$

Here, $d_{\text{hidden}}$ represents the hidden dimension between the two
linear layers. However, when it comes to the non-linearity, caution is needed.
As discussed in Section 3.4, typical non-linearities like ReLU or softmax may
break equivariance. Instead, we require a Fourier space non-linearity. For our
experiments, we utilize a non-linearity proposed by Worrall et al. [40], where
they apply ReLU to the magnitude of the equivariant vectors:

$\sigma(f(\mathbf{x},\rho))=\frac{\text{ReLU}(||f(\mathbf{x},\rho)||+b)}{||f(\mathbf{x},\rho)||}\>f(\mathbf{x},\rho).$

Here, $b$ is a learnable weight. Since the norm of the features in Fourier
space is invariant to rotations, this non-linearity preserves equivariance.
Following the convention in [33], we set
$d_{\text{hidden}}=2d_{\text{model}}$.

###  4.4 Layer Normalization

Another crucial component of the transformer architecture is layer
normalization. Notably, the norm of equivariant vectors remains invariant
under the Fourier transform, as it preserves the $\mathcal{L}_{2}$ norm.
Leveraging this property, we can derive a steerable normalization as follows:

$\text{LN}(f)(\mathbf{x},\rho)=\frac{f(\mathbf{x},\rho)}{\sum_{\rho}||f(\mathbf{x},\rho)||^{2}}\>.$

Indeed, the sum over $\rho$ theoretically extends over all irreps of the
group, but in practical implementations, it is truncated at particular values,
which serves as a hyperparameter that can be tuned. This type of normalization
was employed in the 2D context by Worrall et al. [40].

Datasets | Model |  Frequncy Cutoff | Accuracy |  Parameter ($\sim\times 10^{6}$)
---|---|---|---|---
Rotated MNIST |  Steerable Convolution | $k=4$ | $98.72_{\pm 0.02}$ | $1.98$
$k=8$ | $98.97_{\pm 0.01}$ | $3.94$
Steerable Convolution (Without Transformer) | $k=4$ | $98.70_{\pm 0.05}$ | $0.46$
$k=8$ | $98.97_{\pm 0.05}$ | $0.92$
Steerable Transformer | $k=4$ | $\mathbf{98.82}_{\pm 0.04}$ | $1.13$
$k=8$ | $\mathbf{99.03}_{\pm 0.04}$ | $2.24$
ModelNet10 ($z$ Rotation) ($\mathrm{SO}(3)$ Rotation) |  Steerable Convolution | $\ell=4$ | $90.13_{\pm 0.52}$ | $1.08$
$86.62_{\pm 0.25}$
Steerable Convolution (Without Transformer) | $\ell=4$ | $89.51_{\pm 0.55}$ | $0.48$
$86.58_{\pm 0.92}$
Steerable Transformer | $\ell=4$ | $\mathbf{90.40}_{\pm 0.25}$ | $0.92$
$\mathbf{86.80}_{\pm 0.58}$
Table 1: Comparison of steerable transformers and steerable convolutions. For
each dataset, the first row features an architecture with only steerable
convolutions and an increased number of channels to match the parameter count
of the steerable transformers. The second row presents the same architecture
as the steerable transformer but without the transformer encoder layer. The
third row displays the steerable transformer architecture. The radial
resolution was fixed for these experiments. A more detailed comparison is
available in the Appendix Table 4. The mean and sd are reported for $5$ runs.
For ModelNet10 we have reported both the $z$ rotation and $\mathrm{SO}(3)$
rotation variations.

##  5 Experiments

We evaluate the performance of our steerable transformer architecture in both
two and three dimensions. For all experiments, we used a hybrid architecture
combining steerable convolutions and steerable transformer encoder layers,
trained using the Adam Optimizer [15]. The architecture consists of steerable
CNN blocks and a steerable transformer encoder. Each steerable convolution
block includes two steerable convolutions separated by a Clebsch–Gordan
nonlinearity [1, 16] and followed by layer normalization, as detailed in
Section 4.4. After a convolution block, a transformer encoder is applied. This
pattern is repeated twice, followed by a final convolution block with a kernel
size equal to the input size, downsampling the image to one pixel per Fourier
component and channel. This step is followed by taking the norm, ensuring
rotational and translational invariance due to the properties of the Fourier
transform. The final convolution is followed by two fully connected layers
with a ReLU activation and a dropout layer, using a dropout probability of 0.7
as in [20].

We present results for two datasets: Rotated MNIST (2D) and ModelNet10 (3D).
Table 1 provides a detailed comparison of the steerable architecture with and
without transformers. Results are reported for three configurations: one with
steerable transformers, one with the same architecture minus the transformer
encoder layers, and another with steerable convolutions but with additional
channels to match the parameter count of the steerable transformers.

###  5.1 Rotated MNIST

Method | error%
---|---
$\alpha$-R4 CNN[26] | $1.69$
GSA-Nets[27] | $2.03$
GE-ViT[42] | $1.99$
Our ($k=4$) | ${1.18}$
Our ($k=8$) | $\mathbf{0.97}$
Table 2: Comparison of the performance of the equivariant attentive
architecture on the Rotated MNIST dataset.

Rotated MNIST, a variant of the original MNIST dataset [17], consists of
images subjected to random rotations and includes 12,000 training images and
50,000 testing images. Details on the architecture and training can be found
in the Appendix.

Table 1 compares the performance of steerable transformers and steerable
convolutions. It clearly demonstrates that adding a transformer encoder to
steerable convolutions improves performance, even with a low Fourier cutoff.
Simply increasing the number of parameters in convolutions by adding more
channels does not improve performance and likely results in overfitting. Table
2 presents a comparison of steerable transformers with other attention-based
methods reported for this dataset. Our approach shows significantly superior
performance compared to the others.

Figure 2: The figure demonstrates the equivariance of attention scores in a
trained steerable transformer. For a fixed pixel $i$, we have plotted the
maximum attention score for that pixel $(\max_{j}\alpha_{ij})$. The different
subfigures represent individual heads. The first and last heads appear to
capture the object’s boundary, while the other two heads focus on the object’s
body.

###  5.2 ModelNet10

ModelNet10 [41] is a dataset comprising 3D CAD models representing 10 common
object categories with a train:test split of $3991$:$908$. Each category
contains a collection of triangular meshes representing objects in various
orientations and scales, facilitating the training and evaluation of
algorithms for 3D object recognition and classification tasks. With a
predefined split into training and testing sets, ModelNet10 serves as a
benchmark dataset for assessing the generalization performance of algorithms.
For the experiments, $2,048$ points from the shape surface are uniformly
sampled using the farthest point sampling algorithm. The point cloud data was
then embedded in a grid to get voxel data (see Figure 3 in Appendix).
Additionally, we considered two kinds of perturbations of this dataset : the
images are perturbed by a random rotation along the "gravity up" direction
($z$ rotation) and also random rotations from $\mathrm{SO}(3)$. We also
rescale the voxel values to $[-1,5]$ instead of $[0,1]$ as in [3], who showed
it helps with sparse voxel volumes. The details of architecture and training
are provided in the Appendix.

Table 1 shows the comparison of the steerable transformer architecture with
steerable convolutions for a frequency cutoff of $\ell=4$. The results
indicate a significant improvement due to the introduction of the steerable
transformer encoder layer across both variations of the dataset. Table 3
presents a comparison with other methods (not attention-based) on the
$z$-rotated version of the dataset. Unlike our approach, which does not use
train-time augmentation, all other methods are trained with augmentation
involving 12 uniformly stratified rotations along the $z$-axis. For a fair
comparison, we applied test-time augmentation by averaging the accuracy of all
12 predictions for each test data point. Despite not using train-time
augmentation and having fewer parameters, our method achieves performance
comparable to the other methods.

Figure 3: Examples from the ModelNet10 dataset are shown in various formats:
point cloud, voxel representation, and rotated perturbations of voxels.

| Method | Accuracy |  Parameters ($\sim\times 10^{6}$)
---|---|---|---
Point Cloud | ECC [30] | 90.8 | -
SO-Net [18] | 93.9 | 2.5
Rot-SO-Net [19] | 94.5 | 2.5
Voxel | 3D ShapeNets[41] | 83.5 | 12
VRN [3] | 91.3 | 18
VoxNet [21] | 92.0 | 0.92
FusionNet [14] | 93.1 | 120
ORION [28] | 93.8 | 0.91
Cubenet [39] | 94.6 | 4.5
Our | 91.1 | 0.92

Table 3: Comparison of performance on ModelNet10 with $z$ rotation
perturbation. Other methods used both train and test time augmentation, while
we applied augmentation only during test time, not during training.

##  6 Limitations

The use of transformer architecture comes with its challenges. For instance,
the dot product structure used in attention scores can create a memory
bottleneck for large images, restricting the network size and batch size
during training. Scaling up these architectures is essential to fully
understand the benefits of the attention mechanism in improving performance.
Our experiments on Rotated MNIST and ModelNet10 were conducted with relatively
small Fourier cutoffs ($k=8$ for Rotated MNIST and $\ell=4$ for ModelNet10).
We believe that using higher values for these Fourier cutoffs would enhance
performance, but implementing them in a steerable transformer architecture
would necessitate scaling up to accommodate larger models.

##  7 Conclusions

The transformer framework is increasingly being adopted across various machine
learning fields, including computer vision, due to its exceptional ability to
capture relationships across different parts of the input. In this work, we
introduce steerable transformer encoders that can be used in conjunction with
steerable convolutions. While convolutions excel at capturing local
relationships, incorporating the attention mechanism, which captures global
relationships, enhances performance, as demonstrated in our experiments. We
believe that this architecture has the potential to make a significant impact
in fields such as medical imaging, where equivariance is critical.

## References

  * Anderson et al. [2019] B. Anderson, T. S. Hy, and R. Kondor.  Cormorant: Covariant molecular neural networks.  _Advances in neural information processing systems_ , 32, 2019.
  * Bahdanau et al. [2014] D. Bahdanau, K. Cho, and Y. Bengio.  Neural machine translation by jointly learning to align and translate.  _CoRR_ , abs/1409.0473, 2014.  URL https://api.semanticscholar.org/CorpusID:11212020.
  * Brock et al. [2016] A. Brock, T. Lim, J. M. Ritchie, and N. Weston.  Generative and discriminative voxel modeling with convolutional neural networks.  _ArXiv_ , abs/1608.04236, 2016.  URL https://api.semanticscholar.org/CorpusID:2271164.
  * Byerly [1893] W. E. Byerly.  _An elemenatary treatise on Fourier’s series, and spherical, cylindrical, and ellipsoidal harmonics, with applications to problems in mathematical physics_.  Dover Publicatiions, 1893.
  * Carion et al. [2020] N. Carion, F. Massa, G. Synnaeve, N. Usunier, A. Kirillov, and S. Zagoruyko.  End-to-end object detection with transformers.  In _European conference on computer vision_ , pages 213–229. Springer, 2020.
  * Chatzipantazis et al. [2022] E. Chatzipantazis, S. Pertigkiozoglou, E. Dobriban, and K. Daniilidis.  Se (3)-equivariant attention networks for shape reconstruction in function space.  _arXiv preprint arXiv:2204.02394_ , 2022.
  * Chen and Villar [2022] N. Chen and S. Villar.  Se (3)-equivariant self-attention via invariant features.  In _Machine Learning for Physics NeurIPS Workshop_ , 2022.
  * Cohen and Welling [2016] T. Cohen and M. Welling.  Group equivariant convolutional networks.  In _International conference on machine learning_ , pages 2990–2999. PMLR, 2016.
  * Cohen and Welling [2017] T. S. Cohen and M. Welling.  Steerable CNNs.  In _International Conference on Learning Representations_ , 2017.  URL https://openreview.net/forum?id=rJQKYt5ll.
  * Cohen et al. [2018] T. S. Cohen, M. Geiger, J. Koehler, and M. Welling.  Spherical CNNs.  In _International Conference on Learning Representations_ , 2018.
  * Dosovitskiy et al. [2020] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby.  An image is worth 16x16 words: Transformers for image recognition at scale.  _ArXiv_ , abs/2010.11929, 2020.  URL https://api.semanticscholar.org/CorpusID:225039882.
  * Dwivedi and Bresson [2020] V. P. Dwivedi and X. Bresson.  A generalization of transformer networks to graphs.  _arXiv preprint arXiv:2012.09699_ , 2020.
  * Fuchs et al. [2020] F. Fuchs, D. Worrall, V. Fischer, and M. Welling.  Se (3)-transformers: 3d roto-translation equivariant attention networks.  _Advances in neural information processing systems_ , 33:1970–1981, 2020.
  * Hegde and Zadeh [2016] V. Hegde and R. Zadeh.  Fusionnet: 3d object classification using multiple data representations.  _arXiv preprint arXiv:1607.05695_ , 2016.
  * Kingma and Ba [2014] D. P. Kingma and J. Ba.  Adam: A method for stochastic optimization.  _CoRR_ , abs/1412.6980, 2014.  URL https://api.semanticscholar.org/CorpusID:6628106.
  * Kondor et al. [2018] R. Kondor, Z. Lin, and S. Trivedi.  Clebsch–gordan nets: a fully fourier space spherical convolutional neural network.  In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, _Advances in Neural Information Processing Systems 31_ , pages 10117–10126. Curran Associates, Inc., 2018.
  * LeCun and Cortes [2005] Y. LeCun and C. Cortes.  The mnist database of handwritten digits.  2005\.  URL https://api.semanticscholar.org/CorpusID:60282629.
  * Li et al. [2018] J. Li, B. M. Chen, and G. H. Lee.  So-net: Self-organizing network for point cloud analysis.  In _2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 9397–9406, 2018.  doi: 10.1109/CVPR.2018.00979.
  * Li et al. [2019] J. Li, Y. Bi, and G. H. Lee.  Discrete rotation equivariance for point cloud recognition.  In _2019 International conference on robotics and automation (ICRA)_ , pages 7269–7275. IEEE, 2019.
  * Marcos et al. [2017] D. Marcos, M. Volpi, N. Komodakis, and D. Tuia.  Rotation equivariant vector field networks.  In _Proceedings of the IEEE International Conference on Computer Vision_ , pages 5048–5057, 2017.
  * Maturana and Scherer [2015] D. Maturana and S. Scherer.  Voxnet: A 3d convolutional neural network for real-time object recognition.  In _2015 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)_ , pages 922–928, 2015.  doi: 10.1109/IROS.2015.7353481.
  * Munkres [1974] J. Munkres.  _Topology; a First Course_.  Prentice-Hall, 1974.  ISBN 9780139254956.  URL https://books.google.com/books?id=LtEPAQAAMAAJ.
  * Peter and Weyl [1927] F. Peter and H. Weyl.  Die vollständigkeit der primitiven darstellungen einer geschlossenen kontinuierlichen gruppe.  _Mathematische Annalen_ , 97(1):737–755, 1927.
  * Ramachandran et al. [2019] P. Ramachandran, N. Parmar, A. Vaswani, I. Bello, A. Levskaya, and J. Shlens.  Stand-alone self-attention in vision models.  _Advances in neural information processing systems_ , 32, 2019.
  * Robert [1983] A. Robert.  _Introduction to the representation theory of compact and locally compact groups_ , volume 80.  Cambridge University Press, 1983.
  * Romero et al. [2020] D. Romero, E. Bekkers, J. Tomczak, and M. Hoogendoorn.  Attentive group equivariant convolutional networks.  In _International Conference on Machine Learning_ , pages 8188–8199. PMLR, 2020.
  * Romero and Cordonnier [2020] D. W. Romero and J.-B. Cordonnier.  Group equivariant stand-alone self-attention for vision.  _ArXiv_ , abs/2010.00977, 2020.  URL https://api.semanticscholar.org/CorpusID:222125298.
  * Sedaghat et al. [2016] N. Sedaghat, M. Zolfaghari, E. Amiri, and T. Brox.  Orientation-boosted voxel nets for 3d object recognition.  _arXiv preprint arXiv:1604.03351_ , 2016.
  * Shaw et al. [2018] P. Shaw, J. Uszkoreit, and A. Vaswani.  Self-attention with relative position representations.  In _North American Chapter of the Association for Computational Linguistics_ , 2018.  URL https://api.semanticscholar.org/CorpusID:3725815.
  * Simonovsky and Komodakis [2017] M. Simonovsky and N. Komodakis.  Dynamic edge-conditioned filters in convolutional neural networks on graphs.  In _2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 29–38, 2017.  doi: 10.1109/CVPR.2017.11.
  * Thiede et al. [2020] E. H. Thiede, T. S. Hy, and R. Kondor.  The general theory of permutation equivarant neural networks and higher order graph variational encoders.  _arXiv preprint arXiv:2004.03990_ , 2020.
  * Thomas et al. [2018] N. Thomas, T. E. Smidt, S. M. Kearnes, L. Yang, L. Li, K. Kohlhoff, and P. F. Riley.  Tensor field networks: Rotation- and translation-equivariant neural networks for 3d point clouds.  _ArXiv_ , abs/1802.08219, 2018.  URL https://api.semanticscholar.org/CorpusID:3457605.
  * Vaswani et al. [2017] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.
  * Villar et al. [2021] S. Villar, D. W. Hogg, K. Storey-Fisher, W. Yao, and B. Blum-Smith.  Scalars are universal: Equivariant machine learning, structured like classical physics.  _Advances in Neural Information Processing Systems_ , 34:28848–28863, 2021.
  * Wang et al. [2020] R. Wang, M. Albooyeh, and S. Ravanbakhsh.  Equivariant networks for hierarchical structures.  In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, _Advances in Neural Information Processing Systems_ , volume 33, pages 13806–13817. Curran Associates, Inc., 2020.  URL https://proceedings.neurips.cc/paper_files/paper/2020/file/9efb1a59d7b58e69996cf0e32cb71098-Paper.pdf.
  * Weiler and Cesa [2019] M. Weiler and G. Cesa.  General e (2)-equivariant steerable cnns.  _Advances in neural information processing systems_ , 32, 2019.
  * Weiler et al. [2018] M. Weiler, M. Geiger, M. Welling, W. Boomsma, and T. S. Cohen.  3d steerable cnns: Learning rotationally equivariant features in volumetric data.  _Advances in Neural Information Processing Systems_ , 31, 2018.
  * Wigner [2012] E. Wigner.  _Group theory: and its application to the quantum mechanics of atomic spectra_ , volume 5.  Elsevier, 2012.
  * Worrall and Brostow [2018] D. E. Worrall and G. J. Brostow.  Cubenet: Equivariance to 3d rotation and translation.  In _European Conference on Computer Vision_ , 2018.  URL https://api.semanticscholar.org/CorpusID:4795882.
  * Worrall et al. [2017] D. E. Worrall, S. J. Garbin, D. Turmukhambetov, and G. J. Brostow.  Harmonic networks: Deep translation and rotation equivariance.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 5028–5037, 2017.
  * Wu et al. [2015] Z. Wu, S. Song, A. Khosla, F. Yu, L. Zhang, X. Tang, and J. Xiao.  3d shapenets: A deep representation for volumetric shapes.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 1912–1920, 2015.
  * Xu et al. [2023] R. Xu, K. Yang, K. Liu, and F. He.  $e(2)$-equivariant vision transformer.  In _Uncertainty in Artificial Intelligence_ , pages 2356–2366. PMLR, 2023.

##  Appendix A Background

###  A.1 Groups

#### Definition :

A group is a non-empty set $G$ together with a binary operation (commonly
denoted by “$\cdot$”, that combines any two elements $a$ and $b$ of $G$ to
form an element of $G$, denoted $a\cdot b$, such that the following three
requirements, known as group axioms, are satisfied:

  * •

_Associativity_ : For all $a,b,c\in G$, one has $(a\cdot b)\cdot
c=a\cdot(b\cdot c)$

  * •

_Identity Element_ : There exists as element $e\in G$ such that for every
$a\in G$, one has $e\cdot a=a\cdot e=a$. Such an element is unique. It is
called the _identity element_ of the group.

  * •

_Inverse_ : For each $a\in G$, there exists an element $b\in G$ such that
$a\cdot b=b\cdot a=e$, where $e$ is the identity element. For each $a$, the
element $b$ is unique and it is called the inverse of $a$ and is commonly
denoted by $a^{-1}$.

Examples include set of integers $\mathbb{Z}$ with the addition operation and
the set of non-zero reals $\mathbb{R}\setminus\\{0\\}$ with the multiplication
operation. From here on we will drop “$\cdot$”, for simplicity. The group
operation will be clear from the elements of the group concerned.

#### Group Homomorphism :

Given two groups $G$ and $H$, a function $\phi:G\to H$ is called a group
homomorphism if $\phi(ab)=\phi(a)\phi(b)$ for any $a,b\in G$. If the map
$\phi$ is a bijection, it is called an _isomorphism_ and furthermore, if
$G=H$, then an isomorphism is called _automorphism_. The set of all
automorphisms of a group $G$ with the operation of composition form a group in
itself and is denoted by $\text{Aut}(G)$.

#### Compact Groups :

A topological group is a topological space that is also a group such that the
group operation and the inverse map are continuous. A compact group is a
topological group whose topology realizes it as a compact topological space
(see [22] for definition of topological and compact topological spaces). Some
classic examples of compact groups are the groups $\mathrm{SO}(d)$ (the group
of all real orthogonal matrices in $d$ dimensions with determinant 1),
$\text{U}(d)$ (the group of all complex unitary matrices) and $\text{SU}(d)$
(the group of all complex unitary matrices with determinant 1).

_Special Orthonormal Group $\mathrm{SO}(d)$:_ This group comprises all real
orthogonal matrices in $d$ dimensions with a determinant of $1$. These groups
are associated with rotation matrices in $d$ dimensions. Specifically, we are
focused on the groups $\mathrm{SO}(2)$ and $\mathrm{SO}(3)$ in $2$ and $3$
dimensions, respectively. $\mathrm{SO}(2)$ can be parametrized by a single
angle $\theta$, corresponding to the rotation matrix:

$R(\theta)=\begin{bmatrix}\cos(\theta)&\sin(\theta)\\\
-\sin(\theta)&\cos(\theta)\end{bmatrix}$

Here, $R(\theta)$ signifies rotation in the $x$-$y$ plane by an angle
$\theta\in[0,2\pi)$. The group $\mathrm{SO}(3)$ can be parameterized by the
so-called Euler angles. In our discussions, we adopt the $z$-$y$-$z$
convention:

$R_{z}(\alpha)=\begin{bmatrix}\cos(\alpha)&0&\sin(\alpha)\\\ 0&1&0\\\
-\sin(\alpha)&0&\cos(\alpha)\end{bmatrix}\quad
R_{y}(\beta)=\begin{bmatrix}\cos(\beta)&-\sin(\beta)&0\\\
\sin(\beta)&\cos(\beta)&0\\\ 0&0&1\end{bmatrix}$

Here, $R_{z}(\alpha)$ represents rotation about the $z$-axis by $\alpha$
radians, and $R_{y}(\beta)$ represents rotation about the $y$-axis by $\beta$
radians. An element $R\in\mathrm{SO}(3)$ can then be expressed as
$R=R_{z}(\alpha)R_{z}(\beta)R_{y}(\gamma)$, where $\alpha,\gamma\in[0,2\pi)$
and $\beta\in[0,\pi)$ are the Euler angles.

#### Semi-direct Product Groups :

Given two groups $N$ and $H$ and a group homomorphism
$\phi:H\to\text{Aut}(N)$, we can construct a new group $N\rtimes_{\phi}H$
defined as follows,

  * •

The underlying set is the Cartesian product $N\times H$

  * •

The group operation is given by
$(n_{1},h_{1})(n_{2},h_{2})=(n_{1}\phi_{h_{1}}(n_{2}),h_{1}h_{2})$

_Special Euclidean Group ${\text{SE}(d)}$ :_ The special Euclidean group
$\text{SE}(d)$ encompasses all rotations and translations in $d$ dimensions.
The translation group for an $d$ dimensional real vector is isomorphic to
$\mathbb{R}^{n}$ and the rotation group is given by $\mathrm{SO}(n)$. Then, we
can define a group homomorphism
$\phi:\mathrm{SO}(d)\to\text{Aut}(\mathbb{R}^{d})$, as $[\phi(R)](t)=Rt$, for
any rotation $R\in\mathrm{SO}(d)$ and a translation vector
$t\in\mathbb{R}^{d}$. Then, the special Euclidean group $\text{SE}(d)$ is
defined as $\mathbb{R}^{d}\rtimes_{\phi}\mathrm{SO}(d)$. Hence the group
action is given by

$(t_{1},R_{1})(t_{2},R_{2})=(t_{1}+R_{1}t_{2},R_{1}R_{2})$

for any $(t_{1},R_{1}),(t_{2},R_{2})\in\mathrm{SE}(d)$, the identity is given
by $(0,I)$, and the inverse is given by, $(t,R)^{-1}=(-R^{-1}t,R^{-1})$ for
any $(t,R)\in\mathrm{SE}(d)$.

###  A.2 Group Actions

#### Definition :

If $G$ is a group with identity element $e$, and $X$ is a set, then a (left)
group action $\alpha$ of $G$ on $X$ is a function $\alpha:G\times X\to X$,
that satisfies the following two axioms for all $g,h\in G$ and $x\in X$:

  * •

_Identity_ : $\alpha(e,x)=x$,

  * •

_Compatibility_ : $\alpha(g,\alpha(h,x))=\alpha(gh,x)$.

Often $\alpha(g,x)$ is shortened to $g\cdot x$. Any group $G$ acts on itself
by the group operation $g\cdot h=gh$ for any $g,h\in G$. If $G$ acts on $X$,
then it can also act on any function $f$ defined on $X$, as $(g\cdot
f)(x)=f(g^{-1}\cdot x)$.

_Action of $\mathrm{SE}(d)$:_ The special Euclidean group acts on a vector in
$\mathbb{R}^{d}$ by first applying the rotation component followed by
translation. For $x\in\mathbb{R}^{d}$ and $(t,R)\in\mathrm{SE}(d)$

$(t,R)\cdot x=Rx+t,$

gives us the action of $\mathrm{SE}(d)$ on $\mathbb{R}^{d}$.

###  A.3 Group Representations

#### Definition :

A representation of a group $G$ is a group homomorphism $\rho$ from $G$ to
$\text{GL}(\mathbb{C}^{n})$ (group of invertible linear maps on
$\mathbb{C}^{n}$), i.e,

$\rho(g_{1}g_{2})=\rho(g_{1})\rho(g_{2})\quad\forall g_{1},g_{2}\in G$

Here $n$ is called the dimension of the representation, which can possibly be
infinite. A representation is _unitary_ if $\rho$ maps to unitary linear
transformation of $\mathbb{C}^{n}$.

#### Irreducible Representations :

If we have two representations, $\rho_{1}$ and $\rho_{2}$ of dimensions
$n_{1}$ and $n_{2}$ respectively, then the two can be combined by a direct sum
to give another representation of dimension $n_{1}+n_{2}$,

$\rho_{1}(g)\oplus\rho_{2}(g)=\begin{bmatrix}\rho_{1}(g)&0\\\
0&\rho_{2}(g)\end{bmatrix}.$

A representation is said to be _completely reducible_ if it can be expressed
as a direct sum of other representations after maybe a change of basis, i.e,

$U\rho(g)U^{-1}=\bigoplus_{i}\rho_{i}(g)$

where $U$ is a unitary change of basis matrix and the direct sum extends over
some number of representations. However, for _every_ group there are a some
representations which cannot be broken further into a direct sum of other
representations. These are called the _irreducible representations_ or
_irreps_ of the group. These irreps are the building blocks of the all other
representations of the group, in the sense that any representation can be
written as a direct sum of the irreps:

$\rho(g)=U\left[\bigoplus_{i}\rho^{(i)}(g)\right]U^{-1},$

where again $U$ is a change of basis matrix and $\rho^{(i)}$ are the irreps.
The Peter-Weyl Theorem by [23] tells us that for a compact group $G$, any
unitary representation $\rho$ is completely reducible and splits into direct
sum of irreducible _finite dimensional unitary representations_ of $G$.

#### Irreducible Representations of SO(2) and SO(3) :

$\mathrm{SO}(n)$ being a compact group, all its irreps are finite dimensional
unitary representations. The irreps of SO(2) are all of dimension one, and are
indexed by the set of integers. The group SO(2) can be realized by $2\times 2$
rotation matrices, with rotation angle $\theta$:

$R_{\theta}=\begin{bmatrix}\cos\theta&\sin\theta\\\
-\sin\theta&\cos\theta\end{bmatrix}.$

With this identification of SO(2), the irreducible representations of SO(2)
are given by,

$\rho_{k}(\theta)=e^{ik\theta}\qquad k\in\mathbb{Z}.$

The irreps correspond to the Fourier basis over $\mathbb{R}$. The SO(3) case
is more involved. The irreducible representations are indexed by
$\ell\in\\{0,1,2,\dots\\}$, where the $\ell$’th representation is of dimension
$2\ell+1$:

$\rho_{\ell}(R)=D^{\ell}(R)\qquad\ell\in\\{0,1,2,\dots\\},$

where $D^{\ell}$ are unitary matrices called Wigner D-matrices.

###  A.4 Fourier Transform

#### Haar Measure :

There is, up to a positive multiplicative constant, a unique countably
additive, nontrivial measure $\mu$ on the Borel subsets of $G$ satisfying the
following properties:

  * •

The measure $\mu$ is left-translation-invariant:
${\displaystyle\mu(gS)=\mu(S)}$ for every $g\in G$ and all Borel sets
${\displaystyle S\subseteq G}$.

  * •

The measure $\mu$ is finite on every compact set:
${\displaystyle\mu(K)<\infty}$ for all compact ${\displaystyle K\subseteq G}$.

  * •

The measure $\mu$ on Borel sets ${\displaystyle S\subseteq G}$ is given by

${\displaystyle\mu(S)=\inf\\{\mu(U):S\subseteq U,U{\text{ open}}\\}.}$

  * •

The measure $\mu$ is inner regular on open sets ${\displaystyle U\subseteq
G}$:

${\displaystyle\mu(U)=\sup\\{\mu(K):K\subseteq U,K{\text{ compact}}\\}.}$

Such a measure on G is called a left Haar measure. It can be shown that as a
consequence of the above properties that ${\displaystyle\mu(U)>0}$ for every
non-empty open subset ${\displaystyle U\subseteq G}$. In particular, if $G$ is
compact then $\mu(G)$ is finite and positive, so we can uniquely specify a
left Haar measure on G by adding the normalization condition
${\displaystyle\mu(G)=1}$.

#### Fourier Transform on Compact groups :

A nice property of compact groups is that the set of (isomorphism classes of)
irreducible representations of $G$ is countable (for details see [25]). If we
have a complex valued function supported on a compact group $G$,
$f:G\to\mathbb{C}$, then the Fourier transform of $f$ is given by

$\hat{f}(\rho)=\frac{1}{\mu(G)}\int_{G}f(g)\rho(g)d\mu(G),$

where $\mu$ is the Haar measure.

##  Appendix B Proof

###### Proof of Proposition 1.

We will assume number of input channels $C=1$. If the input is rotated and
translated by $(t,R)\in\mathrm{SO}(d)$ then input to the transformer encoder
transforms as
$f^{\textrm{in}}(\textbf{x},\rho)\mapsto\rho(R)f^{\textrm{in}}(R\textbf{x}+t,\rho)$
and positional encoding transforms as
$P(\textbf{x},\rho)\mapsto\rho(R)P(R\textbf{x},\rho)$. Let
$\mathbf{q}(\textbf{x}_{i},\rho)=\mathbf{q}_{i}^{\rho}$ and similarly,
$\mathbf{k}_{ij}$ and $\mathbf{v}_{i}$. Then,

$\displaystyle\mathbf{q}(\textbf{x}_{i},\rho)\mapsto\rho(R)f^{\textrm{in}}(R\mathbf{x}_{i}+t,\rho)\mathbf{W}_{Q}^{\rho}=\rho(R)\mathbf{q}(R\textbf{x}_{i}+t,\rho)$
$\displaystyle\mathbf{k}(\textbf{x}_{i},\textbf{x}_{j},\rho)\mapsto\rho(R)f^{\textrm{in}}(R\mathbf{x}_{i}+t,\rho)\mathbf{W}_{K}^{\rho}+\rho(R)P(R(\mathbf{x}_{i}-\mathbf{x}_{j}),\rho)=\rho(R)\mathbf{k}(R\textbf{x}_{i}+t,R\textbf{x}_{j}+t,\rho)$
$\displaystyle\mathbf{v}(\textbf{x}_{i},\rho)\mapsto\rho(R)f^{\textrm{in}}(R\mathbf{x}_{i}+t,\rho)\mathbf{W}_{V}^{\rho}=\rho(R)\mathbf{v}(R\textbf{x}_{i}+t,\rho)$

Let $s(\textbf{x}_{i},\textbf{x}_{j},\rho)=s_{ij}^{\rho}$ and
$\alpha(\textbf{x}_{i},\textbf{x}_{j},\rho)=\alpha_{ij}^{\rho}$. Then using
the fact $\rho(R)$ is unitary we have,

$\displaystyle s(\textbf{x}_{i},\textbf{x}_{j},\rho)$
$\displaystyle=\sum_{\rho,\rho^{\prime}}w^{(1)}_{\rho,\rho^{\prime}}\frac{\mathbf{q}(\textbf{x}_{i},\rho)^{\dagger}\mathbf{k}(\textbf{x}_{i},\textbf{x}_{j},\rho^{\prime})}{\sqrt{d_{k}}}$
$\displaystyle\mapsto\sum_{\rho,\rho^{\prime}}w^{(1)}_{\rho,\rho^{\prime}}\frac{\mathbf{q}(R\textbf{x}_{i}+t,\rho^{\prime}){\rho(R)^{\prime}}^{\dagger}\rho(R)^{\prime}\mathbf{k}(R\textbf{x}_{i}+t,\textbf{x}_{j},\rho^{\prime})}{\sqrt{d_{k}}}$
$\displaystyle=\sum_{\rho,\rho^{\prime}}w^{(1)}_{\rho,\rho^{\prime}}\frac{\mathbf{q}(R\textbf{x}_{i}+t,\rho)\mathbf{k}(R\textbf{x}_{i}+t,\textbf{x}_{j},\rho^{\prime})}{\sqrt{d_{k}}}$
$\displaystyle=\sum_{\rho,\rho^{\prime}}w^{(1)}_{\rho,\rho^{\prime}}s(R\textbf{x}_{i}+t,R\textbf{x}_{j}+t,\rho^{\prime})$
$\displaystyle=s(R\textbf{x}_{i}+t,R\textbf{x}_{j}+t,\rho)$

The first equality follows from the fact that since $C=1$ the _vec_ operation
is no longer needed. Since $\alpha$ is obtained by taking the softmax of $s$,
it follows that.

$\displaystyle\alpha(\textbf{x}_{i},\textbf{x}_{j},\rho)\mapsto\alpha(R\textbf{x}_{i}+t,R\textbf{x}_{j}+t,\rho).$

Finally,

$\displaystyle f^{\textrm{out}}(\textbf{x}_{i},\rho)$
$\displaystyle\mapsto\sum_{\rho\rho^{\prime}}\sum_{j}w^{(2)}_{\rho\rho^{\prime}}\alpha(R\textbf{x}_{i}+t,R\textbf{x}_{j}+t,\rho^{\prime})\rho(R)\mathbf{v}(R\textbf{x}_{j}+t,\rho^{\prime})$
$\displaystyle=\rho(R)\left(\sum_{\rho\rho^{\prime}}\sum_{j}w^{(2)}_{\rho\rho^{\prime}}\alpha(R\textbf{x}_{i}+t,R\textbf{x}_{j}+t,\rho^{\prime})\mathbf{v}(R\textbf{x}_{j}+t,\rho^{\prime})\right)$
$\displaystyle=\rho(R)f^{\textrm{out}}(R\textbf{x}_{i}+t,\rho)$

∎

##  Appendix C Experiment Details

Here we provide the details of all the experiments we have conducted. Table 4
shows detailed results on the two datasets.

Datasets | Model |  Radial Resolution |  Frequncy Cutoff | Accuracy |  Parameter ($\sim\times 10^{6}$)
---|---|---|---|---|---
Rotated MNIST |  Steerable Convolution | $r=2$ | $k=4$ | $98.69_{\pm 0.05}$ | $1.00$
$k=8$ | $98.93_{\pm 0.02}$ | $1.98$
$r=4$ | $k=4$ | $98.72_{\pm 0.02}$ | $1.98$
$k=8$ | $98.97_{\pm 0.01}$ | $3.94$
Steerable Convolution (Without Transformer) | $r=2$ | $k=4$ | $98.59_{\pm 0.08}$ | $0.24$
$k=8$ | $98.92_{\pm 0.04}$ | $0.47$
$r=4$ | $k=4$ | $98.70_{\pm 0.05}$ | $0.47$
$k=8$ | $98.97_{\pm 0.05}$ | $0.92$
Steerable Transformer | $r=2$ | $k=4$ | $98.78_{\pm 0.08}$ | $0.91$
$k=8$ | $\mathbf{99.03}_{\pm 0.04}$ | $1.78$
$r=4$ | $k=4$ | $98.82_{\pm 0.04}$ | $1.13$
$k=8$ | $\mathbf{99.03}_{\pm 0.06}$ | $2.24$
ModelNet10 (${z}$ Rotation) ($\mathrm{SO}(3)$ Rotation) |  Steerable Convolution | $r=1$ | $\ell=4$ | $89.40_{\pm 0.53}$ | $0.90$
$\mathbf{86.78}_{\pm 0.62}$
$r=2$ | $\ell=4$ | $90.13_{\pm 0.52}$ | $1.08$
$86.62_{\pm 0.25}$
Steerable Convolution (Without Transformer) | $r=1$ | $\ell=4$ | $89.08_{\pm 0.91}$ | $0.34$
$86.19_{\pm 0.55}$
$r=2$ | $\ell=4$ | $89.51_{\pm 0.55}$ | $0.48$
$86.58_{\pm 0.92}$
Steerable Transformer | $r=1$ | $\ell=4$ | $\mathbf{89.52}_{\pm 0.55}$ | $0.90$
$85.87_{\pm 0.74}$
$r=2$ | $\ell=4$ | $\mathbf{90.40}_{\pm 0.25}$ | $0.92$
$\mathbf{86.80}_{\pm 0.58}$
Table 4: Comparison of Steerable convolution and steerable transformers. The
mean and sd are provided for 5 separate runs

###  C.1 Rotated MNIST

SE2ConvType1(1, 24, 5, r, angle, maxk), CGNonLinearity2D(), SE2ConvType2(24,
48, 5, r, angle, maxk), SteerableBatchNorm2D(), AvgPool2D(2),
SE2TransformerEncoder(18, 5, maxk, layers = 1), SE2ConvType2(48, 48, 5, r,
angle, maxk), CGNonLinearity2D(), SE2ConvType2(48, 96, 5, r, angle, maxk),
SteerableBatchNorm2D(), AvgPool2D(2), SE2TransformerEncoder(96, 4, maxk,
layers = 1), SE2ConvType2(96, 64, 7, r, angle, maxk), NormFlatten(),
Linear(144, 128), BatchNorm1d(128), ReLU(), Dropout(0.7), Linear(128, 10) (a)
Steerable Transformer

SE2ConvType1(1, 24, 5, r, angle, maxk), CGNonLinearity2D(), SE2ConvType2(24,
48, 5, r, angle, maxk), SteerableBatchNorm2D(), AvgPool2D(2), SE2ConvType2(48,
96, 5, r, angle, maxk), CGNonLinearity2D(), SE2ConvType2(96, 192, 5, r, angle,
maxk), SteerableBatchNorm2D(), AvgPool2D(2), SE2ConvType2(192, 192, 7, r,
angle, maxk), NormFlatten(), Linear(192, 128), BatchNorm1d(128), ReLU(),
Dropout(0.7), Linear(128, 10) (b) Steerable Convolution

Figure 4: ModelNet10 Architecture

For this experiment we have used $w^{(1)}_{\rho\rho^{\prime}}$ as learnable
scalars which are same for all $\rho$ and
$w^{(2)}_{\rho\rho^{\prime}}=\mathbbm{1}_{\rho=\rho^{\prime}}$. Using these
weights the steerable attention mechanism boils down to:

$\displaystyle
s_{ij}=\sum_{\rho^{\prime}}w^{(1)}_{\rho^{\prime}}s(\mathbf{q}_{i}^{\rho^{\prime}},\mathbf{k}_{ij}^{\rho^{\prime}})\qquad\alpha_{ij}=\frac{\exp(|s_{ij}|)}{\sum_{j^{\prime}=1}^{N}|s_{ij^{\prime}}|}\qquad
f^{\textrm{out}}(\textbf{x}_{i},\rho)=\sum_{j=1}^{N}\alpha_{ij}\mathbf{v}_{j}^{\rho}$

The network architecture is shown in Figure 4. The networks are trained using
Adam Optimizer [15] with an initial learning rate of $5e$-$3$ which steps down
by a factor of $0.5$ every $20$ epochs and weight decay of $5e$-$4$, for $150$
epochs. The angular resolution for the steerable convolution layers (denoted
by the parameter angle in Figure 4) is fixed at $40$. More detailed
experiments where we have varied the radial resolution (r) and the frequency
cutoff (maxm) are presented in Table 4. We used a batch size of $100$ and it
took $2.5$ hrs to train our largest model on four $16$GB GPUs.

###  C.2 ModelNet10

The point cloud data with $2048$ point can be found here. For this experiment
with $z$ rotation we have used the most general model where both
$w^{(1)}_{\rho\rho^{\prime}}$ and $w^{(2)}_{\rho\rho^{\prime}}$ are learnable
scalars. For the $\mathrm{SO}(3)$ rotation we used the same scalar choice as
Rotated MNIST. The network architecture is shown in Figure 5. The networks are
trained using Adam Optimizer [15] with an initial learning rate of $1e$-$3$
which steps down by a factor of $0.5$ every $20$ epochs, for $50$ epochs. The
angular resolution for the steerable convolution layers (denoted by the
parameter angle in Figure 5) is fixed at $40$. More detailed experiments where
we have varied the radial resolution (r) and the frequency cutoff (maxl) are
presented in Table 4. We used a batch size of $5$ and it took $12$ hrs to
train our largest model on one $16$GB GPU.

SE3ConvType1(1, 9, 5, r, angle, maxl), CGNonLinearity3D(), SE3ConvType2(9, 18,
3, r, angle, maxl), SteerableBatchNorm3D(), AvgPool3D(4),
SE3TransformerEncoder(18, 3, maxl, n_layers = 1), SE3ConvType2(18, 36, 3, r,
angle, maxl), CGNonLinearity3D(), SE3ConvType2(36, 72, 3, r, angle, maxl),
SteerableBatchNorm3D(), AvgPool3D(4), SE3TransformerEncoder(72, 4, maxl,
n_layers = 1), SE3ConvType2(72, 144, 2, r, angle, maxl), NormFlatten(),
Linear(144, 128), BatchNorm1d(128), ReLU(), Dropout(0.7), Linear(128, 10)
(a) Steerable Transformer

SE3ConvType1(1, 9, 5, r, angle, maxl), CGNonLinearity3D(), SE3ConvType2(9, 36,
3, r, angle, maxl), SteerableBatchNorm3D(), AvgPool3D(4), SE3ConvType2(36, 72,
3, r, angle, maxl), CGNonLinearity3D(), SE3ConvType2(72, 72, 3, r, angle,
maxl), SteerableBatchNorm3D(), AvgPool3D(4), SE3ConvType2(72, 144, 2, r,
angle, maxl), NormFlatten(), Linear(144, 128), BatchNorm1d(128), ReLU(),
Dropout(0.7), Linear(128, 10) (b) Steerable Convolution

Figure 5: ModelNet10 Architecture