SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \left\| \boldsymbol {a} _ {t} - \boldsymbol {a} ^ {*} \right\| ^ {2} = \left\| \boldsymbol {a} _ {t - 1} - \boldsymbol {a} ^ {*} - \eta \frac {1}{N} \sum_ {i = 1} ^ {N} \nabla_ {\boldsymbol {a}} L (\boldsymbol {w} _ {t - 1} ^ {i}, \boldsymbol {a} _ {t - 1} ^ {i}) - \eta \epsilon_ {t} \right\| ^ {2} \\ = \left\| \boldsymbol {a} _ {t - 1} - \boldsymbol {a} ^ {*} - \eta \frac {1}{N} \sum_ {i = 1} ^ {N} \nabla_ {\boldsymbol {a}} L (\boldsymbol {w} _ {t - 1} ^ {i}, \boldsymbol {a} _ {t - 1} ^ {i}) \right\| ^ {2} - 2 \eta \langle \tilde {\boldsymbol {a}} _ {t - 1}, \boldsymbol {\epsilon} _ {t - 1} \rangle + \eta^ {2} \| \boldsymbol {\epsilon} _ {t - 1} \| ^ {2} \\ \leq \left(1 - \frac {\eta (\pi - 1)}{2 \pi} + \frac {\eta^ {2} (\pi + k - 1) ^ {2}}{2 \pi^ {2}}\right) \| \boldsymbol {a} _ {t - 1} - \boldsymbol {a} ^ {*} \| ^ {2} \\ - 2 \eta \left\langle \tilde {\boldsymbol {a}} _ {t - 1}, \boldsymbol {\epsilon} _ {t - 1} \right\rangle + \eta^ {2} \| \boldsymbol {\epsilon} _ {t - 1} \| ^ {2} + 6 \eta \epsilon^ {2} K _ {a} ^ {2} + \frac {\eta^ {2} \epsilon^ {4}}{\pi} \| \boldsymbol {a} ^ {*} \| ^ {2} \\ \leq \left(1 - \frac {\eta (\pi - 1)}{4 \pi}\right) \| \boldsymbol {a} _ {t - 1} - \boldsymbol {a} ^ {*} \| ^ {2} - 2 \eta \langle \tilde {\boldsymbol {a}} _ {t - 1}, \boldsymbol {\epsilon} _ {t - 1} \rangle + \eta^ {2} \| \boldsymbol {\epsilon} _ {t - 1} \| ^ {2} + 6 \eta \epsilon K _ {a} ^ {2} + \frac {\eta^ {2} \epsilon^ {2}}{\pi} \| \boldsymbol {a} ^ {*} \| \\ \leq \left(1 - \frac {\eta (\pi - 1)}{4 \pi}\right) ^ {t} \| \boldsymbol {a} _ {0} - \boldsymbol {a} ^ {*} \| ^ {2} + 2 \eta \left| \sum_ {s = 0} ^ {t - 1} \left(1 - \frac {\eta (\pi - 1)}{4 \pi}\right) ^ {s} \langle \tilde {\boldsymbol {a}} _ {s}, \boldsymbol {\epsilon} _ {s} \rangle \right| \\ + \eta^ {2} \sum_ {s = 0} ^ {t - 1} \left(1 - \frac {\eta (\pi - 1)}{4 \pi}\right) ^ {t - s} \| \boldsymbol {\epsilon} _ {s} \| ^ {2} \\ + \sum_ {s = 0} ^ {t - 1} \left(1 - \frac {\eta (\pi - 1)}{4 \pi}\right) ^ {t - 1 - s} \left(6 \eta \epsilon K _ {a} ^ {2} + 1 5 \eta^ {2} \epsilon^ {2} \| \boldsymbol {a} ^ {*} \| ^ {2}\right). \tag {E.17} \\ \end{array}
$$
Using Lemma 12 again, we can verify that
$$
\begin{array}{l} \| \tilde {\boldsymbol {a}} _ {t} \| \leq \| \boldsymbol {a} _ {t} - \boldsymbol {a} ^ {*} \| + \eta \frac {1}{N} \sum_ {i = 1} ^ {N} \| \nabla_ {\boldsymbol {a}} L (\boldsymbol {w} _ {t} ^ {i}, \boldsymbol {a} _ {t} ^ {i}) \| \\ \leq 6 K _ {a} + \frac {\eta}{2 \pi} \left(\| \left(\mathbf {1 1} ^ {\top} + (\pi - 1) \mathbf {I}\right) \left(\boldsymbol {a} _ {t} - \boldsymbol {a} ^ {*}\right) \| + (\pi - g \left(\phi_ {t} ^ {i}\right)) \| \boldsymbol {a} ^ {*} \|\right) \\ \leq 6 K _ {a} + \frac {\eta (\pi + k - 1)}{2 \pi} \cdot 6 \| \boldsymbol {a} ^ {*} \| + \frac {\eta}{2} \| \boldsymbol {a} ^ {*} \| \\ \leq 8 K _ {a}. \tag {E.18} \\ \end{array}
$$
Similar to (A.17) with  $\alpha_{s} = \left(1 - \frac{\eta(\pi - 1)}{4\pi}\right)^{t - s}$ , we can guarantee
$$
\begin{array}{l} \eta \left| \sum_ {s = 0} ^ {t - 1} \left(1 - \frac {\eta (\pi - 1)}{4 \pi}\right) ^ {t - s} \tilde {\boldsymbol {a}} _ {s} ^ {\top} \boldsymbol {\epsilon} _ {s} \right| \leq c \eta K _ {a} ^ {2} \sqrt {\frac {\log (2 / \delta)}{N} \sum_ {s = 0} ^ {t} \left(1 - \frac {\eta (\pi - 1)}{4 \pi}\right) ^ {t - s}} + \frac {c \eta K _ {a} ^ {2} \log (2 / \delta)}{N} \\ \leq 4 c K _ {a} ^ {2} \sqrt {\frac {\eta \log (2 / \delta)}{N}} \tag {E.19} \\ \end{array}
$$
with probability at least  $1 - \delta$ . Now take  $T_{a} = \frac{1}{\eta}\frac{8\pi}{\pi - 1}\log \left(\frac{1}{\epsilon}\right)$ . Plugging (E.19) into (E.17), together with the concentration inequality (A.19), we have for any  $T_{a}\leq t\leq \widetilde{O} (\eta^{-2})$
$$
\begin{array}{l} \left\| \boldsymbol {a} _ {t} - \boldsymbol {a} ^ {*} \right\| ^ {2} \leq \left(1 - \frac {\eta (\pi - 1)}{4 \pi}\right) ^ {t} \left\| \boldsymbol {a} _ {0} - \boldsymbol {a} ^ {*} \right\| ^ {2} + 4 c K _ {a} ^ {2} \sqrt {\frac {\eta \log (2 / \delta)}{N}} \\ + \frac {2 \pi c \eta k \log (t / \delta)}{N} K _ {a} ^ {2} + 1 2 \pi \epsilon K _ {a} ^ {2} + 2 \eta \epsilon^ {2} K _ {a} ^ {2} \\ \leq \epsilon \| \boldsymbol {a} _ {0} - \boldsymbol {a} ^ {*} \| ^ {2} + \sqrt {\epsilon} K _ {a} ^ {2} + \epsilon^ {2} K _ {a} ^ {2} + 1 2 \pi \epsilon K _ {a} ^ {2} + 2 \eta \epsilon^ {2} K _ {a} ^ {2} \\ \leq 2 \sqrt {\epsilon} K _ {a} ^ {2}. \tag {E.20} \\ \end{array}
$$
With this upper bound, we can refine (E.18) by
$$
\| \tilde {\boldsymbol {a}} _ {t} \| ^ {2} \leq 2 \| \boldsymbol {a} _ {t} - \boldsymbol {a} ^ {*} \| ^ {2} + 2 \eta^ {2} \frac {1}{N} \sum_ {i = 1} ^ {N} \| \nabla_ {\boldsymbol {a}} L (\boldsymbol {w} _ {t} ^ {i}, \boldsymbol {a} _ {t} ^ {i}) \| ^ {2} \leq 8 \sqrt {\epsilon} K _ {a} ^ {2}.
$$
Consequently, the bound (E.19) can be improved to
$$
\eta \left| \sum_ {s = 0} ^ {t - 1} \left(1 - \frac {\eta (\pi - 1)}{4 \pi}\right) ^ {t - s} \tilde {\boldsymbol {a}} _ {s} ^ {\top} \boldsymbol {\epsilon} _ {s} \right| \leq 4 c K _ {a} ^ {2} \sqrt {\frac {\eta \epsilon^ {1 / 2} \log (2 / \delta)}{N}}.
$$
Then we can guarantee that for any  $2T_{a}\leq t\leq \widetilde{O} (\eta^{-2})$
$$
\begin{array}{l} \left\| \boldsymbol {a} _ {t} - \boldsymbol {a} ^ {*} \right\| ^ {2} \leq \left(1 - \frac {\eta (\pi - 1)}{4 \pi}\right) ^ {t} \left\| \boldsymbol {a} _ {0} - \boldsymbol {a} ^ {*} \right\| ^ {2} + 4 c K _ {a} ^ {2} \sqrt {\frac {\eta \epsilon \log (2 / \delta)}{N}} \\ + \frac {2 \pi c \eta k \log (t / \delta)}{N} K _ {a} ^ {2} + 1 2 \pi \epsilon K _ {a} ^ {2} + 2 \eta \epsilon^ {2} K _ {a} ^ {2} \\ \leq \epsilon \| \boldsymbol {a} _ {0} - \boldsymbol {a} ^ {*} \| ^ {2} + \epsilon^ {3 / 4} K _ {a} ^ {2} + \epsilon^ {2} K _ {a} ^ {2} + 1 2 \pi \epsilon K _ {a} ^ {2} + 2 \eta \epsilon^ {2} K _ {a} ^ {2} \\ \lesssim \epsilon^ {4 / 3} K _ {a} ^ {2}. \tag {E.21} \\ \end{array}
$$
Repeating the refinement from (E.20) to (E.21) for at most  $\log(1/\epsilon)$  times, we can show finish the proof.
E.3 Conclusion
Theorem 3 restated. Suppose the initial point  $(\pmb{v}_0, \pmb{a}_0)$  satisfies  $\pmb{a}_0^\top \pmb{a}^* \geq \gamma_a / 32$  and  $\phi_0 \leq \tilde{\phi}^\ell$ . For any  $\epsilon > 0$ , we choose  $\eta = \frac{1}{ck^2 \| \pmb{a}^* \|^2} \frac{N\epsilon}{d\log(dN / \epsilon\delta)}$  for some absolute constant  $c > 0$ . If  $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 < \min \{N^{-1}, d^{-1}, dk^{-2}\}$ , then  $\ell(\pmb{v}_T, \pmb{a}_T) = O(\epsilon K_a^2)$  holds with probability at least  $1 - \delta$  where  $T = \widetilde{O}\left(\frac{dk^4}{N\epsilon}\right)$ .
Proof. According to the dynamic of  $\pmb{a}_t$ , we know
$$
\begin{array}{l} \mathbf {1} ^ {\top} (\boldsymbol {a} _ {t} - \boldsymbol {a} ^ {*}) = \left(1 - \frac {\eta (\pi + k - 1)}{2 \pi}\right) \mathbf {1} ^ {\top} (\boldsymbol {a} _ {t - 1} - \boldsymbol {a} ^ {*}) + \frac {\eta}{2 \pi} \left(\frac {1}{N} \sum_ {i = 1} ^ {N} g \left(\phi_ {t - 1} ^ {i}\right) - \pi\right) \mathbf {1} ^ {\top} \boldsymbol {a} ^ {*} + \eta \mathbf {1} ^ {\top} \boldsymbol {\epsilon} _ {t - 1} \\ = \left(1 - \frac {\eta (\pi + k - 1)}{2 \pi}\right) ^ {t} \mathbf {1} ^ {\top} \left(\boldsymbol {a} _ {0} - \boldsymbol {a} ^ {*}\right) + \eta \sum_ {s = 0} ^ {t - 1} \left(1 - \frac {\eta (\pi + k - 1)}{2 \pi}\right) ^ {t - 1 - s} \mathbf {1} ^ {\top} \boldsymbol {\epsilon} _ {s} \\ + \frac {\eta}{2 \pi} \sum_ {s = 0} ^ {t - 1} \left(1 - \frac {\eta (\pi + k - 1)}{2 \pi}\right) ^ {t - 1 - s} \left(\frac {1}{N} \sum_ {i = 1} ^ {N} g \left(\phi_ {s} ^ {i}\right) - \pi\right) \mathbf {1} ^ {\top} \mathbf {a} ^ {*}. \tag {E.22} \\ \end{array}
$$
It follows from (E.14) that
$$
\left| \frac {\eta}{2 \pi} \sum_ {s = 0} ^ {t - 1} \left(1 - \frac {\eta (\pi + k - 1)}{2 \pi}\right) ^ {t - 1 - s} \left(\frac {1}{N} \sum_ {i = 1} ^ {N} g \left(\phi_ {s} ^ {i}\right) - \pi\right) \right| \leq \frac {5 \pi \epsilon}{\pi + k - 1}. \tag {E.23}
$$
Using Lemma 11(1), it holds that
$$
\eta \left| \sum_ {s = 0} ^ {t - 1} \left(1 - \frac {\eta (\pi + k - 1)}{2 \pi}\right) ^ {t - 1 - s} \mathbf {1} ^ {\top} \boldsymbol {\epsilon} _ {s} \right| \leq c K _ {a} \sqrt {\frac {\eta \log (1 / \delta)}{N}}. \tag {E.24}
$$
Plugging (E.23) and (E.24) into (E.22), and taking  $T_{1} = \frac{1}{\eta}\frac{2\pi}{\pi + k - 1}\log \left(\frac{1}{\epsilon}\right)$ , we can have
$$
\begin{array}{l} \left| \mathbf {1} ^ {\top} \left(\boldsymbol {a} _ {t} - \boldsymbol {a} ^ {*}\right) \right| \leq \epsilon \left| \mathbf {1} ^ {\top} \left(\boldsymbol {a} _ {0} - \boldsymbol {a} ^ {*}\right) \right| + \sqrt {\epsilon} K _ {a} + \frac {5 \pi \epsilon}{\pi + k - 1} \left| \mathbf {1} ^ {\top} \boldsymbol {a} ^ {*} \right| \\ \leq \epsilon \sqrt {k} K _ {a} + \sqrt {\epsilon} K _ {a} + \frac {5 \pi \epsilon}{\pi + k - 1} | \mathbf {1} ^ {\top} \mathbf {a} ^ {*} | \\ \leq 2 \sqrt {\epsilon} K _ {a}, \tag {E.25} \\ \end{array}
$$
where we used  $|\mathbf{1}^{\top}\mathbf{a}_{0}| \leq |\mathbf{1}^{\top}\mathbf{a}^{*}|$  and  $\epsilon < 1 / k$ . Invoking Lemma 6 and the normalized network in (2), we know
$$
\begin{array}{l} \ell (\boldsymbol {v}, \boldsymbol {a}) = L (\boldsymbol {w}, \boldsymbol {a}) = \frac {1}{2} \left[ \frac {\pi - 1}{2 \pi} \| \boldsymbol {a} ^ {*} \| ^ {2} + \frac {\pi - 1}{2 \pi} \| \boldsymbol {a} \| ^ {2} - \frac {g (\phi) - 1}{\pi} \boldsymbol {a} ^ {\top} \boldsymbol {a} ^ {*} \right. \\ \left. + \frac {1}{2 \pi} \left(\mathbf {1} ^ {\top} \mathbf {a} ^ {*}\right) ^ {2} + \frac {1}{2 \pi} \left(\mathbf {1} ^ {\top} \mathbf {a}\right) ^ {2} - \frac {1}{\pi} \left(\mathbf {1} ^ {\top} \mathbf {a}\right) \left(\mathbf {1} ^ {\top} \mathbf {a}\right) \right] \\ = \frac {1}{2} \left[ \frac {\pi - 1}{2 \pi} \| \boldsymbol {a} - \boldsymbol {a} ^ {*} \| ^ {2} + \frac {\pi - g (\phi)}{\pi} \boldsymbol {a} ^ {\top} \boldsymbol {a} ^ {*} + \frac {1}{2 \pi} | \mathbf {1} ^ {\top} (\boldsymbol {a} - \boldsymbol {a} ^ {*}) | ^ {2} \right]. \tag {E.26} \\ \end{array}
$$
Similar to (E.14), for any  $t \geq T_v$ , we also have
$$
\pi - g \left(\phi_ {t}\right) = g (0) - g \left(\phi_ {t}\right) \leq \left(\pi - \phi_ {t}\right) \sin \phi_ {t} \cdot \phi_ {t} \leq \pi \epsilon \cdot \phi_ {t} \leq \pi \epsilon . \tag {E.27}
$$
Invoking Lemma 4 and 5 to (E.26), together with (E.27), we conclude that for  $T = \widetilde{O}\left(\frac{dk^4}{N\epsilon}\right)$ ,
$$
\ell (\boldsymbol {v} _ {T}, \boldsymbol {a} _ {T}) \leq \frac {1}{2} \left(\frac {1 5}{2} \epsilon K _ {a} + \epsilon \| \boldsymbol {a} _ {T} \| \| \boldsymbol {a} ^ {*} \| + 4 \epsilon K _ {a} ^ {2}\right) \lesssim \epsilon \| \boldsymbol {a} ^ {*} \| ^ {2},
$$