Trajectory Alignment: Understanding the Edge of Stability Phenomenon via Bifurcation Theory
Authors: Minhak Song
Chulhee Yun
Venue: NeurIPS 2023 poster (2023)
Keywords: non-convex optimization
trajectory alignment of GD
edge of stability
progressive sharpening
bifurcation theory
TL;DR: We empirically demonstrate and provably establish the trajectory alignment phenomenon of gradient descent in the Edge of Stability regime.

Abstract: Cohen et al. (2021) empirically study the evolution of the largest eigenvalue of the loss Hessian, also known as sharpness, along the gradient descent (GD) trajectory and observe the Edge of Stability (EoS) phenomenon. The sharpness increases at the early phase of training (referred to as progressive sharpening), and eventually saturates close to the threshold of $2 / \text{(step size)}$. In this paper, we start by demonstrating through empirical studies that when the EoS phenomenon occurs, different GD trajectories (after a proper reparameterization) align on a specific bifurcation diagram independent of initialization. We then rigorously prove this trajectory alignment phenomenon for a two-layer fully-connected linear network and a single-neuron nonlinear network trained with a single data point. Our trajectory alignment analysis establishes both progressive sharpening and EoS phenomena, encompassing and extending recent findings in the literature.

Trajectory Alignment: Understanding the Edge of Stability Phenomenon via Bifurcation Theory
Minhak Song
Chulhee Yun KAIST AI ac.kr
KAIST ISysE/Math minhaksongOkaiist. ac.
Abstract
Cohen et al. (2021) empirically study the evolution of the largest eigenvalue of the loss Hessian, also known as sharpness, along the gradient descent (GD) trajectory and observe the Edge of Stability (EoS) phenomenon. The sharpness increases at the early phase of training (referred to as progressive sharpening), and eventually saturates close to the threshold of 2/(step size). In this paper, we start by demonstrating through empirical studies that when the EoS phenomenon occurs, different GD trajectories (after a proper reparameterization)) align on a specific bifurcation diagram independent of initialization. We then rigorously prove this trajectory alignment phenomenon for a two-layer fully-connected linear network and a single-neuron nonlinear network trained with a single data point. Our trajectory alignment analysis establishes both progressive sharpening and EoS phenomena, encompassing and extending recent findings in the literature.
1 Introduction
It is widely believed that implicit bias regularization of gradient-based methods plays key role in generalization of deep learning ( (Vardi, 2022). There is growing literature (Gunasekar et al., 2017, 2018; Soudry al., 2018 8; Arora al., 2018, 2019; Ji and Telgarsky, 2019a; Woodworth et al., 2020; Chizat and Bach, 2020; Yun et al., 2021) studying how the choice of optimization methods induces an implicit bias towards specific solutions among the many global minima in overparameterized settings.
Cohen et al. (2021) identify a surprising implicit bias of gradient descent (GD) towards global minima with certain sharpness value depending the step size 7). Specifically, for reasonable choices of 1/. (a) the sharpness of the loss at the GD iterate gradually increases throughout training until it reaches the stability threshold of 2/1 (known as progressive sharpening). and then (b) the sharpness saturates close to or above the threshold for the remainder of training (known as Edge of Stability (EoS)). These findings have sparked a surge of research aimed developing theoretical understanding of the progressive sharpening and EoS phenomena (Arora et al., 2022; Lyu et al.. 2022; Wang et al., 2022; Ahn et al., 2023; Chen and Bruna. 2023; Zhu al., 2023). In this paper, we study these phenomena through the lens of bifurcation theory, both empirically and theoretically.
Motivating observations: Figure illustrates the GD trajectories with different initializations and fixed step sizes trained three types of two-dimensional functions: (a) log(cosh(22) (b) 2 (lanh(xzv)+ and (c) 1 (ELU(r)y)² where : and y are scalars. The functions
I22 IR
have sharpness y²² at the global minimum (0, y) for all three models. These toy models can be viewed as examples of single-neuron models, where (a) represents linear network with log-cosh loss, while (b) and (c) represent nonlinear networks with squared loss. These simple models can capture some interesting aspects of neural network training in the EoS regime, which are summarized below:
Throughout this paper, the term "sharpness" means the maximum eigenvalue of the training loss Hessian. For quadratic loss, GD becomes unstable if the sharpness is larger than threshold of 2/(step size).
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

logs (cosh(ay)
(b) 2 (tanh(x)))*
(BLU(x)))*
2-
1-1
300 000
3000 3009
Figure I: GD trajectories align the bifurcation diagram at the Edge of Stability, We run GD on toy models with step size 1j = 2/25 in (a), and r = 2/ 2/400 in (b) and (c). Distinct colors indicate independent runs of GD with varying initializations. Top row: GD trajectories closely follow the bifurcation diagram of the map x - 1z ((z, y) and asymptotically reaches the bifurcation point (0, / 2/1)). Bottom row: the sharpness reaches 2/1/, and xe typically shows 2-period oscillating dynamics. The theoretical prediction X(q:) (dashed lines, defined in Theorems 4.4 and C.4) in (a) and (b) approximates the sharpness along the GD trajectory and demonstrates progressive sharpening.
FoS phenomenon: GD converges to a global minimum near the point (0, 22/n) with sharpness close to 2/n. During the convergence phase, the training dynamics exhibit period-2 oscillations.
For different initializations, GD trajectories for a given step size align on the same curve. For example, Figure la shows that GD trajectories with differennt initializations closely follow specific U-shaped curve until convergence. We call this phenomenon trajectory alignment.
In Figures 1b and 1c, GD trajectories aligned itatively resembles the bifurcation diagram of a typical polynomial map, such as the logistic map. Particularly, Figure Ic demonstrates a period-halving phase transition the GD dynamics. shifting from period-4 oscillation to period-2 oscillation.
curve with fractal structure that qual-
Surprisingly, the curve that GD trajectories approach and follow coincides with the bifurcation diagram of one-dimensional map æ + x 1) oh ((x, y) with a fixed "control parameter" y.
The stability of its fixed point æ = 0 changes at the bifurcation point (x, y) = (0, 22/7), where period-doubling bifurcation occurs. Note that this point is global minimum with sharpness 2/n-
Interestingly, such striking behaviors can also be observed in more complex models, up to proper reparameterization, we outline in the next subsection.
1.1 Our contributions
In this paper, we discover and study the trajectory alignment behavior of (reparameterizedd GD dynamics in the EoS regime. To our best knowledge, we are the first to identify such an alignment with specific bifurcation diagram independent of initialization. Our empirical findings are rigorously proven for both two-layer fully-con fully-connected networks and single-neuron nonlinear networks. Our main contributions are summarized below:
In Section 2, we introduce a novel canonical reparameterizatiio training parameters, which incorporates the data, network, and GD step size. This reparameterization allows us to study the trajectory alignment phenomenon
unified framework. Through empirical study, Section demonstrates that the alignment property of GD trajectories is not limited toy models but also occurs in wide and deep networks trained on real-world dataset. Furthermore, we find that the alignment trend becomes more pronounced as the network width increases.

---

In Section 4, we use our canonical reparameterization to establish the trajectory alignment phenomenon for a two-layer fully-connected linear network trained with a single data point. Our theoretical analysis rigorously proves both progressive sharpening and the FoS phenomenon, extending the work of Ahn et al. (2023) to a much broader class of networks and also providing more accurate bounds on the limiting sharpness.
Our empirical and theoretical analyses up to Section 4 are applicable convex Lipschitz losses, hence missing the popular squared loss. In Section 5, we take a step towards handling the squared loss. Employing an alternative reparameterization. we prove the same set of theorems as Section for ninllenneuron nonlinear network trained with a single data point under squared loss.
1.2 Related works
The Edge of Stability (EoS) phenomenon has been extensively studied in recent years, with many works seeking to provide a deeper understanding of the evolution of sharpness and the oscillating dynamics of GD. Jastrzębski al. (2019) and Jastrzebski al. (2020) observe that step size affects the sharpness along the optimization trajectory. Cohen et al. (2021) first formalize EoS through empirical study, and subsequent works have built on their findings. Ahn al. (2022) analyze EoS through experiments and identify the relations between the behavior of loss, iterates, and sharpness. Ma al. (2022) suggest that subquadratic growth of the loss landscape is the key factor of oscillating dynamics. Arora et al. (2022) show that (normalized) GD enters the EoS regime, by verifying the convergence to some limiting flow on the manifold of global minimizers. Wang et al. (2022) divide GD trajectory into four phases and explain progressive sharpening and EoS by using the norm of output layer weight as an indicator of sharpness. Lyu et al. (2022) prove that normalization layers encourage GD to reduce sharpness. Damian et al. (2023) use the third-order Taylor approximation of the loss to theoretically analyze EoS, assuming the existence of progressive sharpening. Lee and Jang (2023) propose new sharpness measure using batch gradient distribution and characterize EoS for SGD. Concurrent to our work, Wu et al. (2023) study the logistic regression problem with separable dataset and establish that GD exhibits implicit bias toward the max-marghn solution in the EoS regime, extending prior findings in the small step size regime (Soudry et al., 2018; Ji and Telgarsky, 2019b).
Some recent works rigorously analyze the full GD dynamics for some toy cases and prove that the
limiting sharpness is close to 2/1z. Zhu et al. (2023) study the loss (x, + (x²y²² - 1)² and prove that the sharpness converges close to 2/n with a local convergence guarantee. Notably, Ahn et al. (2023) study the function (I,, 3) {(ry) where l is convex, even, and Lipschitz, and provide global convergence guarantee. The authors prove that when is log-cosh loss or square root loss, the limiting sharpness in the EoS regime is between 2/11 O and 2/1, Our theoretical results extend their results on a single-neuron linear network to a two-layer fully-connecte linear network and provide an improved characterization on the limiting sharpness, tightening the gap between upper and lower bounds to only O1r²)).
The trajectory aligmment phenomenon is closely related to Zhu et al. (2023) which shows empirical evidence of bifurcation-like oscillation in deep neural networks trained on real-world data. However, their empirical results do not show the alignment property of GD trajectory. In comparison, we observe that GD trajectories align on the same bifurcation diagram, independent of initialization.
Very recently, Kreisler et al. (2023) observe a similar trajectory alignment phenomenon for scalar linear networks, employing reparameterization based on the sharpness of the gradient flow solution. However, their empirical findings on trajectory alignment are confined to scalar linear networks, and do not provide a theoretical explanation. In contrast, our work employs a novel canonical reparameterizattin and offers empirical evidence for the aligmment phenomenon across a wide range of networks. Moreover, we provide theoretical proofs for two-layer linear networks and single-neuron nonlinear networks.
2 Preliminaries
Notations, For vectors r and u, we denote the Ep norm of u by || their tensor product as u & v, and u @ u by u 02 For matrix A, we denote the spectral norm by ||]|a Given function L and a parameter O we use Amax ()) == \ cigenvalue of the loss Hessian) at e. We use asymptotic notations with subscripts (c.g.. Oc(-). (s)e()) in order to hide constants that depend on the parameters functions written as subscripts.
to denote the sharpness (i.e., the maximum
3

---

2.1 Problem settings
We study the optimization of neural network f(- O): R² R parameterized by θ. We focus on simple peeepparameterized setting trained on single data point {(r,))) where æ E R and y E IK. We consider the problem of minimizing the empirical risk
(()) = ()) y),
where é is convex, even, and twice-differrentibll with (0) 1. We minimize L using GD with step size v7: 0411 = θ - -)))((())) The gradient and the Hessian of the function are given by
Vel(O)= (( -a:00))-yVeeie v)
C(O( (f((:: () ( v))
Suppose that @* be global minimum of L, i.c.. f(æ; ) y. In this case, the loss Hessian and
the sharpness at @* are simply characterized as
Vo((O') (Vel (æ; 0²)) and
Amax ( = - (ei( (x;
2.2 Canonical reparameterization
Definition 2.1 (canonical reparamctcrizatioo) For given step size the canonical reparameterizatioon of ④ is defined as
(p.g)= = ((x;; θ) v/lV (z; θ)
2
Under the canonical reparameterization, p 0 represents global minima, and Eq. (1) implies that the point (p. q) (0, 1) is global minimum with sharpness 2/n. Note that (p. q) alone does not, in general, uniquely determine the value of @ Rather, the motivation for this reparameterization technique is effectively analyze the complex GD dynamics in the high-dimensional parameter space by reducing it to a 2-dimensional representation. The update of can be written as
Pe+1 = f(x; O - y = fxx; (f(x: θ,) Vof( Q.)) -y
~ f(æ; O2) Vef(#; 0.)
= (f(x: O2) v) n/ (f(æ:
- - V)
28 (pt)
(3)
q1
which can be obtained by first-order Taylor approximation on f for small step size nj.
2.3 Bifurcation analysis
Motivated from the approximated 1-step update rule given by Eq. (3), we conduct the bifurcation analysis on this one-dimensional map, considering as a control parameter. We first review some basic notions used in bifurcation theory (Strogatz, 1994).
Definition 2.2 (stability of fixed point). Let 2a be fixed point of differentiable map f R - R. i.e., f(z)) = 2g. We say Zg is a stable fixed point of f if ff'))) 1, and we say 2o is an unstable fixed point of f if (f))
1.
Definition 2.3 (stability of periodic orbit). A point zo is called period-p point of a map f. - R z0 is the fixed point of fr and f (zo) * 2o for any j S. p The orbit of 20, given by 0, .p 1} is called the period-p orbit of f. A period-p orhit is stable
{2j f 1
(unstable) if its elements are stable (unstable) fixed points of f", i.e., 113=b1 [z))) <1 1).
Now analyze the bifurcation of the one-parameter family of mappings fq IR R given by
[1(1) == ( 2r(p)
where q is a control parameter and r is differentiable function satisfying Assumption 2.4 below.
The approximation is used just to motivate Lemma 2.1: theorems, analyze the exact dynamics.

---

(a)
(b) 10
(c)
r(pe)
and vs. t
20.
18
A=
200
4.00
100 180
2-0
6.0
Figure 2: (a), (b) GD trajectories of two-layer fully-connected linear networks trained with different initialization scale a. Each color corresponds to a single run of GD. Smaller initialization scale falls into the gradient flow regime, whereas larger initialization falls into the EoS regime. (c) In the EoS regime, T(P.)
St_ approaches the early phase of training, whereas Pt converges relatively slowly.
Assumption 2.4. A function r RR + R is even, continuously differentiable, r(0) = L, (0) 0, (p) for any p > O, and p r(p) limp---o r(p) 0. In other words, is a smooth, symmetric bell- bell-shaped function with the maximum value r(0) = 1.
We note that Eq. (3) can be rewritten as Pc+1 = Sa (pt) ) we define r by r(p) ((pp) for p * 0 and
r(0) = 1. Below are examples of é for which the corresponding r's satisfy Assumption 2.4. These loss functions were previously studied by Ahn et al. (2023) explain EoS for (z, y) (rry)
log-cosh loss: ling-ench (p) = log(cosh(p)). Note Mog (p) = lanh(p).
square-root loss: Cayn (p) = + Note (p)
P VI+1²
If r satisfies Assumption 2.4, then for any 0 9 S 1, there exists a nonnegative number p such that r(p) g, and the solution is unique which we denote by f(q). In particular, f (0, - Rz0 is a function satisfying r(())) r(-f(g)) q for any q E 1).
Lemma 2.1 (period-doubling bifurcation of fa). Suppose that r is a function satisfying Assump- r(z)
tion 2.4. Letp" sup(p 01
-م'(x) > for any |x| S SP} and r(p"). If co, we choose
e = O. Then, the one-parameter family of mappings fg + R given by Eq. (4) satisfies
(i) lfg > 1, p. = 0 is the stable fixed point.
(i) 1f E (c, 1), 0 is the unstable fixed point and (++(g)) is the stable period-2 orbit.
Proof. The map fg has the unique fixed point p O for any q > 0. Since f2((0) -11 -1--)122 is stable fixed point if q > 1 and p 0 is an unstable fixed point if 0 < q < Now suppose that q E (c, 1). Then, we have fg(((g)) = -r(g) and f-()((4) = f(g), which implies that (+f(4)) is period-2 orbit of fg. Then, lff(((g))) = 1,g(- is a stable period-2 orbit.
(((g))
1 implies that (4)
According to Lemma 1. the stability of the fixed point 0 undergoes a change at g = 1, resulting in the emergence of stable period-2 orbit. The point (p, q) = (0. 1) is referred to as the bifurcation point, where a asymptotically approached by a system as function of a control parameter. In the case of the map fn. the corresponding bifurcation diagram is represented by p 0 for y Z and p +f(g) (or equivalently, g = r(p)) for q E (c, 1).
bifurcation occurs. A bifurcation diagram illustrates the points
It is worth noting that the period-2 orbit (+)(p)} becomes unstable for q E (0, c). If we choose r to be r(p) f'(p) for p * and r(0) 1, then + pr'(p) ("(p)
r(p) r(p) > 0 for all p, assuming f is convex. Consequently, for log-cosh loss and square root loss we have e = 0, indicating that the period-2 orbit
of f, remains stable for all (0, 1). However, in Section 5, we will consider with > 0, which may lead to additional bifurcations.
5

---

(a) m 64, L = 3
(b) 256, L 3
(C) m 256, L 10
04
04
Figure 3: GD trajectories of tanh-activated neural networks with varying width and depth. Each
color corresponds to single run of GD. We observe that the wider network (m = 256) exhibits stronger trajectory alignment phenomenon compared to the narrower network (m = 64). Figure 3c depicts the trajectories for deeper network (L 10), which also shows the trajectory alignment phenomenon.
3 Trajectory Alignment of GD: An Empirical Study
In this section, we conduct experimental studies on the trajectory alignment phenomenon in GD) dynamics under the canonical reparameterization proposed in Section 2.
We consider a fully-connected L-layer neural network f(:; R + R written as
f(x; )) w PWWL-10(- o( 2 W xx) )),
where @ is an activation function, W1 Radd W e mm for S < S and wt R All L layers have the same width of m. We minimize the empirical risk L(e) = =[f(x;;) y). We visualize GD trajectories under the canonical parameterization, where each plot shows five different randomly initialized weights using Xavier initialization multiplied with a rescaling factor of a. For this analysis, we fix the training data point and hyperparameters as x E1 = (1, 0, o........ 0). (linear) or
y = 1. = 0.01, d 10, and focus the log-cosh loss for &, with either Φ(f) ((t) = tanh(t). We note that the trajectory alignment phenomenon consistently observed in other settings, including square root loss, different activations (e.g.. ELU), and various hyperparameters. in particular for sufficiently wide networks (additional experimental results are provided in Appendix A).
The effect of initialization scale. In Figures 2a and 2b, we examine the effect of the initialization scale on GD trajectories in two-layer fulny-connected linear network with width of Figure 2a, when the weights are initialized with a smaller scale (a = 5), the initial value of g is greater than 1, and converges towards the minimum with only a small change in ge until convergence. In this case, the limiting sharpness is relatively smaller than 2/1), and the FoS phenomenon does not occur. This
256. In
is referred to the gradient flow regime (Ahn et al., 2023). On the other hand, in Figure 2b, when the weights are initialized with a larger scale (a = 10), the initial value of is less than 1, and we observe convergence towards the point (close to) (p. g) (0, 1). This case is referred to as the EoS regime. We note that choosing larger-than trnger-than-ttandard standard scale a is not a necessity for observing EoS; we note that even with a 1, we observe the EoS regime when 1) is larger.
Trajectory alignment on the bifurcation diagram. In order to investigate the trajectory alignment phenomenon on the bifurcation diagram, we plot the bifurcation diagram 9 = r(p) = ('(x) and observe that GD trajectories tend to align with this curve, which depends solely on f. Figure 2b clearly demonstrates this alignment phenomenon. Additionally, we analyze the evolution of dt r(p.)
and Pt in Figure 2c. We observe that the evolution of _ follows two phases. In Phase I, ?t 7(p.J
approaches to 1 quickly. In Phase II, the ratio remains close 1. Notably, the convergence speed of d. towards is much faster than the convergence speed of ps towards O. In Sections 4 and 5, r(p.) we will provide a rigorous analysis of this behavior, focusing on the separation between Phase and Phase II.
The effect of width and depth. In Figure 3, we present the GD trajectories of laul-activated networks with different widths and depths (a = 5). All three cases belong to the FoS regime, where GD converges to a point close (p. q) = (0, 1), resulting in limiting sharpness near 2/1j. However, when comparing Figures 3a and 3b, we observe that the trajectory alignment phenomenon is not

---

(a)
P(z) EE
(b) P(z) |
(c) Loss and t
24
V0a
34
Figure 4: Training subset of CIFAR-10. The GD trajectories trained 50 example subset of CIFAR-10 under the generalized canonical reparameterization (5) with different choices P and network architecture. Each color corresponds to single run of GD. Top row: fully-connected neural network with tanh activation. Bottom row: CNN with tanh activation. Full implementation details and further experimental results given in Appendix A.4.
observed for the narrower network with m = 64, whereas the GD trajectories for the wider network with m 256 are clearly aligned on the bifurcation diagram. This suggests that network width plays role in the trajectory alignment phenomenon, which is reasonable since wide networks are well approximated by their linearized models, hence Eq. (3) is more accurate. Furthermore, we note that the trajectory alignment phenomenon is also observed for deeper network with L = 10, as depicted in Figure 3c
Multiple training data points. In our trajectory alignment analysis, have primarily focused on training with a single data point. However, is important to explore the extension of this phenomenon scenarios with multiple training data points.
To investigate this, we train a neural network on a dataset {x,; y4) H=) where x; E IRd and y E R, by minimizing the empirical risk (()) = + En {(f(x, θ) - y:). Defining X RRadd as the data matrix and y IR" as the label vector, introduce generalized canonical reparameterization:
2n
(p. 9) = (((((x:0) -))
Vof(a, +) 0))6211
(5)
where P. R + IR can be function such as mean value or specific vector norm.
In Figure 4, consider training 50 example subset of CIFAR-10 with only 2 classes and vary the network architecture. We use three-layer fully-comnected network with tanh activation and convo- lutional network (CNN) with tanh activation. Under the generalized canonical reparameterization (5) for various choices of P, including the mean and the & norm, we observe the trajectory alignment phenomenon throughout all settings, indicating common alignment property of the GD trajectories. However, unlike the single data point case, the alignment does not happen on the curve (()) The precise characterization of the coinciding curve is an interesting direction for future research.
Trajectory Alignment of GD: A Theoretical Study
In this section, we study two-layer fully-connected linear network defined as f(x;; +) := v' Un. where ddxmm v E R and O denote the collection of all parameters (U,v), We consider training this network with a single data point {(x, 0)}, where æ E Rdd and | 1. We run GD with step size on the empirical risk
L(O) = l(f(x;; )) 0) l(v Tu))
7

---

where f is loss function satisfying Assumption 4.1. We note that our assumptions on is motivated from the single-neuron linear network analysis (d m 1) by Ahn al. (2023).
Assumption 4.1. The loss & is convex, even, 1-Lipschitz, and twice differentiable with E"(0)
The canonical reparameterization (Definition 2.1) + (U, v) is given by
(p. g) = Un, U 2 00
Under the canonical reparameterization, the 1-step update rule of GD can be written as
2r(p:) + t pir()r(())
pc, a++1 [1 [1-
r(po)) gt. (6)
Pc-1
Qt
where we define the function by r(p) (((p] for p * and r(0) = 1. Note that the sequence (0x) is monotonically increasing if go 2 which is the case our analysis will focus on.
We have an additional assumption on f below, motivated from Lemma 2.1.
Assumption 4.2. The function r(p) p
(p) corresponding the loss l satisfies Assumption 2.4.
We now present our theoretical results on this setting, and defer the proofs to Appendix B.
4.1 Gradient flow regime
We first consider the gradient flow regime, where is initialized with >>
Theorem 4.1 (gradient flow regime), Let 1 E (0, 2 be a fixed step size and be a loss function 33 satisfying Assumptions 4.1 and 4.2. Suppose that the initialization (Pas %)) satisfies lPol 1 and min min iterations (Pi, 4:) converge to the point (0,9*) 1²) such that
r(1) }). Consider the GD trajectory characterized in Eq. (6). Then, the GD
40 1
Qo q- S exp(C1i")) )go S 2go.
where C' 8qo
min
2(go
(()
90
2go
Theorem 4.1 implies that in gradient flow regime, GD with initialization O, (Uo, vo) and step size 17 converges to O which has the sharpness bounded by:
(1 - CrCrr
+ (vo112) S Amax (0)) I (1) + |
Hence, for small step size v,, if the initialization satisfies || ||3 |000 2 1, then the limiting sharpness is slightly below ||Unæ|] - +1 Note that assumed the bound |pol < for simplicity, our proof also works with the assumption |Pol S for any positive constant modulo some changes in numerical constants. Moreover, our assumption on the upper bound of is 1/1) up to constant factor, which covers most realistic choices of initialization.
4.2
EoS regime
We now provide rigorous results in the EoS regime, where the GD trajectory aligns on the bifurcation diagram r(y). To establish these results, we introduce additional assumptions on the loss t.
Assumption 4.3. The function r(2) (x) is CA on R. and satisfies
(i) 2 + r[a]²
(()) is decreasing on IR.
(ii) 2 + (x)
is decreasing on z and increasing on z < 0,
r(2)
(iii) 2 )
ar(s)
is decreasing on 2 > and increasing on 2 < 0.
(())
We note that both the log-cosh loss and the square root loss satisfy Assumptions 4.1, 4.2, and 4.3.

---

Theorem 4.2 (EoS regime, Phase 1). Let 1) be small enough step size and é be a loss function satisfying Assumptions 4.1, 4.2, and 4.3. Let 2g s (r) Let 8 E (0, cn) be any given constant. Suppose that the initialization (po, qo) satisfes lpol S
r(2] - 2 and cu max(r(z)), 2)
and go (co. 1 8). Consider the reparameterized GD trajectory characterized in Eq. (6). We
assume that for all & 2 0 such that q: < L, we have # O. Then, there exists a time step to (log( ))) such that for any to.
1 + M(DE)1 Os,e(1));
r(p:)
pr(p) (p) + p²r((p)) for + 0 and h(p) = 20 for p = 0.
where h((p =
One can check that for log-cosh and square-root losses, the ranges of h are (0, 3/4] and (0, 1/2].
respectively. Theorem 4.2 implies that in the early phase of training (1 S la = Olog(n -1))) ()), GD iterates (pr, Qt) approach closely to the bifurcation diagram r(p) 9. which
called Phase in
Section 3. In Phase II, GD trajectory aligns on this curve in the remaining of the training (t 2 ta). Theorem 4.3 provides an analysis on Phase II stated below.
settings in Theorem 4.2, there exists time step ta O2(1 - -0)' such that qt, S 1 and > 1 for any t > tb. Moreover, the GD iterates (pt, (:) converge the point (0, q²) such that
Theorem 4.3 (EoS regime, Phase II). Under the
o = 1
+ +o,e(11*).
2r"
Theorem 4.3 implies that in EoS regime, GD with step size v converges @" with sharpness
2
Amax (0)
(0)) Ose(11²).
Note that Ahn et al. (2023) study the special case d = m = and prove that the limiting sharpness is between 2/n- O()) and 2/1). Theorem 4.3 provides tighter analysis the limiting sharpness in more general settings, reducing the gap between the upper bound and lower bound to only O Also, our result the first prove that the limiting sharpness in the EoS regime bounded away from 2/n by a nontrivial margin.
We also study the evolution of sharpness along the GD trajectory and prove that progressive sharpening (i.e., sharpness increases) occurs during Phase II.
Theorem 4.4 (progressive sharpening). Under the same setting as in Theorem 4.2, let ta denote the obtained iteration. Define the function À IR50 + IR given by
1 (labr(((a))) if and
A())
otherwise.
Then, the sequence (A(q)) monotonically increasing. Moreover, for any tu the sharpness at GD iterate O closely follows the sequence (A(4:) r by satisfying
Amax(()) - - < 11+0e(m)
The gap between Amax (O.) and À(g:) is bounded by numerical constant, which becomes negligible compared to 2/17 for small 7.. In Figure la, we perform numerical experiments on a single-neum neuron case and observe that A(g.) closely approximates the sharpness.
Note that Cohen et al. (2021) observe increase in sharpness during training in the gradient flow regime, while our work reveals that sharpness increases when exhibiting oscillations the EoS regime. This distinction may be linked to the selection of the loss function, our study focuses on Lipschitz convex losses, while Cohen et al. (2021) examine the squared loss.

---

EoS in Squared Loss: Single-neuron Nonlinear Network
Our canonical reparameterization has a limitation in explaining the EoS phenomenon under squared loss E(p) 22², as the function r(p) = <'(p) does not satisfy Assumption 2.4. However, empirical studies by Cohen et al. (2021) have observed the EoS phenomenon in GD training with squared loss. In this section, we analyze simple toy model to gain insight into the EoS phenomenon and trajectory alignment of GD under squared loss.
We study the GD dynamics on wo-diimensional function L(x, y) = ( where æ, y are scalars and 6 is a nonlinear activation satisfying Assumption 5.1 below.
Assumption 5.1 (sigmoidal activation). The activation function R + R is odd, increasing, 1- Lipschitz and twice continuously differentiable. Moreover, Φ(0) ø' (0) 1, 1 ((x) and -im (xx) -.
One good example of o satisfying Assumption 5.1 is tanh. For this section, we use an alternative reparameterization defined as below.
Definition 5.2. For given step size v, the (p. 9) reparameterizattion of (s, u) E R2 is defined as
2
(p.9) = I,
Under the reparameterizatton, the 1-step update rule be written
Po+1 = (1- 2r(pi)
Pt. G++1 (1 1 (1=))(p))) 29..
(7)
where the function r is given by r(z) ((=) for + and r(0) =
We can observe a notable resemblance between Eq. (7) and Eq. (6). Indeed, our theoretical findings for single-neuron nonlinear network closely mirror those of the two-layer linear network discussed in Section 4. Due to lack of space, we summarize our theorems in this setting as the following:
Theorem 5.1 (informal). Under suitable assumptions on é. step size, and initialization, GD trained on the squared loss L(x.y) == 2 (0(x)y)² exhibits the same gradient flow, EoS (Phase II), and progressive sharpening phenomena as shown in Section 4.
In Theorem C.3, we prove that in the EoS regime, the limiting sharpness is 2 formal statements of the theorems and the proofs, we refer the reader to Appendix C.
+ +O()). For
7 P
6 Conclusion
In this paper, provide empirical evidence and rigorous analysis demonstrate the interesting phenomenon of GD trajectory alignment the EoS regime. Importantly, we show that different GD trajectories, under the canonical reparameterization align bifurcation diagram independent of initialization. This discovery is notable due to the intricate and non-convex nature of neural network optimization, where the algorithm trajectory is heavily influenced by initialization choices. Our theoretical analysis not only characterizes the behavior of limiting sharpness but also establishes pro- gressive sharpening of GI). One immediate future direction is to understand the trajectory alignment behavior when trained on multiple data points. Lastly, will be interesting to extend our analysis to encompass squared loss for general neural network, going beyond the toy single-neuron example.
Acknowledgments
This paper was supported by Institute of Information & communications Technology Planning Evaluation (ITTP) grant (No. 2019-0-00075, Artificial Intelligence Graduate School Program (KAIST)) funded by the Korea government (MSIT), two National Research Foundation of Korea (NRF) grants (No. NRF-2019R 15412833244,, RS-2023-0021 1352) funded by the Korea government (MSIT), and a grant funded by Samsung Electronics Co., Ltd.
10

---

References
Kwangjun Ahn, Jingzhao Zhang, and Suvrit Sra. Understanding the unstable convergence of gradient descent. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning. volume 162 of Proceedings of Machine Learning Research, pages 247-257. PMLR, 17-23 Jul 2022. URL https::/p///. /ttapa:/paae /ptp::::///cceennggss mmr..nng..........mmmmm.m.m..m.m prwppr /atess//v1//aaa 2222222aaaa....a//aa///a/222aaam html.
Kwangjun Ahn, Sébastien Bubeck, Sinho Chewi, Yin Tat Lee, Felipe Suarez, and Yi Zhang. Learning threshold neurons via the "edge of stability". NeurlPS 2023 (arXivv22122.07469) 2023.
Sanjeev Arora, Nadav Cohen, and Elad Hazan, On the optimization of deep networks: Implicit acceleration by overparameterizztion.. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 Proceedings of Machine Learning Research, pages 244-253, PMLR, 10-15 Jul 2018. URL press/v80/ https//v0//rrrra1e0 arora18a. 1ttmm.l..
mll.
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox. and R. Gar- nett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Asso- ciates, Inc., 2019. URL pttttps://proceeeinnge #111//0c78b550orr001111166660602
weurips.cc/paper files/paperr2019/ .df.
Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on the edge of stability in deep learning. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepes Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 948-1024. PMLR, html.
17-23 Jul 2022. URL https: httpp:s//ppocceeeeeeee prwceeding......m mlr. wptpss/vve/aaroeeeeeeeeeeeeee
Lei Chen and Joan Bruna. Beyond the edge of stability via two-step gradient updates. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 4330-4391. PMLR, 23-29 Jul 2023. URL httpp:// http://ppooeeeiin..iliii/ mlr http:iiiwww.l....ee/eeeeeeaeeeee httpss//v22//chhaa .tmml.
Lénaic Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Jacob Abernethy and Shivani Agarwal, editors, Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pages 1305-1338, PMLR, 09-12 Jul 2020. URL https: //rppe/d///cceeingsss mlr. alrrpress/// v125/ 2225/chizat20.. html.
Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability, In International Conference on Learning Representations, 2021. URL /ttps:///penreee hhtp://wwttetooooirtrirarooooooe jh-TwkGe..................
Alex Damian, Eshaan Nichani, and Jason D. Lee. Self-stabilization: The implicit bias of gra- dient descent the edge of stability. The Eleventh International Conference on Learning Representations, 2023, URL https:///ppeneev//
Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro. Implicit regularization in matrix factorization. In Guyon, U. Von Luxburg. S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https: http://///occccepeeeeeeeeaaee neuriinnisniinnnnnooeoe.eeeriiie neurips. ec/ Paper pdf.
Suriya Gunasekar, Jason D Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient de- scent on linear convolutional networks. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URL httptttppe////ooceeeenngeeg neurips. ccpppprrr..pp. Pafer.ddf.

---

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. In International Conference on Learning Representations, 2020. URL https:// www.:/ennnettttommdmrrmrri.
Stanistaw Jastrzębski. Zachary Kenton, Nicolas Ballas, Asja Fischer. Yoshua Bengio, and Amost Storkey. On the relation between the sharpest directions of DNN loss and the SGD step length. In International Conference on Learning Representations, 2019. URL https / /opp//eeppeeeeeeeeeeeeeeeeeeeeee net/ Forum?id=SkgEaj0555
Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. In International Conference on Learning Representations, 2019a. URL https: /owpr//veeee.......//mo/ het/fforum?.ccHHi]]gggoo...
Ziwei Ji and Matus Telgarsky. The implicit bias of gradient descent on nonseparable data. In Alina Beygelzimer and Daniel Hsu, editors, Proceedings of the Shcrty-Second Conference on Learning Theory, volume 99 of Proceedings of Machine Learning Research, pages 1772-1798. PMLR, 25-28 Jun 2019b. URL https::////// http://///occeeee..//oooa.o www.ediinnnmmm......mmm wtw....p..e.///////em ji19a. html.
Itai Kreisler, Mor Shpigel Nacson, Daniel Soudry, and Yair Carmon. Gradient descent monotonically decreases the sharpness of gradient flow solutions in scalar networks and beyond. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 17684-17744. PMLR, 23-29 Jul 2023. URL htttt::://ppoceedinggs mlr www.......................
html.
Sungyoon Lee and Cheongjae Jang. new characterization of the edge of stability based on sharp- ness measure aware of batch gradient distribution. In The Eleventh International Conference on Learning Representations, 2023. URL /ttpps://penret
hett/formmidd-i-- -xCKLLK.K
Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of nor- malization layers: Sharpness reduction. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors. Advances in Neural Information Processing Systems, 2022, URL attps://openrevv// wwwpeerreieeeeeeeeorrreeeeeeeeee
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.11 2.XY::2204.11326, 2022.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. H. Wallach, H. Larochelle, Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran As- sociates, Inc., 2019. URL UtLaattps//pproceedings hwwp://roooorriiin....o.ooo...p. neurips wtt..cccpppppp htcc/pppr_ffeee/peeeeeeeeee sttps//s///ssap//219999 .ddf.
file F111:/acccaBBteff112222000002222
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent separable data. Journal of Machine Learning Research, 19(70):1-57, 2018. URL http: /jmrr./m//////// waw.l..ogg/appeee
org/papers/vv1/118. watpeeer////////-111/ee11 html.
Steven H. Strogatz. Nonlinear dynamics and Chaos: with applications to physics, biology. chemistry, and engineering. Studies in nonlinearity. Addison-Wesley Pub, Reading, Mass, 1994. ISBN 9780201543445.
Gal Vardi. On the implicit bias in deep-learning algorithms. arXiv preprint 2rXhv:2208.12551, 2022.
Zixuan Wang, Zhouzi Li, and Jian Li. Analyzing sharpness along GD trajectory: Progressive sharpening and edge of stability. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URL https::/openreeee wwt..n/ooooooiiiiirggtttaaaaaaa
12

---

Blake Woodworth, Suriya Gunasekar, Jason D. Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes overparametrized models. In Jacob Abernethy and Shivani Agarwal, editors, Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pages 3635-3673. PMIR, 09-12 Jul 2020. URL htwph:tt/p/pppt hatp:s////occeedngsss mlr. www.mllrpprr... press/ press/v122 225///oooorrh20000 hht.m.
Jingfeng Wu, Vladimir Braverman, and Jason D Lee. Implicit bias of gradient descent for logistic regression at the edge of stability. arXiv preprint arXiv:2305.11788 2023.
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), pages 581-590, Dec 2020. doi: 10. (090/BigDDaaa5002222209837888889
Chulhee Yun, Shankar Krishnan, and Hossein Mobahi. A unifying view on implicit bias in training linear neural networks. In International Conference on Learning Representations, 2021, URL https://openrent net/ ftr/m/ormmm// -id.iiiddd--------- 41MQKH.
Xingyu Zhu, Zixuan Wang, Xiang Wang, Mo Zhou, and Rong Ge. Understanding edge-of-stability training dynamics with minimalist example. In The Eleventh International Conference on Learning Representations, 2023. URL https: htttp:://openrevieee net nett/forum'dde-p/Faaaaaaaa.
13

---

Contents
1
Introduction
1.1
Our contributions
1.2
Related works
2
Preliminaries
2.1
Problem settings
2.2
Canonical reparameterriztiion
2.3
Bifurcation analysis
3
Trajectory Alignment of GD: An Empirical Study
Trajectory Alignment of GD: A Theoretical Study
4.1 Gradient flow regime
4.2 EoS regime
8
5 EoS in Squared Loss: Single-neuron Nonlinear Network
10
6
Conclusion
10
Additional Experiments
15
A.1 Experimental setup
15
A.2 Training single data point
16
A.3 Training on multiple synthetic data Points
................................
17
A.4 Training on real-world dataset
21
Proofs for the Twy-Layer Pulne-Connected Linear Network
24
B.1 Proof of Theorem 4.1
24
B.2 Proof of Theorem 4.2
25
B.3 Proof of Theorem 4.3
33
B.4 Proof of Theorem 4.4
C Proofs for the Single-neuron Nonlinear Network
38
C.1 Formal statements of Theorem 5.1
................................
38
C.2 Proof of Theorem C.I
39
C.3 Proof of Theorem C.2
41
C.4 Proof of Theorem C.3
48
C.5 Proof of Theorem C.4
49
14

---

A
Additional Experiments
In this section, we present additional empirical evidence demonstrating the phenomenon of trajectory alignment, which supports the findings discussed in Section 3 of our main paper.
A.1 Experimental setup
Objective function.
We run gradient descent (GD) minimize the objective function defined as (()) (({(;;; θ v),
where represents the parameters, f: R - a loss function, f R. IR is a neural network, and (ax, y denotes a single data point with a E R and E R. We also consider training on multiple data points (()),, W:)) the objective function
with 0; E Rd and yj E IR for each S < YI, where we minimize
( = (f(x;;;) -y:).
In our experiments, we primarily focus on the log-cosh loss function llog-cosh (p) = log(cosh(p)) but we also investigate the square root loss (p) / /1+p2.
Model architecture. We train fully-connected L-layer neural network, denoted as f(::0) Rd R. The network defined as follows:
fax θ) = w (W (e (2 ..))))))))))..
where é R > R is an activation function applied entry-wise, W1 E RRnd W; E Rmxm for 2 < I L 1, and wy E R All L layers have the same width of m, and the biases of all layers are fixed to 0.* We consider three activations: hyperbolic tangent 6(t) = tanh(t). exponential linear unit (f) [t)( and linear <tt) =
Weight initialization. We perform gradient descent (GD) using five different randomly initialized sets of weights. The weights are initialized using Xavier initialization, and each layer is multiplied by
a rescaling factor (gain) of a. the plots presented throughout this section, we mark the initialization points with an 'x" to distinguish them from other points the trajectories.
Canonical reparameterization. We plot GD trajectories after applying the canonical reparameteri- zation introduced in Definition 2.1:
2
(p. q) f)x n v/Vof( θ
where 17 denotes the step size. For training on multiple data points, we employ the generalized canonical reparameterization as defined in Eq. (5):
(p. 9) = (0 (P( - 7 1 (oof(z, θ) 1))
2n
where P: _R + E can represent the mean value or vector norms. Specifically, we mainly focus on the mean value P(z) Σ 2, but we also examine vector norms such as P(z) = ||2]..
P(z) = =| and P(z) I where = (21, 22, 2n) E R..
For large networks, explicitly calculating the lg matrix norm | (V 0))321 is infeasible. Therefore, adopt fast and efficient matrix-free method based power iteration, proposed by Yao et al. (2020). This method allows us to numerically compute the l norm of large-scale symmetric matrices.
'While we fix the biases of all layers to 0 to maintain consistency with our theory, the trajectory alignment phenomenon is consistently observed even for neural networks with bias. In Appendix A.4, we consider training networks with bias.
15

---

(a) 0.5
(b) 1,0
(c) 2.0
(d) - 4.0
12
3.-
35
60-
54
24
60.5
20 pa.
133
83
-08 -04
-2.25 -275
0.02
Figure 5: GD trajectories (1000 iterations) of 3-layer fully connected neural networks with tanh activation and width of m 256. The plots depict the trajectories for different initialization scales: a = 0.5, 1.0, 2.0, 4.0. Smaller initialization scales (a = 0.5, 1.0) lead to trajectories in the gradient flow regime, while larger initialization scales (a 2.0, 4.0) result in trajectories in the EoS regime.
A.2 Training on a single data point
Training data. We conduct experiments on a synthetic single data point (æ, y). where æ = e1 (1.0.....00 0) E Rd and y E IR. Throughout the experiments in this subsection, we keep the data dimension fixed at d 10, the data label at y 1, and the step size 7) 0.01.
The effect of initialization scale. We investigate the impact of the initialization scale a while keeping other hyperparameters fixed. We vary the initialization scale across (0.5, 1.0, 2.0, 4.0). Specifically, in Figure 5, we train a 3-layer fully connected neural network with tanh activation. We observe that smaller initialization scales (a 0.5, 1.0) lead trajectories in the gradient flow regime, while larger initialization scales (a = 2.0, 4.0) result in trajectories in the EoS regime. This behavior primarily due the initial value of g being smaller than for larger initialization scales, which causes the trajectory to fall into the EoS regime.
The effect of network width. We investigate the impact of network width on the trajectory aligmment phenomenon. While keeping other hyperparameters fixed, we control the width m with values of (64, 128, 256, 512]. In Figure 6, we train 3-layer fully connected neural networks with tanh activation and an initialization scale of a = 4. Additionally, in Figure 7, we examine the same setting but with different depth, training 10-layer fully connected neural networks.
is commonly observed that the alignment trend becomes more pronounced as the network width increases. In Figures and 7, all trajectories are in the EoS regime. However, narrower networks (m = 64, 128) do not exhibit the trajectory alignment phenomenon, while wider networks (m = 512) clearly demonstrate this behavior. These results indicate that network width plays significant role in the trajectory alignment property of GD.
The effect of loss and activation functions. The trajectory alignment phenomenon is consistently observed in various settings, including those with square root loss and different activation functions. In Figure 8, we investigate 3-layer fully connected neural network with width of 256 and an initialization scale of a = 2. We explore different activation functions, including tanh, ELU, and linear, and consider both log-cosh loss and square-root loss. Across all these settings, we observe the
trajectory alignment phenomenon, where the GD trajectories align on curve ('(
16

---

(a) = 64
(b) m = 128
(c) m 256
(d) m = 512
10
3.5
0-
1.0
34
0.8
005
0.4
0.4
0.1
0.4
07
0.2
-10
-25 3.0
73
Figure 6: GD trajectories (1000 iterations) of 3-layer fully connected neural networks with tanh activation and an initialization scale of 4. The plots depict the trajectories for different network widths: m = 64, 128, 256, 512. Narrower networks (m = 64, 128) do not exhibit trajectory alignment, while wider networks (m 256, 512) clearly demonstrate the trajectory alignment phenomenon.
(a) 64
(b) 128
(c) m77 255
(d)
512
1.4
12-
pw0
3.0
pw0
0.3
pw0
3-5
2-5
5
0.6
0.4
0.2
0.2
0.0
0.0 -5
5.
2.0
-10 -5
10
-10
-10
-5
10 -20
-10
10
20
Figure 7: GD trajectories (1000 iterations) of 10-layer fully connected neural networks with tanh activation and an initialization scale of a = 4. The plots depict the trajectories for different network widths: 64, 128, 256, 512. Similar to Figure 6, narrower networks (m 64, 128) do not exhibit trajectory alignment, while wider networks (m = 256, 512) clearly demonstrate the trajectory alignment phenomenon.
A.3 Training on multiple synthetic data Points
Training data. We consider training on synthetic dataset consisting of n data points, denoted as (xx, 9:) The input vectors I; are sampled from a standard Gaussian distribution NO, I). where an E RRd, and the corresponding target values yi are sampled from a Gaussian distribution (0, 1), where E R. Throughout our experiments in this subsection, we use fixed data dimension of = 10 and a step size of 17- = 0.01.
The effect of function P. investigate the impact of different choices of the function P, we train 3-layer fully connected neural network with tanh activation. The network has width of m 256 and is initialized with a scale of a = 4. The training is performed on dataset consisting of n = 10 data points. Figure 9 displays the trajectories of GD trajectories under the generalized canonical reparameterization defined Eq. (5) for various choices the function P. These choices include the mean, & norm, &g norm, and loo norm.
We observe that GD trajectories exhibit alignment behavior across different choices of the function P. Notably, when P is selected as the mean, the trajectories align on the curve q = ('()) However, when p
P is based on vector norms, the alignment occurs on different curves. The precise characterization of these curves remains as an interesting open question for further exploration.
The effect of the number of data points. We examine how the size of the training dataset, denoted by n. influences the trajectory alignment behavior of GD. While keeping other hyperparameters constant, we vary n with values (2, 4, 8, 16, 32, 64, 128, 512, 1024). In Figure 10, we train a 3-layer fully connected neural network with tanh activation, width of m 256, and an initialization scale
17

---

(a) log-cosh, tanh
(b) oog-cosh, ELU
(c) log- cosh, linear
1.0
1/2
699
00-
34-
6.
04
05
2.4
03
222
(d) square-root, tanh
square-root, ELU
(1) square-root, linear
1.0
1.
0.9
09
48
08
0.8
0.
60.
0.6
0.4
0.5
P
03
Figure 8: GD trajectories (1000 iterations) of 3-layer fully connected neural networks with different loss and activation functions. The networks have width of m 256 and an initialization scale of = 2. The plots illustrate the trajectories for various combinations activation functions (tanh, ELU, and linear) and loss functions (log-cosh and square-root loss). In all these settings, the trajectory alignment phenomenon observed, where GD trajectories align on a curve trajectories with log-cosh loss. From left to right: tanh, ELU, and linear activations. Bottom row: trajectories with square-root loss. From left right: tanh, ELU, and linear activations.
((p)) Top row:
P
(a) P(z)
EE 25
(b) P(=) = ||=|||
(c) P(2)
|||||2
(d) P(z)
=||=|]
00
0.8
0.5
p
53
p
35
3.6
0.4
2.4
0.4
02
02
60
15
12
Figure 9: GD trajectories (500 iterations) under the generalized canonical reparamcterizztiion (Eq. (5)) for different choices of the function P. The plots depict the training of 3-layer fully connected neural network with tanh activation, width of 256, and initialization scale of a 4. Each plot corresponds to a different parameterization, including the mean, l norm, & norm, and foo norm. GD trajectories commonly exhibit alignment behavior, with the mean parameterization aligning on the curve = e'(p)
o = 4. Additionally, in Figure 11, we investigate the same setting but with different activation function, training ELU-activated fully connected neural networks. The GD trajectories are plotted under the generalized canonical reparameterization using the mean function P(z) = Σm= 24-
We observe consistent trajectory alignment phenomenon across different choices of the number of data points. Interestingly, for small values of n, the trajectories clearly align on the curve & = (p) However, as the number of data points n increases, seems that the trajectories no longer align on this curve but different "narrower" curves. Understanding the underlying reasons for this phenomenon poses an intriguing open question.
18

---

(a) 2
(b) 4
(C)
1.
p
pw0
0.6
24
9-
5.8
35
0.4
0.3
53
0.2
-5.0
-2.5
255
5.0
(d) 16
12
32
(1) 12
1.0
L0-
10-
34
077
57
24
0.2
(g) 128
(h) 256
(i) 512
1.-
6.9
0
p
p
5.8
0.0
2.
%.
0.5
05-
35
2.4
03
52
(j) 73 1024
L.-
03
Figure 10: GD trajectories (500 iterations) under the generalized canonical reparameterization (Eq. (5)) for different choices of the number of data points n. The plots depict the training of 3-layer fully connected neural network with tanh activation, width of 256, and an initialization scale of a = 4. The function P is chosen to be the mean P(z) E The trajectories exhibit
alignment behavior, where the curves followed by the trajectories change depending on the value of
n. For small values of n, the trajectories align on the curve 9 they align distinct curve.
(p) while for large values of n,
19

---

(a) 2
(b) 4
(C)
3.5
3.4
20
p=0
3.2
24
3.0
24
0.0
24
0.4
0.2
-70
-10
10
-20
-10
10
-20
-10
10
(d) 16
(c) 72
32
(f) n = 64
68
2
044
0.
52
-10
-10
10
-10
-5
10
-13
-10
-5
15
(g) 128
(h) 256
(i) 512
P=
p=
20
P.
25
0.2
-10
-5
10
-10
10
-10
-5
10
(i) 7
1024
20
10 20
Figure 11: GD trajectories (500 iterations) under the generalized canonical reparameterization (Eq. (5)) for different choices of the number of data points n. The plots depict the training of 3-layer fully connected neural network with ELU activation, width of 256, and an initialization scale of a = 4. The function P is chosen be the mean P(z) > The trajectories exhibit alignment behavior, where the curves followed by the trajectories change depending on the value of
n. For small values of n, the trajectories align on the curve 9 they align distinct curve.
(p) while for large values of n,
20

---

A.4 Training on real-world dataset
Training data. this subsection, we investigate a binary classification problem using subset of the CIFAR-10 10 image classification dataset. Our dataset consists of 50 samples, with 25 samples from class 0 (airplane) and 25 samples from class (automobile). We assign label of +1 to samples from class O and a label of 11 to samples from class 1. This dataset was used in the experimental setup by Zhu al. (2023).
Architectures. In Figure 4, we examine the training of two types of network architectures: (top row) fully- fully-connected tanh network and (bottom row) convolutional tanh network. The PyTorch fully- fully-connected tanh network is provided as follows:
code for the
nn Sequential(
nn. Flattet cestar__dimsst end_ -id_ddm=-
nn. Ln.Line (3072, 500, bias=True),
nn. Tn.Tanh() (),
nn In..inearr500 500, bias=True), nn. mn.Tan(() (500, 1, bias-True)
Similarly, the PyTorch code for the convolutional tanh network is as follows:
Sequential(
nn. Fla.taa (sttrtt_ssttttddmm11,,. end_ add_dim=-1),
nn. nn. Unflatten (dim=1, nn. (3, 500, kernel_size=(3, 3), stride=(1, 1), padding= 1), bias=True), Tanhanh()
seteetteeed_iieeeieeeeee (3, 32, 32)),
in. mm.Conn2dd500, 500, kernel_size=(3, 3), stride=(1, 1), padding= 1), bias=True),
nn.Tanh(),
nn. in.Faateen
end_dim=-1), L..near 1, bias=True)
)
Note that consider networks with bias. We use a step size of 17 0.01 for the fully-connected network and 1) 0.001 for the CNN. The default PyTorch initialization (Paszke et al., 2019) is applied to all these networks.
In this subsection, we further explore the (reparameterized) GD trajectories of fully-connected networks with different activation functions, network widths, and the choice of function P in (5).
The effect of function P. We investigate the impact of different choices of the function on the GD trajectories. We train 3-layer fully connected neural network with ELU activation, width of m = 256, and an initialization scale of a = 1. Figure 12 illustrates the GD trajectories under the generalized canonical reparameterization defined in Eq. (5) for various choices of the function P. including the mean, & norm, Ea norm, and fo norm.
We observe that the GD trajectories exhibit alignment behavior, which is more pronounced when P is chosen to be the mean or l norm, but less evident for the fo norm. Unlike in Figure 9, the trajectories do not align the curve q f'(2) when P is selected as the mean P(z) 1 E 2z.
The effect of network width. We investigate how the width of the network influences the trajectory alignment phenomenon. We vary the width m using values from { (64, 128, 256, 512] while keeping other hyperparameters constant. In Figure 13, train 3-layer fully connected neural networks with tanh activation and an initialization scale of = 1. Similarly, in Figure 14, we conduct experiments using the same configuration but with ELU activation, training ELU-activated fully connected neural networks.
Consistent with the observations from Figures 6 and 7 in the single data point setting, we commonly find that as the network width increases, the alignment trend becomes more pronounced. In both Figure 13 and Figure 14, all trajectories fall within the EoS regime. However, narrower networks
21

---

(a) P(z) 1. E>
(b) P(z) = ||=||1
(c) P(*) = ||1||
(d) P(z)
||
1.
1.2
0.-
0.3
2.2
0.3
3-
0.7
02
0.0-
0.2
-5.0
-2.5
25
5.0
100
200
30C
20 40
10
Figure 12: GD trajectories (1000 iterations) under the generalized canonical reparameterization (Eq. (5)) for different choices of the function P trained on small subset of CIFAR-10 image dataset. The plots depict the training of a 3-layer fully connected neural network with ELU activation, a width of m = 256, and an initialization scale of a = 1. The function is varied to be the mean, 4 norm, lg norm, and lor norm. The alignment behavior of the trajectories is more prominent when P is chosen as the mean or & norm, but less evident for the &o norm.
(a) 64
(b) 128
(c) 256
(d) 512
3.2
84
1.
5.1
0.
3.5
0.
0.5
0.5
3.4
0.4
0.3
2-0
p=0
02
Figure 13: GD trajectories (1000 iterations) under the generalized canonical reparameterization (Eq. (5)) for different network widths trained on a small subset of CIFAR-10 image dataset. The plots depict the training of 3-layer fully connected neural networks with tanh activation and an initialization scale of a 1. The trajectories are shown for network widths of m 64, 128, m 256, and m- 512.
(a) m =
(b) m = 128
(c) 7FL -256
(d) T/ = 512
1.0
p
p
08-
0.
0.
0.
0.
00.5
00.5
0.5
0.1
0.4
0.2
02
0.2
5.0
-2.5
2.5
5.2
+3.0
+2.5
0.0
2.
5.0
Figure 14: GD trajectories (1000 iterations) under the generalized canonical reparameterization (Eq. (5)) for different network widths trained on a small subset of CIFAR-10 image dataset. The plots depict the training of 3-layer fully connected neural networks with ELU activation and an initialization scale of a = 1. The trajectories are shown for network widths of m = 64, m = 128, m 256 and m 512.
(m 64) show less evidence of the trajectory alignment phenomenon, while wider networks (m = 256, 512) clearly demonstrate this behavior. These findings emphasize the significant impact of network width on the trajectory alignment property of GD.
The effect of data label. In our previous experiments, we assigned labels of +1 and to the dataset. However, in this particular experiment, we investigate the training process dataset with
22

---

zero labels. This means that all samples in the dataset are labeled as zero (3) = 0 for all < S n). Figure 15 visualizes the training of 3-layer fully connected neural networks with tanh activation and an initialization scale of a = 1. The network widths are varied from (256, 512, 1024). Interestingly, the GD trajectories align with the curve y = <'(22 contrast to our observations in Figures 13 and 14. These results suggest that the data label distribution also influences the alignment curve of GD trajectories. As future research direction, it would be intriguing to investigate why
setting the labels as zero leads alignment towards the curve ((m)) which aligns with our theoretical findings in the single data point setting.
(a) 256
512
(c) 1024
5.8
0.-
%
0.8
05-
3.1
322
-5
Figure 15: GD trajectories (1000 iterations) of the training for 3-layer fully connected neural networks with tanh activation and initialization scale of a 1. The networks are trained on small subset of the CIFAR-10 dataset, where all labels are set to zero. The network widths m are varied, including values of 256, 512, and 1024. GD trajectories exhibit alignment behavior, aligning the curve (p)
q
23

---

B Proofs for the Two-Layer
Linear Network
B.1 Proof of Theorem 4.1
We give the proof of Theorem 4.1, restated below for the sake of readability.
Theorem 4.1 (gradient flow regime), Let vj & (0, 33 be fixed step size and be a loss function satisfying Assumptions 4. and 4.2. Suppose thut the initialization (po, 90) satisfies lpol S and -(1)
qo ( min 167 2n }}. Consider the GD trajectory characterized in Eq. (6). Then, the GD
iterations (pt, 9:) converge to the point (0, q²) such that
90 S q- S exp(C1) ) S 2q0.
where C. 8qu
min
2(00 11 (()
0.
2go
We note that in the given interval depending on the value of r(1). However, this does not impact the correctness of the theorem.
2- min 1ôn
r(1) x00)))
}). is possible that the interval is empty,
Proof. By Proposition B.I, pt. converges to as + oo, and for 2 0, we have
40 S qe S exp(Cn))) )00-
Since the sequence (4t = is monotonic increasing and bounded, it converges. Suppose that Ør + g" as co. Then, we can obtain the inequality
qu S - S exp(Cn"))q.
as desired.
Proposition B.1. Suppose that 1j E (0, 2. |Pol $ 1, and qo E ( min
{ 16n / -r21))
}).
Then for
(2-
any 2 0, we have
2(90 1) r(1)
Ipel 1 min
11,
40
2go
and
2(go 1) r(1) |
9o S exp
8n'200 min
40 S 2yo
20
2go
Proof.
We give the proof by induction; namely, if
2(gg - 1) r(1)
2(go 1) r(1)
Ipd < 1 min
4o ge exp
8n7 qa
min
290
go
2qu
9o
240
are satisfied for time steps 0 S t S l for some k then the inequalities are also satisfied for the next time step k 1.
For the base case, the inequalities are satisfied for 1 by assumptions. For the induction step, we assume that the inequalities hold for any 0 < t < k. We will prove that the inequalities are also satisfied for t = +1
By induction assumptions, 4o q - 2go and |px
< 1, so that we have
2 pk+1 2r(pk)
r(1)
<
+ 1 r <
+ n?
90
9k
40
where we used r(1) < r(px) < r(0) 1. Since ạ0 <
r(1)
<
r(1)
we have
2n
2n
Pk+1
1,1 r(1)
< max
pk
[2(go
1) r
min
11
24

---

2(00
1)
r(1)
Σ min
90
2q0
This implies that
| 2 (1
2(9 1)
r(1)
})
2(40
1)
r(1)
k+1
min
|pa| 5
min
240
40
2go
which is the desired bound for |ps+1l- For any S k, we also have
4m
pp = (((()))))))))))))))))))))))))}}) S
9t
qt+1
where we used the induction assumptions to deduce 4go 2 2q. 2r(er(p)) 2 2go 110-13
1 0. This
2-1
gives gk qu 2 go. Furthermore, note that Anj²riq0 I 417² From this, we have
1 I +.2 the ratio _ 1 1].
16n
k
&
qo
qt
qt
log
log
S
2
@k++1
g++1
t=0
t=0
k
< 81)
Pt
2
t=0
r(1) 2(9o
1)
81j 90
2
min
90
2go
1=
2(90 - 1) r(1)
1
$ 817
min
40
2yo
where the second inequality holds since [log(1 2)] S 2|2| if |2|
Moreover, this implies that
2({0
1) r(1)
}
9u+1 I exp
817290
min
9o.
2go
Since do
2 r(1) 12-n² 27
we have
2(go - 1) r(1)
min
2 n7
go
2qo
Therefore, since
1 167
we can conclude that
2(go - 1) r(1)
9o S 1 & exp
8n1200 90
min
go S expp(8n0)0n S 24n-
2go
the desired bounds for qk++1
B.2 Proof of Theorem 4.2
In this subsection, prove Theorem 4.2. From here onwards, we use the following notation:
r(p)
All the lemmas in this subsection are stated in the context of Theorem 4.2.
Lemma B.2. Suppose that the initialization (po, qu) satisfies |pul S 1 and go E (Co, 1 - 5). Then for
any 2 0 such that qi < 1, it holds that
|pel 4, and q: S q(+1 < (1+ )(r²))).
25

---

Proof. We prove by induction. We assume that for some t 2 0, it holds that lpal S 4 and 2 - 9: - 1. We will prove that |pt+1| S 4 and S 2: qt+1 & (1 + )) For the base case, |Po| - 1 4
and 2 S co < ge S 1 holds by the assumptions on the initialization. Now suppose that for some
t 2 0, holds that Ipx| S 4 and 2 S q: S 1. Then for small step size n, we have
20 (pr)
2 2) 2 | (p.)))Pb) 2 (p,,) Pel-
Consequently, by Eq. (6),
|pe+1l = (1 + 1/2 (vt) })p (p((pe)
< max {
2 9t
|ptl
S 4.
where we used 1-Lipshitzness of Moreover.
- STj² S - -2pppell S at
=1 =1 - (())))))) r((e)) $
9-+1
where we used qt E (2: 1] and ppr(pt)) = (e(p))) S from 1-Lipschitzness of t. Hence, gr S 0:+1 S (1 + (1r))))) as desired.
Lemma B.2 implies that pt is bounded by a constant throughout the iterations, and qt monotonically increases slowly, where the increment for each step is O17²). Hence, there exists a time step
T (()r)= = 2)) such that for any S T, it holds that q: S 5. Through out this subsection, we focus on these T early time steps. Note that for all 0 < < T. it holds that q: E (co, -
Intuition on Theorem 4.2. Before we dive in to the rigorous proofs, we provide an intuition on Theorem 4.2. Lemma B.2 establishes that Pe is bounded and q. monotonically increases slowly, with
an increment of O(n²) per step. Lemma 2.1 shows that the map fg. (p) = (1 2r(pe) Pt has a stable 2- 2-period orbit {++( when q. E 10. 1). Consequently, when q: is treated as a fixed value, Pr converges to the orbit {++((:)) leading to 8; converging 1. In the (early) short-term dynamics, qe is nearly fixed for small step size and hence converges to 1. In long-term dynamics perspective, de gradually increases and at the same time, 8, stays near the value 1. In Theorem 4.2, we prove that
it takes only ta = Os (log(1) 1) time steps for st to converge close to 1 (Phase I, t < ta), and after that, 9t stay close to 1 for the remaining iterations (Phase II, t > ta).
We informally summarize the lemmas used in the proof of Theorem 4.2. Lemma B.3 states that in the early phase of training, there exists time step to where St, becomes smaller or equal to 2-r(1)² which is smaller than 2(1 + - Lemma B.5 demonstrates that if is smaller than 2(1 + 17 ) and 1px| 2 f(1 - 4). then |s; 1| decreases exponentially. For the case where |px| < f(1 4) Lemma B.6 proves that 1p:| increases at an exponential rate. Moreover, Lemma B.4 shows that if 8, < 1 at some time step, then 8141 is upper bounded by + O(r)). Combining these findings, Proposition B.7 establishes that in the early phase of training. there exists time step 1 such that $1 = 1+0O Lastly, Lemma B.8 demonstrates that if = + Os (772, then |s: decreases exponentially.
1s:11-100 1- h(p:)
Now we prove Theorem 4.2, starting with the lemma below.
2 Lemma B.3. There exists a time step to = Os,e(1) such that Stg S 2-r(1))
Proof. We start by proving the following statement: for any 0 0 2-r(1) St < 2r(1)-1,
then $141 < 2r(1)-1 and |p++1| S (1 m(U) Mpdl. Suppose that _2 2 2-1(1) <s < 2r -1 Then from
Eq. (6), it holds that
Pt+1
=
1
+ (p:t)
1 -
St + +1
2
2 11- +1² < - r(1)
Pt
St
2
26

---

for small step size 7). Hence, |p++l S (1 -11)))))00 Ippll. Now we prove s1+1 < 2r(1) -1 Assume the
=())
contrary that $1+1 2 2r(1)=1 Then, r(p++)
9++1
9t-1 1 so that [Pt+1| A(1 2)- By
2:11
Mean Value Theorem, there exists p. E ((pet1l, (Pel) such that (recall that r((p)
< 0 for p > 0)
r(p)²
1
1
r'(t))
(pps)
(pe+1l)
r(pet))
r(pvel (1mll pe+
r(pe)
r(o))²
1
r' r'(ptt11)
(r(1)(pel
r(pe)
r(1p++112
2
1
In'( - -21)
r(1) - 2)
r(pe)
(1 - )
2
r(pt) ls.e(1),
where we used Assumption 4.3 (i) and F(1 - 1)) s lpet1| S 5(1 - -1-1))))p )p=) - |pr] Consequently,
(+1) Isse())
9++1 = (1 + 0 s6+1 r(p(+1)
$ r(p:) = 8t < 2r(11))
91
(((p))
for small step size 17. This gives a contradiction to our assumption that 8641 2 2r(1)-1 Hennce, we can conclude that St+1 < 2r((1) as desired.
We proved that for any S 1 S T, if 2ー2 < 8, < 2r(1)-1, holds that 81+1 < <2r(1)-1 and |pe+a/ $ (1 - lp:l At initialization, lpol S and qo < 1, so that 80 < <r(1)-1 If 80 S 2-r(1))
then to = 0 the desired time step. Suppose that 8o --r(1)- Then, we have 81 < 2r(1) and
|pIl S (1 - = ( \pol I 1 - - Then we have either $1 S 2-P(1y) or 2-1(1 < s < 2r(1) In the previous case, tu is the desired time step. In the latter case, we can repeat the same argument and obtain $2 < 2r(1) and Ipal < (1
2 )2 By inductively repeating the same argument, =
can obtain a time step 40 < log(f(1 - / log( = = O5,( such that either 8to S 2-r(1) or |r.. S f(1 - 9). In the latter case, (p..)) 2 1 - 8
2
2 > qo and hence Sto < 11 < 2-r(1) Therefore,
tu (,e(()) (1) is the desired time step satisfying Sta S 2 r1)
According to Lemma B.3, there exists time step 40 O5,(( such that 8t. S 2-1(1] for small step size 17. Now we prove the lemma below. Lemma B.4. Suppose that < 1. Then, it holds that 8++1 S 1 0
2
< 2(1
Proof. For any pe (0),(()), we have r(p) 2 2h so that (p)] = (-11+
2r(p) )p. Hence,
dt
a
2r(p) (1 pr pr'(p)
r(p) ) 1,
fg. (p)
Op
qt
for any p E ()) By Assumption 4.3 (ii) and convexity of E, both r(p) and 1+ (p) t" {'(p)
r(p)
T(p)
are positive, decreasing function on (0, (2)). Consequently, 2p l.. (p) is a decreasing function on (0, (.))).
Now note that qe which means F(1) 0 < <<(())) < <F(%) by the definition of f. Note that 3p Iax (p) at p = f(y:) evaluates to
i Jq (f(2)) = 1 (F(a)) Pp
2F(co)) 0.
r(F(g)) 2 1+
r(f(co))
where the first inequality used Assumption 4.3 (ii) and f(g:) < f(co), which comes from g: > Co ==
max(r(zo), }). The second inequality holds because q: co 2 r(20) where 20 == sup, { [x) } }. from the statement of Theorem 4.2.
>
((a)
27

---

0p (0, F(q:)). )). holds that öp l. (p) 2 O. In other words, If (p) is an increasing function on (0, f(g)))
Therefore, since
8 1 (p)) is decreasing on (0.)(2) and is nonnegative at f(9:). for any p E 0
Since 0 5 - 1, we have Ipal S f(q:) and it holds that
|Pt+1| = 1 2
p ppr((())) Ipal S -1
Ipal = 1f (pt) (f(ae))) = f(a.)
St
Therefore, with this inequality and Lemma B.2, conclude that
9++1 9t r(pe+1) (1 + (1))) ((7)) - r(f(u:)) (1 + F(F(()))(
gt
(1+0
Se+1 r(pe+1)
Using Lemma B.4, we prove the following lemma.
Lemma B.5. For any 0 S t if < 2(1 -
- 1| > o/2 and r(pt) s 1 - 3. then
|st+ - (1- -
where de (0, } is constant which depends on ó and E.
Proof. By Eq. (6) and I-Lipschitzness of L,
PL+1 =
2
('pir(p:)² < - =1²
Pt
St
so that pt and p++1 have opposite signs. By Mean Value Theorem, there exists O2 between -1 and
(1 2 + -/ppr(pp))) ) satisfying
r(po+1)
(-Pp +
2(x
+n7 pppr((1)) Pt)
81
1
r((pp))
12(s - 1)
+ 11 ppppr(pp)) pt
r(-p()
r(O.pe)2
8t
1
r'(Opp))
2 2(s 8t
+ n/ n2ppr(p:)) lpel,
(8)
r(p+)
r(0.pe)2
where the last equality used the fact that and 0.p: have opposite signs and (2) and z have opposite
signs. Note that [0.p.l is between [p.] and [Pt+a]. Consequently, the value [r'(0.p)) between In' (p.)) r(d.p.)²
r(p.)²
and I''(0e+)]l -( by Assumption 4.3 (i). We will prove the current lemma based on Eq. (8). We divide
into following three cases: (1) 8t 2 and 81+1 2 1, (2) 8 2 and 8, and (3) 8, < 1.
Case 1. Suppose that s, 2 1 and $111 2 1. Here, we have Ip:| 2 f(g)) 2 f(1 2) and similarly |Pe41| 2 f(1 - ). By Assumption 4.3 (i),
1''(0.e))) 2
(P(1-4)) (1-
7(0ep))²
Hence, Eq. (8) gives
1
|r((f( - 1))1 2((st - 1)
6 2 1 (1-2)
1p)-1) r(p.)
(1 - 2)2
8t
Consequently, by Lemma B.2,
8141 qe(l + 0(1²))
91
+ 0(1)) < 81
1r'(f(1 5)Il
(((s: 1)
1
6
9: + 0
r(p(41)
r(p(+1)
(1 )
8t
2
|r(((1 - fiil
1
S St
(st 1)f
(1- 2 2
+ 0(72)
(1 )
i(1 3 (r(((((1 [))
S 8t
(s: 1) + +
2(1 2(1 2)²
where we used d: > eg 2 1 inequality:
2 and 8, < <1 + 1 2. Therefore, we can obtain the following
014111 - 15 1 i(1 1 9))
(st 1)+ 1)-1))0007)
2(1- 5)2
28

---

Case 2. Suppose that s; 2 1 and 8141 < 1. Here, we have r(p(+1) > 90+1 2 2 r(p)), so that
|pe+a| Ipel- Consequently,
(0(0p))) r(0px))2 I (r(p))) by Assumption 4.3 (i). Hence, we can deduce from r(p))²
Eq. (8) that
1
(p()
(2(st-1)
+ +npppr((p))* lpd
2
r(p(+1)
r(pi)
r(p))²
8t
1
2lpir' (p:))
(st 1) 1/² (pt)) r(pt]qt
r(t))
1
2/por' (pr))
2 r(p))
(st 1) r2prr(pe) r(pe)qt
1
2per' (pt)
(st 1) O r(pe)q
r(p()
where we used (p:)'(:)) S r(p)) since 1 + per'(p)) ["(p)) > 0 and |p:l S 4 by Lemma B.2.
T(p)) r(ps)
Consequently, by Lemma B.2 (gt ((+1) and Assumption 4.3 (ii),
at 2 + 2ppr' (pr) (st 1) 0(12) 2 81 + r(pt)
8r'(4) r(4) (st -
81+1 2
r(p+11)
Note that 1 > 1+ (1)
= (4)
> 0 holds by convexity of f. Therefore, we can obtain the following
r(4)
r(4)
inequality:
0 s 1 - S1+1 5 - (1+ 84)
(st (1)
r(4)
where -1 < 1 (4) r(4)
< 1.
Case 3. Suppose that 8 < 1. By Lemma B.4, it holds that 8141 S 1 + O(12). Moreover,
assumed r(pr) s 1 - &. so that 1p:l f(1 2t) We also have
1p--1= = (-1 1 2 - pppr(())) Ip:l N -1
2
)
1 0 4
172
Ipdl > Pel 2
n
St
2
where we used the assumption |sg | 1| S5. and |pr(p)]
[k(p)) < 1 due to I-Lipschitzness of E.
Consequently, by Assumption 4.3 (i), holds that (H(ppl)]
|r' (F(1 DM (1 1-
Hence, by Eq. (8), we
-(Oepa))
have
1 2 r(pe+1)
+ |r(((1 - fI)
2(11 (2(1-5)) st)
7 1 T.
r(pt)
8/2
81
|r'(f(1 11)1
2(1 - s:)F (1- 5)
>
+
r(pt)
(1 - 1
4
2f(1 - (f(1 - 11)
+
(1 (t)
r(p))
(1 4)2
and hence, by Lemma B.2 (g: Σ 4++1) and q: > cu 2 1
2' we get
8 + f(1 - 1)) ( (f(1 - f1)
St
r(p++1)
gt
(1 22
(1 st).
Therefore, we can obtain the following inequality:
11 1- f(1 (F fI)l
- < 1 s S
(1-st).
(1 (1- 2)
where we used Lemma B.4 to obtain the first inequality.
29

---

Combining the three cases, we can finally conclude that if we choose
f(1 (f(1 1>>l
2 (1 (1 r(4)
tr((4) f(1 (f 9i1l
(1 (1- }
1 0.
d = min
2' 2(1 2)²
E
2
then 8141 1| S (1 - d)}s: 8t - 11|
Lemma B.5 implies that if st < 2(1 + - and Ipal 2 f(1 - 4) then |st - 1| exponentially
decreases. We prove Lemma B.6 to handle the regime |Pil < f(1 - 2) ) which is stated below.
Lemma B.6. For any 0 S t S T. if (pt) 2 1
6 it holds that
Pr+1
4
Pt
4 4-8
1 0 S t T. Consequently,
Proof. If r(pi) 2
1' then 8t
9t
2
4-26 1-6"
where we used qt
8 for any
r(p))
1-
2
T.
P++1 Pt St
2 1² (ppr((1))
2
2(4 8)
1/
2 2- 8 11 2
4
4 4-26 28
4 - 8¹
for small step size 17.
Now we prove Proposition B.7, which proves that 8, reaches close to 1 with error bound of (17²). Proposition B.7, There exists a time step to Os,, (log(() satisfying
Stg
Os.e(1²))
Proof. By Lemma B.3, there exists a time step to = -.5,e()) such that Sto S z-11 Here, we divide 2
into two possible cases: (1) 8to < 1, and (2) S Sto S 2-1(1]
Suppose that Sto < 1. By Lemma B.6, (ptt.) 2 1 - ] (or equivalently, Ipro S f(1 - 4),, then there exists a time step f I to + log( ((1- 3) / log( 4 = 00,((1)) such that |pt,1 2 f(1 - Q [P..
Case 1.
We denote the first time step satisfying |pt, 2 f(1 2) and 1 2 to by f = (1)) By Lemma B.4, it holds that St, I 1 O(n²) since -t < 1. Consequently, if St, 2 n,² so that t = f is the desired time step. Hence, it suffices to consider the case when st, < 1 T Here, we can apply Lemma B.5 which implies that
2 then |8t, - 1| - (17²)
| - 1| S (1 dd]ss, - 1 + (17?)
where d is a constant which depends on 8 and E. Then, there are two possible cases: either |st, - 1| < Ofr)d or [st++1 1| S (1 4 St - 1/. It suffices to consider the latter case, suppose that |st,+1 - 1| S (1 {) - 1/. Since are considering the case St, 0(1)). |p.+ 1 2 |p,, 2 f(1 4)) must be satisfied unless $t, 1 + O(7") already holds. If s1,+1 2 1 2 then |st,+1 +1 1 s O(7?) that to = 11 + is the desired time step; if not, we can again apply Lemma B.5 and repeat the analogous "" argument. Hence, there exists time step l2 S l1 log( ( )/ log(1 4 = Os.e(log(1n such that |sca 1| On1 Oz,e(1²).
n again by Lemma B.4,
we have 8 S + (1))). Since
|
2
Case 2. Suppose that 1 S sto S 2-r(1)- Then, r(P.)) S 9t < 1 5 so we can apply Lemma B.5. There are two possible cases: either |st++1 1| S 0
2
Os,e(12), or | 1| S (1-5))s, s1, -1.. It suffices to consider the latter case. If sto+1 2 1, we can again apply Lemma B.5 and ,,² repeat the analogous argument. Hence, we can obtain a time step to S tu +og( 1- )/ log(1 () = (log()()( ))) such that either sg or |ssg 1| = 0s,e(1²2) satisfied. If st
we proved
in Case that there exists a time step 1/2 = 1o + Os.e(1o((17))) such that |s1] - 1 s Os.e(1²). and
this is the desired bound.
L
30

---

Now we carefully handle the enror term Onn²) obtained in Proposition B.7 and provide tighter bound on 8, by proving Lemma B.8 stated below.
Lemma B.S. |s 8t 1| (o,e(1²), then it holds that
884+1 - 1 - h(pp+1)0711 - (1+ 2per'
-1 1 -hpp))) + -sell1'p))
r(pt)
where h(p) - pr(pl" ((p) + p>r(p))) for p + 0 and h(p) == - 2 for p = 0.
Proof. Suppose that 8, 1 + Ox.((1²). Then, |pe+1l |
2 + + -lpdl = (1 +
se O,,e(1²)) which satisfies the
Or.e(r"))vll By Eq. (8) proved in Lemma B.5, there exists following:
1
(2(x--1) - -(ppr((p)*)) Pt
+
r(p(+1)
r(ps)
r((1+a))p
St
1
+ (pp.))
+ +Oael1Ppt))
(2(5)-1)
+ 1 ,r(pr(p.))) Pt
r((t))
r(pc)2
8t
1 (pt)
11) + p'prr(pi)) r(p))2 $t
(2(80-
+
pt + 00.e(1'p2)
r(p))
where we used the Taylor expansion on
(p()) with the fact that r(m)²
r'(p) is bounded on [-4, 4] and dp [(p)2
that Ip:| < 4 to obtain the second equality. Note that gt+1 = (1 - 77 p2 (p: (p.)(22 - r(pp)))- gg. by Eq. (6). Consequently.
8t + (prr'(r))
Sc+1 = (1 - (2::
r(pt) (st - 1) + +1 + Oo.((1'13)
2prr'
= (1 + 1/ (p:r - (4)) r(p:))))
(0+ - + v/ (: + One(1'p²)
r(pe)
(p:) (st - 1) + 1/ p2 (2:) r(p:)) + pp)))())) Os.e(n' (pe)
+ 1 r(pt)
Here, since 8t 1 O,,e(1²), we can rewrite
r'pār(m))24 - r(p:)) (
+ + (pr)at
= r2pir(p 1)s; + 12 (p: ((p)ss
12pp) + 0s.
which results in
(pi) 8c+1 1+ (1 (1+ r(pt)
(st - 1 + ( (pr) + Os.((n'Di).
Note that / is even, and twice continuously differentiable function by Lemma B.9. Consequently, k0(0) = 0 and hp) Or(p), since h" is bounded on closed interval. Consequently, l((p+11)
h(11 0s.e(1²))pp) = h(p:) + (n.(n72pp)) Hence, we can obtain the following:
8141 - 1 - A(pr-1)1) )m² 8841 1 - Mpx)) + Ox.((''7))
p.r(p))2
+ +Dir(m))) n72 Ose(nnppi)
S111 1 + 2 (p:))
1 + (pe)
81 1+ 1 ppr(pp)) + vir((n))) 12 Oae(n'p1)
r(pe)
2 ('()))
(1+ (ptr
(st 1 - -h(p)n??) + One(1tp?)
r(pe)
31

---

Note that r(p)) = (1 + 0s.e(1))) 2 (1+022(7))40 2 co 2 r(2)) for small step size 1) where 20 = sup ar (2)
2 - } Consequently, it holds that 1 + (per 2 0. Therefore, we have the r(ps)
r(2)
desired inequality:
18+11 h(p+1))?1 (1 Sper'(pe)
r(pt) |st - MPp)) + Ose(1(p))
We now provide the proof of Theorem 4.2, restated below for the sake of readability.
Theorem 4.2 (FoS regime, Phase 1). Let y be small enough step size and f be loss function
2x (2) satisfying Assumptions 4.1, 4.2, and 4.3. Let 2g == sup-( r(=) --2 and cu = max 2)-
Let 6 E (0, 1 co) be any given constant. Suppose that the initialization (po. qo) satisfies (pol S 1
and go (co, 1 - 8). Consider the reparameterized GD trajectory characterized in Eq. (6).
assume that for all & 2 0 such that 4t < 1, we have Pt + O. Then, there exists time step to OO,g(log(1) ")). such that for any t 2 ta-
4:
= + h(p))nf + Os,e(1²t),
r(p:)
where hpp) -
pr(p) (())
+ p²r(1))) for p + 0 and h(p) 2,"( for p 0.
Proof of Theorem 4.2, By Proposition B.7, there exists time step to Os.e(log(n !)) which
satisfies:
4:
|8,; 1|
Os,e(1²²)
r(pe:)
By Lemma B.8, there exists a constant D > which depends on 6, such that if (s, - 1| = Os,e(1?²) then
[st+1 h(p(+1)) [ 2per' 2per'(pt)
1s: 1 - + Dry'p
(10)
r(p))
Hence, if [sc 1| = (7.e(11²) and 1sc 1- h(pp)1??1 2 (-
(-per(n)) Drj", then r(pa)
1st 156-1- h(pp)) |st-1 - h(p(+1m711 S 1 per'(pe)
(11)
r(p:)
For any / < T, we have < & so that if |st 11 1| Ov,2(1²),, then r(p.) - (1
1 i for small step size y7. From Eq. (11) with = thg we have either
1s: - h(p(g)n71 < (- Pe; rpp:) Dij"
(pt:)) -(pp)))
or
(14 f(1 - ()r 1))
[8 -1 - h(p: +177 <
1+
|80g h(pp; )n²²
(1 4
where we used Assumption 4.3 (ii) and |p. f(1 )) In the later case, |5t++1 = Os.e (T²²) continues to hold and arguments, we can obtain the time step
again Eq. (11) with + 1. By repeating the analogous
log
De ""((0] hPpg nm1
to t +
Os.e(log(n-1)
+ F(1 (P(((((1 1)
log
(1-4)
which satisfies: either
(
P..Tlp ) Dmy,
1st, -1 h((pt))7
('tt.)
32

---

or
1st, -1 - h((p,,) )n (- (0)) Drj² =
r(pt)) Dny² - (- 4r(4) Dry²
((-r(a) (4(
(-Ar(4) (r(4))
rp(pa))
where we used Ipel S 4 from Lemma B.2 and
ar(2) ((2) 2 [ for any by Assumption 4.3 (iii).
By Eq. (10), if (st 1- h(pt))/ - (- 4r(4) (-+(aa
Dry* is satisfied for any time step i, then
1841 - 1 - h(p:-1)m < 11 2per' (pi)
)
(-
4r(4) Dif + Du' s
4r(4)
Dm
r(pt)
r((4)
r4(4)
by 1pl S 4 from Lemma B.2 and Assumption 4.3 (iii).
Hence, by induction, we have the desired bound following: for any 2 ta
8 h(p))n1 - 4r(4)
Drj = 0s,e(1)).
r'(4)
by |Pel < 4 from Lemma B.2 and Assumption 4.3 (iii).
B.3 Proof of Theorem 4.3
In this subsection, we prove Theorem 4.3, We start by proving Lemma B.9 which provides a useful property of h defined in Theorem 4.2.
Lemma B.9, Consider the function defined Theorem 4.2, given by
h(p) = 1 ( p²r(p)))
pr(p)
if + 0, and
= 0.
(0)
Then, h is a positive, even, and bounded twice continuously difserentiable function.
Proof. It is clear that h is even. We first prove that h is positive. For any p & 0, it holds that
M()) = (p) 1+ pr (p)
pp(p)3
0,
r(p)
since pr(p) (P) and + pr (p)
(")p) by Assumption 4.2 and convexity The function
r(p)
r(p)
continuous since limp-0 hbp)
h(0).
Continuous function on compact domain is bounded, so
bounded on the closed interval [-1,1]. 1]. We can rewrite
h(p) 2 p22r(0) (- (p) -1)
1
r(p)
Note that p²r(p)²=
= 2p) 1, and
r(p) -1) is positive, decreasing function p >> by (-prie) pp)
Assumption 4.3 (ii). Hence, h is bounded on [1, oo). Since h is even, h is bounded on (-o0, 1].
Therefore, a bounded function on
We finally prove that h is twice continuously differentiable. Since r is even and C on R, we can check that
(p()))))r"())))) + +pr(p)(5r(p) + 2pr r(p)²
p * and if 0. 2p) r((p)
h'(p) =
Moreover, for any # 0.
2r 2 (p) - r' (p)) pr())3,(3(p)
hp(p) - 2
r(p)3
r((p)²
- -rrpp)² - 7pr(p)r (p)p))p(((()))))) (r((()rr (r)))((p))))
and
(4)
k²(0)= = (0)
Since limp-s0 (0), we can conclude that J is twice continuously differentiable function.
D
33

---

We now give the proof of Theorem 4.3, restated below for the sake of readability.
Theorem 4.3 (EoS regime, Phase II). Under the same settings as in Theorem 4.2, there exists a time step to @ qu)n 2) such that qt I 1 and qt > 1 for any to. Moreover, the GD iterates (pt. qt) converge to the point (0, q²) such that
1
q
2r" Oo,e(11").
Proof. We first prove that there exists a time step to such that qt, >> 1. Assume the contrary that q: S 1 for all t 0. Let to be the time step obtained in Theorem 4.2. Then for any 2 tas we have
r(p:) = (1 h1-h(p))n + 10--0555,5 Os S h(p))n1²
for small step size 1). The function g(p) = r(p) - h(p)n" even, continuous, and the function x/2
value g(0) (0) > Consequently. there exists positive constant for all (-e, e). Then, we have [pxl 2 for all > la, since g(pr) S 0. Moreover, 8, 2 for any t 2 ta by Theorem 4.2 for small step size 1). This implies that for any 2 ta.
such that g
qt = 1 r'p] - 1) = 1 -
1
S - 111---1-1 (e))
g+++1
so qt grows exponentially, which results the existence of a time step 1 2 la such that qrc 1, contradiction.
Therefore, there exists a time step tú such that 9s S 1 and qi for any ta, i.c. qt jumps across the value 1. This holds since the sequence (4: is monotonically increasing. For any t < to, we have Q++1 qt + (n²²) by Lemma B.2, and this implies that to S2((1 qo)n 2). as desired.
Lastly, we prove the convergence GD iterates (pt, gt). Let t > to be given. Then, q 2 9t+11 > and it holds that
Pt+1 2r(p:) 1 12rrpppr(
- - 1<1.
Pt
gt
91s+1
Hence, |Pel is exponentially decreasing for >>> ts- Therefore, converges to 0 as f - cc. Since the sequence (ge) monotonically increasing and bounded (due to Theorem 4.2), converges. Suppose that (Pr, Q:) converges to the point (0, 4-' By Theorem 4.2, we can conclude that
112
9 1
Oa.(11²))
2r"(0)
which is the desired bound.
B.4 Proof of Theorem 4.4
In this subsection, we prove Theorem 4.4. We first prove a useful lemma which bounds the Hessian of the function (U, v) v Uv, stated below.
Lemma B.10. For any O (U,v) with E RRmxd R" and E Rdd with ||0||2 1, the
following equality holds:
Vo |1, S.
Moreover, if X is an eigenvalue of 2 (U.
(U.,)) (u Un), then - > is also an eigenvalue of 2 (tr ( Un..
Proof. We first define the notations. We the operator @ represent tensor product, or Kronecker product between matrices, For example, for any given two matrices A = (a.s) E jpmxx and B, we define A B by
B
ain B
180 B
aml B
amn B
34

---

We use 0mxn to denote m by matrix with all entries filled with zero, and I, denotes n by identity matrix. Now we provide the proof of the original problem.
Let U = (us)) E IRmx t = (us) E IR. and x: = =x²) E TR be given. Then,
U- Σ o.USsTz
We vectorize the parameter O vU, by (V), Um Ul1..... "il....,mmm... Uml..... Uids.... Uma) E Rm+ndd Then, we can represent the Hessian as
V (er (0 Uw)
x Odxd
2 mm
(12)
For any given IR and y I2 with c² + |1yl1
1, we have
0
Oaxd
c
=
1(
x T y
c
)
2
$
|
= 1.
Hence, by definition of matrix operator norm, we have
0 x
1.
Oaxd
2
Therefore, we can conclude that
Viu.o (v Un)){ (U)))))
-
(1
x
| 2
2
I
11.
Oaxd
0 Oaxd
Now suppose that À is an eigenvalue of V (U, Ux). We note that for any given matrices A and B, if A, is eigenvalue of A with the corresponding eigenvector u. and As is an eigenvalue of B with the corresponding eigenvector us then Aς is an eigenvalue of A & B with the corresponding eigenvector u, us, Moreover, any eigenvalue of A B arises such product of eigenvalues of A and B. Hence, using Eq. (12), we have
À is eigenvalue of the matrix
0dxd
We denote the corresponding eigenvector by (c,yT) TT where E IR and E Rd i.e., it holds that
x
C
x
c
O Oaxd
y
C
y
Consequently, we have
-c
y
-
Odxd
C x
and this implies that
0 x
is an eigenvalue of the matrix
x
Oaxd
Therefore, by Eq. (12), -Å is an eigenvalue of Vi V Ux).
Using Lemma B. 10, we prove an important bound on the sharpness value provided by the Proposi- tion B. stated below.
Proposition B.11. For any + = (U, v) with U E TRmxdd E RM, and a RRd with ||m||2 1, the following bound holds:
Amax (0) - (v Ux) (|10 ( <
35

---

Proof. The loss Hessian at θ = (U. can be characterized as:
V Uz) (Ve( Uz)) 62 +6((0 Vz)Va( Ux).
(13)
We first prove that A (V)((0)) [ Note that the largest absolute value of the eigenvalue of symmetric matrix equals to its spectral norm. Hence, [Q)
max Amxx V (()),, -mmin (V (())))
it suffices to prove that Amax (V) >
-Amin ()).. Let w denote the eigenvector of (()( corresponding to the smallest eigen-
value min V V ((e)) with ||||| 1. Then, using Eq. (13), we have
Amin (Q)())) = w V2 = (v Ux) (Ve (Ve(vTUx)) Ux)) 82 w + Ux)) Vā(v Ux)
Ux))u
> (vv"Ux))u V Un/) (e Un) V (v
where we used Lemma B.10 to obtain the last incquality. Note that the matrix ("(vTUx) (Ve Uz)) is PSD, so that Amax (V) 2 (f Un) Vo Ux))
[e'(w T V (v Ua) 2 Ami (Q)( Therefore, Amaax ( (
Now, we have the following triangle inequality:
Dmmax ()) - ( T Ux) (1Ux ||-133)
Ux)
1,
where the last inequality holds by Lemma B.10 and 1-Lipschitzness of f.
We now give the proof of Theorem 4.4, restated below for the sake of readability.
Theorem 4.4 (progressive sharpening). Under the same setting as in Theorem 4.2, let ta denote the obtained iteration. Define the function À R200 - IR given by
À(q) = (1 Flg)r (F(2))
2 if S 1, and
otherwise.
Then, the sequence (X(q:) is monotonically increasing. Moreover, for any t 2 ta the sharpness at GD iterate O closely follows the sequence (X(q))) by satisfying
Xmax (0:) - ((a) $ Or(n)
Proof. By Proposition B.11, we can bound the sharpness \max at time step t by
26" (px)
Amax (..)
S
n9:t
Since (z(x)= rz() + (2), we can rewrite as following:
Per (pt)
Amax (Oi)
8 +
11.
(14)
By Theorem 4.2 and since is bounded function by Lemma B.9, we have s, 1 Oc(n²) for any t tg. Consequently, |s, - 1| Or(1²) and (r(p.) - q.1 Oz(17). Moreover, for any g d ((())r(((())) F'(2)()(((()) + (f())) F(2)r((((())
dg (
1 (1 f((r"(((a))) F(a)r(((())
9
r(F(a))
q²
36

---

so that
( d
(F
(1+ prr((())
lim q->1
lim
2.
dg
q
p-00
r'(p)
(P(a) is bounded on [ 1) and Taylor's theorem gives Therefore, 2q
per (pt) (((a) Oc(Ir(p:) a:l) = Oc(7²),
r(pt)
Gr
for any time step with q:. < Hence, if qe < 1, we have the following bound:
x(qo) $1
per (pt)
2
$ si + (((())
ppr' 2 + = =c(())=(((r)
r(p)) 17
(15)
where we used per(
por((p)) (1+0, per'(p)) r(ps) + +r""""), since 1 p.r'
((( >
r(pe)
r(pe)
r(pe)
implies Iper' (pt) < r(p)) I 1. Now let t be any given time step with qi 2 1. Then, r(p:) =
1 Or(n²), and since r(2) 1 + (0)))² + Ox(22) for small 2, we have [pt Oc()). Hence,
05 + per per'
(vr) + Oz(1) = Oz(1). r(p:)
X(q+)
S
(16)
for any / with qe 1. By Eqs. (14), (15), and (16), can conclude that for any 1 la we have
Amax)(.) X(ao{ $1+0A Or(n).
Finally, we can easily check that the sequence (A(4)))fo is monotonically increasing, since + 2x) is decreasing function by Assumption 4.3 (ii) and the sequence (g:) is monotonically r[=]
increasing.
37

---

C Proofs for the Single-neuron Nonlinear Network
C.1 Formal statements of Theorem 5.1
In this subsection, we provide the formal statements of Theorem 5.1. We study the GD dynamics on two-dimensional function L(x, y) == { where æ, y are scalars and ø nonlinear activation satisfying Assumption 5.1. We consider the reparamcterization given by Definition 5.2,
which is (p.g) I,
We emphasize that the results we present in this subsection closely mirror those of the Section 4. In particular,
Assumption C.I mirrors Assumption 4.2,
Assumption C.2 mirrors Assumption 4.3,
(gradient flow regime) Theorem C.1 mirrors Theorem 4.1.
(EoS regime, Phase I) Theorem C.2 mirrors Theorem 4.2,
(EoS regime, Phase II) Theorem C.3 mirrors Theorem 4.3, and
(progressive sharpening) Theorem C.4 mirrors Theorem 4.4.
The proof strategies also similar. This is mainly because the 1-step update rule Eq. (7) resembles Eq. (6) for small step size We now present our rigorous results contained Theorem 5.1.
Inspired by Lemma 2.1, we have an additional assumption on Φ below.
Assumption C.1. Let r a function defined by r(z) == ((=) for = + 0 and r(0) = 1. The function satisfies Assumption 2.4.
In contrast the function defined in Section 4, the expression + pr'(p) r(p) can be negative, which implies that the constant defined Lemma 2.1 is positive. As result, the dynamics of Pt. may
exhibit period-4 (or higher) oscillation or even chaotic behavior (as illustrated in Figure lb).
We first state our results on the gradient flow regime.
Theorem C.1 (gradient flow regime). Let E (0, )((())22) be fixee step size and ó be sigmoidal function satisfying Assumptions 5.1 and C.I. Suppose that the initialization (po, 90) satisfies lpol S and go E
Consider the reparameterized GD trajectory characterized in Eq. (7), Then, the GD iterations (pe, 9:) converge to the point (0, q²) such that
1-2 r(12
2 2(qu
1) r(1)
11 40 S 2go-
1o S q" exp
2n
min
Qo
Theorem C. implies that in gradient flow regime, GD with initialization (æo, yo) and step size converges to (0, y") which has the sharpness bounded by:
2(g0 1)
r(1) 40
}
1 2m min
Vo S Amax V ((. yo-
go
Now we provide our results on the EoS regime with additional assumption below. Assumption C.2. Let be function defined in Assumption C.I. Then is C* on R. and satisfies:
(i) z r(s)z
(2))
is decreasing on R,
(ii) z + (x)
is decreasing on and increasing on z < 0,
T[2]
(iii) 2 + sr(z) [
is decreasing on > 0 and increasing on 2 < 0, and
(iv) f())r((((()) >
38

---

Note that the function that arise from the activation ø tanh satisfies Assumptions 5.1, C.I, and C.2.
Theorem C.2 (EoS regime, Phase D). Let v be small enough constant and & be an ac- (2)
rivation function satisfying Assumptions 5.1, C.I, and C.2. Let 20 = supzt r(z) 2 - 2
(2) r(=) 2 -1}. and Co max(r(zo), r(2)) 2 (2, Let d (0, - co) be
21 sup,
any given constant. Suppose that the initialization (po, 90) satisfies |pol S and go E (co. - 8) Consider the reparameterized GD trajectory characterized in Eq. (7). We assume that for all t > 0 such that qr < 1, we have Pt + O. Then, there exists time step ta Os.o (log(n such that for any 2 ta
q:
= 1+ 06.
r(p:)
where h - R is a function defined as
((p) ()(
if + 0, and
h(p)
p = 0.
(0)
The main difference between Theorem C.2 and Theorem 4.2 is the error term which is O(r²)) in the former and On)²)) in the latter. This is because the 1-step update rule of di in Theorem C.2 is given by 4++1 (1 + (n))). while in Theorem 4.2 we have 4++1 (1 + O(1)))
Theorem C.3 (EoS regime, Phase II). Under the same settings as Theorem C.2, there exists time step tú = (2((1 )) such that qa, I and qi for any t >> ty, Moreover, the GD iterates (pt, (t) converge the point (0, q²) such that
9º = (")))
+0,
Theorem C.3 implies that in the EoS regime, GD with step size n converges to (0, y²) which has the sharpness approximated as:
Amax V 2 2
(n.
r [r"(0)
Theorem C.4 proves that progressive sharpening (i.e., sharpness increases) occurs during Phase II. Theorem C.4 (progressive sharpening). Under the same setting as Theorem C.2, let to denote the obtained time step. Define the function Ā 1330 - IR given by
(P(a))) 2
if and otherwise.
A(1)
2
Then, the sequence (A
is monotonically increasing. For any t 2 ta the sharpness at GD satisfying the following:
iterate (z,, ve) closely follows the sequence (A(9)) 20
1=0 Amax 2 LLL( v))) (q:) + O1)).
In Figure 1b, we conduct numerical experiments on single neuron model with tanh-activation, demonstrating that X(ge) provides a close approximation the sharpness.
C.2 Proof of Theorem C.1
We give the proof of Theorem C.I. restated below for the sake of readability.
Theorem C.1 (gradient flow regime). Let n E (0, 2()(1)+2) be a fixed step size and 6 be a sigmoidal
r(1)
function satisfying Assumptions 5.1 and C.I. Suppose that the initialization (po, qu) satisfies |pol S 1
39

---

and qo E (1-2)) the GD iterations (pt. 9:) converge the point (0, q") such that
r(3) An
Consider the reparameterized GD trajectory characterized in Eq. (7). Then.
2(go
1) ((1)
go S o S cxp
2n
min
40 S
240
Qo
40
Theorem C.1 directly follows from Proposition C.5 stated below.
r(1) Proposition C.5. Suppose that 17 (0, 2(r(1)+2) ). lpol S and qu E
1 r(1) 1-22n in
Then for any
t t 0, we have
2( 1) r(1)
[pll S 1 min
1,
4u
Qo
and
2(40 1) r(1) } }
S exp
2n min
S 2go
qu
qu
Proof.
We give the proof by induction; namely, if
[pel S 1 min 2(go 1) r(1)
}
2(go 1) r(1)
}
2o S qr S exp
2n
min
40 S 20o
Qo
go
are satisfied for time steps 0 S t S ls for some k. then the incqualities are also satisfied for the next time step & + 1.
For the base case, the inequalities are satisfied for by assumptions. For the induction step, we assume that the inequalities hold for any 0 s t < k. We will prove that the inequalities are also satisfied for t k
240- From Eq. (7), we get By induction assumptions, we have r(1) S r(px) < 1 and S 9b S
Pk++1
2r(px)
< max {1- 2r(1) -1
2
= min - 1), r())
Pk
2go
20
go
Due to the induction assumption, we obtain the desired bound |pt41 following:
{ 2(
1)
r(1)
k+1
|p*+1l < min
go
40
Moreover, for any < < k; by Eq. (7) we have
1 2mp? s (1 mp?
<
gt
= - 76-70(p)3)) )22 < 1,
a++1
where the second inequality comes from the fact that is 1-Lipschitz and x(0) = 0 (Assumption 5.
Hence, we have g 2 ql 2 qo- Note that _t_
[ 1] for small 17. Consequently, we have
ge+1
q0
qt
q
log
log
g6++1
90+1
9t+11
t=
1=0
< 2n
Pt
2
1=0
2(go
r(1)
2
1 min
qo
2(10
1) r(1)
< 2n
min
qu
qu
40

---

where the second inequality holds since [log(1 + 2)1 - desired bound on q as following:
2| if | S
Therefore, we obtain the
2(q0
1) r(1)
9o I 46+1 S cxp
27 min
40-
90
40
Since qu 2 1-2n and 9o I r(1)
we have
2(go- 1)
r(1)
min
2 417.
40
40
This implies that g++1 S exp(2) go I 2go. as desired.
C.3 Proof of Theorem C.2
In this subsection, we prove Theorem C.2. We use the following notation:
rpp)
All the lemmas in this subsection are stated in the context of Theorem C.2. The proof structure resem- bles that of Theorem 4.2. We informally summarize the lemmas used in the proof of Theorem C.2. Lemma C.6 proves that is bounded by constant and g increases monotonically with the increment bounded by (1)). Lemma C.7 states that in the early phase of training, there exists time step to where Stg becomes smaller or equal z-r(IJ+ which is smaller than 2. Lemma C.9 demonstrates that if is smaller than 2 and [pe| 2 f(1 - 4), then (r 1| decreases exponentially. For the case where exponential rate. Moreover, Lemma C.8 shows that if st < at some time step, then 9t+1 is upper bounded by 1 + ((r). Combining these findings, Proposition C.11 establishes that in the early phase of training, there exists time step to such that $t; = 1 + (n)) Lastly, Lemma C.12 demonstrates that if s = 1 + Os,o (1), then 88 - h((p)))) decreases exponentially.
[pel < f(1 (11-3) 2). Lemma C.10 proves that [Pel increases
Lemma C.6. Suppose that the initialization (po: 9o) sutisfies lpol S and go E (co, - 8). Then
for any / 2 0 such that qi 1, it holds that
Ipel s 4, and qi 9441 S (1+ (())))))
Proof. We prove by induction. We assume that for some 2 0, it holds that |P:| S 4 and 5 S q: S 1. We will prove that |pt+a| < 4 and 2 S St S a++1 S (1 + (1))). For the base case, [Po 1
and 1 Co < < 1 holds by the assumptions on the initialization. Now suppose that for some
t 2 0, it holds that |p:| < 4 and //2 S qe < 1. By Eq. (7).
Pet1l = Pt ( (pc)
$ max lpel. 2
}
< 4.
where we used Assumption 5.1 bound (p(p:) (p:) < 1. Moreover,
1 1 2n S (1 - 11)² 5 q: = (1 - ())())))))) < 1,
g++1
since |o| is bounded by 1. Hence, S 91+1 S + (1))). as desired.
Lemma C.6 implies that pt is bounded by a constant throughout the iterations, and qt monotonically increases slowly, where the increment for each step is O(z)). Hence, there exists a time step
T 11) = (2(111) such that for any t < T, it holds that 4 < 1 2" Through out this subsection, we focus on these T' early time steps. Note that for all 0 < qi E (co, {).
< T. it holds that
Lemma C.7. There exists time step to Os.o (1) such that 8t0 S 2-r(1))
2
41

---

Proof. We start by proving the following statement: for any 0 T 2
2-r(1) < St < 2r(1)=',
then $141 < 2r(1)=1 and |pt+1 I (1 (Pe) Suppose that 2 2-r(1) < 8 < 2r(11 Then from
Eq. (7), it holds that
P++1 1 St.
2
S 11 r(1)
Pt
Hence, |pc-1l & (1 (1- r(1)pp)) Now we prove st+1 < 2r(1)-1 Assume the contrary that
81411 2 2r(1) Then, r(p++1) Theorem, there exists pt E (1petal: (pel) such that
2+11
< g++1 < < - 2 so that [Pt-1] 2 f(1 - () By Mean Value
8442
(pt))
(Ipel - (pt+al)
r(P(-1) r( r(Pell (lpel - (p(+11))
r(p()
r(p))²
r' r((petz)) (r(1)lpdl)
$ r(p))
+
r(p+1))
1
In'((11 - -2)11
((()) (1-2) 2) )
r(pe)
(1 - ")²²
r(p()
where we used Assumption C.2 (i) and F(1 - S lpt+il S (1 - -r())))ll Consequently,
1 4++1 (1 + (())4 r(pc+1)
(oso(1)) < q: <t (r(1)
r(pt)
r(xe)
for small step size 17. This gives a contradiction to our assumption that 8141 2 2r(1)-1. Hence, we can conclude that sc+1 < 2r((1) as desired.
We proved that for any s $ T. 2 8t < 2r(1) holds that sc+1 < <2r(1)-1 -1 and
|pt+1| (1 - r(1))pp)) At initialization, [Pol S. 1 and go so that so < <r( If 50 S 2-r(1)
2
2 then ln = 0 is the desired time step. Suppose that 80 > 2-r(1) Then, we have 81 < 2r(1)= and
|pi| S (1 - (1-r(11))pp) < - r(1). Then we have either 81 S 2-1(IJ or 2ーン(T < 81 < 2r(1) In the previous case, to = is the desired time step. In the latter case, we can repeat the same argument and obtain $2 < 2r(1)-1 and |p2l - (1 r(1))² By inductively repeating the same argument, can obtain time step to S log(f(1 2)) log r()) 0s.0 (1) such that either Sto
2
2-r(1). or Therefore,
Pr.. S f(1 - - In the latter case, r(p.,) 2 1 - o/2 > noo and hence Sto <1<
2 2-r(ij
tu = O, (1) is the desired time step satisfying Sto 1 2-r(1)-
According Lemma C.7, there exists a time step to O (1) such that sto S 2-r(1) 2(1 for small step size 17. Now we prove the lemma below.
Lemma C.8. Suppose that st S 1. Then, it holds that s1+1 < 1 O(1).
Proof. For any p E (0)) we have r(p) 2 % so that 1Sa. (p) = (-1 2r(pr)
9t )p.. Hence,
fn. (p) = 2r(p) + pr (p(
1,
Tp
for any p E (0.f(%) )) By Assumption C.2 (ii), (iv) and q: 2 co 2 r()) + 2 2r(21) where (x)
2- -1, both r(p) and (1+ pr'(e)) ) are positive, decreasing function on (0, f(+)).
21 sup, 7(2)
-(p)
Consequently, o0 1f (p)1 is decreasing function on (())))
Now note that 2 gt < which means f(1) < f(m)) f( by the definition of f. Note that 0p Iq (p) at p f(g:) evaluates to
i f (f(4:)) 1+ (f())r'(F())))) r(f(ge)) > 1+
(F(co)) 2 0,
Op
r(f(co))
42

---

where the inequalities used Assumption C.2 (ii) and qi co 2 r(2)) where 20 sup,
(x) r(#)
2
- -2 from the statement of Theorem C.2.
Therefore, since Bp fg. (p) is decreasing on (0, (0)) and is nonnegative at f(g:), for any p E (0, f(q:)), holds that 2 3p f. (p) 2 0. In other words, 1f.. (p) is an increasing function on (0, f(g)).
Since 0 S I 1, we have Pel - f(qe) and it holds that
|p+1|| (-1
2 Ipe = 1f (pt) S [g.. (f())))) f(g.).
Therefore, with this inequality and Lemma C.6, we can conclude that
9++1
(1 + O(n))a
qt r(f(a:)) + = 1 +()))
81+1 r(pt+1) r(pt+1)
S
Using Lemma C.8, prove the following lemma.
Lemma C.9. For any t S T, if 8t 2 and r(pe) 1. then
|st+1 1 S (1 d) 1 + O(n).
where d E (0, [] is constant which depends on and db
Proof.
From Eq. (7) holds that
Pt+1 1 < 0.
2
Pt
St
so that Pt and p1+1 have opposite signs. By Mean Value Theorem, there exists 0g between -1 and (1 - 2) satisfying
1
1
r(pc-1)
(--po+
2(xx-1)
+
(((pp)
(2(sr - 1)
r(- r(-pe)
r(Orpe)2
Ir' (Oepe))
2( - 1)
lpel.
(17)
(p:))
r(0,p,)2
where the last equality used the fact that and dep: have opposite signs and (2(2) and have opposite signs. Note that [0.pe] is between [p.| and [pt+1]. Consequently, the value [r'(0-p)) between (p=))
r(0.p.)²
r(p))*
and Ir' (Pe+aa) 1( by Assumption C.2 (i). We will prove the current lemma based on Eq. (17). We divide into following three cases: (1) s, > and s1+1 2 1, (2) st and st
1, and (3) s < 1.
Case 1. Suppose that 2 I and 8241 2 1. Here, we have lp:| 2 f(g)) 2 r(1 ) and similarly |pe41| f(1 9). By Assumption C.2 (i), Ir' p(())p))) T(B...))
(F((11)))) +)) Hence, Eq. (17) gives
(1-
1
rr [f(( - -))I
(2(s:
1) 11)
1 5)
r(p:11) r(p))
(1 4
St
Consequently, by Lemma C.6,
8141 qc(1 + +0))
(f(11 - b))l
2((s 1)
7. (
8
O(r) S
4 + O())
r(p+41) r(pet1)
(1 8)²
8t
2
|r (f((1 - ())1
8
1
< 81
(8t 1)
1
+ O
(1 1/12
2
2
43

---

f(1 (f(r 9)1
$ 8
(st 1) O(n).
2(1 2)²
where we used qi co > and 8 Therefore, we can obtain the following inequality:
0 S 81+1 15 1 f(1 (F)r - 9011 (s--1)+
2(1 - -1-
Case 2.
Suppose that 8; 2 and 8141 1. Here, we have r(P++1) > 9++1 2 qt 2 r(p:), so that (Orpe)) [r((x)] by Assumption C.2 (i). Hence, deduce from I r(pr)²
|Pe41| [pel- Consequently, r(0,pe)²
Eq. (17) that
2 1
|r((p]]
2((s (():-11)
Ipel
r(pc+1)
r(p))
r(pt)²
2/por (pt)
(8t - 1)
r(p.)
r(pe)gt
r(pp) + 2prr'
(st 1).
r(p))q:
Consequently, by Assumption C.2 (ii),
9: r(pt+1)
2 St + 2prr' (pt)
8111 2
(st - 2) +
(s: 1),
r(p:)
r(
where we used r(p)) 2 z
m 2
Therefore, we can obtain the following inequality:
0 0 8141 S - (1 1+ 4F (
(st -1),
-1 F( (F(
2F (F)
since co 2 r(21) + 2 2r(2)) with 21 sup& (=) 2 -1], and r(21) 1 holds by Assumption C.2 (iv).
where
4o
T(2)
Case 3. Suppose that st < 1. By Lemma C.8, it holds that 8141 I 1 O()). Moreover, we assumed r(pr) $ 1. so that |px| 2 f(1 - 4). We also have
|petll = (11 2 2) 1p: 1p/20 1
11 5)
st
(F(111 +I)
Consequently, by Assumption C.2 (i), it holds that
Hence, by Eq. (17), we
r(0ep.)2
(1- 4)
have 1
|r(( fI)l
(2 (2(1- (2(1-8)) 8t)
8
r(p(+1)
r(pe)
(1 - 8)2
8t
|r'(f(1 f))g
2(1 - 8t) (1- 8
r(pt)
+
(1 - $)2
4
2²(1 (F) - 11)1
(1 8t). (1)8
+
r(p))
(1- 4)2
and hence,
f(1 8 ( 91I
St-1
1(p++1)
gt
St +
(1 - 3)2
(1 st):
where we used gt > 2. Therefore, we can obtain the following inequality:
f(1
(F(1 { ) (1 st).
-O(n) < s141 s
(1-
44

---

where the first inequality is from Lemma C.8.
Combining the three cases, we can finally conclude that if we choose
1 f(1 3) (f(1 -11 .2 (1+ 2r( ()())) 2(1 2)
f(1 (flr - f))
1 0,
d = min 2
E
(1 4
2
then 81+1 1| - (1 - d){s: 8: - 11| O11), as desired.
Lemma C.9 implics that if st < 2 and |pe| 2 f(1
2). then |s 1| exponentially decreases. We
prove Lemma C.10 to handle the regime Ipdl < f(1 ) which is stated below.
Lemma C.10. For any 0 S t T. (p:)) 2 1
o holds that
Pc+1 1
2
Pt
2
Proof. If r(p)) 2 1 f. then St
91
1- +
4-26 4-6
where we used qe < & 2
for any
r(pi)
I-
0 t S T. Consequently,
Pc+1
2
2(4 - 8)
2
12
1
1 =
Pt
8t
4 25
2 8
Now prove Proposition C.11, which proves that 8t reaches close to with error bound of O()).
Proposition C.11. There exists a time step 1. = (log(11))))))) satisfying Stg 1 + (1,)
(18)
Proof. By Lemma C.7, there exists time step 4g = 00,0 (1) such that 8to S 2-r(1) 2 Here, we divide into two possible cases: (1) 8to 1, and (2) 8to 5 2-1(1]-
2
Case 1. Suppose that Sto < 1. By Lemma C.10, if ) 2 1 8 (or equivalentty, IPto S f(1 4) then there exists time step t S to + log( / log( (233) O0.o (1) such that [p., f(1 - f
Pro
We denote the first time step satisfying lpt, 2 f(1 i and f 2 to by 1 Os (1). By Lemma C.8.
it holds that St, I 1 1 O(1) since st,-1 < 1. Consequently, if st, 2 1, then 1 - 1| I O(n) so that f f the desired time step. Hence, it suffices consider the case when St Here, we can apply Lemma C.9 which implies that
|st+1 1/ S (1 -d]ss, 1 + O(n).
where d is a constant which depends on 8 and 6. Then, there are two possible cases: either 1s,, - 1| < O (nd 1), or |st++1 - < (1 |) 1/. It suffices to consider the latter case, suppose that [8t,+1
< (1 1s, 1/. Since we are considering the case 8t, < 1, again by Lemma C.8, we
1, we have |Pt++1l | 2 |pr,] 2 F(1-4) This means C.9 and repeat the analogous argument. Hence, there exists a time step 1g < f1 + log I-80, )/ log 4 = Os.o (log(n ))), such that [sr, - 1 < Ond = (s,)
have +1 - (n) Since Pa st,
Pr.
that we can again apply Lemma n
Case 2. Suppose that 1 S 8to < 2-1(1) Then, r(pt)) S (to S 5. so we can apply Lemma C.9. There are possible cases: either |st,+1 1 < Ond ) Oz,((n). or stot1 1| (1 4) 3)88 1/. It suffices consider the latter case. If St,+1 > 1, we can again apply Lemma C.9 and repeat the analogous argument. Hence, we can obtain a time step to < to tog( 1-er0 )/ log(1 )) =
Os,s (Ro(()) )) such that either $ig < 1 or |sg 1| (5,((1) satisfied. If st 1, we proved in Case that there exists a time step 1/ lo + (log(( such that |81; - 1| S 2x.o (1), and this is the desired bound.
45

---

Now we carefully handle the enror term (n) obtained in Proposition C.11 and provide tighter bound on 8, by proving Lemma C.12 stated below.
Lemma C.12. If (t 11 - (8.(()). then it holds that
18041 1 - M(p+1)ml = (1 2per'
15--1- 11-1-((pe)m) + Os, (1)r))))))
r(pt)
where
r(p) pr(())
ifp + 0, and
h(p) =
i = 0.
Proof. Suppose that |s+ - 11 2 (1. Then, |p++1l = 1 2
|pll = (1+ (8,0())) By
Eq. (17) proved in Lemma C.9, there exists e 0s.0(1) such that
1
((1++e)p))
2(ss -1)
r(p++1)
r(pt)
r((1 + et)pt)2
Pt
(pt))
2((8t [2($t-1]
1) -1)
r(p))
+
r((,)2
+ +Oāp(1pP)
Pt
(pe)
(2(s: 1)
+ 1n (1'pp).
r(p))
r(pt)2
8t
r'(p) with the fact that d (r'(p) is bounded on [-4, and dp r(p)2
where we used the Taylor expansion on r(n)² that Ipel I to obtain the second equality. Note that g++1 (1- (p: )²) -2q: (1+2no(p:)) ?)q (nn))) by Eq (7), Consequently.
r'(p))
(2(s) 1))
8141 (1 + 2n((p:)²)
peg: +06,6(1ipp)
81 r(pe)2
St
2p,r' (pt)
(1 + (n)(p:
(st 1) Oz (1²pp)
r(pt)
=1 1 2p (pt)
(st - 2no(p:)² + (T,
r(p))
Note that h is even, and twice continuously differentiable function by Lemma C.13. Consequently. 10()) = 0 and (p) = (O),), since h" is bounded on closed interval. Consequently, h(p+11) = h((1+ + Os,o )p:) = h1(p)) + Os.0 (np?). Hence, we can obtain the following:
S++1 - h(p(+1)m 91+1 h(p:) + Os (1)xp)
8111 1+ p(p))
n + (1p(11²)))
per (pt)
(prr
81 1 p(p)) ((pp))(())),)) + (1,0(1227)
r(pr)
per (pe) (p))
2ptr' (pt)
|+
(0t hhpp))(p))))) )n) 0r,(7pp??
r(pr)
Note that r(p:) = (1+ Ol,()i)))) 2 (1 +1+00 (8,0(1))40 2 c0 2 r r(z)) for small step size 7, where 20 sup 2r( 2 1] Consequently, it holds that
2per' > 0. Therefore, we can conclude r(p.)
r(2)
that
< (1+
pper'(pe) r(pt)
18t41 - - 1M
|st --(p)pp) + Os.o (1)
as desired.
Now give the proof of Theorem C.2, restated below for the sake of readability.
46

---

Theorem C.2 (EoS regime, Phase 1). Let 1) > O be small enough constant and & be an ac- tivation function satisfying Assumptions 5.1, C.I, and C.2. Let 20 sup= 2r'(s)
r(z) z -22
2 -1}. and co = max(r(zo), r(ai) ( + }} E (Z, 1). Let â E (0, - co) be any given constant. Suppose that the initialization (pu. 40) satisfies lpol I and qo E (co. 8). Consider the reparameterized GD trajectory characterized in Eq. (7). We assume that for all t 2 such that qt < 1, we have Pt O. Then, there exists time step La Os.o (log(nn))))))) such that for any 2 la.
(x) 21 = sup, T(=]
4:
= 1 +
r(p.)
where h : IR + R is a function defined as
o(p)rr(p) (r)
if * 0, and
h(p) =
if = 0.
(O)
Proof Theorem C.2. By Proposition C.11, there exists a time step ty- Os.e(log(( which
satisfies:
4: r(p())
1 = (1).
By Lemma C.12, there exists constant 1) > 0 which depends on ó, such that if 1st 11- Os,((1). then
18141 1 h(pe+1)ml S 1 2p (pt)
1s: 1- -h(p)))) + Dr
(19)
r(pe)
Hence, [st 14--10 1| = (n)) and 1st h(pp)) )ml 2 (- (-perlued ) Drj², then
si41 1 h(p(+1)ml 1 (pe)
1st 1 hhpp)p)
(20)
r(pe)
For any & S T, we have q: < 1 11 so that if |sc- (s,( then & (1+0s, < =1- for small step size 17. From Eq. (20) with = ty have either
h(p(;)ml < ( P(ar(Piz)) Dn²,
180g
r((p:))
or
f(1 f) (f(1 ())
s11+1 - 1 - hppp:+1)) 1+
15:2
-h(pe; )nl,
(1 -
where we used Assumption C.2 (ii and IPll > f(1 f In the latter case, |s 1| Os continues to hold and we can again use Eq. (20) with arguments, we can obtain the time step
= to + 1. By repeating the analogous
Dry² log (rr(()))s -1 1- hipp Jol)
La 5 to +
Oz.e (log(n-)))
11 F(1 (F 11 f))
log
(1-3)
which satisfies: either
Pr. r(Pta)
18t, - 1- h((pt,)))
Dry².
r' (pta))
or
Pr, Pe, r(ts))
Ar(4)
18t. - 1((p.. )l 1 ((-p'())) (0)
Dip <
Dry S
(4)) (
Dn²
r' (pt,)
where we used Ipel S 4 from Lemma C.6 and
2r(z) [(2)
for any by Assumption C.2 (iii). p
47

---

By Eq. (19), if [st 1- h(pp)ml - (
1r(1) ((4))
Dr1² is satisfied for any time step t, then
8111 - -----1-hppp h(ppt))ml & (1+ 2per' 2per
(pe)
(-
4r(4)
Dn + Dnipi - (- ()-(())) Dn¹
4r(4)
r(ps)
by |pel S 4 from Lemma C.6 and Assumption C.2 (iii).
Hence, by induction, we have the desired bound following: for any 2 la.
(1 (-4r(4) Ar(44)
81 1- h(pp))) -
Dij² O,,((1)),
r'(4)
by |pr| < 4 and Assumption C.2 (iii).
C.4 Proof of Theorem C.3
In this subsection, we prove Theorem C.3. We start by proving Lemma C.13 which provides a useful property defined in Theorem C.2.
Lemma C.13. Consider the function h defined Theorem C.2, given by
d(p))((p) (r)(n)
i 0, and
h(p)
=
(O)
Then, h is positive, even, and bounded twice continuously disferentiable function.
Proof. By Assumption C.1. is positive, even function. Moreover, h is continuous since limp-so h(p) [ = h(0). Continuous function on a compact domain is bounded, so
(- r(p) pr' (p)(p) o0). Since h is even, h is
bounded on the closed interval [-1,1], Note that <(p) < 1, and
positive, decreasing
function p O by Assumption C.2 (ii). Hence, h is bounded bounded on (-oo, 1]. Therefore, h is a bounded on R.
We finally prove that h is twice continuously differentiable. Since is even and C on I, we can check that for any p + O,
H' (p(p) = 2r(p)² r(()
p()) (1)))()(((())- - -r)
+
12r2
p
and ('(p) = 0. Moreover, for any # O,
h" h((p) 6r(p) 2x(p) <(p)2r"(p)
((p)2 (())2r(p)r(0))
p² pp'
pr (p)2
(r( (p)" + pr (p)
((p)2 (2r
(p) (2)
pr (p)) (1))()))""))(((()) 2pr (p))
(p)
p2r
and
(0 -1 20 (0)
(0)
(0)
3r" (0)
Since 1im h( (0), conclude that h is twice continuously differentiable function.
L
Now we give the proof of Theorem C.3, restated below for the sake readability.
Theorem C.3 (EoS regime, Phase II). Under the same settings as in Theorem C.2, there exists a time step 4, = 2(( 1)), such that yt, S 1 and qi >> for any I. > l5 Moreover, the GD iterates
(pe, qt) converge to the point (0, q²) such that
9² = (0) + (17.6117).
48

---

Proof. We first prove that there exists a time step t 2 0 such that qi > 1. Assume the contrary that q: S 1 for all t 0. Let ta be the time step obtained in Theorem C.2. Then for any 2 tas we have
r(pe) = -1-11pppp h(p))n O8 (s 1 h(p:)n
2
for small step size 1). The function g(p) == r(p) 1+ k(pin is even, continuous, and has the function 2
value g(0) (0) > Consequently, there exists positive constant e such that g(p) > for all p E (-e,e) e). Then, we have |pt| 2 for all las since g(p)) < 0. This implies that for any
t 2 las it holds that
g++ = (1 - n)((p)))) (1 m)((t) ()) 1,
qt
so 9: grows exponentially, which results in the existence of a time step t/ 2 ta such that qt, > 1, a contradiction.
Therefore, there exists a time step i, such that q, S 1 and @: for any > to, i.e. qt jumps across the value 1. This holds since the sequence (q:) is monotonically increasing. For any / IS la. we have 4++1 S q: + ((n) by Lemma C.6, and this implies that 4, 2 52((1 20))) as desired.
Lastly, we prove the convergence of GD iterates (pt: Let t to be given. Then, q: 2 96+11 > 1 and it holds that
Pc+1 Pt
2r(p:)
2
1 1
96,+1
Hence, |Pll is exponentially decreasing for & t.. Therefore, Pt converges to 0 as f + cc. Since the sequence (9:){= is monotonically increasing and bounded (due to Theorem C.2, it converges. Suppose that (Pr, Qt) converges to the point (0, q²) By Theorem C.2, we can conclude that
9"
1
17 (0())
Os,e(17,,
which is the desired bound.
C.5 Proof of Theorem C.4
In this subsection, prove Theorem C.4. We first prove important bound the sharpness value provided by the Proposition C.14 stated below.
Proposition C.14. For any (z, v) E IR? with r(x) + xr' (x) 2 0, the following inequality holds:
4z2r(z)
(r(x) + xr (x))x² & Amax (v ((x - (r(x) + xr' (x))2² + r(xx) - xrr(
Proof. Let (ar, y) IR.2 be given. The loss Hessian (x. y) is
v²L(x,y) (+()))"()) (z) 2 (z)y
20(x)d(())) o(x)²
Note that the largest absolute value of the eigenvalue of symmetric matrix equals to its spectral
norm. Since trace of the Hessian, tr( (((,)((x, u)) = (ø(x))"(z) + o' d'(x)))²² + o(x)² = (r(x)- - xr'()) )xx)2 > 0, is non non-negative, the spectral norm of the Hessian L y) equals to its largest eigenvalue. Hence, we have
Amaxx (V2((x.
1
2 v2 V> v)
= [(({(x)) + (z) + (2))))
2 (o(x)0" (x) +
= -r(xx) + zr (z))))
49

---

which is the desired lower bound. We also note that for any matrix A, the inequality |A S ||A||e holds, where || F is the Frobenius norm. Hence, we have
Amax (2L(x,)))
|| [zx, y)1
- ||
-(xz)
+or()) + (o) 80' +(x)")
((r60)+ or r(z) r(z)+
4r²r(z)
4x²r(x)
(r(x) + or (z))r²
r(æ) + +r(
which is the desired upper bound.
Now we give the proof of Theorem C.4, restated below for the sake of readability.
Theorem C.4 (progressive sharpening). Under the same setting in Theorem C.2, let to denote the obtained time step. Define the function À R250 - R. given by
1 ((())
3 if S and
A(q) =
otherwise.
Then, the sequence (X(+))
is monotonically increasing. For any t 2 tax the sharpness at GD satisfying the following:
iterate (ze, closely follows the sequence (A(9)) 2 == V L((x,, m:)) (1:) + Oz()).
Amax
Proof. By Theorem C.2 and since is a bounded function by Lemma C.13, we have 5, = 1+ (p()) for any 2 ta. Note that rPt)) (1 (r))gg 2 (1 + (1)))00 co r(zo) for small step
size 7,, where 2g = sup 2z((z)
r(:) 2 } Consequently, it holds that + (par 2 0. and this implies r(p:)
1 (pe) r(m)) 2 By Proposition C.14, bound the sharpness d = Ammax ( 2 LPrr, ve)) by
(1+
(pr)( r(pt)
2r(pe) S s (1+
Per'(pu) r(pe)
2r(p:),
+ +p2 [1 Per' per'(p))) -1
1/tt
r(pt)
By Lemma C.6. [Pll for any time step Moreover, since 1 + per((pe) r(pe)
2 1 holds for any to
with small step size we have 4pp (1 + per'(pe) = =O((). Hence, for any 1 2 la it holds that
A = 1+ Per' (pr)
2r(pt) + O1).
(21)
r(pt)
1/4
For any t 2 ta we have = 1+0 and this implies -- =11=0000) and (r(p:) Os(n).
For any 0 < q < 1, we have
d ((A))(((()
F(()(r'(F()) + +f(g)"
f(a)r'(F(g))
dq
F(())(((()))
((g)r(((()))
r(F(g))
q²
so that
( d
(f())(((())
(1 pr) 2.
lim
lim
dg
9
p00
r'(p)
50

---

Therefore, dg ((())r(((e)) is bounded on [2, 1) and Taylor's theorem gives per' ((a))r(((((e) Qt for any time step with qt < Hence, if gt
I (+((r(p)) qr)) S Oa(i).
r(ps)
< 1, we have the following bound:
A((:)= (1+ par (ve)
por'(ee)
2r(p))
f(y f(y:)r (F((a))
par'(pt)
r(pe)
179
82
r(pe)
+
(22)
where we used por'(ps) per((p)) (1 + (n()) per'(pe) Od()). since (per S r(p:) for r(pi)
4t
r(pi)
any & 2 ta- Now let t be any given time step with 2 1. Then r(p:) = - ( and since
r(2) = -1++" + Ox(zt) for small 2, we have [pel = OolvT) Hence,
A((:)
(1+ (ka))-(1))
(pe)
2r(pt)
Per (pt) 2
< 1 1-
+ = (23)
r(p+)
r(pt)
"
for any t with q: 1, where we again used prr'( Eqs. (21), (22), and (23), we can conclude that for any 1 2 la we have
per'(pe) (1 + (r)r) per((p:) x(pe)
Os(1). By
9a
r(pt)
A {(q.) s s0((),
the desired bound.
Finally, we can easily check that the sequence (A( is monotonically increasing, since =x) is decreasing function by Assumption C.2 (ii) and the sequence (gt) is monotonically r(2)
increasing.
51


Reviews:
[{"review_id": "K6EdRnXlW5", "reviewer": "Reviewer_mAcJ", "summary": "This paper studies the edge of stability phenomenon and proposes a \"trajectory alignment\" phenomenon, which states that with different learning rates and different initializations, if one plots the residuals and the inverse squared gradient norms, these lie on a fixed curve. The paper supports this theory with experiments on small networks with log-cosh loss and rigorously proves it for a two-layer linear net and a single neuron nonlinear net, both trained on a single datapoint.", "questions": "- Why did the authors chose to parameterize $q$ in terms of the inverse squared gradient norm instead of simply the squared gradient norm? At first glance, the reciprocal $1/q$ seems like an easier object to work with?\n- How is the analysis in section 5 for square loss different from the analyses in the rest of the paper?", "limitations": "The authors have adequately addressed the limitations of this work.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- While the bifurcation analysis of GD at EoS is not new, this paper takes this analysis further than previous papers and rigorously characterizes the limiting curve for toy models trained on a single datapoint.\n\n- The paper also provides fine-grained analyses for each of the toy cases it considers including precise predictions for the limiting sharpness which improves over previous work.", "weaknesses": "- The abstract and the introduction state that GD trajectories (after a proper reparameterization) align on a specific bifurcation diagram. However, the (p,q) variables in the paper are only a reparameterization in the toy case when $\\theta$ has dimension 2. In the general case (e.g. in section 3) it is not fair to call this a reparameterization as the variables (p,q) do not determine the GD trajectory.\n\n- The abstract claims that the paper establishes the trajectory alignment phenomenon through empirical studies. However, the experiments in the paper are extremely limited. The experiments in figure 3 use a single datapoint equal to $e_1 \\in \\mathbb{R}^{10}$ and the non-standard log-cosh loss function, while the experiments in figure 4 use 10 datapoints sampled from a Gaussian distribution in $\\mathbb{R}^{10}$. While these settings are certainly more difficult to analyze than the settings in the rest of the paper, they are still very much toy settings. It would be interesting to see if similar phenomena occur on more realistic datasets/models with more realistic loss functions (e.g. square loss or cross entropy loss).\n\n- The paper states a couple of times that the bifurcation diagram is determined solely by the loss function, independent of the network architecture, training data, and step size. However, this strong claim is not demonstrated outside of the toy setting of a single datapoint. Figure 4 does suggest that in the general case it may be independent of the learning rate but there is no evidence for the full claim.\n\n- It seems that all of the results in terms of the \"canonical parameterization\" in eq (2) are specific to assumption 2.4. In particular, the results in section 5 for square loss require defining a different parameterization (definition 5.2) which challenges the generality of the results in the rest of the paper. As the original EoS phenomenon as described in Cohen et al. was observed for square/logistic/cross entropy losses, it is unclear whether the phenomenon being studied is the same as the phenomenon observed in Cohen et al.\n\nMinor Comments:\n- There seems to be a clear separation of colors in Figures 2, 3, 4 but as far as I can tell these colors are never defined. Also, increasing the size of the datapoints might help with visibility, especially in Figure 3.\n- When discussing previous EoS papers with bifurcation diagrams, the paper should cite Chen et al. \"Beyond the Edge of Stability via Two-step Gradient Updates\" (it is in the list of references but not in the related work)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your recognition and feedback on our work. We've integrated your suggestions to enhance our manuscript. For additional experiments and figures, please refer to the attached PDF files in the \"global\" response, which we reference as \"Figure (number) - Additional Experiments.\"\n\nWe are addressing your specific comments as follows:\n\n>(W1) reparameterization do not determine the GD trajectory\n\nWe understand your concern that the reparameterization may not fully determine the GD trajectory. Our use of the the term \"reparameterization\" was intended to encompass not only bijections between coordinates but also surjections, such as projections to a subspace. We will clarify the term \"reparameterization\" in the revised manuscript. If you have suggestions for alternative terms that better represent this concept, we would appreciate your input and consider it for the revision.\n\n>(W2) realistic datasets/models with realistic loss functions\n\nWe have indeed observed the trajectory alignment phenomenon in more realistic settings. In our experiments, we trained large width FC networks on CIFAR-10 and observed trajectory alignment, as illustrated in Appendix A.4 and Figures 13-16. We further conducted additional experiments on CNNs trained on CIFAR-10, as shown in Figure 1 - Additional Experiments. The results consistently demonstrate the trajectory alignment phenomenon when training wide CNNs on a real-world dataset.\n\nRegarding the loss functions, extending the results to squared loss is an important future direction. Our work in Section 5 provides an initial step towards understanding trajectory alignment that arises on nonlinear networks trained with squared loss.\n\nHowever, we want to clarify that cross-entropy loss in fact falls outside the scope of our work's focus. Under cross-entropy loss and GD with a large step size, the sharpness typically decreases towards the end of training, going far below 2/(step size). This phenomenon has been empirically observed and theoretically investigated by [Cohen2021] and [Wu2023]. In contrast, our work concentrates on scenarios where sharpness saturates near 2/(step size), and the trajectory alignment phenomenon occurs in this setting.\n\n>(W3) The claim that bifurcation diagram is independent of the network, data, and step size is too strong.\n\nWe acknowledge the reviewer's concern on our initial claim. We will revise our claim to clarify that different GD trajectories (after a proper reparameterization) align on a specific bifurcation diagram __independent of initialization__.\n\n>(W4) It seems that all of the results in terms of the \"canonical parameterization\" ...  observed in [Cohen2021].\n\nIndeed, the canonical reparameterization applies to a class of convex Lipschitz losses, hence excluding square/logistic/cross-entropy losses. Extending our finding to squared loss is an important future direction. However, cross-entropy loss falls outside the scope of our interest, as explained in our response to (W2). \n\nFurthermore, we want to clarify that our primary focus is to study the EoS phenomenon observed by [Cohen2021]. This phenomenon refers to the oscillating behavior exhibited by GD with a large step size, causing the sharpness to saturate near 2/(step size). Our work is centered around characterizing this phenomenon, as the trajectory alignment phenomenon captures the oscillating dynamics and the sharpness approaching 2/(step size).\n\n>(W5) colors in figures\n\nThe distinct colors used in the figures indicate independent runs of GD with varying initializations. In the revised manuscript, we will enhance the figure captions to provide a explanation of the color scheme and increase the size of the markers in our plots for better visibility.\n\n>(W6) The paper should cite [Chen2023]\n\nWe appreciate your suggestion regarding [Chen2023]. We will include discussions on their study in the related works section of our revised manuscript. Specifically, [Chen2023] observed that in specific settings, GD with a large step size converges to its period-2 orbit, which is a fixed point of a two-step update. In contrast, both [Zhu2023] and our work focus on the phenomenon where GD exhibits oscillatory behavior, eventually converging to a fixed point of the one-step update. Indeed, [Chen2023] studied the EoS phenomenon but did not capture that GD trajectories exhibit bifurcation-like oscillation. We will ensure that these differences are properly discussed in the revised manuscript.\n\n>(Q1) Why did the authors chose $q$, but not $1/q$?\n\nWe chose the inverse squared gradient norm for $q$ mainly for better visualization. This choice allows us to represent the trajectory alignment phenomenon more clearly within the range of $q\\in (0,1)$. Although using the reciprocal $1/q$ is an option, working directly with $q$ offered clearer and more interpretable visualizations for our experiments and analysis.\n\n>(Q2) How is the analysis in section 5 for square loss different from the rest?\n\nThe theoretical analysis in Section 5 for squared loss is similar to the analysis in Section 4. This similarity arises from the resemblance between Eq.(7) and Eq.(6).\n\nHowever, there are some differences in the specifics. For example, Theorem C.3 proves that the limiting sharpness is $\\frac{2}{\\eta}-\\frac{2}{\\lvert r''(0) \\rvert} + O(\\eta)$, while Theorem 4.3 proves that the limiting sharpness is $\\frac{2}{\\eta} - \\frac{\\eta}{\\lvert r''(0) \\rvert} + O(\\eta^3)$. The error term is $O(\\eta)$ in Theorem C.3 and $O(\\eta^3)$ in Theorem 4.3. This difference arises because in Eq. (7), the one-step update of $q$ is $q_{t+1} = (1+O(\\eta))q_t$, while in Eq. (6), it is $q_{t+1} = (1+O(\\eta^2)) q_t$. Nonetheless, the main ideas of the theoretical analysis used in both sections remain the same.\n\n___\n__Reference__\n\n[Chen2023] https://openreview.net/forum?id=AvwlrX9AQr\n\n[Cohen2021] https://openreview.net/forum?id=jh-rTtvkGeM\n\n[Wu2023] https://arxiv.org/abs/2305.11788\n\n[Zhu2023] https://openreview.net/forum?id=p7EagBsMAEO"}, {"review_id": "26gc4Czx3D", "reviewer": "Reviewer_ft7H", "summary": "This paper investigates a trajectory alignment phenomenon during the edge of stability (EoS) regime in gradient descent (GD): the trajectories from different GD dynamics tend to align on a bifurcation diagram. The paper first empirically observes the trajectory alignment in deep neural network training in a canonically reparametrized space that can reveal the bifurcation. With a problem setting of a two-layer linear neural network and single neuron nonlinear network training on a single data point, the alignment is theoretically analyzed, explaining the GD trajectories in the gradient flow regime, two phases of the EoS regime, and the progressive sharpening.", "questions": "In addition to the questions posed in the **Weaknesses** section, I have the following questions:\n- Is any further analysis possible for the period-4 (or more) oscillation shown in Figure 1 (b) and (c)?\n- Can the EoS for GD with momentum also be explained by this theory?\n- How do the experimental conditions differ for each color in Figure 2?\n- In Figure 5 of the appendix, why there seems to be more alignment to the bifurcation diagram for the mean and $L_\\infty$-norm cases than the $L_1$, $L_2$-norm cases?\n- What happens if $q_0 < \\frac{1}{2}$ in line 221?\n- How is the sharpness anticipation for progressive sharpening in **Theorem 4.4** accurate for neural networks larger than those in Figure 1 (a)?\n- Can the progressive sharpening also be interpreted in the bifurcation diagram?\n", "limitations": "Some limitations of the current results are discussed in **Conclusion**. For other possible limitations, please refer to the weaknesses and questions above.", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- Based on simple models, this paper clearly characterizes the gradient flow regime, two phases of the EoS regime, and progressive sharpening using the bifurcation theory. This seems a novel and original way to analyze the phenomenon. The introduction of canonical reparametrization greatly helps capture the crux of the EoS phenomenon. This work will provide useful insight to study the dynamics in the EoS regime of neural network training, which is recently getting a huge interest in the community.\n- This paper is in general very well written. The theories seem to be developed in a mathematically sound manner.\n", "weaknesses": "- The claim that the trajectory alignment on a bifurcation diagram is independent of network architecture and training data seems somewhat excessive. It does not corroborate well with some experimental results in the main text and appendix (see Figures 5, 10, 13, 14, 15), which shows that the degree of alignment depends greatly on the width size and the number of data. Also, the analyses are only performed for specific network architecture and loss function combinations.\n- The theory does not cover some widely used losses, e.g., the cross-entropy loss. Is the loss function necessary to be symmetric for this alignment to occur?\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for recognizing our work and providing valuable feedback. We've carefully incorporated your suggestions to improve our manuscript. Additional experiments and figures can be found in the attached PDF files in  the \"global\" response, referenced as \"Figure (number) - Additional Experiments.\"\n\nHere's how we've addressed your specific comments:\n\n>(W1) The claim that the trajectory alignment on a bifurcation diagram is independent of network architecture and training data seems excessive.\n\nWe acknowledge that our initial claim was excessive. We will revise our claim to clarify that reparameterized GD trajectories align on a specific bifurcation diagram __independent of initialization__, when the EoS phenomenon occurs.\n\n>(W2) The theory does not cover cross-entropy loss. Is the loss function necessary to be symmetric?\n\nWe want to clarify that cross-entropy loss falls outside the scope of our work's main focus. Under cross-entropy loss and GD with a large step size, the sharpness typically decreases towards the end of training, going far below 2/(step size). This phenomenon has been empirically observed and theoretically investigated by [Cohen2021] and [Wu2023]. In contrast, our work concentrates on scenarios where sharpness saturates near 2/(step size), and the trajectory alignment phenomenon occurs in this setting.\n\nThe symmetry assumption on loss is unnecessary for the trajectory alignment phenomenon. While we assume a symmetric loss in our theoretical analysis for simplicity, the phenomenon can also arise in an asymmetric setting. As demonstrated in Figure 1(c), we still observe trajectory alignment on an asymmetric bifurcation diagram when using ELU activation. \n\n>(Q1) period-4 (or more) oscillation shown in Figure 1 (b) and (c)\n\nAnalyzing higher-order oscillations goes beyond our current scope. We want to clarify that we mainly focus on convex Lipschitz losses, where period-4 (or more) oscillations do not occur as discussed in Line 154. The occurrence of higher-order oscillations depends on the specific model and remains an interesting future direction.\n\n>(Q2) EoS for GD with momentum\n\nExplaining the EoS for GD with momentum using our current theory is beyond the scope of our work. Analyzing the dynamics of GD with momentum is technically challenging, and to the best of our knowledge, there is no existing work that theoretically investigates EoS for this variant of GD.\n\nWe conducted additional experiments using a toy model $(x,y)\\mapsto \\log(\\cosh(xy))$ to examine the trajectories of GD with both Polyak and Nesterov momentum (plots are not provided due to the space constraint). We observe that the behavior of GD with momentum is qualitatively different from classical GD. For example, GD with Polyak momentum under different initializations does not necessarily converge to the minimum of same sharpness, but we do observe some level of consistency in Nesterov. These observations highlight the need for further investigation to understand GD with momentum in the EoS regime.\n\n>(Q3) color used in Figure 2\n\nIn the figure, each color represents an independent run of GD with different initializations. In the revised manuscript, we will enhance the figure captions to provide a clear explanation of the color used.\n\n>(Q4) Why there seems to be more alignment for the mean and $L_\\infty$-norm cases than the $L_1$,$L_2$-norm cases?\n\nIn Figure 5 of Appendix, the trajectory alignment phenomenon appears more pronounced in the mean plot than the $L_\\infty$ plot. However, the underlying reasons remain currently unclear. In our experiments, we explored different reparameterization options, including mean, $L_1$, $L_2$, and $L_\\infty$ norms. Notably, the consistent occurrence of trajectory alignment in the mean plot across various settings is depicted in Figures 11-16 of Appendix. Understanding this phenomenon would be an intriguing future work.\n\n>(Q5) What happens if q_0 < 1/2 in line 221?\n\nIn the case where 0 < q_0 < 1/2, the sequence (q_t) may not exhibit a monotonic increase. However, empirical observations indicate that q_t increase over the long run and eventually enters a regime where q_t > 1/2. \n\nWe conducted additional experiments to further investigate the behavior of GD trajectories initialized with q_0 < 1/2 under log-cosh loss. The results, presented in Figure 2 - Additional Experiments, demonstrate that in most cases, these trajectories eventually reach the set $S_A:=[-1,1]\\times (\\frac{3}{4}, 1)$. Note that the assumptions in Theorem 4.2 is satisfied when $(p_0, q_0)\\in S_A$. This indicates that the trajectory alignment phenomenon still occurs (in the long run) when q_0 < 1/2.\n\n>(Q6) Is the sharpness anticipation in Theorem 4.4 accurate for large networks?\n\nYes, the sharpness anticipation is accurate for large networks. We conducted additional experiments on two-layer linear networks with widths of 256 and 1024. The results, illustrated in Figure 3 - Additional Experiments, provide strong evidence that the sharpness prediction derived from Theorem 4.4 effectively approximates the actual sharpness values.\n\n>(Q7) Can progressive sharpening be interpreted in the bifurcation diagram?\n\nCertainly, the progressive sharpening phenomenon can indeed be interpreted within the context of the bifurcation diagram. The relationship between progressive sharpening and the bifurcation diagram is established through our theorems: Theorem 4.4 establishes that as the GD trajectory closely follows the bifurcation diagram, the sharpness increases during the training process. In particular, $\\tilde{\\lambda}(q)$ represents the sharpness value on the bifurcation diagram. Consequently, as the sequence $(\\tilde{\\lambda}(q_t))_{t=0}^\\infty$ increases, the sharpness also increases since $\\tilde{\\lambda}(q_t)$ well-approximates the actual sharpness when the trajectory alignment phenomenon occurs.\n\n---\n__Reference__\n\n[Cohen2021] https://openreview.net/forum?id=jh-rTtvkGeM\n\n[Wu2023] https://arxiv.org/abs/2305.11788"}, {"review_id": "Qo7dUE5pF0", "reviewer": "Reviewer_6YQZ", "summary": "This paper studies the Edge of Stability phenomenon in training neural networks. The authors demonstrate that, on some reparameterization of the loss function, for a small enough stepsize, the trajectories of gradient descent with different initializations align on a bifurcation diagram. They show this for two classes of functions:\n(1) Two-layer fully connected linear neural networks trained on a single data point with a convex and Lipschitz loss.\n(2) Single-neuron nonlinear network trained on a single data point with the squared loss.", "questions": "See the weaknesses section.", "limitations": "N/A.", "rating": 3, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The analysis presented by the authors is novel, and it improves upon prior work in some settings (e.g. tightening the gap on sharpness upper and lower bounds from $\\mathcal{O}(\\eta)$ to $\\mathcal{O}(\\eta^3)$ for two-layer linear neural networks).\n2. The paper imports a lot of ideas from the literature on control theory to studying the optimization of neural nets. This might be useful more generally.", "weaknesses": "1. Why is the whole \"trajectory alignment\" phenomenon is interesting in the first place? What new things does it tell us?\n2. It seems very obvious that for high-dimensional, deep networks, gradient descent doesn't converge to the same minimizers for different initializations. Is this whole phenomenon a product of small $d$ and toy examples? You fix the Xavier initialization for all experiments, and as far as I can tell the widest network used has width 256. This is a far cry from modern neural networks.\n3. The conditions you have on the initialization in Thm 4.2 are very restrictive. Do they hold in practice?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your valuable feedback on our paper. We have carefully considered your comments and improved our manuscript accordingly. We have also conducted additional experiments, with corresponding figures available in the attached PDF files in the \"global\" response (Response to All Reviewers). When citing figures from these additional experiments, we will use the format \"Figure (number) - Additional Experiments\" to differentiate them from the figures in our manuscript. We hope that the reviewer will find our response satisfactory. We kindly ask for your consideration in adjusting the review score, as we believe we have effectively addressed your primary concerns.\n\nRegarding the specific review comments, we would like to respond as follows:\n\n>(W1) Why is the whole \"trajectory alignment\" phenomenon interesting in the first place?\n\nThe trajectory alignment phenomenon holds significant interest as it sheds light on the training dynamics of neural network optimization using GD with large step sizes, particularly in the Edge of Stability (EoS) regime. This phenomenon is intriguing because prior understanding in the EoS regime primarily focused on the behavior of sharpness alone. Our work introduces a novel observation where reparameterized GD trajectories consistently align on a specific bifurcation diagram regardless of their initializations. This discovery is notable due to the intricate and non-convex nature of neural network optimization, where the algorithm trajectory is heavily influenced by initialization choices. Understanding trajectory alignment has the potential to advance our fundamental understanding of optimization processes in neural networks, offering insights into improved optimization strategies and landscape properties.\n\nFurthermore, the existing theoretical understanding of the EoS phenomenon remains limited, with few works rigorously examining the limiting sharpness, even in simplified models. Our work goes beyond this by providing a rigorous characterization of the oscillating dynamics and the evolution of sharpness for a broader range of networks.\n\n>(W2) Is this whole phenomenon a product of small $d$ and toy examples?\n\nWe would like to clarify that the trajectory alignment phenomenon is not a product of small data dimension $d$ or toy examples. Our empirical observations consistently demonstrate its presence across higher data dimensions and wider networks. In fact, this phenomenon becomes more pronounced as network width increases.\n\nTo illustrate this, in Figure 3, we show that the trajectory alignment phenomenon occurs in wide networks with a width of 256, but not as clearly in narrower networks with a width of 64. Furthermore, we conducted experiments on fully connected (FC) networks trained on the CIFAR-10 dataset, employing widths of 512 and 1024. The results, displayed in Figures 14, 15, and 16 in Appendix A.4, indicate the trajectory alignment phenomenon in large-width FC networks trained on a real-world dataset.\n\nFor additional validation, we conducted additional experiments on convolutional neural networks (CNNs) trained on CIFAR-10, and the corresponding results are presented in Figure 1 - Additional Experiments. We observe the trajectory alignment phenomenon in wide convolutional networks trained on a real-world dataset. Therefore, the trajectory alignment phenomenon is not confined to small dimensions or toy examples.\n\nThe use of Xavier initialization was maintained consistently across experiments to ensure fair comparisons and minimize potential biases originating from initialization. It's important to note that the presence of the trajectory alignment phenomenon is not related to the initialization method. Rather, it emerges as a characteristic behavior during optimization using GD with \"large\" step sizes.\n\nIn summary, our empirical evidence robustly demonstrates that the trajectory alignment phenomenon occurs in various experimental settings that include wide FC networks, CNNs (with widths of up to 1024) trained on a real-world dataset, and high data dimensions such as $d=32\\times 32\\times 3=3072$.\n\n>(W3) The conditions you have on the initialization in Theorem 4.2 are very restrictive.\n\nWe want to clarify that the conditions on the initialization in Theorem 4.2 are not as restrictive as they might seem. For simplicity, let us focus on the case of log-cosh loss.\n\nIn Theorem 4.2, the conditions on the initialization $(p_0, q_0)$ are satisfied when $\\lvert p_0 \\rvert \\le 1$ and $q_0 \\in (\\frac{3}{4}, 1)$. Although we assumed the bound $\\lvert p_0 \\rvert \\le 1$ for simplicity, our proof extends to the assumption $\\lvert p_0 \\rvert \\le K$ for any positive constant $K$, modulo some changes in numerical constants. Hence, the only essential assumption is that $q_0 \\in (\\frac{3}{4}, 1)$.  Notably, Theorem 4.1 already covers the range $q_0 > 1$, leaving out only the case of $0 < q_0 < \\frac{3}{4}$.\n\nEmpirically, we have observed that when $q_0$ satisfies $0 < q_0 < \\frac{3}{4}$, $q_t$ increases during early dynamics and eventually reaches the regime $\\frac{3}{4} < q_t < 1$, which is covered by our theory. We conducted additional experiments to support this observation. The results, depicted in Figure 2 - Additional Experiments, demonstrate that for small step sizes, GD trajectories initialized with $q_0<1$ eventually reach the set $S_A := [-1, 1]\\times (\\frac{3}{4},1)$ in most cases. This further emphasizes that the assumption in Theorem 4.2 encompasses a wide range of scenarios."}, {"review_id": "usIpRljryf", "reviewer": "Reviewer_Qm1U", "summary": "This paper discovers and studies a new phenomenon that occurs when training neural nets on 1 data point with large enough learning rates.    The phenomenon concerns the end-stage evolution of the tuple (residual, sharpness), where the residual (defined as \"prediction - target\") is a scalar because there is only 1 data point.   If one fixes the sharpness and considers the evolution of the residual, and if one linearizes the network, then the dynamical system that governs the evolution of the residual has a stable period-2 orbit whose value depends on the fixed sharpness, i.e. residual $\\to \\pm$  function(sharpness).\n\nThe main finding of the paper is that the residual and sharpness rapidly evolve onto the curve \"residual $=\\pm$  function(sharpness)\" and then remain on this curve for the end stage of training.  This finding is made empirically on some real architectures and is also proved rigorously in two highly simplified settings.  This is interesting because it is a richer characterization of the training dynamics than has been previously known (previous knowledge is limited to the behavior of the sharpness and says nothing about the residual).\n", "questions": "- Could the authors please cite these two papers in the related work section?  https://arxiv.org/abs/1807.05031, https://arxiv.org/abs/2002.09572.  These papers gave evidence prior to Cohen et al '21, that the step size affects the sharpness along the optimization trajectory.\n\n- I wonder if the reason why the trajectory alignment phenomenon occurs on wide networks but not narrow ones is that the linearization of the network is more accurate for wide networks.  Therefore, I wonder if one could also see trajectory alignment on narrow nets if one adopted the scaling trick from Chizat and Bach (https://arxiv.org/abs/1812.07956) which makes any network closer to its linearization.\n\n- I believe Lemma 2.1(ii) has a typo, it should be r_hat(q) \n\n", "limitations": "Please see the second point in the 'weaknesses' section.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- I think the identification of the trajectory alignment phenomenon is a solid contribution.\n- It's great that the authors can rigorously establish this phenomenon in several simplified settings\n", "weaknesses": "- The main weakness is that the trajectory alignment phenomenon seems mostly limited to the setting of training on one data point.  The authors do briefly try to extend it to the setting of multiple data points, but the experimental evidence is not highly convincing (especially the $\\ell_\\infty$ plot) and I do not really see why the phenomenon _would_ be true.  For example, I do not see why the generalized sharpness metric (which is kind of like trace(Hessian) / n) would be the right generalization.\n\n- Another weakness is that the EOS setting seems different than the one from Cohen et al '21.  In that paper, training with gradient flow makes the sharpness continually increase, whereas my understanding is that in this paper, training with gradient flow does not make the sharpness continually increase.  Is my understanding correct?  If so, I would prefer if the paper were clearer about the differences from Cohen et al '21, so as not to add confusion to the literature.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate your recognition and constructive feedback on our work. We have carefully considered your comments to enhance our manuscript. We have also conducted additional experiments, and the associated figures can be found in the attached PDF file in the \"global\" response (Response to All Reviewers). When referencing figures from the additional experiments, we will use the format \"Figure (number) - Additional Experiments\" to distinguish them from the figures in the manuscript.\n\nRegarding the specific comments in the review, we would like to address them as follows:\n\n>(W1) The trajectory alignment phenomenon seems mostly limited to the setting of training on one data point.\n\nWe understand the reviewer's concerns regarding the multiple data points setting. Specifically, we initially claimed the trajectory alignment phenomenon throughout all settings, including the mean, $\\ell_1$, $\\ell_2$, and $\\ell_\\infty$ plot. In the revised manuscript, we will tone down and claim that trajectory alignment phenomenon is observed especially at the \"mean\" plot across wide range of setting.\n\nWe want to emphasize that the trajectory alignment is clearly observed in the mean plot. Our experiments in Appendix A.4 depicted in Figures 13-16 consistently demonstrate trajectory alignment for large-width FC networks trained on CIFAR-10. We also conducted additional experiments on CNN trained on CIFAR-10, presented in Figure 1 - Additional Experiments, and the results also show trajectory alignment phenomenon. The empirical evidence consistently supports this phenomenon, although explaining it theoretically remains outside the scope of our work.\n\nRegarding the choice of (p,q) in the multiple data points setting, we choose (p,q) to set (p,q) = (0,1) for parameters corresponding to global minima and sharpness (defined as maximum eigenvalue of loss Hessian, rather than trace(Hessian) / n) equal to 2/(step size). As a conjecture, other careful choices of (p,q) may also exhibit trajectory alignment, warranting further investigation.\n\nFurthermore, we present an extension of Eq.(3) to the multiple data settings, assuming a network can be well approximated by its linearized model, yielding insights into the relationship between residual and sharpness. Specifically, under the linear approximation $f(X;\\Theta_{t+1})\\approx f(X;\\Theta_t)+Df(X;\\Theta_t)(\\Theta_{t+1}-\\Theta_t)$, we obtain\n$$f(X;\\Theta_{t+1})-y \\approx (f(X; \\Theta_t)-y) - \\frac{\\eta}{n} K(X,X;\\Theta_t) \\ell'(f(X; \\Theta_t)-y),$$\nwhere $K(X,X;\\Theta)\\in\\mathbb{R}^{n\\times n}$ is a Gram matrix with (i,j)-th entry given by $\\langle \\nabla_{\\Theta} f(x_i;\\Theta),\\nabla_{\\Theta}f(x_j;\\Theta)\\rangle$, also known as NTK [Jacot2018]. In case of n=1, $K(x_1,x_1;\\Theta)=\\lVert\\nabla_{\\Theta} f(x_1;\\Theta)\\rVert_2^2$ and this matches with Eq.(3). Understanding the trajectory alignment phenomenon in this context remains a topic for future work.\n\n>(W2) In [Cohen2021], training with gradient flow makes the sharpness continually increase, whereas in this paper, training with gradient flow does not make the sharpness continually increase.\n\nAs the reviewer pointed out, our findings on progressive sharpening differ from those in [Cohen2021]. While [Cohen2021] observed an increase in sharpness during training in the gradient flow regime, our work reveals that sharpness increases while exhibiting oscillations in the EoS regime. This distinction could indeed be attributed to differences in the loss functions used in two studies. In the revised version, we will make this point clearer and provide a detailed discussion on the differences between our work and the original literature.\n\n>(Q1) Could the authors please cite [Jastrzebski2018] and [Jastrzebski2020] in the related work section?\n\nWe acknowledge that [Jastrzebski2018] and [Jastrzebski2020] provided evidence of the step size's impact on sharpness along the optimization trajectory before [Cohen2021]. In the revised manuscript, we will include proper citations to these papers in the related work section.\n\n>(Q2) I wonder if one could also see trajectory alignment on narrow nets if one adopted the scaling trick from [Chizat2018] which makes any network closer to its linearization.\n\nThank you for your important question and insightful suggestion. We also share the belief that the trajectory alignment phenomenon occurs on wide networks due to their close approximation to their linearized models, as we adopted a linear approximation of a network to obtain Eq.(3).\n\nTo further investigate the relationship between the trajectory alignment and network linearization, we conducted additional experiments on narrow networks (plots are not provided due to the space constraint). As the reviewer suggested, we adopted the scaling trick from [Chizat2018] to bring the networks closer to their linearization. The empirical evidence clearly demonstrates that increasing the scaling factor induces the trajectory alignment phenomenon, even for narrow networks.\n\nWe will include a detailed description of the additional experiments and their results in the revised manuscript. Additionally, we will provide a comprehensive discussion about the relation between the linearization of the network and the trajectory alignment phenomenon.\n\nLastly, we further note that our theoretical results do not rely on linear approximation of networks; instead, we rigorously analyze the exact dynamics, making our theory applicable to a two-layer linear network with an arbitrary width.\n\n>(Q3) I believe Lemma 2.1(ii) has a typo, it should be $\\hat{r}(q)$.\n\nThank you for pointing out the typo. We apologize for the oversight, and we will correct Lemma 2.1(ii) in the revised manuscript.\n\n---\n__Reference__\n\n[Chizat2018] https://arxiv.org/abs/1812.07956\n\n[Cohen2021] https://openreview.net/forum?id=jh-rTtvkGeM\n\n[Jastrzebski2018] https://arxiv.org/abs/1807.05031\n\n[Jastrzebski2020] https://arxiv.org/abs/2002.09572"}]

Author rebuttal:
Dear Reviewers, thank you for your valuable feedback and insightful comments. We have carefully considered all your questions, concerns, and suggestions, and we greatly appreciate the time and effort you have invested in reviewing our work.

We did our best to address all of your comments and will make the necessary revisions to improve the clarity, rigor, and validity of our manuscript. We kindly ask for your consideration in adjusting the review score if you find that our response has effectively addressed your concerns.

We are delighted to see that our novel characterization of the EoS phenomenon has been recognized by Reviewers Qm1U and ft7H. We extend our sincere gratitude to Reviewers Qm1U, ft7H, and mAcJ for their positive comments on our theoretical establishment of the trajectory alignment phenomenon.

In the revised manuscript, we will make the following major modifications.

- In abstract and main text, we initially claimed "When the EoS phenomenon occurs, different GD trajectories (after a proper reparameterization) align on a specific bifurcation diagram determined solely by the loss function, __independent of the network architecture, training data, and step size__." We will revise this claim to "When the EoS phenomenon occurs, different GD trajectories (after a proper reparameterization) align on a specific bifurcation diagram __independent of initialization__."

- We will replace __Figure 4__ (experiments on synthetic Gaussian dataset) in the main text by __Figure 13 in Appendix__ and __Figure 1 of the attached PDF file__, which demonstrates the trajectory alignment phenomenon of GD trained on a __real-world dataset (CIFAR-10) using FC networks and CNNs__.

Furthermore, we conducted additional experiments to address the reviewers' concerns and the results are presented in Figures 1-3 of the attached PDF file. Below, we provide a detailed description of each figure:

-  __Figure 1: Training CNN on CIFAR-10.__ To demonstrate the trajectory alignment phenomenon in more realistic settings, we trained CNN of width 500 and depth 3 on CIFAR-10, and observe the trajectory alignment phenomenon.

- __Figure 2: Assumptions on initialization in Theorem 4.2.__ To verify that the assumptions made in Theorem 4.2 covers wide range of cases, we considered log-cosh loss and investigated whether GD iterates $(p_t, q_t)$ with different initializations $(p_0, q_0)$ eventually fall into $S_A := [−1, 1] \times (\frac{3}{4} , 1)$ within $10^5$ iterations, under different step sizes ($\eta = 0.1, 0.01$). Note that once the iterate enters $S_A$, the assumption in Theorem 4.2 is satisfied and the theorem can be applied thereafter. Figure 2 shows that for small step sizes, GD iterates with initialization in the entire range of $q_0 \in (0, 1)$ eventually reach $S_A$ in most cases.

- __Figure 3: Sharpness value closely follows $\tilde{\lambda}(q_t)$ predicted by Theorem 4.4 for large width networks.__ The sharpness prediction $\tilde{\lambda}(q_t)$ for progressive sharpening in Theorem 4.4 was accurate for single-neuron case as shown in Figure 1(a) of the main text. We conducted additional experiments on two layer linear network with large widths (256 and 1024) and observe that our sharpness prediction is accurate regardless of the width.

We hope that these additional experiments further strengthen our contributions.

Meta-review: The paper studies the edge of stability phenomenon in training of deep neural networks. The main claim is that different GD trajectories — after proper reparameterization — align on a specific bifurcation diagram independent of initialization. 

**Weaknesses**

- One of the major limitations of the paper, as echoed by reviewers, is the experimental scope. The evidence presented mostly pertains to toy models and a singular data point, making the universal applicability of the findings questionable.
- Another concern lies in the strong claims made regarding the independence of the bifurcation diagram from multiple factors like network architecture and training data. These claims have not been sufficiently demonstrated across broader settings, and the claim was revised during discussions.
- Reviewers also raised questions about the generalizability of the reparameterization variables (p, q) and their utility in larger, more complex scenarios, beyond toy models.
- Most of the results and theory focuses on a perhaps too contrived case of training on a single data point. While the Authors have added additional results concerning training on more data points, it is still fully clear to what extend trajectory alignment also holds in this more realistic setting.

**Strengths**

- The concept of trajectory alignment is a solid and novel contribution that will potentially bring clarity to the field.
- The theoretical contributions such as using control theory or defining the canonical reparametrization were also appreciated. The trajectory alignment is proven rigorously. There is an argument to be made that the trajectory alignment is expected or plausible given results in prior work, however the paper makes concrete predictions (in the simple case) for the exact shape of the trajectory.

Although the paper has limitations, particularly concerning its experimental scope and the strong claims made, its theoretical contributions outweight the weaknesses. Based on this, it is my pleasure to recommend acceptance. I would like to ask the Authors to focus in the camera ready version on clarification of the main claim, and bringing in results on training involving a full dataset.
Decision: Accept (poster)