Provable Sharpness-Aware Minimization with Adaptive Learning Rate 
Authors: Hao Sun
Li Shen
Qihuang Zhong
Liang Ding
Shixiang Chen
Jingwei Sun
Guangzhong Sun
Dacheng Tao
Venue: Submitted to ICLR 2023 (2023)
Keywords: Adaptive learning rate
Sharpness aware minimization
mini-batch linear speedup

Abstract: Sharpness aware minimization (SAM) optimizer has been extensively explored as it can converge fast and train deep neural networks efficiently via introducing extra perturbation steps to flatten the landscape of deep learning models. A combination of SAM with adaptive learning rate (AdaSAM) has also been explored to train large-scale deep neural networks without theoretical guarantee due to the dual difficulties in analyzing the perturbation step and the coupled adaptive learning rate. In this paper, we try to analyze the convergence rate of AdaSAM in the stochastic non-convex setting. We theoretically show that AdaSAM admit a $\mathcal{O}(1/\sqrt{bT})$ convergence rate and show linear speedup property with respect to mini-batch size b. To best of our knowledge, we are the first to provide the non-trivial convergence rate of SAM with an adaptive learning rate.  To decouple the two stochastic gradient steps with the adaptive learning rate, we first introduce the delayed second-order momentum during the convergence to decompose them to make them independent while taking an expectation. Then we bound them by showing the adaptive learning rate has a limited range, which makes our analysis feasible.  At last, we conduct experiments on several NLP tasks and they show that AdaSAM could achieve superior performance compared with SGD, AMSGrad, and SAM optimizer.

Under review as a conference paper at ICLR 2023
PROVABLE SHARPNESS-AWARE MINIMIZATION WITH ADAPTIVE LEARNING RATE
Anonymous authors Paper under double-blind review
ABSTRACT
Sharpness aware minimization (SAM) optimizer has been extensively explored as it can converge fast and train deep neural networks efficiently via introducing extra perturbation steps to flatten the landscape of deep learning models. A combination of SAM with adaptive learning rate (AdaSAM) has also been explored train large-scale deep neural networks without theoretical guarantee due to the dual difficulties in analyzing the perturbation step and the coupled adaptive learning rate. In this paper, we try to analyze the convergence rate of AdaSAM in the stochastic non-convex setting, We theoretically show that AdaSAM admit O(L/v convergence rate and show linear speedup property with respect to mini-batch size b. To the best of our knowledge, we are the first to provide the non-trivial convergence rate of SAM with an adaptive learning rate. To decouple the two stochastic gradient steps with the adaptive learning rate, we introduce the delayed second-order momentum term the convergence to decompose them to make them independent while taking expectation. Then we bound them by showing the adaptive learning rate has a limited range, which makes our analysis feasible. At last, conduct experiments several NLP tasks and they show that AdaSAM could achieve superior performance compared with SGD, AMSGrad, and SAM optimizer.
INTRODUCTION
Sharpness-aware minimization (SAM)Foree et al.] (2021) a powerful optimizer for training large- scale deep learning models by minimizing the gap between the training performance and general- ization performance. It has achieved remarkable results on various deep neural networks, such as ResNet He et al. (2016), vision transformer Dosovitskiy et al. (202T); Xu et al. (202I). language models Devlin et al. (2018): He et al. 2020), on extensive benchmarks.
However, SAM-type methods suffer from several issues during training the deep neural network, especially for huge computation costs and heavily gradients computation compared with classic optimizers, like SGD, Adam & Ba (015)). AMSGrad Reddi et al.] (2018), due to the extra perturbation step. Hence, SAM requires to forward and back propagate twice for one parameter update, resulting one more computation cost than the classic optimizers. Moreover, as there are two steps during the training process, it needs double hyper-parameters, which makes the learning rate tuning unbearable and costly.
tuning procedure. It needs double
Adaptive learning rate optimization methods scale the gradients based on the history gradient in- formation accelerate convergence. These methods, such as Adagrad Duchi al. (201T). Adam, and AMSGrad, have been studied for solving the NLP tasks Zhang et al. (2020). However, it might converge local minima which leads to poor generalization performance. Many studies have investigated how to improve the generalization ability of the adaptive learning rate methods. Recently, several work has tried to ease the learning rate tuning in SAM by combining SAM with adaptive learning rate. For example, Zhuang et al. (2022) trains VIT models with adaptive method and Kwon et al. (2021) training a NLP model. A Although remarkable performance has been achieved, their convergence is still unknown since the adaptive learning rate is used in SAM. Directly analyzing the convergence of this method is difficult due to the two steps of optimization, especially the adaptive learning rate adopted in the second optimization step.

---

Under review as a conference paper at ICLR 2023
In this paper, we analyze the convergence rate of SAM with adaptive learning rate in the non- convex stochastic setting, dubbed AdaSAM. To circumvent the difficulty in the analysis, we develop technique to decouple the two-step training of SAM and the adaptive learning rate. The analysis procedure is mainly divided into two parts. The first part is to analyze the procedure of the SAM. Then we analyze the second step that adopts the adaptive learning rate method. We introduce second-order momentum term from the previous iteration, which is related to the adaptive learning rate and independent of SAM while taking an expectation. Then we can bound the term composed by the SAM and the previous second-order momentum due to the limited adaptive learning rate. We prove that AdaSAM enjoys the property of linear speedup with respect to the batch size. Furthermore, we apply AdaSAM to train RoBERTa model on the GLUE benchmark. We show that AdaS AM achieves the best performance experiments, where wins 6 tasks of tasks, and the linear speedup can be clearly observed.
In the end, we summarize our contributions as follows:
We present the first convergence guarantee of the adaptive SAM method under the non- convex setting. Our results suggest that a large mini-batch can help convergence due the linear speedup with respect to batch size.
We conduct a series of experiments on various tasks. The results show that AdaSAM outperforms the most of state-of-art optimizers and the linear speedup is verified.
2 PRELIMINARY
2.1 PROBLEM SETUP
In this work, focus the following stochastic nonconvex optimization min (xx) == (x)))( ZERd
(I)
where d is dimension of variable D is the unknown distribution of the data samples, fx( smooth and possibly non-convex function, and fs, (x) denotes the objective function at the sampled data point &, according data distribution D. In machine learning. it covers empirical risk minimization as a special case and of is the loss function. When the dataset D cover data points,
ie.,, D = {E,, i 1,2,.... ... N). N.) Problem Saneemmreeiie reduces to the following finite-sum problem:
1 min f(x) = N EEa(())
ZERd
Without additional declaration, we represent fi(x) fe, (x) for simplification.
Notations. f((x) is the i-th loss function while parameter dimension. We denote the la norm as || 1l2- A Hadamard product is denoted as a where a are vectors. For vector IRd, Va is denoted vector that the j-th value, (Va) is equal to the square root of aj
Rd is the model parameter and d is the
2.2 RELATED WORK
Sharpness-aware minimization Previous work shows that sharp minima may lead to poor gen- eralization whereas flat minima perform betterJiang et al. (2020); Keskar et af.) (2017): He et al.) (2019). Therefore, it is popular to consider sharpness to be closely related to the generalization. Sharpness-aware minimization is proposed in Foret et al. (202T), which targets to find flat minimizers by minimizing the training loss uniformly entire neighborhood. Specifically, Foret et al. (2021) aims to solve the following minimax saddle point problem:
min max fx + 6) + >||x1|2-
# [|
where p 2 0 and À > are two hyperparameters. That is, the perturbed loss function of f(x) in neighborhood is minimized instead of the original loss function f(x). By using Taylor expansion of fa 8) with respect to ò, the inner max problem is approximately solved via
v((x)
8" (x)
arg max [(x + 5) arg max fxx + &Vv((x) arg max 8 ()(x)= "|V](|] Sp
2

---

Under review as a conference paper at ICLR 2023
Dropping the quadratic term, L is simplified as the following minimization problem
Vf(z)
min f
(4)
[Vf(2)]
The stochastic gradient of / (+
SAM further approximates the gradient by (
[[v]l on batch data includes the Hessian vector product,
V (ar())
Vafo
x +
vxfo())
fx))
Then, along the negative direction [=
SGD is applied to solve minimization
2+P
problem
44. It is easy to see that SAM requires twice gradient pack-propagation, i.e., V.((z) and Wdyte) Due the existence of hyperparameter p. one needs carefully tune both p PINEEE
(zf((z)
and learning rate in SAM. In practice, p is predefined to control the radius of the neighborhood.
Some variants of SAM proposed improve its performance. Recently, several works ZZuang et al. (022); Kwon et al. (2021) have empirically incorporated adaptive learning rate with SAM and shown impressive generalization accuracy, while their convergence analysis has never been studied.
ESAM Du et al. 2022) proposes an efficient method by sparsifying the gradients will alleviate the double computation cost of backpropagation. ASAM Kwon al. (202T) modifies SAM by adaptively scaling the neighborhood so that the sharpness is invariant to parameters re-scaling, GSAM (Zhuang al. 2022) simultaneously minimizes the perturbed function and new defined surrogate gap function to further improve the flatness of minimizers. Liu et al. Liu et al. (2022) also studies SAM in large-batch training scenario and periodically update the perturbed gradient. the other hand, there some works analyzing the convergence of the SAM such as Andriushchenko & Flammarion (022). where it does not consider normalization step, i.e., the normalization in Vi(a) [V
ddaptive optimizer The adaptive learning rate method can automatically adjust the learning rate based on the history gradients methods. The first adaptive method is Adagrad Duchi et al. (2011). and can achieve a better result than other first-order methods under the convex setting. While training the deep neural network, Adagrad will decrease the learning rate rapidly and the result is not good. Adadelta Zeiler 2012) is proposed to change this situation and introduces a learning rate based on the exponential average history gradients. Adam Kingma & Ba (2015) adds momentum, and it shows great performance in many tasks. However, Reddi et al Reddi et al. (0018) give an example that cannot converge even when the objective function is convex. They propose new method called AMSGrad that can solve this method. Many papers also analyze the convergence of the adaptive methods Zhou al. (2018); Chen et al. (2019); Zaheer al. (2018): Ward et al.) (2019): Défossez You et al. (2019): Chen et al. (2021).
et al.) 2020);
3 ADASAM: SAM WITH ADAPTIVE LEARNING RATE
In this section, we introduce SAM with the adaptive learning rate (AdaSAM) from the AMSGrad optimizer. Then, will present the convergence results of AdaSAM. At last, we give the proof sketch for main theorem.
AdaSAM is described in Algorithm 1 In each iteration, mini-batch gradient estimation Jr point x + c(x) with batchsize b is computed, i.e.,
gt (=f6(z)la,
fz. + 8(x)))
Here, 6(x)) is the extra perturbed gradient step in SAM that is given as follows
St 6(xx) = p [
where 8; = V=.((=)),,
Σ (a..
BEB
Then, exponential averaging of gi and the second-order term [9.]2 accumulatively computed as
m, and Dr, respectively. AdaSAM then updates iterate along -m, with the adaptive learning rate yp.
3

---

Under review as a conference paper at ICLR 2023
Algorithm 1: AdaSAM: SAM with adaptive learning rate
Input: Initial parameters 20, m_11 = O, 6-1 = (a small positive scalar to avoid the denominator diminishing), base learning rate y, neighborhood size and momentum parameters 81, 82.
Output: Optimized parameter 27++1
iteration E [0, 1, ........
1) do
Sample mini-batch = (Et, Stz St. S1
Compute gradient 8: = V v (x) EEeD (zi);
Compute (xre) Pt 2 |
Compute SAM gradient 0: Vafu(()laatdee))
my -1 + (1 g1)
Vi 820=11
0, max(01- 1,
T)t 1/v
10
Tt+1 Zt 7m ④ 17:
end
Remark 3.1. Below, we give several comments on AdaSAM:
When B2 1, the adaptive learning rate reduce to the diminishing SGD. Then, AdaSAM recovers the classic SAM optimizer.
If we drop out the 8-th line & = max(@;-11 v2), then our algorithm becomes the variant of Adam. The counterexample that Adam does converge in the Reddi al. 2018 also holds for the SAM variant, while AdaSAM can converge.
3.1 CONVERGENCE ANALYSIS
Before stating our convergence analysis, we first introduce some useful assumptions. Assumption 3.2. L-smooth. f. and f is differentiable with gradient Lipschitz property:
[=) - Vf(vv) s L|z -11 (Vx()) vf(w)) < L||= yll Vz,y E R 1, 2, .. which also implies the descent inequality, i.e., fy) s f((x) + (Vf.(c).v 2) f - 2/
Assumption 3.3. Bounded variance. The estimator of the gradient is unbiased and the variance of the stochastic gradient is bounded:
EVf(() fx()),
When the mini-batch size b is used, we have - $. Assumption 3.4. Bounded stochastic gradients. The stochastic gradient is uniformly bounded, i.e.,
-
[x) S for any I....,N. N.
Remark 3.5. The above assumptions commonly used in the proof of the convergence for adaptive stochastic gradient methods such as Cutkosky & Orabona (2019): Huang et al. (2021) Zhou et al. (2018): Chen et al. (2019).
Note that AdaSAM is the combination of AMSGrad and SAM. We briefly explain the main idea of analysis in AMSGradReddi et al. (2018) and SAMAndrrushchenko & Flammarion (2022).
We firstly notice that the main step in AMSGrad analysis is estimate the expectation E E[ - 2x] Em, @ V/ = - E(1 31 )9. C E,3 m O T/t which is conditioned on the filtration a (ar.). First, we consider the situation that 81 0 which does not include the momen- tum. Some works Zaheer et al.] (2018); Savarese et al. (021 apply delay technology to split the dependence between Sr and T). that is
Ege T/ Egg
7) - + E[g, @ (n 7-1 v f((x)) o n-1 Elg, + (n: Th-
The second term E [g @ (n 1)) dominated by the first term ff(x,) 2 7h-1- Then, is not difficult to get the convergence result of AMSGrad. When apply the strategy AdaSAM,
4

---

Under review as a conference paper at ICLR 2023
find that Egr 17/-1 cannot be handled similarly because Eg, = afo (+ C Mi * Vf(z))
Va(aa
Inspired by (Andriushchenko & Flammarion 2022 Lemma 16), our key observation is that
V v((z)) [++))
Vf(x)
Vf(z)
EV +++
EV.fb
+
vnf
V
[V[())]
vh(2) and we prove the other terms such as E Vz x + PTVl(()]] have small values that do not dominate the convergence rate.
2 (x+ ( P
Pivsall
V/(a)
π)-1
Then, when we apply the momentum method, we find that the term Em--1 ④ T cannot ig- nored. By introducing an auxiliary sequence 2t = x₁ + I- (x, - xz-1), we have B[ 2t] =
-EI, 8t 1-8, Mm @ (M-- Te) -9: @ n/.. The first term contains the momentum term which has small value and we can remove it without hurting the convergence rate.
Theorem 3.6. Under the assumptions 3.2 3.4 and 7 a fixed number that satisfies that 7 S I6I for the sequence { generated by Algorithm we have the following convergence rate
2G(f(zo) f²)
8GyL o²
E
+ +
(S)
78
be
t=
where
45GL22p
20
8GyL Lit -
+ (1
d(
3T
2(4
1-3;
81
12
LG
d(E-2 + (1-3,)2
Ly dG3
(6)
Corollary 3.7 (mini-batch linear speedup). Under the same conditions of Theorem
3.0
Furthermore,
when we choose the learning rate 0( and neighborhood size P Olv
have
bT
T Σ [ O
( + +
VDT
1 bT
+
(
+0 bETT
T
b
O
T
When T' is sufficiently large, we achieve the linear speedup convergence rate with respect to mini-batch size b, i.e.
T Σ E V(
O
VTT
Remark 3.8. To reach O6) stationary point, when the batch size is it needs T = ( iterations. When the batch size is b, we need to run T O 1 steps. The method with batch size b is b times faster than batch size of 1.
3.2 PROOF SKETCHI
In this part, we will give the proof sketch of the Theorem 3.6 For the complete proof, please see Appendix. We first introduce an auxiliary sequence 2, v, + I-3, (xt I2-1). By applying L-smooth condition, we have
f(z(11) S f(x)) + (Vf(2)) - 2)) + |1 - -2
(8)
2
Applying it to the sequence } and using the delay strategy yield
7:1
+ L - = - =oll?
((x41) - ((x)) < (V/(z),
M4-1
@ (M1-1 116))
1- Br
2
+ (Vf(z)), Σ Vl(a + Pt
St
④ (1)-1 1)))
|
5

---

Under review as a conference paper at ICLR 2023
+ (VI(=))) - V(().)) 7 Σ
Vii pt
81
) 19 (--1
[sell
SEB
+ (V)())) 3 Σ
(z +
3 10-111
2 Σ
Vf((2,
+ Pt
8t
o
mh-1)
| fzz))]
5
KEB
IEB
+ (Vf(z))
f1 +
② 7--1)
(9)
Vf(z))]
REB
From the Lemma B.5 follows
Lemma B.6
Lemma
B.7 in appendix, we can bound the above terms in
9 as
(v((2),
Σ
Vii( + Pt
St
O (Tt-1- - 11)) S 1 nell2
(10)
iEB
y31
731
(V)(=)),
m,-1 2 (Mt-1 - 11)) -
G2 || - 1/ell1
(11)
B1
1-
(v)()))
7
Vi P:
(1(=))
3 T)--1
?
Vi( +0r
8t
1h-1)
b
[+}()+)]
b
|s1
< 7 (xx) @ 11 LL
(12)
2p²2
Then we substitute them into the (9. and take the conditional expectation to get
Ef Ef(x+11- f(x) E 7 Σ
Vfla
+ pt
vf(zz)
① M--1) + 2 E1
L
- E112411-11
zel1?
iEB
+ 2 [V.) ④ / 2 L2
7
G 12 -1-1-16 nell
2/²
c
1- B1
E (Vf(z) - fz()), 7 fi + Pt
8t
1
1(-)))
(13)
|
where μ > is to be determined.
From the Lemma
B.8
Lemma B.10 and Lemmal B.9 in Appendix, we have
ff())
E (Vv(), 1/
Σ
V.((I + Pt [(+'))) (
iEB
12?
$ <>>\if(x)) 2 - 2 +
E
7
[V)()) @ 11 2+
(14)
2a²
2e
4
[ El||z11 -211 < (1 3,)2
E In
カー 2 + L
+ B
E
2 2
Lp
σ²
Be
be
+1++3)G2B/1 1 113)
(15)
E(VS(=)- vf(x)),
Σ
Vf( + Pt ssell
8t
o mp-2)
iEB
7 83
1
1
1
dG2
?PP
$ 2c(1 - 32
+
+
+
o 2n +
(>3 + 423). (16)
X3
2
2
2E
Next, we substitute it into the (3). Taking the expectation over all history information yields
1
3- 33
X
- 5
[[(()() 3 TT=
24²2
2a
Be
2pp2L?p3
?
GPE
1パ-1
1/1 0
+
T3L233
dG2
PL Pt
(2
4x3)
31
2c - /1)2
X
X2
2
2
1222
3
L(1
3)
Lpł +
o²
g2
B1
+
+ 2
)2
+ 1 + 3)
11 7- 112
(17)
2c
3e
bc
81

---

Under review as a conference paper at ICLR 2023
We set n²² a² = 8, = 3, A2 22 = A3
and we choose 274 - 8. Note that d is bounded.
We have
Y [ < 2 @ VI=111
(18)
2G
S -Ef(# + Ef(x))
45, p
7
G - 7ell1
2E
1- 31
+ 4 Lpi + bc
))+ + (4+
31
323 dG2 (1 B1)² e²
- n- 2
(19)
1 - B1
Then, telescoping it from t = 0 = T 1, and assuming Y is constant, it follows that
E [ 2G(f(zn) - f*) 8G>L 45GL
2G
d( 1
T 1=
77
be
(1 B1)T
G
SG Luf
2(4 + ( (1-8,
d(E= - G-2) 6-22 dG3
(20)
T
(1- 81) 3
which completes the proof.
EXPERIMENTS
In this section, we apply AdaSAM to training language models and compare it with SGD, AMSGrad, and SAM to show the effectiveness of AdaSAM. Due to space limitations, more experimental results visualization, task description, implementation details and ablation study are placed in the Appendix.
Table I: Evaluating AMSGrad and AdaSAM on the GLUE benchmark with 8 0.9
AdaSAM AdaSAM AdaSAM AdaSAM rho=0,005 rho=0.001
Tasks
Metric
AMSGrad
rho=0.01
best
CoLA
Mcc.
68
65.29
68.74
67.3
68.74
SST-2
Acc.
96.33
96.33
96.67
96.1
96.67
Acc. 90.2
91.18
90.93
91.18
MRPC
F1
92.72
93.64
93.36
92.96
93.64
Pear.C. 91.72
90.13
91.64
91.9
91.9
STS-B
Spca.C.
91.48
90.36
91.38
91.62
91.62
RTE
Acc. 87.73
84.84
87.73
85.92
87.73
MNLI-m Acc. 90.67
90.97
90.88
90.45
90.97
MNLI-mm
Acc.
90.41
90.42
90.4
90.4
90.42
QNLI
Acc.
94.82
94.65
94.56
94.56
94.65
F1
88.7
88.55
88.69
88.64
88.69
QQP
Acc.
91.41
91.23
91.33
91.27
91.33
Average Scores
89.52
88.97
89.69
89.28
89.8
relative improvements
-0.55
0.17
-0.24
0.28
4.1 EXPERIMENTAL SETUP
Tasks and Datascts. We evaluate the SAM with an adaptive learning on a popular benchmark, i.e. General Language Understanding Evaluation (GLUB) Wang et al. (2018), which consists of several language understand- ing tasks including sentiment analysis, question answering and textual entailment. For a fair comparison, we report the results based on taske without multi-task or ensemble training. We evaluate the performance with Accuracy ("Ace") metric for most tasks, except the Fl scores for QQP and MRPC, the Pearson-Spearman correlations ("Pcor/Scor") for STS-B and the Matthew conrelations ("Mec") for CoI.A.
7

---

Under review as a conference paper at ICLR 2023
Table 2: Evaluating SGD and SAM the GLUE benchmark with B1 = 0.9
AdaSAM rho=0.01
AdaSAM rho=0.005
AdaSAM best
Tasks
Metric
AMSGrad
CoLA
Mcc.
9.25
4.64
66.76
66.76
SST-2
Acc.
50.92
95.87
95.76
95.87
Acc.
68.38
70.58
68.38
70.58
MRPC
F1
81.22
81.98
81.22
81.98
Pear C.
3,22
84.74
2
84.74
STS-B
Spea.C.
1.9
85.57
2
85.57
RTE
Acc.
55.6
52.71
52.71
52.71
MNLI-m
Acc.
84.94
90.5
90.42
90.5
MNLI-mm
Acc.
84.87
90.19
89.74
90.19
QNLI
Acc.
63.61
94.44
94.6
94.6
85.6
84.7
86.72
86.72
QQP
Acc.
80.14
87.88
89.94
89.94
Average Scores
55.8
76.98
68.35
82.51
relative improvements
21.18
12.55
26.71
Implementations. We conduct our experiments using a widely-used pretrained language model, RoBERTa- large] in the open-source toolkit fairsed" with 24 transformer layers, a hidden size of 1024, For fine-tuning each task, epochs, the batch size, etc In particular, for RTE, STS-B and MRPC of GLUE benchmark, the pre-trained RoBERTa-large model on the MNLI dataset and continue fine-tuning Ro model on the corresponding single-task corpus for better performance, as many prior works did did et al. (2019);
different combinations of hyper-parre parameters, including the learning rate, the number of
first fine-tune
He et al. (2020). All models are trained on NVIDIA DGX SuperPOD cluster, in which each machine contains 40GB A100 GPUs.
Table Evaluating AMSGGrd and AdaSAM on the GLUE benchmark without momentum (3) 0)
Tasks
Metric
AMSGrad
AdaSAM(hh=00.0)
AdaSAM(rho=0000)
AdaSAM(best)
CoLAA
Mee.
63.78
69.23
68.47
69.23
SST-2
Acc.
96.44
96.22
96.22
96.22
MRPC
Acc.
89.71
89.96
89.96
89.96
Fl
92.44
92.84
92.82
92.84
STS-B
Spea.C.
Pear. C.
90.35
89.98
89.07
88.83
91.59
91.22
91.59
91.22
RTE
Acc. 87.36
87
73.65
MNLI-m
Acc. 90.65
90.83
90.75
90.83
MNLI-mm
Acc.
90.35
90.41
90.42
90.42
QNLI
Acc.
94.53
94.8
94.73
94.8
QQP
Fl
88.59
88.67
88.72
88.72
Acc.
91.27
91.38
91,46
91.46
Average Scores
88.79
89.1
88.33
89.52
relative improvements
0.32
-0.45
0.74
4.2 RESULTS ON GLUE BENCHMARK
We conduct the experiments the SGD based method and adaptive learning method. respectively. Each method contains SAM and rate. For the SAM with an adaptive learning rate, we tune the neighborhood size from 0.01, 0.005, 0.001. The result shows that SAM with an adaptive learning rate outperforms AMSGrad on tasks of tasks except for QNLI and QQP. Overall, it improves the 0.28 than AMSGrad. Table shows the result of SGD and SAM. We
optimizer. Tsole shows the AMSGrad and SAM with adaptive learning
https: http://.....aappiicffffeeeeeeeee /tt.
roberta www.e..rgaa......
https: / //// thub. com/ facebook www.oookreeeacc//aaaa.
Due the space limitation, we show the details of the datasct and training setting in Appendix
8

---

Under review as a conference paper at ICLR 2023
Table 4: Evaluating SGD and SAM on the GLUE benchmark B1
Tasks
Metric
SGD
SAM(rho00.0)
SAM(hh=00005))
SAM(best)
CoLA
Mce.
41.91
58.79
58.79
SST-2
Acc.
51.72
95.3
81.54
95.3
MRPC
Acc.
68.38
68.38
68.38
68.38
Fl
81.22
81.22
81.22
81.22
STS-B
Spea.C.
Pear..
5,55
7.2
10.38
9.21
13.52
16.6
13.52
16.6
RTE
Acc.
51.27
53.07
53.79
53.79
MNLI-m Acc.
32.51
87.99
88.42
88.42
MNLI-mm
Acc.
32.42
87.8
88.15
88.15
QNLI
Ace.
53.32
51.24
92.95
92.95
QQP
F1
0
83.44
83.84
83,84
Acc.
63.18
87.27
87.7
87.7
Average Scores
37.23
63.1
67.91
69.06
relative improvements
25.87
30.68
31.82
tune the neighborhood size from 0.01, 0.005 for SAM. The results show that SAM is better than SGD on 7 tasks of 8 tasks except for RTE. And SAM can significantly improve the performance.
Comparing the results of Table 1 and 2 we find that the adaptive learning rate method is better than SGD based method. SAM with adaptive learning rate achieves the best metric on 6 tasks. In general, SAM with adaptive learning rate better than the other methods.
In addition, we conduct the experiments the momentum is set to 0 to evaluate the influence of the adaptive learning rate. Table 3 shows that SAM with adaptive learning rate outperforms AMSGrad on 6 tasks of tasks except for SST-2 and RTE, In Table we compare SGD and SAM, and without the momentum, SAM outperforms SGD all tasks. Under this situation, SAM with adaptive learning rate method is better than the other methods.
When comparing the result Table and momentum. When there is no momentum term, SAM with adaptive learning rate improves 0.74 than AMSGrad. With momentum term, SAM with adaptive learning rate improves only 0.28.
and Table adaptive learning rate method is more important than the
4.3 MINI-BATCH SPEEDUP
this part, we the performance with different batch sizes to validate the linear speedup property. The experiments are conducted on MRPC, RTE, and CoLA tasks. batch size set as 4, 16, 32, respectively. We scale the leamning rate VN, which is similar Li et al. (202)). ). N is the batch size. The results show that the training decreases faster the batchsize increases, the loss curve with the batch size of achieves nearly half itcrations as the curve with the batch size of 16.
1.125
1.00
1.00
0.875
0.750
Tranning
0.625
3.77
Training
0,70
0.500
0.60
3.40
0.40
0,250
3.25
0.25
0.195
2000
400
12.1)
603
1000
200
400
12.0 Steps
500
1600
1200
250
2000
750
1003 1250 1560 1750 2000
Stops
Stops
(a) MRPC
(b) RTE
(c) CoLA
Figure 1: The linear speedup of SAM-AMSGrad with the number of batch size of 4, 8, 16, 32.
5 CONCLUSION
this work, we study the convergence rate of Sharpness aware minimization (SAM) optimizer with adaptive learning rate from AMSGrad in the stochastic non-convex setting. To the best of our knowledge, we are the first provide non- non-trivial ((1/bbl) convergence rate of SAM with an adaptive learning rate, which achieves linear speedup property with respect to mini-batck size We have conducted extensive experiments on several NLP tasks, which verifies that AdaSAM could achieve superior performance compared with AMSGrad and SAM optimizers. Future works include extending AM to the distributed setting and reducing the twice gradient back-propagation cost.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-awe aware minimization. In International Conference on Machine Learning, pp. 639-668. PMLR, 2022.
Congliang Chen, Li Shen, Fangyu Zou, and Wei Liu. Towards practical adam: Non-convexity, convergence theory, and mini-batch acceleration. arXiv preprint arXiv:2101.0547) 2021.
Xiangyi Chen, Sijia Liu, Ruoyu Sun, and Mingyi Hong. On the convergence of class of adam-type type algorithms for non-couvex optimization. International Conference on Learning Representations, 2019.
Ashok Cutkosky and Francesco Orabona. Momentum- based variance reduction in non-convex sgd. Advances in neural information processing systems, 32, 2019.
Jacob Devlin, Ming- Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.0480,, 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenbom, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
Jiawei Du, Hanshu Yan, Jiashi Feng, Joey Tianyi Zhou, Liangli Zhen, Rick Siow Mong Goh, and Vincent Tan. Efficient sharpness aware minimization for improved training neural networks. In International Conference on Learning Representations, 2022.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Alexandre Défossez, Léon Bottou, Francis Bach, and Nicolas Usunier. On the convergence of adam and adagrad. CoRR, abs/2003.02395, 2020. URL URL http:::/t///tt http://rvvvvvo.. www.w.wwgwww..////.//////
02395
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In International Conference on Learning Representations, 2021.
Haowei He, Gao Huang, and Yang Yuan. Asymmetric valleys: Beyond sharp and flat local minima. Advances in neural information processing systems, 32, 2019.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Pengcheng He, Xiaodong Liu, Jianfeng Gao, and Weizhu Chen. Deberta: Decoding-enhanced bert with disentangled attention. In ICLR, 2020.
Feihu Huang, Junyi Li, and Heng Huang. Super-adam: faster and universal framework of adaptive gradients. Advances Neural Information Processing Systems, 34, 2021.
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. In International Conference on Learning Representations, 2020.
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. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.ned 2017.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR (Poxter), 2015. URL ht http://arxie wwwpexivvvvggeeeooo awtp://bb////1.1... 6980 6980
Jungmin Kwon, Jeongscop Kim, Hyunsco Park, and In Kwon Choi. Asam: Adaptive sharpness-awan minimiza- tion for scalc-invariant learning of deep neural networks. In International Conference Machine Learning, pp. 9905-5914. PMLR, 2021.
Xiaoyun Li, Belhal Karimi, and Ping Li. On distributed adaptive optimization with gradient compression. In International Conference on Learning Representations, 2021.
Yinhan I Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv, 2019.
10

---

Under review as a conference paper at ICLR 2023
Yong Liu, Siqi Mai, Xiangning Chen, Cho-Jui Hsieh, and Yang You. Towards efficient and scalable sharpness- aware minimization. Proceedings of IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR), pp. 13360-1237,, June 2022.
Sashank Reddi, Satyen Kale, Sanjiv Kumar. On convergence adam beyond. In International Conference on Learning Representations, 2018.
Pedro Savarese, David McAllester, Sudarshan Babu, and Michael Maire. Domain independent dominance of adaptive methods. In IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2021, virtual, June 19-25, 2021. pp. 16286-10000 16295, Computer Vision Foundation IEEE, 2021. doi: 10.1109/CVPR44477. 2021.01602.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. In EMNLP, 2018.
Rachel Ward, Xiaoxia Wu, and Leon Bottou. Adagrad stepsizes: Sharp convergence over nonconvex landscapes. In International Conference on Machine Learning, pp. 6677-6686. PMIR, 2019.
Yufci Xu, Qiming Zhang, Jing Zhang, and Dacheng Tao. Vitae: Vision transformer advanced by exploring intrinsic inductive bias. Advances in Neural Information Processing Systems, 34, 2021.
Manzil Zaheer, Sashank Reddi, Devendra Sachan, Satyen Kale, and Sanjiv Kumar. Adaptive methods for nonconvex optimization. Advances in neural information processing systems, 31, 2018
Matthew D Zeiler. Adadelta: an adaptive learning rate method. arXiv preprint arXiv: 12122.701, 2012.
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems, 33:15383-15528 2020.
Dongruo Zhou, Jinghui Chen, Yuan Cao, Yiqi Tang, Ziyan Yang, and Quanquan Gu. On the convergence of adaptive gradient methods for nonconvex optimization. arXiv preprint urXiv:1808.05671. 2018
Juntang Zhuang, Boqing Gong, Liangzhe Yuan, Yin Cui, Hartwig Adam, Nicha C Dvornek, sekhar tatikonda, James Duncan, and Ting Liu. Surrogate gap minimization improves sharpness aware training. In Interna- tional Conference Learning Representations, 2022.
Fangyu Zou, Li Shen, Zequn Jie, Weizhong Zhang, and Wei Liu. A sufficient condition for convergences of adam and rmsprop. In Proceedings of the JEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 11127-11135, 2019.

---

Under review as a conference paper at ICLR 2023
A EXPERIMENTAL SETTINGS
Table 5: Experimental settings and data divisions upon different downstream tasks. Notably, for each tasks in GLUE benchmark, provide the number of classes ("classes"), ), the learning rate ("Ir"), the batch size ("bsz"), the total number of updates ("total"), the number of warmup updates ("warmup") and the number of GPUs ("GPUs") during fine-tuning, respectively.
MNLI
QNLI
QQP
RTE SST-2
MRPC
CoLA
STS-B
experimental settings upon different downstream lasks
classes
2
2
2
2
2
-Ir
le-5
le-5
le-5
2e-5
le-5
le-5
le-5
2e-5
-bsz
256
128
256
32
64
32
32
32
total
15,484
8,278
14,453
1,018
10,467
1,148
2,668
1,799
warmup
929
496
867
61
628
68
160
107
-PUs
4
4
2
2
2
2
2
data divisions for each dataset
train 392,720
104,743
363,870
2,491
67,350
5,801
8,551
5,749
dev 9,815
5,463
40,431
277
873
4,076
1,043
1,500
test 9,796
5,461
390,956
3,000
1,821
1,725
1,063
1,379
The GLUE benchmark contains 8 tasks, they are RTE, STS-B, B, CoLA, SST-2, MNLI, MRPC, QNLI and QQP. CoLA is single sentence task. Each sentence has a label and -1. represents that it is grammatical sentence, evaluation
while represents that it illegal. Matthews correlation coefficient, dubbed mee is used metric. STS-B is similarity and paraphrase task. Each sample has pair paragraph. People annotated the sample from to based on the similarity between the two paragraphs. dubbed p/s correlation coefficients. RTE is an inference task. Each sample has two sentences. If two sentences have relation of entailment, the RTE task, the metric is the accuracy, dubbed that has classes. They are entailment, contradiction and neutral. MRPC task to classify whether the sentences in the pair are equivalent. QNLI is question question-answeriig task. the sentence contains the answer to the question, then it is positive sample. QQP is a social question-answering task that consists of question pairs from Quora. It determines whether the questions equivalent, The metric of MNLI, MRPC, QNIJ, QQP is accuracy.
metric Pearson Spearman,
view them as positive sample. If not, they compose of negative sample. In SST-2 is single sentence task and its metric the accuracy.
MNLI a sentence-
B PROOF
We set 2e =
26
+
T- (Ze 31-1) ) for 2 0 and we assume x-1 1 = and m =
We have that
B1
B1
2141 - 21 1141 +
(#1+1 -xe)= - 21.
(œ) x1-1)
(21)
1 B1
B1
31
(#e+1
21)
(x) x(-1)
(22)
1
1
31
m C 17
+
(x+ -x--1
2 T7-1
(23)
1 81
1 -1
B1
3, 18 + (1- (1-3))
2 1/4
(Ex xz 7M11111111111111 3 T/-1
(24)
1-
B1
=
31
m O (M)-1 mT) 7 O
(25)
By applying L-smooth, we have
((4+1) 5 f(x) + (VI(z)), - -2)+ 20
(26)
We re-organize it, and have
f(#++1) - f(x))
L 11 2xl2
(V/(z),, z)) +
(27)
12

---

Under review as a conference paper at ICLR 2023
7B1
4
<Vi(#),
m @ (7-11-7)) + (Vf(2)), -19 @m)
2 || =ell?
(28)
1- B1
731
= =V)(a,,
Me-1 (77-1 - -1)) ) + |2++1 - 2x 112
B1
+ (V)())).
Σ
V
Vila + Pr
81
① (1/t-1 v/.))
|
+VV((=), b
Vi( + Pt |1
① 1)
(29)
>B1
(V/(4)) 1- B1
m @ (7--1 - 71))
+ 2
2++ - ze
+ +VV(zz) T ΣVSi +
$1
2 (T/-1 - TT-))
+(V/(=) (x+)).
21
Σ
Viilz +
11-1)
b
|
EB
+(V/(a)). "
Σ
fi( + Pt
81
3 1/-1
(30)
||
K
=V[((x); 81 M -
731
2 || - zell
+VV((=) 2
Σ
( +
81 (Mt-1 71)))
||sl
LEB
21 + +[V((=) V.(s.)
Σ VSila + Pt
Σ (i((())
o (/-1))
b
[Σ V[i(aa)]
21 + +[V((E))
Σ Vi(la vf(x:)
① 10-11
2t b
Σ
Vflla + Pi |s:
2 n-
b
[+())]]
ERR
(V)())) T
Vi(
[V,(( m-1)
(31)
E
From the Lemma B.5
Lemma B.6
Lemma B.7
we have
2 (z))).
Σ
Vii + A [ss]|
(n-1 -m)) S G2 mhh
(32)
b
EEB
7/1
7/8 1-11-11-111
(+)(;));
m 2 (7/-1 -1)) S
(33)
1- 31
(=)( 7
Σ
Vil Pt
[V{(2))]
ff(x))
3 1m-1
2t
b
Σ
Pt |sil
8t
1 @m-1)
EEB
EEB
o パ 2 24 22p
< 2n
(34)
Taking conditional expectation, we have
Ef(x-1)) - f(2))
(35)
< <E(Vf(=)), - 24 Σ Vilai
-
3 me-1) + + Bllec4+ - 20l2
L
[Vf(zx]]
FEB
+
Ft. [V 1 v
2p²² L2L'LLp
G2 12 - mella
2μ
6
31
E (() Vf(m),
2t
Vilaz + Pe
81
Tt-1)
(36)
1BB
where μ > is to be determined.
the term
B(V)(z),
7t b
Σ
Vfl + [[f((.))] ) m)-1)
((()))
(37)
iEB
13

---

Under review as a conference paper at ICLR 2023
the term
L
B Bl - xel?:
(38)
2
and the term
[(V)(=)) - VV((2))
21
Vii
+ Pt [|s]]
8t
o n-
(39)
b
BEB
we introduce the Lemma B. the Lemma B.10 and the Lemma processing and we have Ef(2+11) - -x+(
B.9
We take the expectation over the whole
2h [1,) o VT- 24
7
6
1- B1
2 /m==l ?00?1?0²
2
-VEIV
/n- + 7t
Ellm 11-1
2a²
2
(1 - -1-
13 1 (E] 2 v 12
+
Le
+ (1 B) |1 - -t-
Be
be
+
1
+
dG
y>t
[V)())) 3 /==-=1++
p1²22 př
(22 + 4x3)
(40)
2c(1 - 32
X7
X/
X
2²
2
26
=1
3
+ 9)
Xi
3 v--- [()))
24 2p'noL?p pt
7
c
mell
2n
2a
BE
2
1- 81
26
3
( + 3)
Be
T.p.
+
à²
bc
LG'(( ( )² 1-=0/=²+=/ +1+11B)BBlm m
1-B1
11
2c(1
X7
+ X3
1
X2
dG2
6²
L
2e
pi
(>2 + 422)-
(41)
We set μ² = a² = 8, = 3, Xi
= 1.4 X3
= X3 = 1 and we choose 20c < 1 3.
So we have
Ef(2++1) (xx
16
S (())) 2 11 +
+
Y
BB - 1ella
6
1-3 81
1y + 1riL Lpi. 2
+
o²
+ (4+
33
B El - 111
be
1
1-
Bi
3-2L² L2 o7 dG2 <(1 31)2 e²
+ 5 pz
(42)
2
We re-arrange it and T/ is bounded. We have
7t BIV < 2 2G
f(xe) O 11=
(43)
< <Bf(())+1) + (x() + 45ml2 pi
7
G3 7el12
2x
1 1- 81
4 Lpi
0²
B1
1, dG²
+
+ (4+ ( 1- 31
LG Elie 1- 112 + (1- B1)
(44)
c
be
e
We summary it from = t T - 1, and we assume Vo is a constant, and we have
Σ B[V < 2G Bs(xo) Es(_,)) 45GL² pt
2G3
E
nella
T
(1 (1- BT
(-0
80
Lpi
0² ++
2(4 (1=
T-1
2 32 dGs
+
E.
I - 1----11 + (1- 81)
(45)
6 be
T
e
1=0
(G(
- 8r)
15
203
1 +
8GhrL
Lp? 0²
+
+ (1-3))"
d(
)
+
E
6
e é be
2(1 + (G
L dG3
+
d(e 26-22 + (1 81 e
(46)
T
2G(J(zo) -5")
80 0² 45GL2 ph
(1 20
1 )
8G7L Lpi
TT d
ToT
be
2(4 + (1- ))
6n² L2 dG
d(c 2 - +
(47)
T
(1 3 3
14

---

Under review as a conference paper at ICLR 2023
B.1
TECHNICAL LEMMA
Lemma B.I. Given two vectors a, E Rd, we have (a, 6) I 32 "alll 2i2 |161 for parameter YA
(0,)
Proof.
X² d (a)}
RHS
Σ
_ j=1
(b)]
2 Σ
2
(1)3
2x (b) Σ (4)y1 [0), 2 LHS. (48)
212
2
j=1
j=1
Lemma B.2, For any vector 2.y TRd we have
||x o ul s ||=| x llvl & || x lell
(49)
Proof. The first inequality can derived from that Σ (atvt) S Σ - 11/1))) The second inequality follows from that lull.. S 11
Lemma B.3. oi is bounded, iee, dim S (mhess 5
Proof. As the gradient bounded by G and (m)s ((ee)s
1 dx Follow the update rule, we have
S (m)s S
Lemma B.4. For the term defined the algorithm, we have
s d ②
Σ |n
nell
T
(50)
Prouf. (1p)),, the i-th dimension of f/. deceases as increesess. So we have
T-1
T-1 ((v=1); - (ne)sl
E
||
mill
E
T
T
1=0
1=1 1=
$
Σ ((1
(
G
1
T 22
(51)
Lemma B.5. For the term defined in the algorithm, we have
(V)()) Σ
Va(x + Pt ||s,1| @ (n - n)) S G. mell
(52)
Proof.
(V)(2)),
71 Σ
Vf(e + Pt "1a++]'' P
e (nt-1 7)))
b
iED
1
Vf((z)) o (7-1 1n-1-1))al
<
(V](z))) x
16
file + Pt
(53)
5
V(((z)]]
1=1
iEB
Σ (xxi))
$ GG Σ
Vf((# - P:
@ (7:11- - m)))m
(54)
Σ v((())))]
< G
Σ V((EE))
((V A
) 112-1-m))6n
(55)
b
Σ f((()))]
1=1 iEB
0
Σ fx(e))
(Vfl(z +
)G) x (n--1 - 1
(56)
b
Σ f(())l
1
15

---

Under review as a conference paper at ICLR 2023
1 G² d
Σ
Km--1 m)m1
(57)
5=1 IEB
7x |1--1 7/ll12
(58)
Lemma B.6. For the term defined in the algorithm, we have
(V)(2)) 1 m @ (9t-1 -h)) S 739 11 Dell
7
(59)
1- 31
Proof.
(Vf(ze), 1- B1
m4-1 (m--1 ne))
d
7/81
$ 1
(V](z)))un x (mm- @ (mp-1 m))6ul
(60)
j=1
7/1
(mm--1 @ (n--1 1101117))a m))6o
(61)
1 B1
|=1
7/1
Σ
- (m)-1) x ( -n)c
(62)
1 B1
781 G²
Σ
d
(63)
1
3=1
7/1
G²1 (1 - 7ell2
(64)
1 Bi
Lemma B.7. For the term defined in the algorithm, we have
(V(())) 24 Σ Vi
Pt
[V}(>))]
Vf(r.)
7t Σ
Viilre pe P[
81
① m)-1)
BEB
B
s Vffxx) o " 2μ 21²
(65)
2/
Proof.
(Vf()))) 22
Vilxx + P: [+))]]|]
V((=))
② 17-1
7t
Σ
fi + Pt
St
+ 1--1)
b
EEB
KEB
= (Vf(x))
11 7t
Vf +
Vf(r,) [+f(x))|
Vf( + Pt
ΣiEn Vf((z.)
1) vT--1)
DiED ff()))l
(66)
n
V((Tt) [()()))
EieB V(((z))
Σ
(
- Vf(aE + Pt
Eien (f(())]] 1) 211
26
74
[V(())
(67)
2/2
1
P (M())
MED Vf((z)) Vf(m At
DieD Vf((z)]] @ -
2b
VV @ 1
(68)
2n²
1.
vf(x.) [x()))]
EIED V.(((.))
< +
Viia +
- V/(a + PT PT EiE v((e;)]] )11² x 1lym-1
2b
7t
[V
1T=
(69)
20²
16

---

Under review as a conference paper at ICLR 2023
< Vf()))
7K 2p
3
v
p
Σ
Vf(e))
1EieD V(((a))] 12
EiE vf((et)
(70)
2be
[[(=)]]
LeL
$ (ax))
(71)
2n
Lemma B.8. For the term defined in the algorithm, we have
7tt b EE
Σ V
V f(())
@ m--1)
f((x))
76a2 ph
< <-xIVV((x) n 11 +
7+ 20
Vf(x)) 3 -ll
(72)
2
Proof.
B(V)(z)), Σ EB
vf(xs) Vfil + P: [V+()))]|
10 ()12
[((x))
= ->>>) - [VV
[(V((=)) Vfila, + Pt [V()))] 107-1) (73)
- +(Vf 2
Σ
(V)().)) - Vf(n:
V/(z))
)) 0 m--1 (74)
[V
IEB
" 20
[vf(z)) @ -
+ na?
(f+)
2 Bll
b
- Vf((2 (V.(((a)
[z]))]] )) 1
(75)
EB
S {f |-
E π
f)()) 3 |
2a
ma² 1
(
+ 2e
BlI
(v.( - fx(x
V )1 ))(F
(76)
< Vf(x))
- + E 74
V -カ
2a
ma?
+ 2bc
E
(z)) - V/(
(77)
< <->| ((x) 2 + F. 7
f 3 /
vf(..)
|
(78)
2x
2bc
[V](x.]]
LEB
= = O - + E
(
(79)
2
2e
Lemma B.9. For the term defined the algorithm, we have
BVV (I
Y b
V f(
+
Pt |s m--1)
FE
L2 1
+
1
dG
+
V V[(x))
11
+ 2c
2L'
(A2 + 123).
(80)
20(1 B1)3
x
X3
Proof.
E(V[(=)) Vs(a.), 2/1
Σ
Vi((a + pt |s.l
Ø 1)-1)
(81)
b
EieB fi(z))
=>B((V[(>)) -/(=z) ---
Viilzr + P:
TEEE V/(xz)]] 3 m
(82)
1EB
=>B((VS(<)) vS(x)) - (Zx) 1 MT-
Vf(z)) v(((zs)) @ v
-7E((V/(<)) v.(2,)) R 1 Σ (v +
EB
17

---

Under review as a conference paper at ICLR 2023
EieB Vii(et)
[f(er)
+
Σ E8
(Vii Pt
M |iED Vsi(az)] V
@ v
(83)
7 1
7 + [[][(()))) o // 11
[V o vn
2 X7
X3
2
7/3
Vf(z)) V(f((a)) @ v1=1
1
Σ
f. +
[V{())]
2
23
Σ
fa Pt
1 Vi((ai]] Vf(T
EB V)(m)
V (21)
(84)
2
VEB
1
2
[BI(()(()) - vf(=)) O /
fxx)))))))))))))))){())) @ /
2
X3
2
3AZ12
p3
L2
(85)
26
2
1
111 + 72 IVf(z)) @ @VMTT
+
|B|
2E
X
X2
X
2
3ALL2 pL
22
(86)
2E
L
1
1
1
| 11 -111-11+++
7a; V\(xr) @ v|
2c(1 - 31)2
X3
X3
2
(87)
2c
1
dG
+
7A7 Vf(z)) 2 2 yT'P pi
(A2 + 423).
(88)
2c(1 - 31
x7
X3
X3
6²
2
2e
Lemma B.10. For the term defined in the algorithm, have
L
2 [ 20 - (1 (1-3,) 31) I|ne 1/-1/1
+ +33
1 (BI\V[(z)) R v
Lpi o² a + (1+ 3) - -m--113)
c +
(89)
Be
Proof.
L
El|22411 - 2
2
-2E1 El|
L
7/11
m 2 (n - 7-1) - 7 2 mlll
(90)
B1
<
E
f1| o (m Th 11 + LE Vt
Vill + pt "sslll
)) o m
(91)
(1 B.)=
LO
Blin: n 11 +
LE
st (i At ||
)) nell
(92)
(1 B1)²
b
11 Σ
fiv + Pt |
2 9- + 5
Σ
(i + Pt |
)) C (ne n-1)
28t Bllm 1112
(93)
(1 81)2
288
E Blly Tp 11 + 7 1 +
1
1Bl Σ
1
(Vi((t + + Pt
8t
1) m0
(1 21)²
3
|sa
+ + || ; Σ v fi
St
(n n 1)12)
(94)
s
S + 1 10811 (VS + PC ||sa||
)) n- 2 + (1 + - m
7G (1- 81)2 (1-8))²
Bl1r 11
(95)
18

---

Under review as a conference paper at ICLR 2023
S E + 105 1 Σ (V + Pl ||
8t
}) vi=--1/ x x0o-iIlxxIvvnn-.... 11v1-11
+1+ BB -- + (1 - $1)
F. F.lp 1112
(96)
B11 1
fi + Mr |se
)) 3 / T 2 + (1 B) -1-1)
Be
E 11 1-12
(97)
(I B1)2
L3 1+ L ((V Mi-1 |1( -
- vf(x))) -
BE
EiEB (((()))
+11 (VVii PC "' V.((.,)]]
vfi(z))) VV- + (1 B) -
Bllmc 1112
(98)
(1 B1)²
v. L(3 1+ B
EIEB Vf((T.)) "' V(i(a.)]
(EIVS(() M-
11
[[((x+)))
111
Be
L 3
+ (1
111) 3)
E E 1
(99)
bc
(1 31
L/L(3
1+ (())) 1-
E 1
1
Σ
(VV.iia P
EIED Vfs(.+) DiED V[(((z)] (/((,)))]
3c
8² bc
9) - + LG2 p
+ (1
E Elm 1-1
(100)
(1 31)2
1+ 3
EIED Vf((t))
LB3
f())))()))))) M-
Σ VV + P
"iEB Vf((z:)]] {((x)]
8r
eh
8² be
+ (1 9) - -
E Emm 7-1
(101)
(1 3, )2
1....33 B f(x)) 1
Lpi
2²
+ -
Be
be
2g3 (1 3,)2
|| 11
(102)
C ADDITIONAL EXPERIMENT ILLUSTRATIONS
C.I EXPERIMENT ILLUSTRATIONS
We conduct the experiments on the GLUE benchmark with AdaSAM, AMSGrad, SAM and SGD, respectively. The optimizers not have the momentum part (8, = 0). As a supplement to Table and Table Figure 2 and show the detailed loss and evaluation metrics versus number of steps curves during training. The loss curve of AdaSAM decreases faster than SAM and SGD in all tasks, and it has similar decreasing speed as the AMSGrad. The metric curve of AdaSAM and AMSGrad show that the adaptive learning rate method is better than SGD and SAM. And AdaSAM decrease as faster as the AMSGrad in all tasks.
C.2 ABLATION STUDY
this section, we conduct experiments evaluate the impact momentum (8, 0.9) on different optimizers. We show the experiment results Figure and Figure respectivlly. The results are also illustrated in Table and Table 2 SAM with adaptive learning rate (AdaSAM with 0.9) converges as fast as AMSGrad. Both AdaSAM and AMSGrad are faster than SGD with momentum and SAM with momentum. The generalization ability of AM with momentum much better than SAM with momentum in all tasks. Besides, AdaSAM also better than AMSGrad in GLUE benchmark except for QNLI and QQP tasks.
19

---

Under review as a conference paper at ICLR 2023
1.15
5G0
1.15
1.00
Training Loog
0.85
10.00
Toshaing
2.70
ALL SAME30
70
0.55
2.55
0.40
3120
100
2.0
3.00
500
000
100
2.01
0.00
5.00
100
2/0
0.01
800 Steps
1000
1200
100
Steps
Stege
(a) MRPC
(b) RTE
(c) CoLA
15
02
SED
SED
5GD
SANDMES
SAGAUAR
SANDNAGI
SAGEE (001
20
100
200
3000
4100
600
G80
40
100
2000
2000
400
600
200
480
(200
100
1000
1210
1600
Steps
Steps
Steps
(d) MRPC
RTE
CoLA
1.30
0.85
Training loES
2.15
0.70
2.00
0.50
0.85-
0.95
02.70
0.40
0.20
9.001
A.55
D1.0-
0.25
0.10
1000
20:00
3000
1000
5000
300
200
100
600
800
1000
2500
5000
7500100000
250 50000 750
Steps
Steps
Steps
(g) SST-2
(h) STS-B
(i) MNLI
0.25
0.20
3.15
SAVAL 01
2
SUSSE
Gad53a 86)
.03
1000
2000
3000
2000
5000
300
200
400
603
800
5000
10000
15000
20000
Strps
Steps
Steps
SST-2
(k) STS-B
(I) MNLI
Figure 2: The loss and evaluation metric v.s. steps on MRPC, RTE, CoLA,, SST-2, STS-B and M5 = 0)
20

---

Under review as a conference paper at ICLR 2023
1.00
1.0
Teiiinng lose
0.85
TYAisiNg Jana
0.55
SANTURR)
0.70
AAGIVAAY
0.55
0.40
0.25
0.25
Steps
Steps
(a) QQP
(b) QNLI
303
1.00)
00
Faining 1ook
01.50
89
0.70
85
00.00
500014)
0.25
2.00
1000
6000
5000
10000
40.0
6000
Steps
Steps
(c) QQP
(d) QNLI
Figure 3: The loss and evaluation metric v.s. steps on QQP and QNLI((3 0)
21

---

Under review as a conference paper at ICLR 2023
550
1.00
SAGBE
1.00
1.00
SAGAE
SASISDEN
Tralning boss
0,70
3.70
A.L ALSAA0300
8.70
Todning
Trsholng
0.55
2.55
8.55
0.40
0.20
3.20
0.20
1000
2.00
3.00
500
000
100
2.01
0.00
5.00
100
2/0
0.01
800
1000
1200 1007
Steps
Stepe
Steps
(a) MRPC
(b) RTE
(c) CoLA
90
SCD
SED
SGD
SANDMO
SANDRAH
SAGE1 BG)
3
080
100
2000
30.0
410
600
G80
400
100
2000
2000
400
600
200
480
(200
1000
1000
1210
1600
Steps
Steps
Steps
(d) MRPC
RTE
(f) CoLA
1.10
1.00
0.25
3.20
Training LOES
0.70
2.00
0.55
RL.5
0.85-
0.40
3..0
0.000
0.25
0.25
0.10
1000
20.00
3000
1000
5000
3.0
300
100
600
8:0
1000
2500
5000
7500000000
P50 50100 750000000
Steps
Steps
Steps
(g) SST-2
(h) STS-B
(i) MNLI
0.25
0.20
A33SCaa
0.15
4442141
/050
.03
1000
2000
3000
2000
5000
300
200
400
603
800
5000
10000
15000
20000
Stepp
Steps
Steps
① SST-2
(k) STS-B
(I) MNLI
Figure 4: The loss and evaluation metric v.s. steps MRPC, RTE, CoLA, SST-2, STS-B,
MNLL(3 = 0.9)
22

---

Under review as a conference paper at ICLR 2023
1.154
100
SAMBSE SANTURA)
Teitiing leas
0.85
TYAlning Jana
3.55
0.70
AAREVART
0.00
0.20
0.25
Steps
Stops
(a) QQP
(b) QNLI
85
4000
60
84
2000
1000
6000
5000
10000
%
1000)
3000
Steps
Steps
(c) QQP
(d) QNLI
Figure 5: The loss and evaluation metric v.s. steps on QQP, QNLI( = 0.9)
23


Reviews:
[{"review_id": "yA3bjj02gN", "reviewer": "Reviewer_w96Y", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors propose an adaptive algorithm for sharpness aware minimization, a special type of regularized loss function that has a min-max structure. The adaptive algorithm is reminiscent of Adam. A convergence analysis to a first-order stationary value of the original unregularized loss is presented. The convergence of course requires that the size of the perturbation $\\rho$ decreases to zero, which is weird because it means that no regularization is applied, but is consistent with the fact that the convergence is stated with respect to the unregularized loss. Experiments are presented showing that the algorithm can improve over some baselines.", "summary_of_the_review": "The problem solved is not interesting because it is already solved by SGD (guarantees of the form of equation (7)). The algorithm seems to perform well in experiments but the theoretical results derived have no relation to the problem studied (sharpness aware minimization)", "strength_and_weaknesses": "The main weakness is that the theoretical results are not well motivated. As can be seen from equation (7) in Corollary 3.7. the main claimed contribution is that the algorithm converges (in expectation) to a first order stationary point of the smooth function $f$ which is defined in equation (2) and is a simple empirical risk minimization problem. The issue is that the goal of Sharpness aware minimization is not to minimize the original function but rather the regularized objective that is presented in equation (3). That is, if we wanted to achieve convergence of the form stated in equation (7) then we already have good algorithms for it: plain SGD on the objective $\\sum_i f_i(x)$.\n\nThis is not surprising as the \"convergence guarantee\" is given when the regularization parameter $\\rho$ tends to zero. But in this case what is the point? in this case the algorithm wouldn't be performing sharpness aware minimization at all!\n\nA real interesting question is to obtain first-order optimality of equation (3) or equation (4) which are the real and approximated version of the sharpness aware minimization objective. Of course in this case the argument is much more difficult because (3) and (4) are non-smooth and non-convex objectives given the presence of a min-max structure in the case of (3), and the presence of a $\\| \\nabla f(x)\\|$ in a denominator, in the case of (4). First order optimality guarantees in this case would involve more difficult concepts like the Clarke subdifferential.\n\nWhen $\\rho$ does not tend to zero equation (5) simply states that the algorithm converges to some vicinity of a first order stationary point of the original unregularized function $f$, which might not be interesting or useful. Overall, I believe that the authors are getting convergence rates for \"something\" for the sole purpose of showing some rate, but it is not an interesting problem to get a gurantee like (7) which, again, is already obtained by a simple algorithm like SGD.\n\nThe main strength is that the experiments appear to be extensive and the algorithm shows some practical improvements. However this is all eclipsed by the fact that the convergence analysis is presented as the main contribution.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "mF-chtHnf3", "reviewer": "Reviewer_5Rfi", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes an algorithm that combines adaptive step size (ADAM/AMSGrad) and sharpness-aware minimization (SAM). The author first show the convergence of this algorithm under standard assumptions (Lipschitz functions, finite variance, and uniformly bounded stochastic gradients.) Finally, they try their algorithm experimentally and compare it with some baselines. ", "summary_of_the_review": "Overall. I find the theoretical contribution significant but the experimental validation is under ICLR standards. Since Adam is mainly used because of its empirical success, it is essential to show that AdaSAM can be efficiently used in practice. Because the experimental section could be significantly improved, I consider this work marginally below the acceptance threshold. (this justify my current rating on correctness) \n", "strength_and_weaknesses": "## Strength:\n- The theoretical analysis seems sound. \n- The story and the intention are clear. \n- The experiments are done with large-scale models. \n## Weaknesses:\n- I find the experimental could be significantly improved. Some results seems inconsistent and some comparisons are missing. Overall it is not clear if the comparison has been done fairly and whether adaptivity really improves SAM. (see more details in the reproducibility section)\n- Many sentences are relatively ambiguous. It makes reading the paper more difficult than it could. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Wpwq80ehNaQ", "reviewer": "Reviewer_x2Pp", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "the paper proposes the modification of AdaSAM to SAM (sharpness-aware minimization) that records running averages of first and second moment (based on the \"SAM gradient\"), and adjusts the parameter updated based on these estimates a la Adam and AdaGrad. authors show improved results with AdaSAM on GLUE benchmarks. ", "summary_of_the_review": "the paper proposes the modification of AdaSAM to SAM (sharpness-aware minimization) that records running averages of first and second moment (based on the \"SAM gradient\"), and adjusts the parameter updated based on these estimates a la Adam and AdaGrad. authors show improved results with AdaSAM on GLUE benchmarks. the work has some nice theory, but there are mismatches between the theoretical and empirical results, and empirical results can be improved in terms of including additional tasks and baselines to make a stronger argument.", "strength_and_weaknesses": "strength: the paper studies an extension of SAM which to the best of my knowledge has not been studied before (although the idea seems natural, given precursors like Adam, RMSProp, and AdaGrad). \n\nweaknesses: \n- the paper shows limited empirical results regarding AdaSAM's advantage compared to optimizers (the intro claims that \"The results show that AdaSAM outperforms the most of state-of-art optimizers\"). more concretely, i think that to strengthen the empirical advantage argument, authors could try showing that\n  - AdaSAM is competitive for tasks in other domains (e.g., the simplest image classification)\n  - AdaSAM has an advantage over more \"state-of-the-art\" baselines (e.g., ASAM [1])\n- a discrepancy between the theory and experiments: the theoretical result is about training-time convergence rate, whereas the it seems the results with GLUE are about test-time generalization performance (i wasn't able to pinpoint where the latter point was made, but i'd assume the numbers are from evaluation on the test spilt; unless i'm missing something obvious, authors should also clarify in the main text)\n\n[1] Kwon, Jungmin, et al. \"Asam: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks.\" International Conference on Machine Learning. PMLR, 2021.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "-juHFssfk1", "reviewer": "Reviewer_rzri", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper considered an adaptive variant of Sharpness Aware Minimization (AdaSAM) and theoretically showed the $O(1/\\sqrt{bT})$ convergence rate and linear speedup property with respect to minibatch size $b$. The authors conducted several experiments including GLUE and showed that the AdaSAM with the best-tuned $\\rho$ outperforms Amsgrad and SGD.", "summary_of_the_review": "The paper studied an adaptive variant of SAM with corresponding complexity analysis. The technical contribution is unclear and the experiment is weak. I vote for rejecting the paper in its current form.", "strength_and_weaknesses": "Strength:\n\n\nThe paper is well-motivated due to the good empirical performance of SAM in the deep learning applications, and its extension to adaptive learning rate is an interesting topic.\n\nWeaknesses:\n\nW1: The main theoretical results (Theorem 3.6 and Corollary 3.7) are nothing surprising. It is easy to get the same convergence rate by simply running SGD on the original function $f$ but without a bounded gradient assumption (Assumption 3.4 is not required for SGD). In addition, Corollary 3.7 only considers a small radius $\\rho=O(\\sqrt{1/bT})$, which makes the proof extremely easy because the $\\ell_2$ norm of the gradient of the original function and the gradient of the perturbed function can be bounded by $\\rho L$ due to $L$-smoothness.\n\nW2: The technical contribution of this paper is unclear to me. I do not see any technical difficulty by extending the analysis from SGD to SAM. It is also unclear to me what is the concrete meaning of adaptive learning rate in this paper. It seems to be that the algorithm still requires knowledge of the problem-dependent parameters such as the variance and the gradient upper bound, at least from formula (5) in Theorem 3.6.\n\nW3: The paper lacks clarity. I do not understand what is the meaning of $\\epsilon$ in the statement of Theorem 3.6 (I guess this is the missing constant in the denominator of the line 9 of the Algorithm 1 to make $\\eta_t$ upper bounded by $1/\\epsilon$). The effects of $\\beta_1$ and $\\beta_2$ are not explicitly considered in the statement or the proof.\n\nW4: Experiments are rather weak. The authors did not compare best-tuned SAM versus AdaSAM directly. The paper also did not compare with adaptive variant such as Asam [Kwon et al. ICML 2021].\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject