Paper: Probabilistic Stability of Stochastic Gradient Descent (ICLR-2024)

Reviewer comment (Theory):
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.

Author response:
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.

Reviewer comment (Theory):
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_{min}^2$. Should the largest stable learning rate be $1/x_{max}^2$ since we need the algorithm to be stable for any chosen example, and therefore should choose the smallest step size?

Author response:
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).

Reviewer comment (Theory):
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.

Author response:
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.

Reviewer comment (Theory):
In Eq (46), the paper uses the identity $E[\hat{H}_t\theta_t]=E[\hat{H}_t]E[\theta_t]$. This identity holds if $\hat{H}_t$ and $\theta_t$ are independent. This also seems to be a restrictive condition.

Author response:
This is not a restrictive condition. Within the framework of linear stability, the dynamics only depends linearly on $\theta$. Therefore, in linear stability, $\hat{H}_t$ is never dependent on $\theta_t$. This assumption of linear stability holds when the model parameters are sufficiently close to a stationary point $\theta^*$, and so it is equivalent to assuming that $ \hat{H}_t(\theta)\approx \hat{H}_t(\theta^*)$ for a constant (or approximately constant) $\theta^*$.

That being said, the only way for $\hat{H}_t$ and $\theta_t$ to be dependent is that the random sampling of data points is dependent -- which cannot happen by the standard definition of minibatch sampling. Namely, under SGD, $\theta_t$ depends on $\{h_s, s < t\}$. Therefore, as $h_t$ are i.i.d. as explained in the answer to the previous question, there is no dependence between $\theta_t$ and $h_t$.

As an example, please see the problem we studied in proposition 3. It is clear that the $h_t$ matrices are independent in this case. We have clarified this point in the manuscript.

Reviewer comment (Theory):
Theorem 3 requires $X_i$ to be independent random matrices. The paper applies Theorem 3 to get Eq (54), which requires $\hat{H}_i$ to be independent. However, it seems that these matrices are not independent, and therefore the Theorem 3 cannot be applied?

Author response:
Please first see our answers above. As explained in the answer to Question 4, the Hessian only depends on the sampled data points which are independent.

Reviewer comment (Theory):
I cannot see how Eq (52) holds. In particular, how this identity holds with a small O notation.

Author response:
Eq (52) works in the same way as Eq (39). The little $o$ comes from the fact that $\log \epsilon/t$ vanishes at large $t$. The rest of the proof comes from the independence of sampling data points by the SGD algorithm.

Reviewer comment (Writing):
Typos:  
Eq (6): $\lambda^3$ should be $O(\lambda^3)$  
Below Eq (10): ", This" should be ", this"  
Proposition 2: "but $L_p$ stable" should be ""but not $L_p$ stable""?

Author response:
Thank you for pointing them out. We have fixed the typos.

Reviewer comment (Writing):
This paper is poorly written and it needs substantial improvement.

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

Reviewer comment (Writing):
On Page 3, in the text beneath the caption of Figure 1, the statement "This means that Eq. (1) can be seen as an effective description for only a small subset of all parameters ..." lacks clarity in its flow. Could you provide further elaboration on this?

Author response:
Thanks for this question. Please first see our answer to question 1. Here, the discussion refers to the fact that the dynamics of SGD in many subspaces are quite independent of the dynamics of the rest of the parameters so one can independently study the dynamics in this subspace.

Reviewer comment (Theory):
On Page 6, in the fifth-to-last line, you mention, "The theory also shows that if we fix the learning rate and noise level, increasing the batch size makes it more and more difficult to converge to the low-rank solution...". However, in Proposition 3, it's not clear to me how the batch size explicitly factors into this observation. Could you elaborate more on the role of batch size in this context?

Author response:
Thanks for asking for clarification. The simplest way to see this is to note that using a large batch size leads to smaller gradient noise -- so increasing the batch size leads to a small variance in the gradient. Our original manuscript was implicit about this point. We have now updated proposition 3 and its related discussion to make the dependence on batch size explicit.

Reviewer comment (Writing):
Notations are inconsistent in this paper. For example, $n$ serves as the representation of a norm type in the first paragraph on Page 2, where it is a scalar. Later in Section 3.1, $n$ is used as a fixed unit vector. In addition, there is a disparity in the definition of the loss function between Eq.(2) and Eq.(12), as they take different inputs.

Author response:
Thank you for the comment. We now use p instead of n for the norm. And, we now use semicolons to separate the dependence of parameters and data in the loss function so that Eq.(2) and Eq.(12) can be written in a consistent way.

Reviewer comment (Writing):
The paragraph after Definition 1: "A sequence $z_t$ converges" $\Longrightarrow$ "A sequence $\theta_t$ converges"

Author response:
This is now fixed.

Reviewer comment (Novelty):
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].

Author response:
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.

Reviewer comment (Presentation):
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.

Author response:
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.

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

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

Reviewer comment (Evaluation):
Is sparsity/density the fraction of zero/non-zero elements in the matrix? This is my best guess, but it is surprising that there would be so many "zero" elements here; is there some larger cutoff which determines whether an element is "zero", or is something else considered here?

Author response:
Thanks for this question. The sparsity is indeed the fraction of zero elements in the matrix. First of all, we note that we do observe exact zeros in the experiments -- but this often takes quite long to achieve. More often, we set a threshold, and any quantity smaller than the threshold is considered zero. Typically, the threshold is set to be 1e-8, which is roughly the machine precision level of Pytorch.

Reviewer comment (Presentation):
Is Figure 1 comparable to any other empirical examples? If not, what is the purpose of this figure?

Author response:
Figure 1-left is an example of Proposition 3 (when N=2) studied in Section 4. The phase boundaries are drawn according to proposition 3 and its experimental details are given in A.1. Figure 1-left is the case when there are only two data points for proposition 3. For a larger N, the results are given in Figure 7. For infinite N, the result is shown in Figure 2. This has been clarified in the updated Section 4.

Also, Figure 1-right is more of an illustration of how the phase boundaries are decided by the Lyapunov exponent.

Reviewer comment (Presentation):
Am I to interpret Figures 4 (right) and 5 (right) as displaying the same behavior as Figure 2 (right)? Can the theoretical prediction be overlaid here too?

Author response:
Figure 4 and 5 indeed displays the same qualitative behavior as Figure 2. However, the order of magnitude for the learning rate in the figures is different.  Figure 3 is the multi-dimensional version of Figure 2, and the theory curve in Figure 3 corresponds to the dashed curve in Figure 2. 

Figure 4 and Figure 5, obtained using realistic data and ResNet, demonstrate the same qualitative behavior as Figure 3. However, we do not have a quantitative prediction for the precise phase boundary for ResNet because it is not tractable to compute its per-batch Hessians. In summary, the agreement between Figure 2 and 3 is quantitative, while the agreement between Figure 2 and Figures 4-5 is qualitative.

Reviewer comment (Presentation):
Why is there an arrow from A to B in Figure 6? The text suggests that initialized at B, SGD jumps to C.

Author response:
The red arrow from A to B is the prediction by L2 stability: when the learning rate is large, even if initialized at A, SGD will be unstable there and must jump to an alternative solution -- and this certainly is not what happens in reality. The text referring to this is located in the middle paragraph in section 5. We have added a description for this in the caption of Figure 6 to avoid misunderstanding.

Reviewer comment (Evaluation):
Is Figure 7 converging to Figure 2 as \(N \to \infty\)?

Author response:
Yes.

Reviewer comment (Theory):
I assume convergence to the low-rank saddle should correlate with poor performance?

Author response:
Not necessarily. At its face value, low-rankness means only that the capacity of the model is limited. This could be a bad thing when a high model capacity is desired. This could be good if overfitting is a concern. Both perspectives appear frequently in the context of machine learning and deep learning.

Reviewer comment (Presentation):
Can you outline the phases in the main text? These need to be displayed front and center. It is very confusing to have to refer to a figure legend towards the top of the paper for this.

Author response:
Thanks for this suggestion. We have updated Section 4 to clearly outline the phases.

Reviewer comment (Presentation):
Can you put the SGD solution selection part (including Figure 6) near Section 3.2? Otherwise, the purpose of this section is lost on first read.

Author response:
Thank you for the proposal. However, we prefer to separate the theory and its implications, so the organization of the contents remains roughly the same. To help readers who want to see an immediate example after Section 3.2, we have added a reference to section 5 as well as Fig. 6 in section 3.2 so that the purpose of 3.2 is not lost.

Reviewer comment (Novelty):
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.

Author response:
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.

Reviewer comment (Presentation):
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.

Author response:
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.

Reviewer comment (Theory):
The technical soundness and significance can be improved. Section 3 shows that the linearized dynamics of SGD converges with probability under certain conditions, but it is difficult to establish a connection between this result and how it explains the generalization of SGD in deep learning. Section 4 discusses different phases of SGD learning, but it is unclear how these phases relate to the stability notions proposed in the paper and how they explain the generalization of SGD.

Author response:
Thanks for this important question. To explain and analyze the practical good generalization of SGD, two crucial questions need to be answered in order: which solution does SGD prefer; what is the nature of these solutions (in terms of its generalization power). The main focus of our analysis is on the first question. Our key contribution here is to show that saddle points that feature special, constrained solutions are preferred when the learning rate is large or the gradient noise is strong.

In contrast, the second question is only answered in a qualitative manner in our work. Namely, we show that the learning can be divided into four different phases, and the generalization ability of these four phases is clear -- in the trivial phase, the model generalizes like a trivial model that only outputs a constant; in the correct learning phase, the model generalizes like a model that has reached the global minimum (to which the conventional analysis that links the training error and generalization error can be applied). In the incorrect phase, the model learns the spurious correlations and overfits to noises in the label. 

How are these phases related to probabilistic stability? In the example of Section 4, the trivial phase is when both $u-w$ and $u+w$ are probabilistically stable at $0$. The correct learning phase is when $u-w$ is prob. stable but $u+w$ is p-unstable (at zero). The incorrect phase is when $u+w$ is prob. stable and $u-w$ is p-unstable. The unstable phase is when both quantities are p-unstable. We have made this point clear in the revision (see Section 4). 

That being said, the full quantitative analysis of generalization is, however, beyond the scope of this work.

Reviewer comment (Theory):
Can we understand that the probabilistic stability defined in this paper is more like a convergence guarantee?

Author response:
Not quite. A stable solution is an attractive solution. At the end of training, SGD can only stay close (and sometimes converge) to a solution that is stable. It can also be seen as a convergence guarantee IF one is initialized close enough to such a solution -- that being that a solution is stable does not imply that starting from an arbitrary initialization, one will converge to such a solution. For example, for a small learning rate, all global minima are stable, but SGD is more likely to converge to a minimum that is close to its initialization.

Reviewer comment (Novelty):
What is the literature on variance-based stability?

Author response:
In the context of deep learning, the main purpose for both the variance-based stability (previous works) 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. Also, see our answers above. 

There are other works that are based on the variance-based stability, for example, Wu & Su (2023), Ma & Ying (2021), Mulayoff et al. (2021). These works study the implicit regularization of SGD by linking the desired properties of solutions to the linear stability of SGD.

Reviewer comment (Writing):
In Definition 1 and 2, it would be great to clarify whether θ* is fixed or a random variable. After Definition 1, there is a typo in the convergence in probability, "< ε" should be revised to "> ε."

Author response:
The $\theta^*$ is fixed. It is a stationary point in the loss landscape. Thank you for the observation. We have fixed the typo.

Reviewer comment (Writing):
After equation (2): "The dynamics of w thus obeys Eq. (1)." Here, w is not defined.

Author response:
We have fixed this. It should be $\theta$.

Reviewer comment (Presentation):
Figure 1: It would be great to explain the x-axis and y-axis in the caption (the same in Figure 2). Also, what are w_t and u_t? How are the phases (different colors) calculated based on values on the x-axis and y-axis?

Author response:
Thank you for the question. We use $\lambda$ as the learning rate and $\Lambda$ as the Lyapunov exponent consistently throughout the manuscript, so we didn’t add further explanations for them in the caption due to page limitation. In Figure 1, we added a reference for parameter $a$. In Figure 2, we added explanations for the parameters $\mu$ and $S$. The phase diagram in Figure 1 is calculated based on the results in C.2. The phases are determined by the signs of equations (59) and (60). Please refer to the reply to the third point in Weaknesses for a detailed explanation of how to connect the convergence of $u + w$ and $u-w$ to the phases.

Reviewer comment (Presentation):
Figure 3: It would be great to explain how the color map is calculated and how convergence is calculated.

Author response:
The color map denotes the number of non-zero values of the singular value of matrix $u$ using the model specified in section 4, equation(12). We say the experiment converges if we observe plateaux in the rank-time plots. We have updated Figure 3 to clarify this.