LookHere: Vision Transformers with Directed Attention Generalize and Extrapolate

TL;DR

LookHere introduces a novel position encoding method for vision transformers that enhances their ability to generalize and extrapolate to larger images, significantly improving high-resolution image classification performance.

Contribution

We propose LookHere, a new position encoding that restricts attention heads to fixed, directed fields of view, improving extrapolation and generalization of vision transformers to high-resolution images.

Findings

Improves classification accuracy by 1.6% on ImageNet.

Reduces adversarial attack success rate by 5.4%.

Outperforms 2D-RoPE by 21.7% on extrapolated ImageNet images.

Abstract

High-resolution images offer more information about scenes that can improve model accuracy. However, the dominant model architecture in computer vision, the vision transformer (ViT), cannot effectively leverage larger images without finetuning -- ViTs poorly extrapolate to more patches at test time, although transformers offer sequence length flexibility. We attribute this shortcoming to the current patch position encoding methods, which create a distribution shift when extrapolating. We propose a drop-in replacement for the position encoding of plain ViTs that restricts attention heads to fixed fields of view, pointed in different directions, using 2D attention masks. Our novel method, called LookHere, provides translation-equivariance, ensures attention head diversity, and limits the distribution shift that attention heads face when extrapolating. We demonstrate that LookHere improves performance on classification (avg. 1.6%), against adversarial attack (avg. 5.4%), and decreases calibration error (avg. 1.5%) -- on ImageNet without extrapolation. With extrapolation, LookHere outperforms the current SoTA position encoding method, 2D-RoPE, by 21.7% on ImageNet when trained at $224^2$ px and tested at $1024^2$ px. Additionally, we release a high-resolution test set to improve the evaluation of high-resolution image classifiers, called ImageNet-HR.