SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Proof. Invoking Proposition 2.7.1 in [50], we know
$$
\mathbb {E} \left[ \exp \left(\lambda \mathbf {x} _ {t}\right) \mid \mathcal {F} _ {t} \right] \leq \exp \left(c ^ {- 2} \lambda^ {2} \| \mathbf {x} _ {t} \| _ {\psi_ {1}} ^ {2}\right), \quad \text {i f} | \lambda | \leq \frac {c}{\| \mathbf {x} _ {t} \| _ {\psi_ {1}}}. \tag {A.14}
$$
$$
M _ {t} = \exp \left\{\lambda \sum_ {s = 1} ^ {t} \mathbf {x} _ {s} - c ^ {- 2} \lambda^ {2} \sum_ {s = 1} ^ {t} \| \mathbf {x} _ {s} \| _ {\psi_ {1}} ^ {2} \right\}.
$$
It follows from (A.14) that for any  $t \geq 1$ ,
$$
\mathbb {E} \left[ M _ {t} \mid \mathcal {F} _ {t} \right] = M _ {t - 1} \mathbb {E} \left[ \exp \left\{\lambda \mathbf {x} _ {t} - c ^ {- 2} \lambda^ {2} \| \mathbf {x} _ {t} \| _ {\psi_ {1}} ^ {2} \right\} \mid \mathcal {F} _ {t} \right] \leq M _ {t - 1}.
$$
It further implies that
$$
\mathbb {E} \left[ M _ {t} \right] = \mathbb {E} \left[ \mathbb {E} \left[ M _ {t} \mid \mathcal {F} _ {t} \right] \right] \leq \mathbb {E} \left[ M _ {t - 1} \right] \leq \dots \leq \mathbb {E} \left[ M _ {1} \right] \leq 1.
$$
Using Markov's inequality, for  $0 < \lambda \leq \frac{c}{\max_{1 \leq s \leq t} \| \mathbf{x}_t \|_{\psi_1}}$ , we have
$$
\mathbb {P} \left(\sum_ {s = 1} ^ {t} \mathbf {x} _ {s} \geq c ^ {- 2} \lambda \sum_ {s = 1} ^ {t} \| \mathbf {x} _ {s} \| _ {\psi_ {1}} ^ {2} + \frac {1}{\lambda} \log (1 / \delta)\right) = \mathbb {P} \left(M _ {t} \geq \frac {1}{\delta}\right) \leq \delta . \tag {A.15}
$$
Now taking
$$
\lambda = \min \left\{\sqrt {\frac {c ^ {2} \log (1 / \delta)}{\sum_ {s = 1} ^ {t} \| \mathbf {x} _ {s} \| _ {\psi_ {1}} ^ {2}}}, \frac {c}{\max _ {1 \leq s \leq t} \| \mathbf {x} _ {s} \| _ {\psi_ {1}}} \right\},
$$
and plugging it into (A.15) gives
$$
\mathbb {P} \left(\sum_ {s = 1} ^ {t} \mathbf {x} _ {s} \geq \frac {1}{c} \sqrt {\log (1 / \delta) \sum_ {s = 1} ^ {t} \| \mathbf {x} _ {s} \| _ {\psi_ {1}} ^ {2}} + c \log (1 / \delta) \max  _ {1 \leq s \leq t} \| \mathbf {x} _ {s} \| _ {\psi_ {1}}\right) \leq \delta .
$$
Similarly, we can also prove
$$
\mathbb {P} \left(\sum_ {s = 1} ^ {t} \mathbf {x} _ {s} \leq - \frac {1}{c} \sqrt {\log (1 / \delta) \sum_ {s = 1} ^ {t} \| \mathbf {x} _ {s} \| _ {\psi_ {1}} ^ {2}} - c \log (1 / \delta) \max  _ {1 \leq s \leq t} \| \mathbf {x} _ {s} \| _ {\psi_ {1}}\right) \leq \delta .
$$
Then the proof is finished.
A.3.2 Scales of Sub-exponential Norm for Noises in Stochastic Gradients
For any  $(\pmb{v},\pmb{a})$  with  $\pmb {w} = \pmb {v} / \| \pmb {v}\|$  , we denote
$$
\begin{array}{l} \boldsymbol {\xi} _ {(1)} = \left(\sum_ {i = 1} ^ {k} a _ {i} ^ {*} \sigma \left(\mathbf {Z} _ {i} ^ {\top} \boldsymbol {w}\right)\right) \left(\sum_ {i = 1} ^ {k} a _ {i} \mathbf {Z} _ {i} \mathbb {1} \left\{\mathbf {Z} _ {i} ^ {\top} \boldsymbol {w} \geq 0 \right\}\right) \\ \left. \right. - \mathbb {E} \left[\left(\sum_ {i = 1} ^ {k} a _ {i} ^ {*} \sigma \left(\mathbf {Z} _ {i} ^ {\top} \boldsymbol {v}\right)\right)\left(\sum_ {i = 1} ^ {k} a _ {i} \mathbf {Z} _ {i} \mathbb {1} \left\{\mathbf {Z} _ {i} ^ {\top} \boldsymbol {w} \geq 0 \right\}\right)\right], \\ \end{array}
$$
$$
\begin{array}{l} \boldsymbol {\xi} _ {(2)} = \left(\sum_ {i = 1} ^ {k} a _ {i} ^ {*} \sigma \left(\mathbf {Z} _ {i} ^ {\top} \boldsymbol {w} ^ {*}\right)\right) \left(\sum_ {i = 1} ^ {k} a _ {i} \mathbf {Z} _ {i} \mathbb {1} \left\{\mathbf {Z} _ {i} ^ {\top} \boldsymbol {w} \geq 0 \right\}\right) \\ \left. \right. - \mathbb {E} \left[\left(\sum_ {i = 1} ^ {k} a _ {i} ^ {*} \sigma \left(\mathbf {Z} _ {i} ^ {\top} \boldsymbol {w} ^ {*}\right)\right)\left(\sum_ {i = 1} ^ {k} a _ {i} \mathbf {Z} _ {i} \mathbb {1} \left\{\mathbf {Z} _ {i} ^ {\top} \boldsymbol {w} \geq 0 \right\}\right)\right], \\ \end{array}
$$
$$
\boldsymbol {\epsilon} _ {(1)} = \sigma (\mathbf {Z} \boldsymbol {w}) (\boldsymbol {a}) ^ {\top} \sigma (\mathbf {Z} \boldsymbol {w}) - \mathbb {E} \left[ \sigma (\mathbf {Z} \boldsymbol {w}) (\boldsymbol {a}) ^ {\top} \sigma (\mathbf {Z} \boldsymbol {w}) \right],
$$
$$
\boldsymbol {\epsilon} _ {(2)} = \sigma (\mathbf {Z} \boldsymbol {w}) \left(\boldsymbol {a} ^ {*}\right) ^ {\top} \sigma (\mathbf {Z} \boldsymbol {w} ^ {*}) - \mathbb {E} \left[ \sigma (\mathbf {Z} \boldsymbol {w}) \left(\boldsymbol {a} ^ {*}\right) ^ {\top} \sigma (\mathbf {Z} \boldsymbol {w} ^ {*}) \right].
$$
Lemma 11. Let  $\mathbf{P} = \frac{1}{\|\mathbf{v}\|}\left(\mathbf{I} - \mathbf{w}\mathbf{w}^{\top}\right)$ . The following noise terms satisfy: for  $\ell = 1,2$
(1) For any fixed vector  $\mathbf{u} \in \mathbb{R}^k$ ,  $\mathbf{u}^\top \boldsymbol{\epsilon}_{(\ell)}$  is sub-exponential with  $\max_{\ell=1,2} \| \mathbf{u}^\top \boldsymbol{\epsilon}_{(\ell)} \|_{\psi_1} \leq c \| \mathbf{u} \| \max \{\| \mathbf{a} \|, \| \mathbf{a}^* \|\}$ .
(2) For any fixed vector  $\mathbf{u} \in \mathbb{R}^d$ ,  $\mathbf{u}^\top \mathbf{P}\pmb{\xi}_{(\ell)}$  is sub-exponential with  $\max_{\ell=1,2} \| \mathbf{u}^\top \mathbf{P}\pmb{\xi}_{(\ell)} \|_{\psi_1} \leq c \| \mathbf{a} \| \| \mathbf{a}^* \| / \| \mathbf{v} \|$ .
Proof. Recall the definition of  $\epsilon_{(1)}$ , then we have
$$
\left(\boldsymbol {\epsilon} _ {(1)}\right) ^ {\top} \boldsymbol {u} = \boldsymbol {u} ^ {\top} \sigma (\mathbf {Z} \boldsymbol {w}) (\boldsymbol {a}) ^ {\top} \sigma (\mathbf {Z} \boldsymbol {w}) - \mathbb {E} \left[ \boldsymbol {u} ^ {\top} \sigma (\mathbf {Z} \boldsymbol {w}) (\boldsymbol {a}) ^ {\top} \sigma (\mathbf {Z} \boldsymbol {w}) \right].
$$
Applying Lemma 7 yields
$$
\begin{array}{l} \left. \right.\left\| \boldsymbol {u} ^ {\top} \sigma (\mathbf {Z} \boldsymbol {w}) - \mathbb {E} \left[ \boldsymbol {u} ^ {\top} \sigma (\mathbf {Z} \boldsymbol {w}) \right]\right\| _ {\psi_ {2}} ^ {2} = \left\| \sum_ {j = 1} ^ {k} u _ {j} \left(\sigma \left(\mathbf {Z} _ {j} ^ {\top} \boldsymbol {w}\right) - \mathbb {E} \left[ \sigma \left(\mathbf {Z} _ {j} ^ {\top} \boldsymbol {w}\right)\right]\right)\right\| _ {\psi_ {2}} ^ {2} \\ \stackrel {(i)} {\lesssim} \sum_ {j = 1} ^ {k} \left\| u _ {j} \left(\sigma \left(\mathbf {Z} _ {j} ^ {\top} \boldsymbol {w}\right) - \mathbb {E} \left[ \sigma \left(\mathbf {Z} _ {j} ^ {\top} \boldsymbol {w}\right) \right]\right) \right\| _ {\psi_ {2}} ^ {2} \\ \stackrel {(i i)} {\lesssim} \sum_ {j = 1} ^ {k} \left\| u _ {j} \sigma \left(\mathbf {Z} _ {j} ^ {\top} \boldsymbol {w}\right) \right\| _ {\psi_ {2}} ^ {2} \\ \stackrel {(i i i)} {\lesssim} \| \boldsymbol {u} \| ^ {2} \| \boldsymbol {w} \| ^ {2} \\ \stackrel {(i v)} {=} \| \boldsymbol {u} \| ^ {2}, \tag {A.16} \\ \end{array}
$$
where  $(i)$  and  $(ii)$  comes from Lemma 7 and 8 respectively,  $(iii)$  holds since  $\sigma (\mathbf{Z}_j^\top \mathbf{w})$  is a halfnormal random variable with variance  $\| \pmb {w}\| ^2 /2$ , and  $(iv)$  holds due to  $\| \pmb {w}\| = 1$ . Combining (A.16) and Lemma 8, we can verify that
$$
\left\| \boldsymbol {u} ^ {\top} \sigma (\mathbf {Z} \boldsymbol {w}) \right\| _ {\psi_ {2}} \lesssim \| \boldsymbol {u} \|, \quad \left\| \boldsymbol {a} ^ {\top} \sigma (\mathbf {Z} \boldsymbol {w}) \right\| _ {\psi_ {2}} \lesssim \| \boldsymbol {a} \|, \quad \left\| (\boldsymbol {a} ^ {*}) ^ {\top} \sigma (\mathbf {Z} \boldsymbol {w}) \right\| _ {\psi_ {2}} \lesssim \| \boldsymbol {a} ^ {*} \|.
$$
Applying Lemma 9 gives
$$
\begin{array}{l} \left\| \boldsymbol {u} ^ {\top} \sigma (\mathbf {Z w}) (\boldsymbol {a}) ^ {\top} \sigma (\mathbf {Z w}) \right\| _ {\psi_ {1}} \lesssim \left\| \boldsymbol {u} ^ {\top} \sigma (\mathbf {Z w}) \right\| _ {\psi_ {2}} \left\| \boldsymbol {a} ^ {\top} \sigma (\mathbf {Z w}) \right\| _ {\psi_ {2}} \lesssim \| \boldsymbol {u} \| \| \boldsymbol {a} \| \\ \left\| \boldsymbol {u} ^ {\top} \sigma (\mathbf {Z} \boldsymbol {w}) (\boldsymbol {a} ^ {*}) ^ {\top} \sigma (\mathbf {Z} \boldsymbol {w}) \right\| _ {\psi_ {1}} \lesssim \left\| \boldsymbol {u} ^ {\top} \sigma (\mathbf {Z} \boldsymbol {w}) \right\| _ {\psi_ {2}} \left\| (\boldsymbol {a} ^ {*}) ^ {\top} \sigma (\mathbf {Z} \boldsymbol {w}) \right\| _ {\psi_ {2}} \lesssim \| \boldsymbol {u} \| \| \boldsymbol {a} ^ {*} \|. \\ \end{array}
$$
Invoking Lemma 8, we prove conclusion (1) immediately.
Recall the definitions  $\mathbf{P} = \frac{1}{\|\pmb{v}\|} (\mathbf{I} - \pmb {w}\pmb{w}^{\top})$  . Notice that, for any fixed vector  $\pmb {u}\in \mathbb{R}^{d}$
$$
\begin{array}{l} \boldsymbol {u} ^ {\top} \mathbf {P} \boldsymbol {\xi} _ {(1)} = \left(\sum_ {i = 1} ^ {k} a _ {i} ^ {*} \sigma \left(\mathbf {Z} _ {i} ^ {\top} \boldsymbol {w}\right)\right) \left(\sum_ {i = 1} ^ {k} a _ {i} \boldsymbol {u} ^ {\top} \mathbf {P} \mathbf {Z} _ {i} \mathbb {1} \left\{\mathbf {Z} _ {i} ^ {\top} \boldsymbol {w} \geq 0 \right\}\right) \\ - \mathbb {E} \left[ \left(\sum_ {i = 1} ^ {k} a _ {i} ^ {*} \sigma \left(\mathbf {Z} _ {i} ^ {\top} \boldsymbol {w}\right)\right) \left(\sum_ {i = 1} ^ {k} a _ {i} \boldsymbol {u} ^ {\top} \mathbf {P} \mathbf {Z} _ {i} \mathbb {1} \left\{\mathbf {Z} _ {i} ^ {\top} \boldsymbol {w} \geq 0 \right\}\right) \right]. \\ \end{array}
$$
$$
\begin{array}{l} \left\| \boldsymbol {u} ^ {\top} \mathbf {P} \boldsymbol {\xi} _ {(1)} \right\| _ {\psi_ {1}} ^ {2} \lesssim \left\| \left(\sum_ {i = 1} ^ {k} a _ {i} ^ {*} \sigma (\mathbf {Z} _ {i} ^ {\top} \boldsymbol {w})\right) \left(\sum_ {i = 1} ^ {k} a _ {i} \boldsymbol {u} ^ {\top} \mathbf {P} \mathbf {Z} _ {i} \mathbb {1} \left\{\mathbf {Z} _ {i} ^ {\top} \boldsymbol {w} \geq 0 \right\}\right) \right\| _ {\psi_ {1}} ^ {2} \\ \lesssim \left\| \sum_ {i = 1} ^ {k} a _ {i} ^ {*} \sigma (\mathbf {Z} _ {i} ^ {\top} \boldsymbol {w}) \right\| _ {\psi_ {2}} ^ {2} \left\| \sum_ {i = 1} ^ {k} a _ {i} \boldsymbol {u} ^ {\top} \mathbf {P} \mathbf {Z} _ {i} \mathbb {1} \left\{\mathbf {Z} _ {i} ^ {\top} \boldsymbol {w} \geq 0 \right\} \right\| _ {\psi_ {2}} ^ {2} \\ \lesssim \| \boldsymbol {a} ^ {*} \| ^ {2} \| \boldsymbol {w} \| ^ {2} \| \boldsymbol {a} \| ^ {2} \| \mathbf {P} \boldsymbol {u} \| ^ {2} \\ \leq \| \boldsymbol {a} ^ {*} \| ^ {2} \| \boldsymbol {a} \| ^ {2} / \| \boldsymbol {v} \| ^ {2}, \\ \end{array}
$$
where we used  $\| \pmb{w}\| = 1$  and  $\| \mathbf{P}\| \leq 1 / \| \pmb{v}\|$ . The counterpart for  $\pmb{u}^{\top}\mathbf{P}\pmb{\xi}_{(2)}$  follows from similar arguments.
A.3.3 Corollaries of Lemma 11.
Let  $\{\alpha_{l}\}_{l\geq 0}$  be a sequence of real numbers. Recall that  $\pmb {\xi}_t = \frac{1}{N}\sum_{i = 1}^N\mathbf{P}_t^i\pmb {\xi}_t^i$  and  $\epsilon_{t} = \sum_{i = 1}^{N}\epsilon_{t}^{i} / N$  where  $\{\pmb {\xi}_t^i\}_{i\in [N]}$  and  $\{\pmb {\epsilon}_t^i\}_{i\in [N]}$  are independent random variables given  $\mathcal{F}_t$ . Applying Lemma 10 and Lemma 11(1), we can directly get the following results
$$
\begin{array}{l} \left| \sum_ {l = 0} ^ {s - 1} \alpha_ {l} \boldsymbol {a} _ {l} ^ {\top} \boldsymbol {\epsilon} _ {l} \right| \leq \frac {c}{N} \sqrt {\log (2 / \delta) \sum_ {l = 0} ^ {s - 1} \alpha_ {l} ^ {2} \sum_ {i = 1} ^ {N} \| \boldsymbol {a} _ {l} \| ^ {2} (\| \boldsymbol {a} ^ {*} \| \vee \| \boldsymbol {a} _ {l} ^ {i} \|) ^ {2}} \\ + \frac {c \log (2 / \delta)}{N} \max  _ {l \leq s - 1, i \in [ N ]} \left\{\alpha_ {l} \| \boldsymbol {a} _ {l} \| \left(\| \boldsymbol {a} ^ {*} \| \vee \| \boldsymbol {a} _ {l} ^ {i} \|\right) \right\}, \tag {A.17} \\ \end{array}
$$
$$
\begin{array}{l} \left| \sum_ {l = 0} ^ {s - 1} \alpha_ {l} \mathbf {1} ^ {\top} \boldsymbol {\epsilon} _ {l} \right| \leq \frac {c}{N} \sqrt {k \log (2 / \delta) \sum_ {l = 0} ^ {s - 1} \alpha_ {l} ^ {2} (\| \boldsymbol {a} ^ {*} \| \vee \| \boldsymbol {a} _ {l} \|) ^ {2}} \\ + \frac {c \sqrt {k} \log (2 / \delta)}{N} \max  _ {l \leq s - 1, i \in [ N ]} \left\{\alpha_ {l} \left(\| \boldsymbol {a} ^ {*} \| \vee \| \boldsymbol {a} _ {l} ^ {i} \|\right) \right\}. \tag {A.18} \\ \end{array}
$$
Specially, choosing  $\pmb{u} = \pmb{e}_j$  for  $j\in [k]$ , we also have
$$
\left\| \boldsymbol {\epsilon} _ {s} \right\| ^ {2} \leq k \max  _ {j \in [ k ]} \left| \boldsymbol {e} _ {j} ^ {\top} \boldsymbol {\epsilon} _ {s} \right| ^ {2} \leq \frac {c k \log (k / \delta)}{N} \max  _ {i \in [ N ]} \left\{\left(\left\| \boldsymbol {a} ^ {*} \right\| \vee \left\| \boldsymbol {a} _ {s} ^ {i} \right\|\right) ^ {2} \right\}, \tag {A.19}
$$
hold with probability at least  $1 - \delta$ .
Similarly, using Lemma 11(2), we also have
$$
\left| \sum_ {l = 0} ^ {s - 1} \boldsymbol {\xi} _ {l} ^ {\top} \boldsymbol {w} ^ {*} \right| \leq \frac {c \| \boldsymbol {a} ^ {*} \|}{N \| \boldsymbol {v} \|} \sqrt {\log (2 / \delta) \sum_ {l = 0} ^ {s - 1} \sum_ {i = 1} ^ {N} \| \boldsymbol {a} _ {s} ^ {i} \| ^ {2}} + \frac {c \log (2 / \delta)}{N \| \boldsymbol {v} \|} \| \boldsymbol {a} ^ {*} \| \max  _ {i \in [ N ]} \| \boldsymbol {a} _ {s} ^ {i} \|, \tag {A.20}
$$
$$
\left\| \boldsymbol {\xi} _ {s} \right\| ^ {2} \leq \frac {c d \log (d / \delta)}{N \| \boldsymbol {v} \|} \| \boldsymbol {a} ^ {*} \| ^ {2} \max  _ {i \in [ N ]} \| \boldsymbol {a} _ {s} ^ {i} \| ^ {2}, \tag {A.21}
$$
hold with probability at least  $1 - \delta$
B Bound the Weight's Norm during Training Process
B.1 The Scale of the Second Layer's Weight
Lemma 12 (The scale of the second layer's weight). Suppose the learning rate satisfies that
$$
c \max  \left\{\eta (k \vee I) \log (k / \eta \delta), \sqrt {\frac {\eta k \log (k / \eta \delta)}{N}} \right\} <   1, \tag {B.1}
$$
for a large absolute constant  $c$ . Then for any  $0 \leq t \leq (\eta \| \pmb{a}^* \|^2)^{-2}$ , we are guaranteed that
$$
\max _ {i \in [ N ]} \| \pmb {a} _ {t} ^ {i} \| \leq K _ {a} := 5 \| \pmb {a} ^ {*} \|, A _ {t} \geq - c K _ {a} \sqrt {\frac {\eta \log (1 / \eta \delta)}{N}},
$$
Proof. Recall the definition  $\pmb{a}_t = \frac{1}{N} \sum_{i=1}^N \pmb{a}_t^i$ ,  $\pmb{\epsilon}_t = \frac{1}{N} \sum_{i=1}^N \pmb{\epsilon}_t^i$  and
$$
\nabla_ {\pmb {a}} L (\pmb {w} _ {t} ^ {i}, \pmb {a} _ {t} ^ {i}) = \frac {1}{2 \pi} \left(\mathbf {1 1} ^ {\top} + (\pi - 1) \mathbf {I}\right) \pmb {a} _ {t} ^ {i} - \frac {1}{2 \pi} (\mathbf {1 1} ^ {\top} + (g (\phi_ {t} ^ {i}) - 1) \mathbf {I}) \pmb {a} ^ {*}.
$$
Then notice that
$$
\begin{array}{l} \left\| \frac {1}{N} \sum_ {i = 1} ^ {N} \nabla_ {\boldsymbol {a}} L (\boldsymbol {w} _ {t} ^ {i}, \boldsymbol {a} _ {t} ^ {i}; \boldsymbol {\xi} _ {t} ^ {i}) \right\| ^ {2} \leq 2 \left\| \frac {1}{N} \sum_ {i = 1} ^ {N} \nabla_ {\boldsymbol {a}} L (\boldsymbol {w} _ {t} ^ {i}, \boldsymbol {a} _ {t} ^ {i}) \right\| ^ {2} + 2 \left\| \frac {1}{N} \sum_ {i = 1} ^ {N} \boldsymbol {\epsilon} _ {t} ^ {i} \right\| ^ {2} \\ = 2 \left\| \frac {1}{2 \pi} \left\{(\mathbf {1 1} ^ {\top} + (\pi - 1) \mathbf {I}) \boldsymbol {a} _ {t} - \frac {1}{N} \sum_ {i = 1} ^ {N} (\mathbf {1 1} ^ {\top} + (g (\phi_ {t} ^ {i}) - 1) \mathbf {I}) \boldsymbol {a} ^ {*} \right\} \right\| ^ {2} + 2 \| \boldsymbol {\epsilon} _ {t} \| ^ {2} \\ \leq \frac {(k + \pi - 1) ^ {2}}{\pi^ {2}} \left(\| \boldsymbol {a} _ {t} \| ^ {2} + \| \boldsymbol {a} ^ {*} \| ^ {2}\right) + 2 \| \boldsymbol {\epsilon} _ {t} \| ^ {2}, \\ \end{array}
$$
where the last inequality holds since  $\| \mathbf{11}^{\top} + (\pi - 1)\mathbf{I}\| \leq \pi + k - 1$  and  $g(\phi_t^i) \leq \pi$ . Denote  $A_{t} = (\mathbf{1}^{\top}\mathbf{a}_{t})^{2} - (\mathbf{1}^{\top}\mathbf{a}_{t})(\mathbf{1}^{\top}\mathbf{a}^{*})$ . Together with the definition of  $\mathbf{a}_{t + 1}$ , we have
$$
\begin{array}{l} \left\| \boldsymbol {a} _ {t + 1} \right\| ^ {2} = \left\| \boldsymbol {a} _ {t} - \eta \frac {1}{N} \sum_ {i = 1} ^ {N} \nabla_ {\boldsymbol {a}} L (\boldsymbol {w} _ {t} ^ {i}, \boldsymbol {a} _ {t} ^ {i}; \boldsymbol {\xi} _ {t} ^ {i}) \right\| ^ {2} \\ = \| \boldsymbol {a} _ {t} \| ^ {2} - 2 \eta \frac {1}{N} \sum_ {i \in [ N ]} \left\langle \nabla_ {\boldsymbol {a}} L (\boldsymbol {w} _ {t} ^ {i}, \boldsymbol {a} _ {t} ^ {i}), \boldsymbol {a} _ {t} \right\rangle - 2 \eta \boldsymbol {a} _ {t} ^ {\top} \boldsymbol {\epsilon} _ {t} + \eta^ {2} \left\| \frac {1}{N} \sum_ {i = 1} ^ {N} \nabla_ {\boldsymbol {a}} L (\boldsymbol {w} _ {t} ^ {i}, \boldsymbol {a} _ {t} ^ {i}; \boldsymbol {\xi} _ {t} ^ {i}) \right\| ^ {2} \\ = \| \boldsymbol {a} _ {t} \| ^ {2} - \frac {\eta}{\pi} \boldsymbol {a} _ {t} ^ {\top} \left(\mathbf {1 1} ^ {\top} + (\pi - 1) \mathbf {I}\right) \boldsymbol {a} _ {t} + \frac {\eta}{\pi} \frac {1}{N} \sum_ {i = 1} ^ {N} \boldsymbol {a} _ {t} ^ {\top} \left(\mathbf {1 1} ^ {\top} + \left(g \left(\phi_ {t} ^ {i}\right) - 1\right) \mathbf {I}\right) \boldsymbol {a} ^ {*} \\ - 2 \eta \boldsymbol {a} _ {t} ^ {\top} \boldsymbol {\epsilon} _ {t} + \eta^ {2} \left\| \frac {1}{N} \sum_ {i = 1} ^ {N} \nabla_ {\boldsymbol {a}} L (\boldsymbol {w} _ {t} ^ {i}, \boldsymbol {a} _ {t} ^ {i}; \boldsymbol {\xi} _ {t} ^ {i}) \right\| ^ {2} \\ \leq \left(1 - \frac {\eta (\pi - 1)}{2 \pi}\right) \| \boldsymbol {a} _ {t} \| ^ {2} - \frac {\eta}{\pi} A _ {t} + \frac {\eta (\pi - 1)}{2 \pi} \| \boldsymbol {a} _ {t} \| \| \boldsymbol {a} ^ {*} \| - 2 \eta \boldsymbol {a} _ {t} ^ {\top} \boldsymbol {\epsilon} _ {t} + \eta^ {2} \left\| \frac {1}{N} \sum_ {i = 1} ^ {N} \nabla_ {\boldsymbol {a}} L (\boldsymbol {w} _ {t} ^ {i}, \boldsymbol {a} _ {t} ^ {i}; \boldsymbol {\xi} _ {t} ^ {i}) \right\| ^ {2} \\ \leq \left(1 - \frac {\eta (\pi - 1)}{2 \pi}\right) ^ {t + 1} \| \boldsymbol {a} _ {0} \| ^ {2} + \frac {\eta (\pi - 1)}{2 \pi} \sum_ {s = 0} ^ {t} \left(1 - \frac {\eta (\pi - 1)}{2 \pi}\right) ^ {t - s} \| \boldsymbol {a} _ {s} \| \| \boldsymbol {a} ^ {*} \| \\ - \frac {\eta}{\pi} \sum_ {s = 0} ^ {t} \left(1 - \frac {\eta (\pi - 1)}{2 \pi}\right) ^ {t - s} A _ {s} + 2 \eta \left| \sum_ {s = 0} ^ {t} \left(1 - \frac {\eta (\pi - 1)}{2 \pi}\right) ^ {t - s} \boldsymbol {a} _ {s} ^ {\top} \boldsymbol {\epsilon} _ {s} \right| \\ + 2 \eta^ {2} \sum_ {s = 0} ^ {t} \left(1 - \frac {\eta (\pi - 1)}{2 \pi}\right) ^ {t - s} \left(\frac {(k + \pi - 1) ^ {2}}{\pi^ {2}} \left(\| \boldsymbol {a} _ {s} \| ^ {2} + \| \boldsymbol {a} ^ {*} \| ^ {2}\right) + \| \boldsymbol {\epsilon} _ {s} \| ^ {2}\right). \tag {B.2} \\ \end{array}
$$
Denote the event
$$
\mathcal {E} _ {t} = \left\{\forall s \leq t: \max  _ {i \in [ N ]} \| \boldsymbol {a} _ {s} ^ {i} \| \leq 5 \| \boldsymbol {a} ^ {*} \| \right\}.
$$
Let  $\alpha_{l} = \left(1 - \eta^{\frac{\pi - 1}{2\pi}}\right)^{l}$ , then it holds that  $\sum_{l = 0}^{s - 1}\alpha_l^2\leq \sum_{l = 0}^{s - 1}\alpha_l\leq \frac{2\pi}{\eta(\pi - 1)}$ . Using concentration (A.17), we have
$$
\begin{array}{l} \mathbb {1} _ {\mathcal {E} _ {t}} \eta \left| \sum_ {s = 0} ^ {t} \alpha_ {s} \boldsymbol {a} _ {s} ^ {\top} \boldsymbol {\epsilon} _ {s} \right| \leq \mathbb {1} _ {\mathcal {E} _ {t}} \eta c \sqrt {\frac {\log (1 / \delta)}{N} \sum_ {s = 0} ^ {t} \alpha_ {s} \| \boldsymbol {a} _ {s} \| ^ {2} (\| \boldsymbol {a} ^ {*} \| \vee \max  _ {i} \| \boldsymbol {a} _ {s} ^ {i} \|) ^ {2}} \\ + \mathbb {1} _ {\varepsilon_ {t}} c \frac {\log (1 / \delta)}{N} \max  _ {s \leq t} \alpha_ {s} \| \boldsymbol {a} _ {s} \| (\| \boldsymbol {a} ^ {*} \| \vee \max  _ {i} \| \boldsymbol {a} _ {s} ^ {i} \|) \\ \leq \mathbb {1} _ {\varepsilon_ {t}} \cdot 1 0 c \| \boldsymbol {a} ^ {*} \| ^ {2} \sqrt {\frac {\eta \log (1 / \delta)}{N}}, \tag {B.3} \\ \end{array}
$$
holds with probability at least  $1 - \delta$ . Applying the concentration inequality (A.19) gives
$$
\mathbb {1} _ {\mathcal {E} _ {t}} \eta^ {2} \sum_ {s = 0} ^ {t} \alpha_ {s} \| \boldsymbol {\epsilon} _ {t} \| ^ {2} \leq \mathbb {1} _ {\mathcal {E} _ {t}} \eta^ {2} \sum_ {s = 0} ^ {t} \alpha_ {s} \frac {c k \log (t k / \delta)}{N} (\| \boldsymbol {a} ^ {*} \| \vee \max _ {i} \| \boldsymbol {a} _ {s} ^ {i} \|) ^ {2}
$$
$$
\leq \mathbb {1} _ {\mathcal {E} _ {t}} \cdot \frac {5 0 c k \eta \log (t k / \delta)}{N} \| \boldsymbol {a} ^ {*} \| ^ {2}, \tag {B.4}
$$
holds with probability at least  $1 - \delta$ . In addition, from (A.13), we know
$$
\begin{array}{l} \mathbb {1} _ {\mathcal {E} _ {t}} A _ {s} = \mathbb {1} _ {\mathcal {E} _ {t}} \left\{\left(1 - \eta \frac {\pi + k - 1}{2 \pi}\right) ^ {s} A _ {0} + \frac {\eta \left(\mathbf {1} ^ {\top} \boldsymbol {a} ^ {*}\right) ^ {2}}{2 \pi} \sum_ {l = 0} ^ {s - 1} \left(1 - \eta \frac {\pi + k - 1}{2 \pi}\right) ^ {s - 1 - l} \left(\pi - \frac {1}{N} \sum_ {i = 1} ^ {N} g \left(\phi_ {l} ^ {i}\right)\right) \right. \\ \left. + \eta \sum_ {l = 0} ^ {s - 1} \left(1 - \eta \frac {\pi + k - 1}{2 \pi}\right) ^ {s - 1 - l} \left(\mathbf {1} ^ {\top} \boldsymbol {a} ^ {*}\right) \left(\mathbf {1} ^ {\top} \boldsymbol {\epsilon} _ {l}\right) \right\} \\ \geq \mathbb {1} _ {\mathcal {E} _ {t}} \left\{- \eta \left| \sum_ {l = 0} ^ {s - 1} \left(1 - \eta \frac {\pi + k - 1}{2 \pi}\right) ^ {s - 1 - l} \left(\mathbf {1} ^ {\top} \boldsymbol {a} ^ {*}\right) \left(\mathbf {1} ^ {\top} \boldsymbol {\epsilon} _ {l}\right) \right| \right\} \\ \geq \mathbb {1} _ {\mathcal {E} _ {t}} \left\{- c \sqrt {\frac {\eta k \log (1 / \delta)}{N} \sum_ {l = 0} ^ {s - 1} \left(1 - \eta \frac {\pi + k - 1}{2 \pi}\right) ^ {s - 1 - l} \left(\| \boldsymbol {a} ^ {*} \| \vee \max  _ {i} \| \boldsymbol {a} _ {l} ^ {i} \|\right) ^ {2}} \right\} \\ \geq \mathbb {1} _ {\mathcal {E} _ {t}} \cdot - 5 c \| \boldsymbol {a} ^ {*} \| ^ {2} \sqrt {\frac {\eta \log (1 / \delta)}{N}}, \tag {B.5} \\ \end{array}
$$
holds with probability at least  $1 - \delta$ . Plugging (B.3), (B.4) and (B.5) into (B.2), together with  $t \leq \eta^{-2}$ , we can get
$$
\begin{array}{l} \| \boldsymbol {a} _ {t + 1} \| ^ {2} \mathbb {1} _ {\mathcal {E} _ {t}} \leq \| \boldsymbol {a} _ {0} \| ^ {2} + 5 \| \boldsymbol {a} ^ {*} \| ^ {2} + 5 c \| \boldsymbol {a} ^ {*} \| ^ {2} \sqrt {\frac {\eta k \log (1 / \eta \delta)}{N}} + 1 0 c \| \boldsymbol {a} ^ {*} \| ^ {2} \sqrt {\frac {\eta \log (1 / \eta \delta)}{N}} \\ + 4 \eta \frac {\pi + k - 1}{\pi} \left(2 5 \| \boldsymbol {a} ^ {*} \| ^ {2} + \| \boldsymbol {a} ^ {*} \| ^ {2}\right) + \frac {5 0 c \eta k \log (t / \delta)}{N} \| \boldsymbol {a} ^ {*} \| ^ {2} \\ \leq 9 \left\| \boldsymbol {a} ^ {*} \right\| ^ {2}, \tag {B.6} \\ \end{array}
$$
where the last inequality holds due to our assumption (B.1) and  $\| \pmb{a}_0\| \leq \| \pmb{a}^*\|$ . By setting  $T = t + 1$  and  $K_{a} = 5\| \pmb{a}^{*}\|$  in conclusion (1) of Lemma 15, we have
$$
\begin{array}{l} \mathbb {1} _ {\mathcal {E} _ {t}} \left\| \boldsymbol {a} _ {t + 1} ^ {i} - \boldsymbol {a} _ {t + 1} \right\| \leq \mathbb {1} _ {\mathcal {E} _ {t}} \frac {1}{N} \sum_ {j = 1} ^ {N} \left\| \boldsymbol {a} _ {t + 1} ^ {i} - \boldsymbol {a} _ {t + 1} ^ {j} \right\| \\ \leq \mathbb {1} _ {\varepsilon_ {t}} \cdot 1 0 \| \boldsymbol {a} ^ {*} \| \left(2 \eta I + c \eta \sqrt {k I \log (k / \delta)}\right) \\ \leq \mathbb {1} _ {\mathcal {E} _ {t}} \cdot 4 \| \boldsymbol {a} ^ {*} \|, \tag {B.7} \\ \end{array}
$$
where we used the assumption (B.1). In conjunction with (B.6), we can verify that
$$
\begin{array}{l} \mathbb {P} \left(\mathcal {E} _ {t} \cap \left\{\| \boldsymbol {a} _ {t + 1} \| > 5 \| \boldsymbol {a} ^ {*} \| \right\}\right) \leq \mathbb {P} (\text {O n e o f (B . 3) , (B . 4) , (B . 5) , a n d (B . 7) d o e s n o t h o l d}) \\ \leq 4 \delta . \\ \end{array}
$$
$$
\begin{array}{l} \mathbb {P} \left(\mathcal {E} _ {t + 1}\right) = \mathbb {P} \left(\mathcal {E} _ {t} \cap \left\{\| \boldsymbol {a} _ {t + 1} \| \leq 5 \| \boldsymbol {a} ^ {*} \| \right\}\right) \\ \geq \mathbb {P} \left(\mathcal {E} _ {t}\right) - \mathbb {P} \left(\mathcal {E} _ {t} \cap \left\{\left\| \boldsymbol {a} _ {t + 1} \right\| > 5 \left\| \boldsymbol {a} ^ {*} \right\| \right\}\right) \\ \geq \mathbb {P} \left(\mathcal {E} _ {t}\right) - 4 \delta . \\ \end{array}
$$
Hence by induction, we can verify that  $\mathbb{P}(\mathcal{E}_t) \geq 1 - 4t\delta$  for any  $t \leq (\eta K_a^2)^{-2}$ . By adjusting the level of  $\delta$  and  $\mathbb{P}(\mathcal{E}_0) = 1$  due to  $\| \pmb{a}_0 \| \leq \| \pmb{a}^* \|$ , we can finish the proof.
B.2 The Scale of the First Layer's weight
Lemma 13 (The scale of the first layer's weight). Under the conditions of Lemma 12. Let  $K_{a} = 5\| \pmb{a}^{*}\|$ , we assume the learning rate satisfies
$$
c \eta \left(1 + \sqrt {\frac {I}{d}}\right) \sqrt {d I \log (N d / \delta)} K _ {a} ^ {2} <   1,
$$
for a large absolute constant  $c$ . Then for any  $0 \leq t \leq (800cn^2(\sqrt{Id} + I)\sqrt{\log(Nd / \delta)}K_a^4)^{-1}$ , we are guaranteed that  $1/2 \leq \| \pmb{v}_t \| \leq 3$  with probability at least  $1 - \delta$ .