G ROUP E QUIVARIANT S TAND -A LONE S ELF -A TTENTION FOR V ISION
A BSTRACT
We provide a general self-attention formulation to impose group equivariance to arbitrary symmetry groups.
This is achieved by deﬁning 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
I NTRODUCTION
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 ﬁrst adopted into
CNNs by channel-wise attention (Hu et al., 2018) and non-local spatial modeling (Wang et al., 2018).
More recently, it 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.1), yet self-attention layers are able to express any convolutional layer (Cordonnier et al., 2020).
This ﬂexibility allows leveraging long-range dependencies under a ﬁxed 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 to other symmetry groups to improve generalization and sampleefﬁciency further (see §2).
Translation equivariance is key to the success of CNNs.
It describes the property that if a pattern is translated, its numerical descriptors are also translated, but not modiﬁed.
In this work, we introduce group self-attention , a self-attention formulation that grants equivariance to arbitrary symmetry groups.
This is achieved by deﬁning positional encodings invariant to the action of the group considered.
In addition to generalization and sample-efﬁciency improvements provided by group equivariance, group equivariant self-attention networks ( GSA-Nets ) bring important beneﬁts over group convolutional architectures: ( i ) Parameter efﬁciency: 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 ﬁxed parameter budget, yet it 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 equivariance properties of self-attention (§4).
• We provide a 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
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/ .
R ELATED W ORK
Several approaches exist which provide equivariance to various symmetry groups.
The translation equivariance of CNNs has been extended to additional symmetries ranging from planar rotations (Dieleman et al., 2016; Marcos et al., 2017; Worrall et al., 2017; Weiler et al., 2018b; 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 et al., 2020) to spherical rotations (Cohen et al., 2018; 2019b; Worrall & Brostow, 2018; Weiler et al., 2018a; 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 on which the kernels are deﬁned.
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 attentive group equivariant networks (Romero & Hoogendoorn, 2019; Diaconu & Worrall, 2019; Romero et 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 speciﬁc weights to particular positions on a large non-local neighborhood on the group.
As a result, attention is bound to growth of the model size, and thus, to negative statistical efﬁciency.
Differently, group self-attention is able to attend over arbitrarily large group neighborhoods under a ﬁxed parameter budget.
In addition, group self-attention is steerable by nature (§5.1) a property primarily exhibited by works carefully designed to that end.
Other way to detach weights from particular positions comes by parameterizing convolutional kernels as continuous (constrained) neural networks, e.g., Thomas et al.
; Finzi et al.
Introduced to handle irregularly sampled data, e.g., point-clouds, networks parameterizing convolutional kernels receive relative positions as input and output their values at those positions.
In contrast, our mappings change as a function of the input content.
Most relevant to our work are the SE(3) and Lie Transformers (Fuchs et al., 2020; Hutchinson et al., 2020): two self-attentive architectures equivariant to SE(3) with applications for molecular prediction.
Our work differs in that (i) we consider groups other than SE(3), (ii) focus on theoretical aspects and properties of group equivariant self-attention mappings, and (iii) prioritize applications on visual data.
S TAND -A LONE S ELF -A TTENTION
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.
Deﬁnition.
Let X ∈ R N × C in be an input matrix consisting of N tokens of C in dimensions each.
1 A self-attention layer maps an input matrix X ∈ R N × C in to an output matrix Y ∈ R N × C out as:
[Equation]
(1) with W val ∈ R C in × C h the value matrix , A ∈ R N × N the attention scores matrix , and softmax [ , ∶] ( A ) the attention probabilities .
The matrix A is computed as:
[Equation]
parameterized by query and key matrices W qry , W key ∈ R C in × C h .
In practice, it has been found beneﬁcial to apply multiple self-attention operations, also called heads , in parallel, such that different heads are able to attend to different parts of the input.
In this multi-head self-attention formulation, the output of H heads of output dimension C h are concatenated and projected to C out as:
[Equation]
with a projection matrix W out ∈ R HC h × C out and a bias term b out ∈ R C out .
3. T HE ROLE OF THE POSITIONAL ENCODING
Note that the self-attention operation deﬁned in Eq.
3 is equivariant to permutations of the input rows of X .
That is, a permutation of the rows of X will produce the same output Y up to this permutation.
Hence, self-attention is blind to the order of its inputs, i.e., it is a set operation.
Illustratively, an input image is processed as a 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 set elements.
Absolute positional encoding.
Vaswani et al.
(2017) introduced a (learnable) positional encoding
P ∈ R N × C in for each input position which is added to the inputs when computing the attention scores:
[Equation]
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.
This positional encoding injects additional structural information about the tokens into the model, which makes it 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 inefﬁciency is addressed by relative positional encodings .
Relative positional encoding.
Introduced by Shaw et al.
, relative encodings consider the relative distance between the query token (cid:105) – the token we compute the representation of –, and the key token (cid:106) – the token we attend to –.
The calculation of the attention scores (Eq. 2) then becomes:
[Equation]
[Equation]
where P x ( (cid:106) )− x ( (cid:105) )
∈ R 1 × C in is a vector representation of the relative shift and x ( (cid:105) ) is the position of the token (cid:105) as deﬁned in §3.2.
Consequently, similar patterns can be recognized at arbitrary positions, as relative query-key distances always remain equal.
3. A FUNCTIONAL FORMULATION TO SELF - ATTENTION
Notation.
We denote by [ n ] the set { 1 , 2 ,...,n } .
Given a set (cid:83) and a vector space (cid:86) , L (cid:86) ( (cid:83) ) will denote the space of functions { f ∶ (cid:83) → (cid:86) } .
Square brackets are used when functions are arguments.
Let (cid:83) = { (cid:105) } N (cid:105)
=
1 be a set of N elements.
A matrix X ∈ R N × C in can be interpreted as a vector-valued function f ∶ (cid:83) → R C in that maps element sets (cid:105) ∈ (cid:83) to C in -dimensional vectors: f ∶ (cid:105) ↦ f ( (cid:105) ) .
Consequently, a matrix multiplication, XW ⊺ y , of matrices X ∈ R N × C in and W y ∈ R C out × C in can be represented as a function ϕ y ∶ L (cid:86) C in ( (cid:83) )
→ L (cid:86) C out ( (cid:83) ) , ϕ y ∶ f
( (cid:105) )
↦ ϕ y ( f ( (cid:105) )) , parameterized by W y , between functional spaces L (cid:86) C in ( (cid:83) ) =
{ f ∶ (cid:83) → R C in } and L (cid:86) C out ( (cid:83) )
= { f ∶ (cid:83) → R C out } .
Following this notation, we can represent the position-less attention scores calculation (Eq. 2) as:
[Equation]
The function α [ f ] ∶ (cid:83) ×
(cid:83) → R maps pairs of set elements (cid:105) , (cid:106) ∈ (cid:83) to the attention score of (cid:106) relative to (cid:105) .
Therefore, the self-attention (Eq. 1) can be written as:
[Equation]
f
[Equation]
, (7) where σ (cid:106) =
softmax (cid:106) and ζ [ f ] ∶ (cid:83) → R C h .
Finally, multi-head self-attention (Eq. 3) can be written as:
[Equation]
where ∪ is the functional equivalent of the concatenation operator concat , and (cid:109)
[ f ] ∶ (cid:83) → R C out .
Local self-attention.
Recall that α [ f ] assigns an attention scores to every other set element (cid:106) ∈ (cid:83) relative to the query element (cid:105) .
The computational cost of self-attention is often reduced by restricting its calculation to a local neighborhood (cid:78) ( (cid:105) ) around the query token (cid:105) analogous in nature to the local receptive ﬁeld of CNNs (Fig.
A.1a).
Consequently, local self-attention can be written as:
[Equation]
Note that Eq.
9 is equivalent to Eq.
8 for (cid:78) ( i ) =
(cid:83) , i.e. when considering global neighborhoods.
Absolute positional encoding.
The absolute positional encoding is a function ρ ∶ (cid:83) → R C in that maps set elements (cid:105) ∈ (cid:83) to a vector representation of its position: ρ ∶ (cid:105)
→ ρ ( (cid:105) ) .
Note that this encoding is not dependent on functions deﬁned on the set but only on the set itself.
Consequently, absolute position-aware self-attention (Eq. 4) can be written as:
[Equation]
The function ρ can be decomposed as two functions
ρ P ○ x : ( i ) the position function x ∶ (cid:83) → (cid:88) , which provides the position of set elements in the underlying space (cid:88)
(e.g., pixel positions), and, ( ii ) the positional encoding ρ P ∶ (cid:88)
→ R C in , which provides vector representations of elements in (cid:88) .
This distinction will be of utmost importance when we pinpoint where exactly (group) equivariance must be imposed to the self-attention operation (§4.3, §5).
Relative positional encoding.
Here, positional information is provided in a relative manner.
That is, we now provide vector representations of relative positions ρ ( (cid:105) , (cid:106) )
∶=
ρ P ( x ( (cid:106) )
− x ( (cid:105) )) among pairs ( (cid:105) , (cid:106) ) , (cid:105) ∈ (cid:83) , (cid:106) ∈
(cid:78) ( i ) .
Consequently, relative position-aware self-attention (Eq. 5) can be written as:
[Equation]
E QUIVARIANCE A NALYSIS OF S ELF -A TTENTION
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. G ROUP EQUIVARIANCE AND EQUIVARIANCE FOR FUNCTIONS DEFINED ON SETS
First we provide the general deﬁnition of group equivariance and reﬁne it to relevant groups next.
Additionally, we deﬁne the property of unique equivariance to restrict equivariance to a given group.
Deﬁnition 4.1 ( Group equivariance ) .
Let (cid:71) be a group (Def. C.1), (cid:83) , (cid:83) ′
be sets, (cid:86) , (cid:86)
′ be vector spaces, and (cid:76) g [⋅] , (cid:76) ′
g [⋅] be the induced (left regular) representation (Def. C.4) of (cid:71) on L (cid:86) ( (cid:83) ) and
L (cid:86) ′
( (cid:83) ′ ) , respectively.
We say that a map ϕ
∶ L (cid:86) ( (cid:83) )
→ L (cid:86) ′ ( (cid:83) ′ ) is equivariant to the action of (cid:71) – or (cid:71) -equivariant –, if it commutes with the action of (cid:71) .
That is, if:
[Equation]
Example 4.1.1 ( Permutation equivariance ) .
Let (cid:83) =
(cid:83)
= { (cid:105) } N (cid:105) =
1 be a set of N elements, and (cid:71) =
S N be the group of permutations on sets of N elements.
A map ϕ
∶ L (cid:86) ( (cid:83) )
→ L (cid:86) ′ ( (cid:83) ) is said to be equivariant to the action of S N – or permutation equivariant –, if:
[Equation]
where (cid:76) π
[ f ]( (cid:105) )
∶= f
( π − 1 ( (cid:105) )) , and π ∶ (cid:83) → (cid:83) is a bijection from the set to itself.
The element π ( (cid:105) ) indicates the index to which the (cid:105) -th element of the set is moved to as an effect of the permutation π .
In other words, ϕ is said to be permutation equivariant if it commutes with permutations π ∈ S N .
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 deﬁned on sets.
Luckily, as we consider sets gathered from homogeneous spaces (cid:88) where these transformations are welldeﬁned, e.g., R 2 for pixels, there exists an injective map x ∶ (cid:83) → (cid:88) that associates a position in (cid:88) to each set element, the position function .
In Appx.
D we show that the action of (cid:71) on such a set is well-deﬁned and induces a group representation to functions on it.
With this in place, we are now able to deﬁne equivariance of set functions to groups whose actions are deﬁned on homogeneous spaces.
Deﬁnition 4.2 ( Equivariance of set functions to groups acting on homogeneous spaces ) .
Let (cid:71) be a group acting on two homogeneous spaces (cid:88) and (cid:88) ′ , let (cid:83) , (cid:83) ′ be sets and (cid:86) ,
(cid:86) ′ be vector spaces.
Let x ∶ (cid:83) → (cid:88) and x ′ ∶ (cid:83) ′
→ (cid:88) ′ be injective maps.
We say that a map ϕ
∶ L (cid:86) ( (cid:83) )
→ L (cid:86) ′ ( (cid:83) ′ ) is equivariant to the action of (cid:71) – or (cid:71) -equivariant –, if it commutes with the action of (cid:71) .
That is, if:
[Equation] [Equation]
, where (cid:76) (cid:103) [ f ]( (cid:105) ) ∶= f
( x − 1 ( (cid:103) − 1 x ( (cid:105) )))
, (cid:76) ′ (cid:103) [ f ]( (cid:105) ) ∶= f
( x ′− 1 ( (cid:103) − 1 x ′ ( (cid:105) ))) are the induced (left regular) representation of (cid:71) on L (cid:86) ( (cid:83) ) and L (cid:86) ′
( (cid:83) ′ ) , respectively.
I.o.w., ϕ is said to be (cid:71) -equivariant
if a transformation (cid:103) ∈ (cid:71) on its argument produces a corresponding transformation on its response.
Example 4.2.1 ( Translation equivariance ) .
Let (cid:83) , (cid:83) ′ be sets and let x ∶ (cid:83) → (cid:88) and x ′ ∶ (cid:83) ′
→ (cid:88) ′ be injective maps from the sets (cid:83) , (cid:83) ′ to the corresponding homogeneous spaces (cid:88) , (cid:88)
′ on which they are deﬁned, e.g., R d and (cid:71) .
With ( (cid:88) , +) the translation group acting on (cid:88) , we say that a map ϕ
∶ L (cid:86) ( (cid:83) )
→ L (cid:86) ′ ( (cid:83) ) is equivariant to the action of ( (cid:88) , +) – or translation equivariant –, if: ϕ
[ (cid:76) y
[ f ]]( (cid:105) )
= (cid:76) ′
y
[ ϕ [ f ]]( (cid:105) ) , ∀ f [Equation] L (cid:86) ( (cid:83) ) , ∀ y ∈ (cid:88) , with (cid:76) y
[ f ]( (cid:105) )
∶= f
( x − 1 ( x ( (cid:105) )
− y )) , (cid:76) ′
y
[ f ]( (cid:105) )
∶= f
( x ′− 1 ( x ′ ( (cid:105) )− y )) .
I.o.w., ϕ is said to be translation equivariant if a translation on its argument produces a corresponding translation on its response.
4. E QUIVARIANCE PROPERTIES OF SELF - 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 in Appx.
G.
Proposition 4.1.
The global self-attention formulation without positional encoding (Eqs. 3, 8) is permutation equivariant.
That is, it holds that: (cid:109)
[ (cid:76) π [ f ]]( (cid:105) )
= (cid:76) π
[ (cid:109)
[ f ]]( (cid:105) ) .
Note that permutation equivariance only holds for global self-attention.
The local variant proposed in
Eq.
9 reduces permutation equivariance to a smaller set of permutations where neighborhoods are conserved under permutation, i.e., S (cid:78) =
{ π ∈ S N ∣ (cid:106) ∈ (cid:78) ( (cid:105) )
→ π ( (cid:106) )
∈ (cid:78) ( (cid:105) ) , ∀ (cid:105) ∈ (cid:83) } .
Permutation equivariance induces equivariance to important (sub)groups.
Consider the cyclic group of order 4, (cid:90) 4 = { e,r,r 2 ,r 3 } which induces planar rotations by 90 ○ .
As every rotation in (cid:90)effectively induces a permutation of the tokens positions, it can be shown that (cid:90) 4 is a subgroup of S N , i.e., S N ≥ (cid:90) 4 .
Consequently, maps equivariant to permutations are automatically equivariant to (cid:90) 4 .
However, as the permutation equivariance constraint is harder than that of (cid:90) 4 -equivariance, imposing (cid:90) 4 -equivariance as a result of permutation equivariance is undesirable in terms of expressivity.
Consequently, Ravanbakhsh et al.
(2017) introduced the concept of unique (cid:71) -equivariance to express the family of functions equivariant to (cid:71) but not equivariant to other groups (cid:71) ′ ≥ (cid:71) :
Deﬁnition 4.3 ( Unique (cid:71) -equivariance ) .
Let (cid:71) a subgroup of (cid:71) ′ , (cid:71) ≤ (cid:71) ′
(Def. C.2).
We say that a map ϕ is uniquely (cid:71) -equivariant
iff it is (cid:71) -equivariant
but not (cid:71) ′
-equivariant for any (cid:71) ′ ≥ (cid:71) .
In the following sections, we show that we can enforce unique equivariance not only to subgroups of S
N , e.g., (cid:90) 4 , but also to other interesting groups not contained in S N , e.g., groups of rotations ﬁner than 90 degrees.
This is achieved by enriching set functions with a proper positional encoding.
Proposition 4.2.
Absolute position-aware self-attention (Eqs. 4, 10) is neither permutation nor translation equivariant.
i.e., (cid:109)
[ (cid:76) π [ f ] ,ρ ]( (cid:105) )
≠ (cid:76) π
[ (cid:109)
[ f,ρ ]]( (cid:105) ) and (cid:109)
[ (cid:76) y
[ f ] ,ρ ]( (cid:105) )
≠ (cid:76) y
[ (cid:109)
[ f,ρ ]]( (cid:105) ) .
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: (cid:109) r
[ (cid:76) y
[ f ] ,ρ ]( (cid:105) )
= (cid:76) y
[ (cid:109) r
[ f,ρ ]]( (cid:105) ) .
4. W HERE EXACTLY IS EQUIVARIANCE IMPOSED IN SELF - ATTENTION ?
In the previous section we have seen two examples of successfully imposing group equivariance to self-attention.
Speciﬁcally, we see that no positional encoding allows for permutation equivariance and that a relative positional encoding allows for translation equivariance.
For the latter, as shown in the proof of Prop. 4.3 (Appx. G), this comes from the fact that for all shifts y ∈ (cid:88) , ρ
[Equation]
(cid:106) ) .
That is, from the fact that the relative positional encoding is invariant to the action of the translation group, i.e., (cid:76) y
[ ρ ]( (cid:105) , (cid:106) )
=
ρ ( (cid:105) , (cid:106) ) , ∀ y ∈ (cid:88) .
Similarly, the absence of positional encoding – more precisely, the use of a constant positional encoding –, is what allows for permutation equivariance
(Prop. 4.1, Appx. G).
Speciﬁcally, constant positional encodings ρ c ( (cid:105) )
=
c , ∀ (cid:105) ∈ (cid:83) are invariant to the action of the permutation group, i.e., (cid:76) π [ ρ c ]( (cid:105) )
=
ρ c ( (cid:105) ) , ∀ π ∈ S N .
From these observations, we conclude that (cid:71) -equivariance is obtained by providing positional encodings which are invariant to the action of the group (cid:71) , i.e., s.t., (cid:76) (cid:103)
[ ρ ] = ρ , ∀ (cid:103) ∈ (cid:71) .
Furthermore, unique (cid:71) -equivariance is obtained by providing positional encodings which are invariant to the action of (cid:71) but not invariant to the action of any other group (cid:71) ′ ≥ (cid:71) .
This is a key insight that allows us to provide (unique) equivariance to arbitrary symmetry groups, which we provide next.
G ROUP E QUIVARIANT S TAND -A LONE S ELF -A TTENTION
In §4.3 we concluded that unique (cid:71) -equivariance is induced in self-attention by introducing positional encodings which are invariant to the action of (cid:71) but not invariant to the action of other groups (cid:71) ′ ≥ (cid:71) .
However, this constraint does not provide any information about the expressivity of the mapping we have just made (cid:71) -equivariant.
Let us ﬁrst illustrate why this is important:
Consider the case of imposing rotation and translation equivariance to an encoding deﬁned in R 2 .
Since translation equivariance is desired, a relative positional encoding is required.
For rotation equivariance, we must further impose the positional encoding to be equal for all rotations.
That is (cid:76) θ
[ ρ ]( (cid:105) , (cid:106) ) !
=
ρ ( (cid:105) , (cid:106) ) , ∀ θ ∈ [ 0 , 2 π ] , where (cid:76) θ [ ρ ]( (cid:105) , (cid:106) )
∶=
ρ P ( θ − 1 x ( (cid:106) ) − θ − 1 x ( (cid:105) )) , and θ − 1 depicts a 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 2 to a space where rotations are explicitly encoded (Fig.
B.1).
To this end, one performs self-attention operations for positional encodings (cid:76) θ
[ ρ ] of varying values θ and indexes their responses by the corresponding θ value.
Next, as rotations are now explicitly encoded, a positional encoding can be deﬁned in this space which is able to discriminate among rotations (Fig.
B.2).
This in turn allows for rotation equivariance instead of rotation invariance.
It has been shown both theoretically (Ravanbakhsh, 2020) and empirically (Weiler & Cesa, 2019) that the most expressive class of (cid:71) -equivariant functions is given by functions that follow the regular representation of (cid:71) .
In order to obtain feature representations that behave that way, we introduce a lifting self-attention layer (Fig.
B.1, Eq. 14) that receives an input function on R d and produces a feature representation on (cid:71) .
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 R d can be provided by pooling over (cid:72) .
In short, we provide a pure selfattention 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 in 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 ˜ (cid:104) ∈ (cid:72) 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 ﬁne discrete approximations of the underlying continuous group (Tab. 2).
In fact, (Weiler & Cesa, 2019, Tab. 3) show via extensive experiments that networks using regular representations and ﬁne enough discrete approximations consistently outperform networks handling continuous groups via irreducible representations.
We conjecture this is a 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. G ROUP SELF - ATTENTION IS AN STEERABLE OPERATION
Convolutional ﬁlters are commonly parameterized by weights on a discrete grid, which approximate the function implicitly described by the ﬁlter at the grid locations.
Unfortunately, for groups whose action does not live in this grid, e.g., 45 ○ rotations, the ﬁlter must be interpolated.
This is problematic as these ﬁlters are typically small and the resulting interpolation artifacts can be severe (Fig. 2a).
Steerable CNNs tackle this problem by parameterizing convolutional ﬁlters on a continuous basis on which the action of the group is well-deﬁned, e.g., circular harmonics (Weiler et al., 2018b), B-splines
[Figure]
(Bekkers, 2020).
In group self-attention, the action of the group leaves the content of the image intact and only modiﬁes the positional encoding (Figs.
B.1, B.2).
As the positional encoding lives on a continuous space, it can be transformed at an arbitrary grade of precision without interpolation (Fig. 2b).
5. L IFTING AND GROUP SELF - ATTENTION
Lifting self-attention (Fig. B.1).
Let (cid:71) =
R d ⋊ (cid:72) be an afﬁne group (Def. C.3) acting on R d .
The lifting self-attention (cid:109) r (cid:71) ↑ [ f,ρ ] ∶ L (cid:86) ( R d ) →
L (cid:86) ′
( (cid:71) ) is a map from functions on R d to functions on (cid:71) obtained by modifying the relative positional encoding
ρ
( (cid:105) , (cid:106) ) by the action of group elements (cid:104) ∈ (cid:72) : { (cid:76) (cid:104)
[ ρ ]( (cid:105) , (cid:106) )}
(cid:104) ∈ (cid:72) , (cid:76) (cid:104)
[ ρ ]( (cid:105) , (cid:106) )
=
ρ P ( (cid:104) − 1 x ( (cid:106) ) − (cid:104) − 1 x ( (cid:105) )) .
It corresponds to the concatenation of multiple self-attention operations (Eq. 11) indexed by (cid:104) with varying positional encodings (cid:76) (cid:104)
[ ρ ] : (cid:109) r (cid:71)
↑ [ f,ρ ]( (cid:105) , (cid:104) )
= (cid:109) r
[ f, (cid:76) (cid:104)
[ ρ ]]( (cid:105) )
[Equation] [Equation]
Proposition 5.1.
Lifting self-attention is (cid:71) -equivariant.
That is, it holds that: (cid:109) r (cid:71) ↑ [ (cid:76) [Equation]
[ f ] ,ρ ]( [Equation] , (cid:104) )
= (cid:76) [Equation] [ (cid:109) r (cid:71) ↑ [ f,ρ ]]( (cid:105) , (cid:104) ) .
Group self-attention (Fig.
B.2).
Let (cid:71) =
R d ⋊ (cid:72) be an afﬁne group acting on itself and f ( (cid:105) , ˜ (cid:104) ) ∈
L (cid:86) ( (cid:71) ) , (cid:105) ∈ (cid:83) , ˜ (cid:104) ∈ (cid:72) , be a function deﬁned on a set immersed with the structure of the group (cid:71) .
That is, enriched with a positional encoding ρ
(( (cid:105) , ˜ (cid:104) ) , ( (cid:106) , ˆ (cid:104) ))
∶=
ρ P (( x ( (cid:106) ) − x ( (cid:105) ) , ˜ (cid:104) − 1 ˆ (cid:104) )) , (cid:105) , (cid:106) ∈ (cid:83) , ˜ (cid:104) , ˆ (cid:104) ∈ (cid:72) .
The group self-attention (cid:109) r (cid:71)
[ f,ρ ] ∶ L (cid:86) ( (cid:71) )
→ L (cid:86) ′ ( (cid:71) ) is a map from functions on (cid:71) to functions on (cid:71) obtained by modifying the group positional encoding by the action of group elements (cid:104)
[Equation]
ˆ
[Equation]
It corresponds to the concatenation of multiple self-attention operations (Eq. 11) indexed by h with varying positional encodings (cid:76) (cid:104) [ ρ ] and followed by a summation over the output domain along ˜ (cid:104) :
[Equation]
[Equation]
) , ( (cid:106) , ˆ (cid:104) ))⟩) ϕ
( h ) val ( f ( (cid:106) , ˆ (cid:104) ))) .
In contrast to vanilla and lifting self-attention, the group self-attention neighborhood (cid:78) ( i, ˜ (cid:104) ) is now deﬁned on the group.
This allows distinguishing across group transformations, e.g., rotations.
Proposition 5.2.
Group self-attention is (cid:71) -equivariant.
That is, it holds that: (cid:109) r (cid:71) [
(cid:76) (cid:103) [ f ] ,ρ ]( (cid:105) , (cid:104) )
= (cid:76) (cid:103) [ (cid:109) r (cid:71)
[ f,ρ ]]( (cid:105) , (cid:104) ) .
Non-unimodular groups, i.e., groups that modify the volume of the objects they act upon, such as the dilation group, require a special treatment.
This treatment is provided in Appx.
E.
5. G ROUP SELF - ATTENTION IS A GENERALIZATION OF THE GROUP CONVOLUTION
We have demonstrated that it is sufﬁcient to deﬁne self-attention as a function on the group (cid:71) and ensure that (cid:76) (cid:103)
[ ρ ] = ρ ∀ (cid:103) ∈ (cid:71) in order to enforce (cid:71) -equivariance.
Interestingly, this observation is inline with the main statement of Kondor & Trivedi (2018) for (group) convolutions: “ the group convolution on (cid:71) is the only (unique) (cid:71) -equivariant linear map ”.
In fact, our ﬁnding can be formulated as a generalization of Kondor & Trivedi (2018)’s statement as: “ Linear mappings on (cid:71) whose positional encoding is (cid:71)
-invariant are (cid:71) -equivariant. ” This statement is more general than that of Kondor & Trivedi (2018), as it holds for data structures where (group) convolutions are not well-deﬁned, e.g., sets, and it is equivalent to Kondor & Trivedi
(2018)’s statement for structures where (group) convolutions are well-deﬁned.
It is also congruent with results complementary to Kondor & Trivedi (2018)
(Cohen et al., 2019a; Bekkers, 2020) as well as several 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 ( i ) 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 a multi-head self-attention layer provided enough heads (Cordonnier et al., 2020), yet self-attention often uses larger receptive ﬁelds.
As a 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 (cid:72) in their derivations, and deﬁning neighborhoods in this new space with a proportionally larger number of heads.
Statement ( ii ) stems from the ﬁnding of Ravanbakhsh (2020) that functions induced by regular group representations are equivariant universal approximators provided full kernels , i.e., global receptive ﬁelds.
Global receptive ﬁelds are required to guarantee that the equivariant map is able to model any dependency among input components.
Global receptive ﬁelds 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.
E XPERIMENTS
We perform experiments on three image benchmark datasets for which particular forms of equivariance are desirable.
5 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 a fair view to the yet present gap between self-attention and convolutional architectures in vision tasks, also present in their group equivariant counterparts.
Efﬁcient 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.
Speciﬁcally, we calculate self-attention scores w.r.t.
the content X once, and reuse them for all transformed versions of the positional encoding { (cid:76) (cid:104)
[ ρ ]} (cid:104) ∈ (cid:72) .
Model designation.
We refer to translation equivariant self-attention models as Z2 SA .
Reﬂection equivariant models receive the keyword M , e.g., Z2M SA , and rotation equivariant models the keyword
Rn , where n depicts the angle discretization.
For example, R8 SA depicts a model equivariant to rotations by 45 degrees.
Speciﬁc model architectures are provided in Appx.
F.
RotMNIST.
The rotated MNIST dataset (Larochelle et al., 2007) is a classiﬁcation dataset often used as a standard benchmark for rotation equivariance.
It consists of 62 k gray-scale 28x28 uniformly rotated handwritten digits, divided into training, validation and test sets of 10 k , 2 k and 50 k images.
First, we study the effect of the neighborhood size on classiﬁcation 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 as a 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 vicinities 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 insufﬁcient for all models to converge equally well.
Unfortunately, due to computational constraints, we were not able to perform this experiment for a larger number of epochs.
We consider an in-depth study of this behavior an important direction for future work.
Next, we compare GSA-Nets 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
[Figure]
[Table]
[Figure] [Table] [Figure] [Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
best trade-off between accuracy and convergence time.
Our results show that ﬁner 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 ﬁner angle discretizations to fall within the same pixel.
CIFAR-10.
The CIFAR-10 dataset (Krizhevsky et al., 2009) consists of 60 k real-world 32x32 RGB images uniformly drawn from 10 classes, divided into training, validation and test sets of 40 k , 10 k and 10 k images.
Since reﬂection is a symmetry that appears ubiquitously in natural images, we compare GSA-Nets equivariant to translation and reﬂection in this dataset (Tab. 2).
Our results show that reﬂection equivariance indeed improves the classiﬁcation performance of the model.
PCam.
The PatchCamelyon dataset (Veeling et al., 2018) consists of 327 k 96x96 RGB image patches of tumorous/non-tumorous 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 at arbitrary positions and poses, we compare GSA-Nets equivariant to translation, rotation and reﬂection (Tab. 2).
Our results show that incorporating equivariance to reﬂection in addition to rotation, as well as providing ﬁner group discretization, improve classiﬁcation performance.
D ISCUSSION AND F UTURE W ORK
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 our theoretical results indicate that GSA-Nets can be more expressive than G-CNNs (§ 5.3), further research in 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.
The main drawback of our approach is the quadratic memory and time complexity typical of selfattention.
This is an active area of research, e.g., Kitaev et al.
; Wang et al.
Zaheer et al.
(2020); Choromanski et al.
and we believe that efﬁciency advances to 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 a set operation, GSA-Nets provide straightforward solutions to set-like data types, e.g., point-clouds, graphs, symmetric sets, which may beneﬁt from additional geometrical information, e.g., Fuchs et al.
; Maron et al.
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: a type of positional encoding. A CKNOWLEDGMENTS
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 ﬁnanced as part of the Efﬁcient Deep Learning (EDL) programme (grant number P16-25), partly funded by the Dutch Research Council (NWO) and Semiotic Labs.
JeanBaptiste Cordonnier is ﬁnanced 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 A PPENDIX A C ONVOLUTION AND THE S ELF -A TTENTION : A G RAPHICAL C OMPARISON
[Figure]
Figure A. kernels Fig.
A.1a) and self-attention (Fig. A.1b).
Given a budget of 9 parameters, a convolutional ﬁlter ties these parameters to speciﬁc 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 to any speciﬁc positions at all.
Contrarily, it compares the representations of all tokens falling in its receptive ﬁeld.
As a result, provided enough heads, self-attention is more general than convolutions, as it can represent any convolutional kernel, e.g., Fig.
A.1a, as well as several other functions deﬁned on its receptive ﬁeld.
B L IFTING AND G ROUP S ELF -A TTENTION : A GRAPHICAL D ESCRIPTION
attention analogous to the regular lifting group convolution broadly utilized in group equivariant learning literature, e.g., Cohen & Welling (2016); Romero et al.
(2020a).
is deﬁned as (cid:72) = { e, (cid:104) , (cid:104) 2 , (cid:104) 3 } , where (cid:104) depicts a 90 ○ rotation.
Analogous to lifting self-attention (Fig. B.1), group self-attention corresponds to a concatenation of ∣ (cid:72) ∣ = 4 self-attention operations between the input f and (cid:104) -transformed versions of the positional encoding (cid:76) [ ρ ] , ∀ (cid:104) ∈ (cid:72) .
However, in contrast to lifting self-attention, both f and ρ are now deﬁned on the
Consequently, an additional sum over ˜ (cid:104) is required during the operation (c f., Eq. 16).
Since (cid:90) 4 is a cyclic group, i.e., (cid:104)
= e , functions on (cid:90) 4 are often represented as responses on a ring (right side of the image).
This is a self-attention analogous to the regular group convolution broadly utilized in group equivariant learning literature, e.g., Cohen & Welling (2016); Romero et al.
(2020a).
C C ONCEPTS F ROM G ROUP T HEORY
Deﬁnition C.1 ( Group ) .
A group is an ordered pair ( (cid:71) , ⋅) where (cid:71) is a set and ⋅ ∶ (cid:71) × (cid:71) → (cid:71) is a binary operation on (cid:71) , such that (i) the set is closed under this operation, (ii) the operation is associative,
i.e., ( (cid:103) 1 ⋅ (cid:103) 2 ) ⋅ (cid:103) 3 = (cid:103) 1 ⋅ ( (cid:103) 2 ⋅ (cid:103) 3 ) , (cid:103) 1 , (cid:103) 2 , (cid:103) 3 ∈ (cid:71) , (iii) there exists an identity element (cid:101) ∈ (cid:71) s.t.
∀ (cid:103) ∈ (cid:71) we have (cid:101) ⋅ (cid:103) =
(cid:103) ⋅ (cid:101) =
(cid:103) , and (iv) for each (cid:103) ∈ (cid:71) , there exists an inverse (cid:103)
− 1 s.t.
(cid:103) ⋅ (cid:103) − 1 = (cid:101) .
Deﬁnition C.2 ( Subgroup ) .
Let ( (cid:71) , ⋅) be a group.
A subset (cid:72) of (cid:71) is a subgroup of (cid:71) if (cid:72) is nonempty and closed under the group operation and inverses (i.e., (cid:104) 1 , (cid:104) 2 ∈ (cid:72) implies that (cid:104) − 11 ∈ (cid:72) and (cid:104) 1 ⋅ (cid:104) 2 ∈ (cid:72) ).
If (cid:72) is a subgroup of (cid:71) we write (cid:72) ≤ (cid:71)
Deﬁnition C.3 ( Semi-direct product and afﬁne groups ) .
In practice, one is mainly interested in the analysis of data deﬁned on R d , and, consequently, in groups of the form (cid:71)
=
R d ⋊ (cid:72) , resulting from the semi-direct product ( ⋊ ) between the translation group ( R d , + ) and an arbitrary (Lie) group (cid:72) that acts on R d , e.g., rotation, scaling, mirroring, etc.
This family of groups is referred to as afﬁne groups and their group product is deﬁned as:
[Equation]
, (17) with (cid:103) =
( x 1 , (cid:104) 1 ) , (cid:103) 2 = ( x 2 , (cid:104) 2 ) ∈ G , x 1 , x 2 ∈ R d and (cid:104) 1 , (cid:104) 2 ∈ (cid:72) .
The operator ⊙ denotes the action of (cid:104) ∈ (cid:72) on x ∈ R d , and it describes how a vector x ∈ R d is modiﬁed by elements (cid:104) ∈ (cid:72) .
Deﬁnition C.4 ( Group representation ) .
Let (cid:71) be a group and L 2 ( X ) be a space of functions deﬁned on some vector space X .
The (left) regular group representation of (cid:71) is a linear transformation (cid:76) ∶ (cid:71) ×
L 2 ( X ) → L 2 ( X ) , ( (cid:103) ,f ) ↦ (cid:76) (cid:103) [ f ] ∶= f
( (cid:103) − 1 ⊙ x ) , that shares the group structure via:
[Equation]
f
[Equation]
(18) for any (cid:103) 1 , (cid:103) 2 ∈ G , f ∈ L 2 ( X ) .
That is, concatenating two such transformations, parameterized by (cid:103) 1 and (cid:103) 2 , is equivalent to a single transformation parameterized by (cid:103) 1 ⋅ (cid:103) 2 ∈ (cid:71) .
If the group (cid:71) is afﬁne, the group representation (cid:76) (cid:103) can be split as:
[Equation]
with g = ( x, (cid:104) ) ∈ (cid:71) , x ∈ R d and (cid:104) ∈ (cid:72) .
Intuitively, the representation of (cid:71) on a function f describes how the function as a whole, i.e., f ( x ) , ∀ x ∈ X , is transformed by the effect of group elements (cid:103) ∈ (cid:71) .
D A CTIONS AND R EPRESENTATIONS OF G ROUPS A CTING ON H OMOGENEOUS S PACES FOR F UNCTIONS D EFINED ON S ETS
In this section we show that the action of a group (cid:71) acting on a homogeneous space (cid:88) is well deﬁned on sets (cid:83) gathered from (cid:88) , and that it induces a group representation of functions deﬁned on (cid:83) .
Let (cid:83) = { (cid:105) } be a set and (cid:88) be a homogeneous space on which the action of (cid:71) is well-deﬁned, i.e., (cid:103) x ∈ (cid:88)
, ∀ (cid:103) ∈ (cid:71) , ∀ x ∈ (cid:88) .
Since (cid:83) has been gathered from (cid:88) , there exists an injective map x ∶ (cid:83) → (cid:88) , that maps set elements (cid:105) ∈ (cid:83) to unique elements x (cid:105) ∈ (cid:88) .
That is, there exists a map x ∶ (cid:105) ↦ x (cid:105) that assigns an unique value x (cid:105) ∈ (cid:88) to each (cid:105) ∈ (cid:83) corresponding to the coordinates from which the set element has been gathered.
Since the action of (cid:71) is well deﬁned on (cid:88) , it follows that the left regular representation (Def. C.4) (cid:76) (cid:103) [ f (cid:88) ]
( x (cid:105) )
∶= f (cid:88) ( (cid:103) − 1 x (cid:105) )
∈ L (cid:89) ( (cid:88) ) of
functions f (cid:88) ∈ L (cid:89) ( (cid:88) ) exists and is well-deﬁned.
Since x is injective, the left regular representation (cid:76) (cid:103) [ f (cid:88) ]
( x (cid:105) )
=
f (cid:88) ( (cid:103) − 1 x (cid:105) ) can be expressed uniquely in terms of set indices as (cid:76) (cid:103) [ f (cid:88) ]
( x (cid:105) )
=
f (cid:88) ( (cid:103) − 1 x (cid:105) )
=
f (cid:88) ( (cid:103) − 1 x ( (cid:105) )) .
Furthermore, its inverse x − 1 ∶ (cid:88) → (cid:83) , x − 1 ∶ x (cid:105) → (cid:105) also exist and is well-deﬁned.
As a consequence, points x (cid:105) ∈ (cid:88) can be expressed uniquely in terms of set indices as (cid:105) =
x − 1 ( x (cid:105) ) , (cid:105) ∈
(cid:83) .
Consequently, functions f (cid:88) ∈ L (cid:89) ( (cid:88) ) can be expressed in terms of functions f ∈ L (cid:89) ( (cid:83) ) by means of the equality f (cid:88) ( (cid:105) )
=
f ( x − 1 ( x (cid:105) )) .
Resultantly, we see that the group representation (cid:76) (cid:103) [ f (cid:88) ]( x (cid:105) )
=
f (cid:88) ( (cid:103) − 1 x ( (cid:105) )) can be described in terms of functions f ∈ L (cid:89) ( (cid:83) ) as:
[Equation]
[Equation]
with a corresponding group representation on L (cid:89) ( (cid:83) ) given by (cid:76) (cid:103) [ f ]( (cid:105) )
=
f ( x − 1 ( (cid:103) − 1 x ( (cid:105) ))) , and an action of group elements (cid:103) ∈ (cid:71) on set elements (cid:105) given by (cid:103)(cid:105) ∶= x − 1 ( (cid:103) x ( (cid:105) )) .
E T HE C ASE OF N ON -U NIMODULAR G ROUPS : S ELF -A TTENTION ON THE D ILATION -T RANSLATION G ROUP
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.
Non-unimodular 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 (cid:72) =
( R > 0 , ×) .
To illustrate why this distinction is important, consider the following example:
Imagine we have a circle on R of area πr 2 .
If we rotate, mirror or translate the circle, its size is kept constant.
If we increase its radius by a factor (cid:104) ∈ R > 0 , however, its size would increase by (cid:104) 2 .
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 deﬁne a neighborhood (cid:78) for which the original circle ﬁts perfectly.
Since the size of the circle is not modiﬁed for any translated, rotated or translated versions of it, we would still be able to detect the circle regardless of these transformations.
If we scale the circle by a factor of (cid:104) > 1 , however, the circle would fall outside of our neighborhood (cid:78) and hence, we would not be able to recognize it.
A solution to this problem is to scale our neighborhood (cid:78) in a proportional way.
That is, if the circle is scaled by a factor (cid:104) ∈ R > 0 , we scale our neighborhood by the same factor (cid:104) : (cid:78) → (cid:104)(cid:78) .
Resultantly, the circle would fall within the neighborhood for any scale factor (cid:104) ∈ R > 0 .
Unfortunately, there is a problem: self-attention utilizes summations over its neighborhood.
Since ∑ (cid:105) ∈ (cid:104)(cid:78) (cid:105) >
∑ (cid:105) ∈ (cid:78) (cid:105) , for (cid:104) > 1 , and ∑ (cid:105) ∈ (cid:104)(cid:78) (cid:105) <
∑ (cid:105) ∈ (cid:78) (cid:105) , for (cid:104) < 1 , the result of the summations would still differ for different scales.
Speciﬁcally, this result would always be bigger for larger versions of the neighborhood.
This is problematic, as the response produced by the same circle, would still be different for different scales.
In order to handle this problem, one utilizes a normalization factor proportional to the change of size of the neighborhood considered.
This ensures that the responses are equivalent for any scale (cid:104) ∈ R > 0 .
That is, one normalizes all summations proportionally to the size of the neighborhood.
As a result, we obtain that ∑ (cid:105) ∈ (cid:104) 1 (cid:78) ( (cid:104) 21 ) − 1 (cid:105) =
∑ (cid:105) ∈ (cid:104) 2 (cid:78) ( (cid:104) 22 ) − 1 (cid:105) , ∀ (cid:104) 1 , (cid:104) 2 ∈ R > 0 .
In the example above we have provided an intuitive description of the (left invariant) Haar measure d µ ( (cid:104) ) .
As its name indicates, it is a measure deﬁned on the group, which is invariant over all group elements (cid:104) ∈ (cid:72) .
For several unimodular groups, the Haar measure corresponds to the Lebesgue measure as the volume of the objects the group acts upon is kept equal, i.e., d µ ( (cid:104) )
=
d (cid:104) .
For non-unimodular groups, however, the Haar measure requires a normalization factor proportional to the change of volume of these objects.
Speciﬁcally, the Haar measure corresponds to the Lebesgue measure times a normalization factor (cid:104) d , where d corresponds to the dimensionality of the space R
d the group acts upon (Bekkers, 2020; Romero et al., 2020b), i.e., d µ ( (cid:104) )
= 1 (cid:104)
d
d (cid:104) .
In conclusion, in order to obtain group equivariance to non-unimodular groups, the lifting and group self-attention formulation provided in Eqs. 14, 16 must be modiﬁed via normalization factors proportional to the group elements (cid:104) ∈ (cid:72) .
Speciﬁcally, they are redeﬁned as:
[Equation]
[Table]
The factor d + 1 in Eq.
21 results from the fact that the summation is now performed on the group (cid:71) =
R d ⋊ (cid:72) , an space of dimensionality d + 1 .
An interesting case emerges when global neighborhoods are considered, i [Table]
e., s.t. (cid:78) ( (cid:105) ) = (cid:83) , ∀ (cid:105) ∈ (cid:83) .
Since (cid:104)(cid:78) ( (cid:105) )
= (cid:78) ( (cid:105) )
= (cid:83) for any (cid:104) > 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. C URRENT 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 efﬁcient self-attention in group self-attention, e.g., Kitaev et al.
; Wang et al.
(2020); Zaheer et al.
; Katharopoulos et al.
Choromanski et al.
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.
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 to dilated convolutions via “sparse” dilations of the self-attention neighborhood.
As a 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 j } j max j = 0 , and thus, less general that the scale-equivariant architectures listed before.
F E XPERIMENTAL D ETAILS
In this section we provide extended details over our implementation as well as the exact architectures and optimization schemes used in our experiments.
All our models follow the structure shown in Fig.
F.1 and vary only in the number of blocks and channels.
All self-attention operations utilize 9 heads.
We utilize PyTorch for our implementation.
Any missing speciﬁcation can be safely considered to be the PyTorch default value.
Our code is publicly available at https://github.com/dwromero/g selfatt .
F. R OTATED MNIST
For rotational MNIST we use a 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.1.
We train for 300 epochs and utilize the Adam optimizer, batch size of 8, weight decay of 0.0001 and learning rate of 0.001.
F. CIFAR-
For CIFAR-10 we use a group self-attention network composed of 6 attention blocks withchannels for the ﬁrst two blocks and 192 channels for the rest.
attention dropout rate and value dropout rate are both set to 0.1.
We use dropout on the input with a 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 it made all models diverge.
We perform training for 350 epochs and 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 24, weight decay of 0.0001 and He’s initialization.
F. P ATCH C AMELYON
For PatchCamelyon we use a group self-attention network composed of 4 attention blocks withchannels for the ﬁrst block, 24 channels for the second block, 48 channels for the third and fourth blocks and 96 channels for the last block.
attention dropout rate and value dropout rate are both set to 0.1.
We use an additional max-pooling block after the lifting block to reduce memory 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.
G P ROOFS
Proof of Proposition 4.1.
If the self-attention formulation provided in Eqs.
3, 8, is permutation equivariant, then it must hold that (cid:109) [ (cid:76) π
[ f ]]( (cid:105)
[Equation]
= (cid:76) π
[ (cid:109)
[ f ]]( (cid:105) ) .
Consider a permuted input signal (cid:76) π
[ f ]( (cid:105) )
=
f ( π − 1 ( (cid:105) )) .
The self-attention operation on (cid:76) π
[ f ] is given by:
[Equation] [Equation]
Here we have used the substitution ¯ (cid:105) =
π (
[Equation]
and ¯ (cid:106) =
π ( (cid:106) ) .
Since the summation is deﬁned over the entire set we have that ∑ π ( ¯ (cid:106) )
∈
(cid:83) [⋅] =
∑ ¯ (cid:106) ∈ (cid:83) [⋅] .
Conclusively, we see that (cid:109) [ (cid:76) π
[ f ]]( (cid:105) )
= (cid:76) π
[ (cid:109)
[ f ]]( (cid:105) ) .
Hence, permutation equivariance indeed holds.
Proof of Claim 4.2.
Permutation equivariance.
If the self-attention formulation provided in Eq.is permutation equivariant, then it must hold that (cid:109) [ (cid:76) π
[ f ] ,ρ ]( (cid:105) ) =
(cid:76) π
[ (cid:109)
[ f,ρ ]]( (cid:105) ) .
Consider a permuted input signal (cid:76) π
[ f ]( (cid:105) )
=
f ( π − 1 ( (cid:105) )) .
The self-attention operation on (cid:76) π
[ f ] is given by:
[Equation]
As discussed in §3.2, since there exists permutations in S N able to send elements (cid:106) in (cid:78) ( i ) to elements ˜ (cid:106) outside of (cid:78) ( i ) , it is trivial to show
that Eq.
10 is not equivariant to S N .
Consequently, in order to provide a more interesting analysis, we consider global attention here, i.e., cases where (cid:78) ( i ) =
(cid:83) .
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:
[Equation] [Equation]
[Equation] [Equation]
[Equation]
Here we have used the substitution ¯ (cid:105) =
π ( (cid:105) ) and ¯ (cid:106) =
π ( (cid:106) ) .
Since the summation is deﬁned over the entire set we have that ∑ π ( ¯ (cid:106) )
∈
(cid:83) [⋅] =
∑ ¯ (cid:106) ∈ (cid:83) [⋅] .
Since ρ ( π ( ¯ (cid:105) ))
≠ ρ ( ¯ (cid:105) ) and ρ ( π ( ¯ (cid:106) ))
≠ ρ ( ¯ (cid:106) ) , we are unable to reduce the expression further towards the form of (cid:109)
[ f,ρ ]( ¯ (cid:105) ) .
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 equivariant, then it must hold that (cid:109) ( (cid:76) y
[ f ] ,ρ )( (cid:105) )
= (cid:76) y
[ (cid:109) ( f,ρ )]( (cid:105) ) .
Consider a translated input signal
[Equation]
y [ f ]
[Equation]
) =
f ( x − 1 ( x (
[Equation]
) − y )) .
The self-attention operation on (cid:76) y
[ f ] , is given by:
[Equation]
Here, we have used the substitution ¯ (cid:105) =
x − 1 ( x ( (cid:105) )− y )
⇒ (cid:105) =
x − 1 ( x ( ¯ (cid:105) )
+ y ) and ¯ (cid:106) =
x − 1 ( x ( (cid:106) )− y )
⇒ (cid:106) =
x − 1 ( x ( ¯ (cid:106) )
+ y ) .
Since the area of summation remains equal to any translation y ∈ R d , we have:
[Equation]
Hence, we can further reduce the expression above as:
[Equation]
Since, ρ ( ¯ (cid:105) ) + y ≠ ρ ( ¯ (cid:105) ) and ρ ( ¯ (cid:106) )
+ y ≠ ρ ( ¯ (cid:106) ) , we are unable to reduce the expression further towards the form of (cid:109) ( f,ρ )( ¯ (cid:105) ) .
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 is translation equivariant, then it must hold that (cid:109) r
[ (cid:76) y
[ f ] ,ρ ]( (cid:105) )
= (cid:76) y
[ (cid:109) r
[ f,ρ ]]( (cid:105) ) .
Consider a translated input signal (cid:76) y
[ f ]( (cid:105) )
=
f ( x − 1 ( x ( (cid:105) ) − y )) .
The self-attention operation on (cid:76) y
[ f ] is given by:
[Equation]
Here
[Equation]
we have used the substitution ¯ (cid:105) =
x − 1 ( x ( (cid:105) )− y )
⇒ (cid:105) =
x − 1 ( x ( ¯ (cid:105) )
+ y ) and ¯ (cid:106) =
x − 1 ( x ( (cid:106) )− y )
⇒ (cid:106) =
x − 1 ( x ( ¯ (cid:106) )
+ y ) .
By using the deﬁnition of ρ ( (cid:105) , (cid:106) )
we can further reduce the expression above as:
[Equation]
Since the area of the summation remains equal to any translation y ∈ R d
[Equation]
we have that:
[Equation]
Resultantly
[Equation]
we can further reduce the expression above as:
[Equation]
]( (cid:105) )
We see that indeed (cid:109) r
[ (cid:76) y
[ f ] ,ρ ]( (cid:105) )
= (cid:76) y
[ (cid:109) r
[ f,ρ ]]( (cid:105) ) .
Consequently, we conclude that the relative positional encoding allows for translation equivariance.
We emphasize that this is a consequence of the fact that ρ ( x − 1 ( x ( ¯ (cid:105) ) + y ) ,x − 1 ( x ( ¯ (cid:106) ) + y )))
=
ρ ( ¯ (cid:105) , ¯ (cid:106) ) , ∀ y ∈ R d .
In other words, it comes from the fact that relative positional encoding is invariant to the action of the translation group.
Proof of Claim 5.1.
If the lifting self-attention formulation provided in Eq.
11 is (cid:71) -equivariant
, then it must hold that (cid:109) r (cid:71) ↑ [ (cid:76) (cid:103)
[ f ] ,ρ ]( (cid:105) , (cid:104) ) =
(cid:76) (cid:103) [ (cid:109) r (cid:71) ↑ [ f,ρ ]]( (cid:105) , (cid:104) ) .
Consider a (cid:103) -transformed input signal (cid:76) [Equation]
[ f ]
[Equation]
= [Equation] (cid:76) ˜ (cid:104) [ [Equation]
[Equation]
=
[Equation]
( x
[Equation]
( ˜ (cid:104) − 1 ( x ( (cid:105) ) − y ))) , (cid:103) =
[Equation]
, ˜ (cid:104) ) , y ∈ R d , ˜ (cid:104) ∈ (cid:72) .
The lifting group self-attention operation on (cid:76) (cid:103) [ f ] is given by:
[Equation]
[Table]
we have used
[Table] [Equation] [Table] [Equation]
[Table] [Equation] [Table] [Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
By using the deﬁnition of ρ ( (cid:105) , (cid:106) )
we can further reduce the expression above as:
[Equation]
Since
[Equation]
for unimodular groups
[Equation]
the area of summation remains equal for any (cid:103) ∈ (cid:71)
[Equation]
we have that:
[Equation]
Resultantly, we can further reduce the expression above as:
[Equation]
[Equation]
We see indeed that (cid:109) r (cid:71) ↑ [ (cid:76) (cid:121) (cid:76)
˜ (cid:104)
[ f ] ,ρ ]( (cid:105) , (cid:104) ) =
(cid:76) (cid:121) (cid:76) ˜
(cid:104)
[ (cid:109) r (cid:71) ↑ [ f,ρ ]]( (cid:105) , (cid:104) ) .
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 (cid:76) g
[ ρ ]( (cid:105) , (cid:106) )
=
ρ ( (cid:105) , (cid:106) ) , ∀ (cid:103) ∈ (cid:71) .
In other words, it comes from the fact that the positional encoding used is invariant to the action of elements (cid:103) ∈ (cid:71) .
Proof of Claim 5.2.
If the group self-attention formulation provided in Eq.
11 is (cid:71) -equivariant
, then it must hold that (cid:109) r (cid:71) [ (cid:76) (cid:103) [ f ] ,ρ ]( (cid:105) , (cid:104) )
= (cid:76) (cid:103) [ (cid:109) r (cid:71)
[ f,ρ ]]( (cid:105) , (cid:104) ) .
Consider a (cid:103) -transformed input signal (cid:76) (cid:103) [ f ]( (cid:105) , ˜ h )
=
(cid:76) y (cid:76) ¯ (cid:104)
[ f ]( (cid:105) , ˜ (cid:104) )
=
f ( ρ − 1 ( ¯ (cid:104) − 1 ( ρ ( (cid:105) )
− y ))
, ¯ (cid:104) ˜ (cid:104) ) , (cid:103) =
( y, ¯ (cid:104) ) , y ∈ R d , ¯ (cid:104) ∈ (cid:72) .
The group self-attention operation on (cid:76) (cid:103) [ f ] is given by:
[Equation]
[Table]
[Table]
By using the deﬁnition of ρ (( (cid:105) , ˜ (cid:104) )
[Equation]
( (cid:106) , ˆ (cid:104) ))
we can further reduce the expression above as:
[Equation]
[Table]
[Equation] [Table] [Equation] [Table]
[Equation] [Table] [Equation] [Table]
[Equation]
[Table]
Furthermore
[Table]
since for unimodular groups the area of summation remains equal for any transformation (cid:103) ∈ (cid:71)
[Table]
we have that:
[Equation]
Additionally, we have that ∑ ¯ (cid:104) ˜
(cid:104)
[Equation]
∈ (cid:72) [⋅] =
∑ ˜ (cid:104)
∈ (cid:72) [⋅] .
Resultantly, we can further reduce the expression above as:
[Equation]
We see that indeed (cid:109) r (cid:71) [ (cid:76) y (cid:76) ¯ h
[ f ] ,ρ ]( (cid:105) , (cid:104) ) =
(cid:76) y (cid:76) ¯ h
[ (cid:109) r
(cid:71)
[ f,ρ ]]( (cid:105) , (cid:104) ) .
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 (cid:76) (cid:103)
[ ρ ](( (cid:105) , ˜ (cid:104) ) , ( (cid:106) , ˆ (cid:104) ))
=
ρ (( (cid:105) , ˜ (cid:104) ) , ( (cid:106) , ˆ (cid:104) )) , ∀ (cid:103) ∈ (cid:71) .
In other words, it comes from the fact that the positional encoding used is invariant to the action of elements (cid:103) ∈ (cid:71) .