A N
A DAPTIVE
P OLICY
TO E MPLOY S HARPNESSA WARE M INIMIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Sharpness-aware minimization (SAM), which searches for ﬂat minima by min-max optimization, has been shown to be useful in improving model generalization.
However, since each SAM update requires computing two gradients, its computational cost and training time are both doubled compared to standard empirical risk minimization (ERM).
Recent state-of-the-arts reduce the fraction of SAM updates and thus accelerate SAM by switching between SAM and ERM updates randomly or periodically.
In this paper we design an adaptive policy to employ
SAM based on the loss landscape geometry.
Two efﬁcient algorithms, AE-SAM and AE-LookSAM, are proposed.
We theoretically show that AE-SAM has the same convergence rate as SAM.
Experimental results on various datasets and architectures demonstrate the efﬁciency and effectiveness of the proposed method
I NTRODUCTION
Despite great success in many applications (He et al., 2016; Zagoruyko & Komodakis, 2016; Han et al., 2017), deep networks are often over-parameterized and capable of memorizing all training data.
The training loss landscape is complex and nonconvex with many local minima of different generalization abilities.
Many studies have investigated the relationship between geometry of the loss surface and generalization performance (Hochreiter & Schmidhuber, 1994; McAllester, 1999; Keskar et al., 2017; Neyshabur et al., 2017; Jiang et al., 2020), and found that ﬂatter minima generalize better than sharper minima (Dziugaite & Roy, 2017; Petzka et al., 2021; Chaudhari et al., 2017; Keskar et al., 2017; Jiang et al., 2020).
Sharpness-aware minimization (SAM) (Foret et al., 2021) is the current state-of-the-art to seek ﬂat minima by solving a min-max optimization problem.
In the SAM algorithm, each update consists of two forward-backward computations: one for computing the perturbation and the other for computing the actual update direction.
Since these two computations are not parallelizable, SAM doubles the computational overhead as well as the training time compared to empirical risk minimization (ERM).
Several algorithms (Du et al., 2022a; Zhao et al., 2022b; Liu et al., 2022) have been proposed to improve SAM’s efﬁciency.
ESAM
(Du et al., 2022a) uses fewer samples to compute the gradients and updates fewer parameters, but each update still requires two gradient computations.
Thus,
ESAM does not alleviate the bottleneck of training speed.
Instead of using the SAM update at every iteration, recent state-of-the-arts (Zhao et al., 2022b; Liu et al., 2022) proposed to use SAM randomly or periodically.
Speciﬁcally, SS-SAM (Zhao et al., 2022b) selects SAM or ERM according to a
Bernoulli trial, while LookSAM (Liu et al., 2022) employs SAM at every k step.
Though efﬁciency, random or periodic policy of using SAM is suboptimal as it is not geometry-aware.
Intuitively, the
SAM update is more necessary for sharp regions than for ﬂat regions.
In this paper, we propose an adaptive policy to employ SAM based on the geometry of the loss landscape.
The SAM update is used when the model is in sharp regions, while the ERM update is used in ﬂat regions for reducing the fraction of SAM updates.
To measure sharpness, we use the squared stochastic gradient norm and model it as normal distribution.
We estimate the distribution parameters by exponential moving average.
Experimental results on standard benchmarks demonstrate the superiority of the proposed policy.
Our contributions are summarized as follows: (i)
We propose an adaptive policy to employ SAM based on the loss landscape geometry.
(ii)
We propose an efﬁcient algorithm, called AE-SAM, to reduce the fraction of SAM updates.
We also theoretically study its convergence rates with both batch and stochastic gradient descent.
(iii)
The proposed policy is general and can be combined with any SAM variant.
In this paper, we integrate it with LookSAM and propose AE-LookSAM.
(iv) Experimental results on various datasets and network architectures, with and without label noise, verify the superiority of AE-SAM and AE-LookSAM over existing baselines.
Notations .
Vectors (e.g., x ) and matrices (e.g., X ) are denoted by lowercase and uppercase boldface letters, respectively.
For a vector x , its (cid:96) 2 -norm is (cid:107) x
(cid:107) .
N ( µ ; σ 2 ) is the univariate normal distribution with mean µ and variance σ 2 .
diag( x ) constructs a diagonal matrix with x on the diagonal.
Let I
{+A ( x ) be an indicator function, i.e., I A ( x ) = 
+}
if x ∈ A otherwise 0 .
R ELATED W ORK
We are given a training set D with i.i.d.
samples { ( x i , y i ) : i = 1 , . . . , n } drawn from a data distribution, Let f ( x ; w ) be a model parameterized by w .
Its empirical risk on D is L ( D ; w ) =n (cid:80) ni
=1 (cid:96) ( f ( x i ; w ) , y i ) , where (cid:96) ( · , · ) is a loss (e.g., cross-entropy loss for classiﬁcation).
Model training aims to learn a model from the training data that generalizes well on test data.
Generalization and Flat Minima.
The connection between model generalization and loss landscape geometry has been theoretically and empirically studied in (Keskar et al., 2017; Dziugaite & Roy, 2017; Jiang et al., 2020).
Recently, Jiang et al.
(2020) conducted large-scale experiments and found that sharpness-based measures (ﬂatness) are related to generalization of minimizers.
Although ﬂatness can be characterized by Hessian eigenvalues (Keskar et al., 2017; Dinh et al., 2017), dealing with
Hessian explicitly is computationally prohibitive.
To address this issue, practical algorithms propose to seek ﬂat minima by injecting noise in the optimizers (Zhu et al., 2019;
Zhou et al., 2019; Orvieto et al., 2022; Bisla et al., 2022), introducing regularization (Chaudhari et al., 2017; Zhao et al., 2022a;
Du et al., 2022b), averaging model weights during training (Izmailov et al., 2018;
He et al., 2019; Cha et al., 2021), or sharpness-aware minimization (SAM) (Foret et al., 2021; Kwon et al., 2021;
Zhuang et al., 2022; Kim et al., 2022).
SAM.
The state-of-the-art SAM (Foret et al., 2021) and its variants (Kwon et al., 2021; Zhuang et al., 2022; Kim et al., 2022) search for ﬂat minima by solving the min-max optimization problem:
[Equation]
1) where ρ > 0 is the radius of perturbation.
The above can be rewritten as min w L ( D ; w ) + R
( D ; w ) , where R ( D ; w ) ≡ max (cid:107) (cid:15) (cid:107)≤ ρ L ( D ; w + (cid:15) )
− L ( D ; w ) is a regularizer that penalizes sharp minimizers (Foret et al., 2021).
As solving the inner maximization in Eq.(1) exactly is computationally infeasible for nonconvex optimization, SAM approximately solves it by the ﬁrst-order Taylor approximation and its update rule is
[Equation]
where B t is a mini-batch of data, η is the step size, and ρ t = ρ
[Table]
B t ; w t ) [Table] Although SAM has shown to be effective in improving the generalization of deep networks, a major drawback is that each update in Eq.(2) requires two forward-backward calculations.
Speciﬁcally, SAM ﬁrst calculates the gradient of L ( B t ; w ) at w t to obtain the perturbation, then calculates the gradient of L ( B t ; w ) at w t + ρ t ∇L ( B t ; w t ) to obtain the update direction for w t .
As a result, SAM doubles the computational overhead compared to ERM.
Efﬁcient Variants of SAM.
Several algorithms have been proposed to accelerate the SAM algorithm.
ESAM (Du et al., 2022a) uses fewer samples to compute the gradients and only updates part of parameters in the second step, but still requires to compute most of the gradients.
Another direction is to reduce the number of SAM updates used in the training process.
SS-SAM (Zhao et al., 2022b) randomly selects SAM or ERM update according to a Bernoulli trial, while LookSAM (Liu et al., 2022) employ SAM at every k iteration.
Unlike the random or periodic policy of employing SAM, the proposed policy is adaptive to the loss landscape.
M ETHOD
In this section, we propose an adaptive policy to employ the SAM update.
The general idea is to use
ERM when w t is already in a ﬂat region, and use SAM only when the loss landscape is locally sharp.
We start by introducing a sharpness measure (Section 3.1), then propose an adaptive policy based on the sharpness measure (Section 3.2).
Next, we propose two algorithms (AE-SAM and AE-LookSAM) and study the convergence of AE-SAM.
[Figure]
a
[Figure]
i an
[Figure] [Figure]
3. S HARPNESS M EASURE
Though sharpness can be characterized by the eigenvalues of Hessian (Keskar et al., 2017; Dinh et al., 2017), it is expensive to compute.
A widely-used approximation is based on the gradient magnitude diag([ ∇L ( B ; w t )] 2 ) (Bottou et al., 2018; Khan et al., 2018), where [ v ] 2 denotes the elementwise square of v .
As (cid:107)∇L ( B ; w t ) (cid:107) 2 equals the trace of diag([ ∇L ( B ; w t )] 2 ) , it is reasonable to choose (cid:107)∇L ( B ; w t ) (cid:107) 2 as a sharpness measure.
(cid:107)∇L ( B ; w t ) (cid:107) 2 is also related to another sharpness measure, namely, the variance of gradients, which is deﬁned as
[Equation]
With appropriate smoothness assumptions on the loss L , both SAM and ERM can be shown theoretically to converge to critical points of L ( D ; w ) (i.e., ∇L ( D ; w ) = 0 )
(Reddi et al., 2016;
Andriushchenko & Flammarion, 2022).
Thus, it follows from Eq.(3) that the squared stochastic gradient norm E B (cid:107)∇L ( B ; w ) (cid:107) 2 indicates the variance of gradients.
Jiang et al.
(2020) conducted extensive experiments and empirically shows that the variance of gradients is positively correlated with the generalization gap.
In other words, the smaller the Var ( ∇L ( B ; w )) , the better generalization is the model with parameter w .
This ﬁnding also explains why SAM generalizes better than ERM.
Figure 1 shows the variance of gradients w.r.t.
the number of epochs using SAM and ERM on
CIFAR-100 with various network architectures (experimental details are in Section 4.1).
As can be seen, SAM always has a much smaller variance than ERM.
Figure 2 shows the squared stochastic gradient norm w.r.t.
the number of epochs on CIFAR-100 .
As can be seen, SAM achieves a much smaller E B (cid:107)∇L ( B ; w t ) (cid:107) 2 than ERM.
3. A DAPTIVE P OLICY TO E MPLOY SAM
As E B (cid:107)∇L ( B ; w t ) (cid:107) 2 changes over t (Figure 2), it needs to be estimated at every iteration.
One can sample a large number of mini-batches and compute the mean of the corresponding stochastic gradient norms.
However, this can be computationally expensive as the computational overhead is substantial.
To address this problem, we model (cid:107)∇L ( B ; w t ) (cid:107) 2 with a distribution and estimate the distribution parameters in an online manner.
Figure 3(a) shows (cid:107)∇L ( B ; w t ) (cid:107) 2 of 400 mini-batches at different training stages (epoch = 60 , 120 , and 180 ) on CIFAR-100 using ResNet-18 .
As can be seen, the distribution follows a Bell curve.
Figure 3(b) shows the corresponding quantile-quantile (Q-Q) plot, which is usually used to assess normality: the closer the curve to the line, the distribution is closer to the normal distribution.
Both Figures 3(a) and 3(b) suggest that (cid:107)∇L ( B ; w t ) (cid:107) 2 can be modeled with a normal distribution N ( µ t , σ 2 t ) .
We estimate the mean and variance by an exponential moving average (EMA):
[Equation]
(5) where δ ∈ (0 , 1) controls the forgetting rate.
Empirically, we use δ = 0 .
Since ∇L ( B t ; w t ) is already available, EMA does not involve any additional gradient calculations (the cost for the norm operator is negligible).
[Figure]
Using µ t and σ 2 t , we employ SAM only at iterations where (cid:107)∇L ( B t ; w t ) (cid:107) 2 is relatively large.
Speciﬁcally, when (cid:107)∇L ( B t ; w t ) (cid:107) 2 ≥ µ t + c t
σ t (where c t is a threshold), SAM is used; otherwise,
ERM is used.
When using the SAM update is more desirable, c t is small (or even negative).
When c t → −∞ , it reduces to SAM; when c t → ∞ , it becomes ERM.
Note that during the early training stage, the model is still underﬁtting and w t is far from the region of ﬁnal convergence.
Thus, minimizing the empirical loss is more important than seeking a locally ﬂat region.
Andriushchenko & Flammarion (2022) also empirically observe that the SAM update is more effective in boosting performance towards the end of training.
We therefore design a schedule that linearly decreases c t from λ 2 to λ 1 (which are pre-set values): c t = g λ 1 ,λ 2 ( t ) ≡ tT λ 1 + (cid:0) 1 − tT (cid:1) λ 2 , where T is the total number of iterations.
The whole procedure, called Adaptive policy to Employ SAM ( AE-SAM ), is summarized in Algorithm 1.
AE-LookSAM .
The proposed adaptive policy can be combined with any SAM variant.
Here, we consider integrating it with LookSAM (Liu et al., 2022).
When (cid:107)∇L ( B t ; w t ) (cid:107) 2 ≥ µ t + c t σ t , SAM is used and the update direction for w t is decomposed into two orthogonal directions as in LookSAM:
(i) the ERM update direction to reduce training loss, and (ii) the direction that biases the model to a ﬂat region.
When (cid:107)∇L ( B t ; w t ) (cid:107) 2 < µ t + c t
σ
t
, ERM is performed and the second direction of previous SAM update is reused to compose an approximate SAM direction.
The procedure, called
AE-LookSAM, is also shown in Algorithm 1. Algorithm 1 AE
Require: training set D , stepsize η , radius ρ ; λ 1 and λ 2 for g λ 1 ,λ 2 ( t ) ; w 0 , µ − 1 = 0 , σ 2 − 1 =
e − 10 , and α for AE-LookSAM; 1: for t = 0 , . .
, T − 1 do 2: sample a mini-batch data B t from D ; 3: compute g =
∇L ( B t ; w t ) ; 4: update µ t by (4) and σ 2 t by (5); 5: compute c t = g λ 1 ,λ 2 ( t ) ; 6: if (cid:107)∇L ( B t ; w t ) (cid:107) 2 ≥ µ t + c t σ t then 7: g s =
∇L ( B t ; w t + ρ ∇L ( B t ; w t )) ; 8: if AE-LookSAM: decompose g s as g v =
g s − g (cid:62) g s (cid:107) g (cid:107) 2 g ; 9: else : 10: if AE-SAM: g s = g ; 11: if AE-LookSAM: g s = g + α (cid:107) g (cid:107) (cid:107) g v (cid:107) g v ; 12: end if 13: w t +1 =
w t − η g s ; 14: end for 15: return w T .
3. C ONVERGENCE A NALYSIS
In this section, we study the convergence of any algorithm that mixes SAM and ERM.
Due to the mixing scheme, analyzing convergence is more challenging compared with that of SAM. Let ξ t be an indicator of whether SAM or ERM is used at iteration t (i.e., ξ t = 1 for SAM, and 0 for ERM). For example, ξ t =
I { w : (cid:107)∇L ( B t ; w ) (cid:107) 2 ≥ µ t + c t σ t } ( w t ) for the proposed AE-SAM, and ξ t is sampled from a Bernoulli distribution for SS-SAM (Zhao et al., 2022b).
The following assumptions on smoothness and bounded variance of stochastic gradients are standard in the literature on non-convex optimization (Ghadimi & Lan, 2013; Reddi et al., 2016) and
SAM (Andriushchenko & Flammarion, 2022; Abbas et al., 2022; Qu et al., 2022).
All proofs are in Appendix A.
Assumption 3.1 (Smoothness) .
L ( D ; w ) is β -smooth in w , i.e., (cid:107)∇L ( D ; w ) − ∇L ( D ; v ) (cid:107) ≤ β (cid:107) w − v (cid:107) .
Assumption 3.2 (Bounded variance of stochastic gradients) .
E ( x i ,y i ) ∼D (cid:107)∇ (cid:96) ( f ( x i ; w ) , y i ) − ∇L ( D ; w ) (cid:107) 2 ≤ σ 2 .
Theorem 3.3.
Let b be the mini-batch size.
If η = 1 4 β √ T and ρ = 1 /T 14 , algorithms mixing SAM and ERM satisfy
[Equation]
{+where ζ = 1 T (cid:80) T − 1 t =
+}
ξ t ∈ [0 , 1] .
Note that ζ is the fraction of SAM updates, and a larger ζ leads to a larger upper bound in Eq.(6).
The above theorem can also be used on SAM (which corresponds to ζ = 1 ), and is shown in Corollary A.2 of Appendix A.1.
E XPERIMENTS
In this section, we evaluate the proposed AE-SAM and AE-LookSAM on several standard benchmarks.
As the SAM update doubles the computational overhead compared to the ERM update, the training speed is mainly determined by how often the SAM update is used.
Hence, we evaluate efﬁciency by measuring the fraction of SAM updates used:
[Equation]
where T , the total number of iterations, is the same for all methods.
4. CIFAR-10 AND CIFARSetup.
In this section, experiments are conducted on the CIFAR-10 and CIFARdatasets (Krizhevsky & Hinton, 2009) using three network architectures: ResNet-18 (He et al., 2016), WideResNet-28-10 (denoted WRN-28-10 ) (Zagoruyko & Komodakis, 2016), and PyramidNet110 (Han et al., 2017).
Following the setup in (Liu et al., 2022; Foret et al., 2021), we use batch size 128 , initial learning rate of 0 .
1 , cosine learning rate schedule, SGD optimizer with momentum 0 .
9 and weight decay 0 . 0001 .
The number of training epochs is 300 for PyramidNet-110 , and 200 for ResNet-
and
WideResNet-28-10 .
10% of the training set is used as the validation set.
We perform grid search on the radius ρ over { 0 . 01 , 0 . 02 , 0 .
5 } as in Foret et al.
Similarly, α is selected by grid search over { 0 . 1 , 0 . 3 , 0 . 6 , 0 . 9 } .
For the c t schedule g λ 1 ,λ 2 ( t ) , λ 1 = − 1 and λ 2 = 1 for
AE-SAM; λ 1 = 0 and λ 2 = 2 for AE-LookSAM.
Baselines.
The proposed AE-SAM and AE-LookSAM are compared with the following baselines: (i)
[Equation]
ii
[Equation]
(Foret et al., 2021); and its more efﬁcient variants including (iii) ESAM (Du et al., 2022a) which uses part of the weights to compute the perturbation and part of the samples to compute the SAM update direction.
These two techniques can reduce the computational cost, but may not always accelerate SAM, particularly in parallel training (Li et al., 2020); (iv) SS-SAM (Zhao et al., 2022b), which randomly selects SAM or ERM according to a Bernoulli trial with success probability 0 .
This is the scheme with the best performance in (Zhao et al., 2022b); (v) LookSAM
(Liu et al., 2022) which uses SAM at every k = 5 step.
The experiment is repeated ﬁve times with different random seeds.
Results.
Table 1 shows the testing accuracy and fraction of SAM updates (%SAM), in which methods are grouped based on %SAM.
As can be seen, AE-SAM has higher accuracy than SAM while using only 50% of SAM updates.
SS-SAM and AE-SAM have comparable %SAM (about 50% ), and AESAM achieves higher accuracy than SS-SAM (which is statistically signiﬁcant based on the pairwise t-test at 95% signiﬁcance level).
Finally, LookSAM and AE-LookSAM have comparable %SAM (about 20% ), and AE-LookSAM also has higher accuracy than LookSAM.
These improvements conﬁrm the adaptive policy is better.
4. ImageNet
Setup.
In this section, we perform experiments on the ImageNet (Russakovsky et al., 2015), which contains 1000 classes and 1 .
28 million images.
The ResNet-50 (He et al., 2016) is used.
Following the setup in Du et al.
(2022a), we train the network for 90 epochs using a SGD optimizer with momentum 0 .
9 , weight decay 0 . 0001 , initial learning rate 0 .
1 , cosine learning rate schedule, and batch size 512 .
As in (Foret et al., 2021; Du et al., 2022a), ρ = 0 . 05 .
For the c t schedule g λ 1 ,λ 2 ( t ) , λ [-
-]1 =
Results.
Table 2 shows the testing accuracy and fraction of SAM updates.
As can be seen, with only half of the iterations using SAM, AE-SAM achieves comparable (or even better) performance as
SAM.
Compared with LookSAM, AE-LookSAM has better performance (which is also statistically signiﬁcant), verifying the proposed adaptive policy is more effective than LookSAM’s periodic policy.
[Table]
[Table] [Table]
4. R OBUSTNESS TO L ABEL N OISE
Setup.
In this section, we study whether the more-efﬁcient SAM variants will affect its robustness to training label noise.
Following the setup in Foret et al.
[Table]
we conduct experiments on a corrupted version of CIFAR-
[Table]
with some of its training labels randomly ﬂipped (while its testing set is kept clean).
The ResNet-18 and ResNet-32 networks are used.
They are trained for 200 epochs using
SGD with momentum 0 . 9 , weight decay 0 . 0001 , batch size 128 , initial learning rate 0 .
[Table]
and cosine learning rate schedule.
For LookSAM, the SAM update is used every k = 2 steps.
1 For AE-SAM and
AE-LookSAM, we set λ 1 = − 1 and λ 2 = 1 in their c t schedules g λ 1 ,λ 2 ( t ) , such that their fractions of SAM updates (approximately 50% ) are comparable with SS-SAM and LookSAM.
Experiments are repeated with ﬁve different random seeds.
Results.
Table 3 shows the testing accuracy and fraction of SAM updates 2 .
As can be seen, AELookSAM achieves comparable performance with SAM but is faster as only half of the iterations use the SAM update.
Compared with ESAM, SS-SAM and LookSAM, AE-LookSAM performs better.
The improvement is particularly noticeable at the higher noise levels (e.g., 80% ).
[Table]
Figure 4 shows the training and testing accuracies with number of epochs at a noise level of 80% using ResNet-18 .
Results for other noise levels and ResNet-32 are shown in Appendix B.2.
As can be seen, SAM is robust to the label noise, while ERM and SS-SAM heavily suffer from overﬁtting.
AE-SAM and LookSAM can alleviate the overﬁtting problem to a certain extent.
AE-LookSAM, by combining the adaptive policy with LookSAM, achieves the same high level of robustness as SAM.
[Figure]
4. E FFECTS OF λ 1 AND λ
In this experiment
[Figure]
we study the effects of λ 1 and λ 2 on AE-SAM.
A similar on AE-LookSAM is shown in Appendix B.2.
We use the setup as in Section 4.1, where λ 1 and λ 2 (with λ 1 ≤ λ 2 ) are chosen from { 0 , ± 1 , ± 2 } .
Figure 5 shows the effect on the fraction of SAM updates.
For a ﬁxed λ 2 , increasing λ 1 increases the threshold c t , and the condition (cid:107)∇L ( B t ; w t ) (cid:107) 2 ≥ µ t + c t σ t becomes more difﬁcult to satisfy.
Thus, as can be seen, the fraction of SAM updates is reduced.
The same applies when λ 2 increases.
Figure 6 shows the effect on the testing accuracy.
For λ 2 ∈ { 0 , 1 , 2 } (particularly, λ 2 = 2 ), increasing λ 1 leads to a reduction in testing accuracy as SAM updates are fewer.
For λ 2 ∈ {− 2 , − 1 } , their accuracies are similar as the SAM update is used frequently ( % SAM ≥ 80% as shown in Figure 5).
[Figure] [Figure]
4. C ONVERGENCE
In this experiment
[Figure]
we validate whether w t ’s, where t is the number of epochs, obtained from
AE-SAM can reach critical points of L ( D ; w ) , as suggested in Theorem 3.3.
From the experimental results in Section 4.1, Figure 7 shows (cid:107)∇L ( D ; w t ) (cid:107) 2 w.r.t.
t .
As can be seen, in all settings, (cid:107)∇L ( D ; w t ) (cid:107) 2 converges to 0 .
Figure 8 in Appendix B.1 also veriﬁes the convergence of AE-SAM’s training loss on CIFAR-10 and CIFAR-100 .
[Figure]
C ONCLUSION
In this paper, we proposed an adaptive policy to employ SAM based on the loss landscape geometry.
Using the policy, we proposed an efﬁcient algorithm (called AE-SAM) to reduce the fraction of SAM updates during training.
We theoretically and empirically analyzed the convergence of AE-SAM.
Experimental results on a number of datasets and network architectures verify the effectiveness of the proposed policy.
Moreover, the adaptive policy is general and can be combined with SAM variant, as demonstrated
A P ROOFS A. P ROOF OF T HEOREM 3.
[-i.e., I A ( x ) = 
 if x ∈ A otherwise 0 .-]Theorem 3.3.
Let b be the mini-batch size. If η = 1 4 β √ T and ρ = 1 /T 14 , algorithms mixing SAM and ERM satisfy
[Equation]
where ζ = 1 T (cid:80) T {+− 1 t =
+}
ξ t ∈ [0 , 1] .
Lemma A.1 (Andriushchenko & Flammarion (2022)) . Under Assumptions 3.1 and 3.2 for all t and ρ > 0 , we have
[Equation] [Equation]
Proof.
Let g t ≡ [Equation] [Equation] f
[Equation] [Equation] [Equation] i ;
[Equation] t + [Equation] g t ) , y i ) , and ˆ g t ≡ ∇L ( D ; w t ) .
By Taylor expansion and L ( D ; w ) is β -smooth, we have L
[Equation]
D
[Equation] w t +1 ) [Equation] where we have used ξ t (1 − ξ t ) = 0 as ξ t ∈ { 0 , 1 } , ξ 2 t = ξ t , and (1 − ξ t ) 2 = 1 − ξ t to obtain Eq.(9).
Taking expectation w.r.t. w t on both sides of (10), we have
[Equation]
Claim 1: E (cid:107) g t {+(cid:107) 
+}
=
E (cid:107) g t − ˆ g t (cid:107) 2 + E (cid:107) ˆ g t (cid:107) [-
-]2 =
σ 2 b + E (cid:107) ˆ g t (cid:107) 2 , which follows from Assumption 3.2. Claim 2:
[Equation]
t [Equation] t , which is derived as follows: [Equation]
where Eq.(12) follows from (cid:107) h t − g [Equation] ≤
[Equation]
and Assumption 3.2, Eq.(13) follows from the above Claim
[Equation]
Substituting Claims 1 and 2 into (11), we obtain
E L ( D ; w t +1 )
[Equation]
where Eq.(15) follows from Claims 1 and
[Equation]
Eq.(16) follows from Lemma A.1 and 1 − ηβ > 0 . As η < 14 β , we have 1 + βη − 2 ρβ ≤ 3 / 2 and βη < 1 / 4 , thus, 1 − (1 + βη − 2 ρβ ) ξ t 2 − βη 2 > 0 .
Summing over t on both sides of Eq.(17) and rearranging, we obtain
[Equation]
where γ = 1 + βη − 2 ρβ ≤ 3 / 2 , κ = 1 [-t =
 ξ t-]+ ρ 2 β 2 , ρ 2 = 1 / √ T , and ζ = 1 T (cid:80) T − 1 t =
ξ t ∈ [0 , 1] . We thus ﬁnish the proof.
Corollary A.2. Let b be the mini-batch size. If η = 1 4 β √ T and ρ = 1 /T 14 , SAM (Foret et al., 2021) satisﬁes
[Equation]
Corollary A.3. Let b be the mini-batch size.
If η = √ b 4 β √ T and ρ = 1 /T 14 , algorithms mixing SAM and ERM satisfy
[Equation]
{+where ζ = 1 T (cid:80) T − 1 t =
 ξ t ∈ [0 , 1]+} . Proof. It follows from (18) that
[Equation] A. C ONVERGENCE OF FULL - BATCH GRADIENT DESCENT FOR AE-SAM Theorem A.4. Under Assumption 3.1, with full-batch gradient descent, if ρ < 12 β and η < 1 β , algorithms mixing SAM and ERM satisfy
[Equation]
where ζ = 1 T (cid:80) T − 1 t {+=
 ξ+} t ∈ [0 , 1] . Lemma A.5 (Lemma 7 in Andriushchenko & Flammarion (2022)) .
Let L ( D ; w ) be a β -smooth function. For any ρ > 0 , we have
[Equation]
Proof of Theorem A.4. Let g t ≡ ∇L ( D ; w t ) and h t ≡ ∇L ( D ; w t + ρ ∇L ( D ; w t )) be the update direction of ERM and SAM, respectively.
By Taylor expansion and L ( D ; w ) is β -smooth, we have L ( D ; w t +1 )
[Equation]
where we have used ξ t (1 − ξ t ) = 0 as ξ t ∈ { 0 , 1 } , ξ 2 t = ξ t , and (1 − ξ t ) 2 = 1 − ξ t to obtain Eq.(28).
As (cid:107) h t (cid:107)
=
(cid:107) h t − g t (cid:107) 2 − (cid:107) g t (cid:107) 2 + 2 g (cid:62) t h t , it follows from Eq. (29) that
[Equation]
where we have used (cid:107) h t − g t (cid:107)
[Equation]
β
[Equation]
g
[Equation]
2 to obtain (30), and Lemma A.5 to obtain (31).
Summing over t from t = 0 to T − 1 on both sides of Eq.(32) and rearranging, we have
[Equation]
As ρ < 12 β and η < 1 β , it follows that 1 − βη 2 − βρξ t > 0 for all t [-=
-]. Thus, Eq.
[-ξ t = 1 T (cid:80) T − 1 t =
 I { w : (cid:107)∇L ( B t ; w ) (cid:107)≥ µ t + c t σ t } ( w t ) ∈ [0 , 1] . As 12 − βη ≥ 12 − 14 = 14 and η (1 − βη ) = 1 4 β √ T (1 − β 4 β √ T ) ≥ 3 16 β √ T , the upper bound in Eq.(30) can be simpliﬁed as-](33) implies
[Equation]
{+(34) where ζ = 1 T (cid:80) T − 1 t =
+}
ξ t ∈ [0 , 1] and we ﬁnish the proof.
B A DDITIONAL E XPERIMENTS B. A DDITIONAL E XPERIMENTAL R ESULTS ON CIFAR-10 AND CIFAR-
Figure 8 shows the convergence of AE-SAM’s training loss on CIFAR-10 and CIFAR-100 datasets.
As can be seen, AE-SAM achieves convergence in various network architectures.
B. A DDITIONAL E XPERIMENTAL R ESULTS ON CIFAR-10 WITH L ABEL N OISE
Table 4 reports the testing accuracy and fraction of SAM updates on CIFAR-10 with label noise.
Effects of k on LookSAM.
We conduct an additional experiment on CIFAR-10 with label noise to show that the performance of LookSAM is sensitive to the choice of k .
Table 5 reports the testing accuracy and fraction of SAM updates using LookSAM with different k ∈ { 2 , 3 , 4 , 5 } for ResNet-18 .
As can be seen, k = 2 yields better performance than k = 3 , 4 , or 5 , particularly in the higher noisy levels (e.g., 80% ).
Effects of λ 1 and λ 2 on AE-LookSAM.
We study the effects of λ 1 and λ 2 on AE-LookSAM.
We conduct an experiment on CIFAR-10 with label noise ( 80% noisy level), where AE-LookSAM achieves comparable performance.
Figure 9 shows the effects of λ 1 and λ 2 on fraction of SAM updates.
As can be seen, again, for a ﬁxed λ 2 , increasing λ 1 always reduces the fraction of SAM updates. [Figure] [Figure] [Table]
[Table]
[Figure]
[Table] [Figure] [Figure] Figure 10 shows the effects of λ 1 and λ 2 on testing accuracy of AE-SAM.
We can see that for a ﬁxed λ 2 ∈ { 0 , 1 , 2 } (particularly λ 2 = 2 ), increasing λ 1 always leads a reduction in testing accuracy.
Curves of training and testing accuracy.
For ResNet-18 , Figures 11, 12, and 13 show the curves of accuracies at noise levels of 20% , 40% , and 60% , respectively; For ResNet-32 , Figures 14, 15, , 16, and 17 show the curves of accuracies at noise levels of 20% , 40% , 60% , and 80% , respectively.
As can be seen, in all settings, AE-LookSAM is robust to labels noise like SAM. [Figure] [Figure] [Figure]
[Figure] [Figure] [Figure] [Figure]
B. D ISTRIBUTION OF S TOCHASTIC G RADIENT N ORM
Besides Figure 3(a) for CIFAR-100 using ResNet-
[Figure]
we empirically verify that the distribution of stochastic gradient norms follows a Bell curve for other network architectures and dataset.
Figureshows the distribution for WRN-28-10 and PyramidNet-110 on CIFAR-100 , while Figure 19 shows the distributions for CIFAR-10 .
As can be seen, in all settings, the distribution follows a Bell curve
[Figure]
[Figure]
Table 6 shows the effects of perturbance radius ρ on the testing accuracy of AE-SAM.
As can be seen,
AE-SAM is insensitive over a large range of ρ (e.g., ρ ∈ { 0 . 05 , 0 . 1 } ).
In experiments, ρ = 0 . 05 and 0 .
1 are chosen for CIFAR-10 and CIFAR-100 , respectively, as is common in the SAM literature (e.g.,
SAM and ESAM) [Table] [Table] B. E FFECTS OF BATCH SIZE
We conduct experiments to study the effects of batch size on AE-SAM.
Figure 20 show the curves of variance of gradient on CIFAR-10 and CIFAR-100 with different batch sizes.
As shown, variances are related to batch sizes, which is consistent with the observation in Ko et al.
Table 7 shows the testing accuracy.
As can be seen, AE-SAM is insensitive over a large range of batch size
[Table]
and b = 128 is a good choice and is used in the experiments.
[Table] [Figure] i [Figure]
i
[Figure]
B. R ESULTS ON VISION TRANSFORMER
We conduct experiments on vision transformer using the setting ( ViT-S16 and heavy data augmentation) in Zhao et al.
(2022a).
Table 8 shows the testing accuracy.
As can be seen, AE-SAM performs better than SAM while using only 50% of SAM updates.
SS-SAM and AE-SAM have comparable %SAM (about 50% )
[Figure]
and AE-SAM achieves higher accuracy.
As for LookSAM and AE-LookSAM, both use about 20% of SAM updates but AE-LookSAM is better.
These improvements conﬁrm the effectiveness of the adaptive policy for ViT-S
[Figure] [Table]
C ONVERGENCE CURVES OF AE-SAM AND SS-SAM
Figure 21 shows the training losses w.r.t.
number of epochs for AE-SAM and SS-SAM.
As can be seen, AE-SAM and SS-SAM converge at a comparable speed, which agrees with Theorem 3.3 as both of them have comparable fractions of SAM updates (Table 1)
[Table]
[Figure]