O N THE O RIGIN OF I MPLICIT R EGULARIZATION IN S TOCHASTIC G RADIENT D ESCENT
A BSTRACT
For inﬁnitesimal learning rates, stochastic gradient descent (SGD) follows the path of gradient ﬂow on the full batch loss function.
However moderately large learning rates can achieve higher test accuracies, and this generalization beneﬁt 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 shufﬂing, the mean SGD iterate also stays close to the path of gradient ﬂow if the learning rate is small and ﬁnite, but on a modiﬁed loss .
This modiﬁed 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.
I NTRODUCTION
In the limit of vanishing learning rates, stochastic gradient descent with minibatch gradients (SGD) follows the path of gradient ﬂow on 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 beneﬁt is not explained by convergence rate bounds (Ma et al., 2018; Zhang et al., 2019), because it arises even for large compute budgets for which smaller learning rates often achieve lower training losses (Smith et al., 2020).
Although many authors have studied this phenomenon (Jastrz˛ebski et al., 2018; Smith & Le, 2018; Chaudhari & Soatto, 2018; Shallue et al., 2018;
Park et al., 2019; Li et al., 2019; Lewkowycz et al., 2020), it remains poorly understood, and is an important open question in the theory of deep learning.
In a recent work, Barrett & Dherin (2021) analyzed the inﬂuence of ﬁnite 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 et al., 2006).
The key insight of backward error analysis is that we can describe the bias introduced when integrating an ODE with ﬁnite step sizes by introducing an ancillary modiﬁed ﬂow .
This modiﬁed ﬂow is derived to ensure that discrete iterates of the original ODE lie on the path of the continuous solution to the modiﬁed ﬂow.
Using this technique, the authors show that if the learning rate (cid:15) is not too large, the discrete iterates of GD lie close to the path of gradient ﬂow on a modiﬁed loss (cid:101) C GD ( ω ) =
C ( ω ) + ( (cid:15)/ 4) ||∇ C ( ω ) || 2 .
This modiﬁed loss is composed of the original loss C ( ω ) and an implicit regularizer proportional to the learning rate (cid:15) which penalizes the euclidean norm of the gradient.
However these results only hold for full batch GD, while in practice SGD with small or moderately large batch sizes usually achieves higher test accuracies (Keskar et al., 2017; Smith et al., 2020).
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 ﬁnite learning rates, the mean SGD iterate after one epoch, averaged over all possible sequences of minibatches, lies close to the path of gradient ﬂow on a second modiﬁed loss (cid:101)
C SGD ( ω ) , which we deﬁne in equation 1 .
This new modiﬁed 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 modiﬁed losses can have different local and global minima.
Our analysis therefore helps explain both why ﬁnite 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 conﬁrm 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 sufﬁciently diverse, then the expected magnitude of the implicit regularization term of SGD is proportional to the ratio of the learning rate to the batch size (Goyal et al., 2017; Smith et al., 2018).
We note that many previous authors have sought to explain the generalization beneﬁt of SGD using an analogy between SGD and stochastic differential equations (SDEs) (Mandt et al., 2017; Smith &
Le, 2018; Jastrz˛ebski et al., 2018; Chaudhari & 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 (Yaida, 2019) and therefore misses the generalization beneﬁts of ﬁnite learning rates which we identify in this work.
An important exception is Li et al.
(2017), who applied backward error analysis to identify a modiﬁed SDE which holds when the learning rate is ﬁnite.
However this work still relies on the assumption that minibatches are 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 our main result.
We deﬁne the cost function over parameters ω as
C ( ω ) =
(1 /N )
(cid:80) Nj =
1 C j ( ω ) , which is the mean of the per-example costs C j ( ω ) , where N denotes the training set size.
Gradient ﬂow follows the ODE ˙ ω
= −∇ C ( ω ) , while gradient descent computes discrete updates ω i +1 = ω
i − (cid:15) ∇ C ( ω i ) , where (cid:15) is the learning rate.
For simplicity, we assume that the batch size B perfectly splits the training set such that N % B = 0 , where % denotes the modulo operation, and for convenience we deﬁne the number of batches per epoch m =
N/B .
We can therefore re-write the cost function as a sum over minibatches C ( ω ) =
(1 /m ) (cid:80) m − 1 k =
ˆ C k ( ω ) , where the minibatch cost ˆ C k ( ω ) =
(1 /B )
(cid:80) kB + B j = kB +1 C j ( ω ) .
In order to guarantee that we sample each example precisely once per epoch, we deﬁne SGD by the discrete update
ω
i +1 = ω
i − (cid:15) ∇ ˆ C
i % m ( ω i ) .
Informally, our main result is as follows.
After one epoch, the mean iterate of SGD with a small but ﬁnite learning rate (cid:15) , averaged over all possible shufﬂes of the batch indices, stays close to the path of gradient ﬂow on a modiﬁed loss ˙ ω = −∇ (cid:101)
C SGD ( ω ) , where the modiﬁed loss (cid:101)
C SGD is given by:
[Equation]
We emphasize that our analysis studies the mean evolution of SGD, not the path of individual trajectories.
The modiﬁed loss (cid:101) C SGD ( ω ) is composed of the original loss C ( ω ) and an implicit regularizer C reg ( ω ) =
(1 / 4 m ) (cid:80) m − 1 k =
||∇ ˆ C k ( ω ) || 2 .
The scale of this implicit regularization term is proportional to the learning rate (cid:15) , and it penalizes the mean squared norm of the gradient evaluated on a batch of B examples.
To help us compare the modiﬁed losses of GD and SGD, we can expand,
[Equation]
We arrive at Equation 2 from Equation 1 by noting that (cid:80) m − 1 i =0 ( ∇ ˆ C i ( ω ) − ∇ C ( ω ))
= 0 .
In the limit B → N , we identify the modiﬁed loss of gradient descent, (cid:101) C GD = C ( ω ) + ( (cid:15)/ 4) ||∇ C ( ω ) || 2 , which penalizes “sharp” regions where the norm of the full-batch gradient ( ||∇ C ( ω ) || 2 ) is large.
However, as shown by Equation 2, the modiﬁed 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 ( ||∇ ˆ C ( ω ) − ∇ C ( ω ) || 2 ) are large (Wu et al., 2018).
Although global minima of C ( ω ) are global minima of (cid:101) C GD ( ω ) , global minima of C ( ω ) may not be global (or even local) minima of (cid:101)
C SGD ( ω ) .
Note however that C ( ω ) and (cid:101) C SGD ( ω ) 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 1), and we conﬁrm 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 conﬁrm Equation 1 satisﬁes the linear scaling rule between learning rate and batch size (Goyal et al., 2017).
In Section 4, we provide additional experiments which challenge the prevailing view that the generalization beneﬁt of small batch SGD arises from the temperature of an associated SDE (Mandt et al., 2017; Park et al., 2019).
A B ACKWARD E RROR A NALYSIS OF S TOCHASTIC G RADIENT D ESCENT
Backward error analysis has great potential to clarify the role of ﬁnite learning rates, and to help identify the implicit biases of different optimizers.
We therefore give a detailed introduction to the core methodology in Section 2.1, before deriving our main result in Section 2.2.
In Section 2.3, we conﬁrm empirically that the implicit regularizer can enhance the test accuracy of deep networks.
2. A N I NTRODUCTION TO B ACKWARD E RROR A NALYSIS
In numerical analysis, we often wish to integrate ODEs of the form
˙ ω = f ( ω ) .
This system usually cannot be solved analytically, forcing us to simulate the continuous ﬂow with discrete updates, like the Euler step ω ( t + (cid:15) )
≈ ω ( t )
+ (cid:15)f ( ω ( t )) .
However discrete updates will introduce approximation error when the step size (cid:15) is ﬁnite.
In order to study the bias introduced by this approximation error, we assume the learning rate (cid:15) is relatively small, and introduce a modiﬁed ﬂow
˙ ω =
(cid:101) f ( ω ) , where,
[Equation]
The modiﬁed ﬂow of (cid:101) f
( ω ) is equal to the original ﬂow of f ( ω ) when (cid:15) → 0 , but it differs from the original ﬂow if (cid:15) is ﬁnite.
The goal of backward error analysis is to choose the correction terms f i ( ω ) such that the iterates obtained from discrete updates of the original ﬂow with small ﬁnite step sizes lie on the path taken by the continuous solution to the modiﬁed ﬂow with vanishing step sizes.
The standard derivation of backward error analysis begins by taking a Taylor expansion in (cid:15) of the solution to the modiﬁed ﬂow ω ( t + (cid:15) ) .
We obtain the derivatives of ω ( t + (cid:15) ) recursively using the modiﬁed ﬂow equation
˙ ω = (cid:101) f
( ω ) (see Hairer et al.
, and we identify the correction terms f i ( ω ) by ensuring this Taylor expansion matches the discrete update (
e.g., ω t
+1 = ω
t
+ (cid:15)f ( ω t ) ) for all powers of (cid:15) .
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 modiﬁed ﬂow by considering the combined inﬂuence of an inﬁnite number of discrete steps in the limit of vanishing learning rates, and then we will compare this continuous path to either a single step of GD or a single epoch of SGD.
Imagine taking n Euler steps on the modiﬁed ﬂow (cid:101) f
( ω ) with step size α , ω t + n
[Equation]
We arrived at Equation 6 by taking the Taylor expansion of (cid:101) f
and then counting the number of terms of type ∇ (cid:101) f ( ω t ) (cid:101) f ( ω t ) using the formula for an arithmetic series.
Note that we assume ∇ (cid:101) f exists.
Next, to ensure ω t + n in Equation 6 coincides with the solution ω ( t + (cid:15) ) of the continuous modiﬁed
ﬂow ˙ ω =
(cid:101) f ( ω ) for small but ﬁnite (cid:15) , we let the number of steps n → ∞ while setting α = (cid:15)/n , ω ( t
[Equation]
) =
[Equation]
t )
[Equation]
(cid:15)
[Equation]
We have replaced (cid:101) f ( ω ) with its deﬁnition from Equation 3.
As we will see below, Equation 8 is the key component of backward error analysis, which describes the path taken when integrating the continuous modiﬁed ﬂow (cid:101) f ( ω ) with vanishing learning rates over a discrete time step of length (cid:15) .
Notice that we have assumed that the Taylor expansion in Equation 8 converges, while the higher order terms at O ( (cid:15) 3 ) will contain higher order derivatives of the original ﬂow f ( ω ) .
Backward error analysis therefore implicitly assumes that f ( ω ) is an analytic function in the vicinity of the current parameters ω .
We refer the reader to Hairer et al.
(2006) for a detailed introduction.
Gradient descent: As a simple example, we will now derive the ﬁrst order correction f 1 ( ω ) of the modiﬁed ﬂow for GD.
First, we recall that the discrete updates obey ω
i +1 = ω
i − (cid:15) ∇ C ( ω i ) , and we therefore ﬁx f ( ω ) =
−∇ C ( ω ) .
In order to ensure that the continuous modiﬁed ﬂow coincides with this discrete update, we need all terms at O ( (cid:15) 2 ) and above in Equation 8 to vanish.
At order (cid:15) 2 , this implies that f 1 ( ω ) + (1 / 2) ∇∇
[Equation]
= 0 , which yields the ﬁrst order correction,
[Equation]
We conclude that, if the learning rate (cid:15) is sufﬁciently small such that we can neglect higher order terms in Equation 3, then the discrete GD iterates lie on the path of the following ODE,
[Equation]
Equation 11 corresponds to gradient ﬂow on the modiﬁed loss, (cid:101)
[Equation]
2. B ACKWARD E RROR A NALYSIS AND S TOCHASTIC G RADIENT D ESCENT
We now derive our main result (Equation 1).
As described in the introduction, we assume N % B = 0 , where N is the training set size, B is the batch size, and % denotes the modulo operation.
The number of updates per epoch m = N/B , and the minibatch costs ˆ C k ( ω ) =
[Equation]
1 /B )
(cid:80) kB [Equation] B j = kB [Equation]
C j ( ω ) .
SGD with constant learning rates obeys ω
i +1 = ω
i − (cid:15) ∇ ˆ C
i % m ( ω i ) .
It is standard practice to shufﬂe the dataset once per epoch, but we omit this step here and instead perform our analysis over a single epoch.
In Equation 6 we derived the inﬂuence of n Euler steps on the ﬂow (cid:101) f ( ω ) with step size α .
Following a similar approach, we now derive the inﬂuence of m SGD updates with learning rate (cid:15) ,
[Equation]
The error in Equation 14 is O ( m 3 (cid:15) 3 ) since there are O ( m 3 ) terms in the Taylor expansion proportional to (cid:15) 3 .
Notice that a single epoch of SGD is equivalent to a single GD update with learning rate m(cid:15) up to ﬁrst order in (cid:15) .
Remarkably, this implies that when the learning rate is sufﬁciently small, there is no noise in the iterates of SGD after completing one epoch.
For clarity, this observation arises because we require that each training example is sampled once per epoch.
However the second order correction ξ ( ω ) =
(cid:80) m − 1 j =0 (cid:80) k<j ∇∇ ˆ C j ( ω ) ∇ ˆ C k ( ω ) does not appear in the GD update, and it is a random variable which depends on the order of the mini-batches.
In order to identify the bias introduced by SGD, we will evaluate the mean correction E ( ξ ) , where we take the expectation across all possible sequences of the (non-overlapping) mini-batches { ˆ C 0 , ˆ C 1 , ..., ˆ C m − 1 } .
Note that we hold the composition of the batches ﬁxed, averaging only over their order.
We conclude that,
[Equation]
For clarity, in Equation 15 we exploit the fact that every sequence of batches has a corresponding sequence in reverse order.
Combining Equations 14 and 17, we conclude that after one epoch,
[Equation]
Having identiﬁed the expected value of the SGD iterate after one epoch E ( ω m ) (for small but ﬁnite learning rates) [Equation] we can now use this expression to identify the corresponding modiﬁed ﬂow.
First, we set f ( ω ) =
−∇ C ( ω ) , t = 0 , ω (0) =
ω 0 , and let (cid:15) → m(cid:15) in Equations 3 and 8 to obtain,
[Equation]
Next, we equate Equations 18 and 19 by setting ω ( m(cid:15) )
=
E ( ω m ) .
We immediately identify the ﬁrst order correction to the modiﬁed ﬂow
[Equation]
[Equation]
||
We therefore conclude that, after one epoch, the expected SGD iterate E ( ω m )
= ω ( m(cid:15) )
+ O ( m 3 (cid:15) 3 ) , where ω (0) =
ω 0 and ˙ ω
= −∇ C ( ω ) + m(cid:15)f 1 ( ω ) .
Simplifying, we conclude ˙ ω = −∇ (cid:101)
C SGD ( ω ) , where,
[Equation]
Equation 20 is identical to Equation 1, and this completes the proof of our main result.
We emphasize that (cid:101) C SGD assumes a ﬁxed set of minibatches { ˆ C 0 , ˆ C 1 , ..., ˆ C m − 1 } .
We will evaluate the expected modiﬁed loss after shufﬂing the dataset and sampling a new set of minibatches in Section 3.
R EMARKS ON THE A NALYSIS
The phrase “for small ﬁnite learning rates” has a precise meaning in our analysis.
It implies (cid:15) is large enough that terms of O ( m 2 (cid:15) 2 ) may be signiﬁcant, but small enough that terms of O ( m 3 (cid:15) 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; Jastrz˛ebski 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 ( (cid:15) ) , while the bias arises at O ( (cid:15) 2 ) .
By contrast, in our analysis each example is sampled once per epoch, and both the variance and the bias arise at O ( (cid:15) 2 ) (for simplicity, we assume m = N/B is constant).
We therefore anticipate that the variance will play a less important role than is commonly supposed.
Furthermore, we can construct a speciﬁc sequence of minibatches for which the variance at O ( m 2 (cid:15) 2 ) vanishes, such that the evolution of a speciﬁc training run will coincide exactly with gradient ﬂow on the modiﬁed loss of Equation 1 for small ﬁnite learning rates.
To achieve this, we perform two training epochs with the sequence of minibatches ( ˆ C 0 , ˆ C 1 , ..., ˆ C m − 1 , ˆ C m − 1 , ..., ˆ C 1 , ˆ C 0 ) (i.e., the second epoch iterates through the same set of minibatches as the ﬁrst 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 (cid:80) k<j by a (cid:80) k (cid:54) = j ).
The key limitation of our analysis is that we assume m(cid:15)
=
N(cid:15)/B is small, in order to neglect terms at O ( m 3 (cid:15) 3 ) .
This is an extreme approximation, since we typically expect that N/B is large.
Therefore, while our work identiﬁes the ﬁrst order correction to the bias arising from ﬁnite learning rates, higher order terms in the modiﬁed ﬂow 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 O ( (cid:15) 2 ) (for an exception, see Li et al.
Furthermore, as we show in Section 3, the learning rate often scales proportional to the batch size, such that (cid:15)/B is constant (Goyal et al., 2017; McCandlish et al., 2018).
Therefore the accuracy of our approximations does not necessarily degrade as 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 beneﬁt of ﬁnite learning rate SGD for Wide-ResNets trained on CIFAR-10.
We note that to achieve the highest test accuracies, practitioners usually decay the learning rate during training.
Under this scheme, the modiﬁed loss would change as training proceeds.
However it is widely thought that the generalization beneﬁt 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 inﬂuence of ﬁnite learning rates on training.
The implicit regularization term may not be beneﬁcial in all models and datasets.
2. A N E MPIRICAL E VALUATION OF THE M ODIFIED L OSS
In order to conﬁrm that the modiﬁed loss (cid:101) C SGD ( ω ) can help explain why large learning rates enhance generalization, we now verify empirically that the implicit regularizer inherent in constant learning rate SGD, C reg ( ω ) =
(1 / 4 m ) (cid:80) m − 1 k =
||∇ ˆ C k ( ω ) || 2 , can enhance the test accuracy of deep networks.
To this end, we train the same model with two different (explicit) loss functions.
The ﬁrst loss function C ( ω ) represents the original loss, while the second C mod ( ω ) =
C ( ω ) + λC reg ( ω ) is obtained from the modiﬁed loss (cid:101)
C SGD ( ω ) by replacing the learning rate (cid:15) with an explicit regular-
[Figure]
ization coefﬁcient λ .
Notice that C mod ( ω ) = (1 /m ) (cid:80) m − 1 k =0 (cid:16) ˆ C k ( ω ) + ( λ/ 4) ||∇ ˆ C k ( ω ) || 2 (cid:17) , which ensures that it is straightforward to minimize the modiﬁed loss C mod ( ω ) with minibatch gradients.
Since the implicit regularization term C reg ( ω ) is expensive to differentiate (typically 5-10x overhead), we consider a 10-1 Wide-ResNet model (Zagoruyko & Komodakis, 2016) for classiﬁcation on CIFAR-10.
To ensure close agreement with our theoretical analysis, we train without batch normalization using SkipInit initialization (De & Smith, 2020).
We train for 6400 epochs at batch size 32 without learning rate decay using SGD without Momentum.
We use standard data augmentation including crops and random ﬂips, and we use weight decay with L 2 coefﬁcient 5 × 10 − 4 .
We emphasize that, since we train using a ﬁnite (though very large) compute budget, the ﬁnal 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.
In Figure 1(a), we compare two training runs, one minimizing the modiﬁed loss C mod ( ω ) with λ = 2 − 6 , and one minimizing the original loss C ( ω ) .
For both runs we use a small constant learning rate (cid:15)
= 2 − 9 .
As expected, the regularized training run achieves signiﬁcantly higher test accuracies late in training.
This conﬁrms that the implicit regularizer, which arises as a consequence of using
SGD with ﬁnite learning rates, can also enhance the test accuracy if it is included explicitly in the loss.
In Figure 1(b), 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 we take the highest test accuracy achieved during the entire training run.
We use a ﬁxed learning rate (cid:15) =
2 − 9 for all λ .
For comparison, we also provide the test accuracy achieved with the original loss C ( ω ) for a range of learning rates (cid:15) (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 modiﬁed loss λ opt = 2 −is equal to the optimal learning rate on the original loss (cid:15) opt = 2 − 6 .
Meanwhile when λ → 0 the performance of the modiﬁed loss approaches the performance of the original loss at (cid:15) =
2 − 9 (dotted green line).
We provide the corresponding training accuracies in appendix C. Finally, in Figure 1(c), we provide the values of the implicit regularizer C reg ( ω ) at the end of training.
As predicted by our analysis, training with larger learning rates reduces the value of the implicit regularization term.
In Figure 2, we take the same 10-1 Wide-ResNet model and provide the mean training and test accuracies achieved at a range of learning rates for two regularization coefﬁcients (following the experimental protocol above).
In Figure 2(a), we train on the original loss C ( ω ) ( λ = 0 ), while in Figure 2(b), we train on the modiﬁed loss C mod ( ω ) with regularization coefﬁcient λ = 2 − 6 .
From Figure 2(a), when λ =
there is a clear generalization beneﬁt to large learning rates, as the learning rate that maximizes test accuracy ( 2 − 6 ) is 16 times larger than the learning rate that maximizes training accuracy ( 2 − 10 ).
However in Figure 2(b) with λ = 2 − 6 , the learning rates that maximize the test and training accuracies are equal ( 2 − 8 ).
This suggests that when we include the implicit regularizer explicitly in the loss, the generalization beneﬁt of large learning rates is diminished.
I MPLICIT R EGULARIZATION AND THE B ATCH S IZE
In Section 2.2, we derived the modiﬁed loss by considering the expected SGD iterate after one epoch.
We held the composition of the batches ﬁxed, averaging only over the order in which the batches
[Figure] [Figure]
are seen.
This choice helped make clear how to 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 modiﬁed loss after randomly shufﬂing the dataset and sampling a new set of m non-overlapping minibatches { ˆ C 0 , ˆ C 1 , ..., ˆ C m − 1 } .
Since the minibatch losses ˆ C i ( ω ) are all identically distributed by symmetry, we recall Equation 2 and conclude that,
[Equation]
(21) where ˆ C ( ω ) denotes a batch of B non-overlapping examples, drawn randomly from the full dataset.
To simplify equation 21, we prove in appendix A that E ( ||∇ ˆ C ( ω ) − ∇ C ( ω ) || 2 )
=
( N − B ) ( N − 1) Γ( ω ) B , where Γ( ω ) =
(1 /N )
(cid:80) Ni =
||∇ C
i ( ω ) − ∇ C ( ω ) || 2 .
We therefore obtain,
[Equation]
Note that Γ( ω ) is the trace of the empirical covariance matrix of the per-example gradients.
We have not assumed that the minibatch gradients are Gaussian distributed, however if the per-example gradients are heavy tailed (Simsekli et al., 2019) then Γ( ω ) may diverge, in which case the expected value of the modiﬁed loss is ill-deﬁned.
Equation 22 shows that the implicit regularization term of SGD has two contributions.
The ﬁrst term is proportional to the learning rate (cid:15) , and it penalizes the norm of the full batch gradient.
The second term is proportional to the ratio of the learning rate to the batch size (cid:15)/B (assuming N (cid:29) B ), and it penalizes the trace of the covariance matrix.
To interpret this result, we assume that the minibatch gradients are diverse, such that (Γ( ω ) /B )
(cid:29) ||∇ C ( ω )
|| 2 .
This assumption guarantees that increasing the batch size reduces the error in the gradient estimate.
In this limit, the second term above will dominate, and therefore different batch sizes will experience the same implicit regularization so long as 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 with a constant learning rate, for a range of learning rates (cid:15) , regularization coefﬁcients λ and batch sizes B .
As expected, in Figure 3(a), training on the original loss C ( ω ) for 6400 epochs, we see that different batch sizes achieve similar test accuracies so long as the ratio (cid:15)/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., 2017; Smith & Le, 2018; Jastrz˛ebski et al., 2018;
Zhang et al., 2019).
To conﬁrm that this behaviour is consistent with the modiﬁed loss, in Figure 3(b) we ﬁx the learning rate (cid:15)
= 2 − 9 and train on C mod ( ω ) at a range of regularization strengths λ for 10 million steps.
As expected, different batch sizes achieve similar test accuracy so long as the ratio λ/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 ﬁrst 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 identiﬁed by Barrett &
Dherin (2021).
We expect the optimal learning rate to be independent of the batch size in this limit, as observed by McCandlish et al.
and Smith et al.
Convergence bounds also predict a transition between a small batch regime where the optimal learning rate (cid:15) ∝ 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 compliments these claims by explaining why similar conclusions hold when maximizing test accuracy.
F INITE L EARNING R ATES AND S TOCHASTIC D IFFERENTIAL E QUATIONS
In the previous two sections, we argued that the use of ﬁnite learning rates and small batch sizes introduces implicit regularization, which can enhance the test accuracy of deep networks.
We analyzed 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 be understood by interpreting small batch SGD as the discretization of an SDE (Mandt et al., 2017; Smith & Le, 2018; Jastrz˛ebski et al., 2018; Park et al., 2019).
In this section, we compare this popular perspective with our main results from Sections 2 and 3.
To brieﬂy recap, in the SDE analogy a single gradient update is given by ω i +
= ω
i − (cid:15) ∇ ˆ C ( ω i ) , where ˆ C denotes a 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 (cid:29) B (cid:29) 1 and that the gradients are not heavy tailed, the central limit theorem is applied to model the noise in an update by a Gaussian noise source ξ whose covariance is inversely proportional to the batch size:
[Equation]
This assumes E ( ξ i ) = 0 and E ( ξ i ξ Tj ) =
F ( ω ) δ ij , where F ( ω ) is the covariance matrix of the per example gradients, and we deﬁne the “temperature” T = (cid:15)/B .
The SDE analogy notes that Equation 23 is identical to the Euler discretization of an SDE with step size (cid:15) and temperature T (Gardiner et al., 1985).
Therefore one might expect the SGD iterates to remain close to this underlying SDE in the limit of small learning rates ( (cid:15) → 0 ).
In this limit, the temperature deﬁnes the inﬂuence of mini-batching on the dynamics, and it is therefore often assumed that the temperature also governs the generalization beneﬁt of SGD (Smith & Le, 2018; Jastrz˛ebski 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 ﬁrst order in (cid:15) 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 (cid:15) → 0 , but in this limit the temperature T → 0 and SGD converges to gradient ﬂow (Yaida, 2019).
We must use a ﬁnite learning rate to preserve a ﬁnite temperature, but at any ﬁnite 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 beneﬁt of SGD arises from ﬁnite learning rates, not the temperature of an associated stochastic process.
First, we introduce a modiﬁed SGD update rule: n-step SGD: Apply n gradient descent updates sequentially on the same minibatch with bare learning rate α , effective learning rate (cid:15) =
nα and batch size B .
Sample the next minibatch and repeat.
To analyze n-step SGD, we consider the combined inﬂuence of n updates on the same minibatch:
[Equation]
[Equation]
Equations 23 and 26 are identical up to ﬁrst order in (cid:15) but they differ at O ( (cid:15) 2 ) 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 (cid:15) → 0 , but their dynamics will differ when the learning rate is ﬁnite.
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, as we argued in Section 2, the generalization beneﬁt of SGD arises from ﬁnite learning rate corrections at O ( (cid:15) 2 ) and above, then we 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 ﬁnd that after one epoch, the expected n-step SGD iterate E ( ω m )
=
ω ( m(cid:15) )
+ O ( m 3 (cid:15) 3 ) , where ω (0) =
ω 0 , ˙ ω = −∇ (cid:101) C nSGD ( ω ) and (cid:101) C nSGD ( ω ) =
C ( ω ) +
( (cid:15)/ 4 mn ) (cid:80) m − 1 i =
||∇ C i ( ω ) || 2 .
The scale of the implicit regularizer is proportional to α = (cid:15)/n , which implies that the implicit regularization is suppressed as n increases if we hold (cid:15) constant.
As expected, we recover Equation 1 when n = 1 .
[Figure]
In Figure 4(a), we plot the performance of n-step SGD at a range of bare learning rates α , 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 schedule, whereby we hold the learning rate constant for the ﬁrst half of training, before decaying the learning rate by a factor of 2 every remaining tenth of training, and we provide the mean test accuracy of the best 5 out of 7 training runs for each value of α .
The optimal test accuracy drops from 93 .
5% when n = 1 (standard SGD) to 88 .
8% when n = 16 .
This occurs even though all values of n perform the same number of training epochs, indicating that 16-step SGD performed 16 times more gradient updates.
These results suggest that, at least for this model and dataset, the generalization beneﬁt of SGD is not controlled by the temperature of the associated SDE, but instead arises from the implicit regularization associated with ﬁnite learning rates.
When we increase n we reduce the largest stable bare learning rate α , and this suppresses the implicit regularization beneﬁt, which reduces the test accuracy.
We also verify in Figure 4(b) that similar conclusions arise if we hold the number of parameter updates ﬁxed (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 we conﬁrm in Figure 4(c) that the test accuracy degrades as n 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 n increases.
Note that we provide additional experiments on Fashion-MNIST (Xiao et al., 2017) in appendix D.
D ISCUSSION
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˛ebski et al., 2018;
Chaudhari & Soatto, 2018).
We take a different approach, and note that when the learning rate is ﬁnite, 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 ﬁnite 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 ﬁnite).
It 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 modiﬁed loss, which is deﬁned as a Taylor series in powers of the learning rate, does not converge.
Lewkowycz et al.
(2020) recently argued that the test accuracies of wide networks trained with full batch gradient descent on 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 beneﬁt of small batch SGD, or if these chaotic or oscillatory effects also play a role in some networks. A CKNOWLEDGMENTS
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
A T HE EXPECTED NORM OF A MINIBATCH GRADIENT
To keep the notation clean, we deﬁne X i = ( ∇ C i ( ω ) − ∇ C ( ω )) .
We also recall for clarity that the expectation value E ( ... ) is taken over all possible random shufﬂes of the indices i .
Therefore,
E
[Equation]
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 (cid:80) N i =1 X i =
(cid:80) Ni =1 ( ∇ C i ( ω ) − ∇ C ( ω ))
= 0 .
Therefore,
[Equation] [Equation]
where Γ( ω ) = (1 /N )
(cid:80) Ni =
1 X i · X i =
(1 /N )
(cid:80) Ni =
||∇ C
i ( ω ) − ∇ C ( ω ) || 2 .
We can immediately identify Γ( ω ) as the trace of the empirical covariance matrix of the per-example gradients.
B A B ACKWARD E RROR A NALYSIS FOR N - STEP SGD
Under n-step SGD, we apply n gradient descent updates on the same minibatch with bare learning rate α and batch size B .
After n updates, we sample the next minibatch and repeat.
For convenience, we deﬁne the effective learning rate (cid:15) = nα .
After one minibatch ( n parameter updates),
[Equation]
[Equation] [Equation]
After one epoch (including terms up to second order in (cid:15) ),
[Equation]
−
[Equation]
To simplify this expression, we note that ω i +
= ω
i − (cid:15) ∇ ˆ C i ( ω i )
+ O ( (cid:15) 2 ) .
We can therefore re-use our earlier analysis from Section 2.2 of the main text (see Equation 13 for comparison) to obtain,
[Equation]
[Figure]
Taking the expectation over all possible batch orderings (see Equations 15 to 18), we obtain,
[Equation]
Fixing f ( ω ) =
−∇ C ( ω ) and equating Equation 38 with the continuous modiﬁed ﬂow in Equation 19 by setting E ( ω m )
= ω ( m(cid:15) ) , we identify the modiﬁed ﬂow ˙ ω = −∇ (cid:101) C nSGD
+ O ( m 2 (cid:15) 2 ) , where,
[Equation]
Comparing Equation 39 to Equation 1, we note that the modiﬁed losses of SGD and n-step SGD coincide 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
α = (cid:15)/n , not the effective learning rate (cid:15) .
C T RAINING LOSSES
In Figure 1(b) of Section 2.3 in the main text, we compared the test accuracies achieved when training on the original loss C ( ω ) at a range of learning rates (cid:15) , to the test accuracies achieved when training on the modiﬁed loss C mod ( ω ) at ﬁxed learning rate (cid:15) = 2 − 9 and a range of regularization coefﬁcients λ .
For completeness, in Figure 5, we provide the corresponding training accuracies, as well as the ﬁnal values of the original loss C ( ω ) .
Remarkably, large learning rates and large regularization coefﬁcients achieve similar training accuracies and similar original losses.
This suggests that the implicit regularization term in the modiﬁed loss of SGD ( (cid:101)
C
SGD ( ω ) ) may help explain why the training accuracies and losses often exhibit plateaus when training with large learning rates.
D A DDITIONAL RESULTS ON F ASHION -MNIST
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 in the test set.
We consider a simple fully connected MLP which comprises 3 nonlinear layers, each with width 4096 and
ReLU activations, and a ﬁnal linear softmax layer.
We apply a simple data pipeline which ﬁrst applies per-image standardization and then ﬂattens 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 experiments.
We perform seven training runs for each combination of hyper-parameters and show the mean performance of the best ﬁve (to ensure our results are not skewed by a single failed run). We use a batch size B = 16 unless otherwise speciﬁed, and we do not use weight decay.
We note that this model is highly over-parameterized.
Unlike the Wide-ResNet we considered in the main text we consistently achieve 100% training accuracy if the learning rate is not too large.
[Figure] [Figure]
In Figure 6(a), we train for 400 epochs, and we compare the effect of tuning the learning rate (cid:15) when training on the original loss, to the effect of tuning the explicit regularization strength λ (with (cid:15)
= 2 − 9 ).
As observed in the main text, tuning the explicit regularizer has a similar effect on the test accuracy to tuning the learning rate.
Surprisingly, the optimal values of (cid:15) and λ differ by a factor of 8 .
However we note that the optimal learning rate is (cid:15)
= 2 − 5 , while the explicit regularizer already achieves a very similar test accuracy at λ = 2 − 4 (just a factor of two larger), before reaching a higher maximum test accuracy at λ = 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 in the main text, the test accuracy is determined by the ratio of the learning rate to the batch size.
Meanwhile in Figure 6(c), we plot the test accuracy achieved after training for 1.5 million steps on the modiﬁed loss with learning rate (cid:15) =
2 − 9 and regularization coefﬁcient λ .
Once again, we ﬁnd that the test accuracy achieved is primarily determined by the ratio of the regularization coefﬁcient to the batch size, although smaller batch sizes also achieve slightly higher accuracies.
Finally, in Figure 7 we train using n-step SGD (see Section 4) on the original loss at a range of bare learning rates α .
In Figure 7(a) we train for 400 epochs, while in Figure 7(b) we train for 6 million updates.
We recall that the SDE analogy predicts that the generalization beneﬁt of n-SGD would be determined by the effective learning rate (cid:15) =
nα .
By contrast, backward error analysis predicts that the generalization beneﬁt for small learning rates would be controlled by the bare learning rate α , but that higher order terms may be larger for larger values of n .
We ﬁnd that the test accuracy in both ﬁgures is governed by the bare learning rate α , not the effective learning rate (cid:15) = nα , and therefore these results are inconsistent with the predictions from the SDE analysis in prior work.
Note that Figure 7 has a surprising implication.
It suggests that, for this model, while there is a largest stable bare learning rate we cannot exceed, we can repeatedly apply updates obtained on the same batch of training examples without suffering a signiﬁcant degradation in test accuracy.
We speculate that this may indicate that the gradients of different examples in this over-parameterized model are close to orthogonal (Sankararaman et al., 2020).