SAMPa: Sharpness-aware Minimization Parallelized
Authors: Wanyun Xie
Thomas Pethick
Volkan Cevher
Venue: NeurIPS 2024 poster (2024)
Keywords: Sharpness aware minimization
efficient learning
generalization
supervised learning
optimization
Primary area: optimization_for_deep_networks
TL;DR: We parallelize the two gradient computations in SAM, resulting in both improved efficiency and generalization of the model.

Abstract: Sharpness-aware minimization (SAM) has been shown to improve the generalization of neural networks. However, each SAM update requires _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  _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 https://github.com/LIONS-EPFL/SAMPa.

SAMPa: Sharpness-awand Minimization Parallelized
Wanyun Xie EPFL (LIONS) wanyun. Sanyun.xie@epfl.c
Thomas Pethick EPFL (LIONS)
Volkan Cevher EPFL (LIONS)
thomas. pethi ckdepfl.ck
volkan. colkan..evher@eeff.cc ch
Abstract
Sharpness-aware minimization (SAM) has been shown to improve the general- ization neural networks. However, each SAM update requires 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 SAMM 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, method consistently outperforms SAM across both vision and language tasks. Notably, SAMPa theoretically maintains convergence guarantees even for fixed per- turbation sizes, which is established through novel Lyapunov function. We in fact arrive at SAMPa by treating this convergence guarantee as hard requirement-ai approach believe is promising for developing SAM-based methods in general. Our code is available at https: // thub. com/LIONS WtW./LIIIS-EEPL/AAEEE EPFL/SAMPa.
Introduction
The rise in deep neural network (DNN) usage has spurred a resource examination of training optimiza- tion methods, particularly focusing on bolstering their generalization ability. Generalization refers to a DNN's proficiency effectively processing and responding to new, previously unseen data originat- ing from the same distribution as the training dataset. A DNN with robust generalizability can reliably perform well on real-world tasks, when confronted with novel data instances or when quantized.
Improving generalization poses significant challenge in machine learning. Recent studies suggest that smoother loss landscapes lead to better generalization [Keskar al., 2017, Jiang* et al., 2020]. Motivated by this concept, Sharpness-Awen Aware Minimization (SAM) has emerged as a promising optimization approach [Foret et al., 2021, Zheng et al., 2021, Wu al., 2020b]. It is the current state-of-the-arr to seek flat minima by solving min-max optimization problem, in which the inner maximizer quantifies the sharpness as the maximized change of training loss and the minimizer both the vanilla training loss and the sharpness. As result, SAM significantly improves the generalization ability of the trained DNNs which has been observed across various supervised learning tasks in both vision and language domains [Foret et al., 2021, Bahri et al., 2021, Zhong et al., 2022]. Moreover, some variants of SAM improve its generalization further [Kwon et al., 2021, Kim al., 2022].
Although SAM and some variants achieve remarkable generalization improvement, they increase the computational overhead of the given base optimizers. In SAM algorithm [Foret et al., 2021], computations: one for computing the perturbation
each update consists of two and the other for computing the update direction. Since these two computations are not parallelizable, SAM doubles the training time compared to the standard empirical risk minimization (ERM).
Several variants of SAM have been proposed improve its efficiency. A common strategy involves integrating SAM with base optimizers an alternating fashion like RST [Zhao et al., 2022b], LookSAM [Liu al., 2022], and AE-SAM [Jiang al., 2023]. Moreover, ESAM [Du et al., 2022a]
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

uses fewer samples and updates fewer parameters to decrease the computational cost. However, some of these algorithms are suboptimal and their computational time overhead cannot be ignored completely. Du et al. [2022b] utilize loss trajectory instead of a single ascent step to estimate sharpness, albeit at the expense of memory consumption due to the storage of historical outputs or past models.
Since the runtime of SAM critically depends on the sequential computation of its gradients, we ask
Can we perform these two gradient computations in parallel?
In the sequel, we will answer this question in the affirmative. Note that since the second gradient computation highly depends on the first one seeking the worst case around the neighborhood, is challenging to break the sequential relationship between two gradients in one update.
To this end, we introduce computations completely. Furthermore, we also integrate the optimistic gradient descent method with parallelized version of SAM. Our final algorithm, named SAMPa, not only allows for a theoretical speedup up to 2x when there is no communication overhead but also improves the generalization further. Specifically, make the following contributions:
optimization sequence that allows us parallelize these two gradient
Parallelized formulation of SAM. We propose a novel parallelized solution for SAM, which breaks the sequential nature of two gradient computations in each SAM's update. It enables the simultaneous calculation of both gradients, potentially halving the computational time compared to vanilla SAM. We also integrate this parallelized method with the optimistic gradient descent method, known for its stabilizing properties, finalized to SAMPa.
Convergence guarantees. Our theoretical analysis establishes a novel Lyapunov function, through which prove convergence guarantees of SAMPa even with fixed perturbation size. We arrive at SAMPa by treating this convergence guarantee as a hard requirement, which we believe is promising for developing other SAM-based methods.
Improved generalization and efficiency. Our numerical evidence shows that SAMPa significantly reduces overall computational time even with basic implementation while achieving superior generalization performance. Indeed, SAMPa requires the least computa- tional time compared to the other four efficient SAM variants while enhancing generalization across different tasks. Notably, the relative improvement from SAM to SAMPa is 62.07% on CIFAR- 10 and 3 on CIFAR-I00, comparable to the gains from SGD to SAM. SAMPa also shows benefits on large-scale dataset (ImageNet- (K)) image and NLP fine-tuning tasks, as well as noisy label tasks, with the capability to integrate with other SAM variants.
2 Background and Challenge of SAM
This section starts with a brief introduction to SAM and its sequential nature of gradient computations. Subsequently, we discuss naive attempts including an approach from existing literature and our initial attempt which serve essential motivation for constructing final algorithm in the next section.
2.1 SAM and its challenge
Motivated by the concept of minimizing sharpness to enhance generalization, SAM attempts to enforce small loss around the neighborhood in the parameter space [Foret et al., 2021]. is formalized by a minimax problem
min max f(x +c) cld
where f a model parametrized by a weight vector x, and p is the radius of considered neighborhood.
The inner maximization of Equation (1) seeks for maxima around the neighborhood. To address the inner maximization problem, Foret et al. [2021] employ a first-order Taylor expansion of f(x e)
with respect to e in proximity to 0. This approximation yields:
Vi(r)
e" = arg max f(z + c) Ro arg max f(x) + (Vf(z), <ille|/Sp
-||(( [v)())
(2)
et =<lls<p

---

SAM first obtains the perturbed weight ã = + c* by this approximated worst-case perturbation and then adopts the gradient of ī to update the original weight z. Consequently, the updating rule at each iteration during practical training is delineated as follows:
(xr)
F. xt +
T++1 = 2x T/e (T+)
(SAM)
[+)()))
It is apparent from SAM, that the update requires two gradient computations for each iteration, which are on the clean weight 2, and the perturbed weight T. respectively. These two computations are not parallelizable because the gradient at the perturbed point (āx, highly depends on the gradient v f(zz) through the computation of the perturbation ž₁. Therefore, SAM doubles the computational overhead as well as the training time compared to base optimizers e.g., SGD.
2.2 Naive attempts
The computational overhead of SAM is primarily due to the first gradient for computing the pertur- bation as discussed in Section 2.1. Can we avoid this additional gradient computation? Random perturbation offers alternative the worst-case perturbation in SAM, as made precise below:
Et
x = Tt + P
with et ~ N(,11
Tell
(RandSAM)
x ++ = x² 1/s f(
Unfortunately, has been demonstrated empirically that RandSAM does not perform as well as SAM [Foret et al., 2021, Andriushchenko and Flammarion, 2022]. The poor performance of RandSAM maybe not surprising, considering that RandSAM does not converge even for simple convex quadratics as demonstrated in Figure 1.
We argue that the algorithm we construct should least be able solve the original minimization problem. Recently, Si and Yun [2024, Thm. 3.3] very interestingly proved that SAM converges for convex and smooth objectives even with a fixed perturbation size p. Fixed p is interesting to study, firstly, because it is commonly used and successful in practice [Foret et al., 2021, Kwon et al., 2021] I.. Secondly, convergence results relying on decreasing perturbation size are usually agnostic to the direction of the perturbation [Nam et al., 2023, Khanh et al., 2024), so the results cannot distinguish between RandSAM and SAM, which behaves strikingly different practice.
The fact that SAM uses the gradient direction ((() in the perturbation update, turns out to play an important role when showing convergence. It is thus natural ask whether another gradient could be used instead. Inspired by the reuse of past gradients in the optimistic gradient method [Popov, 1980, Rakhlin and Sridharan, 2013, Daskalakis et al., 2017], an intuitive attempt using the previous gradient at the perturbed model, such that V fy( = fz as outlined in the following update:
(
3 = 2t +
Jt+1 = Tt - Ty (G.))
(OptSAM)
(+
Notice that only one gradient computation is needed for each update. However, the empirical findings detailed in Appendix B.1 reveal that OptSAM fails to match SAM and even performs worse than SGD. In fact, such failure is already apparent in a simple toy example demonstrated in Figure 1, where OptSAM fails to converge. It is not surprising to see its failure. To be specific, in contrast with the optimistic gradient method, 2G in OptSAM represents an ascent step from x; while I++1 denotes a descent step from 2t, making (ī+ poor estimate of V ff(+1)). ). In the subsequent Section detail principled way of correcting this issue by developing SAMPa.
Toy example. We use a toy example f(x) | test if an algorithm can be optimized. We show the convergent perfor- mance of SAM, two naive attempts in this section, and SAMPa-A that is our algorithm proposed in Section The results in Figure demonstrate that RandSAM and OptS AM fail to converge, whereas SAM and SAMPa-X converge successfully.
12:
30
RandSAM
SAM
OALSAM
SAMPa-O
SAMPS-OUS
Iteraiio
Figure 1: Comparison on f(x) = 11:111.
3

---

Algorithm SAM Parallelized (SAMPa)
Input: Initialization 2g E IR, initialization Uo = zy and Iu = Vyuu,, Bu), iterations T, step sizes
{nh | 1
neighborhood size 0, interpolation ratio A.
for ( = to T - do
Keep minibatch B, and sample minibatch B++1- Compute perturbed weight 3. - It. + P P[s:
gt
Compute the auxiliary sequence 1/+++1 - dt - 1h.9t-
Compute gradients 9: = V B.) and g++1 = Vf(U++1,B+) in parallel.
Obtain the final gradient G; = (1 \)g. + Age+1-
Update weights 21+1 2t 7) G-
SAM Parallelized (SAMPa)
As discussed in Section 2.2, we wish to ensure that our developed (parallelizable) SAM variant maintains convergence convex smooth problems even when using a fixed perturbation size. To break the sequential nature of SAM, we seek to replace the gradient /((a))) with another gradient Vs(y) computed at some auxiliary sequence (yi (Wi)eeN- Provided the importance of the gradient direction far; in the convergence proof of SAM, are interested in picking the sequence (gr (e)) such that the difference fxx) V [yr)l can be controlled. Additionally, we need to ensure that fā,) and (V++1) can be computed in parallel.
Considering these design constraints, we arrive at the SAMPa method that is similar to SAM apart from the gradient used in perturbation calculation is computed at the auxiliary sequence (y/ keN, as illustrated in the following update:
V
Fz It + V (v:)l
(SAMPa)
Ut+1 = St (W))
Ft+1 T, T (
where the particular choice 1/++1 a direct consequence of the analysis, as discussed in Appendix C. Importantly, V and (f()++)) can be computed in parallel in this case. Intuitively, if V (yn) and (xx) are not too different then the scheme will behave like SAM. This intuition will be made precise by potential function used in the analysis of Section 4.
In SAMPa the gradient at the auxiliary sequence V ( ) only used for the perturbation update. It is reasonable to ask whether the gradient can be reused elsewhere in the update. As 3/++1 can be viewed as an extrapolated sequence of zt, it is directly related to the optimistic gradient descent method [Popov, 1980, Rakhlin and Sridharan, 2013, Daskalakis al., 2017] as outlined below:
V1+1 = 24 - 7/ V (yt), J++1 = x; n V/
(OpGGD
This celebrated scheme is known for its stabilizing properties as made precise through its ability to converge even for minimax problems. By simply taking convex combination of these two con- vergent schemes, I++1 (1-A) (1-A) SAMPa(z;)) (z)) we arrive at the following update rule:
ff()))
T, 21 +
[
(SAMPa-X)
Vi+1 = no (
21+1 22 (1 ()V (
where À E [0, 1]. Notice that SAMPa is obtained as the special case SAMPa-0 whereas SAMPa-1 recovers OptSAM. Importantly, SAMPa-A still admits parallel gradient computations and requires the same number of gradient computations as SAMPa.
SAMPa with stochasticity. An interesting observation in the SAM implementation is that both gradients for perturbation and correction steps have to be computed on the same batch; otherwise, SAM's performance may deteriorate compared to the base optimizer. This is validated by our

---

empirical observation in Appendix B.2 and supported by [Li and Giannakis, 2024, Li et al., 2024].
Therefore, we need be careful when deploying SAMPa practice.
Considering the stochastic setting, we present the finalized algorithm named SAMPa Algorithm 1. Note that 0: = V B.) represents the stochastic gradient estimate of the model T2 mini-batch B, and similarly 9t+11 V f [Br+1. B.+1) is the gradient of the model Vt+1 on mini-batch B++1- This ensures that the gradient 9t. used calculate the perturbed weight T. (line 3), is computed on the same batch as the gradient Hr. As demonstrated in line 5, SAMPa also requires 2 gradient computations for each update. Despite this, SAMPa only needs half of the computational time of SAM because and g++11 are calculated in parallel.
4 Analysis
In this section, we will show convergence of SAMPa even with a nondecreasing perturbation radius. The analysis relies on the following standard assumptions.
Assumption 4.1. The function f Rd -> TR is convex.
Assumption 4.2. The operator V Rd Id is L-Lipschitz with L (0, oc), i.e.
[z) V f - -|| vll Vx, y E R".
The direction of the gradient used in the perturbation turns Specifically, we will show that the auxiliary gradient (B()) in SAMPa can safely be used as replacement of the original gradient V (xx) in SAM, since we will be able control their difference. This is made precise by the following potential function used in our analysis:
play crucial role in the analysis.
Vi = f(x) + 2 (1 [V (f()))
As the potential function suggests we will be able to telescope the last term, which means that our convergence will remarkably only depend on the initial difference ( (xo)ll whose dependency we can remove entirely by choosing the initialization zo yo- See the proof of Theorem 4.4 for details. In the following lemma we establish descent of the potential function.
Lemma 4.3. Suppose Assumptions 4.1 and 4.2 hold. Then SAMPa satisfies the following descent inequality for > and decreasing sequence (1): (n)))eN with Be E (0, max (1, s/r.)) and c E (0, 1),
V4+1 & V₂- - nc mu (())) 2 + PC
where C'= ='= + L3 1 L*).
Notice that 7) is importantly squared in front of the error term p²C, while this is not the case for the term -I[V 2 This allows us control the error term while still providing convergence in terms of | f((xx))]* as made precise by the following theorem.
[zxx]][
Theorem 4.4. Suppose Assumptions 4.1 and 4.2 hold. Then SAMPa satisfies the following descent inequality for > and decreasing sequence (Ne)tN with T) E (0, max 1/21)),
4=(1-11/2) E EE (1-
V (x,))1 2 ActCp² ET-! of
(3)
E me(1-
where Ag = f(z))= - inf BER f(x) and C =
2
+ 24.
3
For Th = min
PVOT
max
27. 1}}
min:= [f = O LA
VT
Remark 4.5. Convergence follows long as Σi= 7t xo and Σ Ei=o The Jo, since the stepsize allows the right hand side to be made arbitrarily small. Note that Theorem 4.4 even allows for an increasing perturbation radius Pt, since it suffice to assume Et= d) co and Er mipt < co.
5 Experiments
In this section, we demonstrate the benefit of SAMPa across a variety of models, datasets and tasks. It is worth noting that to enable parallel computation of SAMPa, we perform the two gradient calculations across 2 GPUs As shown in Algorithm 1, one GPU computes V Vf B,) while another computes ( B.+1), For implementation guidance, we provide pseudo-code in Appendix E, along with algorithms detailing the integration of SAMPa with SGD and Adam both used base optimizers in this section.
5

---

5.1 Image classification
CIFAR-10/100. We follow the experimental setup of Kwon et [2021]. We use the CIFAR- 10 and CIFAR-100 datasets [Krizhevsky et al., 2009], both consisting of 50 000 training images of size 32 x 32, with 10 and 100 classes, respectively. For data augmentation, we apply the commonly used random cropping after padding with 4 pixels, horizontal flipping, and normalization using the statisties of the training distribution at both train and test time. We train multiple variants of VGG [Simonyan and Zisserman, 2014], ResNet [He et al., 2016], DenseNet [Huang et al., 2017] and WideResNet [Zagoruyko and Komodakis, 2016] (see Tables and 2 for details) using cross entropy loss. All experiments are conducted on NVIDIA A A100 GPU.
The models are trained using stochastic gradient descent (SGD) with a momentum of 0.9 and weight decay of 5 x 10 both as baseline and as the base model for SAM variants. We used batch size of 128 and a cosine learning rate schedule that starts 0.1. The number of epochs is set to 200 for SAM and SAMPa while SGD are given 400 epochs. This is done in order to provide a computational fair comparison as SAM and SAMPa use twice as much gradient computation. Moreover, we show S SAMPa-0.2 trained with 400 epochs as reference in the last column colored gray because SAMPa's theoretical limit of the per iteration cost comparable to SGD. Note that all SAMPa-X in this section use the same number of epochs as SAM only except for the last column of Tables I and 2. Label smoothing with a factor of 0.1 is employed for all methods.
Through a grid search over {0.01, 0.05, 0.1, 0.2, 0.4] using the validation dataset on CIFAR-10 10 with ResNet-56, SAM is assigned p values of 0.05 and 0.1 on CIFAR-10 10 and CIFAR- 100 respectively. which is consistent with existing works [Foret et al., 2021, Kwon et al., 2021]. Moreover, SAMPa-0 shares the same P value as SAM while SAMPa-0.2 is configured with twice the value of SAM's ss p. Additionally, À for SAMPa-A is set at 0.2 through a grid search from 0 to 1, with intervals of 0.1, with results detailed in Appendix B.3.
Training data is randomly partitioned into 90% for training and 10% for validation. To prevent overfitting on the test set, we deviate from Foret al. [2021], Kwon et al. [2021] by selecting the model with the highest validation accuracy to report test accuracy. Results are averaged over 6 independent executions and presented in Tables and 2. Compared with the enhancement from SGD to SAM, the average improvement from SAM to SAMPa-0.2 reaches 62.07% and 32.65% on CIFAR-10 10 and CIFAR- 100, respectively.
Table I: Test accuracies on CIFAR-10. SAMPa-0.2 outperforms SAM across all models with halved total temporal cost." "Temporal cost" represents the number of sequential gradient computations per update. SAMPa-0.2 with 400 epochs is included for comprehensive comparison with SGD and SAM.
Model
SGD
SAM
SAMPa-0
SAMPa-0.2
SAMPa-0.2
Temporal cost/Epochs
X1/100
x2/200
x1/200
x1/200
x1/100
121
96.1420.09
96.49+0.14
96.53+0.11
90.17
96.9 96.92+0.09
Resnet-56
94.20
94.26
94.31 10.43
94.62 10.35
95.43
VGG19-BN
94.76=00.1
95.05+0.17
95.06
95.11 10
95.34+0.07
WRN-28-2
95.71
95.98
96.0610.10
96. 13 10.14
96.31
WRN-28-10
96.77=0.21
97. 97.25+0.09
97.2 97.24=0.11
97.3440.08
97. 97.4640.07
Average
05.52_o.10
95.81 10.15
95.8610.11
95.99 10.08
96.29 10.06
Table 2: Test accuracies on C00. SAMPa-0.2 outperforms SAM across all models with halved total temporal cost. "Temporal cost" represents the number of sequential gradient computations per update. SAMPa-0.2 with 400 epochs is included for a comprehensive comparison.
Model
SGD
SAM
SAMPa-0
SAMPa-0.2
SAMPa-0.2
Temporal cost/Epochs
x1/400
x2/200
x1/200
x1/200
x1/400
12
81.08
82.53
82.50
82.70.1o.2
83.4440.21
Resnet-56
74.09=0.39
75, 75.14+0.15
75.22=0.20 20
75. 49.21
75. 75.8440.27
VGG19-BN
74.85
74.94+0.12
74.94+0.17
75.38+0.31
76. 15.2330..16
WRN-28-2
=0.17
78.5040.24
78.4540.29
78. 78.8241.22
79.464..20
WRN-28-10
81.5620.22
83.37 +0.30
83.4620.25
83, 83,90
83.9140.18
Average
77.92=0.17
78.90
78.91 10.00
79.2244.11
79.7840.09

---

ImageNet-IK. We evaluate SAM and SAMPa-0.2 on ImageNet-IK [Russakovsky et al., 2015], using 90 training epochs, a weight decay of 10-1, and batch size of 256. Other parameters match those CIFAR-10 Each method undergoes independent experiments, with test accuracies detailed in Table 3. Note that we omit SGD experiments due to computational con- straints; however, prior research confirms SAM and its variants outperform SGD [Foret et al., 2021, Kwon et al., 2021].
Table 3: Top1/Top5 maximum
test accuracies ImageNet-IK.
SAM
SAMPa-00.
Topl Top5
77. 77.25
77. 77.44
93.60
93.6940.08
5.2 Efficiency comparison with efficient SAM variants
To comprehensively evaluate the efficiency gains of SAMPa compared to other variants of SAM in practical scenarios, we conduct experiments using five additional SAM variants on the CIFAR-10 dataset with ResNet-56 (detailed configuration in Appendix B.4): LookSAM [Liu et al., 2022], AE-SAM [Jiang et al., 2023], SAF [Du et al., 2022b], MESA [Du et al., 2022b], and ESAM [Du et al., 2022a]. Specifically, LookSAM alternates between SAM and a base optimizer periodically, while AE- SAM selectively employs SAM when detecting local sharpness. SAF and MESA climinate the ascent step and introduce an extra trajectory loss term reduce sharpness. ESAM leverages two strategies, Stochastic Weight Perturbation (SWP) and Sharpness- Data Selection (SDS), for efficiency.
The number of sequentially computed gradients, shown in Figure 2a, serves computational time in an ideal scenario. Notably, SAMPa, SAF, and MESA require the fewest number of sequential gradients, each needing only half of SAM's. Specifically, SAF and MESA necessitate just one gradient computation per update, while SAMPa parallelizes two gradients per update.
metric for
However, real-world computational time encompasses more than just gradient computation; it includes forward and backward pass time, weight revision time, and potential communication overhead in distributed settings. Therefore, present the actual training time in Figure 2b, revealing that SAMPa and SAF serve as the most efficient methods. LookSAM and AE-SAM, unable to entirely avoid com- puting two sequential gradients per update, exhibit greater time consumption than SAMPa as expected. MESA, requiring an additional forward step compared to the base optimizer during implementation, cannot halve the computation time relative to SAM's. Regarding ESAM, solely integrate SWP in this experiment, as no efficiency advantage is observed compared to SAM when SDS is included. The reported time of SAMPa-0.2 in Figure 2b includes 7.5% communication overhead across GPUs. Achieving a nearly 2x speedup in runtime could be possible with faster interconnects between GPUs. In addition, the test accuracies and the wall-clock time per epoch reported in Table 4. SAMPa-0.2 achieves strong performance and meanwhile requires near-minimal computational time.
05
85
Valicatton
Validztionn
SAM
SAMPa-0.2
SAMPPs-.2
LockSAN
LookSAN
AE-SAM
SAF
MESA
MESA
20000
120020 Number Sequentiaa Gradients
40000
00000
85000
100000
140000
1000
1500
2000
2500
3000
3300
Taining Time
(a) Number of sequential gradients
(b) Actual running time
Figure 2: Computational time comparison for efficient SAM variants. SAMPa-0.2 requires
near-minimal computational time in both ideal and practical scenarios.
Table 4: Efficient SAM variants. The best result is in bold and the second best is underlined.
SAM
SAMPa-0.2
LookSAM
AE-SAM
SAF
MESA
ESAM
Accuracy 94.26
94.62
91.42
94.46
93.89
94.23
94.21
Time/Epoch (s)
18.81
10.94
16.28
13.47
10.09
15.43
15.97

---

5.3 Transfer learning
We demonstrate the benefits of SAMPa in transfer learning across vision and language domains.
Image fine-tuning, We conduct transfer learning experiments using the pre-trained ViT-B/16 checkpoint from Visual Transform- ers [Wu et al., 2020a], fine-tuning it on CIFAR-10 and CIFAR-I0W 1-0 datasets. AdamW is employed the base optimizer, with gra- dient clipping applied at global norm of 1. Training runs for 10 epochs, with peak learning rate of 10-4 Other parameters remain consistent with those outlined in Section 5.1. Results in Table 5 show the benefits of SAMPa in image fine-tuning.
Table 5: Image fine-tuning. C10 and C100 represent CIFAR-10 10 and CIFAR- 100 respectively.
SAM
SAMPa-0.2
C10 98.87.0.09 98 10.04
C100 91.7940.12
93 93.0640.16
NLP fine-tuning, To explore if SAMPa can benefit the natural language processing (NLP) domain. we show empirical text classification results in this section. In particular, we use BERT- model and finetune on the GLUE datasets [Wang al., 2018]. We use Adam the base optimizer under a linear learning rate schedule and gradient clipping with global norm 1. We set the peak learning rate to 10-5 and batch size to 32, and run 3 epochs with an exception for MRPC and WNLI which are significantly smaller datasets and where we used 5 epochs. Note that we set p = 0.05 for all datasets except for CoLA with p = 0.01, and RTE and STS-B with p = 0.005. The setting of p uniformly applied across SAM, SAMPa-0 and SAMPa-O.I. We report the results computed over 10 independent executions in the Table 6, which demonstrates that SAMPa also benefits in NLP domain.
Table 6: Test results of BERT-base fine-tuned on GLUE
CoLA
SST-2
MRPC
STS-B
QQP
MNLI
QNLI
RTE
WNLI
Method
GLUE
Mcc.
Acc.
Acc./F1.
Pear/Spea.
Acc./F1.
Acc.
Acc.
Acc.
Acc.
Adam W
74.6
56.6
91.6
85.6/89.9
85.4/85.3
90.2/86.8
82.6
89.8
62.4
26.4
SAM
76.6
58.8
92.3
86.5/90.5
85.0/85.0
90.6/87.3
83.9
90.4
60.6
41.2
-W SAMPa-0
76.9
58.9
92.5
86.4/90.4
85.0/85.0
90.6/87.6
83.8
90.4
60.4
43.2
-w SAMPa-0.1
78.0
58.9
92.5
86.8/90.7
85.2/85.1
90.7/87.7
84.0
90.5
61.3
51.6
5.4
Noisy label task
We test on a task outside the i..dd setting that the method was designed for. Following Foret et al. [2021] we consider label noise, where fraction of the labels in the training
corrupted to another label sampled uniformly at random. Through grid search over f0.005, 0.01, 0.05, 0.1, 0.2). set p 0.1 for SAM, SAMPa-0 and SAMPa-0.2 except for adjusting p 0.01 when the noise rate is 80%. Other experimental setup is the same as in Section 5.1. We find that SAMPa-0.2 enjoys better robustness to label noise than SAM.
Table 7: Test accuracies of ResNet-32 models trained on CIFAR-10 with label noise.
Noise rate
SGD
SAM
SAMPa-0
SAMPa-00.
0%
94.22
94.3644.07
94.3640.12
94.41 94.41+0.08 +0.08
20%
88.65
92.20.40.06
92.2240.10
92.39 +0.09
40%
84.24
89.78
89.75
90.01 10.18
60%
76. 76.2992. 76.29-0.25 =0.25
83.8340.51
83.81 +0.37
84.38 +0.07
80%
44.44+1.20
48.01+1.6
48.2241.71
49.92+1.15
5.5
Incorporation with other SAM variants
We demonstrate the potential of SAMPa enhance generalization further by integrating it with other variants of SAM. Specifically, we examine the results of combining SAMPa with five SAM variants: mSAM [Foret et al., 2021, Behdin al., 2023], ASAM [Kwon al., 2021]. SAM-ON [Mueller et al., 2024], VaSSO [Li and Giannakis, 2024]. and BiSAM [Xie et al., 2024]. Our experiments utilize Resnet-56 on CIFAR-10 trained with SAM and SAMPa-0.2, maintaining the same experimental setup detailed in Section 5. Further specifics on the experimental configuration are provided in Appendix B.4. The results summarized in Table 8 underscore the seamless integration of SAMPa with these variants, leading to notable improvements in both generalization and efficiency.

---

Table 8: Incorporation with variants of SAM. SAMPa in the table denotes SAMPa-0.2. The incorporation of SAMPa with SAM variants enhances both accuracy and efficiency.
mSAM +SAMPa ASAM
+SAMPa SAM-ON +SAMPa
VaSSO
+SAMPa
BiSAM
+SAMPa
94.28
94.71
94.84
94.95
94.44
94.51
94.80
94.97
94.49
95.13
Related Works
SAM. Inspired by the strong correlation between the generalization of model and the flat minima revealed in [Keskar al., 2017, Jiang et al., 2020], Foret al. [2021] propose SAM secking for a flat minima to improve generalization capability. SAM frames minimax optimization problem that aims to achieve a minima whose neighborhoods also have low loss. To solve this minimax problem, the most popular way is using an ascent step to approximate the solution for the inner maximization problem with the fact that SAM with more ascent steps does not significantly enhance generalization [Kim et al., 2023]. Notably, SAM has demonstrated effectiveness across various supervised learning tasks in computer vision [Foret et al., 2021], with studies demonstrating the realm of NLP tasks [Bahri et al., 2021, Zhong et al., 2022].
Efficient variants of SAM. Compared with base optimizers like SGD, SAM doubles compu- tational overhead stemming from its need for an extra gradient computation for perturbation per iteration. Efforts to alleviate SAM's computational burden have yielded several strategies. Firstly, strategies integrating SAM with base optimizers an alternating fashion have been explored. For instance, Randomized Shurpness-Aware Training (RST) [Zhao et al., 2022b] employs Bernoulli trial to randomly alternate between the base optimizer and SAM. Similarly, LookSAM [Liu et al., 2022] periodically computes the ascent step and utilizes the previous direction to promote flatness. Additionally, Adaptive policy to Employ SAM (AE-SAM) [Jiang et al., 2023] selectively applics SAM when detecting local sharpness, indicated by the gradient norm.
Efficiency improvements have also been pursued by other means. Efficient SAM (ESAM) [Du et al., 2022a] enhances efficiency by leveraging less data, employing strategies such as Stochastic Weight Perturbation and sharpness-sensitive Data Selection to subset random variables or mini-batch elements during optimization. Moreover, Sparse SAM (SSAM) [Mi et al., 2022] and SAM-ON [Mueller et al., 2024] achieve computational gains by only perturbing a subset of the model's weights, which enhances efficiency during the backward pass when only sparse gradients are needed. Notably, Du et al. [2022b] offer alternative approaches, SAF and MESA, estimating sharpness using loss trajectory instead of a single ascent step. Nonetheless, SAF requires increased memory consumption due to recording the outputs of historical models and MESA needs one extra forward pass. We compare against these methods in Section 5.2, where we find that SAMPa leads to a smaller wall-clock time.
Conclusion and Limitations
This paper introduces Sharpness-aware Minimization Parallelized (SAMPa) that halves the temporal cost of SAM through parallelizing gradient computations. The method additionally incorporates the optimistic gradient descent method. Crucially, SAMPa beats almost all existing efficient SAM variants regarding computational time practice. Besides efficiency, numerical experiments demonstrate that SAMPa enhances the generalization among various tasks including image classification, transfer learning in vision and language domains, and noisy label tasks. SAMPa can be integrated with other SAM variants, offering both efficiency and generalization improvements. Furthermore, we show convergence guarantees for SAMPa even with a fixed perturbation size through a novel Lyapunov function, which we believe will benefit the development of SAM-based methods.
Although SAMPa achieves 2x speedup along with improved generalization, the computational resources required remain the same as SAM's, as two GPUs with equivalent memory (as discussed in Appendix D) are still needed. Future research could explore reducing costs by either: (i) eliminating the need for additional parallel computation, or (ii) reducing memory usage per GPU, making the resource requirements more affordable. Moreover, we prove convergence for SA MPa only in the specific case of x = 0, lcaving the analysis for general À as an open challenge for our future work.

---

Acknowledgements
We thank the reviewers for their constructive feedback. This work was supported by the Swiss National Science Foundation (SNSF) under grant number 200021 1_12111 This work was supported by Google. This work was supported by Hasler Foundation Program: Hasler Responsible AI (project number 21043). This research was sponsored by the Army Research Office and was accomplished under Grant Number W911NF-24-1-0445
References
Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-aware mini- mization. In International Conference on Machine Learning (ICML), 2022.
Dara Bahri, Hossein Mobahi, and Yi Tay. Sharpness-aware minimization improves language model generalization. In Annual Meeting of the Association for Computational Linguistics, 2021.
Kayhan Behdin, Qingquan Song, Aman Gupta, Ayan Acharya, David Durfee, Borja Ocejo, Sathiya Keerthi, and Rahul Mazumder. msam: Micro-batch-ane sarro-bathh-averaeed sharpness-aware minimization. arXiv preprint arXiv:2302..966 2023.
Constantinos Daskalakis, Andrew Ilyas, Vasilis Syrgkanis, and Haoyang Zeng. Training GANs with optimism. arXiv preprint arXiv:1711.0014: 2017.
Jiawei Du Hanshu Yan, Jiashi Feng, Joey Tianyi Zhou, Liangli Zhen, Rick Siow Mong Goh, and Vincent Tan. Efficient sharpness-aween minimization for improved training of neural networks. In International Conference on Learning Representations (ICLR), 2022a.
Jiawei Du, Daquan Zhou, Jiashi Feng, Vincent Tan, and Joey Tianyi Zhou. Sharpness-aware training for free. Advances in Neural Information Processing Systems (NeurIPS), 2022b.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware aware minimization for efficiently improving generalization. In International Conference on Learning Representations (ICLR), 2021.
Fengxiang He, Tongliang Liu, and Dacheng Tao. Control batch size and learning rate to generalize well: Theoretical and empirical evidence. Advances in Neural Information Processing Systems (NeurIPS), 2019.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Conference on Computer Vision and Pattern Recognition (CVPR), 2016.
Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinberger. Densely connected convolutional networks. In Conference Computer Vision and Pattern Recognition (CVPR), 2017.
Weisen Jiang, Hansi Yang, Yu Zhang, and James Kwok. An adaptive policy to employ sharpness- aware minimization. arXiv preprint arXiv:2304. 14647, 2023.
Yiding Jiang* Behnam Neyshabur* Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantas- tic generalization measures and where to find them. In International Conference on Learning Representations (ICLR), 2020.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations (ICLR), 2017.
Pham Duy Khanh, Hoang-Chan Luong, Boris Mordukhovich, and Dat Ba Tran. Fundamental convergence analysis of sharpness-aware minimization. arXiv preprint arXiv:2401.08M0 2024.
Hoki Kim, Jinseong Park, Yujin Choi, Woojin Lee, and Jaewook Lee. Exploring the effect of multi-step ascent in tharpness-award aware minimization, arXiv preprint arXiv:2302.10181, 2023.
10

---

Minyoung Kim, Da Li, Shell X Hu, and Timothy Hospedales. Fisher SAM: Information geometry and sharpness aware minimisation. In International Conference on Machine Learning (ICML). 2022.
Alex Krizhevsky, Geoffrey Hinton, al. Learning multiple layers of features from tiny images. 2009.
Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. ASAM: Adaptive sharpness- aware minimization for scale-invariant learning of deep neural networks. In International Confer- ence on Machine Learning (ICML), 2021.
Bingcong Li and Georgios Giannakis. Enhancing sharpness-aware optimization through variance suppression. Advances in Neural Information Processing Systems (NeurIPS), 2024.
Tao Li, Pan Zhou, Zhenghao He, Xinwen Cheng, and Xiaolin Huang. Friendly sharpness-aware minimization. Conference on Computer Vision and Pattrn Recognition (CVPR), 2024.
Yong Liu, Siqi Mai. Xiangning Chen, Cho-Jui Hsieh, and Yang You. Towards efficient and scalable sharpness-awar aware minimization. In Conference on Computer Vision and Pattern Recognition (CVPR), 2022.
Peng Mi, Li Shen, Tianhe Ren. Yiyi Zhou, Xiaoshuai Sun, Rongrong Ji, and Dacheng Tao. Make sharpness- awaare aware minimization stronger: A sparsified perturbation approach. Advances in Neural Information Processing Systems (NeurlPS), 2022.
Maximilian Mueller, Tiffany Vlaar, David Rolnick, and Matthias Hein. Normalization layers all that sharpness-awen aware minimization needs. Advances in Neural Information Processing Systems (NeurIPS), 2024.
Kyunghun Nam, Jinseok Chung, and Namhoon Lee. Almost sure last iterate convergence of sharpness- aware minimization. 2023.
Leonid Denisovich Popov. A modification of the arrow-hurwitz method of search for saddle points. Mat. Zamerki, 28, 1980.
Alexander Rakhlin and Karthik Sridharan. Online learning with predictable sequences. In Conference on Learning Theory, pages 991-1019. PMLR, 2013.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, al. Imagenet large scale visual recognition challenge. International Journal of Computer Vision (IJCV), 115, 2015.
Dongkuk Si and Chulhee Yun. Practical sharpness-aware minimization cannot converge all the way to optima. Advances in Neural Information Processing Systems (NeurIPS), 36, 2024.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv: 1409. 1556, 2014.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv::804 07461, 2018.
Bichen Wu, Chenfeng Xu, Xiaoliang Dai, Alvin Wan, Peizhao Zhang, Zhicheng Yan, Masayoshi Tomizuka, Joseph Gonzalez, Kurt Keutzer, and Peter Vajda. Visual transformers: Token- based image representation and processing for computer vision. arXiv preprint 2rXiv:2006.03677, 2020a.
Dongxian Wu, Shu-Tao Tao Xia, and Yisen Wang. Adversarial weight perturbation helps robust general- ization. Advances in Neural Information Processing Systems (NeurIPS). 2020b.
Wanyun Xie, Fabian Latorre, Kimon Antonakopoulos, Thomas Pethick, and Volkan Cevher. Improv- ing SAM requires rethinking its optimization formulation. International Conference on Machine Learning (ICML), 2024.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. arXiv preprint 1605.00 2016.

---

Yang Zhao, Hao Zhang, and Xiuyuan Hu. Penalizing gradient norm for efficiently improving generalization deep learning. In International Conference on Machine Learning (ICML), 2022a.
Yang Zhao, Hao Zhang, and Xiuyuan Hu. Randomized sharpness-aware training for boosting computational efficiency in deep learning. arXiv preprint arXiv:2203.09900 2022b.
Yaowei Zheng, Richong Zhang, and Yongyi Mao. Regularizing neural networks via adversarial model perturbation. Conference on Computer Vision and Pattern Recognition (CVPR), 2021.
Qihuang Zhong, Liang Ding, Li Shen, Peng Mi, Juhua Liu, Bo Du, and Dacheng Tao. Improving sharpness-awesd aware minimization with fisher mask for better generalization on language models. arXiv preprint arXiv:2210.05497, 2022.
12

---

A Proofs for Section 4
Lemma 4.3. Suppose Assumptions 4.1 and 4.2 hold. Then SAMPa satisfies the following descent inequality for p 0 and decreasing sequence (Nt) with Nue E (0, max(1,, and c (0, 1).
V4+1 S Ve -ne(1 )(x=() + p²p²
where C = 21. + L3 + 1- L*).
Proof. Using smoothness we have
f(z141) - f(x)) + (V)() Fe+1 21) + 211² =f(r)) NN (Vf V (z())) |V
[(x)) (1
- (xx))l
-n (T11 (V V(())) vf(z))
(Assumption 4.2) S f(xz) ( - //2 || [())) T
-1 (1 (2))) V) vf(z,))
(5)
The last term of (5):
(Vf(x)).Vf(Gi) vf(x))) (Vf(z,) V((()) + (. Vf(xx)
|1 (5t nı -f(aa)
+ (V;)) V.((),, Vf( vx(x,))
(Assumption 4.1) 2 (V)(z)) - Vf(u)), - Vf(z))
(6)
For the left term, using 2 (a, bbb) = |||| 11b1112 ||a n6ll we have
21
- (Vf(x))Vf((t) (()))))))))))))))--))))))))))-)) (xxx))
= 2({f(zz; - V(())) vf(2.))
-IVV(()) ////:///))///)00----///////>-1>- (17 (
V f( f22)))
/ /V/
(M(()) ||2 + (1 + بسايا |V] v
+[V - 2
S (1+0) [ - - + (n + (1-n,)² [ (
+I[V/(=)) - 2.
(7)
The first inequality is due to Young's 's inequality, -1la - b112 < - Itc ||l| + |6| with e > 0, while the second inequality follows from
(v.) - v ( # | Ve+1 2 1 (V((++11) (
(8)
where the last inequality due to Assumption 4.2. We can pick e
1-m6 - such that
1-nch (1
L. To verify that e > O, use that (πe)ee is decreasing to obtain
1- I- 2 2
where the last inequality uses that Th < 1/2, Rearranging shows that > 0.
With the particular choice of e, (7) reduces to
-2 (V fz,)) - Vf(E), vf(xt))
1- 1- |V - V 1)1 (1 + A)I -
I-7LL
+ f (2
(Assumption 4.2) <
1-
1-meh - f( (1 A2)
+ - (V
(10)
13

---

1,2(1 12)=
with A,
Plugging (6) and (10) back into (5) yields
1-m. 7
1-9111 -n2
f(x++1) + +1(11 + L) L) v
f((e)) 2 [f Vf(m)]]
- m - - + 1 ( (1-77 (1 + A₂) + 1.) 1?20²
(11)
What remains is to bound the latter term of (11) in terms of a constant independent of L. First notice that, due the first inequality of (9) and n) I. by assumption, we have
1-mL
ni 2 1-c.
It follows that
1- qu L((-
(1-mb 1-ca
(12)
1-
1- No+ 1L -1
where the last inequality uses 7/2 < and N) S 1. Expanding the last term of (11), we obtain
((1-m + A.) + L)L2
M7c + 2 11 -n)
+ L)
T-m
1- of
+ L (1 ni 1.) 1m + L + L2
1- M 1 11 12
I-
1+ 1- La 1 =: C2
(12) S
which completes the proof.
Theorem 4.4. Suppose Assumptions 4.1 and 4.2 hold. Then SAMPa satisfies the following descent inequality for p and decreasing sequence (Ne)LN with nu E (0, max 1/2L)).
Et= m=(1-m1/22)
[f(x))1 I Ap ET=' of
(3)
E (1-
E aa(1-11/3)
where Ag = f(zo) - inf FERd f(x) and C = 12+13
2 +
22
For mh = min
max
1}}
3
CT
21
min (f(x))) o O
AgC
Proof. The proof follows directly by telescoping the descent inequality from Lemma 4.3 after subtracting inf f(x) from which we have
(1- E EES (1-0-4/)))
([ Aut
vf
(13)
E. 41()31//3)
Using Lipschitz continuity from Assumption 4.2 have that
Vf(xo) - V 220) < L2
yoll =
(14)
where the last equality follows from picking the initialization yo C' simplifies and 17s 2L. obtain the guarantee in (3).
By picking 2. for which
Picking a fixed stepsize ?/t = 1/. the convergence guarantee (3) reduces to
ET- [V < 4 Tw + Cp2n))
(15)
Optimizing the bound suggests stepsize of /
An Cp
Thus, incorporating the other stepsize require-
ments, set 17 =minty A maxt z 1}}. There are three cases.
Case In 11= V
An for which (15) reduces T
[Vf(1)) S
C VT
(16)
14

---

1 Case 1 22
/V22
S
Ag Cp
for which
& (2 + pv3.C
(17)
T
Case 1I1 15
V
Ag Cp
additionally use that 1 2 2L to again establish (17).
Combining the three cases, we have that for any case
P
(18)
VT
Noting that the minimum is always smaller than the average completes the proof.
B Additional Experiments
B.1 Failure of OptSAM
To demonstrate the failure of our naive attempt, described in Section 2.2 as OptSAM, we provide empirical results using identical settings in Section 5.1. As shown in Table 9, OptSAM fails to outperform SAM and even performs worse than SGD.
Table 9: Test accuracies of OptSAM on CIFAR-10. 10.
Model
SGD
SAM
OptSAM
Resnet-56
94.20
94.26
93.99
WRN-28-2
95.71
95.98
95.41
VGGI9-BN
94.76
95.05
94.32
B.2 SAM with stochasticity
To deploy SAM with stochasticity, we find it imperative to utilize the same batch for both gradient calculations of perturbation and correction steps. Otherwise, the performance of SAM may be even worse than the base optimizer. Our empirical observations on CIFAR-10 are shown in Table 10, which is also validated by [Li and Giannakis, 2024, Li al., 2024].
This observation demonstrates that the same batch for both perturbation and correction steps is essen- sequential batches in SAMPa.
tial. This also justifies the need for parallel gradient computation
Table 10:
Two gradients in SAM computed on the same or different batch on CIFAR-10. SAM computes them the same batch while SAM-dbb is two different batches.
Model
SGD
SAM
SAM-db
Resnet-56
94.20
94.26
93.97
WRN-28-2
95.71
95.98
95.50
VGGI9-BN
94.76
95.05
94.48
B.3 Sweep over for SAMPa
To investigate the impact of different values of À in SAMPa, present test accuracy curves for ResNet-56 and WRN-28-111 10 on CIFAR-10 in Figure 3, covering the range [0, 1] with the interval 0.1. Notably, SAMPa-1 corresponds to OptGD.
our experiments, as reported Table and Table we initially optimize À 0.2 using ResNet-56. This default value is applied consistently across other models to maintain a fair comparison. However, continuous tuning of may lead to improved performance on different model architectures, as demonstrated in Figure 3b.
15

---

SAPY
SAMPS
To2t Accuracy (50
97.
7026 A2curss
4.4
97.1
97.0
04.3
12.9
443
75.8
(a) Resnet-56
(b) WRN-28-10
Figure 3: Test accuracy curve obtained from SAMPa algorithm using range of A.
B.4 Hyperparameters for variants of SAM
We present a comparison between SAMPa and various variants of SAM in Table 4 and Table 8. All algorithms in these tables utilize Resnet-56 CIFAR-10 with hyperparameters mostly consistent with those used in Table 1. However, some variants require additional or different hyperparameters, which are listed below:
LookSAM [Liu et al., 2022]: Update frequency &: = 5, scaling factor 0.7.
AE-SAM [Jiang et al., 2023]: p == 0.2, forgetting rate 0.9.
MESA [Du et al., 2022b]: Starting epoch Estan = 5, coefficients = 0.8, Decay factor 3 = 0.9995.
ESAM [Du et al., 2022a]: SWP probability 0.5.
mSAM [Foret et al., 2021]: Size of micro batch m = 32.
ASAM [Kwon et al., 2021]: = 0.5.
SAM-ON [Mueller et al., 2024]: p 0.5.
VaSSO [Li and Giannakis, 2024]: Lineeaizzaiinn parameter 0 0.4.
BiSAM Xie al. [2024]: BiSAM (-log) with μ =
We adhere the default values specified in the original papers and maintain consistent naming conventions. Following the experimental setup detailed in Section 5.1, we set p x for SAMPa in Section 5.5 when incorporated with the algorithms, while keeping other parameters consistent with their defaults. Note that these parameters are not tuned to ensure a fair comparison and avoid waste of computing resources.
B.5 mSAM with GPUs
Since SAMPa parallelizes two gradient computations across two GPUs, we implement mSAM [Behdin et al., 2023], SAM variant that achieves data parallelism, for fair comparison of runtime. Based on experiments in Section 5.2, mSAM (m=2) uses two GPUs and each computes gradient for 64 samples. While mSAM's ss total computation time for batch sizes of 64 and 128 is similar, its wall-clock time is slightly longer due to the added communication overhead between GPUs. This highlights the need for gradient parallelization.
Table 11: Runtime of SAM variants on 2 GPUs.
SAM
mSAM (m=2)
SAMPa-0.2
Number GPUs
2
Time/Fpoch (s)
18.81
21.17
10.94
We also provide the runtime per batch of SGD across various batch sizes in Table 12. The results show that data parallelism reduces time efficiently only when the batch size is sufficiently large. However, excessively large batch sizes can negatively affect generalization [He et al., 2019].
16

---

Table 12:
Runtime per batch/epoch of different batch sizes.
Batch size
64
128
256
512
1024
2048
4096
Time/Batch (ms)
21.70
22.70
23.08
27.84
32.88
50.00
120.16
B.6 SAMPa-λ v.s. the gradient penalization method
SAMPa-X A takes convex combination of the two convergent schemes x++1 (1 X) SAMPa(1)))
A (zt)) which is similar with a gradient penalization method [Zhao et al., 2022a] doing Ft-1 - X) SAM(2;; + À (z)()))) However, it is important to note that SAMPa-À 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 on CIFAR-I0 using the penalizing method with the same hyperparameters as SAMPa-0.2. The results indicate similar performance in standard classification tasks but show worse outcomes with noisy labels. Further investigation into this discrepancy may provide insights into SAMPa's superior performance.
Table 13: Test accuracy of the gradient penalization method.
SAM
SAMPa-0.2
Penalizing
Resnet-56
94.26
94.62
94.57
Resnet-32 (80% noisy label)
48.01
49.92
48.26
C
The choice of yt+1
The particular choice of M in SAMPa is a direct consequence of the analysis. Specifically, in Equation (8) of the proof, the choice M-+1 = 24 (wn) allows us to produce the term fur)) 2 in Equation (7). This is what we refer to in Section 3, when mentioning that we will pick such that || fxx) (yv))) (i.e. the discrepancy from SAM) can be controlled. This gives a precise guarantee explaining why fxx)) can be replaced by f
ffy++1)ll in order telescope with f - V [
Additionally, the small difference between the perturbations based on fx(x)) and ( suggests that V serves as an effective approximation of v)((x;) (Ix;) in practice. In Figure 4, we track the co- sine similarity and Euclidean distance between ( and V (xxe throughout the training process of ResNet-56 on CIFAR-I0 We find that the cosine similarity keeps above 0.99 during the whole train- ing process, and in most period it's around 0.998, while the end of training it is even close 1. This indicates that SAMPa's 's estimated perturbation is an excellent approximation of SAM's perturbation.
Moreover, the Euclidean distance decreases and is close to zero at the end of training. This matches our theoretical analysis that V fx+) (yr))l eventually becomes small, which lemma 4.3 guarantees in the convex case by establishing the decrease of the potential function V4-
casine
Eucl
1.000
SAMPM-0.2
0.912
0:08
Coafte annterk
0.988
dutidean detarce
0.00
0.997
3.948
0.995
SAMPS-
SAMP9-0.2
25
20.0 spou7hs
125
150
375
230
1100
3.25
1.35
175
(a) Cosine similarity
(b) Euclidean distance
Figure 4: Difference between V fxx) and ()()).
17

---

D Discussion of memory usage
From the implementation perspective, worth discussing the memory usage of SAMPa compared with SAM. As depicted in SAM, SAM necessitates the storage of maximum of two sets of model weights (Iy, Tt), along with one gradient (Vf(z;)) and one mini-batch (B,) for each update. Benefiting from 2 GPUs' deployment, SAMPa-X requires the same memory usage SAM on each GPU, specifically needing two model weights (æt, 3t or yt+1), ) one gradient (V)( or (V(++))) and one mini-batch (B, B.+1).
We present a memory usage comparison in Table 14 for all SAM variants introduced Section 5.2. Notably, SAMPa-0.2 requires slightly less memory GPU, while MESA consumes approximately 23% more memory than SAM. The other three methods have comparable memory usage to SAM. However, it's important to note that memory usage is highly dependent on the size of the model and dataset, particularly for SAF and MESA, which store historical model outputs or weights.
Table 14: Memory usage on each GPU.
SAM
SAMPa-0.2 LookSAM
AE-SAM
SAF
MESA
ESAM
Memory (MiB)
2290
2016
2296
2292
2294
2814
2288
E
Implementation guidelines
Our algorithm SAMPa is deployed across two GPUs to facilitate parallel training. As shown in Algo- rithm I, one GPU calculates V (ī,, B.) and another one takes responsibility for Vf(e+1,, 1/(y+11,B++)) B(+1). For ease of implementation, we provide detailed version in Algorithm 2, with the following key points:
Apart from the synchronization step (line 8), all other operations can be executed in parallel both GPUs.
The optimizer state, m. used in ModelUpdate,,.! O includes necessary elements such as step size, momentum, and weight decay. Crucially, to ensure that 3t+11 (line 6) is close T++1+ the update for Vt+1 uses my, the state associated with zy. Note that the optimizer state is not updated in line 6.
Algorithm 2 SAMPa on two GPUs
Input: Initialization 2n E RR, initialization to Iy and 9o V fyo, Bo), iterations T, step sizes (m)) 1-11 neighborhood size 0, interpolation ratio X, optimizer state mg-
for f = to do
GPU1: Load minibatch B,
GPU1: Compute perturbed weight F, Iz + P Iisll
GPU1: Compute gradient 9t fz,, B.).
GPU2: Load minibatch B.+1.
GPU2: Compute auxiliary sequence y411, ModelUpdate,,, (Weeggt).
GPU2: Compute gradient 9t+1 = V
B,+1).
Both: Communicate h. and 4++1 between GPUI and GPU2.
Synchronization barrier
Both: Compute the final gradient G;
X)Ør + Ag
10
Both: Update weights F1+1, Me+1 = ModelUpdate,,, (Ke, G.).
b Updates optimizer state
18

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Abstract and conclusion in Section clearly demonstrate our algorithm named SAMPa, a parallel version of SAM. We also claim that the theoretical result shows a convergence guarantee. The empirical results present that SAMPa not only enhances efficiency but also improves generalization.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: We discuss the limitation of our method in Section including communication overhead across GPUs and general theoretical analysis.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-tex system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
19

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: We demonstrate main assumptions and theoretical results in Section 4, and the proofs are provided in Appendix
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof' can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: All experimental setups with detailed configurations are provided in Section and Appendix B. Additionally, we give practical guidance in Appendix and public code in https: http///gtttbbbtt COom/LION--BE ELLLONS-EFF/SSSE EPFL/SAMPa
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to model. In general. releasing code and data often one good way to accomplish this, but reproducibilityy can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open- open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibiliity In the case of closed-source models, it may be that access to the model is limited in
20

---

some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We provide code in https: bttps://gttt :/githhb LEWLLLINNNEEEL/AAAEE which realizes parallelized version of SAMPa. We also give practical guidance in Appendix E.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (uttps://nips.cm
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed reproduce the results. See the NeurlPS code and data submission guidelines (https:
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The main experimental setting is shown in Section including data splitting, data augmentation, type of optimizers, and the choice of hyperparameters by grid search or following existing works. Some additional details are provided in Appendix B.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We report the average accuracy of multiple runs along with the variance in Section 5, which provides appropriate information about the statistical significance and reliability of our experimental results.
Guidelines:
21

---

The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We report we use NVIDIA A100 GPU for our experiments in Section 5. Moreover, we show actual running time in Figure 2b for several relative works.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: /nwuuriurrp.rrr cttpiiss..ppuuuue
Answer: [Yes]
Justification: All datasets we use are public datasets like CIFAR-10, CIFAR-100, and Im.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
22

---

Justification: Our paper is foundational research about an efficient optimizer and it is not tied to particular applications.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation. generating fake profiles, surveillance), fairness considerations (c.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our paper has no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-us should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We use public datasets and some pre-trained models, all of which are claimed and cited with original papers.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
23

---

The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g.. CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g.. website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi wewrsswithhodee thcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's ' creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do not release new assets.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the datasct/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing research with human subjects.
Guidelines:
24

---

The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
25


Reviews:
[{"review_id": "DUzsaEE6Dg", "reviewer": "Reviewer_zVXP", "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.", "questions": "1. Is there any theoretical explanation about why SAMPa could outperform SAM in test accuracy?", "limitations": "The limitations are well discussed in this paper.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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. \n\n2. Theoretical analysis is provided for the convergence of the proposed algorithm. \n\n3. 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.\n\n2. 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:\nDu, J., Zhou, D., Feng, J., Tan, V.Y., & Zhou, J.T. (2022). Sharpness-Aware Training for Free. NeurIPS 2022.\n\n--------\nMost of my concerns are addressed according to the author's feedback.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback and address all remaining concerns below:\n\n> Q1. To approach to a twofold speedup of SAM, SAMPa requires double of the computation resources compare to SAM.\n\nA1. Please note that the total computational time for SAMPa across all GPUs is comparable to that of SAM. In addition, we would like to emphasize that SAMPa-$\\lambda$ _consistently_ outperforms SAM across models and datasets (see Section 5) with 2$\\times$ speedup. It can show advantages even without parallelization. Specifically, one run SAMPa-$\\lambda$ on a single device has the same running time as SAM, but importantly _still leads to superior performance_.\n\n> Q2. 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.\n\nA2. We kindly remind that in Section 5.2, we have already considered one of two methods named MESA in \"Sharpness-Aware Training for Free\" as a baseline. For completeness, we also add result for SAF in the table below. Although SAF requires slightly less training time, its performance is worse than the other three methods. SAMPa-0.2 can reach the same accuracy as SAF only with 75% of SAF's epochs. Furthermore, SAMPa's time per epoch is 10.94s including 0.82s of communication overhead, which can be reduced further with faster device communication.\n\n|   | SAM |SAMPa-0.2 | MESA | SAF | \n|:---:|:---:|:---:|:---:|:---:|\n| Time/Epoch (s) |  18.81  |  10.94 | 15.43 | 10.09 |\n| Accuracy (%) | 94.26 | 94.62 | 94.23  | 93.89 | \n\n> Q3. Is there any theoretical explanation about why SAMPa could outperform SAM in test accuracy?\n\nA3. The method that consistently beats SAM is SAMPa-0.2, which is simply a convex combination of SAMPa and OptGD (see l. 138 for details). One hypothesis we currently have for its superior performance is that optimistic gradient descent might be better suited for nonsmooth parts of the optimization landscape, thereby stabilizing the algorithm when Lipschitz continuity assumption is not present (this assumption is crucial for allowing a fixed perturbation size). It is an interesting direction for future work to understand this superior performance better.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FNTCSbO0b1", "reviewer": "Reviewer_c6Fm", "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.", "questions": "See weakness.", "limitations": "The SAMPa requires 2 GPUs to realizing parallel computation, which will suffer communication cost and require more memory storage for auxiliary sequence.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The paper is well writen and mostly free of typos.\n- The motivation of this paper is meaningful for practical application.\n- Theoretical analysis supports the proposed method.\n- 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.\n- A difference between the approximated perturbation based on auxiliary sequence $y_t$ and the real perturbation would be helpful.\n- A Visualization of landscape and A spectrum of the Hessian(Figure 3 in [\"SHARPNESS-AWARE MINIMIZATION FOR EFFICIENTLY\nIMPROVING 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.\n- The convex assumption of analysis is too strong, so I'm concerned that its results would have little significance for practical operation.\n- 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=...$ \"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback and address all remaining concerns below:\n\n> Q1. \"Make Sharpness-Aware Minimization Stronger: A Sparsified Perturbation Approach\" should be considered.\n\nA1. We have conducted additional experiments on Sparsed SAM (SSAM) below and will include them in our paper for completeness, but note that we have already compared against _four_ efficient variants of SAM in Section 5.2. \n\nSSAM perturbs a subset of weights to achieve computational gains. There are two methods in the SSAM paper: SSAM-F uses Fisher information for sparse perturbation, and SSAM-D uses efficient binary mask generation. We choose SSAM-D with 95% sparsity for our experiments due to the high computation cost of SSAM-F's empirical Fisher. Importantly, SAMPa-0.2 still shows a significant speedup compared to SSAM-D.\n\n||SAM|SAMPa-0.2|SSAM-D|\n|:---:|:---:|:---:|:---:|\n|Time/Epoch (s)|18.81|10.94|16.95| \n|Accuracy (%)|94.26|94.62|94.48| \n\nAnother sparse SAM variant, SAM-ON that only perturbs normalization layers, has been discussed in Sections 5.5 and 6. Note that both SSAM and SAM-ON can be incorporated within SAMPa for further improvement.\n\n> Q2. A difference between the approximated perturbation based on auxiliary sequence $y_t$ and the real perturbation would be helpful.\n\nA2. We claim that the similarity between the gradient on the auxiliary sequence $\\nabla f(y_t)$ and the gradient on actual weight $\\nabla f(x_t)$ holds not just in theory (Section 4) but also in practice. \n\nIn the [figures](https://imgur.com/a/6hWQm44), we track the cosine similarity and Euclidean distance between $\\nabla f(y_t)$ and $\\nabla f(x_t)$ during whole training process of Resnet-56 on CIFAR-10. We find that the cosine similarity keeps above 0.99 during the whole training process, and in most period it's around 0.998, while at the end of training it is even close to 1. This indicates that SAMPa's estimated perturbation $\\rho \\frac{\\nabla f(y_t)}{||\\nabla f(y_t)||}$ is an excellent approximation of SAM's perturbation $\\rho \\frac{\\nabla f(x_t)}{||\\nabla f(x_t)||}$.\n\nMoreover, the Euclidean distance decreases and is close to zero at the end of training. This matches our theoretical analysis that $||\\nabla f(x_t) -\\nabla f(y_t)||^2$ eventually becomes small, which Lemma 4.3 guarantees in the convex case by establishing decrease of the potential function $\\mathcal V_t$.\n\n> Q3. A spectrum of the Hessian would help.\n\nA3. Thanks for the suggestion. We provide [Hessian spectra](https://imgur.com/a/SFA1Ndf) for the Resnet-56 trained with SAM, SAMPa-0 and SAMPa-0.2 on CIFAR-10. We focus on the largest eigenvalue $\\lambda_1$ and the ratio of the largest to the fifth largest eigenvalue $\\frac{\\lambda_1}{\\lambda_5}$ as suggested by the reviewer, which reflect the flatness of the solution. As expected, the model trained with SAMPa-0 converges to minima with similar curvature to SAM.\n\n||SAM|SAMPa-0|SAMPa-0.2|\n|:---:|:---:|:---:|:---:|\n|$\\lambda_1$|94.29|87.17|149.63| \n|$\\frac{\\lambda_1}{\\lambda_5}$|2.15|1.84|3.30| \n\nThe model trained with SAMPa-0.2, on the other hand, converges to a sharper minimum while it shows better generalization performance. The existence of a well-performing sharper minimum is maybe not too surprising, considering that the relationship between sharpness and generalization remains unclear. For instance, [Andriushchenko and Flammarion, 2022] shows that a flatter minimum does not always lead to better generalization, which is further supported by [Mueller et al., 2024] and [1].\n\n> Q4. The convex assumption is too strong, so I'm concerned its results have little practical significance.\n\nA4. We kindly disagree that convex analysis is too strong for guiding practical applications. Arguably most of the successful optimizers today for deep learning were developed under the assumption of convexity. This includes AdaGrad, Adam and even more recent developments such as Distributed Shampoo and the Scheduler-free optimizer that won the AlgoPerf competition this month [2].\n\nSimilarly, we found convex analysis extremely helpful in developing the algorithm. In our particular case our algorithmic design _falls out of the analysis_, specifically leading to our choice of the auxiliary sequence $(y_t)_{t\\in \\mathbb N}$ (see Section 3). We subsequently show that SAMPa indeed has strong practical performance in deep learning applications (see Section 5).\n\n> Q5. 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=\\dots$\"\n\nA5. The particular choice of $y_t$ is a direct consequence of the analysis! Specifically, in Eq. 7 of the proof, the choice $y_{t+1} = x_t - \\eta_t \\nabla f(y_t)$ allows us to produce the term $||\\nabla f(x_{t+1})-\\nabla f(y_{t+1})||^2$ in order to telescope with $||\\nabla f(x_{t})-\\nabla f(y_{t})||^2$ in Eq. 6. This is what we refer to in l. 124, when mentioning that we will pick $y_{t}$ such that $||\\nabla f(x_{t})-\\nabla f(y_{t})||^2$ (i.e. the discrepancy from SAM) can be controlled. This gives a precise guarantee explaining why $\\nabla f(x_{t})$ can be replaced by $\\nabla f(y_{t})$.\n\nEmpirically, the small difference between perturbations based on $\\nabla f(x_t)$ and $\\nabla f(y_t)$ (as discussed in Q2), along with similar Hessian spectra between SAM and SAMPa-0 (as shown in Q3) indicates that $\\nabla f(y_t)$ is an excellent approximation of $\\nabla f(x_t)$ in practice as well. Moreover, similar performance between SAMPa-0 and SAM supports our choice of $y_t$.\n\nWe will highlight Eq. 7 in the final version and include a discussion on the additional empirical results above.\n\n[1] Ankit Vani, Frederick Tung, Gabriel L. Oliveira, Hossein Sharifi-Noghabi. Forget sharpness: perturbed forgetting of model biases within SAM dynamics. In International Conference on Machine Learning (ICML), 2024.\n\n[2] https://mlcommons.org/2024/08/mlc-algoperf-benchmark-competition/", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rh6BLzyLK3", "reviewer": "Reviewer_LruU", "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.", "questions": "Q1. In algorithm 1, is $g_0=\\nabla f(y_0)$ or $g_0=\\nabla f(y_0, {\\cal B}_0)$?\n\nQ2. Can the authors comment more on the intuition on the reason that SAMPa outperforms SAM? \n\n\nThere are also typos:\n- line 180, we can trivially allows --> we can trivially allow\n- line 126, identical to SAM except .. --> consider to change the word 'identical'", "limitations": "The authors have discussed the limitations.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "S1. **Novel methodology.** Using a new sequence to break the sequential dependence of SAM gradients is novel.  \n\nS2. **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. \n\nW2. 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?\n\nW3. More on parallelization should be discussed. \nFor 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.\n\n\n\n[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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback and address all remaining concerns below:\n\n> Q1.  My major concern is that the runtime comparison is not fair due to the second GPU. In particular, since SAMPa uses 2 GPUs, can the authors also report data parallel for SAM with 2 GPUs as well? (or mSAM for communication efficiency). For example, if the batch size for SAM is 128, each GPU can calculate gradients for 64 samples in parallel.\n\nA1. We report the runtime for mSAM with 2 GPUs as suggested. Each GPU calculates gradients for 64 samples in parallel, based on experiments in Section 5.2. mSAM requires slightly longer wall-clock time due to similar computation times for batch sizes of 64 and 128, but with added communication overhead. This highlights the need for gradient parallelization. Moreover, SAMPa-0.2 can run on a single device with the same runtime as SAM, yet achieves superior performance.\n\n|   | SAM | mSAM (m=2) | SAMPa-0.2 |\n|:---:|:---:|:---:|:---:|\n| Time/Epoch (s) |  18.81  | 22.43 | 10.94 |\n\nWe additionally provide the time per batch and per epoch for SGD with different batch sizes, showing that doubling the batch size slightly increases computation time per batch, while the total epoch time decreases.\n\n| Batch size  | 64 | 128 | 512 | 2048 |\n|:---:|:---:|:---:|:---:|:---:|\n| Time/Batch (ms) | 21.70 | 22.70 | 27.84 | 50.00 |\n| Time/Epoch (s) | 15.28 | 7.99 | 2.45 | 1.10 |\n\n> Q2.1 Can the authors comment more on the optimality of this convergence rate? \n\nA2.1. The rate matches those of SAM (see Thm. 3.3 Si and Yun [2024]) and recovers the rate of gradient descent when $\\rho=0$, but they are not optimal for the given problem-class for which acceleration is possible.\n\nHowever, note that our goal is not to establish optimal rates, but we rather use the analysis to directly construct our algorithm, by requiring the guarantee to match that of SAM. Specifically in Eq. 7 of the proof, the choice of the auxiliary sequence $y_{t+1} = x_t - \\eta_t \\nabla f(y_t)$ is what allows us to produce the term $||\\nabla f(x_{t+1})-\\nabla f(y_{t+1})||^2$ in order to telescope with $||\\nabla f(x_{t})-\\nabla f(y_{t})||^2$ in Eq. 6. This is what we refer to in l. 124, when mentioning that we will pick $y_{t}$ such that $||\\nabla f(x_{t})-\\nabla f(y_{t})||^2$ (i.e. the discrepancy from SAM) can be controlled.\n\n\n> Q2.2 Does SAMPa-$\\lambda$ also enjoy theoretical guarantees?\n\nA2.2. We have convergence gaurantees for $\\lambda=0$, but not for more general $\\lambda$. However, we believe that it should be possible to generalize the result considering that SAMPa-$\\lambda$ is a convex combination of two convergent schemes as comment on in l. 136-138.\n\nWe would like to emphasize that the empirical contribution is valuable in itself regarding SAMPa-0.2: the method _consistently_ outperforms SAM with 2$\\times$ speedup across models and datasets (see Section 5). Note that this implies that it is possible to benefit from this method even without any parallelization. Specifically, one can run SAMPa-0.2 on a single GPU, in which case it would have the same running time as SAM, but importantly _still lead to superior performance_. \n\n> Q3. 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.\n\nA3. Thanks for the suggestion. We kindly remind that we have discussed the memory usage of SAMPa compared with SAM in Appendix E. We will give more detailed discussions about memory usage and communication overhead below.\n- **Memory usage:** In Appendix E, we mention that the memory usage of SAMPa on each GPU is the same with SAM's and note that SAMPa requires two GPUs while SAM needs one. For empirical results, we give memory usage comparison for all SAM variants in Section 5.2 below. SAMPa-0.2 even requires a little less memory usage on _each_ GPU, MESA requires roughly 23% extra memory compared with SAM, and other three methods have similar memory usage to SAM.\n\n|   | SAM | SAMPa-0.2 | LookSAM | AE-SAM |MESA| ESAM|\n|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n| Each GPU memory usage (MiB)| 2290 | 2016|2296| 2292 | 2814 | 2288 |\n\n- **Communication overhead:** In the experiment of Section 5.2, SAMPa's training time per epoch is 10.94s including 0.82s of communication overhead, compared to SAM's 18.81s. This indicates that SAMPa could approach a 2$\\times$ speedup with optimized communication. Moreover, communication overhead might be overlapped by data loading, and further research into this could be interesting.\n\n> Q4. Can the authors comment more on the intuition on the reason that SAMPa outperforms SAM?\n\nA4. The method that consistently beats SAM is SAMPa-0.2, which is simply a convex combination of SAMPa and OptGD (see l. 138 for details). One hypothesis we currently have for the superior performance is that optimistic gradient descent might be better suited for nonsmooth parts of the optimization landscape, thereby stabilizing the algorithm when Lipschitz continuity assumption is not present (this assumption is crucial for allowing a fixed perturbation size). It is an interesting direction to understand this superior performance better.\n\n\n> Q5. Three typos.\n\nA5. Thanks for noticing the typo. We'll revise them in the final version.\n\n- In Algorithm 1, it should be $g_0=\\nabla f(y_0, \\mathcal{B}_0)$ at initialization.\n- In line 126, revise \"identical\" to \"similar\".\n- In line 180, revise to \"we can trivially allow\".", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Plz7X9IT91", "reviewer": "Reviewer_9hZ5", "summary": "This paper studies the efficiency problem of sharpness-aware minimization (SAM) algorithms.\nSAM requires two gradient calculations: one for computing the perturbation and another for computing the update direction.\nHence, SAM doubles the computation cost compared with ERM.\nFurthermore, these two gradients cannot be parallelized.\nThe authors introduce an auxiliary sequence and use the gradient on the auxiliary variable to approximate the gradient in computing the perturbation. \nExperiments are conducted to evaluate the effectiveness of the proposed SAMPa algorithm.", "questions": "see above weaknesses.", "limitations": "no", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- improving the efficiency of SAM is an important problem\n- moreover, making the two gradient computations parallelized is crucial to accelerate SAM algorithms\n- introducing an auxiliary sequence to approximate the gradient in computing perturbation is interesting\n- the writing is clear and easy to follow.", "weaknesses": "- Though introducing an auxiliary sequence is an interesting idea to parallelize the two gradient calculations,\nI have a concern about the auxiliary sequence:\n    -  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.\n\n- The theoretical analysis in Section 4 requires the loss to be convex.\nHowever, losses are usually non-convex in practice (for example, experiments in the paper).\nThus, the analysis in section 4 is meaningless in practice.\n\n- Main results (Tables 1 and 2) show that SAMPa outperforms SAM in all settings.\nThis observation is odd to me.\nNote that SAMPa uses an approximate method to compute the gradient but achieves better performance than SAM which uses the exact gradient.\nSome explanations are needed here.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback and address all remaining concerns below:\n\n> Q1. Though introducing an auxiliary sequence is an interesting idea to parallelize the two gradient calculations, I have a concern about the auxiliary sequence:\nEstimate the approximate error $\\nabla f(x_t, \\mathcal B_{t+1}) - \\nabla f(y_{t+1}, \\mathcal B_{t+1})$. Establishing a bound for this error is crucial to the proposed algorithm. Otherwise, the method has no theoretical support.\n\nA1. In the proof we do not actually need a bound on $||\\nabla f(x_t) - \\nabla f(y_{t})||^2$, since we can telescope the quantity and entirely remove it from the resulting rate by choice of the initialization (see l. 165-168). The only thing we need is that the quantity eventually becomes small, which Lemma 4.3 guarantees in the convex case by establishing decrease of the potential function $\\mathcal V_t$.\n\nIn practice we observe that the difference between $\\nabla f(x_t, \\mathcal B_t)$ and $\\nabla f(y_t, \\mathcal B_t)$ is in fact small. Specifically, we track the cosine similarity and Euclidean distance between $\\nabla f(x_t, \\mathcal B_t)$ and $\\nabla f(y_t, \\mathcal B_t)$ during the whole training process of Resnet-56 on CIFAR-10 [here](https://imgur.com/a/6hWQm44). We find that the cosine similarity keeps above 0.99 during the whole training process, and in most periods it is around 0.998, while at the end of training it is even close to 1. This indicates that SAMPa's approximated perturbation $\\rho \\frac{\\nabla f(y_t)}{||\\nabla f(y_t)||}$ is an excellent approximation of SAM's perturbation $\\rho \\frac{\\nabla f(x_t)}{||\\nabla f(x_t)||}$.\n\nMoreover, the Euclidean distance decreases and is close to zero at the end of training. This matches our above theoretical analysis that $||\\nabla f(x_t) -\\nabla f(y_t)||^2$ eventually becomes small.\n\n> Q2. 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.\n\nA2. We kindly disagree that convex analysis is \"meaningless\" in guiding nonconvex practice. Arguably most of the successful optimizers today for deep learning were developed under the assumption of convexity. This includes AdaGrad, Adam and even more recent developments such as Distributed Shampoo and the Scheduler-free optimizer that won the AlgoPerf competition this month [1].\n\nSimilarly, we found convex analysis extremely helpful in developing the algorithm. In our particular case, our algorithmic design _falls out of the analysis_, specifically leading to our choice of the auxilliary sequence $(y_t)_{t\\in \\mathbb N}$ (see Section 3).\n\nWe subsequently show that SAMPa indeed has strong practical performance in deep learning applications (see Section 5). As mentioned in Q1, we additionally provide evidence that $y_t$ indeed acts as a good estimator for $x_t$ in the experiments, which builds further confidence in the approach.\n\n> Q3. 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.\n\nA3. There is no conflict between the two observations:\n\n- Our theoretical claim arguing through closeness to SAM is concerned with SAMPa-0 (referred to as simply SAMPa in line 128). SAM and SAMPa-0 have very similar performance in practice, which in fact _confirms_ the theory (see Table 1, 2, 6 and 7).\n- The method that consistently beats SAM is SAMPa-0.2, which is simply a convex combination of SAMPa and OptGD (see l. 138 for details). It is an interesting direction to understand this superior performance better. One hypothesis for the superior performance is that optimistic gradient descent might be better suited for nonsmooth parts of the optimization landscape and thus stabilize the algorithm when Lipschitz continuity assumption is not present.\n\nWe would like to highlight again that SAMPa-0.2 _consistently_ outperforms SAM with 2$\\times$ speedup across models and datasets (see Section 5). It can also be run on a single device with the same runtime as SAM, while _still delivering the same superior performance_. We will remark on this in the final version.\n\n\n[1] https://mlcommons.org/2024/08/mlc-algoperf-benchmark-competition/", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper presents SAMPa, a parallelized variant of Sharpness-Aware Minimization (SAM) aimed at improving the efficiency and generalization of deep learning models. The key innovation is the introduction of an auxiliary sequence to parallelize the two gradient computations in SAM, which traditionally incurs double the per-iteration computational cost compared to standard optimizers like SGD. 
The empirical results are strong, and the method shows significant potential for improving the efficiency of SAM. Though this paper does not provide any theoretical explanation on why SAMPa can have better generalization, SAMPa is a practically valuable contribution to the field of deep learning optimization. The reviewers still unanimously recommends accepting this paper. The authors should include the additional experiments in the rebuttal into the final version, especially those related to mSAM and recent other baseline methods.

Minor comments:
The main theorem (Lemma 4.3 and Theorem 4.4) has a strange coefficient $(1+L+L^2)$. It is strange because $L$ is not unit-free and should not be added directly across different powers. By a quick glance it seems to be caused by typos in proofs, e.g., the third line in equation (6), $(1-\eta_t)$ should be $\eta_t$.
Decision: Accept (poster)