Eigencurve: Optimal Learning Rate Schedule for SGD on Quadratic Objectives with Skewed Hessian Spectrums

1 INTRODUCTION . Many machine learning models can be represented as the following optimization problem : min w f ( w ) , 1 n n∑ i=1 fi ( w ) , ( 1.1 ) such as logistic regression , deep neural networks . To solve the above problem , stochastic gradient descent ( SGD ) ( Robbins & Monro , 1951 ) has been widely adopted due to its computation efficiency in large-scale learning problems ( Bottou & Bousquet , 2008 ) , especially for training deep neural networks . Given the popularity of SGD in this field , different learning rate schedules have been proposed to further improve its convergence rates . Among them , the most famous and widely used ones are inverse time decay , step decay ( Goffin , 1977 ) , and cosine scheduler ( Loshchilov & Hutter , 2017 ) . The learning rates generated by the inverse time decay scheduler depends on the current iteration number inversely . Such a scheduling strategy comes from the theory of SGD on strongly convex functions , and is extended to non-convex objectives like neural networks while still achieving good performance . Step decay scheduler keeps the learning rate piecewise constant and decreases it by a factor after a given amount of epochs . It is theoretically proved in Ge et al . ( 2019 ) that when the objective is quadratic , step decay scheduler outperforms inverse time decay . Empirical results are also provided in the same work to demonstrate the better convergence property of step decay in training neural networks when compared with inverse time decay . However , even step decay is proved to be near optimal on quadratic objectives , it is not truly optimal . There still exists a log T gap away from the minimax optimal convergence rate , which turns out to be non-trivial in a wide range of settings and may greatly impact step decay ’ s empirical performance . Cosine decay scheduler ( Loshchilov & Hutter , 2017 ) generates cosine-like learning rates in the range [ 0 , T ] , with T being the maximum iteration . It is a heuristic scheduling strategy which relies on the observation that good performance in practice can be achieved via slowly decreasing the learning rate in the beginning and “ refining ” the solution in the end with a very small learning rate . Its convergence property on smooth non-convex functions has been shown in Li et al . ( 2021 ) , but the provided bound is still not tight enough to explain its success in practice . Except cosine decay scheduler , all aforementioned learning rate schedulers have ( or will have ) a tight convergence bound on quadratic objectives . In fact , studying their convergence property on quadratic objective functions is quite important for understanding their behaviors in general nonconvex problems . Recent studies in Neural Tangent Kernel ( NTK ) ( Arora et al. , 2019 ; Jacot et al. , 2020 ) suggest that when neural networks are sufficiently wide , the gradient descent dynamic of neural networks can be approximated by NTK . In particular , when the loss function is least-square loss , neural network ’ s inference is equivalent to kernel ridge regression with respect to the NTK in expectation . In other words , for regression tasks , the non-convex objective in neural networks resembles quadratic objectives when the network is wide enough . The existence of log T gap in step decay ’ s convergence analysis , which will be proven to be tight in a wide range of settings , implies that there is still room for improvement in theory . Meanwhile , the existence of cosine decay scheduler , which has no strong theoretical convergence guarantees but possesses good empirical performance in certain tasks , suggests that its convergence rate may depend on some specific properties of the objective determined by the network and dataset in practice . Hence it is natural to ask what those key properties may be , and whether it is possible to find theoretically-optimal schedulers whose empirical performance are comparable to cosine decay if those properties are available . In this paper , we offer an answer to these questions . We first derive a novel eigenvalue-distribution-based learning rate scheduler called eigencurve for the quadratic function . Combining with the different eigenvalue distributions of different types of networks , new neural-network-based learning rate schedulers can be generated based on our proposed paradigm , which achieve better convergence properties than step decay in Ge et al . ( 2019 ) . Specifically , eigencurve closes the log T gap in step decay and reaches minimax optimal convergence rates if the Hessian spectrum is skewed . We summarize the main contributions of this paper as follows . 1 . To the best of our knowledge , this is the first work that incorporates the eigenvalue distribution of objective function ’ s Hessian matrix into designing learning rate schedulers . Accordingly , based on the eigenvalue distribution of the Hessian , we propose a novel eigenvalue distribution based learning rate scheduler named eigencurve . 2 . Theoretically , eigencurve can achieve optimal convergence rate ( up to a constant ) for SGD on quadratic objectives when the eigenvalue distribution of the Hessian is skewed . Furthermore , even when the Hessian is not skewed , eigencurve can still achieve no worse convergence rate than the step decay schedule in Ge et al . ( 2019 ) , whose convergence rate are proven to be sub-optimal in a wide range of settings . 3 . Empirically , on image classification tasks , eigencurve achieves optimal convergence rate for several models on CIFAR-10 and ImageNet if the loss can be approximated by quadratic objectives . Moreover , it obtains much better performance than step decay on CIFAR-10 , especially when the number of epochs is small . 4 . Intuitively , our learning rate scheduler sheds light on the theoretical property of cosine decay and provides a perspective of understanding the reason why it can achieve good performance on image classification tasks . The same idea has been used to inspire and discover several simple families of schedules that works in practice . Problem Setup For the theoretical analysis and the aim to derive our eigenvalue-dependent learning rate schedulers , we mainly focus on the quadratic function , that is , min w f ( w ) , Eξ [ f ( w , ξ ) ] , where f ( w , ξ ) = 1 2 w > H ( ξ ) w − b ( ξ ) > w , ( 1.2 ) where ξ denotes the data sample . Hence , the Hessian of f ( w ) is H = Eξ [ H ( ξ ) ] . ( 1.3 ) Letting us denote b = Eξ [ b ( ξ ) ] , we can obtain the optima of problem ( 1.2 ) w∗ = H −1b . ( 1.4 ) Given an initial iterate w0 and the learning rate sequence { ηt } , the stochastic gradient update is wt+1 = wt − ηt∇f ( wt , ξ ) = wt − ηt ( H ( ξ ) wt − b ( ξ ) ) . ( 1.5 ) We denote that nt = Hwt − b− ( H ( ξ ) wt − b ( ξ ) ) , µ , λmin ( H ) , L , λmax ( H ) , and , κ , L/µ . ( 1.6 ) In this paper , we assume that Eξ [ ntn > t ] σ2H . ( 1.7 ) The reason for this assumption is presented in Appendix G.5 . Related Work In convergence analysis , one key property that separates SGD from vanilla gradient descent is that in SGD , noise in gradients dominates . In gradient descent ( GD ) , constant learning rate can achieve linear convergence O ( cT ) with 0 < c < 1 for strongly convex objectives , i.e . obtaining f ( w ( t ) ) − f ( w∗ ) ≤ in O ( log ( 1 ) ) iterations . However , in SGD , f ( w ( t ) ) can not even be guaranteed to converge to f ( w∗ ) due to the existence of gradient noise ( Bottou et al. , 2018 ) . Intuitively , this noise leads to a variance proportional to the learning rate size , so constant learning rate will always introduce a Ω ( ηt ) = Ω ( η0 ) gap when compared with the con- vergence rate of GD . Fortunately , inverse time decay scheduler solves the problem by decaying the learning rate inversely proportional to the iteration number t , which achieves O ( 1T ) convergence rate for strongly convex objectives , specifically , O ( dσ 2 T · κ ) . However , this is sub-optimal since the minimax optimal rate for SGD is O ( dσ 2 T ) ( Ge et al. , 2019 ; Jain et al. , 2018 ) . Moreover , in practice , κ can be very big for large neural networks , which makes inverse time decay scheduler undesirable for those models . This is when step decay ( Goffin , 1977 ) comes to play . Empirically , it is widely adopted in tasks such as image classification and serves as a baseline for a lot of models . Theoretically , it has been proven that step decay can achieve nearly optimal convergence rateO ( dσ 2 T · log T ) for strongly convex least square regression ( Ge et al. , 2019 ) . Nevertheless , step decay is not always the best choice for image classification tasks . In practice , cosine decay ( Loshchilov & Hutter , 2017 ) can achieve comparable or even better performance , but the reason behind this superior performance is still unknown in theory ( Gotmare et al. , 2018 ) . All the aforementioned results are summarized in Table 1 , along with our results in this paper . It is worth mentioning that the minimax optimal rate O ( dσ 2 T ) can be achieved by iterate averaging methods ( Jain et al. , 2018 ; Bach & Moulines , 2013 ; Défossez & Bach , 2015 ; Frostig et al. , 2015 ; Jain et al. , 2016 ; Neu & Rosasco , 2018 ) , but it is not a common practice to use them in training deep neural networks , so only the final iterate ( Shamir , 2012 ) behavior of SGD is analyzed in this paper , i.e . the point right after the last iteration . Paper organization : Section 2 describes the motivation of our eigencurve scheduler . Section 3 presents the exact form and convergence rate of the proposed eigencurve scheduler , along with the lower bound of step decay . Section 4 shows the experimental results . Section 5 discusses the discovery and limitation of eigencurve and Section 6 gives our conclusion . 2 MOTIVATION . In this section , we will give the main motivation and intuition of our eigencurve learning rate scheduler . We first give the scheduling strategy to achieve the optimal convergence rate in the case that the Hessian is diagonal . Then , we show that the inverse time learning rate is sub-optimal in most cases . Comparing these two scheduling methods brings up the reason why we should design eigenvalue distribution dependent learning rate scheduler . Letting H be a diagonal matrix diag ( λ1 , λ2 , . . . , λd ) and reformulating Eqn . ( 1.5 ) , we have wt+1 − w∗ =wt − w∗ − ηt ( H ( ξ ) wt − b ( ξ ) ) =wt − w∗ − ηt ( Hwt − b ) + ηt ( Hwt − b− ( H ( ξ ) wt − b ( ξ ) ) ) =wt − w∗ − ηt ( Hwt − b− ( Hw∗ − b ) ) + ηt ( Hwt − b− ( H ( ξ ) wt − b ( ξ ) ) ) = ( I − ηtH ) ( wt − w∗ ) + ηtnt . It follows , E [ λj ( wt+1 , j − w∗ , j ) 2 ] E [ nt ] =0 = λj [ ( 1− ηtλj ) 2 ( wt , j − w∗ , j ) 2 + η2tE ‖ [ nt ] j‖ 2 ] ( 1.7 ) ≤ ( 1− ηtλj ) 2 · λj ( wt , j − w∗ , j ) 2 + η2t λ2jσ2 . ( 2.1 ) Since H is diagonal , we can set step size scheduling for each dimension separately . Letting us choose step size ηt coordinately with the step size ηt , j = 1λj ( t+1 ) being optimal for the j-th coordinate , then we have the following proposition . Proposition 1 . Assume that H is diagonal matrix with eigenvalues λ1 ≥ λ2 ≥ . . . , λd ≥ 0 and Eqn . ( 1.7 ) holds . If we set step size ηt , j = 1λj ( t+1 ) , it holds that 2E [ f ( wt+1 ) − f ( w∗ ) ] = E d∑ j=1 λj ( wt+1 , j − w∗ , j ) 2 ≤ ∑dj=1 λj ( w1 , j − w∗ , j ) 2 ( t+ 1 ) 2 + t ( t+ 1 ) 2 ·dσ2 . ( 2.2 ) The leading equality here is proved in Appendix G.1 , with the followed inequality proved in Appendix E. From Eqn . ( 2.2 ) , we can observe that choosing proper step sizes coordinately can achieve the optimal convergence rate ( Ge et al. , 2019 ; Jain et al. , 2018 ) . Instead , if the widely used inverse time scheduler ηt = 1/ ( L + µt ) is chosen , we can show that only a sub-optimal convergence rate can be obtained , especially when λj ’ s vary from each other . Proposition 2 . If we set the inverse time step size ηt = 1 ( L+µt ) , then we have 2E [ f ( wt+1 ) − f ( w∗ ) ] = E d∑ j=1 λj ( wt+1 , j − w∗ , j ) 2 ≤ ( L+ µ L+ µt ) 2 d∑ j=1 λj ( w1 , j − w∗ , j ) 2 + d∑ j=1 ( λ2j 2λj − µ · 1 L+ µt · σ2 + λ2jσ 2 ( L+ µt ) 2 ) . ( 2.3 ) Remark 1 . Eqn . ( 2.2 ) shows that if one can choose step size coordinate-wise with step size ηt , j = 1 λj ( t+1 ) , then SGD can achieve a convergence rate E [ f ( wT+1 ) − f ( w∗ ) ] ≤ O ( d T · σ2 ) . ( 2.4 ) which matches the lower bound ( Ge et al. , 2019 ; Jain et al. , 2018 ) . In contrast , replacing L = λ1 and µ = λd in Proposition 2 , we can obtain that the convergence rate of SGD being E [ f ( wT+1 ) − f ( w∗ ) ] ≤ O 1 T d∑ j=1 λj λd · σ2 . ( 2.5 ) Since it holds that λj ≥ λd , the convergence rate in Eqn . ( 2.4 ) is better than the one in Eqn . ( 2.5 ) , especially when the eigenvalues of the Hessian ( H matrix ) decay rapidly . In fact , the upper bound in Eqn . ( 2.5 ) is tight for the inverse time decay scheduling ( Ge et al. , 2019 ) . Main Intuition Though SGD with coordinately proper step sizes can obtain the optimal convergence rate , since H is commonly non-diagonal in practice , one is not able to choose proper step sizes coordinately . However , the diagonal case still provide an important intuition for designing eigenvalue dependent learning rate scheduling . In fact , letting H = UΛU > be the spectral decomposition of the Hessian and setting w′ = U > w , then the Hessian becomes a diagonal matrix from perspective of updating w′ , with the variance of the stochastic gradient being unchanged since U is a unitary matrix . Proposition 1 and 2 imply that a good learning rate scheduler should decrease the error of each coordinate . The inverse time decay scheduler is only optimal for the coordinate related to the smallest eigenvalue . That ’ s the reason why the inverse time scheduler is sub-optimal . Thus , we should reduce the learning rate gradually such that we can run a optimal learning rate associated to λj to sufficiently drop the error of j-th coordinate . Furthermore , given the total iteration T and the eigenvalue distribution of the Hessian , we should allocate the running time for each optimal learning rate associated to λj .

Summary: The paper studies convergence rates of SGD with different stepsize schemes, in the context of linear regression. For convergence of the last iterate of SGD, the best known result still misses a $\\log T$ factor compared to the minimax rate. This work aims to fill this gap with an improved stepsize scheme that utilizes the eigenvalue distribution of the Hessian. When the true Hessian is known, the proposed method successfully fills the gap, provably and in the sense of the worst problem instances. When the true Hessian is hard to know, practical variants are also proposed, and are shown to be comparable to the state-of-the-art stepsize schemes in standard deep neural network benchmarks.