Surrogate Gap Minimization Improves Sharpness-Aware Training
Authors: Juntang Zhuang
Boqing Gong
Liangzhe Yuan
Yin Cui
Hartwig Adam
Nicha C Dvornek
sekhar tatikonda
James s Duncan
Ting Liu
Venue: ICLR 2022 Poster (2022)
Keywords: generalization
sharpness-aware minimization
surrogate gap
deep learning

Abstract: The recently proposed  Sharpness-Aware  Minimization  (SAM)  improves generalization by minimizing a perturbed loss defined as the maximum loss within a neighborhood in the parameter space. However, we show that both sharp and flat minima can have a low perturbed loss, implying that SAM does not always prefer flat minima. Instead, we define a surrogate gap, a measure equivalent to the dominant eigenvalue of Hessian at a local minimum when the radius of neighborhood (to derive the perturbed loss) is small.  The surrogate gap is easy to compute and feasible for direct minimization during training. Based on the above observations, we propose Surrogate Gap Guided Sharpness-Aware Minimization (GSAM), a novel improvement over SAM with negligible computation overhead.  Conceptually, GSAM consists of two steps:  1) a gradient descent like SAM to minimize the perturbed loss, and 2) an ascent step in the orthogonal direction (after gradient decomposition) to minimize the surrogate gap and yet not affect the perturbed loss. GSAM seeks a region with both small loss (by step 1) and low sharpness (by step 2), giving rise to a model with high generalization capabilities. Theoretically, we show the convergence of GSAM and provably better generalization than SAM.Empirically, GSAM consistently improves generalization (e.g., +3.2% over SAM and +5.4% over AdamW on ImageNet top-1 accuracy for ViT-B/32). Code is released at https://sites.google.com/view/gsam-iclr22/home

Published as a conference paper at ICLR 2022
SURROGATE GAP MINIMIZATION IMPROVES SHARPNESS-AWAR TRAINING
Juntang Zhuang
Boqing Gong", Liangzhe Yuan?, Yin Cui², Hartwig Adam? [bgong, 1zyuan, yincui, hadam @googllo leppooll.ccm com
j. .zzhang@yalee.edu
Nicha C. Dvornek', Sekhar Tatikonda', James
Jamees S. Duncan' aricha.vvorne sekhar. tatikonda, james. duncan)
Ting Liu
com
1 Yale University, 2 Google Research
ABSTRACT
The recently proposed Sharpness-Aware Minimization (SAM) improves gener- alization by minimizing a perturbed loss defined as the maximum loss within a neighborhood in the parameter space. However, show that both sharp and flat minima can have low perturbed loss, implying that SAM does not always prefer flat minima. Instead, define surrogate gap, measure equivalent the dom- inant eigenvalue of Hessian at a local minimum when the radius of neighborhood (to derive the perturbed loss) small. The surrogate gap is easy to compute and feasible for direct minimization during training. Based on the above observations, we propose Surrogate Gap Guided Sharpness- novel improvement over SAM with negligible computation overhead. Conceptu- ally, GSAM consists of two steps: 1) a gradient descent like SAM to minimize the perturbed loss, and 2) an ascent step in the orthogonal direction (after gradi- ent decomposition) to minimize the surrogate gap and yet not affect the perturbed loss. GSAM seeks region with both small loss (by step 1) and low sharpness (by step 2), giving rise to model with high generalization capabilities. Theoretically. we show the convergence of GSAM and provably better generalization than SAM. Empirically, GSAM consistently improves generalization (e.g.. +3.2% over SAM and +5,4% over AdamW on ImageNet top-1 accuracy for ViT-B/32). Code is released at https:::///.......... /tspts://site... www.iiteessooeoeeooooooooooo.o.. cww/ wwwpw//wwwwwaaaaaaaiieoe gtt.//saaammmm 1ea/////
Minimization (GSAM), a
INTRODUCTION
Modern neural networks typically highly over- over- wer-parameterized and easy to overfit training data, yet the generalization performances on unseen data (test set) often suffer gap from the training performance (Zhang al. 201 Many studies try understand the generalization of machine learning models, including the Bayesian perspective (McAllester) 1999 Neyshabur et al.] 2017), the information perspective (Liang al. 2019), the loss surface geometry perspective (Hochreiter & Schmidhuber] 1995 Jiang et al. 2019) and the kernel perspective (Jacot et al. 2018 Wei et al.)
2019). Besides analyzing the properties of model after training, some works study the influence of training and the optimization process, such as the implicit regularization of stochastic gradient descent (SGD) (Bottou, 2010, Zhou et al. [2020), the learning rate's 's regularization effect (Li al.)
2019), and the influence of the batch size (Keskar et al. 2016).
These studies have led various modifications to the training process to improve generalization. Keskar & Socher (2017) proposed to use Adam in early training phases for fast convergence and then switch to SGD in late phases for better generalization. [zmailov et al. (2018) proposed to average weights to achieve a wider local minimum, which is expected to generalize better than sharp minima. A similar idea was later used in Lookahead (Zhang et al. 2019). Entropy-SGD (Chaudhari
Work was done during an internship Google

---

Published as a conference paper at ICLR 2022
et al. 2019) derived the gradient of local entropy to avoid solutions in sharp valleys. Entropy-SCo SGD has a nested Langevin iteration, inducing much higher computation costs than vanilla training.
The recently proposed Sharpness-Awar Minimization (SAM) (Foret et al. 2020) is generic train- ing scheme that improves generalization and has been shown especially effective for Vision Trans- formers (Dosovitskiy et al., 2020) when large-scale pre-training Suppose vanilla training minimizes loss f(w) (e.g., the cross-entropy loss for classification). where w is the parameter. SAM minimizes a perlurbed loss defined as fp(w) - max] f(w + 0), which is the maximum loss within radius p centered the model parameter w. Intuitively, vanilla training seeks a single point with a low loss, while SAM searches for a neighborhood within which the maximum loss is low. However, we show that low perturbed loss fp could appear both flat and sharp minima, implying that only minimizing fp is not always sharpness-aware
unavailable (Chen et al.] 2021
Although the perturbed loss fp(w) might disagree with sharpness, we find surrogate gap defined as h(w) & fp(w) - f(w) agrees with sharpness - Lemma 3.3 shows that the surrogate gap is equívalent measure of the dominant eigenvalue of Hessian at local minimum. Inspired by this observation, we propose the Surrogate Gap Guided Sharpness Aware Minimization (GSAM) which jointly minimizes the perturbed loss fp and the surrogate gap h: low perturbed loss fp indicates a low training loss within the neighborhood, and small surrogate gap h avoids solutions in sharp valleys and hence narrows the generalization gap between training and test performances (Thm. 5.3 When both criteria are satisfied, we find generalizable model with good performances.
GSAM consists of two steps for each update: 1) descend gradient V fy(w) minimize the perturbed loss fp (this step is exactly the same as SAM), and 2) decompose gradient V f(w) of the original loss f(w) into components that are parallel and orthogonal (ur) i.e., Vf(w) = VIf(w) +
V v (ww,, and perform ascent step in V (ww) to minimize the surrogate gap h(w). Note that this ascent step does not change the perturbed loss fp because fi((ww 1 V (p() by construction.
We summarize our contribution as follows:
We define surrogate gap, which measures the sharpness at local minima and is easy compute.
We propose the GSAM method to improve the generalization of neural networks. GSAM is widely applicable and incurs negligible computation overhead compared SAM.
We demonstrate the convergence of GSAM and its provably better generalization than SAM.
We empirically validate GSAM over image classification tasks with various neural architec- tures, including ResNets (He et al. 2016), Vision Transformers (Dosovitskiy et al. 2020), and MLP-Mixers (Tolstikhin al. 2021).
2
PRELIMINARIES
2.1 NOTATIONS
f(w): A loss function f with parameter w IR*, where is the parameter dimension.
Pt I: A scalar value controlling the amplitude of perturbation step
c E R: A small positive constant (to avoid division by 0, < = 10 12 by default).
f((m) wede A we + Pt TV fll The solution to max -a. [ww) when pr small.
fp(we) A max [w, + d) Ro (luado): The perturbed loss induced by f(wx). For each wi, fo(w)) returns the worst possible loss f within a ball of radius Pe centered at wy When Pe is small, by Taylor expansion, the solution the maximization problem is equivalent to gradient ascent from w; to wods."
h(w)= fp(w) f(w): The surrogate gap defined as the difference between fy(w) and fww).
Th E IR: Learning rate at step t.
a E R: A constant value that controls the sealed learning rate of the ascent step in GSAM.
(()) gp E IR: At the t-th step, the noisy observation of the gradients f((w)), (.,(w,) of the original loss and perturbed loss, respectively.
2

---

Published as a conference paper at ICLR 2022
(
Sp (2(2))
Sp(ws)
f(w))
(w2)
(wss)
h(w2
h(ws)
sharpness(w1) sharpmess(w2) > sharpness(ws), h(w)) > h(w2) > h(wa), fp(wa) fp(wa) fp(w2), hence agrees with sharpness while Jp might not.
Figure 1: Consider original loss f (solid line), perturbed loss fp _ |aal) f (w++8) (dashed line), and surrogate gap hww) - fp(w) f(w). Intuitively, f, is approximately max-pooled version of with pooling kernel of width 2p, and SAM minimizes fp. From left to right are the local minima centered W1, W2, Ws, and the valleys become flatter. Since fy(w))
fp(ws) < fp(w2), SAM prefers wi and wy to w2. However, low Jp could appear in both sharp (wn) and flat (Wa) minima, so fp might disagree with sharpness. On the contrary, a smaller surrogate gap indicates flatter W3, the loss surface is flatter, and h is smaller.
loss surface (Lemma 3.3b. From
V)(ww) = vf (w)) + fL (w:): Decompose V (w() into parallel component V [[j((e)) and vertical component fii (w,) by projection f(()) onto fpr)).
2.2 SHARPNESS- MINIMIZATION
Conventional optimization of neural networks typically minimizes the training loss fww) by gradient descent w.r.f. V V((w) and searches for single point with a low loss. However, this vanilla training often falls into sharp valley of the loss surface, resulting in inferior generalization performance (Chaudhari et al. 2019). Instead of searching for single point solution, SAM seeks a region with low losses so that small perturbation to the model weights does not cause significant performance degradation. SAM formulates the problem as:
min fp(w) where fy(w) - maxi f(w + +
where is predefined constant controlling the radius of neighborhood. This perturbed loss fp induced by f(w) is the maximum loss within the neighborhood. When the perturbed loss is minimized, the neighborhood corresponds low losses (below the perturbed loss). For small p. using Taylor expansion around w, the inner maximization in Eq. [ turns into linear constrained optimization with solution
vf(w) [
arg max Sp s(w + 5) arg maxjs ISp f(w) V[(w)+0(p)
(2)
As result, the optimization problem of SAM reduces to
Vs(w)
min fp(w) se min, aw) where u adv A w [V[(w]|+
(3)
where e a scalar (default: le- lc-12) avoid division by 0, and add is the "perturbed weight" with the highest loss within the neighborhood. Equivalently. SAM seeks a solution on the surface of the perturbed loss (w) rather than the original loss (w) (Foret al. 2020).
THE SURROGATE GAP MEASURES THE SHARPNESS AT A LOCAL MINIMUM
3.1 THE PERTURBED LOSS IS NOT ALWAYS SHARPNESS- AWARE
Despite that SAM searches for region of low losses, we show that solution by SAM is not guaranteed to flat. Throughout this paper we measure the sharpness at local minimum of loss f(w) by the dominant eigenvalue Fmaz (eigenvalue with the largest absolute value) of Hessian. For simplicity, we do not consider the influence of reparameterization on the geometry of loss surfaces, which is thoroughly discussed in (Laurent & Massart 2000 Kwon al. 2021

---

Published as a conference paper at ICLR 2022
Algorithm 1 GSAM Algorithm
GGAM
For t= 1 to T 0) schedule: Pr = Pmin +
Pmon - Pmin Mir Irmin)
Imaar Irmin
4 la) Awr = Pt 11 fi0
lb) wadv
+ Aw:
2) Get V V by back proppapppaagtionn at ufde.
3) V V Jil + V Decompose / into compo-
nents that are parallel and orthogonal to f
Vh f
4) Update weights:
Figure 2: V is decomposed into parallel and vertical (Vfl) com- ponents by projection onto V.p. VfGSAM -fp-aVfi
Vanilla Wr++1 Wt 1/e V f())
SAM W/+11 Wy T/ V 1/0
GSAM
W1+1 = we (V - aVi(
Lemma 3.1. For some fixed p. consider two local minima wi and w2, fo(wi) S fz(wz) maa (ma)(w1) S (m)r( where mmax is the dominant eigenvalue of the Hessian.
We leave the proof to Appendix. Fig. T illustrates Lemma 3.1 with an example. Consider three local minima denoted as w to W3, and suppose the corresponding loss surfaces are flatter from w to wyz For some fixed p. we plot the perturbed loss fp and surrogate gap h A Jn / around each solution. Comparing w2 with wy: Suppose their vanilla losses are equal, f(w2) = f(w3), then (w2)) > (p(3)) because the loss surface is flatter around w3, implying that SAM will prefer wy to u2- Comparing wi and w2: fy(w)) < (ww)), and SAM will favor W1 over w2 because it only cares about the perturbed loss fp. even though the loss surface is sharper around Wi than w2.
3.2 THE SURROGATE GAP AGREES WITH SHARPNESS
We introduce the surrogate gap that agrees with sharpness, defined as:
h(w) A maxi f(w + 6) fww) f(wadu) - f(w)
(4)
Intuitively, the surrogate gap represents the difference between the maximum loss within the neigh- borhood and the loss at the center point. The surrogate gap has the following properties.
Lemma 3.2. Suppose the perturbation amplitude is sufficiently small, then the approximation to the surrogate gap in Eq. iss always non-negative, h(w) [wwddv f(w) Z 0, Vw.
Lemma 3.3. For a local minimum w", consider the dominate eigenvalue maan of the Hessian of loss f as a measure of sharpness. Considering the neighborhood centered w" with a small radius the surrogate gap lu("") is an equivalent measure of the sharpness: dmaz Re 2h(w")/
The proof is Appendix. Lemma 3.2 tells that the surrogate gap is non-negative, and Lemma 3.3 shows that the loss surface is flatter as h gets closer to O. The two lemmas together indicate that we can find a region with a flat loss surface by minimizing the surrogate gap h(w).
SURROGATE GAP GUIDED SHARPNESS-AWER MINIMIZATION
4.1 GENERAL IDEA: SIMULTANEOUSLY MINIMIZE THE PERTURBED LOSS AND SURROGATE GAP
Inspired by the analysis in Section 3 we propose Surrogate Gap Guided Sharpness-Awan
Minimzation (GSAM) to simultancously minimize two objectives, the perturbed loss fp and the surrogate gap h:
ming (Fp(w), h(w)
(5)
Intuitively, by minimizng f, we search for a region with a low perturbed loss similar to SAM, and by minimizing h we search for local minimum with flat surface. A low perturbed loss implies
4

---

Published as a conference paper at ICLR 2022
low training losses within the neighborhood, and flat loss surface reduces the generalization gap between training and test performances (Chaudhari et al. 2019). When both are minimized, the solution gives rise to high accuracy and good generalization.
Potential caveat in optimization is tempting and yet sub-optimal to combine the objectives in Eq. too arrive at min, (n(()) where À is some positive scalar. One caveat when solving this weighted combination is the potential conflict between the gradients of the two terms. i.e., V f (w())
and Vh(w). We illustrate this conflict by Fig. 2 where Vh(w) f vf(w) (the grey dashed arrow) has negative inner product with (pw() and V [w). Hence, the gradient descent for the surrogate gap could potentially increase the loss fp. harming the model's performance. We empirically validate this argument in Sec. 6.4
4.2 GRADIENT DECOMPOSITION AND ASCENT FOR MULTI-OBJECTIVE OPTIMIZATION
Our primary goal is to minimize f. because otherwise flat solution of high loss is meaningless, and the minimization of should not increase fn- We propose to decompose V fw,, and Vh into components that parallel and orthogonal to V Vfp(w,) respectively (see Fig. 2
V((w)) V (wx) + (f((w)
Vh(w,) = Vhi (w)) + Thi
(6)
Vhi (wx)) - V/ (w))
The key that updating in the direction Vh. (w,) does not change the value of the perturbed loss (p(w)) because Vh1 Sp by construction. Therefore, we propose to perform descent step in the Th1 (w,) direction, which is equivalent to an ascent step in the Vf_(w; direction (because Th ff by the definition of /t), and achieve two goals simultancously it keeps the value of (we) intact and meanwhile decreases the surrogate gap (w::)) = fp(w)) - f(w) (by increasing fw:) and not affect (w:))
The full GSAM Algorithm is shown in Algo. and Fig. 2 where (()) are noisy observations of (1)
fww) and V fp(ee). respectively, and Oli -1 noisy observations of V (wx) and V (())
respectively, by projecting g( onto 97 We introduce a constant a to scale the stepsize of the ascent step. Steps I) to 2) are the same SAM: At current point w/, step 1) takes gradient ascent to wadv followed by step 2) evaluating the gradient 9p at uedes Step 3) projects g( onto 9 (t) which requires negligible computation compared to the forward and backward passes. In step 4), (t) is the same as in SAM and minimizes the perturbed loss fp(w)) with gradient descent, and -neg
(t) ang"" performs ascent step in the orthogonal direction of 9p to minimize the surrogate gap (w:) equivalently increase [(we) and keep fp (w:) intact). coding, GSAM feeds the "surrogate gradient" V JGSAM
(2)
(t) 9p agg first-order gradient optimizers such as SGD and Adam.
The ascent step along g( does not harm convergence SAM demonstrates that minimizing fp
makes the network generalize better than minimizing f. Even though our ascent step along g(f) increases f(w), it does not affect fw(),, so GSAM still decreases the perturbed loss fp in way similar to SAM. In Thm. T.I. we formally prove the convergence GSAM. In Sec. (eaad Appendix C, we empirically validate that the loss decreases and accuracy increases with training.
Illustration with a toy example We demonstrate different algorithms by a numerical toy example shown in Fig. 3 The trajectory of GSAM is closer to the ridge and tends to find flat minimum. Intuitively, since the loss surface smoother along the ridge than in sharp local minima, the sur- rogate gap h(w) is small near the ridge, and the ascent step in GSAM minimizes h to pushes the trajectory closer to the ridge. More concretely, V (((w) points to a sharp local solution and deviates from the ridge; in contrast, wode is closer the ridge and fw(dd)) is closer to the ridge descent direction than V (w,)) Note that V FCSAM and ( always lie at different sides f., (wx) by construction (see Fig. 2), hence f pushes the trajectory closer to the ridge than (w,) does. The trajectory of GSAM is like descent along the ridge and tends to find flat minima.
5

---

Published as a conference paper at ICLR 2022
ade
SGD
SAM
1
SAM
GSAM
Voluu) (uugdn
GSAM
UIA
V TSGAMM
V faa
Figure 3: Consider the loss surface with few sharp local minima. Left: Overview of the procedures of SGD, SAM and GSAM. SGD takes descent step at W; using V (w)) (orange), which points to sharp local minima. SAM first performs gradient ascent in the direction of V ([(,), to reach wode with a higher loss, followed by descent with gradient V V (wodde) (green) at the perturbed weight. Based on fww) and fu(dd.) GSAM updates in a new direction (red) that points to a flatter region. Right: Trajectories by different methods. SGD and SAM fall into different sharp local minima, while GSAM reaches a flat region. video in the supplement for better visualization.
THEORETICAL PROPERTIES OF GSAM
5.1 CONVERGENCE DURING TRAINING
Theorem 5.1. Consider a non-convex function f(w) with Lipschitz-smooth constant L and lower access a noisy, bounded observation gl (llg(1 S (t) of the
bound fmin. Suppose true gradient (w:) the t-th step. For some constant a, with learning rate πu no/ Vt, and perturbation amplitude Pt proportional to the learning rate, e.g. Pt -Po/Vt we have
C1+ C2logT
C3 + C4 log
E
Σ V f(w))
VT
T
VT
where Ci, Cz. Cg, C4 are some constants.
Thm. 5.1 implies both fp and f converge in GSAM rate O(log T/VT) for non-convex stochastic optimization, matching the convergence rate of first-order gradient optimizers like Adam.
5.2 GENERALIZATION OF GSAM
In this section, we show the surrogate gap in GSAM is provably lower than SAM's, SO GSAM is expected to find smoother minimum with better generalization.
Theorem 5.2 (PAC-Bayesian Theorem (McAllester 2003)). Suppose the training set has m ele-
ments drawn i.idd from the true distribution, and denote the loss the training as f(w) 1 E (w, Di), where we use I; to denote the (imput, target) pair of the i-th element. Let w be learned from the training set. Suppose is drawn from posterior distribution Q. Denote the prior distribution (independent of training) as P, then
Ew Oww (w, Ew-ef(w) 4v (KL()))) + log =) /m with probability at least
2m
Corollary 5.2.1. Suppose perturbation is drawn from distribution NN(0, 127 R* the dimension of w, then with probability least (1-a) 1111--1(7--00 ((22-VRR
Fuw.Bys(w,.t x)) S Sh + (KT(())))+++++++++++++++)))))))) log 2m - /m
maxioills 0) - f(w) 1
max f(w +, (xx)- -(w,
(8)

---

Published as a conference paper at ICLR 2022
where C. f(w) is the empirical training loss, and h is the surrogate gap evaluated on the training set.
implies that minimizing i (right hand side of Eq. 7 is expected to achieve tighter upper bound of the generalization performance (left hand side of Eq. (7. The third term on the right of Eq. is typically hard analyze and often simplified to L2 regularization (Foret et al.) 2020). Note that fp = C' + h only holds when Ptrain (the perturbation amplitude specified by users during training) equals Ptrue (the ground truth value determined by underlying data distribution); when
Corollary 5.2.1
Ptrain # Ptrue. min( fi) more effective than min(fp) in terms of minimizing generalization loss. A detailed discussion is in Appendix A.7
Theorem 5.3 (Unlike SAM, GSAM decreases the surrogate gap). Under the assumption in Thm. 5.7 Thm. 5.2 and Corollary 5.2.1 assume the Hessian has lower-bound (a min the absolute value of eigenvalue, and the variance of noisy observation git) is lower-bounded by c.. The surrogate gap can be minimized by the ascent step along the orthogonal direction g1 During training we minimize the sample estimate of h. We use Ah to denote the amount that the ascent step in GSAM decreases ht for the t-th step. Compared to SAM, the proposed method generates total decrease in surrogate gap Et= Alu, which bounded by
ax22 2xc2p6 la baann
lim T-000
Σ
Ahe S 2.70L 2op6
G
We provide proof in the appendix. The lower-bound of Σ Ah, indicates that GSAM achieves provably non-trivial decrease in the surrogate gap. Combined with Corollary 5.2.1 GSAM provably improves the generalization performance over SAM.
6
EXPERIMENTS
6.1 GSAM IMPROVES TEST PERFORMANCE ON VARIOUS MODEL ARCIITECTURES
We conduct experiments with ResNets (He et al. 2016). Vision Transformers (ViTs) (Dosovitskiy
al. 2020) and MLP-Mixers TTolstihhin al. [022]). ). Following the settings by Chen el (2021).
we train on the ImageNet-Ik (Deng et al. 2009) training set using the Inception-styyle (Szegedy al. 2015) pre-processing without extra training data or strong augmentation. For all models, we search for the best learning rate and weight decay for vanilla training, and then use the same values for the experiments with SAM and GSAM. For ResNets, we search for P from 0.01 to 0.05 with stepsize 0.01. For ViTs and Mixers, we search for p from 0.05 to 0.6 with stepsize 0.05. In GSAM, we search for a in {0.01, 0.02, 0.03) for ResNets and a in {0.1, 0.2, 0.3} for ViTs and Mixers. Considering that each step in SAM and GSAM requires twice the computation of vanilla training, we experiment with the vanilla training for twice the epochs of SAM and GSAM, but we observe no significant improvements from the longer training (Table in appendix). We summarize the best hhper- parameters for each model in Appendix B
We report the performances ImageNet- Real (Beyer et al., 2020 in Table GSAM consistently improves over SAM and vanilla training (with SGD or Adam W): on ViT-B/32, GSAM achieves +5.4% improvement over AdamW and +3.2% over SAM in top-1 accuracy; on Mixer-B/32. GSAM achieves +1 +11.1% over AdamW and +1.2% over SAM. We ignore the standard deviation since it is typically negligible (< 0.1%) compared to the improvements. We also test the generalization performance on off--iiiiitbutiion data (ImageNet-R and ImageNet-C), and the observation is consistent with that on ImageNet, e.g., +5.1% 11 on ImageNet-R and +5.9% on ImageNet-C for Mixer-B/32
ImageNet (Deng al.
2009). ImageNet-v2 (Recht et al. 2019) and
6.2 GSAM FINDS A MINIMUM WHOSE HESSIAN IASS SMALL DOMINANT EIGENVALUES
Lemma 3.3 indicates that the surrogate gap is an equivalent measure of the dominant eigenvalue of the Hessian, and minimizing å equivalently searches for flat minimum. We empirically validate this in Fig. g As shown in the left subfigure, for some fixed p. increasing a decreases the dominant value and improves generalization (test accuracy). In the middle subfigure, plot the dominant
7

---

Published as a conference paper at ICLR 2022
Table Top-1 Accuracy (%) on ImageNet datasets for ResNets, ViTs and MLP-Mixers trained with Vanilla SGD or AdamW, SAM, and GSAM optimizers.
Model
Training
ImageNet-vl
ImageNet- Real
ImageNet-V2
ImageNet-R
ImageNet-C
ResNet
Vanilla (SGD)
76.0
82,4
63.6
22.2
44.6
ResNet50
SAM
76.9
83.3
64.4
23,8
46.5
GSAM
77.2
83,9
64.6
23.6
47.6
Vanilla (SGD)
77.8
83.9
65.3
24.4
48.5
ResNet101
SAM
84.8
66.7
25.9
51.3
GSAM
26.3
Vanilla (SGD)
78.5
84.2
66.3
25.3
50.0
R5sNer152
SAM
79.3
84.9
67.3
25.7
$2.2
GSAM
80.0
85.9
68.6
27.3
54.1
Vision Transformer
Vanilla (AdamW)
75.2
54.3
19.0
43.3
ViT-S/32
SAM
70.5
77.5
56.9
21.4
46.2
GSAM
73.8
80.4
60.4
22.5
48.2
Vanilla (AdamW)
74.4
80.4
61.7
20.0
46.5
VIT-S/16
SAM
78.1
84.1
65.6
24.7
53.0
GSAM
79,5
85.3
67.3
25.3
53.3
Vanilla (AdamW)
71.4
77.5
57.5
23.4
44.0
ViT-B/32
SAM
73.6
80.3
60.0
24.0
50.7
GSAM
76.8
82.7
63.0
25.1
51.7
Vanilla (Adam W)
74.6
79.8
61.3
20.1
46.6
ViT-B/16
SAM
79.9
85.2
67.5
26.4
56.5
GSAM
81.0
86.5
69.2
27.1
55.7
ML.e-Mixer
Vanilla (AdamW)
63.9
70.3
49.5
16.9
35.2
Mixer-S/32
SAM
66.7
73.8
52.4
18.6
39.3
GSAM
68.6
75.8
55.0
22.6
44.6
Vanilla (AdamW)
68.8
75.1
54.8
15.9
35.6
Mixer-S/16
SAM
72.9
79.8
58.9
20.1
42.0
GSAM
81.7
61.9
23.7
48.5
Vanilla (AdamW)
76.2
56.1
34.6
Mixer-S/8
SAM
75.9
82.5
62.3
20.5
42.4
GSAM
76.8
83.4
64.0
24.6
47.8
Vanilla (AdamW)
62.5
68.1
47.6
14.6
33.8
Mixer-B/32
SAM
72.4
79.0
58.0
22.8
46.2
GSAM
73.6
80.2
59.9
27.9
52.1
Vanilla (AdamW)
66.4
72.1
50.8
14.5
33.8
Mixer-B/16
SAM
77.4
83.5
63.9
24.7
48.8
GSAM
77.8
84.0
64.9
28.3
54.4
Top-1 Accuracy
Estimation of dominant eigenvalue
Measured dominant eigenvalue
Top-1 Accuracy (9%)
75.
75
0.15
log ma (H))
0.15
74
log 2
4.0
73
72
0.05
3.0
0.15
0.2
25-
3.05
0.10
0.20
0.25
5.30
5.05
0.10
0.15
0.20
0.30
0.05
0.15
0.20
0.25
0.30
0
Figure 4: Influence of p (set as constant for ease of comparison, other experiments use decayed P: schedule) and a on the training of ViT-B/32. Left: Top-1 accuracy on ImageNet. Middle: Estimation of the dominant eigenvalues from the surrogate gap, Fmaz Rs 2h/p². Right: Dominant eigenvalues of the Hessian calculated via the power iteration. Middle and right figures match in the trend of curves, validating that the surrogate gap can be viewed as proxy of the dominant eigenvalue of Hessian.
8

---

Published as conference paper at ICLR 2022
ImageNet accuracy
ImageNet-Real accuracy
ImageNet-v2 accuracy
Top-1 Accuracy (96)
Toe-1 Accuracy cy (19))
Tap- Accuracy (96)
Entrody
Eneropy
Eneraay
70-
56
T
5AM-anzert
68-
ASA!
ASAM
54
AS4M
ASAM-ascont
ASAM 35cent
66-
52
72-
Figure Top-1 accuracy of Mixer-S/32 trained with different methods. "+ascent" represents
applying the ascent step in Algo. [ to optimizer. SAM+ascen(=GSAAM) for consistency.
Note that
GSAM is described as
Table 2: Results (%) of GSAM and min (Sp + Ah) on ViT-B/32
Table Transfer learning results (top-1 accuracy, %)
VIT-B/16
VIT-S/16
Vanilla
SAM
GSAM
Vanilla
SAM
GSAM
Dataset ImageNet ImageNet-Res
min Jp Ah)
GSAM
Cifarl0
98.1
98.6
98.8
97.6
98.2
98.4
75.4
76.8
C0far100
87.6
89.1
89.7
85.7
87.6
88.1
82.7
Flowers
88.5
91.8
91.2
86.4
90.3
ImageNet-v2
60.9
63.0
Pets
93.1
90.4
93.5
ImageNet-R 23.9
25.1
91.5
93.2
93.5
90.0
92.6
92.6
eigenvalues estimated by the surrogate gap, mmaz R 2h/p² (Lemma 3.3 In the right subfigure, we directly calculate the dominant cigenvalues using the power-iteration (Mises & Pollaczek- Geirzzggrr 1929). The estimated dominant cigenvalues (middle) match the real eigenvalues mmox (right) in terms of the trend that mmax decreases with a and p. Note that the surrogate gap is derived over the whole training set, while the measured eigenvalues
computation. These results show that the ascent step in GSAM minimizes the dominant eigenvalue by minimizing the surrogate loss, validating Thm 5.3
subset
6.3 COMPARISON WITH METHODS IN THE LITERATURE
Section 6.1 compares GSAM to SAM and vanilla training. In this subsection, we further compare GSAM against Entropy-SOH -SDD (Chaudhari et al. 2019) and Adaptive-SAM (ASAM) (Kwon et al. 2021]. which are designed to improve generalization. Note that Entropy SSD uses SGD in the inner Langevin iteration and can be combined with other base optimizers such as AdamW as the outer loop. For Entropy-SGID, we find the prpee-pprammter "scope" from 0.0 and 0.9, and search for the inner-loop iteration number between and 14. For ASAM, we search for between and 7 (10x larger than in SAM) recommended by the ASAM authors. Note that the only difference between ASAM and SAM is the derivation of the perturbation, so both can be combined with the proposed ascent step. As shown in Fig. the proposed ascent step increases test accuracy when combined with both SAM and ASAM and outperforms Entropy-SGD and vanilla training.
6.4 ADDITIONAL STUDIES
GSAM outperforms a weighted combination of the perturbed loss and surrogate gap With an example Fig. 2 we demonstrate that directly minimizing fp + Niww)) as discussed in Sec. 4.1 is sub-optimal because Vh (h(w) could conflict with (()(w) and V fw.). We empirically validate this argument on ViT-B/32. We search for between 0.0 and 0.5 with a step 0. and search for p in the same grid as SAM and GSAM. We report the best accuracy of each method. Top-1 accuracy in Table show the superior performance of GSAM, validating our analysis.
min(fp, h) vs. min(f, h) GSAM solves min(f (pp, h) by descent in V fpx decomposing V onto V fp- and an ascent step in the orthogonal direction to increase ff while keep fp intact. Alternatively, we also optimize min(f, h)) by descent in Vf, decomposing f onto Vf. and descent step in the orthogonal direction to decrease fp while keep f intact. The two GSAM variations perform similarly (see Fig. right). We choose min(f,. h)) mainly make the minimal change to SAM.
GSAM benefits transfer learning Using weights trained on ImageNet- 1k, we finetune models with SGD downstream tasks including the CIFAR 10/CIFARRRRRAA 100 (Krizhevsky al. 2009). Oxford-

---

Published as conference paper at ICLR 2022
Influence of augmentations
Influence of base optimizers
min(fp. h) v.s. min(f. h)
Improvement: Accuracy (86)
mint..mi
Top-1 Accurary (08)
78
COAM
Accuracy (%)
7b
SAM
75
mind..0
26
24
65
Light
Madium
Strong
Adam
Adaöellet
ImageNet ImagoteetRealllmagellee.
Figure 6: Top-1 accuracy of ViT-B/32 for the additional studies (Section 6.4). Left: from left right
are performances under different data augmentations (details in Appendix B.3 where the vanilla method is trained for the epochs. Middle: performance with different base optimizers. Right:
Comparison between min(Sp: h)) and mi)(J.h).
flowers Nilssack Zisserman 2008) and Oxford-IITPed (Parkhi al., 2012). Results in Table shows that GSAM leads to better transfer performance than vanilla training and SAM.
GSAM remains effective under various data augmentations We plot the top-1 accuracy of VIT-B/32 model under various Mixup (Zhang et al. (017b) augmentations in Fig. 6 subfigure). Under different augmentations, GSAM consistently outperforms SAM and vanilla training.
GSAM is compatible with different base optimizers GSAM is generic and applicable to various base optimizers. We compare vanilla training, SAM and GSAM using A (Loshchilov & Hutter 2017) and AdaBelief (Zhuang et al.) 2020) with default hyper-parameters Fig. 6 (middle subfigure) shows that GSAM performs the best, and SAM improves over vanilla training.
7 CONCLUSION
We propose the surrogate gap as an equivalent measure of sharpness which is easy to compute and feasible optimize. We propose the GSAM method, which improves the generalization over SAM at negligible computation cost. We show the convergence and provably better generalization of GSAM compared SAM, and validate the superior performance of GSAM on various models.
ACKNOWLEDGEMENT
We would like to thank Xiangning Chen (UCLA) and Hossein Mobahi (Google) for discussions, Yi Tay (Google) for help with datasets, and Yeqing Li, Xianzhi Du, and Shawn Wang (Google) for help with TensorFlow implementation.
ETHICS STATEMENT
This paper focuses on the development of optimization methodologies and can be applied to the training of different deep neural networks for wide range of applications. Therefore, the ethical impact of our work would primarily be determined by the specific models that are trained using new optimization strategy.
REPRODUCIBILITY STATEMENT
We provide the detailed proof of theoretical results in Appendix A and provide the data pre- processing and paraee-pr settings Appendix B works and public codebases, we believe the paper contains sufficient details to ensure reproducibil- ity. We plan to release the models trained by using GSAM upon publication.
Together with the references existing
REFERENCES
Randall Balestriero, Jerome Pesenti, and Yann LeCun. Learning in high dimension always amounts to extrapolation. arXiv preprint arXiv:2/ 10.0948,, 2021.
10

---

Published as a conference paper at ICLR 2022
Lucas Beyer, Olivier J. Henaff, Alexander Kolesnikov, Xiaohua Zhai, and Aaron van den Oord. Are we done with imagenet? arXiv preprint arXiv:2002.0509, 2020.
Léon Bottou. Large-scate scale machine learning with stochastic gradient descent. In Proceedings of 20MPSTAT pp. 177-186. Springer, 2010.
Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riceardo Zecchina. Entropy-sgd: Biasing gradient descent into wide valleys. Journal of Statistical Mechanics: Theory and Experiment, 2019(12):
124018, 2019.
Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong. When vision transformers outperform resnets without pretraining or strong data augmentations, 2021.
Ekin D Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoc Le. Autoaugment: Learning augmentation policies from data. arXiv preprini arXiv:/805.09501, 2018.
Alex Damian, Tengyu Ma, and Jason Lee. Label noise sgd provably prefers flat global minimizers. arXiv preprint ar arXiv:2106.06530, 2021.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Fei-Fei. Imagenet: A large-scale hi- erarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pp. 248-255. leee, 2009.
Terrance Vries and Graham W Taylor. Improved regularization of convolutional neural networks with cutout. arXiv preprint arXiv: 100.. 045522, 2017.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenbomn Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint 1rXiv:2010.11929. 2020.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(lu):2121-2155) 2011.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimiza- tion for efficiently improving generalization. arXiv preprint arX1::2010.0141222222
Xavier Gastaldi. Shake-shake regularization. arXiv preprint arXiv: 1rXiv:1705.0744: 170.00745.0 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Byeongho Heo, Sanghyuk Chun, Seong Joon Oh, Dongyoon Han, Sangdoo Yun, Gyuwan Kim, Youngjung Uh, and Jung-War Woo Ha. Adamp: Slowing down the slowdown for momentum optimiz- ers on scale-invariant weights. arXiv preprint arXiv:2006.08211 2020.
Sepp Hochreiter and Jürgen Schmidhuber. Simplifying neural nets by discovering flat minima. In Advances in neural information processing systems, pp. 529-536, 1995.
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wil- son. Averaging weights leads to wider optima and better generalization. arXiv preprint arX1v:1803.05407, 2018.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gen- eralization in neural networks. arXiv preprint 1rXiv:1806.07577 2018
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan. and Samy Bengio. Fantastic generalization measures and where to find them. arXiv preprint 192.2:19 2019.
Nitish Shirish Keskar and Richard Socher. Improving generalization performance by switching from adam to sgd. arXiv preprint
2017.

---

Published as a conference paper at ICLR 2022
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On large-batcl training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv:1609.04880 2016.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. Asam: Adaptive sharpness- aware minimization for scale-invariant learning of deep neural networks. arXiv preprint arXiv:2102.1 1600, 2021.
Beatrice Laurent and Pascal Massart. Adaptive estimation of a quadratic functional by model selec- tion. Annals of Statistics, pp. 1302-1338, 2000.
Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks. arXiv preprint (eXivv1907.04530 2019.
Tengyuan Liang, Tomaso Poggio, Alexander Rakhlin, and James Stokes. Fisher-rao metric, ge- ometry, and complexity of neural networks. In The 22nd International Conference Artificial Intelligence and Statistics, pp.
PMLR, 2019.
Tao Lin, Lingjing Kong, Sebastian Stich, and Martin Jaggi. Extrapolation for batge-batch training in deep learning. International Conference on Muchine Learning, pp. 6094-6104. PMLR, 2020.
Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. arXiv preprint arXiv:1908.0322 2019.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.5501, 2017.
Liangchen Luo, Yuanhao Xiong, Yan Liu, and Xu Sun. Adaptive gradient methods with dynamic bound of learning rate. arXiv preprint 192.00990 2019.
David McAllester. Simplified pac-bayesian margin bounds. In Learning theory and Kernel ma- chines, pp. 203-215. Springer, 2003.
David A McAllester. Pac-bayesian model averaging. In Proceedings of the twelfth annual confer- ence on Computational learning theory, pp. 164-170, 1999.
RV Mises and Hilda Pollaczek- Geirinee... Praktische verfahren der glcichungsauflöuung ZAMM- Journal of Applied Mathematics and Mechanics/Zeisscrrjft für Angewandte Mathematik und Mechanik, 9()):58-77, 1929.
Rafael Müllerr Simon Kornblith, and Geoffrey Hinton. When does label smoothing help? arXiv
preprint arXiv:1906.02629 2019.
Behnam Neyshabur, Srinadh Bhojanapalli, and Nathan Srebro. A pac-bayesian approach to spectrally-normalized margin bounds for neural networks. arXiv preprint arXiv:1707.09664, 2017.
Maria-Flent Nilsback and Andrew Zisserman. Automated flower classification over a large number of classes. In 2008 Sixth Indian Conference Computer Vision, Graphics & Image Processing, pp. 722-729. IEEE, 2008
Omkar M Parkhi, Andrea Vedaldi, Andrew Zisserman, and CV Jawahar. Cats and dogs. In 2012 IEEE conference on computer vision and pattern recognition, pp. 3498-3505. IEFF, 2012.
Benjamin Recht, Rebecca Roclofs, Ludwig Schmidt, and Vaishaal Shankar. Do imagenet classifiers generalize to imagenet? In International Conference on Machine Learning, pp. 5389-5400, 2019.
Sashank Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. arXiv preprint arXiv: /904 19t.i09297799.7 2019.
12

---

Published as a conference paper at ICLR 2022
David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams. Learning internal representations by error propagation. Technical report, California Univ San Diego La Jolla Inst for Cognitive Science, 1985.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Iya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning research, 15(1):1929-1958 2014.
Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed, Dragomir Anguelov, Du- mitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich. Going deeper with convolutions. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1-9, 2015.
Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Un- terthiner, Jessica Yung, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, et al. MIp-mixer: An all-mlp architecture for vision. arXiv preprint arXiv:2105.01601 2021.
Colin Wei, Jason Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization of neural nets their induced kernel. 2019.
Zeke Xie, Li Yuan, Zhanxing Zhu, and Masashi Sugiyama. Positive-negative momentum: Manip- ulating stochastic gradient noise improve generalization. arXiv preprint arXiv:2103.17182. 2021.
Xuho Yue, Maher Nouiehed, and Raed Al Kontar. Salr: Sharpness-aware learning rates for improved generalization. arXiv preprint arXiv:201 1.05348, 2020.
Manzil Zaheer, Sashank Reddi, Devendra Sachan, Satyen Kale, and Sanjiv Kumar. Adaptive meth- ods for nonconvex optimization. In Advances in neural information processing systems, pp. 9793- 9803, 2018
Matthew D Zeiler. Adadelta: an adaptive learning rate method. arXiv preprint 12Xiv:1212.5701, 2012.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. 2017a.
Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. arXiv preprint arXiv: 1710.09412 2017b.
Michael Zhang, James Lucas, Jimmy Ba, and Geoffrey Hinton. Lookahead optimizer: k steps forward, step back. In Advances in Neural Information Processing Systems, pp. 9593-9604, 2019.
Yaowei Zheng, Richong Zhang, and Yongyi Mao. Regularizing neural networks via adversarial model perturbation. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pp. 8156-88 1165, 2021.
Pan Zhou, Jiashi Feng, Chao Ma, Caiming Xiong, Steven Hoi, et al. Towards theoretically under- standing why sgd generalizes better than adam deep learning. arXiv preprint arXiv:2010.05667, 2020.
Juntang Zhuang, Tommy Tang, Yifan Ding, Sekhar Tatikonda, Nicha Dvornek, Xenophon Pa- pademetris, and James S Duncan. Adabelief optimizer: Adapting stepsizes by the belief in ob- served gradients. arXiv preprint axXiv20000007468 2020.
13

---

Published as a conference paper at ICLR 2022
A PROOFS
A.1 PROOF OF LEMMA.
Suppose p small, perform Taylor expansion around the local minima w, we have:
f(w 8) - f(w) ( 8+ o" H8 + 0(10)))
(10)
where is the Hessian, and is positive semidefinite at a local minima. At a local minima, Vf(w)= O. hence we have
=
f(w + 8) = f(w) + T Hб + 0(18))))
(11)
2
and
1 fo(w) = max 1ö115p f(w + +) f(w) + 70 max (H) + O(1811*)
(12)
where mer is the dominate eigenvalue (eigenvalue with the largest absolute value). Now consider two local minima W1 and W2 with dominate eigenvalue σ1 and dg respectively, we have
1 (p(w)) f(w1) +
1 =0²00
2 30001
fp(w2) ~ s(w2) +
2
We have fp(w)) > fp(w2) + > 02 and > 02 + fp(wi) > fp(w2) because the relation between f(w1) and f(w2) is undetermined.
A.2 PROOF OF LEMMA
Since p is small, we can perform Taylor expansion around wt
h(w) f(w 6) - f(w)
8 ( + 0(p²)
= plVV - 0(p²) > 0
(13)
where the last line is because is approximated as 8 (l)Va())late hence has the same direction as V f(w).
VS(w)
A.3 PROOF OF LEMMA.
Since is small, we can approximate f(w) with a quadratic model around local minima wn
1 ST H8 O(p²)
f(w + 8) f(w) +
where H is the Hessian w, assumed to be positive semidefinite local minima. Normalize 5 such that ||01|2 la = p, Hence we have:
h(w) fy(w) f(w) maxiía 25 f(w (0) fww((() manp + O(p²)
(14)
where Maa is the dominate cigenvalue of the hessian H, and first order term is 0 because the gradient is at local minima. Therefore, we have omaa Re 2h(w)/p²
A.4 PROOF OF THM.
For simplicity we consider the base optimizer is SGD. For other optimizers such as Adam, we can derive similar results by applying standard proof techniques in the literature to proof.
14

---

Published as a conference paper at ICLR 2022
STEP 1: CONVERGENCE W.R.T FUNCTION fy(w)
For simplicity of notation, we denote the update at step t
de = - + maag (t)
(15)
By I.-smoothness of f and the definition of [p(wx)) = f(wgdo), and definition of de = W++1 - and ut = w + 5g we have
w
(1 f(w(f1) s f f(uudd) + (V (Vf(ufd))
woti uadd) +
L
w wodol
(16)
2 [wwddr) + (V +VV+(V/) (V](w) We+1 8:41 - we - 8c)
+ 2 Wx+1 5z④1 Wt - 5x |
4
(17)
S f fwffdr) + (Vf(w adv) ) de) + L [d))
(18)
+ (Vf(w 8:41 - 5z) + L |6+1 - 6/ 1
||
(19)
STEP 1.0: BOUND EQ.
18
We first bound Eq. 18 Take expectation conditioned on observation up to step (for simplicity of notation, E short for E, denote expectation over all possible data points) conditioned on observations up to step t, also by definition of de, we have
Efp(w++1) - (p(x)) S -m (V fp(w)), Bggh + am (V))) Bg(") (0)
2
+ Ln! E
of6") + +ag 1
(20)
2 + +0+ + (a + 12
< B
(21)
Siince Bgg") is orthogonal to VSp(w)) by construction,
(1 < G by assumption)
STEP 1.1: BOUND EQ. 19
By definition of &,, we have
(0)
6, = Pt
(22)
"[(00
||
+
(1++1)
6.41 Pr+1
(23)
|g( +
where (t) is the gradient of f at w, evaluated with noisy data sample. When learning rate 7) is
small, the update in weight de is small, and expected gradient is
((( VS(w, de) V.(w.) Hd, + O(Idd.13
(24)
where H is the Hessian at wi- Therefore, we have
(t)
g (t+1)
E(V (VS((Wdd)) 8141 - 5z) = (V [V do
P.B.
[9(09 +e
Pc+
||g
+
(25)
g(t)
(((+1)
V f (w ade Ilp
E
(26)
1 +
[g((+1)|++
V f (w Ilero
(27)
where the first inequality is due to (1) Pe is monotonically decreasing with t, and (2) triangle in- equality that b), |||| ||b]]- de is the angle between the unit vector in the direction of V
15

---

Published as a conference paper at ICLR 2022
and V (( The second inequality comes from that (1)
< 1 strictly, so we can replace
tc
5g in Eq. 25 with unit vector corresponding directions multiplied by pr and get the upper bound,
(2) the norm of difference in unit vectors can be upper bounded by the are length on a unit circle.
When learning rate n7 and update stepsize d, is small, Φe is also small. Using the limit that
tan x + O(²²), sin = x O(r2), x x
We have:
tan Oi [ - vf(wi]l O(f))
(28)
|Hd + +111de19111 + 0
(29)
[wv( S <1,((1+a))
(30)
where the last inequality is due to (1) max eigenvalue of H is upper bounded by L because f is L -smoohh,, (2) ||d,l| = [l + ag.))l )ll and Eg. = f
Plug into Eq. 27 also note that the perturbation amplitude is small so we is close to wode, then we have
E ( (Vw(d.
det1 51) S L(1 + aa)
(31)
Similarly, we have
2 S 1
g(
g
2
E 0e41
8z
(32)
[1g00
+
||g [g((+1)|
+
S of
(33)
S p. L3
(34)
STEP 1.2: TOTAL BOUND
Reuse results from Eq. 21 (replace Lp with 2L) and plug into Eq. into Eq. 19 we have
18 and plug Eq. 31 and Eq. 34
2
2L(a
1)2
[.f((++1) - fp(w)) S -. ( (
G2
2
2
21
+
+
1h.
Pi
(35)
2
Perform telescope sum, have
T
T
Ef((wr) - (p())) $
Σ
M E (f)(wll1 + [(1
+ L(1 a)Gon
Σ
T
+ ())
(36)
t
1=1
Hence
22
E
V
f 12 < Σ
R V (p(()) fp(wo) (.f((T)) + D Dog T +
(37)
6
where
D I(1 + G2 + I(1 + a)Goo E= 1.3 + 2220
(38)
Note that n7
have
VT
T E
( fo(wo) smi + FF/6
D log
(39)
T
VT
No
VT
which implies that GSAM enables f, to converge at a rate of Oloog T/VT), and all the constants here well-hounded.
16

---

Published as a conference paper at ICLR 2022
STEP 2: CONVERGENCE W.R.T. FUNCTION f(w)
We prove the risk for f(w) convergences for non-convex stochastic optimization case using SGD.
Denote the update at step as
de -ngg + =mmgg_ (0)
(40)
By smoothness of f. we have
f(w(+1) - f(w))) + (V)()) de) + I.
2 dd
(41)
L 2
= f(wi) + (V((w),
- -7:9p (2) + +mmig
+
de
(42)
2
For simplicity, we introduce a scalar Be such that
V (w.) = B: V (we)
(43)
where fiji (we) is the projection of V (w;) onto ( When perturbation amplitude is small, we expect Be to be very close to 1.
Take expectation conditioned on observations up to step / for both sides of Eq.
42
we have:
(( S f(we)) +
[wm)
Th (V](wx))
- V[[(w)) + an. (2)
+
L
de
2
(44)
3
2
2
f(we)=
1/
V V((w))
2
+
+ +a)
+a)n n V ((.)) Vf (wee)
+ - de
1
2
(45)
3,
2
2
2
= (wee
1/e
[[(w)))
2
+
+ a)) ([(())) V[(w)) sin sin0:)
+ 2
de
2
3e
2
B3
2
(46)
(0, is the angle between (wp)) and vf(u))
f(w))=
Th
Vf(we))
| + B. a)) Vf(we)) ( taa 0,1 0(8))) + _ E
d.
3e
(47)
ssin = x + O lan = x + O(x²) when x + 0.) Also note when perturbation amplitude pe is small, we have (w,) = V f(,, + 5,) = V (w,, + [Vf(w;]l +
Pt
H Hww)) + 0
(48)
f((w) where & = Pt [(m)]l by definition, (wse is the Hessian. Hence we have
tan Itan 8 [I f(
P. |[Vf(w))l
(49)
[
where is the Lipschitz constant of f. and L smoothness of f indicates the maximum absolute
eigenvalue of H is upper bounded by L. Plug Eq. (49 Eq. 47 we have
Ef(we++1 S f(w:)=
7/t
(W))
+ (1 + a m
V (w))
tan 0.l
+
E
d
(50)
B1
B,
2
S f(w))
1/8
V f (we)
+
() + a Lp
|1,
T.
dg
2
(51)
8,
B,
2
$ f(w))
1)
f
! + B, + a LpmnG+
1
E do
12
(52)
Be
2
Assume gradient has bounded norm G.)
(53)
<s((r) Brax (wx)))
1/t
Bmin + Lpen. + BB(a+1)044
1
(54)
8, is close to assuming p is small,
hence it's natural to assume 0 8min S BB S 3maz)
17

---

Published as a conference paper at ICLR 2022
Re-arranging above formula, have
7)8 S f(w)) B[(w++))+ + (-1 +
L +1 (55)
maaz
Brnn
perform telescope sum and taking expectations on each step, we have
2 - f(wo) Ef(WT) +
T Bmin + LG
L
Vf(we))
noPt + 2 {( + 1)
Σ
M
m
Bmaz
(56)
Take the schedule be 7/t
Du and Pt V then we have Vt
T
170
S LHS
(57)
Bina VT
7m1
< RHS
(58)
< f(wo) fmmn + ( + +
Bmin
LGmopo
Σ
T
1
1
+ I
2
G2
70
2
Σ
7
1=1
(59)
- <((w)) - -fmin + 1
-(-1ttt)) + L + Io
L +
7/2 (a + 1)²G211011 +lgg?)
(60)
Hence
T
C3 VT
logT VT
7
V
S
C
(61)
1=1
where C1, C₁ are some constants. This implies the convergence rate w.r.t ((w) is O(log T/vT).
STEP 3: CONVERGENCE W.R.T. SURROGATE GAP (ww)
Note that we have proved convergence for (p(w) in step 1, and convergence for [ww in step 3. Also note that
Vh(w>) 12
VVf(u))))))((
| 2
s
( +2
(62)
Hence
T
Vh(w,)
<
2 T
T (w.)
||
2 T
Σ
f(wr))
2
(63)
also converges at rate Oloog T/ ) because each item in the RHS converges at rate Olog T TT).
A.5 PROOF OF COROLLARY. 5.2.1
Using the results from Thm. 5.2 with probability at least 0, have
E Eng B, (w, x) S f(w) + + KL P) + log 20
(64)
m
Assume 8 ~ (O (x) where & is the dimension of model parameters, hence 82 (element-wise square) follows a a Chi-square distribution. By Lemma. I in Laurent & Massart 2000), we have
P(115113 kb² 20 Vkt. 20b?) S exp(-()
(65)
hence with probability at least 1 - 1-1/vi Vi, we have
2
|| < 62 2og V - k + 2 k log v
5 262
1+
log Vn
<
p²
(66)
h:
18

---

Published as a conference paper at ICLR 2022
Therefore, with probability at least 1 1/V1= -11/v=--eepp( -vz vF)22
Es Ess(w + 8) S maxi]ollass [w +6)
(67)
Combine Eq. 65 and Eq. 67
subtract the same constant C on both sides, and under the same as- sumption as in Foret et al. 2020) that ow (w, x) 2 Eo- + ō, wewe
finish the proof.
A.6 PROOF OF TIM.. 5.3
STEP I: A SUFFICIENT CONDITION THAT THE LOSS GAP IS EXPECTED TO DECREASE FOR
EACH STEP
Take Taylor expansion, then the expected change of loss gap caused by descent step is
E (pV.) Vf(()),, - V VS Wheee Egi = V (
(68)
Th
cos Og
(69)
where 6g is the angle between vector (w:)) and V (wx). The expected change of loss gap caused by ascent step
B(V (w,) ([(w),, omp V (wx)) -ome V (wx)) <0
(70)
Above results demonstrate that ascent step decreases the loss gap, while descent step might increase the loss gap. A sufficient (but not necessary) condition for ((Vh(w,), dt) < 0 requires a to be large or || [ cos 0, S f (w:))) In practice, the perturbation amplitude is small and we can assume 9g is close to 0 and (Pp(ww)) is close to [[/]()))] we can also set the parameter a to be large in order to decrease the loss gap.
STEP 2: UPPER AND LOWER BOUND OF DECREASE IN LOSS GAP (BY THE ASCENT STEP IN ORTHOGONAL GRADIENT DIRECTION) COMPARED TO SAM.
Next we give an estimate of the decrease in π caused by our ascent step. We refer to Eq. 69 and Eq. 70 analyze the change in loss gap caused by the descent and ascent (orthogonally) respectively. can be seen that gradient descent step might not decrease loss gap, in fact they often increase loss gap in practice; while the ascent step is guaranteed to decrease the loss gap.
The decrease in loss gap is:
Abg = (V))())) V[(w), om (/. = amd [[][((([ 2
(71)
=aml comllv/ [ tan
(72)
T
T Σ
A <
aT naPt
(73)
1=1
By Eq,
49
(74)
T
aL" nop aL'nopo f3/2
(75)
1=1
< 2. 7
(76)
Hence derive upper bound for Σ Alr.
19

---

Published as a conference paper at ICLR 2022
Next derive lower bound for Ei= Ahe Note that when ps is small, by Taylor expansion
ffp(wz) f(w,, + 6t) vf(w)) +
[[V(())) (ww)) ((w)) O(p))
Pt
(77)
where H (wx) is the Hessian evaluated on training samples. Also when Pt is small, the angle 8, between (lp(w)) and V.())) is small, by the limit that
tan r = x + 0(22), + sin = x 0
We have
tan 0, | 0: 0 |0,1 + O
Omitting high order term, we have
[0.| |V V((w)ll
|l, 1l HI (w₂) + O(Ft)]l
Prlo
[an 0g|
(78)
[1f(w))]I
| (w()ll
where G the upper-bound on norm of gradient, |o (mán is the minimum absolute eigenvalue of the Hessian. The intuition is that perturbation amplitude decreases, the angle 6₂ decreases at a similar rate, though the scale constant might be different. Hence we have
T
oTp [[(u))) tan 0,1 + O
(79)
Polo mmnn
c
2
(80)
G
t=1
p6 ac²-p5no
ac²
o 2
1
min
(81)
13/2
ac2 povoo l 2 min
(82)
where e² is the lower bound of || (e.g. due to noise in data and gradient observation). Re- sults above indicate that the decrease in loss gap caused by the ascent step is non-trivial, hence proposed method efficiently improves generalization compared with SAM.
A.7 DISCUSSION ON COROLLARY 5.2.
The comment "'The corollary gives bound on the risk in terms of the perturbed training loss if one removes C' from both s is correct. But there is misunderstanding in the statement "'the perturbed training loss is small then the model has small risk"" it's only true when Ptroin for training equals its real value Ptrue determined by the data distribution; in practice, we never know Ptrue. In the following we show that the minimization of both and fp is better than simply minimizing Jp when Ptrae # Ptrain.
First, we re-write the conclusion of Corollary 5.2.1 as
E, Jn + C + R R= C C+ p20/2 p O(p
with probability (1 a) - c vk)y
where R is the regularization term, C is the training loss, is the dominant eigenvalue of Hessian. As in lemma 3.3, we perform Taylor-expansion and can ignore the high-order term Opp²). We focus on
= -+= P20/2
2. When Ptrue + Ptrains minimizing / achieves a lower risk than only minimizing fn- (1) Note that after training, C (training loss) is fixed, but could vary with p (e.g. when training on dataset A and testing on unrelated dataset B, the training loss remains unchanged, but the risk would be huge and a large is required for a valid bound). (2) With an example, we show low fp is insufficient for generalization, and low is necessary:
20

---

Published as a conference paper at ICLR 2022
A Suppose we Ptrein for training, and consider two solutions with C1, 01 (SAM) and C2 02 (GSAM), Suppose they have the same fp during training for some Ptrain: so
fpl = + 01/2 x Perain C2 + 02/2 Perain
fp2
Suppose C1 < C2 so 01 > 02
When Ptrue > Ptrains we have
risk bound.l 01 Perue + +R risk.bound.2 +2+ 02/2 Pirue + R
This implies that small helps generalization, but only low fpi (caused by low C1 and high σ1) is insufficient for good generalization.
Note that Ptraen fixed during training, so minimizing Itrain during training is equiva- lently minimizing by Lemma 3.3
3. Why outside the convex hull of the training set because " "inteppolation almost surely never occurs in high-dimensional (> 100) cases"" [alestriero al.) (202T). As a result, the variability of (train + test) sets is almost surely larger than the variability of (train) set. Since increases with data variability (see point below),
often unlucky to have Ptrue Ptrain (1) First, the test sets are almost surely
have Ptrue >> Pivain_act almost surely. (2) Second, don't know the value of Ptrue and can only guess it. In practice, we often guess small value because training often diverges with large o (as observed Foret al. (2020); Chen al. (202T)).
4. Why increases with data variability. In Corollary 5.2.1, we assume weight perturbation 5 (O (2 The meaning of b is the following. If we can randomly sample fixed number of samples from the underlying distribution, then training the model from scratch (with a fixed seed for random initialization) gives rise to set of weights. Repeating this process, we get many sets of weights, and their standard deviation is b. Since the number of training samples is limited and fixed, more variability data, the more variability in weights, and the larger b. Note that Corollary
- ( -RR) In order for the result to hold with fixed probability, must stay proportional to b, hence also increases with the variability of data.
stated that the bound holds with probability proportional
21

---

Published as a conference paper at ICLR 2022
Table 4: prper-parameters to reproduce experimental results
Model
Irmas
Trmin
Weight Decay
Hase Oprimizer
Epochs
Warmup Steps
T. schedule
ResNet30
0.04
0.02
0.01
1.6
1.6c-2
0.3
SGD
90
5k
Linear
ResNet101
0.04
0.02
0.01
1.6
1.6e-2
0.3
SGD
90
5k
Linear
ResNe1512
0.04
0.02
0.005
1.6
1.6e-2
0.3
Linear
VIT-S732
0.0
0.4
3e-3
3e-3
0.3
Adamm
300
TÖK
Linear
VIT-S/16
0.6
0.0
1.0
3e-3
3e-5
0.3
Adam
300
10k
Lineur
VIT-B/32
0.6
0.1
0.6
3e-3
3c-5
0.3
AdamW
300
10k
Linear
VIT-B/16
0.6
0.2
0.4
3e-3
3e-5
0.3
AdamW
300
10k
Linear
Mixer-S/32
0.5
0.0
0.2
3c-3
3c-5
0,3
AdamW
300
10k
Linear
Mixer-S/16
0.5
0.0
0.6
3e-3
3e-5
0.3
Adam
300
10k
Lineur
Mixer-S/8
0.1
3e-3
3e-5
0.3
AdamW
300
10k
Linear
Mixcr-B/32
0.2
0.05
3e-3
3e-5
0.3
AdamW
300
10k
Lineur
Mixer-B/16
0.5
0.2
0.01
3e-3
3e-5
0.3
A
300
10k
Linear
B EXPERIMENTAL DETAILS
B. TRAINING DETAILS
For ViT and Mixer, we search the learning rate in {le-3, 3e-3, le-2, 3e-3). and search weight decay in {0.003, 0.03, 0.3}. For ResNet, we search the learning rate in {1.6, 0.16, 0.016), and search the weight decay in {0.001, 0.01,0.13. For ViT and Mixer, we use the AdamW optimizer with 31 = 0.9, 32 = 0.999; for ResNet we use SGD with momentum= 0.9. We train ResNets for 90 epochs, and train ViTs and Mixers for 300 epochs following the settings in (Chen al., 202T) and (Dosovitskiy et al. 2020). Considering that SAM and GSAM uses twice the computation of vanilla training for each step, for vanilla training we try 2x longer training, and does not find significant improvement in Table. 5
We first search the optimal learning rate and weight decay for vanilla training, and keep these two hyper-parameters fixed for SAM and GSAM. For ViT and Mixer, search p in {0.1, 0.2, 0.3, 0.4, 0.5, 0.6} for SAM and GSAM; for ResNet, we search p from 0.01 to 0.05 with a stepsize 0.01. For ASAM, amplify p by 10x compared to SAM, recommended by Kwon al. (202T). For GSAM, we search a in 10.1, 0.2, 0.3} throughout the paper. We report the best configuration of each individual model in Table.
B.2 TRANSFER LEARNING EXPERIMENTS
Using weights trained ImageNet- Ik, finetune models with SGD downstream tasks in- cluding the CIFAR 10/CIFAR 100 (Kiihevskyy et al. 2009), Oxford-flowers (Nilsback & Zisserman 2008) and Oxford-IITPets (Parkh et al.) |l.][20222 2012). For all experiments, we use the SGD optimizer with no weight decay under a linear learning rate schedule and gradient clipping with global norm 1. We search the maximum learning rate (0.001. 0.003, 0.01, 0.03). On Cifar datasets, we train models for 10k steps with warmup step of 500; on Oxford datasets, we train models for 500 steps with wamup step of 100.
B.3 EXPERIMENTAL SETUP WITH ABLATION STUDIES ON DATA AUGMENTATION
We follow the settings Tolstikhin et al. 2021) perform ablation studies on data augmenta- tion. In the left subfigure of Fig. "Light" refers to Inception-style data augmentation with random flip and crop of images, "Medium" refers to the mixup augmentation with probability 0.2 and Ran- dAug magnitude 10; "Strong" refers to the mixup augmentation with probability 0.2 and RandAug magnitude 15.
C ABLATION STUDIES AND DISCUSSIONS
C.I INFLUENCE OF AND
We plot the performance of a ViT-B/32 model varying with p (Fig. 7a) and a (Fig. (b). We empir- ically validate that fine-tuning p in SAM achieve comparable performance with GSAM, as
22

---

Published as a conference paper at ICLR 2022
Top-1 accuracy of VIT-B/32 under different p
Top-1 accuracy of ViT-B/32 varying with
SAM
Vanilla
GSAM
SAM
Top1 Accuracy (9s)
Top-1 Accuracy
75.
74
13-
72
71
70
69
0-011
0-02
0-3 0.3
0-0.4
0-
0-0.6
0.05
a=0.10
a=
a=0.0
0=0.
(a) Performance of SAM and GSAM under different p.
(b) Performance of GSAM under different a
Figure 7: Performance of GSAM varying with p and a.
Table 5: Top-1 accuracy of ViT-B/32 on ImageNet with Inception-style data augmentation. For vanilla training report results for training 300 epochs and 600 epochs, for GSAM we report the results for 300 epochs.
Method
Epochs
ImageNet
ImageNet-Real
ImageNet-v2
ImageNet-R
300
71.4
77.5
57.5
23.4
Vanilla
600
72.0
78.2
57.9
23.6
GSAM
300
76.8
82.7
63.0
25.1
shown in Fig. 7a
Considering that GSAM has one more parameter we plot the accuracy varying and show that GSAM consistently outperforms SAM and vanilla training.
with a in Fig. 7b
C.2 CONSTANT V.S. DECAYED, SCHEDULE
Note that Thm. 5.1 assumes Pt to decay with i in order to prove the convergence, while SAM uses constant P during training. To eliminate the influence of Pt schedule, we conduct ablation study as in Table. The ascent step in GSAM can be applied to both constant and improves accuracy for both cases. Without ascent step. constant p and decayed Pt achieve similar performance. Results in Table. implies that the ascent step in GSAM is the main reason for improvement of generalization performance.
decayed Pt schedule,
Cosine
1.0
Influence of a on surrogate gap h
1.0
a=0
0.9
a=0.05
a=0.15
Cosine
Surrogate gap
0.8
d=02
0.7
0.7
2=0
3=0.06
0.6
0.6
200.15
0.5
1000
1500
Training step
500
1000
1250
1500
Training step
Figure 8: The value of cos 0, varying with train- ing steps, where 0g the angle between V fww) and V (w,) as in Fig.
Figure 9: Surrogate gap curve under different values.
23

---

Published as a conference paper at ICLR 2022
Table Top-1 Accuracy on ViT-B/32 on ImageNet. Ablation studies on constant p or a decayed pr.
Vanilla 72.0
Constant (SAM)
Constant ++ ascent
Decayed Pt. Decayed Pr ascent
75.8
76.2
75.8
76.8
C.3 VISUALIZE THE TRAINING PROCESS
In the proof of Thm. 5.3 our analysis relies on assumption that O, is small. We empirically validated this assumption by plotting cos 0e in Fig. 8 where 0g is the angle between (w(w)) and V (wr)) Note that the cosine value is calculated in the parameter space of dimension 8.8 x 10?, and in high-dimensional space two random vectors are highly likely be perpendicular. In Fig. the cosine value is always above 0.9, indicating that V (we)) and (pw( point to very close directions considering the high dimension parameters. This empirically validates our assumption that 8, is small during training.
We also plot the surrogate gap during training in Fig. 9 As a increases, the surrogate gap decreases, validating that the ascent step in GSAM efficiently minimizes the surrogate gap. Furthermore, the surrogate gap increases with training steps for any fixed a, indicating that the training process grad- ually falls into local minimum in order to minimize the training loss.
D RELATED WORKS
Besides SAM and ASAM, other methods were proposed in the literature to improve generalization: Lin et al. (020)) proposed extrapolation of gradient, Xie et al. (202T) proposed to manipulate
noise in gradient, and Damian et al. (021)) proved label noise improves generalization, Yue et al.
(2020) proposed to adjust learning rate according to sharpness, and Zheng et al.] (2021) proposed model perturbation with similar idea to SAM. Izmailov et al. (2018) proposed averaging weights to improve generalization, Heo et al.] (2020) restricted the norm of updated weights to improve
generalization. Many of aforementioned methods be combined with GSAM to further improve generalization.
Besides modified training schemes, there are other types of techniques improve generaliza- tion: data augmentation and model regularization. Data augmentation typically generates new data from training samples; besides standard data augmentation such flipping or rotation of images, recent data augmentations include label smoothing (Müller et al. 2019) and mixup (Müller et al.) 2019) which trains on convex combinations of both inputs and labels, automatically learned augmen- tation (Cubuk et al. 2018), and cutout Vries & Taylor 2017) which randomly masks out parts of an image. Model regularization typically applies auxiliary losses besides the training loss such as weight decay (Loshchilov & Hutter 2017), other methods randomly modify the model architecture during training, such as dropout (Srivastava et al. 2014) and shake-shake regularization (Gastaldi, 2017). Note that the data augmentation and model regularization literature mentioned here typically train with the standard back-propagatin (Rumelhart et al. 1985) and first-order gradient optimizers, and both techniques can be combined with GSAM.
Besides SGD, Adam and AdaBelief, GSAM can be combined with other first-order gradient opti- mizers, such as AdaBound (Luo et al. 2019), RAdam (Liu et al. 2019). Yogi (Zaheer al. 2018), AdaGrad (Duchi et al. 20.1), AMSGrad (Reddi et al. 2019) and AdaDelta [
24


Reviews:
[{"review_id": "hEFFFzFA_JF", "reviewer": "Reviewer_6UYP", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths\n\n+ I really like the idea of this paper that notices that both $f_p$ and $h$ are crucial in boosting's SAM's success. \n\n+ The empirical findings are nice to see. \n\n+ The theoretical findings are also nice to have, but perhaps not as convincing as the empirical findings. \n\nWeaknesses\n\n- I perused some of the proofs of the theoretical results. For the most part, they look fine. However, the authors can be rather sloppy at times. One cannot use $\\approx$ in a proof without quantifying what it means. For example, in the transition from (39) to (40), the authors justify the change of $\\sin\\theta_t$ to $\\tan\\theta_t$ by mentioning that $\\theta_t$ is small. I believe that the authors would do better if they quantified the remainder terms carefully. \n\n- One other deficiency is that the performance plots/tables do not include error bars. ", "summary_of_the_paper": "This paper builds on the main idea of SAM. Namely, the authors noticed that minimizing the perturbed loss $f_p$ is insufficient for guaranteeing flat minima. The authors then defined the notion of surrogate gap $h$, which they propose to minimize together with the perturbed loss. By using GSAM which consists of two gradient descent/ascent steps that include the minimizing $f_p$ and $h$, the authors show theoretically and experimentally that this results in a model with better generalization than SAM.", "summary_of_the_review": "The idea is nice. The empirical performances look promising and the theoretically analysis generally sound (though sloppy). I recommend a weak accept. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "aUc9Kr0WWXI", "reviewer": "Reviewer_ianU", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "### Strengths\n* The paper is well written, and I like the clear motivations and illustrations shown in Figure 1-3. \n* The proposed method is simple, elegant and well-motivated yet appears to work extremely well across a wide range of experiments.  \n\n### Weaknesses\n* The main optimization objective presented at the beginning of Section 4.1 is not entirely clear to me and seems not mathematically rigorous. The quantity (f_p(w), h(w)) is a vector, so is this to be understood as a multi-objective optimization task?  What is the definition of a solution of this problem?  \n* Later it is mentioned that GSAM solves min(f_p, h), but it is not clear whether the convergence theorem actually shows that GSAM finds a \"solution\" of this objective.  Instead of the decomposition of the gradient into two components, one could for example consider an algorithm which directly computes gradients of the min(f_p, h)-function.\n* I am a bit unsure whether \"surrogate gap\" is a good name for h(w). Isn't h(w) simply the sharpness as considered eg. in Foret et al?\n* An important aspect of SAM is what is called \"m-sharpness\" in Foret et al, i.e., computing the perturbed loss on mini batches.   I could not find any mentioning of the \"m\"-parameter in the paper, and since the performance of SAM greatly depends on it, it would be great to confirm that a reasonable value is used in this work. \n* Overall, I am not sure whether the paper really needs to include the theoretical results presented in Section 5.2.  They appear very hard to digest for the reader and I am unsure whether anything can be concluded from it. I believe this is a strong paper even without the theoretical analysis.", "summary_of_the_paper": "The paper presents a new method for improving generalization in deep learning called GSAM which builds upon and extends recent works on adversarial weight perturbation and sharpness-aware minimization (SAM). The paper argues that minimizing the loss function used in SAM (which replaces the loss by its maximum in a neighbourhood) does not always lead to flat minima. Based on this insight, the paper puts forward the following contributions:\n1. A new loss function which does not suffer from these problems, along with a new algorithm derived from the loss function\n2. Theoretical results (generalization bounds and convergence theorems) supporting the algorithm and new loss function\n3. An extensive numerical evaluation, showing that the proposed GSAM method leads to significant gains over SAM\n", "summary_of_the_review": "The paper has some smaller weaknesses (see main review) but if these are addressed in the rebuttal I am inclined to increase my score and recommend acceptance of the paper.  Overall, the work proposes a simple and elegant method, which is well motivated and gives state-of-the-art performance in an extensive numerical evaluation.  Therefore, I believe it will be of great interest to the ICLR community. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Oo2F3rVoh7U", "reviewer": "Reviewer_jYWK", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths\n1. This submission provides a more reasonable way to find the flatten minima rather than the heuristic way utilized in SAM. \n2. With negligible additional computation, this method improves the generalization performance.\n\nWeaknesses\n1. It is well known that SAM-based methods may lose their superior on some more complex DNN architecture, such as EffeicientNet for CNN-based model, Swin for vision transformer. This submission still performs the experiments on the vanilla architectures as the SAM did. We may wish the more exciting results on the advanced DNNs.\n2. Eq.5 seems to do not fit the updating rule provided in this submission.  Although the authors alert the potential caveat for optimizing $f_p + \\lambda h$, the proposed GSAM still aims at minimizing this objective. A minor wise modification is to exclude the component in the gradient of the second term, which may degenerate the values of the first term. Hence, the convergence guarantee may require more efforts to deal with the proposed modification; however, no more discussions are given.\n3. The convergence guarantee is a classical conclusion for non-convex stochastic optimization. If the authors cannot provide the details that show how to treat $f_p$ technically, there is no need to provide the guarantee in the main part, which cannot offer any superiority of GSAM for readers. So does the generalization guarantee, which is a direct application of general PAC-Bayesian bound. Thm. 5.3  is more interesting than the previous ones.\n", "summary_of_the_paper": "It is a good modification to the SAM method, which may further improve the performance.", "summary_of_the_review": "It is a good submission that provides an interesting way to improve the updating rule of SAM. However, some modifications for the technical part of this submission are still needed.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "GQtToqubLln", "reviewer": "Reviewer_jbS6", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The idea of minimizing the surrogate gap while maintaining the perturbed loss with gradient ascent along an orthogonal direction is interesting. The experimental results are comprehensive and seem to be convincing. Here are some issues.\n\n- In the proof of Thm 5.1, the smoothness parameter is $L_p=L+L^2\\rho/\\epsilon$. Note $\\epsilon$ is often very small. Indeed, the default choice of $\\epsilon$ is $10^{-12}$. Then the convergence bound would be very crude since $L_p$ is very large, which may not apply well in practice.\n\n- Eq (40) and (41) use inaccurate expressions such as $\\approx$. Therefore, these inequalities are not precise. The proof based on these inequalities are therefore not rigorous. Furthermore, in the proof there are some $\\beta_{\\min}$ and $\\beta_{\\max}$. These quantities are unknown in practice. They depend on the practical implementations of the algorithm and the authors do not give an upper bound on $\\beta_{\\max}/\\beta_{\\min}$. If the term $\\beta_{\\max}/\\beta_{\\min}$ is large, the derived bound can be crude. The authors should also be careful about the expectation operators. There are missing expectation operators for Eq (37)-(40), Eq (43)-(54).\n\n- The proof of Thm 5.3 is also not rigorous since the authors use several inaccurate inequalities, e.g. ($\\approx$ in (70) and (71))\n\n- In Corollary 5.2.1, the authors tried to show that GSAM can find a smoother minimizer by involving the surrogate gap in the generalization bound. However, this actually does not convey this message. The underlying reason is that both sides of eq (7) involve $C$. Therefore, Corollary 5.2.1 actually gives a bound on the risk in terms of the perturbed training loss if one removes $C$ from both sides: the perturbed training loss is small then the model has a small risk. It does not show the benefit of a small surrogate gap in achieving a better generalization, which is a claim of the paper.\n\nTypos:\n\n$g^{(t)}$ in eq (37) should be $g^{(t)}_\\bot$\n\nthere are two $++$ in eq (49)\n\nthere is a mixing $E_x$ in Thm 5.2", "summary_of_the_paper": "In this paper,the authors proposed a new method for sharpness aware training. The idea is to note a surrogate gap as a measure of sharpness, which motivates to minimize the perturbed loss and the surrogate gap simultaneously. The authors proposed to use a gradient descent to minimize the perturbed loss and gradient ascent in an orthogonal direction to minimize the surrogate gap. The use of the orthogonal direction is to avoid the influence of the perturbed loss in the ascent step. The paper gives theoretical results to show the connection between the proposed surrogate gap and the sharpness. Convergence rates and generalization bounds are also presented. Experimental results are reported which seem to be convincing.", "summary_of_the_review": "The idea of minimizing the surrogate gap and the perturbed loss simultaneously is interesting. The experimental results are sufficient. However, the theoretical analysis is not quite convincing.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes an interesting and well-motivated improvement of Sharpness Aware Minimization.  Overall the AC and reviewers are satisfied by the author feedback in improving the solidity and rigor of the theoretical results. 

The points made by the authors in response to the reviewers initial concerns are essential, especially those regarding interpretation of Corollary 5.2.1, making the proofs rigorous, and fixing the potential for crude convergence bounds. It is therefore critical that the authors incorporate them into their manuscript.
Decision: Accept (Poster)