SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Similar to (D.34), we can guarantee that  $\min_{i\in [N]}(\pmb{a}_s^i)^\top \pmb{a}^* \geq \gamma_a / 64$  holds for any  $s \leq t - 1 \leq \min \{\tau_+(\tilde{\phi}^l), \widetilde{O}(\eta^{-2})\}$  under the event  $\tilde{\mathcal{A}}_t \cap \mathcal{E}_{\mathrm{good}}$ . It follows that
$$
\begin{array}{l} \cos \phi_ {t} = \frac {1}{\| \boldsymbol {v} _ {t} \|} \left\{\boldsymbol {v} _ {0} ^ {\top} \boldsymbol {w} ^ {*} + \eta \sum_ {m = t _ {0}} ^ {\ell - 1} \frac {1}{N} \sum_ {i = 1} ^ {N} \frac {\pi - \phi_ {m} ^ {i}}{2 \pi} \frac {\left(\boldsymbol {a} _ {m} ^ {i}\right) ^ {\top} \boldsymbol {a} ^ {*}}{\| \boldsymbol {v} _ {m} \|} \sin^ {2} \phi_ {m} - \eta \sum_ {m = t _ {0}} ^ {t - 1} \left(\boldsymbol {\xi} _ {m} ^ {\top} \boldsymbol {w} ^ {*} + \boldsymbol {h} _ {m} ^ {\top} \boldsymbol {w} ^ {*}\right) \right\} \\ \geq \frac {1}{\| \boldsymbol {v} _ {t} \|} \left\{\| \boldsymbol {v} _ {0} \| \cos \phi_ {0} + \frac {\eta t}{3} \cdot \frac {\gamma_ {a} \sin^ {2} \tilde {\phi} ^ {u}}{6 4} \cdot \frac {\pi - \tilde {\phi} ^ {u}}{2 \pi} \right. \\ \left. - 2 c \eta K _ {a} ^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} - 2 c K _ {a} \sqrt {\frac {\eta \log (N d / \eta \delta)}{N}} \right\}. \tag {D.39} \\ \end{array}
$$
$$
\tau_ {2, 2} = \frac {9 \cos \tilde {\phi} ^ {l}}{\eta} \frac {3 8 4 \pi}{\gamma_ {a} \sin^ {2} \tilde {\phi} ^ {u}} \frac {1}{\pi - \tilde {\phi} ^ {u}}.
$$
Plugging  $t = t_0 + \tau_{2,2}$  into (D.39), under  $\left\{t_0 + \tau_{2,2} - 1 \leq \tau_+(\tilde{\phi}^l)\right\} \cap \tilde{\mathcal{A}}_{t_0 + \tau_{2,2}} \cap \mathcal{E}_{\mathrm{good}}$ , we can get
$$
\begin{array}{l} \cos \phi_ {t _ {0} + \tau_ {2, 2}} \geq 6 \cos \tilde {\phi} ^ {f} + 3 \cos \tilde {\phi} ^ {l} - 2 c \eta K _ {a} ^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} - 2 c K _ {a} \sqrt {\frac {\eta \log (N d / \eta \delta)}{N}} \\ \geq 2 \cos \tilde {\phi} ^ {l} - 2 c \eta K _ {a} ^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} - 2 c K _ {a} \sqrt {\frac {\eta \log (N d / \eta \delta)}{N}} \geq \cos \tilde {\phi} ^ {l}, \\ \end{array}
$$
where we used the condition (D.6). Invoking (D.38), we can have
$$
\begin{array}{l} \mathbb {P} \left(\tau_ {+} (\tilde {\phi} ^ {l}) \geq t _ {0} + \tau_ {2, 2}\right) = \mathbb {P} \left(\left\{\tau_ {+} (\tilde {\phi} ^ {l}) \geq t _ {0} + \tau_ {2, 2} - 1 \right\} \cap \left\{\cos \phi_ {t _ {0} + \tau_ {2, 2}} \leq \cos \tilde {\phi} ^ {l} \right\}\right) \\ = \mathbb {P} \left(\left\{\tau_ {+} (\tilde {\phi} ^ {l}) \geq t _ {0} + \tau_ {2, 2} - 1 \right\} \cap \tilde {\mathcal {A}} _ {t _ {0} + \tau_ {2, 2}} \cap \left\{\cos \phi_ {t _ {0} + \tau_ {2, 2}} \leq \cos \tilde {\phi} ^ {l} \right\}\right) \\ + \mathbb {P} \left(\left\{\tau_ {+} \left(\tilde {\phi} ^ {l}\right) \geq t _ {0} + \tau_ {2, 2} - 1 \right\} \cap \tilde {\mathcal {A}} _ {t _ {0} + \tau_ {2, 2}} ^ {c} \cap \left\{\cos \phi_ {t _ {0} + \tau_ {2, 2}} \leq \cos \tilde {\phi} ^ {l} \right\}\right) \\ \leq 5 \delta + \mathbb {P} \left(\left\{\tau_ {+} (\tilde {\phi} ^ {l}) \geq t _ {0} + \tau_ {2, 2} - 1 \right\} \cap \tilde {\mathcal {A}} _ {t _ {0} + \tau_ {2, 2}} ^ {c}\right) \\ \leq 1 2 \delta + 5 \delta \cdot \tau_ {2, 2}. \\ \end{array}
$$
Recalling  $t_0 = \tau_+(\tilde{\phi}^f)$  and (D.37), we have
$$
\begin{array}{l} \mathbb {P} \left(\tau_ {+} (\tilde {\phi} ^ {l}) \geq \tau_ {2, 1} + \tau_ {2, 2}\right) \leq \mathbb {P} \left(\left\{\tau_ {+} (\tilde {\phi} ^ {l}) \geq \tau_ {+} (\tilde {\phi} ^ {f}) + \tau_ {2, 2} \right\} \cap \left\{\tau_ {+} (\tilde {\phi} ^ {f}) \leq \tau_ {2, 1} \right\}\right) \\ + \mathbb {P} \left(\left\{\tau_ {+} \left(\tilde {\phi} ^ {l}\right) \geq \tau_ {+} \left(\tilde {\phi} ^ {f}\right) + \tau_ {2, 2} \right\} \cap \left\{\tau_ {+} \left(\tilde {\phi} ^ {f}\right) \geq \tau_ {2, 1} \right\}\right) \\ \end{array}
$$
$$
\begin{array}{l} \leq \mathbb {P} \left(\tau_ {+} (\tilde {\phi} ^ {l}) \geq \tau_ {+} (\tilde {\phi} ^ {f}) + \tau_ {2, 2}\right) + \mathbb {P} \left(\tau_ {+} (\tilde {\phi} ^ {f}) \geq \tau_ {2, 1}\right) \\ \leq 1 7 \delta + 5 \delta \cdot \tau_ {2, 2}. \\ \end{array}
$$
By adjusting the level of  $\delta$ , we can verify that  $\tau_v \leq \tau_{2,1} + \tau_{2,2}$  such that  $\phi_{\tau_v} \leq \tilde{\phi}^l$  and  $\pmb{a}_{\tau_v}^\top \pmb{a}^* \geq \gamma_a / 32$  with high probability. Now recalling that
$$
\gamma_ {a} = \frac {1 6 (\pi - 1) \left(\mathbf {1} ^ {\top} \mathbf {a} ^ {*}\right) ^ {2}}{k (\pi + k - 1)}, \quad \sin^ {2} \tilde {\phi} ^ {u} = \sqrt {1 - \left(\frac {1 4 4 \pi}{\pi + k - 1} \frac {\left(\mathbf {1} ^ {\top} \mathbf {a} ^ {*}\right) ^ {2}}{k} \log (4 + k ^ {2})\right) ^ {2}} \geq \frac {1}{2},
$$
$$
\begin{array}{l} \tau_ {2, 1} + \tau_ {2, 2} = \frac {1}{\eta \gamma_ {a}} \left(\frac {4 0 0 0 0}{\sin^ {3} \phi_ {0}} + \frac {1}{\sin^ {2} \tilde {\phi} ^ {u}} \frac {3 4 6 5 \pi \cos \tilde {\phi} ^ {l}}{\pi - \tilde {\phi} ^ {u}}\right) \\ \lesssim \frac {1}{\eta} \left(\frac {k ^ {2}}{\sin^ {3} \phi_ {0}} + \log (k)\right) \\ \lesssim \frac {1}{\eta} \frac {k ^ {2}}{\sin^ {3} \phi_ {0}}, \\ \end{array}
$$
where the last inequality holds due to (D.8).
E Convergence with Linear Speedup
E.1 Convergence of the First Layer
Lemma 4 restated. Under the settings in Theorem 3. Suppose the initial point satisfies  $\phi_0 \leq \tilde{\phi}^l$  and  $a_0^\top a^* \geq \gamma_a / 32$ . With probability at least  $1 - \delta$ , we can guarantee that  $\sin^2\phi_t \lesssim \epsilon$  holds for any  $\widetilde{O}\left(k^2\eta^{-1}\right) \leq t \leq \widetilde{O}(\eta^{-2})$ .
Proof. Denote the event  $\mathcal{C}_t = \{\pmb{a}_s^\top \pmb{a}^* \geq \gamma_a / 128 \cdot (\| \pmb{a}^*\|^2 \wedge 1) : \forall s \leq t\}$ . For any  $s \leq t - 1$ , with probability at least  $1 - 3\delta$ , we have
$$
\begin{array}{l} \mathbb {1} _ {\mathcal {C} _ {t - 1}} \cos \phi_ {s} = \frac {\mathbb {1} _ {\mathcal {C} _ {t - 1}}}{\| \boldsymbol {v} _ {s} \|} \left\{\boldsymbol {v} _ {0} ^ {\top} \boldsymbol {w} ^ {*} + \eta \sum_ {l = 0} ^ {s - 1} \frac {\pi - \phi_ {l}}{2 \pi} \frac {\boldsymbol {a} _ {l} ^ {\top} \boldsymbol {a} ^ {*}}{\| \boldsymbol {v} _ {l} \|} \sin^ {2} \phi_ {l} - \eta \sum_ {l = 0} ^ {s - 1} (\boldsymbol {\xi} _ {l} ^ {\top} \boldsymbol {w} ^ {*} + \boldsymbol {h} _ {l} ^ {\top} \boldsymbol {w} ^ {*}) \right\} \\ \geq \frac {\mathbb {1} _ {\mathcal {C} _ {t - 1}}}{\| \boldsymbol {v} _ {s} \|} \left\{\frac {\cos \tilde {\phi} ^ {l}}{3} - 2 c \eta K _ {a} ^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} - 2 c K _ {a} \sqrt {\frac {\eta \log (N d / \eta \delta)}{N}} \right\} \\ \geq \mathbb {1} _ {\mathcal {C} _ {t - 1}} \cdot \frac {2 \left(\mathbf {1} ^ {\top} \boldsymbol {a} ^ {*}\right) ^ {2}}{k (\pi + k - 1)}, \tag {E.1} \\ \end{array}
$$
where we used the assumption  $I \lesssim \frac{d}{\epsilon N} \min \left\{1, \frac{k^2 d^{1/2}}{N^{1/2}}, \frac{k^4 d}{N^{2\epsilon}} \right\}$  and  $\epsilon < dk^{-2}$ . Together with (D.4), we can guarantee  $\cos \phi_t^i \geq \frac{(1^\top \mathbf{a}^*)^2}{k(\pi + k - 1)}$  holds for any  $i \in [N]$  with probability at least  $1 - \delta$ . It follows that for any  $s \leq t - 1$ ,
$$
\begin{array}{l} \mathbb {1} _ {\mathcal {C} _ {t - 1}} \frac {1}{N} \sum_ {i = 1} ^ {N} B _ {s} ^ {i} \geq \mathbb {1} _ {\mathcal {C} _ {t - 1}} \left\{\frac {1}{N} \sum_ {i = 1} ^ {N} g \left(\phi_ {s} ^ {i}\right) - 1 \right\} \\ \geq \mathbb {1} _ {\mathcal {C} _ {t - 1}} \left\{g \left(\operatorname {a r c c o s} \left(\frac {\left(\mathbf {1} ^ {\top} \boldsymbol {a} ^ {*}\right) ^ {2}}{k (\pi + k - 1)}\right)\right) - g (\pi / 2) \right\} \\ \stackrel {(i)} {\geq} \mathbb {1} _ {\mathcal {C} _ {t - 1}} \left\{g ^ {\prime} (\pi / 6) \left| \operatorname {a r c c o s} \left(\frac {\left(\mathbf {1} ^ {\top} \boldsymbol {a} ^ {*}\right) ^ {2}}{k (\pi + k - 1)}\right) - \operatorname {a r c c o s} (0) \right| \right\} \\ \geq \mathbb {1} _ {\mathcal {C} _ {t - 1}} \cdot \frac {5 \pi}{1 2} \frac {\left(\mathbf {1} ^ {\top} \mathbf {a} ^ {*}\right) ^ {2}}{k (\pi + k - 1)}, \tag {E.2} \\ \end{array}
$$
where  $(i)$  holds since  $g^{\prime}(x) = -(\pi -x)\sin x$  is decreasing in  $(0,\pi /2]$  and  $\frac{(\mathbf{1}^{\top}\mathbf{a}^{*})^{2}}{k(\pi + k - 1)}\leq \frac{\sqrt{3}}{2}$  due to Assumption 1; and  $(ii)$  holds since  $\arccos^{\prime}(x) = -\frac{1}{\sqrt{1 - x^{2}}}$ . Using the dynamic in Lemma 2, under the event  $\mathcal{C}_{t - 1}$ , we can get
$$
\begin{array}{l} \pmb {a} _ {t} ^ {\top} \pmb {a} ^ {*} \geq \left(1 - \eta \frac {\pi - 1}{2 \pi}\right) ^ {t} \pmb {a} _ {0} ^ {\top} \pmb {a} ^ {*} + \frac {\eta \| \pmb {a} ^ {*} \| ^ {2}}{2 \pi} \sum_ {s = 0} ^ {t - 1} \left(1 - \eta \frac {\pi - 1}{2 \pi}\right) ^ {t - 1 - s} \frac {1}{N} \sum_ {i = 1} ^ {N} B _ {s} ^ {i} - | S (\pmb {\epsilon} _ {0: t - 1}) | \\ \geq \left(1 - \eta \frac {\pi - 1}{2 \pi}\right) ^ {t} \frac {\gamma_ {a}}{3 2} + \left[ 1 - \left(1 - \eta \frac {\pi - 1}{2 \pi}\right) ^ {t} \right] \frac {5 \| \pmb {a} ^ {*} \| ^ {2}}{1 2} \frac {(\pmb {1} ^ {\top} \pmb {a} ^ {*}) ^ {2}}{k (\pi + k - 1)} - 2 c \sqrt {\frac {\eta k \log (1 / \eta \delta)}{N}} K _ {a} ^ {2} \\ \geq \frac {\| \pmb {a} ^ {*} \| ^ {2} \wedge 1}{4} \min \left\{\frac {\gamma_ {a}}{3 2}, \frac {5 (\pmb {1} ^ {\top} \pmb {a} ^ {*}) ^ {2}}{1 2 k (\pi + k - 1)} \right\} - 2 c \sqrt {\frac {\eta k \log (1 / \eta \delta)}{N}} K _ {a} ^ {2} \\ \geq \frac {\gamma_ {a}}{1 2 8} \cdot \left(\| \boldsymbol {a} ^ {*} \| ^ {2} \wedge 1\right). \\ \end{array}
$$
By adjusting the level of  $\delta$ , we can guarantee
$$
\mathbb {P} \left(\mathcal {C} _ {\widetilde {O} (\eta^ {- 2})}\right) = \mathbb {P} \left(\boldsymbol {a} _ {s} ^ {\top} \boldsymbol {a} ^ {*} \geq \frac {\gamma_ {a}}{1 2 8} \cdot \left(\| \boldsymbol {a} ^ {*} \| ^ {2} \wedge 1\right): \forall s \leq t\right) \geq 1 - \delta .
$$
Next we assume  $\mathcal{C}_{\widetilde{O} (\eta^{-2})}$  happens and hide the indicator.
Invoking (E.1), we have for any  $s \leq \widetilde{O}(\eta^{-2})$
$$
\begin{array}{l} \cos \phi_ {s} \lambda_ {s} = \cos \phi_ {s} \frac {1}{N} \sum_ {i = 1} ^ {N} \frac {\pi - \phi_ {t} ^ {i}}{2 \pi} \frac {\left(\boldsymbol {a} _ {t} ^ {i}\right) ^ {\top} \boldsymbol {a} ^ {*}}{\| \boldsymbol {v} _ {t} \| ^ {2}} \geq \frac {2 \left(\mathbf {1} ^ {\top} \boldsymbol {a} ^ {*}\right) ^ {2}}{k (\pi + k - 1)} \cdot \frac {\pi - \pi / 2}{2 \pi} \frac {\gamma_ {a}}{1 2 8} \cdot \left(\| \boldsymbol {a} ^ {*} \| ^ {2} \wedge 1\right) \\ \geq \frac {\gamma_ {a}}{5 1 2} \cdot \left(\| \boldsymbol {a} ^ {*} \| ^ {2} \wedge 1\right) =: \frac {\tilde {\gamma} _ {a}}{1 2}, \tag {E.3} \\ \end{array}
$$
where we used  $\phi_t^i <  \pi /2$  . Invoking Lemma 1, we have
$$
\begin{array}{l} \| \boldsymbol {v} _ {t} \| ^ {2} \sin^ {2} \phi_ {t} \leq \left(1 - \frac {\eta \tilde {\gamma} _ {a}}{1 2}\right) ^ {2} \| \boldsymbol {v} _ {t - 1} \| ^ {2} \sin^ {2} \phi_ {t - 1} - 2 \eta M _ {1, t - 1} + \eta^ {2} M _ {2, t - 1} + H _ {t} \\ \leq \left(1 - \frac {\eta \tilde {\gamma} _ {a}}{1 2}\right) ^ {2 t} \| \boldsymbol {v} _ {0} \| ^ {2} \sin^ {2} \phi_ {0} + 2 \eta \left| \sum_ {s = 0} ^ {t - 1} \left(1 - \frac {\eta \tilde {\gamma} _ {a}}{1 2}\right) ^ {2 (t - 1 - s)} M _ {1, s} \right| \\ + \sum_ {s = 0} ^ {t - 1} \left(1 - \frac {\eta \tilde {\gamma} _ {a}}{1 2}\right) ^ {2 (t - 1 - s)} \left(\eta^ {2} M _ {2, t - 1} + \| \boldsymbol {v} _ {t} \| ^ {2} \sin^ {2} \phi_ {t} - \| \check {\boldsymbol {v}} _ {t} \| ^ {2} \sin^ {2} \check {\phi} _ {t}\right). \tag {E.4} \\ \end{array}
$$
Recall the choices
$$
\eta = \frac {1}{c k ^ {2} \| \boldsymbol {a} ^ {*} \| ^ {2}} \frac {N \epsilon}{d \log (d N / \epsilon \delta)}, \quad T _ {v} = \frac {1}{\eta} \log \left(\frac {1}{\epsilon}\right) \cdot \frac {c k ^ {2}}{\left(\mathbf {1} ^ {\top} \boldsymbol {a} ^ {*}\right) ^ {2}} \frac {1}{\| \boldsymbol {a} ^ {*} \| ^ {2} \wedge 1}. \tag {E.5}
$$
According to (D.35) and (D.26), with probability at least  $1 - \delta$ , we have
$$
\begin{array}{l} \sum_ {s = 0} ^ {t - 1} \left(1 - \frac {\eta \tilde {\gamma} _ {a}}{1 2}\right) ^ {2 (t - 1 - s)} \left| \| \boldsymbol {v} _ {t} \| ^ {2} \sin^ {2} \phi_ {s} - \| \check {\boldsymbol {v}} _ {s} \| ^ {2} \sin^ {2} \check {\phi} _ {s} \right| \\ \leq \sum_ {s = 0} ^ {t - 1} \left(1 - \frac {\eta \tilde {\gamma} _ {a}}{1 2}\right) ^ {2 (t - 1 - s)} \left(6 \eta \sin \phi_ {t} \| \pmb {h} _ {s} \| + \eta^ {2} \| \pmb {h} _ {s} \| ^ {2}\right) \\ \leq \frac {c \eta^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \eta \delta)} K _ {a} ^ {4} + c \eta^ {4} (I d + I ^ {2}) \log (N d / \eta \delta) K _ {a} ^ {8}}{1 - \left(1 - \frac {\eta \tilde {\gamma} _ {a}}{1 2}\right) ^ {2}} \\ \leq \frac {2 4}{\tilde {\gamma} _ {a}} \left\{c \eta (\sqrt {I d} + I) \log^ {1 / 2} (N d / \eta \delta) K _ {a} ^ {4} + c \eta^ {3} (I d + I ^ {2}) \log (N d / \eta \delta) K _ {a} ^ {8} \right\} \\ \leq N \epsilon \left(\sqrt {\frac {I}{d}} + \frac {I}{d}\right) + \frac {N ^ {3} \epsilon^ {3} \left(I ^ {2} + I d\right)}{d ^ {3} k ^ {4}}, \tag {E.6} \\ \end{array}
$$
where we used  $\eta \tilde{\gamma}_a / 12 < 1 / 2$ . Using Lemma 11 and recalling the definition of  $M_{1,t}$ , we have the following concentration
$$
\begin{array}{l} \eta \left| \sum_ {s = 0} ^ {t - 1} \left(1 - \frac {\eta \tilde {\gamma} _ {a}}{1 2}\right) ^ {2 (t - 1 - s)} M _ {1, s} \right| \\ = \eta \left| \sum_ {s = 0} ^ {t - 1} \left(1 - \frac {\eta \tilde {\gamma} _ {a}}{1 2}\right) ^ {2 (t - 1 - s)} \left(\boldsymbol {v} _ {s} - \eta \mathbf {P} _ {s} \frac {1}{N} \sum_ {i = 1} ^ {N} \nabla_ {\boldsymbol {w}} L \left(\boldsymbol {w} _ {s} ^ {i}, \boldsymbol {a} _ {s} ^ {i}\right)\right) ^ {\top} \left(\mathbf {I} - \boldsymbol {w} ^ {*} (\boldsymbol {w} ^ {*}) ^ {\top}\right) \boldsymbol {\xi} _ {s} \right| \\ \leq \frac {1 2 c K _ {a} ^ {2} \eta}{\sqrt {N}} \sqrt {\log (1 / \delta) \sum_ {s = 0} ^ {t - 1} \left(1 - \frac {\eta \tilde {\gamma} _ {a}}{1 2}\right) ^ {2 (t - 1 - s)} \left\| \frac {1}{N} \sum_ {i = 1} ^ {N} \boldsymbol {v} _ {s} ^ {i} - \eta \mathbf {P} _ {s} \nabla_ {\boldsymbol {w}} L (\boldsymbol {w} _ {s} ^ {i} , \boldsymbol {a} _ {s} ^ {i}) \right\| ^ {2}} \\ \leq \frac {c K _ {a} ^ {2} \eta}{\sqrt {N}} \sqrt {\log (1 / \delta) \sum_ {s = 0} ^ {t - 1} \left(1 - \frac {\eta \tilde {\gamma} _ {a}}{1 2}\right) ^ {2 (t - 1 - s)} (9 + \eta^ {2} K _ {a} ^ {2})} \\ \leq \frac {2 c K _ {a} ^ {2} \eta}{\sqrt {N}} \sqrt {\frac {1 0 \log (1 / \delta)}{1 - \left(1 - \frac {\eta \tilde {\gamma} _ {a}}{1 2}\right) ^ {2}}} \leq 1 2 c K _ {a} ^ {2} \sqrt {\frac {\eta \log (1 / \delta)}{N \tilde {\gamma} _ {a}}} \leq \sqrt {\frac {\epsilon}{d}} \tag {E.7} \\ \end{array}
$$
holds with probability at least  $1 - \delta$ . Moreover, it holds that
$$
M _ {2, t} = \boldsymbol {\xi} _ {t} ^ {\top} \left(\mathbf {I} - \boldsymbol {w} ^ {*} (\boldsymbol {w} ^ {*}) ^ {\top}\right) \boldsymbol {\xi} _ {t} \leq \| \boldsymbol {\xi} _ {t} \| ^ {2} \leq \frac {c K _ {a} ^ {4} d \log (d / \eta \delta)}{N},
$$
which yields with probability at least  $1 - \delta$  such that
$$
\eta^ {2} \sum_ {s = 0} ^ {t - 1} \left(1 - \frac {\eta \tilde {\gamma} _ {a}}{1 2}\right) ^ {2 (t - 1 - s)} M _ {2, s} \leq \frac {2 4 \eta}{\tilde {\gamma} _ {a}} \cdot \frac {c K _ {a} ^ {4} d \log (d / \eta \delta)}{N} \leq \epsilon . \tag {E.8}
$$
Plugging (E.6), (E.7) and (E.8) into (E.4), we can guarantee that with probability at least  $1 - 3\delta$ , for any  $T_v \leq t \leq \widetilde{O}(\eta^{-2})$
$$
\begin{array}{l} \sin^ {2} \phi_ {t} \leq \frac {1}{\| \boldsymbol {v} _ {t} \| ^ {2}} \left\{3 \epsilon^ {2} + \sqrt {\frac {\epsilon}{d}} + \left[ N \left(\sqrt {\frac {I}{d}} + \frac {I}{d}\right) + \frac {N ^ {3} \epsilon^ {2} (I ^ {2} + I d)}{d ^ {3} k ^ {4}} + 1 \right] \epsilon \right\} \\ \leq 4 \max  \left\{\sqrt {\frac {\epsilon}{d}} + \frac {(\sqrt {I d} + I) N \epsilon}{d}, \epsilon \right\} \\ = 4 \left(\sqrt {\frac {\epsilon}{d}} + \frac {(\sqrt {I d} + I) N \epsilon}{d}\right), \tag {E.9} \\ \end{array}
$$
where we used  $I \lesssim \frac{d}{\epsilon N} \min \left\{1, \frac{k^2 d^{1/2}}{N^{1/2}}, \frac{k^4 d}{N^{2}\epsilon} \right\}$  and  $\epsilon < d^{-1}$ . Then we further have
$$
\begin{array}{l} \left\| \left(\mathbf {I} - \boldsymbol {w} ^ {*} (\boldsymbol {w} ^ {*}) ^ {\top}\right) \left(\boldsymbol {v} _ {t} - \eta \mathbf {P} _ {t} \frac {1}{N} \sum_ {i = 1} ^ {N} \nabla_ {\boldsymbol {w}} L (\boldsymbol {w} _ {t} ^ {i}, \boldsymbol {a} _ {t} ^ {i})\right) \right\| ^ {2} \\ \leq \boldsymbol {v} _ {t} ^ {\top} \left(\mathbf {I} - \boldsymbol {w} ^ {*} (\boldsymbol {w} ^ {*}) ^ {\top}\right) \boldsymbol {v} _ {t} + \eta^ {2} \frac {1}{N} \sum_ {i = 1} ^ {N} \left\| \mathbf {P} _ {t} \nabla_ {\boldsymbol {w}} L (\boldsymbol {w} _ {t} ^ {i}, \boldsymbol {a} _ {t} ^ {i}) \right\| ^ {2} \\ \leq \| \boldsymbol {v} _ {t} \| ^ {2} \sin^ {2} \phi_ {t} + 4 \eta^ {2} \frac {1}{N} \sum_ {i = 1} ^ {N} \| \nabla_ {\boldsymbol {w}} L (\boldsymbol {w} _ {t} ^ {i}, \boldsymbol {a} _ {t} ^ {i}) \| ^ {2} \\ \leq 3 6 \left(\sqrt {\frac {\epsilon}{d}} + \frac {(\sqrt {I d} + I) N \epsilon}{d}\right) + 4 \eta^ {2} K _ {a} ^ {4} \\ \leq 4 9 \left(\sqrt {\frac {\epsilon}{d}} + \frac {(\sqrt {I d} + I) N \epsilon}{d}\right). \\ \end{array}
$$
With this upper bound, we can improve the concentration (E.7) through
$$
\begin{array}{l} \eta \left| \sum_ {s = 0} ^ {t - 1} \left(1 - \frac {\eta \tilde {\gamma} _ {a}}{1 2}\right) ^ {2 (t - 1 - s)} M _ {1, s} \right| \\ \leq \frac {c K _ {a} ^ {2} \eta}{\sqrt {N}} \sqrt {\log (1 / \delta) \sum_ {s = 0} ^ {t - 1} \left(1 - \frac {\eta \tilde {\gamma} _ {a}}{1 2}\right) ^ {2 (t - 1 - s)} \cdot 4 9 \left(\sqrt {\frac {\epsilon}{d}} + \frac {(\sqrt {I d} + I) N \epsilon}{d}\right)} \\ \leq 4 \sqrt {\frac {\epsilon}{d}} \left(\sqrt {\frac {\epsilon}{d}} + \frac {(\sqrt {I d} + I) N \epsilon}{d}\right) ^ {\frac {1}{2}}. \tag {E.10} \\ \end{array}
$$
Denote  $H = c\eta K_{a}^{4}(\sqrt{Id} + I)\sqrt{\log(Nd / \delta)}$ . Using (E.9), we can further refine the bound (E.6) by
$$
\begin{array}{l} \left| \left\| \boldsymbol {v} _ {t} \right\| ^ {2} \sin^ {2} \phi_ {t} - \left\| \check {\boldsymbol {v}} _ {t} \right\| ^ {2} \sin^ {2} \check {\phi} _ {t} \right| \leq 6 \eta \sin \phi_ {t} \| \boldsymbol {h} _ {t - 1} \| + \eta^ {2} \| \boldsymbol {h} _ {t - 1} \| ^ {2} \\ \leq 4 \left(\sqrt {\frac {\epsilon}{d}} + \frac {(\sqrt {I d} + I) N \epsilon}{d}\right) ^ {\frac {1}{2}} \cdot 6 \eta \| \boldsymbol {h} _ {t - 1} \| + \eta^ {2} \| \boldsymbol {h} _ {t - 1} \| ^ {2} \\ \leq 4 \left(\sqrt {\frac {\epsilon}{d}} + \frac {(\sqrt {I d} + I) N \epsilon}{d}\right) ^ {\frac {1}{2}} \cdot 6 \eta H + \eta^ {2} H ^ {2}. \\ \end{array}
$$
$$
\begin{array}{l} \sum_ {s = 0} ^ {t - 1} \left(1 - \frac {\eta \tilde {\gamma} _ {a}}{1 2}\right) ^ {2 (t - 1 - s)} \left| | | \boldsymbol {v} _ {s} | | ^ {2} \sin^ {2} \phi_ {s} - | | \check {\boldsymbol {v}} _ {s} | | ^ {2} \sin^ {2} \check {\phi} _ {s} \right| \\ \leq 4 \frac {(\sqrt {I d} + I) N \epsilon}{d} \left(\sqrt {\frac {\epsilon}{d}} + \frac {(\sqrt {I d} + I) N \epsilon}{d}\right) ^ {\frac {1}{2}} + \epsilon . \tag {E.11} \\ \end{array}
$$
Combining (E.10), (E.11) and (E.8), we can show that for any  $2T_{v} \leq t \leq (\eta K_{a}^{2})^{-2}$
$$
\sin^ {2} \phi_ {t} \leq 4 \max  \left\{\left(\sqrt {\frac {\epsilon}{d}} + \frac {(\sqrt {I d} + I) N \epsilon}{d}\right) ^ {1 + \frac {1}{2}}, \epsilon \right\}. \tag {E.12}
$$
Repeating (E.10)-(E.12) for  $K = \log(1 / \epsilon)$  times, we can guarantee that for any  $(K + 1)T_v \leq t \leq \widetilde{O}(\eta^{-2})$
$$
\begin{array}{l} \sin^ {2} \phi_ {t} \leq 4 \max  \left\{\left(\sqrt {\frac {\epsilon}{d}} + \frac {(\sqrt {I d} + I) N \epsilon}{d}\right) ^ {1 + \frac {K}{2}}, \epsilon \right\} \\ = 4 \epsilon . \\ \end{array}
$$
E.2 Convergence of the Second Layer
Lemma 5 restated. Under the choice for  $\eta$  and conditions for  $\epsilon$  in Theorem 3. Suppose  $\sin^2\phi_t\leq \epsilon$  holds for any  $0\leq t\leq \widetilde{O} (\eta^{-2})$ . With probability at least  $1 - \delta$ , we can guarantee that  $\| \pmb {a}_t - \pmb {a}^*\| ^2\lesssim \epsilon$  holds for any  $\widetilde{O}\left(\eta^{-1}\right)\leq t\leq \widetilde{O} (\eta^{-2})$ .
Proof. Let  $\pmb{h}_t^i = \pmb{v}_t^i - \pmb{v}_t$ . Similar to (D.35), we can verify that for any  $i \in [N]$
$$
\begin{array}{l} \left| \left\| \boldsymbol {v} _ {t} ^ {i} \right\| ^ {2} \sin^ {2} \phi_ {t} ^ {i} - \left\| \boldsymbol {v} _ {t} \right\| ^ {2} \sin^ {2} \phi_ {t} \right| \leq 2 \eta \sin \phi_ {t} \| \boldsymbol {v} _ {t} \| \| \boldsymbol {h} _ {t - 1} ^ {i} \| + \eta^ {2} \| \boldsymbol {h} _ {t - 1} ^ {i} \| ^ {2} \\ \leq 2 \eta \epsilon \cdot c \eta K _ {a} ^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} + c \eta^ {4} (\sqrt {I d} + I) ^ {2} \log (N d / \delta) \\ \leq 4 \eta \epsilon , \\ \end{array}
$$
where we used Lemma 14 and  $\sin^2\phi_t\leq \epsilon$  . It further implies that
$$
\begin{array}{l} \sin^ {2} \phi_ {t} ^ {i} \leq \frac {\| \boldsymbol {v} _ {t} \| ^ {2} \sin^ {2} \phi_ {t} + \left| \| \boldsymbol {v} _ {t} ^ {i} \| ^ {2} \sin^ {2} \phi_ {t} ^ {i} - \| \boldsymbol {v} _ {t} \| ^ {2} \sin^ {2} \phi_ {t} \right|}{\| \boldsymbol {v} _ {t} ^ {i} \| ^ {2}} \\ \leq \frac {9 \epsilon + 4 \eta \epsilon}{\left(\| \boldsymbol {v} _ {t} \| - \| \boldsymbol {v} _ {t} - \boldsymbol {v} _ {t} ^ {i} \|\right) ^ {2}} \\ \leq 1 0 0 \epsilon . \tag {E.13} \\ \end{array}
$$
Further, for any  $0 \leq t \leq \widetilde{O}(\eta^{-2})$  and  $i \in [N]$  we have
$$
\pi - g \left(\phi_ {t} ^ {i}\right) = g (0) - g \left(\phi_ {t} ^ {i}\right) \stackrel {(i)} {\leq} \left(\pi - \phi_ {t} ^ {i}\right) \sin \phi_ {t} ^ {i} \cdot \phi_ {t} ^ {i} \leq 5 \pi \epsilon^ {1 / 2} \cdot \phi_ {t} ^ {i} \stackrel {(i i)} {\leq} 5 \pi \epsilon , \tag {E.14}
$$
where  $(i)$  holds since  $g^{\prime}(x) = -(\pi - x)\sin x$ ; and  $(ii)$  holds due to  $\arcsin^{\prime}(x) = \frac{1}{\sqrt{1 - x^{2}}}$ . Applying this bound, we can get
$$
\begin{array}{l} \eta \left\langle \frac {1}{N} \sum_ {i = 1} ^ {N} \nabla_ {\boldsymbol {a}} L \left(\boldsymbol {w} _ {t} ^ {i}, \boldsymbol {a} _ {t} ^ {i}\right), \boldsymbol {a} _ {t} - \boldsymbol {a} ^ {*} \right\rangle = \frac {\eta}{2 \pi} \left(\boldsymbol {a} _ {t} - \boldsymbol {a} ^ {*}\right) ^ {\top} \left(\mathbf {1} \mathbf {1} ^ {\top} + (\pi - 1) \mathbf {I}\right) \left(\boldsymbol {a} _ {t} - \boldsymbol {a} ^ {*}\right) \\ + \frac {\eta}{2 \pi} \left(\pi - \frac {1}{N} \sum_ {i = 1} ^ {N} g \left(\phi_ {t} ^ {i}\right)\right) \left(\boldsymbol {a} _ {t} - \boldsymbol {a} ^ {*}\right) ^ {\top} \boldsymbol {a} ^ {*} \\ \geq \frac {\eta (\pi - 1)}{2 \pi} \| \boldsymbol {a} _ {t} - \boldsymbol {a} ^ {*} \| ^ {2} - 3 \eta \epsilon \| \boldsymbol {a} _ {t} - \boldsymbol {a} ^ {*} \| \| \boldsymbol {a} ^ {*} \| \\ \geq \frac {\eta (\pi - 1)}{2 \pi} \| \boldsymbol {a} _ {t} - \boldsymbol {a} ^ {*} \| ^ {2} - 6 \eta \epsilon K _ {a} ^ {2}, \tag {E.15} \\ \end{array}
$$
where the last inequality holds due to Lemma 12 such that  $\| \pmb{a}_t - \pmb{a}^*\| \leq 6\| \pmb{a}^*\|$ . In addition, it holds that
$$
\begin{array}{l} \left\| \frac {1}{N} \sum_ {i = 1} ^ {N} \nabla_ {\boldsymbol {a}} L (\boldsymbol {w} _ {t} ^ {i}, \boldsymbol {a} _ {t} ^ {i}) \right\| ^ {2} = \frac {1}{4 \pi^ {2}} \left\| \left(\mathbf {1 1} ^ {\top} + (\pi - 1) \mathbf {I}\right) (\boldsymbol {a} _ {t} - \boldsymbol {a} ^ {*}) + \left(\frac {1}{N} \sum_ {i = 1} ^ {N} g (\phi_ {t} ^ {i}) - \pi\right) \boldsymbol {a} ^ {*} \right\| ^ {2} \\ \leq \frac {(\pi + k - 1) ^ {2}}{2 \pi^ {2}} \| \boldsymbol {a} _ {t} - \boldsymbol {a} ^ {*} \| ^ {2} + 1 5 \epsilon^ {2} \| \boldsymbol {a} ^ {*} \| ^ {2}, \tag {E.16} \\ \end{array}
$$
where we used the relation (E.14). Denote  $\tilde{\pmb{a}}_t = \pmb{a}_t - \pmb{a}^* - \eta \frac{1}{N}\sum_{i=1}^{N}\nabla_{\pmb{a}}L(\pmb{w}_t^i, \pmb{a}_t^i)$ , then we have