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: In this paper, an implicit bias of SGD and GD in terms of the direction of convergence points is studied. This study shows that, in a setting of linear regression, SGD and GD converge to different directions, which are determined by the largest/smallest eigenvectors of a data matrix when the learning rate is moderately large. Experiments using synthetic data and Fashion MNIST support the theoretical results.