Probabilistic Stability of Stochastic Gradient Descent

Thank you for the detailed and constructive feedback. We answer both the weaknesses and questions below.

Weaknesses:

1. It seems the analysis is not quite rigorous, and I found several gaps in the theoretical analysis.

Thanks for the criticism. We address the concerns below and have clarified these points in the revision. In short, we believe that most of the criticisms are due to misunderstanding of the definitions or of the problem setting.

2. The Lyapunov exponent seems to be a conservative parameter as it needs to take the maximum over all initializations. It seems that this quantity cannot fully illustrate the behavior of SGD since the initialization also has a large impact on the behavior of the algorithm, which cannot be explained by the Lyapunov exponent.

Thank you for the criticism. The Lyapunov exponent is actually more versatile than what we have focused on studying in the manuscript, and it is easy to extend the framework to study the effect of initialization. What we have called the "Lyapunov exponent" is more technically called the "largest Lyapunov exponent." More broadly, the Lyapunov exponent is defined with respect to a specific initialization. If the dynamical variable is $d$-dimensional, a famous result proves that the Lyapunov exponent takes on at most $d$ distinctive values. Therefore, an important future direction is to study the effect of initialization on the Lyapunov exponent and this is not a limitation of the proposed theoretical framework. Also, see footnote 5 in our manuscript for this point.

Questions:

1. Above Eq (16), the paper shows that algorithmic stability holds if $\lambda = 1/x_i^2$. Then, Eq (16) says that the largest stable learning rate is $1/x_\text{min}^2$. Should the largest stable learning rate be $1/x_\text{max}^2$ since we need the algorithm to be stable for any chosen example, and therefore should choose the smallest step size?

Thanks for this question. While this point might be surprising to some readers, we point out that, it holds true that for any $i$, if $\lambda = 1/x_i^2$, then the algorithm is globally stable, meaning that running SGD converges to the global minimum in probability.

When $\lambda=1/x_{min}$, it might appear unlikely that SGD is stable because it is unstable for the rest of the data points. However, in this very intriguing example, SGD being stable only for a single sample is sufficient to prevent its divergence and guarantee a probabilistic convergence to the global minimum (and this demonstrates the power of the probabilistic stability).

2. Above Eq (37), the paper shows that $m=tE_x[\log|1-\lambda h_t|]$. However, this identity only holds if $h_1=h_2=\ldots=h_t$. In this case, the matrix $\hat{H}(x)$ should remain the same over the optimization process. This seems to be a strong requirement.

Thank you for the comment. This is a misunderstanding. $m$ is an auxiliary constant, which we defined as the expected value. Therefore, there is no requirement or assumption here. It is just a definition we used to simplify the proof. The purpose of introducing $m$ is to define $z_t$ as a zero-mean random variable.

3. I cannot see how the argument below Eq (39) holds. That is, how to get the convergence by the law of large numbers?

It follows from the fact that $z_t$ is the sum of independent random variables, and so $z_t / t$ is the sample average of the random variable $\log|1 - \lambda h_t| - \mathbb{E}[\log|1 - \lambda h_t|]$. This random variable has zero mean and is identically distributed across $t$. As we always approximate $h$ using the Hessian at the stationary point, there is no dependence of $h$ on $\theta_t$. As a consequence, $h_t$ are independent and so does $\log|1 - \lambda h_t| - \mathbb{E}[\log|1 - \lambda h_t|$. Thus, the law of large numbers applies. The RHS can be obtained taking the limit of $t\to\infty$, and no law of large numbers is required.

Thank you for the detailed and constructive feedback. We answer both the weaknesses and questions below.

Weaknesses:

1. The analysis in this work is restricted to the assumption that initialization point is near a given stationary point. While I acknowledge that such a limitation is not unique to this paper, it deviates from real-world scenarios. Numerous observations have indicated that the early phase of learning substantially impacts the ultimate generalization performance of deep neural networks. Consequently, centering on the dynamics around stationary points may fall short in elucidating the success of deep learning.

Thanks for this criticism. This is a very important point that we want to clarify. We would like to point out that our theoretical analysis is fully compatible with the observation that the weights do not converge to stationary points but to stationary distributions. We have clarified this point in this revision. See our answer to Question 1 below.

2. In addition, this paper is poorly written and it needs substantial improvement.

Thanks for this criticism. We have now fully updated the manuscript to incorporate the feedback from all the reviewers.

Questions: Major concerns:

1. A recent study [1] shows that neural networks trained by gradient-based methods may not necessarily converge to stationary points. In other words, the gradients (norm) may not even vanish when these networks exhibit satisfactory performance. This observation raises questions about the applicability of probabilistic stability in understanding deep learning. [1] Jingzhao Zhang, et al. "Neural network weights do not converge to stationary points: An invariant measure perspective." ICML 2022.

Thank you for the criticism and this interesting reference. First of all, we would like to point out that Ref. [1] leaves the possibility open for models to converge to stationary points. For example, Ref. [1] states clearly in section A-1 that “the study of overparametrization and convergenece to stationary point may still be true in many cases.” Thus, the analysis in this paper is likely applicable for the smaller models that are able to overfit the data.

Secondly, our analysis can still apply when the model converges to a distribution, but not to a stationary point. The reason is that it is fully possible (and quite common) for SGD to converge to a point in some subspace, and at the same time converge to a nontrivial distribution overall. As a simple example to illustrate our point, let us consider a linear regression problem trained with SGD, with the per-batch loss: $l(w)=(w^Tx -y)^2$. Suppose $x$ takes a low-rank structure such that with probability one, $x$ is orthogonal to a vector $n$. Then, one can show that with probability one, $n^T \nabla_w l(w) =0$. Namely, in the subspace specified by $n$, the model parameters stay at a single point during training (and thus its series converges to a point in all common notions of stability). In contrast, under common assumptions on the learning rate and data distribution, the dynamics of $w$ in the subspaces orthogonal to $n$ do not converge to a point but to a distribution. Another more advanced example is the phenomenon of the loss of plasticity, common in continual learning. Here, more and more neurons become dead (incoming and outgoing weights becoming zero gradients once and for all) as the training proceeds. In these tasks, the model never converges to fixed points, yet a subset of the parameters do converge to stationary points (for example, see https://arxiv.org/abs/2303.07507).

The reason for this is that the gradient covariance matrix of SGD in deep learning is highly structured and contains many degenerate directions. We have now included a numerical example of matrix factorization to show this: the overall gradient norm does not converge to zero, while the projection of the gradient norm in a subspace does converge to zero -- therefore, in this given subspace, SGD has converged to a stationary point. See Section A5 for the experiment. We have also added these references and discussion to the main text.

Thank you for the detailed and constructive feedback. We answer both the weaknesses and questions below.

Weaknesses:

1. Studying dynamics of SGD using Lyapunov-type criteria is hardly new, so the theoretical contributions here are particularly limited in their novelty. These conditions for ergodicity of random linear recurrence relations (which immediately imply probabilistic stability as stated here) are already well known [1]. Such conditions have been considered for stochastic optimizers in the ML literature as well [2,3].

Thank you for pointing this out and for providing the references. We believe this criticism is a misunderstanding of the contribution of our work. First of all, let us clarify that we did not claim to be the first to use the Lyapunov-type criteria to study SGD. We have updated the manuscript to clarify this point.

Our main contribution is to study the stability of saddle points in a neural network landscape -- although the analysis can also be extended to study local minima, this is not the emphasis of our work and so its discussion is only presented in the manuscript. The most important contribution of our work is to show both empirically and theoretically that it is both common and possible for SGD to converge to constrained saddle points -- and a crucial theoretical insight is that this type of convergence can ONLY be studied within the framework of probabilistic stability. We would like to emphasize the second point more: we did not just apply probabilistic stability or Lyapunov-type analysis to these saddles, but we also proved that this is the only correct notion, which is established by proposition 2. Neither of these contributions appear in [2,3] and do not follow trivially from any result in [1].

That being said, these references do relate to ours. However, compared to [2] and [3] and as explained above, our work has a different context and emphasis. Both [2] and [3] focus on studying the dynamics of SGD and its ergodicity near a local minimum. In contrast, the focus of our work is the attractivity and/or repulsiveness of the saddles. Thus, the findings in these papers do not apply immediately to the case we are interested in, that is, the linear stability near the saddle points.

We have cited these works and added discussion regarding their relation to our work.

[1] Diaconis, P., & Freedman, D. (1999). Iterated random functions. SIAM Review, 41(1), 45-76.

[2] Gurbuzbalaban, M., Simsekli, U., & Zhu, L. (2021). The heavy-tail phenomenon in SGD. In International Conference on Machine Learning (pp. 3964-3975). PMLR.

[3] Hodgkinson, L., & Mahoney, M. (2021). Multiplicative noise and heavy tails in stochastic optimization. In International Conference on Machine Learning (pp. 4262-4274). PMLR.

2. The presentation of the phase diagrams is less than ideal, especially since this is perhaps the key contribution of the paper. Figures are confusing as presented: at the very least, they should be closer to where they are described in text, and empirical results should be compared more directly with theoretical findings. Terms are introduced here that do not appear to be explained.

Thanks for this comment. We have improved the readability of the figures and clarified their connections to the main text. Also, see our answers to the questions below.

That being said, while the phase diagrams are indeed a key contribution of our work, we stress that it is not the only crucial contribution we have. As we pointed out in the answer to Weakness 1, the first (and perhaps more important) contribution of our work is to show that the saddle points in deep learning have to be studied within the framework of probabilistic stability -- here, proposition 2 plays a key role as it proves that no other existing stability definition can be used to analyze the attractivity of these saddles.

3. Minor: a few typos throughout, e.g. eqn 6 missing an O in front of the cubic term and under eqn 10.

Thanks for pointing this out. We have fixed this typo in the revision.

Thank you for the detailed and constructive feedback. We answer both the weaknesses and questions below.

Weaknesses:

1. Literature on the variance-based stability notion is not adequately discussed in this paper. The paper presents the definition of the variance-based stability notion in Definition 2, but it is unclear what the current results are regarding this type of stability. Additionally, the reference for this stability notion cannot be found in this paper. It would be beneficial to include more discussions regarding the related work.

Thank you for the detailed comments. We believe that your main misunderstanding comes from our insufficient discussion of why the notion of stability is important.

In the context of deep learning, the main purpose for both the variance-based stability (previous works) and probabilistic stability (this work) is to understand which class of solutions or stationary points of the loss function is preferred by the SGD-based training. The most well-known result of the variance-based stability is that at a large learning rate or small batch-size, SGD prefers flat local minima because sharp minima are unstable, and, thus, repulsive, and so SGD will eventually escape sharp minima because of unstable oscillation and fluctuation (for example, see Wu et al. (2018)). However, a major limitation of the variance-based stability is that it cannot be applied to understand the stability of saddle points -- which is the main topic of our work.

We have added references to the variance-based stability and clarified these points. See the discussion below definition 2.

2. The clarity of the paper can be improved. It is difficult to follow and extract the key points of each section that the paper wants to deliver. It is also hard to connect each section. For example, Section 3.1 shows that rank-1 dynamics are solvable, but how this connects with probabilistic stability is unclear. Section 3.2 jumps to the point that variance-based stability is insufficient, and it is hard to connect these two sections.

Thank you for the comment. Please first see our answer to weakness 1 and note that the main focus of our work is this: why and when are saddle points attractive under SGD. It is to serve this purpose that the manuscript is organized.

While there has been various empirical evidence that saddle points are attractive, there has not been any theory to show and understand how and why SGD can converge to these saddle points in deep learning. Our goal is to fill this important theoretical gap. Section 3 solves rank-1 dynamics under the probabilistic stability -- the main message here is that saddle points can indeed become attractive under the probabilistic stability under many choices of hyperparameter (large learning rate, for example).

The next question is whether we have any other theoretical tool to understand the attractivity of saddles under SGD, and this is the purpose of Section 3.2. Section 3.2 shows that perhaps surprisingly, no notions of stability based on the statistical moments can help us understand the attractivity of saddles. This establishes the probabilistic stability as something quite unique -- because one must study the probabilistic stability to understand the attractivity of the saddles.

In summary, sections 3.1 and 3.2 together establish that probabilistic stability is especially suited and perhaps the ONLY tool for studying why the saddle points are attractive in deep learning.