Title: Trajectory Alignment: Understanding the Edge of Stability Phenomenon via Bifurcation Theory
Venue: NeurIPS 2023
Accepted: yes
Keywords: non-convex optimization

trajectory alignment of GD

edge of stability

progressive sharpening

bifurcation theory
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

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 or regularization of gradient-based methods plays a key role in generalization of deep learning (Vardi, 2022). There is a growing literature (Gunasekar et al., 2017, 2018; Soudry et al., 2018; Arora et 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 a certain sharpness 1 value depending on the step size η . Specifically, for reasonable choices of η , (a) the sharpness of the loss at the GD iterate gradually increases throughout training until it reaches the stability threshold 2 of 2 /η (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 at developing a 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 et al., 2023). In this paper, we study these phenomena through the lens of bifurcation theory , both empirically and theoretically.

Motivating observations: Figure 1 illustrates the GD trajectories with different initializations and fixed step sizes trained on three types of two-dimensional functions: (a) log(cosh( xy )) , (b) 1 2 (tanh( x ) y ) 2 , and (c) 1 2 ( ELU ( x ) y ) 2 , where x and y are scalars. The functions L : R 2 → R have sharpness y 2 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 a 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:

(a)

log(cosh(

xy

))

(b)

1

2

(tanh(

x

)

y

)

2

(c)

1

2

(

ELU

(

x

)

y

)

2

Initalization (xo. yo) = (2. 10)

Initialization (X. yo) = (2,8)

22

Initiliation (Xs- yo) = (3., 34)

Initilization (xs- yo) = (3, 32)

X Bifuration pint

Initialization (xo. yo) = (2., 6)

p

Bifurcation diagram

nitalization (x-yo) = (3, 30)

Bifrcation diogrim

Global minima

XBifecation pont

Glotal minina

Bifurcation diagram

2.0

1.5

2.5

2.0

0

1000

750 -

入=

600

入=

ssudeuss

500

250 -

200

入=

0

200

400

600

800

0001

0-

200

400

600

800

0001

0

1000

2000

3000

4000

5000

0.4 -

0.5

0.2-

0.0

X0.0

0.2

0.5 -

0.4

1.0

200

400

600

800

0001

200

400

600

800

0001

0

1000

2000

000

4000

5000

ss s

Time steps

Time steps

EoS phenomenon: GD converges to a global minimum near the point (0 , √ 2 /η ) with sharpness close to 2 /η . During the convergence phase, the training dynamics exhibit period2 oscillations.

For different initializations, GD trajectories for a given step size align on the same curve. For example, Figure 1a shows that GD trajectories with different initializations closely follow a specific U-shaped curve until convergence. We call this phenomenon trajectory alignment .

In Figures 1b and 1c, GD trajectories are aligned on a curve with a fractal structure that qualitatively resembles the bifurcation diagram of a typical polynomial map, such as the logistic map. Particularly, Figure 1c demonstrates a period-halving phase transition in the GD dynamics, shifting from period4 oscillation to period2 oscillation.

Surprisingly, the curve that GD trajectories approach and follow coincides with the bifurcation diagram of a one-dimensional map x ↦→ x -η ∂ ∂x L ( x, y ) with a fixed 'control parameter' y . The stability of its fixed point x = 0 changes at the bifurcation point ( x, y ) = (0 , √ 2 /η ) , where period-doubling bifurcation occurs. Note that this point is a global minimum with sharpness 2 /η .

Interestingly, such striking behaviors can also be observed in more complex models, up to a proper reparameterization, as we outline in the next subsection.

1.1 Our contributions

In this paper, we discover and study the trajectory alignment behavior of (reparameterized) GD dynamics in the EoS regime. To our best knowledge, we are the first to identify such an alignment with a specific bifurcation diagram independent of initialization. Our empirical findings are rigorously proven for both two-layer fully-connected networks and single-neuron nonlinear networks. Our main contributions are summarized below:

In Section 2, we introduce a novel canonical reparameterization of training parameters, which incorporates the data, network, and GD step size. This reparameterization allows us to study the trajectory alignment phenomenon in a unified framework. Through empirical study, Section 3 demonstrates that the alignment property of GD trajectories is not limited to 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 EoS 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 to 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 4 for a single-neuron 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˛ ebski et al. (2019) and Jastrzebski et 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 et al. (2022) analyze EoS through experiments and identify the relations between the behavior of loss, iterates, and sharpness. Ma et 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 a 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 an implicit bias toward the max-margin 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 /η . Zhu et al. (2023) study the loss ( x, y ) ↦→ 1 4 ( x 2 y 2 -1) 2 and prove that the sharpness converges close to 2 /η with a local convergence guarantee. Notably, Ahn et al. (2023) study the function ( x, y ) ↦→ ℓ ( xy ) where ℓ is convex, even, and Lipschitz, and provide a 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 /η -O ( η ) and 2 /η . Our theoretical results extend their results on a single-neuron linear network to a two-layer fully-connected linear network and provide an improved characterization on the limiting sharpness, tightening the gap between upper and lower bounds to only O ( η 3 ) .

The trajectory alignment 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 a 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 reparameterization and offers empirical evidence for the alignment 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 u and v , we denote the ℓ p norm of u by ∥ u ∥ p , their tensor product as u ⊗ v , and u ⊗ u by u ⊗ 2 . For a matrix A , we denote the spectral norm by ∥ A ∥ 2 . Given a function L and a parameter Θ , we use λ max ( Θ ) := λ max ( ∇ 2 Θ L ( Θ )) to denote the sharpness (i.e., the maximum eigenvalue of the loss Hessian) at Θ . We use asymptotic notations with subscripts (e.g., O ℓ ( · ) , O δ,ℓ ( · ) ) in order to hide constants that depend on the parameters or functions written as subscripts.

2.1 Problem settings

We study the optimization of neural network f ( · ; Θ ) : R d → R parameterized by Θ . We focus on a simple over-parameterized setting trained on a single data point { ( x , y ) } , where x ∈ R d and y ∈ R . We consider the problem of minimizing the empirical risk

where ℓ is convex, even, and twice-differentiable with ℓ ′′ (0) = 1 . We minimize L using GD with step size η : Θ t +1 = Θ t -η ∇ Θ L ( Θ t ) . The gradient and the Hessian of the function are given by

Suppose that Θ ∗ be a global minimum of L , i.e., f ( x ; Θ ∗ ) = y . In this case, the loss Hessian and the sharpness at Θ ∗ are simply characterized as

2.2 Canonical reparameterization

Definition 2.1 (canonical reparameterization) . For given step size η , the canonical reparameterization of Θ is defined as

Under the canonical reparameterization, p = 0 represents global minima, and Eq. (1) implies that the point ( p, q ) = (0 , 1) is a global minimum with sharpness 2 /η . Note that ( p, q ) alone does not, in general, uniquely determine the value of Θ . Rather, the motivation for this reparameterization technique is to effectively analyze the complex GD dynamics in the high-dimensional parameter space by reducing it to a 2-dimensional representation. The update of p can be written as

which can be obtained by first-order Taylor approximation on f for small step size η . 3

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 q t as a control parameter. We first review some basic notions used in bifurcation theory (Strogatz, 1994).

Definition 2.2 (stability of fixed point) . Let z 0 be a fixed point of a differentiable map f : R → R , i.e., f ( z 0 ) = z 0 . We say z 0 is a stable fixed point of f if | f ′ ( z ) | < 1 , and we say z 0 is an unstable fixed point of f if | f ′ ( z ) | > 1 .

̸

Definition 2.3 (stability of periodic orbit) . A point z 0 is called a periodp point of a map f : R → R if z 0 is the fixed point of f p and f j ( z 0 ) = z 0 for any 1 ≤ j ≤ p -1 . The orbit of z 0 , given by { z j = f j ( z 0 ) | j = 0 , 1 , . . . , p -1 } is called the periodp orbit of f . A periodp orbit is stable (unstable) if its elements are stable (unstable) fixed points of f p , i.e., ∏ p -1 j =0 | f ′ ( z j ) | < 1 ( > 1 ).

Now we analyze the bifurcation of the one-parameter family of mappings f q : R → R given by

where q is a control parameter and r is a differentiable function satisfying Assumption 2.4 below.

(a)

α

= 5

(b)

α

= 10

(c)

q

t

r

(

p

t

)

and

p

vs.

t

2.00 -

2.25

q=

Global minima

tanh(p)

P

0.8 -

1.0 -

Globalminima

=b

tanh(p)

P

.0

Q

1.2 -

1.4

rpT

q

=1

1.75 -

0.8 -

0.6 -

- 9°0

q

0

2000

4000

6000

8000

10000

1.25

0.4

5.0

p=0

2.5 -

1.00

p

0.0

0.75

0.2

2.5 -

0.50 -

5.0

20

10

U

01

20

0

2000

000+

6000

8000

10000

Assumption 2.4. A function r : R → R is even, continuously differentiable, r (0) = 1 , r ′ (0) = 0 , r ′ ( p ) < 0 for any p > 0 , and lim p →∞ r ( p ) = lim p →-∞ r ( p ) = 0 . In other words, r is a smooth, symmetric bell-shaped function with the maximum value r (0) = 1 .

̸

We note that Eq. (3) can be rewritten as p t +1 = f q t ( p t ) if we define r by r ( p ) := ℓ ′ ( p ) p 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) to explain EoS for ( x, y ) ↦→ ℓ ( xy ) .

log -cosh loss: ℓ log -cosh ( p ) := log(cosh( p )) . Note ℓ ′ log -cosh ( p ) = tanh( p ) .

square-root loss: ℓ sqrt ( p ) := √ 1 + p 2 . Note ℓ ′ sqrt ( p ) = p √ 1+ p 2 .

If r satisfies Assumption 2.4, then for any 0 < q ≤ 1 , there exists a nonnegative number p such that r ( p ) = q , and the solution is unique which we denote by ˆ r ( q ) . In particular, ˆ r : (0 , 1] → R ≥ 0 is a function satisfying r (ˆ r ( q )) = r ( -ˆ r ( q )) = q for any q ∈ (0 , 1] .

Lemma 2.1 (period-doubling bifurcation of f q ) . Suppose that r is a function satisfying Assumption 2.4. Let p ∗ = sup { p ≥ 0 | xr ′ ( x ) r ( x ) > -1 for any | x | ≤ p } and c = r ( p ∗ ) . If p ∗ = ∞ , we choose c = 0 . Then, the one-parameter family of mappings f q : R → R given by Eq. (4) satisfies

(i) If q > 1 , p = 0 is the stable fixed point.

(ii) If q ∈ ( c, 1) , p = 0 is the unstable fixed point and {± ˆ r ( q ) } is the stable period2 orbit.

Proof. The map f q has the unique fixed point p = 0 for any q > 0 . Since | f ′ q (0) | = | 1 -2 q | , p = 0 is a stable fixed point if q > 1 and p = 0 is an unstable fixed point if 0 < q < 1 . Now suppose that q ∈ ( c, 1) . Then, we have f q (ˆ r ( q )) = -ˆ r ( q ) and f q ( -ˆ r ( q )) = ˆ r ( q ) , which implies that {± ˆ r ( q ) } is a period2 orbit of f q . Then, | f ′ q (ˆ r ( q )) | = | f ′ q ( -ˆ r ( q )) | = ∣ ∣ ∣ 1 + 2ˆ r ( q ) r ′ (ˆ r ( q )) q ∣ ∣ ∣ < 1 implies that {± ˆ r ( q ) } is a stable period2 orbit.

According to Lemma 2.1, the stability of the fixed point p = 0 undergoes a change at q = 1 , resulting in the emergence of a stable period2 orbit. The point ( p, q ) = (0 , 1) is referred to as the bifurcation point , where a period-doubling bifurcation occurs. A bifurcation diagram illustrates the points asymptotically approached by a system as a function of a control parameter. In the case of the map f q , the corresponding bifurcation diagram is represented by p = 0 for q ≥ 1 and p = ± ˆ r ( q ) (or equivalently, q = r ( p ) ) for q ∈ ( c, 1) .

̸

It is worth noting that the period2 orbit {± ˆ r ( p ) } becomes unstable for q ∈ (0 , c ) . If we choose r to be r ( p ) = ℓ ′ ( p ) p for p = 0 and r (0) = 1 , then 1 + pr ′ ( p ) r ( p ) = ℓ ′′ ( p ) r ( p ) > 0 for all p , assuming ℓ is convex. Consequently, for log -cosh loss and square root loss we have c = 0 , indicating that the period2 orbit of f q remains stable for all q ∈ (0 , 1) . However, in Section 5, we will consider r with c > 0 , which may lead to additional bifurcations.

(a)

m

= 64

,

L

= 3

(b)

m

= 256

,

L

= 3

(c)

m

= 256

,

L

= 10

q=

tanh(p)

-p

1.0 -

q=

tanh(p)

1.0 -

q=

tanh(p)

1.0 -

Global minima

Globalminima

Globalminima

0.8

0.8 -

0.8 -

Qo.6

b

0.6

b

0.6

0.4

0.4 -

0.4 -

0.2

0.2

0.2

0.0 -

-10

-15

10

10

10

10

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 d → R written as

where ϕ is an activation function, W 1 ∈ R m × d , W l ∈ R m × m for 2 ≤ l ≤ L -1 , and w L ∈ R m . All L layers have the same width of m . We minimize the empirical risk L ( Θ ) = ℓ ( 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 α . For this analysis, we fix the training data point and hyperparameters as x = e 1 = (1 , 0 , . . . , 0) , y = 1 , η = 0 . 01 , d = 10 , and focus on the log -cosh loss for ℓ , with either ϕ ( t ) = t (linear) or ϕ ( t ) = tanh( t ) . We note that the trajectory alignment phenomenon is 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 a two-layer fully-connected linear network with a width of m = 256 . In Figure 2a, when the weights are initialized with a smaller scale ( α = 5 ), the initial value of q is greater than 1 , and it converges towards the minimum with only a small change in q t until convergence. In this case, the limiting sharpness is relatively smaller than 2 /η , and the EoS phenomenon does not occur. This case is referred to as the gradient flow regime (Ahn et al., 2023). On the other hand, in Figure 2b, when the weights are initialized with a larger scale ( α = 10 ), the initial value of q is less than 1 , and we observe convergence towards the point (close to) ( p, q ) = (0 , 1) . This case is referred to as the EoS regime . We note that choosing larger-than-standard scale α is not a necessity for observing EoS; we note that even with α = 1 , we observe the EoS regime when η 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 q = r ( p ) = ℓ ′ ( p ) p and observe that GD trajectories tend to align with this curve, which depends solely on ℓ . Figure 2b clearly demonstrates this alignment phenomenon. Additionally, we analyze the evolution of q t r ( p t ) and p t in Figure 2c. We observe that the evolution of q t r ( p t ) follows two phases. In Phase I , q t r ( p t ) approaches to 1 quickly. In Phase II , the ratio remains close to 1 . Notably, the convergence speed of q t r ( p t ) towards 1 is much faster than the convergence speed of p t towards 0 . In Sections 4 and 5, we will provide a rigorous analysis of this behavior, focusing on the separation between Phase I and Phase II.

The effect of width and depth. In Figure 3, we present the GD trajectories of tanh -activated networks with different widths and depths ( α = 5 ). All three cases belong to the EoS regime, where GD converges to a point close to ( p, q ) = (0 , 1) , resulting in a limiting sharpness near 2 /η . However, when comparing Figures 3a and 3b, we observe that the trajectory alignment phenomenon is not

(a)

P

(

z

) =

1

n

∑

n

i

=1

z

i

(b)

P

(

z

) =

∥

z

∥

2

(c) Loss and

q

vs.

t

1.0

q= lanng)

1.0 -

qae

0.9-

--p=0

--p= 0

0.8

2

0.8

1

90.6-

0.7

0.6

0

0

1000

2000

3000

4000

5000

1.0

0.5

0.4 -

80

0.4 -

0.2 -

0.6 -

0.3

0.4

0.2

0.0

10

15

P

20

25

30

0.2

0

1000

2000

Time steps

000E

4000

5000

0.9

1.0 -

q=a

p=0

0.8

1.0

一

q= a

p=0

0.8

q

0.6

0.7

9°0

0

0

1000

2000

3000

4000

5000

0.5 -

0.4

8°0

1.0

0.4

0.2

0.6

0.3

0.4

0.2

0°0

10

15

p

20

25

30

0

1000

0000000

Time steps

4000

5000

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 a 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 a deeper network with L = 10 , as depicted in Figure 3c.

Multiple training data points. In our trajectory alignment analysis, we have primarily focused on training with a single data point. However, it is important to explore the extension of this phenomenon to scenarios with multiple training data points.

To investigate this, we train a neural network on a dataset { ( x i , y i ) } n i =1 , where x i ∈ R d and y i ∈ R , by minimizing the empirical risk L ( Θ ) := 1 n ∑ n i =1 ℓ ( f ( x i ; Θ ) -y i ) . Defining X ∈ R n × d as the data matrix and y ∈ R n as the label vector, we introduce a generalized canonical reparameterization :

where P : R n → R can be a function such as a mean value or a specific vector norm.

In Figure 4, we consider training on a 50 example subset of CIFAR10 with only 2 classes and vary the network architecture. We use three-layer fully-connected network with tanh activation and convolutional network (CNN) with tanh activation. Under the generalized canonical reparameterization (5) for various choices of P , including the mean and the ℓ 2 norm, we observe the trajectory alignment phenomenon throughout all settings, indicating a common alignment property of the GD trajectories. However, unlike the single data point case, the alignment does not happen on the curve q = ℓ ′ ( p ) p . The precise characterization of the coinciding curve is an interesting direction for future research.

4 Trajectory Alignment of GD: A Theoretical Study

In this section, we study a two-layer fully-connected linear network defined as f ( x ; Θ ) := v ⊤ Ux , where U ∈ R d × m , v ∈ R m , and Θ denote the collection of all parameters ( U , v ) . We consider training this network with a single data point { ( x , 0) } , where x ∈ R d and ∥ x ∥ 2 = 1 . We run GD with step size η on the empirical risk

where ℓ is a 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 et al. (2023).

Assumption 4.1. The loss ℓ is convex, even, 1 -Lipschitz, and twice differentiable with ℓ ′′ (0) = 1 .

The canonical reparameterization (Definition 2.1) of Θ = ( U , v ) is given by

Under the canonical reparameterization, the 1 -step update rule of GD can be written as

̸

where we define the function r by r ( p ) := ℓ ′ ( p ) p for p = 0 and r (0) := 1 . Note that the sequence ( q t ) ∞ t =0 is monotonically increasing if q 0 ≥ 1 2 , which is the case our analysis will focus on.

We have an additional assumption on ℓ as below, motivated from Lemma 2.1.

Assumption 4.2. The function r ( p ) = ℓ ′ ( p ) p corresponding to the loss ℓ 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 q is initialized with q 0 > 1 .

Theorem 4.1 (gradient flow regime) . Let η ∈ (0 , 2 33 ) be a fixed step size and ℓ be a loss function satisfying Assumptions 4.1 and 4.2. Suppose that the initialization ( p 0 , q 0 ) satisfies | p 0 | ≤ 1 and q 0 ∈ ( 2 2 -η , min { 1 16 η , r (1) 2 η }) . Consider the GD trajectory characterized in Eq. (6) . Then, the GD iterations ( p t , q t ) converge to the point (0 , q ∗ ) such that

Theorem 4.1 implies that in gradient flow regime, GD with initialization Θ 0 = ( U 0 , v 0 ) and step size η converges to Θ ∗ which has the sharpness bounded by:

Hence, for small step size η , if the initialization satisfies ∥ U 0 x ∥ 2 2 + ∥ v 0 ∥ 2 2 < 2 η -1 , then the limiting sharpness is slightly below ∥ U 0 x ∥ 2 2 + ∥ v 0 ∥ 2 2 . Note that we assumed the bound | p 0 | ≤ 1 for simplicity, but our proof also works with the assumption | p 0 | ≤ K for any positive constant K modulo some changes in numerical constants. Moreover, our assumption on the upper bound of q 0 is 1 /η up to a 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 q = r ( p ) . To establish these results, we introduce additional assumptions on the loss ℓ .

Assumption 4.3. The function r ( z ) = ℓ ′ ( z ) z is C 4 on R and satisfies

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 I) . Let η be a small enough step size and ℓ be a loss function satisfying Assumptions 4.1, 4.2, and 4.3. Let z 0 := sup z { zr ′ ( z ) r ( z ) ≥ -1 2 } and c 0 := max { r ( z 0 ) , 1 2 } . Let δ ∈ (0 , 1 -c 0 ) be any given constant. Suppose that the initialization ( p 0 , q 0 ) satisfies | p 0 | ≤ 1 and q 0 ∈ ( c 0 , 1 -δ ) . Consider the reparameterized GD trajectory characterized in Eq. (6) . We assume that for all t ≥ 0 such that q t < 1 , we have p t = 0 . Then, there exists a time step t a = O δ,ℓ (log( η -1 )) , such that for any t ≥ t a ,

̸

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 ( t ≤ t a = O (log( η -1 )) ), GD iterates ( p t , q t ) approach closely to the bifurcation diagram r ( p ) = q , which we called Phase I in Section 3. In Phase II, GD trajectory aligns on this curve in the remaining of the training ( t ≥ t a ). Theorem 4.3 provides an analysis on Phase II stated as below.

Theorem 4.3 (EoS regime, Phase II) . Under the same settings as in Theorem 4.2, there exists a time step t b = Ω((1 -q 0 ) η -2 ) such that q t b ≤ 1 and q t > 1 for any t > t b . Moreover, the GD iterates ( p t , q t ) converge to the point (0 , q ∗ ) such that

Theorem 4.3 implies that in EoS regime, GD with step size η converges to Θ ∗ with sharpness

Note that Ahn et al. (2023) study the special case d = m = 1 and prove that the limiting sharpness is between 2 /η -O ( η ) and 2 /η . Theorem 4.3 provides tighter analysis on the limiting sharpness in more general settings, reducing the gap between the upper bound and lower bound to only O ( η 3 ) . Also, our result is the first to prove that the limiting sharpness in the EoS regime is bounded away from 2 /η by a nontrivial margin.

Wealso 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 t a denote the obtained iteration. Define the function ˜ λ : R > 0 → R given by

Then, the sequence ( ˜ λ ( q t ) ) ∞ t =0 is monotonically increasing. Moreover, for any t ≥ t a , the sharpness at GD iterate Θ t closely follows the sequence ( ˜ λ ( q t ) ) ∞ t =0 by satisfying

The gap between λ max ( Θ t ) and ˜ λ ( q t ) is bounded by a numerical constant, which becomes negligible compared to 2 /η for small η . In Figure 1a, we perform numerical experiments on a single-neuron case and observe that ˜ λ ( q t ) closely approximates the sharpness.

Note that Cohen et al. (2021) observe an increase in sharpness during training in the gradient flow regime, while our work reveals that sharpness increases when exhibiting oscillations in the EoS regime. This distinction may be linked to the selection of the loss function, as our study focuses on Lipschitz convex losses, while Cohen et al. (2021) examine the squared loss.

5 EoS in Squared Loss: Single-neuron Nonlinear Network

Our canonical reparameterization has a limitation in explaining the EoS phenomenon under squared loss ℓ ( p ) = 1 2 p 2 , as the function r ( p ) = ℓ ′ ( p ) p = 1 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 a simple toy model to gain insight into the EoS phenomenon and trajectory alignment of GD under squared loss.

We study the GD dynamics on a two-dimensional function L ( x, y ) := 1 2 ( ϕ ( x ) y ) 2 , where x , y are scalars and ϕ 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 , ϕ ′ (0) = 1 , lim x →∞ ϕ ( x ) = 1 , and lim x →-∞ ϕ ( x ) = -1 .

One good example of ϕ satisfying Assumption 5.1 is tanh . For this section, we use an alternative reparameterization defined as below.

Definition 5.2. For given step size η , the ( p, q ) reparameterization of ( x, y ) ∈ R 2 is defined as

Under the reparameterization, the 1 -step update rule can be written as

̸

where the function r is given by r ( z ) := ϕ ( z ) ϕ ′ ( z ) z for z = 0 and r (0) := 1 .

We can observe a notable resemblance between Eq. (7) and Eq. (6). Indeed, our theoretical findings for a 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 ) := 1 2 ( ϕ ( x ) y ) 2 exhibits the same gradient flow, EoS (Phase I, 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 η -2 | r ′′ (0) | + O ( η ) . For formal statements of the theorems and the proofs, we refer the reader to Appendix C.

6 Conclusion

In this paper, we provide empirical evidence and rigorous analysis to demonstrate the interesting phenomenon of GD trajectory alignment in the EoS regime. Importantly, we show that different GD trajectories, under the canonical reparameterization, align on a 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 progressive sharpening of GD. One immediate future direction is to understand the trajectory alignment behavior when trained on multiple data points. Lastly, it 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 (IITP) 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-2019R1A5A1028324, RS-2023-00211352) funded by the Korea government (MSIT), and a grant funded by Samsung Electronics Co., Ltd.

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://proceedings.mlr.press/v162/ahn22a.html .

Kwangjun Ahn, Sébastien Bubeck, Sinho Chewi, Yin Tat Lee, Felipe Suarez, and Yi Zhang. Learning threshold neurons via the 'edge of stability'. NeurIPS 2023 (arXiv:2212.07469) , 2023.

Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparameterization. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning , volume 80 of Proceedings of Machine Learning Research , pages 244-253. PMLR, 10-15 Jul 2018. URL https://proceedings.mlr. press/v80/arora18a.html .

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. Garnett, editors, Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019. URL https://proceedings.neurips.cc/paper_files/paper/2019/ file/c0c783b5fc0d7d808f1d14a6e9c8280d-Paper.pdf .

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 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 948-1024. PMLR, 17-23 Jul 2022. URL https://proceedings.mlr.press/v162/arora22a.html .

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 https://proceedings.mlr.press/v202/chen23b.html .

Lénaïc 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://proceedings.mlr.press/ v125/chizat20a.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 https://openreview.net/forum?id=jh-rTtvkGeM .

Alex Damian, Eshaan Nichani, and Jason D. Lee. Self-stabilization: The implicit bias of gradient descent at the edge of stability. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=nhKHA59gXz .

Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro. Implicit regularization in matrix factorization. In I. Guyon, U. V on 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://proceedings.neurips.cc/ paper_files/paper/2017/file/58191d2a914c6dae66371c9dcdc91b41-Paper.pdf .

Suriya Gunasekar, Jason D Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent 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 https://proceedings.neurips.cc/paper_ files/paper/2018/file/0e98aeeb54acf612b9eb4e48a269814c-Paper.pdf .

Stanislaw Jastrzebski, Maciej Szymczak, 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:// openreview.net/forum?id=r1g87C4KwB .

Stanisław Jastrz˛ ebski, 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://openreview.net/ forum?id=SkgEaj05t7 .

Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. In International Conference on Learning Representations , 2019a. URL https://openreview. net/forum?id=HJflg30qKX .

Ziwei Ji and Matus Telgarsky. The implicit bias of gradient descent on nonseparable data. In Alina Beygelzimer and Daniel Hsu, editors, Proceedings of the Thirty-Second Conference on Learning Theory , volume 99 of Proceedings of Machine Learning Research , pages 1772-1798. PMLR, 25-28 Jun 2019b. URL https://proceedings.mlr.press/v99/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 https://proceedings.mlr.press/v202/kreisler23a.html .

Sungyoon Lee and Cheongjae Jang. A new characterization of the edge of stability based on a sharpness measure aware of batch gradient distribution. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=bH-kCY6LdKg .

Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normalization layers: Sharpness reduction. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022. URL https://openreview.net/forum?id=xp5VOBxTxZ .

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.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. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019. URL https://proceedings.neurips.cc/paper_files/paper/2019/ file/bdbca288fee7f92f2bfa9f7012727740-Paper.pdf .

Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. Journal of Machine Learning Research , 19(70):1-57, 2018. URL http://jmlr.org/papers/v19/18-188.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 arXiv:2208.12591 , 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://openreview.net/forum?id=thgItcQrJ4y .

Blake Woodworth, Suriya Gunasekar, Jason D. Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in 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. PMLR, 09-12 Jul 2020. URL https://proceedings.mlr.press/v125/woodworth20a.html .

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.1109/BigData50022.2020.9378171.

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://openreview.net/forum?id=ZsZM-4iMQkH .

Xingyu Zhu, Zixuan Wang, Xiang Wang, Mo Zhou, and Rong Ge. Understanding edge-of-stability training dynamics with a minimalist example. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=p7EagBsMAEO .

Contents

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) to minimize the objective function defined as

where Θ represents the parameters, ℓ : R → R is a loss function, f : R d → R is a neural network, and { ( x , y ) } denotes a single data point with x ∈ R d and y ∈ R . We also consider training on multiple data points { ( x i , y i ) } n i =1 with x i ∈ R d and y i ∈ R for each 1 ≤ i ≤ n , where we minimize the objective function

In our experiments, we primarily focus on the log -cosh loss function ℓ log -cosh ( p ) = log(cosh( p )) , but we also investigate the square root loss ℓ sqrt ( p ) = √ 1 + p 2 .

Model architecture. We train a fully-connected L -layer neural network, denoted as f ( · ; Θ ) : R d → R . The network is defined as follows:

where ϕ : R → R is an activation function applied entry-wise, W 1 ∈ R m × d , W l ∈ R m × m for 2 ≤ l ≤ L -1 , and w L ∈ R m . All L layers have the same width of m , and the biases of all layers are fixed to 0 . 4 We consider three activations: hyperbolic tangent ϕ ( t ) = tanh( t ) , exponential linear unit ϕ ( t ) = ELU ( t ) , and linear ϕ ( t ) = t .

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 α . In the plots presented throughout this section, we mark the initialization points with an 'x' to distinguish them from other points on the trajectories.

Canonical reparameterization. We plot GD trajectories after applying the canonical reparameterization introduced in Definition 2.1:

where η denotes the step size. For training on multiple data points, we employ the generalized canonical reparameterization as defined in Eq. (5):

where P : R n → R can represent the mean value or vector norms. Specifically, we mainly focus on the mean value P ( z ) = 1 n ∑ n i =1 z i , but we also examine vector norms such as P ( z ) = ∥ z ∥ 1 , P ( z ) = ∥ z ∥ 2 , and P ( z ) = ∥ z ∥ ∞ , where z = ( z 1 , z 2 , . . . , z n ) ∈ R n .

For large networks, explicitly calculating the ℓ 2 matrix norm ∥ ∥ ∑ n i =1 ( ∇ Θ f ( x i ; Θ )) ⊗ 2 ∥ ∥ 2 is infeasible. Therefore, we adopt a fast and efficient matrix-free method based on power iteration, as proposed by Yao et al. (2020). This method allows us to numerically compute the ℓ 2 norm of large-scale symmetric matrices.

(a)

α

= 0

.

5

(b)

α

= 1

.

0

(c)

α

= 2

.

0

(d)

α

= 4

.

0

50

3.5

X

X

1.0

1.0-

q=

tanh(p)

p

40-

3.0

X

0.9-

0.8-

p=0

20

30

o=d

tanh(p)

p

b

2.5

2.0

p=0

q=

tanh(p)

p

b

0.7

0.8

q0.6

1.5

0.6-

0.4

10-

1.0

0.5

b

tanh(p)

p

0.2

0

-1.0-0.8-0.6-0.4-0.2

0.0

1.25 1.00 0.75 0.50 -0.25 0.00

-2

-1

P=0

0

i

2

-5

0

5

p

p

p

p

A.2 Training on a single data point

Training data. We conduct experiments on a synthetic single data point ( x , y ) , where x = e 1 = (1 , 0 , . . . , 0) ∈ R d and y ∈ R . 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 at η = 0 . 01 .

The effect of initialization scale. We investigate the impact of the initialization scale α 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 ( α = 0 . 5 , 1 . 0 ) lead to trajectories in the gradient flow regime, while larger initialization scales ( α = 2 . 0 , 4 . 0 ) result in trajectories in the EoS regime. This behavior is primarily due to the initial value of q being smaller than 1 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 alignment 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 α = 4 . Additionally, in Figure 7, we examine the same setting but with different depth, training 10 -layer fully connected neural networks.

It is commonly observed that the alignment trend becomes more pronounced as the network width increases. In Figures 6 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 = 256 , 512 ) clearly demonstrate this behavior. These results indicate that network width plays a 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 a 3 -layer fully connected neural network with a width of m = 256 and an initialization scale of α = 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 a curve q = ℓ ′ ( p ) p .

(a)

m

= 64

(b)

m

= 128

(c)

m

= 256

(d)

m

= 512

1.2

q=

tanh(p)

p

1.0

q=

tanh(p)

P

1.0-

q=

(d)que)

p

1.0-

q=

(d)que)

p

1.0-

p=0

0.8

p=0

0.8-

p=0

0.8

p=0

0.8

b

90.6

q0.6

b

0.6-

0.6

0.4.

x

0.4-

0.4

0.4-

0.2

X

0.2

X

X

X

0.2

xx

X

0.2

-10

-5

0

5

-5.0

2.5

0.0

2.5

5.0

7.5

-5

0

-5

-10

0

10

p

p

p

p

(a)

m

= 64

(b)

m

= 128

(c)

m

= 256

(d)

m

= 512

1.4

q=

tanh(p)

p

1.2

q=

tanh(p)

p

1.0-

q=

tanh(p)

p

1.0-

q=

tanh(p)

p

1.2 -

0=d

1.0

0=d

0.8-

p=0

0.8-

p=0

1.0

0.8-

0.6-

0.6-

q0.8

do.6

q

q

0.6

0.4-

0.4-

0.4

0.4

0.2

0.2

0.2

0.2

0.0

-10

-5

0

5

10

0.0

-10

-5

0

5

0.0-

-10

-5

0

5

10

0.0-

-20

-10

0

10

20

p

p

p

p

A.3 Training on multiple synthetic data Points

Training data. We consider training on a synthetic dataset consisting of n data points, denoted as ( x i , y i ) n i =1 . The input vectors x i are sampled from a standard Gaussian distribution N ( 0 , I ) , where x i ∈ R d , and the corresponding target values y i are sampled from a Gaussian distribution N (0 , 1) , where y i ∈ R . Throughout our experiments in this subsection, we use a fixed data dimension of d = 10 and a step size of η = 0 . 01 .

The effect of function P . To investigate the impact of different choices of the function P , we train a 3 -layer fully connected neural network with tanh activation. The network has a width of m = 256 and is initialized with a scale of α = 4 . The training is performed on a dataset consisting of n = 10 data points. Figure 9 displays the trajectories of GD trajectories under the generalized canonical reparameterization defined in Eq. (5) for various choices of the function P . These choices include the mean, ℓ 1 norm, ℓ 2 norm, and ℓ ∞ 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 = ℓ ′ ( p ) p . However, when 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, a width of m = 256 , and an initialization scale

(a)

log

-

cosh

,

tanh

(b)

log

-

cosh

, ELU

(c)

log

-

cosh

, linear

(d) square-root,

tanh

(e) square-root, ELU

(f) square-root, linear

1.0

0.9

0.8

b

0.7

0.6

0.5

=b

tanh(p)

P

P=0

-2

-1

0

1

2

p

1.0

0.9

0.8

b

0.7

0.6

0.5

b

V1 +p²

p=0

0.4

-2

-1

0

1

2

1.0

0.9

0.8

0.7

q

0.6

0.5

0.4

0.3

q

tanh(p)

P

p=0

0.2

-4

-2

0

↓

2

4

p

1.0

0.9

0.8

0.7-

q

0.6

0.5

0.4

0.3

P=0

1 +p²

-4

-2

0

- 2

4

1.0

=b

tanh(p)

p=0

0.8-

q0.6

0.4-

0.2

-10

-5

0

5

P

1.0-

=b

√1+p²

p=0

0.8-

q0.6-

0.4-

0.2

-10

-5

0

5

p

p

p

(a)

P

(

z

) =

1

n

∑

n

i

=1

z

i

(b)

P

(

z

) =

∥

z

∥

1

(c)

P

(

z

) =

∥

z

∥

2

(d)

P

(

z

) =

∥

z

∥

∞

1.0-

=b

tanh(p)

p

1.0

=b

tanh(p)

p

1.0

q=

tanh(p)

P

1.0

=b

tanh(p)

p

0.8-

p=0

0.8

p=0

0.8

p=0

0.8

p=0

90.6

q

0.6

b

0.6

q0.6

0.4

0.4

0.4

0.4

0.2

0.2

0.2

X

0.2

0.0

2

0

2

4

6

0

20

40

60

0

5

10

15

2

4

6

8

10

p

p

p

p

of α = 4 . Additionally, in Figure 11, we investigate the same setting but with a 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 ) = 1 n ∑ n i =1 z i .

We observe a 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 q = ℓ ′ ( p ) p . However, as the number of data points n increases, it 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.

(a)

n

= 2

(b)

n

= 4

(c)

n

= 8

(d)

n

= 16

(e)

n

= 32

(f)

n

= 64

(g)

n

= 128

(h)

n

= 256

(i)

n

= 512

(j)

n

= 1024

1.0-

q=

tanh(p)

P

p=0

0.8-

90.6

0.4

0.2

-5.0

-2.5

0.0

2.5

5.0

p

0.9

1.0

p=0

=b

tanh(p)

d

0.8

0.7

9o.6

0.5

0.4

0.3

0.2

0

2

4

P

1.0

0.9

p=0

q=

tanh(p)

P

0.8

0.7-

90.6

0.5

0.4

0.3

X

0.2

-2

0

2

4

p

1.0

q=

tanh(p)

p

p=0

0.8

90.6-

0.4

X

0.2

2

0

2

4

6

p

1.0

0.9-

p=0

q=

tanh(p)

p

0.8

0.7

9o.6

0.5

0.4

0.3

x

0.2-

2

0

2

4

p

1.0

0.9-

p=0

q=

tanh(p)

p

0.8-

0.7

90.6

0.5

0.4

x

0.3

X

0.2-

2

0

2

4

p

0.9

1.0

q=

p=0

tanh(p)

p

0.8

0.7-

9o.6-

0.5

0.4

X

0.3

x

0.2

-2

0

2

4

p

0.9

1.0

q=

0=d

tanh(p)

p

0.8

0.7

9o.6

0.5

0.4

X

0.3

x

0.2

-2

0

2

4

P

0.9-

1.0

q=

p=0

tanh(p)

p

0.8

0.7

do.6

0.5

0.4

X

0.3

X

0.2

-2

0

2

4

P

1.0

0.9

0=d

=b

tanh(p)

p

0.8-

0.7

9o.6

0.5

0.4

X

0.3

0.2

-2

0

2

4

P

(a)

n

= 2

(b)

n

= 4

(c)

n

= 8

(d)

n

= 16

(e)

n

= 32

(f)

n

= 64

(g)

n

= 128

(h)

n

= 256

(i)

n

= 512

(j)

n

= 1024

1.6 -

1.4

q=

tanh(p)

p=0

1.2

1.0-

ao.8-

0.6-

0.4

0.2

0.0

-20

-10

0

1

10

p

1.0

q=

tanh(p)

p

0.8-

p=0

0.6-

b

0.4

0.2

X

X-X

-20

-10

0

10

p

1.0

q=

tanh(p)

P

0.8

p=0

0.6

b

0.4-

0.2

X

-20

-10

0

10

p

1.0

q=

tanh(p)

p

p=0

0.8

0.6

b

0.4

0.2

x

xx

-10

-5

0

5

10

p

1.0

q=

tanh(p)

p=0

0.8

0.6

b

0.4-

0.2

xx

-20

-10

0

10

p

1.0-

q=

tanh(p)

p

p=0

0.8

0.6

b

0.4 -

0.2 -

X

xx

-15

-10

-5

0

=:

5

10

p

1.0

q=

tanh(p)

P

p=0

0.8-

0.4-

0.2

X

-10

-5

0

-5

10

p

1.0

q=

tanh(p)

p

p=0

0.8

90.6.

0.4 -

0.2-

X

-10

5

0

·5

10

p

1.0

q=

tanh(p)

p

p=0

0.8

90.6.

0.4-

0.2

X

-10

-5

0

5

10

p

1.0

q=

tanh(p)

P

p=0

0.8

90.6.

0.4 -

0.2

X

X

-10

-5

0

5

10

p

A.4 Training on real-world dataset

Training data. In this subsection, we investigate a binary classification problem using a subset of the CIFAR10 image classification dataset. Our dataset consists of 50 samples, with 25 samples from class 0 (airplane) and 25 samples from class 1 (automobile). We assign a label of +1 to samples from class 0 and a label of -1 to samples from class 1 . This dataset was used in the experimental setup by Zhu et al. (2023).

Architectures. In Figure 4, we examine the training of two types of network architectures: (top row) fully-connected tanh network and (bottom row) convolutional tanh network. The PyTorch code for the fully-connected tanh network is provided as follows:

Note that we consider networks with bias. We use a step size of η = 0 . 01 for the fully-connected network and η = 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 P on the GD trajectories. We train a 3-layer fully connected neural network with ELU activation, a width of m = 256 , and an initialization scale of α = 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, ℓ 1 norm, ℓ 2 norm, and ℓ ∞ norm.

We observe that the GD trajectories exhibit alignment behavior, which is more pronounced when P is chosen to be the mean or ℓ 1 norm, but less evident for the ℓ ∞ norm. Unlike in Figure 9, the trajectories do not align on the curve q = ℓ ′ ( p ) p when P is selected as the mean P ( z ) = 1 n ∑ n i =1 z i .

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, we 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

(a)

P

(

z

) =

1

n

∑

n

i

=1

z

i

(b)

P

(

z

) =

∥

z

∥

1

(c)

P

(

z

) =

∥

z

∥

2

(d)

P

(

z

) =

∥

z

∥

∞

1.0

=b

tanh(p)

p

1.0

q=

tanh(p)

d

1.0

tanh(p)

p

1.0

q=

tanh(p)

0.8

p=0

0.8

p=0

0.8

p=0

0.8

p=0

0.6

0.6-

q0.6

q

b

q0.6

0.4-

0.4-

0.4-

0.4

0.2

0.2

0.2

0.2-

Xy

0.0

0.0-

-5.0-2.5

0.0

2.5

5.0

0

100

200

300

0

10

20

30

40

50

0

2

4

6

8

10

p

p

P

p

(a)

m

= 64

(b)

m

= 128

(c)

m

= 256

(d)

m

= 512

1.2 -

1.0

1.0

0.9

1.0-

0.9-

p=0

q=

tanh(p)

p

1.0-

0.9

p=0

=b

tanh(p)

p

0.8

0.8

0.8

b

0.8

q

0.7

0.6

q

0.7-

0.6

do.6

0.7-

6

0.6-

0.5

0.5

0.5

0.4

0.4

0.4

-2

-1

x

0

！

p=0

1

=b

tanh(p)

p

2

0.4

0.3

-2

p=0

q=

-1

tanh(p)

p

xi

0

1

2

0.3

0.2

2

0

A

2

4

0.3

0.2

-4

2

X

0

2

p

p

P

p

(a)

m

= 64

(b)

m

= 128

(c)

m

= 256

(d)

m

= 512

1.0

1.0

0.9

1.0

p=0

q=

tanh(p)

P

1.0-

p=0

q=

tanh(p)

p

0.8-

0.8

0.8-

0.8-

b

0.6

do.6

0.7

q0.6-

90.6

和

0.5

0.4

0.4

0.4-

0.4-

0.2-

p=0

=b

tanh(p)

p

0.2

0.3

P=0

b

tanh(p)

p

0.2

0.2

X

X

0

2

0

!

-2

4

-5.0-2.5

0.0

2.5

5.0

5.02.5

0.0

2.5

5.0

p

p

p

p

( 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 -1 to the dataset. However, in this particular experiment, we investigate the training process on a dataset with

zero labels. This means that all samples in the dataset are labeled as zero ( y i = 0 for all 1 ≤ i ≤ n ). Figure 15 visualizes the training of 3 -layer fully connected neural networks with tanh activation and an initialization scale of α = 1 . The network widths m are varied from { 256 , 512 , 1024 } . Interestingly, the GD trajectories align with the curve q = ℓ ′ ( p ) p , in 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 a future research direction, it would be intriguing to investigate why setting the labels as zero leads to alignment towards the curve q = ℓ ′ ( p ) p , which aligns with our theoretical findings in the single data point setting.

(a)

m

= 256

(b)

m

= 512

(c)

m

= 1024

1.0-

q=

tanh(p)

p=0

0.8

b

0.6-

0.4

0.2

Q

2

P

1.0

0.9

0.8

0.7

do.6

0.5

0.4

0.3

=b

tanh(p)

p

0.2

p=0

4

2

0

-2

4

p

1.0-

q=

tanh(p)

p=0

0.8

q0.6

0.4-

0.2

-5

0

5

p

B Proofs for the Two-Layer Fully-Connected 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 η ∈ (0 , 2 33 ) be a fixed step size and ℓ be a loss function satisfying Assumptions 4.1 and 4.2. Suppose that the initialization ( p 0 , q 0 ) satisfies | p 0 | ≤ 1 and q 0 ∈ ( 2 2 -η , min { 1 16 η , r (1) 2 η }) . Consider the GD trajectory characterized in Eq. (6) . Then, the GD iterations ( p t , q t ) converge to the point (0 , q ∗ ) such that

We note that in the given interval ( 2 2 -η , min { 1 16 η , r (1) 2 η }) , it is possible that the interval is empty, depending on the value of r (1) . However, this does not impact the correctness of the theorem.

Proof. By Proposition B.1, p t converges to 0 as t →∞ , and for all t ≥ 0 , we have

Since the sequence ( q t ) ∞ t =0 is monotonic increasing and bounded, it converges. Suppose that q t → q ∗ as t →∞ . Then, we can obtain the inequality

as desired.

Proposition B.1. Suppose that η ∈ (0 , 2 33 ) , | p 0 | ≤ 1 , and q 0 ∈ ( 2 2 -η , min { 1 16 η , r (1) 2 η }) . Then for any t ≥ 0 , we have

and

Proof. We give the proof by induction; namely, if

are satisfied for time steps 0 ≤ t ≤ k 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 t = 0 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 = k +1 .

By induction assumptions, q 0 ≤ q k ≤ 2 q 0 and | p k | ≤ 1 , so that we have

where we used r (1) ≤ r ( p k ) ≤ r (0) = 1 . Since q 0 ≤ r (1) 2 η ≤ r (1) 2 η 2 , we have

This implies that

which is the desired bound for | p k +1 | .

For any t ≤ k , we also have

where we used the induction assumptions to deduce 4 q 0 ≥ 2 q t -r ( p t ) ≥ 2 q 0 -1 ≥ 4 2 -η -1 > 0 . This gives q k +1 ≥ q k ≥ q 0 . Furthermore, note that 4 η 2 p 2 t q 0 ≤ 4 η 2 · 1 16 η ≤ 1 4 , so the ratio q t q t +1 ∈ [ 3 4 , 1] . From this, we have

where the second inequality holds since | log(1 + z ) | ≤ 2 | z | if | z | ≤ 1 2 . Moreover, this implies that

Since q 0 ∈ ( 2 2 -η , r (1) 2 η ) , we have

Therefore, since q 0 ≤ 1 16 η , we can conclude that

the desired bounds for q k +1 .

B.2 Proof of Theorem 4.2

In this subsection, we prove Theorem 4.2. From here onwards, we use the following notation:

All the lemmas in this subsection are stated in the context of Theorem 4.2.

Lemma B.2. Suppose that the initialization ( p 0 , q 0 ) satisfies | p 0 | ≤ 1 and q 0 ∈ ( c 0 , 1 -δ ) . Then for any t ≥ 0 such that q t ≤ 1 , it holds that

Proof. We prove by induction. We assume that for some t ≥ 0 , it holds that | p t | ≤ 4 and 1 2 ≤ q t ≤ 1 . We will prove that | p t +1 | ≤ 4 and 1 2 ≤ q t ≤ q t +1 ≤ (1 + O ( η 2 )) q t . For the base case, | p 0 | ≤ 1 ≤ 4 and 1 2 ≤ c 0 < q t ≤ 1 holds by the assumptions on the initialization. Now suppose that for some t ≥ 0 , it holds that | p t | ≤ 4 and 1 2 ≤ q t ≤ 1 . Then for small step size η , we have

Consequently, by Eq. (6),

where we used 1 -Lipshitzness of ℓ . Moreover,

where we used q t ∈ [ 1 2 , 1] and | p t r ( p t ) | = | ℓ ′ ( p t ) | ≤ 1 from 1 -Lipschitzness of ℓ . Hence, q t ≤ q t +1 ≤ (1 + O ( η 2 )) q t , as desired.

Lemma B.2 implies that p t is bounded by a constant throughout the iterations, and q t monotonically increases slowly, where the increment for each step is O ( η 2 ) . Hence, there exists a time step T = Ω( δη -2 ) = Ω δ ( η -2 ) such that for any t ≤ T , it holds that q t ≤ 1 -δ 2 . Through out this subsection, we focus on these T early time steps. Note that for all 0 ≤ t ≤ T , it holds that q t ∈ ( c 0 , 1 -δ 2 ) .

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 p t is bounded and q t monotonically increases slowly, with an increment of O ( η 2 ) per step. Lemma 2.1 shows that the map f q t ( p ) = ( 1 -2 r ( p t ) q t ) p t has a stable 2 -period orbit {± ˆ r ( q t ) } when q t ∈ (0 , 1) . Consequently, when q t is treated as a fixed value, p t converges to the orbit {± ˆ r ( q t ) } , leading to s t converging to 1 . In the (early) short-term dynamics, q t is nearly fixed for small step size η , and hence s t converges to 1 . In long-term dynamics perspective, q t gradually increases and at the same time, s t stays near the value 1 . In Theorem 4.2, we prove that it takes only t a = O δ,ℓ (log( η -1 )) time steps for s t to converge close to 1 (Phase I, t ≤ t a ), and after that, s t stay close to 1 for the remaining iterations (Phase II, t > t a ).

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 a time step t 0 where s t 0 becomes smaller or equal to 2 2 -r (1) , which is smaller than 2(1 + η 2 ) -1 . Lemma B.5 demonstrates that if s t is smaller than 2(1 + η 2 ) -1 and | p t | ≥ ˆ r (1 -δ 4 ) , then | s t -1 | decreases exponentially. For the case where | p t | < ˆ r (1 -δ 4 ) , Lemma B.6 proves that | p t | increases at an exponential rate. Moreover, Lemma B.4 shows that if s t < 1 at some time step, then s t +1 is upper bounded by 1 + O ( η 2 ) . Combining these findings, Proposition B.7 establishes that in the early phase of training, there exists a time step t ∗ a such that s t ∗ a = 1+ O δ,ℓ ( η 2 ) . Lastly, Lemma B.8 demonstrates that if s t = 1+ O δ,ℓ ( η 2 ) , then | s t -1 -h ( p t ) η 2 | decreases exponentially.

Now we prove Theorem 4.2, starting with the lemma below.

Lemma B.3. There exists a time step t 0 = O δ,ℓ (1) such that s t 0 ≤ 2 2 -r (1) .

Proof. We start by proving the following statement: for any 0 ≤ t ≤ T , if 2 2 -r (1) < s t < 2 r (1) -1 , then s t +1 < 2 r (1) -1 and | p t +1 | ≤ (1 -r (1) 2 ) | p t | . Suppose that 2 2 -r (1) < s t < 2 r (1) -1 . Then from Eq. (6), it holds that

for small step size η . Hence, | p t +1 | ≤ (1 -r (1) 2 ) | p t | . Now we prove s t +1 < 2 r (1) -1 . Assume the contrary that s t +1 ≥ 2 r (1) -1 . Then, r ( p t +1 ) = q t +1 s t +1 < q t +1 < 1 -δ 2 so that | p t +1 | ≥ ˆ r (1 -δ 2 ) . By Mean Value Theorem, there exists p ∗ t ∈ ( | p t +1 | , | p t | ) such that (recall that r ′ ( p ) r ( p ) 2 < 0 for p > 0 )

where we used Assumption 4.3 (i) and ˆ r (1 -δ 2 ) ≤ | p t +1 | ≤ (1 -r (1) 2 ) | p t | ≤ | p t | . Consequently,

for small step size η . This gives a contradiction to our assumption that s t +1 ≥ 2 r (1) -1 . Hence, we can conclude that s t +1 < 2 r (1) -1 , as desired.

We proved that for any 0 ≤ t ≤ T , if 2 2 -r (1) < s t < 2 r (1) -1 , it holds that s t +1 < 2 r (1) -1 and | p t +1 | ≤ (1 -r (1) 2 ) | p t | . At initialization, | p 0 | ≤ 1 and q 0 < 1 , so that s 0 < r (1) -1 . If s 0 ≤ 2 2 -r (1) , then t 0 = 0 is the desired time step. Suppose that s 0 > 2 2 -r (1) . Then, we have s 1 < 2 r (1) -1 and | p 1 | ≤ (1 -r (1) 2 ) | p 0 | ≤ 1 -r (1) 2 . Then we have either s 1 ≤ 2 2 -r (1) , or 2 2 -r (1) < s 1 < 2 r (1) -1 . In the previous case, t 0 = 1 is the desired time step. In the latter case, we can repeat the same argument and obtain s 2 < 2 r (1) -1 and | p 2 | ≤ (1 -r (1) 2 ) 2 . By inductively repeating the same argument, we can obtain a time step t 0 ≤ log(ˆ r (1 -δ 2 )) / log(1 -r (1) 2 ) = O δ,ℓ (1) such that either s t 0 ≤ 2 2 -r (1) , or | p t 0 | ≤ ˆ r (1 -δ 2 ) . In the latter case, r ( p t 0 ) ≥ 1 -δ 2 > q t 0 , and hence s t 0 < 1 < 2 2 -r (1) . Therefore, t 0 = O δ,ℓ (1) is the desired time step satisfying s t 0 ≤ 2 2 -r (1) .

According to Lemma B.3, there exists a time step t 0 = O δ,ℓ (1) such that s t 0 ≤ 2 2 -r (1) < 2(1+ η 2 ) -1 for small step size η . Now we prove the lemma below.

Lemma B.4. Suppose that s t ≤ 1 . Then, it holds that s t +1 ≤ 1 + O ( η 2 ) .

Proof. For any p ∈ (0 , ˆ r ( q t 2 )) , we have r ( p ) ≥ q t 2 so that | f q t ( p ) | = ( -1 + 2 r ( p ) q t ) p . Hence,

for any p ∈ (0 , ˆ r ( q t 2 )) . By Assumption 4.3 (ii) and convexity of ℓ , both r ( p ) and 1 + pr ′ ( p ) r ( p ) = ℓ ′′ ( p ) r ( p ) are positive, decreasing function on (0 , ˆ r ( q t 2 )) . Consequently, ∂ ∂p | f q t ( p ) | is a decreasing function on (0 , ˆ r ( q t 2 )) .

Now note that q t 2 < q t < 1 , which means ˆ r (1) = 0 < ˆ r ( q t ) < ˆ r ( q t 2 ) by the definition of ˆ r . Note that ∂ ∂p | f q t ( p ) | at p = ˆ r ( q t ) evaluates to

where the first inequality used Assumption 4.3 (ii) and ˆ r ( q t ) < ˆ r ( c 0 ) , which comes from q t > c 0 := max { r ( z 0 ) , 1 2 } . The second inequality holds because q t > c 0 ≥ r ( z 0 ) where z 0 := sup z { zr ′ ( z ) r ( z ) ≥ -1 2 } , from the statement of Theorem 4.2.

Therefore, since ∂ ∂p | f q t ( p ) | is decreasing on (0 , ˆ r ( q t 2 )) and is nonnegative at ˆ r ( q t ) , for any p ∈ (0 , ˆ r ( q t )) , it holds that ∂ ∂p | f q t ( p ) | ≥ 0 . In other words, | f q t ( p ) | is an increasing function on (0 , ˆ r ( q t )) . Since 0 ≤ s t ≤ 1 , we have | p t | ≤ ˆ r ( q t ) and it holds that

Therefore, with this inequality and Lemma B.2, we can conclude that

Using Lemma B.4, we prove the following lemma.

Lemma B.5. For any 0 ≤ t ≤ T , if s t < 2(1 + η 2 ) -1 , | s t -1 | > η 2 2 , and r ( p t ) ≤ 1 -δ 4 , then | s t +1 -1 | ≤ (1 -d ) | s t -1 | + O ( η 2 ) , where d ∈ (0 , 1 2 ] is a constant which depends on δ and ℓ .

Proof. By Eq. (6) and 1 -Lipschitzness of ℓ ,

so that p t and p t +1 have opposite signs. By Mean Value Theorem, there exists θ t between -1 and (1 -2 s t + η 2 p 2 t r ( p t ) 2 ) satisfying

where the last equality used the fact that p t and θ t p t have opposite signs and r ′ ( z ) and z have opposite signs. Note that | θ t p t | is between | p t | and | p t +1 | . Consequently, the value | r ′ ( θ t p t ) | r ( θ t p t ) 2 is between | r ′ ( p t ) | r ( p t ) 2 and | r ′ ( p t +1 ) | r ( p t +1 ) 2 by Assumption 4.3 (i). We will prove the current lemma based on Eq. (8). We divide into following three cases: (1) s t ≥ 1 and s t +1 ≥ 1 , (2) s t ≥ 1 and s t < 1 , and (3) s t < 1 .

Case 1. Suppose that s t ≥ 1 and s t +1 ≥ 1 . Here, we have | p t | ≥ ˆ r ( q t ) ≥ ˆ r (1 -δ 2 ) and similarly | p t +1 | ≥ ˆ r (1 -δ 2 ) . By Assumption 4.3 (i), | r ′ ( θ t p t ) | r ( θ t p t ) 2 ≥ | r ′ (ˆ r (1 -δ 2 )) | (1 -δ 2 ) 2 . Hence, Eq. (8) gives

where we used q t > c 0 ≥ 1 2 and s t < 2(1 + η 2 ) -1 < 2 . Therefore, we can obtain the following inequality:

Case 2. Suppose that s t ≥ 1 and s t +1 < 1 . Here, we have r ( p t +1 ) > q t +1 ≥ q t ≥ r ( p t ) , so that | p t +1 | < | p t | . Consequently, | r ′ ( θ t p t ) | r ( θ t p t ) 2 ≤ | r ′ ( p t ) | r ( p t ) 2 by Assumption 4.3 (i). Hence, we can deduce from Eq. (8) that

where we used | p t r ′ ( p t ) | ≤ r ( p t ) since 1 + p t r ′ ( p t ) r ( p t ) = ℓ ′′ ( p t ) r ( p t ) > 0 and | p t | ≤ 4 by Lemma B.2. Consequently, by Lemma B.2 ( q t ≤ q t +1 ) and Assumption 4.3 (ii),

Note that 1 > 1+ 4 r ′ (4) r (4) = ℓ ′′ (4) r (4) > 0 holds by convexity of ℓ . Therefore, we can obtain the following inequality:

where -1 < 1 + 8 r ′ (4) r (4) < 1 .

Case 3. Suppose that s t < 1 . By Lemma B.4, it holds that s t +1 ≤ 1 + O ( η 2 ) . Moreover, we assumed r ( p t ) ≤ 1 -δ 4 , so that | p t | ≥ ˆ r (1 -δ 4 ) . We also have

where we used the assumption | s t -1 | > η 2 2 , and | pr ( p ) | = | ℓ ′ ( p ) | ≤ 1 due to 1 -Lipschitzness of ℓ . Consequently, by Assumption 4.3 (i), it holds that | r ′ ( θ t p t ) | r ( θ t p t ) 2 ≥ | r ′ (ˆ r (1 -δ 4 )) | (1 -δ 4 ) 2 . Hence, by Eq. (8), we have

and hence, by Lemma B.2 ( q t ≤ q t +1 ) and q t > c 0 ≥ 1 2 , we get

Therefore, we can obtain the following inequality:

where we used Lemma B.4 to obtain the first inequality.

Combining the three cases, we can finally conclude that if we choose

Lemma B.5 implies that if s t < 2(1 + η 2 ) -1 and | p t | ≥ ˆ r (1 -δ 4 ) , then | s t -1 | exponentially decreases. We prove Lemma B.6 to handle the regime | p t | < ˆ r (1 -δ 4 ) , which is stated below.

Lemma B.6. For any 0 ≤ t ≤ T , if r ( p t ) ≥ 1 -δ 4 , it holds that

Proof. If r ( p t ) ≥ 1 -δ 4 , then s t = q t r ( p t ) < 1 -δ 2 1 -δ 4 = 4 -2 δ 4 -δ , where we used q t < 1 -δ 2 for any 0 ≤ t ≤ T . Consequently,

for small step size η .

Now we prove Proposition B.7, which proves that s t reaches close to 1 with error bound of O ( η 2 ) . Proposition B.7. There exists a time step t ∗ = O (log( η -1 )) satisfying a δ,ℓ

Proof. By Lemma B.3, there exists a time step t 0 = O δ,ℓ (1) such that s t 0 ≤ 2 2 -r (1) . Here, we divide into two possible cases: (1) s t 0 < 1 , and (2) 1 ≤ s t 0 ≤ 2 2 -r (1) .

Case 1. Suppose that s t 0 < 1 . By Lemma B.6, if r ( p t 0 ) ≥ 1 -δ 4 (or equivalently, | p t 0 | ≤ ˆ r (1 -δ 4 ) ), then there exists a time step t 1 ≤ t 0 +log( ˆ r (1 -δ 4 ) | p t 0 | ) / log( 4 4 -δ ) = O δ,ℓ (1) such that | p t 1 | ≥ ˆ r (1 -δ 4 ) . We denote the first time step satisfying | p t 1 | ≥ ˆ r (1 -δ 4 ) and t 1 ≥ t 0 by t 1 = O δ,ℓ (1) . By Lemma B.4, it holds that s t 1 ≤ 1+ O ( η 2 ) since s t 1 -1 < 1 . Consequently, if s t 1 ≥ 1 -η 2 2 , then | s t 1 -1 | ≤ O ( η 2 ) so that t ∗ a = t 1 is the desired time step. Hence, it suffices to consider the case when s t 1 < 1 -η 2 2 . Here, we can apply Lemma B.5 which implies that

where d is a constant which depends on δ and ℓ . Then, there are two possible cases: either | s t 1 -1 | ≤ O ( η 2 d -1 ) , or | s t 1 +1 -1 | ≤ (1 -d 2 ) | s t 1 -1 | . It suffices to consider the latter case, suppose that | s t 1 +1 -1 | ≤ (1 -d 2 ) | s t 1 -1 | . Since we are considering the case s t 1 < 1 -η 2 2 , again by Lemma B.4, we have s t 1 +1 ≤ 1 + O ( η 2 ) . Since | p t 1 +1 p t 1 | = 2 s t 1 -1 -O ( η 2 ) , | p t 1 +1 | ≥ | p t 1 | ≥ ˆ r (1 -δ 4 ) must be satisfied unless s t 1 = 1 + O ( η 2 ) already holds. If s t 1 +1 ≥ 1 -η 2 2 , then | s t 1 +1 -1 | ≤ O ( η 2 ) so that t ∗ a = t 1 +1 is the desired time step; if not, we can again apply Lemma B.5 and repeat the analogous argument. Hence, there exists a time step t 2 ≤ t 1 +log( η 2 1 -s t 1 ) / log(1 -d 2 ) = O δ,ℓ (log( η -1 )) , such that | s t 2 -1 | ≤ O ( η 2 d -1 ) = O δ,ℓ ( η 2 ) .

Case 2. Suppose that 1 ≤ s t 0 ≤ 2 2 -r (1) . Then, r ( p t 0 ) ≤ q t 0 ≤ 1 -δ 2 , so we can apply Lemma B.5. There are two possible cases: either | s t 0 +1 -1 | ≤ O ( η 2 d -1 ) = O δ,ℓ ( η 2 ) , or | s t 0 +1 -1 | ≤ (1 -d 2 ) | s t 0 -1 | . It suffices to consider the latter case. If s t 0 +1 ≥ 1 , we can again apply Lemma B.5 and repeat the analogous argument. Hence, we can obtain a time step t ′ 0 ≤ t 0 +log( η 2 1 -s t 0 ) / log(1 -d 2 ) = O δ,ℓ (log( η -1 )) such that either s t ′ 0 < 1 or | s t ′ 0 -1 | = O δ,ℓ ( η 2 ) is satisfied. If s t ′ 0 < 1 , we proved in Case 1 that there exists a time step t ′ 2 = t ′ 0 + O δ,ℓ (log( η -1 )) such that | s t ′ 2 -1 | ≤ O δ,ℓ ( η 2 ) , and this is the desired bound.

Now we carefully handle the error term O ( η 2 ) obtained in Proposition B.7 and a provide tighter bound on s t by proving Lemma B.8 stated below.

Lemma B.8. If | s t -1 | = O δ,ℓ ( η 2 ) , then it holds that

̸

where h ( p ) := -1 2 ( pr ( p ) 3 r ′ ( p ) + p 2 r ( p ) 2 ) for p = 0 and h ( p ) := -1 2 r ′′ (0) for p = 0 .

Proof. Suppose that s t = 1 + O δ,ℓ ( η 2 ) . Then, | p t +1 | = ∣ ∣ ∣ 1 -2 s t + η 2 p 2 t r ( p t ) 2 ∣ ∣ ∣ · | p t | = (1 + O δ,ℓ ( η 2 )) | p t | . By Eq. (8) proved in Lemma B.5, there exists ϵ t = O δ,ℓ ( η 2 ) which satisfies the following:

where we used the Taylor expansion on r ′ ( p ) r ( p ) 2 with the fact that d dp ( r ′ ( p ) r ( p ) 2 ) is bounded on [ -4 , 4] and that | p t | ≤ 4 to obtain the second equality. Note that q t +1 = (1 -η 2 p 2 t r ( p t )(2 q t -r ( p t ))) -1 q t by Eq. (6). Consequently,

Here, since s t = 1 + O δ,ℓ ( η 2 ) , we can rewrite

which results in

Note that h is even, and twice continuously differentiable function by Lemma B.9. Consequently, h ′ (0) = 0 and h ′ ( p ) = O ℓ ( p ) , since h ′′ is bounded on closed interval. Consequently, h ( p t +1 ) = h ((1 + O δ,ℓ ( η 2 )) p t ) = h ( p t ) + O δ,ℓ ( η 2 p 2 t ) . Hence, we can obtain the following:

Note that r ( p t ) = (1 + O δ,ℓ ( η 2 )) q t ≥ (1 + O δ,ℓ ( η 2 )) q 0 ≥ c 0 ≥ r ( z 0 ) for small step size η , where z 0 = sup { zr ′ ( z ) r ( z ) ≥ -1 2 } . Consequently, it holds that 1 + 2 p t r ′ ( p t ) r ( p t ) ≥ 0 . Therefore, we have the desired inequality:

We now provide the proof of Theorem 4.2, restated below for the sake of readability.

̸

Theorem 4.2 (EoS regime, Phase I) . Let η be a small enough step size and ℓ be a loss function satisfying Assumptions 4.1, 4.2, and 4.3. Let z 0 := sup z { zr ′ ( z ) r ( z ) ≥ -1 2 } and c 0 := max { r ( z 0 ) , 1 2 } . Let δ ∈ (0 , 1 -c 0 ) be any given constant. Suppose that the initialization ( p 0 , q 0 ) satisfies | p 0 | ≤ 1 and q 0 ∈ ( c 0 , 1 -δ ) . Consider the reparameterized GD trajectory characterized in Eq. (6) . We assume that for all t ≥ 0 such that q t < 1 , we have p t = 0 . Then, there exists a time step t a = O δ,ℓ (log( η -1 )) , such that for any t ≥ t a ,

̸

where h ( p ) := -1 2 ( pr ( p ) 3 r ′ ( p ) + p 2 r ( p ) 2 ) for p = 0 and h ( p ) := -1 2 r ′′ (0) for p = 0 .

Proof of Theorem 4.2. By Proposition B.7, there exists a time step t ∗ a = O δ,ℓ (log( η -1 )) which satisfies:

By Lemma B.8, there exists a constant D > 0 which depends on δ, ℓ such that if | s t -1 | = O δ,ℓ ( η 2 ) , then

Hence, if | s t -1 | = O δ,ℓ ( η 2 ) and | s t -1 -h ( p t ) η 2 | ≥ ( -p t r ( p t ) r ′ ( p t ) ) Dη 4 , then

For any t ≤ T , we have q t < 1 -δ 2 so that if | s t -1 | = O δ,ℓ ( η 2 ) , then r ( p t ) ≤ (1 + O δ,ℓ ( η 2 )) q t < 1 -δ 4 for small step size η . From Eq. (11) with t = t ∗ a , we have either

or where we used Assumption 4.3 (ii) and | p t | > ˆ r (1 -δ 4 ) . In the later case, | s t ∗ a +1 -1 | = O δ,ℓ ( η 2 ) continues to hold and we can again use Eq. (11) with t = t ∗ a + 1 . By repeating the analogous arguments, we can obtain the time step

which satisfies: either

or

where we used | p t | ≤ 4 from Lemma B.2 and -zr ( z ) r ′ ( z ) ≥ -1 r ′′ (0) for any z by Assumption 4.3 (iii).

By Eq. (10), if | s t -1 -h ( p t ) η 2 | ≤ ( -4 r (4) r ′ (4) ) Dη 4 is satisfied for any time step t , then

by | p t | ≤ 4 from Lemma B.2 and Assumption 4.3 (iii).

Hence, by induction, we have the desired bound as following: for any t ≥ t a ,

by | p t | ≤ 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 h defined in Theorem 4.2, given by

̸

Then, h is a positive, even, and bounded twice continuously differentiable function.

̸

Proof. It is clear that h is even. We first prove that h is positive. For any p = 0 , it holds that

since pr ( p ) r ′ ( p ) < 0 and 1 + pr ′ ( p ) r ( p ) = ℓ ′′ ( p ) r ( p ) > 0 by Assumption 4.2 and convexity of ℓ . The function h is continuous since lim p → 0 h ( p ) = h (0) . Continuous function on a compact domain is bounded, so h is bounded on the closed interval [ -1 , 1] . We can rewrite h as

Note that p 2 r ( p ) 2 = ℓ ′ ( p ) 2 ≤ 1 , and ( -r ( p ) pr ′ ( p ) -1 ) is positive, decreasing function on p > 0 by Assumption 4.3 (ii). Hence, h is bounded on [1 , ∞ ) . Since h is even, h is bounded on ( -∞ , 1] . Therefore, h is a bounded function on R .

We finally prove that h is twice continuously differentiable. Since r is even and C 4 on R , we can check that

̸

̸

Moreover, for any p = 0 ,

and

Since lim p → 0 h ′′ ( p ) = h ′′ (0) , we can conclude that h is a twice continuously differentiable function.

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 t b = Ω((1 -q 0 ) η -2 ) such that q t b ≤ 1 and q t > 1 for any t > t b . Moreover, the GD iterates ( p t , q t ) converge to the point (0 , q ∗ ) such that

Proof. We first prove that there exists a time step t b ≥ 0 such that q t b > 1 . Assume the contrary that q t ≤ 1 for all t ≥ 0 . Let t a be the time step obtained in Theorem 4.2. Then for any t ≥ t a , we have

for small step size η . The function g ( p ) := r ( p ) -1+ h ( p ) η 2 2 is even, continuous, and has the function value g (0) = η 2 4 | r ′′ (0) | > 0 . Consequently, there exists a positive constant ϵ > 0 such that g ( p ) > 0 for all p ∈ ( -ϵ, ϵ ) . Then, we have | p t | ≥ ϵ for all t ≥ t a , since g ( p t ) ≤ 0 . Moreover, s t ≥ 3 4 for any t ≥ t a by Theorem 4.2 for small step size η . This implies that for any t ≥ t a ,

so q t grows exponentially, which results in the existence of a time step t ′ b ≥ t a such that q t ′ b > 1 , a contradiction.

Therefore, there exists a time step t b such that q t b ≤ 1 and q t > 1 for any t > t b , i.e., q t jumps across the value 1 . This holds since the sequence ( q t ) is monotonically increasing. For any t ≤ t b , we have q t +1 ≤ q t + O ( η 2 ) by Lemma B.2, and this implies that t b ≥ Ω((1 -q 0 ) η -2 ) , as desired.

Lastly, we prove the convergence of GD iterates ( p t , q t ) . Let t > t b be given. Then, q t ≥ q t b +1 > 1 and it holds that

Hence, | p t | is exponentially decreasing for t > t b . Therefore, p t converges to 0 as t →∞ . Since the sequence ( q t ) ∞ t =0 is monotonically increasing and bounded (due to Theorem 4.2), it converges. Suppose that ( p t , q t ) converges to the point (0 , q ∗ ) . By Theorem 4.2, we can conclude that

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 ⊤ Ux , stated below.

Lemma B.10. For any Θ = ( U , v ) with U ∈ R m × d , v ∈ R m , and x ∈ R d with ∥ x ∥ 2 = 1 , the following equality holds:

Moreover, if λ is an eigenvalue of ∇ 2 ( U , v ) ( v ⊤ Ux ) , then -λ is also an eigenvalue of ∇ 2 ( U , v ) ( v ⊤ Ux ) .

Proof. We first define the notations. We use the operator ⊗ to represent tensor product, or Kronecker product between matrices. For example, for any given two matrices A = ( a ij ) ∈ R m × n and B , we define A ⊗ B by

We use 0 m × n to denote a m by n matrix with all entries filled with zero, and I n denotes n by n identity matrix. Now we provide the proof of the original problem.

Let U = ( u ij ) ∈ R m × d , v = ( v i ) ∈ R m , and x = ( x j ) ∈ R d be given. Then,

We vectorize the parameter Θ = ( U , v ) by ( v 1 , . . . , v m , U 11 , . . . , U m 1 , . . . , U 1 d , . . . , U md ) ∈ R m + md . Then, we can represent the Hessian as

For any given c ∈ R and y ∈ R d with c 2 + ∥ y ∥ 2 2 = 1 , we have

Hence, by definition of matrix operator norm, we have

Therefore, we can conclude that

Now suppose that λ is an eigenvalue of ∇ 2 ( U , v ) ( v ⊤ Ux ) . We note that for any given matrices A and B , if λ a is an eigenvalue of A with the corresponding eigenvector u a and λ b is an eigenvalue of B with the corresponding eigenvector u b , then λ a λ b is an eigenvalue of A ⊗ B with the corresponding eigenvector u a ⊗ u b . Moreover, any eigenvalue of A ⊗ B arises as such a product of eigenvalues of A and B . Hence, using Eq. (12), we have

We denote the corresponding eigenvector by ( c, y ⊤ ) ⊤ where c ∈ R and y ∈ R d , i.e., it holds that

Consequently, we have

and this implies that

Therefore, by Eq. (12), -λ is an eigenvalue of ∇ 2 ( U , v ) ( v ⊤ Ux ) .

Using Lemma B.10, we prove an important bound on the sharpness value provided by the Proposition B.11 stated below.

Proposition B.11. For any Θ = ( U , v ) with U ∈ R m × d , v ∈ R m , and x ∈ R d with ∥ x ∥ 2 = 1 , the following bound holds:

Proof. The loss Hessian at Θ = ( U , v ) can be characterized as:

We first prove that λ max ( ∇ 2 Θ L ( Θ )) = ∥ ∥ ∇ 2 Θ L ( Θ ) ∥ ∥ 2 . Note that the largest absolute value of the eigenvalue of a symmetric matrix equals to its spectral norm. Hence, ∥ ∥ ∇ 2 Θ L ( Θ ) ∥ ∥ 2 = max { λ max ( ∇ 2 Θ L ( Θ )) , -λ min ( ∇ 2 Θ L ( Θ )) } , so it suffices to prove that λ max ( ∇ 2 Θ L ( Θ )) ≥ -λ min ( ∇ 2 Θ L ( Θ )) . Let w denote the eigenvector of ∇ 2 Θ L ( Θ ) corresponding to the smallest eigenvalue λ min ( ∇ 2 Θ L ( Θ )) with ∥ w ∥ 2 = 1 . Then, using Eq. (13), we have

where we used Lemma B.10 to obtain the last inequality. Note that the matrix ℓ ′′ ( v ⊤ Ux ) ( ∇ Θ ( v ⊤ Ux ) ) ⊗ 2 is PSD, so that λ max ( ∇ 2 Θ L ( Θ )) ≥ λ max ( ℓ ′ ( v ⊤ Ux ) ∇ 2 Θ ( v ⊤ Ux )) = | ℓ ′ ( v ⊤ Ux ) |∥∇ 2 Θ ( v ⊤ Ux ) ∥ 2 ≥ -λ min ( ∇ 2 Θ L ( Θ )) . Therefore, λ max ( ∇ 2 Θ L ( Θ )) = ∥ ∥ ∇ 2 Θ L ( Θ ) ∥ ∥ 2 . Now, we have the following triangle inequality:

where the last inequality holds by Lemma B.10 and 1 -Lipschitzness of ℓ .

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 t a denote the obtained iteration. Define the function ˜ λ : R > 0 → R given by

Then, the sequence ( ˜ λ ( q t ) ) ∞ t =0 is monotonically increasing. Moreover, for any t ≥ t a , the sharpness at GD iterate Θ t closely follows the sequence ( ˜ λ ( q t ) ) ∞ t =0 by satisfying

Proof. By Proposition B.11, we can bound the sharpness λ max ( Θ t ) at time step t by

Since ℓ ′′ ( z ) = r ( z ) + zr ′ ( z ) , we can rewrite as following:

By Theorem 4.2 and since h is a bounded function by Lemma B.9, we have s t = 1 + O ℓ ( η 2 ) for any t ≥ t a . Consequently, | s -1 t -1 | = O ℓ ( η 2 ) and | r ( p t ) -q t | = O ℓ ( η 2 ) . Moreover, for any 0 < q < 1 ,

so that

Therefore, d dq ( ˆ r ( q ) r ′ (ˆ r ( q )) q ) is bounded on [ 1 4 , 1) and Taylor's theorem gives

for any time step t with q t < 1 . Hence, if q t < 1 , we have the following bound:

where we used p t r ′ ( p t ) q t = p t r ′ ( p t ) r ( p t ) (1+ O ℓ ( η 2 )) = p t r ′ ( p t ) r ( p t ) + O ℓ ( η 2 ) , since 1+ p t r ′ ( p t ) r ( p t ) = ℓ ′′ ( p t ) r ( p t ) > 0 implies | p t r ′ ( p t ) | ≤ r ( p t ) ≤ 1 . Now let t be any given time step with q t ≥ 1 . Then, r ( p t ) = 1 -O ℓ ( η 2 ) , and since r ( z ) = 1 + r ′′ (0) z 2 + O ℓ ( z 4 ) for small z , we have | p t | = O ℓ ( η ) . Hence,

for any t with q t ≥ 1 . By Eqs. (14), (15), and (16), we can conclude that for any t ≥ t a , we have

Finally, we can easily check that the sequence ( ˜ λ ( q t )) ∞ t =0 is monotonically increasing, since z ↦→ zr ′ ( z ) r ( z ) is a decreasing function by Assumption 4.3 (ii) and the sequence ( q t ) is monotonically increasing.

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 a two-dimensional function L ( x, y ) := 1 2 ( ϕ ( x ) y ) 2 , where x , y are scalars and ϕ is a nonlinear activation satisfying Assumption 5.1. We consider the reparameterization given by Definition 5.2, which is ( p, q ) := ( x, 2 ηy 2 ) .

We emphasize that the results we present in this subsection closely mirror those of the Section 4. In particular,

Assumption C.1 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 are 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 in Theorem 5.1.

Inspired by Lemma 2.1, we have an additional assumption on ϕ as below.

Assumption C.1. Let r be a function defined by r ( z ) := ϕ ( z ) ϕ ′ ( z ) z for z = 0 and r (0) := 1 . The function r satisfies Assumption 2.4.

̸

In contrast to the function r defined in Section 4, the expression 1 + pr ′ ( p ) r ( p ) can be negative, which implies that the constant c defined in Lemma 2.1 is positive. As a result, the dynamics of p t may exhibit a period4 (or higher) oscillation or even chaotic behavior (as illustrated in Figure 1b).

We first state our results on the gradient flow regime.

Theorem C.1 (gradient flow regime) . Let η ∈ (0 , r (1) 2( r (1)+2) ) be a fixed step size and ϕ be a sigmoidal function satisfying Assumptions 5.1 and C.1. Suppose that the initialization ( p 0 , q 0 ) satisfies | p 0 | ≤ 1 and q 0 ∈ ( 1 1 -2 η , r (1) 4 η ) . Consider the reparameterized GD trajectory characterized in Eq. (7) . Then, the GD iterations ( p t , q t ) converge to the point (0 , q ∗ ) such that

Theorem C.1 implies that in gradient flow regime, GD with initialization ( x 0 , y 0 ) and step size η converges to (0 , y ∗ ) which has the sharpness bounded by:

Now we provide our results on the EoS regime with an additional assumption below.

Assumption C.2. Let r be a function defined in Assumption C.1. Then r is C 4 on R and satisfies:

(iii) z → zr ( z ) r ′ ( z ) is decreasing on z > 0 and increasing on z < 0 , and

Note that the function r that arise from the activation ϕ = tanh satisfies Assumptions 5.1, C.1, and C.2.

̸

Theorem C.2 (EoS regime, Phase I) . Let η > 0 be a small enough constant and ϕ be an activation function satisfying Assumptions 5.1, C.1, and C.2. Let z 0 := sup z { zr ′ ( z ) r ( z ) ≥ -1 2 } , z 1 := sup z { zr ′ ( z ) r ( z ) ≥ -1 } , and c 0 := max { r ( z 0 ) , r ( z 1 ) + 1 2 } ∈ ( 1 2 , 1) . Let δ ∈ (0 , 1 -c 0 ) be any given constant. Suppose that the initialization ( p 0 , q 0 ) satisfies | p 0 | ≤ 1 and q 0 ∈ ( c 0 , 1 -δ ) . Consider the reparameterized GD trajectory characterized in Eq. (7) . We assume that for all t ≥ 0 such that q t < 1 , we have p t = 0 . Then, there exists a time step t a = O δ,ϕ (log( η -1 )) such that for any t ≥ t a ,

where h : R → R is a function defined as

̸

The main difference between Theorem C.2 and Theorem 4.2 is the error term which is O ( η 2 ) in the former and O ( η 4 ) in the latter. This is because the 1 -step update rule of q t in Theorem C.2 is given by q t +1 = (1 + O ( η )) q t , while in Theorem 4.2 we have q t +1 = (1 + O ( η 2 )) q t .

Theorem C.3 (EoS regime, Phase II) . Under the same settings as in Theorem C.2, there exists a time step t b = Ω((1 -q 0 ) η -1 ) , such that q t b ≤ 1 and q t > 1 for any t > t b . Moreover, the GD iterates ( p t , q t ) converge to the point (0 , q ∗ ) such that

Theorem C.3 implies that in the EoS regime, GD with step size η converges to (0 , y ∗ ) which has the sharpness approximated as:

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 in Theorem C.2, let t a denote the obtained time step. Define the function ˜ λ : R > 0 → R given by

Then, the sequence ( ˜ λ ( q t ) ) ∞ t =0 is monotonically increasing. For any t ≥ t a , the sharpness at GD iterate ( x t , y t ) closely follows the sequence ( ˜ λ ( q t ) ) ∞ t =0 satisfying the following:

In Figure 1b, we conduct numerical experiments on single neuron model with tanh -activation, demonstrating that ˜ λ ( q t ) provides a close approximation of the sharpness.

C.2 Proof of Theorem C.1

We give the proof of Theorem C.1, restated below for the sake of readability.

Theorem C.1 (gradient flow regime) . Let η ∈ (0 , r (1) 2( r (1)+2) ) be a fixed step size and ϕ be a sigmoidal function satisfying Assumptions 5.1 and C.1. Suppose that the initialization ( p 0 , q 0 ) satisfies | p 0 | ≤ 1

and q 0 ∈ ( 1 1 -2 η , r (1) 4 η ) . Consider the reparameterized GD trajectory characterized in Eq. (7) . Then, the GD iterations ( p t , q t ) converge to the point (0 , q ∗ ) such that

Theorem C.1 directly follows from Proposition C.5 stated below.

Proposition C.5. Suppose that η ∈ (0 , r (1) 2( r (1)+2) ) , | p 0 | ≤ 1 and q 0 ∈ ( 1 1 -2 η , r (1) 4 η ) . Then for any t ≥ 0 , we have

and

Proof. We give the proof by induction; namely, if

are satisfied for time steps 0 ≤ t ≤ k 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 t = 0 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 = k +1 .

By induction assumptions, we have r (1) ≤ r ( p k ) ≤ 1 and q 0 ≤ q k ≤ 2 q 0 . From Eq. (7), we get

Due to the induction assumption, we obtain the desired bound on | p k +1 | as following:

Moreover, for any 0 ≤ t ≤ k , by Eq. (7) we have

where the second inequality comes from the fact that ϕ is 1-Lipschitz and ϕ (0) = 0 (Assumption 5.1). Hence, we have q k +1 ≥ q k ≥ q 0 . Note that q t q t +1 ∈ [ 1 2 , 1] for small η . Consequently, we have

where the second inequality holds since | log(1 + z ) | ≤ 2 | z | if | z | ≤ 1 2 . Therefore, we obtain the desired bound on q k +1 as following:

Since q 0 ≥ 1 1 -2 η and q 0 ≤ r (1) 4 η , we have

This implies that q k +1 ≤ exp( 1 2 ) q 0 ≤ 2 q 0 , as desired.

C.3 Proof of Theorem C.2

In this subsection, we prove Theorem C.2. We use the following notation:

All the lemmas in this subsection are stated in the context of Theorem C.2. The proof structure resembles that of Theorem 4.2. We informally summarize the lemmas used in the proof of Theorem C.2. Lemma C.6 proves that p t is bounded by a constant and q t increases monotonically with the increment bounded by O ( η ) . Lemma C.7 states that in the early phase of training, there exists a time step t 0 where s t 0 becomes smaller or equal to 2 2 -r (1) , which is smaller than 2 . Lemma C.9 demonstrates that if s t is smaller than 2 and | p t | ≥ ˆ r (1 -δ 4 ) , then | s t -1 | decreases exponentially. For the case where | p t | < ˆ r (1 -δ 4 ) , Lemma C.10 proves that | p t | increases at an exponential rate. Moreover, Lemma C.8 shows that if s t < 1 at some time step, then s t +1 is upper bounded by 1 + O ( η ) . Combining these findings, Proposition C.11 establishes that in the early phase of training, there exists a time step t ∗ a such that s t ∗ a = 1 + O δ,ϕ ( η ) . Lastly, Lemma C.12 demonstrates that if s t = 1 + O δ,ϕ ( η ) , then | s t -1 -h ( p t ) η | decreases exponentially.

Lemma C.6. Suppose that the initialization ( p 0 , q 0 ) satisfies | p 0 | ≤ 1 and q 0 ∈ ( c 0 , 1 -δ ) . Then for any t ≥ 0 such that q t ≤ 1 , it holds that

Proof. We prove by induction. We assume that for some t ≥ 0 , it holds that | p t | ≤ 4 and 1 2 ≤ q t ≤ 1 . We will prove that | p t +1 | ≤ 4 and 1 2 ≤ q t ≤ q t +1 ≤ (1 + O ( η )) q t . For the base case, | p 0 | ≤ 1 ≤ 4 and 1 2 ≤ c 0 < q t ≤ 1 holds by the assumptions on the initialization. Now suppose that for some t ≥ 0 , it holds that | p t | ≤ 4 and 1 2 ≤ q t ≤ 1 . By Eq. (7),

where we used Assumption 5.1 to bound | ϕ ( p t ) ϕ ′ ( p t ) | ≤ 1 . Moreover,

since | ϕ | is bounded by 1 . Hence, q t ≤ q t +1 ≤ (1 + O ( η )) q t

, as desired.

Lemma C.6 implies that p t is bounded by a constant throughout the iterations, and q t monotonically increases slowly, where the increment for each step is O ( η ) . Hence, there exists a time step T = Ω( δη -1 ) = Ω δ ( η -1 ) such that for any t ≤ T , it holds that q t ≤ 1 -δ 2 . Through out this subsection, we focus on these T early time steps. Note that for all 0 ≤ t ≤ T , it holds that q t ∈ ( c 0 , 1 -δ 2 ) .

Lemma C.7. There exists a time step t 0 = O δ,ϕ (1) such that s t 0 ≤ 2 2 -r (1) .

Proof. We start by proving the following statement: for any 0 ≤ t ≤ T , if 2 2 -r (1) < s t < 2 r (1) -1 , then s t +1 < 2 r (1) -1 and | p t +1 | ≤ (1 -r (1)) | p t | . Suppose that 2 2 -r (1) < s t < 2 r (1) -1 . Then from Eq. (7), it holds that

Hence, | p t +1 | ≤ (1 -r (1)) | p t | . Now we prove s t +1 < 2 r (1) -1 . Assume the contrary that s t +1 ≥ 2 r (1) -1 . Then, r ( p t +1 ) = q t +1 s t +1 < q t +1 < 1 -δ 2 so that | p t +1 | ≥ ˆ r (1 -δ 2 ) . By Mean Value Theorem, there exists p ∗ t ∈ ( | p t +1 | , | p t | ) such that

where we used Assumption C.2 (i) and ˆ r (1 -δ 2 ) ≤ | p t +1 | ≤ (1 -r (1)) | p t | . Consequently,

for small step size η . This gives a contradiction to our assumption that s t +1 ≥ 2 r (1) -1 . Hence, we can conclude that s t +1 < 2 r (1) -1 , as desired.

We proved that for any 0 ≤ t ≤ T , if 2 2 -r (1) < s t < 2 r (1) -1 , it holds that s t +1 < 2 r (1) -1 and | p t +1 | ≤ (1 -r (1)) | p t | . At initialization, | p 0 | ≤ 1 and q 0 < 1 , so that s 0 < r (1) -1 . If s 0 ≤ 2 2 -r (1) , then t 0 = 0 is the desired time step. Suppose that s 0 > 2 2 -r (1) . Then, we have s 1 < 2 r (1) -1 and | p 1 | ≤ (1 -r (1)) | p 0 | ≤ 1 -r (1) . Then we have either s 1 ≤ 2 2 -r (1) , or 2 2 -r (1) < s 1 < 2 r (1) -1 . In the previous case, t 0 = 1 is the desired time step. In the latter case, we can repeat the same argument and obtain s 2 < 2 r (1) -1 and | p 2 | ≤ (1 -r (1)) 2 . By inductively repeating the same argument, we can obtain a time step t 0 ≤ log(ˆ r (1 -δ 2 )) / log(1 -r (1)) = O δ,ϕ (1) such that either s t 0 ≤ 2 2 -r (1) , or | p t 0 | ≤ ˆ r (1 -δ 2 ) . In the latter case, r ( p t 0 ) ≥ 1 -δ 2 > q t 0 , and hence s t 0 < 1 < 2 2 -r (1) . Therefore, t 0 = O δ,ϕ (1) is the desired time step satisfying s t 0 ≤ 2 2 -r (1) .

According to Lemma C.7, there exists a time step t 0 = O δ,ϕ (1) such that s t 0 ≤ 2 2 -r (1) < 2(1+ η 2 ) -1 for small step size η . Now we prove the lemma below.

Lemma C.8. Suppose that s t ≤ 1 . Then, it holds that s t +1 ≤ 1 + O ( η ) .

Proof. For any p ∈ (0 , ˆ r ( q t 2 )) , we have r ( p ) ≥ q t 2 so that | f q t ( p ) | = ( -1 + 2 r ( p t ) q t ) p t . Hence,

for any p ∈ (0 , ˆ r ( q t 2 )) . By Assumption C.2 (ii), (iv) and q t ≥ c 0 ≥ r ( z 1 ) + 1 2 ≥ 2 r ( z 1 ) where z 1 = sup z { zr ′ ( z ) r ( z ) ≥ -1 } , both r ( p ) and (1 + pr ′ ( p ) r ( p ) ) are positive, decreasing function on (0 , ˆ r ( q t 2 )) . Consequently, ∂ ∂p | f q t ( p ) | is a decreasing function on (0 , ˆ r ( q t 2 )) .

Now note that q t 2 < q t < 1 , which means ˆ r (1) = 0 < ˆ r ( q t ) < ˆ r ( q t 2 ) by the definition of ˆ r . Note that ∂ ∂p | f q t ( p ) | at p = ˆ r ( q t ) evaluates to

where the inequalities used Assumption C.2 (ii) and q t > c 0 ≥ r ( z 0 ) where z 0 := sup z { zr ′ ( z ) r ( z ) ≥ -1 2 } , from the statement of Theorem C.2.

Therefore, since ∂ ∂p | f q t ( p ) | is decreasing on (0 , ˆ r ( q t 2 )) and is nonnegative at ˆ r ( q t ) , for any p ∈ (0 , ˆ r ( q t )) , it holds that ∂ ∂p | f q t ( p ) | ≥ 0 . In other words, | f q t ( p ) | is an increasing function on (0 , ˆ r ( q t )) . Since 0 ≤ s t ≤ 1 , we have | p t | ≤ ˆ r ( q t ) and it holds that

Therefore, with this inequality and Lemma C.6, we can conclude that

Using Lemma C.8, we prove the following lemma.

Lemma C.9. For any 0 ≤ t ≤ T , if s t < 2 and r ( p t ) ≤ 1 -δ 4 , then

where d ∈ (0 , 1 2 ] is a constant which depends on δ and ϕ .

Proof. From Eq. (7) it holds that

so that p t and p t +1 have opposite signs. By Mean Value Theorem, there exists θ t between -1 and (1 -2 s t ) satisfying

where the last equality used the fact that p t and θ t p t have opposite signs and r ′ ( z ) and z have opposite signs. Note that | θ t p t | is between | p t | and | p t +1 | . Consequently, the value | r ′ ( θ t p t ) | r ( θ t p t ) 2 is between | r ′ ( p t ) | r ( p t ) 2 and | r ′ ( p t +1 ) | r ( p t +1 ) 2 by Assumption C.2 (i). We will prove the current lemma based on Eq. (17). We divide into following three cases: (1) s t ≥ 1 and s t +1 ≥ 1 , (2) s t ≥ 1 and s t < 1 , and (3) s t < 1 .

Case 1. Suppose that s t ≥ 1 and s t +1 ≥ 1 . Here, we have | p t | ≥ ˆ r ( q t ) ≥ ˆ r (1 -δ 2 ) and similarly | p t +1 | ≥ ˆ r (1 -δ 2 ) . By Assumption C.2 (i), | r ′ ( θ t p t ) | r ( θ t p t ) 2 ≥ | r ′ (ˆ r (1 -δ 2 )) | (1 -δ 2 ) 2 . Hence, Eq. (17) gives

Consequently, by Lemma C.6,

where we used q t > c 0 > 1 2 and s t < 2 . Therefore, we can obtain the following inequality:

Case 2. Suppose that s t ≥ 1 and s t +1 < 1 . Here, we have r ( p t +1 ) > q t +1 ≥ q t ≥ r ( p t ) , so that | p t +1 | < | p t | . Consequently, | r ′ ( θ t p t ) | r ( θ t p t ) 2 ≤ | r ′ ( p t ) | r ( p t ) 2 by Assumption C.2 (i). Hence, we can deduce from Eq. (17) that

Consequently, by Assumption C.2 (ii),

where we used r ( p t ) ≥ q t 2 > c 0 2 . Therefore, we can obtain the following inequality:

where -1 < ˆ r ( c 0 2 ) r ′ (ˆ r ( c 0 2 )) r (ˆ r ( c 0 2 )) = 2ˆ r ( c 0 2 ) r ′ (ˆ r ( c 0 2 )) c 0 < 0 , since c 0 ≥ r ( z 1 ) + 1 2 ≥ 2 r ( z 1 ) with z 1 = sup z { zr ′ ( z ) r ( z ) ≥ -1 } , and r ( z 1 ) ≤ 1 2 holds by Assumption C.2 (iv).

Case 3. Suppose that s t < 1 . By Lemma C.8, it holds that s t +1 ≤ 1 + O ( η ) . Moreover, we assumed r ( p t ) ≤ 1 -δ 4 , so that | p t | ≥ ˆ r (1 -δ 4 ) . We also have

Consequently, by Assumption C.2 (i), it holds that | r ′ ( θ t p t ) | r ( θ t p t ) 2 ≥ | r ′ (ˆ r (1 -δ 4 )) | (1 -δ 4 ) 2 . Hence, by Eq. (17), we have

and hence,

where we used q t > 1 2 . Therefore, we can obtain the following inequality:

where the first inequality is from Lemma C.8.

Combining the three cases, we can finally conclude that if we choose

then | s t +1 -1 | ≤ (1 -d ) | s t -1 | + O ( η ) , as desired.

Lemma C.9 implies that if s t < 2 and | p t | ≥ ˆ r (1 -δ 4 ) , then | s t -1 | exponentially decreases. We prove Lemma C.10 to handle the regime | p t | < ˆ r (1 -δ 4 ) , which is stated below.

Lemma C.10. For any 0 ≤ t ≤ T , if r ( p t ) ≥ 1 -δ 4 , it holds that

Proof. If r ( p t ) ≥ 1 -δ 4 , then s t = q t r ( p t ) < 1 -δ 2 1 -δ 4 = 4 -2 δ 4 -δ , where we used q t < 1 -δ 2 for any 0 ≤ t ≤ T . Consequently,

Now we prove Proposition C.11, which proves that s t reaches close to 1 with error bound of O ( η ) .

Proposition C.11. There exists a time step t ∗ a = O δ,ϕ (log( η -1 )) satisfying

Proof. By Lemma C.7, there exists a time step t 0 = O δ,ϕ (1) such that s t 0 ≤ 2 2 -r (1) . Here, we divide into two possible cases: (1) s t 0 < 1 , and (2) 1 ≤ s t 0 ≤ 2 2 -r (1) .

Case 1. Suppose that s t 0 < 1 . By Lemma C.10, if r ( p t 0 ) ≥ 1 -δ 4 (or equivalently, | p t 0 | ≤ ˆ r (1 -δ 4 ) ), then there exists a time step t 1 ≤ t 0 +log( ˆ r (1 -δ 4 ) | p t 0 | ) / log( 2 2 -δ ) = O δ,ϕ (1) such that | p t 1 | ≥ ˆ r (1 -δ 4 ) . Wedenote the first time step satisfying | p t 1 | ≥ ˆ r (1 -δ 4 ) and t 1 ≥ t 0 by t 1 = O δ,ϕ (1) . By Lemma C.8, it holds that s t 1 ≤ 1 + O ( η ) since s t 1 -1 < 1 . Consequently, if s t 1 ≥ 1 , then | s t 1 -1 | ≤ O ( η ) so that t ∗ a = t 1 is the desired time step. Hence, it suffices to consider the case when s t 1 < 1 . Here, we can apply Lemma C.9 which implies that

where d is a constant which depends on δ and ϕ . Then, there are two possible cases: either | s t 1 -1 | ≤ O ( ηd -1 ) , or | s t 1 +1 -1 | ≤ (1 -d 2 ) | s t 1 -1 | . It suffices to consider the latter case, suppose that | s t 1 +1 -1 | ≤ (1 -d 2 ) | s t 1 -1 | . Since we are considering the case s t 1 < 1 , again by Lemma C.8, we have s t 1 +1 ≤ 1+ O ( η ) . Since | p t 1 +1 p t 1 | = 2 s t 1 -1 > 1 , we have | p t 1 +1 | ≥ | p t 1 | ≥ ˆ r (1 -δ 4 ) . This means that we can again apply Lemma C.9 and repeat the analogous argument. Hence, there exists a time step t 2 ≤ t 1 +log( η 1 -s t 1 ) / log(1 -d 2 ) = O δ,ϕ (log( η -1 )) , such that | s t 2 -1 | ≤ O ( ηd -1 ) = O δ,ϕ ( η ) .

Case 2. Suppose that 1 ≤ s t 0 ≤ 2 2 -r (1) . Then, r ( p t 0 ) ≤ q t 0 ≤ 1 -δ 2 , so we can apply Lemma C.9. There are two possible cases: either | s t 0 +1 -1 | ≤ O ( ηd -1 ) = O δ,ϕ ( η ) , or | s t 0 +1 -1 | ≤ (1 -d 2 ) | s t 0 -1 | . It suffices to consider the latter case. If s t 0 +1 ≥ 1 , we can again apply Lemma C.9 and repeat the analogous argument. Hence, we can obtain a time step t ′ 0 ≤ t 0 +log( η 1 -s t 0 ) / log(1 -d 2 ) = O δ,ϕ (log( η -1 )) such that either s t ′ 0 < 1 or | s t ′ 0 -1 | = O δ,ϕ ( η ) is satisfied. If s t ′ 0 < 1 , we proved in Case 1 that there exists a time step t ′ 2 = t ′ 0 + O δ,ϕ (log( η -1 )) such that | s t ′ 2 -1 | ≤ O δ,ϕ ( η ) , and this is the desired bound.

Now we carefully handle the error term O ( η ) obtained in Proposition C.11 and a provide tighter bound on s t by proving Lemma C.12 stated below.

Lemma C.12. If | s t -1 | ≤ O δ,ϕ ( η ) , then it holds that

where

̸

Proof. Suppose that | s t -1 | ≤ O δ,ϕ ( η ) . Then, | p t +1 | = ∣ ∣ ∣ 1 -2 s t ∣ ∣ ∣ | p t | = (1 + O δ,ϕ ( η )) | p t | . By Eq. (17) proved in Lemma C.9, there exists ϵ t = O δ,ϕ ( η ) such that

where we used the Taylor expansion on r ′ ( p ) r ( p ) 2 with the fact that d dp ( r ′ ( p ) r ( p ) 2 ) is bounded on [ -4 , 4] and that | p t | ≤ 4 to obtain the second equality. Note that q t +1 = (1 -ηϕ ( p t ) 2 ) -2 q t = (1+2 ηϕ ( p t ) 2 ) q t + O ( η 2 ) by Eq (7). Consequently,

Note that h is even, and twice continuously differentiable function by Lemma C.13. Consequently, h ′ (0) = 0 and h ′ ( p ) = O ϕ ( p ) , since h ′′ is bounded on closed interval. Consequently, h ( p t +1 ) = h ((1 + O δ,ϕ ( η )) p t ) = h ( p t ) + O δ,ϕ ( ηp 2 t ) . Hence, we can obtain the following:

Note that r ( p t ) = (1 + O δ,ϕ ( η )) q t ≥ (1 + O δ,ϕ ( η )) q 0 ≥ c 0 ≥ r ( z 0 ) for small step size η , where z 0 = sup { zr ′ ( z ) r ( z ) ≥ -1 2 } . Consequently, it holds that 1 + 2 p t r ′ ( p t ) r ( p t ) ≥ 0 . Therefore, we can conclude that

as desired.

Now we give the proof of Theorem C.2, restated below for the sake of readability.

̸

Theorem C.2 (EoS regime, Phase I) . Let η > 0 be a small enough constant and ϕ be an activation function satisfying Assumptions 5.1, C.1, and C.2. Let z 0 := sup z { zr ′ ( z ) r ( z ) ≥ -1 2 } , z 1 := sup z { zr ′ ( z ) r ( z ) ≥ -1 } , and c 0 := max { r ( z 0 ) , r ( z 1 ) + 1 2 } ∈ ( 1 2 , 1) . Let δ ∈ (0 , 1 -c 0 ) be any given constant. Suppose that the initialization ( p 0 , q 0 ) satisfies | p 0 | ≤ 1 and q 0 ∈ ( c 0 , 1 -δ ) . Consider the reparameterized GD trajectory characterized in Eq. (7) . We assume that for all t ≥ 0 such that q t < 1 , we have p t = 0 . Then, there exists a time step t a = O δ,ϕ (log( η -1 )) such that for any t ≥ t a ,

where h : R → R is a function defined as

̸

Proof of Theorem C.2. By Proposition C.11, there exists a time step t ∗ a = O δ,ℓ (log( η -1 )) which satisfies:

By Lemma C.12, there exists a constant D > 0 which depends on δ, ϕ such that if | s t -1 | = O δ,ϕ ( η ) , then

Hence, if | s t -1 | = O δ,ϕ ( η ) and | s t -1 -h ( p t ) η | ≥ ( -p t r ( p t ) r ′ ( p t ) ) Dη 2 , then

For any t ≤ T , we have q t < 1 -δ 2 so that if | s t -1 | = O δ,ϕ ( η ) , then r ( p t ) ≤ (1+ O δ,ϕ ( η )) q t < 1 -δ 4 for small step size η . From Eq. (20) with t = t ∗ a , we have either

or

where we used Assumption C.2 (ii) and | p t | > ˆ r (1 -δ 4 ) . In the latter case, | s t ∗ a +1 -1 | = O δ,ϕ ( η ) continues to hold and we can again use Eq. (20) with t = t ∗ a + 1 . By repeating the analogous arguments, we can obtain the time step

which satisfies: either

or

where we used | p t | ≤ 4 from Lemma C.6 and -zr ( z ) r ′ ( z ) ≥ -1 r ′′ (0) for any z by Assumption C.2 (iii).

By Eq. (19), if | s t -1 -h ( p t ) η | ≤ ( -4 r (4) r ′ (4) ) Dη 2 is satisfied for any time step t , then

by | p t | ≤ 4 from Lemma C.6 and Assumption C.2 (iii).

Hence, by induction, we have the desired bound as following: for any t ≥ t a ,

by | p t | ≤ 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 of h defined in Theorem C.2.

Lemma C.13. Consider the function h defined in Theorem C.2, given by

̸

Then, h is a positive, even, and bounded twice continuously differentiable function.

Proof. By Assumption C.1, h is a positive, even function. Moreover, h is continuous since lim p → 0 h ( p ) = -1 r ′′ (0) = h (0) . Continuous function on a compact domain is bounded, so h is bounded on the closed interval [ -1 , 1] . Note that ϕ ( p ) 2 ≤ 1 , and ( -r ( p ) pr ′ ( p ) ) is positive, decreasing function on p > 0 by Assumption C.2 (ii). Hence, h is bounded on [1 , ∞ ) . Since h is even, h is bounded on ( -∞ , 1] . Therefore, h is a bounded on R .

̸

We finally prove that h is twice continuously differentiable. Since r is even and C 4 on R , we can check that for any p = 0 ,

̸

and h ′ ( p ) = 0 . Moreover, for any p = 0 ,

and

Since lim p → 0 h ′′ ( p ) = h ′′ (0) , we can conclude that h is a twice continuously differentiable function.

Now we give the proof of Theorem C.3, restated below for the sake of readability.

Theorem C.3 (EoS regime, Phase II) . Under the same settings as in Theorem C.2, there exists a time step t b = Ω((1 -q 0 ) η -1 ) , such that q t b ≤ 1 and q t > 1 for any t > t b . Moreover, the GD iterates ( p t , q t ) converge to the point (0 , q ∗ ) such that

Proof. We first prove that there exists a time step t ≥ 0 such that q t > 1 . Assume the contrary that q t ≤ 1 for all t ≥ 0 . Let t a be the time step obtained in Theorem C.2. Then for any t ≥ t a , we have

for small step size η . The function g ( p ) := r ( p ) -1 + h ( p ) η 2 is even, continuous, and has the function value g (0) = η 2 | r ′′ (0) | > 0 . Consequently, there exists a positive constant ϵ > 0 such that g ( p ) > 0 for all p ∈ ( -ϵ, ϵ ) . Then, we have | p t | ≥ ϵ for all t ≥ t a , since g ( p t ) ≤ 0 . This implies that for any t ≥ t a , it holds that

so q t grows exponentially, which results in the existence of a time step t ′ b ≥ t a such that q t ′ b > 1 , a contradiction.

Therefore, there exists a time step t b such that q t b ≤ 1 and q t > 1 for any t > t b , i.e., q t jumps across the value 1 . This holds since the sequence ( q t ) is monotonically increasing. For any t ≤ t b , we have q t +1 ≤ q t + O ( η ) by Lemma C.6, and this implies that t b ≥ Ω((1 -q 0 ) η -1 ) , as desired.

Lastly, we prove the convergence of GD iterates ( p t , q t ) . Let t > t b be given. Then, q t ≥ q t b +1 > 1 and it holds that

Hence, | p t | is exponentially decreasing for t > t b . Therefore, p t converges to 0 as t →∞ . Since the sequence ( q t ) ∞ t =0 is monotonically increasing and bounded (due to Theorem C.2, it converges. Suppose that ( p t , q t ) converges to the point (0 , q ∗ ) . By Theorem C.2, we can conclude that

which is the desired bound.

C.5 Proof of Theorem C.4

In this subsection, we prove Theorem C.4. We first prove an important bound on the sharpness value provided by the Proposition C.14 stated below.

Proposition C.14. For any ( x, y ) ∈ R 2 with r ( x ) + xr ′ ( x ) ≥ 0 , the following inequality holds:

Proof. Let ( x, y ) ∈ R 2 be given. The loss Hessian at ( x, y ) is

Note that the largest absolute value of the eigenvalue of a symmetric matrix equals to its spectral norm. Since trace of the Hessian, tr ( ∇ 2 L ( x, y )) = ( ϕ ( x ) ϕ ′′ ( x ) + ϕ ′ ( x ) 2 ) y 2 + ϕ ( x ) 2 = ( r ( x ) + xr ′ ( x )) y 2 + ϕ ( x ) 2 ≥ 0 , is non-negative, the spectral norm of the Hessian ∇ 2 L ( x, y ) equals to its largest eigenvalue. Hence, we have

which is the desired lower bound. We also note that for any matrix A , the inequality ∥ A ∥ 2 ≤ ∥ A ∥ F holds, where ∥ A ∥ F is the Frobenius norm. Hence, we have

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 as in Theorem C.2, let t a denote the obtained time step. Define the function ˜ λ : R > 0 → R given by

Then, the sequence ( ˜ λ ( q t ) ) ∞ t =0 is monotonically increasing. For any t ≥ t a , the sharpness at GD iterate ( x t , y t ) closely follows the sequence ( ˜ λ ( q t ) ) ∞ t =0 satisfying the following:

Proof. By Theorem C.2 and since h is a bounded function by Lemma C.13, we have s t = 1 + O ϕ ( η ) for any t ≥ t a . Note that r ( p t ) = (1 + O δ,ϕ ( η )) q t ≥ (1 + O δ,ϕ ( η )) q 0 ≥ c 0 ≥ r ( z 0 ) for small step size η , where z 0 = sup { zr ′ ( z ) r ( z ) ≥ -1 2 } . Consequently, it holds that 1+ 2 p t r ′ ( p t ) r ( p t ) ≥ 0 , and this implies 1 + p t r ′ ( p t ) r ( p t ) ≥ 1 2 . By Proposition C.14, we can bound the sharpness λ t := λ max ( ∇ 2 L ( x t , y t )) by

By Lemma C.6, | p t | ≤ 4 for any time step t . Moreover, since 1 + p t r ′ ( p t ) r ( p t ) ≥ 1 2 holds for any t ≥ t a with small step size η , we have 4 p 2 t ( 1 + p t r ′ ( p t ) r ( p t ) ) -1 = O ϕ (1) . Hence, for any t ≥ t a , it holds that

For any t ≥ t a , we have s t = 1+ O ϕ ( η ) and this implies | s -1 t -1 | = O ϕ ( η ) and | r ( p t ) -q t | = O ϕ ( η ) . For any 0 < q < 1 , we have

so that

Therefore, d dq ( ˆ r ( q ) r ′ (ˆ r ( q )) q ) is bounded on [ 1 4 , 1) and Taylor's theorem gives

for any time step t with q t < 1 . Hence, if q t < 1 , we have the following bound:

where we used p r r ′ ( p t ) q t = p t r ′ ( p t ) r ( p t ) (1 + O ϕ ( η )) = p t r ′ ( p t ) r ( p t ) + O ϕ ( η ) , since | p t r ′ ( p t ) | ≤ r ( p t ) for any t ≥ t a . Now let t be any given time step with q t ≥ 1 . Then, r ( p t ) = 1 - O ϕ ( η ) , and since r ( z ) = 1 + r ′′ (0) z 2 + O ϕ ( z 4 ) for small z , we have | p t | = O ϕ ( √ η ) . Hence,

for any t with q t ≥ 1 , where we again used p r r ′ ( p t ) q t = p t r ′ ( p t ) r ( p t ) (1 + O ϕ ( η )) = p t r ′ ( p t ) r ( p t ) + O ϕ ( η ) . By Eqs. (21), (22), and (23), we can conclude that for any t ≥ t a , we have

the desired bound.

Finally, we can easily check that the sequence ( ˜ λ ( q t )) ∞ t =0 is monotonically increasing, since z ↦→ zr ′ ( z ) r ( z ) is a decreasing function by Assumption C.2 (ii) and the sequence ( q t ) is monotonically increasing.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n- 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\nWEAKNESSES:\n- 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\nQUESTIONS:\n- 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\nLIMITATIONS:\nPlease see the second point in the 'weaknesses' section.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

comments: 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:
(1) Two-layer fully connected linear neural networks trained on a single data point with a convex and Lipschitz loss.
(2) Single-neuron nonlinear network trained on a single data point with the squared loss.

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).
2. 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?
2. 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.
3. The conditions you have on the initialization in Thm 4.2 are very restrictive. Do they hold in practice?

QUESTIONS:
See the weaknesses section.

LIMITATIONS:
N/A.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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\nWEAKNESSES:\n- 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\nQUESTIONS:\nIn 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\nLIMITATIONS:\nSome limitations of the current results are discussed in **Conclusion**. For other possible limitations, please refer to the weaknesses and questions above.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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)\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nThe authors have adequately addressed the limitations of this work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-31"}