How Do Vision Transformers Work?

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 . 1 INTRODUCTION . There is limited understanding of multi-head self-attentions ( MSAs ) , although they are now ubiquitous in computer vision . The most widely accepted explanation for the success of MSAs is their weak inductive bias and capture of long-range dependencies ( See , e.g. , ( Dosovitskiy et al. , 2021 ; Naseer et al. , 2021 ; Tuli et al. , 2021 ; Yu et al. , 2021a ; Mao et al. , 2021 ; Chu et al. , 2021 ) ) . Yet because of their over-flexibility , Vision Transformers ( ViTs ) —neural networks ( NNs ) consisting of MSAs—have been known to have a tendency to overfit training datasets , consequently leading to poor predictive performance in small data regimes , e.g. , image classification on CIFAR . However , we show that the explanation is poorly supported . 1.1 RELATED WORK . Self-attentions ( Vaswani et al. , 2017 ; Dosovitskiy et al. , 2021 ) aggregate ( spatial ) tokens with normalized importances : zj = ∑ i Softmax ( QK√ d ) i Vi , j ( 1 ) where Q , K , and V are query , key , and value , respectively . d is the dimension of query and key , and zj is the j-th output token . From the perspective of convolutional neural networks ( CNNs ) , MSAs are a transformation of all feature map points with large-sized and data-specific kernels . Therefore , MSAs are at least as expressive as convolutional layers ( Convs ) ( Cordonnier et al. , 2020 ) , although this does not guarantee that MSAs will behave like Convs . Is the weak inductive bias of MSA , such as modeling long-range dependencies , beneficial for the predictive performance ? To the contrary , appropriate constraints may actually help a model learn strong representations . For example , local MSAs ( Yang et al. , 2019 ; Liu et al. , 2021 ; Chu et al. , 2021 ) , which calculate self-attention only within small windows , achieve better performance than global MSAs not only on small datasets but also on large datasets , e.g. , ImageNet-21K . In addition , prior works observed that MSAs have the following intriguing properties : 1 MSAs improve the predictive performance of CNNs ( Wang et al. , 2018 ; Bello et al. , 2019 ; Dai et al. , 2021 ; ||∆winit|| and θt = cos −1 ( ∆wt·∆winit ||∆wt|| ||∆winit|| ) where ∆wt = wt −woptim and wt is NN weight at tth epoch . winit and woptim are the weights at initialization and optimum . Right : The magnitude of the Hessian eigenvalues of ViT is smaller than that of ResNet . See Fig . 4 for a more detailed analysis . Guo et al. , 2021 ; Srinivas et al. , 2021 ) , and ViTs predict well-calibrated uncertainty ( Minderer et al. , 2021 ) . 2 ViTs are robust against data corruptions , image occlusions ( Naseer et al. , 2021 ) , and adversarial attacks ( Shao et al. , 2021 ; Bhojanapalli et al. , 2021 ; Paul & Chen , 2022 ; Mao et al. , 2021 ) . They are particularly robust against high-frequency noises ( Shao et al. , 2021 ) . 3 MSAs closer to the last layer significantly improve predictive performance ( Graham et al. , 2021 ; Dai et al. , 2021 ) . These empirical observations raise immediate questions : 1 What properties of MSAs do we need to better optimize NNs ? Do the long-range dependencies of MSAs help NNs learn ? 2 Do MSAs act like Convs ? If not , how are they different ? 3 How can we harmonize MSAs with Convs ? Can we just leverage their advantages ? We provide an explanation of how MSAs work by addressing them as a trainable spatial smoothing of feature maps , because Eq . ( 1 ) also suggests that MSAs average feature map values with the positive importance-weights . Even non-trainable spatial smoothings , such as a small 2× 2 box blur , help CNNs see better ( Zhang , 2019 ; Park & Kim , 2021 ) . These simple spatial smoothings not only improve accuracy but also robustness by spatially ensembling feature map points and flattening the loss landscapes ( Park & Kim , 2021 ) . Remarkably , spatial smoothings have the properties of MSAs 1 – 3 . See Appendix B for detailed explanations of MSAs as a spatial smoothing . 1.2 CONTRIBUTION . We address the three key questions : 1 What properties of MSAs do we need to improve optimization ? We present various evidences to support that MSA is generalized spatial smoothing . It means that MSAs improve performance because their formulation—Eq . ( 1 ) —is an appropriate inductive bias . Their weak inductive bias disrupts NN training . In particular , a key feature of MSAs is their data specificity , not long-range dependency . As an extreme example , local MSAs with a 3× 3 receptive field outperforms global MSA because they reduce unnecessary degrees of freedom . How do MSAs improve performance ? MSAs have their advantages and disadvantages . On the one hand , they flatten loss landscapes as shown in Fig . 1 . The flatter the loss landscape , the better the performance and generalization ( Li et al. , 2018 ; Keskar et al. , 2017 ; Santurkar et al. , 2018 ; Foret et al. , 2021 ; Chen et al. , 2022 ) . Thus , they improve not only accuracy but also robustness in large data regimes . On the other hand , MSAs allow negative Hessian eigenvalues in small data regimes . This means that the loss landscapes of MSAs are non-convex , and this non-convexity disturbs NN optimization ( Dauphin et al. , 2014 ) . Large amounts of training data suppress negative eigenvalues and convexify losses . 2 Do MSAs act like Convs ? We show that MSAs and Convs exhibit opposite behaviors . MSAs aggregate feature maps , but Convs diversify them . Moreover , as shown in Fig . 2a , the Fourier analysis of feature maps shows that MSAs reduce high-frequency signals , while Convs , conversely , amplifies high-frequency components . In other words , MSAs are low-pass filters , but Convs are high-pass filters . In addition , Fig . 2b indicates that Convs are vulnerable to high-frequency noise but that MSAs are not . Therefore , MSAs and Convs are complementary . 3 How can we harmonize MSAs with Convs ? We reveal that multi-stage NNs behave like a series connection of small individual models . Thus , applying spatial smoothing at the end of a stage improves accuracy by ensembling transformed feature map outputs from each stage ( Park & Kim , 2021 ) as shown in Fig . 3a . Based on this finding , we propose an alternating pattern of Convs and MSAs . NN stages using this design pattern consists of a number of CNN blocks and one ( or a few ) MSA block as shown in Fig . 3c . The design pattern naturally derives the structure of canonical Transformer , which has one MSA block per MLP block as shown in Fig . 3b . It also provides an explanation of how adding Convs to Transformer ’ s MLP block improves accuracy and robustness ( Yuan et al. , 2021 ; Guo et al. , 2021 ; Mao et al. , 2021 ) . Surprisingly , models using this alternating pattern of Convs and MSAs outperform CNNs not only on large datasets but also on small datasets , such as CIFAR . This contrasts with canonical ViTs , models that perform poorly on small amount of data . It implies that MSAs are generalized spatial smoothings that complement Convs , not simply generalized Convs . 2 WHAT PROPERTIES OF MSAS DO WE NEED TO IMPROVE OPTIMIZATION ? . To understand the underlying nature of MSAs , we investigate the properties of the ViT family : e.g. , vanilla ViT ( Dosovitskiy et al. , 2021 ) ; PiT ( Heo et al. , 2021 ) , which is “ ViT + multi-stage ” ; and Swin ( Liu et al. , 2021 ) , which is “ ViT + multi-stage + local MSA ” . This section shows that these additional inductive biases enable ViTs to learn strong representations . We also use ResNet ( He et al. , 2016a ) for comparison . NNs are trained from scratch with DeiT-style data augmentation ( Touvron et al. , 2021 ) for 300 epochs . The NN training begins with a gradual warmup ( Goyal et al. , 2017 ) for 5 epochs . Appendix A provides more detailed configurations and background information for experiments . The stronger the inductive biases , the stronger the representations ( not regularizations ) . Do models with weak inductive biases overfit training datasets ? To address this question , we provide two criteria on CIFAR-100 : the error of the test dataset and the cross-entropy , or the negative log-likelihood , of the training dataset ( NLLtrain , the lower the better ) . See Fig . 5a for the results . Contrary to our expectations , experimental results show that the stronger the inductive bias , the lower both the test error and the training NLL . This indicates that ViT does not overfit training datasets . In addition , appropriate inductive biases , such as locality constraints for MSAs , helps NNs learn strong representations . We also observe these phenomena on CIFAR-10 and ImageNet as shown in Fig . C.1 . Figure C.2 also supports that weak inductive biases disrupt NN training . In this experiment , extremely small patch sizes for the embedding hurt the predictive performance of ViT . ViT does not overfit small training datasets . We observe that ViT does not overfit even on smaller datasets . Figure 5b shows the test error and the training NLL of ViT on subsampled datasets . In this experiment , as the size of the dataset decreases , the error increases as expected , but surprisingly , NLLtrain also increases . Thanks to the strong data augmentation , ViT does not overfit even on a dataset size of 2 % . This suggests that ViT ’ s poor performance in small data regimes is not due to overfitting . ViT ’ s non-convex losses lead to poor performance . How do weak inductive biases of MSAs disturb the optimization ? A loss landscape perspective provides an explanation : the loss function of ViT is non-convex , while that of ResNet is strongly ( near- ) convex . This poor loss disrupts NN training ( Dauphin et al. , 2014 ) , especially in the early phase of training ( Jastrzebski et al. , 2020 ; 2021 ) . Figure 1c and Fig . 4 provide top-5 largest Hessian eigenvalue densities ( Park & Kim , 2021 ) with a batch size of 16 . The figures show that ViT has a number of negative Hessian eigenvalues , while ResNet only has a few . Figure 4 also shows that large datasets suppress negative Hessian eigenvalues in the early phase of training . Therefore , large datasets tend to help ViT learn strong representations by convexifying the loss . ResNet enjoys little benefit from large datasets because its loss is convex even on small datasets . Loss landscape smoothing methods aids in ViT training . Loss landscape smoothing methods can also help ViT learn strong representations . In classification tasks , global average pooling ( GAP ) smoothens the loss landscape by strongly ensembling feature map points ( Park & Kim , 2021 ) . We demonstrate how the loss smoothing method can help ViT improve performance by analyzing ViT with GAP classifier instead of CLS token on CIFAR-100 . Figure 6 shows the Hessian max eigenvalue spectrum of the ViT with GAP . As expected , the result shows that GAP classifier suppresses negative Hessian max eigenvalues , suggesting that GAP convexify the loss . Since negative eigenvalues disturb NN optimization , GAP classifier improve the accuracy by +2.7 percent point . Likewise , Sharpness-Aware Minimization ( SAM ) ( Foret et al. , 2021 ) , an optimizer that relies on the local smoothness of the loss function , also helps NNs seek out smooth minima . Chen et al . ( 2022 ) showed that SAM improves the predictive performance of ViT . MSAs flatten the loss landscape . Another property of MSAs is that they reduces the magnitude of Hessian eigenvalues . Figure 1c and Fig . 4 shows that the eigenvalues of ViT are significantly smaller than that of CNNs . While large eigenvalues impede NN training ( Ghorbani et al. , 2019 ) , MSAs can help NNs learn better representations by suppressing large Hessian eigenvalues . Figure 1a and Fig . 1b also support this claim . In Fig . 1a , we visualize the loss landscapes by using filter normalization ( Li et al. , 2018 ) . The loss landscape of ViT is flatter than that of ResNet , and this trend is noticeable at the boundary . Similarly , Fig . 1b shows that ResNet follows an irregular trajectory , especially in the early phase of training ; ViT converges to the optimum along a smooth trajectory . In large data regimes , the negative Hessian eigenvalues—the disadvantage of MSAs—disappears , and only their advantages remain . As a result , ViTs outperform CNNs on large datasets , such as ImageNet and JFT ( Sun et al. , 2017 ) . PiT and Swin also flatten the loss landscapes . See Fig . C.4 . A key feature of MSAs is data specificity ( not long-range dependency ) . The two distinguishing features of MSAs are long-range dependency and data specificity , also known as data dependency , as discussed in Section 1.1 . Contrary to popular belief , the long-range dependency hinders NN optimization . To demonstrate this , we analyze convolutional ViT , which consists of two-dimensional convolutional MSAs ( Yang et al. , 2019 ) instead of global MSAs . Convolutional MSAs calculates self-attention only between feature map points in convolutional receptive fields after unfolding the feature maps in the same way as convolutions . Figure 7a shows the error and NLLtrain of convolutional ViTs with kernel sizes of 3× 3 , 5× 5 , and 8×8 ( global MSA ) on CIFAR-100 . In this experiment , 5×5 kernel outperforms 8×8 kernel on both the training and the test datasets . NLLtrain of 3× 3 kernel is worse than that of 5× 5 kernel , but better than that of global MSA . Although the test accuracies of 3× 3 and 5× 5 kernels are comparable , the robustness of 5×5 kernel is significantly better than that of 3×3 kernel on CIFAR-100-C ( Hendrycks & Dietterich , 2019 ) . Figure 7b shows that the strong locality inductive bias not only reduce computational complexity as originally proposed ( Liu et al. , 2021 ) , but also aid in optimization by convexifying the loss landscape . 5× 5 kernel has fewer negative eigenvalues than global MSA because it restricts unnecessary degrees of freedom . 5× 5 kernel also has fewer negative eigenvalues than 3× 3 kernel because it ensembles a larger number of feature map points ( See also Fig . 6 ) . The amount of negative eigenvalues is minimized when these two effects are balanced . It is clear that data specificity improves NNs . MLP-Mixer ( Tolstikhin et al. , 2021 ; Yu et al. , 2021a ) , a model with an MLP kernel that does not depend on input data , underperforms compared to ViTs . Data specificity without self-attention ( Bello , 2021 ) improves performance .

Summary: The authors show that the robustness and better performance of ViTs is attributed to their property of flattening the loss surface. Insights are taken from a parallel ICLR submission, where spatial smoothing has been shown to smooth loss landscape, and help improve robustness. The paper shows that the properties of spatial smoothing are exhibited by MSA as well.