Title: Enhancing Sharpness-Aware Optimization Through Variance Suppression
Venue: NeurIPS 2023
Accepted: yes
Keywords: generalization

optimization

neural networks
Abstract: Sharpness-aware minimization (SAM) has well documented merits in enhancing generalization of deep neural networks, even without sizable data augmentation. Embracing the geometry of the loss function, where neighborhoods of 'flat minima' heighten generalization ability, SAM seeks 'flat valleys' by minimizing the maximum loss caused by an *adversary* perturbing parameters within the neighborhood.
Although critical to account for sharpness of the loss function, such an '*over-friendly* adversary' can curtail the outmost level of generalization. The novel approach of this contribution fosters stabilization of adversaries through *variance suppression* (VaSSO) to avoid such friendliness. VaSSO's *provable* stability safeguards its numerical improvement over SAM in model-agnostic tasks, including image classification and machine translation. In addition, experiments confirm that VaSSO endows SAM with robustness against high levels of label noise. Code is available at https://github.com/BingcongLi/VaSSO.

Enhancing Sharpness-Aware Optimization Through Variance Suppression

Abstract

Sharpness-aware minimization (SAM) has well documented merits in enhancing generalization of deep neural networks, even without sizable data augmentation. Embracing the geometry of the loss function, where neighborhoods of 'flat minima' heighten generalization ability, SAM seeks 'flat valleys' by minimizing the maximum loss caused by an adversary perturbing parameters within the neighborhood. Although critical to account for sharpness of the loss function, such an ' over-friendly adversary' can curtail the outmost level of generalization. The novel approach of this contribution fosters stabilization of adversaries through variance suppression (VaSSO) to avoid such friendliness. VaSSO's provable stability safeguards its numerical improvement over SAM in model-agnostic tasks, including image classification and machine translation. In addition, experiments confirm that VaSSO endows SAM with robustness against high levels of label noise. Code is available at https://github.com/BingcongLi/VaSSO .

1 Introduction

Despite deep neural networks (DNNs) have advanced the concept of 'learning from data,' and markedly improved performance across several applications in vision and language (Devlin et al., 2018; Tom et al., 2020), their overparametrized nature renders the tendency to overfit on training data (Zhang et al., 2021). This has led to concerns in generalization, which is a practically underscored perspective yet typically suffers from a gap relative to the training performance.

Improving generalizability is challenging. Common approaches include (model) regularization and data augmentation (Srivastava et al., 2014). While it is the default choice to integrate regularization such as weight decay and dropout into training, these methods are often insufficient for DNNs especially when coping with complicated network architectures (Chen et al., 2022). Another line of effort resorts to suitable optimization schemes attempting to find a generalizable local minimum. For example, SGD is more preferable than Adam on certain overparameterized problems since it converges to maximum margin solutions (Wilson et al., 2017). Decoupling weight decay from Adam also empirically facilitates generalizability (Loshchilov and Hutter, 2017). Unfortunately, the underlying mechanism remains unclear, and whether the generalization merits carry over to other intricate learning tasks calls for additional theoretical elaboration.

Our main focus, sharpness aware minimization/optimization (SAM), is a highly compelling optimization approach that facilitates state-of-the-art generalizability by exploiting sharpness of loss landscape (Foret et al., 2021; Chen et al., 2022). A high-level interpretation of sharpness is how violently the loss fluctuates within a neighborhood. It has been shown through large-scale empirical studies that sharpness-based measures highly correlate with generalization (Jiang et al., 2019). Several works

have successfully explored sharpness for generalization advances. For example, Keskar et al. (2016) suggests that the batchsize of SGD impresses solution flatness. Entropy SGD leverages local entropy in search of a flat valley (Chaudhari et al., 2017). Different from prior works, SAM induces flatness by explicitly minimizing the adversarially perturbed loss, defined as the maximum loss of a neighboring area. Thanks to such a formulation, SAM has elevated generalization merits among various tasks in vision and language domains (Chen et al., 2022; Zhang et al., 2022). The mechanism fertilizing SAM's success is also theoretically investigated based on arguments of implicit regularization; see e.g., (Andriushchenko and Flammarion, 2022; Wen et al., 2023; Bartlett et al., 2022).

The adversary perturbation, or adversary for short, is central to SAM's heightened generalization because it effectively measures sharpness through the loss difference with original model (Foret et al., 2021; Zhuang et al., 2022; Kim et al., 2022). In practice however, this awareness on sharpness is undermined by what we termed friendly adversary . Confined by the stochastic linearization for computational efficiency, SAM's adversary only captures the sharpness for a particular minibatch of data, and can become a friend on other data samples. Because the global sharpness is not approached accurately, the friendly adversary precludes SAM from attaining its utmost generalizability. The present work advocates variance suppressed sharpness aware optimization (VaSSO 1 ) to alleviate 'friendliness' by stabilizing adversaries. With its provable stabilized adversary, VaSSO showcases favorable numerical performance on various deep learning tasks.

All in all, our contribution is summarized as follows.

v We find that the friendly adversary discourages generalizability of SAM. This challenge is catastrophic in our experiments - it can completely wipe out the generalization merits.

v A novel approach, VaSSO, is proposed to tackle this issue. VaSSO is equipped with what we termed variance suppression to streamline a principled means for stabilizing adversaries. The theoretically guaranteed stability promotes refined global sharpness estimates, thereby alleviating the issue of friendly adversary.

v Aside result is tighter convergence analyses for VaSSO and SAM that i) remove the bounded gradient assumption; and ii) deliver a more flexible choice for hyperparameters.

v Numerical experiments confirm the merits of stabilized adversary in VaSSO. It is demonstrated on image classification and neural machine translation tasks that VaSSO is capable of i) improving generalizability over SAM model-agnostically; and ii) nontrivially robustifying neural networks under the appearance of large label noise.

Notation . Bold lowercase (capital) letters denote column vectors (matrices); ‖ x ‖ stands for glyph[lscript] 2 norm of vector x ; and 〈 x , y 〉 is the inner product of x and y . S ρ ( x ) denotes the surface of a ball with radius ρ centered at x , i.e., S ρ ( x ) := { x + ρ u | ‖ u ‖ = 1 } .

2 The known, the good, and the challenge of SAM

This section starts with a brief recap of SAM (i.e., the known), followed with refined analyses and positive results regarding its convergence (i.e., the good). Lastly, the friendly adversary issue is explained in detail and numerically illustrated.

2.1 The known

Targeting at a minimum in flat basin, SAM enforces small loss around the entire neighborhood in the parameter space (Foret et al., 2021). This idea is formalized by a minimax problem

where ρ is the radius of considered neighborhood, and the nonconvex objective is defined as f ( x ) := E B [ f B ( x )] . Here, x is the neural network parameter, and B is a random batch of data. The merits of such a formulation resides in its implicit sharpness measure max ‖ glyph[epsilon1] ‖≤ ρ f ( x + glyph[epsilon1] ) -f ( x ) , which effectively drives the optimization trajectory towards the desirable flat valley (Kim et al., 2022).

Algorithm 1 Generic form of SAM

1: Initialize: x 0 , ρ 2: for t = 0 , . . . , T -1 do 3: Sample a minibatch B t , and define stochastic gradient on B t as g t ( · ) 4: Find glyph[epsilon1] t ∈ S ρ ( 0 ) via stochastic linearization; e.g., (4) for VaSSO or (3) for SAM 5: Calculate stochastic gradient g t ( x t + glyph[epsilon1] t ) 6: Update model via x t +1 = x t -η g t ( x t + glyph[epsilon1] t ) 7: end for 8: Return: x T

The inner maximization of (1) has a natural interpretation as finding an adversary . Critical as it is, obtaining an adversary calls for stochastic linearization to alleviate computational concerns, i.e.,

where linearization ( a ) relies on the first order Taylor expansion of f ( x t + glyph[epsilon1] ) . This is typically accurate given the choice of a small ρ . A stochastic gradient g t ( x t ) then substitutes ∇ f ( x t ) in ( b ) to downgrade the computational burden of a full gradient. Catalyzed by the stochastic linearization in (2), it is possible to calculate SAM's adversary in closed-form

SAM then adopts the stochastic gradient of adversary g t ( x t + glyph[epsilon1] t ) to update x t in a SGD fashion. A step-by-step implementation is summarized in Alg. 1, where the means to find an adversary in line 4 is presented in a generic form in order to unify the algorithmic framework with later sections.

2.2 The good

To provide a comprehensive understanding about SAM, this subsection focuses on Alg. 1, and establishes its convergence for (1). Some necessary assumptions are listed below, all of which are common for nonconvex stochastic optimization problems (Ghadimi and Lan, 2013; Bottou et al., 2016; Mi et al., 2022; Zhuang et al., 2022).

Assumption 1 (lower bounded loss) . f ( x ) is lower bounded, i.e., f ( x ) ≥ f ∗ , ∀ x .

Assumption 2 (smoothness) . The stochastic gradient g ( x ) is L -Lipschitz, i.e., ‖ g ( x ) -g ( y ) ‖ ≤ L ‖ x -y ‖ , ∀ x , y .

Assumption 3 (bounded variance) . The stochastic gradient g ( x ) is unbiased with bounded variance, that is, E [ g ( x ) | x ] = ∇ f ( x ) and E [ ‖ g ( x ) -∇ f ( x ) ‖ 2 | x ] = σ 2 for some σ > 0 .

The constraint of (1) is never violated since ‖ glyph[epsilon1] t ‖ = ρ holds for each t ; see line 4 in Alg. 1. Hence, the convergence of SAM pertains to the behavior of objective, where a tight result is given below.

Theorem 1 (SAM convergence) . Suppose that Assumptions 1 - 3 hold. Let η t ≡ η = η 0 √ ≤ 2 ,

T 3 L

and ρ = ρ 0 √ T . Then with c 0 = 1 -3 Lη 2 (clearly 0 < c 0 < 1 ), Alg. 1 guarantees that

The convergence rate of SAM is the same as SGD up to constant factors, where the detailed expression hidden under big O notation can be found in Appendix D. Our results eliminate the need for the bounded gradient assumption compared to existing analyses in (Mi et al., 2022; Zhuang et al., 2022). Moreover, Theorem 1 enables a much larger choice of ρ = O ( T -1 / 2 ) relative to (Andriushchenko and Flammarion, 2022), where the latter only supports ρ = O ( T -1 / 4 ) .

A message from Theorem 1 is that any adversary satisfying glyph[epsilon1] t ∈ S ρ ( 0 ) ensures converge. Because the surface S ρ ( 0 ) is a gigantic space, it challenges the plausible optimality of the adversary and poses a natural question is it possible to find a more powerful adversary for generalization advances?

ResNet18

ResNet34

0.99

1.00

1.01

1.02

1.03

1.04

1.0

1.0

1.004

1.028

0.999

1.002

0.999

0.999

SGD

SAM

ResNet18

ResNet34

0.99

1.00

1.01

1.02

1.03

1.04

1.0

1.0

1.003

1.028

1.002

1.023

0.997

1.018

0.998

1.011

SGD

SAM

SAM-var 10

6

SAM-var 10

4

SAM-var 10

2

m=b

m=b/2

m=b/4

m=b/8

34.5

34.6

34.7

34.8

34.9

35.0

m-SAM

(a)

(b)

(c)

2.3 The challenge: friendly adversary

Adversary to one minibatch is a friend of others. SAM's adversary is 'malicious' for minibatch B t but not necessarily for other data because it only safeguards f B t ( x t + glyph[epsilon1] t ) -f B t ( x t ) ≥ 0 for a small ρ . In fact, it can be shown that f B ( x t + glyph[epsilon1] t ) -f B ( x t ) ≤ 0 whenever the stochastic gradients do not align well, i.e., 〈 g t ( x t ) , g B ( x t ) 〉 ≤ 0 . Note that such misalignment is common because of the variance in massive training datasets. This issue is referred to as friendly adversary , and it implies that the adversary glyph[epsilon1] t cannot accurately depict the global sharpness of x t . Note that the 'friendly adversary' also has a more involved interpretation, that is, g t ( x t ) falls outside the column space of Hessian at convergence; see more discussions after (Wen et al., 2023, Definition 4.3). This misalignment of higher order derivatives undermines the inductive bias of SAM, thereby worsening generalization.

To numerically visualize the catastrophic impact of the friendly adversary, we manually introduce one by replacing line 4 of Alg. 1 as ˜ glyph[epsilon1] t = ρ ˜ g t ( x t ) / ‖ ˜ g t ( x t ) ‖ , where ˜ g t denotes the gradient on ˜ B t , a randomly sampled batch of the same size as B t . This modified approach is denoted as SAM-db, and its performance for i) ResNet-18 on CIFAR10 and ii) ResNet-34 on CIFAR100 2 can be found in Fig. 1(a). Note that the test accuracy is normalized relative to SGD for the ease of visualization. It is evident that the friendly ˜ glyph[epsilon1] t in SAM-db almost erases the generalization benefits entirely.

Source of friendly adversary. The major cause to the friendly adversary attributes to the gradient variance, which equivalently translates to the lack of stability in SAM's stochastic linearization (2 b ) . An illustrative three dimensional example is shown in Fig. 2, where we plot the adversary glyph[epsilon1] t obtained from different g t realization in (2 b ) . The minibatch gradient is simulated by adding Gaussian noise to the true gradient. When the signal to noise ration (SNR) is similar to a practical scenario (ResNet-18 on CIFAR10 shown in Fig. 2 (e)), it can be seen in Fig. 2 (c) and (d) that the adversaries almost uniformly spread over the norm ball, which strongly indicates the deficiency for sharpness evaluation.

Friendly adversary in the lens of Frank Wolfe. An additional evidence in supportive to SAM's friendly adversary resides in its connection to stochastic Frank Wolfe (SFW) that also heavily relies on stochastic linearization (Reddi et al., 2016). The stability of SFW is known to be vulnerable - its convergence cannot be guaranteed without a sufficient large batchsize. As thoroughly discussed in Appendix A, the means to obtain adversary in SAM is tantamount to one-step SFW with a constant batchsize. This symbolizes the possible instability of SAM's stochastic linearization.

2.4 A detailed look at friendly adversaries

The gradient variance is major cause to SAM's friendly adversary and unstable stochastic linearization, however this at first glance seems to conflict with an empirical note termed m -sharpness, stating that the benefit of SAM is clearer when glyph[epsilon1] t is found using subsampled B t of size m (i.e., larger variance).

Since m -sharpness highly hinges upon the loss curvature, it is unlikely to hold universally. For example, a transformer is trained on IWSLT-14 dataset, where the test performance (BLEU) decreases with smaller m even if we have tuned ρ carefully; see Fig. 1(c). On the theoretical side, an example is provided in (Andriushchenko and Flammarion, 2022, Sec. 3) to suggest that m -sharpness is not

x

1.0

0.5

0.0

0.5

1.0

y

1.0

0.5

0.0

0.5

1.0

z

1.0

0.5

0.0

0.5

1.0

x

1.0

0.5

0.0

0.5

1.0

y

1.0

0.5

0.0

0.5

1.0

z

1.0

0.5

0.0

0.5

1.0

x

1.0

0.5

0.0

0.5

1.0

y

1.0

0.5

0.0

0.5

1.0

z

1.0

0.5

0.0

0.5

1.0

x

1.0

0.5

0.0

0.5

1.0

y

1.0

0.5

0.0

0.5

1.0

z

1.0

0.5

0.0

0.5

1.0

(a) SNR

= 5

(b) SNR

= 1

(c) SNR

= 0

.

1

(d) SNR

= 0

.

01

(e) SNR in practice

necessarily related with sharpness or generalization. Moreover, there also exists specific choice for m such that the m -sharpness formulation is ill-posed. We will expand on this in Appendix B.

Even in the regime where m -sharpness is empirically observed such as ResNet-18 on CIFAR10 and ResNet-34 on CIFAR100, we show through experiments that m -sharpness is not a consequence of gradient variance, thus not contradicting with the friendly adversary issue tackled in this work.

Observation 1. Same variance, different generalization. Let m = 128 and batchsize b = 128 . Recall the SAM-db experiment in Fig. 1(a). If m -sharpness is a direct result of gradient variance, it is logical to expect SAM-db has comparable performance to SAM simply because their batchzises (hence variance) for finding adversary are the same. Unfortunately, SAM-db degrades accuracy. We further increase the variance of ˜ g t ( x t ) by setting m = 64 . The resultant algorithm is denoted as SAM-db-m/2. It does not catch with SAM and performs even worse than SAM-db. These experiments validate that variance/stability correlates with friendly adversary instead of m -sharpness.

Observation 2. Enlarged variance degrades generalization. We explicitly increase variance when finding adversary by adding Gaussian noise ζ to g t ( x t ) , i.e., ˆ glyph[epsilon1] t = ρ g t ( x t )+ ζ ‖ g t ( x t )+ ζ ‖ . After tuning the best ρ to compensate the variance of ζ , the test performance is plotted in Fig. 1(b). It can be seen that the generalization merits clearly decrease with larger variance on both ResNet-18 and ResNet-34. This again illustrates that the plausible benefit of m -sharpness does not stem from increased variance.

In sum, observations 1 and 2 jointly suggest that gradient variance correlates with friendly adversary rather than m -sharpness, where understanding the latter is beyond the scope of current work.

3 Variance-supressed sharpness-aware optimization (VaSSO)

This section advocates variance suppression to handle the friendly adversary. We start with the design of VaSSO, then establish its stability. We also touch upon implementation and possible extensions.

3.1 Algorithm design and stability analysis

A straightforward attempt towards stability is to equip SAM's stochastic linearization with variance reduced gradients such as SVRG and SARAH (Johnson and Zhang, 2013; Nguyen et al., 2017; Li et al., 2019). However, the requirement to compute a full gradient every a few iterations is infeasible and hardly scales well for tasks such as training DNNs.

The proposed variance suppression (VaSSO) overcomes this computational burden through a novel yet simple stochastic linearization. For a prescribed θ ∈ (0 , 1) , VaSSO is summarized below

Compared with (2) of SAM, the key difference is that VaSSO relies on slope d t for a more stable stochastic linearization as shown in (4b). The slope d t is an exponentially moving average (EMA) of { g t ( x t ) } t such that the change over consecutive iterations is smoothed. Noticing that glyph[epsilon1] t and d t share the same direction, the relatively smoothed { d t } t thus imply the stability of { glyph[epsilon1] t } t in VaSSO. Moreover, as d t processes information of different minibatch data, the global sharpness can be captured in a principled manner to alleviate the friendly adversary challenge.

To theoretically characterize the effectiveness of VaSSO, our first result considers d t as a qualified strategy to estimate ∇ f ( x t ) , and delves into its mean square error (MSE).

Theorem 2 (Variance suppression) . Suppose that Assumptions 1 - 3 hold. Let Alg. 1 equip with i) glyph[epsilon1] t obtained by (4) with θ ∈ (0 , 1) ; and, ii) η t and ρ selected the same as Theorem 1. VaSSO guarantees that the MSE of d t is bounded by

Because SAM's gradient estimate has a looser bound on MSE (or variance), that is, E [ ‖ g t -∇ f ( x t ) ‖ 2 ] ≤ σ 2 , the shrunk MSE in Theorem 2 justifies the name of variance suppression.

Next, we quantify the stability invoked with the suppressed variance. It is convenient to start with necessary notation. Define the quality of a stochastic linearization at x t with slope v as L t ( v ) := max ‖ glyph[epsilon1] ‖≤ ρ f ( x t ) + 〈 v , glyph[epsilon1] 〉 . For example, L t ( d t ) and L t ( g t ( x t ) ) are quality of VaSSO and SAM, respectively. Another critical case of concern is L t ( ∇ f ( x t ) ) . It is shown in (Zhuang et al., 2022) that L t ( ∇ f ( x t ) ) ≈ max ‖ glyph[epsilon1] ‖≤ ρ f ( x t + glyph[epsilon1] ) given a small ρ . Moreover, L t ( ∇ f ( x t ) ) -f ( x t ) is also an accurate approximation to the sharpness (Zhuang et al., 2022). These observations safeguard L t ( ∇ f ( x t )) as the anchor when analyzing the stability of SAM and VaSSO.

Definition 1 ( δ -stability) . A stochastic linearization with slope v is said to be δ -stable if its quality satisfies E [ |L t ( v ) -L t ( ∇ f ( x t )) | ] ≤ δ .

A larger δ implies a more friendly adversary, hence is less preferable. We are now well-prepared for our main results on adversary's stability.

Theorem 3 (Adversaries of VaSSO is more stable than SAM.) . Suppose that Assumptions 1 3 hold. Under the same hyperparameter choices as Theorem 2, the stochastic linearization is [ √ θρσ + O ( ρσ θT 1 / 4 ) ] -stable for VaSSO, while ρσ -stable in SAM.

Theorem 3 demonstrates that VaSSO alleviates the friendly adversary problem by promoting stability. Qualitatively, VaSSO is roughly √ θ ∈ (0 , 1) times more stable relative to SAM, since the term in big O notation is negligible given a sufficiently large T . Theorem 3 also guides the choice of θ -preferably small but not too small, otherwise the term in big O is inversely amplified.

3.2 Additional perspectives of VaSSO

Having discussed about the stability, this subsection proceeds with other aspects of VaSSO for a thorough characterization.

Convergence. Summarized in the following corollary, the convergence of VaSSO can be pursued as a direct consequence of Theorem 1. The reason is that glyph[epsilon1] t ∈ S ρ ( 0 ) is satisfied by (4).

Corollary 1 (VaSSO convergence) . Suppose that Assumptions 1 - 3 hold. Choosing η t and ρ the same as Theorem 1, then for any θ ∈ (0 , 1) , VaSSO ensures that

VaSSO better reflects sharpness around optimum. Consider a near optimal region where ‖∇ f ( x t ) ‖ → 0 . Suppose that we are in a big data regime where g t ( x t ) = ∇ f ( x t ) + ζ for some Gaussian random variable ζ . The covariance matrix of ζ is assumed to be σ 2 I for simplicity, but our discussion can be extended to more general scenarios using arguments from von Mises-Fisher statistics (Mardia and Jupp, 2000). SAM has difficulty to estimate the flatness in this case, since glyph[epsilon1] t ≈ ρ ζ / ‖ ζ ‖ uniformly distributes over S ρ ( 0 ) regardless of whether the neighboring region is sharp. On the other hand, VaSSO has glyph[epsilon1] t = ρ d t / ‖ d t ‖ . Because { g τ ( x τ ) } τ on sharper valley tend to have larger magnitude, their EMA d t is helpful for distinguishing sharp with flat valleys.

Memory efficient implementation. Although at first glance VaSSO has to keep both d t and glyph[epsilon1] t in memory, it can be implemented in a much more memory efficient manner. It is sufficient to store d t together with a scaler ‖ d t ‖ so that glyph[epsilon1] t can be recovered on demand through normalization; see (4b). Hence, VaSSO has the same memory consumption as SAM.

Extensions. VaSSO has the potential to boost the performance of other SAM family approaches by stabilizing their stochastic linearization through variance suppression. For example, adaptive SAM methods (Kwon et al., 2021; Kim et al., 2022) ensure scale invariance for SAM, and GSAM (Zhuang et al., 2022) jointly minimizes a surrogated gap with (1). Nevertheless, these SAM variants leverage stochastic linearization in (2). It is thus envisioned that VaSSO can also alleviate the possible friendly adversary issues therein. Confined by computational resources, we only integrate VaSSO with GSAM in our experiments, and additional evaluation has been added into our research agenda.

4 Numerical tests

To support our theoretical findings and validate the powerfulness of variance suppression, this section assesses generalization performance of VaSSO via various learning tasks across vision and language domains. All experiments are run on NVIDIA V100 GPUs.

4.1 Image classification

Benchmarks. Building on top of the selected base optimizer such as SGD and AdamW (Kingma and Ba, 2014; Loshchilov and Hutter, 2017), the test accuracy of VaSSO is compared with SAM and two adaptive approaches, ASAM and FisherSAM (Foret et al., 2021; Kwon et al., 2021; Kim et al., 2022).

CIFAR10. Neural networks including VGG-11, ResNet-18, WRN-28-10 and PyramidNet-110 are trained on CIFAR10. Standard implementation including random crop, random horizontal flip, normalization and cutout (Devries and Taylor, 2017) are leveraged for data augmentation. The first three models are trained for 200 epochs with a batchsize of 128 , and PyramidNet-110 is trained for 300 epochs using batchsize 256 . Cosine learning rate schedule is applied in all settings. The first three models use initial learning rate 0 . 05 , and PyramidNet adopts 0 . 1 . Weight decay is chosen as 0 . 001 for SAM, ASAM, FisherSAM and VaSSO following (Du et al., 2022a; Mi et al., 2022), but 0 . 0005 for SGD. We tune ρ from { 0 . 01 , 0 . 05 , 0 . 1 , 0 . 2 , 0 . 5 } for SAM and find that ρ = 0 . 1 gives the best results for ResNet and WRN, ρ = 0 . 05 and ρ = 0 . 2 suit best for and VGG and PyramidNet, respectively. ASAM and VaSSO adopt the same ρ as SAM. FisherSAM uses the recommended ρ = 0 . 1 (Kim et al., 2022). For VaSSO, we tune θ = { 0 . 4 , 0 . 9 } and report the best accuracy although VaSSO with both parameters outperforms SAM. We find that θ = 0 . 4 works the best for ResNet-18 and WRN-28-10 while θ = 0 . 9 achieves the best accuracy in other cases.

It is shown in Table 1 that VaSSO offers 0 . 2 to 0 . 3 accuracy improvement over SAM in all tested scenarios except for PyramidNet-110, where the improvement is about 0 . 1 . These results illustrate that suppressed variance and the induced stabilized adversary are indeed beneficial for generalizability.

CIFAR100. The training setups on this dataset are the same as those on CIFAR10, except for the best choice for ρ of SAM is 0 . 2 . The numerical results are listed in Table 2. It can be seen that SAM has significant generalization gain over SGD, and this gain is further amplified by VaSSO. On all tested models, VaSSO improves the test accuracy of SAM by 0 . 2 to 0 . 3 . These experiments once again corroborate the generalization merits of VaSSO as a blessing of the stabilized adversary.

ImageNet. Next, we investigate the performance of VaSSO on larger scale experiments by training ResNet-50 and ViT-S/32 on ImageNet (Deng et al., 2009). Implementation details are deferred to Appendix C. Note that the baseline optimizer is SGD for ResNet and AdamW for ViT. VaSSO is also integrated with GSAM (V+G) to demonstrate that the variance suppression also benefits other SAM type approaches (Zhuang et al., 2022). For ResNet-50, it can be observed that vanilla VaSSO

outperforms other SAM variants, and offers a gain of 0 . 26 over SAM. V+G showcases the best performance with a gain of 0 . 28 on top of GSAM. VaSSO and V+G also exhibit the best test accuracy on ViT-S/32, where VaSSO improves SAM by 0.56 and V+G outperforms GSAM by 0.19. These numerical improvement demonstrates that stability of adversaries is indeed desirable.

4.2 Neural machine translation

Having demonstrated the benefits of a suppressed variance on vision tasks, we then test VaSSO on German to English translation using a Transformer (Vaswani et al., 2017) trained on IWSLT-14 dataset (Cettolo et al., 2014). The fairseq implementation is adopted. AdamW is chosen as base optimizer in SAM and VaSSO because of its improved performance over SGD. The learning rate of AdamW is initialized to 5 × 10 -4 and then follows an inverse square root schedule. For momentum, we choose β 1 = 0 . 9 and β 2 = 0 . 98 . Label smoothing is also applied with a rate of 0 . 1 . Hyperparameter ρ is tuned for SAM from { 0 . 01 , 0 . 05 , 0 . 1 , 0 . 2 } , and ρ = 0 . 1 performs the best. The same ρ is picked for ASAM and VaSSO as well.

The validation perplexity and test BLEU scores are shown in Table 4. It can be seen that both SAM and ASAM have better performance on validation perplexity and BLEU relative to AdamW. Although VaSSO with θ = 0 . 9 has slightly higher validation perplexity, its BLEU score outperforms SAM and ASAM. VaSSO with θ = 0 . 4 showcases the best generalization performance on this task, providing a 0 . 22 improvement on BLEU score relative to AdamW. This aligns with Theorems 2 and 3, which suggest that a small θ is more beneficial to the stability of adversary.

4.3 Additional tests

Additional experiments are conducted to corroborate the merits of suppressed variance and stabilized adversary in VaSSO. In particular, this subsection evaluates several flatness related metrics after training a ResNet-18 on CIFAR10 for 200 epochs, utilizing the same hyperparameters as those in Section 4.1.

Hessian spectrum. We first assess Hessian eigenvalues of a ResNet-18 trained with SAM and VaSSO. We focus on the largest eigenvalue λ 1 and the ratio of largest to the fifth largest eigenvalue λ 1 /λ 5 . These measurements are also adopted in (Foret et al., 2021; Jastrzebski et al., 2020) to reflect the flatness of the solution, where smaller numbers are more preferable. Because exact calculation for Hessian spectrum is too expensive provided the size of ResNet-18, we instead leverage Lanczos algorithm for approximation (Ghorbani et al., 2019). The results can be found in Table 5. It can be seen that SAM indeed converges to a much flatter solution compared with SGD, and VaSSO further improves upon SAM. This confirms that the friendly adversary issue is indeed alleviated by

the suppressed variance in VaSSO, which in turn boosts the generalizability of ResNet-18 as shown earlier in Section 4.1.

Label noise. It is known that SAM holds great potential to harness robustness to neural networks under the appearance of label noise in training data (Foret et al., 2021). As the training loss landscape is largely perturbed by the label noise, this is a setting where the suppressed variance and stabilized adversaries are expected to be advantageous. In our experiments, we measure the performance VaSSO in the scenarios where certain fraction of the training labels are randomly flipped. Considering θ = { 0 . 9 , 0 . 4 , 0 . 2 } , the corresponding test accuracies are summarized in Table 6.

Our first observation is that VaSSO outperforms SAM at different levels of label noise. VaSSO elevates higher generalization improvement as the ratio of label noise grows. In the case of 75% label noise, VaSSO with θ = 0 . 4 nontrivially outperforms SAM with an absolute improvement more than 5 , while VaSSO with θ = 0 . 2 markedly improves SAM by roughly 10 . In all scenarios, θ = 0 . 2 showcases the best performance and θ = 0 . 9 exhibits the worst generalization when comparing among VaSSO. In addition, when fixing the choice to θ , e.g., θ = 0 . 2 , it is found that VaSSO has larger absolute accuracy improvement over SAM under higher level of label noise. These observations coincide with Theorem 3, which predicts that VaSSO is suitable for settings with larger label noise due to enhanced stability especially when θ is chosen small (but not too small).

5 Other related works

This section discusses additional related work on generalizability of DNNs. The possibility of blending VaSSO with other approaches is also entailed to broaden the scope of this work.

Sharpness and generalization. Since the study of Keskar et al. (2016), the relation between sharpness and generalization has been intensively investigated. It is observed that sharpness is closely correlated with the ratio between learning rate and batchsize in SGD (Jastrz˛ ebski et al., 2017). Theoretical understandings on the generalization error using sharpness-related measures can be found in e.g., (Dziugaite and Roy, 2017; Neyshabur et al., 2017; Wang and Mao, 2022). These works justify the goal of seeking for a flatter valley to enhance generalizability. Targeting at a flatter minimum, approaches other than SAM are also developed. For example, Izmailov et al. (2018) proposes stochastic weight averaging for DNNs. Wu et al. (2020) studies a similar algorithm as SAM while putting more emphases on the robustness of adversarial training.

Other SAM type approaches. Besides the discussed ones such as GSAM and ASAM, (Zhao et al., 2022a) proposes a variant of SAM by penalizing the gradient norm based on the observation where sharper valley tends to have gradient with larger norm. Barrett and Dherin (2021) arrive at a similar conclusion by analyzing the gradient flow. Exploiting multiple (ascent) steps to find an adversary is systematically studied in (Kim et al., 2023). SAM has also been extended to tackle the challenges in

domain adaptation (Wang et al., 2023). However, these works overlook the friendly adversary issue, and the proposed VaSSO provides algorithmic possibilities for generalization benefits by stabilizing their adversaries. Since the desirable confluence with VaSSO can be intricate, we leave an in-depth investigation for future work.

Limitation of VaSSO and possible solutions. The drastically improved generalization of VaSSO comes at the cost of additional computation. Similar to SAM, VaSSO requires to backpropagate twice per iteration. Various works have tackled this issue and developed lightweight SAM. LookSAM computes the extra stochastic gradient once every a few iterations and reuses it in a fine-grained manner to approximate the additional gradient (Liu et al., 2022). ESAM obtains its adversary based on stochastic weight perturbation, and further saves computation by selecting a subset of the minibatch data for gradient computation (Du et al., 2022a). The computational burden of SAM can be compressed by switching between SAM and SGD following a predesigned schedule (Zhao et al., 2022b), or in an adaptive fashion (Jiang et al., 2023). SAF connects SAM with distillation for computational merits (Du et al., 2022b). It should be pointed out that most of these works follow the stochastic linearization of SAM, hence can also encounter the issue of friendly adversary. This opens the door of merging VaSSO with these approaches for generalization merits while respecting computational overhead simultaneously. This has been included in our research agenda.

6 Concluding remarks

This contribution demonstrates that stabilizing adversary through variance suppression consolidates the generalization merits of sharpness aware optimization. The proposed approach, VaSSO, provably facilitates stability over SAM. The theoretical merit of VaSSO reveals itself in numerical experiments, and catalyzes model-agnostic improvement over SAM among various vision and language tasks. Moreover, VaSSO nontrivially enhances model robustness against high levels of label noise. Our results corroborate VaSSO as a competitive alternative of SAM.

References

Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-aware minimization. In Proc. Int. Conf. Machine Learning , pages 639-668, 2022.

David GT Barrett and Benoit Dherin. Implicit gradient regularization. In Proc. Int. Conf. Learning Represention , 2021.

Peter L Bartlett, Philip M Long, and Olivier Bousquet. The dynamics of sharpness-aware minimization: Bouncing across ravines and drifting towards wide minima. arXiv:2210.01513 , 2022.

Léon Bottou, Frank E Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. arXiv preprint arXiv:1606.04838 , 2016.

MCettolo, J Niehues, S Stker, L Bentivogli, and M Federico. Report on the 11th iwslt evaluation campaign, iwslt 2014. 2014.

Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-SGD: Biasing gradient descent into wide valleys. In Proc. Int. Conf. Learning Represention , 2017.

Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong. When vision transformers outperform resnets without pre-training or strong data augmentations. In Proc. Int. Conf. Learning Represention , 2022.

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large-scale hierarchical image database. In Proc. Conf. Computer Vision and Pattern Recognition , pages 248-255, 2009.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv:1810.04805 , 2018.

Terrance Devries and Graham W. Taylor. Improved regularization of convolutional neural networks with cutout. abs/1708.04552, 2017.

Jiawei Du, Hanshu Yan, Jiashi Feng, Joey Tianyi Zhou, Liangli Zhen, Rick Siow Mong Goh, and Vincent Y. F. Tan. Efficient sharpness-aware minimization for improved training of neural networks. In Proc. Int. Conf. Learning Represention , 2022a.

Jiawei Du, Daquan Zhou, Jiashi Feng, Vincent Y. F. Tan, and Joey Tianyi Zhou. Sharpness-aware training for free. In Proc. Adv. Neural Info. Processing Systems , 2022b.

Gintare Karolina Dziugaite and Daniel M. Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In Proc. Conf. Uncerntainty in Artif. Intel. , 2017.

Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In Proc. Int. Conf. Learning Represention , 2021.

Saeed Ghadimi and Guanghui Lan. Stochastic first-and zeroth-order methods for nonconvex stochastic programming. SIAM Journal on Optimization , 23(4):2341-2368, 2013.

Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao. An investigation into neural net optimization via hessian eigenvalue density. In Proc. Int. Conf. Machine Learning , pages 2232-2241, 2019.

Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry P. Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. In Proc. Conf. Uncerntainty in Artif. Intel. , pages 876-885, 2018.

Stanisław Jastrz˛ ebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in sgd. arXiv:1711.04623 , 2017.

Stanislaw Jastrzebski, Maciej Szymczak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho, and Krzysztof J. Geras. The break-even point on optimization trajectories of deep neural networks. In Proc. Int. Conf. Learning Represention , 2020.

Weisen Jiang, Hansi Yang, Yu Zhang, and James Kwok. An adaptive policy to employ sharpnessaware minimization. arXiv:2304.14647 , 2023.

Yiding Jiang, Behnam Neyshabur, Dilip Krishnan, Hossein Mobahi, and Samy Bengio. Fantastic generalization measures and where to find them. arXiv:1912.02178 , 2019.

Rie Johnson and Tong Zhang. Accelerating stochastic gradient descent using predictive variance reduction. In Proc. Advances in Neural Info. Process. Syst. , pages 315-323, Lake Tahoe, Nevada, 2013.

Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In Proc. Int. Conf. Learning Represention , 2016.

Hoki Kim, Jinseong Park, Yujin Choi, Woojin Lee, and Jaewook Lee. Exploring the effect of multi-step ascent in sharpness-aware minimization. arXiv:2302.10181 , 2023.

Minyoung Kim, Da Li, Shell Xu Hu, and Timothy M. Hospedales. Fisher SAM: Information geometry and sharpness aware minimisation. In Proc. Int. Conf. Machine Learning , pages 11148-11161, 2022.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Proc. Int. Conf. Learning Represention , 2014.

Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. ASAM: Adaptive sharpnessaware minimization for scale-invariant learning of deep neural networks. In Proc. Int. Conf. Machine Learning , volume 139, pages 5905-5914, 2021.

Bingcong Li, Meng Ma, and Georgios B Giannakis. On the convergence of SARAH and beyond. arXiv preprint arXiv:1906.02351 , 2019.

Bingcong Li, Alireza Sadeghi, and Georgios Giannakis. Heavy ball momentum for conditional gradient. In Proc. Advances in Neural Info. Process. Syst. , 2021.

Yong Liu, Siqi Mai, Xiangning Chen, Cho-Jui Hsieh, and Yang You. Towards efficient and scalable sharpness-aware minimization. In Proc. Conf. Computer Vision and Pattern Recognition , volume 2022, pages 12350-12360, 2022.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In Proc. Int. Conf. Learning Represention , 2017.

V. Mardia and P. E. Jupp. Directional Statistics . Directional statistics, 2000.

Peng Mi, Li Shen, Tianhe Ren, Yiyi Zhou, Xiaoshuai Sun, Rongrong Ji, and Dacheng Tao. Make sharpness-aware minimization stronger: A sparsified perturbation approach. In Proc. Adv. Neural Info. Processing Systems , 2022.

Aryan Mokhtari, Hamed Hassani, and Amin Karbasi. Stochastic conditional gradient methods: From convex minimization to submodular maximization. Journal of Machine Learning Research , 21(1): 4232-4280, 2020.

Behnam Neyshabur, Srinadh Bhojanapalli, David Mcallester, Nathan Srebro, and Nati Srebro. Exploring generalization in deep learning. In Proc. Adv. Neural Info. Processing Systems , volume 30, pages 5947-5956, 2017.

Lam M Nguyen, Jie Liu, Katya Scheinberg, and Martin Takáˇ c. SARAH: A novel method for machine learning problems using stochastic recursive gradient. In Proc. Intl. Conf. Machine Learning , Sydney, Australia, 2017.

Sashank J Reddi, Suvrit Sra, Barnabás Póczos, and Alex Smola. Stochastic Frank-Wolfe methods for nonconvex optimization. In Allerton conference on communication, control, and computing , pages 1244-1251. IEEE, 2016.

Nitish Srivastava, Geoffrey E. Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. J. Mach. Learn. Res. , 15: 1929-1958, 2014.

Brown Tom, Mann Benjamin, Ryder Nick, Subbiah Melanie, Kaplan Jared, Dhariwal Prafulla, Neelakantan Arvind, Shyam Pranav, Sastry Girish, Askell Amanda, Agarwal Sandhini, HerbertVoss Ariel, Krueger Gretchen, Henighan Tom, Child Rewon, Ramesh Aditya, Ziegler Daniel M., Wu Jeffrey, Winter Clemens, Hesse Christopher, Chen Mark, Sigler Eric, Litwin Mateusz, Gray Scott, Chess Benjamin, Clark Jack, Berner Christopher, McCandlish Sam, Radford Alec, Sutskever Ilya, and Amodei Dario. Language models are few-shot learners. In Proc. Adv. Neural Info. Processing Systems , volume 33, pages 1877-1901, 2020.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proc. Adv. Neural Info. Processing Systems , volume 30, 2017.

Pengfei Wang, Zhaoxiang Zhang, Zhen Lei, and Lei Zhang. Sharpness-aware gradient matching for domain generalization. In Proc. Conf. Computer Vision and Pattern Recognition , pages 3769-3778, 2023.

Ziqiao Wang and Yongyi Mao. On the generalization of models trained with sgd: Informationtheoretic bounds and implications. In Proc. Int. Conf. Learning Represention , 2022.

Kaiyue Wen, Tengyu Ma, and Z hiyuan Li. How does sharpness-aware minimization minimizes sharpness. In Proc. Int. Conf. Learning Represention , 2023.

Ashia C. Wilson, Rebecca Roelofs, Mitchell Stern, Nathan Srebro, Benjamin Recht, and Nati Srebro. The marginal value of adaptive gradient methods in machine learning. In Proc. Int. Conf. Machine Learning , volume 30, pages 4148-4158, 2017.

Dongxian Wu, Shu-Tao Xia, and Yisen Wang. Adversarial weight perturbation helps robust generalization. In Proc. Adv. Neural Info. Processing Systems , volume 33, pages 2958-2969, 2020.

Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM , 64(3):107-115, 2021.

Zhiyuan Zhang, Ruixuan Luo, Qi Su, and Xu Sun. GA-SAM: Gradient-strength based adaptive sharpness-aware minimization for improved generalization. In Proc. Conf. Empirical Methods in Natural Language Processing , 2022.

Yang Zhao, Hao Zhang, and Xiuyuan Hu. Penalizing gradient norm for efficiently improving generalization in deep learning. In Proc. Int. Conf. Machine Learning , pages 26982-26992, 2022a.

Yang Zhao, Hao Zhang, and Xiuyuan Hu. SS-SAM: Stochastic scheduled sharpness-aware minimization for efficiently training deep neural networks. arXiv:2203.09962 , 2022b.

Juntang Zhuang, Boqing Gong, Liangzhe Yuan, Yin Cui, Hartwig Adam, Nicha Dvornek, Sekhar Tatikonda, James Duncan, and Ting Liu. Surrogate gap minimization improves sharpness-aware training. In Proc. Int. Conf. Learning Represention , 2022.

Supplementary Document for

'Enhancing Sharpness-Aware Optimization Through Variance Suppression'

A Linking SAM adversary with stochastic Frank Wolfe

A.1 Stochastic Frank Wolfe (SFW)

We first briefly review SFW. Consider the following nonconvex stochastic optimization

where X is a convex and compact constraint set. SFW for solving (6) is summarized below.

Algorithm 2 SFW (Reddi et al., 2016)

1: Initialize: x 0 ∈ X

2: for t = 0 , 1 , . . . , T -1 do

3: draw iid samples { ξ b t } B t b =1

4: let ˆ g t = 1 B t ∑ B t b =1 ∇ h ( x t , ξ b t )

5: v t +1 = arg max v ∈X 〈 ˆ g t , v 〉

6: x t +1 = (1 -γ t ) x t + γ t v t

7: end for

It has been shown in (Reddi et al., 2016, Theorem 2) that one has to use a sufficient large batch size B t = O ( T ) , ∀ t to ensure convergence of SFW. This is because line 5 in Alg. 2 is extremely sensitive to gradient noise.

A.2 The adversary of SAM

By choosing h ( glyph[epsilon1] ) = f ( x t + glyph[epsilon1] ) and X = S ρ ( 0 ) , it is not hard to observe that 1-iteration SFW with γ 0 = 1 gives equivalent solution to the stochastic linearization in SAM; cf. (2) and (3). This link suggests that the SAM adversary also suffers from stability issues in the same way as SFW. Moreover, what amplifies this issue in SAM is the adoption of a constant batch size, which is typically small and far less than the O ( T ) requirement for SFW.

Our solution VaSSO takes inspiration from modified SFW approaches which leverage a constant batch size to ensure convergence; see e.g., (Mokhtari et al., 2020; Li et al., 2021). Even though, coping with SAM's instability is still challenging with two major obstacles. First, SAM uses one-step SFW, which internally breaks nice analytical structures. Moreover, the inner maximization (i.e., the objective function to the SFW) varies every iteration along with the updated x t .

A.3 The three dimensional example in Fig. 2

Detailed implementation for Fig. 2 is listed below. We use ∇ f ( x ) = [0 . 2 , -0 . 1 , 0 . 6] . The stochastic noise is ξ = [ ξ 1 , ξ 2 , ξ 3 ] , where ξ 1 , ξ 2 , ξ 3 are iid Gaussian random variables with variance scaling with 0 . 2 , 1 , 2 , respectively. We scale the variance to change the SNR. We generate 100 adversaries by solving arg max ‖ glyph[epsilon1] ‖≤ ρ 〈∇ f ( x ) + ξ , glyph[epsilon1] 〉 for each choice of SNR. As shown in Fig. 2, the adversaries are unlikely to capture the sharpness information when the SNR is small, because they spread indistinguishably over the sphere.

B More on m -sharpness

m -sharpness can be ill-posed. Our reason for not studying m -sharpness directly is that its formulation (Andriushchenko and Flammarion, 2022, eq. (3)) may be ill-posed mathematically due to the lack of a clear definition on how the dataset S is partitioned. Consider the following example, where the same notation as (Andriushchenko and Flammarion, 2022) is adopted for convenience. Suppose that the loss function is l i ( w ) = a i w 2 + b i w , where ( a i , b i ) are data points and w is the

+1

parameter to be optimized. Let the dataset have 4 samples, ( a 1 = 0 , b 1 = 1) ; ( a 2 = 0 , b 2 = -1) ; ( a 3 = -1 , b 3 = 0) ; and, ( a 4 = 1 , b 4 = 0) . Consider 2-sharpness.

If the data partition is {1,2} and {3,4}, the objective of 2-sharpness i.e., equation (3) in (Andriushchenko and Flammarion, 2022), becomes min w ∑ 2 i =1 max || δ || <ρ 0 .

If the data partition is {1,3} and {2,4}, the objective is min w ∑ 2 i =1 max || δ || <ρ f i ( w,δ ) , where f 1 is the loss on partition {1,3}, i.e., f 1 ( w,δ ) = -( w + δ ) 2 + ( w + δ ) ; and f 2 ( w,δ ) = ( w + δ ) 2 -( w + δ ) is the loss on partition {3,4}.

Clearly, the objective functions are different when the data partition varies. This makes the problem ill-posed - different manners of data partition lead to entirely different loss curvature. In practice, the data partition even vary with a frequency of an epoch due to the random shuffle.

C Details on numerical results

CIFAR10 and CIFAR100. For these small resolution datasets, we slightly change the first convolution layer of ResNet18 and WRN-28-10 to one with 3 × 3 kernel size, 1 stride and 1 padding following Mi et al. (2022). The results on SGD and SAM demonstrate that the accuracy is almost identical to the vanilla model.

ResNet50 on ImageNet. Due to the constraints on computational resources, we report the averaged results over 2 independent runs. For this dataset, we randomly resize and crop all images to a resolution of 224 × 224 , and apply random horizontal flip, normalization during training. The batch size is chosen as 128 with a cosine learning rate scheduling with an initial step size 0 . 05 . The momentum and weight decay of base optimizer, SGD, are set as 0 . 9 and 10 -4 , respectively. We further tune ρ from { 0 . 05 , 0 . 075 , 0 . 1 , 0 . 2 } , and chooses ρ = 0 . 075 for SAM. VaSSO uses θ = 0 . 99 . VaSSO and ASAM adopt the same ρ = 0 . 075 .

ViT-S/32 on ImageNet. We follow the implementation of (Du et al., 2022b), where we train the model for 300 epochs with a batch size of 4096 . The baseline optimizer is chosen as AdamW with weight decay 0 . 3 . SAM relies on ρ = 0 . 05 . For the implementation of GSAM and V+G, we adopt the same implementation from (Zhuang et al., 2022).

D Missing proofs

Alg. 1 can be written as

where ‖ glyph[epsilon1] t ‖ = ρ . In SAM, we have glyph[epsilon1] t = ρ g t ( x t ) ‖ g t ( x t ) ‖ , and in VaSSO we have glyph[epsilon1] t = ρ d t ‖ d t ‖ .

D.1 Useful lemmas

This subsection presents useful lemmas to support our main results.

Lemma 1. Alg. 1 (or equivalently iteration (7) ) ensures that

Proof. To start with, we have that

Taking expectation conditioned on x t , we arrive at

where (a) is because of Assumption 2; and (b) is because x t -x t + 1 2 = -glyph[epsilon1] t and its norm is ρ . This inequality ensures that

where the last inequality is because ρη t ‖∇ f ( x t ) ‖ ≤ 1 2 η 2 t ‖∇ f ( x t ) ‖ 2 + 1 2 ρ 2 . Taking expectation w.r.t. x t finishes the proof.

Lemma 2. Alg. 1 (or equivalently iteration (7) ) ensures that

Proof. The proof starts with bounding ‖ g t ( x t + 1 2 ) ‖ as

where (a) is the result of Assumption 2; and (b) is because x t -x t + 1 2 = -glyph[epsilon1] t and its norm is ρ . Taking expectation conditioned on x t , we have

where the last inequality is because of Assumption 3. Taking expectation w.r.t. the randomness in x t finishes the proof.

Lemma 3. Let A t +1 = αA t + β with some α ∈ (0 , 1) , then we have

Proof. The proof can be completed by simply unrolling A t +1 and using the fact 1+ α + α 2 + . . . + α t ≤ 1 1 -α .

D.2 Proof of Theorem 1

Proof. Using Assumption 2, we have that

Taking expectation, then plugging Lemmas 1 and 2 in, we have

As the parameter selection ensures that η t ≡ η = η 0 √ T ≤ 2 3 L , it is possible to divide both sides with η and rearrange the terms to arrive at

Summing over

where (a) uses Assumption 1, and the last equation is obtained by plugging in the value of ρ and η . This completes the proof to the first part.

For the second part of this theorem, we have that

Averaging over t completes the proof.

D.3 Proof of Theorem 2

Proof. To bound the MSE, we first have that

Now we cope with three terms in the right hind of (9) separately.

The second term can be bounded directly using Assumption 2

For the third term, we have

The first term is bounded through

where (a) is because of Young's inequality. Taking expectation and applying Lemma 2, we have that

[ 2 ] (12)

The last inequality uses the value of η = η 0 √ T and ρ = ρ 0 √ T . In particular, we have η 2 ρ 2 L 4 = O (1 /T 2 ) and η 2 L 2 σ 2 = O ( σ 2 /T ) , and

where the last equation is the result of Theorem 1.

Combining (9) with (12), (10) and (11), and choosing λ = θ 1 -θ , we have

where the last inequality is the result of Lemma 3.

D.4 Proof of Theorem 3

Proof. We adopt a unified notation for simplicity. Let v t := d t for VaSSO, and v t := g t ( x t ) for SAM. Then for both VaSSO and SAM, we have that

For convenience, let glyph[epsilon1] ∗ t = ρ ∇ f ( x t ) / ‖∇ f ( x t ) ‖ . From (13), we have that

Applying triangular inequality ‖ a ‖ - ‖ b ‖ ≤ ‖ a -b ‖ , we arrive at

Combining (14) with (15), we have

which further implies that

The last inequality is because ( E [ a ]) 2 ≤ E [ a 2 ] . This theorem can be proved by applying Assumption 3 for SAM and Lemma 2 for VaSSO.

Peer reviews:
comments: SUMMARY:
This paper proposes a method, which applies EMA on batch gradients, to better solve the maximum problem in SAM and obtain higher test accuracy.

STRENGTHS:
1.The proposed method is easy to implement.

2.The theoretical analysis is sufficient.

WEAKNESSES:
1.The core of this paper is to alleviate the gradient noise when solving the inner maximum problem of SAM. To achieve this goal, it proposes to implement EMA on the batch gradient and gives theoretical proof on the variance suppression. Although experiment results show the improvement of VaSSO, I am somehow confused that why the gradient noise harms the generalization performance of SAM. In the Section 4.1 of SAM[1], it claims that smaller batch size tends to yield models having better generalization ability, which conflicts the core idea of this paper.

2.Following 1, the statements on Line 47-48 and Line 50 need more support or evidence. And I think the introduction takes too much space to describe related works. In contrast, the author should give some experiment results to support the statement on Line 47-50. 

[1] Foret P, Kleiner A, Mobahi H, et al. Sharpness-aware minimization for efficiently improving generalization[J]. arXiv preprint arXiv:2010.01412, 2020.

QUESTIONS:
no question

LIMITATIONS:
no limitation
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-06-27

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes to improve upon sharpness-aware minimization (SAM) with variance-reduced inner perturbation steps. Theoretical analysis demonstrates that the proposed method achieves similar convergence rate with the original SAM. Empirical results on different applications (image classification and neural machine translation) demonstrate the effectiveness of proposed method. In addition, the proposed method also endows SAM with robustness against large label noise.\n\nSTRENGTHS:\n- The proposed method is clear and easy to understand\n- Theoretical results are sound\n- Empirical results demonstrate the superiority of proposed method\n\nWEAKNESSES:\n- Some claims are not fully supported\n- Some baseline methods are missing\n\nQUESTIONS:\n- I am a bit puzzled by the claim that the proposed method can be easily integrated with computational efficient variants of SAM, e.g., (Liu et al, 2022; Zhao et al., 2022b) as well as a missing reference (Jiang et al., 2023). Compared with SAM that only needs to compute stochastic gradient, VaSSO has to keep track of d_t, as in (4a). I am not sure if the update is still possible if we only perform SAM step periodically (Liu et al, 2022) or randomly (Zhao et al., 2022b). The authors need to elaborate more on that. \n- If such integration is not possible, the proposed method may suffer from larger computational cost than these variants. Although it can be regarded as the cost of better final performance, the authors may still need to make this limitation clear. \n- Despite these computation efficient versions, the authors claim that some other existing works on SAM are also orthogonal to their work and can be easily integrated. While I do not find critical problems for these works, the authors may need to add some experiments that integrate the proposed method with these methods, and see if such integration can achieve any improvements. Nevertheless, such experiments are missing in current version. \n- It also surprises me that the experiments contain so few baseline methods. For example, why is Fisher SAM (Kim et al., 2022) not compared in experiments? It is confusing to compare with ASAM and not Fisher SAM. \n\nReferences:\n\nWeisen Jiang, Hansi Yang, Yu Zhang, James Kwok. An adaptive policy to employ sharpness-aware minimization. ICLR 2023\n\nLIMITATIONS:\nPlease see Weakness and Questions part.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a variance suppression approach for SAM in order to account for the sensitivity of stochastic gradients used in SAM inner maximization. The proposed method is shown to provably reduce the MSE of gradient estimation. Some experimental results on benchmark datasets are provided.\n\nSTRENGTHS:\nPaper strengths:\n1. great motivation for variance suppression\n2. simple algorithm modification to standard SAM\n3. theoretical results for VaSSO in Theorem 2 and Corollary 1\n\nWEAKNESSES:\nPaper weaknesses:\n1. The authors claim in section 3.2 that VaSSO can boost performance of other SAM family approaches, but this is not shown in the experimental results\n2. gains are marginal for larger datasets, see Table 3\n3. authors provide convergence rates for Vasso in Corollary 1 but it's not clear that these are sharper than the SAM rates in Theorem 1\n\nQUESTIONS:\nQuestions to address in rebuttal:\n1. does VaSSO translate to gains in domain generalization performance? e.g. WILDS benchmark [1]\n2. can authors cite more Frank-Wolfe papers in connection to min-max problems, e.g. [2], [3]\n3. how do convergence rates in Corollary 1 relate to Frank-Wolfe optimization convergence rates? convergence rates for FW algorithms and various variants have been studied in the literature [4]\n4. what are the limitations of VaSSO? authors should include a discussion of limitations\n5. can the authors add more SAM baselines to the list? Also, ASAM combined with VaSSO is not included, and would be good to include for comparisons\n\nI am willing to increase my score if the authors address most of these concerns.\n\n[1] Koh et al, WILDS: A Benchmark of in-the-Wild Distribution Shifts, https://arxiv.org/abs/2012.07421\n\n[2] Tsiligkaridis et al, Understanding and Increasing Efficiency of Frank-Wolfe Adversarial Training, CVPR 2022, https://arxiv.org/abs/2012.12368\n\n[3] Gidel et al, Frank-Wolfe Algorithms for Saddle Point Problems, https://arxiv.org/pdf/1610.07797.pdf\n\n[4] Huang et al, Accelerated Stochastic Gradient-free and Projection-free Methods, ICML 2020\n\nLIMITATIONS:\nNot discussed.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n**I have read the author's rebuttal, see reply below** This work proposes changing SAM from perturbing the weights by the gradient of the current minibatch, with instead a moving average of the gradients across training iterations. They provide upper bounds that indicate that the moving average may better approximates the inner-maximization with respect to the loss over the entire dataset, which is more difficult for naive SAM due to the minibatch noise. Empirically, SAM with moving average gradients (which they call VaSSO) performs better on CIFAR10, CIFAR100 by around $0.1-0.5$%, and the improvements are more apparent in the presence of heavy label noise with gains up to +10% improvement with 75% label noise. Empirical experiments show that VaSSO minimizes the maximum eigenvalue of the Hessian better than SAM.\n\nSTRENGTHS:\n**Significance**: The proposed algorithm VaSSO is a very simple change that achieves flatter minima and achieves better generalization than SAM. Although it may not necessarily be the case that two effects are tied in a causal manner, its effectiveness at achieving both beyond SAM may suggest that it may be a useful optimizer in practice and for further study in future works. \n\n**Quality**: The paper is clearly written and organized.\n\nWEAKNESSES:\n**1. Lack of important literature review**: The authors motivate VaSSO by identifying problems with SAM in terms of its ability to optimize the original objective _due to the minibatch noise_. The moving average is hypothesized to be better at correctly approximating the full-batch gradient. The authors do not address the fact that works that have analyzed SAM previously have unanimously observed that SAM actually **only observes improvements in generalization if it is paired with minibatch noise**. In particular, Andruschenko et al (https://arxiv.org/abs/2206.06232) showed that n-SAM which utilizes the full-batch gradient directly for the perturbation step actually observed no generalization gain. This is supported by experiments on m-sharpness in the original SAM paper (https://arxiv.org/abs/2010.01412), and they also showed that better approximating the sharpest direction by also taking the second order approximation of the loss instead of the first order approximation lead to worse performance. Wen et al (https://arxiv.org/abs/2211.05729) showed that the minibatch noise is important to minimize the trace of the Hessian instead of the max eigenvalue (though the work does not make any claim about which measure of sharpness is better correlated with generalization). \n\nTo summarize, the authors pose the problem as that SAM cannot achieve the full generalization benefits because it does not minimize its intended objective optimally, but previous literature indicate that this suboptimality is actually what allows SAM to achieve better generalization. The authors do not mention the conclusions made in previous works, but it seems important to address this conflict. \n\n**2. Marginal improvements, and inconsistent numbers in comparison to Foret et al.**: Without label noise VaSSO's improvements range between 0.2-0.5% for CIFAR10 and CIFAR100, and it's not clear whether further hyperparameter tuning could close this gap. In particular, the authors train WideResNet-28-10 on CIFAR10 with cutout data augmentation which was also conducted by Foret et al, where they report an error of 2.3% while the authors report 2.7% and for VaSSO 2.5%. This slight improvement may potentially come from the authors only optimizing the rho hyperparameter, and not m-sharpness (the other hyperparameter mentioned in Foret et al. where the minibatch is sharded first). \n\nFor label noise, there is a much more significant boost in test accuracy, but it is unclear whether the authors are reporting the peak early-stopping accuracy or the final accuracy. The former is more important for label noise, and the difference between VaSSO and SAM may be coming from reporting the latter. Also, Foret et al reported better numbers for SAM on CIFAR10 with similar amounts of label noise although they were using ResNet34 instead of ResNet18. \n\n**3. The upper bounds are too loose, and comparing upper bounds does not lead to any meaningful comparison.** The upper bound derived in Theorem 2 requires very approximate intermediate steps (very loose upper bound in Line 181), and is not a function of $t$ but $T$. The bound implies that if you train for longer, the bound for the MSE between the moving average and the full batch gradient _at every intermediate step_ improves which doesn't seem right. \n\nMore importantly, any meaningful comparison between the MSE for SAM and MSE for VaSSO should compare the upper bound of VaSSO to the _lower bound_ for SAM. Comparing two upper bounds isn't very meaningful.\n\nQUESTIONS:\nSee weaknesses\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-11"}