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: This work proposes OASIS, a second-order method which approximates the diagonal of Hessian matrix and uses the information to rescale the gradient vector. The main difference between OASIS and the existing method AdaHessian (Yao et al., 2020) is on the ways they approximate the diagonal of Hessian, that AdaHessian uses a formula similar to Adam and OASIS uses an exponential average. Moreover, OASIS also incorporates the adaptive stepsize in (Mishchenko & Malitsky, 2020). The authors established the convergence guarantees of OASIS under various settings including convex, strongly convex and nonconvex cases, using various learning rate schedulers such as the adaptive learning rate, fixed learning rate and line search. Empirical results on various machine learning tasks are provided to evaluate OASIS.