SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Proof. In this proof, we suppose the event Lemma 12 holds and hide the indicator function everywhere. To facilitate the technical presentation, denote  $\tilde{\boldsymbol{v}}_t^i = \boldsymbol{v}_t^i - \eta \mathbf{P}_t^i \nabla_{\boldsymbol{w}} L(\boldsymbol{w}_t^i, \boldsymbol{a}_t^i)$ . Define the following event
$$
\mathcal {U} _ {t} = \left\{\forall s \leq t \leq (8 0 0 c \eta^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {4}) ^ {- 1}: 1 / 2 \leq \| \boldsymbol {v} _ {s} \| \leq 3 \right\}.
$$
Step 1: show  $\| \pmb{v}_t\| \geq 1 / 2$  under the event  $\mathcal{U}_{t - 1}$ . From  $\pmb{v}_t = \frac{1}{N}\sum_{i = 1}^N\pmb{v}_t^i$  and  $\mathbf{P}_{t - 1}\pmb{v}_{t - 1} = \mathbf{0}$ , we have
$$
\begin{array}{l} \| \boldsymbol {v} _ {t} \| ^ {2} = \left\| \boldsymbol {v} _ {t - 1} - \eta \frac {1}{N} \sum_ {i = 1} ^ {N} \mathbf {P} _ {t - 1} ^ {i} \nabla_ {\boldsymbol {w}} L (\boldsymbol {w} _ {t - 1} ^ {i}, \boldsymbol {a} _ {t - 1} ^ {i}) \right\| ^ {2} - \frac {2 \eta}{N} \sum_ {i = 1} ^ {N} \left\langle \tilde {\boldsymbol {v}} _ {t - 1} ^ {i}, \boldsymbol {\xi} _ {t - 1} \right\rangle + \eta^ {2} \| \boldsymbol {\xi} _ {t - 1} \| ^ {2} \\ = \left\| \boldsymbol {v} _ {t - 1} - \eta \mathbf {P} _ {t - 1} \frac {1}{N} \sum_ {i = 1} ^ {N} \nabla_ {\boldsymbol {w}} L (\boldsymbol {w} _ {t - 1} ^ {i}, \boldsymbol {a} _ {t - 1} ^ {i}) \right\| ^ {2} - \frac {2 \eta}{N} \sum_ {i = 1} ^ {N} \left\langle \tilde {\boldsymbol {v}} _ {t - 1} ^ {i}, \boldsymbol {\xi} _ {t - 1} \right\rangle + \eta^ {2} \| \boldsymbol {\xi} _ {t - 1} \| ^ {2} \\ + 2 \eta \left\langle \boldsymbol {v} _ {t - 1} - \eta \mathbf {P} _ {t - 1} \frac {1}{N} \sum_ {i = 1} ^ {N} \nabla_ {\boldsymbol {w}} L (\boldsymbol {w} _ {t - 1} ^ {i}, \boldsymbol {a} _ {t - 1} ^ {i}), \frac {1}{N} \sum_ {i = 1} ^ {N} (\mathbf {P} _ {t - 1} - \mathbf {P} _ {t - 1} ^ {i}) \nabla_ {\boldsymbol {w}} L (\boldsymbol {w} _ {t - 1} ^ {i}, \boldsymbol {a} _ {t - 1} ^ {i}) \right\rangle \\ + \eta^ {2} \left\| \frac {1}{N} \sum_ {i = 1} ^ {N} (\mathbf {P} _ {t - 1} - \mathbf {P} _ {t - 1} ^ {i}) \nabla_ {\pmb {w}} L (\pmb {w} _ {t - 1} ^ {i}, \pmb {a} _ {t - 1} ^ {i}) \right\| ^ {2} \\ \geq \| \boldsymbol {v} _ {t - 1} \| ^ {2} - \frac {2 \eta}{N} \sum_ {i = 1} ^ {N} \left\langle \tilde {\boldsymbol {v}} _ {t - 1} ^ {i}, \boldsymbol {\xi} _ {t - 1} \right\rangle - 4 \eta \| \boldsymbol {v} _ {t - 1} \| ^ {2} \frac {1}{N} \sum_ {i = 1} ^ {N} \| \mathbf {P} _ {t - 1} - \mathbf {P} _ {t - 1} ^ {i} \| \| \nabla_ {\boldsymbol {w}} L (\boldsymbol {w} _ {t - 1} ^ {i}, \boldsymbol {a} _ {t - 1} ^ {i}) \| \\ - 2 \eta^ {2} \frac {1}{N} \sum_ {i = 1} ^ {N} \| \mathbf {P} _ {t - 1} \| \left\| \nabla_ {\boldsymbol {w}} L \left(\boldsymbol {w} _ {t - 1} ^ {i}, \boldsymbol {a} _ {t - 1} ^ {i}\right) \right\| \frac {1}{N} \sum_ {i = 1} ^ {N} \left\| \mathbf {P} _ {t - 1} - \mathbf {P} _ {t - 1} ^ {i} \right\| \left\| \nabla_ {\boldsymbol {w}} L \left(\boldsymbol {w} _ {t - 1} ^ {i}, \boldsymbol {a} _ {t - 1} ^ {i}\right) \right\| \tag {B.8} \\ \end{array}
$$
where the inequality follows from dropping positive terms and the inductive assumption  $\| \pmb{v}_{t - 1}\| \geq 1 / 2$ . With probability at least  $1 - \delta$ , it holds that
$$
\begin{array}{l} \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\| \mathbf {P} _ {t - 1} ^ {i} - \mathbf {P} _ {t - 1} \right\| = \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left(\left\| \frac {1}{\| \boldsymbol {v} _ {t - 1} ^ {i} \|} \left(\mathbf {I} - \boldsymbol {w} _ {t - 1} ^ {i} (\boldsymbol {w} _ {t - 1} ^ {i}) ^ {\top}\right) - \frac {1}{\| \boldsymbol {v} _ {t - 1} \|} \left(\mathbf {I} - \boldsymbol {w} _ {t - 1} \boldsymbol {w} _ {t - 1} ^ {\top}\right) \right\|\right) \\ \leq \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\{\left| \frac {1}{\| \boldsymbol {v} _ {t - 1} ^ {i} \|} - \frac {1}{\| \boldsymbol {v} _ {t - 1} \|} \right| + \frac {1}{\| \boldsymbol {v} _ {t - 1} \|} \| \boldsymbol {w} _ {t - 1} ^ {i} (\boldsymbol {w} _ {t - 1} ^ {i}) ^ {\top} - \boldsymbol {w} _ {t - 1} \boldsymbol {w} _ {t - 1} ^ {\top} \| \right\} \\ \stackrel {(i)} {\leq} \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\{\left| \frac {1}{\| \boldsymbol {v} _ {t - 1} ^ {i} \|} - \frac {1}{\| \boldsymbol {v} _ {t - 1} \|} \right| + \frac {2 \| \boldsymbol {w} _ {t - 1} ^ {i} - \boldsymbol {w} _ {t - 1} \|}{\| \boldsymbol {v} _ {t - 1} \|} \right\} \\ \stackrel {(i i)} {\leq} \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\{\left| \frac {1}{\| \boldsymbol {v} _ {t - 1} ^ {i} \|} - \frac {1}{\| \boldsymbol {v} _ {t - 1} \|} \right| + 4 \left\| \frac {\boldsymbol {v} _ {t - 1} ^ {i}}{\| \boldsymbol {v} _ {t - 1} ^ {i} \|} - \frac {\boldsymbol {v} _ {t - 1}}{\| \boldsymbol {v} _ {t - 1} \|} \right\| \right\} \\ \leq \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\{\left| \frac {1}{\| \boldsymbol {v} _ {t - 1} ^ {i} \|} - \frac {1}{\| \boldsymbol {v} _ {t - 1} \|} \right| + \frac {4}{\| \boldsymbol {v} _ {t - 1} ^ {i} \|} \| \boldsymbol {v} _ {t - 1} ^ {i} - \boldsymbol {v} _ {t - 1} \| + 4 \| \boldsymbol {v} _ {t - 1} \| \left| \frac {1}{\| \boldsymbol {v} _ {t - 1} ^ {i} \|} - \frac {1}{\| \boldsymbol {v} _ {t - 1} \|} \right| \right\} \\ \stackrel {(i i i)} {\leq} \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\{\frac {2 \| \boldsymbol {v} _ {t - 1} ^ {i} - \boldsymbol {v} _ {t - 1} \|}{\| \boldsymbol {v} _ {t - 1} ^ {i} \|} + \frac {8 \| \boldsymbol {v} _ {t - 1} ^ {i} - \boldsymbol {v} _ {t - 1} \|}{\| \boldsymbol {v} _ {t - 1} ^ {i} \|} \right\} \\ \stackrel {(i v)} {\leq} \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\{\frac {1 0 c \eta (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {2}}{1 / 2 - 6 c \eta (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {2}} \right\} \\ \stackrel {(v)} {\leq} \mathbb {1} _ {\mathcal {U} _ {t - 1}} \cdot 4 0 c \eta (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {2}, \tag {B.9} \\ \end{array}
$$
where  $(i)$  holds due to  $\| \pmb{w}_{t - 1}^i\| = \| \pmb{w}_{t - 1}\| = 1$ ;  $(ii)$  and  $(iii)$  hold due to the inductive assumption  $\| \pmb{v}_{t - 1}\| \geq 1 / 2$ ;  $(iv)$  follows from the conclusion (1) in Lemma 14; and  $(v)$  is true because  $6c\eta (\sqrt{Id} + I)\sqrt{\log(Nd / \delta)}\leq 1 / 4$ . Substituting (B.9) into (B.8), with probability at least  $1 - 2\delta$ , we have
$$
\mathbb {1} _ {\mathcal {U} _ {t - 1}} \| \boldsymbol {v} _ {t} \| ^ {2} \geq \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\{\| \boldsymbol {v} _ {t - 1} \| ^ {2} - \frac {2 \eta}{N} \sum_ {i = 1} ^ {N} \langle \tilde {\boldsymbol {v}} _ {t - 1} ^ {i}, \boldsymbol {\xi} _ {t - 1} \rangle \right.
$$
$$
\begin{array}{l} - 1 6 0 c \eta^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {2} \cdot \| \boldsymbol {v} _ {t - 1} \| ^ {2} \cdot \frac {1}{N} \sum_ {i = 1} ^ {N} \| \nabla_ {\boldsymbol {w}} L (\boldsymbol {w} _ {t - 1} ^ {i}, \boldsymbol {a} _ {t - 1} ^ {i}) \| \\ \left. - 8 0 \eta^ {3} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {2} \cdot \| \mathbf {P} _ {t - 1} \| \cdot \left(\frac {1}{N} \sum_ {i = 1} ^ {N} \| \nabla_ {\boldsymbol {w}} L (\boldsymbol {w} _ {t - 1} ^ {i}, \boldsymbol {a} _ {t - 1} ^ {i}) \|\right) ^ {2} \right\} \\ \stackrel {(i)} {\geq} \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\{\left(1 - 8 0 0 c \eta^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {4}\right) \| \boldsymbol {v} _ {t - 1} \| ^ {2} - \frac {2 \eta}{N} \sum_ {i = 1} ^ {N} \left\langle \tilde {\boldsymbol {v}} _ {t - 1} ^ {i}, \boldsymbol {\xi} _ {t - 1} \right\rangle \right\} \\ - 1 0 0 0 \eta^ {3} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {6} \\ \end{array}
$$
$$
\begin{array}{l} \stackrel {(i i)} {\geq} \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\{\left(1 - 8 0 0 c \eta^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {4}\right) ^ {t} \| \boldsymbol {v} _ {0} \| \right. \\ - 2 \eta \left| \sum_ {s = 0} ^ {t - 1} \left(1 - 8 0 0 c \eta^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {4}\right) ^ {t - s} \left\langle \frac {1}{N} \sum_ {i = 1} ^ {N} \tilde {\boldsymbol {v}} _ {s} ^ {i}, \boldsymbol {\xi} _ {s} \right\rangle \right| \\ \left. - 1 0 0 0 \eta^ {3} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {6} \sum_ {s = 0} ^ {t - 1} \left(1 - 8 0 0 c \eta^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {4}\right) ^ {t - s} \right\} \\ \end{array}
$$
$$
\geq \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\{0. 3 6 - \frac {5}{4} \eta K _ {a} ^ {2} - 2 \eta \left| \sum_ {s = 0} ^ {t - 1} \left(1 - 8 0 0 c \eta^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {4}\right) ^ {t - s} \left\langle \frac {1}{N} \sum_ {i = 1} ^ {N} \tilde {\boldsymbol {v}} _ {s} ^ {i}, \boldsymbol {\xi} _ {s} \right\rangle \right| \right\}, \tag {B.10}
$$
where  $(i)$  holds because  $\| \nabla_{\boldsymbol{w}}L(\boldsymbol{w}_t^i,\boldsymbol{v}_t^i)\| \leq |(\boldsymbol{a}_t^\intercal\boldsymbol{a}^*)| \leq 5K_a^2$  by using Lemma 12 and  $\|\mathbf{P}_{t-1}\| \leq \frac{1}{\|\boldsymbol{v}_{t-1}\|} \leq 2$  holds under the event  $\mathcal{U}_{t-1}$ ; (ii) holds due to  $\|\boldsymbol{v}_0\| = 1$ ; and (iii) follows from  $t \leq (800c\eta^2(d \vee I)K_a^4\sqrt{\log(Nd/\delta)})^{-1}$  such that
$$
\begin{array}{l} \left(1 - 8 0 0 c \eta^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {4}\right) ^ {t} \\ \geq \left(1 - 8 0 0 c \eta^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {4}\right) ^ {\frac {1}{8 0 0 c \eta^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {4}}} \\ \geq (1 - 0. 0 1) ^ {1 0 0} \geq 0. 3 6. \\ \end{array}
$$
In fact, the bound above holds because  $800c\eta^{2}(\sqrt{Id} + I)\sqrt{\log(Nd / \delta)} K_{a}^{4} \leq 0.01$  and  $(1 - x)^{\frac{1}{x}}$  is decreasing in  $(0,1)$ . According to the definition of  $\tilde{\pmb{v}}_s^i$  and Lemma 12, we know
$$
\mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\| \frac {1}{N} \sum_ {i = 1} ^ {N} \tilde {\boldsymbol {v}} _ {s} ^ {i} \right\| = \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\| \boldsymbol {v} _ {s} - \eta \frac {1}{N} \sum_ {i = 1} ^ {N} \mathbf {P} _ {s} ^ {i} \nabla_ {\boldsymbol {w}} L \left(\boldsymbol {w} _ {s} ^ {i}, \boldsymbol {a} _ {s} ^ {i}\right) \right\| \leq 3 + 5 \eta K _ {a} ^ {2}. \tag {B.11}
$$
Recall that  $\pmb{\xi}_{s} = \frac{1}{N}\sum_{i = 1}^{N}\mathbf{P}_{s}^{i}\pmb{\xi}_{s}^{i}$ . Invoking Lemmas 11 and 10, we have
$$
\begin{array}{l} \mathbb {1} _ {\mathcal {U} _ {t - 1}} \eta \left| \sum_ {s = 0} ^ {t - 1} \left(1 - 8 0 0 \eta^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {4}\right) ^ {t - 1 - s} \left\langle \frac {1}{N} \sum_ {i = 1} ^ {N} \tilde {\boldsymbol {v}} _ {s} ^ {i}, \boldsymbol {\xi} _ {s} \right\rangle \right| \\ \leq \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\{\frac {c \eta K _ {a} ^ {2} (3 + 5 \eta K _ {a} ^ {2})}{\sqrt {N}} \sqrt {\log (2 / \delta) \sum_ {s = 0} ^ {t - 1} \left(1 - 8 0 0 \eta^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {4}\right) ^ {2 (t - 1 - s)}} \right. \\ \left. + \frac {c \eta K _ {a} ^ {2} \log (2 / \delta)}{N} \right\} \\ \leq \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\{\frac {c \left(3 + 5 \eta K _ {a} ^ {2}\right)}{2 0 \sqrt {I \log (N d / \delta)}} \sqrt {\frac {\log (2 / \delta)}{N}} + \frac {c \eta K _ {a} ^ {2} \log (2 / \delta)}{N} \right\} \\ \leq \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\{\frac {c \left(3 + 5 \eta K _ {a} ^ {2}\right)}{2 0 \sqrt {N I}} + \frac {c \eta K _ {a} ^ {2} \log (2 / \delta)}{N} \right\}, \tag {B.12} \\ \end{array}
$$
holds with probability at least  $1 - \delta$ . Plugging (B.12) into (B.10), we get
$$
\begin{array}{l} \| \boldsymbol {v} _ {t} \| ^ {2} \mathbb {1} _ {\mathcal {U} _ {t - 1}} \geq \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\{0. 3 6 - \frac {5}{4} \eta K _ {a} ^ {2} - 2 \left(\frac {c (3 + 5 \eta K _ {a} ^ {2})}{2 0 \sqrt {I \log (N d / \delta)}} \sqrt {\frac {\log (2 / \delta)}{N}} + \frac {c \eta K _ {a} ^ {2} \log (2 / \delta)}{N}\right) \right\} \\ \geq \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\{0. 3 6 - \frac {1}{1 6} - \frac {1}{1 6} \right\} \\ \geq \mathbb {1} _ {\mathcal {U} _ {t - 1}} \cdot \frac {1}{4}, \tag {B.13} \\ \end{array}
$$
where the first inequality holds since  $\sqrt{Id\log(Nd / \delta)}\geq c$
Step 2: show  $\| \pmb{v}_t \| \leq 3$  under the event  $\mathcal{U}_{t-1}$ . From the update rule of  $\pmb{v}_t^i$  and  $\pmb{v}_t = \frac{1}{N} \sum_{i=1}^{N} \pmb{v}_t^i$ , with probability at least  $1 - \delta$ , we have
$$
\begin{array}{l} \mathbb {1} _ {\mathcal {U} _ {t - 1}} \| \boldsymbol {v} _ {t} \| ^ {2} = \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\{\left\| \boldsymbol {v} _ {t - 1} - \eta \mathbf {P} _ {t - 1} \frac {1}{N} \sum_ {i = 1} ^ {N} \nabla_ {\boldsymbol {w}} L (\boldsymbol {w} _ {t - 1} ^ {i}, \boldsymbol {a} _ {t - 1} ^ {i}) \right\| ^ {2} - \frac {2 \eta}{N} \sum_ {i = 1} ^ {N} \left\langle \tilde {\boldsymbol {v}} _ {t - 1} ^ {i}, \boldsymbol {\xi} _ {t - 1} \right\rangle + \eta^ {2} \| \boldsymbol {\xi} _ {t - 1} \| ^ {2} \right. \\ + 2 \eta \left\langle \boldsymbol {v} _ {t - 1} - \eta \mathbf {P} _ {t - 1} \frac {1}{N} \sum_ {i = 1} ^ {N} \nabla_ {\boldsymbol {w}} L (\boldsymbol {w} _ {t - 1} ^ {i}, \boldsymbol {a} _ {t - 1} ^ {i}), \frac {1}{N} \sum_ {i = 1} ^ {N} (\mathbf {P} _ {t - 1} - \mathbf {P} _ {t - 1} ^ {i}) \nabla_ {\boldsymbol {w}} L (\boldsymbol {w} _ {t - 1} ^ {i}, \boldsymbol {a} _ {t - 1} ^ {i}) \right\rangle \\ + \eta^ {2} \left\| \frac {1}{N} \sum_ {i = 1} ^ {N} (\mathbf {P} _ {t - 1} - \mathbf {P} _ {t - 1} ^ {i}) \nabla_ {\boldsymbol {w}} L (\boldsymbol {w} _ {t - 1} ^ {i}, \boldsymbol {a} _ {t - 1} ^ {i}) \right\| ^ {2} \Bigg \} \\ \end{array}
$$
$$
\begin{array}{l} \stackrel {(i)} {\leq} \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\{\| \boldsymbol {v} _ {t - 1} \| ^ {2} + \eta^ {2} \left\| \frac {1}{N} \sum_ {i = 1} ^ {N} \mathbf {P} _ {t - 1} \nabla_ {\boldsymbol {w}} L (\boldsymbol {w} _ {t - 1} ^ {i}, \boldsymbol {a} _ {t - 1} ^ {i}) \right\| ^ {2} \right. \\ + \frac {2 \eta}{3} \| \boldsymbol {v} _ {t - 1} \| ^ {2} \frac {1}{N} \sum_ {i = 1} ^ {N} \| \mathbf {P} _ {t - 1} - \mathbf {P} _ {t - 1} ^ {i} \| \| \nabla_ {\boldsymbol {w}} L (\boldsymbol {w} _ {t - 1} ^ {i}, \boldsymbol {a} _ {t - 1} ^ {i}) \| \\ + 2 \eta^ {2} \left\| \frac {1}{N} \sum_ {i = 1} ^ {N} \mathbf {P} _ {t - 1} \nabla_ {\boldsymbol {w}} L (\boldsymbol {w} _ {t - 1} ^ {i}, \boldsymbol {a} _ {t - 1} ^ {i}) \right\| \left\| \frac {1}{N} \sum_ {i = 1} ^ {N} (\mathbf {P} _ {t - 1} - \mathbf {P} _ {t - 1} ^ {i}) \nabla_ {\boldsymbol {w}} L (\boldsymbol {w} _ {t - 1} ^ {i}, \boldsymbol {a} _ {t - 1} ^ {i}) \right\| \\ + \eta^ {2} \left\| \frac {1}{N} \sum_ {i = 1} ^ {N} \left(\mathbf {P} _ {t - 1} - \mathbf {P} _ {t - 1} ^ {i}\right) \nabla_ {\boldsymbol {w}} L \left(\boldsymbol {w} _ {t - 1} ^ {i}, \boldsymbol {a} _ {t - 1} ^ {i}\right) \right\| ^ {2} - \frac {2 \eta}{N} \sum_ {i = 1} ^ {N} \left\langle \tilde {\boldsymbol {v}} _ {t - 1} ^ {i}, \boldsymbol {\xi} _ {t - 1} \right\rangle + \eta^ {2} \| \boldsymbol {\xi} _ {t - 1} \| ^ {2} \Bigg \} \\ \end{array}
$$
$$
\begin{array}{l} \stackrel {(i i)} {\leq} \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\{\left(1 + 5 \eta K _ {a} ^ {2} \max _ {i} \| \mathbf {P} _ {t - 1} ^ {i} - \mathbf {P} _ {t - 1} \|\right) \| \pmb {v} _ {t - 1} \| ^ {2} + 2 5 \eta^ {2} K _ {a} ^ {4} + 2 5 \eta^ {2} K _ {a} ^ {4} \max _ {i} \| \mathbf {P} _ {t - 1} ^ {i} - \mathbf {P} _ {t - 1} \| \right. \\ \left. + \frac {2 5 \eta^ {2} K _ {a} ^ {4}}{4} \max  _ {i} \| \mathbf {P} _ {t - 1} ^ {i} - \mathbf {P} _ {t - 1} \| ^ {2} - \frac {2 \eta}{N} \sum_ {i = 1} ^ {N} \left\langle \tilde {\boldsymbol {v}} _ {t - 1} ^ {i}, \boldsymbol {\xi} _ {t - 1} \right\rangle + \eta^ {2} \| \boldsymbol {\xi} _ {t - 1} \| ^ {2} \right\} \\ \end{array}
$$
$$
\begin{array}{l} \stackrel {(i i i)} {\leq} \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\{\left(1 + 2 0 0 c \eta^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {4}\right) \| \boldsymbol {v} _ {t - 1} \| ^ {2} + 2 5 \eta^ {2} K _ {a} ^ {4} \right. \\ + 1 0 0 0 (1 + 1 0 \eta K _ {a} ^ {2}) \eta^ {3} K _ {a} ^ {6} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} - \frac {2 \eta}{N} \sum_ {i = 1} ^ {N} \left\langle \tilde {\boldsymbol {v}} _ {t - 1} ^ {i}, \boldsymbol {\xi} _ {t - 1} \right\rangle + \eta^ {2} \| \boldsymbol {\xi} _ {t - 1} \| ^ {2} \Bigg \rbrace \\ \end{array}
$$
$$
\begin{array}{l} \leq \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\{\left(1 + 2 0 0 c \eta^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {4}\right) ^ {t} \| \boldsymbol {v} _ {0} \| ^ {2} \right. \\ + 2 0 0 0 \eta^ {3} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {6} \sum_ {s = 0} ^ {t - 1} \left(1 + 2 0 0 c \eta^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {4}\right) ^ {t - s} \\ + 2 \eta \left| \sum_ {s = 0} ^ {t - 1} \left(1 + 2 0 0 c \eta^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {4}\right) ^ {t - 1 - s} \left\langle \frac {1}{N} \sum_ {i = 1} ^ {N} \tilde {\boldsymbol {v}} _ {s} ^ {i}, \boldsymbol {\xi} _ {s} \right\rangle \right| \\ \end{array}
$$
$$
\left. + \eta^ {2} \sum_ {s = 0} ^ {t - 1} \left(1 + 2 0 0 c \eta^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {4}\right) ^ {t - 1 - s} \| \boldsymbol {\xi} _ {s} \| ^ {2} \right\}. \tag {B.14}
$$
where  $(i)$  comes from  $\mathbf{P}_{t - 1}\pmb{v}_{t - 1} = \mathbf{0}$  and  $\| \pmb{v}_{t - 1}\| \leq 3$  under  $\mathcal{U}_{t - 1}$ ;  $(ii)$  comes from  $\| \mathbf{P}_{t - 1}\| \leq 1 / \| \pmb{v}_{t - 1}\| \leq 2$  under  $\mathcal{U}_{t - 1}$ ; and  $(iii)$  holds due to (B.9). With probability at least  $1 - \delta$ , we have the following concentration
$$
\begin{array}{l} \mathbb {1} _ {\mathcal {U} _ {t - 1}} \eta^ {2} \sum_ {s = 0} ^ {t - 1} \left(1 + 2 0 0 c \eta^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {4}\right) ^ {t - 1 - s} \| \boldsymbol {\xi} _ {s} \| ^ {2} \\ \leq \mathbb {1} _ {\mathcal {U} _ {t - 1}} \frac {c \eta^ {2} d \log (1 / \eta \delta) K _ {a} ^ {4}}{N} \sum_ {s = 0} ^ {t - 1} \left(1 + 2 0 0 c \eta^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {4}\right) ^ {t - 1 - s} \\ \leq \mathbb {1} _ {\mathcal {U} _ {t - 1}} \frac {c \eta^ {2} d \log (1 / \eta \delta) K _ {a} ^ {4}}{N} \frac {\left(1 + 4 0 0 c \eta^ {2} (I \lor d) \sqrt {\log (N d / \delta)} K _ {a} ^ {4}\right) ^ {t}}{4 0 0 c \eta^ {2} (I \lor d) \sqrt {\log (N d / \delta)} K _ {a} ^ {4}} \\ \leq \mathbb {1} _ {\mathcal {U} _ {t - 1}} \frac {c e \log (1 / \eta \delta)}{N}, \tag {B.15} \\ \end{array}
$$
where we used the assumption  $t \leq (400c\eta^2(d \vee I)\sqrt{\log(Nd / \delta)}K_a^4)^{-1}$ . Similar to (B.12), we can also guarantee
$$
\begin{array}{l} \mathbb {1} _ {\mathcal {U} _ {t - 1}} \eta \left| \sum_ {s = 0} ^ {t - 1} \left(1 + 2 0 0 c \eta^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {4}\right) ^ {t - 1 - s} \left\langle \frac {1}{N} \sum_ {i = 1} ^ {N} \tilde {\boldsymbol {v}} _ {s} ^ {i}, \boldsymbol {\xi} _ {s} \right\rangle \right| \\ \leq \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\{\frac {c \eta K _ {a} ^ {2} (3 + 5 \eta K _ {a} ^ {2})}{\sqrt {N}} \sqrt {\log (2 / \delta) \sum_ {s = 0} ^ {t - 1} (1 + 2 0 0 c \eta^ {2} (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {4}) ^ {2 (t - 1 - s)}} \right. \\ \left. + \frac {c \eta K _ {a} ^ {2} \log (2 / \delta)}{N} \right\} \\ \leq \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\{\frac {c e (3 + 5 \eta K _ {a} ^ {2})}{\sqrt {(I \vee d) N}} + \frac {c \eta K _ {a} ^ {2} \log (2 / \delta)}{N} \right\}, \tag {B.16} \\ \end{array}
$$
Plugging (B.15) and (B.16) into (B.14) gives that
$$
\begin{array}{l} \mathbb {1} _ {\mathcal {U} _ {t - 1}} \| \boldsymbol {v} _ {t} \| ^ {2} \leq \mathbb {1} _ {\mathcal {U} _ {t - 1}} \left\{e + 1 0 e \eta K _ {a} ^ {2} + \frac {2 c e \log (1 / \eta \delta)}{N} + 2 \left(\frac {c e \left(3 + 5 \eta K _ {a} ^ {2}\right)}{\sqrt {(I \lor d) N}} + \frac {c \eta K _ {a} ^ {2} \log (2 / \delta)}{N}\right) \right\} \\ \leq \mathbb {1} _ {\mathcal {U} _ {t - 1}} \cdot 9, \\ \end{array}
$$
holds with probability at least  $1 - 3\delta$ . Together with (B.13), we have showed that
$$
\begin{array}{l} \mathbb {P} \left(\mathcal {U} _ {t - 1} \cap \{1 / 2 \leq \| \boldsymbol {v} _ {t} \| \leq 3 \}\right) \geq \mathbb {P} \left(\mathcal {U} _ {t - 1}\right) - \mathbb {P} \left(\mathcal {U} _ {t - 1} \cap \{\| \boldsymbol {v} _ {t} \| \leq 1 / 2 \}\right) - \mathbb {P} \left(\mathcal {U} _ {t - 1} \cap \{\| \boldsymbol {v} _ {t} \| > 3 \}\right) \\ \geq \mathbb {P} \left(\mathcal {U} _ {t - 1}\right) - \mathbb {P} (\text {O n e o f (B . 9) , (B . 1 0) , (B . 1 3) , a n d L e m m a 1 2 d o e s n o t h o l d}) \\ - \mathbb {P} (\text {O n e o f (B . 1 4) , (B . 1 5) , (B . 1 6) , a n d L e m m a 1 2 d o e s n o t h o l d}) \\ \geq \mathbb {P} \left(\mathcal {U} _ {t - 1}\right) - 8 \delta . \\ \end{array}
$$
Using induction and  $\mathbb{P}(\mathcal{U}_0) = 1$  due to  $\| \pmb {v}_0\| = 1$ , we can prove the desired conclusion by adjusting the level of  $\delta$ .
Remark. In the following proofs, we use  $c$  to denote the absolute constant, which does not depend on any variables of the model or training process. For conciseness, we do not distinguish the specific value of  $c$  in some contexts.
C Bound the Discrepancy during Training Process
C.1 The Discrepancy in the First Layer
Lemma 14 (The discrepancy in the first layer). Denote  $O_{r,T}^{i,j} = \max_{t_r\leq t\leq T}\| \pmb {v}_t^i -\pmb {v}_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, with probability at least  $1 - \delta$ , we have
$$
\max  _ {i, j \in [ N ]} O _ {r, T} ^ {i, j} \leq c \eta (\sqrt {I d} + I) \sqrt {\log (N d / \delta)} K _ {a} ^ {2}.
$$
$$
\max _ {t _ {r} \leq t \leq T} | (\boldsymbol {a} _ {t} ^ {i} - \boldsymbol {a} _ {t} ^ {j}) ^ {\top} \boldsymbol {a} ^ {*} | \leq c \eta (\sqrt {k I} + I) \sqrt {\log (1 / \eta \delta)} K _ {a} ^ {2}.
$$
Proof. Denote the discrepancy term by
$$
\begin{array}{l} \boldsymbol {h} _ {t} ^ {i, j} = \mathbf {P} _ {t} ^ {j} \nabla_ {\boldsymbol {w}} L (\boldsymbol {w} _ {t} ^ {j}, \boldsymbol {a} _ {t} ^ {j}) - \mathbf {P} _ {t} ^ {i} \nabla_ {\boldsymbol {w}} L (\boldsymbol {w} _ {t} ^ {i}, \boldsymbol {a} _ {t} ^ {i}) \\ = \left[ \frac {\pi - \phi_ {t} ^ {j}}{2 \pi} \frac {(\pmb {a} _ {t} ^ {j}) ^ {\top} \pmb {a} ^ {*}}{\| \pmb {v} _ {t} ^ {j} \|} (\mathbf {I} - \pmb {w} _ {t} ^ {j} (\pmb {w} _ {t} ^ {j}) ^ {\top}) - \frac {\pi - \phi_ {t} ^ {i}}{2 \pi} \frac {(\pmb {a} _ {t} ^ {i}) ^ {\top} \pmb {a} ^ {*}}{\| \pmb {v} _ {t} ^ {i} \|} (\mathbf {I} - \pmb {w} _ {t} ^ {i} (\pmb {w} _ {t} ^ {i}) ^ {\top}) \right] \pmb {w} ^ {*}. \qquad (C. 1) \\ \end{array}
$$
Since  $\| \boldsymbol{w}^* \| = 1$  and  $\| \boldsymbol{a}_t^i \| \leq 5K_a$  in Lemma 12, we can bound  $h_t^{i,j}$  as