We study the statistical properties of the endpoint of stochastic gradient descent (SGD). We approximate SGD as a stochastic differential equation (SDE) and consider its Boltzmann Gibbs equilibrium distribution under the assumption of isotropic variance in loss gradients.. Through this analysis, we find that three factors – learning rate, batch size and the variance of the loss gradients – control the trade-off between the depth and width of the minima found by SGD, with wider minima favoured by a higher ratio of learning rate to batch size. In the equilibrium distribution only the ratio of learning rate to batch size appears, implying that it’s invariant under a simultaneous rescaling of each by the same amount. 
We experimentally show how learning rate and batch size affect SGD from two perspectives: the endpoint of SGD and the dynamics that lead up to it. For the endpoint, the experiments suggest the endpoint of SGD is similar under simultaneous rescaling of batch size and learning rate, and also that a higher ratio leads to flatter minima, both findings are consistent with our theoretical analysis. We note experimentally that the dynamics also seem to be similar under the same rescaling of learning rate and batch size, which we explore showing that one can exchange batch size and learning rate in a cyclical learning rate schedule. Next, we illustrate how noise affects memorization, showing that high noise levels lead to better generalization. Finally, we find experimentally that the similarity under simultaneous rescaling of learning rate and batch size breaks down if the learning rate gets too large or the batch size gets too small. Despite being massively over-parameterized BID13 , deep neural networks (DNNs) have demonstrated good generalization ability and achieved state-of-the-art performances in many application domains such as image BID13 and speech recognition BID1 . 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 involves minimizing a loss function using SGD and its variants BID4 . In SGD, parameters are updated by taking a small discrete step depending on the learning rate in the direction of the negative loss gradient, which is approximated based on a small subset of training examples (called a mini-batch). Since the loss functions of DNNs are highly non-convex functions of the parameters, with complex structure and potentially multiple minima and saddle points, SGD generally converges to different regions of parameter space depending on optimization hyper-parameters and initialization.Recently, several works BID2 BID0 BID28 have investigated how SGD impacts generalization in DNNs. It has been argued that wide minima tend to generalize better than sharp minima BID15 BID28 . This is entirely compatible with a Bayesian viewpoint that emphasizes targeting the probability mass associated with a solution, rather than the density value at a solution BID21 . Specifically, BID28 find that larger batch sizes correlate with sharper minima. In contrast, we find that it is the ratio of learning rate to batch size which is correlated with sharpness of minima, not just batch size alone. In this vein, while BID9 discuss the existence of sharp minima which behave similarly in terms of predictions compared with wide minima, we argue that SGD naturally tends to find wider minima at higher noise levels in gradients, and such wider minima seem to correlate with better generalization.In order to achieve our goal, we approximate SGD as a continuous stochastic differential equation BID3 BID22 BID19 . Assuming isotropic gradient noise, we derive the Boltzmann-Gibbs equilibrium distribution of this stochastic process, and further derive the relative probability of landing in one local minima as compared to another in terms of their depth and width. Our main finding is that the ratio of learning rate to batch-size along with the gradient's covariances influence the trade-off between the depth and sharpness of the final minima found by SGD, with a high ratio of learning rate to batch size favouring flatter minima. In addition, our analysis provides a theoretical justification for the empirical observation that scaling the learning rate linearly with batch size (up to a limit) leads to identical performance in DNNs BID18 BID12 .We verify our theoretical insights experimentally on different 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 show that a high learning rate to batch size ratio helps prevent memorization. Furthermore , we observe that multiplying each of the learning rate and the batch size by the same scaling factor results in similar training dynamics. Extending this observation, we validate experimentally that one can exchange learning rate and batch size for the recently proposed cyclic learning rate (CLR) schedule BID31 , where the learning rate oscillates between two levels. Finally, we discuss the limitations of our theory in practice. In the theoretical section of this work we treat the learning rate as fixed throughout training. However, in practical applications, the learning rate is annealed to a lower value, either gradually or in discrete jumps. When viewed within our framework, at the beginning with high noise, SGD favors width over depth of a region, then as the noise decreases, SGD prioritizes the depth more stronglythis can be seen from Theorem 3 and the comments that follow.In the theoretical section we made the additional assumption that the covariance of the gradients is isotropic, in order to be able to derive a closed form solution for the equilibrium distribution. We do not expect this assumption to hold in practice, but speculate that there may be mechanisms which drive the covariance towards isotropy, for example one may be able to tune learning rates on a per-parameter basis in such a way that the combination of learning rate and covariance matrix is approximately isotropic -this may lead to improvements in optimization. Perhaps some existing mechanisms such as batch normalization or careful initialization give rise to more equalized covariance -we leave study of this for future work.We note further that our theoretical analysis considered an equilibrium distribution, which was independent of the intermediate dynamics. However, this may not be the case in practice. Without the isotropic covariance, the system of partial differential equations in the late time limit will in general have a solution which will depend on the path through which optimization occurs, unless other restrictive assumptions are made to force this path dependence to disappear . Despite this simplifying assumption, our empirical results are consistent with the developed theory. We leave study of path dependence and dynamics to future work.In experiments investigating memorization we explored how the noise level changes the preference of wide minima over sharp ones. BID2 argues that SGD first learns true labels, before focusing on random labels. Our insight is that in the second phase the high level of noise maintains generalization. This illustrates the trade-off between width of minima and depth in practice. When the noise level is lower, DNNs are more likely to fit random labels better, at the expense of generalizing less well on true ones. We shed light on the role of noise in SGD optimization of DNNs and argue that three factors (batch size, learning rate and gradient variance) strongly influence the properties (loss and width) of the final minima at which SGD converges. The learning rate and batch size of SGD can be viewed as one effective hyper-parameter acting as a noise factor n = η/S. This, together with the gradient covariance influences the trade-off between the loss and width of the final minima. Specifically, higher noise favors wider minima, which in turn correlates with better generalization.Further, we experimentally verify that the noise n = η/S determines the width and height of the minima towards which SGD converges. We also show the impact of this noise on the memorization phenomenon. We discuss the limitations of the theory in practice, exemplified by when the learning rate gets too large. We also experimentally verify that η and S can be simultaneously rescaled as long as the noise η/S remains the same.
Three factors (batch size, learning rate, gradient noise) change in predictable way the properties (e.g. sharpness) of minima found by SGD.
{"src": "scitldr-AIC_train", "paper_id": "rJma2bZCW"}