Title: Learning threshold neurons via edge of stability
Venue: NeurIPS 2023
Accepted: yes
Keywords: Gradient descent

edge of stability

generalization
Abstract: Existing analyses of neural network training often operate under the unrealistic assumption of an extremely small learning rate. This lies in stark contrast to practical wisdom and empirical studies, such as the work of J. Cohen et al. (ICLR 2021), which exhibit startling new phenomena (the "edge of stability"' or "unstable convergence") and potential benefits for generalization in the large learning rate regime. Despite a flurry of recent works on this topic, however, the latter effect is still poorly understood. In this paper, we take a step towards understanding genuinely non-convex training dynamics with large learning rates by performing a detailed analysis of gradient descent for simplified models of two-layer neural networks. For these models, we provably establish the edge of stability phenomenon and discover a sharp phase transition for the step size below which the neural network fails to learn ``threshold-like'' neurons (i.e., neurons with a non-zero first-layer bias). This elucidates one possible mechanism by which the edge of stability can in fact lead to better generalization, as threshold neurons are basic building blocks with useful inductive bias for many tasks.

Learning threshold neurons via the 'edge of stability'

Abstract

Existing analyses of neural network training often operate under the unrealistic assumption of an extremely small learning rate. This lies in stark contrast to practical wisdom and empirical studies, such as the work of J. Cohen et al. (ICLR 2021), which exhibit startling new phenomena (the 'edge of stability' or 'unstable convergence') and potential benefits for generalization in the large learning rate regime. Despite a flurry of recent works on this topic, however, the latter effect is still poorly understood. In this paper, we take a step towards understanding genuinely non-convex training dynamics with large learning rates by performing a detailed analysis of gradient descent for simplified models of two-layer neural networks. For these models, we provably establish the edge of stability phenomenon and discover a sharp phase transition for the step size below which the neural network fails to learn 'threshold-like' neurons (i.e., neurons with a non-zero first-layer bias). This elucidates one possible mechanism by which the edge of stability can in fact lead to better generalization, as threshold neurons are basic building blocks with useful inductive bias for many tasks.

0

.

000

0

.

002

0

.

004

0

.

006

0

.

008

0

.

010

0

.

012

learning rate

η

-

1

.

6

-

1

.

4

-

1

.

2

-

1

.

0

-

0

.

8

-

0

.

6

-

0

.

4

-

0

.

2

final bias

0

.

0000

0

.

0005

0

.

0010

0

.

0015

-

0

.

6

-

0

.

4

-

0

.

2

60

58

0

.

0000

0

.

006

learning rate

0

η

1 Introduction

How much do we understand about the training dynamics of neural networks? We begin with a simple and canonical learning task which indicates that the answer is still 'far too little'.

)

%

final test accuracy (

68

66

64

62

60

58

Sinho Chewi

Institute for Advanced Study Princeton, NJ schewi@ias.edu

Yi Zhang

Microsoft Research Redmond, WA zhayi@microsoft.com

0

.

000

0

.

002

0

.

004

.

0

.

0005

008

0

.

0

.

0010

010

0

.

0

.

0015

012

Motivating example: Consider a binary classification task of labeled pairs ( x ( i ) , y ( i ) ) ∈ R d ×{± 1 } where each covariate x ( i ) consists of a 1 -sparse vector (in an unknown basis) corrupted by additive Gaussian noise, and the label y ( i ) is the sign of the non-zero coordinate of the 1 -sparse vector. Due to rotational symmetry, we can take the unknown basis to be the standard one and write

where y ( i ) ∈ {± 1 } is a random label, j ( i ) ∈ [ d ] is a random index, ξ ( i ) is Gaussian noise, and λ > 1 is the unknown signal strength. In fact, (1) is a special case of the well-studied sparse coding model (Olshausen and Field, 1997; Vinje and Gallant, 2000; Olshausen and Field, 2004; Yang et al., 2009; Koehler and Risteski, 2018; Allen-Zhu and Li, 2022). We ask the following fundamental question:

How do neural networks learn to solve the sparse coding problem (1) ?

In spite of the simplicity of the setting, a full resolution to this question requires a thorough understanding of surprisingly rich dynamics which lies out of reach of existing theory . To illustrate this point, consider an extreme simplification in which the basis e 1 , . . . , e d is known in advance, for which it is natural to parametrize a two-layer ReLU network as

The parametrization (2) respects the latent data structure (1) well: a good network has a negative bias b to threshold out the noise, and has a -< 0 and a + > 0 to output correct labels. We are particularly interested in understanding the mechanism by which the bias b becomes negative, thereby allowing the non-linear ReLU activation to act as a threshold function; we refer to this as the problem of learning 'threshold neurons'. More broadly, such threshold neurons are of interest as they constitute basic building blocks for producing neural networks with useful inductive bias.

1

.

0

1

.

5

η

= 2

.

5

×

10

-

5

η

= 2

.

5

×

10

-

4

0

25

50

75

100

iteration

2

4

6

training loss

η

= 2

.

5

×

10

-

3

0

25

50

75

100

iteration

-

0

.

4

-

0

.

3

-

0

.

2

-

0

.

1

0

.

0

bias

b

0

5

0

25

50

75

100

iteration

-

20

0

d

·

(

a

-

+

a

+

)

We train the parameters a -, a + , b using gradient descent with step size η > 0 on the logistic loss ∑ n i =1 ℓ logi ( y ( i ) f ( x ( i ) ; a -, a + , b )) , where ℓ logi ( z ) := log(1 + exp( -z )) , and we report the results in Figures 1 and 2. The experiments reveal a compelling picture of the optimization dynamics.

■ Large learning rates are necessary, both for generalization and for learning threshold neurons. Figure 1 shows that the bias decreases and the test accuracy increases as we increase η ; note that we plot the results after a fixed time (iteration × η ), so the observed results are not simply because larger learning rates track the continuous-time gradient flow for a longer time.

■ Large learning rates lead to unexpected phenomena: non-monotonic loss and wild oscillations of a -+ a + . Figure 2 shows that large learning rates also induce stark phenomena, such as non-monotonic loss and large weight fluctuations, which lie firmly outside the explanatory power of existing analytic techniques based on principles from convex optimization.

■ There is a phase transition between small and large learning rates. In Figure 1, we zoom in on learning rates around η ≈ 0 . 0006 and observe sharp phase transition phenomena.

We have presented these observations in the context of the simple ReLU network (2), but we emphasize that these findings are indicative of behaviors observed in practical neural network training settings. In Figure 3, we display results for a two-layer ReLU network trained on the full sparse coding model (1) with unknown basis, as well as a deep neural network trained on CIFAR-10. In each case, we again observe non-monotonic loss coupled with steadily decreasing bias parameters. For these richer models, the transition from small to large learning rates is oddly reminiscent of well-known separations between the 'lazy training' or 'NTK' regime Jacot et al. (2018) and the more expressive 'feature learning' regime. For further experimental results, see Appendix B.

0

200

400

600

800

1000

1200

1400

step×

0.50

0.55

0.60

0.65

0.70

0.75

0.80

0.85

0.90

test acc

lr=0.03

lr=0.04

lr=0.05

lr=0.07

lr=0.2

0

200

400

600

800

1000

1200

1400

step×

2.0

1.5

1.0

0.5

0.0

bias

lr=0.03

lr=0.04

lr=0.05

lr=0.07

lr=0.2

0

10

20

30

40

50

step×

50

60

70

80

90

test acc

lr=0.2

lr=0.05

0

10

20

30

40

50

step×

0.0200

0.0175

0.0150

0.0125

0.0100

0.0075

0.0050

0.0025

0.0000

median bias

lr=0.2

lr=0.05

We currently do not have right tools to understand these phenomena. First of all, a drastic change in behavior between the small and the large learning rates cannot be captured through well-studied regimes, such as the 'neural tangent kernel' (NTK) regime (Jacot et al., 2018; Allen-Zhu et al., 2019; Arora et al., 2019; Chizat et al., 2019; Du et al., 2019; Oymak and Soltanolkotabi, 2020) or the mean-field regime Chizat and Bach (2018); Mei et al. (2019); Chizat (2022); Nitanda et al. (2022); Rotskoff and Vanden-Eijnden (2022). In addition, understanding why a large learning rate is required to learn the bias is beyond the scope of prior theoretical works on the sparse coding model (Arora et al., 2015; Karp et al., 2021). Our inability to explain these findings points to a serious gap in our grasp of neural network training dynamics and calls for a detailed theoretical study.

1.1 Main scope of this work

In this work, we do not aim to understand the sparse coding problem (1) in its full generality. Instead, we pursue the more modest goal of shedding light on the following question.

Q. What is the role of a large step size in learning the bias for the ReLU network (2)?

As discussed above, the dynamics of the simple ReLU network (2) is a microcosm of emergent phenomena beyond the convex optimization regime. In fact, there is a recent growing body of work (Cohen et al., 2021; Arora et al., 2022; Ahn et al., 2022; Lyu et al., 2022; Ma et al., 2022; Wang et al., 2022b; Chen and Bruna, 2023; Damian et al., 2023; Zhu et al., 2023) on training with large learning rates, which largely aims at explaining a striking empirical observation called the ' edge of stability (EoS) ' phenomenon.

The edge of stability (EoS) phenomenon is a set of distinctive behaviors observed recently by Cohen et al. (2021) when training neural networks with gradient descent (GD). Here we briefly summarize the salient features of the EoS and defer a discussion of prior work to Subsection 1.3. Recall that if we use GD to optimize an L -smooth loss function with step size η , then the well-known descent lemma from convex optimization ensures monotonic decrease in the loss so long as L < 2 /η . In contrast, when L > 2 /η , it is easy to see on simple convex quadratic examples that GD can be unstable (or divergent). The main observation of Cohen et al. (2021) is that when training neural networks 1 with constant step size η > 0 , the largest eigenvalue of the Hessian at the current iterate (dubbed the 'sharpness') initially increases during training ('progressive sharpening') and saturates near or above 2 /η ('EoS').

A surprising message of the present work is that the answer to our main question is intimately related to the EoS . Indeed, Figure 4 shows that the GD iterates of our motivating example exhibit the EoS during the initial phase of training when the bias decreases rapidly.

0

100

200

300

400

500

iteration

-

0

.

4

-

0

.

3

-

0

.

2

-

0

.

1

0

.

0

b

η

= 2

/

2000

η

= 2

/

1500

η

= 2

/

1000

-

5

0

5

-

10

0

10

0

100

200

300

400

500

iteration

0

20

d

·

(

a

-

+

a

+

)

0

100

200

300

400

500

iteration

0

500

1000

1500

2000

2500

sharpness

Consequently, we first set out to thoroughly understand the workings of the EoS phenomena through a simple example. Specifically, we consider a single-neuron linear neural network in dimension 1 , corresponding to the loss

Although toy models have appeared in works on the EoS (see Subsection 1.3), our example is simpler than all prior models, and we provably establish the EoS for (3) with transparent proofs.

We then use the newfound insights gleaned from the analysis of (3) to answer our main question. To the best of our knowledge, we provide the first explanation of the mechanism by which a large learning rate can be necessary for learning threshold neurons.

1.2 Our contributions

Explaining the EoS with a single-neuron example. Although the EoS has been studied in various settings (see Subsection 1.3 for a discussion), these works either do not rigorously establish the EoS phenomenon, or they operate under complex settings with opaque assumptions. Here, we study a simple two-dimensional loss function, ( x, y ) ↦→ ℓ ( xy ) , where ℓ is convex, even, and Lipschitz. Some examples include 2 ℓ ( s ) = 1 2 log(1 + exp( -s )) + 1 2 log(1 + exp(+ s )) and ℓ ( s ) = √ 1 + s 2 . Surprisingly, GD on this loss already exhibits rich behavior (Figure 5).

En route to this result, we rigorously establish the quasistatic dynamics formulated in Ma et al. (2022).

x

y

GF

EoS

The elementary nature of our example leads to transparent arguments, and consequently our analysis isolates generalizable principles for 'bouncing' dynamics. To demonstrate this, we use our insights to study our main question of learning threshold neurons.

Learning threshold neurons with the mean model. The connection between the single-neuron example and the ReLU network (2) can already be anticipated via a comparison of the dynamics: (i) for the single neuron example, x oscillates wildly while y decreases (Figure 5); (ii) for the ReLU network (2), the sum of weights ( a -+ a + ) oscillates while b decreases (Figure 2).

We study this example in Section 2 and delineate a transition from the 'gradient flow' regime to the 'EoS regime', depending on the step size η and the initialization. Moreover, in the EoS regime, we rigorously establish asymptotics for the limiting sharpness which depend on the higher-order behavior of ℓ . In particular, for the two losses mentioned above, the limiting sharpness is 2 /η + O ( η ) , whereas for losses ℓ which are exactly quadratic near the origin the limiting sharpness is 2 /η + O (1) .

In fact, this connection can be made formal by considering an approximation for the GD dynamics for the ReLU network (2). It turns out (see Subsection 3.1 for details) that during the initial phase of training, the dynamics of A t := d ( a -t + a + t ) and b t due to the ReLU network are well-approximated by the 'rescaled' GD dynamics on the loss ( A,b ) ↦→ ℓ sym ( A × g ( b )) , where the step size for the A -dynamics is multiplied by 2 d 2 , g ( b ) := E z ∼N (0 , 1) ReLU( z + b ) is the 'smoothed' ReLU, and ℓ sym is the symmetrized logistic loss; see Subsection 3.1 and Figure 8. We refer to these dynamics as the mean model . The mean model bears a great resemblance to the single-neuron example ( x, y ) ↦→ ℓ ( xy ) , and hence we can

d

(

a

+

+

a

-

)

b

leverage the techniques developed for the latter in order to study the former.

Our main result for the mean model precisely explains the phase transition in Figure 1. For any δ > 0 ,

if η ≤ (8 -δ ) π/d 2 , then the mean model fails to learn threshold neurons : the limiting bias satisfies | b ∞ | = O δ (1 /d 2 ) .

1.3 Related work

if η ≥ (8 + δ ) π/d 2 , then the mean model enters the EoS and learns threshold neurons : the limiting bias satisfies b ∞ ≤ -Ω δ (1) .

Edge of stability. Our work is motivated by the extensive empirical study of Cohen et al. (2021), which identified the EoS phenomenon. Subsequently, there has been a flurry of works aiming at developing a theoretical understanding of the EoS, which we briefly summarize here.

Properties of the loss landscape. The works (Ahn et al., 2022; Ma et al., 2022) study the properties of the loss landscape that lead to the EoS. Namely, Ahn et al. (2022) argue that the existence of forward-invariant subsets near the minimizers allows GD to convergence even in the unstable regime. They also explore various characteristics of EoS in terms of loss and iterates. Also, Ma et al. (2022) empirically show that the loss landscape of neural networks exhibits subquadratic growth locally around the minimizers. They prove that for a one-dimensional loss, subquadratic growth implies that GD finds a 2 -periodic trajectory.

Limiting dynamics. Other works characterize the limiting dynamics of the EoS in various regimes. (Arora et al., 2022; Lyu et al., 2022) show that (normalized) GD tracks a 'sharpness reduction flow' near the manifold of minimizers. The recent work of Damian et al. (2023) obtains a different predicted dynamics based on self-stabilization of the GD trajectory. Also, Ma et al. (2022) describes a quasi-static heuristic for the overall trajectory of GD when one component of the iterate is oscillating.

Simple models and beyond. Closely related to our own approach, there are prior works which carefully study simple models. Chen and Bruna (2023) prove global convergence of GD for the two-dimensional function ( x, y ) ↦→ ( xy -1) 2 and a single-neuron student-teacher setting; note that unlike our results, they do not study the limiting sharpness. Wang et al. (2022b) study progressive sharpening for a neural network model. Also, the recent and concurrent work of Zhu et al. (2023) studies the two-dimensional loss ( x, y ) ↦→ ( x 2 y 2 -1) 2 ; to our knowledge, their work is the first to asymptotically and rigorously show that the limiting sharpness of GD is 2 /η in a simple setting,

at least when initialized locally. In comparison, in Section 2, we perform a global analysis of the limiting sharpness of GD for ( x, y ) ↦→ ℓ ( xy ) for a class of convex, even, and Lipschitz losses ℓ , and in doing so we clearly delineate the 'gradient flow regime' from the 'EoS regime'.

Effect of learning rate on learning. Recently, several works have sought to understand how the choice of learning rate affects the learning process, in terms of the properties of the resulting minima (Jastrzebski et al., 2018; Wu et al., 2018; Mulayoff et al., 2021; Nacson et al., 2022) and the behavior of optimization dynamics (Xing et al., 2018; Jastrzebski et al., 2019, 2020; Lewkowycz et al., 2020; Jastrzebski et al., 2021).

Li et al. (2019) demonstrate for a synthethic data distribution and a two-layer ReLU network model that choosing a larger step size for SGD helps with generalization. Subsequent works have shown similar phenomena for regression (Nakkiran, 2020; Wu et al., 2021; Ba et al., 2022), kernel ridge regression Beugnot et al. (2022), and linear diagonal networks Nacson et al. (2022). However, the large step sizes considered in these work still fall under the scope of descent lemma, and most prior works do not theoretically investigate the effect of large step size in the EoS regime. A notable exception is the work of Wang et al. (2022a), which studies the impact of learning rates greater than 2 / smoothness for a matrix factorization problem. Also, the recent work of Andriushchenko et al. (2023) seeks to explain the generalization benefit of SGD in the large step size regime by relying on a heuristic SDE model for the case of linear diagonal networks. Despite this similarity, their main scope is quite different from ours, as we (i) focus on GD instead of SGD and (ii) establish a direct and detailed analysis of the GD dynamics for a model of the motivating sparse coding example.

2 Single-neuron linear network

In this section, we analyze the single-neuron linear network model ( x, y ) ↦→ f ( x, y ) := ℓ ( x × y ) .

2.1 Basic properties and assumptions

Basic properties. If ℓ is minimized at 0 , then the global minimizers of f are the x - and y -axes. The GD iterates x t , y t , for step size η > 0 and iteration t ≥ 0 can be written as

Assumptions. From here onward, we assume η < 1 and the following conditions on ℓ : R → R .

(A2) There exist constants β > 1 and c > 0 with the following property: for all s = 0 ,

(A1) ℓ is convex, even, 1 -Lipschitz, and of class C 2 near the origin with ℓ ′′ (0) = 1 .

̸

̸

Assumption (A2) imposes decay of s ↦→ ℓ ′ ( s ) /s locally away from the origin in order to obtain more fine-grained results on the limiting sharpness in Theorem 2. As we show in Lemma 5 below, when ℓ is smooth and has a strictly negative fourth derivative at the origin, then Assumption (A2) holds with β = 2 . See Example 1 for some simple examples of losses satisfying our assumptions.

2.2 Two different regimes for GD depending on the step size

Before stating rigorous results, in this section we begin by giving an intuitive understanding of the GD dynamics. It turns out that for a given initialization ( x 0 , y 0 ) , there are two different regimes for the GD dynamics depending on the step size η . Namely, there exists a threshold on the step size such that (i) below the threshold, GD remains close to the gradient flow for all time, and (ii) above the threshold, GD enters the edge of stability and diverges away from the gradient flow. See Figure 9.

First, recall that the GD dynamics are symmetric in x, y and that the lines y = ± x are invariant. Hence, we may assume without loss of generality that

From the expression (8) for the Hessian of f and our normalization ℓ ′′ (0) = 1 , it follows that the sharpness (the largest eigenvalue of loss Hessian) reached by GD in this example is precisely y 2 ∞ .

Initially, in both regimes, the GD dynamics tracks the continuous-time gradient flow. Our first observation is that the gradient flow admits a conserved quantity, thereby allowing us to predict the dynamics in this initial phase.

Lemma 1 (conserved quantity) . Along the gradient flow for f , the quantity y 2 -x 2 is conserved.

At this point, GD either approximately converges to the gradient flow solution (0 , √ y 2 0 -x 2 0 ) or diverges away from it, depending on whether or not y 2 t 0 > 2 /η . To see this, for | x t 0 y t 0 |≪ 1 , we can Taylor expand ℓ ′ near zero to obtain the approximate dynamics for x (recalling ℓ ′′ (0) = 1 ),

Lemma 1 implies that the gradient flow converges to (0 , y GF ∞ ) = (0 , √ y 2 0 -x 2 0 ) . For GD with step size η > 0 , the quantity y 2 -x 2 is no longer conserved, but we show in Lemma 6 that it is approximately conserved until the GD iterate lies close to the y -axis. Hence, GD initialized at ( x 0 , y 0 ) also reaches the y -axis approximately at the point ( x t 0 , y t 0 ) ≈ (0 , √ y 2 0 -x 2 0 ) .

From (4), we deduce the following conclusions.

(i) If y 2 t 0 < 2 /η , then | 1 -ηy 2 t 0 | < 1 . Since y t is decreasing, it implies that | 1 -ηy 2 t | < 1 for all t ≥ t 0 , and so | x t | converges to zero exponentially fast.

(ii) On the other hand, if y 2 t 0 > 2 /η , then | 1 -ηy 2 t 0 | > 1 , i.e., the magnitude of x t 0 increases in the next iteration, and hence GD cannot stabilize. In fact, in the approximate dynamics, x t 0 +1 has the opposite sign as x t 0 , i.e., x t 0 jumps across the y -axis. One can show that the 'bouncing' of the x variable continues until y 2 t has decreased past 2 /η , at which point we are in the previous case and GD approximately converges to (0 , 2 /η ) .

This reasoning, combined with the expression for the Hessian of f , shows that

Accordingly, we refer to the case y 2 0 -x 2 0 < 2 /η as the gradient flow regime , and the case y 2 0 -x 2 0 > 2 /η as the EoS regime .

See Figure 5 and Figure 9 for illustrations of these two regimes; see also Figure 10 for detailed illustrations of the EoS regime. In the subsequent sections, we aim to make the above reasoning rigorous. For example, instead of the approximate dynamics (4), we consider the original GD dynamics and justify the Taylor approximation. Also, in the EoS regime, rather than loosely asserting that | x t | ↘ 0 exponentially fast and hence the dynamics stabilizes 'quickly' once y 2 t < 2 /η , we track precisely how long this convergence takes so that we can bound the gap between the limiting sharpness and the prediction 2 /η .

2.3 Results

Gradient flow regime. Our first rigorous result is that when y 2 0 -x 2 0 = (2 -δ ) / η for some constant δ ∈ (0 , 2) , then the limiting sharpness of GD with step size η is y 2 0 -x 2 0 + O (1) = (2 -δ ) / η + O (1) , which is precisely the sharpness attained by the gradient flow up to a controlled error term.

In fact, our theorem is slightly more general, as it covers initializations in which δ can scale mildly with η . The precise statement is as follows.

Theorem 1 (gradient flow regime; see Subsection C.2) . Suppose we run GD with step size η > 0 on the objective f , where f ( x, y ) := ℓ ( xy ) , and ℓ satisfies Assumptions (A1) and (A2). Let (˜ x, ˜ y ) ∈ R 2 satisfy ˜ y > ˜ x > 0 with ˜ y 2 -˜ x 2 = 1 . Suppose we initialize GD at ( x 0 , y 0 ) := ( 2 -δ η ) 1 / 2 (˜ x, ˜ y ) , where δ ∈ (0 , 2) and η ≲ δ 1 / 2 ∧ (2 -δ ) . Then, GD converges to (0 , y ∞ ) satisfying where the implied constants depend on ˜ x , ˜ y , and ℓ , but not on δ , η .

The proof of Theorem 1 is based on a two-stage analysis. In the first stage, we use Lemma 6 on the approximate conservation of y 2 -x 2 along GD in order to show that GD lands near the y -axis with y 2 t 0 ≈ (2 -δ ) / η . In the second stage, we use the assumptions on ℓ in order to control the rate of convergence of | x t | to 0 , which is subsequently used to control the final deviation of y 2 ∞ from (2 -δ ) / η .

EoS regime. Our next result states that when y 2 0 -x 2 0 > 2 /η , then the limiting sharpness of GD is close to the EoS prediction of 2 /η , up to an error term which depends on the exponent β in (A2).

Theorem 2 (EoS; see Subsection C.4) . Suppose we run GD on f with step size η > 0 , where f ( x, y ) := ℓ ( xy ) , and ℓ satisfies (A1) and (A2). Let (˜ x, ˜ y ) ∈ R 2 satisfy ˜ y > ˜ x > 0 with ˜ y 2 -˜ x 2 = 1 . Suppose we initialize GD at ( x 0 , y 0 ) := √ (2+ δ ) / η (˜ x, ˜ y ) , where δ > 0 is a constant. Also, assume that for all t ≥ 1 such that y 2 t > 2 /η , we have x t = 0 . Then, GD converges to (0 , y ∞ ) satisfying

̸

where the implied constants depend on ˜ x , ˜ y , δ ∧ 1 , and ℓ , but not on η .

Remarks on the assumptions. The initialization in our results is such that both y 0 and y 0 -x 0 are on the same scale, i.e., y 0 , y 0 -x 0 = Θ(1 / √ η ) . This rules out extreme initializations such as y 0 ≈ x 0 , which are problematic because they lie too close to the invariant line y = x . Since our aim in this work is not to explore every edge case, we focus on this setting for simplicity. Moreover, we imposed the assumption that the iterates of GD do not exactly hit the y -axis before crossing y 2 = 2 /η . This is necessary because if x t = 0 for some iteration t , then ( x t ′ , y t ′ ) = ( x t , y t ) for all t ′ > t , and hence the limiting sharpness may not be close to 2 /η . This assumption holds generically, e.g., if we perturb each iterate of GD with a vanishing amount of noise from a continuous distribution, and we conjecture that for any η > 0 , the assumption holds for all but a measure zero set of initializations.

When β = + ∞ , which is the case for the Huber loss in Example 1, the limiting sharpness is 2 /η + O (1) . When β = 2 , which is the case for the logistic and square root losses in Example 1, the limiting sharpness is 2 /η + O ( η ) . Numerical experiments show that our error bound of O ( η 1 / ( β -1) ) is sharp ; see Figure 11 below.

We make a few remark about the proof. As we outline the proof in Subsection C.3, in turns out in order to bound the gap 2 /η -y 2 ∞ , the proof requires a control of the size | x t y t | , where t is the first iteration such that y 2 t crosses 2 /η . However, controlling the size of | x t y t | is surprisingly delicate as it requires a fine-grained understanding of the bouncing phase. The insight that guides the proof is the observation that during the bouncing phase, the GD iterates lie close to a certain envelope (Figure 9).

As a by-product of our analysis, we obtain a rigorous version of the quasi-static principle from which can more accurately track the sharpness gap and convergence rate (see Subsection C.5). The results of Theorem 1, Theorem 2, and Theorem 5 are displayed pictorially as Figure 9.

3 Understanding the bias evolution of the ReLU network

In this section, we use the insights from Section 2 to answer our main question, namely understanding the role of a large step size in learning threshold neurons for the ReLU network (2). Based on the observed dynamics (Figure 2), we can make our question more concrete as follows.

Q. What is the role of a large step size during the ' initial phase ' of training in which (i) the bias b rapidly decreases and (ii) the sum of weights a -+ a + oscillates?

3.1 Approximating the initial phase of GD with the 'mean model'

Deferring details to Appendix D, the GD dynamics for the ReLU network (2) in the initial phase are well-approximated by

where ℓ sym ( s ) := 1 2 (log(1 + exp( -s )) + log(1 + exp(+ s ))) and g ( b ) := E z ∼N (0 , 1) ReLU( z + b ) is the 'smoothed' ReLU. The GD dynamics can be compactly written in terms of the parameter A t := d ( a -t + a + t ) .

We call these dynamics the mean model . Figure 8 shows that the mean model closely captures the GD dynamics for the ReLU network (2), and we henceforth focus on analyzing the mean model.

The main advantage of the representation (5) is that it makes apparent the connection to the single-neuron example that we studied in Section 2. More specifically, (5) can be interpreted

4

2

0

2

4

0

1

2

3

4

5

as the 'rescaled' GD dynamics on the objective ( A,b ) ↦→ ℓ sym ( Ag ( b )) , where the step size for the A -dynamics is multiplied by 2 d 2 . Due to this resemblance, we can apply the techniques from Section 2.

-

0

.

25

-

0

.

20

-

0

.

15

-

0

.

10

-

0

.

05

0

.

00

b

η

= 2

/

2000

η

= 2

/

1000

-

5

0

5

0

10

20

30

40

50

iteration

0

20

A

-

0

.

4

-

0

.

3

-

0

.

2

-

0

.

1

0

.

0

b

η

= 2

/

2000

η

= 2

/

1000

-

5

0

5

0

100

200

300

400

500

iteration

0

20

A

3.2 Two different regimes for the mean model

Throughout the section, we use the shorthand ℓ := ℓ sym , and focus on initializing wiht a ± 0 = Θ(1 /d ) , a -+ a + = 0 , and b 0 = 0 . This implies A 0 = Θ(1) . We also note the following fact for later use.

Lemma 2 (formula for the smoothed ReLU; see Subsection E.1) . The smoothed ReLU function g can be expressed in terms of the PDF φ and the CDF Φ of the standard Gaussian distribution as g ( b ) = φ ( b ) + b Φ ( b ) . In particular, g ′ = Φ .

̸

Note also that b t is monotonically decreasing. This is because ℓ ′ ( A t g ( b t )) A t g ′ ( b t ) ≥ 0 since ℓ ′ is an odd function and g ( b ) , g ′ ( b ) > 0 for any b ∈ R .

Following Subsection 2.2, we begin with the continuous-time dynamics of the mean model:

Lemma 3 (conserved quantity; see Subsection E.1) . Let κ : R → R be defined as κ ( b ) := ∫ b 0 g/g ′ . Along the gradient flow (6) , the quantity 1 2 A 2 -2 d 2 κ ( b ) is conserved.

Based on Lemma 3, if we initialize the continuous-time dynamics (6) at ( A 0 , 0) and if A t → 0 , then the limiting value of the bias b GF ∞ satisfies κ ( b GF ∞ ) = -1 4 d 2 A 2 0 , which implies that b GF ∞ = -Θ( 1 d 2 ) ; indeed, this holds since κ ′ (0) = g (0) /g ′ (0) > 0 , so there exist constants c 0 , c 1 > 0 such that c 0 b ≤ κ ( b ) ≤ c 1 b for all -1 ≤ b ≤ 0 . Since the mean model (5) tracks the continuous-time dynamics (6) until it reaches the b -axis, the mean model initialized at ( A 0 , 0) also approximately reaches ( A t 0 , b t 0 ) ≈ (0 , -Θ( 1 d 2 )) ≈ (0 , 0) in high dimension d ≫ 1 . In other words, the continuoustime dynamics (6) fails to learn threshold neurons .

Once the mean model reaches the b -axis, we again identify two different regimes depending on the step size. A Taylor expansion of ℓ ′ around the origin yields the following approximate dynamics (here ℓ ′′ (0) = 1 / 4 ): A t 0 +1 ≈ A t 0 -ηd 2 2 A t 0 g ( b t 0 ) 2 = A t 0 ( 1 -ηd 2 2 g ( b t 0 ) 2 ) . We conclude that the condition which now dictates whether we have bouncing or convergence is 1 2 d 2 g ( b t 0 ) 2 > 2 /η .

(i) Gradient flow regime: If 2 /η > d 2 g (0) 2 / 2 = d 2 / (4 π ) (since g (0) 2 = 1 / (2 π ) ), i.e., the step size η is below the threshold 8 π/d 2 , then the final bias of the mean model b MM ∞ satisfies b MM ≈ b GF ≈ 0 . In other words, the mean model fails to learn threshold neurons .

∞ ∞ (ii) EoS regime: If 2 /η < d 2 / (4 π ) , i.e., the step size η is above the threshold 8 π/d 2 , then 1 2 d 2 g 2 ( b MM ∞ ) < 2 /η , i.e., b MM ∞ < g -1 (2 / √ ηd 2 ) . For instance, if η = 10 π d 2 , then b MM ∞ < -0 . 087 . In other words, the mean model successfully learns threshold neurons .

3.3 Results for the mean model

Theorem 3 (mean model, gradient flow regime; see Appendix E) . Consider the mean model (5) initialized at ( A 0 , 0) , with step size η = (8 -δ ) π d 2 for some δ > 0 . Let γ := 1 200 min { δ, 8 -δ, 8 -δ | A 0 | } . Then, as long as η γ/ A , the limiting bias b MM satisfies

≤ | 0 | ∞

In other words, the mean model fails to learn threshold neurons.

Theorem 4 (mean model, EoS regime; see Appendix E) . Consider the mean model initialized at ( A 0 , 0) , with step size η = (8+ δ ) π d 2 for some δ > 0 . Furthermore, assume that for all t ≥ 1 such that 1 2 d 2 g ( b t ) 2 > 2 /η , we have A t = 0 . Then, the limiting bias b MM ∞ satisfies

̸

For instance, if η = 10 π d 2 , then b MM ∞ < -0 . 087 . In other words, the mean model successfully learns threshold neurons.

4 Conclusion

In this paper, we present the first explanation for the emergence of threshold neuron (i.e., ReLU neurons with negative bias) in models such as the sparse coding model (1) through a novel connection with the 'edge of stability' (EoS) phenomenon. Along the way, we obtain a detailed and rigorous understanding of the dynamics of GD in the EoS regime for a simple class of loss functions, thereby shedding light on the impact of large learning rates in non-convex optimization.

Our approach is largely inspired by the recent paradigm of ' physics-style ' approaches to understanding deep learning based on simplified models and controlled experiments (c.f. (Zhang et al., 2022; von Oswald et al., 2023; Abernethy et al., 2023; Allen-Zhu and Li, 2023; Li et al., 2023; Ahn et al., 2023a,b)). We found such physics-style approach quite effective to understand deep learning, especially given the complexity of modern deep neural networks. We hope that our work inspires further research on understanding the working mechanisms of deep learning.

Many interesting questions remain, and we conclude with some directions for future research.

Extending the analysis of EoS to richer models. Although the analysis we present in this work is restricted to simple models, the underlying principles can potentially be applied to more general settings. In this direction, it would be interesting to study models which capture the impact of the depth of the neural network on the EoS phenomenon. Notably, a follow-up work by Song and Yun (2023) uses bifurcation theory to extend our results to more complex models.

The interplay between the EoS and the choice of optimization algorithm. As discussed in Subsection 2.3, the bouncing phase of the EoS substantially slows down the convergence of GD (see Figure 11). Investigating how different optimization algorithm (e.g., SGD, or GD with momentum) interact with the EoS phenomenon could potentially lead to practical speed-ups or improved generalization. Notably, a follow up work by Dai et al. (2023) studies the working mechanisms of a popular modern optimization technique called sharpness-aware minimization (Foret et al., 2021) based on our sparse coding problem.

An end-to-end analysis of the sparse coding model. Finally, we have left open the motivating question of analyzing how two-layer ReLU networks learn to solve the sparse coding model (1). Despite the apparent simplicity of the problem, its analysis has thus far remained out of reach, and we believe that a resolution to this question would constitute compelling and substantial progress towards understanding neural network learning. We are hopeful that the insights in this paper provide the first step towards this goal.

Acknowledgments

We thank Ronan Eldan, Suriya Gunasekar, Yuanzhi Li, Jonathan Niles-Weed, and Adil Salim for initial discussions on this project. KA was supported by the ONR grant (N00014-20-1-2394) and MIT-IBM Watson as well as a Vannevar Bush fellowship from Office of the Secretary of Defense. SC was supported by the NSF TRIPODS program (award DMS-2022448).

References

Jacob Abernethy, Alekh Agarwal, Teodor V. Marinov, and Manfred K. Warmuth. A mechanism for sample-efficient in-context learning for sparse retrieval tasks. arXiv preprint arXiv:2305.17040 , 2023.

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, 7 2022.

Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 2023 (arXiv:2306.00297) , 2023a.

Kwangjun Ahn, Xiang Cheng, Minhak Song, Chulhee Yun, Ali Jadbabaie, and Suvrit Sra. Linear attention is (maybe) all you need (to understand transformer optimization). arXiv 2310.01082 , 2023b.

Zeyuan Allen-Zhu and Yuanzhi Li. Feature purification: how adversarial training performs robust deep learning. In 2021 IEEE 62nd Annual Symposium on Foundations of Computer Science (FOCS) , pages 977-988. IEEE, 2022.

Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: part 1, context-free grammar. arXiv preprint arXiv:2305.13673 , 2023.

Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via overparameterization. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning , volume 97 of Proceedings of Machine Learning Research , pages 242-252. PMLR, 6 2019.

Maksym Andriushchenko, Aditya V. Varre, Loucas Pillaud-Vivien, and Nicolas Flammarion. SGD with large step sizes learns sparse features. 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 903-925. PMLR, 7 2023.

Sanjeev Arora, Rong Ge, Tengyu Ma, and Ankur Moitra. Simple, efficient, and neural algorithms for sparse coding. In Peter Grünwald, Elad Hazan, and Satyen Kale, editors, Proceedings of the 28th Conference on Learning Theory , volume 40 of Proceedings of Machine Learning Research , pages 113-149, Paris, France, 7 2015. PMLR.

Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning , volume 97 of Proceedings of Machine Learning Research , pages 322-332. PMLR, 6 2019.

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, 7 2022.

Jimmy Ba, Murat A. Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. Highdimensional asymptotics of feature learning: how one gradient step improves the representation. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022.

Gaspard Beugnot, Julien Mairal, and Alessandro Rudi. On the benefits of large learning rates for kernel methods. In Po-Ling Loh and Maxim Raginsky, editors, Proceedings of Thirty Fifth Conference on Learning Theory , volume 178 of Proceedings of Machine Learning Research , pages 254-282. PMLR, 7 2022.

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, 7 2023.

Lénaïc Chizat. Mean-field Langevin dynamics: exponential convergence and annealing. Transactions on Machine Learning Research , 2022.

Lénaïc Chizat and Francis Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. CesaBianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 31. Curran Associates, Inc., 2018.

Lénaïc Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. 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.

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.

Yan Dai, Kwangjun Ahn, and Suvrit Sra. The crucial role of normalization in sharpness-aware minimization. NeurIPS 2023 (arXiv:2305.15287) , 2023.

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.

Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. In International Conference on Learning Representations , 2019.

Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In International Conference on Learning Representations , 2021.

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: convergence and generalization in neural networks. In Proceedings of the 32nd Advances in Neural Information Processing Systems , pages 8580-8589, 2018.

Stanislaw Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Width of minima reached by stochastic gradient descent is influenced by learning rate to batch size ratio. In Vˇ era K˚ urková, Yannis Manolopoulos, Barbara Hammer, Lazaros Iliadis, and Ilias Maglogiannis, editors, Artificial Neural Networks and Machine Learning - ICANN 2018 , pages 392-402, Cham, 2018. Springer International Publishing.

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.

Stanislaw Jastrzebski, Devansh Arpit, Oliver Astrand, Giancarlo B. Kerg, Huan Wang, Caiming Xiong, Richard Socher, Kyunghyun Cho, and Krzysztof J Geras. Catastrophic Fisher explosion: early phase Fisher matrix impacts generalization. In International Conference on Machine Learning , pages 4772-4784. PMLR, 2021.

Stanisław Jastrzebski, 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.

Stefani Karp, Ezra Winston, Yuanzhi Li, and Aarti Singh. Local signal adaptivity: provable feature learning in neural networks beyond kernels. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , volume 34, pages 24883-24897. Curran Associates, Inc., 2021.

Frederic Koehler and Andrej Risteski. The comparative power of ReLU networks and polynomial kernels in the presence of sparse latent structure. In International Conference on Learning Representations , 2018.

Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.02218 , 2020.

Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks. Advances in Neural Information Processing Systems , 32, 2019.

Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: towards a mechanistic understanding. 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 19689-19729. PMLR, 7 2023.

Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normalization layers: sharpness reduction. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems , volume 35, pages 34689-34708. Curran Associates, Inc., 2022.

Chao Ma, Daniel Kunin, Lei Wu, and Lexing Ying. Beyond the quadratic approximation: the multiscale structure of neural network loss landscapes. Journal of Machine Learning , 1(3): 247-267, 2022.

Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. 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 2388-2464. PMLR, 6 2019.

Rotem Mulayoff, Tomer Michaeli, and Daniel Soudry. The implicit bias of minima stability: a view from function space. Advances in Neural Information Processing Systems , 34:17749-17761, 2021.

Mor Shpigel Nacson, Kavya Ravichandran, Nathan Srebro, and Daniel Soudry. Implicit bias of the step size in linear diagonal neural networks. In International Conference on Machine Learning , pages 16270-16295. PMLR, 2022.

Preetum Nakkiran. Learning rate annealing can provably help generalization, even for convex problems. arXiv preprint arXiv:2005.07360 , 2020.

Atsushi Nitanda, Denny Wu, and Taiji Suzuki. Convex analysis of the mean field Langevin dynamics. In Gustau Camps-Valls, Francisco J. R. Ruiz, and Isabel Valera, editors, Proceedings of the 25th International Conference on Artificial Intelligence and Statistics , volume 151 of Proceedings of Machine Learning Research , pages 9741-9757. PMLR, 3 2022.

Bruno A Olshausen and David J Field. Sparse coding with an overcomplete basis set: a strategy employed by v1? Vision Research , 37(23):3311-3325, 1997.

Bruno A Olshausen and David J Field. Sparse coding of sensory inputs. Current Opinion in Neurobiology , 14(4):481-487, 2004.

Samet Oymak and Mahdi Soltanolkotabi. Toward moderate overparameterization: global convergence guarantees for training shallow neural networks. IEEE Journal on Selected Areas in Information Theory , 1(1):84-105, 2020.

Grant M. Rotskoff and Eric Vanden-Eijnden. Trainability and accuracy of artificial neural networks: an interacting particle system approach. Comm. Pure Appl. Math. , 75(9):1889-1935, 2022.

Minhak Song and Chulhee Yun. Trajectory alignment: understanding the edge of stability phenomenon via bifurcation theory. NeurIPS 2023 (arXiv:2307.04204) , 2023.

William E. Vinje and Jack L. Gallant. Sparse coding and decorrelation in primary visual cortex during natural vision. Science , 287(5456):1273-1276, 2000.

Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning , pages 35151-35174. PMLR, 2023.

Yuqing Wang, Minshuo Chen, Tuo Zhao, and Molei Tao. Large learning rate tames homogeneity: convergence and balancing effect. In International Conference on Learning Representations , 2022a.

Zixuan Wang, Zhouzi Li, and Jian Li. Analyzing sharpness along GD trajectory: progressive sharpening and edge of stability. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems , volume 35, pages 9983-9994. Curran Associates, Inc., 2022b.

Jingfeng Wu, Difan Zou, Vladimir Braverman, and Quanquan Gu. Direction matters: on the implicit bias of stochastic gradient descent with moderate learning rate. In International Conference on Learning Representations , 2021.

Lei Wu, Chao Ma, and Weinan E. How SGD selects the global minima in over-parameterized learning: a dynamical stability perspective. Advances in Neural Information Processing Systems , 31:8279-8288, 2018.

Chen Xing, Devansh Arpit, Christos Tsirigotis, and Yoshua Bengio. A walk with SGD. arXiv preprint arXiv:1802.08770 , 2018.

Jianchao Yang, Kai Yu, Yihong Gong, and Thomas Huang. Linear spatial pyramid matching using sparse coding for image classification. In 2009 IEEE Conference on Computer Vision and Pattern Recognition , pages 1794-1801. IEEE, 2009.

Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling transformers with lego: a synthetic reasoning task. arXiv preprint arXiv:2206.04301 , 2022.

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.

Appendix

A Additional illustrations for Section 2

In this section, we provide some illustrations of the results presented in Section 2. We first illustrate the two different regimes of GD presented in Subsection 2.2.

(

x

0

, y

0

)

√

2

/η

-

√

y

2

0

-

x

2

0

-

y

(

x

0

, y

0

)

√

2

/η

-

√

y

2

0

-

x

2

0

-

y

x

Next, we next present detailed illustrations of the edge-of-stability regime depending on the choice of step size. Compare this plot with our theoretical results characterized in Theorem 2.

y

√

2

/η

1

-

√

2

/η

2

-

√

2

/η

3

-

√

2

/η

4

-

(

x

0

, y

0

)

y

√

2

/η

1

-

√

2

/η

2

-

√

2

/η

3

-

√

2

/η

4

-

(

x

0

, y

0

)

y

√

2

/η

1

-

√

2

/η

2

-

√

2

/η

3

-

√

2

/η

4

-

(

x

0

, y

0

)

y

√

2

/η

1

-

√

2

/η

2

-

√

2

/η

3

-

√

2

/η

4

-

(

x

0

, y

0

)

0

.

0

2

.

5

5

.

0

7

.

5

10

.

0

12

.

5

15

.

0

time (

η

x iteration)

9

7

5

3

sharpness

Lastly, we present another detailed illustrations of Theorem 2 in terms of its dependence on β .

0

-

5

-

4

-

3

-

2

-

1

log(

η

)

-

7

-

6

-

5

-

4

-

3

-

2

-

1

0

log(2

/η

-

y

2

∞

)

gap

∝

η

2

gap

∝

η

1

gap

∝

η

1

/

2

gap

∝

η

1

/

9

β

=

3

2

β

= 2

β

= 3

β

= 10

-

5

.

0

-

4

.

5

-

4

.

0

-

3

.

5

-

3

.

0

-

2

.

5

-

2

.

log(

η

)

4

6

8

10

12

log(

T

)

T

∝

η

-

4

T

∝

η

-

3

T

∝

η

-

2

log(

η

-

1

)

T

∝

η

-

2

β

=

4

3

β

=

3

2

β

= 2

β

= 4

B Further experimental results

In this section, we report further experimental results which demonstrate that our theory, while limited to the specific models we study (namely, the single-neuron example and the mean model), is in fact

indicative of behaviors commonly observed in more realistic instances of neural network training. In particular, we show that threshold neurons often emerge in the presence of oscillations in the other weight parameters of the network.

B.1 Experiments for the full sparse coding model

We provide the details for the top plot of Figure 3. consider the sparse coding model in the form (1). Compared to (2), we assume that the basis vectors are unknown, and the neural network learn them through additional parameters W = ( w i ) i m =1 together with m different weights a = ( a i ) i m =1 for the second layer as follows:

We show results for d = 100 , m = 2000 . We generate n = 20000 data points according to the aforementioned sparse coding model with λ = 5 . We use the He initialization, i.e., a ∼ N (0 , I m /m ) , w ∼ N (0 , I d /d ) , and b = 0 . As shown in the top plot of Figure 3, the bias decreases more with the large learning rate. Further, we report the behavior of the average of second layer weights in Figure 12 (left), and confirm that the sum oscillates.

0

5

10

15

20

25

30

step

60

50

40

30

20

10

0

10

avg. 2nd layer weight

lr=0.03

lr=0.04

lr=0.05

lr=0.07

lr=0.2

0

10

20

30

40

50

60

step

15

10

5

0

5

10

15

20

25

avg. logit

lr=0.2

lr=0.05

B.2 Experiments on the CIFAR-10 dataset

Next, we provide the details for the bottom plot of Figure 3. We train ResNet-18 on a binarized version of the CIFAR-10 dataset formed by taking only the first two classes; this is done for the purpose of monitoring the average logit of the network. The average logit is measured over the entire training set. The median bias is measured at the last convolutional layer right before the pooling. For the optimizer, we use full-batch GD with no momentum or weight decay, plus a cosine learning rate scheduler where learning rates shown in the plots are the initial values.

Oscillation of expected output (logit) of the network. Bearing a striking resemblance to our two-layer models, as one can see from Figure 12 (right) that the expected mean of the output (logit) of the deep net also oscillates due to GD dynamics. As we have argued in the previous sections, this occurs as the bias parameters are driven towards negative values.

Results for SGD. In Figure 13, we report qualitatively similar phenomena when we instead train ResNet-18 with stochastic gradient descent (SGD), where we use all ten classes of CIFAR-10. Again, the median bias is measured at the last convolutional layer. We further report the average activation which is the output of the ReLU activation at the last convolutional layer, averaged over the neurons and the entire training set. The average activation statistics represent the hidden representations before the linear classifier part, and lower values represent sparser representations. Interestingly, the threshold neuron also emerges with larger step sizes similarly to the case of gradient descent.

0

50

100

150

200

250

300

350

400

step

0.2

0.4

0.6

0.8

avg. activation

0

50

100

150

200

250

300

350

400

step

0.03

0.02

0.01

0.00

median bias

lr=0.2

lr=0.05

lr=0.01

0

50

100

150

200

250

300

350

400

step

0.2

0.3

0.4

0.5

0.6

0.7

0.8

avg. activation

0

50

100

150

200

250

300

350

400

step

0.015

0.010

0.005

0.000

median bias

lr=0.2

lr=0.05

lr=0.01

C Proofs for the single-neuron linear network

We start by describing basic and relevant properties of the model and the assumptions on ℓ .

Basic properties. If ℓ is minimized at 0 , then the global minimizers of f are the x - and y -axes. The gradient and Hessian of f are given by:

This results in GD iterates x t , y t , for step size η > 0 and iteration t ≥ 0 :

Lemma 4 (invariant lines) . Assume that ℓ is even, so that ℓ ′ is odd. Then, the lines y = ± x are invariant for gradient descent on f .

Proof. If y t = ± x t , then

and hence y t +1 = ± x t +1 . Note that the iterates ( x t ) t ≥ 0 are the iterates of GD with step size η on the one-dimensional loss function x ↦→ 1 2 ℓ ( x 2 ) .

We focus instead on initializing away from these two lines. We now state our assumptions on ℓ .

We gather together some elementary properties of ℓ .

Lemma 5 (properties of ℓ ) . Suppose that Assumption (A1) holds.

ℓ is minimized at the origin and ℓ ′ (0) = 0 .

Suppose that ℓ is four times continuously differentiable near the origin. If Assumption (A2) holds, then ℓ (4) (0) ≤ 0 . Conversely, if ℓ (4) (0) < 0 , then Assumption (A2) holds for β = 2 .

Proof. The first statement is straightforward. The second statement follows from Taylor expansion: for s = 0 near the origin,

̸

Since ℓ ′′′ is odd, then Assumption (A2) and (9) imply that ℓ ′′′ is non-positive on (0 , ε ) for some ε > 0 , which in turn implies ℓ (4) (0) ≤ 0 . Conversely, if ℓ (4) (0) < 0 , then there exists ε > 0 such that ℓ ′′′ ( s ) ≤ -εs for s ∈ (0 , ε ) . From (9), we see that ℓ ′ ( s ) /s ≤ 1 -ε ∫ s 0 ( s -r ) d r ≤ 1 -εs 2 / 2 . By symmetry, we conclude that Assumption (A2) holds with β = 2 and some c > 0 .

We give some simple examples of losses satisfying our assumptions.

Example 1. The examples below showcase several functions ℓ that satisfy Assumptions (A1) and (A2) with different values of β .

Rescaled and symmetrized logistic loss. ℓ rsym ( s ) := 1 2 ℓ logi ( -2 s ) + 1 2 ℓ logi (+2 s ) . Note ℓ ′ rsym ( s ) = tanh( s ) , thus ℓ ′ rsym ( s ) /s ≤ 1 and ℓ ′ rsym ( s ) /s ≤ 1 -1 4 | s | 2 , for | s | < 1 4 .

Huber loss. ℓ Hub ( s ) := s 2 2 ✶ { s ∈ [ -1 , 1] } + ( | s | -1 2 ) ✶ { s / ∈ [ -1 , 1] } . Note ℓ ′ Hub ( s ) = s ✶ { s ∈ [ -1 , 1] } +sgn( s ) ✶ { s / ∈ [ -1 , 1] } , thus ℓ ′ Hub ( s ) /s ≤ 1 , i.e., we have Assumption (A2) with β = + ∞ . · Higher-order. For β > 1 let c β := 1 β +1 β β +1 β and r β := β +1 β . We define ℓ β implicitly via its

Square root loss. ℓ sqrt ( s ) := √ 1 + s 2 . Note ℓ ′ sqrt ( s ) = s √ 1+ s 2 , thus ℓ ′ sqrt ( s ) /s ≤ 1 and ℓ ′ sqrt ( s ) /s ≤ 1 -2 5 | s | 2 , for | s | < 2 5 .

( ) derivative ℓ ′ β ( s ) := s ( 1 -c β | s | β ) ✶ { s 2 < r 2 β } +sgn( s ) ✶ { s 2 ≥ r 2 β } . By definition, ℓ ′ β ( s ) /s ≤ 1 and ℓ ′ β ( s ) /s ≤ 1 -c ℓ | s | β , where c ℓ = c β ∧ r β .

We now prove our main results from Subsection 2.3 in order.

C.1 Approximate conservation along GD

We begin by stating and proving the approximate conservation of y 2 -x 2 for the GD dynamics. Lemma 6 (approximately conserved quantity) . Let (˜ x, ˜ y ) ∈ R 2 be such that ˜ y > ˜ x > 0 with ˜ y 2 -˜ x 2 = 1 . Suppose that we run GD on f with step size η with initial point ( x 0 , y 0 ) := √ γ η (˜ x, ˜ y ) , for some γ > 0 . Then, there exists t 0 = O ( 1 η ) such that sup t ≥ t 0 | x t | ≤ O ( √ ( γ -1 ∨ γ ) η ) and where the implied constant depends on ˜ x , ˜ y , and ℓ .

Proof. Let D t := y 2 t -x 2 t and note that

Since ℓ is 1 -Lipschitz, then D t +1 = (1 -O ( η 2 )) D t .

This shows that for t ≲ 1 /η 2 , we have y 2 t -x 2 t = D t ≳ D 0 = y 2 0 -x 2 0 ≍ γ/η . Since ℓ ′′ (0) = 1 , there exist constants c 0 , c 1 > 0 such that ℓ ′ ( | xy | ) ≥ ℓ ′ ( c 0 ) ≥ c 1 whenever | xy | ≥ c 0 . Hence, for all t ≥ 1 such that t ≲ 1 /η 2 , x t > 0 , and | x t y t | ≥ c 0 , we have y 2 t ≳ γ/η and

Since x 0 ≍ √ γ/η , this shows that after at most O (1 /η ) iterations, we must have either x t < 0 or | x t y t | ≤ c 0 for the first time. In the first case, (10) shows that | x t | ≲ √ γη . In the second case, since y 2 t ≳ γ/η , we have | x t | ≲ √ η/γ . Let t 0 denote the iteration at which this occurs.

Next, for iterations t ≥ t 0 , we use the dynamics (10) for x and the fact that ℓ ′ ( x t y t ) has the same sign as x t to conclude that there are two possibilities: either x t +1 has the same sign as x t , in which case | x t +1 | ≤ | x t | , or x t +1 has the opposite sign as x t , in which case | x t +1 | ≤ η | ℓ ′ ( x t y t ) | y t ≤ ηy t ≤ O ( √ γη ) . This implies sup t ≥ t 0 | x t | ≤ O ( √ ( γ -1 ∨ γ ) η ) as asserted.

C.2 Gradient flow regime

In this section, we prove Theorem 1. From Lemma 6, there exists an iteration t 0 such that | x t 0 | ≲ √ η/ (2 -δ ) and

In particular, C := | x t 0 y t 0 | ≲ 1 .

We prove by induction the following facts: for t ≥ t 0 ,

| x t | ≤ | x t 0 | exp( -Ω( α ( t -t 0 ))) , where α := min { δ, 2 -δ } .

| x t y t | ≤ C .

Suppose that these conditions hold up to iteration t ≥ t 0 . By Assumption (A2), we have | ℓ ′ ( s ) | ≤ | s | for all s = 0 . Therefore,

̸

In particular, 1 2 2 -δ η ≤ y 2 t ≤ 2 -δ/ 2 η throughout. In order for these assertions to hold, we require η 2 ≲ α (2 -δ ) , i.e., η ≲ min { √ δ, 2 -δ } .

Next, we would like to show that t ↦→| x t | is decaying exponentially fast. Since it suffices to consider the case when x t > 0 . Assumption (A2) implies that

For the upper bound, we split into two cases. We begin by observing that since ℓ is twice continuously differentiable near the origin with ℓ ′′ (0) = 1 , there is a constant ε 0 such that | s | < ε 0 implies | ℓ ′ ( s ) | ≥ 1 2 | s | . If s t := x t y t ≤ ε 0 , then

Otherwise, if s t ≥ ε 0 , then

Combining these inequalities, we obtain

This shows in particular that | x t | ↘ 0 , i.e., we have global convergence. To conclude the proof, observe that (11) gives a bound on the final sharpness.

This verifies the second statement in the induction. The first statement follows because both t ↦→| x t | and t ↦→ y t are decreasing.

Remark 1 . The proof also gives us estimates on the convergence rate. Namely, from Lemma 6, the initial phase in which we approach the y -axis takes O ( 1 η ) iterations. For the convergence phase, in order to achieve ε error, we need | x t | ≲ √ εη √ 2 -δ ; hence, the convergence phase needs only O ( 1 α log 1 ε ) iterations. Note that the rate of convergence in the latter phase does not depend on the step size η .

C.3 EoS regime: proof outline

We give a brief outline of the proof of Theorem 2: As before, Lemma 6 shows that GD reaches the y -axis approximately at (0 , √ y 2 0 -x 2 0 ) . At this point, x starts bouncing while y steadily decreases, and we argue that unless x t = 0 or y 2 t ≤ 2 /η , the GD dynamics cannot stabilize (see Lemma 7).

To bound the gap 2 /η -y 2 ∞ , we look at the first iteration t such that y 2 t crosses 2 /η . By making use of Assumption (A2), we simultaneously control both the convergence rate of | x t | to zero and the decrease in y 2 t in order to prove that

see Proposition 1. Therefore, to establish Theorem 2, we must bound | x t y t | at iteration t .

Controlling the size of | x t y t | , however, is surprisingly delicate as it requires a fine-grained understanding of the bouncing phase. The insight that guides the proof is the observation that during the bouncing phase, the GD iterates lie close to a certain envelope (Figure 9). This envelope is predicted by the quasi-static heuristic as described in Ma et al. (2022). Namely, suppose that after one iteration of GD, we have perfect bouncing: x t +1 = -x t . Substituting this into the GD dynamics, we obtain the equation

According to Assumption (A2), we have ℓ ′ ( x t y t ) = x t y t (1 -Ω( | x t y t | β )) , Together with (13), if y 2 t = (2 + δ t ) /η ≥ 2 /η , where δ t is sufficiently small, it suggests that

The quasi-static prediction (14) fails when δ t is too small. Nevertheless, we show that it remains accurate as long as δ t ≳ η β/ ( β -1) , and consequently we obtain | x t y t | ≲ η 1 / ( β -1) . Combined with (12), it yields Theorem 2.

C.4 EoS regime: crossing the threshold and the convergence phase

In this section, we prove Theorem 2. We first show that y 2 t must cross 2 / η in order for GD to converge, and we bound the size of the jump across 2 / η once this happens.

Throughout this section and the next, we use the following notation:

s t := x t y t ;

r t := ℓ ′ ( s t ) /s t .

In this notation, we can write the GD equations as

We also make a remark regarding Assumption (A2). If β < + ∞ , then Assumption (A2) is equivalent to the following seemingly strongly assumption: for all r > 0 , there exists a constant c ( r ) > 0 such that

Indeed, Assumption (A2) states that (A2 + ) holds for some r > 0 . To verify that (A2 + ) holds for some larger r ′ > r , we can split into two cases. If | s | ≤ r , then ℓ ′ ( s ) /s ≤ 1 -c | s | β . Otherwise, if | s | > r , then ℓ ′ ( r ) /r < 1 and the 1 -Lipschitzness of ℓ ′ imply that ℓ ′ ( s ) /s < 1 for r ≤ | s | ≤ r ′ , and hence ℓ ′ ( s ) /s ≤ 1 -c ′ | s | β , for a sufficiently small constant c ′ > 0 ; thus we can take c ( r ′ ) = c ∧ c ′ . Later, we will invoke (A2 + ) with r chosen to be a universal constant, so that c ( r ) can also be thought of as universal.

We begin with the following result about the limiting value of y t .

Lemma 7 (threshold crossing) . Let (˜ x, ˜ y ) ∈ R 2 satisfy ˜ y > ˜ x > 0 with ˜ y 2 -˜ x 2 = 1 . Suppose we initialize GD with step size η with initial point ( x 0 , y 0 ) := √ 2+ δ η (˜ x, ˜ y ) , where δ > 0 is a constant. Then either x t = 0 for some t or

Proof. Assume throughout that x t = 0 for all t . Recall the dynamics for y :

̸

By assumption ℓ ′ ( s ) /s → 1 as s → 0 , and ℓ ′ is increasing, so this equation implies that if lim inf t →∞ | x t | > 0 then y 2 t must eventually cross 2 /η .

Suppose for the sake of contradiction that there exists ε > 0 with y 2 t > (2 + ε ) /η , for all t . Let ε ′ > 0 be such that 1 -(2 + ε ) (1 -ε ′ ) < -1 , i.e., ε ′ < ε 2+ ε . Then, there exists δ > 0 such | x t | ≤ δ implies r t > 1 -ε ′ , hence

The above means that | x t | increases until it exceeds δ , i.e., lim inf t →∞ | x t | ≥ δ . This is our desired contradiction and it implies that lim t →∞ y 2 t ≤ 2 /η .

Lemma 8 (initial gap) . Suppose that at some iteration t , we have

Then, it holds that

Proof. We can bound

where we used the fact that | ℓ ′ ( s ) | ≤ | s | for all s ∈ R ,

The above lemma shows that the size of the jump across 2 /η is controlled by the size of | s t | at the time of the crossing. From Lemma 6, we know that | s t | ≲ 1 , where the implied constant depends on δ . Hence, the size of the jump is always O ( η ) .

We now provide an analysis of the convergence phase, i.e., after y 2 t crosses 2 /η .

Proposition 1 (convergence phase) . Suppose that y 2 t < 2 /η ≤ y 2 t -1 . Then, GD converges to (0 , y ∞ ) satisfying

Proof. Write y 2 t = (2 -ρ t ) /η , so that ρ t = 2 -ηy 2 t . We write down the update equations for x and for ρ . First, by the same argument as in the proof of Theorem 1, we have

Next, using r t ≤ 1 , which translates into

Using these two inequalities, we can conclude as follows. Let q > 0 be a parameter chosen later, and let t be the first iteration for which ρ t ≥ q (if no such iteration exists, then ρ t ≤ q for all t ). Note that ρ t ≤ q + O ( η | x t | ) due to (15) and (16). By (15), we conclude that for all t ′ ≥ t ,

Substituting this into (16),

By optimizing this bound over q , we find that for all t ,

Translating this result back into y 2 t yields the result.

Let us take stock of what we have established thus far.

According to Lemma 6, | s t | is bounded for all t by a constant.

Then, from Lemma 7 and Lemma 8, we must have either y 2 t → 2 /η , or 2 /η -O ( η ) ≤ y 2 t ≤ 2 /η for some iteration t .

In the latter case, Proposition 1 shows that the limiting sharpness is 2 /η -O (1) .

Note also that the analyses thus far have not made use of Assumption (A2), i.e., we have established the β = + ∞ case of Theorem 2. Moreover, for all β > 1 , the asymptotic 2 /η -O (1) still shows that the limiting sharpness is close to 2 /η , albeit with suboptimal rate. The reader who is satisfied with this result can then skip ahead to subsequent sections. The remainder of this section and the next section are devoted to substantial refinements of the analysis.

To see where improvements are possible, note that both Lemma 8 and Proposition 1 rely on the size of | s t | at the crossing. Our crude bound of | s t | ≲ 1 does not capture the behavior observed in experiments, in which | s t | ≲ η 1 / ( β -1) . By substituting this improved bound into Lemma 7, we would deduce that the gap at the crossing is O ( η 1+2 / ( β -1) ) , and then Proposition 1 would imply that the limiting sharpness is 2 /η -O ( η 1 / ( β -1) ) . Another weakness of our proof is that it provides nearly no information about the dynamics during the bouncing phase, which constitutes an incomplete understanding of the EoS phenomenon. In particular, we experimentally observe that during the bouncing phase, the iterates lie very close to the quasi-static envelope (Figure 9). In the next section, we will rigorously prove all of these observations.

Before doing so, however, we show that Proposition 1 can be refined by using Assumption (A2), which could be of interest in its own right. It shows that even if the convergence phase begins with a large value of | s t | , the limiting sharpness can be much closer to 2 /η than what Proposition 1 suggests. The following proposition combined with Lemma 6 implies Theorem 2 for all β > 2 , but it is insufficient for the case 1 < β ≤ 2 . From now on, we assume β < + ∞ .

Proposition 2 (convergence phase; refined) . Suppose that y 2 t < 2 /η ≤ y 2 t -1 . Then, GD converges to (0 , y ∞ ) satisfying

Proof. Let y 2 t = (2 -ρ t ) /η as before. We quantify the decrease of | x t | in terms of ρ t and conversely the increase of ρ t in terms of | x t | by tracking the half-life of | x t | , i.e., the number of iterations it takes | x t | to halve. We call these epochs: at the i -th epoch, we have

Let i 0 be the index of the first epoch, i.e., i 0 = ⌊ log 2 ( √ η/ | x t | ) ⌋ . Due to Lemma 6, we know that i 0 ≥ -O (1) . From (15), | x t | is monotonically decreasing and consequently | s t | is decreasing as

well. Also, our bound on the limiting sharpness implies that y 2 t > 1 /η for all t , provided that η is sufficiently small.

Let us now compute the dynamics of ρ t and | x t | . At epoch i , | x t | > 2 -( i +1) √ η hence | s t | > 2 -( i +1) . Assumption (A2 + ) with r = | s t | ≲ 1 implies that

where c = c ( | s t | ) . This allows to refine (15) on the decrease of | x t | to where the first inequality follows from (17) and the second from ρ t = 2 -ηy 2 t < 1 . In turn, this inequality shows that the i -th phase only requires O (2 βi ) iterations.

Hence, if t ( i ) denotes the start of the i -th epoch, then (16) shows that

Summing this up, we have

In the case of β < 2 , the final sharpness satisfies 2 /η -O ( ρ t /η ) -O ( η | s t | 2 -β ) ≤ y 2 ∞ ≤ 2 /η . In the other two cases, suppose that we use this argument until epoch i ⋆ such that 2 -i ⋆ ≍ η γ . Then, we have | x t ( i ⋆ ) | ≍ η γ +1 / 2 , | s t ( i ⋆ ) | ≍ η γ , and by using the argument from Proposition 1 from iteration t ( i ⋆ ) onward we obtain

We optimize over the choice of γ , obtaining γ = 1 / ( β -1) and thus

By collecting together the three cases and using Lemma 8 to bound ρ t , we finish the proof.

Using the crude bound | s t 0 | ≲ 1 from Lemma 6, it yields which is optimal for β > 2 .

C.5 EoS regime: quasi-static analysis

We supplement Assumption (A2) with a corresponding lower bound on ℓ ′ ( s ) /s :

̸

Under these assumptions, we prove the following result which is also of interest as it provides detailed information for the bouncing phase of the EoS.

Theorem 5 (quasi-static principle) . Suppose we run GD on f with step size η > 0 , where f ( x, y ) := ℓ ( xy ) and ℓ satisfies Assumptions (A1), (A2), and (A3) . Write y 2 t := (2 + δ t ) /η and suppose that at some iteration t 0 , we have | x t 0 y t 0 | ≍ δ 1 /β t 0 and δ t 0 ≲ 1 . Then, for all t ≥ t 0 with δ t ≳ η β/ ( β -1) , we have

where all implied constants depend on ℓ but not on η .

In this section, we show that the GD iterates lie close to the quasi-static trajectory and give the full proof of Theorem 2. Recall from (13) that the quasi-static analysis predicts

and that during the bouncing phase, this closely agrees with experimental observations (Figure 9). We consider the phase where y 2 t has not yet crossed the threshold 2 /η and we write y 2 t := (2 + δ t ) /η , thinking of δ t as small. Then, (18) can be written (2 + δ t ) r t ≈ 2 . If we have the behavior ℓ ′ ( s ) /s = 1 -Θ( | s t | β ) near the origin, then r t ≈ 1 -Θ( δ t ) implies that

Our goal is to rigorously establish (19). However, we first make two observations. First, in order to establish Theorem 2, we only need to prove an upper bound on | s t | , which only requires Assumption (A2) (to prove a lower bound on | s t | , we need a corresponding lower bound on ℓ ′ ( s ) /s ). Second, even if we relax (19) to read | s t | β ≲ δ t , this fails to hold when δ t is too small, because the error terms (the deviation of the dynamics from the quasi-static trajectory) begin to dominate. With this in mind, we shall instead prove | s t | β ≲ δ t + C ′ η γ , where the added η γ handles the error terms and the exponent γ > 0 emerges from the proof.

Proposition 3 (quasi-static analysis; upper bound) . For all t such that 0 ≤ δ t -1 ≲ 1 / ( β ∨ 1) (for a sufficiently small implied constant), it holds that

where C, C ′ > 0 are constants which may depend on the problem parameters but not on η .

We first show that Theorem 2 now follows.

Proof of Theorem 2. As previously noted, the β = + ∞ case is handled by the arguments of the previous section, so we focus on β < + ∞ . From Lemma 7, we either have y 2 t → 2 /η and | x t |→ 0 , in which case we are done, or there is an iteration t such that y 2 t < 2 /η ≤ y 2 t -1 . From Proposition 3, since δ t -1 ≥ 0 and δ t ≤ 0 , it follows that | s t | β ≲ η 1 / ( β -1) . The theorem now follows, either from Proposition 1 or from the refined Proposition 2.

We now prove Proposition 3. In the proof, we use asymptotic notation O ( · ) , ≲ , etc. in order to hide constants that depend on ℓ (including β ), but not on δ t and η . However, the proof also involves choosing parameters C, C ′ > 0 , and we keep the dependence on these parameters explicit for clarity.

Proof of Proposition 3. The proof goes by induction; namely, if | s t | β ≤ C ( δ t + C ′ η γ ) and δ t ≥ 0 at some iteration t , we prove that the same holds one iteration later, where the constants C, C ′ > 0 as well as the exponent γ > 0 are chosen later in the proof.

For the base case, observe that the approximate conservation lemma (Lemma 6) gives | s t | ≲ 1 , and δ t ≳ 1 / ( β ∨ 1) at the beginning of the induction, so the bound is satisfied initially if we choose C sufficiently large enough.

Throughout, we also write ˆ δ t := δ t + C ′ η γ as a convenient shorthand. The strategy is to prove the following two statements:

If | s t | β = C t ˆ δ t for some C t > C 2 , then | s t +1 | β ≤ C t +1 ˆ δ t +1 for some C t +1 ≤ C t .

Proof of 1. The dynamics for x give

By Assumption (A2 + ) and | s t | ≲ 1 , and hence

If | s t | β = C t ˆ δ t for some C t ≤ C 2 , then | s t +1 | β ≤ C ˆ δ t +1 .

for large C . Also, ℓ ′′ (0) = 1 and a similar argument as in the proof of Theorem 1 yields the reverse inequality ηy 2 t r t ≳ 1 . We conclude that and hence

Since we need a bound in terms of ˆ δ t +1 , we use the dynamics of y ,

Substituting this in,

Let us show that

From (20), we have ˆ δ t +1 ≥ ˆ δ t -2 η 2 ( C ˆ δ t ) 2 /β , so we want to prove that η 2 ( C ˆ δ t ) 2 /β ≤ ˆ δ t / 8 . If β ≤ 2 this is obvious by taking η small, and if β > 2 then this is equivalent to C 2 /β η 2 ≲ ˆ δ 1 -2 /β t . It suffices to have C 2 /β η 2 ≲ ( C ′ ) 1 -2 /β η γ (1 -2 /β ) , which is achieved by taking C ′ large relative to C and by taking γ ≤ 2 / (1 -2 /β ) ; this constraint on γ will be satisfied by our eventual choice of γ = β/ ( β -1) .

Returning to (21), in order to finish the proof and in light of (22), we want to show that C 2 ˆ δ 2 t ≳ C 1+2 /β η 2 ˆ δ 2 /β t . Rearranging, it suffices to have ˆ δ 2 -2 /β t ≳ C 2 /β -1 η 2 , or ˆ δ 1 -1 /β t ≳ C 1 /β -1 / 2 η . Since by definition ˆ δ t ≥ C ′ η γ , by choosing C ′ large it suffices to have γ ≤ 1 / (1 -1 /β ) = β/ ( β -1) , which leads to our choice of γ .

Proof of 2. Using the simple bound ηy 2 t r t ≤ 2 + δ t , we have

where we used (22). If exp( βδ t ) ≤ 4 / 3 , which holds if δ t ≲ 1 /β , then from C t ≤ C/ 2 we obtain | s t +1 | β ≤ C ˆ δ t +1 as desired.

By following the same proof outline but reversing the inequalities, we can also show a corresponding lower bound on | s t | β , as long as δ t ≳ η β/ ( β -1) . Although this is not needed to establish Theorem 2, it is of interest in its own right, as it shows (together with Proposition 3) that the iterates of GD do in fact track the quasi-static trajectory.

Proposition 4 (quasi-static analysis; lower bound) . Suppose additionally that (A3) holds and that β < + ∞ . Also, suppose that at some iteration t 0 , we have δ t 0 ≲ 1 and that

holds at iteration t = t 0 , where c is a sufficiently small constant (depending on the problem parameters but not on η ). Then, (23) also holds for all iterations t ≥ t 0 such that δ t ≳ η β/ ( β -1) .

Proof. The proof mirrors that of Proposition 3. Let δ t ≳ η β/ ( β -1) for a sufficiently large implied constant. We prove the following two statements:

Throughout the proof, due to Proposition 3, we also have | s t | ≲ δ 1 /β t .

Proof of 1. The dynamics for x give

By Assumption (A3),

If c is sufficiently small, then

Therefore, we obtain

On the other hand, and hence

To conclude, we must prove that η 2 δ 2 /β t ≲ δ t , but since δ t ≳ η β/ ( β -1) (with sufficiently large implied constant), then this holds, as was checked in the proof of Proposition 3.

Proof of 2. Using Assumption (A3),

Therefore, and

Together with the dynamics for x and (24),

Since c t ≥ 2 c , if δ t and η are sufficiently small it implies | s t +1 | ≥ c δ 1 /β t +1 .

Convergence rate estimates. Our analysis also provides estimates for the convergence rate of GD in both regimes. Namely, in the gradient flow regime, we show that GD converges in O (1 /η ) iterations, whereas in the EoS regime, GD typically spends Ω(1 /η ( β/ ( β -1)) ∨ 2 ) iterations ( Ω(log(1 /η ) /η 2 ) iterations when β = 2 ) in the bouncing phase (Figure 11). Hence, the existence of the bouncing phase dramatically slows down the convergence of GD.

Remark 2 . Suppose that at iteration t 0 , we have δ t 0 ≍ 1 . Then, the assumption of Proposition 4 is that | s t 0 | ≳ 1 . If this is not satisfied, i.e., | s t 0 |≪ 1 , then the first claim in the proof of Proposition 4 shows that | s t 0 +1 | ≥ (1 + Ω( δ t )) | s t 0 | = (1 + Ω(1)) | s t 0 | . Therefore, after t ′ = O (log(1 / | s t 0 | )) iterations, we obtain | s t 0 + t ′ | ≳ 1 and then Proposition 4 applies thereafter.

Remark 3 . From the quasi-static analysis, we can also derive bounds on the length of the bouncing phase. Namely, suppose that t 0 is such that δ t 0 ≍ 1 and for all t ≥ t 0 , we have | s t | = δ 1 /β t . If δ t 0 is sufficiently small so that r t ≳ 1 for all t ≥ t 0 , then the equation for y yields

We declare the k -th phase to consist of iterations t such that 2 -k ≤ δ t ≤ 2 -( k -1) . During this phase, δ t +1 ≤ δ t -Θ( η 2 2 -2 k/β ) , so the number of iterations in phase k is ≍ 2 k (2 /β -1) /η 2 . We sum over the phases until δ t ≍ η β/ ( β -1) , since after this point the quasi-static analysis fails and y 2 t crosses over 2 /η shortly afterwards. This yields

The time spent in the bouncing phase increases dramatically as β ↘ 1 .

D Deferred derivations of mean model

In this section, we provide the details for the derivations of the mean model in Subsection 3.1. Recall

where x = λy e j + ξ . We first approximate

In other words, we can ignore the contribution of the signal λy e j . This approximation holds because (i) initially, the bias b is not yet negative enough to threshold out the noise, and hence the summation ∑ d i =1 ReLU( ± ξ [ i ] + b ) is of size O ( d ) , and (ii) the difference between the left- and right-hand sides above is simply ReLU( ± λy ± ξ [ j ] + b ) -ReLU( ± ξ [ j ] + b ) , which is of size O (1) and hence negligible compared to the full summation.

Next, letting g ( b ) := E z ∼N (0 , 1) ReLU( z + b ) be the 'smoothed' ReLU (see Figure 7), concentration of measure implies the following two facts:

∑ ✶ ✶ Indeed, the summations above are sums of d i.i.d. non-negative random variables, and hence its mean is Ω( d ) (as long as b ≥ -O (1) ) and its standard deviation is O ( √ d ) . Now, using these approximations, one can rewrite the GD dynamics on the population loss E [ ℓ logi ( yf ( x ; a -, a + , b ))] .

Using these approximations, the output of the ReLU network (2) can be written as

which in turn leads to an approximation of the GD dynamics on the population loss ( a -, a + , b ) ↦→ E [ ℓ logi ( yf ( x ; a -, a + , b ))] :

where ℓ sym ( s ) := 1 2 (log(1 + exp( -s )) + log(1 + exp(+ s ))) is the symmetrized logistic loss. Hence we arrive at the following dynamics on a ± and b that we call the mean model:

Now, we can write the above dynamics more compactly in terms of the parameter A t := d ( a -t + a + t ) .

E Proofs for the mean model

In this section, we prove the main theorems for the mean model. We first recall the mean model for the reader's convenience.

E.1 Deferred proofs

In this section, we collect together deferred proofs from Subsection 3.2.

Proof of Lemma 2. By definition, g ( b ) = ∫ ∞ -b ( ξ + b ) φ ( ξ ) d ξ = ∫ ∞ -b ξ φ ( ξ ) d ξ + b Φ ( b ) . Recalling φ ′ ( ξ ) = -ξ φ ( ξ ) , the first term equals φ ( b ) . Moreover, g ′ ( b ) = -b φ ( b )+ Φ ( b )+ b φ ( b ) = Φ ( b ) .

Proof of Lemma 3. Note that ∂ t ( 1 2 A 2 ) = A ˙ A = -2 d 2 ℓ ′ ( Ag ( b )) Ag ( b ) and also that ∂ t κ ( b ) = -ℓ ′ ( Ag ( b )) κ ′ ( b ) Ag ′ ( b ) = -ℓ ′ ( Ag ( b )) Ag ( b ) since κ ′ = g/g ′ . Hence, ∂ t ( 1 2 A 2 -2 d 2 κ ( b ) ) = 0 and the proof is completed.

E.2 Gradient flow regime

Proof of Theorem 3. The following proof is analogous to the proof of Theorem 1. We first list several facts we use in the proof:

(i) | g ′ ( b ) | = | Φ ( b ) | ≤ 1 for all b ∈ R .

(ii) ℓ ′ ( s ) = 1 exp( s ) -1 . Hence, ℓ ′ ( s ) 1 for all s R , and we have

(iii) ℓ ′′ (0) = 1 / 4 .

Throughout the proof, we assume that A 0 > 0 without loss of generality. We prove by induction the following claim: for t ≥ 0 and it holds that | A t | ≤ A 0 exp( -γt ) . This clearly holds at initialization.

Suppose that the claim holds up to iteration t . Using the bounds on | g ′ | and | ℓ ′ | , it follows that

In particular, b t ≥ -1 and g ( b t ) > 0 . 08 , since η ≤ γ A 0 . Also, the bound shows that if the claim holds for all t , then we obtain the desired conclusion.

It remains to establish the inductive claim; assume that it holds up to iteration t . For the dynamics of A , by symmetry we may suppose that A t > 0 . From ℓ ′ ( A t g ( b t )) ≤ A t g ( b t ) / 4 and g ( b t ) ≤ g (0) = 1 √ 2 π , it follows that

This shows that A t +1 ≥ -(1 -γ ) A t . Next, we show that A t +1 ≤ (1 -γ ) A t . First, if A t g ( b t ) ≤ 2 , since we have g ( b t ) > 0 . 08 . Next, if A t g ( b t ) ≥ 2 , then

This shows that | A t +1 | ≤ (1 -γ ) | A t | for the case A t > 0 . A similar conclusion is obtained for the case A t < 0 . The induction is complete.

E.3 EoS regime

Proof of Theorem 4. The following proof is analogous to the proof of Lemma 7. Assume throughout that A t = 0 for all t . Recall the dynamics for b :

̸

Since ℓ ′ ( s ) /s → 1 / 4 as s → 0 , and ℓ ′ is increasing, this equation implies that if lim inf t →∞ | A t | > 0 then b t must keep decreasing until 1 2 d 2 g ( b t ) 2 < 2 /η .

Suppose for the sake of contradiction that there exists ε > 0 with 1 2 d 2 g ( b t ) 2 > (2 + ε ) /η , for all t . Let ε ′ > 0 be such that 1 -(2 + ε ) (1 -ε ′ ) < -1 , i.e., ε ′ < ε 2+ ε . Then, there exists δ > 0 such | A t | ≤ δ implies ℓ ′ ( A t g ( b t )) / ( A t g ( b t )) > 1 4 (1 -ε ′ ) , hence

The above means that | A t | increases until it exceeds δ , i.e., lim inf t →∞ | A t | ≥ δ . This is our desired contradiction and it implies that lim t →∞ 1 2 d 2 g ( b t ) 2 ≤ 2 /η .

Remark 4 . A straightforward calculation yields that when ( a -⋆ , a + ⋆ , b ⋆ ) is a global minimizer (i.e., a -⋆ + a + ⋆ = 0 ), then λ max ( ∇ 2 f ( a -⋆ , a + ⋆ , b ⋆ ) ) = 1 2 d 2 g ( b ⋆ ) 2 . The mean model initialized at ( A 0 , 0) approximately reaches (0 , 0) whose sharpness is d 2 g (0) 2 / 2 = d 2 / 4 π . Hence, the bias learning regime 2 /η < d 2 / (4 π ) precisely corresponds to the EoS regime, 2 /η < λ max ( ∇ 2 f ( a -⋆ , a + ⋆ , b ⋆ ) ) .

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper is a study on the dynamics of neural network training, particularly focusing on the edge of stability. The authors explore the behavior of gradient descent in the context of a simple sparse coding model. They find that the dynamics of training exhibit a phase transition at the edge of stability, where the learning dynamics switch from a regime of stable convergence to one of chaotic oscillations.\n\nThe authors also demonstrate that the edge of stability is a critical point where the learned representations transition from being dense to sparse. This transition is shown to be sharp and is characterized by a power-law behavior. The authors further show that these findings are indicative of behaviors observed in practical neural network training settings, including a two-layer ReLU network trained on the full sparse coding model with an unknown basis, as well as a deep neural network trained on CIFAR-10.\n\nSTRENGTHS:\nThe paper addresses a complex problem in the field of neural network learning, focusing on the edge of stability. The authors' exploration of the phase transition in learning dynamics at the edge of stability, and the transition from dense to sparse representations, appears to be a novel contribution to the field.\n\nThe quality of the paper is good. The authors have conducted a thorough investigation into the dynamics of neural network training. Their findings are well-supported by their data and analysis, indicating a high level of methodological soundness.\n\nThe presentation of the paper is also good. The authors have clearly communicated their research, making it easy for readers to follow their methodology, findings, and conclusions. The paper appears to be well-structured and well-written.\n\nWEAKNESSES:\nThe authors assert that a large learning rate is necessary to learn the bias. This assertion holds true in practical scenarios, such as training ResNet-18 on binary CIFAR-10 using various learning rates. However, I am curious about the potential effects of freezing the bias during training. In an experiment I conducted, I found that freezing the bias did not significantly impact the results. I would appreciate it if the authors could investigate this aspect further and discuss its implications.\n\nThis work investigates a simplistic neural network model that consists of only three parameters. While this provides a clear and controlled environment for study, it represents a significantly pared-down model when compared to the complex, over-parameterized models commonly used in deep learning. The insights gained from this study may not fully capture the intricacies and nuances of real-world deep learning applications.\n\nQUESTIONS:\nIn lines 31-32, the authors posit that an effective network maintains a negative bias $b$ to filter out noise, and ensures that $a^-$ is less than 0 and $a^+$ is greater than 0 to output correct labels. I would appreciate if the authors could elaborate on this assertion. From my perspective, if the noise is of opposite sign to the label (for instance, $y_i = 1$ when the output of $x[i]$ is less than 0), assigning a negative value to the bias $b$ may not be beneficial, as both ReLU terms would return 0 in this scenario. Could the authors clarify this point?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

comments: SUMMARY:
The authors analyze the dynamics of pairs of ReLu neurons with an input bias, no weight matrix, and a readout layer. They show that it takes large learning rate for non-zero biases to be learned, and at these large learning rates there are formal guarantees of EOS behavior as well. They use a model with random weights and shared biases to show that the EOS induces a phase transition in the behavior of the biases.

STRENGTHS:
The model is straightforward, and the analysis of the EOS seems correct. The phase transition in the bias is well established by this work and is an interesting phenomenon.

WEAKNESSES:
Though this work claims to study a realistic model, the one-neuron model is still very much a toy model. Many other toy models have already been studied in recent works, and it is unclear what the analysis of this model in particular adds. In addition, many networks are trained without the bias term, so there may be limited utility for theory which focuses mainly on that term.

Additional suggested references:
[EOS in a quadratic regression model with formal convergence guarantees](https://arxiv.org/abs/2210.04860)
[EOS and weight magnitude in quadratic expansion of ReLu network](https://arxiv.org/abs/2205.11787)

QUESTIONS:
Is the bias learning phase transition a good way to describe single hidden layer ReLu models with multiple neurons?

LIMITATIONS:
Yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors studied the problem of edge of stability (EoS) phenomena (training with large learning rate) in simple settings. Specifically, the authors first studied the problem of minimizing $\\ell((xy)$ ($\\ell$ is loss function) and showed that under certain condition of loss function, gradient descent are either in the gradient flow regime or EoS regime depending on the initialization, which means different stepsize could lead to final solution with different maximum eigenvalue of Hessian of loss. Then, to study the problem of learning threshold neurons, they instead considered a simplified problem which they called mean model. They proved that under this mean model, larger stepsize could decrease the bias and therefore learn the neuron, while bias remains small under small stepsize. Empirical experiments showed that the training dynamics of this mean model and the actual training is similar.\n\nSTRENGTHS:\n1.\tThe paper is clearly written and easy-to-follow. The proof sketch and dynamics are given to help the readers to understand the proof easier. Experiments are provided to support the results.\n2.\tUnderstanding the effect of large learning rate in deep learning training (specifically EoS phenomenon) is an important problem and may give more insight into the practice of deep learning training.\n3.\tThis paper provided both theoretical and empirical results on the simple model that learning threshold neurons. The theoretical results, though do not directly answer the question of learning threshold neurons, give some possible insights and seem to be interesting.\n\nWEAKNESSES:\n1.\tThe current paper focuses on the simple setting of minimizing $\\ell(xy)$ as well as a mean model introduced to approximate learning threshold neuron. It would be interesting to extend the analysis to broader settings (e.g., real dynamics of learning threshold neuron and sparse coding model).\n2.\tThe results in the EoS regime has an additional assumption on the dynamic, though it is understandable from the technical perspective.\n\nQUESTIONS:\n1.\tI was wondering what the convergence time for are both gradient flow regime and EoS regime.\n\nLIMITATIONS:\nThe limitations are clearly discussed in the paper. This is a theoretical work and therefore does not seem to have negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies two simple problems (single-neuron linear network and mean ReLU model for sparse coding with unknown basis) to investigate learning of threshold neurons. In particular, the authors discover a threshold for the learning rate below which threshold neurons are not learned and where NN training follows the gradient flow. For larger learning rates, training trajectories are oscillating (edge of stability) and threshold neurons are learned, which is essential for solving the sparse coding problem.\n\nSTRENGTHS:\nThe paper is extremely well written and provides a gentle introduction to concepts such as sparse coding, edge of stability, and threshold neurons, which makes it accessible to readers not familiar with these terms. The results in Sections 2 and 3 seem rigorous and are well presented. Particularly Section 1.2 is an excellent summary of the results that follow and helps to put them into context. The theoretical results seem to agree with the experimental evidence. The fact that the considered examples are simple is not a drawback, as this is expected for works that, for the first time, study a certain phenomenon from the theoretical perspective. Nevertheless, I appreciate the authors effort to back their claims and show the generality of the phenomenon via the experiments on CIFAR-10 in Fig. 3.\n\nWEAKNESSES:\nNone, actually -- but see the questions below.\n\nEDIT: Lowered score since practical implications of results are not fully clear (cf. frozen bias).\n\nIf I have to mention something, then it is the fact that Fig. 1 appears slightly out of context and could be shifted to the second page. Furthermore, there are a few typos that can be removed (139 convergenec, 262 in turns out, 268 from which can more accurately, 289 wiht).\n\nQUESTIONS:\nSome questions came up when reading the paper. Taking them into account could further improve the readability of the manuscript:\n- In Fig. 1, why is the bias initially at -0.2? Was this a randomly chosen initial bias?\n- What is the definition of \"sharpness\", as discussed in Sec. 2.2?\n- It would be good to place the mathematical expression for the gradient flow in Sec. 2, maybe close to the GD iterates in line 177.\n- It is not clear, from the main part of the paper, why the step size for the $A$-dynamic is multiplied by $2d^2$ (e.g., in (5)). It would be good to explain that briefly in the main part of the paper.\n\nLIMITATIONS:\nThe authors have listed several interesting future avenues of research. The limitations are evident from the simple setup. There are no negative societal impacts foreseen.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-12"}