Polygonal Unadjusted Langevin Algorithms: Creating stable and efficient adaptive algorithms for neural networks

1 INTRODUCTION . Modern machine learning models including deep neural networks are successfully trained when they are finely tuned via the optimization of their associated loss functions . Two aspects of such optimization tasks pose significant challenges , namely the non-convex nature of loss functions and the highly nonlinear features of many types of neural networks . Moreover , the analysis in Lovas et al . ( 2020 ) shows that the gradients of such non-convex loss functions typically grow faster than linearly and are only locally Lipschitz continuous . Naturally , stability issues are observed , which are known as the ‘ exploding gradient ’ phenomenon ( Bengio et al. , 1994 ; Pascanu et al. , 2013 ) , when vanilla stochastic gradient descent ( SGDs ) or certain types of adaptive algorithms are used for fine tuning . The sparsity of gradients of neural networks is another challenging issue , which is extensively studied in the literature . For example , momentum methods and adaptive learning rate methods such as AdaGrad ( Duchi et al . ( 2011 ) ) , RMSProp ( Tieleman & Hinton ( 2012 ) ) , Adam ( Kingma & Ba ( 2015 ) ) have been developed to tackle this problem and improve training speed by diagonally scaling the gradient by some function of the past gradients . A family of Langevin based algorithms has been another important stream of literature on the stochastic optimization . They are built on the theoretical fact that the Langevin stochastic differential equation , ( 6 ) , converges to its unique invariant measure , which concentrates on the global minimizers of the objective function as β →∞ , see Hwang ( 1980 ) . Since the convergence property remains true for nonconvex optimization problems , the global convergence of the stochastic gradient Langevin dynamics ( SGLD ) and its variants has been extensively studied in a nonconvex setting ( Raginsky et al. , 2017 ; Xu et al. , 2018 ; Erdogdu et al. , 2018 ; Brosse et al. , 2018 ; Lovas et al. , 2020 ) . Moreover , it is worth noting that Langevin based algorithms have been a key element in statistics and Bayesian learning ( Roberts & Tweedie , 1996 ; Durmus & Moulines , 2017 ; Dalalyan , 2017 ; Brosse et al. , 2019 ; Welling & Teh , 2011 ; Deng et al. , 2020a ; b ) . Motivated by the aforementioned developments in the field , we propose a new class of Langevin algorithms which is based on recent advances of Euler ’ s polygonal approximations for Langevin SDEs . The idea of this new form of Euler ’ s polygonal approximations for SDEs with monotone coefficients originates from the articles Krylov ( 1985 ) and Krylov ( 1990 ) . We name this new class as polygonal unadjusted Langevin algorithms . Moreover , it is versatile enough to incorporate further features to address other known shortcomings of adaptive optimizers . Mathematically , it is described as follows : Given an i.i.d . sequence of random variables { Xn } n≥0 of interest , which typically represent available data , the algorithm follows θλ0 : = θ0 , θ λ n+1 : = θ λ n − λGλ ( θλn , Xn+1 ) + √ 2λβ−1ξn+1 , n ∈ N , ( 1 ) where θ0 is an Rd-valued random variable , λ > 0 denotes the step size of the algorithm , β > 0 is the so-called inverse temperature , ( ξn ) n∈N is an Rd-valued Gaussian process with i.i.d . components and Gλ : Rd × Rm → Rd satisfies the following three properties : 1 . For every λ > 0 , There exist constants Kλ > 0 and ρ1 ≥ 0 such that |Gλ ( θ , x ) | ≤ Kλ ( 1 + |x| ) ρ1 ( 1 + |θ| ) for every θ ∈ Rd and x ∈ Rm . 2 . There exist constants γ ≥ 1/2 , K2 > 0 and ρ2 , ρ3 ≥ 0 such that for all λ > 0 , |Gλ ( θ , x ) −G ( θ , x ) | ≤ λγK2 ( 1 + |x| ) ρ2 ( 1 + |θ| ) ρ3 for every θ ∈ Rd and x ∈ Rm , whereG is the ( unbiased ) stochastic gradient of the objective function of the optimization problem under study . 3 . There exist constants λmax and δ ∈ { 1 , 2 } such that for any λ ≤ λmax , lim inf |θ|→∞ E [ 〈 θ |θ|δ , Gλ ( θ , X0 ) 〉 − 2λ |θ|δ |Gλ ( θ , X0 ) |2 ] > 0 . One obtains our new algorithm THεO POULA by considering the case where Gλ ( θ , x ) is the vector with entries H ( i ) λ , c ( θ , x ) as given by ( 8 ) , for i ∈ { 1 , . . . , d } . Its name is formed from its description , namely Tamed Hybrid ε-Order POlygonal Unadjusted Langevin Algorithm and its full detailed analysis ( including its convergence properties ) are given in Section 3 . We note that THεO POULA and TUSLA ( Lovas et al . ( 2020 ) ) satisfy the above three properties with δ = 2 and γ = 1/2 , whereas TULA ( Brosse et al . ( 2019 ) ) satisfies them with δ = γ = 1 as it assumes only deterministic gradients ( and thus the i.i.d . data sequence reduces to a constant ) . 1.1 RELATED WORK : LANGEVIN BASED ALGORITHMS AND ADAPTIVE LEARNING RATE METHODS . Most research on Langevin based algorithms in the literature has been focused on theoretical aspects . Raginsky et al . ( 2017 ) demonstrated the links between Langevin based algorithms and stochastic optimization in neural networks , stimulating further the development and analysis of such algorithms . Xu et al . ( 2018 ) analyzed the global convergence of GLD , SGLD and SVRG-LD . The incorporation of dependent data streams in the analysis of SGLD algorithms has been achieved in Barkhagen et al . ( 2021 ) and in Chau et al . ( 2019 ) , and local conditions have been studied in Zhang et al . ( 2019 ) . Recently , TUSLA of Lovas et al . ( 2020 ) has been proposed based on a new generation of tamed Euler approximations for stochastic differential equations ( SDEs ) with monotone coefficients in nonconvex optimization problems . See Hutzenthaler et al . ( 2012 ) and Sabanis ( 2013 ) for the rationale of taming techniques . Despite their elegant theoretical results , the use of Langevin based algorithms for training deep learning models has been limited in practice as their empirical performance lacked behind in comparison to other popular adaptive gradient methods . We refer to Appendix F.3 for the reader who is interested in recent progress on sampling and Bayesian neural networks . Adaptive learning rate methods such as AdaGrad ( Duchi et al . ( 2011 ) ) , RMSProp ( Tieleman & Hinton ( 2012 ) ) and Adam ( Kingma & Ba ( 2015 ) ) have been successfully applied to neural network models due to their fast training speed . Since the appearance of Adam , a large number of variants of Adam-type optimizers have been proposed to address the theoretical and practical challenges of Adam . For example , Reddi et al . ( 2018 ) provided a simple example that demonstrates the non-convergence issue of Adam and proposed a simple modification , called AMSGrad , to solve this problem . Chen et al . ( 2019 ) discussed the convergence of Adam-type optimizers in a nonconvex setting . RAdam to rectify the variance of adaptive learning rate has been proposed in Liu et al . ( 2020 ) . Wilson et al . ( 2017 ) revealed that the generalization ability of adaptive learning rate methods is worse than a global learning method like SGD . AdaBound of Luo et al . ( 2019 ) attempts to overcome the drawback by employing dynamic bounds on learning rates . Recently , AdaBelief ( Zhuang et al . ( 2020 ) ) and AdamP ( Heo et al . ( 2021 ) ) demonstrated their fast convergence and good generalization via extensive experiments . Nevertheless , the convergence analysis of these ( and other ) adaptive learning rate methods is still restrictive since it is only guaranteed to converge to a stationary point ( which can be a local minimum or a saddle point ) under strong assumptions . Namely , the stochastic gradient is globally Lipschitz continuous and bounded . Note though that none of these two assumptions hold true in a typical optimization problem involving neural networks . This is particularly evident in complex neural network architectures . 1.2 OUR CONTRIBUTIONS . The proposed algorithm , THεO POULA , tries to combines both advantages : namely , global convergence in Langevin based algorithms and powerful empirical performance in adaptive learning rate methods . To the best of the authors ’ knowledge , our algorithm is the first Langevin based algorithm to outperform popular stochastic optimization methods such as SGD , Adam , AMSGrad , RMSProp , AdaBound and AdaBelief for deep learning tasks . The major strengths of our work over related algorithms are summarized as follows : • ( Global convergence ) We provide a global convergence analysis of THεO POULA for nonconvex optimization where the stochastic gradient of the objective is locally Lipscthiz continuous . Moreover , non-asymptotic estimates for the expected excess risk are derived . • ( Stable and fast training ) THεO POULA achieves a stable and fast training process using the ( element-wise ) taming technique , ( element-wise ) boosting function and averaging , which are theoretically well-designed . Furthermore , we validate the effectiveness of the taming and boosting functions through several empirical experiments . • ( Good generalization ) While THεO POULA behaves like adaptive learning rate methods in the early training phase , it takes an almost global learning rate near an optimal point . That is , THεO POULA is quickly switched from adaptive methods to SGD . As a result , it inherits the good generalization ability of SGD . Our experiments support this fact by showing that THεO POULA outperforms the other optimization methods in generalization measured by test accuracy for various deep learning tasks . 2 MOTIVATING EXAMPLE . The local Lipschitz continuity of gradients and its effect on the performance of optimization methods are relatively under-studied . Most relevant studies assume that the stochastic gradient is global Lipscthiz continuous and bounded ( Kingma & Ba , 2015 ; Xu et al. , 2018 ; Brosse et al. , 2018 ; Duchi et al. , 2011 ; Tieleman & Hinton , 2012 ; Reddi et al. , 2018 ; Chen et al. , 2019 ; Liu et al. , 2020 ; Luo et al. , 2019 ; Zhuang et al. , 2020 ) although it is not true for neural network problems . This section provides a simple , one-dimensional optimization problem that illustrates the convergence issue of popular stochastic gradient methods when the stochastic gradient is locally Lipschitz continuous , i.e. , the gradient can be super-linearly growing 1 . Consider the following optimization problem : min θ u ( θ ) = min θ E [ U ( θ , X ) ] , ( 2 ) where U : R× R→ R is defined as U ( θ , x ) = { θ2 ( 1 + 1x≤1 ) + θ 30 , |θ| ≤ 1 , ( 2|θ| − 1 ) ( 1 + 1x≤1 ) + θ30 , |θ| > 1 , and X is uniformly distributed over ( −2 , 2 ) , that is , fX ( x ) = 141|x|≤2 . Furthermore , the stochastic gradient G : R× R→ R is given by G ( θ , x ) = { 2θ ( 1 + 1x≤1 ) + 30θ 29 , |θ| ≤ 1 , 2 ( 1 + 1x≤1 ) sgn ( θ ) + 30θ 29 , |θ| > 1 , where sgn ( · ) is the sign function . Note that the stochastic gradientG is locally Lipschitz continuous , which satisfies |G ( θ , x ) −G ( θ′ , x ) | ≤ 34 ( 1 + |θ|+ |θ′| ) 28|θ − θ′| 1Lovas et al . ( 2020 ) used a similar example to show the stability of TUSLA with a different taming function . for all x ∈ R and θ , θ′ ∈ R. Also , the optimal value is attained at θ = 0 . See Appendix A for more details . Following Reddi et al . ( 2018 ) , adaptive stochastic gradient methods can be generally written as follows , for n ∈ N , mn = φn ( G1 , · · · , Gn ) , Vn = ψn ( G1 , · · · , Gn ) , θn+1 = θn − λn mn ε+ √ V n ( 3 ) where Gi : = G ( θi , Xi ) is the stochastic gradient evaluated at the i-th iteration , λn is the step size and all operations are applied element-wise . Table 1 provides the details for some of the most popular stochastic optimization methods with corresponding averaging functions φn and ψn . We use SGD , Adam , AMSGrad and RMSprop to solve the optimization problem with initial value θ0 = 5 . For hyperparameters of optimization algorithms , we use their default settings provided in PyTorch . Figure 1 ( a ) shows the trajectories of approximate solutions generated by each optimizer . While SGD , Adam , AMSGrad and RMSProp fail to converge to the optimal solution 0 , the proposed algorithm , THεO POULA , finds the optimal solution with a reasonable step size , say , 0.01 . Intuitively , the undesirable phenomenon occurs because , in the iterating rule ( 3 ) , the denominator√ V n excessively dominates the numerator mn , causing the vanishing gradient problem in the presence of the superlinear gradient . On the contrary , SGD suffers from the exploding gradient problem . Moreover , Figure 1 ( b ) highlights that the problematic behavior can not be simply resolved by adjusting the learning rate within the Adam-type framework , while THεO POULA perform extremely well even in the presence of such violent non-linearities . 3 NEW ALGORITHM : THεO POULA We propose a new stochastic optimization algorithm by combining ideas from taming methods specifically designed to approximate Langevin SDEs with a hybrid approach based on recent advances of polygonal Euler approximations . The latter is achieved by identifying a suitable boosting function ( of order ε 1 ) to efficiently deal with the sparsity of ( stochastic ) gradients of neural networks . In other words , the novelty of our algorithm is to utilize a taming function and a boosting function , rather than designing a new Vn as in Adam-type optimizers . We proceed with the necessary preliminary information , main assumptions and formal introduction of the new algorithm .

Summary: The paper presents a new stochastic optimization algorithm inspired by the tamed Euler discretization of the overdamped Langevin diffusion. Convergence guarantees in the Wasserstein metric are established for this algorithm under mild assumptions, primarily local Lipschitz continuity and dissipativity conditions on the (stochastic) gradient of the loss function. In addition to a taming function, the proposed algorithm "Th$\varepsilon$O POULA", employs a boosting function to improve performance in flat regions of the loss landscape. Along with a synthetic example, the authors demonstrate the effectiveness of Th$\varepsilon$O POULA for training deep neural networks for image classification and language modelling.