SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\overline {{\rho}} _ {j, r, i} ^ {(t, b)} \leq (P - 1) \left| \left\langle \mathbf {w} _ {j, r} ^ {(t, b)} - \mathbf {w} _ {j, r} ^ {(0, 0)}, \boldsymbol {\xi} _ {i} \right\rangle \right| + 5 \sqrt {\frac {\log (6 n ^ {2} / \delta)}{d}} n \alpha .
$$
Then we can give an upper bound for  $\overline{\rho}_{j,r,i}^{(t + 1,b)}$  since we only take one small step further,
$$
\overline {{\rho}} _ {j, r, i} ^ {(t, b + 1)} \leq (P - 1) \big | \langle \mathbf {w} _ {j, r} ^ {(t, b)} - \mathbf {w} _ {j, r} ^ {(0, 0)}, \pmb {\xi} _ {i} \rangle \big | + 5 \sqrt {\frac {\log (6 n ^ {2} / \delta)}{d}} n \alpha + \frac {\eta (P - 1) ^ {2}}{B m} \cdot 2 d \sigma_ {p} ^ {2} \leq 1 / 1 2.
$$
Proof of Proposition D.2. We will use induction to give the proof. The results are obvious hold at  $t = 0$  as all the coefficients are zero. Suppose that there exists  $\widetilde{T} \leq T_1$  such that the results in Proposition D.2 hold for all time  $(0,0) \leq (t,b) \leq (\widetilde{T} - 1,\widetilde{b} - 1)$ . We aim to prove that (61), (62), (63) also hold for iteration  $(\widetilde{T} - 1,\widetilde{b})$ .
First, we prove that (61) holds for hold for iteration  $(\widetilde{T} - 1, \widetilde{b})$ . Notice that
$$
\gamma_ {j, r} ^ {(t, b + 1)} = \gamma_ {j, r} ^ {(t, b)} - \frac {\eta}{B m} \cdot \sum_ {i \in \mathcal {I} _ {t, b}} \ell_ {i} ^ {\prime (t, b)} \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(t, b)}, y _ {i} \cdot \boldsymbol {\mu} \rangle) \cdot \| \boldsymbol {\mu} \| _ {2} ^ {2} \leq \gamma_ {j, r} ^ {(t, b)} + \frac {\eta}{m} \| \boldsymbol {\mu} \| _ {2} ^ {2},
$$
where the last inequality is by the fact that  $|\ell_i^{\prime (t,b + 1)}|\leq 1$  and  $\sigma^{\prime}\leq 1$ . Notice that  $\widetilde{T} -1\leq T_1$ , we can conclude that,
$$
\gamma_ {j, r} ^ {(\widetilde {T}, \widetilde {b})} \leq T _ {1} \cdot (n / B) \cdot \frac {\eta}{m} \| \boldsymbol {\mu} \| _ {2} ^ {2} \leq 1 / 1 2.
$$
Second, by Lemma D.6, we know that (62) holds for  $(\widetilde{T} -1,\widetilde{b})$
Last, we need to prove that (63) holds  $(\widetilde{T} - 1, \widetilde{b})$ . The proof is similar to previous proof without SAM.
When  $\underline{\rho}_{j,r,k}^{(\widetilde{T} -1,\widetilde{b} -1)} < - 0.5(P - 1)\beta -6\sqrt{\frac{\log(6n^2 / \delta)}{d}} n\alpha$  , by (31), we have
$$
\begin{array}{l} \langle \mathbf {w} _ {j, r} ^ {(\widetilde {T} - 1, \widetilde {b} - 1)}, \boldsymbol {\xi} _ {k} \rangle <   \langle \mathbf {w} _ {j, r} ^ {(0, 0)}, \boldsymbol {\xi} _ {k} \rangle + \frac {1}{P - 1} \underline {{\rho}} _ {j, r, k} ^ {(\widetilde {T} - 1, \widetilde {b} - 1)} + \frac {5}{P - 1} \sqrt {\frac {\log (6 n ^ {2} / \delta)}{d}} n \alpha \\ \leq - \frac {1}{P - 1} \sqrt {\frac {\log (6 n ^ {2} / \delta)}{d}} n \alpha , \\ \end{array}
$$
$$
\begin{array}{l} \langle \widehat {\epsilon} _ {j, r} ^ {(\widetilde {T} - 1, \widetilde {b} - 1)}, \boldsymbol {\xi} _ {i} \rangle = \frac {\tau}{m B} \| \nabla_ {\mathbf {W}} L _ {\mathcal {I} _ {\widetilde {T} - 1, \widetilde {b} - 1}} (\mathbf {W} ^ {(\widetilde {T} - 1, \widetilde {b} - 1)}) \| _ {F} ^ {- 1} \sum_ {i \in \mathcal {I} _ {\widetilde {T} - 1, \widetilde {b} - 1}} \sum_ {p \in [ P ]} \ell_ {i} ^ {\prime (\widetilde {T} - 1, \widetilde {b} - 1)} j \cdot y _ {i} \\ \sigma^ {\prime} \left(\left\langle \mathbf {w} _ {j, r} ^ {(\widetilde {T} - 1, \widetilde {b} - 1)}, \mathbf {x} _ {i, p} \right\rangle\right) \left\langle \mathbf {x} _ {i, p}, \boldsymbol {\xi} _ {k} \right\rangle \\ = \frac {\tau}{m B} \| \nabla_ {\mathbf {W}} L _ {\mathcal {I} _ {\tilde {T} - 1, \tilde {b} - 1}} \left(\mathbf {W} ^ {(\tilde {T} - 1, \tilde {b} - 1)}\right) \| _ {F} ^ {- 1} \cdot \left(\sum_ {i \in \mathcal {I} _ {\tilde {T} - 1, \tilde {b} - 1}, i \neq k} \ell_ {i} ^ {\prime (\tilde {T} - 1, \tilde {b} - 1)} j \cdot y _ {i} \right. \\ (P - 1) \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(\widetilde {T} - 1, \widetilde {b} - 1)}, \boldsymbol {\xi} _ {i} \rangle) \langle \boldsymbol {\xi} _ {i}, \boldsymbol {\xi} _ {k} \rangle + \ell_ {k} ^ {\prime (t)} j y _ {k} \cdot (P - 1) \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(t)}, \boldsymbol {\xi} _ {k} \rangle) \langle \boldsymbol {\xi} _ {k}, \boldsymbol {\xi} _ {k} \rangle \\ \left. + \sum_ {i \in \mathcal {I} _ {\widetilde {T} - 1, \widetilde {b} - 1}} \ell_ {i} ^ {\prime (\widetilde {T} - 1, \widetilde {b} - 1)} j \cdot \sigma^ {\prime} \left(\left\langle \mathbf {w} _ {j, r} ^ {(\widetilde {T} - 1, \widetilde {b} - 1)}, y _ {i} \boldsymbol {\mu}\right) \left\langle \boldsymbol {\mu}, \boldsymbol {\xi} _ {k} \right\rangle\right)\right) \\ \leq \frac {\tau}{m C _ {2} P \sigma_ {p} \sqrt {B d}} \left[ 0. 8 B (P - 1) \sigma_ {P} ^ {2} \sqrt {d \log (6 n ^ {2} / \delta)} + 0. 4 B \sigma_ {P} \| \boldsymbol {\mu} \| _ {2} \sqrt {2 \log (6 n ^ {2} / \delta)} \right] \\ \leq C _ {4} \frac {\tau \sqrt {B} \sigma_ {p} \sqrt {\log (6 n ^ {2} / \delta)}}{m} \\ = C _ {4} \frac {B \sqrt {\log (6 n ^ {2} / \delta)}}{C _ {3} P \sqrt {d}} \\ \leq \frac {1}{P} \sqrt {\frac {\log (6 n ^ {2} / \delta)}{d}} n \alpha , \\ \end{array}
$$
and thus  $\langle \mathbf{w}_{j,r}^{(\widetilde{T} -1,\widetilde{b} -1)} + \widehat{\pmb{\epsilon}}_{j,r}^{(\widetilde{T} -1,\widetilde{b} -1)},\pmb {\xi}_i\rangle < 0$  which leads to
$$
\begin{array}{l} \underline {{\rho}} _ {j, r, i} ^ {(\widetilde {T} - 1, \widetilde {b})} = \underline {{\rho}} _ {j, r, i} ^ {(\widetilde {T} - 1, \widetilde {b} - 1)} + \frac {\eta (P - 1) ^ {2}}{B m} \cdot \ell_ {i} ^ {\prime (\widetilde {T} - 1, \widetilde {b} - 1)} \cdot \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(\widetilde {T} - 1, \widetilde {b} - 1)}, \boldsymbol {\xi} _ {i} \rangle) \cdot \| \boldsymbol {\xi} _ {i} \| _ {2} ^ {2}. \\ \mathbb {1} (y _ {i} = - j) \mathbb {1} (i \in \mathcal {I} _ {\widetilde {T} - 1, \widetilde {b} - 1}) \\ = \underline {{\rho}} _ {j, r, i} ^ {(\widetilde {T} - 1, \widetilde {b} - 1)}. \\ \end{array}
$$
$$
\underline {{\rho}} _ {j, r, i} ^ {(\widetilde {T} - 1, \widetilde {b})} = \underline {{\rho}} _ {j, r, i} ^ {(\widetilde {T} - 1, \widetilde {b} - 1)} \geq - (P - 1) \beta - 5 P \sqrt {\frac {\log (6 n ^ {2} / \delta)}{d}} n \alpha .
$$
When  $\underline{\rho}_{j,r,i}^{(\widetilde{T} -1,\widetilde{b} -1)}\geq -0.5(P - 1)\beta -5\sqrt{\frac{\log(6n^2 / \delta)}{d}} n\alpha$  we have that
$$
\begin{array}{l} \underline {{\rho}} _ {j, r, i} ^ {(\widetilde {T} - 1, \widetilde {b})} \geq \underline {{\rho}} _ {j, r, i} ^ {(\widetilde {T} - 1, \widetilde {b} - 1)} + \frac {\eta (P - 1) ^ {2}}{B m} \cdot \ell_ {i} ^ {\prime} ^ {(\widetilde {T} - 1, \widetilde {b} - 1)} \cdot \| \boldsymbol {\xi} _ {i} \| _ {2} ^ {2} \\ \geq \underline {{\rho}} _ {j, r, i} ^ {(\widetilde {T} - 1, \widetilde {b} - 1)} - \frac {0 . 4 \eta (P - 1) ^ {2}}{B m} \cdot 2 d \sigma_ {p} ^ {2} \\ \geq - (P - 1) \beta - 5 P \sqrt {\frac {\log (6 n ^ {2} / \delta)}{d}} n \alpha . \\ \end{array}
$$
Therefore, the induction is completed, and thus Proposition D.2 holds.
Next, we will prove that  $\gamma_{j,r}^{(t)}$  can achieve  $\Omega(1)$  after  $T_1 = mB/(12n\eta\|\pmb{\mu}\|_2^2)$  iterations. By Lemma B.6, we know that there exists  $c_3 \cdot T_1$  epochs such that at least  $c_4 \cdot H$  batches in these epochs, satisfy
$$
\left| S _ {+} \cap S _ {y} \cap \mathcal {I} _ {t, b} \right| \in \left[ \frac {B}{4}, \frac {3 B}{4} \right]
$$
for both  $y = +1$  and  $y = -1$ . For SAM, we have the following update rule for  $\gamma_{j,r}^{(t,b)}$ :
$$
\begin{array}{l} \gamma_ {j, r} ^ {(t, b + 1)} = \gamma_ {j, r} ^ {(t, b)} - \frac {\eta}{B m} \sum_ {i \in \mathcal {I} _ {t, b} \cap S _ {+}} \ell_ {i} ^ {\prime (t, b)} \sigma^ {\prime} \left(\left\langle \mathbf {w} _ {j, r} ^ {(t, b)} + \widehat {\boldsymbol {\epsilon}} _ {j, r} ^ {(t, b)}, y _ {i} \cdot \boldsymbol {\mu} \right\rangle\right) \cdot \| \boldsymbol {\mu} \| _ {2} ^ {2} \\ + \frac {\eta}{B m} \sum_ {i \in \mathcal {I} _ {t, b} \cap S _ {-}} \ell_ {i} ^ {\prime (t, b)} \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(t, b)} + \widehat {\pmb {\epsilon}} _ {j, r} ^ {(t, b)}, y _ {i} \cdot \pmb {\mu} \rangle) \cdot \| \pmb {\mu} \| _ {2} ^ {2}. \\ \end{array}
$$
If  $\langle \mathbf{w}_{j,r}^{(t,b)} + \widehat{\pmb{\epsilon}}_{j,r}^{(t,b)},\pmb {\mu}\rangle \geq 0$  , we have
$$
\begin{array}{l} \gamma_ {j, r} ^ {(t, b + 1)} = \gamma_ {j, r} ^ {(t, b)} - \frac {\eta}{B m} \cdot \left[ \sum_ {i \in \mathcal {I} _ {t, b} \cap S _ {+} \cap S _ {1}} \ell_ {i} ^ {\prime (t)} - \sum_ {i \in \mathcal {I} _ {t, b} \cap S _ {+} \cap S _ {- 1}} \ell_ {i} ^ {\prime (t)} \right] \| \boldsymbol {\mu} \| _ {2} ^ {2} \\ \geq \gamma_ {j, r} ^ {(t, b)} + \frac {\eta}{B m} \cdot \left(0. 3 | \mathcal {I} _ {t, b} \cap S _ {+} \cap S _ {1} | - 0. 7 | \mathcal {I} _ {t, b} \cap S _ {+} \cap S _ {- 1} |\right) \cdot \| \boldsymbol {\mu} \| _ {2} ^ {2}. \\ \end{array}
$$
If  $\langle \mathbf{w}_{j,r}^{(t,b)} + \widehat{\pmb{\epsilon}}_{j,r}^{(t,b)},\pmb {\mu}\rangle < 0$  , we have
$$
\begin{array}{l} \gamma_ {j, r} ^ {(t, b + 1)} = \gamma_ {j, r} ^ {(t, b)} + \frac {\eta}{B m} \cdot \left[ \sum_ {i \in \mathcal {I} _ {t, b} \cap S _ {+} \cap S _ {- 1}} \ell_ {i} ^ {\prime (t)} - \sum_ {i \in \mathcal {I} _ {t, b} \cap S _ {+} \cap S _ {1}} \ell_ {i} ^ {\prime (t)} \right] \| \boldsymbol {\mu} \| _ {2} ^ {2} \\ \geq \gamma_ {j, r} ^ {(t, b)} + \frac {\eta}{B m} \cdot \left(0. 3 | \mathcal {I} _ {t, b} \cap S _ {+} \cap S _ {- 1} | - 0. 7 | \mathcal {I} _ {t, b} \cap S _ {+} \cap S _ {1} |\right) \cdot \| \boldsymbol {\mu} \| _ {2} ^ {2}. \\ \end{array}
$$
$$
\begin{array}{l} \gamma_ {j, r} ^ {(T _ {1}, 0)} \geq \frac {\eta}{B m} \left(0. 3 \cdot c _ {3} T _ {1} \cdot c _ {4} H \cdot 0. 2 5 B - 0. 7 T _ {1} n q\right) \| \boldsymbol {\mu} \| _ {2} ^ {2} \\ = \frac {\eta}{B m} (0. 0 7 5 c _ {3} c _ {4} T _ {1} n - 0. 7 T _ {1} n q) \| \boldsymbol {\mu} \| _ {2} ^ {2} \\ \geq \frac {\eta}{1 6 B m} c _ {3} c _ {4} T _ {1} n \| \boldsymbol {\mu} \| _ {2} ^ {2} \\ = \frac {c _ {3} c _ {4}}{1 9 2} = \Omega (1). \\ \end{array}
$$
Lemma D.7. Suppose Condition 3.1 holds. Then we have that  $\left\| \mathbf{w}_{j,r}^{(T_1,0)}\right\| _2 = \Theta (\sigma_0\sqrt{d})$  and
$$
\langle \mathbf {w} _ {j, r} ^ {(T _ {1}, 0)}, j \boldsymbol {\mu} \rangle = \Omega (1),
$$
$$
\langle \mathbf {w} _ {- j, r} ^ {(T _ {1}, 0)}, j \boldsymbol {\mu} \rangle = - \Omega (1),
$$
$$
\widehat {\beta} := 2 \max _ {i, j, r} \{\left| \langle \mathbf {w} _ {j, r} ^ {(T _ {1}, 0)}, \boldsymbol {\mu} \rangle \right|, (P - 1) \left| \langle \mathbf {w} _ {j, r} ^ {(T _ {1}, 0)}, \boldsymbol {\xi} _ {i} \rangle \right| \} = O (1).
$$
Besides, for  $S_{i}^{(t,b)}$  and  $S_{j,r}^{(t,b)}$  defined in Lemma B.3 and B.4, we have that
$$
\begin{array}{l} \left| S _ {i} ^ {\left(T _ {1}, 0\right)} \right| = \Omega (m), \forall i \in [ n ] \\ \left| S _ {j, r} ^ {(T _ {1})} \right| = \Omega (n), \forall j \in \{\pm 1 \}, r \in [ m ]. \\ \end{array}
$$
D.2 Test Error Analysis
Proof of Theorem 4.1. Recall that
$$
\mathbf {w} _ {j, r} ^ {(t, b)} = \mathbf {w} _ {j, r} ^ {(0, 0)} + j \cdot \gamma_ {j, r} ^ {(t, b)} \cdot \| \boldsymbol {\mu} \| _ {2} ^ {- 2} \cdot \boldsymbol {\mu} + \frac {1}{P - 1} \sum_ {i = 1} ^ {n} \rho_ {j, r, i} ^ {(t, b)} \cdot \| \boldsymbol {\xi} _ {i} \| _ {2} ^ {- 2} \cdot \boldsymbol {\xi} _ {i},
$$
by triangle inequality we have
$$
\begin{array}{l} \left| \left\| \mathbf {w} _ {j, r} ^ {(T _ {1}, 0)} \right\| _ {2} - \left\| \mathbf {w} _ {j, r} ^ {(0, 0)} \right\| _ {2} \right| \leq | \gamma_ {j, r} ^ {(t, b)} | \cdot \| \boldsymbol {\mu} \| _ {2} ^ {- 1} + \frac {1}{P - 1} \left\| \sum_ {i = 1} ^ {n} | \rho_ {j, r, i} ^ {(t, b)} | \cdot \| \boldsymbol {\xi} _ {i} \| _ {2} ^ {- 2} \cdot \boldsymbol {\xi} _ {i} \right\| _ {2} \\ \leq \frac {1}{1 2} \| \boldsymbol {\mu} \| _ {2} ^ {- 1} + \frac {\sqrt {n}}{1 2 (P - 1)} \left(\sigma_ {p} ^ {2} d / 2\right) ^ {- 1 / 2} \\ \leq \frac {1}{6} \| \boldsymbol {\mu} \| _ {2} ^ {- 1}. \\ \end{array}
$$
By the condition on  $\sigma_0$  and Lemma B.2, we have
$$
\left\| \mathbf {w} _ {j, r} ^ {(T _ {1}, 0)} \right\| _ {2} = \Theta \big (\left\| \mathbf {w} _ {j, r} ^ {(0, 0)} \right\| _ {2} \big) = \Theta \left(\sigma_ {0} \sqrt {d}\right).
$$
By taking the inner product with  $\pmb{\mu}$  and  $\xi_{i}$ , we can get
$$
\langle \mathbf {w} _ {j, r} ^ {(t, b)}, \boldsymbol {\mu} \rangle = \langle \mathbf {w} _ {j, r} ^ {(0, 0)}, \boldsymbol {\mu} \rangle + j \cdot \gamma_ {j, r} ^ {(t, b)} + \frac {1}{P - 1} \sum_ {i = 1} ^ {n} \rho_ {j, r, i} ^ {(t, b)} \cdot \| \boldsymbol {\xi} _ {i} \| _ {2} ^ {- 2} \cdot \langle \boldsymbol {\xi} _ {i}, \boldsymbol {\mu} \rangle ,
$$
$$
\begin{array}{l} \langle \mathbf {w} _ {j, r} ^ {(t, b)}, \boldsymbol {\xi} _ {i} \rangle = \langle \mathbf {w} _ {j, r} ^ {(0, 0)}, \boldsymbol {\xi} _ {i} \rangle + j \cdot \gamma_ {j, r} ^ {(t, b)} \cdot \| \boldsymbol {\mu} \| _ {2} ^ {- 2} \cdot \langle \boldsymbol {\mu}, \boldsymbol {\xi} _ {i} \rangle + \frac {1}{P - 1} \sum_ {i ^ {\prime} = 1} ^ {n} \rho_ {j, r, i ^ {\prime}} ^ {(t, b)} \cdot \| \boldsymbol {\xi} _ {i ^ {\prime}} \| _ {2} ^ {- 2} \cdot \langle \boldsymbol {\xi} _ {i ^ {\prime}}, \boldsymbol {\xi} _ {i} \rangle \\ = \langle \mathbf {w} _ {j, r} ^ {(0, 0)}, \pmb {\xi} _ {i} \rangle + j \cdot \gamma_ {j, r} ^ {(t, b)} \cdot \| \pmb {\mu} \| _ {2} ^ {- 2} \cdot \langle \pmb {\mu}, \pmb {\xi} _ {i} \rangle + \frac {1}{P - 1} \rho_ {j, r, i} ^ {(t, b)} \\ + \frac {1}{P - 1} \sum_ {i \neq i ^ {\prime}} \rho_ {j, r, i ^ {\prime}} ^ {(t, b)} \cdot \| \pmb {\xi} _ {i ^ {\prime}} \| _ {2} ^ {- 2} \cdot \langle \pmb {\xi} _ {i ^ {\prime}}, \pmb {\xi} _ {i} \rangle . \\ \end{array}
$$
$$
\begin{array}{l} \langle \mathbf {w} _ {j, r} ^ {(T _ {1}, 0)}, j \boldsymbol {\mu} \rangle = \langle \mathbf {w} _ {j, r} ^ {(0, 0)}, j \boldsymbol {\mu} \rangle + \gamma_ {j, r} ^ {(T _ {1}, 0)} + \frac {1}{P - 1} \sum_ {i = 1} ^ {n} \rho_ {j, r, i} ^ {(T _ {1}, 0)} \cdot \| \boldsymbol {\xi} _ {i} \| _ {2} ^ {- 2} \cdot \langle \boldsymbol {\xi} _ {i}, j \boldsymbol {\mu} \rangle \\ \geq \gamma_ {j, r} ^ {(T _ {1}, 0)} - | \langle \mathbf {w} _ {j, r} ^ {(0, 0)}, \boldsymbol {\mu} \rangle | - \frac {1}{P - 1} \sum_ {i = 1} ^ {n} | \rho_ {j, r, i} ^ {(T _ {1}, 0)} | \cdot \| \boldsymbol {\xi} _ {i} \| _ {2} ^ {- 2} \cdot | \langle \boldsymbol {\xi} _ {i}, \boldsymbol {\mu} \rangle | \\ \geq \gamma_ {j, r} ^ {(T _ {1}, 0)} - \sqrt {2 \log (1 2 m / \delta)} \cdot \sigma_ {0} \| \boldsymbol {\mu} \| _ {2} - \frac {n}{1 2 (P - 1)} (\sigma_ {0} ^ {2} d / 2) ^ {- 1} \| \boldsymbol {\mu} \| _ {2} \sigma_ {p} \cdot \sqrt {2 \log (6 n / \delta)} \\ \geq \frac {1}{2} \gamma_ {j, r} ^ {(T _ {1}, 0)}, \\ \end{array}
$$
$$
\begin{array}{l} \langle \mathbf {w} _ {- j, r} ^ {(T _ {1}, 0)}, j \boldsymbol {\mu} \rangle = \langle \mathbf {w} _ {- j, r} ^ {(0, 0)}, j \boldsymbol {\mu} \rangle - \gamma_ {- j, r} ^ {(T _ {1}, 0)} - \frac {1}{P - 1} \sum_ {i = 1} ^ {n} \rho_ {- j, r, i} ^ {(T _ {1}, 0)} \cdot \| \boldsymbol {\xi} _ {i} \| _ {2} ^ {- 2} \cdot \langle \boldsymbol {\xi} _ {i}, j \boldsymbol {\mu} \rangle \\ \leq - \gamma_ {- j, r} ^ {(T _ {1}, 0)} + | \langle \mathbf {w} _ {- j, r} ^ {(0, 0)}, \boldsymbol {\mu} \rangle | + \frac {1}{P - 1} \sum_ {i = 1} ^ {n} | \rho_ {- j, r, i} ^ {(T _ {1}, 0)} | \cdot \| \boldsymbol {\xi} _ {i} \| _ {2} ^ {- 2} \cdot | \langle \boldsymbol {\xi} _ {i}, \boldsymbol {\mu} \rangle | \\ \leq - \gamma_ {- j, r} ^ {(T _ {1}, 0)} + \sqrt {2 \log (1 2 m / \delta)} \cdot \sigma_ {0} \| \boldsymbol {\mu} \| _ {2} + \frac {n}{1 2 (P - 1)} (\sigma_ {0} ^ {2} d / 2) ^ {- 1} \| \boldsymbol {\mu} \| _ {2} \sigma_ {p} \cdot \sqrt {2 \log (6 n / \delta)} \\ \leq - \frac {1}{2} \gamma_ {j, r} ^ {(T _ {1}, 0)}, \\ \end{array}
$$
where the last inequality is by the condition on  $\sigma_0$  and  $\gamma_{j,r}^{(T_1,0)} = \Omega (1)$ . Thus, it follows that
$$
\langle \mathbf {w} _ {j, r} ^ {(T _ {1}, 0)}, j \boldsymbol {\mu} \rangle = \Omega (1), \langle \mathbf {w} _ {- j, r} ^ {(T _ {1}, 0)}, j \boldsymbol {\mu} \rangle = - \Omega (1).
$$
$$
\begin{array}{l} | \langle \mathbf {w} _ {j, r} ^ {(T _ {1}, 0)}, \boldsymbol {\mu} \rangle | \leq | \langle \mathbf {w} _ {j, r} ^ {(0, 0)}, \boldsymbol {\mu} \rangle | + | \gamma_ {j, r} ^ {(T _ {1}, 0)} | + \frac {1}{P - 1} \sum_ {i = 1} ^ {n} | \rho_ {j, r, i} ^ {(t, b)} | \cdot \| \boldsymbol {\xi} _ {i} \| _ {2} ^ {- 2} \cdot | \langle \boldsymbol {\xi} _ {i}, \boldsymbol {\mu} \rangle | \\ \leq \frac {1}{2} \beta + \frac {1}{1 2} + \frac {n}{P - 1} \cdot \frac {1}{1 2} \left(\sigma_ {p} ^ {2} d / 2\right) ^ {- 1} \cdot \| \boldsymbol {\mu} \| _ {2} \sigma_ {p} \cdot \sqrt {2 \log (6 n / \delta)} \\ = \frac {1}{2} \beta + \frac {1}{1 2} + \frac {n}{6 (P - 1)} \| \boldsymbol {\mu} \| _ {2} \sqrt {2 \log (6 n / \delta)} / (\sigma_ {p} d) \\ \leq \frac {1}{6}, \\ \end{array}
$$
$$
\begin{array}{l} | \langle \mathbf {w} _ {j, r} ^ {(T _ {1}, 0)}, \boldsymbol {\xi} _ {i} \rangle | \leq | \langle \mathbf {w} _ {j, r} ^ {(0, 0)}, \boldsymbol {\xi} _ {i} \rangle | + | \gamma_ {j, r} ^ {(T _ {1}, 0)} | \cdot \| \boldsymbol {\mu} \| _ {2} ^ {- 2} \cdot | \langle \boldsymbol {\mu}, \boldsymbol {\xi} _ {i} \rangle | + \frac {1}{P - 1} | \rho_ {j, r, i} ^ {(T _ {1}, 0)} | \\ + \frac {1}{P - 1} \sum_ {i \neq i ^ {\prime}} \left| \rho_ {j, r, i ^ {\prime}} ^ {(T _ {1}, 0)} \right| \cdot \left\| \boldsymbol {\xi} _ {i ^ {\prime}} \right\| _ {2} ^ {- 2} \cdot | \langle \boldsymbol {\xi} _ {i ^ {\prime}}, \boldsymbol {\xi} _ {i} \rangle | \\ \leq \frac {1}{2} \beta + \frac {1}{1 2} \| \boldsymbol {\mu} \| _ {2} ^ {- 1} \sigma_ {p} \cdot \sqrt {2 \log (6 n / \delta)} + \frac {1}{1 2 (P - 1)} \\ + \frac {n}{1 2 (P - 1)} \left(\sigma_ {p} ^ {2} d / 2\right) ^ {- 1} 2 \sigma_ {p} ^ {2} \cdot \sqrt {d \log \left(6 n ^ {2} / \delta\right)} \\ \leq \frac {1}{2} \beta + \frac {1}{1 2 (P - 1)} + \frac {1}{6} \| \boldsymbol {\mu} \| _ {2} ^ {- 1} \sigma_ {p} \cdot \sqrt {\log (6 n / \delta)} \\ \leq \frac {1}{6}. \\ \end{array}
$$
$$
\widehat {\beta} := 2 \max  _ {i, j, r} \{\left| \langle \mathbf {w} _ {j, r} ^ {(T _ {1}, 0)}, \boldsymbol {\mu} \rangle \right|, (P - 1) \left| \langle \mathbf {w} _ {j, r} ^ {(T _ {1}, 0)}, \boldsymbol {\xi} _ {i} \rangle \right| \} = O (1).
$$
In addition, we also have for  $t \leq T_1$  and  $j = y_i$  that
$$
\begin{array}{l} \left\langle \mathbf {w} _ {j, r} ^ {(t, b)}, \boldsymbol {\xi} _ {i} \right\rangle - \left\langle \mathbf {w} _ {j, r} ^ {(0, 0)}, \boldsymbol {\xi} _ {i} \right\rangle \\ \geq \frac {1}{P - 1} \rho_ {j, r, i} ^ {(t, b)} - \gamma_ {j, r} ^ {(t, b)} \cdot \| \boldsymbol {\mu} \| _ {2} ^ {- 2} \cdot | \langle \boldsymbol {\mu}, \boldsymbol {\xi} _ {i} \rangle | - \frac {1}{P - 1} \sum_ {i \neq i ^ {\prime}} | \rho_ {j, r, i ^ {\prime}} ^ {(t, b)} | \cdot \| \boldsymbol {\xi} _ {i ^ {\prime}} \| _ {2} ^ {- 2} \cdot | \langle \boldsymbol {\xi} _ {i ^ {\prime}}, \boldsymbol {\xi} _ {i} \rangle | \\ \geq - \gamma_ {j, r} ^ {(t, b)} \cdot \| \boldsymbol {\mu} \| _ {2} ^ {- 2} \cdot | \langle \boldsymbol {\mu}, \boldsymbol {\xi} _ {i} \rangle | - \frac {1}{P - 1} \sum_ {i \neq i ^ {\prime}} \left| \rho_ {j, r, i ^ {\prime}} ^ {(t, b)} \right| \cdot \| \boldsymbol {\xi} _ {i ^ {\prime}} \| _ {2} ^ {- 2} \cdot | \langle \boldsymbol {\xi} _ {i ^ {\prime}}, \boldsymbol {\xi} _ {i} \rangle | \\ \geq - \frac {1}{1 2} \| \boldsymbol {\mu} \| _ {2} ^ {- 2} \cdot | \langle \boldsymbol {\mu}, \boldsymbol {\xi} _ {i} \rangle | - \frac {n}{1 2 (P - 1)} \| \boldsymbol {\xi} _ {i ^ {\prime}} \| _ {2} ^ {- 2} \cdot | \langle \boldsymbol {\xi} _ {i ^ {\prime}}, \boldsymbol {\xi} _ {i} \rangle | \\ \geq - \frac {1}{1 2} \| \boldsymbol {\mu} \| _ {2} ^ {- 1} \sigma_ {p} \cdot \sqrt {2 \log (6 n / \delta)} - \frac {n}{1 2 (P - 1)} (\sigma_ {p} ^ {2} d / 2) ^ {- 1} 2 \sigma_ {p} ^ {2} \cdot \sqrt {d \log (6 n ^ {2} / \delta)} \\ = - \frac {1}{1 2} \| \boldsymbol {\mu} \| _ {2} ^ {- 1} \sigma_ {p} \cdot \sqrt {2 \log (6 n / \delta)} - \frac {n}{3 (P - 1)} \sqrt {\log (6 n ^ {2} / \delta) / d} \\ \geq - \frac {1}{6} \| \boldsymbol {\mu} \| _ {2} ^ {- 1} \sigma_ {p} \cdot \sqrt {\log (6 n / \delta)}. \\ \end{array}
$$
Now let  $\bar{S}_i^{(0,0)}$  denote  $\{r: \langle \mathbf{w}_{y_i,r}^{(0,0)}, \pmb{\xi}_i \rangle > \sigma_0 \sigma_p \sqrt{d} \rangle\}$  and let  $\bar{S}_{j,r}^{(0,0)}$  denote  $\{i \in [n] : y_i = j, \langle \mathbf{w}_{y_i,r}^{(t,b)}, \pmb{\xi}_i \rangle > \sigma_0 \sigma_p \sqrt{d} \}$ . By the condition on  $\sigma_0$ , we have for  $t \leq T_1$  that
$$
\langle \mathbf {w} _ {j, r} ^ {(t, b)}, \boldsymbol {\xi} _ {i} \rangle \geq \frac {1}{\sqrt {2}} \langle \mathbf {w} _ {j, r} ^ {(0, 0)}, \boldsymbol {\xi} _ {i} \rangle ,
$$
for any  $r \in \bar{S}_i^{(0,0)}$  or  $i \in \bar{S}_{j,r}^{(0,0)}$ . Therefore, we have  $\bar{S}_i^{(0,0)} \subseteq S_i^{(T_1,0)}$  and  $\bar{S}_{j,r}^{(0,0)} \subseteq S_{j,r}^{(T_1,0)}$  and hence
$$
0. 8 \Phi (- \sqrt {2}) m \leq | \bar {S} _ {i} ^ {(0, 0)} | \leq | S _ {i} ^ {(T _ {1}, 0)} | = \Omega (m),
$$
$$
0. 2 5 \Phi (- \sqrt {2}) n \leq | \bar {S} _ {j, r} ^ {(0, 0)} | \leq | S _ {j, r} ^ {(T _ {1}, 0)} | = \Omega (n),
$$
where  $\Phi (\cdot)$  is the CDF of the standard normal distribution.
Now we can give proof of Theorem 4.
Proof of Theorem 4. After the training process of SAM after  $T_{1}$ , we get  $\mathbf{W}^{(T_1,0)}$ . To differentiate the SAM process and SGD process. We use  $\widetilde{\mathbf{W}}$  to denote the trajectory obtained by SAM in the proof, i.e.,  $\widetilde{\mathbf{W}}^{(T_1,0)}$ . By Proposition D.2, we have that
$$
\widetilde {\mathbf {w}} _ {j, r} ^ {(T _ {1}, 0)} = \widetilde {\mathbf {w}} _ {j, r} ^ {(0, 0)} + j \cdot \widetilde {\gamma} _ {j, r} ^ {(T _ {1}, 0)} \cdot \frac {\boldsymbol {\mu}}{\| \boldsymbol {\mu} \| _ {2} ^ {2}} + \frac {1}{P - 1} \sum_ {i = 1} ^ {n} \widetilde {\rho} _ {j, r, i} ^ {(T _ {1}, 0)} \cdot \frac {\boldsymbol {\xi} _ {i}}{\| \boldsymbol {\xi} _ {i} \| _ {2} ^ {2}} + \frac {1}{P - 1} \sum_ {i = 1} ^ {n} \widetilde {\rho} _ {j, r, i} ^ {(T _ {1}, 0)} \cdot \frac {\boldsymbol {\xi} _ {i}}{\| \boldsymbol {\xi} _ {i} \| _ {2} ^ {2}}, \tag {69}
$$
where  $\widetilde{\gamma}_{j,r}^{(T_1,0)} = \Theta (1),\widetilde{\overline{\rho}}_{j,r,i}^{(T_1,0)}\in [0,1 / 12],\widetilde{\underline{p}}_{j,r,i}^{(T_1,0)}\in [-\beta -10\sqrt{\log(6n^2 / \delta) / dn} n,0]$ . Then the SGD start at  $\mathbf{W}^{(0,0)}\coloneqq \widetilde{\mathbf{W}}^{(T_1,0)}$ . Notice that by Lemma D.7, we know that the initial weights of SGD (i.e., the end weight of SAM)  $\mathbf{W}^{(0,0)}$  still satisfies the conditions for Subsection C.1 and C.2. Therefore, following the same analysis in Subsection C.1 and C.2, we have that there exist  $t = \widetilde{O} (\eta^{-1}\epsilon^{-1}mnd^{-1}P^{-2}\sigma_p^{-2})$  such that  $L_{S}(\mathbf{W}^{(t,0)})\leq \epsilon$ . Besides,
$$
\mathbf {w} _ {j, r} ^ {(t, 0)} = \mathbf {w} _ {j, r} ^ {(0, 0)} + j \cdot \gamma_ {j, r} ^ {(t, 0)} \cdot \frac {\boldsymbol {\mu}}{\| \boldsymbol {\mu} \| _ {2} ^ {2}} + \frac {1}{P - 1} \sum_ {i = 1} ^ {n} \bar {\rho} _ {j, r, i} ^ {(t, 0)} \cdot \frac {\boldsymbol {\xi} _ {i}}{\| \boldsymbol {\xi} _ {i} \| _ {2} ^ {2}} + \frac {1}{P - 1} \sum_ {i = 1} ^ {n} \underline {{\rho}} _ {j, r, i} ^ {(t, 0)} \cdot \frac {\boldsymbol {\xi} _ {i}}{\| \boldsymbol {\xi} _ {i} \| _ {2} ^ {2}}
$$
for  $j \in [\pm 1]$  and  $r \in [m]$  where
$$
\gamma_ {j, r} ^ {(t, 0)} = \Theta (\mathrm {S N R} ^ {2}) \sum_ {i \in [ n ]} \bar {\rho} _ {j, r, i} ^ {(t, 0)}, \quad \bar {\rho} _ {j, r, i} ^ {(t, 0)} \in [ 0, \alpha ], \quad \underline {{\rho}} _ {j, r, i} ^ {(t, 0)} \in [ - \alpha , 0 ]. \tag {70}
$$
Next, we will evaluate the test error for  $\mathbf{W}^{(t,0)}$ . Notice that we use  $(t)$  as the shorthand notation of  $(t,0)$ . For the sake of convenience, we use  $(\mathbf{x},\widehat{y},y)\sim \mathcal{D}$  to denote the following: data point  $(\mathbf{x},y)$  follows distribution  $\mathcal{D}$  defined in Definition 2.1, and  $\widehat{y}$  is its true label. We can write out the test error as