A DAPTIVE I NERTIA : D ISENTANGLING THE E FFECTS OF A DAPTIVE L EARNING R ATE AND M OMENTUM
Anonymous authors
Paper under double-blind review
A BSTRACT
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.
Speciﬁcally, we disentangle the effects of Adaptive Learning Rate and Momentum of the Adam dynamics on saddlepoint escaping and ﬂat minima selection.
We prove that Adaptive Learning Rate can escape saddle points efﬁciently, but cannot select ﬂat 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 ﬂat 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 ﬂat minima as well as SGD.
Our extensive experiments demonstrate that the proposed adaptive inertia method can generalize signiﬁcantly better than SGD and conventional adaptive gradient methods.
I NTRODUCTION
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 ﬁnds sharper minima than SGD (Wilson et al., 2017) for popular convolutional neural networks, where the ﬂat 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 ﬁnd ﬂatter 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 reﬂect 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 efﬁciently 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 ﬂat 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 ﬂat minima, while Momentum helps escape saddle point and matters little to escaping sharp minima 1 . 2)
We propose a novel Adaptive Inertia (Adai) optimization framework
[Table]
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 ﬂat minima well.
From the empirical results (see Table 1), Adai signiﬁcantly outperforms SGD and popular
Adam variants.
The paper is organized as follows.
In Section 2, we ﬁrst 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.
Sectionconcludes the paper with remarks.
SGD AND D IFFUSION
In this section, we introduce some preliminaries about the SGD diffusion and then present the saddle-point escaping property of SGD dynamics.
2. P REREQUISITES FOR SGD D IFFUSION
We ﬁrst 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 as B , 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
[Equation]
where dW t ∼ N (0 , Idt ) , I is the identity matrix, and C ( θ ) is the gradient noise covariance matrix.
The gradient noise is deﬁned 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
[Equation]
where ∇· is the divergence operator and D ( θ )
=
ηC ( θ )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
[Equation]
near a critical point c , where L j ( θ ) 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.
We further veriﬁed 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.
and veriﬁed by Xie et al.
(2021b)
and Daneshmand et al.
Please refer to Appendix C for the detailed analysis of the stochastic gradient noise.
If we have eigendecomposition H
=
U diag( H 1 , . . .
, H n − 1 , H n ) U (cid:62) , then we use [ · ] + to denote the transformation that [ H ] + = U diag( | H 1 | , . . .
, | H n − 1 | , | H n | ) U (cid:62) .
The i -th column vector of U is the eigenvector of H corresponding to the eigenvalue H i .
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
[Equation]
2. SGD D IFFUSION NEAR S ADDLE P OINTS
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 , . . .
, σ 2 n − 1 , σ 2 n ) U (cid:62) and σ
[Equation]
where D i is the i -th eigenvalue of the diffusion matrix D and H i 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 | H i | ηT (cid:28) 1 near saddle points, we have
[Equation]
The proof is relegated to Appendix A.1.
We note that 1) if H i > 0 , the distribution of θ along the direction i converges to a Gaussian distribution with constant variance, N (cid:0) c i , η 2 B (cid:1) ; and 2) if H i < 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-ﬁxed Gaussian distribution, the mean squared displacement is equivalent to the variance, namely (cid:104) ∆ θ 2 i ( t ) (cid:105) =
σ 2 i ( t ) .
The result means that SGD escapes saddle points very slowly ( (cid:104) ∆ θ 2 i (cid:105)
=
O ( | H i | ) )
if H i 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.
A NALYSIS OF M OMENTUM D YNAMICS
In this section, we analyze Momentum in terms of saddle-point escaping and ﬂat minima selection.
The continuous-time Momentum dynamics.
The Heavy Ball Method (Zavriev and Kostyuk, 1993) can be written as
[Equation]
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 coefﬁcient γ as
[Equation]
where r t = − m t , F = g t , dt = η , 1 − γdt =
β 1 , and dtM = β 3 .
Thus, we obtain the differential form of the motion equation as
[Equation]
where ¨ θ =
d 2 θ dt 2 and ˙ θ = dθdt .
The left-hand side as the inertia force is equal to the mass M times the acceleration ¨ θ , the ﬁrst term in the right-hand side as the damping force is equal to the damping coefﬁcient γ 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 coefﬁcient γ = 1 − β 1 η by comparing (5) and (6).
As F corresponds to the stochastic gradient term, we obtain
[Equation]
Its Fokker-Planck Equation in the phase space (the θ - ˙ θ space) is well known as
[Equation]
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 ﬁnite-inertia Langevin Dynamics, Equation (7).
We contribution is the ﬁrst 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
[Equation]
where ∆ θ ( t ) = θ ( t ) − θ (0) is the displacement of θ , and (cid:104)·(cid:105) denote the mean value.
In the right-hand side, this ﬁrst term is the momentum drift effect and the second term is the diffusion effect.
As | H i | ηT (cid:28) 1 near ill-conditioned saddle points, it can be written in terms of Momentum notations as
[Equation]
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 ﬁrst 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 ﬁrst 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 ﬁelds’ 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
[Equation]
where the subscript e indicates the escape direction, s ∈ (0 , 1) is a path-dependent parameter, and
H ae and H be 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 |
[Equation]
1 , it reduces to SGD.
In terms of Momentum notations, we have log( τ ) =
O (cid:16)
[Equation] [Equation]
(cid:17) .
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 (cid:16) 2 B ∆ L ηH ae (cid:17) has been obtained for SGD (Xie et al., 2021b).
We see that Momentum does not affect ﬂat minima selection in terms of the mean escape time, if we properly choose the learning rate, i.e., η Momentum = 1 − β 1 β 3 η SGD .
A NALYSIS OF A DAM D YNAMICS
Algorithm 1: Adam
[Equation]
Algorithm 2: Adai
[Equation]
In this section, we analyze the effects of Adaptive Learning Rate in terms of saddle-point escaping and ﬂat 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 ∝ | H i | − 12 .
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 [ g t g (cid:62) t ] (Staib et al., 2019).
Following Staib et al.
(2019), our analysis considers the “idealized” Adam where v = E [ g t g (cid:62) t ] is a full matrix.
Note that E [ g t g (cid:62) 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 coefﬁcient γ = 1 − β 1 ˆ η .
We emphasize that the learning rate is not a real number but the diagonal approximation of the ideal learning rate matrix ˆ η = ηC − 12 in practice.
We apply the adaptive time continuation d ˆ t i
=
ˆ η i in the i -th dimension, where ˆ t i of T iterations is deﬁned 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 − 12 .
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
[Equation]
Under | H i | ηT (cid:28) 1 , it can be written in terms of Adam notations as
[Equation]
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.
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 ﬂat 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:
[Equation]
, where τ is the mean escape time.
Thus, we have log( τ
[Equation]
= O (cid:16) 2 √ B
[Equation]
L η √ H ae (cid:17) .
In comparison, Adam has log( τ ) =
O ( H − 12 ae ) , while SGD and Momentum both have log( τ ) =
O ( H − 1 ae ) .
We say that Adam is weaker Hessian-dependent than SGD in terms of sharp minima escaping, which means that if H ae 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 reﬂected by H ae , the eigenvalue of the Hessian corresponding to the escape direction along an eigenvector.
The weaker Hessian-dependent diffusion term of Adam hurts the efﬁciency of selecting ﬂat minima in comparison with SGD.
In summary, Adam is better at saddle-point escaping but worse at ﬂat 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 ﬁne-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 ﬂat minima due to Adaptive
Learning Rate.
A DAPTIVE I NERTIA
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 ﬂat minima selection.
The total momentum drift effect during passing through a saddle point in Theorem 2 is given by
[Equation]
We generalize the scalar β 1 in Momentum to a vector β 1 in Adai.
The ideal adaptive updating is to keep β 1 = 1 − β 0 ¯ v 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 (cid:15) is for avoiding extremely large inertia.
The default setting (cid:15) = 0 .
001 means the maximal inertia is 1000 times the minimal inertia, as
M max = η (1 − β 1 , max ) − 1 = η(cid:15) − 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 coefﬁcient and the damping coefﬁcient 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 efﬁciently 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
[Equation]
where (cid:80) ni
=
| H i | is the trace norm of the Hessian at c .
Intuitively, the momentum drift effect of Adai can be signiﬁcantly larger than Adam by allowing larger inertia, which can be veriﬁed 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 ﬂat 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:
[Equation]
where τ is the mean escape time.
Thus, we have log( τ ) =
O (cid:16) 2 B [Equation] L ηH [Equation]
(cid:17) .
We can see that both SGD and Adai have log( τ ) =
O ( H − 1 ae ) , while Adam has log( τ ) =
O ( H − 12 ae ) .
Thus, Adai favors ﬂat 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 function 2 , L is lower bounded as L ( θ ) ≥ L (cid:63) , E [ ξ ] = 0 , E [ (cid:107) g
( θ, ξ ) − ∇ L ( θ ) (cid:107) 2 ] ≤ δ 2 , and (cid:107)∇ L ( θ ) (cid:107) ≤ G for any θ , where ξ represents the gradient noise of sub-sampling.
Let Adai run for t + 1 iterations and β 1 , max = 1 − (cid:15) ∈ [0 , 1) for any t ≥ 0 .
If η ≤
C √ t +1 , we have min k =0 ,...,t E
[ (cid:107)∇
[Equation]
where C 1 =
L ( θ 0 ) − L (cid:63) (1 − β 1 , max ) C , C 2 = β 1 , max C 2(1 − β 1 , max ) 2 G 2 , and C 3 = L C 2(1 − β 1 , max ) 2 ( G 2 + δ 2 ) .
θ
∇ (cid:107)
[Figure] [Figure]
[Figure]
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.
E MPIRICAL A NALYSIS
In this section, we ﬁrst conduct experiments to compare Adai, Adam, and SGD with Momentum in terms of convergence speed and generalization, and then empirically analyze ﬂat 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.
[Figure]
r n i n g
[Figure]
a t e
[Figure]
Adai has a much deeper and wider blue region near dark points ( ≤ 4 . 83% ) than SGD with Momentum and Adam.
[Figure]
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 Classiﬁcation on ImageNet.
Figure 3 shows that Adai generalizes signiﬁcantly 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 ( θ )
= ∇ 2 f ( θ ) → H ( θ )
=
k ∇ 2 f ( θ ) .
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 − 1 ae ) in Adai and SGD, while we have log( τ ) =
O ( H − 12 ae ) in Adam.
Moreover, Adam is signiﬁcantly 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.
and empirically veriﬁed again that Adai and Momentum can learn ﬂatter minima than Adam.
C ONCLUSION
To the best of our knowledge, we are the ﬁrst to theoretically disentangle the effects of Adaptive
Learning Rate and Momentum in terms of saddle-point escaping and ﬂat minima selection.
Under reasonable assumptions, our theory explains why Adam is good at escape saddle points but not good at selecting ﬂat 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 ﬂat minima well at the same time.
Our empirical analysis demonstrates that Adai generalizes signiﬁcantly better than popular Adam variants and SGD.
A P ROOFS A. P ROOF OF T HEOREM
Proof.
It is easy to validate that the probability density function
[Equation]
is the solution of the Fokker-Planck Equation (2).
Without losing generality, we only validate one-dimensional solution, such as Dimension
i
The ﬁrst term in Equation (2) can be written as
[Equation]
The second term in Equation (2) can be written as
[Equation]
The third term in Equation (2) can be written as
[Equation]
By Term 1 = Term 2 + Term 3 , we have
[Equation]
[Equation]
The initial condition of σ 2 is given by σ 2 (0)
= 0 .
We can validate that σ 2 satisﬁes
[Equation]
It is true along all eigenvectors’ directions.
By D = η 2 B H , we can get the results of SGD diffusion:
[Equation]
A. P ROOF OF T HEOREM
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
[Equation]
(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
[Equation]
and
[Equation]
) respectively, where Z r is the partition function for P eq ( r ) and Z θ is the partition function for P eq ( θ ) .
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 P eq ( 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
[Equation]
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 r 2 eq,i =
D i γM 2 .
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
[Equation]
The time-dependent components in P ( θ, t ) are caused by two effects.
The ﬁrst 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
[Equation]
We also have ignored the conservative force − H [ c ( t ) − c (0)] near saddle points, as | H i | (cid:28) γ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) =
r eq , and ¨ c (0) = − γr eq .
Then we can obtain the solution c ( t ) as
[Equation]
[Equation]
Step II (b): the diffusion effect.
We can get the dynamics of the diffusion effect as
[Equation]
This is equivalent to SGD dynamics with ˆ η = ηγM .
The expression of P ( θ, t ) and σ 2 i ( t ) is directly given by Theorem 1 as
[Equation]
[Equation]
We combine the momentum drift effect and the diffusion effect together, and then obtain the mean squared displacement of θ as
[Equation]
We respectively introduce the notations of Adam-style Momentum and SGD-style Momentum, and apply the second order Taylor expansion in case of small − 2 H i t γM .
Then we obtain
[Equation]
for Adam-style Momentum, and (cid:104)
[Equation]
for SGD-style Momentum.
A. P ROOF OF T HEOREM
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 ﬁrst 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 ﬂux respectively.
According to Gauss Divergence Theorem, we may rewrite the Fokker-Planck Equation (8) as
[Equation]
We will take similar forms of the proof in Xie et al.
(2021b) as follows.
The mean escape time is written as
[Equation]
To compute the mean escape time, we decompose the proof into two steps: 1) compute the probability of locating in valley a, P ( θ ∈ V a ) , and 2) compute the probability ﬂux j = (cid:82) S a J · dS .
The deﬁnition of the probability ﬂux 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
[Equation]
Under Quasi-Equilibrium Assumption, the distribution around minimum a is P ( θ )
[Equation]
We use the T notation as the temperature parameter in the stationary distribution, and use the D notation as the diffusion coefﬁcient in the dynamics, for their different roles. [Equation]
[Equation]
This result of P ( θ ∈ V a ) 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:
[Equation]
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:
[Equation]
where T = D .
According to (Kalinay and Percus, 2012), in case of ﬁnite inertia, we can transform the phase-space equation into the position-space Smoluchowski-like form with the effective diffusion correction:
[Equation]
where T = DγM , and ˆ D deﬁned 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 T a dominates the path a → s , while temperature T b dominates the path s
→ b .
So we have
[Equation]
We integrate the equation from Valley a to the outside of Valley a along the most possible escape path
[Equation]
We move J to the outside of integral based on Gauss’s Divergence Theorem, because J is ﬁxed on the escape path from one minimum to another.
As there is no ﬁeld source on the escape path, (cid:82)
V ∇ · J ( θ ) dV = 0 and ∇ J ( θ ) = 0 .
Obviously, probability sources are all near minima in deep learning.
So we obtain
[Equation]
Near saddle points, we have
[Equation]
Besides the temperature correction T = DγM [Equation] this result of J in Momentum also differs from SGD by the effective diffusion correction ˆ D b D b .
The effective diffusion correction coefﬁcient is given by
[Equation]
Based on the formula of the one-dimensional probability current and ﬂux
[Equation]
we obtain the highdimensional ﬂux escaping through b:
[Equation]
[Equation] where [ · ] ⊥ e indicates the directions perpendicular to the escape direction e .
Based on the results of Step 1 and Step
[Equation]
we have
[Equation]
We have replaced the eigenvalue of H b along the escape direction by its absolute value.
Finally, by introducing γ and M , we obtain the log-scale expression as
[Equation]
A. P ROOF OF P ROPOSITION
Proof.
The proof closely relates to the proof of 3.
We only need replace the standard learning rate by the adaptive learning rate ˆ η = ηv − [Equation] and set γM = 1 .
Particularly,
[Equation]
We introduce ˆ η = ηv − 12 into the proof of
[Equation]
and obtain
[Equation]
A. P ROOF OF P ROPOSITION
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 ofas:
[Equation]
As M = η ( I − β 1 ) − 1 , γ = η − 1 ( I − β 1 ) , and I − β 1 = β
¯ v v , we obtain the result:
[Equation]
A. P ROOF OF T HEOREM
Without loss of generality, we assume the dimensionality is one and rewrite the main updating rule of
Adai as
[Equation]
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 ﬁrst prove four useful lemmas.
Lemma 1.
Under the conditions of Theorem 4, for any t ≥ 0 , we have
[Equation]
where
[Equation]
Then we have
[Equation]
Proof.
Recall that
[Equation]
Then we have
[Equation]
Let q k,t
= (1 − β k ) (cid:81) ti = 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 :
[Equation]
Note that 0 ≤ β 1 ,k ≤ β 1 , max .
Then we have
[Equation]
Recursively, we can calculate the derivatives with respect to β 1 , 1 , β 1 , 2 , . .
. , β 1 ,t .
Then we obtain max( (cid:80) tk =0 q k,t )
= 1 by letting β 1 ,k = 0 for all k .
Similarly, we obtain min( (cid:80) t k =0 q k,t ) = 1 by letting β 1 ,t = β 1 , max for all k .
Then we obtain
[Equation]
The proof is now complete
[Equation]
Lemma 2.
Under the conditions of Theorem [Equation] for any t ≥ [Equation]
we have
[Equation]
Proof.
As L ( θ ) is L -smooth, we have
[Equation]
Taking expectation on both sides gives
[Equation]
By the Cauchy-Schwarz Inequality
[Equation]
we have
[Equation]
The proof is now complete
[Equation]
Lemma 3.
Under the conditions of Theorem [Equation] for any t ≥ [Equation]
we have
[Equation]
Proof.
For analyzing the maximum [Equation] we calculate the derivatives with respect to β 1 ,k for any 0 ≤ k ≤ t .
With respect to β 1 , 1 , we have
[Equation]
Then we let β 1 , 0 = 0 and obtain the derivative with respect to β 1 , 1 as
[Equation]
Then we let β 1 ,j = β 1 , max for 1 ≤ j ≤ k − 1 and obtain the derivative with respect to β 1 ,k as
[Equation]
Thus, we have β 1 , 0 = 0 and β 1 ,k = β 1 , max for 1 ≤ j ≤ t to maximize (cid:80) [Equation] 1 k[-=
-]
[Equation]
q k,t ( t − k ) .
We may write it as
[Equation]
Note that (1 − β 1 , max ) β k 1 , max k is an arithmetico-geometric sequence.
Thus, we have
[Equation]
The proof is now complete
[Equation]
Lemma 4.
Under the conditions of Theorem [Equation] for any t ≥ 0 , we have
[Equation]
1 − β 1 , max )( G 2 + δ 2 ) .
Proof.
As L ( θ ) is L -smooth, we have (cid:107)∇ L ( θ j +1 )
− L ( θ j ) (cid:107)(98) ≤L(cid:107) θ j +1 − θ j (cid:107)(99) = L(cid:107) − η j (cid:88) i = q i,j g i (cid:107) 2 .
By (cid:80) ji =
q i,j ≤ 1 in Lemma 1, we have
E [ (cid:107)∇ L ( θ j +1 ) − L ( θ j ) (cid:107) 2 ] (101) ≤L η 2 ( G 2 + δ 2 ) .
(102) and
L η2 (cid:107) t (cid:88) k
= q k,t g k (cid:107)(103) ≤L η2 ( G 2 + δ 2 ) .
By Lemma 3, we have  t (cid:88) k
= q k,t E [ t − 1 (cid:88)
j = k (cid:107)∇ L ( θ j +1 ) − L ( θ j ) (cid:107) 2 ] (105) ≤ η2 ( G 2 + δ 2 ) t (cid:88) k = q k,t ( t − k ) (106) ≤ η 2 β 1 , max 2(1 − β 1 , max ) ( G 2 + δ 2 ) .
Then we obtain
[Equation]
Proof.
The proof of Theorem 4 is organized as follows [Equation]
By Lemma 2 and Lemma
[Equation]
we have
[Equation]
Thus, we have
[Equation]
)] +
[Equation]
By Lemma 1 and Lemma 3, we have
[Equation]
By summing the above inequality for t = 0 , . . .
, t , we have
[Equation]
Then min
[Equation]
Let η ≤
C √ [Equation]
We have
[Equation]
The proof is now complete.
B C LASSICAL A PPROXIMATION A SSUMPTIONS
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 ﬂux 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 justiﬁed 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 = η 2 B H .
In statistical physics, if ∆ LD is large, then we call it Low Temperature Approximation.
Note that ∆ LD 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 ﬁnal 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 ∆
LH > 2 .
3 × 10 − 4 .
Empirically, the condition ∆ LH > 2 .
3 × 10 − 4 holds well in SGD dynamics.
C S TOCHASTIC G RADIENT N OISE A NALYSIS
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 ) ∈ [1 e − 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 ﬂat 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 veriﬁed that SGN is highly similar to
Gaussian noise instead of heavy-tailed Lévy noise.
Xie et al.
(2021b) recovered the experiment of
[Figure]
The noise covariance is approximately proportional to the Hessian and inverse to the batch size B even not around critical points
[Figure]
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.
Figureholds well at least when the batch size B is larger than 16 , which is common in practice.
D E XPERIMENTAL D ETAILS
Computational environment.
The experiments are conducted on a computing cluster with GPUs of
NVIDIA R (cid:13) Tesla TM P100 16GB and CPUs of Intel R (cid:13)
Xeon R (cid:13) CPU E5-2640 v3 @ 2.60GHz.
D. I MAGE C LASSIFICATION
Data Preprocessing: For CIFAR-10/CIFAR-100, we perform the per-pixel zero-mean unit-variance normalization, horizontal random ﬂip, 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 ﬂip, 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 43 ) 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 − −           ab c a
Energy Landscape (a) 1-Dimensional Escape ab a
[Figure] 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 L 2 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 L 2 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.
D. L ANGUAGE M ODELING
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, andhidden dimensions, which has 14 million model parameters and is similar to “medium LSTM” in
Zaremba et al.
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 .
Learning Rate: η = 0 .
The dropout probability is set to 0 .
We clipped gradient norm to 1 .
We select optimal learning rates from { 10 , 1 , 0 . 1 , 0 . 01 , 0 .
for each optimizer.
E T HE M EAN E SCAPE T IME A NALYSIS
Mean Escape Time: The mean escape time is the expected time for a particle governed by Equation 1 to escape from Sharp Valley a 1 to Flat Valley a 2 , 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 , 4 I ) .
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
[Equation]
We use 10-dimensional Styblinski-Tang Function as the test function, and Gaussian samples as training data.
[Equation]
where data samples x ∼ N (0 , 4 I )
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 a 1 and Valley a 2 .
For a n-dimensional Styblinski-Tang
Function, we initialize parameters θ t =0 = 1 √ k ( − 2 .
. , − 2 . 903534) , and set the valley’s boundary as θ i < 1 √ k 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 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
[Equation]
The 95% conﬁdence interval of this estimator is
[Equation]
F S UPPLEMENTARY E MPIRICAL R ESULTS
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 ﬁxed learning rate scheduler forepochs 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 deﬁned as E ζ [ L ( θ (cid:63) + ζ ) − L ( θ (cid:63) )] , where ζ is Gaussian noise and θ (cid:63) is the empirical minimizer learned by a training algorithm.
If the loss landscape near θ (cid:63) is sharp, the weight-perturbed loss E ζ [ L ( θ (cid:63) + ζ )] will be much larger than L ( θ (cid:63) ) .
Figure 13 empirically supports that Adai and
Momentum can learn signiﬁcantly ﬂatter minima than Adam.
[Figure]
[Figure] [Figure] [Figure] [Figure]
In this setting, some Adam variants may compare with SGD mainly because the baseline performance of SGD is lower than the baseline performance in Table 1. The test errors of AdaiW, Adai, and Momentum in middle columns is the original results in Table 1.
[Table] G A DAI WITH S TABLE /D ECOUPLED W EIGHT D ECAY g t = ∇ ˆ L ( θ t − 1 ) ; v t = β 2 v t − 1 + (1 − β 2 ) g 2 t ; ˆ v t
= v t 1 − β t 2 ; ¯ v t = mean (ˆ v t ) ; β [Equation] t = (1 − β 0 ¯ v t ˆ v t ) .Clip
(0 , 1 − (cid:15) ) ;
[Equation]
H E XPRESSIONS OF A DAM D YNAMICS
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
[Equation]
where β 3 = 1 − β 1 , β 1 is the hyperparameter, and ˆ η is the adaptive learning rate.
For simplicity, it is ﬁne 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 coefﬁcient γ as
[Equation]
where r t = − m t , F = g t , dt = ˆ η , 1 − γdt =
β 1 , and dtM =
β 3 = 1 − β 1 .
Thus, we obtain the differential-form motion equation of Adam as
[Equation]
[Equation]
where the mass M =
[Equation]
1 and the damping coefﬁcient
[Equation]
=
[Equation]
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 v t 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.
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.