SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $\overline{z}^{(i)}\coloneqq \frac{z^{(i)}}{\|\pmb{z}^{(i)}\|}$  is distributed uniformly over  $\mathbb{S}^{d - 1}$ . From Lemma 24 we have  $\| \mathbf{I}_d - \hat{\pmb{\Sigma}}^{-1 / 2}\pmb{\Sigma}^{1 / 2}\| \lesssim \sqrt{\frac{d}{n'}}$  with probability at least  $1 - 2e^{-d}$ . On the other hand, from Lemma 17 we know with probability at least  $1 - e^{-d}$ , for any  $\overline{w}\in \mathbb{S}^{d - 1}$  at most  $\tilde{\mathcal{O}} (d)$  of the labeled samples have  $|\langle \overline{w},z^{(i)}\rangle |\lesssim \sqrt{d} /n$ . Recall that  $n^{\prime}\gtrsim n^{2}$  when using the ReLU activation. This is precisely why we make this choice for the ReLU activation, as we need to balance the RHS of (D.7) which is of order  $\sqrt{d / n'}$  with the LHS of (D.7) which should at most be of order  $\sqrt{d} /n$  if we want to ensure only  $\tilde{\mathcal{O}} (d)$  samples satisfy the bound. When  $n^{\prime} = n^{2}$  we can balance these two terms, thus with probability at least  $1 - 3e^{-d}$  the sign change can occur for at most  $\tilde{\mathcal{O}} (d)$  many samples, and
$$
\frac {1}{n} \sum_ {i = 1} ^ {n} \left\{\phi^ {\prime} \left(\left\langle \overline {{\boldsymbol {w}}} ^ {t}, \boldsymbol {z} ^ {(i)} \right\rangle\right) - \phi^ {\prime} \left(\left\langle \overline {{\boldsymbol {w}}} ^ {t}, \tilde {\boldsymbol {z}} ^ {(i)} \right\rangle\right) \right\} ^ {2} \leq \tilde {\mathcal {O}} \left(\frac {d}{n}\right).
$$
In this case, we end up with
$$
\mathcal {E} _ {2} \geq - K \| \boldsymbol {\nu} (\overline {{\boldsymbol {w}}} ^ {t}) \| \tilde {\mathcal {O}} (\sqrt {d / n}).
$$
Bounding  $\mathcal{E}_3$  for ReLU and Lipschitz  $\phi^\prime$  is identical. In both cases, by Cauchy-Schwartz,
$$
\begin{array}{l} \mathcal {E} _ {3} \geq - \sqrt {\frac {1}{n} \sum_ {i = 1} ^ {n} \phi^ {\prime} (\langle \overline {{\boldsymbol {w}}}, \boldsymbol {z} ^ {(i)} \rangle) ^ {2} y _ {K} ^ {(i)}} ^ {2} \sqrt {\frac {1}{n} \sum_ {i = 1} ^ {n} \left\langle \boldsymbol {z} ^ {(i)} - \tilde {\boldsymbol {z}} ^ {(i)} , \boldsymbol {\nu} (\overline {{\boldsymbol {w}}} ^ {t}) \right\rangle^ {2}} \\ \gtrsim - K \| \frac {1}{n} \sum_ {i = 1} ^ {n} \pmb {z} ^ {(i)} \pmb {z} ^ {(i) ^ {\top}} \| \| \mathbf {I} _ {d} - \hat {\pmb {\Sigma}} ^ {- 1 / 2} \pmb {\Sigma} ^ {1 / 2} \| \| \pmb {\nu} (\overline {{\pmb {w}}} ^ {t}) \| \\ \gtrsim - K \| \boldsymbol {\nu} (\overline {{\boldsymbol {w}}} ^ {t}) \| \sqrt {d / n ^ {\prime}}, \\ \end{array}
$$
which holds on the intersection of event  $\mathcal{G}$  and of Lemma 23. At last, using the bound on  $\| \pmb {\nu}(\overline{\pmb{w}}^t)\|$  from (D.6), we obtain
$$
\mathcal {E} _ {2} \wedge \mathcal {E} _ {3} \geq - \lambda_ {\max } (\boldsymbol {\Sigma}) \tilde {\mathcal {O}} (\sqrt {d / n}),
$$
with probability at least  $1 - \mathcal{O}(d^{-q})$ .
Step 3. Analyzing the Convergence. As a result of the previous steps, we have established
$$
\frac {\mathrm {d} \left\langle \overline {{\boldsymbol {w}}} ^ {t} , \overline {{\boldsymbol {u}}} \right\rangle}{\mathrm {d} t} \geq \left\langle \boldsymbol {\nu} (\overline {{\boldsymbol {w}}} ^ {t}), \mathbb {E} \left[ \phi^ {\prime} (\langle \overline {{\boldsymbol {w}}} ^ {t}, \boldsymbol {z} \rangle) y \boldsymbol {z} \right] \right\rangle - \lambda_ {\max } (\boldsymbol {\Sigma}) \tilde {\mathcal {O}} (\sqrt {d / n}).
$$
Thanks to Lemma 11, we can write
$$
\begin{array}{l} \mathbb {E} \left[ \phi^ {\prime} (\langle \overline {{\boldsymbol {w}}}, \boldsymbol {z} \rangle) y \boldsymbol {z} \right] = \mathbb {E} \left[ \phi^ {\prime} (\langle \overline {{\boldsymbol {w}}}, \boldsymbol {z} \rangle) g (\langle \overline {{\boldsymbol {u}}}, \boldsymbol {z} \rangle) \boldsymbol {z} \right] \\ = - \zeta_ {\phi , g} (\langle \overline {{\boldsymbol {w}}}, \overline {{\boldsymbol {u}}} \rangle) \overline {{\boldsymbol {u}}} - \psi_ {\phi , g} (\langle \overline {{\boldsymbol {w}}} ^ {t}, \overline {{\boldsymbol {u}}} \rangle) \overline {{\boldsymbol {w}}}, \\ \end{array}
$$
where  $\zeta_{\phi ,g}$  and  $\psi_{\phi ,g}$  were introduced in (D.2) and (D.3) respectively. Recall the definition of  $\pmb {\nu}(\overline{\pmb{w}}^t)$
$$
\boldsymbol {\nu} (\overline {{\boldsymbol {w}}} ^ {t}) := \left(\mathbf {I} _ {d} - \overline {{\boldsymbol {w}}} ^ {t} \overline {{\boldsymbol {w}}} ^ {t ^ {\top}}\right) \hat {\boldsymbol {\Sigma}} \left(\mathbf {I} _ {d} - \overline {{\boldsymbol {w}}} ^ {t} \overline {{\boldsymbol {w}}} ^ {t ^ {\top}}\right) \overline {{\boldsymbol {u}}}.
$$
$$
\begin{array}{l} \frac {\mathrm {d} \left\langle \overline {{\boldsymbol {w}}} ^ {t} , \overline {{\boldsymbol {u}}} \right\rangle}{\mathrm {d} t} \geq - \zeta_ {\phi , g} \left(\left\langle \overline {{\boldsymbol {w}}} ^ {t}, \overline {{\boldsymbol {u}}} \right\rangle\right) \overline {{\boldsymbol {u}}} ^ {\top} \left(\mathbf {I} _ {d} - \overline {{\boldsymbol {w}}} ^ {t} \overline {{\boldsymbol {w}}} ^ {t} ^ {\top}\right) \hat {\boldsymbol {\Sigma}} \left(\mathbf {I} _ {d} - \overline {{\boldsymbol {w}}} ^ {t} \overline {{\boldsymbol {w}}} ^ {t ^ {T}}\right) \overline {{\boldsymbol {u}}} - \lambda_ {\max } (\boldsymbol {\Sigma}) \tilde {\mathcal {O}} \left(\sqrt {d / n}\right) \\ \geq c \left\langle \overline {{\boldsymbol {w}}} ^ {t}, \overline {{\boldsymbol {u}}} \right\rangle^ {s - 1} \left\langle \overline {{\boldsymbol {u}}} _ {\perp} ^ {t}, \hat {\boldsymbol {\Sigma}} \overline {{\boldsymbol {u}}} _ {\perp} ^ {t} \right\rangle - \lambda_ {\max } (\boldsymbol {\Sigma}) \tilde {\mathcal {O}} (\sqrt {d / n}) \quad \text {(B y A s s u m p t i o n 2)} \\ \geq c \lambda_ {\min } \left(\hat {\boldsymbol {\Sigma}}\right) \left\langle \overline {{\boldsymbol {w}}} ^ {t}, \overline {{\boldsymbol {u}}} \right\rangle^ {s - 1} \left(1 - \left\langle \overline {{\boldsymbol {w}}} ^ {t}, \overline {{\boldsymbol {u}}} \right\rangle^ {2}\right) - \lambda_ {\max } (\boldsymbol {\Sigma}) \tilde {\mathcal {O}} (\sqrt {d / n}), \\ \end{array}
$$
where  $\overline{\pmb{u}}_{\perp}^{t}\coloneqq \overline{\pmb{u}} -\langle \overline{\pmb{u}},\overline{\pmb{w}}^{t}\rangle \overline{\pmb{w}}^{t}$ . Moreover, from Lemma 23, we have
$$
\begin{array}{l} \lambda_ {\min } \left(\hat {\boldsymbol {\Sigma}}\right) \geq \lambda_ {\min } (\boldsymbol {\Sigma}) \left(\frac {1}{2} - \sqrt {\frac {\operatorname {t r} (\boldsymbol {\Sigma})}{n ^ {\prime} \lambda_ {\min } (\boldsymbol {\Sigma})}}\right) \\ \geq \lambda_ {\min } (\boldsymbol {\Sigma}) \left(\frac {1}{2} - \sqrt {\frac {d \varkappa (\boldsymbol {\Sigma})}{n ^ {\prime}}}\right), \\ \end{array}
$$
with probability at least  $1 - e^{-n' / 8}$ . Hence, for  $n' \gtrsim d\varkappa(\Sigma)$  we have  $\lambda_{\min}(\hat{\Sigma}) \gtrsim \lambda_{\min}(\Sigma)$ , and consequently,
$$
\frac {\mathrm {d} \left\langle \overline {{\boldsymbol {w}}} ^ {t} , \overline {{\boldsymbol {u}}} \right\rangle}{\mathrm {d} t} \geq c ^ {\prime} \lambda_ {\min } (\boldsymbol {\Sigma}) \left\langle \overline {{\boldsymbol {w}}} ^ {t}, \overline {{\boldsymbol {u}}} \right\rangle^ {s - 1} \left(1 - \left\langle \overline {{\boldsymbol {w}}} ^ {t}, \overline {{\boldsymbol {u}}} \right\rangle^ {2}\right) - \lambda_ {\max } (\boldsymbol {\Sigma}) \tilde {\mathcal {O}} (\sqrt {d / n}),
$$
where  $c'$  is a universal constant. Notice that the first term in the RHS above denotes the signal, while the second term denotes the noise. We want to ensure the noise remains smaller than the signal throughout the trajectory, which leads to the convergence of  $\overline{\boldsymbol{w}}^t$  to  $\overline{\boldsymbol{u}}$ . Notice that the signal term is first increasing, then decreasing for  $\langle \overline{\boldsymbol{w}}^t, \overline{\boldsymbol{u}} \rangle \in [0,1]$ . Thus, it suffices to ensure the noise is smaller than the signal on the two ends of the interval, i.e. at time  $t = 0$  and at time  $t = T$  where  $\langle \overline{\boldsymbol{w}}^T, \overline{\boldsymbol{u}} \rangle = 1 - \varepsilon$ . At initialization, this condition leads to
$$
n \gtrsim C d \varkappa (\boldsymbol {\Sigma}) ^ {2} \left\langle \overline {{\boldsymbol {w}}} ^ {0}, \overline {{\boldsymbol {u}}} \right\rangle^ {2 (1 - s)},
$$
and at time  $t = T$ , leads to
$$
n \gtrsim \frac {C d \varkappa (\pmb {\Sigma}) ^ {2}}{\varepsilon^ {2}},
$$
where  $C$  hides constant depending only on  $s$  and at most polylogarithmic factors of  $d$ . Thus, we have established the sample complexity as presented by Theorem 5.
It remains to obtain the convergence time. With the above sample complexity, we have
$$
\frac {\mathrm {d} \left\langle \overline {{\boldsymbol {w}}} ^ {t} , \overline {{\boldsymbol {u}}} \right\rangle}{\mathrm {d} t} \geq c ^ {\prime \prime} \lambda_ {\min } (\boldsymbol {\Sigma}) \left\langle \overline {{\boldsymbol {w}}} ^ {t}, \overline {{\boldsymbol {u}}} \right\rangle^ {s - 1} (1 - \left\langle \overline {{\boldsymbol {w}}} ^ {t}, \overline {{\boldsymbol {u}}} \right\rangle^ {2}),
$$
where  $c''$  is a universal constant. The rest of the proof follows by integration and is identical to the proof of Proposition 4 in Appendix D.1.
D.4 Proof of Corollary 6
The proof follows immediately from Theorem 5 and the following lemma which describes how  $\langle \overline{\boldsymbol{w}}^0,\overline{\boldsymbol{u}}\rangle$  behaves under different regimes of  $r_1$  and  $r_2$ .
Lemma 19. Suppose  $\pmb{\Sigma}$  follows the  $(\kappa, \theta)$ -spiked model,  $\mathbf{w}^0$  is sampled uniformly from  $\mathbb{S}^{d-1}$ ,  $n' \gtrsim d$ , and there exist universal constants  $C_2, C_2', C_3, C_3' > 0$  such that
$$
C _ {2} d ^ {r _ {2}} \leq \kappa \leq C _ {2} ^ {\prime} d ^ {r _ {2}} \quad a n d \quad C _ {3} d ^ {- r _ {1}} \leq \langle \boldsymbol {u}, \boldsymbol {\theta} \rangle \leq C _ {3} ^ {\prime} d ^ {- r _ {1}}
$$
for  $r_1 \in [0, 1/2]$  and  $r_2 \in [0, 1]$ . Then, conditioned on  $\langle \overline{\boldsymbol{w}}^0, \overline{\boldsymbol{u}} \rangle > 0$ , with any arbitrarily large constant probability  $1 - \delta$ , for sufficiently large  $d$  (that depends on  $\delta$ ) we have
$$
\left\langle \overline {{\boldsymbol {w}}} ^ {0}, \overline {{\boldsymbol {u}}} \right\rangle \gtrsim \left\{ \begin{array}{l l} d ^ {- 1 / 2} & 0 \leq r _ {2} <   r _ {1} \\ d ^ {r _ {2} - r _ {1} - 1 / 2} & r _ {1} <   r _ {2} <   2 r _ {1}. \\ d ^ {(r _ {2} - 1) / 2} & 2 r _ {1} <   r _ {2} <   1 \end{array} \right. \tag {D.8}
$$
$$
\left\langle \overline {{\boldsymbol {w}}} ^ {0}, \overline {{\boldsymbol {u}}} \right\rangle = \frac {\left\langle \hat {\boldsymbol {\Sigma}} ^ {1 / 2} \boldsymbol {w} ^ {0} , \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {u} \right\rangle}{\| \hat {\boldsymbol {\Sigma}} ^ {1 / 2} \boldsymbol {w} ^ {0} \| \| \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {u} \|}.
$$
Recall that we are conditioning our arguments on  $\langle \overline{\boldsymbol{w}}^0,\overline{\boldsymbol{u}}\rangle >0$ , hence the numerator of the above fraction is positive. To translate the sample complexities of Theorems 5 and 7 to the spiked model, our goal is to lower bound  $\langle \overline{\boldsymbol{w}}^0,\overline{\boldsymbol{u}}\rangle$  in terms of  $d$ ,  $r_1$ , and  $r_2$ .
We begin by observing that
$$
\| \hat {\boldsymbol {\Sigma}} ^ {1 / 2} \boldsymbol {w} \| \leq \| \hat {\boldsymbol {\Sigma}} ^ {1 / 2} \boldsymbol {\Sigma} ^ {- 1 / 2} \| \| \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {w} \| \lesssim \| \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {w} \|,
$$
where the last inequality holds on the event of Lemma 24, which happens with probability at least  $1 - 2e^{-d}$ . Consequently,
$$
\left\langle \overline {{\boldsymbol {w}}} ^ {0}, \overline {{\boldsymbol {u}}} \right\rangle \gtrsim \frac {\left\langle \hat {\boldsymbol {\Sigma}} ^ {1 / 2} \boldsymbol {w} ^ {0} , \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {u} \right\rangle}{\| \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {w} ^ {0} \| \| \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {u} \|} = \frac {\left\langle \boldsymbol {w} ^ {0} , \boldsymbol {\Sigma} \boldsymbol {u} \right\rangle + \left\langle \boldsymbol {w} ^ {0} , (\hat {\boldsymbol {\Sigma}} ^ {1 / 2} - \boldsymbol {\Sigma} ^ {1 / 2}) \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {u} \right\rangle}{\| \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {w} ^ {0} \| \| \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {u} \|}.
$$
Furthermore, due to the Markov inequality,
$$
\mathbb {P} \left(\left\langle \boldsymbol {w} ^ {0}, \boldsymbol {\theta} \right\rangle^ {2} \geq \frac {C _ {1}}{d}\right) \leq 1 / C _ {1}.
$$
Similarly (by conditioning on  $\hat{\Sigma}$ )
$$
\mathbb {P} \left(\left\langle \boldsymbol {w} ^ {0}, \left(\hat {\boldsymbol {\Sigma}} ^ {1 / 2} - \boldsymbol {\Sigma} ^ {1 / 2}\right) \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {u} \right\rangle^ {2} \geq \frac {C _ {1} \| \left(\hat {\boldsymbol {\Sigma}} ^ {1 / 2} - \boldsymbol {\Sigma} ^ {1 / 2}\right) \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {u} \| ^ {2}}{d}\right) \leq 1 / C _ {1}.
$$
Additionally, on the event of Lemma 24,
$$
\| (\hat {\pmb {\Sigma}} ^ {1 / 2} - \pmb {\Sigma} ^ {1 / 2}) \pmb {\Sigma} ^ {1 / 2} \pmb {u} \| \leq \| \hat {\pmb {\Sigma}} ^ {1 / 2} \pmb {\Sigma} ^ {- 1 / 2} - \mathbf {I} _ {d} \| \| \pmb {\Sigma} \pmb {u} \| \lesssim \sqrt {d / n ^ {\prime}} \| \pmb {\Sigma} \pmb {u} \|.
$$
Therefore, on the above events, for some absolute constant  $C > 0$
$$
\begin{array}{l} \langle \overline {{\boldsymbol {w}}} ^ {0}, \overline {{\boldsymbol {u}}} \rangle \gtrsim \frac {\left\langle \boldsymbol {w} ^ {0} , \boldsymbol {\Sigma} \boldsymbol {u} \right\rangle - C \| \boldsymbol {\Sigma} \boldsymbol {u} \| \sqrt {1 / n ^ {\prime}}}{\sqrt {\frac {1 + C _ {2} ^ {\prime} C _ {1}}{1 + \kappa}} \| \boldsymbol {\Sigma} \boldsymbol {u} \|} \\ \gtrsim \frac {\left\langle \boldsymbol {w} ^ {0} , \boldsymbol {u} \right\rangle + \kappa \left\langle \boldsymbol {w} ^ {0} , \boldsymbol {\theta} \right\rangle \left\langle \boldsymbol {u} , \boldsymbol {\theta} \right\rangle - C (1 + \kappa | \left\langle \boldsymbol {u} , \boldsymbol {\theta} \right\rangle |) \sqrt {1 / n ^ {\prime}}}{\sqrt {1 + C _ {2} ^ {\prime} C _ {1}} \sqrt {1 + \kappa \left\langle \boldsymbol {u} , \boldsymbol {\theta} \right\rangle^ {2}}}. \\ \end{array}
$$
Recall that  $C_2 d^{r_2} \leq \kappa \leq C_2' d^{r_2}$  and  $C_3 d^{-r_1} \leq \langle \pmb{u}, \pmb{\theta} \rangle \leq C_3' d^{-r_1}$  (notice that changing  $\pmb{\theta}$  to  $-\pmb{\theta}$  does not change the spiked model of Assumption 1, thus we can assume  $\langle \pmb{u}, \pmb{\theta} \rangle \geq 0$  without loss of generality). Then,
$$
\langle \overline {{\boldsymbol {w}}} ^ {0}, \overline {{\boldsymbol {u}}} \rangle \gtrsim \frac {\left\langle \boldsymbol {w} ^ {0} , \boldsymbol {u} \right\rangle + \kappa \left\langle \boldsymbol {w} ^ {0} , \boldsymbol {\theta} \right\rangle \left\langle \boldsymbol {u} , \boldsymbol {\theta} \right\rangle - C (1 + \kappa \left\langle \boldsymbol {u} , \boldsymbol {\theta} \right\rangle) \sqrt {1 / n ^ {\prime}}}{\sqrt {1 + C _ {2} ^ {\prime} C _ {1}} \sqrt {1 + C _ {2} ^ {\prime} C _ {3} ^ {\prime 2} d ^ {r _ {2} - 2 r _ {1}}}}. \tag {D.9}
$$
The last term in the numerator can be made arbitrarily small by sufficiently large  $n'$ , hence we focus on other terms for now. Intuitively, when  $r_2 < 2r_1$ , the denominator is of constant order. If additionally  $r_2 < r_1$ , the dominant term in the numerator is  $\langle \pmb{w}^0, \pmb{u} \rangle$  and  $\langle \overline{\pmb{w}}^0, \overline{\pmb{u}} \rangle \asymp 1 / \sqrt{d}$ , otherwise the dominant term is  $\kappa \langle \pmb{w}^0, \pmb{\theta} \rangle \langle \pmb{u}, \pmb{\theta} \rangle$  and  $\langle \overline{\pmb{w}}^0, \overline{\pmb{u}} \rangle \asymp d^{r_2 - r_1 - 1/2}$ . On the other hand, when  $r_2 > 2r_1$ , the denominator is of order  $d^{r_2/2 - r_1}$ , and once again the dominant term of the numerator is  $\kappa \langle \pmb{w}^0, \pmb{\theta} \rangle \langle \pmb{u}, \pmb{\theta} \rangle$ , therefore  $\langle \overline{\pmb{w}}^0, \overline{\pmb{u}} \rangle \asymp d^{(r_2 - 1)/2}$ . Using this intuition, we analyze each of the following regimes separately.
Case 1.  $0 < r_2 < r_1$ : In this case, by [BBSS22, Lemma A.7] we have  $\left|\left\langle \pmb{w}^0,\pmb{u}\right\rangle \right| \geq c / \sqrt{d}$  with probability at least  $1 - 4c$ . On the intersection of all considered events with  $\left\langle \overline{\pmb{w}}^0,\overline{\pmb{u}}\right\rangle > 0$ , and for sufficiently large  $d$  and  $n' \gtrsim d$ , we must have  $\left\langle \pmb{w}^0,\pmb{u}\right\rangle > 0$  (otherwise  $\left\langle \overline{\pmb{w}}^0,\overline{\pmb{u}}\right\rangle < 0$ ). Thus by plugging the values in (D.9),
$$
\left\langle \overline {{\boldsymbol {w}}} ^ {0}, \overline {{\boldsymbol {u}}} \right\rangle \gtrsim \frac {c - \sqrt {C _ {1}} C _ {2} ^ {\prime} C _ {3} ^ {\prime} d ^ {r _ {2} - r _ {1}} - C \left(1 + C _ {2} ^ {\prime} C _ {3} ^ {\prime} d ^ {r _ {2} - r _ {1}}\right) \sqrt {d / n ^ {\prime}}}{\sqrt {d} \sqrt {1 + C _ {2} ^ {\prime} C _ {1}} \sqrt {1 + C _ {2} ^ {\prime} C _ {3} ^ {\prime 2} d ^ {r _ {2} - 2 r _ {1}}}} \gtrsim \frac {1}{\sqrt {d}}. \tag {D.10}
$$
The intersection of all desired events and  $\langle \overline{\boldsymbol{w}}^0,\overline{\boldsymbol{u}}\rangle >0$  happens with probability at least  $\frac{1}{2} -4c - 2 / C_1 - 2e^{-d}$ , thus conditioned on  $\langle \overline{\boldsymbol{w}}^0,\overline{\boldsymbol{u}}\rangle$  the probability is at least  $1 - 8c - 4 / C_{1} - 4e^{-d}$ . Choosing sufficiently small  $c$ , large  $C_1$ , and respectively large  $d$  and  $n^\prime \gtrsim d$  with sufficiently large absolute constant, we can arbitrarily increase the (constant) probability of success. Thus the analysis of this regime is complete.
Case 2.  $r_1 < r_2 < 2r_1$ : This time we use the fact that  $|\langle \pmb{w}^0, \pmb{u} \rangle| \leq \sqrt{C_1 / d}$  with probability at least  $1 - 1 / C_1$ , and  $|\langle \pmb{w}^0, \pmb{\theta} \rangle| \geq c / \sqrt{d}$  with probability at least  $1 - 4c$ . By an argument similar to the previous case, for sufficiently large  $d$  and  $n' \gtrsim d$ ,  $\langle \overline{\pmb{w}}^0, \overline{\pmb{u}} \rangle > 0$  implies  $\langle \pmb{w}^0, \pmb{\theta} \rangle > 0$ , hence by (D.9)
$$
\langle \overline {{\boldsymbol {w}}} ^ {0}, \overline {{\boldsymbol {u}}} \rangle \gtrsim \frac {- \sqrt {C _ {1}} + c C _ {2} C _ {3} d ^ {r _ {2} - r _ {1}} - C \left(1 + C _ {2} ^ {\prime} C _ {3} ^ {\prime} d ^ {r _ {2} - r _ {1}}\right) \sqrt {d / n ^ {\prime}}}{\sqrt {d} \sqrt {1 + C _ {2} ^ {\prime} C _ {1}} \sqrt {1 + C _ {2} ^ {\prime} C _ {3} ^ {\prime 2} d ^ {r _ {2} - 2 r _ {1}}}} \gtrsim d ^ {r _ {2} - r _ {1} - 1 / 2}, \tag {D.11}
$$
with probability at least  $1 - 8c - 4 / C_1 - 4e^{-d}$  when conditioned on  $\langle \overline{\boldsymbol{w}}^0, \overline{\boldsymbol{u}} \rangle > 0$ .
Case 3.  $2r_{1} < r_{2} < 1$ : Once again recall (D.9). To bound the numerator, we repeat the exact same argument as in the previous case, thus
$$
\left\langle \overline {{\boldsymbol {w}}} ^ {0}, \overline {{\boldsymbol {u}}} \right\rangle \geq \frac {- \sqrt {C _ {1}} + c C _ {2} C _ {3} d ^ {r _ {2} - r _ {1}} - C \left(1 + C _ {2} ^ {\prime} C _ {3} ^ {\prime} d ^ {r _ {2} - 1}\right) \sqrt {d / n ^ {\prime}}}{\sqrt {\left(1 + C _ {2} ^ {\prime} C _ {1}\right) C _ {2} ^ {\prime} C _ {3} ^ {\prime 2}} d ^ {\frac {1 + r _ {2} - 2 r _ {1}}{2}}} \gtrsim d ^ {(r _ {2} - 1) / 2}, \tag {D.12}
$$
which finishes the proof of the lemma.
E Proofs of Section 4
We recall from (D.5) that
$$
\frac {\mathrm {d} \overline {{\boldsymbol {w}}} ^ {t}}{\mathrm {d} t} = \frac {\left(\mathbf {I} _ {d} - \overline {{\boldsymbol {w}}} ^ {t} \overline {{\boldsymbol {w}}} ^ {t} ^ {\top}\right) \hat {\boldsymbol {\Sigma}} ^ {1 / 2}}{\| \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {w} \|} \frac {\mathrm {d} \boldsymbol {w} ^ {t}}{\mathrm {d} t}.
$$
Furthermore, the preconditioned dynamics of  $\boldsymbol{w}^t$  given by (4.1) reads
$$
\frac {\mathrm {d} \pmb {w} ^ {t}}{\mathrm {d} t} = \frac {\eta (\pmb {w} ^ {t}) \hat {\pmb {\Sigma}} ^ {- 1 / 2} (\mathbf {I} _ {d} - \overline {{\pmb {w}}} ^ {t} \overline {{\pmb {w}}} ^ {t} ^ {\top})}{\| \hat {\pmb {\Sigma}} ^ {1 / 2} \pmb {w} \|} \Bigg \{\frac {1}{n} \sum_ {i = 1} ^ {n} \phi^ {\prime} \Big (\big <   \overline {{\pmb {w}}} ^ {t}, \tilde {\pmb {z}} ^ {(i)} \big > \Big) y ^ {(i)} \tilde {\pmb {z}} ^ {(i)} \Bigg \},
$$
where we recall  $\tilde{\pmb{z}}^{(i)}\coloneqq \hat{\pmb{\Sigma}}^{-1 / 2}\pmb{x}^{(i)}$ . Plugging in  $\eta (\pmb {w}^t) = \| \hat{\pmb{\Sigma}}^{1 / 2}\pmb {w}\| ^2$  yields
$$
\begin{array}{l} \frac {\mathrm {d} \overline {{\boldsymbol {w}}} ^ {t}}{\mathrm {d} t} = (\mathbf {I} _ {d} - \overline {{\boldsymbol {w}}} ^ {t} \overline {{\boldsymbol {w}}} ^ {t} ^ {\top}) ^ {2} \left\{\frac {1}{n} \sum_ {i = 1} ^ {n} \phi^ {\prime} \left(\left\langle \overline {{\boldsymbol {w}}} ^ {t}, \tilde {\boldsymbol {z}} ^ {(i)} \right\rangle\right) y ^ {(i)} \tilde {\boldsymbol {z}} ^ {(i)} \right\} \\ = \left(\mathbf {I} _ {d} - \overline {{\boldsymbol {w}}} ^ {t} \overline {{\boldsymbol {w}}} ^ {t} ^ {\top}\right) \left\{\frac {1}{n} \sum_ {i = 1} ^ {n} \phi^ {\prime} \left(\left\langle \overline {{\boldsymbol {w}}} ^ {t}, \tilde {\boldsymbol {z}} ^ {(i)} \right\rangle\right) y ^ {(i)} \tilde {\boldsymbol {z}} ^ {(i)} \right\} \\ \end{array}
$$
The rest of the analysis is identical to that of the proof of Theorem 5 in Appendix D.3. Specifically, by defining
$$
\overline {{\boldsymbol {u}}} _ {\perp} ^ {t} := \overline {{\boldsymbol {u}}} - \left\langle \overline {{\boldsymbol {w}}} ^ {t}, \overline {{\boldsymbol {u}}} \right\rangle \overline {{\boldsymbol {w}}} ^ {t},
$$
$$
\begin{array}{l} \frac {\mathrm {d} \left\langle \overline {{\boldsymbol {w}}} ^ {t} , \overline {{\boldsymbol {u}}} \right\rangle}{\mathrm {d} t} = \left\langle \overline {{\boldsymbol {u}}} _ {\perp} ^ {t}, \mathbb {E} _ {\boldsymbol {z}, y} \left[ \phi^ {\prime} (\left\langle \overline {{\boldsymbol {w}}} ^ {t}, \boldsymbol {z} \right\rangle) y \boldsymbol {z} \right] \right\rangle \\ + \underbrace {\left\langle \overline {{\boldsymbol {u}}} _ {\perp} ^ {t} , \frac {1}{n} \sum_ {i = 1} ^ {n} \phi^ {\prime} \left(\left\langle \overline {{\boldsymbol {w}}} ^ {t} , \boldsymbol {z} ^ {(i)} \right\rangle\right) y ^ {(i)} \boldsymbol {z} ^ {(i)} - \mathbb {E} _ {\boldsymbol {z} , y} \left[ \phi^ {\prime} \left(\left\langle \overline {{\boldsymbol {w}}} ^ {t} , \boldsymbol {z} \right\rangle\right) y \boldsymbol {z} \right] \right\rangle} _ {=: \mathcal {E} _ {1}} \\ + \underbrace {\frac {1}{n} \sum_ {i = 1} ^ {n} \left\{\phi^ {\prime} \left(\left\langle \overline {{\boldsymbol {w}}} ^ {t} , \tilde {\boldsymbol {z}} ^ {(i)} \right\rangle\right) - \phi^ {\prime} \left(\left\langle \overline {{\boldsymbol {w}}} ^ {t} , \boldsymbol {z} ^ {(i)} \right\rangle\right) \right\} y ^ {(i)} \left\langle \boldsymbol {z} ^ {(i)} , \overline {{\boldsymbol {u}}} _ {\perp} ^ {t} \right\rangle} _ {=: \mathcal {E} _ {2}} \\ + \underbrace {\frac {1}{n} \sum_ {i = 1} ^ {n} \phi^ {\prime} \left(\left\langle \overline {{\boldsymbol {w}}}, \tilde {\boldsymbol {z}} ^ {(i)} \right\rangle\right) y ^ {(i)} \left\langle \tilde {\boldsymbol {z}} ^ {(i)} - \boldsymbol {z} ^ {(i)} , \overline {{\boldsymbol {u}}} _ {\perp} ^ {t} \right\rangle} _ {=: \mathcal {E} _ {3}}. \\ \end{array}
$$
As long as  $n \gtrsim d$ ,  $n' = n$  for the smooth case, and  $n' \gtrsim n^2$  for the ReLU case, the first two steps of the proof of Theorem 5 in Appendix D.3 imply that
$$
\mathcal {E} _ {1} \wedge \mathcal {E} _ {2} \wedge \mathcal {E} _ {3} \geq - \| \overrightarrow {\boldsymbol {u}} _ {\perp} ^ {t} \| \tilde {\mathcal {O}} (\sqrt {d / n}) \geq - \tilde {\mathcal {O}} (\sqrt {d / n}).
$$
Once again, we apply Lemma 11 to obtain
$$
\mathbb {E} \left[ \phi^ {\prime} (\langle \overline {{\boldsymbol {w}}}, \boldsymbol {z} \rangle) y z \right] = - \zeta_ {\phi , g} (\langle \overline {{\boldsymbol {w}}}, \overline {{\boldsymbol {u}}} \rangle) \overline {{\boldsymbol {u}}} - \psi_ {\phi , g} (\langle \overline {{\boldsymbol {w}}} ^ {t}, \overline {{\boldsymbol {u}}} \rangle) \overline {{\boldsymbol {w}}},
$$
with  $\zeta_{\phi,g}$  and  $\psi_{\phi,g}$  given in (D.2) and (D.3) respectively. As a result,