Direction Matters: On the Implicit Bias of Stochastic Gradient Descent with Moderate Learning Rate

1 INTRODUCTION . Stochastic gradient descent ( SGD ) and its variants play a key role in training deep learning models . From the optimization perspective , SGD is favorable in many aspects , e.g. , scalability for large-scale models ( He et al. , 2016 ) , parallelizability with big training data ( Goyal et al. , 2017 ) , and rich theory for its convergence ( Ghadimi & Lan , 2013 ; Gower et al. , 2019 ) . From the learning perspective , more surprisingly , overparameterized deep nets trained by SGD usually generalize well , even in the absence of explicit regularizers ( Zhang et al. , 2016 ; Keskar et al. , 2016 ) . This suggests that SGD favors certain “ good ” solutions among the numerous global optima of the overparameterized model . Such phenomenon is attributed to the implicit bias of SGD . It remains one of the key theoretical challenges to characterize the algorithmic bias of SGD , especially with moderate and annealing learning rate as typically used in practice ( He et al. , 2016 ; Keskar et al. , 2016 ) . In the small learning rate regime , the regularization effect of SGD is relatively well understood , thanks to the recent advances on the implicit bias of gradient descent ( GD ) ( Gunasekar et al. , 2017 ; 2018a ; b ; Soudry et al. , 2018 ; Ma et al. , 2018 ; Li et al. , 2018 ; Ji & Telgarsky , 2019b ; a ; Ji et al. , 2020 ; Nacson et al. , 2019a ; Ali et al. , 2019 ; Arora et al. , 2019 ; Moroshko et al. , 2020 ; Chizat & Bach , 2020 ) . According to classical stochastic approximation theory ( Kushner & Yin , 2003 ) , with a sufficiently small learning rate , the randomness in SGD is negligible ( which scales with learning rate ) , and as a consequence SGD will behave highly similar to its deterministic counterpart , i.e. , GD . Based on this fact , the regularization effect of SGD with small learning rate can be understood through that of GD . Take linear models for example , GD has been shown to be biased towards maxmargin/minimum-norm solutions depending on the problem setups ( Soudry et al. , 2018 ; Gunasekar et al. , 2018a ; Ali et al. , 2019 ) ; correspondingly , follow-ups show that SGD with small learning rate has the same bias ( up to certain small uncertainty governed by the learning rate ) ( Nacson et al. , 2019b ; Gunasekar et al. , 2018a ; Ali et al. , 2020 ) . The analogy between SGD and GD in the small learning rate regime is also demonstrated in Figures 1 ( a ) and 3 . However , the regularization theory for SGD with small learning rate can not explain the benefits of SGD in the moderate learning rate regime , where the initial learning rate is moderate and followed by annealing ( Li et al. , 2019 ; Nakkiran , 2020 ; Leclerc & Madry , 2020 ; Jastrzebski et al. , 2019 ) . In particular , empirical studies show that , in the moderate learning rate regime , ( small batch ) SGD generalizes much better than GD/large batch SGD ( Keskar et al. , 2016 ; Jastrzębski et al. , 2017 ; Zhu et al. , 2019 ; Wu et al. , 2020 ) ( see Figure 3 ) . This observation implies that , instead of imitating the bias of GD as in the small learning rate regime , SGD in the moderate learning rate regime admits superior bias than GD — it requires a dedicated characterization for the implicit regularization effect of SGD with moderate learning rate . In this paper , we reveal a particular regularization effect of SGD with moderate learning rate that involves convergence direction . In specific , we consider an overparameterized linear regression model learned by SGD/GD . In this setting , SGD and GD are known to converge to the unique minimumnorm solution ( Zhang et al. , 2016 ; Gunasekar et al. , 2018a ) ( see also Section 2.1 ) . However , with a moderate and annealing learning rate , we show that SGD and GD favor different convergence directions : SGD converges along the large eigenvalue directions of the data matrix ; in contrast , GD goes after the small eigenvalue directions . The phenomenon is illustrated in Figure 1 ( b ) . To sum up , we make the following contributions in this work : 1 . For an overparameterized linear regression model , we show that SGD with moderate learning rate converges along the large eigenvalue directions of the data matrix , while GD goes after the small eigenvalue directions . To our knowledge , this result initiates the regularization theory for SGD in the moderate learning rate regime , and complements existing results for the small learning rate . 2 . Furthermore , we show the particular directional bias of SGD with moderate learning rate benefits generalization when early stopping is used . This is because converging along the large eigenvalue directions ( SGD ) leads to nearly optimal solutions , while converging along the small eigenvalue directions ( GD ) can only give suboptimal solutions . 3 . Finally , our results explain several folk arts for tuning SGD hyperparameters , such as ( 1 ) linearly scaling the initial learning rate with batch size ( Goyal et al. , 2017 ) ; and ( 2 ) overrunning SGD with high learning rate even when the loss stops decreasing ( He et al. , 2016 ) . 2 PRELIMINARY . Let ( x , y ) ∈ Rd×R be a pair of d-dimensional feature vector and 1-dimensional label . We consider a linear regression problem with square loss defined as ` ( x , y ; w ) : = ( w > x− y ) 2 , where w ∈ Rd is the model parameter . LetD be the population distribution over ( x , y ) , then the test loss is LD ( w ) : = E ( x , y ) ∼D [ ` ( x , y ; w ) ] . Let S : = { ( xi , yi ) } ni=1 be a training set of n data points drawn i.i.d . from the population distributionD . Then the training/empirical loss is defined as the average of the individual loss over all training data points , LS ( w ) : = 1 n n∑ i=1 ` i ( w ) , where ` i ( w ) : = ` ( xi , yi ; w ) = ( w > xi − yi ) 2 . We use { ηk } to denote a learning rate scheme ( LR ) . Then gradient descent ( GD ) iteratively performs the following update : wk+1 = wk − ηk∇LS ( wk ) = wk − 2ηk n n∑ i=1 xi ( x > i wk − yi ) . ( GD ) Next we introduce mini-batch stochastic gradient descent ( SGD ) .1 Let b be the batch size . For simplicity suppose n = mb for an integer m ( number of mini-batches ) . Then at each epoch , SGD first randomly partitions the training set into m disjoint mini-batches with size b , and then sequentially performs m updates using the stochastic gradients calculated over the m mini-batches . Specifically , at the k-th epoch , let the mini-batch index sets be Bk1 , Bk2 , . . . , Bkm , where |Bkj | = b and⋃m j=1 Bkj = { 1 , 2 , . . . , n } , then SGD takes m updates as follows wk , j+1 = wk , j− ηk b ∑ i∈Bkj ∇ ` i ( wk , j ) = wk , j− 2ηk b ∑ i∈Bkj xi ( x > i wk , j−yi ) , j = 1 , . . . , m. ( SGD ) We also write wk+1 = wk , m+1 and wk = wk,1 to be consistent with notations in ( GD ) . 2.1 THE MINIMUM-NORM BIAS . Before presenting our results on the directional bias , let us first recap the well-known minimumnorm bias for SGD/GD optimizing linear regression problem ( Zhang et al. , 2016 ; Gunasekar et al. , 2018a ; Belkin et al. , 2019 ; Bartlett et al. , 2020 ) . We rewrite the training loss as LS ( w ) = 1 n ∥∥X > w − Y ∥∥2 2 , where X = ( x1 , . . . , xn ) ∈ Rd×n and Y = ( y1 , . . . , yn ) > ∈ Rn . Then its global minima are given byW∗ : = { w ∈ Rd : Pw = w∗ , w∗ : = X ( X > X ) −1Y } , where P is the projection operator onto the data manifold , i.e. , the column space of X . We focus on overparameterized cases whereW∗ contains multiple elements . Notice that every gradient∇ ` i ( w ) = 2xi ( x > i w−yi ) is spanned in the data manifold , thus ( GD ) and ( SGD ) can never move along the direction that is orthogonal to the data manifold . In other words , ( GD ) and ( SGD ) implicitly admit the following hypothesis class : HS = { w ∈ Rd : P⊥w = P⊥w0 } , ( 1 ) where w0 is the initialization and P⊥ = I − P is the projection operator onto the orthogonal complement to the column space of X . Putting things together , for any global optimum w ∈ W∗ ( hence Pw = w∗ ) , we have ‖w − w0‖22 = ‖Pw − Pw0‖ 2 2 + ‖P⊥w − P⊥w0‖ 2 2 = ‖w∗ − Pw0‖ 2 2 + ‖P⊥w − P⊥w0‖ 2 2 , where the right hand side is minimized when P⊥w = P⊥w0 , i.e. , w ∈ HS , thus w is the solution found by SGD/GD in the non-degenerated cases ( when the learning rate is set properly so that the algorithms can find a global optimum ) . In sum , SGD/GD is biased to find the global optimum that is closest to the initialization , which is referred as the “ minimum-norm ” bias in literature since the initialization is usually set to be zero . 1In this paper we focus on SGD without replacement , nonetheless our results and techniques are ready to be extended to SGD with replacement as well . 3 WARMING UP : A 2-DIMENSIONAL CASE STUDY . In this section we conduct a 2-dimensional case study to motivate our understanding on the directional bias of SGD in the moderate learning rate regime . Let us consider a training set consisting of two orthogonal points , S = { ( x1 , y1 = 0 ) , ( x2 , y2 = 0 ) } where x1 = √ κ · e1 = ( √ κ , 0 ) > , x2 = e2 = ( 0 , 1 ) > , κ > 2 . Clearly w∗ = 0 is the unique minimum of LS ( w ) . The Hessian of the empirical loss is∇2 LS ( w ) = x1x > 1 + x2x > 2 = diag ( κ , 1 ) , which has two eigenvalues : the smaller one 1 is contributed by data x2 , and the larger one κ contributed by data x1 . Hence LS ( w ) is κ-smooth . Similarly the Hessian of the individual losses are∇2 ` 1 ( w ) = 2x1x > 1 = diag ( 2κ , 0 ) and∇2 ` 2 ( w ) = 2x2x > 2 = diag ( 0 , 2 ) . Thus ` 2 ( w ) is 2-smooth , but ` 1 ( w ) , the individual loss for data x1 , is only 2κ-smooth , which is more ill-conditioned compared to LS ( w ) and ` 2 ( w ) . Next we consider a moderate initial learning rate η ∈ ( 1 κ , 2 1+κ ) . According to convex optimization theory ( Boyd et al. , 2004 ) , gradient step with such learning rate is convergent for LS ( w ) and ` 2 ( w ) , but oscillating for ` 1 ( w ) . In other words , ( GD ) is convergent ; and ( SGD ) is convergent along direction x2 ( or e2 ) , but oscillating along direction x1 ( or e1 ) . We also see this by analytically solving ( GD ) and ( SGD ) for this example : wgdk = ( ( 1− ηκ ) k ( 1− η ) k ) w0 , w sgd k = ( ( 1− 2ηκ ) k ( 1− 2η ) k ) w0 , ( 2 ) where |1− ηκ| < |1− η| < 1 and |1− 2η| < 1 < |1− 2ηκ| . By Eq . ( 2 ) , with moderate learning rate GD is convergent for both directions e1 and e2 . Moreover , GD fits e1 faster since the contraction parameter is smaller , i.e. , |1− ηκ| < |1− η| < 1 . Thus observing the entire optimization path , GD approaches the minimum w∗ = 0 along e2 , which corresponds to the smaller eigenvalue direction of ∇2 LS ( w ) . This is verified by the blue dots in Figure 1 ( b ) . We note this directional bias for GD also holds in the small learning rate regime , as shown in Figure 1 ( a ) . As for SGD in the initial phase where the learning rate is moderate , Eq . ( 2 ) shows it converges along e2 but oscillates along e1 since |1− 2η| < 1 < |1− 2ηκ| . In other words , SGD can not fit e1 before the learning rate decays ; however when this happens , e2 is already well fitted . Overall , SGD fits e2 first then fits e1 , i.e. , SGD converges to the minimum w∗ = 0 along e1 , which corresponds to the larger eigenvalue direction of∇2 LS ( w ) . This is verified by the red dots in Figure 1 ( b ) . We note this particular directional bias for SGD is dedicated to the moderate learning rate regime ; in the small learning rate regime , as discussed before , SGD behaves similar to GD thus goes after the smaller eigenvalue direction , which is illustrated in Figure 1 ( a ) . The above idea can be carried over to more general cases : the training loss usually has relatively smooth curvature because of the empirical averaging ; yet some individual losses can possess bad smoothness condition , corresponding to the data points that contribute to the large eigenvalues of the Hessian/data matrix . Then with a moderate learning rate , while GD is convergent , SGD is convergent for the smooth individual losses but oscillating for the ill-conditioned individual losses . Thus SGD can only fit the latter losses after the learning rate anneals . Therefore , in the moderate learning rate regime , SGD tends to converge along the large eigenvalue directions while GD tends to go after the small eigenvalue directions . We will rigorously justify the above intuitions in the following section .

Summary: This paper analyzes the differences in the convergence of SGD and GD when using "small" and "moderate" learning rates to shed light on why SGD with "moderate" and annealing learning rates perform well in practice. Focusing on an overparametrized linear regression problem, the paper claims that SGD with a "moderate" learning converges differently (along the large eigenvalue directions) compared to SGD/GD with a "small" learning rate (along the small eigenvalue directions).  They further show analytically for this problem that there exists a learning rate schedule (with moderate learning rates) such that SGD will perform + generalize well, while there exist small learning rate schedules such that GD will perform + generalize poorly.