Reviving Shift Equivariance in Vision Transformers

TL;DR

This paper introduces an adaptive polyphase anchoring method to restore shift-equivariance in vision transformers, significantly improving their robustness and prediction consistency under input shifts and transformations.

Contribution

It proposes a novel adaptive polyphase anchoring algorithm that ensures shift-equivariance in vision transformers, addressing a key limitation of existing models.

Findings

Achieves 100% shift consistency in predictions.

Demonstrates robustness to cropping, flipping, and affine transformations.

Maintains high accuracy even under input shifts that reduce baseline models' performance.

Abstract

Shift equivariance is a fundamental principle that governs how we perceive the world - our recognition of an object remains invariant with respect to shifts. Transformers have gained immense popularity due to their effectiveness in both language and vision tasks. While the self-attention operator in vision transformers (ViT) is permutation-equivariant and thus shift-equivariant, patch embedding, positional encoding, and subsampled attention in ViT variants can disrupt this property, resulting in inconsistent predictions even under small shift perturbations. Although there is a growing trend in incorporating the inductive bias of convolutional neural networks (CNNs) into vision transformers, it does not fully address the issue. We propose an adaptive polyphase anchoring algorithm that can be seamlessly integrated into vision transformer models to ensure shift-equivariance in patch embedding and subsampled attention modules, such as window attention and global subsampled attention. Furthermore, we utilize depth-wise convolution to encode positional information. Our algorithms enable ViT, and its variants such as Twins to achieve 100% consistency with respect to input shift, demonstrate robustness to cropping, flipping, and affine transformations, and maintain consistent predictions even when the original models lose 20 percentage points on average when shifted by just a few pixels with Twins' accuracy dropping from 80.57% to 62.40%.