When Vision Transformers Outperform ResNets without Pre-training or Strong Data Augmentations
Authors: Xiangning Chen
Cho-Jui Hsieh
Boqing Gong
Venue: ICLR 2022 Spotlight (2022)
Keywords: Vision Transformers
Optimization

Abstract: Vision Transformers (ViTs) and MLPs signal further efforts on replacing hand-wired features or inductive biases with general-purpose neural architectures. Existing works empower the models by massive data, such as large-scale pre-training and/or repeated strong data augmentations, and still report optimization-related problems (e.g., sensitivity to initialization and learning rates). Hence, this paper investigates ViTs and MLP-Mixers from the lens of loss geometry, intending to improve the models' data efficiency at training and generalization at inference. Visualization and Hessian reveal extremely sharp local minima of converged models. By promoting smoothness with a recently proposed sharpness-aware optimizer, we substantially improve the accuracy and robustness of ViTs and MLP-Mixers on various tasks spanning supervised, adversarial, contrastive, and transfer learning (e.g., +5.3\% and +11.0\% top-1 accuracy on ImageNet for ViT-B/16 and Mixer-B/16, respectively, with the simple Inception-style preprocessing). We show that the improved smoothness attributes to sparser active neurons in the first few layers. The resultant ViTs outperform ResNets of similar size and throughput when trained from scratch on ImageNet without large-scale pre-training or strong data augmentations. Model checkpoints are available at \url{https://github.com/google-research/vision_transformer}.

Published as a conference paper at ICLR 2022
WHEN VISION TRANSFORMERS OUTPERFORM RESNETS WITHOUT TRAINIGGNING OR STRONG DATA AUGMENTATIONS
Xiangning Chen Cho-Jui Hsieh", Boqing Gong'
Google Research, "Department of Computer Science UCLA
xi xxiangning,, bgong) 8google.com. chohsich@cs wcla a.aaedd edu
ABSTRACT
Vision Transformers (ViTs) and ML.Ps signal further efforts on replacing hand- wired features or inductive biases with ceneral-purpose neural architectures. Ex- isting works empower the models by massive data, such as large-scale pre-training and/or repeated strong data augmentations, and still report optimization-relate problems (e.g., sensitivity to initialization and learning rates). Hence, this pa- per investigates ViTs and MLP-Mixers from the lens of loss geometry, intend- ing to improve the models' data efficiency at training and generalization at in- ference. Visualization and Hessian reveal extremely sharp local minima of con- verged models. By promoting smoothness with a recently proposed sharpness- aware optimizer, we substantially improve the accuracy and robustness of ViTs and MLP-Mixers on various tasks spanning supervised, adversarial, contrastive, and transfer learning (e.g., +5.3% and +11.0% top-1 accuracy on ImageNet for ViT-B/16 and Mixer-B/16. respectively, with the simple Inception-style prepro- cessing). We show that the improved smoothness attributes to sparser active neu- rons in the first few layers. The resultant ViTs outperform ResNets of similar size and throughput when trained from scratch on ImageNet without large-scale pre- pre-training or strong data augmentations. Model checkpoints available at httpps:/ /tttps://githb... com ccm/gooo e resea r ch / v s toiiio_____rreeree trans former
INTRODUCTION
Transformers (Vaswani 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 sodd-o--nd
Transformers (Dosovitskiy al., 2021 Touvron et al. 2021b
Liu et al. 2021b Fan et al. 2021
Amab et al.] [202]] Bertasius et al. 202 Akbari et al. 2021
and MLPs (Tolstikhin et al.
2021 Touvron al.) al.[[221: Liu et al.] [021a]] Melas- Melas-Kyriazi 2021).
prompting the efforts to replace hand-wired features or inductive biases with general-purpon
neu-
ral 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 biases 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 al. (002T) To boost the performance, existing works resort to large-scale pre-training (Dosovitskiy et al. 2021 Arnab et al. [202]] Akbari al. 2021) and repeated strong data augmentations (Touvron
exces- al. 2021b). resulting sive demands of data, computing, and sophisticated tuning of many hyperparameters. For instance, Dosovitskiy et al. (Dosovitskiy et al. (022]) pre-train ViTs using 304M labeled images, and Touvron et al. b) repeatedly stack four strong image augmentations.
WWrk done student researcher at Google.

---

Published as a conference paper at ICLR 2022
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 large- scale pre- training or strong data augmentations. MLP-Mixers (Tolstikhin et al. 202I) 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 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 MI.P-Mixers are more prone to overfitting than ViTs of more pa- rameters (because of the difference in self-attention). Third, ViT's 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 sharp local min- imum when comes to the convolution-fres ViTs and MLP-Mixers. The first-order optimizers (e.g., SGD and Adam (Kingma & Ba. 2013)) only seek the model parameters that minimize the training error. They dismiss the higher-order information such as flatness that correlates with gen- eralization (Keskar et al. 2017: Kleinberg et al. 2018 Jastrzębski et al. 2019 Smith & Le. 2018 Chaudhari et 2017).
The above study and reasoning lead us to the recently proposed sharpness-awan minimizer (SAM) (Foret et al., (222)) that explicitly smooths the loss geometry during model training. SAM strives to find solution whose entire neighborhood has low losses rather than focus on any single- ton point. We show that the resultant models exhibit smoother loss landscapes, and their general- ization capabilities improve tremendously across different tasks including supervised, adversarial, contrastive, and transfer learning (c.g., +5.3% and +1 +11.0% top-1 accuracy on ImageNet for ViT-B/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 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 (cspecially 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 reg- ularization alone. A side observation is that, unlike ResNets and MI.--Mieers ViT's have extremely sparse active neurons (see Figure 2 (right)), revealing the potential for network pruning (Akbari jet al. 2021). Another interesting finding is that the improved ViTs appear to have 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 al. (2021) show that a pure Transformer architecture (Vaswani al. 2017) can achieve of 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 normaliza- tion (Ba et al. 2016). and MLP layer. ViT first embeds an input image x E " WXC XC into a se- quence of features E R W by applying linear projection over N nonoverlapping image patches Ip E (P2--C) where D is the feature dimension, P is the patch resolution, and N HW/P" p2 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 pre- training over massive datasets (Dosovitskiy et al. 2021
Akbari et al. 2021 Arnab al. 2021) or
2

---

Published as a conference paper at ICLR 2022
Table 1: Number of parameters, NTK condition number K, Hessian dominate eigenvalue Amaz. training еттоr at convergence Ltrain, average flatness Ly accuracy on ImageNet, and accu- racy/robustness on ImageNet-C. ViT and MLP-Mixer suffer divergent and converge at sharp re- gions; SAM rescues that and leads to better generalization.
155 SAM
VIT-B/I6- SAM
Mixer-B/16- SAM
ResNet-152
VIT-B/16
Mixer-B/16
#Params
60M
87M
59M
NTK &
2801.6
4205.3
14468.0
Hesslan Amar
179.8
42.0
738.8
20.9
1644.4
22.5
Laguen
0.86
0.90
0.65
0.82
0.45
0.97
Darain
2.39
2.16
6.66
0.96
7.78
1.01
ImageNet (%)
78.5
79.3
74.6
79.9
66.4
77.4
ImageNet-C (%)
50.0
52.2
46.6
56.5
33.8
48.8
1 As prohibitive compute exact NTK, approximate the value by averaging over its sub- diagonal blocks (sec Appendix G for details). We averuge the results for 1,000 random noises when calculating L.
(a) ResNet
(b) VIT
(c) Mixer
(d) VIT-SAM
(e) Mixer-SAM
Figure I: Cross-entropy loss landscapes of R55 ViT-B/16, and Mixer-B/16. ViT and MLP- Mixer converge sharper regions than ResNet when trained on ImageNet with the basic Inception- style preprocessing, SAM, sharpness-awan optimizer, significantly smooths the landscapes.
strong data augmentations (Tuvrron et al. 2021b). Some works specialize the ViT architectures for visual data (Liu et al. 2021b Yuan al. 2021 Fan et al. 2021 Bertasius et al. 2022)).
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. 202 Liu et al. 2021a Melas-Kyriazi 202T). Here we take MLP-Mixer (Tolstikhin et al., (221) as an example. MLP-Mixer shares the same input layer VIT; namely, partitions image into sequence of nonoverlapping patches/tokens. 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 sequence of 16 x 16 patches. Appendices contain more details.
3 VITS AND MLP-MIXERS CONVERGE AT SHARP LOCAL MINIMA
The current training recipe of ViTs,, MLP-Mixers, and related convolution- free architectures relies heavily on massive pre-training (Dosovisskiy al. [021 Arnab et al. al.| 2021 Akbari et al. [222]) or bag of strong data augmentations (Touvron et al. 202 Tolstikhin et al. | Cubuk et al. 2019 2020 Zhang al.) 2018 Yun al. 2019). lighly demands data and computing, and leads many hyperparameters to tune. Existing works report that ViTs yield inferior accuracy to the ConvNets of similar size and throughput when trained from seratch on ImageNet without the combination of those advanced data augmentations, despite using various regularization techniques (e.g., large weight decay, Dropout (Srivastava al.) 2014), etc.). For instance, ViT-B/16 (Dosovitskiy et al.) 2021) gives rise to 74.6% top-1 accuracy the ImageNet validation set (224 image resolution), compared with 78.5% ResNet- (He et al. (2016). Mixer-B/16 (Tolstikhin et al. (221)) performs worse (66.4%). There also exists large gap between ViTs and ResNets in robustness tests (see Table 2 for details).
Moreover, Chen et al. 2021c) find that the gradients spike and cause sudden accuracy dip when training ViTs, and Touvron et al. (021b) report the training is sensitive to initialization and hyperparameters. These all point optimization problems. In this paper, investigate the loss
3

---

Published as a conference paper at ICLR 2022
IT3CE
15
4atho Heages
Bexzann
Figure 2: Left and Middle: ImageNet training error and validation accuracy vs. iteration for ViTs and MLP-Mixers. Right: Percentage of active neurons for ResNet-152 ViT-B/16, and Mixer-B/1e
landscapes of ViTs and MLP-Mixers to understand them from the optimization perspective, intend- ing to reduce their dependency on the large-scale pre-training or strong data augmentations.
VITs and MI. converge at extremely sharp local minima. It has been extensively studied
that the convergence to a flat region whose curvature small benefits the generalization of neural
networks (Keskar et al., 2017: Kleinberg et al. 2018 Jastrzębski et al. 2019 Chen & Hsich 2020
Smith & Le 2018 Zela et al. Chaudhhaanhar et al. 2017). 7). Following Li et al. (2018). we plot the
loss landscapes convergence when ResNets, ViTs, and MLP-Mixers are trained from scratch
ImageNet with the basic Inception- -style preprocessing (Szegedy et al. 2016) (see Appendices for
details). As shown in Figures ((a) to [(e] ViTs and MLP-Mixers converge much sharper regions than ResNets, Besides, we calculate the training error under Gaussian perturbations on the model parameters L LEewW[Lrraiie (w (w+++) e)] in Table which reveals the average flatness. Although ViT-B/16 and Mixer-B/16 achieve lower training eπror Ltrain than that of ResNet-152, their loss values after random weight perturbation become much higher. We further validate the results by computing the dominate Hessian eigenvalue Amaz. which is mathematical evaluation of the worst- case landscape curvature. The Amaz 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 sharp regions coincides with the training dynamics shown in Figure (left. Although Mixer-B/16 has fewer parameters than ViT-B/16 (59M vs. 87M), it has smaller training error (also Lirain in Table 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 Lirain (0.56 vs. 0.65), the obtained ViT- B/16-Free performs much worse than the original ViT-B/16 (70.5% vs. 74.6%). Its L and Amaz are 7.01 and 1236.2, revealing that ViT-B/16-Free converges to sharper region than ViT- B/16 train is 6.66 and \max is 738.8) both on average and in the worst-case direction. Such difference probably explains why it is easier for MLP-Mixers get stuck in sharp local minima.
VITs and MLP-Mixers have worse trainability. Furthermore, we discover that ViTs and MLP- Mixers suffer poor trainabilities, defined as the effectiveness of a network to be optimized by gradi- ent descent (Xiao et al. [220 Burkholz & Dubatovka 2019: [Shin & Karniadakis 2020). Xiao et al.) neural network can be characterized by the condition number of the associated neural tangent kernel (NTK), Oz x²) (zz)) where J is the Jacobian matrix. Denoting by A 2 ... 2 Am the cigenvalues converges exponentially at rate given by the condition number k >1/Am- If n. diverges then the network will become untrainable (Xiao et al.) 2020 Chen et al. 2021a). As shown in Table [ is pretty stable for ResNets, echoing previous results that ResNets enjoy superior trainability re- gardless of the depth (Yang & Schoenholz 2017 Li et al. 2018). However, we observe that the condition number diverges when it comes
(2200) show that the trainability
NTK Otrain the smallest eigenvalue Am
ViT and MLP P-xier, confirming that the training of ViTs desires extra care (Chen et al.) 2021c Touvron et al. 2021b).
4 A PRINCIPLED OPTIMIZER FOR CONVOLUTION-FRED ARCHITECTURES
The commonly used first-order optimizers (c.g., SGD (Nesterov) 1983), Adam Kingmaa & Ba 2015)) only seek to minimize the training loss Ltrrin (w). They usually dismiss the higher-order
4

---

Published as a conference paper at ICLR 2022
information such as curvature that correlates with the generalization (Keskar al. 2017 Chaud- hari et al. 2017 Dziugaite & Roy 2017). However, the objective Lerain for deep neural networks highly non-convex, making it easy reach near-zero training enror but high generalization ror Tasest during evaluation, let alone their robustness when the test sets have different distribu- tions (Hendrycks & Dietterich] 20T9: [Hendrycks al. 2020). ViT's and MLPs amplify such draw- backs of first-order optimizers due to the lack of inductive bias for visual data, resulting in exces- sively sharp loss landscapes and poor generalization, as shown the previous section. We hypothe- size that smoothing the loss landscapes at convergence can significantly improve the generalization ability of those convolution- free architectures, leading us to the recently proposed sharpness-awant minimizer (SAM) (Foret et al. 2021) that explicitly avoids sharp minima.
4.1
SAM: OVERVIEW
Intuitively, SAM Foret et al. 202T) seeks to find the parameter to whose entire neighbours have low training loss Lerain by formulating a minimax objective:
min max Ietrai (w c)2 |ellaSSp
where p is the size of the neighhourhood ball. Without loss of generality, here use lz norm for its strong empirical results (Foret et al. (202)) and omit the regularization term for simplicity. Since the exact solution of the inner maximization e* arg max Lerain (w + e) is hard obtain, they employ an efficient first-order approximation:
ε(w) arg max Ltrain (w) + E (..brrai(() pv aewerain (w) (w)) (2)ll2 1k
(2)
Under the l2 norm, ((w) is simply scaled gradient of the current weight w. After computing é SAM updates w based on the sharpness-ave aware gradient V Levain (w)lw-((w)-
4.2 SHARPNESS-AWAR OPTIMIZATION IMPROVES VITS AND MLP-MIXERS
We train ViTs and MLP-Mixers with no large-scal scale pre- pre-training or strong data augmentations. We directly apply SAM to the original ImageNet training pipeline of ViTs (Dosovitskiy et al.) (222)) without changing any hyperparameters. The pipeline employs the basic Inception-sty style preprocess- ing (Szegedy et al. (016).. The original training setup of MLP-Mixers (Tolstikhin et al. 2021) includes a combination of strong data augmentations, and we replace it with the same Inception- style preprocessing for a fair comparison. Note that perform grid search for the learning rate, weight decay, Dropout before applying SAM. Please see Appendices for training details.
Smoother regions around the local minima. Thanks to SAM, both ViTs and MLP-Mixers con- verge at much smoother regions, as shown in Figures (dd) and [(e] Moreover, both the average and the worst-case curvature, i.e., LN and Xmaz, decrease dramatically (see Table
Higher accuracy. What comes along is tremendously improved generalization performance. On ImageNet, SAM boosts the top-1 accuracy of V from 74.6% 79.9%, and Mixer-B/16 from 66.4% to 77.4%. For comparison, the improvement on a similarly sized ResNet-152 is 0.8%. Empirically, the degree of improvement negatively correlates with the constraints of inductive biases built into the architecture. ResNets with inherent translation equivalence and locality benefit less from landscape smoothing than the attention-based ViTs. MLP-Mixers gain the most from the smoothed loss geometry. In Table we further train two hybrid models (Dosovitskiy et al. 2021) to validate this observation, where the Transformer takes the feature map extracted from a ResNet-50 as the input sequence. The improvement brought by SAM decreases after we introduce the convolution to VIT, for instance, +2.7% for R50-B/16 compared to +5.3% for ViT-B/16. Moreover, SAM brings larger improvements to the models of larger capacity (c.g.. +4.1% for Mixer-S/16 vs. +1 1.0% for Mixer-B/16 and longer patch sequence (c.g., +2.1% for ViT-S/32 vs. +5.3% for ViT-S/8). Please see Table 2 for more results.
SAM can be easily applied to common base optimizers. Besides Adam, we also apply SAM on top of the (momentum) SGD that usually performs much worse than Adam when training Transform- ers (Zhang et al. 2020). As expected, we find that under the same training budget (300 epochs), the VTT-B/16 trained with SGD only achieves 71.5% accuracy on ImageNet, whereas Adam achieves
5

---

Published as a conference paper at ICLR 2022
Table 2: Performance of ResNets, ViTs, and MLP-Mixers trained from scratch on ImageNet with SAM (improvement over the vanilla model is shown in the parentheses). We use the Inception-style preprocessing (with resolution 224) rather than combination of strong data augmentations.
Throughput (img/sec/core)
Model
#params
ImageNet
ReaL
V2
ImageNet-R
ImageNet-C
ResNet
ResNet-50-AAM
25M
2161
76.7
83.1 (+0.7)
64.6(+1.0)
23.3 (+L.))
46.5
ResNet-101-SAr
44M
1334
78.6 (-0.8)
84.8 (40.9)
66.7 (+1.4)
(+1.5)
51.38)2.)
ResNet-15-SAAM
60M
935
79.3 (+0.8)
84.9 (40.7)
67.3
25.7 (+0.4)
(222))
ResNet-50x2-SAM
9SM
891
79.6
85.3 (41.6)
(+1.7)
(42.9)
50.7 (43.9)
ResNa--10122-AM
173M
519
80.9 (+2.4)
86.4 (+2.4)
69.1 (+2.8)
27.8 (+3.2)
(4.0
ResNe--15x2-SAM
236M
356
81.1 (+1.8)
86.4 (+1.9)
(42.3)
28.1 (42.8)
(442))
Vision Transformer
VTT-S/32-SM
23M
6888
70.5 (+2.1)
77.5
(42.6)
21.4
(42.9)
VIT-S/IG-SAM
22M
2043
78.1 (+3.7)
84.1 (43.7)
(5.69)
24.7 (+4.7)
(3.0
VIT-S/14-SA
22M
1234
78.8 (+4.0)
84.8 (44.5)
(45.2)
24.4 (44.7)
54.2
S/8
22M
333
81.3 (+5.3)
86.7 (+5.5)
70.4 (+6.2)
25.3 (+6.1)
55.6 (+8.5
VTT-B/32-SA
SSM
2805
73.6 (+4.1)
80.3 (45.1)
60.0 (+4.7)
24.0 (44.1)
50.7 (46.7)
87M
863
79.9 (+5.3)
85.2 (+5.4)
67.5 (+6.2)
26.4 (+6.3)
(+9.9)
MLP-Mixer
S/xer-S/32-AAM
19M
11401
66.7 (+2.8)
73.8 (+3.5)
52.4 (+2.9)
(+2.7)
39.3 (+4.1)
Mixer-S/16-SAM
18M
4005
72,9 (+4.1)
79.8 (44.7)
58.9 (44.1)
20.1 (14.2)
(46.4)
Mixer-S/8-Se SAM
S/ee--S8--AMM
75.9 (+5.7)
(46.3)
(26.2)
20.5 (+5.1)
Mixer-B/32-SAM
60M
4209
72.4
(-.0
(4.044
22.8 (48.2)
46.2
Mixer B/M6-AAM
59M
1390
(+110))
(+15
(+13.1)
(+10.)
48.8 (+15.0)
Mixer-B/8-SAM.
64M
466
(+10
(4.40-
(+11
(49.2)
48.9 (+16.9)
74.6%. Surprisingly, SGD + SAM can push the result to 79.1% 19,1% which is a huge +7.6% absolute improvement. Although Adam + SAM is still higher (79.9%), their gap largely shrinks.
Better robustness. We also evaluate the models' robustness using ImageNet-R (Hendrycks al. 2020) and ImageNet-C (Hendrycks & Dietterich 2019) and find even bigger impacts of the smoothed loss landscapes. On ImageNet-C, which corrupts images by noise, bad weather, blur, etc., we report the average accuracy against 19 corruptions across five levels. As shown in Ta- bles and 2 the accuracies of ViT-B/16 and Mixer-B/16 increase by 9.9% and 15.0% (which are 21.2% and 44.4% relative improvements), after SAM smooths their converged local regions. In comparison, SAM improves the accuracy ResNet-152 by 2.2% (4.4% relative improvement). We that SAM enhances the robustness even more than the relative clean accuracy improvements (7.1%, 16.6%, and 1.0% for ViT-B/16, Mixer- B/16, and ResNet- 152, respectively).
4.3 VITS OUTPERFORM RESNETS WITHOUT PRE- TRAINING OR STRONG AUGMENTATIONS
The performance of an architecture is often conflated with the training strategies (Bello ct al. 202I), where data augmentations play key role (Cubuk et al. 2019 2020 Zhang al. 2018 Xie et af. 2020) Chen et al. (002bb). However, the design of augmentations requires substantial domain expertise and may not trans-
Table 3: Accuracy and robustness of two hybrid architectures.
ImageNet (%)
ImageNet-CC (%)
Model
#params
R50-S/16
79.8
53.4
34M
R50-S/16-SAM
(+1.2)
57.2 (+3.8)
R50-B/16
79.7
54.4
late between images and videos, for instance. Thanks to the principled sharpness- mizer, we can remove the advanced augmentations and focus on the architectures themselves.
R50-B/16-SAM
99M
82.4742.7)
61.0 (+6.6)
opti-
When trained from scratch on ImageNet with SAM, ViTs outperform ResNets of similar and greater sizes (also comparable throughput at inference) regarding both clean accuracy (on ImageNet (Deng et al. 2009), ImageNet-Real. (Beyer et al.] (020). and ImageNet V2 (Recht et al. (019)) and robustness (on ImageNet-R (Hendrycks et al. 2020) and ImageNet-C (Hendrycks & Dietterich) 2019 ViT-B/16 achieves 79.9%, 26.4%, and 56.6% top-1 accuracy ImageNet, ImageNet-R, and ImageNet-C. while the counterpart numbers for ResNet-152 are 79.3%, 25.7%, and 52.2%, re- spectively (see Table The gaps between ViTs and ResNets
wider for small architectures. ViT-S/16 outperforms a similarly sized ResNet-50 by 1.4% on ImageNet, and 6.5% on ImageNet-C. SAM also significantly improves MI Mix--Mixers' results.

---

Published as a conference paper at ICLR 2022
Table Dominant eigenvalue Amas of the sub-diagonal Hessians for different network components, and norm of the model parameter w and the post-activation of block k. Each ViT block consists of MSA and MLP, and MLP-Mixer alternates between token MLP a channel MLP. Shallower layers have larger Xmaz.
SAM smooths every component.
Model
Amuu of disgonal blocks of Hessian
||il|=
||ax||
||ad
Jevil
Embedding
MSA Token MLP
Channel MLP
MLP/
Block1
Blockó
Block12
Whole
VIT-B/16
300.4
179.8
281.4
44.4
32.4
26.9
738.8
269.3
104.9
104.3
138.1
VIT-B/16-SA
20.9
353.8
117.0
97.2
Mixer D/16
1042.3
95.8
417.9
239.3
41.2
5.1
1644.4
197.6
96.7
135.1
74.9
Bixer-B/16-SAM
22.5
389.9
110.9
176.0
216.1
4.4 INTRINSIC CHANGES AFTER SAM
We take a deeper look into the models to understand how they intrinsically change to reduce the Hessian" eigenvalue Amaz and what the changes imply in addition to the enhanced generalization.
Smoother loss landscapes for every network component. In Table we break down the Hessian of the whole architecture into small diagonal blocks of Hessians concerning each set of parameters, attempting to analyze what specific components cause the blowing up of Amaz in the models trained without SAM. We observe that shallower layers have larger Hessian eigenvalues Amaz. and the first linear embedding layer incurs the sharpest geometry. This agrees with the finding in (Chen et al. 2021c) that spiking gradients happen early in the embedding layer. Additionally, the multi-head self-attention (MSA) in ViTs and the Token MLPs in MLP-Mixers, both of which mix information across spatial locations, have comparably lower Xmar than the other network components. SAM consistently reduces the Amaz of all network blocks.
We can gain insights into the above findings by the recursive formulation of Hessian matrices for MLPs (Botev et al., 2017). Let fy and ag be the pre-activation and post-activation values for layer k. respectively. They satisfy hyg Wgag=1 and ag = (hx), where W₁ is the weight matrix and fa is the activation function (GELU (Hendrycks & Gimpel, 2020) in MLP-Mixers). Here we omit the bias term for simplicity. The diagonal block of Hessian matrix Hg with respect to W can be recursively calculated as:
H (ak-1ak 11 Hi, Hes B&WT k+1 Ho+ Wh+1 Ba Dis
(3)
B₁ diag(
D2 = diag(S (ha)
(4)
Dax"
where @ is the Kronecker product, H. is the pre-activation Hessian for layer k. and L is the objective function. Therefore, the Hessian norm accumulates the recursive formulation backpropagates to shallow layers, explaining why the first block has much larger Amaz than the last block Table 4
Greater weight norms. After applying SAM, find that in most cases, the norm of the post- activation value (h-1 and the weight WK+1 become even bigger (see Table 4 indicating that the commonly used weight decay may effectively regularize ViT's and MLP-Mixers (see Appendix] [ for further verification when we vary the weight decay strength).
Sparser active neurons MLP-Mixers Given the recursive formulation Equation 3 we identify another intrinsic measure of MLP-Mixers that contribute to the Hessian: the number of activated neurons. Indeed, B₁ determined by the activated neurons whose values are greater than zero, since the first-order derivative of GELU becomes much smaller when the input is negative. As a result, the number of active GELU neurons directly connected to the Hessian norm. Figure g (right) shows the proportion of activated neurons for each block, counted using 10% of the ImageNet training set. We can see that SAM greatly reduces the proportion of activated neurons for the first few layers of the Mixer-B/16 16, pushing them to much sparser states. This result also suggests the potential redundancy of image patches.
VITs' active neurons are highly sparse. Although Equations (3) and 4 only involve MLPs, we still observe a decrease of activated neurons in the first layer of ViTs (but not as significant as in MiP-Mixers). More interestingly, we find that the proportion of active neurons in ViT is much smaller than another two architectures - given an input image, less than 10% neurons have values greater than zero for most layers (see (riht). In other words, ViTs offer huge potential for

---

Published as a conference paper at ICLR 2022
Table 5: Data augmentations, SAM, and their combination applied to different model architectures trained on ImageNet and its subsets from scratch.
ResNet-152
ViT-B/16
Mixer-B/16
Dataset
SAM + AUG
SAM + AUG
SAM
Vanilla
SAM
AUG
Vanilla
SAM
AUG
Vanilla
SAM
AUG
+ AUG
ImageNet
78.5
79.3
78.8
78.9
74.6
79.9
79.6
81.5
66.4
77.4
76.5
78.1
ilk (1/2)
74.2
75.6
75.1
75.5
64.9
75.4
73.1
75.8
53.9
71.0
70.4
73.1
ilk (1/4)
68.0
70.3
70.2
70.6
52.4
66.8
63.2
65.6
37.2
62.8
61.0
65.8
i1k (1/10)
57.1
59.5
32.8
46.1
38.5
45.7
43.5
51.0
Figure 3: Raw images (Left) and attention maps of ViT-S/16 with (Right) and without (Middle) sharpness-aware optimization.
network pruning. This sparsity may also explain why one Transformer can handle multi-mout modality signals (vision, text, and audio) (Akbari al. 202T
Visually improved attention maps in ViTs. We visualize ViT-S/16's attention map of the classifi- cation token averaged over the last multi-head attentions in Figure following Caron et al.) (202T). Interestingly, the ViT model optimized with SAM appears to possess visually improved attention map compared with the one trained via the vanilla AdamW optimizer.
4.5 SAM VS. STRONG AUGMENTATIONS
Previous sections show that SAM can improve the generalization (and robustness) of ViT's and MLP-Mixers. Meanwhile, another paradigm to train these models on ImageNet from scratch is to stack multiple strong augmentations (Touvron et al. 202 Tolstikhin et al. 2021). Hence, it is interesting to study the differences and similarities between the models trained by SAM and by using strong data augmentations. For the augmentation experiments, follow Tolstikhin al. (2021)ss pipeline that includes mixup (Zhang et al. 2018) and RandAugment (Cubuk et al. (020).
Generalization. Table the results of strong data augmentation, SAM, and their combination on ImageNet. Each row corresponds to a training set a different fraction of ImageNet-1k: SAM benefits ViT-B/16 and Mixer-B/16 more than the strong data augmentations, especially when the training set is small. For instance, when the training set contains only 1/10 of ImageNet training images, ViT-B/I6-SAM outperforms ViT-B/16-AUG by 7.6%. Apart from the improved validation accuracy, we also observe that both SAM and strong augmentations increase the training error (see Figure E (Middle) and Table 6 indicating their regularization effects. However, they have distinct training dynamics as the loss curve for ViT-B/16-AUG is much nosier than ViT-B/16-SAM.
Sharpness at convergence. Another intriguing question is as follows. Can augmentations also smooth the loss geometry similarly to SAM? To answer it, we also plot the landscape of VIT- B/16-AUG (see Figure 5 in the Appendix) and compute its Hessian Amaz together with the erage flatness L Table 6 Surprisingly. strong augmentations enlarge the Xmaz-
Table 6: Comparison between ViT-B/I6-SAM and ViT-B/16-AUG. R denotes the missing rate under linear interpolation.
Model
Xmar
Ltratn
R(.)
VIT-B/16
738.8
6.66
57.9%
VIT-B/16-SAM
20,9
0.96
39.6%
VTTBB/6-AUG
1659.3
1.23
21.4%
8

---

Published as a conference paper at ICLR 2022
However, like SAM, augmentations make ViT-B/16-AUG smoother and achieve significantly smaller training error under random Gaussian perturbations than ViT-B/16. These results show that both SAM and augmentations make the loss landscape flat average. The difference is that SAM enforces the smoothness by reducing the largest curvature via a minimax formulation to optimize the worst-case scenario, while augmentations ignore the worse-case curvature and instead smooth the landscape over the directions induced by the augmentations.
Interestingly, besides the similarity in smoothing the loss curvature on average, we also discover that SAM-trained models possess "linearality" resembling the property manually injected by the mixup augmentation. Following Zhang et al. (2018), we compute the prediction error in-between training data in T where prediction y counted as a miss if it does not belong to {yt. y evaluated at x = 0.5x, 0.5.r We observe that SAM greatly reduces the missing rate (R) compared with the vanilla baseline, showing similar effect to mixup that explicitly encourages such linearity.
5 ABLATION STUDIES
In this section, we provide more comprehensive study about SAM's effect on various vision models and under different training setups. We refer to Appendices B D for the adversarial, contrastive and transfer learning results.
5.1 WHEN SCALING THE TRAINING SET SIZE
Previous studies scale up training data to show massive pre-training trumps inductive biases (Doso- vitskiy et al. 2021 Tolstikhin et al. 202T). Here show SAM further enables ViTs and MLP- Mixers to handle small-scale training data well. We randomly sample 1/4 and 1/2 images from each ImageNet class to compose two smaller-scale training sets, i.e., ilk (1/4) and ilk (1/2) with 320,291 and 640.583 images, respectively. We also use ImageNet-21k to pre-train the models with SAM, followed by fine-tuning on ImageNet- Ik without SAM. The ImageNet validation set remains intact. SAM can still bring improvement when pre-trained on ImageNet-21k (+0.3% +1.4%, and 2.3% for ResNet- 152, VIT-B/16, and Mixer-B/16 respectively).
As expected, fewer training examples amplify the drawback of ViTs and MLP-Mixers lack of the convolutional inductive bias - their accuracies decline much faster than ResNets' (see Figure f in the Appendix and the corresponding numbers in Table 5) However, SAM can drastically rescue ViTs and MI.eMixens performance decrease on smaller training sets. Figure (right)) (right) shows that the improvement brought by SAM over vanilla SGD training is proportional to the number of train- ing images. When trained on ilk (1/4), it boosts ViT-B/16 and Bixer-B/166 by 14.4% and 25.6%, escalating their results 66.8% and 62.8%, respectively. It also tells that ViT-B/16-SAN matches the performance of ResNett-152-5 152-SAM even with only 1/2 ImageNet training data.
CONCLUSIONS AND LIMITATIONS
This paper presents a detailed analysis of the convolution-fiee ViT's and MLP-Mixers from the lens of the loss landscape geometry, intending reduce the models' dependency on massive pre-training and/or strong data augmentations. We arrive at the sharpness-aware minimizer (SAM) after ob- serving sharp local minima of the converged models. By explicitly regularizing the loss geometry through SAM, the models enjoy much flatter loss landscapes and improved generalization regard- ing accuracy and robustness. The resultant ViT models outperform ResNets of comparable size and throughput when learned with no pre-training or strong augmentations. Further investigation reveals that the smoothed loss landscapes attribute to much sparser activated neurons the first few lay- Last but not least, we discover that SAM and strong augmentations share certain similarities to enhance the generalization. They both smooth the average loss curvature and encourage linearity.
Despite achieving better generalization. training VITs with SAM has the following limitations which could lead to potential future work. First, SAM incurs another round of forward and backward propagations to update e, which will lead to around 2x computational cost per update. Second, notice that the effect of SAM diminishes as the training dataset becomes larger, so is vital to develop learning algorithms that can improve/accelerate the large-scale pre-training process.

---

Published as a conference paper at ICLR 2022
ETHICS STATEMENT
We are not aware of any immediate ethical issues in our work. We hope this paper can provide new insights into the convolution-fre neural architectures and their interplay with optimizers, hence benefiting future developments of advanced neural architectures that are efficient in data and com- putation. Possible negative societal impacts mainly hinge on the applications of convolution-fies architectures, whose societal effects may translate to this work.
ACKNOWLEDGEMENT
This work is partially supported by NSF under IIS-1901527, IIS-2008173, IIS-2048280 and by Army Research Laboratory under agreement number W911NF-20-2-0158.
REPRODUCIBILITY STATEMENT
We provide comprehensive experimental details and references existing works and codebases to ensure reproducibility. The specification of all the architectures used in this paper is available in Appendix A The instructions for plotting the landscape and the attention map are detailed in Appendix E We also present our approach to approximating Hessian's 's dominant eigenvalue Amaz and the NTK condition number
Appendices[Jand[e respectively. Finally, Appendix Appendix[H]deserrete describes all the necessary training configurations, data augmentations, and SAM hyperparameters to ensure the reproducibility our results.
REFERENCES
Hassan Akbari, Liangzhe Yuan, Rui Qian, Wei-Hong Chuang, Shih-Fu Chang, Yin Cui, and Boqing Gong. Vatt: Transformers for multimodal sulf-super learning from raw video, audio and 2021. 2021.
text. arXiv preprint arXhv:2104.1117,,
Anurag Arnab, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Lučić, and Cordelia Schmid. Vivit: A video vision transformer. arXiv preprint arXiv:2103. 150.111,, 2021.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint vrXiv:1607.06450 2016.
Irwan Bello, William Fedus, Xianzhi Du, Ekin D. Cubuk, Aravind Srinivas, Tsung Yi Lin, Jonathon Shlens, and Barret Zoph. Revisiting resnets: Improved training and scaling strategies, 2021.
Gedas Bertasius, Heng Wang, and Lorenzo Torresani. Is space-time attention all you need for video understanding? arXiv preprint arXiv:2102.05099, 2021.
Lucas Beyer, Olivier Hénaff, Alexander Kolesnikov, Xiaohua Zhai, and Aäron den Oord. Are we done with imagenet? arXiv preprint arXiv:2006.07/552 2020.
Aleksandar Botev, Hippolyt Ritter, and David Barber. Practical Gauss-Newton optimisation for deep learning. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 557-565. PMLR, 06-11 Aug 2017. URL http://proccednng....... www. wtt.sssss..ppp.s000....
Notev17a.htme
Rebekka Burkholz and Alina Dubatovka. Initialization of relus for dynamical isometry. In H. Wallach, Н. Larochelle, A. Beygelzimer, F. d'Alché-Bue, E. Fox, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Asso- ciates, Inc., 2019. URL
pdf
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers, 2021.
10

---

Published as a conference paper at ICLR 2022
Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-sgd Biasing gradient descent into wide valleys. In International Conference Learning Representations. 2017.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Pro- ceedings of Machine Learning Research, pp. 1597-1607. PMLR, 13-18 Jul 2020. URL http:
\/proceeding mlr press v119/chen2011 html
Wuyang Chen, Xinyu Gong, and Zhangyang Wang. Neural architecture search on imagenet in four GPU hours: A theoretically inspired perspective. In International Conference on Learning Representations, 2021a. URL
hitpa/vveeee.eettterooeiioeeoooo www.rrrmmmmdooooeeeeeememe uMHtu
Xiangning Chen and Cho-Jui Hsieh. Stabilizing differentiable architecture search via perturbation- based regularization. In Hal Daumé III and Aarti Singh (eds.). Proceedings of the 37th Interna- tional Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Re- search, pp. 1554-1565, PMLR, 13-18 Jul 2020. TRU/tttp://pooceesee v119/chen20E html
mlr. ..s.s.
Xiangning Chen, Cihang Xie, Mingxing Tan, Li Zhang, Cho-Jui Hsieh, and Boqing Gong. Robust and accurate object detection via adversarial learning. In Proceedings of the IEFF/CVF Confer- ence on Computer Vision and Pattern Recognition (CVPR), pp. 16622-16631, June 2021b.
Xinlei Chen, Saining Xie, and Kaiming He. An empirical study of training self-supervised vision transformers, 2021c.
Ekin D. Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoc V. Le. Autoaugment: Learning augmentation strategies from data. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2019.
Ekin D. Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V. Le. Randaugment: Practical au- tomated data augmentation with reduced search space. In 2020 IEEE/CVF Conference Computer Vision and Puttern Recognition Workshops (CVPRW), pp. 3008-3017, 2020. doi: 10. 10.1109/CVPRW50009. 20200000900000
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hier- archical image database. In 2009 JEEE Conference on Computer Vision and Pattern Recognition, pp. 248-255, 2009. doi: 10.
Jacob Devlin, Ming- Wei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805. 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recogni- tion at scale. In International Conference Learning Representations, 2021. URL https: //openreview. www.:/ttt.//oddieeeemomomm cib.dttt
Gintare Karolina Dziugaite and Daniel M. Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In Gal Elidan, Kristian Kersting, and Alexander T. Ihler (eds.), Proceedings of the Thirty-Third Conference on Uncertainty in Artificial Intelligence. UAI 2017. Sydney, Australia, August 11-15, 2017. AUAI pdd
Press, 2017. URL http::///aaaa..crraaaa///////oo/
Haoqi Fan, Bo Xiong, Karttikeya Mangalam, Yanghao Li, Zhicheng Yan, Jitendra Malik, and Christoph Feichtenhofer. Multiscale vision transformers. arXiv preprint arXiv:2104. 12221104../22... 2021.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimiza- tion for efficiently improving generalization. In International Conference on Learning Represen- fations, 2021. wtt..ntpe:///geereeeee...eeeeeed
6 rM

---

Published as a conference paper at ICLR 2022
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp- 770-778, 2016.
doi: 10.1109/CVPR.2016.990
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR), June 2020.
Dan Hendryeks and Thomas Dietterich. Benchmarking neural network robustness to common cor- ruptions and perturbations. In International Conference on Learning Representations, 2019, URL
httpss://ppenr:
iCqYm
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus), 2020.
Dan Hendrycks, Steven Basart, Norman Mu, Saurav Kadavath, Frank Wang. Evan Dorundo, Rahul Desai, Tyler Zhu, Samyak Parajuli, Mike Guo, Dawn Song, Jacob Steinhardt, and Justin Gilmer. The many faces of robustness: A critical analysis of -ittr--dittribution generalization. 2020.
Stanistaw Jastrzæbski, Zachary Kenton, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amost Storkey. On the relation between the sharpest directions of DNN loss and the SGD step length. In International Conference on Learning Representations, 2019. URL http:::/
net woruforrmmmmm
d http://oommmmimkkkpe. 10
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2017. URL
Prannay Khosla, Piotr Teterwak, Chen Wang, Aaron Sarna, Yonglong Tian, Phillip Isola, Aaron Maschinot, Ce Liu, and Dilip Krishnan. Supervised contrastive learning. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 18661-18673. Curran Asso- mewpoeeeeeddins.e.eiee..ee..eee cc/papeer2022///1ee d (99aa667880222111bbbbbaaaaaaaaa2 Paper pddf
ciates, Inc., 2020. URL
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015.
Bobby Kleinberg, Yuanzhi Li, and Yang Yuan. An alternative view: When does SGD escape lo- cal minima? In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 2698 2707. PMLR, 10-15 Jul 2018. URL http:/ http::/pproceeiingg http://pooceeinns.......... mir. prs.pvees/v.//
kleinberg 8a. htm.
Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Joan Puigcerver, Jessica Yung, Sylvain Gelly, and Neil Houlsby. Big transfer (bit): General visual representation learning. In Andrea Vedaldi, Horst Bischof, Thomas Brox, and Jan-Michael Frahm (eds.). Computer Vision - ECCV 2020, pp- 491-507, Cham, 2020. Springer International Publishing. ISBN 978-3-030-58555-7.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss land- scape of neural nets. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran As- sociates, Inc., 2018. URL
1411:000366600000000600666666663 Papee pdf
Hanxiao Liu, Zihang Dai, David R So, and Quoc V Le. Pay attention to mlps. arXiv preprint 1:Xiv:2105.08000 202
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows, 2021b.
12

---

Published as a conference paper at ICLR 2022
Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu. To- wards deep learning models resistant to adversarial attacks. In International Conference on Learn ing Representations, 2018. URL /ttptttt/////eteeeeeee htwperreeewweeeeeeeeeeeeeeeeeeee 2ttp://orrmmi1d122122222222
Luke Syriazi Do you even need attention? stack of feed-forward layers does surprisingly well on imagenet. arXiv preprint arXiv:2105.02725, 2021.
Y. Nesterov. A method for solving the convex programming problem with convergence rate (l1/k²). Proceedings of the USSR Academy of Sciences, 269:543-555 -99444,477,, 1983.
Maria-Elena Nilsback and Andrew Zisserman. Automated flower classification over a large number of classes. In 2008 Sixth Indian Conference on Computer Vision, Graphics Image Processing, pp. 722- 729, 2008. doi: 0..119//CVGIP.2200.441
Omkar M Parkhi, Andrea Vedaldi, Andrew Zisserman, and C. V. Jawahar. Cats and dogs. In 2012 IEEE Conference on Computer Vision and Pattern Recognition, pp. 3498-3505, 2012. doi: 10.1 2..119//CVPR.22222228888
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under- standing by generative pre-training, 2018.
Benjamin Recht, Rebecca Roclofs, Ludwig Schmidt, and Vaishaal Shankar. Do imagenet classifiers generalize to imagenet? In International Conference on Machine Learning, pp. 5389-5400.
PMLR, 2019.
Ali Shafahi, Mahyar Najibi, Mohammad Amin Ghiasi, Zheng Xu, John Dickerson, Christoph Studer, Larry S Davis, Gavin Taylor, and Tom Goldstein. Adversarial training for free! In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Asso- ciates, Inc., 2019. URL
neuripss httpppscccccpppe//////1e Pape: .poperrpdfff.
Yeonjong Shin and George Em Karniadakis. Trainability of relu networks and data-dependent ini- tialization. Journal of Machine Learning for Modeling and Computing, 1(1):39-74, 2020. ISSN 2689-3967.
Samuel I.. Smith and Quoe Le. bayesian perspective on generalization and stochastic gradient descent. In International Conference Learning Representations, 2018. URL https://
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: A simple way to prevent neural networks from overfitting. Journal of Machine Learning Research, 15(56):1929-195, 2014. URL http: http:/mw.l..ppppp../eeeep
srivastaval4a.hter
Chen Sun, Abhinav Shrivastava, Saurabh Singh, and Abhinav Gupta. Revisiting unreasonable effee- tiveness of data in deep learning era. In 20/7 IEEE International Conference Computer Vision (ICCV), pp. 843-852, 2017. doi: 10. 1011109C..2001..99
Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna. Rethinking the inception architecture for computer vision. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 2818-2826 2016. doi: 10. 1119/CVP..0116.000
T. Tieleman and G. Hinton. Lecture 6.5 R-s-mmpprop: Divide the gradient by running average of its recent magnitude. COURSERA: Neural Networks for Machine Learning, 2012.
llya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Un- terthiner, Jessica Yung, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, and Alexey Dosovitskiy. Mlp-mixer: An all-mlp architecture for vision, 2021.
Hugo Touvron, Piotr Bojanowski, Mathilde Caron, Matthieu Cord, Alaaeldin El-Nouby, Edouard Grave, Armand Joulin, Gabriel Synnaeve, Jakob Verbeek, and Hervé Jégou. Resmlp: Feedforward networks for image classification with data-efficient training, 2021a.
13

---

Published as a conference paper at ICLR 2022
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention, 2021b.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. arXiv preprint arXiv:1706.0376. 2017.
Eric Wong, Leslie Rice, and Zico Kolter. Fast is better than free: Revisiting adversarial training. In International Conference on Learning Representations, 2020. UtL.htttps://ppeteevi
net / wtt/://romnndddx00000000
Dongxian Wu, Shu-Tao Xia, and Yisen Wang. Adversarial weight perturbation helps robust gen- eralization. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Ad- vances in Neural Information Processing Systems, volume 33, pp. 2958-228888 2969. Curran Asso- ciates, Inc., 2020. URL NetfSiccI..c00001111116633322222
neurips wwtprriss.cccppppeeee cc/ppperr/222//ilee ppp
Lechao Xiao, Jeffrey Pennington, and Samuel Schoenholz. Disentangling trainabilityy and general- ization in deep neural networks. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learn- ing Research, pp. 10462-10472. PMLR, 13-18 Jul 2020. URL httphhtt:/pppe http::::/prccceee.gggs..... mmr.
press/v119/xiao20bb html
Cihang Xie, Mingxing Tan, Boqing Gong, Jiang Wang, Alan L. Yuille, and Quoc V. Le. Adversarial examples improve image recognition. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition (CVPR), June 2020.
Ge Yang and Samuel Schoenholz. Mean field residual networks: On the edge of chaos. In Guyon, U. V. Luxburg. S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Asso- ciates, Inc., 2017. URL http:::tt//stt.tt.ttt http:ss:///rooceedgggs..oo Pape.t
httpps.ccpppppeeee cct/ppppe//211//1eee
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes. In International Conference on Learning Representations, 2020. URLlhttps:
http///eeeevveeeeeeeeeeeee /
net/ wtw.ummmmd.Sxxxmtttt
Li Yuan, Yunpeng Chen, Tao Wang, Weihao Yu, Yujun Shi, Zihang Jiang, Francis EH Tay, Jiashi Feng, and Shuicheng Yan. Tokens-to-token vit: Training vision transformers from scratch on imagenet, 2021.
Sangdoo Yun, Dongyoon Han, Sanghyuk Chun, Seong Joon Oh, Youngjoon Yoo, and Junsuk Choe. Cutmix: Regularization strategy to train strong classifiers with localizable features. In 2019 IEEE/CVF International Conference Computer Vision (ICCV), pp. 6022 6031, 2019. 10. N0.1109//CCV.2219.0001
doi:
Arber Zela, Thomas Elsken, Tonmoy Saikia, Yassine Marrakchi, Thomas Brox, and Frank Hut- ter. Understanding and robustifying differentiable architecture search. In International Confer- ence on Learning Representations, 2020. URL
www. fttt:ffrrrmmidd
H1 @DNyrkDD
Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, and David Lopez-Paz. mixup: Beyond em- pirical risk minimization. In International Conference Learning Representattonss 2018. URL
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank J. Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models?
In NeurlPS, 2020. URL / neuriissssnn.n.uu....e.......... wttrrrissccccaaae cc/pppper/22//hsh// 3Db51TT8400000010000006000003333
14

---

Published as a conference paper at ICLR 2022
APPENDICES
A ARCHITECTURES
Table 8 specifies the ViT (Dosovitskiy et al. 2021 Vaswani et et 2017) and MLP-Mixer (Tolstikhin ct al., 2021) architectures used in this paper. "S" and "B denote the small and base model scales following (Dosovitskiy et al.) 202 Touvron et al. 202 Tolstikhin et al. 202ID. followed by the size of each image patch. For instance, "B/16" means the model of base scale with non-overlapping image patches of resolution 16 16. We use the input resolution 224 x 224 throughout the paper. Following Tolstikhin et al. (2021). we sweep the batch sizes in {32, 64, 8192} on TPU-v3 and report the highest throughput for each model.
Table 7: Comparison under the adversarial training framework on ImageNet (numbers in the paren- theses denote the improvement over the standard adversarial training without SAM). With similar model size and throughput, ViTs-SAM can still outperform ResNets-SAM for clean accuracy and adversarial robustness.
Throughpot (img/se/corr)
Model
*params
ImageNet
V2
PGD-10
ImageNet-RR
ImageNet-C
ResNet
ResNet-50-SAM
25M
2161
(-0.7)
77.9
56.6 (-0.8)
54.1 (40.9)
(+0.9
42.7
ResNet- 101 SAM
44M
1334
73.6 (-0.4)
81.0 (+0.1)
60.4 60.4 (-0.6)
58.8 (+1.4)
29.5 (+0.6)
46.9 (+0.3)
ResNet-152-SAN
6IM
935
75.1 (40.4)
82.3 (+0.2)
62.2 (-0.4)
61.0(+1.88
30.8 (+1.4)
49.1 (+0.6)
Vision Transformer
SIT-S/16-SA
22M
2043
73.2 (+1.2)
80.7 (+1.7)
60.2 (+1.4)
58.0 (+5.2)
28.4 (+2.4)
47.5 (+1.6)
VIT-B/22-SM
88M
2805
69.9 (+3.0)
76.9 (+3.4)
55.7 (+2.5)
(464)
26.0 (+3.0)
46.4 (+3.0)
VIT-B/IG-SM
S7M
863
76.7 (+3.9)
82.9 (+4.1)
63.6 (+4.3)
62.0 (+7.7)
30.0 (+4.9)
$1.4 (+5.0)
MLP-Mixer
Siner-S/16-SAM
18M
4005
67.1 (+2.2)
74.5 (+2.3)
52.8 (+2.5)
50.1 (14.1)
22.9 (+2.6)
37.9 (+2.5)
Mixer-8/32-SAM
60M
4209
(+9.1)
76.4 (+10.2)
54.7 (+9.4)
54.5 (+13.9)
26.3 (48.0)
43.7
Mixer-B/16-SAM
59M
1390
73.9 73.9 (+11.1)
80.8 (+11.8)
(0.2(+1119)
59.8 (.17.3)
29.0 (+10.5)
45.9 12.5)
Table 8: Specifications of the ViT and MLP-Mixer architectures used in this paper. We train all the architectures with image resolution 224 224.
Throughput (img/sec/core)
Patch Resolution
Sequence Length
Token MLP
Channel MLP
Model
#params
Hidden Size
Hheads
#layers
Dimension
Dimension
VIT-8/32
23M
6888
32 32
384
VIT:S/16
22M
2043
16
196
384
6
12
VIT-S/14
22M
1234
14 14
256
384
6
12
VIT-S/8
22M
333
384
6
12
V/T-B/32
S&M
2805
32 x
49
768
12
12
VIT-B/16
Mixer-S/32
19M
1140]
32 32
49
512
8
255
2048
Mixer-S/16
18M
4005
x6xx
196
512
2048
Mixer-S/8
20M
1498
x
784
512
8
256
2048
Mixer B/32
60M
4209
32 32
768
384
3072
Mixer-B/16
5UM
1390
16 x6xx11
196
768
384
3072
Mixer-B/8
64M
466
8
784
768
12
384
3072
B WHEN SAM MEETS ADVERSARIAL TRAINING
Interestingly, SAM and adversarial training are both minimax problems except that SAM's inner maximization is with respect to the network weights, while the latter concerns about the input for
Table 9: Hyperparameters for downstream tasks. All models are fine-tuned with 224 x 224 resolu- tion, a batch size of 512, cosine learning rate decay, no weight decay, and grad clipping at global norm 1.
Dataset
Total steps
Warmup steps
Base LR
CIFAR-10
10K
500
CIFAR-100
10K
500
[0.001, 0.003, 0.01, 0.03)
Flowers
500
100
500
100
15

---

Published as a conference paper at ICLR 2022
Imageteet Top-2 Improvemeett (%)
ResNet-152
ImageNet Top-1 Accuracy (Sa
VIT-8/16
20
25
ResNet-152
Rasteet-52--AA
10-
50-
VIT-8/16
VIT B/15 SAM
Mixer8/16
80
Mixer-8/16--00
13k (1/4)
13k (1/2)
1tk 121K
Figure 4: ImageNet accuracy (Left) and improvement (Right) brought by SAM.
defending contrived attack (Madry al. 2018 Wong al. 2020). Moreover, similar to SAM, Shafahi et al. (2019) suggest that adversarial training can flatten and smooth the loss landscape. In light of these connections,
study ViTs and MIP--Mixers under the adversarial training frame-
work (Wu et al.) 2020 Madry et al. 2018). We use the fast adversarial training (Wong et al.) 2020) (FGSM with random start) with the to norm and maximum per-pixel change 2/255 during training. All the hyperparameters remain the same as the vanilla supervised training. When evaluating the ad- versarial robustness, use the PGD attack (Madry al. 2018) with the same maximum per-pixel
change 2/255. The total number of attack steps is 10, and the step size is 0.25/255. To incorporate SAM, we formulate a three-level objective:
min max max Lerain (w + E, + 6,
(5)
10
cE
BES,do
where Spam and Sade denote the allowed perturbation norm balls for the model parameter and in- put image x, respectively. Note that we can simultaneously obtain the gradients for computing and 5 by backpropagation only once. To lower the training cost, we use fast adversarial training (Wong) cet al. 2020) with the Lx0 norm for 6, and the maximum per-pixel change is set as 2/255.
Table see Appendices) evaluates the models' clean accuracy, real-world robustness, and adver- sarial robustness (under 10-step PGD attack (Madry et al., 2018)). It is clear that the landscape smoothing significantly improves the convolution-free architectures for both clean and adversarial accuracy. However, we observe a slight accuracy decrease on clean images for ResNets despite gain for robustness. Similar to our previous observations, ViT's surpass similar-size ResNets when adversarially trained on ImageNet with Inception-style preprocessing for both clean accuracy and adversarial robustness.
C WHEN SAM MEETS CONTRASTIVE LEARNING
In addition to data augmentations and large-scale pre- pre-training, another notable way of improving a neural model's generalization is (supervised) contrastive learning (Chen et al. 2020 He et al.
2020 Caron et al. 2021 Khosla al. 2020). We couple SAM with the supervised contrastive learning (Khosla et al. 2020) for 350 epochs, followed by fine-tuning the classification head by 90 epochs for both ViT-S/16 and ViT-B/16. We train ViTs under the supervised contrastive learning framework (Khosla et al. 2020). We take the classification token output from the last layer as the encoded representation and retain the structures of the projection and classification heads (Khosla et al.) 2020). We employ a batch size 2048 without memory bank (He et al. 2020) and use Au- touAumeent (Cubuk et al. 2019) with strength 1.0 following Khosla et al. (2020).
). For the 350-epoch pretraining stage, the contrastive loss temperature is set as 0. and we use the LAMB optimizer (You) al. 2020) with learning rate 0.001 x 256 along with cosine decay schedule. For the second stage, train the classification head for 90 epochs via a RMSProp optimizer (Tieleman & Hinton, 2012) with base learning rate 0.05 and exponential decay. The weight decays are set as 0.3 and 1e-6 small SAM perturbation strength p 0.02.
for the first and second stages, respectively. We
Compared to the training procedure without SAM, we find considerable performance gain thanks to SAM's s smoothing of the contrastive loss geometry. improving the ImageNet top-1 accuracy of ViT- S/16 from 77.0% to 78.1% %. and ViT-B/16 from 77.4% to 80.0%. In comparison, the improvement on ResNet-152 is less significant (from 79.7% to 80.0% after using SAM).
16

---

Published as a conference paper at ICLR 2022
Table 10: Accuracy on downstream tasks of the models pre-trained on ImageNet. SAM improves
VITs and MLP-Mixers' transferabilities. ViT's transfer better than ResNets of similar sizes.
ResNet- S0-SAM
VIE://16
VIT-R/16-
Mixer-V/16- SAM
Mixer-R/16-
152-SAM
ViT-S/16
SAM
ViT-8/16
SAM
Mixar-N/16
Mixcr-B/16
SAM
CIFAR-10
97.4
97.6
98.2
98.1
98.5
94.1
95.4
97.8
CIFAR-100
85.2
85.7
87.6
77.9
82.4
86.4
Flowers
90.0
91.1
85.4
91.5
89.55
91.8
83.3
87.9
82.8
90.0
Pets
91.6
93.3
90.4
92.9
91.9
93.1
86.1
88.7
86.1
92.5
Averuge
91.5
(a) VIT
(b) VIT-SAM
(c) VIT-AUG
(d) VIT-2Ik
Figure 5: Cross-entropy loss landscapes of ViT-B/16, ViT-B/16-SAM, ViT-B/I6-AUG, and VIT-
B/16-21k. Strong augmentations and large-scale pre-training also smooth the eurvature.
D WHEN SAM MEETS TRANSFER LEARNING
We also study the role of smoothed loss geometry in transfer learning. We select four datasets test ViTs and MI-P-Mieess transferabilities: CIFAR-10/100 (Krizhevsky 2009). Oxford-IIIT Pets (Parkhi et al. 2012). and Oxford Flowers-1002 (Nilsback & Zisserman (008). We use image resolution 224 224 during fine-tuning on downstream tasks, other settings exactly follow Doso- vitskiy et al. (202T) Tolstikhin et al. (02TT) (see Table Note that we do not employ SAM during fine-tuning. We perform a grid search over the base learning rates on small sub-splits of the train- ing sets (10% for Flowers and Pets, 2% for CIFAR-10/100), After that, we fine-tune on the entire training sets and report ResNet-50-SAM and ResNet- 102-152-AAMM in the experiments. confirm that the enhanced models also perform better after fine-tuning and that MI.P-Mixers gain the most from the sharpness-aware optimization.
results on the respective test sets. For comparison, we also include
the results, which
E VISUALIZATION
E LOSS LANDSCAPE
We use the "filter normalization" method (Li al. 2018) visualize the loss function curvature in Figure and For fair comparison, we use the cross-entropy loss when plotting the landscapes for all architectures, although the original training objective the sigmoid loss for ViT's and MI.P- Mixers. Note that their sigmoid loss geometry
sharper. We equally sample 2,500 points the 2D projection space and compute the losses using 10% of the ImageNet training images (Chen al. 2020 i.., the ilk (1/10) subset in the main text save computation.
E.2 ATTENTION MAP
The visualization of the VIT's attention maps (Figure Binn the main text) follows (Caron al. 1021) We average the self-attention scores of the "classification token" from the last MSA layer to obtain a matrix E H//Px//P where H. W, P are the image height, width, and the patch resolution, respectively. Then upsample A to the image shape H W before generating the figure.
17

---

Published as a conference paper at ICLR 2022
Table 11: The SAM perturbation strength p for training on ImageNet. ViTs and MLP-Mixers favor larger than ResNets does. Larger models with longer patch sequences need stronger strengths.
Model
Task
SAM
ResNet
RosNct-50-SAM
supervised
0.02
ResNet-101-SAM
supervised
0.05
ResNei-12-SAAM
supervised
0.02
ResNet-50x2-SAM
supervised
0.05
ResNet-10122SAA
supervised
0.05
ResNet-15222-SAM
supervised
0.05
ResNet-50-SAM
adversurial
0.05
ResNet-101-SAM
adversarial
0.05
ResNct-152-SAM
adversarial
0.05
VIT
VIT-S/32-SAM
supervised
0.05
ViT-S/I6-SAM
supervised
0.1
S/T-S/14-SAM
supervised
0.1
ViT-S/8-AM
supervised
0.15
VIT-B/32-SAM
supervised
0.15
VIT-B/I6-SAM
supervised
0.2
SUG-B/I6AAGG-AM
supervised
0.05
VIT-S/I6-SAM
adversarial
0.1
ViT-B/32-SAM
adversarial
0.1
ViT-B/IG-SAM
adversarial
0.1
ViT-S/I6-SAM
supervised contrastive
0.02
ViT-B/16-SAM supervised contrastive
0.02
MixeMixerr
Mixer-S/32-SAM
supervised
Mixer-S/66-SAM
supervised
0.15
Mixer-S8-SAA
supervised
0.2
Mixcr-
B/xcr-B/32-SAM
supervised
0.35
Mixer-B/16-SAM
supervised
0.6
Mixcr-B/8-SAM
supervised
0.6
supervised
0.2
Mixer-S/16-SAM
adversarial
0.05
Mixer-B/32-SAM
adversarial
0.25
Mixer-B/16.n B/Ie-ABAI6
adversarial
0.25
F HESSIAN EIGENVALUE
The Hessian matrix requires second-order derivative, so we compute the Hessian (and all the sub- diagonal Hessian) Amar using 10% of the ImageNet training images (i.e., ilk (1/10) via power iteration where we use 100 iterations to ensure convergence.
G NTK CONDITION NUMBER
We approximate the neural tangent kernel on the ilk (1/10) subset by averaging over block diagonal entries (with block size 48 48) in the full NTK. Notice that the computation based on the archi- tecture at initialization without training. As the activation plays an important role when computing NTK we find that smoother activation functions enjoy smaller condition numbers, we replace the GELU in ViT and MLP-Mixer with ReLU for fair comparison with ResNet.
H TRAINING DETAILS
We use image resolution 224 x 224 during fnn-ttuning on downstream tasks, other settings exactly
follow (Dosovitskiy et al. 2021 Tolstikhin et al. (021) (see Table Note that we do not employ SAM during fine-tuning. We perform a grid search over the base learning rates on small sub-splits of the training sets (10% for Flowers and Pets, 2% for CIFAR- CIFAR-10/100), After that, we fine-tune the entire training sets and report the results on the respective test sets.
htttss///en.wiiieedia.oom htwp://wwwwiceeeeieeeeeee/eeeee/
18

---

Published as a conference paper at ICLR 2022
Table 12: Hyperparameters for training from scratch on ImageNet with basic Inception-style pre- processing and 224 x 224 image resolution.
ResNet
ViT
MLP-Mixer
Data augmentation
Inception-style
Input resolution
221 x 221
Batch size
4,096
Epoch
90
300
300
Warmup steps
5K
10K
10K
Peak learning rate
0.1
3e-3
3e-3
258
Learning rate decay
cosine
cosine
linear
Optimizer
SGD
AdamW
AdamW
SGD Momentum
0.9
Adam (B,, 82)
(0.9, 0.999)
(0.9, 0.999)
Weight decay
lc-3
0.3
0.3
Dropout rate
0.0
0.1
0.0
Stochastic depth
0.1
Gradient clipping
1.0
1.0
Table 13: ImageNet top-1 accuracy (%) of ViT-B/16 and Mixer-B/16 when trained from scratch with different perturbation strength p in SAM.
SAM
0.0
0.05
0.1
0.2
0.25
0.35
0.4
0.5
0.6
0.65
ViT-B/16
74.6
77.5
78.8
79.9
79.3
Mixer-Bil B/16
66.4
69.5
74.1
74.7
75.6
77.4
77.1
Except for the experiments in Section 4.5 (SAM with strong data augmentations) and Appendix C (contrastive learning), we train all the models from scratch on ImageNet with the basic Inception- style preprocessing (Szegedy al. 2016), i.e., random image crop and horizontal flip with probability 50%. Please see Table for the detailed training settings. We simply follow the original training settings of ResNet and ViT (Kolesnikov al. 2020 Dosovitskiy al., 2021). For MLP- Mixer, we remove the strong augmentations in its original training pipeline and perform a grid search over the learning rate {0.003, 0.001}, weight decay in {0.3, 0.1, 0.03), Dropout rate in (0.1, 0.0), and stochastic depth in {0.1. 0.03. Note that training for 90 epochs is enough for ResNets to converge, and longer schedule brings almost no effect. For the experiments, we use 128 TPU- v3 cores (2 per chip), resulting in 32 images per core. The SAM computation for is conducted on each core independently.
H.1 PERTURBATION STRENGTH IN SAM
Different architecture species favor different strengths of perturbation p. We perform grid search p and report the best results Table 11 reports the corresponding strengths used in Ima- geNet experiments. Besides, we show the results when varying p in Table Similar to (Foret et all] 2021, J, also find that relative small م [0.02, 0.05) works the best for ResNets. However, larger gives rise to the best results for ViTs and MLP-Mixers. We also observe that architectures with larger capacities and longer input sequences prefer stronger perturbation strengths. Interestingly, the choice of coincides with our previous observations. Since MLP-Mixers suffer the sharpest landscapes, they need the largest perturbation strength. As strong augmentations and contrastive learning already improve generalization, the suitable p becomes significantly smaller. Note that we do not re-tune any other hyperparameters when using SAM.
H.2 TRAINING ON IMAGENET SUBSETS
In Section we train the models on ImageNet subsets, and the hyperparameters have to be ad- justed accordingly. We simply change the batch size to maintain similar total iterations and keep all other settings the same, i.c., 2048 for ilk (1/2), 1024 for ilk (1/4), and 512 for ilk (1/10). We do not scale the learning rate find the scaling harms the performance.
19

---

Published as a conference paper at ICLR 2022
H.3 TRAINING WITH STRONG AUGMENTATIONS
We tune the learning rate and regularization when using strong augmentations (mixup with probabil- ity 0.5, RandAugment with layers and magnitude 15) in Section followiig (Tolstikhin al., 2021). For VIT, we use le-3 peak learning rate, 0.1 weight decay. 0.1 Dropout, and 0.1 stochastic depth; For MLP-Mixer, those hyperparameters are exactly the same as (Tolstikhin al. (02I), I), peak learning rate as Ic-3, weight decay as 0.1, Dropout as 0.0, and stochastic depth as 0.1. Other settings are unchanged (Table 12
LONGER SCHEDULE OF VANILLA SGD
Since SAM needs another forward and backward propagation to compute è, its training overhead is ~ 2x of the vanilla baseline. We also experiment with 2x schedule vanilla training (600 epochs). We observe that training longer brings no effect on both clean accuracy and robustness, indicating that the current 300 training epochs for ViTs and MLP-Mixers are enough for them to converge.
J VARYING WEIGHT DECAY STRANGTH
Table 14: ImageNet accuracy and curvature analysis for ViT-B/16 when we vary the weight decay strength in Adam (AdamW).
Model
Weight decay
ImageNet (%)
|||||
Lrrain
LA train
Amer
0.2
74.2
339.8
0.51
4.22
507.4
0.3
74.6
269.3
0.65
6.66
738.8
VIT-B/16
0,4
74.7
236.7
0.77
1548.9
0.5
74.4
211.8
7.21
2251.7
0,2
79.9
461.4
0.69
0.72
13.1
0,3
79.9
353.8
0.82
0.96
20.9
ViT-B/16-SAM
0.4
79.4
301.1
0.85
0.98
26.1
0.5
78.7
259.6
0.95
1.33
45.5
In this section, we vary the strength of weight decay and see the effects of this commonly used regularization approach. shown Table 14 weight decay helps improve the accuracy on Im- ageNet when training without SAM, the weight norm also decreases when we enlarge the decay strength as expected. However, enlarging the weight decay aggravates the problem of converging to sharper region measured by both Lyyain and Amaz. Another observation is that ||w]|2 consis- tently increases after applying SAM for every weight decay strength Table together with the improved ImageNet accuracy and smoother landscape curvature.
20


Reviews:
[{"review_id": "z0F5EKQpDty", "reviewer": "Reviewer_1uek", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Update:** I have increased my score from a 3 to an 8 after the the authors' response\n\nSAM is a new and promising method, so it’s important to examine and understand its effects. And the benefits of SAM for ViTs and MLP-Mixers appear significant. I applaud the breadth of the authors’ attempts: they examine the effects of SAM across a wide range of image classification scenarios and using a variety of analysis methods.\n\nI am convinced that SAM works well for ViTs and MLP-Mixers. But I’m not convinced of why it works. My main concern about this paper is that it is insufficiently rigorous. The authors provide a number of explanations that are not justified by the evidence, and none of the experiments are run with replicates, making it difficult to determine whether observed differences are meaningful. I am not confident that this work is suitable for publication in its current state. However, I am optimistic that it could be suitable for publication after revision.\n\n# Detailed Feedback\n\n## Replicates and measures of variability\n\nIdeally, all the experiments should be run in replicate with 10 different seeds. However, I understand that this is computationally costly, so I would accept running replicates (three at a bare minimum) for the six core model configurations: ViT-B/16, MLP-Mixer-B/16, and ResNet-152 with and without SAM. Results with these models should be accompanied by measures of variability (or simply listing all three numbers, if only 3 replicates are used).\n\n## Unjustified claims\n\nThere are numerous unjustified scientific claims made, which I detail as follows:\n\nEnd of first paragraph in **Section 3**:\n>“There also exists a large gap between ViTs and ResNets in robustness tests. \n\nRequires a reference.\n\nIn **Section 3. Small training errors** the authors claim that\n\n>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.\n\nThis claim seems speculative, and should be tested. The authors could train a ViT with a different attention normalization scheme, or without one altogether. Learning Q, K, and V projections as opposed to learning MLP. At the very least, a convincing theoretical explanation should be provided.\n\nIn **Section 4.2. Higher Accuracy**, the authors state: \n>Empirically, the degree of improvement negatively correlates with the level of inductive biases built into the architecture. ResNets with inherent translation equivalence and locality benefit less from landscape smoothing than the attention-based ViTs.\n\nThis is an interesting claim, but there are multiple issues with it. First, “level of inductive bias” is vague. I would suggest the language be changed to “correlates with expressivity” or “negatively correlates with the constraints of the inductive biases“. The “level” or “severity” of an inductive bias can be orthogonal to the constraints it imposes on expressivity. ConViT (d’Ascoli et al., 2021), for example, has a convolutional inductive bias but is completely unconstrained in that it can learn to be a vanilla ViT. Second, this claim can easily be tested. The authors could examine the effects of SAM when varying model architecture choices. For example, exchanging layer norm for batch norm in ResNets would have little effect on the model’s expressivity, and consequently the authors’ hypothesis predicts that the effect of SAM should not vary significantly. The authors could examine the effect of SAM on other architectures with varying (and controllable) inductive biases, such as ConViT (d’Ascoli et al., 2021), CvT (Wu et al., 2021), CMT (Guo et al., 2021), a ViT with a convolutional stem (Xiao et al., *Early Convolutions Help Transformers See Better*, 2021), or a CNN-reparameterized into an MLP á la (d’Ascoli et al., *Finding the Needle in the Haystack with Convolutions*, 2020).\n\nThe effects of SAM on activation norm (Table 3) in the ViT are not consistent, so I think the authors should avoid making the claim that “we find that the norm of the post-activation value...become even bigger” (*Section 4.4 Greater Weight Norms*).\n\nThe authors state\n\n>Interestingly, the ViT model optimized with SAM can encode plausible segmentation information, giving rise to better interpretability than the one trained via the conventional SGD optimization\n\nClaims regarding segmentation should be evaluated with a segmentation task (e.g. results using a transformer backbone for segmentation with vs. without SAM), and claims about interpretability should be evaluated with controlled human experiments (see Leavitt and Morcos, *Towards Falsifiable Interpretability*). I think the strongest claim that could be made without appropriate experiments is “it appears more interpretable”.\n\nTable 11 shows that the ViT model is trained using Adam, not SGD, but you write \"...giving rise to better interpretability than the one trained via the conventional SGD optimization\".\n\nTransformers are notoriously unstable during training. ADAM is thought to compensate for this instability. It would be interesting to see whether SAM allows vision transformers to be trained with SGD.\n\nDoes Figure 2a depict training with or without SAM? I think it would be informative to plot both to depict the effect of SAM over the course of training.\n\nThe following claims are made in the introduction, some of which I address in my above comments:\n\n>We conjecture that the convolution-induced translation equivariance and locality help ResNets escape from bad local minima when trained on visual data.\n\nThis claim is not justified: the effects of translation equivariance and locality on SAM are never tested directly, nor is the concept of a “bad” local minimum clearly defined. As suggested before, the authors should repeat their analyses using vision transformers with convolutional/local inductive biases. The authors should also clearly define what they mean by a “bad” local minimum.\n\n>By analyzing some intrinsic model properties, we find that the models after SAM reduce the Hessian eigenvalues by activating sparser neurons (on ImageNet), especially in the first few layers.\n\nThis language implies a causal chain that SAM increases sparsity, and the increased sparsity reduces Hessian eigenvalues. The current experiments can only justify the claim that “SAM increases sparsity and reduces Hessian eigenvalues”. The authors need to conduct experiments showing that SAM causes sparsity and sparsity causes a reduction in hessian eigenvalues. This could be done by regularizing to increase/decrease sparsity with vs. without SAM, and examining the effects on Hessian eigenvalues and accuracy. Otherwise, the authors should amend this claim to remove the implied SAM->sparsity->Hessian causality.\n\n>The weight norms increase, implying the commonly used weight decay may not be an effective regularization alone.\n\nChanges in weight norms alone are insufficient to justify this claim. The authors should repeat their analyses with varying amounts of weight decay with vs. without SAM.\n\n>A side observation is that, unlike ResNets and MLP-Mixers, ViTs have extremely sparse active neurons, revealing the redundancy of input image patches and the capacity for network pruning.\n\nIt’s not clear to me how sparsity translates to image patch redundancy. This is a very interesting idea, and should be better explained.\n\n>Another interesting finding is that ViTs’ performance gain also translates to plausible attention maps containing more perspicuous information about semantic segmentation. \n\nI address this in an earlier comment: evaluate this claim with a segmentation task and/or human interpretability experiments.\n\n## Other Feedback\n\nAnalyzing a vision transformer with an inductive bias for locality such as ConViT (d’Ascoli et al., 2021), CvT (Wu et al., 2021), and/or CMT (Guo et al., 2021) could strengthen a lot of the claims in the paper and bridge the results from ViTs to ResNets.\n\nThe authors state that\n\n>These results show that both SAM and augmentations make the loss landscape flat on average. The difference is that SAM enforces the smoothness by reducing the largest curvature via a minimax formulation to optimize the worst-case scenario, while augmentations ignore the worse-case curvature and instead smooth the landscape over the directions concerning the inductive biases induced by the augmentations.\n\nThis distinction between average- and worst-case curvature is an interesting result. I think the authors should extend their analysis and assess the average flatness of other models with vs. w/o SAM.\n\nIn *Section 4.2. Better Robustness*, the effects of SAM on robustness in ResNets should be mentioned. Furthermore, percent changes in evaluation metrics such as accuracy must be specified as relative or absolute (i.e. percentage point). Finally, the authors need to control for the effect of the clean accuracy improvement: is the change in ImageNet-C accuracy larger than expected given the baseline change in ImageNet accuracy?\n\nAll the plots in Figure 2 should also contain a ResNet-152, with and without SAM.\n\nI understand that there are space constraints, but the ViT sparsity results should be depicted in a figure. Perhaps consider combining Figs. 2c&d to make more space.\n\nThe presentation of the contrastive learning results (Section 5.2) should include results for the baseline, ResNet-152.\n\nFigure 4a should show values both with and without SAM\n\nIt’s not clear to me how the sparsity of Mixer activations “suggests the potential redundancy\nof image patches.”\n\nThe authors should include a discussion of limitations of their work.\n", "summary_of_the_paper": "The authors analyze the effects of sharpness-aware minimization (SAM) when applied to vision transformers (ViTs) and MLP-Mixers. They find that the converged loss surface of ViTs and MLP-Mixers is sharp compared to ResNets, and that SAM ameliorates this issue, yielding improved validation accuracy on ImageNet. They then show that SAM improves the performance of ViTs and MLP-Mixers in a variety of image classification scenarios including adversarial attacks, naturalistic corruptions, contrastive training, and transfer learning. They also examine the effect of SAM on activation sparsity, activation maps, and the relationships between different model architecture components and loss surface sharpness.", "summary_of_the_review": "MLP-Mixers and ViTs are more difficult to train than CNNs. This study shows that SAM is very effective for training MLP-Mixers and ViTs, and attempts to understand why. These could all be important contributions to computer vision. I am convinced that SAM works well for ViTs and MLP-Mixers. But this paper is insufficiently rigorous, leaving me unconvinced of why it works. I don't think its suitable for publication in its current state, but it certainly could be after revisions.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "u37YePG1ovA", "reviewer": "Reviewer_y3Uf", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "+ The paper has conducted an abundance of experiments on different tasks to demonstrate the effectiveness of the sharpness-aware minimization (SAM). \n+ The authors present a comprehensive analysis on SAM. The differences and similarities between SAM and data augmentation are well explained. \n\n- This paper uses the existing idea of sharpness-aware minimization to alleviate the dependency of massive pre-training and strong augmentation, which makes the novelty limited.\n\n", "summary_of_the_paper": "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. ", "summary_of_the_review": "Overall, this paper is well writen. The overview of SAM and the explanation of some observations are clearly presented. Besides, the authors demonstrate the effectivenss of SAM in various tasks by showing a large amount of experimental results. The experimental part is comprehensive and convincing. The major drawback of this paper is that the SAM is an existing idea, even though the authors use it to overcome the issue brought by heavy pre-training and data augmentations. Therefore, this paper is like proving the effectiveness of an existing idea by doing many experimental validations. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "n3hRBpOuHMY", "reviewer": "Reviewer_CmdC", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper is well written and organized.  The motivation comes from the loss landscapes comparison for different network architecture such as ResNet, ViT and Mixer.  The visualization has revealed the problem of the current ViT and Mixer as both of them tend to converge at sharp minima which will limit the generalization capabilities.  As the issue has been clearly pointed out, the solution is straightforward: avoid the sharp local minima when training ViT.  Such that this paper adopts SAM as an answer.  Although, SAM is from existing work, this paper gives a good explanation on why among different optimizer, SAM could be a good choice. \n\nIn addition to the method, this paper has shown that on image classification,  contrastive learning and adversarial training ViT can be consistently improved with SAM.  Especially, the visualization of attention map w/ SAM, the ViT feature map does contain more semantics.   \n\nConcerns: I agree with that current ViT tends to converge at sharp local minima.  It could be the problem from the usage of existing optimizer.  On the other hand, the problem may also comes from the nature of current ViT design.  Recent vision transformers such as SWIN, PVTv2, VOLO, all of them has clearly shown that by optimizing the architecture of ViT, it can also avoid the local minima.  Based on current manuscript , it is unknown that how SAM can be used for SWIN or PVTv2 and whether SAM can be used as a general optimizer for arbitrary vision transformer training. \n\n[1] Swin Transformer: Hierarchical Vision Transformer using Shifted Windows\n\n[2] PVTv2: Improved Baselines with Pyramid Vision Transformer\n\n[3] VOLO: Vision Outlooker for Visual Recognition", "summary_of_the_paper": "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.", "summary_of_the_review": "In overall, I vote for acceptance for this paper but I think it can be further improved. \nIn the experiments, ResNet-SAM results are promising, the claim could be stronger if other vision transformers can be evaluated as ViT has been widely argued for its inefficient architecture design.  \n\n\n\n ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "jGlBRduuCXY", "reviewer": "Reviewer_4nyX", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is experimental, non-theoretical kind, mostly studyng the effect of SAM on ResNet, ViT and MLP-Mixer training and its interplay with dataset size and augmentation. Also paper studies how SAM influences different aspects of the trained network, namely - attention masks, eigenvalues of the Hessian of models blcoks and loss landscape.\n\nThe main stated result is that SAM can more or less replace the heavy image augmentations for ViT and MLP-Mixer (Result1) and improve accuracy on clean and corrupted test set. \n\nStrong points:\n\n- Impact of the Result1. I think, that despite of the paper simplicity (let’s apply SAM to the non-CNN architectures), the message that one can replace (domain-specific) augmentations with (general-purpose) optimizer is of great importance. All sota methods for image classification rely on the augmentations  to provide good results and it is not easy to come up with new augmentations. The most papers, which present “learned augmentations” in fact just learn some parameters and combinations of _known_ augmentation. However, for other domains, like NLP, or some kind of medical data, it is not clear, how one could augment the data. If one can instead just take general-purpose architecture (transformer) and train it with general-purpose optimizer (SAM), that is a big deal.\n- The experiments are done rigorously and with care. The only possibly missing experiment would be to train a model without an inductive bias at all, like vanilla MLP and check if SAM would greatly improve its results or not.\n- The paper also states the limitations of the approach, the main of which is computational cost\n\nWeak points: \n\n- In the current form (I was asked to review a paper after the rebuttal stage), I do not see any major weaknesses.", "summary_of_the_paper": "The paper presents an application of sharpness-aware minimization (SAM) to training of visual transformers (ViT) and MLP-Mixer. The idea of the paper is that models with weaker inductive priors (transformers and MLP-Mixer) suffer from sharp local minima more, than CNN-models and SAM fixes this issue, also significantly improving test accuracy on various version of ImageNet.", "summary_of_the_review": "The paper tackles a single problem and studies it extensively. The problem in my opinion is significant and the paper would likely have an impact. Therefore I vote for clear acceptance. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "Czu0wWYp0y", "reviewer": "Reviewer_hGsf", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The main contributions of the paper are: 1 ViTs with existing sharpness-aware minimizer (SAM) outperform ResNets of similar size and throughput when trained from scratch on ImageNet without large-scale pre-training or strong data augmentations. 2 The motivation and experiments, including the ablation studies are well-organized and adequate. The paper reveals that the ViTs and MLP-Mixers have extremely sharp local minima of converged models through loss landscape visualization and Hessian dominate value.\n\nThe impact of the paper may be limited. The paper only compares with ResNet-152 with no augmentation and large dataset pre-training. Currently, the augmentation and large dataset pre-training are vastly employed in the entire community. The impact can be enlarged by conducting experiments on ViT-L and comparing ResNet152x4 pretrained on BiT-L.", "summary_of_the_paper": "The paper is clearly written with extensive experiments. The paper compares dataset size, data augmentation, and different network architectures, ViTs, ResNets and MLP-Mixers by using loss landscape visualization, number of parameters, neural tangent kernel (NTK), Hessian matrix, training dynamics, percentage of activated neurons, throughput, norm of weights and activations, attention map visualizations and missing rate under linear interpolation.", "summary_of_the_review": "The paper is pretty well-written with extensive results and figures. The impact of the paper can be largely improved if the large-scale experiments can be conducted and a better accuracy than the current state-of-the-art methods can be achieved.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: ### Description 
The paper demonstrates that efficient architectures such as transformers and MLP-mixers, which do not utilize translational equivariance in the design, when regularized with SAM (sharpness aware minimization) can achieve same or better performance as convolutional networks, in the vision problems where the convolutional networks were traditionally superior (with data augmentation or not, regularized or not). The paper demonstrates it very thoroughly through many experiments and analysis of the loss surfaces.

### Decision + Discussion

I find the paper to be very timely in its context. It has a remarkable coverage of experimental studies and different use cases: SAM + augmentation, +contrastive, +adversarial, +transfer learning; as well as ablation studies such as keeping first layers convolutional. The reviewers have asked further questions, and the authors were able to conduct respective experiments within the discussion period fully addressing all concerns and making the findings of the paper even more comprehensive and convincing.

After the rebuttal 3 reviewers were for "accept", one "marginally above" and one "marginally below". In the latter case the concern was that the paper is an experimental study of a known method, SAM. While I understand that many researchers are expecting theoretical and innovative results from ICLR papers, I find that it does not prevent acceptance. Indeed, the experimental findings in this paper are on a "hot" topic, could be of wide interest and could lead to a change of paradigm in designing models towards more generic ones. On the other hand, it could just indicate that CNNs are not fully exploiting their potential, e.g. not exploiting the context well enough in the hidden layers?

To get more insight, I am still wondering, how the predictions behave if the input is shifted by a few pixels in CNN and Transformers? It seems counterintuitive that making the first layers in ViT just an MLP of image patches is a good design. Furthermore, fully convolutional models allow to take input of an arbitrary size and average the predictions on the output if it happened to be larger than 1x1. 

Since convolutions are also used for e.g. semantic segmentation and generative models, one should not (and the authors do not in the paper) discard them too fast. See also a recent work combining transformers and convolutional networks,
Chen et al. (ICCV 2021) Visformer: The Vision-friendly Transformer.
Decision: Accept (Spotlight)