Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 8: accept, good paper
review_content: 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.
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.

The 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. 

Strong points:

- 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.
- 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.
- The paper also states the limitations of the approach, the main of which is computational cost

Weak points: 

- 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_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.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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.
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.

The 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_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.
review_title: null

final_score: rating: 8: accept, good paper
review_content: 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.
main_review: **Update:** I have increased my score from a 3 to an 8 after the the authors' response

SAM 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.

I 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.

# Detailed Feedback

## Replicates and measures of variability

Ideally, 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).

## Unjustified claims

There are numerous unjustified scientific claims made, which I detail as follows:

End of first paragraph in **Section 3**:
>“There also exists a large gap between ViTs and ResNets in robustness tests. 

Requires a reference.

In **Section 3. Small training errors** the authors claim 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.

This 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.

In **Section 4.2. Higher Accuracy**, the authors state: 
>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.

This 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).

The 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*).

The authors state

>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

Claims 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”.

Table 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".

Transformers 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.

Does 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.

The following claims are made in the introduction, some of which I address in my above comments:

>We conjecture that the convolution-induced translation equivariance and locality help ResNets escape from bad local minima when trained on visual data.

This 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.

>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.

This 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.

>The weight norms increase, implying the commonly used weight decay may not be an effective regularization alone.

Changes 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.

>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.

It’s not clear to me how sparsity translates to image patch redundancy. This is a very interesting idea, and should be better explained.

>Another interesting finding is that ViTs’ performance gain also translates to plausible attention maps containing more perspicuous information about semantic segmentation. 

I address this in an earlier comment: evaluate this claim with a segmentation task and/or human interpretability experiments.

## Other Feedback

Analyzing 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.

The authors state that

>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.

This 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.

In *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?

All the plots in Figure 2 should also contain a ResNet-152, with and without SAM.

I 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.

The presentation of the contrastive learning results (Section 5.2) should include results for the baseline, ResNet-152.

Figure 4a should show values both with and without SAM

It’s not clear to me how the sparsity of Mixer activations “suggests the potential redundancy
of image patches.”

The authors should include a discussion of limitations of their work.

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.
review_title: null

final_score: rating: 8: accept, good paper
review_content: 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.
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. 

In 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.   

Concerns: 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. 

[1] Swin Transformer: Hierarchical Vision Transformer using Shifted Windows

[2] PVTv2: Improved Baselines with Pyramid Vision Transformer

[3] VOLO: Vision Outlooker for Visual Recognition
summary_of_the_review: In overall, I vote for acceptance for this paper but I think it can be further improved. 
In 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.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: 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. 
main_review: + The paper has conducted an abundance of experiments on different tasks to demonstrate the effectiveness of the sharpness-aware minimization (SAM). 
+ The authors present a comprehensive analysis on SAM. The differences and similarities between SAM and data augmentation are well explained. 

- 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.


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.
review_title: null

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)