On the Origin of Implicit Regularization in Stochastic Gradient Descent
Authors: Samuel L Smith
Benoit Dherin
David Barrett
Soham De
Venue: ICLR 2021 Poster (2021)
Keywords: SGD
learning rate
batch size
optimization
generalization
implicit regularization
backward error analysis
SDE
stochastic differential equation
ODE
ordinary differential equation

Abstract: For infinitesimal learning rates, stochastic gradient descent (SGD) follows the path of gradient flow on the full batch loss function. However moderately large learning rates can achieve higher test accuracies, and this generalization benefit is not explained by convergence bounds, since the learning rate which maximizes test accuracy is often larger than the learning rate which minimizes training loss. To interpret this phenomenon we prove that for SGD with random shuffling, the mean SGD iterate also stays close to the path of gradient flow if the learning rate is small and finite, but on a modified loss. This modified loss is composed of the original loss function and an implicit regularizer, which penalizes the norms of the minibatch gradients. Under mild assumptions, when the batch size is small the scale of the implicit regularization term is proportional to the ratio of the learning rate to the batch size. We verify empirically that explicitly including the implicit regularizer in the loss can enhance the test accuracy when the learning rate is small.

Published as a conference paper at ICLR 2021
ON THE ORIGIN OF IMPLICIT REGULARIZATION IN STOCHASTIC GRADIENT DESCENT
Samuel L. Smith Benoit Dherin², David G. T. Barrett' and Soham De DeepMinn, GGoogle (s (slsmith, cherin, barretidavid, schamde (@google. com
ABSTRACT
For infinitesimal learning rates, stochastic gradient descent (SGD) follows the path of gradient flow on the full batch loss function. However moderately large learn- ing rates can achieve higher test accuracies, and this generalization benefit is not explained by convergence bounds, since the learning rate which maximizes test accuracy often larger than the learning rate which minimizes training loss. To interpret this phenomenon we prove that for SGD with random shuffling, the mean SGD iterate also stays close the path of gradient flow if the learning rate is small and finite, but on a modified loss. This modified loss is composed of the original loss function and an implicit regularizer, which penalizes the norms of the mini- batch gradients. Under mild assumptions, when the batch size is small the scale of the implicit regularization term is proportional to the ratio of the learning rate to the batch size. We verify empirically that explicitly including the implicit reg- ularizer in the loss can enhance the test accuracy when the learning rate is small.
INTRODUCTION
In the limit of vanishing learning rates, stochastic gradient descent with minibatch gradients (SGD) follows the path of gradient flow the full batch loss function (Yaida) 2019). However in deep
networks, SGD often achieves higher test accuracies when the learning rate is moderately large
(LeCun et al. 2012 Keskar et al.] 2017). This generalization benefit is not explained by convergence rate bounds (M et 20188 Zhang et al.] (201)), because it arises even for large compute budgets
for which smaller learning rates often achieve lower training losses (Smith al.) 2020). Although
many authors have studied this phenomenon (Jastrzebski et al. 2018 Smith & Le 2018 Chaudhari
& Soatto 2018; (Shallue et al. 2018 Park et al. 2019: Li al. 2019 Lewkowycz et al.) 2020),
remains poorly understood, and is an important open question in the theory of deep learning.
In a recent work, Barrett & Dherin (202)) analyzed the influence of finite learning rates on the iterates of gradient descent (GD). Their approach is inspired by backward error analysis, a method for the numerical analysis of ordinary differential equation (ODE) solvers (Hairer al. (006). The key insight of backward error analysis is that we can describe the bias introduced when integrating an ODE with finite step sizes by introducing an ancillary modified flow. This modified flow is derived to ensure that discrete iterates of the original ODE lie on the path of the continuous solution to the modified flow. Using this technique, the authors show that if the learning rate e is not too large, the discrete iterates of GD lie close to the path of gradient flow on a modified loss CGD = C(w)+ + (E/ ()w) This modified loss is composed of the original loss C(w) and an implicit regularizer proportional to the learning rate which penalizes the euclidean norm of the gradient.
However these results only hold for full batch GD, while in practice SGD with small moderately large batch sizes usually achieves higher test accuracies (Keskar et al. 2017 Smith et al.) 20200. In this work, we devise an alternative approach to backward error analysis, which accounts for the correlations between minibatches during one epoch of training. Using this novel approach, we prove that for small finite learning rates, the mean SGD iterate after one epoch, averaged over all possible sequences of minibatches, lies close to the path of gradient flow a second modified loss (sCD(W), which define in equation This new modified loss is also composed of the full batch loss function and an implicit regularizer, however the structure of the implicit regularizers for GD and SGD differ, and their modified losses can have different local and global minima. Our analysis

---

Published as a conference paper at ICLR 2021
therefore helps explain both why finite learning rates can aid generalization, and why SGD can achieve higher test accuracies than GD. We assume that each training example is sampled once per epoch, in line with best practice (Bottou 2012). and we confirm empirically that explicitly including the implicit regularization term of SGD in the training loss can enhance the test accuracy when the learning rate is small. Furthermore, we prove that if the batch size is small and the gradients are sufficiently diverse, then the expected magnitude of the implicit regularization term of SGD is
proportional the ratio of the learning rate to the batch size (Goyal et al. 2017 Smith al.
2018).
We note that many previous authors have sought to explain the generalization benefit of SGD using an analogy between SGD and stochastic differential equations (SDEs) (Mandt al., 2017 |mmth & Le. 2018 Jastrzębski et al. 2018:; CChuuhaaii & Soatto 2018). However this SDE analogy assumes that each minibatch is randomly sampled from the full dataset, which implies that some examples will be sampled multiple times in one epoch. Furthermore, the most common SDE analogy holds only for vanishing learning rates (aidaa (019) and therefore misses the generalization benefits of finite learning rates which we identify in this work. An important exception is Li et al. (0017), who applied backward error analysis to identify a modified SDE which holds when the learning rate is finite. However this work still relies on the assumption that minibatches sampled randomly. It also focused on the convergence rate, and did not discuss the performance of SGD on the test set.
Main Result. We now introduce main result. We define the cost function over parameters w as C(w) - (1/N) Cy(w), which is the mean of the per-example costs Cj(w), where N denotes the training set size. Gradient flow follows the ODE = (W((), while gradient descent computes discrete updates w;-1 = wi e (V((4)) where is the learning rate. For simplicity, we assume that the batch size B perfectly splits the training set such that N%B = O, where % denotes the modulo operation, and for convenience we define the number of batches per epoch m = N/B. We can therefore re-write the cost function as a sum over minibatches C(w) (1/m) Σk= Ck (w(w)) where the minibatch cost Čx (w) = (L/B) Σ k Cý(w). In order to guarantee that we sample each example precisely once per epoch, we define SGD by the discrete update wi+1 = w - (w:).
Informally, our main result is as follows. After one epoch, the mean iterate of SGD with a small but finite learning rate e, averaged over all possible shuffles of the batch indices, stays close to the path of gradient flow on a modified loss ف = -VCsa((w), where the modified loss CSGD is given by:
Čsan(w) = C(w) + 4vn Σ K= [|VC(())]
(1)
We emphasize that our analysis studies the mean evolution of SGD, not the path of individual tra- jectories. The modified loss ČsaD (w) is composed of the original loss C(w) and an implicit reg- ularizer Creg (w) = (1/4m) Σx= ||VC& (2) The scale of this implicit regularization term is proportional to the learning rate e, and it penalizes the mean squared norm of the gradient evaluated on a batch of B examples. To help us compare the modified losses of GD and SGD, we can expand,
(san(w) = C(w) + { [[V[C())]
+
E E m 4m == [|VĆi(w) VC(W)]
(2)
We arrive at Equation 2 from Equation by noting that E: ( (VC((w) VC(w) = O. In the limit B + N, we identify the modified loss of gradient descent, CGD = C(w) + (e (e/ which penalizes "sharp" regions where the norm of the full-batch gradient (I/VO(<)) is large. However, as shown by Equation [2] the modified loss of SGD penalizes both sharp regions where the full-batch gradient is large, and also "non-uniform" regions where the norms of the errors in the minibatch gradients (I/V((w) - VC( are large (wu et al. 2018). Although global minima of C(w) are global minima of CGD(W), global minima of C(w) may not be global (or even local) minima of (san(w). Note however that C(w) and (san(w) do share the same global minima on over-parameterized models which can interpolate the training set (Ma et al., 2018). We verify in our experiments that the implicit regularizer can enhance the test accuracy of models trained with SGD.
Paper structure. In Section 2 we derive our main result (Equation []] and we confirm empirically that we can close the generalization gap between small and large learning rates by including the implicit regularizer explicitly in the loss function. In Section 3 we confirm Equation I satisfies the linear scaling rule between learning rate and batch size (Goyal et al.) (017). In Section we provide additional experiments which challenge the prevailing view that the generalization benefit of small batch SGD arises from the temperature of associated SDE (Mandt et al. 2017 Park al. 2019.
2

---

Published as a conference paper at ICLR 2021
2 A BACKWARD ERROR ANALYSIS OF STOCHASTIC GRADIENT DESCENT
Backward error analysis has great potential to clarify the role of finite learning rates, and to help identify the implicit biases of different optimizers. We therefore give detailed introduction to the core methodology in Section 2.1] before deriving our main result in Section 2.2 In Section [.3] we confirm empirically that the implicit regularizer can enhance the test accuracy of deep networks.
2.1 AN INTRODUCTION TO BACKWARD ERROR ANALYSIS
In numerical analysis, we often wish to integrate ODEs of the form ه == (Ww). This system usually cannot be solved analytically, forcing us simulate the continuous flow with discrete updates, like the Euler step (tt + +) R (t) + e (w).). However discrete updates will introduce approximation error when the step size is finite. In order study the bias introduced by this approximation error, we assume the learning rate is relatively small, and introduce a modified flow ت = f(w), where,
F(w) f(w) + +ff((w) + e² f2(WW +....
The modified flow of f(w) is equal to the original flow of f(w) when e c 0, but it differs from the original flow if e is finite. The goal of backward error analysis is to choose the correction terms fe(w) such that the iterates obtained from discrete updates of the original flow with small finite step sizes lie on the path taken by the continuous solution to the modified flow with vanishing step sizes.
The standard derivation of backward error analysis begins by taking a Taylor expansion in c of the solution to the modified flow wt + e). We obtaín the derivatives of wt++ + +) recursively using the modified flow equation ت = f(w) (see Hairer et al. (2006)), and we identify the correction terms fi(w) by ensuring this Taylor expansion matches the discrete update (e.g.. Wt+1 Wi + ef for all powers of e. However, this approach does not clarify why these correction terms arise. To build our intuition for the origin of the corrections terms, and to clarify how we might apply this analysis to SGD, we take a different approach. First, we will identify the path taken by the continuous modified flow by considering the combined influence of an infinite number of discrete steps in the limit of vanishing learning rates, and then we will compare this continuous path to either single step of GD or single epoch of SGD. Imagine taking n Euler steps on the modified flow f(w) with step size a,
Whon = wi + af(w)) + a (1( + af(w))2) +...
(4)
= wi +a](w( + alw - al(w)) + aflwr + +a((w)) + allur + af(w))) +... (5)
= wi+ (n+nnJ(()) + (n/2)(n 1)a ((m)) + O(n²a²)
(6)
We arrived at Equation&by taking the Taylor expansion of and then counting the number of terms of type V[(w)) (wx) using the formula for an arithmetic series. Note that we assume V exists. coincides with the solution w(t + e) of the continuous modified flow à = -[(w) for small but finite c, we let the number of steps n - oc while setting a = c/n,
Next, to ensure Wi+n Equation 6
+ w(t c) w(t) ( (2/22 + +(2)
w())+e + ((((())))))()))))))))))()))))))) (1/))
+ O
(8)
We have replaced F(w) with its definition from Equation 3 the key component of backward error analysis, which describes the path taken when integrating the continuous modified flow f(w) with vanishing learning rates over discrete time step of length e.
As we will see below, Equation g is
Notice that we have assumed that the Taylor expansion in Equation S converges, while the higher order terms at (c)a) will contain higher order derivatives of the original flow f(w). Backward error analysis therefore implicitly assumes that f(w) is an analytic function in the vicinity of the current parameters w. We refer the reader to Hairer et al. (006) for a detailed introduction.
Gradient descent: As simple example, we will now derive the first order correction (w) of the modified flow for GD. First, we recall that the discrete updates obey Wi+1 = w. - cVVC(w;; and we therefore fix (w) = (C(w). In order to ensure that the continuous modified flow coincides with this discrete update, we need all terms at (c²²) and above in Equation g to vanish. At order e", this implies that fi (w: + (1/2) V VC(W) O, which yields the first order correction,
[w( - -(12))VC())VCC)
(1/44)V (I/VC(W))))
99)
3

---

Published as a conference paper at ICLR 2021
We conclude that, if the learning rate e is sufficiently small such that can neglect higher order terms in Equation 3 then the discrete GD iterates lie on the path of the following ODE,
= - VVC(W) - (c(w))(c/4)V) (IICCWW)I
(10)
-VC
(11)
Equation 11
corresponds to gradient flow on the modified loss, Can(w) ((w)+(//4]]]//))]])
2.2 BACKWARD ERROR ANALYSIS AND STOCHASTIC GRADIENT DESCENT
We now derive main result (Equation I.. As described in the introduction, assume N%B 0, where N is the training set size, B the batch size, and % denotes the modulo operation. The number of updates per epoch m = N/B, and the minibatch costs C₂ (w) (1/B) jj=kB+1 Cj(w).
SGD with constant learning rates obeys wi+1 = W/- vV (w,). It is standard practice to shuffle the dataset once per epoch, but omit this step here and instead perform our analysis over single epoch. In Equation we derived the influence of Euler steps on the flow f(w) with step size a. Following a similar approach, we now derive the influence of m SGD updates with learning rate e,
Wm
wo - <VColun) (VCi(un) <VC2(w2)
(12)
m-1
m-1
Wp c Σ VŪj(ub) + +
Σ
(wb) VCL(o))
(13)
1=0
1=0
Wo MeVC(up) + (((Wo) + O(m²e))
(14)
The error in Equation 14 (2) since there O(m²) terms in the Taylor expansion propor- tional to e². Notice that a single epoch of SGD is equivalent to single GD update with learning rate up to first order in c. Remarkably, this implies that when the learning rate is sufficiently small, there is no noise in the iterates of SGD after completing one epoch. For clarity, this observation arises because require that each training example is sampled once per epoch. However the sec- ond order correction {(w) E Exey VVC, (V)(()VVC(())) does not appear in the GD update, and it is random variable which depends on the order of the mini-batches. In order identify the bias introduced by SGD, we will evaluate the mean correction B(E), where we take the expectation across all possible sequences of the (nonooverlapping mini-batches (Co, Č1,... Om-1). Note that we hold the composition of the batches fixed, averaging only over their order. We conclude that,
E(E(w)) = 2
E
m-1
Σ Kマj
VVCJW))CA(()
(15)
j=0
m²
VVC(w)VC(()
m Σ
VVĆTVC,
(16)
2
==
m²
(IVC(W))
1
E j
IV(((W))))
(17)
m
For clarity, in Equation 15 we exploit the fact that every sequence of batches has corresponding and
sequence in reverse order. Combining Equations and
17
we conclude that after one epoch.
E(wm) = wi mcVC((u) m²²
+ v (1VCC((0))1-- (1/m²)
E,- m-
(1CC,(0))1)) -0 Om) (18)
Having identified the expected value of the SGD iterate after one epoch E(wm) (for small but finite learning rates), we can now use this expression to identify the corresponding modified flow. First, set (w) VV((w, / 0, (0) wn. and let - me in Equations 3 and to obtain,
w(me) = We - Mc + 222 (wi(w)) + (1/4)V + O(m²))) (19)
Next, we equate Equations 18 and 19 by setting (me) = [(wm). ). We immediately identify the first order correction to the modified flow [I( -(1/(4m²)) V (1)), We therefore
4

---

Published as a conference paper at ICLR 2021
conclude that, after one epoch, the expected SGD iterate E(wm) = w(me) + (m) where w0) = wo and = -[((w + mef((w). Simplifying, we conclude i -VCsGD (w), where,
ČSaD(w) = C(w) + (c/4m)
Σ -- 1 [VV(
(20)
k=0
Equation 20 identical to Equation[] and this completes the proof of our main result. We emphasize that CSGD assumes a fixed set of minibatches {Co, Cj,..,, Cm-1). We will evaluate the expected modified loss after shuffling the dataset and sampling a new set of minibatches in Section 3
REMARKS ON THE ANALYSIS
The phrase "for small finite learning rates" has a precise meaning in our analysis. It implies c is large enough that terms of (m) may be significant, but small enough that terms of 3 are negligible. Our analysis is unusual, because we consider the mean evolution of the SGD iterates but ignore the variance of individual training runs. Previous analyses of SGD have usually focused on the variance of the iterates in limit of vanishing learning rates (Mandt et al., 2017] Smith & Le, 2018; Jastrzebski et al. 2018). However these works assume that each minibatch is randomly sampled from the full dataset. Under this assumption, the variance in the iterates arises at O(c), while the bias arises at O By contrast, in our analysis each example is sampled once per epoch, and both the variance and the bias arise at (c²)) (for simplicity, we assume = N/B B is constant). We therefore anticipate that the variance will play a less important role than is commonly supposed.
Furthermore, we can construct specific sequence of minibatches for which the variance at e2) vanishes, such that the evolution of a specific training run will coincide exactly with gradient flow on the modified loss of Equation I for small finite learning rates. To achieve this, perform two training epochs with the sequence of minibatches (Cu, C1,..., Cm- Cm-1,.... C2, Co) (i.e., the second epoch iterates through the same set of minibatches as the first but in the opposite order). If one inspects Equations 13 to 15 one will see that reversing the second epoch has the same effect as taking the mean across all possible sequences of minibatches (it replaces the Σx<j by a Exzz).
The key limitation of our analysis is that we assume me Ne/B is small, in order to neglect terms (2))²)) This is an extreme approximation, since we typically expect that N/B B is large. There- fore, while our work identifies the first order correction to the bias arising from finite learning rates, higher order terms the modified flow may also play an important role at practical learning rates. We note however that previous theoretical analyses have made even more extreme assumptions. For instance, most prior work studying SGD in the small learning rate limit neglects all terms at (c²)) (for exception, see Li et al. (017)). Furthermore, scales proportional to the batch size, such that e/ is constant (Goyal et al. 2017 McCandlish et al.) 2018 Therefore the accuracy of our approximations does not necessarily degrade the batch size falls, but higher order terms may play an increasingly important role as the dataset size increases. Our experimental results in Section 2.3 and Section 4 suggest that our analysis can explain most of the generalization benefit of finite learning rate SGD for Wide-ResNets trained on CIFAR-10.
show in Section 3 the learning rate often
We note that to achieve the highest test accuracies, practitioners usually decay the learning rate during training. Under this scheme, the modified loss would change as training proceeds. However is widely thought that the generalization benefit of SGD arises from the use of large learning rates early in training (Smith et al. 2018 Li et al.) 2019 Jastrzebski et al. 2020 Lewkowycz et al., 2020), and popular schedules hold the learning rate constant or approximately constant for several epochs.
Finally, we emphasize that our primary goal in this work is to identify the influence of finite learning rates on training. The implicit regularization term may not be beneficial in all models and datasets.
2.3 AN EMPIRICAL EVALUATION OF THE MODIFIED Loss
In order to confirm that the modified loss (s) can help explain why large learning rates en- hance generalization, we now verify empirically that the implicit regularizer inherent in constant learning rate SGD, Crey(w)= = (1 /4m) Σmt [[VC( can enhance the test accuracy of deep networks. To this end, we train the same model with two different (explicit) loss functions. The first loss function ((w) represents the original loss, while the second (w) = C(w) + XCrey (w) is obtained from the modified loss (sen(w) by replacing the learning rate with explicit regular-
5

---

Published as a conference paper at ICLR 2021
C.90
0.89
~ MM
101-
(A-0)
Accurscy
Test sccurazy
(c-2-1)
0.06
Codo))
2-2-9 (0-0)
101
0.84-
0.59
(A-0)
A-0 A-2 6
0.5
le-2 5)
2- (A (4
10a-
2000
40D0
coco
3.53
2:23
2-11
3-3
3"7
3
2:33
2-11
2-11
2-1
7-1
2.
Training cpochs
Learning rate/lequiaaiattinn scale
Learning waw.//olguarraaaaaaaa scale
(a)
(b)
(c)
Figure 1: (a) Explicitly including the implicit regularizer in the loss improves the test accuracy when training with small learning rates. (b) The optimal regularization coefficient Aopt = 2 is equal to the optimal learning rate Copt 2 2-6 (c) Increasing either the learning rate the regularization coefficient À reduces the value of the implicit regularization term Creg (w) at the end of training.
ization coefficient A. Notice that (w) (1/m) EE= (Ce(w) + ( which
ensures that it is straightforward to minimize the modified loss Cmod(w) with minibatch gradients. Since the implicit regularization term (w) expensive to differentiate (typically 5-10x over- head), we consider 10-1 Wide-ResNee model (Zagoruyko & Komodakis 2016) for classification on CIFAR-10. To ensure close agreement with our theorctical analysis, we train without batch nor- malization using SkipInit initialization (De & Smith 2020). We train for 6400 epochs batch size 32 without learning rate decay using SGD without Momentum. We use standard data augmenta- tion including crops and random flips, and we use weight decay with L2 coefficient emphasize that, since we train using a finite (though very large) compute budget, the final networks may not have fully converged. This is particularly relevant when training with small learning rates. Note that we provide additional experiments on Fashion-MNIST (Xiao et al. 2017) in appendix D
10-4. We
In Figure I(a) we compare two training runs, one minimizing the modified loss (m) with A = 2-6 and one minimizing the original loss C(w)). For both runs we use a small constant learning rate c 2-9 As expected, the regularized training run achieves significantly higher test accuracies late in training. This confirms that the implicit regularizer, which arises as a consequence of using SGD with finite learning rates, can also enhance the test accuracy if it is included explicitly in the loss. In Figure [(bb) we provide the test accuracy for a range of regularization strengths À (orange line). We provide the mean test accuracy of the best 5 out of 7 training runs at each regularization strength, and for each run take the highest test accuracy achieved during the entire training run. We use a fixed learning rate e = 2 for all X. For comparison, we also provide the test accuracy achieved with the original loss C(w) for range of learning rates e (blue line). In both cases, the test accuracy rises initially, before falling for large regularization strengths or large learning rates. Furthermore, in this network the optimal regularization strength on the modified loss Xope 2 2-6 is equal to the optimal learning rate the original loss Copt 2-6. Meanwhile when X - 0 the performance of the modified loss approaches the performance of the original loss at 2-9 (dotted green line). We provide the corresponding training accuracies in appendix[C] Finally, in Figure [[c] provide the values of the implicit regularizer Creg (w) at the end of training. As predicted by analysis, training with larger learning rates reduces the value of the implicit regularization term.
In Figure take the same 10-1 RedeeRReNNet model and provide the mean training and test accuracies achieved at a range of learning rates for two regularization coefficients (following the experimental protocol above). In Figure 2(a)) we train the original loss C(w) (A = 0), while in Figure 2(b). we train on the modified loss Cmoalw) with regularization coefficient X = 2-6 From Figure 2(a) when À= = 0 there is clear generalization benefit to large learning rates, as the learning rate that maximizes test accuracy (2 (2-b) is 16 times larger than the learning rate that maximizes training accuracy (2 (2-19). However in Figure 2(b) with X 2-6, the learning rates that maximize the test and training accuracies are equal (2 (2-3). This suggests that when we include the implicit regularizer explicitly in the loss, the generalization benefit of large learning rates is diminished.
3 IMPLICIT REGULARIZATION AND THE BATCH SIZE
In Section 2.2] we derived the modified loss by considering the expected SGD iterate after one epoch. We held the composition of the batches fixed, averaging only over the order in which the batches

---

Published as a conference paper at ICLR 2021
6:00
2.50
3.54
1.00
1.40
0.52
C.0
2.00
0.50
C00
6.00
3.01
A brain
A bes0
Aw 2-5 (ralk
6.8
41-33
201
3-a
2
16-
16-1
1011 11
10'1
20"
10"1
Learning oste Ic)
Loamiin g rate 10
42
718
(a)
(b)
(a)
(b)
Figure 2: (a) There is a clear generalization
Figure 3: (a) Different batch sizes achieve the
benefit large learning rates when training on
same test accuracy if the ratio of the learning rate
the original loss C(w) with À = O. (b) When we include the implicit regularizer explicitly in the generalization benefit of large learning rates is diminished.
to the batch size (e/ B) is constant and B is not too large. (b) The test accuracy is independent of the batch size if the ratio of the regularization coefficient to the batch size (A/B) is constant.
(w) and set À = 2 -
seen. This choice helped make clear how to explicitly explicitly include the implicit regularizer in the loss function in Section 2.3 However, in order to clarify how the implicit regularizer term depends on the batch size, we now evaluate the expected modified loss after randomly shuftling the dataset and sampling a new set of m non-overlapping minibatches {Co, Ci,,.., Cm-1h Since the minibatch losses (((w) are all identically distributed by symmetry, we recall Equation 2 and conclude that, B(CscD(w)) = C(w) + (e/4) || ||[[( + (c/4) E(IVV - VC(W)]]²),
(21)
where C(w) denotes batch of B non-overlapping examples, drawn randomly from the full dataset. (N-B) Tw)
To simplify equation 21 we prove in appendix A that [[][V((w) VC(w)]l3) = (N-1) B
where ((w) = (1/N) Σ [[VC,(W) - V. We therefore obtain,
-N-B) I(w). B(CscD(w)) = C(w) + / ||CC(W)| + IN-1) 4B
(22)
Note that ((w) is the trace of the empirical covariance matrix of the per-example gradients. We have not assumed that the minibatch gradients Gaussian distributed, however if the per-example gradients are heavy tailed (Simsekli et al. 2019) then I' ('(w) may diverge, in which case the expected value of the modified loss is ill-defined. Equation 22 shows that the implicit regularization term of SGD has two contributions. The first term is proportional to the learning rate and it penalizes the norm of the full batch gradient. The second term is proportional batch size e/ B (assuming N > B), and it penalizes the trace of the covariance matrix. To interpret this result, we assume that the minibatch gradients are diverse, such that (T(w)/ B) This assumption guarantees that increasing the batch size reduces the error in the gradient estimate.
the ratio of the learning rate the
>> [|VC(w)]
In this limit, the second term above will dominate, and therefore different batch sizes will experience the same implicit regularization so long the ratio of the learning rate to the batch size is constant.
To verify this claim, in Figure 2.3 we plot the mean test accuracies achieved on a 10-1 Wide-ResNet, trained on CIFAR-10 10 with a constant learning rate, for range of learning rates c, regularization coefficients À and batch sizes B. As expected, in Figure 3(a) training on the original loss C(w) for 6400 epochs, we that different batch sizes achieve similar test accuracies long the ratio c/B B is constant and the batch size is not too large. We note that this linear scaling rule is well known and has been observed in prior work (Goyal et al., [017 /mmih & Le, 2018 Jastrzębski al. 2018 Zhang et al. 2019). To confirm that this behaviour is consistent with the modified loss, in Figure 3(b) we fix the learning rate e = 2- and train on Cmod (w) range of regularization strengths for 10 million steps. As expected, different batch sizes achieve similar test accuracy so long as the ratio A/B is constant. We note that we expect this phenomenon to break down for very large batch sizes, however we were not able to run experiments in this limit due to computational constraints.
For very large batch sizes, the first implicit regularization term in Equation 22 dominates, the linear scaling rule breaks down, and the bias of SGD is similar to the bias of GD identified by Barrett & Dherin 2021 l). We expect the optimal learning rate to be independent of the batch size in this limit, as observed by McCandlish et al. (2018) and Smith et al. (2020). Convergence bounds also predict a transition between a small batch regime where the optimal learning rate x B and a large batch regime where the optimal learning rate is constant (Ma et al.] 2018 Zhang et al.) 2019). However these analyses identify the learning rate which minimizes the training loss. Our analysis compli- ments these claims by explaining why similar conclusions hold when maximizing test accuracy.
7

---

Published as a conference paper at ICLR 2021
4 FINITE LEARNING RATES AND STOCHASTIC DIFFERENTIAL EQUATIONS
In the previous two sections, we argued that the use of finite learning rates and small batch sizes introduces implicit regularization, which can enhance the test accuracy of deep networks. We ana- lyzed this effect using backward error analysis (Hairer et al. 2006 Li et al. 2017 Barrett & Dherin 2021]. but many previous papers have argued that this effect can understood by interpreting small batch SGD as the discretization of an SDE (Mandt et al., 2017 Smith & Le 2018 Jastrzębski et al.)
2018 Park et al. 2 In this section, we compare this popular perspective with our main results from Sections 2 and To briefly recap, in the SDE analogy single gradient update is given by
Wi++1 W/: V (w:) ). where Ö' denotes random batch of B non-overlapping training examples. Notice that in the SDE analogy, since examples are drawn randomly from the full dataset, there is no guarantee that each training example is sampled once per epoch. Assuming N > B > 1 and that the gradients are not heavy tailed, the central limit theorem is applied to model the noise in an update by Gaussian noise source & whose covariance is inversely proportional to the batch size:
Wi+1 = wi - <(VCC(::++ + +3//B) = W/- - cVC((w:) + VeTī.
(23)
This assumes E(,)) = 0 and B(G,E))) - F(w)o,3, where F(w) is the covariance matrix of the per example gradients, and we define the "temperature" T e/ B. The SDE analogy notes that Equation 23 is identical to the Fuler discretization of an SDE with step size e and temperature T (Gardiner ctt al., 1985). Therefore one might expect the SGD iterates to remain close to this underlying SDE in the limit of small learning rates (E + 0). In this limit, the temperature defines the influence of mini mini-batching on the dynamics, and it is therefore often assumed that the temperature also governs
the generalization benefit of SGD (Smith & Le 2018: Dastrzzebbi et al. 2018 Park et al.) 2019).
However this conclusion from the SDE analogy is inconsistent with our analysis in Section 2 To see this, note that in Section 2 we assumed that each training example is sampled once per epoch, as recommended by practitioners (Bottou 2012), and showed that under this assumption there is no noise in the dynamics of SGD up to first order in after one epoch of training. The SDE analogy therefore relies on the assumption that minibatches are sampled randomly from the full dataset. Furthermore, SGD only converges to the underlying SDE when the learning rate + O, but in this limit the temperature T - O and SGD converges to gradient flow Yaida 2019 We must use a finite learning rate to preserve finite temperature, but at any finite learning rate the distributions of the SGD iterates and the underlying SDE may differ. We now provide intriguing empirical evidence to support our contention that the generalization benefit of SGD arises from finite learning rates, not the temperature of an associated stochastic process. First, we introduce a modified SGD update rule:
n-step SGD: Apply gradient descent updates sequentially on the same minibatch with bare learn- ing rate a, effective learning rate e = na and batch size B. Sample the next minibatch and repeat.
To analyze n-step SGD, we consider the combined influence of updates the same minibatch:
Wi+1
Wi avV((w)) - a VV ...
(24)
W; navŌ((::) + 0(n²a²) W; - cVC(w;: + VETE++((2))
(25)
(26)
Equations 23 and 26 are identical up to first order in e but they differ at O(e2) and above. Therefore, if minibatches are randomly sampled from the full dataset, then the dynamics of standard SGD and n-step SGD should remain close to the same underlying SDE in the limit e + 0, but their dynamics will differ when the learning rate is finite. We conclude that if the dynamics of SGD is close to the continuum limit of the associated SDE, then standard SGD and n-step SGD ought to achieve similar test accuracies after the same number of training epochs. However if, generalization benefit of SGD arises from finite learning rate corrections at O(e2) and above, then should expect the performance of standard SGD and n-step SGD to differ. For completeness, we provide a backward error analysis of n-step SGD in appendix B In line with Section 2 (and best practice), we assume each training example is sampled once per epoch. We find that after one epoch, the expected n-step SGD iterate B(wm) = w(me) + O(m²e*), where w()) = wo. w -VCnsGD(w) and ČnSGD (w) = C(w) + (e/4mn) EnD (w) The scale of the implicit regularizer is proportional to a = c/n, which implies that the implicit regularization is suppressed n increases if we hold constant. As expected, we recover Equation T when 1.
argued in Section 2 the
8

---

Published as a conference paper at ICLR 2021
1950
3440
0.550
F=1
c.023
acuurcyy
Test accurecy
2.900
2090
c:000
16
2815
n=2
n-1 [3600 epochs)
[800 =pox.o))
cara
J030
(400 spachs)
3825
n (200 spachs)
n-16 1100 epochs)
0.25
n=8
2800
2
2,000
80000
25
52
100
70n
400 800
1600
Pare fwwring nnte (X)
Parr leaming rate (e)
Training epocts
(a)
(b)
(c)
(a) When training for 400 epochs. smaller values of n are stable at larger bare learning rates a, and this enables them to achieve higher test accuracies. (b) Similar conclusions hold when training for fixed number of updates. (c) We show the test accuracy at the optimal learning rate for a range of epoch budgets. We find that smaller values of n consistently achieve higher test aceuracy.
Figure 4:
In Figure 4(a) we plot the performance of n-step SGD at range of bare learning rates a, when training a 16-4 Wide-ResNet on CIFAR-10 for 400 epochs using SkipInit (De & Smith, 2020) at batch size 32. Each example is sampled once per epoch. We introduce a learning rate decay sched- ule, whereby we hold the learning rate constant for the first half of training, before decaying the learning rate by factor of 2 every remaining tenth of training, and we provide the mean test accu- racy of the best 5 out of 7 training runs for each value of a. The optimal test accuracy drops from 93.5% when n (standard SGD) to 88.8% when n 16. This occurs even though all values of perform the same number of training epochs, indicating that 16-step SGD performed 16 times more gradient updates. These results suggest that, least for this model and dataset, the generalization benefit of SGD is not controlled by the temperature of the associated SDE, but instead arises from the implicit regularization associated with finite learning rates. When we increase n we reduce the largest stable bare learning rate a, and this suppresses the implicit regularization benefit, which re- duces the test accuracy. We also verify in Figure [bb that similar conclusions arise if we hold the number of parameter updates fixed (such that the number of training epochs is inversely proportional to n). Smaller values of n are stable at larger bare learning rates and achieve higher test accuracies.
Finally confirm in Figure [(c) that the test accuracy degrades as increases even if one tunes both the learning rate and the epoch budget independently for each value of n, thus demonstrating that n-step SGD consistently achieves lower test accuracies as increases. Note that we provide additional experiments on Fashion-MNIST (Xiao et al., (017)) in appendix D.
5 DISCUSSION
Many authors have observed that large learning rates (Li et al. 2019: Lewkowycz et al., 2020).
and small batch sizes (Keskar et al. 2017 Smith et al. 2020), can enhance generalization. Most theoretical work has sought to explain this by observing that increasing the learning rate, or reducing the batch size, increases the variance of the SGD iterates (Smith & Le| [2018] Jastrzębski al.) 2018 Chaudhari & Soatto) 2018). We take different approach, and note that when the learning rate is finite, the SGD iterates are also biased (Roberts) 2018). Backward error analysis (Hairer et al. 2006 Li et al., 2017; Barrett & Dherin, 2021) provides a powerful tool that computes how this bias accumulates over multiple parameter updates. Although this work focused on GD and SGD, we anticipate that backward error analysis could also be used to clarify the role of finite learning rates in adaptive optimizers like Adam (Kingma & Ba 2015) or Natural Gradient Descent (Amari 1998).
We note however that backward error analysis assumes that the learning rate is small (though finite). therefore does not capture the chaotic or oscillatory dynamics which arise when the learning rate is close to instability. At these very large learning rates the modified loss, which is defined as Taylor series in powers of the learning rate, does not converge. Lewkowyez et al. (2020) recently argued that the test accuracies of wide networks trained with full batch gradient descent quadratic losses are maximized for large learning rates close to divergence. In this "catapult" regime, the GD iterates oscillate along high curvature directions and the loss may increase early in training. It remains an open question to establish whether backward error analysis fully describes the generalization benefit of small batch SGD, or if these chaotic or oscillatory effects also play a role in some networks.

---

Published as a conference paper at ICLR 2021
ACKNOWLEDGMENTS
We would like to thank Jascha Sohl-Dickstein, Razvan Pascanu, Alex Botev, Yee Whye Teh and the anonymous reviewers for helpful discussions and feedback on earlier versions of this manuscript.
REFERENCES
Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2):251- 276, 1998.
David Barrett and Benoit Dherin. Implicit gradient regularization. 9th International Conference Learning Representations, ICLR, 2021.
Léon Bottou. Stochastic gradient descent tricks. In Neural networks: Tricks of the trade, pp. 421- 436. Springer, 2012.
Pratik Chaudhari and Stefano Soatto. Stochastic gradient descent performs variational inference, converges to limit cycles for deep networks. 2018 Information Theory and Applications Work- shop (ITA), pp. 1-10. IEEE, 2018.
Soham De and Sam Smith. Batch normalization biases residual blocks towards the identity function in deep networks. Advances in Neural Information Processing Systems, 33, 2020.
Crispin W Gardiner et al. Handbook of stochastic methods, volume springer Berlin, 1985.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, An- drew Tulloch, Yangqing Jia, and Kaiming He. Accurate, Large Minibatch SGD: Training Ima- genet in Hour. arXiv preprint arXiv:1706.02677, 2017.
Ernst Hairer, Christian Lubich, and Gerhard Wanner. Geometric numerical integration: structure- preserving algorithums for ordinary differential equations, volume 31. Springer Science & Busi- ness Media, 2006.
Stanistaw Jastrzębski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Ben- gio, and Amos Storkey. Three Factors Influencing Minima SGD. In Artificial Neural Networks and Machine Learning, ICANN, 2018.
Stanislaw Jastrzebski, Maciej Szymezak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho, and Krzysztof Geras. The break-even point on optimization trajectories of deep neural networks. arXiv preprint arXiv:2002.09571 095000 2020.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, ICLR, 2017.
Diederik P. Kingma and Jimmy Ba. Adam: method for stochastic optimization. In 3rd Interna- tional Conference on Learning Representations, ICLR, 2015.
Yann A LeCun, Léon Bottou, Genevieve B Orr, and Neural networks: Tricks of the trade, pp. 9-48. Springer, 2012.
Müller. Efficient backprop. In
Aitor Lewkowyez, Yasaman Bahri, Ethan Dyer, Jascha Sohl- Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.02218. 2020.
Qianxiao Li, Cheng Tai, et al. Stochastic modified equations and adaptive stochastic gradient algo- rithms. In Proceedings of the 34th International Conference on Machine Learning-Volus Volume 70, pp. 2101-2110. JMLR. org, 2017.
Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks. In Advances in Neural Information Processing Systems, pp. 1669-1 1680, 2019.
10

---

Published as a conference paper at ICLR 2021
Siyuan Ma, Raef Bassily, and Mikhail Belkin. The power of interpolation: Understanding the effec- tiveness of SGD in modern overmmaraet learning. In Proceedings of the 35th International Conference Machine Learning. JCML, 2018.
Stephan Mandt, Matthew Hoffman, and David M Blei. Stochastic gradient descent approximate bayesian inference. The Journal of Machine Learning Research, 18(1):487344907, 2017.
Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAl Dota Team. An empirical model of large-batch training. arXiv preprint /rXiv:/8/2.06//00 2018.
Daniel Park, Jascha Sohl-Dickstein, Quoc Lc, and Samuel Smith. The effect of network width on stochastic gradient descent and generalization: an empirical study. In International Conference on Machine Learning, pp- 5042-5051, 2019.
Daniel A Roberts. SGD implicitly regularizes generalization error. In Integration of Deep Learning Theories Workshop, NeurIPS, 2018.
Karthik Sankararaman, Soham De, Zheng Xu, Ronny Huang, and Tom Goldstein. The impact of neural network overparameterrization on gradient confusion and stochastic gradient descent. In Teirty-seventh International Conference on Machine Learning, ICMI., 2020.
Christopher J Shallue, Jaehoon Lee, Joe Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George E Dahl. Measuring the effects of data parallelism neural network training. arXiv preprint 18X 1.03600, 2018.
Umut Simsekli, Levent Sagun, and Mert Gurbuzbalaban. A tail-index analysis of stochastic gradient noise in deep neural networks. In International Conference on Machine Learning, pp. 5827-5837, 2019.
Samuel L Smith and Quoe V Le. A bayesian perspective on generalization and stochastic gradient descent. International Conference on Learning Representations, 2018.
Samuel L Smith, Pieter-Jan Kindermans, Chris Ying, and Quoe Le. Don't decay the learning rate, increase the batch size. In International Conference on Learning Representations, 2018.
Samuel L Smith, Erich Elsen, and Soham De. On the generalization benefit of noise in stochastic gradient descent. In International Conference on Machine Learning, 2020.
Lei Wu, Chao Ma, and E Weinan. How sgd selects the global minima in pver-parameterized learning: A dynamical stability perspective. Advances Neural Information Processing Systems, pp- 2018
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-MNIST: Novel Image Dataset for Bench- marking Machine Learning Algorithms. arXiv preprint arXiv:/708.00) 1708.07747, 2017.
Sho Yaida. Fluctuation Fluctuation-dissiptioon relations for stochastic gradient descent. In 7th International Conference on Learning Representations, ICLR, 2019.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. In Proceedings of the British Machine Vision Conference (BMVC), pp. 87.1-87.12, September 2016.
Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George E. Dahl, Christopher Shallue, and Roger B. Grosse. Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Model. In Advances in Neural Information Processing Systems, 2019.
11

---

Published as a conference paper at ICLR 2021
A THE EXPECTED NORM OF A MINIBATCH GRADIENT
To keep the notation clean, we define X, = (VC)(w) - vC(w)). We also recall for clarity that the expectation value E(...) is taken over all possible random shuffles of the indices i. Therefore,
B(I(Vv(()- -
E
X²)
(27)
1=
B(B - 1)
B(X
X.) +
B(X; -jzi))
(28)
B
B²
X; X,+ (B
1
Σ
X; Xj (29)
NB
B
N(N
1)
E=1
jzi
1 N
Σ X, X1 (B-1) N N
(X; X,(1 - -.5))
NB
BN(N 1) (=1
Note that we obtain Equation 28 by counting the number of diagonal and off-diagonal terms in the
sum in Equation 27 Next, we recall that Σ X;
- VC(w)) Therefore,
N
(B - 1)
N
- vC(w)) = NB i=1
X1 X₁
X₁ X
(30)
BN(N -
i=1
(B -1)
1
Σ
X; X;
(31)
NB
(N 11)
1=1
(N - B) ((w) (N-1) B
(32)
where I(w) = (1/N) E X, (1/N) E [[V VC(w)]] We can immediately
identify l(W) as the trace the empirical covariance matrix of the per-example gradients.
B A BACKWARD ERROR ANALYSIS FOR N-STEP SGD
Under n-step SGD, we apply n gradient descent updates on the same minibatch with bare learning rate a and batch size B. After n updates, we sample the next minibatch and repeat. For convenience, we define the effective learning rate na. After minibatch (n parameter updates).
Wi+1 Wi VV,(w)) - avŌ,(w; VV,((
(33)
= W: ((.( + (n/2)(n - 1))=VVC + O(n²a")
(34)
Wi - -VC,(w)) + (1/4)(1-1/n))) 100(1-1/n))) V (IVC((W)))))++) + +O(E)).
(35)
After epoch (including terms up second order in c),
eVColwo) + (1/4)(1 V V (IVCC((IIV) ( + 1/ (1VC)(
Wm wi
cVCm-1(um-)) + (1/4)(1-1/n)) -1
(IVCm-((4mm)))I)) + O(()). (36)
To simplify this expression, we note that Wi+1 = w; <vC((w: + (c2)). We can therefore re-use our earlier analysis from Section 2.2 of the main text (see Equation 13 for comparison) to obtain,
Wm
wo mCVC(() +
+ (1/4)(1 - 1//nc²
V (IVC(((b))1P + +((2).
(37)
12

---

Published as a conference paper at ICLR 2021
1.00
14
00-0)
Traning accrrsyy
(Ctull
68
2
0.00
=M=0)
64
0.5
te-
2")
69
2-2-1
(1-20
2nd1 2-1
gm1
2
2-11
201
2-1
Learning ratafiouubrttatoo scale
Leariing rato/Ragu arizotion scale
(a)
(h)
Figure 5: (a) When training on the original loss C(w), the training accuracy achieved within 6400 epochs is maximized when the learning rate f. 2-9 Smaller learning rates have not yet converged, while larger learning rates reach plateau. When training on the modified loss (m)d(w) with fixed learning rate = 2-9 and regularization coefficient A, we achieve high training accuracies when is small, but are unable to achieve high training accuracies when À is large. (b) We plot the value of the original loss ((w) at the end of training. Remarkably, the training losses obtained when training on the original loss with large learning rate similar to the training losses achieved when training
on the modified loss with small learning rate (e = 2 9) and a large regularization coefficient X.
Taking the expectation over possible batch orderings (see Equations 15 to 18), we obtain,
((wm) = wo (wV)(())) + m²c² 4
(IVCC(o))1
m
[VC1(W))11 + O(m3 e). (38)
m²n
Fixing ((w) = (WW and equating Equation 38 with the continuous modified flow in Equation 19 by setting E(wm) = w(me), we identify the modified flow i VCnson + O(m² where,
Casan(w) = C(w) 4mn
(39)
note that the modified losses of SGD and n-step SGD co- incide when n = 1. However for n-step SGD when n > 1, the strength of the implicit regularization term is proportional to the scale of the bare learning rate c/n, not the effective learning rate c.
Comparing Equation 39 to Equation 1
C TRAINING LOSSES
In Figure [(b) of Section 2.3 in the main text, we compared the test accuracies achieved when train- ing on the original loss ((w) range of learning rates c, the test accuracies achieved when training on the modified loss (m) at fixed learning rate e = 2 and range of regularization Figure 5 we provide the corresponding training accuracies, as well as the final values of the original loss (ww). Remarkably, large learning rates and large regu- larization coefficients achieve similar training accuracies and similar original losses. This suggests that the implicit regularization term in the modified loss of SGD (CsGn (w)) may help explain why the training accuracies and losses often exhibit plateaus when training with large learning rates.
coefficients For completeness,
D ADDITIONAL RESULTS ON FASHION-MNISS
In this section we provide additional experiments on the Fashion-MNIST dataset (Xiao et al. 2017), which comprises 10 classes, 60000 training examples and 10000 examples the test set. We con- sider simple fully connected MIP which comprises 3 nonlinear layers, each with width 4096 and ReLU activations, and a final linear softmax layer. We apply a simple data pipeline which first ap- plies per-image standardization and then flattens the input to a 784 dimensional vector. We do not apply data augmentation and we train using vanilla SGD without learning rate decay for all exper- iments. We perform seven training runs for each combination of hyper-parameted and show the mean performance of the best five (to ensure our results are not skewed by a single failed run). We batch size B 16 unless otherwise specified, and we do not weight decay.
We note that this model is highly over parameterized.. Unlike the Wide-ResNet we considered in the main text consistently achieve 100% training accuracy if the learning rate is not too large.
13

---

Published as a conference paper at ICLR 2021
0.92
0.92
0.92
eccuracy
0.91
Test Accuracy
Test Accurace
0.91
0.40
0.90
B
32
0.89
16
B
0.88
(c-2-1)
0.88
0.87
2'
0.87
10"1
10
0.82
10 :
10 130
learning rate/Regullatbattinn scate
FJB
(a)
(b)
(c)
Figure 6: (a) Tuning the regularization scale À has similar influence on test accuracy to tuning the learning rate e, although unlike our CIFAR-IO/Wide-RR ResNet experiments the optimal values of and do not coincide. (b) We obtain similar accuracies at different batch sizes if the ratio e/ B/B is constant. (c) We obtain similar accuracies at different batch sizes if the ratio (A/ B) constant.
699
SCCHRRCK
0.91
0.93
0.50
2.50
1-1
n-1 (1600 epacts)
D=1
0.59
2
0.89
D-4 (400 opocha)
n-2 (822 epochs)
Test
n=2
(200 opachs)
0.50
(100 opachss
n-1E
2
2
2
2
2
2 2 2 2nd 2 2
2
0.00
Pare rale (d)
Bern lewming rabm (o)
Training opocns
(a)
(b)
(c)
Figure 7: (a) We train with n-step SGD for 400 epochs, and achieve similar test accuracies for
different values of n, so long as the bare learning rate a is the same. This is consistent with backward error analysis of n-step SGD, but contradicts the predictions of the SDE analogy.
(b) We observe (c) Different values of achieve similar test accuracies at their optimal learning rate, across range of epoch budgets.
similar phenomenon when training for fixed number of parameter updates.
In Figure 6(a) we train for 400 epochs, and compare the effect of tuming the learning rate when training on the original loss, to the effect of tuning the explicit regularization strength À (with 2 2-)), As observed in the main text, tuning the explicit regularizer has similar effect on the test accuracy to tuning the learning rate. Surprisingly, the optimal values of e and differ by a factor of 8. However note that the optimal learning rate is = 2 2-5 while the explicit regularizer already achieves a very similar test accuracy at = 2-1 (just a factor of two larger), before reaching a higher maximum test accuracy À 2-2 In Figure 6(b) we train for 400 epochs on the original loss and compare the test accuracies achieved for a range of batch sizes at different learning rates. As observed the main text, the test accuracy is determined by the ratio of the learning rate the batch size. Meanwhile in Figure we plot the test accuracy achieved after training for 1.5 million steps on the modified loss with learning rate we find that the test accuracy achieved is primarily determined by the ratio of the regularization coefficient to the batch size, although smaller batch sizes also achieve slightly higher accuracies.
2-9 and regularization coefficient Once again,
Finally, in Figure 7 we train using n-step SGD (see Section 4 on the original loss at a range of bare learning rates a. In Figure [[] we train for 400 epochs, while in Figure 77b) we train for million updates. We recall that the SDE analogy predicts that the generalization benefit of n-SGD would be determined by the effective learning rate = na. By contrast, backward error analysis predicts that the generalization benefit for small learning rates would be controlled by the bare learning rate a, but that higher order terms may be larger for larger values of n. We find that the test accuracy in both figures is governed by the bare learning rate a, not the effective learning rate therefore these results are inconsistent with the predictions from the SDE analysis in prior work.
na, and
Note that Figure 7 has surprising implication. It suggests that, for this model, while there is largest stable bare learning rate we cannot exceed, we can repeatedly apply updates obtained on the same batch of training examples without suffering significant degradation in test accuracy. We speculate that this may indicate that the gradients of different examples in this pver-parameterized model are close to orthogonal (Sankararaman et al. 2020).
14


Reviews:
[{"review_id": "uJhk1VSHz-P", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "## Summary\n\nUsing backward error analysis, the paper argues that SGD with small but finite step sizes stays on the path of a gradient flow ODE of a modified loss, which penalizes the squared norms of the mini-batch gradients. This offers a possible explanation of the empirically observed positive effect of (relatively) large step sizes on generalization performance. The paper further contests previous findings based on a vanishing step size assumption.\n\n\n## Rating\n\nSimilar to several recent works, this paper tries to explain certain aspects of stochastic gradient descent using a continuous time approximation. In contrast to existing works, it explicitly accounts for the effect of finite step sizes, which I think is a very interesting direction and surfaces several interesting aspects. I also welcome and endorse the critical discussion of prior work based on infinitesimal step size assumptions. Overall, the paper was interesting and pleasant to read. To the very best of my knowledge, all mathematical derivations are technically correct.\n\nHowever—as the authors themselves note in their critique of SDE approximations to SGD—the devil is in the details with continuous time approximations. In my opinion, that makes is absolutely crucial to discuss the scope of the results carefully and transparently, including a critical discussion on assumptions made and simplifications that go into the continuous-time model. In my opinion, this paper fails to deliver that, which is why I recommend rejection. Below, I am asking for clarification on various points and would encourage the authors to respond to the major points in the rebuttal phase.\n\n\n## Major Comments\n\n\n1) The main result says that the *expected* SGD iterate after a *single* epoch lands close to the path of a gradient flow ODE on a modified loss. Unless I am missing something, this fundamentally fails to capture the behavior over multiple epochs. The analysis only guarantees that, from any given starting point $\\omega_0$, the expected iterate after one epoch of SGD ends up close to the ODE path starting from $\\omega_0$. Unless I am missing something, this does *not* imply that two epochs of SGD starting from $\\omega_0$ end up on that path. We can not simply chain two epochs together: The first epoch only stays on the path in expectation, but any realization of that random variable will deviate from the path, which affects the initial condition of the next epoch. Intuitively, one needs to get a handle on the variance of the iterate as well in order to give guarantees for multiple epochs. Is this understanding correct? If so, to what extent can insights about a single epoch of SGD be transferred to practical settings?\n\n2) Comment (1) hints at a larger (but vague) point that the paper is trying to characterize a *stochastic* optimization procedure with a solution of a *deterministic* gradient flow ODE. It does so by focusing on the *expectation* of the iterate, which might be an approach to highlight certain aspects, but it will never give a full picture. Why wouldn’t we also be interested in the covariance of the iterates? The limitations of this characterization should be discussed thoroughly in the paper.\n\n3) In Section 2, the composition of the minibatches is assumed to be fixed and the randomness only comes from their ordering. The paper says: \"It is standard practice to shuffle the dataset once per epoch, but this step does not affect our analysis and we omit it for brevity.“ I don’t think that statement is justified with respect to the result in Eq. (1), given that the modified loss depends on the minibatch composition. Therefore, would we reshuffle the dataset after each epoch, the modified loss would change from one epoch to the next. Later, in Section 3, the expectation is additionally taken over the composition of the batches. Why is the result presented in these two distinct steps? None of the key findings of the paper seems to rely on the intermediate fixed-composition result. It also doesn’t reflect the common practice of reshuffling the entire dataset and then traversing it, which simultaneously randomizes the composition and ordering of batches. So why not give the result of Eq. (22) directly? It is also the more intuitive result, invoking the trace of the gradient covariance matrix, which also appears in prior work on continuous time approximations of SGD.\n\n4) While the analysis tries to account for finite step sizes, it still seems to assume step sizes that are orders of magnitude smaller than those used in practice. In particular, when going from Eq. (12) to Eq. (13), each minibatch cost function is equated with its second-order Taylor approximation around the starting point $\\omega_0$. This is a *drastic* approximation and I don’t see any justification for why this should be anywhere near accurate for practical settings. For large datasets and moderate batch sizes, the number of updates in one epoch will be in the thousands. For realistic step size choices, a second-order Taylor expansion around the starting point will probably be rather poor after a handful of SGD updates, no?\n\n5) The paper strongly emphasizes the assumption of sampling data points without replacement. While sampling without replacement is indeed the usual setting in practice, most of the stochastic optimisation literature builds on the assumption of sampling with replacement. And to my knowledge, no major differences (in terms of generalization performance) have been reported in the literature between the two approaches.\n    a) Can the analysis presented in the paper be extended to setting of sampling with replacement? It seems to me that this should be straight-forward. Equations (12) and (13) should hold also when each minibatch is obtained from sampling with replacement. In that case, the expectation of the second-order correction term should directly give a result akin to Eq. (22). If that is in fact possible, it should definitely be added to the paper.\n    b) If that is not possible, what prevents the application and is this a technicality or would you actually expect substantially different behavior in terms of generalization?\n    c) It would also have been nice to see the experiments repeated with sampling with replacement to check empirically whether the findings hold in that case?\n\n6) Something that bugs me from an optimization perspective is that the smoothness properties of the problem do not enter this analysis at all. For example, you write (near the bottom of page 4) that “our analysis assumes $m\\epsilon = N\\epsilon / B$ is small.” However, any given loss function $C(w)$ can be rescaled by a constant $M\\gg 1$ while scaling the step size with $1/M$. This leaves the behavior of SGD unaffected while making the step size arbitrarily small. Why does that not enter into the analysis? It probably relates to my comment (4), seeing that the step sizes are assumed to be so small that they are not restricted by the smoothness of the function.\n\n\n## Minor Comments\n\n7) The paper derives the implicit regularizer and provides empirical evidence that it can partially explain the benefits of large step sizes for generalization. However, very little attention is given to the regularization term itself and to the question *why* this regularizer might be beneficial. The only comment speaking to that is that the regularizer penalizes “sharp” regions. I would like to see this discussion expanded and connected to the recent literature.\n\n8) At the end of page 6, you write about the large batch size regime and say that the “we expect the optimal learning rate to be independent of the batch size in this limit.” It would have been great to substantiate that conjecture with an experiment and/or to refer to specific experiments done in prior work.\n\n9) You repeatedly use the phrase “small but finite learning rates”. If my understanding is correct, that has phrase has a very precise meaning in the context of this work, namely that terms of order $O(\\epsilon^3)$ are vanishingly small while terms that a quadratic or linear in $\\epsilon$ can not be ignored. (This is in contrast to prior work that also ignores quadratic terms.) Maybe this could be stated clearly the first time you use this phrase.\n\n\n## Typos / Style\n\n- I think you should capitalize references to sections, equations, figures, et cetera.\n- The bib file could really need some love. You are citing the arXiv versions for several papers that have been published in peer-reviewed venues. Capitalization in paper titles is messed up (e.g., “sgd”).\n\n## Edit after Rebuttal\n\nI thank the authors for their engagement with my review. Many of my comments and questions have been resolved and, consequently, I have increase my score and **recommend accepting this paper.**", "extra_scores": {}, "extra_text": {}}, {"review_id": "f2ookzGQv_r", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: To analyze why the generalization error of SGD with larger learning rates achieves better test error, this paper analyzes the implicit regularization of SGD (with a finite step size) via a first order backward error analysis. Under this analysis the paper shows that the mean position of SGD with $m$ minibatches effectively follows the flow according to Eq (20) for a small but finite step size, while GD effectively follows the last inline equation in section 2.1. The paper shows empirically on an image classification task that by explicitly including the (implicit SGD) regularizer, SGD on the modified loss behaves similarly to using a larger learning rate when evaluating on the test set. The paper then extends this results to consider varying the batch size in section 3, showing that for small batchsizes the implicit regularization scales with the ratio of learning rate and batchsize $\\epsilon/B$. Finally in section 4, the paper analyzes SGD when for each sampled minibatch in an epoch, we apply $n$ gradient steps with a stepsize $\\epsilon/n$ and show that performance degrades as $n$ increases, suggesting that the benefit of SGD with larger learning rates is due to the implicit regularizer and not the temperature of an associated SDE. \n\nThis paper is clearly written and well edited. I find the main result and the analysis technique interesting and novel. Although the experiments are well explained and help support the theory developed, there is only one experiment setting making it difficult to believe strong general claims such as those in section 4. I do have concerns about equating the \"mean\" behavior of SGD with the actual behavior of SGD and. \n \nRecommendation:\nI recommend accepting this paper. As it currently stands, this paper is borderlin on the acceptance threshold for me. I like the novel use of the backward error analysis to gain insight into the behavior of SGD and I believe it would be of interest to ICLR readers. My main concerns are the papers' narrow focus on the mean behavior of SGD and the single experiment setting used to validate results. I would much more strongly support this paper if the theoretical analysis was stronger (e.g. analyzing the variance of individual SGD flows/regularizers to the mean SGD flow/regularizer) or if more experiments (in different settings) supported the results.  \n\nQuestions: \nIf we don't take the expectation over $\\xi(m)$ in Section 2.2, the theory suggests that there exist a (random) modified flow for each (random) ordering of minibatches $\\hat{C}_0, \\ldots, \\hat{C}_m$ by equating equations (14) and (19). The main result Eq (20) would correspond to the expected value over the (random) modified flow. I believe this paper would be much stronger if there was some discussion of how the variance / deviations of these random flows from the mean flow (i.e the variance of $\\xi(m)$) affects the implicit regularization and how this scales with batch size and properties of the loss. Would the implicit regularization break down for some experiments? \n\nIs the assumption that $m \\epsilon$ is small reasonable (so that we can ignore the higher order $O(m^3 \\epsilon^3)$ terms in the analysis)? Isn't $m = N/B$ the number of updates per epochs very large in practice since $N >> B$?\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "N0BP1f9P5ZW", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper studies an implicit regularization mechanism of finite learning rate SGD by introducing explicitely a regularization term, using the framework of backward analysis.\n\nThey theoretically motivate their analysis, then empirically demonstrate it on CIFAR-10 using a Wide ResNet architecture.\n\nThis extends a previous (Barrett and Dherin, preprint) analysis of GD using the same framework, but limited to full batch GD. Noticeably, this new analysis using minibatch GD highlights an additional regularization of the trace of the covariance of per-example gradients.\n\nIn sec 2, however, I think it should be made clear that the setup is slightly different from minibatch GD, even when trained for a single epoch, in that there is an expectation accross permutations of sequences of minibatches. Can you discuss this assumption a bit more?\n\nIn terms of experiments, it would be useful to include other architecture/tasks, even toyish, in order to appreciate the generality of the empirical evaluation.\n\nOverall, I think this contributes new interesting insights which are very relevant for studying minibatch GD in deep learning.", "extra_scores": {}, "extra_text": {}}, {"review_id": "55Fkhakb3Ah", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper analyzes the implicit regularization in SGD with finite learning rates via backward error analysis. The modified flow introduced in this paper better approximates the practical behavior of SGD as it does not require vanishing learning rates and it allows to use random shuffling in stead of i.i.d sampling. The numerical experiments validates the existence of the implicit regularization and how it affects the generalization of the model trained by SGD. The difference from SDE analysis is also discussed.\n\nReason for score:\n1. The paper is well organized. Specially, I enjoy reading section II. The tool of backward error analysis and the derivation of the implicit regularization in SGD flow are introduced clearly and concisely. The analysis is based on random shuffling instead of i.i.d sampling matches the practical use of SGD.\n2. The numerical experiments are very convincing. The consistency of SGD with larger lr and SGD with smaller lr plus explicit regularization validates the results of theoretical analysis. The numerical experiments also provide some insights into tuning hyper parameters such as learning rate and batch size.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Dear authors,

all reviewers found many interesting contributions in your paper and also pointed out some minor/major issues. In your rebuttal discussions, you addressed most of them to their satisfaction and I hope you will incorporate them carefully also in your final submission.

I hence recommend accepting this paper


Decision: Accept (Poster)