Implicit Gradient Regularization

1 INTRODUCTION . The loss surface of a deep neural network is a mountainous terrain - highly non-convex with a multitude of peaks , plateaus and valleys ( Li et al. , 2018 ; Liu et al. , 2020 ) . Gradient descent provides a path through this landscape , taking discrete steps in the direction of steepest descent toward a sub-manifold of minima . However , this simple strategy can be just as hazardous as it sounds . For small learning rates , our model is likely to get stuck at the local minima closest to the starting point , which is unlikely to be the most desirable destination . For large learning rates , we run the risk of ricocheting between peaks and diverging . However , for moderate learning rates , gradient descent seems to move away from the closest local minima and move toward flatter regions where test data errors are often smaller ( Keskar et al. , 2017 ; Lewkowycz et al. , 2020 ; Li et al. , 2019 ) . This phenomenon becomes stronger for larger networks , which also tend to have a smaller test error ( Arora et al. , 2019a ; Belkin et al. , 2019 ; Geiger et al. , 2020 ; Liang & Rakhlin , 2018 ; Soudry et al. , 2018 ) . In addition , models with low test errors are more robust to parameter perturbations ( Morcos et al. , 2018 ) . Overall , these observations contribute to an emerging view that there is some form of implicit regularization in gradient descent and several sources of implicit regularization have been identified . We have found a surprising form of implicit regularization hidden within the discrete numerical flow of gradient descent . Gradient descent iterates in discrete steps along the gradient of the loss , so after each step it actually steps off the exact continuous path that minimizes the loss at each point . Instead of following a trajectory down the steepest local gradient , gradient descent follows a shallower path . We show that this trajectory is closer to an exact path along a modified loss surface , which can be calculated using backward error analysis from numerical integration theory ( Hairer et al. , 2006 ) . Our core idea is that the discrepancy between the original loss surface and this modified loss surface is a form of implicit regularization ( Theorem 3.1 , Section 3 ) . We begin by calculating the discrepancy between the modified loss and the original loss using backward error analysis and find that it is proportional to the second moment of the loss gradients , which we call Implicit Gradient Regularization ( IGR ) . Using differential geometry , we show that IGR is also proportional to the square of the loss surface slope , indicating that it encourages optimization paths with shallower slopes and optima discovery in flatter regions of the loss surface . Next , we ∗equal contribution explore the properties of this regularization in deep neural networks such as MLP ’ s trained to classify MNIST digits and ResNets trained to classify CIFAR-10 images and in a tractable two-parameter model . In these cases , we verify that IGR effectively encourages models toward minima in the vicinity of small gradient values , in flatter regions with shallower slopes , and that these minima have low test error , consistent with previous observations . We find that IGR can account for the observation that learning rate size is correlated with test accuracy and model robustness . Finally , we demonstrate that IGR can be used as an explicit regularizer , allowing us to directly strengthen this regularization beyond the maximum possible implicit gradient regularization strength . 2 THE MODIFIED LOSS LANDSCAPE INDUCED BY GRADIENT DESCENT . The general goal of gradient descent is to find a weight vector θ̂ in parameter space Rm that minimizes a loss E ( θ ) . Gradient descent proceeds by iteratively updating the model weights with learning rate h in the direction of the steepest loss gradient : θn+1 = θn − h∇θE ( θn ) ( 1 ) Now , even though gradient descent takes steps in the direction of the steepest loss gradient , it does not stay on the exact continuous path of the steepest loss gradient , because each iteration steps off the exact continuous path . Instead , we show that gradient descent follows a path that is closer to the exact continuous path given by θ̇ = −∇θẼ ( θ ) , along a modified loss Ẽ ( θ ) , which can be calculated analytically using backward error analysis ( see Theorem 3.1 and Section 3 ) , yielding : Ẽ ( θ ) = E ( θ ) + λRIG ( θ ) , ( 2 ) where λ ≡ hm 4 ( 3 ) and RIG ( θ ) ≡ 1 m m∑ i=1 ( ∇θiE ( θ ) ) 2 ( 4 ) Immediately , we see that this modified loss is composed of the original training loss E ( θ ) and an additional term , which we interpret as a regularizer RIG ( θ ) with regularization rate λ . We call RIG ( θ ) the implicit gradient regularizer because it penalizes regions of the loss landscape that have large gradient values , and because it is implicit in gradient descent , rather than being explicitly added to our loss . Definition . Implicit gradient regularization is the implicit regularisation behaviour originating from the use of discrete update steps in gradient descent , as characterized by Equation 2 . We can now make several predictions about IGR which we will explore in experiments : Prediction 2.1 . IGR encourages smaller values of RIG ( θ ) relative to the loss E ( θ ) . Given Equation 2 and Theorem 3.1 , we expect gradient descent to follow trajectories that have relatively small values of RIG ( θ ) . It is already well known that gradient descent converges by reducing the loss gradient so it is important to note that this prediction describes the relative size of RIG ( θ ) along the trajectory of gradient descent . To expose this phenomena in experiments , great care must be taken when comparing different gradient descent trajectories . For instance , in our deep learning experiments , we compare models at the iteration time of maximum test accuracy ( and we consider other controls in the appendix ) , which is an important time point for practical applications and is not trivially determined by the speed of learning ( Figures 1 , 2 ) . Also , related to this , since the regularization rate λ is proportional to the learning rate h and network size m ( Equation 3 ) , we expect that larger models and larger learning rates will encourage smaller values of RIG ( θ ) ( Figure 2 ) . Prediction 2.2 . IGR encourages the discovery of flatter optima . In section 3 we will show that RIG ( θ ) is proportional to the square of the loss surface slope . Given this and Prediction 2.1 , we expect that IGR will guide gradient descent along paths with shallower loss surface slopes , thereby encouraging the discovery of flatter , broader optima . Of course , it is possible to construct loss surfaces at odds with this ( such as a Mexican-hat loss surface , where all minima are equally flat ) . However , we will provide experimental support for this using loss surfaces that are of widespread interest in deep learning , such as MLPs trained on MNIST ( Figure 1 , 2 , 3 ) . Prediction 2.3 . IGR encourages higher test accuracy . Given Prediction 2.2 , we predict that IGR encourages higher test accuracy since flatter minima are known empirically to coincide with higher test accuracy ( Figure 2 ) . Prediction 2.4 . IGR encourages the discovery of optima that are more robust to parameter perturbations . There are several important observations to make about the properties of IGR : 1 ) It does not originate in any specific model architecture or initialization , although our analysis does provide a formula to explain the influence of these model properties through IGR ; 2 ) Other sources of implicit regularization also have an impact on learning , alongside IGR , and the relative importance of these contributions will likely depend on model architecture and initialization ; 3 ) In defining λ and RIG we chose to set λ proportional to the number of parameters m. To support this choice , we demonstrate in experiments that the test accuracy is controlled by the IGR rate λ . 4 ) The modified loss and the original loss share the same global minima , so IGR vanishes when the gradient vanishes . Despite this , the presence of IGR has an impact on learning since it changes the trajectory of gradient descent , and in over-parameterized models this can cause the final parameters to reach different solutions . 5 ) Our theoretical results are derived for full-batch gradient descent , which allows us to isolate the source of implicit regularisation from the stochasticity of stochastic gradient descent ( SGD ) . Extending our theoretical results to SGD is considerably more complicated , and as such , is beyond the scope of this paper . However , in some of our experiments , we will demonstrate that IGR persists in SGD , which is especially important for deep learning . Next , we will provide a proof for Theorem 3.1 , and we will provide experimental support for our predictions . 3 BACKWARD ERROR ANALYSIS OF GRADIENT DESCENT . In this section , we show that gradient descent follows the gradient flow of the modified loss Ẽ ( Equation 2 ) more closely than that of the original loss E. The argument is a standard argument from the backward error analysis of Runge-Kutta methods ( Hairer et al. , 2006 ) . We begin by observing that gradient descent ( Equation 1 ) can be interpreted as a Runge-Kutta method numerically integrating the following ODE : θ̇ = −∇θE ( θ ) ( 5 ) In the language of numerical analysis , gradient descent is the explicit Euler method numerically integrating the vector field f ( θ ) = −∇E ( θ ) . The explicit Euler method is of order 1 , which means that after one gradient descent step θn = θn−1 − h∇E ( θn−1 ) , the deviation from the gradient flow ‖θn−θ ( h ) ‖ is of orderO ( h2 ) , where θ ( h ) is the solution of Equation 5 starting at θn−1 and evaluated at time h. Backward error analysis was developed to deal with this discrepancy between the discrete steps of a Runge-Kutta method and the continuous exact solutions ( or flow ) of a differential equation . The main idea is to modify the ODE vector field θ̇ = f ( θ ) with corrections in powers of the step size f̃ ( θ ) = f ( θ ) + hf1 ( θ ) + h 2f2 ( θ ) + · · · ( 6 ) so that the numerical steps θn approximating the original Equation 5 now lie exactly on the solutions of the modified equation θ̇ = f̃ ( θ ) . In other words , backward error analysis finds the corrections fi in Equation 6 such that θn = θ̃ ( nh ) for all n , where θ̃ ( t ) is the solution of the modified equation starting at θ0 . In theory , we can now precisely study the flow of the modified equation to infer properties of the numerical method because its steps follow the modified differential equation solutions perfectly in a formal sense . The following result is a direct application of backward error analysis to gradient descent : Theorem 3.1 . Let E be a sufficiently differentiable function on a parameter space θ ∈ Rm . The modified equation for gradient flow ( Equation 5 ) is of the form θ̇ = −∇Ẽ ( θ ) +O ( h2 ) ( 7 ) where Ẽ = E + λRIG is the modified loss introduced in Equation 2 . Consider gradient flow with the modified loss θ̇ = −∇Ẽ ( θ ) and its solution θ̃ ( t ) starting at θn−1 . Now the local error ‖θn − θ̃ ( h ) ‖ between θ̃ ( h ) and one step of gradient descent θn = θn−1 − h∇E ( θn−1 ) is of order O ( h3 ) , while it is of order O ( h2 ) for gradient flow with the original loss . Proof . We begin by computing f1 for which the first two orders in h of the Taylor ’ s Series of the modified equation solution θ ( t ) at t = h coincide with one gradient descent step . Since θ′ ( t ) = f̃ ( θ ) , we see that θ′′ ( t ) = f̃ ′ ( θ ) f̃ ( θ ) and we find that θ + hf ( θ ) = θ + hf ( θ ) + h2 ( f1 ( θ ) + 1 2 f ′ ( θ ) f ( θ ) ) , yielding f1 ( θ ) = −f ′ ( θ ) f ( θ ) /2 . Now , when f is a gradient vector field with f = −∇E , we find : f1 ( θ ) = − 1 2 ( D2θE ) ∇θE ( θ ) = − 1 4 ∇‖∇E ( θ ) ‖2 , where D2θE is the Hessian of E ( θ ) . Putting this together , we obtain the first order modified equation : θ̇ = f + hf1 +O ( h2 ) = −∇ ( E ( θ ) + h 4 ‖∇E ( θ ) ‖2 ) +O ( h2 ) , which is a gradient system with modified loss Ẽ ( θ ) = E ( θ ) + h 4 ‖∇E ( θ ) ‖2 . As for the local error , if θ ( h ) is a solution of gradient flow starting at θn−1 , we have in general that θ ( h ) = θn + O ( h2 ) . The correction f1 is constructed so that it cancels out the O ( h2 ) term in the expansion of its solution , yielding θ̃ ( h ) = θn +O ( h3 ) . Remark 3.2 . A direct application of a standard result in backward error analysis ( Hairer & Lubich ( 1997 ) , Thm . 1 ) indicates that the learning rate range where the gradient flow of the modified loss provides a good approximation of gradient descent lies below h0 = CR/M , where∇E is analytic and bounded by M in a ball of radius R around the initialization point and where C depends on the Runge-Kutta method only , which can be estimated for gradient descent . We call this the moderate learning rate regime . For each learning rate below h0 , we can provably find an optimal truncation of the modified equation whose gradient flow is exponentially close to the steps of gradient descent , so the higher term corrections are likely to contribute to the dynamics . Given this , we see that the exact value of the upper bound for the moderate regime will correspond to a setting where the optimal truncation is the first order correction only . Calculating this in general is difficult and beyond the scope of this paper . Nonetheless , our experiments strongly suggest that this moderate learning rate regime overlaps substantially with the learning rate range typically used in deep learning . Next , we give a purely geometric interpretation of IGR , supporting Prediction 2.2 . Consider the loss surface S associated with a loss function E defined over the parameter space θ ∈ Rm . This loss surface is defined as the graph of the loss : S = { ( θ , E ( θ ) ) : θ ∈ Rm } ⊂ Rm+1 . We define α ( θ ) to be the angle between the tangent space TθS to S at θ and the parameter plane , i.e. , the linear subspace { ( θ , 0 ) : θ ∈ Rm } in Rm+1 . We can compute this angle using the inner product between the normal vectorN ( θ ) to S at θ and the normal vector ẑ to the parameter plane : α ( θ ) = arccos 〈N ( θ ) , ẑ〉 . Now we can define the loss surface slope at θ as being the tangent of this angle : slope ( θ ) : = tanα ( θ ) . This is a natural extension of the 1-dimensional notion of slope . With this definition , we can now reformulate the modified loss function in a purely geometric fashion : Proposition 3.3 . The modified loss Ẽ in Equation 2 can be expressed in terms of the loss surface slope as Ẽ ( θ ) = E ( θ ) + h4 slope 2 ( θ ) . This proposition is an immediate consequence of Theorem 3.1 and Corollary A.6.1 in Appendix A.2 . It tells us that gradient descent with higher amounts of implicit regularization ( higher learning rate ) will implicitly minimize the loss surface slope locally along with the original training loss . Prediction 2.2 claims that this local effect of implicit slope regularization accumulates into the global effect of directing gradient descent trajectories toward global minima in regions surrounded by shallower slopes - toward flatter ( or broader ) minima . Remark 3.4 . It is important to note that IGR does not help gradient descent to escape from local minima . In the learning rate regime where the truncated modified equation gives a good approximation for gradient descent , the steps of gradient descent follow the gradient flow of the modified loss closely . As Proposition A.10 shows , the local minima of the original loss are still local minima of the modified loss so gradient descent within this learning rate regime remains trapped within the basin of attraction of these minima . IGR does not lead to an escape from local minima , but instead , encourages a shallower path toward flatter solutions close to the submanifold of global interpolating minima , which the modified loss shares with the original loss ( Proposition A.10 ) .

Summary: The authors show the discrete steps of gradient descent implicitly regularize models by penalizing trajectories that have large loss-gradients, which is called Implicit Gradient Regularization in the paper. The authors adopt a standard argument from the backward error analysis of Runge-Kutta methods to show this phenomenon. In the paper, the authors also provide some empirical results which indicate gradient descent leads to flat minima where test errors are small and solutions are robust to noisy parameter perturbations.