T HREE FACTORS INFLUENCING MINIMA IN SGD
Anonymous authors
Paper under double-blind review
A BSTRACT
We focus on the importance of noise in stochastic gradient descent (SGD) based training of deep neural networks (DNNs).
We develop theory that studies SGD training as a stochastic differential equation and show that its stationary distribution is related to the loss surface.
Our analysis suggests that the combination of batch size, learning rate and the variance of the true loss gradients acts as a hyperparameter steering the behavior of SGD and determines the trade-offs between the depth and width of the minima that SGD converges to.
In a nutshell, a higher ratio of learning rate to batch size leads to wider minima.
We validate our theory by examining the correlation between these three factors and the ﬁnal performance and sharpness of the minimum found.
As a veriﬁcation of our theory, we empirically demonstrate that the learning dynamics is similar between experiments with different learning rates and batch sizes in SGD if the ratio of learning rate to batch size is the same.
I NTRODUCTION
Despite being massively over-parameterized models (Zhang et al., 2016), deep neural networks (DNNs) have demonstrated good generalization ability and achieved state-of-the-art performances in many application domains such as image (He et al., 2016) and speech recognition (Amodei et al., 2016).
The reason for this success has been a focus of research recently but still remains an open question.
Our work provides new theoretical insights and useful suggestions for deep learning practitioners.
The standard way of training DNNs corresponds to minimizing a loss function using stochastic gradient descent (SGD) and its variants (Bottou, 1998).
In SGD, parameters are updated by taking a small discrete step depending on the learning rate (LR) in the direction of the negative loss gradient, which is approximated based on a small subset of training examples (called mini-batch).
Since
DNNs are highly non-convex functions with multiple minima, SGD in general converges to different local minima depending on optimization hyper-parameters, initialization and the loss curvature.
Recently, several works (Arpit et al., 2017; Advani & Saxe, 2017; Shirish Keskar et al., 2016) have investigated how SGD impacts generalization in DNNs.
It has been argued that local minima with low curvature (wide minima) can generalize better than minima whose curvature is large (sharp minima)
(Hochreiter & Schmidhuber, 1997; Shirish Keskar et al., 2016).
Speciﬁcally, (Shirish Keskar et al., 2016) found that batch size (BS) strongly affects the curvature of minima that are reached by SGD (i.e., smaller batch sizes usually lead to wider minima).
We generalize this notion to talk about the controllable noise level in SGD which we theoretically show is determined by the ratio of learning rate over batch size.
As such, we discuss how SGD with appropriate noise level tends to converge to deeper and wider minima.
In this vein, while (Dinh et al., 2017) discuss the existence of minima with different widths but which behave similarly in terms of predictions, we argue that
SGD naturally tends to ﬁnd wider minima when appropriate noise levels are used.
We approximate SGD as a continuous stochastic differential equation (Bottou, 1998;
Mandt et al., 2017).
Based on this, we derive the stationary distribution of this stochastic process, and further derive the relative probability of landing in one local minima as compared to another depending on their depth and width.
Our main ﬁnding is that the ratio of learning rate and batch-size along with noise about the true gradient (due to the choice of batch sample) inﬂuences the trade-off between the depth and sharpness of the ﬁnal minima, with a high ratio of learning rate to batch size favouring ﬂatter minima.
In addition, our analysis provides a theoretical justiﬁcation for the empir- ical observation that scaling the learning rate linearly with batch size leads to good performance of
DNNs (Krizhevsky, 2014; He et al., 2016).
We verify our theoretical insights experimentally on several models and datasets.
In particular, we demonstrate that high learning rate to batch size ratio (due to either high learning rate or low batchsize) leads to wider minima and correlates well with better validation performance.
We also observe that a high learning rate to batch size ratio helps to prevent memorization.
Furthermore, we show that multiplying each of the learning rate and the batch size by the same scalar factor results in similar training dynamics, but if the scalar factor gets too large, then the assumptions of the theory are violated and performance drops.
R ELATED WORK
The relationship between stochastic gradient descent (SGD) and sampling a posterior distribution via stochastic Langevin methods has been the subject of discussion in a number of papers (Chen et al., 2014;
Ding et al., 2014; Vollmer et al., 2015; Welling & Teh, 2011; Shang et al., 2015; Sato & Nakagawa, 2014).
In particular, Mandt et al.
(2017) describe the dynamics of stochastic gradient descent (SGD) as a stochastic process that can be divided into three distinct phases.
In the ﬁrst phase, weights diffuse and move away from the initialization.
In the second phase the gradient magnitude dominates the noise in the gradient estimate.
In the ﬁnal phase, the weights are near the optimum.
(Shwartz-Ziv & Tishby, 2017) make related observations from an information theoretic point of view and suggest the diffusion behavior of the parameters in the last phase leads to the minimization of mutual information between the input and hidden representation.
In a similar vein, we relate the SGD dynamics to the stationary distribution of the stochastic differential equation.
Our derivation bears similarity with Mandt et al.
However, while Mandt et al.
(2017) study SGD as an approximate Bayesian inference method in the ﬁnal phase of optimization in a locally convex setting, our end goal is to analyse the stationary distribution over the entire parameter space reached by SGD.
Further, our analysis allows us to compare the probability of SGD ending up in one minima over another, which is novel in our case.
Our work is closely related to the ongoing discussion about the role of large batch size and sharpness of found minima in generalization (Shirish Keskar et al., 2016).
Goyal et al.
; Hoffer et al.
(2017) empirically show that scaling up the learning rate, and training for more epochs, leads to good generalization using large batch size.
Our novelty is in explaining from the theoretical point of view the importance of the speciﬁc ratio of learning rate to batch size.
We also observe that the ﬁnal minima found by large batch size can be wide if a larger learning rate is used.
Our work is also related to the importance of noise in SGD, which has been previously explored.
The main inspiration behind learning rate schedule has been shown to be noise annealing (Bottou, 1998).
Neelakantan et al.
(2015) observe empirically that adding noise can aid optimization of very deep networks.
Our analysis allows us to derive the impact of the gradient’s noise in the SGD stationary distribution.
Additionally, our work also provides intuitions toward explaining the recently proposed Cyclic learning rate (CLR) schedule (Smith, 2015).
Cyclic learning rate schedules have demonstrated good optimization and generalization performances, but are grounded on empirical observation.
We also show that one can replace learning rate annealing with an equivalent batch size schedule.
It suggests that the beneﬁt of cyclic learning rate relates to the noise that it induces.
T HEORETICAL R ESULTS
Our focus in this section is on ﬁnding the relative probability with which we end optimization in a certain minimum, as compared to another minimum.
We will ﬁnd that the relative probability depends on the local geometry of the loss function at each minimum, and on the BS, the LR and the covariance of the stochastic gradients.
3. S ETUP
We follow a theoretical setup similar to Mandt et al.
(2017), approximating SGD with a continuoustime stochastic process, which we now outline.
Let us consider a model parametrized by θ = { θ 1 , . .
, θ q } .
For N training examples x i , i ∈ { 1 , ..., N } , the loss function, L ( θ ) , and the corresponding gradient g ( θ ) , are deﬁned based on the sum over the loss values for all training examples.
Stochastic gradients g ( S )
( θ ) arise when we consider a batch B of size S < N of random indices drawn uniformly from { 1 , ..., N } and form an (unbiased) estimate of loss and gradient based on the corresponding subset of training examples
[Equation]
We consider stochastic gradient descent (SGD) with learning rate η , as deﬁned by the update rule
[Equation]
We now make the following assumptions:
(1) By the central limit theorem (CLT), we assume the gradient noise is Gaussian with covariance matrix 1 S C ( θ )
[Equation]
We note that the covariance is symmetric positive-semideﬁnite, and so can be decomposed into the product of two matrices C ( θ )
=
B ( θ ) B (cid:62) ( θ ) .
(2) We assume the discrete process of SGD can be approximated by the continuous time limit of the following stochastic differential equation (known as a Langevin equation)
[Equation]
where f ( t ) is a normalized Gaussian time-dependent stochastic term.
3. M AIN R ESULTS
The Langevin equation is a stochastic differential equation, and we are interested in its stationary distribution to gain insights into the behavior of SGD and the quality of minima it converges to.
It can be derived by ﬁnding the stationary solutions of a partial differential equation known as the
Fokker-Planck equation, which governs the evolution of the probability density of the value of the parameters with time 1 .
The Fokker-Planck equation and its derivation can be found in Appendix A in equation (5).
We arrive at the following theorem (all the notations are deﬁned in the setup section above):
Theorem 1 (Stationary Distribution) .
Assume that the gradient covariance is proportional to the identity, i.e. C ( θ )
=
σ 2 ( θ ) I and that |∇ θ ( σ 2 ( θ ))
| (cid:28) 2 S | g | η .
Then the stationary distribution of the stochastic differential equation 6 is given by
[Equation]
where n ≡ σ 2 ( θ ) η 2 S and P 0 is a normalization constant.
Here P ( θ ) deﬁnes the density over the parameter space.
The above result says that if we run SGD long enough (under the assumptions made), then the likelihood of the parameters being in a particular state asymptotically follows the above density.
Note, that n ≡ σ 2 ( θ ) η 2 S is a measure of the noise in the system and it depends on two parts: n c ≡ η/S , which is the controllable noise set by the choice of hyper-parameters, and the gradient variance σ 2 ( θ ) , which is not a tunable hyper-parameter.
Given the probability density P ( θ ) , we are now interested in deriving the probability of ending at a given minimum, θ A , which we will denote by lowercase p A = ˜ p
A C , where C is a normalization constant which is the same for every mimnima (the unnormalized probability ˜ p
A is all we are interested in when estimating the relative probability of ﬁnishing in a given minimum compared to another one).
This probability is derived in Appendix D, and given in the following theorem, which is the core result of our theory.
context of the stochastic gradient langevin dynamics (SGLD) algorithm in the Bayesian learning setting (Sato & Nakagawa, 2014).
While in this setting it was found that the stochastic noise can be ignored, we arrive at a different conclusion.
Theorem 2 (Probability of ending in minima A) .
Assume the loss is locally convex with Hessian
H A and loss L A at the minimum θ A .
Then the unnormalized probability of ending in minima θ A is
[Equation]
where q is the number of parameters and n A ≡ n c σ 2 ( θ A ) is the noise at the minimum.
Here we see that the probability of landing in a given minimum depends on the geometry of the loss surface (captured by the height L A and the determinant of the Hessian det H A , which determines the width of the minimum) as well as on the noise n A at θ A .
For larger noise, n A , the factor exp (cid:16) − L A n A (cid:17) is down-weighted compared to (cid:113) (2 πn A ) q det H A , which grows with larger n A .
Thus, with larger noise the width gets more important than the height of the minimum, while the converse holds if the noise is low enough.
Also notice, that with n A the probability ˜ p
A directly depends on the ratio of learning rate and batch size n c = η/S .
We will explore the implications of this experimentally in Section 4.
To see which kind of minima is preferred, it is instructive to consider the ratio of probabilities p A and p B given by
[Equation]
Let us now look at the special case where n A = n = n B (other cases can be found in Appendix E).
Without loss of generality we can set L A ≥ L B .
If L A = L B the minimum with lower determinant of the Hessian (i.e. the ﬂatter minima) is more probable.
If L A ≥ L B , it holds that p
A ≥ p B if and only if
[Equation]
That is, there is an upper bound on the inverse of the noise for θ A to be favored in the case that its loss is higher than at θ B , and this upper bound depends on the difference in the heights compared to the ratio of the widths.
In particular we can see that if det H B < det H A , then Y < 0 , and so no amount of noise will result in θ A being more probable than θ B – that is, if the minimum at θ A is both higher and sharper than the minimum at θ B , it is never reached with higher probability than θ B , regardless of the amount of noise.
However, if det H B > det H A then Y > 0 , and there is a lower bound on the noise
[Equation]
to make θ A more probable than θ B – that is if the minimum at θ A is higher but ﬂatter than the minimum at θ B , it is favored over θ B , as long as the noise is large enough, as deﬁned by eq.
To summarize, the presented theory shows that the noise level in SGD (which is deﬁned by the learning rate, batch size and gradient covariance) controls the extent to which optimization favors wider over lower minima.
Increasing the noise by increasing the ratio of learning rate and batch size, increases the probability of wider compared to lower optima.
E XPERIMENTS 4. I MPACT OF SGD ON MINIMA
In this section, we empirically study the impact of the learning rate η , batch size S on the local minimum that SGD ﬁnds.
We ﬁrst focus on an MLP with 20 layers with ReLU activation functions (similar to one used in (Neelakantan et al., 2015)) and trained on FashionMNIST.
We study how the controllable noise
[Figure] [Figure]
ratio n c = ηS leads to minima with different curvatures and validation accuracy.
To measure the curvature at that minimum, we compute the norm of its Hessian using a ﬁnite difference (Wu et al., 2017) (higher Hessian norm imples higher sharpness of minima).
In Figure 1a, we report the norm of the Hessian for local minima obtained by SGD for different n c with ηS (where η ∈ [5 e − 3 , 5 e − 2] and S ∈ [50 , 1000] ).
As n c = ηS grows, we observe that the norm of the Hessian at the minima also decreases, suggesting that higher ηS pushes the optimization towards a ﬂatter minimum.
This agrees with Theorem 2, (3), that higher ηS favors ﬂatter over sharper minima.
In Figure 1b, we explore the impact of n c = ηS on the ﬁnal validation performance and conﬁrm that better generalization correlates with higher n c .
Taken together, Fig.
1a and Fig.
1b imply wider minima correlate well with better generalization.
As n c = ηS increases, SGD ﬁnds local minima that generalize better.
In the Appendix, we report similar results for Resnet applied on CIFAR10 and the 20 layer ReLU network with bad initialization.
To further illustrate the behavior of SGD with different noise levels, we train three Resnet56 models on CIFAR10 using naive SGD, but with different ηS .
Our baseline model uses η =0 .
1 S =128 .
We deﬁne a large batch model η =0 .
1 S =1024 and a small learning rate model η =0 .
01 S =128 .
The large batch and the small learning rate model have approximately the same ηS ratio.
In Figure 2, we follow (Goodfellow et al., 2014) and look at the loss between interpolation of different pairs of models.
More speciﬁcally, let f 1 and f 2 be the early stop models of two SGD optimizers using different ηS .
In Fig.
2, we report
L ((1 − α )
f 1 + αf 2 where L is the loss function and α ∈ [ − 1 , 2] .
We observe that both model with large batch or low learning rate end in a sharper minimum relatively to the baseline model, both having a lower ηS than the baseline.
Each of the three plots in Figureadds empirical weight to our theoretical prediction that higher n c = η/S gives preference to wider minima over sharper minima.
Finally we perform a similar experiment on VGG-11 on CIFAR-10, but in this case, we train all the models with the same noise level but different values of learning rate and batch size.
Speciﬁcally, in this case we use η =0 .
1 × β S =50 × β , where we use β = 0 .
We then interpolate between the model
[Figure]
[Figure] [Figure]
our theory predicts the minima
[Figure]
parameters of the pair of models and record loss and accuracy (shown in ﬁgure 3).
We ﬁnd that all the minima in this case have similar width and depth qualitatively showing that for the same noise ratio, SGD ends up in minima of similar quality.
4. ηS RATIO DETERMINES LEARNING DYNAMICS OF SGD
A prediction of the theory is that the learning rate η and batch size S only appear in the ratio n c = η/S in the expressions for the probability of ﬁnishing in a given minimum.
Based on this, the endpoint of SGD with a learning rate schedule η → η/a , for some a > 0 , and a constant batch size S , should be the same as the endpoint of SGD with a constant learning rate and a batch size schedule S → aS .
We see in Fig.
4 that exchanging a learning rate schedule for a batch size schedule leads to similar train and test accuracy, for both cyclic and constant schedules, in the case where the controllable noise n c = η/S is the same in each case.
In particular we notice very similar dynamics in the case of a cyclic learning rate when compared to a cyclic batch size.
Full loss curves and learning rate and batch size schedules for this experiment can be found in Appendix G.2 in Fig.
The test accuracy for cyclic batch size and cyclic learning rate was 89.39% and 89.24% respectively.
For constant (batch size, learning rate) =
(128, 0.001) and constant (batch size, learning rate) =
(640, 0.005) the test accuracy was 87.25% and 86.92% respectively.
Note that exchanging batch size for learning rate such that the ratio η/S remains constant is different to the common assumption that one should trade them so as to keep the ratio η/ √ S constant.
This commonly used heuristic of scaling the learning rate with the square root of the batch size is used for example in (Hoffer et al., 2017), as a way of keeping the covariance matrix of the parameter update step the same for any batch size.
However, our point here is that our theory and experiments
[Figure]
suggest exchanging learning rate and batch size in a way that keeps the ratio of controllable noise n c = η/S constant will result in the same stationary distribution.
4. I MPACT OF SGD ON M EMORIZATION
To generalize well, a model must identify the underlying pattern in the data instead of simply perfectly memorizing each training example.
An empirical approach to test for memorization is to see if a deep neural network can ﬁt a training set with random noise labels rather than the true labels (Zhang et al., 2016; Arpit et al., 2017).
In this section, we highlight that SGD with a sufﬁcient amount of noise reduces the amount of memorization in a network.
Experiments are performed using the MNIST dataset and an MLP similar to the one in (Arpit et al., 2017), but with 256 hidden units.
We train the MLP on the dataset with different amounts of random labels in the training set.
For each label noise level, we evaluate the impact of ηS , controlling the amount of noise in the SGD and its impact on the generalization performances.
Speciﬁcally, we run a grid of batch size in 50 , 100 , 200 , 400 , 800 , learning rate in 0 .
05 , 0 . 1 , 0 . 2 and momentum in 0 . 0 and 0 .
Models are trained for 300 epochs.
We report in Fig.
5 the MLP performances for both the noisy training set and the clean validation set (without random label, to measure the model generalization ability).
Our analysis highlights that SGD with low controllable noise n c = ηS steers the endpoint of optimization towards a minima with low generalization ability, associated with a sharp minima, as predicted by our theory.
We observe in Fig.
5 that larger noise in SGD (regardless if controlled using smaller batch size or larger learning rate) leads to solutions which generalize better for the same amount of memorized random labels on the training set.
We also reproduce the observation reported in (Arpit et al., 2017): that randomization roughly starts after reaching maximum generalizationfor more details see the curves in Fig.
11 included in Appendix G.3.
They show that lower controllable noise ηS optimize much slower and do not reach larger than 20% accuracy on random labels.
For runs with momentum we exclude higher learning rates than 0 .
02 as they lead to divergence (see appendix).
4. B REAKING P OINT OF THE T HEORY IN P RACTICE
Our analysis relies on the fact that the gradient step is sufﬁciently small so that ﬁrst order approximation of a Taylor expansion is a good estimate of the loss function.
In the case where the learning rate becomes too high, the ﬁrst order approximation of a Taylor expansion is no longer suitable, in which case the continuous limit of the discrete SGD update equation will no longer be valid.
In this case, the stochastic differential equation doesn’t hold, and hence neither does the Fokker-Planck equation, and so our theory will not be predictive.
In particular, we don’t expect to arrive at the same stationary distribution as governed just by the ratio η/S .
This is exempliﬁed in Fig 6, where similar learning dynamics and ﬁnal performance is seen when simultaneously multiplying the learning rate and batch size by a factor β up to a certain limit.
This is done for different train-set size to see if the breaking point depends on this size.
This plots seem to suggest that this happens for smaller β
[Figure]
values when dataset size is smaller.
We highlight other limitations due to our theory assumption in appendix F. A similar experiment is performed on Resnets in ﬁgure 7.
4. C YCLICAL B ATCH AND L EARNING R ATE S CHEDULES
It has been observed that a cyclic learning rate (CLR) schedule leads to better generalization (Smith, 2015).
In Sec.
4.2 we demonstrated that one can exchange cyclic learning rate schedule (CLR) with batch size (CBS) and approximately preserve the practical beneﬁt of CLR.
This inspired us to hypothesize that by changing between controllable noise levels ( ηS ) CLR switches between sharp/deep and wide/shallow minima.
To validate that, we run VGG-11 on CIFAR10 using 4 training schedules: CLR with stepsize of 4 and 15 epochs in each stage, CBS with stepsize 4 and 15 epochs in each stage.
Each run is repeated 8 times and for each variant we track evolution of sharpness and accuracy.
We observe that CBS and CLR with longer stages lead to 89 .
9 ± 0 .
20% and 90 .
2 ± 0 .
10% test accuracy, in both cases 0 .
2 − 0 .
3% above variance with shorter range.
CBS and CLR with long stepsize each seem to be promising schedules for training DNNs.
Finally, we validate that CBS and CLR switch between sharp/deep and wide/shallow minima, suggesting that CLR improves convergence time to stationary distribution, as seen in Fig.
13 of Appendix G.4.
C ONCLUSIONS
We shed light on the role of noise in SGD optimization of DNNs and argued that three factors (batch size, learning rate and gradient variance) strongly inﬂuence the properties (loss and width) of the ﬁnal minima at which SGD converges.
Learning rate and batch size of SGD can be viewed as one effective hyper-parameter that act as a factor controllable noise n c = η/S , which, together with the gradient covariance inﬂuences the trade-off between the loss and width of the ﬁnal minima (high noise favors wide minima), which in turn tunes the robustness of the prediction function.
Further, we experimentally verify that the controllable noise n c = η/S determines the width and height of the minima towards which SGD converges.
We also show the impact of this controllable noise on the memorization phenomenon.
We discussed the limitations of the theory and in what situations it breaks down, exempliﬁed by when the learning rate gets too large.
We also experimentally verify that η and S can vary in linear proportion as long as the controllable noise η/S remains the same.
In addition, our experiments suggest that cyclical learning rates oscillate between sharp/deep and wide/shallow minima as long as the stage of increased noise is long enough to allow for mixing.
A PPENDIX A D ERIVATION OF THE F OKKER -P LANCK E QUATION
In this appendix we derive the Fokker-Planck equation.
For any stochastic differential equation, since the evolution is noisy we can’t say exactly where in parameter space the parameter values will be at any given time.
But we can talk about the probability P ( θ , t | θ 0 , t 0 ) that a parameter takes a certain value θ at a certain time t , given that it started at θ 0 , t 0 .
That is captured by the
Fokker-Planck equation, which reads
[Equation]
In this appendix we will derive the above equation from the stochastic differential equation (6).
We will not be interested in pure mathematical rigour, but intend the proof to add intuition for the machine learning audience.
For brevity we will sometimes write the probability just as P ( θ , t ) .
We will sometimes make use of tensor index notation, where a tensor is denoted by its components, (for example θ i are the components of the vector θ ), and we will use the index summation convention, where a repeated index is to be summed over.
We start with the stochastic differential equation
[Equation]
A formal expression for the probability for a given noise function f is given by P ( θ , t )
=
δ ( θ − θ f ) , but since we don’t know the noise, we instead average over all possible noise functions to get
[Equation]
While this is just a formal solution we will make use of it later in the derivation.
We now consider a small step in time δt , working at ﬁrst order in δt , and ask how far the parameters move, denoted δ θ , which is given by integrating (6)
[Equation]
where we’ve assumed that δ θ is small enough that we can evaluate g at the original θ .
We now look at expectations of δ θ .
Using the fact that the noise is normalized Gaussian, E ( f ( t ))
= 0 , we get (switching to index notation for clarity)
[Equation]
) and using that the noise is normalized Gaussian again, we have that E ( f ( t ) f ( t (cid:48) ))
=
I δ ( t − t (cid:48) ) , leading to
[Equation]
If at time t we are at position θ (cid:48) and end up at position θ = θ (cid:48)
+ δ θ at time t + δt , then we can take (7) with θ f = θ (cid:48) and Taylor expand it in δθ
i
[Equation]
We deal with the derivatives of the delta functions in the following way.
We will use the following identity, called the Chapman-Kolmogorov equation, which reads
[Equation]
(12) for any t (cid:48) , such that t 0 ≤ t (cid:48) ≤ t + δt .
This identity is an integral version of the chain rule of probability, stating there are multiple paths of getting from an initial position θ 0 to θ at time t + δt and one should sum all of these paths.
We will now substitute (11) into the ﬁrst term on the right hand side of (12) with t (cid:48) set to be t , and apply integration by parts (assuming vanishing boundary conditions at inﬁnity) to move the derivatives off of the delta functions and onto the other terms.
We end up with
[Equation]
We can then take the ﬁrst term on the right hand side of (13) to the other side, insert (9) and (10), divide by δt and take the limit δt → 0 , getting a partial derivative with respect to time on the left hand side, leading directly to the Fokker-Planck equation quoted in the text (5) (where we have reverted back to vector notation from index notation for conciseness)
[Equation]
B I NTUITIONS FROM F OKKER -P LANCK
In this appendix we will give some supplementary comments about the intuition we can gain from the Fokker-Planck equation (5).
If the learning rates, batch size and covariance all change slowly, and the covariance is proportional to the identity C
= σ
I then we can rewrite the Fokker-Planck equation in the following form
[Equation]
15) where we have rescaled the time coordinate to be t η ≡ tη .
One can now see that in terms of this rescaled time coordinate, the ratio between the drift and diffusion terms is governed by the following quantity
[Equation]
which we call the noise.
In terms of the balance between drift and diffusion, we see that higher n gives rise to a more diffusive evolution, while a lower n allows the potential drift term to dominate.
In the next section we will see how n controls the stationary distribution at which SGD converges to.
For now we highlight that in terms of this rescaled coordinate, only n controls the evolution towards this stationary distribution (not just its endpoint), in terms of the rescaled time t η .
That is, learning rate and batch size are interchangable in the sense that n is invariant under transformations S → aS , η → η/a , for a > 0 .
But note that the time it takes to reach the stationary distribution depends on η as well, because of the rescaled time variable.
For example, for a higher learning rate, but constant n , one arrives at the same stationary distribution, but in a quicker time by a factor of 1 /η .
However, a caution here is necessary.
The ﬁrst order SGD update equation only holds for small enough η that a ﬁrst order approximation to a Taylor expansion is valid, and hence we expect the ﬁrst order approximation of SGD as a continuous stochastic differential equation to break down for high η .
Thus, we expect learning rate and batch size to be interchangable up to a maximum value of η at which the approximation breaks.
C S TATIONARY D ISTRIBUTION
In this appendix we will prove equation (2) quoted in the main text, which claims that the stationary (ie late time, independent of time) solution of the Fokker-Planck equation has the following form
[Equation]
(17) where n ≡ σ 2 η 2 S , under the assumptions that the covariance is proportional to the identity C ( θ ) =
σ 2 ( θ ) I and further that |∇ θ ( σ 2 ) | (cid:28) 2 S | g | η .
Here P 0 is a normalization constant.
In order to prove this is the stationary distribution, we need to solve the Fokker-Planck equation (5) with the left hand side set equal to zero (since we are looking for the stationary solution).
To do this, we begin by writing the Fokker-Planck equation in a slightly different form, making use of a probability current J , deﬁned [Equation]
[Equation]
in which the Fokker-Planck equation becomes
[Equation]
At this point we use the assumptions that C ( θ ) =
σ 2 ( θ ) I and that |∇ θ ( σ 2 ) | (cid:28) 2 S | g | η in order to bring C ( θ ) outside of the divergence in (18) and turn the divergence into a gradient of P .
We get
[Equation]
Now we notice we can write this in the following way
[Equation]
21) where we’ve deﬁned n = ησ 2 2 S .
Then we consider that ∇ θ · J = 0 and hence
[Equation]
This equation has the rather easy to spot solution, simply
[Equation]
23) which is the desired stationary solution we intended to ﬁnd.
D D ERIVATION OF P ROBABILITY OF E NDING IN A GIVEN M INIMA
In this appendix we derive the discrete set of probabilities of ending at each minima, as given in (3).
We work locally near θ A and take the following approximation of the loss function, since it is near a minimum,
[Equation]
where H A is the Hessian, and is positive deﬁnite for a minimum.
The distribution P A ( θ ) is a probability density, while we are interested in the discrete set of probabilities of ending at a given minimum, which we will denote by lowercase p A .
To calculate this discrete set of probabilities, for each minimum we need to integrate this stationary distribution over an interval containing the minimum.
We take σ 2 to take the constant value σ 2 A within this minimum and hence denote n by n A .
Integrating (2) in some region R A around θ A and using (24)
[Equation]
where in the last line we assume the region is large enough that an approximation to the full Gaussian integral can be used (i.e. that the tails don’t contribute, which is ﬁne as long as the region is sufﬁciently larger than det H A ) – note that the region can’t be too large, otherwise we would invalidate our local assumption.
The picture is that the minima are sufﬁciently far apart that the region can be taken sufﬁciently large for this approximation to be valid.
Since we are interested in relative probabilities between different minima, so we can consider the unnormalized probability, dropping factors that are common amongst all minima we get
[Equation]
This is the required expression given in the main text.
As an aside, we note that the derivation used above talked about true minima, with a positive deﬁnite
Hessian, i.e. a minima in all directions.
In the real world application to deep neural networks with a large number of parameters, it is unrealistic to expect the endpoint of training to be in a true minimum.
Instead, it is more likely to be a point in which the Hessian has some positive eigenvalues λ
Ai in just a few directions, say d (cid:28) q , with the other eigenvalues (approximately) zero.
In this case, the relative probabilities are given by
[Equation]
In the main paper we only consider the simplest case of true minima, though much of what we say can be extended to the case of points with some positive Hessian eigenvalues with some ﬂat directions.
E F URTHER S PECIAL C ASES
In this appendix we explore other special cases in which the ratio of probabilities for ending in one or another minima simpliﬁes, allowing for an intuitive picture of the way SGD favours some minima over others.
We begin this appendix with the special case of equal height minima, and then look at equal width minima.
[Figure]
E. E QUAL H EIGHT M INIMA
Consider now the case where the minima are the same height
L A =
L = L B .
The simplest case is to look at equal width minima, in which case A is favoured over B if and only if n A > n B , with the relation being exponential – ie a noisier minima is exponentially more probable than a less noisy one.
If on the other hand the minima are not equal width, but do have equal height, we can without loss of generality take det H B > det H A .
We ﬁnd that p A > p B if and only if the following inequality is satisﬁed
[Equation]
This places an interesting constraint on the noises and the height L .
E. E QUAL W IDTH M INIMA
The case of equal width minima has been covered in the previous sections, but with either equal width or equal height presupposed.
Here, if we don’t assume the noise or height are equal, but that the widths are the same, then p A > p B if and only if
[Equation]
It is not clear how to simplify this further, but one can see that if L A < L B then this inequality will be satisﬁed for most reasonable values of n A and n B , unless n B is exponentially larger than n A .
F L IMITATION OF OUR A NALYSIS
Due to our assumptions, we expect the theory to become unreliable when the discrete to continuous approximation fails, when the variance in gradients changes as fast as the loss, when the covariance in gradients cannot be written as proportional to the identity matrix, the ﬁnite size of the training set and when momentum is considered.
We discussed limit of the discrete to continous approximation in 4.4 (further illustrated in 7).
When the gradient variance changes rapidly compared to the loss, there will be a different, more complicated stationary solution that arises when solving the stationary Fokker-Planck equation.
We expect this to happen in certain datasets and in using certain architectures, though it is hard to say an exact relation between datasets, architectures and gradient variances.
When the covariance in gradients can’t be written as proportional to the identity matrix, the stationary solution to the Fokker-Planck equation is the solution to a complicated partial differential
[Figure]
V a li d a t i o n A cc u r a c y equation, and one can’t easily spot the solution by inspection as we do in Appendix C. We expect this approximation to break down especially in the case of complicated architectures where different gradient directions will be have very different gradient covariances.
Our theory does not involve the ﬁnite size of the training set, which is a drawback of the theory.
This may be especially apparent when the batch size becomes large, compared to the training set size, and we expect the theory to break down at this point.
Finally, we mention that momentum is used in practical deep learning optimization algorithms.
Our theory does not consider momentum, which is a drawback of the theory and we expect it to break down in models in which momentum is important.
We can write a Langevin equation for the case of momentum, with momentum damping coefﬁcient µ ,
[Equation]
where the velocity is v = d θ /dt .
The Fokker-Planck equation corresponding to this Langevin equation with momentum is
[Equation]
[Equation]
where P is now a function of velocity v as well as θ and t , and we’ve again assumed the gradient covariance varies slowly compared to the loss.
From this more complicated Fokker-Planck equation, it is hard to spot the stationary distribution.
We leave the study of this for further work.
For now, we just note that a factor of η can be taken from the right hand side, and then the ratio between the diffusion and drift terms will again be the noise n = ησ 2 / (2 S )
[Equation]
O THER
E MPIRICAL R ESULTS G. M ORE SGD I MPACTS ON M INIMA
Figure 8 reports the validation accuracy for Resnet models trained on CIFAR with different LR/BS ratio.
Figure 9 reports 20 layer ReLU experiment variant with scaled initialization by 0 . 7 , which explores case in which model is driven to bad local minima (and is not addressed by our theory).
G. M ORE L EARNING R ATE & B ATCH S IZE E XCHANGEABILITY
In this appendix we show in more detail the experiments of Section 4.2 which show the exchangeability of learning rate and batch size.
In Fig.
10 we show the log cross entropy loss, the epochaveraged cross entropy loss, train, validation and test accuracy as well as the batch size schedule,
[Figure]
learning rate schedule.
We see in the lr/bs plot that the orange and blue lines (cyclic schedules) have the same ratio of lr/bs as each other throughout, and that their dynamics are very similar to each other.
The same holds for the green and red lines, of constant batch size and learning rate.
This supports the theoretical result of this paper, that the ratio S/η governs the stationary distribution of
SGD.
We note that it is not just in the stationary distribution that the exchangeability holds in these plots - it appears throughout training, highlighted especially in the cyclic schedules.
We postulate that due to the scaling relation in the Fokker-Planck equation, the exchangeability holds throughout learning as well as just at the end, as long as the learning rate does not get so high as to ruin the approximations under which the Fokker-Planck equation holds.
G. D ETAILS ON MEMORIZATION EXPERIMENT
We report learning curves from memorization experiment with 0 . 0 momentum, see Fig.
We additionally conﬁrm similar result as in previous experiments and show correlation between batch size and learning rate ratio and norm of Hessian, see Fig.
G. C YCLIC BATCH AND LEARNING SCHEDULE
We report sharpness and accuracy evolution for learning schedules discussed in Sec. 4.5, see Fig. 13.
[Figure]
[Figure] [Figure]