A Precise Characterization of SGD Stability Using Loss Surface Geometry
Authors: Gregory Dexter
Borja Ocejo
Sathiya Keerthi
Aman Gupta
Ayan Acharya
Rajiv Khanna
Venue: ICLR 2024 poster (2024)
Keywords: SGD
linear dynamics
sharpness
implicit regularization
Primary area: optimization
TL;DR: Improved necessary and sufficient conditions for SGD to be linearly stable at a point when minimizing additively decomposable loss functions.

Abstract: Stochastic Gradient Descent (SGD) stands as a cornerstone optimization algorithm with proven real-world empirical successes but relatively limited theoretical understanding. Recent research has illuminated a key factor contributing to its practical efficacy: the implicit regularization it instigates. Several studies have investigated the linear stability property of SGD in the vicinity of a stationary point as a predictive proxy for sharpness and generalization error in overparameterized neural networks (Wu et al., 2022; Jastrzebski et al., 2019; Cohen et al., 2021). In this paper, we delve deeper into the relationship between linear stability and sharpness. More specifically, we meticulously delineate the necessary and sufficient conditions for linear stability, contingent on hyperparameters of SGD and the sharpness at the optimum. Towards this end, we introduce a novel coherence measure of the loss Hessian that encapsulates pertinent geometric properties of the loss function that are relevant to the linear stability of SGD. It enables us to provide a simplified sufficient condition for identifying linear instability at an optimum. Notably, compared to previous works, our analysis relies on significantly milder assumptions and is applicable for a broader class of loss functions than known before, encompassing not only mean-squared error but also cross-entropy loss.

Published as a conference paper at ICLR 2024
A PRECISE CHARACTERIZATION OF SGD STABILITY USING Loss SURFACE GEOMETRY
Gregory Dexter', Borja Ocejo', Sathiya Keerthi", Aman Gupta',
Ayan Acharya & Rajiv Khanna
I Purdue University
2 Linkedln Corporation
ABSTRACT
Stochastic Gradient Descent (SGD) stands as a cornerstone optimization algo- rithm with proven real-world empirical successes but relatively limited theoretical understanding. Recent research has illuminated a key factor contributing to its practical efficacy: the implicit regularization it instigates. Several studies have investigated the linear stability property of SGD in the vicinity of a stationary point as a predictive proxy for sharpness and generalization error in overparam- eterized neural networks (Wu et al. 2022 Jastrzebski et al., 2019; Cohen et al., 2021). In this paper, we delve deeper into the relationship between linear stability and sharpness. More specifically, we meticulously delineate the necessary and sufficient conditions for linear stability, contingent hyperparameters of SGD) and the sharpness at the optimum. Towards this end, we introduce a novel coherence measure of the loss Hessian that encapsulates pertinent geometric properties of the loss function that are relevant to the linear stability of SGD. It enables us to provide a simplified sufficient condition for identifying linear instability at an optimum. Notably, compared to previous works, our analysis relies on significantly milder assumptions and is applicable for a broader class of loss functions than known before, encompassing not only mean-squared error but also cross-entropy loss.
INTRODUCTION
Stochastic Gradient Descent (SGD) is fundamental optimization algorithm widely used in practice. In addition to its computational efficiency, there is irrefutable evidence of its superior generalization performance even on non-convex functions, including neural networks (Bottou) 1991). For large over-param treepparamttriieed neural networks, the number of points to fit is often much less than the number of free parameters in the model. In this case, there is often digh-dimessional manifold of model weights that can perfectly fit the data (Cooper) 2021); hence, focusing solely on the ability of an optimizer to minimize the loss function ignores a central part of training such networks. The primary goal in model is not to achieve high performance generalization performance on previously unseen data. Although we currently lack a comprehensive theoretical explanation for the empirical success of SGD in these models, promising hypothesis suggests that SGD naturally applies a form of implicit regu regularization (Zhang et al., [017]] Neyshabur al. 2015) when multiple optima are present (Keskar et al., 2017 Liu et al. 2020). This phenomenon guides the iterative process towards more favorable optima purely through algorithmic choices.
training data set but rather to achieve strong
In order to measure this distinguishing favorability between more and less desirable optima, prior work has proposed the concept of sharpness at a minimum as an indicator of the generalization performance of the trained model. Lower sharpness is often indicative of better generalization performance (Hochreiter & Schmidhuber 1997). There is wealth of empirical work exploring the relationship between sharpness and generalization performance, particularly in networks trained with SGD, e.g., (Jiang et al. 2019) JJastrzebski et al.] 2019; AAnrrishchenkoo et al.) 2023 Wu cet al., [0177] Chaudhari et al. 2017] Izmailov et al., 2018). Furthermore, these ideas have led to new optimizers which deliberately reduce sharpness and are observed to attain improved empirical performance (Behdin et al. 2023 Foret et al. 2020). Although the connection between sharpness
"Correspondence Rajiv Khanna (rajivak purdhneeuurdueeeduu

---

Published as a conference paper at ICLR 2024
and generalization performance isn't precise completely understood, the partial achievements of this theory has inspired several works, including ours, to investigate how SGD implicitly tends to converge to flatter optima.
Sharpness has been defined in several ways in prior literature, but most commonly, the sharpness of a trained neural network minimum is the maximum eigenvalue of the Hessian of the loss with respect to weights. Intuitively, one can see that w is a stationary point of a smooth function f(w) with
Hessian Hw),, then for perturbation v such that ||v||2 l22 = e, f(w* +v) < ==((")==(=")>> (H (w")). where X1 (-) denotes the maximum eigenvalue. This relation follows from the Taylor expansion of f()) around w", and we see that the sharpness at w i.e. A1 (H(w*)) determines how rapidly small perturbations to the weights w can increase the value of f (ww). In other words, model sharpness measures how robust the loss of the trained model to small perturbations of the model parameters.
In this paper, our focus is on providing a precise characterization of how the SGD hyperparameters and properties of the loss function affect its implicit regularization of model sharpness. Towards this goal, we consider the linearized dynamics of SGD (defined in Section 2 close to the optimum. When w is close to w", it allows us to make a useful simplification by focusing on the quadratic approximation of the loss function. In particular, we consider mean-squared stability, that is, w is considered unstable iterates of SGD diverge from w under the l2- fo-norm in expectation. Unlike differential equation- based approaches, which liken the SGD dynamics to continuous flow for eliciting implicit regularization properties (Li et al. 2017 Xie et al. 2021). the linear stability analysis does not break down in the regime of large step sizes. Moreover, the linearized dynamics in gradient descent (GD) has been empirically validated to predict the sharpness of overparameterizee neural networks due to the Edge-of-Stabilic phenomenon (Wu et al.
20h8/Cohen Cohen et al.) 2021). This
behavior has also been observed in other optimizers (Cohen et al. 1022|Jstrzebsb el al. 2019
Bartlett et al.) 2022 Wen et al.] [022]] [Ujááry et al. (022). lending further weight to this theoretical
framework (Agarwala & Dauphin|2000 (202)). While prior work has already considered the linear stability of SGD (Wu et al. 2022 2018 Ma Ying 2021 Ziyin et al.] 2023 Agarwala Dauphin 2023).
our analysis provides substantial advancement over these prior results as we detail below.
Contributions:
We offer an interpretable yet rigorously established sufficient condition to determine the instability of a point w under the lincarized dynamics of SGD (Theorem [ Importantly, unlike previous works, our bound applies to any additively decomposable loss function.
Our sufficient condition hinges on a coherence measure a, which we introduce to capture the relevant geometric characteristics of the loss surface around minimum. This measure is intuitively connected to the Gram matrix of point-wise loss Hessians. We provide additional context and rationale for introducing this measure in Section
We demonstrate that our bound is nearly optimal across a natural range of SGD hyperparameters (Theorem 2) This implies that our analysis, which offers a sufficient condition for the stability of linearized SGD dynamics, is precise and closely aligned with the behavior of SGD across various choices for the coherence measure o, batch size, and learning rate.
In the course of deriving Theorem 4.1). This lemma provides sufficient conditions for (i) the divergence of linearized SGD dynamics and (ii) the convergence of linearized SGD dynamics toward the specified minima. An intriguing aspect of this lemma is that it suggests a multiplicative decoupling effect between the impact of batch size and learning rate instability and the instability introduced by the geometry of the Hessian.
present independently useful technical lemma (Lemma
Finally, we corroborate the validity of our theoretical findings through a series of experiments conducted on additively decomposable quadratic loss functions. Our experimental results align with our theory and underscore the significance of the Hessian coherence measure a in determining the conditions under which SGD dynamics diverge.
Related Work: While extensive research investigates the intricate relationship between optimization methods, generalization error, and sharpness, the prior work most relevant to ours focuses on a linear stability analysis of SGD. In this section, we briefly compare our results to related research. However, we defer a detailed comparison of our results until Section 3.2.I] which follows the formal introduction of the problem setup and the presentation of our primary theorem.
2

---

Published as a conference paper at ICLR 2024
An important line of work in this area is that of Wu et al. (2018 2022 Wu & Su (2023), which progressively distill more theoretical insight into how the linear dynamics SGD affect final sharpness of neural network. Our work goes beyond this in multiple ways. For sake of comparison, the result in this line of work most related to our contribution is Theorem 3.3 in Wu et al. (2022). and restate this result in Appendix [ Our results are significantly more general than this theorem, in that our results apply to any general additively decomposable loss function, which answers the question raised in Agarwala & Dauphin (023) on the behavior of SGD for cross-entropy loss A Additionally, we guarantee a stronger form of divergence. Even with relaxed conditions and stronger implications, the condition of our theorem is practically easier satisfy than Theorem 3.3 in Wu et al. 2022).
Other research delves into related questions, although our study may not align directly with them.
For example, Jastrzzbski et al. (2019) combine previous analysis by Wu et al. (2018) with additional assumptions about how the optimizer behaves. The paper demonstrates the impact of learning rate and batch size sharpness throughout the training trajectories of SGD, Agarwala & Dauphin (2023) examine how batch size affects sharpness within SGD training trajectories, particularly in the context of second-order regression models. Ma & Ying 2021) provide meticulous characterization of linear stability. However, this characterization might not be immediately interpretable and is primarily used to draw connections between behaviours of SGD and Sobolev regularization. Ziyin al.] (2023) focuses on the convergence and divergence of linearized dynamics in probability rather than in expected distance from the optimum, as considered by the other work we have mentioned.
PROBLEM FORMULATION
We consider the case where SGD is used to minimize an additively decomposable loss function L(w) D E (w(). where each & (w() is twice-differentiable and w E TR.. Given learning rate O and batch size B E [m]. the dynamics of SGD are defined by the recurrence W++1 w, B Eies (w.)). where S uniformly sampled from all sized subsets of [n]]. To facilitate our probabilistic analysis, we apply two standard simplifications. First, we consider Bernoulli sampling rather than sampling with replacement so that E S with probability B/n and the event E S is independent of the event j E S for all i * j. Second, we consider the quadratic approximation to the loss around fixed point w so that (w)
(,(w") + (w - "')'V((("")) + { w*) WWu et al. 2022 Ma & Ying 2021). Since the dynamics of SGD are shift-invariant, can assume w 0 without loss of generality. We restrict our attention to the case where w' is local minimum of (-) for all E [n]. This assumption is particularly relevant
w*)² (
in the context of overparameterized neural networks, where is common for data fit the model almost perfectly (Allen-Zhu et al., 2019). In the described linearized setting, VQ((w) (ew( + W w))
(2(((w")
Define H, (,,,(")) which is the Hessian of 4(-) at w Note that H; E Rddd is a Positive- Semidefinitie matrix (PSD) since w' is a local minimum of (-))) (we refer the reader to Appendix A. for notation and necessary background). The linearized dynamics of SGD in our setting of interest follows below.
Definition Linearized SGD Dynamics: Let (Hi)) be set of d x d PSD matrices, and let H Σ H. Let 1) >> denote the learning rate and B E [n] be the batch size. The linearized SGD dynamics defined by the recurrence relation:
Wi+1 1
IH H W., SES
(1)
B and the event E S is independent from the event E S for all i + j. n We will refer to J = I 7/H and I - iie H; as the Jacobians of GD and SGD respectively. Note that using B = n recovers the gradient descent dynamics.
where E S with probability
Note that linear stability is not meaningful for pure eross-entropy loss on perfectly fit data, since || is not finite. However, our theory holds when using label smoothing (Szegedy al. 2016), as commonly done in practice.
3

---

Published as a conference paper at ICLR 2024
3 THE ROLE OF HESSIAN GEOMETRY IN SGD INSTABILITY
This section introduces and motivates the Hessian coherence measure o (HH,) Subsequently, utilize this measure to present our primary result, Theorem | This theorem furnishes sufficient condition for E|] wall2 to diverge as k + cc. Following this, we demonstrate that our established sufficient condition in Theorem is nearly optimal across broad range of hyperparameters. We formally state this optimality result in Theorem 2
3.1 HESSIAN COHERENCE MEASURE
Note that, to understand the behavior of linearized SGD dynamics around a point w " suffices to consider how they operate under various configurations of 1). B, and (Hi)e(n) To illustrate the effect of HH hee[n) let us explore two extreme scenarios.
Suppose we have H; = e, Vi, where represents the i-th canonical basis vector. In this scenario, with all H. being identical, we anticipate that the stochastic dynamics will closely resemble the deterministic dynamics. This expectation arises from the fact that Eies H, should exhibit strong concentration around H. Furthermore, when the elements of S are sampled from [n] without replacement, the SGD dynamics coincide with the GD dynamics. Therefore, we expect no difference in the characterization of stability of the respective linearized dynamics.
First Setting:
Now, let us consider the opposite extreme, where all H; matrices are orthogonal, meaning that their inner products satisfy tr Hr] 0 Vi & j. In this scenario, we anticipate that randomness exerts a substantial influence on the steps taken by SGD, the context a full linear GD step, the component projected onto the subspace defined by H. is /H.w/fo However, in the stochastic setting, if H, is not selected in the sampling process, no step is taken in this particular subspace. Conversely, if it is selected, the step taken is n/,ww/ B, which significantly overshoots the deterministic step by a factor of n
Second Setting:
These extreme metric arrangement within the set Hej in determining the stability of the dynamics (alongside the learning rate and batch size). While establish both sufficient and necessary conditions to address this geometric aspect in Lemma 4.1 our aim is also to offer an intuitive characterization that captures the significance of He without resorting to complex analytical expressions. To that end, we introduce the following measure, which succinctly captures the geometric structure within (e(m))
illustrative examples, highlighting the importance of the relative geo-
Definition 2. Coherence Measure: For a given set of PSD matrices (H:)ve(n) define S E R
such that S.y -|H|//H /11r. Equivalently, we may define as the entry-wise square root of the Gram matrix of (H, [H.) under the trace inner product. The coherence measure a is then defined as:
(S)
max4e[m) À (H,)
To provide some insight into this measure, we can consider (H, (H.)eejn) as collection of n vectors in Raxd, endowed with the trace inner-product. The Gram matrix of of vectors compiles the pairwise inner-products among the vectors, thus representing the relative alignments and magnitudes of these vectors within the space, and the matrix S is entry-wise renormalization of the Gram matrix. In the case where rank(H,) = 1 for all E [n]. x (H,) is the i-th diagonal entry of S, and a measure how close to being diagonally dominant. Due the construction of a then measures the iross-interactions within [H, Jie relative to the magnitude of the individual Hessians in (H; hejm under the Frobenius norm. The case where all H; are rank one is particularly important since it occurs when Vf(w") = under loss functions such as cross-entropy loss.
Let us examine the two extreme cases mentioned earlier. In the first case, where H; = e Vi, it Vi,j j [n]. Consequently, S becomes an n X n matrix consisting of all ones, yielding X (S) Vn. Meanwhile, in the scenario, where H, matrices are mutually
follows that [| H ; |
!/2
orthogonal, ||| H 1/2 1/2 [| 0 Vi / j. Therefore, S becomes the identity matrix I, and \s 1.
In both cases, X (H.) 1 for all [n). Consequently, in the first case, a vn, and in the
4

---

Published as a conference paper at ICLR 2024
second case, 1. This demonstrates that coherence measure, 0, effectively distinguishes between these two extreme scenarios and increases as the alignment among the H, matrices grows stronger. Below, show how this measure allows us to establish natural sufficient condition for the divergence of these linear dynamics.
3.2 SIMPLIFIED DIVERGENCE CONDITION
We present our sufficient condition for the linear dynamics diverge, which relies solely on the values of Åι (H),7 1), B, nt and our coherence measure a. Note that this bound aligns with our intuitive expectations based the extreme cases we examine. When all H, elef and
vn, the second condition in the theorem cannot be met. In such cases, we must resort to the GD condition for instability, namely, (H) > " Conversely, when H, e Vi and o 1, the theorem asserts that the linear dynamics will diverge even for small values of X (H), especially when B is small relative to n.. This behavior aligns with our expectations based on the different scenarios we consider.
Theorem 1. Let (JtheN be sequence of i.i.d. copies of defined in Definition | Let {H, (Hs)))))) have coherence measure a. If.
7) - -1-)- then, lim E [ 1/22-1112 00.
-1/2
A1 (H) or X1 (H)
k-+00
We defer all proof's to Appendix Note that the quantity EJ Ej...... ille = E max wooll=/ || [/wwlle ll2 where Wa is the random vector determined by the SGD dynamics in Definition ! when the dynamics start from Wo- In other words, if Je..J
- oc, the linearized SGD dynamics diverge from almost every starting point Wo E R. We highlight two observations of the condition in Theorem [ First, this analysis supports "squared-scaling" between batch size and learning rate, that is, if B is increased proportional to nj², the stability will not increase. Second, the Hessian alignment (captured by o) can cause instability even when 17 is small and B is on the order of n.
diverges
3.2. COMPARISON TO PRIOR WORK
Based on the formal introduction of our result, we now provide more detailed comparison to the result of Wu et al. (2022). The condition outlined in Theorem 3.3 of Wu et al. (2022) represents one of the most recent findings in this research field, which we restate in Appendix B The contrapositive
of this theorem provides sufficient condition, namely |||] F > + po to guarantee instability.
Here, Ho serves as an alignment metric, empirically argued have a practical lower bound.
Theorem T has multiple advantages over these prior results. First, the analysis in Wu et al. (022 is confined solely to the MSE loss function. Second, their definition of stability entails that BTT. W+)
C. [ Vk: E N. where C' > is a constant. This definition is notably weaker than our notion of stability. Finally, despite Theorem [ holding for more general loss functions and guaranteeing a stronger notion of stability/instability it is also easier to satisfy under the typical setting where H has low stable rank, i.c., when |||| | is small. Let us ignore the effects of the measures μo
and by considering both equal to one, then Theorem 3.3 (Wu et al. 2022) guarantees instability if
H " /B while Theorem guarantees instability if X (H)
which is more general
condition when the stable rank of H is less than n, as is typical in practical settings (Xie et al. 2022).
3.3 OPTIMALITY OF THEOREM
Theorem Dnnoormmprrovinee provides a sufficient condition for the linear dynamics diverge, where the condition is of the form Aı (H) > [(1,, n, B), where [n, n, B) = 7 (# 1) -1/2 The next theorem shows that our condition is optimal in the sense that, for natural range of parameters (when Ø, B = (()), the function f f17, o, n, BB is within a constant factor of its lowest possible value. This shows that our sufficient condition cannot be significantly relaxed without relying on other information about the set }H.) as we do in Lemma 4.1
Overall, the following theorem demonstrates that sufficient condition for divergent dynamics approaches optimality among all sufficient conditions that rely solely on n) o, B, X (H), and n. However, it does not rule out further improvement in the important regime where B <
5

---

Published as a conference paper at ICLR 2024
Theorem 2. For every choice of X1
0, n N, B [n]. 0, and [n], that satisfies:
20 A
(o -
7
There exists a set of PSD matrices (Hip) such that X1 (H) A and lim E E|J J......1. 11113- < n.
SHARP STABILITY CONDITIONS OF LINEARIZED SGD
In the previous section, we provide measure of the geometric coherence in (H; he|n along with theorem that provides sufficient conditions for the dynamics of SGD (Definition I to diverge. The proof of this theorem relies on part (i) of the following technical lemma. Aside from its utility in establishing the aforementioned theorem, the statement of the following lemma also imparts valuable insights into the behavior of linearized SGD dynamics.
The proof of the following lemma relies on the observation that | < |M < d for all M E RRdxd, where || ||x denotes the Frobenius norm. Hence, we may focus on divergence in the Frobenius norm of the k-step linearized dynamics. Now, E| [|.JJ || = [[tr[J,....] J?.. Jx]l
tr J JR]] by linearity. By operator monotonicity of the trace, further have tr(N&] 2 tr
T2...
tr[E[J, J P & 2 tr[M,], where N₂ = E[J P.... > M2 under the Loewner ordering (see Appendix A.I.. The technical challenge in our proof lies in composing an inductive argument to define matrices N₂ and M2. We opt for an approach that directly bounds the matrix of the k-step linear dynamics under the Loewner ordering, which does introduce greater technical complexity compared to using norm incqualities, as seen in previous work. However, this added complexity is essential to accurately account for the alignment in the unstable eigenvectors of each J., allowing us to provide thorough characterization of the instability of SGD dynamics.
Lemma 4.1. Let J. be independent Jacobians of SGD dynamics described in Definition (i) If
2 A1 (H)
11²
r
lim k+000
H
H
Do,
1
nB
n
V1 Ul
then lime E| J
oc.
(ii) ff. some E (0, 1).
c (H) < 2 7
for all i E [d] and
lim 1--00 &
1
1)2
972 n
Σ
H 91 0,
11.1
7
nB
then lim E g.... J 112. = 0.
Notice that part (i) and part (ii) of this theorem complementary in the sense that the condition of part (ii) is nearly the negation of the condition in part (i) except for the additional factor. In sense, the parameter& captures the balance of how close we to instability in GD dynamics, i.e. and how much additional instability is added by the stochasticity in the dynamics of
X (H)
SGD.F
To provide a more detailed explanation of this intuition, let us consider the setting where B < The second term from part (ii) of the above Lemma can be approximated as:
2)2
2k
Σ V1...
H Ub
H
Σ
H y
E
A
A A11
nB
n
n
B
V1
B
=1
We believe that the requirement part (ii) that (H could likely be removed by more carcfully
accounting for alignment of the negligible eigenvectors of J and J and relaxing the theorem to imply boundedness of the limit. However, given that the role of part (ii) in the theorem is only to contrast with part (i), we do not think this is high priority for the purpose of this paper.

---

Published as a conference paper at ICLR 2024
where A, is independently sampled uniformly from the set (H, Heejs) Interestingly, this implies that the effect of the parameters B and can be decoupled from the structure within (H Ne In other words, if we knew the minimal learning rate 1j at which linearized SGD with a fixed batch size diverges, then would immediately be able to determine which parameter pairs (17, B) are divergent at the given point w", since the term [S| A 1/:- does not change with these hyperparameters.
Note that determining whether the quantity " AR.. l diverges for arbitrary inputs of 17, B, and set of arbitrary symmetrie matrices (He would be an NP-Hard problem (see Section 3.3 of Huang et al. (2022), ). Even in our case, where we have the additional constraint that each H, is PSD,
are unaware of an efficient method to determine whether n² E A|
diverges, motivating
our simplified sufficient condition in Theorem
EXPERIMENTS
In this section, we support prior theorems by empirically evaluating the behavior of SGD synthetic optimization problems with additively decomposable loss functions. The high-level points our experiments support are:
Parameter tuples that Theorem guaranttes divergence do indeed diverge.
Parameter tuples that Theorem guarantees no divergence indeed do not diverge.
The coherence measure a has an important effect on the instability of SGD.
Our theoretical results hold when using SGD that samples without replacement.
The first two points outlined above serve validation for the accuracy and soundness of our theoretical results and proofs. The third point enhances the rationale for adopting the Hessian coherence measure we introduce. Lastly, the fourth point offers justification for employing SGD with Bernoulli sampling in our theoretical analysis, as its behavior mirrors that of the more prevalent SGD approach that samples without replacement. To ensure reproducibility, we include all our implementations in the supplementary material.
5.1 EXPERIMENT SETUP
We leverage the construction used in the proof of Theorem 2 to verify our predictions empirically, which offers two advantages: 1) we may apply the analysis of Theorem 2 for a condition that guarantees no divergence, and 2) the construction is parameterized by a, and so we may easily test the effect of varying In this construction, we set H, m eree for all E [o] and H, = m e; @e-o+tee otherwise, with m 2n We set the dimension of the space to n a 1, so there is unique minimizer of the loss, as this does not affect divergence. Notice that this construction essentially interpolates the two extreme settings in Section 3.1 as a varies from a = 1 to o = n. Additionally, note that A (H) = by construction. In experiments, < 1. hence the first condition of Theorem 1 i.e., An (H) > 2/1 will not hold or be relevant to characterizing stability.
The loss function that corresponds to the set of Hessians (He)e is given by the additively decomposable quadratic function L(w) nt Σi= ((w). where (w) w"H... Note that, for this construction of (en)) ((w) is particularly easy to compute, as it is equivalent to squaring and rescaling a single entry of w.
Across all experiments, we set n = 100. For each set of parameters (B, O), we determine whether the combination leads divergence or not by executing SGD for a maximum of 1000 steps. Specifically, we classify tuple divergent if, in the majority of the five repetitions, the norm of the parameter vector increases by factor of 1000. Conversely, we terminate SGD prematurely and classify the point as not divergent if the norm of decreases by a factor of 1000 during the course of the SGD trajectory.
It is possible to construct regression problem with mean-squared error that results in exactly the same optimization we describe. However, we note that Theorem 3.3 in Wu et al. (2022) does not directly apply to this problem, since the corresponding loss is zero at the optimum, and hence the loss-scaled alignment factor (see Equation Equation44below beloow of Wu et al. (2022) is undefined at the optimum. Therefore, the condition of Theorem 3.3 in Wu et al. (2022 (see Theorem cannot be satisfied.
7

---

Published as a conference paper at ICLR 2024
50
50-
30
Sigme
20
10
Batch Size
(a) Learning rate sct to rj 0.8.
(b) Learning rate is set to = 0.5.
Figure 1: The red area indicates where SGD diverges and blue where it does not diverge among parameter pairs (a, B). The solid black line is where the condition of Theorem I attains equality and the dashed line is where the condition of Theorem 2 attains equality.
5.2
EXPERIMENTAL RESULTS
5.2.1 EFFECT OF COHERENCE MEASURE AND BATCH SIZE
First, we look at how the stability of SGD changes as we vary coherence measure a and batch size B. We show results for two fixed values of the learning rate, 1 0.8 and 0.5.
Two key observations Figure 1 are that all tuples (a, B) that are below the boundary given by Theorem indeed diverge. This is visually shown as all points below the solid black line are red (besides some aberration due to visual smoothing). Additionally, the fact that all points above the dashed line are blue indicate that tuples (a, .B) which the proof of Theorem guarantees converge indeed converge.
An intriguing observation is the pattern where the gap between the upper and lower bounds diminishes as the batch size increases. Specifically, we notice that the lower bound more closely aligns with the actual boundary between divergence and convergence across all batch sizes when ! 0.5. However, the upper bound is closer to the true boundary when v = 0.8.
Finally, observe that the coherence measure exerts substantial influence the stability of SGD. For small values of a, SGD demonstrates instability even at high values of B. This observation underscores the importance of considering the geometry of the loss surface understanding the behavior of SGD. Furthermore, it highlights that the coherence measure is an effective tool for capturing and accounting for the contribution of loss surface geometry to the stability of SGD.
5.2.2 EFFECT OF BATCH SIZE AND LEARNING RATE
Next, we examine how the stability of SGD evolves when we manipulate the batch size B and the learning rate 7,. For this analysis, maintain fixed value of 5, which provides clear boundary given the granularity of the point grid. We show both the iog-scale plot, since learning rate generally varies on scalsc and the linear-scale plot since we expect the relationship between and r7² to be roughly linear in the boundary.
The plots displayed in Figure further corroborate the validity of Theorem and Theorem Additionally, the pattern continues to support the facts that the lower bound condition more closely approximates the true boundary as learning rate decreases, and the upper bound provides tighter approximation to the true boundary as learning rate increases.
CONCLUSION
We present precise yet interpretable, necessary and sufficient conditions to determine when a point w is stable under linearized SGD dynamics. The sufficient condition in Theorem rrllie novel
8

---

Published as a conference paper at ICLR 2024
10
80
60
Batch 520
Ealdon size
10--
20
10y
1
10P
0.2
0.4
0.6
0.8
1.0
Learning
Squared Learning Rate
Log-scale.
(b) Linear-scale.
Figure 2: The red area indicates where SGD diverges and grey where it does not diverge among parameter pairs (1). B). We plot the squared value 17² make the linear relation between B and 17 clearer. The solid black line is where the condition of Theorem I atttain equality and the dashed line where the condition of Theorem attains equality.
coherence measure a that summarizes relevant information in the loss surface geometry. We next list some open questions our work raises:
In future research endeavors, it would be intriguing to close the gap between Theorem [ and Theorem to establish the actual dependency of SGD stability on o and the hyperparameters of SGD.
Additionally, it would be interesting to empirically measure the value of the coherence measure in realistic neural networks. Acquiring knowledge about the practical range of values that a can assume would enhance the utility of the theoretical contributions provided here for predicting the behavior of SGD in real-world scenarios. Developing efficient approaches to approximate a in large neural networks would represent a valuable step toward achieving this objective.
We may also consider extending the same proof techniques to characterize the stability sharpness-ave -wwre aware methods (Behdin et al. [2033] Foret et al. 2020 Zhuang et al. 2022 Liu al.) 2022] Kwon al.) Kim Kim al. 2022). which commonly employed for
training many overparameterizes models.
Along these lines, would also be useful to consider whether the stability of SGD with momentum or other adaptive gradient methods could be analyzed with this approach.
The convergence analysis in Lemma [..mma4..1]aa could possibly used to derive fine-grained local convergence rates of SGD depending Hessian alignment.
ACKNOWLEDGMENTS
GD was partially supported by NSF AF 1814041, NSF FRG 1760353, and DOE-SC0022085. RK would like to acknowledge support from AnalytiXIN Indiana.
REFERENCES
A. Agarwala and Y. Dauphin. Sam operates far from home: eigenvalue regularization as dynamical phenomenon. In International Conference on Machine Learning. 2023.
Allen-Zhu, Li, and Song. convergence theory for deep learning over-paramettriiation.. In International conference on machine learning, pp. 242-252. PMLR, 2019.
M. Andriushchenko, F. Croce, M. Müller, M. Hein, and N. Flammarion. A modern look the relationship between sharpness and generalization. arXiv preprint 1.X1v:2300200011,,2000
P. I.. Bartlett, P. M. Long, and O. Bousquet. The dynamics of sharpness-aware minimization Bouncing across ravines and drifting towards wide minima, 2022. URL R.K.l//.s..//a.vvoog///////////

---

Published as a conference paper at ICLR 2024
K. Behdin, Q. Song, A. Gupta, A. Acharya, D. Durfee, B. Ocejo, S. Keerthi, and R. Mazumder. mSAM: Micro-batch
berro-bacch-averaged averaged tharpness-aware minimization. arXiv preprint arXiv:2302. 09699999.9999. 2023.
R. Bhatia. Matrix analyxis, volume 169. Springer Science Business Media, 2013.
L. Bottou. Stochastic gradient learning in neural networks. In Proceedings of Newro-Nimes 91, Nimes, France, 1991. EC2. URL pettp://lecnt bottou. wwwpr//appp////e wotttetr//ttoottoooo 91c
P. Chaudhari, A. Choromanska, S. Soatto, Y. LeCun, C. Baldassi, C. Borgs, J. Chayes, L. Sagun, and R. Zecchina. Entropy-SGD: Biasing gradient descent into wide valleys. In International Conference Learning Represen- tations, 2017. URL
not http://wwwww.oooooooooooooo.eaa. wwweee/ooomniddateeeeeoe
J. Cohen, S. Kaur, Y. Li, Kolter, and A. Talwalkar. Gradient descent on neural networks typically occurs at edge of stability. arXiv preprint arXiv:2103.00006 00060, 2021.
J. M. Cohen, B. Ghorbani, S. Krishman, N. Agarwal, S. Medapati, M. Badura, D. Suo, D. Cardoze, Z. Nado,
G. E. Dahl, and J. Gilmer. Adaptive gradient methods at the edge of stability, 2022. httpp::////xx org arg/abs// abs/
0atss/207..1. 14484
Cooper. Global minima of overparameterized neural networks. SIAM Journal on Mathematics of Data Science, 3(2):676-66) 101667669112220 2021.
P. Foret, A. Kleiner, H. Mobahi, and B. Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In International Conference Learning Representations, 2020.
S. Hochreiter and J. Schmidhuber. Flat Minima. Neural Computation, 9(1):1 -2, 01 1997. ISSN 1899-7667.
D. Huang, J. WillesWeed J. Tropp, and R. Ward. Matrix concentration for products. Foundations Computa- tional Mathematics, 1766)177 2022.
Izmailov, D. Podoprikhin, T. Garipov, D. Vetrov, and A. Wilson. Averaging weights leads to wider optima and better generalization. In 34th Conference on Uncertainty in Artificial Intelligence 2018, UAI 2018, pp. 876-885, 2018.
S. Jastrzebski, M. Szymezak, S. Fort, D. Arpit, J. Tabor, K. Cho, and K. Geras. The break-even point optimization trajectories of deep neural networks. In International Conference on Learning Representations, 2019.
Y. Jiang. B. Neyshabur, H. Mobahi, D. Krishnan, and S. Bengio. Fantastic generalization measures and where to find them, 2019.
N. Keskar, D. Mudigere, J. Nocedal, M. Smelyanskiy, and P. Tang. On large-batch training for deep learning: Generalization gap and sharp minima. JCI.R, 2017.
M. Kim, D. Li, S. X. Hu, and T. M. Hospedales. Fisher sam: Information geometry and sharpness aware minimisation, 2022. Utthssts:://axx http:////w/wwwwloo////////////// C4920
J. Kwon, J. Kim, H. Park, and I. K. Choi. Asam: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks. In Proc. of ICML, volume 139, pp. 5905 5914, 2021.
Q. Li, C. Tai, and F. Weinan. Stochastic modified equations and adaptive stochastic gradient algorithms. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference Machine Learning, volume 70 Proceedings of Machine Learning Rexearch, pp. 2101-2110. PMLR. 06-11 Aug 2017.
S. Liu, D. Papailiopoulos, and D. Achlioptas. Bad global minima exist and sgd can reach them. Advances in Neural Information Processing Systems, 33:8543-8552, 2020.
Y. Liu, S. Mai, X. Chen, C. Hsieh, and Y. You. Towards efficient and scalable sharpness-aware minimization. In Proceedings of JEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12360-12370. 2022.
C. Ma and L.. Ying. On linear stability of sgd and input-smoothness of neural networks. Advances in Neural Information Processing Systems, 34:16805-16817,, 2021.
B. Neyshabur, R. Tomioka, and N. Srebro. In search of the real inductive hias: On the role of implicit regularization deep learning, 2015.
10

---

Published as a conference paper at ICLR 2024
C. Szegedy, V. Vanhoucke, S. Ioffe. Shlens, and Z. Wojna. Rethinking the inception architecture for computer vision. 2016 JEEE Conference Computer Vision and Pattern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016, pp. 2818-28200 IEEE Computer Society, 2016.
Ujváry, Telek, Kerekes, A. Mészáros, Huszár. Rethinking sharpness-aware minimization as variational inference, 2022. URL
K. Wen, T. Ma, and Li. How does sharpness-aware minimization minimize sharpness?, 2022. URL https:/7 //wrxivarvvvvv...a http//aaiv.....aaaaa..... httpg//bb//2/..... 05729
Wu Su. The implicit regularization dynamical stability in stochastic gradient descent. In International Conference on Machine Learning, 2023.
Wu, Zhu, and Weinan, Towards understanding generalization of deep learning: Perspective of loss landscapes. ArXiv, abs/1706.10299, 2017.
Wu, Ma, et al. How sgd selects the global minima in over-parameterizer learning: dynamical stability perspective. Advances Neural Information Processing Systems, 31, 2018.
Wu, Wang, and W. Su. The alignment property sgd noise and how helps select flat minima: stability analysis. Advances in Neural Information Processing Systems, 35:4680-4693,, 2022.
Xie, I. Sato, M. Sugiyama. A diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima. In International Conference on Learning Representations, 2021.
Xie, Q. Tang. Y. Cai, M. Sun, and Li. power-law hessian spectrums in deep learning. arXiv preprint arXiv:2201.13011 2022.
Zhang. S. Bengio, M. Hardt, B. Recht, and Vinyals. Understanding deep learning requires rethinking generalization. 2017.
Zhuang, Gong, Yuan, Y. Cui, H. Adam, N. Dvomek, Tatikonda, J. Duncan, and T. Liu. Surrogate gap minimization improves sharpness-awan training, arXiv preprint S:Xii:2203008066 2022.
Ziyin, Li, Galanti, and M. Ueda. The probabilistic stability stochastic gradient descent, arXiv preprint arXiv:2303. 13093, 2023.

---

Published as a conference paper at ICLR 2024
A PROOFS
A. PROOF PRELIMINARIES
this section, we provide notation and necessary background for the following proofs.
Notation: Let [r]] denote 1, 2......: 71, Let e; denote i-th canonical basis vector. Let Bern(p) Bernoulli distribution with parameter p. Let O denote the all-zero matrix or vector, where the dimension will be clear from context. Let 1 be defined similarly as the all-ones matrix or vector. Let ||A||2 denote the spectral norm of matrix A. Let (A) denote the i-th largest eigenvalue of the matrix A. Let | ||| denote the Frobenius norm matrix A. Let | denote p-Schatten norm matrix A. that is, p-norm the vector singular values A.
Fact 1. f1- Norm Inequality: For any vector RR, ||x||2 ||x|1
vd[|x|12-
Fact 2. Recursive Formula for Binomial Coefficients: For all n, k N such that ls S the binomial coefficients satisfy the following recursive formula:
n2
n--1
1:
k
The notion of PSD matrices can be used to define a partial order on the set of symmetric matrices as follows.
Definition 3. Loewner Order: The Loewner order is a partial order on the set of positive semidefinite symmetric matrices. For two positive semidefinite matrices A and B, we write A = B to denote that B A is positive semidefinite and A < B to denote that B A is positive definite.
Note that if A is a PSD matrix, then for any symmetric B, BAB must also be PSD.
We frequently use the following properties of the trace in our derivations.
Properties of the Trace:
tr(A]=E. (A).
Invariance to cyclic permutation: tr [ABC] tr (CAB..
Linearity of the trace: tr(cA + B] = etr[A] + tr( where c IR.
Lemma A.1. For any matrix M E x=*x" ||M]| < M |s I /M|M
Proof. This follows from the fact that || M s. and |M is the p-norm of the spectrum of M along
with applying the & form incquality (Fact
Lemma A.2. For any length & sequence of square matrices An...u E gaxd, Lr.A,A2...A.) S Vol A Aeel
Proof. First, it follows from Weyl's Majorant Theorem that tr Ar. S AA Iss (see Section I11.5 Bhatia 2013)). Then, applying Lemma A.1 limpiies
tr Ar
A A2.
Finally, conclude the lemma statment by submultiplicativity of the Frobenius norm.
A.2 PROOFS
Proof of Theorem 1
Proof. Note that [ < |M] d| for all M Rdxd Therefore, by Lemma 4.1 suffices show
that the condition theorem implies that the following quantity diverges towards infinity as k - co:
11
H 11:
nB
n²
12

---

Published as a conference paper at ICLR 2024
From here, lower bound it by the following. Note that |M] norm (scc Lemma A.
2
||
where || is the I-Schatten
172
2
2
H
Σ
1 Cr H 2k. H
H M2
Hy
B
n2
1B
n
91
M1 2/
17'
11
Σ
nB
n
31 D
11'
11
H
nB
m
11
Σ
H
nB
n?
nd
The last line follows by again applying the LI-L2 norm inequality. Now we show that the complexity measure satisfics:
H
d² ok
First, bound the left-hand term follows.
By
Lemma
A.2
know that Ar...A) S
d|| ||||||/||||...... | A || [r. Therefore.
Lr(H"]
H Hyy
n
Yi Vi-
tr
Wi- (.l._,... ( (HH/21//2))
1
Σ
H 1
////2/1 ve- 11
| 1 H 3/
n
SI. S SuL-n .... Sy
n
df S
ir(S*]
n (s)*
n
Therefore,
n² max:e|m| (H,,)
@ oF [r[H]]
(r
S max sE|n]
An (H.)* - Σ [
d3 A1 (S)
Starting from where left off before,
E Lr J7
112
11
Σ
Lr [
BBB
n
nd
1=
11
11
n
Lr (r(H")
nBB
?
d2
17'
17
X (H) 2E
nBB
?
ndo
2k
(H)
0-
B
nd.5
We see that by the condition of the theorem that A (H) towards infinity as le + oo. Hence, we conclude the theorem statement.
1.
(
) 1/2
that this above equation diverges
Proof of Theorem 2
13

---

Published as a conference paper at ICLR 2024
Prouf. Construct (H, (e(m) that H, eref if [o] and H, 0 Otherwise, where Aue Note that A1 (H)
A1. Furthermore, note that || ||H,2] H || m if i, E [o] and zero otherwise. Therefore, S, the entry-wise square root of the Gram matrix (i.c, Sy ||H || (12)), has all of the first x o entries equal to and rest equal to zero. Therefore, (s m a. Meanwhile, max4e[n] (H..) m. Hence, we see the Hessian coherence measure for this constructed problem is indeed equal to the chosen value of a.
We now show that lima-ro F.l.......j random variables sampled from Bermoulli distribution with parameter p B/n. A key observation here that all J are diagonal (since I and all H, are diagonal) and hence JI...Jx commute. Therefore,
31 12 Note that J -- H, where the as are i.i.d.
J + Etr J........Jf..................... Jf.. E Et J
tr [(]]
where in the last step we use the fact that all of the J are independent and identically distributed. Recall that the trace is the sum the diagonal entries a matrix and hence tr [r[]]]}]*] = - of E[ Note that
ef [[]]]]]e e, = when i * 1, so tr [[]] = (n -1-1 eñ tr(E[Ji]*]e. To show that this quantity is bounded by ns, we must show that of [[] e1 (of e( R bounded. Following the base case of part (i)
Lemma 4.1 we have that:
n72
1² n² H
BJJ = 1- 21/H + 2²²
nB
We can then write ei E[E[]e]e as:
2² n2
1 n
of [T = 1 211=
+
m
nB
where we used the fact that of He A (H) in construction and ej Hfe; m² if [o] and zero otherwise. show convergence, we must show that previous equation is than This is equivalent to the following condition:
222 n
m
+
112
r
T
<
21)
0
m
YB
1
1/20
n
1)
D
12
3
21
0
(#
1)
(
1
A (H)
< 2n (H)
2
n/B
-
(H
[+
20
n
X1
0
1
B
Therefore, conclude the theorem statement
Proof of Lemma 4.1
Proof. First, note that | = Lr(M²M). Since J and J are always symmetric, to prove the theorem conclusion, we must show that limg-+2o proving a PSD lower bound on the expectation of this product of matrices and then leveraging the operator monotonicity of the trace. To that end, we define the matrix:
co. will achieve this by inductively
M. 22
17
Σ
H
H H.
nB
14

---

Published as a conference paper at ICLR 2024
We show that B[Jx J. JI.1..JJ = M. Vk E N. First, let us start with the base 1. Recall that H Eisi H. and H B Σi= H,H,, where 2; ~ Beru(P/a)
E[J,J,] [[( nf )l H1H)
H
+
n
E
1-2
11
24 H: Σ 1=1
+
7
H
B
B
2n
H
E Emray H,H,
11=1
At this point, we have simply substituted relevant definitions and rearranged the terms. Now, note that E[2,2,]
D² if i j and E[,,,,] otherwise. We substitute this in, and use the fact that n H² > H.)
B[ = - 217 H.
H
H.H,
nB
11
(F]
-- 21 1
Σ
H
+
H.H
11
H
n
nB
n²
1=1
=1 =I- 1-22HH4
H
H
nB
n2
1=1
M
Thus, we have proved the base case B[J ž M1, which also implies BBJ > M1. since they are identically
distributed. We next prove inductive step under the assumption that [[ k- ..... J JR... t MA-1. By the Law of Total Expectation and linearity,
B| Jf.... J k = E[E[J [[E]J...JJi........ J J... [ EIJ. B JR-1] Ju] > BJJ,Ma JN).
The arithmetic in the next part of the argument is analogous the argument in the base case.
E[JaMa-IJ. TE( M -7]
M - M HH
H. M Hy
B
f,)=1
M
M
HHM
H.M H
iJ=1
+1²
H. M - H.
nB
JJM Ju +
H H
nB
n2
Next, we substitute the definition M Recall that if A is symmetric and PSD, then ABA is also PSD.
EJ Jx)
J
[2(k-1)
17a
n
Σ
H Hy
H²
Hy.
Jr
nB
n2
+-²
1
n
H,
12(-11)
n
Vo H H Hy 9e-1
H.
nBB
nB
i=1
Vi co W
Σ J 2k
Σ
H
H
H
nB
n2
11
Mi Vi....
MM
15

---

Published as a conference paper at ICLR 2024
Note that able drop the cross-terms in the first step since ABA O. whenever A is symmetric and B is PSD. Therefore, we have show that E[J JI..... 2 M2 for all l E (n). By the operator monotonicity and linearity of the trace,
tr
[l] tr[
J..... [R..J.]]] 2 tr
Mr(M)]
77
of²
lr
Σ
H
H
H
nB
n
91
77
H
nB
n
Hence, we have proved part (i) of the lemma. We will now prove part (ii). Define the following matrix:
N.
17
Hy. H Hy
nB
n²
Mi-...Br=1
and define No I. We will prove that:
k N. E 2(k--)) r=0
k
E J
J&l
for all & N. where O. First, note that the condition of the theorem implies -(1 (1-c)1 (1 <) for some e First, note that the base case:
إاJ < (1 -1- No Ni (1-c)))))
π
H
nB
n2
holds from computations in part (i). since J2 1 (1 c) Now for the inductive step. Under the
assumption that Equation hholds for l: 1. show also hold for k. By the argument in part (i), find that:
J
J&l
1 E
J&
Σ (1 5) 2(6-1-1)
k 6-11
N,
Ju
Ju
R-1
Σ
5) 2(2-1-1) 1-r)
k-
N₂
P=0
k-
k
11
H.
Σ
2(k
1-r)
N
H:
nB
r
1=1
r=0
(11-6)
Σ
(1 e)
h
h-
Nr
r
11² 1
K-
+
H.
Σ
[1
2(k
1-r)
N,
H:
nB
n
r
k-1
Σ
(1 )) (k--r)
k-1
N
0=0
Σ D=d
k-1
e) 2k6-1)
k-1
1
H.N.H.
nB
n²
[m]
(1 8) 2(k-r)
b 6-
(1 5) 2(k-1-r) k
N
r
16

---

Published as a conference paper at ICLR 2024
where the last line follows from to go from r t k by replacing
Ei=l H.N.H, Nr+1. We can rewritc the second summation
+ with to combine the terms like so:
--1 Σ ==0
-1--)
k 1-11 N
c) (k-r)
k
N.
r r-
r=1
(1 (1 No
+
k-1
e)at 2((
k-1
+ +(1 -2) -2--))ae 2)
k 8-1]
N. + No
r r-1)
r
k (+)
(+)
= (1 s) No +
Σ (1 - c) 2(()))))
N. + N
Σ (1 - (1-c) ()((()))
h
N
Hence, we have completed proof by induction. Note to go from the second to the third line the previous equation block, we applied the recursive formula for binomial coefficients (see Fact 2). Now, doing the same as before, we can conclude that:
& J&J S El
(M)
E
c)2(k-1)
Ir[N)..
r=0
By the condition of the theorem that lim
n²
H
H 117- = O, we have that:
nB
tr IN[N]
11
21
Σ
tr
H
Hyr
nB
n²
11"
1
Σ
H
n
goes to zero as r + Dc. Let n., [r and = (). To prove the theorem statement, it suffices to show that lima- Et=o 5t ( (*)n, n, = 0, when E (0, 1) and (n,) is sequence of positive real numbers converging zero. By the theorem statement, we know that + tr(N,) constant C such that + tr[N,] S C all r [m]. which implies tr [N S c.C. Therefore, by the Binomial formula:
0. Therefore, there exists some
Σ
(1 (1-2) 211-2 -c' ((1-c)² +c)*
B4
r=0
This term must go zero as k + 0o, since (1 e)² + 1 1
B RESULT OF WU ET AL. (2022)
Here, we summarize the relevant information needed to state Theorem 3.3 Wu et al. (2022). which considers the settings where one has access dataset of the form (x,, (0) where x, IR", ys R. f(x, w) function parameterized by w such that f E" x Rd -> R. The loss function applied to each sample is MSE, denoted by & (ww = []((x; w) yil² and the total loss over all the samples is given by L(w) = = Σ (ww). except that S C [x]] is sampled from all
The linearized SGD dynamics considered is the same as Definition [ sizc B subsets of [n] uniformly.
One can now define two critical matrices. Let E(w) = 1 Σ. VE(()V VL(W)VL(w)) be the noise covariance matrix, and let G(w) = 1 Σr Vf(x, w V f(x;, w)" the Fisher matrix that characterizes the local geometry of the loss landscape. Next, Wu et al. 2022) define the loss-scaled alignment factor (Equation 3 in Wu et al. (2022)):
tr 2 G
((w)
(4)
The last step before stating the theorem is the introduction of the definition of linear stability considered in Wu
al. 2022
17

---

Published as a conference paper at ICLR 2024
Definition 4. (Linear stability in WW al. (2022)) global minimum w" is said to be linearly stable if there a such that holds for linearized dynamics that E[L(w:)] S [[]]) Vt 2 O, with We being sifficiently close to w".
Finally, the theorem follows.
Theorem (Theorem in Wu et al. (2022) ) Let w a global minima that is linearly stable (Definition
Denote by p(w) the aligumentt factors for linearized SGD (Equation 4). If p(w) > po, then [[( )l[F S
P4
18


Reviews:
[{"review_id": "Cmetfbx872", "reviewer": "Reviewer_asFh", "summary": "This paper analyzes the linear stability of SGD through the lens of the loss surface geometry. Assuming a scenario where the model perfectly fits the data, the necessary and sufficient conditions for linear stability is characterized by learning rates, batch sizes, sharpness (the maximum eigenvalue of the Hessian), and a coherence measure of the individual loss Hessians, which is newly proposed in this work. The theoretical findings are validated through experiments on a synthetic optimization problem.", "questions": "- What is the exact notion of stability considered in the paper? It is not clearly explained in the manuscript.\n- Is the coherence measure easily computable for typical neural networks? How complex would it be to compute in practice?\n- On page 4, the term $x_i$ is used but not defined. \n- It seems that the second last paragraph on page 4 assumes that $H_i$ is a matrix of rank one. Is this the case?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper relaxes certain assumptions in characterizing the stability of SGD, e.g., the restriction to MSE loss, from existing works.\n- The mathematical derivations appear to be sound and accurate.", "weaknesses": "- The paper lacks intuitive and qualitative explanations of the characterized linear stability, which could enhance its accessibility.\n- The experiments are limited to engineered quadratic losses without considering actual neural networks or real-world data.\n- 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? \n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RoL4dJv2IV", "reviewer": "Reviewer_mpeb", "summary": "In this paper, the authors considered quadratic approximations of the loss function around the zero-loss manifold that allows the definition of the 'linearized SGD dynamics' (given in Equation 1). This quadratic approximation allows the definition of the coherence measure (Definition 1), which is used in the statement of Theorem 1: SGD dynamics diverge when the coherence measure lower bounds the first eigenvalue of the Hessian of the loss function. In Theorem 2, they provide a partial converse to their proven divergence results in Theorem 1. The paper concludes with some experiments.", "questions": "- the font used in the paper looks weird; I think the authors have changed something in the latex code\n\n- In Theorem 1, $\\hat{J}_i$ is used, while it is never defined (the reference to Definition 1 only defines $\\hat{J}$).\n\n- After Theorem 1, why the expectation and the arg max are equal? Also, how does divergence allow you to conclude that for almost all initial parameters, SGD convergences? These claims are not clear and are vague.\n\n\n\n\n\n------------------------------------------------------------------------------------------------------------\n\n\nAfter the rebuttal: I appreciate the authors for their response. They partially answered some of my concerns but this paper is still not well-written, in my opinion. The authors only referred me to another reviewer for this part of my comments which I think is not an approrpiate response.", "limitations": "", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 1, "contribution": 1, "strengths": "- an important problem on the stability of SGD\n- having experiments", "weaknesses": "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. But this is probably one of the most necessary things one needs to know to follow the paper. Moreover, the section 'contributions' is also vague. Instead of long sentences, it's better to use a concise way to deliver the message. It's fairly impossible to identify the contributions of the paper based on that section (before reading the whole paper). \n\n\nSection 2: the approximation is called 'linearized dynamics,' but I think this is not the right word since you are essentially approximating the loss function with a quadratic function. Moreover, this dynamics only happens if you project to the zero-loss manifold at each step; otherwise, there is a term relating to the gradient. As a result, the setting is quite limited to only quadratic loss functions. \n\n\nThe word 'stability of dynamics' is used frequently in the paper while not being explained in the early pages.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9HBusHX1cW", "reviewer": "Reviewer_ELtf", "summary": "The paper analyzes the linear stability of SGD of any additively decomposable loss function around the minimum $w^*$. The paper then derives a necessary and sufficient condition for the (in)stability, which relies on a novel *coherence measure* $\\sigma$, which is, in turn, intuitively connected to the alignment property of the collection of Hessians (of individual loss function). This is verified experimentally.", "questions": "1. 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?\n2. 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?\n3. 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?\n4. 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?\n5. 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$?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. Very well-written. I especially liked how the theoretical discussions are backed with intuitions and specific examples. I very much enjoyed reading this work.\n2. Although the mathematics used here are rather elementary, the theoretical discussions are complete; a necessary-sufficient condition for stability is provided, and the motivation and intuition behind the coherence measure $\\sigma$ are well-explained. (Elementaryness is a plus for me)\n3. Clear contribution and a good advancement in the linear stability analysis of SGD", "weaknesses": "1. 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?\n2. 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?)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DyojlzrUsO", "reviewer": "Reviewer_SPXN", "summary": "This manuscript investigates the linear stability of SGD and obtain a sufficient condition for instability (equivalently, a necessary condition for stability).\nThe authors introduce a coherence measure $\\sigma$ to measure the strength of  alignment among Hessian matrices. Using this measure, they derive their main result Theorem 1, which, as they claim, is more general than Theorem 3.3 in [Wu et al, 2022]. The authors also show that Theorem 1 is nearly optimal given that $\\sigma$ and $n/B$ are $O(1)$ quantities. Some experiments are carried out to support the theoretical results.", "questions": "- In paragraph 1, 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). \n - In paragraph 2, when citing empirical works on relating sharpness to generalization, I think the important comprehensive investigations by (Jiang et al., ICLR2020) is missed.  \n - In paragraph 4, 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).\n - In Section 2\n \t- In 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. \n \t- In Definition 1,  it is unclear whether the sampling is done with or without replacement.\n- In Section 3\n\t- What is the $\\frac{1}{B}\\sum_{i=1}^n x_i H_i$ in the second paragraph. \n\t- In Theorem 1, what does the subscript in $\\hat{J}_i$ stand for? Complexity measure => coherence measure. \n\t- In Theorem 2, what do you mean $\\sigma\\in [n]$? Is $\\sigma$ a real value?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "On a whole I think these results are neat, novel and have their own advantages.  The characterization seems to be rather precise and Lemma 4.1 is of particular interest.", "weaknesses": "- The writing of this manuscript needs to be improved. In particular, the many details are present in an unclear way and I find it very hard to follow them smoothly. The citations and references also need to be re-organized. See Question section for more details. \n\n- The definition of ``Coherence measure'' is not intuitive.  Personally, I do not understand why the proposed definition can quantify the coherence among Hessian matrices. In particular, the authors claim that $\\lambda_1(H_i)$ is the i-th diagonal entry of $S$. This is obviously not true unless that $H_i$ is rank-1.\n\n- 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.\n\n - In Section 3.2.1 the authors compare Theorem 1 to Theorem 3.3 in [Wu et al, 2022], and stated the advantages of their result. Among these stated advantages,\n   - The first point makes sense to me.\n   - In the second point, why do you say ``This definition is notably weaker than our notion of stability''?\n   - By the third point, you seem to imply that the bound in Theorem 1 is sharper in the low-rank case. But what is the point in considering $\\sigma$ equal to one? To me, the third point is an unclear comparison between two results, which cannot prove the advantage of Theorem 1.\n\n\n- Theorem 2, the optimality of Theorem 1, strongly relies on the condition that $\\sigma, n/B = O(1)$. There are two concerns:\n   - 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.\n   - Also, I do not think it is natural to assume $n/B=O(1)$ as usually $B\\ll n$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)