🤖 machine learning

REViT: Roto-reflection Equivariant Convolutional Vision Transformer

This paper introduces REViT, a novel discrete roto-reflection equivariant vision transformer that incorporates convolutional attention to effectively preserve rotational, flip, and positional symmetries, demonstrating superior performance in image classification compared to existing equivariant neural network approaches.

Original authors: Sheir A. Zaheer, Alexander C. Holston, Chan Y. Park

Published 2026-06-25

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Sheir A. Zaheer, Alexander C. Holston, Chan Y. Park

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are looking at a photograph of a cat. If you rotate the photo 90 degrees, flip it upside down, or turn it sideways, you still know it's the same cat. However, most standard computer vision models (the "brains" behind image recognition) are like people who have never seen a cat from the side; they get confused when the picture changes orientation. They might think a sideways cat is a completely different animal.

To fix this, scientists build "equivariant" models. Think of these as models that have a built-in understanding of symmetry. If you rotate the input, the model's internal "thought process" rotates along with it, ensuring the final answer stays consistent.

This paper introduces a new model called REViT (Roto-reflection Equivariant Convolutional Vision Transformer). Here is how it works, broken down into simple concepts:

1. The Problem with the Old Way

Previously, to make a model understand rotations, researchers used two main tools:

Convolutional Neural Networks (CNNs): These are like a grid of tiny flashlights scanning an image. They are good at seeing patterns but can be rigid.
Vision Transformers (ViTs): These are like a team of detectives looking at the whole picture at once, connecting the dots between different parts. They are very powerful but usually struggle with rotation unless you add a lot of complex "position tags" (like GPS coordinates for every pixel) to tell them where things are.

The old method for making Transformers rotation-aware was like trying to teach a detective to rotate by giving them a massive, complicated map for every single turn. It worked, but it was slow and computationally expensive.

2. The REViT Solution: A New Kind of "Lift"

The authors of this paper came up with a simpler, more elegant way to build a rotation-aware Transformer. They removed the need for those complicated "position tags" entirely.

Instead, they use a "Lifting Layer."

The Analogy: Imagine a flat, 2D map of a city. Now, imagine you "lift" that map into a 3D tower. The bottom floor is the map as it is. The floors above are the same map, but rotated 45 degrees, 90 degrees, 135 degrees, and so on.
How it works: The REViT model takes an image and instantly creates this 3D "tower" of rotated versions of that image. It doesn't just look at the image; it looks at the image and all its possible rotations simultaneously.

3. The "Group Convolutional Self-Attention"

Once the image is "lifted" into this 3D tower, the model uses a special type of attention mechanism called Group Convolutional Self-Attention (G-CSA).

The Analogy: In a normal Transformer, the "detectives" look at different parts of the image to see how they relate. In REViT, the detectives are looking at the 3D tower. They can see how a "cat ear" in the original image relates to a "cat ear" in the 90-degree rotated image, all at once.
Because they are looking at all rotations together, the model naturally learns that "this is a cat, no matter how it's turned." It doesn't need to be told "this is the top" or "this is the bottom" because the 3D structure handles that for them.

4. What They Found (The Results)

The researchers tested this new model on three different "games" (datasets):

Rotated MNIST: Recognizing handwritten numbers that have been spun around.
PatchCamelyon: Identifying tumor cells in medical tissue samples (which can appear in any orientation).
CIFAR-10 & ImageNet: Recognizing everyday objects like cars, dogs, and airplanes.

The Results:

Better Accuracy: REViT beat the previous best methods for rotation-aware models. It got more questions right on the tests.
Simpler and Faster: Even though it was more accurate, it used fewer computer "steps" (parameters) and less memory than the older, more complicated methods.
Scalability: They showed that this model can handle huge, complex datasets (like ImageNet), proving it's not just a toy for small experiments but a robust tool for real-world use.

5. The Catch (Limitations)

The paper is honest about one downside: Because the model has to process the image in multiple rotated states at once (that 3D tower), it requires more computing power and memory than a standard model that doesn't care about rotation. It's like having a team of detectives working together instead of one; they get the job done better, but you need more office space and coffee for the whole team.

Summary

In short, REViT is a new type of AI that understands images from every angle without needing complex instructions on where things are located. By "lifting" the image into a 3D space of rotations and using a special attention mechanism, it recognizes objects more accurately and efficiently than previous methods, making it a strong candidate for tasks where orientation matters, like medical imaging or robotics.

Technical Summary: REViT – Roto-reflection Equivariant Convolutional Vision Transformer

Problem Statement
While equivariant neural networks have proven effective for tasks requiring rotational and reflectional symmetry (e.g., molecular analysis, medical pathology, and robotics), their application in Vision Transformers (ViTs) remains challenging. Traditional ViTs rely heavily on absolute or relative positional encodings to maintain spatial structure, which complicates the implementation of group equivariance. Specifically, standard positional encodings break equivariance under transformations like rotation or reflection because they assign unique vectors to fixed coordinates. Previous attempts to achieve group equivariance in ViTs, such as Group Self-Attention (G-SA) with relative position encoding (RPE), require complex modifications to attention mechanisms and introduce significant computational overhead. Furthermore, existing research on roto-reflection equivariance has predominantly focused on Convolutional Neural Networks (G-CNNs), leaving a gap in leveraging the global context modeling capabilities of transformers for these tasks.

Methodology
The authors propose REViT (Roto-reflection Equivariant Convolutional Vision Transformer), a discrete roto-reflection group equivariant architecture that eliminates the need for explicit positional encodings. The core methodology involves two primary innovations:

Convolutional Patch Embedding and Lifting:
Instead of linear patch embedding, REViT employs a lifting layer that transforms the 2D input image $f: \mathbb{Z}^2 \to \mathbb{R}^C$ into a feature map defined over a discrete group space $F: \mathbb{Z}^2 \times C_N \to \mathbb{R}^{C'}$ . This is achieved via rotational convolutions where the input is convolved with a set of rotated filters corresponding to the elements of the discrete rotation group $C_N$ (e.g., $p4$ , $p8$ ). This lifts the data from Euclidean space to a position-orientation space ( $\mathbb{R}^2 \times S^1$ ), inherently encoding spatial and orientational information without learnable positional parameters.
Group Convolutional Self-Attention (G-CSA):
The standard self-attention mechanism is replaced with Group Convolutional Self-Attention (G-CSA). In this module, the Query, Key, and Value projections are computed using group-equivariant convolutions rather than linear projections. The attention mechanism operates over both spatial dimensions and the group dimension (orientation). The attention weights are calculated based on the dot product of queries and keys across the lifted feature space, ensuring that the operation remains equivariant to the group actions (rotations and reflections). By utilizing convolutional projections and relative spatial relationships inherent in the convolution, REViT achieves shift and rotation equivariance without the per-layer complexity of computing relative position encodings for every attention head.

Key Contributions

Discrete Group Convolutional Self-Attention (G-CSA): The paper introduces G-CSA as a mechanism to achieve roto-reflection ( $E(2,N)$ ) equivariance in transformers without relying on position encoding.
Implementation of REViT: The authors implement REViT blocks utilizing G-CSA and evaluate them across multiple discrete roto-translation ($pN$) and roto-reflection ($pNm$) group configurations.
Scalability: The work demonstrates that REViT can be scaled to large-scale datasets, specifically ImageNet-1k, serving as a potential backbone for downstream equivariant tasks.

Experimental Results
The authors evaluated REViT on three datasets: Rotated MNIST, PatchCamelyon, and CIFAR-10, and scaled the model to ImageNet-1k.

Performance vs. G-SA with RPE: On Rotated MNIST and PatchCamelyon, REViT models (specifically $p4$ -CSA and $p8$ -CSA) matched or outperformed Group Self-Attention models enriched with Relative Position Encoding (G-SA+RPE). Crucially, REViT achieved this with significantly fewer parameters and lower computational complexity (fewer multiplication-addition operations and reduced memory footprint).
Performance vs. G-CNNs: REViT consistently outperformed corresponding Group Equivariant CNN (G-CNN) baselines across all datasets and symmetry groups. Notably, on PatchCamelyon, REViT achieved higher accuracy than higher-order symmetry G-CNN variants while using fewer parameters.
ImageNet Scaling: On ImageNet-1k, a $p4m$ -equivariant REViT achieved a Top-1 accuracy of 79.27%, outperforming both a vanilla ViT trained with augmentation (72.08%) and an equivariant ResNet baseline (77.37%).
Equivariance Error: The models exhibited near-zero equivariance error after the lifting and G-CSA layers, with minor errors in higher-order groups (e.g., $p8$ ) attributed to interpolation artifacts during rotation.

Significance and Claims
The paper claims that REViT offers a simpler and more efficient pathway to achieving discrete roto-reflection equivariance in vision transformers compared to existing methods. By removing the reliance on explicit positional encodings and leveraging convolutional self-attention, the authors demonstrate that it is possible to maintain the global context modeling strengths of transformers while adhering to strict symmetry constraints. The results suggest that G-CSA improves group equivariant classification performance without necessitating a significant increase in model size or complexity. The authors position REViT as a viable, scalable backbone for tasks where orientation is critical, such as medical imaging and oriented object detection, while acknowledging that the computational cost still scales with the cardinality of the symmetry group, presenting a challenge for latency-sensitive applications.