Enhancing Sharpness-Aware Optimization Through Variance Suppression
Authors: Bingcong Li
Georgios B. Giannakis
Venue: NeurIPS 2023 poster (2023)
Keywords: generalization
optimization
neural networks
TL;DR: We propose variance suppression (VaSSO), a new sharpness-aware minimization approach, for improving the generalizability of 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
Bingcong Li
Georgios B. Giannakis
University of Minnesota -Twin Cities Minneapolis, MN, USA (1ixx5599, georgi os ) @umn.edde
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 min- ima" heighten generalization ability, SAM seeks 'flat valleys' by minimizing the maximum loss caused by an adversary perturbing parameters within the neigh- borhood. 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 safe- guards 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 ht https: /ittbbb.. com/ /Biggggggggg////....
Introduction
Despite deep neural networks (DNNs) have advanced the concept of "Iearning from data," and markedly improved performance across several applications vision and language (Devlin 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 generalization, which 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 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 suitable optimization schemes attempting 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 other intricate learning tasks calls for additional theoretical elaboration.
Our main focus, sharpness tion approach that facilitates sfate-of-the-arr generalizability by exploiting sharpness of loss landscape (Foret et al., 2021; Chen et al., 2022). high-level interpretation of sharpness is how violently the loss fluctuates within a neighborhood. It has been shown through large-scale empirical studies that Charpness-based measures highly correlate with generalization (Jiang al., 2019). Several works
minimization/opiimizztii (SAM), is highly compelling optimiza-
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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 neighboring area. Thanks to such 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 undermined by what we termed friendly adversary. Confined by the stochastic linearization for computational efficiency, SAM's adversary only captures the sharpness for particular minibatch of data, and can become a friend on other data samples. Because the global sharpness is not approached accurately, the friendly ad adversary precludes SAM from attaining its utmost generalizability. The present work advocates variance suppressed sharpness aware optimization (VaSSO') 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.
We find that the friendly adversary discourages generalizability of SAM. This challenge is catastrophic in our experiments it completely wipe out the generalization merits.
A novel approach, VaSSO, is proposed to tackle this issue. VaSSO is equipped with what we termed variance suppression to streamline principled means for stabilizing adversaries. The theoretically guaranteed stability promotes refined global sharpness estimates, thereby alleviating the issue of friendly adversary.
A side 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.
Numerical experiments confirm the merits of stabilized adversary in VaSSO. It is demon- strated 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.
Bold lowercase (capital) letters denote column vectors (matrices); ||x|| stands for &o norm of vector x; and (x, y) is the inner product of x and y. (x) denotes the surface of ball with radius p centered at x, i.e., (x) == {x + pu |u| = 1}.
Notation.
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
min max f(x + lellSp
(1)
where p is the radius of considered neighborhood, and the nonconvex objective is defined f(x) :=
Eg [ee[f((()]] Here, x is the neural network parameter, and B is a random batch of data. The merits of such formulation resides in its implicit sharpness measure max]elll f(x E) fxx), which effectively drives the optimization trajectory towards the desirable flat valley (Kim al., 2022).
Vasso coincides with the Greek nickname for Vasiliki.

---

Algorithm Generic form of SAM
1: Initialize: X0. p
2: for 0, ......, do
3: Sample minibatch B,, and define stochastic gradient B3 (x())
4: Find E E (0) via stochastic linearization; e.g., (4) for VaSSO or (3) for SAM
5: Calculate stochastic gradient (x, + Ex)
6: Update model via X = X 1/B Xt + Et)
7:
end for
8:
Return: XT
The inner maximization of (1) has natural interpretation as finding an adversary. Critical as it is, obtaining an adversary calls for stochastic linearization to alleviate computational concerns, i.e.,
c = arg max f(x + c) (a)
arg max f(x+) (VI())), arg max f(x,) + (g(())), (2)
|
||ellSP
|ellsp
where linearization (a) relies on the first order Taylor expansion of f(x+ + c). This is typically accurate given the choice of small p. A stochastic gradient g: (x;) then substitutes V)((x) in (b) to
downgrade the computational burden of full gradient. Catalyzed by the stochastic linearization in (2), is possible to calculate SAM's adversary closed-form
ge (x))
SAM: P
gr (x,)
SAM then adopts the stochastic gradient adversary g (x, + Et) to update x in a SGD fashion. A step-by-step implementation is summarized in Alg. I, where the means to find adversary in line is presented in generic form in order to unify the algorithmic framework with later sections.
2.2 The good
To provide comprehensive understanding about SAM, this subsection focuses on Alg. 1, and establishes its convergence for (1). Some necessary assumptions are listed below, 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*, Vx
Assumption 2 (smoothness). The stochastic gradient g(x) is I--Lipschitz iee. llg(x) - g(y)l S
L|x -|l, yx,y.
Assumption 3 (bounded variance). The stochastic gradient g(x) is unbiased with bounded variance, that [Eg(()) xx] fx( and [x]lg(x) - V [xxxx) 0² for some a
The constraint of (1) is never violated since ||e,| = p holds for each I; see line in Alg. 1. Hence, the convergence SAM pertains the behavior of objective, where tight result is given below. Theorem 1 (SAM convergence). Suppose that Assumptions 1 - 3 hold. Let 2 (clearly 0 < co < 1), Alg. I guarantees that
= 2h < 2
vT
37
and p
Po VT.
Then with =
3Lm
T-1 E 11
0² VT
T-1
o²
50
50
and
[[]V + SO
T
T 0
0=
The convergence rate of SAM is the same as SGD up to constant factors, where the detailed expression hidden under big O notation 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 = (1/) relative to (Andriushchenko and Flammarion, 2022), where the latter only supports (T 1/4).
be found in Appendix D. Our results eliminate the need for the
A message from Theorem is that any adversary satisfying E (0) ensures converge. Because the surface (0) a gigantic space, it challenges the plausible optimality of the adversary and poses natural question is possible to find more powerful adversary for generalization advances?
3

---

1.04
1.04
35.0
SGD
SGD
. SAM
1.03
SAM
1.08
SAM
SAMvvsr 10m-
34.9
022
34.7
3.0
L9950 996
1.0
1.0
34.5 m-/p3
Rechwe18
ReshheiM
Rm
Re-Nei34
m-
m-bi
(a)
(b)
(c)
Figure I: (a) A friendly adversary erases the generalization merits of SAM; (b) m-sharppess may not directly correlate with variance since noisy gradient degrades generalization; and (c) m-sharpness may not hold universally. Note that test accuracies in (a) and (b) are normalized to SGD.
2.3 The challenge: friendly adversary
Adversary one minibatch is a friend of others. SAM's adversary is 'malicious' for minibatch B, but not necessarily for other data because it only safeguards f8, (x, + +2) - (xx) > 0 for small In fact, it can be shown that fa (X; Et) f& (xx) (x, S 0 whenever the stochastic gradients do not align well, i.e., (gr (x,), gs(x:)) I 0. Note that such misalignment is common because of the variance massive training datasets. This issue is referred to as friendly adversary, and implies that the adversary 6, cannot accurately depict the global sharpness of x,. Note that the 'friendly adversary' also has more involved interpretation, that is, g: (x,) 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 é; = .g: (x, / gr (xx)],, where gr denotes the gradient on B,, randomly sampled batch of the same and its performance for i) ResNet-18 on CIFARIO and ii) ResNet-34 on CIFAR 100 can be found in Fig. (a)) Note that the test accuracy is normalized relative SGD for the ease of visualization. evident that the friendly é in SAM-db almost erases the generalization benefits entirely.
as B. This modified approach is denoted as SAM-db,
Source of friendly adversary. The major cause the friendly adversary attributes to the gradient variance, which equivalently translates to the lack of stability in SAM's ss stochastic linearization (2b). An illustrative three dimensional example is shown in Fig. 2, where we plot the adversary e, obtained from different g, realization in (2b). The minibatch gradient is simulated by adding Gaussian noise to the true gradient. When the signal noise ration (SNR) similar a practical scenario (ResNet-18 on CIFAR 10 shown in Fig. (c)), 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 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 e; is found using subsampled B, of size (i.e., larger variance).
Since m-sharpness highly hinges upon the loss curvature, is unlikely to hold universally. For example, transformer is trained on IWSLT-14 dataset, where the test performance (BI.U) decreases with smaller m even we have tuned P carefully; see Fig. ((c). On the theoretical side, an example is provided in (Andriushchenko and Flammarion, 2022, Sec. 3) suggest that m-sharpness is not

---

(a) SNR = 5
(b) SNR =1
(c) SNR = 0.1 (d) SNR = 0.01 (e) SNR in practice
Figure 2: (a) (d) SAM's adversaries spread over the surface; (e) SNR is in [0.01, 0.1] when training
a ResNet-18 on CIFARIO, where the SNR is calculated at the first iteration of every epoch.
necessarily related with sharpness or generalization. Moreover, there also exists specific choice for m such that the m-sharpness formulation is ill- ill-posed. We will expand on this Appendix B.
Even in the regime where m-sharpness is empirically observed such as ResNet- 18 on CIFARIO and ResNet-34 on CIFAR 100, we show through experiments that m-sharpness is not consequence of gradient variance, thus not contradicting with the friendly adversary issue tackled in this work.
Observation 1. Same variance, different generalization. Let
128 and batchsize b = 128.
Recall the SAM-db experiment in Fig. (a). If m-sharpness is a direct result of gradient variance, is logical to expect SAM-db has comparable performance SAM simply because their batchzises (hence variance) for finding adversary are the same. Unfortunately, SAM-db degrades accuracy. We further increase the variance of gi (xx) by setting
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--sharpeess.
Observation 2. Enlarged variance degrades generalization. We explicitly increase variance when
finding adversary by adding Gaussian noise Ç to g: (x.), i.e., & = P [[g. (x, )+(++ 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- and ResNet-34. This again illustrates that the plausible benefit of m-sharpness does not stem from increased variance.
5: (xx))x(
In sum, observations and 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-supreesed sharpness-aware optimization (VaSSO)
This section advocates variance suppression 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; 7; Li et al., 2019). However, the requirement compute a full gradient every few iterations is infeasible and hardly scales well for tasks such as training DNNs.
The proposed variance suppression (VSSSO) overcomes this computational burden through novel yet simple stochastic lincarization. For a prescribed e (0, 1). VaSSO is summarized below
VaSSO: d. (1 0)d, (x:)
(4a)
de [d
arg max f(x)) + (d,, ||e||Sp
(4b)
Compared with (2) of SAM. the key difference is that VaSSO relies on slope d, for more stable stochastic linearization as shown in (4b). The slope d, is an exponentially moving average (EMA) of (g:()))): J such that the change over consecutive iterations is smoothed. Noticing that e, and d, share the same direction, the relatively smoothed {d,}, thus imply the stability of {e.}, 2 in VaSSO. Moreover, as d, processes information of different minibatch data, the global sharpness can be captured in principled manner to alleviate the friendly adversary challenge.
5

---

To theoretically characterize the effectiveness of VaSSO, our first result considers d, qualified strategy estimate V fxx). and delves into its mean square error (MSE).
Theorem 2 (Variance suppression), Suppose that Assumptions 1 - hold. Let Alg. 1 equip with i) e obtained by (4) with 0 E (0, 1); and, ii) ni and selected the same as Theorem 1. VaSSO guarantees that the MSE of d, is bounded by
0)00 -1-0)000 =0000 OT
E - s 5 500
(5)
Because SAM's gradient estimate has a looser bound on MSE (or variance), that is, Blllg v((x)]] S o², 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; with slope v as (v( == maxijelis> fxx) (xx) + (v, For example, (((d)) and Li (x.)) are quality of VaSSO and SAM, respectively. Another critical case of concern Le (xxx)) is shown in (Zhuang et al., 2022) that L: (xx))) & max]]el|<s [xx + 6) given a small p. Moreover, L. (V - f(xx) is also an accurate approximation to the sharpness (Zhuang et al., 2022). These observations safeguard (xxx)) as the anchor when analyzing the stability SAM and VaSSO.
Definition (ô-stability). A stochastic linearization with slope is said to be ô-stable if its quality satisfies E (v) - (.V((())) S 6.
larger implies a more friendly adversary, hence is less preferable. We are now well-prepared for our main results on adversary's 's stability.
Theorem 3 (Adversaries of VaSSO is more stable than SAM.). Suppose that Assumptions 3 hold. Under the sume hyperparameter choices Theorem 2, the stochastic linearization is [VOpo + O 8T174) -stabble for VaSSO, while po-stable in SAM.
Theorem demonstrates that VaSSO alleviates the friendly adversary problem by promoting stability. Qualitatively, VaSSO is roughly v0 E (0, 1) times more stable relative to SAM, since the term in big O notation is negligible given sufficiently large T. Theorem 3 also guides the choice of 0 preferably small but not too small, otherwise the term 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 direct consequence of Theorem 1. The reason is that e, E (0) is satisfied by (4).
Corollary (VaSSO convergence). Suppose that Assumptions I - 3 hold. Choosing Th and the same as Theorem 1, then for any 0 E (0, 1). VaSSO ensures that
and T Σ E[IIVf( so 0²
T [[IV)(x))1]] < o
VaSSO better reflects sharpness around optimum. Consider a near optimal region where [Vx()l) + O. Suppose that we are in big data regime where (x.) = Vx(x,) + Ç for some Gaussian random variable Ç. The covariance matrix of Ç is assumed to be o 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 E, oc/||c|| uniformly distributes over (0) regardless of whether the neighboring region is sharp. On the other hand, VaSSO has 4 = pd:/ldd larger magnitude, their EMA d, is helpful for distinguishing sharp with flat valleys.
Because (x+))r on sharper valley tend to have
Memory efficient implementation. Although at first glance VaSSO has to keep both d, and e; in memory, it can be implemented in much more memory efficient manner. It is sufficient to store d, together with a scaler |dell so that e, can be recovered on demand through normalization; see (4b). Hence, VaSSO has the same memory consumption as SAM.

---

CIFAR10
SGD
SAM
ASAM
FisherSAM
VaSSO
VGG-1I-BN
93.20
93.82+0.05
93.47
93.60
94.1040.07
ResNet-18
06.2544.00
96.5810.10
96.3310.09
96.72
96.77.1009
WRN-28-10
97. 97.08
97.32+t011
97.15
97.46=0.18
97.46
97.5440.12
PyramidNet-110
97.39+0.09
97.85+0.14
97.5640.11
97.84
97.93+0.08
Table 1: Test accuracy (%) of VaSSO on various neural networks trained on CIFARIO.
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 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).
CIFAR1O. Neural networks including VGG-11, ResNet-18, 18, WRN-28-10 and PyramidNet- 110 trained on CIFARIO. 10. 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 batchsize of 128, and PyramidNet-11 10 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 p from (0.01, 0.05, 0.1, 0.2, 0.5) for SAM and find that p = 0.1 gives the best results for ResNet and WRN, respectively. ASAM and VaSSO adopt the same p as SAM. FisherSAM uses the recommended = 0.1 (Kim et al., 2022), For VaSSO, we tune 0 = {0.44 0.9) and report the best accuracy although VaSSO with both parameters outperforms SAM. We find that 0 0.4 works the best for ResNet- -8 and WRN-28- 10 while f = 0.9 achieves the best accuracy in other cases.
0.05 and p = 0.2 suit best for and VGG and PyramidNet,
SAM in all tested It is shown in Table that VaSSO offers 0.2 to 0.3 accuracy improvement scenarios except for PyramidNet-1 I0, where the improvement is about 0.1. These results illustrate that suppressed variance and the induced stabilized adversary indeed beneficial for generalizability.
CIFAR100. The training setups on this dataset are the same as those on CIFARIO, except for the best choice for p 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
7

---

CIFARI00
SGD
SAM
ASAM
FisherSAM
VaSSO
ResNet-18
77.90
80.9640.12
79.91
80.99
81.30+0..13
WRN-28-10
81.71
84.88+0.10
83.54
84.91
85.0640.05
PyramidNet-110
83.50+0.12
85.60+0.11
83.7240.09
85.55
85.85
Table 2: Test accuracy (%) of VaSSO on various neural networks trained on CIFAR100
ImageNet
vanilla
SAM
ASAM
GSAM
VaSSO
V+G
ResNet-50
76.6220.12
77.16
77.1040.16
77.20
77.4240.13
77.48
ViT-S/32
(68.1210.05
68.98
68.744..11
69.42
69.54
69.61 10.11
Table 3: Test accuracy (%) of VaSSO ImageNet, where V+G is short for VaSSO + GSAM.
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 suppressed variance vision tasks, then test VaSSO German to English translation using a Transformer (Vaswani et al., 2017) trained on IWSLT-I 14 dataset (Cettolo et al., 2014). The fairseq implementation is adopted. AdamW is chosen base optimizer in SAM and VaSSO because of its improved performance over SGD. The learning rate of Adam W is initialized to 5 10-4 and then follows inverse square root schedule. For momentum, we choose 31 = 0.9 and 32 = 0.98. Label smoothing is also applied with rate of 0.1. Hyperparameter is tuned for SAM from {0.01, 0.05, 0.1, 0.2), and p 0.1 performs the best. The same p picked for ASAM and VaSSO as well.
The validation perplexity and test BLEU scores shown in Table 4. It can be that both SAM and ASAM have better performance on validation perplexity and BLEU relative to AdamW. Although VaSSO with 0 0.9 has slightly higher validation perplexity, its BLEU score outperforms SAM and ASAM. VaSSO with f = 0.4 showcases the best generalization performance on this task, providing a BLEU score relative AdamW. This aligns with Theorems 2 and 3, which suggest that a small / is more beneficial to the stability of adversary.
0.22 improvement
4.3 Additional tests
Additional experiments are conducted to cor- roborate the merits of suppressed variance and stabilized adversary in VaSSO. In particular, this subsection evaluates several flatness related met- rics after training a ResNet- 18 on CIFAR 10 for 200 epochs, utilizing the same hyperparameters as those in Section 4.1.
SGD
SAM
VaSSO
A
82.52
26.40
23.32
A1/A5
16.63
2.12
1.86
Table 5:
Hessian spectrum of a ResNet- trained on CIFARIO.
Hessian spectrum. We first assess Hessian eigenvalues of ResNet-18 trained with SAM and VaSSO. We focus on the largest eigenvalue x3 and the ratio of largest to the fifth largest eigenvalue A These measurements are also adopted in (Foret et al., 2021; Jastrzebski al., 2020) reflect the flatness of the solution, where smaller numbers are more preferable. Because exact calculation for Hessian spectrum too expensive provided the size of ResNet-18, we instead leverage Lanezos algorithm for approximation (Ghorbani et al., 2019). The results can be found in Table 5. It can be seen that SAM indeed converges to much flatter solution compared with SGD, and VaSSO further improves upon SAM. This confirms that the friendly adversary issue is indeed alleviated by
8

---

VaSSO (0 0.9)
VaSSO (0 0.4)
A
SAM
ASAM
val. ppl.
1..03
5.00 10.04
4.99 10.03
1.003
4.99 10.03
BLEU
34.66+0.06
34.75.0.04
34.7640.04
34.81
34.88
Table 4: Performance of VaSSO for training Transformer IWSLT- 14 dataset.
VaSSO = 0.9)
VaSSO (0 = 0.4)
VaSSO (0 = 0.2)
SAM
25% % label noise 96.39+0.12
96.3640.11
96.4220012
96.48
50% % label noise 93.93
44.00
94.6340.21
94.93
75% label noise
75.3610.42
77.40
80.9410.40
85.02
Table 6: Test accuracy (%) of VaSSO on CIFARIO under different levels of label noise.
the suppressed variance in VaSSO, which in turn boosts the generalizability of ResNet- 18 as shown earlier in Section 4.1.
Label noise. 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 largely perturbed by the label noise, this is setting where the suppressed variance and stabilized adversaries expected to be advantageous. In experiments,
measure the performance VaSSO in the scenarios where certain fraction of the training labels are randomly flipped. Considering i {0.9, 0.4, 0.2], the corresponding test accuracies summarized in Table 6.
Our first observation is that VaSSO outperforms SAM at different levels of label noise. VaSSO elevates higher generalization improvement the ratio of label noise grows. In the case of 75% label noise, VaSSO with 0 = 0.4 nontrivially outperforms SAM with an absolute improvement more than 5. while VaSSO with 0 = 0.2 markedly improves SAM by roughly 10. In all scenarios, 0 = 0.2 showcases the best performance and 0 0.9 exhibits the worst generalization when comparing among VaSSO. In addition, when fixing the choice to 0, e.g.. 0 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 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 flatter valley to enhance generalizability. Targeting at 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 SAM while putting more emphases on the robustness adversarial training.
SGD (Jastrzębski al.,
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 systematically studied in (Kim 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 the cost of additional computation. Similar to SAM, VaSSO requires backpropagate twice per iteration. Various works have tackled this issue and developed lightweight SAM. AM computes the extra stochastic gradient once every a few iterations and reuses it in a fine- 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 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 predesigned schedule (Zhao et al., 2022b), or in an adaptive fashion (Jiang et al., 2023). SAF connects SAM with distillation for computational merits (Du 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 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 Moreover, VaSSO nontrivially enhances model robustness against high levels of label noise. Our results corroborate VaSSO competitive alternative of SAM.
SAM among various vision and language tasks.
Acknowledgement
This research is supported by NSF grants 2128593, 2126052, 2212318, 2220292, and 2312547. The authors would also like thank anonymous reviewers for their feedback.
References
Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-aware mini- mization. In Proc. Int. Conf. Machine Learning, pages 639-668, 2022.
David GT Barrett and Benoit Dherin. Implicit gradient regularization. Proc. Int. Conf. Learning Represention, 2021.
Peter L Bartlett, Philip M Long, and Olivier Bousquet. The dynamics of sharpness-aware minimiza- tion: Bouncing across ravines and drifting towards wide minima. arXiv:2210.01573, 2022.
Léon Bottou, Frank Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. arXiv preprint 1XXXv:6666004800 2016.
M Cettolo, J Niehues, $ Stker, L Bentivogli, and M Federico. Report 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 Hsich, and Boqing Gong. When vision transformers outperform resnets without pre-training or strong data augmentations. 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. Proc. Conf. Computer Vision and Pattern Recognition, pages 248-255, 2009.
10

---

Jacob Devlin, Ming- bidirectional transformers for language understanding. arXiv:1810.04805, 2018.
Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-traing fre-training of deep
Terrance Devries and Graham W. Taylor. Improved regularization convolutional neural networks with cutout. abs/1 708.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. Leurning 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.
Saced 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. Ini. 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
Stanislaw Jastrzebski, 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 Szymezak. 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 sharpness aware minimization. WrXiv:2304.14644 2023.
Yiding Jiang, Behnam Neyshabur, Dilip Krishnan, Hossein Mobahi, and Samy Bengio. Fantastic generalization measures and where to find them. urXiv:/9/2.02177 2019.
Rie Johnson and Tong Zhang, Accelerating stochastic gradient descent using predictive variance reduction. In Pruc. 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. Ini. 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. w:Xiv:2302.10188) 2023.
-waare
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 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 sharpness- aware minimization for scale-invariant learning of deep neural networks. 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 (006::9900022333. 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 Hsich, and Yang You. Towards efficient and scalable sharpness-awars minimization. In Proc. Conf. Computer Vision and Pattern Recognition, volume 2022, pages 12350-12360, 2022.
Hya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In Proc. Inf. Conf. Learning Represention, 2017.
K. 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. Ex- ploring generalization deep learning. In Proc. Ady. Neural Info. Processing Systems, volume 30, pages 5947-5956, 2017.
Lam M Nguyen, Jie Liu, Katya Scheinberg, and Martin Takáč. SARAH: A novel method for machine learning problems using stochastic recursive gradient. In Proc. Intl. Conf. Machine Learning, Sydney, Australia, 2017.
Sashank 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: 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, Herbert- Voss 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 llya, and Amodei Dario. I Language models are few-shot learners. In Proc. Ady. Neural Info. Processing Systems, volume 33, pages 1877-1901, 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz 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: Information- theoretic bounds and implications. 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 machine learning. In Proc. Int. Conf. Machine Learning, volume 30, pages 4148-4155 2017.
12

---

Dongxian Wu, Shu- Tao Xia, and Yisen Wang. Adversarial weight perturbation helps robust general- ization. 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-ave 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 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 minimiza- tion for efficiently training deep neural networks. arXiv:2203.09962, 2022b.
Juntang Zhuang, Boging Gong, Liangzhe Yuan, Yin Cui, Hartwig Adam, Nicha Dvorekk, Sekhar Tatikonda, James Duncan, and Ting Liu. Surrogate gap minimization improves sharpness-aware training. In Proc. Int. Conf. Learning Represention, 2022.
13

---

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
max (xx) := Eg [x(,,,))))
(6)
xEX
where X' is convex and compact constraint SFW for solving (6) is summarized below.
Algorithm 2 SFW (Reddi et al., 2016)
1: Initialize: X E X
2: for = 0, ,....,
do
3: draw iid samples | Br
4: let g: = B, Eb (1i(x,,55
5: Vi+1 = arg max (8) v)
6: X(+11 (1 X1 V++1
7: end for
has been shown (Reddi et al., 2016, Theorem 2) that one has use a sufficient large batch size B, = OTT), vi to ensure convergence of SFW. This is because line 5 in Alg. 2 is extremely sensitive gradient noise.
A.2 The adversary of SAM
By choosing h(e) = f(x, + E) and x== (0), it is not hard to observe that I-iteration SFW with To = 1 gives equivalent solution to the stochastic lincarization in SAM; ef. (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 SAM is the adoption of a constant batch size, which is typically small and far less than the ((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 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 the SFW) varies every iteration along with the updated x,-
A.3 The three dimensional example in Fig.
Detailed implementation for Fig. 2 is listed below. We Vf(x) 00.2. 00.1, 0.6). The stochastic noise is & [L1, E2, Ea]. where ε1, ε2, E3 are iid Gaussian random variables with variance scaling with 0.2, 1, respectively. We scale the variance change the SNR. We generate 100 adversaries by solving arg max]ellssp (V fxx + &, for each choice SNR. As shown 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 formu-
lation (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 (ww am' b,w, where (a,, b.) are data points and w is the
14

---

parameter to be optimized. Let the dataset have 4 samples, (a1 = 0,,b = 1); (a2 0, b2 = -1); (og -1 bg 0); and, (a4 1, b4 0). Consider 2-sharpness.
If the data partition is {1,2} and (3,4), the objective of 2-sharpness i.c. equation (3) in (Andriushchenko and Flammarion, 2022), becomes min Ei=1 max|]0)]<
If the data partition is { [1,3] and {2,4), the objective min, Σi= max fi 8). where fi the loss on partition {1,3), i.c., fi (w, 6) -(w + 8)2 + (w + 8); and
fz(w, 6) (w + 5)2 (w 5) is the loss 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 the random shuffle.
C Details on numerical results
CIFARIO and CIFARI00. For these small resolution datasets, we slightly change the first con- volution layer of ResNet18 and WRN-28-100 10 to one with 3 x kernel size, 1 stride and padding following Mi et al. (2022). The results on SGD and SAM demonstrate that the accuracy is almost. identical 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 resolution of 224 x 224, and apply random horizontal flip, normalization during training. The batch size is chosen momentum and weight decay of base optimizer, SGD, are set as 0.9 and 10-4, respectively. We further tune p from (0.05, 0.075, 0,1, 0.2), and chooses p = 0.075 for SAM. VaSSO uses 0.99. VaSSO and ASAM adopt the same p = 0.075.
128 with cosine learning rate scheduling with an initial step size 0.05. The
VIT-S/32 ImageNet. We follow the implementation of (Du al., 2022b), where 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 p = 0.05. For the implementation of GSAM and V+G, adopt the same implementation from (Zhuang et al., 2022).
D Missing proofs
Alg. can be written as
X X, + Ee
(7a)
X++1 = X; - (x
(7b)
[x(x,] where ||e: p. In SAM, have e, P g: [x,x]]] and in VaSSO we have E; p PTá.T
d.
D.1 Useful lemmas
This subsection presents useful lemmas to support our main results.
Lemma 1. Alg. I (or equivalently iteration (7)) ensures that
Lnf [[I]S(>>)) + Lp²
(VS()) Be (x(44))
Proof. To start with, we have that
(V)())), Vf(x)) g: (x(x++))
(V](x,), gx(x)) (x:))) gx(x (8)
15

---

Taking expectation conditioned on we arrive at
f(x))( Vf(x)) (x / )> xx,]
[((V/(x.), Vx(x)) (xx))) xx|] + (x))V
Bx (xc+ Mx|
(V(V ge (x)) go( |X]]
[[IV)(()))) (x.)) gi (x++ }
l9 19 LE [E[[VV/(X])) ||x, - x ||x
[xx)]
where (a) is because of Assumption 2; and (b) is because x; -x44 = C₁ and its norm is This inequality ensures that
Luf V
Lp
The (VI(x)) - g: (x [x|]]] S
2
where the last inequality is because pT) [x(x)]l 1n X₁ finishes the proof.
2p².. Taking expectation w.r.t.
Lemma 2. Alg. 1 (or equivalently iteration (7)) ensures that
E E||lax((+ < 2222" p² + 2E [[V((()) + 20².
Proof. The proof starts with bounding |g: / )ll as
(x )| |1g:( x++)) - (xx,) + (xx)]
< 2||g: (x++)) g (x.)l1 + (x
(a)
I [|1 - X1- 1| + 2x:)
(b) 2L p² + gx vs(x)) + (xx)))]1
where (a) is the result of Assumption 2; and (b) is because x and its norm is
Taking expectation conditioned on we have
E g (x++
)11
< 2222p² + 2E (x) - vf(x;)) + f
< 2T. 2 + 20
where the last inequality is because of Assumption 3. Taking expectation w.r.. the randomness in finishes the proof.
Lemma 3. Let A,11 A, + 8 with some a E (0, 1), then we huve
B A++1 S a Ao 1-
Proof.
The proof can completed by simply unrolling A4t1 and using the fact 1+a+ +...+++a <
I-a.
16

---

D.2 Proof of Theorem
Proof. Using Assumption 2, we have that
s(x+11) - -xxx)
$ (V]();), X - x,) + L
|x - xt|1
2
-V[(V((x,) 8c ) Lip
||B: (x )11
=<m(Vf(<))) gc (x++)) v fxx,) vf((x)) + Lnn
2 ( )1
f((xx)) + Lif
- -
(x 111 2
=
Taking expectation, then plugging Lemmas and in, we have
-( - -(11 3LMi
LØ²
[[IVS()))1+] + 2
1= + Info²
2
As the parameter selection ensures that 7/t 1) and rearrange the terms to arrive at
90
2 31
it is possible to divide both sides with
VT
(1- 3Lm IE [[IV)(())]]] 2
2 [[f(x,) f(x(+1)]
Lp²
+ L' + LIjo
+
2n
Summing over I, we have
1
3L4 2
E [[((x)) -((TT]
Lp²
+ L'npp² + Lnjz
+
1/T
21
(a) f(xo) -5
LPp
+
+ L L''p'² + lnno
nT
2n
f(xo) -
Lph
L'mppö
Lipo²
2npTTT
73/2
VT
where (a) uses Assumption 1, and the last equation 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
E + (
E V fxx
c.) Vf(x)) vx(x))]
21E
V (x+l ||
[/V Ex) vx(x))111
< 21E
| (x+
[|
22
2L'P6
2IE
(IV]
T
Averaging over completes the proof.
D.3 Proof of Theorem
Proof. To bound the MSE, we first have that
|d V fx
(9)
= ||(1 0)d,-1 + -gg (x.) - (xx(x))))))((xx) 0V /x,()1))1
= (1 -1 - fx(()))) 2 {2
+ +20 () ) gx())) (x())))
Now cope with three terms in the right hind of (9) separately.
17

---

The second term can be bounded directly using Assumption 2
E (x.))
(x(x))l 50²
(10)
For the third term, we have
E B (f(x)), V [xx)) = O.
(11)
The first term is bounded through
d
[((x)]]
|| ( -/ - -Vx(x))
V + (1+ + 1) | [V -
(1
$ (1 + A)ild,-1 - V )) + (1+ + x
= -1 +1+ d) V/(x,- )l + +1+ x)
+)1
where (a) is because of Young's 's inequality. Taking expectation and applying Lemma 2, we have that
E [[|dd-1 - [x())]]
(12)
< (1 + x)E[Ild,-1 [x(xx=1)11 + (1 )
21'pp² p² + 2E [ 20
1 x 0
$ (1 + A)B[ v + +
The last inequality uses the value of Ti /T
and p
In particular, have "²² = O(1/Te) vT
po.
and 0² O(o2/((2))))))) (T), and
L2 T 1 - S LELL Σ E = 0²
LB [V.
where the last equation is the result of Theorem 1. Combining (9) with (12), (10) and (11), and choosing À
0 I-Ø have
(1 0)2
0²
E [|d (xx))) ?
< (1 - 1)
002
S 00 + (1
0))00
0T
where the last inequality is the result of Lemma
D.4 Proof of Theorem 3
Proof. We adopt a unified notation for simplicity. Let V. = de for VaSSO, and v, = gx(x)) for
SAM. Then for both VaSSO and SAM, have that
f(x.) + (v), = fxx) pllv.ll f(x)) + pllv,
VS(x))
(13)
For convenience, let pVf(xx//ll From (13), we have that
f(x,) (Vo, f(x:) + p|vv; (x(x)) (x(x))]
(14)
- f(x)) + flx pllv, - vf(x))]
f(x.) + (Vf(x,), + pl|v. v)(x)))
Applying triangular inequality |||||| ||b||| < |a - b/I. we arrive
f(x:) + (ve, Ex)
= [(x+) + / (())) (VS(x)) v:)ll
(15)
2 [(xx) + -IV/((x)) - v/(x))]
= f(x,) + (Vf(x,), 4) - -plv. vf(x))l
18

---

Combining (14) with (15), we have
- L((V =
pliv - vs(x)]]
which further implies that
B[|Ci(ve) - (((f(()))] S pB fx(x)]|] - PV -
The last incquality is because (E[a])² < [B[aa] This theorem can be proved by applying Assumption 3 for SAM and Lemma 2 for VaSSO.
19


Reviews:
[{"review_id": "wJPbSnEcO0", "reviewer": "Reviewer_zZon", "summary": "**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.", "questions": "See weaknesses", "limitations": "N/A", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "**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.", "weaknesses": "**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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the time devoted to this review. The issues raised are addressed one-by –one next. \n\n**W1.** It is prudent to clarify that *m-sharpness is only tested for SAM, which may not necessarily generalize to other SAM variants*. One of the key differences with m-sharpness in SAM is that VaSSO introduces **bias** in $d_t$. As shown in our **general response,**  VaSSO works in a regime where m-sharpness does not necessarily hold. Thus, the role of noise in SAM does not imply that for VaSSO.\n\nIn addition to that, there are some other concerns that we hope to settle in our response.\n\n>  Andruschenko et al .. n-SAM .. actually observed no gain\n\n- VaSSO does not reduce to SAM even when using full gradient when finding $\\epsilon_t$; hence, the results of the aforementioned citation may not generalize to VaSSO. Moreover, n-SAM does not account for the normalization step in SAM. This can make a great difference especially when the stochastic gradient noise is pronounced. In addition, the implementation in this citation does not fully match the claim, as n-SAM is implemented by performing the `ascent step on a different batch compared to the descent step,' as shown in Appendix D of the aforementioned citation.\n\n- The experiment in this citation only observes that SAM has an accuracy of 95.8 using ResNet18 on CIFAR10, which is lower than our implementation of SGD (96.25) or SAM (96.58). This is perhaps  because a piecewise linear learning rate is adopted in this citation or no data augmentation (such as cutout) has been applied. Therefore, the results in this citation may not necessarily generalize to our setting since the *cutout* might have markedly changed the loss curvature. \n\n>  pose the problem as that SAM cannot achieve the full benefits because it does not minimize its intended objective optimally\n\nThere may have been a misunderstanding. We do not claim that solving the inner maximization optimally is helpful. We only point out that the approximation in the current SAM derivation is not perfect. And this is a chance for making SAM stronger. In fact, even if the inner maximization uses the full gradient, it is unlikely to solve the inner maximization optimally in a single step.\n\n**W2.** \n\n> Marginal improvements\n\nWe will respectfully disagree with this assessment. Most of our numerical results show a clear improvement over SAM, especially those in  Tables 3 and 6.\n\n> inconsistent numbers in comparison to Foret et al \n\nThere are several reasons for the difference on WRN. The first is that [Foret et al] uses JAX, while we work with pytorch. The second reason is that  [Foret et al] works with 8 GPUs, which allows calculation of 8 adversarial models $ \\epsilon_t^i$ for $i \\in 1, 2, .., 8$ and computes $g_t^i(x_t + \\epsilon_t^i)$ separately on each GPU. Unfortunately, this approach requires each GPU to backpropagate twice (16 times in the 8-GPU setup), which is unaffordable by our single GPU setting.\n\n> peak early-stopping accuracy or the final accuracy\n\nAs suggested, we report the peak early stopping accuracy when the noise level is 75%. SAM achieves 76.02, while VaSSO exhibits an accuracy of 83.63, which is still a considerable improvement.\n\n>  Foret et al reported better numbers for SAM on CIFAR10 although they were using ResNet34 instead of ResNet18.\n\nOur difference is that [Foret et al] trains a **ResNet32,** which has only 0.46M parameters (there is a **typo** in the review, not ResNet34). We are training a ResNet18, which is an 11M model. Our model is more than 20x larger than that in [Foret et al]; hence, the difference is reasonable. \n\n**W3.** \n\n> Theorem 2 dependent on T but not t. \n\nThis T dependence comes from the choice of $\\rho = O(1/\\sqrt{T})$ and $\\eta = O(1/\\sqrt{T})$. This is mainly to simplify the proof. However, extending the choice of hyperparameters to $\\rho_t = O(1/\\sqrt{t})$ and $\\eta_t = O(1/\\sqrt{t})$ is rather straightforward based on standard optimization techniques. For such a parameter choice, Theorem 2 can be readily modified to exhibit dependence of order $O(1 / \\sqrt{t} )$. \n\n> Comparing upper bounds does not lead to any meaningful comparison\n\nUnfortunately, the reviewer *missed the fact that VaSSO is compared with a lower bound*. \n\nRecall that a lower bound means that there exists an instance, such that $\\mathbb{E}[ || g_t(x_t) - \\nabla f(x_t) ||^2]= \\sigma^2$. The simple 1d example provided next shows that this is indeed a lower bound. \n\nLet $f(x, \\xi) = h(x) + \\xi x$, where h(x) is a deterministic loss function, and $\\xi$ is a Gaussian random variable with 0 mean and variance $\\sigma^2$. For such a loss function, it can be readily verified that \n$\\mathbb{E}[| | g_t(x_t) - \\nabla f(x_t) ||^2] = \\sigma^2$. Hence, our bound on SAM is indeed a lower bound, and the comparison is certainly meaningful."}, {"review_id": "7le3SfANkw", "reviewer": "Reviewer_BP9H", "summary": "This 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.", "questions": "Questions 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", "limitations": "Not discussed.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Paper 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", "weaknesses": "Paper 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\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the time devoted for this work. We find the comments helpful to improve the quality of our work, and we are happy to modify our manuscript accordingly. \n\n**W1.**  We have combined VaSSO with ASAM, and here are our results for CIFAR10. It can be seen that VaSSO+ASAM outperforms ASAM.\n\n|                       |   ASAM   |          VaSSO+ASAM              \n| -------------    | ---------------    | -------------   \n| ResNet18      | 96.33 $\\pm$ 0.09  | 96.52 $\\pm$ 0.12|\n| WRN-28-10   | 97.15 $\\pm$ 0.05  |  97.46 $\\pm$ 0.08 |\n\n**W2.** We respectfully disagree with this assessment. The accuracy is not easy to improve when using ResNet50 on the ImageNet. In addition, SAM improves over SGD by 0.54, and VaSSO further improves over SAM by 0.28. This is a 52% extra improvement compared to SAM’s merit over SGD, which is not marginal.\n\n**W3.** The convergence rates are the same.\n\n**Q1.** Thanks for pointing out another possible domain for VaSSO. While domain adaptation is not our central theme, it will be of interest to investigate a possible extension of VaSSO to this end. In fact, we have noticed that other sharpness aware optimization approaches have been applied to domain adaptation; see e.g., [5].  We have this direction in our future research agenda, and look forward to leverage the problem structure of domain adaption for further improvement on top of VaSSO.\n\n[5] P. Wang, et al. \"Sharpness-aware gradient matching for domain generalization.\" In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 3769-3778. 2023.\n\n\n**Q2.** Thanks for pointing out missing references. We will include them in the updated manuscript.\n\n**Q3.** Our rate is not directly comparable to [4]. This is because the inner maximization of SAM/VaSSO is changing over iterations, i.e., $\\max_{||\\epsilon|| \\leq \\rho}f(x_t + \\epsilon)$; and SAM/VaSSO only uses 1-step FW for solving its inner maximization problem per iteration; see more in Appendix 1. These differences make it difficult to compare our rates with [4]. However, it may be possible to apply methods in [4] to our framework for better bounds. We will discuss these issues when outlining our future work in the revised paper. \n\n**Q4.** The limitation of VaSSO is that it has to compute gradient twice per iteration, which we have discussed in line 329 - 331. We have also included potential solutions in the same paragraph, with numerical results shown in our responses to Q1 and Q2 of Reviewer 65jc.\n\n**Q5.** The combination of VaSSO and ASAM is shown earlier in our response to W1. For other SAM baselines, we adopt FisherSAM (Kim et al, 2022). The test accuracy on CIFAR10 is shown below, where VaSSO also has numerical merits over FisherSAM\n\n\n|                       |   FisherSAM  |          VaSSO |            \n| -------------    | ---------------          | -------------   \n| ResNet18      | 96.72 $\\pm$ 0.03   | 96.77 $\\pm$ 0.09|\n| WRN-28-10   | 97.46 $\\pm$ 0.18   |  97.54 $\\pm$ 0.12 |\n| PyramidNet110| 97.84 $\\pm$ 0.11  |  97.93 $\\pm$ 0.08|\n\n\nWe hope our responses address your concern. Let us know if there are further comments and we are happy to discuss.\n\n"}, {"review_id": "rHkcp5A6fn", "reviewer": "Reviewer_65jc", "summary": "This 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.", "questions": "- 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", "limitations": "Please see Weakness and Questions part. ", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The proposed method is clear and easy to understand\n- Theoretical results are sound\n- Empirical results demonstrate the superiority of proposed method\n", "weaknesses": "- Some claims are not fully supported\n- Some baseline methods are missing\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the time spent for coping with our submission. We hope your concerns can be addressed after reading our responses.\n\n**Q1 & Q2.**  Since updating $d_t$ in (4a) only needs the gradient on the original model , i.e., $g_t(x_t)$, which is computed every iteration, our work can be adopted jointly with the computational efficient variants  (Liu et al, 2022; Zhao et al., 2022b). For example, we combine our VaSSO with (Zhao et al., 2022b) in the pseudocode below.\n\n- For t = 0 ... T\n     - Draw a Bernoulli random variable $B_t$ ($B_t = 0$ with probability p)\n     - Calculate g_t(x_t)\n     - Update $d_t = (1- \\theta) d_{t-1} + \\theta g_t(x_t)$\n     - If $B_t = 1$, update with SGD, i.e., $x_{t+1} = x_t - \\eta_t  g_t(x_t)$\n     - If $B_t = 0$, update with VaSSO, i.e., $\\epsilon_t = \\rho d_t /|| d_t ||$ and $x_{t+1} = x_t - \\eta_t  g_t(x_t + \\epsilon_t)$\n- EndFor\n\nWe also include some numerical results with $p=0.3$ on CIFAR10  to demonstrate the efficiency of VaSSO in this case. \n\n|                       |    (Zhao et al., 2022b)   |          VaSSO+ (Zhao et al., 2022b) |    \n| -------------    | ---------------    | -------------   \n| ResNet18      | 96.37  $\\pm$ 0.13  | 96.50 $\\pm$ 0.16|\n\n\n**Q3.** We test VaSSO aided ASAM on CIFAR10, and the results show that VaSSO helps ASAM.\n\n|                       |   ASAM   |          VaSSO+ASAM              \n| -------------    | ---------------    | -------------   \n| ResNet18      | 96.33 $\\pm$ 0.09  | 96.52 $\\pm$ 0.12|\n| WRN-28-10   | 97.15 $\\pm$ 0.05  |  97.46 $\\pm$ 0.08 |\n\n**Q4.** The numerical tests for FisherSAM on CIFAR10 are shown below. FisherSAM performs slightly worse than VaSSO.\n\n|                       |   FisherSAM  |          VaSSO |            \n| -------------    | ---------------          | -------------   \n| ResNet18      | 96.73 $\\pm$ 0.03   | 96.77 $\\pm$ 0.09|\n| WRN-28-10   | 97.46 $\\pm$ 0.18   |  97.54 $\\pm$ 0.12 |\n| PyramidNet110| 97.84 $\\pm$ 0.11  |  97.93 $\\pm$ 0.08|\n\n"}, {"review_id": "gi19pKwVZG", "reviewer": "Reviewer_988S", "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.", "questions": "no question", "limitations": "no limitation", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1.The proposed method is easy to implement.\n\n2.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.\n\n2.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. \n\n[1] Foret P, Kleiner A, Mobahi H, et al. Sharpness-aware minimization for efficiently improving generalization[J]. arXiv preprint arXiv:2010.01412, 2020.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for reviewing this submission. Responses to the issues raised are provided next. \n\n**W1.** Most of existing works on m-sharpness only test SAM itself, but not SAM variants. Hence, m-sharpness does not necessarily generalize to our setting, simply because we introduce bias in $d_t$. Please see more experiments and explanation in the general response, where we numerically confirm that i) m-sharpness depends on the choice of neural network; ii) with the appearance of bias, m-sharpness even may not hold; and, iii) m-sharpness heavily depends on the specific means of updating SAM, which may not hold for different choices of $\\epsilon_t$.\n\n**W2.** \nWe will update our manuscript based on these responses."}]

Author rebuttal:
**General response to Reviewers zZon and 988S**: Our results do not conflict with existing works.

Specifically, m-sharpness does not conflict with the contribution of this submission for three reasons: 

- the behavior of m-sharpness may depend on the adopted dataset and neural network; see Experiment 1.
- m-sharpness may not hold when the gradient estimator is biased, which is the case of VaSSO; see Experiment 2.
- m-sharpness depends strongly on the specific SAM update, and its generalizability to other approaches is still well supported; see Experiment 2.

Reasons above explain why our results do not conflict with m-sharpness, simply because the present submission deals with a regime that m-sharpness may not hold. We confirm this with the following experiments.  

**Experiment 1**: m-sharpness on transformers

We conduct experiments on a transformer following the settings in section 4.2. We use a fixed batch size of B and vary the choices of m in {B/2, B/4}. For each m, we tune $\rho$ from {0.025, 0.05, 0.1, 0.2, 0.5}, respectively. The best BLEU score are reported below (larger is better).

|         |   SAM    |    B/2-SAM  |   B/4-SAM   |
| ---    | ---         | ---          | ---          |
| BLEU  | 34.75 $\pm$ 0.04  | 34.73 $\pm$ 0.02 |  34.69 $\pm$ 0.04  |

It can be seen that a smaller m actually worsens the test performance, suggesting that m-sharpness depends on the adopted neural networks and datasets.

**Experiment 2**: m-sharpness under the appearance of bias

A critical difference between m-sharpness in SAM and VaSSO is that the latter uses a biased gradient estimator $d_t$ to find $\epsilon_t$. To test the impact of bias in m-sharpness, consider the following experiment. First, fix a bias vector $\beta$. $\beta$ with the same size of gradients, and with each entry sampled from a distribution of ${\cal N}(0.01,1)$. Next, normalize $\beta$ to ensure $||\beta|| = 0.1$. We use $\beta$ as the bias to find the adversarial model, i.e., $\epsilon_t = \rho (g_t(x_t) + \beta)/ || g_t(x_t) + \beta ||$. We find that when bias is present, the m-sharpness may not hold true for ResNet18 on CIFAR10. Moreover, a decreasing trend on test accuracy is observed in this case, which illustrates how helpful is to reduce the gradient variance.

|     |   m=128    |    m=64  |   m=32   | m=16 |
| ---  | ---    | ---   | --- |  --- |
| BLEU  | 96.33 $\pm$ 0.04  | 96.26 $\pm$ 0.07 |  96.26 $\pm$ 0.10  |  96.18 $\pm$ 0.13|

Experiment 2 demonstrates also m-sharpness is highly related to the specific updateSAM is updated. Changing the way one finds $\epsilon_t$ can reverse the m-sharpness.

In sum, our Experiment 1 shows that m-sharpness depends on the neural network SAM is applied to. The exact reasoning behind m-sharpness is unclear, and it is clearly possible that m-sharpness is caused by a careless optimization step. In fact, most of the existing works on m-sharpness only test SAM, but do not consider alternative choices for inner maximization. 

Lastly, it is important to stress the reason for not studying m-sharpness directly. We find that m-sharpness formulation, e.g., eq. (3) in Andruschenko et al 2022 *may be ill-posed mathematically due to the lack of a clear definition on how the dataset ${\cal S}$ is partitioned*. Using their notation,  suppose for instance that the loss function is $l_i(x) = a_i x^2 + b_i x $, where $(a_i, b_i)$ are data points. Consider a dataset with 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)$. Let us consider 2-sharpness below under different partitions of the dataset.

- If the data partition is (1,2) and (3, 4), the objective of 2-sharpness i.e., (3) in Andruschenko et al 2022, becomes $\min_w \sum \max_{||\delta|| < \rho} 0$.
- If the data partition is (1,3) and (2,4), the objective is $\min_w \sum_{i=1}^2 \max_{||\delta|| < \rho} f_i(w,\delta)$, where $f_1$ is the loss on partition (1,3), i.e., $f_1(w,\delta) = -(w+\delta)^2 + (w+\delta)$; and $f_2(w,\delta) = (w + \delta)^2 - (w + \delta)$ is the loss on partition (2,4).

Clearly, the loss functions are not the same when the data partition varies. 

Meta-review: This paper studies SAM an points out a potential issue with the original algorithm which can adversely affect its behavior. SAM attempts to solve a minimax optimization, where the inner optimization is solved over a mini-batch. The paper points out that, the behavior of SAM can be very brittle due to sensitivity of the inner optimization to the draw of the mini-batch; one sample can drastically change the max output if the mini-batch is very small. To stabilize the inner optimization, the paper proposes a variance reduction scheme based on exponential moving average.

The paper received 4 reviews, and after the rebuttal two reviewers increased their initial score, and the paper ended up being a borderline score-wise. The main concern from the reviewers who maintained their initial scores, reviewers 988S and zZon, is some conflicting evidence in the literature that having a smaller mini-batch actually can help with generalization of SAM, an observation that is referred to as M-Sharpness. Authors' response to this concern is that, while there is some evidence about the role of M-Sharpness in generalization, the exact connection is still far from understood, and the effect of M-Sharpness on generalization can depend on add the choice of the dataset and architecture. This is backed up by a toy example in their general rebuttal response where they authors show how the different partition of the data can lead to suboptimal effect of M-Sharpness on generalization. I tend to agree with this response, and I find the main line of arguments in the paper (around the brittleness of the inner maximization) reasonable. The proposes scheme shows one cheap way to stabilize the inner optimization, and the empirical evidence in the paper suggests that this can indeed help in certain scenarios.

I encourage authors to incorporate the M-Sharpness arguments they presented in the rebuttal into the final version of the paper. Also, regarding related works on convergence of SAM, I encourage the authors to consider including a brief discussion about "Bouncing Across Ravines and Drifting Towards Wide Minima" by Bartlett et al.
Decision: Accept (poster)