On the Origin of Implicit Regularization in Stochastic Gradient Descent

1 INTRODUCTION . In the limit of vanishing learning rates , stochastic gradient descent with minibatch gradients ( SGD ) follows the path of gradient flow 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 benefit 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 ( Jastrzębski 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 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 et al. , 2006 ) . 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 is not too large , the discrete iterates of GD lie close to the path of gradient flow on a modified loss C̃GD ( ω ) = C ( ω ) + ( /4 ) ||∇C ( ω ) ||2 . This modified loss is composed of the original loss C ( ω ) 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 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 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 on a second modified loss C̃SGD ( ω ) , which we define in equation 1 . 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 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 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 benefit of SGD using an analogy between SGD and stochastic differential equations ( SDEs ) ( Mandt et al. , 2017 ; Smith & Le , 2018 ; Jastrzębski 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 benefits of finite learning rates which we identify in this work . An important exception is Li et al . ( 2017 ) , 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 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 define the cost function over parameters ω as C ( ω ) = ( 1/N ) ∑N j=1 Cj ( ω ) , which is the mean of the per-example costs Cj ( ω ) , where N denotes the training set size . Gradient flow follows the ODE ω̇ = −∇C ( ω ) , while gradient descent computes discrete updates ωi+1 = ωi − ∇C ( ωi ) , where 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 define 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 ) ∑m−1 k=0 Ĉk ( ω ) , where the minibatch cost Ĉk ( ω ) = ( 1/B ) ∑kB+B j=kB+1 Cj ( ω ) . In order to guarantee that we sample each example precisely once per epoch , we define SGD by the discrete update ωi+1 = ωi− ∇Ĉi % m ( ωi ) . Informally , our main result is as follows . After one epoch , the mean iterate of SGD with a small but finite learning rate , averaged over all possible shuffles of the batch indices , stays close to the path of gradient flow on a modified loss ω̇ = −∇C̃SGD ( ω ) , where the modified loss C̃SGD is given by : C̃SGD ( ω ) = C ( ω ) + 4m ∑m−1 k=0 ||∇Ĉk ( ω ) ||2 . ( 1 ) We emphasize that our analysis studies the mean evolution of SGD , not the path of individual trajectories . The modified loss C̃SGD ( ω ) is composed of the original loss C ( ω ) and an implicit regularizer Creg ( ω ) = ( 1/4m ) ∑m−1 k=0 ||∇Ĉk ( ω ) ||2 . The scale of this implicit regularization term is proportional to the learning rate , 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 , C̃SGD ( ω ) = C ( ω ) + 4 ||∇C ( ω ) ||2 + 4m ∑m−1 i=0 ||∇Ĉi ( ω ) −∇C ( ω ) ||2 . ( 2 ) We arrive at Equation 2 from Equation 1 by noting that ∑m−1 i=0 ( ∇Ĉi ( ω ) − ∇C ( ω ) ) = 0 . In the limit B → N , we identify the modified loss of gradient descent , C̃GD = C ( ω ) + ( /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 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 ( ||∇Ĉ ( ω ) − ∇C ( ω ) ||2 ) are large ( Wu et al. , 2018 ) . Although global minima of C ( ω ) are global minima of C̃GD ( ω ) , global minima of C ( ω ) may not be global ( or even local ) minima of C̃SGD ( ω ) . Note however that C ( ω ) and 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 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 1 satisfies 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 benefit of small batch SGD arises from the temperature of an associated SDE ( Mandt et al. , 2017 ; Park et al. , 2019 ) . 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 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 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 ω̇ = f ( ω ) . This system usually can not be solved analytically , forcing us to simulate the continuous flow with discrete updates , like the Euler step ω ( t+ ) ≈ ω ( t ) + f ( ω ( t ) ) . However discrete updates will introduce approximation error when the step size is finite . In order to study the bias introduced by this approximation error , we assume the learning rate is relatively small , and introduce a modified flow ω̇ = f̃ ( ω ) , where , f̃ ( ω ) = f ( ω ) + f1 ( ω ) + 2f2 ( ω ) + ... . ( 3 ) The modified flow of f̃ ( ω ) is equal to the original flow of f ( ω ) when → 0 , but it differs from the original flow if is finite . The goal of backward error analysis is to choose the correction terms fi ( ω ) 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 of the solution to the modified flow ω ( t + ) . We obtain the derivatives of ω ( t + ) recursively using the modified flow equation ω̇ = f̃ ( ω ) ( see Hairer et al . ( 2006 ) ) , and we identify the correction terms fi ( ω ) by ensuring this Taylor expansion matches the discrete update ( e.g. , ωt+1 = ωt + f ( ωt ) ) for all powers of . 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 a single step of GD or a single epoch of SGD . Imagine taking n Euler steps on the modified flow f̃ ( ω ) with step size α , ωt+n = ωt + αf̃ ( ωt ) + αf̃ ( ωt+1 ) + αf̃ ( ωt+2 ) + ... ( 4 ) = ωt + αf̃ ( ωt ) + αf̃ ( ωt + αf̃ ( ωt ) ) + αf̃ ( ωt + αf̃ ( ωt ) + αf̃ ( ωt + αf̃ ( ωt ) ) ) + ... ( 5 ) = ωt + nαf̃ ( ωt ) + ( n/2 ) ( n− 1 ) α2∇f̃ ( ωt ) f̃ ( ωt ) +O ( n3α3 ) . ( 6 ) We arrived at Equation 6 by taking the Taylor expansion of f̃ and then counting the number of terms of type ∇f̃ ( ωt ) f̃ ( ωt ) using the formula for an arithmetic series . Note that we assume ∇f̃ exists . Next , to ensure ωt+n in Equation 6 coincides with the solution ω ( t+ ) of the continuous modified flow ω̇ = f̃ ( ω ) for small but finite , we let the number of steps n→∞ while setting α = /n , ω ( t+ ) = ω ( t ) + f̃ ( ω ( t ) ) + ( 2/2 ) ∇f̃ ( ω ( t ) ) f̃ ( ω ( t ) ) +O ( 3 ) ( 7 ) = ω ( t ) + f ( ω ( t ) ) + 2 ( f1 ( ω ( t ) ) + ( 1/2 ) ∇f ( ω ( t ) ) f ( ω ( t ) ) ) +O ( 3 ) . ( 8 ) We have replaced f̃ ( ω ) with its definition 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 modified flow f̃ ( ω ) with vanishing learning rates over a discrete time step of length . Notice that we have assumed that the Taylor expansion in Equation 8 converges , while the higher order terms at O ( 3 ) will contain higher order derivatives of the original flow 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 first order correction f1 ( ω ) of the modified flow for GD . First , we recall that the discrete updates obey ωi+1 = ωi− ∇C ( ωi ) , and we therefore fix f ( ω ) = −∇C ( ω ) . In order to ensure that the continuous modified flow coincides with this discrete update , we need all terms at O ( 2 ) and above in Equation 8 to vanish . At order 2 , this implies that f1 ( ω ) + ( 1/2 ) ∇∇C ( ω ) ∇C ( ω ) = 0 , which yields the first order correction , f1 ( ω ) = − ( 1/2 ) ∇∇C ( ω ) ∇C ( ω ) = − ( 1/4 ) ∇ ( ||∇C ( ω ) ||2 ) . ( 9 ) We conclude that , if the learning rate is sufficiently 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 , ω̇ = −∇C ( ω ) − ( /4 ) ∇ ( ||∇C ( ω ) ||2 ) ( 10 ) = −∇C̃GD ( ω ) . ( 11 ) Equation 11 corresponds to gradient flow on the modified loss , C̃GD ( ω ) = C ( ω ) + ( /4 ) ||∇C ( ω ) ||2 .

Summary: Using 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.