K-SAM: Sharpness-Aware Minimization at the Speed of SGD
Authors: Renkun Ni
Ping-yeh Chiang
Jonas Geiping
Micah Goldblum
Andrew Gordon Wilson
Tom Goldstein
Venue: Submitted to ICLR 2023 (2023)
Keywords: deep learning
efficient training

Abstract: Sharpness-Aware Minimization (SAM) has recently emerged as a robust technique for improving the accuracy of deep neural networks. However, SAM incurs a high computational cost in practice, requiring up to twice as much computation as vanilla SGD. The computational challenge posed by SAM arises because each iteration requires both ascent and descent steps and thus double the gradient computations. To address this challenge, we propose to compute gradients in both stages of SAM on only the top-k samples with highest loss. K-SAM is simple and extremely easy-to-implement while providing significant generalization boosts over vanilla SGD at little to no additional cost.

Under review as a conference paper at ICLR 2023
K-SAM: SHARPNESS-AWARE M AT THE SPEED OF SGD
Anonymous authors Paper under double-blind review
ABSTRACT
Sharpness-Aware Minimization (SAM) has recently emerged as a robust tech- nique for improving the accuracy of deep neural networks. However, SAM incurs a high computational cost in practice, requiring up to twice as much computa- tion as vanilla SGD. The computational challenge posed by SAM arises because each iteration requires both ascent and descent steps and thus double the gradi- ent computations. To address this challenge, we propose to compute gradients in both stages of SAM only the top-k samples with highest loss. K-SAM is sim- ple and extremely easy-to-implen to-implement while providing significant generalization boosts over vanilla SGD at little no additional cost.
INTRODUCTION
Methods for promoting good generalization are of tremendous value to deep learning practitioners and a point of fascination for deep learning theorists. While machine learning models can easily achieve perfect training accuracy will generalize well to test data. Improving the ability of models to extrapolate their knowledge, learned during training, to held out test samples is the key to performing well in the wild.
any dataset given enough parameters, it is unclear when a model
Recently, there has been a line of work that argues for the geometry of the loss landscape as a major contributor to generalization performance for deep learning models. A number of researchers have argued that flatter minima lead to models that generalize better (Keskar et al., 2016 Xing et al.) 2018 Jiang al. 2019, Smith al. 202I). Underlying this work the intuition that small changes in parameters yield perturbations to decision boundaries so that flat minima yield wide-margin deci- sion boundaries (Huang et al. 2020). Motivated by these investigations, Foret et al. (2020) propose an effective algorithm Sharpness Aware Minimization (SAM) to optimize models toward flatter minima and better generalization performance. The proposed algorithm entails performing one-step adversarial training in parameter space, finding loss function minimum that is "flat" in the sense that perturbations to the network parameters in worst-case directions still yield low training loss. This simple concept achieves impressive performance
wide variety of tasks. For example, Foret ct al. (2020) achieved notable improvements on various benchmark vision datasets (c.g., CIFAR-10 -00, IeageNet) by simply swapping out the optimizer. Later, Chen al. (2021) found that SAM im- proves sample complexity and performance of vision transformer models so that these transformers are competitive with ResNets even without pre-training. Moreover, further innovations the SAM setup, such as modifying the radius of the adversarial step to be invariant to parameter re-scaling (Kwon al. 202T), yield additional improvements to generalization on vision tasks. In addition, Bahri et al. (2021) recently found that SAM not only works in the vision domain, but also improves the performance of language models on GLUE SuperGLUE, Web Questions, Natural Questions,
Trivia QA, and TyDiQA (Wang et al. 2018 2019 Joshi et al. 2017: Clark et al.) 2020).
Despite the simplicity of SAM, the improved performance comes with a steep cost of twice as much compute, given that SAM requires two forward and backward passes for each optimization step: one for the ascent step and another for the descent step. The additional cost may make SAM too expen- sive for widespread adoption by practitioners, thus motivating studies to decrease the computational cost of SAM. For example, Efficient SAM (Du et al., 2021) decreases the computational cost of SAM by using examples with the largest increase in losses for the descent step. Bahri et al. (2021) randomly select a subset of examples for the ascent step making the ascent step faster. However, both methods still require full forward and backward pass for either the ascent descent step on

---

Under review as a conference paper at ICLR 2023
all samples in the batch, so they more expensive than vanilla SGD. In comparison these works, we develop a version of SAM that is as fast as SGD so that practitioners can adopt our variant at no cost.
In this paper, we propose K-SAM, a simple modification to SAM that reduces the computational costs that of vanilla SGD, while achieving comparable performance to the original SAM op- timizer. K-SAM exploits the fact that a small subset of training examples is sufficient for both gradient computation steps (Du et al. 2021] Fan et al.! 2017; Bahri et al. 2021), and the examples with largest losses dominate the average gradient over a large batch. To decrease computational cost, we only use the K examples with the largest loss values in the training batch for both gradient computations. When K is chosen properly, our proposed K-SAM can be as fast as vanilla SGD and meanwhile improves generalization by seeking flat minima similarly to SAM.
We empirically verify the effectiveness of our proposed approach across datasets and models. We demonstrate that small number of samples with high loss produces gradients that are well-aligned with the entire batch loss. Moreover, we show that our proposed method can achieve comparable performance to the original (non-accelerated) SAM for vision tasks, such as image classification on CIFAR-(10, (10, 100), and language models the GLUE benchmark while keeping the training cost roughly as low as vanilla SGD. On the other hand, we observe that on large-scale many-class image classification tasks, such as ImageNet, the average gradient within batch is broadly distributed, so that a very small number of samples with highest losses are not representative of the batch. This phenomenon makes it hard to simultaneously achieve training efficiency and strong generalization by subsampling, Nonetheless, we show that K-SAM can achieve comparable generalization to SAM with around 65% training cost.
2 RELATED WORK
2.1 SHARPNESS-AWARE MINIMIZATION
In this section, we briefly introduce how SAM simultaneously minimizes the loss while also de- creasing loss sharpness, and we detail why it requires additional gradient steps that in turn double the computational costs during training. Instead of finding parameters yielding low training loss only, SAM attempts to find the parameter vector whose neighborhood possesses uniformly low loss, thus leading to flat minimum. Formally, SAM achieves this goal by solving the mini-max opti- mization problem,
min LAAM LSA(ww) (w), where (SA(w)) = max Ls(w + +e,, Se
(1)
where are the parameters of the neural network, L is the loss function, S is the training set, and is a small perturbation within an l2 ball of norm p. In order to solve the outer minimization problem, SAM applics a first-order approximation to solve the inner maximization problem.
arg max Ls + c),
arg max Ls (w) + ef Lsws(w).
pV Lsww / Ls())ll2.
(2)
After computing the approximate maximizer é, SAM obtains "sharpness aware" gradient for descent:
VLSAM (w) x v.,Ls(w))ltt
(3)
In short, given a base optimizer, i.e., SGD, instead of computing the gradient of the model at the current parameter vector w, SAM updates model parameters using the gradient with respect to the perturbed model at w + è. Therefore a SAM update step requires two forward and backward passes on each sample in a batch, namely a gradient ascent step to achieve the perturbation é and gra- dient descent step to update the current model, which doubles the train time compared to the base optimizer.
2.2 EFFICIENT SHARPNESS-AWAR MINIMIZATION
Recently, several works improve the efficiency of SAM while retaining its performance benefits. Bahri al. (2021) improve the efficiency of SAM by reducing the computational cost of the gradient
2

---

Under review as a conference paper at ICLR 2023
ascent step. Instead of computing the ascent gradient the whole mini-batch they estimate the gradient using a random subset. This tweak makes the computational cost only about 25% slower than the vanilla SGD training routine when using 1/4 of the training batch for the ascent gradient, while maintaining performance comparable to SAM. However, reducing the computational cost of only the ascent step cannot make SAM as fast or faster than the hase optimizer. In addition, we show in Section 4 that using random subset of the mini-batch for both gradient computations in SAM has a negative impact on performance, which suggests that we should seek a better selection method than random selection.
Du et al. (202I) propose an efficient SAM (ESAM) algorithm employing two strategies, namely stochastic weight perturbation and sharpness senssiiv--eesitive data selection. Stochastic weight perturbation updates only part of the weights during the gradient ascent step which offers limited speed-ups. During the descent step, the gradient is only computed using the examples whose loss values increase the most after the parameter perturbation. With these two strategies, ESAM achieves up to 40.3% acceleration compared to SAM. However, since ESAM selects the examples with the largest loss differences, it has to compute the gradient over all samples in the batch for the ascent step and then must similarly perform forward pass over all samples in the batch for the descent step, which limits the possible speed-ups from this approach.
2.3 TOP-K K OPTIMIZATION
Top-k optimization has been applied to vanilla SGD Fan et al. 2017 Kawaguchi & Lu 2020) Given training mini-batch, Ordered SGD selects the K examples with the largest losses on which to perform the minimization and computes subgradient based only on these examples. This work theoretically shows that on convex loss functions, Ordered-SGD is guaranteed to converge sub- linearly to a global optimum and to critical point with weakly convex losses. In addition, they empirically show that Ordered-SGD can achieve comparable results to vanilla SGD multiple ma- chine learning models such as SVM and deep neural networks. In this paper, we show that top-k optimization can be effectively applied to both gradient computation steps in SAM and is actually essential to achieving good performance when is small.
3 K-SAM: EFFICIENT SHARPNESS- MINIMIZATION BY SUBSAMPLING
We now introduce proposed method, K-SAM, where we select K1, K2 examples with the largest loss values to estimate the gradients in the ascent and descent steps of a SAM update, respectively. When Kj and K₂ are small, the computational complexity of SAM will be vastly reduced since large proportion of the compute required for neural network training is concentrated in gradient computations.
Recall that in the Section 2.1 given a mini-batch of examples, B, SAM first computes the ascent gradient with respect to the current parameters w find the perturbation by equation 2 Then, the final gradient descent direction is formed by equation 3 In order to decrease the computational cost, we approximate both gradients by using subsets M M of the mini-batch B for gradients calculation, where each subset contains training examples with the largest K1, K2 loss values, re- spectively. Formally, given the losses, 1 = Ls(w). of batch B and its index set I {1, 2, [ of the samples, the subsets can be generated as following,
M(KI,K.) = ={ (0,,, v:) E B: Qp.2), where Q(1.2) = arg max
Σdj. QCLıQI[(K K₂) FEQ
Given these subsets, we can efficiently obtain the ascent gradients, and then achieve adversarial perturbation in parameter space by
pV LMK, (w)/ LMn, (w)ll2-
Finally, we get the actual descent gradient GSAM of SAM via
gSAM = V.LLMK, (w) w
The detailed formulation can be found in Algorithm 1
3

---

Under review as a conference paper at ICLR 2023
Algorithm K-SAM
1: Input: training set S, loss function L, batch size b, ascent subset size K1 descent subset size K2, neighborhood size p. model parameter w, learning rate v,, iterations
2: for - to 7 do
3: Sample mini-batch B E S with size
4: Get the loss values for the mini-batch: LB = LB(W).
5: Generate ascent batch MK1 and descent batch MMx with largest K1, K2 losses respec-
tively.
V.LMa (w)
6:
Compute perturbation for ascent step : PTV=,LM,, (w)]l2
7:
Compute descent gradient of SAM: 9SAM VuLiMNs (w)lw+e. Update parameters: w 9SAM.
8:
3.1 GRADIENTS OF K-SAM ARE ACCURATE AND EFFICIENT APPROXIMATIONS
K-SAM is motivated by the observation that gradients of both the ascent and descent steps in SAM can be accurately approximated by gradients calculated from the subsets with the largest & losses (top-k losses). The intuition underlying this phenomenon is that the gradient data is an average of the individual per-sample gradients, and this average dominated by a small cheaply select samples with high
mini-batch of
number of terms with large gradient magnitudes. Since loss, these will serve as effective surrogates. We will investigate both key components, ascent and descent gradient approximation, in detail to validate this hypothesis.
3.1.1 GRADIENT APPROXIMATION ASCENT STEPS
We first analyze how accurately the gradients, Vw LM,, (w), of the subset with top-k losses ap- proximate the gradients, V Lww in ascent steps of SAM, by measuring the cosine similarity of the two gradients during training. We train WideResNet-Iee 16-8 (Zagoruyko & Komodakis) 2016) with SAM and batch size 128 on CIFAR-10. In the left plot of Figure we observe that the align- ments between the top-k gradients and those computed on the full mini-batch the ascent steps are very high, especially at the beginning and the middle of the training. For example, we see cosine similarity of around 0.9 when & 64. In addition, compared to gradients computed on randomly chosen subset of & examples (random-k) (Bahri et al. 2021). we find that top-k gradients ascent steps far better. When we use the same
approximate the gradients of the full mini-batch number of examples k; top-32 is almost twice well aligned with the full batch ascent gradients as random-32. Even when we only use 16 examples with largest losses, we are still able to estimate the ascent gradients up 50% more accurately compared to random-32. This difference between the two approximations is especially prominent during the central parts of training. On the other hand, if we want to estimate the full mini-hatch gradients as accurately top-16, would have to use times as many examples with random selection, as visualized in the middle plot of Figure [] In summary, top-k much more efficient and accurate for estimating gradients in the ascent steps compared to random subsets.
Alggement
Alfgnement
06
Top-l
Tip-18
Top
Zand 32
Top-26
Hand-15
S01
Hand 222
Hand-og
K-SAM
Epushs
Epocles
Figure 1: (Left) Gradient alignment between full mini-batch and subsets of largest K₁ (16,32 32, 64) losses. (Middle) Gradient Aligument between full mini-batch and subsets of random K, = {16. 32, 64) losses. (Right) Gradient Alignment between SAM updates and updates of SGD and K-SAM.
4

---

Under review as a conference paper at ICLR 2023
Table 1: Using subset update the model with SGD on CIFAR-IO and CIFAR-100. We find that using top-k subsets yields similar performance to SGD, while using random-k subsets exhibits a clear performance drop.
Model
Accuracy(%)
R8
CIFAR-10
CIFAR-100
SGD
96.29 0.09
79.08 + 0.18
top-64
96.26 + 0.14
79.04 + 0.13
random-64
95.92 + 0.06
78.53 + 0.20
Wide-28-10
CIFAR-10
CIFAR-100
SGD
96.99 + 0.09
82.05 + 0.15
top-64
97.02 0.14
82.29 = 0.22
random-64
96.66 + 0.09
82.31 = 0.17
3.1.2 GRADIENT APPROXIMATION IN DESCENT STEPS
Now that we can reasonably approximate the gradients in ascent steps with fewer samples, we analyze whether not
estimate the gradients in descent steps accurately with similar strategy. We show in Table T that comparable performance can be achieved when we compute the top-k gradients and perform descent using this approximation vanilla SGD (rather than SAM) for a range of settings. Using top-k descent gradients is thus a sensible strategy, even for the updates of deep neural networks, an extension of theoretical investigations of convex optimization Fan et al. (217)) and Kawaguchi & Lu (2020). We conduct these experiments on both CIFAR-10 and CIFAR- 100 datasets with ResNet-18 and WideResNet-266 10 models. the forward pass, we feed the full mini-batch of size 128 into the network. In the backward pass, we compute the gradients with 1) the full training batch (SGD), 2) half of the training batch with the largest loss values (top-64) and 3) half of the training batch with random sampling (random-64). As shown in Table [ training with only the 64 samples with the largest losses yields performance comparable, and sometimes even superior, to vanilla SGD on both network architectures and datasets. However, when we do gradient descent based on a random subset of 64 examples, we see a clear accuracy drop in most experiments. The accuracy drop suggests that using top-k subsets is a better approximation of the descent gradients of the full mini-batch.
In addition, we measure the cosine similarity between the gradients in descent steps of K-SAM and SAM, where the same ascent gradients are used for both methods. In the right plot of Figure I we observe that throughout training. K-SAM's estimated descent gradients are very similar to SAM's descent gradients with a cosine similarity of 0.9. We also show the similarity between the gradients in descent steps of SGD and SAM for comparison, where SGD's descent gradient is quite similar to SAM in the first few epochs but drops dramatically late in training. This analysis reveals that the K-SAM's approximation of SAM systemically points in the intended direction and toward flatter minima.
3.2 EFFICIENCY GAINS FROM K-SAM
We implement K-SAM in the following way: we first compute forward pass on the full mini- batch, and then perform forward and backward propagation on both subsets Mr, and MKg. Theoretically, if we want K-SAM to produce exactly the same computational cost as vanilla SGD (ignoring the cost of loading data into memory). we have to choose K2 such that K₁ + K₂ B/4, where B is the batch size. However, In practice, due to the fact that we do not need to compute the gradient for the first forward pass, we can benefit from significant speed-up and memory savings, as analyzed in T In addition, this first forward pass can be further accelerated by computing it with lower precision, given that only a ranking of the output values is necessary. For example, on a conventional NVIDIA 2080t i card, forward pass in mixed lower precision and inference mode (26ms) can be twice as fast as an (already fast) normal forward pass (60ms) as described in Table [ As result, we do not need to select K₃ and K₂ as small as fou-tourth of the batch size. Indeed, empirically find that when we set K; = B/8 and K₂ = B/2, the proposed K-SAM is nearly as fast vanilla SGD (See 4.2) In addition, can further increase the efficiency of K-SAM by gradually
5

---

Under review as a conference paper at ICLR 2023
Table 2: Time comparison between full brckkrd pass, regular forward pass, forward pass in inference mode with mixed precision, and a forward pass on a subset of the batch. The model a ResNet-50 and batch size is 48. All numbers are based on the average of 250 measurements on NVIDIA 2080ti card each after CUDA synchronization. Note that larger batch sizes have more potential for efficiency gains on smaller reductions here, we focus on a conventional device and scenario where we find that the 16-fold reduction only halves the runtime compared to the 4-lold reduction, due to GPU undersaturation however, the proposed approach still succeeds in reaching SGD runtimes.
Operation
Time (ms)
Regular forward pass
56.81
Forward and backward pass
189.92
Mixed stixed-precision forward pass
25.93
16-fold reduced forward- hrward-backward pass
26.16
4-fold reduced forward-hackwand pass
55.31
decreasing K2 during training, thus leading to an algorithm that often runs even faster than vanilla SGD conventional hardware. We include PyTorch-like pseudo-code in Appendix 2 where we summarize all steps discussed so far in detail.
EXPERIMENTS
This section contains experimental evaluations of the proposed K-SAM optimization scheme. We all experiments using conventional widely available GPUs and evaluate realistic PyTorch code without additional customized modifications for speed. We note that although we could make the first forward pass faster by using mixed- -eeeedprecision arithmetic as described in Section 3.2, we keep the precision unchanged for all experiments for fair comparisons. We evaluate the effectiveness of K- SAM on image classification tasks CIFAR-10, CIFAR-100, and ImageNet, and on the General Lan- guage Understanding Evaluation (GLUE) benchmark. In addition, ablation studies are conducted to further show the robust effectiveness of our proposed method under hyperparameter selections.
4.1 BENCHMARK EVALUATIONS
CIFAR-10 and CIFAR-100. We first show that K-SAM can achieve comparable performance to SAM on both CIFAR-10 and CIFAR-100 but at a comparable runtime to vanilla SGD. We eval- uate the proposed K-SAM for three different network architectures: ResNet-18 18 (He et al. 2016). WideResNNet-28-10 10 (Zagoruyko
Komodakis 2016) and PyramidNet-1 10 (Han et al., 2017). Un- less otherwise stated, use the Inception -teppe (Szegedy et al. 2016) image pre- pre-processing and Cutout Vrevs & Taylor 2017) as default data augmentations. For each experiment, we evaluate over 5 trials and report the
well as confidence intervals with width of one standard error. In Table 3, we compare test classification accuracy as well as training time among optimizers: SGD, SAM, ESAM, and our proposed K-SAM. Within the same architecture, we use the same hyper- parameters for all optimizers. For ResNet-18 and WideResNett28-110 we train 200 epochs with a cosine learning rate scheduler and peak learning rate 0.05. For PyramidNet-1] 10, we train for 300 epochs with cosine learning rate scheduler and peak learning rate 0.1. We use perturbation ra- dius of p 0.05 for ResNet-18 and p 0.1 for WideResNet-29 10 and PyramidNet-110 Unless otherwise stated, we train all models with the same batch size 128 on a single GPU.
Starting from the default batch size of 128, we set K₁ 16, K₂ = 64 for K-SAM. To accelerate training, we further decrease K₂ by half midway through training, denoted by K-SAM". As can be seen in T the test accuracy of K-SAM comparable to SAM while the training speed remains almost at the level of SGD. On average, without K-SAM", we add ~ 6% runtime top of vanilla SGD, which is significantly less than ESAM which is ~ 30% slower than SGD. When we switch to K-SAM", we can achieve test accuracy par with SAM at even faster speeds than SGD in 5 out of 6 experiments. For example, for Wett 10 trained on CIFAR- 10, K-SAM* is ~ 10% faster than SGD while obtaining comparable accuracy SAM (within 0.04%). In addition, since

---

Under review as a conference paper at ICLR 2023
Table Classification accuracy and training efficiency on the CIFAR-I0 and CIFAR-100 datasets. The efficiency is evaluated by total training time in hours and percentage of GPU memory allocation during training. Classification accuracy mean is calculated from runs. * denotes the faster version of proposed K-SAM, where K2 gradually decays during training. K-SAM can achieve comparable performance as SAM while adding little computational cost on the base optimizer, i.e., SGD.
CIFAR-10
CIFAR-100
Method
Backbone
Accuracy (%)
Train Time(h)
Memory($6 (+)
Accuracy (6)
Train Timeth)
Memory( (6)
SGD
ResNet-18
96.29 + 0.09
1.07
21.50
79.08 :: 0.18
1.28
21.20
SAM
ResNect-18
0.13
2.03
20.85
80.08
2.04
20.85
ESAM
ResNet-18
96.55 + 0.13
1.48
21.92
79.34 + 0.24
1.59
20.75
K-SAM
ResNet-18
96.54 + 0.14
1.13
17.48
79.24 = 0.29
1.18
17.48
K-SAM
ResNet-18
96.47 = 0.05
17.48
79.00
17.48
SGD
Wide-28-10
96.99 + 0.09
6.19
43.52
82.05 = 0.15
5.96
45.05
SAM
Wide-28-10
97.45 = 0.07
12.02
51.45
84.16 0.16
11.44
47.90
ESAM
Wide-28-10
97.29 + 0.06
8.15
38.25
84.21 + 0.14
7.98
38.28
K-SAM
Wide-28-10
97.45 + 0.07
6,28
36.48
84.01 = 0.29
6.21
36.50
K-SAM
Wide-28-10
0.05
PyrmidNeet-l 110
97.07 0.11
83.45
11.22
SAM
PyramiNNet- 110
97.82 + 0.11
22.50
70.56
85.95 0.21
23.18
70.57
ESAM
PyramidNet- 110
97.71 + 0.07
19.65
64.68
85.41 = 0.32
19.78
65.07
K-SAM
PyramidNet-1
97.60 + 0.06
12.56
41.20
85.38 0.18
12.82
41.25
K-SAM
PyramidNet-110
97.62 0.10
11.53
40.03
84.60 0.22
11.90
40.07
we only compute forward passes on the full mini-batch and never compute a full backward pass, which is restricted merely the subsets , the memory required by K-SAM is actually smaller than both SGD and SAM. Table 3 shows the memory allocation a single GPU during training. On average, K-SAM saves around 20% of GPU memory compared to vanilla SGD.
ImageNet To verify our proposed method on ImageNet, we train ResNet-50 models for 90 epochs with a cosine annealing learning rate scheduler and peak learning rate 0.05. For both SAM and K-SAM, we set perturbation radius = 0.05 and train all the models with the same batch size 512 on 8 GPUs parallel (64 images per GPU). We vary the size of both subsets Mr, and MM, from (B/8, B B/2, B}. where B is the batch size. Contrary to our observations on CIFAR-10 and CIFAR-100, here we observe a clear trade-off between the size K1, K2 the subsets and the test accuracy. In Appendix A (sce Table 7 and Table 8). severely when we have fewer samples to estimate the gradients, especially for the descent steps. To understand this phenomenon,
see that the test accuracy drops
measure the gradient alignments for ImageNet models as well. In Figure 2 we show that in general, ImageNet models have worse alignment between gradients of SAM and K-SAM, especially when the size of subsets is small, i.e., K B/8.
Although we do not achieve the same speed as SGD while enjoying the same generalization power SAM, show in Table 4 that of its training cost. In addition, we can achieve comparable performance to SAM (0.1% 1% accuracy drop), but with only 65% of its training cost.
achieve better results than SAM yet with only 80%
Paad-334
Alganzoo
Opack
Figure 2: Gradient alignments between full mini-batch and subsets of Topp RBnd B/4, B/8} losses for ImageNet, where B is the batch size. In general, the alignments are small, especially when K = B/8. Again, gradients calculated on samples with top-k losses align better compared to random sampling.

---

Under review as a conference paper at ICLR 2023
Table 4: Performance of SGD, SAM and K-SAM on ImageNet. K-SAM achieves comparable testing accuracy to SAM, while reducing the computational cost greatly.
ResNet-50
Accuracy(%)
Images/s
ResNet-50
Accuracy(%)
Images/s
SGD
75.98
1664.57
K-256/360
77.08
1192.67
SAM
77.19
775.80
K-256/512
77.28
935.86
GLUE benchmark In this section, test the effectiveness of K-SAM the GLUE benchmark. We fine-tune the BERT-BASE pre-trained model (Devlin et al.) 2018) on the NLP tasks with a batch size of 56 for 250k steps. We use perturbation radius p 0.02 for all experiments. For our proposed K-SAM, we set K₃ and K₂ to 8 and 32, respectively. We find that on most tasks, SAM improves the performance of the model comparing to SGD, and K-SAM yields similar improvement to SAM. When excluding the rte dataset, a special case in which SGD significantly outperforms SAM and its variants, SAM and K-SAM improve the performance of SGD by 0.25% on average (see Table 5). Notably, our observations concerning the rte dataset are consistent with observations in (Bahri al.] 2021). where the performances of both SAM and K-SAM are very poor, and the performance gap dominates the average over datasets.
Table 5: Performance of SAM and K-SAM on the GLUE benchmark. Comparing to the base opti- mizer SGD, K-SAM yields similar improvement to SAM on most tasks, except on the rte dataset.
Overall
Overall
(exclude rte)
ss12
mppc
qqp
stsh
mnli
quli
SGD
83.51
85.12
61.40
92.20
88.24
91.31
91.27
88.20
88.
87.85
83.62
91.14
72.20
SAM K-SAM
83,18
85,37
61.36
92.66
87.50
91.75
91.53
88.60
88.68
88.29
84.08
91.32
67.87
83.15
85.38
63.06
92.32
86.03
90.36
91.21
88.
89,80
89,51
83.35
91.38
67.51
4.2 EFFECTIVENESS OF K-SAM ON DIFFERENT SETTINGS
If selecting samples with high loss is to be useful, the gradients of these samples must be more representative of the full batch than a random subset. And moreover, we need to know exactly how big our selected subsets must be. We now answer these questions by performing series of ablations, and we additionally put K-SAM to the test in the distributed training to show that the same efficiency benefits can be achieved in such setting.
4.2.1 K-SAM PERFORMS BETTER THAN RANDOMLY SAMPLED SUBSETS
Bahri et al. 2021 suggest that SAM be made more efficient by randomly sampling 1/4 of the batch for the computation of the model perturbation. To compare this variant with K-SAM, we train ResNet- 18 models both CIFAR-10 and CIFAR-100 with both optimizers and batch size 128. We evaluate both methods across various K2 and K2. We see Table random subsets can achieve comparable results to K-SAM only when either K₁ or K₂ is large. However, when we use smaller K1, K2. that is one-eighth and half of the batch size, accuracy noticeably drops when using the random subsets for computing ascent directions, which prevents further subsampling acceleration for descent steps.
4.2.2 DIFFERENT SIZES OF TOP-K SURSETS
The size of both subsets K2 and K2 offers trade-off between training speed and the fidelity with which we approximate SAM. To evaluate the effect of K1, K2, we train K-SAM with both
ResNet-18 18 and WideResNet-28-100 10 on CIFAR-10 and 100 with various K1, K2 from set {16, 64, 128). where batch size equals to 128. We provide the results in Appendix C where Table shows the performance and efficiency across combinations of K₁, K2. Interestingly, we usually achieve the highest classification accuracy (even higher than SAM on WideResNett28-110 10) when K1 K₂ 64, which offers 30% speed-up compared SAM.
8

---

Under review as a conference paper at ICLR 2023
Table 6: Comparison between sampling methods for subsets MM, used for parameter perturba- tions (ascent). "Rand-Ky K,/K," denotes random sampling and "K-K,/K²" denotes K-SAM. K-SAM outperforms random sampling when both 1 and K₂ are small.
CIFAR-I0
CIFAR-108 100
ResNet-18
Accuracy(%)
Time(h)
Accuracy(%)
Time(h)
SGD
96.29 + 0.09
1.07
79.08 + 0.18
1.28
Rand- 16/1-
96.41 + 0.09
1.35
79.35 + 0.22
1.28
Rand-16/64
96.22 + 0.15
1.15
79.01 = 0.31
1.15
Rund-32/128
96.59 + 0.07
1.39
79.53 0.38
1.28
Rand-32/64
96.34 + 0.15
1.19
79.49 + 0.24
1.16
K-16/128
96.58 0.12
1.53
79.17 + 0.32
1.62
K-16/64
96.54
1.13
79.24 0.29
1.18
K-32/128
96.61 + 0.09
1.64
79.55 0.20
1.64
K-32/64
96.61 + 0.03
1.17
79.45 + 0.25
1.17
4.2.3 DISTRIBUTED LEARNING
In this section, we verify that the effectiveness and efficiency of K-SAM will hold in the distributed setting as well when the batch size larger and multiple GPUs are used. In the distributed setting, the ascent gradients are calculated by the training examples on each card as described in Foret et al. (2020). The gradients are only aggregated after the descent gradients are calculated. We train the WideResNet-16-8 model with different optimizers, using batch size 512 on 4 GPUs. The results are provided in Appendix (Table TO). We find that K-SAM can still achieve comparable performance and great efficiency improvements in this distributed setting. For some combination of K₃ and K2, K-SAM even outperforms the baseline SAM but with significantly lower computational cost. On CIFAR-100, for example, the runtime of vanilla SGD.
obtain performance greater than that of SAM within minute of
CONCLUSION
In this paper, we explore mechanism for improving the training efficiency Sharpness- Minimization (SAM). We find that simple top-k selection strategy for both ascent and descent steps is a succinet and effective way for accelerating SAM training. We validate the underlying intuition that gradients calculated on training examples with the largest several losses are good ap- proximations for the gradients both steps SAM by measuring the cosine similarity between the gradients in SAM and their approximations. Compared to random subsampling, we show that gradients calculated on top-k samples align better. In addition, we empirically evaluate our method, K-SAM, on multiple vision and language benchmarks, and show that method can achieve com- parable performance to SAM while reducing the computational cost significantly. Although K-SAM can accelerate the sharpness-aware aware training, it still needs an additional forward pass on the full mini- batch. How to avoid this pass or even avoid the ascent step in SAM will be an interesting future direction.
REPRODUCIBILITY STATEMENT
We ran all the experiments with specific random seeds and multiple runs. We provide the source code the supplementary material, and all the results are reproducible.
REFERENCES
Dara Bahri, Hossein Mobahi, and Yi Tay. Sharpness-awar aware minimization improves language model generalization. arXiv preprint arXiv:21 2021.
Xiangning Chen, Cho-Jui Hsich, and Boqing Gong. When vision transformers outperform resnets without pre-training strong data augmentations. arXiv preprint arXiv:2106.0155 2021.

---

Under review as a conference paper at ICLR 2023
Jonathan H Clark, Eunsol Choi, Michael Collins, Dan Garrette, Tom Kwiatkowski, Vitaly Nikolaev, and Jennimaria Palomaki. Tydi qa: A benchmark for information-seeking question answering in typologically diverse languages. Transactions of the Association for Computational Linguistics, 8:454-470, 2020.
Jacob Devlin, Ming-Wes Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 18100000,, 2018.
Terrance De Vries and Graham W Taylor. Improved regularization of convolutional neural networks with cutout. arXiv preprint arXiv: 04555,, 2017.
Jiawei Du, Hanshu Yan, Jiashi Feng, Joey Tianyi Zhou, Liangli Zhen, Rick Siow Mong Goh, and Vincent YF Tan. Efficient sharpness-aware minimization for improved training of neural net- works, arXiv preprint arXiv:21 2021.
Yanbo Fan, Siwei Lyu, Yiming Ying, and Baogang Hu. Learning with average top-k loss. Advances in neural information processing systems, 30, 2017.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimiza- tion for efficiently improving generalization. arXiv preprint arXiv:2010.01412, 2020.
Dongyoon Han, Jiwhan Kim, and Junmo Kim. Deep pyramidal residual networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 5927-5935, 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 770-778, 2016.
W Ronny Huang, Zeyad Emam, Micah Goldblum, Liam Fowl, Justin K Terry, Furong Huang, and Tom Goldstein. Understanding generalization through visualizations. 2020.
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. arXiv preprint 1rXiv:192.002111 2019.
Mandar Joshi, Eunsol Choi, Daniel S Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. arXiv preprint arXiv: 170..iv:100.03330. 005555 2017.
Kenji Kawaguchi and Haihao Lu. Ordered sgd: A new stochastic optimization framework for em- pirical risk minimization. In International Conference on Artificial Intelligence and Statistics, pp- 669-679. PMLR, 2020.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv: 10XX:06609 2016.
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 Con- ference Machine Learning, pp. 5905-5914. PMLR, 2021.
Samuel L. Smith, Benoit Dherin, David GT Barrett, and Soham De. On the origin of implicit regu- larization in stochastic gradient descent. arXiv preprint arXhv:2101.1217, 2021.
Christian Szegedy. Vincent Vanhoucke, Sergey loffe, Jon Shlens, and Zbigniew Wojna. Rethink- ing the inception architecture for computer vision. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2818-2826. 2016.
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:1804.0746]. 2018.
Alex Wang, Yada Pruksachatkun. Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. Superglue: A stickier benchmark for general-purpost language understanding systems. Advances in neural information processing systems, 32, 2019.
10

---

Under review as a conference paper at ICLR 2023
Chen Xing, Devansh Arpit, Christos Tsirigotis, and Yoshua Bengio. A walk with sgd. arXiv preprint arXiiv1802.08770, 2018.
Sergey Zagoruyko and Nikos Komodakis. arXiv:1605.00116, 2016.
Wide residual networks
arXiv
preprint

---

Under review as a conference paper at ICLR 2023
K-SAM: Sharpness-Aware Minimization at the Speed pf SGD (Appendix)
A TRADE-OFF BETWEEN THE SIZE OF SUBSETS AND K-SAM PERFORMANCE ON IMAGENET
We provide results for different combinations of K1 and K₂ for K-SAM ImageNet with model ResNet-50. We observe a clear trade-offs between the size of subsets and the testing accuracy, where when have smaller K, especially K2, the performance addition, we show that this happens to SGD as well, where the accuracy drops more than 1 1% when only backpropagate through half of the mini-batches.
ImageNet will drop significantly. In
Table 7: K-SAM has clear trade-offs ImageNet between the size of the subsets and testing accuracy.
K1/K₂
256
360
512
64
75.37
76.43
76.55
128
75.99
76.81
76.73
256
76.08
77.09
77.29
512
76.34
76.90
77.19
Table 8: Performance drops significantly on ImageNet K₂ is small when applied to vanilla SGD.
K₂
256
360
512
SGD
74.866
75.516
75.978
12

---

Under review as a conference paper at ICLR 2023
B PYTORCH-LIKE PSEUDO CODE FOR K-SAM
The PyTorch-Like pseudo code is proved in Alg
Algorithm PyTorch-like implementation for K-SAM
Pytorch-like training loop
for inputs, targets in data.loader:
with torch.no.grad no-grad ()
outputs = model (inputs)
loss = criterion (outputs, targets)
idx kl torch. (loss, K1) # Get MK.
idx_k2 torch. (loss, K2) & Get Mr.
Get the model perturbation
outputs.k model (inputs [idx_kl])
loss kl criterion (outputs k, targets [idx k1)) loss k1. mean () ()
optimizer. fit.srsttrrstttppppe
# Get the SAM updates
outputs.k model (i 1oss_k2 criterion (outputs.k, targets (idx.k2))
loss k2. 122.mee. mean () ) opt zer second step()
C ABLATION STUDIES
Different Subset Sizes Contrary to ImageNet, K-SAM consistently outperforms SGD with different K1, K2 on CIFAR-10 and CIFAR-100. In addition, the best performance is usually achieved when Ki K2 B/2, where B is the batch size, which is almost good even better than the original SAM.
Table Classification accuracy and total training time for K-SAM across combinations of K,, K2- Best accuracy is usually achieved when both of K1, K2 equal to half of the batch size.
CIFAR- 10
CIFAR-100
ResNet-18
Accuracy($%)
Time(h)
Accuracy($3)
Time(h)
SGD
96.29 = 0.09
1.07
79.08 = 0.18
1.28
K-128/64
96.59 = 0.03
1.64
79.83 = 0.16
1.77
K-64/64
96.64 = 0.04
1.36
79.75 0.22
1.26
K-64/128
96.63 0.08
1.74
79.86 0.28
1.84
K-16/64
96.54 = 0.14
1.13
79.24 = 0.29
1.18
Wide-28-10
Accuracy(96)
Time(h)
Accuracy(%6)
Time(h)
SGD
96.91
0.07
6.19
82.05 0.15
5.96
K-128/64
97.46 = 0.10
10.84
84.39 = 0.20
10.61
K-64/64
97.46 = 0.04
8.13
84,52 = 0.15
8.36
K-64/128
97.45 = 0.05
10.64
84.17 = 0.15
10.79
K-16/64
97.45 = 0.07
6.28
84.01 = 0.29
6.21
13

---

Under review as a conference paper at ICLR 2023
Effectiveness and Efficiency under distributed learning In Table 10 tributed setting, where vanilla SGD and SAM will be faster, we can still achieve the same efficiency improvements by K-SAM. In addition, K-SAM will achieve comparable results to SAM as well.
show that in the dis-
Table 10: Top-k SAM in the distributed larger batch setting. In this table, observe that results in smaller batch setting extend to larger batch sizes and distributed computation. Experiments this table batch size of 512 and p 0.02 train WideResNet-16-8.
CIFAR-10
CIFAR-108 100
Wide-16-8
Accuracy(%)
Time (minutes)
Accuracy(%)
Time(minutes)
SGD
96.50 0.10
47
80.01
0.09
55
SAM
96.70 0.13
93
80.35 0.23
91
K-64/512
96.79 + 0.07
62
80.17 + 0.28
60
K-64/384
96.53 0.10
55
80.45 0.28
56
K-64/256
96.28 = 0.12
47
80.30 = 0.08
46
14


Reviews:
[{"review_id": "znrtlKDF6kL", "reviewer": "Reviewer_gf5z", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "This work proposes K-SAM as an alternative to SAM to improve algorithmic efficiency. The basic idea is to reduce the number of training samples used to approximate the stochastic gradients in both inner and outer optimization steps based on their loss values. This idea is originally proposed in Ordered-SGD, a biased stochastic optimization method. The authors argue that the proposed methodology can reduce the amount of computations by choosing K well without losing much on the generalization effect by the original SAM.\n", "summary_of_the_review": "Unfortunately this work does not provide non-trivial addition to algorithmic extension to the original SAM or empirical/theoretical evidence to support the realization of the proposed idea.\n\n(after rebuttal) It is my regret that my concerns are not effectively addressed by the authors. I'm still not convinced that the paper is ready for publication.", "strength_and_weaknesses": "Strength\n* clarity; the main idea of the paper is described straightforward so K-SAM is read easily.\n\nWeakness\n* very limited novelty; this work naively applies the idea of Ordered-SGD to SAM, not more really.\n* overclaiming based on shallow analysis; Figure 1 is clearly showing that the gradient estimate based on less number of samples gets more and more dissimilar to full mini-batch gradient, but the authors drive discussion into a direction where this result unfairly gives grounds for their proposed idea. This is an obvious trade-off between approximation accuracy and computational efficiency. \n* lack of theoretical grounds; the authors borrow some explanations from Ordered-SGD on the convergence analysis, but without any extension on mini-max as in the SAM setting.\n* insufficient empirical evidence; the experiments do not show any significant improvements on both performance or efficiency, except for the expected trade-off.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "tQRiXWACmYv", "reviewer": "Reviewer_wYVX", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "The authors study the effectiveness of an iterative stochastic minimization algorithm for SAM (sharpness aware minimization): at each iteration the batch is further reduced in size using a Top-K rule based on the value of the loss. The goal of this choice is to reduce the computational expense, taking the cost closer to vanilla SGD while hopefully retaining the benefits of full-batch SAM. Extensive experiments show that the approach is feasible and it is possible to keep almost the same performance of SAM while reducing the computation time to a similar level to that of SGD.", "summary_of_the_review": "The authors present an application of Ordered SGD (Kawaguchi and Lu 2020) in the context of SAM, showing through extensive and convincing experimental evidence that it allows almost a halving of the training time, while retaining almost the same accuracy numbers from the vanilla SAM method. I don't see a strong reason to reject the paper. The lack of technical novelty might be compensated by the experimental evaluation. It might be useful for practitioners in need of halving the cost of performing SAM.", "strength_and_weaknesses": "The main weakness of this work is that it lacks novelty and originality. The idea of choosing the top-k elements of the batch with highest loss is natural and has been explored before. In particular the work of Kawaguchi and Lu 2020 present the idea in its most generic form and the presented algorithm is almost a simple application of the generic method in the context of SAM. The only difference being that because SAM is a two-step process, one has to choose a subset of the batch both for the ascent step and the descent step. I wouldn't be surprised if there were even older references before 2020 exploring the exact same idea of Top-k elements in a batch. Authors should be more transparent about the fact that their algorithm is a special case of that presented in Kawaguchi and Lu 2020, applied to SAM (perhaps with some minor modifications).\n\nAnother weakness is that the authors overstate the computational cost of SAM compared to SGD. It is known that SAM has double the cost of vanilla SGD, but traditionally a constant factor (of 2 in this case) is not necessarily a reason to worry or discard a method as \"too slow\". This is reserved for methods whose complexity increases beyond constants. For example SGD which needs constant time to compute an update, vs GD which has linear complexity in the size of the data. In the third paragraph of the introduction statements like \"the additional cost may make SAM too expensive for widespread adoption\" are simply an exaggeration and should be rewritten.\n\nEven though I started stating the weaknesses, I had a positive impression of the work, mainly due to the extensive experimental evaluation. In my view, the strengths of the paper are twofold:\n1. The problem is relevant and significant to some extent: it is great to halve the time required to achieve the benefits of SAM\n2. The extensive experimental results are convincing and confirm that it is indeed possible to apply Top-K in the context of SAM without a noticeable degradation of the results. The experiments presented cover many different settings like (a) different architectures (b) different and relevant baselines (c) different tasks like supervised learning with images and NLP tasks (d) different hyperparameter settings and (e) different computational infrastructure (single node vs distributed training). Overall I find the authors did a great job in this regard.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "UrZUuidV2X1", "reviewer": "Reviewer_Y4Jo", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper proposes to sample two subsets from each batch (choosing examples with the largest loss) to perform the SAM algorithm. The authors justify their method by showing the cosine similarity between the gradients obtained before and after the subsample process. Empirical results on CIFAR-10/100, ImageNet, and BERT finetuning show the effectiveness of the proposes method.", "summary_of_the_review": "This paper has limited novelty, and I've got some questions for the authors (refer to the weaknesses part). Empirical results is not strong enough to overcome the drawbacks. So I vote for rejection for now.", "strength_and_weaknesses": "### Strengths\n- Motivation of this paper is clear as SAM usually doubles the computational cost of SGD.\n\n### Weaknesses\n- How do you obtain the K1 and K2 example specifically? Do you just randomly choose from the K1+K2 examples that have the largest loss? Will their order affect the result? The authors should elaborate more on the algorithm.\n- To me, it's surprising to see in Table 1 that just using top-K examples in SGD bring little performance drop. One possible reason is that the   batch size 128 used here is pretty small. Can the authors also provide analysis with a larger batch size.\n- The reported training time shows that K-SAM can achieve speedup compared to SGD, while on ImageNet it's much slower. Can the authors provide metrics like FLOPs, which is not impacted by the hardware.\n- It could be better if the authors can analyze whether we can save more compute in the ascent or the descent steps? Probably you can control the training time but vary the ratio of K1/K2. I assume that the ascent step can be performed with fewer example thus more savings.\n- For the BERT finetuning result, the authors should include the variance as it's know to be pretty high.\n- There are some other efficient SAM approaches that the authors can consider to compare with. For example, only performing SAM every k step[1], etc.\n\nLiu et al. \"Towards Efficient and Scalable Sharpness-Aware Minimization.\" CVPR 2022.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "SoNNyFNU2Hp", "reviewer": "Reviewer_T2UG", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper investigates the application of the \"top-k trick\" known from SGD, where only the $k$ examples with the highest loss values are used in backward pass, to the sharpness-aware minimization (SAM) method. It shows via extensive experiments that this can reduce the computational cost of SAM to the level of vanilla SGD while maintaining the generalization-boosting effects of SAM.", "summary_of_the_review": "This is a well-written paper providing a small but significant improvement to a relevant technique such as SAM and demonstrates that improvement in extensive experiments. In terms of originality and scientific value, this is certainly a very incremental paper. That, together with the missing baseline of top-$K$-SGD (weakness 2), puts this paper below the bar for a top conference such as ICLR.", "strength_and_weaknesses": "### Strengths\n\n1) The paper takes a highly relevant technique, SAM, and improves its computational efficiency. In my opinion that is commendable work with a direct positive impact on practitioners.\n\n2) The paper is well-written. It does a good job of positioning itself in the context of existing work and summarizing the work it builds upon. Overall, the claims are scoped and phrased adequately.\n\n3) The experimental comparison is well done. The selected benchmark problems offer a good amount of variety. The most relevant ablation/sensitivity studies (random examples instead of top-$k$, effect of $K_1$ and $K_2$) are included.\n\n\n### Weaknesses\n\n1) The paper is of limited originality. It combines the known optimization method SAM with the known (from vanilla SGD) speed-up method of selecting the examples with the top-k losses. Personally, I am fine with such a rather incremental paper, but it requires excellent execution.\n\n2) My main criticism is that I sorely miss one baseline to compare against in the experiments: top-$K$-SGD (e.g. with $K=K_1+K_2$). The paper evaluates methods with regards to their trade-off between generalization and computational cost. For SAM, it uses the \"top-$k$ trick\", for SGD it does not. In my view, that is an incomplete picture. Assuming that top-$K$-SGD can maintain the generalization performance of SGD for reasonable values of $K$ it would naturally be the fastest method. This would make the claim of providing \"significant generalization boosts over [...] SGD at little to no additional costs\" rather vacuous. (Importantly, I don't think that would reduce the value of the paper. It is still worthwhile investigating the application of the top-k trick to SAM. But I _do_ think that this baseline is needed to give the full picture.)\n\n3) The proposed methods computes the ranking based on the losses computed for the _ascent_ step and then uses that ranking for both the ascent and the descent step. In my view, the most straight-forward way of transferring the top-k trick from vanilla SGD to SAM would be to select the top-k examples _separately_ for the ascent step and the descent step. Since these loss values would be obtained at $w$ and $w+\\varepsilon_\\ast$, respectively, this could lead to slightly different results. Using the ranking from the ascent step might be a totally adequate strategy, but I think the first option should be mentioned/discussed in the paper and it would be nice to see an ablation study on that.\n\n\n### Update after Rebuttal\n\nThank you for your response. I see now that Table 8 in the appendix contains numbers for top-K-SGD on ImageNet. (However, that is a _single_ experiment.) If we compare top-K-SGD with K-SAM such that $K=K_1 + K_2$ for (roughly) equal computational cost, the test performance achieved by the two methods is actually quite close. In any case, I think this is an absolutely crucial baseline to include for all experiments and contrast the results for K-SAM with. Therefore, unfortunately this doesn't resolve my concern outlined in the original review and I will keep my recommendation as is.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject