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: This paper alleviate the dependency on massive pre-training and data augmentation. It promotes the smoothness using a recently proposed sharpness-aware optimizer to improve the accuracy and robustness of the ViTs and MLP-Mixers. This paper has demonstrated that the sharpness-aware optimizer could be leveraged to boost the performance of ViTs and MLP-Mixers without pretraining and strong augmentaitons on different tasks including supervised, adversarial, contrastive, and transfer learning.