SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\frac {\mathrm {d} \langle \boldsymbol {w} ^ {t} , \boldsymbol {u} \rangle}{\mathrm {d} t} \approx \left\{s \alpha_ {s} \beta_ {s} (\sqrt {\kappa} \langle \boldsymbol {w} ^ {0}, \boldsymbol {u} \rangle) ^ {s - 1} - \frac {\langle \boldsymbol {w} ^ {0} , \boldsymbol {u} \rangle}{2} \right\}.
$$
Hence the dynamics is trapped at  $|\langle \pmb{w}^t, \pmb{u} \rangle| = \mathcal{O}(1 / \sqrt{d})$  for all  $t > 0$  as long as  $\kappa = \mathcal{O}(d^{1 - 1 / (s - 1)})$ .
To remove the repulsive force in the spherical dynamics, we can directly normalize the input of  $\phi$ . As demonstrated by (3.1), once again the only term that varies with  $\boldsymbol{w}$  would be the correlation loss. Specifically, using the result of Lemma 3, in the population limit we can track  $\langle \overline{\boldsymbol{w}}^t,\overline{\boldsymbol{u}}\rangle$  via
$$
\frac {\mathrm {d} \left\langle \overline {{\boldsymbol {w}}} ^ {t} , \overline {{\boldsymbol {u}}} \right\rangle}{\mathrm {d} t} = \mathbb {E} \left[ \phi^ {\prime} (\langle \overline {{\boldsymbol {w}}}, \boldsymbol {z} \rangle) g ^ {\prime} (\langle \overline {{\boldsymbol {u}}}, \boldsymbol {z} \rangle) \right] \left\langle \overline {{\boldsymbol {u}}} _ {\perp} ^ {t}, \boldsymbol {\Sigma} \overline {{\boldsymbol {u}}} _ {\perp} ^ {t} \right\rangle , \tag {5.2}
$$
where  $\overline{\boldsymbol{u}}_{\perp}^{t} := \overline{\boldsymbol{u}} - \langle \overline{\boldsymbol{w}}^{t}, \overline{\boldsymbol{u}} \rangle \overline{\boldsymbol{w}}^{t}$ . Thus, the strength of the signal at initialization is of order  $\langle \overline{\boldsymbol{w}}^{0}, \overline{\boldsymbol{u}} \rangle^{s-1} / \varkappa(\boldsymbol{\Sigma})$ , which after controlling the error in the estimate of  $\hat{\boldsymbol{\Sigma}}$  and in the estimate of population gradient using finitely many samples, leads to the sample complexity  $n \asymp d\varkappa(\boldsymbol{\Sigma})^{2} \langle \overline{\boldsymbol{w}}^{0}, \overline{\boldsymbol{u}} \rangle^{2(1-s)}$ . Importantly,  $\boldsymbol{\Sigma}$  can incude a much stronger initial alignment  $\langle \overline{\boldsymbol{w}}^{0}, \overline{\boldsymbol{u}} \rangle$  than the isotropic case  $\langle \boldsymbol{w}^{0}, \boldsymbol{u} \rangle$ , which is emphasized in Corollary 6. Using preconditioning will further remove the dependency on the condition number of  $\boldsymbol{\Sigma}$ .
We studied the dynamics of gradient flow to learn single index models when the input data covariance may contain additional structure. Under a spiked model for the covariance matrix, we showed that using spherical gradient flow, as an example of a covariance-agnostic training mechanism employed in the recent literature, is unable to learn the target direction of the single index model even when the spike and the target directions are identical. In contrast, we showed that an appropriate weight normalization removes this problem and successfully recovers the target direction. Moreover, depending on the alignment between the covariance structure and the target direction, the sample complexity can improve upon the isotropic setting, while also outperforming lower bounds for rotationally-invariant kernels. This phenomenon is due to the additional information about the target direction contained in the covariance matrix which improves the effective alignment at initialization. Additionally, we showed that a simple preconditioning of the gradient flow using the inverse empirical covariance can improve the sample complexity, achieving almost linear rate in certain settings.
We outline a few limitations of our current work and discuss directions for future research.
- While studying single index models provides a pathway to a general understanding of feature learning with structured covariance, considering multi-index models can provide a more complete picture [PSE22], e.g. by establishing incremental learning dynamics [ABAM23]. We leave the problem of learning multi-index models under structured input as an interesting future direction.
- Gradient flow under squared loss can be seen as an example of a Correlational Statistical Query (CSQ) algorithm [BF02, Rey20], i.e. an algorithm that only accesses noisy estimates of expected correlation queries from the model. Understanding the limitations of learning single index models under a structured input through a CSQ lower bound perspective is another important direction that would complement our results in this paper.
- When training the first layer, we considered a somewhat unconventional initialization and relied on the symmetry it induces. It is interesting to consider cases where we train a network with multiple neurons starting from a more standard initialization which can help relax Assumption 2.
The authors thank Alberto Bietti and Zhichao Wang for discussions and feedback on the manuscript. TS was partially supported by JSPS KAKENHI (20H00576) and JST CREST (JPMJCR2015). MAE was partially supported by NSERC Grant [2019-06167], CIFAR AI Chairs program, CIFAR AI Catalyst grant.
[AAM22] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz, The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks, Conference on Learning Theory, 2022.
[ABAM23] Emmanuel Abbe, Enric Boix-Adsera, and Theodor Misiakiewicz, Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics, arXiv preprint arXiv:2302.11055 (2023).
[ASKL23] Luca Arnaboldi, Ludovic Stephan, Florent Krzakala, and Bruno Loureiro, From high-dimensional & mean-field dynamics to dimensionless odes: A unifying approach to sgd in two-layers networks, arXiv preprint arXiv:2302.05882 (2023).
[ATV22] Pranjal Awasthi, Alex Tang, and Aravindan Vijayaraghavan, Agnostic learning of general relu activation using gradient descent, arXiv preprint arXiv:2208.02711 (2022).
[BAGJ21] Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath, Online stochastic gradient descent on non-convex losses from high-dimensional inference., J. Mach. Learn. Res. 22 (2021), 106-1.
[BBAP05] Jinho Baik, Gérard Ben Arous, and Sandrine Péché, Phase transition of the largest eigenvalue for nonnull complex sample covariance matrices.
[BBSS22] Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song, Learning single-index models with shallow neural networks, Advances in Neural Information Processing Systems, 2022.
[  $\mathrm{BEG^{+}22}]$  Boaz Barak, Benjamin L Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang, Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit, arXiv preprint arXiv:2207.08799 (2022).
[BES+19] Jimmy Ba, Murat Erdogdu, Taiji Suzuki, Denny Wu, and Tianzong Zhang, Generalization of two-layer neural networks: An asymptotic viewpoint, International Conference on Learning Representations, 2019.
$\left[\mathrm{BES}^{+}22\right]$  Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang, High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation, arXiv preprint arXiv:2205.01445 (2022).
[BES+23] Jimmy Ba, Murat A. Erdogdu, Taiji Suzuki, Zhichao Wang, and Denny Wu, Learning in the presence of low-dimensional structure: a spiked random matrix perspective, Thirty-seventh Conference on Neural Information Processing Systems (NeurIPS 2023), 2023.
[Nader H Bshouty and Vitaly Feldman, On using extended statistical queries to avoid membership queries, Journal of Machine Learning Research 2 (2002), no. Feb, 359-395.
[BMZ23] Raphaël Berthier, Andrea Montanari, and Kangjie Zhou, Learning time-scales in two-layers neural networks, arXiv preprint arXiv:2303.00055 (2023).
[BPVZ23] Joan Bruna, Loucas Pillaud-Vivien, and Aaron Zweig, On single index models beyond gaussian data, arXiv preprint arXiv:2307.15804 (2023).
[BS06] Jinho Baik and Jack W Silverstein, Eigenvalues of large sample covariance matrices of spiked population models, Journal of multivariate analysis 97 (2006), no. 6, 1382-1408.
[RVE18] Grant M Rotskoff and Eric Vanden-Eijnden, Neural networks as Interacting Particle Systems: Asymptotic convexity of the Loss Landscape and Universal Scaling of the Approximation Error, arXiv preprint arXiv:1805.00915 (2018).
[Sol17] Mahdi Soltanolkotabi, Learning relus via gradient descent, Advances in neural information processing systems 30 (2017).
[SWON23] Taiji Suzuki, Denny Wu, Kazusato Oko, and Atsushi Nitanda, Feature learning via mean-field Langevin dynamics: classifying sparse parities and beyond, Thirty-seventh Conference on Neural Information Processing Systems (NeurIPS 2023), 2023.
[Tel23] Matus Telgarsky, Feature selection and low test error in shallow low-rotation relu networks, The Eleventh International Conference on Learning Representations, 2023.
[Ver18] Roman Vershynin, High-dimensional probability: An introduction with applications in data science, Cambridge University Press, 2018.
[VH16] Ramon Van Handel, Probability in high dimension, 2016.
[Wai19] Martin J. Wainwright, High-dimensional statistics: A non-asymptotic viewpoint, Cambridge University Press, 2019.
[WLLM19] Colin Wei, Jason D Lee, Qiang Liu, and Tengyu Ma, Regularization matters: Generalization and optimization of neural nets vs their induced kernel, Advances in Neural Information Processing Systems 32 (2019).
[YO20] Gilad Yehudai and Shamir Ohad, Learning a single neuron with gradient methods, Proceedings of Thirty Third Conference on Learning Theory, Proceedings of Machine Learning Research, vol. 125, PMLR, 2020, pp. 3756-3786.
[YS19] Gilad Yehudai and Ohad Shamir, On the Power and Limitations of Random Features for Understanding Neural Networks, Advances in Neural Information Processing Systems, 2019.
A Concurrent Works
In this paragraph we briefly summarize a few relevant results concurrent to our submission. [BMZ23] provided a precise analysis of the two-timescale dynamics in learning a single index model with information exponent  $s = 1$ . [MHD $^{+}$ 23] considered the learning of a single index target with  $s = 2$  using a neural network in the mean-field regime. [BPVZ23] extended the information exponent-based characterization of online SGD to input data beyond Gaussian. [DNGL23] showed that a gradient-smoothed dynamics can improve the sample complexity and match the CSQ lower bound. Finally, beyond the single-index setting, [DKL $^{+}$ 23, SWON23, CWPPS23] considered learning low-dimensional target functions supported on  $k > 1$  dimensions via gradient-based feature learning.
B Background on Hermite Expansion
The normalized Hermite polynomials  $\{h_j\}_{j\geq 0}$  given by (2.4) provide an orthonormal basis for  $L^2 (\gamma)$ , thus for every  $f\in L^{2}(\gamma)$  we have
$$
f = \sum_ {j = 0} ^ {\infty} \langle f, h _ {j} \rangle_ {\gamma} h _ {j},
$$
where  $\langle f,h_j\rangle_{\gamma}\coloneqq \mathbb{E}_{z\sim \mathcal{N}(0,1)}[f(z)h_j(z)]$ . We will commonly invoke the following well-known properties of Hermite polynomials. If  $j\geq 1$ , then  $h_j^\prime = \sqrt{j} h_{j - 1}$ , where  $h_j^\prime$  stands for the derivative of  $h_j$ . Furthermore, if  $z_{1}$  and  $z_{2}$  are two standard Gaussian random variables with  $\mathbb{E}[z_1z_2] = \rho$ , then  $\mathbb{E}[h_i(z_1)h_j(z_2)] = \delta_{ij}\rho^j$  where  $\delta_{ij}$  is the Kronecker delta. We refer the interested reader to [O'D14, Chapter 11.2] for additional discussions and properties of these polynomials.
We will now discuss how our Definition 1 relates to the original definition of information exponent of [BAGJ21]. In their setting, they assume the true data distribution  $\mathbb{P}_{\boldsymbol{u}}$  is parameterized by some unit vector  $\boldsymbol{u} \in \mathbb{S}^{d-1}$ , and we know the parametric family  $\{\mathbb{P}_{\boldsymbol{w}}\}_{\boldsymbol{w} \in \mathbb{S}^{d-1}}$ ; thus the problem is to estimate the direction  $\boldsymbol{u}$ . Furthermore, they assume the population loss, which is the expectation of some per-sample loss, has spherical symmetry, i.e. the population loss  $R(\boldsymbol{w})$  can be written as  $R(\boldsymbol{w}) = \tilde{R}(\langle \boldsymbol{w}, \boldsymbol{u} \rangle)$ . Then, [BAGJ21, Definition 1.2] defines the information exponent to be the degree of the first non-zero coefficient of  $\tilde{R}$  in its Taylor expansion around the origin. In other words, we say  $R$  has information exponent  $s$  if
$$
\left\{ \begin{array}{l l} \frac {\mathrm {d} ^ {k} \tilde {R}}{\mathrm {d} z ^ {k}} (0) = 0 & 1 \leq \tilde {k} <   s \\ \frac {\mathrm {d} ^ {k} \tilde {R}}{\mathrm {d} z ^ {k}} (0) = - c <   0 & k = s \\ \left| \frac {\mathrm {d} ^ {k} \tilde {R}}{\mathrm {d} z ^ {k}} (z) \right| \leq C & k > s, \forall z \in [ - 1, 1 ] \end{array} \right.,
$$
where  $C, c > 0$  are universal constants. To specialize the above abstract definition to the Gaussian case, consider the setting where the input data is standard Gaussian  $\boldsymbol{x} \sim \mathcal{N}(0, \mathbf{I}_d)$  and the problem is to estimate  $\boldsymbol{u} \in \mathbb{S}^{d-1}$  given a response variable  $y = f(\langle \boldsymbol{u}, \boldsymbol{x} \rangle)$  with known  $f$ . Via the Hermite expansion of  $f$ , one can write
$$
\tilde {R} (\langle \boldsymbol {w}, \boldsymbol {u} \rangle) := \frac {1}{2} \mathbb {E} \left[ (f (\langle \boldsymbol {w}, \boldsymbol {x} \rangle) - f (\langle \boldsymbol {u}, \boldsymbol {x} \rangle)) ^ {2} \right] = - \sum_ {j \geq 1} \langle f, h _ {j} \rangle_ {\gamma} ^ {2} \langle \boldsymbol {w}, \boldsymbol {u} \rangle^ {j} + \text {c o n s t .}
$$
Thus, the information exponent of  $\tilde{R}$  is indeed the degree of the first non-zero term in the Hermite expansion of  $f$ .
Now consider the general case where  $x \sim \mathcal{N}(0, \Sigma)$ . The spherical symmetry assumed in [BAGJ21] no longer holds. However, after proper normalization of weights, if we consider the population loss
$$
R (\pmb {w}) := \frac {1}{2} \mathbb {E} \left[ \left(f \left(\frac {\langle \pmb {w} , \pmb {x} \rangle}{\| \pmb {\Sigma} ^ {1 / 2} \pmb {w} \|}\right) - f \left(\frac {\langle \pmb {u} , \pmb {x} \rangle}{\| \pmb {\Sigma} ^ {1 / 2} \pmb {u} \|}\right)\right) ^ {2} \right],
$$
then  $R(\boldsymbol{w}) = \tilde{R}\Big(\frac{\langle\boldsymbol{w},\boldsymbol{\Sigma}\boldsymbol{u}\rangle}{\|\boldsymbol{\Sigma}^{1 / 2}\boldsymbol{w}\|\|\boldsymbol{\Sigma}^{1 / 2}\boldsymbol{u}\|}\Big)$ . Indeed, a close examination of the arguments of [BAGJ21] reveals that for their results to hold, the proper symmetry to consider is the ellipsoidal symmetry, and
the proper definition of information exponent is the degree of the first non-zero term in the Hermite expansion of  $\tilde{R}$ , which reads
$$
\tilde {R} (z) = - \sum_ {j \geq 1} \langle f, h _ {j} \rangle^ {2} z ^ {j} + \text {c o n s t}.
$$
Once again, we can consistently define the information exponent to be the degree of the first non-zero term in the Hermite expansion of  $f$ , as long as the input is Gaussian (potentially anisotropic).
C Proofs of Section 2.1
Before beginning our main discussions, we state the following lemma which is a generalization of Stein's lemma (Gaussian integration by parts), and will help obtain a closed-form expression for the population gradient. We refer to [Erd15, MHPG+23] for similar statements.
Lemma 11. Let  $f, g: \mathbb{R} \to \mathbb{R}$  with  $g$  weakly differentiable. Suppose  $z \sim \mathcal{N}(0, \mathbf{I}_d)$ . Then, for any  $w, u \in \mathbb{S}^{d-1}$ , we have
$$
\mathbb {E} [ f (\langle \boldsymbol {w}, \boldsymbol {z} \rangle) g (\langle \boldsymbol {u}, \boldsymbol {z} \rangle) \boldsymbol {z} ] = \mathbb {E} [ f (\langle \boldsymbol {w}, \boldsymbol {z} \rangle) g ^ {\prime} (\langle \boldsymbol {u}, \boldsymbol {z} \rangle) ] \boldsymbol {u} + \mathbb {E} [ f (\langle \boldsymbol {w}, \boldsymbol {z} \rangle) \{g (\langle \boldsymbol {u}, \boldsymbol {z} \rangle) \langle \boldsymbol {w}, \boldsymbol {z} \rangle - g ^ {\prime} (\langle \boldsymbol {u}, \boldsymbol {z} \rangle) \langle \boldsymbol {u}, \boldsymbol {w} \rangle \} ] \boldsymbol {w}.
$$
Proof. Consider the conditional distribution  $z|\langle \pmb{w},\pmb{z}\rangle \sim \mathcal{N}(\overline{\pmb{\mu}},\overline{\pmb{\Sigma}})$ , where
$$
\bar {\boldsymbol {\mu}} = \langle \boldsymbol {w}, \boldsymbol {z} \rangle \boldsymbol {w} \quad \text {a n d} \quad \bar {\boldsymbol {\Sigma}} = \mathbf {I} _ {d} - \boldsymbol {w} \boldsymbol {w} ^ {\top}.
$$
Recall that Stein's lemma (Gaussian integration by parts) states that when  $\overline{z} \sim \mathcal{N}(\overline{\mu}, \overline{\Sigma})$ , then
$$
\mathbb {E} [ g (\overline {{z}}) \overline {{z}} ] = \mathbb {E} [ g (\overline {{z}}) ] \overline {{\boldsymbol {\mu}}} + \overline {{\boldsymbol {\Sigma}}} \mathbb {E} [ \nabla g (\overline {{z}}) ].
$$
$$
\mathbb {E} [ g (\langle \boldsymbol {u}, \boldsymbol {z} \rangle) \boldsymbol {z} | \langle \boldsymbol {w}, \boldsymbol {z} \rangle ] = \mathbb {E} [ g (\langle \boldsymbol {u}, \boldsymbol {z} \rangle) | \langle \boldsymbol {w}, \boldsymbol {z} \rangle ] \langle \boldsymbol {w}, \boldsymbol {z} \rangle \boldsymbol {w} + (\mathbf {I} _ {d} - \boldsymbol {w} \boldsymbol {w} ^ {\top}) \mathbb {E} [ g ^ {\prime} (\langle \boldsymbol {u}, \boldsymbol {z} \rangle) | \langle \boldsymbol {w}, \boldsymbol {z} \rangle ] \boldsymbol {u}.
$$
Applying the tower property of conditional expectation and rearranging the terms yields the desired result.
We are now ready to state and prove the expression for the population gradient when using the ReLU activation.
Lemma 12. Suppose  $\phi$  is the ReLU activation and  $g$  is weakly differentiable. Let  $z\sim \mathcal{N}(0,\mathbf{I}_d)$ . Define  $\overline{\boldsymbol{w}}\coloneqq \frac{\boldsymbol{\Sigma}^{1 / 2}\boldsymbol{w}}{\|\boldsymbol{\Sigma}^{1 / 2}\boldsymbol{w}\|}$  and similarly define  $\overline{\boldsymbol{u}}$ . Then
$$
\nabla R (\boldsymbol {w}) = \boldsymbol {\Sigma} \left\{\tilde {\psi} _ {\phi , g} (\boldsymbol {w}) \boldsymbol {w} + \tilde {\zeta} _ {\phi , g} (\boldsymbol {w}) \boldsymbol {u} \right\}
$$
$$
\tilde {\psi} _ {\phi , g} (\boldsymbol {w}) := \frac {\mathbb {E} [ - \phi (\langle \overline {{\boldsymbol {w}}} , \boldsymbol {z} \rangle) g (\langle \overline {{\boldsymbol {u}}}, \boldsymbol {z} \rangle) + \phi^ {\prime} (\langle \overline {{\boldsymbol {w}}}, \boldsymbol {z} \rangle) g ^ {\prime} (\langle \overline {{\boldsymbol {u}}}, \boldsymbol {z} \rangle) \langle \overline {{\boldsymbol {u}}}, \overline {{\boldsymbol {w}}} \rangle ]}{\| \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {w} \|} + \frac {1}{2},
$$
$$
\tilde {\zeta} _ {\phi , g} (\boldsymbol {w}) := - \frac {\mathbb {E} [ \phi^ {\prime} (\langle \overline {{\boldsymbol {w}}} , \boldsymbol {z} \rangle) g ^ {\prime} (\langle \overline {{\boldsymbol {u}}} , \boldsymbol {z} \rangle) ]}{\| \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {u} \|}.
$$
Proof. Notice that the population risk is given by
$$
R (\boldsymbol {w}) = \frac {1}{2} \mathbb {E} \left[ \left(\phi (\langle \boldsymbol {w}, \boldsymbol {x} \rangle) - g \left(\frac {\langle \boldsymbol {u} , \boldsymbol {x} \rangle}{\| \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {u} \|}\right)\right) ^ {2} \right] + \frac {\mathbb {E} [ \epsilon^ {2} ]}{2} = \frac {1}{2} \mathbb {E} \big [ \phi (\langle \boldsymbol {w}, \boldsymbol {x} \rangle) ^ {2} \big ] - \mathbb {E} [ \phi (\langle \boldsymbol {w}, \boldsymbol {x} \rangle) g ] + \frac {\mathbb {E} [ y ^ {2} ]}{2},
$$
Notice that  $\pmb{x} = \pmb{\Sigma}^{1/2}\pmb{z}$ . By the homogeneity of ReLU, we can rewrite the first term as
$$
\mathbb {E} \left[ \phi (\langle \boldsymbol {w}, \boldsymbol {x} \rangle) ^ {2} \right] = \boldsymbol {w} ^ {\top} \boldsymbol {\Sigma} \boldsymbol {w} \| \phi \| _ {\gamma} ^ {2} = \frac {\boldsymbol {w} ^ {\top} \boldsymbol {\Sigma} \boldsymbol {w}}{2}.
$$
$$
\nabla R (\boldsymbol {w}) = \frac {\boldsymbol {\Sigma} \boldsymbol {w}}{2} - \underbrace {\mathbb {E} \left[ \phi (\langle \boldsymbol {w} , \boldsymbol {x} \rangle) ^ {\prime} g \left(\frac {\langle \boldsymbol {u} , \boldsymbol {x} \rangle}{\| \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {u} \|}\right) \boldsymbol {x} \right]} _ {=: v (\boldsymbol {w})}
$$
Note that  $\phi'(\langle \pmb{w}, \pmb{x} \rangle) = \phi'(\langle \overline{\pmb{w}}, \pmb{z} \rangle)$ . Then,
$$
\begin{array}{l} v (\boldsymbol {w}) = \Sigma^ {1 / 2} \mathbb {E} \left[ \phi^ {\prime} \left(\langle \overline {{\boldsymbol {w}}}, \boldsymbol {z} \rangle\right) ^ {\prime} g \left(\langle \overline {{\boldsymbol {u}}}, \boldsymbol {z} \rangle\right) \boldsymbol {z} \right] \\ = \boldsymbol {\Sigma} ^ {1 / 2} \left\{\mathbb {E} \left[ \phi^ {\prime} (\langle \overline {{\boldsymbol {w}}}, \boldsymbol {z} \rangle) g ^ {\prime} (\langle \overline {{\boldsymbol {u}}}, \boldsymbol {z} \rangle) \right] \overline {{\boldsymbol {u}}} + \mathbb {E} \left[ \phi (\langle \overline {{\boldsymbol {w}}}, \boldsymbol {z} \rangle g (\langle \overline {{\boldsymbol {u}}}, \boldsymbol {z} \rangle) - \phi^ {\prime} (\langle \overline {{\boldsymbol {w}}}, \boldsymbol {z} \rangle) g ^ {\prime} (\langle \overline {{\boldsymbol {u}}}, \boldsymbol {z} \rangle) \langle \overline {{\boldsymbol {u}}}, \overline {{\boldsymbol {w}}} \rangle ] \overline {{\boldsymbol {w}}} \right] \right. \\ = \boldsymbol {\Sigma} \left\{\frac {\mathbb {E} \left[ \phi^ {\prime} (\langle \overline {{\boldsymbol {w}}} , \boldsymbol {z} \rangle) g ^ {\prime} (\langle \overline {{\boldsymbol {u}}} , \boldsymbol {z} \rangle) \right]}{\| \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {u} \|} \boldsymbol {u} + \frac {\mathbb {E} \left[ \phi (\langle \overline {{\boldsymbol {w}}} , \boldsymbol {z} \rangle g (\langle \overline {{\boldsymbol {u}}} , \boldsymbol {z} \rangle) - \phi^ {\prime} (\langle \overline {{\boldsymbol {w}}} , \boldsymbol {z} \rangle) g ^ {\prime} (\langle \overline {{\boldsymbol {u}}}, \boldsymbol {z} \rangle) \langle \overline {{\boldsymbol {u}}}, \overline {{\boldsymbol {w}}} \rangle ] \boldsymbol {w} \right.}{\| \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {w} \|} \right\} \\ \end{array}
$$
where we used Lemma 11 and the fact that  $\phi'(z)z = \phi(z)$ . Therefore,
Particularly, the above lemma yields the following corollary for the spherical dynamics in the population limit.
Corollary 13. Suppose  $\{\pmb{w}^t\}_{t\geq 0}$  is a solution to the population spherical gradient flow (2.5),  $\phi$  is the ReLU activation, and  $\pmb{\Sigma}$  follows the  $(\kappa ,\pmb {\theta})$  -spiked model. Then,
$$
\begin{array}{l} \frac {\mathrm {d} \langle \boldsymbol {w} ^ {t} , \boldsymbol {u} \rangle}{\mathrm {d} t} = - \frac {\tilde {\zeta} _ {\phi , g} (\boldsymbol {w} ^ {t}) (1 - \langle \boldsymbol {w} ^ {t} , \boldsymbol {u} \rangle^ {2})}{1 + \kappa} \\ - \frac {\kappa \left\{\tilde {\psi} _ {\phi , g} \left(\boldsymbol {w} ^ {t}\right) \langle \boldsymbol {w} ^ {t} , \boldsymbol {\theta} \rangle + \tilde {\zeta} _ {\phi , g} \left(\boldsymbol {w} ^ {t}\right) \langle \boldsymbol {u} , \boldsymbol {\theta} \rangle \right\}}{1 + \kappa} (\langle \boldsymbol {\theta}, \boldsymbol {u} \rangle - \langle \boldsymbol {w} ^ {t}, \boldsymbol {\theta} \rangle \langle \boldsymbol {w} ^ {t}, \boldsymbol {u} \rangle). \tag {C.1} \\ \end{array}
$$
Plugging in  $\pmb{\theta} = \pmb{u}$  in Corollary 13, we obtain
$$
\frac {\mathrm {d} \langle \boldsymbol {w} ^ {t} , \boldsymbol {u} \rangle}{\mathrm {d} t} = - \left\{\tilde {\zeta} _ {\phi , g} (\boldsymbol {w} ^ {t}) + \frac {\kappa \tilde {\psi} _ {\phi , g} (\boldsymbol {w} ^ {t}) \langle \boldsymbol {w} ^ {t} , \boldsymbol {u} \rangle}{1 + \kappa} \right\} (1 - \langle \boldsymbol {u}, \boldsymbol {w} ^ {t} \rangle^ {2}). \tag {C.2}
$$
To prove the statement of the theorem, we will show that whenever
$$
| \langle \boldsymbol {w} ^ {t}, \boldsymbol {u} \rangle | \leq \frac {C}{\sqrt {d}},
$$
we will have  $\frac{\mathrm{d}\left\langle\boldsymbol{w}^t,\boldsymbol{u}\right\rangle^2}{\mathrm{d}t} < 0$ , thus when initialized from  $\mathcal{O}(1 / \sqrt{d})$ ,  $\langle \boldsymbol{w}^t,\boldsymbol{u}\rangle$  can never escape the saddle point near the equator  $\langle \boldsymbol{w},\boldsymbol{u}\rangle = 0$ .
Recall from the properties of the Hermite expansion in Appendix B that