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 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