Efficient Sharpness-aware Minimization for Improved Training of Neural Networks
Authors: Jiawei Du
Hanshu Yan
Jiashi Feng
Joey Tianyi Zhou
Liangli Zhen
Rick Siow Mong Goh
Vincent Tan
Venue: ICLR 2022 Poster (2022)
Keywords: Efficient learning
gengeralization
training algorithm

Abstract: Overparametrized Deep Neural Networks (DNNs) often achieve astounding performances, but may potentially result in severe generalization error. Recently, the relation between the sharpness of the loss landscape and the generalization error has been established by Foret et al. (2020), in which the Sharpness Aware Minimizer (SAM) was proposed to mitigate the degradation of the generalization. Unfortunately, SAM’s computational cost is roughly double that of base optimizers, such as Stochastic Gradient Descent (SGD). This paper thus proposes Efficient Sharpness Aware Minimizer (ESAM), which boosts SAM’s efficiency at no cost to its generalization performance. ESAM includes two novel and efficient training strategies—StochasticWeight Perturbation and Sharpness-Sensitive Data Selection. In the former, the sharpness measure is approximated by perturbing a stochastically chosen set of weights in each iteration; in the latter, the SAM loss is optimized using only a judiciously selected subset of data that is sensitive to the sharpness. We provide theoretical explanations as to why these strategies perform well. We also show, via extensive experiments on the CIFAR and ImageNet
datasets, that ESAM enhances the efficiency over SAM from requiring 100% extra computations to 40% vis-`a-vis base optimizers, while test accuracies are preserved or even improved.

Published as a conference paper at ICLR 2022
EFFICIENT SHARPNESS-AWARE MINIMIZATION FOR IMPROVED TRAINING OF NEURAL NETWORKS
Jiawei Du Hanshu Yan? Jiashi Feng² Joey Tianyi Zhou Zhou'; :Liangli Zhen",
Rick Siow Mong Goh*, Vincent Y. F. Tan 3.2
1 CCentre for Frontier AT Research (CFAR), ASTAR, Singapore,
DDpaarment of Electrical and Computer Engineering, National University of Singapore
3 'Department of Mathematics, National University of Singapore
IInstitute of High Performance Computing (IHPC), A*STAR, Singapore {dujiawei, hanshu. yan nus. edu, vtan.nus edu.sg jshfengegmail. com, Joey. rianyi. zhou@ amail..mmall....... com (zhen liangli, gohsm} ihpc. a-star edu.
ABSTRACT
Overparametrized Deep Neural Networks (DNNss often achieve astounding per- formances, but may potentially result in severe generalization error. Recently, the relation between the sharpness of the loss landscape and the generalization error has been established by Foret et al. (2020), in which the Sharpness Aware Minimizer (SAM) was proposed mitigate the degradation of the generaliza- tion. Unfortunately, SAM's computational cost is roughly double that of base optimizers, such as Stochastic Gradient Descent (SGD). This paper thus proposes Efficient Sharpness Aware Minimizer (ESAM), which boosts SAM's efficiency at no cost to its generalization performance. ESAM includes two novel and effi- cient training Data Selection. In the former, the sharpness measure is approximated by per- turbing stochastically chosen set of weights in each iteration; in the latter, the SAM loss optimized using only a judiciously selected subset data that is sensitive the sharpness. We provide theoretical explanations to why these strategies perform well. We also show, via extensive experiments on the CIFAR and ImageNet datasets, that ESAM enhances the efficiency over SAM from re- quiring 100% extra computational overhead to 40% vis-à-vis base optimizers, while test accuracies preserved even improved. Our codes are avaliable at https:: //ttttttttttt huub. com/ wttp//fciientteee.ee
SStcchstic Weight Perturbation and Shaspness-Sensitive
INTRODUCTION
Deep learning has achieved astounding performances in many fields by relying larger numbers of parameters and increasingly sophisticated optimization algorithms. However, DNNs with far more parameters than training samples are more prone poor generalization. Generalization is arguably the most fundamental and yet mysterious aspect of deep learning.
Several studies have been conducted better understand the generalization of DNNs and to train DNNs that generalize well across the natural distribution (Keskar et al., 2017; Neyshabur et al., 2017; Chaudhari al., 2019; Zhang et al., 2019; Wu et al., 2020; Foret et al., 2020; Zhang al., 2021). For example, Keskar et al. (2017) investigate the effect of batch size on neural networks' generalization ability. Zhang al. (2019): Zhou et (2021) propose optimizers for training DNNs with improved generalization ability. Specifically, Hochreiter & Schmidhuber (1995), Li et al. (2018 1) and Dinh al. (2017) argue that the geometry of the loss landscape affects generalization and DNNs with flat minimum can generalize better. The recent work by Foret et al. (2020) proposes an effective training algorithm Sharpness Aware Minimizer (SAM) for obtaining a flat minimum. SAM employs base optimizer such Stochastic Gradient Descent (Nesterov, 1983) Adam (Kingma & Ba, 2015) to minimize both the vanilla training loss and the sharpness. The sharpness, which describes
corresponding author

---

Published as a conference paper at ICLR 2022
the flatness of minimum, is characterized using cigenvalues of the Hessian matrix by Keskar al. (2017). SAM quantifies the sharpness as the maximized change of training loss when a constraint perturbation is added current weights. As result, SAM leads to flat minimum and significantly improves the generalization ability of the trained DNNs. SAM and its variants have been shown to outperform the state-of-the-art across variety of deep learning benchmarks (Kwon et al., 2021; Chen et al., 2021; Galatolo et al., 2021; Zheng et al., 2021). Regrettably though, SAM and its vari- ants achieve such remarkable performance at the expense of doubling the computational overhead of the given base optimizers, which minimize the training loss with single forward and backward propagation step. SAM requires an additional propagation step compared to the base optimizers to resolve the weight perturbation for quantifying the sharpness. The extra propagation step requires the same computational overhead as the single propagation step used by base optimizers, resulting in SAM's 's computational overhead being doubled (2x). As demonstrated in Figure 1, SAM achieves higher test accuracy (i.c., 84.46% vs. 81.89%) at the expense of sacrificing half of the training speed of the base optimizer (i.e., 276 imgs/s vs. 557 imgs/s).
In this paper, we aim to improve the efficiency of SAM but preserve its superior performance in general- ization. We propose Efficient Sharpness Aware Mini- mizer (ESAM), which consists of two training strategies Stochastic Weight Perturbation (SWP) and Sharpness- sensitive Data Selection (SDS), both of which reduce computational overhead and preserve the performance of SAM. On the one hand, SWP approximates the sharp- ness by searching weight perturbation within stochasti- cally chosen neighborhood of the current weights. SWP preserves the performance by ensuring that the expected weight perturbation is identical to that solved by SAM. On the other hand, SDS improves efficiency by ap- proximately optimizing weights based on the sharpness- sensitive subsets of batches. These subsets consist of samples whose loss values increase most w.r.t. the weight perturbation and consequently sharpness of DNNs. As a result, the sharpness calculated over the subsets can serve as an upper bound of the SAM's sharpness, ensuring that SDS's performance is comparable to that of SAM's.
700
Acceracy
Training Speed
Acturaty
-000
02
300
ESSAM
Figure 1: Training Speed vs. Accuracy SGD, SAM and ESAM evaluated by Pyra- midNet on CIFAR100. ESAM improves the efficiency with better accuracy compared to SAM.
better quantify the
We verify the effectiveness of ESAM on the CIFARIO, CIFAR 100 (Krizhevsky et al., 2009) and Im- ageNet (Deng al., 2009) datasets with five different DNN architectures. The experimental results demonstrate that ESAM obtains flat minima at cost of only 40% (vs. SAM's 100%) extra com- putational overhead over base optimizers. More importantly, ESAM achieves better performance in terms of the test accuracy compared to SAM. In nutshell, our contributions are as follows:
We propose two novel and effective training strategies Stochastic Weight Perturbation (SWP) and Sharpness-vensitive Data Selection (SDS). Both strategies are designed to im- prove efficiency without sacrificing performance. The empirical results demonstrate that both of the proposed strategies can improve both the efficiency and effectiveness of SAM.M.
We introduce the ESAM, which integrates SWP and SDS. ESAM improves the generaliza- tion ability of DNNs with marginally additional computational cost compared to standard training.
The rest of this paper is structured in this way. Section 2.1 introduces SAM and its computational issues. Section 2.2 and Section 2.3 discuss how the two proposed training strategies SWP and SDS designed respectively. Section 3 verifies the effectiveness of ESAM across variety of datasets and DNN architectures. Section 4 presents the related work and Section 5 concludes this paper.
2 METHODOLOGY
We start with recapitulating how SAM achieves flat minimum with small sharpness, which is quantified by resolving a maximization problem. To compute the sharpness, SAM requires addi- tional forward and backward propagation and results in the doubling of the computational overhead
2

---

Published as a conference paper at ICLR 2022
Algorithm 1 Efficient SAM (ESAM)
Input: Network fo with parameters 8 (01, 02...................... .a.....,yy. 0N); Training set S; Batch size b; Learning
rate > O; Neighborhood size p > 0; Number of iterations A; SWP hyperparameter 3; SDS
hyperparameter 7,
Output: A flat minimum solution 0.
1: for = to A do
2: Sample a mini-batch B S with size
3: for
to N do if 0n chosen by probability 3 then én c I-3 Vo, (Jo)
4:
5:
D. SWP in B
6:
else
7:
En 0
8:
& EN)
b Assign Weight Perturbation
9:
Compute ((fo+e: F4, y4) and construct B+ with selection ratio y (Equation 6)
10:
Compute gradients = Vol.s (Jo+z)
D SDS in Bg
11:
Update weights A t 0 199
compared to base optimizers. Following that, we demonstrate how we derive and propose ESAM, which integrates SWP and SDS, to maximize efficiency while maintaining the performance. We in- troduce SWP and SDS in Sections 2.2 and 2.3 respectively. Algorithm 1 shows the overall proposed ESAM algorithm.
Throughout this paper, denote neural network with weight parameters 0 fo. The weights are contained in the vector f (01, 0...... .., 02), where is the number of weight units in the neural network. Given training dataset $ that contains samples i..d. drawn from distribution D, the network trained obtain optimal weights / via empirical risk minimization (ERM), i...
min min Lo(fo)
1
Σ (#:.)))EB
e(fe. ui)
6
where can be an arbitrary loss function. We take l to be the cross entropy loss in this paper. In each training iteration, optimizers sample mini-batch B C S with size b update parameters.
The population loss is defined as LD(fa) - ((e,,)))D [C[[s xi, x., v:)].
.
2.1 SHARPNESS-AWANE MINIMIZATION AND ITS COMPUTATIONAL DRAWBACK
To improve the generalization capability of DNNs, Foret et al. (2020) proposed the SAM training strategy for searching flat minima. SAM trains DNNs by solving the following min-max optimiza- tion problem,
min max Le(fe+e). e:|el||SS
Given 0, the inner optimization attempts to find weight perturbation in Euclidean ball with radius that maximizes the empirical loss. The maximized loss at weights is the sum of the empirical loss and the sharpness, which defined be Rs(fa) = maxe:llela<<p) [a((e++)) max=:leela<p_[aa[[[e. - La( This sharpness is quantified by the maximal change of empirical loss when perturbation (whose norm constrained by p) is added to The min max problem encourages SAM to find flat minima.
For certain set of weights 0, Foret et al. (2020) theoretically justifies that the population loss of DNNs can be upper-bounded by the sum of sharpness, empirical loss, and a regularization term on the norm of weights (refer to Equation 3). Thus, by minimizing the sharpness together with the empirical loss, SAM produces optimized solutions for DNNs with flat minima, and the resultant models can thus generalize better (Foret et al., 2020; Chen et al., 2021; Kwon al., 2021). Indeed, we have
(o(fo) s Ra(fe) + (e(o)) + 2lllr = max Ls(fete) + 2|101
(3)
In practice, SAM first approximately solves the inner optimization by means of a single-step gradient descent method, i.e.,
arg max Ls z pVoLsffo)
<liclla<p
3

---

Published as a conference paper at ICLR 2022
The sharpness at weights 0 is approximated by Rs(fo) Ls((o+2) - Ls(f)). Then, a base opti- mizer, such as SGD (Nesterov. 1983) or Adam (Kingma & Ba, 2015), updates the DNNs' weights to minimize Ls(fo+)) We refer to (f(++)) the SAM loss. Overall, SAM requires two forward and two backward operations to update weights once. refer to the forward and backward propagation for approximating as Fi and B1 and those for updating weights by base optimizers respectively, Although SAM can effectively improve the generalization of DNNs, it additionally re- quires one forward and one backward operation (Fi and B1) in each training iteration. Thus, SAM results in doubling of the computational overhead compared to the use of base optimizers.
F2 and B2
To improve the efficiency of SAM, we propose ESAM, which consists of two strategies SWP and SDS, to accelerate the sharpness approximation phase and the weight updating phase. Specif- ically, on the one hand, when estimating around weight vector 0, SWP efficiently approximates ? by randomly selecting each parameter with given probability to form subset of weights to be perturbed. The reduction of the number of perturbed parameters results in lower computational overhead during the backward propagation. SWP rescales the resultant weight perturbation so as to assure that the expected weight perturbation equals to é, and the generalization capability thus will not be significantly degraded. On the other hand, when updating weights via base optimizers, instead of computing the upper bound (fo++) over a whole batch samples, SDS selects a subser of samples, IB+,, whose loss values increase the most with respect to the perturbation Optimizing the weights based on a fewer number of samples decreases the computational overhead (in linear fashion). We further justify that LB(fe+d can be upper bounded by Lg+ (fe+z) and consequently the generalization capability can be preserved. In general, ESAM works much more efficiently and performs as well as SAM in terms of the generalization capability.
2.2 STOCHASTIC WEIGHT PERTURBATION
This section elaborates on the first efficiency enhancement strategy. SWP, and explains why SWP can effectively reduce computational overhead while preserving the generalization capability.
To efficiently approximate É(O, S) during the sharpness estimation phase, SWP randomly chooses a subset Ô = 10 .1,.-..> from the original set of weights 0 = (01...., 0v) to perform backprop- agation B1. Each parameter is selected to be in the subvector 0 with some probability B, which can be tuned as a hyperparameter. SWP approximates the weight perturbation with pVIIs(fo)-
i.dd To be formal, we introduce a gradient mask m = (mi..... my) where m, ~ Bern(3) for all i E (l,..., Nj. Then, we have pV g turbation of SWP equals é, we scale p (((o by factor of z. Finally, SWP produces an approximate solution of the inner maximization as
E(0, B). To ensure the expected weight per-
()
a00 BB)
(5)
|
Computation Ideally, SWP reduces the overall computational overhead in proportion to 1 - B in B1. However, there exists some parameters not included in Ø that are still required to be updated in the backpropagation step. This additional computational overhead is present due to the use of the chain rule, which calculates the entire set of gradients with respect to the parameters along a prop- agation path. This additional computational overhead slightly increases in deeper neural networks. Thus, the amount of reduction in the computational overhead is positively correlated to 1- - B. In practice, is tuned maximize SWP's efficiency while maintaining generalization performance comparable to SAM's.
Generalization We will next argue that SWP's generalization performance can be preserved when compared to SAM by showing that the expected weight perturbation a(,, IB) of SWP equals to the original SAM's ss perturbation 2(0, B) the sense of the &g norm and direction. We denote the expected SWP perturbation by a(0, B), where
1 3
a(0, IB) M
Ela B)
3e(0, B) = (),)))))
for i E (1....,,, N7. Thus, it holds that
(0 B2| = ||ã(0, B)]l and CosSim (),, (), BB) = 1,
4

---

Published as a conference paper at ICLR 2022
showing that the expected weight perturbation of SWP is the same as that of SAM's.
2.3 SHARRPESSS-
DATA SELECTION
In this section, we introduce the second efficiency enhance- ment technique, SDS, which reduces computational overhead of SAM linearly as the number of selected samples decreases. We also explain why the generalization capability of SAM is preserved by SDS.
B (fo+
In the sharpness estimation phase, we obtain the approximate solution é of the inner maximization. Perturbing weights along this direction significantly increases the average loss over a batch B. To improve the efficiency but still control the upper bound Læ((f+)), we select subset of samples from the whole batch. The loss values of this subset of samples increase most when the weights are perturbed by é. To specific, SDS splits the mini-batch B into the following two subsets
a= 0
LB- (fo+e)
Illustration the loss changes of samples B and B along the weight perturbation é. The average loss of samples in B in- creases the most along the perturbation direction
Figure 2:
B = [ Mi) 6 B e(foofooe....... Fi mi) (.fo. a),
(6)
B =- [(.,,,) E B ((Jo+e, m. f.fo. m²) < <a),
where B+ is termed as the sharpness-sensiteee threshold controls the size of B+. We 7 |B+|/|B| be the ratio of the number of selected samples with respect to the batch size. In practice, determines the exact value of a and serves as a predefined hyperparameter of SDS. As illustrated in Figure 2, when weights evaluated on B+ aligns with the direction of and the loss values of the samples in I will increase with respect the weight perturbation
subset and the
0, the gradient of the
Computation SDS reduces the computational overhead in F2 and Bg. The reduction is linear in 7- The hyperparameter 7y can be tuned to meet up distinct requirements in efficiency and performance. SDS is configured the same as SWP for maximizing efficiency with comparable per- formance to SAM.
Generalization For the generalization capability, we now justify that the SAM loss computing over the batch BB, [n(J)))), can approximately upper bounded by the corresponding loss evalu- ated only on B+, Lz+ (fo+z). From Equation 3, we have
Lz(fe I (fete) + (1 ((1-))Ln LB (fo+e)
Lat (fo+z) + -
(1->) [n (fotz) - La+ (fo+z)l
LB (fote) + - (1) (fo) + LB (fo) Rg (fa) - LB (fe)].
On the one hand, since RR(Ja) 1 E(er Den [((S0zz Iz, xz, - [(Jo, v:3) represents the aver- age sharpness of the batch BB by Equation 6, we have R3-(fo) (fo) < RB(fo) I (fo), and
[3|
Rs (fo) R31 (fo) S 0.
(8)
On the other hand, B+ and B are constructed by sorting E(fore, 2i, 9:) - e(fe, T, vi). which is positively correlated to 1(fo, vi) (Li al., 2019) (more details can be found in Appendix A.2). Thus, we have
(f- La+ (fo) S O.
(9)
Therefore, by Equation 8 and Equation 9, we have
La(fote) S LB+ (fo+z).
(10)
Experimental results in Figure corroborate that Rg- (fo) Rg- (fa) < and Lf- (fo) < 0. Besides, Figure 6 verifies that the selected batch B+ is sufficiently representative to mimic the gradients of since B+ has significantly higher cosine similarity with B compared to B in terms of the computed gradients. According to Equation 10, one utilize Lat (fo+e) as a proxy to the real objective to minimize of the overall loss (n(f)))) with a smaller number of samples. As a result, SDS improves SAM's efficiency without performance degradation.
5

---

Published as a conference paper at ICLR 2022
3 EXPERIMENTS
This section demonstrates the effectiveness of our proposed ESAM algorithm. We conduct experi- ments on several benchmark datasets: CIFAR- 10 (Krizhevsky al., 2009), CIFAR- (Krizhevsky et al., 2009) and ImageNet (Deng et al., 2009), using various model architectures: ResNet (He et al., 2016), Wide ResNet (Zagoruyko & Komodakis, 2016), and PyramidNet (Han et al., 2017). We demonstrate the proposed ESAM improves the efficiency of vanilla SAM by speeding up to 40.3% computational overhead with better generalization performance. We report the main results in Ta- ble 1 and Table 2. Besides, perform an ablation study on the two proposed strategies of ESAM (i.e.. SWP and SDS). The experimental results in Table 3 and Figure 3 indicate that both strategies improve SAM's efficiency and performance.
3.1 RESULTS
CIFARIO and CIFAR100 We start from evaluating ESAM the CIFAR-10 and CIFAR-100 image classification datasets. The evaluation is carried out on three different model architectures: ResNet- 18 (He et al., 2016), WideResNet-28-100 (Zagoruyko & Komodakis, 2016) and PyramidNet-1 10 (Han et al., 2017). We set all the training settings, including the maximum number of training epochs, iterations per epoch, and data augmentations, the same for fair comparison among SGD, SAM and ESAM. Additionally, the other hyperparameters of SGD, SAM and ESAM have been tuned separately for optimal test accuracies using grid search.
We train all the models with 3 different random seeds using a batch size of 128, weight decay 10-4 and cosine learning rate decay (Loshchilov Hutter, 2017). The training epochs are set be 200 for ResNet-18 18 (He et al., 2016), WideResNNet28-UD (Zagoruyko & Komodakis, 2016), and 300 for PyramidNet-1 (Han et al., 2017). We set 3 = 0.6 and y 0.5 for ResNet-18 and PyramidNet- 110 models; and set 3 = 0.5 and 7 0.5 for WideResNett22-1U The above-mentioned B and are optimal for efficiency with comparable performance compared to SAM. The details of training setting are listed in Appendix A.7. We record the best test accuracies obtained by SGD, SAM and ESAM in Table 1.
The experimental results indicate that our proposed ESAM increase the training speed by up to 40.30% in comparison with SAM. Concerning the performance, ESAM outperforms SAM in the six sets of experiments. The best efficiency of ESAM is reported in CIFARIO trained with ResNet- 18 (Training speed 140.3% vs. SAM 100%). The best accuracy is reported in CIFARIOR 100 trained with PyramidNetl 10 (Accuracy 85.56% vs. SAM 84.46%). ESAM improves efficiency and achieves better performance compared to SAM in CIFARI0/100 benchmarks.
Table I:
Classification accuracies and training speed on the CIFAR-10 10 and CIFAR-100 datasets. Compu- tational overhead quantified by #images processed per second (images/s). The numbers in parentheses (-) indicate the ratio of ESAM's training speed w.r.t. SAM.
CIFAR-10
CIFAR-100
ResNet-18
Accuracy
images/s
Accuracy
images/s
SGD
95.41 0.03
3,387
78.17 0.05
3,483
SAM
96.52
1,717(100.0%
80.17 0.17
1,730 (100.0%)
ESAM
96.56
2,409 (140.3%)
80.41
2,423 (140.0%)
Wide-28-10
Accuracy
images/s
Accuracy
images/s
SGD
96.34 0.12
801
81.56 0.13
792
SAM
97.27
396 (100.0%)
83,42
391 (100.0%)
ESAM
97.29 0.11
550 (138.9%)
84.51 0.01
545 (139.4%)
PyramidNet-110
Accuracy
images/s
Accuracy
images/s
SGD
96.62
580
81.89
555
SAM
97.30 0.10
289 (100.0%)
84,46 0.04
276 (100.0%)
ESAM
97.81 0:01
401 (138.7%)
85.56 0.85
381 (137.9%)
ImageNet To evaluate ESAM's effectiveness on large-scale benchmark dataset, we conduct ex- periments ImageNet Datasets.Thh 1000 class ImageNet dataset contains roughly 1.28 million

---

Published as a conference paper at ICLR 2022
Table 2: Classification accuracies and training speed on the ImageNet dataset. The numbers in parentheses (-)
indicate the ratio of ESAM's training speed w.r.t. SAM's. Results with * are referred to Chen et al. (2021)
ResNet-50
ResNet-101
ImageNet
Accuracy
images/s
Accuracy
images/s
SGD
76,00*
1,327
77.80*
891
SAM
76.70*
654 (100.0%)
78. 78.60*
438 (100.0%)
ESAM
77.05
846 (129.3%)
79.09
564 (128.7%)
training images and 50, 000 validation images with 469 387 averaged resolution. The ImageNet dataset is more representative (of real-world scenarios) and persuasive (of a method's s effectiveness) than CIFAR datasets. We resize the images on ImageNet to 224 224 resolution to train ResNet-50 and 101 models. We train 90 epochs and set the optimal hyperparameters for SGD, SAM and ESAM suggested by Chen al. (2021), and the details are listed in appendix A.7. We | = 0.6 and 7 = 0.7 for ResNet-50 and ResNet-101. We employ the m--sarppeess strategy for both SAM and ESAM with m 128, which is the same that suggested in Zheng al. (2021).
The experimental results are reported in Table 2. The results indicate that the performance of ESAM on large-scale datasets is consistent with the two (smaller) CIFAR datasets. ESAM outperforms SAM by 0.35% to 0.49% in accuracy and, more importantly, enjoys 28, 7% faster training speed compared to SAM. As the y we used here larger than the one used in CIFAR datasets, the training speed of ESAM here is slightly slower than that in the CIFAR datasets.
These experiments demonstrate that ESAM outperforms SAM on a variety of benchmark datasets for widely-used DNNs' architectures in terms of training speed and classification accuracies.
3.2 ABLATION AND PARAMETER STUDIES
To better understand the effectivvnnss of SWP and SDS in improving the performance and efficiency compared to SAM, we conduct four sets of ablation studies on CIFAR- 10 and CIFAR- 100 datasets using Re and WideResNeet28-- 10 models, respectively. We consider two variants of ESAM: (i) only with SWP, (ii) only with SDS. The rest of the experimental settings are identical to the settings described in Section 3.1. We conduct grid search over the interval [0.3, 0.9) for / and the interval 00.3, 0.9) for with a same step size of 0.1. We report the grid search results in Figure 3. We use 3 0.6.7 0.5 for ResNet-18; and set = 0.5, 7 0.5 for WideResNet-28-- 10 in the four sets of ablation studies. The ablation study results are reported in Table 3.
Table 3: Ablation Study ESAM on CIFAR-10 and CIFARI00. The numbers brackets |:] represent the accuracy improvement in comparison to SGD. numbers in parentheses (-) indicate the ratio of ESAM's training speed to SAM's. Green color indicates improvement compared to SAM, whereas red color suggests a degradation.
CIFAR-10
CIFAR-100
R8sNet-18
Accuracy
images/s
Accuracy
images/s
SGD
95.41
3,387
78.17
3,438
SAM
96.52 [+1.11]
1,717 (100.0%)
80.17 [+2.00]
1,730 (100.0%)
ESAM-SWP
96.74
[+1.33]
1,896 (110.5%)
80.53 [+2.36]
1,887 (109.1%)
ESAM-SDS
96.45 [+1.04]
2,105 (122.6%)
80.38 [+2.21]
2,103 (121.5%)
ESAM
96.56 [+1.15]
2,409 (140.3%)
80.41 [+2.24]
2,423 (140.9%)
Wide-28-10
Accuracy
images/s
Accuracy
images/s
SGD
96.34
801
81.56
792
SAM
97.27 [40.93]
396 (100.0%)
83.42 [+1.86]
391 (100.0%)
+ ESAM-SWP
97.37
[+1.03]
430 (108.5%)
84.44 [+2.88]
423 (108.3%)
+ ESAM-SDS
97.24 [+0,90]
495 (124.8%)
84.46 [+2.90]
492 (125.8%)
ESAM
97.29 [+0.95]
551 (138.9%)
84.51 [+2.95]
545 (139.4%)
7

---

Published as a conference paper at ICLR 2022
8esnet18 on CIFAR10
Resnet18 CIFAR100
97.25
560
81.0
86.75
87.90
SAM
ESAM SDS
FSSAM-SWPP
80.5
SAM
ESAM-5DS
ESAAA-SH SWP
8:06
Accuracy
Y-
Accuracy
80.0
06.25
03
79.5
96,00
79.0
96.75
0.5
96.50
78.0
14
16
18
20
22
24
26
28
Elapsed Time Per Epoch
Elapsed
Epoch
Wide-28-10 on CIFAR10
Wide-28- CIFAR100
SCD
SAM
55
ISM
97.4
0.5
EBAM-SWPP
y=
0
ESAM
Accuracy
Accuracy
97.0
96.6
0-3
82
130
60
70
80
90
100
110
120
150
Elapsed Time Per Epoch
Elapsed Time Per Epoch
Figure 3: Parameter study SWP SDS. The connected dots refer to SWP and SDS with different param- eters; the isolated dots refer to the final results of SGD, SAM, and ESAM.
ESAM-SWP As shown in Table 3, SWP improves SAM's training speed by 8.3% to 10.5%, and achieves better performance at the same time. SWP can further improve the efficiency by using a smaller B. The best performance of SWP is obtained when 0.6 for ResNet-18 and 3 0.5 for WideResNee-28- 10. The four sets of experiments indicate that 3 is consistent among different architectures and datasets. Therefore, ResNet on ImageNet datasets.
/ 0.6 for PyramidNet CIFAR10/100 datasets and
ESAM-SDS SDS also significantly improves the efficiency 21.5% 25.8% compared SAM. outperforms SAM's performance on CIFAR 100 datasets, and achieves comparable performance on CIFARIO 10 datasets. SDS outperform SAM on both datasets with both architectures with little degradation to the efficiency, as demonstrated in Figure 3. Across all experiments, = 0.5 is the smallest value that is optimal for efficiency while maintaining comparable performance to SAM.
Visualization of Loss Landscapes To visualize the sharpness of the flat minima obtained by ESAM, we plot the loss landscapes trained with SGD, SAM and ESAM on the ImageNet dataset. We display the loss landscapes in Figure 4, following the plotting algorithm in Li et al. (2018). The and y- axes represent two random sampled orthogonal Gaussian perturbations. We sampled 100 100 points for 10 groups random Gaussian perturbations. The displayed loss landscapes are the results we obtained by averaging over ten groups of random perturbations. can be clearly seen that both SAM and ESAM improve the sharpness significantly in comparison to SGD.
To summarize, SWP and SDS both reduce the computational overhead and accelerate training com- pared to SAM. Most importantly, both these strategies achieve a comparable better performance than SAM. In practice, by configuring the 3 and 7, ESAM can meet a variety of user-defined effi- ciency and performance requirements.
RELATED WORK
The concept of regularizing sharpness for better generalization dates back to (Hochreiter Schmid- huber, 1995). By using MDL--based argument, which clarifies that statistical model with fewer
8

---

Published as a conference paper at ICLR 2022
Training Loss
Training Loss
Training Loss
0.5
0.00..20
0.0
0.5
1.0 9. 0.0-0.5
1.0
0.5
0.0
9,5
-0.5
110
1.0
0.5
1.0
-0.5
-1.0
10
0.5
SGD
SAM
ESAM
Figure 4: Cross-entropy loss landscapes of the ResNet50 model on the ImageNet datasct trained with SGD, SAM, and ESAM.
bits to describe can have better generalization ability, Hochreiter & Schmidhuber (1995) claim that flat minimum alleviate overfitting issues. Following that, more studies were proposed in- vestigate the connection between the flat minima with the generalization abilities (Keskar et al., 2017: Dinh et al., 2017: Liu al., 2020; Li al., 2018; Dziugaite & Roy, 2017; Jiang et al., 2019; Moosavi-Dezfooli et al., 2019). Keskar et al. (2017) starts by investigating the phenomenon that training with a larger batch size results in worse generalization ability. The authors found that the sharpness of the minimum is critical in accounting for the observed phenomenon. Keskar et al. (2017) and Dinh et al. (2017) both argue that the sharpness can be characterized using the eigenval- ues of the Hessian. Although they also define specific notions and methods to quantify sharpness, they do not propose complete training strategies to find minima that are relative "flat".
SAM (Foret et al., 2020) leverages the connection between "flat" minima and the generalization error to train DNNs that generalize well across the natural distribution. Inspired by Keskar et al. (2017) and Dinh et al. (2017). SAM first proposes the quantification of the sharpness, which is achieved by solving a maximization problem. Then. SAM proposes a complete training algorithm to improve the generalization abilities of DNNs. SAM is demonstrated to achieve state-of seee-of-thee off-the-ue art performance in a variety of deep learning benchmarks, including image classification, natural language processing, and noisy learning (Foret et al., 2020; Chen al., 2021; Kwon et al., 2021; Pham et al., 2021; Yuan et al., 2021; Jia et al., 2021).
A series of SAM-related works has been proposed. A work that was done contemporancously SAM (Wu et al., 2020) also regularizes the sharpness term in adversarial training and achieves much more robust generalization performance against adversarial attacks. Many works focus on combining SAM with other training strategies or architectures (Chen et al., 2021; Wang et al., 2022; Tseng et al., 2021), or apply SAM on other tasks (Zheng et al., 2021; Damian et al., 2021; Galatolo al., 2021). Kwon et al. (2021) improves SAM's sharpness by adaptively scaling the size of the nearby search space P in relation to the size of parameters. Liu et al. (2022) leverages the past calculated weight perturbations to save SAM's computations. However, most of these works overlook the fact that SAM improves generalization at the expense of the doubling the computational overhead. As result, most of the SAM--related works suffer from the same efficiency drawback SAM. This computational cost prevents SAM from being widely used in large-scale datasets and architectures, particularly in real-world applications, which motivates to propose ESAM esficiently improve the generalization ability of DNNs.
5 CONCLUSION
In this paper, we propose the Efficient Sharpness Aware Minimizer (ESAM) to enhance the efficiency of vanilla SAM. The proposed ESAM integrates two novel training strategies, namely, SWP and SDS, both of which are derived based on theoretical underpinnings and are evaluated over a variety of datasets and DNN architectures. Both SAM and ESAM are two-step training strategies consisting of sharpness estimation and weight updating. In each step, gradient back propapropagttion is performed to compute the weight perturbation updating. In future research, we will explore how to combine the two steps into one by utilizing the information of gradients in previous iterations so that the computational overhead of ESAM can be reduced to the same base optimizers.

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGEMENT
Jiawei Du and Joey Tianyi Zhou are suppored by Joey Tianyi Zhou's A*STAR SERC Central Re- search Fund.
Hanshu Yan and Vincent Tan are funded by Singapore National Research Foundation (NRF) Fel- lowship (--26--000--02--81) and Singapore Ministry of Education AcRF Tier grant (R-263- 000-E80-114).
We would like to express our special thanks of gratitude to Dr. Yuan Li for helping us conduct experiments on ImageNet, and Dr. Wang Yangzihao for helping us implement Distributed Data Parallel codes.
REFERENCES
Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-SGD, Biasing gradi- ent descent into wide valleys. Journal of Statistical Mechanics: Theory and Experiment, 2019
(12):124018,, 2019.
Xiangning Chen, Cho-Jui Hsich, and Boqing Gong. When vision transformers outperform resnets without pretraining strong data augmentations. arXiv preprint arXiv:2/06.01544, 2021.
Alex Damian, Tengyu Ma, and Jason Lee. Label noise sgd provably prefers flat global minimizers. Advances in Neural Information Processing Systems, 34, 2021.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hi- erarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255. Icee, 2009.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In International Conference Machine Leurning, pp. 1019-1028. PMLR, 2017.
Jiawei Du, Hu Zhang, Joey Tianyi Zhou, Yi Yang, and Jiashi Feng. Query-efficient meta attack to deep neural networks. In International Conference Learning Representations, 2019.
Gintare Karolina Dziugaite and Daniel M Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data, arXiv preprint arXiv:1703.11008 2017.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimiza- tion for efficiently improving generalization. In International Conference Learning Represen- tations, 2020.
Alessio Galatolo, Alfred Nilsson, Roderick Karlemstrand, and Yineng Wang. Using carly-learning regularization to classify real-world noisy data. arXiv preprint arXiv:2105.13244 2021.
Dongyoon Han, Jiwhan Kim, and Junmo Kim. Deep pyramidal residual networks. In Proceedings of the IEEE conference 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.
Sepp Hochreiter and Jürgen Schmidhuber. Simplifying neural nets by discovering flat minima. In Advances in neural information processing systems, pp. 529-536, 1995.
Chao Jia, Yinfei Yang, Ye Xia, Yi-Ting Chen, Zarana Parekh, Hieu Pham, Quoc Le, Yun-Hsuan Sung, Zhen Li, and Tom Duerig. Scaling up visual and vision-language representation learning with noisy text supervision. International Conference on Machine Learning, pp. 4904-4916. PMLR, 2021.
10

---

Published as a conference paper at ICLR 2022
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, 2019.
Nitish Shirish Keskar, Jorge Nocedal, Ping Tak Peter Tang, Dheevatsa Mudigere, and Mikhail Smelyanskiy. On large-batch training for deep learning: Generalization gap and sharp minima. In 5th International Conference on Learning Representations, ICLR 2017, 2017.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. In ICLR (Poster), 2015.
Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. CIFAR-10 and CIFAR-100 datasets. URI: https://www. cs. toronto. edu/kriz/ifar. html, 6(1):1, 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 Con- ference on Machine Learning, pp. 5905-5914. PMLR, 2021.
Buyu Li, Yu Liu, and Xiaogang Wang. Gradient harmonized single-stage detector. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 8577-8584, 2019.
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss land- scape of neural nets. Advances neural information processing systems, 31, 2018.
Chen Liu, Mathieu Salzmann, Tao Lin, Ryota Tomioka, and Sabine Stisstrunk. On the loss landscape of adversarial training: Identifying challenges and how to overcome them. Advances in Neural Information Processing Systems, 33:21476-21487, 2020.
Yong Liu, Siqi Mai, Xiangning Chen, Cho-Jui Hsieh, and Yang You. Towards efficient and scalable sharpness-awe aware minimization. arXiv preprint arXiv:2203.0277 2022.
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. 2017.
Seyed-Mohsen Moosavi-Dezfoole Alhussein Fawzi, Jonathan Uesato, and Pascal Frossard. Robust- ness via curvature regularization, and vice versa. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 9078-9086, 2019.
Yurii E Nesterov. A method for solving the convex programming problem with convergence rate O(1/k?). Dokl. akad. nauk Sssr, volume 269, pp. 543-547, 1983.
Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nati Srebro. Exploring general- ization deep learning. Advances in neural information processing systems, 30, 2017.
Hieu Pham, Zihang Dai, Qizhe Xie, and Quoe V Le. Meta pseudo labels. In Proceedings of the JEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11557-11566. 2021.
Chongli Qin, James Martens, Sven Gowal, Dilip Krishnan, Krishnamurthy Dvijotham, Alhussein Fawzi, Soham De, Robert Stanforth, and Pushmeet Kohli. Adversarial robustness through local linearization. Advances in Neural Information Processing Systems. 32, 2019.
Ching- Ching-Hsun Tseng, Shin-Jye Lee, Jia-Nan Feng, Shengzhong Mao, Yu-Ping Wu, Jia- Yu Shang, Mou-Chung Tseng, and Xiao-Jun Zeng. Upanets: Learning from the universal pixel attention networks. arXiv preprini arXiv: (0t::2000308440) 2021.
Pichao Wang, Xue Wang, Hao Luo, Jingkai Zhou, Zhipeng Zhou, Fan Wang, Hao Li, and Rong Jin. Scaled relu matters for training vision transformers. 2022.
Dongxian Wu, Shu-Tao Xia, and Yisen Wang. Adversarial weight perturbation helps robust gener- alization. Advances in Neural Information Processing Systems, 13:2958-2969 2020.
Hanshu Yan, Jiawei Du, Vincent Tan, and Jiashi Feng. On robustness of neural ordinary differential equations. In International Conference Learning Representations, 2019.

---

Published as a conference paper at ICLR 2022
Hanshu Yan, Jingfeng Zhang, Gang Niu, Jiashi Feng, Vincent YF Tan, and Masashi Sugiyama. CIFS: Improving adversarial robustness of CNNs via channel-wise importance- -bssed feature se- lection. International Conference Machine Learning, 2021.
Li Yuan, Qibin Hou, Zihang Jiang, Jiashi Feng, and Shuicheng Yan. Volo: Vision outlooker for visual recognition. arXiv preprint arXiv:2106. 13111222..111,, 2021.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. In British Machine Vision Conference 2016. British Machine Vision Association, 2016.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):107- 115, 2021.
Michael Zhang, James Lucas, Jimmy Ba, and Geoffrey E Hinton. Lookahead optimizer: k steps forward, step back. Advances in Neural Information Processing Systems, 32, 2019.
Yaowei Zheng, Richong Zhang, and Yongyi Mao. Regularizing neural networks via adversarial model perturbation. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pp. 8156-8165, 2021.
Pan Zhou, Hanshu Yan, Xiaotong Yuan, Jiashi Feng, and Shuicheng Yan. Towards understanding why lookahead generalizes better than SGD and beyond. In Proc. Conf. Neural Information Processing Systems, 2021.
12

---

Published as a conference paper at ICLR 2022
A APPENDIX
A.1
THE ALGORITHM OF SAM
The algorithm of SAM is demonstrated in Algorithm
Algorithm 2 SGD vs. SAM
Input: Network Se. (01, 0 ...., 0N), Training set S, Batch size Learning 7, Neighbor- hood size p. Iterations A.
Output: A minimum solution 0.
A do 2: Sample B with size b that IB C S,
I: for
3: if SGD then
- 0 5: else if SAM then
4:
No addtional computational overhead for è
6:
t (olo((o)
> additional F1 and B1 to compute é
7: Compute g VoLL((fota)
> F2 and Bg
8:
update the weights & 0 119
A.2
OPTIMIZING OVER SUBSET B REPRESENTATIVE
The sharpness-sensitive subset B+ is constructed by sorting {(Jo+e, 4i, v:) E(Jo, mi, Wi), which is positively correlated to 1(So, v4, U1). By a first-order Taylor series approximation,
E(fo+e: 24, yı) E(fo, v4, yi) = Vol(fe, T4, ys) + o(l[@]).
By Equation 4, is the aggregated gradients of each instance in the complete dataset IB, i.e.,
P
غ = arg max L (VLLs(fo) = clcba<
Tel
1=1
which indicates that E(f+e, Di. vi) E(fo. F.. Ui) is positively correlated to the gradient Val(fo. T4, 8:). Li et al. (2019) claims that the difficult examples in deep learning (the training sam- ples with high training loss) produce gradients with larger magnitudes. Therefore, {(fe+e, 24, V1) {(Ja, Fs y1) positively correlated ([So, Ft, y:). We also demonstrate the correlation empirically.
We conduct experiments to verify Equation 9 and Equation 10. In Figure 5, We plot the four losses, Lg+ (fe), LB- (fa), Lg+ (fe+z), and (f=(f)+=) w.r. the epochs. The experimental results verify that Equation 9 and Equation 10 hold for every training epoch.
Moreover, we conduct experiments to demonstrate that optimizing over the subset IB+ is much more representative than the subset B We compare the updating gradients computed from B _ and random subset Brand that [Baand = [B | B to those computed from by calculating the cosine similarity inspired by (Du et al., 2019), i.e.
CosSim CosSim(VoLa (fote), VoLs(fous)),
CosSim (oiinvVoolgaae (Jo+e), V Voln(Jo+e)),
CosSim CosSim(VoLL..n
(fonef))))
In Figure 6, we plot the cosine similarities in each training epoch evaluated with ResNer-18, Wide- 28-10 10 on CIFARIO. terms of the computed gradients. the experimental results show that B has the highest cosine similarities with B than B and the random set Bandd
A.3 LINEARITY MEASUREMENT OF SWP
As the experimental results in Figure 3 demonstrated, SWP can also improve the accuracy of ESAM compared SAM. We will investigate the advantage of SWP in terms of generalization in the
13

---

Published as a conference paper at ICLR 2022
Resnet18 on CIFAR10
Resnet18 on CIFAR100
3.0
La (fa)
Lar (fa)
La (fa)
Lg (fa)
2.0
(f0+e)
Training Lons
La (fo+2)
Lg (f++
Lg (f++2)
Training
1.5
0.5
0.0
25
50
76
100 Epochs
125
160
175
200
26
76
103 Epochs
125
160
176
200
Wide-28-10 10 CIFAR10
Wide-28-10 CIFAR100
Ly (fa)
2.0
La-
Lg (fo)
La (fo+e)
La (fo+
1.5
La (o++2))
Lg (o+
Training
Training
5.0
200
Epochs
Epochs
Figure 5: The SAM loss and the empirical loss calculated over the selected subsets B+ B w.r.t the epochs, evaluated with G8sNet-18, Wide-28- 10 CIFARIO and CIFAR100. The subset TB+ selected by SDS has much higher SAM loss and empirical loss than among all the four groups of experiments.
Resnet18 on CIFAR10
Resnet18 on CIFAR100
1.0
0.8
Simianty
0.6
Cosine Similatty
0.8
0.4
200
100 Epochs
Epochs
Wide-28-10 CIFAR10
Wide-28-10 CIFAR100
1.0
Similaryy
Codine Smillart
0.2
0.0
0.0
50
75
100
125
190
175
200
25
75
100
150
175
200
Epochs
Epochs
Figure 6: The cosine similarity between the gradients computed from subsets B B and Brnd with B. evaluated with R8sNet-18. Wide-28-10 CIFARIO CIFAR100. The gradients from subset B selected by SDS has much higher cosine similarity with the gradients from B than the gradients from B and Brnd among all the four groups of experiments.
14

---

Published as a conference paper at ICLR 2022
future research. Here provide discussion about the accuracy improvement contributed by SWP. A plausible reason for such improvement that SWP leads to a better inner maximization solved in equation 2. The current solution of é is approximated by assuming Ls(fo) is linear function. Therefore, the would result in a better inner maximization Lg(Jo) is "more linear" with respect to 0. Inspired by (Qin al., 2019; Yan al., 2019; 2021), measure the linearity of the loss function by
S(c, B) = (f (ffo) - VaLa(fo)l-
We conduct experiments the CIFARIO dataset with ResNet-18 model to verify that SWP can improve the linearity [(e, IB) of the loss function Lg(fo). We compare the linearity of ESAM with the ranging from {0.2, 0.3,...... 0.9) the SAM. The results are demonstrated in Figure It can
Linearity Measure SWP
0.080
0.028
0.004
Linearty
SAM
SWP 0.9
0.020
5t0m 8-
SWM
0.016
SWP. 8
30
75
100
125
150
175
200
Epochs
Figure 7: The linearity measurement of SWP evaluated with the ResNet18 model the CIFARIO dataset compared with SAM. The experimental results indicate that a smaller can result better linearity.
be shown that SWP will result in a better linearity as the 3 decreases. However, decreasing will also reduce the magnitude of and thus result in a worse inner maximization in equation 2. We observe that 3 = (0.5, 0.6} is optimal to balance the accuracy and efficiency of ESAM.
A.4 REDUCED COMPUTATIONAL OVERHEAD CONTRIBUTED BY SWP
Formulation We formulate the saved computational overhead contributed by SWP here. We dis- cuss and examine the computational overhead in the PyTorch framework. Suppose that the NN discussed here has layers. The most unit of parameters (D x x H x W) is the entire param- eters of layer in NN, where D is the number of kernels, C is the number of channels, H and W are the height and width the input. SWP select each basic parameter unit to compute gradients measure the saved computational
(i.e. requries grad True) with probability 3. We use gNN, overhead in terms of percentage contributed by SWP compared to the vanilla SAM.
The computational overhead (N, (8) is reduced not just by calculating gradients, but also by the storing and hooking gradients. Because the storing and hooking operations of each parameter's gradients are independent each other, the computational overhead saved from storing and hooking are proportional to 3 and irrelevant to the depth N of NN. In addition, the storing and hooking gradients are the dominant factor that result in the reduced computational overhead according to our toy example in the following.
Next we discuss the saved computational overhead that stems from the calculation operations in the general case. Some layers in the DNNs with complicated architectures such DenseNet and VIT, may have multiple basic parameter units in the same layer and be connected across any other layers. Suppose the nshh layer has K ((n) basic parameter units, let p" be the calculation- rate of a certain parameter unit in Pihh layer, we have 7d
-3) K()) n- Therefore, pd (1 -3) KO) x
15

---

Published as a conference paper at ICLR 2022
(1 11-3) 3)n* where K = = E K()) Assumed that the computations of each layer the same, by summing up the saved computational overhead of all parameters, we have
p
(N = ka B) k2
N
n=
12 - 33
2 k1 -B) +
N
12 1 3)
2 k1 - B)
(1I)
N
where ki, 22 are determined by the computing time of calculation, R/2
kg
[1
However, the commonly used DNNs' architectures such as ResNet only have one basic parameter unit in each layer. Besides, each layer of them is connected in serial with the next layer. Then,
have y" = (1 - 3) 3) p==11 Therefore, y" = (1 - 3) The saved calculation contributed by the parameter unit is pp² By summing up the saved computational overhead of all parameters, have
pd
9(N, 11 3)
K2
N
n=i
k2 1
B
(1
3)
=h1(1
N
B
h2
B)
k1 3)
(12)
N
where k1, 22 are determined by the computing time of calculation, storing and hooking gradients.
Toy Example We conducted toy example the CIFARIO dataset with two MLPs. Each fully connected layer of MLP is the same with a size 3, 000 for both in and out features. The first mlp is basic parameter unit and is connected
for the special case that K where each layer has only in serial with the next layer. We examined = {50, 75, 100, 125] and 3 = {0.1, 0.9, 1.0) to record the saved computational overhead in percentage. Part of the results By linear regression, we have k = 0.3185,4 l2 = 0.1 0.1310, and the returned π² 0.9983. The second mip is for the general case that K 1, where each layer may have multiple basic parameter units in the same layer and be connected across any other layers. We examined N = {35, 50, 65, 75) and B (0.1,... ..., 0.9, 1.0] to record the saved computational overhead in percentage. Part of the results are reported in Table A.4. By linear regression, we have k 0.3143, 2₂ 0.0737, and the returned R 0.9989. The above experimental results verify the formulation of the reduced computational overhead contributed by SWP in equation 11 and equation 12.
reported in Table A.4.
3
g(N B)
3
g (3)
50
0.9
3.42%
25
0.9
2.88%
50
0.5
16.28%
25
0.5
15.54%
0.1
30.08%
25
0.1
30.03%
75
0.9
3.44%
50
0.9
2.34%
75
0.5
15.79%
50
0.5
15.10%
75
0.1
30.42%
50
0.1
29.83%
100
0.9
3.37%
65
0.9
2.23%
100
0.5
15.86%
65
0.5
14.95%
100
0.1
29.56%
0.1
29.22%
125
0.9
2.94%
75
0.9
2.22%
125
0.5
14.87%
75
0.5
15.10%
125
0.1
29.42%
75
0.1
28.68%
Table 4: The special case that K 1. By linear regression, h: = 0.3185, 12 0.1310, and the returned R² 0.9983.
Table 5: The general case that K > By linear
regression, k 0.3143, Ki 0.0737, and the returned R² = 0.9989.
16

---

Published as a conference paper at ICLR 2022
A.5 VISUALIZATION Loss LANDSCAPES WITH RESPECT ADVERSARIAL WEIGHT PERTURBATIONS
We visualize the sharpness of the flat minima with respect adversarial weight perturbations of SGD,SAM and ESAM on the Cifarl0 dataset. The x- and y-axes represent two orthogonal ad- versarial weight perturbations, which are /VLL3, (fe) and r7V&LB, (fe) respectively, where ry is the learning rate during training. E and BB are the randomly sampled subsets of batch IP, and |B= [B y [B] B U., I. We display the loss landscape in Figure 8, which demonstrates that both SAM and ESAM improve the sharpness significantly in comparison SGD.
0.10
0.08
0.06
Trining
0.06
1.005 0. 0.5.1. 10
-0.5-1. 1.0
1.5 00 0.5..11
0.5-
1.0.5 0.0.
.05...
-0.5
10
10
SGD
SAM
ESAM
Figure 8: Cross cross-entropy loss landscapes of the ResNet18 model respect to adversarial weight perturbations on the CIFARIO dataset trained with SGD, SAM, and ESAM.
A.6 EVALUATION OF ESAM ON VIT-S/16
SAM has also been demonstrated to be effective on the new vision Transformer( ViT) architecture (Chen al., 2021). Therefore, also evaluate ESAM with ViT-S/16 on ImageNet Datasets. We use 3 = 0.5 and = 0.7 for ViT-S/16, which share the same hyperparameters as ResNet-50 and ResNet- 101 in section 3.1. The results are reported in Table 6, which indicate that ESAM can still be effective to improve efficiency in ViT-S/16 architecture. In particular, ESAM-SWP achieves much better accuracy than SAM (80.8% v.s. 80.34%).
Table 6: Classification accuracies and training speed of VIT-S/16 on the ImageNet dataset.
VIT-S/16
ImageNet
Accuracy
images/s
SGD
79.72
1.133
SAM
80.34
581
SSAM-SWP
80.88
616
ESAM-SDS
79.97
693
ESAM
80.46
734
A.7 TRAINING DETAILS
We tune the training parameters of SGD, SAM, and ESAM, by using grid searches. The learning rate is chosen from the set {0.01, 0.05, 0.1, 0.2). the weight decay from the set (5 x 10-4, 4, x 10-3), and the batch size from the (64, 128, 256). This done attain the best accuracies. The exact training hyperparameters are reported in Table 7. On the ImageNet datasets, limited by the computing resource, we follow and slightly modify the optimal hyperparameters as suggested by Chen et al. (2021) for SGD, SAM and ESAM. The exact training hyperparameters are reported in Table
17

---

Published as a conference paper at ICLR 2022
Table 7: Hyperparameters for training from scratch on CIFARIO and CIFARI00
CIFAR-10
CIFAR-100
ResNet-18
SGD
SAM
ESAM
SGD
SAM
ESAM
Epoch
200
200
Batch size
128
128
Data augmentation
Basic
Basic
Peak learning rate
0.05
0.05
Learning rate decay
Cosine
Cosine
Weight decay
5 x 10
1 x 10
X 10
x 10
x 10
1 x 10
0.05
0.05
0.05
0.05
Wide-28-10
SGD
SAM
ESAM
SGD
SAM
ESAM
Epoch
200
200
Batch size
256
256
Data augmentation
Basic
Basic
Peak learning
0.05
0.05
Learning rate decay
Cosine
Cosine
Weight decay
5 x 10
1 x 10
x 10-3 0.1
5 X 10
1 x 10
x 10-3
0.1
0.1
0.1
PyramidNet-110
SGD
SAM
ESAM
SGD
SAM
ESAM
Epoch
300
300
Batch size
256
256
Data augmentation
Basic
Basic
Peak learning rale
0.1
Learning rate decay
Cosine
Cosine
Weight decay
5 10
5 x 10
P
0.2
0.2
0.2
0.2
Table 8: Hyperparameters for training from scratch ImageNNtt
5esNet-50
ResNet-110
ImageNet
SGD
SAM
ESAM
SGD
SAM
ESAM
Epoch Batch size
90
90
512
512
Data augmentation
Inception-style
Inception-style
Peak learning rate Learning rate decay
0.2
0.2
Cosine
Cosine 1 10
Weight decay
1 10 10-4
P Input resolution
0.05
0.05
0.05
224 x 224
224 x 224
18


Reviews:
[{"review_id": "aDbCbC7zT3H", "reviewer": "Reviewer_mQEh", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "This is a nice paper, that is well written and addressing a very clear problem: reducing the computational complexity of SAM. It proposes two simple tricks that both make sense, and shows that they improve efficiency without hurting performance.\n\n## SWP Implementation\n\nI wanted to ask a question about SWP: how do you implement it in practice? I am a bit surprised that not computing the gradients for a random subset (e.g. 50%) of weights leads to a significant improvement in performance. Do you have to do any tricks to achieve this improvement?\n\n## Alignment of SSP gradient with full-batch gradient\n\nYou mention in page 5 that the gradient on the subset selected by SSP is well aligned with the gradient on the full batch. Did you try to actually measure this alignment, e.g. in terms of cosine similarity. It could be interesting to compare it to e.g. a random subset and $\\mathbb{B}^{-}$.\n\n## Why does ESAM outperform SAM?\n\nIn some of your experiments ESAM outperforms SAM in accuracy. Do you have ideas for why this could be happening? \n\n## ImageNet results\n\nOne concern I have with the paper is that the original SAM paper [1] seems to report better results for SAM in some of the same settings than you do. In particular, on ImageNet in Table 2 you report 76.7 for ResNet-50 with SAM and 77.05 with ESAM. However, [1] reports 77.5 with SAM in Table 2. Similarly, their results for ResNet-101 are better, and the results for PyramidNet on CIFAR are better.\n\nIs there a difference in the setting you use? \n\n## References\n\n[1] Sharpness-Aware Minimization for Efficiently Improving Generalization\nPierre Foret, Ariel Kleiner, Hossein Mobahi, Behnam Neyshabur", "summary_of_the_paper": "The paper proposes two simple modifications of the SAM optimizer that allow to significantly improve its computational efficiency without sacrificing performance.", "summary_of_the_review": "Overall, this is a very nice paper. My only concern is about the results relative to the original SAM paper. If the authors address this concern in the rebuttal, I am happy to recommend this paper for acceptance.", "strength_and_weaknesses": "", "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": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "X9UYPp3GZyS", "reviewer": "Reviewer_jSHm", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "Strengths\n+ simple modifications that improves over existing SAM results\n+ Wide range of baselines\n\nWeakness\n- lack of ablations across batch size (especially with the data-selection algorithm)                          \n\n== On cost of SAM ==\nCost of SAM is mentioned as 100% in the text of the paper. Typical SAM application when batch sizes are larger, involve data parallel training across a flock of GPUs. The efficient version of SAM (that works in practice) do not communicate the gradients across the cores, and hence cost is not exactly 100% but lower. It would be good to know results against this standard version of SAM.\n\n== Hyperparameter tuning == \nIt would be ideal if the authors normalized for the training cost. What would be the accuracy of SGD with 40% more epochs, with appropriate tuning of learning rate schedules?\n\n== Choice of gamma and batch size==\nIs the gamma=0.5 mean 50% of examples in a batch are discarded from the gradient step? This is quite fascinating that this method was able to improve accuracy while discarding half the batch! It would be extremely valuable to know (see comment on cost of SAM) to much larger batches (for example 2048 for CIFAR10/100) if this continues to remain true.\n", "summary_of_the_paper": "Paper proposes techniques to improve the efficiency of Sharpness aware minimization method. They are Stochastic Weight Perturbation (Select subset of the parameters at any step) and Sharpness-sensitive Data Selection.  Results demonstrates efficacy over SAM at small batch sizes on multiple models.\n\n", "summary_of_the_review": "Modifications are simple and efficient in practice, and may improve the practical usage of SAM across workloads.", "strength_and_weaknesses": "", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "SXI_ri9UjRi", "reviewer": "Reviewer_cAHr", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "Strengths:\n* This paper is well written and easy to follow.\n* While SAM has proven to be helpful for many applications, this technique can be helpful in practice.\n\nWeakness:\n*  For SWP, is the gradient mask the same for the samples in one batch, or different gradient masks are applied for each sample in the batch?\n* For SDS, an interesting baseline to be compared is estimating $\\epsilon$ with a subset of randomly sampled data, which could match the origin $\\epsilon$ in terms of empirical estimation.\n*  Recent paper [1] points that SAM can be very useful in ViT and MLP-Mixer. It would be more convincing if the author could show ESAM can also work on those two architectures.\n\nDetailed comments:\n* First row in Page 8,  \"As shwon in Table 3, SWP improves\". shwon -> shown\n\n[1] Chen, Xiangning, Cho-Jui Hsieh, and Boqing Gong. \"When Vision Transformers Outperform ResNets without Pretraining or Strong Data Augmentations.\" arXiv preprint arXiv:2106.01548 (2021).", "summary_of_the_paper": "This paper presents a method called ESAM for improving the efficiency of SAM. ESAM has two components: SWP and SDS. SWS accelerates the estimation of $\\epsilon$ via random sampling a subset of parameters in backpropagation. SDS further improves the efficiency via sampling a subset of data points that is enough for calculating the upper bound of $L$. Combining these two techniques, they achieve a speed improvement for SAM while yielding comparable or even better performance.", "summary_of_the_review": "This paper is the pioneer in exploring a more efficient SAM method. This technique can be useful since SAM has been proven to be useful for many networks. Therefore, I would recommend this paper as weak accept. I would be more convincing the practicability of this paper if the authors can show the proposed ESAM works on ViT and MLP mixer, where the SAM has been proven to be very effective.", "strength_and_weaknesses": "", "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": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "1dasJvCRdW0", "reviewer": "Reviewer_wp3y", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "Pros: \n1.\tThe proposed approach is well-motivated and practical in real-world applications. The enhanced efficiency does not lead to a degradation of accuracy. The experimental setup is satisfactory and convincing. The proposed method is verified on multiple datasets to show the effectiveness. The ablation studies and the supplementary experiments in the appendix are appropriate.\n2.\tThe idea of picking representative samples for updating parameters is interesting and effective. Why not use the $\\alpha$ directly as the hyperparameter for data selection?\n3.\tThe paper is well-organized. Figures 3 and 4 are good for demonstrating. \nCons:\n1.\tIt is mentioned that the computations of SWP slightly increase in the deeper DNNs. What is the exact impact of using SWP in large-scale deep neural networks? \n2.\tThe explanation of SWP's computation in Section 2.2 is not clear and may be made shorter. For example, L8 \"Decreasing $\\beta$ ... may degrade\" is superfluous; The claimed positive correlation between saved computations with $\\beta$ is vague and should be more precise. \n3.\tThe gaussian perturbations used in Figure 4 are not representative enough. There is no difference between the x-axis perturbation and the y-axis perturbation. Consider visualizing the loss landscape with the adversarial perturbations. \nAdditional comments:\n1.\tBetter to use \"hyperparameters of SGD\" to replace \"parameters of SGD\" in sec 3.1 line 6. \n2.\tBetter to add \"(Training speed 140.3% vs. SAM 100%)\" for the efficiency comparison in section 3.1 P3L4 \"(140.3% vs. SAM 100%)\".\n", "summary_of_the_paper": "This paper investigates the crucial efficiency issue of the sharpness-aware minimizer (SAM). SAM improves the generalization of DNNs but results in a double training time compared to vanilla training. By analyzing the min-max procedure of SAM, the authors observe the computational redundancy and then propose a method ESAM to improve the efficiency from the data and parameters perspectives. The authors argue that SAM can be approximated properly with fewer computations. Empirical results show that ESAM can reduce the extra training time in CIFAR10/100 and ImageNet datasets with improved accuracy compared to SAM. ", "summary_of_the_review": "This work raises an important question of SAM where the improved training algorithm requires doubled training cost. This paper addresses and improves the efficiency drawback of SAM. The proposed approach is well-motivated and has been verified to improve the efficiency and accuracy of SAM. The authors also provide codes for reproducibility. However, the degradation of SWP's effectiveness in large-scale neural networks, as well as the justifications for the improved accuracy contributed by SWP and SDS, have not been clearly addressed in this paper. I expect the author can answer the points above, and then I can adjust my final score.\n", "strength_and_weaknesses": "", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper focuses on improving the efficiency of sharpness-aware minimization method for training neural networks. The proposals are stochastic weight perturbation, namely selecting subset of the parameters at any step, and sharpness-sensitive data selection. The philosophy behind sounds quite interesting to me, namely, sharpness-aware minimizer can be approximated properly with fewer computations after analyzing the min-max procedure. This philosophy leads to a novel algorithm design I have never seen.

The clarity and novelty are clearly above the bar of ICLR. While the reviewers had some concerns on the significance, the authors did a particularly good job in their rebuttal. Thus, all of us have agreed to accept this paper for publication! Please include the additional experimental results in the next version.
Decision: Accept (Poster)