Paper: Momentum-SAM: Sharpness Aware Minimization without Computational Overhead (ICLR-2024)

Reviewer comment (Theory):
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.

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Evaluation):
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?

Author response:
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.

Reviewer comment (Novelty):
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)

Author response:
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.

Reviewer comment (Theory):
Could the author provide a formulation similarly to NAG on which lookahead direction is calculated.

Author response:
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})$.

Reviewer comment (Experiments):
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.

Author response:
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).

Reviewer comment (Theory):
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.

Author response:
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.

Reviewer comment (Theory):
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.

Author response:
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.).

Reviewer comment (Experiments):
From the results in Table 1, 2, 3, it can be seen that MSAM achieves relatively low performance than SAM except in 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.

Author response:
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.

Reviewer comment (Theory):
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}}$.

Author response:
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.

Reviewer comment (Reproducibility):
Is the momentum applied to SAM as well in the paper? I understand that SAM has no momentum term.

Author response:
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.

Reviewer comment (Theory):
The intuition of why the negative momentum direction is a good estimate of the gradient ascent direction is still not clear, and understanding the phenomenon is more important and interesting since there are already quite a lot of methods that can efficiently simulate SAM.

Author response:
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.

Reviewer comment (Experiments):
Can we have some language model related experiments?

Author response:
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.

Reviewer comment (Evaluation):
Any hypotheses or intuitions of why the momentum direction is so different from the current batch?

Author response:
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.