SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \left\| \boldsymbol {h} _ {t} ^ {i, j} \right\| \leq \frac {\left| \phi_ {t} ^ {j} - \phi_ {t} ^ {i} \right|}{2 \pi} \frac {\left(\boldsymbol {a} _ {t} ^ {i}\right) ^ {\top} \boldsymbol {a} ^ {*}}{\left\| \boldsymbol {v} _ {t} ^ {i} \right\|} + \frac {\pi - \phi_ {t} ^ {j}}{2 \pi} \frac {\left| \left(\boldsymbol {a} _ {t} ^ {i} - \boldsymbol {a} _ {t} ^ {j}\right) ^ {\top} \boldsymbol {a} ^ {*} \right|}{\left\| \boldsymbol {v} _ {t} ^ {i} \right\|} + \frac {\pi - \phi_ {t} ^ {j}}{2 \pi} \left(\boldsymbol {a} _ {t} ^ {j}\right) ^ {\top} \boldsymbol {a} ^ {*} \left| \frac {1}{\left\| \boldsymbol {v} _ {t} ^ {i} \right\|} - \frac {1}{\left\| \boldsymbol {v} _ {t} ^ {j}\right) \|} \right| \\ + \frac {\pi - \phi_ {t} ^ {i}}{\pi} \frac {\left(\boldsymbol {a} _ {t} ^ {j}\right) ^ {\top} \boldsymbol {a} ^ {*}}{\| \boldsymbol {v} _ {t} ^ {j} \|} \left| \frac {1}{\| \boldsymbol {v} _ {t} ^ {j} \|} - \frac {1}{\| \boldsymbol {v} _ {t} ^ {j} \|} \right| \\ \leq 3 K _ {a} ^ {2} + \left| \left(\boldsymbol {a} _ {t} ^ {i} - \boldsymbol {a} _ {t} ^ {j}\right) ^ {\top} \boldsymbol {a} ^ {*} \right| + 6 K _ {a} ^ {2} \left\| \boldsymbol {v} _ {t} ^ {i} - \boldsymbol {v} _ {t} ^ {j} \right\| + 5 K _ {a} ^ {2} \left(\frac {\left\| \boldsymbol {v} _ {t} ^ {i} - \boldsymbol {v} _ {t} ^ {j} \right\|}{\left\| \boldsymbol {v} _ {t} ^ {i} \right\|} + \frac {\left\| \boldsymbol {v} _ {t} ^ {i} - \boldsymbol {v} _ {t} ^ {j} \right\|}{2 \left\| \boldsymbol {v} _ {t} ^ {i} \right\| \left\| \boldsymbol {v} _ {t} ^ {j} \right\|}\right) \\ \leq 3 K _ {a} ^ {2} + \left| \left(\boldsymbol {a} _ {t} ^ {i} - \boldsymbol {a} _ {t} ^ {j}\right) ^ {\top} \boldsymbol {a} ^ {*} \right| + 1 0 K _ {a} ^ {2} \| \boldsymbol {v} _ {t} ^ {i} - \boldsymbol {v} _ {t} ^ {j} \|, \tag {C.2} \\ \end{array}
$$
where we also used the fact  $\max_{i\in [N]}\| \mathbf{I} - \pmb{w}_t^i (\pmb{w}_t^i)^\top \| \leq 1$  and
$$
\begin{array}{l} \left\| \boldsymbol {w} _ {t} ^ {i} \left(\boldsymbol {w} _ {t} ^ {i}\right) ^ {\top} \boldsymbol {w} ^ {*} - \boldsymbol {w} _ {t} ^ {j} \left(\boldsymbol {w} _ {t} ^ {j}\right) ^ {\top} \boldsymbol {w} ^ {*} \right\| \leq \left\| \boldsymbol {w} _ {t} ^ {i} \left(\boldsymbol {w} _ {t} ^ {i} - \boldsymbol {w} _ {t} ^ {j}\right) ^ {\top} \boldsymbol {w} ^ {*} \right\| + \left\| \left(\boldsymbol {w} _ {t} ^ {i} - \boldsymbol {w} _ {t} ^ {j}\right) \left(\boldsymbol {w} _ {t} ^ {j}\right) ^ {\top} \boldsymbol {w} ^ {*} \right\| \\ \leq 2 \left\| \boldsymbol {w} _ {t} ^ {i} - \boldsymbol {w} _ {t} ^ {j} \right\| \\ = 2 \left| \frac {1}{\| \boldsymbol {v} _ {t} ^ {i} \|} - \frac {1}{\| \boldsymbol {v} _ {t} ^ {j} \|} \right|. \\ \end{array}
$$
From the definition of  $B_s^i$  and  $B_s^j$  in Lemma 2, we know for any  $i,j\in [N]$ ,
$$
\begin{array}{l} \left| B _ {s} ^ {i} - B _ {s} ^ {j} \right| \leq \left| g (\phi_ {s} ^ {j}) - g (\phi_ {s} ^ {i}) \right| + \frac {\eta}{2 \pi} \frac {(\mathbf {1} ^ {\top} \mathbf {a} ^ {*}) ^ {2}}{\| \mathbf {a} ^ {*} \| ^ {2}} \sum_ {l = t _ {r}} ^ {s - 1} \left(1 - \eta \frac {\pi + k - 1}{2 \pi}\right) ^ {s - 1 - l} \left| g (\phi_ {l} ^ {j}) - g (\phi_ {l} ^ {i}) \right| \\ \leq \pi + \frac {\eta I}{2} \frac {\left(\mathbf {1} ^ {\top} \mathbf {a} ^ {*}\right) ^ {2}}{\| \mathbf {a} ^ {*} \| ^ {2}}, \\ \end{array}
$$
where the last inequality holds due to  $g(\phi) \in [0, \pi]$  for  $\phi \in [0, \pi]$ . Using concentration (A.18) and  $t \leq k\eta^{-1}$ , with probability at least  $1 - \delta$ , we can guarantee
$$
\begin{array}{l} | S (\pmb {\epsilon} _ {t _ {r}: s - 1}) | = \frac {\eta^ {2} | \mathbf {1} ^ {\top} \pmb {a} ^ {*} |}{2 \pi} \left| \sum_ {l = t _ {r}} ^ {s - 1} \left(1 - \frac {2 \eta (\pi - 1)}{\pi}\right) ^ {s - l - 1} \sum_ {m = t _ {r}} ^ {l - 1} \left(1 - \eta \frac {\pi + k - 1}{2 \pi}\right) ^ {l - 1 - m} (\mathbf {1} ^ {\top} \pmb {\epsilon} _ {m}) \right| \\ + \eta \left| \sum_ {j = t _ {r}} ^ {s - 1} \left(1 - \frac {2 \eta (\pi - 1)}{\pi}\right) ^ {s - 1 - j} \boldsymbol {\epsilon} _ {j} ^ {\top} \boldsymbol {a} ^ {*} \right| \\ \leq c \eta I \sqrt {\frac {\log (2 t / \delta)}{N}} K _ {a} ^ {2} \cdot \left(\eta \frac {\left| \mathbf {1} ^ {\top} \boldsymbol {a} ^ {*} \right|}{\left\| \boldsymbol {a} ^ {*} \right\|} \sum_ {l = t _ {r}} ^ {s - 1} \left(1 - \frac {2 \eta (\pi - 1)}{\pi}\right) ^ {s - l - 1} + 1\right) \\ \leq 2 c \eta I \sqrt {\frac {k \log (2 t / \delta)}{N}} K _ {a} ^ {2}. \tag {C.3} \\ \end{array}
$$
Using the dynamic of  $\pmb{a}_t^\top \pmb{a}^*$  in Lemma 2, with probability at least  $1 - \delta$ , we also have
$$
\begin{array}{l} \left| \left(\boldsymbol {a} _ {t - 1} ^ {j} - \boldsymbol {a} _ {t - 1} ^ {i}\right) ^ {\top} \boldsymbol {a} ^ {*} \right| = \left| \frac {\eta \| \boldsymbol {a} ^ {*} \| ^ {2}}{2 \pi} \sum_ {s = t _ {r}} ^ {t - 2} \left(1 - \eta \frac {\pi - 1}{2 \pi}\right) ^ {t - 2 - s} \left(B _ {s} ^ {j} - B _ {s} ^ {i}\right) + S (\boldsymbol {\epsilon} _ {t _ {r}: t - 2} ^ {j}) - S (\boldsymbol {\epsilon} _ {t _ {r}: t - 2} ^ {i}) \right| \\ \leq \frac {\eta I \| \boldsymbol {a} ^ {*} \| ^ {2}}{2 \pi} \left(\pi + \frac {\eta I}{2} \frac {\left(\mathbf {1} ^ {\top} \boldsymbol {a} ^ {*}\right) ^ {2}}{\| \boldsymbol {a} ^ {*} \| ^ {2}}\right) + \left| S \left(\boldsymbol {\epsilon} _ {t _ {r}: t - 2} ^ {j}\right) - S \left(\boldsymbol {\epsilon} _ {t _ {r}: t - 2} ^ {i}\right) \right| \\ \leq \frac {\eta I}{2 \pi} \left(\pi \| \boldsymbol {a} ^ {*} \| ^ {2} + \frac {\eta I}{2} \left(\mathbf {1} ^ {\top} \boldsymbol {a} ^ {*}\right) ^ {2}\right) + 2 c \eta \sqrt {k I \log (t / \delta)} K _ {a} ^ {2} \\ \leq 2 \eta (\sqrt {k I} + I) \sqrt {\log (1 / \eta \delta)} K _ {a} ^ {2}, \tag {C.4} \\ \end{array}
$$
where the second last inequality holds due to the concentration (C.3). Plugging (C.4) into (C.2) gives
$$
\begin{array}{l} \| \boldsymbol {v} _ {t} ^ {i} - \boldsymbol {v} _ {t} ^ {j} \| \leq \eta \left\| \sum_ {s = t _ {r}} ^ {t - 1} \boldsymbol {h} _ {s} ^ {i, j} \right\| + \eta \left\| \sum_ {s = t _ {r}} ^ {t - 1} \mathbf {P} _ {s} ^ {i} \boldsymbol {\xi} _ {s} ^ {i} - \mathbf {P} _ {s} ^ {j} \boldsymbol {\xi} _ {s} ^ {j} \right\| \\ \leq \eta \sum_ {s = t _ {r}} ^ {t - 1} \left(3 K _ {a} ^ {2} + | (\boldsymbol {a} _ {s} ^ {i} - \boldsymbol {a} _ {s} ^ {j}) ^ {\top} \boldsymbol {a} ^ {*} | + 1 0 K _ {a} ^ {2} \| \boldsymbol {v} _ {s} ^ {i} - \boldsymbol {v} _ {s} ^ {j} \|\right) + \eta \left\| \sum_ {s = t _ {r}} ^ {t - 1} \mathbf {P} _ {s} ^ {i} \boldsymbol {\xi} _ {s} ^ {i} \right\| + \eta \left\| \sum_ {s = t _ {r}} ^ {t - 1} \mathbf {P} _ {s} ^ {j} \boldsymbol {\xi} _ {s} ^ {j} \right\| \\ \leq \eta I \left(3 K _ {a} ^ {2} + 2 \eta (I + \sqrt {k I}) K _ {a} ^ {2} + 1 0 K _ {a} ^ {2} O _ {r, T} ^ {i, j}\right) + 4 c \eta \sqrt {d I \log (N / \delta)} K _ {a} ^ {2}, \tag {C.5} \\ \end{array}
$$
where we used the following concentration
$$
\max  _ {i} \left\| \sum_ {s = t _ {r}} ^ {t - 1} \mathbf {P} _ {s} ^ {i} \boldsymbol {\xi} _ {s} ^ {i} \right\| \leq \sqrt {d} \max  _ {i \in [ N ], \ell \in [ d ]} \left| \left(\mathbf {P} _ {s} ^ {i} \boldsymbol {\xi} _ {s} ^ {i}\right) _ {\ell} \right| \leq c \sqrt {d I \log (d N / \delta)} K _ {a} ^ {2}.
$$
Taking maximum on the left hand side of (C.5) over  $t_r \leq t \leq T$ , we can get
$$
\begin{array}{l} O _ {r, T} ^ {i, j} \leq \left(1 - 1 0 \eta I K _ {a} ^ {2}\right) \left[ \eta I \left(3 K _ {a} ^ {2} + 2 \eta (I + \sqrt {k I}) K _ {a} ^ {2}\right) + 4 c \eta \sqrt {d I \log (N / \delta)} K _ {a} ^ {2} \right] \\ \leq 2 \left(4 \sqrt {\frac {I}{d \log (N d / \delta)}} + 4 c\right) \eta \sqrt {d I \log (N d / \delta)} K _ {a} ^ {2} \\ \leq 2 \left(5 c + 4 \sqrt {\frac {I}{d}}\right) \eta \sqrt {d I \log (N d / \delta)} K _ {a} ^ {2}, \\ \end{array}
$$
where we also used the assumption  $\eta (I + \sqrt{kI}) < 1$ .
C.2 The Discrepancy in the Second Layer
Lemma 15 (The discrepancy in the second layer). Let  $\Xi_{r,T}^{i,j} = \max_{t_r\leq t\leq T}\| \pmb {a}_t^i -\pmb {a}_t^j\|$  for  $t_r\leq T\leq t_{r + 1} - 1$  . Suppose the learning rate satisfies the conditions in Lemmas 12 and 13. If  $\max_{i\in [N],t_r\leq t\leq T - 1}\| \pmb {a}_t^i\| \leq K_a$  holds for  $K_{a} > \| \pmb{a}^{*}\|$  and  $c\eta kI < 1$  , with probability at least  $1 - \delta$  , we have
$$
\max  _ {i, j} \Xi_ {r, T} ^ {i, j} \leq 2 K _ {a} \left(2 \eta I + c \eta \sqrt {k I \log (k / \delta)}\right).
$$
Proof. Recall the closed form of  $\nabla_{\pmb{a}}L(\pmb{w},\pmb{a})$ , then for  $i,j\in [N]$ , we have
$$
\nabla_ {\boldsymbol {a}} L (\boldsymbol {w} _ {t} ^ {i}, \boldsymbol {a} _ {t} ^ {i}) - \nabla_ {\boldsymbol {a}} L (\boldsymbol {w} _ {t} ^ {j}, \boldsymbol {a} _ {t} ^ {j}) = \frac {\eta}{2 \pi} \left(\mathbf {1 1} ^ {\top} + (\pi - 1) \mathbf {I}\right) (\boldsymbol {a} _ {t} ^ {i} - \boldsymbol {a} _ {t} ^ {j}) - \frac {\eta}{2 \pi} \left(g (\phi_ {t} ^ {i}) - g (\phi_ {t} ^ {j})\right) \boldsymbol {a} ^ {*}.
$$
Let  $t_r$  be the previous synchronization time before  $t$ . Using the fact  $\boldsymbol{a}_{t_r}^i = \boldsymbol{a}_{t_r}^j = \boldsymbol{a}_{t_r}$ , we have
$$
\left\| \boldsymbol {a} _ {t} ^ {i} - \boldsymbol {a} _ {t} ^ {j} \right\| \leq \eta \left\| \sum_ {s = t _ {r}} ^ {t - 1} \nabla_ {\boldsymbol {a}} L (\boldsymbol {w} _ {s} ^ {i}, \boldsymbol {a} _ {s} ^ {i}) - \nabla_ {\boldsymbol {a}} L (\boldsymbol {w} _ {s} ^ {j}, \boldsymbol {a} _ {s} ^ {j}) \right\| + \eta \left\| \sum_ {s = t _ {r}} ^ {t - 1} \boldsymbol {\epsilon} _ {s} ^ {i} - \boldsymbol {\epsilon} _ {s} ^ {j} \right\|
$$
$$
\begin{array}{l} \leq \frac {\eta (\pi + k - 1)}{2 \pi} \sum_ {s = t _ {r}} ^ {t - 1} \| \boldsymbol {a} _ {s} ^ {i} - \boldsymbol {a} _ {s} ^ {j} \| + \frac {\eta \| \boldsymbol {a} ^ {*} \|}{2 \pi} \sum_ {s = t _ {r}} ^ {t - 1} \left| g \left(\phi_ {t - 1} ^ {i}\right) - g \left(\phi_ {t - 1} ^ {j}\right) \right| + \eta \left\| \sum_ {s = t _ {r}} ^ {t - 1} \boldsymbol {\epsilon} _ {s} ^ {i} - \boldsymbol {\epsilon} _ {s} ^ {j} \right\| \\ \leq \frac {\eta (\pi + k - 1)}{2 \pi} \sum_ {s = t _ {r}} ^ {t - 1} \| \boldsymbol {a} _ {s} ^ {i} - \boldsymbol {a} _ {s} ^ {j} \| + \frac {\eta \| \boldsymbol {a} ^ {*} \|}{2 \pi} \pi I + c \sqrt {k I \log (k / \delta)} K _ {a}, \tag {C.6} \\ \end{array}
$$
where the last inequality holds due to Lemma 11(1). Taking maximum on the both sides over  $t_r \leq t \leq T$ , we get
$$
\begin{array}{l} \Xi_ {r, T} ^ {i, j} \leq \left(1 - \frac {\eta (\pi + k - 1) I}{2 \pi}\right) ^ {- 1} \left[ \frac {\pi \eta I}{2} \| \boldsymbol {a} ^ {*} \| + c \eta \sqrt {k I \log (k / \delta)} K _ {a} \right] \\ \leq 2 K _ {a} \left(2 \eta I + c \eta \sqrt {k I \log (k / \delta)}\right). \\ \end{array}
$$
D Self-correction of Signals
Recalling the dynamic in Lemma 2,
$$
\begin{array}{l} \boldsymbol {a} _ {t} ^ {\top} \boldsymbol {a} ^ {*} = \left(1 - \eta \frac {\pi - 1}{2 \pi}\right) ^ {t} \boldsymbol {a} _ {0} ^ {\top} \boldsymbol {a} ^ {*} + \frac {1 - \left(\frac {2 \pi - \eta (k + \pi - 1)}{2 \pi - \eta (\pi - 1)}\right) ^ {t}}{k} \left(1 - \eta \frac {\pi - 1}{2 \pi}\right) ^ {t} A _ {0} \\ + \frac {\eta \| \boldsymbol {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 \left(\epsilon_ {0: t - 1}\right), \tag {D.1} \\ \end{array}
$$
$$
B _ {s} ^ {i} = g \left(\phi_ {s} ^ {i}\right) - 1 + \frac {\eta}{2 \pi} \frac {\left(\mathbf {1} ^ {\top} \mathbf {a} ^ {*}\right) ^ {2}}{\| \mathbf {a} ^ {*} \| ^ {2}} \sum_ {l = 0} ^ {s - 1} \left(1 - \eta \frac {\pi + k - 1}{2 \pi}\right) ^ {s - 1 - l} (\pi - g \left(\phi_ {l} ^ {i}\right)). \tag {D.2}
$$
Using concentration (A.18), with probability at least  $1 - \delta$ , we can guarantee
$$
\begin{array}{l} \max  _ {t} | S (\boldsymbol {\epsilon} _ {0: t - 1}) | = \frac {\eta^ {2} | \mathbf {1} ^ {\top} \boldsymbol {a} ^ {*} |}{2 \pi} \max  _ {t} \left\{\left| \sum_ {l = 0} ^ {t - 1} \left(1 - \frac {2 \eta (\pi - 1)}{\pi}\right) ^ {t - l - 1} \sum_ {m = 0} ^ {l - 1} \left(1 - \eta \frac {\pi + k - 1}{2 \pi}\right) ^ {l - 1 - m} (\mathbf {1} ^ {\top} \boldsymbol {\epsilon} _ {m}) \right| \right. \\ + \eta \left| \sum_ {j = 0} ^ {t} \left(1 - \frac {2 \eta (\pi - 1)}{\pi}\right) ^ {t - 1 - j} \boldsymbol {\epsilon} _ {j} ^ {\top} \boldsymbol {a} ^ {*} \right| \\ \leq c \sqrt {\frac {\eta \log (2 / \delta)}{N}} K _ {a} ^ {2} \cdot \max  _ {t} \left[ \eta \frac {| \mathbf {1} ^ {\top} \mathbf {a} ^ {*} |}{\| \mathbf {a} ^ {*} \|} \sum_ {l = 0} ^ {t - 1} \left(1 - \frac {2 \eta (\pi - 1)}{\pi}\right) ^ {t - l - 1} + 1 \right] \\ \leq 2 c \sqrt {\frac {\eta k \log (2 / \delta)}{N}} K _ {a} ^ {2}. \tag {D.3} \\ \end{array}
$$
For any  $0 \leq s \leq \widetilde{O}(\eta^{-2})$  and any  $i \in [N]$ , with probability at least  $1 - \delta$ , it holds that
$$
\begin{array}{l} \left| \cos \phi_ {s} ^ {i} - \cos \phi_ {s} \right| = \left| \frac {\left(\boldsymbol {v} _ {s} ^ {i}\right) ^ {\top} \boldsymbol {w} ^ {*}}{\left\| \boldsymbol {v} _ {s} ^ {i} \right\|} - \frac {\boldsymbol {v} _ {s} ^ {\top} \boldsymbol {w} ^ {*}}{\left\| \boldsymbol {v} _ {s} \right\|} \right| \\ \leq \frac {\left| \left(\boldsymbol {v} _ {s} ^ {i} - \boldsymbol {v} _ {s}\right) ^ {\top} \boldsymbol {w} ^ {*} \right|}{\left\| \boldsymbol {v} _ {s} ^ {i} \right\|} + \left| \boldsymbol {v} _ {s} ^ {\top} \boldsymbol {w} ^ {*} \right| \left| \frac {1}{\left\| \boldsymbol {v} _ {s} ^ {i} \right\|} - \frac {1}{\left\| \boldsymbol {v} _ {s} \right\|} \right| \\ \stackrel {(i)} {\leq} \frac {\| \boldsymbol {v} _ {s} ^ {i} - \boldsymbol {v} _ {s} \|}{\| \boldsymbol {v} _ {s} \| - \| \boldsymbol {v} _ {s} ^ {i} - \boldsymbol {v} _ {s} \|} + \frac {\| \boldsymbol {v} _ {s} ^ {i} - \boldsymbol {v} _ {s} \|}{\| \boldsymbol {v} _ {s} \| - \| \boldsymbol {v} _ {s} ^ {i} - \boldsymbol {v} _ {s} \|} \\ \stackrel {(i i)} {\leq} \frac {2 c \eta K _ {a} ^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)}}{1 / 2 - c \eta K _ {a} ^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)}} \\ \end{array}
$$
$$
\leq 1 2 c \eta K _ {a} ^ {- 2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)}, \tag {D.4}
$$
where  $(i)$  follows from  $\| \pmb{w}^{*}\| = 1$  and the triangle inequality; and  $(ii)$  holds due to Lemma 14. Define a good event
$$
\mathcal {E} _ {\text {g o o d}} = \cap_ {\ell = 0} ^ {4} \mathcal {E} _ {\ell}, \tag {D.5}
$$
$$
\mathcal {E} _ {0} = \left\{\max  _ {t \geq 0} | S (\epsilon_ {0: t - 1}) | \leq 2 c \sqrt {\frac {\eta k \log (2 / \delta)}{N}} K _ {a} ^ {2} \right\},
$$
$$
\mathcal {E} _ {1} = \left\{\max _ {i \in [ N ], s \leq \widehat {O} (\eta^ {- 2})} \left| \boldsymbol {a} _ {m} ^ {\top} \boldsymbol {a} ^ {*} - (\boldsymbol {a} _ {m} ^ {i}) ^ {\top} \boldsymbol {a} ^ {*} \right| \leq c \eta^ {2} (I + \sqrt {k I}) \sqrt {\log (N / \eta \delta)} K _ {a} ^ {2} \right\},
$$
$$
\mathcal {E} _ {2} = \left\{\max  _ {s \leq \widetilde {O} (\eta^ {- 2}), 0 \leq \tau \leq \widetilde {O} (\eta^ {- 2})} \left| \sum_ {m = \tau} ^ {s - 1} \left(\boldsymbol {\xi} _ {m} ^ {\top} \boldsymbol {w} ^ {*} + \boldsymbol {h} _ {m} ^ {\top} \boldsymbol {w} ^ {*}\right) \right| \leq c K _ {a} \sqrt {\frac {\eta^ {- 1} \log (1 / \eta \delta)}{N}} \right\},
$$
$$
\mathcal {E} _ {3} = \left\{\max  _ {i \in [ N ], s \leq \widetilde {O} (\eta^ {- 2})} | \cos \phi_ {s} ^ {i} - \cos \phi_ {s} | \leq 1 2 c \eta K _ {a} ^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} \right\},
$$
$$
\mathcal{E}_{4} = \left\{\max_{0\leq \tau \leq \widehat{O} (\eta^{-2})}\| \boldsymbol{a}_{\tau}\| \leq K_{a},\quad \min_{0\leq \tau \leq \tau_{1,1}}A_{\tau}\geq -c\eta \sqrt{\frac{\eta\log(1 / \eta\delta)}{N}}\right\} .
$$
According to Lemmas 14, 12 and relations (D.3), (D.4), (A.20), we know  $\mathbb{P}\left(\mathcal{E}_{\ell}^{c}\right) \leq \delta$  for  $0 \leq \ell \leq 4$ .
D.1 Self-correction of the Second Layer
Theorem 1 restated. For any initial point  $(\pmb{v}_0, \pmb{a}_0)$  with  $\phi_0 \in [0, \pi)$ , we denote  $\tau_a = \inf \{t \geq 0 : \pmb{a}_t^\top \pmb{a}^* \geq \gamma_a\}$  the first time, where  $\gamma_a = \frac{16(\pi - 1)|\mathbf{1}^\top \pmb{a}^*|^2}{k(\pi + k - 1)}$ . Under Assumption 1, if
$$
c k ^ {2} \sqrt {\log (N d k / \delta)} \max  \left\{\eta (\sqrt {I k} + I), \eta (\sqrt {I d} + I), \sqrt {\frac {\eta k}{N}} \right\} \frac {\| \boldsymbol {a} ^ {*} \| ^ {2}}{\left| \mathbf {1} ^ {\top} \boldsymbol {a} ^ {*} \right| ^ {2}} <   1, \tag {D.6}
$$
for a sufficiently large constant  $c$ , then  $\tau_{a} \leq \widetilde{O}(\eta^{-2})$  with probability at least  $1 - \delta$ .
Proof. We first define a sequence of events as
$$
\mathcal {A} _ {t} := \left\{\forall s \leq t: \boldsymbol {a} _ {s} ^ {\top} \boldsymbol {a} ^ {*} \leq \gamma_ {a} \right\}, \quad \text {w h e r e} \gamma_ {a} = \frac {1 6 (\pi - 1) \left(\boldsymbol {1} ^ {\top} \boldsymbol {a} ^ {*}\right) ^ {2}}{k (\pi + k - 1)}.
$$
Notice that once we show for some  $\tau > 0$ ,
$$
\mathbb {P} \left(\mathcal {A} _ {\tau - 1} \cap \left\{\boldsymbol {a} _ {\tau} ^ {\top} \boldsymbol {a} ^ {*} <   \gamma_ {a} \right\}\right) \leq \delta ,
$$
then we can conclude that  $\mathbb{P}(\mathcal{A}_{\tau}) \leq \delta$ . According to the definition of  $\tau_{a}$ , we can guarantee  $\tau_{a} \leq \tau$  with probability at least  $1 - \delta$ .
For simplicity, we introduce the following stopping time w.r.t. the angle  $\phi$  during the training process,
$$
\tau_ {-} (\phi) = \inf  _ {t \geq 0} \left\{t: \phi_ {t} \geq \phi \right\} \quad \text {f o r} \quad \phi \in [ 0, \pi ].
$$
Define  $\tilde{\phi}^l$  and  $\tilde{\phi}^u$  such that
$$
\cos \tilde {\phi} ^ {l} = \frac {1 4 4 \pi}{\pi + k - 1} \frac {\left(\mathbf {1} ^ {\top} \mathbf {a} ^ {*}\right) ^ {2}}{k} \log (4 + k ^ {2}), \quad \cos \tilde {\phi} ^ {u} = - \cos \tilde {\phi} ^ {l}.
$$
Choose  $\tilde{\phi}^o = \arccos (1 / 5)$  such that  $\cos \tilde{\phi}^l\leq 1 / 5 = \cos \tilde{\phi}^o$  due to Assumption 1 such that
$$
\left(\mathbf {1} ^ {\top} \mathbf {a} ^ {*}\right) ^ {2} \leq \frac {k (\pi + k - 1)}{7 2 0 \pi \log \left(4 + k ^ {2}\right)}. \tag {D.7}
$$
The correction condition (6) is equivalent to
$$
\left(1 - \frac {3 2 (\pi - 1)}{k}\right) \frac {\pi - 1}{\pi + k - 1} \frac {\left(\mathbf {1} ^ {\top} \mathbf {a} ^ {*}\right) ^ {2}}{\| \mathbf {a} ^ {*} \| ^ {2}} > 1. \tag {D.8}
$$
$$
\mathcal {Z} _ {1} = \left\{\tau_ {-} \left(\tilde {\phi} _ {0}\right) > \tau_ {1, 1} - 1 \right\},
$$
$$
\mathcal {Z} _ {2} = \left\{\tau_ {-} (\pi / 2) > \tau_ {-} \left(\tilde {\phi} ^ {o}\right) + \tau_ {1, 2} - 1 \right\},
$$
$$
\mathcal {Z} _ {3} = \left\{\tau_ {-} (\tilde {\phi} ^ {u}) > \tau_ {-} (\pi / 2) + \tau_ {1, 3} - 1 \right\},
$$
where  $\tau_{1,1},\tau_{1,2}$  and  $\tau_{1,3}$  are deterministic and will be defined later.
Initial region:  $(0, \tilde{\phi}^o)$ . For any  $s \leq \min \left\{\tau_{-}(\tilde{\phi}^o), \widetilde{O}(\eta^{-2})\right\}$ , under the event  $\mathcal{E}_{\mathrm{good}}$ , it holds that
$$
\begin{array}{l} \cos \phi_ {s} ^ {i} \geq \cos \phi_ {s} - | \cos \phi_ {s} ^ {i} - \cos \phi_ {s} | \\ \geq \cos \tilde {\phi} ^ {o} - 1 2 c \eta K _ {a} ^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} \\ \geq \frac {\cos \tilde {\phi} ^ {o}}{2} = \frac {1}{1 0} > \arccos (0. 4 6 \pi). \\ \end{array}
$$
Since  $g(\cdot) \in [0, \pi]$  is decreasing in  $[0, \pi]$ , plugging it into (D.2), we can get the following lower bound
$$
\frac {1}{N} \sum_ {i = 1} ^ {N} B _ {s} ^ {i} \geq \frac {1}{N} \sum_ {i = 1} ^ {N} [ g (\phi_ {s} ^ {i}) - 1 ] \geq g (0. 4 6 \pi) - 1,
$$
Under the event  $\mathcal{E}_{\mathrm{good}}$ , using the lower bound of  $A_0$  in (B.5), for  $t \leq \min \left\{\tau_{-}(\tilde{\phi}^o), \widetilde{O}(\eta^{-2})\right\}$ , we can guarantee
$$
\begin{array}{l} \boldsymbol {a} _ {t} ^ {\top} \boldsymbol {a} ^ {*} \geq \eta \| \boldsymbol {a} ^ {*} \| ^ {2} \frac {g (0 . 4 6 \pi) - 1}{2 \pi} \sum_ {s = 0} ^ {t - 1} \left(1 - \frac {\eta (\pi - 1)}{2 \pi}\right) ^ {t - 1 - s} \left[ 1 - \left(1 - \eta \frac {\pi + k - 1}{2 \pi}\right) ^ {s} \right] \\ - \left\| \boldsymbol {a} ^ {*} \right\| ^ {2} \left(1 - \frac {\eta (\pi - 1)}{2 \pi}\right) ^ {t} - 2 c \sqrt {\frac {\eta \log (2 / \delta)}{N}} K _ {a} ^ {2} \\ = \frac {g (0 . 4 6 \pi) - 1}{\pi - 1} \| \pmb {a} ^ {*} \| ^ {2} \left[ 1 - \left(1 - \frac {\eta (\pi - 1)}{2 \pi}\right) ^ {t} \right] \\ - \frac {1 - \left(\frac {2 \pi - \eta (k + \pi - 1)}{2 \pi - \eta (\pi - 1)}\right) ^ {t}}{k} \left(1 - \eta \frac {\pi - 1}{2 \pi}\right) ^ {t} [ g (0. 4 6 \pi) - 1 ] \| \boldsymbol {a} ^ {*} \| ^ {2} \\ - \left\| \boldsymbol {a} ^ {*} \right\| ^ {2} \left(1 - \frac {\eta (\pi - 1)}{2 \pi}\right) ^ {t} - 2 c \sqrt {\frac {\eta k \log (2 / \delta)}{N}} K _ {a} ^ {2}. \\ \end{array}
$$
Here we take
$$
\tau_ {1, 1} = \frac {1}{\eta} \frac {\pi - 1}{2 \pi} \log \left(4 + \frac {4 (\pi - 1)}{g (0 . 4 6 \pi) - 1}\right) \leq \frac {1}{\eta} \frac {\pi - 1}{2 \pi} \log (4 + 2 0 (\pi - 1)), \tag {D.9}
$$
where we used the fact  $g(0.46\pi) - 1 \geq 0.2$ . Since  $\tau_{1,1} - 1 \leq \widetilde{O}(\eta^{-2})$ , under the event  $\left\{\tau_{1,1} - 1 < \tau_{-}(\tilde{\phi}^{o})\right\} \cap \mathcal{E}_{\mathrm{good}}$ , it holds that
$$
\begin{array}{l} \boldsymbol {a} _ {\tau_ {1, 1}} ^ {\top} \boldsymbol {a} ^ {*} \geq \frac {(g (0 . 4 6 \pi) - 1) \| \boldsymbol {a} ^ {*} \| ^ {2}}{2 (\pi - 1)} - 2 c \sqrt {\frac {\eta k \log (2 t / \delta)}{N}} K _ {a} ^ {2} \\ \stackrel {(i)} {\geq} \frac {0 . 2 (\mathbf {1} ^ {\top} \mathbf {a} ^ {*}) ^ {2}}{2 k (\pi - 1)} - 2 c \sqrt {\frac {\eta k \log (2 t / \delta)}{N}} K _ {a} ^ {2} \\ \end{array}
$$
$$
\begin{array}{l} \stackrel {(i i)} {\geq} \frac {3 2 (\pi - 1) \left(\mathbf {1} ^ {\top} \mathbf {a} ^ {*}\right) ^ {2}}{k (\pi + k - 1)} - 2 c \sqrt {\frac {\eta k \log (2 t / \delta)}{N}} K _ {a} ^ {2} \\ \stackrel {(i i i)} {\geq} \frac {1 6 (\pi - 1) \left(\mathbf {1} ^ {\top} \boldsymbol {a} ^ {*}\right) ^ {2}}{k (\pi + k - 1)} = \gamma_ {a}, \\ \end{array}
$$
where  $(i)$  holds due to  $\| \pmb{a}^{*}\| \geq |\mathbf{1}^{\top}\pmb{a}^{*}| / \sqrt{k}$ ;  $(ii)$  follows from  $k\geq 320(\pi -1)^2$ ; and  $(iii)$  holds due to the assumption (D.6). We can conclude that
$$
\begin{array}{l} \mathbb {P} \left(\left\{\tau_ {1, 1} - 1 <   \tau_ {-} \left(\tilde {\phi} ^ {o}\right) \right\} \cap \left\{\tau_ {a} > \tau_ {1, 1} \right\}\right) \leq \mathbb {P} \left(\left\{\tau_ {1, 1} - 1 <   \tau_ {-} \left(\tilde {\phi} ^ {o}\right) \right\} \cap \left\{\boldsymbol {a} _ {\tau_ {1, 1}} ^ {\top} \boldsymbol {a} ^ {*} \leq \gamma_ {a} \right\}\right) \\ \leq \mathbb {P} \left(\left\{\tau_ {1, 1} - 1 <   \tau_ {-} \left(\tilde {\phi} ^ {o}\right) \right\} \cap \left(\mathcal {E} _ {0} ^ {c} \cup \mathcal {E} _ {3} ^ {c}\right)\right) \\ \leq \mathbb {P} \left(\mathcal {E} _ {0} ^ {c} \cup \mathcal {E} _ {3} ^ {c}\right) \leq 2 \delta , \tag {D.10} \\ \end{array}
$$
where the second last inequality holds since
$$
\left\{\tau_ {1, 1} - 1 <   \tau_ {-} \left(\tilde {\phi} ^ {o}\right) \leq \widetilde {O} \left(\eta^ {- 2}\right) \right\} \cap \mathcal {E} _ {\text {g o o d}} \subseteq \left\{\boldsymbol {a} _ {\tau_ {1, 1}} ^ {\top} \boldsymbol {a} ^ {*} \geq \gamma_ {a} \right\},
$$
$$
\left\{\tau_ {1, 1} - 1 <   \tau_ {-} (\tilde {\phi} ^ {o}) \right\} \cap \left\{\boldsymbol {a} _ {\tau_ {1, 1}} ^ {\top} \boldsymbol {a} ^ {*} \leq \gamma_ {a} \right\} \subseteq \left\{\tau_ {1, 1} - 1 <   \tau_ {-} (\tilde {\phi} ^ {o}) \right\} \cap \mathcal {E} _ {\mathrm {g o o d}} ^ {c}.
$$
If  $\tau_{1,1} \geq \tau_{-}(\tilde{\phi}^o)$ , we regard  $(\pmb{v}_{\tau_{-}(\tilde{\phi}^{o})}, \pmb{a}_{\tau_{-}(\tilde{\phi}^{o})})$  as the new initial point and step into the analysis of the next region.
Initial region:  $[\tilde{\phi}^o, \pi/2)$ . Denote  $t_0 = \tau_{-}(\tilde{\phi}^o)$ . From now on, we assume the event  $\mathcal{Z}_1^c$  happens and hide the indicator. Recalling the definition of  $h_t$ , we have