SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Next, we next present detailed illustrations of the edge-of-stability regime depending on the choice of step size. Compare this plot with our theoretical results characterized in Theorem 2.
Figure 10: We plot the GD trajectory for  $\ell(s) = \sqrt{1 + s^2}$  and sharpness for step sizes  $2/\eta_1 = 9$ ,  $2/\eta_2 = 7$ ,  $2/\eta_3 = 5$ , and  $2/\eta_4 = 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  $\beta$ .
Figure 11: (Left) Log-log plot of the sharpness gap as a function of  $\eta$ , for  $\ell_{\beta}$  in Example 1 and  $\beta = \frac{3}{2}$ , 2, 3, 10. (Right) Log-log plot of the iteration count for the bouncing region with  $y_t^2 \in [\frac{2}{\eta}, \frac{3}{\eta}]$  as a function of  $\eta$ , for  $\ell_{\beta}$  in Example 1 and  $\beta = \frac{4}{3}$ ,  $\frac{3}{2}$ , 2, 4. The dashed lines show the predicted sharpness gap and iteration count with an offset computed via linear regression of the data for  $\eta < e^{-2}$ .
B Further experimental results
In this section, we report further experimental results which demonstrate that our theory, while limited to the specific models we study (namely, the single-neuron example and the mean model), is in fact
indicative of behaviors commonly observed in more realistic instances of neural network training. In particular, we show that threshold neurons often emerge in the presence of oscillations in the other weight parameters of the network.
B.1 Experiments for the full sparse coding model
We provide the details for the top plot of Figure 3. consider the sparse coding model in the form (1). Compared to (2), we assume that the basis vectors are unknown, and the neural network learn them through additional parameters  $\pmb{W} = (\pmb{w}_i)_{i=1}^m$  together with  $m$  different weights  $\pmb{a} = (a_i)_{i=1}^m$  for the second layer as follows:
$$
f (\boldsymbol {x}; \boldsymbol {a}, \boldsymbol {W}, b) = \sum_ {i = 1} ^ {m} a _ {i} \operatorname {R e L U} \left(\langle \boldsymbol {w} _ {i}, \boldsymbol {x} \rangle + b\right). \tag {7}
$$
We show results for  $d = 100$ ,  $m = 2000$ . We generate  $n = 20000$  data points according to the aforementioned sparse coding model with  $\lambda = 5$ . We use the He initialization, i.e.,  $\pmb{a} \sim \mathcal{N}(0, I_m / m)$ ,  $\pmb{w} \sim \mathcal{N}(0, I_d / d)$ , and  $b = 0$ . As shown in the top plot of Figure 3, the bias decreases more with the large learning rate. Further, we report the behavior of the average of second layer weights in Figure 12 (left), and confirm that the sum oscillates.
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 is simply a scalar.
B.2 Experiments on the CIFAR-10 dataset
Next, we provide the details for the bottom plot of Figure 3. We train ResNet-18 on a binarized version of the CIFAR-10 dataset formed by taking only the first two classes; this is done for the purpose of monitoring the average logit of the network. The average logit is measured over the entire training set. The median bias is measured at the last convolutional layer right before the pooling. For the optimizer, we use full-batch GD with no momentum or weight decay, plus a cosine learning rate scheduler where learning rates shown in the plots are the initial values.
Oscillation of expected output (logit) of the network. Bearing a striking resemblance to our two-layer models, as one can see from Figure 12 (right) that the expected mean of the output (logit) of the deep net also oscillates due to GD dynamics. As we have argued in the previous sections, this occurs as the bias parameters are driven towards negative values.
Results for SGD. In Figure 13, we report qualitatively similar phenomena when we instead train ResNet-18 with stochastic gradient descent (SGD), where we use all ten classes of CIFAR-10. Again, the median bias is measured at the last convolutional layer. We further report the average activation which is the output of the ReLU activation at the last convolutional layer, averaged over the neurons and the entire training set. The average activation statistics represent the hidden representations before the linear classifier part, and lower values represent sparser representations. Interestingly, the threshold neuron also emerges with larger step sizes similarly to the case of gradient descent.
Figure 13: SGD dynamics of ResNet-18 on (multiclass) CIFAR-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  $\ell$ .
Basic properties. If  $\ell$  is minimized at 0, then the global minimizers of  $f$  are the  $x$ - and  $y$ -axes. The gradient and Hessian of  $f$  are given by:
$$
\nabla f (x, y) = \ell^ {\prime} (x y) \left[ \begin{array}{c} y \\ x \end{array} \right],
$$
$$
\nabla^ {2} f (x, y) = \ell^ {\prime \prime} (x y) \left[ \begin{array}{l} y \\ x \end{array} \right] ^ {\otimes 2} + \ell^ {\prime} (x y) \left[ \begin{array}{c c} 0 & 1 \\ 1 & 0 \end{array} \right]. \tag {8}
$$
This results in GD iterates  $x_{t}, y_{t}$ , for step size  $\eta > 0$  and iteration  $t \geq 0$ :
$$
x _ {t + 1} = x _ {t} - \eta \ell^ {\prime} \left(x _ {t} y _ {t}\right) y _ {t},
$$
$$
y _ {t + 1} = y _ {t} - \eta \ell^ {\prime} \left(x _ {t} y _ {t}\right) x _ {t}.
$$
Lemma 4 (invariant lines). Assume that  $\ell$  is even, so that  $\ell'$  is odd. Then, the lines  $y = \pm x$  are invariant for gradient descent on  $f$ .
Proof. If  $y_{t} = \pm x_{t}$ , then
$$
y _ {t + 1} = y _ {t} - \eta \ell^ {\prime} (x _ {t} y _ {t}) x _ {t} = \pm x _ {t} \mp \eta \ell^ {\prime} (x _ {t} ^ {2}) x _ {t},
$$
$$
x _ {t + 1} = x _ {t} - \eta \ell^ {\prime} \left(x _ {t} y _ {t}\right) y _ {t} = x _ {t} - \eta \ell^ {\prime} \left(x _ {t} ^ {2}\right) x _ {t},
$$
and hence  $y_{t + 1} = \pm x_{t + 1}$ . Note that the iterates  $(x_{t})_{t\geq 0}$  are the iterates of GD with step size  $\eta$  on the one-dimensional loss function  $x\mapsto \frac{1}{2}\ell (x^{2})$ .
We focus instead on initializing away from these two lines. We now state our assumptions on  $\ell$ .
We gather together some elementary properties of  $\ell$ .
Lemma 5 (properties of  $\ell$ ). Suppose that Assumption (A1) holds.
1.  $\ell$  is minimized at the origin and  $\ell'(0) = 0$ .
2. Suppose that  $\ell$  is four times continuously differentiable near the origin. If Assumption (A2) holds, then  $\ell^{(4)}(0) \leq 0$ . Conversely, if  $\ell^{(4)}(0) < 0$ , then Assumption (A2) holds for  $\beta = 2$ .
Proof. The first statement is straightforward. The second statement follows from Taylor expansion: for  $s \neq 0$  near the origin,
$$
\frac {\ell^ {\prime} (s)}{s} = \frac {\ell^ {\prime} (0) + \ell^ {\prime \prime} (0) s + \int_ {0} ^ {s} (s - r) \ell^ {\prime \prime \prime} (r) \mathrm {d} r}{s} = 1 + \int_ {0} ^ {s} \left(1 - \frac {r}{s}\right) \ell^ {\prime \prime \prime} (r) \mathrm {d} r. \tag {9}
$$
Since  $\ell'''$  is odd, then Assumption (A2) and (9) imply that  $\ell''''$  is non-positive on  $(0, \varepsilon)$  for some  $\varepsilon > 0$ , which in turn implies  $\ell^{(4)}(0) \leq 0$ . Conversely, if  $\ell^{(4)}(0) < 0$ , then there exists  $\varepsilon > 0$  such that  $\ell'''(s) \leq -\varepsilon s$  for  $s \in (0, \varepsilon)$ . From (9), we see that  $\ell'(s)/s \leq 1 - \varepsilon \int_{0}^{s}(s-r)\mathrm{d}r \leq 1 - \varepsilon s^2/2$ . By symmetry, we conclude that Assumption (A2) holds with  $\beta = 2$  and some  $c > 0$ .
We give some simple examples of losses satisfying our assumptions.
Example 1. The examples below showcase several functions  $\ell$  that satisfy Assumptions (A1) and (A2) with different values of  $\beta$ .
- Rescaled and symmetrized logistic loss.  $\ell_{\mathrm{rsym}}(s) \coloneqq \frac{1}{2}\ell_{\mathrm{logi}}(-2s) + \frac{1}{2}\ell_{\mathrm{logi}}(+2s)$ .
Note  $\ell_{\mathrm{rsym}}^{\prime}(s) = \tanh (s)$ , thus  $\ell_{\mathrm{rsym}}^{\prime}(s) / s\leq 1$  and  $\ell_{\mathrm{rsym}}^{\prime}(s) / s\leq 1 - \frac{1}{4} |s|^2$ , for  $|s| < \frac{1}{4}$ .
- Square root loss.  $\ell_{\mathrm{sqrt}}(s) \coloneqq \sqrt{1 + s^2}$ .
Note  $\ell_{\mathrm{sqrt}}'(s) = \frac{s}{\sqrt{1 + s^2}}$ , thus  $\ell_{\mathrm{sqrt}}'(s) / s \leq 1$  and  $\ell_{\mathrm{sqrt}}'(s) / s \leq 1 - \frac{2}{5} |s|^2$ , for  $|s| < \frac{2}{5}$ .
- Huber loss.  $\ell_{\mathrm{Hub}}(s) \coloneqq \frac{s^2}{2} \mathbb{1}\{s \in [-1, 1]\} + \left(|s| - \frac{1}{2}\right) \mathbb{1}\{s \notin [-1, 1]\}$ .
Note  $\ell_{\mathrm{Hub}}^{\prime}(s) = s\mathbb{1}\{s\in [-1,1]\} +\operatorname {sgn}(s)\mathbb{1}\{s\notin [-1,1]\}$ , thus  $\ell_{\mathrm{Hub}}^{\prime}(s) / s\leq 1$ , i.e., we have Assumption (A2) with  $\beta = +\infty$
- Higher-order. For  $\beta > 1$  let  $c_{\beta} \coloneqq \frac{1}{\beta + 1}\left(\frac{\beta}{\beta + 1}\right)^{\beta}$  and  $r_{\beta} \coloneqq \frac{\beta + 1}{\beta}$ . We define  $\ell_{\beta}$  implicitly via its derivative
$$
\ell_ {\beta} ^ {\prime} (s) := s \left(1 - c _ {\beta} | s | ^ {\beta}\right) \mathbb {1} \{s ^ {2} <   r _ {\beta} ^ {2} \} + \operatorname {s g n} (s)   \mathbb {1} \{s ^ {2} \geq r _ {\beta} ^ {2} \}.
$$
By definition,  $\ell_{\beta}^{\prime}(s) / s\leq 1$  and  $\ell_{\beta}^{\prime}(s) / s\leq 1 - c_{\ell}|s|^{\beta}$ , where  $c_{\ell} = c_{\beta}\wedge r_{\beta}$ .
We now prove our main results from Subsection 2.3 in order.
C.1 Approximate conservation along GD
We begin by stating and proving the approximate conservation of  $y^{2} - x^{2}$  for the GD dynamics.
Lemma 6 (approximately conserved quantity). Let  $(\tilde{x},\tilde{y})\in \mathbb{R}^2$  be such that  $\tilde{y} >\tilde{x} >0$  with  $\tilde{y}^2 -\tilde{x}^2 = 1$ . Suppose that we run  $GD$  on  $f$  with step size  $\eta$  with initial point  $(x_0,y_0)\coloneqq \sqrt{\frac{\gamma}{\eta}} (\tilde{x},\tilde{y})$ , for some  $\gamma >0$ . Then, there exists  $t_0 = O(\frac{1}{\eta})$  such that  $\sup_{t\geq t_0}|x_t|\leq O(\sqrt{(\gamma^{-1}\lor\gamma)\eta})$  and
$$
y _ {t _ {0}} ^ {2} - x _ {t _ {0}} ^ {2} = \left(1 - O (\eta)\right) \left(y _ {0} ^ {2} - x _ {0} ^ {2}\right),
$$
where the implied constant depends on  $\tilde{x},\tilde{y}$  ,and  $\ell$
Proof. Let  $D_{t} \coloneqq y_{t}^{2} - x_{t}^{2}$  and note that
$$
\begin{array}{l} D _ {t + 1} = \left(y _ {t} - \eta \ell^ {\prime} \left(x _ {t} y _ {t}\right) x _ {t}\right) ^ {2} - \left(x _ {t} - \eta \ell^ {\prime} \left(x _ {t} y _ {t}\right) y _ {t}\right) ^ {2} \\ = \left(1 - \eta^ {2} \ell^ {\prime} \left(x _ {t} y _ {t}\right) ^ {2}\right) D _ {t}. \\ \end{array}
$$
Since  $\ell$  is 1-Lipschitz, then  $D_{t + 1} = (1 - O(\eta^2))D_t$ .
This shows that for  $t \lesssim 1/\eta^2$ , we have  $y_t^2 - x_t^2 = D_t \gtrsim D_0 = y_0^2 - x_0^2 \asymp \gamma/\eta$ . Since  $\ell''(0) = 1$ , there exist constants  $c_0, c_1 > 0$  such that  $\ell'(|xy|) \geq \ell'(c_0) \geq c_1$  whenever  $|xy| \geq c_0$ . Hence, for all  $t \geq 1$  such that  $t \lesssim 1/\eta^2$ ,  $x_t > 0$ , and  $|x_t y_t| \geq c_0$ , we have  $y_t^2 \gtrsim \gamma/\eta$  and
$$
x _ {t + 1} = x _ {t} - \eta \ell^ {\prime} \left(x _ {t} y _ {t}\right) y _ {t} = x _ {t} - \Theta (\eta y _ {t}) = x _ {t} - \Theta \left(\sqrt {\gamma \eta}\right). \tag {10}
$$
Since  $x_0 \asymp \sqrt{\gamma / \eta}$ , this shows that after at most  $O(1 / \eta)$  iterations, we must have either  $x_t < 0$  or  $|x_t y_t| \leq c_0$  for the first time. In the first case, (10) shows that  $|x_t| \lesssim \sqrt{\gamma \eta}$ . In the second case, since  $y_t^2 \gtrsim \gamma / \eta$ , we have  $|x_t| \lesssim \sqrt{\eta / \gamma}$ . Let  $t_0$  denote the iteration at which this occurs.
Next, for iterations  $t \geq t_0$ , we use the dynamics (10) for  $x$  and the fact that  $\ell'(x_t y_t)$  has the same sign as  $x_t$  to conclude that there are two possibilities: either  $x_{t+1}$  has the same sign as  $x_t$ , in which case  $|x_{t+1}| \leq |x_t|$ , or  $x_{t+1}$  has the opposite sign as  $x_t$ , in which case  $|x_{t+1}| \leq \eta |\ell'(x_t y_t)| y_t \leq \eta y_t \leq O(\sqrt{\gamma \eta})$ . This implies  $\sup_{t \geq t_0} |x_t| \leq O(\sqrt{(\gamma^{-1} \vee \gamma) \eta})$  as asserted.
C.2 Gradient flow regime
In this section, we prove Theorem 1. From Lemma 6, there exists an iteration  $t_0$  such that  $|x_{t_0}| \lesssim \sqrt{\eta/(2-\delta)}$  and
$$
\frac {2 - \delta}{\eta} - O (2 - \delta) \leq y _ {t _ {0}} ^ {2} \leq \frac {2 - \delta}{\eta} + x _ {t _ {0}} ^ {2} \leq \frac {2 - \delta}{\eta} + O \left(\frac {\eta}{2 - \delta}\right).
$$
In particular,  $C\coloneqq |x_{t_0}y_{t_0}|\lesssim 1$
We prove by induction the following facts: for  $t \geq t_0$ ,
1.  $|x_{t}y_{t}| \leq C$ .
2.  $|x_{t}| \leq |x_{t_{0}}| \exp(-\Omega(\alpha(t - t_{0})))$ , where  $\alpha \coloneqq \min\{\delta, 2 - \delta\}$ .
Suppose that these conditions hold up to iteration  $t \geq t_0$ . By Assumption (A2), we have  $|\ell'(s)| \leq |s|$  for all  $s \neq 0$ . Therefore,
$$
\begin{array}{l} y _ {t + 1} = y _ {t} - \eta \ell^ {\prime} \left(x _ {t} y _ {t}\right) x _ {t} \geq \left(1 - \eta x _ {t} ^ {2}\right) y _ {t} \\ \geq \exp \left(- O \left(\frac {\eta^ {2}}{2 - \delta}\right) \exp \left(- \Omega (\alpha (t - t _ {0}))\right)\right) y _ {t} \\ \geq \exp \Bigl (- O \bigl (\frac {\eta^ {2}}{2 - \delta} \bigr) \sum_ {s = t _ {0}} ^ {t} \exp \bigl (- \Omega (\alpha (s - t _ {0})) \bigr) \Bigr) y _ {t _ {0}} \geq \exp \Bigl (- O \bigl (\frac {\eta^ {2}}{\alpha (2 - \delta)} \bigr) \Bigr) y _ {t _ {0}}, \\ \end{array}
$$
$$
y _ {t + 1} ^ {2} \geq \frac {2 - \delta}{\eta} - O (2 - \delta) - O \left(\frac {\eta}{\alpha}\right). \tag {11}
$$
In particular,  $\frac{1}{2}\frac{2 - \delta}{\eta} \leq y_t^2 \leq \frac{2 - \delta / 2}{\eta}$  throughout. In order for these assertions to hold, we require  $\eta^2 \lesssim \alpha(2 - \delta)$ , i.e.,  $\eta \lesssim \min\{\sqrt{\delta}, 2 - \delta\}$ .
Next, we would like to show that  $t \mapsto |x_t|$  is decaying exponentially fast. Since
$$
\left| x _ {t + 1} \right| = \left| x _ {t} - \eta \ell^ {\prime} \left(x _ {t} y _ {t}\right) y _ {t} \right| = \left| \left| x _ {t} \right| - \eta \ell^ {\prime} \left(\left| x _ {t} \right| y _ {t}\right) y _ {t} \right|,
$$
it suffices to consider the case when  $x_{t} > 0$ . Assumption (A2) implies that
$$
x _ {t + 1} \geq \left(1 - \eta y _ {t} ^ {2}\right) x _ {t} \geq - \left(1 - \frac {\delta}{2}\right) x _ {t}.
$$
For the upper bound, we split into two cases. We begin by observing that since  $\ell$  is twice continuously differentiable near the origin with  $\ell''(0) = 1$ , there is a constant  $\varepsilon_0$  such that  $|s| < \varepsilon_0$  implies  $|\ell'(s)| \geq \frac{1}{2} |s|$ . If  $s_t \coloneqq x_t y_t \leq \varepsilon_0$ , then
$$
x _ {t + 1} \leq \left(1 - \frac {\eta}{2} y _ {t} ^ {2}\right) x _ {t} \leq \left(1 - \frac {2 - \delta}{4}\right) x _ {t}.
$$
Otherwise, if  $s_t \geq \varepsilon_0$ , then
$$
x _ {t + 1} \leq x _ {t} - \eta \ell^ {\prime} (\varepsilon_ {0}) y _ {t} \leq x _ {t} - \eta \ell^ {\prime} (\varepsilon_ {0}) \frac {y _ {t} ^ {2}}{s _ {t}} \leq x _ {t} - \eta \ell^ {\prime} (\varepsilon_ {0}) \frac {2 - \delta}{2 C \eta} x _ {t} \leq (1 - \Omega (2 - \delta)) x _ {t}.
$$
Combining these inequalities, we obtain
$$
\left| x _ {t + 1} \right| \leq \left| x _ {t} \right| \exp \bigl (- \Omega (\alpha) \bigr).
$$
This verifies the second statement in the induction. The first statement follows because both  $t \mapsto |x_t|$  and  $t \mapsto y_t$  are decreasing.
This shows in particular that  $|x_{t}| \searrow 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\left(\frac{1}{\eta}\right)$  iterations. For the convergence phase, in order to achieve  $\varepsilon$  error, we need  $|x_{t}| \lesssim \frac{\sqrt{\varepsilon\eta}}{\sqrt{2 - \delta}}$ ; hence, the convergence phase needs only  $O\left(\frac{1}{\alpha} \log \frac{1}{\varepsilon}\right)$  iterations. Note that the rate of convergence in the latter phase does not depend on the step size  $\eta$ .
C.3 EoS regime: proof outline
We give a brief outline of the proof of Theorem 2: As before, Lemma 6 shows that GD reaches the  $y$ -axis approximately at  $(0, \sqrt{y_0^2 - x_0^2})$ . At this point,  $x$  starts bouncing while  $y$  steadily decreases, and we argue that unless  $x_t = 0$  or  $y_t^2 \leq 2 / \eta$ , the GD dynamics cannot stabilize (see Lemma 7).
To bound the gap  $2 / \eta -y_{\infty}^{2}$ , we look at the first iteration  $\mathfrak{t}$  such that  $y_{\mathfrak{t}}^{2}$  crosses  $2 / \eta$ . By making use of Assumption (A2), we simultaneously control both the convergence rate of  $|x_{t}|$  to zero and the decrease in  $y_{t}^{2}$  in order to prove that
$$
y _ {\infty} ^ {2} \geq \frac {2}{\eta} - O \left(\left| x _ {\mathbf {t}} y _ {\mathbf {t}} \right|\right), \tag {12}
$$
see Proposition 1. Therefore, to establish Theorem 2, we must bound  $|x_{\mathbf{t}}y_{\mathbf{t}}|$  at iteration  $\mathbf{t}$ .
Controlling the size of  $|x_{\mathbf{t}}y_{\mathbf{t}}|$ , however, is surprisingly delicate as it requires a fine-grained understanding of the bouncing phase. The insight that guides the proof is the observation that during the bouncing phase, the GD iterates lie close to a certain envelope (Figure 9). This envelope is predicted by the quasi-static heuristic as described in Ma et al. (2022). Namely, suppose that after one iteration of GD, we have perfect bouncing:  $x_{t + 1} = -x_t$ . Substituting this into the GD dynamics, we obtain the equation
$$
\eta \ell^ {\prime} \left(x _ {t} y _ {t}\right) y _ {t} = 2 x _ {t}. \tag {13}
$$
According to Assumption (A2), we have  $\ell'(x_t y_t) = x_t y_t (1 - \Omega(|x_t y_t|^\beta))$ , Together with (13), if  $y_t^2 = (2 + \delta_t) / \eta \geq 2 / \eta$ , where  $\delta_t$  is sufficiently small, it suggests that
$$
\left| x _ {t} y _ {t} \right| \lesssim \delta_ {t} ^ {1 / \beta}. \tag {14}
$$
The quasi-static prediction (14) fails when  $\delta_t$  is too small. Nevertheless, we show that it remains accurate as long as  $\delta_t \gtrsim \eta^{\beta / (\beta - 1)}$ , and consequently we obtain  $|x_{\mathbf{t}}y_{\mathbf{t}}| \lesssim \eta^{1 / (\beta - 1)}$ . Combined with (12), it yields Theorem 2.
C.4 EoS regime: crossing the threshold and the convergence phase
In this section, we prove Theorem 2. We first show that  $y_{t}^{2}$  must cross  $2 / \eta$  in order for GD to converge, and we bound the size of the jump across  $2 / \eta$  once this happens.
Throughout this section and the next, we use the following notation:
$s_t\coloneqq x_ty_t$
-  $r_t \coloneqq \ell'(s_t) / s_t$ .
In this notation, we can write the GD equations as
$$
x _ {t + 1} = \left(1 - \eta r _ {t} y _ {t} ^ {2}\right) x _ {t},
$$
$$
y _ {t + 1} = \left(1 - \eta r _ {t} x _ {t} ^ {2}\right) y _ {t}.
$$
We also make a remark regarding Assumption (A2). If  $\beta < +\infty$ , then Assumption (A2) is equivalent to the following seemingly strongly assumption: for all  $r > 0$ , there exists a constant  $c(r) > 0$  such that
$$
\frac {\ell^ {\prime} (s)}{s} \leq 1 - c (r) | s | ^ {\beta}, \quad \text {f o r a l l} 0 <   | s | \leq r. \tag {A2}
$$
Indeed, Assumption (A2) states that  $(\mathrm{A}2^{+})$  holds for some  $r > 0$ . To verify that  $(\mathrm{A}2^{+})$  holds for some larger  $r' > r$ , we can split into two cases. If  $|s| \leq r$ , then  $\ell'(s)/s \leq 1 - c|s|^{\beta}$ . Otherwise, if  $|s| > r$ , then  $\ell'(r)/r < 1$  and the 1-Lipschitzness of  $\ell'$  imply that  $\ell'(s)/s < 1$  for  $r \leq |s| \leq r'$ , and hence  $\ell'(s)/s \leq 1 - c'|s|^{\beta}$ , for a sufficiently small constant  $c' > 0$ ; thus we can take  $c(r') = c \wedge c'$ . Later, we will invoke  $(\mathrm{A}2^{+})$  with  $r$  chosen to be a universal constant, so that  $c(r)$  can also be thought of as universal.
We begin with the following result about the limiting value of  $y_{t}$ .
Lemma 7 (threshold crossing). Let  $(\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$  with step size  $\eta$  with initial point  $(x_0,y_0)\coloneqq \sqrt{\frac{2 + \delta}{\eta}} (\tilde{x},\tilde{y})$ , where  $\delta >0$  is a constant. Then either  $x_{t} = 0$  for some  $t$  or
$$
\lim  _ {t \to \infty} y _ {t} ^ {2} \leq \frac {2}{\eta}.
$$
Proof. Assume throughout that  $x_{t} \neq 0$  for all  $t$ . Recall the dynamics for  $y$ :
By assumption  $\ell^{\prime}(s) / s\to 1$  as  $s\to 0$ , and  $\ell^\prime$  is increasing, so this equation implies that if  $\lim \inf_{t\to \infty}|x_t| > 0$  then  $y_{t}^{2}$  must eventually cross  $2 / \eta$ .
Suppose for the sake of contradiction that there exists  $\varepsilon > 0$  with  $y_t^2 > (2 + \varepsilon) / \eta$ , for all  $t$ . Let  $\varepsilon' > 0$  be such that  $1 - (2 + \varepsilon)(1 - \varepsilon') < -1$ , i.e.,  $\varepsilon' < \frac{\varepsilon}{2 + \varepsilon}$ . Then, there exists  $\delta > 0$  such that  $|x_t| \leq \delta$  implies  $r_t > 1 - \varepsilon'$ , hence
$$
\frac {| x _ {t + 1} |}{| x _ {t} |} = | 1 - \eta r _ {t} y _ {t} ^ {2} | > | (2 + \varepsilon) (1 - \varepsilon^ {\prime}) - 1 | > 1.
$$
The above means that  $|x_{t}|$  increases until it exceeds  $\delta$ , i.e.,  $\lim \inf_{t\to \infty}|x_t|\geq \delta$ . This is our desired contradiction and it implies that  $\lim_{t\to \infty}y_t^2\leq 2 / \eta$ .
Lemma 8 (initial gap). Suppose that at some iteration  $\mathbf{t}$ , we have
$$
y _ {\mathbf {t} + 1} ^ {2} <   \frac {2}{\eta} \leq y _ {\mathbf {t}} ^ {2}.
$$
$$
y _ {\mathbf {t} + 1} ^ {2} \geq \frac {2}{\eta} - 2 \eta s _ {\mathbf {t}} ^ {2}.
$$
Proof. We can bound
$$
y _ {\mathbf {t} + 1} ^ {2} = y _ {\mathbf {t}} ^ {2} - 2 \eta \ell^ {\prime} (x _ {\mathbf {t}} y _ {\mathbf {t}}) x _ {\mathbf {t}} y _ {\mathbf {t}} + \eta^ {2} \ell^ {\prime} (x _ {\mathbf {t}} y _ {\mathbf {t}}) ^ {2} x _ {\mathbf {t}} ^ {2} \geq y _ {\mathbf {t}} ^ {2} - 2 \eta | x _ {\mathbf {t}} y _ {\mathbf {t}} | ^ {2},
$$
where we used the fact that  $|\ell^{\prime}(s)|\leq |s|$  for all  $s\in \mathbb{R}$
The above lemma shows that the size of the jump across  $2 / \eta$  is controlled by the size of  $|s_{\mathbf{t}}|$  at the time of the crossing. From Lemma 6, we know that  $|s_{\mathbf{t}}| \lesssim 1$ , where the implied constant depends on  $\delta$ . Hence, the size of the jump is always  $O(\eta)$ .
We now provide an analysis of the convergence phase, i.e., after  $y_{t}^{2}$  crosses  $2 / \eta$ .
Proposition 1 (convergence phase). Suppose that  $y_{\mathbf{t}}^2 < 2 / \eta \leq y_{\mathbf{t} - 1}^2$ . Then,  $GD$  converges to  $(0, y_{\infty})$  satisfying
$$
\frac {2}{\eta} - O (| s _ {\mathbf {t}} |) \leq y _ {\infty} ^ {2} \leq \frac {2}{\eta}.
$$
Proof. Write  $y_{t}^{2} = (2 - \rho_{t}) / \eta$ , so that  $\rho_t = 2 - \eta y_t^2$ . We write down the update equations for  $x$  and for  $\rho$ . First, by the same argument as in the proof of Theorem 1, we have
$$
\left| x _ {t + 1} \right| \leq \left| x _ {t} \right| \exp \left(- \Omega \left(\rho_ {t}\right)\right). \tag {15}
$$
Next, using  $r_t \leq 1$ ,
$$
y _ {t + 1} = \left(1 - \eta r _ {t} x _ {t} ^ {2}\right) y _ {t} \geq \left(1 - \eta x _ {t} ^ {2}\right) y _ {t},
$$
$$
y _ {t + 1} ^ {2} \geq \left(1 - 2 \eta x _ {t} ^ {2}\right) y _ {t} ^ {2},
$$
which translates into
$$
\rho_ {t + 1} \leq \rho_ {t} + 2 \eta^ {2} x _ {t} ^ {2} y _ {t} ^ {2} \leq \rho_ {t} + 4 \eta x _ {t} ^ {2}. \tag {16}
$$