Original: E FFICIENT S HARPNESS - AWARE M INIMIZATION FOR I MPROVED T RAINING OF N EURAL N ETWORKS
Revised: E FFICIENT S HARPNESS - AWARE M INIMIZATION FOR I MPROVED T RAINING OF N EURAL N ETWORKS

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Overparametrized Deep Neural Networks (DNNs) often achieve astounding performances, but may potentially result in severe generalization error.
Revised: Overparametrized Deep Neural Networks (DNNs) often achieve astounding performances but may potentially result in severe generalization error.

Original: Recently, the relation between the sharpness of the loss landscape and the generalization error has been established by Foret et al.
Revised: Recently, the relation between the sharpness of the loss landscape and the generalization error has been established by Foret et al.

Original: (2020), in which the Sharpness Aware
Revised: (2020) in which the Sharpness Aware

Original: Minimizer (SAM) was proposed to mitigate the degradation of the generalization.
Revised: Minimizer (SAM) was proposed to mitigate the degradation of the generalization.

Original: Unfortunately, SAM’s computational cost is roughly double that of base optimizers, such as Stochastic Gradient Descent (SGD).
Revised: Unfortunately, SAM’s computational cost is roughly double that of base optimizers such as Stochastic Gradient Descent (SGD).

Original: This paper thus proposes
Revised: This paper thus proposes

Original: Efﬁcient Sharpness Aware Minimizer (ESAM), which boosts SAM’s efﬁciency at no cost to its generalization performance.
Revised: Efﬁcient Sharpness Aware Minimizer (ESAM) which boosts SAM’s efﬁciency at no cost to its generalization performance.

Original: ESAM includes two novel and efﬁcient training strategies—Stochastic Weight Perturbation and Sharpness-Sensitive Data
Revised: ESAM includes two novel and efﬁcient training strategies—Stochastic Weight Perturbation and Sharpness-Sensitive Data

Original: Selection.
Revised: Selection.

Original: 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.
Revised: 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.

Original: We provide theoretical explanations as to why these strategies perform well.
Revised: We provide theoretical explanations as to why these strategies perform well.

Original: We also show, via extensive experiments on the CIFAR and ImageNet datasets, that ESAM enhances the efﬁciency over SAM from requiring 100% extra computational overhead to 40% vis-`a-vis base optimizers, while test accuracies are preserved or even improved.
Revised: We also show, via extensive experiments on the CIFAR and ImageNet datasets, that ESAM enhances the efﬁciency over SAM from requiring 100% extra computational overhead to 40% vis- ` a-vis base optimizers, while test accuracies are preserved or even improved

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Deep learning has achieved astounding performances in many ﬁelds by relying on larger numbers of parameters and increasingly sophisticated optimization algorithms.
Revised: Deep learning has achieved astounding performances in many ﬁelds by relying on larger numbers of parameters and increasingly sophisticated optimization algorithms.

Original: However, DNNs with far more parameters than training samples are more prone to poor generalization.
Revised: However, DNNs with far more parameters than training samples are more prone to poor generalization.

Original: Generalization is arguably the most fundamental and yet mysterious aspect of deep learning.
Revised: Generalization is arguably the most fundamental and yet mysterious aspect of deep learning.

Original: Several studies have been conducted to better understand the generalization of DNNs and to train
Revised: Several studies have been conducted to better understand the generalization of DNNs and to train

Original: 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).
Revised: DNNs that generalize well across the natural distribution (Keskar et al., 2017; Neyshabur et al., 2017; Chaudhari et al., 2019; Zhang et al., 2019; Wu et al., 2020; Foret et al., 2020; Zhang et al., 2021).

Original: For example, Keskar et al.
Revised: For example, Keskar et al.

Original: (2016) investigate the effect of batch size on neural networks’ generalization ability.
Revised: (2017) investigate the effect of batch size on neural networks’ generalization ability.

Original: Zhang et al.
Revised: Zhang et al.

Revised: ; Zhou et al.

Original: (2019) propose an optimizer for training DNNs with improved generalization ability.
Revised: (2021) propose optimizers for training DNNs with improved generalization ability.

Original: Speciﬁcally, Hochreiter & Schmidhuber (1995), Li et al.
Revised: Speciﬁcally, Hochreiter & Schmidhuber (1995), Li et al.

Original: (2017) and Dinh et al.
Revised: and Dinh et al.

Original: (2017) argue that the geometry of the loss landscape affects generalization and DNNs with a ﬂat minimum can generalize better.
Revised: (2017) argue that the geometry of the loss landscape affects generalization and DNNs with a ﬂat minimum can generalize better.

Original: The recent work by Foret et al.
Revised: The recent work by Foret et al.

Original: (2020) proposes an effective training algorithm Sharpness Aware Minimizer (SAM) for obtaining a ﬂat minimum.
Revised: (2020) proposes an effective training algorithm Sharpness Aware Minimizer (SAM) for obtaining a ﬂat minimum.

Original: 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.
Revised: SAM employs a base optimizer such as Stochastic Gradient Descent (Nesterov, 1983) or Adam (Kingma & Ba, 2015) to minimize both the vanilla training loss and the sharpness.

Original: The sharpness, which describes the ﬂatness of a minimum, is characterized using eigenvalues of the Hessian matrix by Keskar et al.
Revised: The sharpness, which describes the ﬂatness of a minimum, is characterized using eigenvalues of the Hessian matrix by Keskar et al.

Original: SAM quantiﬁes the sharpness as the maximized change of training loss when a constraint perturbation is added to current weights.
Revised: SAM quantiﬁes the sharpness as the maximized change of training loss when a constraint perturbation is added to current weights.

Original: As a result, SAM leads to a ﬂat minimum and signiﬁcantly improves the generalization ability of the trained DNNs.
Revised: As a result, SAM leads to a ﬂat minimum and signiﬁcantly improves the generalization ability of the trained DNNs.

Original: 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;
Revised: 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;

Original: Chen et al., 2021; Galatolo et al., 2021; Zheng et al., 2021).
Revised: Chen et al., 2021; Galatolo et al., 2021; Zheng et al., 2021).

Original: 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.
Revised: 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.

Original: SAM requires an additional propagation step compared to the base optimizers to resolve the weight perturbation for quantifying the sharpness.
Revised: SAM requires an additional propagation step compared to the base optimizers to resolve the weight perturbation for quantifying the sharpness.

Original: 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 × ).
Revised: 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 × ).

Original: As demonstrated in Figure 1, SAM achieves higher test accuracy (i.e., 84 .
Revised: As demonstrated in Figure 1, SAM achieves higher test accuracy (i.e., 84 .

Original: 46% vs. 81 . 89% ) at the expense of sacriﬁcing half of the training speed of the base optimizer (i.e.,imgs/s vs. 557 imgs/s).
Revised: 46% vs. 81 . 89% ) at the expense of sacriﬁcing half of the training speed of the base optimizer (i.e., 276 imgs/s vs. 557 imgs/s).

Original: [Figure]
Revised: [Figure]

Original: In this paper, we aim to improve the efﬁciency of
Revised: In this paper, we aim to improve the efﬁciency of

Original: SAM but preserve its superior performance in generalization.
Revised: SAM but preserve its superior performance in generalization.

Original: We propose Efﬁcient Sharpness Aware Minimizer (ESAM) , which consists of two training strategies
Revised: We propose Efﬁcient Sharpness Aware Minimizer (ESAM) , which consists of two training strategies

Original: Stochastic Weight Perturbation (SWP) and Sharpnesssensitive Data Selection (SDS) , both of which reduce computational overhead and preserve the performance of
Revised: Stochastic Weight Perturbation (SWP) and Sharpnesssensitive Data Selection (SDS) , both of which reduce computational overhead and preserve the performance of

Original: SAM.
Revised: SAM.

Original: On the one hand, SWP approximates the sharpness by searching weight perturbation within a stochastically chosen neighborhood of the current weights.
Revised: On the one hand, SWP approximates the sharpness by searching weight perturbation within a stochastically chosen neighborhood of the current weights.

Original: SWP preserves the performance by ensuring that the expected weight perturbation is identical to that solved by SAM.
Revised: SWP preserves the performance by ensuring that the expected weight perturbation is identical to that solved by SAM.

Original: On the other hand, SDS improves efﬁciency by approximately optimizing weights based on the sharpnesssensitive subsets of batches.
Revised: On the other hand, SDS improves efﬁciency by approximately optimizing weights based on the sharpnesssensitive subsets of batches.

Original: These subsets consist of samples whose loss values increase most w.r.t.
Revised: These subsets consist of samples whose loss values increase most w.r.t.

Original: the weight perturbation and consequently can better quantify the sharpness of DNNs.
Revised: the weight perturbation and consequently can better quantify the sharpness of DNNs.

Original: 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.
Revised: 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.

Original: We verify the effectiveness of ESAM on the CIFAR10, CIFAR100 (Krizhevsky et al., 2009) and ImageNet (Deng et al., 2009) datasets with ﬁve different DNN architectures.
Revised: We verify the effectiveness of ESAM on the CIFAR10, CIFAR100 (Krizhevsky et al., 2009) and ImageNet (Deng et al., 2009) datasets with ﬁve different DNN architectures.

Original: The experimental results demonstrate that ESAM obtains ﬂat minima at a cost of only 40% (vs. SAM’s 100%) extra computational overhead over base optimizers.
Revised: The experimental results demonstrate that ESAM obtains ﬂat minima at a cost of only 40% (vs. SAM’s 100%) extra computational overhead over base optimizers.

Original: More importantly, ESAM achieves better performance in terms of the test accuracy compared to SAM.
Revised: More importantly, ESAM achieves better performance in terms of the test accuracy compared to SAM.

Original: In a nutshell, our contributions are as follows:
Revised: In a nutshell, our contributions are as follows:

Original: • We propose two novel and effective training strategies Stochastic Weight Perturbation (SWP) and Sharpness-sensitive Data Selection (SDS) .
Revised: • We propose two novel and effective training strategies Stochastic Weight Perturbation (SWP) and Sharpness-sensitive Data Selection (SDS) .

Original: Both strategies are designed to improve efﬁciency without sacriﬁcing performance.
Revised: Both strategies are designed to improve efﬁciency without sacriﬁcing performance.

Original: The empirical results demonstrate that both of the proposed strategies can improve both the efﬁciency and effectiveness of SAM.
Revised: The empirical results demonstrate that both of the proposed strategies can improve both the efﬁciency and effectiveness of SAM.

Original: • We introduce the ESAM, which integrates SWP and SDS.
Revised: • We introduce the ESAM, which integrates SWP and SDS.

Original: ESAM improves the generalization ability of DNNs with marginally additional computational cost compared to standard training.
Revised: ESAM improves the generalization ability of DNNs with marginally additional computational cost compared to standard training.

Original: The rest of this paper is structured in this way.
Revised: The rest of this paper is structured in this way.

Original: Section 2.1 introduces SAM and its computational issues.
Revised: Section 2.1 introduces SAM and its computational issues.

Original: Section 2.2 and Section 2.3 discuss how the two proposed training strategies SWP and SDS are designed respectively.
Revised: Section 2.2 and Section 2.3 discuss how the two proposed training strategies SWP and SDS are designed respectively.

Original: Section 3 veriﬁes the effectiveness of ESAM across a variety of datasets and DNN architectures.
Revised: Section 3 veriﬁes the effectiveness of ESAM across a variety of datasets and DNN architectures.

Original: Section 4 presents the related work and Section 5 concludes this paper.
Revised: Section 4 presents the related work and Section 5 concludes this paper.

Original:  M ETHODOLOGY
Revised:  M ETHODOLOGY

Original: We start with recapitulating how SAM achieves a ﬂat minimum with small sharpness, which is quantiﬁed by resolving a maximization problem.
Revised: We start with recapitulating how SAM achieves a ﬂat minimum with small sharpness, which is quantiﬁed by resolving a maximization problem.

Original: To compute the sharpness, SAM requires additional forward and backward propagation and results in the doubling of the computational overhead compared to base optimizers.
Revised: To compute the sharpness, SAM requires additional forward and backward propagation and results in the doubling of the computational overhead compared to base optimizers.

Original: Following that, we demonstrate how we derive and propose ESAM, which integrates SWP and SDS, to maximize efﬁciency while maintaining the performance.
Revised: Following that, we demonstrate how we derive and propose ESAM, (cid:15) n ← ρ 1 − β ∇ θ n L B ( f θ ) 8: ˆ ←

Revised: Compute

Revised: Compute gradients g = ∇ θ L B + f θ +ˆ (cid:15)

Revised: Update weights θ ← θ − ηg which integrates SWP and SDS, to maximize efﬁciency while maintaining the performance.

Original: We introduce SWP and SDS in Sections 2.2 and 2.3 respectively.
Revised: We introduce SWP and SDS in Sections 2.2 and 2.3 respectively.

Original: Algorithm 1 shows the overall proposed
Revised: Algorithm 1 shows the overall proposed

Original: ESAM algorithm.
Revised: ESAM algorithm.

Original: Throughout this paper, we denote a neural network f with weight parameters θ as f θ .
Revised: Throughout this paper, we denote a neural network f with weight parameters θ as f θ .

Original: The weights are contained in the vector θ =
Revised: The weights are contained in the vector θ =

Original: ( θ 1 , θ 2 , . .
Revised: ( θ 1 , θ 2 , . .

Original: , θ N ) , where N is the number of weight units in the neural network.
Revised: , θ N ) , where N is the number of weight units in the neural network.

Original: Given a training dataset S that contains samples i.i.d.
Revised: Given a training dataset S that contains samples i.i.d.

Original: drawn from a distribution D ,

Original: Algorithm

Original: Input: Network f θ , θ =

Original: ( θ 1 , θ 2 , . .

Original: , θ N ) , Training set S , Batch size b , Learning rate η > 0 ,

Original: Neighborhood size ρ > 0 , Iterations A , SWP hyperparameter β , SDS hyperparameter γ .

Original: Output: A ﬂat minimum solution ˆ θ .

Original: 1: for a = 1 to A do 2:

Original: Sample a mini-batch B ⊂ S with size b .

Original: 3: for n = 1 to N do 4: if θ n is chosen by probability β then 5: (cid:15) n ← ρ 1 − β ∇ θ n L B ( f θ ) (cid:46) SWP in B6:

Original: else 7: (cid:15) n ←8: ˆ (cid:15) ← ( (cid:15) 1 , ..., (cid:15) N ) (cid:46)

Original: Assign Weight Perturbation 9:

Original: Compute (cid:96) ( f θ +ˆ (cid:15) , x i , y i ) and construct B + with selection ratio γ (Equation 6) 10:

Original: Compute gradients g = ∇ θ L B + ( f θ +ˆ (cid:15) )

Original: (cid:46) SDS in B11:

Original: Update weights θ ← θ − ηg the network is trained to obtain optimal weights ˆ θ via empirical risk minimization (ERM), i.e.,ˆ [Equation] = arg [Equation]
Revised: drawn from a distribution D , the network is trained to obtain optimal weights ˆ θ via empirical risk minimization (ERM), i.e.,

Original: θ (cid:26)

Original: [Equation]
Revised: [Equation]

Original: ( f

Original: [Equation] = [Equation] (cid:88) ( x i ,y i ) ∈ S (cid:96) ( f θ , x i , y i ) (cid:27) (1) where (cid:96) can be an arbitrary loss function.
Revised: where (cid:96) can be an arbitrary loss function.

Original: We take (cid:96) to be the cross entropy loss in this paper.
Revised: We take (cid:96) to be the cross entropy loss in this paper.

Original: The population loss is deﬁned as L D ( f θ ) (cid:44) E
Revised: The population loss is deﬁned as L D ( f θ ) (cid:44) E

Original: ( x i ,y i ) ∼D ( (cid:96) ( f θ , x i , y i )) .
Revised: ( x i ,y i ) ∼D (cid:2) (cid:96) ( f θ , x i , y i ) (cid:3) .

Original: In each training iteration, optimizers sample a mini-batch B ⊂ S with size b to update parameters.
Revised: In each training iteration, optimizers sample a mini-batch B ⊂ S with size b to update parameters.

Original: 2. S HARPNESS - AWARE MINIMIZATION AND ITS COMPUTATIONAL DRAWBACK
Revised: 2. S HARPNESS - AWARE MINIMIZATION AND ITS COMPUTATIONAL DRAWBACK

Original: To improve the generalization capability of DNNs, Foret et al.
Revised: To improve the generalization capability of DNNs, Foret et al.

Original: (2020) proposed the SAM training strategy for searching ﬂat minima.
Revised: (2020) proposed the SAM training strategy for searching ﬂat minima.

Original: SAM trains DNNs by solving the following min-max optimization problem,
Revised: SAM trains DNNs by solving the following min-max optimization problem,

Original: [Equation]
Revised: [Equation]

Original: Given θ , the inner optimization attempts to ﬁnd a weight perturbation (cid:15) in Euclidean ball with radius ρ that maximizes the empirical loss.
Revised: Given θ , the inner optimization attempts to ﬁnd a weight perturbation (cid:15) in Euclidean ball with radius ρ that maximizes the empirical loss.

Original: The maximized loss at weights θ is the sum of the empirical loss and the sharpness, which is deﬁned to be R S ( f θ )
Revised: The maximized loss at weights θ is the sum of the empirical loss and the sharpness, which is deﬁned to be R S ( f θ )

Original: = max (cid:15) : (cid:107) (cid:15) (cid:107) 2 <ρ
Revised: = max (cid:15) : (cid:107) (cid:15) (cid:107) 2 <ρ

Original: [ L S ( f θ + (cid:15) )
Revised: [ L S ( f θ + (cid:15) )

Original: − L S ( f θ )] .
Revised: − L S ( f θ )] .

Original: This sharpness is quantiﬁed by the maximal change of empirical loss when aconstraint perturbation (cid:15)  is added to θ .
Revised: This sharpness is quantiﬁed by the maximal change of empirical loss when a perturbation (cid:15) (whose norm is constrained by ρ ) is added to θ .

Original: The min-max problem encourages SAM to ﬁnd ﬂat minima.
Revised: The min-max problem encourages SAM to ﬁnd ﬂat minima.

Original: For a certain set of weights θ , Foret et al.
Revised: For a certain set of weights θ , Foret et al.

Original: (2020) theoretically justiﬁes that the population loss of
Revised: (2020) theoretically justiﬁes that the population loss of

Original: 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).
Revised: 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).

Original: Thus, by minimizing the sharpness together with the empirical loss, SAM produces optimized solutions for DNNs with ﬂat minima, and the resultant models can thus generalize better (Foret et al., 2020; Chen et al., 2021; Kwon et al., 2021). 
Revised: Thus, by minimizing the sharpness together with the empirical loss, SAM produces optimized solutions for DNNs with ﬂat minima, and the resultant models can thus generalize better (Foret et al., 2020; Chen et al., 2021; Kwon et al., 2021). Indeed, we have

Original: [Equation]
Revised: [Equation]

Original: In practice, SAM ﬁrst approximately solves the inner optimization by means of a single-step gradient descent method, i.e.,
Revised: In practice, SAM ﬁrst approximately solves the inner optimization by means of a single-step gradient descent method, i.e.,

Original: [Equation]
Revised: [Equation]

Original: The sharpness at weights θ is approximated by R S ( f θ ) =
Revised: The sharpness at weights θ is approximated by R S ( f θ ) =

Original: L S ( f θ +ˆ (cid:15) )
Revised: L S ( f θ +ˆ (cid:15) )

Original: − L S ( f θ ) .
Revised: − L S ( f θ ) .

Original: Then, a base optimizer, such as SGD (Nesterov, 1983) or Adam (Kingma & Ba, 2014), updates the DNNs’ weights to minimize L S ( f θ +ˆ (cid:15) ) .
Revised: Then, a base optimizer, such as SGD (Nesterov, 1983) or Adam (Kingma & Ba, 2015), updates the DNNs’ weights to minimize L S ( f θ +ˆ (cid:15) ) .

Original: We refer to L S ( f θ +ˆ (cid:15) ) as the SAM loss .
Revised: We refer to L S ( f θ +ˆ (cid:15) ) as the SAM loss .

Original: Overall, SAM requires two forward and two backward operations to update weights once.
Revised: Overall, SAM requires two forward and two backward operations to update weights once.

Original: We refer to the forward and backward propagation for approximating ˆ (cid:15) as F 1 and B 1 and those for updating weights by base optimizers as F 2 and B respectively.
Revised: We refer to the forward and backward propagation for approximating ˆ (cid:15) as F 1 and B 1 and those for updating weights by base optimizers as F 2 and Brespectively.

Original: Although SAM can effectively improve the generalization of DNNs, it additionally requires one forward and one backward operation ( F 1 and B 1 ) in each training iteration.
Revised: Although SAM can effectively improve the generalization of DNNs, it additionally requires one forward and one backward operation ( F 1 and B 1 ) in each training iteration.

Original: Thus, SAM results in a doubling of the computational overhead compared to the use of base optimizers.
Revised: Thus, SAM results in a doubling of the computational overhead compared to the use of base optimizers.

Original: To improve the efﬁciency of SAM, we propose ESAM, which consists of two strategies—SWP and SDS, to accelerate the sharpness approximation phase and the weight updating phase.
Revised: To improve the efﬁciency of SAM, we propose ESAM, which consists of two strategies—SWP and SDS, to accelerate the sharpness approximation phase and the weight updating phase.

Original: Specifically, on the one hand, when estimating ˆ (cid:15) around weight vector θ , SWP efﬁciently approximates ˆ (cid:15) by randomly selecting each parameter with a given probability to form a subset of weights to be perturbed.
Revised: Specifically, on the one hand, when estimating ˆ (cid:15) around weight vector θ , SWP efﬁciently approximates ˆ (cid:15) by randomly selecting each parameter with a given probability to form a subset of weights to be perturbed.

Original: The reduction of the number of perturbed parameters results in lower computational overhead during the backward propagation.
Revised: The reduction of the number of perturbed parameters results in lower computational overhead during the backward propagation.

Original: SWP rescales the resultant weight perturbation so as to assure that the expected weight perturbation equals to ˆ (cid:15) , and the generalization capability thus will not be signiﬁcantly degraded.
Revised: SWP rescales the resultant weight perturbation so as to assure that the expected weight perturbation equals to ˆ (cid:15) , and the generalization capability thus will not be signiﬁcantly degraded.

Original: On the other hand, when updating weights via base optimizers, instead of computing the upper bound L B ( f θ +ˆ (cid:15) ) over a whole batch of samples, SDS selects a subset of samples, B + , whose loss values increase the most with respect to the perturbation ˆ (cid:15) .
Revised: On the other hand, when updating weights via base optimizers, instead of computing the upper bound L B ( f θ +ˆ (cid:15) ) over a whole batch of samples, SDS selects a subset of samples, B + , whose loss values increase the most with respect to the perturbation ˆ (cid:15) .

Original: Optimizing the weights based on a fewer number of samples decreases the computational overhead (in a linear fashion).
Revised: Optimizing the weights based on a fewer number of samples decreases the computational overhead (in a linear fashion).

Original: We further justify that L B ( f θ +ˆ (cid:15) ) can be upper bounded by L B + ( f θ +ˆ (cid:15) ) and consequently the generalization capability can be preserved.
Revised: We further justify that L B ( f θ +ˆ (cid:15) ) can be upper bounded by L B + ( f θ +ˆ (cid:15) ) and consequently the generalization capability can be preserved.

Original: In general, ESAM works much more efﬁciently and performs as well as SAM in terms of its generalization capability.
Revised: In general, ESAM works much more efﬁciently and performs as well as SAM in terms of the generalization capability.

Original: 2. S TOCHASTIC W EIGHT P ERTURBATION
Revised: 2. S TOCHASTIC W EIGHT P ERTURBATION

Original: This section elaborates on the ﬁrst efﬁciency enhancement strategy, SWP, and explains why SWP can effectively reduce computational overhead while preserving the generalization capability.
Revised: This section elaborates on the ﬁrst efﬁciency enhancement strategy, SWP, and explains why SWP can effectively reduce computational overhead while preserving the generalization capability.

Original: To efﬁciently approximate ˆ (cid:15) ( θ, S ) during the sharpness estimation phase, SWP randomly chooses a subset ˜ θ
Revised: To efﬁciently approximate ˆ (cid:15) ( θ, S ) during the sharpness estimation phase, SWP randomly chooses a subset ˜ θ

Original: = {
Revised: = {

Original: θ I 1 , θ I 2 , . . .
Revised: θ I 1 , θ I 2 , . . .

Original: } from the original set of weights
Revised: } from the original set of weights

Original: θ =
Revised: θ =

Original: ( θ 1 , . .
Revised: ( θ 1 , . .

Original: , θ N ) to perform backpropagation B 1 .
Revised: , θ N ) to perform backpropagation B 1 .

Original: Each parameter is selected to be in the subvector ˜
Revised: Each parameter is selected to be in the subvector ˜

Original: θ with some probability β , which can be tuned as a hyperparameter.
Revised: θ with some probability β , which can be tuned as a hyperparameter.

Original: SWP approximates the weight perturbation with ρ ∇ ˜ θ
Revised: SWP approximates the weight perturbation with ρ ∇ ˜ θ

Original: L S ( f θ ) .
Revised: L S ( f θ ) .

Original: To be formal, we introduce a gradient mask m = ( m 1 , . . .
Revised: To be formal, we introduce a gradient mask m = ( m 1 , . . .

Original: , m N ) where m i i .
Revised: , m N ) where m i i .

Original: i .
Revised: i .

Original: d . ∼ Bern( β ) for all i ∈ { 1 , . . .
Revised: d . ∼ Bern( β ) for all i ∈ { 1 , . . .

Original: , N } .
Revised: , N } .

Original: Then, we have ρ ∇ ˜ θ
Revised: Then, we have ρ ∇ ˜ θ

Original: L S ( f θ ) = m (cid:62) ˆ (cid:15) ( θ, B ) .
Revised: L S ( f θ ) = m (cid:62) ˆ (cid:15) ( θ, B ) .

Original: To ensure the expected weight perturbation of SWP equals to ˆ (cid:15) , we scale ρ ∇ ˜ θ L S ( f θ ) by a factor of 1 β .
Revised: To ensure the expected weight perturbation of SWP equals to ˆ (cid:15) , we scale ρ ∇ ˜ θ L S ( f θ ) by a factor of 1 β .

Original: Finally, SWP produces an approximate solution of the inner maximization as
Revised: Finally, SWP produces an approximate solution of the inner maximization as

Original: [Equation]
Revised: [Equation]

Original: Computation Ideally, SWP reduces the overall computational overhead in proportion to 1 − β in B 1 .
Revised: Computation Ideally, SWP reduces the overall computational overhead in proportion to 1 − β in B 1 .

Original: However, there exists some parameters not included in ˜ θ that are still required to be updated in the backpropagation step.
Revised: However, there exists some parameters not included in ˜ θ that are still required to be updated in the backpropagation step.

Original: 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.
Revised: 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.

Original: This additional computational overhead slightly increases in deeper neural networks.
Revised: This additional computational overhead slightly increases in deeper neural networks.

Original: Thus, the amount of reduction in the computational overhead is positively correlated to 1 − β .
Revised: Thus, the amount of reduction in the computational overhead is positively correlated to 1 − β .

Original: In practice, β is tuned to maximize SWP’s efﬁciency while maintaining a generalization performance comparable to SAM’s.
Revised: In practice, β is tuned to maximize SWP’s efﬁciency while maintaining a generalization performance comparable to SAM’s.

Original: 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 ˆ (cid:15) ( θ, B ) in the sense of the (cid:96) 2 norm and direction.
Revised: 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 ˆ (cid:15) ( θ, B ) in the sense of the (cid:96) 2 norm and direction.

Original: We denote the expected SWP perturbation by ¯ a ( θ, B ) , where¯ a ( θ, B )
Revised: We denote the expected SWP perturbation by ¯ a ( θ, B ) , where

Original: [ i ] =
Revised: [Equation]

Original: E [ a ( θ, B ) [ i ] ] = 1 β · β ˆ (cid:15) ( θ, B )
Revised: for i ∈ { 1 , . . . , N } . Thus, it holds that

Original: [ i ] =
Revised: [Equation]

Original: ˆ (cid:15) ( θ, B )
Revised: and CosSim (cid:16)

Original: [ i ] , for i ∈ { 1 , . . . , N } .
Revised: [Equation]

Original: Thus, it holds that (cid:107) ¯ a ( θ, [-B ) (cid:107) 
-]
Revised: θ, [-B ) (cid:107) 
-][Equation]

Original: = (cid:107) ˆ (cid:15) ( θ, B ) (cid:107)and CosSim (cid:16) ¯ a ( θ, B ) , ˆ (cid:15) ( θ, B ) (cid:17) = 1 , showing that the expected weight perturbation of SWP is the same as that of SAM’s.
Revised: showing that the expected weight perturbation of SWP is the same as that of SAM’s.

Original: 2. S HARPNESS - SENSITIVE D ATA S ELECTION
Revised: 2. S HARPNESS - SENSITIVE D ATA S ELECTION

Original: L + [Equation] f 
Revised: L + ( f + )

Original: [Equation]

Original: L ( f + ) =0 .0 .1 .1 .2 .2 .2 .
Revised: L ( f + ) =0 .0 .1 .1 .2 .2 .2 .

Original: In this section, we introduce the second efﬁciency enhancement technique, SDS, which reduces computational overhead of SAM linearly as the number of selected samples decreases.
Revised: In this section, we introduce the second efﬁciency enhancement technique, SDS, which reduces computational overhead of SAM linearly as the number of selected samples decreases.

Original: We also explain why the generalization capability of SAM is preserved by SDS.
Revised: We also explain why the generalization capability of SAM is preserved by SDS.

Original: In the sharpness estimation phase, we obtain the approximate solution ˆ (cid:15) of the inner maximization.
Revised: In the sharpness estimation phase, we obtain the approximate solution ˆ (cid:15) of the inner maximization.

Original: Perturbing weights along this direction signiﬁcantly increases the average loss over a batch B .
Revised: Perturbing weights along this direction signiﬁcantly increases the average loss over a batch B .

Original: To improve the efﬁciency but still control the upper bound L B ( f θ +ˆ (cid:15) ) , we select a subset of samples from the whole batch.
Revised: To improve the efﬁciency but still control the upper bound L B ( f θ +ˆ (cid:15) ) , we select a subset of samples from the whole batch.

Original: The loss values of this subset of samples increase most when the weights are perturbed by ˆ (cid:15) .
Revised: The loss values of this subset of samples increase most when the weights are perturbed by ˆ (cid:15) .

Original: To be speciﬁc, SDS splits the mini-batch B into the following two subsets
Revised: To be speciﬁc, SDS splits the mini-batch B into the following two subsets

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where B + is termed as the sharpness-sensitive subset and the threshold α controls the size of B + .
Revised: where B + is termed as the sharpness-sensitive subset and the threshold α controls the size of B + .

Original: We let γ = | B + | / | B | be the ratio of the number of selected samples with respect to the batch size.
Revised: We let γ = | B + | / | B | be the ratio of the number of selected samples with respect to the batch size.

Original: In practice, γ determines the exact value of α and serves as a predeﬁned hyperparameter of SDS.
Revised: In practice, γ determines the exact value of α and serves as a predeﬁned hyperparameter of SDS.

Original: As illustrated in Figure 2, when α = 0 , the gradient of the weights evaluated on B + aligns with the direction of ˆ (cid:15) and the loss values of the samples in B + will increase with respect to the weight perturbation ˆ (cid:15) .
Revised: As illustrated in Figure 2, when α = 0 , the gradient of the weights evaluated on B + aligns with the direction of ˆ (cid:15) and the loss values of the samples in B + will increase with respect to the weight perturbation ˆ (cid:15) .

Original: Computation SDS reduces the computational overhead in F 2 and B 2 .
Revised: Computation SDS reduces the computational overhead in F 2 and B 2 .

Original: The reduction is linear in 1 − γ .
Revised: The reduction is linear in 1 − γ .

Original: The hyperparameter γ can be tuned to meet up distinct requirements in efﬁciency and performance.
Revised: The hyperparameter γ can be tuned to meet up distinct requirements in efﬁciency and performance.

Original: SDS is conﬁgured the same as SWP for maximizing efﬁciency with comparable performance to SAM.
Revised: SDS is conﬁgured the same as SWP for maximizing efﬁciency with comparable performance to SAM.

Original: Generalization For the generalization capability, we now justify that the SAM loss computing over the batch B , L B ( f θ +ˆ (cid:15) ) , can be approximately upper bounded by the corresponding loss evaluated only on B + , L B + ( f θ +ˆ (cid:15) ) .
Revised: Generalization For the generalization capability, we now justify that the SAM loss computing over the batch B , L B ( f θ +ˆ (cid:15) ) , can be approximately upper bounded by the corresponding loss evaluated only on B + , L B + ( f θ [Equation] ˆ [Equation]

Original: From Equation 3, we have
Revised: From Equation 3, we have

Original: [Equation]
Revised: [Equation]

Original: On the one hand, since R B ( f θ ) = 1 | B | (cid:80) ( x i ,y i ) ∈
Revised: On the one hand, since R B ( f θ ) = 1 | B | (cid:80) ( x i ,y i ) ∈

Original: B
Revised: B

Original: [ (cid:96) ( f θ +ˆ (cid:15) , x i , y i ) − (cid:96) ( f θ , x i , y i )] represents the average sharpness of the batch B , by Equation 6, we have R B − ( f θ ) ≤ R B ( f θ ) ≤ R B
Revised: [ (cid:96) ( f θ +ˆ (cid:15) , x i , y [Equation] ) − (cid:96) ( f θ [Equation]

Original: + ( f θ ) , and
Revised: x i [Equation] y

Original:  R B − ( f θ ) − R B+ ( f θ ) ≤ 0 .
Revised: [Equation] )] represents the average sharpness of the batch B , by Equation 6, we have R B − ( f θ ) ≤ R B ( f θ ) ≤ R B + ( f θ ) , and [Equation]

Original: On the other hand, B + and B − are constructedsorting by  (cid:96) ( f θ +ˆ (cid:15) , x i , y i ) − (cid:96) ( f θ , x i , y i ) , which is positively correlated to l ( f θ , x i , y i ) (Li et al., 2019) (more details can be found inthe Appendix A.2).
Revised: On the other hand, B + and B − are constructed by sorting (cid:96) ( f θ +ˆ (cid:15) , x i , y i ) − (cid:96) ( f θ , x i , y i ) , which is positively correlated to l ( f θ , x i , y i ) (Li et al., 2019) (more details can be found in Appendix A.2).

Original: Thus, we have
Revised: Thus, we have

Original: [Equation]
Revised: [Equation]

Original: Therefore, by Equation 8 and Equation 9, we have
Revised: Therefore, by Equation 8 and Equation 9, we have

Original: [Equation]
Revised: [Equation]

Original: Experimental results in Figure 5 corroborate that R B − ( f θ ) − R B + ( f θ ) < 0 and L B − ( f θ ) − L B + ( f θ ) < 0 .
Revised: Experimental results in Figure 5 corroborate that R B − ( f θ ) − R B + ( f θ ) < 0 and L B − ( f θ ) − L B + ( f θ ) < 0 .

Original: Besides, Figure 6 veriﬁes that the selected batch B + is sufﬁciently representative to mimic the gradients of B since B + has a signiﬁcantly higher cosine similarity with B compared to B − in terms of the computed gradients.
Revised: Besides, Figure 6 veriﬁes that the selected batch B + is sufﬁciently representative to mimic the gradients of B since B + has a signiﬁcantly higher cosine similarity with B compared to B − in terms of the computed gradients.

Original: According to Equation 10, one can utilize L B + ( f θ +ˆ (cid:15) ) as a proxy to the real objective to minimize of the overall loss L B ( f θ +ˆ (cid:15) ) with a smaller number of samples.
Revised: According to Equation 10, one can utilize L B + ( f θ +ˆ (cid:15) ) as a proxy to the real objective to minimize of the overall loss L B ( f θ +ˆ (cid:15) ) with a smaller number of samples.

Original: As a result, SDS improves SAM’s efﬁciency without performance degradation.
Revised: As a result, SDS improves SAM’s efﬁciency without performance degradation.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: This section demonstrates the effectiveness of our proposed ESAM algorithm.
Revised: This section demonstrates the effectiveness of our proposed ESAM algorithm.

Original: 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:
Revised: 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:

Original: ResNet (He et al., 2016), Wide ResNet (Zagoruyko & Komodakis, 2016), and PyramidNet (Han et al., 2017).
Revised: ResNet (He et al., 2016), Wide ResNet (Zagoruyko & Komodakis, 2016), and PyramidNet (Han et al., 2017).

Original: We demonstrate the proposed ESAM improves the efﬁciency of vanilla SAM by speeding up to 40 .
Revised: We demonstrate the proposed ESAM improves the efﬁciency of vanilla SAM by speeding up to 40 .

Original: 3% computational overhead with better generalization performance.
Revised: 3% computational overhead with better generalization performance.

Original: We report the main results in Table 1 and Table 2.
Revised: We report the main results in Table 1 and Table 2.

Original: Besides, we perform an ablation study on the two proposed strategies of ESAM (i.e., SWP and SDS).
Revised: Besides, we perform an ablation study on the two proposed strategies of ESAM (i.e., SWP and SDS).

Original: The experimental results in Table 3 and Figure 3 indicate that both strategies improve SAM’s efﬁciency and performance.
Revised: The experimental results in Table 3 and Figure 3 indicate that both strategies improve SAM’s efﬁciency and performance.

Original: 3. R ESULTS
Revised: 3. R ESULTS

Original: CIFAR10 and CIFAR

Revised: CIFAR10 and CIFAR

Original: We start from evaluating ESAM on the CIFAR-10 and CIFAR-100 image classiﬁcation datasets.
Revised: We start from evaluating ESAM on the CIFAR-10 and CIFAR-100 image classiﬁcation datasets.

Original: The evaluation is carried out on three different model architectures: ResNet18 (He et al., 2016), WideResNet-28-10 (Zagoruyko & Komodakis, 2016) and PyramidNet- (Han et al., 2017).
Revised: The evaluation is carried out on three different model architectures: ResNet18 (He et al., 2016), WideResNet-28-10 (Zagoruyko & Komodakis, 2016) and PyramidNet- (Han et al., 2017).

Original: 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,
Revised: 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,

Original: SAM and ESAM.
Revised: SAM and ESAM.

Original: Additionally, the other hyperparameters of SGD, SAM and ESAM have been tuned separately for optimal test accuracies using grid search.
Revised: Additionally, the other hyperparameters of SGD, SAM and ESAM have been tuned separately for optimal test accuracies using grid search.

Original: We train all the models with 3 different random seeds using a batch size of 128 , weight decay 10 −and cosine learning rate decay (Loshchilov & Hutter, 2016).
Revised: We train all the models with 3 different random seeds using a batch size of 128 , weight decay 10 −and cosine learning rate decay (Loshchilov & Hutter, 2017).

Original: The training epochs are set to befor ResNet-18 (He et al., 2016), WideResNet-28-10 (Zagoruyko & Komodakis, 2016), and 300 for
Revised: The training epochs are set to befor ResNet-18 (He et al., 2016), WideResNet-28-10 (Zagoruyko & Komodakis, 2016), and 300 for

Original: PyramidNet-110 (Han et al., 2017).
Revised: PyramidNet-110 (Han et al., 2017).

Original: We set β = 0 .
Revised: We set β = 0 .

Original: 6 and γ = 0 .
Revised: 6 and γ = 0 .

Original: 5 for ResNet-18 and PyramidNet110 models; and set β = 0 .
Revised: 5 for ResNet-18 and PyramidNet110 models; and set β = 0 .

Original: 5 and γ = 0 .
Revised: 5 and γ = 0 .

Original: 5 for WideResNet-28-10.
Revised: 5 for WideResNet-28-10.

Original: The above-mentioned β and γ are optimal for efﬁciency with comparable performance compared to SAM.
Revised: The above-mentioned β and γ are optimal for efﬁciency with comparable performance compared to SAM.

Original: The details of training setting are listed in Appendix A.7.
Revised: The details of training setting are listed in Appendix A.7.

Original: We record the best test accuracies obtained by SGD, SAM and
Revised: We record the best test accuracies obtained by SGD, SAM and

Original: ESAM in Table
Revised: ESAM in Table

Original: [Table]
Revised: [Table]

Original: The experimental results indicate that our proposed ESAM can increase the training speed by up to 40 .
Revised: The experimental results indicate that our proposed ESAM can increase the training speed by up to 40 .

Original: 30% in comparison with SAM.
Revised: 30% in comparison with SAM.

Original: Concerning the performance, ESAM outperforms SAM in the six sets of experiments.
Revised: Concerning the performance, ESAM outperforms SAM in the six sets of experiments.

Original: The best efﬁciency of ESAM is reported in CIFAR10 trained with ResNet- (Training speed 140 .
Revised: The best efﬁciency of ESAM is reported in CIFAR10 trained with ResNet- (Training speed 140 .

Original: 3% vs. SAM 100% ).
Revised: 3% vs. SAM 100% ).

Original: The best accuracy is reported in CIFAR100 trained with
Revised: The best accuracy is reported in CIFAR100 trained with

Original: PyramidNet110 (Accuracy 85 .
Revised: PyramidNet110 (Accuracy 85 .

Original: 56% vs. SAM 84 .
Revised: 56% vs. SAM 84 .

Original: ESAM improves efﬁciency and achieves better performance compared to SAM in CIFAR10/100 benchmarks
Revised: ESAM improves efﬁciency and achieves better performance compared to SAM in CIFAR10/100 benchmarks

Original: [Table] [Table]
Revised: [Table] [Table]

Original: ImageNet To evaluate ESAM’s effectiveness on a large-scale benchmark dataset
Revised: ImageNet To evaluate ESAM’s effectiveness on a large-scale benchmark dataset

Original: [Table]
Revised: [Table]

Original: we conduct experiments on ImageNet Datasets.
Revised: we conduct experiments on ImageNet Datasets.

Original: The 1000 class ImageNet dataset contains roughly 1 .
Revised: The 1000 class ImageNet dataset contains roughly 1 .

Original: 28 million
Revised: 28 million

Revised: [Table]

Original: [Table]
Revised: [Table]

Original: training images and 50 , 000 validation images with 469 × 387 averaged resolution.
Revised: training images and 50 , 000 validation images with 469 × 387 averaged resolution.

Original: The ImageNet dataset is more representative (of real-world scenarios) and persuasive (of a method’s effectiveness) than CIFAR datasets.
Revised: The ImageNet dataset is more representative (of real-world scenarios) and persuasive (of a method’s effectiveness) than CIFAR datasets.

Original: We resize the images on ImageNet to 224 × 224 resolution to train ResNetand ResNet-110 models.
Revised: We resize the images on ImageNet to 224 × 224 resolution to train ResNetand ResNet-110 models.

Original: We train 90 epochs and set the optimal hyperparameters for SGD, SAM and ESAM as suggested by Chen et al.
Revised: We train 90 epochs and set the optimal hyperparameters for SGD, SAM and ESAM as suggested by Chen et al.

Original: , and the details are listed in appendix A.
Revised: , and the details are listed in appendix A.7.

Original: We use β = 0 .
Revised: We use β = 0 .

Original: 6 and γ = 0 .
Revised: 6 and γ = 0 .

Original: 7 for ResNet-50 and ResNet-101.
Revised: 7 for ResNet-50 and ResNet-101.

Original: We employ the m -sharpness strategy for both
Revised: We employ the m -sharpness strategy for both

Original: SAM and ESAM with m = 128 , which is the same as that suggested in Zheng et al.
Revised: SAM and ESAM with m = 128 , which is the same as that suggested in Zheng et al.

Original: The experimental results are reported in Table 2.
Revised: The experimental results are reported in Table 2.

Original: The results indicate that the performance of ESAM on large-scale datasets is consistent with the two (smaller) CIFAR datasets.
Revised: The results indicate that the performance of ESAM on large-scale datasets is consistent with the two (smaller) CIFAR datasets.

Original: ESAM outperforms
Revised: ESAM outperforms

Original: SAM by 0 .
Revised: SAM by 0 .

Original: 35% to 0 .
Revised: 35% to 0 .

Original: 49% in accuracy and, more importantly, enjoys 28 .
Revised: 49% in accuracy and, more importantly, enjoys 28 .

Original: 7% faster training speed
Revised: 7% faster training speed

Original: compared to SAM.
Revised: compared to SAM.

Original: 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.
Revised: 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.

Original: These experiments demonstrate that ESAM outperforms SAM on a variety of benchmark datasets for widely-used DNNs’ architectures in terms of training speed and classiﬁcation accuracies.
Revised: These experiments demonstrate that ESAM outperforms SAM on a variety of benchmark datasets for widely-used DNNs’ architectures in terms of training speed and classiﬁcation accuracies.

Original: 3. A BLATION AND P ARAMETER S TUDIES
Revised: 3. A BLATION AND P ARAMETER S TUDIES

Original: To better understand the effectiveness of SWP and SDS in improving the performance and efﬁciency 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.
Revised: To better understand the effectiveness of SWP and SDS in improving the performance and efﬁciency 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.

Original: We consider two variants of ESAM:
Revised: We consider two variants of ESAM:

Original: (i) only with SWP, (ii) only with SDS.
Revised: (i) only with SWP, (ii) only with SDS.

Original: The rest of the experimental settings are identical to the settings described in Section 3.1.
Revised: The rest of the experimental settings are identical to the settings described in Section 3.1.

Original: 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 .
Revised: 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 .

Original: We report the grid search results in Figure 3.
Revised: We report the grid search results in Figure 3.

Original: We use β = 0 .
Revised: We use β = 0 .

Original: 6 , γ = 0 .
Revised: 6 , γ = 0 .

Original: 5 for ResNet-18; and set β = 0 .
Revised: 5 for ResNet-18; and set β = 0 .

Original: 5 , γ = 0 .
Revised: 5 , γ = 0 .

Original: 5 for WideResNet-28-10 in the four sets of ablation studies.
Revised: 5 for WideResNet-28-10 in the four sets of ablation studies.

Original: The ablation study results are reported in Table 3.
Revised: The ablation study results are reported in Table 3.

Revised: [Table]

Revised: [Table]

Original: [Table]
Revised: [Table]

Original: + ESAM-SWP 96.74 [+1.33]
Revised: + ESAM-SWP 96.74 [+1.33]

Original: [Table]
Revised: [Table]

Original:  + ESAM-SDS 96.45 [+1.04] 2,105 (
Revised: [Table] [Table] [Table] + ESAM-SDS 96.45 [+1.04] 2,105 (122.6%)

Original: [Table]
Revised: [Table]

Original: [Table] [Table]
Revised: 38 [+2.21]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: ESAM-SWP As shown in Table 3, SWP improves SAM’s training speed by 8 . 3% to 10 .
Revised: ESAM-SWP As shown in Table 3, SWP improves SAM’s training speed by 8 . 3% to 10 .

Original: 5% , and achieves better performance at the same time.
Revised: 5% , and achieves better performance at the same time.

Original: SWP can further improve the efﬁciency by using a smaller β .
Revised: SWP can further improve the efﬁciency by using a smaller β .

Original: The best performance of SWP is obtained when β = 0 .
Revised: The best performance of SWP is obtained when β = 0 .

Original: 6 for ResNet-18 and β = 0 .for
Revised: 6 for ResNet-18 and β = 0 .for

Original: WideResNet-28-10.
Revised: WideResNet-28-10.

Original: The four sets of experiments indicate that β is consistent among different architectures and datasets.
Revised: The four sets of experiments indicate that β is consistent among different architectures and datasets.

Original: Therefore, we set β = 0 .
Revised: Therefore, we set β = 0 .

Original: 6 for PyramidNet in CIFAR10/100 datasets and
Revised: 6 for PyramidNet on CIFAR10/100 datasets and

Original: ResNet on ImageNet datasets.
Revised: ResNet on ImageNet datasets.

Original: ESAM-SDS SDS also signiﬁcantly improves the efﬁciency by 21 .
Revised: ESAM-SDS SDS also signiﬁcantly improves the efﬁciency by 21 .

Original: 5% to 25 .
Revised: 5% to 25 .

Original: 8% compared to SAM.
Revised: 8% compared to SAM.

Original: It outperforms SAM’s performance in CIFAR100 datasets, and achieves comparable performance in CIFAR10 datasets.
Revised: It outperforms SAM’s performance on CIFAR100 datasets, and achieves comparable performance on CIFAR10 datasets.

Original: SDS can outperform SAM in both datasets with both architectures with little degradation to the efﬁciency, as demonstrated in Figure 3.
Revised: SDS can outperform SAM on both datasets with both architectures with little degradation to the efﬁciency, as demonstrated in Figure 3.

Original: Across all experiments, γ = 0 .
Revised: Across all experiments, γ = 0 .

Original: 5 is the smallest value that is optimal for efﬁciency while maintaining comparable performance to SAM.
Revised: 5 is the smallest value that is optimal for efﬁciency while maintaining comparable performance to SAM.

Original: Visualization of Loss
Revised: Visualization of Loss

Original: Landscapes To visualize the sharpness of the ﬂat minima obtained by ESAM, we plot the loss landscapes trained with SGD, SAM and ESAM on the ImageNet dataset.
Revised: Landscapes To visualize the sharpness of the ﬂat minima obtained by ESAM, we plot the loss landscapes trained with SGD, SAM and ESAM on the ImageNet dataset.

Original: We display the loss landscapes in Figure 4, following the plotting algorithm in Li et al.
Revised: We display the loss landscapes in Figure 4, following the plotting algorithm in Li et al.

Original: The x - and yaxes represent two random sampled orthogonal Gaussian perturbations.
Revised: The x - and yaxes represent two random sampled orthogonal Gaussian perturbations.

Original: We sampled 100 ×points for 10 groups random Gaussian perturbations.
Revised: We sampled 100 ×points for 10 groups random Gaussian perturbations.

Original: The displayed loss landscapes are the results we obtained by averaging over ten groups of random perturbations.
Revised: The displayed loss landscapes are the results we obtained by averaging over ten groups of random perturbations.

Original: It can be clearly seen that both
Revised: It can be clearly seen that both

Original: SAM and ESAM improve the sharpness signiﬁcantly in comparison to SGD.
Revised: SAM and ESAM improve the sharpness signiﬁcantly in comparison to SGD.

Original: To summarize, SWP and SDS both reduce the computational overhead and accelerate training compared to SAM.
Revised: To summarize, SWP and SDS both reduce the computational overhead and accelerate training compared to SAM.

Original: Most importantly, both these strategies achieve a comparable or better performance than SAM.
Revised: Most importantly, both these strategies achieve a comparable or better performance than SAM.

Original: In practice, by conﬁguring the β and γ , ESAM can meet a variety of user-deﬁned efﬁciency and performance requirements.
Revised: In practice, by conﬁguring the β and γ , ESAM can meet a variety of user-deﬁned efﬁciency and performance requirements.

Original:  R ELATED WORK
Revised:  R ELATED WORK

Original: The concept of regularizing sharpness for better generalization dates back to (Hochreiter & Schmidhuber, 1995).
Revised: The concept of regularizing sharpness for better generalization dates back to (Hochreiter & Schmidhuber, 1995).

Original: By using an MDL-based argument, which clariﬁes that a statistical model with fewer
Revised: By using an MDL-based argument, which clariﬁes that a statistical model with fewer

Original: [Figure]
Revised: [Figure]

Original: bits to describe can have better generalization ability, Hochreiter & Schmidhuber (1995) claim that a ﬂat minimum can alleviate overﬁtting issues.
Revised: bits to describe can have better generalization ability, Hochreiter & Schmidhuber (1995) claim that a ﬂat minimum can alleviate overﬁtting issues.

Original: Following that, more studies were proposed to investigate the connection between the ﬂat minima with the generalization abilities (Keskar et al., 2016;
Revised: Following that, more studies were proposed to investigate the connection between the ﬂat minima with the generalization abilities (Keskar et al., 2017; Dinh et al., 2017;

Original: Dinh et al., 2017; Liu et al., 2020; Li et al., 2017; Dziugaite & Roy, 2017; Jiang et al., 2019;
Revised: Liu et al., 2020; Li et al., 2018; Dziugaite & Roy, 2017; Jiang et al., 2019;

Original: Moosavi-Dezfooli et al., 2019).
Revised: Moosavi-Dezfooli et al., 2019).

Original: Keskar et al.
Revised: Keskar et al.

Original: (2016) starts by investigating the phenomenon that training with a larger batch size results in worse generalization ability.
Revised: (2017) starts by investigating the phenomenon that training with a larger batch size results in worse generalization ability.

Original: The authors found that the sharpness of the minimum is critical in accounting for the observed phenomenon.
Revised: The authors found that the sharpness of the minimum is critical in accounting for the observed phenomenon.

Original: Keskar et al.
Revised: Keskar et al.

Original: (2016) and Dinh et al.
Revised: (2017) and Dinh et al.

Original: (2017) both argue that the sharpness can be characterized using the eigenvalues of the Hessian.
Revised: (2017) both argue that the sharpness can be characterized using the eigenvalues of the Hessian.

Original: Although they also deﬁne speciﬁc notions and methods to quantify sharpness, they do not propose complete training strategies to ﬁnd minima that are relative “ﬂat”.
Revised: Although they also deﬁne speciﬁc notions and methods to quantify sharpness, they do not propose complete training strategies to ﬁnd minima that are relative “ﬂat”.

Original: SAM (Foret et al., 2020) leverages the connection between “ﬂat” minima and the generalization error to train DNNs that generalize well across the natural distribution.
Revised: SAM (Foret et al., 2020) leverages the connection between “ﬂat” minima and the generalization error to train DNNs that generalize well across the natural distribution.

Original: Inspired by Keskar et al.
Revised: Inspired by Keskar et al.

Original:  and Dinh et al.
Revised: (2017) and Dinh et al.

Original: (2017), SAM ﬁrst proposes the quantiﬁcation of the sharpness, which is achieved by solving a maximization problem.
Revised: (2017), SAM ﬁrst proposes the quantiﬁcation of the sharpness, which is achieved by solving a maximization problem.

Original: Then, SAM proposes a complete training algorithm to improve the generalization abilities of DNNs.
Revised: Then, SAM proposes a complete training algorithm to improve the generalization abilities of DNNs.

Original: SAM is demonstrated to achieve state-of-the-art performance in a variety of deep learning benchmarks, including image classiﬁcation, natural language processing, and noisy learning (Foret et al., 2020;
Revised: SAM is demonstrated to achieve state-of-the-art performance in a variety of deep learning benchmarks, including image classiﬁcation, natural language processing, and noisy learning (Foret et al., 2020;

Original: Chen et al., 2021; Kwon et al., 2021; Pham et al., 2021; Yuan et al., 2021; Jia et al., 2021).
Revised: Chen et al., 2021; Kwon et al., 2021; Pham et al., 2021; Yuan et al., 2021; Jia et al., 2021).

Original: A series of SAM-related works has been proposed.
Revised: A series of SAM-related works has been proposed.

Original: 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.
Revised: 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.

Original: Many works focus on combining
Revised: Many works focus on combining

Original: 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).
Revised: 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 et al., 2021).

Original: Kwon et al.
Revised: Kwon et al.

Original: (2021) improves SAM’s sharpness by adaptively scaling the size of the nearby search space ρ in relation to the size of parameters.
Revised: (2021) improves SAM’s sharpness by adaptively scaling the size of the nearby search space ρ in relation to the size of parameters.

Revised: Liu et al.

Revised: (2022) leverages the past calculated weight perturbations to save SAM’s computations.

Original: However, most of these works overlook the fact that SAM improves generalization at the expense of the doubling the computational overhead.
Revised: However, most of these works overlook the fact that SAM improves generalization at the expense of the doubling the computational overhead.

Original: As a result, most of the SAM-related works suffer from the same efﬁciency drawback as SAM.
Revised: As a result, most of the SAM-related works suffer from the same efﬁciency drawback as SAM.

Original: 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 efﬁciently improve the generalization ability of DNNs.
Revised: 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 efﬁciently improve the generalization ability of DNNs.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this paper, we propose the Efﬁcient Sharpness Aware Minimizer (ESAM) to enhance the efﬁciency of vanilla SAM.
Revised: In this paper, we propose the Efﬁcient Sharpness Aware Minimizer (ESAM) to enhance the efﬁciency of vanilla SAM.

Original: The proposed ESAM integrates two novel training strategies, namely, SWP and
Revised: The proposed ESAM integrates two novel training strategies, namely, SWP and

Original: SDS, both of which are derived based on theoretical underpinnings and are evaluated over a variety of datasets and DNN architectures.
Revised: SDS, both of which are derived based on theoretical underpinnings and are evaluated over a variety of datasets and DNN architectures.

Original: Both SAM and ESAM are two-step training strategies consisting of sharpness estimation and weight updating.
Revised: Both SAM and ESAM are two-step training strategies consisting of sharpness estimation and weight updating.

Original: In each step, gradient back-propagation is performed to compute the weight perturbation or updating.
Revised: In each step, gradient back-propagation is performed to compute the weight perturbation or updating.

Original: 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.
Revised: 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.

Original: 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.
Revised: A CKNOWLEDGEMENT

Original: 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.
Revised: Jiawei Du and Joey Tianyi Zhou are suppored by Joey Tianyi Zhou’s A*STAR SERC Central Research Fund. Hanshu Yan and Vincent Tan are funded by a Singapore National Research Foundation (NRF) Fellowship (R-263-000-D02-281) and a Singapore Ministry of Education AcRF Tier 1 grant (R-263000-E80-114)

Original: We will also release our ESAM codes that embedded into TIMM (Wightman, 2019) framework to the public.
Revised: 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.

Original: A A PPENDIX A. T HE ALGORITHM OF SAM
Revised: A A PPENDIX A. T HE ALGORITHM OF SAM

Original: The algorithm of SAM is demonstrated in Algorithm 2.
Revised: The algorithm of SAM is demonstrated in Algorithm 2.

Original: Algorithm 2 SGD vs. SAM
Revised: Algorithm 2 SGD vs. SAM

Original:  do 
Revised: 1: for a do Sample

Original: 4: ˆ (cid:15) ←
Revised: 4: ˆ (cid:15) ←(cid:46)

Original: 5: else  6: ˆ (cid:15) ← ∇ θ L B 
Revised: 5: else then 6: ˆ (cid:15) ← ∇ θ L B ( f θ ) (cid:46) additional F

Original: ( f compute 
Revised: B 1 to compute ˆ

Original: Compute g = ∇ θ L B
Revised: Compute g = ∇ θ L B

Original: ( f θ +ˆ (cid:15) )
Revised: ( f θ +ˆ (cid:15) )

Revised: (cid:46) F and B

Original: 8: update the weights  − 
Revised: 8: update the weights θ ← θ − ηg

Original: A. O PTIMIZING OVER SUBSET B + IS REPRESENTATIVE
Revised: A. O PTIMIZING OVER SUBSET B + IS REPRESENTATIVE

Original: The sharpness-sensitive subset B + are constructedsorting by  (cid:96) ( f θ +ˆ (cid:15) , x i , y i ) − (cid:96) ( f θ , x i , y i ) , which is positively correlated to l ( f θ , x i , y i ) .
Revised: The sharpness-sensitive subset B + is constructed by sorting (cid:96) ( f θ +ˆ (cid:15) , x i , y i ) − (cid:96) ( f θ , x i , y i ) , which is positively correlated to l ( f θ , x i , y i ) .

Original: By the ﬁrst-order Taylor series approximation, (cid:96) ( f θ +ˆ (cid:15) , x i , y i ) − (cid:96) ( f θ , x i , y i ) =
Revised: By a ﬁrst-order Taylor series approximation, (cid:96) ( f θ +ˆ (cid:15) , x i , y i ) − (cid:96) ( f θ , x i , y i ) =

Original: ˆ (cid:15) · ∇ θ (cid:96) ( f θ , x i , y i )
Revised: ˆ (cid:15) · ∇ θ (cid:96) ( f θ , x i , y i )

Original: + o ( (cid:107)
Revised: + o ( (cid:107)

Original: ˆ (cid:15) (cid:107) ) 
Revised: ˆ (cid:15) (cid:107) ) .

Original: By Equation 4, ˆ (cid:15) is the aggregated gradients of each instance in the complete dataset B , i.e.,
Revised: By Equation 4, ˆ (cid:15) is the aggregated gradients of each instance in the complete dataset B , i.e.,

Original: [Equation]
Revised: [Equation]

Original:  which indicates that (cid:96) ( f θ +ˆ (cid:15) , x i , y i ) − (cid:96) ( f θ , x i , y i ) is positively correlated to the gradient ∇ θ (cid:96) ( f θ , x i , y i ) .
Revised: , which indicates that (cid:96) ( f θ +ˆ (cid:15) , x i , y i ) − (cid:96) ( f θ , x i , y i ) is positively correlated to the gradient ∇ θ (cid:96) ( f θ , x i , y i ) .

Original: Li et al.
Revised: Li et al.

Original: (2019) claims that the hard examples in deep learning (the training samples with high training loss) produce gradients with larger magnitudes.
Revised: (2019) claims that the difﬁcult examples in deep learning (the training samples with high training loss) produce gradients with larger magnitudes.

Original: Therefore, (cid:96) ( f θ +ˆ (cid:15) , x i , y i ) − (cid:96) ( f θ , x i , y i ) is positively correlated to l ( f θ , x i , y i ) .
Revised: Therefore, (cid:96) ( f θ +ˆ (cid:15) , x i , y i ) − (cid:96) ( f θ , x i , y i ) is positively correlated to l ( f θ , x i , y i ) .

Original: We also demonstrate the correlation empirically.
Revised: We also demonstrate the correlation empirically.

Original: We conduct experiments to verify Equation 9 and Equation 10.
Revised: We conduct experiments to verify Equation 9 and Equation 10.

Original: In Figure 5, We plot the four losses,
Revised: In Figure 5, We plot the four losses,

Original: L B + ( f θ ) , L B − ( f θ ) , L B + ( f θ +ˆ (cid:15) ) , and L B − ( f θ +ˆ (cid:15) ) w.r.t the epochs.
Revised: L B + ( f θ ) , L B − ( f θ ) , L B + ( f θ +ˆ (cid:15) ) , and L B − ( f θ +ˆ (cid:15) ) w.r.t the epochs.

Original: The experimental results verify that Equation 9 and Equation 10 hold for every training epoch.
Revised: The experimental results verify that Equation 9 and Equation 10 hold for every training epoch.

Original: Moreover, we conduct experiments to demonstrate that optimizing over the subset B + is much more representative than the subset B − .
Revised: Moreover, we conduct experiments to demonstrate that optimizing over the subset B + is much more representative than the subset B − .

Original: We compare the updating gradients computed from B + , B − and a random subset B rand that | B rand | = | B + | = | B − | to those computed from B by calculating the cosine similarity, i.e.
Revised: We compare the updating gradients computed from B + , B − and a random subset B rand that | B rand | = | B + | = | B − | to those computed from B by calculating the cosine similarity inspired by (Du et al., 2019), i.e.

Original: [Equation]
Revised: [Equation]

Revised: θ L B ( f θ +ˆ (cid:15) )) .

Original: In Figure 6, we plot the cosine similarities in each training epoch evaluated with ResNet-18, Wide28-10 on CIFAR10.
Revised: In Figure 6, we plot the cosine similarities in each training epoch evaluated with ResNet-18, Wide28-10 on CIFAR10.

Original: 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 B rand .
Revised: 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 B rand .

Original: A. L INEARITY MEASUREMENT OF SWP
Revised: A. L INEARITY MEASUREMENT OF SWP

Original: As the experimental results in Figure 3 demonstrated, SWP can also improve the accuracy of ESAM compared to SAM.
Revised: As the experimental results in Figure 3 demonstrated, SWP can also improve the accuracy of ESAM compared to SAM.

Original: We will investigate the advantage of SWP in terms of generalization in thefuture research.
Revised: We will investigate the advantage of SWP in terms of generalization in the

Original: Here we provide a discussion about the accuracy improvement contributed by SWP.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: C o s i ne
Revised: C o s i ne

Original: S i m il a i t y
Revised: S i m il a i t y

Original: [Figure]
Revised: [Figure]

Revised: future research.

Revised: Here we provide a discussion about the accuracy improvement contributed by SWP.

Original: A plausible reason for such improvement is that SWP leads to a better inner maximization solved in equation 2.
Revised: A plausible reason for such improvement is that SWP leads to a better inner maximization solved in equation 2.

Original: The current solution of ˆ (cid:15) is approximated by assuming L S ( f θ ) is a linear function.
Revised: The current solution of ˆ (cid:15) is approximated by assuming L S ( f θ ) is a linear function.

Original: Therefore, the ˆ (cid:15) would result in a better inner maximization if L S ( f θ ) is more linear with respect to θ .
Revised: Therefore, the ˆ (cid:15) would result in a better inner maximization if L S ( f θ ) is “more linear” with respect to θ .

Original: Inspired by (Qin et al., 2019), we measure the linearity of the loss function by
Revised: Inspired by (Qin et al., 2019; Yan et al., 2019; 2021), we measure the linearity of the loss function by

Original: [Equation]
Revised: [Equation]

Original: We conduct experiments on the CIFAR10 dataset with ResNet-18 model to verify that SWP can improve the linearity ζ ( (cid:15), B ) of the loss function L S ( f θ ) .
Revised: We conduct experiments on the CIFAR10 dataset with ResNet-18 model to verify that SWP can improve the linearity ζ ( (cid:15), B ) of the loss function L S ( f θ ) .

Original: We compare the linearity of ESAM with the β ranging from { 0 . 2 , 0 .
Revised: We compare the linearity of ESAM with the β ranging from { 0 . 2 , 0 .

Original: 3 , ..., 0 . 9 } to the SAM.
Revised: 3 , ..., 0 . 9 } to the SAM.

Original: The results are demonstrated in Figure 7.
Revised: The results are demonstrated in Figure 7.

Original: It can
Revised: It can

Original: [Figure]
Revised: [Figure]

Original: be shown that SWP will result in a better linearity as the β decreases.
Revised: be shown that SWP will result in a better linearity as the β decreases.

Original: However, decreasing β will also reduce the magnitude of ˆ (cid:15) and thus result in a worse inner maximization in equation 2.
Revised: However, decreasing β will also reduce the magnitude of ˆ (cid:15) and thus result in a worse inner maximization in equation 2.

Original: We observe that β
Revised: We observe that β

Original: = { 0 . 5 , 0 . 6 } is optimal to balance the accuracy and efﬁciency of ESAM.
Revised: = { 0 . 5 , 0 . 6 } is optimal to balance the accuracy and efﬁciency of ESAM.

Original: A. R EDUCED COMPUTATIONAL OVERHEAD CONTRIBUTED BY SWP
Revised: A. R EDUCED COMPUTATIONAL OVERHEAD CONTRIBUTED BY SWP

Original: Formulation We formulate the saved computational overhead contributed by SWP here.
Revised: Formulation We formulate the saved computational overhead contributed by SWP here.

Original: We discuss and examine the computational overhead in the PyTorch framework.
Revised: We discuss and examine the computational overhead in the PyTorch framework.

Original: Suppose that the NN we discussed here has N layers.
Revised: Suppose that the NN we discussed here has N layers.

Original: 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.
Revised: 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.

Original: SWP select each basic parameter unit to compute gradients (i.e. requries grad =
Revised: SWP select each basic parameter unit to compute gradients (i.e. requries grad =

Original: True) with probability β .
Revised: True) with probability β .

Original: We use g ( N, β ) to measure the saved computational overhead in terms of percentage contributed by SWP compared to the vanilla SAM.
Revised: We use g ( N, β ) to measure the saved computational overhead in terms of percentage contributed by SWP compared to the vanilla SAM.

Original: The computational overhead g ( N, β ) is reduced not just by calculating gradients, but also by the storing and hooking gradients.
Revised: The computational overhead g ( N, β ) is reduced not just by calculating gradients, but also by the storing and hooking gradients.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Next we discuss the saved computational overhead that stems from the calculation operations in the general case.
Revised: Next we discuss the saved computational overhead that stems from the calculation operations in the general case.

Original: 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.
Revised: 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.

Original: Suppose the n th layer has K ( n ) basic parameter units,  we have p n = (1 − β ) K ( n ) · p n − 1 .
Revised: Suppose the n th layer has K ( n ) basic parameter units, let p n be the calculation-free rate of a certain parameter unit in n th layer, we have p n = (1 − β ) K ( n ) · p n − 1 .

Original: Therefore, p n = (1 − β ) (cid:80) n j =1 K ( j ) ≈ (1 − β ) n ¯ K , where ¯ K = 1 N (cid:80) Nj =1 K ( j ) .
Revised: Therefore, p n = (1 − β ) (cid:80) n j =1 K ( j ) ≈ (1 − β ) n ¯ K , where ¯ K = 1 N (cid:80) Nj =1 K ( j ) .

Original: Assumed that the computations of each layer is the same, by summing up the saved computational overhead of all parameters, we have
Revised: Assumed that the computations of each layer are the same, by summing up the saved computational overhead of all parameters, we have

Original: [Equation]
Revised: [Equation]

Original: where k 1 , k 2 are determined by the computing time of calculation, k (cid:48)
Revised: where k 1 , k 2 are determined by the computing time of calculation, k (cid:48)

Original: 2 =
Revised: 2 =

Original: [Equation]
Revised: [Equation]

Original: However, the commonly used DNNs’ architectures such as ResNet only have one basic parameter unit in each layer.
Revised: However, the commonly used DNNs’ architectures such as ResNet only have one basic parameter unit in each layer.

Original: Besides, each layer of them is connected in serial with the next layer.
Revised: Besides, each layer of them is connected in serial with the next layer.

Original: Then,let p n be the calculation-free rate of a certain parameter unit in n th layer, we have p n = (1 − β ) · p n − 1 .
Revised: Then, we have p n = (1 − β ) · p n − 1 .

Original: Therefore, p n = (1 − β ) n .
Revised: Therefore, p n = (1 − β ) n .

Original: The saved calculation contributed by the parameter unit is 1 N p n .
Revised: The saved calculation contributed by the parameter unit is 1 N p n .

Original: By summing up the saved computational overhead of all parameters, we have
Revised: By summing up the saved computational overhead of all parameters, we have

Original: [Equation]
Revised: [Equation]

Original: where k 1 , k 2 are determined by the computing time of calculation, storing and hooking gradients.
Revised: where k 1 , k 2 are determined by the computing time of calculation, storing and hooking gradients.

Original: Toy Example We conducted a toy example in Cifar10 dataset with two MLPs.
Revised: Toy Example We conducted a toy example on the CIFAR10 dataset with two MLPs.

Original: Each fully connected layer of MLP is the same with a size of 3 , 000 for both in and out features.
Revised: Each fully connected layer of MLP is the same with a size of 3 , 000 for both in and out features.

Original: The ﬁrst 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.
Revised: The ﬁrst 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.

Original: We examined N = { 50 , 75 , 100 , 125 } and β = { 0 . 1 , ..., 0 .
Revised: We examined N = { 50 , 75 , 100 , 125 } and β = { 0 . 1 , ..., 0 .

Original: 9 , 1 . 0 } to record the saved computational overhead in percentage.
Revised: 9 , 1 . 0 } to record the saved computational overhead in percentage.

Original: Part of the results are reported in Table A.4.
Revised: Part of the results are reported in Table A.4.

Original: By linear regression, we have k 1 = 0 .
Revised: By linear regression, we have k 1 = 0 .

Original: 3185 , k 2 = 0 .
Revised: 3185 , k 2 = 0 .

Original: 1310 , and the returned R 2 = 0 . 9983 .
Revised: 1310 , and the returned R 2 = 0 . 9983 .

Original: 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.
Revised: 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.

Original: We examined N = { 35 , 50 , 65 , 75 } and β = { 0 . 1 , ..., 0 .
Revised: We examined N = { 35 , 50 , 65 , 75 } and β = { 0 . 1 , ..., 0 .

Original: 9 , 1 . 0 } to record the saved computational overhead in percentage.
Revised: 9 , 1 . 0 } to record the saved computational overhead in percentage.

Original: Part of the results are reported in Table A.4.
Revised: Part of the results are reported in Table A.4.

Original: By linear regression, we have k 1 = 0 .
Revised: By linear regression, we have k 1 = 0 .

Original: 3143 , k 2 = 0 .
Revised: 3143 , k 2 = 0 .

Original: 0737 , and the returned R 2 = 0 . 9989 .
Revised: 0737 , and the returned R 2 = 0 . 9989 .

Original: The above experimental results verify the formulation of the reduced computational overhead contributed by SWP in equation 12 and equation 11.
Revised: The above experimental results verify the formulation of the reduced computational overhead contributed by SWP in equation 11 and equation 12.

Original: [Table] [Table]
Revised: [Table] [Table]

Original: A. V ISUALIZATION OF L OSS L ANDSCAPES WITH RESPECT TO ADVERSARIAL WEIGHT PERTURBATIONS
Revised: A. V ISUALIZATION OF L OSS L ANDSCAPES WITH RESPECT TO ADVERSARIAL WEIGHT PERTURBATIONS

Original: We visualize the sharpness of the ﬂat minima with respect to adversarial weight perturbations of
Revised: We visualize the sharpness of the ﬂat minima with respect to adversarial weight perturbations of

Original: SGD,SAM and ESAM on the Cifar10 dataset.
Revised: SGD,SAM and ESAM on the Cifar10 dataset.

Original: The x - and y -axes represent two orthogonal adversarial weight perturbations, which are η ∇ θ L B x ( f θ ) and η ∇ θ L B y ( f θ ) respectively, where η is the learning rate during training.
Revised: The x - and y -axes represent two orthogonal adversarial weight perturbations, which are η ∇ θ L B x ( f θ ) and η ∇ θ L B y ( f θ ) respectively, where η is the learning rate during training.

Original: B x and B y are the randomly sampled subsets of batch B, and | B x | =
Revised: B x and B y are the randomly sampled subsets of batch B

Revised: [Figure]

Original:  | B y | = 12 | B | , B x ∪ B y = B .
Revised: and | B x | = | B y | = 12 | B | , B x ∪ B y = B .

Original: We display the loss landscape in Figure 8, which demonstrates that both SAM and ESAM improve the sharpness signiﬁcantly in comparison to SGD.
Revised: We display the loss landscape in Figure [Figure] which demonstrates that both SAM and ESAM improve the sharpness signiﬁcantly in comparison to SGD

Original: [Figure]
Revised: [Figure]

Original: A. E VALUATION OF ESAM ON V I T-S/ 
Revised: A. E VALUATION OF ESAM ON V I T-S/

Original: SAM has also been demonstrated to be effective on the new vision Transformer(ViT) architecture (Chen et al., 2021).
Revised: SAM has also been demonstrated to be effective on the new vision Transformer(ViT) architecture (Chen et al., 2021).

Original: Therefore, we also evaluate ESAM with ViT-S/16 on ImageNet Datasets.
Revised: Therefore, we also evaluate ESAM with ViT-S/16 on ImageNet Datasets.

Original: We use β = 0 .
Revised: We use β = 0 .

Original: 5 and γ = 0 .
Revised: 5 and γ = 0 .

Original: 7 for ViT-S/16, which share the same hyperparameters as ResNet-50 and
Revised: 7 for ViT-S/16, which share the same hyperparameters as ResNet-50 and

Original: ResNet-101 in section 3.1.
Revised: ResNet-101 in section 3.1.

Original: The results are reported in Table [Figure] which indicate that ESAM can still be effective to improve efﬁciency in ViT-S/16 architecture.
Revised: The results are reported in Table 6, which indicate that ESAM can still be effective to improve efﬁciency in ViT-S/16 architecture.

Original: In particular, ESAM-SWP achieves much better accuracy than SAM(80.88% v.s. 80.34%).
Revised: In particular, ESAM-SWP achieves much better accuracy than SAM (80.88% v.s. 80.34%).

Original: [Table]
Revised: [Table]

Original: A. T RAINING D ETAILS
Revised: A. T RAINING D ETAILS

Original: We tune the training parameters of SGD, SAM, and ESAM, by using grid searches.
Revised: We tune the training parameters of SGD, SAM, and ESAM, by using grid searches.

Original: 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 } .
Revised: 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 } .

Original: This is done to attain the best accuracies.
Revised: This is done to attain the best accuracies.

Original: The exact training hyperparameters are reported in Table 7.
Revised: The exact training hyperparameters are reported in Table 7.

Original: On the ImageNet datasets, limited by the computing resource, we follow and slightly modify the optimal hyperparameters as suggested by
Revised: On the ImageNet datasets, limited by the computing resource, we follow and slightly modify the optimal hyperparameters as suggested by

Original: Chen et al.
Revised: Chen et al.

Original: (2021) for SGD, SAM and ESAM.
Revised: (2021) for SGD, SAM and ESAM.

Original: The exact training hyperparameters are reported in
Revised: The exact training hyperparameters are reported in

Original: Table 8.
Revised: Table 8.

Original: [Table] [Table]
Revised: [Table] [Table]