Understanding Edge-of-Stability Training Dynamics with a Minimalist Example
Authors: Xingyu Zhu
Zixuan Wang
Xiang Wang
Mo Zhou
Rong Ge
Venue: ICLR 2023 poster (2023)
Keywords: edge of stability
nonconvex optimization
gradient descent
training dynamics
scalar network

Abstract: Recently, researchers observed that gradient descent for deep neural networks operates in an ``edge-of-stability'' (EoS) regime: the sharpness (maximum eigenvalue of the Hessian) is often larger than stability threshold $2/\eta$ (where $\eta$ is the step size). Despite this, the loss oscillates and converges in the long run, and the sharpness at the end is just slightly below $2/\eta$. While many other well-understood nonconvex objectives such as matrix factorization or two-layer networks can also converge despite large sharpness, there is often a larger gap between sharpness of the endpoint and $2/\eta$. In this paper, we study EoS phenomenon by constructing a simple function that has the same behavior. We give rigorous analysis for its training dynamics in a large local region and explain why the final converging point has sharpness close to $2/\eta$. Globally we observe that the training dynamics for our example has an interesting bifurcating behavior, which was also observed in the training of neural nets.

Published as a conference paper at ICLR 2023
UNDERSTANDING EDGE-OF-STABILITY TRAINING DYNAMICS WITH A MINIMALIST EXAMPLE
Xingyu Zhu"l, Zixuan Wang"? Xiang Wang', Mo Zhou', Rong Ge
Duke University, "Tsinghua University {xingyu. zhu, zw270 (aw200@duuke @duke. edu, {xwang, mozhou, rongge)ecs.
ABSTRACT
Recently, researchers observed that gradient descent for deep neural networks op- erates in an "edge-of-stabilit" (EoS) regime: the sharpness (maximum eigen- value of the Hessian) is often larger than stability threshold 2/n (where v the step size). Despite this, the loss oscillates and converges in the long run, and the sharpness the end just slightly below 2/n,. While many other well-understood nonconvex objectives such matrix factorization or two-layer networks can also converge despite large sharpness, there is often a larger gap between sharpness of the endpoint and 2/1.. In this paper, study EoS phenomenon by constructing a simple function that has the same behavior. We give rigorous analysis for its train- ing dynamics in large local region and explain why the final converging point has sharpness close to 2/1- Globally we observe that the training dynamics for example have an interesting bifurcating behavior, which also observed in the training of neural nets.
INTRODUCTION
Many works tried to understand how simple gradient-based methods can optimize complicated neu- ral network objectives. However, recently some empirical observations show that optimization for deep neural networks may operate in more surprising regime. In particular, Cohen et al. (2021) observed that when running gradient descent on neural networks with fixed step-size n), the sharp- ness (largest eigenvalue of the Hessian) of the training trajectory often oscillates around the stability threshold of 2/n',, while the loss still continues to decrease in the long run. This phenomenon is called "edge-of-stability" and has received lot of attention (see Section 1.2 for related works).
While many works try to understand why (variants of) gradient descent can still converge despite that the sharpness is larger than 2/n, empirically gradient descent for deep neural networks has even stronger properties. As shown in Fig. la, for a fixed initialization. if one changes the step size y,, the final converging point has sharpness very close to the corresponding 2/1,. We call this phenomenon "sharpness adaptivity". Another perspective on the same phenomenon is that for wide range of initializations. for a fixed step-size 17. their final converging points all have sharpness very close to 2/1. We call this phenomenon "sharpness concentration".
Surprisingly, both sharpness adaptivity and sharpness concentration happen on deeper networks, while for shallower models of non-convex optimization such as matrix factorization or 2-layer neural networks, the gap between sharpness and 2/n often much larger (see Fig. 1b). This suggests that these phenomena are related to network depth. What is the mechanism for sharpness adaptivity and concentration, and how does that relate to the number of layers? answer these questions, this paper consider minimalist example of edge-of-stabilit
More specifically, we construct an objective function (4-layer sealar network with coupling entries), such that gradient descent this objective has similar empirical behavior deeper networks. We give rigorous analysis for the training dynamics of this objective function in a large local region, which proves that the dynamics satisfy both sharpness adaptivity and sharpness concentration. The global training dynamics for our objective exhibit complicated fractal behavior (which is also why rigorous results are local), and such behavior has been observed in training of neural networks.
Equual Contribution.
'The value 2/1 is called the stability threshold, because if the objective has fixed Hessian, the gradient descent trajectory will become unstable if the largest eigenvalue of the Hessian larger than 2/1-

---

Published as a conference paper at ICLR 2023
30
2/300
1174
2/6
2/300
1.00
2/400
2/40
10
1071
100
Stappese (A)
37.5
2.
Sharpone
300
Sharpurs
20.11
A. 10
300
201
2/200
2/8
2/300
250
2200
1- 2/30)
1/400
2010
2/19
225
2000
3000
0000
5000
10000
20
60
100
200
6:0
1000
Eterations
Eoorations
Eterations
(a) ReLU 5-layer FC network with neurons per layer. = 2/1)
(b) Linear 2-layer FC network with neurons per layer (A 2/m))
(c) 4-layer scalar network.
(A 2/r)
EoS Phenomena in NN Training. We consider three models including a 5-layer ReLU activated fully connected network, a 2-layer fully connected linear network with asymmetric initial- ization factor (4, 0.1) (see Appendix A. for explanation), and 4-layer scalar network equivalent to mill,,, {1 x 22 For each model we run gradient descent from the same initialization using different learning rates. For (a) and (c), the sharpness converges very close to 2/17 with loss contin- uing to decrease. For (b), the sharpness decreases to be significantly lower than f/1
Figure 1:
1.1 OUR RESULTS
aal-xyye xy One can view The objective function we consider is very simple: C(x, v. [1 this as a 4-layer scalar network (each layer has single neuron). We even couple the initialization so that æ 2, = so effectively becomes an objective on two variables L(x; y) & [1 22) For this objective function we prove its convergence and sharpness concentration properties:
Theorem 1.1 (Sharpness Concentration, Informal). For any learning rate n smaller than some con- stant, there is constant size region S. such that the GD trajectory with step size 11 from all initial- izutions in S. converge to global minimum with sharpness within (2/n 311 2/1)-
As direct corollary, we can also prove that has the sharpness adaptivity property.
Corollary 1.1 (Sharpness Adaptivity, Informal). There exists constant size region S and responding range of step sizes that for all n E K, the GD trajectory with step size 1) from any initialization in S converges a global minimum with sharpness within (2/n) 3 2/n).
The training dynamics are illustrated in Fig. 2. To analyze the training dynamics, we reparametrize the objective function and show that the 2-step dynamics of gradient descent roughly follow parabola trajectory. The extreme point of this parabola is the final converging point which has sharpness very close to 2/n. Intuitively, the parabola trajectory comes from a cubic term in the ap- proximation of the training dynamics (see Section 3.1 for detailed discussions). We can also extend our result to setting where v, y are replaced by vectors, see Theorem 3.2 Section 3.3.
In Section 4 we explain the difference between the dynamics of our degree-4 model with degree-2 models (which are more similar to matrix factorizations or 2-layer neural networks). We show that the dynamics for degree-2 models do not have the higher order terms, and their trajectories form an ellipse instead of a parabola.
In Section 5 we show why it is difficult to extend Theorem 1 to global convergence the training trajectory exhibits fractal behavior globally. Such behaviors can be qualitatively approximated by simple lew-ee nonlinear dynamics standard in chaos theory, but still very difficult to analyze.
Finally, in Section we present the similarity between our minimalist model and the GD trajectory of some over-vorreeee parameeeettre deep neural networks trained real-world dataset. Toward the end of convergence, the trajectory of the deep networks mostly lies on a 2-dimensional subspace and can be well characterized by parabola as in the scalar case.
2

---

Published as a conference paper at ICLR 2023
1.2 RELATED WORKS
The phenomenon of gradient descent on the Edge of Stability (EoS) was first formalized and empir- ically demonstrated in Cohen et al. (2021). They show that the loss can non-monotonically decrease even when the sharpness À >> 2/1z. The non-monotone property of the loss has also been observed in many other settings (Jastrzebski et al., 2020; Xing et al., 2018; Lewkowycz al., 2020; Wang et al., 2022; Arora et al., 2018; Li et al., 2022a).
Recently several works try understand the mechanism behind EoS with different loss functions under various assumptions (Ahn et al., 2022; Ma et al., 2022; Arora et al., 2022; Lyu et al., 2022: Li et al., 2022b). Ahn et al. (2022) studied the non-monotonic decreasing behavior of gradient descent (which they call unstable convergence) and discussed the possible causes of this phenomenon. From a landscape perspective, Ma et al. (2022) defined special subquadratic property of the loss function, and proved that EoS occurs based this assumption. Despite the simplicity, their model displayed the EoS phenomenon without sharpness adaptivity. Instead, our model focuses on a minimalist scalar network and proves the convergence results together with the sharpness adaptive phenomenon.
Arora et al. (2022) and Lyu et al. (2022) studied the implicit bias on the sharpness of gradient descent in some general loss function. Both works focus on the regime where the parameter is close to the manifold of minimum loss. Arora al. (2022) proved that with modified loss VL or using normalized GD, gradient descent enters the EoS regime and has a sharpness reduction effect around the manifold of minima. Lyu al. (2022) provably showed how GD enters EoS regime and keeps reducing spherical sharpness on scale-invariant objective. In both works, the effective step-size n changes throughout the training process, so sharpness adaptivity and concentration do not apply. Our results start from a simpler example without normalization, whereas the above works focus on general functions with normalized gradient or
property.
Another line of works (Lewkowycz et al., 2020; Wang et al., 2022) focuses on the implicit bias introduced by large learning rate. Lewkowyez et al. (2020) first proposed "catapult phase", a regime similar to the EoS, where loss does not diverge even if sharpness is larger than 2/1,. Wang et al. (2022) provided a convergence analysis on the matrix factorization problem for large learning rate beyond 2/A where À is the sharpness. Their results include two stages: in the first phase, the loss may oscillate but never diverge; the sharpness decreases to enter the second phase, where the loss decreases monotonically. Recently Li et al. (2022b) provided theoretical analysis on sharpness along the GD trajectory in a two-layer linear network setting under some assumptions during the training process. These works mostly foeus the degree-2 setting which does not have the sharpness adaptivity and sharpness concentration properties.
PRELIMINARIES AND NOTATIONS
In this section, we introduce the minimalist model which exhibits both sharpness adaptivity and sharpness concentration.
2.1 GRADIENT DESCENT ON PRODUCT OF 4 SCALARS
We focus the simple objective C(x, y, w) 21 xyzw)²² Let the learnable parameters x, y. 2, W E E to be trained using gradient descent with a fixed step size 17 E that
(Ir-1, Bt+1,
w++1) = (It, Ut, 2t, we) - 1/ L(F,, 3/s 2t, Wt). 1 2t4
(1)
Here H, denotes the value of parameter x after the 1-th update. To further simplify the problem, we consider the symmetric initialization of 20 To. wo = yo- Note that due to symmetry of objective, the identical entries will remain identical throughout the training process, so the training dynamics reduces to two dimensional and the 1-step update of x and y follows
It+1 = T₁ 2
1). Mi+1 = 3t- = 1).
It's easy to show that the set of global minima for this function form the hyperbola ry = 1. Without loss of generality we focus the case when x, y > 0, and in most of the analysis we also focus on the side where x > y. As shown in Fig. 2, with GD running on such a minimal model, we observe convergence on EoS for a wide range of initializations. Eventually all such trajectories converge to minima that are just slightly flatter than the "EoS minima" (the minima whose sharpness is exactly 2/n, see Definition 1).
3

---

Published as a conference paper at ICLR 2023
I-Step Trajestrry (1000 Inerations)
Giibed Minena
Intrabottke
Meneee 1 -b) Mexnum 15 10)
Shoaster (3 229
Glabd MisaA
Thapsters (3 22100
Stedean 1A 12
Digletery (3 2/12)
14-
Starjese
11
0.45
a5-
sa
(luttalaattio))
2/3
2022
033-1
Chnneggtty Sherppone (50333
(a) Evolution of training loss, sharpness, and trajectory of GD on the scalar example from the same initialization with different learning rates.
(b) Initializations converging close to À = 2/n (1 = 0.2)
Figure 2:
EoS phenomenon on degree-4 model. In (a) we demonstrate sharpness adaptivity by 2
running GD with learning rate 7) = 8 10 12 from the same initialization. The sharpness of all trajectories converges around their corresponding stability threshold 2/n while the loss decreases exponentially, In the 2D trajectory, the 2-step movement quickly converges to some smooth curves ending very close the EoS minimum. In (b) we demonstrate sharpness concentration by running GD with constant learning rate 17 = 0.2 for 50000 iterations from a dense grid of initializations and plot the sharpness of their converging minima. Initializations in the red shaded area all converge to minima with sharpness in (2/n) 0.1,2/1)
2
2.2 EOS MINIMA AND REPARAMETERIZZTTIN
Given that wide range of initializations all converge very close the "EoS minima" with sharpness 2/n, we want to concretely characterize those points. The complete calculations are deferred to Appendix B.1. Denote 7 zy. the Hessian of the objective C (z, x y, y) admits eigenvalues
A
(0² + y²) (37 1)
+ 22 3/ + 4 10,² 7744)
(3)
A2 3 (a + )) (37² - 1) V 12 3 + 10,2+ 771)
and Aa = 22 7), A4 = (1- When (z,y) converges to any minimum, 7 = xy so X2, A3, A4 all vanishes. Therefore it is A that corresponds to the EoS phenomenon people observe.
When 77 < 1, solving A = 2/17 with zzy² 1 gives = + ( ?) + + 1),, #,
+V2((-44+ + 2) 1 and their multiplicative inverses. These solutions correspond to the minima with sharpness exactly equal to the EoS threshold of 2/1j. Since they are all symmetric with each other, without loss of generality we pick the minimum of interest as follows.
Definition (n-EoS Minimum). For any step size 1) E (0, }), the E-- minimum under the (x; y)- parameterization is
(i, ()
(1 ( (12 (1
.2 2)
(4)
Though we are able to obtain a closed-form expression for the EoS minimum, its z-y coordinate could still be tricky to analyze. Thus we consider the following reparameterization: For any (z, y) e
((x.)) v) R+ x R x > v}. define & (=²-y²) =y²) and d & zy. This gives a bijective continuous mapping between {(x.g) y) E RR x R+ : x > y} and {(c. E R x Rt.. This a natural reparameterization since intuitively the basis in the new coordinate system are the two orthogonal family of hyperbolas xy = C and x² y² = C. The former captures the movement orthogonal to the manifold of minima zy 1 while the latter captures the movement along the manifold of minima. Note that similar separation of dynamics was also used in Arora et al. (2022).
With c, d as defined, the EoS minimum simplifies to (č. - ((7 2 - 4) To expand the dynamics near the EoS minimum, we let a 2 c -1- 4) and b A d - 1 to be the offset from (Z, d). Our analysis will primarily be using the (a, al-pparameterization.
4

---

Published as a conference paper at ICLR 2023
Definition 2 (17-EoS Reparameterrization) For any step size 17 > 0. for any (x,y) E R x R such that x > y, the (a, b) reparameterization (z, y) are respectively given by
(a, - ((=* (²)) (n (m= - 4) xy -1
(5)
Let VTH following Eq. (2), the 1-step update under the reparameterization becomes
a++1 = (x 4 4) + (a + (x (x-+- - -4)) (1
4+1 = 4 + ((1 + b.)3 2(1 bb) + +1++b))))
(6)
+ ((1 bτ) - (1+ (4(1 (1 6,) xt + (axk (1-4x=)+)))
Now we can proceed to analyze the dynamics of this simple example.
3 DYNAMICS OF GRADIENT DESCENT ON DEGREE-4 MODEL
In this section, we will rigorously analyze the training dynamics characterized by Eq. (6). First will introduce the approximation of one and two-step update and build up intuition the dy- namics. Then we will present our main theoretical results that the degree-4 model exhibits both characterizations EoS training.
3.1 APPROXIMATING 1-STEP AND 2-STEP UPDATES
Here we introduce the informal approximation Eq. (6) and the corresponding two-step updates. For cleanness of presentation we will use x to hide all dominated terms. The rigorous statements of the approximations and corresponding proofs are deferred to Appendix B.3. When we are only describing the one/two-step dynamics, we use a, a', a" to denote de; a++1, a++2 and b, U, 8" to de- note b,, be+1, b/42. Denoting x A VT, when Ks |al, are all not too large (see precise ranges in condition B.1), we have
a' 2b²x' a" 4b2 8" = b 8abs 16b3.
b' 2 4abk 36² b²:
In the approximation, is monotonically decreasing at steady rate of 26 per step. The one step update is flipping signs and contains second and third order terms of For the two-step approximation however, the oscillation behavior and the even-order terms of b all cancels. This consistent with the analysis in (Arora et al., 2022) that the two step dynamics travels along a sharpness reducing flow.
Before proceeding analyze the discrete GD movement, we first get intuition by approximating the two-step dy- namics with simple ODE
Global
<T-8S Minimum
db bb 166 Sabx da
8 lone + he!
a"
xb
(8)
1.0
X Connerging Minmuun is -
This would be the limit when & going to 0 and the move- ment of two-step dynamics become very small. The gen- eral solution for Eq. (8) is given by
b² = = anx + to xt + +C +Cexp(8ax
for some constant C R. As a decreases following Eq. (7), the trajectory converges toward the parabola b² jan + * Note that the convergence to the parabola exponential with respect to a, so if is initialized positive and not too small, it will converge to minima that is very close to = x as shown in Fig. 3. This is minimum that is just slightly flatter than the x² -.oS minimum.
00.5
0..25
0.00
0.25
0.50
0.775
1.00
Figure 3: Solutions of Eq. (8) (K 1)
5

---

Published as a conference paper at ICLR 2023
3.2 CONVERGENCE ON EoS FOR THE DEGREE-4 MODEL
Now we state our convergence result on the 4 scalar objective under (a, b/-parameterization.
Theorem 3.1 (Sharpness Concentration). For large enough absolute constant K, suppose 2000v2 K- and the initialization (ag, bu) satisfies ag E (12k3, 1 - and by E (-K-1, K Kr MO... Consider the GD trajectory characterized in Eq. (6) with fixed step size 6² from (0o, bo), for any >> 0 there existy T = OKK-2x # +ogg(=-)) (o)((b/l))x== such that for all t > T. [be] < E and az E (- x²))
Under the context of x, y coordinate and sharpness, Theorem 3.1 gives the following corollary: Corollary 3.1 (Sharpness Concentration under (x, y) Parameterization). For large enough ab- solute constant K, suppose < 8000000 K -2 and the initialization (ro, yo) satisfies Do E (IE + 13m9,, I - 4) and |2oV0 1 (0, K-1) where (ě. is the n-EoS minima de- fined in Definition 1. The GD trajectory characterized in Eq. (2) with fixed step size 1) from (No. yo)
will converge global minimum with sharpness X E (7 20 1) 2)
Note that when the step size 7) (and hence k) is relatively small, the final sharpness is very close to 2/1. The range of initialization that satisfies the requirement is quite large: in the original (z, y)- ł) and height O (kK-11g )))) Many of the initial
parameterization it contains a box of width O K -2------ points can be far from the mosimmmmmm.
The complete proofs are deferred to Appendix B.6. Here we discuss the proof sketch of Theorem 3.1. Our convergence analysis focuses on the 2-step update. It contains two phases:
Phase 1. (Convergence to near parabola) We consider initializations in region I, II, and III.
8 = 2ax
Jas 1er
In I, 6" - 6 is dominated by -63 and (a, 6) follows exponential trajectory. We show that [b] decreases exponen- tially with respect and enters region II (Lemma 8).
6 for
Jou 1:
stot
|6|
2
In III, bs dominated by abs and (a, b) follows an el- liptic trajectory centered at (0, 0). We show that [b] increases at superlinearly with respect to a and enters II (Lemma 9),
1 4
1 1
8
Iv
We also show that once (a, b) enters II, it will stay in II until it exits from the left and enters IV (Lemma 11). Thus after Phase 1, all initializations will be in IV.
Phase 2. (Convergence along parabola)
MI
I
After (a, b) enters IV, we show that it will further con-
verge to the parabola that [b2 - To x²1 < 200 200** x will be
X- Minimim (0.0)
satisfied before a decreases to x and enters (Lemma 13).
Figure 4: Convergence Diagram for GD the degree-4 example. The quiver arrows indicate the directions of local 2-step movement. This dia- gram is only for demonstration pur- pose and ratios are not exact.
Then we show that the inequality will be preserved in V while it moves left until it enters VI (Lemma 14).
In VI, the dynamics is again similar to III, but with being negative. We conclude our proof by showing [b] will converge to 0 superlinearly with respect to a (Lemma 15).
Following Theorem 1. we can also formally characterize the sharpness adaptive phenomenon for a local region using the following corollary. The proof is deferred to Appendix B.6.2.
Corollary 3.2 (Sharpness Adaptivity). For a large enough constant K, fix any a < 2000 /2
K
For all initialization (5a, yo) in the region characterized by
To E b -a
(10)
and |zuyl 1| E (0, K ), the GD trajectory from (Zo, Yu) characterized by Eq. (2) with any step
size 11 E (a2 TO K 2 a²) will converge to a minima with sharpness X E (2 3 n, 3)
20
3.3 CONVERGENCE ON EOS FOR RANK-1 FACTORIZATION OF ISOTROPIC MATRIX
Inspired by the scalar factorization problem, we extend it to a rank-1 factorization of an isotropic matrix. In particular, we consider the following optimization problem:

---

Published as a conference paper at ICLR 2023
3 Taxd ay ay T T| 2 2.yER"
min
(11)
Similar to the under-parrameterized case in Wang et al. (2022), this problem also guarantees the alignment between and y if (ar, y) is global minimum, i.e., cy for some E R. To prove the convergence for Eq. (11) at the edge of stability, we first prove the alignment can be soon achieved. After the alignment, we prove the equivalence between this problem and the degree-4 scalar model, and prove the convergence of this problem.
We directly give the final theorem and the proof is deferred Appendix C. The experiments demon- strate similar EoS phenomenon (See Appendix A.5).
Theorem 3.2. For a large enough absolute constant K, with all the initialization (zo: yo) satisfying
2o ~ 6g Unif($tt=)) Yo ~ by Unif(St-1), d = 1.. Or E (t + 5)) k 2, + - -
K K-4 if step size 7) < min 5000000 20000-2000( 1ogioo3; b. and a multiplicative perturbation 3/- = (1(1 2K 2k is performed at time = 1p for some lp > (1-log()) + log(d) log + K for any > O, with probability p > 26g 2 expl ( there exists = O 2x 1
K-
log(e) log( on)) such that for all >>> f(x, v) and ||=|| + || E ( 3 1) 4
10
Note that we require additional perturbation because need guarantee that the trajectory does not converge to an unstable point (where sharpness À > 2/n). This was proved without perturbation for the scalar case but more challenging higher dimensions. The objective will still converge to minimum very close to an 17-EoS minimum. The experiment results are available in Appendix A.5.
4 DIFFERENCES IN DEGREF-2 2 AND HIGHER DEGREE MODELS
In this section, we will look at some similar models of lower degree, and explain why for degree-2 models the sharpness of final converging point is often farther from 2/1 compared to higher degree models. We will use similar methods as in Section 2 and Section 3 to gain intuition for the dynamics.
Previous works including (Chen & Bruna, 2022) and (Wang al., 2022) have studied the dynamics of beyond EoS training on the problem of factorizing single scalar or an isotropic matrix into two components. The objectives studied includes ming NER- (/4 - x y) min VER I - zy l|?. and the corresponding scalar case min ( 2y)= They were able show that for initializations with sharpness greater than 2/n, GD with constant learning rate 7/ provably converges to global minimum with sharpness less or equal to 2/1,. Empirically, the sharpness reduction process on these 2 2-component objectives will usually "overshoot" the EoS threshold and converge to minima that is significantly flatter than the EoS minimum, and one does not observe the oscillation of sharpness around the EoS threshold (see Appendix A.3).
In this section we consider the scalar objective min - zy)² since it is able to captures the major dynamical properties of those more complex objectives discussed in Wang et al. (2022). As shown in Fig. 5, initializations with sharpness exceeding the EoS threshold will converge to minima that is distinguishably flatter than the EoS minimum, and globally there not region of initialization that gives EoS convergence. Unlike the parabola for the degree-4 case, the 2-step update travels in roughly circular trajectory centered the 2-oS minimum shown in Fig. 5a (right). Therefore locally we observe that sharper initializations tend to converge to flatter minima.
The difference between the degree-2 and degree-4 case can be easily explained by local expansion. Using the same (c, d) reparameterization and setting (a, b) to be the offset of (c. d' from the EoS minimum, the two step update of (a. b) under learning rate n² can be approximated by
d" ~ a 22 0b*bb442abb
(12)
This is very similar to Eq. (7) except that we no longer have the -b3 term for the 2-step update on b which was attracting b close to 0. In this case, the ODE approximation db/ aa/da- 4a bxx gives the general solution 8² 4(C a²)/n² for E Ik+,, which corresponds to the family of ellipses centered at (0, 0) and matches the two step trajectory in Fig. 5a. 5a.
50 Un denote the uniform distribution over (d -iddimensio sphere with radius 6o-

---

Published as a conference paper at ICLR 2023
1-Step Trajectory (100 Itenation))
Liitel Mirrrs
X
SBaRAA Glab Minzz
Heceney (s-m 1000
Incertay (n- 201000
Midenas (k as
Tacwtay to 2003)
Misenas (Am am
teccetiag to 2000
(Cattabrattrv)
Qurgge Qll
x
(bilialbuckon k1
Oonoegiin Bers)
(a) Evolution of training loss, sharpness, and trajectory of GD on the scalar example from the same initialization with different learning rates.
(b) Converging sharpness of initializations (1) 0.2)
Figure 5: Beyond EoS training on product of two scalars. We the same experiment in
Fig. 2 except for using objective ming, VER (1 - -ry)². Note that in this case the two-step trajectories form circular curves and converge points that are farther from EoS minima.
In Appendix A.2.3, we discuss a degree-3 model exhibiting mixed behavior around different EoS minima, which further verifies our explanation above. We also empirically note that the coupling of entries will naturally arise when training general scalar networks from non-coupling initializations. Thus it is not an artifact we have to impose on the model to observe EoS (see Appendix A.4).
5 GLOBAL TRAJECTORY AND CHAOS
There exists very limited global convergence analysis for constant step size gradient descent training beyond EoS on complicated non-convex objectives. Even for the product of scalars, the boundary separating converging and diverging initializations (Fig. 6a) exhibits complicated fractal structures.
Moreover, we observe that for initializations close to such boundary. their GD training trajectories usually begin with a phase of chaotic oscillation which eventually "de- tieebifurrates"" and converges to the parabolic two-step trajectory as discussed in Section 3. Similar oscillation phenomenon has also been empirically observed by Ruiz-Garcia et al. (2021) in neural networks when they increase the learning rate and destabilize the network from a local trajectory.
So what is causing the bifurcation? Previously, Ruiz-Garcia et al. (2021) attributed the phenomenon to the cascading effect of oscillation along multiple large eigendirections of the network. Yet this explanation is quite unsatisfying for our simple model as there is only one oscillating direction.
Gabbel Miaias (Cacoogei
04
308
00
101-
Dd Malata
-0
K. S.K. starettteed
all
12 [lataaleatas
33
1.0
(a) Converging and diverging initialization (17 = 0.2).
(b) GD step-wise trajectory with asymmetric init. (1) 0.01).
(c) GD trajectory and bifurcation diagram of approximated 8"
Figure 6: Bifurcation behavior of GD on the degree-4 model. In (a) we show the zoomed in version of the lower right part of Fig. 2b, the fractal boundary can be clearly observed. In (b), we run GD with 7 = 0.01 starting from the asymmetric initialization (Irg, Mo) (12.5, 0.05) close to the boundary of divergence until it converge close to the EoS minimum at around (10, 0.1). In (c), we plot the trajectory with bifurcation under (a, b) rerepaaaatttiinanaionnn and compare it with the bifurcation diagram of the approximated dynamical system characterized by b" b(1 + Sax 1662).
8

---

Published as a conference paper at ICLR 2023
Looking closely to the trajectory (Fig. 6b), one will find it very similar to the bifurcation diagram of self-recurrent polynomial maps (such as the famous logistic map 21+1 = rIx(1 Iz) parameterized by r). In the degree-4 model, the existence of such reffrrecurrent map is explicit since following Eq. (7), the approximate 2-step update of b can be rewritten as b" = b(1 + 8ox - 1612).
If consider self-recurrent 1D nonlinear dynamical system - = b,(1 (1 + 8ax 160,)) parameterized by a. In Fig. 6c, we compute the bifurcation diagram for the recurrent map numerically and see that they are qualitatively similar. Following this analogy, one may instantly relate the first bifurcating point with the EoS minima that the trajectory eventually converges to, and the con-bifurcating regime for the polynomial maps with the sub-EoS regime" on the left (in Fig. 6b) of the EoS minima.
be relatively stationary, the trajectory of b will be locally characterized by the
CONNECTION TO REAL-WORLDY MODELS
In this section we show how the degree-4 model analyzed above resembles the converging dynamics of pee-parraeterizee regression models trained on real-worid dataset. We train a 5-layer ELU- activated fully connected network on a 2-class small subset of CIFAR-10 (Krizhevsky et al., 2009) with GD. The loss converges to 0 and the sharpness converges to just slightly below 2/17-
We visualize the dynamics by projecting the trajectory onto the subspace spanned by the top eigen- vector of minimum (oscillation direction) and the movement direction of parameters orthogonal to oscillation (see Definition 4 Appendix A. for exact characterization). As shown Fig. 7 (mid), after some initial bifurcation-like oscillation, the 2-step trajectory stabilizes and moves along some smooth curves toward the minimum. Near the minimum (Fig. 7, right), the trajectory in fact lies mostly in this 2-dimensional subspace (see Fig. 23c in Appendix) and can be very well-captured by a parabola, which is very similar to our minimalist example. More experimental results on real- world models are available in Appendix A.6.
238
Projacial Taiistocy (leration
Projeced Trajectry (Beratton 3000 18500)
CIKEE
330
Qutlalte Docetin
294
shorptone
200
-Largs
-0,00
180
1.50
06
6.3
C5
000
3.50
01001
0213
Figure 7: Training trajectory of 5-layer ELU-activated FC Network. We train the model using with 77 = 0.01 for 18500 iterations. The sharpness converges to 199.97 while 2/n = 200. The local trajectory (right) can be very well approximated by the parabola 7500y².
7 DISCUSSION AND CONCLUSION
In this paper we proposed a simple degree-4 model that captures the sharpness adaptivity and sharp- ness concentration phenomena that happen in gradient descent training of deep neural networks. The simplicity of the model allowed us to perform rigorous analysis on the training dynamics for large local region. The analysis gives new insights on why the training dynamics of the degree-4 model inherently different from the training dynamics of degree-2 models. Finally we show that the over- ver-paramterized deep networks trained on real data exhibits similar parabolic converging tra- jectory as the scalar example. We hope many of these observations can be generalized to highlight the difference between training dynamics of deeper networks and the shallower models.
There are still many open problems. Can we identify the hidden dynamics of the real world model that yields the parabolic converging trajectory? Can we theoretically understand the automatic pling of small entries as discussed in Appendix A.4? Is there a way to understand and leverage the fractal/bifurcatioo behavior in Section 5 toward global dynamics analysis?

---

Published as a conference paper at ICLR 2023
8 ACKNOWLEDGMENTS
This work is supported by NSF Award DMS-2031849, CCF-1845171 (CAREER), CCF- 1934964 (Tripods) and Sloan Research Fellowship.
REFERENCES
Kwangjun Ahn, Jingzhao Zhang, and Suvrit Sra. Understanding the unstable convergence of gradi- ent descent. arXiv preprint arXiv:2204.0100 2022.
Sanjeev Arora, Zhiyuan Li, and Kaifeng Lyu. Theoretical analysis of auto rate-tuning by batch normalization. arXiv preprint arXiv:1812.0399// 2018.
Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on edge of stability in deep learning. arXiv preprint arXiv:2205.09745, 2022.
Lei Chen and Joan Bruna. On gradient descent convergence beyond the edge of stability. arXiv preprint arXiv:2206.04 2022.
Jeremy M Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. arXiv preprint arXiv:2103.00055, 2021.
Alex Damian, Tengyu Ma, and Jason D Lee. Label noise sgd provably prefers flat global minimizers. Advances in Neural Information Processing Systems, 34:27449 224411, 2021.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 249-25696 JMLR Workshop and Conference Proceedings, 2010.
Stanislaw Jastrzebski, Maciej Szymezak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho, and Krzysztof Geras. The break-even point on optimization trajectories of deep neural networks. arXiv preprint arXiv:2002.09.. 02977700 2020.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.022/8, 2020.
Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after sgd reaches zero loss?-a math- ematical framework. arXiv preprint arXiv: 10.06914, 2021.
Zhiyuan Li, Srinadh Bhojanapalli, Manzil Zaheer, Sashank Reddi, and Sanjiv Kumar. Robust train- ing of neural networks using scale invariant architectures. In International Conference on Ma- chine Learning, pp. 22656-12684 PMLR, 2022a.
Zhouzi Li, Zixuan Wang, and Jian Li. Analyzing sharpness along gd trajectory: Progressive sharp- ening and edge of stability. arXiv preprint arXiv:2207.12673 2022b.
Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normal- ization layers: Sharpness reduction. arXiv preprint arXiv:2206.0708 2022.
Chao Ma, Lei Wu, and Lexing Ying. The multiscale structure of neural network loss functions: The effect on optimization and origin. arXiv preprint arXiv:2204.11322 2022.
Vardan Papyan. The full spectrum of deepnet hessians at scale: Dynamics with sgd training and sample size. arXiv preprint arXiv: 181 1000622, 2018.
Miguel Ruiz-Garcia, Ge Zhang, Samuel S Schoenholz, and Andrea J Liu. Tilting the playing field: Dynamical loss functions for machine learning. In International Conference on Machine Leam- ing, pp. 9157-9167. PMLR, 2021.
10

---

Published as a conference paper at ICLR 2023
Levent Sagun, Leon Bottou, and Yann LeCun. Eigenvalues of the hessian in deep learning: Singu- larity and beyond. arXiv preprint 1rXiv:1611.0777,, 2016.
Roman Vershynin. Hligh-dimensional probability: An introduction with applications in data science, volume 47, Cambridge university press, 2018.
Yuging Wang, Minshuo Chen, Tuo Zhao, and Molei Tao. Large learning rate tames homogeneity: Convergence and balancing effect. International Conference on Learning Representations, 2022.
Yikai Wu, Xingyu Zhu, Chenwei Wu, Annie Wang, and Rong Ge. Dissecting hessian: Understand- ing common structure of hessian neural networks. arXiv preprint arXiv:2010.04261, 2020.
Chen Xing, Devansh Arpit, Christos Tsirigotis, and Yoshua Bengio. A walk with sgd. arXiv preprint arXiv: 1802:0877780 2018.
Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael W Mahoney. Pyhessian: Neural networks through the lens of the hessian. In 2020 IEEE international conference on big data (Big data), pp. 581-590, IEEE, 2020.

---

Published as a conference paper at ICLR 2023
Supplementary Materials for Understanding Edge-of-Stability
Training Dynamics with a Minimalist Example
A ADDITIONAL EXPERIMENTS
In this section, we provide more empirical evidences supporting the main text.
In Appendix I. we will first introduce our experiment setup including the structure and initializa- tion for the neural network models as well as the generative model for the synthetic datasets.
In Appendix A.2, will present some additional figures demonstrating the training dynamics near the EoS minima for the degree-2 and degree-4 examples discussed in Section 3 and Section 4. We will also discuss a degree-3 example exhibiting different behavior around different EoS minima. We will explain the phenomenon using our understanding of the 2 and degree-4 models.
In Appendix A.3, we will provide additional empirical evidence that shallow neural networks usually does not converge to the exact EoS threshold.
In Appendix A.4, we will present some results on the training dynamics scalar networks without the coupling initialization. We will empirically show that the coupling of entries will arise along the training process.
In Appendix A.5, we will demonstrate the EoS phenomenon on the rank-1 factorization.
In Appendix A.6, we will introduce the experiment on learning real-world images (as presented in Section 6) in more detail. We will also present additional experiment results on networks with different activations and local trajectory with perturbation.
In Appendix A.7, we will present some experiments on the edge of stability phenomenon when the model optimized with stochastic gradient descent.
A.1 EXPERIMENT SETTINGS
A.1.1 CALCULATION NUMERICAL SHARPNESS
For the scalar network examples the closed-form Hessian is simple. We compute the exact parameter Hessian and use numerical packages compute its top eigenvalue.
For neural networks, we use the PyHess an package by (Yao al., 2020), which compute the top eigenvector eigenvalue pair by inferencing the Hessian vector product and do power iteration. For all numerical sharpness computed for neural networks, we set tol=le--eeeee and max.iter=10000.
A.1.2 SYNTHETIC EXPERIMENTS
For all experiments involving neural networks on synthetic datasets (Fig. 1a, Fig. 1b, Fig. 15), we use fully connected networks with the same dimension for input, output, and all hidden layers. The bias of all layers are fixed Formally, L-layer width d network can be modeled by f Rdd Rd such that for input vector E Rd.
f(x) =
(13)
where a d + Rd some entry-wise activation function and W, E Raxd for all E [L]. For this paper only considered a being the ReLU activation ozz) z1,20 or the identity oxx)
Initialization of Neural Networks
We use Xavier initialization (Glorot & Bengio, 2010) with gain of I to initialize the all weight matri- ces. For shallow two-layer networks that will not enter the EoS regime if using completely random initialization, we will asymmetrically re-scale the layers after random initialization by multiplying constant to all entries of the same layer. When we present results for the re-scaled experiments, we will state the re-scale factor.
12

---

Published as a conference paper at ICLR 2023
Synthetic Dataset and Loss Function
For the experiments involving neural networks, we use synthetic datasets very similar to the linear network experiment in section L.3 of Cohen et al. (2021). For neural network described above with dimension d, we consider the problem of mapping n inputs 21,..., &n E I to n outputs 91.....; Yn E Rd Let X Rdxd and Y d denote the vertically stack imputs and outputs respectively (Here n = d). We generate X as a whitened matrix such that XX = dIa and generate Y by Y X A where A diag(1, - ..... z)
For all experiments with neural networks on synthetic datasets, we consider the simple squared loss
IIf(;)) -3.122
(14)
1=1
which matches our analysis on scalar networks.
A.1.3 WOALDWORLD DATA EXPERIMENTS
Here we provide the detailed setting for the experiment results shown in Section 6 in the main text as well as Appendix A.6. We consider binary classification problem on subset of CIFAR-10 image classification dataset (Krizhevsky et al., 2009).
Dataset
To study the training process in an werrpaarmettrizee setting (in which the loss can converge close 0), we take a binary 50-sample subset from CIFAR-10 containing the first 25 samples class (airplane) and class (automobile). Then we label samples from class O by -1 and samples from class by +1.
Here we are consider binary classification problem since for networks with output dimension larger than 2, there typically not strong eigengap between the first eigenvalue and the other eigenval- ues (Sagun et al., 2016; Papyan, 2018; Wu et al., 2020). The dynamics with multiple eigenval- ues around the stability threshold may exhibits cascading oscillation along different eigendirections Cruiz-Garcia et al., 2021), ),, and could be complicated to analyze.
Network Structure
We conduct the experiment fully-connected neural networks with four hidden layers of width 200. We consider tanh and ELU as activations. In Table we provide the structure of fully-connected ELU-activated architecture. This architecture follows the experiments in Li al. (2022b).
Table 1: Structure of fully-connected network
#
Name
Module
In Shape
Out Shape
1
Flatten()
(32,32.3)
3072
2
fcl
nn Linear (3072, 200, bias=False)
3072
200
3
nn. ()
200
200
4
fc2
nn. Linear (200, 200, bias=False)
200
200
5
nn. nn.ELU ()
200
200
fc3
Ln.Linea (200, 200, bias=False)
200
200
7
nn. En.ELU ()
200
200
8
fc4
In.eine (200, 200, bias-False)
200
200
9
nn. En. In.Linear (200, 1, bias=False)
200
200
10
fc5
200
Loss Function For the CIFAR- 10 subset experiment ((#1, where 50, m, 3072 and
Y E -13, we consider the mean squared loss
((..))) 1ill
(15)
which matches theoretical analysis the scalar example.
13

---

Published as a conference paper at ICLR 2023
A.2 ADDITIONAL EXPERIMENTS FOR SCALAR NETWORK EXAMPLES
In this section we show some additional figures demonstrating the local training dynamics and con- vergence boundary for the cases analyzed in Section and Section 4.
A.2.1
4-LAYER SCALAR NETWORK
Init: 0.4001)
1-Step Trajectory (400 Iterations)
(2.2: 0.3)
1150
Init:(2.35,
10-4
Shar pRsS Contons
0.40
Threshold (2/v)
(2.6, 0.2001)
0.35
Slarnnss
16-
Inir: 0.5)
0.30
5-
X
(2.5, 0.4001)
X
(2.2, 0.3)
X
[235, 0.5)
Minima
Global Minima
0.25
300
400
2.25
2.36
2.40
2.45
2,00
Iterations
2
Figure 8: Sharpness concentration for the degree-4 example. We run GD with 0.2 from initializations that are above, below, and very close the line of global minima. The sharpness of all trajectories converges to around 2/n while the loss decreases exponentially. In the 2D trajectory, the 2-step movement quickly converges the parabolic curves ending very close to the EoS minimum.
Normalized Two-Step Update
EoS Minima
Global Minima
0.55
Converging Sharpness
0.15
7
0.40
6
0.35
0.30
1.9
2.0
2.1
2.2
2.3
2.4
2.5
Figure 9: Local two-step movement for the degree-4 example. We record the two-step movement with 7 0.2 from a grid of initializations near the EoS minima. At each point, the arrow points toward the direction of two-step movement from that point and the color of arrow indicates the converging sharpness of trajectories passing that point. We can see the parabolic trajectory and how initialization to the right of the EoS minima all tend to converge to it.
14

---

Published as a conference paper at ICLR 2023
4.0
4.0
EaS Minima
Ec8 Minima
3.5
Global Mininut
3.5
Global Minima
EoS
EoS
3.0
Diverge
3.0
Diverge
(Iiitiaizaaionn)
2.5
(Initialization
2.5
2.0
2.0
1.5
1.5
0.5
0.5
(Initialiaation)
or (Initialiaation)
4.0
4.2
5.0
Converging Sharpness (50000 iters)
Converging Sharpness (50000 iters)
(a) 0.3
(b) 1
Figure 10: Converging Sharpness of Initializations Under Different Step Sizes. Please see the caption of Fig. 2b for detailed description.
15

---

Published as a conference paper at ICLR 2023
A.2.2 2-LAYER SCALAR NETWORK
10-1
1-Step Trajoctory (100 Iterations)
10-4-
10"
42
1181
10
Init-(3.5,
0.6
Sharpane
EoS Threshold (2/6)
(8.2 (125)
0.4
12
Init (3.5,
9
10
X
13.2. 0.25)
X
[3.35, 0.25)
X
(3.5. 0.25
Minina
Minima
Iterations
Figure 11: Beyond EoS training on product of two scalars. We run the identical experiment as in Fig. 8 except for the objective minz, x.y) Note that in this case do observe both sharpness concentration and sharpness adaptivity. The two-step trajectory follows an elliptical trajectory centered the EoS minima. In this context, sharper initialization (e.g. the purple curve) will eventually converge to flatter minima (as shown in the left figure).
Normalized Twoo-Step Update
EoS Minima
Global Mininn
-9
0.8
Converging Sharppess
2.6
2.8
3.0
3..2
3.4
3.6
Figure 12: Local two-step movement for the degree-2 example. This is the same figure Fig. except for the degree-2 example. We use the same step size 77 = 0.2. There is no longer the concentration behavior to the EoS minimum will converge to a sharpness near the stability threshold (which is 10 with ? 0.2).
for the degree-4 case. Locally, only the initialization very close
16

---

Published as a conference paper at ICLR 2023
A.2.3 3-LAYER SCALAR NETWORK
Now we look into an interesting example with different behaviors around different EoS minima.
We consider a 3-layer scalar network with objective
min 2 -xyzyy) w.N.ERR
(16)
To make the dynamics two dimensional, we consider the initialization with y. The equality of the last two entries will be preserved through training the dynamics two dimensional in terms of x and y. In the positive quadrant, the global minima is vzy = and there are two EoS minima.
In Fig. 13, we plot the converging sharpness from different initializations in comparison with Fig. 2b and Fig. 5b in the main text. Around the EoS minimum that the single entry æ is small and the duplicated entries y are large (upper left of Fig. 13), the behavior is similar to the 2 scalar (Fig. 5b) with no sharpness concentration. Around the EoS minima with large single entry and small duplicating entries (lower right of Fig. 13), we have a region of initialization (the red shaded area) with sharpness concentration similar to the 4 scalar case (Fig. 2b).
2.5
EoS Vinima
Global Minima
EoS
2,0
Divergo
(Initialization
1.5
1.0
0.0
0.5
1.0
1.5
22.0
2.5
3.0
3.5
V (Initialization)
Converging Sharpness (50000 iters)
Figure 13: Converging sharpness of (ar, y. parameterized initializations (17 0.2).
heuristic explanation this difference lies in the difference in the degree of the small entries. Around the EoS minima that the single entry is small, the local two-step approximation is similar to Eq. (12) and gives us elliptical two-step trajectories. Around the minima with small duplicating entries, the two-step approximation would contain the cubic term in Eq. (7), which gives both sharpness concentration and adaptivity.
Such heuristics can be further verified by visualizing the local dynamics around the minima. As shown in Fig. 14, the local dynamics around the minima with small duplicating entries is similar to the of the 4 example with convergence toward parabolic trajectory. On the other hand, the local dynamics around the minima with only one small entry is similar to the case of the 2 example where parameters follow an locally elliptic trajectory centered at the EoS minima.
17

---

Published as a conference paper at ICLR 2023
Normalized Stwp Update (FoS)
Normalized Two-Step Update (No EoS)
10
EaS Minima
FoS Minima
Global Minina
-9
0.30
Coaverging Sharpuess
0.45
Convenging Shar
-7
0.40
-G
0.35
1.6
0.30
1.5
2.15
2.20
2.25
2:30 2:35
0.3
0.4
0.5 0.6
58 0.7
0.8
V
Figure 14: Local two-step movement for the degree-3 example. In this figure, we show the local dynamics the two EoS minima in the positive quadrant of the degree-3 example. The left figure corresponds to the EoS minima at the lower right of Fig. 13. At this EoS minima, the duplicated entry y small, and the local behavior is very similar to the case of degree-4 example (Fig. 9) for which we have provable sharpness concentration. The right figure corresponds to the EoS minima the top left of Fig. 13. The local behavior very similar the case of degree-2 example (Fig. 12), for which we do not have EoS behaviors.
A.3 ADDITIONAL EXPERIMENTS FOR 2- 2-COMPONENT SCALAR FACTORIZATION
In this section we present the experiment results for 2-component scalar factorization deferred from Section 4. The dynamics as shown in Fig. 15 is very similar the degree-2 example in Fig. 11.
10
Loss 4.5
10-1
10 9
Sharpness
EoS Threshold (2/n)
Sharpness
45
EoS Threshold (2/1)
40
40
35
35
10
20
30
40
5
10
15
20
25
30
Iterations
Iterations
(a) min a wcdd (1
-x y)
(b) min DERd
La zy
Figure 15: GD on 2-component scalar factorization problems. We run gradient descent with 1) 0.05 for two different objectives. Both models are asymmetrically initialized with factor (5, 0.1) so that they have an initial sharpness larger than 2/1). For both cases, the converging sharpness is distinguishably smaller than the stability threshold.
18

---

Published as a conference paper at ICLR 2023
A.4 EXPERIMENTS FOR GENERAL SCALAR NETWORKS
In this section, will present some empirical observations on training dynamics of more general scalar networks related to the sharpness concentration and adaptation phenomena. A n-layer scalar network is defined to be the model parameterized by entries 21, .....,.......................... 2n. E R with objective
f(x1....,
T
(17)
2
A.4.1 INITIALIZATION WITHOUT DUPLICATED ENTRIES
We first consider variant of the degree-3 3 example as discussed in Appendix A.2.3. In particular, we initialize the two small entries differently and record their values throughout the training trajectory.
10
10a
18
EuS Threshold 2/1
10
16
10-2-
10
10-4
Sharpness
12
(2)
10
10
[
10
10
10
2000
1000
2000
1000
50
100
150
200
Iterations
Iterations
Iterations
This figure records the training loss (left) and sharpness (middle) of a degree-3 scalar network with initialization (6. 0,1, 0.4) optimized by gradient descent with fixed step size y) 0.2. In (right) we plot the distance of the last entry to other entries.
Figure 16:
As we can in Fig. 16, the very beginning of the training, the second entry converges to 33 geometrically, then the dynamics reduced to the case of duplicated entries, which we know that the sharpness concentration behavior would happen for sufficiently asymmetric initialization. In Fig. 17 we consider a 7-layer scalar network with different large entries and different small entries. We observe similar behavior 24, Øz. Zg all converges 2t7 geometrically, and we observe concentration of sharpness with 4 duplicated small entries.
10
35
EuS Threshold 2/q
10a
30
10
10-1
Sharpness (X)
22 27
Loss
10
20
10-2
8g 27
24 ory
15
10
10
90
5-
10
2000
4000
2000
4000
100
150
200
Iterations
Iterations
Iterations
Figure 17: This figure records the training loss (left) and sharpness (middle) of 7-layer scalar network with initialization (2, 2.5, 3, 0.2, 0.3, 0.4) optimized by gradient descent with fixed step size 77 0.2. In (right) plot the distance of the last entry 27 to other entries.
19

---

Published as a conference paper at ICLR 2023
To probe into the detailed training dynamics of general scalar networks, plot the pairwise dy- namics of the entries as shown in Fig. 18.
30
8 27
24
1.75
2,00
2/2
0.-
0.5
0.5
0.1
0.8
0.1
1.75
3.00
3.50
2.4
2.6
3.0
272
373
0.7
0.7
0.5
0.4
0.4
0.6
0.7
272
J7g
Figure 18: Pairwise Training Dynamics. This set of figures record the pairwise training dynamics for entries 22, 43, #4, 25 in the same experiment as Fig. 17. æ22 and 2g are initialized large while 24 and are initialized small. We see that within the small entries and the large entries, the pairwise dynamics are all approximately linear while the cross comparisons across the small entries and large entries gives the parabolic two-step trajectory (and also some bifurcation behavior).
07
30
23
28
27
28
0.
24
0.1
100
1007
Eterations
Ieerationas
Figure 19: Single Entry Movement. We plot the value of 23 (initialized to 3) and ø3 (initialized to 0.1) along the training process. The larger entry decreases in approximately monotone manner and the small entry increase while oscillating.
A.4.2 ON "LARGE" AND "SMALL" INITIALIZATIONS
In the experiments shown above, we have seen that there are mainly two classes of behaviors for the entries: the entries that were initialized to be large moves slowly with little oscillation while the entries that were initialized to be small has significant oscillation along the trajectory. Intuitively, is the decreasing large entry that decreases the sharpness and stabilizes the oscillating small entries and result in the final convergence close to the EoS minimum. A natural question to ask is whether there exists clear boundary separating the "small" and "large" entries.
20

---

Published as a conference paper at ICLR 2023
We consider 4-layer scalar network with initialization 06, 0.3, 0.2) optimized with GD with step size 17 = 0.2. In Fig. 20 and Fig. 21 we visualize the training loss, sharpness, and the pairwise dynamics. The mixed behaviors suggests clear boundary between the "large" and "small" entries does not exists, and the complexity of this problem is beyond this simple heuristics.
10-1
EoS Threshold 2/9
10-3
100
Sharpness (A)
10-3
10-1
10
6-
10"2
10-11
4-
2000
4000
2000
4000
50
100
150
200
Iterations
Tterations
Iterations
This figure records the training loss (left) and sharpness (middle) of a 4-layer scalar network with initialization (6, 0.7, 0.3, 0.2) optimized by gradient descent with fixed step size 1) 0.2. In (right) we plot the distance of the last entry 24 to other entries. In this example, the small entries did not converge be exactly the same value, yet the loss still decreased geometrically and the sharpness concentration phenomenon still occurred.
Figure 20:
0.80
0.78
0.05
0.00
5.00
1.36
6.00
21
0.45
0.40
0.00
0.35
0.80
0.31
0.5
0.5
0.5
0.5
0.1
0.3
5.00
3.35
6.00
0.6
0.7
0.6
0.3
Figure 21: Pairwise Training Dynamics. This set of figures record the pairwise training dynamics for entries T1, X3, in the same experiment as Fig. 20. was initialized large (at 6), æ3, were initialized small (at 0.2, 0.3), and 22 was initialized moderately small (at 0.7). We see that the two-step pairwise dynamics between Ø2 and 23 roughly follows parabolic trajectory, yet the pairwise dynamics between Ø1 and D2 also exhibits similar features while still following a roughly linear relation.
21

---

Published as a conference paper at ICLR 2023
A.5 ADDITIONAL EXPERIMENTS FOR RANK- FACTORIZATION OF ISOTROPIC MATRIX
In this section, we will demonstrate the EoS phenomenon on the rank-1 factorization of isotropic matrix in Section 3.3.
We first show that the loss, the sharpness and the trajectory of GD is very similar to the degree-4 scalar network case. For each different learning rate, the sharpness concentrates to a tiny interval close to the stability threshold 2/17 when trained with gradient descent.
Also, we consider the 2D trajectory of the two vectors æ. y. We plot the the trajectory in the norm of each vector, i.c. || ||lll//ll, and get a similar figure to the scalar case. Actually, we can prove that the dynamics of this training objective will eventually be reduced to our degree-4 scalar network. That is because all the global minima of this optimization problem requires that 2 is aligned with y, i.e. x cy. After the two vectors are aligned, the training dynamics of || |ly/| will be exactly equivalent to those of the scalar network.
The following figure shows how GD enters BoS the rank-1 factorization problem, and how fast the alignment of the two vectors is achieved. Here we consider an alignment indicator ||>/||||yy| (x y)² showing how the two vectors are aligned. If æ is parallel y, i.c. æ = cy, then the variable becomes O. Detailed analysis for this problem is deferred to Appendix C.
1-Step Thrjectory (1030 Iteration))
53
X Dxtiiieriee Tralener 220
13.1
(a) Evolution of training loss sharpness and the trajectory of GD rank-1 factorization of isotropic matrix.
(b) Evolution of the alignment
indicator ||>|? (a" y)²
Figure 22: EoS phenomenon the rank-1 factorization of isotropic matrix. In (a), similar to the degree-4 scalar case, we demonstrate sharpness adaptivity by running GD with learning rate 7 - TO 13 from the same initialization (1 1/2). All the sharpness of each trajectory converges to around their corresponding stability threshold 2/1 while the loss decreases exponen- tially. In the 2D trajectory, GD quickly converges along some smooth curves ending near the EoS minimum. In (b). we see that in the first 30 iterations, the alignment indicator |||=||²||| |2]| (a ")²) decreases geometrically, and stabilizes at its numerical minimal value. Theoretically, we character- izes the geometric decay for the alignment variable by Lemma 20 in Appendix C.
A.6 EOS CONVERGENCE FOR DEEP NEURAL NETWORKS TRAINED ON REAL DATA
In this section, we present a more comprehensive description and additional results for the exper- iment of learning 2-class small subset of CIFAR-10 with 50 images in pver-parameterized regression setting. The details for the network structures and dataset construction are available in Appendix A.1.3.
In this experiment, we train two 5-layer fully connected (fc) networks of width 200 with ELU and tanh activation using (full-batch) gradient descent the binary dataset with mean squared loss. We chose these two activation functions following the empirical experiments in Cohen et al. (2021). ReLU is not being used since its training dynamics as the loss converges to 0 is very unstable.
We record the training loss and sharpness of the two training processes. To better visualize the training trajectory, we consider the following projection mechanism.
22

---

Published as a conference paper at ICLR 2023
A.6.1 TRAJECTORY PROJECTION
Inspired by the observation on the scalar example, note that the dynamics toward the end of the convergence has two distinctive directions: an "oscillation direction" which is aligned with the first eigenvector of the Hessian, and "movement direction" which the 2-step average of the model moves along and converges to the final minimum.
In the context of our (a, re-reparameterization for theoretical analysis (Definition 2), the oscillation direction corresponds to b and the movement direction corresponds to a. In local region around the converging minima, (a. b) constitutes a parabolic trajectory that can be well captured by the solution of the ODE in Eq. (8). In a high-dimensions top eigenvector at minimum, but we have to manually pick movement direction to project onto.
setting, the oscillation direction is still naturally the
To be concrete, consider trajectory of the parameters 101,02..... -r parameter vector for the model after the t-th iteration. We define the oscillation direction Vosc as the first eigenvector of the parameter Hessian H (0r) and the movement direction Unove (f) as follows:
fr), where 0g E IR' is the
Definition 3 (Movement direction). For some iteration i E [T). define Unove (f)
move ( - (0p-1 + 01)
+ Or)).
(18)
Fix some iteration É, Umooe (f) captures the non-oscillatory movement of the parameters from step 8 to step T. We orthonormalize the basis by projecting off from Umove (1) and get
T Umove (1)
Umooe (t) e (Vosc);
Umove
Omove (f)/ @oooe (i)l
(19)
Tose
Uoso
|ose
Now with the orthonormal basis, we define the novement-oscillaaiin projection of 0g to be the projection of its offset from the minima (which we approximate by the mean of the last two steps in the trajectory) onto Umove (f) and Tose.
Definition 4 (Movement-Ossillaaee Projection). Fix iteration i for determining the movement direction, the movement--oscillation projection of 0g is
(enove(t) (t)T (0: (0r-1+ + 0r)) T (0:
(20)
When doing the projection in practice (as in Fig. 24 and Fig. 26), fix f 5000, which when the 2-step trajectory becomes relatively stable. We also record the norm of the component of the offset 0, (0r-1 + fxr) that is orthogonal to the subspace spanned by Umove (f) and Vorc- These results are shown in Fig. 23c and Fig. 25c.
A.6.2 ELU-ACTIVATED FULLY CONNECTED NETWORK
Here we present the experiment results for training a 5-layer ELU-activated FC network on the binary subset of CIFAR-10. In Fig. 23, we show the evolution of loss and sharpness along the training process. The sharpness eventually converge to just slightly below the 2/n threshold. We also observe that the dynamics toward the end of the converging process is mainly happening in the 2-dimensional subspace spanned by the oscillation and movement directions (Fig. 23c).
In Fig. 24 we plot the projected trajectory of the training process. Toward the end of the training be very accurately characterized by parabola and the converging sharp-
process, the trajectory ness is just slightly below the stability threshold. This is identical to what we observe (and proved) for the scalar network case.
23

---

Published as a conference paper at ICLR 2023
10
2/20
0.000110-
3- 3/200
141
EL000000
08
Heddbad of PPujecko
214
204775
Sharpoos (A)
Starpeee (X)
221120
200
2141
300
100.76
198
100.00
100.00
7,700
7/730
00-
10010
5000
10000
10000
5000
9000
9303
5000
10000
Eterations
Eerrations
Iterations
(a) Training loss and sharpness. (iteration 0- 18500)
(b) Training loss and sharpness. (iteration 5000-5030)
(c) Residual orthogonal to oscillation movement subspace.
Figure 23: Training Statistics for ELU-activated 5-layer FC network. (1) = 0.01, 2/1 200) (a) identical to Fig. (left) in the main text. We see that the model is capable memorizing all data as the loss decreases exponentially to O. Toward convergence, the sharpness oscillates very close to the stability threshold and eventually converges to 199.97. (b) we show section of (a) between iteration 5000 and 5030. We can clearly observe two distinetive features of the EoS regime: the loss decreases non-monotonically and the sharpness oscillates around 2/17. In (c) we plot the norm of the offset from minima that orthogonal the movement-oscillation projection. After 3000 iterations the residual becomes very small, suggesting that dynamics is mainly happening in the 2 dimensional subspace and hence the projection captures the dynamics quite well.
Projected Trajectory (Toeration 18500)
Projectod Trajectory (Iteaation 5000 18500)
Minico
2900y"
330
005
200
Ocillation Disetba
000
150
200 0G
200 00
0.000)
-0-05
100 DC
1.30
-0 -0,0005
210 M0
-0.99
-0,0000
-00
Monrment
(a) Projected Trajectory. (iteration 0- (-18500)
(b) Projected Trajectory. (iteration 5000- 5000-18500)
Figure 24: Projected Trajectory ELU-activated 5-layer FC network. (1 0.01, 200) In (a), plot the projected trajectory for the entire training process. After some large bifurcation like oscillation, the 2-step trajectory quickly stabilizes and moves toward the minimum along the movement direction. In (b), we show the tip of the trajectory, which can be very well captured by a parabola. These figures are identical to Fig. the main text. The color of the dots reflects the local numerical sharpness.
24

---

Published as a conference paper at ICLR 2023
A.6.3 TANH-ACTIVATED FULLY CONNECTED NETWORK
Here we show the results for the same experiment on tanh-activated 5-layer FC network. The phenomena are qualitatively identical to the ELU case described above.
x100
10
2/30
101
1011
101
Booklnd of Projectonn
201.00
0.-
237.
200.00
Staryeins (A)
809.5-
0.4
Shappone
20330)
137.5-
099.70
0.2
38540
e 272M)
19925
2:30
00-
1300
10300
5000
13000
5301
5003
1000
5005
10000
Berations
Iterations
(a) Training loss and sharpness. (iteration O- 0-50)00)
(b) Training loss and sharpness. (iteration 5000-5030)
(c) Residual orthogonal to
sullaation-oovmennt oscillation movement subspace.
Figure 25: Training Statistics for tanh-activated 5-layer FC network. (17 = 0.01, 2/n = 200) Please refer the caption of Fig. 23 for detailed explanation.
Projected Trajectory (lteaaton 0 18500)
Minzen
Projected Trajestory [Iveration 18500)
0.33
200
2004
Oxilletoo Doctim
0.3661
0000
-0 30.3005
-03000
0.94
30
041
0.1
0.931
0.004
Movement
Dinvction
Mosemmee
Direction
(a) Projected trajectory. (iteration 0-18500)
(b) Projected trajectory. (iteration 5000-5030)
Figure 26: Projected Trajectory of tanh-activated 5-layer FC network. (1 = 0.01, 2/1 200) Please refer the caption of Fig. 24 for detailed explanation.
A.7 EDGE OF STABILITY AND STOCHASTIC GRADIENT DESCENT
In this section, we will briefly discuss some empirical observations of EoS in stochastic gradient descent (SGD). In Appendix A.7.1, we will first empirically present the effects of different forms of noise made on real world models trained with mini- batch gradient descent. Finally, we will discuss the limitations of our scalar model in explaining what people observe about EoS when the model is trained with SGD.
scalar model. Then in Appendix A.7.2 will compare with the observations
A.7.1 GD WITH NOISE ON SCALAR NETWORK
We first look into the training trajectory of our degree-4 scalar network example with noise injected to the gradient descent process. We consider label noise, which perturbs the target by small amount per iteration, and gradient noise, which perturbs the gradient by a small amount per iteration before updating the parameter according to it.
25

---

Published as a conference paper at ICLR 2023
Label Noise:
To simulate the existence of label noise, at each iteration we compute the gradient for the objective
[N y, 8) = [1+
+ 8 z²
where 5 is sampled from a zero-mean Gaussian AV(0,, o²) for each iteration. This is equivalent to adding a perturbation of 5 to the label (which is 1 in our original model). We start from the same initialization as in Fig. 8 and plot the trajectory in Fig. 27.
As shown in Fig. 27a and Fig. 27b, the trajectory first roughly follows parabolic boundary and reaches close to the set of global minima near the EoS- --S-mmmimmm relatively quickly (for around 200 iterations). This part of the trajectory resembles our analysis for the case without label noise.
After the model reaches the tip of the parabola, the dynamics is mainly dominated by the label noise. The gradient dominated by its noise component of özy(y, z), which is orthogonal to the manifold of global minima zy = 1, thus the model starts oscillating around the global minima. As shown in Fig. 27c and Fig. 27d, the sharpness further decreases very slowly (for 10 iterations) and eventually reaches the flattest global minimum at (1, 1) with sharpness of 4. We believe this is within the regime of the sharpness reduction flow near the manifold of minima, which is comprehensively studied by
Damian et al. (2021); Li et al. (2021; 2022b); Lyu al. (2022).
1-Step Trajentory (1000 Iterattons)
Ise (10, (.)1)
10-1
1000
0.45
10'1
Ead Thmaadd (2/0
0.35
12 5, 6:4)
Misiza
0.5
Micias
X Ota, (al)
(a) Training loss and sharpness. (iteration 0-1000)
(b) Training trajectory. (iteration 0-1000)
1-Step Trajectory (100000 leccatinns)
10-
Eab Mistaa
10
Clobel Misies
08-
X
(1.5, (:41)
16
05-
10
Slarpae
(360
Shapppo
0.0
12
1.4
1.5
Iecrationn
(c) Training loss and sharpness. (iteration 0-1000000)
(d) Training trajectory. (iteration 0-1000000)
Figure 27: Training Trajectory of Degrec-4 Model with Label Noise. (1) = 0.2)
We plot the loss, sharpness, and training trajectory of the label noise model with a 0.01. The label noise model first follows trajectory similar to the original GD training trajectory and reaches near the EoS minimum as shown in (a, b). Then it follows the sharpness reduction flow along the manifold of minima and reaches the flattest minima as shown in (c, d).
26

---

Published as a conference paper at ICLR 2023
Gradient Noise:
For gradient noise model, we sample a perturbation of (6)8") from 2-dimensional spherical Gaus- sian (0, o² I2) at each iteration and apply this perturbation to the gradient before we update the parameter. The one-step dynamics with gradient noise (6. 8') is then:
1++1 = x4 - n( (x,n² 1) + +), Vi+1 = yt n( - 1) + 8)).
With gradient noise, the initial parabolic trajectory can still be observed as shown in Fig. 28b. As the model reaches close to the manifold of global minima near the tip of the parabola, it no longer follows a monotone sharpness reduction flow (as in Fig. 27c for the label noise case) but instead randomly oscillates along the manifold of global minima between the two EoS minima. We believe this is due to the component of the gradient noise parallel to the minima manifold, which dominates the sharpness reduction effect.
1-Step Trajectary (1000 Iterations)
I4k a.l))
10'4.
0.5
225, 0.61
10'1
10-1-
0.43
1011
0.46
Thomiedl
0.35
Stampoe
15-
0.32
0.1
200
1090
10
2.0
Iranationn
(a) Training loss and sharpness. (iteration 0-500000)
(b) Training trajectory. (iteration 0-500000)
1-Step Trajertory (500000 Iterations)
S Mlains
Glénd Miniza (25,0.4)
L00
1.2-
Fid Tlackk P/0
(2.. Mal)
0.55
800
Demnidess
(c) Training loss and sharpness. (iteration 0-500000)
(d) Training trajectory. (iteration 0-500000)
Figure 28: Training Trajectory of Degrec-4 Model with Gradient Noise. (1) - 0.2)
We plot the loss, sharpness, and training trajectory of the gradient noise model with o = 0.01. Like the label noise model, the gradient noise model first follows a trajectory similar to the original GD training trajectory and reaches near the EoS minimum as shown in (a, b). After getting around the FoS-minimum, the parameter begins to randomly oscillate and traverse around the manifold of global minima between the two EoS minima as shown in (c,d). It is likely that the gradient noise finally converges to some distribution along the manifold of global minima.
27

---

Published as a conference paper at ICLR 2023
A.7.2 MINIBATCH SGD FOR OVER-OVRAMEERR PARAMETERIZED MODELS
In this section, we empirically investigate what happens to the converging sharpness when over- parameterized network are trained with minibatch SGD. We use the same 5-layer FC models and dataset as used in Section 6, Other than full-batch gradient descent, we also train the models with mini-bateh gradient descent with varying batchsizes and record their converging sharpness.
2/v- (X)
2/v dl1
Cuperrgaa Sharpuse
300
2/v
Oanerning Sharpees
2/v
230
200
150
300
100
Barchsize
(a) 5-layer FC network with ELU activation
(b) 5 5-layer FC network with tanh activation
Figure 29: FC networks trained with SGD of varying batchsizes. 0.005, 0.01, 0.02)
For each learning rate and batchsize, we train 10 models from different random initialization for 20000 epochs and record their converging sharpness. The standard deviation of sharpness is repre- sented by the shaded area. (The loss of all models converges to lower than 10-*). When the batch size of SGD is large, the converging sharpness close to 2/1,, which very similar to the gradient descent cases. On the other hand, the converging sharpness is significantly lower when the batch size small compared with the dataset. worth noting that the converging sharpness for each batch size is quite concentrated.
Instead of going to the flattest minima (as the label noise case) or randomly oscillating below the EoS threshold (as in the gradient noise case), the converging minima of overparameterized deep networks trained with mini-batch SGD have highly concentrated sharpness that correlated to the batch size.
We note that a key difference between the werrparraeterriee mini-batch SGD and the noisy GD experiment discussed in Appendix A.7.1 is that the loss for mini-batch SGD converge to fixed point with loss 0 (i.c. the model can memorize all training data) while the models with fixed additive noise will not converge to fixed point. Currently the minimalist scalar model example analyzed can only memorize one data point. We believe is an interesting future direction to generalize the model to memorize more data and understand why mini-batch SGD converges below the EoS threshold.
28

---

Published as a conference paper at ICLR 2023
B THEORETICAL ANALYSIS ON THE DEGREE-4 EXAMPLE
In this section we present the complete rigorous analysis on the training dynamics of the degree-4 example discussed in Section 3. The section will be organized in the following way:
In Appendix B.1 we will first define the problem and two reparameterizations we used for analysis, this serves as a more comprehensive version of Section 2 in the main text.
In Appendix B.2 we will first restate our main theorem for the degree-4 example (Theorem 3. along with two corollaries (Corollary 3.1, and Corollary 3.2). These theoretical results characterized the sharpness concentration and sharpness adaptivity phenomenon of the degree-4 example. Then we will provide more comprehensive proof sketch for Theorem 3.1 that is similar to the discussion in Section 3.2 of the main text.
Then we will provide the lemmas for dynamics approximation (Appendix B.3), phase I convergence (Appendix B.4), and phase II convergence (Appendix B.5).
Finally, in Appendix B.6 we will use the lemmas to complete the proof for the main theorem along with its corollaries.
B.1 PRELIMINARIES
We consider a simple objective function C(x, y, (xyzz - 1)?. Denote y == ryzw, then
VL(I,,, = (7² - 1),,, V. w]-
(21)
2 1/22
(2r² - /
(272 7)/x2 1z2
(272 7) rr
(27 7) ry
/2/ /v² (2-² -
(272 7) 7)
(2 2 7) you
(0,,,,,,, (T2 7)
(272 - - 7 /2
721
(272
zw
(272 - 7//
(272 - 7)/
(2 - / 2m
72
(22)
Let the parameter [z, y, 2, w| to be optimized by gradient descent with step size 17, that
[1(441): 1(++1); 1) 1({+ (m<ce)> ((t)) =(t), w(o)] 17 () U(e): 2(1) w(e))-
(23)
To further simplify the problem, we consider the symmetric initialization of 2g = To. w = yo- Note that due to symmetry of objective, the identical entries will remain identical throughout the training process, so the training dynamics reduces to two dimensional, and the global minima is simply S = {(1, y) TR. x²y² 1}. Computing the closed-form of the gradient, know the 1-step update of x and y follows
It+1 It- F.y nl (1312 1), V/+1 = y/ - =n-nyynn=>i 1).
(24)
Denote = xy, the parameter Hessian of the objective C at (Jr, v,,,,) admits eigenvalues λ = 2² 7), '2 y2(1 7) and
3 d 1/2 (1. + 5²) (37² 1) V (x2 3+) 4 10,
(25)
\ } (x² + y²) (37² - 1) + v 2x2 (1 (1-33221 22 -0 774)
2
When (z. y) converges any minimum, و y² 1, so \1 A2, As all vanishes. Therefore it is 1
A that corresponds to the EoS phenomenon people observe. When < 2. solving d = 2/n with
x: x²²² = 1 gives
+ (-
11)
(26)
y -= +1
and their multiplicative inverses. These solutions correspond to the minima with sharpness exactly equal to the EoS threshold of 2/17. Since they are all symmetric with each other, without loss of generality we pick the minimum of interest as
(z. Φ) e 2 (1-4+
2) +
- -).).
(27)
29

---

Published as a conference paper at ICLR 2023
To better analyze the dynamics under natural coordinate, consider the reparameterizzation that For any (I. y) E {(I, y) E R
: y), define
(x² - , d 2 xy.
(28)
This gives a bijective continuous mapping between {(x, v) E IRt x RR= x > y) and {(c, d) E R+ x RR} Intuitively,
taking the lower half of y 1/2 on the positive quadrant d 1.
With c, as defined, the 7-EoS minimum simplifies to (č, d) ((n -4)a 1). The inverse map can be computed as
V 2 V 4d
V2d
V
(29)
V2
2 2 ve+v++ 47
To expand the dynamics near the 1--EoS minimum, we define
c- 2 4) b2d-1
(30)
to the offset from (2, d). Our amalysis will primarily be using the (a, pa-aammtttrrzation. To summarize, the (c, d) and (a, b) reparameterization of (ar, v) respectively given by
(c, a ((x² (x)~ myy) (a, 52) - (0)- -4)
-4)= .zy
(31)
Let VT, under the reparameterization Eq. (24) becomes.
a++1 = (x 1-4--43- + (a + 4-- 4) (1 1 101+ - (1+ bi)) x²))
4)
b1+1 = b, + ((1 + b,)3 2(1 + b,)B + (1 + )r)
(32)
+ ((1+ +
- - (1) (1+ (4(1 + b,)3%e + (aik + (1 - (1x4)))))))
B.2 THEORETICAL RESULTS AND PROOF SKETCH
Now with the reparamcterization defined, we restate our convergence result on the 4 scalar objective and discuss the proof sketch.
Theorem 3.1 (Sharpness Concentration). For large enough absolute constant K, suppose
< 2000 / K-1, and the initialization (o),bo) satisfies ao E (12k² -K-2==1) and bo E (-K Mo0.. Consider the GD trajectory characterized in Eq. (6) with fixed step size x² from (an, bo). for any > 0 there exists T = O(K-KKK-00 x + log(=-1) log((bol-1)x) >)xx=) -B such that for all T. [b.] < e and ag E (-5
e 5 10 x²).
B.2.1 PROOF SKETCH OF THEOREM 3.1
Our analysis begins with approximating the local movement using primarily Taylor expansion around the 2-ES Minimum (Appendix B.3). We show that for initialization within local re- gion of width 2 and height - centered at the x 2.2-EoS minimum (condition B. the local two-step update of a and b can be characterized by
a" = -a 2b2 Ra.
bb 1663 + Sabk + Rs.
(33)
Where R,, R3 remainders that effectively bound (Corollary B.1). We note that in the region we are considering, a is always monotonically decreasing at 12,'3 per 2 steps (Lemma 5).
With the approximation ready, will conduct our convergence analysis with 2 phases.
30

---

Published as a conference paper at ICLR 2023
3(1 +3
2am
6²
1ox
8/
aax
18a
jax
60
IV
VI
22 Eo8 Minimum
b
Figure 30: Convergence Diagram GD on the 4 scalar example. The horizontal directions represents and the vertical direction represents b. The arrows indicate the directions of local 2-step movement. This diagram for demonstration and ratios are not exact.
In Phase 1 (Appendix B.4), we consider all possible initializations (ao- bo) such that do E (12x , f and bn E (-K-1, K-o-,KK-JJ_ (){O.. We partition the region of initializations into three parts separated by 62 2ax and b² aax (shown region I, II, III in Fig. 30).
For initializations in region where 1/ > 2ax (condition B.3), we show that the cubic term b in the expression of 8" in Eq. (33) dominates the abk term as well as the remainder, so that the two step update on [6] monotonically decreasing with at least an additive update of -16%). Combining with slow movement of we show that initializations in region I will quickly enter region II (Lemma 8).
For initializations in region III where 122 (0, fax) (condition B.4), we show that the abx term will dominate the 63 term and other remainders. Thus the two-step update of [6] will monotonically increase with a multiplicative rate of at least (1 ax). Combining with slow movement of a, we show that initializations region will also quickly enter region (Lemma 9).
For the last step in Phase 1, we show that the 2-step trajectories entering region will stay in the region in the sense that 82 E aa,, 2nx). (Lemma 11) We also show that will keep decreasing and enter ( { x 2x²) ). In the diagram (Fig. 30) this corresponds to the phase of entering region IV from II.
At the end of Phase I, we would have shown that all trajectories starting from the required initializa- tion will converge to near the parabola 62 = an and enter region from the right.
In Phase 2 (Appendix B.5), we begin with initialization (ao bu) such that do E (gx 2 2x and o3 (1 aax, 2ayk).. In phase we have shown rough convergence result close to the parabola
62 fan with the extreme point of (0, 0). In phase we will change the parabola of interest be 6² fax + 1a x which is characterized by the ODE approximation as discussed Section 3.1. In particular we will focus on the residual å 6² fax To 1xx The phase 2 convergence has the following three stages. Throughout the analysis we fix a small constant 0.04
), we will show that the two-step update of & follows I5" (1- 32c²) < 862 (Lemma 12). Then we show that [g] will further decrease to less than o before decreases to less than & and the trajectory will enter region V from the right (Lemma 13).
Stage After the trajectory enters a E ( 2x ) and 1/2 E (Gan2, 2,xx,),
Stage 2. After the trajectory enters region V, we will show that it will remain close to 1o Rd that [[] will remain less than $ 28 while decreases
the parabola A 12 jan
31

---

Published as a conference paper at ICLR 2023
into the interval (- - 3ōx²), (1 + 26)x²) and the trajectory enters region VI (Lemma 14).
Stage 3. Finally we conclude the proof by a convergence analysis in region VI. The two- step dynamics approximation region VI is very similar region III. that the abx term in the two-step update of will dominate. Since is now negative, |6| will follow the multiplicative update [b" < (1 + ax}[b|. We will also show that the movement of will be small, and the final converging minima will not be far from the extreme point 1 k3 (Lemma 15).
B.3 AND 2-STEP DYNAMICS OF AND
Now we begin our rigorous analysis on the dynamics of (de.be).
For simplicity of notations, when analyzing the 1-step and 2-step dynamics of (ae, be), we a. a', a" to denote 42+11 4c+2 and b, b, 6" to denote b,, bet1, be42. For simplicity of calculation, we consider the change of variable K VT..
In the following analysis, use operator (-) to only hides absolute constants that are independent of e, K. a, and no asymptotic limits are taken. Concretely, for monomial x and polynomial y of some variables, we denote y = (z)) if there exists some absolute constant K independent of the variables such that for any parameterization of the variables, ly/ K] [//]. Note that this stronger than the usual big-O notations. Throughout the analysis we will use K to represent the absolute constant that uniformly upper bounds all the absolute constants the (-)) terms. This well defined as we
will only be considering finite number of (-) terms.
B.3.1 ONE STEP DYNAMICS APPROXIMATION OF AND b
Lemma 1. Fix any positive e that < 0.5, for any (0, for all (a, b) such that |a| ex- and [b] minn1 =ex- we have
d' = d'=a-28"
2b + +0(e6*2)+
(34)
Proof of Lemma 1. Recall from Eq. (32) that
a' (a (n (x-1_ 4) (1- - (1 {
(35)
Since (x-4 -4) 4)+ will approach infinity as x goes to 0 we instead analyze
ka' = [sa + (1 - 4nx)) (1-4x*)))
- ()
& (1 4x*)#
(36)
Note that (1- - () will be close to 1 - 26 (1+ + +) for not so large b and 1 - 4x will be close to for small K, we will leverage these two properties to approximate a'.
First observe that for any x (0.8e/ + 2c)²) we have 1 1- - cz v1
g. Since
6 < 0.5, (1 2c)² 4, so it is sufficient to let x < for the inequality to hold.
Now we substitute by (1 + 6) + b6 Since |b| < 1, (1 () + b)2 < 36. Thus when
[b] < fex=" we have 80 + 6)2 + b) < 3612 4
< 2e, and therefore
(1-
= =1 (01 -(e)?)))))))))))
(37)
-1 262x* + 0(b2x²5) 0leb2xt)
Since 4x4 > 0, by Bernoulli inequality we have (1 x k4)a 1- 4 4x combining with the requirement of x < e we have (1 - 4n:4) 2 x1 > - 6. Meanwhile since we required |d| < e1 |xd] < e, so
(sd + (1 1 1 0())
(38)
32

---

Published as a conference paper at ICLR 2023
Combining Eq. (37) and Eq. (38) we have
ka'
K0 + (1 4x=*) (1 -11-1xx 4x=)
ka + (1 - 4x {1 -0=- - (1- - 4x²))
ka (1 (1-2b²k* +(+ (1-24))) + + O(e)2x)))
Na (Ka+ (-27
Ka - (1+ (1+) (-2b2 +0(3m)))))+++ +1++0()))==+++++000+0+++++++++ +0(Pp(²))))))))))) = ka - + O(
(39)
Dividing both sides by x completes the proof.
Lemma 2. For any (0, for all (a, b) such that |a| x-1 and |6| 1,
= --bs - 3b²² 13 + O(ab2n)) m) + O(a²bn)) + x(b2x²) + O(br)).
(40)
Proof of Lemma 2. For simplicity of notation, let
8 - Anp (1+ 6)2 + (an² + (1 - -17)))) V 4x4 b)2 (ax + (1 -4nt))²²²
(41)
Plugging Eq. (41) into Eq. (32), since |6| < 1 we have
4b2 + 16b²R² 256 x² + 196" + = =0'===(((()))
(42)
Thus
b' = - 2-28- - 313 - 138 + x + 166 + 25b + 196 + 70" + b²
(43)
= b- 26ō 362 688 Ob2xx.)
Since x < 0.1, |a| < - and |b] < 1, by Lemma 18 we have
8 = + 2ak a a²x² + (4b + 20²) + O(x²).
(44)
Thus
= b 2b 4abx - 312 - Gab - 13 2nb3x
+ (2b + 3b2 + +b²) (O(a² + (bxx') 0(x²))
(45)
b 4abx 36² - 63 + Olab² k) 0a22bx3) O((k")
Combining the conditions required by Lemma and Lemma 2 we have Condition B.1 (One-step Dynamics Approximation Condition).
c < 0.5,
min (0.1,e },
(46)
[a| ek e!
|b| < min [1, Jex 2}.
When condition B.1 satisfied, we have
a' = a - 262 + 0(eb2x3) + (bb2x))
(47)
6' = 4abx - - 362 - 63 (ab²xk) + 0(ab² 0(ab)+))+0004+ + O(bx2).
33

---

Published as a conference paper at ICLR 2023
B.3.2 Two STEP DYNAMICS APPROXIMATION OF AND
Now we approximate the 2-step dynamics for (a, 6).
Lemma 3. Fix some positive constant < 0.1, with K, a, b satisfying condition B.I. 8" = b 16b² + Sabk + O(b*) + (aa²²k)) + O(a + O(b²xc*) + Obbic)) + O 1²)) (48)
Proof. Combining the one step dynamics characterized in Lemma and Lemma 2 we have
b" = -b - 4a'b'n H - 3b2 - 8 + x)d
(49)
We will analyze these terms one by one.
d'b'k x 26²x=² +((6(x))+((())))) (-6-
(-b
-
- 6²
+0 +00%
= an (-b - 4abk- - (- +10
4adx - 3b2 - + +0 0(br)) (50)
-2 (-b 4-bn-
+O(eb2n23) + xOb2x(*)
abn + x) O (n²bnn) +0x
2 (-b 4abx 3b² 1/ +0000
(b (-1 4an 30 --0----------00000000000000----- --0-0000000000000000000000000000
(-b (-b-4atre - 362 000---0+-00000000000000000000000 (51)
b² + 66" +000))00000
13 8 4abis -00
(b (1 4ax 30 - -00)))))0))))
3
(52)
(-b 4abx 362 +00
+0(bk²)
63 + 066*) + +0(x)=))))00)()+=+=))=))))))+))) Obbk).
Note that 1/ ((b) + (abn.) and a' + 0(12x3),
Old'b²k) O k) + O(b'nt), (aa26k²²) O(d²bk?) 0(22bx))
(53)
0112x2) O(B2x4 + Oab3xx²),
0(bx²) = O(b)²) + Oabn.))
Hence
-,w.000000000000000oco.000010000 (54)
Combining above we have
8" (-b 4abn
3b² - b²)
- -00 (-b²)
+ + +
= +00_0)0))--0(((00000000000000000 x3)
+ Sabi + +00-000))0-000000000000000000++0 + O(b'')+(('''")))
+ +
(ct2
(55)
Lemma 4, Fix some positive constant < 0.5, with k, a, satisfying condition B.
a" 2 O(eb² 0(12AA).
(56)
34

---

Published as a conference paper at ICLR 2023
Proof of Lemma 4. Combining Lemma 1 and Lemma 2 we have ="=a' 20 + Oeeb +(b2x'
(57)
From Eq. (51) in the proof of Lemma 3 we have
b²= = 12 + 0(b²) + Ola't Hence the last three terms of Eq. (57) can be calculated that
(58)
+
0((9)))
(59)
O(eb9x²) e(
= =(eb2x))+ + +00
(60)
= +
0( O(b9)k*
+ 0000
(61)
= -(((xx +10k))))))
Combining above we have that
a" a' 21 +0
a = 0(e3)) -00))
+ 0016))))00(())))--0000)000000000 O(b2n>) (((R)) ( ((en((2))))
(62)
+ +000))))0003)++000300+000---0000
= a - 4122 +
This completes the proof.
Combining Lemma 4 and Lemma 3, the 2-step dynamics approximation can be summarized as: Corollary B.1. For any positive constant < 0.5, with k, a, b satisfying condition B. we have 6" = - 1653 + Sabx + O(b") + (ab²kk) + Oa?bkx)) + O(b2xt) + O(bx*) (eb²² x3), a" " = a 4b2 1 O(eb² +((9))")++(()))) O(b2..t.
(63)
B.4 PHASE I: CONVERGENCE TO NEAR PARABOLA
Now we show that under the (a, b) parameterization, any initialization (4o, bo) such that do E (3k², 1 -1) and bn (-K-1, K-1) will converge near the parabola very fast. As men- tioned above, there are mainly three regimes of interest. We will first determine the region in which the two step update 6" b is solely dominated by -b3 or abs.
To formally analyze the different dynamics and characterizing the regimes for them, in this section we use K denote the uniform upper bound over the absolute constants hidden by the O-)) operator in the 2-step dynamics approximation characterized by Corollary B.1. Since there is only finite terms with O-notation, such constant K well defined and independent of e, K, a, Without loss of generality we assume K >> 512.
Rewriting Corollary B.I with the uniform upper bound K, we have the following corollary: Corollary B.2. (2-step dynamics approximation of a and b) There exists some absolute constants 0.5, for all k. a, b satisfying condition B. 1, the 2-step update of
K such that for any constant (a, b) can be characterized by
d" == a - 2 1.3 + Ray b" -b 16b + Sabn + Ro-
(64)
where the remainder R₂ and R., have upper bound:
RRal < K lec² 3
+ K xx3 + K x22x
(65)
[Ra] K
| K 2 dxl
K cdd ?
+
K 1 2
K |ex.
K |ec2x3]
Now we establish the conditions to characterize the work zones in order for the analysis to be more tractable in different regimes.
35

---

Published as a conference paper at ICLR 2023
B.4.1 WHEN a" - IS CLOSE TO x4
Here we formalize the observation that when and R are not large, the two-step movement of is monotone and always close to 462 k², Concretely we have the following lemma:
Lemma 5. Fix K-1 for all ks a, satisfying condition B.I well the extra condition of and [6] K -1 we have Ja" - -(a (a-44b2 4b2xx²) ) < 3b² k²
< K |K|<
Proof of Lemma 5. Since we assume K > 512, we have K- 0.5 as required by Corol- lary B.2. Since a, b satisfies condition B.1, by Corollary B.2 we have a" 4b²m" n + Ra where
R
K <b K |R
K |b2x*1
(66)
Thus to prove the lemma we only need to bound every term on RHS of Eq. (66) by b
Since fixed K
K
cb2 = b²
(ii) Since [6] < K K |6
< K k < b2...
(iii) Since K -1 K 13 n
K K
8 x
< kR
Therefore |Ra < 362x2 which completes the proof.
Note that when we fix e K"1, condition B.1 becomes
K min K- $},
|a| < K- k-
(67)
|b| < min (
Combining with the additional condition of Lemma 5, we summarize the condition for Lemma to hold as
K < min (0.1, K-1,,
(1)
|a| < K- x
(2)
(68)
[b] < min (1, K 15
K K 2
(3)
With K > 512, K-1 < 0.1, so (1) can reduced to k K-1, Since < K-1, fx- K2
so K- K > 1 K-1, and thus (3) can be reduced to |b| < K-1 In conclusion, the following condition is sufficient for la" - (a - 4b² x3) < 362x3.
Condition B.2 (Condition for (4 3)62x> x3 movement of a).
K- |a| [6|
K
(69)
< K-1
B.4.2 WHEN -03 DOMINATES THE DYNAMICS OF
Lemma 6. For all K. a, satisfying
N < K-1, lal E (N K x=xx)) |6| (Via)K, K-1).
(70)
We have 16" - 116--661166100 < 14).
Proof of Lemma 6. First it is straightforward to check that the condition K, a, b is stronger than
condition B.1 if we set e = 0.1, thus by Corollary B.2 we have 1/" = b 1613 + Sabs + Rg where
Ry K ba| + K |ab² Nc| K ja²bn? + b0 + K |bx²| + K leb²
(71)
Thus to prove the claim, is sufficient to bound |Sabx| by 8|63| and every term on RHS of Eq. (71) by [b3]. We will now bound them term by term.
36

---

Published as a conference paper at ICLR 2023
Since v |a|k [6], |ax] < 62, so 8|bk]] 8|b||ax] < =./b8.
(ii) Since (6]
< K -- , K|b < K K 10 = -b%
(iii Since |a| < x multiplying K² [ax² on both side gives k² |a jc. it follows by
taking square root that K |a|k < Vlals. Since |b| > Vlalis we have K K|a|s < 16/. Multiply
b² on both side gives K| K|ab²x| < -b%
(iv) Since |a| < K R- < K x-''_=,, multiply K|aa i on both side gives Ka²k²² < |a/re and hence v aKaa/ < Vaan < 16.. Squaring both sides and multiply by b gives K bn < |6²).
(v) Since |a| > x² and |b| > v alk,, we have |b| > 2 Kis Multiplying on both sides, we have K 16 1691.
(vi) In (V) observe that |6|
x² so |b| VKx and hence b² Ks². Multiplying |6) on both
side gives K] K|bx** < 1b%)
(vii) Since we fixed c == 0.1 while < K-1 < 0.1, we have oR < 1. Since |b] > n² we have [b] > x3c Kea3. Multiply 6² on both side gives K| K|ce2k0| |c²).
Therefore we have [8abx] + Rol < 14 ea which completes the proof.
We restate the sufficient condition for [b" (b - 16b²) S 14|6*| as follow Condition B.3 (Condition for - Dominated b Movement).
K-1, lal E (m²3 K-2 xx
(72)
|b| E (Vla)m, K K-1).
B.4.3 WHEN abs DOMINATES THE 8"
Lemma 7, For any K, a, satisfying
K
la| (k5,, K 2 x1 1 2/2 Vlalia K-1)
(73)
|b| < minf
We have 16" - (b + 8abx)] < abbx].
Proof of Lemma 7. First e = 0.5. is straightforward to check that for all Ks a, satisfying the given condition, condition B.1 is also satisfied, so by Corollary B.2 we have b" = b 1663 + Sabk + Rg where
Rs
< K
6" + K |ab2N k6| K bn + K b0xx
K |b.* + K leb2x.1.
(74)
Thus to prove the claim, sufficient to bound 16 by 2 2labx]. two terms in RHS Eq. (74) by ألوا and the remaining 4 terms by |b3). We will now bound them term by term.
(i) Since (6| < 2/2 Vlalk, 6² < Halk. Multiply 16|b] on both side gives 16|b²| < 2|abri).
iii) Since |a| < K-2kp 2k, we have Ja/N < K-2 Multiplying a²x² on both side gives |a <
K k-2a²m² If we take the 6-th root, we have Vlal < (K-+)a)x)3. Since [6] < Vlal.
we have [b] (K-*(a)w)) and hence |b31 Ka//a Multiply K|b| on both side gives
K [114 labk].
(iii) Since bb < K-1, multiply both side by K Klabel gives KlatPx| < |abe].
(iv) Since |a| < --=x-11 < - (as assumed K > 512), multiply both side by lab|*?
gives K|a < } |abb).
37

---

Published as a conference paper at ICLR 2023
(v) Since |a] > H² and < K-1, we have |a] > R² > K~2,3 > K2x Multiplying K K-2/aair both side gives
k- |a|s. It follows by taking the square root that Vla iss K- {|=|a//>>> Since we have [6] 2v2 V Vala,, combining with above gives |b] <
} |K| Multiply K|b|=* both side gives K|b2x*4| Habs).
(vi) Since |a| > x², we have |a| > Ks Multiplying [b/rc on both side gives K K|bx²1 < label-
(vii) Since |a| n and e 0.5, have |a| k2e2k5. Taking the multiplicative inverse and multiply a²x on both side gives |als < K 2-2 " It follows by taking the square root that Vlalk < KK-1c=1x=9a,, Since [b] < Vlaln, we have |b| < - Finally
multiply Ke Ke/b/e on both side, we have K k)bb2x < labx).
From (i) we have 16|b²| < 2/abs], from (ii) (vii) we have Rs = K |644 + K Jab2n + K ba²bn? +
K |b2x*] + K bbr0| + K |cb < 5/abs]. Therefore Іы" - (b + 8abr)) - 16
+ Rs I 7 |abbs]
which completes the proof.
We restate the sufficient condition for |b" - (b + 8abx)] S 7|abx| follow: Condition B.4 (Condition for abis Dominated b" Movement).
k K-
la| E (R²2, K Kー2 n |b| mint 2/2 Valls, K-1).
(75)
B.4.4 CONVERGENCE WHEN -b² DOMINATES abx
Now are ready to analyze the dynamics when e² dominates the movement of c. For E N, abuse the notation to let (ct, de) denote e and d after the 1-th 2-step update with step size 1j = x² from the initialization (Co, do).
Lemma 8 (Convergence to near parabola from large 0). For any K < K-1, for any initialization (ag, bu) satisfying ap E (12kg,, 1 K KK-2 -n=x) and [bul E 122/aaole K-1), there exists some < K such that or (2x², K -K2xx and |br| (Vark, 2/TTK).
Proof of Lemma 8. We will prove the claim using induction.
Consider the inductive hypothesis for l E 1,, [s= that
P(k): lax| E (2x², 1K1 -xxxxx111 and |bx| E (2 min(K- ,,(,, (k + 1)-1)}.
Since [ba] < K- <1= and ag E (3x2, - c (2x², 1 - as required by the initialization, the base case P(O) holds trivially. Now can proceed to the inductive step.
Assume P(() holds for all / < k, we want to show that P(k + 1) holds unless |6x+1| 2 |as+1|n. In. By the inductive hypothesis, (ax, by) satisties condition B.2 and condition B.3. Thus by Lemma and Lemma we have
0g41 (ak - 762,53, 12,,,ax, as - kz
[bx+1| (lbxl 301b21 |ba] 2-12))
(76)
By the strong inductive hypothesis, these properties also holds when substituting & by any < k.
38

---

Published as a conference paper at ICLR 2023
We first check the lower bound for ax+1 under the assumption that & < Since at+1 > ai 7b x² for all l I k. we have
01++1 > 40
7b k
1=0
> 00 Σ
(1+ 1)-11)89) ) n²
(Since by < (1+ 1)- by IH)
k
(77)
> 12A 7
Σ
1+1
1=0
k++
> =1k - -7103
1+
- dr
12x 7 + log(k: + 1)).
Under the assumption that k < n and K < K-1 < 1 it is not hard to check that
512'
1 + log(k + 1) < 2 + log(k x 4) =
log(k)
<
10 7 Ks
(78)
The last inequality holds since when バ. 512 we have 2 4 log(
512 )
< 5 and
512
2 + log(n= 2 log(k) is monotonically increasing when K <
Plugging back into
Eq. (77), we know that if P(1) holds for all S k and k < K x then 512'
ag + 12k 7 [10x=
12k 10 = 2
(79)
The upper bound of ax < - always holds since is monotonically decreasing by Eq. (76).
Now check the upper bound for |bx+2l- Consider (1: --= 1, we have f'(x) 1
(x:)3 and (*(x) 4 - Note that (*)x) > 0 for all x > 0, so by a first order Taylor expansion around
h + 1 we have f(k + 2) > f(k + 1) f(k + 1) Combined with
|b&+1| < [6x| [62] . it follows
f(k + 2) - | 2 f(k + 1) 2 f(k + 1)3 - |byl + 1621
(f(k+ + 1) - 1bbl) 1
( + 1)3 -162) 2
1
2 (f(k+ + 1) - 1b+]) 2 (f(A: +1) - -1xx + 1bblf + 1) + f(k + 1)3)
-f(k + 1) (bel) (1 1 (0) |bba](kk 1) - f(k: [(++1)))
(80)
Since f(k + 1) < 1 and |b < lbul < K-1 1.4 for all 1: 2 1, we have by |b f(k+
1)² < 2, and hence (1 - 2 (bg + [bllf + f(k 1)² > 0. Since |bxl < (k + 1)
f(k + 1) by the induction hypothesis, we have f(k 2 - | 1| >> O. so |bę+1| < (k- + 2) The The
other upper bound [b&+1| < K-1 holds trivially since [bol < K as required by the initialization
and [by] is monotonically decreasing according to Eq. (76).
Now we will show that there exists some T < K 4 that |b,| < 2/ a, jk. Assume toward contra- diction that there is no such т, then the induction may proceed to h & | so that on > 2n |ba| < (h + 1) < (x= < x² and [bx| > 2 > 2v 2V 2 2/2x The last two inequal- ities lead contradiction 2 2 > x², so the assumption does not hold and there exists some T < R- for |b~| < 2
Let T < [ be the smallest such т, then the induction may proceed to k = =T T 1, which guarantees P(t) for all t < T. Thus for all t < T, a, > (2x K -1= and |b,| E (2VOTK, K -)). Since P(T 1) holds, following Eq. (77) we also have ar 2
39

---

Published as a conference paper at ICLR 2023
Now still need to show |brl > y Vark. Since ag is monotonically decreasing, it is sufficient to
show |br| / -I-- From Eq. (76) have |br| > b 30 30 b 1br-11(1 - 306)
Since bt < K- 5T2 (1 306
Combined with |bT-1] / shown
above, we have |brl > (2 (12 VTTIR which completes the proof.
B.4.5 CONVERGENCE WHEN abs DOMINATES 33
Lemma 9 (Convergence to near parabola from small b). For any K < K-1, for any initial- ization (0o, bo) satisfying ao E (12x², 1K-2 and lbol E (0, (VnR), there exists some T < 1 log( (olbbull)))))) -1)x 1 such that [by] (ivarR 2 and ar (2x², - 1
Proof of Lemma Lemma 9. We will prove this claim using induction.
Consider the inductive hypothesis for le E N that
P(k) :bul E (160) (1 + ) AVagk)) ao E (2n², x
The base case when and if l: 2 1, (lbxl - (bol)/(as ax) > -
holds from the initialization, so we proceed to the inductive step, Assume P(I) holds for all I S k we want to show P(k + 1) holds unless [bg| > 4 ag
First note that by the inductive hypothesis, since ag < 1 K-x we have bbl /Vogg 1 vV VVK-2 < K-1, Combining with conditions on las| and | from the inductive hypothesis we have (ax, by) satisfying condition B.2 and condition B.4. Thus by Lemma 5 and Lemma 7 we have
at+1 (ax - 7bp a2
bRx²)),
1bx+1| E (lb&l + 2/axbbxl), lbx| 14|axbyx().
(81)
Observe that when is not too small, the movement of b is significantly larger than the movement of a. From Eq. (81) we have Og
ax+1 < 812 and |bx+1l - [bx| > 2uglbale (x, so
b lbx
2ax lbulN
ak
0x Gx+1
863
[b&| 4x
4 aj
(since lbel (vark)
(82)
4x²
ak
4.
(since a xτ²)
x
When & = O, we directly have (1621 - /bll)//c/ (as > When & Z 1, from the inductive hypothesis we have (lbyl - (bul)/ ax) > x : combining with (1b++1l 1b 1//x/ 4x+1) >
k- i, we have (16++1l= [6ol)/ (ao ax+1) > 3 by the mediant inequality.
Now we check the bounds on |a*+1| 1| and [b++1]-
Since is monotonically decreasing from Eq. (81), we have ox+1 < ao < 1 K-
Since (lb₂ bol))(ag ax) > x according to the inductive hypothesis and (as ax) > by monotonicity of a, ag - ag < (l6xl - [bol) < |b k < A 1 V < 1 _ vanx 7 Here the last step holds again by monotonicity of Reorganizing the inequality we have 0₁ >
an / VoalVan t Since og 12x 3 by the initialization condition, vag 3x and Vag RE > 2n... Thus a₁ > vaolVao } > 6r Note that since |by]
< K < 512"
7b2 < 7K-2, < R² From Eq. (81) we have a ax 8bf 6x x3 2x This
gives the desired lower bound for as+1
Since 1b4-1| > |byl + 2ax|buln from Eq. (81), combining with as > 2x we have |bs+1l > [b,1(1 (1 + 20 > bbx1(1 + 2x²a)) x). Since [bx| [bol 4x by the inductive hypothesis, have |bx+1| [bol(1 + 4-41
40

---

Published as a conference paper at ICLR 2023
With the guarantees on |ax+1l. |6k+1|. |. and ([bg|-Ibb) -bbll)///// (ag -ax) as shown above, if we additionally assume that |bb+1| < i /ak-1k, P(k + 1) will hold and the induction can proceed.
Now claim that there exists some contradiction that there is no such т, then the induction can proceed for any t E N.
1/ log((bel-'))== that |b,| 2 Va Assume toward
Consider Σ 2 log( [uu 1)xx 4, we have
2 log( ( & + 3)
log( 10o (1 + 1 2)
(Since N x 1 2)
> (2 log( K -x log [bol 1) (1 + 1x=²)
(Since log( K 2) < 0)
G log((K )k²) + log( ")) (1 11+ 2 B
(83)
> ( log(a**) + (0)((()))))))))))))))))))))))( (1+ xxx
(Since K K-2 > >0>> ax)
= & log(an*) + log((6n) -)) (1 + (1+4x3)/44 //
(2 log(ax*) + log |00 -) log(1 + 3
(Since log(1 + x)) 2 20(1 x)).
It follows that
t (3 log(axx) + log((bol log(1 + 4x22
(log(vas)) log( 1)) log(1 + -xx3)-1
(log() (( + (0al log(1 + 4x2))-1
(84)
log log (14
|bol)
Hence > (g5//bbo) and therefore |b,| > [oo(1+4x'2 1o2 Vazk, which contradicts that P(t) holds and lead to contradiction.
Therefore there must exists some r such that bb, > Va, Let T = log 1)) 2 be the smallest such т, then the induction will proceed to h: = T - 1. Moreover since P(T 1) holds, following the previous analysis, the bounds on |uT| also holds, so we have |ar| (2x 1 K-2 11))
Now to complete the proof we only need to show that |6T| < 2 VaT..
From Eq. (81) we have |BT| < |bT-1| | + 14a7-1 110- |n = (1 + 14aT 140T-Ik) (b) Since ar-1 <
K x and K 512, (1 + 14ar-1k) (1 + 14K 1 1.1. So
|br| < 1.1 1.1 bb- < 1. -T)-
(85)
On the other side, note that ar ar-1 - 76% 3 by Eq. (81), where 67-1 (1 (1)=-=1))= 1
16 dr-1K. We have or > 6r-1
16 0r- > (1 x) Since x < K-1 512' we have
7
ar >> 0.99ar-1 and hence
V < 1.1( VaTk).
(86)
Combining Eq. (85) and Eq. (86) we have |by| < 1.21 (aTk) 1 Vark. This concludes the proof for this lemma.
B.4.6 |6] STAYS IN C1 VaK, 2/OK) AS a DECREASES WHEN IS NOT TOO SMALL
After showing that b will enter (1 Vax, 2 a. we now show that will not leave this region unless is very small. To do so, we first determine regime in which we can effectively bound the two-step movement of b.
Lemma 10. For any a, b satisfying
< K
1 |a|
< K
K-22x
(87)
7
|6] < 2
41

---

Published as a conference paper at ICLR 2023
have 18" < 16 Vak.
Proof of Lemma 10. Fix = 0.1, it is straightforward to check that K, b satisfies condition B.1. Hence from Eq. (64) and Eq. (65) we have that
18" b bl 16|6 | + 8labs] + Kb K|ab>mi x| + K K|a + KK2 + K|bx>| + K|elb>x>. (88)
To prove the claim it is sufficient to bound all terms on RHS under 1 T28 V ak. Note that with b <
2aa and
K x-2x-1 Vak vK-2= =
and hence < 2K-1
128
(i) Since K 128(ax)
512, ak < K-2 16384
Multiplying / ak on both sides we have 128 aak. Since |6| < 2 |6%| < 8(ax) thus 16 16|6*| < +128(ak)² <
128 Vas.
(ii) Since [6] < 2a
ak, 8labn|| < 16(ax)$ < 1024 ak. The last inequality holds directly from (i).
(iii) Since [b
< 2/an, Kb 16Ka2 16K(ax) Van. Since Van K-1
16K (ak)
3 vak 16K VVan 16281 V an < 128 Vak. Therefore Kb* 128 Vabk.
(iv) Since [6| < 2K K K|ab>k| < 2labss,, which is less than 512 Van from (ii).
(v) Since [ax] < K-1, Klab2 x2 < |abxl < 1024 Vabe.
(vi) Note that x < 0.1 in condition B.I, xt -
Since |6| 22aa < 2K-1,
1, 6²
4K'1 Van, and thus Kt2x² 4yank*
128 Van.
(vii) Since < 2Vak, K|bxo1 < 2K Vak. Given x < K-1 < 1 K- we have n < 1024
K-1
and hence 2K Vax Vax
1 Vak. Thus K K|bb>*
1
Vabk.
512
512
(viii) Since we fixed 0.1, K-1 on both sides we have 4 /EVaKk* Keb < ((AKK-VaaR)s 46 ark
612 < _ ; 13 size Multiplying 4evox
Vak. Since 8² 4K Van from (vi), we have
125
T2% ax.
Now that we have bounded every monomial term on RHS of Eq. (88) by 128 Vak, we have 1 bl < 1
1
16 v ax, which completes the proof.
Here we restate the condition for Lemma 10 Condition B.5 (Condition for small movement).
k < K-1
|a| < K 22 x-
(89)
[b] < 2v dr.
With the two-step movement of e bounded above, we may proceed to state the lemma that guarantees e will not leave (LVdk, 2v dx) unless is very small.
Lemma 11. For any k and initialization ao, by satisfying
K < T6 K
ag E (K xx-2x-1)
(90)
|bol E (1 Vagri 2oun,,/auns 2ayk)..
There exists some T S 160 # such that ar < n and for all t < T, a, > x² and b, e (1 York, 2
42

---

Published as a conference paper at ICLR 2023
Proof of Lemma 11. First check that the region defined is not empty. This is true since given K- we have - > 4K'1 K
To prove the claim we consider the inductive hypothesis
P(k): ag E (I²² 00 - Te (k 1)xx 3 and |by| E (1 vaRk, 2 (akk).
1 1 - xx_,,,1, we have |by| < 2/ < Assume that P(k) holds for some k. Since ak < ao <
2 K -k = K- With K and as < ag K- x-_ we have bu, ag satisfying condition B.2 and condition B.5. Thus a++1 ag k3
3, 0 0g 812 (by Lemma 5), and
[bk+1 bul < 16 v Vagk (by Lemma 10).
Observe that since ax++ > 8bg
2vakt1k3 2/((0 (as 802
22/((a) (ax 8(40Hx))
(sinbe by 22/ak)
(91)
= 2v (a<x)))(1 - 32x²)
22 - 32n² VaLi
Note that with & <
1 K- where we assume K > 128, we have v1 -1- > 0.99 and hence
T6
V > /a99 Now we will show that b will not leave LiV / There are three cases to consider:
1. When |b& E (VORk, 2 (x along with 1bg < K-1 we have condition B.3 satisfied and thus |bx+1| 1bx1(1 bb) by Lemma Since x. < 16 K- 1 we have x²
2018
1v2'
and thus x3 > 4v2x2. Moreover, since a& > xÿ, we have |by| > Vaga > na 4
Squaring both sides we have b2 > 32x*. Hence 1 32x > bg > 1 26 + b
(1-12)². The last inequality holds since b3 >> by as |bxl < 1. Now taking the square root on both sides we have 1-b2 < VI 4 Since 16xl < 2 agk combining with Eq. (91) and
1 b- vI 32x² we have |bx+1| 2//ak((11 -b) 221 32x Vark < < /ax+1h which gives the desired upper bound to |bx+1].
Now we prove the lower bound for [bx.1l. Since we know |6x+1 bul < L K by 15 16 V YoRk > VaRk > / vab+1k.
16
triangle inequality, |bx| > Vax implies [bx+1|
2. When |byl E 2//2 V aatt VaKRe, since |bb+1 - byl < 16 Vark, by triangle inequality we have [bk+1l
1( 2
1 /akk, 16 Vagkl. Since ( 1
2222 1 > 1 VaaK >
7 VaR+IR and Yakk 2 VaRtiR LavakK 2..KK,2//ttt 2//
0.99 /ank, we have lbx+1|
3. When |bx| E ( 2V2 V (akk) along with lbul < K- we have condition B.4 satisfied, and thus |b&+1| > lb&l + |bba4,sl > (bal. Since Var+ik < vark, |bs+1l > _/2 /v
On the other side, since [be+1 - bul < To Vank, and |bal < 2/2 Vask, by triangle in- equality we have lbx+il < (16 + 2VZ2 V < 2 The last step holds since Vag+Ik > 0.99 a.K.
I
I
Summarizing the three cases, we know that be+1 (1 ( 201x+115))
By the assumption of P(k) we also have as > σ² and 0% < 0p - to (k 1) 1)x 4 Since we know by > 1 | /agk, have 13 1o takkt 1 Thus a ag bg ag -/ -
ag 16 kk H. Therefore unless ba+1 < xł. P(k + 1) holds. Note that there must be some t such that 04 N since when t. > 16ag 1 + 1, ag (1 1)x 0. We induct on k from 1, the
43

---

Published as a conference paper at ICLR 2023
base case holds by the initialization of by and ag. Let T be the smallest such that ds R²,, which we terminate the induction. Then for all t < T, Ot (K3,,0) and b, E (ivarh (ark) This concludes the proof.
Corollary B.3.
Following the initialization condition and notation Lemmaa 11, if 2x, there exists some < T such that a, E =gxx3 2x
Proof of Corollary B.3. We will follow the notations defined in the proof of Lemma 11. By defi- nition of T. P(t) holds for all < T. Then for all & < T, we have > x², b < K-1, and
a++1 < 04 12 Hence |a++1 a,l k-2 _ since we assumed K > 128.
Since 0g 2x and aT < <x² < 3x3, combining with |4++1 - arl < 2 we know that there must exist some such that or E (2 2x²).
Corollary B.4. Following Lemma 11, if E (2 2,,,2,x)... T
1
128 K
Proof of Corollary B.4. We will follow the notations defined in the proof of Lemma 11. By def-
inition of T, P(t) holds for all t < T. Then for all < T, we have a ag
2x
and by < 2 < 2v2x1.
follows that at+1 > ar 862k3 > a - 64x 2 Since
i3
ao ar > 3 k = j we must have T > fx2 / 128
B.5 PHASE II: CONVERGENCE ALONG THE PARABOLA
In Appendix B.4 we have shown that for a certain range of initializations, (a. b) converges close to the parabola 2b very fast. In this section, will show that (a, b) will slowly move along the parabola, and will eventually converge to a point with sharpness just below the EoS threshold 2/1- 2/k².
to and consider small pertur- To facilitate the analysis, we define the residual & & b² -ax
bation constant threshold 6 = 0.04.
Follow from Corollary B.1 we have that for any 0.5, for any is a, b satisfying condition B.
s" = 822 1
aa
K
2
16"
=(1² + 8abs - 166)3) 2
O +(a + +((P)r)) + +Ob( +000==+00000000000000000)))++00) 0(6))
+
1
-ax + 262 +
16 か
2
= 62 -2b4 + 1662 ax 2 + 2/2 +22
16 か
+ +Ob6' + O(ab2>) k) + Ola2b3 x2 + 0(b3x* +(0++)))
= (1 11-3201 3263) (82
2 ak
16 k² +
+ +
+ 00000440000000000000000000220330
(1-3263) & 0(6²) + 0lab3 k) +0033000000000000000000000000000
Oeb3rat).
(92)
When |b| < ki the above expression can be further reduced
O(b²) s"= (1-32))) < 3
(abb*)) + (a²t? (bb)xt))))))))))))))))))))))(( + O
(93)
Hence there exists absolute constants K such that for all c < 0.5, for all a, b, ss satisfying condi- tion B. and b| < K. we have ę" = (1 -2022 ){ + Rs where
|Rg K | + K lab3 Nx| K 1a=0">=" K 16 K K(bdkn
(94)
44

---

Published as a conference paper at ICLR 2023
Now fix 8 0.04, will determine the regime such that [Rg] is less than 6b2nx
Lemma 12. For any b satisfying
R< 80V2
SK
[6] < 2//22, |a| < 2m
(95)
where ö = 0.04, we have
5 - (2c)) 5bb2x4
(96)
Proof of Lemma 12. c __ claim that K, a in the given regime satisfies condition B.1. We check the conditions one by one:
(i) Since we assume K 128, fixing = 5 j satisfies < 0.5
(ii) With both å and K- less than 0, k 80.2
oK < 0.1. Also 80V2
1 oK < _0OK
ca. Thus A < min
(iii) With = _ and < 1 5K-1 Ex
16v22 > 2x >
|a|. Thus |a| < <ex-
80.2
(iv) Since en > 16v2 and & < 1, Ex 2/5 min(1,ex 3/5}.
V2 > > 2/2x1, Thus [6]
16 5
<
Thus Eq. (94) applies, and we only need to bound every term on its RHS by gob2 to complete the proof. We will do that term by term.
(i) Since & < 1 5K -1 < 1, we have x²
5K Multiplying 16/2xc* on both sides 80v2
80V2
we have 16 4 <
1 Note that since [b| < 2v/kž, |6 < 16 so |ba <
1 {KK wA. Multiplying KT² on both sides gives K|b*| < 5 555b.44
5
(ii) Since x
20 OK- multiplying 4/2x3 on both sides we have 4 {KK Note that since |b| < 2v2Kī and |a| 2x labi 4v we have labl <
80.2 30120K 1LKK-
5 6K 18K-1k" k3 Multiplying Kb2K on both sides gives K|ab" x| 66b24.
1 _ < 1, have 33 <R O0 Multiplying 4x² both side gives (iii) Since 80v2
_ K 1-x2 Multiplying Kb2k² on 4x -1x22 Since lal < 2xd, we have a² < 4x
both side, we have 1a2 fo0224.
(iv) Since x < 2. 2x < K. Thus ||| < 2n < k <
1 - FoK
Multiplying Kb2k
80
on both side gives K1b3<* x4
1 k2.
(v) Since fixed 5 Keb? 1ob2
Therefore have
Rgg < K |5 K lab'k| K x22
K 133
K leb2x=1 <
5b2
(97)
Plugging back to g" = (1 )2 + Rg completes the proof.
Here restate the condition for Lemma 13: Condition B.6. With 5 = 0.04 and K > 512,
1
8K-1 2v |a| 3
K
(98)
2
Corollary B.5. For any K, a, satisfying condition B.6, if(E
160 have [5" < <(1- )[]
45

---

Published as a conference paper at ICLR 2023
Proof of Corollary B.5. Since condition B.6 holds, by Lemma 13 we have
1s"1 < [(1 10-233222 )5l 5b2p'.
(99)
Thus if |5| >
Sns, it follows that
16
|s"| < [11 - 326 + 8b2kr²
= (1 326²) ){l + 8b2m=* n²
(Since 326² < 1)
= |gl 16b²|g (16b² [g| - 5bb2x²)
< [&| 1662| [ 6²66 1/ 5k ) 5k")
(Since [<]
1/ 6xx))
(100)
І&| 16b² k
< K<| 4x 7 [|]
(Since [6] > 1x²t)
(1 4x $ |[()
B.5.1 PHASE II STAGE
In this stage we will show that after b) gets close to Vak/2 while a decreases to around 2x from Phase of the convergence, the residual of (b,a) a) to the parabola will further decrease to below jōx²
Lemma 13. For any K < I
80V2 5K- all initialization (bo, an) such that ap E (ō 2x and
lbol E (orK, 2/oKk). Let T be the time that exits (1:3, 2x as characterized in Lemma 11 and Corollary B.4. There exists some T < T such that & < j
Proof of Lemma 13. First note that since 80V2 SK-1 I6 K the initialization condition given is a subset of the valid initialization for Lemma 11. Thus for all < T, we have a
(N:3, 2x and |be] E (3va (0)K). Hence [be] < 2V 25 = 2/2x 2 and [b,| > V An2n=
f Therefore (b,, a.) satisfies condition B.6 and |5+t1 - -2 < Ob, Also note that with K > 512 and 5 = 0.05, b7 Sk3 8k < 1 1-8K-1 < 1.
For all / < T that [6+1 initialization, we have |Ei] S 16 + Hood + ttxe < Jaon + t < 9x* + 16 1 < 10x3, Thus
I6 0x by Corollary B.5 we have |et+1| (1 4< At the
for all F < T' such that for all t t, |&r L ōnx,, we have [6rl < (1 4x 10 Now we only need to show that & will decrease sufficiently fast.
Consider
]
log( 8n2) 80
1og(805-1
==
log (og(1-an7)
3)
1
8
log(1 - (x - log(1 - 4x
(101)
Since log(1 - 4x <
_4x3 by a second order Taylor expansion, we have 7 log(806-14) x 112 3x Substituting 8 = 0.04 in the expression, log(800 -xx x log(2000), +
log(s= 4) 8 + log(A (N- t). Observe that for all 175, we have log(x) < Vo.
Since
1 8K 80120K
2 we have 8 log(k +) < x 7 175 let
80V2
512x200062
Hence log(806 K $) < x - and < K - //4 fNx - +. Recall from Corollary B.4 we have 324 and thus - > 32 and
T > R -4. If we assume K > 512, then K < 512x2000
1
T> A > 1
Since 7 < T. a, > #2. If for all t < r. [sel > 1 6x4, then following the analysis above, by definition of T we have |rl < fōnt. there exists some < 7 that [5el I 5r64, then setting r = directly completes the proof.
46

---

Published as a conference paper at ICLR 2023
B.5.2 PHASE II STAGE 2
In this phase, we show that once [<] is smaller than & jont, a will decrease slowly while |<| does not increase beyond 1
Lemma 14. For all <
80.2 - 20=70K-, for all initialization (ao, bu) satisfying 01 E (x3 2xd) and |& { 5ntt there exists some T < 486-1861 1000--x3__ + such that or -3( (1-36)mm² and fnd all < T,
ag (1 36 and [Ex1 < _
We will prove the claim using induction. Consider the inductive hypothesis
P(A): Ex < _ 5md and dg E ( (- (1- - -35)m²
15 kk]
1
Note that P(O) holds directly from construction, so we proceed to the inductive step. Assuming P(k) holds, We will show that either ag < - 35) or P(k: + 1) holds.
First we verify that dy. ag satisfies condition B.6. Since ō = 0.05, | - (1-35))))) n5 < x1 < 2 Also since a1 2x as required, we have |a,| < max - [ (1-335)%)),0 |, 1 k|} < 2n²
Since |5x1 < 8 1 we have by < 2. 16 nd + 1 dn < 2 (2x + < <2c3.. Hence 1b91 < 2 2
8
as required. Therefore we have [6x-1 (1 - 321 )&x < 50g.n.
Next we establish the lower bounds for [bu]. which will give lower bound for the movement of a. Since ag > -1 35)rc3 and |Exl = 10 (LaR + * < Lont we have
of > 2/2 dgk + Nu 5 j 1/2 (-)(11 -
38)x²) (01-388)w33 + 1 1 = 16 1n.t
(102)
Since 1b4l < K-1 and laxl K- condition B.2 is satisfied and we have ax+1 ag bRn
ag-
16 5n 3.3 Given that 0g I a1 1o 5k² is by the inductive hypothesis, ax+1 S a. I6 ōx
(k- +1)
What remains show for the inductive step is that K6+1l O There are two cases to consider.
When [& E (1=5 , 5xx"), by Corollary B.5 we know [St+1l < (1 -4x3)|EA < KL < ōōx
When [xl < 16
1 dn', since [5x+1 (1 10-330 3212 )xl < obgn' we have
|Ex+1 Ex| < 662r5* 4 + +32b2 Ex| S 662 + ( - 330B
(103)
Since b% < 2n² < 1 1Ex+1 - Exl < 1 35 = 1 dn Since (Exl S
1 inn, we have [sttal <
48.
18
10
|Ex+1 Exl + KEal < Lōx² desired.
In summary we have P(k) implies P(k: + 1) unless ax+1 < < k5 Note that there must
K exists some t such that a++1 < - { (1 - k5 since for any r 2 486 486-1k 3 + 1, if the induction proceed to P(T), then a1 16
1 5x (r - 1) < a1 3rd < <kà - 35)k², which violates P(r). Let T < 485-1 - + be the first t such that ag < - {/( (1 3,)n², then by construction we have for all < T, P(t) holds. This completes the proof of the lemma.
Corollary B.6. Following Lemma 14, ar E (- - (1-
to (1 + +0))x²).
Proof of Corollary B.6. Denote T by r, by definition of T, we know P(T) holds and therefore
have b2 2n², a > {1 )² and ar > a, 5b2 Combining above have
or > 0r - 10k2 3.3 so a < ar + 10m # < (1 - k5) + 10 Note that since we set
8 0.04, we have
1/ (1 + 28)n² - (- { (1 36)
= (40
( + 1)))
1 500 K.3
(104)
Since - > K3 > 5000, we have 10x 18
<
500
1 3 = 10
1
(1 + 20) -
- 36 n² Adding
(1 - k² on both sides, we have a, <
-1- 38)
+ 10 < < (1 26) Combining
with a, - (1 30 concludes the proof.
47

---

Published as a conference paper at ICLR 2023
B.5.3 PHASE II STAGE
Here state the lemma which proves the final convergence of the step trajectory. The proof is very similar to that of Lemma except a is negative now and [6] is decreasing.
Lemma 15 (Final Convergence). For all k < 1 kKk for all do, bo satisfying op (- k6) - 10 + 26))²) and ISol = |b 1 agg 16" k < 1 f6x,, for all O. there exists some T' < 25 log((-1) such that for all t 2 T, 16g < e and a, E (- -- 16))
8
Proof. We will prove the claim using induction. Consider the inductive hypothesis
P(k): [bxl < boo to + 26)) ag as E [0,8v/2(((b)) [bxl)).
When l = O, the statement holds trivially, so we proceed to the inductive step. Assume P(k) holds for some N, we want to show that P(k 1) holds well.
First we check that (ak. bk) with K. < 1 K- satisfies condition B.4. 16
Since (1 1/ (1 + 26)) < 1, by the inductive hypothesis and the initialization condition on [bo] have 6xl lbo]. also from the inductive hypothesis that |ax laol. thus to show |6g 2V2
_ /Vasaal k. one only need to show for k = case, which equivalent bū agt.
From the initialization condition, must have j0 + 16 1x <0 S by. follows that
|6ğ F 1c 5kk. Since an 1/ (1 + 25)
bg < 2 ann + 1 + A²²
= 5
8 agn + 11 + 26) aax
< $ 10( (1+ 1-+200 21xx+)x++++ 1/ (1 + 25)x** 5ao
8
(Since an < - (1 +2)
(105)
= (1 + 25 + (1 25) g00
1 ag k.
From the initialization condition we have do - -1 (1 + 28) 3 -0.108 < <- -1 lax| >
lool 16 1.3. For upper-bound on Jax| note that ag > 11 36) 3 - by initialization,
combining with the inductive hypothesis we have as > -x 2x((bol - (b/)
-x² //b/ 1x3 and hence
Since bg < |ao|/s < (1 35)x² < 25 |bl < $m², so as > -n²
|ax| < k K-2x-1 since we assumed к < 16 K-1 Therefore we have shown that K, ag, by
satisfies condition B.4 and by Lemma we have [b&+1 < |b, + agb.
512 we may write the update of by as 1b++1 16 (1++ Since dk ao (1+2))) we have |b < 1b 1 (1 26)) Combining with the inductive hypothesis that [bx| < bol +1 20) we have |ba+1|
Since 1 + agK > 0 as laæl < fs²a and K < K <
[bol(1 (1 + 33++11
Since condition B.4 4 is stronger than condition B.2. by Lemma we have as ak+1 < 863,13. Combining with by - |ba+1] |au bbs we have
862
a2 ax+1
1x |b++1|
8) 8|bx lx?
8 2/3 v |
(Since |bal < 2/2 Mlacls))
(106)
lakl
22/2 aall 2
2V22 1/ -
(Sincelax) 1 1xx
< 8..x..
48

---

Published as a conference paper at ICLR 2023
Since the inductive hypothesis gives (au 0 ((bbll - lbxl) < 8/22, it follows by the mediant inequality that (ag 0x+11 (0++)/((660 1)++ 8v2k.
Thus we have shown P(k) + P(k + 1), and by induction we know P(k) holds for any t E N. Now we can wrap up the convergence analysis leveraging this property.
Since for all t, [by] - [bo[(1 - t0 (1+2 for any e > 0 we may pick > log ds (1+2))
such that for all & T, lbe| e. Note that since |bo| A < 1 and 1 to 26), we have
log(t)
T < log (1426)) (e/bo)))//ll)))))) log(1 (1.08))
25 log((-1)
(107)
For the region of final convergence, for any t we know that from P(t) that ag - E (0,8v2x((bb))
[bel)), so we have a; > Ag- 8/2n: (uul. Since we know |bul < g it follows that at - v2x² 5 n3
1 x² and Ag > n by initialization,
The upper bound of a to + 28) 10 R is
trivial since a is monotonically decreasing.
Thus in summary we have shown that for any O, there exists some T 25 log((-))) such that for all t > T, [bx] < and ar E (- 5 3R.
10 (²))
1
3
B.6 PROOF OF THEOREM 3.1 AND ITS COROLLARIES
With all the lemmas ready, we may now prove Theorem 3.1 and its corollaries.
We first restate Theorem 3.1 here.
Theorem 3.1 (Sharpness Concentration). For a large enough absolute constant K, suppose A < 2000 K-1, and the initialization (ao, bo) satisfies ag E (12x3, -[=-===) and by E (-K-I, KK-))O Consider the GD trajectory characterized in Eq. (6) with fixed step size x² from (an, bo). for any 6. > O there exists T = 2K + +))((=1)++ +lg())lb log((bol -1)===) such that for all t > T, [b,] < 6 and ag E (- k 10 x3).
The proof for the main theorem is very simple after we have all the lemmas discussed above.
Proof of Theorem 3.1. We consider any initialization (ao, bo) satisfying ap E (12k², --xxx-11 and bg E (-K-1, (-K-),K-1)\ (00.. We abuse the notation to let a and b, be the value of a and b after the 1-th two step update from ag and by.
If lbol= 2 2 then by Lemma there exists some TI < -- such that on E (2x 1 ---)) and 1b,, E (var, R, 2 R). If lbul I VVloous by Lemma 9 there exists some T2 < log((bol-))ee= such that br2 E (yaz,, 2 (ar, and an E (2x² 1 x=xxx1)) Thus there exists some Ti O(x + log((bol-1)<- -=)=)=))))))))-) such that or, E (2x3 1 x-1x=11 and bn, E
2
Now by Lemma 11 and Corollary B.3, we know that there exists some T3 S ap 1 such that within Ty two-step updates from (ar, br,) we have aT E Ggx2 2x²) and by E (1 Var 2 +T. Let T. = Ti + 73-
This completes phase of convergence.
For phase 2, since ar2 E ( 2x and |bral E (1 R, k)rzk), by Lemma 13, there exists < lary |n |-- < 26 such that within steps from (ar,, br;), we have ort E (N 4, 2x- and |Erzt7a1 < 6nn where we fixed 0.04. Let T3 T2 + 74-
After the residual [g] decreases to less than 1 Snd with T3 two-step updates, by Lemma 14 and Corollary B.6 there exists some T5 < 488 such that aTs+7 E (- (- (1-3)xx (1+26))) while rol < fōx Let T4 Ta + 73.
49

---

Published as a conference paper at ICLR 2023
Finally, by Lemma 15 we have that starting from
(ara byg ) there exists some TO 25 log(c-)) that
for any t >> Tı + To. lbr| < E and Ge E
(- 3 kx 10 n. Thus we have the trajectory converging to
some minima with a E - x))))
-
Finally we bound the total number of steps required for convergence. Since lool <
1 K 2x
we have T3 < 1 K K 2xx Y. Moreover, since R K-1, we have x=4 O(K xxx-2x 4). Since
T3 = Ti + 73 + T4 + Ts + To. we have
T= T-o L xx # log((bol='))) + log((=)))
(108)
This completes the proof of the theorem.
B.6.1 PROOF OF COROLLARY 3.1
Before we proceed prove Corollary 3.1, we first show a simple lemma on the approximity of and c(x, y) _ /x² - y² when ø is large and |1 -11-a zy small. Recall that was previously defined
in Eq. (31) and the (a, b) coordinate that we have been focusing on is the offset from the - minima in the (c, d) coordinate.
Lemma 16 (Approximity of r). For any large constant K > 512, fix any K. 2000v2 2
K
For any x E (2/ =- ,,,==1 and any y such that |1 ryl < K- we have c(r,y) 4 V y² E (x 32k²3, x).
Proof of Lemma 16. Since xy < 1+ K-1 and x > 2 xK we must have y < 2(1 + K-1)) <
2v2K, where the last step holds since we assumed K > 512. Mcanwhile 1 < K also K --1 > 0, y O.
implies ry
Thus we have
c(z 3) = v yy² y²
(109)
v x2 - (2V2x)2 = x x/1 8n2x > xV (2n2 > sx11 16x²).
where the last two inequality holds since x. > -2 and V > x when x E 10, Thus c(x, >
I 16 > 32x² since we assume x 2x Since y > 0, V2² x2 y² x. so cl, y)
(xr - 3 x).
Now we can proceed to proving Corollary 3.1. We first restate the result here:
Corollary 3.1 (Sharpness Concentration under (x, y)- Parameterizaaion). For a large enough ab- solute constant K, suppose < 5000000 K - and the initialization (zo, yo) satisfies To e
(I + 13md, x 1 3==2 b) and |zaY/e - 1 E (0, K-1) where (ø. () is the n- n-EoS minima de- fined in Definition 1. The GD trajectory characterized in Eq. (2) with fixed step size n, from (Io. yo) will converge to global minimum with sharpness X E ( 3 20,, 2).
Proof of Corollary 3.1. Following the convergence proof for Theorem 3.1, to prove this corollary we only need show that all initializations wo. yo satisfies the initialization condition of Theorem 3.1 after re- re-parameterized to (a, b), and the sharpness À of the minima will satisfy X E ( 3
20 2))
We first check the initialization in the (I,,) coordinate satisfies the initialization condition in The- orem 3. Due to the different contexts, we will use λ = / and 7) itself interchangeably.
Recall from Eq. (27) that # and 3 < x where k. VT.. Since |0050
-1) It is not hard check that J K-1 as required by the initialization, from Lemma 16 we have cp A c(2o, yo) E (xo 32x² 32k,, By the same reasoning we also have
--x-11
8 c(ă Ø) E (i - 32k3, #). Thus zo > 2 + 13x implies co > č + 13 a 32xx3 Since we
K- where K > 512. 32x² k², and hence co > غ č 12x2. Therefore
assume K < 2000
an Co- 8 > 3
50

---

Published as a conference paper at ICLR 2023
On the other end, since 전 < k= as shown above, have + - 2x < 2k_1 ! So again म + 2 b. Since ė > I 32x we :. The last step holds since 32xº < K 3 3.
apply Lemma 16 so that co < c(t + 3 2,
v0)
have co ē K 2n 1 32kb < 5 22+
Therefore dn - Co ė < 1 1 3. Since [bb| = =| 11- zyl E (0, ()) by construction, we know (ao, bo) satisfies the initialization condition of Theorem 3.1, and we can have the trajectory converging to global minima with a E [- TØl and = 0.
20
Now we show that for global minima with satisfies ds E [ 3 10. and by = O, the sharpness
20 satisfies À E ? 3 () -)
Recall from Eq. (3) that the sharpness of (I, near the global minima is given by
X = 1 (x² +") (37² - 1) + V (x² + v1) - 3/ + 4 10,² + 7711) where y 2 xy. When (z, y) is a global minima, 7 1, and Eq. (110) reduces
(110)
A (1=²+
(x²+ 3)²
12((x² ()) 2(x2 + y²))
= =2(x² -2(x²+ +1²)
= =22 + x 2/2
(111)
= 2/ (x² x 2-2
= 2/ 4 4 (V y224
= 2/4+ 4 cc.
Since = (x-4 - 4) + where
| 5 k3
1/ we have < (x-4 - 4) and hence
2/ di 2 + x-
2n
Now we prove the lower bound for A. Follow from (N- - 4) gk², we have
(K-- 4) (n:))
(k-4 - -4) (x=
- -4) - g (x= 4) 4))+) (n (x=4- -4) 4)
1 (n -1 4)- )
4
(x- - 4)
2
20
( (n= - 4) 2) (15 4 4)
(Bernoulli inequality)
3
- 4
20 xx (x- 4)
3
k 4__ - 4
3 (k-
= 4
20
(112)
Thus
= 2 c c²
2 2
4 -
20 3
2 -- 20
(113)
2 1
20
k
-2
3
2 2(1
20 3
Kk
2k
20 3
37.
20
Hence in conclusion, the converging minima has sharpness À E ( proof for the corollary.
n, ;), which completes the
51

---

Published as a conference paper at ICLR 2023
B.6.2 PROOF OF COROLLARY 3.2
Before proving Corollary 3.2, we first show simple lemma on the approximity of i' and x- where I is the x-coordinate for the x²-EoS minima.
Lemma 17 (Approximity of # and k= For any large constant K >> 512, any B < 2000 K- With & - v2 (1 4 + n 7) + + 1)å we have l K ==1 <33x)
Proof of Lemma 17. Since the condition for x is identical to that of Lemma 16, we have | E < 32x² from Lemma 16 where č = (x 14 - 4) from the calculation in Appendix B.1. Note that
č= (K 4) = =x= (1 - 4x +)= > x -11---0 4x x- 4xc².
(114)
It is straightforward that é < x=1 so |é - K -ー < 4x². Combining with | E) < 32 we have
[ - x 1 < 36x², which completes the proof.
Now can proceed to prove Corollary 3.2. We first restate the corollary.
Corollary 3.2 (Sharpness Adaptivity). For large enough constant K, fix any a < 2000 K For all initialization (Io, yo) in the region characterized by
To E (a + a 1 t
(10)
and |oo0 1| E (0, K ), the GD trajectory from (zo. yo) characterized by Eq. (2) with any step
size 1j E (a² 2 TO K a2 a²) will converge to a minima with sharpness X E ( 3 17 7
20
Proof of Corollary 3.2. To prove this corollary, we only need to show that for all n in the required range, the initialization region characterized by the corollary is a subset of the initialization region required by Corollary 3.1 for that particular 17.
For the ease of derivation, we will use x² to substitute for n7. Since are dealing with different step sizes, we augment our notation to let (H,2, denote the =²-S minimum. Note that the condition for you namely [ooW 1| E (0, K K-1 is identical to what is required by Corollary 3.1 so we only need to show for any learning rate n²² E (a²
1 1 aa?
10
(a- Ia' 1 a aa _ a
1) C (Øxa + 13k H 3
(115)
We will first show X. + 1 <a +
+
K- a
T5
Since x² > a² - 1 K aa² we have K > (1- To K -))) > (1 - = 321 K- = (1-
H K k-² where the last inequality holds since may assume K 512. Taking the multiplicative inverse, we have -= < (1 1/ 1K- ---K-2-aa"
Now note that since K-1 < 1, (1+
16 1 )1 18
1 K 144
1 1, so
238
(1 - 1 -- < 1 1 K2 and hence A
< (1 - 1 -ー
(1+ 1/
))a -1 From
T6
Lemma 17 we know ش22 < k *+ 36xx3 so
# + 13x 8
< K + 36 3 + 13k < (1 + 1 Kー 3a -1 + 36x 3 + 13x7. 8
(116)
I6
Note that since < a < /
K-1 where K 512, we have
36x + 13k a < x²
22 < 1 K- <
15 - 1
16 k 2
(117)
5000000
480
Thus combining with ش. + 13k < (1 + 1 K 2)a + 36 + 1, we have
3x- + 13k3 < (1
1 K
aa
- + ( 1 K -1 < +
1
2
(118)
16
15
The other side is much simpler to show. Since π a, we have -1 > -l
From Lemma 17
we have 2hs2 > x= 36x so we know 36x²3 < 480
شو + K k -1 > a 36x3 3 + 5 K 2x-a0
aa From Eq. (117)
1 K- K- = ( - k Therefore ش. + K- >
a + g K- -1 This concludes the proof for the corollary.
52

---

Published as a conference paper at ICLR 2023
B.7 OTHER AUXILIARY LEMMAS
B.7.1 CONSTANT BOUND ON Ra(K)
Lemma 18. For all < O.1, for all (a,b) b) E (-x (-1,1), there exists some absolute
constant K independent of a, b, to such that
- V (1) an + (1 - 4x21) = =1 2ak + a²n²² 2b(2 + b)n² + KK.. (119)
+))
Proof of Lemma 18. By explicitly computing the derivatives for ō with respect K, know that is O with respect to R. and have the Taylor expansion
8 = 1 + 2ax + x²2x² + 26(2 + b)x* + (K(K))
(120)
where RR(K:) is the Lagrangian remainder that Rs (k) = 0²6 the lemma we only need to bound the Rs some absolute constants. For simplicity of notation,
arcs (z) 120 for some x E [0, x]. To prove
denote e (1 4x B + az, 7 1 + b, and (472** x4 + 8²)² Moreover, let
4
P1 = a
12k²
48
P2
a²
24k
432k
1344x
(121)
P3
a
all
24
2448
24192k
59136x
P4
q
a
a
a5
10080A3
262080
1774080
3548160e
P5
a
al
a5
19
By doing some tedious calculation have
058
105 (16-2*9
4,33
75 (16-223 + 4,33p1) (48-2R2
12,8 1 + 4 p2)
(x)
0k5
320
8o37
45 216 36
+ 483 P1)
2 + 12 p + 4 p2)
+
805
15 2 + 4.3 pi))
(9672
24 3 + 363 1033303pp222
P1P2 + 4,3%
+
40
5 (487 + 12 p
4/3 P2
(96 + 24 368821102
P1 + 4,33ps))
203
5 + (96- 24p + 1448
36,3
48,3 1188p11P
PIP3 + -.33pa) p4)
3
240pfp2 + 3603 2403pips
2408pp + 1203² P2p3
603
603
PIp + 433
2Φ
(122)
0.99 and = (472x* + +/34) 32 (a+ dr)² > a > 0.9. Thus a and Φ are bounded from below by some constants. Since k is bounded from above, we have P5. bounded from above by some constants.
Since i < 0.1, we have = -1- 4x
Now give upper bounds for and y. Since |x| < 0.1, we have (1 (1=4zt) à < 1. Since |a < k= we have 3 = a + az < a + ak < + < 2. Since |b] < 1, we have y = 1 + 2. Note that is straightforward from construction that all these terms are positive.
53

---

Published as a conference paper at ICLR 2023
From Eq. (122) have 8 (z) as some finite degree polynomial of B, 7/ Pl,..... Po. and ø- Since the absolute value of x, we have completes the proof of the lemma.
2xe all of these terms are bounded above by some constants independent
8 0r. (z) uniformly bounded above by some constant K for all [0, k.). This
54

---

Published as a conference paper at ICLR 2023
C
THEORETICAL ANALYSIS ON RANK-1 APPROXIMATION OF ISOTROPIC MATRIX
In this section we prove the convergence of the vector
C.I PRELIMINARIES
Model: We consider generalized model from the scalar product case.
min 2.yE
- xy ay 11
(123)
The normalization factor added show the equivalence this rank-1 isotropic matrix factor- ization problem and the scalar case considered in Appendix B.6. We will show how the equivalence achieved due the alignment in the next section.
Update Rule: We use gradient descent to optimize this problem. By computing the closed-form of the gradient, we know the 1-step update of æ and y follows:
1
Dt+1 It 7) (xr y:
2 a ycl
yz
2 (a. y 11 12
12:
(124)
1
1 (2. y1) 11
Ui++ yt + 17
(x. yo y1)
a vil ye
x,
2
2
With the gradient descent update above, we can have the dynamics of the following four quantities: | /1 112, and (x: y)².
|
|,,,| F 2n ((x) y² (a) 3e ( 11 [|.11111/111 3/
4 (4(x)) y.) | + (x:7 wr ) 2||
11
3(x y "1>,1
A ui) ||a:e|| I | 4 (1
(125)
||
l1 3 + 2n ((x)) y1) (2, v. P1ac, |11l3r:l)
+ n (4(x)) y.) 2 || + (nn 91 |11a
|1 + 3(xr ye) *1,r 1l
- (xx T yd) ) 11 11 4(x, y1) 11
1
1 [xx 2²)3)
(o T (x14131112
(x, y²
+ n[
11 (1)
||
2
2 + 4 (x) y1)² [4(1]=..1|
2
11 11)² ||<
(1
|| (2r yz)
(126)
Furthermore, we define an alignment notation for the variable: & := 11 - (xx, 31)2 By Eq. (126) and Eq. (125), we can derive the following update rule of g.
E++1 = E (1- (x,, (xx 21
||a 2
|| +
+ ((2
||: y 12)2
|| 1 (x ye
2
4
(127)
C.2 PROOF OF THEOREM 3.2
We restate the theorem of the vector case.
Theorem C.I. For large enough absolute constant K, with all the initialization (wo, yo) satisfying
æo d1 Yo Sn.nn(((3d-1), dzd 2. 82 (i J/ - + K-2 -, +), if
55

---

Published as a conference paper at ICLR 2023
step size 17 < min K-1
[o)] od)] }. and multiplicatiive perturbation y/ )²= is performed at time t = tp for some ip > P ((-log(n) + log(d) log( 6o) + K²),, (dd) there exists T O 2k 4
K-
8000000 20000+2000 log
yi(1 + 2K 2
then for any > O, with probability p - 2do 2
log(c) ( such that for all t > T, f(x, y) < c and | l| + llvell E ( 3 7) 1
10 1
In the following analysis, we still use the operator (-) to only hides absolute constants, and K to represent the absolute constant that uniformly upper bounds all the absolute constants of the (-) terms. Also K > 512. We still use the notation of 17-EoS minimum (ž, Ø) of the scalar case in the vector case.
We first prove some properties of the global minimizers. The following lemma guarantees the align- ment of the two vectors I, y at any global minimizer (z, v).
Lemma 19. (Global mininizers) For all global minimizers (a, y) of the optimization problem Eq. (123), we have a = cy for some c R. and |ل
Proof. We directly consider the objective || xy xayyy" T TIF
| - xy zy 1| Tr((I- xy xy ) (I xy xy })
= Tr(I yx yz T - -xy xy (x 11
d -2(x² T y)) + (x v<)y)>((1 /1/11
d + ((x 3)² 1)a + (a v) (|= 2)
2 d - 1.
The global minimizer takes value 1 and the equality holds when
y)= 1)² (ax T ([(=r] || (xx v)) O.
which is equivalent
cy for some e R, and ||>||| | = 1.
Now we consider the formal proof. To prove the theorem we have four steps: (i) we prove the two vectors 2, y will decay geometrically after T O(-log())) - log(6)) + log(d) time (Lemma 20, Lemma 21): (ii) we prove the norm of the two vectors | |/3/|| will satisfy the initialization condi- tion in T < 6 + 4K time (Lemma 22); To escape from some sharp minima, we add some deterministic perturbation, and we prove that after the perturbation gradient descent will re-enter the feasible regime, while at the same time keep constant distance the manifold of the minimizers; (iv) after entering the feasible region and & is small enough, the re-parameterized dynamics of ||>||| and || can be captured by the scalar dynamics (Lemma 25). Then, we can reduce this vector case to the scalar product case and finish the proof.
Here we present the following key lemmas to prove the convergence at the edge of stability. For simplicity, we denote att)= := 1| + || [la) -.
Lemma 20. (Alignment) Assume 0 < a()) <
8. 0 x yo < || < 1.2 17 < 0.01.
Then as long as 100 < a(t) < holds for t E [0, for some t1, then for all t E [0, + 1].
0 < (x. y <
13 and Ei+1 < E.: Moreover, there exists some to < +log2( 20m0 (0g2(20m0) yo
for all
10
t to. + 11.
7 (z. y1)
13 and Eu+1 < 0.7E..
20
10
Proof. We use induction. For iteration O, the induction basis holds. Now we consider iteration ++1 when assuming the conclusion holds at iteration 0, 0, 2,,,,..... ...,...
Consider Equation (126). We denote G (x. y:) And from the induction hypothesis, & < |||0|| < 1/4 for all time. Thus ||x;; || = C. +
< 31 Then we have
20
1001
C+1 < (1 + + a()) cx) 20m ca(
+a(t) (1 + a())a.)
500
V. (1 + 2K 2k at time tp means at this iteration, we multiply (1 + 2K 2K-1) to the vector y.
56

---

Published as a conference paper at ICLR 2023
Since
C
13/10 and -0.01 a 1/8, this function takes maximal value when C₁
1001
1600(1-af9))
1001
1001
1001
+ a(t))²
13
Ct+1 < ce(
+ a(I) - -1 (1+
(1)) (1
750
500
1500
1 a(((1)
10
The lower bound 0 of c is straightforward since have
c++1 ce(1 + (1 + a())) Ct - (2/2))
due c + & 31/20 and -0.01 o(I)) 1/8.
Then we consider the tighter bound and a faster decaying rate after some to- If co > 7/20 and the lower bound holds at 0, then to logz 20ke No Consider 2 to.
Then we begin the induction from to-
& E (xo )n))n((1 2oll + |1 13)223)/22 1
and the lower bound of c+1 becomes (since 7/20 < < 1.3)
C++1 > min (1+ a(t) - ca(11 + at Er/2)) {//2))²
> mi 9 (1 c 1/8))² ae (1 + (1 -111/8))²
8
> 7/200.
The last inequality is because of the monotonic decrement of the last function, which takes minimal value at G 1o
13 This proves the first statement.
Then consider the alignment dynamics by Equation (127). The factor can be bounded as follow:
(-1+ (x T 1) (1 1/221/222222 + |/1 /2 (2 ||> 13) || (z //4)
3
< mmxx(1
20 x 2
x
10,/2)² < 0.7
2
By induction, we finish the proof.
If co 7/20, then we prove that a will eventually become larger than 7/20 after some time to, and then apply the same induction process above to finish the proof.
Still consider the lower bound of the dynamics of cc. If have Ga < <7/20.
C+1 > (1(1 + (1 - Et/2))) > (1 + - -11))) > 20,
Then it at most takes to = [log2( 20mm yo
7
to satisfy the condition. Now we finish the proof.
We first consider the time & takes to become smaller than n
Lemma 21. (Alignment convergence time) Suppose all the conditions in Theorem C.1 holds. Then
with probability 26, 2 exp( there exists some time To loogo.7 (17²) ) + logg 21d
205
O log() + log(d) log(ōu)) such that & < 172 and ||æ,||| (3 + 200 K-2, 2n 3. + 1 K -2 n +)
6
for E (oo, To + 6K3 + 4K].
Proof. We first prove that the bound of (x:T²)² y)² can be reached with probability 1 25g
2 2exp and meanwhile the alignment & begins to shrink after to [log2 21d
)7.
206g
For the initialization, 855, 2. and by symmetry we know
82
83
Pr[(2]] y²
= Pr
Pr(x
3d
1.5d
57

---

Published as a conference paper at ICLR 2023
It is equivalent to consider sampling from from.N(O. (0, I) and then divide it by its norm. By the initialization condition, we apply the Gaussian concentration bound and Theorem 3.1.1 in Vershynin (2018) and get
% 3d
Pr[( yo)²
< Pr(zf
< o9 + Pr|||=]| < 1.5d] < 260 +
Then with probability p 1 - 280 - expl f ((d)),
o d (x) yo)² During t E [0,T+6K*
4K], we can apply the first argument in Lemma 20 and induction prove that (x) yo)² 10 and
0 < att) <
13K
For = 0, since 62 E (
( , 1
and 10. Then we suppose for time t E [0. t - 1] the statement is 1] the condition of Lemma 20 holds.
20 K < a(0) < 1 Kー
(x yo) < ||00| |_=o||||||||? correct. By the induction hypothesis, we know for t E f0. f1
4 <
13
Therefore, by Lemma 20, for all E [0, 41. (at, y.)²
13
I0
With the upper bound of (a) y,) the one step movement of ||0|| + || can be bounded:
| 112 | 2|2 | 1|2 1 1
= (1<=,112 + ++ (77 yn ((11=r 2lyell - 1)2
3
+ (2 1
11 || welll111 11)
- 1 (2x ye ))) 4n(x, y:) y:) (1
13
< (1 + (t))
(1+ 2m < 5n7-
10
Fort < To + 66 + +4K, the total movement of ||x, + |1 12 is smaller than 5n + 6K* 4K),
and the corresponding movement of a(t) is smaller than 57 + 6K3 + 4K) < K-2
200 (since
1 min 5000000 2000+2000log
K-4
K-3 50 Then by the second argument of Lemma 20, we have
}). Thus by induction, (x) 31) <
13 and 0 < a(t) <
10
K- K < 13K
20 < (क) y:)² 13
100 4
and E < 0.75, for fort= > log2( 21d = (log(d) - log(6o)). 206g
After t = [log2( (
200g
21d )1, we can calculate the time when & < 172 It needs at most loga loga.7(17²)
=
(log((n) for & to become smaller than r/2. Now we know with To = loga. + log2 2065
21d
O (og)))() + log(d) - log()))), within 1 E (To, To + 6K3 + 4K.. & < 1/ always holds.
On the other hand, we also have the lower bound of the || 12 || The one step movement has the lower bound
| 2| + | 1| ||>e|| 11
-1|=,|1
|B: ()) (zx y. (( |l 1)²
3
+ (2
|
|l (xx yr) 1)) + 4n(2) y1) -
> 0 - 2m > -17-
So for 1 < To + +KK + 4K, the total decrement of ||,, 1|2 ||/.|2 2| is larger than (To + 6k +
4K)
--
200 " 1. Because x, Y: (Z) (20 13),
|| 13
?
10/
1//l|= /| < 1000 looK K22 1. Therefore,
| > 200 K K-2 +
After we have the alignment guarantee, we can prove that we will enter the feasible regime for the scalar case with high probability. Still, we denote c, (z) y.)
Lemma 22. (Feasible regime guarantee) & < 17² < 300000 K 1
ToO a(t) < 1 k
for all E [t,, t + 6K3 + 4K], then there exists some T < 6K3 3 + 4K such that for all t e [t1 + T,4 + 6K + 4KJ. xx 1llll 1| K-1
Proof. To prove the statement, we try to prove a stronger statement:
e E
58

---

Published as a conference paper at ICLR 2023
If this statement holds, because & n/,, the product of the norms
| 11 y 11? E (1
K ,1+ K K +
So have | (1 3 K } K - and the lemma is proved. So in the rest of the proof, we will prove this stronger statement.
First we prove that there exists some time t < 4K + 1 that the square of the inner product c, will be larger than K-1.
If c E 7 1/2), then
7
c++1 > (1-1411111/1------11 - >
x 1.42 1/2. 20
Thus C will provably enter > 1/2 in one step. And if 1/2 2 < 4 < 1- - K-1, we have
1 k > a 1 K- a1 + -1-11---00 (1))- c1 +
C+41
which means it takes at most 4K steps for c to become larger than K-1
Now we prove if G > 1 - K-1, then it will take at most T = 6 steps s.t. for t' > + T, K-1 K k-2), which finishes the proof. We first prove there exists some
or E (1 - K1 1 1
E (1 - 1+ K-1 K K- afterwards we prove that once gets in, it will never get out of
the region.
If at first E - - k-2), the proof is done. Otherwise, we suppose et > 1 K K-2 Next, we consider the two step dynamics of c and prove will decay by constant factor every steps.
First we pick out the relatively small terms, and find out the main part of the two step dynamics. att) ao 1/8, and 1 + K - K- c 1.3, we have:
c+1 < (1 (1+ a) + 10m2)²
G41 (2(12(++++1)) (1(0)( a)) 2)
C42 (cn[((((+((((+ c(+) 10m
cc(1+(1+a()) q) 10n (1 a(t 1))(1 +++1)++++++++++)) 10n 107
Then we upper bound the difference att + 1) att) in one step. Consider the dynamics of ||x:|| + ? The difference each step is at most
|| + 11 - || |
3 ||
-[|=|| + 11)l²) ( ((1=)|])) 1)a (2
1 19) x |illl 12 (x (xTT))) 2))
+ 4n(x (1)(xT))(11)))))))) - ||=|| 1))l))
(1
13
< (1-(1+ at))
(1 1) 2n 5n
10
Then the corresponding update of a(t) will less than 5x/2. That means we have
C+4 <42 -()))(((((()))))))) c.) (+)
1)) c+1) 10n2)²
- +
-
[000==0000000))00000000000000000 - -x))((())))))))))))))))))) 1007?
< -
+ +200m7
(128)
59

---

Published as a conference paper at ICLR 2023
Similarly, have the lower bound of the 2-step dynamics.
- Cx+2 > + - - 21
(1)( -2+1) -
>
- (10_())))00000-22(2((((+1++01111
- )/
>
100m?
>
+1++a())) - 100 - 100n?
+(1 =((+((++a(t) -+)))))((1)00000000) -0)) ((( a1) -))) -
(129)
After bound all small terms to 200y/² we consider the main part
(1 + (1 a()) (1 (1 a()) - c11 +
(1 (1 a - cc))))
To prove it will decrease, need prove the factor
f(a,c) =
a) (1 (1-+a)(1cc))) c)) <1
when a E (- 1
1.3.
We first prove that the function f(a,c) c) monotonically increases when
100 increases, i.e.
af(ac) > O. We directly give the simplified expression of this partial derivative. d0
Af(a c)
15
17
(c 1)( 2a c(1 a
+ (2c(1 + +a)
22)
da
16
4
15
17
= (c 1)(- 20 + (1+ (1
a
c (2c(1 +a)
- 2c)
16
4
> (c -4
15
c+
81
(1
-r)²))) (since
1
100 < a < 1/8)
16
4
2
13
-
(CE (1,
10'
>.
(130)
Therefore, f(a, c) < f( for all 1+K-1-K-2 < < 1.3. Denote 3 = { 444) and suppose c= 1+ t/B for some > V2 (since + K-1 K-2), 3). Mcanwhile, since t/B 0.3, & 0
-K-2-2K-3
Then we plug in e = 1 + 1/3 and expand (3,
= =
f(8,1+tv/))
11+2253/2 -8 -
5/3 -6t2 + -3 37/2
- -312 +3 3
< 1 + (2t - 311---1-8010000-1-0000-3033331"3
5t /2/2
6t
+ 1.2t3/352
< "","
0.09t
<1 3/2
(131)
Thus we have
C+2 < a(1 -
2 + 200m² <a - <a
K
10
Thus it takes most 2 x 0.3/ 10 K K-3 3K3s steps for e get into the region E (1 (1-K-1,
K-1 K- 2-
Finally we prove that if e, E (1 -1 1, 1 K- K-2) for some 8, then for all { 8,
4 E K- 1 K- -
60

---

Published as a conference paper at ICLR 2023
We use induction and suppose a satisfies the condition above. Note that we have the upper and
lower bound of C++1 above:
c++1 + (1 +1+ -)) (00+0000000000000000000000000000
100r?²
C++1 2 -c(()(((((((00))))=-=-)33)---3 2 c + (1+ -c))
Meanwhile c + (1+) is monotonically decreasing with ot when - Too < a(t) <
K and 4 E 1- 2- Thus we have:
9+1 sc( =²)² + 100n
+ K 1 (1+ K 2) 100T?²
1
=1 -1+K - 2 +
-
K 212
+ K K-3 3 K k
+ 2
2-2 + 100
K
2
<1+ -1=K-1KK-
941 2 a(1 (1 a (t) (1 - ct)) 100m?
> (1+K K K )1 (K-1 - K- 31
K
2)2 100n)
K-8 2
15
11
1 - K + 2K 3K -4 + +4K -5
K-6
+
Kー
4
4
3
+ K-9
K-10 100n?
4 4
> K
Therefore, by induction we prove that for all s. 1--K-- In this way, || B - K-1 for > t T.
K-2).
Entering the feasible region not enough for establishing the equivalence the sealar case. Fur- thermore, we need the alignment variable & to be small enough, such that some O(-) notation term can contain all the terms with & in the (a, pa-paaammttrizotionn dynamics.
However, we need to guarantee that the trajectory does not converge to an unstable point near the minima. We can prove that for the scalar case if the initialization is not exactly on the manifold of minimizers, but it becomes more challenging in higher dimensions. Therefore, we require an additional perturbation to escape from any unstable point. We pick the time fp = To + 6K3 + 4K = log(t)) log(d) - (6)) k3) 3 to guarantee that when the perturbation happens, æ and y are aligned and (xr ")² is not large enough to cause instability. After the perturbation at lp. the gradient descent dynamics prevent the objective from hitting the manifold of minimizers.
=| + The following lemma proves that the properties of the bound of e == TxTy)² v) [|y]l3) and & are still valid after the perturbation.
Lemma 23. After the perturbation at time Im we have the following properties hold in l Mp, 21 +
2): After the perturbation d, (1 + K-1 2K 3K +2K-2), and e E ( 20 19); (i)
7
K-2
< <t(() 1 K-2 (ii) & (0.7)
100
Proof. We prove this property by induction. Firstly, we prove the basis of induction at t tp.
For (i), before perturbation we have Ct E (1- 1 K-1- K-2), so after the perturbation we
have d,, E =1+ - 2K 2 1+ 3K + 2K 2) C (00 1
For (ii), the initial value of (x({) satisfies that 1. k a()) < 1 K But the total movement of a(t) before perturbation is bounded within (-5z f,. tp)
C (- 200
k- K=) ) and the perturbation
introduce a movement of |//|2 by ||y.1|2 < 13 / x. 1000
K 2 1. Due to the upper bound of
n, (1)
K-2 K-2
100
61

---

Published as a conference paper at ICLR 2023
For (iii), since K 512, thus tp. & (0.7)ok*g & (0.7)6k² rp. 172.. So all three statement holds for tp-
After the perturbation
St = ([1+2K-1)2E
Then we suppose for i E Un lp (1] all three statement holds and prove them for hp + l + 1. First by the dynamics of or with condition o E ( 13 and < a(t) < K 2. we have C++11 < TO
ce(1+ (1+ - -1))))) 13/10 and c++1 > c (1 - ca) - 10n > 20
Statement (i) is
proved.
For (ii), since et and & are bounded for 1 S 1p + 4 + we can still have the total movement of a(t) lies in (-5nj"t, (p) C (- K K 3) ). Combine with the movement of y perturbation 200
6K ok-2), which proves the bound (< Tobo K -= b), the total movement is still bounded by (- 1000
in (ii).
Finally for (iii), as long as < proved by induction.
13 1o Eet1 < 0.75, < 0.7 n/2 Therefore all three statements are
After reclaiming all the bounds after the perturbation, we need prove that the alignment variable & will be small enough to approximate the scalar case. The following lemma proves that b stays at constant level when (tl) some constant.
Lemma 24. Suppose & 17², 1) < 8000000
K-1 K-2
100 <a(t) < K-2 for all [11 261 2]
for some t1. If (2,,, Nt, ) (1 (1+1) K 2K-2,
3K-1 +2K 2K- then for all E [t, 211}
(x) 1 +2 +2x)
) E (1+ + 20 -1 1 + 3K-1 + +2K 2)- (x)
Ur, +28-
1 K
20
N, k 11/2
Proof. Denote Gt := (a. yt) Here we consider the two step dynamics of the inner product (Eq. (128) and Eq. (129)).
C++2 (1 + (1 (1)) at) (1 (1+ (1 (1 (1))))))))))))) 200m) C++2 > c((1 (1+ (() (1++ (()))))))))))) (1 (1+)) c)))))) 2007
We first consider the lower bound of the (a) e, +2k Vt, +2x) We prove the sequence of the c+22k k 1, 2, ... by induction, and then use this conclusion to prove the upper bound of C+42R-1- We know that k = 0 the statement holds. Then suppose the lower bound holds for k S. kn, and we prove it for k Kg + 1.
Here we consider the function as we do in Lemma 22:
f(a,c) = (1 + (1 + (1 ()) + (1 + a1 a) c(1 (1 +1++ali a) c))3))
Denote 7 = 100
1 K-2 By Eq. (130), we know
f(a, > f(%,c)
Suppose +t+
1 + +.9.7.9 (g,30 40)- Then the expression f(1, c) becomes:
=
= 1 -1+ 2g7 3/2
2/333/2 -
392 + -²
+ + 29 - 5/2/ 6g2n3 + 4g4 - -39377/2
-3927* + 6 33 9/2 20²73 3 11/2 7/76
Notice that when e > 1, f(7,e) decrease as e increase. Now we consider the range of q: [ ). we have:
2 f(7,6) 2 f6,,1
(7)
30/3/2
5/2
7/2
8-99/2
11/2
+
27
27
9
27
27
62

---

Published as a conference paper at ICLR 2023
167
647
4y
152
92
+1
81
81
27
81
81
10
K
> +
3/2
(since K >> 512,
71/²
27
10
That means +2,22 ((1 K 2)2 200nn))c;+22 +2 (1
k_) and the proof is
27
20
done. Otherwise, we have E (5. 30 + 4), also we have the lower bound of this function:
2 f(1,c) f(%, 3 3(
127920 3/2 319920, 192000, 7/2 + 64000 9/2 + 64000,
+ 25600007 + 102400007 + 153552007 + 10230400y 2555200y +
(since > 512, 22/2 K-1
100000 >
10
That means
Ct 22+2 > ((1 1000000 2200m )ct1-
> (1 -
2001)))(1 +
3 3v
200000
> (1 26214400
1/2 200m)) +
3 v7)
1
K x=1)(1-- K-1
> (1 1000
K
4000
15
K
+
Then the lower bound of +2k is proved.
As for the upper bound, since when t t1 2k, y satisfies all the conditions in Lemma 22. Ct,22R < + K-1 -K-2, we have proved proved that it will never be larger than + K K-2
Otherwise if 42 (1 K-1 K-2 3K-1 + 2K' We apply the inequality Eq. (131) and know Ct +2-+2 < Ct +2k- In this way, by induction we prove the bound (z) +2 E
(1 + 20 -- 3K + 2k
As for the upper bound of (x: 1 2k 3t 421 1), we directly apply the upper bound of 1-step dy-
T
namics (if c 1 20 K ))
C+1 S (((1 (1 o())) 10
(2 - C + ))
(1 + K-1
K-1
(1
2 100,7²
20
20
K-1
1
K 2
K 3 + 100,)
20
400
8000
1
K-1
20
Therefore the upper bound (x)
[
+2
22
2
20
K-1 holds for all l < 1/22.
1. With all the lemmas above, we prove the equivalence between the dynamics of (a, b) and the one step dynamics of the scalar case (Lemma 1, Lemma 2). For simplicity of notations, when analyzing the 1-step and 2-step dynamics of (a,, by). we use a, a' to denote a,, a2+1 and 6, D' to denote b,, b etc. For simplicity of calculation, consider the change of variable k /n).
Finally, we denote
V |
b) -
|
63

---

Published as a conference paper at ICLR 2023
Lemma 25. (Equivalence with scalar updates) If min(n" . ||, E (Ø+ 200 K
200K 21 2, .
(K--)1]] and ||| well - 1| < K K-, the following equations hold for some fixed constant e = min (0.5, 1
125 }:
a' = a 2b²x2 + O k²) + +O(2k(")
(132)
8' = - - 4abx - 36² - + + O x) +0(2()) O(bR))
Proof. For simplicity, we denote x = a and y = |//| We first prove that:
x' d'= x =
R2 xy - -²yy² + O(K*b2).
(133)
y' = y x² 2yy22yy111111 y(1 -1==>²²" OKK66')
We suppose y = ca + 0 for some 0 E Rd and (0, = 0. In this way, we have x: aTy co? co²,,²²
exx2 + 10012, x²1 02 Since { < 1 and 0xx ||||| = (
Then check the dynamics of and plug in y ca O. 1
x' = x + n1(x>> y)y (x y) y
2 x T 3)
(1 + (1+ - npry x + (-nl 2 + np 2 2² / j + (m(x v) 2 2n T y)
Since (a: T y) and ay are both bounded as constant, we can directly take the norm of both sides and with triangle inequality we have:
= (1-"))
Similarly, we have the dynamics of y.
- k (1=+
We now reparameterize the dynamics in (a. b) parameterizzatin.
a' = (a ( - -- b)²xx 4) 0((km))
8' = b - 26ō - -31225 135--33 -88 + 402mx + 16 + 250 + 1965 + 70% + + + O(ab2x2)
where 8 A (n b)² (an' + (1 4n
(134)
And by the bound of ||ax||| and ||| || 1/. we can have |a| < en, [b| < min(1, gex which satisfies the condition in Lemma and Lemma 2. Then we follow the proof of Lemma I and Lemma 2 to finish the proof (since the only difference is the two O(-) notation).
After we can reduce the dynamics of vector case the scalar case dynamics, need prove that in the scalar case (in all stages), for any & 2 (bet1) > 0.7, which means & shrinks faster than bg. In that case, we can conclude that &, > n69 always holds along the scalar case trajectory.
Lemma 26. Suppose all conditions in Theorem 3.1 hold. Then for all t 0, (bLL))+ 0.7.
Proof. From the proof of Theorem 3.1, we know that for all > 0, |b9l < K-1 and |a.l < k x1
Thus Lemma 2 holds and for all t > 0 have
b++1 bb 40,b,x 30² b (a.bbkk) + Oa?bin) + 00 + (birs)).
(135)
Since we know all constants hidden by the O(.) operator are upper bounded by K, we have the following lower bound on the multiplicative update of b:
> (1 4agk 3b, 67 - -Kabbrn - Kakn² Kban Kk3)
b4
64

---

Published as a conference paper at ICLR 2023
2 (1 11- -33K"
K-²
K
R² k²)
[502
2 (1 - 10K 10K-
> 0.7.
512
Finally conclude the proof of Theorem C.I.
Proof of Theorem C.I. By Lemma 21, we know with probability 1 - 25g 2 - there exists some time To = O log(n) - log(ón) + -logdd) that & < nf for all To- Also for I E [To, To - 6k + 4KK] || || E (x++ 200 K 2 - -n) This bound of | guarantees
that for E [To. To + 6K + 4K]. a(l) (0, x22 ). which satisfies the condition of Lemma 22.
Then by Lemma 22, we know for some < To + 6 4K <p. |||aeel Bl - 1| < K-1 for t [t", tp. After entering and staying in the region, we add the perturbation at Ep. By Lemma 23, have the bounds before perturbation still hold: (i) After the perturbation cxp, E (1+K~11 -- 1+
3K + 2K--)) and E (Z) 13 (ii) K-2 < (t) < - (iii) & < (0.7)* Therefore, 100
the condition of Lemma 22 and Lemma 24 are both met. So after another 6K + 4K steps, we have
1b,1 b2 well 1 K-1 and (x) y1)²
1 20 The second expression lead to 4T This means & < (0.7) "k,
[br]
K (1 | 1
x/² < (0.7)20001)
5000000 < <n/b,,
K
By Lemma 25, we know the dynamics of the norm the vectors (I| Il. |/ can be captured by the scalar case (including the initialization condition and the one step update rules). And by Lemma 26, we know the alignment will be kept and the dynamics of the vectors will always be true.
Finally, we apply Theorem 3.1 and finish the proof.
65


Reviews:
[{"review_id": "wZD8H__fDnb", "reviewer": "Reviewer_8B9h", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper studies the \"Edge-of-Stability\" effect on a toy model of deep linear neural network with single neuron per layer (with a further extension to another toy model of rank-1 matrix factorization). The authors demonstrate that under suitable conditions this model exhibits \"sharpness adaptivity\". Then, the paper shows that, due to the absence of high order terms, this adaptivity need not hold for shallow networks, in agreement with empirical observations. Finally, the paper discusses the fractal structure of the global convergence pattern and draws parallels with chaotic nonlinear systems such as the well-known logistic map. ", "summary_of_the_review": "An interesting paper that clarifies some aspects of the EoS effect, but has several weaknesses, especially in comparison to existing alternative explanations (only toy models; no experiments with real world data; somewhat too technical exposition; no rigorous results confirming the global fractal picture - probably the most interesting part of the paper)", "strength_and_weaknesses": "EoS is a popular topic, and there are several concurrent submissions explaining this effect from somewhat different points of view (e.g. https://openreview.net/forum?id=nhKHA59gXz, https://openreview.net/forum?id=R2M14I9LEwW). Compared to these other works, the present submission has a rather narrow scope limited to one or two toy models. Accordingly, there are no empirical results involving real world data. I also have the impression that, relative to the other submissions, the exposition is a little bogged down by technical details and conceptually not very illuminating. Though the model is a toy one, its analysis involves cumbersome computations and does not convey easy intuition applicable in more general settings.\n\nOn the other hand, the paper does discuss some interesting points not found elsewhere, particularly the effect of the model degree on EoS and a fractal structure of converged sharpness.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "ryXfjWmKzBF", "reviewer": "Reviewer_DD1S", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The papers studies the gradient descent dynamics on the loss $\\frac{1}{2}(1-x^2y^2)^2$ as a model for deep neural networks (it can be viewed as a depth 4 scalar network with a specific initialization). \n\nThey show two (related) results: that for small enough learning rates $\\eta$, there is a region of initializations $\\mathbb{S}_\\eta$ where gradient descent converges to a global minimum whose largest eigenvalue of the Hessian is close to $2 / \\eta$. Secondly, there is a range of learning rates and a region $\\mathbb{S}$ such that gradient descent starting in $\\mathbb{S}$ with any learning rate in the range will converge to a solution with almost minimal stability (i.e. a largest eigenvalue close to $2 / \\eta$).\n\nThe authors then show that the same does not apply to shallow networks, suggesting that this phenomenon is characteristic to deeper networks.\n\nFinally to motivate the requirement to start within a specific region the authors show empirically (as well as motivating theoretically) that the same loss can lead to chaotic behavior in the trajectory of the gradient. This might explain the difficulty in proving any global convergence.", "summary_of_the_review": "The paper shows the appearance of Edge of Stability dynamics in a simplified model of deep neural networks. The results are interesting and well explained.", "strength_and_weaknesses": "The phenomenon of 'edge of stability' is the subject of much interest and while there are a number of results that assume convergence at a $\\eta$-stable global minima, it is important to also understand how the network converges to such points.\n\nDespite its simplicity, the model exhibit similar behavior to the ones observed in practice. The reason behind the simplifications are quite well explained and motivated. The authors have done a great job explaining the ideas behind the proofs, both visually and mathematically. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "ZDWmT9Wgman", "reviewer": "Reviewer_wrPK", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper describes a simple non-convex function in four dimensions (a four-layer scalar network):\n\n1. which empirically demonstrates the *\"edge-of-stability\"* phenomenon while training with gradient descent, \n2. is qualitatively more similar to deeper neural networks than previously proposed model functions in terms of *\"sharpening\"*, \n3. and is analytically simple enough to show a convergence theoretically to the sharpest global minimum attainable for a given step size. \n\nThe paper provides a local convergence guarantee (which they call sharpness concentration), i.e., for any learning rate, provided the initialization is in a ball around the sharpest global minima (attainable for that learning rate), the optimizer will converge to that minima. And it discusses the roadblock in extending it to a global convergence guarantee (i.e, convergence to the edge of stability for any initialization):\n\n> the boundary separating converging and diverging initializations (Fig. 7a) exhibits complicated fractal structures.\n\nTo underline the limitations of previous works trying to model the phenomenon, the paper discusses the difference between two, three, and four-layered scalar networks. It shows, empirically and analytically, that the phenomenon gets more pronounced, i.e., there is less of a sharpness gap w.r.t. sharpest attainable minima, when depth increases.  The theoretical results are also extended to the vector case, i.e., rank-1 factorization of an isotropic matrix.  ", "summary_of_the_review": "I believe the proposed model captures the essence of the edge-of-stability phenomenon while training with gradient descent. The paper is well-written and highlights why even their simple model has useful insights about different aspects of training.  I recommend accepting the paper. ", "strength_and_weaknesses": "The paper is well-written and precise. The flow of the paper gives a good intuition for different re-parameterizations, convergence properties of the system, and a rough proof sketch outlining different training phases. Since the paper is supposed to model a complex phenomenon in neural networks, one desirable outcome is to provide an intuition about the edge of stability phenomenon. I believe it does so well. The paper also highlights the theoretical result's limitation and the roadblock, thus providing a balanced view of the contribution. It compares itself to the previous attempts at modeling the problem, highlighting their limitations. \n\nI believe the paper would benefit from some more experiments for completeness. For instance, try to minimize loss (11) and demonstrate edge-of-stability. It would also be good to see some experiments to check if the hardness of showing global convergence is an analysis issue or a genuine artifact. To do this, the authors could explicitly run synthetic experiments on their function and vary the size of the ball for initializing their models. Repeating this for several different step sizes and dimensional problems would be insightful. \n\nIn my understanding, training with SGD may or may not demonstrate the same sharpening depending on the noise level. I am curious if the authors can extend some of their simpler results to incorporate stochasticity, perhaps just simple additive noise. It would be good to discuss how the noise scale qualitatively changes the phenomenon. Intuitively the noise might also smoothen the boundary between initialization that do and do not converge. At least some experiments with SGD can shed light on its behavior. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "XqAZwN2M6xy", "reviewer": "Reviewer_sSz2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper studies the the phenomenon of \"Edge-of-Statility\" (EOS) training dynamics and analyses the gradient dynamics on a simple example (the example is simple but the analysis is quite complicated).\n\nThe EOS of training dynamics, initially observed by [Cohen et al. 2021], refers the phenomenon that when running gradient descent for deep neural networks with learning rate $\\eta$, despite the fact that the sharpness (maximum eigenvalue of the Hessian) is often larger than stability threshold $2/\\eta$, the loss converges in the long run, and the sharpness at the end is around $2/\\eta$.\n\nThe paper examines the function of $f(x, y, z, w) = \\frac{1}{2}(1- xyzw)^2$ and proves the EOS occurs in a local region. The paper also highlights the importance of considering degree 4 function versus degree 2 (or degree 3), and argues that even for this simple example, global analysis is hard since it corresponds to some problems in chaos theory.\n\n", "summary_of_the_review": "The paper provides a simple example and analyze the GD dynamics on it, however it is unclear how does it contribute to the broad understanding of EOS.", "strength_and_weaknesses": "Strength. I believe this paper has some contribution towards  understanding the EOS phenomenon, especially by pointing out the difference between degree 2 construction and the connection between chaos theory. The analysis seems non-trivial.\n\nWeakness. The paper works on some toy case, as pointed out by the paper, even the toy case seems to be complicated and the paper provides theoretical guarantee only over local region. The major unsatisfactory point for me is that it is hard to generalize the finding of this paper to general regime.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "VbCM_Sf9Sf", "reviewer": "Reviewer_jBkM", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This work is motivated by the “edge of stability” phenomenon, observed by Cohen et al.: when training neural networks with step size η, the largest eigenvalue of the Hessian (dubbed the “sharpness”) along the training trajectory often stabilizes near 2/η, which is the stability threshold in convex optimization. The paper carefully analyzes the GD dynamics on the loss function (w, x, y, z) -> (wxyz - 1)^2 (actually because the initialization is taken to be w = y and z = x, it is really a 2-variable analysis, see comment below) and shows that for a region of initializations and step sizes, the limiting sharpness of the GD trajectory is 2/η + O(η).", "summary_of_the_review": "I have issues with the way that the results are presented; namely, I believe that the paper should not mention the four-variable objective if it is not used in the analysis. Overall, I think that the contribution is solid, the problem setting (EoS) is important, and the work is technically difficult, and hence it merits acceptance.", "strength_and_weaknesses": "The strength of the work is that it rigorously establishes the EoS phenomenon with a very careful analysis. Despite only considering a simple objective function, the dynamics are quite intricate and it is technically challenging to establish the result. As it is important to understand non-convex dynamics, especially those capturing the behavior of neural network training, this work makes a valuable contribution to the literature.\n\nHowever, due to the initialization w = y and z = x, the authors really end up analyzing a 2-variable function. In fact, up to rescaling the step size, it is easy to see that the dynamics they analyze are equivalent to GD on the 2-variable function (x, y) -> (x^2 y^2 - 1)^2. Hence I believe that the introduction of the 4-variable function is highly misleading and I ask that the authors reformulate the discussion in terms of the 2-variable objective that they actually analyze.\n\nAlso, it is difficult to claim that this example is minimalistic considering that the analysis and dynamics are so complicated, but I believe it is out of scope of the present work to greatly simplify the analysis.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster