SURROGATE GAP MINIMIZATION IMPROVES SHARPNESS-AWARE TRAINING
Year: 2022

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).

heading: 1 INTRODUCTION
text: Modern neural networks are typically highly over-parameterized and easy to overfit to training data, yet the generalization performances on unseen data (test set) often suffer from a gap from the training performance (Zhang et al., 2017a). Many studies try to understand the generalization of machine learning models, including the Bayesian perspective (McAllester, 1999; Neyshabur et al., 2017), the information perspective (Liang et 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 a model after training, some works study the influence of training and the optimization process, such as implicit regularization of stochastic gradient descent (SGD) (Bottou, 2010; Zhou et al., 2020), the learning rate’s regularization effect (Li et al., 2019), and the influence of the batch size (Keskar et al., 2016).
These studies have led to various modifications to the training process to improve generalization. Keskar & Socher (2017) proposed to use Adam in early training phases for fast convergence, then switch to SGD in late phases for better generalization. Izmailov 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 et al.,
∗Work was done during an internship at Google
2019) derived the gradient of local entropy to avoid solutions in sharp valleys. Entropy-SGD had a nested Langevin iteration, inducing much higher computation costs than vanilla training.
The recently proposed Sharpness-Aware Minimization (SAM) (Foret et al., 2020) is a generic training scheme that improves generalization and has been shown especially effective for Vision Transformers (Dosovitskiy et al., 2020) when large-scale pre-training is unavailable (Chen et al., 2021). Suppose vanilla training minimizes loss f(w) (e.g., the cross-entropy loss for classification), where w is the parameter. SAM minimizes perturbed loss defined as fp(w) , max||δ||≤ρ f(w+ δ), which is the maximum loss within radius ρ centered at 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 a low perturbed loss fp could appear in both flat and sharp minima, implying that only minimizing fp is not always sharpness-aware.
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 h is an equivalent measure of the dominant eigenvalue of Hessian at a 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: a low perturbed loss fp indicates a low training loss within the neighborhood, and a 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 a generalizable model with good performance.
GSAM consists of two steps for each update: 1) descend gradient∇fp(w) to minimize the perturbed loss fp (this step is exactly the same as SAM), and 2) decompose gradient ∇f(w) of the original loss f(w) into components that are parallel and orthogonal to ∇fp(w), i.e., ∇f(w) = ∇‖f(w) + ∇⊥f(w), and perform an ascent step in ∇⊥f(w) to minimize the surrogate gap h(w). Note that this ascent step does not change the perturbed loss fp because∇f⊥(w) ⊥ ∇fp(w) by construction. We summarize our contribution as follows:
• We define surrogate gap, which measures the sharpness at local minima and is easy to compute. • We propose the GSAM method to improve generalization of neural networks. GSAM is widely
applicable and incurs negligible computation overhead compared to 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 et al., 2021).

heading: 2 PRELIMINARIES

heading: 2.1 NOTATIONS
text: • f(w): A loss function f with parameter w ∈ Rk, where k is the dimension of parameters. • ρt ∈ R: A scalar value controlling the amplitude of perturbation at step t. • ∈ R: A small positive constant (to avoid division by 0, = 10−12 by default).
• wadvt , wt + ρt ∇f(wt) ||∇f(wt)||+ : The solution to max||w′−wt||≤ρt f(w ′) when ρt is small.
• fp(wt) , max||δ||≤ρt f(wt + δ) ≈ f(wadvt ): The perturbed loss induced by f(wt). For each wt, fp(wt) returns the worst possible loss f within a ball of radius ρt centered at wt. When ρt is small, by Taylor expansion, the solution to the maximization problem is equivalent to a gradient ascent from wt to wadvt .
• h(w) , fp(w)− f(w): The surrogate gap defined as the difference between fp(w) and f(w). • ηt ∈ R: Learning rate at step t. • α ∈ R: A constant value that controls the scaled learning rate of the ascent step in GSAM.
• g(t), g(t)p ∈ Rk: At the t-th step, the noisy observation of the gradients ∇f(wt), ∇fp(wt) of the original loss and perturbed loss, respectively.
• ∇f(wt) = ∇f‖(wt) +∇f⊥(wt): Decompose ∇f(wt) into parallel component ∇f‖(wt) and vertical component∇f⊥(wt) by projection∇f(wt) onto∇fp(wt).

heading: 2.2 SHARPNESS-AWARE MINIMIZATION
text: Conventional optimization of neural networks typically minimizes the training loss f(w) by gradient descent w.r.t. ∇f(w) and searches for a single pointw with a low loss. However, this vanilla training often falls into a sharp valley of the loss surface, resulting in inferior generalization performance (Chaudhari et al., 2019). Instead of searching for a 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:
minw fp(w) where fp(w) , max||δ||≤ρ f(w + δ) (1)
where ρ is a predefined constant controlling the radius of a 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 to low losses (below the perturbed loss). For a small ρ, using Taylor expansion around w, the inner maximization in Eq. 1 turns into a linear constrained optimization with solution
arg max||δ||≤ρ f(w + δ) = arg max||δ||≤ρ f(w) + δ >∇f(w) +O(ρ2) = ρ ∇f(w)
||∇f(w)|| (2)
As a result, the optimization problem of SAM reduces to
minw fp(w) ≈ minw f(wadv) where wadv , w + ρ ∇f(w)
||∇f(w)||+ (3)
where is a scalar (default: 1e-12) to avoid division by 0, and wadv is the “perturbed weight” with the highest loss within the neighborhood. Equivalently, SAM seeks a solution on the surface of the perturbed loss fp(w) rather than the original loss f(w) (Foret et al., 2020).

heading: 3 THE SURROGATE GAP MEASURES THE SHARPNESS AT A LOCAL MINIMUM

heading: 3.1 THE PERTURBED LOSS IS NOT ALWAYS SHARPNESS-AWARE
text: Despite that SAM searches for a region of low losses, we show that a solution by SAM is not guaranteed to be flat. Throughout this paper we measure the sharpness at a local minimum of loss f(w) by the dominant eigenvalue σmax (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 et al., 2021).
Lemma 3.1. For some fixed ρ, consider two local minima w1 and w2, fp(w1) ≤ fp(w2) 6=⇒ σmax(w1) ≤ σmax(w2), where σmax is the dominant eigenvalue of the Hessian.
We leave the proof to Appendix. Fig. 1 illustrates Lemma 3.1 with an example. Consider three local minima denoted as w1 to w3, and suppose the corresponding loss surfaces are flatter from w1 to w3. For some fixed ρ, we plot the perturbed loss fp and surrogate gap h around each solution. Comparing w2 with w3: Suppose their vanilla losses are equal, f(w2) = f(w3), then fp(w2) > fp(w3) because the loss surface is flatter around w3, implying that SAM will prefer w3 to w2. Comparing w1 and w2: fp(w1) < fp(w2), and SAM will favor w1 over w2 because it only cares about the perturbed loss fp, even though the loss surface is sharper around w1 than w2.

heading: 3.2 THE SURROGATE GAP AGREES WITH SHARPNESS
text: We introduce the surrogate gap that agrees with sharpness, defined as:
h(w) , f(wadv)− f(w) (4)
Intuitively, the surrogate gap represents the difference between the maximum loss within the neighborhood 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. 4 is always non-negative, h(w) = f(wadv)− f(w) ≥ 0,∀w. Lemma 3.3. For a local minimum w∗, consider the dominate eigenvalue σmax of the Hessian of loss f as a measure of sharpness. Considering the neighborhood centered at w∗ with a small radius ρ, the surrogate gap h(w∗) is an equivalent measure of the sharpness: σmax ≈ 2h(w∗)/ρ2.
The proof is in 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 0. The two lemmas together indicate that we can find a region with a flat loss surface by minimizing the surrogate gap h(w).

heading: 4 SURROGATE GAP GUIDED SHARPNESS-AWARE MINIMIZATION

heading: 4.1 GENERAL IDEA: JOINTLY MINIMIZE THE PERTURBED LOSS AND SURROGATE GAP
text: Inspired by the analysis in Section 3, we propose Surrogate Gap Guided Sharpness-Aware Minimzation (GSAM) to minimize both the perturbed loss fp and surrogate gap h:
minw ( fp(w), h(w) ) . (5)
Intuitively, by minimizng fp we search for a region with a low perturbed loss similar to SAM, and by minimizing h we search for a local minimum with a flat surface. A low perturbed loss implies low training losses within the neighborhood, and a 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 It is tempting and yet sub-optimal to combine the objectives in Eq. 5 to arrive at minw fp(w)+λh(w), 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., ∇fp(w) and ∇h(w). We illustrate this conflict by Fig. 2, where ∇h(w) = ∇fp(w) − ∇f(w) (the grey dashed arrow) has a negative inner product with ∇fp(w) and ∇f(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.

heading: 4.2 GRADIENT DECOMPOSITION AND ASCENT FOR THE MULTI-OBJECTIVE OPTIMIZATION
text: Our primary goal is to minimize fp because otherwise a flat solution of high loss is meaningless, and the minimization of h should not increase fp. We propose to decompose ∇f(wt) into components that are parallel and orthogonal to∇fp(wt), respectively (see the black and blue arrows in Fig. 2):
∇f(wt) = ∇f‖(wt) +∇f⊥(wt). (6)
Figure 2: Illustration of GSAM.
Algorithm 1 GSAM Algorithm For t = 1 to T
0) ρt schedule: ρt = ρmin + (ρmax−ρmin)(lr−lrmin)
lrmax−lrmin
1) ∆wt = ρt g (t)
||g(t)||+ , If adaptive: ∆wt ← ∆wt × abs(wt) wadvt = wt + ∆wt
2) Get g(t)p by back-propagation at wadvt . 3) g(t) = g(t)‖ + g (t) ⊥ Decompose g (t) into components
that are parallel and orthogonal to g(t)p . 4) Update weights:
Vanilla wt+1 = wt − ηtg(t) SAM wt+1 = wt − ηtg(t)p GSAM wt+1 = wt − ηt(g(t)p − αg(t)⊥ )
The key is that updating in the direction of∇f⊥(wt) does not change the value of the perturbed loss fp(wt) because ∇f⊥ ⊥ ∇fp by construction. Therefore, we propose to perform an ascent step in the ∇f⊥(wt) direction, which achieves two goals simultaneously — it keeps the value of fp(wt) intact and meanwhile decreases the surrogate gap h(wt) = fp(wt)− f(wt) (by increasing f(wt)).
The full GSAM Algorithm is shown in Algo. 1 and Fig. 2, where g(t), g(t)p are noisy observations of ∇f(wt) and∇fp(wt), respectively, and g(t)‖ , g (t) ⊥ are noisy observations of∇f‖(wt) and∇f⊥(wt), respectively, by projecting g(t) onto g(t)p . We introduce a constant α to scale the stepsize of the ascent step. Steps 1) to 2) are the same as SAM: At current point wt, step 1) takes a gradient ascent to wadvt followed by step 2) evaluating the gradient g (t) p at wadvt . Step 3) projects g
(t) onto g(t)p , which requires negligible computation compared to the forward and backward pass. In step 4), −ηtg(t)p is the same as in SAM and minimizes the perturbed loss fp(wt) with gradient descent, and αηtg (t) ⊥ performs an ascent step in the orthogonal direction of g (t) p to minimize the surrogate gap h(wt) ( equivalently increase f(wt) and keep fp(wt) intact). In coding, GSAM feeds the “surrogate gradient”∇fGSAMt , g (t) p − αg(t)⊥ to first-order gradient optimizers such as SGD and Adam.
The ascent step along g(t)⊥ does not harm convergence SAM demonstrates that minimizing fp successfully trains the network and generalizes better than minimizing f . Even though our ascent step along g(t)⊥ increases f(w), it does not affect fp(w), so GSAM still decreases the perturbed loss fp in a way similar to SAM. In Thm. 5.1, we formally prove the convergence of GSAM. In Sec. 6 and 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 a flat minimum. Intuitively, since the loss surface is smoother along the ridge than in sharp local minima, the surrogate 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,∇f(wt) points to a sharp local solution and deviates from the ridge; in contrast, wadvt is closer to the ridge and ∇f(wadvt ) is closer to the ridge descent direction than ∇f(wt). Note that ∇fGSAMt and ∇f(wt) always lie at different sides of ∇fp(wt) by construction (see Fig. 2), hence∇fGSAMt pushes the trajectory closer to the ridge than∇fp(wt) does. The trajectory of GSAM is like descent along the ridge and tends to find flat minima.

heading: 5 THEORETICAL PROPERTIES OF GSAM

heading: 5.1 CONVERGENCE DURING TRAINING
text: Theorem 5.1. Consider a non-convex function f(w) with Lipschitz-smooth constant L and lower bound fmin. Suppose we can access a noisy, bounded observation g(t) (||g(t)||2 ≤ G,∀t) of the true gradient ∇f(wt) at the t-th step. For some constant α, with learning rate ηt = η0/ √ t, and
perturbation amplitude ρt proportional to the learning rate, e.g., ρt = ρ0/ √ t, we have
1
T T∑ t=1 E ∣∣∣∣∣∣∇fp(wt)∣∣∣∣∣∣2 2 ≤ C1 + C2 log T√ T ,
1
T T∑ t=1 E ∣∣∣∣∣∣∇f(wt)∣∣∣∣∣∣2 2 ≤ C3 + C4 log T√ T
where C1, C2, C3, C4 are some constants. Thm. 5.1 implies both fp and f converge in GSAM at rate O(log T/ √ T ) for non-convex stochastic optimization, matching the convergence rate of first-order gradient optimizers like Adam.

heading: 5.2 GENERALIZATION OF GSAM
text: In this section, we show the surrogate gap in GSAM is provably lower than SAM’s, so GSAM is expected to find a smoother minimum with better generalization. Theorem 5.2 (PAC-Bayesian Theorem (McAllester, 2003)). Suppose the training set has m elements drawn i.i.d. from the true distribution, and denote the loss on the training set as f̂(w) = 1 m ∑m i=1 f(w, xi), where we use xi to denote the (input, target) pair of the i-th element. Let w be learned from the training set. Suppose w is drawn from posterior distribution Q. Denote the prior distribution (independent of training) as P , then
Ew∼QExf(w, x) ≤ Ew∼Qf̂(w) + 4 √(
KL(Q||P) + log 2m a
) /m with probability at least 1− a
Corollary 5.2.1. Suppose perturbation δ is drawn from distribution δ ∼ N (0, b2Ik), δ ∈ Rk, k is the dimension of w, then with probability at least ( 1− a )[ 1− e− ( ρ√ 2b − √ k )2]
Ew∼QExf(w, x)− C ≤ max||δ||2≤ρ f̂(w + δ)− C︸ ︷︷ ︸ ĥ +4
√( KL(Q||P) + log 2m
a
) /m (7)
ĥ , max||δ||2≤ρ f̂(w + δ)− f̂(w) = 1
m m∑ i=1 [ max||δ||2≤ρ f(w + δ, xi)− f(w, xi) ] (8)
where C = f̂(w) is the empirical training loss, and ĥ is the surrogate gap evaluated on the training set.
Corollary 5.2.1 implies that minimizing ĥ (right hand side of Eq. 7) is expected to achieve a tighter upper bound of the generalization gap (left hand side of Eq. 7). The second term on the right of Eq. 7 is typically hard to analyze and often simplified to L2 regularization (Foret et al., 2020). A detailed discussion is in Appendix A.7.
Theorem 5.3 (Unlike SAM, GSAM decreases the surrogate gap). Under the assumption in Thm. 5.1, Thm. 5.2 and Corollary 5.2.1, we assume the Hessian has a lower-bound |σ|min on the absolute value of eigenvalue, and the variance of noisy observation g(t) is lower-bounded by c2. The surrogate gap h can be minimized by the ascent step along the orthogonal direction g(t)⊥ . During training we minimize the sample estimate of h. We use ∆ĥt to denote the amount that the ascent step in GSAM decreases ĥ for the t-th step. Compared to SAM, the proposed method generates a total decrease in surrogate gap ∑T t=1 ∆ĥt, which is bounded by
αc2ρ20η0|σ|2min G2 ≤ lim T→∞ T∑ t=1 ∆ĥt ≤ 2.7αL2η0ρ20 (9)
We provide proof in the appendix. The lower-bound of ∑T t=1 ∆ĥt indicates that GSAM achieves a provably non-trivial decrease in the surrogate gap. Combined with Corollary 5.2.1, GSAM provably improves the generalization performance over SAM.

heading: 6 EXPERIMENTS

heading: 6.1 GSAM IMPROVES TEST PERFORMANCE ON VARIOUS MODEL ARCHITECTURES
text: We conduct experiments with ResNets (He et al., 2016), Vision Transformers (ViTs) (Dosovitskiy et al., 2020) and MLP-Mixers (Tolstikhin et al., 2021). Following the settings by Chen et al. (2021), we train on the ImageNet-1k (Deng et al., 2009) training set using the Inception-style (Szegedy et 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 ρ from 0.01 to 0.05 with a stepsize 0.01. For ViTs and Mixers, we search for ρ from 0.05 to 0.6 with a stepsize 0.05. In GSAM, we search for α in {0.01, 0.02, 0.03} for ResNets and α 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 5 in appendix). We summarize the best hyper-parameters for each model in Appendix B.
We report the performances on ImageNet (Deng et al., 2009), ImageNet-v2 (Recht et al., 2019) and ImageNet-Real (Beyer et al., 2020) in Table 1. GSAM consistently improves over SAM and vanilla training (with SGD or AdamW): 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 +11.3% over AdamW and +1.4% 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 out-of-distribution data (ImageNet-R and ImageNet-C), and the observation is consistent with that on ImageNet, e.g., +5.0% on ImageNet-R and +6.2% on ImageNet-C for Mixer-B/32.

heading: 6.2 GSAM FINDS A MINIMUM WHOSE HESSIAN HAS SMALL DOMINANT EIGENVALUES
text: Lemma 3.3 indicates that the surrogate gap h is an equivalent measure of the dominant eigenvalue of the Hessian, and minimizing h equivalently searches for a flat minimum. We empirically validate this in Fig. 4. As shown in the left subfigure, for some fixed ρ, increasing α decreases the dominant value and improves generalization (test accuracy). In the middle subfigure, we plot the dominant eigenvalues estimated by the surrogate gap, σmax ≈ 2h/ρ2 (Lemma 3.3). In the right subfigure, we directly calculate the dominant eigenvalues using the power-iteration (Mises & Pollaczek-Geiringer, 1929). The estimated dominant eigenvalues (middle) match the real eigenvalues σmax (right) in
Table 2: Results (%) of GSAM
and min(fp+λh) on ViT-B/32
Dataset min(fp + λh) GSAM ImageNet 75.4 76.8 ImageNet-Real 81.1 82.8 ImageNet-v2 60.9 62.8 ImageNet-R 23.9 25.3
ViT-B/16 Mixer-B/16 Mixer-S/16
Vanilla SAM GSAM Vanilla SAM GSAM Vanilla SAM GSAM CIFAR10 98.1 98.6 98.9 95.4 97.8 98.5 94.1 96.1 98.4 CIFAR100 87.6 89.1 89.7 80.0 86.4 88.0 77.9 82.4 87.8 Flowers 88.5 91.8 91.7 82.8 90.0 90.0 83.3 87.9 90.5 Pets 91.9 93.1 94.1 86.1 92.5 93.5 86.1 88.7 93.5 mean 91.5 93.2 93.6 86.1 91.7 92.5 85.4 88.8 92.6
terms of the trend that σmax decreases with α and ρ. Note that the surrogate gap h is derived over the whole training set, while the measured eigenvalues are over a subset to save computation. These results show that the ascent step in GSAM minimizes the dominant eigenvalue by minimizing the surrogate loss, validating Thm 5.3.

heading: 6.3 COMPARISON WITH METHODS IN THE LITERATURE
text: Section 6.1 compares GSAM to SAM and vanilla training. In this subsection, we further compare GSAM against Entropy-SGD (Chaudhari et al., 2019) and Adaptive-SAM (ASAM) (Kwon et al., 2021), which are designed to improve generalization. Note that Entropy-SGD uses SGD in the inner Langevin iteration and can be combined with other base optimizers such as AdamW as the outer loop. For Entropy-SGD, we find the hyper-parameter “scope” from 0.0 and 0.9, and search for the inner-loop iteration number between 1 and 14. For ASAM, we search for ρ between 1 and 7 (10× larger than in SAM) as 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. 5, the proposed ascent step increases test accuracy when combined with both SAM and ASAM and outperforms Entropy-SGD and vanilla training.

heading: 6.4 ADDITIONAL STUDIES
text: GSAM outperforms a weighted combination of the perturbed loss and surrogate gap With an example in Fig. 2, we demonstrate that directly minimizing fp(w) + λh(w) as discussed in Sec. 4.1 is sub-optimal because ∇h(w) could conflict with ∇fp(w) and ∇f(w). We empirically validate this argument on ViT-B/32. We search for λ between 0.0 and 0.5 with a step 0.1 and search for ρ in the same grid as SAM and GSAM. We report the best accuracy of each method. Top-1 accuracy in Table 2 show the superior performance of GSAM, validating our analysis.
min(fp, h) vs. min(f, h) GSAM solves min(fp, h) by descent in∇fp, decomposing∇f onto∇fp, and an ascent step in the orthogonal direction to increase f while keep fp intact. Alternatively, we can also optimize min(f, h) by descent in ∇f , decomposing ∇fp onto ∇f , and a descent step in the orthogonal direction to decrease fp while keep f intact. The two GSAM variations perform similarly (see Fig. 6, right). We choose min(fp, h) mainly to make the minimal change to SAM.
GSAM benefits transfer learning Using weights trained on ImageNet-1k, we finetune models with SGD on downstream tasks including the CIFAR10/CIFAR100 (Krizhevsky et al., 2009), Oxfordflowers (Nilsback & Zisserman, 2008) and Oxford-IITPets (Parkhi et al., 2012). Results in Table 3 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 a ViT-B/32 model under various Mixup (Zhang et al., 2017b) augmentations in Fig. 6 (left 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 AdamW (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.

heading: 7 CONCLUSION
text: We propose the surrogate gap as an equivalent measure of sharpness which is easy to compute and feasible to 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 to SAM, and validate the superior performance of GSAM on various models.

heading: ETHICS STATEMENT
text: This paper focuses on the development of optimization methodologies and can be applied to the training of different deep neural networks for a wide range of applications. Therefore, the ethical impact of our work would primarily be determined by the specific models that are trained using our new optimization strategy.

heading: REPRODUCIBILITY STATEMENT
text: We provide the detailed proof of theoretical results in Appendix A and provide the data preprocessing and hyper-parameter settings in Appendix B. Together with the references to existing works and public codebases, we believe the paper contains sufficient details to ensure reproducibility. We plan to release the models trained by using GSAM upon publication.

heading: A PROOFS

heading: A.1 PROOF OF LEMMA. 3.1
text: Suppose ρ is small, perform Taylor expansion around the local minima w, we have:
f(w + δ) = f(w) +∇f(w)>δ + 1 2 δ>Hδ +O(||δ||3) (10)
whereH is the Hessian, and is positive semidefinite at a local minima. At a local minima,∇f(w) = 0, hence we have
f(w + δ) = f(w) + 1
2 δ>Hδ +O(||δ||3) (11)
and
fp(w) = max||δ||≤ρ f(w + δ) = f(w) + 1
2 ρ2σmax(H) +O(||δ||3) (12)
where σmax is the dominate eigenvalue (eigenvalue with the largest absolute value). Now consider two local minima w1 and w2 with dominate eigenvalue σ1 and σ2 respectively, we have
fp(w1) ≈ f(w1) + 1
2 ρ2σ1 fp(w2) ≈ f(w2) +
1 2 ρ2σ2
We have fp(w1) > fp(w2) 6=⇒ σ1 > σ2 and σ1 > σ2 6=⇒ fp(w1) > fp(w2) because the relation between f(w1) and f(w2) is undetermined.

heading: A.2 PROOF OF LEMMA. 3.2
text: Since ρ is small, we can perform Taylor expansion around w,
h(w) = f(w + δ)− f(w) = δ>∇f(w) +O(ρ2) = ρ||∇f(w)||2 +O(ρ2) > 0 (13)
where the last line is because δ is approximated as δ = ρ ∇f(w)||∇f(w)||2+ , hence has the same direction as∇f(w).

heading: A.3 PROOF OF LEMMA. 3.3
text: Since ρ is small, we can approximate f(w) with a quadratic model around a local minima w:
f(w + δ) = f(w) + 1
2 δ>Hδ +O(ρ3)
where H is the Hessian at w, assumed to be positive semidefinite at local minima. Normalize δ such that ||δ||2 = ρ, Hence we have:
h(w) = fp(w)− f(w) = max||δ||2≤ρ f(w + δ)− f(w) = 1
2 σmaxρ
2 +O(ρ3) (14)
where σmax is the dominate eigenvalue of the hessian H , and first order term is 0 because the gradient is 0 at local minima. Therefore, we have σmax ≈ 2h(w)/ρ2.

heading: A.4 PROOF OF THM. 5.1
text: 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 our proof.
STEP 1: CONVERGENCE W.R.T FUNCTION fp(w)
For simplicity of notation, we denote the update at step t as
dt = −ηtg(t)p + ηtαg (t) ⊥ (15)
By L−smoothness of f and the definition of fp(wt) = f(wadvt ), and definition of dt = wt+1 − wt and wadvt = wt + δt we have
fp(wt+1) = f(w adv t+1) ≤ f(wadvt ) + 〈∇f(wadvt ), wadvt+1 − wadvt 〉+
L
2 ∣∣∣∣∣∣wadvt+1 − wadvt ∣∣∣∣∣∣2 (16) = f(wadvt ) + 〈∇f(wadvt ), wt+1 + δt+1 − wt − δt〉
+ L
2 ∣∣∣∣∣∣wt+1 + δt+1 − wt − δt∣∣∣∣∣∣2 (17) ≤ f(wadvt ) + 〈∇f(wadvt ), dt〉+ L
∣∣∣∣∣∣dt∣∣∣∣∣∣2 (18) + 〈∇f(wadvt ), δt+1 − δt〉+ L
∣∣∣∣∣∣δt+1 − δt∣∣∣∣∣∣2 (19) STEP 1.0: BOUND EQ. 18
We first bound Eq. 18. Take expectation conditioned on observation up to step t (for simplicity of notation, we use E short for Ex to denote expectation over all possible data points) conditioned on observations up to step t, also by definition of dt, we have
Efp(wt+1)− fp(wt) ≤ −ηt〈∇fp(wt),Eg(t)p 〉+ αηt〈∇fp(wt),Eg (t) ⊥ 〉 + Lη2tE ∣∣∣∣∣∣− g(t)p + αg(t)⊥ ∣∣∣∣∣∣2
2 (20) ≤ −ηtE ∣∣∣∣∣∣∇fp(wt)∣∣∣∣∣∣2
2 + 0 + (α+ 1)2G2η2t (21)(
Since Eg(t)⊥ is orthogonal to ∇fp(wt) by construction, ||g(t)|| ≤ G by assumption )
STEP 1.1: BOUND EQ. 19
By definition of δt, we have
δt = ρt g(t)
||g(t)||+ (22)
δt+1 = ρt+1 g(t+1)
||g(t+1)||+ (23)
where g(t) is the gradient of f at wt evaluated with a noisy data sample. When learning rate ηt is small, the update in weight dt is small, and expected gradient is
∇f(wt+1) = ∇f(wt + dt) = ∇f(wt) +Hdt +O(||dt||2) (24) where H is the Hessian at wt. Therefore, we have
E〈∇f(wadvt ), δt+1 − δt〉 = 〈∇f(wadvt ), ρtE g(t)
||g(t)||+ − ρt+1E
g(t+1)
||g(t+1)||+ 〉 (25)
≤ ||∇f(wadvt )||ρt ∣∣∣∣∣∣E g(t)||g(t)||+ − E g(t+1)||g(t+1)||+ ∣∣∣∣∣∣ (26)
≤ ||∇f(wadvt )||ρtφt (27) where the first inequality is due to (1) ρt is monotonically decreasing with t, and (2) triangle inequality that 〈a, b〉 ≤ ||a|| · ||b||. φt is the angle between the unit vector in the direction of ∇f(wt)
and∇f(wt+1). The second inequality comes from that (1) ∣∣∣∣∣∣ g||g||+ ∣∣∣∣∣∣ < 1 strictly, so we can replace δt in Eq. 25 with a unit vector in corresponding directions multiplied by ρt and get the upper bound, (2) the norm of difference in unit vectors can be upper bounded by the arc length on a unit circle.
When learning rate ηt and update stepsize dt is small, φt is also small. Using the limit that
tanx = x+O(x2), sinx = x+O(x2), x→ 0 We have:
tanφt = ||∇f(wt+1)−∇f(wt)||
||∇f(wt)|| +O(φ2t ) (28)
= ||Hdt +O(||dt||2)|| ||∇f(wt)|| +O(φ2t ) (29)
≤ ηtL(1 + α) (30) where the last inequality is due to (1) max eigenvalue of H is upper bounded by L because f is L−smooth, (2) ||dt|| = ||ηt(g‖ + αg⊥)|| and Egt = ∇f(wt).
Plug into Eq. 27, also note that the perturbation amplitude ρt is small so wt is close to wadvt , then we have E〈∇f(wadvt ), δt+1 − δt〉 ≤ L(1 + α)Gρtηt (31) Similarly, we have
E ∣∣∣∣∣∣δt+1 − δt∣∣∣∣∣∣2 ≤ ρ2tE∣∣∣∣∣∣ g(t)||g(t)||+ − g(t+1)||g(t+1)||+ ∣∣∣∣∣∣2 (32)
≤ ρ2tφ2t (33) ≤ ρ2tη2tL2(1 + α)2 (34)
STEP 1.2: TOTAL BOUND
Reuse results from Eq. 21 (replace Lp with 2L) and plug into Eq. 18, and plug Eq. 31 and Eq. 34 into Eq. 19, we have
Efp(wt+1)− fp(wt) ≤ −ηtE ∣∣∣∣∣∣∇fp(wt)∣∣∣∣∣∣2
2 +
2L(α+ 1)2
2 G2η2t
+ L(1 + α)Gρtηt + 2L3(1 + α)2
2 η2t ρ 2 t (35)
Perform telescope sum, we have
Efp(wT )− fp(w0) ≤ − T∑ t=1 ηtE||∇fp(wt)||2 + [ L(1 + α)2G2η20 + L(1 + α)Gρ0η0 ] T∑ t=1 1 t
+ L3(1 + α)2η20ρ 2 0 T∑ t=1 1 t2 (36)
Hence
ηT T∑ t=1 E||∇fp(wt)||2 ≤ T∑ t=1 ηtE||∇fp(wt)||2 ≤ fp(w0)− Efp(wT ) +D log T + π2E 6 (37)
where D = L(1 + α)2G2η20 + L(1 + α)Gρ0η0, E = L 3(1 + α)2η20ρ 2 0 (38) Note that ηT = η0√T , we have
1
T T∑ t=1 E||∇fp(wt)||2 ≤ fp(w0)− fmin + π2E/6 η0 1√ T + D η0 log T√ T
(39)
which implies that GSAM enables fp to converge at a rate of O(log T/ √ T ), and all the constants here are well-bounded.
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 t as
dt = −ηtg(t)p + αηtg (t) ⊥ (40)
By smoothness of f , we have
f(wt+1) ≤ f(wt) + 〈∇f(wt), dt〉+ L
2 ∣∣∣∣∣∣dt∣∣∣∣∣∣2 2
(41)
= f(wt) + 〈∇f(wt),−ηtg(t)p + αηtg (t) ⊥ 〉+
L
2 ∣∣∣∣∣∣dt∣∣∣∣∣∣2 2
(42)
For simplicity, we introduce a scalar βt such that ∇f‖(wt) = βt∇fp(wt) (43)
where∇f‖(wt) is the projection of∇f(wt) onto∇fp(wt). When perturbation amplitude ρ is small, we expect βt to be very close to 1.
Take expectation conditioned on observations up to step t for both sides of Eq. 42, we have:
Ef(wt+1) ≤ f(wt) + 〈 ∇f(wt),−
ηt βt
( ∇f(wt)−∇f⊥(wt) ) + αηtEg(t)⊥ 〉 + L 2 E ∣∣∣∣∣∣dt∣∣∣∣∣∣2 2 (44)
= f(wt)− ηt βt ∣∣∣∣∣∣∇f(wt)∣∣∣∣∣∣2 2 + ( 1 βt + α ) ηt 〈 ∇f(wt),∇f⊥(wt) 〉 + L 2 E ∣∣∣∣∣∣dt∣∣∣∣∣∣2 2 (45)
= f(wt)− ηt βt ∣∣∣∣∣∣∇f(wt)∣∣∣∣∣∣2 2 + ( 1 βt + α ) ηt 〈 ∇f(wt),∇f(wt) sin θt 〉 + L 2 E ∣∣∣∣∣∣dt∣∣∣∣∣∣2 2
(46)( θt is the angle between∇fp(wt) and ∇f(wt) ) = f(wt)−
ηt βt ∣∣∣∣∣∣∇f(wt)∣∣∣∣∣∣2 2 + ( 1 βt + α ) ηt ∣∣∣∣∣∣∇f(wt)∣∣∣∣∣∣2 2 (| tan θt|+O(θ2t )) + L 2 E ∣∣∣∣∣∣dt∣∣∣∣∣∣2 2
(47)( sinx = x+O(x2), tanx = x+O(x2) when x→ 0. ) Also note when perturbation amplitude ρt is small, we have
∇fp(wt) = ∇f(wt + δt) = ∇f(wt) + ρt
||∇f(wt)||2 + H(wt)∇f(wt) +O(ρ2t ) (48)
where δt = ρt ∇f(wt) ||∇f(wt)||2 by definition, H(wt) is the Hessian. Hence we have
| tan θt| ≤ ||∇fp(wt)−∇f(wt)|| ||∇f(wt)|| ≤ ρtL ||∇f(wt)||
(49)
where L 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 into Eq. 47, we have
Ef(wt+1) ≤ f(wt)− ηt βt ∣∣∣∣∣∣∇f(wt)∣∣∣∣∣∣2 2 + ( 1 βt + α ) ηt ∣∣∣∣∣∣∇f(wt)∣∣∣∣∣∣2 2 | tan θt|+ L 2 E ∣∣∣∣∣∣dt∣∣∣∣∣∣2 2 (50)
≤ f(wt)− ηt βt ∣∣∣∣∣∣∇f(wt)∣∣∣∣∣∣2 2 + ( 1 βt + α ) Lρtηt ∣∣∣∣∣∣∇f(wt)∣∣∣∣∣∣ 2 + L 2 E ∣∣∣∣∣∣dt∣∣∣∣∣∣2 2 (51)
≤ f(wt)− ηt βt ∣∣∣∣∣∣∇f(wt)∣∣∣∣∣∣2 2 + ( 1 βt + α ) LρtηtG+ L 2 E ∣∣∣∣∣∣dt∣∣∣∣∣∣2
2 (52)( Assume gradient has bounded norm G. ) (53)
≤ f(wt)− ηt
βmax
∣∣∣∣∣∣∇f(wt)∣∣∣∣∣∣2 2 + ( 1 βmin + α ) LρtηtG+ L
2 E(α+ 1)2G2η2t (54)(
βt is close to 1 assuming ρ is small, hence it’s natural to assume 0 < βmin ≤ βt ≤ βmax )
Re-arranging above formula, we have ηt
βmax
∣∣∣∣∣∣∇f(wt)∣∣∣∣∣∣2 2 ≤ f(wt)− Ef(wt+1) + ( 1 βmin + α ) LGηtρt + L 2 (α+ 1)2G2η2t (55)
perform telescope sum and taking expectations on each step, we have
1
βmax T∑ t=1 ηt ∣∣∣∣∣∣∇f(wt)∣∣∣∣∣∣2 2 ≤ f(w0)− Ef(wT ) + ( 1 βmin + α ) LG T∑ t=1 ηtρt + L 2 (α+ 1)2G2 T∑ t=1 η2t (56) Take the schedule to be ηt = η0√t and ρt = ρ0√ t , then we have
η0 βmax 1√ T T∑ t=1 ∣∣∣∣∣∣∇f(wt)∣∣∣∣∣∣2 2 ≤ LHS (57)
≤ RHS (58)
≤ f(w0)− fmin + ( 1 βmin + α ) LGη0ρ0 T∑ t=1 1 t + L 2 (α+ 1)2G2η20 T∑ t=1 1 t
(59) ≤ f(w0)− fmin + ( 1 βmin + α ) LGη0ρ0(1 + log T )
+ L
2 (α+ 1)2G2η20(1 + log T ) (60)
Hence 1
T T∑ t=1 ∣∣∣∣∣∣∇f(wt)∣∣∣∣∣∣2 2 ≤ C3√ T + C4 log T√ T
(61)
where C1, C4 are some constants. This implies the convergence rate w.r.t f(w) is O(log T/ √ T ).
STEP 3: CONVERGENCE W.R.T. SURROGATE GAP h(w)
Note that we have proved convergence for fp(w) in step 1, and convergence for f(w) in step 3. Also note that ∣∣∣∣∣∣∇h(wt)∣∣∣∣∣∣2 2 = ∣∣∣∣∣∣∇fp(wt)−∇f(wt)∣∣∣∣∣∣2 2 ≤ 2 ∣∣∣∣∣∣∇fp(wt)∣∣∣∣∣∣2 2 + 2 ∣∣∣∣∣∣∇f(wt)∣∣∣∣∣∣2 2 (62)
Hence 1
T T∑ t=1 ∣∣∣∣∣∣∇h(wt)∣∣∣∣∣∣2 2 ≤ 2 T T∑ t=1 ∣∣∣∣∣∣∇fp(wt)∣∣∣∣∣∣2 2 + 2 T T∑ t=1 ∣∣∣∣∣∣∇f(wt)∣∣∣∣∣∣2 2
(63)
also converges at rate O(log T/ √ T ) because each item in the RHS converges at rate O(log T √ T ).

heading: A.5 PROOF OF COROLLARY. 5.2.1
text: Using the results from Thm. 5.2, with probability at least 1− a, we have
Ew∼QExf(w, x) ≤ Ew∼Qf̂(w) + 4
√ KL(Q||P) + log 2ma
m (64)
Assume δ ∼ N (0, b2Ik) where k is the dimension of model parameters, hence δ2 (element-wise square) follows a a Chi-square distribution. By Lemma.1 in Laurent & Massart (2000), we have
P ( ||δ||22 − kb2 ≥ 2b2 √ kt+ 2tb2 ) ≤ exp(−t) (65)
hence with probability at least 1− 1/ √ n, we have
||δ||22 ≤ b2 ( 2 log √ n+ k + 2 √ k log √ n ) ≤ 2b2k ( 1 + √ log √ n
k
)2 ≤ ρ2 (66)
Therefore, with probability at least 1− 1/ √ n = 1− exp ( − ( ρ√ 2b − √ k )2)
Eδ f̂(w + δ) ≤ max||δ||2≤ρ f̂(w + δ) (67)
Combine Eq. 65 and Eq. 67, subtract the same constant C on both sides, and under the same assumption as in (Foret et al., 2020) that Ew∼QExf(w, x) ≤ Eδ∼N (0,b2Ik)Ew∼QExf(w + δ, x)we finish the proof.

heading: A.6 PROOF OF THM. 5.3
text: STEP 1: A SUFFICIENT CONDITION THAT THE LOSS GAP IS EXPECTED TO DECREASE FOR

heading: EACH STEP
text: Take Taylor expansion, then the expected change of loss gap caused by descent step is
E〈∇fp(wt)−∇f(wt),−ηt∇fp(wt)〉 (68)( where Eg⊥ = ∇f⊥(wt) ) = ηt [ − ∣∣∣∣∇fp(wt)∣∣∣∣22 + ∣∣∣∣∇fp(wt)∣∣∣∣2∣∣∣∣∇f(wt)∣∣∣∣2 cos θt ] (69)
where θt is the angle between vector ∇fp(wt) and ∇f(wt). The expected change of loss gap caused by ascent step is
E〈∇fp(wt)−∇f(wt), αηt∇f⊥(wt)〉 = −αηt ∣∣∣∣∇f⊥(wt)∣∣∣∣22 < 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 E〈∇h(wt), dt〉 ≤ 0 requires α to be large or | ∣∣∇f(wt)∣∣∣∣2 cos θt ≤ ∣∣∣∣∇fp(wt)∣∣∣∣. In practice, the perturbation amplitude ρ is small and we can
assume θt is close to 0 and ∣∣∣∣∇fp(wt)∣∣∣∣ is close to ∣∣∣∣∇f(wt)∣∣∣∣, we can also set the parameter α 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 ĥ caused by our ascent step. We refer to Eq. 69 and Eq. 70 to analyze the change in loss gap caused by the descent and ascent (orthogonally) respectively. It 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: ∆ĥt = −〈∇f̂p(wt)−∇f̂(wt), αηt∇f̂⊥(wt)〉 = αηt ∣∣∣∣∇f̂⊥(wt)∣∣∣∣22 (71)
= αηt ∣∣∣∣∇f̂(wt)∣∣∣∣22| tan θt|2 (72)
T∑ t=1 ∆ĥt ≤ T∑ t=1
αL2ηtρ 2 t (73)(
By Eq. 49 )
(74)
≤ T∑ t=1 αL2η0ρ 2 0 1 t3/2 (75)
≤ 2.7αL2η0ρ20 (76)
Hence we derive an upper bound for ∑T t=1 ∆ĥt.
Next we derive a lower bound for ∑T t=1 ∆ĥt Note that when ρt is small, by Taylor expansion
∇f̂p(wt) = ∇f̂(wt + δt) = ∇f̂(wt) + ρt
||∇f̂(wt)|| Ĥ(wt)∇f̂(wt) +O(ρ2t ) (77)
where Ĥ(wt) is the Hessian evaluated on training samples. Also when ρt is small, the angle θt between ∇f̂p(wt) and ∇f̂(wt) is small, by the limit that
tanx = x+O(x2), x→ 0 sinx = x+O(x2), x→ 0
We have | tan θt| = | sin θt|+O(θ2t ) = |θt|+O(θ2t ) Omitting high order term, we have
| tan θt| ≈ |θt| = ||∇f̂p(wt)−∇f̂(wt)|| ||f̂(wt)|| = ||ρtĤ(wt) +O(ρ2t )|| ||∇f̂(wt)|| ≥ ρt|σ|min G (78)
where G is the upper-bound on norm of gradient, |σ|min is the minimum absolute eigenvalue of the Hessian. The intuition is that as perturbation amplitude decreases, the angle θt decreases at a similar rate, though the scale constant might be different. Hence we have
T∑ t=1 ∆ĥt = T∑ t=1 αηt ∣∣∣∣∇f̂(wt)∣∣∣∣22| tan θt|2 +O(θ4t ) (79)
≥ T∑ t=1 αηtc 2 (ρt|σ|min G )2 (80)
= αc2ρ20η0|σ|2min
G2
T∑ t=1 1 t3/2 (81)
≥ αc 2ρ20η0|σ|2min
G2 (82)
where c2 is the lower bound of ||∇f̂ ||2 (e.g. due to noise in data and gradient observation). Results above indicate that the decrease in loss gap caused by the ascent step is non-trivial, hence our proposed method efficiently improves generalization compared with SAM.

heading: A.7 DISCUSSION ON COROLLARY 5.2.1
text: The comment “‘The corollary gives a bound on the risk in terms of the perturbed training loss if one removes C from both sides”’ is correct. But there is a misunderstanding in the statement “‘the perturbed training loss is small then the model has a small risk”’: it’s only true when ρtrain for training equals its real value ρtrue determined by the data distribution; in practice, we never know ρtrue. In the following we show that the minimization of both h and fp is better than simply minimizing fp when ρtrue 6= ρtrain. 1. First, we re-write the conclusion of Corollary 5.2.1 as
EwExf(w, x) ≤ fp +R = C + ĥ+R = C + ρ2σ/2 +R+O(ρ3)
with probability (1− a)[1− e−( ρ√ 2b − √ k)2 ]
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 termO(ρ3). We focus on
fp = C + ĥ = C + ρ 2σ/2
2. When ρtrue 6= ρtrain, minimizing h achieves a lower risk than only minimizing fp. (1) Note that after training, C (training loss) is fixed, but h could vary with ρ (e.g. when training on dataset A and testing on an 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 a low fp is insufficient for generalization, and a low σ is necessary:
A Suppose we use ρtrain for training, and consider two solutions with C1, σ1 (SAM) and C2, σ2 (GSAM). Suppose they have the same fp during training for some ρtrain, so
fp1 = C1 + σ1/2× ρ2train = C2 + σ2/2× ρ2train = fp2 Suppose C1 < C2 so σ1 > σ2.
B When ρtrue > ρtrain, we have
risk bound 1 = C1 +σ1/2× ρ2true +R > risk bound 2 = C2 +σ2/2× ρ2true +R
This implies that a small σ helps generalization, but only a low fp1 (caused by a low C1 and high σ1) is insufficient for a good generalization.
C Note that ρtrain is fixed during training, so minimizing htrain during training is equivalently minimizing σ by Lemma 3.3
3. Why we are often unlucky to have ρtrue > ρtrain (1) First, the test sets are almost surely outside the convex hull of the training set because “‘interpolation almost surely never occurs in high-dimensional (> 100) cases”’ Balestriero et al. (2021). 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 4 below), we have ρtrue > ρtrain set almost surely. (2) Second, we don’t know the value of ρtrue and can only guess it. In practice, we often guess a small value because training often diverges with large ρ (as observed in Foret et al. (2020); Chen et al. (2021)).
4. Why ρ increases with data variability. In Corollary 5.2.1, we assume weight perturbation δ ∼ N (0, b2Ik). The meaning of b is the following. If we can randomly sample a fixed number of samples from the underlying distribution, then training the model from scratch (with a fixed seed for random initialization) gives rise to a 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, the more variability in data, the more variability in weights, and the larger b. Note that Corollary stated that the bound holds with probability proportional to [1 − e−( ρ√ 2b − √ k)2
]. In order for the result to hold with a fixed probability, ρ must stay proportional to b, hence ρ also increases with the variability of data.

heading: B EXPERIMENTAL DETAILS

heading: B.1 TRAINING DETAILS
text: For ViT and Mixer, we search the learning rate in {1e-3, 3e-3, 1e-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.1}. For ViT and Mixer, we use the AdamW optimizer with β1 = 0.9, β2 = 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 et al., 2021) 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 2× longer training, and does not find significant improvement as 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, we search ρ in {0.1, 0.2, 0.3, 0.4, 0.5, 0.6} for SAM and GSAM; for ResNet, we search ρ from 0.01 to 0.05 with a stepsize 0.01. For ASAM, we amplify ρ by 10× compared to SAM, as recommended by Kwon et al. (2021). For GSAM, we search α in {0.1, 0.2, 0.3} throughout the paper. We report the best configuration of each individual model in Table. 4.

heading: B.2 TRANSFER LEARNING EXPERIMENTS
text: Using weights trained on ImageNet-1k, we finetune models with SGD on downstream tasks including the CIFAR10/CIFAR100 (Krizhevsky et al., 2009), Oxford-flowers (Nilsback & Zisserman, 2008) and Oxford-IITPets (Parkhi et al., 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 in {0.001, 0.003, 0.01, 0.03}. On Cifar datasets, we train models for 10k steps with a warmup step of 500; on Oxford datasets, we train models for 500 steps with a wamup step of 100.

heading: B.3 EXPERIMENTAL SETUP WITH ABLATION STUDIES ON DATA AUGMENTATION
text: We follow the settings in (Tolstikhin et al., 2021) to perform ablation studies on data augmentation. In the left subfigure of Fig. 6, “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 RandAug magnitude 10; “Strong” refers to the mixup augmentation with probability 0.2 and RandAug magnitude 15.

heading: C ABLATION STUDIES AND DISCUSSIONS
text: C.1 INFLUENCE OF ρ AND α
We plot the performance of a ViT-B/32 model varying with ρ (Fig. 7a) and α (Fig. 7b). We empirically validate that fine-tuning ρ in SAM can not achieve comparable performance with GSAM, as
shown in Fig. 7a. Considering that GSAM has one more parameter α, we plot the accuracy varying with α in Fig. 7b, and show that GSAM consistently outperforms SAM and vanilla training.
C.2 CONSTANT ρ V.S. DECAYED ρt SCHEDULE
Note that Thm. 5.1 assumes ρt to decay with t in order to prove the convergence, while SAM uses a constant ρ during training. To eliminate the influence of ρt schedule, we conduct ablation study as in Table. 6. The ascent step in GSAM can be applied to both constant ρ or a decayed ρt schedule, and improves accuracy for both cases. Without ascent step, constant ρ and decayed ρt achieve similar performance. Results in Table. 6 implies that the ascent step in GSAM is the main reason for improvement of generalization performance.
C.3 VISUALIZE THE TRAINING PROCESS
In the proof of Thm. 5.3, our analysis relies on assumption that θt is small. We empirically validated this assumption by plotting cos θt in Fig. 8, where θt is the angle between ∇f(wt) and ∇fp(wt). Note that the cosine value is calculated in the parameter space of dimension 8.8 × 107, and in high-dimensional space two random vectors are highly likely to be perpendicular. In Fig. 8 the cosine value is always above 0.9, indicating that∇f(wt) and∇fp(wt) point to very close directions considering the high dimension of parameters. This empirically validates our assumption that θt is small during training.
We also plot the surrogate gap during training in Fig. 9. As α 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 α, indicating that the training process gradually falls into local minimum in order to minimize the training loss.

heading: D RELATED WORKS
text: Besides SAM and ASAM, other methods were proposed in the literature to improve generalization: Lin et al. (2020) proposed extrapolation of gradient, Xie et al. (2021) proposed to manipulate the noise in gradient, and Damian et al. (2021) 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, and Heo et al. (2020) restricted the norm of updated weights to improve generalization. Many of aforementioned methods can be combined with GSAM to further improve generalization.
Besides modified training schemes, there are other two types of techniques to improve generalization: data augmentation and model regularization. Data augmentation typically generates new data from training samples; besides standard data augmentation such as flipping or rotation of images, recent data augmentations include label smoothing (Müller et al., 2019) and mixup (Müller et al., 2019) which trains on convex combinations of both inputs and labels, automatically learned augmentation (Cubuk et al., 2018), and cutout (DeVries & 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-propagation (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 optimizers, such as AdaBound (Luo et al., 2019), RAdam (Liu et al., 2019), Yogi (Zaheer et al., 2018), AdaGrad (Duchi et al., 2011), AMSGrad (Reddi et al., 2019) and AdaDelta (Zeiler, 2012).

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

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.

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: 1. A new loss function which does not suffer from these problems, along with a new algorithm derived from the loss function 2. Theoretical results (generalization bounds and convergence theorems) supporting the algorithm and new loss function 3. An extensive numerical evaluation, showing that the proposed GSAM method leads to significant gains over SAM