Normalization Layers Are All That Sharpness-Aware Minimization Needs

Summary

The paper proposes an adversarial perturbation method linked to SAM for the affine normalization parameters, in contrast to perturbing the full set of parameters. The results show this approach improves upon standard SAM and prior sparse SAM approaches.

Strengths

The SAM-ON proposed method yields several benefits: 1. improves upon SAM-all and carious other SAM-variants 2. achieves this by perturbing only the normalization layers, which correspond to a small percentage of the total parameters

Weaknesses

Weaknesses of this paper include: 1. gains in Table 1 are marginal 2. originality of approach is lacking given prior literature on SAM variants 3. justification of SAM-ON is poor; the experiments in 5.2 and Table 7 try to poke at this but fall short as the authors do not provide a solid explanation as to what this can be attributed to; they only show that SAM-ON increases sharpness in L-inf sense; but what about other metrics of flatness? is there a more universal metric that can explain the effectiveness of SAM-ON? 4. it is unclear whether SAM-ON is truly better than other SAM variants in out of distribution robustness, e.g. WILDS benchmark 5. results are only shown for CIFAR and ImageNet datasets; more datasets and benchmarks need to be considered

Questions

see weaknesses above One of the questions this paper fails to answer is why sharpness is higher for SAM-ON, even though its formulation is motivated to minimize sharpness. This is not a complete enough story and opens up lots of questions that need to be addressed.

Rating

3: Reject: For instance, a paper with technical flaws, weak evaluation, inadequate reproducibility and incompletely addressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

Somewhat discussed but limited. Would like to see a more thorough limitations discussion.

Summary

As Sharpness-Aware Minimization (SAM) aims to regularize the flatness of the loss landscape for better generalization, this paper shows that only perturbing the normalization layers is sufficient to achieve this. To prove this, the authors first propose a method called SAM-ON (SAM-OnlyNorm). Then, they conduct experiments on CIFAR datasets across different models and variants of SAM, comparing the test accuracy with and without applying SAM-ON. The experiments (in most cases) show that SAM-ON outperforms vanilla SAM, as well as SAM variants with ON outperforming SAM variants. Furthermore, since SAM-ON only needs to perturb a few parameters, it can save computational resources. Finally, this paper conducts numerous experiments to understand SAM-ON and concludes that sharpness may not be the key to SAM. ## post-rebuttal I've updated my score since my concerns are adequately addressed.

Strengths

1. This paper demonstrates a new understanding of the underlying mechanism of SAM, which is both novel and informative. 2. The proposed SAM-ON can improve efficiency and generalization in most cases. 3. The experiments are adequate to support the claims. 4. The code is provided, which guarantees reproducibility.

Weaknesses

1. SAM-ON cannot outperform SAM (or its variants) in some cases, showing that SAM-ON may not always be effective and decreasing the soundness of the claims that support SAM-ON. 2. The improvement of SAM-ON on ImageNet seems to be consistently lower than on CIFAR datasets for different architectures. Therefore, this reviewer hypothesizes that SAM-ON only works for small datasets and is not very effective for large datasets. 3. As the author acknowledges, this paper only investigates SAM and SAM on vision data. Thus, the analysis cannot support the proposed understanding that generalizes to other tasks, such as language tasks.

Questions

See Weaknesses.

Rating

9: Very Strong Accept: Technically flawless paper with groundbreaking impact on at least one area of AI/ML and excellent impact on multiple areas of AI/ML, with flawless evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

N/A

Summary

This paper relates the normalization layer with the Sharpness-Awareness Minimization (SAM). Surprisingly, this paper finds that in the perturbation stage, only perturbing the affine parameters of normalization layers in the networks leads to a better generalization performance. Later, the authors investigate the reason behind such surprising phenomenon, and find that SAM-ON(only perturbing the normalization layer) even increase the sharpness, which doubts the intuition that the effectiveness of SAM comes from sharpness minimization. Extensive experiments demonstrate the effectiveness of the methods this paper proposed.

Strengths

1. This paper identifies an important yet interesting problem: Is Sharpness-Awareness minimization effective because of minimizing sharpness? This paper found that only perturbing the normalization layer in the perturbation stage of SAM leads to both superior generalization performance and higher sharpness. The interact relationship between sharpness and generalization are quite doubtful. 2. This paper investigates the SAM from the normalization layer and found the normalization might be the reason for the success of SAM, which is quite surprising. 3. SAM with only perturbing the normalization layer largely saves the computational cost than the original SAM which needs perturbing all the parameters.

Weaknesses

1. These experiments should be conducted on more network architectures, including VGG-Net and etc. 2. The observation that SAM with only perturbing the normalization layer holds the higher sharpness than original SAM might be trivial. SAM are designed to reduce the sharpness, while a degraded version of SAM should hold a higher sharpness in principle. 3. In this paper, it is clear that there is deep relationship between normalization layer and generalization, and there are indeed some papers on the relationship between normalization layer and generalization. Therefore, a related discussion should be included.

Questions

1. In fact, the authors don't explain why SAM with only perturbing normalization layer works better than original SAM. Therefore, I would like the authors to include some discussion or theoretical analysis on that.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes. This paper discusses the limitation of this paper, where they only focus on the vision data.

Summary

This paper tries to analyze the effects of different layers for sharpness-aware minimization (SAM). The authors find that normalization layer plays an important role in the improvements of SAM and only perturbing the normalization layer can obtain a comparable result. Therefore, this paper proposes SAM-ON and the experimental results illustrate that SAM-ON can obtain a great performance on many different tasks with different models.

Strengths

Strengths: 1. This paper focus on an important problem, SAM is very important for improving generalization. 2. The experimental finding is very interesting and can efficiently reduce the computation. 3. The authors try to evaluate the performance of proposed method on various datasets and models.

Weaknesses

Weakness: 1. Although the proposed method can improve the performance on CIFAR. But CIFAR is too simple for SAM and that is weak to illustrate the improvement of SAM-ON. However, the results about resnet on imagenet is too close for me and the improvement is minimal. I'm not sure whether you run the experiments multiple times with different seeds. 2. The experiments on ViT is also a little weak for me. For ViT, cifar is too simple and you should train ViT form scratch on imagenet. Since SAM can achieve significant improvement on ViT+ImageNet and maybe you need to analyze it to obtain more convincing results. 3. Although SAM-ON can reduce the computation, the training efficiency cannot be significantly improved.

Questions

Question: 1. Do you try to run the experiments for multiple times with different seeds. 2. You say only perturbing normalization layers can obtain similar results with perturbing all parameters. My question is do you try to only perturb the parameters outside the normalization layer?

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

N/A

Thank you for response

I have thoroughly read the response and appreciate the detailed response to my questions. - I have noticed the additional experiments on VGG-nets and DenseNets and the experimental results are sound to me. - I agree that it is worth investigating that SAM-ON could still be able to reduce sharpness. - I agree that it is important to study the underlying relationship between SAM and BN and would like to any future work on that. Regarding the authors' response and other reviewers' comments, I have raised my rating.

Thanks for your effort to make this paper more clear!

Thanks for your response and effort to make this paper more clear! I still have several concerns: 1. ImageNet results: I find the performance gain about imagenet training from scratch in your attached pdf is a little limited. For example: top 1 ACC: 71.33 -> 71.41. In my past experience about SAM, the improvement is usually significant for ViT training, such as gsam. So I think maybe the proposed method is not vet suitable for vit training. In addition, I find the most experiments are about ResNet, such as table 1,2,3. However, the accuracy on imagent is very close (SAM vs SAM-ON), although the improvement on cifar is clear. So I think maybe the proposed method can work well on simple tasks, but it is still not very clear for more complex tasks. Finally, I think big model may be easier to converge to sharp minima and we need to focus more on these big models and datasets. CIFAR is not enough for me. 2. Only perturb the parameters outside the normalization layer: Thanks for your clarification. As you mentioned, no-norm typically performs similarly to SAM-all. Based on that, whether I can say the parameters outside the normalization layer provide the most properties that SAM can improve the performance? So that means the parameters outside the normalization layer are all you need and the normalization layer has some special properties. If that is true, maybe you need to provide more analysis. Overall, I think these paper mainly focus on the experimental analysis about different layers for SAM. So I think the author need to provide more results and analysis to make us convincing.

Summary of new experiments

Dear all, We would like to thank you all for your time and helpful remarks and suggestions. In order to facilitate your decision, we would like to summarize how we addressed your comments from both the reviews and the discussion phase: 1. We trained ViTs from scratch on ImageNet with both the AdamW and Lion optimizer and found consistent gains of SAM-ON over SAM-all and the vanilla optimizer (Table 2 in rebuttal pdf). To the best of our knowledge, we are the first to investigate the Lion+SAM combination. (Ftd6 asked for _”ViT form scratch on imagenet”_) 2. We additionally evaluated those models on the OOD datasets ImageNet-R and ImageNet-C, where again SAM-ON showed the best performance (reviewer sERj asked for _”out-of-distribution robustness”_) 3. We trained a range of VGG architectures and a DenseNet on CIFAR10 and CIFAR100 (Table 1 in rebuttal pdf), where SAM-ON outperformed SAM and the vanilla optimizer consistently for all investigated SAM-variants. (reviewer wExQ asked for more _“more network architectures, including VGG-Net”_ ) 4. We provided results with multiple seeds for the ResNet on ImageNet numbers. For ViTs, the runs haven’t finished yet, and we will provide them in the revised version. (reviewer Ftd6 asked for _”multiple seeds [for ImageNet]”_, Table is reported in response to their comment) 5. We report a range of other sharpness metrics in Table 3 (rebuttal pdf) and find in alignment with our previous findings that SAM-ON is sharper than SAM-all with respect to most metrics (especially worst-case sharpness and the metrics that are optimized during SAM), although there exist some exceptions. (reviewer sERj asked for _”other metrics of flatness”_) 6. We show that further computational gains may be achieved for our approach by applying SAM to only part of the normalization layers, cutting SAM’s additional cost roughly in half. Although this leads to some reduction in test accuracy, the approach still strongly outperforms the base optimizer (see Figure 1, bottom right in rebuttal pdf, half of additional cost corresponds to layer3 in purple). (reviewer Ftd6 commented on the training efficiency) 7. We find that in order to get SAM-like improvements for adversarial robustness (as shown in arxiv:2305.05392) it is enough to only perturb the normalization layers in SAM, illustrating again their special role as outlined in our paper. (reviewer dJUQ pointed us at a possible connection to adversarial robustness, Table is shown in response to their comment) 8. We provide experiments in Figure 1 in the rebuttal pdf that explore potential connections of SAM-ON with weight decay and dropout. We find that SAM-ON’s improvements are not due to its interaction with weight decay (Figure 1, bottom left), nor can its success be mimicked by applying dropout solely to the normalization layers (Figure 1, top). (Reviewer Ftd6 and wExQ asked for additional analysis) Overall, we believe that SAM-ONs success is highly surprising and an important finding that should be shared with the community as it provides new insights into SAM: The original motivation behind SAM is a sharpness measure defined with respect to all parameters, and minimizing the respective measure only for an extremely small fraction of the total parameters should intuitively not lead to the same or even improved generalization benefits. We would further like to stress that the mechanism behind vanilla SAM’s success is still not well understood and that our work lends support to the idea that sharpness-reduction may not be the sole source of it, which fits well into recent findings (arXiv:2302.07011 and arxiv:2206.06232). Our work hence contributes to a search for alternative explanations by illustrating the success of only applying SAM to the normalization layers. Thanks again for all your time and useful suggestions! We believe the new experiments provided in the rebuttal have made the paper even stronger, and we hope you will take this into consideration when making the final decision. The authors

Copy: Official Comment by Reviewer dJUQ

Dear Authors, Thanks for the rebuttal. I believe that most of my concerns have been addressed, and I would like to raise my score when review editing is allowed. Here are some further questions for curious. In the context of adversarial robustness, a recent work aims to understand SAM as implicit adversarial training and revealed that SAM is beneficial for improving robustness [1]. On the other hand, there are several works attributing adversarial vulnerabilities of DNNs through normalization layers, and designed normalization layer-based adversarial training methods to improve adversarial robustness [2,3]. My questions are: 1. Can SAM-ON and the corresponding observations in this paper be the connection between normalization layers and adversarial robustness? 2. Can using SAM-ON only improve adversarial robustness? Feel free to dismiss these questions if you are not familiar with adversarial robustness. However, I believe discussing these questions in this paper can help you strengthen the understanding of the effectiveness of SAM-ON. [1] Sharpness-Aware Minimization Alone can Improve Adversarial Robustness. ICML 2023 workshop [2] Intriguing properties of adversarial training at scale. ICLR 2020 [3] Batch Normalization Increases Adversarial Vulnerability and Decreases Adversarial Transferability: A Non-Robust Feature Perspective. ICCV 2021 --- Dear authors, please paste the response below this comment. Thanks