A Unified Framework for Vision Transformers Equivariant to Discrete Subgroups of O(2) Tı̄kun Ông1 and Georg Bökman2 1 Independent Researcher
arXiv:2606.27864v1 [cs.CV] 26 Jun 2026
2 University of Amsterdam
Abstract Vision transformers have become a dominant architecture for visual recognition. However, standard models do not explicitly encode the planar symmetries that arise in many vision domains. We introduce a family of vision transformers equivariant to arbitrary discrete subgroups of O(2), providing a unified framework that generalizes prior flipping- and 𝐷 4 -equivariant transformer architectures. Our construction yields equivariant analogues of the core transformer components, together with expressivity guarantees for the resulting layers. In particular, we show that whenever 𝐻 ≤ 𝐺, the class of 𝐺-equivariant ViTs embeds naturally into the class of 𝐻-equivariant ViTs. We also prove that, in the single-head setting, the corresponding equivariant self-attention layer realizes every 𝐺-equivariant self-attention map representable by ordinary self-attention. We further construct a 𝐷 6 -equivariant model based on hexagonal patches, making the architecture compatible with six-fold rotational symmetries. We evaluate the resulting models on the PatternNet aerial image dataset in artificially data-scarce regimes across subgroups of 𝐷 4 and 𝐷 6 . Our experiments compare two equivariant attention mechanisms and analyze how the choice of homogeneous-space configurations used in the nonlinearities affects performance. Preliminary results under matched parameter budgets indicate that equivariance can improve recognition accuracy, motivating further study of how discrete symmetry groups shape transformer-based visual recognition models.
1
Introduction
Geometric deep learning is concerned with designing model architectures that systematically incorporate symmetry and geometry of a learning task as inductive bias [6, 9, 15]. An important class of such models are the so-called group-equivariant neural networks [10, 34], which enjoy layer-wise group equivariance. That is, the map represented by each layer, mapping input features to output features, commutes with a priori specified group actions. Using a group-equivariant neural network exploits the existence of a group action on the data, for instance translations and rotations of images, and aim to simplify the learning task by hard-coding this symmetry into the neural network architecture. An important special case of equivariance is invariance, where the group action on the output of the network is trivial. Image classification of aerial imagery is a prototypical invariant task, which we will consider in the experiments. So far, apart from graph neural networks with equivariant features [1, 3, 4, 13, 26, 32], which have been popular in particular due to their applications in chemistry [43], the most prominent examples of equivariant or invariant neural networks are generalizations of convolutional neural networks (CNNs) [5, 10, 11, 17, 18, 22– 24, 37, 39, 40]. These often involve convolutions over groups, with the usual convolution being the special case for the group R𝑛 or Z𝑛 , equivariant to translations. CNNs equivariant to discrete subgroups of the roto-reflection group O(2) have been widely studied and often outperform ordinary CNNs in the equal parameter setting [39]. Recently, CNNs have been replaced by vision transformers (ViTs) in many state-of-the-art computer vision models [12, 28, 38]. There are multiple reasons for preferring ViTs, including ease of capturing long-range relationships between different parts of an image (or multiple images) and architectural alignment 1
with networks used for other data modalities, such as large language models. Given the success of equivariant CNNs, it is natural to consider equivariant ViTs, which are the main objects of study in this paper. We take a representation-theoretic view of equivariant vision transformers for discrete subgroups 𝐺 ≤ O(2). This viewpoint subsumes the flipping- and 𝐷 4 -equivariant ViTs of Refs. [7, 27], while making it possible to reason about the resulting model classes independently of any particular symmetry group. An important part of our analysis is to compare these classes as the symmetry group varies. We show, for instance, that imposing a larger symmetry group does not lead to an unrelated architecture: a 𝐺-equivariant ViT can be regarded naturally as an 𝐻-equivariant ViT for every subgroup 𝐻 ≤ 𝐺. We also analyze the self-attention layer itself and prove that, at least for a single attention head, the equivariant parameterization loses no expressive power relative to standard self-attention once one restricts to maps that are 𝐺-equivariant. At the same time, this formalism brings several architectural choices into focus, including nonlinearities constructed from arbitrary 𝐺-sets and several possibilities for equivariant self-attention mechanisms. The experiments in Section 5 are designed to probe these choices in controlled, small-scale settings, rather than to optimize for large-scale benchmark performance. By separating the representation-theoretic structure from group-specific implementation choices, the framework provides a common basis for constructing, comparing, and analyzing equivariant ViTs across different planar symmetry groups.
2
Related Works
Our work is a generalization of equivariant transformers presented in Refs. [7, 27], rendering these architectures as special cases for 𝐺 = 𝐷 1 (mirror symmetry) and 𝐺 = 𝐷 4 . These architectures in turn closely follow the original Vision Transformer [12], which can be seen as the 𝐺 = {𝑒} (trivial group) case. Other types of equivariant vision transformers have been considered in the literature. Most notably, Ref. [42] uses a “lifting self-attention” layer in the very beginning to lift token features to functions on the group 𝐺 (i.e., spatial domain features). In Ref. [19], discrete subgroups 𝐺 of O(2) (as well as O(3)) are considered, where spatial domain features are used in conjunction with group convolutions for equivariant linear layers. There are also several works on equivariant transformers for point cloud data. In Ref. [2], where an SO(3)-equivariant attention mechanism for 3D point clouds is presented, the token feature vectors transform in the defining (fundamental/three-dimensional) representation of SO(3). In Refs. [8, 14], higher-order SO(3)-tensors (higher-dimensional irreps) are also used. We would also like to note that there is another line of work on equivariant architectures which aims to achieve equivariance by having the model learn to rotate an input image to its “canonical” orientation [20, 21, 35]. These models are frequently called (spatial) transformers in the literature, but their approach is completely distinct to what is commonly referred to as a Transformer following the landmark Attention Is All You Need paper, Ref. [36]. Our vision transformers are transformers in the sense of Ref. [36].
3
Group theory preliminaries
We will assume some familiarity with group theory, a good reference is Serre’s textbook on representation theory, Ref. [33]. Let 𝐺 be a finite group acting on two sets 𝑋, 𝑌 . A map 𝜙 : 𝑋 → 𝑌 is said to be 𝐺-equivariant if 𝜙 commutes with the 𝐺-action. The main goal of Section 4 is to construct Vision Transformer layers that are 𝐺-equivariant, where 𝐺 is a discrete subgroup of O(2). An important form of group actions are group representations, which are linear group actions on vector spaces. Since representation theory will play a central role in our construction of equivariant layers, we briefly recall some basic definitions and mathematical results here. If not otherwise specified, all vector spaces are over R. Definition 1. Let 𝐺 be a finite group. A representation of 𝐺 on a real vector space 𝑉 is a group homomorphism 𝜌 : 𝐺 → GL(𝑉). 2
(i) A representation (𝜌, 𝑉) is irreducible if there exists no nontrivial subspace 𝑈 ⊂ 𝑉 that is invariant under the 𝐺-action. (ii) A representation (𝜌, 𝑉) on an inner product space 𝑉 is orthogonal if 𝜌(𝐺) ⊂ O(𝑉). By standard abuse of notation, we will sometimes write 𝑉 or 𝜌 instead of (𝜌, 𝑉) to refer to a group representation. We often use the shorthand irrep to refer to irreducible representations. Every group has a one-dimensioanl irrep, called the trivial representation and denoted by 𝜌triv , by sending all elements to the identity 1 × 1 matrix. Two representations are said to be isomorphic if there exists a 𝐺-equivariant linear bijection between them. By Maschke’s theorem, any representation (𝜌, 𝑉) of 𝐺 is isomorphic to a direct sum of irreps. In practice, the vector space 𝑉 will always come with a natural inner product, and we always take representations to be orthogonal, which facilitates the construction of equivariant self-attention (see Section 4). For two 𝐺-representations (𝜌, 𝑉), (𝜎, 𝑊), we denote by Hom𝐺 (𝑉, 𝑊) the vector space of 𝐺-equivariant linear maps 𝑉 → 𝑊. We will also write End𝐺 (𝑉) := Hom𝐺 (𝑉, 𝑉), which is an algebra over R. If 𝑉 is any real irrep, by Schur’s lemma, End𝐺 (𝑉) is a division algebra (all nonzero elements are invertible), so End𝐺 (𝑉) R, C, or H by Frobenius’ theorem on division algebras, where H is the algebra of quaternions. The real irrep 𝑉 is then of real type, complex type, and quaternionic type respectively. In this paper, as we consider discrete subgroups 𝐺 of O(2), real irreps are either of real type or of complex type and either oneor two-dimensional. We provide an overview of the irreps of 𝐺 in the Supplementary Material. For our construction of nonlinearities, we will need the following notion: Definition 2. A homogeneous space of a group 𝐺 is a set 𝑋 with a transitive 𝐺-action. Here, transitive means that each element 𝑥 ∈ 𝑋 can be taken to any other element 𝑦 ∈ 𝑋 by the 𝐺-action. For any 𝑥 ∈ 𝑋, we denote by Stab𝐺 (𝑥) := {𝑔 ∈ 𝐺 | 𝑔𝑥 = 𝑥} the stabilizer subgroup. It is then straightforward to show that 𝑋 𝐺/Stab𝐺 (𝑥) as 𝐺-sets (i.e., there is a 𝐺-equivariant bijection). Conversely, for any subgroup 𝐻 ≤ 𝐺, the coset space 𝐺/𝐻 is naturally a homogeneous space. For any two subgroups 𝐻, 𝐻 ′ ≤ 𝐺, the homogeneous spaces 𝐺/𝐻 and 𝐺/𝐻 ′ are isomorphic as 𝐺-spaces if and only if 𝐻 and 𝐻 ′ are conjugate to each other. Thus, 𝐺/𝐻, with one subgroup 𝐻 from each subgroup conjugacy class, exhaust all possible homogeneous spaces of 𝐺 up to isomorphism. Finally, we would like to fix the notation for a construction that is ubiquitous in our work and in geometric deep learning in general. For a set 𝑋 and a vector space 𝑉, we denote by 𝐶 (𝑋, 𝑉) the vector space of all maps 𝑋 → 𝑉. If there is a 𝐺-action on 𝑋, the vector space 𝐶 (𝑋, 𝑉) is in addition a 𝐺-representation, with a group element 𝑔 ∈ 𝐺 acting on a function 𝑓 : 𝑋 → 𝑉 by (𝑔 · 𝑓 ) (𝑥) = 𝑓 (𝑔 −1 𝑥). Clearly, 𝐶 (𝑋, 𝑉) R𝑋 ⊗ 𝑉 canonically. If 𝑉 also carries a 𝐺-representation 𝜌, then a natural 𝐺-representation on 𝐶 (𝑋, 𝑉) is given by 𝑔 · 𝑓 (𝑥) = 𝜌(𝑔) 𝑓 (𝑔 −1 𝑥).
4
Method
We start by setting up the underlying geometric structure on which the equivariant ViT will operate. Recall that the Minkowski sum of two subsets 𝑋, 𝑌 of a vector space is defined as 𝑋 + 𝑌 := {𝑥 + 𝑦|𝑥 ∈ 𝑋, 𝑦 ∈ 𝑌 }. Definition 3. Let 𝐺 be a discrete subgroup of O(2). A 𝐺-patchified grid is a set H0 ⊂ R2 that can be written as the Minkowski sum of two 𝐺-stable finite subsets 𝑈, H ⊂ R2 . 𝑈 is called the base patch, and its translates 𝑈𝑎 := 𝑈 + 𝑎, where 𝑎 ∈ H , are the patches of H0 . Note that we do not require the patches 𝑈𝑎 to be disjoint. Since 𝑈 and H are stable under 𝐺, so is H0 , implying that all three subsets of R2 are 𝐺-sets. For example, for 𝐺 = 𝐷 4 acting on R2 by reflections and 90◦ rotations, we can take 2 2 𝑃−1 𝑃−3 𝑃−1 𝑞−1 𝑞−3 𝑞−1 (1) , H= − 𝑈= − ,− ,··· , 𝑃, − 𝑃, · · · , 𝑃 . 2 2 2 2 2 2 3
Then H0 = 𝑈 + H is a usual square grid with (𝑞𝑃) 2 pixels and 𝑞 2 disjoint patches, with each patch having 𝑃2 pixels. An RGB image is then an element of 𝐶 (H0 , R3 ), and each image patch is an element of 𝐶 (𝑈𝑎 , R3 ) 𝐶 (𝑈, R3 ). Before discussing the details of each equivariant layer, we would like to clarify the space in which a single token in an intermediate layer of our model lives. In order for equivariance to make sense at all, a token feature 𝑥 must be an element of a space on which 𝐺 acts. A simple and natural assumption is that 𝑥 belongs to a finite-dimensional 𝐺-representation 𝑉 over R. By Maschke’s theorem, we can take Ê 𝑉= R𝐶𝜌 ⊗ 𝑉𝜌 , (2) b 𝜌∈ 𝐺
b denotes the set of (equivalence classes of) irreps of 𝐺, 𝑉𝜌 is the irrep space of 𝜌, and 𝐶𝜌 is the where 𝐺 multiplicity of the irrep. A 𝐺-equivariant ViT (without class tokens) of depth 𝛿 is the composition 3
patch embed &pos. encoding
Block1
Block2
Block 𝛿
𝐶 (H0 , R ) −−−−−−−−−−−→ 𝐶 (H , 𝑉) −−−−−→ 𝐶 (H , 𝑉) −−−−−→ · · · −−−−−→ 𝐶 (H , 𝑉),
(3)
where each map is 𝐺-equivariant. Here, Block𝑖 is the 𝑖-th transformer block, which consists of a multi-head self attention layer followed by a multilayer perceptron (MLP), both with residual connections. In the rest of this section, we will elaborate the construction of each layer in Eq. (3). In practice, elements of the vector space 𝑉 are stored as a tuple of tensors of shape (𝐶𝜌 , 𝑑 𝜌 ), where 𝑑 𝜌 is the dimension of the irrep 𝜌. For example, the dihedral group 𝐷 6 of order 12 has 6 irreps, labeled by A1 , A2 , B1 , B2 (one-dimensional), and E1 , E2 (two-dimensional). A token feature is then represented by 𝑥 = (𝑥 A1 , 𝑥 A2 , 𝑥 B1 , 𝑥 B2 , 𝑥 E1 , 𝑥 E2 ), where 𝑥 A1 has shape (𝐶A1 , 1), 𝑥 E1 has shape (𝐶E1 , 2), etc. A visualization of feature maps of a 𝐷 6 -equivariant ViT is shown in Figure 1. In principle, we allow arbitrary choices of irrep multiplicities 𝐶𝜌 . A common choice involves 𝐶 copies of the regular representation, where 𝐶𝜌 = 𝐶𝑑 𝜌 for 𝜌 of real type and 𝐶𝜌 = 𝐶𝑑 𝜌 /2 for 𝜌 of complex type.
Figure 1: Feature maps of a 𝐷 6 -equivariant vision transformer after four transformer blocks in a trained c6 , we select a single channel from the irrep component 𝑥 𝜌 ∈ classification model. For each irrep 𝜌 ∈ 𝐷 H 𝐶 𝜌 R ⊗ R ⊗ 𝑉𝜌 . The features in the two-dimensional irreps E1 and E2 are represented by encoding the polar angle and length of a vector in R2 using the hue and a combination of saturation and brightness respectively (see the color wheel for reference). For the one-dimensional irreps (A1 , A2 , B1 , B2 ), red and blue indicate positive and negative values respectively, with gray representing zero.
4
4.1
Linear Layer
We first present a generic equivariant linear layer, which is used in our patch embedding layer as well as in the MLP and self-attention in each transformer É layers É block.𝐶 ′′ 𝐶𝜌′ ′′ = 𝜌 ⊗ 𝑉 are two 𝐺-representations decomposed into Suppose 𝑉 ′ = R ⊗ 𝑉 and 𝑉 𝜌 𝜌 b bR 𝜌∈ 𝐺 𝜌∈ 𝐺 ′ ′′ irreps. Characterizing Hom𝐺 (𝑉 , 𝑉 ), the space of 𝐺-equivariant linear maps 𝑉 ′ → 𝑉 ′′ , is straightforward by Schur’s lemma: Ê Ê ′ ′′ Hom𝐺 (𝑉 ′ , 𝑉 ′′ ) = Hom𝐺 (R𝐶𝜌 ⊗ 𝑉𝜌 , R𝐶𝜌 ⊗ 𝑉𝜌 ) = Mat𝐶𝜌′′ ×𝐶𝜌′ (R) ⊗ End𝐺 (𝑉𝜌 ). (4) b 𝜌∈ 𝐺
b 𝜌∈ 𝐺
Here, Mat𝑚×𝑛 (R) denotes the set of 𝑚 × 𝑛 matrices with entries in R. For practical implementations, this means that a general 𝐺-equivariant linear map 𝑊 ∈ Hom𝐺 (𝑉 ′ , 𝑉 ′′ ) is an irrep-wise linear map, which we Í End𝐺 (𝑉 ) denote by 𝑊 = (𝑊 (𝜌) )𝜌∈ 𝐺b , with 𝑊 (𝜌) = dim 𝑤 𝑖 ⊗ 𝐿 𝑖 , where (𝐿 𝑖 )𝑖 is a chosen basis for End𝐺 (𝑉𝜌 ) 𝑖=1 ′′ ′ and 𝑤 𝑖 contains 𝐶𝜌 × 𝐶𝜌 learnable parameters. If the irrep 𝑉𝜌 is of real type, then (by definition) End𝐺 (𝑉𝜌 ) R, so a linear map is equivariant if and ′ only if it only acts on the channel dimension in R𝐶𝜌 ⊗ 𝑉𝜌 . To implement 𝑊 (𝜌) for a complex-type irrep 𝑉𝜌 , which is always two-dimensional in our case (see Section A of the Supplementary Material), a convenient choice is 𝐿 1 = 10 01 and 𝐿 2 = 01 −1 2 × 2 matrices isomorphic to C. 0 , spanning a subalgebra of 𝐶𝜌′′triv Finally, we allow the possibility of adding a bias 𝑏 ∈ R ⊗ 𝑉𝜌triv only for the trivial representation. That is, the complete linear layer is given by © Ê ª © Ê ª 𝑥 𝜌triv ⊕ 𝑥 𝜌 ® ↦→ (𝑊 (𝜌triv ) 𝑥 𝜌triv + 𝑏) ⊕ 𝑊 (𝜌) 𝑥 𝜌 ® . b b «𝜌∈ 𝐺\{𝜌 ¬ «𝜌∈ 𝐺\{𝜌 ¬ triv } triv }
(5)
This exhausts the space of all 𝐺-equivariant affine maps 𝑉 ′ → 𝑉 ′′ . In general, we denote this space by Aff 𝐺 (𝑉 ′ , 𝑉 ′′ ).
4.2
Patch Embedding and Positional Encoding
Given a discrete subgroup 𝐺 ≤ O(2) and a 𝐺-patchified grid H0 = 𝑈 + H (see Definition 3), the patch embedding layer together with added positional encodings is a 𝐺-equivariant affine map 𝐶 (H0 , R3 ) → 𝐶 (H , 𝑉). We describe their construction in the following. 4.2.1
Patch embedding
For each 𝑎 ∈ H , a patch 𝑥(𝑎) of an input color image I ∈ 𝐶 (H0 , R3 ) R H0 ⊗ R3 is nothing but the restriction of I to 𝑈𝑎 . That is, 𝑥(𝑎) = I|𝑈𝑎 ∈ 𝐶 (𝑈𝑎 , R3 ). Because the 𝑈𝑎 are translates of 𝑈, we can naturally identify 𝑈𝑎 and 𝑈. The patchification (unfold) map is given by P : 𝐶 (H0 , R3 ) → 𝐶 (H , 𝐶 (𝑈, R3 )) R H ⊗ R𝑈 ⊗ R3 ,
I ↦→ (𝑎 ↦→ 𝑥(𝑎)).
(6)
Note that this map is 𝐺-equivariant with respect to the natural 𝐺-action on both sides. An equivariant patch embedding layer is then defined as the composition of P together with an equivariant linear map 𝑙 ∈ Hom𝐺 (𝐶 (𝑈, R3 ), 𝑉): IdH ⊗𝑙
P
𝐶 (H0 , R3 ) −→ 𝐶 (H , 𝐶 (𝑈, R3 )) −−−−→ 𝐶 (H , 𝑉)
(7)
In practice, one fixes an orthonormal basis 𝑒 1 , · · · , 𝑒 𝑑𝜌 for each irrep 𝑉𝜌 , and precomputes an orthonormal 𝜌 basis (𝐸 𝛼 𝑗 )𝜌∈ 𝐺, b 𝛼∈ [𝜈𝜌 ], 𝑗 ∈ [𝑑𝜌 ] for 𝐶 (𝑈, R), where 𝜈 𝜌 is the multiplicity of 𝜌 in 𝐶 (𝑈, R), such that for each 𝜌 and 5
irrep filters
hexagonal patch
E2
C E2
E1
C E1
B2
CB2
B1
CB1
A2
C A2
A1
C A1 H
Figure 2: Illustration of the equivariant patch embedding layer with 𝐺 = 𝐷 6 . 𝜌
𝜌
𝜌
𝛼 ∈ [𝜈𝜌 ], the linear map defined by 𝐸 𝛼 𝑗 ↦→ 𝑒 𝑗 is 𝐺-equivariant from span{𝐸 𝛼1 , · · · , 𝐸 𝛼𝑑𝜌 } onto 𝑉𝜌 . Then, the Í 𝜌 𝜌 𝜌-th component of the patch embedding layer is given by 𝑦 𝜌 (𝑎) 𝑐 = 𝑐′ ,𝑘,𝑙, 𝛼, 𝜇 𝐾 𝛼𝑐𝑐′ ;𝜇 (𝐿 𝜇 ) 𝑘𝑙 ⟨𝐸 𝛼𝑙 , 𝑥(𝑎) 𝑐′ ⟩ 𝑒 𝑘 , Í 𝜌 𝜌 𝜌 where 𝐾 𝜌 ∈ R𝜈𝜌 ×𝐶𝜌 ×3×dim End𝐺 (𝑉𝜌 ) is a learnable tensor. We can intepret 𝐹𝑐𝑐′ 𝑘 := 𝛼,𝜇 𝐾 𝛼𝑐𝑐′ ;𝜇 (𝐿 𝜇 ) 𝑘𝑙 𝐸 𝛼𝑙 ∈ 𝐶 (𝑈, R) as a filter for the 𝑗-th component of the irrep 𝜌 for output channel 𝑐 and input color channel 𝑐 ′ . Fig. 2 illustrates the patch embedding layer for 𝐺 = 𝐷 6 , where both 𝑈 and H are taken to be regular hexagons. 4.2.2
Positional encoding
In this work, we employ learnable absolute positional encodings. That is, a position-dependent learnable element of 𝑉 is added to the token features after the patch embedding layer: PosEnc : 𝐶 (H , 𝑉) ∋ 𝑥 ↦→ 𝑥 + 𝑝 ∈ 𝐶 (H , 𝑉), (8) É 𝐶𝜌 ⊗ 𝑉 where 𝑝 ∈ 𝐶 (H , 𝑉) = R H ⊗ R 𝜌 denotes the positional encodings. As noted in Ref. [27], b 𝜌∈ 𝐺 the map PosEnc is equivariant if and only if 𝑝 is invariant under the 𝐺-action. That is, we require b In practice, we precompute a basis for the 𝐺-invariant 𝜌(𝑔) 𝑝 𝜌 (𝑔 −1 · 𝑎) = 𝑝 𝜌 (𝑎). for all irreps 𝜌 ∈ 𝐺. H subspace of R ⊗ 𝑉𝜌 , and linearly combine them using learnable weights during training.
4.3
Nonlinearity
While nonlinearity is straightforward to implement if the features are represented in the “spatial domain” of the group, it is significantly more complex in our case, as our features are numerically represented as tuples of irrep components. In this section, we describe a type of nonlinearity that first performs a Fourier transform (more precisely, a generalization thereof) of the input features, applies a pointwise nonlinearity, and then transforms back. As will become clear, not only does this procedure generalize the constructions in Refs. [7, 27] to any finite group, it also allows strictly more freedom in constructing nonlinearities. In the second part of this section, we argue that this is the most general class of equivariant nonlinearities for MLPs under certain natural assumptions. 4.3.1
The construction
Fix a (finite) 𝐺-set 𝑋. The set 𝑉˜ := 𝐶 (𝑋, R) of real-valued functions is naturally a 𝐺-representation. If 𝜎 : R → R is any function, then the entrywise application of 𝜎, i.e., 𝐶 (𝑋, R) ∋ (𝑦 𝑚 ) 𝑚∈𝑋 ↦→ (𝜎(𝑦 𝑚 )) 𝑚∈𝑋 ,
6
is 𝐺-equivariant. By Maschke’s theorem, the representation 𝑉˜ is isomorphic to a direct sum of copies of irreps of 𝐺. Let FT denote such an isomorphism: Ê ′ ∼ FT : 𝑉˜ − → R𝐶𝜌 ⊗ 𝑉𝜌 =: 𝑉 ′ (9) b 𝜌∈ 𝐺
˜ The Here, FT stands for Fourier transform. The 𝐶𝜌′ are the multiplicities of the irreps appearing in 𝑉. composition −1
entrywise 𝜎
FT FT 𝑉 ′ −−−−→ 𝑉˜ −−−−−−−−−→ 𝑉˜ −−→ 𝑉 ′
is then equivariant and not linear (if 𝜎 is not linear). For fixed 𝜎 and up to 𝐺-equivariant linear bijections 𝑉 ′ → 𝑉 ′ , the map represented by Eq. (10) depends only on the orbit structure of 𝑋. That is, we can decompose 𝑋 into a disjoint union of copies of homogeneous 𝐺-spaces, 𝑋=
Ä
𝑛𝛼 Ä
(10)
H = {e}
E1
H = {e, t}
A2
𝑋𝛼,
H = C3
(11)
𝛼∈Sub(G)/∼ 𝑠=1
A1 H = D3 and the 𝐺-equivariant MLP constructed using 𝑋 depends only on the integers (𝑛 𝛼 ) 𝛼∈Sub(𝐺)/∼ . Here, entrywise nonlinearity Sub(𝐺)/∼ is the set of equivalence classes of subgroups of 𝐺 with respect to conjugation, and 𝑋 𝛼 Figure 3: Illustration of our equivariant nonlinearity for is the homogeneous space obtained by taking the 𝐺 = 𝐷 3 . In this case, 𝐶A1 = 8, 𝐶A2 = 4, 𝐶E1 = 4, 𝑛 {𝑒} = quotient by any subgroup 𝐻 ∈ 𝛼 in the equivalence 1, 𝑛 {𝑒,𝑡 } = 2, 𝑛𝐶3 = 3, and 𝑛 𝐷3 = 2. class. Hence, we can think of the 𝑋 𝛼 as the “elemenatry lego blocks” for constructing the equivariant Í nonlinear layer. The irrep multiplicities in Eq. (9) can be related to the 𝑛 𝛼 by 𝐶𝜌′ = 𝛼 Γ𝜌𝛼 𝑛 𝛼 , where Γ𝜌𝛼 is the multiplicity of irrep 𝜌 in 𝐶 (𝑋 𝛼 , R). For example, take 𝐺 to be the dihedral group 𝐷 3 of order 6. It has, in total, 4 subgroups up to conjugacy. These are given by (12) cyclic: {𝑒}, {𝑒, 𝑟, 𝑟 2 } = 𝐶3 dihedral: {𝑒, 𝑡}, 𝐷 3
The homogeneous space 𝐷 3 /{𝑒} is simply the regular group action (this is true for any group), which decomposes as A1 ⊕ A2 ⊕ 2 · E1 . The homogeneous space 𝐷 3 /⟨𝑡⟩ is the action on the three vertices of the triangle. Thus, 𝐶 (𝐷 3 /⟨𝑡⟩ , R) is three-dimensional, and it is an easy exercise to verify that this representation decomposes as A1 ⊕ E1 . If we take one copy of 𝐷 4 /{𝑒} and two copies of 𝐷 4 /⟨𝑡𝑟⟩ in Eq. (11), we get in the right hand side of Eq. (9) 𝑉 ′ = (R2 ⊗ 𝑉A1 ) ⊕ (R1 ⊗ 𝑉A2 ) ⊕ (R3 ⊗ 𝑉E1 ), (13) and any nonlinear function 𝜎 gives rise to an equivariant nonlinearity 𝑉 ′ → 𝑉 ′ via Eq. (10). See Fig. 3 for an illustration for a more general choice of homogeneous spaces for 𝐷 3 . 4.3.2
How general is this nonlinearity?
We will now argue that the construction presented in Section 4.3 is the most general type of nonlinearity for an MLP that is 𝐺-equivariant, given some natural assumptions on the nature of the nonlinearity. We assume that an equivariant MLP layer takes the form 𝑙2 ◦ 𝑓 ◦ 𝑙1 , where 𝑙1 : 𝑉 → R𝑛 and 𝑙2 : R𝑛 → 𝑉 are 𝐺-equivariant affine maps, R𝑛 carries a 𝐺-representation, and 𝑓 : R𝑛 → R𝑛 is the entrywise application 7
of any activation function 𝜎 : R → R. The following lemma then implies that this class of MLPs coincides with the class of functions representable by MLPs with 𝐺-equivariant affine maps together with nonlinearity constructed according to Eq. (10): Lemma 1. Suppose a matrix representation 𝜌 : 𝐺 → GL(𝑛) commutes with entrywise application of any function 𝜎 : R → R on R𝑛 , then every 𝜌(𝑔) is a permutation matrix. In other words, the action of 𝐺 on R𝑛 is induced from some action of 𝐺 on the set 𝑋 := {1, · · · , 𝑛}. Lemma 1 follows from known results in the literature [16, 29, 41], we provide a self-contained simple proof in the Supplementary Material. Note that it has been shown [31] that for universal approximation of 𝐺-equivariant maps using MLPs with Ã𝑛 one hidden layer, it is enough to take 𝑋 = 𝑠=1 𝐺. That is, the regular 𝐺-set alone is enough. It is also known that not all 𝐺-sets yield universal approximation [30]. It would be of independent interest to understand which combination of homogeneous spaces realizes approximations of 𝐺-equivariant functions most efficiently.
4.4
Multi-head Self-Attention
Our equivariant attention layers will be maps attn : R H ⊗ 𝑉 → R H ⊗ 𝑉 that are equivariant to the 𝐺-action on R H ⊗ 𝑉. Note that 𝐺 acts on both factors of the tensor product, but only equivariance on the second factor is nontrivial, since self-attention is permutation-equivariant on H . Contrary to Refs. [7, 27], we will outline a general construction for equivariant self-attention and then describe two special cases that are in some sense opposite to each other. The underlying principle that guarantees equivariance is to compute invariant attention scores [2, 25, 27]. In fact, our construction can be summarized as ordinary multi-head self-attention with respect to a 𝐺-invariant inner product and a 𝐺-stable orthogonal decomposition of 𝑉. This is elaborated in the following. Equip 𝑉 with a 𝐺-invariant inner product ⟨·, ·⟩, and suppose 𝑉 = 𝑉1 ⊕ · · · ⊕ 𝑉ℎ is an orthogonal decomposition of 𝑉 into subspaces stable under 𝐺. We will refer to these subspaces as attention heads. Let 𝜙𝑞 , 𝜙 𝑘 , 𝜙 𝑣 ∈ Aff 𝐺 (𝑉, 𝑉) be learnable 𝐺-equivariant affine maps. The raw attention scores 𝛼 (𝑟 ) : H ×H → R in the 𝑟-th head are computed according to 𝛼 (𝑟 ) (𝑎, 𝑏) = ⟨Π𝑟 𝜙𝑞 𝑥(𝑎), Π𝑟 𝜙 𝑘 𝑥(𝑏)⟩ ,
(14)
where Π𝑟 : 𝑉 → 𝑉𝑟 is the orthogonal projection onto the 𝑟-th head. The output token 𝑦 ∈ 𝐶 (H , 𝑉) Í in the 𝑟-th head is given by 𝑦 (𝑟 ) (𝑎) = 𝑏∈ H 𝑠 (𝑟 ) (𝑎, 𝑏)Π𝑟 𝜙 𝑣 𝑥(𝑏), where 𝑠 (𝑟 ) (𝑎, 𝑏) are the attention probabilities, obtained by taking softmax of 𝛼 (𝑟 ) over the second entry. The output token at 𝑎 ∈ H is simply 𝜙𝑜 (𝑦 (1) (𝑎) ⊕ · · · ⊕ 𝑦 (ℎ) (𝑎)), where 𝜙𝑜 ∈ Aff 𝐺 (𝑉, 𝑉) is a learnable output projection. There are several possibilities for the choice of the orthogonal decomposition of 𝑉. For example, we could É Éℎ𝜌 𝐶 /ℎ 𝜌 𝜌 ⊗ 𝑉 , called irrep-wise attention. Another decompose each irrep into ℎ𝜌 heads, 𝑉 = 𝜌 b 𝜌∈ 𝐺 𝑟=1 R É 𝐶 /ℎ 𝜌 natural choice is to take 𝑉𝑟 = ⊗ 𝑉𝜌 , which we call coupled attention. bR 𝜌∈ 𝐺 How expressive is our construction? More precisely, can all functions that are 𝐺-equivariant and expressible using an ordinary self-attention layer be expressed using a 𝐺-equivariant self-attention layer presented in this section? We answer this question affirmatively in the single-head case (see the Supplementary Material for the proof): Theorem 1. Let 𝑉 be an orthogonal 𝐺-representation. Let attn : R 𝐿 ⊗ 𝑉 → R 𝐿 ⊗ 𝑉 be a 𝐺-equivariant function representable by one layer of single-head ordinary self-attention, potentially with biases in the query, key, and value maps. Then attn is representable by one layer of 𝐺-equivariant single-head self-attention. That is, the query, key, and value maps can be taken to be 𝐺-equivariant.
4.5
Class Token and Invariantization
For classification, which is the task considered in our experiments in Section 5, we append a class token cls ∈ 𝑉 right before the first transformer block. Mathematically, this means we replace H with H ⊔ {★}, 8
where ★ is a single point that is invariant under the 𝐺-action, and the class token is the token at ★. For this procedure to be 𝐺-equivariant, the initial class token itself has to belong to the invariant subspace of 𝑉. That is, it can only be nonzero in the trivial representation. Note that after passing through an attention layer, the non-trivial parts of the class token will in general be nonzero via interaction with the other tokens. The class token is plucked out after the final transformer block and its features are used in a final linear layer for classification. Since one expects the classification model to be invariant, that is, the output logits should remain the same if the input image is transformed by a group element, the class token must be invariantized before the linear classification head. Following Ref. [27], we use the following map for invariantization: Ê 𝑉 ∋ cls ↦→ clstrivial ⊕ ∥cls𝜌 ∥ 𝑉𝜌 , (15) b 𝜌∈ 𝐺,𝜌≠trivial
which is followed by concatenation along the channel dimension. Here, ∥·∥ 𝑉𝜌 is any 𝐺-invariant norm on 𝑉𝜌 . Since we choose all representation matrices to be orthogonal in practice, we can simply use the 𝐿 2 norm.
4.6
The Embedding Theorem
In this section, we will show that, roughly speaking, “in a fixed network architecture, more equivariance means less expressivity”. In slightly more technical terms, the map that takes a discrete group 𝐺 ≤ O(2) and sends it to the h set of functions expressible byi a 𝐺-equivariant transformer of fixed depth is inclusion-reversing. Let F𝐺 𝛿, ℎ; (𝐶𝜌 )𝜌∈ 𝐺b , (𝑛 𝛼 ) 𝛼∈Sub(𝐺)/∼ denote the set of functions E-ViT : R
H0
3
⊗R →R
H
Ê 𝐶𝜌 ⊗ R ⊗ 𝑉𝜌 𝜌∈ 𝐺b
(16)
expressible as a composition Block 𝛿 ◦ · · · ◦ Block1 ◦ PosEnc ◦ PE, where PE and PosEnc are the patch embedding and positional encoding layers described in Sec. 4.2, and each Block 𝑘 is a transformer block with ℎ-head coupled self-attention (see Sec. 4.4) whose MLP involves 𝑛 𝛼 copies of the 𝛼-th homogeneous space (see Sec. 4.3). Theorem 2. Let 𝐻 ≤ 𝐺 be a subgroup. Fix an 𝐻-equivariant linear isometric bijection Ê Ê ∼ 𝐶𝜌 Res𝐺 : R ⊗ 𝑉 − → R𝐷 𝜎 ⊗ 𝑊 𝜎 . 𝜌 𝐻 |
(17)
b 𝜎∈𝐻
b 𝜌∈ 𝐺
{z
=:𝑉
}
|
{z
=:𝑊
}
Note that the multiplicities 𝐷 𝜎 are uniquely determined. Then h i (Id ⊗ Res𝐺 ) ◦ F 𝛿, ℎ; (𝐶 ) , (𝑛 ) 𝐺 𝜌 𝛼 𝛼∈Sub(𝐺)/∼ b 𝐻 𝜌∈ 𝐺 ⊂ F𝐻 𝛿, ℎ; (𝐷 𝜎 ) 𝜎 ∈ 𝐻b , (𝑚 𝛽 ) 𝛽 ∈Sub(𝐻 )/∼ ,
(18)
where the numbers of homogeneous space copies 𝑚 𝛽 are determined as follows: the 𝐺-space 𝑋 :=
Ä
𝑛𝛼 Ä
𝑋𝛼
(19)
𝛼∈Sub(𝐺)/∼ 𝑠=1
decomposes into a disjoint union of 𝐻-orbits. 𝑚 𝛽 is then the number of times the 𝛽-th homogeneous 𝐻-space appears in 𝑋. Moreover, if dim Hom 𝐻 (R𝑈 , 𝑊) > dim Hom𝐺 (R𝑈 , 𝑉), then the inclusion is strict. 9
The proof of Theorem 2 is in the Supplementary Material. Apart from providing a rigorous formulation of the bias-expressivity tradeoff in the context of equivariant ViTs, Theorem 2 is also of practical value: it allows one to view a 𝐺-equivariant ViT as an 𝐻-equivariant ViT, opening the door to gradual symmetry-breaking, for example, by enforcing 𝐺-equivariance in early training epochs, and then procedurally relaxing to smaller subgroups at later stages.
5
Experiments
In this section, we carry out experiments with equivariant ViTs for subgroups of 𝐷 4 and 𝐷 6 on the PatternNet dataset. For the models equivariant to 𝐷 6 and its subgroups, we consider the hexagonal lattice structure for the underlying 𝐺-patchified grid (see Fig. 2). The PatternNet data set [44] consists of 30,400 aerial images divided into 38 classes, each of which has 800 images. We split the images within each class into 80% training and 20% validation data. All models are trained using AdamW with (unweighted) cross entropy as loss function, and evaluated using mean accuracy as the primary metric. We refer the reader to the Supplementary Material for further experiment details.
5.1
Does equivariance matter?
Our first experiment aims to test the hypothesis that equivariance (more precisely invariance since the task is classification) implies better sample efficiency. This is intuitively plausible: for a 𝐺-invariant model, a single labeled example (𝑥, 𝑦) is equivalent to |𝐺 | examples, namely {(𝑔𝑥, 𝑦)|𝑔 ∈ 𝐺}. In other words, the model has built-in data augmentation. To this end, we train 𝐺-invariant classifiers on 10%, 40%, and 100% of the training data for 𝐺 = 𝐷 4 , 𝐶4 , 𝐷 2 , 𝐶1 on the square grid and 𝐺 = 𝐷 6 , 𝐶6 , 𝐷 3 , 𝐶1 on the hexagonal grid. The same set of data is always used at each sampling fraction. We adjust the maximum training epoch and early stopping patience to roughly compensate for the reduced number of training examples. Coupled equivariant attention with ℎ = 3 attention heads is employed in this experiment. For all models, we take the feature space 𝑉 to be 3𝑞 copies of the regular representation with 𝑞 = 1, 2, · · · , so that each head in the attention layer has 𝑞 copies of the regular representation. The results are shown in Fig. 4. Note that there are significantly more runs with sample ratio = 0.1 due to high variance.
5.2
Comparison of attention mechanisms
Here, we compare the two attention mechanisms detailed in Section 4.4: irrep-wise and coupled. For this experiment, we take 𝐺 = 𝐶4 with 4𝑞 copies of the regular representation, with 𝑞 = 1, 2, 3, · · · , and we use four attention heads in total for both attention types (see Table 1). The attention heads are configured such that there are always 4𝑞 features per token per head (note that E1 is two-dimensional). Two sets of experiments are carried out: one with only 10% of the training data, and one with the complete training data. The results are shown in Fig. 5.
channel dimension irreps in each head
coupled irrepwise 𝐶A = 𝐶B = 𝐶E1 = 4𝑞 A⊕4𝑞 B⊕4𝑞 A⊕𝑞 ⊕ B⊕𝑞 ⊕ E1⊕𝑞 (all heads) E⊕2𝑞 E⊕2𝑞
Table 1: Chosen configuration of attention heads for the two attention mechanisms for the 𝐶4 -equivariant transformer.
10
validation accuracy (%) validation accuracy (%)
sample ratio = 0.1
100
sample ratio = 0.4
100
90
90
90
80
80
80
70
70
70 60
D4 D2
4
C4 C1
5
log10 (#params) sample ratio = 0.1
100
60
4
60
5 sample ratio = 0.4
100 90
90
80
80
80
70
70
60
D6 D3
4
5
log10 (#params)
C6 C1
60
4
60
5
4
5 sample ratio = 1.0
100
90
70
sample ratio = 1.0
100
4
5
validation accuracy (%)
Figure 4: Top: equivariant models on the square grid. Bottom: equivariant models on the hexagonal grid.
sample ratio = 0.1
100 90
90
80
80 irrep-wise coupled
70 4
sample ratio = 1.0
100
70
5
4
log10 (#params)
5
Figure 5: Comparison of irrep-wise and coupled attention with 10% (left) and 100% (right) of the training data.
11
5.3
Comparison of homogeneous space combinations in MLP
validation accuracy (%)
Finally, we study the performance differences resulting from different choices of the 𝐺-set 𝑋 in our construction of the nonlinear layer (see Section 4.3). We take 𝐺 = 𝐷 4 with 3𝑞 copies of the regular representation as the feature space 𝑉, and consider here three families of 𝐺-sets, each formed by 𝑛 copies of 𝐷 4 (regular action), 𝑛 copies of 𝐷 4 plus 8 points (trivial action), and 𝑛 copies of 𝐷 4 /⟨𝑟⟩ ⊔ 𝐷 4 /⟨𝑡⟩ ⊔ 𝐷 4 /⟨𝑡𝑟, 𝑟 2 ⟩ (see the Supplementary Material for a list of homogeneous spaces of 𝐷 4 ). These are chosen so that all irreps appear at least once in 𝑋, preventing information loss in the MLP layer. Note that the second and third choices contain 9 higher ratios of A1 features in the hidden layer ( 16 and 38 respectively, as opposed to 18 ). 100
ml _ratio = 0.333
100
98 96 D4 D4 ⊔ (8 ∙ ) (D4/⟨r⟩) ⊔(D4/⟨t⟩)
94 92 90 4.0
⊔(D4/⟨tr⟩ r 2⟩ 4.5
5.0
5.5
ml _ratio = 1
100
98
98
96
96
94
94
92
92
90 4.0
4.5
5.0
5.5
90 4.0
ml _ratio = 2
4.5
5.0
5.5
log10 (#params) Figure 6: Performance comparison of 𝐷 4 -equivariant ViTs with different combinations of homogeneous spaces in the MLP layers.
Instead of varying the number of homogeneous spaces 𝑛 independently of 𝑞, we fix three “MLP ratios”, 𝑛 defined as 3𝑞 . The findings are summarized in Fig. 6.
5.4
Discussion
Our first experiment (Section 5.1) confirms the intuition that the significance of equivariance is magnified in the low-data regime: while the nonequivariant ViT performs significantly worse than the equivariant counterparts when trained on 10% of the data, all ViT variants are practically indistinguishable when all of the training data is used. In fact, the 𝐷 2 -equivariant model seems to perform slightly worse than the nonequivariant one. However, it is not clear whether more equivariance always entails better performance: at 10% sample ratio, the 𝐷 3 -equivariant model achieves slightly better accuracies than the 𝐷 6 -equivariant counterpart at fixed parameter counts. The last two experiments show that the choice of attention type (Section 5.2) and of 𝐺-set 𝑋 for the MLP layers (Section 5.3) do not affect performance in a drastic way. We can only conclude that marginal accuracy gain is achieved with higher fractions of A1 features in the MLP hidden layer (Fig. 6) when the MLP ratio is low, and with irrep-wise attention as opposed to coupled attention (see Fig. 5). Nevertheless, one cannot rule out the possibility that perturbations in other hyperparameters might change this conclusion. One potential explanation for the observed worse performance with larger symmetry groups is that for classification tasks, especially for “low-frequency” images, in the sense of Fourier transforms, like those in PatternNet (e.g. beaches, baseball fields, runways), non-trivial irrep features are not as crucial. This leads us to the following conjecture: Conjecture. A 𝐺-invariant classification model with a 𝐺-equivariant ViT backbone mostly uses features from the trivial representation. If this is true, then it might be more difficult for a 𝐷 4 -equivariant model to learn invariant features compared to a 𝐷 2 -equivariant one if the respective regular representations are used as feature spaces in both models, as the fraction of invariant features in 𝑉 is 1/|𝐺 |. This could also explain the apparent advantage of irrep-wise attention over coupled attention (Section 5.2): if little useful information is stored in nontrivial irreps, mixing these irreps with the trivial one in the same attention head could result in noisy attention scores. 12
6
Limitations and Future Work
Due to large hyperparameter search space and limited computational resources, all experimental results presented here involve very small models (≲ 0.5M parameters), trained on a relatively small dataset (PatternNet). It is not clear whether similar results hold at scale. Indeed, as our experiments suggest, equivariance could be most important when data are scarce. However, in Ref. [7], the argument is made that equivariance can also be practical in large dataset regimes, due to the increased sparsification of the linear layers when increasing the size of the group 𝐺 while keeping the total channel dimension fixed. At the model sizes used in this paper, the computational benefit is not visible, but we hope that our implementations will be useful to further characterize the scaling properties of equivariant ViTs in future work. In addition to scaling up, there are multiple orthogonal avenues for future exploration. First, it is in principle possible to construct ViTs equivariant to groups 𝐺 ≤ O(2) that are not subgroups of 𝐷 4 or 𝐷 6 . However, their construction will inevitably involve more irregular grids (c.f. Definition 3), and it would be interesting to understand the trade-off between exact discrete equivariance and approximate continuous equivariance. Second, as already remarked after Theorem 2, one can break the symmetry, either “statically”, by concatenating transformer blocks in a way that later blocks have smaller symmetry groups , or “dynamically”, by imposing a large symmetry group at early training stages, and then gradually relaxing the symmetry group to smaller ones. We expect this kind of architectures or training schedule to be advantageous for datasets that do not respect rotational symmetries exactly. Finally, for a fixed group 𝐺 ≤ O(2), each choice of feature space 𝑉 must be independently tested in the current formulation. It would be of great benefit for the practical use of the networks presented in this paper, and for group equivariant neural networks in general, to find a principled way to pick 𝑉 for a given task, or to optimize the choice as part of training a network. Code Availability. The code used in this study will be made publicly available in a permanent online repository upon publication of the article.
References [1] Brandon Anderson, Truong-Son Hy, and Risi Kondor. Cormorant: Covariant Molecular Neural Networks, November 2019. URL http://arxiv.org/abs/1906.04015. arXiv:1906.04015 [physics]. [2] Serge Assaad, Carlton Downey, Rami Al-Rfou, Nigamaa Nayakanti, and Ben Sapp. VN-Transformer: Rotation-Equivariant Attention for Vector Neurons, January 2023. URL http://arxiv.org/abs/ 2206.04176. arXiv:2206.04176 [cs]. [3] Ilyes Batatia, Dávid Péter Kovács, Gregor N. C. Simm, Christoph Ortner, and Gábor Csányi. MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields, January 2023. URL http://arxiv.org/abs/2206.07697. arXiv:2206.07697 [stat]. [4] Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P. Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature Communications, 13(1):2453, May 2022. ISSN 2041-1723. doi: 10.1038/s41467-022-29939-5. URL https://www.nature.com/articles/ s41467-022-29939-5. [5] Erik J. Bekkers, Maxime W. Lafarge, Mitko Veta, Koen AJ Eppenhof, Josien PW Pluim, and Remco Duits. Roto-Translation Covariant Convolutional Networks for Medical Image Analysis, April 2018. URL https://arxiv.org/abs/1804.03393v3. 13
[6] Michael M. Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges, May 2021. URL http://arxiv.org/abs/2104.13478. arXiv:2104.13478 [cs]. [7] Georg Bökman, David Nordström, and Fredrik Kahl. Flopping for FLOPs: Leveraging equivariance for computational efficiency, June 2025. URL http://arxiv.org/abs/2502.05169. arXiv:2502.05169 [cs]. [8] Evangelos Chatzipantazis, Stefanos Pertigkiozoglou, Edgar Dobriban, and Kostas Daniilidis. SE(3)Equivariant Attention Networks for Shape Reconstruction in Function Space, February 2023. URL http://arxiv.org/abs/2204.02394. arXiv:2204.02394 [cs]. [9] Taco Cohen, Mario Geiger, and Maurice Weiler. A General Theory of Equivariant CNNs on Homogeneous Spaces, January 2020. URL http://arxiv.org/abs/1811.02017. arXiv:1811.02017 [cs]. [10] Taco S. Cohen and Max Welling. Group Equivariant Convolutional Networks, February 2016. URL https://arxiv.org/abs/1602.07576v3. [11] Taco S. Cohen and Max Welling. Steerable CNNs, December 2016. URL http://arxiv.org/abs/ 1612.08498. arXiv:1612.08498 [cs]. [12] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale, June 2021. URL http://arxiv.org/abs/2010.11929. arXiv:2010.11929 [cs]. [13] Floor Eijkelboom, Rob Hesselink, and Erik Bekkers. E(n) Equivariant Message Passing Simplicial Networks, October 2023. URL http://arxiv.org/abs/2305.07100. arXiv:2305.07100 [cs]. [14] Fabian B. Fuchs, Daniel E. Worrall, Volker Fischer, and Max Welling. SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks, November 2020. URL http://arxiv.org/abs/ 2006.10503. arXiv:2006.10503 [cs]. [15] Jan E. Gerken, Jimmy Aronsson, Oscar Carlsson, Hampus Linander, Fredrik Ohlsson, Christoffer Petersson, and Daniel Persson. Geometric deep learning and equivariant neural networks. Artificial Intelligence Review, 56(12):14605–14662, December 2023. ISSN 1573-7462. doi: 10.1007/s10462-023-10502-7. URL https://doi.org/10.1007/s10462-023-10502-7. [16] Charles Godfrey, Davis Brown, Tegan Emerson, and Henry Kvinge. On the symmetries of deep learning models and their internal representations. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 11893–11905. Curran Associates, Inc., 2022. URL https://proceedings.neurips.cc/paper_files/paper/ 2022/file/4df3510ad02a86d69dc32388d91606f8-Paper-Conference.pdf. [17] Lingshen He, Yuxuan Chen, zhengyang shen, Yiming Dong, Yisen Wang, and Zhouchen Lin. Efficient Equivariant Network. In Advances in Neural Information Processing Systems, volume 34, pages 5290–5302. Curran Associates, Inc., 2021. URL https://proceedings.neurips.cc/paper/ 2021/hash/2a79ea27c279e471f4d180b08d62b00a-Abstract.html. [18] Emiel Hoogeboom, Jorn W. T. Peters, Taco S. Cohen, and Max Welling. HexaConv, March 2018. URL http://arxiv.org/abs/1803.02108. arXiv:1803.02108 [cs].
14
[19] Mohammad Mohaiminul Islam, Rishabh Anand, David R. Wessels, Friso de Kruiff, Thijs P. Kuipers, Rex Ying, Clara I. Sánchez, Sharvaree Vadgama, Georg Bökman, and Erik J. Bekkers. Platonic Transformers: A Solid Choice For Equivariance, October 2025. URL http://arxiv.org/abs/2510.03511. arXiv:2510.03511 [cs]. [20] Max Jaderberg, Karen Simonyan, Andrew Zisserman, and Koray Kavukcuoglu. Spatial Transformer Networks, February 2016. URL http://arxiv.org/abs/1506.02025. arXiv:1506.02025 [cs]. [21] Sékou-Oumar Kaba, Arnab Kumar Mondal, Yan Zhang, Yoshua Bengio, and Siamak Ravanbakhsh. Equivariance with Learned Canonicalization Functions, July 2023. URL http://arxiv.org/abs/ 2211.06489. arXiv:2211.06489 [cs]. [22] Risi Kondor and Shubhendu Trivedi. On the Generalization of Equivariance and Convolution in Neural Networks to the Action of Compact Groups, November 2018. URL http://arxiv.org/abs/1802. 03690. arXiv:1802.03690 [stat]. [23] Risi Kondor, Zhen Lin, and Shubhendu Trivedi. Clebsch-Gordan Nets: a Fully Fourier Space Spherical Convolutional Neural Network, November 2018. URL http://arxiv.org/abs/1806. 09231. arXiv:1806.09231 [stat]. [24] Soumyabrata Kundu and Risi Kondor. A Geometric Approach to Steerable Convolutions, October 2025. URL http://arxiv.org/abs/2510.18813. arXiv:2510.18813 [cs]. [25] Soumyabrata Kundu and Risi Kondor. Steerable Transformers for Volumetric Data, October 2025. URL http://arxiv.org/abs/2405.15932. arXiv:2405.15932 [cs]. [26] Haggai Maron, Heli Ben-Hamu, Nadav Shamir, and Yaron Lipman. Invariant and Equivariant Graph Networks, April 2019. URL http://arxiv.org/abs/1812.09902. arXiv:1812.09902 [cs]. [27] David Nordström, Johan Edstedt, Fredrik Kahl, and Georg Bökman. Octic Vision Transformers: Quicker ViTs Through Equivariance, September 2025. URL http://arxiv.org/abs/2505.15441. arXiv:2505.15441 [cs]. [28] Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Marc Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, Mahmoud Assran, Nicolas Ballas, Wojciech Galuba, Russell Howes, Po-Yao Huang, Shang-Wen Li, Ishan Misra, Michael Rabbat, Vasu Sharma, Gabriel Synnaeve, Hu Xu, Hervé Jegou, Julien Mairal, Patrick Labatut, Armand Joulin, and Piotr Bojanowski. DINOv2: Learning Robust Visual Features without Supervision, February 2024. URL http://arxiv.org/abs/2304.07193. arXiv:2304.07193 [cs]. [29] Marco Pacini, Xiaowen Dong, Bruno Lepri, and Gabriele Santin. A characterization theorem for equivariant networks with point-wise activations. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.net/forum?id=79FVDdfoSR. [30] Marco Pacini, Gabriele Santin, Bruno Lepri, and Shubhendu Trivedi. On universality classes of equivariant networks. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. URL https://openreview.net/forum?id=V4YAS7NLXi. [31] Siamak Ravanbakhsh. Universal Equivariant Multilayer Perceptrons. In Proceedings of the 37th International Conference on Machine Learning, pages 7996–8006. PMLR, November 2020. URL https://proceedings.mlr.press/v119/ravanbakhsh20a.html.
15
[32] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) Equivariant Graph Neural Networks, February 2022. URL http://arxiv.org/abs/2102.09844. arXiv:2102.09844 [cs]. [33] Jean-Pierre Serre. Linear representations of finite groups. Graduate Texts in Mathematics. Springer, New York, NY, September 1977. [34] J. Shawe-Taylor. Building symmetries into feedforward networks. In 1989 First IEE International Conference on Artificial Neural Networks, (Conf. Publ. No. 313), pages 158–162, October 1989. URL https://ieeexplore.ieee.org/document/51951. [35] Kai Sheng Tai, Peter Bailis, and Gregory Valiant. Equivariant Transformer Networks, May 2019. URL http://arxiv.org/abs/1901.11399. arXiv:1901.11399 [cs]. [36] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention Is All You Need, December 2017. URL http: //arxiv.org/abs/1706.03762. arXiv:1706.03762 [cs]. [37] Bastiaan S. Veeling, Jasper Linmans, Jim Winkens, Taco Cohen, and Max Welling. Rotation Equivariant CNNs for Digital Pathology, June 2018. URL http://arxiv.org/abs/1806.03962. arXiv:1806.03962 [cs]. [38] Jianyuan Wang, Minghao Chen, Nikita Karaev, Andrea Vedaldi, Christian Rupprecht, and David Novotny. Vggt: Visual geometry grounded transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 5294–5306, June 2025. [39] Maurice Weiler and Gabriele Cesa. General $E(2)$-Equivariant Steerable CNNs, April 2021. URL http://arxiv.org/abs/1911.08251. arXiv:1911.08251 [cs]. [40] Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma, and Taco Cohen. 3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data, October 2018. URL http: //arxiv.org/abs/1807.02547. arXiv:1807.02547 [cs]. [41] Jeffrey Wood and John Shawe-Taylor. Representation theory and invariant neural networks. Discrete Applied Mathematics, 69(1):33–60, 1996. ISSN 0166-218X. doi: https://doi.org/10. 1016/0166-218X(95)00075-3. URL https://www.sciencedirect.com/science/article/pii/ 0166218X95000753. [42] Renjun Xu, Kaifan Yang, Ke Liu, and Fengxiang He. $E(2)$-Equivariant Vision Transformer, July 2023. URL http://arxiv.org/abs/2306.06722. arXiv:2306.06722 [cs]. [43] Xuan Zhang, Limei Wang, Jacob Helwig, Youzhi Luo, Cong Fu, Yaochen Xie, Meng Liu, Yuchao Lin, Zhao Xu, Keqiang Yan, Keir Adams, Maurice Weiler, Xiner Li, Tianfan Fu, Yucheng Wang, Alex Strasser, Haiyang Yu, YuQing Xie, Xiang Fu, Shenglong Xu, Yi Liu, Yuanqi Du, Alexandra Saxton, Hongyi Ling, Hannah Lawrence, Hannes Stärk, Shurui Gui, Carl Edwards, Nicholas Gao, Adriana Ladera, Tailin Wu, Elyssa F. Hofgard, Aria Mansouri Tehrani, Rui Wang, Ameya Daigavane, Montgomery Bohde, Jerry Kurtin, Qian Huang, Tuong Phung, Minkai Xu, Chaitanya K. Joshi, Simon V. Mathis, Kamyar Azizzadenesheli, Ada Fang, Alán Aspuru-Guzik, Erik Bekkers, Michael Bronstein, Marinka Zitnik, Anima Anandkumar, Stefano Ermon, Pietro Liò, Rose Yu, Stephan Günnemann, Jure Leskovec, Heng Ji, Jimeng Sun, Regina Barzilay, Tommi Jaakkola, Connor W. Coley, Xiaoning Qian, Xiaofeng Qian, Tess Smidt, and Shuiwang Ji. Artificial Intelligence for Science in Quantum, Atomistic, and Continuum Systems. Foundations and Trends® in Machine Learning, 18(4):385–912, 2025. ISSN 1935-8237, 1935-8245. doi: 10.1561/2200000115. URL http://arxiv.org/abs/2307.08423. arXiv:2307.08423 [cs]. 16
[44] Weixun Zhou, Shawn Newsam, Congmin Li, and Zhenfeng Shao. PatternNet: A Benchmark Dataset for Performance Evaluation of Remote Sensing Image Retrieval. ISPRS Journal of Photogrammetry and Remote Sensing, 145:197–209, November 2018. ISSN 09242716. doi: 10.1016/j.isprsjprs.2018.01.004. URL http://arxiv.org/abs/1706.03424. arXiv:1706.03424 [cs].
Discrete subgroups of O(2)
A
Since our principal focus is vision transformers equivariant to discrete subgroups of O(2), it is natural to first discuss the properties of these subgroups. Due to the compactness of O(2), any discrete subgroup is finite, and is isomorphic to either a cyclic group 𝐶𝑛 or a dihedral group 𝐷 𝑛 , with the former generated by rotation by 2𝜋/𝑛, and the latter by rotation by 2𝜋/𝑛 and a reflection. In the following, we will fix the notation used to denote elements of 𝐶𝑛 and 𝐷 𝑛 and summarize the basic results on the representation theory over R of these groups. In general, we adopt the point of view that the groups 𝐶𝑛 and 𝐷 𝑛 exist as abstract groups themselves, independently of the natural injective homomorphisms 𝐶𝑛 ↩→ O(2), 𝐷 𝑛 ↩→ O(2). We will always use 𝑒 to denote the identity element of a group. For notational simplicity, we will write cos 𝜃 − sin 𝜃 R(𝜃) := (20) sin 𝜃 cos 𝜃 for the rotation matrix by 𝜃.
A.1
Cyclic Groups
For a positive integer 𝑛, the cyclic group 𝐶𝑛 is generated by the symbol 𝑟, subject to the relation 𝑟 𝑛 = 𝑒. It is of order (cardinality) 𝑛, and is isomorphic to Z/𝑛Z. Proposition 1. (i) If 𝑛 is even, the group 𝐶𝑛 has 𝑛2 + 1 real irreps, labeled by A, B, E1 , E2 , · · · , E 𝑛2 −1 , where 2𝜋𝑘 𝜌A (𝑟) = 1 𝜌B (𝑟) = −1 𝜌E𝑘 (𝑟) = R . 𝑛
(21)
(ii) If 𝑛 is odd, the group 𝐶𝑛 has 𝑛+1 , and the representation 2 real irreps, labeled by A, E1 , E2 , · · · , E 𝑛−1 2 matrices are the same as the ones given in Eq. (21). (iii) A real irrep of a cyclic group is of real type if it is one-dimensional (A 𝑘 or B 𝑘 ), otherwise it is of complex type (E 𝑘 ).
A.2
Dihedral Groups
The dihedral group 𝐷 𝑛 is generated by two symbols, 𝑡 and 𝑟, subject to the relations 𝑟 𝑛 = 𝑒 and 𝑡𝑟 = 𝑟 −1 𝑡. Proposition 2. (i) If 𝑛 is even, the group 𝐷 𝑛 has 𝑛2 + 3 real irreps, labeled by A1 , A2 , B1 , B2 , E1 , E2 , · · · , E 𝑛2 −1 , where 𝜌A𝑘 (𝑟) = 1 𝜌A𝑘 (𝑡) = (−1) 𝑘+1 𝜌B𝑘 (𝑟) = −1 𝜌B𝑘 (𝑡) = (−1) 𝑘+1 (22) 2𝜋𝑘 1 0 𝜌E𝑘 (𝑟) = R 𝜌E𝑘 (𝑡) = . 0 −1 𝑛 (ii) If 𝑛 is odd, the group 𝐶𝑛 has 𝑛+3 , and the representation 2 real irreps, labeled by A1 , A2 , E1 , E2 , · · · , E 𝑛−1 2 matrices are the same as the ones given in Eq. (22). (iii) All real irreps of 𝐷 𝑛 are of real type for any 𝑛. 17
B
Irrep Multiplicities of Homogeneous Spaces of 𝐷 4 and 𝐷 6
Here, we provide the irrep multiplicities of 𝐶 (𝑋 𝛼 , R) for all homogeneous spaces 𝑋 𝛼 of 𝐺 = 𝐷 4 and 𝐷 6 , which are used in the MLP layer of each transformer block. For each homogeneous space, we indicate the stabilizer subgroup 𝐻 𝛼 as well as the number of elements |𝑋 𝛼 |. 𝐻𝛼 {𝑒} ⟨𝑟 2 ⟩ ⟨𝑟⟩ ⟨𝑡⟩ ⟨𝑡𝑟⟩ ⟨𝑡, 𝑟 2 ⟩ ⟨𝑡𝑟, 𝑟 2 ⟩ ⟨𝑡, 𝑟⟩ = 𝐷 4
ΓA𝛼1 1 1 1 1 1 1 1 1
ΓA𝛼2 1 1 1 0 0 0 0 0
ΓB𝛼1 1 1 0 1 0 1 0 0
ΓB𝛼2 1 1 0 0 1 0 1 0
ΓE𝛼1 2 0 0 1 1 0 0 0
|𝑋 𝛼 | 8 4 2 4 4 2 2 1
Table 2: Homogeneous spaces of 𝐷 4
𝐻𝛼 {𝑒} ⟨𝑟 3 ⟩ ⟨𝑟 2 ⟩ ⟨𝑟⟩ ⟨𝑡⟩ ⟨𝑡𝑟⟩ ⟨𝑡, 𝑟 3 ⟩ ⟨𝑡, 𝑟 2 ⟩ ⟨𝑡𝑟, 𝑟 2 ⟩ ⟨𝑡, 𝑟⟩ = 𝐷 6
ΓA𝛼1 1 1 1 1 1 1 1 1 1 1
ΓA𝛼2 1 1 1 1 0 0 0 0 0 0
ΓB𝛼1 1 0 1 0 1 0 0 1 0 0
ΓB𝛼2 1 0 1 0 0 1 0 0 1 0
ΓE𝛼1 2 0 0 0 1 1 0 0 0 0
ΓE𝛼2 2 2 0 0 1 1 1 0 0 0
|𝑋 𝛼 | 12 6 4 2 6 6 3 2 2 1
Table 3: Homogeneous spaces of 𝐷 6
C
Proof of Lemma 1
Let 𝜎 : R → R be some function satisfying 𝜎(0) = 0. We will use the same symbol to denote the entrywise application R𝑛 → R𝑛 of 𝜎. Let (𝑒 1 , 𝑒 2 , · · · , 𝑒 𝑛 ) be the standard basis for R𝑛 , with respect to which we define the representation matrices 𝐷 (𝑔): 𝑛 ∑︁ 𝑔𝑒 𝑖 = 𝐷 𝑗𝑖 (𝑔)𝑒 𝑗 (23) 𝑗=1
Equivariance at 𝑒 𝑖
∈ R𝑛 reads 𝑔𝜎(𝑒 𝑖 ) = 𝜎(𝑔𝑒 𝑖 ).
(24)
That is, 𝑛 𝑛 ©∑︁ ª ∑︁ 𝐷 𝑗𝑖 (𝑔)𝑒 𝑗 = 𝜎 𝐷 𝑗𝑖 (𝑔)𝑒 𝑗 ® = 𝜎 𝐷 𝑗𝑖 (𝑔) 𝑒 𝑗 . 𝑗=1 « 𝑗=1 ¬ 𝑗=1 Comparing the coefficient of 𝑒 𝑗 gives
𝜎(1)
𝑛 ∑︁
𝜎(1)𝐷 𝑗𝑖 (𝑔) = 𝜎(𝐷 𝑗𝑖 (𝑔)). 18
(25)
(26)
If 𝐷 𝑗𝑖 (𝑔) ∉ {0, 1}, then we can find a function 𝜎 with 𝜎(0) = 0 that violates Eq. (26). Hence, it must be that 𝐷 𝑗𝑖 (𝑔) ∈ {0, 1}. Note that this holds for any 𝑔 ∈ 𝐺. Now, consider the equation 𝐷 (𝑔)𝐷 (𝑔 −1 ) = 1𝑛 , (27) where 1𝑛 is the 𝑛 × 𝑛 identity matrix. We will now show by contradiction that both 𝐷 (𝑔) and 𝐷 (𝑔 −1 ) are permutation matrices. Suppose the 𝑗-th row of 𝐷 (𝑔) has at least two 1’s. Without loss of generality, we may assume 1 1 ★ ··· ★ © ª ★ ★ ★ · · · ★® ®, 𝐷 (𝑔) = . . . . (28) . . ... ®® .. .. .. «★ ★ ★ · · · ★¬
where each ★ can be either 0 or 1. Then, the first two entries of every column of 𝐷 (𝑔 −1 ) except for the first one must be zero, since otherwise 𝐷 (𝑔)𝐷 (𝑔 −1 ) = 1 would not hold. That is, 𝑎1 0 · · · © 𝑎2 0 · · · 𝐷 (𝑔 −1 ) = 𝑎 2 ★ · · · .. .. . . . . . «𝑎 𝑛 ★ · · ·
0 ª 0® ® ★® . .. ®® .®
(29)
★¬
Since 𝐷 (𝑔 −1 ) is invertible, we have 𝑎 1 = 𝑎 2 = 1 (otherwise 𝐷 (𝑔 −1 ) would not have full rank). But then (𝐷 (𝑔)𝐷 (𝑔 −1 ))11 ≥ 2, contradicting 𝐷 (𝑔)𝐷 (𝑔 −1 ) = 1𝑛 . □
D
Proof of Theorem 1
The claim is trivial for 𝐿 = 1, so we assume 𝐿 > 1. If 𝑤 𝑣 = 0, then attn sends all token sequences to (𝑏 𝑣 , · · · 𝑏 𝑣 ). In this case, attn is equivariant if and only if 𝑏 𝑣 ∈ 𝑉 is invariant, and the claim follows easily. In the following, we assume 𝑤 𝑣 ≠ 0. Let 𝜙𝑞 , 𝜙 𝑘 , 𝜙 𝑣 : 𝑉 → 𝑉 be the query, key, and value maps, which are affine by assumption. We will write 𝜙𝑞 (𝑧) = 𝑤 𝑞 𝑧 + 𝛽𝑞 etc. Let 𝑥 = (𝑥 1 , · · · , 𝑥 𝐿 ) ∈ R 𝐿 ⊗ 𝑉 be a token sequence. The attention layer acts according to Í𝐿 ⟨ 𝜙𝑞 ( 𝑥𝑎 ) , 𝜙 𝑘 ( 𝑥𝑏 ) ⟩ 𝜙 (𝑥 ) 𝑣 𝑏 𝑏=1 𝑒 attn(𝑥) 𝑎 = Í𝐿 ⟨ 𝜙 ( 𝑥 ) , 𝜙 ( 𝑥 )⟩ 𝑞 𝑎 𝑘 𝑏 𝑏=1 𝑒 Í𝐿 𝑡 ⟨𝑤𝑞 𝑥𝑎 ,𝑤𝑘 𝑥𝑏 ⟩+⟨𝑤𝑘 𝛽𝑞 , 𝑥𝑏 ⟩+⟨ 𝑥𝑎 ,𝑤𝑞𝑡 𝛽𝑘 ⟩+⟨𝛽𝑞 ,𝛽𝑘 ⟩ 𝜙 𝑣 (𝑥 𝑏 ) 𝑏=1 𝑒 (30) = 𝑡 𝑡 𝑒 ⟨𝑤𝑞 𝑥𝑎 ,𝑤𝑘 𝑥𝑏 ⟩+⟨𝑤𝑘 𝛽𝑞 , 𝑥𝑏 ⟩+⟨ 𝑥𝑎 ,𝑤𝑞 𝛽𝑘 ⟩+⟨𝛽𝑞 ,𝛽𝑘 ⟩ Í𝐿 ⟨ 𝑥𝑎 , 𝑀 𝑥𝑏 ⟩+⟨𝑤𝑘𝑡 𝛽𝑞 , 𝑥𝑏 ⟩ 𝜙 𝑣 (𝑥 𝑏 ) 𝑏=1 𝑒 , = 𝑡 𝑒 ⟨ 𝑥𝑎 , 𝑀 𝑥𝑏 ⟩+⟨𝑤𝑘 𝛽𝑞 , 𝑥𝑏 ⟩ where 𝑀 := 𝑤 𝑡𝑞 𝑤 𝑘 . 𝐺-equivariance reads Í𝐿 ⟨ 𝑥𝑎 , 𝑀 𝑥𝑏 ⟩+⟨𝑤𝑘𝑡 𝛽𝑞 , 𝑥𝑏 ⟩ ⟨ 𝑥𝑎 ,𝑔 −1 𝑀𝑔𝑥𝑏 ⟩+⟨𝑔 −1 𝑤𝑘𝑡 𝛽𝑞 , 𝑥𝑏 ⟩ 𝑔𝜙 𝑣 (𝑥 𝑏 ) 𝜙 𝑣 (𝑔𝑥 𝑏 ) 𝑏=1 𝑒 𝑏=1 𝑒 = Í𝐿 Í𝐿 𝑡 −1 −1 ⟨ 𝑥𝑎 , 𝑀 𝑥𝑏 ⟩+⟨𝑤𝑘𝑡 𝛽𝑞 , 𝑥𝑏 ⟩ ⟨ 𝑥𝑎 ,𝑔 𝑀𝑔𝑥𝑏 ⟩+⟨𝑔 𝑤𝑘 𝛽𝑞 , 𝑥𝑏 ⟩ 𝑏=1 𝑒 𝑏=1 𝑒
Í𝐿
(31)
for all 𝑔 ∈ 𝐺. For any vector 𝑧 ∈ 𝑉, the above equation for the token sequence (𝑧, 𝑧, · · · , 𝑧) reads 𝑔𝜙 𝑣 (𝑧) = 𝜙 𝑣 (𝑔𝑧). 19
(32)
That is, 𝜙 𝑣 is 𝐺-equivariant. This in particular implies 𝑔𝛽𝑣 = 𝛽𝑣 for all 𝑔 ∈ 𝐺. Multiplying Eq. (31) by 𝑔 −1 and cancelling out the value bias, we get Í𝐿 ⟨ 𝑥𝑎 , 𝑀 𝑥𝑏 ⟩+⟨𝑤𝑘𝑡 𝛽𝑞 , 𝑥𝑏 ⟩ ⟨ 𝑥𝑎 ,𝑔 −1 𝑀𝑔𝑥𝑏 ⟩+⟨𝑔 −1 𝑤𝑘 𝛽𝑞 , 𝑥𝑏 ⟩ 𝑤 𝑥 𝑤 𝑣 𝑥𝑏 𝑣 𝑏 𝑏=1 𝑒 𝑏=1 𝑒 = Í𝐿 . Í𝐿 −1 𝑀𝑔𝑥 ⟩+⟨𝑔 −1 𝑤 𝛽 , 𝑥 ⟩ ⟨ 𝑥𝑎 , 𝑀 𝑥𝑏 ⟩+⟨𝑤𝑘𝑡 𝛽𝑞 , 𝑥𝑏 ⟩ ⟨ 𝑥 ,𝑔 𝑎 𝑞 𝑏 𝑘 𝑏 𝑏=1 𝑒 𝑏=1 𝑒
Í𝐿
(33)
Let 𝑧, 𝑧 ′ ∈ 𝑉 be arbitrary, and let 𝑥 be the token sequence for which 𝑥 𝑎 = 𝑧, 𝑥 𝑏 = 𝑧 ′ and 𝑥 𝑏′ = 0 for all 𝑏 ′ ∉ {𝑎, 𝑏}. Then ∀𝑧, 𝑧 ′ ∈ 𝑉, 𝑔 ∈ 𝐺 : ′
𝑡
′
′
𝑡
𝑒 ⟨𝑧, 𝑀 𝑧⟩+⟨𝑤𝑘 𝛽𝑞 ,𝑧⟩ 𝑤 𝑣 𝑧 + 𝑒 ⟨𝑧, 𝑀 𝑧 ⟩+⟨𝑤𝑘 𝛽𝑞 ,𝑧 ⟩ 𝑤 𝑣 𝑧 ′ 𝑡
′
(34)
𝐿 − 2 + 𝑒 ⟨𝑧, 𝑀 𝑧⟩+⟨𝑤𝑘 𝛽𝑞 ,𝑧⟩ + 𝑒 ⟨𝑧, 𝑀 𝑧 ⟩+⟨𝑤𝑘 𝛽𝑞 ,𝑧 ⟩ −1 −1 𝑡 −1 ′ −1 𝑡 ′ 𝑒 ⟨𝑧,𝑔 𝑀𝑔𝑧⟩+⟨𝑔 𝑤𝑘 𝛽𝑞 ,𝑧⟩ 𝑤 𝑣 𝑧 + 𝑒 ⟨𝑧,𝑔 𝑀𝑔𝑧 ⟩+⟨𝑔 𝑤𝑘 𝛽𝑞 ,𝑧 ⟩ 𝑤 𝑣 𝑧 ′ . = ′ −1 ′ −1 𝑡 −1 −1 𝑡 𝐿 − 2 + 𝑒 ⟨𝑧,𝑔 𝑀𝑔𝑧⟩+⟨𝑔 𝑤𝑘 𝛽𝑞 ,𝑧⟩ + 𝑒 ⟨𝑧,𝑔 𝑀𝑔𝑧 ⟩+⟨𝑔 𝑤𝑘 𝛽𝑞 ,𝑧 ⟩ 𝑡
Taking 𝑧 ′ = 0 in this equation gives 𝑒 ⟨𝑧, 𝑀 𝑧⟩+⟨𝑤𝑘 𝛽𝑞 ,𝑧⟩
𝑒 ⟨𝑧,𝑔
𝑡
∀𝑧 ∈ 𝑉 \ ker(𝑤 𝑣 ) :
𝐿 − 1 + 𝑒 ⟨𝑧, 𝑀 𝑧⟩+⟨𝑤𝑘 𝛽𝑞 ,𝑧⟩ 𝑡
=
−1 𝑀𝑔𝑧⟩+⟨𝑔 −1 𝑤 𝑡 𝛽 ,𝑧⟩ 𝑘 𝑞
𝐿 − 1 + 𝑒 ⟨𝑧,𝑔
−1 𝑀𝑔𝑧⟩+⟨𝑔 −1 𝑤 𝑡 𝛽 ,𝑧⟩ 𝑘 𝑞
.
(35)
The set 𝑉 \ ker(𝑤 𝑣 ) is open and nonempty (because 𝑤 𝑣 ≠ 0). Since Eq. (35) is analytic in 𝑧 and holds on an 𝑒𝑢 nonempty open set, it holds for all 𝑧 ∈ 𝑉. By the monotonicity of the function 𝑢 ↦→ 𝐿−1+𝑒 𝑢 , we conclude ∀𝑧 ∈ 𝑉, 𝑔 ∈ 𝐺 : ⟨𝑧, 𝑀 𝑧⟩ + ⟨𝑤 𝑡𝑘 𝛽𝑞 , 𝑧⟩ = ⟨𝑧, 𝑔 −1 𝑀𝑔𝑧⟩ + ⟨𝑔 −1 𝑤 𝑡𝑘 𝛽𝑞 , 𝑧⟩ .
(36)
Now, go back to Eq. (34) and again use analyticity and monotonicity to arrive at ∀𝑧, 𝑧 ′ ∈ 𝑉, 𝑔 ∈ 𝐺 : ⟨𝑧, 𝑀 𝑧 ′ ⟩ + ⟨𝑤 𝑡𝑘 𝛽𝑞 , 𝑧 ′ ⟩ = ⟨𝑧, 𝑔 −1 𝑀𝑔𝑧 ′ ⟩ + ⟨𝑔 −1 𝑤 𝑡𝑘 𝛽𝑞 , 𝑧 ′ ⟩ .
(37)
Setting 𝑧 = 0 gives 𝑤 𝑡𝑘 𝛽𝑞 = 𝑔 −1 𝑤 𝑡𝑘 𝛽𝑞 , which in turn implies 𝑀 = 𝑔 −1 𝑀𝑔. We can now take 𝑤˜ 𝑘 = 1, 𝑤˜ 𝑞 = 𝑀, and 𝛽˜𝑞 = 𝑤 𝑡𝑘 𝛽𝑞 , so that 𝑤˜ 𝑞 , 𝑤˜ 𝑘 are 𝐺-equivariant, 𝛽˜𝑞 is 𝐺-invariant, and ( 𝑤˜ 𝑞 , 𝛽˜𝑞 , 𝑤˜ 𝑘 , 0, 𝑤 𝑣 , 𝛽𝑣 ) yields the same self-attention layer. □
E
Proof of Theorem 2
In principle, the theorem could be proved in an abstract and almost trivial way, more or less by noting that any 𝐺-representation is also an 𝐻-representation and any 𝐺-equivariant map is also 𝐻-equivariant. Here, we choose to provide a much more verbose proof, because it has the additional advantage of giving a recipe to map a 𝐺-equivariant ViT to an 𝐻-equivariant one. Step 1: patch embedding and positional encoding. Let 𝑙 ∈ Hom𝐺 (R𝑈 ⊗ R3 , 𝑉) denote the operation of patch embedding followed by the addition of positional encodings. Lemma 2 applied to R𝑈 ⊗ R3 , 𝑉, and 𝑊 with 𝑗 = IdR𝑈 ⊗R3 implies that there exists an 𝑙˜ ∈ Hom 𝐻 (R𝑈 ⊗ R3 , 𝑊) such that the following diagram commutes: 𝑉 𝑙
R𝑈 ⊗ R3
Res𝐺 𝐻 𝑙˜
𝑊 20
(38)
The map 𝑙˜ is then expressible in terms of an 𝐻-equivariant patch embedding layer. H Let 𝑝 ∈ R H ⊗𝑉 be the positional encodings for the 𝐺-equivariant model. Then (Id ⊗ Res𝐺 𝐻 ) ( 𝑝) ∈ R ⊗ 𝑊 is 𝐻-invariant. Step 2: multi-head self-attention. First, note that the number of heads ℎ divides each irrep multiplicity 𝐷 𝜎 (because it is a linear combination of the 𝐶𝜌 with integer coefficients). We will the 𝐻-equivariant model can output the same attention scores. Let 𝜙 : 𝑉 → Rℎ ⊗ 𝑉1 , Éshow that 𝐶 /ℎ 𝜌 with 𝑉1 = ⊗ 𝑉𝜌 denoting the “reshaping” linear isomorphism, which is 𝐺-equivariant. Similarly, bR 𝜌∈ 𝐺 we have an 𝐻-equivariant linear isomorphism 𝜓 : 𝑊 → Rℎ ⊗ 𝑊1 . Note that 𝑉1 and 𝑊1 are isomorphic as 𝐻-representations, so we can find an 𝐻-equivariant isometric linear isomorphism Φ : 𝑉1 → 𝑊1 . Consider the following diagram: 𝑥↦→𝑞⊕𝑘
𝑉
𝜙
R2 ⊗ 𝑉
Res𝐺 𝐻
R2 ⊗ Rℎ ⊗ 𝑉1 Id⊗Id⊗Φ
𝑊
𝑓˜𝑘𝑞
R2 ⊗ 𝑊
𝜓
(39)
R2 ⊗ R ℎ ⊗ 𝑊 1 .
By Lemma 2, there exists an affine 𝐻-equivariant map 𝑓˜ : 𝑊 → R2 ⊗ Rℎ ⊗ 𝑊1 making the diagram commute. Since 𝜓 is an equivariant linear isomorphism, there is an affine map 𝑓˜𝑘𝑞 : 𝑊 → R2 ⊗ 𝑊 such that this diagram commutes. The map 𝑓˜𝑘𝑞 then computes the key and query for the 𝐻-equivariant model. Since Φ is an isometry, it is clear that the following diagram commutes: R H ⊗ R2 ⊗ Rℎ ⊗ 𝑉1 compute raw attention scores
R H ⊗ R H ⊗ Rℎ
Id⊗Id⊗Id⊗Φ
(40)
compute raw attention scores
R H ⊗ R2 ⊗ R ℎ ⊗ 𝑊
1
Finally, consider the following diagram: RH ⊗ 𝑉
Res𝐺 𝐻
Id⊗ 𝑓˜𝑣
Id⊗ ( 𝑥↦→𝑣)
RH ⊗ 𝑉
RH ⊗ 𝑊
Id⊗ 𝜙 Id⊗Id⊗Φ H R ⊗ Rℎ ⊗ 𝑉1 𝐵⊗Id
R H ⊗ Rℎ ⊗ 𝑉1
Id⊗ 𝜓
R H ⊗ Rℎ ⊗ 𝑊
1
𝐵⊗Id Id⊗Id⊗Φ
Id⊗ 𝜙 −1
(41)
R H ⊗ R ℎ ⊗ 𝑊1 Id⊗ 𝜓 −1
RH ⊗ 𝑉
RH ⊗ 𝑊
Id⊗p
RH ⊗ 𝑉
RH ⊗ 𝑊
Id⊗ p̃ Res𝐺 𝐻
R H ⊗ 𝑊,
where we applied Lemma 2 twice to get the maps 𝑓˜𝑣 , p̃ ∈ Aff 𝐻 (𝑊, 𝑊), and 𝐵 ∈ End(R H ⊗ Rℎ ) is the block-diagonal matrix that aggregates the value vectors according to the attention scores. The middle block commutes because 𝐵 and Φ act on different factors of the tensor product. 21
Step 3: MLP. Let
𝑚𝛽 Ä
Ä
𝑌 :=
(42)
𝑌𝛽 ,
𝛽 ∈Sub(𝐻 )/∼ 𝑠=1
where 𝑌𝛽 is the 𝛽-th homogeneous space of 𝐻. Clearly, 𝑋 𝑌 as 𝐻-sets. Any isomorphism induces an ∼ isomorphism of 𝐻-representations 𝐶 (𝑋, R) − → 𝐶 (𝑌 , R) that commutes with maps that acts entrywise. Thus, by Lemma 2, there exist 𝐻-equivariant affine maps 𝑙˜1 , 𝑙˜2 such that the following diagram commutes: 𝑉
𝑙1
𝐶 (𝑋, R)
𝑙˜1
𝑙2
𝐶 (𝑋, R)
∼
Res𝐺 𝐻
𝑊
𝜎
∼
𝐶 (𝑌 , R)
𝜎
𝐶 (𝑌 , R)
𝑉 Res𝐺 𝐻
𝑙˜2
(43)
𝑊
Step 4: residual connections. It remains to show that residual connections do not ruin any of the expressivity proofs above, which is true since the following diagram commutes: RH ⊗ 𝑉
𝑥↦→ 𝑥+ 𝑓 ( 𝑥 ) H R ⊗𝑉
Id⊗Res𝐺 𝐻
Id⊗Res𝐺 𝐻
RH ⊗ 𝑊
(44)
RH ⊗ 𝑊
𝑦↦→ 𝑦+ 𝑓˜ ( 𝑦)
as long as the same diagram without residual connections also commutes. Step 5: strictness of the inclusion. Assume dim Hom 𝐻 (R𝑈 , 𝑊) > dim End𝐺 (R𝑈 , 𝑉). To show that the inclusion is strict, it suffices to find one function in F𝐻 [𝛿, ℎ; (𝐷 𝜎 ) 𝜎 ∈ 𝐺b , (𝑚 𝛽 ) 𝛽 ∈Sub(𝐻 )/∼ ] that is not 𝐺-equivariant. Consider an 𝐻-equivariant Vision transformer E-ViT defined by setting all parameters except the ones in the patch embedding layer to zero. Then, PosEnc = Block 𝑘 = Id for all 𝑘, and E-ViT = PE. Since Hom 𝐻 (R𝑈 , 𝑊) is bigger than Hom𝐺 (R𝑈 , 𝑉), it is possible to choose PE to be 𝐻-equivariant but not 𝐺-equivariant. □ Lemma 2. Let 𝑉, 𝑉 ′ be 𝐺-representations and suppose 𝑊, 𝑊 ′ are 𝐻-representations such that there exist isomorphisms 𝑗 ∈ Hom 𝐻 (𝑉, 𝑊) and 𝑗 ′ ∈ Hom 𝐻 (𝑉 ′ , 𝑊 ′ ). Then for all 𝑓 ∈ Aff 𝐺 (𝑉, 𝑉 ′ ), there exists an 𝑓˜ ∈ Aff 𝐻 (𝑊, 𝑊 ′ ) such that the following diagram commutes: 𝑓
𝑉
𝑗′
𝑗
𝑊
𝑉′
𝑓˜
(45)
𝑊′
Proof. Write 𝑓 (𝑥) = 𝐴𝑥 + 𝑏. Then 𝑏 is 𝐺-invariant and 𝐴 ∈ Hom𝐺 (𝑉, 𝑉 ′ ). We can take 𝑓˜(𝑦) = 𝑗 ′ 𝐴 𝑗 −1 𝑦 + 𝑗 ′ 𝑏. □
F
Experimental Setup Details
Unless otherwise stated, the following hyperparameters are always chosen in our experiments (Section 5 of the main text): 22
depth attention type number of attention heads attention dropout transformer block drop path classification head dropout optimizer learning rate betas weight decay loss function
6 coupled 3 0.1 0.05 0.1 AdamW 0.001 0.9, 0.999 0.05 unweighted cross-entropy
Table 4: Default hyperparameters for all experiments.
For models operating on square grids (equivariant to 𝐷 4 or its subgroups), we choose the image size to be 2562 with patch size 162 . For models operating on hexagonal patches (equivariant to 𝐷 6 or its subgroups), each patch (𝑈) is a hexagonal lattice restricted to a regular hexagon with 𝑁2 = 9 grid points on each side, and the patches themselves (H ) form a hexagonal lattice with 𝑁1 = 9 grid points on each side (see Fig. 2 in the main text for the case 𝑁1 = 5, 𝑁2 = 9). This results in 217 hexagonal pixels per patch and 217 hexagonal patches, with a total of 42, 073 pixels in the whole image (note that the patches have nonempty overlap). Each image is preprocessed by normalizing the input RGB values of each pixel to [−1, 1]. No data augmentation is performed. For the 𝐷 6 family, we perform bilinear interpolation to convert a square image to one defined on H0 (union of hexagonal patches). For the first experiment (Section 5.1 of the main text), the maximum number of epochs and early stopping patience are 600/60 for 10% sample ratio, 150/50 for 40% sample ratio, and 50/10 for 100% sample ratio. For the second experiment (Section 5.2 of the main text), we perform at least three runs for each combination of (sample ratio, attention type, feature dimension). The maximum number of epochs and early stopping patience are 600/60 for 10% sample ratio and 200/30 for 100% sample ratio. For the third experiment (Section 5.3 of the main text), The maximum number of epochs and early stopping patience are 200/30.
23