Lookahead Sharpness-Aware Minimization

Thank you for your constructive reviews. We summarize the reviewer's questions and present our responses below.

Q1: I am curious about the motivation behind the proposed method. ... i) Can you clarify why EG (and correspondingly OG) are straightforward solutions in terms of SAM for escaping saddle points? ii) Or the purpose of EG is to increase the loss of $L(w+\varepsilon)$ by perturbing $\varepsilon$ ? The loss maximization does not seem to align with performance gain [1].

i) The lookahead mechanism is designed to help to avoid saddle point by gathering more information about the landscape, as mentioned in Sec. 1. Additionally, another more intuitive explanation is that, as also pointed out by Reviewer hY3b, the EG-SAM update scheme can be written as $w_t =w_{t-1}-\eta_t \nabla L( w_{t-1}+\epsilon^\prime_t)$ where $\epsilon^\prime_t = (\rho /||\nabla L(w_{t-1})||-\eta_t) \nabla L(w_{t-1})$. Note that SAM's update scheme is $w_t =w_{t-1}-\eta_t \nabla L( w_{t-1}+\epsilon_t)$, where $\epsilon_t = (\rho/||\nabla L(w_{t-1})||)\nabla L(w_{t-1})$. EG-SAM can be interpreted as suppressing perturbation in the early stages and allowing more perturbation in the final stages (as $\epsilon_t$ decreases from large to small values). Since a larger perturbation can make the model more prone to being trapped in saddle points, as indicated by Kim et al. (2023), our method is more effective in avoiding saddle points than SAM.

ii) Our EG does not aim to increase the loss of $L(w+\varepsilon)$. As discussed in 1), our work is designed to avoid saddle points, which is different from aiming to increase the loss value of $L(w+\varepsilon)$.

Q2: To push further the question 1, I am not convinced that SAM+EG and EG-SAM (also OG) are not highly correlated because $\epsilon$ in EG-SAM is not calculated to enlarge the loss with respect to $\epsilon$.

We agree that SAM+EG can enlarge the losses, as $\epsilon$ is no longer a one-step gradient ascent. However, as discussed in A1, different from SAM+EG, EG-SAM is not designed for enlarging the loss. Thus, SAM+EG and EG-SAM are different.

Specifically, recall SAM can be formulated as a minimax problem (Eq. (1)) $\min_{\boldsymbol{w}} \max_{\boldsymbol{\epsilon}:\|\boldsymbol{\epsilon}\| \leq \rho} L(\boldsymbol{w}+\boldsymbol{\epsilon})$. As EG is used for addressing minimax problem, SAM+EG can be directly implemented as the solver. However, for EG-SAM, it firstly uses a max oracle $\max_{\boldsymbol{\epsilon}:\|\boldsymbol{\epsilon}\| \leq \rho} L(\boldsymbol{w}+\boldsymbol{\epsilon}) \approx L(\boldsymbol{w}+\frac{\rho \nabla_{\boldsymbol{w}} L(\boldsymbol{w})}{\|\nabla_{\boldsymbol{w}} L(\boldsymbol{w})\|})$ following the original SAM (Foret et al., 2020), then (Eq. 1) is converted to a minimization problem $\min_{\boldsymbol{w}} L(\boldsymbol{w}+\frac{\rho \nabla_{\boldsymbol{w}} L(\boldsymbol{w})}{\|\nabla_{\boldsymbol{w}} L(\boldsymbol{w})\|})$, and finally EG is designed to solve this problem.

Q3: Rather, even though I did not calculate the ODE accurately, the behavior of EG-SAM seems somewhat related to implicit methods in ODEs (thus stable) rather than using explicit methods by using $\hat{w}_t$ obtained from the next step in SGD terms.

Regarding the ODE, we thank the reviewer for sharing the ODE perspective to help explain our method. Indeed, we agree with the reviewer that using future gradients to update the current model is somewhat similar to the implicit methods in ODE. We will further explore this connection in our future works.

Q4: I agree that EG-SAM can optimize better in quadratic loss, as shown in Section 4 with ODE. However, the optimization itself is discrete, and some behaviors (such as SGD or catapult effects) cannot be explained in ODE settings. Thus, could the authors demonstrate the effects of EG-SAM in real situations with SGD to escape saddle points, such as in Fig. 7 in [2]?

The reviewer might have some misunderstandings. We have already conducted a toy experiment at the beginning of Sec. 3, which uses quadratic objective to show the discrete time behaviors of SAM and our methods around saddle point following (Compagnoni et al., 2023). Results reveal that discrete time EG-SAM can escape saddle point faster than SAM. Fig. 7 in [2], as referred to by the reviewer, similarly illustrates the discrete time behavior of SAM around saddle point.

Besides, ODE is also widely used to analyse SAM (Kim et al., 2023; Compagnoni et al., 2023), and we follow the same analysis scheme.

Thank you for your all valuable suggestions. Below we respond to your key concerns point by point:

Q1: The idea of multi-step SAM is not new and has been explored before. The EG and OG versions proposed here are new, but I think the contribution of the work is marginal.

i) The reviewer might have some misunderstandings. Our EG-SAM significantly diverges from that of multi-step SAM methods like those in [1]:

[1] aims to use multi-step SAM to yield a more precise solution of max oracle $\boldsymbol{\epsilon}^* (\boldsymbol{w}):= \arg \max_{\boldsymbol{\epsilon}:\|\boldsymbol{\epsilon}\| \leq \rho} L(\boldsymbol{w}+\boldsymbol{\epsilon})$. This method focuses on refining the maximization step within the SAM framework.

In contrast, our EG-SAM is primarily focused on avoiding saddle points and finding a better trajectory for convergence. This is achieved by incorporating a look-ahead mechanism. Unlike multi-step SAM (e.g., [1]), it is not designed specifically to refine the maximization step but rather to enhance overall convergence by avoiding saddle points.

ii) Regarding the contributions,

1. we are the first to improve SAM by using look-ahead mechanism to avoid saddle points.

2. We design different variants of look-ahead mechanism, e.g. EG-SAM, OG-SAM and AO-SAM, to help further reduce computational cost.

3. We provide theoretical guarantees.

4. Our method outperforms other SOTA baselines consistently with different structures and in different datasets.

Q2: The theory is nice to have, but is not really insightful and the proofs are fairly standard.

i) Regarding our proof is not really insightful, we provide two distinct properties of EG-SAM: 1) EG-SAM can escape saddle points faster than original SAM (Proposition 4.1), and 2) EG-SAM has more chance to converge at flatter minima (Proposition 4.2). These findings are novel and were not present in any previous work related to SAM. The ability to escape saddle points quickly and the increased probability of converging to flatter minima provide valuable insights into the advantages of our proposed methods.

ii) Regarding our proof is fairly standard, note that original SAM proof cannot be directly used to prove our proposed methods due to the following reasons:

1. Recall in OG-SAM, we have $\hat{w_t}=w_{t-1}-\eta_t g_{t-1}$, which includes past gradient $g_{t-1}$. This makes the standard SAM analysis (Andriushchenko & Flammarion, 2022) difficult to apply directly, as their analytical approach does not include a technique to address past gradients.

2. In the proof of AO-SAM, we use an adaptive policy with OG (Jiang et al., 2023). Although (Jiang et al., 2023) provides a proof, it is only established for unnormalized SAM. Since ours is normalized SAM with OG, it inherits the difficulties of both the adaptive policy in normalized SAM and past gradient analysis.

Q3: I'm not fully convinced with the numerical experiments. They consider CIFAR + ResNet, and only one example of ImageNet + ResNet. No experiments on transformers or language models are done, which makes the experiments rather limited. Also, important version of SAM such as gSAM and mSAM are not included in the experiments. Also, I didn't find runtimes of methods, which I might've missed.

1. We added an experiment using vision transformer ViT-S16 on CIFAR-10 and CIFAR-100 in our response to Reviewer rtXV Q3. As can been seen, our proposed method also outperforms the baselines.

2. We also added the performance comparison between GSAM [3], MSAM [4], and our AO-SAM in CIFAR-10 and CIFAR-100 with WideResNet-28-10 backbone, as suggested by the reviewer. From the below table, we find that our method also outperforms GSAM and MSAM.

3. We have also added the running time comparisons in our response to Reviewer 1P81's Q5. As can be seen, our proposed OG-SAM has relatively same training time as SAM, while AO-SAM is faster than SAM.

Q4: The presentation can improve. There are so many versions and abbreviations (EG, OG, AO, AMO) that it's difficult to see the points the authors are trying to make.

Thank you for pointing out. we have included the full names alongside their abbreviations wherever necessary to enhance readability.

Reference:

[1] Kim, Hoki, et al. "Exploring the effect of multi-step ascent in sharpness-aware minimization." arXiv preprint arXiv:2302.10181 (2023).

[2] Dosovitskiy, Alexey, et al. "An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale." ICLR. 2020.

[3] Zhuang, Juntang, et al. "Surrogate Gap Minimization Improves Sharpness-Aware Training." ICLR. 2021.

[4] Behdin, Kayhan, et al. "mSAM: Micro-Batch-Averaged Sharpness-Aware Minimization." arXiv preprint arXiv:2302.09693 (2023).

We thank the reviewer’s helpful feedback on our submission. We summarize the reviewer’s questions and the following is our responses:

Q1: The EG-SAM update can be rewritten as: ... When the stepsize is taken decreasing (as in the experiments) it is simply a scheme that perturbs less adversarially in the early training phase and more later.

Our scheme is designed based on some non-trivial insights:

i) This scheme is based on incorporating the lookahead mechanism into SAM, which has not been proposed previously.

ii) The scheme is also explainable. By suppressing perturbation during the early stages, our method is more effective at avoiding saddle points than the original SAM. This is because larger perturbations make the model more susceptible to being trapped in saddle points, as noted by Kim et al. (2023). Additionally, by allowing more perturbation in the final stages, our method encourages the model to find flat minima.

Q2: It would be instructive to plot the gradient norm at $w_t$ throughout training.

We added gradient norm at $\nabla L(w_t)$ in Figure 5, Sec. B, as suggested by the reviewer. As can be seen by the figure, the gradient of EG-SAM is larger than that of SAM in the early stages, but it converges to the same level as SAM in the final stage.

Q3: How does EG-SAM compare against SAM with an (increasing) stepsize schedule for $\rho$?

We added an extra experiment regarding using increasing stepsize schedule for $\rho$, i.e., $\rho_t = \rho_0+ \frac{t(\rho_{T}- \rho_0)}{T}$, where $\rho_{T}$ is the final value $\rho_{T}$ (1.0 in our experiment), while $\rho_0$ is the initial value of $\rho$ (0.1 in our experiment). As can been from the table, our method with increasing $\rho$ still outperforms the SAM, and is comparable to original EG-SAM, while SAM with increasing $\rho$ perform worse than original SAM method.

Q4: OG-SAM appears unnecessary since it does not save on gradient computations: EG-SAM simultaneous computes the extrapolation and perturbation computation so it already only requires two gradient computations. What is the motivation for OG-SAM in not saving gradient computations?

In EG-SAM, we need to compute three variables at each step: $\hat{w}$, $\hat{\epsilon}$, and $w$. In contrast, standard SAM and OG-SAM require only $\hat{\epsilon}$ and $w$ per step (OG-SAM reuses the gradient of $\hat{w}$), and thus we are saving gradients. Besides this computational efficiency, another advantage of OG-SAM is that it contains the past gradient $\boldsymbol{g}_{t-1}$ at current step $t$. As discussed in Sec. 3.4, by leveraging extra information from the past gradient, it improves the performance (Rakhlin & Sridharan, 2013).

Q5: The convergence result in Section 4.2 requires strong assumptions on both the i) problem class and ii) the a decreasing perturbation radius $\rho$. ... More importantly, by decreasing $\rho$ you essentially avoid the difficulty of the perturbation. One indication that your assumptions are too strong is that your argument would hold for any kind of (feasible) perturbation.

i) Regarding the problem class assumptions, note that we do not introduce the function Lipschitz assumption. Our three assumptions are bounded variance, smoothness (gradient Lipschitz), and uniformly bounded gradient (Assumptions 4.3-4.5). These assumptions are widely used in the analysis of SAM and its variants (Mi et al., 2022; Dai et al., 2023; Zhang et al., 2023; Yue et al., 2023). Beyond SAM, these assumptions are utilized in analyzing different optimizers [3].

ii) Regarding decreasing $\rho$, a recent paper [1] theoretically reveals that SAM with constant $\rho$ CANNOT converge to stationary points. Therefore, assuming a decreasing $\rho$ is necessary in theoretical analysis. Moreover, related works with decreasing $\rho$ assumptions under stochastic gradient settings are also found in related works (Mi et al., 2022; Dai et al., 2023; Zhang et al., 2023; Yue et al., 2023).

We much appreciate your positive and insightful comments on our paper. Here is our response to your questions point by point:

Q1: Those proposed methods are only some variations of SAM with small changes. Hence, the authors did not encounter further difficulties to obtain those convergence rate results.

i) Our proposed methods include the following contributions:

1. we improve SAM by using look-ahead mechanism to avoid saddle points.
2. We design different variants of the look-ahead mechanism, e.g., EG-SAM, OG-SAM, and AO-SAM, to further reduce computational cost.
3. We provide theoretical analysis. Experiments in Sec. 5 reveal that these modifications all help improve the performance of SAM.

ii) Adding a look-ahead mechanism introduces new difficulties to the proof:

1. Recall in OG-SAM, we have $\hat{w_t}=w_{t-1}-\eta_t g_{t-1}$, which includes past gradient $g_{t-1}$. This makes the standard SAM analysis (Andriushchenko & Flammarion, 2022) difficult to apply directly, as their analytical approach does not include a technique to address past gradients.

2. In the proof of AO-SAM, we use an adaptive policy (Jiang et al., 2023) with OG. Although (Jiang et al., 2023) provides a proof, it is only established for unnormalized SAM. Since ours is normalized SAM with OG, it inherits the difficulties of both the adaptive policy and past gradient analysis.

Q2: The experiment results show incremental enhancement to SAM.

i) Regarding the experimental results, while our method does not exhibit a dramatic improvement over the SOTA baselines, it consistently outperforms them across various environments and with different backbone models. Additionally, the baselines represent the current best-performing methods in the field. Therefore, it is common for the observed improvements to be relatively modest (Jiang et al., 2023; Du et al., 2022a; Mi et al., 2022).

ii) To further demonstrate that our improvement is NOT merely incremental, we have added a test of significance (T-test) on the noise-labeled CIFAR-100 dataset (Table 6). Our method shows statistical significance with a p-value less than 0.05. This supports our claim that our improvement is statistically significant, particularly for the noise-labeled dataset, which requires a higher degree of generalization ability.

Q3: In comparison with SAM paper, the authors had done less experiments for comparisons, i.e. ResNet-101, ResNet-152 etc.

Due to resource limitations, we are sorry that our experiments did not include ResNet-101 and ResNet-152 on ImageNet.

As an alternative, we incorporated an extra experiment with a large neural network, Vision Transformer ViT-S16 [1], on the CIFAR-10 and CIFAR-100 dataset. The results from this experiment demonstrate that our method outperform other baselines continually.

Q4: There is also a little intuitive explanation/proof for advantages of the proposed methods. It is decent work, not something fancy to read.

i) As discussed in Section 1, the introduction of the extra gradient in our model allows for gathering more information about the landscape by looking further ahead. This facilitates finding a better trajectory for convergence, which is the motivation for incorporating the extra-gradient into SAM.

ii) Additionally, as also pointed out by Reviewer hY3b, the EG-SAM update scheme can be written as $w_t =w_{t-1}-\eta_t \nabla L\left( w_{t-1}+\epsilon^\prime_t\right)$ where $\epsilon^\prime_t = (\rho /||\nabla L(w_{t-1})||-\eta_t)\nabla L(w_{t-1})$. In comparison, SAM's update scheme is $w_t =w_{t-1}-\eta_t \nabla L\left( w_{t-1}+\epsilon_t\right)$, where $\epsilon_t = (\rho /||\nabla L(w_{t-1})||)\nabla L(w_{t-1})$. EG-SAM can thus be interpreted as it suppresses perturbation at the early stage (as $\epsilon_t$ decreases from large to small values), and allows more perturbation in the final stage. Since larger perturbations make the model more prone to being trapped in saddle points (Kim et al., 2023), our method more effectively avoids these points compared to SAM by reducing perturbation. Furthermore, the increased perturbation in the final stage enables the model to find flat minima.

References:

[1] Dosovitskiy, Alexey, et al. "An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale." ICLR. 2020.