SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Varying different starting points for SAM. In Section 4, we show that the SAM algorithm can effectively prevent noise memorization and thus improve feature learning in the early stage of training. Is SAM also effective if we add the algorithm in the middle of the training process? We conduct experiments on the ImageNet dataset with ResNet50. We choose the batch size as 1024, and the model is trained for 90 epochs with the best learning rate in grid search  $\{0.01, 0.03, 0.1, 0.3\}$ . The learning rate schedule is 10k steps linear warmup, then cosine decay. As shown in Table 1, the earlier SAM is introduced, the more pronounced its effectiveness becomes.
SAM with additive noises. Here, we conduct experiments on the CIFAR dataset with WRN-16-8. We add Gaussian random noises to the image data with variance  $\{0.1, 0.3, 1\}$ . We choose the batch size as 128 and train the model over 200 epochs using a learning rate of 0.1, a momentum of 0.9, and a weight decay of  $5e - 4$ . The SAM hyperparameter is chosen as  $\tau = 2.0$ . As we can see from Table 2, SAM can consistently prevent noise learning and get better performance, compared to the SGD, which varies from different additive noise levels.
A.2 Discussion on the Stochastic Gradient Descent
Here, we include additional empirical results to show how the batch size and step size influence the transition phase of the benign/harmful regimes.
Our synthetic experiment is only performed with gradient descent instead of SGD in Figure 2. We add an experiment of SGD with a mini-batch size of 10 on the synthetic data. If we compare Figure 3 and Figure 2, the comparison results should remain the same, and both support our main Theorems 3.2 and 4.1.
In Figure 4, we also conducted an extended study on the learning rate to check whether tuning the learning rate can get better generalization performance and achieve SAM's performance. Specifically, we experimented with learning rates of 0.001, 0.01, 0.1, and 1 under the same conditions described
(a) SGD Test Error Heatmap with minibatch 10
(b) SAM Test Error Heatmap with minibatch 10
Figure 3: (a) is a heatmap illustrating test error on synthetic data for various dimensions  $d$  and signal strengths  $\mu$  when trained using stochastic gradient descent. High test errors are represented in blue, while low test errors are shown in yellow. (b) displays a heatmap of test errors on the synthetic data under the same conditions as in (a), but trained using SAM instead. The y-axis represents a normal scale with a range of 1000-21000.
in Section 5. The results show that for all learning rates, the harmful and benign overfitting patterns are quite similar and consistent with our Theorem 3.2.
We observed that larger learning rates, such as 0.1 and 1, can improve SGD's generalization performance. The benign overfitting region is enlarged for learning rates of 0.1 and 1 when contrasted with 0.01. This trend resonates with recent findings that large LR helps SGD find better minima (Li et al., 2019, 2021a; Ahn et al., 2022; Andriushchenko et al., 2023). Importantly, even with this expansion, the benign overfitting region remains smaller than what is empirically observed with SAM. Our conclusion is that while SGD, with a more considerable learning rate, exhibits improved generalization, it still falls short of matching SAM's performance.
Lemma B.1 (Lemma B.4 in Kou et al. (2023)). Suppose that  $\delta > 0$  and  $d = \Omega (\log (6n / \delta))$ . Then with probability at least  $1 - \delta$ ,
$$
\sigma_ {p} ^ {2} d / 2 \leq \| \boldsymbol {\xi} _ {i} \| _ {2} ^ {2} \leq 3 \sigma_ {p} ^ {2} d / 2,
$$
$$
\left| \left\langle \boldsymbol {\xi} _ {i}, \boldsymbol {\xi} _ {i ^ {\prime}} \right\rangle \right| \leq 2 \sigma_ {p} ^ {2} \cdot \sqrt {d \log (6 n ^ {2} / \delta)},
$$
$$
| \langle \boldsymbol {\xi} _ {i}, \boldsymbol {\mu} \rangle | \leq \| \boldsymbol {\mu} \| _ {2} \sigma_ {p} \cdot \sqrt {2 \log (6 n / \delta)}
$$
Lemma B.2 (Lemma B.5 in Kou et al. (2023)). Suppose that  $d = \Omega(\log(mn/\delta))$ ,  $m = \Omega(\log(1/\delta))$ . Then with probability at least  $1 - \delta$
$$
\sigma_ {0} ^ {2} d / 2 \leq \| \mathbf {w} _ {j, r} ^ {(0, 0)} \| _ {2} ^ {2} \leq 3 \sigma_ {0} ^ {2} d / 2,
$$
$$
| \langle \mathbf {w} _ {j, r} ^ {(0, 0)}, \boldsymbol {\mu} \rangle | \leq \sqrt {2 \log (1 2 m / \delta)} \cdot \sigma_ {0} \| \boldsymbol {\mu} \| _ {2},
$$
$$
| \langle \mathbf {w} _ {j, r} ^ {(0, 0)}, \boldsymbol {\xi} _ {i} \rangle | \leq 2 \sqrt {\log (1 2 m n / \delta)} \cdot \sigma_ {0} \sigma_ {p} \sqrt {d}
$$
for all  $r \in [m]$ ,  $j \in \{\pm 1\}$  and  $i \in [n]$ . Moreover,
$$
\sigma_ {0} \| \boldsymbol {\mu} \| _ {2} / 2 \leq \max  _ {r \in [ m ]} j \cdot \left\langle \mathbf {w} _ {j, r} ^ {(0, 0)}, \boldsymbol {\mu} \right\rangle \leq \sqrt {2 \log (1 2 m / \delta)} \cdot \sigma_ {0} \| \boldsymbol {\mu} \| _ {2},
$$
$$
\sigma_ {0} \sigma_ {p} \sqrt {d} / 4 \leq \max _ {r \in [ m ]} j \cdot \langle \mathbf {w} _ {j, r} ^ {(0, 0)}, \boldsymbol {\xi} _ {i} \rangle \leq 2 \sqrt {\log (1 2 m n / \delta)} \cdot \sigma_ {0} \sigma_ {p} \sqrt {d}
$$
for all  $j\in \{\pm 1\}$  and  $i\in [n]$
Lemma B.3. Let  $S_{i}^{(t,b)}$  denote  $\{r : \langle \mathbf{w}_{y_i, r}^{(t,b)}, \boldsymbol{\xi}_i \rangle > \sigma_0 \sigma_p \sqrt{d} / \sqrt{2} \rangle\}$ . Suppose that  $\delta > 0$  and  $m \geq 50 \log (2n / \delta)$ . Then with probability at least  $1 - \delta$ ,
$$
\left| S _ {i} ^ {(0, 0)} \right| \geq 0. 8 \Phi (- 1) m, \forall i \in [ n ].
$$
(a) Full-batch GD Test Error Heatmap with learning rate  $\eta = 0.001$
(b) Full-batch GD Test Error Heatmap with learning rate  $\eta = 0.01$
(c) Full-batch GD Test Error Heatmap with learning rate  $\eta = 0.1$
(d) Full-batch GD Test Error Heatmap with learning rate  $\eta = 1$
Figure 4: (a) is a heatmap illustrating test error on synthetic data for various dimensions  $d$  and signal strengths  $\mu$  when trained using learning rate 0.001. High test errors are represented in blue, while low test errors are shown in yellow. (b)(c)(d) displays a heatmap of test errors on the synthetic data under the same conditions as in (a), but trained using GD with different learning rates 0.01, 0.1, 1. The y-axis represents a normal scale with a range of 1000-21000.
Proof of Lemma B.3. Since  $\langle \mathbf{w}_{y_i,r}^{(0,0)},\pmb {\xi}_i\rangle \sim \mathcal{N}(0,\sigma_0^2\| \pmb {\xi}_i\| _2^2)$ , we have
$$
P (\langle \mathbf {w} _ {y _ {i}, r} ^ {(0, 0)}, \pmb {\xi} _ {i} \rangle > \sigma_ {0} \sigma_ {p} \sqrt {d} / \sqrt {2}) \geq P (\langle \mathbf {w} _ {y _ {i}, r} ^ {(0, 0)}, \pmb {\xi} _ {i} \rangle > \sigma_ {0} \| \pmb {\xi} _ {i} \| _ {2}) = \Phi (- 1),
$$
where  $\Phi (\cdot)$  is CDF of the standard normal distribution. Note that  $|S_i^{(0,0)}| = \sum_{r = 1}^{m}\mathbb{1}[\langle \mathbf{w}_{y_i,r}^{(0,0)},\pmb {\xi}_i\rangle >$ $\sigma_0\sigma_p\sqrt{d} /\sqrt{2} ]$  and  $P\big(\langle \mathbf{w}_{y_i,r}^{(0,0)},\pmb {\xi}_i\rangle >\sigma_0\sigma_p\sqrt{d} /\sqrt{2}\big)\geq \Phi (-1)$ , then by Hoeffding's inequality, with probability at least  $1 - \delta /n$ , we have
$$
\frac {\left| S _ {i} ^ {(0 , 0)} \right|}{m} \geq \Phi (- 1) - \sqrt {\frac {\log (2 n / \delta)}{2 m}}.
$$
Therefore, as long as  $0.2\sqrt{m}\Phi(-1) \geq \sqrt{\frac{\log(2n/\delta)}{2}}$ , by applying union bound, with probability at least  $1 - \delta$ , we have
$$
\left| S _ {i} ^ {(0)} \right| \geq 0. 8 \Phi (- 1) m, \forall i \in [ n ].
$$
Lemma B.4. Let  $S_{j,r}^{(t,b)}$  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} / \sqrt{2}\}$ . Suppose that  $\delta > 0$  and  $n \geq 32 \log (4m / \delta)$ . Then with probability at least  $1 - \delta$ ,
$$
\left| S _ {j, r} ^ {(0)} \right| \geq n \Phi (- 1) / 4, \forall j \in \{\pm 1 \}, r \in [ m ].
$$
Proof of Lemma B.4. Since  $\langle \mathbf{w}_{j,r}^{(0,0)},\pmb{\xi}_i\rangle \sim \mathcal{N}(0,\sigma_0^2\| \pmb {\xi}_i\| _2^2)$ , we have
$$
P (\langle \mathbf {w} _ {j, r} ^ {(0, 0)}, \pmb {\xi} _ {i} \rangle > \sigma_ {0} \sigma_ {p} \sqrt {d} / \sqrt {2}) \geq P (\langle \mathbf {w} _ {j, r} ^ {(0, 0)}, \pmb {\xi} _ {i} \rangle > \sigma_ {0} \| \pmb {\xi} _ {i} \| _ {2}) = \Phi (- 1),
$$
where  $\Phi (\cdot)$  is CDF of the standard normal distribution.
Note that  $|S_{j,r}^{(0,0)}| = \sum_{i=1}^{n} \mathbb{1}[y_i = j] \mathbb{1}[\langle \mathbf{w}_{j,r}^{(0)}, \pmb{\xi}_i \rangle > \sigma_0 \sigma_p \sqrt{d} / \sqrt{2}]$  and  $\mathbb{P}(y_i = j, \langle \mathbf{w}_{j,r}^{(0)}, \pmb{\xi}_i \rangle > \sigma_0 \sigma_p \sqrt{d} / \sqrt{2}) \geq \Phi(-1) / 2$ , then by Hoeffding's inequality, with probability at least  $1 - \delta / 2m$ , we have
$$
\frac {\left| S _ {j , r} ^ {(0)} \right|}{n} \geq \Phi (- 1) / 2 + \sqrt {\frac {\log (4 m / \delta)}{2 n}}.
$$
Therefore, as long as  $\Phi(-1)/4 \geq \sqrt{\frac{\log(4m/\delta)}{2n}}$ , by applying union bound, we have with probability at least  $1 - \delta$ ,
Lemma B.5 (Lemma B.3 in Kou et al. (2023)). For  $|S_{+} \cap S_{y}|$  and  $|S_{-} \cap S_{y}|$  where  $y \in \{\pm 1\}$ , it holds with probability at least  $1 - \delta (\delta > 0)$  that
$$
\left| | S _ {+} \cap S _ {y} | - \frac {(1 - p) n}{2} \right| \leq \sqrt {\frac {n}{2} \log \left(\frac {8}{\delta}\right)}, \left| | S _ {-} \cap S _ {y} | - \frac {p n}{2} \right| \leq \sqrt {\frac {n}{2} \log \left(\frac {8}{\delta}\right)}, \forall y \in \{\pm 1 \}.
$$
Lemma B.6. It holds with probability at least  $1 - \delta$ , for all  $T \in \left[\frac{\log(2T^{*} / \delta)}{c_{3}^{2}}, T^{*}\right]$  and  $y \in \{\pm 1\}$ , there exist at least  $c_{3} \cdot T$  epochs among  $[0, T]$ , 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]. \tag {16}
$$
$$
\mathcal {E} _ {1, t} := \left\{\text {I n e p o c h} t, \text {t h e r e a r e a t l e a s t} c _ {2} \cdot \frac {n}{B} \text {b a t c h e s s u c h t h a t 1 6 h o l d s f o r} y = 1 \right\},
$$
$$
\mathcal {E} _ {1, t, b} := \left\{\text {I n e p o c h} t \text {n a t c h} b, 1 6 \text {h o l d s f o r} y = 1 \right\}.
$$
First let  $n$  big enough, then we have  $S_{+} \cap S_{y} \in \left[\frac{3(1 - p)n}{8}, \frac{5(1 - p)n}{8}\right]$ . We consider the first  $c_{1}H$  batches. At the time we are starting to sample  $h$ -th batch in the first  $c_{1}H$  batches, suppose there are  $n_{1}$  samples that belong to  $S_{+} \cap S_{y}$ , and there are  $n_{2}$  samples that don't belong to  $S_{+} \cap S_{y}$ . Then  $n_{1} \geq \frac{3(1 - p)n}{8} - c_{1}n \geq \frac{5(1 - p)n}{16}$  and  $n_{2} \geq \frac{3(1 - p)n}{8} - c_{1}n \geq \frac{5(1 - p)n}{16}$ .
$$
\begin{array}{l} \mathbb {P} (\mathcal {E} _ {1, t, h}) = \frac {\sum_ {l = B / 4} ^ {3 B / 4} C _ {B} ^ {l} C _ {n _ {1}} ^ {l} C _ {n _ {2}} ^ {B - l}}{C _ {n} ^ {B}} \\ \geq \frac {\sum_ {l = 1 B / 4} ^ {3 B / 4} C _ {B} ^ {l} C _ {\frac {5 (1 - p) n}{1 6}} ^ {l} C _ {\frac {5 (1 - p) n}{1 6}} ^ {B - l}}{C _ {n} ^ {B}} \\ \geq \frac {\frac {B}{2} C _ {B} ^ {B / 4} (\frac {9 (1 - p) n}{3 2}) ^ {B} / B !}{n ^ {B} / B !} \\ = \frac {B}{2} C _ {B} ^ {B / 4} \left(\frac {9 (1 - p)}{3 2}\right) ^ {B} := 2 c _ {2}. \\ \end{array}
$$
Then, the probability that there are less than  $c_{1}c_{2}H$  batches in first  $c_{1}H$  batches such that 16 holds is:
$$
\sum_ {i = 0} ^ {c _ {1} c _ {2} H - 1} \sum_ {\sum l _ {h} = i} \mathbb {P} [ \mathbb {1} (\mathcal {E} _ {1, t, 0}) = l _ {0} ] \mathbb {P} [ \mathbb {1} (\mathcal {E} _ {1, t, 1}) = l _ {1} | \mathbb {1} (\mathcal {E} _ {1, t, 0}) = l _ {0} ] \dots
$$
$$
\begin{array}{l} \mathbb {P} \left[ \mathbb {1} \left(\mathcal {E} _ {1, t, c _ {1} H - 1}\right) = l _ {c _ {1} H - 1} \mid \mathbb {1} \left(\mathcal {E} _ {1, t, 0}\right) = l _ {0}, \dots , \mathbb {1} \left(\mathcal {E} _ {1, t, c _ {1} H - 2}\right) = l _ {c _ {1} H - 2} \right] \\ \leq \sum_ {i = 0} ^ {c _ {1} c _ {2} H} C _ {c _ {1} H} ^ {i} (1 - 2 c _ {2}) ^ {c _ {1} H - i} \\ \leq c _ {1} c _ {2} H \cdot \left(2 c _ {2}\right) ^ {c _ {1} c _ {2} H} \left(1 - 2 c _ {2}\right) ^ {c _ {1} H - c _ {1} c _ {2} H}. \\ \end{array}
$$
Choose  $H_0$  such that  $c_{1}c_{2}H_{0}\cdot (2c_{2})^{c_{1}c_{2}H_{0}}(1 - 2c_{2})^{c_{1}H_{0} - c_{1}c_{2}H_{0}} = 1 - 2c_{3}$ , then as long as  $H\geq H_0$  with probability  $c_{3}$ , there are at least  $c_{1}c_{2}H$  batches in first  $c_{1}H$  batches such that 16 holds. Then  $\mathbb{P}[\mathcal{E}_{1,t}]\geq 2c_3$ . Therefore,
$$
\mathbb {P} \left(\sum_ {t} \mathbb {1} \left(\mathcal {E} _ {t}\right) - 2 T c _ {3} \leq - t\right) \leq \exp \left(- \frac {2 t ^ {2}}{T}\right).
$$
Let  $T \geq \frac{\log(2T^{*} / \delta)}{2c_{3}^{2}}$ . Then, with probability at least  $1 - \delta /(2T^{*})$
$$
\sum_ {t} \mathbb {1} \left(\mathcal {E} _ {1, t}\right) \geq c _ {3} T.
$$
Let  $c_{4} = c_{1}c_{2}$ . Thus, there are at least  $c_{3}T^{*}$  epochs, such that they have at least  $c_{4}H$  batches satisfying EquationB.6. This also holds for  $y = -1$ . Taking a union bound to get the result.
C Proofs for SGD
In this section, we prove the results for SGD. We first define some notations. Define  $H = n / B$  as the number of batches within an epoch. For any  $t_1, t_2$  and  $b_1, b_2 \in [\overline{H}]$ , we write  $(t_1, b_1) \leq (t, b) \leq (t_2, b_2)$  to denote all iterations from  $t_1$ -th epoch's  $b_1$ -th batch (included) to  $t_2$ -th epoch's  $b_2$ -th batch (included). And the meanings change accordingly if we replace  $\leq$  with  $<$ .
C.1 Signal-noise Decomposition Coefficient Analysis
This part is dedicated to analyzing the update rule of Signal-noise Decomposition Coefficients. It is worth noting that
$$
F _ {j} (\mathbf {W}, \mathbf {X}) = \frac {1}{m} \sum_ {r = 1} ^ {m} \sum_ {p = 1} ^ {P} \sigma (\langle \mathbf {w} _ {j, r}, \mathbf {x} _ {p} \rangle) = \frac {1}{m} \sum_ {r = 1} ^ {m} \sigma (\langle \mathbf {w} _ {j, r}, \widehat {y} \boldsymbol {\mu} \rangle) + (P - 1) \sigma (\langle \mathbf {w} _ {j, r}, \boldsymbol {\xi} \rangle).
$$
Let  $\mathcal{I}_{t,b}$  denote the set of indices of randomly chosen samples at epoch  $t$  batch  $b$ , and  $|\mathcal{I}_{t,b}| = B$ , then the update rule is:
$$
\begin{array}{l} \text {f o r} b \in \overline {{[ H ]}} \quad \mathbf {w} _ {j, r} ^ {(t, b + 1)} = \mathbf {w} _ {j, r} ^ {(t, b)} - \eta \cdot \nabla_ {\mathbf {w} _ {j, r}} L _ {\mathcal {I} _ {t, b}} \left(\mathbf {W} ^ {(t, b)}\right) \\ = \mathbf {w} _ {j, r} ^ {(t, b)} - \frac {\eta (P - 1)}{B m} \sum_ {i \in \mathcal {I} _ {t, b}} \ell_ {i} ^ {\prime (t, b)} \cdot \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(t, b)}, \boldsymbol {\xi} _ {i} \rangle) \cdot j y _ {i} \boldsymbol {\xi} _ {i} \\ - \frac {\eta}{B m} \sum_ {i \in \mathcal {I} _ {t, b}} \ell_ {i} ^ {(t, b)} \cdot \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(t, b)}, \widehat {y} _ {i} \boldsymbol {\mu} \rangle) \cdot j y _ {i} \widehat {y} _ {i} \boldsymbol {\mu}, \\ \end{array}
$$
$$
\text {a n d} \quad \mathbf {w} _ {j, r} ^ {(t + 1, 0)} = \mathbf {w} _ {j, r} ^ {(t, H)}. \tag {17}
$$
C.1.1 Iterative Expression for Decomposition Coefficient Analysis
Lemma C.1. The coefficients  $\gamma_{j,r}^{(t,b)},\overline{\rho}_{j,r,i}^{(t,b)},\underline{\rho}_{j,r,i}^{(t,b)}$  defined in Lemma 3.3 satisfy the following iterative equations:
$$
\begin{array}{l} \gamma_ {j, r} ^ {(0, 0)}, \bar {\rho} _ {j, r, i} ^ {(0, 0)}, \underline {{\rho}} _ {j, r, i} ^ {(0, 0)} = 0, \tag {18} \\ \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 _ {+}} \ell_ {i} ^ {\prime (t, b)} \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(t, b)}, \widehat {y} _ {i} \cdot \boldsymbol {\mu} \rangle) \right. \\ \end{array}
$$
$$
\left. - \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 {y} _ {i} \cdot \boldsymbol {\mu} \right\rangle\right) \right] \cdot \| \boldsymbol {\mu} \| _ {2} ^ {2}, \tag {19}
$$
$$
\bar {\rho} _ {j, r, i} ^ {(t, b + 1)} = \bar {\rho} _ {j, r, i} ^ {(t, b)} - \frac {\eta (P - 1) ^ {2}}{B m} \cdot \ell_ {i} ^ {\prime (t, b)} \cdot \sigma^ {\prime} \left(\left\langle \mathbf {w} _ {j, r} ^ {(t, b)}, \boldsymbol {\xi} _ {i} \right\rangle\right) \cdot \| \boldsymbol {\xi} _ {i} \| _ {2} ^ {2} \cdot \mathbb {1} \left(y _ {i} = j\right) \mathbb {1} (i \in \mathcal {I} _ {t, b}), \tag {20}
$$
$$
\underline {{\rho}} _ {j, r, i} ^ {(t, b + 1)} = \underline {{\rho}} _ {j, r, i} ^ {(t, b)} + \frac {\eta (P - 1) ^ {2}}{B m} \cdot \ell_ {i} ^ {\prime (t, b)} \cdot \sigma^ {\prime} \left(\left\langle \mathbf {w} _ {j, r} ^ {(t, b)}, \boldsymbol {\xi} _ {i} \right\rangle\right) \cdot \| \boldsymbol {\xi} _ {i} \| _ {2} ^ {2} \cdot \mathbb {1} \left(y _ {i} = - j\right) \mathbb {1} \left(i \in \mathcal {I} _ {t, b}\right), \tag {21}
$$
Proof. First, we iterate the gradient descent update rule  $t$  epochs plus  $b$  batches and get
$$
\begin{array}{l} \mathbf{w}_{j,r}^{(t,b)} = \mathbf{w}_{j,r}^{(0,0)} - \frac{\eta(P - 1)}{Bm}\sum_{(t^{\prime},b^{\prime}) <   (t,b)}\sum_{i\in \mathcal{I}_{t^{\prime},b^{\prime}}}\ell_{i}^{\prime (t^{\prime},b^{\prime})}\cdot \sigma^{\prime}(\langle \mathbf{w}_{j,r}^{(t^{\prime},b^{\prime})},\boldsymbol {\xi}_{i}\rangle)\cdot jy_{i}(P - 1)\boldsymbol{\xi}_{i} \\ - \frac {\eta}{B m} \sum_ {(t ^ {\prime}, b ^ {\prime}) <   (t, b)} \sum_ {i \in \mathcal {I} _ {s, k}} \ell_ {i} ^ {\prime (t ^ {\prime}, b ^ {\prime})} \cdot \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(t ^ {\prime}, b ^ {\prime})}, \widehat {y} _ {i} \boldsymbol {\mu} \rangle) \cdot y _ {i} \widehat {y} _ {i} j \boldsymbol {\mu}. \\ \end{array}
$$
According to the definition of  $\gamma_{j,r}^{(t)}$  and  $\rho_{j,r,i}^{(t)}$ ,
$$
\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}.
$$
Since  $\xi_{i}$  and  $\pmb{\mu}$  are linearly independent with probability 1, we have the unique representation as follows:
$$
\begin{array}{l} \rho_ {j, r, i} ^ {(t, b)} = - \frac {\eta (P - 1) ^ {2}}{B m} \sum_ {(t ^ {\prime}, b ^ {\prime}) <   (t, b)} \ell_ {i} ^ {\prime \left(t ^ {\prime}, b ^ {\prime}\right)} \cdot \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(t ^ {\prime}, b ^ {\prime})}, \boldsymbol {\xi} _ {i} \rangle) \cdot \| \boldsymbol {\xi} _ {i} \| _ {2} ^ {2} \mathbb {1} (i \in \mathcal {I} _ {s, k}) y _ {i} j, \\ \gamma_ {j, r} ^ {(t, b)} = - \frac {\eta}{B m} \sum_ {(t ^ {\prime}, b ^ {\prime}) <   (t, b)} \left[ \sum_ {i \in \mathcal {I} _ {t ^ {\prime}, b ^ {\prime}} \cap S _ {+}} \ell_ {i} ^ {\prime} ^ {(t ^ {\prime}, b ^ {\prime})} \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(t ^ {\prime}, b ^ {\prime})}, y _ {i} \cdot \boldsymbol {\mu} \rangle) \right. \\ \left. - \sum_ {i \in \mathcal {I} _ {t ^ {\prime}, b ^ {\prime}} \cap S _ {-}} \ell_ {i} ^ {\prime \left(t ^ {\prime}, b ^ {\prime}\right)} \sigma^ {\prime} \left(\left\langle \mathbf {w} _ {j, r} ^ {\left(t ^ {\prime}, b ^ {\prime}\right)}, y _ {i} \cdot \boldsymbol {\mu} \right\rangle\right) \right] \| \boldsymbol {\mu} \| _ {2} ^ {2}. \\ \end{array}
$$
Since we define  $\overline{\rho}_{j,r,i}^{(t,b)}\coloneqq \rho_{j,r,i}^{(t,b)}\mathbb{1}(\rho_{j,r,i}^{(t,b)}\geq 0),\underline{\rho}_{j,r,i}^{(t)}\coloneqq \rho_{j,r,i}^{(t,b)}\mathbb{1}(\rho_{j,r,i}^{(t,b)}\leq 0)$ , we obtain
$$
\overline {{\rho}} _ {j, r, i} ^ {(t, b)} = - \frac {\eta (P - 1) ^ {2}}{B m} \sum_ {(t ^ {\prime}, b ^ {\prime}) <   (t, b)} \ell_ {i} ^ {\prime (t ^ {\prime}, b ^ {\prime})} \cdot \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(t ^ {\prime}, b ^ {\prime})}, \boldsymbol {\xi} _ {i} \rangle) \cdot \| \boldsymbol {\xi} _ {i} \| _ {2} ^ {2} \mathbb {1} (i \in \mathcal {I} _ {t ^ {\prime}, b ^ {\prime}})   \mathbb {1} (y _ {i} = j),
$$
$$
\underline {{\rho}} _ {j, r, i} ^ {(t, b)} = \frac {\eta (P - 1) ^ {2}}{B m} \sum_ {(t ^ {\prime}, b ^ {\prime}) <   (t, b)} \ell_ {i} ^ {\prime} (t ^ {\prime}, b ^ {\prime}) \cdot \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(t ^ {\prime}, b ^ {\prime})}, \pmb {\xi} _ {i} \rangle) \cdot \| \pmb {\xi} _ {i} \| _ {2} ^ {2} \mathbb {1} (i \in \mathcal {I} _ {t ^ {\prime}, b ^ {\prime}}) \mathbb {1} (y _ {i} = - j).
$$
And the iterative update equations (19), (20), and (21) follow directly.
C.1.2 Scale of Decomposition Coefficients
We first define  $T^{*} = \eta^{-1}\mathrm{poly}(\epsilon^{-1},d,n,m)$  and
$$
\alpha := 4 \log \left(T ^ {*}\right), \tag {22}
$$
$$
\beta := 2 \max  _ {i, j, r} \left\{\left| \left\langle \mathbf {w} _ {j, r} ^ {(0, 0)}, \boldsymbol {\mu} \right\rangle \right|, (P - 1) \mid \left\langle \mathbf {w} _ {j, r} ^ {(0, 0)}, \boldsymbol {\xi} _ {i} \right\rangle \right| \}, \tag {23}
$$
$$
\operatorname {S N R} := \frac {\| \boldsymbol {\mu} \| _ {2}}{(P - 1) \sigma_ {p} \sqrt {d}}. \tag {24}
$$
By Lemma B.2 and Condition 3.1,  $\beta$  can be bounded as
$$
\beta = 2 \max  _ {i, j, r} \left\{\left| \left\langle \mathbf {w} _ {j, r} ^ {(0, 0)}, \boldsymbol {\mu} \right\rangle \right|, (P - 1) \left| \left\langle \mathbf {w} _ {j, r} ^ {(0, 0)}, \boldsymbol {\xi} _ {i} \right\rangle \right| \right\}
$$
$$
\begin{array}{l} \leq 2 \max  \left\{\sqrt {2 \log (1 2 m / \delta)} \cdot \sigma_ {0} \| \boldsymbol {\mu} \| _ {2}, 2 \sqrt {\log (1 2 m n / \delta)} \cdot \sigma_ {0} (P - 1) \sigma_ {p} \sqrt {d} \right\} \\ = O \left(\sqrt {\log (m n / \delta)} \cdot \sigma_ {0} (P - 1) \sigma_ {p} \sqrt {d}\right), \\ \end{array}
$$
Then, by Condition 3.1, we have the following inequality:
$$
\left. \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 \right\} \leq \frac {1}{1 2}. \right. \tag {25}
$$
We first prove the following bounds for signal-noise decomposition coefficients.
Proposition C.2. Under Assumption 3.1, for  $(0,0) \leq (t,b) \leq (T^{*},0)$ , we have that
$$
\gamma_ {j, r} ^ {(0, 0)}, \bar {\rho} _ {j, r, i} ^ {(0, 0)}, \underline {{\rho}} _ {j, r, i} ^ {(0, 0)} = 0 \tag {26}
$$
$$
0 \leq \bar {\rho} _ {j, r, i} ^ {(t, b)} \leq \alpha , \tag {27}
$$
$$
0 \geq \underline {{\rho}} _ {j, r, i} ^ {(t, b)} \geq - \beta - 1 0 \sqrt {\frac {\log (6 n ^ {2} / \delta)}{d}} n \alpha \geq - \alpha , \tag {28}
$$
and there exists a positive constant  $C'$  such that
$$
- \frac {1}{1 2} \leq \gamma_ {j, r} ^ {(t, b)} \leq C ^ {\prime} \hat {\gamma} \alpha , \tag {29}
$$
for all  $r \in [m]$ ,  $j \in \{\pm 1\}$  and  $i \in [n]$ , where  $\widehat{\gamma} \coloneqq n \cdot \mathrm{SNR}^2$ .
We will prove Proposition C.2 by induction. We first approximate the change of the inner product by corresponding decomposition coefficients when Proposition C.2 holds.
Lemma C.3. Under Assumption 3.1, suppose (27), (28) and (29) hold after  $b$ -th batch of  $t$ -th epoch. 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 {30}
$$
$$
\left| \left\langle \mathbf {w} _ {j, r} ^ {(t, b)} - \mathbf {w} _ {j, r} ^ {(0, 0)}, \boldsymbol {\xi} _ {i} \right\rangle - \frac {1}{P - 1} \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 {31}
$$
$$
\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 {32}
$$
Proof of Lemma C.3. First, for any time  $(t,b)\geq (0,0)$ , we have from the following decomposition by definitions,