SAMPa: Sharpness-aware Minimization Parallelized

Sharpness-aware minimization (SAM) has been shown to improve the generalization of neural networks. However, each SAM update requires \emph{sequentially} computing two gradients, effectively doubling the per-iteration cost compared to base optimizers like SGD. We propose a simple modification of SAM, termed SAMPa, which allows us to fully parallelize the two gradient computations. SAMPa achieves a twofold speedup of SAM under the assumption that communication costs between devices are negligible. Empirical results show that SAMPa ranks among the most efficient variants of SAM in terms of computational time. Additionally, our method consistently outperforms SAM across both vision and language tasks. Notably, SAMPa theoretically maintains convergence guarantees even for \emph{fixed} perturbation sizes, which is established through a novel Lyapunov function. We in fact arrive at SAMPa by treating this convergence guarantee as a hard requirement -- an approach we believe is promising for developing SAM-based methods in general. Our code is available at \url{https://github.com/LIONS-EPFL/SAMPa}.

Paper

Similar papers

Peer review

Reviewer 9hZ55/10 · confidence 4/52024-06-19

Summary

This paper studies the efficiency problem of sharpness-aware minimization (SAM) algorithms. SAM requires two gradient calculations: one for computing the perturbation and another for computing the update direction. Hence, SAM doubles the computation cost compared with ERM. Furthermore, these two gradients cannot be parallelized. The authors introduce an auxiliary sequence and use the gradient on the auxiliary variable to approximate the gradient in computing the perturbation. Experiments are conducted to evaluate the effectiveness of the proposed SAMPa algorithm.

Strengths

- improving the efficiency of SAM is an important problem - moreover, making the two gradient computations parallelized is crucial to accelerate SAM algorithms - introducing an auxiliary sequence to approximate the gradient in computing perturbation is interesting - the writing is clear and easy to follow.

Weaknesses

- Though introducing an auxiliary sequence is an interesting idea to parallelize the two gradient calculations, I have a concern about the auxiliary sequence: - estimate the approximate error $\nabla f(x_t, B_{t+1}) - \nabla f(y_{t+1}, B_{t+1}) $. Establishing a bound for this error is crucial to the proposed algorithm. Otherwise, the method has no theoretical support. - The theoretical analysis in Section 4 requires the loss to be convex. However, losses are usually non-convex in practice (for example, experiments in the paper). Thus, the analysis in section 4 is meaningless in practice. - Main results (Tables 1 and 2) show that SAMPa outperforms SAM in all settings. This observation is odd to me. Note that SAMPa uses an approximate method to compute the gradient but achieves better performance than SAM which uses the exact gradient. Some explanations are needed here.

Questions

see above weaknesses.

Rating

5

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

no

Reviewer 9hZ52024-08-13

reply to rebuttal

Thanks for the authors' reply. Indeed, I mean a bound on the approximate error $\nabla f(x_t, B_{t+1}) - \nabla f(y_{t+1}, B_{t+1}) $ in the **non-convex** case! For Q2, As SAM is designed for non-convex problems, I do not think it is meaningful to focus on analyzing the convergence in the convex case.

Authorsrebuttal2024-08-14

Further response to Reviewer 9hZ5

We sincerely thank the reviewer for the positive response. Regarding the focus on convex analysis, we do not understand why convex analysis should be less suitable for SAM than e.g. the mentioned AdaGrad or Adam. SAM has mostly seen success in overparameterized settings (where the training data is eventually perfectly fitted), where it has been exploited in several works that a convexification can occur. Additionally, the solution set for the convex problem is not necessarily a singleton, so biasing the obtained solution still seems reasonable. Moreover, we do not see a bound on the approximate error as crucial, since it is not our primary goal. At least in the convex case, convergence can be established _without_ such a bound as we demonstrated in the rebuttal. We are unsure why it would become essential in the non-convex setting. We would like to explore the convergence in the non-convex case in our future work. Thank you again for your engagement, and we remain available for any further questions.

Reviewer LruU6/10 · confidence 4/52024-06-28

Summary

This paper studies a parallelized variant of sharpness aware minimization (SAM). This is achieved by introducing a sequence of auxiliary variables to break down the sequential dependence for the two gradients in every SAM iteration. The resultant approach SAMPa has convergence guarantee for convex problems when $\rho$ is chosen as a constant independent of $T$. A heuristic (SAMPa-$\lambda$) combining SAMPa with optimistic gradient descent method is also proposed. Numerical results show that SAMPa and SAMPa-$\lambda$ generalization merits and runtime benefits.

Strengths

S1. **Novel methodology.** Using a new sequence to break the sequential dependence of SAM gradients is novel. S2. **Flexible framework and improved test performance.** As shown in Section 5.5, the proposed approach can be integrated with different variants of SAM. Moreover, the proposed approach seems to improve the test performance on various downstream tasks.

Weaknesses

W1. My major concern is that the runtime comparison is not fair due to the second GPU. In particular, since SAMPa uses 2 GPU, can the authors also report data parallel for SAM with 2GPUs as well? (or m-SAM for communication efficiency). For example, if the batch size for SAM is 128, each GPU can calculate gradients for 64 samples in parallel. W2. While theoretical results for deterministic convex functions are provided, the metric is gradient norm. It is known that gradient norm is easier to be optimized compared with function values [1]. Can the authors comment more on the optimality of this convergence rate? Does SAMPa-$\lambda$ also enjoy theoretical guarantees? W3. More on parallelization should be discussed. For example, how does the memory consumption of SAMPa, and how does it compare to other efficient SAM variants? How is the communication of SAMPa? Can communication be overlapped with computation? These questions should also be discussed in detail as the parallelization in SAMPa is the key novelty. [1] Allen-Zhu, Zeyuan. "How to make the gradients small stochastically: Even faster convex and nonconvex sgd." Advances in Neural Information Processing Systems 31 (2018).

Questions

Q1. In algorithm 1, is $g_0=\nabla f(y_0)$ or $g_0=\nabla f(y_0, {\cal B}_0)$? Q2. Can the authors comment more on the intuition on the reason that SAMPa outperforms SAM? There are also typos: - line 180, we can trivially allows --> we can trivially allow - line 126, identical to SAM except .. --> consider to change the word 'identical'

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

The authors have discussed the limitations.

Authorsrebuttal2024-08-13

Further response to Reviewer LruU

We sincerely thank the reviewer for their positive feedback and for raising the score. Regarding the two comments: - **Convergence of gradient norm:** We would like to emphasize that the goal is not to establish optimal rates for e.g. function value, but rather to match the rates of SAM, which we have demonstrated (please see the rebuttal for more elaboration). If convergence of function value is deemed important, note that it converges asymptotically due to the established descent like inequality. - **Convexity and SAM:** We do not understand why convex analysis should be less suitable for SAM than e.g. the mentioned AdaGrad or Adam. SAM has mostly seen success in overparameterized settings (where the training data is eventually perfectly fitted), where it has been exploited in several works that a convexification can occur. Additionally, the solution set for the convex problem is not necessarily a singleton, so biasing the obtained solution still seems reasonable. We thank again the reviewer for the positive engagement and remain available if there are any further questions.

Reviewer c6Fm5/10 · confidence 4/52024-07-08

Summary

This paper proposes a parallelized algorithm of Sharpness Aware Minimization(SAM) named SAMPa, which aims at making optimization more efficient by parallelizing the computation of one update in SAM. Plain SAM requires a 2x computational cost since 2 gradient computations (one for the perturbation and one for updating parameters) are needed at each training step. This paper utilizes an auxiliary sequence and another GPU device to calculate the perterbutation and gradient for updating parameters at the same time. Moreover, the convex analysis guarantees the convergence with a fixed perturbation size. SAMPa is shown to perform nearly as well as SAM on benchmark tasks such as, CIFAR10, CIFAR100, ImageNet-1k, and GLUE.

Strengths

- The paper is well writen and mostly free of typos. - The motivation of this paper is meaningful for practical application. - Theoretical analysis supports the proposed method. - The experiments are abundant and relatively comprehensive.

Weaknesses

- This paper mentions the efficiency and effectiveness of SAM, the following paper ["Make Sharpness-Aware Minimization Stronger: A Sparsified Perturbation Approach"](https://proceedings.neurips.cc/paper_files/paper/2022/hash/c859b99b5d717c9035e79d43dfd69435-Abstract-Conference.html) should be considered to discuss in this paper. A test accuracy of comparison between SAMPa and this sparsified approach and a discussion in introduction would be helpful. - A difference between the approximated perturbation based on auxiliary sequence $y_t$ and the real perturbation would be helpful. - A Visualization of landscape and A spectrum of the Hessian(Figure 3 in ["SHARPNESS-AWARE MINIMIZATION FOR EFFICIENTLY IMPROVING GENERALIZATION"](https://arxiv.org/pdf/2010.01412) and Figure 5 in ["Make Sharpness-Aware Minimization Stronger: A Sparsified Perturbation Approach"](https://proceedings.neurips.cc/paper_files/paper/2022/hash/c859b99b5d717c9035e79d43dfd69435-Abstract-Conference.html)) would be helpful for presenting the generalization superiority of SAMPa. - The convex assumption of analysis is too strong, so I'm concerned that its results would have little significance for practical operation. - The main reason why I hesitate to give a high score is that the paper lacks some intuitive or numerical explanation or analysis of "why SAMPa works" or "why we choose $y_t=...$ "

Questions

See weakness.

Rating

5

Confidence

4

Soundness

3

Presentation

2

Contribution

2

Limitations

The SAMPa requires 2 GPUs to realizing parallel computation, which will suffer communication cost and require more memory storage for auxiliary sequence.

Authorsrebuttal2024-08-14

Further response to Reviewer c6Fm

We sincerely thank the reviewer for the response. We would like to address the remaining concerns: > 1. Motivation We kindly ask in what way our response for Q5 in the rebuttal did not clarify the motivation. In rebuttal, we explained that the particular choice of $y_{t+1} = x_t - \eta_t \nabla f(y_t)$ is a direct consequence of the analysis, which ensures that the discrepancy between $\nabla f(x_{t})$ and $\nabla f(y_{t})$ can be controlled. Empirically, the small difference between $\nabla f(x_t)$ and $\nabla f(y_t)$, along with similar Hessian spectra and test accuracy between SAM and SAMPa-0, confirms that $\nabla f(y_t)$ is an excellent approximation of $\nabla f(x_t)$. Additionally, Section 2.2 and Section 3 provide detailed explanations of our intuition and the development of SAMPa. Key points include: - *Perturbation direction matters:* RandSAM, which uses random perturbation, performs worse than SAM, highlighting the importance of perturbation direction. - *Using past gradients is ineffective:* OptSAM, using past gradients for perturbation, underperforms than SAM and it even fails to converge in a toy example due to its reliance on a poor gradient estimate from an ascent step. - *Introduce the auxiliary sequence:* We finally propose an auxiliary sequence $y_{t+1} = x_t - \eta_t \nabla f(y_t)$ as a direct result of our theoretical analysis, which aligns with the descent direction from $x_t$. > 2. Theory limitation We do not understand why convex analysis should be less suitable for SAM than e.g. the mentioned AdaGrad or Adam. SAM has mostly seen success in overparameterized settings (where the training data is eventually perfectly fitted), where it has been exploited in several works where convexification can occur. Additionally, the solution set for the convex problem is not necessarily a singleton, so biasing the obtained solution still seems reasonable. We hope our clarification addresses all your concerns, and we are happy to discuss them if you have any further questions. We would like to emphasize that SAMPa significantly advances SAM by breaking the sequential nature of its two gradient computations, enabling parallel execution. SAMPa not only maintains convergence guarantees in the convex setting but also, as shown in our empirical results, achieves the fastest computational time among five efficient SAM variants (including SSAM introduced in the rebuttal) while consistently improving generalization across multiple tasks. We kindly ask the reviewer to consider these contributions. Thank you again for your engagement!

Reviewer zVXP7/10 · confidence 3/52024-07-13

Summary

This paper proposes a modification of SAM, named SAMPa, which enables to fully parallelize the two gradient computations in SAM, in order to accelerating the training. By doubling the computation resources, parallelized SAM could approach to a twofold speedup of SAM. Theoretical analysis is provided for the convergence of the proposed algorithm. Empirical results show that SAMPa can match or even surpass SAM in test accuracy.

Strengths

1. This paper proposes a modification of SAM, named SAMPa, which enables to fully parallelize the two gradient computations in SAM, in order to accelerating the training. By doubling the computation resources, parallelized SAM could approach to a twofold speedup of SAM. 2. Theoretical analysis is provided for the convergence of the proposed algorithm. 3. Empirical results show that SAMPa can match or even surpass SAM in test accuracy.

Weaknesses

1. To approach to a twofold speedup of SAM, SAMPa requires double of the computation resources compare to SAM. 2. There are previous works actually reduce the overall computation overhead of SAM. I think these works should also be compared to SAMPa as baselines in the experiments. For example: Du, J., Zhou, D., Feng, J., Tan, V.Y., & Zhou, J.T. (2022). Sharpness-Aware Training for Free. NeurIPS 2022. -------- Most of my concerns are addressed according to the author's feedback.

Questions

1. Is there any theoretical explanation about why SAMPa could outperform SAM in test accuracy?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The limitations are well discussed in this paper.

Authorsrebuttal2024-08-12

Happy to provide further clarification

Dear reviewers, Thank you for your constructive feedback. During the rebuttal, we addressed the following key points, with detailed responses in the rebuttal: - **Choice of $y_t$:** The choice of $y_{t+1} = x_t - \eta_t \nabla f(y_t)$ enables telescoping in our analysis. Empirically, $\nabla f(y_t)$ closely approximates $\nabla f(x_t)$, supported by small differences and similar Hessian spectra between SAM and SAMPa-0. - **Why SAMPa-0.2 outperforms SAM:** We hypothesize that the inclusion of optimistic gradient descent in SAMPa-0.2 stabilizes the algorithm in nonsmooth optimization landscapes, improving performance when Lipschitz continuity is not assumed. - **Convex analysis:** We emphasize that convex analysis has been foundational for many successful deep learning optimizers and was instrumental in our design of the auxiliary sequence $y_t$. - **Efficient SAM variants:** Additional experiments comparing SAMPa with _SAM for Free_ and _Sparsified SAM_ further demonstrate its improvements in both generalization and efficiency. - **Extra resource:** SAMPa-$\lambda$ consistently outperforms SAM across models and datasets, achieving a 2$\times$ speedup with parallelization, and still delivers superior performance even on a single device. - **Data parallelization:** Our results show that data parallelization slows down training, underscoring the importance of gradient parallelization. - **Memory and communication:** SAMPa's memory usage per GPU is comparable to SAM's. Its communication overhead could potentially be reduced or hidden. As the discussion window is closing soon, if you have any additional questions, we would be happy to elaborate on them. Otherwise, we kindly request you to consider increasing the score. Thank you! Best regards, Authors

Reviewer LruU2024-08-12

Thanks for the detailed responses. I am still skeptical about A2.1, because i) the gradient norm might be too weak to measure the convergence in convex case; ii) convex problems are not the most suitable testbed for SAM. On the other hand, I do appreciate the additional experiments. This is reflected through the increased score.

Reviewer zVXP2024-08-12

I thank the authors for their feedback. Most of my concerns are addressed and I would like to raise the score. After reading the author's feedback, it brings to my attention that SAMPa-0.2 is the overall best which is basically a convex combination of ordinary gradients and SAM gradients. Note that there is actually a variant of SAM that does sth. similar: Zhao, Y., Zhang, H., & Hu, X. Penalizing gradient norm for efficiently improving generalization in deep learning. ICML 2022. I think the gradient penalization method above with tuned interpolation ratio would be a better baseline than SAM (and I'm sorry for bringing this up so late). However, since SAMPa-0 also beats SAM in some cases, it seems that the improvement of SAMPa comes not only from the interpolation.

Authorsrebuttal2024-08-13

Further response to Reviewer zVXP

We sincerely thank the reviewer for the positive response and for increasing the score. Thank you for bringing the paper about penalizing gradient norms to our attention. The convex combination of gradients on ordinary and perturbed weights in that method is indeed similar to SAMPa-$\lambda$, and we agree this may contribute to the performance improvement. However, it's important to note that SAMPa-$\lambda$ differs in a key aspect: it computes gradients for each update on _two different batches_ (as shown in line 6 of Algorithm 1), while the penalizing method combines gradients from the same batch. We conducted preliminary experiments using the penalizing method on CIFAR-10 with the same hyperparameters as SAMPa-0.2. The results suggest similar performance in standard classification tasks, but worse outcomes with noisy labels. We would like to explore this further in future work. | | SAM |SAMPa-0.2 | Penalizing | |:---:|:---:|:---:|:---:| | Resnet-56 | 94.26 | 94.62 | 94.57 | | Resnet-32 (80% noisy label) | 48.01 | 49.92 | 48.26 | As mentioned in the rebuttal, another potential reason for SAMPa’s enhanced performance is that optimistic gradient descent may be better suited for nonsmooth optimization landscapes. We will include this discussion in the final version. Thank you again for your positive engagement, and we remain available for any further questions.

Reviewer c6Fm2024-08-14

Thanks for the author's reply, I appreciate the detailed response. However, based on the concern about motivation and theory limitation, I would give my final score as Boardline Accept.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC