When Vision Transformers Outperform ResNets without Pre-training or Strong Data Augmentations

1 INTRODUCTION . Transformers ( Vaswani et al. , 2017 ) have become the de-facto model of choice in natural language processing ( NLP ) ( Devlin et al. , 2018 ; Radford et al. , 2018 ) . In computer vision , there has recently been a surge of interest in end-to-end Transformers ( Dosovitskiy et al. , 2021 ; Touvron et al. , 2021b ; Liu et al. , 2021b ; Fan et al. , 2021 ; Arnab et al. , 2021 ; Bertasius et al. , 2021 ; Akbari et al. , 2021 ) and MLPs ( Tolstikhin et al. , 2021 ; Touvron et al. , 2021a ; Liu et al. , 2021a ; Melas-Kyriazi , 2021 ) , prompting the efforts to replace hand-wired features or inductive biases with general-purpose neural architectures powered by data-driven training . We envision these efforts may lead to a unified knowledge base that produces versatile representations for different data modalities , simplifying the inference and deployment of deep learning models in various application scenarios . Despite the appealing potential of moving toward general-purpose neural architectures , the lack of convolution-like inductive bias also challenges the training of vision Transformers ( ViTs ) and MLPs . When trained on ImageNet ( Deng et al. , 2009 ) with the conventional Inception-style data preprocessing ( Szegedy et al. , 2016 ) , Transformers “ yield modest accuracies of a few percentage points below ResNets of comparable size ” ( Dosovitskiy et al. , 2021 ) . To boost the performance , existing works resort to large-scale pre-training ( Dosovitskiy et al. , 2021 ; Arnab et al. , 2021 ; Akbari et al. , 2021 ) and repeated strong data augmentations ( Touvron et al. , 2021b ) , resulting in excessive demands of data , computing , and sophisticated tuning of many hyperparameters . For instance , Dosovitskiy et al . ( Dosovitskiy et al. , 2021 ) pre-train ViTs using 304M labeled images , and Touvron et al . ( 2021b ) repeatedly stack four strong image augmentations . In this paper , we show ViTs can outperform ResNets ( He et al. , 2016 ) of even bigger sizes in both accuracy and various forms of robustness by using a principled optimizer , without the need for largescale pre-training or strong data augmentations . MLP-Mixers ( Tolstikhin et al. , 2021 ) also become on par with ResNets . We first study the architectures fully trained on ImageNet from the lens of loss landscapes and draw the following findings . First , visualization and Hessian matrices of the loss landscapes reveal that Transformers and MLP-Mixers converge at extremely sharp local minima , whose largest principal curvatures are almost an order of magnitude bigger than ResNets . Such effect accumulates when the gradients backpropagate from the last layer to the first , and the initial embedding layer suffers the largest eigenvalue of the corresponding sub-diagonal Hessian . Second , the networks all have very small training errors , and MLP-Mixers are more prone to overfitting than ViTs of more parameters ( because of the difference in self-attention ) . Third , ViTs and MLP-Mixers have worse “ trainabilities ” than ResNets following the neural tangent kernel analyses ( Xiao et al. , 2020 ) . Therefore , we need improved learning algorithms to prevent the convergence to a sharp local minima when it comes to the convolution-free ViTs and and MLP-Mixers . The first-order optimizers ( e.g. , SGD and Adam ( Kingma & Ba , 2015 ) ) only seek the model parameters that minimize the training error . They dismiss the higher-order information such as flatness that correlates with the generalization ( Keskar et al. , 2017 ; Kleinberg et al. , 2018 ; Jastrzębski et al. , 2019 ; Smith & Le , 2018 ; Chaudhari et al. , 2017 ) . The above study and reasoning lead us to the recently proposed sharpness-aware minimizer ( SAM ) ( Foret et al. , 2021 ) that explicitly smooths the loss geometry during model training . SAM strives to find a solution whose entire neighborhood has low losses rather than focus on any singleton point . We show that the resultant models exhibit smoother loss landscapes , and their generalization capabilities improve tremendously across different tasks including supervised , adversarial , contrastive , and transfer learning ( e.g. , +5.3 % and +11.0 % top-1 accuracy on ImageNet for ViTB/16 and Mixer-B/16 , respectively , with the simple Inception-style preprocessing ) . The enhanced ViTs achieve better accuracy and robustness than ResNets of similar and bigger sizes when trained from scratch on ImageNet , without large-scale pre-training or strong data augmentations . Moreover , we demonstrate that SAM can even enable the ViT to be effectively trained with ( momentum ) SGD , which usually lies far behind Adam when training Transformers ( Zhang et al. , 2020 ) . By analyzing some intrinsic model properties , we observe that SAM increases the sparsity of active neurons ( especially for the first few layers ) , which contribute to the reduced Hessian eigenvalues . The weight norms increase , implying the commonly used weight decay may not be an effective regularization alone . A side observation is that , unlike ResNets and MLP-Mixers , ViTs have extremely sparse active neurons ( see Figure 2 ( right ) ) , revealing the potential for network pruning ( Akbari et al. , 2021 ) . Another interesting finding is that ViTs ’ performance gain also translates to visually more interpretable attention maps . Finally , we draw similarities between SAM and strong augmentations ( e.g. , mixup ) in that they both smooth the average loss geometry and encourage the models to behave linearly between training images . 2 BACKGROUND AND RELATED WORK . We briefly review ViTs , MLP-Mixers , and some related works in this section . Dosovitskiy et al . ( 2021 ) show that a pure Transformer architecture ( Vaswani et al. , 2017 ) can achieve state-of-the-art accuracy on image classification by pre-training it on large datasets such as ImageNet-21k ( Deng et al. , 2009 ) and JFT-300M ( Sun et al. , 2017 ) . Their vision Transformer ( ViT ) is a stack of residual blocks , each containing a multi-head self-attention , layer normalization ( Ba et al. , 2016 ) , and a MLP layer . ViT first embeds an input image x ∈ RH×W×C into a se- quence of features z ∈ RN×D by applying a linear projection overN nonoverlapping image patches xp ∈ RN× ( P 2·C ) , where D is the feature dimension , P is the patch resolution , and N = HW/P 2 is the sequence length . The self-attention layers in ViT are global and do not possess the locality and translation equivariance of convolutions . ViT is compatible with the popular architectures in NLP ( Devlin et al. , 2018 ; Radford et al. , 2018 ) and , similar to its NLP counterparts , requires pretraining over massive datasets ( Dosovitskiy et al. , 2021 ; Akbari et al. , 2021 ; Arnab et al. , 2021 ) or strong data augmentations ( Touvron et al. , 2021b ) . Some works specialize the ViT architectures for the visual data ( Liu et al. , 2021b ; Yuan et al. , 2021 ; Fan et al. , 2021 ; Bertasius et al. , 2021 ) . More recent works find that the self-attention in ViT is not vital for performance , resulting in several architectures exclusively based on MLPs ( Tolstikhin et al. , 2021 ; Touvron et al. , 2021a ; Liu et al. , 2021a ; Melas-Kyriazi , 2021 ) . Here we take MLP-Mixer ( Tolstikhin et al. , 2021 ) as an example . MLP-Mixer shares the same input layer as ViT ; namely , it partitions an image into a sequence of nonoverlapping patches/tokens . It then alternates between token and channel MLPs , where the former allows feature fusion from different spatial locations . We focus on ViTs and MLP-Mixers in this paper . We denote by “ S ” and “ B ” the small and base model sizes , respectively , and by an integer the image patch resolution . For instance , ViT-B/16 is the base ViT model taking as input a sequence of 16×16 patches . Appendices contain more details . 3 VITS AND MLP-MIXERS CONVERGE TO SHARP LOCAL MINIMA . The current training recipe of ViTs , MLP-Mixers , and related convolution-free architectures relies heavily on massive pre-training ( Dosovitskiy et al. , 2021 ; Arnab et al. , 2021 ; Akbari et al. , 2021 ) or a bag of strong data augmentations ( Touvron et al. , 2021b ; Tolstikhin et al. , 2021 ; Cubuk et al. , 2019 ; 2020 ; Zhang et al. , 2018 ; Yun et al. , 2019 ) . It highly demands data and computing , and leads to many hyperparameters to tune . Existing works report that ViTs yield inferior accuracy to the ConvNets of similar size and throughput when trained from scratch on ImageNet without the combination of those advanced data augmentations , despite using various regularization techniques ( e.g. , large weight decay , Dropout ( Srivastava et al. , 2014 ) , etc. ) . For instance , ViT-B/16 ( Dosovitskiy et al. , 2021 ) gives rise to 74.6 % top-1 accuracy on the ImageNet validation set ( 224 image resolution ) , compared with 78.5 % of ResNet-152 ( He et al. , 2016 ) . Mixer-B/16 ( Tolstikhin et al. , 2021 ) performs even worse ( 66.4 % ) . There also exists a large gap between ViTs and ResNets in robustness tests ( e.g. , against 19 corruptions in ImageNet-C ( Hendrycks & Dietterich , 2019 ) ) ( see Table 2 for ImageNet-R and ImageNet-C results ) . Moreover , Chen et al . ( 2021c ) find that the gradients can spike and cause a sudden accuracy dip when training ViTs , and Touvron et al . ( 2021b ) report the training is sensitive to initialization and hyperparameters . These all point to optimization problems . In this paper , we investigate the loss landscapes of ViTs and MLP-Mixers to understand them from the optimization perspective , intending to reduce their dependency on the large-scale pre-training or strong data augmentations . ViTs and MLP-Mixers converge to extremely sharp local minima . It has been extensively studied that the convergence to a flat region whose curvature is small benefits the generalization of neural networks ( Keskar et al. , 2017 ; Kleinberg et al. , 2018 ; Jastrzębski et al. , 2019 ; Chen & Hsieh , 2020 ; Smith & Le , 2018 ; Zela et al. , 2020 ; Chaudhari et al. , 2017 ) . Following Li et al . ( 2018 ) , we plot the loss landscapes at convergence when ResNets , ViTs , and MLP-Mixers are trained from scratch on ImageNet with the basic Inception-style preprocessing ( Szegedy et al. , 2016 ) ( see Appendices for details ) . As shown in Figures 1 ( a ) to 1 ( c ) , ViTs and MLP-Mixers converge to much sharper regions than ResNets . Besides , we calculate the training error under Gaussian perturbations on the model parameters LNtrain = E ∼N [ Ltrain ( w + ) ] in Table 1 , which reveals the average flatness . Although ViT-B/16 and Mixer-B/16 achieve lower Ltrain than ResNet-152 , their loss values after random weight perturbation become much higher . We further validate the results by computing the dominate Hessian eigenvalue λmax , which is a mathematical evaluation of the worst-case landscape curvature . The λmax values of ViT and MLP-Mixer are orders of magnitude larger than that of ResNet , and MLP-Mixer suffers the largest curvature among the three species ( see Section 4.4 for a detailed analysis ) . Small training errors . This convergence to sharp regions coincides with the training dynamics shown in Figure 2 ( left ) . Although Mixer-B/16 has fewer parameters than ViT-B/16 ( 59M vs. 87M ) , it has a smaller training error ( also seeLtrain in Table 1 ) but much worse test accuracy , implying that using the cross-token MLP to learn the interplay across image patches is more prone to overfitting than ViTs ’ self-attention mechanism whose behavior is restricted by a softmax . To validate this statement , we simply remove the softmax in ViT-B/16 , such that the query and key matrices can freely interact with each other . Although having lower Ltrain ( 0.56 vs. 0.65 ) , the obtained ViTB/16-Free performs much worse than the original ViT-B/16 ( 70.5 % vs. 74.6 % ) . Its LNtrain and λmax are 7.01 and 1236.2 , revealing that ViT-B/16-Free converges to a sharper region than ViTB/16 ( LNtrain is 6.66 and λmax is 738.8 ) both on average and in the worst-case direction . Such a difference probably explains why it is easier for MLP-Mixers to get stuck in sharp local minima . ViTs and MLP-Mixers have worse trainability . Furthermore , we discover that ViTs and MLPMixers suffer poor trainability , defined as the effectiveness of a network to be optimized by gradient descent ( Xiao et al. , 2020 ; Burkholz & Dubatovka , 2019 ; Shin & Karniadakis , 2020 ) . Xiao et al . ( 2020 ) show that the trainability of a neural network can be characterized by the condition number of the associated neural tangent kernel ( NTK ) , Θ ( x , x′ ) = J ( x ) J ( x′ ) T , where J is the Jacobian matrix . Denoting by λ1 ≥ · · · ≥ λm the eigenvalues of NTK Θtrain , the smallest eigenvalue λm converges exponentially at a rate given by the condition number κ = λ1/λm . If κ diverges then the network will become untrainable ( Xiao et al. , 2020 ; Chen et al. , 2021a ) . As shown in Table 1 , κ is pretty stable for ResNets , echoing previous results that ResNets enjoy superior trainability regardless of the depth ( Yang & Schoenholz , 2017 ; Li et al. , 2018 ) . However , we observe that the condition number diverges when it comes to ViT and MLP-Mixer , confirming that the training of ViTs desires extra care ( Chen et al. , 2021c ; Touvron et al. , 2021b ) .

Summary: As the success of Vision Transformer ViT has shown its potentials for computer vision tasks, this paper investigates a more effective way of training a ViT understand a standard ImageNet pre-training setting such as no extra training data and no strong data augmentation.  In general, without those conditions, a typical ViT can not perform as good as widely convolutional based network architectures such as ResNet.  This paper addresses the issue from the point of loss landscape and then propose to use a better optimization strategy named Sharp-aware minimizer (SAM) for ViT related architecture optimization.  With the proposed SAM, ViT can achieve better accuracy significantly understand standard ImageNet training/testing protocol.  In addition to the improved performance on ImageNet, this paper further shows the visualization of the attention map and the improved performance on other applications such as contrastive learning and adversarial training.