Vanilla Group Equivariant Vision Transformer: Simple and Effective

Incorporating symmetry priors as inductive biases to design equivariant Vision Transformers (ViTs) has emerged as a promising avenue for enhancing their performance. However, existing equivariant ViTs often struggle to balance performance with equivariance, primarily due to the challenge of achieving holistic equivariant modifications across the diverse modules in ViTs-particularly in harmonizing the Self-Attention mechanism with Patch Embedding. To address this, we propose a straightforward framework that systematically renders key ViT components, including patch embedding, self-attention, positional encodings, and Down/Up-Sampling, equivariant, thereby constructing ViTs with guaranteed equivariance. The resulting architecture serves as a plug-and-play replacement that is both theoretically grounded and practically versatile, scaling seamlessly even to Swin Transformers. Extensive experiments demonstrate that our equivariant ViTs consistently improve performance and data efficiency across a wide spectrum of vision tasks.

💡 Research Summary

The paper introduces a simple yet powerful framework for constructing vision transformers (ViTs) that are equivariant to the dihedral group consisting of 90‑degree rotations and reflections. Existing equivariant ViTs either sacrifice performance or require complex modifications, largely because the patch‑embedding stage breaks the inherent permutation equivariance of self‑attention. To overcome this, the authors propose a “vanilla” approach that systematically makes every major component of a ViT equivariant: patch embedding, self‑attention, positional encoding, and down/up‑sampling.

The key technical contribution is the introduction of a group dimension of size t = |G| (e.g., t = 8 for the rotation‑reflection group). An equivariant patch embedding (EQ‑PE) lifts the image into a feature map with an extra group channel by applying the same convolution kernel to all transformed versions of the input (π_g x) and stacking the results. This yields a feature tensor \hat{z} ∈ ℝ^{H_s × W_s × C × t} that shifts cyclically along the group axis when the input is transformed, guaranteeing equivariance at the very first stage.

Self‑attention is then adapted (EQ‑SA) by using equivariant linear layers. For each group slice z_g, a separate weight matrix W_q^g, W_k^g, W_v^g is learned, but the matrices are tied together by a tiling operation that creates a large block‑circulant matrix W ∈ ℝ^{Ct × Ct}. Multiplying this with the reshaped token matrix Z ∈ ℝ^{N × Ct} implements query, key, and value transformations that respect the group structure. Consequently, the attention map A = Softmax(QKᵀ/√Ct) and the resulting output are equivariant. Multi‑head attention follows the same principle.

Positional encodings are made equivariant by sharing embeddings across an orbit of positions. For absolute encodings, each patch coordinate p is mapped to its canonical representative p_c (the lexicographically smallest element of its orbit), and the same embedding is used for all members of the orbit. Relative encodings (as used in Swin‑Transformer) are similarly tied across the group, ensuring that the bias added to the attention scores transforms consistently.

Down‑sampling and up‑sampling layers, crucial for hierarchical transformers like Swin, are re‑designed to operate independently on each group slice and then re‑align the group dimension, preserving equivariance across resolution changes.

The authors provide rigorous proofs that each module satisfies equivariance with respect to the chosen group and that the composition of equivariant modules yields a globally equivariant network. They also argue that parameter sharing reduces model capacity, which theoretically tightens generalization bounds.

Empirically, the method is evaluated on a broad set of tasks: ImageNet‑1K classification, CIFAR‑10/100, COCO detection, Cityscapes segmentation, and DIV2K super‑resolution. When applied to ViT‑Base, ViT‑Large, and Swin‑Tiny/Small/Base, the equivariant versions consistently outperform their vanilla counterparts by 0.5‑1.2 % in top‑1 accuracy on ImageNet, while requiring only a quarter of the training data to achieve comparable performance. Detection and segmentation mAP/mIoU improve by roughly 1 %, and super‑resolution PSNR gains are around 0.1‑0.3 dB. Importantly, FLOPs and memory usage remain on par with the original models because the group dimension is handled with shared weights and efficient reshaping.

Overall, the paper delivers a plug‑and‑play recipe for end‑to‑end equivariant vision transformers that does not demand radical architectural changes. By making the entire pipeline equivariant, it reduces reliance on massive data augmentation, improves data efficiency, and offers a clear path to extending equivariance to newer ViT variants. The work opens avenues for further exploration of larger symmetry groups (continuous rotations, scaling) and for integrating equivariance into multimodal transformer architectures.