Adaptive Inertia: Disentangling the Effects of Adaptive Learning Rate and Momentum
Authors: Zeke Xie
Xinrui Wang
Huishuai Zhang
Issei Sato
Masashi Sugiyama
Venue: ICLR 2022 Submitted (2022)
Keywords: Learning Dynamics
Diffusion
Stochastic Optimization
Momentum

Abstract: Adaptive Momentum Estimation (Adam), which combines Adaptive Learning Rate and Momentum, would be the most popular stochastic optimizer for accelerating the training of deep neural networks. However, empirically Adam often generalizes worse than Stochastic Gradient Descent (SGD). We unveil the mystery of this behavior in the diffusion theoretical framework. Specifically, we disentangle the effects of Adaptive Learning Rate and Momentum of the Adam dynamics on saddle-point escaping and flat minima selection. We prove that Adaptive Learning Rate can escape saddle points efficiently, but cannot select flat minima as SGD does. In contrast, Momentum provides a drift effect to help the training process pass through saddle points, and almost does not affect flat minima selection. This partly explains why SGD (with Momentum) generalizes better, while Adam generalizes worse but converges faster. Furthermore, motivated by the analysis, we design a novel adaptive optimization framework named Adaptive Inertia, which uses parameter-wise adaptive inertia to accelerate the training and provably favors flat minima as well as SGD. Our extensive experiments demonstrate that the proposed adaptive inertia method can generalize significantly better than SGD and conventional adaptive gradient methods.

Under review as a conference paper at ICLR 2022
ADAPTIVE INERTIA: DISENTANGLING THE EFFECTS OF ADAPTIVE LEARNING RATE AND MOMENTUM
Anonymous authors Paper under double-hlind review
ABSTRACT
Adaptive Momentum Estimation (Adam), which combines Adaptive Learning Rate and Momentum, would be the most popular stochastic optimizer for accelerating the training of deep neural networks. However, empirically Adam often generalizes worse than Stochastic Gradient Descent (SGD). We unveil the mystery of this behavior in the diffusion theoretical framework. Specifically, we disentangle the effects of Adaptive Learning Rate and Momentum of the Adam dynamics on saddle- point escaping and flat minima selection. We prove that Adaptive Learning Rate can escape saddle points efficiently, but cannot select flat minima as SGD does. In contrast, Momentum provides a drift effect to help the training process pass through saddle points, and almost does not affect flat minima selection. This partly explains why SGD (with Momentum) generalizes better, while Adam generalizes worse but converges faster. Furthermore, motivated by the analysis, we design a novel adaptive optimization framework named Adaptive Inertia, which uses parameter- wise adaptive inertia to accelerate the training and provably favors flat minima well as SGD. Our extensive experiments demonstrate that the proposed adaptive inertia method can generalize significantly better than SGD and conventional adaptive gradient methods.
INTRODUCTION
Adam (Kingma and Ba 2015). which combines Adaptive Learning Rate and Momentum, would be the most popular optimizer for accelerating the training of deep networks. However, Adam often generalizes worse and finds sharper minima than SGD [Wilson et al. 2017) for popular convolutional neural networks, where the flat minima have been argued to be closely related with good generalization
(Hochreiter and Schmidhuber 1995 1997a Hardt et al. 2016 Zhang et al. 2017).
Meanwhile, the diffusion theory has been used as tool study how SGD selects minima (Jas- trzkebski et al., 2017; Li et al. 2017 Wu et al., 2018 Xu et al., 2018: Hu et al. 2019 Nguyen
al.1 Zhu et al. 2019; Xie et al. 2021b; Li al. 2021). This line of research also suggests
that injecting or enhancing SGD-like sharpness searpness-depedenn gradient noise may effectively help find
flatter minima (An 1996, Neelakantan al. 2015) Zhou et al. 2019 Xie al. 2021 1a HaoChen et al. 2021). Especially, Zhou et al. (2020) argued the better generali lization performance of SGD over Adam by showing that SGD enjoys smaller escaping time than Adam from a basin of the same local minimum. However, this argument does not reflect the whole picture of the dynamics of SGD and Adam. Empirically, does not explain why Adam converges faster than SGD. Moreover, theoretically, all previous works have not touched the saddle-point escaping property of the dynamics. which is considered as an important challenge of efficiently training deep networks (Dauphin et al.) 2014 Staib et al.] 2019 Jin al. 2017 Reddi et al. 2018).
Our work mainly has two contributions. 1) We disentangle the effects of Adaptive Learning Rate and Momentum in Adam learning dynamics and characterize their behaviors in terms of saddle-point escaping and flat minima selection, which explains why Adam usually converges fast but does not generalize well. Particularly, we prove that, Adaptive Learning Rate is good at escaping saddle points but not good at selecting flat minima, while Momentum helps escape saddle point and matters little to escaping sharp minima") 2) We propose a novel Adaptive Inertia (Adai) optimization framework
When we talk about saddle points, we mean strict saddle points in this paper.

---

Under review as a conference paper at ICLR 2022
Table I: Test performance comparison of optimizers. We report the mean and the standard deviations (as the subscripts) of the optimal test errors computed over three runs. Our methods, including Adai and AdaiW, consistently outperforms all other popular optimizers.
AMSGRAD
RESNETLS
1.
4.74c.3
6.53
6.16m.s
c7
5.8fa.12
6.01a.1
5.81u.r
6.00uus
6.42
7:33
7.346.24
6.48x33
6.76 12
6.15xus
6.00b.22
6.36x.c
CIBAR-
RESNEY24
21.05.ae
70uss
21.52x55
27.16u.5
25.536.ca
22.90l...
22.876.33
22.72x.10
20.570.12
21.41:.0
DENSrNar1221
19.44um
..
19.8
24-43c.s
22.00x.is
21.30ss
22.27.20
GooGLENNT
20.50..s
20,60,.2
21.21wa
26.12aa4
25.606.47
23.29
23.15usa
21.82 v
24.24y.0
22.23,.15
which is conceptually orthogonal to the existing adaptive gradient framework. Adai does not wenaeetr-wwiell adjust learning rates, but parameter-wisely adjusts the momentum hyperparameter, called inertia. We show that Adai can provably escape saddle points fast while learning flat minima well. From the empirical results (see Table []. Adai significantly outperforms SGD and popular Adam variants.
The paper is organized as follows. In Section we first introduce the dynamics of SGD and analyze its behavior saddle-point escaping. which serves basis to compare with the behavior of Adaptive Learning Rate and Momentum. Section B we analyze the behavior of Momentum. In Section we analyze the dynamics of Adam. In Section | we introduce the new optimizer Adai. In Section 6 we empirically compare the performance of Adai, Adam variants, and SGD. Section concludes the paper with remarks.
SGD AND DIFFUSION
In this section, introduce some preliminaries about the SGD diffusion and then present the saddle-point escaping property of SGD dynamics.
2.1 PREREQUISITES FOR SGD DIFFUSION
We first review the SGD diffusion theory for escaping minima proposed by Xie et al. (2021b). We denote the model parameters as 0, the learning rate as 1), the batch size as B, and the loss function over one minibatch and the whole training dataset as (()) and L(0), respectively. A typical optimization problem can be formulated as ming L(0). We may write the stochastic differential equation/Langevin Equation that approximates SGD dynamics (Mandt al. 2017 Li et al. 2019
d8 VL(O)dt + [nC()))
where dW, (0, Idt), / is the identity matrix, and C(0) is the gradient noise covariance matrix. The gradient noise is defined through the difference of the stochastic gradient over one minibatch and the true gradient over the whole training dataset, & = vi - ( is well known that the Fokker-Planck Equation describes the probability density governed by Langevin Equation (Risken) 1996; Sato and Nakagawa 2014). The Fokker-Planck Equation is
OP(,, v [P(O:))VL(0)) -- ((O0))
where V. is the divergence operator and D(0) x0(0) is the diffusion matrix (Xie et al. 2021b). We note that the dynamical time is equal to the product of the number of iterations 7 and the learning rate Ty: = T/..
As the gradient variance dominates the gradient expectation near critical points, we have
D(0)
(C 2
n
(L, j=1
N
7 FIM(A) 2B
17 [H]]) 2B
(3)
2B
near critical point c, where L, (0) is the loss function of the j-th training sample, ( is the Hessian of the loss function at 6, and FIM(0) is the observed Fisher Information matrix, referring to Chapter 8 of Pawitan (2001) and Zhu al. (2019). We further verified that Equation approximately
2

---

Under review as a conference paper at ICLR 2022
holds even not around critical points in Figure of Appendix C Equation 33)
also proposed by
Jastrzkebski et al. (0017) and Zhu et al. (2019) and verified by Xie et al. (2021b) and Daneshmand et al. (2018). Please refer to Appendix for the detailed analysis of the stochastic gradient noise.
If we have eigendecomposition H = U diag .......... Hn-1. 1, Hy H then we use [-]+t to denote the transformation that [H]+ U diag( H1 , ..... Hn 1, Hr, I)U The i-th column vector of U is the eigenvector H corresponding to the eigenvalue H,-
In the following analysis, assumption is common and mild, when we focus on the behaviors near critical points (Mandt et al. 2017 Zhang al., 2019a: Xie al., 202Tb). Note that, by Equation (3) and Assumption [ the diffusion matrix D is independent of f near critical points.
the second-order Taylor approximation near critical points. This
Assumption 1. The loss function around the critical point can be upproximately written as
L(O) - L(c) + 1 (0) H()((0
2.2 SGD DIFFUSION NEAR SADDLE POINTS
In this subsection, we establish the saddle-point escaping property of SGD diffusion as Theorem [
Theorem (SGD Escapes Saddle Points). Suppose c is critical point, Assumption 7 and Equation
3 hold, the dynamics is governed by Equation (I), and the initial parameter the saddle point
0 Then the probability density function of f after time t is given by the Gaussian distribution (c. E()), where E(t) = ..... 0 and
D: D exp(-2H,t)),
H;
where Dj is the i-th eigenvalue of the diffusion matrix D and H; is the i-th eigenvalue of the Hessian matrix at The column vectors of U are exactly the eigenvectors of H. The dynamical time near saddle points, we have
t =. In terms of SGD notations and |n <
2T
(T)
+ HH
B
The proof is relegated to Appendix A.1 We note that 1) if H1 > O, the distribution of 0 along the direction i converges to a Gaussian distribution with constant variance, / (c., 2;: ); and 2) if H; < O, the distribution is Gaussian distribution with the variance exponentially increasing with time t.
As the displacement Ad, from the saddle point c can be modeled as a center-fixed Gaussian distribu- tion, the mean squared displacement is equivalent to the variance, namely (A) of (t). The result means that SGD escapes saddle points very slowly ((A??) O(HH; H, is close to zero. Note that, in the diffusion analysis, the direction denotes the direction of an eigenvector instead of coordinate's ' direction. While SGD updates model parameters along the coordinates, we do not need to treat the coordinates' directions specially in the continuous-tin analysis.
ANALYSIS OF MOMENTUM DYNAMICS
In this section, we analyze Momentum in terms of saddle-point escaping and flat minima selection.
The continuous-time Momentum dynamics. The Heavy Ball Method (Zavriev and Kostyuk 1993) can be written as
{
m. = -3m + Bagt:
(4)
0c+1 0, 1711
where 6n and 33 are the hyperparameters. We note there are two popular choices, which are. respectively, f corresponding to SGD-style Momentum and /33 = - 19 corresponding to Adam-style Momentum, namely the exponentially moving average.
We can write the motion equation in physics with the mass M and the damping coefficient y as
r; = (1 - -d + F dt
37
(5)
0 = 02 + r;dt,
3

---

Under review as a conference paper at ICLR 2022
where -, F gt. dt 1,. ydt B1, and dt form of the motion equation as
= B3. Thus, we obtain the differential d
MÒ -M +
where 0 = d'
and 0 σ The left-hand side the inertia force is equal to the mass M times the acceleration 8, the first term in the right-hand side as the damping force is equal to the damping coefficient y times the physical momentum M0, and the second term in the right-hand side equal to and the damping coefficient the external force F in physics. We can easily obtain the mass M Pa
dez
by comparing 5 and (6
As F corresponds to the stochastic gradient term, we obtain
aL(O)
Mdè = ~ddo
dt + [2D]
(7)
00
Its Fokker-Planme Equation in the phase space (the 0-0 space) well known as 0P(0,r,t)
V |rP(0,r, v, [+r M M P(O,r,))
01
V M 2 P,,,O
(8)
where r = Ó. Equation 8 is not specialized for learning dynamics but a popular result in Langevin Dynamics literatures (Risken 1996 Risken and Eberly 1985 Zhou 2010 Kalinay and Percus 2012).
Equation (57) of Radpay (2020) gives an exactly same form for describing finite- initeiinertia Langevin
Dynamics, Equation 77. We contribution is the first apply it to deep learning dynamics.
Momentum escapes saddle points. We formulate how Momentum escapes saddle points as Theorem whose proof is relegated to Appendix A.2
Theorem (Momentum Escapes Saddle Points). Suppose is critical point, Assumption 7 and Equation (3) hold, the dynamics is governed by Momentum, and the initial parameter at the saddle point 0 Then the mean squared displacement near saddle points is
D; (A) = 733 73 -1-
D: + YMH,
2H,t yM )
-1eppp
(9)
where AB(t) = ftt) - (()) is the displacement of 9, and (.) denote the mean value. In the right-hand side, this first term is the momentum drift effect and the second term is the diffusion effect. As & HT
near ill-conitiooned saddle points, it can be written in terms of Momentum notations as
13n
2
(AOP) 2 - 3133B [1 exp (- - --0-
B(1 B1) +O H2TTTT)
By comparing Theorems M and 2
we notice that, SGD escapes saddle points only due to the diffusion
effect (similar to the second term in Equation ()), but Momentum provides an additional momentum drift effect (the first term in Equation ( for passing through saddle points (Wang et al.) 2019). This momentum drift effect has not been mathematically revealed before, the best of our knowledge.
Momentum escapes minima. We first introduce two classical assumptions.
Assumption 2 (Quasi-Equilibruum Approximation). The system is in quasi-equilibrium near minima. Assumption (Low Temperature Approximation). The system is under low temperature (small
gradient noise).
Xie et al. 2202bb) modeled the process of SGD escaping minima as Kramers Escape Problem in
deep learning. Recent machine learning papers (Jastrzkebski et al. 2017 Xie et al. 202 Zhou et al. 2020) also used Assumptions and as the background implicitly explicitly. Quasi-Equilibrium
Approximation and Low Temperature Approximation have been widely used in many fields' Kramers
Escape Problems for state transition/minima selection, including statistical physics(Kramee 1940)
Hanggi 1986). chemistry(Eyring! 1935 Hänggi et al.) 1190), biology(Zhou 2010). electrical engi-
neering and Kalmykov 2012). and stochastic process Van Kampen 1992 Berglund 2013)
4

---

31,0)90 Under review as a conference paper at ICLR 2022
Assumptions 2 and mean that the diffusion theory is good at modeling "slow" escape processes that cost more iterations. As this class of "slow" escape processes takes main computational time compared with "fast" escape processes, this class of "slow" escape process is more interesting for training of deep neural networks. Empirically, Xie et al. (2021b) reported that the escape processes in the wide range of iterations (50 100,000 iterations) very well. Our empirical results in Section 6 support this point again. Thus, Assumptions 2 and are reasonable in practice (see more discussions Appendix B
be modeled Kramers Escape Problem
We next formulate how Momentum escapes minima as Theorem 3 which is based on Theorem 3.2 in Xie et al. (2021b) and the effective diffusion correction for the phase-space Fokker Blakkrr-laank Equation (8) in Kalinay and Percus (012)) We analyze the mean escape time T (Van Kampen 1992 Xie et al. 2021b) required for the escaping process from a loss valley. Suppose that the saddle point b exit of Loss Valley a, and AT, = L(bb L(a) is the barrier height (See more details in Appendix ED.
Theorem (Momentum Escapes Minima). Suppose Assumptions 2 and and the dynamics governed by Momentum. Then the mean escape time from Loss Valley a to the outside of Loss Valley a is given by
2YMBAL (
(1 (=+01-)))
|Hbll
exp
17
Har |Heel
where the subscript indicates the escape direction, s E (0,1) is patt-dependent parameter, and Ha and Hee are the eigenvalues of the Hessians of the loss function at the minimum and the saddle
< 1, it reduces to SGD. In terms of Momentum point b along the escape direction When 4]Hhel
3
notations, we have log(T) = (2(1-))8AA
SsnHaa
Note that the escape direction is aligned with some eigenvector in the diffusion theory (Xie et al. 2021b). We leave the proof in Appendix A.3 We note that log(T) = O 2BAL
has been obtained
for SGD (Xie al., 2021b). We see that Momentum does not affect flat minima selection in terms of 1-8h
the mean escape time, if we properly choose the learning rate, i.e., 77Momentum
9g
7SGD.
ANALYSIS OF ADAM DYNAMICS
Algorithm 1: Adam
Algorithm 2: Adai
gr = vi(O:);
gt v(():);
me 31m4-1
+ (1-31)9:
04 = 32 + - 92)
Vt 328 1 + 82)9?
9t I-
me T-
T. = mean (ū:):
0t I-35
31.t t = (1 Oz) Clip
n1
8.+1 O₃ Voete m4s
m. 3
my = 1- 9
m/
0 = 0, m
In this section, we analyze the effects of Adaptive Learning Rate in terms of saddle-point escaping and flat minima selection.
A motivation behind Adaptive Learning Rate. The previous theoretical results naturally point us good way to help escape saddle points: adaptively adjust the learning rates for different parameter as 7) x - However, estimating the Hessian expensive in practice. In Adam (Algorithm 1). the diagonal U be regarded
diagonal approximation of the full matrix E Eig,9,"] 1 (Staib et al. 2019
Following Staib et al. 2019). our analysis considers the "idealized" Adam where v = Elg.9. is full matrix. Note that E.g.9. = C(#) = t approximately holds near critical points.
!!
The continuous-time Adam dynamics. The continuous-time dynamics of Adam can also be written Equation (6), except the mass M 1-4 and the damping coefficient y 1-8, We emphasize that the learning rate is not real number but the diagonal approximation of the ideal learning rate
5

---

Under review as a conference paper at ICLR 2022
matrix f) -nC in practice. We apply the adaptive time continuation dt; f) in the i-th dimension, where t2 of iterations is defined as the sum of f): of each iteration. Thus, the Fokker-Plane Equation for Adam can still be written Equation (8) We leave more details in Appendix H
Adam escapes saddle points. Similarly to Theorem we formulate how Adam escapes saddle point as Proposition 1 which can be obtained from Theorem 2 and ij = =
Proposition (Adam Escapes Saddle Points), Suppose e is a critical point, Assumption | and Equation ( hold, the dynamics is governed by Adam, and the initial parameter is at the saddle point C. Then the mean squared displucement is wrilten as
- -1pp(((--)))) +
D; 72M
D: Di_11 -xp(expp-- 7MH,
2H,t
(E.)
exp(- )].
yM
Under |HilnjT < 1, it can be written in terms of Adam notations as
(A0?) = 2(1 - 31) ex- (-((1 3()T)) + OVV
112
From Proposition] ! we can see that Adam escapes saddle points fast, because both the momentum drift and the diffusion effect are approximately gressian-independeee and isotropic near saddle points, which is also supported by (Saiib et al. (2019). Proposition indicates that one advantage of Adam over RMSprop (Hinton et al. 2012) comes from the additional momentum drift effect on saddle-point escaping.
Adam escapes minima. We next formulate how Adam escapes minima as Proposition 2 which shows that Adam cannot learn flat minima as well as SGD. We leave the proof in Appendix A.4 Proposition 2 (Adam Escapes Minima). The mean escape time of Adam only exponentially depends on the square root of the eigenvalues of the Hessian al a minimum:
4n B Hbel
[ Hb)]
2 BAL
+ (1 8)
(1-8)
) v |Hbb
=
+1
exp
1 - B1
Hobel
VHHae
where T is the mean escape time. Thus, we have log(7) = =
(2/BAL) nVHaa
In comparison, Adam has log(T) = O(Hee)), ), while SGD and Momentum both have log(7) = O.H=)). We say that Adam weaker Ieessian-dependent than SGD in terms of sharp minima escaping, which means that if Hue increases, i.e., the escaping direction becomes steeper, the escaping time log(r) for SGD decreases more dramatically than (T) for Adam. Note that, in the diffusion theory, the minima sharpness is reflected by Hac. the eigenvalue of the Hessian corresponding to the escape direction along an eigenvector. The weaker Hessian -deeedede diffusion term of Adam hurts the efficiency of selecting flat minima in comparison with SGD. In summary, Adam is better at saddle-point escaping but worse at flat minima selection than SGD.
Adam variants. Many variants of Adam have been proposed to improve performance, such as AdamW (Loshchilov and Hutter 2018), AdaBound (Loo et al. 2019). Padam (Chen and Gu 2018). RAdam (Liu et al. (2019), AMSGrad (Reddi et al., 2019)) Yogi (Zaheer al. 2018). and others (Shi et al. 2021 Défossez et al.) 2020 Zou et al. 2019). Many of them introduced extra hyperparameters that require tuning effort. Most variants often generalize better than Adam, while they may still not generalize as well as fine-tuned SGD (Zhang et af.] 2019b), which we also discuss in Appendix Moreover, they do not have theoretical understanding of minima selection. Loosely speaking, our analysis suggests that Adam variants may also be poor at selecting flat minima due to Adaptive Learning Rate.
ADAPTIVE INERTIA
In this section, we propose novel adaptive inertia optimization framework. Adaptive Inertia Optimization (Adai). The basic idea of Adai (Algorithmp) 2 comes from Theorems 2 and from which can that parameter-wise adaptive inertia achieve the approximately

---

Under review as a conference paper at ICLR 2022
Hessian independent momentum drift without damaging flat minima selection. The total momentum drift effect during passing through a saddle point in Theorem is given by
12i
(A0,)² 2(1 -B)
(10)
We generalize the scalar 81 in Momentum to a vector 31 in Adai. The ideal adaptive updating e v, where the rescaling factor namely the mean of all elements in the estimated 0, could be used to make the optimizer more robust in practice. The default values of 8o and 82 are 0.1 and 0.99, respectively. The hyperparameter is for avoiding extremely large inertia. The default setting = 0.001 means the maximal inertia is 1000 times the minimal inertia, as Mmax = 711 maxx)
is to keep 31 1 80
r7E Compared with Adam, Adai does not increase the number of hyperparameters. Note that an existing "adaptive momentum" method (Wang and Ye (020) is not parameter-wisely adaptive and essentially different from Adai.
The Fokker-Planck Equation associated with Adai dynamics in the phase space can be written as Equation (8) except that we replace the mass coefficient and the damping coefficient by the mass matrix and the damping matrix: M = 7/( diag((1)) and 7 - (iag(
Adai escapes saddle points. Proposition 3 shows that Adai can escape saddle points efficiently due momentum drift, while the diffusion effect is
to an isotropic and approximately bressian-independee
the same as Momentum. Proposition iis a direct result of Theorem 2
Proposition 3 (Adai Escapes Saddle Points). Suppose e is a critical point, Assumption Z and Equation 3 hold, the dynamics is governed by Adai, and the initial parameter the saddle point c. Then the total momentum drift in the procedure of passing through the saddle point is given by
-- 0n)²
11=1 HHalm
(20,)²
Bu
BunB
where Ei H.] is the trace norm of the Hessian at c.
Intuitively, the momentum drift effect of Adai can be significantly larger than Adam by allowing larger inertia, which can be verified in experiments. We do not rigorously claim that Adai must converge faster than fast as Adam. Instead, we will empirically compare Adai with Adam on various datasets.
Adai escapes minima. We formulate how Adai escapes minima in Proposition 4 which shows Adai can learn flat minima better than Adam. We leave the proof in Appendix A.5 Proposition 4 (Adai Escapes Minima). The mean escape time of Adai exponentially depends on the eigenvalues of the Hessian at a minimum:
4y >m MMall
1
2BAL
(1-8))
exp
3on
Hbel
7
Hac
Hbe
where T the mean escape time. Thus, we have log(7) = =0
(BBL nH.ec )
We can see that both SGD and Adai have log(r) = O(H while Adam has log(r) (H Thus, Adai favors flat minima as well as SGD and better than Adam.
Convergence Analysis. Theorem 4 proves that Adai has similar convergence guarantees SGD Momentum Yan et al. 2018 Ghadimi and Lan 2013). The proof is relegated to Appendix A. Theorem 4 (Convergence of Adai). Assume that I.(O) is an L-smooth function L is lower bounded as L()) 2 L", [[]]] = 0, [[]lg((, E) ( < 8², and [[)] < G for any 0, where represents the gradient noise of sub-sampling. Let Adai run t t+ 1 iterations and m1, 1-66 [0, for any 2 If S VIT we have
1
min E L[V)(()))11
VT+ +01+C2
+0++C+Ca))
k=
where C
[1- -i.moox y C2 m G2, and C3
I(0u) 1.-
CC
(G2 + 82).
2(1-Bt.maas)
2 Bmmma)
"I means that [VV(( (L(())] S L1)0, 0b]l holds for any 0, and 0s.
7

---

Under review as a conference paper at ICLR 2022
3.13
2.33
0200
Figure 1: The learning curves CIFAR-10 and SIFAR-100, Left: R8 on CIFAR-10. Right: ResNet34 CIFAR-100.
Figure 2: Adai shows better Generalization in the comparison with Momentum and Adam un- der similar convergence speed. DenseNet121 on CIFAR-100.
Adai
SGD M
Adam
Trkto
Top1
23.20
23.51
27.13
Top5
6.62
6.82
9.18
412
Figure 3: ResNet50 on ImageNet. Left Subfigure: Top Test Error. Right Subfigure: Training Loss.
Table: Top-1 and top-5 test errors. Note that the popular SGD bascline performance of ResNet50 on ImageNet has the test errors as 23.85% in Py PyTorch and 24.9% in He et al. (2016), which are both weaker than our SGD baseline.
We note that Adai is the base optimizer in the adaptive inertia framework, while Adam is the base optimizer in the adaptive gradient framework. We can also combine Adai with stable/decoupled weight decay (AdaiS/AdaiW) (Loshchilov and Hutter 2018 Xie et al.) (0200 (see more discussions in Appendix G or other techniques from adaptive gradient methods.
EMPIRICAL ANALYSIS
In this section, first conduct experiments compare Adai, Adam, and SGD with Momentum in terms of convergence speed and generalization, and then empirically analyze flat minima selection.
Datasets: CIFAR-10, CIFAR-1000
and Hinton 2009). ImageNet(DDnng et al. 2009). and (Marcus et al.) 1993). Models: SeaNNel1/ResNNe33//Ree
Penn TreeBank (Mareuu
(He et al. (2016).
VGG16 (Simonyan and Zisserman 2014), 1 (Huang et al., (2017). GoogLeNer (Szegedy et al. 2015), and Long Short- Term Memory (LSTM) (Hochreiter and Schmidhuber 1997b). More details and results can be found in Appendix and Appendix
MSN
Momestum
Adkm
5.985
5.985
5.985
10-2
5.655
5.655
5. 5.655
5.490
5 5.450
5. 5.490
Lg
10
5.150
3.150
3.180
4.955
4.955
4.995
10-4
46655
4 4.665
4.665
10
10
10
Decay
Weight Becay
Figure 4: The test emror of ResNet18 CIFAR-10 under various learning rates and weight decay. Adai has a much deeper and wider blue region near dark points (< 4.83%) than SGD with Momentum and Adam.
8

---

Under review as a conference paper at ICLR 2022
deel Farrart 5064
sdmm Feavord 3844
Jan Partoodd
8--
03
54
07
88
63 10
(a) Adai: - log(T) = O(k (b) Adam: - log((T) # O(k-t) (c) Adam: - log(T) = 0(k
Figure 5: Flat Minima Selection: Adai = Momentum >> Adam. The scale mean escape time - log(T) with the 95% confidence interval is displayed. We empirically verify that - log(T) = O(k holds for Momentum and Adai but does not hold for Adam. Instead, we observe that log(T) (k- holds better for Adam. While Adai and Momentum favor flat minima similarly, Adai may escape loss valleys slightly faster than Momentum.
Generalization and Convergence Speed. We display the test performance of all models on CIFAR in Table [ Table [ and Figure [ shows that Adai has excellent generalization compared with CIFAR- ITFAR-0//CIFAR-O0 Figure 2 further demonstrates that Adai can
other popular optimizers consistently generalize better than SGD with Momentum and Adam, while maintaining similarly fast or faster convergence, respectively.
Image Classification on ImageNet. Figure 3 shows that Adai generalizes significantly better than SGD and Adam by training ResNet50 on ImageNet.
Robustness to the hyperparameters. test performance, but also has wider test error basin than SGD with Momentum and Adam. It means that Adai is more robust to the choices of learning rates and weight decay.
gemuonsratttttteee that Adai not only has better optimal
The Mean Escape Time Analysis. We empirically study how the escape rate r, which equals the inverse mean escape time, depends on the minima sharpness for different optimizers in Figure 3 We Styblinski- Function as the test function which has clear boundaries between loss valleys. Our method for adjusting the minima sharpness is multiply rescaling factor vk each parameter, and
the Hessian will be proportionally rescaled by factor k. If we let L(0) s()) - L(O) I(VK0),
then (0) = (() + (()) = AV2 & Thus, we can use k to indicate the relative minima sharpness. We leave more details in Appendix
From Figure 5 it is easy to see that we have log(T) O(H)!)) in Adai and SGD, while we have log(r) = O (H,,.) in Adam. Moreover, Adam significantly less dependent on the minima sharpness than both Momentum and Adai.
Supplementary Experiments. (1) It is known that training of language models requires Adaptive Learning Rate (Zhang et al., 2019b). We will leave fusing Adaptive Learning Rate and Adaptive Inertia together as future work. which may be interesting for Transformer-ben based models (Vaswani al. 2017). We leave the results of training LSTM on Penn TreeBank in Figure 12 of Appendix F (2) In Appendix [F] we used the expected minima sharpness in Neyshabur et al. (017)) and empirically verified again that Adai and Momentum can learn flatter minima than Adam.
CONCLUSION
To the best of our knowledge, we are the first to theoretically disentangle the effects of Adaptive Learning Rate and Momentum terms of saddle-point escaping and flat minima selection. Under reasonable assumptions, our theory explains why Adam is good at escape saddle points but not good at selecting flat minima. We further propose a novel optimization framework, Adai, which can parameter-
seieeter-wiellyyyyee wisely adjust the momentum hyperparameter. Supported by good theoretical properties, Adai can accelerate training and favor flat minima well at the same time. Our empirical analysis demonstrates that Adai generalizes significantly better than popular Adam variants and SGD.

---

Under review as a conference paper at ICLR 2022
REFERENCES
An, G. (1996). The effects of adding noise during backpropagation training a generalization performance. Neural computation, 8(3):443-61) 6744.
Balakrishnan, V. (2008). Elements of nonequilibrium statistical mechanics, volume 3. Springer.
Berglund, N. (2013). Kramers' law: Validity, derivations and generalisations. Markov Processes and
Related Fields, 19(3):459-49.
Chen, and Gu, Q. (2 Closing the generalization gap adaptive gradient methods in training deep neural networks. arXiv preprint arXiv: 1806.06763.
Coffey, W. and Kalmykov, Y. (2012). The Langevin equation: with applications to stochastic problems in physics, chemistry and electrical engineering, volume 27. World Scientific.
Daneshmand, H., Kohler, J., Lucchi, A., and Hofmann, (2018). Escaping saddles with stochastic gradients. In International Conference on Machine Learning, pages 1155-1164.
Dauphin, N., Pascanu, R., Gulcehre, C., Cho, K., Ganguli, S., and Bengio, (2014). Identifying and attacking the saddle point problem in high-dimensional non-convex optimization. Advances in Neural Information Processing Systems, 27:2933-2941
Défosse7, A., Bottou, I.., Bach, F., and Usunier, N. (2020). On the convergence of adam and adagrad. arXiv preprint arXiv:2003.02333
Deng, J., Dong. W. Socher, R., Li, L.-J., Li, K., and Fei-Fei, L. (2009). Imagenet: A large- hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248-255. lece.
Eyring. H. (1935). The activated complex in chemical reactions. The Journal of Chemical Physics, 3(2):107-115.
Ghadimi, S. and Lan, G. (2013). Stochastic first-and zeroth-order methods for nonconvex stochastic programming. SIAM Journal Optimization, 23(4):2341-255 2368.
Hanggi, (1986). Escape from a metastable state. Journal of Statistical Physics, 1255--):
Hänggi, P., Talkner, P., and Borkovec, M. (1990). Reaction-rate theory: fifty years after kramers. Reviews of modern physics, 62(2):251.
HaoChen, J. Z., Wei. C., Lee, J., and Ma, T. (2021). Shape matters: Understanding the implicit bias of the noise covariance. In Conference on Learning Theory, pages 2315-2357. PMLR.
Hardt, M., Recht, B., and Singer, Y. (2016). Train faster, generalize better: Stability of stochastic gradient descent. In International Conference on Machine Learning, pages 1225-1234.
He, K., Zhang, X., Ren, S., and Sun, J. (2016). Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778.
Hinton, G.. Srivastava, N., and Swersky, K. (2012). Neural networks for machine learning lecture 6a
overview of mini-batch gradient descent.
Hochreiter, S. and Schmidhuber, J. (1995). Simplifying neural nets by discovering flat minima. In Advances in neural information processing systems, pages 529-536.
Hochreiter, S. and Schmidhuber, J. (1997a). Flat minima. Neural Computation, 9(1):1-42.
Hochreiter, S. and Schmidhuber. (1997b). Long short-term memory. Neural computation, 98): 1780.
Hu, W., Li, J., Li, L., and Liu, J.-G. (2019). On the diffusion approximation of nonconvex stochastic gradient descent. Annaly of Mathematical Sciences and Applications, 4(1):3-32.
10

---

Under review as a conference paper at ICLR 2022
Huang, G., Liu, Z., Van Der Maaten, L.,, and Weinberger, K. Q. (2017). Densely connected con- volutional networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 4700-4708.
Jastrzkebski, S., Kenton, Z., Arpit, D., Ballas, N., Fischer, A., Bengio, Y., and Storkey, A. (2017). Three factors influencing minima sgd. arXiv preprint arXiv:1711.04622
Jin, C., Ge, R., Netrapalli, P., Kakade, S. M., and Jordan, M. (2017). How to escape saddle points efficcietly. In International Conference on Machine Learning, pages 1724-1732. PMLR.
Kalinay, P. and Percus, J. K. (2012). Phase space reduction of the one-dimensional fokker- fokker-planck (kramers) equation. Journal of Statistical Physics, 148(6):1135-1155
Kingma, D. P. and Ba, J. (2015). Adam: A method for stochastic optimization. 3rd International Conference Learning Representations, ICLR 2015.
Kramers, H. A. (1940). Brownian motion in a field of force and the diffusion model of chemical reactions. Physica, 7(4):284-304.
Krizhevsky, A. and Hinton, G. (2009). Learning multiple layers of features from tiny images.
LeCun, Y. (1998). The mnist database of handwritten digits. http://yann. lecun. com/exdl/mmistt/
Li. Q. Tai, C., al. (2017). Stochastic modified equations and adaptive stochastic gradient algorithms. In Proceedings of the 34th International Conference on Machine Learning- learning-Volume 70, pages 2101- 21 2110. JMLR. org.
Li, Q.. Tai, C., and Weinan, E. (2019). Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations. J. Mach. Learn. Res., 20:40-1.
Li, Z., Malladi, S., and Arora, S. (2021). On the validity of modeling sgd with stochastic differential equations (sdes). arXiv preprint 1rXiv:2102.1244:
Liu. L., Jiang, H., He. P., Chen, W., Liu, X., Gao, J., and Han, (2019). On the variance of the adaptive learning rate and beyond. In International Conference on Learning Representations.
Loshchilov. I. and Hutter, F. (2018). Decoupled weight decay regularization. In International Conference on Learning Representations.
Luo, L.. Xiong, Y., Liu, Y., and Sun, X. (2019). Adaptive gradient methods with dynamic bound of learning rate. 7th International Conference on Learning Representations, ICLR 2019.
Mandt, S., Hoffman, M. D., and Blei, D. M. (2017). Stochastic gradient descent as approximate bayesian inference. The Journal of Machine Learning Research, 18(1):487344907.
Marcus, M., Santorini, B., and Marcinkiewicz, M. A. (1993). Building a large annotated corpus of english: The penn treebank.
Neelakantan, A., Vilnis, L., Le, Q. V., Sutskever, I., Kaiser, L., Kurach, K., and Martens, J. (2015). Adding gradient noise improves learning for very deep networks. arXiv preprint arXiv:1511.06680
Neyshabur, B., Bhojanapalli, S., McAllester, D., and Srebro, N. (2017). Exploring generalization in deep learning. In Advances in Neural Information Processing Systems, pages 5947-5956.
Nguyen, T. H., Simsekli, U., Gurbuzbalaban, M., and Richard, G. (2019). First exit time analysis of stochastic gradient descent under heavy-tailed gradient noise. In Advances in Neural Information Processing Systems, pages 273-283.
Pawitan, Y. (2001). In all likelihood: statistical modelling and inference using likelihood. Oxford University Press.
Radpay, P. (2020). Langevin equation and fokker-planck equation.
Reddi, S., Zaheer, M.. Sra, S., Poczos, B., Bach, F., Salakhutdinov, R., and Smola, A. (2018). A generic approach for escaping saddle points. International Conference on Artificial Intelligence and Statistics, pages 1233-1242. PMLR.

---

Under review as a conference paper at ICLR 2022
Reddi, S. J., Kale, S., and Kumar, S. (2019). On the convergence of adam and beyond. 6th International Conference on Learning Representations, ICLR 2018.
Risken, H. (1996). Fokker-planck equation. In The Fokker-Planck Equation, pages 63-95. Springer.
Risken, H. and Eberly, J. (1985). The fokker-planck equation, methods of solution and applications. Journal of the Optical Society of America B Optical Physics, 2(3):508.
Sato, I. and Nakagawa, H. (2014). Approximation analysis of stochastic gradient langevin dynamics by using fokker-planck equation and ito process. In International Conference on Machine Learning, pages 982-990.
Shi, N., Li, D., Hong, M., and Sun, (2021). (RMS can converge with proper hyper-parameter. In International Conference on Learning Representations.
Simonyan, K. and Zisserman, A. (2014). Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409.
Simsekli, U., Sagun, L.,, and Gurbuzbalaban, M. (2019). A tail-index analysis of stochastic gradient noise in deep neural networks. In International Conference on Machine Learning, pages 5827- 5837.
Staib, M., Reddi, S., Kale, S., Kumar, S., and Sra, S. (2019). Escaping saddle points with adaptive gradient methods. In International Conference on Machine Learning, pages 5956-5965. PMLR.
Szegedy, C., Liu, W., Jia, Y., Sermanet, P., Reed, S., Anguelov, D., Erhan, D., Vanhoucke, V., and Rabinovich, A. (2015). Going deeper with convolutions. Proceedings of the IEEE conference computer vision and pattern recognition, pages 1-9.
Van Kampen, N. G. (1992). Stochastic processes in physics and chemistry, volume 1. Elsevier.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. (2017). Attention is all you need. In Advances in neural information processing systems, pages
Wang, B. and Ye, Q. (2020). Stochastic gradient descent with nonlinear conjugate gradient-style adaptive momentum. arXiv preprini arXiv:2012.02188.
Wang, J.-K., Lin, C.-H., and Abernethy, J. (2019). Escaping saddle points faster with stochastic momentum. In International Conference on Learning Representations.
Welling, M. and Teh, Y. W. (2011). Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11), pages 68 681-688.
Wilson, A. C., Roelofs, R., Stern, M., Srebro, N., and Recht, B. (2017). The marginal value of adaptive gradient methods in machine learning. In Advances Neural Information Processing Systems, pages 4148-411 44584158.
Wu, L., Ma. C., and Weinan, E. (2018). How sgd selects the global minima in over wer-parameterized learning: A dynamical stability perspective. In Advances Neural Information Processing Systems, pages 8279-88788 8288.
Xie, Z., He, F., Fu, S., Sato, I., Tao, D., and Sugiyama, M. (2021a). Artificial neural variability for deep learning: On overfitting, noise memorization, and catastrophic forgetting. Neural Computation.
Xie, Z., Sato, I., and Sugiyama, M. (2020). Stable weight decay regularization. arXiv preprint arXiv:201..11152.
Xie. Z., Sato, I., and Sugiyama, M. (2021b). A diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima. In International Conference on Learning Representattons.
Xie, Z., Yuan, L.,, Zhu, Z., and Sugiyama, M. (2021c). Resitive-negative momentum: Manipulating stochastic gradient noise to improve generalization. In International Conference on Machine Learning, volume 139 of Proceedings Machine Learning Research, pages 11448-1 11458. PMLR.
12

---

Under review as a conference paper at ICLR 2022
Xu, P.. Chen, J., Zou, D., and Gu, Q. (2018). Global convergence of langevin dynamics based algorithms for nonconvex optimization. In Advances in Neural Information Processing Systems, pages 3122-3133.
Yan, Y., Yang, T., Li. Z., Lin, Q., and Yang. Y. (2018). A unified analysis of stochastic momentum methods for deep learning. In IJCAI International Joint Conference on Artificial Intelligence.
Zaheer, M., Reddi, S., Sachan, D., Kale, S., and Kumar, S. (2018). Adaptive methods for nonconvex optimization. In Advances in neural information processing systems, pages 9793-9803.
Zaremba, W., Sutskever, I., and Vinyals, O. (2014). Recurrent neural network regularization, arXiv preprint arXiv:
Zavriev, S. and Kostyuk, F. (1993). Heavy-ball method in nonconvex optimization problems. Compu- tational Mathematics and Modeling, 4(4):336-34) 341.
Zhang, C., Bengio, S., Hardt, M., Recht, B., and Vinyals, O. (2017). Understanding deep learning requires rethinking generalization. In International Conference Machine Learning.
Zhang. G., Li, L., Nado, Z.., Martens, J., Sachdeva, S., Dahl, G., Shallue, C., and Grosse, R. B. (2019a). Which algorithmic choices matter at which batch sizes? insights from a noisy quadratic model. In Advances Neural Information Processing Systems, pages 8 8196-8207.
Zhang. J., He, T., Sra, S., and Jadbabaic, A. (2019b). Why gradient clipping accelerates training: A theoretical justification for adaptivity. In International Conference Learning Representations.
Zhou, H.-X. (2010). Rate theories for biologists. Quarterly reviews of biophysics, 43(():2119229
Zhou, M., Liu, T., Li, Y., Lin, D., Zhou, E., and Zhao, T. (2019). Toward understanding the importance of noise in training neural networks. In International Conference Machine Learning.
Zhou, P., Feng, Ma, C.,, Xiong, C., Hoi, S. C. H., et al. (2020). Towards theoretically understanding why sgd generalizes better than adam in deep learning. Advances in Neural Information P Processing Systems, 33.
Zhu, Z., Wu, J., Yu, B., Wu, L.,, and Ma, J. (2019). The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects. In ICML., pages 7654-7663
Zou, F. Shen, I.., Jie, Z., Zhang, W., and Liu, W. (2019). A sufficient condition for convergences of adam and rmsprop. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11127-11135.
A PROOFS
A. PROOF OF THEBREM
1
Proof. is easy validate that the probability density function
P(0,t) ((2))" (Ett) exp (-) (00--cTT((1
T(()()()-)))) (-)
(11)
is the solution of the Fokker-Planck Equation one-dimensional solution, such as Dimension
Without losing generality, we only validate
The first term in Equation 2 can be written as
0P(0,t) Ø1
1
1
(- 02
0
00²
2
2ma2
exp
20
01.
+
v 2m
exp
2
2
al.
(12)
0
1
P(0,t)
(13)
gt
Øt
13

---

Under review as a conference paper at ICLR 2022
The second term in Equation 2 can be written
0²
(
0
(P(O,))VLIO) P(O)H + 110
exp
(14)
V2n0²
2
o²
8
H P(0,t).
(15)
The third term in Equation
can be written
0²
8
A
P(O,)( D
exp
(16)
σb
2T
20
02
P(O,t).
(17)
0²
By Terml= Term2 + Term3, we have
Og2
(02
0²)
Ho²
2 0²) + D(02 02)
(18)
0t
Oa²
= 2D 2Ho
(19)
01
The initial condition of 2 is given by o0) 0. We validate that 0² satisfies
Dj [1 exp(- (H,)) H;
(t)
(20)
true along eigenvectors' directions.
By D = 2 H, we can get the results of SGD diffusion:
of (T) sign(H.:) 201 exp(-211,711)
(21)
A.2 PROOF OF THBOREM 2
Proof. Without losing the generality, we consider the one-dimensional case that aligns with an cigenvector the Hessian. The stationary solution in equilibrium to the phase-space Fokker-Planck Equation near critical point is given by canonical ensemble
(0, r) = Z exp(- (0,
Mr3)],
(22)
where Z is partition function and the inverse temperature MD-1 This result is famous in large number of physics literature (Van Kampen] 1992 Risken 1996 Balakrishnan 2008). Thus, we have the equilibrium distribution in velocity space and position space
rM
Poς(r) = 7, exp(
(23)
2
and
(0) = 2 exp(-3L(0)),
(24)
respectively, where Z, is the partition function for (r)(r) and Zo is the partition function for Pa..
known that the phase-space probability density solution P(0, r, t) can be obtained from the position-space solution P(O,t) t) and the selocity-space solution P(r, t). Thus, may compute the velocity-space solution (in Step 1) and the position-space solution (in Step II), respectively.
Step I:
14

---

Under review as a conference paper at ICLR 2022
We further argue that the velocity- space solution P(r.t) has reached its equilibrium Boltzmann distribution (r) when the particle passing saddle points. This is reasonable because the equilibrium solution in the velocity space is stable under Assumption [ The velocity distribution t(r,t) must be the solution of the free diffusion equation in the velocity space. We can approximately
ignore the gradient expectation term #L(9) in dynamical equations near critical points, the gradient expectation is much smaller the gradient noise scale near critical points. Near critical points, the velocity r obeys the equilibrium distribution
P(r, 2 (r) (2T) det(BM)3 exp( BMr²
(25)
2
where 1 in one-dimensional case. This is classical Boltzmann distribution. The expected velocity (also called "equilibrium velocity") along dimension can be given by reg. 7.
D:
Step II:
As we discussed Theorem | the position-space distribution equilibrium time-dependent near saddle points. Following the form of the solution to the position-space Fokker-Plann Planck Equation (Seen in Appendix I) the ansatz solution of P(0, t) is given by
P(O,t) t) ((2m)" (e)( exp ( dt)) (()) c(t)))
Σ() U diag(off() (.., (1),
The time-dependent components in P(O,t) t) are caused by two effects. The first effect is the momentum drift effect (due to equilibrium velocity), which decides c(t), the center position of the probability density. As studied in Theorem [ the second effect is the diffusion effect (due random noise), which decides covariance of the probability density, E(t).
The momentum drift effect is governed by the motion equation without random noise, while the diffusion effect is governed by the diffusion equation where gradient noise dominates gradient mean.
Step II (a): the momentum drift effect.
We can write the dynamics of the momentum drift effect as
Më(() = -)M((()).
(26)
We also have ignored the conservative force [c[((t) <(0)] near saddle points, |H,| exists for ill-condition Hessian eigenvalues. We focus on the behaviors near ill-conditioned saddle points, where Hessian eigenvalues small along the escape directions.
MM
The initial condition is given by c(0) = c. c(0) = Teg. and ē(0) = -TYeg- Then we can obtain the solution c(t) as
(t = Gi + T2,0g [1 - (-xxpp-xtt))
(27)
Step II (b): the diffusion effect.
We get the dynamics of the diffusion effect as
aT(()) di + [2D
Md = 30
(28)
This is equivalent SGD dynamics with i = TM The expression of P(O.t) and (f directly given by Theorem Tas
n
2H, of(() = Dj [1-cxp(
7M )
(29)
>HH
We combine the momentum drift effect and the diffusion effect together, and then obtain the mean squared displacement of 0
D; [1 exp(->t))2
Dj
2He exp(
(A))()) (t) c² =f(()= M²
)] (30)
7MH;
15

---

Under review as a conference paper at ICLR 2022
We respectively introduce the notations of Adym-style Momentum and SGD-style Momentum, and
apply the second order Taylor expansion in case of small
2/1
Then we obtain
7M
1 2(1 B
H. 2T T
(A+T)
[1 cxp -
3) +
H((
(31)
B
for Adam-style Momentum, and
|
(A 2(1 - B1) [1 exp (-(1-
B(1 3,)3 00
(32)
for SGD-style Momentum.
A.3 PROOF OF THEOREM
3
Proof. The proof closely relates to the proof of Theorem 3.2 in Xie et al. (2021b) and a physics work (Kalinay and Percus 2012).
We first discover how SGD dynamics differs from Momentum dynamics in terms of escaping loss valleys. In this approach, may transform the proof for Theorem 3.2 Xie et al. (2021b) into the proof for Theorem with the effective diffusion correction. We use J and j to denote the probability current and the probability flux respectively. According to Gauss Divergence Theorem, we may rewrite the Likker-Plancle Equation (8) as
@P(,,r,t)
VoP VoL()) M (V,((())
01
V, M DD()) Peg (r(
(33)
V ((0,t).
(34)
will take similar forms of the proof in Xie et al. 2021b as follows. The mean escape time written
P(O E Va) Ss. J ds'
(35)
To compute the mean escape time, we decompose the proof into steps: compute the probability of locating in valley P(O Va), and 2) compute the probability flux j = fs. J The definition of the probability flux integral may refer to Gauss Divergence Theorem.
Step
Fortunately, the stationary probability distribution inside Valley in Momentum dynamics is also given by a Gaussian distribution in Theorem as
D.
(t 7MH,
(36)
Under Quasi- Equuii-Equil Assumption, the distribution around minimum a is P(0)
P((a exp M (0 (D= H. a)} We the T notation the temperature param- eter in the stationary distribution, and use the D notation as the diffusion coefficient in the dynamics, for their different roles.
(37)
-
=
Joev.
P(O)dV
(38)
YM (0 a) (D, H. $ (6 a) dv
-P(a)
exp
(39)
JBVV
YM (0 a) (D. T) H.D. 10 a) dV
PP(a)
exp
(40)
(BE(-00+t00)
(2T7:M)
=P(a)
(41)
dot(Dr, Ha)
16

---

Under review as a conference paper at ICLR 2022
This result of P(O Va) in Momentum only differs from SGD by the temperature correction ydM.
Step 2: We directly introduce the effective diffusion result from (Kalinay and Percus 2012) into our analysis. KKalinay and Percus (2012) proved that the phase-space Fokker-Planck Equation (8) can be reduced to a space-dependent Smoluchowski-lik equation, which is extended by an effective diffusion correction:
D. = D (1
4H,(0) M
2H. M
(42)
As only employ the Smoluchowski Equation along the escape direction, use the one- dimensional expression along the escape direction (an eigenvector direction) for simplicity. Without losing clarity, we use the commonly used to denote the temperature in the proof.
In case of SGD (ii et al. 2021b we can obtain Smoluchowski Equation in position space:
L(0)
L(0)
P())
D(0) exp
exp
(43)
where T D. According to (Kalinay and Percus 2012), in case of finite inertia, we transform the phase-space equation into the position-space Smoluchowskiilike form with the effective diffusion correction:
L(0)
L(0)
J D(O) exp
exp
P(0)
(44)
where T = D 7M
and D defined by Equation 42 replaces standard D.
We assume the point is midpoint on the most possible path between and b, where L(s) (1- L(a) + (b). The temperature TĄ dominates the path a + s, while temperature Tj dominates the path b. So we have
L(0) I(s)
P(B) = JD- exp
[[(0) L(s)
V
exp
(45)
T
T
We integrate the equation from Valley a to the outside of Valley a along the most possible escape path
L(0) L(s)
Left
lexp
P(O)Jd0
(46)
00
T
0
L(0) L(s)
exp
P(O) d0
(47)
d0
To
((0)=
L(s))
+
exp
P(O) d0
(48)
00
Tb
[(a)= L(s)
=(P(s) exp
P(a)) + 10 P())
(49)
To
[(aa) L(s)
exp
P(a)
(50)
Ta
Right = = 1 D- exp
((0)) L(s)
DD-exp
d8
(51)
T
We move J to the outside of integral based on Gauss's Divergence Theorem, because / is fixed on the escape path from one minimum to another. As there is no field source on the escape path, fy V J(O)dV 0 and (()) 0. Obviously, probability sources are all near minima in deep learning. So we obtain
exp
((a)-L(3) T.
P(a)
(52)
exp
((0) L(s)
d8
17

---

Under review as a conference paper at ICLR 2022
Near saddle points, we have
[(0) L(s)
D exp
d0
(53)
T
L(b) - L(s) + + - b)
Ho(0 - b)
D exp
d0
(54)
Tb
400
L(b) L(s) + 50 b) Ho - b)
=D
exp
d0
(55)
Tb
L(bb I(s)
TT HH
= exp
(56)
Tb
Besides the temperature correction the effective diffusion correction
T = 7M this result of J in Momentum also differs from SGD by
D. The effective diffusion correction coefficient is given by
(
4H,(0) 7"M
(57)
D.
(H,
7-
Based on the formula of the one-dimensional probability current and flux, we obtain the high- dimensional flux escaping through b:
f
dS
(58)
S
7M (0 b) DD,
= -Jid
cxp
H,D, ( b) dS
(59)
Sb
2
(2 n=1
=Jid (Ilipte (Dzr Hb:))
(60)
[aa)-L(s)) P(a) (e (Pud Hoa))
(2 "
exp
Toe
(61)
D
L(b) (((b)-((s)
2 Haal
exp
The
where [i]e indicates the directions perpendicular to the escape direction
Based on the results of Step 1 and Step 2, we have
P(O Va)
T Ss, J J.d
(62)
Dre exp
[(b))
2x Hoal
M2)
Toe
=P(a) dot( Di. Ha)
(63)
((a))(()) P(a)
(2 "
exp Tac
(DTe Ho.))
1 2nT
Doe
2yMBAL
(1-s)
Doe
|Hell
exp
Has
|Hbel
(64)
4 Hel
(
Hae + (1 8) )
=
exp
(65)
22
Hbel
Hoel
We have replaced the eigenvalue of H, along the escape direction by its absolute value. Finally, by introducing and M, we obtain the log-scale expression
log(T) 0 ( 2((1 B1)BAL
(66)
18

---

Under review as a conference paper at ICLR 2022
A.4 PROOF OF PROPOSITION 2
Proof. The proof closely relates to the proof of We only need replace the standard learning rate by the adaptive learning rate = 170 2, and set yM = 1. Particularly,
D Dodam = De nH
H
2B
2
1
B
(67)
We introduce f -pp into the proof of 3 and obtain
r P(O E Va)
(68)
Js. J.
[det(D, Hb)l
4 Hel
1
MYMBAL
(1 8)
exp
det(Da 1iD=-VV Vo² Ho)
72
HHbel
Vac Hac
V 2 Hbol
(69)
4n Hhe
det det HoOl
2 BAL
(1 (1-s)
exp
(70)
31
|Hbll
7
VHae
V|Hbel
A.5 PROOF OF PROPOSITION 4
Proof. The proof closely relates to the proof of We only need introduce the mass matrix M and the dampening matrix 7. Fortunately, 7M I. Thus we directly have the result from the proof of
as:
P(O V.)
(71)
Ss. J.d
= V 22M
Hhel
exp
2yMBAL
8
+
(1 8)
(72)
|Hell
Hon
Hbel
As M 711 B1) -1 1.7 = 11 7-'(1 B1). and 1- B1
de we obtain the result:
4n Σi HHb
2BAL
3
+ (1-s) 3)
T=T
+1
exp
(73)
Bon
Hon
7)
Har
HHe
A.6 PROOF OF THEOREM 4
Without loss generality, we assume the dimensionality is one and rewrite the main updating rule of Adai as
0.+1 0g - 101-n(1- 31,c)g. + BI, 0.-1).
(74)
In the convergence proof, we do not need specify how to update BI, but just let 31, E (0, 1). We denote that 0_1 0g
Before presenting the main proof, we first prove four useful lemmas.
Lemma 1. Under the conditions of Theorem for any t 2 O, we have
T1+1 Tt.
19

---

Under review as a conference paper at ICLR 2022
where
Th,t Ix = (1- -kk
II BIg
i=++1
Then we have
3
1.max
Dk 1.
k=0
Proof. Recall that
0.+1 =0; n111 -3,1)0t (1, 02-1)
(75)
Oct1 0g (0, 0t-1) 17(1 - 31,t)9r-
(76)
Then we have
0011 - 0: 11 Σ (1 31,k)gR
II 81. 1=k+1
(77)
k=
Let Iz.1 (1 Bx) IT= $1,5- For analyzing the maximum and the minimum, we calculate the derivatives with respect to 31,x for any 0 < t:
i E Uhit
II k=++1
B1. S.
(78)
0B.0
Note that 0 S B1, S 31max. Then we have
qhel qeela max
Th
S
(79)
k=
Om
k=
Recursively, we can calculate the derivatives with respect to $1 31,2
B1.
Then obtain max tt (k,t) 1 by letting B1, for all k.
Similarly, obtain min (Ei= (k,e) by letting 31,t Mm, for all k.
Then we obtain
1 1.max
Σ
T S 1
(80)
km
The proof is now complete.
Lemma 2. Under the conditions of Theorem 4 for any t t2 0, we have
E[L(O(42) -(.))) <
Is
(L(0)+2)
L(0,)
j=k
t k
Ln²
11 11²
Σ
< 12
2
k=
k=0
20

---

Under review as a conference paper at ICLR 2022
Proof. As L(0) is C-smooth, we have
I(((+1) L(0s)
(81)
SVL(0,) (01+1 - -6,)+ _ |10c+11 1111
L
(82)
Σ Thtt VL((,) 9k +
E
1k
(83)
2
k=
k=0
(LL(6s) + Ex)
1)
qeet
k=0
j=k
LT²
g ((())) ((V()) + +x) +
9k,19k11?
(84)
k=0
k=0
Taking expectation on both sides gives
B[L(0++1) [()))]
(85)
(86)
2C
k=0
j=k
1 9k,19111 112
(87)
By the Cauchy-Scheee Schhar- Inequality, we have
B|L(O(+1) L(O.)]
(88)
- (L.0))))V((M0)P L((x))1
LT2
B1I
Th. 1k.9kk ||
(89)
2
k=
g
VL(0)+1) - ((
2
k=
j=k
k
11
- -
D 11?
(90)
2
2 k=0
The proof is now complete.
Lemma 3. Under the conditions of Theorem f for any t 2 O, we have
31 max
Ct k) S 1 Bi,max
0=0
Proof. For analyzing the maximum, we calculate the derivatives with respect to B1,8 for any 0 <
k l.
With respect to S1.1, we have
i Et Tt.e - k)
II
i < O.
(91)
08
k=1
Then we let 31,0 = and obtain the derivative with respect to 31,1 as
Ek queett - k)
BI. (t (t- 1)) 2 0.
(92)
031,1
21

---

Under review as a conference paper at ICLR 2022
Then we let Bi,5 max for 1 j $ 1 and obtain the derivative with respect to 91,6
i σE 0 gk, k)
2 (t k) 2 2k
R= gket
O.
(93)
Thus, we have 31,0 and $1,4 ma, for j S to maximize Ek=o qr - 1s). We may write it as
Σ
gt, - k) S
g
1,maxx
+
Σ (1 31 maxx f.max
(94)
k=
k=0
Note that (1 - mmax) ) maxk is an ceithmetico-geomeetn sequence. Thus, we have
87 P.mmax (1 BT,
max
qk.elt - k)) S +
|1.max - 31,ttax t +
(95)
1 - 31,max
k=
Ff,aa
B1 max
(96)
m1,max
11.max
(97)
81,max
The proof is now complete.
Lemma 4. Under the conditions of Theorem 4 for any t 2 O, we have
CT²
LaF
Tr E
IV((0++) L(0,)) +
g S
(G² +8²).
2
2(1 m1,
j=k
k=0
Proof. As L(O) is C-smooth, we have
[V - L(0,)112
(98)
$2110,+1 - 0311
(99)
=11 Σ 9..
(100)
==0
By Ei Ri $ 1 in Lemma | have
B[[] VIIVTT(((+))) - [(0)))13)
(101)
<Cnp (G. +
(102)
and
LT²
1k.99kl ||2
(103)
k=
2n
(G² + 82).
(104)
By Lemma we have
1
ak.
E
VL(0)+1) - 1(8))119)
(105)
2 k=0
j=k
G22 +
qt.((t - k)
(106)
2
k=n
? 11max
(G2 +2).
(107)
2(1
3 31,max)
22

---

Under review as a conference paper at ICLR 2022
Then obtain
LT²
[V(00++1 - -(0)))
ch !|?
(108)
2
2
j=k
k=
C1j² 31 1,max
Laj
(G2 + 02)
GG +2)
(109)
2
f1,maxx
2
LT2
(G2 + (G2
(110)
2(1
mmxx
Proof. The proof of Theorem is organized as follows.
By Lemma
and Lemma
we have
B[L(0(+1) - L(Oe)]
2 k=0
j
Ca
Σ Yhee
11² - B(
(111)
2
2
k=0
k=
- -
LTp
ahael
(G² ( 52).
(112)
2
2(1 m,,
k=
Thus, we have
yk, (B
1 h "1
LT}
<B(L((.)) I((++1)]
Ohet 2
"
(G2 +2).
(113)
2(1 11,max)
k=
By Lemma aan Lemma
3 we have
(1 1.max ) min
E| VL
31 1)
La²
<E|L(O.) (()+ + 2(1 Bl,ma
2 +
(G2 +8))
(114)
2(1 m,,m
By summing the above inequality for 1 = O....,,, I, we have
(t + 1-(1111 - 31,max) )7 min
E[l V[(+))111
k k
Ca²
<L(00)
+ 2(1 B,,max)
G (t + 1)+ 1)+
(G2 + (t
(115)
2(1 31,max)
Then
min (Bx)L A ==0.... O....!
L((b)= - 1*
31 maxx
Ln
+
62
(G2 1 +2).
(116)
3l,max) (1 1)n
2(1 91, }
2(1 31,max) )
23

---

Under review as a conference paper at ICLR 2022
C Let 17 VF+I
We have
min [[:
$ T((0)) I,*
max C
max Cv++1 + 2(1 - max /21+1+ G2+
LC
2(1 m1, 12 (G+88 +82)
(117)
(118)
The proof is now complete.
B CLASSICAL APPROXIMATION ASSUMPTIONS
Assumption 2 iindicates that the dynamical system is in equilibrium near minima but not necessarily near saddle points. It means that @P(0.t) R 0 holds near minima, but not necessarily holds near saddle point b. Quasi-Equilibrim Assumption is actually weaker but more useful than the conven- tional stationary assumption for deep learning (Welling and Teh 2011: Mandt et al. 2017). Under Assumption 2 the probability density P can behave like a stationary distribution only inside valleys, but density transportation through saddle points can be busy. Quasi Buuuii--uiliiiium is more like: stable lakes (loss valleys) is connected by rapid Rivers (escape paths). In contrast, the stationary assumption requires strictly zero flux between lakes (loss valleys). Little knowledge about density motion be obtained under the stationary assumption.
Low Temperature Assumption is common (Van Kampen 1992, Zhou 2010 Berglund 2013 Jas-
trzkebski et al. (017)), and is always justified when 1 is small. Under Assumption the probability densities will concentrate around minima and MPPs. Numerically, the 6-sigma rule may often provide good approximation for a Gaussian distribution. Assumption will make the second order Taylor approximation, Assumption [. even more reasonable in SGD diffusion.
Here, we try to provide more intuitive explanation about Low Temperature Assumption in the domain of deep learning. Without loss of generality, we discuss temperature can be interpreted as real number D. In SGD, we have the temperature as D 2B H. In statistical physics, if AL is large, then we call Low Temperature Approximation. Note that AL
one-dimensional dynamics. The
D appears insides an exponential function in the theoretical analysis. People usually believe that, numerically, AL >> 6 can make good approximation, for a similar reason of the 6-sigma rule in statistics. In the final training phase of deep networks, common setting r) 0.01 and B 128. Thus, we may safely apply Assumption 3 to the loss valleys which satisfy the very mild condition AL 2.3 10-4 Empirically, the condition
2.3 10-4 holds well in SGD dynamics.
C STOCHASTIC GRADIENT NOISE ANALYSIS
In this section, we empirically discussed the covariance of stochastic gradient noise (SGN) and why SGN is approximately Gaussian in common settings.
In Figure 6. following Xie et al. 2b221b) we again validated the relation between gradient noise covariance and the Hessian. We particularly choose randomly initialized mode so that the model is not near critical points. We display all elements H())) E [le 4, 0.5] of the Hessian matrix and the corresponding elements C(t.g) of gradient covariance matrix the space spanned by the eigenvectors of Hessian. Even if the model is far from critical points, the SGN covariance is still approximately proportional to the Hessian and inverse to the batch size B. The correlation is especially high along the directions with small- maginnd cigenvalues of the Hessian. Fortunately, small- maal--mag cigenvalues the Hessian indicate the flat directions which we care most in saddle-point escaping.
We also note that the SGN we study is introduced by minibatch training, aL(0.) 09,
aL(0.), which is 50,
the difference between gradient descent and stochastic gradient descent.
In Figure 7 of the Appendix, Xie et al. (021b) empirically verified that SGN highly similar to Gaussian noise instead of heavy-tailed Levy noise. Xie al. (2021b) recovered the experiment of
24

---

Under review as a conference paper at ICLR 2022
0.5
Bris Co Narion C.kk
Bris Castaaal Narsos CaRd
0.4
0.4
Covar x Batch Size
0.3
0.2
0.2
0.0
0.0
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
Hessian
Hessian
0.5
Coato PeMEO E
0.5
2-24 Coslannss 8U677
64108 Out huae 2i0
0.3
Covar
0.2
0.1
0.0
0.0
0.2
0.4
Hessian
Hesslan
Figure 6: We verified Equation 3 by training pretrained and random three-layer fully-connected networks on MNIST (LeCun 1998 Top Row: Pretrained Models. Bottom Row: Random Models. The noise covariance is approximately proportional to the Hessian and inverse to the batch size B even not around critical points.
25

---

Under review as a conference paper at ICLR 2022
10
10
207
160
0.25
0.22
0.55
3.20
3.00
0.02
6.50
0.15
0.20
Asta 9urm
1a so Norm
(a) Gradient Noise of one minibatch across parameters
(b) Lévy noise
201
Centss
0.00
6.02
0.66
6.36
3.00
0.50
0.22
5.000
0.303
3.00
0.005
0.003
6.300
(c) Gradient Noise one parameter across minibatches
(d) Gaussian noise
Figure 7: The Stochastic Gradient Noise Analysis (Xie et al., 2021b). The histogram of the norm of the gradient noises computed with ResNet18 on CIFAR-10 10. Subfigure (a) follows Simsekli et al. (2019) and computes "stochastic gradient noise" across parameters. Subfigure (c) follows the usual definition and computes stochastic gradient noise across minibatches. Obviously, SGN computed over minibatches is more like Gaussian noise rather than Lévy noise.
Simsekli al. 2019) to show that gradient noise is upproximately Lévy noise only if it is computed across parameters. Figure the Appendix actually suggests that the contradicted observations are from the different formulations of gradient noise. [Simsekli al. (2019) computed "SGN" across model parameters and regarded "SGN" as n samples drawn from a single-variant distribution. In Xic al. (0211)) SGN computed over different minibatches obeys N --variant Gaussian distribution, which can be d-dependent and anisotropic. Simsekli et al. (2019) studied the distribution of SGN as Linglevvarint distribution, while Xie al. (2021b)) relaxed as n-vaiiant distribution. Figure holds well at least when the batch size B is larger than 16, which is common in practice.
D EXPERIMENTAL DETAILS
Computational environment. The experiments are conducted on a computing cluster with GPUs of NVIDIA @ Tesla P100 16GB and CPUs of Intel Xeon CPU E5-2640 v3 @ 2.60GHz.
D.1 IMAGE CLASSIFICATION
For CIFAR-IO/CIFAN 100, we perform the per-pixel zero-mean unit-variance normalization, horizontal random flip, and 32 x 32 random crops after padding with 4 pixels each side. For ImageNet, we perform the per-pixel zero-mean uaii-variance normalization, horizontal random flip, and the resized random crops where the random size (of 0.08 to 1.0) of the original size and a random aspect ratio (of 3 to 1 3 of the original aspect ratio made.
Data Preprocessing:
Hy Hlyeepaaraeeter Settings for CIFAR-14 R-10 and CIFAR-100: We select the optimal learning rate for each experiment from (0.00001, 0.0001, 0.001, 0.01, 0.1, 1, 10) for non-adaptive gradient methods and use the default learning rate in original papers for adaptive gradient methods. The settings of
26

---

Under review as a conference paper at ICLR 2022
(a) 1- I-Dimensional Escape
(b) High- Bigh--imeessioal Escape
Figure 8: The illustration of Kramers Escape Problems (Xie et al. 2021b). Assume there are two valleys, Sharp Valley ai and Flat Valley 02. Also Col b is the boundary between two valleys. aj and aa are minima of two neighboring valleys. b is the saddle point separating the two valleys. c locates outside of Valley 01
learning rates: 1/ for Adai:n 0.1 for SGD with Momentum and AdaiW; 17 0.001 for Adam, AMSGrad, AdamW, AdaBound, Yogi, and RAdam; schedule, the learning rate is divided by 10 the epoch of {80, 160} for CIFAR-IO and {100, 150) for CIFAR-100. The batch size is set to 128 for CIFAR- 10 and CIFAR-100, The L2 regularization hyperparameter set to 0.0005 for CIFAR-10 10 and CIFAR- 100. Considering the linear scaling rule of decoupled weight decay and initial learning rates (Loshchilov and Hutter (018). we chose decoupled weight decay hyperparameters as: 0.5 for Adam on CIFAR- 10 and CIFAR-100; 0.005 for AdaiW on CIFAR-10 and CIFAR-100. We set the momentum hyperparameter 31 = 0.9 for SGD with Momentum. As for other optimizer hyperparameters, we apply the default
0.01 for Padam. For the learning rate
hyperparameter settings directly.
Hyperparameter Settings for ImageNet: We select the optimal learning rate for each experiment from f0.00001, 0.0001, 0.001, 0.01, 0.1, 10) for Adai, SGD with Momentum, and Adam. The settings of learning rates: 7 = 1 for Adai, = 0.1 for SGD with Momentum; = 0.0001 for Adam. For the learning rate schedule, the learning rate is divided by 10 at the epoch of (30, 60,,60, 60, 90). The batch size is set to 256. The L22 regularization hyperparameter is set to X = 0.0001. We set the momentum hyperparameter 31 0.9 for SGD with Momentum. As for other training hyperparameters, we apply the default hyperparameter settings directly.
Some papers often choose 0.0001 the default weight decay setting for CIFAR-10 and CIFAR-100. We study the weight decay setting in Appendix
D.2 LANGUAGE MODELING
Moreover, we present the learning curves for language modeling experiments in Figure 12 We empirically compare three base optimizers, including Adai, SGD with Momentum, and Adam, for language modeling experiments. We use classical language model, Long Short-Term Memory (LSTM) (Hochreiter and Schmidhuber) 1997b) with 2 layers, 512 embedding dimensions, and 512 hidden dimensions, which has 14 million model parameters and similar to "medium LSTM" in Zaremba et al. (2014). The benchmark task is the word-level Penn TreeBank (Marcus al. 1993).
Hyperparameter Settings. Batch Size: B = 20. BPTT Size: bptt = 35. Weight Decay: 0.00005. Learning Rate: 0.001. The dropout probability is set to 0.5. We clipped gradient norm to We select optimal learning rates from {10, 1, 0.1, 0.01, 0.001, 0.0001, 0.00001} for each optimizer.
E THE MEAN ESCAPE TIME ANALYSIS
Mean Escape Time: The mean escape time is the expected time for a particle governed Equation
to escape from Sharp Valley a to Flat Valley 42, seen in Figure 3 The mean escape time is widely used in related statistical physics and stochastic process (Van Kampen 1992 Nguyen et al. 2019
27

---

Under review as a conference paper at ICLR 2022
Related machine learning papers (Xie et al., 2021b Hu et al., 2019 Nguyen al. 2019) also studied how SGD selects minima by using the concept of the mean escape time.
Data Set: We generate 50000 Gaussian samples as the training data set, where (0, 4/).
Hyperparameters: The batch size is set 10. No weight decay. The learning rate: 0.001 for Adai, 0.0001 for Momentum (with / 0.9), and 0.03 for Adam.
Test Function: Styblinski-Tang Function is a commonly used function nonconvex optimization, written
S(0) 00 2 160 + 50,).
i=1
We use 10-dimensional Styblinski-Tany Function as the test function, and Gaussian samples as training data.
((0) 1(0 xx). x);
where data samples x ~ (0, 47). The one-dimensional Styblinski-Tane Function has one global minimum located at = -2.903534, one local minimum located d. and one saddle point b 0.156731 as the boundary separating Valley aj and Valley ag. For a n-dimensional Styblinski- Tang Function, we initialize parameters 0z=0- 1 2229993534...... .... 2.903334) and set the valley's
Vi
boundary as 0; < vR 0.156731, where i is the dimension index. We record the number of iterations required to escape from the valley to the outside of valley.
Observation: we observe the number of iterations from the initialized position to the terminated position. As are more interested in the number of iterations than "dynamical time" in practice, use the number of iterations to denote the mean escape time and ignore the time unit rj in "dynamical time". We repeat experiments 100 times estimate the escape rate I' and the mean escape time T. As the escape time is approximately a random variable obeying an exponential distribution, t Exponential(()) the estimated escape rate can be written as
100
r
(119)
Σ
100 D==1
t
The 95% confidence interval of this estimator is
1.96
1.96
r(1-
I T(1+
(120)
V100
V100
F SUPPLEMENTARY EMPIRICAL RESULTS
Some papers (Luo al. 2019 Chen and Gu 2018 argued that their proposed Adam variants may generalize as well as SGD. But we found that this argument is contracted with our comparative experimental results, such as Table The main problem may lie in weight decay. SGD with weight decay À = 0.0001, a common setting in related papers, is not good baseline on CIFAR- 10 and CIFAR- 100, as À 0.0005 often shows better generalization, seen Figures g We also conduct comparative experiments with X = 0.0001, seen in Table 2 While some Adam variants under this setting sometimes may compare with SGD due the lower baseline performance of SGD, Adai and SGD with fair weight decay still show superior test performance.
We display all learning curves of Adai, SGD, and Adam on CIFAR-10 10 and CIFAR- 100 Figure I0 We further compare convergence of Adai and Adam with the fixed learning rate scheduler for 1000 epochs in Figure
For Language Modeling, we display the results of LSTM in Figure 12
Finally, we use the measure of the expected minima sharpness proposed by Neyshabur et al. (017)) to compare the sharpness of minima learned by Adai, Momentum, and Adam. The expected minima sharpness is defined as Eg (6 - () - L(Ø*)], where ( is Gaussian noise and 6* is the empirical minimizer learned by training algorithm. If the loss landscape near 6* is sharp, the weight-pertuubed loss Eş (6* + {)] will be much larger than L(#*). Figure 13 empirically supports that Adai and Momentum learn significantly flatter minima than Adam.
28

---

Under review as a conference paper at ICLR 2022
0.22
0.22
0.31
0.2
0.24
0.89
A3
Figure 9: The test errors of ResNet18 on CIFAR-10 and VGG16 on CIFAR- 10 under various weight decay. Left: ResNet18. Right: VGG16. The optimal test performance corresponds to À 0.0005. Obviously, Adai has better optimal test performance than SGD with Momentum.
0.44
V
0.30
Are
1200
Acee
Ares
420
400m
(a) ResNet18
(b) VGG16
(c) DenseNet121
(d) GoogLeNet
150
12
156
12
151
122
Figure 10: Generalization and Convergence Comparison. Subfigures (a)-(b): ResNet18 and VGG16 on CIFAR-10. Subfigures (c)-(d): DenseNet121 and GoogLeNet on CIFAR-IOK 100. Top Row: Test curves, Bottom Row: Training curves, Adai with 17 = 1 and = 0.5 converge similarly fast to SGD with Momentum and Adam, respectively, and Adai generalizes significantly better than SGD with Momentum and Adam.
Adal LRO.0 01 Adam LRO.0003
10P
Ad5m UR3c 05
800
800
1000
Spochs
Figure 11: Convergence comparison by training VGG16 on CIFAR-10 for 1000 epochs with the fixed
learning rate. When they converge similarly fast, Adai converges in a lower training loss the end. When they converge in similarly low training loss, Adai converges faster during training.
29

---

Under review as a conference paper at ICLR 2022
15
(a) Test Perplexity
(b) Training Perplexity
Figure 12: Language Modeling. The learning curves of Adai, SGD (with Momentum), and Adam for
LSTM on Penn TreeBank. The optimal test perplexity of Adai, SGD, and Adam are 74.3, 74.9, and 74.3, respectively. Adai and Adam outperform SGD, while Adai may lead to lower training loss than Adam and SGD.
107
Adai
Momeetum
1.0P
10
10-4
10-1
0,00030. 00250 Dww.0...00ooo wwt.:/o0/.........oo www.0011.000..... 2tt. 110.090000..002000000
Weight Noise Scale
Figure 13: The expected minima sharpness analysis of the weight-perturhed training loss landscape of ResNet18 on CIFAR-10. The weight noise scale is the standard deviation of the injected Gaussian noise. The minima learned by Adai and SGD are more robust weight noise. Obviously, Adai and Momentum can learn much flatter minima than Adam in terms of the expected minima sharpness.
30

---

Under review as a conference paper at ICLR 2022
Table 2: Test performance comparison of optimizers with the weight decay hyperparmeter À 0.0001. In this setting, some Adam variants may compare with SGD mainly because the baseline performance of SGD is lower than the baseline performance Table T The test errors of AdaiW, Adai, and Momentum in middle columns is the original results in Table
DATASET
MOZEL
ADAIW
ABAI
SGD M
SGD M
ADAM
AMSGRAD
ADAMW
ADABOUND
PADAN
YoGi
RADAN
CIFAR-10
REsNari8
4.59s.1
4.
1.58
6:35
5.72
0.33
6,57
5.53
5.43
0.81
VGG16
5.81
6,30g.s
6,420us
6.32
7.34
0.68
6,40
7.33
8.74
6.69
6.73
CIEAR 100
RESNE124
21.08g.a
20.79e.2
21.32,.'
24.92
25.36
24.74
23.51
25.57
25.30
28.72
25.60
DINSENET121
19.44 19.44
19.50
19.81e33
20.96
24.39
22.60
22.25
26.25
22.26
22. 60
22.40
GoogLaNer
20.50,.s
20.55
21.21pa
21
24.00
24.00
21.71
25.33
26.68
22.56
22.35
G ADAI WITH STABLE/DDECOUPLEED WEIGHT DECAY
Algorithm 3: AdaiS/AdaiW
vi(O,-1);
= 3201-1 + (1 B2)94; Ft I-
Gt
meam((:):
I1. (1
L0 (e c);
Dz
m 31. + 31,
th 1-1li=1 B1.5"
0. 0,- mp14 A100-1;
H EXPRESSIONS OF ADAM DYNAMICS
In this section, discuss why Adam dynamics can also be expressed as Equation 7 similarly to Momentum dynamics.
The derivation that generalizes Momentum dynamics to Adam dynamics is trivial. We only need to replace by the adaptive f in Equations (4, (5). and (6). We write the updating rule of Adam as
{ my --m--1 + Bagh:
(121)
0c+1 0, fyme.
where 83 1
31. 31 is the hyperparameter, and á is the adaptive learning rate. For simplicity, it fine consider the updating rules as element- wise operation. We can also write the Newtonian motion equations of Adam with the mass M and the damping coefficient 7 as
= (1- - + F dt
(122)
0,11
ridt,
where Ft. = -m == It, dt = f7 1 - ydt = 31. and
= B3 31. Thus, we obtain the M
di
differential-for motion equation of Adam as
MØ M0 + F,
(123)
where the mass M - and the damping coefficient 7 1-8g
This shows that the analysis can applied to both Adam dynamics and Momentum dynamics. This is not surprising, because the Newtonian motion equations, including Equations (5) and (6). are universal. The basic updating rules, Equations (4). generally holds for any optimizer that uses Momentum, including Adam. As the terms in Equations (4) can correspond to the terms in Equations (5) one by one, any dynamics governed by Equations (4) can also corresponds to differential-fom Equation (6), where M and may have different expressions.
We note that Equations (5) in our paper is not contradictory to the SDEs of Adam in Zhou et al. 2200)), as long let the expressions of M and y follow Adam dynamics. In fact, the updating
31

---

Under review as a conference paper at ICLR 2022
rule of of the SDEs in Zhou et al. (2020) can be incorporated into the expressions of M and Equation (5) in our work. which is a single stochastic differential equation, may be more concise for analyzing optimization dynamics than the SDEs in Zhou et al. (2020).
We also point out that the most important property of using Momentum is to employ the phase- dynamics for training of deep neural networks. If we let 31 = 0 in Adam, Adam will be reduced to RMSprop. Similarly SGD, RMSprop has no the momentum drift effect on saddle-poing escaping but has isotropic diffusion effect. Our theoretical analysis about Momentum and Adam helps us understand how manipulate element-wise momentum and learning rate separately to improve saddle pointt escaping we want.
32


Reviews:
[{"review_id": "lwcvMabPi1", "reviewer": "Reviewer_1sdV", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "I believe that the main ideas of the paper are interesting. However I find that the presentation of this work is not very clear and somehow confusing. In particular, the structure of the paper and the results presented in sections 2 and 3 are difficult to absorb. See my comments below.\n\nI understand the motivation of the authors and what they tried to communicate but i find that there is no satisfactory explanation of the results presented in sections 2 and 3. The authors assumed that the reader is familiar with the closely related recent work on the SGD diffusion theory and do not provide enough details on the framework. For example they use terminology like \"Fokker-Planck equation\",  \"divergence operator\", and \"diffusion matrix\" that are not really standard in the area of adaptive methods. In addition, Assumption 1 on the second order Taylor approximation near critical points, is given without providing some interesting problems where it is satisfied.\n\nAlso in section 3, Assumptions 2 and 3 are used without further explanation of what exactly they mean. The authors provide a few details in the appendix on what is the quasi-equilibrium approximation and low temperature approximation but this is not sufficient.  How these assumptions are related to standard concepts in the area? What are the mathematical expressions of these assumptions? how are related to stochastic gradients and the noise? \n\nAlso i find it a bit surprising that there is no formal presentation of the problem that we are interested to solve and the assumptions that one requires to be able to prove convergence. A statement of the minimization (or maximation) problem under study with the main assumptions is missing from the paper. \n\nOne of the most important contributions of the paper is the analysis of the new algorithm, Adaptive Inertia Optimization (Adai) proposed in Section 5. However if one focuses on Theorem 4 which provides the convergence guarantees of the Adai it is clear that the analysis hold under very strong conditions / assumptions. For example the authors assumed both bounded variance, and bounded gradient of the objective function, which rarely hold in practical scenarios. Note that these conditions have already been proved to contradict special classes of non-convex problems like functions satisfying the Polyak-Lojasiewicz condition (the combination of these assumptions lead to an empty set of problems). Thus the theorem cannot hold for all non-convex smooth problems.", "summary_of_the_paper": "The paper focuses on the understanding of the effects of adaptive learning rate and momentum. In particular it proves that the adaptive learning rate can escape saddle points efficiently and cannot select flat minima as SGD does. It also shows that momentum helps the training process by passing through saddle points and without affecting the minima selection.The paper also proposes a new adaptive algorithm, named Adai (Algorithm 2), which uses parameter-wise adaptive intertia to accelerate the training and finds flat minima as well as SGD.  Finally, the paper provides extensive numerical testing showing the benefits of Adai.", "summary_of_the_review": "As i mentioned in the main review, I believe that the main ideas of the paper are interesting. However I find that the presentation of this work is not very clear and somehow confusing. In particular, the structure of the paper and the results presented in sections 2 and 3 are difficult to absorb. See my comments below. ", "strength_and_weaknesses": "", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Ak5Uw9oI5c2", "reviewer": "Reviewer_MYxM", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "The theoretical analysis of the Adam optimizer is based on the SGD diffusion theory, and the results confirms and explains the observation that Adam can sometimes converge faster but generalize worse than SGDM. The proposed Adai optimizer is theoretically sound, and demonstrates slightly better generalization performance than SGDM (and significantly better than Adam) on image classification tasks.\n\nDespite estimating the moments in a similar way as Adam, the proposed Adai optimizer seems more akin to SGDM, with the only difference being its adaptive momentum; and it doesn't use adaptive learning rates, which is a main feature of Adam. Moreover, as shown in Figs. 1 to 3, and 10, the training curves of Adai very much resemble those of SGDM. Therefore, to further improve this work, more comparisons should be made between Adai and SGDM rather than between Adai and Adam. In particular, it would be interesting to see if the performance gap between Adai and SGDM results from faster convergence (as suggested by the theory), and therefore a convergence comparison between Adai and SGDM as the one conducted between Adai and Adam (Fig. 11) should be helpful.", "summary_of_the_paper": "This paper disentangles the effects of adaptive learning rate and momentum in Adam learning dynamics, and proves that adaptive learning rate is good at escaping saddle points but not good at selecting ﬂat minima, while momentum helps escape saddle point and matters little to escaping sharp minima. Based on the analysis, the authors propose a novel optimizer, Adai. Compared to SGDM, Adai parameter-wisely adapts the momentum hyperparameter to the (approximated) Hessians of saddle points, and is proved to fast escape saddle points and sharp minima.", "summary_of_the_review": "This paper provides new insights into the performance of Adam, and proposes a novel optimizer that both converges fast and generalizes well. Further improvements can be made by comparing the proposed method to SGDM more thoroughly.", "strength_and_weaknesses": "", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "4JRxRkpvNOp", "reviewer": "Reviewer_EqDN", "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": "Here are some of the questions and concerns:\n\nFrom my perspective, the theoretical analysis is not rigorous and skips a lot of details. For example, in Section 4 (Page 5), the paper claims that the continuous-time dynamic of Adam can be written as equation (6). However, a detail derivation to support the claim and the proof of showing that the SDE is a valid one for Adam are missing in the paper. Indeed, it can be seen from Page 5, where the authors said they are analyzing an \"idealized\" Adam. It is not clear why this is an idealized one and the writing is like hiding something under the rug. We can see that some statements about Adam are made, but these conclusions are drawn from a SDE which does not really correspond to Adam. On the other hand, the SDE of Adam does exist in the literature, e.g. https://arxiv.org/pdf/2010.05627.pdf. If this paper really wants to claim/argue something about Adam, then a careful analysis regarding the discretization error between the solution of the proposed SDE and the discrete-time Adam should be provided in the paper. \n\nFor another example, on page 3, it states QUOTE the diffusion matrix D is independent of $\\theta$ near critical points UNQUOTE. But the proof is not provided. How is it independent? It looks like some approximations were used there. The paper should prove the \"independence\".\n\nAnother concern is about (8). The paper should provide a detailed derivation of showing that (8) is really about how the distribution evolves when the underlying dynamic is (7). Currently the equation (8) is like jumping out of nowhere. How does Assumption 1 help to show (8)?\n\nAlso, some places in the paper are not clear:\n\n(a) (Second paragraph on Page 1) QUOTE all previous works have not touched the saddle-point escaping property of the dynamics UNQUOTE Apparently there are quite a few works regarding saddle-point escaping by SGD, SGD with momentum, and Adam. The authors might want to explain what they meant here.\n\n(b) There are two approximations on (3). But it would be more helpful to explain how the approximations are made in detail. There are some descriptions below (3) but are not very clear. \n\n(c) (Last paragraph in Section 3) QUOTE The momentum does not affect flat minima selection in terms of the escape time UNQUOTE This is another confusing statement. What does \"momentum does not affect flat minima selection\" mean? Does it mean momentum and SGD without momentum converge to the same point? What is the definition of \"flat minima\"?\n\n(d) (Second to last paragraph in Section 4) QUOTE Adam has $\\log(\\tau)=O(H_{ae}^{-1/2})$ ... SGD and SGD+momentum both have $log(\\tau) = O(H_{ae}^{-1})$ UNQUOTE It seems that the conclusion right below this sentence would be reversed if $|H_{ae}|>1$. Something wrong?\n\n(e) (Theorem 1 and 2) The authors show some guarantees about the variance at time $t$ when the iterate of the algorithm follows the Gaussian distribution. But does a higher variance of the Gaussian distribution imply a faster saddle point escape? The authors might want to add some discussions about the connection to the notion of saddle point escape in the literature (e.g  Jin et al. 2017).\n\n(f) After reading the paper, I am still not sure how the effect of the learning rate and momentum was \"disentangled\" in the analysis. I see some analysis about the behavior of SGD, ADAM, SGD+momentum at critical points.  It would be more helpful if the authors can explain why the effect of the learning rate and momentum can be isolated.\n\n", "summary_of_the_paper": "This paper studies the behaviors of some algorithms when the iterate is at a critical point via SDEs. A variant of Adam is given in the end. The paper draws some conclusions about adaptive learning rate and momentum, claiming that \n\n- QUOTE momentum matters little to escaping sharp minima UNQUOTE\n- QUOTE adaptive learning rate is not good at selecting flat minima UNQUOTE\n", "summary_of_the_review": "The presentation and statements are confusing in my opinion. Some steps in the analysis are not transparent. ", "strength_and_weaknesses": "", "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": "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": "36DZfnqI2q7", "reviewer": "Reviewer_A3Wq", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "This paper is generally well-written. The diffusion theoretical analysis does provide some insight on the empirical performance of momentum SGD and Adam. The authors also put in efforts to conduct numerical verifications to their theoretical statements, which is highly appreciated. However, I think that this work does not completed ''disentangle'' the effects of adaptive learning rate and momentum since the work analyzes Adam, which fuses these two algorithmic components. It would be better to discuss the effect of each component in Adam separately (probably by setting some parameters to zero). \n\nThe authors then propose Adai, which achieves approximately Hessian-independent momentum drift without damaging flat minima selection (BTW, the proof of Proposition 3 is missing, if it is a direct consequence of Theorem 2, it is better to mention it somehow). The construction of Adai is interesting, and its effectiveness is justified by the empirical experiments. However, it seems to me that this contribution is a bit disconnected to the main story as Adai does not use adaptive learning rate. Some revision (probably changing the title?) might be good to make the story clearer and more fluent.\n\nTypos:\n- Missing reference on page 7, \"Note that an existing “adaptive momentum” method (?) \"\n- The last sentence \"better than popular Adam and SGD\".", "summary_of_the_paper": "This work analyzes the dynamics of momentum SGD and Adam on escaping saddle points and sharp minima, which is based on the diffusion theoretical framework proposed in (Xie et al. 2021b). The authors prove that momentum provides a drift effect around saddle points and does not affect flat minima selection (for SGD), and while Adam escapes saddle points efficiently, it does not favor flat minima as well as SGD. The analysis explains some empirical observation of SGDM and Adam. Motivated by the analysis, the authors propose adaptive inertia (Adai) method, which can approximately achieve Hessian-independent momentum drift (escapes saddle points fast) and favors flat minima as well as (momentum) SGD.  ", "summary_of_the_review": "This work provides some new theoretical insights for momentum SGD and Adam, which are interesting and important. The authors then propose adaptive inertia based on the insights, which shows good performance. Some revision is needed to make the story clearer (see main review).", "strength_and_weaknesses": "", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper is aimed at providing an explaining the perceived lack of generalization results for Adam as compared to SGD. To this end the paper decouples the effect of adaptive per parameter learning rate and the momentum aspect of Adam. The paper shows that the while adaptive rates help escape saddle points faster - they are worse when consider the flatness of minima being selected. Further momentum has no effect on the flatness of minima but again leads to better optimization by providing a drift leading to saddle point evasion. They also provide a new algorithm Adai (based on inertia) targeted at better generalization of adaptive methods. 

The paper definitely provides an interesting perspective and the approach to decouple the effect of momentum and adaptive LR and study their efficacy in escaping saddle points and flatness of minima seems a very useful perspective. The primary reason for my recommendation is the presentation of the paper in terms of the rigor its assumptions to establish the results. These aspects have been highlighted by the reviewers in detail. I suggest the authors to carefully revisit the paper and improve the presentation of the assumptions, adding rigor to the presentation as well as adding justifications where appropriate especially in light of non-standardness of these assumptions in optimization literature.
Decision: Reject