A Unified Framework for Vision Transformers Equivariant to Discrete Subgroups of $\mathrm{O}(2)$

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 $\mathrm{O}(2)$, providing a unified framework that generalizes prior flipping- and $D_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 $H \le G$, the class of $G$-equivariant ViTs embeds naturally into the class of $H$-equivariant ViTs. We also prove that, in the single-head setting, the corresponding equivariant self-attention layer realizes every $G$-equivariant self-attention map representable by ordinary self-attention. We further construct a $D_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 $D_4$ and $D_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.

• The paper introduces a unified mathematical framework that generalizes vision transformers to be equivariant to any discrete subgroup of O(2).
• The methodology leverages group representation theory to design G-equivariant patch embeddings, MLP nonlinearities, and attention mechanisms.
• Empirical results demonstrate that enhanced equivariance boosts data efficiency in low-data regimes while enabling flexible symmetry-breaking strategies.

Unified Framework for Discrete $\mathrm{O}(2)$-Equivariant Vision Transformers

Motivation and Context

This work develops a unified mathematical and algorithmic formalism for constructing vision transformers (ViTs) that are equivariant to any discrete subgroup $G \leq \mathrm{O}(2)$, capturing planar symmetries such as rotations and reflections. While group-equivariant models for convolutional neural networks (CNNs) are established and widely adopted, vision transformers have largely ignored explicit geometric inductive bias despite their dominance in visual recognition. Existing efforts have focused on special cases (e.g., $D_1$ or $D_4$ equivariance); the present framework generalizes and systematizes the construction of equivariant ViTs for an arbitrary discrete planar symmetry group. The analysis draws on representation theory, enabling precise control of model bias, expressivity, and architectural flexibility.

Mathematical Framework and Architecture

The equivariant ViT framework is grounded in rigorous group representation theory. Tokens and feature spaces are modeled as direct sums of irreducible $G$-representations. The key architectural modules—patch embedding, positional encoding, linear layers, MLP nonlinearities, and multi-head self-attention—are parameterized to be $G$-equivariant by construction.

Patchification and Feature Representations

Input images are patchified with respect to a $G$-stable grid structure, parametrized by a Minkowski sum $\mathcal{H}_0 = U + \mathcal{H}$ of two $G$-stable finite subsets of the plane (base patch set $U$ and patch positions $G \leq \mathrm{O}(2)$0). Token features live in a $G \leq \mathrm{O}(2)$1-representation $G \leq \mathrm{O}(2)$2, where $G \leq \mathrm{O}(2)$3 ranges over real irreducible representations of $G \leq \mathrm{O}(2)$4. The representation type structure enables precise decomposition and control over equivariant feature channels.

Equivariant Layers

• Linear Maps and Bias: By Schur’s lemma, $G \leq \mathrm{O}(2)$5-equivariant linear maps between feature spaces decompose irrep-wise; bias is only permitted in the trivial representation.
• Patch Embedding: Composed of an equivariant map from patches to token features, using a basis of $G \leq \mathrm{O}(2)$6-equivariant filters.
• Positional Encoding: Only $G \leq \mathrm{O}(2)$7-invariant encodings are allowed, ensuring equivariance to $G \leq \mathrm{O}(2)$8.
• Nonlinearities: The class of permissible $G \leq \mathrm{O}(2)$9-equivariant nonlinearities is systematized using Fourier-type transforms on homogeneous $D_1$0-spaces. Entrywise activation is equivariant if and only if the representation is induced from a $D_1$1-set (i.e., permutation representation), and nonlinear layers are shown to be maximally general under natural assumptions.
• Self-Attention: Attention is formulated as ordinary multi-head self-attention with $D_1$2-invariant inner products and $D_1$3-stable orthogonal head decompositions. Coupled and irrep-wise heads are supported. The framework guarantees that under single-head attention, equivariant parametrization is maximally expressive among $D_1$4-equivariant functions implementable by standard self-attention.

Expressivity Results

A striking theoretical claim is formalized: if $D_1$5, every $D_1$6-equivariant ViT is naturally an $D_1$7-equivariant ViT, including strict embeddings. Therefore, enforcing larger symmetry groups induces stricter inductive bias and provable function class inclusion. This quantifies the expressivity–bias tradeoff in equivariant networks and enables staged symmetry breaking for practical modeling.

Algorithmic Innovations

• Hexagonal and Non-Square Grids: Construction generalizes to $D_1$8-equivariant ViTs using hexagonal patchification, enabling compatibility with six-fold symmetry.
• Homogeneous Space Parameterization of MLP Layers: Nonlinear invariance is achieved via pointwise nonlinearities on $D_1$9-sets decomposed into homogeneous spaces, allowing refined control over the layerwise representation structure and information flow between irreps.

Empirical Results

The primary empirical domain is aerial image classification with PatternNet under data-scarce regimes and various discrete symmetry groups ($D_4$0, $D_4$1, $D_4$2, etc. on square grids; $D_4$3, $D_4$4, $D_4$5, etc. on hexagonal grids). Main empirical findings:

• Equivariance boosts data efficiency: $D_4$6-equivariant ViTs consistently outperform nonequivariant and lower-symmetry models when training data is limited, validating the inductive bias hypothesis. Performance gaps largely vanish when training with full datasets.
• Tradeoff in symmetry and expressivity: More symmetry is not always better; in some setups, classifiers with intermediate symmetry ($D_4$7) outperform those with maximal symmetry ($D_4$8) at matched parameter budgets.
• Attention and MLP structure: Irrep-wise attentions are marginally superior to coupled attention in some low-data settings. Choice of $D_4$9-set for MLP nonlinearities can modestly impact performance, particularly favoring configurations with higher fractions of trivial irrep components in the hidden layers.
• Architectural flexibility: The unified formalism supports systematic design and analysis of equivariant ViT architectures, including new group symmetries and symmetry breaking strategies across training.

Implications, Limitations, and Future Directions

This framework provides the theoretical and practical foundation for systematic design, optimization, and analysis of equivariant vision transformers for arbitrary discrete planar symmetries. The results have significant implications for data-efficient learning, robustness to geometric transformations, and principled control of model bias versus expressivity. The embedding result has clear potential for staged symmetry-breaking in training (dynamically or blockwise), making it possible to balance inductive bias and representational flexibility in complex real-world tasks.

Numerical experiments highlight that the benefits of equivariant architectures are most pronounced in data-limited regimes; further exploration at scale and with more diverse datasets is warranted, especially considering the computational benefits of equivariant factorization for large $G$0.

Open directions include:

• Extending the framework to non-dihedral and irregular $G$1 (requiring more exotic grid structures).
• Continuous symmetry groups, approximate equivariance, and equivariance-breaking schedules.
• Automated and task-adaptive selection of feature space decompositions and group actions.
• Applications to detection, segmentation, and other structured vision tasks.
• Empirical exploration of the scaling laws for equivariant architectures in the large-data regime [see for instance (Brehmer et al., 2024)].

Conclusion

This paper introduces a unified, representation-theoretic framework for constructing, analyzing, and generalizing vision transformers equivariant to all discrete subgroups of $G$2. Through mathematical results and controlled experiments, it establishes principled methods for incorporating planar symmetries into ViTs, quantifies tradeoffs between bias and expressivity, and enables new architectural flexibility. The demonstrated gains in data efficiency and robust modeling motivate further study on the role of symmetry in large-scale visual recognition and structured data domains (2606.27864).