SlowGuess/NeurIPS_2023_OCR huggingface dataset:
(i) If  $y_{t_0}^2 < 2 / \eta$ , then  $|1 - \eta y_{t_0}^2| < 1$ . Since  $y_t$  is decreasing, it implies that  $|1 - \eta y_t^2| < 1$  for all  $t \geq t_0$ , and so  $|x_t|$  converges to zero exponentially fast.
(ii) On the other hand, if  $y_{t_0}^2 > 2 / \eta$ , then  $|1 - \eta y_{t_0}^2| > 1$ , i.e., the magnitude of  $x_{t_0}$  increases in the next iteration, and hence GD cannot stabilize. In fact, in the approximate dynamics,  $x_{t_0 + 1}$  has the opposite sign as  $x_{t_0}$ , i.e.,  $x_{t_0}$  jumps across the  $y$ -axis. One can show that the "bouncing" of the  $x$  variable continues until  $y_t^2$  has decreased past  $2 / \eta$ , at which point we are in the previous case and GD approximately converges to  $(0, 2 / \eta)$ .
This reasoning, combined with the expression for the Hessian of  $f$ , shows that
$$
\begin{array}{l} \operatorname {s h a r p n e s s} \left(0, y _ {\infty}\right) := \lambda_ {\max } \left(\nabla^ {2} f \left(0, y _ {\infty}\right)\right) \approx \min  \left\{y _ {0} ^ {2} - x _ {0} ^ {2}, 2 / \eta \right\} \\ = \min  \left\{\text {g r a d i e n t f l o w s h a r p e n s s}, \text {E o S} \right. \\ \end{array}
$$
Accordingly, we refer to the case  $y_0^2 - x_0^2 < 2 / \eta$  as the gradient flow regime, and the case  $y_0^2 - x_0^2 > 2 / \eta$  as the EoS regime.
See Figure 5 and Figure 9 for illustrations of these two regimes; see also Figure 10 for detailed illustrations of the EoS regime. In the subsequent sections, we aim to make the above reasoning rigorous. For example, instead of the approximate dynamics (4), we consider the original GD dynamics and justify the Taylor approximation. Also, in the EoS regime, rather than loosely asserting that  $|x_{t}| \searrow 0$  exponentially fast and hence the dynamics stabilizes "quickly" once  $y_{t}^{2} < 2 / \eta$ , we track precisely how long this convergence takes so that we can bound the gap between the limiting sharpness and the prediction  $2 / \eta$ .
2.3 Results
Gradient flow regime. Our first rigorous result is that when  $y_0^2 - x_0^2 = (2 - \delta) / \eta$  for some constant  $\delta \in (0, 2)$ , then the limiting sharpness of GD with step size  $\eta$  is  $y_0^2 - x_0^2 + O(1) = (2 - \delta) / \eta + O(1)$ , which is precisely the sharpness attained by the gradient flow up to a controlled error term.
In fact, our theorem is slightly more general, as it covers initializations in which  $\delta$  can scale mildly with  $\eta$ . The precise statement is as follows.
Theorem 1 (gradient flow regime; see Subsection C.2). Suppose we run  $GD$  with step size  $\eta >0$  on the objective  $f$ , where  $f(x,y)\coloneqq \ell (xy)$ , and  $\ell$  satisfies Assumptions (A1) and (A2). Let  $(\tilde{x},\tilde{y})\in \mathbb{R}^2$  satisfy  $\tilde{y} >\tilde{x} >0$  with  $\tilde{y}^2 -\tilde{x}^2 = 1$ . Suppose we initialize  $GD$  at  $(x_0,y_0)\coloneqq (\frac{2 - \delta}{\eta})^{1 / 2}(\tilde{x},\tilde{y})$ , where  $\delta \in (0,2)$  and  $\eta \lesssim \delta^{1 / 2}\wedge (2 - \delta)$ . Then,  $GD$  converges to  $(0,y_{\infty})$  satisfying
$$
\frac {2 - \delta}{\eta} - O (2 - \delta) - O \Big (\frac {\eta}{\min \{\delta , 2 - \delta \}} \Big) \leq \lambda_ {\max} \big (\nabla^ {2} f (0, y _ {\infty}) \big) \leq \frac {2 - \delta}{\eta} + O \Big (\frac {\eta}{2 - \delta} \Big),
$$
where the implied constants depend on  $\tilde{x},\tilde{y}$  , and  $\ell$  , but not on  $\delta ,\eta$
The proof of Theorem 1 is based on a two-stage analysis. In the first stage, we use Lemma 6 on the approximate conservation of  $y^{2} - x^{2}$  along GD in order to show that GD lands near the  $y$ -axis with  $y_{t_0}^2 \approx (2 - \delta) / \eta$ . In the second stage, we use the assumptions on  $\ell$  in order to control the rate of convergence of  $|x_t|$  to 0, which is subsequently used to control the final deviation of  $y_{\infty}^{2}$  from  $(2 - \delta) / \eta$ .
EoS regime. Our next result states that when  $y_0^2 - x_0^2 > 2 / \eta$ , then the limiting sharpness of GD is close to the EoS prediction of  $2 / \eta$ , up to an error term which depends on the exponent  $\beta$  in (A2).
Theorem 2 (EoS; see Subsection C.4). Suppose we run  $GD$  on  $f$  with step size  $\eta >0$ , where  $f(x,y)\coloneqq \ell (xy)$ , and  $\ell$  satisfies (A1) and (A2). Let  $(\tilde{x},\tilde{y})\in \mathbb{R}^2$  satisfy  $\tilde{y} >\tilde{x} >0$  with  $\tilde{y}^2 -\tilde{x}^2 = 1$ . Suppose we initialize  $GD$  at  $(x_0,y_0)\coloneqq \sqrt{(2 + \delta) / \eta} (\tilde{x},\tilde{y})$ , where  $\delta >0$  is a constant. Also, assume that for all  $t\geq 1$  such that  $y_{t}^{2} > 2 / \eta$ , we have  $x_{t}\neq 0$ . Then,  $GD$  converges to  $(0,y_{\infty})$  satisfying
$$
2 / \eta - O \left(\eta^ {1 / (\beta - 1)}\right) \leq \lambda_ {\max } \left(\nabla^ {2} f (0, y _ {\infty})\right) \leq 2 / \eta ,
$$
where the implied constants depend on  $\tilde{x},\tilde{y},\delta \wedge 1$  , and  $\ell$  , but not on  $\eta$
Remarks on the assumptions. The initialization in our results is such that both  $y_0$  and  $y_0 - x_0$  are on the same scale, i.e.,  $y_0, y_0 - x_0 = \Theta(1 / \sqrt{\eta})$ . This rules out extreme initializations such as  $y_0 \approx x_0$ , which are problematic because they lie too close to the invariant line  $y = x$ . Since our aim in this work is not to explore every edge case, we focus on this setting for simplicity. Moreover, we imposed the assumption that the iterates of GD do not exactly hit the  $y$ -axis before crossing  $y^2 = 2 / \eta$ . This is necessary because if  $x_t = 0$  for some iteration  $t$ , then  $(x_{t'}, y_{t'}) = (x_t, y_t)$  for all  $t' > t$ , and hence the limiting sharpness may not be close to  $2 / \eta$ . This assumption holds generically, e.g., if we perturb each iterate of GD with a vanishing amount of noise from a continuous distribution, and we conjecture that for any  $\eta > 0$ , the assumption holds for all but a measure zero set of initializations.
When  $\beta = +\infty$ , which is the case for the Huber loss in Example 1, the limiting sharpness is  $2 / \eta + O(1)$ . When  $\beta = 2$ , which is the case for the logistic and square root losses in Example 1, the limiting sharpness is  $2 / \eta + O(\eta)$ . Numerical experiments show that our error bound of  $O(\eta^{1 / (\beta - 1)})$  is sharp; see Figure 11 below.
We make a few remark about the proof. As we outline the proof in Subsection C.3, in turns out in order to bound the gap  $2 / \eta -y_{\infty}^{2}$ , the proof requires a control of the size  $|x_{\mathbf{t}}y_{\mathbf{t}}|$ , where  $\mathbf{t}$  is the first iteration such that  $y_{\mathbf{t}}^2$  crosses  $2 / \eta$ . However, controlling the size of  $|x_{\mathbf{t}}y_{\mathbf{t}}|$  is surprisingly delicate as it requires a fine-grained understanding of the bouncing phase. The insight that guides the proof is the observation that during the bouncing phase, the GD iterates lie close to a certain envelope (Figure 9).
As a by-product of our analysis, we obtain a rigorous version of the quasi-static principle from which can more accurately track the sharpness gap and convergence rate (see Subsection C.5). The results of Theorem 1, Theorem 2, and Theorem 5 are displayed pictorially as Figure 9.
3 Understanding the bias evolution of the ReLU network
In this section, we use the insights from Section 2 to answer our main question, namely understanding the role of a large step size in learning threshold neurons for the ReLU network (2). Based on the observed dynamics (Figure 2), we can make our question more concrete as follows.
Q. What is the role of a large step size during the "initial phase" of training in which  $(i)$  the bias  $b$  rapidly decreases and  $(ii)$  the sum of weights  $a^{-} + a^{+}$  oscillates?
3.1 Approximating the initial phase of GD with the "mean model"
Deferring details to Appendix D, the GD dynamics for the ReLU network (2) in the initial phase are well-approximated by
$$
\text {G D} (a ^ {-}, a ^ {+}, b) \mapsto \ell_ {\text {s y m}} \left(d \left(a ^ {-} + a ^ {+}\right) g (b)\right),
$$
where  $\ell_{\mathrm{sym}}(s) \coloneqq \frac{1}{2} (\log (1 + \exp (-s)) + \log (1 + \exp (+s)))$  and  $g(b) \coloneqq \mathbb{E}_{z \sim \mathcal{N}(0,1)} \mathrm{ReLU}(z + b)$  is the 'smoothed' ReLU. The GD dynamics can be compactly written in terms of the parameter  $A_t \coloneqq d(a_t^- + a_t^+)$ .
$$
A _ {t + 1} = A _ {t} - 2 d ^ {2} \eta \ell_ {\mathrm {s y m}} ^ {\prime} (A _ {t} g (b _ {t})) g (b _ {t}), \qquad b _ {t + 1} = b _ {t} - \eta \ell_ {\mathrm {s y m}} ^ {\prime} (A _ {t} g (b _ {t})) A _ {t} g ^ {\prime} (b _ {t}). \qquad (5)
$$
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)\mapsto \ell_{\mathrm{sym}}(Ag(b))$ , where the step size for the  $A$ -dynamics is multiplied by  $2d^{2}$ . Due to this resemblance, we can apply the techniques from Section 2.
Figure 7: The 'smoothed' ReLU  $g(b)$
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  $\ell \coloneqq \ell_{\mathrm{sym}}$ , and focus on initializing with  $a_0^{\pm} = \Theta (1 / d)$ ,  $a^{-} + a^{+}\neq 0$ , and  $b_{0} = 0$ . This implies  $A_0 = \Theta (1)$ . We also note the following fact for later use.
Lemma 2 (formula for the smoothed ReLU; see Subsection E.1). The smoothed ReLU function  $g$  can be expressed in terms of the PDF  $\varphi$  and the CDF  $\Phi$  of the standard Gaussian distribution as  $g(b) = \varphi(b) + b\Phi(b)$ . In particular,  $g' = \Phi$ .
Note also that  $b_{t}$  is monotonically decreasing. This is because  $\ell^{\prime}(A_t g(b_t)) A_t g'(b_t) \geq 0$  since  $\ell'$  is an odd function and  $g(b), g'(b) > 0$  for any  $b \in \mathbb{R}$ .
Following Subsection 2.2, we begin with the continuous-time dynamics of the mean model:
$$
\dot {A} = - 2 d ^ {2} \ell^ {\prime} (A g (b)) g (b), \quad \dot {b} = - \ell^ {\prime} (A g (b)) A g ^ {\prime} (b). \tag {6}
$$
Lemma 3 (conserved quantity; see Subsection E.1). Let  $\kappa : \mathbb{R} \to \mathbb{R}$  be defined as  $\kappa(b) \coloneqq \int_0^b g / g'$ . Along the gradient flow (6), the quantity  $\frac{1}{2} A^2 - 2d^2\kappa(b)$  is conserved.
Based on Lemma 3, if we initialize the continuous-time dynamics (6) at  $(A_0,0)$  and if  $A_{t}\rightarrow 0$ , then the limiting value of the bias  $b_{\infty}^{\mathrm{GF}}$  satisfies  $\kappa (b_{\infty}^{\mathrm{GF}}) = -\frac{1}{4d^{2}} A_{0}^{2}$ , which implies that  $b_{\infty}^{\mathrm{GF}} = -\Theta (\frac{1}{d^2})$ ; indeed, this holds since  $\kappa '(0) = g(0) / g'(0) > 0$ , so there exist constants  $c_{0},c_{1} > 0$  such that  $c_{0}b\leq \kappa (b)\leq c_{1}b$  for all  $-1\leq b\leq 0$ . Since the mean model (5) tracks the continuous-time dynamics (6) until it reaches the  $b$ -axis, the mean model initialized at  $(A_0,0)$  also approximately reaches  $(A_{t_0},b_{t_0})\approx (0, - \Theta (\frac{1}{d^2}))\approx (0,0)$  in high dimension  $d\gg 1$ . In other words, the continuous-time dynamics (6) fails to learn threshold neurons.
Once the mean model reaches the  $b$ -axis, we again identify two different regimes depending on the step size. A Taylor expansion of  $\ell'$  around the origin yields the following approximate dynamics (here  $\ell''(0) = 1/4$ ):  $A_{t_0 + 1} \approx A_{t_0} - \frac{\eta d^2}{2} A_{t_0} g(b_{t_0})^2 = A_{t_0}\left(1 - \frac{\eta d^2}{2} g(b_{t_0})^2\right)$ . We conclude that the condition which now dictates whether we have bouncing or convergence is  $\frac{1}{2} d^2 g(b_{t_0})^2 > 2/\eta$ .
(i) Gradient flow regime: If  $2 / \eta >d^{2}g(0)^{2} / 2 = d^{2} / (4\pi)$  (since  $g(0)^2 = 1 / (2\pi)$ ), i.e., the step size  $\eta$  is below the threshold  $8\pi /d^2$ , then the final bias of the mean model  $b_{\infty}^{\mathrm{MM}}$  satisfies  $b_{\infty}^{\mathrm{MM}}\approx b_{\infty}^{\mathrm{GF}}\approx 0$ . In other words, the mean model fails to learn threshold neurons.
(ii)  $\widetilde{EoS}$  regime: If  $2 / \eta < d^2 / (4\pi)$ , i.e., the step size  $\eta$  is above the threshold  $8\pi / d^2$ , then  $\frac{1}{2} d^2 g^2(b_{\infty}^{\mathrm{MM}}) < 2 / \eta$ , i.e.,  $b_{\infty}^{\mathrm{MM}} < g^{-1}(2 / \sqrt{\eta d^2})$ . For instance, if  $\eta = \frac{10\pi}{d^2}$ , then  $b_{\infty}^{\mathrm{MM}} < -0.087$ . In other words, the mean model successfully learns threshold neurons.
3.3 Results for the mean model
Theorem 3 (mean model, gradient flow regime; see Appendix E). Consider the mean model (5) initialized at  $(A_0,0)$ , with step size  $\eta = \frac{(8 - \delta)\pi}{d^2}$  for some  $\delta > 0$ . Let  $\gamma \coloneqq \frac{1}{200} \min \{\delta, 8 - \delta, \frac{8 - \delta}{|A_0|}\}$ . Then, as long as  $\eta \leq \gamma / |A_0|$ , the limiting bias  $b_{\infty}^{\mathrm{MM}}$  satisfies
$$
0 \geq b _ {\infty} ^ {\mathrm {M M}} \geq - (\eta / \gamma) | A _ {0} | = - O _ {A _ {0}, \delta} (1 / d ^ {2}).
$$
In other words, the mean model fails to learn threshold neurons.
Theorem 4 (mean model, EoS regime; see Appendix E). Consider the mean model initialized at  $(A_0,0)$ , with step size  $\eta = \frac{(8 + \delta)\pi}{d^2}$  for some  $\delta > 0$ . Furthermore, assume that for all  $t \geq 1$  such that  $\frac{1}{2} d^2 g(b_t)^2 > 2/\eta$ , we have  $A_t \neq 0$ . Then, the limiting bias  $b_{\infty}^{\mathrm{MM}}$  satisfies
$$
b _ {\infty} ^ {\mathrm {M M}} \leq g ^ {- 1} \left(2 / \sqrt {(8 + \delta) \pi}\right) \leq - \Omega_ {\delta} (1).
$$
For instance, if  $\eta = \frac{10\pi}{d^2}$ , then  $b_{\infty}^{\mathrm{MM}} < -0.087$ . In other words, the mean model successfully learns threshold neurons.
In this paper, we present the first explanation for the emergence of threshold neuron (i.e., ReLU neurons with negative bias) in models such as the sparse coding model (1) through a novel connection with the "edge of stability" (EoS) phenomenon. Along the way, we obtain a detailed and rigorous understanding of the dynamics of GD in the EoS regime for a simple class of loss functions, thereby shedding light on the impact of large learning rates in non-convex optimization.
Our approach is largely inspired by the recent paradigm of "physics-style" approaches to understanding deep learning based on simplified models and controlled experiments (c.f. (Zhang et al., 2022; von Oswald et al., 2023; Abernethy et al., 2023; Allen-Zhu and Li, 2023; Li et al., 2023; Ahn et al., 2023a,b)). We found such physics-style approach quite effective to understand deep learning, especially given the complexity of modern deep neural networks. We hope that our work inspires further research on understanding the working mechanisms of deep learning.
Many interesting questions remain, and we conclude with some directions for future research.
- Extending the analysis of EoS to richer models. Although the analysis we present in this work is restricted to simple models, the underlying principles can potentially be applied to more general settings. In this direction, it would be interesting to study models which capture the impact of the depth of the neural network on the EoS phenomenon. Notably, a follow-up work by Song and Yun (2023) uses bifurcation theory to extend our results to more complex models.
- The interplay between the EoS and the choice of optimization algorithm. As discussed in Subsection 2.3, the bouncing phase of the EoS substantially slows down the convergence of GD (see Figure 11). Investigating how different optimization algorithms (e.g., SGD, or GD with momentum) interact with the EoS phenomenon could potentially lead to practical speed-ups or improved generalization. Notably, a follow up work by Dai et al. (2023) studies the working mechanisms of a popular modern optimization technique called sharpness-aware minimization (Foret et al., 2021) based on our sparse coding problem.
- An end-to-end analysis of the sparse coding model. Finally, we have left open the motivating question of analyzing how two-layer ReLU networks learn to solve the sparse coding model (1). Despite the apparent simplicity of the problem, its analysis has thus far remained out of reach, and we believe that a resolution to this question would constitute compelling and substantial progress towards understanding neural network learning. We are hopeful that the insights in this paper provide the first step towards this goal.
We thank Ronan Eldan, Suriya Gunasekar, Yuanzhi Li, Jonathan Niles-Weed, and Adil Salim for initial discussions on this project. KA was supported by the ONR grant (N00014-20-1-2394) and MIT-IBM Watson as well as a Vannevar Bush fellowship from Office of the Secretary of Defense. SC was supported by the NSF TRIPODS program (award DMS-2022448).
Jacob Abernethy, Alekh Agarwal, Teodor V. Marinov, and Manfred K. Warmuth. A mechanism for sample-efficient in-context learning for sparse retrieval tasks. arXiv preprint arXiv:2305.17040, 2023.
Kwangjun Ahn, Jingzhao Zhang, and Suvrit Sra. Understanding the unstable convergence of gradient descent. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 247-257. PMLR, 7 2022.
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 2023 (arXiv:2306.00297), 2023a.
Kwangjun Ahn, Xiang Cheng, Minhak Song, Chulhee Yun, Ali Jadbabaie, and Suvrit Sra. Linear attention is (maybe) all you need (to understand transformer optimization). arXiv 2310.01082, 2023b.
Zeyuan Allen-Zhu and Yanzhi Li. Feature purification: how adversarial training performs robust deep learning. In 2021 IEEE 62nd Annual Symposium on Foundations of Computer Science (FOCS), pages 977-988. IEEE, 2022.
Zeyuan Allen-Zhu and Yanzhi Li. Physics of language models: part 1, context-free grammar. arXiv preprint arXiv:2305.13673, 2023.
Zeyuan Allen-Zhu, Yanzhi Li, and Zhao Song. A convergence theory for deep learning via overparameterization. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 242-252. PMLR, 6 2019.
Maksym Andriushchenko, Aditya V. Varre, Loucas Pillaud-Vivien, and Nicolas Flammarion. SGD with large step sizes learns sparse features. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 903-925. PMLR, 7 2023.
Sanjeev Arora, Rong Ge, Tengyu Ma, and Ankur Moitra. Simple, efficient, and neural algorithms for sparse coding. In Peter Grünwald, Elad Hazan, and Satyen Kale, editors, Proceedings of the 28th Conference on Learning Theory, volume 40 of Proceedings of Machine Learning Research, pages 113-149, Paris, France, 7 2015. PMLR.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 322-332. PMLR, 6 2019.
Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on the edge of stability in deep learning. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szeptesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 948-1024. PMLR, 7 2022.
Samet Oymak and Mahdi Soltanolkotabi. Toward moderate overparameterization: global convergence guarantees for training shallow neural networks. IEEE Journal on Selected Areas in Information Theory, 1(1):84-105, 2020.
Grant M. Rotskoff and Eric Vanden-Eijnden. Trainability and accuracy of artificial neural networks: an interacting particle system approach. Comm. Pure Appl. Math., 75(9):1889-1935, 2022.
Minhak Song and Chulhee Yun. Trajectory alignment: understanding the edge of stability phenomenon via bifurcation theory. NeurIPS 2023 (arXiv:2307.04204), 2023.
William E. Vinje and Jack L. Gallant. Sparse coding and decorrelation in primary visual cortex during natural vision. Science, 287(5456):1273-1276, 2000.
Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vlademyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pages 35151-35174. PMLR, 2023.
Yuqing Wang, Minshuo Chen, Tuo Zhao, and Molei Tao. Large learning rate tames homogeneity: convergence and balancing effect. In International Conference on Learning Representations, 2022a.
Zixuan Wang, Zhouzi Li, and Jian Li. Analyzing sharpness along GD trajectory: progressive sharpening and edge of stability. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 9983-9994. Curran Associates, Inc., 2022b.
Jingfeng Wu, Difan Zou, Vladimir Braverman, and Quanquan Gu. Direction matters: on the implicit bias of stochastic gradient descent with moderate learning rate. In International Conference on Learning Representations, 2021.
Lei Wu, Chao Ma, and Weinan E. How SGD selects the global minima in over-parameterized learning: a dynamical stability perspective. Advances in Neural Information Processing Systems, 31:8279-8288, 2018.
Chen Xing, Devansh Arpit, Christos Tsirigotis, and Yoshua Bengio. A walk with SGD. arXiv preprint arXiv:1802.08770, 2018.
Jianchao Yang, Kai Yu, Yihong Gong, and Thomas Huang. Linear spatial pyramid matching using sparse coding for image classification. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pages 1794-1801. IEEE, 2009.
Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling transformers with lego: a synthetic reasoning task. arXiv preprint arXiv:2206.04301, 2022.
Xingyu Zhu, Zixuan Wang, Xiang Wang, Mo Zhou, and Rong Ge. Understanding edge-of-stability training dynamics with a minimalist example. In The Eleventh International Conference on Learning Representations, 2023.
A Additional illustrations for Section 2 15
B Further experimental results 16
B.1 Experiments for the full sparse coding model 17
B.2 Experiments on the CIFAR-10 dataset 17
C Proofs for the single-neuron linear network 18
C.1 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 21
C.5 EoS regime: quasi-static analysis 24
D Deferred derivations of mean model 28
E Proofs for the mean model 29
E.1 Deferred proofs 29
E.2 Gradient flow regime 29
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.
Figure 9: Two regimes for GD. We run GD on the square root loss with step size  $\frac{1}{4}$ . The gradient flow regime is illustrated on the left for  $(x_0, y_0) = (3, 4)$ . GD (blue) tracks the gradient flow (green) when  $\eta < 2 / (y_0^2 - x_0^2)$ . Otherwise, as illustrated on the right for  $(x_0, y_0) = (3, 6)$ , GD is in the EoS regime and goes through a gradient flow phase (blue), an intermediate bouncing phase (orange) that tracks the quasi-static envelope (purple), and a converging phase (red).