How Do Vision Transformers Work? (2202.06709v4)

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.

• The paper demonstrates that multi-head self-attention improves accuracy and generalization by smoothing loss landscapes.
• It reveals that vision transformers and CNNs exhibit complementary filtering properties, paving the way for effective hybrid architectures.
• It introduces AlterNet, a hybrid model that replaces Conv blocks with MSA blocks to achieve superior performance on benchmark datasets.

How Do Vision Transformers Work?

The paper presents a detailed exploration of the mechanics behind Multi-Head Self-Attentions (MSAs) and Vision Transformers (ViTs) in computer vision. It delivers robust empirical evidence to discuss three primary properties of MSAs, contrasting with conventional convolutional neural networks (CNNs) and proposing a new model, AlterNet, which combines Convolutional Networks (Convs) with MSAs.

Key Insights

1. MSAs and Loss Landscapes:
   • The research demonstrates that MSAs not only enhance accuracy but also improve generalization by smoothing loss landscapes. The beneficial impact is largely attributed to data specificity rather than their ability to model long-range dependencies.
   • ViTs are shown to face challenges with non-convex loss functions. However, when trained on large datasets or with loss landscape smoothing techniques, they overcome these barriers, achieving competitive performance.
2. Comparison with Convs:
   • The paper reveals opposing behaviors between MSAs and Convs. While MSAs function as low-pass filters, reducing high-frequency signals, Convs operate as high-pass filters.
   • The complementary nature of MSAs and Convs suggests opportunities for hybrid architectures. The paper explores how Convs and MSAs can be harmonized, showcasing that each has unique attributes beneficial in different contexts.
3. Multi-Stage Networks and AlterNet:
   • It is posited that multi-stage networks operate like a series of connected individual models. MSAs, particularly those at the end of a stage, contribute significantly to performance.
   • AlterNet is proposed, substituting Conv blocks at a stage's conclusion with MSA blocks. This design achieves superior performance across both large and small dataset scenarios compared to traditional CNNs.

Methodology and Results

The researchers conduct a suite of experiments on CIFAR and ImageNet datasets, employing strong data augmentation techniques. They rigorously analyze the spectral density of Hessian eigenvalues, revealing that MSAs tend to flatter the loss landscapes compared to CNNs. This flattening is associated with improved generalization capabilities.

Through Fourier analysis, the paper highlights the distinctive frequency-domain behaviors of MSAs and Convs. ViTs demonstrate robustness against high-frequency perturbations, unlike the texture-biased ResNets vulnerable to such noise.

Implications and Future Directions

The implications of this research are manifold. On a theoretical level, it challenges the accepted understanding of MSAs as primarily beneficial for modeling long-range dependencies, shifting focus towards their data-specific filtering capabilities. Practically, the paper sets the stage for more nuanced integration of MSAs and Convs in neural architectures, suggesting AlterNet as a viable path forward.

Looking ahead, the authors encourage further exploration into the loss landscape properties of MSAs, especially concerning ViTs’ non-convexity challenges. Another avenue is refining hybrid models like AlterNet to exploit the best of both MSAs and Convs, potentially influencing a broad spectrum of vision tasks.

In summary, this research not only deepens understanding of MSAs and ViTs but also innovates in architectural design for enhanced performance in diverse data regimes. The insights offered reaffirm the potential of transformers while paving the way for pragmatic improvements in AI vision systems.