SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} = \sum_ {i: y _ {i} = j} \bar {\rho} _ {j, r, i} ^ {(\widetilde {T}, 0)} \\ = \sum_ {i: y _ {i} = j} \bar {\rho} _ {j, r, i} ^ {(\widetilde {T} - 1, 0)} - \frac {\eta (P - 1) ^ {2}}{B m} \cdot \sum_ {i: y _ {i} = j} \ell_ {i} ^ {\prime} ^ {(\widetilde {T} - 1, b _ {i} ^ {(\widetilde {T} - 1)})} \cdot \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(\widetilde {T} - 1, b _ {i} ^ {(\widetilde {T} - 1)})}, \boldsymbol {\xi} _ {i} \rangle) \| \boldsymbol {\xi} _ {i} \| _ {2} ^ {2} \\ = \sum_ {i = 1} ^ {n} \overline {{\rho}} _ {j, r, i} ^ {(\widetilde {T} - 1, 0)} - \frac {\eta (P - 1) ^ {2}}{B m} \cdot \sum_ {i \in \widetilde {S} _ {j, r} ^ {(\widetilde {T} - 1, \widetilde {b} _ {i} ^ {(\widetilde {T} - 1)})}} \ell_ {i} ^ {\prime} ^ {(\widetilde {T} - 1, b _ {i} ^ {(\widetilde {T} - 1)})} \| \pmb {\xi} _ {i} \| _ {2} ^ {2} \\ \geq \sum_ {i = 1} ^ {n} \bar {\rho} _ {j, r, i} ^ {(\widetilde {T} - 1)} + \frac {\eta (P - 1) ^ {2} \sigma_ {p} ^ {2} d H \Phi (- 1)}{8 m} \cdot \min  _ {i \in \widetilde {S} _ {j, r} ^ {(\widetilde {t}, \widetilde {b} - 1)} \cap \mathcal {I} _ {\widetilde {t}, \widetilde {b} - 1}} | \ell_ {i} ^ {\prime (\widetilde {T} - 1, b _ {i} ^ {(\widetilde {T} - 1)})} |, \tag {52} \\ \end{array}
$$
where the last equality is by the definition of  $S_{j,r}^{(\widetilde{T} -1)}$  as  $\{i\in [n]:y_i = j,\langle \mathbf{w}_{j,r}^{(\widetilde{T} -1)},\pmb {\xi}_i\rangle >0\}$ ; the last inequality is by Lemma B.1 and the fifth statement of Lemma C.8. And
$$
\begin{array}{l} \gamma_ {j ^ {\prime}, r ^ {\prime}} ^ {(\widetilde {T}, 0)} \leq \gamma_ {j ^ {\prime}, r ^ {\prime}} ^ {(\widetilde {T} - 1, 0)} - \frac {\eta}{B m} \cdot \sum_ {i \in S _ {+}} \ell_ {i} ^ {\prime (\widetilde {T} - 1, b _ {i} ^ {(\widetilde {T} - 1)})} \sigma^ {\prime} \left(\left\langle \mathbf {w} _ {j ^ {\prime}, r ^ {\prime}} ^ {(\widetilde {T} - 1, b _ {i} ^ {\widetilde {T} - 1})}, \widehat {y} _ {i} \cdot \boldsymbol {\mu} \right\rangle\right) \cdot \| \boldsymbol {\mu} \| _ {2} ^ {2} \\ \leq \gamma_ {j ^ {\prime}, r ^ {\prime}} ^ {(\widetilde {T} - 1, 0)} + \frac {H \eta \| \boldsymbol {\mu} \| _ {2} ^ {2}}{m} \cdot \max  _ {i \in S _ {+}} \left| \ell_ {i} ^ {\prime} (\widetilde {T} - 1) \right|. \tag {53} \\ \end{array}
$$
According to the third statement of Lemma C.8, we have  $\max_{i\in S_+}|\ell_i'(\widetilde{T} -1,b_i^{(\widetilde{T} -1)})|\leq$ $C_2\min_{i\in S_{j,r}^{(\widetilde{T} -1,b_i^{(\widetilde{T} -1)})}}|\ell_i'^{(\widetilde{T} -1)}|$ . Then by combining (52) and (53), we have
$$
\frac {\sum_ {i = 1} ^ {n} \bar {\rho} _ {j , r , i} ^ {(\widetilde {T} , 0)}}{\gamma_ {j ^ {\prime} , r ^ {\prime}} ^ {(\widetilde {T} , 0)}} \geq \min  \left\{\frac {\sum_ {i = 1} ^ {n} \bar {\rho} _ {j , r , i} ^ {(\widetilde {T} - 1 , 0)}}{\gamma_ {j ^ {\prime} , r ^ {\prime}} ^ {(\widetilde {T} - 1 , 0)}}, \frac {(P - 1) ^ {2} \sigma_ {p} ^ {2} d}{1 6 C _ {2} \| \boldsymbol {\mu} \| _ {2} ^ {2}} \right\} = \Theta (\mathrm {S N R} ^ {- 2}). \tag {54}
$$
On the other hand, we will now show  $\frac{\sum_{i=1}^{n} \overline{\rho}_{j,r,i}^{(t,0)}}{\gamma_{j',r'}^{(t,0)}} \leq \Theta(\mathrm{SNR}^{-2})$  for  $t \geq T_2$  by induction. By Lemma B.1 and (52), we have
$$
\begin{array}{l} \sum_ {i = 1} ^ {n} \overline {{\rho}} _ {j, r, i} ^ {(T _ {2}, 0)} \leq \sum_ {i = 1} ^ {n} \overline {{\rho}} _ {j, r, i} ^ {(T _ {2} - 1, 0)} + \frac {3 \eta (P - 1) ^ {2} \sigma_ {p} ^ {2} d n}{2 B m} \\ \leq \frac {3 \eta (P - 1) ^ {2} \sigma_ {p} ^ {2} d n T _ {2}}{2 B m}. \\ \end{array}
$$
And, by (46), we know that at  $t = T_2$ , we have
$$
\gamma_ {j ^ {\prime}, r ^ {\prime}} ^ {(T _ {2}, 0)} \geq \frac {\eta c _ {3} c _ {4} C T _ {2} n \| \mu \| _ {2} ^ {2}}{8 B m}.
$$
$$
\frac{\sum_{i = 1}^{n}\overline{\rho}_{j,r,i}^{(T_{2},0)}}{\gamma_{j^{\prime},r^{\prime}}^{(T_{2},0)}}\leq \Theta (\mathrm{SNR}^{-2}).
$$
Suppose  $\frac{\sum_{i=1}^{n} \vec{\rho}_{j,r,i}^{(T,0)}}{\gamma_{j',r'}^{(T,0)}} \leq \Theta(\mathrm{SNR}^{-2})$ . According to the decomposition, we have:
$$
\begin{array}{l} \langle \mathbf {w} _ {j, r} ^ {(T, b)}, \widehat {y} _ {i} \boldsymbol {\mu} \rangle = \langle \mathbf {w} _ {j, r} ^ {(0, 0)}, \widehat {y} _ {i} \boldsymbol {\mu} \rangle + j \cdot \gamma_ {j, r} ^ {(T, b)} \cdot \widehat {y} _ {i} \\ + \frac {1}{P - 1} \sum_ {i = 1} ^ {n} \bar {\rho} _ {j, r, i} ^ {(T, b)} \cdot \| \boldsymbol {\xi} _ {i} \| _ {2} ^ {- 2} \langle \boldsymbol {\xi} _ {i}, \widehat {y} _ {i} \boldsymbol {\mu} \rangle + \frac {1}{P - 1} \sum_ {i = 1} ^ {n} \underline {{\rho}} _ {j, r, i} ^ {(T, b)} \cdot \| \boldsymbol {\xi} _ {i} \| _ {2} ^ {- 2} \langle \boldsymbol {\xi} _ {i}, \widehat {y} _ {i} \boldsymbol {\mu} \rangle . \tag {55} \\ \end{array}
$$
$$
\begin{array}{l} | \langle \mathbf {w} _ {j, r} ^ {(0, 0)}, \widehat {y} _ {i} \boldsymbol {\mu} \rangle + \frac {1}{P - 1} \sum_ {i = 1} ^ {n} \overline {{\rho}} _ {j, r, i} ^ {(T, b)} \cdot \| \boldsymbol {\xi} _ {i} \| _ {2} ^ {- 2} \langle \boldsymbol {\xi} _ {i}, \widehat {y} \boldsymbol {\mu} \rangle + \frac {1}{P - 1} \sum_ {i = 1} ^ {n} \underline {{\rho}} _ {j, r, i} ^ {(T, b)} \cdot \| \boldsymbol {\xi} _ {i} \| _ {2} ^ {- 2} \langle \boldsymbol {\xi} _ {i}, \widehat {y} \boldsymbol {\mu} | \\ \leq \beta / 2 + | \sum_ {i = 1} ^ {n} \overline {{\rho}} _ {j, r, i} ^ {(T, b)} | \frac {4 \| \mu \| _ {2} \sqrt {2 \log (6 n / \delta)}}{\sigma_ {p} d (P - 1)} \\ \leq \beta / 2 + \frac {\Theta (\mathrm {S N R} ^ {- 1}) \gamma_ {j , r} ^ {(T , b)}}{\sqrt {d}} \\ \leq \gamma_ {j, r} ^ {(T, 0)}, \\ \end{array}
$$
where the first inequality is due to triangle inequality and Lemma B.1, the second inequality is due to induction hypothesis, and the last inequality is due to Condition 3.1.
Thus, the sign of  $\langle \mathbf{w}_{j,r}^{(T,b)},\widehat{y}_i\pmb {\mu}\rangle$  is persistent through out the epoch. Then, without loss of generality, we suppose  $\langle \mathbf{w}_{j,r}^{(T,b)},\pmb {\mu}\rangle >0$  . Thus, the update rule of  $\gamma$  is:
$$
\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, b)} - \sum_ {i \in \mathcal {I} _ {T, b} \cap S _ {-} \cap S _ {1}} \ell_ {i} ^ {\prime (T, b)} \right] \| \boldsymbol {\mu} \| _ {2} ^ {2} \\ \geq \gamma_ {j, r} ^ {(T, b)} + \frac {\eta}{B m} \cdot \left[ \min  _ {i \in \mathcal {I} _ {T, b}} \ell_ {i} ^ {(T, b)} | \mathcal {I} _ {T, b} \cap S _ {+} \cap S _ {1} | - \max  _ {i \in \mathcal {I} _ {T, b}} | \mathcal {I} _ {T, b} \cap S _ {-} \cap S _ {- 1} | \right] \cdot \| \boldsymbol {\mu} \| _ {2} ^ {2}. \\ \end{array}
$$
$$
\gamma_ {j, r} ^ {(T + 1, 0)} \geq \gamma_ {j, r} ^ {(T, b)} + \frac {\eta}{B m} \cdot \left[ \min  \ell_ {i} ^ {(T, b _ {i} ^ {(T)})} | S _ {+} \cap S _ {1} | - \max  \ell_ {i} ^ {(T, b _ {i} ^ {(T)})} | S _ {-} \cap S _ {- 1} | \right] \cdot \| \boldsymbol {\mu} \| _ {2} ^ {2}. \tag {56}
$$
And, by (52), we have
$$
\sum_ {i = 1} ^ {n} \bar {\rho} _ {j, r, i} ^ {(T + 1, 0)} \leq \sum_ {i = 1} ^ {n} \bar {\rho} _ {j, r, i} ^ {(T, 0)} + \frac {\eta (P - 1) ^ {2} \sigma_ {p} ^ {2} d H \Phi (- 1)}{8 m} \cdot \max  \left| \ell_ {i} ^ {\prime (T, b _ {i} ^ {(T)})} \right|. \tag {57}
$$
Thus, combining (56) and (57), we have
$$
\begin{array}{l} \frac{\sum_{i = 1}^{n}\overline{\rho}_{j,r,i}^{(T + 1,0)}}{\gamma_{j,r}^{(T + 1,0)}} \\ \leq \max \left\{\frac {\sum_ {i = 1} ^ {n} \overline {{\rho}} _ {j , r , i} ^ {(T , 0)}}{\gamma_ {j , r} ^ {(T , 0)}}, \frac {(P - 1) ^ {2} \sigma_ {p} ^ {2} d n \Phi (- 1) \cdot \max | \ell_ {i} ^ {\prime (T , b _ {i} ^ {(T)})} |}{8 \big [ \min \ell_ {i} ^ {(T , b _ {i} ^ {(T)})} | S _ {+} \cap S _ {1} | - \max \ell_ {i} ^ {(T , b _ {i} ^ {(T)})} | S _ {-} \cap S _ {- 1} | \big ] \cdot \| \boldsymbol {\mu} \| _ {2} ^ {2}} \right\} \\ \leq \Theta (\mathrm {S N R} ^ {- 2}), \tag {58} \\ \end{array}
$$
where the last inequality is due to the induction hypothesis, the third statement of Lemma C.8, and Lemma B.5. Thus, by induction, we have for all  $T_{1} \leq t \leq T^{*}$  that
$$
\frac {\sum_ {i = 1} ^ {n} \overline {{\rho}} _ {j , r ; i} ^ {(t , 0)}}{\gamma_ {j ^ {\prime} , r ^ {\prime}} ^ {(t , 0)}} \leq \Theta (\mathrm {S N R} ^ {- 2}).
$$
And for  $(T_{1},0)\leq (t,b)\leq (T^{*},0)$  , we can bound the ratio as follows:
$$
\frac {\sum_ {i = 1} ^ {n} \overline {{\rho}} _ {j , r , i} ^ {(t , b)}}{\gamma_ {j ^ {\prime} , r ^ {\prime}} ^ {(t , b)}} \leq \frac {4 \sum_ {i = 1} ^ {n} \overline {{\rho}} _ {j , r , i} ^ {(t , 0)}}{\gamma_ {j ^ {\prime} , r ^ {\prime}} ^ {(t , 0)}} \leq \Theta (\mathrm {S N R} ^ {- 2}),
$$
where the first inequality is due to the update rule of  $\overline{\rho}_{j,r,i}^{(t,b)}$  and  $\overline{\rho}_{j,r,i}^{(t,b)}$ . Thus, we have completed the proof.
C.3 Test Error Analysis
In this section, we present and prove the exact upper bound and lower bound of test error in Theorem 3.2. Since we have resolved the challenges brought by stochastic mini-batch parameter update, the remaining proof for test error is similar to the counterpart in Kou et al. (2023).
C.3.1 Test Error Upper Bound
First, we prove the upper bound of test error in Theorem 3.2 when the training loss converges.
Theorem C.17 (Second part of Theorem 3.2). Under the same conditions as Theorem 3.2, then there exists a large constant  $C_1$  such that when  $n\|\pmb{\mu}\|_2^2 \geq C_1(P - 1)^4\sigma_p^4 d$ , for time  $t$  defined in Lemma C.15, we have the test error
$$
\mathbb {P} _ {(\mathbf {x}, y) \sim \mathcal {D}} \big (y \neq \mathrm {s i g n} (f (\mathbf {W} ^ {(t, 0)}, \mathbf {x})) \big) \leq p + \exp \left(- n \| \boldsymbol {\mu} \| _ {2} ^ {4} / (C _ {2} (P - 1) ^ {4} \sigma_ {p} ^ {4} d)\right),
$$
where  $C_2 = O(1)$
Proof. The proof is similar to the proof of Theorem E.1 in Kou et al. (2023). The only difference is substituting  $\xi$  in their proof with  $(P - 1)\xi$ .
C.3.2 Test Error Lower Bound
In this part, we prove the lower bound of the test error in Theorem 3.2. We give two key Lemmas.
Lemma C.18. For  $(T_1,0)\leq (t,b) < (T^*,0)$ , denote  $g(\pmb {\xi}) = \sum_{j,r}j(P - 1)\sigma (\langle \mathbf{w}_{j,r}^{(t,b)},\pmb {\xi}\rangle)$ . There exists a fixed vector  $\mathbf{v}$  with  $\| \mathbf{v}\| _2\leq 0.06\sigma_p$  such that
$$
\sum_ {j ^ {\prime} \in \{\pm 1 \}} \left[ g \left(j ^ {\prime} \boldsymbol {\xi} + \mathbf {v}\right) - g \left(j ^ {\prime} \boldsymbol {\xi}\right) \right] \geq 4 C _ {6} \max  _ {j \in \{\pm 1 \}} \left\{\sum_ {r} \gamma_ {j, r} ^ {(t, b)} \right\}, \tag {59}
$$
for all  $\pmb {\xi}\in \mathbb{R}^d$
Proof of Lemma C.18. The proof is similar to the proof of Lemma 5.8 in Kou et al. (2023). The only difference is substituting  $\pmb{\xi}$  in their proof with  $(P - 1)\pmb{\xi}$ .
Lemma C.19 (Proposition 2.1 in Devroye et al. (2018)). The TV distance between  $\mathcal{N}(0, \sigma_p^2 \mathbf{I}_d)$  and  $\mathcal{N}(\mathbf{v}, \sigma_p^2 \mathbf{I}_d)$  is smaller than  $\| \mathbf{v} \|_2 / 2\sigma_p$ .
Then, we can prove the lower bound of the test error.
Theorem C.20 (Third part of Theorem 3.2). Suppose that  $n\|\pmb{\mu}\|_2^4 \leq C_3 d(P - 1)^4 \sigma_p^4$ , then we have that  $L_{\mathcal{D}}^{0-1}(\mathbf{W}^{(t,0)}) \geq p + 0.1$ , where  $C_3$  is an sufficiently large absolute constant.
Proof. The proof is similar to the proof of Theorem 4.3 in Kou et al. (2023). The only difference is substituting  $\xi$  in their proof with  $(P - 1)\xi$ .
D Proofs for SAM
D.1 Noise Memorization Prevention
The following lemma shows the update rule of the neural network
Lemma D.1. We denote  $\ell_i^{\prime (t,b)} = \ell^{\prime}[y_i\cdot f(\mathbf{W}^{(t,b)},\mathbf{x}_i)]$ , then the adversarial point of  $\mathbf{W}^{(t,b)}$  is  $\mathbf{W}^{(t,b)} + \widehat{\epsilon}^{(t,b)}$ , where
$$
\widehat {\pmb {\epsilon}} _ {j, r} ^ {(t, b)} = \frac {\tau}{m} \frac {\sum_ {i \in \mathcal {I} _ {t , b}} \sum_ {p \in [ P ]} \ell_ {i} ^ {\prime (t , b)} j \cdot y _ {i} \sigma^ {\prime} (\langle \mathbf {w} _ {j , r} ^ {(t , b)} , \mathbf {x} _ {i , p} \rangle) \mathbf {x} _ {i , p}}{\| \nabla_ {\mathbf {W}} L _ {\mathcal {I} _ {t , b}} (\mathbf {W} ^ {(t , b)}) \| _ {F}}.
$$
Then the training update rule of the parameter is
$$
\mathbf {w} _ {j, r} ^ {(t + 1, b)} = \mathbf {w} _ {j, r} ^ {(t, b)} - \frac {\eta}{B m} \sum_ {i \in \mathcal {I} _ {t, b}} \sum_ {p \in [ P ]} \ell_ {i} ^ {\prime (t, b)} \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(t, b)} + \widehat {\epsilon} _ {j, r} ^ {(t, b)}, \mathbf {x} _ {i, p} \rangle) j \cdot \mathbf {x} _ {i, p}
$$
$$
\begin{array}{l} = \mathbf {w} _ {j, r} ^ {(t, b)} - \frac {\eta}{B m} \sum_ {i \in \mathcal {I} _ {t, b}} \sum_ {p \in [ P ]} \ell_ {i} ^ {\prime (t, b)} \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(t, b)}, \mathbf {x} _ {i, p} \rangle + \langle \widehat {\boldsymbol {\epsilon}} _ {t, j, r}, \mathbf {x} _ {i, p} \rangle) j \cdot \mathbf {x} _ {i, p} \\ = \mathbf {w} _ {j, r} ^ {(t, b)} - \frac {\eta}{B m} \sum_ {i \in \mathcal {I} _ {t, b}} \ell_ {i} ^ {\prime (t, b)} \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(t, b)}, y \boldsymbol {\mu} \rangle + \langle \widehat {\boldsymbol {\epsilon}} _ {t, j, r}, y \boldsymbol {\mu} \rangle) j \boldsymbol {\mu} \\ \underbrace{-\frac{\eta(P - 1)}{Bm}\sum_{i\in\mathcal{I}_{t,b}}\ell^{\prime}_{i}^{(t,b)}\sigma^{\prime}(\langle\mathbf{w}^{(t,b)}_{j,r},\boldsymbol{\xi}_{i}\rangle + \langle\widehat{\boldsymbol{\epsilon}}^{(t,b)}_{j,r},\boldsymbol{\xi}_{i}\rangle)jy_{i}\boldsymbol{\xi}_{i}}_{\text{NoiseTerm}}. \\ \end{array}
$$
We will show that the noise term will be small if we train with the SAM algorithm. We consider the first stage where  $t \leq T_1$  where  $T_1 = mB / (12nn\|\pmb{\mu}\|_2^2)$ . Then, the following property holds.
Proposition D.2. Under Assumption 3.1, for  $0 \leq t \leq T_{1}$ , we have that
$$
\gamma_ {j, r} ^ {(0, 0)}, \bar {\rho} _ {j, r, i} ^ {(0, 0)}, \underline {{\rho}} _ {j, r, i} ^ {(0, 0)} = 0, \tag {60}
$$
$$
0 \leq \gamma_ {j, r} ^ {(t, b)} \leq 1 / 1 2, \tag {61}
$$
$$
0 \leq \bar {\rho} _ {j, r, i} ^ {(t, b)} \leq 1 / 1 2, \tag {62}
$$
$$
0 \geq \underline {{\rho}} _ {j, r, i} ^ {(t, b)} \geq - \beta - 1 0 \sqrt {\frac {\log (6 n ^ {2} / \delta)}{d}} n. \tag {63}
$$
Besides,  $\gamma_{j,r}^{(T_1,0)} = \Omega (1)$
Lemma D.3. Under Assumption 3.1, suppose (27), (28) and (29) hold at iteration  $t$ . Then, for all  $r \in [m]$ ,  $j \in \{\pm 1\}$  and  $i \in [n]$ ,
$$
\left| \left\langle \mathbf {w} _ {j, r} ^ {(t, b)} - \mathbf {w} _ {j, r} ^ {(0, 0)}, \boldsymbol {\mu} \right\rangle - j \cdot \gamma_ {j, r} ^ {(t, b)} \right| \leq \operatorname {S N R} \sqrt {\frac {3 2 \log (6 n / \delta)}{d}} n \alpha , \tag {64}
$$
$$
\left| \left\langle \mathbf {w} _ {j, r} ^ {(t, b)} - \mathbf {w} _ {j, r} ^ {(0, 0)}, \boldsymbol {\xi} _ {i} \right\rangle - \frac {1}{P - 1} \underline {{\rho}} _ {j, r, i} ^ {(t, b)} \right| \leq \frac {5}{P - 1} \sqrt {\frac {\log (6 n ^ {2} / \delta)}{d}} n \alpha , j \neq y _ {i}, \tag {65}
$$
$$
\left| \left\langle \mathbf {w} _ {j, r} ^ {(t, b)} - \mathbf {w} _ {j, r} ^ {(0, 0)}, \boldsymbol {\xi} _ {i} \right\rangle - \frac {1}{P - 1} \bar {\rho} _ {j, r, i} ^ {(t, b)} \right| \leq \frac {5}{P - 1} \sqrt {\frac {\log (6 n ^ {2} / \delta)}{d}} n \alpha , j = y _ {i}. \tag {66}
$$
Proof of Lemma D.3. Notice that  $1 / 12 < \alpha$ , if the condition (61), (62), (63) holds, (27), (28) and (29) also holds. Therefore, by Lemma C.3, we know that Lemma D.3 also holds.
Lemma D.4. Under Assumption 3.1, suppose (61), (62), (63) hold at iteration  $t, b$ . Then, for all  $j \in \{\pm 1\}$  and  $i \in [n]$ ,  $F_{j}(\mathbf{W}_{j}^{(t,b)}, \mathbf{x}_{i}) \leq 0.5$ . Therefore  $-0.3 \geq \ell_{i}' \geq -0.7$ .
Proof. Notice that  $1 / 12 < \alpha$ , if the condition (61), (62), (63) holds, (27), (28) and (29) also holds. Therefore, by Lemma C.4, we know that for all  $j \neq y_i$  and  $i \in [n]$ ,  $F_j(\mathbf{W}_j^{(t,b)}, \mathbf{x}_i) \leq 0.5$ . Next we will show that for  $j = y_i$ ,  $F_j(\mathbf{W}_j^{(t,b)}, \mathbf{x}_i) \leq 0.5$  also holds.
According to Lemma D.3, we have
$$
\begin{array}{l} F _ {j} \left(\mathbf {W} _ {j} ^ {(t, b)}, \mathbf {x} _ {i}\right) = \frac {1}{m} \sum_ {r = 1} ^ {m} \left[ \sigma \left(\left\langle \mathbf {w} _ {j, r} ^ {(t, b)}, y _ {i} \boldsymbol {\mu} \right\rangle\right) + (P - 1) \sigma \left(\left\langle \mathbf {w} _ {j, r} ^ {(t)}, \boldsymbol {\xi} _ {i} \right\rangle\right) \right] \\ \leq 2 \max  \left\{ \right.\left|\left\langle \mathbf {w} _ {j, r} ^ {(t, b)}, y _ {i} \boldsymbol {\mu} \right\rangle\right|, (P - 1) \mid \left\langle \mathbf {w} _ {j, r} ^ {(t)}, \boldsymbol {\xi} _ {i} \right\rangle\left. \right|\left. \right\} \\ \leq 6 \max  \left\{\left| \langle \mathbf {w} _ {j, r} ^ {(0)}, \widehat {y} _ {i} \boldsymbol {\mu} \rangle \right|, (P - 1) \left| \langle \mathbf {w} _ {j, r} ^ {(0)}, \boldsymbol {\xi} _ {i} \rangle \right|, \mathrm {S N R} \sqrt {\frac {3 2 \log (6 n / \delta)}{d}} n \alpha , \right. \\ \left. 5 \sqrt {\frac {\log (6 n ^ {2} / \delta)}{d}} n \alpha , | \gamma_ {j, r} ^ {(t, b)} |, | \underline {{\rho}} _ {j, r, i} ^ {(t, b)} | \right\} \\ \leq 6 \max  \left\{\beta , \operatorname {S N R} \sqrt {\frac {3 2 \log (6 n / \delta)}{d}} n \alpha , 5 \sqrt {\frac {\log (6 n ^ {2} / \delta)}{d}} n \alpha , | \gamma_ {j, r} ^ {(t, b)} |, | \bar {\rho} _ {j, r, i} ^ {(t, b)} | \right\} \\ \end{array}
$$
$$
\leq 0. 5,
$$
where the second inequality is by (64), (65) and (66); the third inequality is due to the definition of  $\beta$ ; the last inequality is by (25), (61), (62).
Since  $F_{j}(\mathbf{W}_{j}^{(t,b)},\mathbf{x}_{i})\in [0,0.5]$  we know that
$$
- 0. 3 \geq - \frac {1}{1 + \exp (0 . 5)} \geq \ell_ {i} ^ {\prime} \geq - \frac {1}{1 + \exp (- 0 . 5)} \geq - 0. 7.
$$
Based on the previous foundation lemmas, we can provide the key lemma of SAM which is different from the dynamic of SGD.
Lemma D.5. Under Assumption 3.1, suppose (61), (62) and (63) hold at iteration  $t$ ,  $b$ . We have that if  $\langle \mathbf{w}_{j,r}^{(t,b)}, \pmb{\xi}_k \rangle \geq 0$ ,  $k \in \mathcal{I}_{t,b}$  and  $j = y_k$ , then  $\langle \mathbf{w}_{j,r}^{(t,b)} + \widehat{\pmb{\epsilon}}_{j,r}^{(t,b)}, \pmb{\xi}_k \rangle < 0$ .
Proof. We first prove that there for  $t \leq T_1$ , there exists a constant  $C_2$  such that
$$
\left\| \nabla_ {\mathbf {W}} L _ {\mathcal {I} _ {t, b}} \left(\mathbf {W} ^ {(t, b)}\right) \right\| _ {F} \leq C _ {2} P \sigma_ {p} \sqrt {d / B}.
$$
$$
L _ {\mathcal {I} _ {t, b}} (\mathbf {W} ^ {(t, b)}) = \frac {1}{B} \sum_ {i \in \mathcal {I} _ {t, b}} \ell (y _ {i} f (\mathbf {W} ^ {(t, b)}, x _ {i})),
$$
$$
\begin{array}{l} \nabla_ {\mathbf {w} _ {j, r}} L _ {\mathcal {I} _ {t, b}} (\mathbf {W} ^ {(t, b)}) = \frac {1}{B} \sum_ {i \in \mathcal {I} _ {t, b}} \nabla_ {\mathbf {w} _ {j, r}} \ell (y _ {i} f (\mathbf {W} ^ {(t, b)}, \mathbf {x} _ {i})) \\ = \frac {1}{B} \sum_ {i \in \mathcal {I} _ {t, b}} y _ {i} \ell^ {\prime} (y _ {i} f (\mathbf {W} ^ {(t, b)}, \mathbf {x} _ {i})) \nabla_ {\mathbf {w} _ {j, r}} f (\mathbf {W} ^ {(t, b)}, \mathbf {x} _ {i}) \\ = \frac {1}{B m} \sum_ {i \in \mathcal {I} _ {t, b}} y _ {i} \ell_ {i} ^ {\prime (t, b)} [ \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(t, b)}, \boldsymbol {\mu} \rangle) \cdot \boldsymbol {\mu} + \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(t, b)}, \boldsymbol {\xi} _ {i} \rangle) \cdot (P - 1) \boldsymbol {\xi} _ {i} ]. \\ \end{array}
$$
$$
\begin{array}{l} \left\| \nabla_ {\mathbf {w} _ {j, r}} L _ {\mathcal {I} _ {t, b}} \left(\mathbf {W} ^ {(t, b)}\right) \right\| _ {2} \\ \leq \frac {1}{B m} \left\| \sum_ {i \in \mathcal {I} _ {t, b}} | \ell_ {i} ^ {\prime (t, b)} | \cdot \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(t)}, \boldsymbol {\mu} \rangle) \cdot \boldsymbol {\mu} \right\| _ {2} + \frac {1}{B m} \left\| \sum_ {i \in \mathcal {I} _ {t, b}} | \ell_ {i} ^ {\prime (t, b)} | \cdot \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(t)}, \boldsymbol {\xi} _ {i} \rangle) \cdot (P - 1) \boldsymbol {\xi} _ {i} \right\| _ {2} \\ \leq 0. 7 m ^ {- 1} \| \boldsymbol {\mu} \| _ {2} + 1. 4 (P - 1) m ^ {- 1} \sigma_ {p} \sqrt {d / B} \\ \leq 2 P m ^ {- 1} \sigma_ {p} \sqrt {d / B}, \\ \end{array}
$$
$$
\| \nabla_ {\mathbf {W}} L _ {\mathcal {I} _ {t, b}} (\mathbf {W} ^ {(t, b)}) \| _ {F} ^ {2} = \sum_ {j, r} \| \nabla_ {\mathbf {w} _ {j, r}} L _ {\mathcal {I} _ {t, b}} (\mathbf {W} ^ {(t, b)}) \| _ {2} ^ {2} \leq 2 m (2 P m ^ {- 1} \sigma_ {p} \sqrt {d / B}) ^ {2},
$$
$$
\left\| \nabla_ {\mathbf {W}} L _ {\mathcal {I} _ {t, b}} \left(\mathbf {W} ^ {(t, b)}\right) \right\| _ {F} \leq 2 \sqrt {2} P \sigma_ {p} \sqrt {d / B m}.
$$
$$
\begin{array}{l} \langle \widehat {\boldsymbol {\epsilon}} _ {j, r} ^ {(t, b)}, \boldsymbol {\xi} _ {k} \rangle = \frac {\tau}{m B} \| \nabla_ {\mathbf {W}} L _ {\mathcal {I} _ {t, b}} (\mathbf {W} ^ {(t, b)}) \| _ {F} ^ {- 1} \sum_ {i \in \mathcal {I} _ {t, b}} \sum_ {p \in [ P ]} \ell_ {i} ^ {\prime (t)} j \cdot y _ {i} \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(t)}, \mathbf {x} _ {i, p} \rangle) \langle \mathbf {x} _ {i, p}, \boldsymbol {\xi} _ {k} \rangle \\ = \frac {\tau}{m B} \| \nabla_ {\mathbf {W}} L _ {\mathcal {I} _ {t, b}} \left(\mathbf {W} ^ {(t, b)}\right) \| _ {F} ^ {- 1} \cdot \left(\sum_ {i \in \mathcal {I} _ {t, b}, i \neq k} \ell_ {i} ^ {\prime (t, b)} j y _ {i} \cdot (P - 1) \sigma^ {\prime} \left(\left\langle \mathbf {w} _ {j, r} ^ {(t, b)}, \boldsymbol {\xi} _ {i} \right\rangle\right) \left\langle \boldsymbol {\xi} _ {i}, \boldsymbol {\xi} _ {k} \right\rangle \right. \\ + \ell_ {k} ^ {\prime (t)} j y _ {k} \cdot (P - 1) \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(t)}, \pmb {\xi} _ {k} \rangle) \langle \pmb {\xi} _ {k}, \pmb {\xi} _ {k} \rangle \\ \end{array}
$$
$$
\begin{array}{l} \left. + \sum_ {i \in \mathcal {I} _ {t, b}} \ell_ {i} ^ {\prime (t, b)} j \cdot \sigma^ {\prime} \left(\left\langle \mathbf {w} _ {j, r} ^ {(t, b)}, 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. \\ \left. - 0. 1 5 (P - 1) \sigma_ {p} ^ {2} d \right] \\ <   - C \frac {\tau \sigma_ {p} \sqrt {d}}{m \sqrt {B}} \\ = - \frac {1}{4 (P - 1)}, \tag {67} \\ \end{array}
$$
where we the last equality is by choosing  $\tau = \frac{m\sqrt{B}}{C_3P\sigma_p\sqrt{d}}$ . Now we give an upper bound of  $\langle \mathbf{w}_{j,r}^{(t)},\pmb{\xi}_k\rangle$ , by (66) we have that
$$
\langle \mathbf {w} _ {j, r} ^ {(t)}, \boldsymbol {\xi} _ {k} \rangle \leq 3 \max  \left\{\left| \langle \mathbf {w} _ {j, r} ^ {(0)}, \boldsymbol {\xi} _ {i} \rangle \right|, 5 \sqrt {\frac {\log (6 n ^ {2} / \delta)}{d}} n \alpha , | \bar {\rho} _ {j, r, i} ^ {(t, b)} | \right\} \leq 1 / (4 (P - 1)). \tag {68}
$$
Combining (67) and (68) completes the proof.
Lemma D.6. Under Assumption 3.1, suppose (61), (62), (63) hold at iteration  $t$ ,  $b$ . Then (62) also holds for  $t$ ,  $b + 1$ .
Proof. Now consider the SAM algorithm. Recall that
$$
\overline {{\rho}} _ {j, r, i} ^ {(t, b + 1)} = \overline {{\rho}} _ {j, r, i} ^ {(t, b)} - \frac {\eta (P - 1) ^ {2}}{B m} \cdot \ell_ {i} ^ {\prime (t, b)} \cdot \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(t)} + \widehat {\pmb {\epsilon}} _ {j, r} ^ {(t)}, \pmb {\xi} _ {i} \rangle) \cdot \| \pmb {\xi} _ {i} \| _ {2} ^ {2} \cdot \mathbb {1} (y _ {i} = j) \mathbb {1} (i \in \mathcal {I} _ {t, b}).
$$
Case 1:  $i \notin \mathcal{I}_{t,b}$ . In this case, clearly we have that  $\overline{\rho}_{j,r,i}^{(t,b+1)} = \overline{\rho}_{j,r,i}^{(t,b)} \leq 1/12$ .
Case 2:  $i \in \mathcal{I}_{t,b}$  and  $\langle \mathbf{w}_{j,r}^{(t,b)}, \pmb{\xi}_i \rangle \geq 0$ , then by Lemma D.5, we have that  $\langle \mathbf{w}_{j,r}^{(t)} + \hat{\epsilon}_{j,r}^{(t)}, \pmb{\xi}_k \rangle < 0$ , therefore we have that  $\overline{\rho}_{j,r,i}^{(t,b+1)} = \overline{\rho}_{j,r,i}^{(t,b)} \leq 1/12$ .
Case 3:  $i \in \mathcal{I}_{t,b}$  and  $\langle \mathbf{w}_{j,r}^{(t,b)}, \pmb{\xi}_i \rangle \leq 0$  then by (66) and triangle inequality, we can conclude that  $\overline{\rho}_{j,r,i}^{(t,b)}$  can not reach a constant order,