I MPLICIT G RADIENT R EGULARIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Gradient descent can be surprisingly good at optimizing deep neural networks without overﬁtting and without explicit regularization.
We ﬁnd that the discrete steps of gradient descent implicitly regularize models by penalizing gradient descent trajectories that have large loss gradients.
We call this Implicit Gradient Regularization (IGR) and we use backward error analysis to calculate the size of this regularization.
We conﬁrm empirically that implicit gradient regularization biases gradient descent toward ﬂat minima, where test errors are small and solutions are robust to noisy parameter perturbations.
Furthermore, we demonstrate that the implicit gradient regularization term can be used as an explicit regularizer, allowing us to control this gradient regularization directly.
More broadly, our work indicates that backward error analysis is a useful theoretical approach to the perennial question of how learning rate, model size, and parameter regularization interact to determine the properties of overparameterized models optimized with gradient descent.
I NTRODUCTION
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.
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 ﬂatter 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 identiﬁed (see related work, Section 7).
We have found a surprising form of implicit regularization hidden within the discrete numerical ﬂow 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 modiﬁed 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 modiﬁed loss surface is a form of implicit regularization (Theorem 3.1, Section 3).
We begin by calculating the discrepancy between the modiﬁed loss and the original loss using backward error analysis and ﬁnd 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 ﬂatter regions of the loss surface.
Next, we 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 ﬂatter regions with shallower slopes, and that these minima have low test error, consistent with previous observations.
We ﬁnd 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.
T HE MODIFIED LOSS LANDSCAPE INDUCED BY GRADIENT DESCENT
The general goal of gradient descent is to ﬁnd a weight vector ˆ θ in parameter space R m 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:
[Equation]
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 ˙
θ = −∇ θ (cid:101) E
( θ ) , along a modiﬁed loss (cid:101)
E ( θ ) , which can be calculated analytically using backward error analysis (see Theorem 3.1 and Section 3), yielding:
[Equation]
where
[Equation]
and
[Equation]
Immediately, we see that this modiﬁed loss is composed of the original training loss E ( θ ) and an additional term, which we interpret as a regularizer R IG ( θ ) with regularization rate λ .
We call
R IG ( θ ) 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.
Deﬁnition 2.1.
Implicit gradient regularization (IGR) in gradient descent is the impact of taking discrete update steps along a loss gradient in comparison to following the exact continuous ﬂow along the same loss gradient.
IGR is characterized in Equation 2 and can be strengthened using larger learning rate h or larger network size m (Equation 3).
We can now make several predictions about IGR which we will explore in experiments:
Prediction 2.2.
IGR encourages smaller values of R IG ( θ ) 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 R IG ( θ ) .
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
R IG ( θ ) 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 R IG ( θ ) (Figure 2).
Prediction 2.3.
IGR encourages the discovery of ﬂatter optima.
In section 3 we will show that R IG ( θ ) is proportional to the square of the loss surface slope.
Given this and Prediction 2.2, we expect that IGR will guide gradient descent along paths with shallower loss surface slopes, thereby encouraging the discovery of ﬂatter, 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 ﬂat).
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.4.
IGR encourages higher test accuracy.
Given Prediction 2.3, we predict that IGR encourages higher test accuracy since ﬂatter minima are known empirically to coincide with higher test accuracy (Figure 2).
Prediction 2.5.
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 speciﬁc model architecture or initialization, although our analysis does provide a formula to explain the inﬂuence 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 deﬁning λ and R IG 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 modiﬁed 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 ﬁnal parameters to reach different solutions.
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.
B ACKWARD ERROR ANALYSIS OF GRADIENT DESCENT
In this section, we show that gradient descent follows the gradient ﬂow of the modiﬁed loss (cid:101) 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:
[Equation]
In the language of numerical analysis, gradient descent is the explicit Euler method numerically integrating the vector ﬁeld 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 ﬂow (cid:107) θ n − θ ( h ) (cid:107) is of order
O ( h 2 ) , 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 ﬂow) of a differential equation.
The main idea is to modify the ODE vector ﬁeld ˙ θ = f ( θ ) with corrections in powers of the step size
[Equation]
so that the numerical steps θ n approximating the original Equation 5 now lie exactly on the solutions of the modiﬁed equation ˙ θ = ˜
f ( θ ) .
In other words, backward error analysis ﬁnds the corrections f
i in
Equation 6 such that θ n = ˜
θ ( nh ) for all n , where ˜ θ ( t ) is the solution of the modiﬁed equation starting at θ 0 .
In theory, we can now precisely study the ﬂow of the modiﬁed equation to infer properties of the numerical method because its steps follow the modiﬁed 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 sufﬁciently differentiable function on a parameter space θ ∈ R m .
The modiﬁed equation for gradient ﬂow (Equation 5) is of the form
[Equation]
where (cid:101) E
= E + λR IG is the modiﬁed loss introduced in Equation 2.
Consider gradient ﬂow with the modiﬁed loss ˙ θ = −∇ (cid:101) E ( θ ) and its solution ˜ θ ( t ) starting at θ n − 1 .
Now the local error (cid:107) θ n − ˜ θ
( h ) (cid:107) between ˜ θ ( h ) and one step of gradient descent θ n = θ n − 1 − h ∇ E ( θ n − 1 ) is of order
O ( h 3 ) , while it is of order O ( h 2 ) for gradient ﬂow with the original loss.
Proof.
The idea is to compute f 1 so that the ﬁrst two orders in h of the Taylor’s Series of the modiﬁed equation solution θ ( t ) at t = h coincide with one gradient descent step.
Since θ (cid:48) ( t ) =
˜ f
( θ ) , we see that θ (cid:48)(cid:48) ( t ) =
˜ f
(cid:48)
( θ ) ˜
f
( θ ) .
We thus want to ﬁnd f 1 so that
[Equation]
which yields f 1 ( θ )
=
− f (cid:48)
( θ ) f ( θ ) / 2 .
Now, when f is a gradient vector ﬁeld −∇ E , we see that:
[Equation]
where D 2 θ E is the Hessian of E ( θ ) .
Putting all together, we obtain the modiﬁed equation at order 1:
[Equation]
which is again a gradient system with the modiﬁed loss
[Equation]
As for the local error, if θ ( h ) is a solution of gradient ﬂow starting at θ n − 1 , we have in general that θ ( h )
= θ n
+ O ( h 2 ) .
The correction f 1 to gradient ﬂow is constructed so that it cancels out the O ( h 2 ) term in the expansion of its solution, yielding ˜ θ ( h )
= θ n
+ O ( h 3 ) .
Remark 3.2.
A direct application of a standard result in backward error analysis (Hairer & Lubich (1997), Thm. 1) indicates that the moderate learning rate range where the gradient ﬂow of the modiﬁed loss provides a good approximates of gradient descent lies below h
= 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.
For each learning rate below h 0 , we can provably ﬁnd an optimal truncation of the modiﬁed equation whose gradient ﬂow 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 ﬁrst order correction only.
Calculating this in general is difﬁcult and beyond the scope of this paper.
Nonetheless, our experiments strongly suggest that the moderate regime overlaps substantially with the learning rate regimes typically used in deep learning, which are neither very large or inﬁnitesimally small.
Next, we give a purely geometric interpretation of IGR, supporting Prediction 2.3.
Consider the loss surface S associated with a loss function E deﬁned over the parameter space θ ∈ R m .
This loss surface is deﬁned as the graph of the loss:
S = { ( θ, E ( θ )) : θ ∈ R m } ⊂ R m +1 .
We deﬁne α ( θ ) to be the angle between the tangent space T θ S to S at θ and the parameter plane, i.e., the linear subspace { ( θ, 0) : θ ∈ R m } in R m +1 .
We can compute this angle using the inner product between the normal vector N ( θ ) to S at θ and the normal vector ˆ z to the parameter plane: α ( θ ) =
arccos (cid:104) N ( θ ) , ˆ z (cid:105) .
Now we can deﬁne 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 deﬁnition, we can now reformulate the modiﬁed loss function in a purely geometric fashion:
Proposition 3.3.
The modiﬁed loss ˜ E in Equation 2 can be expressed in terms of the loss surface slope as ˜ E ( θ ) =
E ( θ ) + h 4 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.3 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 ﬂatter (or broader) minima.
Remark 3.4.
It is important to note that IGR does not help escape local minima.
Namely, in the learning rate regime where the truncated modiﬁed equation gives a good approximation for gradient descent (and where IGR has an impact), the steps of gradient descent follow the gradient ﬂow of the modiﬁed loss closely.
As Proposition A.10 shows, the local minima of the original loss are still local minima of the modiﬁed loss, and gradient descent at this learning rate regime remains trapped within the basin of attraction of these minima.
To be clear, the impact of IGR is not to escape local minima, but to encourage a shallower path toward ﬂatter solutions close to the submanifold of global interpolating minima, which the modiﬁed loss shares with the original loss (Proposition A.10).
E XPLICIT G RADIENT R EGULARIZATION
For overparameterized models, we predict that the strength of IGR relative to the original loss can be controlled by increasing the learning rate h (Prediction 2.2).
However, gradient descent becomes unstable when the learning rate becomes too large.
For applications where we wish to increase the strength of IGR beyond this point, we can take inspiration from implicit gradient regularization to motivate Explicit Gradient Regularization (EGR), which we deﬁne as:
[Equation]
where, µ is the explicit regularization rate, which is a hyper-parameter that we are free to choose, unlike the implicit regularization rate λ (Equation 3) which can only by controlled indirectly.
Now, we can do gradient descent on E µ with small learning rates and large µ .
Although EGR is not the primary focus of our work, we will demonstrate the effectiveness of EGR for a simple two parameter model in the next section (Section 5) and for a ResNet trained on Cifar- (Figure 3c).
Our EGR experiments act as control study in this work, to demonstrate that the R IG term arising implicitly in gradient descent can indeed improve test accuracy independent of any possible confounding effects that may arise when we control IGR implicitly through the learning rate.
Namely, if we could not observe a signiﬁcant boost in model test accuracy by adding the RIG term explicitly, our prediction that implicit regularization helps to boost test accuracy would be in doubt.
Remark 4.1.
Explicit regularization using gradient penalties has a long history dating back to
Drucker & Le Cun (1992) where a penalty w.r.t.
the input (and not the model parameters) is proposed.
Closer to EGR, in Hochreiter & Schmidhuber (1997) a regularization penalty with properties reminiscent of R IG was proposed to guide gradient descent toward ﬂat minima.
EGR is also reminiscent of other regularizers, such as dropout, which similarly encourage robust parameterizations (Morcos et al., 2018; Srivastava et al., 2014; Tartaglione et al., 2018).
The success of these explicit regularizers demonstrates the importance of this type of regularization in deep learning.
More recently, loss gradient penalties with respect to the discriminator or generator parameters have been proposed to help stabilize the training of GANs (Nagarajan & Kolter, 2017;
Balduzzi et al., 2018; Mescheder et al., 2017;
Qin et al., 2020).
IGR AND EGR IN A 2- D LINEAR MODEL
In our ﬁrst experiment we explore implicit and explicit gradient regularization in a simple twoparameter model with a loss given by E ( a, b ) =
( y − f ( x ; a, b )) 2 / 2 , where f ( x ; a, b ) = abx is our model, a, b ∈ R are the model parameters and x, y ∈ R is the training data.
We have chosen this model because we can fully visualize the gradient descent trajectories in 2-d space.
For a single data point, this model is overparameterized with global minima located along a curve attractor deﬁned by the hyperbola ab =
y/x .
For small learning rates, gradient descent follows the gradient ﬂow of the loss from an initial point ( a 0 , b 0 ) toward the line attractor.
For larger learning rates, gradient descent follows a longer path toward a different destination on the line attractor (Figure 1a).
We can understand these observations using Theorem 3.1, which predicts that gradient descent is closer to the modiﬁed ﬂow given by ˙ a = −∇ a (cid:101) E ( a, b ) and ˙ b = −∇ b (cid:101) E ( a, b ) where (cid:101) E ( a, b ) =
E ( a, b ) + λR IG ( a, b ) is the modiﬁed loss from Equation 2, R IG ( a, b )
= (cid:0) | a | 2 + | b | 2 (cid:1) x 2 E ( a, b ) is the implicit regularization term from Equation 4 and λ =
h/ 2 is the implicit regularization rate from Equation 3, with learning rate h .
Although the modiﬁed loss and the original loss have the same global minima, they generate different ﬂows.
Solving the modiﬁed ﬂow equations numerically starting from the same initial point ( a 0 , b 0 ) as before, we ﬁnd that the gradient descent trajectory is closer to the modiﬁed ﬂow than the exact ﬂow, consistent with Theorem 3.1 (Figure 1a).
Next, we investigate Prediction 2.2, that the strength of the implicit gradient regularization R IG ( a, b ) relative to the original loss E ( a, b ) can be controlled by increasing the regularization rate λ .
In this case, this means that larger learning rates should produce gradient descent trajectories that lead to minima with a smaller value of R IG ( a, b ) /E
( a, b ) = x 2 (cid:0) | a | 2 + | b | 2 (cid:1) .
It is interesting to note that this is proportional to the parameter norm, and also, to the square of the loss surface slope.
In our numerical experiments, we ﬁnd that larger learning rates lead to minima with smaller L2 norm (Figure 1b), closer to the ﬂatter region in the parameter plane, consistent with Prediction 2.2 and 2.3.
The extent to which we can strengthen IGR in this way is restricted by the learning rate.
For excessively
[Figure]
large learning rates, gradient descent ricochets from peak to peak, until it either diverges or lands in the direct vicinity of a minimum, which is sensitively dependent on initialization (Figure A.2)
To go beyond the limits of implicit gradient regularization we can explicitly regularize this model using Equation 8 to obtain a regularized loss E µ ( a, b ) =
E ( a, b ) + µ (cid:0) | a | 2 + | b | 2 (cid:1) x 2 E ( a, b ) .
Now, if we numerically integrate ˙ a = −∇ a E µ ( a, b ) and ˙ b = −∇ b E µ ( a, b ) starting from the same initial point ( a 0 , b 0 ) , using a very large explicit regularization rate µ (and using gradient descent with a very small learning rate h for numerical integration, see Appendix A.4)
we ﬁnd that this ﬂow leads to global minima with a small L2 norm (Figure 1a) in the ﬂattest region of the loss surface.
This is not possible with IGR, since it would require learning rates so large that gradient descent would diverge.
IGR AND EGR IN DEEP NEURAL NETWORKS
Next, we empirically investigate implicit gradient regularization and explicit gradient regularization in deep neural networks.
We consider a selection of MLPs trained to classify MNIST digits and we also investigate Resnet-18 trained to classify CIFAR-10 images.
All our models are implemented using Haiku (Hennigan et al., 2020).
To begin, we measure the size of implicit regularization in MLPs trained to classify MNIST digits with a variety of different learning rates and network sizes (Figure 2).
Speciﬁcally, we train 5-layer MLPs with n l units per layer, where n l ∈ { 50 , 100 , 200 , 400 , 800 , 1600 } , h ∈ { 0 . 5 , 0 . 1 , 0 . 05 , 0 . 01 , 0 .
001 , 0 . 0005 } , using ReLu activation functions and a cross entropy loss (see Appendix A.5 for further details and see Figures A.3, A.4 and A.5 for training and test data curves).
We report R IG and test accuracy at the time of maximum test accuracy for each network that ﬁts the training data exactly.
We choose this time point for comparison because it is important for practical applications.
We ﬁnd that R IG is smaller for larger learning rates and larger networks (Figure 2a), consistent with Theorem 3.1 and Prediction 2.2.
Next, we measure the loss surface slope in 5-layer MLPs, with 400 units per layer, trained to classify MNIST digits with a range of different learning rates.
We ﬁnd that neural networks with larger learning rates, and hence, with stronger IGR have smaller slopes at the time of maximum test accuracy (Figure 3a).
We also measure the loss surface slope in the vicinity of these optima.
To do this, we add multiplicative Gaussian noise to every parameter according to θ p
= θ (1 + η ) , where θ are the parameters of a fully trained
[Figure]
model and θ p are the parameters after the addition of noise, where η ∼ N (0 , σ ) .
We ﬁnd that neural networks trained with larger learning rates have ﬂatter slopes and these slopes remain small following larger perturbations (Figure 3a).
These numerical results are consistent with our prediction that IGR encourages the discovery of ﬂatter optima (Prediction 2.3)
Next, we observe that improvements in test set accuracy are correlated with increases in regularization rate (Figure 2b), and also with increases in learning rate and network size (Figure A.6).
This is consistent with Prediction 2.4.
Furthermore, the correlation between test set accuracy and network size m supports our use of network size scaling in Equation 3 and 4.
Next, we explore the robustness of deep neural networks in response to parameter perturbations.
In previous work, it has been reported that deep neural networks are robust to a substantial amount of parameter noise, and that this robustness is stronger in networks with higher test accuracy (Morcos et al., 2018).
We measure the degradation in classiﬁcation accuracy as we increase the amount of multiplicative Gaussian noise and ﬁnd that neural networks with larger learning rates, and hence, with stronger IGR, are more robust to parameter perturbations after training (Figure 3c), consistent with Prediction 2.5.
This may explain the origin, in part, of deep neural network robustness.
We also explore IGR in several other settings.
For ResNet-18 models trained on CIFAR-10, we ﬁnd that R IG is smaller and test accuracy is higher for larger learning rates (at the time of maximum test accuracy)
(Figure A.7, A.8), consistent again with Theorem 3.1 and Predictions 2.2 and 2.4.
We also explore IGR using different stopping time criteria (other than the time of maximum test accuracy), such as ﬁxed iteration time (Figures A.3, A.4), and ﬁxed physical time (Figure A.5) (where iteration time is rescaled by the learning rate, see Appendix A.5 for further information).
We explore IGR for full batch gradient descent and for stochastic gradient descent (SGD) with a variety of different batch sizes (Figure A.6) and in all these cases, our numerical experiments are consistent with Theorem 3.1.
These supplementary experiments are designed to control for the presence, and absence, of other sources of implicit regularisation - such as model architecture choice, SGD stochasticity and the choice of stopping time criteria.
Finally, we provide an initial demonstration of explicit gradient regularization (EGR).
Speciﬁcally, we train a ResNet-18 using our explicit gradient regularizer (Equation 8) and we observe that EGR produces a boost of more than 12% in test accuracy (see Figure 3c).
This initial experiment indicates that EGR may be a useful tool for training of neural networks, in some situations, especially where
IGR cannot be increased with larger learning rates, which happens, for instance, when learning rates are so large that gradient descent diverges.
However, EGR is not the primary focus of our work here, but for IGR, which is our primary focus, this experiment provides further evidence that IGR may play an important role as a regularizer in deep learning.
[Figure]
R ELATED WORK
Implicit regularization: Many different sources of implicit regularization have been identiﬁed, including early-stopping (Hardt et al., 2016), model initialization (Glorot & Bengio, 2010; Li &
Liang, 2018; Nagarajan & Kolter, 2019; Gunasekar et al., 2018; Zhang et al., 2019; Zou et al., 2019), model architecture (Li et al., 2018; Lin & Tegmark, 2016; Ma et al., 2020), stochasticity (Keskar et al., 2017; Soudry et al., 2018; Roberts, 2018; Ali et al., 2020; Chaudhari & Soatto, 2018; De &
Smith, 2020; Mandt et al., 2017; Park et al., 2019; Sagun et al., 2017; Smith & Le, 2018; Wilson et al., 2017; Jastrzebski et al., 2020), implicit L2 regularity (Soudry et al., 2018; Neyshabur et al., 2015; Ali et al., 2019; Ji & Telgarsky, 2019; Nacson et al., 2019; Poggio et al., 2019; Suggala et al., 2018), low rank biases (Arora et al., 2019a; Gunasekar et al., 2017; Razin & Cohen, 2020) among other possibilities.
There are many possible interactions between these types of implicit regularization and IGR.
For instance, initialization in a ﬂat region of the loss surface will reduce the impact of IGR on the training.
Model architecture regularization like the skip connections in Resnet analyzed in Li et al.
(2018), which tend to produce smoother loss surfaces, will probably help IGR.
IGR probably doesn’t interact much with forms of implicit regularization arising for very small learning rate where the identiﬁcation with gradient ﬂow holds as in (Poggio et al., 2017; 2019; Saxe et al., 2014), because in this regime IGR is negligible.
A number of works have identiﬁed implicit regularization in the discrete steps of gradient descent for speciﬁc datasets, losses, or architectures (Soudry et al., 2018; Neyshabur et al., 2015; Gidel et al., 2019).
IGR may be connected to these implicit forms of regularization when the location of the ﬂat regions in our modiﬁed loss surface correspond to the locations where the proposed regularizer is minimized, in the same way that the ﬂat region in our example in Section 5 coincides with the minimum of implicit L2 norm regularization.
IGR might also be useful for understanding implicit regularization in deep matrix factorization with gradient descent (Arora et al., 2019a; Gunasekar et al., 2017; Razin & Cohen, 2020), where gradient descent seems to have a low-rank bias. Our work may also provide a useful perspective on the ’break-even’ point in deep learning (Jastrzebski et al., 2020). At a break-even point our backward analysis suggests that gradient descent with large learning rates will move toward ﬂatter regions, consistent with this work. Stochastic effects are also likely to contribute to the trajectory at break-even points.
Learning rate schedules and regimes:
Implicit gradient regularization can be used to understand the role of learning schedules, since learning rate controls the relative strength of implicit regularization and loss optimization.
For example, for a cyclical learning rate schedule (Smith, 2017), cyclically varying learning rates between large and small learning rates can be interpreted as a cyclical variation between large and small amounts of IGR (i.e., alternate phases of optimization and regularization).
For learning rate schedulers with fast changing learning rates, the dynamics of the learning rate scheduler cannot be analyzed separately from the dynamics of gradient descent.
In such cases, we can again use backward analysis to reveal additional implicit regularization terms that arise from the learning rate scheduler dynamics.
A number of studies have identiﬁed various learning rates regimes characterized by different convergence and generalization properties.
For instance Li et al.
(2019) identiﬁes a small learning rate regime where a network tends to memorize and a large learning rate regime characterized by increased generalization power.
This is consistent with IGR which we believe is most useful at the beginning of training, orienting the search toward ﬂatter regions and less important in later stages of the training, when a ﬂatter region has been reached, and where convergence to any of the ﬂatter minima is more important. This is also consistent with
Jastrzebski et al.
(2020) that showed the importance of large learning rate at the beginning of training in encouraging more favourable optimization trajectories.
Also note that Lewkowycz et al.
(2020) identiﬁes a lazy phase, a catapult phase, and a divergent phase, which may be related to the range of backward error analysis applicability.
Neural Tangent Kernel:
The Neural Tangent Kernel (NTK) is especially interesting (Arora et al.
, 2019b;c; Chizat & Bach, 2019; Jacot et al., 2018; Lee et al., 2019; Oymak & Soltanolkotabi, 2019;
Woodworth et al., 2020; Cao & Gu, 2019) since, in the case of the least square loss, the IGR term R IG can be related to the NTK (see Appendix A.3).
This is particularly interesting because it suggests that the NTK may play a role beyond the kernel regime , into the rich regime .
In this context, IGR is also related to the trace of the Fisher Information Matrix (Karakida et al., 2019).
Runge-Kutta methods: Recently, Runge-Kutta methods have been used to understand old (and devise new) gradient-based optimization methods (Betancourt et al., 2018;
Scieur et al., 2017; Zhang et al., 2018; França et al., 2020).
However, implicit regularization and backward error analysis has not been explored.
A stochastic version of the modiﬁed equation was used (Li et al., 2017; Feng et al., 2020) to study stochastic gradient descent in the context of stochastic differential equations and diffusion equations for the study of convergence and adaptive learning schemes, and very recently (França et al., 2020) used backward analysis to device new optimizers to control their convergence and stability.
To the best of our knowledge, it is the ﬁrst time backward analysis is proposed as a general method to uncover inductive bias and regularization in gradient based optimizers.
In future work, it would be interesting to use this approach to understand implicit regularisation across a number of standard optimizer like momentum, Adam, or RMSProp.
It would also be interesting to explore the properties of higher order modiﬁed loss corrections.
Although this is outside the scope of our work here, we have provided formulae for several higher order terms in the appendix.
D ISCUSSION
Following our backward error analysis, we now understand gradient descent as an algorithm that effectively optimizes a modiﬁed loss with an implicit regularization term arising through the discrete nature of gradient descent.
This leads to several predictions that we conﬁrm experimentally: (i) IGR penalizes the second moment of the loss gradients (Prediction 2.2), and consequently, (ii) it penalizes minima in the vicinity of large gradients and encourages ﬂat broad minima in the vicinity of small gradients (Prediction 2.3); (iii) these broad minima are known to have low test errors, and consistent with this, we ﬁnd that IGR produces minima with low test error (Prediction 2.4); (iv) the strength of regularization is proportional to the learning rate and network size (Equation 3), (iv) consequently, networks with small learning rates or fewer parameters or both will have less IGR and worse test error, and (v) solutions with high IGR are more robust to parameter perturbations (Prediction 2.5).
It can be difﬁcult to study implicit regularization experimentally because it is not always possible to control the impact of various alternative sources of implicit regularization.
Our analytic approach to the study of implicit regularization in gradient descent allows us to identify the properties of implicit gradient regularization independent of other sources of implicit regularization.
In our experimental work, we take great care to choose models and datasets that were sufﬁciently simple to allow us to clearly expose implicit gradient regularization, yet, sufﬁciently expressive to provide insight into larger, less tractable settings.
For many state-of-the-art deep neural networks trained on large real-world datasets, the IGR is likely to be just one component of a more complex recipe of implicit and explicit regularization.
However, given that many of the favourable properties of deep neural networks such as low test error capabilities and parameter robustness are consistent with IGR, it is possible that IGR is an important piece of the regularization recipe.
A A PPENDIX A. B ACKWARD E RROR A NALYSIS OF THE E XPLICIT E ULER M ETHOD
In this section, we provide formulae for higher order backward analysis correction terms for the explicit Euler method, including the ﬁrst order correction which is required to complete the proof of
Theorem 3.1.
We start by restating the general problem addressed by backward error analysis.
To begin, consider a ﬁrst order differential equation
[Equation]
with vector ﬁeld f : R m → R m .
The explicit Euler method
[Equation]
with step size h produces a sequence of discrete steps θ 0 , θ 1 , . .
, θ n , . .
approximating the solution θ ( t ) of Equation A.1 with initial condition θ 0 .
(In other word, θ n approximates θ ( nh ) for n ≥ 0 .)
However, at each discrete step the Euler method steps off the continuous solution θ ( t ) with a onestep error (or local error) (cid:107) θ 1 − θ ( h ) (cid:107) of order O
( h 2 ) .
Backward error analysis was introduced in numeric integration to study the long term error (or global error) (cid:107) θ n − θ ( nh ) (cid:107) of the numerical method.
More generally, backward error analysis is useful for studying the long term behavior of a discrete numeric method using continuous ﬂows, such as its numerical phase portrait near equilibrium points, its asymptotic stable orbits, and conserved quantities among other properties (see Hairer &
Lubich (1997); Hairer et al. (2006) for a detailed exposition).
It stems from the work of Wilkinson (Wilkinson, 1960) in numerical linear algebra where the general idea in the study of a numeric solution via backward error analysis is to understand it as an exact solution for the original problem but with modiﬁed data.
When applied to numerical integration of ODEs this idea translates into ﬁnding a modiﬁed vector ﬁeld ˜ f with corrections f
i ’s to the original vector ﬁeld f in powers of the step-size
[Equation]
so that the numeric method steps now exactly follow the (formal) solution of the modiﬁed equation :
[Equation]
In other words, if θ ( t ) is the solution of the modiﬁed equation (A.4) and θ n is the n th discrete step of the numerical method, now one has:
[Equation]
In general, the sum of the corrections f
i ’s in (A.3) diverges, and the modiﬁed vector ﬁeld ˜ f is only a formal power series.
For practical purposes, one needs to truncate the modiﬁed vector ﬁeld.
If we truncate the modiﬁed vector ﬁeld up to order n (i.e., we discard the higher corrections f l for l ≥ n + 1 ), the one-step error between the numeric method and the solution of the truncated modiﬁed equation is now of order
O ( h n +1 ) .
It is also possible to bound the long term error, but in this section we will only formally derive the higher corrections f
i for the explicit Euler method.
(We refer the reader to Hairer & Lubich (1997), for instance, for precise bounds on the error for the truncated modiﬁed equation.)
To derive the corrections
f i ’s for the explicit Euler method, it is enough to consider a single step of gradient descent θ + hf ( θ ) and identify the corresponding powers of h in
[Equation]
by expanding the solution θ ( h ) of the modiﬁed equation (A.4) starting at θ into its Taylor series at zero:
[Equation]
Lemma A.1.
In the notation above, we have
[Equation]
where d n − 2 dt n − 2 ddθ g ( θ ) is shorthand to denote the operator d n − 2 dt n − 2 (cid:12)(cid:12) t =0 ddθ (cid:12)(cid:12) θ
= θ ( t ) g
( θ ) and θ ( t ) is the solution of the modiﬁed equation A.4.
We will will use this shorthand notation throughout this section.
Proof.
By deﬁnition θ ( t ) is the solution of (A.4) with initial condition θ (0)
= θ , hence ˙ θ (0) =
˜ f
( θ ) .
Differentiating both sides of (A.4) with respect to t , we obtain
[Equation]
Now the higher derivatives are obtained by differentiating both sides of the last equation n − 2 times with respect to t and setting t = 0 .
The previous lemma gives a formula for the derivatives θ ( n ) (0) .
However in order to compare the powers of h we need to expand these derivatives into power of h , which is what the next lemma does:
Lemma A.2.
In the notation above, we have
[Equation]
where we deﬁne
[Equation]
with f 0 ( θ ) being the original vector ﬁeld f ( θ ) , and (cid:104) , (cid:105) denoting the inner product of two vectors.
Proof.
This follows immediately from (A.9) and by expanding (cid:107) ˜
f ( θ ) (cid:107) 2 in powers of h :
[Equation]
Putting everything together, we now obtain a Taylor series for the solution of the modiﬁed equation as a formal power series in h :
Lemma A.3.
In the notation above, we have
[Equation]
where f 0 is the original vector ﬁeld f , H
=
and we deﬁne recursively
[Equation]
Proof.
Replacing θ ( n ) (0) by their expression in (A.10) in the Taylor series (A.6), we obtain
[Equation]
which ﬁnishes the proof.
Now comparing the Taylor series for the modiﬁed equation solution in its last form in (A.12) with one step of the Euler method
[Equation]
for each order of h , we obtain the following proposition:
Proposition A.4.
The corrections f
i ’s for the Euler method modiﬁed equation in (A.3) are given by the general recursive formula:
[Equation]
where the L n,k are deﬁned by Equation A.11.
Let us use (A.14) to compute the ﬁrst order correction in the modiﬁed equation for the Euler explicit method:
Example A.5.
Suppose f = −∇ E , then the ﬁrst order correction is
[Equation]
which is the only correction we use in Theorem 3.1.
The remarkable thing is that now the ﬁrst two terms of ( A. 3 ) can be understood as the gradient of a modiﬁed function, yielding the following form for the modiﬁed equation (A.4):
[Equation] [Equation]
with ˜ E
= E
+ h 4 (cid:107)∇ E (cid:107) 2 , which is the main result of Theorem 3.1.
A. G EOMETRY OF IMPLICIT GRADIENT REGULARIZATION
In this section, we provide all the details for a proof of Proposition 3.3 and for our claim concerning the relationship between the loss surface slope and the implicit gradient regularizer, which we package in Corollary A.6.1.
The geometry underlying implicit gradient regularization makes it apparent that gradient descent has a bias toward ﬂat minima (Keskar et al., 2017; Hochreiter & Schmidhuber, 1997).
To begin with, consider a loss E over the parameter space θ ∈ R m .
The loss surface is deﬁned as the graph of the loss:
[Equation]
It is a submanifold of R m +1 of co-dimension 1, which means that the space of directions orthogonal to S at a given point ( θ, E ( θ )) is spanned by a single unit vector, the normal vector N ( θ ) to S at ( θ, E ( θ )) .
There is a natural parameterization for surfaces given by the graphs of functions, the
Monge parameterization, where the local chart is the parameter plane: θ → ( θ, E ( θ )) .
Using this parameterization it is easy to see that the tangent space to S at ( θ, E ( θ )) is spanned by the tangent vectors:
[Equation]
for i = 1 , . . .
, m and where the 1 is at the i th position.
Now that we have the tangent vectors, we can verify that the following vector is the normal vector, since its inner product with all the tangent vectors is zero and its norm is one:
[Equation]
We can compute the cosine of the angle between the normal vector N ( θ ) at ( θ, E ( θ )) and the vector ˆ z
= (0 , . .
. , 0 , 1) that is perpendicular to the parameter plane by taking the inner product between these two vectors, immediately yielding the following Proposition:
Proposition A.6.
Consider a loss E and its loss surface S as above.
The cosine of the angle between the normal vector N ( θ ) to the loss surface at ( θ, E ( θ )) and the vector ˆ z
= (0 , . .
. , 0 , 1) perpendicular to the parameter plane
can be expressed in terms of the implicit gradient regularizer as follows:
[Equation]
Now observe that if (cid:104) N ( θ ) , ˆ z (cid:105) is zero this means that the tangent plane to S at ( θ, E ( θ )) is orthogonal to the parameter space, in which case the loss surface slope is maximal and inﬁnite at this point!
On the contrary, when (cid:104) N ( θ ) , ˆ z (cid:105) is equal to 1, the tangent plane at ( θ, E ( θ )) is parallel to the parameter plane, making S look like a plateau in a neighborhood of this point.
Let us make precise what we mean by loss surface slope.
First notice that the angle between ˆ z (which is the unit vector normal to the parameter plane) and N ( θ ) (which is the vector normal to the tangent plane to S ) coincides with the angle between this two planes.
We denote by α ( θ ) this angle:
[Equation]
Now, we deﬁne the loss surface slope at ( θ, E ( θ )) by the usual formula
[Equation]
As we expect, when the loss surface slope is zero this means that the tangent plane to the loss surface is parallel to the parameter plane (i.e., α ( θ ) = 0 ), while when the slope goes to inﬁnity it means the tangent plane is orthogonal to the parameter plane (i.e., α ( θ )
= π/ 2 ).
The following corollary makes it clear that implicit gradient regularization in gradient descent orients the parameter search for minima toward ﬂatter regions of the parameter space, or ﬂat minima, which have been found to be more robust and to possess more generalization power (see Keskar et al.
Hochreiter & Schmidhuber (1997)):
Corollary A.6.1.
The slope of the loss surface S at ( θ, E ( θ )) can be expressed in terms of the implicit gradient regularizer as follows:
[Equation]
Proof.
From (A.15) and the fact that cos α ( θ ) = (cid:104) N ( θ ) , ˆ z (cid:105) , we have that
[Equation]
Now basic trigonometry tells us that in general 1 / cos 2 α = 1 + tan 2 α , which implies here that tan 2 α ( θ ) = mR IG ( θ ) .
Taking the square root of this last expression ﬁnishes the proof.
Remark A.7.
Corollary A.6.1 gives us a very clear understanding of implicit gradient regularization.
Namely, the quantity that is regularized is nothing other than the square of the slope R IG ( θ ) =m slope 2 ( θ ) and the modiﬁed loss becomes ˜ E ( θ ) =
E ( θ ) + h 4 slope 2 ( θ ) .
For explicit gradient regularization, we can now also understand the explicitly regularized loss in terms of the slope:
[Equation]
This makes it clear that this explicit regularization drives the model toward ﬂat minima (with zero slope).
Remark A.8.
There is another connection between IGR and the underlying geometry of the loss surface through the metric tensor.
It is a well-known fact from Riemannian geometry that the metric tensor g ( θ ) for surfaces in the Monge parameterization θ → ( θ, E ( θ )) has the following form:
[Equation]
where δ ij is the Kronecker delta.
Now the determinant | g | , which deﬁnes the local inﬁnitesimal volume element on the loss surface, can also be expressed in terms of the implicit gradient regularizer:
Namely, | g ( θ ) |
= 1 + (cid:107)∇ E ( θ ) (cid:107) 2 = 1 + mR IG ( θ ) .
Solving this equation above for R IG , we obtain a geometric deﬁnition for the implicit gradient regularizer:
[Equation]
which incidentally is zero when the surface looks like an Euclidean space.
We conclude this section by showing that the increase in parameter norm can be bounded by the loss surface slope at each gradient descent step.
Proposition A.9.
Let θ n be the parameter vector after n gradient descent updates.
Then the increase in parameter norm is controlled by the loss surface slope as follows:
[Equation]
A.
[Equation]
Proof.
The triangle inequality applied to one step of gradient descent (cid:107) θ n +1 (cid:107)
=
(cid:107) θ n − h ∇ E ( θ n ) (cid:107)
yields ( (cid:107) θ n +1 (cid:107)−(cid:107) θ n (cid:107) )
≤ h (cid:107)∇ E ( θ n ) (cid:107) , which concludes the proof, since the gradient norm coincides with the loss surface slope.
We now prove a proposition that relates the geometry of the minima of E and ˜ E .
Proposition A.10.
Let E be a non-negative loss.
Then local minima of E are local minima of the modiﬁed loss (cid:101) E .
Moreover, the two losses have the same interpolating solutions (i.e., locus of zeros).
Proof.
Since ∇ (cid:101) E = ∇ E + h 2 D 2 E ∇ E , where D 2 E is the Hessian of E , it is clear that a critical point of E is a critical point of (cid:101) E .
Suppose now that θ ∗ is a local minimum of E .
This means that there is a neighbourhood of θ ∗ where E ( θ ) ≥ E ( θ ∗ ) .
We can add h 4 (cid:107)∇ E ( θ ) (cid:107) 2 on the left of this inequality, since it is a positive quantity, and we can also add h 4 (cid:107)∇ E ( θ ∗ ) (cid:107) 2 on the right of the inequality, since it is zero.
This shows that in a neighborhood of θ ∗ , we also have that (cid:101) E ( θ ) ≥ (cid:101)
E ( θ ∗ ) .
This means that θ ∗ is also a local minimum of (cid:101) E .
Finally, let us see that E and (cid:101) E share the same locus of zeros.
Let θ be a zero of E .
Since E is non-negative and E ( θ ) =
then θ is a global minima, which implies that ∇ E ( θ ) =
also, and hence (cid:101) E ( θ ) = 0 .
Now for positive E , (cid:101) E
( θ )
= 0 trivially implies E ( θ ) = 0 .
A. T HE NTK CONNECTION
In the case of the least square loss, the modiﬁed equation as well as the implicit gradient regularizer take a very particular form, involving the Neural Tangent Kernel (NTK) introduced in Jacot et al.
Proposition A.11.
Consider a model f θ : R d → R c with parameters θ ∈ R m and with least square loss E ( θ
[Equation]
f θ
[Equation]
x i
[Equation]
− y i
[Equation]
The modiﬁed loss can then be expressed as
[Equation]
where K θ is the Neural Tangent Kernel deﬁned by
[Equation]
where (cid:15) k
( θ ) =
f θ ( x k ) − y k ∈ R c is the error vector on data point x k .
In the particular case when the model output is one-dimensional, we can write the modiﬁed loss compactly as follows:
[Equation]
Proof.
Let us begin by computing the gradient
[Equation]
−
[Equation] [Equation]
∇
[Equation]
since ∇ (cid:15) i ( θ )
= ∇ θ f θ ( x i ) is a matrix.
Using that result, we can compute the implicit gradient regularizer in this case:
[Equation]
which concludes the ﬁrst part of the proof.
Now when the model output is one-dimensional, then the (cid:15) i ( θ ) are no longer vectors but scalars.
We can then collect them into a single vector (cid:15) ( θ ) =
( (cid:15) 1 ( θ ) , . . .
, (cid:15) n ( θ )) .
Similarly, the terms K θ ( x i , x j ) are no longer matrices but scalars, allowing us to collect them into a single matrix K θ .
In this notation we now see that the original loss can be written as E = (cid:15) T (cid:15) , yielding ˜ E
= (cid:15) T (1 + K θ ) (cid:15) for the modiﬁed loss, concluding the proof.
For the least square loss, we see that the IGR is expressed in terms of the NTK.
Therefore the NTK entries will tend to be minimized during gradient descent.
In particular, the cross terms K θ ( x i , x j ) will be pushed to zero.
This means that gradient descent will push the maximum error direction ∇ (cid:15)
k
( θ ) at different data points to be orthogonal to each other (i.e., ∇ (cid:15)
k ( θ ) T ∇ (cid:15) l ( θ ) (cid:39) 0 ).
This is good, since a gradient update is nothing other than a weighted sum of these error directions.
If they are orthogonal, this means that the gradient update contribution at point x k will not affect the gradient update contribution at point x l , so the individual data point corrections are less likely to nullifying each other as gradient descent progresses.
A. E XPERIMENT DETAILS FOR THE 2- D LINEAR MODEL
In this section, we provide supplementary results (Figure A.2), hyper-parameter values (Table A.1) and modiﬁed loss derivations for the two parameter model described in Section 5.
This model has a loss given by:
[Equation]
24) where a, b ∈ R are the model parameters and x, y ∈ R is the training data.
The implicit regularization term for this model can be calculated using Equation 4, yielding:
[Equation]
(A.25)
The implicit regularization rate can be calculated using Equation 3, yielding: λ = mh/ 4 = h/ 2 .
(A.26)
[Table]
( a 0 , b 0 ) (2 . 8 , 3 . 5) (75 , 74 . 925) ( x, y ) (1 , 0 . 6) (1 , 0 . 6) h S 10 −10 −h M 2 .
5 × 10 −
[Table]
5 ×
[Table]
−h L 1 . 8 ×
[Table]
n/a µ 0 .4 .
1 × 10 −h Euler 10 −5 × 10 −
[Figure]
The modiﬁed loss can be calculated using Equation 2, yielding:
[Equation]
A.27)
Here, we can see that the global minima for E ( a, b ) (i.e. the zeros) are the same as the global minima for (cid:101) E ( a, b ) since 1 + λ (cid:0) a 2 + b 2 (cid:1) x 2 is positive.
However, as we will see, the corresponding gradient ﬂows are different.
The exact modiﬁed gradient ﬂow for the modiﬁed loss is given by:
[Equation]
The exact gradient ﬂow for the original loss is given by
[Equation]
The exact numerical ﬂow of gradient descent is given by
[Equation]
where ( a n , b n ) are the parameters at iteration step n .
For this model, we have ∇ a E = − bx
( y − abx )
and ∇ b E = − ax
( y − abx ) .
Remark A.12.
In vector notation, we see that the modiﬁed gradient ﬂow equation is
[Equation]
with θ = ( a, b ) .
The last term − (2 λx 2 E ) θ is a central vector ﬁeld re-orienting the original vector ﬁeld ∇ E away from the steepest slopes and toward the origin which coincides in this example with ﬂatter regions where the minimal norm global minima are located.
This phenomenon becomes stronger for parameters further away from the origin, where, coincidentally the slopes are the steepest.
Speciﬁcally, slope
[Equation]
[Figure]
l
[Figure]
In Figure 1a, we plot the trajectories of four different ﬂows starting from the same initial point ( a 0 , b 0 ) , to illustrate the impact of IGR.
We look at two initial points, chosen to illustrate the behaviour of gradient descent in different settings.
The full set of hyper-parameters used for this experiment is given in Table A.1.
First, we calculate the numerical ﬂow of gradient descent with a small learning rate h S using Equation A.30.
Next, we plot the numerical ﬂow of gradient descent with a moderate learning rate h M using Equation A.30.
We then calculate the modiﬁed gradient ﬂow by solving
Equation A.28 numerically using the Euler method, starting from initial point ( a 0 , b 0 ) and using λ = h M / 2 .
For this numerical calculation, we use a very small Euler method step size h Euler so that the Euler method follows the gradient ﬂow of the modiﬁed loss accurately.
We observe that this modiﬁed ﬂow is close to the numerical ﬂow of gradient descent, consistent with Theorem 3.1.
We also plot the trajectory of gradient descent for a large learning rate h L , where backward analysis is no longer applicable (Fig.
A.2) and observe that gradient descent ricochets across the loss surface, stepping over line attractors until it lands in a low gradient region of the loss surface where it converges toward a global minima.
This large learning rate regime can be unstable.
For larger learning rates, or for different initial positions, we observe that gradient descent can diverge in this regime.
In Figure 1b, we explore the asymptotic behaviour of gradient descent by measuring R/E after convergence for a range of models, all initialized at a 0 = 2 . 8 and b 0 = 3 . 5 , with a range of learning rates.
We also explore the impact of explicit gradient regularization, using Equation 8 to deﬁne the explicitly regularized modiﬁed loss for our two-parameter model:
[Equation]
We use this modiﬁed loss for gradient descent:
[Equation]
Here, we have used a very small learning rate, h Euler (Table A.1) and a very large value of µ (Table
A.1).
This allows us to achieve stronger regularization, since µ can be increased to a large value where gradient descent with h = 2 µ would diverge.
A. D EEP NEURAL NETWORK EXPERIMENT DETAILS
In this section, we provide further details for the calculation of R IG ( θ ) in a deep neural network (Figure 2, 3, A.3, A.4, A.5, A.6, A.7, A.8).
For all these experiments, we use JAX (Bradbury et al.,
[Figure]
[Figure] [Figure]
[Figure] [Figure]
and Haiku (Hennigan et al., 2020) to automatically differentiate and train deep neural networks for classiﬁcation.
Conveniently, the loss gradients that we compute with automatic differentiation are the same loss gradients that we need for the calculation of R IG ( θ ) .
We calculate the size of implicit gradient regularization R IG ( θ ) , during model training, using
Equation 4.
We observe that R IG ( θ ) , the loss E ( θ ) and the ratio R IG /E
( θ ) all decrease as training progresses, for all learning rates considered (Figure A.3).
We also observe that the parameter magnitudes grow during training, and this growth slows as R IG ( θ ) becomes small, in agreement with Proposition A.9.
After a sufﬁciently large ﬁxed number of training steps, we see that models with larger learning rates have much smaller values of R IG ( θ ) relative to E ( θ ) , which appears to be consistent with Prediction 2.2.
However, the speed of learning clearly depends on the learning rate h
so it may not be reasonable to compare models after a ﬁxed number of training iterations.
Instead of stopping after a ﬁxed number of iterations, we could stop training after n = T/h iterations, where T is the ﬁxed physical time that naturally occurs in our backward analysis (Equation A.5).
Again, we ﬁnd that models with larger learning rates have lower values of R IG ( θ ) and E ( θ ) after a sufﬁciently large amount of physical training time T (Figure A.5).
However, even for ﬁxed physical time comparisons, we still need to choose an arbitrary physical time point T for making comparisons between models.
The choice of stopping time is effectively an unavoidable form of implicit regularization.
Instead of ﬁxed iteration time or ﬁxed physical time, we use the time of maximum test accuracy as the stopping time for model comparison in Figure 2, 3 and A.6.
We choose this option because it is the most useful time point for most real-world applications.
For each model, we calculate E ( θ ) , R IG ( θ ) and the test accuracy at the time of maximum test accuracy (which will be a different iteration time for each model)
(Figure A.4).
The observation that (i) ﬁxed iteration stopping time, (ii) ﬁxed physical stopping time, and (iii) maximum test accuracy stopping time all have smaller values of R IG ( θ ) /E
( θ ) for larger values of λ , consistent with Prediction 2.2, indicates that the relationships between these quantities cannot be trivially explained to be a consequence of a particular choice stopping time regularization.
In these examples, we use n l = 400 (corresponding to ∼ 9 .
6 × 10 6 parameters) with batch size 32.
In Figure 2 and Figure A.6 we report R IG ( θ ) and test accuracy at the time of maximum test accuracy for a selection of networks of different size, trained with different learning rates.
For models with sufﬁcient capacity to solve the training task and simultaneously minimize R IG ( θ ) , we expect
R IG ( θ ) /E and test error to decrease as λ increases (Prediction 2.2 and 2.4).
To expose this behaviour, we exclude models that fail to reach 100% MNIST training accuracy, such as models that diverge (in the large learning rate regime), and models with excessively small learning rates, which fail to solve the task, even after long training periods.
We observe that test error is strongly correlated with the size of R IG after training (Figure A.6).
We also conﬁrm this for a range of batch sizes, including full batch gradient descent (Figure A.6, top right) with n l = 400 , and for SGD with batch size 32 across a range of learning rates and network sizes (Figure A.6, bottom right).
Finally, to explore IGR and EGR for a larger model we trained a ResNet-18 to classify CIFARimages (using Haiku (Hennigan et al., 2020)).
We used stochastic gradient descent for the training with a batch size of 512 for a range of learning rates l ∈ { 0 . 005 , 0 . 01 , 0 .
We observe the same behaviour as in the MNIST experiments: as the learning rate increases, the values of R IG decrease (Prediction 2.2), the test accuracy increases (Prediction 2.4) and the optimization paths follow shallower slopes leading to broader minima (Prediction 2.3).
The experimental results are summarized by the training curves displayed in Figure A.7 and in Figure A.8, where we plot the relation between learning rate, R IG , and test accuracy taken at the time of maximum test accuracy for each training curve.
[Figure]