Year: 2021

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 saddlepoint 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 parameterwise 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.

heading: 1 INTRODUCTION
text: 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 a tool to study how SGD selects minima (Jastrzkebski et al., 2017; Li et al., 2017; Wu et al., 2018; Xu et al., 2018; Hu et al., 2019; Nguyen et al., 2019; Zhu et al., 2019; Xie et al., 2021b; Li et al., 2021). This line of research also suggests that injecting or enhancing SGD-like sharpness-dependent gradient noise may effectively help find flatter minima (An, 1996; Neelakantan et al., 2015; Zhou et al., 2019; Xie et al., 2021a;c; HaoChen et al., 2021). Especially, Zhou et al. (2020) argued the better generalization 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, it 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 et 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 minima1. 2) We propose a novel Adaptive Inertia (Adai) optimization framework
1When we talk about saddle points, we mean strict saddle points in this paper.
which is conceptually orthogonal to the existing adaptive gradient framework. Adai does not parameter-wisely 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 1), Adai significantly outperforms SGD and popular Adam variants.
The paper is organized as follows. In Section 2, we first introduce the dynamics of SGD and analyze its behavior on saddle-point escaping, which serves as a basis to compare with the behavior of Adaptive Learning Rate and Momentum. In Section 3, we analyze the behavior of Momentum. In Section 4, we analyze the dynamics of Adam. In Section 5, we introduce the new optimizer Adai. In Section 6, we empirically compare the performance of Adai, Adam variants, and SGD. Section 7 concludes the paper with remarks.

heading: 2 SGD AND DIFFUSION
text: In this section, we introduce some preliminaries about the SGD diffusion and then present the saddle-point escaping property of SGD dynamics.

heading: 2.1 PREREQUISITES FOR SGD DIFFUSION
text: We first review the SGD diffusion theory for escaping minima proposed by Xie et al. (2021b). We denote the model parameters as θ, the learning rate as η, the batch size asB, and the loss function over one minibatch and the whole training dataset as L̂(θ) and L(θ), respectively. A typical optimization problem can be formulated as minθ L(θ). We may write the stochastic differential equation/Langevin Equation that approximates SGD dynamics (Mandt et al., 2017; Li et al., 2019) as
dθ = −∇L(θ)dt+ [ηC(θ)] 12 dWt, (1)
where dWt ∼ N (0, Idt), I is the identity matrix, and C(θ) 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, ξ = ∇L̂(θ)−∇L(θ). It 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
∂P (θ, t)
∂t = ∇ · [P (θ, t)∇L(θ)] +∇ · ∇D(θ)P (θ, t), (2)
where∇· is the divergence operator and D(θ) = ηC(θ)2 is the diffusion matrix (Xie et al., 2021b). We note that the dynamical time t is equal to the product of the number of iterations T and the learning rate η: t = ηT .
As the gradient variance dominates the gradient expectation near critical points, we have
D(θ) = ηC(θ) 2 ≈ η 2B  1 N N∑ j=1 ∇Lj(θ)∇Lj(θ)>  = η 2B FIM(θ) ≈ η 2B [H(θ)]+ (3)
near a critical point c, where Lj(θ) is the loss function of the j-th training sample,H(θ) is the Hessian of the loss function at θ, and FIM(θ) is the observed Fisher Information matrix, referring to Chapter 8 of Pawitan (2001) and Zhu et al. (2019). We further verified that Equation (3) approximately
holds even not around critical points in Figure 6 of Appendix C. Equation (3) was also proposed by Jastrzkebski et al. (2017) and Zhu et al. (2019) and verified by Xie et al. (2021b) and Daneshmand et al. (2018). Please refer to Appendix C for the detailed analysis of the stochastic gradient noise.
If we have eigendecomposition H = U diag(H1, . . . ,Hn−1, Hn)U>, then we use [·]+ to denote the transformation that [H]+ = U diag(|H1|, . . . , |Hn−1|, |Hn|)U>. The i-th column vector of U is the eigenvector of H corresponding to the eigenvalue Hi.
In the following analysis, we use the second-order Taylor approximation near critical points. This assumption is common and mild, when we focus on the behaviors near critical points (Mandt et al., 2017; Zhang et al., 2019a; Xie et al., 2021b). Note that, by Equation (3) and Assumption 1, the diffusion matrix D is independent of θ near critical points. Assumption 1. The loss function around the critical point c can be approximately written as
L(θ) = L(c) + 1
2 (θ − c)>H(c)(θ − c).

heading: 2.2 SGD DIFFUSION NEAR SADDLE POINTS
text: In this subsection, we establish the saddle-point escaping property of SGD diffusion as Theorem 1. Theorem 1 (SGD Escapes Saddle Points). Suppose c is a critical point, Assumption 1 and Equation (3) hold, the dynamics is governed by Equation (1), and the initial parameter is at the saddle point θ = c. Then the probability density function of θ after time t is given by the Gaussian distribution θ ∼ N (c,Σ(t)), where Σ(t) = U diag(σ21 , . . . , σ2n−1, σ2n)U> and
σ2i (t) = Di Hi [1− exp(−2Hit)],
where Di is the i-th eigenvalue of the diffusion matrix D and Hi is the i-th eigenvalue of the Hessian matrix H at c. The column vectors of U are exactly the eigenvectors of H . The dynamical time t = ηT . In terms of SGD notations and |Hi|ηT 1 near saddle points, we have
σ2i (T ) = |Hi|η2T
B +O(B−1H2i η3T 2).
The proof is relegated to Appendix A.1. We note that 1) if Hi > 0, the distribution of θ along the direction i converges to a Gaussian distribution with constant variance,N ( ci, η 2B ) ; and 2) if Hi < 0, the distribution is a Gaussian distribution with the variance exponentially increasing with time t.
As the displacement ∆θi from the saddle point c can be modeled as a center-fixed Gaussian distribution, the mean squared displacement is equivalent to the variance, namely 〈∆θ2i (t)〉 = σ2i (t). The result means that SGD escapes saddle points very slowly (〈∆θ2i 〉 = O(|Hi|)) if Hi is close to zero. Note that, in the diffusion analysis, the direction i denotes the direction of an eigenvector instead of a coordinate’s direction. While SGD updates model parameters along the coordinates, we do not need to treat the coordinates’ directions specially in the continuous-time analysis.

heading: 3 ANALYSIS OF MOMENTUM DYNAMICS
text: 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 {
mt = β1mt−1 + β3gt, θt+1 = θt − ηmt, (4)
where β1 and β3 are the hyperparameters. We note there are two popular choices, which are, respectively, β3 = 1 corresponding to SGD-style Momentum and β3 = 1 − β1 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 γ as{ rt = (1− γdt)rt−1 + FM dt θt+1 = θt + rtdt, (5)
where rt = −mt, F = gt, dt = η, 1 − γdt = β1, and dtM = β3. Thus, we obtain the differential form of the motion equation as
Mθ̈ = −γMθ̇ + F, (6)
where θ̈ = d 2θ dt2 and θ̇ = dθ dt . The left-hand side as the inertia force is equal to the mass M times the acceleration θ̈, the first term in the right-hand side as the damping force is equal to the damping coefficient γ times the physical momentum Mθ̇, and the second term in the right-hand side is equal to the external force F in physics. We can easily obtain the mass M = ηβ3 and the damping coefficient γ = 1−β1η by comparing (5) and (6).
As F corresponds to the stochastic gradient term, we obtain
Mdθ̇ = −γMdθ − ∂L(θ) ∂θ dt+ [2D] 1 2 dWt. (7)
Its Fokker-Planck Equation in the phase space (the θ-θ̇ space) is well known as
∂P (θ, r, t)
∂t =−∇θ · [rP (θ, r, t)] +∇r ·
[ γr +M−1∇θL(θ) ] P (θ, r, t)
+∇r ·M−2D · ∇rP (θ, r, t) (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-inertia Langevin Dynamics, Equation (7). We contribution is the first to apply it to deep learning dynamics.
Momentum escapes saddle points. We formulate how Momentum escapes saddle points as Theorem 2, whose proof is relegated to Appendix A.2.
Theorem 2 (Momentum Escapes Saddle Points). Suppose c is a critical point, Assumption 1 and Equation (3) hold, the dynamics is governed by Momentum, and the initial parameter is at the saddle point θ = c. Then the mean squared displacement near saddle points is
〈∆θ2i (t)〉 = Di γ3M2 [1− exp(−γt)]2 + Di γMHi [1− exp(−2Hit γM )], (9)
where ∆θ(t) = θ(t)− θ(0) is the displacement of θ, 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 |Hi|ηT 1 near ill-conditioned saddle points, it can be written in terms of Momentum notations as
〈∆θ2i 〉 = |Hi|β23η2
2(1− β1)3B [1− exp (−(1− β1)T )]2 + |Hi|β23η2T B(1− β1)2 +O(B−1H2i η3T 2).
By comparing Theorems 1 and 2, we notice that, SGD escapes saddle points only due to the diffusion effect (similar to the second term in Equation (9)), but Momentum provides an additional momentum drift effect (the first term in Equation (9)) for passing through saddle points (Wang et al., 2019). This momentum drift effect has not been mathematically revealed before, to the best of our knowledge.
Momentum escapes minima. We first introduce two classical assumptions. Assumption 2 (Quasi-Equilibrium Approximation). The system is in quasi-equilibrium near minima. Assumption 3 (Low Temperature Approximation). The system is under low temperature (small gradient noise).
Xie et al. (2021b) modeled the process of SGD escaping minima as a Kramers Escape Problem in deep learning. Recent machine learning papers (Jastrzkebski et al., 2017; Xie et al., 2021b; Zhou et al., 2020) also used Assumptions 2 and 3 as the background implicitly or 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(Kramers, 1940; Hanggi, 1986), chemistry(Eyring, 1935; Hänggi et al., 1990), biology(Zhou, 2010), electrical engineering(Coffey and Kalmykov, 2012), and stochastic process(Van Kampen, 1992; Berglund, 2013).
Assumptions 2 and 3 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 to 100,000 iterations) can be modeled as a Kramers Escape Problem very well. Our empirical results in Section 6 support this point again. Thus, Assumptions 2 and 3 are reasonable in practice (see more discussions in Appendix B).
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-Planck Equation (8) in Kalinay and Percus (2012). We analyze the mean escape time τ (Van Kampen, 1992; Xie et al., 2021b) required for the escaping process from a loss valley. Suppose that the saddle point b is the exit of Loss Valley a, and ∆L = L(b)− L(a) is the barrier height(See more details in Appendix E.). Theorem 3 (Momentum Escapes Minima). Suppose Assumptions 1, 2, and 3 hold, and the dynamics is governed by Momentum. Then the mean escape time from Loss Valley a to the outside of Loss Valley a is given by
τ = π [√ 1 +
4|Hbe| γ2M + 1
] 1
|Hbe| · exp
[ 2γMB∆L
η
( s
Hae + (1− s) |Hbe|
)] ,
where the subscript e indicates the escape direction, s ∈ (0, 1) is a path-dependent parameter, and Hae and Hbe are the eigenvalues of the Hessians of the loss function at the minimum a and the saddle point b along the escape direction e. When 4|Hbe|γ2M 1, it reduces to SGD. In terms of Momentum
notations, we have log(τ) = O (
2(1−β1)B∆L β3ηHae
) .
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(τ) = O ( 2B∆L ηHae ) has been obtained for SGD (Xie et al., 2021b). We see that Momentum does not affect flat minima selection in terms of the mean escape time, if we properly choose the learning rate, i.e., ηMomentum = 1−β1β3 ηSGD.

heading: 4 ANALYSIS OF ADAM DYNAMICS
text: Algorithm 1: Adam
gt = ∇L̂(θt); mt = β1mt−1 + (1− β1)gt; vt = β2vt−1 + (1− β2)g2t ; m̂t =
mt 1−βt1 ; v̂t =
vt 1−βt2 ; θt+1 = θt − η√v̂t+ m̂t;
Algorithm 2: Adai
gt = ∇L̂(θt); vt = β2vt−1 + (1− β2)g2t ; v̂t =
vt 1−βt2 ; v̄t = mean(v̂t); β1,t = (1− β0v̄t v̂t).Clip(0, 1− ); mt = β1,tmt−1 + (1− β1,t)gt; m̂t =
mt 1− ∏t z=1 β1,z
; θt+1 = θt − ηm̂t;
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 a good way to help escape saddle points: adaptively adjust the learning rates for different parameter as ηi ∝ |Hi|− 1 2 . However, estimating the Hessian is expensive in practice. In Adam (Algorithm 1), the diagonal v can be regarded as a diagonal approximation of the full matrix E[gtg>t ] (Staib et al., 2019). Following Staib et al. (2019), our analysis considers the “idealized” Adam where v = E[gtg>t ] is a full matrix. Note that E[gtg>t ] = C(θ) = HB approximately holds near critical points.
The continuous-time Adam dynamics. The continuous-time dynamics of Adam can also be written as Equation (6), except the mass M = η̂1−β1 and the damping coefficient γ = 1−β1 η̂ . We emphasize that the learning rate is not a real number but the diagonal approximation of the ideal learning rate
matrix η̂ = ηC− 1 2 in practice. We apply the adaptive time continuation dt̂i = η̂i in the i-th dimension, where t̂i of T iterations is defined as the sum of η̂i of each iteration. Thus, the Fokker-Planck Equation for Adam can still be written as Equation (8). We leave more details in Appendix H.
Adam escapes saddle points. Similarly to Theorem 2, we formulate how Adam escapes saddle point as Proposition 1, which can be obtained from Theorem 2 and η̂ = ηC− 1 2 .
Proposition 1 (Adam Escapes Saddle Points). Suppose c is a critical point, Assumption 1 and Equation (3) hold, the dynamics is governed by Adam, and the initial parameter is at the saddle point θ = c. Then the mean squared displacement is written as
〈∆θ2i (t̂i)〉 = Di γ2M [1− exp(−γt̂i)]2 + Di γMHi [1− exp(−2Hit̂i γM )].
Under |Hi|ηT 1, it can be written in terms of Adam notations as
〈∆θ2i 〉 = η2
2(1− β1) [1− exp (−(1− β1)T )]2 + η2T +O(
√ B|Hi|η3T 2).
From Proposition 1, we can see that Adam escapes saddle points fast, because both the momentum drift and the diffusion effect are approximately Hessian-independent and isotropic near saddle points, which is also supported by Staib et al. (2019). Proposition 1 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 at a minimum:
τ = π √1 + 4η√B|Hbe| 1− β1 + 1  |det(H−1a Hb)| 14 |Hbe| exp [ 2 √ B∆L η ( s√ Hae + (1− s)√ |Hbe| )] ,
where τ is the mean escape time. Thus, we have log(τ) = O ( 2 √ B∆L
η √ Hae
) .
In comparison, Adam has log(τ) = O(H− 1 2
ae ), while SGD and Momentum both have log(τ) = O(H−1ae ). We say that Adam is weaker Hessian-dependent than SGD in terms of sharp minima escaping, which means that ifHae increases, i.e., the escaping direction becomes steeper, the escaping time log(τ) for SGD decreases more dramatically than log(τ) for Adam. Note that, in the diffusion theory, the minima sharpness is reflected by Hae, the eigenvalue of the Hessian corresponding to the escape direction along an eigenvector. The weaker Hessian-dependent 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 (Luo et al., 2019), Padam (Chen and Gu, 2018), RAdam (Liu et al., 2019), AMSGrad (Reddi et al., 2019), Yogi (Zaheer et 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 al., 2019b), which we also discuss in Appendix F. 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.

heading: 5 ADAPTIVE INERTIA
text: In this section, we propose a novel adaptive inertia optimization framework.
Adaptive Inertia Optimization (Adai). The basic idea of Adai (Algorithm 2) comes from Theorems 2 and 3, from which we can see that parameter-wise adaptive inertia can achieve the approximately
Hessian-independent momentum drift without damaging flat minima selection. The total momentum drift effect during passing through a saddle point in Theorem 2 is given by
〈∆θi〉2 = |Hi|η2
2(1− β1)B . (10)
We generalize the scalar β1 in Momentum to a vector β1 in Adai. The ideal adaptive updating is to keep β1 = 1 − β0v̄ v, where the rescaling factor v̄, namely the mean of all elements in the estimated v̂, could be used to make the optimizer more robust in practice. The default values of β0 and β2 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 = η(1 − β1,max)−1 = η −1. Compared with Adam, Adai does not increase the number of hyperparameters. Note that an existing “adaptive momentum” method (Wang and Ye, 2020) 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 = η(I − diag(β1))−1 and γ = η−1(I − diag(β1)). Adai escapes saddle points. Proposition 3 shows that Adai can escape saddle points efficiently due to an isotropic and approximately Hessian-independent momentum drift, while the diffusion effect is the same as Momentum. Proposition 3 is a direct result of Theorem 2. Proposition 3 (Adai Escapes Saddle Points). Suppose c is a critical point, Assumption 1 and Equation (3) hold, the dynamics is governed by Adai, and the initial parameter is at the saddle point θ = c. Then the total momentum drift in the procedure of passing through the saddle point is given by
〈∆θi〉2 = v̄η2
β0 =
∑n i=1 |Hi|η2
β0nB , where ∑n i=1 |Hi| 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 or as 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:
τ = π √1 + 4η∑ni=1 |Hbi| β0n + 1  1 |Hbe| exp [ 2B∆L η ( s Hae + (1− s) |Hbe| )] ,
where τ is the mean escape time. Thus, we have log(τ) = O (
2B∆L ηHae
) .
We can see that both SGD and Adai have log(τ) = O(H−1ae ), while Adam has log(τ) = O(H − 12 ae ). Thus, Adai favors flat minima as well as SGD and better than Adam.
Convergence Analysis. Theorem 4 proves that Adai has similar convergence guarantees to SGD Momentum(Yan et al., 2018; Ghadimi and Lan, 2013). The proof is relegated to Appendix A.6. Theorem 4 (Convergence of Adai). Assume that L(θ) is an L-smooth function2, L is lower bounded as L(θ) ≥ L?, E[ξ] = 0, E[‖g(θ, ξ) − ∇L(θ)‖2] ≤ δ2, and ‖∇L(θ)‖ ≤ G for any θ, where ξ represents the gradient noise of sub-sampling. Let Adai run for t+ 1 iterations and β1,max = 1− ∈ [0, 1) for any t ≥ 0. If η ≤ C√
t+1 , we have
min k=0,...,t E[‖∇L(θk)‖2] ≤ 1√ t+ 1 (C1 + C2 + C3),
where C1 = L(θ0)−L? (1−β1,max)C , C2 = β1,maxC 2(1−β1,max)2G 2, and C3 = LC2(1−β1,max)2 (G 2 + δ2).
2It means that ‖∇L(θa)−∇L(θb)‖ ≤ L‖θa − θb‖ holds for any θa and θb.
0 25 50 75 100 125 150 175 200 Epochs
0.04
0.06
0.08
0.10
0.12
0.14
Te st
E rro
r
Adai Momentum Adam AMSGrad AdamW AdaBound Padam Yogi RAdam
0 25 50 75 100 125 150 175 200 Epochs
0.175
0.200
0.225
0.250
0.275
0.300
0.325
0.350
0.375
Te st
E rro r Adai Momentum Adam AMSGrad AdamW AdaBound Padam Yogi RAdam
Figure 1: The learning curves on CIFAR-10 and CIFAR-100. Left: ResNet18 on CIFAR-10. Right: ResNet34 on CIFAR-100.
0 20 40 60 80 100 Epochs
0.22
0.24
0.26
0.28
0.30
0.32
Te st
E rro
r
Adai Momentum Adam
0 20 40 60 80 100 Epochs
100
101
Tr ai
ni ng
L os
s
Adai Momentum Adam
Adai SGD M Adam
Top1 23.20 23.51 27.13 Top5 6.62 6.82 9.18
Figure 3: ResNet50 on ImageNet. Left Subfigure: Top 1 Test Error. Right Subfigure: Training Loss. Table: Top-1 and top-5 test errors. Note that the popular SGD baseline performance of ResNet50 on ImageNet has the test errors as 23.85% in 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., 2020) (see more discussions in Appendix G.) or other techniques from adaptive gradient methods.

heading: 6 EMPIRICAL ANALYSIS
text: In this section, we first conduct experiments to 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-100(Krizhevsky and Hinton, 2009), ImageNet(Deng et al., 2009), and Penn TreeBank(Marcus et al., 1993). Models: ResNet18/ResNet34/ResNet50 (He et al., 2016), VGG16 (Simonyan and Zisserman, 2014), DenseNet121 (Huang et al., 2017), GoogLeNet (Szegedy et al., 2015), and Long Short-Term Memory (LSTM) (Hochreiter and Schmidhuber, 1997b). More details and results can be found in Appendix D and Appendix F.
Generalization and Convergence Speed. We display the test performance of all models on CIFAR in Table 1. Table 1 and Figure 1 shows that Adai has excellent generalization compared with other popular optimizers on CIFAR-10/CIFAR-100. Figure 2 further demonstrates that Adai can 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. Figure 4 demonstrates that Adai not only has better optimal test performance, but also has a 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.
The Mean Escape Time Analysis. We empirically study how the escape rate Γ, which equals to the inverse mean escape time, depends on the minima sharpness for different optimizers in Figure 5. We use Styblinski-Tang Function as the test function which has clear boundaries between loss valleys. Our method for adjusting the minima sharpness is to multiply a rescaling factor √ k to each parameter, and the Hessian will be proportionally rescaled by a factor k. If we let L(θ) = f(θ)→ L(θ) = f( √ kθ), then H(θ) = ∇2f(θ) → H(θ) = k∇2f(θ). Thus, we can use k to indicate the relative minima sharpness. We leave more details in Appendix E.
From Figure 5, it is easy to see that we have log(τ) = O(H−1ae ) in Adai and SGD, while we have log(τ) = O(H− 1 2
ae ) in Adam. Moreover, Adam is 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-based models (Vaswani et 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. (2017) and empirically verified again that Adai and Momentum can learn flatter minima than Adam.

heading: 7 CONCLUSION
text: To the best of our knowledge, we are the first to theoretically disentangle the effects of Adaptive Learning Rate and Momentum in 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-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.

heading: A PROOFS

heading: A.1 PROOF OF THEOREM 1
text: Proof. It is easy to validate that the probability density function
P (θ, t) = 1√
(2π)n det(Σ(t)) exp
( −1
2 (θ − c)>Σ(t)(θ − c)
) (11)
is the solution of the Fokker-Planck Equation (2). Without losing generality, we only validate one-dimensional solution, such as Dimension i.
The first term in Equation (2) can be written as
∂P (θ, t) ∂t = −1 2 1√ 2πσ2 1 σ2 exp
( − θ 2
2σ2
) ∂σ2
∂t + 1√ 2πσ2 exp
( − θ 2
2σ2
) θ2
2σ4 ∂σ2 ∂t (12)
= 1
2
( θ2
σ4 − 1 σ2
) P (θ, t) ∂σ2
∂t . (13)
The second term in Equation (2) can be written as
∇ · [P (θ, t)∇L(θ)] = P (θ)H +Hθ 1√ 2πσ2 exp
( − θ 2
2σ2 )( − θ σ2 ) (14)
= H ( 1− θ 2
σ2
) P (θ, t). (15)
The third term in Equation (2) can be written as
D∇2P (θ, t) = −Dσ 2 − θ2
σ5 √ 2π exp
( − θ 2
2σ2
) (16)
= D
( θ2
σ4 − 1 σ2
) P (θ, t). (17)
By Term1 = Term2 + Term3, we have
1 2 (θ2 − σ2)∂σ
2
∂t = Hσ2(σ2 − θ2) +D(θ2 − σ2) (18)
∂σ2
∂t = 2D − 2Hσ2. (19)
The initial condition of σ2 is given by σ2(0) = 0. We can validate that σ2 satisfies
σ2i (t) = Di Hi [1− exp(−2Hit)]. (20)
It is true along all eigenvectors’ directions.
By D = η2BH , we can get the results of SGD diffusion:
σ2i (T ) = sign(Hi) η
2B [1− exp(−2HiηT )] (21)

heading: A.2 PROOF OF THEOREM 2
text: Proof. Without losing the generality, we consider the one-dimensional case that aligns with an eigenvector of the Hessian. The stationary solution in equilibrium to the phase-space Fokker-Planck Equation near a critical point is given by a canonical ensemble
Peq(θ, r) = Z −1 exp[−β(L(θ) + 1
2 Mr2)], (22)
where Z is a partition function and the inverse temperature β = γMD−1. This result is famous in a large number of physics literature (Van Kampen, 1992; Risken, 1996; Balakrishnan, 2008). Thus, we have the equilibrium distribution in velocity space and position space as
Peq(r) = Z −1 r exp(−
βMr2
2 ) (23)
and
Peq(θ) = Z −1 r exp(−βL(θ)), (24)
respectively, where Zr is the partition function for Peq(r) and Zθ is the partition function for Peq(θ).
It is known that the phase-space probability density solution P (θ, r, t) can be obtained from the position-space solution P (θ, t) and the velocity-space solution P (r, t). Thus, we may compute the velocity-space solution (in Step I) and the position-space solution (in Step II), respectively.
Step I:
We further argue that the velocity-space solution P (r, t) has reached its equilibrium Boltzmann distribution Peq(r) when the particle is passing saddle points. This is reasonable because the equilibrium solution in the velocity space is stable under Assumption 1. The velocity distribution P (r, t) must be the solution of the free diffusion equation in the velocity space. We can approximately ignore the gradient expectation term −∂L(θ)∂θ in dynamical equations near critical points, as 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, t) ≈ Peq(r) = (2π)− n 2 det(βM) 1 2 exp(−βMr
2
2 ), (25)
where n = 1 in one-dimensional case. This is a classical Boltzmann distribution. The expected velocity (also called “equilibrium velocity”) along dimension i can be given by r2eq,i = Di γM2 .
Step II:
As we discussed in Theorem 1, the position-space distribution in equilibrium is time-dependent near saddle points. Following the form of the solution to the position-space Fokker-Planck Equation (Seen in Appendix A.1), the ansatz solution of P (θ, t) is given by{
P (θ, t) = 1√ (2π)n det(Σ(t)) exp ( − 12 (θ − c(t)) >Σ(t)(θ − c(t)) ) Σ(t) = U diag(σ21(t), . . . , σ 2 n−1(t), σ 2 n(t))U >
The time-dependent components in P (θ, 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 we studied in Theorem 1, the second effect is the diffusion effect (due to random noise), which decides the covariance of the probability density, Σ(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
Mc̈(t) = −γMċ(t). (26) We also have ignored the conservative force −H[c(t) − c(0)] near saddle points, as |Hi| γM exists for ill-condition Hessian eigenvalues. We focus on the behaviors near ill-conditioned saddle points, where Hessian eigenvalues are small along the escape directions.
The initial condition is given by c(0) = c, ċ(0) = req, and c̈(0) = −γreq. Then we can obtain the solution c(t) as
ci(t) = ci + ri,eq γ [1− exp(−γt)]. (27)
Step II (b): the diffusion effect.
We can get the dynamics of the diffusion effect as
γMdθ = −∂L(θ) ∂θ dt+ [2D(θ)] 1 2 dWt. (28)
This is equivalent to SGD dynamics with η̂ = ηγM . The expression of P (θ, t) and σ 2 i (t) is directly given by Theorem 1 as
σ2i (t) = Di γMHi [1− exp(−2Hit γM )]. (29)
We combine the momentum drift effect and the diffusion effect together, and then obtain the mean squared displacement of θ as
〈∆θ2i (t)〉 = (ci(t)− ci)2 + σ2i (t) = Di γ3M2 [1− exp(−γt)]2 + Di γMHi [1− exp(−2Hit γM )]. (30)
We respectively introduce the notations of Adam-style Momentum and SGD-style Momentum, and apply the second order Taylor expansion in case of small − 2HitγM . Then we obtain
〈∆θ2i 〉 = |Hi|η2
2(1− β1)B [1− exp (−(1− β1)T )]2 + |Hi|η2T B +O(B−1H2i η3T 2) (31)
for Adam-style Momentum, and
〈∆θ2i 〉 = |Hi|η2
2(1− β1)3B [1− exp (−(1− β1)T )]2 + |Hi|η2T B(1− β1)2 +O(B−1H2i η3T 2) (32)
for SGD-style Momentum.

heading: A.3 PROOF OF THEOREM 3
text: 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, we may transform the proof for Theorem 3.2 of Xie et al. (2021b) into the proof for Theorem 3 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 Fokker-Planck Equation (8) as
∂P (θ, r, t)
∂t =− r · ∇θP (θ, r, t) +∇θL(θ) ·M−1∇rP (θ, r, t)
+∇r ·M−2D(θ) · Peq(r)∇r[Peq(r)−1P (θ, r, t)] (33) =−∇ · J(θ, t). (34)
We will take similar forms of the proof in Xie et al. (2021b) as follows. The mean escape time is written as
τ = P (θ ∈ Va)∫ Sa J · dS . (35)
To compute the mean escape time, we decompose the proof into two steps: 1) compute the probability of locating in valley a, P (θ ∈ Va), and 2) compute the probability flux j = ∫ Sa J · dS. The definition of the probability flux integral may refer to Gauss Divergence Theorem.
Step 1:
Fortunately, the stationary probability distribution inside Valley a in Momentum dynamics is also given by a Gaussian distribution in Theorem 1 as
σ2i (t) = Di
γMHi . (36)
Under Quasi-Equilibrium Assumption, the distribution around minimum a is P (θ) = P (a) exp [ −γM2 (θ − a) >(D − 12 a HaD − 12 a )(θ − a) ] . We use the T notation as the temperature parameter in the stationary distribution, and use the D notation as the diffusion coefficient in the dynamics, for their different roles.
P (θ ∈ Va) (37)
= ∫ θ∈Va P (θ)dV (38)
=P (a) ∫ θ∈Va exp [ −γM 2 (θ − a)>(D− 1 2 a HaD − 12 a )(θ − a) ] dV (39)
≈P (a) ∫ θ∈(−∞,+∞) exp [ −γM 2 (θ − a)>(D− 1 2 a HaD − 12 a )(θ − a) ] dV (40)
=P (a) (2πγM)
n 2
det(D−1a Ha) 1 2
(41)
This result of P (θ ∈ Va) in Momentum only differs from SGD by the temperature correction γM . Step 2: We directly introduce the effective diffusion result from (Kalinay and Percus, 2012) into our analysis. Kalinay and Percus (2012) proved that the phase-space Fokker-Planck Equation (8) can be reduced to a space-dependent Smoluchowski-like equation, which is extended by an effective diffusion correction:
D̂i(θ) = Di(θ)
( 1− √ 1− 4Hi(θ)
γ2M
)( 2Hi(θ)
γ2M
)−1 . (42)
As we only employ the Smoluchowski Equation along the escape direction, we use the onedimensional expression along the escape direction (an eigenvector direction) for simplicity. Without losing clarity, we use the commonly used T to denote the temperature in the proof.
In case of SGD (Xie et al., 2021b), we can obtain Smoluchowski Equation in position space:
J = D(θ) exp ( −L(θ) T ) ∇ [ exp ( L(θ) T ) P (θ) ] , (43)
where T = D. According to (Kalinay and Percus, 2012), in case of finite inertia, we can transform the phase-space equation into the position-space Smoluchowski-like form with the effective diffusion correction:
J = D̂(θ) exp ( −L(θ) T ) ∇ [ exp ( L(θ) T ) P (θ) ] , (44)
where T = DγM , and D̂ defined by Equation (42) replaces standard D.
We assume the point s is a midpoint on the most possible path between a and b, where L(s) = (1−s)L(a)+sL(b). The temperature Ta dominates the path a→ s, while temperature Tb dominates the path s→ b. So we have
∇ [ exp ( L(θ)− L(s)
T
) P (θ) ] = JD−1 exp ( L(θ)− L(s)
T
) . (45)
We integrate the equation from Valley a to the outside of Valley a along the most possible escape path
Left = ∫ c a ∂ ∂θ [exp ( L(θ)− L(s) T ) P (θ)]dθ (46)
= ∫ s a ∂ ∂θ [ exp ( L(θ)− L(s) Ta ) P (θ) ] dθ (47)
+ ∫ c s ∂ ∂θ [ exp ( L(θ)− L(s) Tb ) P (θ) ] dθ (48)
=[P (s)− exp ( L(a)− L(s)
Ta
) P (a)] + [0− P (s)] (49)
=− exp ( L(a)− L(s)
Ta
) P (a) (50)
Right =− J ∫ c a D−1 exp ( L(θ)− L(s) T ) dθ (51)
We move J to the outside of integral based on Gauss’s Divergence Theorem, because J is fixed on the escape path from one minimum to another. As there is no field source on the escape path,∫ V ∇ · J(θ)dV = 0 and ∇J(θ) = 0. Obviously, probability sources are all near minima in deep learning. So we obtain
J = exp
( L(a)−L(s)
Ta ) P (a)∫ c
a D̂−1 exp
( L(θ)−L(s)
T
) dθ . (52)
Near saddle points, we have∫ c a D̂−1 exp ( L(θ)− L(s) T ) dθ (53)
≈ ∫ c a D̂−1 exp [ L(b)− L(s) + 12 (θ − b) >Hb(θ − b) Tb ] dθ (54)
≈D̂−1b ∫ +∞ −∞ exp [ L(b)− L(s) + 12 (θ − b) >Hb(θ − b) Tb ] dθ (55)
=D̂−1b exp
( L(b)− L(s)
Tb
)√ 2πTb |Hb| . (56)
Besides the temperature correction T = DγM , this result of J in Momentum also differs from SGD by
the effective diffusion correction D̂bDb . The effective diffusion correction coefficient is given by
D̂i(θ) Di(θ) =
1− √
1− 4Hi(θ)γ2M 2Hi(θ) γ2M . (57)
Based on the formula of the one-dimensional probability current and flux, we obtain the highdimensional flux escaping through b:∫
Sb
J · dS (58)
=J1d ∫ Sb exp [ −γM 2 (θ − b)>[D− 1 2 b HbD − 12 b ] ⊥e(θ − b) ] dS (59)
=J1d (2πγM)
n−1 2 ( ∏ i 6=e(D −1 bi Hbi)) 1 2
(60)
=
exp ( L(a)−L(s)
Tae
) P (a) (2πγM) n−1 2
( ∏
i6=e(D −1 bi Hbi)) 1 2 D̂−1be exp ( L(b)−L(s)
Tbe )√ 2πTbe |Hbe|
(61)
where [·]⊥e indicates the directions perpendicular to the escape direction e. Based on the results of Step 1 and Step 2, we have
τ = P (θ ∈ Va)∫ Sb J · dS
(62)
=P (a) (2πγM)
n 2
det(D−1a Ha) 1 2
D̂−1be exp ( L(b)−L(s)
Tbe )√ 2πTbe |Hbe|
exp ( L(a)−L(s)
Tae
) P (a) (2πγM) n−1 2
( ∏
i6=e(D −1 bi Hbi)) 1 2
(63)
= 1
D̂be 2π Dbe |Hbe| exp
[ 2γMB∆L
η
( s
Hae + (1− s) |Hbe|
)] (64)
=π [√ 1 +
4|Hbe| γ2M + 1
] 1
|Hbe| exp
[ 2γMB∆L
η
( s
Hae + (1− s) |Hbe|
)] . (65)
We have replaced the eigenvalue of Hb along the escape direction by its absolute value.
Finally, by introducing γ and M , we obtain the log-scale expression as log(τ) = O (
2(1− β1)B∆L β3ηHae
) (66)

heading: A.4 PROOF OF PROPOSITION 2
text: Proof. The proof closely relates to the proof of 3. We only need replace the standard learning rate by the adaptive learning rate η̂ = ηv− 1 2 , and set γM = 1. Particularly,
Dadam = Dv − 12 =
η[H]+v− 1 2
2B =
1 2 η
√ [H]+
B . (67)
We introduce η̂ = ηv− 1 2 into the proof of 3, and obtain
τ = P (θ ∈ Va)∫ Sb J · dS
(68)
=
[ |det(D−1b V 1 2 b Hb)|
det(D−1a V 1 2 a Ha)
] 1 2
π [√ 1 +
4|Hbe| γ2M + 1
] 1
|Hbe| exp
[ 2γMB∆L
η
( s
V − 12 ae Hae
+ (1− s)
V − 12 be |Hbe| )] (69)
=π √1 + 4η√B|Hbe| 1− β1 + 1  |det(H−1a Hb)| 14 |Hbe| exp [ 2 √ B∆L η ( s√ Hae + (1− s)√ |Hbe| )] . (70)

heading: A.5 PROOF OF PROPOSITION 4
text: Proof. The proof closely relates to the proof of 3. We only need introduce the mass matrix M and the dampening matrix γ. Fortunately, γM = I . Thus we directly have the result from the proof of 3 as:
τ = P (θ ∈ Va)∫ Sb J · dS
(71)
=π [√ 1 +
4|Hbe| γ2M + 1
] 1
|Hbe| exp
[ 2γMB∆L
η
( s
Hae + (1− s) |Hbe|
)] . (72)
As M = η(I − β1)−1, γ = η−1(I − β1), and I − β1 = β0v̄ v, we obtain the result:
τ = π √1 + 4η∑ni=1 |Hbi| β0n + 1  1 |Hbe| exp [ 2B∆L η ( s Hae + (1− s) |Hbe| )] (73)

heading: A.6 PROOF OF THEOREM 4
text: Without loss of generality, we assume the dimensionality is one and rewrite the main updating rule of Adai as
θt+1 = θt − η(1− β1,t)gt + β1,t(θt − θt−1). (74)
In the convergence proof, we do not need to specify how to update β1,t but just let β1,t ∈ [0, 1). We denote that θ−1 = θ0.
Before presenting the main proof, we first prove four useful lemmas.
Lemma 1. Under the conditions of Theorem 4, for any t ≥ 0, we have
xt+1 − xt = −η t∑
k=0
qk,tgt,
where
qk,t = (1− βk) t∏
i=k+1
β1,i.
Then we have
1− βt+11,max ≤ t∑
k=0
qk,t ≤ 1.
Proof. Recall that
θt+1 =θt − η(1− β1,t)gt + β1,t(θt − θt−1) (75) θt+1 − θt =β1,t(θt − θt−1)− η(1− β1,t)gt. (76)
Then we have
θt+1 − θt = −η t∑
k=0
(1− β1,k)gk t∏
i=k+1
β1,i. (77)
Let qk,t = (1− βk) ∏t i=k+1 β1,i.
For analyzing the maximum and the minimum, we calculate the derivatives with respect to β1,k for any 0 ≤ k ≤ t:
∂ ∑t k=0 qk,t ∂β1,0 = − t∏
i=k+1
β1,i ≤ 0. (78)
Note that 0 ≤ β1,k ≤ β1,max. Then we have
t∑ k=0 qk,t|β1,0=β1,max ≤ t∑ k=0 qk,t ≤ t∑ k=0 qk,t|β1,0=0. (79)
Recursively, we can calculate the derivatives with respect to β1,1, β1,2, . . . , β1,t. Then we obtain max( ∑t k=0 qk,t) = 1 by letting β1,k = 0 for all k.
Similarly, we obtain min( ∑t k=0 qk,t) = 1 by letting β1,t = β1,max for all k.
Then we obtain
1− βt+11,max ≤ t∑
k=0
qk,t ≤ 1 (80)
The proof is now complete.
Lemma 2. Under the conditions of Theorem 4, for any t ≥ 0, we have
E[L(θt+1)− L(θt)] ≤ 1
2 t∑ k=0 qk,tE[ t−1∑ j=k ‖∇L(θj+1)− L(θj)‖2]+ t∑ k=0 qk,t( t− k 2 η2 − η)E[‖∇L(θk)‖2] + Lη2 2 ‖ t∑ k=0 qk,tgk‖2.
Proof. As L(θ) is L-smooth, we have L(θt+1)− L(θt) (81)
≤∇L(θt)>(θt+1 − θt) + L 2 ‖θt+1 − θt‖2 (82)
=− η t∑
k=0
qk,t∇L(θt)>gk + Lη2
2 ‖ t∑ k=0 qk,tgk‖2 (83)
=− η t∑
k=0
qk,t t−1∑ j=k (∇L(θj+1)−∇L(θj))>(∇L(θk) + ξk)−
η t∑ k=0 qk,t∇L(θk)>(∇L(θk) + ξk) + Lη2 2 ‖ t∑ k=0 qk,tgk‖2. (84)
Taking expectation on both sides gives
E[L(θt+1)− L(θt)] (85)
≤ 1 2L t∑ k=0 qk,tE[ t−1∑ j=k (∇L(θj+1)−∇L(θj))>∇L(θk)]− (86)
t∑ k=0 qk,tηE[‖∇L(θk)‖2] + Lη2 2 ‖ t∑ k=0 qk,tgk‖2. (87)
By the Cauchy-Schwarz Inequality, we have
E[L(θt+1)− L(θt)] (88)
≤ t∑
k=0
qk,tE[ 1
2 ‖∇L(θt)−∇L(θk‖2 +
η2
2 ‖∇L(θk)‖2]−
η t∑ k=0 qk,tE[‖∇L(θk)‖2] + Lη2 2 E[‖ t∑ k=0 qk,tgk‖2] (89)
= 1
2 t∑ k=0 qk,tE[ t−1∑ j=k ‖∇L(θj+1)− L(θj)‖2]+ t∑ k=0 qk,t( t− k 2 η2 − η)E[‖∇L(θk)‖2] + Lη2 2 ‖ t∑ k=0 qk,tgk‖2. (90)
The proof is now complete.
Lemma 3. Under the conditions of Theorem 4, for any t ≥ 0, we have t∑
k=0
qk,t(t− k) ≤ β1,max
1− β1,max
Proof. For analyzing the maximum, we calculate the derivatives with respect to β1,k for any 0 ≤ k ≤ t. With respect to β1,1, we have
∂ ∑t k=0 qk,t(t− k) ∂β1,0 = − t∏
k=1
β1,kt ≤ 0. (91)
Then we let β1,0 = 0 and obtain the derivative with respect to β1,1 as ∂ ∑t k=0 qk,t(t− k) ∂β1,1 = t∏ k=2 β1,k(t− (t− 1)) ≥ 0. (92)
Then we let β1,j = β1,max for 1 ≤ j ≤ k − 1 and obtain the derivative with respect to β1,k as ∂ ∑t k=0 qk,t(t− k) ∂β1,k ≥ (t− k) ∂ ∑t k=0 qk,t ∂β1,k = 0. (93)
Thus, we have β1,0 = 0 and β1,k = β1,max for 1 ≤ j ≤ t to maximize ∑t−1 k=0 qk,t(t− k). We may write it as t∑
k=0
qk,t(t− k) ≤ βt+11,maxt+ t∑
k=0
(1− β1,max)βk1,maxk (94)
Note that (1− β1,max)βk1,maxk is an arithmetico-geometric sequence. Thus, we have
max( t∑ k=0 qk,t(t− k)) ≤βt+11,maxt+ [ β1,max − βt+11,maxt+ β21,max(1− βt−11,max) 1− β1,max ] (95)
≤β1,max + β21,max
1− β1,max (96)
= β1,max
1− β1,max (97)
The proof is now complete.
Lemma 4. Under the conditions of Theorem 4, for any t ≥ 0, we have
1
2 t∑ k=0 qk,tE[ t−1∑ j=k ‖∇L(θj+1)− L(θj)‖2] + Lη2 2 ‖ t∑ k=0 qk,tgk‖2 ≤ Lη2 2(1− β1,max) (G2 + δ2).
Proof. As L(θ) is L-smooth, we have
‖∇L(θj+1)− L(θj)‖2 (98) ≤L‖θj+1 − θj‖2 (99) =L‖ − η j∑ i=0 qi,jgi‖2. (100)
By ∑j i=0 qi,j ≤ 1 in Lemma 1, we have
E[‖∇L(θj+1)− L(θj)‖2] (101) ≤Lη2(G2 + δ2). (102)
and
Lη2
2 ‖ t∑ k=0 qk,tgk‖2 (103)
≤Lη 2
2 (G2 + δ2). (104)
By Lemma 3, we have
1
2 t∑ k=0 qk,tE[ t−1∑ j=k ‖∇L(θj+1)− L(θj)‖2] (105)
≤η 2
2 (G2 + δ2) t∑ k=0 qk,t(t− k) (106)
≤ η 2β1,max
2(1− β1,max) (G2 + δ2). (107)
Then we obtain
1
2 t∑ k=0 qk,tE[ t−1∑ j=k ‖∇L(θj+1)− L(θj)‖2] + Lη2 2 ‖ t∑ k=0 qk,tgk‖2 (108)
≤ Lη 2β1,max
2(1− β1,max) (G2 + δ2) +
Lη2
2 (G2 + δ2) (109)
≤ Lη 2
2(1− β1,max) (G2 + δ2). (110)
Proof. The proof of Theorem 4 is organized as follows.
By Lemma 2 and Lemma 4, we have
E[L(θt+1)− L(θt)]
≤1 2 t∑ k=0 qk,tE[ t−1∑ j=k ‖∇L(θj+1)− L(θj)‖2]+
t∑ k=0 qk,t( t− k 2 η2 − η)E[‖∇L(θk)‖2] + Lη2 2 ‖ t∑ k=0 qk,tgk‖2 (111)
≤ t∑
k=0
qk,t( t− k
2 η2 − η)E[‖∇L(θk)‖2] +
Lη2
2(1− β1,max) (G2 + δ2). (112)
Thus, we have
t∑ k=0 qk,tηE[‖∇L(θk)‖2]
≤E[L(θt)− L(θt+1)] + t∑
k=0
qk,t t− k
2 η2E[‖∇L(θk)‖2] +
Lη2
2(1− β1,max) (G2 + δ2). (113)
By Lemma 1 and Lemma 3, we have
(1− βt+11,max)η min k=0,...,t E[‖∇L(θk)‖2]
≤E[L(θt)− L(θt+1)] + β1,maxη
2
2(1− β1,max) G2 +
Lη2
2(1− β1,max) (G2 + δ2). (114)
By summing the above inequality for t = 0, . . . , t, we have
(t+ 1)(1− β1,max)η min k=0,...,t E[‖∇L(θk)‖2]
≤L(θ0)− L? + β1,maxη
2
2(1− β1,max) G2(t+ 1) +
Lη2
2(1− β1,max) (G2 + δ2)(t+ 1). (115)
Then
min k=0,...,t
E[‖∇L(θk)‖2]
≤ L(θ0)− L ?
(1− β1,max)(t+ 1)η +
β1,maxη
2(1− β1,max)2 G2 + Lη 2(1− β1,max)2 (G2 + δ2). (116)
Let η ≤ C√ t+1 . We have
min k=0,...,t
E[‖∇L(θk)‖2]
≤ L(θ0)− L ?
(1− β1,max)C √ t+ 1
+ β1,maxC
2(1− β1,max)2 √ t+ 1 G2+
LC 2(1− β1,max)2 √ t+ 1 (G2 + δ2) (117)
(118) The proof is now complete.

heading: B CLASSICAL APPROXIMATION ASSUMPTIONS
text: Assumption 2 indicates that the dynamical system is in equilibrium near minima but not necessarily near saddle points. It means that ∂P (θ,t)∂t ≈ 0 holds near minima, but not necessarily holds near saddle point b. Quasi-Equilibrium Assumption is actually weaker but more useful than the conventional 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-Equilibrium 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 can be obtained under the stationary assumption.
Low Temperature Assumption is common (Van Kampen, 1992; Zhou, 2010; Berglund, 2013; Jastrzkebski et al., 2017), and is always justified when ηB is small. Under Assumption 3, the probability densities will concentrate around minima and MPPs. Numerically, the 6-sigma rule may often provide good approximation for a Gaussian distribution. Assumption 3 will make the second order Taylor approximation, Assumption 1, even more reasonable in SGD diffusion.
Here, we try to provide a more intuitive explanation about Low Temperature Assumption in the domain of deep learning. Without loss of generality, we discuss it in one-dimensional dynamics. The temperature can be interpreted as a real number D. In SGD, we have the temperature as D = η2BH . In statistical physics, if ∆LD is large, then we call it Low Temperature Approximation. Note that ∆L D appears insides an exponential function in the theoretical analysis. People usually believe that, numerically, ∆LD > 6 can make a good approximation, for a similar reason of the 6-sigma rule in statistics. In the final training phase of deep networks, a common setting is η = 0.01 and B = 128. Thus, we may safely apply Assumption 3 to the loss valleys which satisfy the very mild condition ∆L H > 2.3× 10 −4. Empirically, the condition ∆LH > 2.3× 10 −4 holds well in SGD dynamics.

heading: C STOCHASTIC GRADIENT NOISE ANALYSIS
text: 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. (2021b), we again validated the relation between gradient noise covariance and the Hessian. We particularly choose a randomly initialized mode so that the model is not near critical points. We display all elements H(i,j) ∈ [1e− 4, 0.5] of the Hessian matrix and the corresponding elements C(i,j) of gradient covariance matrix in 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-magnitude eigenvalues of the Hessian. Fortunately, small-magnitude eigenvalues of 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, ∂L(θt)∂θt − ∂L̂(θt) ∂θt , which is the difference between gradient descent and stochastic gradient descent.
In Figure 7 of the Appendix, Xie et al. (2021b) empirically verified that SGN is highly similar to Gaussian noise instead of heavy-tailed Lévy noise. Xie et al. (2021b) recovered the experiment of
Simsekli et al. (2019) to show that gradient noise is approximately Lévy noise only if it is computed across parameters. Figure 7 of the Appendix actually suggests that the contradicted observations are from the different formulations of gradient noise. Simsekli et al. (2019) computed “SGN” across n model parameters and regarded “SGN” as n samples drawn from a single-variant distribution. In Xie et al. (2021b), SGN computed over different minibatches obeys a N -variant Gaussian distribution, which can be θ-dependent and anisotropic. Simsekli et al. (2019) studied the distribution of SGN as a single-variant distribution, while Xie et al. (2021b) relaxed it as a n-variant distribution. Figure 7 holds well at least when the batch size B is larger than 16, which is common in practice.

heading: D EXPERIMENTAL DETAILS
text: Computational environment. The experiments are conducted on a computing cluster with GPUs of NVIDIA R© Tesla TM P100 16GB and CPUs of Intel R© Xeon R© CPU E5-2640 v3 @ 2.60GHz.
D.1 IMAGE CLASSIFICATION
Data Preprocessing: For CIFAR-10/CIFAR-100, we perform the per-pixel zero-mean unit-variance normalization, horizontal random flip, and 32× 32 random crops after padding with 4 pixels on each side. For ImageNet, we perform the per-pixel zero-mean unit-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 34 to 4 3 ) of the original aspect ratio is made.
Hyperparameter Settings for CIFAR-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
learning rates: η = 1 for Adai;η = 0.1 for SGD with Momentum and AdaiW; η = 0.001 for Adam, AMSGrad, AdamW, AdaBound, Yogi, and RAdam; η = 0.01 for Padam. For the learning rate schedule, the learning rate is divided by 10 at the epoch of {80, 160} for CIFAR-10 and {100, 150} for CIFAR-100. The batch size is set to 128 for CIFAR-10 and CIFAR-100. The L2 regularization hyperparameter is set to λ = 0.0005 for CIFAR-10 and CIFAR-100. Considering the linear scaling rule of decoupled weight decay and initial learning rates (Loshchilov and Hutter, 2018), we chose decoupled weight decay hyperparameters as: λ = 0.5 for AdamW on CIFAR-10 and CIFAR-100; λ = 0.005 for AdaiW on CIFAR-10 and CIFAR-100. We set the momentum hyperparameter β1 = 0.9 for SGD with Momentum. As for other optimizer hyperparameters, we apply the default hyperparameter settings directly.
Hyperparameter Settings for ImageNet: We select the optimal learning rate for each experiment from {0.00001, 0.0001, 0.001, 0.01, 0.1, 1, 10} for Adai, SGD with Momentum, and Adam. The settings of learning rates: η = 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, 90}. The batch size is set to 256. The L2 regularization hyperparameter is set to λ = 0.0001. We set the momentum hyperparameter β1 = 0.9 for SGD with Momentum. As for other training hyperparameters, we apply the default hyperparameter settings directly.
Some papers often choose λ = 0.0001 as the default weight decay setting for CIFAR-10 and CIFAR-100. We study the weight decay setting in Appendix F.

heading: D.2 LANGUAGE MODELING
text: 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 a 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 is similar to “medium LSTM” in Zaremba et al. (2014). The benchmark task is the word-level Penn TreeBank (Marcus et 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 1. We select optimal learning rates from {10, 1, 0.1, 0.01, 0.001, 0.0001, 0.00001} for each optimizer.

heading: E THE MEAN ESCAPE TIME ANALYSIS
text: Mean Escape Time: The mean escape time is the expected time for a particle governed by Equation 1 to escape from Sharp Valley a1 to Flat Valley a2, seen in Figure 8. The mean escape time is widely used in related statistical physics and stochastic process (Van Kampen, 1992; Nguyen et al., 2019).
Related machine learning papers (Xie et al., 2021b; Hu et al., 2019; Nguyen et 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 x ∼ N (0, 4I). 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 in nonconvex optimization, written as
f(θ) = 1
2 N∑ i=1 (θ4i − 16θ2i + 5θi).
We use 10-dimensional Styblinski-Tang Function as the test function, and Gaussian samples as training data.
L(θ) = f(θ − x), where data samples x ∼ N (0, 4I). The one-dimensional Styblinski-Tang Function has one global minimum located at a = −2.903534, one local minimum located at d, and one saddle point b = 0.156731 as the boundary separating Valley a1 and Valley a2. For a n-dimensional Styblinski-Tang Function, we initialize parameters θt=0 = 1√k (−2.903534, . . . ,−2.903534), and set the valley’s boundary as θi < 1√k0.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 we are more interested in the number of iterations than “dynamical time” in practice, we use the number of iterations to denote the mean escape time and ignore the time unit η in “dynamical time”. We repeat experiments 100 times to estimate the escape rate Γ and the mean escape time τ . As the escape time is approximately a random variable obeying an exponential distribution, t ∼ Exponential(Γ), the estimated escape rate can be written as
Γ = 100− 2∑100 i=1 ti . (119)
The 95% confidence interval of this estimator is
Γ(1− 1.96√ 100 ) ≤ Γ ≤ Γ(1 + 1.96√ 100 ). (120)

heading: F SUPPLEMENTARY EMPIRICAL RESULTS
text: Some papers (Luo et 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 1. The main problem may lie in weight decay. SGD with weight decay λ = 0.0001, a common setting in related papers, is not a good baseline on CIFAR-10 and CIFAR-100, as λ = 0.0005 often shows better generalization, seen in Figures 9. We also conduct comparative experiments with λ = 0.0001, seen in Table 2. While some Adam variants under this setting sometimes may compare with SGD due to 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 and CIFAR-100 in Figure 10. We further compare convergence of Adai and Adam with the fixed learning rate scheduler for 1000 epochs in Figure 11.
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. (2017) to compare the sharpness of minima learned by Adai, Momentum, and Adam. The expected minima sharpness is defined as Eζ [L(θ? + ζ)− L(θ?)], where ζ is Gaussian noise and θ? is the empirical minimizer learned by a training algorithm. If the loss landscape near θ? is sharp, the weight-perturbed loss Eζ [L(θ? + ζ)] will be much larger than L(θ?). Figure 13 empirically supports that Adai and Momentum can learn significantly flatter minima than Adam.
0 25 50 75 100 125 150 175 200 Epochs
0.04
0.06
0.08
0.10
0.12
0.14
Te st
E rro
r
Adai LR1.0 Momentum Adai LR0.5 Adam
(a) ResNet18
0 25 50 75 100 125 150 175 200 Epochs
0.04
0.06
0.08
0.10
0.12
0.14
Te st
E rro
r
Adai LR1.0 Momentum Adai LR0.5 Adam
(b) VGG16
0 25 50 75 100 125 150 175 200 Epochs
0.175
0.200
0.225
0.250
0.275
0.300
0.325
0.350
0.375
Te st
E rro
r
Adai LR1.0 Momentum Adai LR0.5 Adam
(c) DenseNet121
0 25 50 75 100 125 150 175 200 Epochs
0.175
0.200
0.225
0.250
0.275
0.300
0.325
0.350
0.375
Te st
E rro
r
Adai LR1.0 Momentum Adai LR0.5 Adam
(d) GoogLeNet

heading: DATASET MODEL ADAIW ADAI SGD M SGD M ADAM AMSGRAD ADAMW ADABOUND PADAM YOGI RADAM

heading: G ADAI WITH STABLE/DECOUPLED WEIGHT DECAY
text: Algorithm 3: AdaiS/AdaiW
gt = ∇L̂(θt−1); vt = β2vt−1 + (1− β2)g2t ; v̂t =
vt 1−βt2 ; v̄t = mean(v̂t); β1,t = (1− β0v̄t v̂t).Clip(0, 1− ); mt = β1,tmt−1 + (1− β1,t)gt; m̂t =
mt 1− ∏t z=1 β1,z
; θt = θt−1 − ηm̂t − ληθt−1;

heading: H EXPRESSIONS OF ADAM DYNAMICS
text: In this section, we 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 η̂ in Equations (4), (5), and (6). We write the updating rule of Adam as{
mt = β1mt−1 + β3gt, θt+1 = θt − η̂mt, (121)
where β3 = 1− β1, β1 is the hyperparameter, and η̂ is the adaptive learning rate. For simplicity, it is fine to 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 γ as{
rt = (1− γdt)rt−1 + FM dt θt+1 = θt + rtdt,
(122)
where rt = −mt, F = gt, dt = η̂, 1 − γdt = β1, and dtM = β3 = 1 − β1. Thus, we obtain the differential-form motion equation of Adam as
Mθ̈ = −γMθ̇ + F, (123)
where the mass M = η̂1−β1 and the damping coefficient γ = 1−β1 η̂ .
This shows that the analysis can be 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 a differential-form 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. (2020), as long as we let the expressions of M and γ follow Adam dynamics. In fact, the updating
rule of vt 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-time dynamics for training of deep neural networks. If we let β1 = 0 in Adam, Adam will be reduced to RMSprop. Similarly SGD, RMSprop has no the momentum drift effect on saddle-point escaping but has an 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-point escaping as we want.

Review summaries:
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.

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.

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.  

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   - QUOTE momentum matters little to escaping sharp minima UNQUOTE - QUOTE adaptive learning rate is not good at selecting flat minima UNQUOTE