How SGD Selects the Global Minima in Over-parameterized Learning: A Dynamical Stability Perspective
The question of which global minima are accessible by a stochastic gradient decent (SGD) algorithm with specific learning rate and batch size is studied from the perspective of dynamical stability. The concept of non-uniformity is introduced, which, together with sharpness, characterizes the stability property of a global minimum and hence the accessibility of a particular SGD algorithm to that global minimum. In particular, this analysis shows that learning rate and batch size play different roles in minima selection. Extensive empirical results seem to correlate well with the theoretical findings and provide further support to these claims.
This paper examines, theoretically and empirically, how does SGD select the global minima it converges to. It first defines two properties ("sharpness" and "non-uniformity") of a fixed point, and how these determine, together with batch size, the maximal learning rate in which the fixed point is stable under SGD dynamics (both in mean and variance). It is then demonstrated numerically how these results relate affect the learning rate and batch size affect the selection of minima, and the dynamics of "escape" from sharp minima".   Clarity: This paper is nicely written, and quite clear. Quality: Seems correct, except some fixable errors (see below), and the numerical results seem reasonably convincing.  Originality: The results are novel to the best of my knowledge. Significance: The results shed light on the connections between sharpness, learning rate, batch size, and highlight the importance of "non-uniformity". These connections are not well understood and have received attention since the recent work by Keskar et al. From the empirical side, I especially found it interesting that GD selects solutions with maximal sharpness and the strong correlation between sharpness and non-uniformity. I also agree with the authors that the fast escape from sharp minima (e.g., Figure 7, left column) favors this paper's "instability" explanation over the SDE interpretations (e.g., Jastrzebski et al.).  Main Comments: 1) The linearization x_{t+1}=x_t - \eta A_{\xi_t}x_t is wrong: the zeroth order term G(0;\xi) is not necessarily zero (x^* is only a fixed point of G(x), not G(x;\xi)). In other words we need to keep constant term  G(0;\xi) in the dynamics.  These constants do not affect the dynamics of E[x_t] (since their mean is zero), but they do affect E[x^2_t] . This implies that the variance generally does not go to zero. An indeed, SGD generally doesn't converge to a fixed point without decreasing the learning rate to zero.  A way to fix this would be to assume we converge to a global minima where \nabla G(x;\xi) for all \xi . This seems reasonable in overparameterized neural networks.  2) The numerical results focus on quadratic loss, since the authors claim that the Hessian vanishes for cross-entropy. This indeed happens, but only for minima at infinity. However, if we use cross-entropy with L2 regularization (as is commonly done), then all minima are finite and the Hessian would not vanish.  Therefore, I'm not sure I understand this argument. Do the reported results also hold for cross-entropy with L2 regularization?  Minor Comments: 1) I think eqs. 13 and 14 would be more helpful before eq. 8 (with A_i=H_i) than in their current location. 2) It would be helpful to see the derivation of eq. 6 in the appendix.  3) Figure 3 could use larger fonts (especially in labels) 4) Figure 6(a): "2/\eta" should not be there, right? 5) I recommend the authors to make their code available online to enable other to easily reproduce their results.  Some Typos: line 217 - "the extend" line 224 - "is easily" line 292 - "generlizes" equation below line 314 - equality should be \leftarrow (since both sides have v^k).  %% After author feedback %% I thank the authors for their reply. Some futher comments to improve the quality of the paper:  First, I suggest the authors add complete proofs (in the SM), and the main Theorem they mention in the rebuttal.  Second, please include citation and discussion related to the works mentioned by Reviewer 3. For example, with regard to Ma et al., which seems most relevant (the results in the first two papers assume strong convexity, which makes them generally inapplicable to over-parameterized models), I think the main differences are: 1) Ma et al. focused on how to select the learning rate and batch size to achieve an optimal convergence rate. In contrast, this paper focuses on how the batch size and learning rate determine when do we get dynamical instability and SGD escapes away from the minimum.  2) I think the resulting expressions are rather different. For example, the bounds in Ma et al. (e.g. Theorem 2) depend on the eigenvalues of the Hessian, while in this paper they depend on the maximal eigenvalues of H and Sigma. The latter expressions expose nicely the important factors affecting stability (sharpness and non-uniformity). Also, (and this is more minor) in Ma et al. the derived upper bound is for the expected loss, while here the bound is the first two moments of the optimized parameters. 3) I believe the results in this paper are more easily generalizable (to realistic models) since we only care about dynamical stability near the fixed point, and not the optimal learning rate/batch size (Which depend on the entire optimization path).  4) The derivation in this paper is quite simple and intuitive (at least for me).  Third, the definitions of non-uniformity vary somewhat in the paper. Specifically: a. line 24 (intro): "the variance of the Hessian (see below for details)" b. line 66 (1D case): non-uniformity = s = std of the (1D) Hessian c. lines 90-91 (General case, matrix version): \Sigma (variance of Hessian matrix) is a "measure of the non-uniformity"  d. lines 97-98 (General case, scalar version): non-uniformity = s, where s^2 is the largest eigenvalue of \Sigma e. Appendix C (what they used in practice): same as the previous definition (up to the approximation accuracy of the power iteration method). So definitions (b),(d), and (e) are consistent, and I think it is pretty clear that this what they actually mean by "non-uniformity".  (c) is different, but they only say this is a "measure of .." and not the definition. I think this matrix version was mainly pointed out to give some intuition. (a) indeed seems to be inaccurate, and should be corrected, but I think it is a small mistake. I think these issues should be corrected, and a formal definition should be given, as non-uniformity is one of the key ingredients of this paper.  Fourth, I think the authors should try to emphasize and clarify more the important parts in their experiments. I'll try to summarize the interesting observations I found in the experiments: 1) The sharpness value is nearly always near the maximum value (2/eta) for GD (Table 2, Figure 6a); Non-uniformity also has values with a similar magnitude as the bound (Figure 6a). 2) The linear relation between sharpness and non-uniformity (Figure 6) 3) Observations (1-2) above (as well as figure 3), suggest that smaller mini-batches decrease sharpness due to the non-uniformity bound (Figure 6). 4) Non-uniformity seems to predict well when we escape from sharp minima (Figure 7). 5) The escape from sharp minima is fast (Figure 7). I think this is quite a lot of new interesting information that I've learned here. Previous high impact papers suggested that * Keskar et al.: (a) "large batch size implies large sharpness"  (b) "small sharpness is important for generalization" * Hoffer et al., Goyal et al., Smith and Le: (c) "changing the learning rate at large batch sizes can improve generalization" However, so far it was not clear why these statements should be true (e.g. counterexample in Appendix A for the statement b). I think the observations (1-4) above make significant progress in understanding statements (a-c), and suggests we should also focus on non-uniformity (instead of just sharpness). Also, observation 5 suggests that previous SDE models for this phenomena are wrong. Other previous explanations (e.g., Yin et al.) seems not be applicable to over-parameterized case, as explained in the discussion. In terms of math, I am not aware of other works that examined this specific over-parameterized case, where all sample gradient are zero at the minimum, except for Ma et al. (which focused on different questions).   Fifth, I think that there is an issue in the conditions relying only on the maximal eigenvalues. For example, in at line 98 the author state that "a sufficient condition for (9) to hold is eq. (10)", and condition (10) only depends on the maximum eigenvalue of H.  However, I don't see why small (possibly even zero) eigenvalues of H can't affect stability. For example, suppose that B=1 and \eta=1, that v is the eigenvector corresponding the max eigenvalue of (9), and also that H*v=c*v, \Sigma*v = q*v for two positive scalars c and q.  In this case, we get an eigenvalue of the matrix in (9) which is equal to (1-c)^2+q. This value can be larger than 1, if c is sufficiently small (e.g. c=0), relatively to q. The authors should explain how this issue could be avoided, or correct their statement. In any case, I recommend adding more mathematical details for all calculations in the supplementary.