Apollo: An Adaptive Parameter-wised Diagonal Quasi-Newton Method for Nonconvex Stochastic Optimization

1 INTRODUCTION . Nonconvex stochastic optimization is of core practical importance in many fields of machine learning , in particular for training deep neural networks ( DNNs ) . First-order gradient-based optimization algorithms , conceptually attractive due to their linear efficiency on both the time and memory complexity , have led to tremendous progress and impressive successes . A number of advanced first-order algorithms have emerged over the years to pursue fast and stable convergence , among which stochastic gradient descent ( SGD ) ( Robbins & Monro , 1951 ; LeCun et al. , 1998 ) , equipped with momentum ( Rumelhart et al. , 1985 ; Qian , 1999 ; Bottou & Bousquet , 2008 ) , has stood out for its simplicity and effectiveness across a wide range of applications ( Hinton & Salakhutdinov , 2006 ; Hinton et al. , 2012 ; Graves , 2013 ) . However , one disadvantage of SGD is that the gradients in different directions are scaled uniformly , resulting in limited convergence speed and sensitive choice of learning rate , and thus has spawned a lot of recent interests in accelerating SGD from the algorithmic and practical perspectives . Recently , many adaptive first-order optimization methods have been proposed to achieve rapid training progress with element-wise scaled learning rates , and we can only mention a few here due to space limits . In their pioneering work , Duchi et al . ( 2011 ) proposed AdaGrad , which scales the gradient by the square root of the accumulative square gradients from the first iteration . While AdaGrad works well for sparse settings , its performance significantly degrades for dense settings , primarily due to the monotonic increase of the accumulation . Subsequently , several methods have been proposed with the intuition to limit the accumulation to a small window of past iterations , and in particular exponentially reduce the weight of earlier iterations . Notable works incorporating this method are RMSProp ( Tieleman & Hinton , 2012 ) , AdaDelta ( Zeiler , 2012 ) , and Adam ( Kingma & Ba , 2015 ) , among which Adam has become the default optimization algorithm across many deep learning applications because of its fast convergence speed and relatively consistent selections of hyper-parameters ( Ruder , 2016 ; Zhang et al. , 2020 ) . However , it has been observed that these adaptive optimization methods may converge to bad/suspicious local optima , resulting in worse generalization ability than their non-adaptive counterparts ( Wilson et al. , 2017 ) , or fail to converge due to unstable and extreme learning rates ( Luo et al. , 2019 ) . Quasi-Newton methods have been widely used in solving convex optimization problems , due to their efficient computation and fast convergence rate ( Broyden , 1967 ; Dennis & Moré , 1977 ) . However , the stochastic , high-dimensional and nonconvex nature of many machine learning tasks , such as training deep neural networks , has rendered many classical quasi-Newton methods ineffective and/or inefficient ( Keskar & Berahas , 2016 ; Wang et al. , 2017 ; Yao et al. , 2020 ) . Indeed , in many natural language processing ( NLP ) and computer vision ( CV ) tasks ( He et al. , 2016 ; Ma & Hovy , 2016 ; Luo et al. , 2019 ) , SGD ( with momentum ) is chosen as the optimizer , benefiting from its stable and efficient training and outstanding generalization . In this work , we develop APOLLO , a quasi-Newton method for nonconvex stochastic optimization to simultaneously tackle the aforementioned challenges of stochastic variance , nonconvexity and inefficiency . Algorithmically , APOLLO dynamically incorporates the curvature of the objective function with diagonally approximated Hessian . It only requires first-order gradients and updates the approximation of the Hessian diagonally so that it satisfies a parameter-wise version of the weak secant condition ( Wolfe , 1959 ) . To handle nonconvexity , we replace the Hessian with its rectified absolute value , the computation of which is also efficient under our diagonal approximation , yielding an efficient optimization algorithm with linear complexity for both time and memory ( §3 ) . Experimentally , through three tasks on CV and NLP with popular deep neural networks , including ResNets ( He et al. , 2016 ) , LSTMs ( Hochreiter & Schmidhuber , 1997 ) and Transformers ( Vaswani et al. , 2017 ) , we demonstrate that APOLLO significantly outperforms SGD and variants of Adam , in terms of both convergence speed and generalization performance ( §4 ) . 2 BACKGROUNDS . In this section , we set up the notations on nonconvex stochastic optimization , briefly review the ( quasi- ) Newton methods , and discuss the problems of applying quasi-Newton methods to nonconvex stochastic optimization that we attempt to study in the rest of the paper . 2.1 NONCONVEX STOCHASTIC OPTIMIZATION . In this paper , we consider the following nonconvex stochastic optimization problem : min θ∈Rd f ( θ ) = E [ l ( θ ; Γ ) ] ( 1 ) where l : Rd ×Rn → R is a continuously differentiable ( and possible nonconvex ) function , θ ∈ Rd denotes the parameter to be optimized , Γ ∈ Rn denotes a random variable with distribution function P , and E [ · ] denotes the expectation w.r.t Γ . Intuitively , Γ incorporates noises in f , leading to a stochastic objective function . A special case of ( 1 ) that arises frequently in machine learning is the empirical risk minimization problem : min θ∈Rd f ( θ ) = 1 N N∑ i=1 li ( θ ) ( 2 ) where li : Rd → R is the loss function corresponding to the i-th data , and N is the number of data samples that is assumed to be extremely large . Objective functions may also have other sources of noise than data subsampling , such as dropout ( Srivastava et al. , 2014 ) in deep neural networks . Decoupled Parameters . In this work , we consider a setting of decoupled parameters : θ = { θ ( l ) , l = 1 , . . . , L } . Intuitively , under this setting the parameter θ is decoupled into a sequence of parameters serving different functionalities . For example , in neural network training the parameters of a neural network can be naturally decoupled into the parameters of different layers or modules . 2.2 NEWTON AND QUASI-NEWTON METHODS . Newton ’ s method usually employs the following updates to solve ( 1 ) : θt+1 = θt −H−1t gt ( 3 ) where gt = ∇f ( θt ) is the gradient at θt and Ht = ∇2f ( θt ) is the Hessian matrix . The convergence rate of Newton ’ s method is quadratic under standard assumptions ( Nocedal & Wright , 2006 ) . However , major challenges with this method are i ) the expensive computation of the inverse Hessian at every iteration and the corresponding quadratic memory complexity ; and ii ) the limitation to convex functions ( nonconvexity results in negative curvature of Ht and misleads the update directions ) . A standard alternative to Newton ’ s method is a class of quasi-Newton methods , which have been widely used in solving convex deterministic optimization problem : θt+1 = θt − ηtB−1t gt ( 4 ) where ηt is the stepsize ( a.k.a learning rate ) , Bt is an approximation to the Hessian matrix ∇2f ( θt ) at θt , which is updated based on the well-known secant equation : Bt+1 = argmin B ‖B −Bt‖ s.t . Bt+1st = yt ( secant equation ) ( 5 ) where st = θt+1−θt and yt = gt+1−gt . Bt+1 is , in the sense of some matrix norm , the closest toBt among all symmetric matrices that satisfy the secant equation . Each choice of the matrix norm results in a different update formula , such as DFP ( Davidon , 1991 ; Fletcher , 1987 ) and BFGS ( Broyden , 1970 ; Fletcher , 1970 ; Goldfarb , 1970 ; Shanno , 1970 ) . The popularity of this method is due to the fact that only the gradient of the objective function is required at each iteration . Since no second derivatives ( Hessian ) are required , quasi-Newton methods are sometimes more efficient than Newton ’ s method , especially when the computation of Hessian is expensive . To further reduce memory cost , one seminal work is the limited memory BFGS ( L-BFGS ) ( Liu & Nocedal , 1989 ; Byrd et al. , 1995 ) that achieves desirable linear computational and memory complexity by approximating the Hessian as a series of sum of first order information from previous iterations . 2.3 PROBLEMS OF QUASI-NEWTON METHODS . Despite their impressive successes on convex deterministic optimization , quasi-Newton methods suffer from their own problems in more challenging scenarios . In this section , we mainly discuss three problems preventing quasi-Newton methods from being applied to the scenario of largescale nonconvex stochastic optimization . Due to these problems , no quasi-Newton methods ( to our best knowledge ) designed for nonconvex optimization consistently outperform adaptive firstorder algorithms w.r.t convergence speed and generalization performance . The main goal of this work is to algorithmically design and experimentally demonstrate a novel quasi-Newton method , in hope of improving the convergence speed and generalization performance of nonconvex stochastic optimization eventually . Stochastic Variance . One challenge of quasi-Newton methods on nonconvex stochastic optimization ( 1 ) is the variance introduced by the stochastic nature of the problem . At each iteration , only the stochastic gradient gt is available , which is an unbiased estimation of the gradient∇f ( θt ) and may lead to an erroneous approximation of Hessian ( Byrd et al. , 2011 ) . Nonconvexity . Another key challenge in designing such quasi-Newton methods lies in the difficulty of preserving the positive-definiteness of Bt in ( 5 ) , due to the nonconvexity of the objective function . What is worse is that performing line search is infeasible in the stochastic setting , due to the presence of noise in the stochastic gradients ( Wang et al. , 2017 ) . Computational and Memory Efficiency . Even though quasi-Newton methods are more efficient than Newton ’ s method , the time and memory complexities are still relatively large compared with adaptive first-order methods . For instance , L-BFGS requires to store first-order information from m previous iterations with commonly m ≥ 5 , which is still too expensive for deep neural networks containing millions of parameters . Moreover , adapting quasi-Newton methods to nonconvex stochastic optimization probably introduces additional computation , further slowing down these methods . 3 ADAPTIVE PARAMETER-WISE DIAGONAL QUASI-NEWTON . With the end goal of designing an efficient quasi-Newton method to solve the problem in ( 1 ) in mind , we first propose to approximate the Hessian with a diagonal matrix , whose elements are determined by the variational approach subject to the parameter-wise weak secant equation ( §3.1 ) . Then , we explain our stepsize bias correction technique to reduce the stochastic variance in §3.2 . To handle nonconvexity , we directly use the rectified absolute value of the diagonally approximated Hessian as the preconditioning of the gradient ( §3.3 ) . The initialization technique of APOLLO allows us to eliminate one hyper-parameter ( §3.4 ) . At last , we provide a theoretical analysis of APOLLO ’ s convergence in both convex optimization and nonconvex stochastic optimization ( §3.5 ) . The pseudo-code is shown in Algorithm 1 . 3.1 QUASI-NEWTON METHODS WITH DIAGONAL HESSIAN APPROXIMATION . As discussed in Bordes et al . ( 2009 ) , designing an efficient stochastic quasi-Newton algorithm involves a careful trade-off between the sparsity of the approximation matrix Bt and the quality of its approximation of the Hessian Ht , and diagonal approximation is a reasonable choice ( Becker et al. , 1988 ; Zhu et al. , 1999 ) . If B is chosen to be a diagonal matrix satisfying ( 5 ) , one can obtain a formula similar to the SGD-QN algorithm ( Bordes et al. , 2009 ) . An alternative of the secant equation in the updating formula ( 5 ) , as first introduced by Nazareth ( 1995 ) , is the weak secant equation ( Dennis & Wolkowicz , 1993 ) : Bt+1 = argmin B ‖B −Bt‖ s.t . sTt Bt+1st = s T t yt ( weak secant equation ) ( 6 ) The motivation of using the weak secant condition in diagonal quasi-Newton method is straightforward : the standard mean-value theorem might not necessarily hold for vector-valued functions expressed in the secant equation , Bt+1st = yt ≈ ∇2f ( θt ) st . Thus , we do not know whether there exists a vector θ̃ ∈ Rd such that yt = ∇2f ( θ̃ ) st ( Dennis & Moré , 1977 ) . On the other hand , the Taylor theorem ensures that there exists such θ̃ that sTt yt = s T t ∇2f ( θ̃ ) st , leading to the reasonable assumption of the weak secant condition ( 6 ) . Based on the variational technique ( Zhu et al. , 1999 ) , the solution of ( 6 ) with Frobenius norm is : Λ , Bt+1 −Bt = sTt yt − sTt Btst ‖st‖44 Diag ( s2t ) ( 7 ) where s2t is the element-wise square vector of st , Diag ( s 2 t ) is the diagonal matrix with diagonal elements from vector s2t , and ‖ · ‖4 is the 4-norm of a vector . Parameter-Wise Weak Secant Condition . However , in optimization problems with highdimensional parameter space , such as training deep neural networks with millions of parameters , the weak secant condition might be too flexible to produce a good Hessian approximation . In the setting of decoupled parameters ( §2.1 ) , we propose a parameter-wise version of the weak secant equation to achieve a trade-off between the secant and weak secant conditions : for each parameter θ ( l ) ∈ θ , we update B corresponding to θ ( l ) by solving ( 6 ) individually . Remarkably , the secant condition restricts B with an equation of a d-dimensional vector , while the weak secant condition relaxes it with a 1-dimensional scalar . The parameter-wise weak secant condition expresses the restriction as a l-dimension vector ( 1 < l < d ) , resulting in a reasonable trade-off . The updating formula is the same as ( 7 ) for each parameter-wise B .

Summary: This work proposes a quasi-Newton method APOLLO for nonconvex stochastic optimization. Based on a parameter-wise weak secant condition, a diagonal approximation of the Hessian is constructed, and rectified absolute value is adopted on the approximation. A step-size bias correction technique is used to mitigate stochastic gradient variance. Theoretical convergence analysis is provided under a convex online setting and a nonconvex stochastic setting. Experiments on CV and NLP tasks demonstrate the effectiveness and stability of the proposed method.