Group Equivariant Stand-Alone Self-Attention For Vision
Authors: David W. Romero
Jean-Baptiste Cordonnier
Venue: ICLR 2021 Poster (2021)
Keywords: group equivariant transformers
group equivariant self-attention
group equivariance
self-attention
transformers

Abstract: We provide a general self-attention formulation to impose group equivariance to arbitrary symmetry groups. This is achieved by defining positional encodings that are invariant to the action of the group considered. Since the group acts on the positional encoding directly, group equivariant self-attention networks (GSA-Nets) are steerable by nature. Our experiments on vision benchmarks demonstrate consistent improvements of GSA-Nets over non-equivariant self-attention networks.

Published as a conference paper at ICLR 2021
GROUP EQUIVARIANT
STAND-ALONE SELF-ATTENTION FOR VISION
David W. Romero
Jean-Baptiste Cordonnier
Vrije Universiteit Amsterdam
École Polytechnique Fédérale de Lausanne (EPFL)
rome rogurranuuaauu
Jean-baptiste wow.etteeccooonnieee.eeeee
ABSTRACT
We provide a general self-attention formulation to impose group equivariance to arbitrary symmetry groups. This achieved by defining positional encodings that are invariant to the action of the group considered. Since the group acts on the posi- tional encoding directly, group equivariant self atff-attentonn networks (GSA-Nets) are steerable by nature. Our experiments on vision benchmarks demonstrate consistent improvements of GSA-Nets over non-equivariant self-attention networks.
INTRODUCTION
Recent advances in Natural Language Processing have been largely attributed to the rise of the Transformer (Vaswani et al., 2017). Its key difference with previous methods, e.g., recurrent neural networks, convolutional neural networks (CNNs), is its ability to query information from all the input words simultaneously. This is achieved via the self-attention operation (Bahdanau et al., 2015; Cheng et al., 2016), which computes the similarity between representations of words in the sequence in the form of attention scores. Next, the representation of each word is updated based on the words with the highest attention scores. Inspired by the capacity of transformers to learn meaningful inter-word dependencies. researchers have started applying self-attention in vision tasks. It was first adopted into CNNs by channel-win attention (Hu et al., 2018) and non-local spatial modeling (Wang et al., 2018). More recently. has been proposed to replace CNNs with self-attention networks either partially (Bello et al., 2019) or entirely (Ramachandran et al., 2019). Contrary to discrete convolutional kernels, weights in self-attention are not tied to particular positions (Fig. A. I), yet self-attention layers are able to express any convolutional layer (Cordonnier al., 2020). This flexibility allows leveraging long-range dependencies under a fixed parameter budget.
An arguable orthogonal advancement to deep learning architectures is the incorporation of symmetries into the model itself. The seminal work by Cohen & Welling (2016) provides a recipe to extend the translation equivariance of CNNs other symmetry groups improve generalization and sample- efficiency further (see $2). Translation equivariance is key to the success of CNNs. It describes the property that if pattern is translated, its numerical descriptors are also translated, but not modified.
In this work, we introduce group self-attention a self-attention formulation that grants equivariance to arbitrary symmetry groups This is achieved by defining positional encodings invariant to the action of the group considered. In addition to generalization and sample-efficiency improvements provided by group equivariance, group equivariant self-attention networks (GSA-Nets) bring important benefits over group convolutional architectures: (i) Parameter efficiency: contrary to conventional discrete group convolutional kernels, where weights are tied to particular positions of neighborhoods on the group, group equivariant self-attention leverages long-range dependencies on group functions under a fixed parameter budget, yet is able to express any group convolutional kernel. This allows for very expressive networks with low parameter count. (ii) Steerability: since the group acts directly on the positional encoding, GSA-Nets are steerable (Weiler et al., 2018b) by nature. This allows us to go beyond group discretizations that live in the grid without introducing interpolation artifacts.
Contributions:
We provide an extensive analysis on the cquivariance properties of self-attention ($4).
We provide general formulation to impose group equivariance to self-attention ($5).
We provide instances of self-attentive architectures equivariant to several symmetry groups ($6).
Our results demonstrate consistent improvements of GSA-Nets over non-equivariant ones ($6).

---

Published as a conference paper at ICLR 2021
Figure 1: Behavior of feature represen-
Grorp Scl-Auerrice
tations in group self-attention networks. An input rotation induces a rotation plus a cyclic permutation to the intermediary feature representations of the network. Additional examples for all the groups used in this work as well as their usage are provided in repo/ demo/.
Netwokk
Uhopp Sti-Moesiio
Netrcrs
RELATED WORK
Several approaches exist which provide equivariance various symmetry groups. The translation cquivariance of CNNs has been extended to additional symmetries ranging from planar rotations (Dieleman et al., 2016; Marcos et al., 2017; Worrall al., 2017: Weiler al., 201 8b; Li et al., 2018; Cheng et al., 2018; Hoogeboom et al., 2018; Bekkers et al., 2018; Veeling et al., 2018; Lenssen et al., 2018; Graham al., 2020) to spherical rotations (Cohen al., 2018; 2019b; Worrall & Brostow, 2018; Weiler et al., 201 8a; Esteves et al., 2019a;b; 2020), scaling (Marcos et al., 2018; Worrall & Welling, 2019; Sosnovik et al., 2020; Romero et al., 2020b) and more general symmetry groups (Cohen& & Welling. 2016; Kondor & Trivedi, 2018; Tai et al., 2019; Weiler & Cesa, 2019; Cohen et al., 2019a; Bekkers, 2020; Venkataraman et al., 2020). Importantly, all these approaches utilize discrete convolutional kernels, and thus, tie weights to particular positions in the neighborhood kernels are defined. As group neighborhoods are (much) larger than conventional ones, the number of weights discrete group convolutional kernels require proportionally increases. This phenomenon is further exacerbated by atlentive group equivariant networks (Romero & Hoogendoorn, 2019; Diaconu & Worrall, 2019; Romero al., 2020a). Since attention is used to leverage non-local information to aid local operations, non-local neighborhoods are required. However, as attention branches often rely on discrete convolutions, they effectively tie specific weights particular positions on large non-local neighborhood on the group. As result, attention is bound to growth of the model size, and thus, to negative statistical efficiency. Differently, group self-attention able to attend over arbitrarily large group neighborhoods under fixed parameter budget. In addition, group self-attention is steerable by nature 1) a property primarily exhibited by works carefully designed to that end.
which the
Other way to detach weights from particular positions comes by parameterizing convolutional kernels as (constrained) neural networks (Thomas et al., 2018; Finzi et al., 2020). Introduced to handle trregularly-sampll
data, c.g., point-clouds, networks parameterizing convolutional kernels receive relative positions as input and output their values at those positions. In contrast, our mappings change function of the input content. Most relevant to our work the SE(3) and Lie Transformers (Fuchs et al., 2020; Hutchinson et al., 2020). However, we obtain group equivariance via a generalization of positional encodings. Hutchinson al. (2020) does so via operations on the Lie algebra, and Fuchs et al. (2020) via irreducible representations. In addition, our work prioritizes applications on visual data and extensively analyses theoretical aspects and properties of group equivariant oeff-attention.
STAND-ALONE SELF-ATTENTION
In this section, we recall the mathematical formulation of self-attention and emphasize the role of the positional encoding. Next, we introduce a functional formulation to self-attention which will allow us to analyze and generalize its equivariance properties.
Definition. Let X IR NxC be input matrix consisting of N tokens of Cin dimensions each.' A self-attention layer maps an input matrix X E NxC. to an output matrix e Nx as:
Y SA(X) softmax
(A) val.
(1)
with W val E IR Chah the value matrix, A E R N*N the attention scores matrix, and softmax (.: the attention probubilities. The matrix A is computed as:
(A)
A XW (XWkey)
parameterized by query and key matrices Wgry. Wkey RCxxC, In practice, it has been found beneficial to apply multiple self-attention operations, also called heads, in parallel, such that different
We consequently consider image as of N discrete objects i (1,2...,,N) N).
2

---

Published as a conference paper at ICLR 2021
heads are able to attend to different parts of the input. In this multi-head self- atff-atttntion formulation, the output H heads of output dimension C₁, are concatenated and projected to Cout as:
MHSA(X) = cconcat [SA (1) (X)] oot + baur, he["]]
(3)
with projection matrix W TR. HCC and bias term bout R TR
3.1 THE ROLE OF THE POSITIONAL ENCODING
Note that the self-attention operation defined in Eq. 3 is equivariant to permutations of the input rows of X. That is, permutation of the rows of X will produce the same output Y up this permutation. Hence, self-attention blind to the order its inputs, i.e., it is a set operation. Illustratively, an input image is processed bag of pixels and the structural content is not considered. To alleviate this limitation, the input representations in self-attention are often enriched with a positional encoding that provides positional information of the elements.
bsolute positional encoding. Vaswani et al. (2017) introduced a (learnable) positional encoding P E NxC for each input position which is added the inputs when computing the attention scores:
A : (X P) (o PX Weey)".
(4)
More generally, P can be substituted by any function that returns a vector representation of the position and can be incorporated by means of addition or concatenation, e.g.. Zhao et al. (2020). This positional encoding injects additional structural information about the tokens into the model, which makes susceptible to changes in the token's ' positions. Unfortunately, the model must learn to recognize similar patterns at every position independently as absolute positional encodings are unique to each position. This undesired data inefficiency addressed by relative positional encodings.
Relative positional encoding. Introduced by Shaw al. (2018), relative encodings consider the relative distance between the query token the token we compute the representation of -, and the key token the token we attend The calculation of the attention scores (Eq. 2) then becomes:
rel X, Woy ( + Px()-z(1) W
(5)
where P Pa)))a(t) E R11CC. is vector representation of the relative shift and x<() is the position of the token as defined in $3.2. Consequently, similar patterns can be recognized at arbitrary positions, relative query-key distances always remain equal.
3.2 A FUNCTIONAL FORMULATION TO SELF-ATTENTION
Notation. We denote by [n] the set (1,2....,n) Given set S and a vector space V. (8) will denote the space of functions { 8 - V}. Square brackets are used when functions are arguments.
Let ( be set of N elements. A matrix X NxC, can be interpreted as a vector-valued function S C that maps element sets е 8 to Cin-dimensional vectors: / i f().. Consequently, matrix multiplication, XWT. of matrices X NxC, and W, I2 CC can be represented a function Gy Lve,, (S) LVc.,, (S), Py f(i) (f))),, parameterized by Wy, wy between functional spaces LVc, (8) = (f S + C and Lvc.. (8) = {f: 8 + RC Following this notation, we can represent the position-less attention scores calculation (Eq. 2) as:
Aij aa]=((((( = (Pay (f)), :Pky (f()))
The function a[f] : x 8 + IR maps pairs of set elements j E 8 to the attention score of relative to í. Therefore, the self-attention (Eq. 1) can be written as:
Y., <[f]())= [o,(alf((t/))
Σ (((()) (44)))())))))) JES
(7)
jeS
where o, softmax, and <[f] RG. Finally, multi-head self-attention (Eq. 3) can be written as:
MHSA(X),, m() = Pour (UneMn (())))())
= Pout Une(A) Σ (10 (f()), Bh (F(N)) (f()))
(8)
jeB
where U is the functional equivalent of the concatenation operator concat, and m[ S C...
Local self-attention. Recall that a[f] assigns an attention scores to every other set element j 8 relative to the query element é. The computational cost of self-attention is often reduced by restricting
3

---

Published as a conference paper at ICLR 2021
its calculation to local neighborhood ()()) around the query token analogous in nature the local receptive field of CNNs (Fig. A. Aa).). Consequently, local self-attention can be written as:
m [[]
Pout U (1 (S(()), Pkey (IUM)) (SG))-
Note that Eq. 9 is equivalent to Eq. 8 for 711)= 8, i.e. when considering global neighborhoods.
Absolute positional encoding. The absolute positional encoding is a function p: S - RC that maps set elements e 8 a vector representation of its position: p + p(t). Note that this encoding is not dependent functions defined on the set but only the set itself." Consequently, absolute position-awae aware self- atteettenni (Eq. 4) can be written as:
iPout
U jen())
(N) (1, (1) (f(i) p(i)), (key (f() + (()) (f()))
(10)
The function p can be decomposed as two functions of x: (i) the position function x X. which provides the position of set elements in the underlying space I' (e.g., pixel positions), and, (ii) the positional encoding p² X CC, which provides vector representations of elements in X(. This distinction will be of utmost importance when we pinpoint where exactly (group) equivariance must be imposed to the self- eeffaattetiion operation ($4.3, $5).
Relative positional encoding. Here, positional information is provided in relative manner. That is, we now provide vector representations of relative positions p(i. = p (=((j - xi)) among pairs (i,j), is 8 j (i). Consequently, relative position-awan self-attention (Bq. 5) can be written as: [f p1)) Fou
U jen(s)
[ (F(()))()))))))))))))) 4 (I()) + +((0)+(((,())))))) ((c.)))) (1) (f()))
(11)
4 EQUIVARIANCE ANALYSIS OF ATTENTIN
In this section we analyze the equivariance properties of self-attention. Since the analysis largely relies on group theory, we provide all concepts required for proper understanding in Appx. C.
4.1 GROUP EQUIVARIANCE AND EQUIVARIANCE FOR FUNCTIONS DEFINED ON SETS
First we provide the general definition of group equivariance and refine it to relevant groups next. Additionally, we define the property of unique equivariance to restrict equivariance to given group. Definition 4.1 (Group equivariance). Let G be a group (Def. C.I), S,S' be sets, V, V be vector spaces, and L. be the induced (left regular) representation (Def. C.4) o on LV(S) and (8'),, respectively. We say that map y Ly(8) + Ly (88) is equivariant to the action of G - or G G-quivvarann -, if commutes with the action of G. That is, if:
+[fg[/]]] L[[v[[]], Vf Lv(), Vge
( be a set of N elements, and G = =NN be the group of permutations on sets of N elements. A map ip: Le(8) + Lo (8) is said to be equivariant to the action of Sy - or permutation equivariant if:
Example 4.1.1 (Permutation equivariance). Lei = 8' 8'
= =[==)
e V E SN Vi
where - I(m and n: + 8 a bijection from the set itself. The element Ti) indicates the index to which the i-th element of the set is moved to as an effect of the permutation т. other words, &p said to be permutation equivariant if commutes with permutations π Sy.
That is, if permutations in its argument produce equivalent permutations on its response.
Several of the transformations of interest, e.g., rotations, translations, are not defined on sets. Luckily, as we consider sets gathered from homogeneous spaces J where these transformations are well- defined, e.g., R for pixels, there exists an injective map x: : 8 .CC that associates a position in X to each set element, the position function. In Appx. D we show that the action of G on such set is well-defined and induces a group representation to functions on it. With this in place, we are now able to define equivariance of set functions to groups whose actions are defined on homogeneous spaces. Definition 4.2 (Equivariance of set functions to groups acting on homogeneous spaces). Let G be a group acting on two homogeneous spaces X. and X.', let 8, 8' be sets and V. v be vector
Illustrativlll, one can think of this as function returning vector representation of pixel positions in grid. Regardless of any transformation performed the image, the laheling of the grid itself remains exactly equal.
4

---

Published as a conference paper at ICLR 2021
spaces, Let X and x' ; 8' + equivariant to the action of - or -equivarrann -, if commutes with the action of G. That is. if. [f [[v[[]], Lo(), Vg G.
- X' be injective maps. We say that map 9 (x) (5))
where f[(f(())== = f(= (g)()))))))))))))))))))))))))))) (()))
are the induced (left regular) representation of G on L)( and Lor(8') respectively. I.o.w., ip is said be G-equivariant if a transformation, E & on its argument produces a corresponding transformation on its response.
Example 4.2.1 (Translation equivariance). Let S, 8' be sets and let x :: + and x' S'' + X" be injective maps from the sets S, 8,8' s to the corresponding homogeneous spaces XX' on which they are defined, e.g., Ra and G. With (X,+ ()) the translation group acting on I, we say that a map ip (S( + Ly (8) is equivariant to the action of (X,+)) -or translation equivariant -, if: <[<g[f](())= [[lpf]](() e
Vye X,
Low.,, ip is said to be translation equivariant if translation on its argument produces a corresponding translation on its response.
with
(x)
4.2 EQUIVARIANCE PROPERTIES OF BELE-ATTENTION
In this section we analyze the equivariance properties of the self-attention The proofs to all the propositions stated in the main text are provided Appx. G.
Proposition 4.1. The global self- attentton formulation without positional encoding (Eqs. 3, 8) is permulation equivariant. That is, it holds that: [[[f([]](( [.((m(]](())....
Note that permutation equivariance only holds for global self-attention. The local variant proposed in Eq. reduces permutation equivariance conserved under permutation, i.e.. Sn { E SN j (1(1) - xi) ((),, Vi V).
smaller set of permutations where neighborhoods are
Permutation equivariance induces equivariance to important (sub) (sub)groups. Consider the cyclic group of order 4, Za {e,r,, r²: r which induces planar rotations by 90" As every rotation in Zz effectively induces permutation of the tokens positions, can be shown that I, is subgroup of Sy i.e., Sx 2 Z4. Consequently, maps equivariant to permutations are automatically equivariant to Z4- However, as the permutation equivariance constraint harder than that of Z,-equivariance, imposing Zı-equivariance as a result of permutation equivariance is undesirable in terms of expressivity. Consequently, Ravanbakhsh et al. (2017) introduced the concept of unique Geequivariance to express the family of functions equivariant to & but not equivariant to other groups G' G:
Definition 4.3 (Unique G-equivariance), Let G a subgroup of G', G G (Def. C.2). We say that map yp is uniquely G Gequivvariant iff it is G-equivariant but not G -qquvaarian for any G' 2 G.
In the following sections, SN. e.g., Z4, but also to other interesting groups not contained in SN e.g., groups of rotations finer than 90 degrees. This is achieved by enriching set functions with proper positional encoding.
show that we enforce unique equivariance not only to subgroups of
Proposition 4.2. Absolute position-aware self-altention (Eqs. 4, 10) is neither permutation nor trans- lation equivariant. i..,, nz m=[d''[[)
# L' Lmm[fp]ll and mz m[[',
# [mm[f.]](())
Though absolute positional encodings do disrupt permutations equivariance, they are unable to provide translation equivariance. We show next that translation equivariance is obtained via relative encodings.
Proposition 4.3. Relative position-aware self-attention (Eq. 11) is translation equivariant. That is, it holds that: m² [f)[ L [f.))((i.
4.3 WHERE EXACTLY IS EQUIVARIANCE IMPOSED IN SELF-ATTENTION?
In the previous section we have seen two examples of successfully imposing group equivariance to self- aeff-attention Specifically, we see that no positional encoding allows for permutation equivariance and that a relative positional encoding allows for translation cquivariance, For the latter, as shown in the proof of Prop. 4.3 (Appx. G), this comes from the fact that for all shifts y 6 X,,
(<x-1<(()) v) (=(j) + v)) p (x()+++)+++)+++)+) y (=(1) + v)) = (x(s) x()) f). (12)
That is, from the fact that the relative positional encoding is invariant to the action of the translation group, i.c., (i, = f), Vy e X. Similarly, the absence of positional encoding - more precisely, the use of a constant positional encoding -, is what allows for permutation equivariance
" represents 0 rotation, i.c., the identity. The remaining elements r represent rotations by (90-))
5

---

Published as a conference paper at ICLR 2021
(Prop. 4.1, Appx. G). Specifically, constant positional encodings pc) V & are invariant to the action of the permutation group, i.e., (i)) pc(1), V EN..
From these observations, we conclude that Geequivariance is obtained by providing positional en- codings which are invariant to the action of the group G. i.c., s.t., [ p. Vg G. Furthermore, unique G-equivariance is obtained by providing positional encodings which are invariant to the action of & but not invariant to the action of any other group G' 2 G. This is a key insight that allows us to provide (unique) equivariance arbitrary symmetry groups, which we provide next.
GROUP EQUIVARIANT STAND-ALONE SELF-ATTENTE ATTENTION
In $4.3 we concluded that unique G Geequivariance is induced in self-attention by introducing positional encodings which are invariant to the action of & but not invariant to the action of other groups G' 2 G. However, this constraint does not provide any information about the expressivity of the mapping we have just made G-equivariant. Let us first illustrate why this is important:
Consider the case of imposing rotation and translation equívariance to an encoding defined in R². Since translation equivariance is desired, relative positional encoding is required. For rotation equivariance, we must further impose the positional encoding to be equal for all rotations. That is
() -( j(i,)), v0 [0,2 where Lo(o] == of (0)=x(( (-)x( and 0-1 depicts rotation by - degrees. This constraint leads to an isotropic positional encoding unable to discriminate among orientations, which in turn enforces rotation invariance instead of rotation equivariance. This is alleviated by lifting the underlying function on R to space where rotations are explicitly encoded (Fig. 1. To this end, one performs self-attentio operations for positional encodings [o[p] of varying values O and indexes their responses by the corresponding / value. Next, as rotations are now explicitly encoded, a positional encoding can be defined in this space which is able to discriminate among rotations (Fig. B.2). This in turn allows for rotation equivariance instead of rotation invariance.
has been shown both theoretically (Ravanbakhsh, 2020) and empirically (Weiler & Cesa, 2019) that the most expressive class of equiivaarannt functions is given by functions that follow the regular representation of G. In order to obtain feature representations that behave that way, we introduce lifting self-attention layer (Fig. 1. Eq. 14) that receives input function on T.d and produces a feature representation on G. Subsequently. arbitrarily many group self-attention layers (Fig. B.2, Eq. 16) interleaved with optional point-wise non-linearities can be applied. At the end of the network a feature representation on 8d can be provided by pooling over H. In short, we provide a pure self- attention analogous to Cohen & Welling (2016). However, as the group acts directly on the positional encoding. our networks are steerable as well (Weiler et al., 2018b). This allows us to go beyond group discretizations that live the grid without introducing interpolation artifacts ($5.1).
Though theoretically sound, neural architectures using regular representations are unable to handle continuous groups directly in practice. This is a result of the summation over elements á e H in Eq. 15, which becomes an integral for continuous groups. Interestingly, using discrete groups does not seem to be detrimental in practice. Our experiments indicate that performance saturates for fine discrete ap- proximations of the underlying continuous group (Tab. 2). fact, (Weiler & Cesa, 2019, Tab. 3) show via extensive experiments that networks using regular representations and fine enough discrete approx- imations consistently outperform networks handling continuous groups via irreducible representations. We conjecture this is result of the networks receiving discrete signals as input. As the action of several group elements fall within the same pixel, no further improvement can be obtained.
5.1 GROUP SELF-ATTENTION IS AN STEERABLE OPERATION
Convolutional filters are commonly parameterized by weights on a discrete grid, which approximate the function implicitly described by the filter at the grid locations. Unfortunately, for groups whose action does not live in this grid, e.g., 45 rotations, the filter must be interpolated. This is problematic as these filters are typically small and the resulting interpolation artifacts can be severe (Fig. 2a). Steerable CNNs tackle this problem by parameterizing convolutional filters on a continuous basis on which the action of the group is well-defined, e.g., circular harmonics (Weiler et al., 2018b), B-splines
"This phenomenon arises from the fact that ? is a quotient of the roto ato-ttrasslation group. Consequently, imposing group equivariance in the quotient space is equivalent to imposing an additional homomorphism of constant value over its cosets. Conclusively, the resulting map is of constant value over the rotation elements and, thus, not able to discriminate among them. See Ch. 3.1 of Dummit Foote (2004) for an intuitive deseription.

---

Published as a conference paper at ICLR 2021
Conv. Filter
Content
Pos. Encending
f
Iva.
Group Acsion 45"
Group Action
45'
(a) Discrete convolution
(b) Group self-attention
Figure 2: Steerability analysis discrete convolutions and group self-attention.
(Bekkers, 2020). In group self-attention, the action of the group leaves the content of the image intact and only modifies the positional encoding (Figs. B.1, B.2). As the positional encoding lives con- tinuous space, it can be transformed at an arbitrary grade of precision without interpolation (Fig. 2b).
5.2 LIFTING AND GROUP SELF-ATTENTION
Lifting self-attention (Fig. B.1). Let G R2 x H be an affine group (Def. C.3) acting on Rd The Lu (RR) L (G) is map from functions on Rd functions on G obtained by modifying the relative positional encoding p(i.j) by the action of group elements A Ht: {LL[p]( ()) ( = (f -
lifting self-attention [ (tf..p.. p]
f x()). It corresponds to the concatenation of multiple self self-attention operations (Eq. 11) indexed by f with varying positional encodings <G[p]]
mi m² [f[f
(13)
PPou (U
(I()) (1) (S()) Llll)( (N) Me[H)]jnn(
(1())) (14)
Proposition 5.1. Lifting self-attention is G-equivariant. That is, holds that: m/y [ (f ((,
Group self-attention (Fig, B.2). Let G RRd H be an affine group acting on itself and f(i,h) Ly(G), i 8, h H. be a function defined on a set immersed with the structure of the group G. That is, enriched with a positional encoding p((t,)(, (i. ) = (p:===ppaa>aa ((=()) x(i), fr) i, e 8, R e He. The group self-attention mi [G[[p) Lv( Le (G (G) is map from functions on G to functions on G obtained by modifying the group positional encoding by the action of group elements A e He: { [p]((i, (i, (i) heeat, Li ((), (i,A)) - (x(() ( (n-11)) corresponds the concatenation of multiple self-attention operations (Eq. 11) indexed by h with varying positional encodings La [e] and followed by summation over the output domain along h: m((,) Eñer m² [[ [A (I])((t))))))))
(15)
= =Poout U JUE Σ E0 he[H] Rede (1A))n(()))
( (y ()(
+ +IIIPl((i,)) (.
(16)
In contrast to vanilla and lifting self-attention, the group self-attention neighborhood n A) is now defined the group. This allows distinguishing across group transformations, e.g. rotations.
Proposition 5.2. Group self-attention is G-equivariant. That is, it holds that: (((-A))) [[lm[[f.p](C..]
Non-unimodular groups, i.e., groups that modify the volume of the objects they act upon, such the dilation group, require a special treatment. This treatment is provided in Appx. E.
5.3 GROUP ATTENITNO
GENERALIZATION OF THE GROUP CONVOLUTION
We have demonstrated that it is sufficient to define self-attention as function on the group G and en- sure that [g(p)
Vg e G order to enforce G E-quiivariance. Interestingly, this observation inline with the main statement of Kondor & Trivedi (2018) for (group) convolutions: "the group convolution on G is the only (unique) G Geequivariant linear map"". In fact, our finding can formulated as generalization of Kondor & Trivedi (2018)'s statement
"Linear mappings on G whose positional encoding is G-invariant are G-equivariant"
This statement is more general than that Kondor Trivedi (2018), as it holds for data structures where (group) convolutions not well-defined, e.g., sets, and it is equivalent to Kondor & Trivedi
7

---

Published as a conference paper at ICLR 2021
(2018)'s statement for structures where (group) convolutions well- is also congruent with results complementary to Kondor& & Trivedi (2018) (Cohen et al., 2019a; Bekkers, 2020) as well as sev- eral works on group equivariance handling set-like structures like point-clouds (Thomas et al., 2018; Defferrard et al., 2020; Finzi et al., 2020; Fuchs et al., 2020) and symmetric sets (Maron et al., 2020).
In addition, we can characterize the expressivity of group self-attention. It holds that (7) group self-attention generalizes the group convolution and (ii) regular global group self-attention is an equivariant universal approximator. Statement (i) follows from the fact that any convolutional layer can be described as multi-head self-attention layer provided enough heads (Cordonnier et al., 2020), yet self-attention often uses larger receptive fields. As result, self-attention is able to describe a larger set of functions than convolutions, e.g., Fig. A.1. Cordonnier et al. (2020)'s statement can be seamlessly extended to group self-attention by incorporating an additional dimension corresponding to H. in their derivations, and defining neighborhoods in this new space with proportionally larger number of heads. Statement (ii) stems from the finding of Ravanbakhsh (2020) that functions induced equivariant universal approximators provided fill kernels, i.e.,
by regular group representations global receptive fields. Global receptive fields are required to guarantee that the equivariant map is able to model any dependency among input components. Global receptive fields are readily utilized by our proposed regular global group self-attention and, provided enough heads, one can ensure that any such dependencies is properly modelled.
6 EXPERIMENTS
We perform experiments on three image benchmark datasets for which particular forms of equiv- ariance are desirable." We evaluate our approach by contrasting GSA-Nets equivariant to multiple symmetry groups. Additionally, we conduct an study on rotMNIST to evaluate the performance of GSA-Nets as a function of the neighborhood size. All our networks follow the structure shown in Fig. F.1 and vary only in the number of blocks and channels. We emphasize that both the architecture and the number of parameters in GSA-Nets is left unchanged regardless of the group used. Our results illustrate that GSA-Nets consistently outperform equivalent non-equivariant attention networks. We further compare GSA-Nets with convolutional architectures. Though our approach does not build upon these networks, this comparison provides fair view to the yet present gap between self-attention and convolutional architectures in vision tasks, also present in their group equivariant counterparts.
Efficient implementation of lifting and group self-attention. Our self-attention implementation takes advantage of the fact that the group action only affects the positional encoding P to reduce the total computational cost of the operation. Specifically, we calculate self-attention scores w.r.t. the content X once, and reuse them for all transformed versions of the positional encoding { [
Model designation. We refer to translation equivariant self-attention models as Z2 SA. Reflection equivariant models receive the keyword M, e.g., Z2M SA, and rotation equivariant models the keyword Rn, where depicts the angle discretization. For example, R8 R8_SA depicts a model equivariant to rotations by 45 degrees. Specific model architectures are provided in Appx. F.
The rotated MNIST dataset (Larochelle et al., 2007) is a classification dataset often used as a standard benchmark for rotation equivariance. It consists of 62k gray-scale 28x28 uni- formly rotated handwritten digits, divided into training, validation and test sets of 10k, 2k and 50k images. First, we study the effect of the neighborhood size on classification accuracy and convergence time. We train R4 SA networks for 300 epochs with vicinities NxN of varying size (Tab. 1, Fig. 3). Since GSA-Nets optimize where to attend, the complexity of the optimization problem grows function of N. Consequently, models with big vicinities are expected to converge slower. However, as the family of functions describable by big vicinities contains those describable by small ones, models with big vicinities are expected to be at least as good upon convergence. Our results show that models with small vicinities do converge much faster (Fig. 3). However, though some models with large vicini- ties do outperform models with small ones, e.g., 7x7 vs. 3x3, a trend of this behavior is not apparent. We conjecture that 300 epochs are insufficient for all models to converge equally well. Unfortunately, due to computational constraints, we were not able to perform this experiment for larger number of epochs. We consider an in-depth study of this behavior an important direction for future work.
RotMNIST.
Next, we compare GSA-Neta equivariant to translation and rotation at different angle discretizations (Tab. 2). Based on the results of the previous study, we select a 5x5 neighborhood, as it provides the
"Our code is publicly available ht thub. com wttp://m//////oot wwwm///womree////oooooooooooooce ro/ sel fatt..
8

---

Published as a conference paper at ICLR 2021
Table 1: Accuracy vs. neighborhood size.
1.0
ROTMNIST
R4,5A,5K5
MODEL
NORD. sizB
Acc. (%)
TRAIN, TIME
0.8
84 SA 9x9
90,9
EEPCC
R4 1_SA_2x7
90.0
R4,5A_11k11
3x3
96.56
04:53 IGPU
84 SA. 15x15
89.1
0.6
R4, 14_5A_19x19
86.7
7x7
97.33
09:03 IGPU
R4_5A_IK3
9x9
97.42
09:16 IGPU
0.5
R4 SA, 23x23
83.0
84
11x11
97.17
12:09 IGPU
SA. 28x28
81.3
15x15
96.89
10:27 2GPU
19x19
96.86
14:27 2GPU
10
15
20
25
23x23
97.05
06:13 3GPU
spach
Figure 3: Test accuracy in early training stage.
Table 2: Classification results. All convolutional architectures use 3x3 filters.
BOTMNIST
CIFARIO
PATCHCAMELYON
MODEL
Acc. (8)
PARAMS.
MODEL
ACC. (%)
PARAMS.
MODEL
Acc. (X)
PARAMS.
22_SA
96.37
Z2_SA
82.3
Z2_SA
83.04
R4. SA
97.46
Z2M SA
83.72
2.99M
R4. SA
83.44
205.66
B8 SA
97.90
44.67%
R8 SA
83.58
R12.8A
97.97
90.56
1.37M
R4M_SA
84.76
R16_SA
97.66
Cohen Welling (2016).
Z2_CNN $4.07
60k
22_C
94.97
21.75k
R4_CNN
87.55
129.65K
R4 CUN
77.54x
R4M CNN
88.36
124.21K
R-R4. ONN
98.31
73.13K
Ou R4_CNN
88.
140.45K
Cohen Welling (2016)
89.12
141,22K
"Romero et al. (2020ks
Romero et (2020s).
best trade-off between accuracy and convergence time. Our results show that finer discretizations lead to better accuracy but saturates around R12. We conjecture that this is due to the discrete resolution of the images in the dataset, which leads finer angle discretizations to fall within the same pixel.
CIFAR-10. The CIFAR-10 dataset (Krizhevsky et al., 2009) consists of 60k: real-world 32x32 RGB images uniformly drawn from 10 classes, divided into training, validation and test sets of 40k, 10k: and 10k images. Since reflection is symmetry that appears ubiquitously in natural images, compare GSA-Nets equivariant to translation and reflection in this dataset (Tab. 2). Our results show that reflection equivariance indeed improves the classification performance of the model.
PCam. The PatchCamelyon dataset (Veeling et al., 2018) consists of 327k 96x96 RGB image patches of cettorous/non-tumorroo breast tissues extracted from Camelyon16 (Bejnordi et al., 2017). Each patch is labeled as tumorous if the central region (32x32) contains at least one tumour pixel. As cells appear arbitrary positions and poses, we compare GSA-Nets equivariant translation, rotation and reflection (Tab. 2). Our results show that incorporating equivariance to reflection in addition to rotation, well as providing finer group discretization, improve classification performance.
DISCUSSION AND FUTURE WORK
Though GSA-Nets perform competitively to G-CNNs for some tasks, G-CNNs still outperforms our approach in general. We conjecture that this is due to the harder nature of the optimization problem in GSA-Nets and the carefully crafted architecture design, initialization and optimization procedures developed for CNNs over the years. Though more expressive than G-CNNs ($ 5.3), further research terms of design, optimization, stability and generalization is required. These are in fact open questions for self-attention in general (Xiong et al., 2020; Liu et al., 2020; Zhao et al., 2020) and developments in this direction are of utmost importance.
theoretical results indicate that GSA-Nets can be
The main drawback of our approach is the quadratic memory and time complexity typical of self- attention. This is an active area of research, e.g., Kitaev et al. (2020); Wang et al. (2020); Zaheer et al. (2020); Choromanski al. (2020) and we believe that efficiency advances vanilla self attention can be seamlessly integrated in GSA-Nets. Our theoretical results indicate that GSA-Nets have the potential to become the standard solution for applications exhibiting symmetries, e.g., medical imagery. In addition, as self-attention is set operation, GSA-Nets provide straightforward solutions to set-like data types, e.g., point-clouds, graphs, symmetric sets, which may benefit from additional geometrical information, e.g., Fuchs et al. (2020); Maron et al. (2020). Finally, we hope our theoretical insights serve as a support point to further explore and understand the construction of equivariant maps for graphs and sets, which often come equipped with spatial coordinates: type of positional encoding.
9

---

Published as a conference paper at ICLR 2021
ACKNOWLEDGMENTS
We gratefully acknowledge Michael Hutchinson, Charline Le Lan, Sheheryar Zaidi, Emilien Dupont, and Hyunjik Kim for useful discussions, Robert-Jan Bruintjes, Fabian Fuchs, Erik Bekkers, Andreas Loukas, Mark Hoogendoorn and our anonymous reviewers for their valuable comments on early versions of this work which largely helped us to improve the quality of our work.
David W. Romero is financed as part of the Efficient Deep Learning (EDL) programme (grant number P16-25), partly funded by the Dutch Research Council (NWO) and Semiotic Labs. Jean- Baptiste Cordonnier is financed by the Swiss Data Science Center (SDSC), Both authors thankfully acknowledge everyone involved in funding this work. This work was carried out on the Dutch national e-infrastructure with the support of SURF Cooperative.
REFERENCES
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXiv preprint arXiv:1607.06490 2016.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL http: / wwxiv.orog........... abs 1409.049440 0473.
Babak Ehteshami Bejnordi, Mitko Veta, Paul Johannes Van Diest, Bram Van Ginneken, Nico Karssemeijer, Geert Litjens, Jeroen AWM Van Der Laak, Meyke Hermsen, Quirine F Manson, Maschenka Balkenhol, al. Diagnostic assessment of deep learning algorithms for detection of lymph node metastases in women with breast cancer. Jama, 318(22)2219922)) 2017.
Erik Bekkers. B-spline [cnn]s on lie groups In International Conference on Learning Representa- tions, 2020. URL /tttttps:///pereiviieii www.ee..eeeeeeeeeeeeeeeeeeeeeeee / Eo uorum? 2 1d=:1 gBhk
Erik Bekkers, Maxime W Lafarge, Mitko Veta, Koen AJ Eppenhof, Josien PW Pluim, and Remco Duits. Roto-translation covariant convolutional networks for medical image analysis. In Interna- tional Conference Medical Image Computing and Computer-Assister Intervention, pp. 440-448. Springer, 2018
Irwan Bello, Barret Zoph, Ashish Vaswani, Jonathon Shlens, and Quoe V. Le. Attention augmented convolutional networks, 2019.
Jianpeng Cheng, Li Dong, and Mirella Lapata. Long short-term metory-networks for machine reading. arXiv preprint arXiv:1601.0673: 2016.
Xiuyuan Cheng, Qiang Qiu, Robert Calderbank, and Guillermo Sapiro. Rotdef: Decomposition of convolutional filters for rotation-equivariant deep networks. arXiv preprint arXiv: hiti::/805.008.. 2018.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreca Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv: 14794, 2020.
Taco Cohen and Max Welling. Group equivariant convolutional networks. In International confere machine learning, pp. 2990-2999, 2016.
Taco Cohen, Mario Geiger, Jonas Köhler, and Max Welling. Spherical cnns. CoRR, 1010000..00 1801.1111000 10130.
2018. URL http: htttt::://xxiv.vrrgg/ee
Taco S Cohen, Mario Geiger, and Maurice Weiler. A general theory of equivariant cnns ho- mogeneous spaces. In Advances in Neural Information Processing Systems, pp. 9142-9153, 2019a.
Taco S Cohen, Maurice Weiler, Berkay Kicanaoglu, and Max Welling. Gauge equivariant convolu- tional networks and the icosahedral enn. arXiv preprint arXiv: arXiv: 1902.04615. 2019b.
10

---

Published as a conference paper at ICLR 2021
Jean- Cean-Baptiste Cordonnier, Andreas Loukas, and Martin Jaggi. On the relationship between self- attention and convolutional layers. In International Conference on Learning Representations, 2020.
URL https: http:///eeeeeeeeeeeeeeeeeeeeeeee net httppp/rwwww..ntttffooooo/iioooo Eorum?21c-HUnnC11333
Michaël Defferrard, Martino Milani, Frédérick Gusset, and Nathanaël Perraudin. Deepsphere: graph- based spherical cnn. In International Conference on Learning Representations, 2020. URL
https:/ /ttp///perree//e ew net forum? e.01sttPP
Nichita Diaconu and Daniel E Worrall. Learning to convolve: A generalized weight-tying approach. 2019.
Sander Dieleman, Jeffrey De Fauw, and Koray Kavukcuoglu. Exploiting cyclic symmetry in convolu- tional neural networks. arXiv preprint 1xXi::660..22660, 2016.
David Steven Dummit and Richard M Foote. Abstract algebra, volume 3. Wiley Hoboken, 2004.
Carlos Esteves, Avneesh Sud, Zhengyi Luo, Kostas Danillidis, and Ameesh Makadia. Cross-domain 3d equivariant image embeddings. In International Conference on Machine Learning, pp. 1812- 1822. PMLR, 2019a.
Carlos Esteves, Yinshuang Xu, Christine Allen-Blanchette, and Kostas Daniilidis. Equivariant multi- view networks. In Proceedings of the IEEE International Conference Computer Vision, pp. 1568-1577, 2019b.
Carlos Esteves, Ameesh Makadia, and Kostas Danilidis. Spin-weighted spherical cnns. Advances in Neural Information Processing Systems, 33, 2020.
Mare Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolutional neural networks for equivariance lie groups arbitrary continuous data. arXiv preprint arXiv:2002.128800 2020.
Fabian B Fuchs, Daniel E Worrall, Volker Fischer, and Max Welling, Se (3)-transformers: 3d antotrrnnslatonn equivariant attention networks. arXiv preprint arXiv:2006. 10503, 2020.
Simon Graham, David Epstein, and Nasir Rajpoot. Dense steerable filter cnns for exploiting rotational symmetry in histology images. arXiv preprint arXiv:2004.03003 2020.
Emiel Hoogeboom, Jorn WT Peters, Taco Cohen, and Max Welling, Hexaconv. arXiv preprint 1rXiv:1803.021000 2018.
Jie Hu, Li Shen, and Gang Sun. Squeeze-and-exciaation networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 7132-7141, 2018.
Michael Hutchinson, Charline Le Lan, Sheheryar Zaidi, Emilien Dupont, Yee Whye Teh, and Hyunjik Kim. Lietransformer: Equivariant self-attention for lie groups arXiv preprint arXiv:2012. 10885, 2020.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns:: Fast autoregressive transformers with linear attention, 2020.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer, 2020.
Risi Kondor and Shubhendu Trivedi. On the generalization of equivariance and convolution in neural networks the action of compact groups. arXiv preprint
2018.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. Technical report, Citeseer, 2009.
Hugo Larochelle, Dumitru Erhan. Aaron Courville. James Bergstra, and Yoshua Bengio. An empirical
evaluation of deep architectures on problems with many factors of variation. In Proceedings of the 24th international conference on Machine learning, pp. 473-480. ACM, 2007.
Jan Eric Lenssen, Matthias Fey, and Pascal Libuschewski. Group equivariant capsule networks, In Advances in Neural Information Processing Systems, pp. 8844-8853, 2018.
11

---

Published as a conference paper at ICLR 2021
Junying Li, Zichen Yang, Haifeng Liu, and Deng Cai. Deep rotation equivariant network. Neurocom- puting, 290:26-33, 2018.
Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the difficulty of training transformers. arXiv preprint arXiv:2004.08224 2020.
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:I608.03940 2016.
Diego Marcos, Michele Volpi, Nikos Komodakis, and Devis Tuia. Rotation equivariant vector field networks. In Proceedings of the IEEE International Conference Computer Vision, pp. 5048-5057, 2017.
Diego Marcos, Benjamin Kellenberger, Sylvain Lobry, and Devis Tuia. Scale equivariance in cnns with vector fields. arXiv preprint arXiv:/807.11788 2018.
Haggai Maron, Or Litany, Gal Chechik, and Ethan Fetaya. On learning sets of symmetric elements. arXiv preprint (rXXi:20022.8599,, 2020.
Prajit Ramachandran, Barret Zoph, and Quoe V Le. Searching for activation functions. arXiv preprint arXiv:1710.05941, 2017.
Prajit Ramachandran, Niki Parmar, Ashish Vaswani, Irwan Bello, Anselm Levskaya, and Jonathon Shlens. Stand-alone self-attention in vision models. arXiv preprint arXiv: 19066059099 2019.
Siamak Ravanbakhsh. Universal equivariant multilayer perceptrons. arXiv preprint arXiv:2002.02912, 2020.
Siamak Ravanbakhsh, Jeff Schneider, and Barnabas Poczos. Equivariance through parameter-sharing. arXiv preprint arXiv: 1702.08389, 2017.
David W Romero and Mark Hoogendoorn. Co-attentive equivariant neural networks: Focusing
equivariance on transformations co-occurring in data. arXiv preprint arXiv:1911.07880 2019.
David W Romero, Erik J Bekkers, Jakub M Tomezak, and Mark Hoogendoorn. Attentive group equivariant convolutional networks. arXiv preprint arXiv:2002.03830, 2020a.
David W Romero, Erik J Bekkers, Jakub M Tomezak, and Mark Hoogendoorn. Wavelet networks:
Scale equivariant learning from raw waveforms. arXiv preprini arXiv:2006.05259, 2020b.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representations. In Marilyn A. Walker, Heng Ji, and Amanda Stent (eds.), Proceedings of the 2018 Conference of the North American Chupter of the Association for Computational Linguisties: Human Language Technologies, NAACL-HLT; New Orleans, Louisiana, USA, June 1-6, 2018, Volume (Short Papers), pp. 464-468. Association for Computational Linguistics, 2018. doi: 10.18653/v1/ 1174.074. URL https: /doi. owg/ 11.11..866.00 1ww53//1/11118 2074.
Ivan Sosnovik, Michal Szmaja, and Amold Smeulders. Scale-equivariant steerable networks. In International Conference on Learning Representations, 2020. URL htttps://opereeee
Kai Sheng Tai, Peter Bailis, and Gregory Valiant, Equivariant transformer networks. arXiv preprint 19X1v:1901.11399 2019.
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang. Li Li, Kai Kohlhoff, and Patrick Riley, Tensor Field Networks: Rotation-ant counslation-Equivariin Neural Networks for 3D Point Clouds. arXiv preprint /00.::/80002////// 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Bastiaan S Veeling, Jasper Linmans, Jim Winkens. Taco Cohen, and Max Welling, Rotation equiv- ariant cnns for digital pathology. In International Conference on Medical image computing and computer-assisted intervention, pp. 210-218. Springer, 2018
12

---

Published as a conference paper at ICLR 2021
Sai Raam Venkataraman, S. Balasubramanian, and R. Raghunatha Sarma. Building deep equivariant capsule networks. In International Conference on Learning Representations, 2020. URL httpa:
/ ///oppereview... net httpaa//eweww.nt//ooooooooooooa- ho
Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint (rXiv:2006.04768, 2020.
Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaiming He. Non-local neural networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 7794-7803, 2018.
Maurice Weiler and Gabriele Cesa. General (2)-equivariant steerable enns. In Advances in Neural Information Processing Systems, pp. 14334-14345, 2019.
Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma, and Taco S Cohen. 3d steerable cnns: Learning rotationally equivariant features in volumetric data. Advances in Neural Information Processing Systems, pp. 10381-10392, 2018a.
Maurice Weiler, Fred A Hamprecht, and Martin Storath. Learning steerable filters for rotation equivariant cnns. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 849-858, 2018b.
Daniel Worrall and Gabriel Brostow. Cubenet: Equivariance to 3d rotation and translation. In
Proceedings of the European Conference on Computer Vision (ECCV), pp. 567-584, 2018.
Daniel E Worrall and Max Welling. Deep scale-spaces: Equivariance over scale. arXiv preprint
arXiv: 1669, 2019.
Daniel E Worrall, Stephan J Garbin, Daniyar Turmukhambetor, and Gabriel J Brostow. Harmonic networks: Deep translation and rotation equivariance. In Proceedings of the IEEE Conference Computer Vision and Puttern Recognition, pp. 5028-5037, 2017.
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang. and Tie Yan Liu. On layer normalization in the transformer architecture. arXiv preprint arXiv:2002.04745. 2020.
Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang. Li Yang. al. Big bird: Transformers for longer sequences. arXiv preprint arXiv:2007. /4062, 2020.
Hengshuang Zhao, Jiaya Jia, and Vladlen Koltun. Exploring self-attention for image recognition. In The IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2020.
13

---

Published as a conference paper at ICLR 2021
APPENDIX
A CONVOLUTION AND THE AELE-ATTNNIIN: A GRAPHICAL COMPARISON
Figure A.1: Parameter usage in convolutional kernels (Fig. A.la) and self-attention (Fig. 1b). Given budget of 9 parameters, a convolu- tional filter ties these parameters to specific positions. Subsequently, these parameters remain static regardless of (i) the query input position and (ii) the input signal itself. Self-attention, on the other hand, does not tie parameters any specific positions at all. Contrarily, it pares the representations of all tokens falling in its receptive field. As result, provided enough heads, self-aftention more general than con- volutions, as it can represent any convolutional kernel, e.g., Fig. A.la, well several other functions defined on its receptive field.
(a)
(b)
B LIFTING AND GROUP SELF-ATTENTIO A GRAPHICAL DESCRIPTION
mlg F.pl
Content:
L(pl
nu' [s. [ell]
Positionsl Encoding:
Group Action
Multi-Hol Head Sell-Attention
Figure B.1: Lifting self- colf-attention on the roto-translation group for discrete rotations by 90 degrees (also called the Z4 group). The Z₁ group defined as H (e, li2 n²), where á depicts a 90" rotation. The lifting self-attention corresponds to the concatenation of |3(] self-attentio operations between the input f and f-trranffrrmed versions of the positional encoding L[p]. Vy H. As result, the model "sees" the input each of the rotations in the group once. Since Za is cyclic group, i.e., ft = e, functions on this group are often represented as responses on ring (right side of the image). This is self-attention analogous the regular lifting group convolution broadly utilized in group equivariant learning literature, e.g., Cohen & Welling (2016); Romero et al. (2020a).
14

---

Published as a conference paper at ICLR 2021
Contect
Poactiooa Faceding
Grour Artion
Malli Head Sett-tteeriinn
Figure B.2: Lifting self-attention on the roto-translation group for discrete rotations by 90 degrees (also called the Z1 group). The Z, group is defined as H = (ee ft, h², (,²), h², tion. Analogous to lifting self-attention (Fig. B.1), group self-attention corresponds to concatenation of [ = 4 self-attention operations between the input f and fi-transformed versions of the positional encoding [[). ví E H. However, in contrast to lifting self-attention, both and p are now defined on the group G. Consequently, an additional sum over A is required during the operation (c.f., Eq. 16). Since Z, is a cyclic group. i.e., ft =e, e, functions on Z₁ are often represented as responses on ring (right side of the image). This is self-attention analogous to the regular group convolution broadly utilized in group equivariant learning literature, e.g.. Cohen & Welling (2016); Romero et al. (2020a).
where ft depicts a 90º rota-
C CONCEPTS FROM GROUP THEORY
Definition C.1 (Group). A group is an ordered pair (G.-) where G is a set and G G G is a binary operation on G. such that (i) the set is closed under this operation, (ii) the operation is associative, i.e., (g1 92) - gs 91 (g2- ga), 91.92. gs E G. (iii) there exists an identity element e s.t. Vg e G we have e. g-e g, and (iv) each g E G. there exists an inverse g²1 s.t. g.g =e. Definition C.2 (Subgroup). Let (G.-) be a group. A subset H of G is a subgroup of G if H is nonempty and closed under the group operation and inverses (i.e., fú1, fiz H implies that fi e H and hy fu2 e Ht). If H is a subgroup of G we write H S G
Definition C.3 (Semi-direct product and affine groups). In practice, one is mainly interested in the analysis of data defined on Rdd, and, consequently, in groups of the form G R x H. resulting from the semi-diree product (x) between the translation group (IR +) and an arbitrary (Lie) group H that acts on RRd,, e.g., rotation, scaling, mirroring, etc. This family of groups referred to as affine groups and their group product is defined as:
91 92 = (I1, (I2,62) = (x1 + h1 2 21
(17)
with g (II, -2= (zz, h2) E G. I1, D2 E Rd and fi,, h2 E H. The operator @ denotes the action h E He on & E Rd, and it describes how a vector I RR is modified by elements e H. Definition C.4 (Group representation). Let G be a group and L'( ) be space of functions de- fined on some vector space X. The (left) regular group representation of G is linear transformation L G L2 (X) + L (x). (9,f) + Lg[f]= f(g o x)), that shares the group structure via:
Lo [f] Lff
(18)
for any g2.92 G. e (x).()). That is, concatenating two such transformations, parameterized by
15

---

Published as a conference paper at ICLR 2021
91 and 92. is equivalent single transformation parameterized by 91 92 G. If the group G is affine, the group representation Lg can be split as:
-
(19)
with = (ar, f)) & G. I E TRd and E H. Intuitively, the representation of G on a function I describes how the function as whole, ie., fxx), Vx X, is transformed by the effect of group elements g e G.
D ACTIONS AND REPRESENTATIONS OF GROUPS ACTING ON HOMOGENEOUS SPACES FOR FUNCTIONS DEFINED ON SETS
In this section we show that the action of a group G acting on a homogeneous space .t is well defined on sets 8 gathered from X', and that it induces a group representation of functions defined on 8.
Let 8 {i} be set and X' be a homogeneous space on which the action of G is well-defined, i.e. gr _, Vg G. Vz é X Since 8 has been gathered from Xt there exists an injective map r: S _C, that maps set elements i é 8 to unique elements 2; é X. That is, there exists a map a: i - that assigns an unique value .C to each i 8 corresponding to the coordinates from which the set element has been gathered.
Since the action of G is well defined on .X, it follows that the left regular representation (Def. C.4) [(;) = f (a 'x;) E Ly(')) of functions f & (y(')) exists and is well-defined. Since x is injective, the left regular representation L. (x:)) = f (g-1x;) can be expressed uniquely in terms of set indices as Izzf">((;;)) =
f x()). Furthermore, its inverse xx : X + S, x=1 Ij + i also exist and is well-defined. As a consequence, points I; e X can be expressed uniquely in terms of set indices as i x (3t), S. Consequently, functions fx (Xy) can be expressed in terms of functions f Ly( by means of the equality f*)()= (x:). Resultantly, we that the group representation Lg[ [xx) f (g- xx()) can be described in terms of functions f f (S) as:
(x:) f (( f(x
(0'
(())) Li(((()
with a corresponding group representation on (8) given by (() (( (g-t()) and an action of group elements g G on set elements given by gi = r (gr(t))
E THE CASE OF NON-UNIMODULAR GROUPS: SELF-ATTENTION ON THE DILATION-TANNSLANN
GROUP
The lifting and group self-attention formulation provided in $5.2 are only valid for unimodular groups. That is, for groups whose action does not change the volume of the objects they act upon, e.g., rotation, mirroring, etc.
groups, however, do modify the volume of the acted objects (Bekkers, 2020). The most relevant non-unimodular group for this work is the dilation group H (R.o, x). To illustrate why this distinction is important, consider the following example:
Imagine we have a circle on IR² of area 2 we rotate, mirror or translate the circle, its size is kept constant. If we increase its radius by factor A Rso, however, its size would increase by f?? Imagine that we have an application for which we would like to recognize this circle regardless of any of these transformations by means of self-attention. For this purpose, we define a neighborhood N for which the original circle fits perfectly. Since the size of the circle is not modified for any translated, rotated or translated versions of it, we would still be able detect the circle regardless of these transformations. If we scale the circle by a factor of f > 1, however, the circle would fall outside of our neighborhood 71 and hence, we would not be able to recognize it.
A solution to this problem is to scale our neighborhood 71 in proportional way. That is, if the circle is scaled by a factor A E Rso, we scale our neighborhood by the same factor fi: 71 + fill. Resultantly, the circle would fall within the neighborhood for any scale factor fi Es.. Unfortunately, there is problem: self-attention utilizes summations over its neighborhood. Since Σienn i > Σxen i, for A 1, and Σian i Σien i, for f < 1. the result of the summations would still differ for different scales. Specifically, this result would always be bigger for larger versions of the neighborhood. This is problematic, the response produced by the same circle, would still be different for different scales.
16

---

Published as a conference paper at ICLR 2021
In order to handle this problem, one utilizes normalization factor proportional to the change of size of the neighborhood considered. This ensures that the responses are equivalent for any scale A Rss- That is, one normalizes all summations proportionally to the size of the neighborhood. As result, we obtain that En n (A?) (Ai)~1i Eieng (n3) -1, Whi, A2 e Ryo- 6
In the example above we have provided an intuitive description of the (left invariant) Haar measure df((A). As its name indicates, it is a measure defined on the group, which is invariant over all group elements fi H. For several unimodular groups, the Haar measure corresponds the Lebesgue measure as the volume of the objects the group acts upon is kept equal, i.e., du((f) = d/. For non non-unimodular groups. however, the Haar measure requires normalization factor proportional to the change of volume of these objects. Specifically, the Haar measure corresponds to the Lebesgue measure times a normalization factor Ad, where d corresponds to the dimensionality of the space TR the group acts upon (Bekkers, 2020; Romero et al., 2020b), i.e., dp(At) fu dh.
In conclusion, in order to obtain group equivariance to con-unimodular groups, the lifting and group self-attention formulation provided in Eqs. 14, 16 must be modified via normalization factors proportional to the group elements e H. Specifically, they are redefined as:
"(]L) = yoout
U
Σ of (12 (S(1)) Play (S(i) + La P of ((()) (20)
bel jek
mE[S,p](i, = Poul
U hel (1 hezt (J.8))nn((.))
F (1009 (I(1, Pkay (((i.)) Σ
Σ
+ (;)
(21)
The factor d + in Eq. 21 results from the fact that the summation is now performed on the group : an space of dimensionality 1+1. An interesting case emerges when global neighborhoods are considered, i.e., s.t. 17(1)= S, Vi ε S. Since (i) (i) S for any > 1, approximation artifacts are introduced. It is not clear if it is better to introduce normalization factors in these situations or not. An in-depth investigation of this phenomenon is left for future research.
E. CURRENT EMPIRICAL ASPECTS OF SCALE EQUIVARIANT SELF-ATTENTION
Self-attention suffers from quadratic memory and time complexity proportional to the size of the neighborhood considered. This constraint is particularly important for the dilation group, for which these neighborhoods grow as a result of the group action. We envisage two possible solutions to this limitation left out for future research:
The most promising solution is given by incorporating recent advances in efficient self-attention in group self-attention, e.g., Kitaev et al. (2020); Wang et al. (2020); Zaheer et al. (2020); Katharopoulos et al. (2020); Choromanski et al. (2020). By reducing the quadratic complexity of self-attention, the current computational constraints of scale equivariant self-attention can be (strongly) reduced. Importantly, resulting architectures would be comparable to Bekkers (2020); Sosnovik et al. (2020); Romero et al. (2020b) in terms of their functionality and the group discretizations they can manage.
The second option is to draw a self-attention analogous to Worrall & Welling (2019), where scale equivariance is implemented via dilated convolutions. One might consider an analogous dilated convolutions via "sparse" dilations of the self-attention neighborhood. As result, scale equivariance can be implemented while retaining an equal computational cost for all group elements. Importantly however, this strategy is viable for a dyadic set of scales only, i.e., a set of scales given by a set (2) Jms and thus, less general that the scale-equivariant architectures listed before.
F EXPERIMENTAL DETAILS
In this section we provide extended details over our implementation as well the exact archi- tectures and optimization schemes used in our experiments. All our models follow the structure shown in Fig. F.l and vary only in the number of blocks and channels. All self-attention oper- ations utilize 9 heads. We utilize PyTorch for our implementation. Any missing specification can be safely considered be the PyTorch default value. Our code is publicly available at https: : / httpt/ttttttttttt///////tt
thub. com/ wwt.meoo/ooooooooooommmo sel fatt.
"The squared factor in Af and fi appears as a result that the neighborhood growth is quadratic in Ie? "This is why this subtlety is often left out in group equivariance literature.
17

---

Published as a conference paper at ICLR 2021
LiftingSelrAtt
LayerNorm
Layerforr
Swish
Syish
:
Gzoup
LayezHoora
AttentionBlock
Nq
Syich
LayerNorm
AttectionProbe
xL
Swish
MaxPoolRd
Layar Norr
DropOut
- -----------
DropOut
Seliih
LayarMoorn
Linear
GlobalPooling
Figure F.1: Graphical description of group self- aelf-attention networks. Dot-lined blocks depict optional blocks. Linear layers applied point-wise across the feature map. Swish non-linearities (Ramachan- dran al., 2017) and layer normalization (Ba et al., 2016) are used all across the network. The GlobalPooling block consists of max-pool over group elements followed by spatial mean-pool.
F.1 ROTATED MNIST
For rotational MNIST we use group self-attention network composed of 5 attention blocks with 20 channels. We utilize automatic mixed precision during training to reduce memory requirements. attention dropout rate and value. dropout rate are both set to 0.I. We train for 300 epochs and utilize the Adam optimizer, batch size of 8, weight decay 0.0001 and learning rate of 0.001.
F.2 CIFAR-10
For CIFAR-10 we group self-attention network composed of attention blocks with 96 channels for the first two blocks and 192 channels for the rest. value wruuueddroouutrrre rate are both set 0.1. We use dropout on the input with rate of 0.3 and additional dropout blocks of rate 0.2 followed by spatial max-pooling after the second and fourth block. We did not use automatic mixed precision training for this dataset as made all models diverge. We perform training for 350 epochs and utilize stochastic gradient descent with momentum of 0.9 and cosine learning rate scheduler with base learning rate 0.01 (Losschilov Hutter, 2016). We utilize batch size of 24, weight decay of 0.0001 and He's 's initialization.
dropoutrttttee and
F.3 PATCHCAMELYON
For PatchCamelyon we use group self-attention network composed 4 attention blocks with 12 channels for the first block, 24 channels for the second block, 48 channels for the third and fourth blocks and 96 channels for the last block. attention on dropout rate and value_de wwr.uuudrrouuu.... pate both
O. We requirements. We did not use automatic mixed precision training for this dataset as it made all models diverge. We perform training for 100 epochs, utilize stochastic gradient descent with a momentum of 0.9 and cosine learning rate scheduler with base learning rate 0.01 (Loshchilov & Hutter, 2016). We utilize a batch size of 8, weight decay of 0.0001 and He's initialization.
additional max-pooling block after the lifting block reduce memory
G PROOFS
Proof of Proposition 4.1. If the self-attention formulation provided in Eqs. 3, 8, is permutation equivariant, then it must hold that m[{+[/]]()) = [=[m[[][() Consider permuted input signal Lx(f((t)= (()). The self aelf-attention operation on Ix[f]] is given by:
m[{()[f]](1) Four
(4, "
U he[M]] Je8
(1)
U Σ (12 ( (f(T ( (47")) ((T
=xPout
18

---

Published as a conference paper at ICLR 2021
U Σ x()
(F()))
=ouut
U (1(() je
-Poul
mL [f((t) - [f[ (0)
= (m)
Here have used the substitution -= 7(i) and j= Since the summation is defined over the entire set we have that [-))sss-]== = Ejes("] Conclusively, we see that m² m[Lx[f]](() = [=[mm[[[[]
Hence, permutation equivariance indeed holds.
Proof of Claim 4.2. Permutation equivariance. If the self-attention formulation provided in Eq. 10 is permutation equivariant, then it must hold that m[{f{[f),pI(i) = [m] pl))) Consider a permuted input signal Lx[f((t)= ((r=((()) The self-attention operation on [x] is given by:
m (())
U [1 (1+L)(())+++))000)000))))) ( (4M()))) jen(i)
= PPou
As discussed in $3.2, since there exists permutations in Sy able to send elements j in 71(1) to elements j outside of 711), it is trivial to show that Eq. 10 is not equivariant to Sy. Consequently, in order to provide a more interesting analysis, we consider global attention here, i.e. cases where N())= 8. As shown for Proposition 4.1, this self-attention instantiation is permutation equivariant. Consequently, by considering this particular case, we are able to explicitly analyze the effect of introducing absolute positional encodings into the self-attention formulation. We have then that:
m2 (T),
U AE
Pout
U
Pour
U (() he8
(1) wyy
(f()))
Pour
Here we have used the substitution 7 (i) and == Since the summation is defined over the entire set we have that Ene) [+] Ejs({] Since (()(1)) # p(7) and ((=(f)) * p()), we are unable to reduce the expression further towards the form of m[ p](). Consequently, we conclude that absolute position-aware self-attention is not permutation equivariant.
Translation equivariance. If the self-attention formulation provided in Eq. 10, is translation equiv- ariant, then it must hold that m( (t) (f p)(() Consider a translated input signal L,Lff(()= (a("(((()- y)). The self-attention operation on Lf is given by:
m/ (((
U
(1) (1) (() (4LT))))))((((((())))))))))))))) 44 (GLAOO))
= Pour he[M] (eN(s)
U
E
p())) (()
U)
1 (1(x - (1)+
Pour U
Eor (2())))))))))))))))))))))))))))) (z(1)+))))
he[m] x (0))
Plyy (f() p(= (z() (f())
Fou U hc[H]
(1)) (<(()))(((())))))))))))))(()))>) +)))))+)))))))))))))))))))
Alyy
(1)
(F())
19

---

Published as a conference paper at ICLR 2021
Here, we have used the substitution i= (z)))))))-)) i==x>< +y) +
j x (x()) v). Since the area of summation remains equal to any translation y IR we have:
ΣH (a) Ma (a())))) (a (x(i))
ΣDI
E1I
jen(i)
Hence, we can further reduce the expression above as:
mL (( Poou
g)), U (1 (f(1) o² (x(()= + 9) Pkey ( (1 d² (x() x))) (r()) jen()
Since, p(i) y + p(i) and p() + # p()), we are unable to reduce the expression further towards the form of (i). Consequently, we conclude that the absolute positional encoding does not allow for translation equivariance either.
Proof of Claim 4.3. If the self-attention formulation provided in Eq. 11 translation equivariant, then it must hold that m² [n] Ly]mm"[f.P][C Consider a translated input signal L f(x (x(1)- - y)). The self-attention operation on Ly[s] is given by:
m² [=((f))))())
U 34ll
-Pout
Je((()
U
Pour
Me[H)jeN((
*((,)))) (F()))))))(((((((((((((((
Σot(OS (1P4))(0))) U he[l] M(a(+ty)ell (=(()=0)) + (x()) + u) x -x(i)=x()+++) + w
(f(i), Pkay (s())
Pour
(1))
Here, we have used the substitution i= j x (x(i) y). By using the definition of p(i,)) f) we can further reduce the expression above as:
and j= x (x
=aall U Eo (())) (1) (f(i)
he[H] -- (a())))) (x(0+t))
( (((0))))((((00)))))))00 ((())
Σom ((0)))()))() =Poont U
he[l]] (((()t))enz (x(1)+0)
U Earca (10 (1()) (h) (f(i) (())) (f())
Pout
he[m] x- (M=(((())))((( (x()+z))
Since the area of the summation remains equal to any translation e Rdd, we have that:
Σl. (a(() * <y)ell( (a(()))))
Σ1 (a (((i))
Σ1I. jen())
Resultantly, we can further reduce the expression above as:
m² [[oLLll.p]()) = Poul he[m] jen(()
U Eg((s (((1)) ((.))))))
=m](f,]] m" m" px (x()) (x(1)-)))
= =[,[("(U.)](()
We see that indeed m² r=]== mm" [f,]mmmm[[[]]]]"" p])(.. Consequently, we conclude that the relative positional encoding allows for translation equivariance. We emphasize that this is a consequence of = (), Vy e Rd In other words, it comes from the fact that relative positional encoding is invariant the action of the translation group.
the fact that ( v) xx(() + +)))
20

---

Published as a conference paper at ICLR 2021
Proof of Claim 5.1. If the lifting self-attention formulation provided in Eq. 11 is Geeuuivariant, then it must hold that http Consider a g g-transformed input signal
fg(f((1) -y-1)(((())) f ((==((- 9= (v. y E Rd, h e H. The lifting group
self attentionno operation on [g[ff] is given by:
U (o (fg_aaa(111)))11011)0000 (1)
Poul
jen()
Poul U
[o (1) (f(x (h- (1))))))) (h) (f
je(()
+ (ar
-)))))
Poor U he[H] x (
En ()())))))))))))))))))))))))))( k ( (J()
L (fx(() + +, ( + +)))))
(s())
Here we have used the substitution i x i (I(1) f) = i= x (hx(() + y) and j a- (A (x([) y)) j= (ha([)++ + y) By using the definition of p(i, we can further reduce the expression above as:
U - (n (1 (f()),
Pkey (1()
-Paul
he[H] x (ha(() (ha()))))))))))) + +
())(())))))))))))))) (I())
Poort U
Σ Eo
Pry (()
he[m]] x ((a()+))) (1a()))))
((()) (())
(f())
U he[H] x
(0 (ha()))))
(1o (J(O) Plyy ( (1O) u P
(h) (f())
Poul
Since, for unimodular groups the area of summation remains equal for any g & G. we have that:
Σ [.]
E
Σ H Σ
-(xx(()) (e)(( ((z()-o)) (bx (fx(()) = (xt))
jen())
Resultantly, we can further reduce the expression above as: m6g (f (i,,A)
(U [o (f(O), (1O)) (:0 (F(O)) jEn()
Pool
Lp h 21) Pe (x())-
We see indeed that [Gg [/], = Ly [m Consequently, we conclude
that the lifting group self-attention operation is group equivariant. We emphasize once more that this is a consequence of the fact that [g]p](i.)) p((,,) j). Vg G. In other words, it comes from the fact that the positional encoding used is invariant to the action of elements E G.
Li
Proof of Claim 5.2. If the group self-attention formulation provided in Eq. 11 is G-equivariant, then must hold that mf fi, Lf f (( -))(())))))).)) v)), g (y, y Rd, h H. The group
= L.I Consider a g-transformed input signal
self-attention operation on Lg[s]] is given by: [Lll),
PEO (fGCA(1(A))) U Σ Eo (fofall) (
Pour
Reae (1
U heMe (i.))
PPou
-)))
21

---

Published as a conference paper at ICLR 2021
Pout (UT (UEE (he1(h(()))) &n (1(,(f((, (= (i ((a))+))))))))((=>=>>()))))))))) (ka())))))
(1)()
(x (hx()) +), (FG.R))) ()
Here have used the substitutions i 7= (x(1) - =i= ( +
=² h-
and j x- (x)) y))
(hx( y)), h' h By using the definition of
p((i,)),(i (i,h)) we can further reduce the expression above as:
= -Poul U EE (he())))) bF (1p by (f Kkyy (f(Thh) he[H] Aite
(w- (f (E (60 + p (hx(s) + +y (()))))))))))),)--))))))))))))) (h=())x))
U
Σ
Σ
4 key (JG..')
= PPout
x (ha KR (a- (60) he
(ba())
(A (h- (
U ΣΣ oF ((an()))) (1 (f(C, Pray (f(j.h) he[H] AitE (MM))((())))))))M)(())))(+(( Mi' La (...I((..,,,...))))))220
PPout
Furthermore, since for unimodular groups the area of summation remains equal for any transformation g E G. we have that:
Σ F
Σ
E F
E.. F..
(x-1 (ha(( +V) ( (k (x- (ea (z)) ( (IASen(iA)))
Additionally, we have that Exve E. above as:
Resultantly, we can further reduce the expression
m (() Poml
E he[m)] Nest (T.Ae)n La (. (T(I,))))
((C Play (F
=mālf.pl(.h A) pf (x(1) - b)),
- (m (()
We see that indeed [[ = Ly f pll(i, Consequently, we conclude that the group self-attention operation on is group equivariant. We emphasize once more that this is a consequence of the fact that (g(p)((i, ()) = (((i,)). (j. f)). Vg G. In other words, comes from the fact that the positional encoding used invariant to the action of elements g e C. O
22


Reviews:
[{"review_id": "xEOs98b6CI_", "reviewer": "AnonReviewer5", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper develops self-attention mechanisms for group equivariant networks. This is a natural and potentially useful piece of architecture, that, drawing inspiration from its success for text and image analysis, is likely to introduce an effective inductive bias to the group convolution networks. \nThe paper is based on the following series of observations. First, represent an image as a set of $N$ feature vectors (e.g.,  pixels). Second, consider group actions $\\mathcal{G}$ that can be represented as sub-groups of the permutation group $ֿ\\mathbb{S}_N$ acting on the $N$ image pixels. Third, the standard self-attention layer is permutation equivariant. A way to further restrict this equivariance to smaller sub-group of the permutations  $\\mathcal{G}<\\mathbb{S}_N$ is by adding (or concatenating?) to the set of feature vectors, used as keys in the attention layer, functions which are invariant to this subgroup but not to any larger group. This will lead to operators that are equivariant to the desired sub-group $\\mathcal{G}$. Lastly, since doing that directly on the representation of an image leads to larger invariances than desired, the authors work with signals defined over the group and consider regular group action (i.e., change of variables of functions). For example, for the roto-translation group that adds $90$ degrees rotations around the center of the image to the standard translations, this means representing four copies of the image and defining the self-attention layer from the original image and such signals as well as self-attention maps between pairs of such signals. \n\nOverall, I feel this work offers some interesting observations relating symmetry and equivariance/invariance and self-attention. However, I find this work not fully ready in terms of exposition, ideas delivery and elaboration. The paper is not well balanced between developing and explaining its main contribution (starting page 6), and self-attention background. I understand there is much to cover, but explaining the main constructions in half a page is too short. I think the paper keeps reintroducing positional encoding and self-attention in several formulations and there is much space to save there. Furthermore, the explanations and notations are somewhat not coherent and clear, for example (eqs. 13, 14) are not sufficiently explained and demonstrated. Although i dont think it is a necessary condition for publication of such paper, the results don't show an improvement over \"standard\" group convolution methods, although the authors claim expressiveness of their self-attention networks is at-least of that of generalized convolutions. Lastly, the last paragraph before section 6, describing the expressive power of these network is not sufficiently developed and justified. \n\nSome more comments are:\n- There are different notations for concat (word and union) and softmax (word and sigma).\n- Equation 7 has some repeating unclear part.\n- $X_i$ is not defined in equation 5.\n- What happens when the group action cannot be represented as subgroup of permutations? (Page 4, lower part). That is, when there is no such injective map as $x$...\n- What is $m$ in equation in Proposition 4.1?\n- Can $\\mathcal{X}$ be continuous space or should is be discrete as well? If continuous how does $y\\in\\mathcal{X}$ in example 4.2.1 work?\n- Sentence after proposition 4.2 is not clear. \n- After definition 4.3 - \"we shown\" - where?\n- What is $m^r$?\n- \"This allows us to go beyond group discretizations that live in the grid without introducing interpolation artifacts.\" --> How?\n- Before section 5.1 - what is $\\mathcal{H}$?\n- In Proposition 5.2: shouldn't $\\mathcal{L}_g$ also be applied to $\\rho$?\n- Page 6: \"an space\".\n- Page 7: what are linear mappings on $\\mathcal{G}$?\n- Last paragraph before section 6: why is this layer maximally expressive for *sub-groups* of $\\mathbb{S}_N$.\n\nUPDATE:\nI thank the authors for their detailed response and edits. Overall, the authors have addressed my comments, but I have failed to understand some parts. I am still (slightly) positive about this work. >> Please check again for typos.", "extra_scores": {}, "extra_text": {}}, {"review_id": "w-_srXVwjr", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "***Summary***\n\nI would firstly like to thank the authors for an interesting read. I enjoyed going through the submission very much.\n\nThe submission outlines a general framework to make self-attention equivariant under the action of ‘arbitrary’ groups. The authors go about laying out the mathematical foundations of self-attention from a functional standpoint and then find the conditions that must be satisfied by the individual components of the attention mechanism to make it adhere to the imposed equivariance constraints. Namely, they note that absolute positional encodings must be done away with and that activations must be lifted to a homogenous space of the group in question in order for non-trivial (non-constant) relative positional encodings to be used. They experiment on the MNIST-rot, CIFAR10 and PatchCameylon datasets, but unfortunately the results fall significantly behind standard baselines.\n\n\n\n***Pros***\n\nThe paper is written very clearly and from what I could see the mathematics appears to be technically sound. \n\nThe authors also provide an in-depth appendix, which contains a lot of details as to group theory and their experiments.\n\nI think the primary objective of making self-attention equivariant is nice. I also think that the conclusion the authors come to makes sense and in someways (with hindsight, having read the paper) is to be expected. \n\nIn the experiments, the choices of which datasets to use make sense from the point of view of the equivariance literature.\n\nI like it how the authors explain that the networks are steerable by design because the group can act on the positional encodings.\n\n\n\n***Cons and constructive feedback***\n\nI would like to see a discussion comparing the proposed method with the works of “Affine Self-Convolution” by (Diaconu and Worrall, 2019), and “Attentive group equivariant convolutional networks” (Romero et al., 2019), which are both very close works in this area. What are the differences?\n\nThe results are quite disappointing. After the main theory I would have been happy with results en par with existing works, but CIFAR performance of, for instance, 83.7% compared to a baseline of 91.1% really isn’t great. For me this is the main drawback of this work. Furthermore, how are the models comparable with the baselines? Are they matched in terms of numbers of parameters, or size of the activations, or something else? Perhaps this could be a source of the discrepancy between the results.\n\nI would also like to see comparisons with the work of Romero and Diaconu, who have results in their papers.\n\nIn the abstract you state that you can impost equivariance to arbitrary groups? Is this really so? Surely just compact, discrete groups?\n\nEquation 1: you introduce the notation ‘sigma’ for the softmax nonlinearity but then immediately use \\text{softmax} in eqn 1.\n\nEquation 7: what is the difference between the right hand part of the first line and the second line?\n\nDefinition 4.1: is this definition not overly restrictive by focussing on left-regular representations for L and L’? It also precludes the use of irreducible representations as commonly used in the works on group equivariant networks.\n\nExample 4.2.1 I think this could have been stated more simply by specifying the group and action.\n\nDefinition 4.3 shown -> show\n\nProposition 4..2 Nor permutation nor translation -> Neither permutation nor translation\n\nSection 5.2: Good references here would be Theorem 3.1 of “A General Theory of Equivariant CNNs on Homogeneous Spaces”, (Cohen et al., 2018) and Theorem 1 of “B-Spline CNNs on Lie Groups” (Bekkers, 2020). Also, is this statement proven? I couldn’t find it in the submission.\n\nSection 6. Efficient implementation…: I would like to see more detail here since it is not clear in which way the implementation is efficient.\n\nJust a small note: you use American spelling throughout the paper, but missed out behaviour and neighbourhood, which should read behavior and neighborhood.\n\n***Post-rebuttal review***\n\nI have increasing my recommendation to a 6. This is on account of the improvements to the submission by the authors, a detailed rebuttal, and somewhat to align with the recommendations of the other reviewers. I still do find that the experimental results could be improved somewhat, but as another reviewer pointed out, this is not a huge concern within the majority of the equivariance literature. I thank the authors for an interesting read and thank you for a good rebuttal.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "vQMVbMs-7eh", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### Summary\nThe paper proposes an approach for building group-equivariant self-attention networks. The authors use group-invariant positional encodings, and thus the output of a self-attention layer is equivariant under the actions from the considered group. They empirically demonstrate that group-equivariant self-attention networks have an improvement over their non-equivariant counterparts.\n\n### Pros\n1. I enjoyed reading the paper. It has a clear title. The abstract, the intro, and the related work are well-written and give a clear understanding of the content of the paper. The technical sections are coherent and well-structured.\n2. The authors propose a general formulation of group-equivariant self-attention networks and rigorously prove their statements.\n3. The experiments demonstrate the improvements achieved by the added group-equivariance in image classification on several datasets.\n\n### Cons\nI did not find any major weaknesses in the presented paper. I did find some minor issues which, however, do not change the general impression and my decision.\n1. While the authors claim that the proposed formulation is general, they only consider 2 compact groups. A group of reflections and a discrete group of rotations. The authors also claim that current computational constraints restrict the possibility of training models with big vicinities until they converge. This issue can be especially important if a group of dilations (scaling group) is considered. The current scale-equivariant models of Bekkers 2020, Sosnovik et al. 2020, Worrall and Welling 2019, Romero et al. 2020 require one to process images with big filters. Is an implementation of a scale-equivariant self-attention network feasible given the current constraints? A discussion of these limitations would make the contribution of the current paper more clear.\n2. The authors use long equations with a very nested structure. Reading these equations is complicated given the current notation. The main source of confusion comes from repetitive sequences of brackets. A minor rethinking of the notation may improve the readability. For example, using square brackets when a function is an argument i.e. $\\alpha(f)(i, j) \\rightarrow \\alpha[f](i, j)$\n\nIt is a well-written paper. The theory is coherent. The text is easy-to-follow. The experimental section demonstrates a thorough evaluation of the proposed method. \n\n### After-rebuttal comments\nThe authors andwered my questions. My decision stays the same\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "jiKkjHqyeaf", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors describe their contributions in the introduction: the analysis of equivariance of self-attention, and how group invariance in the relative positional encoding enables group equivariance of the self-attention.\n\n\nThe authors have a very condensed related work section without going into any detail but with a lot of citations on (non attention) papers about equivariance. Works missed include Equivariant transformer networks by Tai et al., Equivariant multi-view networks (Esteves et al.), SO(3)-equivariant representations (Esteves et al. ), and early work on equivariant function spaces by Hel-Or and Teo (Canonical Decomposition of Steerable Functions).  \n\nBut probably most relevant is the missing discussion on equivariance on set networks (incl point cloud networks) and graph networks (Maron and Lipman). There is no positional encoding in point clouds but the value at each element are the coordinates and some parallels can be drawn to positional encodings (for example first few layers of pointnet). \n\nThe authors clearly define self-attention, first with matrices, and then with a functional formulation. The functional formulation is elegant but difficult to follow. The concatenation in the multi-head case is an example of where the vector space formulation allows replacing concatenation with a union. The authors might want to explain the benefits compared to tensor formulation.\n\nAuthors first prove permutation equivariance of global self-attention without positional encoding.\nThe proof in G.4.1 is kind of convoluted and might be clearer using only matrices (if $\\Pi$ is permutation than $\\sigma(\\Pi Q K^T \\Pi^T)\\Pi V = \\Pi \\sigma(Q K^T) V$.\n\nThen they prove translational equivariance for relative positional encoding. It is easy to see that relative positional encoding is translation invariant. The step to equivariance has to be followed in the appendix and it would be easier for the reader to provide at least a sketch in the main text. \n\nThe observation about translation paves the ground for generalizing to any group if the positional encoding is invariant to this group. Unfortunately, at this point the discussion becomes very confusing compared to the original \"lifting\" by Cohen and Welling (2016). While this paper makes the lifting appear as a trick, the main idea of Cohen and Welling is that when one applies group convolution on a quotient space, for example SE(2)-convolution on R^2, the result is automatically invariant in SO(2) since R^2=SE(2)/SO(2). Instead, one performs a group correlation where the output is a function of the group action (not of the quotient space) and after this step one applies a group convolution (also appearing in the spherical CNNs (Cohen)as well as in the icosahedral multi-view networks (Esteves)).\n\nThe paper concludes with the claim that linear mappings whose positional encoding is G-invariant are G-equivariant. This is easy to see in the definition of convolution and one can imagine this for linear mappings but it is difficult to see that self-attention is a linear mapping if one looks at eq. 5 and possible definitions of the encoding function (12). I see it through the equivalence proof in Cordonnier but not through the self-attention definition.\n\nExperiments do not show any advantage of equivariant self-attention in z2 or r4 CNNS.\n\nTo summarize the paper is interesting but quite difficult to follow. I wish the paper would follow the tensor notation like in the SE(3)-transformers. Authors should justify the superiority of their formalism vs tensors.\n\nRelated work should not be condensed with mere listing of citations like \\cite{*).\n\nLast: Steerability is claimed in the abstract and the introduction but never mentioned again in the paper. One can somehow see how the positional encoding implies it but a section would be worth to be dedicated to it. \n\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper introduces group equivariant self attention networks constructed by defining positional encoding that are invariant to the group action considered. This is related to equivariance in set networks . The work is sound and the idea of infusing the inductive bias via the positional encoding is  interesting and leads to improvement  results when comparing transformers with equivariance and without it, nevertheless more work needs to be done to bridge the gap in terms of performance with CNNs as pointed by the reviewers.  Authors made an admirable efforts in the rebuttal and in the revision of the paper clarifying most of the reviewers questions and concerns.   Accept 


Decision: Accept (Poster)