Learning threshold neurons via edge of stability
Authors: Kwangjun Ahn
Sebastien Bubeck
Sinho Chewi
Yin Tat Lee
Felipe Suarez
Yi Zhang
Venue: NeurIPS 2023 poster (2023)
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"
Kwangjun Ahn MIT EECS Cambridge, MA kjahn@mit.. edu
Sébastien Bubeck Microsoft Research Redmond, WA btbubeekk@iicoosof.ccom
Sinho Chewi
Institute Advanced Study
Princeton, NJ schewi Sachww@ias.co
Yin Tat Lee Microsoft Research Redmond, WA yintatQuw.edi edu
Felipe Suarez Carnegie Mellon University Pittsburgh, PA felipeac@mit.
Yi Zhang Mierosoft Research Redmond, WA zhayi saayyiimmcorsooft.comm
Abstract
Existing analyses of neural network training often operate under the unrealistic sumption of an extremely small learning rate. This lies in stark contrast to practical wisdom and empirical studies, such the work of J. Cohen et al. (ICLR 2021), which exhibit startling new phenomena (the "edge stability" or "unstable conver- gence") 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 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 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.
-12
Eroal taxt accurncy (K)
0.803
01011
1.35
-12
-12
10812
1002 1.01
14000
0.00
0002
Figure 1: Large step sizes are necessary to learn the "threshold neuron" of a ReLU network (2) for a simple binary classification task (1). We choose d = 200, 300, À = 3, and run gradient descent with the logistic loss. The weights initialized iteration number thcoretical prediction of the phase transition phenomenon (precise threshold at r = 87 /d)).
a A² 1/(2d) and 0. For each learning 7/, total elapsed (iteration x 17) 10. The vertical dashed lines indicate our
the
Introduction
How much do we understand about the training dynamics of neural networks? We begin with simple and canonical learning task which indicates that the answer is still "far too little".
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Motivating example: Consider binary classification task of labeled pairs (z() y E R x {+1}
where each covariate σ consists of 1-sparse vector (in an unknown basis) corrupted by additive Gaussian noise, and the label ( 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
x (4) Xy Eg(() g(() TR
where yy0) {+1} is random label, j(} E [d] is random index, { is Gaussian noise, and is the unknown signal strength. fact, (1) 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-Zhe 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 under- standing of surprisingly rich dynamics which lies out of reach of existing theory. To illustrate this point, consider an extreme simplification in which the basis ei...., Ed is known in advance, for which it natural parametrize two-layer ReLU network as
[(x; a a .b) b) = Σ ReLU (-a x|1| + b) + at Σ ReLUU+=[[]] +
The parametrization (2) respects the latent data structure (1) well: a good network has negative bias b to threshold out the noise, and has a < 0 and at >> 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 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.5
26x
e 26x
training loss
1.0
bias
25x
-03
-01
-31
26
20
75
100
20
50
100
25
80
75
100
iteration
heration
iteration
pheeppemedphenmmeeet phenomena: non-monotonic loss and wild oscillations of weights. We choose the same setting Figure 1. With small learning rate (1) 2.5 10-5), the bias does not decrease noticeably, and the same is true even when increase the learning rate by ten times (1) 2.5 10 10-4). When we increase the learning rate by another ten times (1)
Figure 2: Large learning rates lead to unexpected
10 10-"),, we finally see noticeable decrease the bias, but with this we observe unexpected behavior: the loss decreases non monotonically and the sum of second layer weights d (a + a oscillates wildly.
We train the parameters a a -,,b using gradient descent with step size 1/ > 0 on the logistic loss E ()) ( at, where (iogi(2) = log(1 + exp(-2)), and we report the results in Figures! 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 shows that the bias decreases and the test accuracy increases as we increase ry. note that we plot the results after a fixed time (iteration X 11), so the observed results are not simply because larger learning rates track the continuous time gradient flow for longer time.
Large learning rates lead to unexpected phenomena: non-monotonic loss and wild oscillations of a + at. 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 phase transition between small and large learning rates. In Figure 1, we zoom in on learning rates around 17 se 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 ReI.U 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.50
[1w0:33
0.0
-0.5
805
0.75
1est
0.70
bia
-1.0
0.00
-1.3
F-003
-2.0
F-607
0.50
200
5:00
800
1.00
1200
1400
620
1000
1230
1400
stepXn
step x
0.0000
0.0025
med ain bllas
test aCC
0.0075
70
00.0000
60
-0.1150
P-43
-0.0175
2-3.5
-0.0200
Stepx.
stepxi
Figure 3: (Top) Results for training an over-parametrized two-layer neural network f(x; a, W,b) Ems a; ReLU(w; (w, æ + b) with m >> for the full sparse coding model (1); in this setting, the basis vectors unknown, and the neural network learn them through additional parameters W = )= Also, we use m different weights a = (a=) 1 second layer. (Bottom) Full-hatch gradient descent dynamics of 18 on (binary) CIFAR-10 with various learning rates. Details are deferred to Appendix B.
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 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 detailed theoretical study.
Main scope of this work
In this work, we do not aim to understand the sparse coding problem (1) its full generality. Instead, 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 al., 2022; Lyu et al., 2022: Ma 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.
3

---

The edge of stability (EoS) phenomenon is set of distinctive behaviors observed recently by Cohen al. (2021) when training neural networks with gradient descent (GD). Here we briefly summarize the salient features of the FoS and defer a discussion of prior work to Subsection 1.3. Recall that we use GD to optimize an [-smooth loss function with step size 7) then the well-known descent lemma from convex optimization ensures monotonic decrease in the loss so long L < 2/1-. In contrast, when L > 2/1,, 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' with constant step size 1) > O, 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/n ("EoS").
A surprising message of the present work is that the answer to our main question is intimately related the EoS. Indeed, Figure shows that the GD iterates of our motivating example exhibit the EoS during the initial phase of training when the bias decreases rapidly.
2,000
614
2000
-42
2000
64-
103
2.00
2011
300
130 3.0
Figure 4: Understanding our main question is surprisingly related the EoS. Under the same setting as Figure 1, we report the largest eigenvalue of the Hessian ("sharpness") " and observe that GD iterates lie in the EoS during the initial phase of training when there is a fast drop in the bias.
Consequently, we first out to thoroughly understand the workings of the EoS phenomena through simple example. Specifically, we consider single-neure neuron linear neural network in dimension 1, corresponding to the loss
TR (x. v) - E(ry), where is convex, even, and Lipschitz.
Although toy models have appeared works on the EoS (see Subsection 1 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. Al- though the EoS has been studied in various settings (see Subsection 1.3 for a discussion), these works either do rigorously establish the EoS phenomenon, they op- erate under complex settings with opaque assumptions. Here, we study simple two-dimensional loss function, (x, v) + ((ry), where E is convex, even, and Lipschitz. Some examples include? E(s) 3 log(1 exp(-s))
GF
log(1 + exp(+s)) and {(s) = vi + s². Surprisingly, GD this loss already exhibits rich behavior (Figure 5).
Illustration of two different regimes (the "gradient flow" regime and the "EoS" regime) of the GD dynamics.
Figure 5:
En route to this result, we rigorously establish the quasi- static dynamics formulated in Ma et al. (2022).
'The phenomenon in Cohen et al. (2021) is most clearly observed for tanh activations, although the appendix of Cohen et al. (2021) contains thorough experimental results for various neural network architectures. "Suppose that we have a single-layer linear neural network fæ; a, b) = abz, and that the data is drawn according to x = 1. -ll Then, the population loss under the logistic loss is (a, b) + Laym(ab)
with leym(s (s) log(1 + exp(-s)) } log(1 + exp(+s)).

---

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, 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 transition from the "gradient flow" regime to the "EoS regime", depending on the step size 17 and the initialization. Moreover, in the EoS regime, rigorously establish asymptotics for the limiting sharpness which depend the higher- behavior of f. 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/n O(1)
In fact, this connection can 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 Ag = d (aj + a;") and b, due to the ReLU network well- well-approximated by the "rescaled" GD dynamics on the loss (A, 6) + laym x g(b)). where the step size for the A-dynamics is multiplied by 2dh, g(b) == 1 ReLU(z + 6) is the "smoothed" ReLU, and laym 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-neum neuron example (z, y) + (ary), and hence we can leverage the techniques developed for the latter in order to study the former.
An* +**)
Figure 6: Illustration of GD dynamies on the ReLU network (2). The sum of weights (a + a²) oscillates while b decreases.
Our main result for the mean model precisely explains the phase transition in Figure 1. For any ō >> O,
if r S (8- /dd/ then the mean model fails learn threshold neurons: the limiting bias satisfies 1
= (s(L/d²).
if 2 (8 + /dd, then the mean model enters the EoS and learns threshold neurons: the limiting bias satisfies b I (1).
1.3 Related work
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 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 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 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 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 (I, v) + (ry 1)² and single-neuron suudent-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 wwo-dimensional loss (x, y) >> (x²r² 1)2: to our knowledge, their work is the first to asymptotically and rigorously show that the limiting sharpness of GD is 2/n in simple setting,
5

---

at least when initialized locally. In comparison, in Section 2, we perform global analysis of the limiting sharpness of GD for (z, + ((zy) for a class of convex, even, and Lipschitz losses L, and in doing so we clearly delineate the "gradient flow regime" from the "EoS regime".
Effecc of learning rate 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 al., 2021).
Li et al. (2019) demonstrate for a synthethic data distribution and a two-layer ReLU network model that choosing 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 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. notable exception is the work of Wang et al. (2022a), which studies the impact of learning rates greater than 2/ 2/smoothness for a matrix factorization problem. Also, the recent work of Andriushchenko 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 (z. v) + f(I = f(x x v).
2.1 Basic properties and assumptions
Basic properties. If is minimized at O, then the global minimizers of f are the x-- and y-axes. The GD iterates Ft, yt. for step size 17 > 0 and iteration 2 0 can be written as
T++1 xt l' (xx33)) Vrr
Vt++1 = 1/: 1) E' (x.y))) Zz.
Assumptions. From here onward, we assume v) < 1 and the following conditions on f: R - R.
(A1) is convex, even, 1-Lipschitz, and of class C2 near the origin with ("( = 1.
(A2) There exist constants B > and > 0 with the following property: for all + 0.
E' (s)/s < - c|s|@ 1|1s] S c].
We allow +00, in which case we simply require that (()) < 1 for all s + 0.
Assumption (A2) imposes decay s e' (x)/ / locally away from the origin in order to obtain more fine-grat nee-grained results on the limiting sharpness in Thcorem 2. As we show in Lemma below, when E is smooth and has strictly negative fourth derivative at the origin, then Assumption (A2) holds with | = 2. See Example 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 (Fo. yo). there are two different regimes for the GD dynamics depending on the step size 17. Namely, there exists 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
yo To 0, yo | for all 2 1, and GD converges to (0, yoc) for you 0.
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 va.

---

Initially, in both regimes, the GD dynamics tracks the continuous-ting gradient flow. Our first observation is that the gradient flow admits a conserved quantity, thereby allowing us predict the dynamics in this initial phase.
Lemma (conserved quantity). Along the gradient flow for S, the quantity 3² - x²2 is conserved.
Proof. Differentiating x² with respect to gives 2y, (- (-((2)))) x1) 2x, ( (-((x136) Bo)
Lemma implies that the gradient flow converges to (0, Moc (0. v yū xg)). For GD with step size 1 > O, the quantity y² x²² 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 (To, yo)
also reaches the y-axis approximately at the point (Ito lt.) R (0, V - -2))
At this point, GD either approximately converges to the gradient flow solution (0, V Vo 2 or diverges away from it, depending on whether or not vi, > 2/n. To see this, for |oto Sto < 1, we can
Taylor expand l' near zero to obtain the approximate dynamics for (recalling (0()) = 1),
1to+1 R Tto 1f tho (1 mhho) Tig-
(4)
From (4), we deduce the following conclusions.
(i) If vi < 2/n, then |1 n < 1. Since y is decreasing, it implies that [1 - --nll < 1 for all I 12 lo. and so |ız]] converges to zero exponentially fast.
(ii) On the other hand, if vž, > 2/n, then | nyf > 1, i.e., the magnitude of Tioo increases in the next iteration, and hence GD cannot stabilize. In fact, in the approximate dynamics, Ita+1 has the opposite sign as Itg i.e., 2tg jumps across the y-axis. One can show that the "bouncing" of the x variable continues until has decreased past 2/1,, at which point we are in the previous case and GD approximately converges to (0, 2/n).
This reasoning, combined with the expression for the Hessian of f. shows that
sharpness(1) (0, yxx) == Amax ((0, = min v zg. 2/n} min (gradient flow sharpness, EoS prediction).
Accordingly, we refer to the case vo x < 2/1 as the gradient flow regime, and the case 36 x > 2/ 2/1 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 ahove 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 (xel >ا 0 exponentially fast and hence the dynamics stabilizes "quickly" once si 2/1, we track precisely how long this convergence takes so that we can bound the gap between the limiting sharpness and the prediction 2/m.
2.3 Results
Gradient flow regime. Our first rigorous result is that when y0 - x0 = (2-61/n for some constant
8 E (0, 2), then the limiting sharpness of GD with step size 17 is vo 2g 0(1) (2-6%/9 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, it covers initializations in which ó can scale mildly with 17. The precise statement is as follows.
Theorem (gradient flow regime; see Subsection C.2). Suppose we run GD with step size 1) >> on the objective f. where f(x, y) == E(zy), and satisfies Assumptions (Al) and (A2). Let (ă. D) E R² satisfy ÿ > ã= >> 0 with y²² - x² = 1. Suppose we initialize GD at (æo. Mo) = ( (2 (i, D) where
E (0,2) and 17 S 81/2 A (2 6). Then, GD converges to (0, yoo) satisfying
2 0(2 8) 0
min 83) S Amax ( (2 <
1
2 8
0
17
")
2 5
where the implied constants depend on ā, ÿ. and e, but not on 8, 17.
7

---

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² x² along GD in order show that GD lands near the y-axis with via ~ (2-8)/n. In the second stage, we use the assumptions on l in order to control the rate of convergence of [ayl to O, which is subsequently used to control the final deviation of v from (2 (2-6)/n-
EoS regime. Our next result states that when vā = 2/1,, then the limiting sharpness of GD is close to the EoS prediction of 2/ 17, up to an error term which depends on the exponent 3 in (A2).
Theorem 2 (EoS; see Subsection C.4). Suppose we run GD on with step size 1) > 0, where fz, v) == E(xy), and E satisfies (Al) and (A2). Let (ž. ū) E IR2 satisfy ÿ > I > 0 with y² - ====1 1. Suppose we initialize GD at (ro, yo) == v2+0%/ (ā., ū). where 6 > 0 is constant. Also, assume that for all 2 such that of 2/4, we have 24 # O. Then, GD converges to Uzo) satisfying
2/n O(n 1(8-1)) I Amax 1 yso)) S 2/m
where the implied constants depend on ē, y. 6 A 1, and E, but not on 17-
Remarks on the assumptions. The initialization in our results is such that both yo and yo Ty are on the same scale, i.e., yo. Yo To = ((1/vT) This rules out extreme initializations such as yo ~ To- which are problematic because they lie too close to the invariant line y 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² = 22/1. This is necessary because if o = 0 for some iteration t, then (æy', ye) = (zz, Ve) for all > t. and hence the limiting sharpness may not be close to 2/n. This assumption holds generically, c.g., if we perturb each iterate GD with a vanishing amount of noise from a continuous distribution, and we conjecture that for any 1) > 0, the assumption holds for all but measure zero set of initializations.
When = +OC, which is the case for the Huber loss in Example 1, the limiting sharpness is 2/n - O()). When / = 2. which is the case for the logistic and square root losses in Example 1, the limiting sharpness is 2/1 + O()). Numerical experiments show that our error bound of (1z 1(8-1)) 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 bound the gap 2/1 - yy² the proof requires a control of the size [rev/]]. where is the first iteration such that V/ - crosses 2/1j. However, controlling the size of [*ttel is surprisingly delicate as it
requires 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 certain envelope (Figure 9).
As by- produuu of our analysis, obtain 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 Theorem 2, and Theorem 5 are displayed pictorially 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 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-approximmate by
GD dynamics (a= ,b) + (d + a+))(b)).
where fsym( (s) = (log(1 1 + exp(-s)) + log(1 + exp(+s))) and g(b) == Ez_,((0 ReLU(z + b) the 'smoothed' ReLU. The GD dynamics can be compactly written in terms of the parameter A, = d (a, + a.).
A A 2 E fy (Arg(be)) g(be),
b = be - n E (Arg(be)) Arg (be). (5)
8

---

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 as the "rescaled" GD dynamics on the objective (A,b) + (Ag(())) where the step size for the A-dynamics is multiplied by 2d2 Due to this resemblance, we can apply the techniques from Section 2.
Figure 7: The "smoothed" ReLU g())
0.00
031
2/33
-210
-51-
2/000
-12-
-2.15
-0.20
2/300
-15
-2.25
V- 2/100
-1.4
5.
-00
2041
590
Zeration
Figure 8: Under the same setting as Figure 1, we compare the mean model with the GD dynamics of the ReLU network. The mean model is plotted with black dashed line. Note that the mean model tracks the GD dynamics quite well during the initial phase of training.
3.2 Two different regimes for the mean model
Throughout the section, we use the shorthand E == faym: and focus on initializing wiht oj = O(1dd). a + at * 0, and by = O. This implies An = ((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 can be expressed in terms of the PDF ip and the CDF & of the standard Gaussian distribution as g(b) = (b) + b (b). In particular, g' = Φb
Note also that b, is monotonically decreasing. This is because E' (()) Agg' 2 0 since e' is an odd function and g(b), g' (b) > for any b E TR.
Following Subsection 2.2, we begin with the continuous-time dynamics of the mean model:
À = 2d l' (Ag(b)) 9(b),
b b= (Ag(()))) (b).
(6)
Lemma 3 (conserved quantity; see Subsection E.1). Let k: R 12 be defined as k(b) = So g/g.. Along the gradient flow (6), the quantity 1 A2 (bb) is conserved.
Based on Lemma 3, if we initialize the continuous-time dynamics (6) at (Ao, 0) and if Ae - 0, then the limiting value of the bias &CF satisfies x(6CF) = 4. AG. which implies that JCF = - );
indeed, this holds since n²' 20()) g(0)/g (0) > 0, so there exist constants co. e > such that cob x(b) < c/ for all -1 < b < 0. Since the mean model (5) tracks the continuous-time dynamics (6) until reaches the b-axis, the mean model initialized at (Ag, 0) also approximately
reaches (Ac,, by,)) z (0, )))) Re (0, 0) in high dimension d >> 1. In other words, the continuous- time dynamics 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 l around the origin yields the following approximate dynamics (here e" (0) = 1/4): Atot1 R A 2 At, g(br,)) = At, (1 - ud 9(be,))) We conclude that the condition which now dictates whether we have bouncing convergence is de g(ba) > 2/n.
nd?

---

(i) Gradient flow regime: If 2/n > / (4T) (since g(0)² 1/(2T] D) i.e. the step size 1 is below the threshold 8 then the final bias of the mean model bMM satisfies UM Re 1GF 2 O. In other words, the mean model fails to learn threshold neurons.
(ii) EOS regime: If 2/n < (4) i.e., the step size yj above the threshold 8T then 10r 1 (b < 2/m, ie., 6% < g (2/ / /nd²). For instance, if = d² then bMM < -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 (Ao, 0), with step size 11 {8-8)ı for some > O. Let == 200 min(ó, 8 - 6, 8-6
d²
[.o
Then, as long as rj S 7/ Aol the limiting bias BMM satisfies
0 bMM (n/n) Aol A (1/dP).
In other words, the mean model learn threshold neurons.
Theorem 4 (mean model, EoS regime; see Appendix E). Consider the mean model initialized at (Ag, 0), with step size n (8+5) for some 5 > O. Furthermore, assume that for all 2 such that dl gl,))² 2/1, we have A, # 0. Then, the limiting bias IMM satisfies
BMM S و (2/v(8+8))) S -15(1).
For instance, in threshold neurons.
10a de
then MM < 0.0887 In other words, the mean model successfully learns
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 the sparse coding model (1) through 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 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 under- standing 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., 20)3a,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, follow up work by Dai et al. (2023) studies the working mech- anisms of a popular modern optimization technique called sharpness-ave -aware minimization (Foret al., 2021) based on 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 ReL.U 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 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.
10

---

Acknowledgments
We thank Ronan Eldan, Suriya Gunasekar, Yuanzhi Li, Jonathan Need, and Adil Salim for initial discussions this project. KA was supported by the ONR grant (N00014-20-1-25) 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 (XXi:23305.70440 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, 2022.
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. NeurlPS 2023 (arXin: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 Anmual Symposium on Foundations of Computer Science (FOCS), pages 977-988. IEEE, 2022.
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: part I, context-free grammar. arXiv preprint Xiv:22300 /3673, 2023.
Zeyuan Aluen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterizzatoon 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- 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, 2015. PMLR.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-graine 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, 2022.
11

---

Jimmy Ba, Murat A. Erdogdu. Tajji Suzuki. Zhichao Wang, Denny Wu, and Greg Yang. High- dimensional asymptoties 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 Muchine Learning Research, pages 284-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 Leurning Reseurch, 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- paramarrmmteee models using optimal transport. In S. Bengio. H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
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é-Bus Buc. E. Fox, and R. Gamett, 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. NeurlPS 2023 (arXiv:2305. 15287), 2023.
Alex Damian, Eshaan Nichani, and Jason D. Lee. solfistabilization: the implicit bias of gradient descent at the edge of stability. In The Kleventh International Conference on Learning Representa- fions, 2023.
Simon S. Du Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes premttaraeter neural networks. In International Conference on Learning Representations, 2019.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-awart 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ēra Kürková, Yannis Manolopoulos, Barbara Hammer, Lazaros Iliadi, and Ilias Maglogiannis, editors, Artificial Neural Networks and Machine Learning ICANN 2018, pages 392-402, Cham, 2018 Springer International Publishing.
Stanislaw Jastrzebski, Maciej Szymezak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho, and Krzysztof Geras. The break-even point on optimization trajectories of deep neural networks. In International Conference on Learning Representations, 2020.
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.
12

---

Stanistaw 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 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 04883-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.022 022218. 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 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 nor- malization layers: sharpness reduction. S. Koyejo. S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 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 Leurning, 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.07336 0000000000000,.0000000 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 Staristics, volume 151 of Proceedings of Machine Learning Research, pages 9741-9757. PMLR, 3 2022.
Bruno A Olshausen and David Field. Sparse coding with an overcomplete basis set: strategy employed by v1? Vision Research, 37(23):3311-322 1997.
Bruno A Olshausen and David J Field. Sparse coding of sensory inputs. Current Opinion in Neurobiology, 14(4):4814477, 2004.
13

---

Samet Oymak and Mahdi Soltanolkotabi. Toward moderate overparameterrizaton:: global convergence guarantees for training shallow neural networks. IEEE Journal on Selected Areas in Information Theory, (10:84-105, 2020.
Grant M. Rotskoff and Eric Vanden- 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 phe- nomenon via bifurcation theory. NeurlPS 2023 (arXiv:2307.0422) 2023.
William Vinje and Jack L. Gallant. Sparse coding and decorrelation in primary visual cortex during natural vision. Science, 287(5456): 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 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. 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-822 82888,888888008 2018
Chen Xing, Devansh Arpit, Christos Tsirigotis, and Yoshua Bengio. A walk with SGD, arXiv preprint arXiv: 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.04001, 2022.
Xingyu Zhu, Zixuan Wang, Xiang Wang, Mo Zhou, and Rong Ge. Understanding edge-of-stabili training dynamics with minimalist example. In The Eleventh International Conference on Learning Representations, 2023.
14

---

Appendix
A Additional illustrations for Section 2
15
Further experimental results
16
B.1 Experiments for the full sparse coding model
17
B.2 Experiments on the CIFAR-10 dataset
17
Proofs for the single-neuron linear network
18
C.I Approximate conservation along GD
19
C.2 Gradient flow regime
20
C,3 EoS regime: proof outline
21
C.4 EoS regime: crossing the threshold and the convergence phase
C.5 EoS regime: quasi-static analysis
24
D Deferred derivations of mean model
28
Proofs for the mean model
29
E.1 Deferred proofs
29
E,2 Gradient flow regime
................................
E.3 EoS regime
30
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.
(ato, Mo)
(No, Uo)
V 2/1
20
2/4
Figure 9: Two regimes for GD. We run GD the square root loss with step size 1. The gradicnt flow regime is illustrated on the left for (xo. yo) (3, 4). GD (bluc) tracks the gradient flow (green) when ') < 2/( xn) Otherwise, illustrated on the right for (zo, yo) = (3, 6), GD is in the EoS regime and goes through gradient flow phase (blue), intermediate bouncing phase (orange) that tracks the quasi-static envelope (purple). and converging phase (red).
15

---

Next, we next present detailed illustrations of the odge-of-stabilie regime depending on the choice of step size. Compare this plot with our theoretical results characterized Theorem 2.
V
y
y
(2o- vo.
(20-30)
( 86]
(sto Mo)
2,00
77
9
7
sharpness
5
3
0.0
2.5
5.0
7.5
10.0
12.5
15.0
time (17 x iteration)
Figure 10: We plot the GD trajectory for E(s) = v1+ x and sharpness for step sizes 2/m = 9, 2/1² 7. 2/1s = 5, and 2/1a = 3. In the EoS regime, the final sharpness is close to 2/(step size).
Lastly, we present another detailed illustrations of Theorem 2 in terms of its dependence on 3.
Eap
800 1777
Tay Taurl (adlor))
800 of
-1
logt?")
50p X
10
Figure (Left) Log-log plot the sharpness gap as a function of 17, for E. in Example and 3 3. 2, 3, 10. (Right) Log-log plot of the iteration count for the bouncing region with yf [2, }] function of 1, for E, in Example and 3. 2, 4. dashed lines show the predicted sharpness gap and iteration count with offset computed via linear regression of the data for < -
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
16

---

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 consider the sparse coding model 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,) together with different weights a (a:) for the second layer as follows:
m
(1x; a, W,b)
Σ a ReI.UU(W, x))
We show results for d 100, 2000. We generate aforementioned sparse coding model with X = 5. We use the He initialization, i.e., a Wo0, Im/m). (0, Ia/d), and b = 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.
20000 data points according to the
avg. 2nd layer weight
20
-10
15
-20
DVg
-40
1-0 00
-00
15
15 step
20
10
20
30
50
60
step
Figure 12: (Left) The average of the second layer weights of the ReLU network (7). Note that the average value oscillates similarly to our findings for the mean model. (Right) Oscillation of logit of ResNet18 model averaged over the (binary) CIFAR-10 training set. Since the dataset is binary, the logit simply scalar.
B.2 Experiments on the CIFAR-10 dataset
Next, provide the details for the bottom plot of Figure 3. We train ResNet-18 on binarized version of the CIFAR-10 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 measured over the entire training set. The median bias is measured at the last convolutional layer right before the pooling. For the optimizer, scheduler where learning rates shown in the plots are the initial values.
full-batch GD with momentum weight decay, plus cosine learning rate
Oscillation of expected output (logit) of the network. Bearing a striking resemblance 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 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 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.
17

---

avg. avg. activation
0.8
median blas
0.00
0.6
1-1.
-0.02
-0.03
0.2
100
150
200
250 300
350
400
50
100
150
200
250 300 350 400
step
step
avg. activation
median bias
0.000
1-11 1-2.05
0.6
Whyy Why
0.3
-0.015
100
150 200
250 300
350 400
50 100
150
250 300 step
400 200
350
step
Figure 13: SGD dynamics of ResNet-18 on (multiclass) CIFAR-IU 10 with various learning rates and batch sizes.
(Top) batch size 100; (Bottom) batch size 1000. The results are consistent across different batch sizes.
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 0, then the global minimizers of are the and y-axes. The gradient and Hessian of f are given by:
V ff v) (zyy
y
v's(x,)) = E" (ay)
y
00
+ l' (xy)
0 1
(8)
x
1 0
This results in GD iterates It, Vt. for step size 17 0 and iteration 2 0:
Fz4 T, 77 (xx3e) Ve.
Ve+1 yv n (x((t) zt.
Lemma 4 (invariant lines), Assume that E even, so that e' odd. Then, the lines y fr are invariant for gradient descent on f.
Proof. If yt +r,, then
Me+1 = Vt 7 (I((1)) zt = +xa F 17 (x+) xe.
Tz+1 It n (I+3))) Br = It El (अ+² Zt.
and hence V1+1 1 Note that the iterates (zz)120 are the iterates of GD with step size r the one-dimensional loss function x + 1 {xx2)
We focus instead on initializing away from these two lines. We now state our assumptions on f. We gather together some elementary properties of L.
Lemma (properties of f). Suppose that Assumption (Al) holds.
1. E is minimized the origin and (0) =.
2. Suppose that E four times continuously differentiable near the origin. If Assumption (A2) holds, then g(4) (0) O. Conversely, if ((4) (0) then Assumption (A2) holds for 3 = 2.
Proof. The first statement is straightforward. The second statement follows from Taylor expansion: for * near the origin,
(s) (0) + (d (0) 8 + ( - r) (" dr 1+ ( - (r)
(9)
18

---

Since e" is odd, then Assumption (A2) and (9) imply that ed is non con-positive on (0, for some
E > O, which in turn implies (0) I O. Conversely, if ((4) (0) < 0, then there exists > 0 such that (s) V -Es for 8 E (0, From (9), we see that (s) S 1 É ( - r) dr I 1 - es /2. and some c > 0.
By symmetry, we conclude that Assumption (A2) holds with
We give some simple examples of losses satisfying our assumptions.
Example 1. The examples below showcase several functions/ e that satisfy Assumptions (A1) and (A2) with different values of B.
Rescaled and symmetrized logistic loss. frsym (s) == 1 flogi (-2s) + 2 llogi (+2+).
Note (ssym( (s) = tanh(s), thus ('sym (s)/8 S 1 and C 1sym (s)/8= S 1 - 1 |s|² for |s| < 1
Square root loss. lagnt() == v1+ s2.
Note "sget (8) V1+=z thus l ( (s)/s S 1 and (ggee()) (5)/s S 1 1- |s|2 for |s| <
Huber loss. (Hbb(s) == d² {s E [-1, + ([s) - - 1{s [-1,
2
Note (igh (8) = 8 81 E [11 + sgn(s) (8) (s & [-1, 1]]. thus Hus (5)/8 S 1, i.e. we have Assumption (A2) with 8 = to0.
Higher-order. For > let c3 = B+I ($41) 9 and rB = S+1 We define Lg implicitly via its
derivative
(5) = (1 - cи 8 1 < < + sgn( 8) 1 2 r3]
By definition, fg($)/s < and (s)s 1 1 - ct | where ce = Ca A rg.
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 -2² for the GD dynamics.
Lemma 6 (approximately conserved quantity). Let (ī, E I22 be such that ÿ > > 0 with
g² 2 L. Suppose that we run GD on with step size n with initial point (Io, yo) = / (ã. 0).
for some > 0. Then, there exists to = =0()) such that supezto |x S Ol V Vn) n) and
v 2 O(n)) (3ă =8), -
where the implied constant depends on न. g. and e.
Proof. Let Dt v2- - = and note that
Di+1 = =(v - 11 (((r1)) xz)² - (xt: - (((e))))
= =(1 (2 D..
Since f is 1-Lipschitz, then Do+1 = (1- 0(1))) D..
This shows that for t S 1/17². we have v²² - x² = Dr 2 Do = vi 20 x 7/n. Since (0) = 1,
there exist constants co: C1 > such that (((zyy) 2 (((o) 2 c1 whenever |ayl 2 co. Hence, for all
t 2 1 such that t S 1 1/n2 24 > 0, and [1 2 Co. we have vo2 2 7/n and
Tt+1 = xz- - 77 (x(e)) y = Hi - ( = z; (Vm)).
(10)
Since Fo x Vn/nj, this shows that after at most O(1/n) iterations, we must have either 2; < 0 or
| I cu for the first time. In the first case, (10) shows that | S Vyn. In the second case, since
2² Z 7/n. we have | S vn/7. Let to denote the iteration at which this occurs.
Next, for iterations 2 to. we use the dynamics (10) for x and the fact that (((x)))) has the same sign as xe to conclude that there are two possibilities: either a+++1 has the same sign as It, in which case |x++ll= < |zz]. or De+1 has the opposite sign as Fe, in which case | S n 1e (2(1)) yi S T8 S
O Mm) This implies sup laa| <
V 7) v) as asserted.
19

---

C.2
Gradient flow regime
In this section, we prove Theorem From Lemma there exists an iteration ty such that 1 1 v (2/ 5) and
2 O - 6) S vio - 2
2-
+ 0 " 8
+ Fta
"
In partieular, C = 1 Mro
prove by induction the following facts: for 2 lo-
1. [ < C.
2. ral S exp(=xxp(2((e 2(a (t to))), where min (6, 83.
Suppose that these conditions hold up to iteration 2 to- By Assumption (A2), we have [e((s) - [8] for all + 0. Therefore,
M+++1 1/2 71 (I+Mle) Te (1 )p²t)
exp (-0(, 112
exp -S2(cc exp( a (t to))) y
2 8
2 exp (-0(,
11
Σ
exp( sap(-SQ(a')) -l(a (s to)))) uo 2 exp(-0( 1/2 (2 - -))) Moo
2
8=1o
vi41 2 2-6 0 5) 0 "
(11)
a
In particular, 2- < vi < 2-8/2 throughout. In order for these assertions hold, we require v² S a i.e, 1 S min{ [V8,2 5).
Next, we would like show that + |x decaying exponentially fast. Since
|x++1| = x - 17 (#)()) ya! = | 77 ([arel yı) ye
suffices to consider the case when F, 0. Assumption (A2) implies that
Le+1 2 (1 - mnf) v. 2 - )) Ft.
For the upper bound, split into two cases. We begin by observing that since is twice continuously differentiable near the origin with (0(0) = 1, there is a constant Eo such that |s] < Eo implics
[e'(s) |s]. If s; Tzfee I Eg. then
Fe41 < (1 -
7 91
xt S (1
- 5
Fe.
Otherwise, if 2 8o. then
Vi < (E0)
2 2-6
Te+1 S xz (sn) /1 < d1 7)
S (1 922 -8)) xt
8t
20n
Combining these inequalities, we obtain
| S [zl| exp(-
This verifies the second statement in the induction. The first statement follows because both t + | and Ve. are decreasing.
This shows particular that | > 0, i.e., we have global convergence. To conclude the proof, observe that (11) gives a bound on the final sharpness.
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()) iterations. For the convergence phase, in
order to achieve error, we need |x,|| S
/2- hence, the convergence phase needs only O log !
JEN
iterations. Note that the rate of convergence in the latter phase does not depend the step size m.
20

---

C.3 EoS regime: proof outline
We give brief outline of the proof of Theorem 2: As before, Lemma shows that GD reaches the y-axis approximately at (0, v y/ z At this point, a starts bouncing while y steadily decreases, and we argue that unless 24 or 2 - 2/1,, the GD dynamics cannot stabilize (see Lemma 7).
To bound the gap 2/11 V we look at the first iteration t such that of crosses 2/1,. By making use of Assumption (A2), we simultaneously control both the convergence rate of [are] to zero and the decrease in of in order to prove that
2 0((xe+)))
(12)
see Proposition Therefore, to establish Theorem 2, we must bound | at iteration t.
Controlling the size of [ettll however, is surprisingly delicate as it requires a fine-grained under- standing of the bouncing phase. The insight that guides the proof is the observation that during the bouncing phase, the GD iterates lie close to 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 Zt. Substituting this into the GD dynamics, we obtain
of GD, we have perfect bouncing: T++1 = the equation
n (xete) Vt 2xç
(13)
According Assumption (A2), we have e' (3444) Sett (1)))))))) Together with (13), if y/ = (2 + 6,)/n 2 2/17, where & is sufficiently small, it suggests that
|r S 8. 1/8
(14)
The quasi-static prediction (14) fails when 6, is too small. Nevertheless, we show that it remains accurate as long as 5g 2 178/((-1), and consequently we obtain |xtgel Z 17 1/8-1) Combined with (12), yields Theorem
C.4 EoS regime:
crossing the threshold and the convergence phase
In this section, we prove Theorem 2. We first show that y² must cross 2/4 in order for GD to converge, and we bound the size of the jump across 2/1 once this happens.
Throughout this section and the next, we use the following notation:
8t == Telle: * == (s)))/
In this notation, can write the GD equations
xt-1 (1 11-11 m/,11/ It>
9t-1 = (1 - -);;),]) y/-
We also make a remark regarding Assumption (A2). 3 < Do, then Assumption (A2) equivalent to the following seemingly strongly assumption: for all r > 0, there exists constant c(r) > 0 such that
(s)s < -1- c(r |s|
for all 0 < |s| Sr.
(A2+)
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, |s| > r. then (rr < 1 and the I-Lipschitzness of E imply that (s)s < 1 for r S [s] < r', and hence (s(s// < 1 - d' |s| for sufficiently small constant > O; thus we can take c(r') = C Later, will invoke (A2+) with r chosen be universal constant, so that c(r) can also be thought of as universal.
We begin with the following result about the limiting value of yt-
21

---

Lemma 7 (threshold crossing). Let (2,0) R2 satisfy ÿ > 1 > 0 with 8² - x² = 1. Suppose we
initialize GD with step size ry with initial point (zo, Vo) = v Then either 24 = 0 for some t or
2+6 (Ø, f), where 8 > 0 is constant.
lim yi - - 1-000
Proof. Assume throughout that 22 + 0 for all t. Recall the dynamics for y:
Mi+1 = 11 e (r(t)) x
By assumption (ss)/s - 1 s - 0, and is increasing, so this equation implies that lim inf [2 > 0 then yž must eventually cross 2/n.
Suppose for the sake of contradiction that there exists >>>> 0 with vē > (2 s)/n, for all Let
e' > O be such that - (2+ ) (1 - E') < i,, E'
< 24 Then, there exists 5 > 0 such | S
implies r, -E-,, E', hence
F141|
m > 12 (2+ (1
[xc
The above means that |ze| increases until it exceeds 8, i.e., lim infe-+so [al 2 8. This is our desired contradiction and it implies that lim yi s 2/n-
Lemma (initial gap). Suppose that at some iteration t, we have
2
2
Vt
yt
11
Then, it holds that
2
2
211s
Proof. We can bound
v v22 2m El (Ittit sthht + 1)2 (xtttt 2² 2 3/2 12n-2aae lt
where we used the fact that (s)) |s| for all s E RR
The above lemma shows that the size of the jump across 2/1) is controlled by the size of [st] at the time of the crossing. From Lemma 6, we know that |st| S 1, where the implied constant depends on 8. Hence, the size of the jump is always O(1).
We now provide an analysis of the convergence phase, i.e., after yi crosses 2/4. Proposition (convergence phase). Suppose that v/2 2/n S v Then, GD converges to (0, yoo) satisfying
2 -(((s) S 1o2 <
2 17
7
Proof. Write y = (2 pc)/n, that P: 2 ny?. We write down the update equations for z and for p. First, by the same argument as in the proof of Theorem 1, we have
|m44a| S |
(15)
Next, using 5 1,
V441 = - -1-1ma 3 2 (1 nn yt:
viz1 2 (1 - 2111²) vh=
which translates into
Prt1 S + 2 - + 4n
(16)
22

---

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 pt 2 q (if no such iteration exists. then pt S q for all t). Note that Pr S + O (xe)) due to (15) and (16). By (15), we conclude that for all t' 2 1
[re| - |2x|
(t' f))) - | cxp( expp(-9(q9) (t' - t)))
Substituting this into (16),
Pe Pt + 4n
2 S 9 0 (+)|=t)))))))))))))))))))))))( + On mr
Σ ==1
exp( (9 t))
s=t
2 + On (xl 0 7 | 2
q
By optimizing this bound over 9, we find that for all t,
Pr S v l S n [stl.
Translating this result back into vf yields the result.
Let us take stock of what we have established thus far.
According to Lemma 6, |sel is bounded for all t by a constant.
Then, from Lemma 7 and Lemma 8, we must have either yř - 2/n, or 2/11 - 0()) S s 2/n for some iteration t. In the latter case, Proposition shows that the limiting sharpness is 2/1 O(1).
Note also that the analyses thus far have not made use of Assumption (A2), i.e., we have established the 3 = + case of Theorem 2. Moreover, for all / > 1, the asymptotic 2/n - O()) still shows that the limiting sharpness is close 2/n, 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 |st] at the crossing. Our crude bound of |stl S does not capture the behavior observed in experiments, in which |st| S 17 /(0-1) By substituting this improved bound into Lemma 7, would deduce that the gap the crossing is -2))111(-- and then Proposition 1 would imply that the limiting sharpness is 2/n (11 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, 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 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 |st]. the limiting sharpness can be much closer to 2/1 than what Proposition 1 suggests. The following proposition combined with Lemma 6 implies Theorem 2 for all B 2, but is insufficient for the case 1 < 3 S 2. From now on, we assume 3 < +00.
Proposition (convergence phase; refined), Suppose that y2 < 2/n 2r Then, GD converges to (0, yoo) satisfying
3 2, Onn (ogls((ss /n)),,
(0-11))
2 n
2 Vox 2
-O(n (st
3 = 2,
Ol (-)))
B <<< 2.
Proof. Let y/ (2 -2)pp)) / n before. We quantify the decrease of [2,] in terms of Pt and conversely the increase of pt in terms of |z,| by tracking the half-life of | i.c., the number of iterations it takes [ar| to halve. We call these epochs: at the i-th epoch, we have
2 (-++) v |z $ 2 Vn.
Let in be the index of the first epoch, i.e., in log, /VT/ / Due to Lemma 6, we know that i 2 -O()). From (15), |z,| is monotonically decreasing and consequently |se| is decreasing as
23

---

well. Also, our bound on the limiting sharpness implies that v > 1/1 for all t, provided that is sufficiently small.
Let us compute the dynamics of Pr and (z)). At epoch x > 2 (i+2)) v hence |sp| > 2 (++1).
Assumption (A2+) with r | X 1 implies that
(s)) S -1-c2 <1-c2-
(17)
St
where c c c((st]). This allows to refine (15) on the decrease of [ry to
|041 grev? S (2 pt) pt) (1 c2 ( - 1 1 Pt -
where the first inequality follows from (17) and the second from Pt. = myy? < 1. In tum, this inequality shows that the i-th phase only requires ((2*) iterations. Hence, if t(i) denotes the start of the i-th epoch, then (16) shows that
Pc(i+1) S P(()) + 4n²² 2-22 0(2%)) S Pa() + +00
Summing this up, we have
(
B> 2,
PUO) S Pt 1/
O - i0),
3=2,
0
0((st12-8),
<<2.
In the case of 3 < 2, the final sharpness satisfies 2/n O(pt/)) - (nl S v2. S 2/n.
In the other two cases, suppose that we use this argument until epoch i, such that 2- x 17" Then, we have [x(i,)) 1/7+1/2, [sGG.)) x n7,,, and by using the argument from Proposition from iteration t(t.) onward we obtain
Poo = Po(s.) + Px Po(t.) $ Pt + 0 + +1² x (0(2(8-2) 2);.) = On ( B 2, 00i - in),
B = 2.
We optimize over the choice of y, obtaining 7 1/ and thus [O
3>3 2,
Pro S Pt + O log((stl/n)),
B = 2.
By collecting together the three cases and using Lemma 8 to bound Pt: we finish the proof.
Using the crude bound [stpl S from Lemma 6, yields
((-1())1)),
3> 2,
2 2 0
Ux
3= 2, O(1log(1/n),
"
1)
O(n),
3 < 2,
which is optimal for 3 >
C.5
EoS regime: quasi-static analysis
We supplement Assumption (A2) with corresponding lower bound on (()))
(((s) 2 C [s|
There exists C > 0 such that
for all + 0.
(A3)
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 nj > 0, where f(z y) = Exry) and l satisfies Assumptions (Al), (A2), and (A3). Write y/ = (2 8t)nn and suppose that at
some iteration to, we have [taytaal x 81/8 and dra S 1. Then, for all t 2 to with o, Z 2/8(11)) we have
|onael 1/8
where all implied constants depend on but not n.

---

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
n 2,
(18)
and that during the bouncing phase, this closely agrees with experimental observations (Figure 9). We consider the phase where yē has not yet crossed the threshold 2/1 and write ył == (2 + 6/)/n, thinking of 6, as small. Then, (18) can be written (2 + 5e) r A 2. If we have the behavior (s))/s = 1 - @(sx11)) near the origin, then re R 1 (6r) implics that
|00| ~ 6,
(19)
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 |sy|. which only requires Assump- tion (A2) (to prove a lower bound on |sy]. we need a corresponding lower bound on (s(s)ss. Second, even if we relax (19) read |s 18 S 6c, this fails to hold when & 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 8 < 6, + C" n?", where the added 17? handles the error terms and the exponent y > 0 emerges from the proof.
Proposition 3 (quasi-static analysis; upper bound). For all such that 0 S 6 S 1/13 1 (for a sufficientll small implied constant), it holds that
| < C (ōe
+ CC' 17 1/81(-11),
where C, C' > 0 are constants which may depend on the problem parameters but not on r).
We first show that Theorem now follows.
Proof of Theorem 2. As previously noted, the 3 = too case is handled by the arguments of the previous section, so we focus on B < +oc. From Lemma 7, we either have y - 2/n and |xy| + in which case we are done, or there is an iteration t such that v/ < 2/n S vt-1 From Proposition 3,
since 5 2 0 and 6 I 0, it follows that |st| < 17 //0011)) 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 0(), S. etc. in order to hide constants that depend on l (including (3), but not 6n and n.. 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. 1 < C (6, + +'')") and 6, 2 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 y are chosen later in the proof.
For the base case, observe that the approximate conservation lemma (Lemma 6) gives | S 1, and 5, 2 1/33 V 1) at the beginning of the induction, so the bound is satisfied initially if we choose C sufficiently large enough.
Throughout, we also write ô, = 5, + C'11²' as a convenient shorthand. The strategy is to prove the following two statements:
1. If |s [8 Cjå; for some C > C then |se+1| 1 < C441de+1 for some C++1 S C..
2
2. If |s [ &r& for some C. S 2 |se41| | < Cd
then
Proof of 1. The dynamics for x give
[e+1| = [1 ny8r |2te
By Assumption (A2t) and |sr| S 1,
r $ ((1ss11) ) =
2((C8)) -
and hence
nvrr. (2 + 5,) ((C8))) = = -((())
25

---

for large C. Also, (0) and a similar argument as in the proof of Theorem yields the reverse inequality nyfr, 2 1. We conclude that
0241/=(1110 (1-((((8)) |oca
and hence
|s4+1| 1 - (1 (1-11110 (1-)())))) se = c &,
Since we need bound in terms of dc+1, we use the dynamics of y.
yi+1 = (1 - -- rz) V/ 2 (1 naa) y
vi+1 2 (1 - 2)prr vr
6141 v/L+1 2 z 8. -21/² 2 6, - 227 (C8,)
(20)
Substituting this in,
|8+1| 1 C. (1- ((C8,)) (6.41 + 212² (C3,) 2/0) C,8++1 - (,++1)))++1) + +2'7/ (C8,) 2/9
(21)
Let us show that
3 8 8.t1 2 4
(22)
From (20), have 8e41 2 8, 20p2 (C8.)
2/3 we want prove that x/2 (C8,) 2/8
< 8 If
|3 < 2 this is obvious by taking small, and if 3 > 2 then this is equivalent to C2/8,12 2 8, 81-2/0
1C"1-2/3 77> (1-2/0), which is achieved by taking C" large relative to It suffices to have , <
C' and by taking S (1 2/3); this constraint on will be satisfied by our eventual choice of
y = (B 1).
Returning to (21), in order to finish the proof and in light of (22), want to show that C 2 C1+2//8778,11 Rearranging, suffices to have 8 82-2/8 2 C2/8 '72 or 8, 8.-1/8 1/8 2 11/8-11277 Since by definition ò, > C'1I", by choosing C" large it suffices to have 7 S 1/(1 1/3) = 1/ which leads to our choice of y.
Proof of 2. Using the simple bound ngfre s 2 + 8g, we have
st41| (1 + 6x) |scl,
4 C. cx)((3)) 8,42
$1+1 8 v cxp( | = C. cxp( 8, S
where we used (22). If exp( - 4/3, which holds if 6, x 1/3, then from C, S C/2 we obtain | | s Cd desired.
By following the same proof outline but reversing the inequalities, we can also show a corresponding lower bound on |s, as long as 6, Z 1/166-)), 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 GD do in fact track the quasi-static trajectory.
Proposition 4 stuasi-static analysis; lower bound). Suppose additionally that (A3) holds and that +00. Also, suppose that at some iteration to, we have dt, and that
|scl 2 co.10 1/B
(23)
holds at iteration t = to. where e is sufficiently small constant (depending on the problem parameters but not on 1). Then, (23) also holds for all iterations 2 to such that ô, Z 17 1/8((-11)
Proof. The proof mirrors that of Proposition 3. Let d, 2 1//10-11 for a sufficiently large implied constant. We prove the following two statements:
26

---

1. If | C₁ 6/ 2/18 for some e < 2c, then |s++1| 2 c++1 041 for some c++1 c
62/9
2. If [8t] C2 6, 2/8 for some 2 2c, then |s441| 2 c
Throughout the proof, due Proposition 3, we also have |st| 2 8//30
Proof of I. The dynamics for x give
Fe41 = [1 11-my/r my/rrl [
By Assumption (A3),
2 1- 2 11-20e0 2=210000a221100(())) -
If c is sufficiently small, then
myrri 2 (2+ (2) (1 - 0(c5)) 22 22(6.).
Therefore, we obtain
|2++1 2 ((())) 102]-
On the other hand,
Vt+1 (1 (1²)t y 2 (1 - 0(12)))) v1 (1 0(170,1))) y
(24)
and hence
|9 2 (1+)()))) (1 0 |sa| 2 (1 ((6.) 06720/10) 61/3
24 (1 + ((,,) 01126718) )) 62/3
To conclude, must prove that 17 52 implied constant), then this holds, as was checked in the proof of Proposition Proof of 2. Using Assumption (A3),
2/8 2 8t, but since 5, > 1/ (with sufficiently large
1 0 < 0(1s ()) S $1.
Therefore,
2 2 0(6) s (2+ (1- S noîre S 2 dc
and
-1+ 2 nr/r. 2 -1 -,. -r.
Together with the dynamics for x and (24),
| 2 (1 - O(6.)) (1 - 0(n (; |sel 2 a - 0(8:)) (1- 06720 (2/2)/1))) 6142.
Since a 2 2c, if 8 and r7 are sufficiently small implies | 2 c0
1/3
Convergence rate estimates. Our analysis also provides estimates for the convergence rate of GD in both regimes. Namely, in the gradient flow regime, show that GD converges in O(1/n) iterations, whereas in the EoS regime, GD typically spends 2(11), ( iterations (S((og(1/n)/1) iterations when 3 = 2) in the bouncing phase (Figure 11). Hence, the existence of the bouncing phase dramatically slows down the convergence of GD.
Remark Suppose that at iteration fo, we have 8r, x 1. Then, the assumption of Proposition is that |sto 2 1. If this is not satisfied, i.e., |st, & 1, then the first claim in the proof of Proposition shows that |sty+1| 2 (1 + ((6,)) 18t = (1 + ((1)) |st, Therefore, after l' = 0(lo((//)s) l)) iterations, we obtain [ | Z 1 and then Proposition 4 applies thereafter.
27

---

Remark 3.
From the quasi quasi-static analysis, we can also derive bounds on the length of the bouncing phase. Namely, suppose that to such that òt, and for all t to. we have |st| = 1//0 If å is sufficiently small so that Fr Z for all & 2 to. then the equation for yields
6c+1 S 6ç - ( s.) = 5. - 0 6//
We declare the k-th phase to consist of iterations such that 2 S 6, S 2k-1),, During this phase, 61+1 S 62 - (17² 26/3), so the number of iterations in phase k is x 2 22(2/3-1) /n2. We sum over the phases until 6, x 1y 1/8/(8-1), since after this point the quasi-static analysis fails and yf crosses over 2/1 shortly afterwards. This yields
L/n².
8> 2,
Σ KER
(2/--1)
log(1/n)/rf (1/12/(-11)
==2,
<<2.
1110-1) $2-
The time spent in the bouncing phase increases dramatically as 3 x 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
f(x; a a aaabi) b) a Σ ReLUU-=[[]] + b) + ReLUU++<[]]++ + b).
where x = Aye;; + &. We first approximate
d
Σ
~ ReLU ReLU + +)
In other words, we can ignore the contribution of the signal Aye,, This approximation holds because (i) initially, the bias b is not yet negative enough threshold out the noise, and hence the summation Ei ReLU(=E[]] + b) is of size (d)) and (ii) the difference between the left- and right-hand sides above is simply ReLU(= + E[j] + b) RcLU(=E(5] b), which is of size O(1) and hence negligible compared to the full summation.
Next, letting g(b) == E2-N(0.1 ReLU(= + b) be the 'smoothed' ReLU (see Figure 7), concentration of measure implies the following two facts:
ReLU(+ i + b) d'Egs ReLU(E + =: dg(b) and
DE 1(=o( + b z 0} d dE LE 0) = dg
Indeed, the summations above are sums of i..d. non-negative random variables, and hence its
mean is (2d (as long as N -0(1) and its standard deviation is O(vd). Now, using these approximations, one can rewrite the GD dynamics on the population loss E (u a a b))]
Using these approximations, the output of the ReLU network (2) can be written as
f(z; a att b) (a + a gbb),
which in turn leads to an approximation of the GD dynamics on the population loss (a= aat + Elltogi (ff a aaatt, attb)]]
x att a.- - nn Enggi (u f f(z; ar a bb))
ReLU + b.)
ss da; ta! glbe)
dd((b))
ar l' ey (d at ((b))) d (bbb)),
d
d
b14 = b, - 17 E Mlogi (y f(x; aj a bb)) x (0z Σ 1[-x + 2 07 + a. Σ 1 1(+x
+ +b 2 0))]
d +at g(bb))
(de)
d (1.(
2 b 17 Eyym (d (a, + g(br)) d(aj + + gbb)),
28

---

where (sy = 2 (1og(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:
aH+1 = or sym (d (ar + ( g(be)) dg(b)),
be+1 b, 77 aya (d (az + + d (aH + g')))(()))
Now, we can write the above dynamics more compactly terms of the parameter A, d (aj at).
A++1 A, 2nd2 l. (bx)) g(be),
be+1 = b, 10-7-- Lym (Aig(be)) Arg (be).
E Proofs for the mean model
In this section, we prove the main theorems for the mean model. We first recall the model for the reader's convenience.
A4+1 = A. 2n (Aeg(be)) g(b.), b = b, ( A.g (b.)
E.1 Deferred proofs
In this section, we collect together deferred proofs from Subsection 3.2.
By definition, g(b) = [ (5 + b) ((E) de dE (E) = -E)(E), the first term equals (bb). Moreover, gbb = -b((())+ = (bb.. O
Proof of Lemma 2.
bb(b). Recalling
Proof of Lemma 3. Note that 0, A?) = AÅ 22d e' (Ag(b)) Ag(b) and also that dcx(b) =
-t' (Ag())b) x² (b) Ag' (b) -f" (Ag(b) Ag(b) since r = g/g'. Hence, ( - 2d x(()) 0
and the proof completed.
E.2 Gradient flow regime
Proof of Theorem 3. The following proof is analogous the proof of Theorem We first list several facts we use in the proof:
(i) I (b(b)) (b)) < for all b TR.
exp(x)-1 (ii) ('(s) 2 exp(s)+1
Hence, 16((s) for all R, and we have
(((s 1 1
if | < 2, if > 2.
X 12/ls),
8
(iii) (0) = 1/4.
(iv) (s)()) exp(s) (oxp(s)-1) Hence, (*(() for s > and (*)(s) 0 for (+x
< In
particular, ((((s) |8| for all E
Throughout the proof, we assume that Ag > 0 without loss of generality. We prove by induction the following claim: for t and
min
min(d {8, 8 5, Ao
200
it holds that |Ar S Ao exp(- (xp((7tt). This clearly holds at initialization.
Suppose that the claim holds up iteration t. Using the bounds on Ig'| and [ć], follows that
b+t 2 bt - 1e" (A.gg(()))) A (be) 2 be 2 11
2 b, 2 10 exp(->t) 2 ... 2 bo 2 nAo Σ exp exp(-7s) 2 .1A0
==0
29

---

In particular, b, 2 -1 and ( > 0.08, since 1j - 2 for all t, then we obtain the desired conclusion.
Also, the bound shows that if the claim holds
Ao
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, > 0. From (A() Σ AAg(bb)/4 and g(bi) I g(0) V2=² it follows that
As+1 14-2nd El (Aig(be)) 2 (1 nd6
g(b:)²) A:
ndf
2 o(0)²) A = - à) Ac.
This shows that Art1 2 -(1 -(1 77) A.. Next, we show that Ac+1 S (1-7)) As. First, if Arg(be) s 2.
1
A141 A. 2nd el (Aeg(b:)) g(be) S A,
nd A g(b,))
4
(8 6) g(b,) )²) A, S (1 (8-5)
a 0.082) A, S (1- Ac:
since we have g(be)= > 0.08. Next, A.g(b:) 2 2, then
A-1 = A₁ 2nd E' (Aeg(be)) g(bb)) 1 A, 2 dd g(bb) = =(1
(8 8) g(be) Ae
2 A
$ 1 (8 - 8) π
0.08
At S (1 7) Ac.
2
Ao
This shows that | < (1-1) 7) |r for the case A, > 0. A similar conclusion is obtained for the case A, < O. 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 # for all t. Recall the dynamics for
bi+1 b 7 f' (Aig((b)) Aig (b,).
Since (s)/s 1/4 + 0, and l' is increasing, this equation implies that if lim inf |Agl > 0 then by must keep decreasing until dd (b.) < 2/n.
Suppose for the sake of contradiction that there exists > 0 with 1 gl > (2 + s)/n, for all t. Let e' > 0 be such that 1 - (2 + e) (1 - =E) < -1, i.e., e' < 24e Then, there exists 6 > 0 such
|Ael S 8 implies E' (((A.g(t)))((AAs(()))) > (1-
E'), hence
|A411 1-14 1 (1-
1 ddl g(b,)² +(22+s) - E') - 1|>1.
|A
The above means that |A,] increases until it exceeds i.e. lim ino |Ag| 2 This is our desired contradiction and it implies that lim;- dl g(b)) 2/n.
Remark A straightforward calculation yields that when (a, at b.) is a global minimizer (i.e.. at, b.)) = 1 d2 gb,))² The mean model initialized at (Ao, 0) approximately reaches (0, 0) whose sharpness is d (0)) /2 d² /m. Hence, the bias learning regime 2/1 < d² /44T) precisely corresponds to the EoS regime, 2/n < Amax (V2 (a ot b.)).
a_ + a = 0), then Amax
30


Reviews:
[{"review_id": "5nxp9HjyMY", "reviewer": "Reviewer_gYkS", "summary": "The 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.", "questions": "Some 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.", "limitations": "The authors have listed several interesting future avenues of research. The limitations are evident from the simple setup. There are no negative societal impacts foreseen.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The 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.", "weaknesses": "None, 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your generous review and we are glad that you enjoyed our paper. We will follow your suggestion to move Figure 1 to page 2 and we will correct the typos that you mentioned. Below, we answer your questions:\n- In Fig. 1, the bias was indeed initialized at 0. Note that Figure 1 plots the final bias after training. Over a long time scale, the bias of the ReLU model does become slightly negative, even for small step sizes (around -0.2 in the plot); we believe this is because the GD dynamics for the real network (as opposed to the mean model approximation) is noisier. However, it is remarkable that the real network still exhibits a sharp phase transition. We will include an explanation of this point in the paper.\n- The sharpness is defined as the largest eigenvalue of the Hessian at the current iterate. We will clarify this in the document.\n- We will make this change.\n- We will include a brief explanation of this point in the main text.\n"}, {"review_id": "0n0fEarI7N", "reviewer": "Reviewer_5SQG", "summary": "In 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.", "questions": "1.\tI was wondering what the convergence time for are both gradient flow regime and EoS regime.", "limitations": "The limitations are clearly discussed in the paper. This is a theoretical work and therefore does not seem to have negative societal impact.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1.\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", "weaknesses": "1.\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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We address some of your points below.\n- We agree that extending our analysis to more general settings is of great interest, although it’s beyond the scope of this work. In fact, a recent work [1] was able to prove similar results to ours for more general settings of neural networks, building on the analytical techniques developed in this work.\n\n[1] Song and Yun  ``Trajectory Alignment: Understanding the Edge of Stability Phenomenon via Bifurcation Theory’’ (https://arxiv.org/abs/2307.04204)\n\n\n- Regarding the convergence time in the gradient flow and EoS regimes, please see the paragraph titled “Convergence rate estimates” in Appendix B, as well as Figure 9. In particular, we show that the iteration complexity transitions from $\\eta^{-1}$ in the gradient flow regime, to $\\eta^{-\\min(\\beta/(\\beta-1), 2}$ ($\\log(1/\\eta)/\\eta^2$ when $\\beta = 2$) in the EoS regime, and that these convergence rates agree well with experiments (Figure 9). Unfortunately, we were not able to include this in the main text due to space constraints, but we will try to include a pointer. We hope that this addresses your question.\n"}, {"review_id": "B81gglneU4", "reviewer": "Reviewer_C6cz", "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.", "questions": "Is the bias learning phase transition a good way to describe single hidden layer ReLu models with multiple neurons?", "limitations": "Yes", "rating": 3, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 1, "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.\n\nAdditional suggested references:\n[EOS in a quadratic regression model with formal convergence guarantees](https://arxiv.org/abs/2210.04860)\n[EOS and weight magnitude in quadratic expansion of ReLu network](https://arxiv.org/abs/2205.11787)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review, and for the additional references; we will cite them. We address your points below.\n- Firstly, regarding the advantage of our toy model analysis over prior works: this is a fair criticism. We argue that compared to prior models, our model and analysis is the simplest demonstration of EoS, which is important for intuition. Due to this simplicity, we are able to obtain refined results on how the shape of the loss affects the limiting sharpness and convergence time (see Figure 9). \n\n- Furthermore, we note that the simplicity of our approach inspired a follow-up work (Song and Yun [1]) where the authors extend our main results to more general neural networks.\n\n\n[1] Song and Yun  ``Trajectory Alignment: Understanding the Edge of Stability Phenomenon via Bifurcation Theory’’ (https://arxiv.org/abs/2307.04204)\n\n- In addition to the simplicity of our model, another one of our contributions is to clearly link the phenomenon we study with practical implications for learning ``threshold neurons.’’\n\n\n- Secondly, regarding your question on whether the bias learning describes the behavior of ReLU networks with multiple neurons: in Appendix A.1, we perform experiments for the full sparse coding model with a ReLU network with multiple neurons and we observe qualitatively similar behavior, indicating that our findings do generalize beyond the toy model (see also our experiments for ResNet18 in Appendix A.2).\n\nThank you once again for taking the time to review our paper. Given that you have not raised any concerns regarding the correctness or novelty of our work, if you believe that we have addressed your points (in particular, that we provide conceptual and generalizable insights on the relationship of neural network optimization and generalization), we would appreciate it if you would consider raising your score.\n"}, {"review_id": "SiWPVEdyLv", "reviewer": "Reviewer_poXm", "summary": "The 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", "questions": "In 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?", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The 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.", "weaknesses": "The 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We address some of your points below.\n- Regarding freezing the biases: we acknowledge that even with frozen biases, it is possible for neural networks to succeed at learning binary CIFAR-10. The full story of generalization in deep learning is complex and we do not claim to fully explain it in our paper. Our scope is more limited in scope, as we aim to identify one possible mechanism in which optimization impacts generalization; nevertheless, we see that the phenomenon we observe is indeed borne out in some practical training examples. Although your question is interesting and important, it is beyond the scope of our work.\n- We agree that the insights we reveal are still far from the complex intricacies of training bigger networks. However, our insights carry over to more general models, for example see the interesting recent work [1].\n\n[1] Song and Yun  ``Trajectory Alignment: Understanding the Edge of Stability Phenomenon via Bifurcation Theory’’ (https://arxiv.org/abs/2307.04204)\n- Finally, regarding your last question: we consider the regime of sparse coding in which the noise level is not too high (so that the problem is still statistically solvable); in particular, we expect that for most of the examples, the label is of the same sign as the signal coordinate (by definition of the model). In this case, the architecture that we proposed will correctly threshold out the noise. Please let us know if you have further questions regarding this point.\n \n"}]

Author rebuttal:

Meta-review: The paper brings attention to an intriguing dynamic in the training of deep networks caused by reaching the edge of stability/the break-even point. It is rigorously proven for a simple model that regularization of sharpness enables learning thresholding function implemented by a nonzero bias in ReLU network. 

One of the issues brought by reviewers is that using biases is not required in many successful neural networks. However, as other reviewers have noted, it is a very solid contribution to describe and prove a specific example of how implicit regularization of sharpness changes model activation patterns. 

The key strength of the paper is clarity and focus on a simple but still interesting phenomenon. All in all, it's my pleasure to recommend acceptance of the paper. 

One of the reviewers remarked in the discussion that the paper might bring confusion to the field by overemphasizing the importance of biases in the training of deep nets. I would like to ask the Authors to clarify this point in the camera-ready version.
Decision: Accept (poster)