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.
The paper investigates how the learning rate and mini-batch size in SGD impacts the optima that the SGD algorithm finds.
Empirically, the authors argue that it was observed that larger learning rates converge to minima which are more wide,
and that smaller learning rates more often lead to convergence to minima which are narrower, i.e. where the Hessian has large Eigenvalues. In this paper, the authors derive an analytical theory that aims at explaining this phenomenon.
Point of departure is an analytical theory proposed by Mandt et al., where SGD is analyzed in a continuous-time stochastic
formalism. In more detail, a stochastic differential equation is derived which mimicks the behavior of SGD. The advantage of
this theory is that under specific assumptions, analytic stationary distributions can be derived. While Mandt et al. focused
on the vicinity of a local optima, the authors of the present paper assumed white diagonal gradient noise, which allows to
derive an analytic, *global* stationary distribution (this is similar as in Langevin dynamics).
Then, the authors focus again on individual local optima and "integrate out" the stationary distribution around a local optimum, using again a Gaussian assumption. As a result, the authors obtain un-normalized probabilities of getting trapped in a given local optimum. This un-normalized probability depends on the strength of the value of the loss function in the vicinity of the optimum, the gradient noise, and the width of the optima. In the end, these un-normalized probabilities are taken as
probabilities that the SGD algorithm will be trapped around the given optimum in finite time.
Overall assessment:
I find the analytical results of the paper very original and interesting. The experimental part has some weaknesses. The paper could be drastically improved when focusing on the experimental part.
Detailed comments:
Regarding the analytical part, I think this is all very nice and original. However, I have some comments/requests:
1. Since the authors focus around Gaussian regions around the local minima, perhaps the diagonal white noise assumption could be weakened. This is again the multivariate Ornstein-Uhlenbeck setup examined in Mandt et al., and probably possesses an analytical solution for the un-normalized probabilities (even if the noise is multivariate Gaussian). Would the authors to consider generalizing the proof for the camera-ready version perhaps?
2. It would be nice to sketch the proof of theorem 2 in the main paper, rather than to just refer to the appendix. In my opinion, the theorem results from a beautiful and instructive calculation that should provide the reader with some intuition.
3. Would the authors comment on the underlying theoretical assumptions a bit more? In particular, the stationary distribution predicted by the Ornstein-Uhlenbeck formalism is never reached in practice. When using SGD in practice, one is in the initial mode-seeking phase. So, why is it a reasonable assumption to still use results obtained from the stationary (equilibrated) distribution which is never reached?
Regarding the experiments: here I see a few problems. First, the writing style drops in quality. Second, figures 2 and 3 are cryptic. Why do the authors focus on two manually selected optima? In which sense is this statistically significant? How often were the experiments repeated? The figures are furthermore hard to read. I would recommend overhauling the entire experiments section.
Details:
- Typo in Figure 2: ”with different with different”.
- “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.” This is clearly wrong as there are many local minima, and running teh algorithm twice results in different local optima.  Maybe add something that this only true on average, like “the characteristics of these minima ... should be the same”.