SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathcal {L} (w _ {t}) - \mathcal {L} (w ^ {\star}) \leq \left(1 - \alpha \eta (2 - \eta \beta)\right) ^ {t} (\mathcal {L} (w _ {0}) - \mathcal {L} (w ^ {\star})) + \frac {\eta \beta^ {3} \rho^ {2}}{2 \alpha (2 - \eta \beta)}.
$$
Proof. We first claim the following analog of descent lemma, which we state as Lemma 9:
$$
\mathcal {L} \left(w _ {t + 1}\right) \leq \mathcal {L} \left(w _ {t}\right) - \frac {1}{2} \eta (2 - \eta \beta) \| \nabla \mathcal {L} \left(w _ {t}\right) \| ^ {2} + \frac {\eta^ {2} \beta^ {3} \rho^ {2}}{2}.
$$
By definition of strong convexity, we have
$$
\mathcal {L} (w _ {t}) - \mathcal {L} (w ^ {\star}) \leq \langle \nabla \mathcal {L} (w _ {t}), w _ {t} - w ^ {\star} \rangle - \frac {\alpha}{2} \| w _ {t} - w ^ {\star} \| ^ {2} \leq \frac {1}{2 \alpha} \| \nabla f (w _ {t}) \| ^ {2},
$$
where the last inequality uses the fact that  $\langle a,b\rangle -\frac{1}{2}\| b\|^2\leq \frac{1}{2}\| a\|^2$ . Thus, combining the two inequalities above, we obtain
$$
\mathcal {L} (w _ {t + 1}) \leq \mathcal {L} (w _ {t}) - \alpha \eta (2 - \eta \beta) (\mathcal {L} (w _ {t}) - \mathcal {L} (w ^ {\star})) + \frac {\eta^ {2} \beta^ {3} \rho^ {2}}{2},
$$
which after rearrangement becomes
$$
\mathcal {L} (w _ {t + 1}) - \mathcal {L} (w ^ {\star}) \leq (1 - \alpha \eta (2 - \eta \beta)) (\mathcal {L} (w _ {t}) - \mathcal {L} (w ^ {\star})) + \frac {\eta^ {2} \beta^ {3} \rho^ {2}}{2}.
$$
Unrolling this recursion, we obtain
$$
\mathcal {L} (w _ {t}) - \mathcal {L} (w ^ {\star}) \leq (1 - \alpha \eta (2 - \eta \beta)) ^ {t} (\mathcal {L} (w _ {0}) - \mathcal {L} (w ^ {\star})) + \frac {\eta^ {2} \beta^ {3} \rho^ {2}}{2} \sum_ {k = 0} ^ {t - 1} (1 - \alpha \eta (2 - \eta \beta)) ^ {k}
$$
$$
\begin{array}{l} \leq \left(1 - \alpha \eta (2 - \eta \beta)\right) ^ {t} (\mathcal {L} (w _ {0}) - \mathcal {L} (w ^ {\star})) + \frac {\eta^ {2} \beta^ {3} \rho^ {2}}{2} \sum_ {k = 0} ^ {\infty} (1 - \alpha \eta (2 - \eta \beta)) ^ {k} \\ = (1 - \alpha \eta (2 - \eta \beta)) ^ {t} (\mathcal {L} (w _ {0}) - \mathcal {L} (w ^ {\star})) + \frac {\eta^ {2} \beta^ {3} \rho^ {2}}{2 \alpha \eta (2 - \eta \beta)}. \\ \end{array}
$$
Lemma 9 (SAM Descent lemma). For a convex loss  $\mathcal{L}$  that is  $\beta$ -smooth, SAM iterates  $w_{t}$  satisfy the following when the learning rate  $\eta < \frac{2}{\beta}$  and  $\rho \geq 0$ :
$$
\mathcal {L} (w _ {t + 1}) \leq \mathcal {L} (w _ {t}) - \frac {1}{2} \eta (2 - \eta \beta) \| \nabla \mathcal {L} (w _ {t}) \| ^ {2} + \frac {\eta^ {2} \beta^ {3} \rho^ {2}}{2}, \quad \forall t.
$$
Proof. Let  $v_{t} \triangleq \nabla \mathcal{L}(w_{t}) / \| \nabla \mathcal{L}(w_{t}) \|$  and  $w_{t+1/2} \triangleq w_{t} + \rho v_{t}$  so we have  $w_{t+1} = w_{t} - \eta \nabla \mathcal{L}(w_{t+1/2})$ . Then the  $\beta$ -smoothness of  $\mathcal{L}$  yields
$$
\mathcal {L} \left(w _ {t + 1}\right) \leq \mathcal {L} \left(w _ {t}\right) - \eta \langle \nabla \mathcal {L} \left(w _ {t}\right), \nabla \mathcal {L} \left(w _ {t + 1 / 2}\right) \rangle + \frac {\eta^ {2} \beta}{2} \| \nabla \mathcal {L} \left(w _ {t + 1 / 2}\right) \| ^ {2}.
$$
We start with upper bounding the norm of  $\nabla \mathcal{L}(w_{t + 1 / 2})$ :
$$
\begin{array}{l} \left\| \nabla \mathcal {L} \left(w _ {t + 1 / 2}\right) \right\| ^ {2} = \left\| \nabla \mathcal {L} \left(w _ {t + 1 / 2}\right) - \nabla \mathcal {L} \left(w _ {t}\right) \right\| ^ {2} - \left\| \nabla \mathcal {L} \left(w _ {t}\right) \right\| ^ {2} + 2 \langle \nabla \mathcal {L} \left(w _ {t + 1 / 2}\right), \nabla \mathcal {L} \left(w _ {t}\right) \rangle \\ \leq \beta^ {2} \rho^ {2} - \| \nabla \mathcal {L} (w _ {t}) \| ^ {2} + 2 \langle \nabla \mathcal {L} (w _ {t + 1 / 2}), \nabla \mathcal {L} (w _ {t}) \rangle , \\ \end{array}
$$
Hence, as long as  $\eta < \frac{2}{\beta}$ , we have the following upper bound on  $\mathcal{L}(w_{t + 1})$ :
$$
\begin{array}{l} \mathcal {L} (w _ {t + 1}) \leq \mathcal {L} (w _ {t}) - \eta \langle \nabla \mathcal {L} (w _ {t + 1 / 2}), \nabla \mathcal {L} (w _ {t}) \rangle + \frac {\eta^ {2} \beta}{2} \| \nabla \mathcal {L} (w _ {t + 1 / 2}) \| ^ {2} \\ \leq \mathcal {L} (w _ {t}) - \eta \langle \nabla \mathcal {L} (w _ {t + 1 / 2}), \nabla \mathcal {L} (w _ {t}) \rangle + \frac {\eta^ {2} \beta}{2} \left(\beta \rho^ {2} - \| \nabla \mathcal {L} (w _ {t}) \| ^ {2} + 2 \langle \nabla \mathcal {L} (w _ {t + 1 / 2}), \nabla \mathcal {L} (w _ {t}) \rangle\right) \\ = \mathcal {L} (w _ {t}) - \frac {\eta^ {2} \beta}{2} \| \nabla \mathcal {L} (w _ {t}) \| ^ {2} - (\eta - \eta^ {2} \beta) \langle \nabla \mathcal {L} (w _ {t + 1 / 2}), \nabla \mathcal {L} (w _ {t}) \rangle + \frac {\eta^ {2} \beta^ {3} \rho^ {2}}{2}. \\ \end{array}
$$
Now we lower bound  $\langle \nabla \mathcal{L}(w_{t + 1 / 2}),\nabla \mathcal{L}(w_t)\rangle$  . Note that
$$
\begin{array}{l} \langle \nabla \mathcal {L} (w _ {t} + \rho v _ {t}), \nabla \mathcal {L} (w _ {t}) \rangle = \langle \nabla \mathcal {L} (w _ {t} + \rho v _ {t}) - \nabla \mathcal {L} (w _ {t}), \nabla \mathcal {L} (w _ {t}) \rangle + \| \nabla \mathcal {L} (w _ {t}) \| ^ {2} \\ = \frac {\left\| \nabla \mathcal {L} \left(w _ {t}\right) \right\|}{\rho} \left\langle \nabla \mathcal {L} \left(w _ {t} + \rho v _ {t}\right) - \nabla \mathcal {L} \left(w _ {t}\right), \rho v _ {t} \right\rangle + \left\| \nabla \mathcal {L} \left(w _ {t}\right) \right\| ^ {2} \geq \left\| \nabla \mathcal {L} \left(w _ {t}\right) \right\| ^ {2}, \\ \end{array}
$$
where the last inequality uses the following standard fact about convex functions: for any  $w_{1},w_{2}$ ,  $\langle \nabla \mathcal{L}(w_1) - \nabla \mathcal{L}(w_2),w_1 - w_2\rangle \geq 0$ . Hence, we arrive at
$$
\begin{array}{l} \mathcal {L} \left(w _ {t + 1}\right) \leq \mathcal {L} \left(w _ {t}\right) - \frac {\eta^ {2} \beta}{2} \left\| \nabla \mathcal {L} \left(w _ {t}\right) \right\| ^ {2} - \left(\eta - \eta^ {2} \beta\right) \left\langle \nabla \mathcal {L} \left(w _ {t + 1 / 2}\right), \nabla \mathcal {L} \left(w _ {t}\right) \right\rangle + \frac {\eta^ {2} \beta^ {3} \rho^ {2}}{2} \\ \leq \mathcal {L} (w _ {t}) - \frac {\eta^ {2} \beta}{2} \| \nabla \mathcal {L} (w _ {t}) \| ^ {2} - (\eta - \eta^ {2} \beta) \| \nabla \mathcal {L} (w _ {t}) \| ^ {2} + \frac {\eta^ {2} \beta^ {3} \rho^ {2}}{2} \\ \leq \mathcal {L} (w _ {t}) - \frac {1}{2} \eta (2 - \eta \beta) \| \nabla \mathcal {L} (w _ {t}) \| ^ {2} + \frac {\eta^ {2} \beta^ {3} \rho^ {2}}{2}, \\ \end{array}
$$
and thus finishing the proof.
C.2 USAM Diverges on Quadratic Losses
Theorem 10 (USAM over Quadratic Losses). Following (Bartlett et al., 2022), consider the quadratic loss  $\mathcal{L}$  induced by a PSD matrix  $\Lambda$ . Without loss of generality, let  $\mathcal{L}$  minimize at the origin. Formally,
$$
\mathcal {L} (w) = \frac {1}{2} w ^ {\top} \Lambda w, \quad \text {w h e r e} \Lambda = \operatorname {d i a g} \left(\lambda_ {1}, \dots , \lambda_ {d}\right), \tag {8}
$$
where  $\lambda_{\mathrm{max}} = \lambda_1 > \dots \geq \lambda_d > 0$  and  $v_{\mathrm{max}} = e_1, e_2, \ldots, e_d$  are the eigenvectors corresponding to  $\lambda_1, \lambda_2, \ldots, \lambda_d$ , respectively. Then the iterates of USAM Equation 2 applied to Equation 8 satisfy:
1. If  $\eta (\lambda_{\mathrm{max}} + \rho \lambda_{\mathrm{max}}^2) < 2$ , then the iterates converge to the global minima exponentially fast.
2. If  $\eta (\lambda_{\mathrm{max}} + \rho \lambda_{\mathrm{max}}^2) > 2$  and if  $\langle w_0,v_{\mathrm{max}}\rangle \neq 0$ , then the iterates diverge.
Moreover, such a loss function  $\mathcal{L}$  is  $\lambda_{1}$ -smooth and  $\lambda_{d}$ -strongly-convex.
Proof. As  $\nabla \ell (w) = \Lambda w$ , the USAM update Equation 2 reads
$$
w _ {t + 1} = w _ {t} - \eta \nabla \mathcal {L} (w _ {t} + \rho \nabla \mathcal {L} (w _ {t})) = \left(I - \eta \Lambda - \eta \rho \Lambda^ {2}\right) w _ {t}.
$$
Hence, if  $\eta (\lambda_{\mathrm{max}} + \rho \lambda_{\mathrm{max}}^2) < 2$ , then since  $\| I - \eta \Lambda - \eta \rho \Lambda^2 \| < 1$ , we have that  $\| w_t \| \to 0$  exponentially fast. On the other hand, if  $\eta (\lambda_{\mathrm{max}} + \rho \lambda_{\mathrm{max}}^2) > 2$ , then  $|1 - \eta \lambda_1 - \eta \lambda_1^2| > 1$ . Since  $\| w_t \| \geq |1 - \eta \lambda_1 - \eta \lambda_1^2|^t |\langle w_0, e_1 \rangle|$ , it follows that the iterates diverge as long as  $\langle w_0, v_{\mathrm{max}} \rangle \neq 0$ . Finally, as  $\nabla^2 \mathcal{L}(w) = \Lambda$ , we directly know that  $\mathcal{L}$  is  $\lambda_1$ -smooth and  $\lambda_d$ -strongly-convex.
D Omitted Proof for Scalar Factorization Problems
Theorem 11 (Formal Version of Theorem 2). Consider the scalar factorization loss  $\mathcal{L}(x,y) = \frac{1}{2} (xy)^2$ . If  $(x_0,y_0)$  satisfies  $2x_0 > y_0 > x_0 \gg 0$  and  $\eta = (x_0^2 + y_0^2)^{-1}$ , then i) SAM finds a neighborhood of the origin with radius  $\mathcal{O}(\rho)$  for all  $\rho$ , and ii) USAM diverges as long as  $\rho \geq 15\eta$ .
This theorem is a combination of Theorem 12 and Theorem 15 that we will show shortly.
D.1 SAM Always Converges on Scalar Factorization Problems
Let  $w_{t} = [x_{t} \quad y_{t}]^{\top}$  be the  $t$ -th iterate. Denote  $\nabla_{t} = \nabla \mathcal{L}(w_{t}) = [x_{t}y_{t}^{2} \quad x_{t}^{2}y_{t}]^{\top}$ . For each step  $t$ , the actual update of SAM is therefore the gradient taken at
$$
\tilde {w} _ {t} = w _ {t} + \rho \frac {\nabla_ {t}}{\| \nabla_ {t} \|} = \left[ \begin{array}{c} x _ {t} + \rho y _ {t} z _ {t} ^ {- 1} \\ y _ {t} + \rho x _ {t} z _ {t} ^ {- 1} \end{array} \right], \quad \text {w h e r e} z _ {t} ^ {- 1} = \frac {\operatorname {s i g n} (x _ {t} y _ {t})}{\sqrt {x _ {t} ^ {2} + y _ {t} ^ {2}}}.
$$
By denoting  $\tilde{\nabla}_t = \nabla \mathcal{L}(\tilde{w}_t)$ , the update rule of SAM is
$$
\left[ \begin{array}{c} x _ {t + 1} \\ y _ {t + 1} \end{array} \right] = w _ {t + 1} = w _ {t} - \eta \tilde {\nabla} _ {t} = \left[ \begin{array}{c} x _ {t} - \eta \big (x _ {t} + \rho y _ {t} z _ {t} ^ {- 1} \big) \big (y _ {t} + \rho x _ {t} z _ {t} ^ {- 1} \big) ^ {2} \\ y _ {t} - \eta \big (x _ {t} + \rho y _ {t} z _ {t} ^ {- 1} \big) ^ {2} \big (y _ {t} + \rho x _ {t} z _ {t} ^ {- 1} \big) \end{array} \right].
$$
We make the following claim:
Theorem 12 (SAM over Scalar Factorization Problems). Under the setting of Theorem 11, there exists some threshold  $T$  for SAM such that  $|x_{t}|, |y_{t}| \leq 5\rho$  for all  $t \geq T$ .
Proof. We first observe that SAM over  $\mathcal{L}(x,y)$  always pushes the iterate towards the minima (which are all the points on the  $x$ - and  $y$ -axes). Formally:
- If  $x_{t} \geq 0$ , then  $x_{t + 1} \leq x_{t}$ . If  $x_{t} \leq 0$ , then  $x_{t + 1} \geq x_{t}$ .
- If  $y_{t} \geq 0$ , then  $y_{t + 1} \leq y_{t}$ . If  $y_{t} \leq 0$ , then  $y_{t + 1} \geq y_{t}$ .
This observation can be verified by noticing  $\mathrm{sign}(\rho y_t z_t^{-1}) = \mathrm{sign}(y_t) \mathrm{sign}(x_t y_t) = \mathrm{sign}(x_t)$  and similarly  $\mathrm{sign}(\rho x_t z_t^{-1}) = \mathrm{sign}(y_t)$ . In other words,  $\mathrm{sign}(x_t + \rho y_t z_t^{-1}) = \mathrm{sign}(x_t)$  and thus always pushes  $x_t$  towards the  $y$ -axis. The same also holds for  $y_t$  by symmetry.
In analog to the descent lemma for GD, we can show the following lemma:
Lemma 13. When picking  $\eta = (x_0^2 + y_0^2)^{-1}$ , we have  $\mathcal{L}(w_{t+1}) - \mathcal{L}(w_t) < 0$  as long as  $x_t, y_t \geq 5\rho$ .
Proof. As  $\nabla^2\mathcal{L}(x,y) = \left[ \begin{array}{cc}y^2 & 2xy\\ 2xy & x^2 \end{array} \right]$ , we know  $\mathcal{L}$  is  $\beta \triangleq (x_0^2 +y_0^2)$ -smooth inside the region  $\{(x,y):x^2 +y^2\leq \beta \}$ . Then we have (recall that  $\eta = (x_0^2 +y_0^2)^{-1} = 1 / \beta$ )
$$
\mathcal {L} \left(w _ {t + 1}\right) - \mathcal {L} \left(w _ {t}\right) \leq \langle \nabla \mathcal {L} \left(w _ {t}\right), w _ {t + 1} - w _ {t} \rangle + \frac {\beta}{2} \| w _ {t + 1} - w _ {t} \| ^ {2}
$$
$$
= - \eta \langle \nabla_ {t}, \tilde {\nabla} _ {t} \rangle + \frac {\beta}{2} \eta^ {2} \| \tilde {\nabla} _ {t} \| ^ {2} = - \eta \Big (\langle \nabla_ {t} - \frac {1}{2} \tilde {\nabla} _ {t}, \tilde {\nabla} _ {t} \rangle \Big).
$$
To make sure that it's negative, we simply want
$$
\begin{array}{l} 0 \leq \langle \nabla_ {t} - \frac {1}{2} \tilde {\nabla} _ {t}, \tilde {\nabla} _ {t} \rangle = \left(x _ {t} y _ {t} ^ {2} - \frac {1}{2} \big (x _ {t} + \rho y _ {t} z _ {t} ^ {- 1} \big) \big (y _ {t} + \rho x _ {t} z _ {t} ^ {- 1} \big) ^ {2}\right) \big (x _ {t} + \rho y _ {t} z _ {t} ^ {- 1} \big) \big (y _ {t} + \rho x _ {t} z _ {t} ^ {- 1} \big) ^ {2} + \\ \left(x _ {t} ^ {2} y _ {t} - \frac {1}{2} \left(x _ {t} + \rho y _ {t} z _ {t} ^ {- 1}\right) ^ {2} \left(y _ {t} + \rho x _ {t} z _ {t} ^ {- 1}\right)\right) \left(x _ {t} + \rho y _ {t} z _ {t} ^ {- 1}\right) ^ {2} \left(y _ {t} + \rho x _ {t} z _ {t} ^ {- 1}\right), \\ \end{array}
$$
which can be ensured once
$$
\left(x _ {t} + \rho y _ {t} z _ {t} ^ {- 1}\right) \left(y _ {t} + \rho x _ {t} z _ {t} ^ {- 1}\right) ^ {2} \leq 2 x _ {t} y _ {t} ^ {2}, \quad \left(x _ {t} + \rho y _ {t} z _ {t} ^ {- 1}\right) ^ {2} \left(y _ {t} + \rho x _ {t} z _ {t} ^ {- 1}\right) \leq 2 x _ {t} ^ {2} y _ {t}.
$$
If we have  $x_{t},y_{t}\geq 5\rho$ , then as  $z_{t}^{-1}\leq \min \{x_{t}^{-1},y_{t}^{-1}\}$ , we have
$$
\left(x _ {t} + \rho y _ {t} z _ {t} ^ {- 1}\right) \left(y _ {t} + \rho x _ {t} z _ {t} ^ {- 1}\right) ^ {2} \leq \left(x _ {t} + \rho\right) \left(y _ {t} + \rho\right) ^ {2} \leq 1. 2 ^ {3} x _ {t} y _ {t} ^ {2} <   2 x _ {t} y _ {t} ^ {2},
$$
which shows the first inequality. The second one follows from symmetry.
Therefore, SAM will progress until  $x_{t} \leq 5\rho$  or  $y_{t} \leq 5\rho$ . Without loss of generality, assume that  $x_{t} \leq 5\rho$ ; we then claim that  $y_{t}$  will soon decrease to  $\mathcal{O}(\rho)$ .
Lemma 14. Suppose that  $|x_{t}| \leq 5\rho$  but  $y_{t} \geq 5\rho$ . Then  $|x_{t + 1}| \leq 5\rho$  but  $y_{t + 1} \leq (1 - \frac{1}{2}\eta \rho^{2})y_{t}$ .
Proof. First, show that  $|x_{t}|$  remains bounded by  $5\rho$ . Assume  $x_{t} \geq 0$  without loss of generality:
$$
\begin{array}{l} x _ {t + 1} = x _ {t} - \eta \left(x _ {t} + \rho y _ {t} z _ {t} ^ {- 1}\right) \left(y _ {t} + \rho x _ {t} z _ {t} ^ {- 1}\right) ^ {2} \\ \geq x _ {t} - \eta (5 \rho + \rho) (y _ {t} + \rho) ^ {2} \\ \geq - \eta 6 \rho \left(\frac {4}{5 \eta} + \rho^ {2} + 2 \rho \sqrt {\frac {4}{5 \eta}}\right) \geq - 5 \rho , \tag {9} \\ \end{array}
$$
where the second last line uses  $y_{t}^{2} \leq y_{0}^{2} = \frac{y_{0}^{2}}{x_{0}^{2} + y_{0}^{2}}\eta^{-1} \leq \frac{4}{5}\eta^{-1}$  and the last one uses  $\eta = (x_0^2 + y_0^2)^{-1} \ll 1$ . Meanwhile, we see that  $y_{t}$  decreases exponentially fast by observing the following:
$$
\begin{array}{l} y _ {t + 1} = y _ {t} - \eta \left(y _ {t} + \rho x _ {t} z _ {t} ^ {- 1}\right) \left(x _ {t} + \rho y _ {t} z _ {t} ^ {- 1}\right) ^ {2} \\ \leq y _ {t} - \eta y _ {t} \left(\rho y _ {t} z _ {t} ^ {- 1}\right) ^ {2} \\ \leq \left(1 - \frac {1}{2} \eta \rho^ {2}\right) y _ {t}, \\ \end{array}
$$
where the last line uses  $z_{t}^{-1} = (x_{t}^{2} + y_{t}^{2})^{-1 / 2} \geq (2y_{t}^{2})^{-1 / 2} = 2^{-1 / 2}y_{t}^{-1}$  as  $y_{t} \geq 5\rho \geq |x_{t}|$ .
So eventually we have  $|x_{t}|, |y_{t}| \leq 5\rho$ . Recall that Equation 9 infers  $|x_{t + 1}| \leq 5\rho$  from  $|x_{t}| \leq 5\rho$ . Hence, by symmetry, we conclude that  $|x_{t + 1}|, |y_{t + 1}| \leq 5\rho$  hold as well. Therefore, SAM always finds an  $\mathcal{O}(\rho)$ -neighborhood of the origin, i.e., it is guaranteed to converge regardless of  $\rho$ .
D.2 USAM Diverges with Small  $\rho$
For USAM, the dynamics can be written as
$$
\left[ \begin{array}{l} x _ {t + 1} \\ y _ {t + 1} \end{array} \right] = \left[ \begin{array}{l} x _ {t} \\ y _ {t} \end{array} \right] - \eta \nabla \mathcal {L} \binom {x _ {t} + \rho x _ {t} y _ {t} ^ {2}} {y _ {t} + \rho x _ {t} ^ {2} y _ {t}} = \left[ \begin{array}{l} x _ {t} - \eta (x _ {t} + \rho x _ {t} y _ {t} ^ {2}) (y _ {t} + \rho x _ {t} ^ {2} y _ {t}) ^ {2} \\ y _ {t} - \eta (x _ {t} + \rho x _ {t} y _ {t} ^ {2}) ^ {2} (y _ {t} + \rho x _ {t} ^ {2} y _ {t}) \end{array} \right]. \tag {10}
$$
We make the following claim which is similar to Theorem 10:
Theorem 15 (USAM over Scalar Factorization Problems). Under the setup of Theorem 11, for any  $\rho \geq 15\eta$ ,  $|x_{t+1}| \geq 2|x_t|$  and  $|y_{t+1}| \geq 2|y_t|$  for all  $t \geq 1$ ; in other words, USAM diverges exponentially fast.
Proof. Prove by induction. From Equation 10, our conclusion follows once
$$
\eta \big | (x _ {t} + \rho x _ {t} y _ {t} ^ {2}) (y _ {t} + \rho x _ {t} ^ {2} y _ {t}) ^ {2} \big | \geq 3 | x _ {t} |, \quad \eta \big | (x _ {t} + \rho x _ {t} y _ {t} ^ {2}) ^ {2} (y _ {t} + \rho x _ {t} ^ {2} y _ {t}) \big | \geq 3 | y _ {t} |.
$$
According to our setup that  $\eta = (x_0^2 + y_0^2)^{-1}$ ,  $y_0 \leq 2x_0$ , and the induction statement that, we have  $\eta \geq (5x_t^2)^{-1}$ . The second inequality then holds as long as
$$
\left| \left(\rho x _ {t} y _ {t} ^ {2}\right) ^ {2} \left(\rho x _ {t} ^ {2} y _ {t}\right) \right| \geq 1 5 x _ {t} ^ {2} | y _ {t} |, \quad \text {i . e .}, \rho^ {3} x _ {t} ^ {2} y _ {t} ^ {4} \geq 1 5,
$$
which is true as  $x_{t}^{2} \geq x_{0}^{2}, y_{t}^{4} \geq y_{0}^{4}, y_{0} \geq x_{0}$ , and  $\rho \geq 15\eta \geq 3x_0^{-2}$ . Note that the bounds on  $\rho$  are very loose, and we made no effort to optimize it; instead, we only aimed to show that USAM starts to diverge from a  $\rho = \Theta (\eta) \ll 1$ .
E Assumptions in the Single-Neuron Linear Network Model
Assumptions on  $\ell$ . Following Ahn et al. (2023a), we make the following assumptions about  $\ell$ :
(A1)  $\ell$  is a convex, even, 1-Lipschitz function that is minimized at 0.
(A2)  $\ell$  is twice continuously differentiable near the origin with  $\ell''(0) = 1$ , which infers the existence of a constant  $c > 0$  such that  $|\ell'(s)| \leq |s|$  for all  $|s| \leq c$ .
(A3) We further assume a "linear tail" away from the minima, i.e.,  $|\ell'(s)| \geq c/2$  for all  $|s| \geq c$  and  $|\ell'(s)| \geq |s|/2$  for  $|s| \leq c$ .
Some concrete example of loss functions satisfying the above assumption include a symmetrized logistic loss  $\frac{1}{2}\log (1 + \exp (-2s)) + \frac{1}{2}\log (1 + \exp (2s))$  and a square root loss  $\sqrt{1 + s^2}$ . One may refer to their paper for more details.
F Omitted Proof of SAM Over Single-Neuron Linear Networks
Theorem 16 (Formal Version of Theorem 4). 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:
$$
y _ {0} \geq x _ {0} \gg 0, y _ {0} ^ {2} - x _ {0} ^ {2} = \frac {\gamma}{\eta}, y _ {0} ^ {2} = C \frac {\gamma}{\eta},
$$
where  $\gamma \in [\frac{1}{2}, 2]$  and  $C \geq 1$  is constant, then for all hyper-parameter configurations  $(\eta, \rho)$  such that
$$
\eta \rho + \sqrt {C \gamma \eta} \leq \min  \left\{\frac {1}{2}, C \right\} \sqrt {\frac {\gamma}{\eta}}, \quad \frac {4 \sqrt {C}}{c} \eta^ {- 1} = \mathcal {O} (\min  \{\eta^ {- 1. 5} \rho^ {- 1} \gamma^ {1 / 2}, \eta^ {- 2} \}),
$$
we can decompose the trajectory of SAM (defined in Equation 1) into three phases, whose main conclusions are stated separately in three theorems and are informally summarized here:
1. (Theorem 20) 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 26) After Initial Phase and until  $y_{t} \leq |x_{t}|, |x_{t}| = \mathcal{O}(\eta \rho + \sqrt{\eta})$  still holds. Meanwhile,  $y_{t+1} \leq y_{t} - \min \{\Omega(\eta \rho^{2})y_{t}, \Omega(\eta \rho)\}$  (i.e.,  $y_{t}$  either drops by  $\Omega(\eta \rho)$  or decays by  $\Omega(\eta \rho^{2})$ ).
3. (Theorem 32) After Middle Phase, we always have  $|x_{t}|, |y_{t}| = \mathcal{O}(\eta \rho + \sqrt{\eta})$ .
F.1 Basic Properties and Notations
Recall the update rule of SAM in Equation 1:  $w_{t + 1} \gets w_t - \eta \nabla \mathcal{L}(w_t + \rho \frac{\nabla \mathcal{L}(w_t)}{\|\nabla \mathcal{L}(w_t)\|})$ . By writing  $w_t$  as  $\begin{bmatrix} x_t \\ y_t \end{bmatrix}$ , substituting  $\mathcal{L}(x,y) = \ell(xy)$ , and utilizing the expressions of  $\nabla \mathcal{L}(x,y)$  in Equation 6, we have:
$$
w _ {t} + \rho \frac {\nabla \mathcal {L} \left(w _ {t}\right)}{\| \nabla \mathcal {L} \left(w _ {t}\right) \|} = \left[ \begin{array}{l} x _ {t} \\ y _ {t} \end{array} \right] + \rho \frac {\ell^ {\prime} \left(x _ {t} y _ {t}\right)}{\left| \ell^ {\prime} \left(x _ {t} y _ {t}\right) \right| \sqrt {x _ {t} ^ {2} + y _ {t} ^ {2}}} \left[ \begin{array}{l} y _ {t} \\ x _ {t} \end{array} \right] = \left[ \begin{array}{l} x _ {t} \\ y _ {t} \end{array} \right] + \rho \frac {\operatorname {s i g n} \left(x _ {t} y _ {t}\right)}{\sqrt {x _ {t} ^ {2} + y _ {t} ^ {2}}} \left[ \begin{array}{l} y _ {t} \\ x _ {t} \end{array} \right], \tag {11}
$$
As  $y_0 \gg 0$ , we must have  $\eta \ll 1$ ; thus, these conditions automatically hold when  $\eta = o(1)$  and  $\rho = \mathcal{O}(1)$ .
where the second step uses the fact that  $\ell$  is a even function so  $\ell'(t)$  has the same sign with  $t$ . Define  $z_t = \frac{\sqrt{x_t^2 + y_t^2}}{\operatorname{sign}(x_t y_t)}$ . Then  $w_t + \rho \frac{\nabla \mathcal{L}(w_t)}{\|\nabla \mathcal{L}(w_t)\|} = \begin{bmatrix} x_t + \rho y_t z_t^{-1} \\ y_t + \rho x_t z_t^{-1} \end{bmatrix}$ . Further denoting  $\ell' \left( (x_t + \rho y_t z_t^{-1})(y_t + \rho x_t z_t^{-1}) \right)$  by  $\ell'_t$ , one can simplify the update rule Equation 1 as follows:
$$
\begin{array}{l} w _ {t + 1} = \left[ \begin{array}{c} x _ {t} \\ y _ {t} \end{array} \right] - \eta \ell^ {\prime} \big ((x _ {t} + \rho y _ {t} z _ {t} ^ {- 1}) (y _ {t} + \rho x _ {t} z _ {t} ^ {- 1}) \big) \left[ \begin{array}{c} y _ {t} + \rho x _ {t} z _ {t} ^ {- 1} \\ x _ {t} + \rho y _ {t} z _ {t} ^ {- 1} \end{array} \right] \\ = \left[ \begin{array}{l} x _ {t} \\ y _ {t} \end{array} \right] - \eta \ell_ {t} ^ {\prime} \left[ \begin{array}{l} y _ {t} + \rho x _ {t} z _ {t} ^ {- 1} \\ x _ {t} + \rho y _ {t} z _ {t} ^ {- 1} \end{array} \right]. \tag {12} \\ \end{array}
$$
Assumption. Following footnote 1, we assume that  $x_{t}, y_{t} \neq 0$  for all  $t$ .
We are ready to give some basic properties of Equation 12. First, we claim that the sign of  $\ell_t^\prime$  is the same as the product  $x_{t}y_{t}$ . Formally, we have the following lemma:
Lemma 17 (Sign of Gradient in SAM). If  $x_{t} \neq 0, y_{t} \neq 0$ , then  $\mathrm{sign}(x_{t}) = \mathrm{sign}(x_{t} + \rho y_{t}z_{t}^{-1})$  and  $\mathrm{sign}(y_{t}) = \mathrm{sign}(y_{t} + \rho x_{t}z_{t}^{-1})$ . In particular, if  $x_{t} \neq 0, y_{t} \neq 0$ , we must have  $\mathrm{sign}(\ell_t') = \mathrm{sign}(x_ty_t)$ .
Proof. Note that  $\mathrm{sign}(z_t) = \mathrm{sign}(x_t y_t)$ , so  $\mathrm{sign}(y_t z_t^{-1}) = \mathrm{sign}(x_t)$ . Similarly,  $\mathrm{sign}(x_t z_t^{-1}) = \mathrm{sign}(y_t)$ . Thus, we have  $\mathrm{sign}(x_t) = \mathrm{sign}(x_t + \rho y_t z_t^{-1})$  and  $\mathrm{sign}(y_t) = \mathrm{sign}(y_t + \rho x_t z_t^{-1})$ , which in turn shows that  $\mathrm{sign}(\ell_t') = \mathrm{sign}\bigl (\mathrm{sign}(x_t + \rho y_t z_t^{-1})\mathrm{sign}(y_t + \rho x_t z_t^{-1})\bigr) = \mathrm{sign}(x_t y_t)$ .
Using Lemma 17, the SAM update can be equivalently written as
$$
\left[ \begin{array}{l} x _ {t + 1} \\ y _ {t + 1} \end{array} \right] = \left[ \begin{array}{l} x _ {t} \\ y _ {t} \end{array} \right] - \eta \ell_ {t} ^ {\prime} \left[ \begin{array}{l} y _ {t} \\ x _ {t} \end{array} \right] - \eta | \ell_ {t} ^ {\prime} | \left[ \begin{array}{l} \rho x _ {t} \left(x _ {t} ^ {2} + y _ {t} ^ {2}\right) ^ {- 1 / 2} \\ \rho y _ {t} \left(x _ {t} ^ {2} + y _ {t} ^ {2}\right) ^ {- 1 / 2} \end{array} \right]. \tag {13}
$$
We then claim the following property, which can be intuitively interpolated as SAM always goes towards the minimizes (i.e., both axes) for each step, although sometimes it may overshoot.
Lemma 18. Suppose that  $x_{t}, y_{t} \neq 0$ . Then, the following basic properties hold: