Paper: A Precise Characterization of SGD Stability Using Loss Surface Geometry (ICLR-2024)

Reviewer comment (Theory):
The authors might make a better interpretation of Theorem 1. What does this result imply about the implicit regularization of SGD (beyond that of GD)? How stability is related to the hyperparameters, e.g. $\eta$, $B$, and alignment of $H_i$? Some relevant discussion can be found in the experiment part Section 5.2, but I think it would be better to provide some intuition right after Theorem 1.

Author response:
We agree it is best to immediately explain the key takeaways after the theorem statement.

We discuss how our analysis implies a "squared scaling rule" between $B$ and $\eta$ after Lemma 4.1, but we now also mention this after Theorem 1. Another point to note regarding how SGD is regularized beyond GD is that GD does not depend on the coherence measure $\sigma$. We give examples in the paragraph above Theorem 1 demonstrating the importance of the Hessian alignment (captured by $\sigma$) in characterizing SGD stability. We now emphasize this more explicitly after the theorem statement.

Reviewer comment (Theory):
Theorem 2, the optimality of Theorem 1, strongly relies on the condition that $\sigma, n/B = O(1)$. There are two concerns: In Theorem 2 the authors assume that $\sigma \in [n]$. Is there anything to guarantee $\sigma \leq n$? I don't think it is clear from the definition of $\sigma$. Is $\sigma$ inherently bounded? Moreover, do you mean $\sigma\leq n$? Hence, it is unclear what the assumption $\sigma=O(1)$ means. Also, I do not think it is natural to assume $n/B=O(1)$ as usually $B\ll n$.

Author response:
Theorem 2 is provided to give an idea of how sharp our current analysis is and what approaches would be needed to improve it. To clarify, we do not assume that $\sigma$ is a specific value, rather we state that there exists an input problem with Hessians $\\{H\_i\\}\_{i \in [n]}$ such that $\sigma$ equals this value and the linearized dynamics don't diverge. This means any improvement to Theorem 1 cannot imply that these problems diverge.

The discussion above Theorem 2 where we consider $\sigma, n/B = O(1)$ is just pointing out that the conditions of Theorem 1 and Theorem 2 are close in this parameter range, and so we cannot expect much improvement in this regime. However, Theorem 2 does not make any assumption regarding this. From a theoretical perspective, we think it is natural to consider this regime to understand its behavior. However, we agree that the case where $B \ll n$ is important to understand in future work, and so we have added a sentence pointing this out above Theorem 2.

Reviewer comment (Novelty):
In paragraph 1 of the introduction, when introducing the concept of "implicit bias", instead of citing (Li et al., 2022), I think it is more appropriate to cite the seminar works (Neyshabur et al., arXiv:1412.6614) and (Zhang et al., ICLR 2017).

Author response:
Thank you, we have included these references in the updated manuscript.

Reviewer comment (Novelty):
In paragraph 2 of the introduction, when citing empirical works on relating sharpness to generalization, I think the important comprehensive investigations by (Jiang et al., ICLR2020) is missed.

Author response:
Of course, this is a seminal paper that is worth referring to. Thank you for pointing this out.

Reviewer comment (Novelty):
In paragraph 4 of the introduction, when stating "GD ... empirically validated to predict the sharpness of overparameterized neural network", the author cites (Cohen et al., 2021). However, this empirical predictability of linear stability analysis has been observed in (Wu et al., NeurIPS2018).

Author response:
We apologize for missing this important reference in this place. We have referred to the same paper elsewhere in other contexts.

Reviewer comment (Theory):
In Section 2, paragraph 1, when stating the rationale for assuming over-parameterization, the authors cite the work (Allen-Zhu et al., 2019). This seems quite strange to me.

Author response:
We are happy to add additional references, e.g. Zhang et al 2016.

Reviewer comment (Writing):
In Section 2, Definition 1, it is unclear whether the sampling is done with or without replacement.

Author response:
We are using Bernoulli sampling which is not the same as sampling with or without replacement. The text under eqn. (1) in Definition 1 describes the probabilistic sampling model we work with. To clarify $\mathcal{S}$ is a set (so there are no repetitions) and $i \in \mathcal{S}$ being independent from $j \in \mathcal{S}$ implies it cannot be sampling without replacement.

Reviewer comment (Writing):
In Section 3, second paragraph, what is the $\frac{1}{B}\sum_{i=1}^n x_i H_i$?

Author response:
Thanks for pointing this typo out. We have changed it to $\frac{1}{B} \sum_{i \in \mathcal{S}} H_i$. We reformulate this later in the proof as $\frac{1}{B}\sum_{i=1}^n x_i H_i$ where the $x_i$ are Bernoulli random variables, but you are correct we have not defined it at this point.

Reviewer comment (Writing):
In Section 3, Theorem 1, what does the subscript in $\hat{J}_i$ stand for? Complexity measure => coherence measure.

Author response:
We have changed "Let $\hat{J}\_i$ be independent Jacobians of SGD dynamics" to "Let $\\{\hat{J}\_i\\}\_{i\in\mathbb{N}}$ be a sequence of i.i.d. copies of $\hat{J}$ defined in Definition 1" to make it clearer we consider an infinite sequence of i.i.d. samples of the stochastic Jacobian.

Reviewer comment (Writing):
In Section 3, Theorem 2, what do you mean $\sigma\in [n]$? Is $\sigma$ a real value?

Author response:
For an arbitrary set of Hessians $\{H_i\}_{i\in [n]}$, the coherence measure $\sigma$ can be a general real value. However, Theorem 2 is an existential result that says for any chosen $\sigma \in \{1,2,...,n\}$, there exists a set of Hessians $\{H_i\}$ satisfying the properties described in the theorem.

Reviewer comment (Theory):
The analysis relies on the Bernoulli sampling model, which, although is in expectation the same as with replacement sampling or uniformly sampling from all $B$-sets, still is a bit different as the size of $\mathcal{S}$ itself now becomes random. Have the authors tried to consider multinomial distribution as the sampling distribution?

Author response:
We are optimistic that improved understanding of the Bernoulli sampling model can provide useful insight to other sampling models, as was the case for the previous line of work on matrix completion by Candes, Tao, and others which culminated in Benjamin Recht's seminal paper [Recht, 2011]. We also note that sampling with/without replacement are also just theoretical approximations of the typical true sampling behavior, where the data is shuffled and iterated over in an epoch. In reality, the sampling is often not even independent between iterations.

Lemma 4.1 could likely be rewritten using the sampling without replacement formulation for $\hat{J}$ used in [Ma and Ying, 2021]. However, there doesn't seem to be an obvious way to simplify the expressions, and we would be stuck with a characterization that is no more interpretable than [Ma and Ying, 2021].

Another possibility would be to rewrite Lemma 4.1 for sampling with replacement. We expect that the coefficients depending on $\eta$ and $B$ would change, but the proof would generally stay the same. That is, let $x_1,...,x_B$ each be uniformly sampled from $[n]$. Then, $\mathbb{E}[H_{x_i}H_{x_j}] = \frac{1}{n}\sum_k {H_k^2}$ when $i = j$ and $\mathbb{E}[H_{x_i}H_{x_j}] = \frac{1}{n^2}H^2$ when $i \neq j$. Therefore, the proof of Lemma 4.1 would start the same way when $\hat{H}=\sum_{i=1}^B H_{x_i}$. Results for the sampling with replacement case would immediately imply asymptotic results for sampling without replacement, since the total variation distance between the two distributions goes to zero as $n \rightarrow \infty$ and $B/n$ converges to a constant.

Reviewer comment (Theory):
Moreover, without replacement sampling where the event $i \in \mathcal{S}$ is dependent on $j \in \mathcal{S}$ (depending on the order), there should be some theoretical discussions on the effect of using these two (most-widely-used) random sampling schemes. (I appreciate that the paper has experiments on this, but then, theorem-wise, what would possibly change?)

Author response:
We are optimistic that improved understanding of the Bernoulli sampling model can provide useful insight to other sampling models, as was the case for the previous line of work on matrix completion by Candes, Tao, and others which culminated in Benjamin Recht's seminal paper [Recht, 2011]. We also note that sampling with/without replacement are also just theoretical approximations of the typical true sampling behavior, where the data is shuffled and iterated over in an epoch. In reality, the sampling is often not even independent between iterations.

Lemma 4.1 could likely be rewritten using the sampling without replacement formulation for $\hat{J}$ used in [Ma and Ying, 2021]. However, there doesn't seem to be an obvious way to simplify the expressions, and we would be stuck with a characterization that is no more interpretable than [Ma and Ying, 2021].

Another possibility would be to rewrite Lemma 4.1 for sampling with replacement. We expect that the coefficients depending on $\eta$ and $B$ would change, but the proof would generally stay the same. That is, let $x_1,...,x_B$ each be uniformly sampled from $[n]$. Then, $\mathbb{E}[H_{x_i}H_{x_j}] = \frac{1}{n}\sum_k {H_k^2}$ when $i = j$ and $\mathbb{E}[H_{x_i}H_{x_j}] = \frac{1}{n^2}H^2$ when $i \neq j$. Therefore, the proof of Lemma 4.1 would start the same way when $\hat{H}=\sum_{i=1}^B H_{x_i}$. Results for the sampling with replacement case would immediately imply asymptotic results for sampling without replacement, since the total variation distance between the two distributions goes to zero as $n \rightarrow \infty$ and $B/n$ converges to a constant.

Reviewer comment (Theory):
The analyses presented here are solely on the stability of the iterates, i.e., whether they diverge or not. Is there any chance that this gives some insight into whether they converge? Even further, depending on the alignment of the Hessians, can we say something about the local convergence rate?

Author response:
Lemma 4.1 provides a sufficient condition to guarantee the linearized dynamics converge $w^*$. As we simplified part (i) of the Lemma using the Hessian alignment to prove Theorem 1, it is likely that similar manipulations could be used to provide sufficient conditions for convergence of the linear dynamics depending on the Hessian alignment.

We did not pursue this direction further since there seems to be less interest in this convergence behavior in prior work on linear stability of SGD. However, relating our analysis to local convergence rates of SGD is an interesting idea. For quadratic functions, this should follow directly. Extending this further would be a very interesting future direction, so we have added it to our conclusion as an open direction.

Reviewer comment (Theory):
The relationship between batch size and learning rate that I'm more familiar with (e.g., starting from Goyal et al. (2017)) is the linear scaling rule, but here it is shown to be squared, which has also been reported in the past (e.g., Krizhevsky (2014)). Can the authors elaborate on why this stability analysis leads to squared? Then, at least locally, is squared scaling law the way to go, given that the Taylor expansion is accurate?

Author response:
The derivation of $\mathbb{E}[\hat{J}^2]$ at the beginning of the proof for Lemma 4.1 may give an idea of how the squared scaling happens, particularly by considering the case where $H_i$ are $1 \times 1$. This emphasizes a connection to the additivity versus scalar multiplication in the variance of a random variable, i.e., $\operatorname{Var}(X+Y) = \operatorname{Var}(X) + \operatorname{Var}(Y)$ and $\operatorname{Var}(cX) = c^2\operatorname{Var}(X)$, where $X,Y$ are independent random variables and $c \in \mathbb{R}$ is a constant. We can interpret $\eta$ as a scalar multiplier in a sum of random variables and $B$ as the number of random variables being summed. This seems to imply that when the Taylor series truncation and Bernoulli dynamics are good approximations, we expect squared scaling to hold.

Reviewer comment (Evaluation):
In Figure 1, for small batch sizes, there is a large gap between Theorem 2 and the red boundary for $\eta = 0.8$. Any particular reason for this?

Author response:
We conjecture that this may be due to the discrepancy between Bernoulli sampling used in the theorem and sampling without replacement used in the experiments. Specifically when $B$ is small, differences between the two sampling methods may have more of an effect due to the higher relative variance in the effective batch size. In the case where $\eta=0.8$ is near the limit of GD instability, the stability condition seems more sensitive to this difference in effective batch size.

Reviewer comment (Experiments):
How would this be extended to momentum gradient descent or any of the adaptive gradient methods (e.g., Adam)? If the time allows, can the authors provide experiments for these as well in the same setting?

Author response:
Unfortunately, our experiments are coded in a way to run SGD on the described quadratic functions using minimal space/time, and as such, we cannot easily change the optimization method during the discussion period. However, we have added this idea as an open direction in our conclusion, as we agree it would be interesting to understanding how momentum affects the stability of SGD.

Reviewer comment (Novelty):
In the prior works on linear stability, were there any quantities that resemble $\sigma$ in their role? If there were, can the authors compare those to the proposed $\sigma$?

Author response:
The alignment factor $\mu$ of [Wu et. al., 2022] is a measure that is similar in spirit to our $\sigma$ measure. We have described this measure in Appendix B of our paper (see eqn. (4)). 

While this measure is also used to quantify the stability of SGD, its mathematical formulation is significantly different from ours. It is specific to MSE loss, and it measures the alignment of the gradients for the regression output and the empirical covariance of the MSE loss gradients. Meanwhile, our measure is defined through the Hessians of an additively decomposable loss function, and it applies to this general class of functions. It is difficult to push the comparison much beyond this.

Reviewer comment (Writing):
The introduction section (i.e., the first two pages) is poorly written. For example, the definition of 'linearized dynamics of SGD' is missing and is just referred to in Section 2. This is probably one of the most necessary things one needs to know to follow the paper.

Author response:
We believe it is more appropriate to reference Section 2 for the formal definition of the linearized dynamics, as we have done, rather than interrupt the introduction. Our paper advances a long line of work on characterizing the linear stability of SGD. As such, we believe readers who would be most interested in our paper likely have some familiarity with linear stability analysis. We believe our choice in this tradeoff is worthwhile, and Reviewer ELtf finds our paper to be "very well-written".

Reviewer comment (Writing):
The section 'contributions' is vague. Instead of long sentences, it's better to use a concise way to deliver the message. It is fairly impossible to identify the contributions of the paper based on that section (before reading the whole paper).

Author response:
Our contributions give a high-level description of our contributions that separates our results from prior work in the area. We have made an effort to describe our contributions as concisely as possible while including all important information by leaving the technical details to later sections. We would be happy to make changes if you have any specific recommendations for how the contributions could be rephrased.

Reviewer comment (Writing):
The phrase 'stability of dynamics' is used frequently in the paper without being explained in the early pages.

Author response:
In the paragraph before our contributions section, we have added the sentence: "We consider mean-squared stability, that is, $w^*$ is considered unstable if iterates of SGD diverge from $w^*$ under the $\ell_2$-norm in expectation".

Reviewer comment (Presentation):
The font used in the paper looks weird; it appears that the authors have changed something in the LaTeX code.

Author response:
We have double-checked our paper to ensure it follows all ICLR formatting guidelines and made any necessary changes.

Reviewer comment (Theory):
After Theorem 1, it is unclear why the expectation and the arg max are equal. Additionally, it is unclear how divergence allows the conclusion that for almost all initial parameters, SGD converges. These claims are not clear and are vague.

Author response:
We have changed the "argmax" to "max". For any matrix $M \in \mathbb{R}^{n \times n}$, $\||M\||_2 = \max_w \||Mw\||_2$ such that $\||w\||_2 = 1$ by definition of the $\ell_2$-induced vector norm. Hence, our statement immediately follows from our definition of the linearized dynamics, i.e., $w_t = \hat{J}_t....\hat{J}_1w_0$.  

If you are referring to our statement that $\||w_t\||_2 \rightarrow \infty$ for almost all starting $w_0$ points under the conditions of Theorem 1, this follows due to the fact that the set of vectors orthogonal to the maximum singular vector of matrix has measure zero.

Reviewer comment (Writing):
The paper lacks intuitive and qualitative explanations of the characterized linear stability, which could enhance its accessibility.

Author response:
We have made an effort to explain our intuitive understanding of our results within the page limits, and we note that Reviewer ELtf states: "Very well-written. I especially liked how the theoretical discussions are backed with intuitions and specific examples.". However, we understand that what is intuitive to us may not be to everyone. Hence, we could appreciate any more details on what specifically could be improved to make it accessible to a wider audience.

Reviewer comment (Experiments):
The experiments are limited to engineered quadratic losses without considering actual neural networks or real-world data.

Author response:
We note that our main contributions are theoretical and consider the general problem of minimizing additively decomposable loss functions. By restricting ourselves to these quadratic losses, we are able to avoid introducing unknown structure into the problem. This allows us to validate our general theory and get an idea of its tightness in the worst case.  

Computing the coherence measure naively would take about $O(n^2d^3)$ time, which is prohibitive for realistic neural networks. However, it likely can be approximated more efficiently using iterative methods, for example, the main bottleneck is computing $\||H_i^{1/2}H_j^{1/2}\||_F$ for all $i,j \in [n]$. However, we can estimate the Frobenius norm of a matrix product through matrix-vector products much more efficiently than it can be computed exactly.

We do agree that extending this understanding to neural networks and real-world data is a critical next step, as we noted in our conclusion. However, we believe that handling these numerical challenges is best left as future work since our current paper has a different focus.

Reviewer comment (Theory):
The paper should address scenarios where the condition $\nabla_w l_i (w) = 0$ is violated. There may be various cases that $\| \nabla_w l_i \|$ is small but non-zero, e.g., cross-entropy loss without label smoothing, early stopping, and so on. How can the proposed analysis accommodate these cases?

Author response:
We agree that the understanding the behavior of the case where the sample-wise gradients do not completely vanish is an important problem. At the same time, we would like to note that our assumption that $\nabla \ell_i(w^*) = 0$ for all $i \in [n]$ is a generally accepted assumption in this line of work (see e.g., [Ziyin at. al. 2023] and references within); and this is an important and practical regime to study.

Reviewer comment (Novelty):
There are existing works to characterize the stability of SGD considering the noise covariance matrix $\Sigma = \frac{1}{n}\sum_i \nabla l_i(w) \nabla l_i(w)^T - \nabla L(w) \nabla L(w)^T$, without assuming $\nabla_w l_i (w) = 0$. The paper should clarify how its results relate to these existing works.

Author response:
The work most relevant to ours that takes the approach of looking at the noise covariance matrix is [Wu et. al., 2022]. We have restated their main results in Appendix B and compare our results to theirs extensively in the related work section and in Section 3.2.1.

Some important observations we repeat here are that [Wu et. al., 2022] is restricted to MSE loss, and for the conditions of their main theorem to hold, the optimal loss must be non-zero. This makes further comparison somewhat difficult, since their result does not apply to perfectly fit data as we consider and as is more common to consider in overparameterized networks.

Reviewer comment (Writing):
What is the exact notion of stability considered in the paper? It is not clearly explained in the manuscript.

Author response:
In the paragraph before our contributions section, we have added the sentence: ``We consider mean-squared stability, that is, $w^*$ is considered unstable if iterates of SGD diverge from $w^*$ under the $\ell_2$-norm in expectation".

Reviewer comment (Reproducibility):
Is the coherence measure easily computable for typical neural networks? How complex would it be to compute in practice?

Author response:
Computing the coherence measure naively would take about $O(n^2d^3)$ time, which is prohibitive for realistic neural networks. However, it likely can be approximated more efficiently using iterative methods, for example, the main bottleneck is computing $\||H_i^{1/2}H_j^{1/2}\||_F$ for all $i,j \in [n]$. However, we can estimate the Frobenius norm of a matrix product through matrix-vector products much more efficiently than it can be computed exactly.

Reviewer comment (Writing):
On page 4, the term $x_i$ is used but not defined.

Author response:
Thanks for pointing this out. We have fixed the expression.

Reviewer comment (Theory):
It seems that the second last paragraph on page 4 assumes that $H_i$ is a matrix of rank one. Is this the case?

Author response:
Thank you for pointing out this omission. We have added this assumption in the example along with the following explanation for why it is reasonable.

For binary classification with $f$ denoting a scalar scoring function of weights and loss being an $\mathbb{R} \rightarrow \mathbb{R}$ function of $f$, the Hessian is the sum of a rank-one term and a full second order term. However, in over-parameterized systems, each $f_i$ attains its least value and hence the full second order term is zero, and so $H_i$ being rank-one is a good assumption. (See eqn. (13) for additional details here: https://www.cs.toronto.edu/~rgrosse/courses/csc2541_2021/readings/L02_Taylor_approximations.pdf)