The recently proposed optimization algorithm for deep neural networks Sharpness Aware Minimization (SAM) suggests perturbing parameters before gradient calculation by a gradient ascent step to guide the optimization into parameter space regions of flat loss. While significant generalization improvements and thus reduction of overfitting could be demonstrated, the computational costs are doubled due to the additionally needed gradient calculation, making SAM unfeasible in case of limited computationally capacities. Motivated by Nesterov Accelerated Gradient (NAG) we propose Momentum-SAM (MSAM), which perturbs parameters in the direction of the accumulated momentum vector to achieve low sharpness without significant computational overhead or memory demands over SGD or Adam. We evaluate MSAM in detail and reveal insights on separable mechanisms of NAG, SAM and MSAM regarding training optimization and generalization. Code is available at https://github.com/MarlonBecker/MSAM.
Paper
Similar papers
Peer review
Summary
In this paper, the use of momentum in Sharpness Aware Minimization (SAM) is investigated and then a speed-up method is proposed. The main idea is to use the capability of momentum on lookahead to reduce one backward calculation in each iteration.
Strengths
+ The computational time of SAM is almost twice of that of SGD, which is due to twice backward calculation used in SAM. Thus, the idea of using momentum to estimate the gradient and to reduce one backward calculation is interesting. + The paper is generally well organized and some explanation, e.g., Fig. 1, is clear.
Weaknesses
- the mechanism of momentum has not been clearly presented.If we look two successive iterations together, we could find that using a momentum in the next iteration almost equals to use a larger learning rate in the previous iteration. Then I can understand that why the author claim that "the loss is expected to increase" when "a well-chosen learning rate" is used, which has similar effect on attack used in SAM. However, if this is the reason, then there will be some other unclear points, see question 1 and 2. - the numerical experiments need improvements: (1) the general performance is not very good, comparing with the performance reported in recent papers. (2) as a method that focusing on speeding up SAM by cutting one backward calculation, it is should be compared with random perturbation based methods, which have almost the same speed-up performance and mechanism.
Questions
- Using momentum in the next iteration has similar performance as taking a larger learning rate in the previous iteration. But it is based on the fact that the "a well-chosen learning rate" is used. In practice, the learning rate is not necessarily optimal, thus I do not think increasing it can always increase the loss value, as attack procedure in SAM. In Section 4.1, the authors "empirically show that the momentum descent step actually results in an ascent on the per-batch loss", however, by showing the cos similarity (which is not directly related to value, due to the heavy non-linearity of DNN's training landscape). Can the author report the value change directly, also for different learning rate? - In a paper, https://arxiv.org/pdf/2309.15639v2.pdf, recently accepted in NeurIPS2023, an opposite conclusion is given: one need to add momentum to enhance the attack in SAM. Following their idea, even when we do not want to perform backward calculation, we need to add a momentum to increase not to decrease the value. So can the author discuss the link to that paper? (generally I do not want to ask the authors to compare with recent papers, but they two are strongly linked, so I think it is better to discuss here) - Could the author provide a formulation similarly to NAG on which lookahead direction is calculated. - Could the author include comparison with other speed-up method for SAM and also please notice that the current reported performance for other method also the baseline SAM is not very good.
Rating
3: reject, not good enough
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
2 fair
Presentation
3 good
Contribution
2 fair
***Weaknesses*** 1. We agree, that the momentum factor $\mu$ is heavily linked to the learning rate. However, the momentum buffer not only spreads the application of gradients over multiple iterations. Instead, in the moving average of the gradients directions of different batches/iterations can cancel out or accumulate, forming new (smoothed) optimization directions, which are not directly reducible to a larger learning. 2. Please note that we use the Simple-ViT variant of Li et al. 2021 which slightly differs from the original ViT implementations. We base our code on a common PyTorch implementation (https://github.com/lucidrains/vit-pytorch) and are on par with the reported results. Please also note, that we did not perform heavy hyperprameter optimization for our method and used the default parameters suggested by the authors. We release code for all experiments and report full results for our $\rho$-search in the appendix to ascertain maximal transparency.\ We agree, that comparing against random perturbations is of key importance. You can find regarding experiments in Sec. A.7 of our initial submission, showing that MSAM outperforms random perturbations and last gradient perturbations significantly. ***Questions*** 1. By "well-chosen" we do not mean optimal, but only in a suitable magnitude for training, as it is the case in practice. Crucial is the relation to the magnitude of the perturbation which is typically at least one magnitude larger than the learning rate. We reformulated this in the main text.\ Thanks for suggesting the additional experiment. We added it to the appendix (Fig. A.9) and referred to it in the main text. Explicitly showing the value change of the loss due to the momentum perturbation is a great empirical addition to our work. For all suggested learning rates, we observe clear loss increases, validating that the negative momentum step causes a local loss ascent. 2. Thanks for pointing out the interesting paper, we absolutely agree: it is now included in our related work section. The method proposed by Li et al. is in fact closely related to our approach. However, the authors are calculating perturbations in momentum direction after updating the momentum with the current gradient, which results in the momentum buffer being dominated by the current gradient and thus not having a negative inner product with the gradient (which we exploit to find efficient perturbations). Thus, the perturbations suggested by the authors are in positive momentum direction and do not contradict our findings. While the authors seek to improve SAM by reducing the variance of perturbations, no speedup over SAM can be achieved. 3. Yes: We calculate the perturbation direction as $\epsilon^\text{MSAM} = - \rho v_{t}/||v_{t}||$ (Eq. 5) where the momentum buffer $v_{t}$ is an exponential moving average over past gradients, i.e. $v_{t+1} = \mu v_{t} + \nabla L_{B_t}(w_{t})$. 4. We show extensive comparisons against five related methods in Tab. 1 and 2. We adopted the original implementations where available and performed $\rho$-searches for all methods with equal sampling to ascertain a fair comparison (see Sec. A.5 and A.6).
thanks for the rebuttal
Thanks for the discussions, which partially address my concern, e.g., the discussion to Li et al. However, I still think the presentation and experiments should be largely enhanced. So I would like to keep my score unchanged.
Summary
The authors are striving to reduce the computational overhead incurred in the SAM algorithms, where they propose Momentum-SAM (MSAM). The proposed method utilizes the momentum direction to approximate sharpness computations. In such a manner, nearly no extra computation cost would be needed compared to the vanilla SAM. The authors empirically show the effectiveness of their method via experiments on CIFAR100 and ImageNet. The results show that the proposed method can reduce the computational overhead, but may harm the model performance.
Strengths
**Strengths** 1. The paper is clearly written and easy to follow. Also, the presentation of the paper is good. 2. The paper present some interesting results, where it not only gives empirical results on the conventional benchmarks, also presents detailed study in regards to its properties.
Weaknesses
**Weakness** The paper is an empirical paper aiming to reducing the computational overhead incurred by the vanilla SAM algorithm. The proposed method is a quite heuristic method that substitutes the accent direction in SAM which is the positive gradient direction of the current weight model by the history gradient direction. In this way, the computational overhead could indeed be reduced. I would discuss the paper from several perspectives. 1. I think the novelty of the paper is moderate. There are quite a number of papers targeting on the very topic. The core idea is somewhat quite similar to LookSAM. And the improvement of the proposed method is marginal. In fact, it is quite difficult to contribute further on this very topic, given there are a number of efficient methods proposed. And it may be more important to ensure the accuracy first. Given the results of the proposed method is good, I think the proposed method is acceptable. But the authors should present discussions in regards to the following paper, which contribute the topic from another perspective. [1] Bahri, Dara, Hossein Mobahi, and Yi Tay. "Sharpness-aware minimization improves language model generalization." arXiv preprint arXiv:2110.08529 (2021). [2] Ni, Renkun, et al. "K-SAM: Sharpness-Aware Minimization at the Speed of SGD." arXiv preprint arXiv:2210.12864 (2022). 2. The significance of this paper is limited. To my understanding, the core of NAG is to accelerate convergence. It should be noted that here accelerating convergence does not mean to reduce the computation cost at each training batch. The authors have not provided any demonstration regarding the convergence of the proposed heuristic method. Given that the support of the proposed method is too weak, in my opinion, the minimum requirement of ICLR acceptance here is to discuss to what extent the proposed method can convergence from a theoretical perspective. See [3] and other related papers for reference. [3] Andriushchenko, Maksym, and Nicolas Flammarion. "Towards understanding sharpness-aware minimization." International Conference on Machine Learning. PMLR, 2022. 3. The proposed method is not well supported. I could understand the authors' intention. However, the rationality of the proposed method is not clearly demonstrated. Regarding SAM, the ascent direction is solved explicitly and has a clear meaning. Minimization following the ascent direction can force the loss within the neighbourhood region to be low (i.e. flat). This is why SAM works. However, in my opinion, the proposed method cannot provide such a guarantee, given that the ascent direction could not characterize the maximum loss within the neighbourhood region. I notice that the authors give some simple empirical demonstration regarding the similarity between the SAM and MSAM. However, such a empirical demonstration could not give enough support and moreover, a similarity in direction could not give direct connection that the corresponding loss of the two models are close. It is not quite appropriate that one just give some other directions and claim that such directions could lead training to flat minimum.
Questions
See weakness.
Rating
5: marginally below the acceptance threshold
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
2 fair
Presentation
3 good
Contribution
2 fair
Ethics concerns
I have not found any discussions about the limitations and potential negative societal impact. But in my opinion, this may not be a problem, since the work only focuses on the learning method in machine learning. Still, it is highly encouraged to add corresponding discussions.
1. Thanks for suggesting the other related methods which we included to our paper now. While MSAM outperforms all other suggested efficient versions of SAM, it also not only relies on exploiting the exact same phenomenon as SAM and making it more efficient (as e.g. by sparsification by ESAM, K-SAM, SSAM, LookSAM, Bahri et al.), but instead shows that other directions (momentum directions) are causing similar generalization improvements. We think this yields additional perspectives on sharpness minimization in general and might motivate other researches to further investigate the reasons for SAM/MSAM's benefits. Furthermore, we show multiple experiments to empirically analyze the effects of MSAM and its relations to NAG. 2. We fully agree, the main idea of NAG ist to improve convergence. However, we might missed to emphasise that that momentum perturbations can also be used to improve generalization, which is shown in our manuscript. The runtime benefits of our method are not directly related to NAG but instead exploit the generalization capabilities of momentum perturbations to circumvent the additional network passes needed for SAM. Thus, we do not aim for faster or better convergence, but for better generalization. We proof this strictly by deriving a closer generalization bound under the occurring setting in Sec. A.1. 3. Foret et al. assume to find the maximum in a local $\epsilon$-ball. They approximate the maximization by first order on single batches and thereby already violate their theoretical justification heavily. However, we show that approximating this maximization with the momentum direction also yields great performance increases while being less computationally costly. While we show the similarity between resulting gradients if perturbed with MSAM or SAM, we also extensively validate that our method efficacious reduces the loss sharpness. This not only holds for the sharpness in momentum direction, but also for gradient directions and random directions as we show by numerically analysing the direct loss difference (Fig. 4; especially A) and analytically calculating the curvature from the expectation of the hessian eigenvalues (Fig. A.1.).
Thanks for the kind response.
I have read the authors' rebuttal. And thanks for the response. Firstly, I could understand the claim of A.1. The claims are quite hard to support that the proposed method could be better than SAM, or its variants. Then, since the authors have changed the ascent gradient computation, I think the authors should demonstrate the convergnce analysis of your algorithm based on the paper [3], which has also been done by many efficient papers. Secondly, SAM is clearly motivated. And solving the maximization or minimization using the first-order approach is a widely accepted technique. But using the gradient of last batch to solve the maximization of current batch seems not quite reasonable in theory. Overall, although the authors address my concerns to an extent, my major concerns still exist.
Summary
The authors propose Momentum-SAM (MSAM) that removes the computational overhead. The method utilizes momentum inspired by NAG instead of gradient ascent step, and therefore has almost the same speed as conventional optimizers, such as SGD and Adam. The experimental results show that the method yields comparable performance within the same computational budget.
Strengths
- The paper is easy to follow and provides clear motivation. - The proposed method does not introduce additional hyperparameter except the momentum coefficient. - The paper provides experimental results from various perspectives.
Weaknesses
- From the results in Table 1, 2, 3, it can be seen that MSAM achieves relatively low performance that SAM except several cases. When the number of epochs of MSAM is doubled to make the same budget just like in A.9, can MSAM achieve higher performance than SAM? More generally, if MSAM yields higher final performance than that of SAM, the method could be more convincing. It would be nice to add more results to Figure A.8, for example, other methods like ESAM or other models/datasets. - The claims in chapter 4 are not convincing. Is there any theoretical reasons why positive momentum direction doesn't make the performance better? Also, it is unclear why $\rho_0$ is close to $\rho^{\textrm{opt}}$. - Is the momentum applied to SAM as well in the paper? I understand that SAM has no momentum term. - "efficient implementation" of the method described in page 4 is quite straightforward, it seems to be unnecessary to describe it separately.
Questions
See Weaknesses.
Rating
6: marginally above the acceptance threshold
Confidence
5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.
Soundness
3 good
Presentation
3 good
Contribution
2 fair
1. Given limited computational budget we can confirm that MSAM outperforms SAM as shown in Fig. A.8. We further support this by training a ViT-S/32 with MSAM with doubled number of epochs (180) on ImageNet where we reach $70.1\%$ test accuracy and thus clearly outperform SAMs $69.1\%$ (cf. Tab. 2). We added these results to Sec. A.9. Furthermore, we added ESAM to Fig. A.8 and demonstrate that MSAM clearly outperforms ESAM. For comparison against other models and datasets please refer to Tab. 1. 2. Given the negative inner product of momentum and gradient (cf. Fig. 2C), positive momentum directions do not yield a lower generalization bound as theoretically shown in Sec. A.1. Furthermore, we empirically show in Fig. 2B that positive momentum directions do significantly harm training accuracies and thus the loss minimization in general. We hypothesize this to be caused by the optimization slowing down and getting stuck, since positive momentum directions perturb towards parameters of past iterations.\ Furthermore, the similarity between $\rho_0$ and $\rho^\text{opt}$ yields evidence that the momentum perturbations of MSAM efficaciously approximate the perturbations of SAM. We hope this explanations clarify this observation and we changed the main text to emphasis this hypothesis and to strengthen the claims in Chapter 4. 3. We might missed to clarify this in the initial submission: SAM only changes the position of gradient calculation, but does not change the underlying optimizer. Thus, if combined with SGD, momentum is typically used during the descent step (default momentum factor $\mu = 0.9$). We also did this for all our experiments. However, SAM does not utilize momentum for perturbations as we did for MSAM. 4. We absolutely agree. To make our paper accessible to readers who are not directly familiar with SAM and NAG we, however, would like to keep the "efficient implementation" as an explanation of our algorithm.
Thanks for the author's response.
I have read the author’s response, and some of my minor questions are addressed. But I still have a concern, especially in terms of final performance as mentioned in my first comment. Even within the sufficient computational budget, for example 200 epochs (ResNet) and 600 epochs (ViT-S) in Table 2, can MSAM outperform SAM? In Figure A.8, I found that SAM achieves almost the same or slightly higher performance at the end. Of course, MSAM has an strength that it can achieve better performance in early stages, but someone may care about the final performance as well.
Summary
The paper proposed an efficient sharpness-aware-minimization optimizer that does not need to calculate the inner optimization during the training which accelerates the training progress while keeping the model performance close to the original SAM. In order to obtain a good estimator of the gradient ascent direction in SAM, this paper uses the negative momentum direction as the estimation and shows that the model reaches a low sharpness and high performance. Since for most optimizers such as SGD and ADAM, momentum is used during training, there will be no computation overhead for the proposed method. MSAM achieves good empirical results on both cifar10 and ImageNet.
Strengths
The paper is well-written and easy to follow. The problem is important and the phenomenon that the direction of momentum is negatively correlated to the gradient direction of the current batch is very interesting. The empirical results are good at both cifar10 and ImageNet.
Weaknesses
To me, the intuition of why the negative momentum direction is a good estimate of the gradient ascent direction is still not clear. To be honest, the phenomenon why this happens is more important and interesting to me since there are quite a lot of methods that can efficiently simulate SAM. The paper does have some analysis showing that MSAM actually minimizes the sharpness, but I would love to have more intuitions or hypotheses. For example, is this dataset related? Is this task related, like only for classification tasks or only for vision problems?
Questions
Can we have more experiments on finetuning results as the original SAM? Can we have some language model related experiments? Any hypotheses or intuitions of why the momentum direction is so different from the current batch? For Section 4.3, can you make it clear why it is related to m-sharpness? Can we have some ImageNet NAG results?
Rating
6: marginally above the acceptance threshold
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
Thanks four your feedback. Given that the dot product of the momentum buffer and the current gradient is negative (Fig. 3C), the validity of the negative momentum direction as an efficient proxy to calculate SAM-perturbations can be proven strictly as we show in Sec. A.1. We absolutely agree that this phenomenon is counterintuitive at a first glance, however, we think it can be understood heuristically based on the the high variance among input samples and the little shared components of gradients of independent batches. Moving in negative gradient directions of single batches for multiple iterations first causes a loss decrease, but quickly overshoots local minima and results in an increase again (cf. loss minima on local batches shown in Fig. 4A). Furthermore, gradients of consecutive batches are typically nearly orthogonal with a slightly positive inner product (especially for early iterations). Thus, after moving in the momentum direction and overshooting local minima, the shared components of the gradients accumulated in the momentum direction turn negative and cause the observed phenomenon. Due to the high variance among input samples the next gradient added to the momentum buffer causes the optimization direction to turn towards a decreasing direction again and thus nevertheless allows loss minimization. Since it is very hard to formulate a closed theoretical framework for SOTA-model architectures to cover these phenomenons from a theoretical perspective, we decided to not only focus on explaining the negative correlations between momentum buffer and the gradients, but share this observation and additionally exploit it to propose an efficient algorithm. We agree that this specificity of MSAM is an interesting phenomenon and we tried to explain our intuition better in the main text now. With respect to the requested language models, we performed machine translation (English to Romanian) on the WMT 2016 dataset by finetuning a T5-tiny model (efficient version by Tay et al. Scale Efficiently: Insights from Pretraining and Finetuning Transformers, ICLR 2022) from a publicly available checkpoint pretrained on the C4 dataset. We scanned $\rho \in \{0.01, 0.03, 0.1, 0.3, 1\}$ and found MSAM and SAM both to perform best at $\rho = 0.1$. The resulting BLEU scores are 23.35 for AdamW, 23.57 for SAM and 23.64 for MSAM. MSAM slightly outperforms SAM (as well as AdamW) while requiring two times less computations. We added these results to the Appendix. We could validate the suitability of mometum perturbations for sharpness minimization for all investigated datasets. For NAG for ResNet-50 on ImageNet we reached $76.5\%$ test accuracy, which is on par with MSAM. The comparably low performance gain of SAM and MSAM under these settings might be caused by the fact that no pronounced overfitting occurs, i.e. the training accuracy only reaches $78.4\%$, and thus the improved generalization capabilities of MSAM and SAM do not increase the performance substantially. However, larger models like ViT-S/32, where overfitting occurs (train: $85.4\%$; test: $67.0\%$ for adamW and 90 epochs) show significant performance increases by MSAM. Regarding $m$-sharpness, MSAM can be seen as an opposite phenomenon. While $m$-sharpness calculates perturbations on small sub-batches, MSAM utilizes a moving average (momentum buffer) over many batches. However, the empirical results are in contradiction with the theoretical motivation of SAM (Foret at al.) where the generalization bounds where theoretically derived on the full training set. Following this idea, MSAM should yield the best generalization, followed by SAM and $m$-SAM, but we actually observe the opposite order where MSAMs advantage over SAM only stems from the reduced runtime. By this we show evidence that the derivation of Foret et al. can not fully explain the mechanisms of SAM/MSAM/$m$-MSAM.
Meta-review
This paper proposes a very interesting approach to improve computational efficiency of SAM. Each iteration of SAM involves two gradient computations: an ascent perturbation and then a decent update from the perturbed point; this doubles the computational cost of SAM compared to SGD. The idea in this paper is to eliminate the extra cost of the ascent (gradient) computation by replacing it with negative momentum which is cheap to compute. This is a very interesting idea, as also recognized by all reviewers. Building on the negative momentum instead of the ascent step, the authors propose a new variant of SAM called MSAM (Momentum-SAM). They evaluate MSAM empirically and report speedups compared to SAM. I personally like many aspects of the paper, but Figure 1 is one of my favorite parts: it communicates the proposed ideas and illustrates its difference with SGD and SAM using a simple and clear diagram. The major issue that reviewers had with the submission was lack of clarity around why negative momentum can be used as a proxy for the ascent step. As this is the core contribution of the paper, the authors could have made more effort to elaborate this connection. The paper currently talks about the interaction of the momentum with learning rate and how this setup can lead to an overshoot of the minimum (hence increasing the loss in lieu of an ascent perturbation), but the reviewers believe more rigorous justification is needed. Authors attempted to address this during rebuttal, but at the end reviewers still felt more evidence (either theoretical or empirical) is needed to support the main claim of the paper. Unfortunately, with the current state of the review, the paper cannot be accepted. However, in concordance with the reviewers, I DO believe that this is a very very interesting work and I hope authors are not discouraged by the outcome of this round of review and they consider improving the draft and resubmitting it.
Why not a higher score
The core contribution of the paper, using negative momentum instead of SAM's ascent step, is not fully justified. Additional evidence (either theoretical and empirical) is needed to be convincing.
Why not a lower score
N/A