Venue: ICLR 2021 Conference

Reviews:

final_score: rating: 7: Good paper, accept
review_content: Motivated by the connection between the flatness of minima and its generalization ability, the authors propose Sharpness-aware Minimization (SAM), which explicitly minimizes both loss value and loss sharpness during training deep neural networks. They find SAM improves generalization for a range of image classification tasks and provide robustness to label noise as well. They also introduce a new notion of sharpness named m-sharpness.

Strength: 
* The paper is overall well written with clear motivation.
* The experiments are comprehensive and the results show clear improvement over non-SAM approaches or previous SOTA.

Weakness:
* There is no clear definition of the “sharpness“ that the algorithm tries to optimize. Given many existing definitions of the sharpness (e.g., [1]), it is not clear how the proposed measurement connects or differs with previous works. 

* My major concern is about the usage of hyperparameter $\rho$:

a) The introduction of the dataset and model dependent hyperparameter $\rho$ and the need of grid-search before training makes the algorithm more tricky to work and sensitive to other hyperparameters and scale of $w$, e.g., when weight decay is applied, the norm of $w$ usually shrinks during training, and the same radius $\rho$ could be too large for a small scaled $w$ at the end of training in comparison with the $w$ at the beginning. This discrepancy would become larger when the number of epochs training gets larger. 

b) The details for how to obtain the optimal $rho$ is not quite clear, e.g., smaller $\rho$ in sec 3.3. An ablation study on the sensitivity of $\rho$ regarding different dataset, model and noisy level would be useful. 

c) The wall-clock training time of the SAM method is not discussed. A comprehensive of the cost (including hyperparameter search for $rho$.) would be helpful to have for evaluating the complexity of the method. 


* The message conveyed in section 4.1 is not quite clear. Does each accelerator perform independent $\epsilon$ estimation? Is $epsilon$ obtained on each accelerator synchronized after their estimation? Does it indicate the SAM training is done better in model-parallel in small batches rather than data-parallel with large batches? 


Suggestions:
1) To avoid the scaling issue of $\rho$, one suggestion would be considering optimizing the sharpness metric on the normalized loss landscape as described in [2]. In Figure 1, the authors adopt [2] for comparing the landscape of minimas obtained by non-SAM and SAM, so it might be intuitive to optimize this normalized sharpness directly, in which $\rho$ can be fixed and random direction is sufficient?


2) The benefit of flatness to the robustness to label noise is not well discussed. What is the performance when the label noise is over 90% or even 100%. Eventually all models should not generalize given 100% corruption but it would be interesting to know where the limit of SAM is.

Minor:
* Some figures are not well described, e.g., the meaning of Figure 1 left is not quite clear.  Figure 2 is not intuitive as the loss contour value is not clear. It is not straightforward to know why w_{t+1}^{SAM} is a better or “flatter” move. The notion $w_{adv}$ is also not defined anywhere.

[1] Keskar et al, On large-batch training for deep learning: Generalization gap and sharp minima, ICLR 2017
[2] Li et al, Visualizing the Loss Landscape of Neural Nets, NIPS 2018

======
After Rebuttal

Thanks for the detailed reply and additional experiments. I increased my score accordingly and I hope the authors could further address following issues:

- While the results in C.3 shows default $\rho$ improves over SGD on most experiments (may also add SVHN and Fashion), I can still see its sensitivity to datasets, architecture, noise level and number of accelerators as shown in Table 6, 7, 8 and Fig. 3. For example, 0.05 is not close to optimal with labe noise 20%~60% in Table 8. It is unclear whether $\rho$ is robust to other hyperparameter changes (e.g., weight decay that controls weight scales).  So an ablation study on the sensitivity of $\rho$ and further explanation would be necessary and much valuable for practitioners. 

- It would be also helpful if the authors can provide more details about how to get the flat minima of Fig .1 (right) when optimizing deep non-residual networks, such as $\rho$ and other hyperparameters.

- Minor: Table 8 should be validation errors rather than accuracy.

review_title: Interesting work with good results, concern is on selecting the right $\rho$

final_score: rating: 6: Marginally above acceptance threshold
review_content: - Overview

This paper proposed a learning algorithm using the idea of flatness.
Basically, the algorithm was constructed by a first-order Taylor approximation of flatness and using only the main term.
The developed algorithm uses the gradient of the loss with steepest direction when updating the parameters.
The proposed algorithm is simple and fast to work.
It is also easy to incorporate into existing networks and algorithms and has updated accuracy in many models.

- Comments.

This is an interesting study, as it is the first application of flatness to a practical algorithm.
It is also excellent in that it is easy to implement and immediately applicable. The performance is also nice.

The validity of this approximation is a concern.
When the authors perform a first-order Taylor approximation, they are looking only at the gradient of the parameter to determine the adversarial direction.
So, if the loss function is very spiky, there is a concern that the approximation will not work because the local gradient alone will not find the appropriate direction.
The nice networks used in their experiments is supposed to use smoothing techniques for the loss surface, such as the batch normalization and the skip connections, but the concern is whether the proposed method works properly in neural networks with a more less smooth loss surface.

To address this concern, my suggestion is to use the algorithm in without the Taylor approximation make comparisons on a more primitive network with a less smooth loss surface.
That way, users would feel more comfortable using the method.
review_title: Interesting work. It would be good to investigate its property.

final_score: rating: 8: Top 50% of accepted papers, clear accept
review_content: ## Summary ##
The paper proposes a modified loss function for supervised learning, in which the original loss at w is replaced with a maximum of the loss in a small p-norm ball around w. An approximate way to compute gradients for this loss is presented, and evaluated in high detail on a variety of supervised learning problems where it is shown to consistently improve the overall generalization error. Furthermore, based on PAC-Bayes theory, a generalization bound for learning under that loss is presented. 

## Explanation of Rating ## 
The main strengths of the paper are the simplicity and convincing evaluation of the method. Another strength is the sound theoretical generalization bound. The paper is also very well written, and I therefore clearly recommend acceptance. 

Perhaps a weakness of the paper is the lack of a "broader/high-level perspective" on the method (see detailed comment #1) and its comparison to variational inference methods which also optimize PAC-Bayes bounds (see detailed comment #2). 

## Detailed Comments ##

1. In the past, Gaussian convolution smoothings of loss functions have been considered a lot in the context of homotopy continuation methods or variational inference. To me, the proposed loss function \max_\eps L(w + \eps) - \delta_{|\eps| < \rho} can be seen as a convolution of the loss, but on the tropical semiring (max, +) with the convolution kernel being the indicator function \delta_{|\eps| < \rho}. Such types of convolution have been heavily studied in the field of convex analysis, where they are known under the name infimal convolution / epi-addition. This opens up the question on why this specific convolution kernel has been chosen. The practical results in this paper are quite strong and have been elusive so far for variational inference and Gaussian smoothing methods, even though they have been around for a long time. Do these good results mainly stem from the efficiently implementable algorithm, or are they more due to favorable geometrical properties of the loss due to performing the convolution in the tropical geometry? 

2. What is the main advantage of the proposed approach over a variational inference (VI) method (e.g. with mean-field Gaussian approximation) to minimize the PAC-Bayes bound? In variational inference, the free parameter \rho is also be optimized by minimizing the right-hand side of the PAC-Bayes bound. Have you tried minimizing over \rho, or perhaps consider a diagonal approximation of \rho? 

Minor comments / typos:
- p.12 "Adding a Gaussian perturbation should increase the test error" -> "Adding a Gaussian perturbation should not decrease the test error."  This is expected to hold in practice at a minimum. But it is not clear if it holds for any w. 
- p.12 Typo: "Then we KL divergence" -> "Then the KL divergence". 
- The steps (12) and (13) in the proof are not easy to follow. In particular, it was unclear to me what meant with "each bound" before Eq. 12, and why the bounds should hold with this specific choice of probability. Furthermore, I couldn't follow what happens from (12) to (13). 
- I did not notice where the assumption ||w|| >= 1 is used, perhaps it can be mentioned in the proof. 
- \lambda is another hyperparameter, but it is not mentioned later on how it is chosen. 
review_title: Clean and simple method, strong empirical results. 

final_score: rating: 8: Top 50% of accepted papers, clear accept
review_content: ## Summary

This paper proposes and empirically evaluates SAM, an optimization method that is designed to seek out regions of uniformly low training loss. The method is derived from a bound on the generalization performance of parameters $w$ in terms of the maximal training loss in a region around $w$. After various approximations, minimizing this upper bound gives rise to a simple method, which first performs a (normalized) gradient ascent step; computes the gradient at that perturbed location; and uses that gradient to update the weights. The empirical evaluation shows that SAM improves generalization performance across a wide range of settings.

## Rating

The paper is well-structured and easy to read. The proposed method is motivated from a PAC-Bayesian generalization bound, but involves approximations which are only vaguely justified. The empirical evaluation is extensive and rigorous. Given its simplicity, scalability, and the convincing results, I believe that this method could have significant impact. With a small exception (see below) all derivations are, to the best of my knowledge, mathematically sound. Overall, this is a good paper and I recommend acceptance, but there are various aspects that could be improved. If these are addressed during the rebuttal, I will consider increasing my score.

## Major Comments

1) The proposed method is motivated by a generalization bound, but this bound is never evaluated numerically. Prior work has shown that many generalization bounds are numerically vacuous in deep learning settings. The bound in Theorem 1 should be really straight-forward to evaluate after the end of training and I think this would be a vital addition to this paper. The paper shows convincing practical benefits of SAM, but it is important to know whether this benefit actually stems from the (approximate) minimization of a non-vacuous bound. This is my single biggest concern with this work.

2) The derivation of the SAM update relies on various approximations to turn the intractable bi-level optimization problem of Eq. (1) into a practical method. My concern here is that the paper does not even attempt to justify these approximations and/or to discuss their limitations and possible pitfalls. I’m not saying that these approximations are not justified; I just think that a bit of discussion would significantly strengthen the paper. A few thoughts:
    a) The linear approximation that facilitates the closed-form solution of the inner optimization problem (top of page 4) is really drastic. It would be accurate if and only if the curvature in a $\rho$-sized region around w would be negligibly small. That, however, is exactly the type of region the method wants to find, so this is somewhat circular. Contrasting the first-order with a second-order approximation may be a useful exercise. An interesting ablation would be to solve the inner loop optimization in Eq. (1) more accurately, say, by performing multiple iterations of (projected) gradient descent. (When dropping the second-order term, this wouldn’t be exceedingly costly and could provide further insight.)
    b) Later on, the “second order” term, which derives through the solution $\hat{\epsilon}$ of the inner-loop optimization, is dropped. The authors write that this is done to “further accelerate the computation” and “does not adversely affect model training and in fact improves numerical stability”. However, as shown in Appendix C.3, dropping the second-order term yields a quite considerable *improvement*. It always makes me a bit nervous, when a (relatively crude) approximation works better than the quantity you actually want to compute. I don’t think it is adequate to sweep this under the rug with a vague reference to numerical stability. At least, I thank that (i) the fact should be stated clearly in the main text, and (ii) the reference to numerical stability should be marked as speculation, unless there’s evidence to back it up.

3) Even if the approximations involved are only vaguely justified, the SAM update itself makes a lot of intuitive sense, and I think the reader might benefit from discussing that in a bit more detail. For example, one simple observation would be that, for a sharp minimum whose basin of attraction is of radius smaller than rho, the ascent step of SAM (by virtue of the normalisation to length rho) would take you outside of that basin. If the basin of attraction is larger than rho, the SAM update stays inside. Maybe this intuition could be formalized?

4) Connection to Entropy-SGD:
    a) Is it really fair to characterize EntropySGD as “only suitable to small models and datasets”? Of course, it requires a number L of additional gradient steps to obtain an MC-estimate of the gradient of the local entropy. But, if my understanding is correct, these could be parallelized. The original paper uses L=20 but, to my knowledge, nobody has tested the (lower) limits of that number. If L<=5 does something useful, it wouldn’t be too outlandish in terms of computational cost.
    b) On that note: It would have been great to see an empirical comparison to EntropySGD for at least one of the experiments. 
    c) Looking more closely at the proof of the generalization bound, the term involving maximization over a local region actually originates from a bound on the *integral* over that region, which makes the connection to EntropySGD even closer. Is it fair to say that the maximization step in SAM is actually just a clever and cheap surrogate for integrating the gradient the over the local region as done in EntropySGD?
    d) An interesting piece of related work that you might want to add is [1], which shows that EntropySGD (in a specific sense) also optimizes a generalization bound.

5) You clearly show the benefits of SAM in pushing the state-of-the-art, which is fantastic! Personally, I would have also valued experiments that are more illustrative and maybe test the limits of the method:
    a) What does SAM do in a convex setting, say, a logistic regression on some (maybe synthetic) dataset that is easily overfit? Does it converge to the unique minimizer of the training loss? Does it plateau before reaching that?
    b) What happens if you apply SAM in the setting of [2] with random labels assigned to images? Will it converge to a point that closely fits each (nonsense) training sample?
    c) With regards to the finding on m-sharpness, it would be great to see a small-scale experiments with full-batch optimization.

6) The finding on m-sharpness is intriguing, though somewhat counter to the motivation for SAM laid out in the paper. The results in Figure 3 make me wonder, whether SAM would even gives practical improvements with large m? This is an important question, but I totally agree with the authors decision to leave it for future work.

## Minor Comments

7) SAM is motivated from the generalization bound in Theorem 1. The bound includes a term $h(||w||^2 / \rho^2)$, which is later replaced with a simple L2 regularization term with coefficient lambda. From the Theorem, that coefficient $\lambda$ should be coupled with the neighborhood size $\rho$ in an inversely proportional manner. Of course, for a fixed choice of $\rho$, this can be subsumed into the regularization parameter lambda. However, it would make a difference when comparing various choices for rho, as is done in the experiments. Did you keep the regularization parameter lambda constant in those experiments or did it vary with rho, as suggested by the Theorem? Furthermore, it would be interesting to know to what extent SAM depends on that L2 regularization. Did all experiments use L2 regularization and did you, by any chance, test SAM’s performance without the regularization parameter?

8) The bib file could really need some love:
    a) You cite arXiv versions of of multiple papers that have been published at peer-reviewed venues.
    b) Capitalize properly: “nesterov momentum”, “pac-bayesian”, “journal of machine learning research”, etc…
    c) Cite identical venues consistently.

9) An interesting connection (or non-connection) to discuss in the paper would be adversarial training. Some of these methods perform an adversarial perturbation step to the *data*, followed by a weight update computed on that perturbed data. This is an interesting parallel to SAM, where the first step is sort of an adversarial perturbation of the weights.

10) If I am not mistaken, there is a minor flaw in the proof of Theorem 1. The very last equation gives a high-probability bound on the norm of $||\epsilon||^2$, which is then plugged into the high-probability generalization bound of Eq. (). Both bounds hold with probability 1-delta individually, but that does not imply that the combined bound holds with probability 1-delta. However, that should easily be fixed with a slight adaptation of the constants.

## Typos / Style

- At least according to some style guides, I think that you should write “Sharpness-Aware Minimization” instead of “Sharpness-aware Minimization”.
- Capitalize references to sections, equations, figures: “in Appendix C.3”, “using Equation 3”, et cetera…
- Typo “opefrations” near the bottom of page 4
- Footnote markers that refer to an entire sentence go after the punctuation mark (e.g., for footnotes 3 and 5).
- In Table 3, the boldface is assigned to the wrong method in the Stanford Cars row.

## References

[1] Dziugaite, G. K., & Roy, D. (2018). Entropy-SGD optimizes the prior of a PAC-Bayes bound: Generalization properties of Entropy-SGD and data-dependent priors. In International Conference on Machine Learning (pp. 1377-1386). PMLR.

[2] Zhang, C., Bengio, S., Hardt, M., Recht, B., & Vinyals, O. (2016). Understanding deep learning requires rethinking generalization. arXiv preprint arXiv:1611.03530.

## Update after Rebuttal

Thanks for the detailed replies to my questions and comments. I think the paper has been improved substantially and I have increased my rating. Congratulations on the good work!
review_title: Good paper, but considerable room for improvement

Meta-review: The paper proposes to minimize the loss while regularizing its sharpness: so that the minimum will lie in a region with uniformly low loss.
The reviewers uniformly appreciated the paper. They have made a number of suggestion for improving the paper, which the authors should consider incorporating in their final version.


Decision: Accept (Spotlight)