How Do Vision Transformers Work?

Abstract

The success of multi-head self-attentions (MSAs) for computer vision is now indisputable. However, little is known about how MSAs work. We present fundamental explanations to help better understand the nature of MSAs. In particular, we demonstrate the following properties of MSAs and Vision Transformers (ViTs): (1) MSAs improve not only accuracy but also generalization by flattening the loss landscapes. Such improvement is primarily attributable to their data specificity, not long-range dependency. On the other hand, ViTs suffer from non-convex losses. Large datasets and loss landscape smoothing methods alleviate this problem; (2) MSAs and Convs exhibit opposite behaviors. For example, MSAs are low-pass filters, but Convs are high-pass filters. Therefore, MSAs and Convs are complementary; (3) Multi-stage neural networks behave like a series connection of small individual models. In addition, MSAs at the end of a stage play a key role in prediction. Based on these insights, we propose AlterNet, a model in which Conv blocks at the end of a stage are replaced with MSA blocks. AlterNet outperforms CNNs not only in large data regimes but also in small data regimes. The code is available at https://github.com/xxxnell/how-do-vits-work.

ICLR 2022 (Spotlight)

Cited by in corpus (10)

• TransMorph: Transformer for unsupervised medical image registration
• MedViT: A Robust Vision Transformer for Generalized Medical Image Classification
• Point Cloud Classification Using Content-based Transformer via Clustering in Feature Space
• Fredformer: Frequency Debiased Transformer for Time Series Forecasting
• On the Role of ViT and CNN in Semantic Communications: Analysis and Prototype Validation
• Clinically-Inspired Multi-Agent Transformers for Disease Trajectory Forecasting from Multimodal Data
• LGViT: Dynamic Early Exiting for Accelerating Vision Transformer
• A Simple and Robust Framework for Cross-Modality Medical Image Segmentation applied to Vision Transformers
• STAA: Spatio-Temporal Alignment Attention for Short-Term Precipitation Forecasting
• Architecture Analysis and Benchmarking of 3D U-shaped Deep Learning Models for Thoracic Anatomical Segmentation