SlowGuess/NeurIPS_2023_OCR huggingface dataset:
- If  $x_{t} > 0$ , then  $x_{t + 1} < x_{t}$ . If  $x_{t} < 0$ , then  $x_{t + 1} > x_{t}$ .
- If  $y_{t} > 0$ , then  $y_{t + 1} < y_{t}$ . If  $y_{t} < 0$ , then  $y_{t + 1} > y_{t}$ .
Proof. We only show the case that  $x_{t} > 0$  as the other cases are similar. Recall the dynamics of  $x_{t}$ :
$$
x _ {t + 1} = x _ {t} - \eta \ell_ {t} ^ {\prime} \cdot \rho z _ {t} ^ {- 1} x _ {t} - \eta \ell_ {t} ^ {\prime} \cdot y _ {t}.
$$
Since  $x_{t} \neq 0$ , Lemma 17 implies that  $\mathrm{sign}(\ell_t^\prime) = \mathrm{sign}(x_t y_t)$ , and so
$$
\begin{array}{l} x _ {t + 1} = x _ {t} - \eta \ell_ {t} ^ {\prime} \cdot \rho z _ {t} ^ {- 1} x _ {t} - \eta \ell_ {t} ^ {\prime} \cdot y _ {t} \\ = x _ {t} - \eta \mathrm {s i g n} (x _ {t} y _ {t}) | \ell_ {t} ^ {\prime} | \cdot \rho \mathrm {s i g n} (x _ {t} y _ {t}) (x _ {t} ^ {2} + y _ {t} ^ {2}) ^ {- 1 / 2} x _ {t} - \eta \mathrm {s i g n} (x _ {t} y _ {t}) | \ell_ {t} ^ {\prime} | \cdot \mathrm {s i g n} (y _ {t}) | y _ {t} | \\ = x _ {t} - \eta \left| \ell_ {t} ^ {\prime} \right| \cdot \rho \left(x _ {t} ^ {2} + y _ {t} ^ {2}\right) ^ {- 1 / 2} x _ {t} - \eta \left| \ell_ {t} ^ {\prime} \right| \cdot \left| y _ {t} \right| \leq x _ {t}, \\ \end{array}
$$
where the last line uses the assumption that  $x_{t} > 0$ .
F.2 Initial Phase:  $x_{t}$  Decreases Fast while  $y_{t}$  Remains Large
As advertised in Theorem 16, we group all rounds until  $x_{t} \leq \frac{c}{4}\sqrt{\gamma\eta}$  as Initial Phase. Formally:
Definition 19 (Initial Phase). Let  $t_1$  be the largest time such that  $x_t > \frac{c}{4}\sqrt{\gamma\eta}$  for all  $t \leq t_1$ . We call the iterations  $[0, t_1]$  Initial Phase.
Theorem 20 (Main Conclusion of Initial Phase; SAM Case). For  $t_0 = \Theta(\min\{\eta^{-1.5}\rho^{-1}\gamma^{1/2}, \eta^{-2}\})$ , we have  $y_t \geq \frac{1}{2}\sqrt{\gamma/\eta}$  for all  $t \leq \min\{t_0, t_1\}$  under the conditions of Theorem 16. Moreover, we have  $x_{t+1} \leq x_t - \frac{c}{4}\sqrt{\gamma\eta}$  for all  $t \leq \min\{t_0, t_1\}$ , which consequently infers  $t_1 \leq t_0$  under the conditions of Theorem 16, i.e.,  $\min\{t_0, t_1\}$  is just  $t_1$ . This shows the first claim of Theorem 16.
Remark. This theorem can be intuitively explained as follows: At initialization,  $x_0, y_0$  are both  $\Theta(\sqrt{1/\eta})$ . However, after  $t_1$  iterations, we have  $|x_{t_1+1}| = \mathcal{O}(\sqrt{\eta})$ ; meanwhile,  $y_t$  is still of order  $\Omega(\sqrt{1/\eta})$  (much larger than  $\mathcal{O}(\sqrt{\eta})$ ). Hence,  $|x_t|$  and  $|y_t|$  get widely separated in Initial Phase.
Proof. This theorem is a direct consequence of Lemma 21 and Lemma 24.
Lemma 21. There exists  $t_0 = \Theta(\min\{\eta^{-1.5}\rho^{-1}\gamma^{1/2}, \eta^{-2}\})$  such that  $|y_t| \geq \frac{1}{2}\sqrt{\gamma/\eta}$  for  $t \leq t_0$ . If  $\eta$  is sufficiently small s.t.  $\eta\rho + \sqrt{C\gamma\eta} \leq \frac{1}{2}\sqrt{\gamma/\eta}$ , then  $y_t \geq \frac{1}{2}\sqrt{\gamma/\eta}$  for  $t \leq \min\{t_1, t_0\}$ .
Proof. Let  $t_0$  be defined as follows:
$$
t _ {0} \triangleq \max  \left\{t: \left(1 - \frac {8}{\sqrt {C}} \eta^ {1. 5} \gamma^ {- 1 / 2} \rho - \eta^ {2}\right) ^ {t} \geq \frac {1}{4} \right\}. \tag {14}
$$
Then, it follows that  $t_0 = \Theta (\min \{\eta^{-1.5}\rho^{-1}\gamma^{1 / 2},\eta^{-2}\})$  . We first prove the following claim.
Claim 22. For all  $t \leq t_0$ ,  $y_t^2 \geq \frac{1}{4}\gamma/\eta$ .
Proof. We prove this claim by induction. We trivially have  $y_0^2 = C\gamma/\eta \geq \frac{1}{4}C\gamma/\eta$ . Now suppose that the conclusion holds up to some  $t < t_0$ , i.e.,  $y_{t'}^2 \geq \frac{1}{4}\gamma/\eta$  for all  $t' \leq t$ . We consider  $y_{t+1}^2 - x_{t+1}^2$ :
$$
\begin{array}{l} y _ {t + 1} ^ {2} - x _ {t + 1} ^ {2} = (1 - \eta \ell_ {t} ^ {\prime} \cdot \rho z _ {t} ^ {- 1}) ^ {2} (y _ {t} ^ {2} - x _ {t} ^ {2}) - (\eta \ell_ {t} ^ {\prime}) ^ {2} (y _ {t} ^ {2} - x _ {t} ^ {2}) \\ = (1 - 2 \eta \ell_ {t} ^ {\prime} \cdot \rho z _ {t} ^ {- 1} + \eta^ {2} (\ell_ {t} ^ {\prime} \cdot \rho z _ {t} ^ {- 1}) ^ {2} - \eta^ {2} (\ell_ {t} ^ {\prime}) ^ {2}) \cdot (y _ {t} ^ {2} - x _ {t} ^ {2}) \\ \stackrel {(a)} {\geq} \left(1 - 2 \eta \rho z _ {t} ^ {- 1} - \eta^ {2}\right) \cdot \left(y _ {t} ^ {2} - x _ {t} ^ {2}\right) \\ \geq \left(1 - 2 \eta \rho y _ {t} ^ {- 1} - \eta^ {2}\right) \cdot \left(y _ {t} ^ {2} - x _ {t} ^ {2}\right) \\ \stackrel {(b)} {\geq} \left(1 - \frac {8}{\sqrt {C}} \eta^ {1. 5} \gamma^ {- 1 / 2} \rho - \eta^ {2}\right) \cdot \left(y _ {t} ^ {2} - x _ {t} ^ {2}\right), \\ \end{array}
$$
where (a) used  $\ell_t^\prime \leq 1$  (thanks to Assumption (A1)) and (b) used the induction hypothesis  $y_{t}^{2}\geq$ $\frac{1}{4} C\gamma /\eta$  . This further implies that
$$
y _ {t + 1} ^ {2} - x _ {t + 1} ^ {2} \geq \left(1 - \frac {8}{\sqrt {C}} \eta^ {1. 5} \gamma^ {- 1 / 2} \rho - \eta^ {2}\right) ^ {t + 1} (y _ {0} ^ {2} - x _ {0} ^ {2}) = \left(1 - \frac {8}{\sqrt {C}} \eta^ {1. 5} \gamma^ {- 1 / 2} \rho - \eta^ {2}\right) ^ {t + 1} \frac {\gamma}{\eta}.
$$
Thus, by the definition of  $t_0$ , we must have  $y_{t + 1}^2 \geq y_{t + 1}^2 - x_{t + 1}^2 \geq \frac{1}{4}\gamma/\eta$ , which proves the claim.
Next, we prove the second conclusion.
Claim 23. If  $\eta \rho +\sqrt{C\gamma\eta}\leq \frac{1}{2}\sqrt{\gamma / \eta}$ , then  $y_{t}\geq \frac{1}{2}\sqrt{\gamma / \eta}$  for  $t\leq \min \{t_0,t_1\}$ .
Proof. Still show by induction. Let  $y_{t} \geq \frac{1}{2}\sqrt{\gamma / \eta} > 0$  for some  $t < \min\{t_0, t_1\}$ . Consider  $y_{t + 1}$ .
By Definition 19,  $x_{t}$  is positive for all  $t \leq t_1$ . Thus, using Lemma 18, we have  $x_{t} \leq x_{0} \leq y_{0} = \sqrt{C\gamma / \eta}$ . Since  $x_{t}, y_{t} > 0$ , we have  $\mathrm{sign}(\ell_t') = \mathrm{sign}(x_t y_t) = \mathrm{sign}(x_t) > 0$  and hence (13) gives
$$
\begin{array}{l} y _ {t + 1} = y _ {t} - \eta \left| \ell_ {t} ^ {\prime} \right| \rho \left(x _ {t} ^ {2} + y _ {t} ^ {2}\right) ^ {- 1 / 2} \left| y _ {t} \right| - \eta \left| \ell_ {t} ^ {\prime} \right| \left| x _ {t} \right| \\ \geq y _ {t} - \eta \rho - \eta | x _ {t} |, \\ \end{array}
$$
where we used the facts that  $|\ell_t^{\prime}|\leq 1$  and  $(x_{t}^{2} + y_{t}^{2})^{-1 / 2}\leq \min \{|x_{t}|^{-1},|y_{t}|^{-1}\}$  . Hence, we get
$$
y _ {t + 1} \geq y _ {t} - \eta \rho - \sqrt {C \gamma \eta} \geq 0
$$
since  $\eta \rho +\sqrt{C\gamma\eta}\leq \frac{1}{2}\sqrt{\gamma / \eta}\leq y_t$ . By Claim 22,  $y_{t + 1}\geq 0$  implies  $y_{t + 1}\geq \frac{1}{2}\sqrt{\gamma / \eta}$  as well.
Combining Claim 22 and Claim 23 finishes the proof of Lemma 21.
Lemma 24. For any  $t \leq \min\{t_0, t_1\}$ , we have
$$
x _ {t + 1} \leq x _ {t} - \frac {c}{4} \sqrt {\gamma \eta}.
$$
In particular, if  $\eta$  is sufficiently small s.t.  $\frac{4\sqrt{C}}{c}\eta^{-1} \leq t_0$ , then we must have  $t_1 \leq \frac{4}{c}\eta^{-1} - 1 < t_0$ .
Proof. Since  $x_{t} > 0$  for all  $t \leq t_1$ , we have  $\mathrm{sign}(\ell_t') = \mathrm{sign}(x_t y_t) = \mathrm{sign}(y_t)$  from Lemma 17 and Lemma 21, so the SAM update (13) becomes
$$
x _ {t + 1} = x _ {t} - \eta | \ell_ {t} ^ {\prime} | \cdot \rho (x _ {t} ^ {2} + y _ {t} ^ {2}) ^ {- 1 / 2} x _ {t} - \eta | \ell_ {t} ^ {\prime} | \cdot | y _ {t} |.
$$
Since  $x_{t} > 2c\sqrt{\frac{\eta}{C\gamma}}$  we have
$$
(x _ {t} + \rho y _ {t} z _ {t} ^ {- 1}) (y _ {t} + \rho x _ {t} z _ {t} ^ {- 1}) > x _ {t} y _ {t} \geq 2 c \sqrt {\frac {\eta}{\gamma}} \cdot \frac {1}{2} \sqrt {\frac {\gamma}{\eta}} \geq c,
$$
which implies that  $\ell_t^\prime \geq c / 2$  from Assumption (A3). Together with Lemma 21, we have
$$
\begin{array}{l} x _ {t + 1} = x _ {t} - \eta | \ell_ {t} ^ {\prime} | \cdot \rho (x _ {t} ^ {2} + y _ {t} ^ {2}) ^ {- 1 / 2} z _ {t} ^ {- 1} x _ {t} - \eta | \ell_ {t} ^ {\prime} | \cdot | y _ {t} | \\ \leq x _ {t} - \eta | \ell_ {t} ^ {\prime} | \cdot | y _ {t} | \leq x _ {t} - \eta \frac {c}{2} \frac {1}{2} \sqrt {\frac {\gamma}{\eta}} = x _ {t} - \frac {c}{4} \sqrt {\gamma \eta}. \\ \end{array}
$$
Let  $t_1' \triangleq \frac{\sqrt{C\gamma/\eta}}{\frac{c}{4}\sqrt{\gamma\eta}} = \frac{4\sqrt{C}}{c}\eta^{-1}$ . Since  $x_0 \leq \sqrt{C\gamma/\eta}$ , we have  $x_{t_1'} < \frac{c}{4}\sqrt{\gamma\eta}$  as long as  $t_1 \leq t_0$ . Thus, it follows that  $t_1 \leq t_1' - 1 < t_0$ , as desired.
F.3 Middle Phase:  $y$  Keeps Decreasing Until Smaller Than  $|x_{t}|$
Then we move on to the second claim of Theorem 16. We define all rounds until  $y_{t} < |x_{t}|$  that are after Initial Phase as Middle Phase. Formally, we have the following definition.
Definition 25 (Middle Phase). Let  $t_2$  be the first time that  $y_t < |x_t|$ . We call  $(t_1, t_2]$  Middle Phase.
Before presenting the main conclusion of Middle Phase, we first define a threshold  $B$  that we measure whether a variable is "small enough", which we will use throughout this section. Formally,
$$
B \triangleq \max  \left\{\frac {2}{c} \sqrt {\frac {\eta}{C \gamma}}, \eta \rho + \sqrt {C \eta \gamma} \right\}.
$$
Theorem 26 (Main Conclusion of Middle Phase; SAM Case). Under the conditions of Theorem 16,  $|y_t| \leq \sqrt{C\gamma / \eta}$  and  $|x_t| \leq B$  throughout Middle Phase. Under the same conditions, we further have  $y_{t + 1} \leq y_t - \min \left\{\frac{1}{2}\eta \rho^2 y_t, \frac{c}{2\sqrt{2}}\eta \rho \right\}$  for all  $t_1 \leq t \leq t_2$  showing the second claim of Theorem 16.
Remark. This theorem can be understood as follows. Upon entering Middle Phase,  $|x_{t}|$  is bounded by  $\mathcal{O}(\eta \rho + \sqrt{\eta})$ . This then gets preserved throughout Middle Phase. Meanwhile, in  $t_2$  iterations,  $y_{t}$  drops rapidly such that  $y_{t_2 + 1} \leq |x_{t_2 + 1}| = \mathcal{O}(\eta \rho + \sqrt{\eta})$ . In other words,  $x_{t}$  and  $y_{t}$  are both "close enough" to 0 after Middle Phase; thus, SAM finds the flattest minimum (which is the origin).
Proof. The claims are shown in Lemma 27, Lemma 28 and Lemma 30, respectively.
Lemma 27. If  $\eta \rho + \sqrt{C\gamma\eta} < \sqrt{C\gamma / \eta}$ , then during the Middle Phase,  $|y_t| \leq y_0 \leq \sqrt{C\gamma / \eta}$ .
Proof. When  $\operatorname{sign}(y_{t+1}) = \operatorname{sign}(y_t)$ , then Lemma 18 implies that  $|y_{t+1}| \leq |y_t|$ . Now consider the case  $\operatorname{sign}(y_{t+1}) = -\operatorname{sign}(y_t)$ . For simplicity, say  $y_t > 0$  and  $y_{t+1} < 0$ . Since  $y_t > 0$ , we have  $\operatorname{sign}(\ell_t') = \operatorname{sign}(x_t y_t) = \operatorname{sign}(x_t)$  and hence Equation 13 gives
$$
\begin{array}{l} y _ {t + 1} = y _ {t} - \eta | \ell_ {t} ^ {\prime} | \rho \left(x _ {t} ^ {2} + y _ {t} ^ {2}\right) ^ {- 1 / 2} | y _ {t} | - \eta | \ell_ {t} ^ {\prime} | | x _ {t} | \\ \geq - \eta \rho - \eta | x _ {t} | \geq - \eta \rho - \eta y _ {t} \\ \geq - \eta \rho - \eta \sqrt {\frac {C \gamma}{\eta}} = - \eta \rho - \sqrt {C \gamma \eta}, \\ \end{array}
$$
which shows that  $|y_{t + 1}| \leq \eta \rho + \sqrt{C\gamma\eta}$ . This proves the statement.
Lemma 28. Suppose that  $|x_{t}| \leq B = \max \left\{\frac{2}{c}\sqrt{\frac{\eta}{C\gamma}}, \eta \rho + \sqrt{C\eta\gamma}\right\}$ . Then we have  $|x_{t+1}| \leq B$ .
Proof. By symmetry, we only consider the case where  $x_{t} > 0$ . If  $x_{t + 1} \geq 0$ , then we have  $0 \leq x_{t + 1} \leq x_{t}$  due to Lemma 24. If  $x_{t + 1} < 0$ , then since  $x_{t} > 0$ , it follows that
$$
\begin{array}{l} x _ {t + 1} = \operatorname {s i g n} (x _ {t}) x _ {t + 1} = | x _ {t} | - \eta | \ell_ {t} ^ {\prime} | \rho \left(x _ {t} ^ {2} + y _ {t} ^ {2}\right) ^ {- 1 / 2} | x _ {t} | - \eta | \ell_ {t} ^ {\prime} | | y _ {t} | \\ \geq \left| x _ {t} \right| - \eta \rho - \eta \sqrt {C \gamma / \eta} \geq - \eta \rho - \sqrt {C \eta \gamma}, \tag {15} \\ \end{array}
$$
where the last inequality is because  $|y_{t}| \leq y_{0} \leq \sqrt{C\gamma / \eta}$ . This concludes the proof.
Corollary 29. Note that, by definition of Initial Phase in Definition 19, we already have  $|x_{t_1 + 1}| \leq B$ . Hence, this lemma essentially says  $|x_t| \leq B = \max \left\{\frac{2}{c}\sqrt{\frac{\eta}{C\gamma}}, \eta \rho + \sqrt{C\eta\gamma}\right\}$  for all  $t_1 < t \leq t_2$ .
Lemma 30. If  $y_{t} \geq |x_{t}|$ , then we must have
$$
y _ {t + 1} \leq y _ {t} - \min  \left\{\frac {1}{2} \eta \rho^ {2} y _ {t}, \frac {c}{2 \sqrt {2}} \eta \rho \right\}.
$$
Proof. Since  $y_{t} \geq |x_{t}|$ , we have  $(x_{t}^{2} + y_{t}^{2})^{-1 / 2} \geq (2y_{t}^{2})^{-1 / 2} = (\sqrt{2} y_{t})^{-1}$ . Consider two cases:
i) if  $|(x_{t} + \rho y_{t}z_{t}^{-1})(y_{t} + \rho x_{t}z_{t}^{-1})| \geq c$ , then  $|\ell_t^{\prime}| \in [\frac{c}{2}, 1]$  according to Assumption (A3). Moreover, since  $\mathrm{sign}(\ell_t^\prime) = \mathrm{sign}(x_ty_t) = \mathrm{sign}(x_t)$ , it follows that
$$
\begin{array}{l} y _ {t + 1} = y _ {t} - \eta | \ell_ {t} ^ {\prime} | \cdot \rho \left(x _ {t} ^ {2} + y _ {t} ^ {2}\right) ^ {- 1 / 2} y _ {t} - \eta \ell_ {t} ^ {\prime} \cdot x _ {t} \\ = \left(1 - \eta | \ell_ {t} ^ {\prime} | \rho \left(x _ {t} ^ {2} + y _ {t} ^ {2}\right) ^ {- 1 / 2}\right) y _ {t} - \eta | \ell_ {t} ^ {\prime} | | x _ {t} | \\ \leq \left(1 - \eta c _ {t} \rho (\sqrt {2} y _ {t}) ^ {- 1}\right) y _ {t} \\ \leq y _ {t} - \eta \rho \frac {c}{2} \frac {1}{\sqrt {2}} = y _ {t} - \Omega (\eta \rho). \tag {16} \\ \end{array}
$$
ii) otherwise, it follows from (A3) that
$$
\left| \ell_ {t} ^ {\prime} \right| \geq \frac {1}{2} \left| \left(x _ {t} + \rho y _ {t} z _ {t} ^ {- 1}\right) \left(y _ {t} + \rho x _ {t} z _ {t} ^ {- 1}\right) \right|.
$$
As  $\mathrm{sign}(x_t) = \mathrm{sign}(\rho y_t z_t^{-1})$  and  $\mathrm{sign}(y_t) = \mathrm{sign}(\rho x_t z_t^{-1})$ , it follows that
$$
\begin{array}{l} | \ell_ {t} ^ {\prime} | \cdot \rho z _ {t} ^ {- 1} y _ {t} \geq \frac {1}{2} \left| \left(x _ {t} + \rho y _ {t} z _ {t} ^ {- 1}\right) \left(y _ {t} + \rho x _ {t} z _ {t} ^ {- 1}\right) \right| \cdot \rho z _ {t} ^ {- 1} y _ {t} \\ \geq \frac {1}{2} (\rho y _ {t} z _ {t} ^ {- 1}) ^ {2} y _ {t} \\ \end{array}
$$
$$
y _ {t + 1} \leq y _ {t} - \frac {1}{2} \eta \left(\rho y _ {t} z _ {t} ^ {- 1}\right) ^ {2} \left(y _ {t}\right) \leq y _ {t} - \frac {1}{2} \eta \rho^ {2} y _ {t} = (1 - \Omega \left(\eta \rho^ {2}\right)) y _ {t}, \tag {17}
$$
where we used  $z_{t}^{-1} \geq (\sqrt{2} y_{t})^{-1}$  and thus  $y_{t}z_{t}^{-1} \geq 1 / \sqrt{2}$ .
Combining item 16 and item 17, we obtain the desired conclusion.
F.4 Final Phase: Both  $x_{t}$  and  $y_{t}$  Oscillates Around the Origin
It only remains to show that the iteration never escapes the origin.
Definition 31 (Final Phase). We denote by "Final Phase" all iterations after  $t_2$ .
Theorem 32 (Main Conclusion of Final Phase; SAM Case). For all  $t > t_2$ , we have  $|x_t|, |y_t| \leq B$  where  $B = \max \left\{\frac{2}{c} \sqrt{\frac{\eta}{C\gamma}}, \eta \rho + \sqrt{C\eta\gamma}\right\}$ .
Remark. As we will see shortly, when entering Final Phase, both  $|x_{t}|$  and  $|y_{t}|$  are bounded by  $\mathcal{O}(\eta \rho + \sqrt{\eta})$ . This theorem essentially says that they can never exceed this bound in Final Phase. In other words, in Final Phase, both  $x_{t}$  and  $y_{t}$  are oscillating around 0.
Proof. We first show the following lemma, ruling out the possibility of  $|y_t| > B$  after Middle Phase:
Lemma 33. If  $\eta \rho +\eta B\leq B,5$  then  $|y_{t_2 + 1}|\leq B = \max \Bigl \{\frac{2}{c}\sqrt{\frac{\eta}{C\gamma}},\eta \rho +\sqrt{C\eta\gamma}\Bigr \} .$
Proof. The proof will be very similar to Lemma 28. Recall that by Definition 25, we must have  $y_{t_2} > 0$ . If  $y_{t_2 + 1} > 0$  as well, then  $|y_{t_2 + 1}| = y_{t_2 + 1} \leq |x_{t_2 + 1}| \leq B$ . Otherwise,
$$
y _ {t _ {2} + 1} = \operatorname {s i g n} \left(y _ {t _ {2}}\right) y _ {t _ {2} + 1} = \left| y _ {t _ {2}} \right| - \eta \left| \ell_ {t _ {2}} ^ {\prime} \right| \rho \left(x _ {t _ {2}} ^ {2} + y _ {t _ {2}} ^ {2}\right) ^ {- 1 / 2} \left| y _ {t _ {2}} \right| - \eta \left| \ell_ {t _ {2}} ^ {\prime} \right| \left| x _ {t _ {2}} \right| \geq - \eta \rho - \eta B,
$$
where we used  $|x_{t_2}| \leq B$ . We proved our claim as  $\eta \rho + \eta B \leq B$ .
According to Lemma 28 Lemma 33, we have
$$
\left| x _ {t _ {2} + 1} \right|, \left| y _ {t _ {2} + 1} \right| \leq B = \max  \left\{\frac {2}{c} \sqrt {\frac {\eta}{C \gamma}}, \eta \rho + \sqrt {C \eta \gamma} \right\}. \tag {18}
$$
Hence, it only remains to do a reduction, stated as follows.
Lemma 34. If  $\eta \rho +\eta B\leq B$  , then  $|x_{t + 1}|,|y_{t + 1}|\leq B$  as long as  $|x_{t}|,|y_{t}|\leq B$
Proof. Without loss of generality, let  $x_{t} > 0$  for some  $t > t_2$  and show that  $|x_{t + 1}| \leq B -$  which is identical to the proof of Lemma 27. By symmetry, the same also holds for  $|x_{t}|$ .
Our conclusion thus follows from an induction based on Equation 18 and Lemma 34.
G Omitted Proof of USAM Over Single-Neuron Linear Networks
Theorem 35 (Formal Version of Theorem 5). For a loss  $\mathcal{L}$  over  $(x,y)\in \mathbb{R}^2$  defined as  $\mathcal{L}(x,y) = \ell(xy)$  where  $\ell$  satisfies Assumptions (A1), (A2), and (A3), if the initialization  $(x_0,y_0)$  satisfies:
$$
\eta \leq \frac {1}{2}, \quad \eta \rho \leq \min  \left\{\frac {1}{2}, C ^ {- 1} \gamma^ {- 1} \right\}, \quad 2 \sqrt {C} \eta^ {- 1} = \mathcal {O} (\min  \left\{\eta^ {- 1} \rho^ {- 1}, \eta^ {- 2} \right\}), \quad C \gamma \left(1 + \rho C \frac {\gamma}{\eta}\right) \geq 1 6,
$$
we can characterize the initial and final phases of the trajectory of USAM (defined in Equation 2) by the following two theorems, whose main conclusions are informally summarized below:
1. (Theorem 39) Until  $x_{t} = \mathcal{O}(\sqrt{\gamma\eta})$ , we must have  $y_{t} = \Omega (\sqrt{\gamma / \eta})$ , and  $x_{t + 1}\leq x_t - \Omega (\sqrt{\gamma\eta})$ .
2. (Theorem 44) Once  $(1 + \rho y_t^2)y_t^2 \lesssim 2/\eta$ ,  $|x_t|$  decays exponentially and thus USAM gets stuck.
Different from Theorem 16, there is no characterization of Middle Phase here, which means the iterates can also stop above the threshold  $\tilde{y}_{\mathrm{USAM}}^2$  defined in Equation 7 (the technical reason is sketched in subsubsection 3.1.3, i.e., SAM gradients non-negligible when  $y_{t}$  is large, while USAM gradients vanish once  $|x_{t}|$  is small). However, we remark that the main takeaway of Theorem 35 is to contrast SAM (which always attains  $y_{\infty}^{2} \approx 0$ ) with USAM (which must stop once  $(1 + \rho y_{t}^{2})y_{t}^{2} \lesssim 2 / \eta$ ).
G.1 Basic Properties and Notations
Recall the update rule of USAM in Equation 2:  $w_{t + 1} \gets w_t - \eta \nabla \mathcal{L}(w_t + \rho \nabla \mathcal{L}(w_t))$ . Still writing  $w_t$  as  $[x_t, y_t]^\top$ , we have
$$
w _ {t + 1} = \left[ \begin{array}{c} x _ {t} \\ y _ {t} \end{array} \right] - \eta \nabla \mathcal {L} \left(\left[ \begin{array}{c} x _ {t} \\ y _ {t} \end{array} \right] + \rho \ell^ {\prime} (x _ {t} y _ {t}) \left[ \begin{array}{c} y _ {t} \\ x _ {t} \end{array} \right]\right)
$$
Recall that  $B = \max \left\{\frac{2}{c}\sqrt{\frac{\eta}{C\gamma}},\eta \rho +\sqrt{C\eta\gamma}\right\}$ , we only need to ensure that  $\eta B\leq \sqrt{C\eta\gamma}$ , which can be done by  $\eta \leq \min \{cC\gamma ,\frac{1}{4} C\gamma \rho^{-2},\frac{1}{2}\}$ . As the RHS is of order  $\Omega (1)$ ,  $\eta = \mathcal{O}(1)$  again suffices.
Similar to Theorem 16, these conditions are satisfied once  $\eta = o(1)$  and  $\rho = \mathcal{O}(1)$ .
$$
= \left[ \begin{array}{l} x _ {t} \\ y _ {t} \end{array} \right] - \eta \ell^ {\prime} \big ((x _ {t} + \rho \ell^ {\prime} (x _ {t} y _ {t}) y _ {t}) (y _ {t} + \rho \ell^ {\prime} (x _ {t} y _ {t}) x _ {t})) \left[ \begin{array}{l} y _ {t} + \rho \ell^ {\prime} (x _ {t} y _ {t}) x _ {t} \\ x _ {t} + \rho \ell^ {\prime} (x _ {t} y _ {t}) y _ {t} \end{array} \right]. \tag {19}
$$
Due to the removal of normalization,  $\ell^{\prime}$  are taken twice at different points in Equation 19. For simplicity, we denote  $\tilde{\ell}_t^\prime = \ell^\prime (x_ty_t)$  and  $\ell_t^\prime = \ell^\prime ((x_t + \rho \tilde{\ell}_t^\prime y_t)(y_t + \rho \tilde{\ell}_t^\prime x_t))$ . The update rule can be rewritten as:
$$
x _ {t + 1} = x _ {t} - \eta \ell_ {t} ^ {\prime} \cdot \left(y _ {t} + \rho \tilde {\ell} _ {t} ^ {\prime} x _ {t}\right), \quad y _ {t + 1} = y _ {t} - \eta \ell_ {t} ^ {\prime} \cdot \left(x _ {t} + \rho \tilde {\ell} _ {t} ^ {\prime} y _ {t}\right). \tag {20}
$$
Similar to the SAM case, we shall approximate  $\tilde{\ell}_t^\prime$  and  $\ell_t^\prime$  according to the magnitude of  $x_{t}y_{t}$  and  $(x_{t} + \rho \tilde{\ell}_{t}^{\prime}y_{t})(y_{t} + \rho \tilde{\ell}_{t}^{\prime}x_{t})$ . We first have the following lemma similar to Lemma 17:
Lemma 36 (Sign of Gradient in USAM). If  $x_{t} \neq 0, y_{t} \neq 0$ , then  $\mathrm{sign}(x_{t}) = \mathrm{sign}(x_{t} + \rho \tilde{\ell}_{t}'y_{t})$  and  $\mathrm{sign}(y_{t}) = \mathrm{sign}(y_{t} + \rho \tilde{\ell}_{t}'x_{t})$ . In particular, if  $x_{t} \neq 0, y_{t} \neq 0$ , then  $\mathrm{sign}(\ell_t') = \mathrm{sign}(\tilde{\ell}_t') = \mathrm{sign}(x_ty_t)$ .
Proof. First of all,  $\mathrm{sign}(\tilde{\ell}_t^{\prime}) = \mathrm{sign}(x_t y_t)$  according to Assumption (A1). Therefore,
$$
\operatorname {s i g n} \left(x _ {t} + \rho \bar {\ell} _ {t} ^ {\prime} y _ {t}\right) = \operatorname {s i g n} \left(x _ {t} + \rho \operatorname {s i g n} \left(x _ {t} y _ {t}\right) y _ {t}\right) = \operatorname {s i g n} \left(x _ {t}\right),
$$
$$
\operatorname {s i g n} \left(y _ {t} + \rho \tilde {\ell} _ {t} ^ {\prime} x _ {t}\right) = \operatorname {s i g n} \left(y _ {t} + \rho \operatorname {s i g n} \left(x _ {t} y _ {t}\right) x _ {t}\right) = \operatorname {s i g n} \left(y _ {t}\right),
$$
giving our first two claims. The last conclusion follows by definition.
Moreover, we also have the following lemma analog to Lemma 18:
Lemma 37. Suppose that  $x_{t}, y_{t} \geq 0$ . Then, the following basic properties hold:
Proof. We only prove the statement for  $x_{t} > 0$  as the proof is similar for other cases. Recall the dynamics of  $x_{t}$ :
$$
x _ {t + 1} = x _ {t} - \eta \ell_ {t} ^ {\prime} \cdot \left(y _ {t} + \rho \tilde {\ell} _ {t} ^ {\prime} x _ {t}\right).
$$
Since  $x_{t} \neq 0$ , Lemma 36 implies that  $\mathrm{sign}(\ell_t^{\prime}) = \mathrm{sign}(\tilde{\ell}_t^{\prime}) = \mathrm{sign}(x_t y_t)$ , and so
$$
\begin{array}{l} x _ {t + 1} = x _ {t} - \eta \ell_ {t} ^ {\prime} \cdot \left(y _ {t} + \rho \tilde {\ell} _ {t} ^ {\prime} x _ {t}\right) \\ = x _ {t} - \eta \operatorname {s i g n} \left(x _ {t} y _ {t}\right) \left| \ell_ {t} ^ {\prime} \right| \cdot \left(y _ {t} + \rho \operatorname {s i g n} \left(x _ {t} y _ {t}\right) \mid \tilde {\ell} _ {t} ^ {\prime} \mid x _ {t}\right) \\ = x _ {t} - \eta \left| \ell_ {t} ^ {\prime} \right| \cdot \left(\left| y _ {t} \right| + \rho \right| \tilde {\ell_ {t} ^ {\prime}} \left| x _ {t} \right|) \leq x _ {t}, \\ \end{array}
$$
G.2 Initial Phase:  $x_{t}$  Decreases Fast while  $y_{t}$  Remains Large
The definition of Initial Phase is very similar to the one of SAM – and the conclusion is also analogue, although the proofs are slightly different because of the different update rule.
Definition 38 (Initial Phase). Let  $t_1$  be the largest time such that  $x_t > \frac{1}{2}\sqrt{\gamma\eta}$  for all  $t \leq t_1$ . We denote by "Initial Phase" the iterations  $[0, t_1]$ .
Theorem 39 (Main Conclusion of Initial Phase; USAM Case). For  $t_0 = \Theta(\min\{\eta^{-1}\rho^{-1}, \eta^{-2}\})$ , we have  $y_t \geq \frac{1}{2}\sqrt{\gamma/\eta}$  for all  $t \leq \min\{t_0, t_1\}$  under the conditions of Theorem 35. Moreover, we have  $x_{t+1} \leq x_t - \frac{1}{2}\sqrt{C\gamma\eta}$  for all  $t \leq \min\{t_0, t_1\}$ , which consequently infers  $t_1 \leq t_0$  under the conditions of Theorem 35, i.e.,  $\min\{t_0, t_1\}$  is just  $t_1$ . This shows the first claim of Theorem 35.
Proof. This theorem is a combination of Lemma 40 and Lemma 43.
Similar to Lemma 21,  $y_{t}$  in USAM also cannot be too small in the first few iterations:
Lemma 40. There exists  $t_0 = \Theta(\min\{\eta^{-1}\rho^{-1}, \eta^{-2}\})$  such that for  $t \leq t_0$ , we have  $|y_t| \geq \frac{1}{2}\sqrt{\gamma/\eta}$ . Assuming  $\eta \leq \frac{1}{2}$ , then  $y_t \geq \frac{1}{2}\sqrt{\gamma/\eta}$  for  $t \leq \min\{t_1, t_0\}$ .