Steerable Transformers for Volumetric Data
Authors: Soumyabrata Kundu
Risi Kondor
Venue: ICML 2025 poster (2025)
Keywords: Equivariance
transformers
vision transformers
steerable
Primary area: applications->computer_vision
TL;DR: SE(d) equivariant transformfers for volumetric data

Abstract: 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 steerable transformer layers to steerable convolutional networks enhances performance.

Steerable Transformers for Volumetric Data
Soumyabrata Kundu Risi Kondor 1.2
Abstract
of the graph remains invariant under permutations of node
labels (Thiede al.) 2020 Wang et al.) 2020). Similarly, in vision tasks, it is desirable to have equivariance to rigid body transformations such as rotations and translations. Recent studies by Cohen & Welling 2016 2017); Worrall et al. (2017); Weiler et al. (2018a); Weiler & Cesa (2019); Thiede et al. (2020); Anderson et al. (201)), and many more, have demonstrated the remarkable efficacy of these equivariant architectures in enforcing symmetry, without relying brute force techniques like data augmentation.
We introduce Steerable Transformers, an exten- sion of the Vision Transformer that is equivari- ant to the action of the Special Fuclidean group SE(d). We propose an steerable self-attention mechanism that operates on features extracted by steerable convolutions. Our experiments in both two and three dimensions show augmenting steerable convolutional networks with steerable transformer leads improved performance.
In vision tasks, steerable convolutions have proven to be powerful tool for enforeing equivariance to rotations and translations (Cohen & Welling 2016 Worrall et al. 2017: Cohen & Welling 2017 Cohen et al. 2018 Weiler et al. 2018a Weiler & Cesa 2019). These networks primarily operate in Fourier space, leveraging representations of the underlying group extract rotation and translation equiv- ariant features from the data. These architectures are ef- fective at learning relevant features in local neighborhoods. Conversely, transformer architectures excel at learning re- lationships between distant regions of an image. Our con- tribution lies in combining these two concepts, yielding equivariant architectures that are capable of capturing both local and global patterns in images.
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 remarkable success, has led to their widespread adoption across various domains, including image classification (Dosovitskiy et al.) 2021], object detection (Carion et al.
2020), and graph
based problems (Dwivedi & Bresson 2020). The self-
attention mechanism (Bahdanau et al. transformer architectures has proven to be crucial for cap- turing relationships between different parts of the input se- quence. Dosovitskiy et al. (202T) introduced transformers as an alternative to traditional convolutional architectures for vision tasks. Unlike convolutional neural networks (CNNs), which focus on local neighborhoods, transform- ers excel at capturing relations across different parts of the input.
2014) employed in
1.1. Related Work
The concept of attention was first proposed by Bahdanau et al. (2014) in the context of sequence modeling. Building this, Vaswani et al. (2017) developed the transformer architecture for natural language translation tasks. Later, Dosovitskiy et al. (2021) extended the transformer frame- work to visual domains, proposing a hybrid model that in- tegrates features extracted from conventional CNNs. In contrast, Ramachandran et al. (2019) introduced fully at- tention based architecture that entirely replaces traditional convolutional networks.
Equivariant neural network architectures have gained sig- nificant popularity in recent years due their inherent abil- ity to comprehend the underlying symmetries of problems, making them highly effective tools for real world tasks. For instance, achieving equivariance to the permutation group S, is crucial in graph based problems, since the structure
Transformer architectures are inherently nonlinear in their inputs, and applying them directly into equivariant archi- tectures can break equivariance. This has motivated in- creasing interest in developing transformer models that are designed preserve equivariance. |Romero & Cordonnier (202 demonstrated that vision transformers using relative positional encoding (Shaw al. 2018) exhibit translation
'Dpartment of Statisties, University of Chicago, Chicago, USA DDepartment of Computer Science, University of Chicago, Chicago, USA. Correspondence to: Soumyabrata Kundu souumyabratakundue @ uchicago.edu>.
Proceedings of the 1nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Stecrable Transformers for Volumetric Data
equivariance. They extended the attention based frame- work of [Ramachandran et al. (2019) by formulating the attention mechanism over the group SE(2), thereby achiev- ing both rotation and translation equivariance in two di- mensions. Building on this approach. Xu et al. (2023) pro- posed a novel positional encoding operator to further im- prove performance. Additionally, Romero et al. (2020) in-
structured nature of the underlying grid. Furthermore, treating volumetric data as point clouds and applying these methods can lead to excessive memory usage, often result- ing in out-of-memory errors.
To address these challenges, we present a novel steerable transformer architecture specifically designed for volumet- ric data in d dimensions. Our proposed steerable trans- former integrates on top of a steerable convolutional en- coder and, as demonstrated through our experiments, en- hances the performance of steerable convolutions. Draw- ing inspiration primarily from|Dosovitskiy et al. (2021) and Fuchs et al. (2020), our method operates in the Fourier do- main, by using group representations to learn equivariant features. While Vaswani et al. (2017) employ fixed posi- tional functions and Dosovviskiy al. (202I) use learnable positional encodings, our method integrates both ideas by modulating a fixed positional function with learnable scal- ing parameters.
corporated attention directly into convolution by defining the attention scores over the equivariance group.
In three dimensions, attention based methods that enforce cquivariance to rotations and translations have been ex- plored. However, to the best of our knowledge, these ef- forts have been limited to point cloud data. |uucs et al.
(020)) introduced an -E3()--euivaaiaat architecture for point clouds by incorporating the self-attention mecha- nism into the Tensor Field Network (Thomas et al., 2018). Liao & Smidi (2023) proposed Equiformer, an equivari- ant extension of the Graph Attention Network (Velickovid 2018) for 3D point clouds, and subsequently, Liao
let al.
jet al. (2024) presented EquiforeerV2. a more computa- tionally efficient variant that reduces S3(3) convolutions to (2), leveraging ideas from Passaro & Zitnick (2023). Chen & Villar (0222 achieved equivariant attention by pro- and invariant features, based on equivariant feature maps developed by Villar et al. (202T) Hutchinson et al.] (202T) proposed a transformer framework with attention mechanism that is equivariant un- der the action of general Lie groups, using regular rep- resentations. For shape reconstruction tasks, Chatzipan tazis et al. (2023) introduced -Equi-equi atten- tion model.
2. Background
In this section we introduce the relevant background re- quired to design the steerable transformer architecture.
jecting data onto
2.1. The Multihead Self-Attention Mechanism
The attention mechanism, originally introduced by Bah- danau et al. (2014) for eqquence-to-sequenee learning, re- lies on three core components for each element in the input sequence: query vector qi, key vector ky. and value vector vy. The scores sij measure the compatibility between the query and key vectors using a scaled dot product. These scores are passed through a softmax function to produce attention weights Glay, which quantify the contribution of each value vector to the output. The output is then com- puted as weighted sum of the value vectors:
1.2. Our Contribution
Romero et al.] (2020) explored equivariant transformer architectures for two dimensional images by integrating the self-attention mechanism within convolutional layers. While effective, this comes at the cost of significantly higher memory usage. In contrast, the architectures intro- Xu et al. (2023), which build on the model of Ra- duced bu
ATTN(g;, (kyl, (u,))
asjyy-
machandran et al. 2019). entirely eliminate convolutional components in favor of self-attention based mechanisms. Nevertheless, as noted by Xiao al. (2021), incorporating convolutional encoder prior to the transformer leads to better performance, suggesting a more balanced and effec- tive design.
The attention weights aij are given by
of kg 8:3 Vdx
exp(8ty)
Ey= exp(s1t)
where dæ denotes the dimension of the query and key vec- tors, and N represents the number of elements in the input sequence. Vaswani et al. (017)) proposed a Multihead ver- sion of this attention mechanism which involves stacking multiple attention mechanisms, concatenating their outputs and applying a linear transformation to get the final result:
In three dimensions, Fuchs et al. 2020 Hutchinson et al. 2 Liao & Smidt (2023) proposed SE(3)- eequivariant transformer architectures tailored for point cloud data. Al- though volumetric data, in principle, can be conceptual- ized as point clouds arranged on regular grid, using these point cloud based methods on dense volumetric inputs can impose inappropriate inductive biases, as overlooks the
Concat(ArTN) ,.... "
Wo E ggdy Xdaaai
Here, dy denotes the dimension of the value vectors, and dmaatel denotes the dimension of the output. Typically, the
2

---

Stecrable Transformers for Volumetric Data
query, key, and value vectors are obtained by multiplying the input sequence with learnable weight matrices followed by adding positional encoding.
Inmuttively, this means that applying the group action be- fore after the map À leads the same outcome, guar- antecing that the resulting features respect the symmetry structure imposed by the group. Equivariance naturally arises as a desirable property in many real world applica- tions that involve underlying symmetries, such as permu- tations in graph structured data or rigid body transforma- tions in images. There extensive research on equivari- ant neural networks, particularly in the setting where the map À is linear (Worrall et al.] 20117] Cohen et al. 2018
Positional Encoding: Positional encoding in transform- ers is used to embed information about the position of words or tokens within a sequence into their input repre- sentations. In vision applications, these encodings help the model retain spatial awareness by indicating the original
locations of image patches. Vaswani et al. (0017) used sine and cosine functions at different frequencies to encode po- sitions. A refinement of this approach, known as relative positional encoding, was proposed by [Shaw et al. (2018). Unlike fixed encodings, this method learns representations that capture relative distances between elements in the in- put. As we will discuss later, incorporating relative posi- tional encoding is essential for preserving equivariance.
Weiler et al.
2018a Cohen & Welling 2017 2016 Weiler
& Cesa 2019 Kondor et al. 2018 Anderson et al. 2019
Thomas et al.
2018).
2.4. Steerable Convolutions
Steerable neural networks are designed to exhibit equivari- ance to the action of the Special Euclidean group SE(d) in d dimensions (see Appendix A.T for definition). Steer- able convolutions represent the most general class of lin- ear maps that are equivariant with respect the action of SE(d) (Kondor & Trivedi 2018 Although practical neural networks work with finite resolution rasterized im- ages, we adopt continuous formulation for the network's s inputs and filters to simplify the mathematical exposition and provide a clear overview of steerable convolutions. In this setting, d dimensional image with a single input channel is represented a compactly supported function f Rd + R, and analogously. filter is also repre- sented as a compactly supported function w IRd - R. The output of a steerable convolution layer is a function rour : :E(d) + R that can be expressed as
2.2. Vision Transformers
Dosovitskiy et al. (202T) adapted the transformer architec- ture for vision tasks. Their method involves partitioning an image into equal sized patches, which are then processed by transformer. The resulting feature representations are used for downstream tasks like image classification. As- suming a single input channel, the operations of a vision transformer layer can be summarized as follows:
20= = Ex,,, .... E] + Epos
2/ MULTT-ATIN(L((2MM-))) + ZM-1 m = I....M 2m (L)(L)((m-11) + 2m
m 1, ....... M.
you (x, R) = | gd
Here, X; E IR represents flattened image patch, E E Rexduute is the learnable linear weights matrix, Epos E Nxdeats is the positional encoding, and indexes the various blocks of the transformer. IN denotes a normal- ization layer, and MLP stands for a Multilayer Perceptron. dmodel is the dimension of the features fed into the trans- former, which remains fixed throughout the entire trans- former layer. Typically, the MLP layer comprises two hid- den layers separated by non-linearity (Vaswani al.) 2017 Dosovitskiy et al. 2021).
fi + Ry) w(y) dy,
(I)
1
| ( + Rty RR')
fou (x, Rt)
[so(d) JRt
w(y, R') dy dR)
Here, x E TR.d R (O(d) and jr denotes the Haar measure on SO(d) (see Appendix A.4 for definition). The presence of two distinct formulas is due to the fact that the input to the first layer is a vector field on Rd, whereas the input to subsequent layers is vector field on SE(d). Equation [I] represents the formula for the first layer, while (2) repre- sents subsequent layers. These layers are equivariant under the action of the group SE(d). Specifically, if the input function f is transformed by an element (t, R) E SE(d) according to
2.3. Equivariance
Equivariance describes property in which map between vector spaces commutes with the action of group G' (see Appendix A.2 for definition of group action). Specifically, for each element g E G, suppose we have linear operators Tg : V- + V and T : W- > W that represent the action of g on vector spaces V and W. respectively, A map À V + Wi is said to be equívariant if, for all v E V and g E G,
[tt, (2) (x = ( + t),
(3)
then the corresponding output fout transforms as
=Tx))))
[t., R1. [x)) R') = root j tt, RR').
(4)
3

---

Stecrable Transformers for Volumetric Data
For compact groups like SO(d), where discretization is not straightforward, it is advantageous to compute the rotation component of the output feature map of ( in Fourier space (Worrall et al., 2017). The Fourier transform of a functioon f on a compact group can be computed by inte- grating it against it's irreducible representations or irreps (see Appendix A.. |A.3 for definition):
tensor product of Fourier vectors followed by Clebsch- Gordan decomposition. Weiler al. (2018a) introduce other steerable convolution filter to serve as non-linearity, These methods enable effective non- non-linear transformations in Fourier space, while preserving equivariance.
3. Method
S(R) (() f()) f (so(a)
The methodology behind steerable transformers builds on the foundation of vision transformers (Dosovitskiy et al. 2021). which operate on "patchified" images, i.e., linear projections of fixed sized patches extracted from the input image. Conceptually, this patchification process is anal- ogous to applying strided convolutions, where the stride is equal to the kernel size. [Dosoviskkiy et al. (202T) also propose hybrid architecture that extends this idea by em- ploying deeper convolutional encoder extract more pressive features, which are then used as input to the trans- former layers. In steerable transformers, we build on this idea by using feature maps extracted from steerable convo- lution encoder as input to the transformer architecture.
(5)
Leveraging the Convolution Theorem, equations - in Fourier space becomes a matrix product of the Fourier transforms of the input and the filter. By the translation property of the Fourier transform, the vector fields satisfy the property that when the imput transforms under the action of (F.(d) as in 3 the corresponding output transforms as
[[t, R)
fie"l ((., p) = p(R) foul (Rx + p).
(6)
In the two dimensions, the irreps of S0( are just complex exponential indexed by integers, Pa: (f) ea for & E Z.
Therefore, the hthh frequency component of the output fea- ture map is a function
In the discussion of steerable convolutions in Section 2.4 the input, filter and consequently, the output were modeled as compactly supported functions defined over continuous domains. We adopt this continuous perspective in the dis- cussion that follows, while acknowledging that, in practical implementations, the output of a steerable convolutional encoder for a given irrep is represented as a collection of complex-valued vectors arranged on discrete grid. Ac- cordingly, any integrals in the theoretical formulation can be approximated by summations over this grid. While the number of channels associated with each irrep may vary in practice, for brevity, we consider the simplified case in which all irreps share the same number of channels. The generalization to a variable channel setting is straightfor- ward.
fouil (.A) C.
(7)
In three dimensions, the irreps of SO(3) are indexed by non-negative integers & E Zzo. and are given by (R) DR).. the so called Wigner D-matrices (Wigner, 1932). which are unitary matrices of size 1.+ Since the columns of the output can be interpreted as separate channels, the 4th Fourier component of the output feature map for single channel can be represented as function
poul( (. (): : R + C24+1
(8)
Non-linearities: In steerable neural networks, applying non-linearities like ReLU in Fourier space can disrupt equivariance. One approach to mitigate this issue is to tran- sition back time domain, apply the desired non-linearity there, and then return to Fourier domain (Cohen al.) 2018). However, this back and forth transformation can be computationally expensive and error prone, especially on groups like 53(3), where the uniform grid gives rise to singularities.
Assume that the input has dmodel channels for each irrep. Then, the input to the transformer corresponding to an irrep p can be represented as function
f (x. p) E Cd
which has compact support ((fr)) for all p. and satisfies the equivariance constraint (6). Here, Rd, and dp de- notes the dimension of p. In two dimensions, all irreps are one dimensional 77. whereas in three dimensions, the gh representation has dimension 2 +
As an alternative, Fourier space non-linearities are pre- ferred in stecrable neural networks. These non-linearities act directly Fourier space while preserving equivari- ance. Several such non-linearities can be found in the liter- ature. Worrall et al. (2017) apply non-linearity to the norm of the Fourier vector, as the norm remains invariant under rotations. Kondor et al. (2018); Anderson et al. (0119) use the Gordan non-linearity, which involves a
3.1. Steerable Self-Attention Mechanism
The query, key, and value vectors are computed by apply- ing learnable weight matrices to the input frn in ( fol- lowed by the addition of positional encodings. Romero & Cordonnier (202T) showed that incorporating relative po-

---

Steerable Transformers for Volumetric Data
Scaled Product
Query
Equirariant Scores
Noumalizon
Soweratle
Ataptionn
Input
Sizerable Pusilionl Erizding
Softmax the Norm
Key
Faslure Maps for Components
Linear Combiration
Weight Musspiictioon
Value
Equaanant Attention Scores
Wiigght
Murapiccatoo
Outpul
Culpul
Vision
Steerable Altention Mechanism
Transformer Layer
Figure The schematic illustrates the steerable self-attention mechanism for a single head (h = 1) and one query dimension (dx- = 1) (left) and steerable transformer encoder layer (right); c.f. Figure by Dosovitskiy al. (2021).
sitional encodings (Shaw al., 2018) into self-attention enforces translation equivariance. Our design of steerable positional encodings is inspired by this relative positional encoding. We employ a hybrid strategy that blends learn- able weights and fixed functions to encode positional infor- mation.
this be expressed
x(x,y) = exp([s(x,)) y)l
(14)
Jai exp()s(x, 2)))
Since we are dealing with complex numbers, we use the absolute value of the scores, and the dot product involves the conjugate transpose, denoted by -f, instead of reg- ular transpose. Finally, the output for each position and Fourier component is computed as a linear combination of the value vectors and attention scores:
In the context of steerable neural networks, it is essential that positional encodings also encode rotational structure, which is naturally expressed in Fourier space. Accordingly, for each irrep p. we define the steerable positional encoding as a continuous function
fTm (x, = - a(x, (( (x,
(())
Rd
x + C4.
R
(15)
We elaborate on the exact forms of these encodings in the next section. For now, we focus on the design of the steer- able self-attention mechanism. Within this framework, the query, key, and value vectors corresponding to an input map fin are given by
It is easy to verify that if the steerable positional encodings are set to zero, then by the unitary nature of the irreps, the output of the self-attention mechanism satisfies the equiv- ariance constraint (6). However, to encode the spatial in- formation of the input patches, we want use non-trivial functions as steerable positional encodings. The following theorem establishes the necessary and sufficient condition that p(p) must satisfy for the self-attention mechanism to preserve equivariance.
(x) (x
(10)
k(o) (x,)) (x. () p(e)((x.))))
(11)
(() (x, () W + (xl))(x.))
(12)
Theorem 1. For any f that transforms under the action of SE(d) according to the equivariance constraint (6), the output of the self-attention mechanism also satisfies (6) for any collection of weight matrices if and only if. for every (t, R) E SE(d), the positional encodings p(e) satisfy
(p) E Chuss xdx and w!()) V E Calasa xd Next, the score function is calculated by taking the scaled dot product between the query and key vectors:
where (())
1:: w
s(x,y) = Σ vec(q( (y))) veo ())
(13)
+ tt, RR +t)
t)= Ry
p( p(e) (x.y), (16)
Vdk
for all x,y E Idd and for all irreps
Here, the vec operation flattens the matrix into a vector. The attention scores are then obtained by applying the soft- max function to the raw scores. In continous formulation
In the following section, we elaborate on the construction of (P)) that satisfies the constraint (6).
5

---

Steerable Transformers for Volumetric Data
3
6
9
2
Rotation by 90 degrees
5
8
5
7
Figure 2: Visual representation of steerable positional encoding. Arrows denote directional components, while the color gradient indicates magnitude, decaying proportionally to r22
In practice, the integrals appearing in equations (14 and (15) are approximated by discrete summations over a finite of grid points. Specifically, the attention weights and output computed as
for E Z. Here, o(r, k) modulates the encoding strength for nearby and distant points. Similarly, in three dimen- sions, we have
r sin cos ø r rsi 0 sin Φ cos
xx)()s(x;,x)))) x)) 1 exp((s(xx.xxx))]
plo (r, x
(19)
for Zzu- Here y denotes the spherical harmonics of the unit sphere S² in three dimensions (Byerly 1893).
(u (xi, p) Σ a (x; x, ) (x, x,);
The stecrability property of positional encoding (17 im- plies p()) (0) = for any p except for the constant rep- resentation (k: 0 and sions, respectively). since (()) p R) (() for all R SO(d). This condition can be enforced by setting ó to zero when = 0. For our experiments, we used
respectively, where {x: denotes set of discrete spa- tial locations grid. For multihead self self-attention, this mechanism is repeated independently for h heads. Finally, the outputs are concatenated, and each Fourier component scaled by another matrix W(() Chdy xdmat Consis- tent with Vaswani al. (2017) choose dy dk
0 in two and three dimen-
Φ(T, p) = wpr - 21,50
dm
where wp is a learnable scalar. This kind of inverse square modulation of the radial component results in positional en- coding assigning higher weight to neighboring points, and the weights decrease as points move further apart (see Fig- ure 2}. In our implementation, we have different learnable sealars, not only for different Fourier components, but also for different heads and query dimensions.
3.2. Steerable Positional Encoding
Theorem I establishes necessary and sufficient conditions that ( must satisfy to ensure that the self-attention mechanism preserves equivariance. In particular, for any translation vector t E Rdd the positional encoding must satisfy the condition
p(e) +x = p(P) -x,))x..).
3.3. Multilayer Perceptron
This implies that ((e) depends only on the relative posi- tion, and can thus be expressed as a function of single argument p(e) (x y). As result, the encoding is in- variant under translations. Furthermore, for any rotation R E SO(d), it has the property that
In a transformer architecture, the self-attention layer is fol- lowed by an MLP, which consists of two linear layers sep- arated by non-linearity (Vaswani et al. (017).. The linear layers themselves do not affect equivariance they just mix together the channels for each irrep:
p(o] (Rx) p(R) p(P)
(17)
pour (x, p) = ( w/()
natural and expressive choice for such function is the Spherical Harmonic basis functions (Frye Efthimiou 2012). In two dimensions, this is takes the form
T E Chunsa Xehaats and dhideen represents the hidden dimension between the two linear layers. When selecting the non- con-lincarity a, it is important to proceed with caution, as commonly used functions like ReL.U or sigmoid
Here, W(() (p)
cos rsin 0
p(k) (x) = o(r, kk =
(18)
6

---

Stecrable Transformers for Volumetric Data
5
15
5
Steerable Transforee
Rotalion
Rolation
Steerable Transtormer
Figure 3:
The figure demonstrates the equivariance of attention scores in a trained model. For a fixed pixel x,, we have plotted the maximum attention score for that pixel (max, a(x;, x,)). )) 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.
may break equivariance, as discussed in Section 2.4 For experiments, we use Fourier space non-linearity pro- posed by Worrall et al. (2017), where they apply ReLU to the magnitude of the equivariant vectors:
d,C' channels, leading a complexity of O(N Analogously, vanilla convolutions in d dimensions, pro- cessing N pixels with C' input and output channels with kernel size kd has a complexity of O( In steerable convolutions, replacing C' with d,C'' yields O(NeCC²k))
o(f(x,)) [eL.U(() bb) p(x,p).
[l,(x, p)ll2
In summary, with a fixed kernel size, the computational complexity of steerable self-attention matches that of steer- able convolution. Additionally, the increase in complexity from standard to steerable methods is similar for both self- attention mechanisms and convolutions when using many channels and small kernels, as is common in practice. This parallel implies that the additional complexity introduced by stcerable methods scales proportionally across both self- attention mechanisms and convolution operations.
Here, à is a learnable bias. Since the norm of the fea- tures in Fourier space is invariant to rotations, this non- linearity preserves equivariance. Following the convention in Vaswani et al. (2017). we set dhidden = 2dmodel
3.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 L2 norm. Leveraging this property, we can derive a steerable normalization as
4. Experiments
We evaluate the performance of our steerable transformer architecture in both two and three dimensions. Specifically, we focus on two vision tasks: image classification and se- mantic segmentation. In all experiments, we employed a hybrid architecture that integrates steerable convolutions with steerable transformer, and the trained using the Adam Optimizer (Kingma & Ba 2014).
p(x
(x())xx
|I(x,
Indeed, the sum over p theoretically extends over all irreps trun-
of the group, but practical implementations, cated at particular values, which serves as hyperparameter that can be tuned. This type of normalization was employed in the two dimensional case by Worrall et al. (2017).
The steerable convolution encoder is composed of multi- ple steerable convolution blocks. A steerable convolution block consists of two steerable convolution layers, sepa- rated by non-linearity, and is followed by layer normaliza- tion. Each steerable convolution block followed by an average pooling layer. In the baseline models for the clas- sification tasks, this encoder is followed by final convolu- tion layer with kernel size matching the input size, effec- tively serving as a steerable analogue to flattening layer. In the steerable transformer model, a steerable transformer with two blocks is inserted before this flattening layer. A norm operation is then applied enforce invariance to both
3.5. Complexity
The computational complexity of transformer block mainly stems from its self-attention mechanism and MLP. For sequence of length N and model dimension C, the self-attention mechanism has complexity of O(NC2 + N²C), and the MLP adds O(NC2), totaling O + C²C). When N x C, this becomes ONN²). In steerable self-attentio mechanisms, each component effectively has

---

Steerable Transformers for Volumetric Data
Datasets
Model
Frequency Cutoff
Accuracy Dice Score
Parameter (~ x x10")
Avg Runtime Per Batch
Train (s)
Inference (s)
Steerable
k = 4
98.72+0.02
1.18
0.14
0.08
Rotated MNIST
Convolution
k = 8
98.97 :0..07
2.54
0.31
0.15
Steerable
A: 4
98. 98.8211.04
1.13
0.13
0.08
Transformer
f 8
99.03+0.04
2.24
0.30
0.14
Steerable
90.1340.52
ModelNet10 (z Rotation)
Convolution
i = 3
86.6240.25
1.08
1.22
0.46
(SO(3) Rotation)
Steerable
l = 3
90.40++025
0.92
1.20
0.45
Transformer
10.58
Steerable
h = 4
89.31
0.35
0.27
0.09
Convolution
h 8
89. 89.6340030
0.96
0.67
0.19
PH2
Steerable
E = 4
90.61 16.10.90
0.44
0.28
0.10
Transformer
A: 8
90.72
1.22
0.72
0.20
73.29:40.76
Steerable Convolution
BraTS (Enhancing Tumor)
(=
51.07
0.12
5.04
1.00
74, 44.47
(Tumor Core)
75. 00.332
(Whole Tumor)
Steerable Transformer
t= 2
54.89 10.7
0.16
5.37
1.01
76.37 10.46
Table 1: Comparison of steerable transformers and steerable convolutions. The mean and sd are reported for For ModelNet10 have reported both the rotation and SO(3) rotation variations. For PH2 dataset we have reported the dice score for segmentation of the binary mask. For the BraTS dataset have reported the dice score individually for each tumor category.
rotation and translation. The output is passed through two fully connected layers, with a ReLU activation in between and dropout layer with a probability of 0.7 (Marcos al.) 2017]. We evaluate performance on datasets: Rotated MNIST (2D) and ModelNet10 (3D).
other methods, refer to Tables 2 and 3 Appendix C Be- low, we briefly describe the datasets used in our experi- ments. Further details regarding the architecture and hy- perparameters are provided in Appendix C The code for all experiments is available at futttps::// ttwpb/tttttut..oeeg com/ Sounyabrat aKundu SeeeStteeebbler Tetterrabe-Trrafffrrroe Transformer
For the segmentation tasks, we used U-net architec- ture (Ronneberger et al. 2015), with steerable convolution blocks and average pooling to downsample the input im- age, and combination of steerable convolution blocks and interpolation to upsample the features back to the input res- olution. The model includes skip connections, where fea- tures from the downsampling path added to the corre- sponding resolution features in the upsampling path. In the baseline model, an MLP is placed between the downsam- pling and upsampling paths, while in the steerable trans- former model, this is replaced by a steerable transformer with two blocks. We report results on two datasets: PH2 (2D) and BraTS (3D).
Rotated MNIST: Rotated MNIST, variant of the orig-
inal MNIST dataset (LeCun et al. 2010). includes gray- scale images handwritten digits with resolution 28 x 28 that have been randomly rotated. The dataset contains 12,000 training images and 50,000 testing images.
The ModelNet10 dataset (Wu et al. 2015) consists of 3D CAD models from 10 common object categories, with train/test split of 3991:908. Each category includes triangular meshes represent- ing objects at various orientations and scales. For our experiments, we used the point cloud version of the dataset available at [attp:://ittuu. com/ where the authors employed the farthest point sampling algorithm sam- ple 2048 points from the surface of each object. These point clouds were then embedded into grids to create vox- elized representations. Additionally, rotated versions the dataset was generated by applying random rotations using nearest-neighhor interpolation (see Figure
ModelNet10:
Table D compares the baseline steerable convolutional architecture with the proposed steerable transformer ar- chitecture. For classification tasks, we report accuracy, while for segmentation tasks, we present the dice score for each class (excluding the background). The results in Table demonstrate consistent performance improve- ments with the addition of the steerable transformer en- coder layer across various datasets. For comparisons with
8

---

Steerable Transformers for Volumetric Data
Input
Ground Truth
Probability
Predicted
PH2
BraTs
Figure 4: Illustration ground truth and predicted segmentation, along with the predicted probability for the true class for two examples from each dataset. The red the PH2 example represent the binary mask. The red, blue and green the BraTS dataset represent enhancing tumor, tumor core and whole tumor respectively.
PH2: The PH2 dataset (Mendonça et al. 2013) is col- lection of 200 dermoscopic images used for research in classification and segmentation of skin lesions. We ran- domly split the dataset into 100:50:50 for training, testing, and validation, respectively. Each image has a resolution of 578x xBxx770 770 with RGB channels, and accompanied by binary masks of the same resolution that delineate lesion regions.
mance. In current experiments, employed rela- tively modest Fourier cutoffs, which served computa- tional compromise. However, we hypothesize that increas- ing these cutoff values could lead to better results. Realiz- ing this improvement within steerable transformer frame- work would, however, require the model to be scaled up significantly handle the increased complexity and dimen- sionality of the representation.
BraTS: The Brain Tumor Segmentation (BraTSS) dataset (Menze et al.) 2015) is used for developing and evaluat- ing algorithms for brain tumor segmentation in MRI scans. includes pre-operative MRI scans across four modali- ties, with each image having resolution of 240x240x 155 and annotated delineate three tumor regions. We used train/validation/test split of 243:96:145
6. Conclusion
The transformer architecture has rapidly gained traction across a wide range machine learning domains, owing to its remarkable capacity to model long range dependen- cies and contextual relationships within the data. In this work, present steerable transformers, novel architec- tural component designed work in tandem with steer- able convolutions. While convolutions are highly effective at encoding local geometric structures, they are inherently limited their ability to capture interactions across dis- tant spatial regions. By integrating the transformer's self- attention mechanism, we create hybrid architecture that combines the strengths of both local and global represen- tations. Our experimental results demonstrate that this in- tegration leads to measurable performance improvements. Furthermore, the steerable transformer framework holds significant promise for high stakes applications like med- ical imaging, where equivariance rotations and transla- tions is essential for robust and interpretable analysis.
5. Limitations
While transformer architectures offer powerful modeling capabilities, they also present significant computational challenges. The dot product structure the attention mech- anism can lead substantial memory consumption, espe- cially when applied to high-resolution images. This lim- its both the size of the network and the batch size dur- ing training, thereby constraining the ability to fully lever- age the potential attention based models. To address this, scaling up the architecture crucial for unlocking the full benefits of attention in improving model perfor-

---

Stecrable Transformers for Volumetric Data
Impact Statement
Cohen, T. S. and Welling, M. Steerable CNNs. In Interna- tional Conference Learning Representations, 2017.
This paper contributes to the ongoing advancement of ma- chine learning, with a particular focus on computer vision and the development of equivariant neural network archi- tectures. By introducing novel techniques that respect the symmetries inherent in visual data, our work pushes the boundaries of how neural networks can effectively and ef- ficiently process complex spatial information. Beyond the technical contributions, this research carries the potential for meaningful societal impact. particular, the proposed architecture can be applied to biomedical imaging, field where accurate interpretation of spatial patterns is critical. Enhanced models for tasks such as tumor detection, dis- ease classification, and anatomical structure segmentation could significantly improve the capabilities of Al-assisted diagnosis and precision medicine. By enabling more reli- able and interpretable machine learning systems in health- care, our work supports the broader goal of making medical technologies more effective, accessible, and equitable.
Cohen, T. S., Geiger, M., Koehler, J., and Welling, M. Spherical CNNs. In International Conference on Learn- ing Representations, 2018.
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D.,, Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. An image is worth 16x16 words: Transformers for image recognition scale. In International Conference Learning Representations, 2021.
Dwivedi, V. and Bresson, X. A generalization of trans- former networks to graphs. In AAAI Workshop on Deep Learning on Graphs: Methods and Applications, 2020.
Frye, C. and Efthimiou, C. J. Spherical harmonics in p dimensions. arXiv preprint
2012.
Fuchs, F., Worrall, D., Fischer, V., and Welling, M. Se trantfassfomm 3d ttto-translation equivariant atten- tion networks. Advances in neural information process- ing systems, 33:1970-1981, 2020.
References
Anderson, B., Hy, T. S., and Kondor, R. Cormorant: Co- variant molecular neural networks. Advances in neural information processing systems, 32, 2019.
Hegde, V. and Zadeh, R. Fusionnet: 3d object classifica- tion using multiple data representations. arXiv preprint 2016.
Bahdanau, D., Cho, K., and Bengio, Y. Neural machine translation by jointly learning to align and translate. CoRR, abs/ 0bss/1499.043,, 2014.
Hutchinson, M. J., Le Lan, C., Zaidi, S., Dupont, E., Teh, Y. W., and Kim, H. Lietransformer: Equivariant self- attention for lie groups In International Conference on Machine Learning, pp. 1133-4543. PMLR, 2021.
Brock, A., Lim, T., Ritchie, J. M., and Weston, N. Gen- erative and discriminative voxel modeling with convolu- tional neural networks. ArXiv, hass/660.042284 2016.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. CoRR, abs/1412 111.1111266.99. 2014.
Byerly, W. An elemenatary treatise on Fourier's 's series, and spherical, cylindrical, and ellipsoidal harmonics, with applications problems in mathematical physics. Dover Publicatiions, 1893.
Kondor, R. and Trivedi, S. On the generalization of equiv- ariance and convolution in neural networks to the action of compact groups. In Dy, J. and Krause, A. (eds.), Pro- ceedings of the 35th International Conference chine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 2747-2755 PMLR, 10-15 Jul 2018.
Carion, N., Massa, F., Synnaeve, G., Usunier, N., Kirillov, A., and Zagoruyko, S. End-to-end object detection with transformers. In European conference on computer vi- sion, pp. 213-229. Springer, 2020.
Ma-
Chatzipantazis, E.,, Pertigkiozoglou. S., Dobriban, E., and Daniliidis, K. S networks for shape reconstruction in function space. In The Eleventh International Conference Learning Representations, 2023.
Kondor, R., Lin, Z., and Trivedi, S. Clebsch-gorde nets: fully fourier space spherical convolutional neural net- work. In Bengio, S., Wallach, H., Larochelle, H., Grau- man, K., Cesa-Bianchi, N., and Garnett, R. (eds.), Ad- vances in Neural Information Processing Systems 31, pp. 10117-10..6. Curran Associates, Inc., 2018
( attention
Chen, N. and Villar, S. Se (3)-equivariant self-attention via invariant features. In Machine Learning for Physics NewrlPS Workshop, 2022.
Laptev, D., Savinov, N., Buhmann, J. M., and Pollefeys, M. Ti-pooling: Transfornation-inaarinn pooling for feature learning in convolutional neural networks. In Proceed- ings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2016.
Cohen, T. and Welling, M. Group equivariant convolutional networks. In International conference on machine learn- ing, pp. 2990-2999, PMLR, 2016
10

---

Steerable Transformers for Volumetric Data
Larochelle, H., Erhan, D., Courville, A. C., Bergstra, and Bengio, Y. An empirical evaluation of deep archi- tectures on problems with many factors of variation. In International Conference on Machine Learning, 2007.
Passaro, S. and Zitnick, C. L. Reducing (3) convolu- tions (2) for efficient equivariant gnns. In Inter- national Conference on Machine Learning, pp. 27420-
27438. PMLR, 2023.
LeCun, Y., Cortes, C., and Burges, C. Mnist hand- written digit database. ATT Labs [Online]. Available: http://yan.lecum www....coo/ed////mnii.... 2, 2010.
Peter, F. and Weyl, H. Die vollständigkeit der primi- tiven darstellungen einer geschlossenen kontinuierlichen gruppe. Mathematische Annalen, 97(1):737-755, 1927.
Li, Chen, B. M., and Lee, G. H. So-net: Self-organizing network for point cloud analysis. In 2018 IEEE/CVF Conference on Computer Vision and Pattern Recogni- tion, pp. 9397-9406, 2018. doi: 10.1109/CVPR.2115 00979.
Ramachandran, P., Parmar, N., Vaswani, A., Bello, I., Lev-
skaya, A., and Shlens, J. Stand-alone self-attention in vision models. In Advances in neural information pro- cessing systems, volume 32, 2019.
Robert, Introduction to the representation theory of compact and locally compact groups, volume 80. Cam- bridge University Press, 1983.
Li, J., Bi, Y., and Lee, G. H. Discrete rotation equivariance for point cloud recognition. In 2019 International con- ference on robotics and automation (ICRA), pp. 7269- 7275. IEEE, 2019.
Romero, D., Bekkers, E., Tomczak, J., and Hoogendoorn, M. Attentive group equivariant convolutional networks. International Conference on Machine Learning, pp. 8188-8199, PMLR, 2020.
Liao, Y.-L. and Smidt, Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. In The Eleventh International Conference Learning Repre- sentations, 2023.
Romero, D. W. and Cordonnier, J.-B. Group equivariant stand-alone self-attention for vision. In International Conference on Learning Reprexentations, 2021.
Liao, Y-L., Wood, B. M., Das, A., and Smidt, T. Equiformerv2: Improved equivariant transformer for scaling to higher-degree representations. In The Twelfth International Conference on Learning Representations, 2024.
Ronneberger, O., Fischer, P., and Brox, T. U-net: Con- volutional networks for biomedical image segmentation. Medical image computing and computer-assisted intervention-MICCAl 2015: 18th international confer- ence. Munich. Germany, October 5-9, 2015, proceed- ings, part III 18, pp. 144-241.2 Springer, 2015.
Marcos, D., Volpi, M., Komodakis, N., and Tuia, D. Rota- tion equivariant vector field networks. In Proceedings of the IEEE International Conference on Computer Vision, pp. 5048-5051 2017.
Sedaghat, N., Zolfaghari, M., Amiri, E., and Brox, T. Orientation- voxel nets for 3d object recognition. In British Machine Vision Conference (BMVC) 2017, 2016.
Maturana, D. and Scherer, S. Voxnet: 3d convolu- tional neural network for real-time object recognition. In 2015 IEEE/RSJ International Conference on Intelligntt Robots and Systems (IROS), pp. 922-928, 2015. doi: 10.1 109/0000 201.RO5,2015.33344
Shaw, P., Uszkoreit, J., and Vaswani, A. Self-attention with relative position representations. North Ameri- can Chapter of the Association for Computational Lin- guisties, 2018.
Mendonça, T.. Ferreira, P. M., Marques, J., Marcal, A. R. S., and Rozeira, J. Ph2 a dermoscopic image database for research and benchmarking. In 35th An- nual International Conference of the IEEE Engineering in Medicine and Biology Society (EMBC), 5437- 5440. IEEE, 2013.
Simonovsky, M. and Komodakis, N. Dynamic edge- conditioned filters in convolutional neural networks on graphs. In 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 29 29-38, 2017. doi: 10. N0.100/CVPR.201..11
Menze, H., Jakab, A., Bauer, S., Kalpathy- Salpathy-Cramer, J., Farahani, K., and et al., K. The multimodal brain tumor image segmentation benchmark (brats). IEEE Transactions Medical Imaging, 34910):1993-2024, 2015. doi: N0.1109/TM..2114237776
Sohn, K. and Lee, H. Learning invariant representations with local transformations. International Conference on Machine Learning. 2012.
Thiede, E. H., Hy, T. S., and Kondor, R. The general theory of permutation equivarant neural networks and higher order graph variational encoders. arXiv preprint urXiv:2004.03990 2020.
Munkres, J. Topology: a First Course. Prentice-Hall 1974. ISBN 97801 39254956.

---

Steerable Transformers for Volumetric Data
Thomas, N., Smidt, T. E., Kearnes, S. M., Yang, L., Li, L.,, Kohlhoff, K., and Riley, P. F. Tensor field net- works: Rotation- and translation-equivarian neural net- works for 3d point clouds. ArXiv, abs/1802.08219, 2018.
Wu, Z.. Song. S., Khosla, A., Yu, F., Zhang, L., Tang, X., and Xiao, J. 3d shapenets: A deep representation for vol- umetric shapes. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1912- 1920, 2015.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, At- tention is all you need. Advances in neural information processing systems, 30, 2017.
Xiao, T., Singh, M., Mintun, E., Darrell, T., Dollár, P., and Girshick, R. Early convolutions help transformers better. Advances in neural information processing sys- lems, 14:30392-30400, 2021.
Velicković, P., Cucurull, G., Casanova, A., Romero, A., Liò, P.. and Bengio, Y. Graph attention networks. In International Conference Learning Representations, 2018.
Xu, R., Yang, K., Liu, K., and He, e(2)-equivariant vision transformer. In Uncertainty in Artificial Intelli- gence, pp. 2356-2366. PMLR, 2023.
Villar, S., Hogg, D. W., Storey-Fisher, K., Yao, W., and Blum-Smith, B. Scalars are universal: Equivariant ma- chine learning, structured like classical physies. Ad- vances in Neural Information Processing Systems, 34: 2021.
Wang, R., Albooyeh, M., and Ravanbakhsh, S. Equivari- ant networks for hierarchical structures. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Sys- tems, volume 33, pp. 13806-1317.. Curran Associates, Inc., 2020.
Weiler, M. and Cesa, G. General e (2)-equivariant steer- able cmns. Advances in neural information processing systems, 32, 2019.
Weiler, M., Geiger, M., Welling, M., Boomsma, W., and Cohen, T. S. 3d steerable cnns: Learning rotationally equivariant features in volumetric data. Advances in Neurul Information Processing Systems, 31, 2018a.
Weiler, M., Hamprecht, F, A., and Storath, M. Learning steerable filters for rotation equivariant enns. In Pro- ceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2018b.
Wigner, Gruppentheorie und ihre Anwendung auf die Quantenmechanik der Atomspektren, volume 39. Monatsh, Mathematik und Physik, 1932.
Worrall, D. E. and Brostow, G. J. Cubenet: Equivariance to 3d rotation and translation. European Conference on Computer Vision, 2018.
Worrall, D. E., Garbin, S. J., Turmukhambetov, D., and Brostow, G. J. Harmonic networks: Deep translation and rotation equivariance. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 5028 5028-5037, 2017.
2017.
12

---

Steerable Transformers for Volumetric Data
A. Prerequisites
A.1. Groups
Definition: A group G is non-empty set together with a binary operation (commonly denoted by ".", that combines any two elements and G form an element of G, denoted b, such that the following three requirements, known as group axioms, satisfied:
Associativity: For all c G, one has (a b)
(b
Identity Element There exists as element e G such that for every a E G, one has e a a. Such an element is unique. It is called the identity element of the group.
Inverse: For each a E G, there exists an element G such that a b b. e, where is the identity element. For each a, the element b is unique and it called the inverse of a and is commonly denoted by
Examples include set of integers Z with the addition operation and the set of non-zero reals I {0) with the multiplication operation. From here on we will drop ".", for simplicity. The group operation will be clear from the elements of the group concerned.
Group Homomorphism: Given two groups G and H, function G + H is called group homomorphism if é(ab) = (()))(b) for any a, E G. If the map & is bijection, it is called an isomorphism. 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 Aut(G).
Compact Groups: A topological group is topological space that is also group such that the group operation and the inverse map are continuous. A compact group is a topological group whose topology realizes it compact topological space (see (Munkres) 1974) for definition of compact topological spaces). Some classic examples of compact groups are the groups (O)d) (the group of all real orthogonal matrices in d dimensions with determinant 1), U(d) (the group of all complex unitary matrices) and (U) (the group of all complex unitary matrices with determinant 1).
Special Orthonormal Group (O))): This group comprises all real orthogonal matrices dimensions with determinant of These groups are associated with rotation matrices in d dimensions. in two dimensions, the group SO(2) can be parametrized by single angle corresponding the rotation matrix
cos (0) R(0)
sin (0) sin (0) cos (0).
Here, R(0) signifies rotation in the x-y plane by an angle 0 E 2T). In three dimensions, the group SO(3) can be parameterized using the so called Euler angles. Consider the rotation matrices (o) and R. (3), defined by
(oo((a) sin(a)
sin(a) cas
0
(B)
s3n
(2(a))
R3(
1
sin (3)
cos
which represent rotations by angles a and B about the z-axis and y-axis, respectively. Using the 2-y-2 convention, any rotation matrix R SO(3) can be expressed R (3) where a,7 E [0, 2Tr) and 8 E [0, T) the Euler angles.
Semi-direct Product Groups: Given two groups N and H and a group homomorphism o И - Aut(N), we can construct a new group N As H defined follows:
The underlying set is the Cartesian product N x H.
The group operation is given by (n1, (1) h2) = (nio,, (nz), h2h2).
13

---

Steerable Transformers for Volumetric Data
Special Euclidean Group SE(d): The Special Euclidean group SB(d) consists of all combinations of rotations and transla- tions in dimensions. Translations in IR form group isomorphic to Rd itself, while rotations represented by the group (O(d). We define a group homomorphism Φ SO(d) - Aut(Rd) by setting @R(() Rt, which describes how a rotation acts on a translation vector. Using this, the group SE(d) can be constructed as the semidirect product R X SO(d). The group operation for two elements (t,, R1) and (t/, R2) in SE(d) is given by
(t),2,) R2) = (t1 + R,t2, R2)2).
The identity element of SE(d) is (0,4) I), and the inverse of an element (t, R) is given by (t, R) (-R R.
A.2. Group Actions
Definition: is a group with identity element e, and is a set, then a (left) group action of G on X is a function G + X, that satisfies the following two axioms for all g, h E G and X:
Identity: afe,x) = I,
Compatibility; alg, ah,, = a(gh, x).
Often alg, x)) is shortened to g. Any group G acts itself by the group operation. If G acts X, then also naturally acts any function f defined on X, as (g (z) f(g zz)
Action of SE(d) on Rd: The special Euclidean group acts on vector in Rd by first applying the rotation component followed by translation. For x E Id and (t, Rt) E SE(d).
(t. R) x= R== +t
gives us the action SE(d) on Rd.
A.3. Group Representations
Definition: A representation of group a group homomorphism from G to GL(C") (group of invertible linear maps on C"). Here called the dimension of the representation, which can possibly be infinite. representation unitary P maps to unitary linear transformation of C
we have two representations, P1 and P2 of dimensions n1 and my respectively, then the two can be combined by direct sum give another representation of dimension n1 n22
Irreducible Representations:
(1 @ p2(g) = (p(()
0
p2(o)
representation is said to be completely reducible if it can be expressed direct sum of other representations after maybe change basis, i.e,
Up(o)U- θ pc(9)
where U is unitary change of basis matrix and the direct sum extends for every group there are some representations which cannot be broken further into 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:
number of representations. However,
p(g) U
p
(g)
U-!
where again U a change of basis matrix and (0) are the irreps. The Peter-Weyy Theorem by (Peter & Weyl 1927 tells us that for a compact group G, any unitary representation p is completely reducible and splits into direct sum irreducible finite dimensional unitary representations of G.
14

---

Steerable Transformers for Volumetric Data
Irreducible Representations of (0( and SO(3): SO(d) being a compact group, all its irreps finite dimensional unitary representations. In the case of SO(2), every irrep is one-dimensional and indexed by integer. The group (O can be parameterized by angle 0 E [0, 2T), corresponding to rotation in the z-y plane by 0 radians. Under this parameterization, the irreps of (2(2) take the form
(ah) (0) e
k E Z.
The irreps of S3(3) are indexed by positive integers E Ezo, where the l'th representation is of dimension 2 + and are given by the called Wigner D-matrices [Wigner 1932):
(R)=
E Zzo-
A.4. Fourier Transform
Haar Measure: There is, up positive multiplicative constant, unique countably additive, nontrivial measure on the Borel subsets of G satisfying the following properties:
j is eef-trranslationnnnaaianttt μ(9S) y=(S) for every g G and all Borel sets S G.
is finite on every compact set: v(K) < oo for all compact K C G.
μ is outer regular Borel sets 8 S G: ((S) (n(()(U) $ S U,U open).
/ is inner regular on open sets U C G: p(U) = sup(u(K): K C U, K compact).
Fourier Transform on Compact groups: A notable and useful property of compact groups is that the set of (iso- morphism classes of) their irreps is countable (Robert 1983). This fact underpins powerful generalization of classical Fourier analysis to the setting of compact groups, enabling square-integrable functions to be decomposed into frequency components indexed by group representations. Let f E C2(G), where L2(G) denotes the space of complex- valuudd square- integrable functions G with respect to the (normalized) Haar measure. For each irrep p. the Fourier transform of f at p defined as
F(e) = Io ((0) p(g) dp(g),
where F(p) is complex matrix of size dp dp- The inverse Fourier transform allows one reconstruct the original function f from its Fourier coefficients via the formula
f(g) = - Tr (f())
where the convergence is in the Ce-sense. This formula justified by the Peter-Weyl theorem, which states that the matrix coefficients of all irreps form complete orthonormal basis for (2) Consequently, any square integrable function G can be expressed as (generalized) Fourier series in terms of these basis functions.
15

---

Stecrable Transformers for Volumetric Data
B. Proof of Theorem 1
Proof. Throughout the proof we will assume a single input channel (dmaddl = 1). Fix (t, It) SE(d). By equation ( under the action of (t, RR any imput f transforms as
(")() - [tt, It) (') - p (XXx +t).
Analogously, the query, key and value vectors equations
12) [0)-
transform as
( (x) [t. R) - f w = p()) fd (Rx +t p() q (0) (Rx +t)
(20)
k(l) (x (tt, (
+ p()) (x - (()) nT W) K p()) (x,y) ple) (x,y),
(21)
v( (x.y) [tt (x)) + +
(x,y) - p(R) (T (Rx
(22)
respectively. The score function is computed using a scaled dot product. Using the fact that p unitary, the score function transforms as
sxx,
Σ dx 1/2
vec (x)) vec (Kl0) (x!)(x).))
(by equation
[3)
Σ dx vec (((R))q(P) (Rx t)) vee ((R))) + +
(by equation
20
21
=x 1/2 ((R) ( (eX+ (O(R) ( + (())(())))
(since we assumed dmodel = 1)
1xx (Rx (R))))))(()))) +t) + + 12 ( ())
- dx 1/2 Edx// q (p) (Rx ( (RX + 2 1)² ())
(since p(R) is unitary)
K 1/2 (0) (0) (Rx+ t)+ ( ( + t, +t))
Σ dx q
1/2 (p) (Rx+ + t)²t (( (,,t
+
P
Σ
1/2 (())
Rxx (X Ry +t)+ + Σ 1/2
(by equation
11
Σ d 1/2 vec (q (RX) vec (kl) vec(kkk)
t, Ry + Σ 1 (R
(since we assumed dmnddl =
=s(Rx+ + t, Ry +
dx
Ry+
t. + t),
(23)
(by equation
13)
where (0) Rx t, Rly t) = p( (R)P(e) (x,y) p(P) + tt Ry + tt). The attention scores are obtained by taking the softmax of the scores 23 Therefore, the attention scores transform
(xxy)) y)
16

---

Steerable Transformers for Volumetric Data
exp (1s(x, y)l)
Secr exp (1s (x. dz
(by equation
(1)
s( RX + t, Ry + ++++++++ Ep 1/2 (BX Rty + +t))
exp
S1((,, exp
Pix + t, Rz + +1 + Ep 1 (RX Pxx + t, Rl + t))) de
(by equation 23)
s + Ry + t) Ep 1/2 (P) +B Ryy t))
cxp
(24)
(0)m exp ssx + t, Rz + t) E, dx 1/2 ( [Rx (Rx+ ( Rx + t, Rz + t))) dz
-1 frm) (if is the support of f then (t, R) (2( is the support of Rt) R)
Define the expression in 24 8(xx + t, Ry + t). Using these attention scores, the output of the attention mechanism frm for any input frh transforms as
(x)
= -Sacm y)(e)))
(by equation
15
3(xx+ t, Ry +) ((Rx+ (x)) dy
(t,R) 1_(2(())
(by equation
22
=p([(}
8[Rx + t, Ry + |
t)] dy
1
(4. -(c.
|Rxx t, Ry +t) (PP
dy
(.00f)))00))))))))))
=p(R)'
8(Rx tt t, Ry + t
((.P)-
+ +R(
B(Rx+ + t, Ry + t) (Rx + t. Ry + tt
(25)
(t. 1
(by equation
12 and definition A (1))
According to the statement of the proposition, the output foul should transform as
(ou p) + p(R)t yor Rx + t. - p(R)
a{Rx + +ty))(()) (Rx+tty)) t,y) dy
2((
a(Px + t, + t))((( (Rx t, + +t) dy.
+
(26)
(E.R) (.)))))00))))))
The transformations (25) and 26 are equivalent iff
J J(.0)- 8(xx + t, Ry t)y(() [xx +t, Ry + t) /
3(Rx t, Ry t)A + t. Ry + ty
(0.,0))1-000
a(Rx + t, Ry + t)y (0) fxx t, Ry t) dy
(27)
((. i
for all inputs fü and for all weights w (((e) w(()) for all irreps p. Note that, if (p) 0 for all irreps p. then a by equation (24 and hence, the equality in 27 holds for all fum w(e) w(r) wofe) Conversely, suppose the equality 27
17

---

Steerable Transformers for Volumetric Data
holds for any f W W W In particular, setting W 00 by equations 20 and 24 therefore equation 27 implies
we get a = B x 1, and
J mm)m))))
alRx t, Ryy + ()xx + dy + /
a(1x nyy t)A (Rx + t, Ryy + t)
(t.R) 1.004))
a(Rx + t, Ry + t))(l) Ry t) dy
(1.R)-1--1("")
a(Rx + Ry + + (2) (Rx tt t) dy 0
(t.R)-1-9(())))
A + Ry tt
(28)
(t.R))-1-2(9)) (since x 1)
The equality (28) holds for any compactly supported fm. In particular, consider a class of functions supported on
ball of radius I/n around Ry, t, with n N and fixed No E Rd, denoted by B1/n + t). Then, setting ((f"") B1/n( (Ryo + t) in 28 we have
A ( (Rx tt, Ry + dy
for all n nEN
(L, BI/
A (Rx+,t + dy =0
for all n n N
(Ix +t, Ry t) dy
for all N
(N)1 Ba/n(No)
lim 100 Bi B1/n(t0)
A + Ry + dy
(e) (Rx + t. Ryo + +t) =0
(29)
(using continuity of p(p) and Lemma 1
Since Yo Rd was arbitrarily fixed, the equality in 29 hold for all yo E Rd. Hence, the equality in 27 holds iff
(60 (Rx + Ry p()))(() (x,))
for all x, Rd and irreps p. Since (t, R) SE(d) was chosen arbitrarily, the result holds for all (t, R) SE(d), which completes the proof.
Lemma 1. Suppose f: Rd + C is a function continuous at xo E Rd. Let (xo)), defined as
Bx(xo) = (x E Idd [ - -xl22 < sr),
denote ball of radius around Xo- Then,
lim (xox Br(xo)
f(x) dx f(xo)
where (xa)) is the Lebesgue measure of (xo)
Proof. By continuity of ff at Xg, for 0. there exists 8 > 0 such that, for any r
[lf(xx) - fxx 12 Se
(30)
for all x E (xo). Therefore, we have
1 f B.(x)
[x,)) / B,(x))
f(x) dx f(xo)
f(x) - f(xo) dx
[xxx)
2
18

---

Steerable Transformers for Volumetric Data
1
1lf(x)
f(xo)l12 dx
(xn)
Br(xa)
ex
[x,(x Br(xa)
(by equation (30)
(31)
The inequality in (11 holds for 0 S d. Hence,
lim
d(x) ((x0)
= 0.
[x,)] B,o)
C. Experiment Details
this section, we provide additional details on the datasets, architectural variations, and hyperparameters all experi- ments conducted in this work.
C.1. Rotated MNIST
Type
Method
Error%
Parameters
SVM (Larochelle et al.
2007)
11.11
TIRBM (Sohn Lee 2012)
4.2
Miscellaneous
Tl-Pooling (Laptev al. 2016
1.2
P4CNN (Cohen & Welling
2016
2.28
25k
Harmonic Net (Worrall al.
2017
1.69
33k
Equivariant Convolution
RotEqNet (Marcos et al. 2017
1.09
Weiler et al. (Weiler et al.
(018b)
0.71
6M
E2CNN (Weiler & Cesa
2019)
0.69
6M
a-R4 CNN (Romero et al.
2020)
1.69
73k
GSA-Nets (Romero Cordonnier
2021
2.03
44k
Equivariant Attention
GE- VIT (Xu et al. 2023
1.99
44k
Our (k 4)
1.18
1.13M
Our (k 8)
0.97
2.24M
Table 2: Comparison of the performance of the equivariant attentive architecture on the Rotated MNIST dataset.
The network architecture in this experiment consists of three convolutional blocks with progressively increasing numbers of channels. These blocks are separated by average pooling layers, which down-sample the feature maps improve com- putational efficiency. The output is flattened and passed through fully connected layers that apply linear transformations. batch normalization, ReLU activation, and dropout to prevent overfitting. The final layer maps the output to ten classes. In the case of the steerable transformer, a transformer block with a single layer follows the convolutional blocks.
The networks were trained using the Adam optimizer (Kingma Ba 2014), starting with a learning rate of 5 x 10-3 which was reduced by factor 0.5 every 20 epochs. along with weight decay of 5 10-4 for 150 epochs. A batch size of 25 used, and training the largest model took 4 hours 16GB GPU.
Table 2 compares the performance of the steerable transformers with other attention-based methods reported for this dataset. Our approach significantly outperforms the other attention-based methods, which are standalone, while ours is built on top of convolutions. The only methods that surpass our results, by Weiler et al. (2018b) and Weiler & Cesal (2019). use a Fourier cutoff 16, while we use a cutoff of 8. We believe this accounts for the performance difference, which is due to implementation and resource limitations, as higher Fourier cutoffs lead out-of-memory issues.
19

---

Stecrable Transformers for Volumetric Data
Method
Accuracy
Parameters x x10%)
Point CLoud
Voxels
ECC (Simonovsky & Komodakis
2017
90.8
Point
Cloud
SO-Net (Li ct al. 2018
93.9
2.5
Rot-SO-Net (Li et al. 2019
94.5
2.5
3D ShapeNets (Wu et al. 2015)
83.5
12
VRN (Brock et al. 2016
91.3
VoxNet (Maturana & Scherer
2015
92.0
0.92
Voxel
Rotation
SO(3) Rotation
FusionNet (Hegde & Zadeh
2016)
93.1
120
ORION (Sedaghat et al.
2016)
93.8
0.91
Figure Examples from the ModelNet10 dataset are shown in various formats: point cloud, voxel representa- tion, and rotated perturbations of voxels.
Cubenet (Worrall & Brostow 2018)
94.6
4.5
Our
91.1
0.92
Table 3: Comparison of performance ModelNet10 with rotation pertur- bation. Other methods used both train and test time augmentation, while we applied augmentation only during test time, not during training.
C.2. ModelNet10
The point cloud data with 2048 points is available at htltttss//ggtthue Similar to the Rotated MNIST experiment, we utilized three convolutional blocks. The features for each irrep are in 1:1:1 ratio. The number of features per irrep increases with each convolutional block, and the data is downsampled using average pooling. For the steerable transformer, this convolutional encoder is supplemented with a transformer block containing single transformer layer. After additional normalization step, the output is flattened and passed through fully connected layers that reduce the dimensionality to 128 features. To ensure training stability and prevent overfitting, batch normalization, ReLU activation, and dropout are applied, with the final output layer classifying the data into ten classes
The networks were trained using the Adam optimizer Kingma & Ba 2014), with an initial learning rate of 1 x 10-3 which decreased by factor of 0.5 every epochs, for a total epochs. A batch size of 5 was used, and training the largest model took 12 hours on a 16GB GPU.
Table compares method to other con-attention-basel methods the z-rotated version of the dataset. Unlike approach, which does not use augmentation during training, all other methods were trained with augmentation involving 12 uniformly stratified rotations along the z-axis. To ensure 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 these other methods.
C.3. PH2
The data for the experiment is available at work begins with two convolutional blocks, each composed of steerable convolutional layers, non-linearities, and batch normalization, which progressively downsample the input and extract relevant features. Then, a transformer encoder used capture global dependencies the data. After this, the model uses two convolutional blocks to upsample and refine the extracted features. In between each block, the features are upsampled using bilinear interpolation. Finally, the features are restored to the original resolutions. In the final step, the vector norm is computed, converting the complex-valued feature maps into real values to produce the logits of each class.
httm The net-
The networks were trained using the Adam optimizer (Kingma & Ba 2014), starting with an initial learning rate of x x which decreased by a factor of 0.5 every 20 epochs, for a total 100 epochs. We used a batch size of 1, and training the largest model took 2 hours on 16GB GPU. larger batch size could not be used due out-of-memory errors.
20

---

Steerable Transformers for Volumetric Data
Input
Ground Truth
Probability
Predicted
PHZ
BraTS
Figure 6: More segmented examples from the test datasets of PH2 and BraTS.
C.4. BraTS
The data for the experiment is available at for this task, consisting 484 images. We split this dataset into training, validation, and test sets for our experiments.
Only the training data is available
The network begins with two convolutional stem blocks that use steerable convolutions extract multi-resolution features. These blocks include non-linearities and batch normalization for stable training, well average pooling layers for downsampling. The features for each irrep in a ratio of 8:4:2, with the number of features increasing in subsequent layers according to this ratio. The encoded features are then passed through transformer encoder, which captures long- range dependencies and spatial relationships using multiple layers and positional encodings. The decoder consists of convolutional head blocks and trilinear interpolation kernel that upsample the features. The final output generated by taking the absolute value of these complex features, corresponding to the constant representation.
The networks were trained using the Adam optimizer (Kingma & Ba 2014), with initial learning rate of 1 10-?, which reduced by factor of 0.5 every 20 epochs, for total of 100 epochs. A batch size of was used, and training the largest model took 40 hours 16GB GPU. A larger batch size could not be used due to out-of-memooy errors.
21


Reviews:
[{"review_id": "g0bhFNd1WW", "reviewer": "Reviewer_skHN", "summary": "This work introduces a steerable equivariant transformer operating in the spectral domain. The method calculates attention scores exclusively between Fourier embeddings associated with matching irreducible representations (irreps). These scores are then combined linearly to construct the final attention matrix, which is used to generate output tokens. The architecture is mathematically proven to maintain equivariance by design and demonstrates superior performance over baseline methods in both 2D and 3D tasks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed review and thoughtful suggestions. We have carefully considered the questions and concerns raised, and provide detailed responses below.\n\n**Equivariance Error**\n\nWe thank the reviewer for the insightful suggestion about equivariance error. In response, we have included a visual analysis of the equivariance error, following a methodology similar to that of Sosnovik et al. (2020). Specifically, we evaluated our model on randomly selected samples from the MNIST and ModelNet10 datasets, along with their rotated counterparts. The equivariance error is measured using the following expression:\n\n$$\\Delta =\\frac{||M(T_\\theta[f]) - T_\\theta[M(f)]||^2}{||M(T_\\theta[f])||    ||T_\\theta[M(f)]||},$$\n\nwhere $M$ denotes the model and $T_\\theta$ represents a rotation by $\\theta^\\circ$. For each input image $f$, we computed $\\Delta$ across $\\theta \\in [0^\\circ, 90^\\circ]$ in $5^\\circ$ increments, averaging the results over 50 random samples per dataset. For the ModelNet10 dataset, we considered two rotation types: around the $Y$-axis and around the $Z$-axis.\n\nThe resulting plots of $\\Delta$ as a function of $\\theta$ for varying numbers of transformer layers are available [at this anonymized GitHub link](https://anonymous.4open.science/r/Rebuttal-7B8B/Equivariance_error.ipynb). We observe that the equivariance error grows slightly with the number of layers, eventually stabilizing at approximately 1.5\\% for Rotated MNIST and 10\\% for ModelNet10.\n\nSosnovik at al., (2020). Scale-Equivariant Steerable Networks. The International Conference on Learning Representations.\n\n**Positional Encoding Intuition**\n\nWe would like to clarify the intuition behind the positional encoding formulation. The central idea is to encode not only the relative distance between two patches but also their relative direction. Since we use relative positional encoding, the encoding between two patches $x_i$ and $x_j$ is based on the difference vector, i.e., $P(x_i - x_j)$. To incorporate directional information, in a steerable setting, we extend this to include an angular component: $P(x_i - x_j, \\theta)$ ($\\theta\\in (0,2\\pi)$ in 2D case). This formulation is inherently translation invariant by construction. Additionally, we require the encoding to behave consistently under rotations. That is, if the input image is rotated by an angle $\\phi$, $x \\mapsto R_\\phi x$, the relative position vector undergoes the same rotation, and in this case, the angular component of $P$ should shift accordingly, resulting in the transformation $P(x, \\theta) \\mapsto P(R_\\phi x, \\theta + \\phi)$. In steerable neural networks, we operate in the Fourier domain, where such angular shifts correspond to modulations by Fourier basis functions. Specifically, the positional encoding $P(x, k)$ transforms under rotation as $P(x, k) \\mapsto e^{-\\iota k\\phi} P(R_\\phi x, k)$, reflecting the equivariance properties of the Fourier basis. This insight motivates our design choice in Equation (6), where we define the encoding as $P(x, k) = \\phi(r, k) e^{-\\iota k \\theta}$, with $r = ||x||$ and $\\theta$ denoting the polar angle of $x$. The function $\\phi(r, k) \\propto r^{-2}$ modulates the influence of distance, assigning greater weight to nearby patches. This formulation extends naturally to higher dimensions, leading to the 3D generalization presented in Equation (7).\n\nWe appreciate the reviewer’s suggestion and have created a visualization to better convey the intuition behind our equivariant positional encoding (available [at this anonymized GitHub link](https://anonymous.4open.science/r/Rebuttal-7B8B/positional.png)).", "claims_and_evidence": "**Claim 1: Design of Steerable Transformer:**  Authors provided proof of the equivariance of the model. However, this claim will be bolstered if authors can empirically report the equivariance/invariance error (see Sec 6 [1]).\n\n**Claim 2: Performance gain:** Works shows performance gain against steerable convolution. However, authors should also report the consistency of the model over the orbit of an input for each dataset. For example, in the case of ModelNet10, any rotation should not change the predicted probability scores.\n\n\n[1] Scale-Equivariant Steerable Networks", "methods_and_evaluation_criteria": "The chosen datasets are valid. However, as mentioned earlier, additional evaluation metrics, such as equivariance error", "theoretical_claims": "Yes. I went through the proof, and I found it to be correct.", "experimental_designs_or_analyses": "The experiments are conducted on four different datasets. As I mentioned earlier, additional metrics, such as equivariance error and consistency over the orbit of an input, should be reported for completeness.", "relation_to_broader_scientific_literature": "The key contribution of the paper is the proposal of an equivariant transformer for volumetric data. Even though there exist models of equivariant transformers for vision, i.e., 2D data, the work proposes a more general framework for such transformers, which I believe will be of great interest to the machine learning community.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "The positional encoding design would benefit from greater clarity and intuition. While the authors provide a mathematical proof validating their approach, this formal validation should be preceded by a more accessible conceptual explanation.\n\nI strongly recommend including a visual representation—specifically, a figure illustrating the positional encoding mechanism for a simple 2D case.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vHjGbrZrYb", "reviewer": "Reviewer_M9Pw", "summary": "This paper explores the steerable (SE3 equivariant) vision transformer for volumetric data. The proposed framework was built upon existing steerable convolution and vision transformers. Steerable transformers was evaluated on 2D and 3D classification (rotated MNIST and ModelNet10) and segmentation (skin images and brain MRI) and outperformed the original steerable convolution. In addition to that, additional comparisons with previous studies spanning from 2015 to 2019 were conducted on ModelNet10 with z rotation perturbation. \n\nThe contribution was claimed to integrate ViT and steerable convolutions to yield steerable transformers for 3D volumetric image data (instead of a point cloud, which was extensively studied).", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed review and thoughtful suggestions. We have carefully considered the questions and concerns raised, and provide detailed responses below.\n\n**Complexity Calculations**\n\nWe acknowledge the reviewer’s concern regarding the complexity calculation. It is true that the assumption $N \\ll C$ does not generally apply to raw image inputs. However, in our case, $N$ denotes the number of patches in the *downsampled* feature map produced by the convolutional encoder—not the number of pixels in the original image. For example, in our BraTs experiments, two pooling layers (with sizes 8 and 4) reduce the input to a feature map of dimensions $7 \\times 7 \\times 4$, yielding $N = 196$ patches. The transformer then operates on this representation using 64, 32, and 16 channels for the $\\ell = 0, 1, 2$ components, respectively, totaling 240 channels (each $\\ell$ component is a $2\\ell + 1$-dimensional vector). In this configuration, $N < C$, and our complexity analysis reflects this setting, where the computational costs of attention and convolution are comparable. Please refer to our response to reviewer URnE for a more detailed discussion.\n\nIn line with the reviewer’s suggestion, we have also updated Table 1 to report both training and inference times for each model. The revised table is available in Table 1 [at this anonymized GitHub link](https://anonymous.4open.science/r/Rebuttal-7B8B/table.pdf). The results show that steerable transformers and steerable convolutions with similar parameter counts have comparable runtimes.\n\n**Experimental Design**\n\nThe framework presented in our manuscript is designed to operate in general $d$ dimensions, encompassing both 2D and 3D settings. While our discussion of group representations are motivated by the 3D setting, we believe that the proposed approach is also novel in the 2D case. To our knowledge, prior work on equivariant vision transformers for image data has not employed this formulation. For this reason, we have demonstrated our method in both 2D and 3D contexts. \n\nWe agree with the reviewer that medical imaging applications is particularly well-suited for this architecture. In response to this comment, we will expand the manuscript to include more such experiments and discussions regarding relevant works, including those by Sangalli et al. (2023) and Kuipers et al. (2023). We appreciate the reviewer for pointing us in this direction.\n\n**Question 1 and 2**\n\nWe thank the reviewer for the feedback on our visualization. The BraTs dataset has resolution $4\\times 240\\times 240\\times 155$, where the first dimension corresponds to four MRI modalities (T1, post-contrast T1-weighted, T2, and T2-FLAIR) rather than traditional channels. In our initial submission, we averaged over the modality and the first spatial dimension, followed by a 90-degree rotation:\n```python\ninputs_plot = torch.rot90(torch.mean(inputs, dim=(0,1)))\nprobs_plot = torch.rot90(torch.max(probs, dim=dim)[0])\n```\nHowever, we acknowledge that this view offered limited insight. To improve interpretability, we revised the visualization as follows:\n- Display only the T1 modality.\n- Use the transverse (z-axis) view by averaging over the third spatial dimension.\n- For probabilities, take the mean across the last spatial dimension.\n\nUpdated code:\n```python\ninputs_plot = torch.rot90(torch.mean(inputs[0], dim=2))\nprobs_plot = torch.rot90(torch.mean(probs, dim=2))\n```\nThe improved visualizations are available [at this anonymized GitHub link](https://anonymous.4open.science/r/Rebuttal-7B8B/brats.png). We appreciate the reviewer’s suggestion, which helped enhance the clarity of our presentation.\n\n\n**Question 3**\n\nWhile some equivariant Vision Transformers (ViTs) report lower parameter counts, they are typically much more memory-intensive. For instance, Romero et al. (2020) noted that \n> “the $\\alpha$-p4 All-CNN requires approximately 72GB of CUDA memory, as opposed to 5GB for the p4-All-CNN. This is due to the storage of the intermediary convolution responses required for the calculation of the attention weights.”\n\nRomero et al. (2021) also reported slow training times, even with multiple GPUs. Since GE-ViT (Xu et al., 2023) builds on this framework, it expected to face similar challenges. We believe the reduced parameter count in such models is a consequence of significant memory and compute overhead. In contrast, our method strikes a balance by augmenting steerable convolutions with a transformer applied only to down-sampled features, greatly reducing computational demands. GGiven this, we compare our model with both equivariant attention-based and convolutional methods. For example, Weiler et al. (2021) uses ~6M parameters, which we consider a fair baseline. We acknowledge that omitting the parameter count in the original manuscript was an oversight. The revised table is available in Table 2 [at this anonymized GitHub link](https://anonymous.4open.science/r/Rebuttal-7B8B/table.pdf).", "claims_and_evidence": "1. The claim of equivariance was supported by Proposition 3.1. \n2. The claim regarding complexity analysis (Section 3.5) was not supported by convincing evidence; please refer to the Theoretical Claims section. \n3. Introducing transformers' attention mechanism enhances the baseline performance of steerable convolution.", "methods_and_evaluation_criteria": "While the proposed methods make sense, the main discovery of this paper seems to be how to maintain the equivariance of the output from a steerable convolution throughout a vision transformer via some slight modifications to position encoding and nonlinear activation functions. \n\nThe evaluation criteria of classification and segmentation performance are standard.", "theoretical_claims": "I checked the correctness of theoretical claims. I found the complexity analysis (Sec 3.5, lines 310-329) has flaws. Specifically, when analyzing the complexity of the proposed steerable attention, the author ignores a significant item using an assumption that almost does not exist in the real world. The original manuscript writes, and I quote (lines 314-316), \"SA has a complexity $\\mathcal{O}(NC^2+N^2C)$, and FFN adds $\\mathcal{O}(NC^2)$, totaling $\\mathcal{O}(NC^2+N^2C)$. **When $N<<C$, the $\\mathcal{O}(NC^2)$ term dominates.**\" The authors then conclude that (Lines 323-326), \"In summary, with a fixed kernel size, the computational complexity of steerable attention matches that of convolution.\"\n\nThis **When $N<<C$ assumption**, unfortunately, does not exist in many real-world 3D volumetric data, especially in medical images. The BraTs brain MRI dataset used in this study has a matrix size of 240x240x155, and real-world high-resolution chest CT may have a size of 512x512x300+; in either case, the final sequence length fed into ViT after the steerable convolution would not have an $N<<C$.\n\nThe correct time complexity of steerable transformer, waiving that unrealistic assumption, would be $\\mathcal{O}(Nd_{\\rho}^2C^2+N^2d_{\\rho}C)$, not as claimed in Line 318 $\\mathcal{O}(Nd_{\\rho}^2C^2)$. Compared to $\\mathcal{O}(Nd_{\\rho}^2C^2k^d)$ of steerable convolution (note that N for ViT is sequence length and N for conv is pixel number), it is clear that steerable transformer has higher complexity for real-world volumetric data. \n\nTherefore, the time complexity analysis has flaws, and to better evaluate the complexity, it is better to use other practical metrics, such as training and inference throughput time per fixed batch, floating-point operations per second, or multiply-accumulate operations.", "experimental_designs_or_analyses": "The experimental designs and analyses are not comprehensive. \n\nI am unsure why the authors included two 2D and two 3D datasets in the experiments. The main aim of this paper is to integrate steerable conv with ViT for volumetric data. Therefore, there should be a focus on experiments on 3D volumetric datasets (which were dominated by medical images in the real world). Additionally, a few studies focused on volumetric medical images should be compared, including [1] and [2]. \n\nI also want to point out that, when comparing to other Equivariant Attention methods in Table 3, α-R4 CNN (Romero et al., 2020) has only around 75k params, and GE-ViT (Xu et al., 2023) has only 45k params, compared to around 1M params of the proposed steerable transformer. While the authors argued in lines 901-902 that the performance difference among E2CNN (Weiler & Cesa, 2019), Weiler et al. (Weiler et al., 2018b), and the proposed steerable transformers come from the computation limitation, it is not explained in the current manuscript that the proposed framework is also considerable larger than GE-ViT and α-R4 CNN. \n\n[1]: Sangalli, Mateus, et al. \"Moving frame net: SE (3)-equivariant network for volumes.\" NeurIPS Workshop on Symmetry and Geometry in Neural Representations. PMLR, 2023.\n\n[2]: Kuipers, Thijs P., and Erik J. Bekkers. \"Regular se (3) group convolutions for volumetric medical image analysis.\" International Conference on Medical Image Computing and Computer-Assisted Intervention. Cham: Springer Nature Switzerland, 2023.", "relation_to_broader_scientific_literature": "Previous studies mostly focused on equivariant networks on 2D and 3D point cloud analysis. This paper focuses on designing SE(3)-equivariant ViT for volumetric image data, which has practical value in volumetric medical image analysis.", "essential_references_not_discussed": "Some highly relevant and domain-specific previous studies are not discussed. \n\n[1]: Sangalli, Mateus, et al. \"Moving frame net: SE (3)-equivariant network for volumes.\" NeurIPS Workshop on Symmetry and Geometry in Neural Representations. PMLR, 2023.\n[2]: Kuipers, Thijs P., and Erik J. Bekkers. \"Regular se (3) group convolutions for volumetric medical image analysis.\" International Conference on Medical Image Computing and Computer-Assisted Intervention. Cham: Springer Nature Switzerland, 2023.", "other_strengths_and_weaknesses": "Other strengths:\n1. The presentation was good \n2. The discussion of existing works was extensive and comprehensive while lacking acknowledgment of a few domain-specific papers.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. The visualization of the BraTs dataset looks weird and does not look like a brain MRI; how was it generated? \n2. How do we interpret the probability map for BraTs datasets shown in Fig. 3 and 5? It seems to have no correlation with predicted labels. \n3. Why is the parameter number of other comparing methods not listed in Table 3? The α-R4 CNN (Romero et al., 2020) has only 73.13k params compared to ~ 1M params of the proposed Steerable Transformer. And GE-ViT (Xu et al., 2023) has only around 45k parameters. \nIs the performance gain coming from the difference in parameter number in this case (i.e., proposed vs. GE-ViT and α-R4 CNN).", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HbXPpsEOoH", "reviewer": "Reviewer_GYnN", "summary": "The authors present an SE(3) steerable attention layer. Specifically, input features are assumed \nto be equivariant tensor features in the specific form as is output by SE(3) steerable networks \nof (Weiler et al. 2018-2019). The authors key contribution could be understood as a novel form of \npositional embedding which when passed through the standard attention layer (modified only to compute \nthe dot product between complex valued features), preserves SE(3) equivariance. Comparisons are performedonly against the relatively dated steerable convolutions proposed by (Weiler et al. 2018-2019)", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed review and thoughtful suggestions. We have carefully considered the questions and concerns raised, and provide detailed responses below.\n\n**Patchification**\n\nWe would like to clarify the reviewer’s concern regarding \"patchification\". In ViTs, patchification refers to dividing an image into patches, flattening them, and projecting them via a learnable linear embedding. This operation is equivalent to using a convolutional layer with stride equal to kernel size. In fact, most ViT implementations, implement patchification using a strided convolution. This perspective aligns with the original ViT paper (Dosovitskiy et al., 2020), which states:\n\n>\"As an alternative to raw image patches, the input sequence can be formed from feature maps of a CNN...As a special case, the patches can have spatial size 1x1, which means that the input sequence is obtained by simply flattening the spatial dimensions of the feature map and projecting to the Transformer dimension.\"\n\nSo, patchification can be thought of as a special case of a convolutional encoder. These downsample the volumetric input to a grid of lower resolution tokens which are then passed into the transformer. However, we did not emphasize this in the manuscript as it is not a novel take on ViTs. Thus, contrary to the reviewer’s concern, our method does support patchification in a manner consistent with standard ViTs, adapted appropriately to the steerable and volumetric setting.\n\n**Comparison with Weiler et al**\n\nWe appreciate the reviewer’s feedback and the opportunity to clarify our contributions. Our use of steerable features is not a design choice but is motivated by theoretical foundations. As shown by Cohen et al. (2017) and Kondor et al. (2018), steerable convolutions (Weiler et al., 2018–19) represent the most general class of equivariant linear maps. Our work builds on this foundation by introducing a transformer architecture that operates on these structured feature fields. While steerable convolutions may face scalability limitations, they remain theoretically sound and widely adopted. Our aim is to enhance, not replace, their capabilities with global attention.\n\nWe would like to clarify the reviewer’s comment regarding the performance of steerable convolutions.Table 3 of our supplementary material compares our method with other equivariant approaches on Rotated MNIST. To the best of our knowledge, Weiler et al. (2019) still represents the state of the art on this benchmark. It is unclear which specific work the reviewer is referring to in stating that Weiler et al. (2018–19) are far from state of the art. If the reviewer is aware of more recent methods with better performance, we would be grateful for references and will gladly include them in the revision.\n\n**Vector Neurons**\n\nWe thank the reviewer for pointing out Vector Neurons (Deng et al., 2021), but we must emphasize that the methodology proposed in that work operates in a fundamentally different setting from ours. Vector Neurons are specifically designed for point cloud data, where inputs are of the form $V\\in\\mathbb{R}^{N\\times 3}$, ($N$=number of points). In contrast, our work focuses on volumetric data of the form $\\mathbb{R}^{C\\times K\\times K\\times K}$, ($K$=resolution and $C$=channels).\n\nIt is not clear to us how the Vector Neurons framework can be applied to the volumetric setting. For instance, their linear layer (Sec. 3.1) involves right-multiplication of the input by a learnable matrix. A comparable operation in our setting would correspond to a convolution with filter os size $1^3$-known to have limited expressiveness. Furthermore, their nonlinearity (Sec 3.2) operates directly on the outputs of such a linear transformation, and is defined as:\n\n$v'=q\\text{ if }q^Tk\\geq0$\n\n$=q-(q^Tk)k/||k||^2\\text{ else where}\\quad q=WV,k=UV$\n\nThis relies on the comparison of the inner product value with zero, which is not meaningful in our setting where features lie in complex space due to the Fourier-based representation. Consequently, applying this nonlinearity in our framework would not be straightforward,. We are unsure by what the reviewers mean by small modifications in this context. It is also unclear to us why this nonlinearity is expected to outperform norm-based ones (Worrall et al., 2017), which have been well validated for 2D data. In contrast, to the best of our knowledge, the Vector Neurons nonlinearity has only been evaluated on point cloud datasets (ModelNet40, ShapeNet) and has not been demonstrated in datasets we evaluate. This makes direct comparison difficult and potentially misleading. Given these differences, we do not see a viable path for incorporating the Vector Neurons nonlinearity without substantially redefining the core of our framework. We would be appreciate if the reviewer could clarify how they envision Vector Neurons being adapted to our setting, and what specific advantages are expected over established volumetric methods.", "claims_and_evidence": "The authors claim to present an attention layer which preserves the specific type of SE(3) equivariance\nenjoyed by the features of the type in (Weiler et al. 2018-2019). The authors provide a proof which to my\ninspection is sound and does verify their claim.\n\nThe authors also claim that their steerable attention layer improves the performance of existing \nSE(3) steerable convolution networks. While the reported results make this claim technically correct,\nthe performance increase is at best marginal and several of the datasets considered are toy. More generally,\nin the supplement, the results of additional methods not based on the method of (Weiler et al. 2018-2019) are shown and the authors proposed approach performs significantly worse.", "methods_and_evaluation_criteria": "The authors proposed method appears to be technically sound, if limited. The evaluation is\nextensive enough, though the results are not compelling.\n\nOne of the chief limitations it that it appears the method cannot \"patchify\" volumetric inputs\nas is standard in vision transformers -- i.e. that it is assumed to operate with a patch size of 1^3\nand therefore likely scales poorly. Patchification is key to the success of vision transformers, and the \ninability of the proposed method to do this is a significant weakness.", "theoretical_claims": "The authors main theoretical claim is that their proposed attention layer is equivariant, the the \nprovided proof appears to be correct.", "experimental_designs_or_analyses": "See above.\n\nIn general, the results are not compelling and do not convince me that the proposed method offers a \nmeaningful improvement over existing methods.", "relation_to_broader_scientific_literature": "The key contributions are unclear. What the authors present is a specialized attention mechanism\ndesigned to work on a specific set of features output by the now dated methods of (Weiler et al. 2018-19)\nwhich have scalability issues and are now far from state of the art. Furthermore, the authors proposed approach\ndoes not improve the performance steerable convolutions to be comparable to existing state of the art methods.", "essential_references_not_discussed": "The authors do not discuss Vector Neurons (Deng et al. 2021) which have been show to far outperform\nother equivariant nonlinearities, and in particular those discussed and employed by the authors. \nIt appears that with a small modification, this nonlinearity could be used with the \nauthors approach (since \\rho(R) are orthogonal).", "other_strengths_and_weaknesses": "This paper is well written, but has has several important weaknesses: \n\nThe novelty of the method is limited -- the authors real contribution appears to be the positional embedding,\nit is already well known that orthogonally-equivariant features make the standard attention mechanism equivariant.\n\nMoreover, while the title of the paper implies that the authors are presenting a novel transformer, they are really just \npresenting a specialized positional embedding. Important practical considerations, like an equivariant patchification method, are not considered and limit the scalability of the method.\n\nIn general, the results are not compelling. The steerable convolutions proposed by (Weiler et al. 2018-19) are now \nfar from state of the art (as shown in the results reported by the authors in the supplement). This method offers what appears to bea way to marginally improve their performance, but does not come close to state-of-the-art.\n\nIn addition see the comments on nonlinearities above.", "other_comments_or_suggestions": "See above.", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UEGjEHPuAu", "reviewer": "Reviewer_URnE", "summary": "This paper introduces a new equivariant transformer architecture in two variants, with symmetry group either SE(2) or SE(3), utilizing a steerable (Fourier) basis. In numerical experiments, performance gains are shown if some layers in steerable CNNs are replaced by the novel attention layers.\n\n## Update after rebuttals\nI thank the authors for their thorough rebuttal and additional clarification on my question. That the scaling behavior of this setup is independent of the steerable nature of their architecture and therefore not a benefit of the proposed model is an important point which should be mentioned in the corresponding section. However, this is not a major contribution of this work.\n\nI agree with the other reviewers that the tasks investigated in the experimental section are not very difficult, I see this paper more of a proof of concept of using steerable layers in transformers. This fills a hole in the literature, and does not need to be a new competitive architecture for some specific application domain.\n\nFor these reasons, I’m happy to maintain my score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed review and thoughtful suggestions. We have carefully considered the questions and concerns raised and provide detailed responses below.\n\n**Weaknesses**\n\nWe have included comparisons with other methods on both the Rotated MNIST and ModelNet10 datasets in the supplementary material. Specifically, Table 3 of the Supplementary Material compares our approach against a range of equivariant methods on Rotated MNIST, including both convolutional and transformer-based architectures. For ModelNet10, while we were unable to find any existing equivariant transformer-based methods for direct comparison, Table 4 presents results alongside other equivariant approaches, covering both volumetric and point cloud-based techniques.\n\n\nIn response to the reviewer’s critique regarding the lack of comparison between pure transformer-based and convolutional architectures, we conducted preliminary experiments on the Rotated MNIST dataset. The results are shown in Table 3 [at this anonymized GitHub link](https://anonymous.4open.science/r/Rebuttal-7B8B/table.pdf). \n\nIn these models, we patchified the image and fed the resulting patches into a 3-layer transformer. The transformer output was then pooled and passed through fully connected layers for classification. The number of parameters for these standalone transformer models are approximately $1.10$M and $2.58$M for Fourier cutoffs of \\(k = 4\\) and \\(k = 8\\), respectively—comparable to the models in Table 1. Our findings show that reducing patch size improves accuracy but increases the runtime, and still falls short of the performance achieved by our hybrid model. This aligns with previous observations in literature (e.g., Xiao et al., 2021) that standalone ViTs require large datasets to outperform CNNs, and that convolutional encoders enhance transformer performance. These additional results offer a clearer comparison between standalone steerable transformers, steerable convolutional models, and our hybrid approach in the volumetric setting. We thank the reviewer for this insightful suggestion.\n\nXiao et al. (2021). Early convolutions help transformers see better. Advances in Neural Information Processing Systems.\n\n**Questions**\n\nWe agree with the reviewer that the complexity analysis warrants further clarification. The theoretical computational complexity of the transformer is $O(N^2C + NC^2)$, where $N$ is the number of patches and $C$ is the number of channels. The term $NC^2$ dominates in the regime $N \\ll C$. As reviewer M9Pw correctly noted, this condition does not hold for raw image inputs—e.g., the BraTs brain MRI dataset has a resolution of $240 \\times 240 \\times 155$, so $N \\ll C$ is clearly not satisfied at the pixel level.\n\nHowever, $N$ does not refer to the total number of pixels. In ViTs, $N$ represents the number of image patches after patchification: $N = \\frac{ \\text{Number of pixels}}{k^d}$, where $k$ is the patch size and $d$ is the image dimensionality. While $N$ can still be large for moderate $k$, our approach is not based on applying the transformer directly to raw image patches.\n\nInstead, we use a convolutional encoder to down-sample the input image before applying the transformer. This design is motivated by practical constraints: directly applying a transformer to high-resolution images requires a large patch size to stay within memory limits, which sacrifices local detail; conversely, using small patches leads to memory bottlenecks. Our hybrid convolution-transformer model addresses this by letting convolutions capture local features and the transformer focus on global context.\n\nFor example, in our BraTs experiments, the image is down-sampled using two pooling layers (of sizes 8 and 4), resulting in a feature grid of $7 \\times 7 \\times 4$, i.e., $N = 196$ patches. The transformer operates on these features using $64$, $32$ and $16$ channels for the $\\ell = 0, 1, 2$ components, respectively. Since each $\\ell$-th component is a $(2\\ell + 1)$-dimensional vector, this results in a total of 240 channels. In this setting, $N < C$, and our complexity calculations reflect this regime. Under these conditions, the computational cost of applying a transformer is comparable to that of a convolution.\n\nFinally, we also note that applying a sparse attention mask (restricting attention to nearby patches, as in graph-based models) is indeed a viable way to reduce complexity. However, such an approach undermines the main strength of transformers—capturing long-range dependencies. If only local interactions are needed, convolutions are a more natural and efficient choice.", "claims_and_evidence": "The paper claims that replacing steerable CNN layers by the novel steerable attention layers improves performance. This claim is tested on four diverse datasets. The performance benefits seem small but do lie outside of the standard deviation bands based on 5 runs.", "methods_and_evaluation_criteria": "The datasets seem appropriate for evaluating the performance of this architecture.", "theoretical_claims": "The claimed equivariance of the layers proposed is straightforward to see in the steerable basis.", "experimental_designs_or_analyses": "The authors evaluate their architecture by replacing certain layers of steerable CNNs with their novel steerable attention layers. These combined architectures are not as common as pure CNN or pure attention architectures and I would find it more convincing to see a comparison of the latter.", "relation_to_broader_scientific_literature": "Equivariant transformer architectures are an active area of research. Using (steerable) Fourier bases to construct equivariant architectures is not novel, but it seems that the most straightforward implementation of this, without tensor products and by using norm-nonlinearities as realized here, has not been done before. For this reason, I think this is a valuable baseline to consider.", "essential_references_not_discussed": "None.", "other_strengths_and_weaknesses": "Strengths:\n- Straightforward implementation of an equivariant transformer in a steerable basis\n- Background is well-introduced\n- Diverse datasets in numerical experiments\n\nWeaknesses:\n- No comparison to other equivariant transformers\n- No comparison of pure steerable CNN vs pure steerable transformers", "other_comments_or_suggestions": "I found these typos:\n- No integral over SO(d) in (1)\n- Are the indices correct on the RHS of $s_{ij}$ in section 3.2 (compare indices in $s_{ij}$ in section 2.1)?", "questions_for_authors": "Since transformers scale quadratically in the number of input tokens, I do not understand how the computational complexity of the steerable transformers can match the scaling of convolutions as claimed in section 3.5. Could you clarify?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The manuscript received four high-quality reviews, the authors provided rebuttals, and the reviewers acknowledged and reflected on the rebuttals. None of the reviewers came to a clear negative or positive verdict, the average is slightly positive. The most negative reviewer M9Pw criticizes in particular the insufficient experiments, but acknowledges not being an expert in the field. Otherwise, the majority of comments point out a good but somehow limited contribution and a correct proof with potential detailed issues regarding the complexity calculation. Although the experimental results could be even more convincing, the manuscript is accepted.
Decision: Accept (poster)