SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Kwangjun Ahn  
MIT EECS
kjahn@mit.edu
Sebastien Bubeck
Redmond, WA
sebubeck@microsoft.com
Sinho Chewi
Princeton, NJ
schewi@ias.edu
Yin Tat Lee
yintat@uw.edu
Felipe Suarez
felipesc@mit.edu
Yi Zhang
zhayi@microsoft.com
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.
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$ ,  $n = 300$ ,  $\lambda = 3$ , and run gradient descent with the logistic loss. The weights are initialized as  $a^{-}$ ,  $a^{+} \sim \mathcal{N}(0,1/(2d))$  and  $b = 0$ . For each learning rate  $\eta$ , we set the iteration number such that the total time elapsed (iteration  $\times \eta$ ) is 10. The vertical dashed lines indicate our theoretical prediction of the phase transition phenomenon (precise threshold at  $\eta = 8\pi / d^2$ ).
How much do we understand about the training dynamics of neural networks? We begin with a simple and canonical learning task which indicates that the answer is still "far too little".
Motivating example: Consider a binary classification task of labeled pairs  $(\pmb{x}^{(i)},y^{(i)})\in \mathbb{R}^d\times \{\pm 1\}$  where each covariate  $\pmb{x}^{(i)}$  consists of a 1-sparse vector (in an unknown basis) corrupted by additive Gaussian noise, and the label  $y^{(i)}$  is the sign of the non-zero coordinate of the 1-sparse vector. Due to rotational symmetry, we can take the unknown basis to be the standard one and write
$$
\boldsymbol {x} ^ {(i)} = \lambda y ^ {(i)} \boldsymbol {e} _ {j (i)} + \boldsymbol {\xi} ^ {(i)} \in \mathbb {R} ^ {d}, \tag {1}
$$
where  $y^{(i)} \in \{\pm 1\}$  is a random label,  $j(i) \in [d]$  is a random index,  $\pmb{\xi}^{(i)}$  is Gaussian noise, and  $\lambda > 1$  is the unknown signal strength. In fact, (1) is a special case of the well-studied sparse coding model (Olshausen and Field, 1997; Vinje and Gallant, 2000; Olshausen and Field, 2004; Yang et al., 2009; Koehler and Risteski, 2018; Allen-Zhu and Li, 2022). We ask the following fundamental question:
How do neural networks learn to solve the sparse coding problem (1)?
In spite of the simplicity of the setting, a full resolution to this question requires a thorough understanding of surprisingly rich dynamics which lies out of reach of existing theory. To illustrate this point, consider an extreme simplification in which the basis  $e_1, \ldots, e_d$  is known in advance, for which it is natural to parametrize a two-layer ReLU network as
$$
f \left(\boldsymbol {x}; a ^ {-}, a ^ {+}, b\right) = a ^ {-} \sum_ {i = 1} ^ {d} \operatorname {R e L U} \left(- \boldsymbol {x} [ i ] + b\right) + a ^ {+} \sum_ {i = 1} ^ {d} \operatorname {R e L U} \left(+ \boldsymbol {x} [ i ] + b\right). \tag {2}
$$
The parametrization (2) respects the latent data structure (1) well: a good network has a negative bias  $b$  to threshold out the noise, and has  $a^{-} < 0$  and  $a^{+} > 0$  to output correct labels. We are particularly interested in understanding the mechanism by which the bias  $b$  becomes negative, thereby allowing the non-linear ReLU activation to act as a threshold function; we refer to this as the problem of learning "threshold neurons". More broadly, such threshold neurons are of interest as they constitute basic building blocks for producing neural networks with useful inductive bias.
Figure 2: Large learning rates lead to unexpected phenomena: non-monotonic loss and wild oscillations of weights. We choose the same setting as Figure 1. With a small learning rate  $(\eta = 2.5 \cdot 10^{-5})$ , the bias does not decrease noticeably, and the same is true even when we increase the learning rate by ten times  $(\eta = 2.5 \cdot 10^{-4})$ . When we increase the learning rate by another ten times  $(\eta = 2.5 \cdot 10^{-3})$ , we finally see a noticeable decrease in the bias, but with this we observe unexpected behavior: the loss decreases non-monotonically and the sum of second-layer weights  $d \cdot (a^- + a^+)$  oscillates wildly.
We train the parameters  $a^{-}$ ,  $a^{+}$ ,  $b$  using gradient descent with step size  $\eta > 0$  on the logistic loss  $\sum_{i=1}^{n} \ell_{\mathrm{logi}}(y^{(i)} f(\boldsymbol{x}^{(i)}; a^{-}, a^{+}, b))$ , where  $\ell_{\mathrm{logi}}(z) \coloneqq \log (1 + \exp(-z))$ , and we report the results in Figures 1 and 2. The experiments reveal a compelling picture of the optimization dynamics.
- Large learning rates are necessary, both for generalization and for learning threshold neurons. Figure 1 shows that the bias decreases and the test accuracy increases as we increase  $\eta$ ; note that we plot the results after a fixed time (iteration  $\times \eta$ ), so the observed results are not simply because larger learning rates track the continuous-time gradient flow for a longer time.
Large learning rates lead to unexpected phenomena: non-monotonic loss and wild oscillations of  $a^{-} + a^{+}$ . Figure 2 shows that large learning rates also induce stark phenomena, such as non-monotonic loss and large weight fluctuations, which lie firmly outside the explanatory power of existing analytic techniques based on principles from convex optimization.
There is a phase transition between small and large learning rates. In Figure 1, we zoom in on learning rates around  $\eta \approx 0.0006$  and observe sharp phase transition phenomena.
We have presented these observations in the context of the simple ReLU network (2), but we emphasize that these findings are indicative of behaviors observed in practical neural network training settings. In Figure 3, we display results for a two-layer ReLU network trained on the full sparse coding model (1) with unknown basis, as well as a deep neural network trained on CIFAR-10. In each case, we again observe non-monotonic loss coupled with steadily decreasing bias parameters. For these richer models, the transition from small to large learning rates is oddly reminiscent of well-known separations between the "lazy training" or "NTK" regime Jacot et al. (2018) and the more expressive "feature learning" regime. For further experimental results, see Appendix B.
Figure 3: (Top) Results for training an over-parametrized two-layer neural network  $f(\pmb{x}; \pmb{a}, \pmb{W}, b) = \sum_{i=1}^{m} a_i \mathrm{ReLU}\big(\pmb{w}_i^\top \pmb{x} + b\big)$  with  $m \gg d$  for the full sparse coding model (1); in this setting, the basis vectors are unknown, and the neural network learn them through additional parameters  $\pmb{W} = (\pmb{w}_i)_{i=1}^m$ . Also, we use  $m$  different weights  $\pmb{a} = (a_i)_{i=1}^m$  for the second layer. (Bottom) Full-batch gradient descent dynamics of ResNet-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 on the sparse coding model (Arora et al., 2015; Karp et al., 2021). Our inability to explain these findings points to a serious gap in our grasp of neural network training dynamics and calls for a detailed theoretical study.
1.1 Main scope of this work
In this work, we do not aim to understand the sparse coding problem (1) in its full generality. Instead, we pursue the more modest goal of shedding light on the following question.
Q. What is the role of a large step size in learning the bias for the ReLU network (2)?
As discussed above, the dynamics of the simple ReLU network (2) is a microcosm of emergent phenomena beyond the convex optimization regime. In fact, there is a recent growing body of work (Cohen et al., 2021; Arora et al., 2022; Ahn et al., 2022; Lyu et al., 2022; Ma et al., 2022; Wang et al., 2022b; Chen and Bruna, 2023; Damian et al., 2023; Zhu et al., 2023) on training with large learning rates, which largely aims at explaining a striking empirical observation called the "edge of stability (EoS)" phenomenon.
The edge of stability (EoS) phenomenon is a set of distinctive behaviors observed recently by Cohen et al. (2021) when training neural networks with gradient descent (GD). Here we briefly summarize the salient features of the EoS and defer a discussion of prior work to Subsection 1.3. Recall that if we use GD to optimize an  $L$ -smooth loss function with step size  $\eta$ , then the well-known descent lemma from convex optimization ensures monotonic decrease in the loss so long as  $L < 2 / \eta$ . In contrast, when  $L > 2 / \eta$ , 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<sup>1</sup> with constant step size  $\eta > 0$ , the largest eigenvalue of the Hessian at the current iterate (dubbed the "sharpness") initially increases during training ("progressive sharpening") and saturates near or above  $2 / \eta$  ("EoS").
A surprising message of the present work is that the answer to our main question is intimately related to the EoS. Indeed, Figure 4 shows that the GD iterates of our motivating example exhibit the EoS during the initial phase of training when the bias decreases rapidly.
Figure 4: Understanding our main question is surprisingly related to 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 set out to thoroughly understand the workings of the EoS phenomena through a simple example. Specifically, we consider a single-neuron linear neural network in dimension 1, corresponding to the loss
$$
\mathbb {R} ^ {2} \ni (x, y) \mapsto \ell (x y), \quad \text {w h e r e} \ell \text {i s c o n v e x , e v e n , a n d L i p s c h i t z .} \tag {3}
$$
Although toy models have appeared in works on the EoS (see Subsection 1.3), our example is simpler than all prior models, and we provably establish the EoS for (3) with transparent proofs.
We then use the newfound insights gleaned from the analysis of (3) to answer our main question. To the best of our knowledge, we provide the first explanation of the mechanism by which a large learning rate can be necessary for learning threshold neurons.
Explaining the EoS with a single-neuron example. Although the EoS has been studied in various settings (see Subsection 1.3 for a discussion), these works either do not rigorously establish the EoS phenomenon, or they operate under complex settings with opaque assumptions. Here, we study a simple two-dimensional loss function,  $(x,y)\mapsto \ell (xy)$  , where  $\ell$  is convex, even, and Lipschitz. Some examples include  $^2\ell (s) = \frac{1}{2}\log (1 + \exp (-s)) + \frac{1}{2}\log (1 + \exp (+s))$  and  $\ell (s) = \sqrt{1 + s^2}$  . Surprisingly, GD on this loss already exhibits rich behavior (Figure 5).
En route to this result, we rigorously establish the quasi-static dynamics formulated in Ma et al. (2022).
Figure 5: Illustration of two different regimes (the "gradient flow" regime and the "EoS" regime) of the GD dynamics.
<sup>1</sup>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(x;a,b) = abx$ , and that the data is drawn according to  $x = 1$ ,  $y \sim \mathrm{unif}(\{\pm 1\})$ . Then, the population loss under the logistic loss is  $(a,b) \mapsto \ell_{\mathrm{sym}}(ab)$  with  $\ell_{\mathrm{sym}}(s) = \frac{1}{2}\log (1 + \exp (-s)) + \frac{1}{2}\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,  $x$  oscillates wildly while  $y$  decreases (Figure 5);  $(ii)$  for the ReLU network (2), the sum of weights  $(a^{-} + a^{+})$  oscillates while  $b$  decreases (Figure 2).
We study this example in Section 2 and delineate a transition from the "gradient flow" regime to the "EoS regime", depending on the step size  $\eta$  and the initialization. Moreover, in the EoS regime, we rigorously establish asymptotics for the limiting sharpness which depend on the higher-order behavior of  $\ell$ . In particular, for the two losses mentioned above, the limiting sharpness is  $2 / \eta + O(\eta)$ , whereas for losses  $\ell$  which are exactly quadratic near the origin the limiting sharpness is  $2 / \eta + O(1)$ .
In fact, this connection can be made formal by considering an approximation for the GD dynamics for the ReLU network (2). It turns out (see Subsection 3.1 for details) that during the initial phase of training, the dynamics of  $A_{t} \coloneqq d(a_{t}^{-} + a_{t}^{+})$  and  $b_{t}$  due to the ReLU network are well-approximated by the "rescaled" GD dynamics on the loss  $(A,b) \mapsto \ell_{\mathrm{sym}}(A \times g(b))$ , where the step size for the  $A$ -dynamics is multiplied by  $2d^{2}$ ,  $g(b) \coloneqq \mathbb{E}_{z \sim \mathcal{N}(0,1)} \mathrm{ReLU}(z + b)$  is the "smoothed" ReLU, and  $\ell_{\mathrm{sym}}$  is the symmetrized logistic loss; see Subsection 3.1 and Figure 8. We refer to these dynamics as the mean model. The mean model bears a great resemblance to the single-neuron example  $(x,y) \mapsto \ell(xy)$ , and hence we can leverage the techniques developed for the latter in order to
Figure 6: Illustration of GD dynamics 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  $\delta >0$
- if  $\eta \leq (8 - \delta)\pi /d^2$ , then the mean model fails to learn threshold neurons: the limiting bias satisfies  $|b_{\infty}| = O_{\delta}(1 / d^{2})$ .
- if  $\eta \geq (8 + \delta)\pi /d^2$ , then the mean model enters the EoS and learns threshold neurons: the limiting bias satisfies  $b_{\infty} \leq -\Omega_{\delta}(1)$ .
Edge of stability. Our work is motivated by the extensive empirical study of Cohen et al. (2021), which identified the EoS phenomenon. Subsequently, there has been a flurry of works aiming at developing a theoretical understanding of the EoS, which we briefly summarize here.
Properties of the loss landscape. The works (Ahn et al., 2022; Ma et al., 2022) study the properties of the loss landscape that lead to the EoS. Namely, Ahn et al. (2022) argue that the existence of forward-invariant subsets near the minimizers allows GD to convergence even in the unstable regime. They also explore various characteristics of EoS in terms of loss and iterates. Also, Ma et al. (2022) empirically show that the loss landscape of neural networks exhibits subquadratic growth locally around the minimizers. They prove that for a one-dimensional loss, subquadratic growth implies that GD finds a 2-periodic trajectory.
Limiting dynamics. Other works characterize the limiting dynamics of the EoS in various regimes. (Arora et al., 2022; Lyu et al., 2022) show that (normalized) GD tracks a "sharpness reduction flow" near the manifold of minimizers. The recent work of Damian et al. (2023) obtains a different predicted dynamics based on self-stabilization of the GD trajectory. Also, Ma et al. (2022) describes a quasi-static heuristic for the overall trajectory of GD when one component of the iterate is oscillating.
Simple models and beyond. Closely related to our own approach, there are prior works which carefully study simple models. Chen and Bruna (2023) prove global convergence of GD for the two-dimensional function  $(x,y)\mapsto (xy - 1)^{2}$  and a single-neuron student-teacher setting; note that unlike our results, they do not study the limiting sharpness. Wang et al. (2022b) study progressive sharpening for a neural network model. Also, the recent and concurrent work of Zhu et al. (2023) studies the two-dimensional loss  $(x,y)\mapsto (x^{2}y^{2} - 1)^{2}$ ; to our knowledge, their work is the first to asymptotically and rigorously show that the limiting sharpness of GD is  $2 / \eta$  in a simple setting,
at least when initialized locally. In comparison, in Section 2, we perform a global analysis of the limiting sharpness of GD for  $(x,y)\mapsto \ell (xy)$  for a class of convex, even, and Lipschitz losses  $\ell$ , and in doing so we clearly delineate the "gradient flow regime" from the "EoS regime".
Effect of learning rate on learning. Recently, several works have sought to understand how the choice of learning rate affects the learning process, in terms of the properties of the resulting minima (Jastrzebski et al., 2018; Wu et al., 2018; Mulayoff et al., 2021; Nacson et al., 2022) and the behavior of optimization dynamics (Xing et al., 2018; Jastrzebski et al., 2019, 2020; Lewkowycz et al., 2020; Jastrzebski et al., 2021).
Li et al. (2019) demonstrate for a synthetic data distribution and a two-layer ReLU network model that choosing a larger step size for SGD helps with generalization. Subsequent works have shown similar phenomena for regression (Nakkiran, 2020; Wu et al., 2021; Ba et al., 2022), kernel ridge regression Beugnot et al. (2022), and linear diagonal networks Nacson et al. (2022). However, the large step sizes considered in these work still fall under the scope of descent lemma, and most prior works do not theoretically investigate the effect of large step size in the EoS regime. A notable exception is the work of Wang et al. (2022a), which studies the impact of learning rates greater than  $2/$  smoothness for a matrix factorization problem. Also, the recent work of Andriushchenko et al. (2023) seeks to explain the generalization benefit of SGD in the large step size regime by relying on a heuristic SDE model for the case of linear diagonal networks. Despite this similarity, their main scope is quite different from ours, as we  $(i)$  focus on GD instead of SGD and  $(ii)$  establish a direct and detailed analysis of the GD dynamics for a model of the motivating sparse coding example.
2 Single-neuron linear network
In this section, we analyze the single-neuron linear network model  $(x,y)\mapsto f(x,y)\coloneqq \ell (x\times y)$
2.1 Basic properties and assumptions
Basic properties. If  $\ell$  is minimized at 0, then the global minimizers of  $f$  are the  $x$ - and  $y$ -axes. The GD iterates  $x_{t}, y_{t}$ , for step size  $\eta > 0$  and iteration  $t \geq 0$  can be written as
$$
x _ {t + 1} = x _ {t} - \eta \ell^ {\prime} \left(x _ {t} y _ {t}\right) y _ {t}, \quad y _ {t + 1} = y _ {t} - \eta \ell^ {\prime} \left(x _ {t} y _ {t}\right) x _ {t}.
$$
Assumptions. From here onward, we assume  $\eta < 1$  and the following conditions on  $\ell : \mathbb{R} \to \mathbb{R}$ .
(A1)  $\ell$  is convex, even, 1-Lipschitz, and of class  $\mathcal{C}^2$  near the origin with  $\ell''(0) = 1$ .
(A2) There exist constants  $\beta > 1$  and  $c > 0$  with the following property: for all  $s \neq 0$ ,
$$
\ell^ {\prime} (s) / s \leq 1 - c | s | ^ {\beta} \mathbb {1} \{| s | \leq c \}.
$$
We allow  $\beta = +\infty$ , in which case we simply require that  $\frac{\ell'(s)}{s} \leq 1$  for all  $s \neq 0$ .
Assumption (A2) imposes decay of  $s \mapsto \ell'(s)/s$  locally away from the origin in order to obtain more fine-grained results on the limiting sharpness in Theorem 2. As we show in Lemma 5 below, when  $\ell$  is smooth and has a strictly negative fourth derivative at the origin, then Assumption (A2) holds with  $\beta = 2$ . See Example 1 for some simple examples of losses satisfying our assumptions.
2.2 Two different regimes for GD depending on the step size
Before stating rigorous results, in this section we begin by giving an intuitive understanding of the GD dynamics. It turns out that for a given initialization  $(x_0,y_0)$ , there are two different regimes for the GD dynamics depending on the step size  $\eta$ . Namely, there exists a threshold on the step size such that  $(i)$  below the threshold, GD remains close to the gradient flow for all time, and  $(ii)$  above the threshold, GD enters the edge of stability and diverges away from the gradient flow. See Figure 9.
First, recall that the GD dynamics are symmetric in  $x, y$  and that the lines  $y = \pm x$  are invariant. Hence, we may assume without loss of generality that
$$
y _ {0} > x _ {0} > 0, \quad y _ {t} > | x _ {t} | \text {f o r a l l} t \geq 1, \quad \text {a n d G D c o n v e r g e s t o (0 , y _ {\infty}) f o r y _ {\infty} > 0}.
$$
From the expression (8) for the Hessian of  $f$  and our normalization  $\ell''(0) = 1$ , it follows that the sharpness (the largest eigenvalue of loss Hessian) reached by GD in this example is precisely  $y_{\infty}^2$ .
Initially, in both regimes, the GD dynamics tracks the continuous-time gradient flow. Our first observation is that the gradient flow admits a conserved quantity, thereby allowing us to predict the dynamics in this initial phase.
Lemma 1 (conserved quantity). Along the gradient flow for  $f$ , the quantity  $y^2 - x^2$  is conserved.
Proof. Differentiating  $y_{t}^{2} - x_{t}^{2}$  with respect to  $t$  gives  $2y_{t}\left(-\ell^{\prime}(x_{t}y_{t})x_{t}\right) - 2x_{t}\left(-\ell (x_{t}y_{t})y_{t}\right) = 0$ .
Lemma 1 implies that the gradient flow converges to  $(0, y_{\infty}^{\mathrm{GF}}) = (0, \sqrt{y_0^2 - x_0^2})$ . For GD with step size  $\eta > 0$ , the quantity  $y^2 - x^2$  is no longer conserved, but we show in Lemma 6 that it is approximately conserved until the GD iterate lies close to the  $y$ -axis. Hence, GD initialized at  $(x_0, y_0)$  also reaches the  $y$ -axis approximately at the point  $(x_{t_0}, y_{t_0}) \approx (0, \sqrt{y_0^2 - x_0^2})$ .
At this point, GD either approximately converges to the gradient flow solution  $(0, \sqrt{y_0^2 - x_0^2})$  or diverges away from it, depending on whether or not  $y_{t_0}^2 > 2 / \eta$ . To see this, for  $|x_{t_0}y_{t_0}| \ll 1$ , we can Taylor expand  $\ell'$  near zero to obtain the approximate dynamics for  $x$  (recalling  $\ell''(0) = 1$ ),
$$
x _ {t _ {0} + 1} \approx x _ {t _ {0}} - \eta x _ {t _ {0}} y _ {t _ {0}} ^ {2} = \left(1 - \eta y _ {t _ {0}} ^ {2}\right) x _ {t _ {0}}. \tag {4}
$$
From (4), we deduce the following conclusions.