SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathcal {R} (\boldsymbol {w}) := \frac {1}{2} \mathbb {E} \left[ \left(\hat {y} (\boldsymbol {x}; \overline {{\boldsymbol {w}}}) - y\right) ^ {2} \right] = \frac {1}{2} \| \phi \| _ {\gamma} ^ {2} + \frac {1}{2} \mathbb {E} \left[ y ^ {2} \right] - \mathbb {E} [ \phi (\langle \overline {{\boldsymbol {w}}}, \boldsymbol {z} \rangle) y ]. \tag {3.1}
$$
In (3.1), the only term that depends on the weights  $\pmb{w}$  is the correlation term and the source of the repulsive force in (2.7) is eliminated; we have  $\nabla_{\pmb{w}}\mathcal{R}(\pmb{w}) = -\nabla_{\pmb{w}}\mathbb{E}[\phi (\langle \overline{\pmb{w}},\pmb {z}\rangle)y]$ . Based on this, we use the following normalized gradient flow for training
$$
\frac {\mathrm {d} \boldsymbol {w} ^ {t}}{\mathrm {d} t} = - \eta (\boldsymbol {w} ^ {t}) \nabla_ {\boldsymbol {w}} \mathcal {R} (\boldsymbol {w} ^ {t}) \text {w h e r e} \eta (\boldsymbol {w}) = \| \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {w} \| ^ {2}. \tag {3.2}
$$
We remark that, though not identical, this normalization is closely related to batch normalization which is commonly employed in practice [IS15]. Under the invariance provided by the current normalization, minimizing  $\mathcal{R}(\boldsymbol{w})$  corresponds to maximizing  $\mathbb{E}[\phi (\langle \overline{\boldsymbol{w}},\boldsymbol {z}\rangle)y]$ . Thus, instead of  $\boldsymbol{w}$ , it will be more useful to track the dynamics of its normalized counterpart  $\overline{\boldsymbol{w}}$ , which is made possible by the following intermediary result that follows from Stein's lemma; also see e.g. [EDB16, MHPG  $^{+}23$ ].
Lemma 3. Suppose we train  $\boldsymbol{w}^t$  using the gradient flow (3.2). Then  $\overline{\boldsymbol{w}}^t$  solves the following ODE
$$
\frac {\mathrm {d} \overline {{\boldsymbol {w}}} ^ {t}}{\mathrm {d} t} = - \zeta_ {\phi , g} (\left\langle \overline {{\boldsymbol {w}}} ^ {t}, \overline {{\boldsymbol {u}}} \right\rangle) \left(\mathbf {I} _ {d} - \overline {{\boldsymbol {w}}} ^ {t} \overline {{\boldsymbol {w}}} ^ {t} ^ {\top}\right) \boldsymbol {\Sigma} \left(\mathbf {I} _ {d} - \overline {{\boldsymbol {w}}} ^ {t} \overline {{\boldsymbol {w}}} ^ {t} ^ {\top}\right) \overline {{\boldsymbol {u}}}, \tag {3.3}
$$
where  $\zeta_{\phi ,g}(\langle \overline{\mathbf{w}},\overline{\mathbf{u}} \rangle)\coloneqq -\mathbb{E}[\phi^{\prime}(\langle \overline{\mathbf{w}},\mathbf{z}\rangle)g^{\prime}(\langle \overline{\mathbf{u}},\mathbf{z}\rangle)]$
We will investigate if the modified flow (3.3) achieves alignment; in this context, alignment corresponds to  $\langle \overline{\boldsymbol{w}}^t,\overline{\boldsymbol{u}}\rangle \approx 1$ . Towards that end, we make the following assumption.
Assumption 2. Let  $g = \sum_{j\geq 0}\alpha_jh_j$  and  $\phi = \sum_{j\geq 0}\beta_jh_j$  be the Hermite decomposition of  $g$  and  $\phi$  respectively. Let  $s$  be the information exponent of  $g$ . For some universal constant  $c > 0$ , we assume
$$
\zeta_ {\phi , g} (\omega) = - \sum_ {j > 0} j \alpha_ {j} \beta_ {j} \omega^ {j - 1} \leq - c \omega^ {s - 1}, \quad \forall \omega \in (0, 1).
$$
There are several important examples that readily satisfy Assumption 2. The obvious example is when the link function is known as in [BAGJ21], i.e.  $\phi = g$ . A more interesting example is when  $\phi$  is an activation with degree  $s$  non-zero Hermite coefficient (e.g. ReLU when  $s$  is even, see [GKK19, Claim 1]) and  $g$  is a degree  $s$  Hermite polynomial, which for  $s = 2$  corresponds to the phase retrieval problem. In this case, the assumption is satisfied if  $\alpha_{s}$  and  $\beta_{s}$  have the same sign, which occurs with probability 0.5 if we randomly choose the sign of the second layer.
Under this condition, the following result shows that the population flow (3.3) can achieve alignment.
Proposition 4. Suppose Assumption 2 holds and consider the gradient flow given by (3.3) with initialization satisfying  $\langle \overline{\boldsymbol{w}}^0,\overline{\boldsymbol{u}}\rangle >0$ . Then, we have  $\langle \overline{\boldsymbol{w}}^T,\overline{\boldsymbol{u}}\rangle \geq 1 - \varepsilon$  as soon as
$$
T \asymp \frac {\tau_ {s} \left(\left\langle \overline {{\boldsymbol {w}}} ^ {0} , \overline {{\boldsymbol {u}}} \right\rangle\right) + \ln (1 / \varepsilon)}{\lambda_ {\min } (\boldsymbol {\Sigma})} \text {w h e r e} \tau_ {s} (z) := \left\{ \begin{array}{l l} 1 & s = 1 \\ \ln (1 / z) & s = 2. \\ (1 / z) ^ {s - 2} & s > 2 \end{array} \right. \tag {3.4}
$$
We remark that the information exponent enters the rate in (3.4) through the function  $\tau_{s}$ , and time needed to achieve  $\varepsilon$  alignment gets worse with larger information exponent. Indeed, it is understood that this quantity serves as a measure of complexity for the target function being learned.
3.2 Empirical dynamics and sample complexity
Given  $n$  i.i.d. samples  $\{(\pmb{x}^{(i)},y^{(i)})\}_{i = 1}^{n}$  from the single index model (2.1), we consider the flow
$$
\frac {\mathrm {d} \boldsymbol {w} ^ {t}}{\mathrm {d} t} = - \eta (\boldsymbol {w} ^ {t}) \nabla \hat {\mathcal {R}} (\boldsymbol {w} ^ {t}) \text {w i t h} \nabla \hat {\mathcal {R}} (\boldsymbol {w}) := - \nabla_ {\boldsymbol {w}} \left\{\frac {1}{n} \sum_ {i = 1} ^ {n} \phi \left(\frac {\left\langle \boldsymbol {w} , \boldsymbol {x} ^ {(i)} \right\rangle}{\| \hat {\boldsymbol {\Sigma}} ^ {1 / 2} \boldsymbol {w} \|}\right) y ^ {(i)} \right\}, \tag {3.5}
$$
where we estimate the covariance matrix  $\pmb{\Sigma}$  using the sample mean  $\hat{\pmb{\Sigma}}\coloneqq \frac{1}{n'}\sum_{i = 1}^{n'}\pmb{x}^{(i)}\pmb{x}^{(i)}^\top$  over  $n^{\prime}$  i.i.d. samples; the above dynamics defines an empirical gradient flow. Notice that we ignored the gradient associated with the term  $\phi^2$  since the population dynamics ensures that its gradient will concentrate around zero; thus, it is redundant to estimate this term. Below, we will use  $n^{\prime} = n$  for smooth activations, i.e. the same dataset can be used for covariance estimation; For ReLU, we require a more accurate covariance estimator, thus, we use  $n^{\prime}\gtrsim n^{2}$  by assuming access to an additional  $n^{\prime} - n$  unlabeled data points. Similar to the previous section, we track the dynamics of normalized  $\pmb{w}$  by defining  $\overline{\boldsymbol{w}}\coloneqq \hat{\pmb{\Sigma}}^{1 / 2}\pmb {w} / \| \hat{\pmb{\Sigma}}^{1 / 2}\pmb {w}\|$  (and leave  $\overline{u}$  unchanged from Section 3.1). The same arguments as in Lemma 3 allow us to track the evolution of  $\overline{w}$ , which ultimately yields the following alignment result under general covariance structure.
Theorem 5. Let  $s$  be the information exponent of  $g$ , and assume it satisfies  $|g(\cdot)| \lesssim 1 + |\cdot|^p$  for some  $p > 0$ . For  $\phi$  denoting either the ReLU activation or a smooth activation satisfying  $|\phi'| \vee |\phi''| \lesssim 1$ , suppose Assumption 2 holds. For any  $\varepsilon > 0$ , suppose we run the finite sample gradient flow (3.5) with  $\eta(\boldsymbol{w}) = \| \hat{\Sigma}^{1/2}\boldsymbol{w} \|^2$ , initialized such that  $\langle \overline{\boldsymbol{w}}^0, \overline{\boldsymbol{u}} \rangle > 0$ , and with number of samples
$$
n \gtrsim d \varkappa (\boldsymbol {\Sigma}) ^ {2} \left\{\left\langle \overline {{\boldsymbol {w}}} ^ {0}, \overline {{\boldsymbol {u}}} \right\rangle^ {2 (1 - s)} \vee \varepsilon^ {- 2} \right\},
$$
where  $\varkappa(\Sigma)$  is the condition number of  $\Sigma$ . Then, for  $T \asymp \frac{\tau_s\left(\left\langle\overline{w}^0,\overline{u}\right\rangle + \ln(1 / \varepsilon)}{\lambda_{\min}(\Sigma)}$ , we have
$$
\left\langle \overline {{\boldsymbol {w}}} ^ {T}, \overline {{\boldsymbol {u}}} \right\rangle \geq 1 - \varepsilon , \tag {3.6}
$$
with probability at least  $1 - c_{1}d^{-c_{2}}$  for some universal constants  $c_{1},c_{2} > 0$  over the randomness of the dataset. Here,  $\tau_{s}$  is defined in (3.4) and  $\gtrsim$  hides poly-logarithmic factors.
Remark. We make the following remarks on the above theorem.
- The initial condition  $\langle \overline{\boldsymbol{w}}^0, \overline{\boldsymbol{u}} \rangle > 0$  is required when we have odd information exponent. When  $\boldsymbol{w}^0$  is initialized uniformly over  $\mathbb{S}^{d - 1}$ , the condition holds with probability 0.5 over the initialization. See [BAGJ21, Remark 1.8] for further discussion on this condition.
- Although  $\overline{w}$  is defined using the empirical covariance unlike  $\overline{u}$  which is defined by population covariance, this definition is the suitable choice to approximate the target function  $g$  (c.f. Theorem 9), since it ensures the arguments of  $\phi$  and  $g$  are sufficiently close when  $\overline{w}$  recovers  $\overline{u}$ .
The intuition behind the proof of Theorem 5 is presented in Section 5 with the complete proof in the appendix. We highlight that the improvement in the sample complexity compared to the isotropic setting occurs whenever the covariance structure induces a stronger initial alignment and consequently stronger signal. The following corollary demonstrates a concrete example of such improvement by specializing Theorem 5 for a spiked covariance model.
Corollary 6. Consider the setting of Theorem 5 with  $\pmb{\Sigma}$  following the  $(\kappa, \theta)$ -spiked model, where  $\langle \pmb{u}, \pmb{\theta} \rangle \asymp d^{-r_1}$  and  $\kappa \asymp d^{r_2}$  with  $r_1 \in [0, 1/2]$  and  $r_2 \in [0, 1]$ . Suppose  $\pmb{w}^0$  is sampled uniformly from  $\mathbb{S}^{d-1}$ . Then, when conditioned on  $\langle \overline{\pmb{w}}^0, \overline{\pmb{u}} \rangle > 0$ , the sample complexity in Theorem 5 reads
$$
n \gtrsim \left\{ \begin{array}{l l} d ^ {1 + 2 r _ {2}} \left(d ^ {s - 1} \vee \varepsilon^ {- 2}\right) & 0 <   r _ {2} <   r _ {1} \\ d ^ {1 + 2 r _ {2}} \left(d ^ {(s - 1) (1 - 2 (r _ {2} - r _ {1}))} \vee \varepsilon^ {- 2}\right) & r _ {1} <   r _ {2} <   2 r _ {1}, \\ d ^ {1 + 2 r _ {2}} \left(d ^ {(s - 1) (1 - r _ {2})} \vee \varepsilon^ {- 2}\right) & 2 r _ {1} <   r _ {2} <   1 \end{array} \right. \tag {3.7}
$$
where  $\gtrsim$  hides poly-logarithmic factors of  $d$ .
Remark. We have the following observations on the above sample complexity.
- Corollary 6 demonstrates that structured data can lead to better sample complexity when the right normalization is used during training. This complements Theorem 2 where we recall that spherical training dynamics ignores the structure in data and the target direction cannot be recovered.
- When  $g$  is a polynomial of degree  $p$ , the lower bound for rotationally invariant kernels (including the neural tangent kernel at initialization) implies a complexity of at least  $d^{\Omega((1 - r_2)p)}$  [DWY21]. Thus the sample complexity of Corollary 6 can always outperform the kernel lower bound when  $p$  is sufficiently large and  $s$  remains constant.
Three-step phase transition. Recall that in the isotropic setting  $\Sigma = \mathbf{I}_d$ , the sample complexity of learning  $g$  with information exponent  $s$  using full-batch gradient flow is  $\tilde{\mathcal{O}}(d^s)$  for  $s \geq 3$  [BBSS22]. The sample complexity in Corollary 6 is strictly smaller than  $\tilde{\mathcal{O}}(d^s)$  as soon as  $(s - 1)r_1 / (s - 2) < r_2$ . Furthermore, for any  $\nu > 0$  it is at most  $\tilde{\mathcal{O}}(d^{3 + \nu})$  as soon as  $r_2 \geq 1 - \nu / (s - 3)$  and  $2r_1 < r_2$ , in which case the sample complexity becomes independent of the information exponent. Interestingly, the complexity becomes independent of  $r_1$  when  $r_2 > 2r_1$  or  $r_2 < r_1$ , i.e. the direction of the spike becomes irrelevant when the spike magnitude is sufficiently large or small.
The three-stage phase transition of Corollary 6 is due to the different behaviour of the inner product  $\langle \overline{\boldsymbol{w}}^0,\overline{\boldsymbol{u}}\rangle$  in different regimes of  $r_1$  and  $r_2$ . When  $r_2 < r_1$ , we have  $\langle \overline{\boldsymbol{w}}^0,\overline{\boldsymbol{u}}\rangle \asymp \langle \boldsymbol{w}^0,\boldsymbol {u}\rangle$ , thus the initial alignment is just as uninformative as the isotropic case providing no improvement. Moreover, a potentially large condition number may hurt the sample complexity in this case. On the other hand, when  $r_1 < r_2 < 2r_1$  we have  $\langle \overline{\boldsymbol{w}}^0,\overline{\boldsymbol{u}}\rangle \asymp \kappa \langle \boldsymbol {u},\boldsymbol {\theta}\rangle \langle \boldsymbol{w}^0,\boldsymbol {\theta}\rangle$ , and  $r_2 > 2r_1$  leads to  $\langle \overline{\boldsymbol{w}}^0,\overline{\boldsymbol{u}}\rangle \asymp \sqrt{\kappa}\langle \boldsymbol{w}^0,\boldsymbol {\theta}\rangle$ , thus large  $\kappa$  or  $\langle \boldsymbol {u},\boldsymbol {\theta}\rangle$  in this regime may improve the sample complexity.
4 Implications to Neural Networks and Further Improvements
4.1 Improving Sample Complexity via Preconditioning
We now demonstrate that preconditioning the training dynamics with  $\hat{\Sigma}^{-1}$  can remove the dependency on  $\varkappa(\Sigma)$ , ultimately improving the sample complexity. Consider the preconditioned gradient flow
$$
\frac {\mathrm {d} \boldsymbol {w} ^ {t}}{\mathrm {d} t} = - \eta (\boldsymbol {w} ^ {t}) \hat {\boldsymbol {\Sigma}} ^ {- 1} \nabla \hat {\mathcal {R}} (\boldsymbol {w} ^ {t}) \quad \text {w i t h} \quad \eta (\boldsymbol {w}) = \| \hat {\boldsymbol {\Sigma}} ^ {1 / 2} \boldsymbol {w} \| ^ {2}. \tag {4.1}
$$
We have the following alignment result.
Theorem 7. Consider the same setting as Theorem 5, and assume we run the preconditioned empirical gradient flow (4.1) with number of samples
$$
n \gtrsim d \left\{\left\langle \overline {{\boldsymbol {w}}} ^ {0}, \overline {{\boldsymbol {u}}} \right\rangle^ {2 (1 - s)} \vee \varepsilon^ {- 2} \right\},
$$
where  $\gtrsim$  hides poly-logarithmic factors of  $d$ . Then, for  $T \asymp \tau_s\left(\langle \overline{\boldsymbol{w}}^0, \overline{\boldsymbol{u}} \rangle\right) + \ln(1 / \varepsilon)$ , we have
$$
\left\langle \overline {{\boldsymbol {w}}} ^ {T}, \overline {{\boldsymbol {u}}} \right\rangle \geq 1 - \varepsilon ,
$$
with probability at least  $1 - c_{1}d^{-c_{2}}$  for some universal constants  $c_{1},c_{2} > 0$ .
Preconditioning removes the condition number dependence, which is particularly important in the spiked model case where this quantity can be large.
Corollary 8. Consider the setting of Theorem 7, and assume we run the preconditioned empirical gradient flow (4.1) for the  $(\kappa, \pmb{\theta})$ -spiked model where  $\langle \pmb{u}, \pmb{\theta} \rangle \asymp d^{-r_1}$  and  $\kappa \asymp d^{r_2}$  with  $r_1 \in [0,1/2]$  and  $r_2 \in [0,1]$ . Suppose  $\pmb{w}^0$  is sampled uniformly from  $\mathbb{S}^{d-1}$ . Then, when conditioned on  $\langle \overline{\pmb{w}}^0, \overline{\pmb{u}} \rangle > 0$ , the sample complexity of Theorem 7 reads
$$
n \gtrsim \left\{ \begin{array}{l l} d \left(d ^ {s - 1} \vee \varepsilon^ {- 2}\right) & 0 <   r _ {2} <   r _ {1} \\ d \left(d ^ {(s - 1) (1 - 2 \left(r _ {2} - r _ {1}\right))} \vee \varepsilon^ {- 2}\right) & r _ {1} <   r _ {2} <   2 r _ {1}, \\ d \left(d ^ {(s - 1) (1 - r _ {2})} \vee \varepsilon^ {- 2}\right) & 2 r _ {1} <   r _ {2} <   1 \end{array} \right. \tag {4.2}
$$
The above result improves upon Corollary 6; thus, making a case for preconditioning in practice. The complexity results also strictly improve upon the  $\tilde{\mathcal{O}}(d^s)$  complexity in the isotropic case [BBSS22] when  $r_2 > r_1$ . Further, for any  $\nu > 0$ , we can obtain the complexity of  $\tilde{\mathcal{O}}(d^{1 + \nu})$  (nearly linear in dimension) when  $r_2 > 1 - \nu/(s-1)$  and  $r_2 > 2r_1$  or  $r_1 + 1/2(1 - \nu/(s-1)) < r_2 < 2r_1$ . In addition to the remarks of Corollary 6, we note that the complexity is independent of both  $r_1$  and  $r_2$  when  $r_2 < r_1$  (cf. Figure 1 hard regime), i.e. the spike magnitude and the spike-target alignment have no effect on the complexity unless  $r_2 \geq r_1$ .
Under the spiked covariance model, one could improve the above results by instead using spectral initialization, i.e. initializing at  $\pmb{\theta}$ , which can be estimated from unlabeled data. Assuming perfect access to  $\pmb{\theta}$ , using the statement of Theorems 5 and 7, this initialization would imply a sample complexity of  $\tilde{\mathcal{O}}(d^{1 + 2r_2 + ((s - 1)(2r_1 - r_2) \vee 0)})$  without and  $\tilde{\mathcal{O}}(d^{1 + ((s - 1)(2r_1 - r_2) \vee 0)})$  with preconditioning.
4.2 Two-layer neural networks and learning the link function
Our main focus so far was learning the target direction  $\mathbf{u}$ . Next, we consider learning the unknown link function with a neural network, providing a complete learnability result for single index models.
We use Algorithm 1 and train the first-layer of the neural network with either the empirical gradient flow (3.5) or the preconditioned version (4.1). Then, we randomly choose the bias units and minimize the second layer weights using another gradient flow. Our goal is to track the sample complexity  $n$  needed to learn the single index target which we compare against the results of [BBSS22]. We highlight that layer-wise training in Algorithm 1 is frequently employed in the literature [BES+22, BBSS22, DLS22, MHPG+23] and in particular [BBSS22] also used gradient flow for training.
Algorithm 1 Layer-wise training of a two-layer ReLU network with gradient flow (GF).
Input:  $\pmb{w}^0\in \mathbb{R}^d,T,T',\Delta ,\lambda \in \mathbb{R}_+$  and data  $\{(\pmb{x}^{(i)},y^{(i)})\}_{i = 1}^{n}$ .
1: Train the first layer weights  $\mathbf{W}_j^T$  using the GF (3.5) or the preconditioned GF (4.1).
2:Normalize the weights  $\pmb{W}_{j}^{T}\coloneqq \pmb{W}_{j}^{T} / \| \hat{\pmb{\Sigma}}^{1 / 2}\pmb{W}_{j}^{T}\|$  for every  $1\leq j\leq m$
3: Let  $b_{j} \stackrel{i.i.d.}{\sim} \mathrm{Unif}(-\Delta, \Delta)$  and  $a_{j}^{0} = 1 / m$  for  $1 \leq j \leq m$ .
4: Train the second layer weights  $\pmb{a}^{T'}$  via the gradient flow
$$
\frac {\mathrm {d} \boldsymbol {a} ^ {t}}{\mathrm {d} t} = - \nabla_ {\boldsymbol {a}} \left\{\frac {1}{2 n} \sum_ {i = 1} ^ {n} (\hat {y} (\boldsymbol {x} ^ {(i)}; \boldsymbol {W} ^ {T}, \boldsymbol {a} ^ {t}, \boldsymbol {b}) - y ^ {(i)}) ^ {2} + \frac {\lambda \| \boldsymbol {a} ^ {t} \| ^ {2}}{2} \right\}.
$$
5: return  $(\pmb{W}^T, \pmb{a}^{T'}, \pmb{b})$ .
Theorem 9. Let  $g$  be twice weakly differentiable with information exponent  $s$  and assume  $g''$  has at most polynomial growth. Suppose  $\phi$  is the ReLU activation, Assumption 2 holds and we run Algorithm 1 with  $w^0$  initialized uniformly over  $\mathbb{S}^{d-1}$ . For any  $\varepsilon > 0$ , let  $n$  and  $T$  be chosen according to Theorem 5 when we run the gradient flow (3.5) and Theorem 7 when we run the preconditioned gradient flow (4.1). Then, for  $\Delta \asymp \sqrt{\ln(nd)}$ , some regime of  $\lambda$  given by (E.3) and sufficiently large  $T'$  given by (E.4), we have
$$
\mathbb {E} _ {(\boldsymbol {x}, y)} \left[ \left(\hat {y} (\boldsymbol {x}; \boldsymbol {W} ^ {T}, \boldsymbol {a} ^ {T ^ {\prime}}, \boldsymbol {b}) - y\right) ^ {2} \right] \leq C _ {1} \mathbb {E} \left[ \epsilon^ {2} \right] + C _ {2} (\varepsilon + 1 / m), \tag {4.3}
$$
conditioned on  $\langle \overline{\mathbf{w}}^0, \overline{\mathbf{u}} \rangle > 0$  with probability at least 0.99 over the randomness of the dataset, biases, and initialization, where  $C_1$  is a universal constant and  $C_2$  hides polylog $(m, n, d)$  factors.
The next result immediately follows from the previous theorem together with Corollaries 6 & 8.
Corollary 10. In the setting of Theorem 9, if  $\pmb{\Sigma}$  follows the  $(\kappa, \theta)$ -spiked model, the sample complexity  $n$  is given by (3.7) if we use the empirical gradient flow and (4.2) if we use the preconditioned version.
We remark that for fixed  $\varepsilon$ , the sample complexity to learn  $g$  in the isotropic case is  $\tilde{\mathcal{O}}(d^s)$  [BBSS22]. Under the spiked model, if we assume that  $r_2$  is sufficiently large and  $r_1$  is sufficiently small as discussed in the previous section, Corollary 10 improves this rate to either (3.7) when the empirical gradient flow is used without preconditioning or to (4.2) with preconditioning.
5 Technical Overview
In this section, we briefly discuss the key intuitions that lead to the proof of our main results. We first review the case  $\pmb{\Sigma} = \mathbf{I}_d$ , where we have the following decomposition for population loss
$$
R (\boldsymbol {w}) := \frac {1}{2} \mathbb {E} \left[ \left(\phi \left(\langle \boldsymbol {w}, \boldsymbol {x} \rangle\right) - y\right) ^ {2} \right] = \frac {1}{2} \| \phi \| _ {\gamma} ^ {2} + \frac {1}{2} \mathbb {E} \left[ y ^ {2} \right] - \mathbb {E} \left[ \phi \left(\langle \boldsymbol {w}, \boldsymbol {x} \rangle\right) g \left(\langle \boldsymbol {u}, \boldsymbol {x} \rangle\right) \right]. \tag {5.1}
$$
Notice that the only term contributing to the population gradient is the last term which measures the correlation between  $\phi$  and  $g$ . Following the gradient flow and applying Stein's lemma yields
$$
\frac {\mathrm {d} \langle \boldsymbol {w} ^ {t} , \boldsymbol {u} \rangle}{\mathrm {d} t} = \mathbb {E} \left[ \phi^ {\prime} (\langle \boldsymbol {w} ^ {t}, \boldsymbol {x} \rangle) g ^ {\prime} (\langle \boldsymbol {u}, \boldsymbol {x} \rangle) \right] (1 - \langle \boldsymbol {w} ^ {t}, \boldsymbol {u} \rangle^ {2}) = (1 - \langle \boldsymbol {w} ^ {t}, \boldsymbol {u} \rangle^ {2}) \sum_ {j \geq s} j \alpha_ {j} \beta_ {j} \langle \boldsymbol {w} ^ {t}, \boldsymbol {u} \rangle^ {j - 1},
$$
where the second identity follows from the Hermite expansion; see also [EDB16, EBD19]. Assume  $\alpha_{s}\beta_{s} > 0$  to ensure that the population dynamics will move towards  $\pmb{u}$  at least near initialization. When replacing the population gradient with a full-batch gradient, we need the estimation noise to be smaller than the signal existing in the gradient. When  $\langle \pmb{w}^0,\pmb {u}\rangle \ll 1$ , this signal is roughly of the order  $\langle \pmb {w}^0,\pmb {u}\rangle^{s - 1}$ . As the uniform concentration error over  $\mathbb{S}^{d - 1}$  scales with  $\sqrt{d / n}$ , we need  $n\asymp d\langle \pmb{w}^0,\pmb {u}\rangle^{2(s - 1)}$  to ensure the signal remains dominant and  $\pmb{w}^t$  moves towards  $\pmb{u}$ . When  $\pmb{w}^0$  is initialized uniformly over  $\mathbb{S}^{d - 1}$  this translates to a sample complexity of  $n\asymp d^s$ , which is indeed obtained by [BBSS22] via similar arguments.
However, the behavior of the spherical dynamics entirely changes when we move to the anisotropic case. Suppose  $\pmb{\Sigma}$  follows a  $(\kappa ,\pmb {u})$  -spiked model and  $\phi$  is ReLU. Using Lemma 12, it is easy to show that with the spherical gradient flow, the alignment obeys the following ODE
$$
\frac {\mathrm {d} \langle \boldsymbol {w} ^ {t} , \boldsymbol {u} \rangle}{\mathrm {d} t} = \left\{\mathbb {E} \big [ \phi^ {\prime} (\langle \overline {{\boldsymbol {w}}} ^ {t}, \boldsymbol {z} \rangle) g ^ {\prime} (\langle \overline {{\boldsymbol {u}}}, \boldsymbol {z} \rangle) \big ] - \frac {\kappa \tilde {\psi} _ {\phi , g} (\boldsymbol {w} ^ {t})}{1 + \kappa} \big <   \boldsymbol {w} ^ {t}, \boldsymbol {u} \big > \right\} (1 - \langle \boldsymbol {w} ^ {t}, \boldsymbol {u} \rangle^ {2}),
$$
where  $\tilde{\psi}_{\phi,g}(\boldsymbol{w}^t)$  is introduced in Lemma 12. The additional  $\tilde{\psi}_{\phi,g}(\boldsymbol{w}^t)$  term creates a repulsive force towards the equator  $\langle \boldsymbol{w}^t, \boldsymbol{u} \rangle = 0$ . The presence of this term is due to the fact that unlike (5.1), the term  $\mathbb{E}\big[\phi (\langle \boldsymbol{w}^t, \boldsymbol{u} \rangle)^2\big]$  is no longer independent of  $\boldsymbol{w}$  and cannot be replaced by  $\| \phi \|_{\gamma}^{2}$ . When  $\boldsymbol{w}^0$  is initialized uniformly over  $\mathbb{S}^{d-1}$  and  $\Omega(1) \leq \kappa \leq \mathcal{O}(d)$ , we have  $\langle \overline{\boldsymbol{w}}^0, \overline{\boldsymbol{u}} \rangle \asymp \sqrt{\kappa} \langle \boldsymbol{w}^0, \boldsymbol{u} \rangle$ . Furthermore, at this initialization  $\tilde{\psi}_{\phi,g}(\boldsymbol{w}^0) \approx 1/2$ . Therefore,