S URROGATE G AP M INIMIZATION I MPROVES S HARPNESS -A WARE T RAINING
Anonymous authors
Paper under double-blind review
A BSTRACT
The recently proposed Sharpness-Aware Minimization (SAM) improves generalization by minimizing a perturbed loss deﬁned as the maximum loss within a neighborhood in the parameter space.
However, we show that both sharp and ﬂat minima can have a low perturbed loss, implying that SAM does not always prefer ﬂat minima.
Instead, we deﬁne 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 G ap Guided S harpness-
A ware M inimization (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).
I NTRODUCTION
Modern neural networks are typically highly over-parameterized and easy to overﬁt 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 inﬂuence 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 inﬂuence of the batch size (Keskar et al., 2016).
These studies have led to various modiﬁcations 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., 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 classiﬁcation), where w is the parameter.
SAM minimizes perturbed loss deﬁned as f p ( w ) (cid:44) 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 max- imum loss is low.
However, we show that a low perturbed loss f p could appear in both ﬂat 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 ﬁnd surrogate gap deﬁned as h ( w ) (cid:44) 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 G ap Guided S harpness
A ware M inimization (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 satisﬁed, we ﬁnd a generalizable model with good performance.
GSAM consists of two steps for each update: 1) descend gradient ∇ f p ( w ) to minimize the perturbed loss f p (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 ∇ f p ( w ) , i.e., ∇ f ( w ) = ∇ (cid:107) 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 f p because ∇ f ⊥ ( w ) ⊥ ∇ f p ( w ) by construction.
We summarize our contribution as follows:
• We deﬁne 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 classiﬁcation tasks with various neural architectures, including ResNets (He et al., 2016),
Vision Transformers (Dosovitskiy et al., 2020), and MLP-Mixers (Tolstikhin et al., 2021).
P RELIMINARIES 2. N OTATIONS
• f ( w ) :
A loss function f with parameter w ∈ R k , where k is the dimension of parameters.
• ρ t ∈ R :
A scalar value controlling the amplitude of perturbation at step t .
• (cid:15) ∈ R :
A small positive constant (to avoid division by 0, (cid:15) = 10 − 12 by default).
• w advt (cid:44) w t + ρ t ∇ f ( w t )
||∇ f ( w t ) ||
+ (cid:15) :
The solution to max || w (cid:48) − w t ||≤ ρ t f ( w (cid:48) ) when ρ t is small.
• f p ( w t ) (cid:44) max || δ ||≤ ρ t f ( w t + δ ) ≈ f ( w advt ) :
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 ρ t centered at w t .
When ρ t is small, by Taylor expansion, the solution to the maximization problem is equivalent to a gradient ascent from w t to w advt .
• h ( w ) (cid:44) f p ( w ) − f ( w ) :
The surrogate gap deﬁned as the difference between f p ( 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 ∈ R
k
At the t -th step, the noisy observation of the gradients ∇ f ( w t ) , ∇ f p ( w t ) of the original loss and perturbed loss, respectively.
• ∇ f
( w t ) =
∇ f (cid:107) ( w t ) + ∇ f ⊥ ( w t ) :
Decompose ∇ f ( w t ) into parallel component ∇ f (cid:107) ( w t ) and vertical component ∇ f ⊥ ( w t ) by projection ∇ f ( w t ) onto ∇ f p ( w t ) .
2. S HARPNESS -A WARE M INIMIZATION
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 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 signiﬁcant performance degradation.
SAM formulates the problem as:
[Equation]
where ρ is a predeﬁned 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 ρ , using Taylor expansion around w , the inner maximization in Eq.
1 turns into a linear constrained optimization with solution
[Equation]
As a result, the optimization problem of SAM reduces to
[Equation]
f ( w adv ) where w adv (cid:44) w + ρ
[Equation]
where (cid:15) 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).
T HE SURROGATE GAP MEASURES THE SHARPNESS AT A LOCAL MINIMUM 3. T HE 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 ﬂat.
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 inﬂuence 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 ﬁxed ρ , consider two local minima w 1 and w 2 , f p ( w 1 ) ≤ f p ( w 2 ) (cid:54) =
⇒ σ max ( w 1 ) ≤ σ max ( w 2 ) , 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 w 1 to w 3 , and suppose the corresponding loss surfaces are ﬂatter from w 1 to w 3 .
For some ﬁxed ρ , we plot the perturbed loss f p and surrogate gap h 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 ﬂatter 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. T HE SURROGATE GAP AGREES WITH SHARPNESS
We introduce the surrogate gap that agrees with sharpness, deﬁned as:
[Equation]
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 sufﬁciently small, then the approximation to the surrogate gap in Eq.
4 is always non-negative, h ( w ) =
f ( w adv ) − 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 ≈ 2 h ( w ∗ ) /ρ 2 .
The proof is in Appendix.
Lemma 3.2 tells that the surrogate gap is non-negative, and Lemma 3.
shows that the loss surface is ﬂatter as h gets closer to 0.
The two lemmas together indicate that we can ﬁnd a region with a ﬂat loss surface by minimizing the surrogate gap h ( w ) .
S URROGATE G AP G UIDED S HARPNESS -A WARE M INIMIZATION 4. G ENERAL IDEA : J OINTLY MINIMIZE THE PERTURBED LOSS AND SURROGATE GAP
Inspired by the analysis in Section 3, we propose Surrogate G ap Guided S harpness-
A ware
M inimzation (GSAM) to minimize both the perturbed loss f p and surrogate gap h : [Equation]
[Equation]
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 ﬂat surface.
A low perturbed loss implies low training losses within the neighborhood, and a ﬂat 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 min w f p ( w )+ λh ( w ) , where λ is some positive scalar.
One caveat when solving this weighted combination is the potential conﬂict between the gradients of the two terms, i.e., ∇ f p ( w ) and ∇ h ( w ) .
We illustrate this conﬂict by Fig.
2, where ∇ h ( w ) = ∇
f p ( w ) − ∇ f ( w ) (the grey dashed arrow) has a negative inner product with ∇ f p ( w ) and ∇ 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. G RADIENT DECOMPOSITION AND ASCENT FOR THE MULTI - OBJECTIVE OPTIMIZATION
Our primary goal is to minimize f p because otherwise a ﬂat solution of high loss is meaningless, and the minimization of h should not increase f p .
We propose to decompose ∇ f ( w t ) into components that are parallel and orthogonal to ∇ f p ( w t ) , respectively (see the black and blue arrows in Fig. 2):
[Equation]
The key is that updating in the direction of ∇ f ⊥ ( w t ) does not change the value of the perturbed loss f p ( w t ) because ∇ f ⊥ ⊥ ∇ f p by construction.
Therefore, we propose to perform an ascent step in the ∇ f ⊥ ( w t ) direction , which achieves 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 ) ).
The full GSAM Algorithm is shown in Algo.
1 and Fig.
2, where g ( t ) , g ( t ) p are noisy observations of ∇ f ( w t ) and ∇ f p ( w t ) , respectively, and g ( t ) (cid:107) , g ( t ) ⊥ are noisy observations of ∇ f (cid:107) ( w t ) and ∇ f ⊥ ( w t ) , 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 w t , step 1) takes a gradient ascent to w adv t followed by step 2) evaluating the gradient g ( t ) p at w adv t .
Step 3) projects g ( t ) onto g ( t ) p ,
[Figure] [Figure]
[Figure] [Figure]
[Figure]
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” ∇ f GSAMt (cid:44) g ( t ) p − αg ( t ) ⊥ to ﬁrst-order gradient optimizers such as SGD and Adam.
The ascent step along g ( t ) ⊥ does not harm convergence SAM demonstrates that minimizing f p 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 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.
The trajectory of GSAM is closer to the ridge and tends to ﬁnd a ﬂat 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 ( w t ) points to a sharp local solution and deviates from the ridge; in contrast, w advt is closer to the ridge and ∇ f ( w advt ) is closer to the ridge descent direction than ∇ f ( w t ) .
Note that ∇ f GSAMt and ∇ f ( w t ) always lie at different sides of ∇ f p ( w t ) by construction (see Fig. 2), hence ∇ f GSAMt pushes the trajectory closer to the ridge than ∇ f p ( w t ) does.
The trajectory of GSAM is like descent along the ridge and tends to ﬁnd ﬂat minima.
T HEORETICAL PROPERTIES OF GSAM 5. C ONVERGENCE 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 ≤ G, ∀ t ) of the true gradient ∇ f ( w t ) 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
[Equation]
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/ √ T ) for non-convex stochastic optimization, matching the convergence rate of ﬁrst-order gradient optimizers like Adam.
5. G ENERALIZATION OF GSAM
In this section, we show the surrogate gap in GSAM is provably lower than SAM’s, so GSAM is expected to ﬁnd 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 (cid:98) f ( w ) =m (cid:80) mi =1 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 Q .
Denote the prior distribution (independent of training) as P , then
[Equation]
Corollary 5.2.1.
Suppose perturbation δ is drawn from distribution δ ∼ N (0 , b 2 I k ) , δ ∈ R k , k is the dimension of w , then with probability at least (cid:16)
[Equation]
where C = (cid:98) f ( w ) is the empirical training loss, and (cid:98) h is the surrogate gap evaluated on the training set.
Corollary 5.2.1 implies that minimizing (cid:98) 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 simpliﬁed to L 2 regularization (Foret et al., 2020).
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 c 2 .
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 ∆ (cid:98) h t to denote the amount that the ascent step in GSAM decreases (cid:98) h for the t -th step.
Compared to SAM, the proposed method generates a total decrease in surrogate gap (cid:80)
Tt =1 ∆ (cid:98) h t , which is bounded by
[Equation]
[Table]
We provide proof in the appendix.
The lower-bound of (cid:80) Tt =1 ∆ (cid:98) 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.
E XPERIMENTS 6. 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 ρ 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 .
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 signiﬁcant improvements from the longer training (Table 5 in appendix).
We summarize the best hyper-parameters for each model in Appendix B.
[Figure] [Figure]
Note that our GSAM is described as
SAM+ascent( =GSAM ) for consistency.
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 .
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.
6. GSAM FINDS A MINIMUM WHOSE H ESSIAN 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 ﬂat minimum.
We empirically validate this in Fig.
As shown in the left subﬁgure, for some ﬁxed ρ , increasing α decreases the dominant value and improves generalization (test accuracy).
In the middle subﬁgure, we plot the dominant eigenvalues estimated by the surrogate gap, σ max ≈ 2 h/ρ 2 (Lemma 3.3).
In the right subﬁgure, 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 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.
6. C OMPARISON 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 ﬁnd the hyper-parameter “scope” from 0.0 and 0.9, and search for the
[Table] [Table]
[Figure]
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.
6. A DDITIONAL 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 ) + λh ( w ) as discussed in Sec. 4.
is sub-optimal because ∇ h ( w ) could conﬂict with ∇
f p ( 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( f p , h ) min( f p , h ) min( f p , h ) vs. min( f, h ) min( f, h ) min( f, h ) GSAM solves min( f p , h ) by descent in ∇ f p , decomposing ∇ f onto ∇ f p , and an ascent step in the orthogonal direction to increase f while keep f p intact.
Alternatively, we can also optimize min( f, h ) by descent in ∇ f , decomposing ∇ f p onto ∇ 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 min( f p , h ) mainly to make the minimal change to SAM.
GSAM beneﬁts transfer learning Using weights trained on ImageNet-1k, we ﬁnetune models with
SGD on downstream tasks including the CIFAR10/CIFAR100 (Krizhevsky et al., 2009), Oxfordﬂowers (Nilsback & Zisserman, 2008) and Oxford-IITPets (Parkhi et al., 2012).
Results in Tableshows 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 subﬁgure).
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 subﬁgure) shows that GSAM performs the best, and SAM improves over vanilla training.
C ONCLUSION
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.
E THICS S TATEMENT
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 speciﬁc models that are trained using our new optimization strategy.
R EPRODUCIBILITY S TATEMENT
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 sufﬁcient details to ensure reproducibility.
We plan to release the models trained by using GSAM upon publication.
A P ROOFS A. P ROOF OF L EMMA . 3.
Suppose ρ is small, perform Taylor expansion around the local minima w , we have: [Equation]
[Equation]
where H is the Hessian, and is positive semideﬁnite at a local minima.
At a local minima, ∇ f ( w ) = 0 , hence we have
[Equation]
11) and
[Equation]
12) where σ max is the dominate eigenvalue (eigenvalue with the largest absolute value).
Now consider two local minima w 1 and w 2 with dominate eigenvalue σ 1 and σ 2 respectively, we have
[Equation]
We have f p ( w 1 ) >
f p ( w 2 ) (cid:54) = ⇒ σ 1 > σ 2 and σ 1 > σ [Equation]
⇒ f p [Equation] w 1 ) >
f p ( w 2 ) because the relation between f ( w 1 ) and f ( w 2 ) is undetermined.
[Equation]
A. P ROOF OF L EMMA . 3.
Since ρ is small, we can perform Taylor expansion around w ,
[Equation]
13) where the last line is because δ is approximated as δ
= ρ
[Equation]
||∇ f ( w ) || 2 + (cid:15) , hence has the same direction as ∇ f ( w ) .
(cid:3)
A. P ROOF OF L EMMA . 3.
Since ρ is small, we can approximate f ( w ) with a quadratic model around a local minima w :
[Equation]
where H is the Hessian at w , assumed to be positive semideﬁnite at local minima.
Normalize δ such that
||
δ ||
= ρ , Hence we have:
[Equation]
(14) where σ max is the dominate eigenvalue of the hessian H , and ﬁrst order term is 0 because the gradient is 0 at local minima.
Therefore, we have σ max ≈ 2 h ( w ) /ρ 2 .
(cid:3)
A. P ROOF OF T HM . 5.
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.
[Equation]
TEP 1:
[Equation]
ONVERGENCE W . R . T FUNCTION f p (
[Equation]
By L − smoothness of f ( w ) , we have
[Equation]
Hence we can derive the smoothness of f p
[Equation]
w )
[Equation]
Lipschitz
[Equation]
continuous with constant
[Equation]
Hence we prove the smoothness of f p .
For simplicity, we denote L p =
L + L 2 ρ(cid:15) as the Lipschitz constant of f p ( w ) .
For simplicity of notation, we denote the update at step t as
[Equation]
[Equation]
By smoothness of f p ( w ) , we have
[Equation]
Take expectation conditioned on observation up to step t (for simplicity of notation, we use E short for E x to denote expectation over all possible data points) conditioned on observations up to step t , we have
[Equation]
(25) (cid:16) Since E g ( t ) ⊥ is orthogonal to ∇ f p ( w t ) by construction , || g ( t ) || ≤ G by assumption (cid:17)
Hence we have
[Equation]
By telescope sum, we have:
[Equation] [Equation]
Since η t = η 0 / √ t , we have
[Equation]
Hence
[Equation]
where C 1 , C 2 are some constants.
This implies that the regret w.r.t f p ( w ) converges at rate
O (log T/ √ T ) .
S TEP
[Equation] C ONVERGENCE W [Equation]
R . T
[Equation]
FUNCTION
[Equation] [Equation] We prove the risk for f ( w ) convergences for non-convex stochastic optimization case using SGD.
Denote the update at step t as
[Equation]
By smoothness of f , we have
[Equation]
For simplicity, we introduce a scalar β t such that
[Equation]
[Equation]
where ∇ f (cid:107) ( w t ) is the projection of ∇ f ( w t ) onto ∇ f p ( w t ) .
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.
35, we have:
E f
[Equation]
t +
[Equation] ≤ [Equation] [Equation]
2 2 + (cid:16) 1 β t + α (cid:17) η t (cid:68) ∇ f ( w t ) , ∇ f
[Equation]
( w t ) (cid:69) + L 2 E (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) d t (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) 2(38) = f ( w t ) − η t [Equation] (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) ∇ f ( w t ) (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
(cid:16) 1 β [Equation] (cid:17) η [Equation]
(cid:68)
[Equation]
( w
[Equation]
[Equation]
( w t ) sin θ
[Equation]
(cid:69) + L 2 E (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) d t (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) 2(39) (cid:16) θ t is {+the angle between ∇ f p ( w t ) and ∇ f ( w t ) (cid:17) = f ( w t ) − η t β t (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) ∇ f ( w t ) (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) 2 
 + (cid:16) 1 β t + α (cid:17) η t (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) ∇ f ( w t ) (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) [Equation]+}
(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation] [Equation] [Equation] where δ t = ρ t ∇ f ( w t )
||∇ f ( w t ) || 2 by deﬁnition, H ( w t ) is the Hessian.
Hence we have
[Equation]
where L is the Lipschitz constant of f [Equation] and L − smoothness of f indicates the maximum absolute eigenvalue of H is upper bounded by L .
Plug Eq.
42 into Eq.
40, we have
[Equation]
η t (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation] E (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) d t (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) 2(44) ≤ f [Equation] (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
(cid:16)
[Equation]
(cid:17)
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
t (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation] (45) (cid:16) Assume gradient has bounded norm G. (cid:17) (46) ≤ [Equation] w
[Equation]
−
[Equation] (cid:12)(cid:12) [Equation] (47) (cid:16) β t is close to 1 assuming ρ is small, hence it’s natural to assume 0 < β min ≤ β t ≤ β max (cid:17)
Re-arranging above formula, we have
[Equation]
β
[Equation]
48) perform telescope sum and taking expectations on each step, we have β max
[Equation]
[Equation]
Take the schedule to be η t = η 0 √ t and ρ t = ρ 0 √ t , then we have ηβ max
[Equation]
[Equation] [Equation] [Equation]
[Equation]
Hence
[Equation]
(54) where C 1 , C 4 are some constants.
This implies the convergence rate w.r.t f ( w ) is O (log T/ √ T ) .
S TEP 3: C ONVERGENCE 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
[Equation]
Hence
[Equation]
56) also converges at rate O (log T/ √ T ) because each item in the RHS converges at rate O (log T √ T ) .
(cid:3)
A. P ROOF OF C OROLLARY . 5.2.
Using the results from Thm.
5.2, with probability at least 1 − a , we have
[Equation] [Equation]
Assume δ ∼ N (0 , b 2 I k ) 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
[Equation]
(58) hence with probability at least 1 − 1 / √ n , we have
[Equation]
Therefore, with probability at least 1 − 1 / √ n = 1 − exp
[Equation]
Combine Eq.
58 and Eq.
60, subtract the same constant C on both sides, and under the same assumption as in (Foret et al., 2020)
that E w ∼Q E x f ( w, x ) ≤ E δ ∼N (0 ,b 2 I k ) E w ∼Q E x f ( w + δ, x ) we ﬁnish the proof.
[Equation]
A. P ROOF OF T HM . 5. S TEP 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
[Equation]
61) (cid:16) where E g ⊥ = ∇ f ⊥ ( w t ) (cid:17) =
[Equation]
where θ t is the angle between vector ∇ f p ( w t ) and ∇ f ( w t ) .
The expected change of loss gap caused by ascent step is
[Equation]
Above results demonstrate that ascent step decreases the loss gap, while descent step might increase the loss gap.
A sufﬁcient (but not necessary) condition for E (cid:104)∇ h ( w t ) , dt (cid:105) ≤ 0 requires α to be large or | (cid:12)(cid:12) ∇ f ( w t ) (cid:12)(cid:12)(cid:12)(cid:12) 2 cos θ t ≤ (cid:12)(cid:12)(cid:12)(cid:12) ∇ f p ( w t ) (cid:12)(cid:12)(cid:12)(cid:12) .
In practice, the perturbation amplitude ρ is small and we can assume θ t is close to 0 and (cid:12)(cid:12)(cid:12)(cid:12) ∇ f p ( w t ) (cid:12)(cid:12)(cid:12)(cid:12) is close to (cid:12)(cid:12)(cid:12)(cid:12) ∇ f ( w t ) (cid:12)(cid:12)(cid:12)(cid:12) , we can also set the parameter α to be large in order to decrease the loss gap.
Next we give an estimate of the decrease in (cid:98) h caused by our ascent step.
We refer to Eq.
62 and
Eq.
63 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:
[Equation]
[Equation]
Hence we derive an upper bound for (cid:80)
Tt =1 ∆ (cid:98) h t .
Next we derive a lower bound for (cid:80) Tt =1 ∆ (cid:98) h t Note that when ρ t is small, by Taylor expansion ∇ (cid:98)
[Equation]
where (cid:98) H ( w t ) is the Hessian evaluated on training samples.
Also when ρ t is small
[Equation]
the angle θ t between ∇ (cid:98) f p ( w t ) and ∇ (cid:98) f ( w t ) is small [Equation] by the limit that [Equation]
We have
[Equation]
Omitting high order term, we have
[Equation]
where G is the upper-bound on norm of gradient [Equation] | σ | 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
[Equation]
where c 2 is the lower bound of ||∇ (cid:98) f
||
(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 efﬁciently improves generalization compared with SAM.
[Equation]
A. P ROOF FOR CONVERGENCE OF GSAM WITHOUT RELYING ON THE L- SMOOTHNESS OF f p
In this section, we provide a proof for the convergence of GSAM without replying on the Lsmoothness of f p .
By L − smoothness of f and the deﬁnition of f p ( w t ) =
f ( w advt )
[Equation]
and deﬁnition of d t = w t +1 − w t and w advt =
w t + δ t we have
[Equation]
It’s trivial to see that Eq.
78 has the same form as Eq. 22, except the constant L p is replaced with a constant 2 L which is even better bounded, so we can reuse RHS of Eq.
by replacing L p with 2 L .
Now we focus on the difference caused by Eq. 79.
By deﬁnition of δ t , we have
[Equation]
where g ( t ) is the gradient of f at w t evaluated with a noisy data sample.
When learning rate η t is small, the update in weight d t is small
[Equation]
and expected gradient is
[Equation]
where H is the Hessian at w t .
Therefore
[Equation]
we have
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) (84) ≤
[Equation]
(85) where the ﬁrst inequality is due to (1) ρ t is monotonically decreasing with t , and (2) triangle inequality that (cid:104) a, b (cid:105) ≤ || a || · || b || .
φ t is the angle between the unit vector in the direction of ∇ f ( w t ) and ∇ f ( w t +1 ) .
The second inequality comes from that (1) (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
g
|| g
||
+ (cid:15) (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) < 1 strictly, so we can replace δ t in Eq.
83 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 d t is small, φ t is also small.
Using the limit that
[Equation]
We have:
[Equation]
[Equation]
[Equation]
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 ||
=
||
η t ( g (cid:107) + αg ⊥ ) ||
and E g t = ∇ f ( w t ) .
Plug into Eq. 85, also note that the perturbation amplitude ρ t is small so w t is close to w advt , then we have
[Equation]
Similarly, we have
[Equation]
Reuse results from Eq.
{+25 (replace L p with 2 L ) and plug into Eq. 78, and plug Eq. 
+}
and Eq.into Eq.
79, we have
[Equation]
Perform telescope sum, we have
[Equation]
Hence η T
[Equation]
95) where
[Equation]
Note that η T = η 0 √ T , we have
[Equation]
) which implies that GSAM converges at a rate of O (log T/ √ T ) , and all the constants here are wellbounded.
[Table]
B E XPERIMENTAL D ETAILS B. T RAINING DETAILS
For ViT and Mixer, we search the learning rate in { 1e-3, 3e-3, 1e-2, 3e-3 }
[Table]
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 }
[Table]
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 .
We train ResNets forepochs, 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 ﬁnd signiﬁcant improvement as in Table.
We ﬁrst search the optimal learning rate and weight decay for vanilla training, and keep these two hyper-parameters ﬁxed 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.
For
GSAM, we search α in { 0.1, 0.2, 0.3 } throughout the paper.
We report the best conﬁguration of each individual model in Table.
B. T RANSFER LEARNING EXPERIMENTS
Using weights trained on ImageNet-1k, we ﬁnetune models with SGD on downstream tasks including the CIFAR10/CIFAR100 (Krizhevsky et al., 2009), Oxford-ﬂowers (Nilsback & Zisserman, 2008) and Oxford-IITPets (Parkhi et al., 2012).
For all experiments, we use the SGD optimizer with no weight decay under a cosine 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. E XPERIMENTAL 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 subﬁgure of Fig. 6, “Light” refers to Inception-style data augmentation with random ﬂip 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.
=0. =0.
=0. =0.
=0. =0.
T o p 1 - A cc u r a c y ( % )
Top-1 accuracy of ViT-B/32 under different
SAM GSAM (a) Performance of SAM and GSAM under different ρ .
=0.
=0. =0.
=0. =0.
T o pA cc u r a c y ( % )
Top-1 accuracy of ViT-B/32 varying with
Vanilla SAM (b) Performance of GSAM under different
α
[Table]
C A BLATION STUDIES AND DISCUSSIONS C. I NFLUENCE OF ρ AND α
We plot the performance of a ViT-B/32 model varying with ρ (Fig. 7a) and α (Fig.
7b).
We empirically validate that ﬁne-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.
C ONSTANT ρ 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 inﬂuence of ρ t schedule, we conduct ablation study as in Table.
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. V ISUALIZE 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 ( w t ) and ∇ f p ( w t ) .
Note that the cosine value is calculated in the parameter space of dimension 8 .
8 × 10 7 , 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 ( w t ) and ∇ f p ( w t ) 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.
As α increases, the surrogate gap decreases, validating that the ascent step in GSAM efﬁciently minimizes the surrogate gap.
Furthermore, the surrogate gap increases with training steps for any ﬁxed α , indicating that the training process gradually falls into local minimum in order to minimize the training loss.
[Table]
[Figure] [Figure]
Influence of
[Figure]
surrogate
[Figure]
D R ELATED 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 modiﬁed 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 ﬂipping or rotation of images, recent data augmentations include label smoothing (M ¨ uller et al., 2019) and mixup (M ¨ uller 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 ﬁrst-order gradient optimizers, and both techniques can be combined with GSAM.
Besides SGD, Adam and AdaBelief, GSAM can be combined with other ﬁrst-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).