STOCHASTIC GRADIENT DESCENT WITH MODERATE LEARNING RATE
Year: 2021

Abstract:
Understanding the algorithmic bias of stochastic gradient descent (SGD) is one of the key challenges in modern machine learning and deep learning theory. Most of the existing works, however, focus on very small or even infinitesimal learning rate regime, and fail to cover practical scenarios where the learning rate is moderate and annealing. In this paper, we make an initial attempt to characterize the particular regularization effect of SGD in the moderate learning rate regime by studying its behavior for optimizing an overparameterized linear regression problem. In this case, SGD and GD are known to converge to the unique minimum-norm solution; however, with the moderate and annealing learning rate, we show that they exhibit different directional bias: SGD converges along the large eigenvalue directions of the data matrix, while GD goes after the small eigenvalue directions. Furthermore, we show that such directional bias does matter when early stopping is adopted, where the SGD output is nearly optimal but the GD output is suboptimal. Finally, our theory explains several folk arts in practice used for SGD hyperparameter tuning, such as (1) linearly scaling the initial learning rate with batch size; and (2) overrunning SGD with high learning rate even when the loss stops decreasing.

heading: 1 INTRODUCTION
text: 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 cannot 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).

heading: 2 PRELIMINARY
text: 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).

heading: 2.1 THE MINIMUM-NORM BIAS
text: 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.

heading: 3 WARMING UP: A 2-DIMENSIONAL CASE STUDY
text: 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 cannot 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.

heading: 4 MAIN RESULTS
text: In this section we present our main theoretical results. The proofs are deferred to Appendix B.
We specify the population distribution of (x, y) ∈ Rd×R in the following manner. (1) We consider the feature vector as x = ζ · ξ, where ζ and ξ are two independent random variables that represent the magnitude and angle of x, respectively. That is, ζ ∈ R is bounded in (0, 1], and ξ ∈ Rd obeys a sphere uniform distribution, U(Sd−1). (2) We consider a realizable setting where the label is given by y = w>∗ x, i.e., there exists a true parameter w∗ ∈ Rd that generates the label from the feature
vector2. Then the test loss is LD(w) = E(x,y)∼D [ (w − w∗)>xx>(w − w∗) ] = µ ‖w − w∗‖22 , where µ = E[ζ2]/d. For an i.i.d. generated training set S = {(xi, yi)}ni=1, the training loss and the individual losses are
LS(w) = 1 n (w − w∗) > XX> (w − w∗), `i(w) = (w − w∗)> xix>i (w − w∗) , i = 1, . . . , n,
where X = (x1, . . . , xn). We denote by P the projection operator onto the column space of X (the data manifold). For i ∈ [n], we denote λi := ‖xi‖22 = ζ2i ∈ (0, 1]. Without loss of generality, we assume {λi}i∈[n] are sorted in a descending order, i.e., λ1 ≥ λ2 ≥ · · · ≥ λn. With the these preparations, we are ready to state our main theorems.

heading: 4.1 THE DIRECTIONAL BIAS OF SGD
text: We first present Theorems 1 and 2 that characterize the different directional biases of SGD and GD in the moderate learning rate regime. Theorem 1 (The directional bias of SGD with moderate LR, informal). Suppose d ≥ poly (n)3. Denote ν = n/ √ d (which is small). Then with high probability it holds that λ1 > λ2 +Θ (ν) , λn−1 > λn + Θ (ν) , λn > Θ (ν) . Suppose the initialization is set such that x>i (w0 − w∗) 6= 0 for every i ∈ [n] 4. Consider (SGD) with the following moderate learning rate scheme
ηk =
{ η ∈ ( b
λ1−Θ(ν) , b λ2+Θ(ν)
) , k = 1, . . . , k1;
η′ ∈ ( 0, b2λ1 ) , k = k1 + 1, . . . , k2,
(3)
then for such that poly ( ) > ν, there exist k1 = O ( log 1 + k2 ) and k2 > 0 such that with high probability the output of SGD wsgd := wk2 satisfies
(1− ) · γ1 ≤ ( P (wsgd − w∗) )> ·XX> · P (wsgd − w∗) ‖P (wsgd − w∗)‖22 ≤ γ1, (4)
where γ1 is the largest eigenvalue of the data matrix XX>. Theorem 2 (The directional bias of GD with moderate or small LR, informal). Under the same conditions as Theorem 1, consider (GD) with the following moderate or small learning rate scheme
ηk ∈ ( 0, n
2λ1 + Θ (ν)
) , k = 1, . . . , k2, (5)
then for any > 0, if k2 > O ( log 1 ) , then with high probability the output of GD wgd := wk2 satisfies
γn ≤ ( P (wgd − w∗) )> ·XX> · P (wgd − w∗) ‖P (wgd − w∗)‖22 ≤ (1 + ) · γn, (6)
where γn is the smallest eigenvalue of the data matrix XX> restricted in the column space of X . Remark 1. As the Rayleigh quotient (4) (resp. (6)) converges to its maximum (resp. minimum), the vector gets closer to the eigenvector of the largest (resp. smallest) eigenvalue (Trefethen & Bau III, 1997). Thus Theorem 1 and 2 suggest that, when projected onto the data manifold, SGD and GD converge to the optimum along the largest and smallest eigenvalue direction respectively. Here we are only interested in the projection onto the data manifold, since SGD/GD cannot move along the direction that is orthogonal to the data manifold as discussed in Section 2.1. Remark 2. In Theorem 1 we use the gap between λ1 and λ2 to show a learning rate scheme such that SGD converges along the largest eigenvalue direction. This can be extended by considering the gap between λr and λr+1 and the learning rate scheme defined similarly, then SGD converges along the subspace spanned by the eigenvectors of the top r eigenvalues. Similar extension applies to Theorem 2 as well.
2This is for the conciseness of presentation. Our results can be easily generalized to linear regression with well-specified noise, i.e., noise that is independent of the feature vector.
3For two sequences {xn ≥ 0} and {yn ≥ 0}: xn = O (yn) if there exist constants C > 0 and N such that xn ≤ Cyn for every n ≥ N ; xn = Θ (yn) if xn = O (yn) and yn = O (xn). xn = o (yn) if for every > 0 there exists a positive constant N( ) > 0 such that xn ≤ Cyn for every n ≥ N( ); xn = poly (yn) if there exists large absolute constant D > 0 such that xn = Θ ( yDn ) .
4This holds with probability 1 if w0 is initialized randomly and follows, e.g., Gaussian distribution.
Remark 3. We legitimately assume b < n/2 − Θ (ν) since it is not very meaningful to discuss SGD that uses more than (roughly) half of the training set as a mini-batch. Then the learning rate schedule in (3) intersects with that in (5), i.e., (5) covers both moderate and small learning rate schemes. Their intersection determines a moderate learning rate scheme, where SGD converges along the large eigenvalue directions while GD goes after the small eigenvalue directions. This justifies the regularization effect of SGD with moderate learning rate. Remark 4. Technically, in Theorem 1 one can set b = n to include GD as a special case, so that GD also follows the large eigenvalue directions. However, the initial learning rate in (3) needs to be at least nλ1 ≥ n ignoring the small order term, which is way too large to be even numerically stable in practical big data circumstances. This observation is also directly supported by Theorem 2, where (5) specifies the range of learning rate such that GD converges along small eigenvalue directions. Note the upper bound in (5) linearly scales with n, and is large enough to include all learning rate that can be adopted in practice. Thus one cannot have a legitimate learning rate for GD to converge along the large eigenvalue directions as SGD does with moderate learning rate.
Note GD with small learning rate also converges along the small eigenvalue directions, since (5) covers the small learning rate scheme. In complement, the following Theorem 3 shows that in the small learning rate regime, SGD is imitating GD and converges along the small eigenvalue directions as well. Theorems 1, 2 and 3 together show that, converging along the large eigenvalue directions is a distinct regularization effect that is unique to SGD with moderate learning rate. Theorem 3 (The directional bias of SGD with small LR, informal). Theorem 2 applies to (SGD) with the following small learning rate scheme
ηk = η ′ ∈ ( 0, b
2λ1 + Θ (ν)
) , k = 1, . . . , k2. (7)
Experiment Figure 2 shows two experiments for verifying Theorems 1, 2 and 3. The details of the experimental setup are deferred to Appendix D. In Figure 2(a), we run SGD and GD in both moderate and small learning rate regimes, and directly compare their Rayleigh quotients as defined in (4) and (6). We can see that the Rayleigh quotient reaches its maximum for SGD with moderate learning rate, and reaches its minimum for both GD and SGD with small learning rate, which verifies Theorems 1, 2 and 3. In Figure 2(b), we run the algorithms to optimize a neural network on a subset of the FashionMNIST dataset. Since the neural network is non-convex and have multiple local minima, we compare the relative Rayleigh quotients, i.e., the Rayleigh quotients of the convergence directions divided by the maximum absolute eigenvalue of the Hessian (see Appendix D.3). Figure 2(b) shows that SGD with moderate learning rate converges along relatively large eigenvalue directions while GD/SGD with small learning rate converges along relatively small eigenvalue directions. This distinguishes the directional bias of SGD and GD in the moderate learning rate regime and provides evidence from neural network training to support our theory.

heading: 4.2 EFFECTS OF THE DIRECTIONAL BIAS
text: Next we justify the benefit of the particular directional bias of SGD with moderate learning rate.
Recall the hypothesis class HS (Eq. (1)) for SGD and GD. Then for an algorithm output walg, we have the following generalization error decomposition (Shalev-Shwartz & Ben-David, 2014),
LD(w alg)− inf
w LD(w) = LD(w alg)− inf w′∈HS LD(w ′)︸ ︷︷ ︸
∆(walg), estimation error
+ inf w′∈HS
LD(w ′)− inf
w LD(w)︸ ︷︷ ︸
approximation error
.
The approximate error is an intrinsic error determined by the hypothesis class, and is not improvable unless enlarging the hypothesis class. In contrast, the estimation error ∆(walg) is determined by the algorithm as well as its hyperparameters. Thus, in the following theorem, we use the estimation error to compare the generalization performance of the SGD and GD outputs in different learning rate regimes. Theorem 4 (Effects of the directional bias, informal). Let Wα := { w ∈ HS : LS(w) = α } be an α-level set of the training loss LS(w). Let ∆∗α := infw∈Wα ∆(w) be the minimum estimation error within the α-level set Wα. Under the same conditions as Theorems 1- 3 and assuming that b < n/2−Θ (ν), then the following holds with high probability:
• The output of (SGD) with moderate LR (3) in Theorem 1 satisfies ∆(wsgd) < (1 + ) ·∆∗α, where α is the training loss of wsgd and is a small constant; • The output of (GD) with moderate or small LR (5) in Theorem 2 satisfies ∆(wgd) > M · ∆∗α, where α is the training loss of wgd and M = (1− ) · γ1/γn is a constant larger than 1 + ; • The output of (SGD) with small LR (7) in Theorem 3 satisfies ∆(wsgd) > M ·∆∗α, where α is the training loss of wsgd and M = (1− ) · γ1/γn is a constant larger than 1 + .
Theorem 4 suggests that: (1) in the moderate learning rate regime, there is a separation between the test error of SGD and that of GD. In detail, early stopped SGD finds a nearly optimal solution thanks to its particular directional bias. In contrast, early stopped GD can only find a suboptimal one; and (2) in the small learning rate regime, however, SGD no longer admits the dedicated directional bias for moderate learning rate. Instead it behaves similarly as GD, and hence outputs suboptimal solutions when early stopping is adopted.
Remark 5. In practice it is usually intractable and unnecessary to achieve the exact global minima of the training loss; instead we often early stop the algorithm once obtaining a small enough training loss, i.e., reaching an α-level set. In this spirit, Theorem 4 compares the generalization ability of SGD with moderate learning rate vs. GD/SGD with small learning rate within a level set.
Remark 6. We note the second conclusion in Theorem 4 is also obtained by Nakkiran (2020) for a different purpose. In specific, Nakkiran (2020) show the separation between the test error of GD with “large” and annealing learning rate, and test error of GD with small learning rate. However, the “large” learning rate for GD in their analysis is linear in the training sample size and is not practical as we have discussed in Remark 4. In contrast, we show that under the practically used moderate learning rate, there is a separation between the generalization abilities of SGD and GD. To our knowledge, our work gives the first theoretical justification of the phenomenon that SGD outperforms GD when the learning rate is moderate.
Experiment Figure 3 shows the generalization performance of a neural network trained by SGD and GD in both moderate and small learning rate regimes. The setup details are included in Appendix D. We can observe that (1) SGD with moderate learning rate generalizes the best, and (2) GD and SGD with small learning rate perform similarly, but both are worse than SGD with moderate learning rate. The empirical results suggest SGD with moderate learning has certain benign regularization effect. This is explained by the distinct directional bias for SGD with moderate learning rate shown in previous theorems.
5 RELATED WORK
In this section, we review related works and discuss their similarities and differences to our work.
The implicit bias of GD The implicit bias of GD has been extensively studied in recent years. We summarize several representative results as follows. For homogeneous model with exponentially-tailed loss, GD converges along the max-margin direction (Gunasekar et al., 2017; 2018a;b; Soudry et al., 2018; Ma et al., 2018; Ji & Telgarsky, 2019a; Ji et al., 2020; Nacson et al., 2019a). For least square problem and its variants, GD is biased towards minimum-norm solution (Gunasekar et al., 2018a; Ali et al., 2019; Suggala et al., 2018). Note that this is also the foundation for the learning theory in the interpolation regime such as double descent (Belkin et al., 2019) and benign overfitting (Bartlett et al., 2020). For matrix factorization and linear network, Gunasekar et al. (2017); Li et al. (2018) show the GD solution minimizes nuclear norm in special cases; more generally, Arora et al. (2019); Ji & Telgarsky (2019b) show GD balances/aligns layers; Moroshko et al. (2020) suggest the GD bias relies on initialization scale and training accuracy. For infinite-width network, Chizat & Bach (2020) show GD finds a max-
margin classifier in a functional space. The regularization theory for GD is fruitful; While some of them can be applied to SGD with small learning rate as have discussed (Nacson et al., 2019b; Gunasekar et al., 2018a; Ali et al., 2020), none of them can be carried over to SGD in the moderate learning rate regime. As far as we know, our paper is the first to study the regularization effect of SGD with moderate learning rate.
The stability of SGD Stability is another approach to justify the generalization ability of SGD, where a stable algorithm is guaranteed to have small generalization error (Bousquet & Elisseeff, 2002). Along this line, several works show that SGD is stable under certain assumptions and therefore generalizes well (Hardt et al., 2016; Kuzborskij & Lampert, 2018; Charles & Papailiopoulos, 2018; Bassily et al., 2020). More interestingly, Charles & Papailiopoulos (2018) show a simple example where SGD is stable but GD is not, which partly explains the empirically superior performance of SGD. We also aim to justify the benefits of SGD, but take a different approach from the algorithmic regularization.
The escaping behavior of SGD A popular theory (Jastrzębski et al., 2017; Zhu et al., 2019; Simsekli et al., 2019) attributes the regularization effect of SGD to its behavior of escaping from sharp minima. These works are built upon the continuous approximation of SGD via stochastic differential equations (Li et al., 2017; Hu et al., 2017; Xu et al., 2018; Simsekli et al., 2019). However it requires a small learning rate for the approximation to hold. In contrast, our main interest is in the moderate learning rate regime. Another related work in this line is (Wu et al., 2018), where they study the dynamical stability of minima and how SGD chooses them, but they do not show a directional bias introduced in our work.
Non-small learning rate The regularization effect of non-small learning rate has received increasing attentions recently. Theoretically, Li et al. (2019); HaoChen et al. (2020) study the generalization of certain stochastic dynamics equipped with annealing learning rate scheme. However, their results cannot cover vanilla SGD. Lewkowycz et al. (2020) study the role of initial large learning rate for infinity-width network. Our work is motivated by Nakkiran (2020), which shows that a large initial learning rate helps GD generalize. This result can be recovered from our theorems; more importantly, we characterize the directional bias of SGD. Empirically, Jastrzebski et al. (2019); Leclerc & Madry (2020) investigate the impact of two-phase learning rate for training with SGD, but they do not provide any theoretical analysis.

heading: 6 DISCUSSIONS
text: Linear scaling rule Linearly enlarging the initial learning rate according to batch size, or the linear scaling rule (Goyal et al., 2017), is an important folk art for paralleling SGD with large batch size while maintaining a good generalization performance. Interestingly, the linear scaling rule arises naturally from Theorem 1, where LR (3) suggests the initial learning rate η to scale linearly with batch size b to guarantee the desired directional bias. Our theory thus partly explains the mechanism of the linear scaling rule.
Overrunning SGD with high learning rate Besides the moderate initial learning rate, another key ingredient in Theorem 1 is a sufficiently large k1, i.e., SGD needs to run with moderate learning rate for sufficiently long time (to fit small eigenvalue directions). This requirement coincidentally agrees with the folk art to overrun SGD with high learning rate. For example, see Figure 4 in (He et al., 2016): from the∼ 2×105-th to the∼ 3×105-th iteration, even the training error seems to make no progress, practitioners let SGD run with a relatively high learning rate for nearly 105 iterates. Our theory sheds light on understanding the hidden benefits in this “overrunning” phase: indeed the loss is not decreasing since SGD with high learning rate cannot fit the large eigenvalue directions, but on the other hand the overrunning lets SGD fit the small eigenvalue directions better, which in the end leads to the directional bias that SGD converges along the large eigenvalue directions according to Theorem 1. Thus overrunning SGD with high learning rate is helpful.
Key technical challenges With non-small learning rate, analyzing SGD is usually hard since measure concentration turns vacuous and as a result one cannot relate the SGD iterates to that of GD. Alternatively, we control the SGD iterates epoch by epoch, then bound their composition to characterize the long run behavior of SGD. However, controlling the epoch-wise update of SGD is highly non-trivial, since in different epochs the sequence of stochastic gradient steps varies, and they do not commute due to the issue of matrix multiplication. To overcome this difficulty, we adopt techniques from matrix perturbation theory (Horn & Johnson, 2012) and conduct an analysis in the overparameterized regime. We believe these techniques are of independent interest.

heading: 7 CONCLUSION AND FUTURE WORK
text: We characterize a distinct directional regularization effect of SGD with moderate learning rate, where SGD converges along the large eigenvalue directions of the data matrix. In contrast, neither GD nor SGD with small learning rate can achieve this effect. Moreover, we show this directional bias benefits generalization when early stopping is adopted. Finally, our theory explains several folk arts used in practice for SGD hyperparameter tuning.
As an initial attempt, our results are limited to overparameterized linear models, and we ignore other factors that may contribute to the good generalization of SGD for training neural networks, e.g., network structures and explicit regularization. It is left as a future work to extend our results to nonlinear/nonconvex models (with explicit regularization).

heading: ACKNOWLEDGEMENT
text: We would like to thank the anonymous reviewers for their helpful comments. QG is partially supported by the National Science Foundation CAREER Award 1906169, IIS-2008981 and Salesforce Deep Learning Research Award. DZ is supported by the Bloomberg Data Science Ph.D. Fellowship. VB is supported in part by NSF CAREER grant 1652257, ONR Award N00014-18-1-2364 and the Lifelong Learning Machines program from DARPA/MTO. JW is supported by ONR Award N00014-18-1-2364. The views and conclusions contained in this paper are those of the authors and should not be interpreted as representing any funding agencies.

heading: A PRELIMINARIES

heading: A.1 ADDITIONAL NOTATIONS
text: We adopt the notations and settings in main text. In addition we make the following notations.
For a vector x ∈ Rd, denote its direction as x̄ := x‖x‖2 . For simplicity assume the training data {x1, . . . , xn} are linear independent. For training data xi, i ∈ [n], we denote λi = ‖xi‖22, then by construction we have λi ∈ (0, 1]. Without loss of generality let λ1 ≥ · · · ≥ λn. We define
X := (x1, . . . , xn) ∈ Rd×n, X−1 := (x2, x3, . . . , xn) ∈ Rd×(n−1).
Then based on the above definitions, we define the following two projection operators
P = X(X>X)−1X>,
P⊥ = I − P.
Clearly for any v ∈ Rd, Pv projects v onto subspace span {x1, . . . , xn}, while P⊥v projects v onto the orthogonal complement of span {x1, . . . , xn}. Furthermore we introduce two more projection operators
P−1 = X−1(X > −1X−1) −1X>−1,
P1 = P − P−1 = I − P⊥ − P−1.
For any v ∈ Rd, P−1v projects v onto subspace span {x2, . . . , xn}, while P1v projects v into the orthogonal complement of span {x2, . . . , xn} with respect to span {x1, . . . , xn}. In the following, we often write the column space of P , which refers to { Pv : v ∈ Rd } , similarly for P−1, P1 and P⊥ as well. Clearly the column space of P is also span {x1, . . . , xn}; the column space of P−1 is also the data manifold span {x2, . . . , xn}. We highlight that the total space Rd can be decomposed as the direct sum of the column space of P−1, P1 and P⊥, i.e., I = P−1 + P1 + P⊥. By definition, it is easy to verify that
P⊥X = 0,
PX = X,
P1X = (P1x1, 0, . . . , 0) ,
P−1X = (P−1x1, x2, . . . , xn) .
Then we define the following matrices which will be repeatedly used in the subsequent proof.
H := XX>,
H−1 := (P−1X)(P−1X) >,
H1 := (P1X)(P1X) >,
Hc := (P−1x1)(P1x1) > + (P1x1)(P−1x1) >.
Based on the above definitions, it is easy to show that
H = (P1X + P−1X)(P1X + P−1X) >
= (P−1X)(P−1X) > + (P1X)(P1X) > + (P1X)(P−1X) > + (P−1X)(P1X) >
= H−1 +H1 +Hc.

heading: A.2 LEMMAS
text: We present the following theorems and lemmas as preparation for our analysis. Theorem (Gershgorin circle theorem, restated for symmetric matrix). Let A ∈ Rn×n be a symmetric matrix. Let Aij be the entry in the i-th row and the j-th column. Let
Ri(A) := ∑ j 6=i |Aij | , i = 1, . . . , n.
Consider n Gershgorin discs
Di(A) := {z ∈ R, |z −Aii| ≤ Ri(A)} , i = 1, . . . , n. The eigenvalues of A are in the union of Gershgorin discs
G(A) := n⋃ i=1 Di(A).
Furthermore, if the union of k of the n discs that comprise G(A) forms a set Gk(A) that is disjoint from the remaining n−k discs, then Gk(A) contains exactly k eigenvalues of A, counted according to their algebraic multiplicities.
Proof. See, e.g., Horn & Johnson (2012), Chap 6.1, Theorem 6.1.1.
Theorem (Hoffman-Wielandt theorem, restated for symmetric matrix). Let A,E ∈ Rn×n be symmetric. Let λ1, . . . , λn be the eigenvalues of A, arranged in decreasing order. Let λ̂1, . . . , λ̂n be the eigenvalues of A+ E, arranged in decreasing order. Then
n∑ i=1 ∣∣λ̂i − λi∣∣2 ≤ ‖E‖2F . Proof. See, e.g., Horn & Johnson (2012), Chapter 6.3, Theorem 6.3.5 and Corollary 6.3.8.
Lemma 1. Let d ≥ 4 log(2n2/δ) for some δ ∈ (0, 1). Then with probability at least 1− δ, we have |〈x̄i, x̄j〉| < ι := Õ (
1√ d
) , i 6= j.
Proof. See Section C.1.
By Lemma 1 we can assume d ≥ poly (n) such that nι is sufficiently small depends on requirements. The following two lemmas characterize the projected components of each training data onto the column space of P1, P−1, and P⊥. Lemma 2. For xj 6= x1, we have
• P−1xj = xj ;
• P1xj = 0;
• P⊥xj = 0.
Proof. These are by the construction of the projection operators. Lemma 3. Assume √ nι ≤ 1/4. With probability at least 1− δ, we have
• 0 ≤ ‖P−1x̄1‖2 ≤ 2 √ nι; • √
1− 4nι2 ≤ ‖P1x̄1‖2 ≤ 1;
• P⊥x1 = 0.
Proof. See Section C.2.
The following four lemmas characterize the spectrum of the matrices H , H−1, H1 and Hc. Lemma 4. Let γ1, . . . , γn be the n non-zero eigenvalues of H := XX> in decreasing order. then
λn − nι ≤ γ1, . . . , γn ≤ λ1 + nι. Furthermore, if there exist λr and λr+1 such that λr > λr+1 + 2nι, then
λn − nι ≤ γr+1, . . . , γn ≤ λr+1 + nι < λr − nι ≤ γ1, . . . , γr ≤ λ1 + nι.
Proof. See Section C.3.
Lemma 5. Assume λn ≥ 3nι. Consider the symmetric matrix H−1 := P−1X(P−1X)> ∈ Rd×d.
• 0 is an eigenvalue of H−1 with algebraic multiplicity being d − n + 1, and its corresponding eigenspace is the column space of P1 + P⊥.
• Restricted in the column space of P−1, the n− 1 eigenvalues of H−1 belong to
(λn − nι, λ2 + nι).
Proof. See Section C.4.
Lemma 6. Consider matrix H1 := P1X(P1X)> ∈ Rd×d. We have H1 has only one non-zero eigenvalue, which belongs to
[λ1 ( 1− 4nι2 ) , λ1].
Moreover, the corresponding eigenspace is the column space of P1, which is 1-dim.
Proof. Clearly H1 is rank-1 since the column space of P1 is 1-dim. Thus it has only one non-zero eigenvalue, which is given by
tr(H1) = n∑ i=1 ‖P1xi‖22 = ‖P1x1‖ 2 2 ∈ [λ1 ( 1− 4nι2 ) , λ1],
where the last equality follows from Lemma 3.
Lemma 7. Consider matrix Hc := (P−1x1)(P1x1)> + (P1x1)(P−1x1)> ∈ Rd×d.
‖Hc‖2 ≤ 2λ1 ‖P−1x̄1‖2 ≤ 4 √ nι.
Proof. ‖Hc‖2 ≤ 2 ‖P−1x1‖2 · ‖P1x1‖2 ≤ 2λ1 ‖P−1x̄1‖2 ≤ 4 √ nι,
where the last equality follows from Lemma 3 and λ1 ≤ 1.

heading: B MISSING PROOFS FOR THE THEOREMS IN MAIN TEXT

heading: B.1 THE DIRECTIONAL BIAS OF SGD WITH MODERATE LEARNING RATE
text: Reloading notations Let πk := { Bk1 , . . . ,Bkm } be a randomly chosen uniform m-partition of [n], where n = mb. Then the SGD iterates at the k-th epoch can be formulated as:
wk,j+1 = wk,j − 2ηk b ∑ i∈Bkj xix > i (wk,j − w∗), j = 1, . . . ,m,
where we assume that the learning rate is fixed within each epoch. Note here πk is independently and randomly chosen at each epoch. For simplicity we often ignore the epoch-indicator k, and write the uniform partition as π := {B1, . . . ,Bm} . It is clear from context that π is random over epochs. For a mini-batch Bj ∈ π, we denote H(Bj) := ∑ i∈Bj xix > i .
Considering translating the variable by
v = w − w∗,
then we can reformulate the SGD update rule as
vk,j+1 = vk,j − 2ηk b H(Bj)vk,j =
( I − 2ηk
b H(Bj)
) vk,j , j = 1, . . . ,m. (8)
Let
Mπ := m∏ j=1 ( I − 2ηk b H(Bj) )
:= ( I − 2ηk
b H(Bm)
) · ( I − 2ηk
b H(Bm−1)
) · · · ( I − 2ηk
b H(B1)
) .
Here the matrix production over a sequence of matrices { Mi ∈ Rd×d }m j=1
is defined from the left to right with descending index,
m∏ j=1 Mj := Mm ×Mm−1 × · · · ×M1.
Let vk+1 = vk,m+1 and vk = vk,1. Then we can further reformulate Eq. (8) and obtain the epochwise update of SGD as
vk+1 =
( I − 2ηk
b H(Bm)
) · ( I − 2ηk
b H(Bm−1)
) · · · ( I − 2ηk
b H(B1)
) · vk =Mπvk. (9)
In light of the notion of v, the following lemma restates the related notations of loss functions, hypothesis class, level set, and estimation error defined in Section 2 and 4. Lemma 8 (Reloading SGD notations). Regarding repramaterization v = w − w∗, we can reload the following related notations:
• Empirical loss and population loss are
LS(v) = 1
n (P1v)
>H1(P1v) + 1
n (P−1v)
>H−1(P−1v) + 1
n (Pv)>Hc(Pv),
LD(v) = µ ‖v‖22 .
• The hypothesis class is HS = { v ∈ Rd : P⊥v = P⊥v0 } .
• The α-level set is V = { v ∈ HS : LS(v) = α } .
• For v ∈ HS , the estimation error is
∆(v) = µ ‖Pv‖22 .
Moreover, ∆∗ = inf
v∈V ∆(v) =
µnα
γ1 .
Proof. See Section C.5.
Based on the above definitions, the following lemma characterizes the one-step update of SGD. Lemma 9 (One step SGD update). Consider the j-th SGD update at the k-th epoch as given by Eq. (8). Set the learning rate be constant η during that epoch. Then for j = 1, . . . ,m we have(
P1vk,j+1 P−1vk,j+1
) = ( I − 2ηb P1H(Bj)P1 − 2η b P1H(Bj)P−1
− 2ηb P−1H(Bj)P1 I − 2η b P−1H(Bj)P−1
) · ( P1vk,j P−1vk,j ) Moreover, if 1 /∈ Bj , i.e., x1 is not used in the j-the step, then(
P1vk,j+1 P−1vk,j+1
) = ( I 0 0 I − 2ηb P−1H(Bj)P−1 ) · ( P1vk,j P−1vk,j ) Proof. See Section C.6.
Eq. (9) indicates the key to analyze the convergence of vk+1 is to characterize the spectrum of the matrixMπ . In particular the following lemma bounds the spectrum ofMπ when projected onto the column space of P−1.
Lemma 10. Suppose 3nι < λn. Suppose 0 < η < bλ2+3nι . Let π := {B1, . . . ,Bm} be a uniform m partition of index set [n], where n = mb. Consider the following d× d matrix
M−1 := m∏ j=1 ( I − 2η b P−1H(Bj)P−1 ) ∈ Rd×d.
Then for the spectrum ofM>−1M−1 we have:
• 1 is an eigenvalue ofM>−1M−1 with multiplicity being d − n + 1; moreover, the corresponding eigenspace is the column space of P1 + P⊥.
• Restricted in the column space of P−1, the eigenvalues of M>−1M−1 are upper bounded by (q−1(η)) 2 < 1, where
q−1(η) := max {∣∣∣∣1− 2ηb (λ2 + nι) ∣∣∣∣+ 3nηιb , ∣∣∣∣1− 2ηb (λn − nι) ∣∣∣∣+ 3nηιb } < 1.
Proof. See Section C.7.
Consider the projections of vk onto the column space of P−1 and P1. For simplicity let
Ak := ‖P−1vk‖2 , Bk := ‖P1vk‖2 .
The following lemma controls the update of Ak and Bk.
Lemma 11 (Update rules for Ak and Bk). Suppose 3nι < λn. Suppose 0 < η < bλ2+3nι . Consider the k-th epoch of SGD iterates given by Eq. (9). Set the learning rate in this epoch to be constant η. Denote
ξ(η) := 4η √ nι
b ,
q1(η) := ∣∣∣∣1− 2ηλ1b ‖P1x̄1‖22 ∣∣∣∣ ,
q−1(η) := max {∣∣∣∣1− 2ηb (λ2 + nι) ∣∣∣∣+ 3nηιb , ∣∣∣∣1− 2ηb (λn − nι) ∣∣∣∣+ 3nηιb } < 1.
Then we have the following:
• Ak+1 ≤ q−1(η) ·Ak + ξ(η) ·Bk.
• Bk+1 ≤ q1(η) ·Bk + ξ(η) ·Ak.
• Bk+1 ≥ q1(η) ·Bk − ξ(η) ·Ak.
Proof. See Section C.8.
Note we can rephrase the update rules for Ak and Bk as( Ak+1 Bk+1 ) ≤ ( q−1(η) ξ(η) ξ(η) q1(η) ) · ( Ak Bk ) ,
where “≤” means “entry-wisely smaller than”. The following two lemmas characterize the long run behaviors of Ak and Bk with different learning rate.
Lemma 12 (The long run behavior of SGD with moderate LR). Suppose 3nι < λn, and λ2 +4nι < λ1. Suppose v0 is far away from 0. Consider the first k1 epochs of SGD iterates given by Eq. (9). Set the learning rate during this stage to be constant, i.e., ηk = η for 0 ≤ k < k1. Suppose
b
λ1 − 3 √ nι
< η < b
λ2 + 3nι .
Then for 0 < < 1 and 0 < β < β0 < B0 satisfying √ nι ≤ poly ( β) , there exists k1 ≥ O (
log 1 β
) such that
• Ak1 ≤ · β.
• Bk1 ≤ ‖Pv0‖2 · ρ k1 1 + 2 · β = poly ( 1 β ) .
• For all k = 0, 1, . . . , k1, Bk > β0.
Proof. See Section C.9.
Lemma 13 (The long run behavior of SGD with small LR). Suppose 3nι < λn, and λ2 +4nι < λ1. Suppose v0 is far away from 0. Consider another k2 − k1 epochs of SGD iterates given by Eq. (9). Set the learning rate to be constant during the updates, i.e., ηk = η′ for k1 ≤ k < k2. Suppose
0 < η′ < b
2λ1 .
Consider the and β given in Lemma 12. Then for k ≥ k1, we have
• Ak ≤ · β.
• Bk ≤ { q ·Bk−1, Bk−1 > β, β, Bk−1 < β. where q ∈ (0, 1) is a constant.
Proof. See Section C.10.
Theorem 5 (Theorem 1, formal version). Suppose 3nι < λn and λ2 + 4nι < λ1. Suppose v0 is away from 0. Consider the SGD iterates given by Eq. (9) with the following moderate learning rate scheme
ηk = η ∈ ( b λ1−3 √ nι , bλ2+3nι ) , k = 1, . . . , k1; η′ ∈ (
0, b2λ1
) , k = k1 + 1, . . . , k2.
Then for 0 < < 1 such that √ nι ≤ poly ( ) , there exist k1 > O ( log 1 ) and k2 such that
(1− ) · γ1 ≤ v>k2Hvk2
‖Pvk2‖ 2 2
≤ γ1.
Proof. We choose k1 and k2 as in Lemma 12 and Lemma 13 with β set as a small constant, then we are guaranteed to have
Ak2 ≤ · β ≤ ·Bk2 ,
from where we have ‖P1vk2‖ 2 2
‖Pvk2‖ 2 2
= B2k2
A2k2 +B 2 k2
≥ 1 1 + 2 ≥ 1− 2.
Then we have
v>k2Hvk2 ‖Pvk2‖ 2 2 = (P1vk2) >H1(P1vk2) ‖Pvk2‖ 2 2 + (P−1vk2) >H−1(P−1vk2) ‖Pvk2‖ 2 2 + (Pvk2) >Hc(Pvk2) ‖Pvk2‖ 2 2
≥ λ1 ( 1− 4nι2 ) · ‖P1vk2‖ 2 2
‖Pvk2‖ 2 2
+ 0− 4 √ nι
≥ λ1(1− 4nι2) · (1− 2)− 4 √ nι ≥ (γ1 − nι)(1− 4nι2) · (1− 2)− 4 √ nι (since γ1 ≤ λ1 + nι by Lemma 4) = γ1(1− 4nι2)(1− 2)− nι(1− 4nι2)(1− 2)− 4 √ nι ≥ γ1(1− 0.5 )− 0.5γ1 (since √ nι ≤ poly ( ))
= γ1(1− ).
Theorem 6 (Theorem 4 first part, formal version). Suppose 3nι < λn and λ2 + 4nι < λ1. Suppose v0 is away from 0. Consider the SGD iterates given by Eq. (9) with the following moderate learning rate schedule
ηk = η ∈ ( b λ1−3 √ nι , bλ2+3nι ) , k = 1, . . . , k1; η′ ∈ (
0, b2λ1
) , k = k1 + 1, . . . , k2.
Then for 0 < < 1 satisfying √ nι ≤ poly ( ) , there exist k1 and k2 such that SGD outputs an
-optimal solution.
Proof. We set
β =
√ nα
γ1 , (10)
β0 =
√ nα
γn > β, (11)
and apply Lemma 12 to choose a k1 such that ‖P−1vk1‖2 ≤ · β = · √ nα
γ1 ; (12) ‖P1vk‖2 ≥ β0 = √ nα
γn , ∀ 0 ≤ k ≤ k1. (13)
Thus for all 0 ≤ k ≤ k1,
LS(vk) = 1
n (Pvk)
>XX>(Pvk)
≥ γn n ‖Pvk‖22 (γn is the smallest eigenvalue of XX > in the column space of P ) ≥ γn n ‖P1vk‖22
> α, (by Eq. (13))
which implies SGD cannot reach the α-level set during the iteration of first stage, i.e., SGD does not terminate in this stage.
We thus consider the second stage. From Lemma 13 we know ‖P1vkn‖2 will keep decreasing before being smaller than β, and ‖P−1vk‖2 stays small during this period, i.e., SGD fits P1v while in the same time does not mess up P−1v. Mathematically speaking, there exists k2 and α such that
Ak2 := ‖P−1vk2‖2 ≤ · β = · √ nα
γ1 ,
LS(vk2) = α,
which implies SGD terminates at the k2-th epoch. Then by Lemmas 3 and 8, we have
nα = nLS(vk2)
= (P1vk2) >H1(P1vk2) + (P−1vk2) >H2(P−1vk2) + (Pvk2) >Hc(Pvk2) ≥ (P1vk2n)>H1(P1vk2n)− ‖P−1x̄1‖ 2 2 · ‖Pvk2n‖ 2 2 ≥ (λ1 − nι)B2k2 − 4nι 2(A2k2 +B 2 k2) ≥ (γ1 − 3nι)B2k2 − 4nι 2A2k2 ,
which yields
B2k2 ≤ nα+ 4nι2A2k2 γ1 − 3nι
≤ ( 1 +
2 ) · nα γ1 .
Then we can bound the estimation error as
∆(vk2) = µ ‖Pvk2‖ 2 2
= µ(B2k2 +A 2 k2) ≤ ( 1 +
2 ) · µnα γ1 + 2 · µnα γ1
≤ (1 + ) · µnα γ1
= (1 + ) ·∆∗, where we use the fact that ∆∗ = µnα/γ1 from Lemma 8. Hence SGD is -near optimal.

heading: B.2 THE DIRECTIONAL BIAS OF GD WITH MODERATE OR SMALL LEARNING RATE
text: Reloading notations Denote the eigenvalue decomposition of XX> as
XX> = GΓG>, Γ := diag (γ1, . . . , γn, 0, . . . , 0) , G = (g1, . . . , gn, . . . , gd) ,
where G ∈ Rd×d is orthonormal, and γ1, . . . , γn are given by Lemma 4. Clearly, span {g1, . . . , gn} = span {x1, . . . , xn} . Let
G‖ = (g1, . . . , gn) , G⊥ = (gn+1, . . . , gd) ,
then P = G‖G > ‖ , P⊥ = G⊥G > ⊥.
Recall the GD iterates at the k-th epoch:
wk+1 = wk − 2ηk n XX> (wk − w∗) .
Considering translating then rotating the variable as,
u = G>(w − w∗), then we can reformulate the GD iterates as
uk+1 = uk − 2ηk n Γuk =
( I − 2ηk
n Γ
) uk. (14)
We present the following lemma to reload the related notations regarding the parameterization u = G>(w − w∗). Lemma 14 (Reloading GD notations). Regarding reparametrization u = G>(w − w∗), we can reload the following related notations:
• Empirical loss and population loss are
LS(u) = 1
n n∑ i=1 γi ( u(i) )2 , LD(u) = µ ‖u‖22 .
• The hypothesis class is HS = { u ∈ Rd : u(i) = u(i)0 , for i = n+ 1, . . . , d } .
• The α-level set is U = { u ∈ Hu : LS(u) = α } .
• For u ∈ HS , the estimation error is
∆(u) = µ n∑ i=1 ( u(i) )2 .
Moreover, ∆∗ = µnα
γ1 .
Proof. See Section C.11.
The following lemma sovles GD iterates in Eq. (14). Lemma 15. For t = 0, . . . , T ,
u (i) k =
{∏k−1 t=0 ( 1− 2ηtγin ) · u(i)0 , 1 ≤ i ≤ n;
u (i) 0 , n+ 1 ≤ i ≤ d.
Proof. This is by directly solving Eq. (14) where Γ is diagonal.
Theorem 7 (Theorem 2, formal version). Suppose λn + 2nι < λn−1. Suppose u0 is away from 0. Consider the GD iterates given by Eq. (14) with learning rate scheme
ηk ∈ ( 0, n
2λ1 + 2nι
) .
Then for ∈ (0, 1), if k ≥ O ( log 1 ) , then we have
γn ≤ u>k Γuk∑n i=1 ( u (i) k )2 ≤ (1 + ) · γn. Proof. For i = 1, . . . , n, denote qi(η) = 1 − 2γin · η, where η ∈ ( 0, n2λ1+2nι ) . Then we have 0 < qi(η) < 1 since η < n
2(λ1 + nι) <
n 2γ1 ≤ n 2γi ,
where the second inequality follows from γ1 < λ1 +nι by Lemma 4. Furthermore, since λn+nι < λn−1 − nι, Lemma 4 gives us
0 < γn < γn−1 ≤ · · · ≤ γ1 < 1, (15)
which implies 1 > qn(η) > qn−1(η) ≥ · · · ≥ q1(η) > 0. (16)
Moreover,
f(η) := qn−1(η)
qn(η) = 1− 2γn−1n η 1− 2γnn η
is increasing, let q = maxη< n2λ1+2nι f(η), then q < 1 by our assumption on the learning rate.
From Lemma 15 we have
u (i) k = k−1∏ t=0 qi(ηt) · u(i)0 , i = 1, . . . , n. (17)
By the assumption that
k > 1
2 ·
log γn (u
(n) 0 ) 2 γ1n ∑n i=1 ( u (i) 0 )2 log q = O ( 1 ) , (18)
we have ∑n i=1 ( u (i) k )2 (u
(n) k )
2 = 1 + n−1∑ i=1 (u (i) k ) 2 (u (n) k ) 2
= 1 + n−1∑ i=1 ∏k−1 t=0 qi(ηt) 2 · (u(i)0 )2∏k−1 t=0 qn(ηt) 2 · (u(n)0 )2 (by Eq. (17))
≤ 1 +
∑n i=1 ( u (i) 0 )2 (u
(n) 0 ) 2 · n−1∑ i=1 k−1∏ t=0 qi(ηt) 2 qn(ηt)2
≤ 1 +
∑n i=1 ( u (i) 0 )2 (u
(n) 0 )
2 · n · k−1∏ t=0 qn−1(ηt) 2 qn(ηt)2 (by Eq. (16))
≤ 1 +
∑n i=1 ( u (i) 0 )2 (u
(n) 0 )
2 · n · q2k
≤ 1 + γn γ1 , (by Eq. (18))
which further yields
1 ≥ (u
(n) k ) 2∑n i=1 ( u (i) k )2 ≥ 11 + γnγ1 ≥ 1− γnγ1 . (19) By the above inequalities we have
u>k Γuk∑n i=1 ( u (i) k )2 = n∑ i=1
(u (i) k ) 2∑n i=1 ( u (i) k )2 · γi = (u (n) k )
2∑n i=1 ( u (i) k )2 · γn + n−1∑ i=1
(u (i) k ) 2∑n i=1 ( u (i) k )2 · γi ≤ γn +
n−1∑ i=1
(u (i) k ) 2∑n i=1 ( u (i) k )2 · γ1 (by Eq. (15)) = γn +
1− (u(n)k )2∑n i=1 ( u (i) k )2  · γ1
≤ γn + γn γ1 · γ1 (by Eq. (19)) = γn · (1 + ).
Finally we note that u > k Γuk∑n
i=1 ( u (i) k )2 ≥ γn since γn is the smallest in {γi}ni=1. Theorem 8 (Theorem 4 second part, formal version). Suppose λn + 2nι < λn−1. Suppose u0 is away from 0. Consider the GD iterates given by Eq. (14) with learning rate scheme
ηk ∈ ( 0, n
2λ1 + 2nι
) .
Then for ∈ (0, 1), if k ≥ O ( log 1 ) ,, then GD outputs an M -suboptimal solution, where M = γ1 γn (1− ) > 1 is a constant.
Proof. Consider an α-level set where
α = LS(uk) = 1
n u>k Γuk. (20)
From Lemma 15 we know LS(uk) is monotonic decreasing thus GD cannot terminate before the k-epoch, i.e., the output of GD is uk.
Thus
∆(u)
∆∗ = γ1
∑n i=1 ( u (i) k )2 nα
(by Lemma 14)
= γ1 ·
∑n i=1 ( u (i) k )2 u>k Γuk
(by Eq. (20))
≥ γ1 · 1
(1 + )γn (by Theorem 7)
≥ γ1 γn (1− )
=: M,
where we have M > 1 by letting < 1− γnγ1 .

heading: B.3 THE DIRECTIONAL BIAS OF SGD WITH SMALL LEARNING RATE
text: We analyze SGD with small learning rate by repeating the arguments in previous two sections.
Let us denote X−n := (x1, x2, . . . , xn−1) and
P−n = X−n(X > −nX−n) −1X>−n
Pn = P − P−n. That is, P−n is the projection onto the column space of X−n and Pn is the projection onto the orthogonal complement of the column space of X−n with respect to the column space of X .
Let us reload
H := XX>,
H−n := (P−nX)(P−nX) >,
Hn := (PnX)(PnX) >,
Hc := (P−nxn)(Pnxn) > + (Pnxn)(P−nxn) >.
Then H = H−n +Hn +Hc.
Following a routine check we can reload the following lemmas. Lemma 16 (Variant of Lemma 10). Suppose 3nι < λn. Suppose 0 < η < bλ1+3nι . Let π := {B1, . . . ,Bm} be a uniform m partition of index set [n], where n = mb. Consider the following d× d matrix
M−n := m∏ j=1 ( I − 2η b P−nH(Bj)P−n ) ∈ Rd×d.
Then for the spectrum ofM>−nM−n we have:
• 1 is an eigenvalue ofM>−nM−n with multiplicity being d− n+ 1; moreover, the corresponding eigenspace is the column space of Pn + P⊥.
• Restricted in the column space of P−n, the eigenvalues of M>−nM−n are upper bounded by (q−n(η)) 2 < 1, where
q−n(η) := max {∣∣∣∣1− 2ηb (λ1 + nι) ∣∣∣∣+ 3nηιb , ∣∣∣∣1− 2ηb (λn−1 − nι) ∣∣∣∣+ 3nηιb } < 1.
Proof. This is by a routine check of the proof of Lemma 10.
Consider the projections of vk onto the column space of P−n and Pn. For simplicity we reload the following notations
Ak := ‖P−nvk‖2 , Bk := ‖Pnvk‖2 .
The following lemma controls the update of Ak and Bk.
Lemma 17 (Variant of Lemma 11). Suppose 3nι < λn. Suppose 0 < η < bλ1+3nι . Consider the k-th epoch of SGD iterates given by Eq. (9). Set the learning rate in this epoch to be constant η. Denote
ξ(η) := 4η √ nι
b ,
qn(η) := ∣∣∣∣1− 2ηλnb ‖Pnx̄n‖22 ∣∣∣∣ ,
q−n(η) := max {∣∣∣∣1− 2ηb (λ1 + nι) ∣∣∣∣+ 3nηιb , ∣∣∣∣1− 2ηb (λn−1 − nι) ∣∣∣∣+ 3nηιb } < 1.
Then we have the following:
• Ak+1 ≤ q−n(η) ·Ak + ξ(η) ·Bk.
• Bk+1 ≤ qn(η) ·Bk + ξ(η) ·Ak.
• Bk+1 ≥ qn(η) ·Bk − ξ(η) ·Ak.
Proof. This is by a routine check of the proof of Lemma 11.
Lemma 18 (Variant of Lemma 13). Suppose 3nι < λn and λn + 4nι < λn−1. Consider the SGD iterates given by Eq. (9) with the following small learning rate scheme
ηk = η ′ ∈ ( 0, b
2λ1 + 2nι
) , k = 1, . . . , k2.
Then for 0 < < 1 satisfying √ nι ≤ poly ( ) , if k2 ≥ O ( log 1 ) , then Ak2Bk2 ≤ .
Proof. From the assumption we have η′ < b2(λ1+nι) and η ′ < b2λn , thus
ξ := ξ(η′) = 4η′ √ nι
b ,
qn := qn(η ′) = ∣∣∣∣1− 2η′λnb ‖Pnx̄n‖22 ∣∣∣∣
= 1− 2η ′λn b ‖Pnx̄n‖22 ≤ 1− 2λn(1− 4nι 2)
b η′, (since ‖Pnx̄n‖22 ≥ 1− 4nι 2 by reloading Lemma 3 )
< 1
q−n := q−n(η ′) = max {∣∣∣∣1− 2η′b (λ1 + nι) ∣∣∣∣+ 3nη′ιb , ∣∣∣∣1− 2η′b (λn−1 − nι) ∣∣∣∣+ 3nη′ιb } = max { 1− 2η ′
b (λ1 + nι) +
3nη′ι
b , 1− 2η
′
b (λn−1 − nι) +
3nη′ι
b } = 1− 2η ′
b (λn−1 − nι) +
3nη′ι
b
= 1− 2(λn−1 − nι)− 3nι b η′ ∈ (0, 1).
Moreover, by the gap assumption λn + 4nι < λn−1 we have qn − q−n ≥ η′ (
2(λn−1 − nι)− 3nι b − 2λn(1− 4nι 2) b ) ≥ 2η ′
b (λn−1 − λn − 3nι)
> 0.
Therefore 0 < q−n < qn < 1. Thus we can set ξ = 4η′ √ nι
b to be small such that
0 < q := q−n
qn − ξ ·A0/B0 < 1. (21)
Moreover, since √ nι ≤ poly ( ) and ξ = 4η ′√nι b , we have
ξ qn − ξ ·A0/B0 ≤ (1− q) 2 . (22)
Now we recursively show AkBk ≤ A0 B0 . Clearly it holds for k = 0. Suppose AkBk ≤ A0 B0 , we consider Ak+1 Bk+1 in the following
Ak+1 Bk+1 ≤ q−n ·Ak + ξ ·Bk qn ·Bk − ξ ·Ak
(by Lemma 17)
= q−n · AkBk + ξ qn − ξ · AkBk
≤ q−n
Ak Bk + ξ
qn − ξ ·A0/B0 (by inductive assumption)
= q−n qn − ξ ·A0/B0 Ak Bk + ξ qn − ξ ·A0/B0
≤ q · Ak Bk + (1− q) 2 (by Eq. (21) and (22))
≤ q · A0 B0 + (1− q) 2
≤ A0 B0 ,
where in the last inequality we assume 2 < A0 B0 .
Moreover, from the above we have
Ak+1 Bk+1 ≤ q · Ak Bk + (1− q) 2 ,
which implies
Ak2 Bk2 ≤ qk2 · A0 B0 + 1 1− q · (1− q) 2 ,
≤ 2 + 2 = ,
where we set k2 ≥ O ( log 1 ) .
Next we prove the directional bias of SGD with small learning rate. Theorem 9 (Theorem 3, formal version). Suppose 3nι < λn and λn + 4nι < λn−1. Suppose v0 is away from 0. Consider the SGD iterates given by Eq. (9) with the following small learning rate scheme
ηk = η ′ ∈ ( 0, b
2λ1 + 2nι
) , k = 1, . . . , k2.
Then for 0 < < 1 satisfying √ nι ≤ poly ( ) , if k2 ≥ O ( log 1 ) , then
γn ≤ v>k2Hvk2
‖Pvk2‖ 2 2
≤ (1 + ) · γn.
Proof. First by Lemma 18 we have
B2k2 A2k2 +B 2 k2 = 1
A2k2 B2k2
+ 1 ≥ 1 2 + 1 ≥ 1− 2. (23)
Next by H = Hn +H−n +Hc we obtain
v>k2Hvk2 ‖Pvk2‖ 2 2 = (Pnvk2) >Hn(Pnvk2) ‖Pvk2‖ 2 2 + (P−nvk2) >H−n(P−nvk2) ‖Pvk2‖ 2 2 + (Pvk2) >Hc(Pvk2) ‖Pvk2‖ 2 2
≤ λn · ‖Pnvk2‖ 2 2
‖Pvk2‖ 2 2
+ (λ1 + nι) · ‖P−nvk2‖ 2 2
‖Pvk2‖ 2 2
+ 4 √ nι by reloading Lemma 5, 6, 7
≤ λn + (λ1 + nι) · A2k2
A2k2 +B 2 k2
+ 4 √ nι
≤ γn + nι+ (λ1 + nι) · 2 + 4 √ nι by reloading Lemma 4 and Eq. (23) ≤ γn + γn · . since √ nι ≤ poly ( )
Finally, we note v>k2 Hvk2
‖Pvk2‖ 2
2
≥ γn since γn is the smallest eigenvalue of H restricted in the column
space of P .
Theorem 10 (Theorem 4 third part, formal version). Suppose 3nι < λn and λn + 4nι < λn−1. Suppose v0 is away from 0. Consider the SGD iterates given by Eq. (9) with the following small learning rate scheme
ηk = η ′ ∈ ( 0, b
2λ1 + 2nι
) , k = 1, . . . , k2.
Then for 0 < < 1 such that √ nι ≤ poly ( ) , if k2 ≥ O ( log 1 ) , then SGD outputs an M - suboptimal solution where M = γ1γn (1− ) > 1 is a constant.
Proof. From Eq. (9) and η′ < 12λ1 we know that restricted in the column space of P , the eigenvalues ofMπ is smaller than 1, thus vk indeed converges to 0. Consider an α-level set where
α = LS(vk) = 1
n v>k2Hvk2 . (24)
Then ∆(u)
∆∗ = γ1 ‖Pvk‖22 nα
(by Lemma 8)
= γ1 · ‖Pvk‖22 v>k Hvk
(by Eq. (24))
≥ γ1 · 1
(1 + )γn (by Theorem 9)
≥ γ1 γn (1− )
=: M,
where we have M > 1 by letting < 1− γnγ1 .

heading: C PROOF OF AUXILIARY LEMMAS IN SECTIONS A AND B

heading: C.1 PROOF OF LEMMA 1
text: Proof of Lemma 1. Note that x̄i follows uniform distribution on the sphere Sd−1. Therefore, let ξ be a random variable following distribution χ2d distribution and define zi = ξ · x̄i, we have zi follows standard normal distribution in the d-dimensional space. Then it suffices to prove that |〈zi, zj〉|/(‖zi‖2‖zj‖2) ≤ ι for all i 6= j. First we will bound the inner product 〈zi, zj〉. Note that we have each entry in zi is 1-subgaussian, it can be direcly deduced that
〈zi, zj〉 = d∑ k=1 z (k) i z (k) j = d∑ k=1 (z(k)i + z(k)j 2 )2 − ( z (k) i − z (k) j 2 )2 is d-subexponential, where z(k)i denotes the k-th of the vector zi. Then if follows that
P (|〈zi, zj〉| ≥ t) ≤ 2 exp ( − t 2
d
) .
Next we will lower bound ‖zi‖2. Note that
‖zi‖22 − d = d∑ k=1 (( z (k) i )2 − 1 ) .
Since z(k)i is 1-subgaussian, we have ‖zi‖ 2 − d is d-subexpoential, then P (∣∣∣‖zi‖22 − d∣∣∣ ≥ t) ≤ 2 exp(− t2d ) .
Finally, applying the union bound for all possible i, j ∈ [n], we have with probability at least 1− δ, the following holds for all i 6= j,
|〈zi, zj〉| ≤ √ d log 2n2
δ , ‖zi‖22 ≥ d− √ d log 2n2
δ .
Assume d ≥ 4 log(2n2/δ), we have ‖zi‖2 ≥ d/2. Then it follows that
|〈x̄i, x̄j〉| = |〈zi, zj〉| ‖zi‖2‖zj‖2 < 2
√ 1
d log
2n2
δ =: ι.
This completes the proof.

heading: C.2 PROOF OF LEMMA 3
text: Proof of Lemma 3. Similar to the proof of Lemma 1, we consider translating x1, . . . , xn to z1, . . . , zn by introducing χ2d random variables. Let Z−1 = (z2, . . . , zn) ∈ Rd×(n−1), in which each entry is i.i.d. generated from Gaussian distribution N (0, 1). Then we have
P−1x̄1 = X−1(X > −1X−1) −1X>−1x̄1 = Z−1(Z > −1Z−1) −1Z>−1x̄1.
Then conditioned on x̄1, we have each entry in Z>−1x̄1 i.i.d. follows N (0, 1). Then it is clear that ‖Z>−1x̄1‖22 follows from χ2n−1 distribution, implying that with probability at least 1− δ′, we have
‖Z>−1x̄1‖22 ≤ (n− 1) + √
(n− 1) log(2/δ′). Then by Corollary 5.35 in Vershynin (2010), we know that with probability at least 1 − δ′ it holds that√
d− √ n− 1− √ 2 log(2/δ′) ≤ σmin(Z−1) ≤ σmax(Z−1) ≤ √ d+ √ n− 1 + √ 2 log(2/δ′).
Therefore, assume √ (n− 1) + √
2 log(2/δ′) ≤ √ d/8, we have with probability at least 1− δ′∥∥Z−1(Z>−1Z−1)−1∥∥2 ≤ √ d+ √ n− 1 + √ 2 log(2/δ′)(√
d− √ n− 1− √ 2 log(2/δ′) )2 ≤ 1√
d
( 1 + 4 (√ n− 1 d + √ 2 log(2/δ′) d )) .
Combining with the upper bound of ‖Z>−1x̄1‖2, set δ′ = δ/2, we have with probability at least 1− δ that ‖P−1x̄1‖2 ≤
∥∥Z−1(Z>−1Z−1)−1∥∥2 · ‖X>−1x̄1‖2 ≤ ( 1 + 4 (√ n− 1 d + √ 2 log(4/δ) d )) · √n− 1 d + √√ (n− 1) log(4/δ) d
 ≤ (1 + 4 √ nι) · √ nι,
where the last inequality follows from the definition of ι. Then assume √ nι ≤ 1/4, we are able to completes the proof of the first argument. Note that ‖P1x̄1‖22 + ‖P−1x̄1‖22 = ‖x̄1‖22 = 1, we have
‖P−1x̄1‖2 = √ 1− ‖P1x̄1‖22 ≥ √ 1− 4nι2 ≥ 1− 4nι2.
This completes the proof of the second argument. The third argument holds trivially by the construction of P⊥.

heading: C.3 PROOF OF LEMMA 4
text: Proof of Lemma 4. Clearly XX> ∈ Rd×d is of rank n and symmetric, thus XX> has n real, nonzero (potentially repeated) eigenvalues, denoted as γ1, . . . , γn in non-decreasing order. Moreover, γ1, . . . , γn are also eigenvalues of X>X ∈ Rn×n, thus it is sufficient to locate the eigenvalues of X>X , where ( X>X ) ij = x>i xj . We first calculate the diagonal entry ( X>X ) ii = x>i xi = λi. Then we bound the off diagonal entries. For j 6= i,( X>X ) ij = x>i xj = √ λiλj〈x̄i, x̄j〉 ∈ (−ι, ι) ,
where we use 0 < λ1, . . . , λn ≤ 1. Thus we have
Ri(X >X) = ∑ j 6=i ∣∣∣(X>X) ij ∣∣∣ ≤ nι, i = 1, . . . , n, Finally our conclusions hold by applying Gershgorin circle theorem.

heading: C.4 PROOF OF LEMMA 5
text: Proof of Lemma 5. The first conclusion is clear since by construction, we have P−1P1 = P−1P⊥ = 0.
Note that H−1 is a rank n− 1 symmetric matrix. Let τ2, . . . , τn be the n− 1 non-zero eigenvalues of H−1. Clearly, τ2, . . . , τn with τ1 := 0 give the spectrum of
H ′−1 := (P−1X) >P−1X ∈ Rn×n.
We then bound τ2, . . . , τn by analyzing H ′−1. From Lemma 3 we have ‖P−1x̄1‖2 ≤ 2 √ nι. From Lemma 2 we have
P−1X = (P−1x1, P−1x2, . . . , P−1xn) = (P−1x1, x2, . . . , xn) .
Then we calculate the diagonal entries:
( H ′−1 ) ii = { ‖P−1x1‖22 ≤ λ1 · 4nι2 ≤ 4nι2, i = 1; ‖xi‖22 = λi, i 6= 1.
Then we bound the off diagonal entries. Let j 6= i. Then at least one of them is not 1. Without loss of generality let i 6= 1, which yields xi = P−1xi by Lemma (2). Thus 〈xi, P1xj〉 = 〈P−1xi, P1xj〉 = 0. Thus we have (
H ′−1 ) ij = (P−1xi) >P−1xj
= x>i P−1xj = x>i xj − x>i P1xj = x>i xj
= √ λiλj · 〈x̄i, x̄j〉
∈ (−ι, ι) .
Thus we have Ri(H ′ −1) = ∑ j 6=i ∣∣∣(H ′−1)ij∣∣∣ ≤ nι, i = 1, . . . , n. Finally, we set 4nι2 + 2nι < λn, so that the first Geoshgorin disc does not intersect with the others, then Gershgorin circle theorem gives our second conclusion.

heading: C.5 PROOF OF LEMMA 8
text: Proof of Lemma 8. For the empirical loss, it is clear that
LS(v) = 1 n (w − w∗)>XX> (w − w∗) = 1 n v>XX>v = 1 n v>Hv
= 1
n (Pv)>H(Pv)
= 1
n (Pv)>H1(Pv) +
1 n (Pv)>H−1(Pv) + 1 n (Pv)>Hc(Pv)
= 1
n (P1v)
>H1(P1v) + 1
n (P−1v)
>H−1(P−1v) + 1
n (Pv)>Hc(Pv),
where we use Lemma 4, Lemma 5, and Lemma 6. For the population loss,
LD(v) = µ ‖w − w∗‖22 = µ ‖v‖ 2 2 .
For the hypothesis classHS = { w ∈ Rd : P⊥w = P⊥w0 } , applyingw−w∗ = v andw0−w∗ = v0, we obtain HS = { v ∈ Rd : P⊥v = P⊥v0 } .
For the α-level set, we note the optimal training loss is L∗S = infv∈HS LS(v) = 0.
As for the estimation error, we note that infv∈HS LD(v) = infP⊥v=P⊥v0 µ ‖v‖ 2 2 = µ ‖P⊥v0‖ 2 2 . thus for v ∈ HS , we have
∆(v) = LD(v)− inf v′∈V LD(v ′) = µ ‖v‖22 − µ ‖P⊥v0‖ 2 2 = µ ‖Pv‖ 2 2 .
Finally, consider v ∈ V , i.e., nα = v>XX>v, thus
∆∗ = inf v∈V ∆(v) = inf nα=v>XX>v
µ ‖Pv‖22 = µnα
γ1 ,
where γ1 is the largest eigenvalue of the matrix XX> and the inferior is attended by setting v parallel to the first eigenvector of XX>.

heading: C.6 PROOF OF LEMMA 9
text: Proof of Lemma 9. From Eq. (8) we have
vk,j+1 =
( I − 2η
b H(Bj)
) vk,j , j = 1, . . . ,m. (25)
Recall the following property of projection operators:
P1 = P1P1, P−1 = P−1P−1
0 = P1P−1 = P−1P1.
Moreover since x>i P⊥v = 0, we have H(Bj)P⊥v = ∑ i∈Bj xix > i P⊥v = 0.
Applying P1 to Eq. (25) we have
P1vk,j+1 = P1
( I − 2η
b H(Bj)
) vk,j
= P1
( I − 2η
b H(Bj)
) (P1vk,j + P−1vk,j + P⊥vk,j)
= P1
( I − 2η
b H(Bj)
) P1vk,j + P1 ( I − 2η
b H(Bj)
) P−1vk,j
= ( I − 2η
b P1H(Bj)P1
) · P1vk,j − ( 2η
b P1H(Bj)P−1
) · P−1vk,j .
Similarly applying P−1 to Eq. (25) we have
P−1vk,j+1 = P−1
( I − 2η
b H(Bj)
) vk,j
= P−1
( I − 2η
b H(Bj)
) (P1vk,j + P−1vk,j + P⊥vk,j)
= P−1
( I − 2η
b H(Bj)
) P1vk,j + P−1 ( I − 2η
b H(Bj)
) P−1vk,j
= − ( 2η
b P−1H(Bj)P1
) · P1vk,j + ( I − 2η
b P−1H(Bj)P−1
) · P−1vk,j .
To sum up we have( P1vk,j+1 P−1vk,j+1 ) = ( I − 2ηb P1H(Bj)P1 − 2η b P1H(Bj)P−1
− 2ηb P−1H(Bj)P1 I − 2η b P−1H(Bj)P−1
) · ( P1vk,j P−1vk,j ) Notice that if 1 /∈ Bj , i.e., x1 is not used in the j-th step, then we claim
P1H(Bj) = H(Bj)P1 = 0,
sinceH(Bj) = ∑ i∈Bj xix > i is composed by the data belonging to the column space of P−1. Therefore if 1 /∈ Bj we have( P1vk,j+1 P−1vk,j+1 ) = ( I 0 0 I − 2ηb P−1H(Bj)P−1 ) · ( P1vk,j P−1vk,j )

heading: C.7 PROOF OF LEMMA 10
text: Proof of Lemma 10. Clearly for each component in the production, the column space of P1 + P⊥, which is (n − d + 1)-dimensional, belongs to its eigenspace of eigenvalue 1, which yields the first claim.
In the following, we restrict ourselves in the column space of P−1. Let us expandM−1:
M−1 = m∏ j=1 ( I − 2η b P−1H(Bj)P−1 ) = ( I − 2η
b P−1H(Bm)P−1
) · · · ( I − 2η
b P−1H(B1)P−1 ) = I − 2η
b m∑ j=1
P−1H(Bj)P−1︸ ︷︷ ︸ H−1
+
( 2η
b )2 ∑ 1≤i<j≤n
P−1H(Bj)P−1H(Bi)P−1 + . . .︸ ︷︷ ︸ C .
We first analyze matrix H−1. Since H(Bj) = ∑ i∈Bj xix > i and π = {B1, . . . ,Bm} is a partition for index set [n], we have
H−1 = m∑ j=1 P−1H(Bj)P−1
= m∑ j=1 P−1 ∑ i∈Bj xix > i P−1
= P−1 n∑ i=1 xix > i P−1 = P−1XX >P−1,
which is exactly the matrix we studied in Lemma 5, and from where we have H−1 has eigenvalue zero (with multiplicity being n−d+1) in the column space of P1 +P⊥, and restricted in the column space of P−1, the eigenvalues of H−1 belong to (λn − nι, λ2 + nι). Then we analyze matrix C.
P−1H(Bj)P−1H(Bi)P )−1 = ( P−1
∑ i′∈Bi xi′x > i′P−1 )P−1 ∑ j′∈Bj xj′x > j′P−1  = ∑ i′∈Bi ∑ j′∈Bj (P−1xi′)〈P−1xi′ , P−1xj′〉(P−1xj′)>. (26)
Remember that Bi ∩ Bj = ∅ for i 6= j, thus xi′ 6= xj′ for i′ ∈ Bi and j′ ∈ Bj . Then from Lemma 1 we have, |〈P−1xi′ , P−1xj′〉| ≤ |〈xi′ , xj′〉| ≤ √ λi′λj′ · ι ≤ ι.
Inserting this into Eq. (26) we obtain
‖P−1H(Bj)P−1H(Bi)P )−1‖F ≤ b 2 ·max |〈P−1xi′ , P−1xj′〉|2 ≤ b2ι2.
We can bound the Frobenius norm of the higher degree terms in matrix C in a similar manner; in sum for the Frobenius norm of C, we have
‖C‖F ≤ m∑ s=2 ( 2η b )s · bs · ιs · ( m s )
= m∑ s=2 (2ηι) s · ( n s )
= m∑ s=0 (2ηι)s · ( n s ) − 1− 2mηι = (1 + 2ηι)m − 1− 2mηι ≤ 1 +m · 2ηι+ m 2 √ e
2 · (2ηι)2 − 1− 2mηι (for 2ηι < 1 2m )
≤ 4m2η2ι2,
where for the second to the last inequality we notice that for f(t) = (1 + t)m and t ∈ [0, 12n ], we have f
′′ (t) = m(m−1)(1+t)m−2 ≤ m(m−1)(1+ 12m )
m−2 ≤ m(m−1)· √ e,which implies f(t)
is (m2 √ e)-smooth for t ∈ [0, 12m ]; moreover, by the assumption that 3nι < λn and η < b λn+3nι
, we can indeed verify that
2ηι < 2bι λn + 3nι ≤ 2bι 6nι ≤ 1 2m . (27)
Now we rephraseM>−1M−1 as M>−1M−1 = ( I − 2η
b H−1 + C
> ) · ( I − 2η
b H−1 + C ) = ( I − 2η
b H−1
)2 + C> ( I − 2η
b H−1
) + ( I − 2η
b H−1 ) C + C>C︸ ︷︷ ︸
D
. (28)
Restricting ourselves in the column space of P−1, the eigenvalues of H−1 belong to (λn − nι, λ2 + nι), thus the eigenvalues of ( I − 2ηb H−1 )2 are upper bounded by
max {( 1− 2η
b (λ2 + nι)
)2 , ( 1− 2η
b (λn − nι)
)2} < 1, (29)
where the last inequality is guaranteed by our assumptions on η and ι. For simplicity we defer the verification to the end of the proof.
Consider the following eigen decomposition I − 2ηH−1 = U diag (µ1, . . . , µn−1, 1, . . . , 1)U>, where µ1, . . . , µn−1 ∈ (−1, 1) by Eq. (29). Then we have
‖(I − ηH−1)C‖F = ∥∥diag (µ1, . . . , µn−1, 1, . . . , 1)U>CU∥∥F
≤ ∥∥U>CU∥∥
F = ‖C‖F .
Therefore we can bound the Frobenius norm of D by
‖D‖F ≤ 2 ‖(I − 2ηH−1)C‖F + ‖C‖ 2 F
≤ 2 ‖C‖F + ‖C‖ 2 F ≤ 8m2η2ι2 + 16m4η4ι4
≤ 9m2η2ι2, (30)
where the last inequality follows from 2ηι ≤ 1/(2m) proved in Eq. (27). Finally, applying Hoffman-Wielandt theorem with Eq. (28), (29) and (30), we conclude that, restricted in the column space of P−1, the eigenvalues ofM>−1M−1 are upper bounded by
max {( 1− 2η
b (λ2 + nι)
)2 + 9m2η2ι2, ( 1− 2η
b (λn − nι)
)2 + 9m2η2ι2 }
≤ max {(∣∣∣∣1− 2ηb (λ2 + nι) ∣∣∣∣+ 3mηι)2 , (∣∣∣∣1− 2ηb (λn − nι) ∣∣∣∣+ 3mηι)2 }
:= (q−1(η)) 2 . (31)
At this point we left to verify Eq. (29) and
q−1(η) := max {∣∣∣∣1− 2ηb (λ2 + nι) ∣∣∣∣+ 3nηιb , ∣∣∣∣1− 2ηb (λn − nι) ∣∣∣∣+ 3nηιb } < 1. (32)
Clearly it suffices to verify Eq. (32).∣∣∣∣1− 2ηb (λ2 + nι) ∣∣∣∣+ 3nηιb < 1
⇔ 3nι b η − 1 < 1− 2(λ2 + nι) b η < 1− 3nι b η
⇔ { 2λ2−nι b η > 0
2λ2+5nι b η < 2
⇐  η > 0
2λ2 − nι > 0 η < bλ2+2.5nι
⇐ { 3nι < λn (since λ2 ≥ λn) 0 < η < bλ2+3nι
Similarly, we verify that ∣∣∣∣1− 2ηb (λn − nι) ∣∣∣∣+ 3nηιb < 1
⇔ 3nι b η − 1 < 1− 2(λn − nι) b η < 1− 3nι b η
⇔ { 2λn−5nι b η > 0
2λn+nι b η < 2
⇐  η > 0
2λn − 5nι > 0 η < bλn+0.5nι
⇐ { 3nι < λn 0 < η < bλ2+3nι (since λ2 ≥ λn)
These complete our proof.

heading: C.8 PROOF OF LEMMA 11
text: Proof of Lemma 11. Note that during one epoch of SGD updates, x1 is used for only once. Without loss of generality, assume SGD uses x1 at the l-th step, i.e., 1 ∈ Bl and 1 /∈ Bj for j 6= l. Recursively
applying Lemma 9, we have( P1vk,m+1 P−1vk,m+1 ) = ( I 0 0 ∏m j=l+1 ( I − 2ηb P−1H(Bj)P−1 ))×( I − 2ηb P1H(Bl)P1 − 2η b P1H(Bl)P−1
− 2ηb P−1H(Bl)P1 I − 2η b P−1H(Bj)P−1 ) ×(
I 0 0 ∏l−1 j=1 ( I − 2ηb P−1H(Bj)P−1 ))× ( P1vk,1 P−1vk,1 ) Let vk+1 = vk,m+1, vk = vk,1 and
Ml := I − 2η
b P−1H(Bl)P−1
M>l := m∏
j=l+1
( I − 2η
b P−1H(Bj)P−1
)
M<l := l−1∏ j=1 ( I − 2η b P−1H(Bj)P−1 )
M−1 :=M>l · Ml · M<l = m∏ j=1 ( I − 2η b P−1H(Bj)P−1 ) then we have(
P1vk+1 P−1vk+1
) = ( I 0 0 M>l )( I − 2ηb P1H(Bl)P1 − 2η b P1H(Bl)P−1
− 2ηb P−1H(Bl)P1 Ml
)( I 0 0 M<l )( P1vk P−1vk ) = ( I − 2ηb P1H(Bl)P1 − ( 2η b P1H(Bl)P−1 ) M<l
−M>l ( 2η b P−1H(Bl)P1 ) M−1
)( P1vk P−1vk ) (33)
In the following we bound the norm of each entries in the above coefficient matrix.
According to Lemma 5, we have the eigenvalues of P−1H(Bj)P−1 are upper bounded by λ2 + nι. Thus the assumption η < bλ2+2nι yields∥∥∥∥I − 2ηb P−1H(Bj)P−1 ∥∥∥∥ 2 ≤ 1,
which further yields ‖M>l‖2 ≤ 1, ‖M<l‖2 ≤ 1. (34) On the other hand notice that P1xi = 0 for i 6= 1, thus
P1H(Bl)P−1 = P1 ∑ i∈Bl xix > i P−1 = P1x1x > 1 P−1,
P−1H(Bl)P1 = P−1 ∑ i∈Bl xix > i P1 = P−1x1x > 1 P1,
which yield
max {‖P1H(Bl)P−1‖2 , ‖P−1H(Bl)P1‖2} ≤ ‖P1x1‖2 · ‖P−1x1‖2 ≤ 2 √ nι, (35)
where the last inequality is from Lemma 3 and λ1 = ‖x1‖2 ≤ 1. Eq. (34) and (35) imply
max {∥∥∥∥(2ηb P1H(Bl)P−1 ) M<l ∥∥∥∥ 2 , ∥∥∥∥M>l(2ηb P−1H(Bl)P1 )∥∥∥∥ 2 } ≤ 4η √ nι b =: ξ(η) (36)
Next, by P1xi = 0 for i 6= 1 we have P1H(Bl)P1 = P1 ∑ i∈Bl xix > i P1 = P1x1x > 1 P1 = (P1x1)(P1x1) >,
from where we know ‖P1x1‖22 is the only non-zero eigenvalue of the rank-1 matrix P1H(Bl)P1, and the corresponding eigenspace is the column space of P1. Therefore 1− 2ηb ‖P1x1‖ 2 2 is an eigenvalue of the matrix I− 2ηb P1H(Bl)P1, and the corresponding eigenspace is the column space of P1, which implies ∥∥∥∥(I − 2ηb P1H(Bl)P1 ) P1vk ∥∥∥∥ 2 = ∥∥∥∥(1− 2ηb ‖P1x1‖22 ) P1vk
∥∥∥∥ = ∣∣∣∣1− 2ηb ‖P1x1‖22 ∣∣∣∣ · ‖P1vk‖2
=: q1(η) · ‖P1vk‖2 . (37)
Finally, according to Lemma 10, we have, restricted in the column space of P−1, the right eigenvalues ofM−1 is upper bounded by (q−1(η))2 , which implies
‖M−1P−1vk‖2 ≤ q−1(η) · ‖P−1vk‖2 . (38)
Note we have q−1(η) < 1 by Lemma 10.
Combining Eq. (33) with Eq. (36), (37), (38), and letting Bk := ‖P1vk‖2 , Ak := ‖P−1vk‖2, we obtain
Bk+1 ≤ q1(η) ·Bk + ξ(η) ·Ak Bk+1 ≥ q1(η) ·Bk − ξ(η) ·Ak Ak+1 ≤ q−1 · (η)Ak + ξ(η) ·Bk.

heading: C.9 PROOF OF LEMMA 12
text: Proof of Lemma 12. Let
ξ := ξ(η) = 4η √ nι
b ,
q1 := q1(η) = ∣∣∣∣1− 2ηλ1b ‖P1x̄1‖22 ∣∣∣∣ ,
q−1 := q−1(η) = max {∣∣∣∣1− 2ηb (λ2 + nι) ∣∣∣∣+ 3nηιb , ∣∣∣∣1− 2ηb (λn − nι) ∣∣∣∣+ 3nηιb } .
(39)
Then for 0 < k ≤ k1, Lemma 11 gives us
Bk ≥ q1Bk−1 − ξAk−1, (40)( Ak Bk ) ≤ ( q−1 ξ ξ q1 ) · ( Ak−1 Bk−1 ) , (41)
where “≤” means “entry-wisely smaller than”. Let θ, ρ−1, ρ1 determine the eigen decomposition of the coefficient matrix, i.e.,(
q−1 ξ ξ q1
) = ( cos θ sin θ − sin θ cos θ )( ρ−1 0 0 ρ1 )( cos θ − sin θ sin θ cos θ ) . (42)
Then Eq. (41) and Eq. (42) yield( Ak Bk ) ≤ ( q−1 ξ ξ q1 )k · ( A0 B0 ) = ( cos θ sin θ − sin θ cos θ )( ρk−1 0 0 ρk1 )( cos θ − sin θ sin θ cos θ )( A0 B0
) = ( ρk−1 + (ρ k 1 − ρk−1) sin2 θ (ρk1 − ρk−1) cos θ sin θ
(ρk1 − ρk−1) cos θ sin θ ρk1 − (ρk1 − ρk−1) sin2 θ )( A0 B0 ) = ( A0 · ρk−1 + ( ρk1 − ρk−1 ) (A0 sin θ +B0 cos θ) sin θ
B0 · ρk1 + ( ρk1 − ρk−1 ) (A0 cos θ −B0 sin θ) sin θ ) ≤ ( A0 · ρk−1 +
∣∣ρk1 − ρk−1∣∣√A20 +B20 sin θ B0 · ρk1 + ∣∣ρk1 − ρk−1∣∣√A20 +B20 sin θ )
=
( A0 · ρk−1 + ∣∣ρk1 − ρk−1∣∣ · ‖Pv0‖2 · sin θ B0 · ρk1 + ∣∣ρk1 − ρk−1∣∣ · ‖Pv0‖2 · sin θ ) . (43)
We claim the following inequalities hold by our assumptions:
0 < ρ−1 < 1 < ρ1 ≤ q1 + ξ (44a)
ρk1−1 ‖Pv0‖2 ≤ 2 · β (44b)
ρk11 ‖Pv0‖2 sin θ ≤ 2 · β, (44c) ξ · ( A0 +
β0 2
) < (q1 − 1)β0. (44d)
The verification of Eq. (44) is left later. In the following we prove the conclusions using Eq. (44).
We first bound Ak1 using Eq. (43) and Eq. (44):
Ak1 ≤ A0 · ρ k1 −1 + ∣∣∣ρk11 − ρk1−1∣∣∣ · ‖Pv0‖2 · sin θ ≤ ‖Pv0‖2 · ρ k1 −1 + ρ k1 1 · ‖Pv0‖2 · sin θ
≤ 2 · β + 2 · β = · β,
which justifies the first conclusion. In addition we can obtain an uniform upper bound for Ak for k = 0, 1, . . . , k1:
Ak ≤ A0 · ρk−1 + ∣∣ρk1 − ρk−1∣∣ · ‖Pv0‖2 · sin θ
≤ A0 + ρk1 · ‖Pv0‖2 · sin θ
≤ A0 + 2 · β. (45)
Next we bound Bk1 using Eq. (43) and Eq. (44):
Bk1 ≤ B0 · ρ k1 1 + ∣∣∣ρk11 − ρk1−1∣∣∣ · ‖Pv0‖2 · sin θ ≤ ‖Pv0‖2 · ρ k1 1 + ρ k1 1 · ‖Pv0‖2 · sin θ
≤ ‖Pv0‖2 · ρ k1 1 + 2 · β,
which justifies the second conclusion.
We proceed to derive the uniform lower bound for Bk for k = 0, 1, . . . , k1. We do it by induction. For k = 0, by assumption we have B0 ≥ β0. Suppose Bk−1 ≥ β0, then by Eq. (40), (44) and (45)
we have Bk ≥ q1 ·Bk−1 − ξ ·Ak−1
≥ q1 · β0 − ξ · ( A0 + 2 β )
≥ q1 · β0 − ξ · ( A0 +
2 β0 ) ≥ β0,
which justifies the third conclusion.
Verification of Eq. (44)
From Eq. (42) and Gershgorin circle theorem we have
q1 − ξ ≤ ρ1 ≤ q1 + ξ, q−1 − ξ ≤ ρ−1 ≤ q−1 + ξ.
(46)
Moreover, reformatting Eq. (42) as( q−1 ξ ξ q1 ) = ( cos θ sin θ − sin θ cos θ )( ρ−1 0 0 ρ1 )( cos θ − sin θ sin θ cos θ ) = ( ρ−1 cos 2 θ + ρ1 sin 2 θ (ρ1 − ρ−1) cos θ sin θ
(ρ1 − ρ−1) cos θ sin θ ρ−1 sin2 θ + ρ1 cos2 θ ) = ( ρ−1 + (ρ1 − ρ−1) sin2 θ (ρ1 − ρ−1) cos θ sin θ (ρ1 − ρ−1) cos θ sin θ ρ1 − (ρ1 − ρ−1) sin2 θ ) ,
we then have ξ
q1 − q−1 = (ρ1 − ρ−1) cos θ sin θ (ρ1 − ρ−1)(1− 2 sin2 θ) = 1 2 tan 2θ. (47)
For Eq. (44a), using Eq. (46) it suffices to show
0 < q1 − ξ, (48a) q−1 + ξ < 1, (48b) 1 < q1 − ξ. (48c)
Notice the definitions of q1, q−1 and ξ are given in Eq. (39). Firstly, Eq. (48c) holds trivially when√ n > 4/3. Secondly, for Eq. (48b), noticing that ξ = 4η √ nι b ≤ ηnι b when n ≥ 16, it suffices to show
max {∣∣∣∣1− 2ηb (λ2 + nι) ∣∣∣∣+ 4nηιb , ∣∣∣∣1− 2ηb (λn − nι) ∣∣∣∣+ 4nηιb } < 1
⇔
{ 4nι b η − 1 < 1− 2(λ2+nι) b η < 1− 4nι b η
4nι b η − 1 < 1− 2(λn−nι) b η < 1− 4nι b η
⇔  2λ2−2nι b η > 0 2λ2+6nι b η < 2 2λn−6nι
b η > 0 2λn+2nι
b η < 2
⇐  η > 0 λ2 − nι > 0 λn − 3nι > 0 η < bλ2+3nι η < bλn+nι
⇐ { 3nι < λn 0 < η < bλ2+3nι
which are given in assumptions. Thirdly, for Eq. (48c) it suffices to show
2ηλ1 b ‖P1x̄1‖22 − 1−
4η √ nι
b > 1
⇐ 2λ1(1− 4nι 2) b η − 4
√ nι
b η > 2 (by Lemma 3)
⇐ η > b λ1(1− 4nι2)− 2 √ nι
⇐ η > b λ1 − 3 √ nι , (since nι < 1)
which are given in assumptions.
For Eq. (44b), it suffices to show set
k1 = 1 + log 0.5 β‖Pv0‖2
log ρ−1 = O
( log 1
β
) ,
as given in assumptions.
For Eq. (44c), using Eq. (44b) it suffices to show
sin θ ≤ ( ρ−1 ρ1 )k1 = ρ−1 ρ1 · ( 0.5 β ‖Pv0‖2 )1− log ρ1log ρ−1 ⇐ sin θ ≤ q−1 − ξ
q1 + ξ · ( 0.5 β ‖Pv0‖2 )1− log(q1+ξ) log(q−1−ξ)
⇐ ξ ≤ 0.9 (q1 − q−1) · q−1 − ξ q1 + ξ
· ( 0.5 β
‖Pv0‖2
)1− log(q1+ξ) log(q−1−ξ)
(by Eq. (47))
⇐ √ nι ≤ poly ( β) . (by Eq. (39))
For Eq. (44c), it suffices to show
ξ ≤ (q1 − 1)β0 A0 + 0.5 β0
⇐ √ nι ≤ O (1) . (by Eq. (39))

heading: C.10 PROOF OF LEMMA 13
text: Proof of Lemma 13. Let
ξ′ := ξ(η′) = 4η′ √ nι
b ,
q′1 := q1(η ′) = ∣∣∣∣1− 2η′λ1b ‖P1x̄1‖22 ∣∣∣∣ ,
q′−1 := q−1(η ′) = max {∣∣∣∣1− 2η′b (λ2 + nι) ∣∣∣∣+ 3nη′ιb , ∣∣∣∣1− 2η′b (λn − nι) ∣∣∣∣+ 3nη′ιb } . (49)
Then for k1 < k ≤ k2, Lemma 11 gives us( Ak Bk ) ≤ ( q′−1 ξ ′ ξ′ q′1 ) · ( Ak−1 Bk−1 ) , (50)
where “≤” means “entry-wisely smaller than”. Denote
B := ‖Pv0‖2 · ρ k1 1 + 2 · β = poly
( 1
β
) . (51)
We claim the following inequalities hold by our assumptions:
0 < q′1 < q ′ −1 < 1, (52a)
ξ′ · ≤ q′−1 − q′1, (52b) ξ′ ·B ≤ (1− q′−1) · · β. (52c)
The verification of Eq. (52) is left later. In the following we prove the main conclusions in the lemma using Eq. (52). We proceed by induction. Clearly the conclusions are true for k = k1. Suppose for k1, . . . , k − 1, the conclusions are also true. Then the induction assumptions give us
Ak−1 ≤ · β, (53) Bk−1 ≤ Bk1 ≤ B, (54)
where the last inequality is due to B ≥ Bk1 ≥ β0 > β. Then by Eq. (50) we have
Bk ≤ q′1 ·Bk−1 + ξ′ ·Ak−1 ≤ q′1 ·Bk−1 + ξ′ · · β (by Eq. (53)) ≤ q′1 ·Bk−1 + (q′−1 − q′1) · β (by Eq. (52b))
≤ { q′−1 ·Bk−1, Bk−1 > β, β, Bk−1 ≤ β.
(by Eq. (52a))
Also by Eq. (50) we have
Ak ≤ q′−1 ·Ak−1 + ξ′ ·Bk ≤ q′−1 · · β + ξ′ ·B (by Eq. (53) and (54)) ≤ q′−1 · · β + (1− q′−1) · · β (by Eq. (52c)) = · β.
Verification of Eq. (52) Notice the definitions of q′1, q′−1 and ξ′ are given in Eq. (39). Recall q′−1 < 1 is already justified by the choice of learning rate η
′ < 1λ2+3nι (e.g., see Lemma 10), thus for Eq. (52a), it suffices to show
0 < 1− 2η ′λ1 b ‖P1x̄1‖22 < 1− 2η′ b (λn − nι) + 3nη′ι b
⇐
{ 1− 2η
′λ1 b > 0
2λ1(1− 4nι2) > 2(λn − nι) + 3nι (by Lemma 3)
⇐ { η′ < b2λ1 λ1 > λn + 2nι
which are given in assumptions.
For Eq. (52b), it suffices to show
ξ′ ≤ 1 · ( q′−1 − q′1 ) ⇐ √ nι ≤ O ( 1 ) ,
which is implied by nι ≤ poly ( β). For Eq. (52c), it suffices to show
ξ′ ≤ 1 B · ·
( 1− q′−1 ) β
⇐ √ nι ≤ poly ( β) . (by Eq. (51))
We complete our proof.

heading: C.11 PROOF OF LEMMA 14
text: Proof of Lemma 14. For the empirical loss,
LS(u) = 1 n (w − w∗)>XX> (w − w∗) = 1 n u>G>XX>Gu = 1 n u>Γu = 1 n n∑ i=1 γi ( u(i) )2 .
For the population loss,
LD(u) = µ ‖w − w∗‖22 = µ ‖Gu‖ 2 2 = µ ‖u‖ 2 2 .
For the hypothesis classHS = { w ∈ Rd : P⊥w = P⊥w0 } , Note P⊥G = diag (0, . . . , 0, 1, . . . , 1). Apply w − w∗ = Gu and notice w0 − w∗ = Gu0, then we obtain
HS = { u ∈ Rd : P⊥Gu = P⊥Gu0 } = { u ∈ Rd : u(i) = u(i)0 , for i = n+ 1, . . . , d } .
For the level set, we only need to note that L∗S = infu∈HS LS(u) = 0.
As for the estimation error, we note that
inf u∈U LD(u) = µ d∑ i=n+1 ( u (i) 0 )2 ,
thus for u ∈ U , we have
∆(u) = L(u)− inf u′∈U
L(u′) = µ ‖u‖22 − µ d∑
i=n+1
( u
(i) 0 )2 = µ
n∑ i=1 ( u(i) )2 + µ d∑ i=n+1 ( u(i) )2 − µ d∑ i=n+1 ( u (i) 0 )2 = µ
n∑ i=1 ( u(i) )2 .
Now consider u ∈ U , i.e., 1n ∑n i=1 γi ( u(i) )2 = α, then
∆∗ = inf u∈U ∆(u) = inf nα= ∑n i=1 γi(u(i)) 2 µ n∑ i=1 ( u(i) )2 = µnα γ1 ,
where the inferior is attended when, e.g., u(1) = ± √
nα γ1 and u(2) = · · · = u(n) = 0.

heading: D DETAILS OF THE EXPERIMENTS
text: In this section, we describe the details for our experiments.
D.1 2-D EXAMPLE
This part corresponds to Section 3 and Figure 1.
The two training data points are
x1 = ( √ κ, 0)>, x2 = (0, 1) >, κ = 4,
and the corresponding individual losses are
`1(w) = w >x1x > 1 w, `2(w) = w >x2x > 2 w.
Then the training loss is LS(w) = 0.5(`1(w) + `2(w)).
We initialize the algorithms from w0 = (0.6, 0.6)>. Two kinds of learning rate regime are considered. In the small learning rate regime, the learning rate is
ηk = 0.1/κ, k = 1, . . . , 800.
In the moderate learning rate regime, the learning rate is
ηk =
{ 1.1/κ, k = 1, . . . , 100;
0.1/κ, k = 101, . . . , 800.
For SGD, the mini-batch size is 1.

heading: D.2 LINEAR REGRESSION ON SYNTHETIC DATA
text: This part corresponds to Section 4 and Figure 2(a).
The model is an overparameterized linear model, with d = 104 and n = 100. The true parameter w∗ is randomly drawn from an d-dimensional Gaussian distribution, N (0, 0.1 · Id×d). We then randomly draw n = 100 samples from the d-dimensional space as described in Section 4, where ζ ∼ U([0.5, 1]). We initialize the algorithms from zero. We consider two kinds of learning rate regimes. The small learning rate scheme is specified by
ηk = 0.2, k = 1, . . . , 10 4.
The moderate learning rate scheme is specified by
ηk = { 1.05, k = 1, . . . , 2× 103; 0.1, k = 1 + 2× 103, . . . , 3× 103.
For SGD, the mini-batch size is 1.

heading: D.3 NEURAL NETWORK ON A SUBSET OF FASHIONMNIST
text: This part corresponds to Figure 2(b) and Figure 3.
Model We use a LeNet-alike convolutional network:
input⇒ conv1⇒ ReLU⇒ max_pool⇒ conv2⇒ ReLU⇒ max_pool⇒ fc1⇒ ReLU⇒ fc2⇒ ReLU⇒ linear⇒ output.
The first convolutional layer uses 5 × 5 kernels with 10 channels and no padding and the second convolutional layer uses 5×5 kernels with 16 channels and no padding. The number of hidden units between the two fully connected layers are 60.
Dataset https://github.com/zalandoresearch/fashion-mnist
We randomly choose 2, 000 original test data as our training set, and use the 60, 000 original training data as our test set. Thus we have 2, 000 training data and 60, 000 test data. We scale the image data to [0, 1].
Algorithms We randomly initialize the algorithms from a Gaussian distribution with zero mean and standard deviation 0.02. We consider two kinds of learning rate regimes. The small learning rate scheme is specified by
ηk = 10 −3, k = 1, . . . , 104.
The moderate learning rate scheme is specified by
ηk = { 10−2, k = 1, . . . , 2.5× 103; 10−3, k = 1 + 2.5× 103, . . . , 104.
For SGD, the mini-batch size is 25. For both GD and SGD, the weight decay parameter is set as 0.002.
Relative Rayleigh quotient We discuss the relative Rayleigh quotients calculated in Figure 2(b). Unlike the linear regression model where the Hessian is a constant, for neural networks the loss function is non-convex. In other words, not only there are multiple local minima, but also the Hessian varies at different points. Therefore, a direct comparison in terms of the Rayleigh quotients for the iterates of different algorithms makes little sense. Instead, we consider the relative Rayleigh quotient, i.e., the Rayleigh quotient normalized by the maximum absolute eigenvalue of the Hessian at that point. Mathematically, the relative Rayleigh quotient is defined as
RRQ(w) := ∇L(w)> ‖∇L(w)‖2 · ∇2L(w) · ∇L(w)‖∇L(w)‖2 ‖∇2L(w)‖2 ,
where∇L(w)/ ‖∇L(w)‖2 is the convergence direction of gradient methods, and ‖∇2L(w)‖2 is the operator norm of the Hessian, i.e., its maximum absolute eigenvalue. Note that it is computationally hard in practice to project a parameter onto the data manifold, thus we use the vanilla convergence direction instead of the projected one in the above definition of the relative Rayleigh quotient. Since our goal is to compare the relative Rayleigh quotient between different algorithms, this simplification would not affect our conclusions. We obtain the maximum absolute eigenvalue of the Hessian by running power method for 5 iterates.
Additional experiments: GD and SGD with increasing learning rates We conduct numerical experiments of training neural networks on a subset of FashionMNIST dataset for SGD and GD with different learning rates η ∈ {0.001, 0.01, 0.02, 0.04, 0.08, 0.16}. The test accuracy results are displayed in Figure 4.
Several conclusions can be drawn from the plots. (1) For a learning rate over η = 0.08, SGD cannot converge (thus only gives about 10% test accuracy). SGD generalizes best with a moderate learning rate η = 0.01. (2) GD fails to converge when η = 0.16. Also, as the learning rate increases, the test accuracy of GD first increases then decreases; but even at its peak (η = 0.01), GD performs worse than SGD with a moderate learning rate.
Paired t-test for Figure 3 We also conduct statistical test to show SGD with moderate learning rate is significantly better than the other baselines. Recall that the experiments are repeated for 10 runs, at each run, we first fix a random seed, then run the four algorithms (GD/SGD with small/moderate learning rate) under the same seed. In Table 1, we report the complete results from the 10 runs. By running a paired t-test at the 5% significance level, we find that SGD with moderate learning rate is significantly better than GD with moderate learning rate (p-value = 0.0043). Similarly, SGD with moderate learning rate is significantly better than SGD with small learning rate (p-value = 0.012), and is also significantly better than GD with small learning rate (p-value = 0.0095).

Review summaries:
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.

In the paper the authors analyzed the convergence dynamics toward a minimum of gradient descent (GD) and stochastic gradient descent (SGD). The algorithms are considered to start in the basin of attraction of a minimum, and using discrete steps they approach the bottom. The main result of the paper concerns the fact that, with moderate learning rate, SGD approaches the minimum along the steepest direction, contrary to GD.

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.