LookHere: Vision Transformers with Directed Attention Generalize and Extrapolate
This addresses a key limitation in computer vision models for researchers and practitioners, offering a drop-in solution to enhance ViT generalization and extrapolation, though it is incremental as it builds on existing ViT architectures.
The paper tackles the problem of vision transformers (ViTs) failing to effectively leverage larger images without finetuning, proposing a new position encoding method called LookHere that restricts attention heads to fixed fields of view; it improves performance on classification, adversarial robustness, and calibration error, and outperforms the state-of-the-art by 21.7% when extrapolating to higher resolutions.
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.