# Surrogate Gap Minimization
Improves Sharpness-Aware Training

Juntang Zhuang1
<EMAIL_ADDRESS>
&Boqing Gong2, Liangzhe Yuan2, Yin Cui2, Hartwig Adam2
{bgong, lzyuan, yincui<EMAIL_ADDRESS>&Nicha C. Dvornek1, Sekhar
Tatikonda1, James S. Duncan1
{nicha.dvornek, sekhar.tatikonda<EMAIL_ADDRESS>
&Ting Liu2
<EMAIL_ADDRESS>1 Yale University, 2 Google Research Work was done during an
internship at Google

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

##  1 Introduction

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 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 the 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 and 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., 2019) derived the gradient of local entropy to
avoid solutions in sharp valleys. Entropy-SGD has 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 a perturbed loss
defined as
$f_{p}(w)\triangleq\operatorname{max}_{||\delta||\leq\rho}f(w+\delta)$, which
is the maximum loss within radius $\rho$ centered at 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 a low perturbed loss $f_{p}$ could appear in both flat and sharp
minima, implying that only minimizing $f_{p}$ is not always sharpness-aware.

Although the perturbed loss $f_{p}(w)$ might disagree with sharpness, we find
a surrogate gap defined as $h(w)\triangleq f_{p}(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 $f_{p}$ and the
surrogate gap $h$: a low perturbed loss $f_{p}$ 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 performances.

GSAM consists of two steps for each update: 1) descend gradient $\nabla
f_{p}(w)$ to minimize the perturbed loss $f_{p}$ (this step is exactly the
same as SAM), and 2) decompose gradient $\nabla f(w)$ of the original loss
$f(w)$ into components that are parallel and orthogonal to $\nabla f_{p}(w)$,
i.e., $\nabla f(w)=\nabla_{\parallel}f(w)+\nabla_{\perp}f(w)$, and perform an
ascent step in $\nabla_{\perp}f(w)$ to minimize the surrogate gap $h(w)$. Note
that this ascent step does not change the perturbed loss $f_{p}$ because
$\nabla f_{\perp}(w)\perp\nabla f_{p}(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 the 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 architectures, including ResNets (He et al., 2016), Vision Transformers
(Dosovitskiy et al., 2020), and MLP-Mixers (Tolstikhin et al., 2021).

##  2 Preliminaries

###  2.1 Notations

  * •

$f(w)$: A loss function $f$ with parameter $w\in\mathbb{R}^{k}$, where $k$ is
the parameter dimension.

  * •

$\rho_{t}\in\mathbb{R}$: A scalar value controlling the amplitude of
perturbation at step $t$.

  * •

$\epsilon\in\mathbb{R}$: A small positive constant (to avoid division by 0,
$\epsilon=10^{-12}$ by default).

  * •

$w_{t}^{adv}\triangleq w_{t}+\rho_{t}\frac{\nabla f(w_{t})}{||\nabla
f(w_{t})||+\epsilon}$: The solution to
$\operatorname{max}_{||w^{\prime}-w_{t}||\leq\rho_{t}}f(w^{\prime})$ when
$\rho_{t}$ is small.

  * •

$f_{p}(w_{t})\triangleq\operatorname{max}_{||\delta||\leq\rho_{t}}f(w_{t}+\delta)\approx
f(w_{t}^{adv})$: The perturbed loss induced by $f(w_{t})$. For each $w_{t}$,
$f_{p}(w_{t})$ returns the worst possible loss $f$ within a ball of radius
$\rho_{t}$ centered at $w_{t}$. When $\rho_{t}$ is small, by Taylor expansion,
the solution to the maximization problem is equivalent to a gradient ascent
from $w_{t}$ to $w_{t}^{adv}$.

  * •

$h(w)\triangleq f_{p}(w)-f(w)$: The surrogate gap defined as the difference
between $f_{p}(w)$ and $f(w)$.

  * •

$\eta_{t}\in\mathbb{R}$: Learning rate at step $t$.

  * •

$\alpha\in\mathbb{R}$: A constant value that controls the scaled learning rate
of the ascent step in GSAM.

  * •

$g^{(t)},g_{p}^{(t)}\in\mathbb{R}^{k}$: At the $t$-th step, the noisy
observation of the gradients $\nabla f(w_{t})$, $\nabla f_{p}(w_{t})$ of the
original loss and perturbed loss, respectively.

  * •

$\nabla f(w_{t})=\nabla f_{\parallel}(w_{t})+\nabla f_{\perp}(w_{t})$:
Decompose $\nabla f(w_{t})$ into parallel component $\nabla
f_{\parallel}(w_{t})$ and vertical component $\nabla f_{\perp}(w_{t})$ by
projection $\nabla f(w_{t})$ onto $\nabla f_{p}(w_{t})$.

###  2.2 Sharpness-Aware Minimization

Conventional optimization of neural networks typically minimizes the training
loss $f(w)$ by gradient descent w.r.t. $\nabla f(w)$ and searches for a single
point $w$ 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:

$\operatorname{min}_{w}f_{p}(w)\text{ where
}f_{p}(w)\triangleq\operatorname{max}_{||\delta||\leq\rho}f(w+\delta)$ (1)

where $\rho$ is a predefined constant controlling the radius of a
neighborhood. This perturbed loss $f_{p}$ 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
$\rho$, using Taylor expansion around $w$, the inner maximization in Eq. 1
turns into a linear constrained optimization with solution

$\arg\operatorname{max}_{||\delta||\leq\rho}f(w+\delta)=\arg\operatorname{max}_{||\delta||\leq\rho}f(w)+\delta^{\top}\nabla
f(w)+O(\rho^{2})=\rho\frac{\nabla f(w)}{||\nabla f(w)||}$ (2)

As a result, the optimization problem of SAM reduces to

$\operatorname{min}_{w}f_{p}(w)\approx\operatorname{min}_{w}f(w^{adv})\text{
where }w^{adv}\triangleq w+\rho\frac{\nabla f(w)}{||\nabla f(w)||+\epsilon}$
(3)

where $\epsilon$ is a scalar (default: 1e-12) to avoid division by 0, and
$w^{adv}$ is the “perturbed weight” with the highest loss within the
neighborhood. Equivalently, SAM seeks a solution on the surface of the
perturbed loss $f_{p}(w)$ rather than the original loss $f(w)$ (Foret et al.,
2020).

##  3 The surrogate gap measures the sharpness at a local minimum

Figure 1: Consider original loss $f$ (solid line), perturbed loss
$f_{p}\triangleq\operatorname{max}_{||\delta||\leq\rho}f(w+\delta)$ (dashed
line), and surrogate gap $h(w)\triangleq f_{p}(w)-f(w)$. Intuitively, $f_{p}$
is approximately a max-pooled version of $f$ with a pooling kernel of width
$2\rho$, and SAM minimizes $f_{p}$. From left to right are the local minima
centered at $w_{1},w_{2},w_{3}$, and the valleys become flatter. Since
$f_{p}(w_{1})=f_{p}(w_{3})<f_{p}(w_{2})$, SAM prefers $w_{1}$ and $w_{3}$ to
$w_{2}$. However, a low $f_{p}$ could appear in both sharp ($w_{1}$) and flat
($w_{3}$) minima, so $f_{p}$ might disagree with sharpness. On the contrary, a
smaller surrogate gap $h$ indicates a flatter loss surface (Lemma 3.3). From
$w_{1}$ to $w_{3}$, the loss surface is flatter, and $h$ is smaller.

###  3.1 The perturbed loss is not always sharpness-aware

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
$\sigma_{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 $\rho$, consider two local minima $w_{1}$ and $w_{2}$,
$f_{p}(w_{1})\leq
f_{p}(w_{2})\centernot\implies\sigma_{max}(w_{1})\leq\sigma_{max}(w_{2})$,
where $\sigma_{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 $w_{1}$ to $w_{3}$, and suppose the
corresponding loss surfaces are flatter from $w_{1}$ to $w_{3}$. For some
fixed $\rho$, we plot the perturbed loss $f_{p}$ and surrogate gap
$h\triangleq f_{p}-f$ around each solution. Comparing $w_{2}$ with $w_{3}$:
Suppose their vanilla losses are equal, $f(w_{2})=f(w_{3})$, then
$f_{p}(w_{2})>f_{p}(w_{3})$ because the loss surface is flatter around
$w_{3}$, implying that SAM will prefer $w_{3}$ to $w_{2}$. Comparing $w_{1}$
and $w_{2}$: $f_{p}(w_{1})<f_{p}(w_{2})$, and SAM will favor $w_{1}$ over
$w_{2}$ because it only cares about the perturbed loss $f_{p}$, even though
the loss surface is sharper around $w_{1}$ than $w_{2}$.

###  3.2 The surrogate gap agrees with sharpness

We introduce the surrogate gap that agrees with sharpness, defined as:

$h(w)\triangleq\operatorname{max}_{||\delta||\leq\rho}f(w+\delta)-f(w)\approx
f(w^{adv})-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 $\rho$ is sufficiently small, then the
approximation to the surrogate gap in Eq. 4 is always non-negative,
$h(w)\approx f(w^{adv})-f(w)\geq 0,\forall w$.

######  Lemma 3.3.

For a local minimum $w^{*}$, consider the dominate eigenvalue $\sigma_{max}$
of the Hessian of loss $f$ as a measure of sharpness. Considering the
neighborhood centered at $w^{*}$ with a small radius $\rho$, the surrogate gap
$h(w^{*})$ is an equivalent measure of the sharpness: $\sigma_{max}\approx
2h(w^{*})/\rho^{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)$.

Figure 2:  $\nabla f$ is decomposed into parallel and vertical ($\nabla
f_{\perp}$) components by projection onto  $\nabla f_{p}$.
${\color[rgb]{1,0,0}\nabla f^{GSAM}}={\color[rgb]{0,0.546875,0.1953125}\nabla
f_{p}}-\alpha{\color[rgb]{0,0,1}\nabla f_{\perp}}$

For  $t=1$ to $T$ 0) $\rho_{t}$ schedule:
$\rho_{t}=\rho_{min}+\frac{(\rho_{max}-\rho_{min})(lr-
lr_{min})}{lr_{max}-lr_{min}}$ 1a) $\Delta
w_{t}={\color[rgb]{0,0,0}\rho_{t}}\frac{\color[rgb]{1,.5,0}\nabla
f^{(t)}}{||{\color[rgb]{1,.5,0}\nabla f^{(t)}}||+\epsilon}$ 1b)
${\color[rgb]{0,0.546875,0.1953125}w_{t}^{adv}}=w_{t}+\Delta w_{t}$ 2) Get
$\color[rgb]{0,0.546875,0.1953125}\nabla f_{p}^{(t)}$ by back-propagation at
$\color[rgb]{0,0.546875,0.1953125}w_{t}^{adv}$. 3) ${\color[rgb]{1,.5,0}\nabla
f^{(t)}}=\nabla f_{\parallel}^{(t)}+\color[rgb]{0,0,1}\nabla f_{\perp}^{(t)}$
Decompose $\color[rgb]{1,.5,0}\nabla f^{(t)}$ into components that are
parallel and orthogonal to $\color[rgb]{0,0.546875,0.1953125}\nabla
f_{p}^{(t)}$. 4) Update weights: Vanilla
$w_{t+1}=w_{t}-\eta_{t}\color[rgb]{1,.5,0}\nabla f^{(t)}$ SAM
$w_{t+1}=w_{t}-\eta_{t}{\color[rgb]{0,0.546875,0.1953125}\nabla f_{p}^{(t)}}$
GSAM $w_{t+1}=w_{t}-\eta_{t}({\color[rgb]{0,0.546875,0.1953125}\nabla
f_{p}^{(t)}}-{\color[rgb]{0,0,1}\alpha\nabla f_{\perp}^{(t)}})$

Algorithm 1 GSAM Algorithm

##  4 Surrogate Gap Guided Sharpness-Aware 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-Aware Minimzation (GSAM) to simultaneously minimize two objectives,
the perturbed loss $f_{p}$ and the surrogate gap $h$:

$\operatorname{min}_{w}\big{(}f_{p}(w),h(w)\big{)}$ (5)

Intuitively, by minimizng $f_{p}$ 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
$\operatorname{min}_{w}f_{p}(w)+\lambda h(w)$, where $\lambda$ is some
positive scalar. One caveat when solving this weighted combination is the
potential conflict between the gradients of the two terms, i.e., $\nabla
f_{p}(w)$ and $\nabla h(w)$. We illustrate this conflict by Fig. 2, where
$\nabla h(w)=\nabla f_{p}(w)-\nabla f(w)$ (the grey dashed arrow) has a
negative inner product with $\nabla f_{p}(w)$ and $\nabla f(w)$. Hence, the
gradient descent for the surrogate gap could potentially increase the loss
$f_{p}$, harming the model’s performance. We empirically validate this
argument in Sec. 6.4.

###  4.2 Gradient decomposition and ascent for the multi-objective
optimization

Our primary goal is to minimize $f_{p}$ because otherwise a flat solution of
high loss is meaningless, and the minimization of $h$ should not increase
$f_{p}$. We propose to decompose $\nabla f(w_{t})$ and $\nabla h$ into
components that are parallel and orthogonal to $\nabla f_{p}(w_{t})$,
respectively (see Fig. 2):

$\displaystyle\nabla f(w_{t})$ $\displaystyle=\nabla
f_{\parallel}(w_{t})+\nabla f_{\perp}(w_{t})$ $\displaystyle\nabla h(w_{t})$
$\displaystyle=\nabla h_{\parallel}(w_{t})+\nabla h_{\perp}(w_{t})$ (6)
$\displaystyle\nabla h_{\perp}(w_{t})$ $\displaystyle=-\nabla
f_{\perp}(w_{t})$

The key is that updating in the direction of $\nabla h_{\perp}(w_{t})$ does
not change the value of the perturbed loss $f_{p}(w_{t})$ because $\nabla
h_{\perp}\perp\nabla f_{p}$ by construction. Therefore, we propose to perform
a descent step in the $\nabla h_{\perp}(w_{t})$ direction, which is equivalent
to an ascent step in the $\nabla f_{\perp}(w_{t})$ direction (because $\nabla
h_{\perp}=-\nabla f_{\perp}$ by the definition of $h$), and achieve two goals
simultaneously — it keeps the value of $f_{p}(w_{t})$ intact and meanwhile
decreases the surrogate gap $h(w_{t})=f_{p}(w_{t})-f(w_{t})$ (by increasing
$f(w_{t})$ and not affect $f_{p}(w_{t})$).

The full GSAM Algorithm is shown in Algo. 1 and Fig. 2, where
$g^{(t)},g_{p}^{(t)}$ are noisy observations of $\nabla f(w_{t})$ and $\nabla
f_{p}(w_{t})$, respectively, and $g_{\parallel}^{(t)},g_{\perp}^{(t)}$ are
noisy observations of $\nabla f_{\parallel}(w_{t})$ and $\nabla
f_{\perp}(w_{t})$, respectively, by projecting $g^{(t)}$ onto $g_{p}^{(t)}$.
We introduce a constant $\alpha$ to scale the stepsize of the ascent step.
Steps 1) to 2) are the same as SAM: At current point $w_{t}$, step 1) takes a
gradient ascent to $w_{t}^{adv}$ followed by step 2) evaluating the gradient
$g_{p}^{(t)}$ at $w_{t}^{adv}$. Step 3) projects $g^{(t)}$ onto $g_{p}^{(t)}$,
which requires negligible computation compared to the forward and backward
passes. In step 4), $-\eta_{t}g_{p}^{(t)}$ is the same as in SAM and minimizes
the perturbed loss $f_{p}(w_{t})$ with gradient descent, and
$\alpha\eta_{t}g_{\perp}^{(t)}$ performs an ascent step in the orthogonal
direction of $g_{p}^{(t)}$ to minimize the surrogate gap $h(w_{t})$ (
equivalently increase $f(w_{t})$ and keep $f_{p}(w_{t})$ intact). In coding,
GSAM feeds the “surrogate gradient” $\nabla f_{t}^{GSAM}\triangleq
g_{p}^{(t)}-\alpha g_{\perp}^{(t)}$ to first-order gradient optimizers such as
SGD and Adam.

The ascent step along $g_{\perp}^{(t)}$ does not harm convergence SAM
demonstrates that minimizing $f_{p}$ makes the network generalize better than
minimizing $f$. Even though our ascent step along $g_{\perp}^{(t)}$ increases
$f(w)$, it does not affect $f_{p}(w)$, so GSAM still decreases the perturbed
loss $f_{p}$ 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, $\nabla f(w_{t})$ points
to a sharp local solution and deviates from the ridge; in contrast,
$w_{t}^{adv}$ is closer to the ridge and $\nabla f(w_{t}^{adv})$ is closer to
the ridge descent direction than $\nabla f(w_{t})$. Note that $\nabla
f_{t}^{GSAM}$ and $\nabla f(w_{t})$ always lie at different sides of $\nabla
f_{p}(w_{t})$ by construction (see Fig. 2), hence $\nabla f_{t}^{GSAM}$ pushes
the trajectory closer to the ridge than $\nabla f_{p}(w_{t})$ does. The
trajectory of GSAM is like descent along the ridge and tends to find flat
minima.

##  5 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 bound $f_{min}$. Suppose we can access a noisy, bounded observation
$g^{(t)}$ ($||g^{(t)}||_{2}\leq G,\forall t$) of the true gradient $\nabla
f(w_{t})$ at the $t$-th step. For some constant $\alpha$, with learning rate
$\eta_{t}=\eta_{0}/\sqrt{t}$, and perturbation amplitude $\rho_{t}$
proportional to the learning rate, e.g., $\rho_{t}=\rho_{0}/\sqrt{t}$, we have

$\displaystyle\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}\Big{|}\Big{|}\nabla
f_{p}(w_{t})\Big{|}\Big{|}_{2}^{2}\leq\frac{C_{1}+C_{2}\log
T}{\sqrt{T}},\hskip
11.38109pt\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}\Big{|}\Big{|}\nabla
f(w_{t})\Big{|}\Big{|}_{2}^{2}\leq\frac{C_{3}+C_{4}\log T}{\sqrt{T}}$

where $C_{1},C_{2},C_{3},C_{4}$ are some constants.

Thm. 5.1 implies both $f_{p}$ and $f$ converge in GSAM at rate $O(\log
T/\sqrt{T})$ for non-convex stochastic optimization, matching the convergence
rate of first-order gradient optimizers like Adam.

Figure 3: Consider the loss surface with a few sharp local minima. Left:
Overview of the procedures of SGD, SAM and GSAM. SGD takes a descent step at
$w_{t}$ using $\nabla f(w_{t})$ (orange), which points to a sharp local
minima. SAM first performs gradient ascent in the direction of $\nabla
f(w_{t})$ to reach $w_{t}^{adv}$ with a higher loss, followed by descent with
gradient $\nabla f(w_{t}^{adv})$ (green) at the perturbed weight. Based on
$\nabla f(w_{t})$ and $\nabla f(w_{t}^{adv})$, 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. A video is in the supplement for better visualization.

###  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 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
$\widehat{f}(w)=\frac{1}{m}\sum_{i=1}^{m}f(w,x_{i}),$ where we use $x_{i}$ 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
$\mathcal{Q}$. Denote the prior distribution (independent of training) as
$\mathcal{P}$, then

$\mathbb{E}_{w\sim\mathcal{Q}}\mathbb{E}_{x}f(w,x)\leq\mathbb{E}_{w\sim\mathcal{Q}}\widehat{f}(w)+4\sqrt{\Big{(}KL(\mathcal{Q}||\mathcal{P})+\log\frac{2m}{a}\Big{)}/m}\text{
with probability at least }1-a$

######  Corollary 5.2.1.

Suppose perturbation $\delta$ is drawn from distribution
$\delta\sim\mathcal{N}(0,b^{2}I^{k}),\delta\in\mathbb{R}^{k}$, $k$ is the
dimension of $w$, then with probability at least
$\Big{(}1-a\Big{)}\Big{[}1-e^{-\big{(}\frac{\rho}{\sqrt{2}b}-\sqrt{k}\big{)}^{2}}\Big{]}$

$\displaystyle\mathbb{E}_{w\sim\mathcal{Q}}\mathbb{E}_{x}f(w,x)\leq\widehat{h}+C+4\sqrt{\Big{(}KL(\mathcal{Q}||\mathcal{P})+\log\frac{2m}{a}\Big{)}/m}$
(7)
$\displaystyle\widehat{h}\triangleq\operatorname{max}_{||\delta||_{2}\leq\rho}\widehat{f}(w+\delta)-\widehat{f}(w)=\frac{1}{m}\sum_{i=1}^{m}\Big{[}\operatorname{max}_{||\delta||_{2}\leq\rho}f(w+\delta,x_{i})-f(w,x_{i})\Big{]}$
(8)

where $C=\widehat{f}(w)$ is the empirical training loss, and $\widehat{h}$ is
the surrogate gap evaluated on the training set.

Corollary 5.2.1 implies that minimizing $\widehat{h}$ (right hand side of Eq.
7) is expected to achieve a tighter upper bound of the generalization
performance (left hand side of Eq. 7). The third term on the right of Eq. 7 is
typically hard to analyze and often simplified to $L2$ regularization (Foret
et al., 2020). Note that $f_{p}=C+\widehat{h}$ only holds when $\rho_{train}$
(the perturbation amplitude specified by users during training) equals
$\rho_{true}$ (the ground truth value determined by underlying data
distribution); when $\rho_{train}\neq\rho_{true}$, $min(f_{p},\widehat{h})$ is
more effective than $min(f_{p})$ 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.1, Thm. 5.2 and Corollary 5.2.1, we assume the
Hessian has a lower-bound $|\sigma|_{min}$ on the absolute value of
eigenvalue, and the variance of noisy observation $g^{(t)}$ is lower-bounded
by $c^{2}$. The surrogate gap $h$ can be minimized by the ascent step along
the orthogonal direction $g_{\perp}^{(t)}$. During training we minimize the
sample estimate of $h$. We use $\Delta\widehat{h}_{t}$ to denote the amount
that the ascent step in GSAM decreases $\widehat{h}$ for the $t$-th step.
Compared to SAM, the proposed method generates a total decrease in surrogate
gap $\sum_{t=1}^{T}\Delta\widehat{h}_{t}$, which is bounded by

$\frac{\alpha
c^{2}\rho_{0}^{2}\eta_{0}|\sigma|_{min}^{2}}{G^{2}}\leq\lim_{T\to\infty}\sum_{t=1}^{T}\Delta\widehat{h}_{t}\leq
2.7\alpha L^{2}\eta_{0}\rho_{0}^{2}$ (9)

We provide proof in the appendix. The lower-bound of
$\sum_{t=1}^{T}\Delta\widehat{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.

##  6 Experiments

Table 1: 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-v1 | ImageNet-Real | ImageNet-V2 | ImageNet-R | ImageNet-C
---|---|---|---|---|---|---
ResNet
ResNet50 | Vanilla (SGD) | 76.0 | 82.4 | 63.6 | 22.2 | 44.6
SAM | 76.9 | 83.3 | 64.4 | 23.8 | 46.5
GSAM | 77.2 | 83.9 | 64.6 | 23.6 | 47.6
ResNet101 | Vanilla (SGD) | 77.8 | 83.9 | 65.3 | 24.4 | 48.5
SAM | 78.6 | 84.8 | 66.7 | 25.9 | 51.3
GSAM | 78.9 | 85.2 | 67.3 | 26.3 | 51.8
ResNet152 | Vanilla (SGD) | 78.5 | 84.2 | 66.3 | 25.3 | 50.0
SAM | 79.3 | 84.9 | 67.3 | 25.7 | 52.2
GSAM | 80.0 | 85.9 | 68.6 | 27.3 | 54.1
Vision Transformer
ViT-S/32 | Vanilla (AdamW) | 68.4 | 75.2 | 54.3 | 19.0 | 43.3
SAM | 70.5 | 77.5 | 56.9 | 21.4 | 46.2
GSAM | 73.8 | 80.4 | 60.4 | 22.5 | 48.2
ViT-S/16 | Vanilla (AdamW) | 74.4 | 80.4 | 61.7 | 20.0 | 46.5
SAM | 78.1 | 84.1 | 65.6 | 24.7 | 53.0
GSAM | 79.5 | 85.3 | 67.3 | 25.3 | 53.3
ViT-B/32 | Vanilla (AdamW) | 71.4 | 77.5 | 57.5 | 23.4 | 44.0
SAM | 73.6 | 80.3 | 60.0 | 24.0 | 50.7
GSAM | 76.8 | 82.7 | 63.0 | 25.1 | 51.7
ViT-B/16 | Vanilla (AdamW) | 74.6 | 79.8 | 61.3 | 20.1 | 46.6
SAM | 79.9 | 85.2 | 67.5 | 26.4 | 56.5
GSAM | 81.0 | 86.5 | 69.2 | 27.1 | 55.7
MLP-Mixer
Mixer-S/32 | Vanilla (AdamW) | 63.9 | 70.3 | 49.5 | 16.9 | 35.2
SAM | 66.7 | 73.8 | 52.4 | 18.6 | 39.3
GSAM | 68.6 | 75.8 | 55.0 | 22.6 | 44.6
Mixer-S/16 | Vanilla (AdamW) | 68.8 | 75.1 | 54.8 | 15.9 | 35.6
SAM | 72.9 | 79.8 | 58.9 | 20.1 | 42.0
GSAM | 75.0 | 81.7 | 61.9 | 23.7 | 48.5
Mixer-S/8 | Vanilla (AdamW) | 70.2 | 76.2 | 56.1 | 15.4 | 34.6
SAM | 75.9 | 82.5 | 62.3 | 20.5 | 42.4
GSAM | 76.8 | 83.4 | 64.0 | 24.6 | 47.8
Mixer-B/32 | Vanilla (AdamW) | 62.5 | 68.1 | 47.6 | 14.6 | 33.8
SAM | 72.4 | 79.0 | 58.0 | 22.8 | 46.2
GSAM | 73.6 | 80.2 | 59.9 | 27.9 | 52.1
Mixer-B/16 | Vanilla (AdamW) | 66.4 | 72.1 | 50.8 | 14.5 | 33.8
SAM | 77.4 | 83.5 | 63.9 | 24.7 | 48.8
GSAM | 77.8 | 84.0 | 64.9 | 28.3 | 54.4

###  6.1 GSAM improves test performance on various model architectures

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 $\rho$ from 0.01 to 0.05 with a stepsize 0.01. For
ViTs and Mixers, we search for $\rho$ from 0.05 to 0.6 with a stepsize 0.05.
In GSAM, we search for $\alpha$ in $\\{0.01,0.02,0.03\\}$ for ResNets and
$\alpha$ 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.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 out-of-distribution data (ImageNet-R and ImageNet-C), and the
observation is consistent with that on ImageNet, e.g., +5.1% on ImageNet-R and
+5.9% on ImageNet-C for Mixer-B/32.

Figure 4: Influence of $\rho$ (set as constant for ease of comparison, other
experiments use decayed $\rho_{t}$ schedule) and $\alpha$ on the training of
ViT-B/32. Left: Top-1 accuracy on ImageNet. Middle: Estimation of the dominant
eigenvalues from the surrogate gap, $\sigma_{max}\approx 2h/\rho^{2}$. 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 a proxy of the dominant eigenvalue of Hessian.

###  6.2 GSAM finds a minimum whose Hessian has small dominant eigenvalues

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 $\rho$, increasing $\alpha$ decreases the
dominant value and improves generalization (test accuracy). In the middle
subfigure, we plot the dominant eigenvalues estimated by the surrogate gap,
$\sigma_{max}\approx 2h/\rho^{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 $\sigma_{max}$ (right) in terms of the trend that
$\sigma_{max}$ decreases with $\alpha$ and $\rho$. 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.

Figure 5: Top-1 accuracy of Mixer-S/32 trained with different methods.
“+ascent” represents applying the ascent step in Algo. 1 to an optimizer. Note
that our GSAM is described as SAM+ascent(=GSAM) for consistency.

###  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-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 $\rho$ between 1 and 7 ($10\times$ 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.

Table 2: Results (%) of GSAM and $\operatorname{min}(f_{p}+\lambda h)$ on
ViT-B/32

Dataset | $\operatorname{min}(f_{p}+\lambda h)$ | GSAM
---|---|---
ImageNet | 75.4 | 76.8
ImageNet-Real | 81.1 | 82.7
ImageNet-v2 | 60.9 | 63.0
ImageNet-R | 23.9 | 25.1

Table 3: Transfer learning results (top-1 accuracy, %)

| ViT-B/16 | ViT-S/16
---|---|---
| Vanilla | SAM | GSAM | Vanilla | SAM | GSAM
Cifar10 | 98.1 | 98.6 | 98.8 | 97.6 | 98.2 | 98.4
Cifar100 | 87.6 | 89.1 | 89.7 | 85.7 | 87.6 | 88.1
Flowers | 88.5 | 91.8 | 91.2 | 86.4 | 91.5 | 90.3
Pets | 91.9 | 93.1 | 94.4 | 90.4 | 92.9 | 93.5
mean | 91.5 | 93.2 | 93.5 | 90.0 | 92.6 | 92.6

Figure 6: Top-1 accuracy of ViT-B/32 for the additional studies (Section 6.4).
Left: from left to right are performances under different data augmentations
(details in Appendix B.3) , where the vanilla method is trained for $2\times$
the epochs. Middle: performance with different base optimizers. Right:
Comparison between $\operatorname{min}(f_{p},h)$ and
$\operatorname{min}(f,h)$.

###  6.4 Additional studies

GSAM outperforms a weighted combination of the perturbed loss and surrogate
gap With an example in Fig. 2, we demonstrate that directly minimizing
$f_{p}(w)+\lambda h(w)$ as discussed in Sec. 4.1 is sub-optimal because
$\nabla h(w)$ could conflict with $\nabla f_{p}(w)$ and $\nabla f(w)$. We
empirically validate this argument on ViT-B/32. We search for $\lambda$
between 0.0 and 0.5 with a step 0.1 and search for $\rho$ 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.

$\boldsymbol{\operatorname{min}(f_{p},h)}$ vs.
$\boldsymbol{\operatorname{min}(f,h)}$ GSAM solves
$\operatorname{min}(f_{p},h)$ by descent in $\nabla f_{p}$, decomposing
$\nabla f$ onto $\nabla f_{p}$, and an ascent step in the orthogonal direction
to increase $f$ while keep $f_{p}$ intact. Alternatively, we can also optimize
$\operatorname{min}(f,h)$ by descent in $\nabla f$, decomposing $\nabla f_{p}$
onto $\nabla f$, and a descent step in the orthogonal direction to decrease
$f_{p}$ while keep $f$ intact. The two GSAM variations perform similarly (see
Fig. 6, right). We choose $\operatorname{min}(f_{p},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), Oxford-flowers (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.

##  7 Conclusion

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.

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

## Reproducibility Statement

We provide the detailed proof of theoretical results in Appendix A and provide
the data pre-processing 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.

## References

  * Balestriero et al. (2021) Randall Balestriero, Jerome Pesenti, and Yann LeCun.  Learning in high dimension always amounts to extrapolation.  _arXiv preprint arXiv:2110.09485_ , 2021.
  * Beyer et al. (2020) Lucas Beyer, Olivier J. Henaff, Alexander Kolesnikov, Xiaohua Zhai, and Aaron van den Oord.  Are we done with imagenet?  _arXiv preprint arXiv:2002.05709_ , 2020.
  * Bottou (2010) Léon Bottou.  Large-scale machine learning with stochastic gradient descent.  In _Proceedings of COMPSTAT’2010_ , pp. 177–186. Springer, 2010\.
  * Chaudhari et al. (2019) Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina.  Entropy-sgd: Biasing gradient descent into wide valleys.  _Journal of Statistical Mechanics: Theory and Experiment_ , 2019(12):124018, 2019.
  * Chen et al. (2021) Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong.  When vision transformers outperform resnets without pretraining or strong data augmentations, 2021.
  * Cubuk et al. (2018) Ekin D Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoc V Le.  Autoaugment: Learning augmentation policies from data.  _arXiv preprint arXiv:1805.09501_ , 2018.
  * Damian et al. (2021) Alex Damian, Tengyu Ma, and Jason Lee.  Label noise sgd provably prefers flat global minimizers.  _arXiv preprint arXiv:2106.06530_ , 2021.
  * Deng et al. (2009) Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.  Imagenet: A large-scale hierarchical image database.  In _2009 IEEE conference on computer vision and pattern recognition_ , pp. 248–255. Ieee, 2009.
  * DeVries & Taylor (2017) Terrance DeVries and Graham W Taylor.  Improved regularization of convolutional neural networks with cutout.  _arXiv preprint arXiv:1708.04552_ , 2017.
  * Dosovitskiy et al. (2020) Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  _arXiv preprint arXiv:2010.11929_ , 2020.
  * Duchi et al. (2011) John Duchi, Elad Hazan, and Yoram Singer.  Adaptive subgradient methods for online learning and stochastic optimization.  _Journal of machine learning research_ , 12(Jul):2121–2159, 2011.
  * Foret et al. (2020) Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur.  Sharpness-aware minimization for efficiently improving generalization.  _arXiv preprint arXiv:2010.01412_ , 2020.
  * Gastaldi (2017) Xavier Gastaldi.  Shake-shake regularization.  _arXiv preprint arXiv:1705.07485_ , 2017.
  * He et al. (2016) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 770–778, 2016.
  * Heo et al. (2020) Byeongho Heo, Sanghyuk Chun, Seong Joon Oh, Dongyoon Han, Sangdoo Yun, Gyuwan Kim, Youngjung Uh, and Jung-Woo Ha.  Adamp: Slowing down the slowdown for momentum optimizers on scale-invariant weights.  _arXiv preprint arXiv:2006.08217_ , 2020.
  * Hochreiter & Schmidhuber (1995) Sepp Hochreiter and Jürgen Schmidhuber.  Simplifying neural nets by discovering flat minima.  In _Advances in neural information processing systems_ , pp. 529–536, 1995.
  * Izmailov et al. (2018) Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson.  Averaging weights leads to wider optima and better generalization.  _arXiv preprint arXiv:1803.05407_ , 2018.
  * Jacot et al. (2018) Arthur Jacot, Franck Gabriel, and Clément Hongler.  Neural tangent kernel: Convergence and generalization in neural networks.  _arXiv preprint arXiv:1806.07572_ , 2018.
  * Jiang et al. (2019) Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio.  Fantastic generalization measures and where to find them.  _arXiv preprint arXiv:1912.02178_ , 2019.
  * Keskar & Socher (2017) Nitish Shirish Keskar and Richard Socher.  Improving generalization performance by switching from adam to sgd.  _arXiv preprint arXiv:1712.07628_ , 2017.
  * Keskar et al. (2016) Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang.  On large-batch training for deep learning: Generalization gap and sharp minima.  _arXiv preprint arXiv:1609.04836_ , 2016.
  * Krizhevsky et al. (2009) Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images.  2009\.
  * Kwon et al. (2021) 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.11600_ , 2021.
  * Laurent & Massart (2000) Beatrice Laurent and Pascal Massart.  Adaptive estimation of a quadratic functional by model selection.  _Annals of Statistics_ , pp. 1302–1338, 2000.
  * Li et al. (2019) Yuanzhi Li, Colin Wei, and Tengyu Ma.  Towards explaining the regularization effect of initial large learning rate in training neural networks.  _arXiv preprint arXiv:1907.04595_ , 2019.
  * Liang et al. (2019) Tengyuan Liang, Tomaso Poggio, Alexander Rakhlin, and James Stokes.  Fisher-rao metric, geometry, and complexity of neural networks.  In _The 22nd International Conference on Artificial Intelligence and Statistics_ , pp. 888–896. PMLR, 2019.
  * Lin et al. (2020) Tao Lin, Lingjing Kong, Sebastian Stich, and Martin Jaggi.  Extrapolation for large-batch training in deep learning.  In _International Conference on Machine Learning_ , pp. 6094–6104. PMLR, 2020.
  * Liu et al. (2019) 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.03265_ , 2019.
  * Loshchilov & Hutter (2017) Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  _arXiv preprint arXiv:1711.05101_ , 2017.
  * Luo et al. (2019) Liangchen Luo, Yuanhao Xiong, Yan Liu, and Xu Sun.  Adaptive gradient methods with dynamic bound of learning rate.  _arXiv preprint arXiv:1902.09843_ , 2019.
  * McAllester (2003) David McAllester.  Simplified pac-bayesian margin bounds.  In _Learning theory and Kernel machines_ , pp. 203–215. Springer, 2003.
  * McAllester (1999) David A McAllester.  Pac-bayesian model averaging.  In _Proceedings of the twelfth annual conference on Computational learning theory_ , pp. 164–170, 1999.
  * Mises & Pollaczek-Geiringer (1929) RV Mises and Hilda Pollaczek-Geiringer.  Praktische verfahren der gleichungsauflösung.  _ZAMM-Journal of Applied Mathematics and Mechanics/Zeitschrift für Angewandte Mathematik und Mechanik_ , 9(1):58–77, 1929\.
  * Müller et al. (2019) Rafael Müller, Simon Kornblith, and Geoffrey Hinton.  When does label smoothing help?  _arXiv preprint arXiv:1906.02629_ , 2019.
  * Neyshabur et al. (2017) Behnam Neyshabur, Srinadh Bhojanapalli, and Nathan Srebro.  A pac-bayesian approach to spectrally-normalized margin bounds for neural networks.  _arXiv preprint arXiv:1707.09564_ , 2017.
  * Nilsback & Zisserman (2008) Maria-Elena Nilsback and Andrew Zisserman.  Automated flower classification over a large number of classes.  In _2008 Sixth Indian Conference on Computer Vision, Graphics & Image Processing_, pp. 722–729. IEEE, 2008.
  * Parkhi et al. (2012) 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. IEEE, 2012.
  * Recht et al. (2019) Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar.  Do imagenet classifiers generalize to imagenet?  In _International Conference on Machine Learning_ , pp. 5389–5400, 2019.
  * Reddi et al. (2019) Sashank J Reddi, Satyen Kale, and Sanjiv Kumar.  On the convergence of adam and beyond.  _arXiv preprint arXiv:1904.09237_ , 2019.
  * Rumelhart et al. (1985) 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.
  * Srivastava et al. (2014) Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya 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.
  * Szegedy et al. (2015) Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed, Dragomir Anguelov, Dumitru 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.
  * Tolstikhin et al. (2021) Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, et al.  Mlp-mixer: An all-mlp architecture for vision.  _arXiv preprint arXiv:2105.01601_ , 2021.
  * Wei et al. (2019) Colin Wei, Jason Lee, Qiang Liu, and Tengyu Ma.  Regularization matters: Generalization and optimization of neural nets vs their induced kernel.  2019\.
  * Xie et al. (2021) Zeke Xie, Li Yuan, Zhanxing Zhu, and Masashi Sugiyama.  Positive-negative momentum: Manipulating stochastic gradient noise to improve generalization.  _arXiv preprint arXiv:2103.17182_ , 2021.
  * Yue et al. (2020) Xubo Yue, Maher Nouiehed, and Raed Al Kontar.  Salr: Sharpness-aware learning rates for improved generalization.  _arXiv preprint arXiv:2011.05348_ , 2020.
  * Zaheer et al. (2018) Manzil Zaheer, Sashank Reddi, Devendra Sachan, Satyen Kale, and Sanjiv Kumar.  Adaptive methods for nonconvex optimization.  In _Advances in neural information processing systems_ , pp. 9793–9803, 2018.
  * Zeiler (2012) Matthew D Zeiler.  Adadelta: an adaptive learning rate method.  _arXiv preprint arXiv:1212.5701_ , 2012.
  * Zhang et al. (2017a) Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals.  Understanding deep learning requires rethinking generalization.  2017a.
  * Zhang et al. (2017b) Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz.  mixup: Beyond empirical risk minimization.  _arXiv preprint arXiv:1710.09412_ , 2017b.
  * Zhang et al. (2019) Michael Zhang, James Lucas, Jimmy Ba, and Geoffrey E Hinton.  Lookahead optimizer: k steps forward, 1 step back.  In _Advances in Neural Information Processing Systems_ , pp. 9593–9604, 2019.
  * Zheng et al. (2021) Yaowei Zheng, Richong Zhang, and Yongyi Mao.  Regularizing neural networks via adversarial model perturbation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 8156–8165, 2021.
  * Zhou et al. (2020) Pan Zhou, Jiashi Feng, Chao Ma, Caiming Xiong, Steven Hoi, et al.  Towards theoretically understanding why sgd generalizes better than adam in deep learning.  _arXiv preprint arXiv:2010.05627_ , 2020.
  * Zhuang et al. (2020) Juntang Zhuang, Tommy Tang, Yifan Ding, Sekhar Tatikonda, Nicha Dvornek, Xenophon Papademetris, and James S Duncan.  Adabelief optimizer: Adapting stepsizes by the belief in observed gradients.  _arXiv preprint arXiv:2010.07468_ , 2020.

##  Appendix A Proofs

###  A.1 Proof of Lemma. 3.1

Suppose $\rho$ is small, perform Taylor expansion around the local minima $w$,
we have:

$f(w+\delta)=f(w)+\nabla
f(w)^{\top}\delta+\frac{1}{2}\delta^{\top}H\delta+O(||\delta||^{3})$ (10)

where $H$ is the Hessian, and is positive semidefinite at a local minima. At a
local minima, $\nabla f(w)=0$, hence we have

$f(w+\delta)=f(w)+\frac{1}{2}\delta^{\top}H\delta+O(||\delta||^{3})$ (11)

and

$f_{p}(w)=\operatorname{max}_{||\delta||\leq\rho}f(w+\delta)=f(w)+\frac{1}{2}\rho^{2}\sigma_{max}(H)+O(||\delta||^{3})$
(12)

where $\sigma_{max}$ is the dominate eigenvalue (eigenvalue with the largest
absolute value). Now consider two local minima $w_{1}$ and $w_{2}$ with
dominate eigenvalue $\sigma_{1}$ and $\sigma_{2}$ respectively, we have

$\displaystyle f_{p}(w_{1})\approx f(w_{1})+\frac{1}{2}\rho^{2}\sigma_{1}$
$\displaystyle f_{p}(w_{2})\approx f(w_{2})+\frac{1}{2}\rho^{2}\sigma_{2}$

We have $f_{p}(w_{1})>f_{p}(w_{2})\centernot\implies\sigma_{1}>\sigma_{2}$ and
$\sigma_{1}>\sigma_{2}\centernot\implies f_{p}(w_{1})>f_{p}(w_{2})$ because
the relation between $f(w_{1})$ and $f(w_{2})$ is undetermined. $\square$

###  A.2 Proof of Lemma. 3.2

Since $\rho$ is small, we can perform Taylor expansion around $w$,

$\displaystyle h(w)$ $\displaystyle=f(w+\delta)-f(w)$
$\displaystyle=\delta^{\top}\nabla f(w)+O(\rho^{2})$
$\displaystyle=\rho||\nabla f(w)||_{2}+O(\rho^{2})>0$ (13)

where the last line is because $\delta$ is approximated as
$\delta=\rho\frac{\nabla f(w)}{||\nabla f(w)||_{2}+\epsilon}$, hence has the
same direction as $\nabla f(w)$. $\square$

###  A.3 Proof of Lemma. 3.3

Since $\rho$ is small, we can approximate $f(w)$ with a quadratic model around
a local minima $w$:

$f(w+\delta)=f(w)+\frac{1}{2}\delta^{\top}H\delta+O(\rho^{3})$

where $H$ is the Hessian at $w$, assumed to be positive semidefinite at local
minima. Normalize $\delta$ such that $||\delta||_{2}=\rho$, Hence we have:

$\displaystyle
h(w)=f_{p}(w)-f(w)=\operatorname{max}_{||\delta||_{2}\leq\rho}f(w+\delta)-f(w)=\frac{1}{2}\sigma_{max}\rho^{2}+O(\rho^{3})$
(14)

where $\sigma_{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
$\sigma_{max}\approx 2h(w)/\rho^{2}$. $\square$

###  A.4 Proof of Thm. 5.1

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 $f_{p}(w)$

For simplicity of notation, we denote the update at step $t$ as

$d_{t}=-\eta_{t}g_{p}^{(t)}+\eta_{t}\alpha g^{(t)}_{\perp}$ (15)

By $L-$smoothness of $f$ and the definition of $f_{p}(w_{t})=f(w^{adv}_{t})$,
and definition of $d_{t}=w_{t+1}-w_{t}$ and $w^{adv}_{t}=w_{t}+\delta_{t}$ we
have

$\displaystyle f_{p}(w_{t+1})$ $\displaystyle=f(w^{adv}_{t+1})\leq
f(w^{adv}_{t})+\langle\nabla
f(w^{adv}_{t}),w^{adv}_{t+1}-w^{adv}_{t}\rangle+\frac{L}{2}\Big{|}\Big{|}w^{adv}_{t+1}-w^{adv}_{t}\Big{|}\Big{|}^{2}$
(16) $\displaystyle=f(w_{t}^{adv})+\langle\nabla
f(w_{t}^{adv}),w_{t+1}+\delta_{t+1}-w_{t}-\delta_{t}\rangle$
$\displaystyle+\frac{L}{2}\Big{|}\Big{|}w_{t+1}+\delta_{t+1}-w_{t}-\delta_{t}\Big{|}\Big{|}^{2}$
(17) $\displaystyle\leq f(w^{adv}_{t})+\langle\nabla
f(w^{adv}_{t}),d_{t}\rangle+L\Big{|}\Big{|}d_{t}\Big{|}\Big{|}^{2}$ (18)
$\displaystyle+\langle\nabla
f(w_{t}^{adv}),\delta_{t+1}-\delta_{t}\rangle+L\Big{|}\Big{|}\delta_{t+1}-\delta_{t}\Big{|}\Big{|}^{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 $\mathbb{E}$ short for
$\mathbb{E}_{x}$ to denote expectation over all possible data points)
conditioned on observations up to step $t$, also by definition of $d_{t}$, we
have

$\displaystyle\mathbb{E}f_{p}(w_{t+1})-f_{p}(w_{t})$
$\displaystyle\leq-\eta_{t}\langle\nabla
f_{p}(w_{t}),\mathbb{E}g_{p}^{(t)}\rangle+\alpha\eta_{t}\langle\nabla
f_{p}(w_{t}),\mathbb{E}g^{(t)}_{\perp}\rangle$
$\displaystyle+L\eta_{t}^{2}\mathbb{E}\Big{|}\Big{|}-g_{p}^{(t)}+\alpha
g_{\perp}^{(t)}\Big{|}\Big{|}_{2}^{2}$ (20)
$\displaystyle\leq-\eta_{t}\mathbb{E}\Big{|}\Big{|}\nabla
f_{p}(w_{t})\Big{|}\Big{|}_{2}^{2}+0+(\alpha+1)^{2}G^{2}\eta_{t}^{2}$ (21)
$\displaystyle\Big{(}\text{Since }\mathbb{E}g^{(t)}_{\perp}\text{ is
orthogonal to }\nabla f_{p}(w_{t})\text{ by construction},$
$\displaystyle||g^{(t)}||\leq G\text{ by assumption}\Big{)}$

#### Step 1.1: Bound Eq. 19

By definition of $\delta_{t}$, we have

$\displaystyle\delta_{t}$
$\displaystyle=\rho_{t}\frac{g^{(t)}}{||g^{(t)}||+\epsilon}$ (22)
$\displaystyle\delta_{t+1}$
$\displaystyle=\rho_{t+1}\frac{g^{(t+1)}}{||g^{(t+1)}||+\epsilon}$ (23)

where $g^{(t)}$ is the gradient of $f$ at $w_{t}$ evaluated with a noisy data
sample. When learning rate $\eta_{t}$ is small, the update in weight $d_{t}$
is small, and expected gradient is

$\nabla f(w_{t+1})=\nabla f(w_{t}+d_{t})=\nabla
f(w_{t})+Hd_{t}+O(||d_{t}||^{2})$ (24)

where $H$ is the Hessian at $w_{t}$. Therefore, we have

$\displaystyle\mathbb{E}\langle\nabla
f(w^{adv}_{t}),\delta_{t+1}-\delta_{t}\rangle$ $\displaystyle=\langle\nabla
f(w^{adv}_{t}),\rho_{t}\mathbb{E}\frac{g^{(t)}}{||g^{(t)}||+\epsilon}-\rho_{t+1}\mathbb{E}\frac{g^{(t+1)}}{||g^{(t+1)}||+\epsilon}\rangle$
(25) $\displaystyle\leq||\nabla
f(w_{t}^{adv})||\rho_{t}\Big{|}\Big{|}\mathbb{E}\frac{g^{(t)}}{||g^{(t)}||+\epsilon}-\mathbb{E}\frac{g^{(t+1)}}{||g^{(t+1)}||+\epsilon}\Big{|}\Big{|}$
(26) $\displaystyle\leq||\nabla f(w_{t}^{adv})||\rho_{t}\phi_{t}$ (27)

where the first inequality is due to (1) $\rho_{t}$ is monotonically
decreasing with $t$, and (2) triangle inequality that $\langle
a,b\rangle\leq||a||\cdot||b||$. $\phi_{t}$ is the angle between the unit
vector in the direction of $\nabla f(w_{t})$ and $\nabla f(w_{t+1})$. The
second inequality comes from that (1)
$\Big{|}\Big{|}\frac{g}{||g||+\epsilon}\Big{|}\Big{|}<1$ strictly, so we can
replace $\delta_{t}$ in Eq. 25 with a unit vector in corresponding directions
multiplied by $\rho_{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 $\eta_{t}$ and update stepsize $d_{t}$ is small, $\phi_{t}$
is also small. Using the limit that

$\operatorname{tan}x=x+O(x^{2}),\ \ \ \ \operatorname{sin}x=x+O(x^{2}),\ \ \ \
x\to 0$

We have:

$\displaystyle\operatorname{tan}\phi_{t}$ $\displaystyle=\frac{||\nabla
f(w_{t+1})-\nabla f(w_{t})||}{||\nabla f(w_{t})||}+O(\phi_{t}^{2})$ (28)
$\displaystyle=\frac{||Hd_{t}+O(||d_{t}||^{2})||}{||\nabla
f(w_{t})||}+O(\phi_{t}^{2})$ (29) $\displaystyle\leq\eta_{t}L(1+\alpha)$ (30)

where the last inequality is due to (1) max eigenvalue of $H$ is upper bounded
by $L$ because $f$ is $L-$smooth, (2)
$||d_{t}||=||\eta_{t}(g_{\parallel}+\alpha g_{\perp})||$ and
$\mathbb{E}g_{t}=\nabla f(w_{t})$.

Plug into Eq. 27, also note that the perturbation amplitude $\rho_{t}$ is
small so $w_{t}$ is close to $w^{adv}_{t}$, then we have

$\mathbb{E}\langle\nabla f(w^{adv}_{t}),\delta_{t+1}-\delta_{t}\rangle\leq
L(1+\alpha)G\rho_{t}\eta_{t}$ (31)

Similarly, we have

$\displaystyle\mathbb{E}\Big{|}\Big{|}\delta_{t+1}-\delta_{t}\Big{|}\Big{|}^{2}$
$\displaystyle\leq\rho_{t}^{2}\mathbb{E}\Big{|}\Big{|}\frac{g^{(t)}}{||g^{(t)}||+\epsilon}-\frac{g^{(t+1)}}{||g^{(t+1)}||+\epsilon}\Big{|}\Big{|}^{2}$
(32) $\displaystyle\leq\rho_{t}^{2}\phi_{t}^{2}$ (33)
$\displaystyle\leq\rho_{t}^{2}\eta_{t}^{2}L^{2}(1+\alpha)^{2}$ (34)

#### Step 1.2: Total bound

Reuse results from Eq. 21 (replace $L_{p}$ with $2L$) and plug into Eq. 18,
and plug Eq. 31 and Eq. 34 into Eq. 19, we have

$\displaystyle\mathbb{E}f_{p}(w_{t+1})-f_{p}(w_{t})$
$\displaystyle\leq-\eta_{t}\mathbb{E}\Big{|}\Big{|}\nabla
f_{p}(w_{t})\Big{|}\Big{|}_{2}^{2}+\frac{2L(\alpha+1)^{2}}{2}G^{2}\eta_{t}^{2}$
$\displaystyle+L(1+\alpha)G\rho_{t}\eta_{t}+\frac{2L^{3}(1+\alpha)^{2}}{2}\eta_{t}^{2}\rho_{t}^{2}$
(35)

Perform telescope sum, we have

$\displaystyle\mathbb{E}f_{p}(w_{T})-f_{p}(w_{0})$
$\displaystyle\leq-\sum_{t=1}^{T}\eta_{t}\mathbb{E}||\nabla
f_{p}(w_{t})||^{2}+\Big{[}L(1+\alpha)^{2}G^{2}\eta_{0}^{2}+L(1+\alpha)G\rho_{0}\eta_{0}\Big{]}\sum_{t=1}^{T}\frac{1}{t}$
$\displaystyle+L^{3}(1+\alpha)^{2}\eta_{0}^{2}\rho_{0}^{2}\sum_{t=1}^{T}\frac{1}{t^{2}}$
(36)

Hence

$\displaystyle\eta_{T}\sum_{t=1}^{T}\mathbb{E}||\nabla
f_{p}(w_{t})||^{2}\leq\sum_{t=1}^{T}\eta_{t}\mathbb{E}||\nabla
f_{p}(w_{t})||^{2}\leq f_{p}(w_{0})-\mathbb{E}f_{p}(w_{T})+D\log
T+\frac{\pi^{2}E}{6}$ (37)

where

$D=L(1+\alpha)^{2}G^{2}\eta_{0}^{2}+L(1+\alpha)G\rho_{0}\eta_{0},\ \ \ \
E=L^{3}(1+\alpha)^{2}\eta_{0}^{2}\rho_{0}^{2}$ (38)

Note that $\eta_{T}=\frac{\eta_{0}}{\sqrt{T}}$, we have

$\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}||\nabla
f_{p}(w_{t})||^{2}\leq\frac{f_{p}(w_{0})-f_{min}+\pi^{2}E/6}{\eta_{0}}\frac{1}{\sqrt{T}}+\frac{D}{\eta_{0}}\frac{\log
T}{\sqrt{T}}$ (39)

which implies that GSAM enables $f_{p}$ to converge at a rate of $O(\log
T/\sqrt{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

$d_{t}=-\eta_{t}g_{p}^{(t)}+\alpha\eta_{t}g_{\perp}^{(t)}$ (40)

By smoothness of $f$, we have

$\displaystyle f(w_{t+1})$ $\displaystyle\leq f(w_{t})+\langle\nabla
f(w_{t}),d_{t}\rangle+\frac{L}{2}\Big{|}\Big{|}d_{t}\Big{|}\Big{|}_{2}^{2}$
(41) $\displaystyle=f(w_{t})+\langle\nabla
f(w_{t}),-\eta_{t}g_{p}^{(t)}+\alpha\eta_{t}g_{\perp}^{(t)}\rangle+\frac{L}{2}\Big{|}\Big{|}d_{t}\Big{|}\Big{|}_{2}^{2}$
(42)

For simplicity, we introduce a scalar $\beta_{t}$ such that

$\nabla f_{\parallel}(w_{t})=\beta_{t}\nabla f_{p}(w_{t})$ (43)

where $\nabla f_{\parallel}(w_{t})$ is the projection of $\nabla f(w_{t})$
onto $\nabla f_{p}(w_{t})$. When perturbation amplitude $\rho$ is small, we
expect $\beta_{t}$ to be very close to 1.

Take expectation conditioned on observations up to step $t$ for both sides of
Eq. 42, we have:

$\displaystyle\mathbb{E}f(w_{t+1})$ $\displaystyle\leq
f(w_{t})+\Bigg{\langle}\nabla
f(w_{t}),-\frac{\eta_{t}}{\beta_{t}}\Big{(}\nabla f(w_{t})-\nabla
f_{\perp}(w_{t})\Big{)}+\alpha\eta_{t}\mathbb{E}g^{(t)}_{\perp}\Bigg{\rangle}+\frac{L}{2}\mathbb{E}\Big{|}\Big{|}d_{t}\Big{|}\Big{|}_{2}^{2}$
(44) $\displaystyle=f(w_{t})-\frac{\eta_{t}}{\beta_{t}}\Big{|}\Big{|}\nabla
f(w_{t})\Big{|}\Big{|}_{2}^{2}+\Big{(}\frac{1}{\beta_{t}}+\alpha\Big{)}\eta_{t}\Big{\langle}\nabla
f(w_{t}),\nabla
f_{\perp}(w_{t})\Big{\rangle}+\frac{L}{2}\mathbb{E}\Big{|}\Big{|}d_{t}\Big{|}\Big{|}_{2}^{2}$
(45) $\displaystyle=f(w_{t})-\frac{\eta_{t}}{\beta_{t}}\Big{|}\Big{|}\nabla
f(w_{t})\Big{|}\Big{|}_{2}^{2}+\Big{(}\frac{1}{\beta_{t}}+\alpha\Big{)}\eta_{t}\Big{\langle}\nabla
f(w_{t}),\nabla
f(w_{t})\operatorname{sin}\theta_{t}\Big{\rangle}+\frac{L}{2}\mathbb{E}\Big{|}\Big{|}d_{t}\Big{|}\Big{|}_{2}^{2}$
(46) $\displaystyle\Big{(}\theta_{t}\text{ is the angle between }\nabla
f_{p}(w_{t})\text{ and }\nabla f(w_{t})\Big{)}$
$\displaystyle=f(w_{t})-\frac{\eta_{t}}{\beta_{t}}\Big{|}\Big{|}\nabla
f(w_{t})\Big{|}\Big{|}_{2}^{2}+\Big{(}\frac{1}{\beta_{t}}+\alpha\Big{)}\eta_{t}\Big{|}\Big{|}\nabla
f(w_{t})\Big{|}\Big{|}_{2}^{2}(|\operatorname{tan}\theta_{t}|+O(\theta_{t}^{2}))+\frac{L}{2}\mathbb{E}\Big{|}\Big{|}d_{t}\Big{|}\Big{|}_{2}^{2}$
(47)
$\displaystyle\Big{(}\operatorname{sin}x=x+O(x^{2}),\operatorname{tan}x=x+O(x^{2})\text{
when }x\to 0.\Big{)}$

Also note when perturbation amplitude $\rho_{t}$ is small, we have

$\nabla f_{p}(w_{t})=\nabla f(w_{t}+\delta_{t})=\nabla
f(w_{t})+\frac{\rho_{t}}{||\nabla f(w_{t})||_{2}+\epsilon}H(w_{t})\nabla
f(w_{t})+O(\rho_{t}^{2})$ (48)

where $\delta_{t}=\rho_{t}\frac{\nabla f(w_{t})}{||\nabla f(w_{t})||_{2}}$ by
definition, $H(w_{t})$ is the Hessian. Hence we have

$|\operatorname{tan}\theta_{t}|\leq\frac{||\nabla f_{p}(w_{t})-\nabla
f(w_{t})||}{||\nabla f(w_{t})||}\leq\frac{\rho_{t}L}{||\nabla f(w_{t})||}$
(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

$\displaystyle\mathbb{E}f(w_{t+1})$ $\displaystyle\leq
f(w_{t})-\frac{\eta_{t}}{\beta_{t}}\Big{|}\Big{|}\nabla
f(w_{t})\Big{|}\Big{|}_{2}^{2}+\Big{(}\frac{1}{\beta_{t}}+\alpha\Big{)}\eta_{t}\Big{|}\Big{|}\nabla
f(w_{t})\Big{|}\Big{|}_{2}^{2}|\operatorname{tan}\theta_{t}|+\frac{L}{2}\mathbb{E}\Big{|}\Big{|}d_{t}\Big{|}\Big{|}_{2}^{2}$
(50) $\displaystyle\leq
f(w_{t})-\frac{\eta_{t}}{\beta_{t}}\Big{|}\Big{|}\nabla
f(w_{t})\Big{|}\Big{|}_{2}^{2}+\Big{(}\frac{1}{\beta_{t}}+\alpha\Big{)}L\rho_{t}\eta_{t}\Big{|}\Big{|}\nabla
f(w_{t})\Big{|}\Big{|}_{2}+\frac{L}{2}\mathbb{E}\Big{|}\Big{|}d_{t}\Big{|}\Big{|}_{2}^{2}$
(51) $\displaystyle\leq
f(w_{t})-\frac{\eta_{t}}{\beta_{t}}\Big{|}\Big{|}\nabla
f(w_{t})\Big{|}\Big{|}_{2}^{2}+\Big{(}\frac{1}{\beta_{t}}+\alpha\Big{)}L\rho_{t}\eta_{t}G+\frac{L}{2}\mathbb{E}\Big{|}\Big{|}d_{t}\Big{|}\Big{|}_{2}^{2}$
(52) $\displaystyle\Big{(}\text{Assume gradient has bounded norm }G.\Big{)}$
(53) $\displaystyle\leq
f(w_{t})-\frac{\eta_{t}}{\beta_{max}}\Big{|}\Big{|}\nabla
f(w_{t})\Big{|}\Big{|}_{2}^{2}+\Big{(}\frac{1}{\beta_{min}}+\alpha\Big{)}L\rho_{t}\eta_{t}G+\frac{L}{2}\mathbb{E}(\alpha+1)^{2}G^{2}\eta_{t}^{2}$
(54) $\displaystyle\Big{(}\beta_{t}\text{ is close to 1 assuming }\rho\text{
is small,}$ $\displaystyle\text{hence it's natural to assume
}0<\beta_{min}\leq\beta_{t}\leq\beta_{max}\Big{)}$

Re-arranging above formula, we have

$\frac{\eta_{t}}{\beta_{max}}\Big{|}\Big{|}\nabla
f(w_{t})\Big{|}\Big{|}_{2}^{2}\leq
f(w_{t})-\mathbb{E}f(w_{t+1})+\Big{(}\frac{1}{\beta_{min}}+\alpha\Big{)}LG\eta_{t}\rho_{t}+\frac{L}{2}(\alpha+1)^{2}G^{2}\eta_{t}^{2}$
(55)

perform telescope sum and taking expectations on each step, we have

$\frac{1}{\beta_{max}}\sum_{t=1}^{T}\eta_{t}\Big{|}\Big{|}\nabla
f(w_{t})\Big{|}\Big{|}_{2}^{2}\leq
f(w_{0})-\mathbb{E}f(w_{T})+\Big{(}\frac{1}{\beta_{min}}+\alpha\Big{)}LG\sum_{t=1}^{T}\eta_{t}\rho_{t}+\frac{L}{2}(\alpha+1)^{2}G^{2}\sum_{t=1}^{T}\eta_{t}^{2}$
(56)

Take the schedule to be $\eta_{t}=\frac{\eta_{0}}{\sqrt{t}}$ and
$\rho_{t}=\frac{\rho_{0}}{\sqrt{t}}$, then we have

$\displaystyle\frac{\eta_{0}}{\beta_{max}}\frac{1}{\sqrt{T}}\sum_{t=1}^{T}\Big{|}\Big{|}\nabla
f(w_{t})\Big{|}\Big{|}_{2}^{2}$ $\displaystyle\leq LHS$ (57)
$\displaystyle\leq RHS$ (58) $\displaystyle\leq
f(w_{0})-f_{min}+\Big{(}\frac{1}{\beta_{min}}+\alpha\Big{)}LG\eta_{0}\rho_{0}\sum_{t=1}^{T}\frac{1}{t}+\frac{L}{2}(\alpha+1)^{2}G^{2}\eta_{0}^{2}\sum_{t=1}^{T}\frac{1}{t}$
(59) $\displaystyle\leq
f(w_{0})-f_{min}+\Big{(}\frac{1}{\beta_{min}}+\alpha\Big{)}LG\eta_{0}\rho_{0}(1+\log
T)$ $\displaystyle+\frac{L}{2}(\alpha+1)^{2}G^{2}\eta_{0}^{2}(1+\log T)$ (60)

Hence

$\frac{1}{T}\sum_{t=1}^{T}\Big{|}\Big{|}\nabla
f(w_{t})\Big{|}\Big{|}_{2}^{2}\leq\frac{C_{3}}{\sqrt{T}}+C_{4}\frac{\log
T}{\sqrt{T}}$ (61)

where $C_{1},C_{4}$ are some constants. This implies the convergence rate
w.r.t $f(w)$ is $O(\log T/\sqrt{T})$.

#### Step 3: Convergence w.r.t. surrogate gap $h(w)$

Note that we have proved convergence for $f_{p}(w)$ in step 1, and convergence
for $f(w)$ in step 3. Also note that

$\Big{|}\Big{|}\nabla h(w_{t})\Big{|}\Big{|}_{2}^{2}=\Big{|}\Big{|}\nabla
f_{p}(w_{t})-\nabla f(w_{t})\Big{|}\Big{|}_{2}^{2}\leq 2\Big{|}\Big{|}\nabla
f_{p}(w_{t})\Big{|}\Big{|}_{2}^{2}+2\Big{|}\Big{|}\nabla
f(w_{t})\Big{|}\Big{|}_{2}^{2}$ (62)

Hence

$\frac{1}{T}\sum_{t=1}^{T}\Big{|}\Big{|}\nabla
h(w_{t})\Big{|}\Big{|}_{2}^{2}\leq\frac{2}{T}\sum_{t=1}^{T}\Big{|}\Big{|}\nabla
f_{p}(w_{t})\Big{|}\Big{|}_{2}^{2}+\frac{2}{T}\sum_{t=1}^{T}\Big{|}\Big{|}\nabla
f(w_{t})\Big{|}\Big{|}_{2}^{2}$ (63)

also converges at rate $O(\log T/\sqrt{T})$ because each item in the RHS
converges at rate $O(\log T\sqrt{T})$. $\square$

###  A.5 Proof of Corollary. 5.2.1

Using the results from Thm. 5.2, with probability at least $1-a$, we have

$\mathbb{E}_{w\sim\mathcal{Q}}\mathbb{E}_{x}f(w,x)\leq\mathbb{E}_{w\sim\mathcal{Q}}\widehat{f}(w)+4\sqrt{\frac{KL(\mathcal{Q}||\mathcal{P})+\log\frac{2m}{a}}{m}}$
(64)

Assume $\delta\sim\mathcal{N}(0,b^{2}I_{k})$ where $k$ is the dimension of
model parameters, hence $\delta^{2}$ (element-wise square) follows a a Chi-
square distribution. By Lemma.1 in Laurent & Massart (2000), we have

$\mathbb{P}\big{(}||\delta||_{2}^{2}-kb^{2}\geq
2b^{2}\sqrt{kt}+2tb^{2}\big{)}\leq exp(-t)$ (65)

hence with probability at least $1-1/\sqrt{n}$, we have

$||\delta||_{2}^{2}\leq
b^{2}\Bigg{(}2\log\sqrt{n}+k+2\sqrt{k\log\sqrt{n}}\Bigg{)}\leq
2b^{2}k\Bigg{(}1+\sqrt{\frac{\log\sqrt{n}}{k}}\Bigg{)}^{2}\leq\rho^{2}$ (66)

Therefore, with probability at least
$1-1/\sqrt{n}=1-exp\Big{(}-\big{(}\frac{\rho}{\sqrt{2}b}-\sqrt{k}\big{)}^{2}\Big{)}$

$\mathbb{E}_{\delta}\widehat{f}(w+\delta)\leq\operatorname{max}_{||\delta||_{2}\leq\rho}\widehat{f}(w+\delta)$
(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
$\mathbb{E}_{w\sim\mathcal{Q}}\mathbb{E}_{x}f(w,x)\leq\mathbb{E}_{\delta\sim\mathcal{N}(0,b^{2}I^{k})}\mathbb{E}_{w\sim\mathcal{Q}}\mathbb{E}_{x}f(w+\delta,x)$we
finish the proof. $\square$

###  A.6 Proof of Thm. 5.3

#### Step 1: 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

$\displaystyle\mathbb{E}\langle\nabla f_{p}(w_{t})-\nabla
f(w_{t}),-\eta_{t}\nabla f_{p}(w_{t})\rangle$ (68)
$\displaystyle\Big{(}\textit{where }\mathbb{E}g_{\perp}=\nabla
f_{\perp}(w_{t})\Big{)}$ $\displaystyle=\eta_{t}\Bigg{[}-\big{|}\big{|}\nabla
f_{p}(w_{t})\big{|}\big{|}_{2}^{2}+\big{|}\big{|}\nabla
f_{p}(w_{t})\big{|}\big{|}_{2}\big{|}\big{|}\nabla
f(w_{t})\big{|}\big{|}_{2}\operatorname{cos}\theta_{t}\Bigg{]}$ (69)

where $\theta_{t}$ is the angle between vector $\nabla f_{p}(w_{t})$ and
$\nabla f(w_{t})$.
The expected change of loss gap caused by ascent step is

$\mathbb{E}\langle\nabla f_{p}(w_{t})-\nabla f(w_{t}),\alpha\eta_{t}\nabla
f_{\perp}(w_{t})\rangle=-\alpha\eta_{t}\big{|}\big{|}\nabla
f_{\perp}(w_{t})\big{|}\big{|}_{2}^{2}<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 $\mathbb{E}\langle\nabla h(w_{t}),dt\rangle\leq 0$ requires
$\alpha$ to be large or $|\big{|}\nabla
f(w_{t})\big{|}\big{|}_{2}\operatorname{cos}\theta_{t}\leq\big{|}\big{|}\nabla
f_{p}(w_{t})\big{|}\big{|}$. In practice, the perturbation amplitude $\rho$ is
small and we can assume $\theta_{t}$ is close to 0 and $\big{|}\big{|}\nabla
f_{p}(w_{t})\big{|}\big{|}$ is close to $\big{|}\big{|}\nabla
f(w_{t})\big{|}\big{|}$, we can also set the parameter $\alpha$ 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 $\widehat{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:

$\displaystyle\Delta\widehat{h}_{t}=-\langle\nabla\widehat{f}_{p}(w_{t})-\nabla\widehat{f}(w_{t}),\alpha\eta_{t}\nabla\widehat{f}_{\perp}(w_{t})\rangle$
$\displaystyle=\alpha\eta_{t}\big{|}\big{|}\nabla\widehat{f}_{\perp}(w_{t})\big{|}\big{|}_{2}^{2}$
(71)
$\displaystyle=\alpha\eta_{t}\big{|}\big{|}\nabla\widehat{f}(w_{t})\big{|}\big{|}_{2}^{2}|\operatorname{tan}\theta_{t}|^{2}$
(72) $\displaystyle\sum_{t=1}^{T}\Delta\widehat{h}_{t}$
$\displaystyle\leq\sum_{t=1}^{T}\alpha L^{2}\eta_{t}\rho_{t}^{2}$ (73)
$\displaystyle\Big{(}\text{By Eq.~{}\ref{eq:bound_tan}}\Big{)}$ (74)
$\displaystyle\leq\sum_{t=1}^{T}\alpha
L^{2}\eta_{0}\rho_{0}^{2}\frac{1}{t^{3/2}}$ (75) $\displaystyle\leq 2.7\alpha
L^{2}\eta_{0}\rho_{0}^{2}$ (76)

Hence we derive an upper bound for $\sum_{t=1}^{T}\Delta\widehat{h}_{t}$.

Next we derive a lower bound for $\sum_{t=1}^{T}\Delta\widehat{h}_{t}$ Note
that when $\rho_{t}$ is small, by Taylor expansion

$\nabla\widehat{f}_{p}(w_{t})=\nabla\widehat{f}(w_{t}+\delta_{t})=\nabla\widehat{f}(w_{t})+\frac{\rho_{t}}{||\nabla\widehat{f}(w_{t})||}\widehat{H}(w_{t})\nabla\widehat{f}(w_{t})+O(\rho_{t}^{2})$
(77)

where $\widehat{H}(w_{t})$ is the Hessian evaluated on training samples. Also
when $\rho_{t}$ is small, the angle $\theta_{t}$ between
$\nabla\widehat{f}_{p}(w_{t})$ and $\nabla\widehat{f}(w_{t})$ is small, by the
limit that

$\displaystyle\operatorname{tan}x=x+O(x^{2}),x\to 0$
$\displaystyle\operatorname{sin}x=x+O(x^{2}),x\to 0$

We have

$|\operatorname{tan}\theta_{t}|=|\operatorname{sin}\theta_{t}|+O(\theta_{t}^{2})=|\theta_{t}|+O(\theta_{t}^{2})$

Omitting high order term, we have

$|\operatorname{tan}\theta_{t}|\approx|\theta_{t}|=\frac{||\nabla\widehat{f}_{p}(w_{t})-\nabla\widehat{f}(w_{t})||}{||\widehat{f}(w_{t})||}=\frac{||\rho_{t}\widehat{H}(w_{t})+O(\rho_{t}^{2})||}{||\nabla\widehat{f}(w_{t})||}\geq\frac{\rho_{t}|\sigma|_{min}}{G}$
(78)

where $G$ is the upper-bound on norm of gradient, $|\sigma|_{min}$ is the
minimum absolute eigenvalue of the Hessian. The intuition is that as
perturbation amplitude decreases, the angle $\theta_{t}$ decreases at a
similar rate, though the scale constant might be different. Hence we have

$\displaystyle\sum_{t=1}^{T}\Delta\widehat{h}_{t}$
$\displaystyle=\sum_{t=1}^{T}\alpha\eta_{t}\big{|}\big{|}\nabla\widehat{f}(w_{t})\big{|}\big{|}_{2}^{2}|\operatorname{tan}\theta_{t}|^{2}+O(\theta_{t}^{4})$
(79)
$\displaystyle\geq\sum_{t=1}^{T}\alpha\eta_{t}c^{2}\Big{(}\frac{\rho_{t}|\sigma|_{min}}{G}\Big{)}^{2}$
(80) $\displaystyle=\frac{\alpha
c^{2}\rho_{0}^{2}\eta_{0}|\sigma|_{min}^{2}}{G^{2}}\sum_{t=1}^{T}\frac{1}{t^{3/2}}$
(81) $\displaystyle\geq\frac{\alpha
c^{2}\rho_{0}^{2}\eta_{0}|\sigma|_{min}^{2}}{G^{2}}$ (82)

where $c^{2}$ is the lower bound of $||\nabla\widehat{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.
$\square$

###  A.7 Discussion on Corollary 5.2.1

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 $\rho_{train}$
for training equals its real value $\rho_{true}$ determined by the data
distribution; in practice, we never know $\rho_{true}$. In the following we
show that the minimization of both $h$ and $f_{p}$ is better than simply
minimizing $f_{p}$ when $\rho_{true}\neq\rho_{train}$.

1\. First, we re-write the conclusion of Corollary 5.2.1 as

$\displaystyle\mathbb{E}_{w}\mathbb{E}_{x}f(w,x)\leq
f_{p}+R=C+\widehat{h}+R=C+\rho^{2}\sigma/2+R+O(\rho^{3})$
$\displaystyle\textit{with probability
}(1-a)[1-e^{-(\frac{\rho}{\sqrt{2}b}-\sqrt{k})^{2}}]$

where $R$ is the regularization term, $C$ is the training loss, $\sigma$ is
the dominant eigenvalue of Hessian. As in lemma 3.3, we perform Taylor-
expansion and can ignore the high-order term $O(\rho^{3})$. We focus on

$f_{p}=C+\widehat{h}=C+\rho^{2}\sigma/2$

2\. When $\rho_{true}\neq\rho_{train}$, minimizing $h$ achieves a lower risk
than only minimizing $f_{p}$. (1) Note that after training, $C$ (training
loss) is fixed, but $h$ could vary with $\rho$ (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 $\rho$ is required for a valid bound).
(2) With an example, we show a low $f_{p}$ is insufficient for generalization,
and a low $\sigma$ is necessary:

  1. A

Suppose we use $\rho_{train}$ for training, and consider two solutions with
$C_{1},\sigma_{1}$ (SAM) and $C_{2},\sigma_{2}$ (GSAM). Suppose they have the
same $f_{p}$ during training for some $\rho_{train}$, so

$f_{p1}=C_{1}+\sigma_{1}/2\times\rho_{train}^{2}=C_{2}+\sigma_{2}/2\times\rho_{train}^{2}=f_{p2}$

Suppose $C_{1}<C_{2}$ so $\sigma_{1}>\sigma_{2}$.

  2. B

When $\rho_{true}>\rho_{train}$, we have

$\texttt{risk\\_bound\\_1}=C_{1}+\sigma_{1}/2\times\rho_{true}^{2}+R>\texttt{risk\\_bound\\_2}=C_{2}+\sigma_{2}/2\times\rho_{true}^{2}+R$

This implies that a small $\sigma$ helps generalization, but only a low
$f_{p1}$ (caused by a low $C_{1}$ and high $\sigma_{1}$) is insufficient for a
good generalization.

  3. C

Note that $\rho_{train}$ is fixed during training, so minimizing $h_{train}$
during training is equivalently minimizing $\sigma$ by Lemma 3.3

3\. Why we are often unlucky to have $\rho_{true}>\rho_{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 $\rho$ increases with data variability (see point 4 below), we have
$\rho_{true}>\rho_{train\\\ \\_set}$ almost surely. (2) Second, we don’t know
the value of $\rho_{true}$ and can only guess it. In practice, we often guess
a small value because training often diverges with large $\rho$ (as observed
in Foret et al. (2020); Chen et al. (2021)).

4\. Why $\rho$ increases with data variability. In Corollary 5.2.1, we assume
weight perturbation $\delta\sim\mathcal{N}(0,b^{2}I^{k})$. 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^{-(\frac{\rho}{\sqrt{2}b}-\sqrt{k})^{2}}]$. In order for the result to
hold with a fixed probability, $\rho$ must stay proportional to $b$, hence
$\rho$ also increases with the variability of data.

##  Appendix B Experimental Details

###  B.1 Training details

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
$\beta_{1}=0.9,\beta_{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\times$ 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 $\rho$ in {0.1, 0.2, 0.3, 0.4, 0.5, 0.6} for SAM and
GSAM; for ResNet, we search $\rho$ from 0.01 to 0.05 with a stepsize 0.01. For
ASAM, we amplify $\rho$ by $10\times$ compared to SAM, as recommended by Kwon
et al. (2021). For GSAM, we search $\alpha$ in {0.1, 0.2, 0.3} throughout the
paper. We report the best configuration of each individual model in Table. 4.

Table 4: Hyper-parameters to reproduce experimental results

Model | $\rho_{max}$ | $\rho_{min}$ | $\alpha$ | $lr_{max}$ | $lr_{min}$ | Weight Decay | Base Optimizer | Epochs | Warmup Steps | LR schedule
---|---|---|---|---|---|---|---|---|---|---
ResNet50 | 0.04 | 0.02 | 0.01 | 1.6 | 1.6e-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
ResNet512 | 0.04 | 0.02 | 0.005 | 1.6 | 1.6e-2 | 0.3 | SGD | 90 | 5k | Linear
ViT-S/32 | 0.6 | 0.0 | 0.4 | 3e-3 | 3e-5 | 0.3 | AdamW | 300 | 10k | Linear
ViT-S/16 | 0.6 | 0.0 | 1.0 | 3e-3 | 3e-5 | 0.3 | AdamW | 300 | 10k | Linear
ViT-B/32 | 0.6 | 0.1 | 0.6 | 3e-3 | 3e-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 | 3e-3 | 3e-5 | 0.3 | AdamW | 300 | 10k | Linear
Mixer-S/16 | 0.5 | 0.0 | 0.6 | 3e-3 | 3e-5 | 0.3 | AdamW | 300 | 10k | Linear
Mixer-S/8 | 0.5 | 0.1 | 0.1 | 3e-3 | 3e-5 | 0.3 | AdamW | 300 | 10k | Linear
Mixer-B/32 | 0.7 | 0.2 | 0.05 | 3e-3 | 3e-5 | 0.3 | AdamW | 300 | 10k | Linear
Mixer-B/16 | 0.5 | 0.2 | 0.01 | 3e-3 | 3e-5 | 0.3 | AdamW | 300 | 10k | Linear

###  B.2 Transfer learning experiments

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.

###  B.3 Experimental setup with ablation studies on data augmentation

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.

##  Appendix C Ablation studies and discussions

((a)) Performance of SAM and GSAM under different $\rho$.

((b)) Performance of GSAM under different $\alpha$

Figure 7: Performance of GSAM varying with $\rho$ and $\alpha$. Table 5: Top-1 accuracy of ViT-B/32 on ImageNet with Inception-style data augmentation. For vanilla training we 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
---|---|---|---|---|---
Vanilla | 300 | 71.4 | 77.5 | 57.5 | 23.4
600 | 72.0 | 78.2 | 57.9 | 23.6
GSAM | 300 | 76.8 | 82.7 | 63.0 | 25.1

###  C.1 Influence of $\rho$ and $\alpha$

We plot the performance of a ViT-B/32 model varying with $\rho$ (Fig. 7(a))
and $\alpha$ (Fig. 7(b)). We empirically validate that fine-tuning $\rho$ in
SAM can not achieve comparable performance with GSAM, as shown in Fig. 7(a).
Considering that GSAM has one more parameter $\alpha$, we plot the accuracy
varying with $\alpha$ in Fig. 7(b), and show that GSAM consistently
outperforms SAM and vanilla training.

Table 6: Top-1 Accuracy on ViT-B/32 on ImageNet. Ablation studies on constant $\rho$ or a decayed $\rho_{t}$. Vanilla | Constant $\rho$ (SAM) | Constant $\rho$ \+ ascent | Decayed $\rho_{t}$ | Decayed $\rho_{t}$ \+ ascent
---|---|---|---|---
72.0 | 75.8 | 76.2 | 75.8 | 76.8

###  C.2 Constant $\rho$ v.s. decayed $\rho_{t}$ schedule

Note that Thm. 5.1 assumes $\rho_{t}$ to decay with $t$ in order to prove the
convergence, while SAM uses a constant $\rho$ during training. To eliminate
the influence of $\rho_{t}$ schedule, we conduct ablation study as in Table.
6. The ascent step in GSAM can be applied to both constant $\rho$ or a decayed
$\rho_{t}$ schedule, and improves accuracy for both cases. Without ascent
step, constant $\rho$ and decayed $\rho_{t}$ achieve similar performance.
Results in Table. 6 implies that the ascent step in GSAM is the main reason
for improvement of generalization performance.

Figure 8: The value of $\operatorname{cos}\theta_{t}$ varying with training
steps, where $\theta_{t}$ is the angle between $\nabla f(w_{t})$ and $\nabla
f_{p}(w_{t})$ as in Fig. 2.

Figure 9: Surrogate gap curve under different $\alpha$ values.

###  C.3 Visualize the training process

In the proof of Thm. 5.3, our analysis relies on assumption that $\theta_{t}$
is small. We empirically validated this assumption by plotting
$\operatorname{cos}\theta_{t}$ in Fig. 9, where $\theta_{t}$ is the angle
between $\nabla f(w_{t})$ and $\nabla f_{p}(w_{t})$. Note that the cosine
value is calculated in the parameter space of dimension $8.8\times 10^{7}$,
and in high-dimensional space two random vectors are highly likely to be
perpendicular. In Fig. 9 the cosine value is always above 0.9, indicating that
$\nabla f(w_{t})$ and $\nabla f_{p}(w_{t})$ point to very close directions
considering the high dimension of parameters. This empirically validates our
assumption that $\theta_{t}$ is small during training.

We also plot the surrogate gap during training in Fig. 9. As $\alpha$
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 $\alpha$, indicating that the
training process gradually falls into local minimum in order to minimize the
training loss.

##  Appendix D Related works

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 (Müller et al., 2019) and mixup (Mü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).