Adaptive Inertia: Disentangling the Effects of Adaptive Learning Rate and Momentum

1 INTRODUCTION . Adam ( Kingma and Ba , 2015 ) , which combines Adaptive Learning Rate and Momentum , would be the most popular optimizer for accelerating the training of deep networks . However , Adam often generalizes worse and finds sharper minima than SGD ( Wilson et al. , 2017 ) for popular convolutional neural networks , where the flat minima have been argued to be closely related with good generalization ( Hochreiter and Schmidhuber , 1995 ; 1997a ; Hardt et al. , 2016 ; Zhang et al. , 2017 ) . Meanwhile , the diffusion theory has been used as 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 . 2 SGD AND DIFFUSION . In this section , we introduce some preliminaries about the SGD diffusion and then present the saddle-point escaping property of SGD dynamics . 2.1 PREREQUISITES FOR SGD DIFFUSION . We first review the SGD diffusion theory for escaping minima proposed by Xie et al . ( 2021b ) . We denote the model parameters as θ , 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 ) . 2.2 SGD DIFFUSION NEAR SADDLE POINTS . 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 . 3 ANALYSIS OF MOMENTUM DYNAMICS . In this section , we analyze Momentum in terms of saddle-point escaping and flat minima selection . The continuous-time Momentum dynamics . The Heavy Ball Method ( Zavriev and Kostyuk , 1993 ) can be written as { 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 .

Summary: 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.