Year: 2021

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—Stochastic Weight 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 computational overhead to 40% vis-à-vis base optimizers, while test accuracies are preserved or even improved.

heading: 1 INTRODUCTION
text: Deep learning has achieved astounding performances in many fields by relying on larger numbers of parameters and increasingly sophisticated optimization algorithms. However, DNNs with far more parameters than training samples are more prone to poor generalization. Generalization is arguably the most fundamental and yet mysterious aspect of deep learning.
Several studies have been conducted to better understand the generalization of DNNs and to train DNNs that generalize well across the natural distribution (Keskar et al., 2016; Neyshabur et al., 2017; Chaudhari et al., 2019; Zhang et al., 2019; Wu et al., 2020; Foret et al., 2020; Zhang et al., 2021). For example, Keskar et al. (2016) investigate the effect of batch size on neural networks’ generalization ability. Zhang et al. (2019) propose an optimizer for training DNNs with improved generalization ability. Specifically, Hochreiter & Schmidhuber (1995), Li et al. (2017) and Dinh et al. (2017) argue that the geometry of the loss landscape affects generalization and DNNs with a 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 a base optimizer such as Stochastic Gradient Descent (Nesterov, 1983) or Adam (Kingma & Ba, 2014) to minimize both the vanilla training loss and the sharpness. The sharpness, which describes the flatness of a minimum, is characterized using eigenvalues of the Hessian matrix by Keskar et al. (2016). SAM quantifies the sharpness as the maximized change of training loss when a constraint perturbation is added to current weights. As a result, SAM leads to a 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 a 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 variants achieve such remarkable performance at the expense of doubling the computational overhead of the given base optimizers, which minimize the training loss with a 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 computational overhead being doubled (2×). As demonstrated in Figure 1, SAM achieves higher test accuracy (i.e., 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).
SGD SAM ESAM 80
81
82
83
84
85
86
A cc ur ac y (%
)
100
200
300
400
500
600
700
Tr ai ni ng
S pe
ed (i m ag
es /s )
Accuracy Training Speed
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.
We verify the effectiveness of ESAM on the CIFAR10, CIFAR100 (Krizhevsky et al., 2009) and ImageNet (Deng et al., 2009) datasets with five different DNN architectures. The experimental results demonstrate that ESAM obtains flat minima at a cost of only 40% (vs. SAM’s 100%) extra computational overhead over base optimizers. More importantly, ESAM achieves better performance in terms of the test accuracy compared to SAM. In a nutshell, our contributions are as follows:
• We propose two novel and effective training strategies Stochastic Weight Perturbation (SWP) and Sharpness-sensitive Data Selection (SDS). Both strategies are designed to improve efficiency without sacrificing performance. The empirical results demonstrate that both of the proposed strategies can improve both the efficiency and effectiveness of SAM.
• We introduce the ESAM, which integrates SWP and SDS. ESAM improves the generalization 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 are designed respectively. Section 3 verifies the effectiveness of ESAM across a variety of datasets and DNN architectures. Section 4 presents the related work and Section 5 concludes this paper.

heading: 2 METHODOLOGY
text: We start with recapitulating how SAM achieves a flat minimum with small sharpness, which is quantified by resolving a maximization problem. To compute the sharpness, SAM requires additional forward and backward propagation and results in the doubling of the computational overhead 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 introduce SWP and SDS in Sections 2.2 and 2.3 respectively. Algorithm 1 shows the overall proposed ESAM algorithm.
Throughout this paper, we denote a neural network f with weight parameters θ as fθ. The weights are contained in the vector θ = (θ1, θ2, . . . , θN ), where N is the number of weight units in the neural network. Given a training dataset S that contains samples i.i.d. drawn from a distribution D,
Algorithm 1 Efficient SAM (ESAM) Input: Network fθ, θ = (θ1, θ2, . . . , θN ), Training set S, Batch size b, Learning rate η > 0,
Neighborhood size ρ > 0, Iterations A, SWP hyperparameter β, SDS hyperparameter γ. Output: A flat minimum solution θ̂.
1: for a = 1 to A do 2: Sample a mini-batch B ⊂ S with size b. 3: for n = 1 to N do 4: if θn is chosen by probability β then 5: n ← ρ1−β∇θnLB(fθ) . SWP in B1 6: else 7: n ← 0 8: ̂← ( 1, ..., N ) . Assign Weight Perturbation 9: Compute `(fθ+̂, xi, yi) and construct B+ with selection ratio γ (Equation 6) 10: Compute gradients g = ∇θLB+(fθ+̂) . SDS in B2 11: Update weights θ ← θ − ηg
the network is trained to obtain optimal weights θ̂ via empirical risk minimization (ERM), i.e.,
θ̂ = arg min θ
{ LS(fθ) = 1 |S| ∑
(xi,yi)∈S
`(fθ, xi, yi)
} (1)
where ` can be an arbitrary loss function. We take ` to be the cross entropy loss in this paper. The population loss is defined as LD(fθ) , E(xi,yi)∼D(`(fθ, xi, yi)). In each training iteration, optimizers sample a mini-batch B ⊂ S with size b to update parameters.

heading: 2.1 SHARPNESS-AWARE MINIMIZATION AND ITS COMPUTATIONAL DRAWBACK
text: 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 optimization problem,
min θ max :‖ ‖2≤ρ LS(fθ+ ). (2)
Given θ, the inner optimization attempts to find a 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 is defined to be RS(fθ) = max :‖ ‖2<ρ[LS(fθ+ ) − LS(fθ)]. This sharpness is quantified by the maximal change of empirical loss when a constraint perturbation is added to θ. The min-max problem encourages SAM to find flat minima.
For a certain set of weights θ, 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 et al., 2021).
LD(fθ) ≤ RS(fθ) + LS(fθ) + λ‖θ‖22 = max :‖ ‖2≤ρ LS(fθ+ ) + λ‖θ‖22 (3)
In practice, SAM first approximately solves the inner optimization by means of a single-step gradient descent method, i.e.,
̂ = arg max :‖ ‖2<ρ
LS(fθ+ ) ≈ ρ∇θLS(fθ). (4)
The sharpness at weights θ is approximated by RS(fθ) = LS(fθ+̂) − LS(fθ). Then, a base optimizer, such as SGD (Nesterov, 1983) or Adam (Kingma & Ba, 2014), updates the DNNs’ weights to minimize LS(fθ+̂). We refer to LS(fθ+̂) as the SAM loss. Overall, SAM requires two forward and two backward operations to update weights once. We refer to the forward and backward propagation for approximating ̂ as F1 and B1 and those for updating weights by base optimizers as F2 and B2
respectively. Although SAM can effectively improve the generalization of DNNs, it additionally requires one forward and one backward operation (F1 and B1) in each training iteration. Thus, SAM results in a doubling of the computational overhead compared to the use of base optimizers.
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. Specifically, on the one hand, when estimating ̂ around weight vector θ, SWP efficiently approximates ̂ by randomly selecting each parameter with a given probability to form a 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 LB(fθ+̂) over a whole batch of samples, SDS selects a subset of samples, B+, 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 a linear fashion). We further justify that LB(fθ+̂) can be upper bounded by LB+(fθ+̂) and consequently the generalization capability can be preserved. In general, ESAM works much more efficiently and performs as well as SAM in terms of its generalization capability.

heading: 2.2 STOCHASTIC WEIGHT PERTURBATION
text: 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 ̂(θ,S) during the sharpness estimation phase, SWP randomly chooses a subset θ̃ = {θI1 , θI2 , . . .} from the original set of weights θ = (θ1, . . . , θN ) to perform backpropagation B1. Each parameter is selected to be in the subvector θ̃ with some probability β, which can be tuned as a hyperparameter. SWP approximates the weight perturbation with ρ∇θ̃LS(fθ). To be formal, we introduce a gradient mask m = (m1, . . . ,mN ) where mi
i.i.d.∼ Bern(β) for all i ∈ {1, . . . , N}. Then, we have ρ∇θ̃LS(fθ) = m>̂(θ,B). To ensure the expected weight perturbation of SWP equals to ̂, we scale ρ∇θ̃LS(fθ) by a factor of 1 β . Finally, SWP produces an approximate solution of the inner maximization as
a(θ,B) = m>̂(θ,B)
β . (5)
Computation Ideally, SWP reduces the overall computational overhead in proportion to 1 − β 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 propagation 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 − β. In practice, β is tuned to maximize SWP’s efficiency while maintaining a 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(θ,B) of SWP equals to the original SAM’s perturbation ̂(θ,B) in the sense of the `2 norm and direction. We denote the expected SWP perturbation by ā(θ,B), where
ā(θ,B)[i] = E[a(θ,B)[i]] = 1
β · β̂(θ,B)[i] = ̂(θ,B)[i],
for i ∈ {1, . . . , N}. Thus, it holds that ‖ā(θ,B)‖2 = ‖̂(θ,B)‖2 and CosSim ( ā(θ,B), ̂(θ,B) ) = 1,
showing that the expected weight perturbation of SWP is the same as that of SAM’s.
2.3 SHARPNESS-SENSITIVE DATA SELECTION
In this section, we introduce the second efficiency enhancement 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.
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 LB(fθ+̂), we select a subset of samples from the whole batch. The loss values of this subset of samples increase most when the weights are perturbed by ̂. To be specific, SDS splits the mini-batch B into the following two subsets
B+ := { (xi, yi) ∈ B : `(fθ+̂, xi, yi)− `(fθ, xi, yi) > α } ,
B− := { (xi, yi) ∈ B : `(fθ+̂, xi, yi)− `(fθ, xi, yi) < α } , (6)
where B+ is termed as the sharpness-sensitive subset and the threshold α controls the size of B+. We let γ = |B+|/|B| be the ratio of the number of selected samples with respect to the batch size. In practice, γ determines the exact value of α and serves as a predefined hyperparameter of SDS. As illustrated in Figure 2, when α = 0, the gradient of the weights evaluated on B+ aligns with the direction of ̂ and the loss values of the samples in B+ will increase with respect to the weight perturbation ̂.
Computation SDS reduces the computational overhead in F2 and B2. The reduction is linear in 1 − γ. The hyperparameter γ can be tuned to meet up distinct requirements in efficiency and performance. SDS is configured the same as SWP for maximizing efficiency with comparable performance to SAM.
Generalization For the generalization capability, we now justify that the SAM loss computing over the batch B, LB(fθ+̂), can be approximately upper bounded by the corresponding loss evaluated only on B+, LB+(fθ+̂). From Equation 3, we have
LB(fθ+̂) = γLB+(fθ+̂) + (1− γ)LB−(fθ+̂) = LB+(fθ+̂) + (1− γ)[RB−(fθ)−RB+(fθ) + LB−(fθ)− LB+(fθ)].
(7)
On the one hand, since RB(fθ) = 1|B| ∑
(xi,yi)∈B[`(fθ+̂, xi, yi)− `(fθ, xi, yi)] represents the average sharpness of the batch B, by Equation 6, we have RB−(fθ) ≤ RB(fθ) ≤ RB+(fθ), and
RB−(fθ)−RB+(fθ) ≤ 0. (8)
On the other hand, B+ and B− are constructed sorting by `(fθ+̂, xi, yi) − `(fθ, xi, yi), which is positively correlated to l(fθ, xi, yi) (Li et al., 2019) (more details can be found in the Appendix A.2). Thus, we have
LB−(fθ)− LB+(fθ) ≤ 0. (9)
Therefore, by Equation 8 and Equation 9, we have
LB(fθ+̂) ≤ LB+(fθ+̂). (10)
Experimental results in Figure 5 corroborate thatRB−(fθ)−RB+(fθ) < 0 andLB−(fθ)−LB+(fθ) < 0. Besides, Figure 6 verifies that the selected batch B+ is sufficiently representative to mimic the gradients of B since B+ has a significantly higher cosine similarity with B compared to B− in terms of the computed gradients. According to Equation 10, one can utilize LB+(fθ+̂) as a proxy to the real objective to minimize of the overall loss LB(fθ+̂) with a smaller number of samples. As a result, SDS improves SAM’s efficiency without performance degradation.

heading: 3 EXPERIMENTS
text: This section demonstrates the effectiveness of our proposed ESAM algorithm. We conduct experiments on several benchmark datasets: CIFAR-10 (Krizhevsky et al., 2009), CIFAR-100 (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 Table 1 and Table 2. Besides, we 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.

heading: 3.1 RESULTS
text: CIFAR10 and CIFAR100 We start from evaluating ESAM on the CIFAR-10 and CIFAR-100 image classification datasets. The evaluation is carried out on three different model architectures: ResNet18 (He et al., 2016), WideResNet-28-10 (Zagoruyko & Komodakis, 2016) and PyramidNet-110 (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, 2016). The training epochs are set to be 200 for ResNet-18 (He et al., 2016), WideResNet-28-10 (Zagoruyko & Komodakis, 2016), and 300 for PyramidNet-110 (Han et al., 2017). We set β = 0.6 and γ = 0.5 for ResNet-18 and PyramidNet110 models; and set β = 0.5 and γ = 0.5 for WideResNet-28-10. The above-mentioned β 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 can 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 CIFAR10 trained with ResNet-18 (Training speed 140.3% vs. SAM 100%). The best accuracy is reported in CIFAR100 trained with PyramidNet110 (Accuracy 85.56% vs. SAM 84.46%). ESAM improves efficiency and achieves better performance compared to SAM in CIFAR10/100 benchmarks.
ImageNet To evaluate ESAM’s effectiveness on a large-scale benchmark dataset, we conduct experiments on ImageNet Datasets.The 1000 class ImageNet dataset contains roughly 1.28 million
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 effectiveness) than CIFAR datasets. We resize the images on ImageNet to 224× 224 resolution to train ResNet-50 and ResNet-110 models. We train 90 epochs and set the optimal hyperparameters for SGD, SAM and ESAM as suggested by Chen et al. (2021), and the details are listed in appendix A. We use β = 0.6 and γ = 0.7 for ResNet-50 and ResNet-101. We employ the m-sharpness strategy for both SAM and ESAM with m = 128, which is the same as that suggested in Zheng et 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 γ we used here is 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.

heading: 3.2 ABLATION AND PARAMETER STUDIES
text: To better understand the effectiveness 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 ResNet-18 and WideResNet-28-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 [0.3, 0.9] for γ, with a same step size of 0.1. We report the grid search results in Figure 3. We use β = 0.6,γ = 0.5 for ResNet-18; and set β = 0.5, γ = 0.5 for WideResNet-28-10 in the four sets of ablation studies. The ablation study results are reported in Table 3.
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 β. The best performance of SWP is obtained when β = 0.6 for ResNet-18 and β = 0.5 for WideResNet-28-10. The four sets of experiments indicate that β is consistent among different architectures and datasets. Therefore, we set β = 0.6 for PyramidNet in CIFAR10/100 datasets and ResNet on ImageNet datasets.
ESAM-SDS SDS also significantly improves the efficiency by 21.5% to 25.8% compared to SAM. It outperforms SAM’s performance in CIFAR100 datasets, and achieves comparable performance in CIFAR10 datasets. SDS can outperform SAM in 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. (2017). The x- and yaxes 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. It 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 compared to SAM. Most importantly, both these strategies achieve a comparable or better performance than SAM. In practice, by configuring the β and γ, ESAM can meet a variety of user-defined efficiency and performance requirements.

heading: 4 RELATED WORK
text: The concept of regularizing sharpness for better generalization dates back to (Hochreiter & Schmidhuber, 1995). By using an MDL-based argument, which clarifies that a statistical model with fewer
bits to describe can have better generalization ability, Hochreiter & Schmidhuber (1995) claim that a flat minimum can alleviate overfitting issues. Following that, more studies were proposed to investigate the connection between the flat minima with the generalization abilities (Keskar et al., 2016; Dinh et al., 2017; Liu et al., 2020; Li et al., 2017; Dziugaite & Roy, 2017; Jiang et al., 2019; Moosavi-Dezfooli et al., 2019). Keskar et al. (2016) 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. (2016) and Dinh et al. (2017) both argue that the sharpness can be characterized using the eigenvalues 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. (2016) 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-the-art performance in a variety of deep learning benchmarks, including image classification, natural language processing, and noisy learning (Foret et al., 2020; Chen et 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 contemporaneously 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., 2021; Tseng et al., 2021), or apply SAM on other tasks (Zheng et al., 2021; Damian et al., 2021; Galatolo et al., 2021). Kwon et al. (2021) improves SAM’s sharpness by adaptively scaling the size of the nearby search space ρ in relation to the size of parameters. However, most of these works overlook the fact that SAM improves generalization at the expense of the doubling the computational overhead. As a result, most of the SAM-related works suffer from the same efficiency drawback as SAM. This computational cost prevents SAM from being widely used in large-scale datasets and architectures, particularly in real-world applications, which motivates us to propose ESAM to efficiently improve the generalization ability of DNNs.

heading: 5 CONCLUSION
text: 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-propagation is performed to compute the weight perturbation or 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 as base optimizers.
Reproducibility Statement To ensure the reproducibility of our proposed ESAM formulation and algorithm, we will submit our source codes implemented in the PyTorch framework (Paszke et al., 2017) together with the supplementary material. The source codes contains the ESAM training algorithm codes, documents on how to apply ESAM, and examples of the application of ESAM on the CIFAR benchmark datasets. We will also release our ESAM codes that embedded into TIMM (Wightman, 2019) framework to the public.

heading: A APPENDIX

heading: A.1 THE ALGORITHM OF SAM
text: The algorithm of SAM is demonstrated in Algorithm 2.
Algorithm 2 SGD vs. SAM Input: Network fθ, θ = (θ1, θ1, . . . , θN ), Training set S, Batch size b, Learning rate η, Neighbor-
hood size ρ, Iterations A. Output: A minimum solution θ̃.
1: for a = 1 to A do 2: Sample B with size b that B ⊂ S, 3: if SGD then 4: ̂← 0 . No addtional computational overhead for ̂ 5: else if SAM then 6: ̂← ∇θLB(fθ) . additional F1 and B1 to compute ̂ 7: Compute g = ∇θLB(fθ+̂) . F2 and B2 8: update the weights θ ← θ − ηg

heading: A.2 OPTIMIZING OVER SUBSET B+ IS REPRESENTATIVE
text: The sharpness-sensitive subset B+ are constructed sorting by `(fθ+̂, xi, yi) − `(fθ, xi, yi), which is positively correlated to l(fθ, xi, yi). By the first-order Taylor series approximation,
`(fθ+̂, xi, yi)− `(fθ, xi, yi) = ̂ · ∇θ`(fθ, xi, yi) + o(‖̂‖)
By Equation 4, ̂ is the aggregated gradients of each instance in the complete dataset B, i.e.,
̂ = arg max :‖ ‖2<ρ LS(fθ+ ) ≈ ρ∇θLS(fθ) = |B|∑ i=1 ∇θ`(fθ, xi, yi)
which indicates that `(fθ+̂, xi, yi) − `(fθ, xi, yi) is positively correlated to the gradient ∇θ`(fθ, xi, yi). Li et al. (2019) claims that the hard examples in deep learning (the training samples with high training loss) produce gradients with larger magnitudes. Therefore, `(fθ+̂, xi, yi)− `(fθ, xi, yi) is positively correlated to l(fθ, xi, yi). We also demonstrate the correlation empirically.
We conduct experiments to verify Equation 9 and Equation 10. In Figure 5, We plot the four losses, LB+(fθ), LB−(fθ), LB+(fθ+̂), and LB−(fθ+̂) w.r.t 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 B+ is much more representative than the subset B−. We compare the updating gradients computed from B+,B− and a random subset Brand that |Brand| = |B+| = |B−| to those computed from B by calculating the cosine similarity, i.e.
CosSim(∇θLB+(fθ+̂),∇θLB(fθ+̂)) CosSim(∇θLB−(fθ+̂),∇θLB(fθ+̂)) CosSim(∇θLBrand(fθ+̂),∇θLB(fθ+̂))
In Figure 6, we plot the cosine similarities in each training epoch evaluated with ResNet-18, Wide28-10 on CIFAR10. In terms of the computed gradients, the experimental results show that B+ has the highest cosine similarities with B than B− and the random set Brand.

heading: A.3 LINEARITY MEASUREMENT OF SWP
text: As the experimental results in Figure 3 demonstrated, SWP can also improve the accuracy of ESAM compared to SAM. We will investigate the advantage of SWP in terms of generalization in the future research. Here we provide a discussion about the accuracy improvement contributed by SWP.
A plausible reason for such improvement is that SWP leads to a better inner maximization solved in equation 2. The current solution of ̂ is approximated by assuming LS(fθ) is a linear function. Therefore, the ̂ would result in a better inner maximization if LS(fθ) is more linear with respect to θ. Inspired by (Qin et al., 2019), we measure the linearity of the loss function by
ζ( ,B) = |LB(fθ+ )− LB(fθ)− >∇θLB(fθ)|
We conduct experiments on the CIFAR10 dataset with ResNet-18 model to verify that SWP can improve the linearity ζ( ,B) of the loss function LS(fθ). We compare the linearity of ESAM with the β ranging from {0.2, 0.3, ..., 0.9} to the SAM. The results are demonstrated in Figure 7. It can
be shown that SWP will result in a better linearity as the β decreases. However, decreasing β will also reduce the magnitude of ̂ and thus result in a worse inner maximization in equation 2. We observe that β = {0.5, 0.6} is optimal to balance the accuracy and efficiency of ESAM.

heading: A.4 REDUCED COMPUTATIONAL OVERHEAD CONTRIBUTED BY SWP
text: Formulation We formulate the saved computational overhead contributed by SWP here. We discuss and examine the computational overhead in the PyTorch framework. Suppose that the NN we discussed here has N layers. The most unit of parameters (D × C ×H ×W ) is the entire parameters of a layer in NN, where D is the number of kernels, C is the number of channels, H and W are the height and width of the input. SWP select each basic parameter unit to compute gradients (i.e. requries grad = True) with probability β. We use g(N, β) to measure the saved computational overhead in terms of percentage contributed by SWP compared to the vanilla SAM.
The computational overhead g(N, β) 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 to each other, the computational overhead saved from storing and hooking are proportional to 1− β 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 as DenseNet and ViT, may have multiple basic parameter units in the same layer and be connected across any other layers. Suppose the nth layer hasK(n) basic parameter units, we have pn = (1−β)K(n) ·pn−1. Therefore, pn = (1− β) ∑n j=1 K(j) ≈ (1− β)nK̄ , where K̄ = 1N ∑N j=1K(j). Assumed that the computations
of each layer is the same, by summing up the saved computational overhead of all parameters, we have
g(N, β) = k1(1−β)+k2 N∑ n=i pn N ≈ k1(1−β)+ k2(1− β) N [1− (1− β)K̄ ] ≈ k1(1−β)+ k ′ 2(1− β) Nβ (11)
where k1, k2 are determined by the computing time of calculation, k ′ 2 = k2β [1−(1−β)K̄ ] .
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, let pn be the calculation-free rate of a certain parameter unit in nth layer, we have pn = (1−β) · pn−1. Therefore, pn = (1 − β)n. The saved calculation contributed by the parameter unit is 1N p
n. By summing up the saved computational overhead of all parameters, we have
g(N, β) = k1(1− β) + k2 ∑N n=i pn N = k1(1− β) + k2 N 1−β−(1−β)N β ≈ k1(1− β) + k2(1−β) Nβ (12)
where k1, k2 are determined by the computing time of calculation, storing and hooking gradients.
Toy Example We conducted a toy example in Cifar10 dataset with two MLPs. Each fully connected layer of MLP is the same with a size of 3, 000 for both in and out features. The first mlp is for the special case that K̄ = 1, where each layer has only one basic parameter unit and is connected in serial with the next layer. We examined N = {50, 75, 100, 125} and β = {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 k1 = 0.3185, k2 = 0.1310, and the returned R2 = 0.9983. The second mlp 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 examinedN = {35, 50, 65, 75} and β = {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 k1 = 0.3143, k2 = 0.0737, and the returned R2 = 0.9989. The above experimental results verify the formulation of the reduced computational overhead contributed by SWP in equation 12 and equation 11.
A.5 VISUALIZATION OF LOSS LANDSCAPES WITH RESPECT TO ADVERSARIAL WEIGHT PERTURBATIONS
We visualize the sharpness of the flat minima with respect to adversarial weight perturbations of SGD,SAM and ESAM on the Cifar10 dataset. The x- and y-axes represent two orthogonal adversarial weight perturbations, which are η∇θLBx(fθ) and η∇θLBy (fθ) respectively, where η is the learning rate during training. Bx and By are the randomly sampled subsets of batch B, and |Bx| = |By| = 12 |B|, Bx ∪ By = B. We display the loss landscape in Figure 8, which demonstrates that both SAM and ESAM improve the sharpness significantly in comparison to SGD.

heading: A.6 EVALUATION OF ESAM ON VIT-S/16
text: SAM has also been demonstrated to be effective on the new vision Transformer(ViT) architecture (Chen et al., 2021). Therefore, we also evaluate ESAM with ViT-S/16 on ImageNet Datasets. We use β = 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.88% v.s. 80.34%).

heading: A.7 TRAINING DETAILS
text: 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 × 10−4, 1 × 10−3}, and the batch size from the set {64, 128, 256}. This is done to 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 8.

Review summaries:
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.  

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.

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.