Source-linked AI summary
REViT-v2: Hierarchical Windowed Roto-reflection Equivariant ViT for Equivariant Feature Extraction
Sheir A. Zaheer, Jihwan Moon, Chan Y. Park
TL;DR
Standard vision transformers do not explicitly encode image symmetries, and global group-equivariant attention has quadratic spatial cost. REViT-v2 combines local windowed group-equivariant attention with a hierarchical feature architecture. It outperforms two reported ImageNet-1K baselines and reduces resource costs on Rotated MNIST while maintaining comparable accuracy.
Problem
Standard ViTs do not explicitly encode image symmetries, while global group-equivariant attention has quadratic spatial cost that is especially important when representations allocate multiple orientation channels.
Source: §1 · p. 1Method
REViT-v2 combines windowed group-convolutional self-attention with a multistage equivariant feature hierarchy.
Source: §1 · p. 2, §2 · p. 2, §3.3 · p. 4Results
REViT-v2 outperforms the augmented vanilla ViT and p4m-equivariant RE-ResNet on ImageNet-1K; on Rotated MNIST, wG-CSA reduces resource costs while maintaining comparable accuracy.
Source: §4 · p. 4, Complexity Comparison with Global G-CSA · p. 4Takeaways & Limitations
The authors demonstrate REViT-v2 as a scalable group-equivariant feature-extractor backbone and provide code and pretrained weights for further research.
Source: §5 · p. 4
Abstract
from arXiv · showhide
We propose a scalable roto-reflection-group-equivariant vision transformer based on windowed group-convolutional self-attention and a hierarchical feature architecture. We demonstrate that our approach can be scaled to group-equivariant vision transformers (ViTs) with millions of parameters and large datasets with practically sized images, i.e., ImageNet. The code and pretrained weights for the proposed Hierarchical Windowed Roto-reflection Equivariant ViTs (REViT-v2) are available at https://github.com/kc-ml2/revit.
1. Introduction
Standard ViTs do not explicitly encode image symmetries, while prior equivariant-attention approaches face scaling costs. REViT-v2 combines local group-equivariant attention with a hierarchy to target scalable recognition while maintaining equivariance.
1. Introduction
- Standard ViTs do not explicitly encode rotation and reflection symmetry, whereas group-equivariant networks incorporate these relationships into their architectures.Group convolutions share filters across translations, rotations, and reflections.Source: §1 · p. 1
- Equivariant attention based on invariant positional encodings does not scale beyond low-resolution images, while global G-CSA retains quadratic spatial cost, amplified by orientation channels.REViT removed explicit positional encoding through convolutional projections, but its global attention remained costly.Source: §1 · p. 1
- REViT-v2 combines windowed group-convolutional self-attention with a multistage equivariant feature pyramid, reducing attention complexity from quadratic to linear in image area.The design targets large-scale recognition while maintaining equivariance throughout the feature extractor.Source: §1 · p. 2
2. Windowed Group Convolutional Self-Attention
wG-CSA extends group-convolutional self-attention with local windows, using equivariant projections and complete representation fields to retain equivariance while reducing spatial attention cost.
2. Windowed Group Convolutional Self-Attention
- wG-CSA is a local-window extension of G-CSA intended to reduce its quadratic complexity in image size to linear complexity.Source: §2 · p. 2
2.1. Equivariant Projections
- Group-equivariant convolutional projections produce queries, keys, and values from feature maps organized by a representation of a finite planar symmetry group.The attention output is followed by another group-equivariant projection.Source: §2.1 · p. 2
2.2. Local Window Attention
- Attention is restricted to each position’s local window, and head outputs are concatenated along complete representation fields before projection.The method introduces no absolute positional embedding; spatial structure is preserved by the equivariant stem and convolutional Q/K/V projections.Source: §2.2 · pp. 2–3
2.3. Linear Complexity of Windowed G-CSA
- For fixed window size M, windowed attention costs O(NM^2C), linear in image area N, instead of global attention’s O(N^2C).Here N is the number of feature-map positions; windows contain M × M positions.Source: §2.3 · p. 3
3. REViT-v2
REViT-v2 combines windowed equivariant attention with a multistage feature hierarchy. Equivariant lifting, transformer blocks, and progressive downsampling build capacity while addressing the costs of high-resolution attention.
3. REViT-v2
- The hierarchy uses early layers to preserve spatial resolution for local structures and deeper layers to increase representation capacity and receptive fields.Source: §3 · p. 3
3.1. Equivariant Patching and Lifting
- A group-equivariant stem lifts the image, whose RGB channels are treated as trivial representations, into regular representations and reduces resolution by a factor of four.Source: §3.1 · p. 3
3.2. Equivariant Transformer Block
- Each transformer block combines pre-normalized wG-CSA with an equivariant feed-forward network using 1 × 1 group convolutions.Stochastic depth regularizes residual branches without independently dropping orientation channels, which would violate their representation structure.Source: §3.2 · p. 3
3.3. Multistage Hierarchy
- After each of the first three stages, strided group convolution reduces resolution and increases representation fields; together with windowing, this avoids propagating high-resolution maps through the full network.At deeper stages, fixed-size windows cover progressively larger regions of the input.Source: §3.3 · p. 4
4. Experimental Results
REViT-v2 outperforms the two reported ImageNet-1K baselines. On Rotated MNIST, windowed G-CSA substantially reduces resource costs while maintaining comparable accuracy.
4. Experimental Results
- REViT-v2 outperforms both the rotation- and flip-augmented vanilla ViT and the p4m-equivariant RE-ResNet on ImageNet-1K.Table 1 compares the p4m-equivariant REViT-v2 with these two baselines.Source: §4 · p. 4
Complexity Comparison with Global G-CSA
- On Rotated MNIST, wG-CSA substantially reduces model size, forward FLOPs, latency, and peak memory versus global G-CSA while maintaining comparable accuracy.The comparison also evaluates global attention after the same downsampling stem to distinguish early downsampling from windowing effects.Source: Complexity Comparison with Global G-CSA · p. 4
5. Conclusion
REViT-v2 is presented as a scalable group-equivariant vision-transformer backbone, with code and pretrained weights provided for further research.
5. Conclusion
- REViT-v2 is presented as a scalable group-equivariant vision transformer that can serve as a backbone feature extractor like conventional ViTs.The paper also provides code and pretrained weights to support further research.Source: §5 · p. 4
- The proposed models are intended to serve as backbone feature extractors in the same manner as conventional ViTs.Source: §5 · p. 4
- Code and pretrained weights are provided to support further research.Source: §5 · p. 4
A.1. Theoretical Analysis
The analysis assumes orthogonal group representations within each attention head. Under that assumption, attention coefficients transform consistently, provided transformations map complete windows to complete windows.
A.1. Theoretical Analysis
- Orthogonality of the group representations within each head is assumed, as with the permutation-based regular representations used in this work.Source: §A.1 · p. 6
- Under the stated assumption, attention coefficients are permuted consistently with transformed spatial positions, and the output satisfies the equivariance relation.Source: §A.1 · p. 6
- The equivariance condition requires transformations to map complete windows to complete windows; square windows meet it for 90° rotations and axis-aligned reflections when feature-map dimensions are divisible by the window size.Source: §A.1 · p. 6
A.2. Quantitative Analysis
The appendix evaluates architectural equivariance at the stem and pre-classification stages on sampled ImageNet validation images. It reports that p8 errors begin at lifting, which the authors attribute to interpolation at off-grid rotations.
A.2. Quantitative Analysis
- Equivariance was tested on 1,024 ImageNet validation images at the lifting and pre-classification outputs, with the model randomly initialized to isolate design-induced equivariance.Source: §A.2 · p. 7
- For p8, equivariance errors accumulate starting at the lifting layer, attributed to interpolation artifacts at 45° and 135° rotations that place pixels between grid locations.Source: §A.2 · p. 7
- The analysis compares responses for each sampled image under every non-identity element of the p4, p4m, and p8 groups.Source: §A.2 · p. 7
Appendix B. Training Setup
REViT-v2 is trained on ImageNet for 300 epochs using AdamW and a cosine-decay schedule after warmup, with distributed training across four RTX 4090 GPUs.
Appendix B. Training Setup
- REViT-v2 is trained for 300 epochs with AdamW, an initial learning rate of 3e-4, and weight decay of 0.05.Source: Appendix B · p. 7
- Training uses four Nvidia GeForce RTX 4090 GPUs with per-GPU batch size 128, for an effective batch size of 512.Source: Appendix B · p. 7
- The learning-rate schedule uses 20 warmup epochs followed by cosine decay; attention projections use a 3x3 equivariant kernel and window size 7.Source: Appendix B · p. 7