A BAYESIAN PERSPECTIVE ON GENERALIZATION AND STOCHASTIC GRADIENT DESCENT
We consider two questions at the heart of machine learning; how can we predict if a minimum will generalize to the test set, and why does stochastic gradient descent find minima that generalize well? Our work responds to Zhang et al. (2016), who showed deep neural networks can easily memorize randomly labeled training data, despite generalizing well on real labels of the same inputs. We show that the same phenomenon occurs in small linear models. These observations are explained by the Bayesian evidence, which penalizes sharp minima but is invariant to model parameterization. We also demonstrate that, when one holds the learning rate fixed, there is an optimum batch size which maximizes the test set accuracy. We propose that the noise introduced by small mini-batches drives the parameters towards minima whose evidence is large. Interpreting stochastic gradient descent as a stochastic differential equation, we identify the "noise scale" g = (
where is the learning rate, N the training set size and B the batch size. Consequently the optimum batch size is proportional to both the learning rate and the size of the training set, B opt ∝ N . We verify these predictions empirically.
Summary:
This paper presents a very interesting perspective on why deep neural networks may generalize well, in spite of their high capacity (Zhang et al, 2017). It does so from the perspective of "Bayesian model comparison", where two models are compared based on their "marginal likelihood" (aka, their "evidence" --- the expected probability of the training data under the model, when parameters are drawn from the prior).  It first shows that a simple weakly regularized (linear) logistic regression model over 200 dimensional data can perfectly memorize a random training set with 200 points, while also generalizing well when the class labels are not random (eg, when a simple linear model explains the class labels); this provides a much simpler example of a model generalizing well in spite of high capacity, relative to the experiments presented by Zhang et al (2017). It shows that in this very simple setting, the "evidence" of a model correlates well with the test accuracy, and thus could explain this phenomena (evidence is low for model trained on random data, but high for model trained on real data).
The paper goes on to show that if the evidence is approximated using a second order Taylor expansion of the cost function around a minimia $w_0$, then the evidence is controlled by the cost at the minimum, and by the logarithm of the ratio of the curvature at the minimum compared to the regularization constant (eg, standard deviation of gaussian prior).  Thus, Bayesian evidence prefers minima that are both deep and broad.  This provides a way of comparing models in a way which is independent of the model parametrization (unfortunately, however, computing the evidence is intractable for large networks). The paper then discusses how SGD can be seen as an algorithmic way of finding minima with large "evidence" --- the "noise" in the gradient estimation helps the model avoid "sharp" minima, while the gradient helps the model find "deep" minima.  The paper shows that SGD can be understood using stochastic differential equations, where the noise scale is approximately aN/((1-m)B) (a = learning rate, N = size of training set, B = batch size, m = momentum).  It argues that because there should be an optimal noise scale (which maximizes test performance), the batch size should be taken proportional to the learning rate, as well as the training set size, and proportional to 1/(1-m).  These scaling rules are confirmed experimentally (DNN trained on MNIST).  Thus, this Bayesian perspective can also help explain the observation that models trained with smaller batch sizes (noisier gradient estimates) often generalize better than those with larger batch sizes (Kesker et al, 2016). These scaling rules provide guidance on how to increase the batch size, which is desirable for increasing the parralelism of SGD training.
Review:
Quality: The quality of the work is high.  Experiments and analysis are both presented clearly.
Clarity: The paper is relatively clear, though some of the connections between the different parts of the paper felt unclear to me:
1) It would be nice if the paper were to explain, from a theoretical perspective, why large evidence should correspond to better generalization, or provide an overview of the work which has shown this (eg, Rissanen, 1983).
2) Could margin-based generalization bounds explain the superior generalization performance of the linear model trained on random vs. non-random data?  It seems to me that the model trained on meaningful data should have a larger margin.
3) The connection between the work on Bayesian evidence, and the work on SGD, felt very informal. The link seems to be purely intuitive (SGD should converge to minima with high evidence, because its updates are noisy).  Can this be formalized?  There is a footnote on page 7 regarding Bayesian posterior sampling -- I think this should be brought into the body of the paper, and explained in more detail.
4) The paper does not give any background on stochastic differential equations, and why there should be an optimal noise scale 'g', which remains constant during the stochastic process, for converging to a minima with high evidene.  Are there any theoretical results which can be leveraged from the stochastic processes literature? For example, are there results which prove anything regarding the convergence of a stochastic process under different amounts of noise?
5) It was unclear to me why momentum was used in the MNIST experiments.  This seems to complicate the experimental setting.  Does the generalization gap not appear when no momentum is used?  Also, why is the same learning rate used for both small and large batch training for Figures 3 and 4?  If the learning rate were optimized together with batch size (eg, keeping aN/B constant), would the generalization gap still appear?  Figure 5a seems to suggest that it would not appear (peaks appear to all have the same test accuracy).
6) It was unclear to me whether the analysis of SGD as a stochastic differential equation with noise scale aN/((1-m)B) was a contribution of this paper.  It would be good if it were made clearer which part of the mathematical analysis in sections 2 and 5 are original.
7) Some small feedback: The notation $< x_i > = 0$ and $< x_i^2 > = 1$ is not explained.  Is each feature being normalized to be zero mean, unit variance, or is each training example being normalized?
Originality: The works seems to be relatively original combination of ideas from Bayesian evidence, to deep neural network research.  However, I am not familiar enough with the literature on Bayesian evidence, or the literature on sharp/broad minima, and their generalization properties, to be able to confidently say how original this work is.
Significance: I believe that this work is quite significant in two different ways:
1) "Bayesian evidence" provides a nice way of understanding why neural nets might generalize well, which could lead to further theoretical contributions.
2) The scaling rules described in section 5 could help practitioners use much larger batch sizes during training, by simultaneously increasing the learning rate, the training set size, and/or the momentum parameter.  This could help parallelize neural network training considerably.
Some things which could limit the significance of the work:
1) The paper does not provide a way of measuring the (approximate) evidence of a model.  It simply says it is prohibitively expensive to compute for large models.  Can the "Gaussian approximation" to the evidence (equation 10) be approximated efficiently for large neural networks?
2) The paper does not prove that SGD converges to models of high evidence, or formally relate the noise scale 'g' to the quality of the converged model, or relate the evidence of the model to its generalization performance.
Overall, I feel the strengths of the paper outweight its weaknesses.  I think that the paper would be made stronger and clearer if the questions I raised above are addressed prior to publication.