Doubly Adaptive Scaled Algorithm for Machine Learning Using Second-Order Information

1 INTRODUCTION . This paper presents an algorithm for solving empirical risk minimization problems of the form : minw∈RdF ( w ) : = 1 n ∑n i=1f ( w ; x i , yi ) = 1n ∑n i=1fi ( w ) , ( 1 ) where w is the model parameter/weight vector , { ( xi , yi ) } ni=1 are the training samples , and fi : Rd → R is the loss function . Usually , the number of training samples , n , and dimension , d , are large , and the loss function F is potentially nonconvex , making this minimization problem difficult to solve . In the past decades , significant effort has been devoted to developing optimization algorithms for machine learning . Due to easy implementation and low per-iteration cost , ( stochastic ) first-order methods ( Robbins & Monro , 1951 ; Duchi et al. , 2011 ; Schmidt et al. , 2017 ; Johnson & Zhang , 2013 ; Nguyen et al. , 2017 ; 2019 ; Kingma & Ba , 2014 ; Jahani et al. , 2021a ; Recht et al. , 2011 ) have become prevalent approaches for many machine learning applications . However , these methods have several drawbacks : ( i ) they are highly sensitive to the choices of hyperparameters , especially learning rate ; ( ii ) they suffer from ill-conditioning that often arises in large-scale machine learning ; and ( iii ) they offer limited opportunities in distributed computing environments since these methods usually spend more time on “ communication ” instead of the true “ computation. ” The main reasons for the aforementioned issues come from the fact that first-order methods only use the gradient information for their updates . On the other hand , going beyond first-order methods , Newton-type and quasi-Newton methods ( Nocedal & Wright , 2006 ; Dennis & Moré , 1977 ; Fletcher , 1987 ) are considered to be a strong family of optimizers due to their judicious use of the curvature information in order to scale the gradient . By exploiting the curvature information of the objective function , these methods mitigate many of the issues inherent in first-order methods . In the deterministic regime , it is known that these methods are relatively insensitive to the choices of the hyperparameters , and they handle ill-conditioned problems with a fast convergence rate . Clearly , this does not come for free , and these methods can have memory requirements up to O ( d2 ) with computational complexity up to O ( d3 ) ( e.g. , with a naive use of the Newton method ) . There are , of course , efficient ways to solve the Newton system with significantly lower costs ( e.g. , see Nocedal & Wright ( 2006 ) ) . Moreover , quasi-Newton methods require lower memory and computational complexities than Newton-type methods . Recently , there has been shifted attention towards stochastic second-order ( Roosta-Khorasani & Mahoney , 2018 ; Byrd et al. , 2011 ; Martens , 2010 ; Jahani et al. , 2020a ; Xu et al. , 2017 ; Roosta et al. , 2018 ; Yao et al. , 2018 ) and quasi-Newton methods ( Curtis , 2016 ; Berahas et al. , 2016 ; Mokhtari & Ribeiro , 2015 ; Jahani et al. , 2021b ; Berahas et al. , 2019 ; Jahani et al. , 2020b ) in order to approximately capture the local curvature information . These methods have shown good results for several machine learning tasks ( Xu et al. , 2020 ; Berahas et al. , 2020 ; Yao et al. , 2019 ) . In some cases , however , due to the noise in the Hessian approximation , their performance is still on par with the first-order variants . One avenue for reducing the computational and memory requirements for capturing curvature information is to consider just the diagonal of the Hessian . Since the Hessian diagonal can be represented as a vector , it is affordable to store its moving average , which is useful for reducing the impact of noise in the stochastic regime . To exemplify this , AdaHessian Algorithm ( Yao et al. , 2020 ) uses Hutchinson ’ s method ( Bekas et al. , 2007 ) to approximate the Hessian diagonal , and it uses a second moment of the Hessian diagonal approximation for preconditioning the gradient . AdaHessian achieves impressive results on a wide range of state-of-the-art tasks . However , its preconditioning matrix approximates the Hessian diagonal only very approximately , suggesting that improvements are possible if one can better approximate the Hessian diagonal . 1 0 0 0 0 S a m p le s In this paper , we propose the dOubly Adaptive Scaled algorIthm for machine learning using Second-order information ( OASIS ) . OASIS approximates the Hessian diagonal in an efficient way , providing an estimate whose scale much more closely approximates the scale of the true Hessian diagonal ( see Figure 1 ) . Due to this improved scaling , the search direction in OASIS contains gradient information , in which the components are well-scaled by the novel preconditioning matrix . Therefore , every gradient component in each dimension is adaptively scaled based on the approximated curvature for that dimension . For this reason , there is no need to tune the learning rate , as it would be updated automatically based on a local approximation of the Lipschitz smoothness parameter ( see Figure 2 ) . The well-scaled preconditioning matrix coupled with the adaptive learning rate results in a fully adaptive step for updating the parameters . Here , we provide a brief summary of our main contributions : • Novel Optimization Algorithm . We propose OASIS as a fully adaptive method that preconditions the gradient information by a well-scaled Hessian diagonal approximation . The gradient component in each dimension is adaptively scaled by the corresponding curvature approximation . • Adaptive Learning Rate . Our methodology does not require us to tune the learning rate , as it is updated automatically via an adaptive rule . The rule approximates the Lipschitz smoothness parameter , and it updates the learning rate accordingly . • Comprehensive Theoretical Analysis . We derive convergence guarantees for OASIS with respect to different settings of learning rates , namely the case with adaptive learning rate for convex and strongly convex cases . We also provide the convergence guarantees with respect to fixed learning rate and line search for both strongly convex and nonconvex settings . • Competitive Numerical Results . We investigate the empirical performance of OASIS on a variety of standard machine learning tasks , including logistic regression , nonlinear least squares problems , and image classification . Our proposed method consistently shows competitive or superior performance in comparison to many first- and second-order state-of-the-art methods . Notation . By considering the positive definite matrix D , we define the weighted Euclidean norm of vector x ∈ Rd with ‖x‖2D = xTDx . Its corresponding dual norm is shown as ‖ · ‖∗D . The operator is used as a component-wise product between two vectors . Given a vector v , we represent the corresponding diagonal matrix of v with diag ( v ) . 2 RELATED WORK . In this paper , we analyze algorithms with the generic iterate updates : wk+1 = wk − ηkD̂k−1mk , ( 2 ) where D̂k is the preconditioning matrix , mk is either gk ( the true gradient or the gradient approximation ) or the first moment of the gradient with momentum parameter β1 or the bias corrected first moment of the gradient , and ηk is the learning rate . The simple interpretation is that , in order to update the iterates , the vector mk would be rotated and scaled by the inverse of preconditioning matrix D̂k , and the transformed information would be considered as the search direction . Due to limited space , here we consider only some of the related studies with a diagonal preconditioner . For more general preconditioning , see Nocedal & Wright ( 2006 ) . Clearly , one of the benefits of a well-defined diagonal preconditioner is the easy calculation of its inverse . There are many optimization algorithms that follow the update in ( 2 ) . A well-known method is stochastic gradient descent ( SGD ) .The idea behind SGD is simple yet effective : the preconditioning matrix is set to be D̂k = Id , for all k ≥ 0 . There are variants of SGD with and without momentum . The advantage of using momentum is to smooth the gradient ( approximation ) over the past iterations , and it can be useful in the noisy settings . In order to converge to the stationary point ( s ) , the learning rate in SGD needs to decay . Therefore , there are many important hyperparameters that need to be tuned , e.g. , learning rate , learning-rate decay , batch size , and momentum . Among all of them , tuning the learning rate is particularly important and cumbersome since the learning rate in SGD is considered to be the same for all dimensions . To address this issue , one idea is to use an adaptive diagonal preconditioning matrix , where its elements are based on the local information of the iterates . One of the initial methods with a non-identity preconditioning matrix is Adagrad ( Duchi et al. , 2011 ) . In Adagrad , the momentum parameter is set to be zero ( mk = gk ) , and the preconditioning matrix is defined as : D̂k = diag ( √∑k i=1gk gk ) . ( 3 ) As is clear from the preconditioning matrix D̂k in ( 3 ) , every gradient component is scaled with the accumulated information of all the past squared gradients . It is advantageous in the sense that every component is scaled adaptively . However , a significant drawback of D̂k in ( 3 ) has to do with the progressive increase of its elements , which leads to rapid decrease of the learning rate . To prevent Adagrad ’ s aggressive , monotonically decreasing learning rate , several approaches , including Adadelta ( Zeiler , 2012 ) and RMSProp ( Tieleman & Hinton , 2012 ) , have been developed . Specifically , in RMSProp , the momentum parameter β1 is zero ( or mk = gk ) and the preconditioning matrix is as follows : D̂k = √ β2D̂k−12 + ( 1− β2 ) diag ( gk gk ) , ( 4 ) where β2 is the momentum parameter used in the preconditioning matrix . As we can see from the difference between the preconditioning matrices in ( 3 ) and ( 4 ) , in RMSProp an exponentially decaying average of squared gradients is used , which prevents rapid increase of preconditioning components in ( 3 ) . Another approach for computing the adaptive scaling for each parameter is Adam ( Kingma & Ba , 2014 ) . Besides storing an exponentially decaying average of past squared gradients like Adadelta and RMSprop , Adam also keeps first moment estimate of gradient , similar to SGD with momentum . In Adam , the bias-corrected first and second moment estimates , i.e. , mk and D̂k in ( 2 ) , are as follows : mk = 1−β1 1−βk1 ∑k i=1β k−i 1 gi , D̂k = √ 1−β2 1−βk2 ∑k i=1β k−i 2 diag ( gi gi ) . ( 5 ) There have been many other first-order methods with adaptive scaling ( Loshchilov & Hutter , 2017 ; Chaudhari et al. , 2019 ; Loshchilov & Hutter , 2016 ; Shazeer & Stern , 2018 ) . The methods described so far have only used the information of the gradient for preconditioning mk in ( 2 ) . The main difference of second-order methods is to employ higher order information for scaling and rotating the mk in ( 2 ) . To be precise , besides the gradient information , the ( approximated ) curvature information of the objective function is also used . As a textbook example , in Newton ’ s method D̂k = ∇2F ( wk ) and mk = gk with ηk = 1 . Diagonal Approximation . Recently , using methods from randomized numerical linear algebra , the AdaHessian method was developed ( Yao et al. , 2020 ) . AdaHessian approximates the diagonal of the Hessian , and it uses the second moment of the diagonal Hessian approximation as the preconditioner D̂k in ( 2 ) . In AdaHessian , Hutchinson ’ s method1 is used to approximate the Hessian diagonal as follows : Dk ≈ diag ( E [ zk ∇2F ( wk ) zk ] ) , ( 6 ) where zk is a random vector with Rademacher distribution . Needless to say,2 the oracle ∇2F ( wk ) zk , or Hessian-vector product , can be efficiently calculated ; in particular , for AdaHessian , it is computed with two back-propagation rounds without constructing the Hessian explicitly . In a nutshell , the first momentum for AdaHessian is the same as ( 5 ) , and its second order momentum is : D̂k = √ 1−β2 1−βk2 ∑k i=1β k−i 2 D 2 i . ( 7 ) The intuition behind AdaHessian is to have a larger step size for the dimensions with shallow loss surfaces and smaller step size for the dimensions with sharp loss surfaces . The results provided by AdaHessian show its strength by using curvature information , in comparison to other adaptive first-order methods , for a range of state-of-the-art problems in computer vision , natural language processing , and recommendation systems ( Yao et al. , 2020 ) . However , even for AdaHessian , the preconditioning matrix D̂k in ( 7 ) does not approximate the scale of the actual diagonal of the Hessian particularly well ( see Figure 1 ) . One might hope that a better-scaled preconditioner would enable better use of curvature information . This is one of the main focuses of this study . Adaptive Learning Rate . In all of the methods discussed previously , the learning rate ηk in ( 2 ) is still a hyperparameter which needs to be manually tuned , and it is a critical and sensitive hyperparameter . It is also necessary to tune the learning rate in methods that use approximation of curvature information ( such as quasi-Newton methods like BFGS/LBFGS , and methods using diagonal Hessian approximation like AdaHessian ) . The studies ( Loizou et al. , 2020 ; Vaswani et al. , 2019 ; Chandra et al. , 2019 ; Baydin et al. , 2017 ; Malitsky & Mishchenko , 2020 ) have tackled the issue regarding tuning learning rate , and have developed methodologies with adaptive learning rate , ηk , for first-order methods . Specifically , the work ( Malitsky & Mishchenko , 2020 ) finds the learning rate by approximating the Lipschitz smoothness parameter in an affordable way without adding a tunable hyperparameter which is used for GD-type methods ( with identity norm ) . Extending the latter approach to the weighted-Euclidean norm is not straightforward . In the next section , we describe how we can extend the work ( Malitsky & Mishchenko , 2020 ) for the case with weighted Euclidean norm . This is another main focus of this study . In fact , while we focus on AdaHessian , any method with a positive-definite preconditioning matrix and bounded eigenvalues can benefit from our approach .

Summary: The paper designs and analyses an algorithm for minimizing a separable function. It provides deterministic as well as stochastic versions, which either fully compute the gradient or sample it. The algorithm estimates the diagonal elements of the Hessian via Hessian-vector products. The algorithm makes use of this information for finding a better search direction and the step length, eliminating the need for a line search. It provides convergence guarantees and a number of experiments on classical ML problems as well as on deep nets.