Large Learning Rate Tames Homogeneity: Convergence and Balancing Effect

Recent empirical advances show that training deep models with large learning rate often improves generalization performance . However , theoretical justifications on the benefits of large learning rate are highly limited , due to challenges in analysis . In this paper , we consider using Gradient Descent ( GD ) with a large learning rate on a homogeneous matrix factorization problem , i.e. , minX , Y ‖A−XY > ‖2F . We prove a convergence theory for constant large learning rates well beyond 2/L , where L is the largest eigenvalue of Hessian at the initialization . Moreover , we rigorously establish an implicit bias of GD induced by such a large learning rate , termed ‘ balancing ’ , meaning that magnitudes of X and Y at the limit of GD iterations will be close even if their initialization is significantly unbalanced . Numerical experiments are provided to support our theory . 1 INTRODUCTION . Training machine learning models such as deep neural networks involves optimizing highly nonconvex functions . Empirical results indicate an intimate connection between training algorithms and the performance of trained models ( Le et al. , 2011 ; Bottou et al. , 2018 ; Zhang et al. , 2021 ; Soydaner , 2020 ; Zhou et al. , 2020 ) . Especially for widely used first-order training algorithms ( e.g. , GD and SGD ) , the learning rate is of essential importance and has received extensive focus from researchers ( Smith , 2017 ; Jastrzebski et al. , 2017 ; Smith , 2018 ; Gotmare et al. , 2018 ; Liu et al. , 2019 ; Li & Arora , 2019 ) . A recent perspective is that large learning rates often lead to improved testing performance compared to the counterpart trained with small learning rates ( Smith & Topin , 2019 ; Yue et al. , 2020 ) . Towards explaining the better performance , a common belief is that large learning rates encourage the algorithm to search for flat minima , which often generalize better and are more robust than sharp ones ( Seong et al. , 2018 ; Lewkowycz et al. , 2020 ) . Despite abundant empirical observations , theoretical understandings of the benefits of large learning rate are still limited for non-convex functions , partly due to challenges in analysis . For example , the convergence ( of GD or SGD ) under large learning rate is not guaranteed . Even for globally smooth functions , very few general results exist if the learning rate exceeds certain threshold ( Kong & Tao , 2020 ) . Besides , popular regularity assumptions such as global smoothness for simplified analyses are often absent in homogeneous models , including commonly used ReLU neural networks . This paper theoretically studies the benefits of large learning rate in a matrix factorization problem min X , Y 1 2 ∥∥A−XY > ∥∥2 F , where A ∈ Rn×n , X , Y ∈ Rn×d . ( 1 ) We consider Gradient Descent ( GD ) for solving ( 1 ) : at the k-th iteration , we have Xk+1 = Xk + h ( A−XkY > k ) Yk and Yk+1 = Yk + h ( A > − YkX > k ) Xk , where h is the learning rate . Despite its simple formula , problem ( 1 ) serves as an important foundation of a variety of problems , including matrix sensing ( Chen & Wainwright , 2015 ; Bhojanapalli et al. , 2016 ; Tu et al. , 2016 ) , matrix completion ( Keshavan et al. , 2010 ; Hardt , 2014 ) , and linear neural networks ( Ji & Telgarsky , 2018 ; Gunasekar et al. , 2018 ) . Problem ( 1 ) possesses several intriguing properties . Firstly , the objective function is non-convex , and critical points are either global minima or saddles ( see e.g. , Baldi & Hornik ( 1989 ) ; Li et al . ( 2019b ) ; Valavi et al . ( 2020a ) ; Chen et al . ( 2018 ) ) . Secondly , problem ( 1 ) is homogeneous in X and Y , meaning that rescaling X , Y to aX , a−1Y for any a 6= 0 will not change the objective ’ s value . This property is shared by commonly used ReLU neural networks . A direct consequence of homogeneity is that global minima of ( 1 ) are non-isolated and can be unbounded . The curvatures at these global minima are highly dependent on the magnitudes of X , Y . When X , Y have comparable magnitudes , the largest eigenvalue of Hessian is small , and this corresponds to a flat minimum ; on the contrary , unbalanced X and Y give a sharp minimum . Last but not the least , the homogeneity impairs smoothness conditions of ( 1 ) , rendering the gradient being not Lipschitz continuous unless X , Y are bounded . See a formal discussion in Section 2 . Existing approaches for solving ( 1 ) often uses explicit regularization ( Ge et al. , 2017 ; Tu et al. , 2016 ; Cabral et al. , 2013 ; Li et al. , 2019a ) , or infinitesimal ( or diminishing ) learning rates for controlling the magnitudes of X , Y ( Du et al. , 2018 ; Ye & Du , 2021 ) . In this paper , we go beyond the scope of aforementioned works , and analyze GD with a large learning rate for solving ( 1 ) . In particular , we allow the learning rate h to be as large as approximately 4/L ( see more explanation in Section 2 ) , where L denotes the largest eigenvalue of Hessian at GD initialization . In connection to empirical observations , we provide positive answers to the following two questions : Does GD with large learning rate converge at least for some cases of ( 1 ) ? Does larger learning rate biases toward flatter minima ( i.e. , X , Y with comparable magnitudes ) ? We theoretically show the convergence of GD with large learning rate for the two situations n = 1 , d ∈ N+ or d = 1 , n ∈ N+ with isotropic A . We also observe a , perhaps surprising , “ balancing effect ” for general matrix factorization ( i.e. , any d , n , and A ) , meaning that when h is sufficiently large , the difference between X and Y shrinks significantly at the convergence of GD compared to its initial , even if the initial point is close to an unbalanced global minimum . In fact , with a proper large learning rate h , ‖Xk − Yk‖2F may decrease by an arbitrary factor at its limit . The following is a simple example of our theory for n = 1 ( i.e . scalar factorization ) , and more general results will be presented later with a precise bound for h depending on the initial condition and A. Theorem 1.1 ( Informal version of Thm.3.1 & 3.2 ) . Given scalar A and initial condition X0 , Y0 ∈ R1×d chosen almost everywhere , with learning rate h . 4/L , GD converges to a global minimum ( X , Y ) satisfying ‖X‖2F + ‖Y ‖2F ≤ 2h . Consequently , its extent of balancing is quantified by ‖X − Y ‖2F ≤ 2h − 2A . We remark that having a learning rate h ≈ 4/L is far beyond the commonly analyzed regime in optimization . Even for globally L-smooth objective , traditional theory requires h < 2/L for GD convergence and h = 1/L is optimal for convex functions ( Boyd et al. , 2004 ) , not to mention that our problem ( 1 ) is never globally L-smooth due to homogeneity . Modified equation provides a tool for probing intermediate learning rates ( see Hairer et al . ( 2006 , Chapter 9 ) for a general review , and Kong & Tao ( 2020 , Appendix A ) for the specific setup of GD ) , but the learning rate here is too large for modified equation to work ( see Appendix C ) . In fact , besides blowing up , GD with large learning rate may have a zoology of limiting behaviors ( see e.g. , Appendix B for convergence to periodic orbits under our setup , and Kong & Tao ( 2020 ) for convergence to chaotic attractors ) . Our analyses ( of convergence and balancing ) leverage various mathematical tools , including a proper partition of state space and its dynamical transition ( specifically invented for this problem ) , stability theory of discrete time dynamical systems ( Alligood et al. , 1996 ) , and geometric measure theory ( Federer , 2014 ) . The rest of the paper is organized as : Section 2 provides the background of studying ( 1 ) and discusses related works ; Section 3 presents convergence and balancing results for scalar factorization problems ; Section 4 generalizes the theory to rank-1 matrix approximation ; Section 5 studies problem ( 1 ) with arbitrary A and its arbitrary-rank approximation ; Section 6 summarizes the paper and discusses broadly related topics and future directions . 2 BACKGROUND AND RELATED WORK . Notations . ‖v‖ is ` 2 norm of a column or row vector v. ‖M‖F is the Frobenius norm of a matrix M . Sharp and flat minima in ( 1 ) We discuss the curvatures at global minima of ( 1 ) . To ease the presentation , consider simplified versions of ( 1 ) with either n = 1 or d = 1 . In this case , X and Y become vectors and we denote them as x , y , respectively . We show the following proposition characterizing the spectrum of Hessian at a global minimum . Proposition 2.1 . When n = 1 , d ∈ N+ or d = 1 , n ∈ N+ in ( 1 ) , the largest eigenvalue of Hessian at a global minimum ( x , y ) is ‖x‖2 + ‖y‖2 and the smallest eigenvalue is 0 . Homogenity implies global minimizers of ( 1 ) are not isolated , which is consistent with the 0 eigenvalue . On the other hand , if the largest eigenvalue ‖x‖2 + ‖y‖2 is large ( or small ) , then the curvature at such a global minimum is sharp ( or flat ) , in the direction of the leading eigenvector . Meanwhile , note this sharpness/flatness is an indication of the balancedness between magnitudes of x , y at a global minimum . To see this , singular value decomposition ( SVD ) yields that at a global minimum , ( x , y ) satisfies ∥∥xy > ∥∥2 F = σ2max ( A ) . Therefore , large ‖x‖ 2 + ‖y‖2 is obtained when |‖x‖ − ‖y‖| is large , i.e. , x and y magnitudes are unbalanced , and small ‖x‖2 + ‖y‖2 is obtained when balanced . Large learning rate We study smoothness properties of ( 1 ) and demonstrate that our learning rate is well beyond conventional optimization theory . We first define the smoothness of a function . Definition 2.2 ( L-smooth ) . A function f ∈ C1 defined on RN is L-smooth if for all u1 , u2 ∈ RN , ‖∇f ( u1 ) −∇f ( u2 ) ‖ ≤ L‖u1 − u2‖ . ( 2 ) If further f ∈ C2 , then∇2f LI . Moreover , if we have ( 2 ) for u1 , u2 ∈ X ⊆ RN , we call it locally L-smooth . In traditional optimization ( Nesterov , 2003 ; Polyak , 1987 ; Nesterov , 1983 ; Polyak , 1964 ; Beck & Teboulle , 2009 ) , most analyzed objective functions often satisfy ( i ) ( some relaxed form of ) convexity or strong convexity , and ( ii ) L-smoothness . Choosing a step size h < 2/L guarantees the convergence of GD to a minimum by the existing theory ( reviewed in Appendix F.3 ) . Our choice of learning rate h ≈ 4/L ( more precisely , 4/L0 ; see below ) goes beyond the classical analyses . Besides , in our problem ( 1 ) , the regularity is very different . Even simplified versions of ( 1 ) , i.e. , with either n = 1 or d = 1 , suffer from ( i ) non-convexity and ( ii ) unbounded eigenvalues of Hessian , i.e. , no global smoothness ( see Appendix F.2 for more details ) . As shown in Du et al . ( 2018 ) ; Ye & Du ( 2021 ) ; Ma et al . ( 2021 ) , decaying or infinitesimal learning rate ensures that the GD trajectory stays in a locally smooth region . However , the gap between the magnitudes of X , Y can only be maintained in that case . We show , however , that larger learning rate can shrink this gap . More precisely , if initial condition is well balanced , there is no need to use large learning rate ; otherwise , we can use learning rate as large as approximately 4/L , and within this range , larger h provides smaller gap between X and Y at the limit ( i.e . infinitely many GD iterations ) . Related work Matrix factorization problems in various forms have been extensively studied in the literature . The version of ( 1 ) is commonly known as the low-rank factorization , although here d can arbitrary and we consider both d ≤ rank ( A ) and d > rank ( A ) cases . Baldi & Hornik ( 1989 ) ; Li et al . ( 2019b ) ; Valavi et al . ( 2020b ) provide landscape analysis of ( 1 ) . Ge et al . ( 2017 ) ; Tu et al . ( 2016 ) propose to penalize the Frobenius norm of ∥∥X > X − Y > Y ∥∥2 F to mitigate the homogeneity and establish global convergence guarantees of GD for solving ( 1 ) . Cabral et al . ( 2013 ) ; Li et al . ( 2019a ) instead penalize individual Frobenius norms of ‖X‖2F + ‖Y ‖ 2 F. We remark that penalizing ‖X‖2F + ‖Y ‖ 2 F is closely related to nuclear norm regularization , since the variational formula for nuclear norm ‖Z‖2∗ = minZ=XY > ‖X‖ 2 F + ‖Y ‖ 2 F. More recently , Liu et al . ( 2021 ) consider using injected noise as regularization to GD and establish a global convergence ( see also Zhou et al . ( 2019 ) ; Liu et al . ( 2022 ) . More specifically , by perturbing the GD iterate ( Xk , Yk ) with Gaussian noise , GD will converge to a flat global optimum . On the other hand , Du et al . ( 2018 ) ; Ye & Du ( 2021 ) ; Ma et al . ( 2021 ) show that even without explicit regularization , when the learning rate of GD is infinitesimal , i.e. , GD approximating gradient flow , X and Y maintain the gap in their magnitudes . Such an effect is more broadly recognized as implicit bias of learning algorithms ( Neyshabur et al. , 2014 ; Gunasekar et al. , 2018 ; Soudry et al. , 2018 ; Li et al. , 2020 ; 2021 ) . Built upon this implicit bias , Du et al . ( 2018 ) further prove that GD with diminishing learning rates converges to a bounded global minimum of ( 1 ) , and this conclusion is recently extended to the case of a constant small learning rate ( Ye & Du , 2021 ) . Our work goes beyond the scopes of these milestones and considers matrix factorization with much larger learning rates . Additional results exist that demonstrate large learning rate can improve performance in various learning problems . Most of them involve non-constant learn rates . Specifically , Li et al . ( 2019c ) consider a two-layer neural network setting , where using learning rate annealing ( initially large , followed by small ones ) can improve classification accuracy compared to training with small learning rates . Nakkiran ( 2020 ) shows that the observation in Li et al . ( 2019c ) even exists in convex problems . Lewkowycz et al . ( 2020 ) study constant large learning rates , and demonstrate distinct algorithmic behaviors of large and small learning rates , as well as empirically illustrate large learning rate yields better testing performance on neural networks . Their analysis is built upon the neural tangent kernel perspective ( Jacot et al. , 2018 ) , with a focus on the kernel spectrum evolution under large learning rate . Worth noting is , Kong & Tao ( 2020 ) also study constant large learning rates , and show that large learning rate provides a mechanism for GD to escape local minima , alternative to noisy escapes due to stochastic gradients .

Summary: This paper studies the properties of gradient descent with large learning rate in the matrix factorization problem. The goal is to understand when gradient descent converges to a global minimum where the two factors are roughly balanced in norm, which is a shallow minimum that may be hoped to generalize better. For small learning rates (i.e. where classical optimization results guarantee convergence to a critical point), this does not happen. Hence, this paper considers learning rates beyond this regime.  For two special cases (when the matrices are 1xd or dx1), it is shown that under a large learning rate, gradient descent still converges to a global minimum, and this minimum is roughly "balanced". For the general case, it is shown that for almost any initial point, if the learning rate is such that gradient descent converges to a global minimum, then this minimum is roughly balanced.