Learning Orthogonal Multi-Index Models: A Fine-Grained Information Exponent Analysis

## Abstract

Abstract The information exponent ( Ben Arous et al. [ 2021 ] ) — which is equivalent to the lowest degree in the Hermite
expansion of the link function for Gaussian single-index models — has played an important role in predicting the
sample complexity of online stochastic gradient descent (SGD) in various learning tasks. In this work, we
demonstrate that, for multi-index models, focusing solely on the lowest degree can miss key structural details
of the model and result in suboptimal rates. Specifically, we consider the task of learning target functions of form f ∗ ⁢ ( 𝒙 ) = ∑ k = 1 P ϕ ⁢ ( 𝒗 k ∗ ⋅ 𝒙 ) subscript 𝑓 𝒙 superscript subscript 𝑘 1 𝑃 italic-ϕ ⋅ superscript subscript 𝒗 𝑘 𝒙 f_{*}({\bm{x}})=\sum_{k=1}^{P}\phi({\bm{v}}_{k}^{*}\cdot{\bm{x}}) italic_f start_POSTSUBSCRIPT ∗ end_POSTSUBSCRIPT ( bold_italic_x ) = ∑ start_POSTSUBSCRIPT italic_k = 1 end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_P end_POSTSUPERSCRIPT italic_ϕ ( bold_italic_v start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT start_POSTSUPERSCRIPT ∗ end_POSTSUPERSCRIPT ⋅ bold_italic_x ) ,
where P ≪ d much-less-than 𝑃 𝑑 P\ll d italic_P ≪ italic_d , the ground-truth directions { 𝒗 k ∗ } k = 1 P superscript subscript superscript subscript 𝒗 𝑘 𝑘 1 𝑃 \{{\bm{v}}_{k}^{*}\}_{k=1}^{P} { bold_italic_v start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT start_POSTSUPERSCRIPT ∗ end_POSTSUPERSCRIPT } start_POSTSUBSCRIPT italic_k = 1 end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_P end_POSTSUPERSCRIPT are orthonormal, and only the second and 2 ⁢ L 2 𝐿 2L 2 italic_L -th Hermite coefficients of the link function ϕ italic-ϕ \phi italic_ϕ can be nonzero.
Based on the theory of information exponent, when the lowest degree is 2 ⁢ L 2 𝐿 2L 2 italic_L , recovering the directions requires d 2 ⁢ L − 1 ⁢ poly ⁡ ( P ) superscript 𝑑 2 𝐿 1 poly 𝑃 d^{2L-1}\operatorname{poly}(P) italic_d start_POSTSUPERSCRIPT 2 italic_L - 1 end_POSTSUPERSCRIPT roman_poly ( italic_P ) samples, and when the lowest degree is 2 2 2 2 , only the relevant subspace (not the exact
directions) can be recovered due to the rotational invariance of the second-order terms. In contrast, we show that
by considering both second- and higher-order terms, we can first learn the relevant space via the second-order
terms, and then the exact directions using the higher-order terms, and the overall sample and complexity of online
SGD is d ⁢ poly ⁡ ( P ) 𝑑 poly 𝑃 d\operatorname{poly}(P) italic_d roman_poly ( italic_P ) .

## 1 Introduction

In many learning problems, the target function exhibits or is assumed to exhibit a low-dimensional structure.
A classical model of this type is the multi-index model, where the target function depends only on a
$P$-dimensional subspace of the ambient space $\mathbb{R}^{d}$, with $P$ typically much smaller than $d$.
When the relevant dimension $P=1$, the model is known as the single-index model, which dates back to at least
Ichimura [].
Both single- and multi-index models have been widely studied, especially in the context of neural network
and stochastic gradient descent (SGD) in recent years, sometimes under the name “feature learning”(Ben Arous et al. [], Bietti et al. [], Damian et al. [], Abbe et al. [, ], Damian et al. [], Oko et al. [], Dandi et al. []).

In Ben Arous et al. [], the authors show that for single-index models, the behavior of online SGD can be split
into two phases: an initial “searching” phase, where most of the samples are used boost the correlation with the
relevant (one-dimensional) subspace to a constant, and a subsequent “descending” phase, where the correlation further
increases to $1$.
They introduce the concept of the information exponent (IE), defined as the index of the first nonzero coefficient
in the Taylor expansion of the population loss around 0 0, which corresponds to the lowest degree in the Hermite
expansion of the link function in Gaussian single-index models.
They prove that the sample complexity of online SGD is $\tilde{O}(d)$ when $\mathrm{IE}=2$ and $\tilde{O}(d^{k-1})$ when $\mathrm{IE}=k\geq 3$.
After that, various lower and upper bounds have been established for single-index models in Bietti et al. [], Damian et al. [, ].
Similar results for certain multi-index models have also been derived in Abbe et al. [, ], Oko et al. []. In all cases, the sample complexity of online SGD scales with $d^{\mathrm{IE}-1}$ when $\mathrm{IE}\geq 3$.(^1^11This bound can be significantly improved if we allow SGD to reuse batches (Dandi et al. [], Lee et al. []). For simplicity, we will not consider this technique and focus on online SGD in this work.)

For multi-index models of form $f_{*}({\bm{x}})=\sum_{k=1}^{P}\phi_{k}({\bm{v}}_{k}^{*}\cdot{\bm{x}})$, another layer of complexity arises.
In this setting, there are two types of recovery: recovering each direction ${\bm{v}}_{k}^{*}$ (strong recovery) and recovering
the subspace spanned by $\{{\bm{v}}_{k}^{*}\}_{k}$. The former notion is stronger, because once the directions are known, the learning
task essentially reduces to learning the one-dimensional $\phi_{k}:\mathbb{R}\to\mathbb{R}$ for each $k\in[P]$.
However, strong recovery is not always possible. To see this, consider the case
$\phi_{k}(z)=h_{2}(z)$, where $h_{L}$ is the $L$-th (normalized) Hermite polynomial. One can show that this corresponds
to decomposing the projection matrix (a second-order tensor) of the subspace $\mathrm{span}\{{\bm{v}}_{k}^{*}\}_{k}$.
If the model is isotropic in the relevant subspace, recovering the directions is impossible due to the rotational
invariance (see Section 3.1 for more discussion).
In contrast, when $\phi_{k}(z)=h_{2}(z)+h_{4}(z)$, the identifiability property of the fourth-order
tensor decomposition problem allows strong recovery via tensor power method or (stochastic) gradient
descent (Ge et al. [], Li et al. [], Ge et al. []). Note that in both examples,
the information exponent is $2$, indicating that information exponent alone does not distinguish between these
two scenarios.

This leads to a natural question: Can we combine the above results for orthogonal multi-index models by
first using the second-order terms to recover the subspace and then using the higher-order terms to learn the directions?
Ideally, the first stage would require at most $\tilde{O}(d\operatorname{poly}(P))$ samples, consistent with the case $\mathrm{IE}=2$, and
once the subspace is recovered, later steps would also cost at most $d\operatorname{poly}(P)$ samples.(^2^22 The $d$ factor in the second stage comes from the fact that the typical squared norm of the noise is $d$, so we have to choose the step size to be $O(d^{-1})$ for the noise to be reasonably small.)
This would yield an overall $\tilde{O}(d\operatorname{poly}(P))$ sample (and also time) complexity for strong recovery of the ground-truth directions.
Note that the $d$-denpendence matches the $\mathrm{IE}=2$ case and the strong recovery guarantee aligns with the results for
$\mathrm{IE}>2$. In this work, we prove the following theorem, providing a positive answer to this question.

###### Theorem 1.1 (Informal version of Theorem 2.1 ) .

Suppose that the target function is $f_{*}({\bm{x}})=\sum_{k=1}^{P}\phi({\bm{v}}_{k}^{*}\cdot{\bm{x}})$ where $\phi=h_{2}+h_{2L}$
($L\geq 2$) and $\{{\bm{v}}_{k}^{*}\}_{k=1}^{P}$ are orthonormal, and the input ${\bm{x}}$ follows the standard Gaussian distribution
$\mathcal{N}\left(0,\bm{I}_{d}\right)$. Then, we can use online SGD (followed by a ridge regression step) to train a two-layer network
of width $\operatorname{poly}(P)$ to learn (with high probability) this target function using $\tilde{O}(d\operatorname{poly}(P))$ samples and steps.

###### Remark .

For simplicity, we assume the link function is $\phi=h_{2}+h_{2L}$. Our results can be extended to more
general even link function, provided their Hermite coefficients decay sufficiently fast.
See Section 2 (in particular Lemma 2.1 and Lemma 2.2) for further discussion.
∎

###### Theorem 1.1 (Informal version of Theorem 2.1 ) .

###### Remark .

##### Organization

The rest of the paper is organized as follows. First, we review the related works and summarize our contributions.
Then, we describe the detailed setting and state the formal version of the main theorem in Section 2.
In Section 3, we discuss the easier case where the training algorithm is population gradient
flow. Then, in Section 4, we show how to convert the gradient flow analysis to an online SGD one.
Finally, we conclude in Section 5.
The proofs, simulation results, and a table of contents can be found in the appendix.

### 1.1 Related work

In this subsection, we discuss works that are directly related to ours or were not covered earlier in the introduction.

Along the line of information exponent, the paper most related to ours is (Oko et al. []). They show that for
near orthogonal multi-index models, the sample complexity of recovering all ground-truth directions using online SGD is
$\tilde{O}(Pd^{\mathrm{IE}-1})$ when $\mathrm{IE}\geq 3$. However, their results do not apply to the case $\mathrm{IE}=2$ for the reason
we have discussed earlier. Our result considers the situation where both $\mathrm{IE}=2$ and $\mathrm{IE}\geq 3$ terms are present
and show that in this case, the sample complexity of online SGD is $\tilde{O}(d\operatorname{poly}(P))$.

During the writing of this manuscript, we became aware of the concurrent work (Ben Arous et al. []).
Our main results are not directly comparable since the settings are different. In addition, only the lowest degree is
considered in their work. However, they also show (in a different setting) that when the second order term is isotropic,
the initial randomness can be preserved throughout training. A similar idea is used in our analysis of
Stage 1.1 (cf. Section 3.1).

Another related line of research is learning two-layer networks in the teacher-student setting
(Zhong et al. [], Li and Yuan [], Tian [], Li et al. [], Zhou et al. [], Ge et al. []).
Among them, the ones most relevant to this work are (Li et al. []) and the follow-up (Ge et al. []),
both of which consider orthogonal models similar to ours and use similar ideas in the analysis of the population process.
However, they do not assume a low-dimensional structure and only provide very crude $\operatorname{poly}(d)$-style sample complexity
bounds.

### 1.2 Our contributions

We summarize our contributions as follows:

- •
We demonstrate that information exponent alone is insufficient to characterize certain structures in the learning
task and show that for a specific orthogonal multi-index model, if we consider both the lower- and higher-order
terms, the sample complexity of strong recovery using online SGD can be greatly improved over the vanilla
information exponent-based analysis.
- •
In the analysis, we prove that when the second-order term is isotropic, the initial randomness can be preserved
during training and the relevant subspace can be recovered using $\tilde{O}(d\operatorname{poly}(P))$ samples. To the best of our knowledge,
this has only been shown by the concurrent work (Ben Arous et al. []) in a different setting.
- •
As a by-product, we provide a collection of user-friendly technical lemmas to analyze difference between noisy
one-dimensional processes and their deterministic counterparts, which may be of independent interests
(see Section 4.1 and Section F.2).

## 2 Setup and main result

In this section, we describe the setting of our learning task and the training algorithm. Then we formally state our
main result. We will also convert the problem to an orthogonal tensor decomposition task using the standard Hermite
argument (Ge et al. []).

##### Notations

We use $\left\|\cdot\right\|_{p}$ to denote the $p$-norm of a vector. When $p=2$, we often drop the subscript and simply write
$\left\|\cdot\right\|$. For $a,b,\delta\in\mathbb{R}$, $a=b\pm\delta$ means $|a-b|\leq|\delta|$ and
$a\vee b=\max\{a,b\}$ and $a\wedge b=\min\{a,b\}$. Beside the standard
asymptotic (big $O$) notations, we also use the notation $f_{d}=O_{L}(g_{d})$, which means there exists a constant $C_{L}>0$
that can depend only on $L$ such that $f_{d}\leq C_{L}g_{d}$ for all large enough $d$. Sometimes we also write
$f_{d}\lesssim_{L}g_{d}$ for $f_{d}=O_{L}(g_{d})$. The actual value of $C_{L}$ can vary between lines, but we will typically
point this out when it does.

### 2.1 Input and target function

We assume the input ${\bm{x}}$ follows the standard Gaussian distribution $\mathcal{N}\left(0,\bm{I}_{d}\right)$ and the target function has
form
$f_{*}({\bm{x}})=\sum_{k=1}^{P}\phi({\bm{v}}_{k}^{*}\cdot{\bm{x}})$,
where $\log^{C}d\leq P\leq d$ for a large universal constant $C>0$,
$\{{\bm{v}}_{k}^{*}\}_{k=1}^{P}$ are orthonormal and $\phi(z)=h_{2}(z)+h_{2L}(z)$ with $L\geq 2$ and $h_{l}:\mathbb{R}\to\mathbb{R}$ being the
$l$-th (normalized) Hermite polynomial.

Our target model and algorithm will all be invariant under rotation. Hence, we may assume without
loss of generality that ${\bm{v}}_{k}^{*}=\bm{e}_{k}$ where $\{\bm{e}_{k}\}_{k}$ is the standard basis of $\mathbb{R}^{d}$. For now, we continue
writing ${\bm{v}}_{k}^{*}$ since most of the results in this section do not depend on the orthonormality of $\{{\bm{v}}_{k}^{*}\}_{k}$.

### 2.2 Learner model, loss function and its gradient

Our learner model is a width-$m$ two-layer network

$$ $f({\bm{x}}):=f({\bm{x}};\bm{a},\bm{V}):=\sum_{k=1}^{m}a_{k}\phi({\bm{v}}_{k} \cdot{\bm{x}}),$ $$

where $\bm{a}=(a_{1},\dots,a_{m})\in\mathbb{R}^{m}$ and $\bm{V}=({\bm{v}}_{1},\dots,{\bm{v}}_{m})\in(\mathbb{S}^{d-1})^{m}$ are the trainable parameters.
We will call $\{{\bm{v}}_{k}\}_{k\in[m]}$ the first-layer neurons. We measure the difference between the learner and
the target model using the mean-square error (MSE). Given a sample $({\bm{x}},f_{*}({\bm{x}}))$, we define the per-sample
loss as

$$ $l({\bm{x}}):=l({\bm{x}};\bm{a},\bm{V}):=\frac{1}{2}\left(f_{*}({\bm{x}})-f({ \bm{x}})\right)^{2}.$ $$

For convenience, we denote the population MSE loss with $\mathcal{L}:=\mathcal{L}(\bm{a},\bm{V}):=\mathop{\mathbb{E\/}}_{{\bm{x}}}l({
\bm{x}};\bm{a},\bm{V})$.
With Hermite expansion, one can rewrite $\mathcal{L}$ as a tensor decomposition loss as in the following lemma.
The proof of this lemma is standard and can be found in, for example, Ge et al. []. We also provide a proof
in Appendix A for completeness.

###### Lemma 2.1 (Population loss) .

Consider the setting described above. For $l\in\mathbb{N}_{\geq 0}$, let $\hat{\phi}_{l}$ denote the $l$-th Hermite
coefficient of $\phi$ (with respect to the normalized Hermite polynomials). Then, for the population loss, we have

$$ $\mathcal{L}=\mathrm{Const.}-\sum_{l=0}^{\infty}\sum_{k=1}^{P}\sum_{j=1}^{m}a_{ j}\hat{\phi}_{l}^{2}\left\langle{\bm{v}}_{k}^{*},{\bm{v}}_{j}\right\rangle^{l} +\frac{1}{2}\sum_{l=0}^{\infty}\sum_{j_{1},j_{2}=1}^{m}a_{j_{1}}a_{j_{2}}\hat{ \phi}_{l}^{2}\left\langle{\bm{v}}_{j_{1}},{\bm{v}}_{j_{2}}\right\rangle^{l},$ (1) $$

where $\mathrm{Const.}$ is a real number that does not depend on $\bm{a}$ nor $\bm{V}$.

###### Remark .

The lemma does not require $\{{\bm{v}}_{k}^{*}\}_{k}$ to be orthonormal nor $\phi=h_{2}+h_{2L}$. All we need is $\phi\in L^{2}(\mathcal{N}\left(0,\bm{I}_{d}\right))$ so that the Hermite expansion is well-defined.
∎

For the per-sample and population gradients, we have the following lemma, the proof of which can also be found in
Appendix A.

###### Lemma 2.2 (First-layer gradients) .

Consider the setting described above. Suppose that $\phi=h_{2}+h_{2L}$ and $|a_{k}|\leq a_{0}$ for some
$a_{0}>0$ and all $k\in[m]$. Then, for each $i\in[m]$, we have

$$ $\nabla_{{\bm{v}}_{i}}\mathcal{L}=-2a_{i}\sum_{k=1}^{P}\left\langle{\bm{v}}_{k} ^{*},{\bm{v}}_{i}\right\rangle{\bm{v}}_{k}^{*}-2La_{i}\sum_{k=1}^{P}\left \langle{\bm{v}}_{k}^{*},{\bm{v}}_{i}\right\rangle^{2L-1}{\bm{v}}_{k}^{*}\pm_{2 }2Lma_{0}^{2},$ (2) $$

where ${\bm{z}}={\bm{z}}^{\prime}\pm_{2}\delta$ means $\left\|{\bm{z}}-{\bm{z}}^{\prime}\right\|_{2}\leq\delta$.

Moreover, for ${\bm{x}}\sim\mathcal{N}\left(0,\bm{I}_{d}\right)$ and every direction $\bm{u}\in\mathbb{S}^{d-1}$ that is independent of ${\bm{x}}$, there
exists a constant $C_{L}>0$ that can depend only on $L$ such that

$$ $\displaystyle\mathop{\mathbb{P\/}}\left(a_{0}^{-1}\left|\left\langle\nabla_{{ \bm{v}}_{i}}l({\bm{x}})-\nabla_{{\bm{v}}}\mathcal{L},\bm{u}\right\rangle\right |\geq s\right)$ $\displaystyle\leq C_{L}\exp\left(-\frac{1}{C_{L}}\left(\frac{s}{P}\right)^{1/( 2L)}\right),$ $\displaystyle\mathop{\mathbb{P\/}}\left(a_{0}^{-1}\left\|\nabla_{{\bm{v}}_{i}} l({\bm{x}})-\nabla_{{\bm{v}}}\mathcal{L}\right\|\geq s\right)$ $\displaystyle\leq C_{L}\exp\left(\log d-\frac{1}{C_{L}}\left(\frac{s}{P\sqrt{d }}\right)^{1/(2L)}\right),$ $\displaystyle a_{0}^{-2}\mathop{\mathbb{E\/}}_{{\bm{x}}}\left\langle\nabla_{{ \bm{v}}_{i}}l({\bm{x}}),\bm{u}\right\rangle^{2}$ $\displaystyle\leq C_{L}P^{2}.$ $$

###### Remark on the population gradient .

Note that (2) implies that when $a$ is small, the dynamics of different neurons are
approximately decoupled. This allows us to consider each neuron separately. The same is also true when we consider
the per-sample gradient. Hence, we can often drop the subscript $i$ and say ${\bm{v}}:={\bm{v}}_{i}$ is an arbitrary first-layer
neuron and the (population) gradient with respect to it is given by (2).
∎

###### Remark on the tail bounds .

We will choose $m=\operatorname{poly}(P)$. In this case, in order for the RHS of the bounds to be $o(1)$ (after applying the union
bound over all $m$ neurons), it suffices to choose $s=\omega(P\log^{2L}P)$ and $s=\omega(Pd^{1/2}\log^{2L}d)$.
Up to some logarithmic terms, this matches what one should expect when $\nabla_{{\bm{v}}_{i}}l({\bm{x}})$ is a $P^{2}$-subgaussian
random vector.
∎

###### Remark on possible extensions .

The formula (2) and the tail and variance bounds in this lemma are essentially
all the structures we need (besides the orthonormality) to establish our results.
To extend our results to general even link function whose Hermite coefficients decay sufficiently fast, first note that
the second-order and then the $2L$-th order (the lowest even order that is larger than $2$) terms dominate the gradient.
Moreover, since $\{{\bm{v}}_{k}^{*}\}_{k}$ are assumed to be orthonormal, for any fixed even order (that is larger than $4$), the
minimizer of the corresponding terms matches the ground-truth directions, and the gradient will always push the
neurons toward one of the ground-truth directions. In other words, they only help the model recover the directions.
We consider only the lowest order since it determines the overall complexity (as in the theory of information exponent).

Our tail bound is based on Theorem 1.3 of Adamczak and Wolff [] (cf. Theorem A.1), which deals with polynomials of a fixed degree. Theorem 1.2 of Adamczak and Wolff []
deals with general functions with controlled higher-order derivatives and can be used to extend our result to
non-polynomial link functions. See Appendix G for an empirical evidence.
∎

###### Lemma 2.1 (Population loss) .

###### Remark .

###### Lemma 2.2 (First-layer gradients) .

###### Remark on the population gradient .

###### Remark on the tail bounds .

###### Remark on possible extensions .

### 2.3 Training algorithm

Now, we describe the training algorithm. First, we initialize each output weight $a_{i}$ to be $a_{0}$ where
$a_{0}>0$ is a hyperparameter to be determined later and ${\bm{v}}_{i}\sim\mathrm{Unif}(\mathbb{S}^{d-1})$ independently.
Then, we fix the output weights $\bm{a}$ and train the first-layer weight ${\bm{v}}_{i}$ using online (spherical) SGD with step size
$\eta/a_{0}$ ($\eta>0$) for $T$ iterations. Then, we fix the first-layer weights and use ridge regression to train the
output weights $\bm{a}$.

Let $\{({\bm{x}}_{t},f_{*}({\bm{x}}_{t}))\}_{t\in\mathbb{N}}$ be our samples where $\{{\bm{x}}_{t}\}$ are i.i.d. standard Gaussian vectors,
and let $\tilde{\nabla}_{{\bm{v}}}=(\bm{I}-{\bm{v}}{\bm{v}}^{\top})\nabla_{{\bm{v}}}$ denote the spherical gradient.
Then, we can formally describe the training procedure as follows:

$$ Initialization: $\displaystyle a_{0,k}=a_{0},\quad{\bm{v}}_{0,k}\overset{\mathrm{i.i.d.}}{\sim} \mathrm{Unif}(\mathbb{S}^{d-1}),$ $\displaystyle\forall k\in[m];$ (3) Stage 1: $\displaystyle{\bm{v}}_{t+1,k}={\bm{v}}_{t,k}-\frac{\eta}{a_{0}}\tilde{\nabla}_ {{\bm{v}}_{k}}l({\bm{x}}_{t};\bm{a}_{0},\bm{V}_{t}),$ $\displaystyle\forall k\in[m],t\in[T];$ Stage 2: $\displaystyle\bm{a}=\operatorname*{argmin}_{\bm{a}^{\prime}}\frac{1}{2N}\sum_{ n=1}^{N}l({\bm{x}}_{T+n};\bm{a}^{\prime},\bm{V}_{T})+\lambda\left\|\bm{a}^{ \prime}\right\|^{2}.$ $$

Here, the hyperparameters are the initialization scale $a_{0}>0$, network width $m>0$, step size $\eta>0$,
time horizon $T>0$, the number of samples $N$ in Stage 2, and the regularization strength $\lambda>0$.

Before move on, we make some remarks here on the training algorithm. As we have seen in Lemma 2.1 and Lemma 2.2, when the second-layer
weights are small, the dynamics of the first-layer weights are roughly decoupled. Hence, we choose to initialize each
$a_{i}$ small and fix them at $a_{0}$ in Stage 1. We rescale the learning rate with $1/a_{0}$ to compensate the
fact that the first-layer gradients are proportional to $a_{0}$.

We will show that after the first stage, for each ground truth direction
${\bm{v}}_{k}^{*}$, there will be some neurons ${\bm{v}}_{i}$ that converge to that direction.
As a result, in the second stage, we can use ridge regression to pick out those neurons and use them to fit the target
function. The analysis of this stage is standard and has been done in (Damian et al. [], Abbe et al. [], Ba et al. [], Lee et al. [], Oko et al. []). Hence, we will not further discuss this stage in the
main text and defer the proofs for this stage to Appendix D.

### 2.4 Main result

The following is our main result. The proof of it can be found in Appendix E.

###### Theorem 2.1 (Main Theorem) .

Consider the setting and algorithm described above. Let $C>0$ be a large universal constant.
Suppose that $\log^{C}d\leq P\leq d$ and $\{{\bm{v}}^{*}_{k}\}_{k=1}^{P}$ are orthonormal.
Let $\delta_{\mathop{\mathbb{P\/}}}\in(\exp(-\log^{C}d),1)$ and $\varepsilon_{*}>0$ be given. Suppose that we choose $a_{0},\eta,T,N$
satisfying

$$ $\displaystyle m=\Omega\left(P^{8}\log^{1.5}(P\vee 1/\delta_{\mathop{\mathbb{P \/}}})\right),\quad a_{0}=O_{L}\left(\frac{\varepsilon_{*}^{2}}{mdP^{2L+2}\log ^{3}d\log(1/\varepsilon_{*})}\right),\quad N=\Omega_{L}\left(\frac{Pm}{ \varepsilon_{*}^{2}\delta_{\mathop{\mathbb{P\/}}}^{2}}\right),$ $\displaystyle\eta=O_{L}\left(\frac{\varepsilon_{*}^{4}\delta_{\mathop{\mathbb{ P\/}}}}{dP^{L+8}\log^{4L+1}(d/\delta_{\mathop{\mathbb{P\/}}})}\right)=\tilde{O }_{L}\left(\frac{\varepsilon_{*}^{4}\delta_{\mathop{\mathbb{P\/}}}}{dP^{L+8}} \right),$ $\displaystyle T=O_{L}\left(\frac{\log d+P^{L-1}+\log(P/\varepsilon_{*})}{\eta} \right)=\tilde{O}_{L}\left(\frac{dP^{2L+7}}{\delta_{\mathop{\mathbb{P\/}}} \varepsilon_{*}^{4}}\right).$ $$

Then, there exists some $\lambda>0$ such that at the end of training, we have $\mathcal{L}(\bm{a},\bm{V})\leq\varepsilon$
with probability at least $1-O(\delta_{\mathop{\mathbb{P\/}}})$.

###### Theorem 2.1 (Main Theorem) .

## 3 The gradient flow analysis

In this section, we consider the situation where the training algorithm in Stage 1 is gradient flow over the population
loss instead of online SGD. The discussion here is non-rigorous and our formal proof does not rely on anything in
this section. Nevertheless, this gradient flow analysis will provide valuable intuition on the behavior of online SGD
and also lead to rough guesses on the time complexity.

For notational simplicity, we will assume without loss of generality that ${\bm{v}}_{k}^{*}=\bm{e}_{k}$. Let ${\bm{v}}$ be an arbitrary
first-layer neuron. By Lemma 2.2, when we rescale the time by $a_{0}^{-1}$, the dynamics of
${\bm{v}}$ are controlled by(^3^33We use $\tau$ to index the time in this continuous-time process (as $t$ has been used to index the steps in the discrete-time process) and will often omit it when it is clear from the context.)

$$ $\dot{{\bm{v}}}_{\tau}\approx 2\sum_{k=1}^{P}v_{k}(\bm{I}-{\bm{v}}{\bm{v}}^{ \top})\bm{e}_{k}+2L\sum_{k=1}^{P}v_{k}^{2L-1}(\bm{I}-{\bm{v}}{\bm{v}}^{\top}) \bm{e}_{k}.$ $$

The second term on the RHS comes from the normalized/projection. For each $k\in[d]$, we have

$$ $\frac{\mathrm{d}}{\mathrm{d}\tau}v_{k}^{2}\approx 4\mathbbm{1}\{k\leq P\}\left (1+Lv_{k}^{2L-2}\right)v_{k}^{2}-4\left(\left\|{\bm{v}}_{\leq P}\right\|^{2}+L \left\|{\bm{v}}_{\leq P}\right\|_{2L}^{2L}\right)v_{k}^{2}.$ (4) $$

We further split Stage 1 into two substages. In Stage 1.1, the second-order terms dominate and
$\left\|{\bm{v}}_{\leq P}\right\|^{2}/\left\|{\bm{v}}_{>P}\right\|^{2}$ grows from $\Theta(P/d)$ to $\Theta(1)$. In Stage 1.2,
${\bm{v}}$ converges to one ground-truth direction.

The direction to which ${\bm{v}}$ will converge depends on the index of the largest $v_{k}^{2}$ at the beginning of Stage 1.2.
With some standard concentration/anti-concentration argument, one can show that $\max_{k\in[P]}v_{k}^{2}$ is at least $1+c$ times larger
than the second-largest $v_{k}^{2}$ for a small constant $c>0$ with probability at least $1/\operatorname{poly}(P)$ at initialization
(of Stage 1.1). Hence, as long as this gap can be preserved throughout Stage 1, we can choose $m=\operatorname{poly}(P)$
to ensure all ground-truth directions can be found after Stage 1.2.

### 3.1 Stage 1.1: learning the subspace and preservation of the gap

In this substage, we track $\left\|{\bm{v}}_{\leq P}\right\|^{2}/\left\|{\bm{v}}_{>P}\right\|^{2}$ and $v_{p}^{2}/v_{q}^{2}$(^4^44A slightly different quantity will be used in the online SGD analysis, but the intuition remains the same.) where $p,q\in[P]$
are arbitrary. The goal is to show that $\left\|{\bm{v}}_{\leq P}\right\|^{2}/\left\|{\bm{v}}_{>P}\right\|^{2}$ will grow to a constant while
$v_{p}^{2}/v_{q}^{2}$ stay close to its initial value.

For the norm ratio, by (4), we have

$$ $\displaystyle\frac{\mathrm{d}}{\mathrm{d}\tau}\frac{\left\|{\bm{v}}_{\leq P} \right\|^{2}}{\left\|{\bm{v}}_{>P}\right\|^{2}}$ $\displaystyle=\frac{\frac{\mathrm{d}}{\mathrm{d}\tau}\left\|{\bm{v}}_{\leq P} \right\|^{2}}{\left\|{\bm{v}}_{>P}\right\|^{2}}-\frac{\left\|{\bm{v}}_{\leq P} \right\|^{2}}{\left\|{\bm{v}}_{>P}\right\|^{2}}\frac{\frac{\mathrm{d}}{\mathrm {d}\tau}\left\|{\bm{v}}_{>P}\right\|^{2}}{\left\|{\bm{v}}_{>P}\right\|^{2}}$ $\displaystyle=\frac{4\left\|{\bm{v}}_{\leq P}\right\|^{2}}{\left\|{\bm{v}}_{>P }\right\|^{2}}+\frac{4L\left\|{\bm{v}}_{\leq P}\right\|_{2L}^{2L}}{\left\|{\bm {v}}_{>P}\right\|^{2}}-\cancel{\frac{4\left(\left\|{\bm{v}}_{\leq P}\right\|^{ 2}+L\left\|{\bm{v}}_{\leq P}\right\|_{2L}^{2L}\right)}{\left\|{\bm{v}}_{>P} \right\|^{2}}}$ $\displaystyle\qquad+\cancel{\frac{\left\|{\bm{v}}_{\leq P}\right\|^{2}}{\left \|{\bm{v}}_{>P}\right\|^{2}}\frac{4\left(\left\|{\bm{v}}_{\leq P}\right\|^{2}+ L\left\|{\bm{v}}_{\leq P}\right\|_{2L}^{2L}\right)\left\|{\bm{v}}_{>P}\right\| ^{2}}{\left\|{\bm{v}}_{>P}\right\|^{2}}}.$ $$

In particular, note that the terms coming from normalization cancel with each other. Moreover,
this implies $\frac{\mathrm{d}}{\mathrm{d}\tau}\frac{\left\|{\bm{v}}_{\leq P}\right\|^{2}}{
\left\|{\bm{v}}_{>P}\right\|^{2}}\geq 4\frac{\left\|{\bm{v}}_{\leq P}\right\|^
{2}}{\left\|{\bm{v}}_{>P}\right\|^{2}}$,
and therefore, it takes only at most $\frac{1+o(1)}{4}\log(d/P)=\Theta(\log(d/P))$ amount of time for the ratio to grow from $\Theta(P/d)$ to $\Theta(1)$.
If we choose a small step size $\eta$ so that online SGD closely tracks the gradient flow, then the number of
steps one should expect is $O(\log(d/P)/\eta)$.

Meanwhile, for any $p,q\in[P]$, we have

$$ $\displaystyle\frac{\mathrm{d}}{\mathrm{d}\tau}\frac{v_{p}^{2}}{v_{q}^{2}}$ $\displaystyle=4\left(1+Lv_{p}^{2L-2}\right)\frac{v_{p}^{2}}{v_{q}^{2}}-4\left( \left\|{\bm{v}}_{\leq P}\right\|^{2}+L\left\|{\bm{v}}_{\leq P}\right\|_{2L}^{2 L}\right)\frac{v_{p}^{2}}{v_{q}^{2}}$ $\displaystyle\qquad-\frac{v_{p}^{2}}{v_{q}^{2}}\left(4\left(1+Lv_{q}^{2L-2} \right)-4\left(\left\|{\bm{v}}_{\leq P}\right\|^{2}+L\left\|{\bm{v}}_{\leq P} \right\|_{2L}^{2L}\right)\right)=4L\left(v_{p}^{2L-2}-v_{q}^{2L-2}\right)\frac {v_{p}^{2}}{v_{q}^{2}}.$ $$

Note that not only those terms coming from normalization cancel with each other, but also the second-order terms.
In particular, this also implies that we cannot learn the directions using only the second-order terms.
At initialization, it is unlikely that some $v_{k}^{2}$ are significantly larger than all other $v_{l}^{2}$. Hence, if
we assume the induction hypothesis $v_{p}^{2}/v_{q}^{2}\approx v_{0,p}^{2}/v_{0,q}^{2}$, we will have $v_{k}^{2}\leq\tilde{O}(1/P)$
and the above will become $\frac{\mathrm{d}}{\mathrm{d}\tau}v_{p}^{2}/v_{q}^{2}\leq\tilde{O}(L/P)v_{p}^{2
}/v_{q}^{2}$. As a result,
$v_{t,p}^{2}/v_{t,q}^{2}\leq(1+o(1))v_{0,p}^{2}/v_{0,q}^{2}$ for any $t\leq\Theta(\log(d/P))$,
as long as $P\geq\operatorname{poly}\log d$.

### 3.2 Stage 1.2: learning the directions

Let ${\bm{v}}$ be a first-layer neuron with $v_{1}^{2}\geq(1+c)\max_{2\leq k\leq P}v_{k}^{2}$ for some small constant
$c>0$ at initialization. By our previous discussion, we know at the end of Stage 1.1, the above bound still holds
with a potentially smaller constant $c>0$. In addition, since $\left\|{\bm{v}}_{\leq P}\right\|^{2}=\Theta(1)$, we also have $v_{1}^{2}\geq\Omega(1/P)$ at the end of Stage 1.1. We claim that ${\bm{v}}$ will converge to $\bm{e}_{1}$. The argument here is similar to
the proofs in Li et al. [] and Ge et al. [].

Again, by (4), we have

$$ $\frac{\mathrm{d}}{\mathrm{d}\tau}v_{1}^{2}\approx 4\left(1-\left\|{\bm{v}}_{ \leq P}\right\|^{2}+Lv_{1}^{2L-2}-L\left\|{\bm{v}}_{\leq P}\right\|_{2L}^{2L} \right)v_{1}^{2}\geq 4L\left(v_{1}^{2L-2}-\left\|{\bm{v}}_{\leq P}\right\|_{2L }^{2L}\right)v_{1}^{2}.$ $$

Assume the induction hypothesis $v_{1}^{2}\geq(1+c)\max_{2\leq k\leq P}v_{k}^{2}$ and write

$$ $v_{1}^{2L-2}-\left\|{\bm{v}}_{\leq P}\right\|_{2L}^{2L}=v_{1}^{2L-2}\left(1-v_ {1}^{2}\right)-\left(\left\|{\bm{v}}_{\leq P}\right\|^{2}-v_{1}^{2}\right)\sum _{k=2}^{P}\frac{v_{k}^{2}}{\left\|{\bm{v}}_{\leq P}\right\|^{2}-v_{1}^{2}}v_{k }^{2L-2}.$ $$

Note that the summation is a weighted average of $\{v_{k}^{2L-2}\}_{k\geq 2}$ and therefore is upper bounded by
$\left(v_{1}^{2}/(1+c)\right)^{L-1}\leq(1-c_{L})v_{1}^{2L-2}$ for some constant $c_{L}>0$ that can only depend on $L$.
Thus, we have

$$ $\frac{\mathrm{d}}{\mathrm{d}\tau}v_{1}^{2}\gtrsim 4L\left(1-v_{1}^{2}-\left( \left\|{\bm{v}}_{\leq P}\right\|^{2}-v_{1}^{2}\right)(1-c_{L})\right)v_{1}^{2L }\geq 4c_{L}L\left(1-v_{1}^{2}\right)v_{1}^{2L}.$ $$

When $v_{1}^{2}\leq 3/4$, this implies $\frac{\mathrm{d}}{\mathrm{d}\tau}v_{1}^{2}\geq c_{L}Lv_{1}^{2L}$. As a result, it takes at most
$O_{L}(P^{L-1})$ amount of time for $v_{1}^{2}$ to grow from $\Omega(1/P)$ to $3/4$.
It is important that $v_{1}^{2}=\Omega(1/P)$ instead of $\Omega(1/d)$ at the start of Stage 1.2, since otherwise the time
needed will be $O_{L}(d^{L-1})$. After $v_{1}^{2}$ reaches $3/4$, we have $\frac{\mathrm{d}}{\mathrm{d}\tau}(1-v_{1}^{2})\leq-4c_{L}L(3/4)^{2L}\left(1-v_
{1}^{2}\right)$. Thus, $v_{1}^{2}$ will converge linearly to $1$ afterwards.

## 4 From gradient flow to online SGD

In this section, we discuss how to convert the previous gradient flow analysis to an online SGD one. Our actual proof
will be based directly on the online SGD analysis, but the overall idea is still proving that the online SGD dynamics
of certain important quantities closely track their population gradient descent (GD) counterparts. Our choice of
learning rate $\eta$ will be much smaller than what needed for GD to track GF — the bottleneck comes from the
GD-to-SGD conversion, not the GF-to-GD one. In other words, provided that SGD tracks GD well, the number of
steps/samples it needs to finish each substage is roughly the amount of time GF needs, divided by the step size $\eta$.

The rest of this section is organized as follows. In Section 4.1, we collect a few useful lemmas for controlling the difference between noisy dynamics
and their deterministic counterparts. The idea behind them has appeared in Ben Arous et al. [] and is also
used in Abbe et al. []. Here, we simplify and slightly generalize their argument and provide
a user-friendly interface. When used properly, it reduces the GD-to-SGD proof to routine calculus. Then, in
Section 4.2, we discuss how to apply those general results to analyze the
dynamics of online SGD in our setting.

### 4.1 Technical lemmas for analyzing general noisy dynamics

We start with the lemma that will be used to analyze $\left\|{\bm{v}}_{\leq P}\right\|^{2}/\left\|{\bm{v}}_{>P}\right\|^{2}$. The proof of it and
all other lemmas in this subsection can be found in Section F.2.

###### Lemma 4.1 .

Let $(\Omega,\mathcal{F},(\mathcal{F}_{t})_{t\in\mathbb{N}},\mathop{\mathbb{P\/}})$ be a filtered probability space. Suppose that $(X_{t})_{t}$ is an
$(\mathcal{F}_{t})_{t}$-adapted real-valued process satisfying

$$ $X_{t+1}=X_{t}+\alpha X_{t}+\xi_{t+1}+Z_{t+1},\quad X_{0}=x_{0}>0,$ (5) $$

where $\alpha>0$ is fixed, $(\xi_{t})_{t}$ is an $(\mathcal{F}_{t})_{t}$-adapted process, and $(Z_{t})_{t}$ is an $(\mathcal{F}_{t})_{t}$-adapted
martingale difference sequence. Define its deterministic counterpart as $x_{t}=(1+\alpha)^{t}x_{0}$.

Let $T>0$ and $\delta_{\mathop{\mathbb{P\/}}}\in(0,1)$ be given. Suppose that there exists some $\delta_{\mathop{\mathbb{P\/}},\xi}\in(0,1)$ and
$\Xi,\sigma_{Z}>0$ such that for every $t\leq T$, if $X_{t}=(1\pm 0.5)x_{t}$, then we have $|\xi_{t+1}|\leq(1+\alpha)^{t}\Xi$ with probability at least $1-\delta_{\mathop{\mathbb{P\/}},\xi}$ and $\mathop{\mathbb{E\/}}[Z_{t+1}^{2}\mid\mathcal{F}_{t}]\leq(1+\alpha)^{t}\sigma_
{Z}^{2}$. If

$$ $\Xi\leq\frac{x_{0}}{4T}\quad\text{and}\quad\sigma_{Z}^{2}\leq\frac{\delta_{ \mathop{\mathbb{P\/}}}\alpha x_{0}^{2}}{16},$ (6) $$

then we have $X_{t}=(1\pm 0.5)x_{t}$ for all $t\in[T]$ with probability at least $1-T\delta_{\mathop{\mathbb{P\/}},\xi}-\delta_{\mathop{\mathbb{P\/}}}$.

###### Remark on condition ( 6 ) .

One may interpret $Z_{t+1}$ as those terms coming from the difference between the population and mini-batch gradients
and $\xi_{t+1}$ as the higher-order error terms. $\alpha$ is usually small. In our case, it is proportional to the
step size $\eta$. $T$ is usually the time needed for $X_{t}$ to grow from a small $x_{0}>0$ to $\Theta(1)$, which is
roughly $\alpha^{-1}\log(1/x_{0})$. In other words, we have $\alpha=\tilde{O}(1/T)$. As a result, in order for
(6) to hold, it suffices to have $\Xi=O(x_{0}/T)$ and $\sigma_{Z}=O(x_{0}/\sqrt{T})$.
Note that the condition on $\sigma_{Z}$ is much weaker than the condition on $\Xi$. Meanwhile, since $\xi_{t+1}$ models
the higher-order error terms, we should expect it to be able to satisfy the stronger condition $\Xi\leq O(1/T)$.
∎

###### Remark on stochastic induction .

One important feature of this lemma is that it only requires the bounds $|\xi_{t+1}|\leq(1+\alpha)^{t}\Xi$ and
$\mathop{\mathbb{E\/}}[Z_{t+1}^{2}\mid\mathcal{F}_{t}]\leq(1+\alpha)^{t}\sigma_
{Z}^{2}$ to hold when $X_{t}=(1\pm 0.5)x_{t}$. This can be viewed
as a form of induction. This is particularly useful when considering the dynamics of, say, $v_{k}^{2}$.
Similar to how the RHS of $\frac{\mathrm{d}}{\mathrm{d}\tau}v_{\tau,k}^{2}=2v_{\tau,k}\dot{v}_{\tau,k}$ depends on $v_{\tau,k}$,
the size of $\xi_{t+1}$ and $Z_{t+1}$ will usually depend on $X_{t}$. Hence, we will not be able to bound them without
an induction hypothesis on $X_{t}$.
∎

###### Remark on the dependence on $\delta_{\mathop{\mathbb{P\/}}}$ .

The dependence on $\delta_{\mathop{\mathbb{P\/}}}$ can be improved to $\operatorname{poly}\log(1/\delta_{\mathop{\mathbb{P\/}}})$ if we have tail bounds on $Z_{t+1}$
similar to the ones in Lemma 2.2.
We state this lemma in this simpler form because we will only take union bound over $\operatorname{poly}(P)$ events, and we are not
optimizing the dependence on $P$. We include in Section F.2 an example
(cf. Lemma F.9 and Lemma F.10) where this
improvement is made (though that result will not be used in the proof).
∎

###### Proof sketch of Lemma 4.1 .

For the ease of presentation, we assume that $|\xi_{t+1}|\leq(1+\alpha)^{t}\Xi$ with probability at least
$1-\delta_{\mathop{\mathbb{P\/}},\xi}$ and $\mathop{\mathbb{E\/}}[Z_{t+1}^{2}\mid\mathcal{F}_{t}]\leq(1+\alpha)^{t}\sigma_
{Z}^{2}$ always hold. Recursively
expand the RHS of (5), and we obtain

$$ $X_{t+1}=(1+\alpha)^{t+1}x_{0}+\sum_{s=1}^{t}(1+\alpha)^{t-s}\xi_{s+1}+\sum_{s= 1}^{t}(1+\alpha)^{t-s}Z_{s+1}.$ $$

Divide both sides with $(1+\alpha)^{t+1}$ and replace $t+1$ with $t$. Then, the above becomes

$$ $X_{t}(1+\alpha)^{-t}=x_{0}+\sum_{s=1}^{t}(1+\alpha)^{-s}\xi_{s}+\sum_{s=1}^{t} (1+\alpha)^{-s}Z_{s}.$ $$

The second term is bounded by $T\Xi$ (uniformly over $t\leq T$) with probability at least $1-T\delta_{\mathop{\mathbb{P\/}},\xi}$.
Note that $(1+\alpha)^{-s}Z_{s}$ is still a martingale difference sequence. Hence, by Doob’s $L^{2}$-submartingale
inequality, the third term is bounded by $x_{0}/4$ with probability at least $16\sigma_{Z}^{2}/(\alpha x_{0}^{2})$.
Thus, when (6) holds, the RHS is $(1\pm 0.5)x_{0}$ with probability at
least $1-T\delta_{\mathop{\mathbb{P\/}},\xi}-\delta_{\mathop{\mathbb{P\/}}}$. Multiply both sides with $(1+\alpha)^{t}$, and we complete the proof.
∎

Using the same strategy, one can prove a similar lemma (cf. Lemma F.8) that
deals with the case $\alpha=0$, which will be used to show the preservation of the gap in Stage 1.1. Another
interesting case is where the growth is not linear but polynomial. This is the case of Stage 1.2 in our setting.
For this case, we have the following lemma.

###### Lemma 4.2 .

Suppose that $(X_{t})_{t}$ satisfies

$$ $X_{t+1}=X_{t}+\alpha X_{t}^{p}+\xi_{t+1}+Z_{t+1},\quad X_{0}=x_{0}>0,$ $$

where $p>1$, the signal growth rate $\alpha>0$ and initialization $x_{0}>0$ are given and fixed, $(\xi_{t})_{t}$ is
an adapted process, and $(Z_{t})_{t}$ is a martingale difference sequence. Let $\hat{x}_{t}$ be the solution to
the deterministic recurrence relationship

$$ $\hat{x}_{t+1}=\hat{x}_{t}+\alpha\hat{x}_{t}^{p},\quad\hat{x}_{0}=x_{0}/2.$ $$

Fix $T>0,\delta_{\mathop{\mathbb{P\/}}}\in(0,1)$. Suppose that there exist $\Xi,\sigma_{Z}>0$ and $\delta_{\mathop{\mathbb{P\/}},\xi}\in(0,1)$
such that when $X_{t}\geq\hat{x}_{t}$, we have $|\xi_{t}|\leq\Xi$ with probability at least $1-\delta_{\mathop{\mathbb{P\/}},\xi}$ and
$\mathop{\mathbb{E\/}}[Z_{t+1}\mid\mathcal{F}_{t}]\leq\sigma_{Z}^{2}$. Then, if

$$ $\Xi\leq\frac{x_{0}}{4T}\quad\text{and}\quad\sigma_{Z}^{2}\leq\frac{x_{0}^{2} \delta_{\mathop{\mathbb{P\/}}}}{16T},$ $$

we have $X_{t}\geq\hat{x}_{t}$ for all $t\leq T$.

The proof is essentially the same as the previous one, except that we need to replace $(1+\alpha)^{t}$ with
$\prod_{s=0}^{t-1}(1+\alpha X_{s}^{p-1})$. Let $x_{t}$ be the version of $\hat{x}_{t}$ with the initial value being
$x_{0}$ instead of $x_{0}/2$. Unlike the linear case, here it is generally difficult to ensure $X_{t}\geq x_{t}/2$ since
this type of polynomial systems exhibits sharp transitions and blows up in finite time. In fact, the difference between
the deterministic processes $\hat{x}_{t}$ and $x_{t}/2$ can be large. However, if one is only interested in the time needed
for $X_{t}$ to grow from a small value to a constant, then results obtained from $\hat{x}_{t}$ and $x_{t}$ differ only by
a multiplicative constant, and when $\alpha>0$ is small, both of them can be estimated using their continuous-time
counterpart $\dot{x}_{\tau}=x_{\tau}^{p}$ (cf. Lemma F.12).

###### Lemma 4.1 .

###### Remark on condition ( 6 ) .

###### Remark on stochastic induction .

###### Remark on the dependence on δ ℙ subscript 𝛿 ℙ \delta_{\mathop{\mathbb{P\/}}} italic_δ start_POSTSUBSCRIPT blackboard_P end_POSTSUBSCRIPT .

###### Proof sketch of Lemma 4.1 .

###### Lemma 4.2 .

### 4.2 Sample complexity of online SGD

In this subsection, we demonstrate how to use the previous results to obtain results for online SGD and discuss why
the sample complexity is $\tilde{O}(d\operatorname{poly}(P))$ instead of $\tilde{O}(d^{2L-1})$ even though we are relying on the
$2L$-th order terms to learn the directions.

#### 4.2.1 A simplified version of Stage 1.1

As an example, we consider the dynamics of $Pv_{p}^{2}/(dv_{q}^{2})$ where $p\leq P$ and $q>P$ and assume both of $v_{p}$ and $v_{q}$ are
small and $Pv_{p}^{2}/(dv_{q}^{2})\leq 1$.
This can be viewed as a simplified version of the analysis of $\left\|{\bm{v}}_{\leq P}\right\|^{2}/\left\|{\bm{v}}_{>P}\right\|^{2}$ in Stage 1.1.
The analysis of other quantities/stages is essentially the same — we rewrite the update rule to single out martingale
difference terms and the higher-order error terms, and apply a suitable lemma from the previous
subsection (or Section F.2) to complete the proof.

For the ease of presentation, in this subsection, we ignore the higher-order terms. In particular, we assume the approximation

$$ $\hat{v}_{t+1,k}\approx v_{t,k}+2\eta\left(\mathbbm{1}\{k\leq P\}-\left\|{\bm{v }}_{\leq P}\right\|^{2}\right)+\eta Z_{t+1,k},\qquad\forall k\in[d],$ $$

where $Z_{t+1,k}$ represents the difference between the population and mini-batch gradients. Then, we compute

$$ $\hat{v}_{t+1,k}^{2}\approx\left(1+4\eta\left(\mathbbm{1}\{k\leq P\}-\left\|{ \bm{v}}_{\leq P}\right\|^{2}\right)\right)v_{k}^{2}+2\eta v_{k}Z_{k}\pm C_{L} \eta^{2}(1\vee Z_{k}^{2}).$ $$

Here, the last term is the higher-order term and will eventually be included in $\xi$. For simplicity,
we will also ignore them in the following discussion.
The second term is the martingale difference term. Its (conditional) variance depend on $v_{k}$, and this necessitates
the induction-style conditions in Lemma F.6. Note that
$v_{t+1,p}^{2}/v_{t+1,q}^{2}=\hat{v}_{t+1,p}^{2}/\hat{v}_{t+1,q}^{2}$. Hence, we have

$$ $\frac{v_{t+1,p}^{2}}{v_{t+1,q}^{2}}\approx\frac{\left(1+4\eta\left(1-\left\|{ \bm{v}}_{\leq P}\right\|^{2}\right)\right)v_{p}^{2}+2\eta v_{p}Z_{p}}{\left(1- 4\eta\left\|{\bm{v}}_{\leq P}\right\|^{2}\right)v_{q}^{2}+2\eta v_{q}Z_{q}}.$ $$

For any small $a>0$ and small $\delta>0$, we have the following elementary identity:
$\frac{1}{a+\delta}=\frac{1}{a}\left(1-\frac{\delta}{a}\left(1-\frac{\delta}{a+
\delta}\right)\right)\approx\frac{1}{a}\left(1-\frac{\delta}{a}\right).$
Repeatedly use this identity, and we can rewrite the above equation as

$$ $\displaystyle\frac{Pv_{t+1,p}^{2}}{dv_{t+1,q}^{2}}$ $\displaystyle\approx\frac{P\left(1+4\eta\left(1-\left\|{\bm{v}}_{\leq P}\right \|^{2}\right)\right)v_{p}^{2}}{d\left(1-4\eta\left\|{\bm{v}}_{\leq P}\right\|^ {2}\right)v_{q}^{2}}\left(1-\frac{2\eta v_{q}Z_{q}}{\left(1-4\eta\left\|{\bm{v }}_{\leq P}\right\|^{2}\right)v_{q}^{2}}\right)$ $\displaystyle\qquad+\frac{2P\eta v_{p}Z_{p}}{d\left(1-4\eta\left\|{\bm{v}}_{ \leq P}\right\|^{2}\right)v_{q}^{2}}\left(1-\frac{2\eta v_{q}Z_{q}}{\left(1-4 \eta\left\|{\bm{v}}_{\leq P}\right\|^{2}\right)v_{q}^{2}}\right)$ $\displaystyle\approx\left(1+4\eta\right)\frac{Pv_{p}^{2}}{dv_{q}^{2}}-\frac{Pv _{p}^{2}}{dv_{q}^{2}}\frac{2\eta v_{q}Z_{q}}{v_{q}^{2}}+\frac{2P\eta v_{p}Z_{p }}{dv_{q}^{2}}.$ $$

Suppose that $v_{p}^{2}\approx v_{q}^{2}$ at initialization and assume the induction hypothesis $Pv_{p}^{2}/(dv_{q}^{2})=(1\pm 0.5)(1+4\eta)^{t}Pv_{0,p}^{2}/(dv_{0q}^{2})$. Then, by Lemma 2.2, the conditional variance of the martingale
difference terms (the last two terms) is bounded by $O_{L}((1+4\eta)^{t}\eta^{2}P^{4}/d)$. Using the language of
Lemma 4.1, this means $\sigma_{Z}^{2}\leq O_{L}(\eta^{2}P^{4}/d)$. Hence, in order
for (the second condition of) (6) to hold, it suffices to choose
$\eta\lesssim_{L}\delta_{\mathop{\mathbb{P\/}}}/(dP^{2})$. By our gradient flow analysis, the number steps Stage 1.1 needs is roughly
$\log d/\eta$. In other words, for Stage 1.1, the sample complexity is $\tilde{O}_{L}(dP^{2}/\delta_{\mathop{\mathbb{P\/}}})$ (if we
ignore the higher-order error terms).

#### 4.2.2 The improved sample complexity for Stage 1.2

To see why the existence of the second-order terms can reduce the sample complexity from $d^{\mathrm{IE}-1}$ to $d\operatorname{poly}(P)$,
first note that after Stage 1.1, $\max_{p\in[P]}v_{p}^{2}$ will be $\Omega(1/P)$. Also
note that the conditions in Lemma 4.2 depend on the initial value.
With the initial value being $\Omega(1/P)$ instead of $\tilde{O}(1/d)$, the largest possible step size we can choose
will be $O(1)/(d\operatorname{poly}(P))$, which is much larger than the usual $O(1/d^{L-1})$ requirement from the vanilla
information exponent argument. Meanwhile, by our gradient flow analysis, we know the number of iterations needed
is $O(P^{L-1}/\eta)$. Combine these and we obtain the $d\operatorname{poly}(P)$ sample complexity.

## 5 Conclusion and future directions

In this work, we study the task of learning multi-index models of form $f_{*}({\bm{x}})=\sum_{k=1}^{P}\phi({\bm{v}}_{k}^{*}\cdot{\bm{x}})$
with $P\ll d$, $\{{\bm{v}}_{k}^{*}\}_{k}$ be orthogonal and $\phi=h_{2}+h_{2L}$. By considering both the lower- and
higher-order terms, we prove an $\tilde{O}(d\operatorname{poly}(P))$ bound on the sample complex for strong recovery of
directions using online SGD, which improve the results one can obtain using vanilla information exponent-based analysis.

One possible future direction of our work is to generalize our results to more general link functions and assume the
learner model is a generic two-layer network with, say, ReLU activation. Another interesting but more challenging
direction is to consider the non-(near)-orthogonal case. We conjecture when the target model has a hierarchical
structure across different orders, online SGD can gradually learn the directions using those terms of different
order sequentially.

## Acknowledgements

JDL acknowledges support of NSF CCF 2002272, NSF IIS 2107304, NSF CIF 2212262, ONR Young Investigator Award,
and NSF CAREER Award 2144994.

## Appendix A From multi-index model to tensor decomposition

In this section, we show that the task of learning the multi-index target function $f_{*}({\bm{x}})=\sum_{k=1}^{P}\phi({\bm{v}}_{k}^{*}\cdot{\bm{x}})$
can be reduced to tensor decomposition. We will need the following classical result on Hermite polynomials
(cf. Chapter 11.2 of O’Donnell []) and correlated Gaussian variables.

###### Lemma A.1 (Proposition 11.31 of O’Donnell [ 2014 ] ) .

For $k\in\mathbb{N}_{\geq 0}$ denote the normalized Hermite polynomials.
Let $\rho\in[-1,1]$ and $z,z^{\prime}$ be $\rho$-correlated standard Gaussian variables. Then, we have

$$ $\mathop{\mathbb{E\/}}_{z,z^{\prime}}\left[h_{k}(z)h_{j}(z^{\prime})\right]= \mathbbm{1}\{k=j\}\rho^{k}.$ $$

See 2.1

###### Proof.

By definition, we have

$$ $\displaystyle\mathcal{L}$ $\displaystyle=\frac{1}{2}\mathop{\mathbb{E\/}}_{{\bm{x}}\sim\mathcal{N}\left(0 ,\bm{I}_{d}\right)}\left(\sum_{k=1}^{P}\phi({\bm{v}}_{k}^{*}\cdot x)-\sum_{j=1 }^{m}a_{j}\phi({\bm{v}}_{j}\cdot{\bm{x}})\right)^{2}$ $\displaystyle=\frac{1}{2}\sum_{k_{1},k_{2}=1}^{P}\mathop{\mathbb{E\/}}_{{\bm{x }}\sim\mathcal{N}\left(0,\bm{I}_{d}\right)}\left\{\phi({\bm{v}}_{k_{1}}^{*} \cdot x)\phi({\bm{v}}_{k_{2}}^{*}\cdot x)\right\}-\sum_{k=1}^{P}\sum_{j=1}^{m} a_{j}\mathop{\mathbb{E\/}}_{{\bm{x}}\sim\mathcal{N}\left(0,\bm{I}_{d}\right)} \left\{\phi({\bm{v}}_{k}^{*}\cdot x)\phi({\bm{v}}_{j}\cdot{\bm{x}})\right\}$ $\displaystyle\qquad+\frac{1}{2}\sum_{j_{1},j_{2}=1}^{m}a_{j_{1}}a_{j_{2}} \mathop{\mathbb{E\/}}_{{\bm{x}}\sim\mathcal{N}\left(0,\bm{I}_{d}\right)}\left \{\phi({\bm{v}}_{j_{1}}\cdot x)\phi({\bm{v}}_{j_{2}}\cdot{\bm{x}})\right\}.$ $$

The first term is independent of $\bm{a}$ and $\bm{V}$. For the other two terms, we now use Lemma A.1 to evaluate the expectation. Let $\phi=\sum_{k=0}^{\infty}\hat{\phi}_{k}h_{k}$ be the Hermite
expansion of $\phi$ where the convergence is in $L^{2}$ sense. For any $\rho\in[-1,1]$ and $\rho$-correlated
standard Gaussian variables $z,z^{\prime}$, we have

$$ $\mathop{\mathbb{E\/}}_{z,z^{\prime}}\left\{\phi(z)\phi(z^{\prime})\right\}= \sum_{k,l=0}^{\infty}\hat{\phi}_{k}\hat{\phi}_{l}\mathop{\mathbb{E\/}}_{z,z^{ \prime}}\left\{h_{k}(z)h_{l}(z^{\prime})\right\}=\sum_{k=0}^{\infty}\hat{\phi} _{k}^{2}\rho^{k},$ $$

where the first equality comes from the Dominated Convergence Theorem and the second from Lemma A.1. Note that ${\bm{v}}_{k}^{*}\cdot{\bm{x}}$ and ${\bm{v}}_{j}\cdot{\bm{x}}$ are $\left\langle{\bm{v}}_{k}^{*},{\bm{v}}_{j}\right\rangle$-correlated
standard Gaussian variables. Hence, by applying the above identity to the second term, and we obtain

$$ $\sum_{k=1}^{P}\sum_{j=1}^{m}a_{j}\mathop{\mathbb{E\/}}_{{\bm{x}}\sim\mathcal{N }\left(0,\bm{I}_{d}\right)}\left\{\phi({\bm{v}}_{k}^{*}\cdot x)\phi({\bm{v}}_{ j}\cdot{\bm{x}})\right\}=\sum_{l=0}^{\infty}\sum_{k=1}^{P}\sum_{j=1}^{m}a_{j} \hat{\phi}_{l}^{2}\left\langle{\bm{v}}_{k}^{*},{\bm{v}}_{j}\right\rangle^{l}.$ $$

Similarly, for the last term, we have

$$ $\frac{1}{2}\sum_{j_{1},j_{2}=1}^{m}a_{j_{1}}a_{j_{2}}\mathop{\mathbb{E\/}}_{{ \bm{x}}\sim\mathcal{N}\left(0,\bm{I}_{d}\right)}\left\{\phi({\bm{v}}_{j_{1}} \cdot x)\phi({\bm{v}}_{j_{2}}\cdot{\bm{x}})\right\}=\frac{1}{2}\sum_{l=0}^{ \infty}\sum_{j_{1},j_{2}=1}^{m}a_{j_{1}}a_{j_{2}}\hat{\phi}_{l}^{2}\left \langle{\bm{v}}_{j_{1}},{\bm{v}}_{j_{2}}\right\rangle^{l}.$ $$

∎

Then, we consider the population and per-sample gradient. It is well-known that any Lipschitz function of a Gaussian
variable is still subgaussian. Similar tail bounds can still be obtained when the function is not Lipschitz but has
a bounded higher-order derivative. To estimate the tail of the per-sample gradient, we need the following result from
Adamczak and Wolff []. As a side note, Theorem 1.2 of Adamczak and Wolff [] is a more general
result that deals with general non-Lipschitz functions with controlled higher-order derivatives. That result can
be used to extend our setting to link functions with infinitely many nonzero higher-order Hermite coefficients, given
that they decay sufficiently fast.

###### Theorem A.1 (Theorem 1.3 of Adamczak and Wolff [ 2015 ] ) .

Let $\bm{Z}\sim\mathcal{N}\left(0,\bm{I}_{d}\right)$ and $f:\mathbb{R}^{d}\to\mathbb{R}$ be a polynomial of degree $Q$. Then, for any $t\geq 0$,
we have

$$ $\mathop{\mathbb{P\/}}\left[|f(Z)-\mathop{\mathbb{E\/}}f(\bm{Z})|\geq t\right] \leq C_{Q}\exp\left(-C_{Q}^{-1}\min_{q\in[Q]}\min_{J\in P_{q}}\left(\frac{t}{ \left\|\mathop{\mathbb{E\/}}\nabla^{q}f(\bm{Z})\right\|_{J}}\right)^{2/|J|} \right),$ (7) $$

where $C_{Q}>0$ is a constant that depends only on the degree $Q$, $P_{q}$ is the collection of partitions of $[q]$,
and for any $J\in P_{q}$ and ${\bm{A}}\in(\mathbb{R}^{d})^{\otimes q}$,

$$ $\left\|{\bm{A}}\right\|_{J}:=\sup\left\{\sum_{\bm{i}\in[d]^{q}}A_{\bm{i}}\prod _{l=1}^{|J|}X^{(l)}_{\bm{i}_{J_{l}}}\;:\;\bm{X}^{(l)}\in(\mathbb{R}^{d})^{ \otimes|J_{l}|},\left\|\bm{X}^{(l)}\right\|_{F}\leq 1,\forall l\in[|J|]\right\}.$ $$

###### Remark on the definition of $\left\|\cdot\right\|_{J}$ .

The definition of $\left\|{\bm{A}}\right\|_{J}$ might look bizarre, but it has a natural functional interpretation. Given a
partition $J\in P_{q}$, we can treat a tensor ${\bm{A}}\in(\mathbb{R}^{d})^{\otimes q}$ as a multilinear function by grouping
the indices according to $J$ as follows.
For each $J_{l}\in J$, we take $\bm{X}^{(l)}\in(\mathbb{R}^{d})^{|J_{l}|}$ and feed them
into ${\bm{A}}$ to obtain a real number. Similar to how the induced norm is defined for matrices, we restrict
the norm of each $\bm{X}^{(l)}$ to be at most $1$ to obtain this definition of $\left\|{\bm{A}}\right\|_{J}$. As an example,
consider ${\bm{A}}\in(\mathbb{R}^{d})^{\otimes 3}$ and $J=\{\{1,2\},\{3\}\}$. In this case, $\bm{X}^{(1)}$ is a matrix and
$\bm{X}^{(2)}$ is a vector, and we have

$$ $\left\|{\bm{A}}\right\|_{\{1,2\},\{3\}}=\sup\left\{\sum_{i,j,k\in[d]}A_{i,j,k} X^{(1)}_{i,j}X^{(2)}_{k}\;:\;\left\|\bm{X}^{(1)}\right\|_{F}\leq 1,\left\|\bm{ X}^{(2)}\right\|_{2}\leq 1\right\}.$ $$

∎

###### Remark on the RHS of ( 7 ) .

Fix ${\bm{z}}\in\mathbb{R}^{d}$ and $f$ be a polynomial with degree at most $Q$. Suppose that the coefficients of monomials of $f$
are all bounded by some constant $A_{Q}>0$ that may depend on $Q$. Note that $f$ can contain at most $d^{Q}$ monomials.
Meanwhile, for each $q\in[Q]$ and $\bm{i}\in[d]^{q}$, $[\nabla^{q}f({\bm{z}})]_{\bm{i}}$ is nonzero only if
$[\nabla^{q}m({\bm{z}})]_{\bm{i}}$ for some monomial $m:\mathbb{R}^{d}\to\mathbb{R}$ contained in $f$. Since $m$ has degree at most $Q$,
$\nabla^{q}m({\bm{z}})$ can have at most $Q!$ nonzero entries (across all different ${\bm{z}}$). Thus, the total number of
possible nonzero entries in $\nabla^{q}f({\bm{z}})$ is bounded by $Q!d^{Q}$ and all entries of it are bounded by
$Q!A_{Q}$. Thus, we have $\left\|\mathop{\mathbb{E\/}}\nabla^{q}f(\bm{Z})\right\|_{J}\leq C_{Q}^{\prime}
d^{Q}$ for some constant $C_{Q}^{\prime}>0$ that can depend only
on $Q$. In other words, for the RHS of (7) to be $o(1)$, we need
$t=\omega(C_{Q}^{\prime}d^{Q})$.

The above bound might seem to be bad. Fortunately, in our case, we only need to consider $f:\mathbb{R}^{d}\to\mathbb{R}$ of form
$f({\bm{x}})=F(\bm{u}_{1}\cdot{\bm{x}},\bm{u}_{2}\cdot{\bm{x}},\bm{u}_{3}\cdot{
\bm{x}})$ where $F$ is a polynomial and $\bm{u}_{1},\bm{u}_{2},\bm{u}_{3}\in\mathbb{S}^{d-1}$ are
three arbitrary directions. Suppose that ${\bm{x}}\sim\mathcal{N}\left(0,\bm{I}_{d}\right)$ and define $\mathbf{\Sigma}\in\mathbb{R}^{3\times 3}$ via
$\Sigma_{i,j}=\left\langle\bm{u}_{i},\bm{u}_{j}\right\rangle$. Then, we have

$$ $f({\bm{x}})\overset{d}{=}F\left(\mathbf{\Sigma}^{1/2}{\bm{z}}\right)\quad\text {where}\quad{\bm{z}}\sim\mathcal{N}\left(0,\bm{I}_{3}\right).$ $$

When $F:\mathbb{R}^{3}\to\mathbb{R}$ is a degree-$Q$ polynomial with coefficients being constants that can depend only on $Q$,
so ${\bm{z}}\mapsto F\left(\mathbf{\Sigma}^{1/2}{\bm{z}}\right)$. Thus, we can apply this theorem (with dimension being $3$) and our
previous discussion to obtain

$$ $\mathop{\mathbb{P\/}}\left[|f(Z)-\mathop{\mathbb{E\/}}f(\bm{Z})|\geq t\right] \leq C_{Q}\exp\left(-\frac{t^{2/Q}}{C_{Q}}\right),$ $$

where $C_{Q}>0$ is a constant that can depend only on $Q$.
∎

Now, we are ready to prove Lemma 2.2, which we also restate bellow.

See 2.2

###### Proof.

Fix $i\in[m]$. First, by Lemma 2.1, we have

$$ $\displaystyle\nabla_{{\bm{v}}_{i}}\mathcal{L}$ $\displaystyle=-\sum_{k=1}^{P}a_{i}\nabla_{{\bm{v}}_{i}}\left\langle{\bm{v}}_{k }^{*},{\bm{v}}_{i}\right\rangle^{2}-\sum_{k=1}^{P}a_{i}\nabla_{{\bm{v}}_{i}} \left\langle{\bm{v}}_{k}^{*},{\bm{v}}_{i}\right\rangle^{2L}+\frac{1}{2}\sum_{l \in\{2,2L\}}\sum_{j=1}^{m}a_{i}a_{j}\nabla_{{\bm{v}}_{i}}\left\langle{\bm{v}}_ {i},{\bm{v}}_{j}\right\rangle^{l}$ $\displaystyle=-2a_{i}\sum_{k=1}^{P}\left\langle{\bm{v}}_{k}^{*},{\bm{v}}_{i} \right\rangle{\bm{v}}_{k}^{*}-2La_{i}\sum_{k=1}^{P}\left\langle{\bm{v}}_{k}^{* },{\bm{v}}_{i}\right\rangle^{2L-1}{\bm{v}}_{k}^{*}$ $\displaystyle\qquad+\frac{1}{2}a_{i}\sum_{l\in\{2,2L\}}\left(l\sum_{j\in[m] \setminus\{i\}}a_{j}\left\langle{\bm{v}}_{i},{\bm{v}}_{j}\right\rangle^{l-1}{ \bm{v}}_{j}+2la_{i}\left\langle{\bm{v}}_{i},{\bm{v}}_{i}\right\rangle^{l-1}{ \bm{v}}_{i}\right).$ $$

Note that the last line is bounded by $2Lma_{0}^{2}$. In other words,

$$ $\nabla_{{\bm{v}}_{i}}\mathcal{L}=-2a_{i}\sum_{k=1}^{P}\left\langle{\bm{v}}_{k} ^{*},{\bm{v}}_{i}\right\rangle{\bm{v}}_{k}^{*}-2La_{i}\sum_{k=1}^{P}\left \langle{\bm{v}}_{k}^{*},{\bm{v}}_{i}\right\rangle^{2L-1}{\bm{v}}_{k}^{*}\pm_{2 }2Lma_{0}^{2}.$ $$

Now, consider the per-sample gradient. We write

$$ $\displaystyle\nabla_{{\bm{v}}_{i}}l({\bm{x}})$ $\displaystyle=-\left(f_{*}({\bm{x}})-f({\bm{x}};\bm{a},\bm{V})\right)\nabla_{{ \bm{v}}_{i}}f({\bm{x}};\bm{a},\bm{V})$ $\displaystyle=-a_{i}\left(f_{*}({\bm{x}})-f({\bm{x}};\bm{a},\bm{V})\right)\phi ^{\prime}({\bm{v}}_{i}\cdot{\bm{x}}){\bm{x}}$ $\displaystyle=-a_{i}\sum_{k=1}^{P}\phi({\bm{v}}_{k}^{*}\cdot{\bm{x}})\phi^{ \prime}({\bm{v}}_{i}\cdot{\bm{x}}){\bm{x}}+a_{i}\sum_{k=1}^{m}a_{k}\phi({\bm{v }}_{k}\cdot{\bm{x}})\phi^{\prime}({\bm{v}}_{i}\cdot{\bm{x}}){\bm{x}}$ $\displaystyle=:\bm{g}_{i,1}+\bm{g}_{i,2}.$ $$

Let $\bm{u}\in\mathbb{S}^{d-1}$ be an arbitrary direction. We now estimate the tail of $\left\langle\nabla_{{\bm{v}}_{i}}l,\bm{u}\right\rangle$.
By Theorem A.1 (and the second remark following it), we have

$$ $\mathop{\mathbb{P\/}}\left(\left|\phi({\bm{v}}_{k}^{*}\cdot{\bm{x}})\phi^{ \prime}({\bm{v}}_{i}\cdot{\bm{x}})\left\langle{\bm{x}},\bm{u}\right\rangle- \mathop{\mathbb{E\/}}_{{\bm{x}}^{\prime}}\phi({\bm{v}}_{k}^{*}\cdot{\bm{x}}^{ \prime})\phi^{\prime}({\bm{v}}_{i}\cdot{\bm{x}}^{\prime})\left\langle{\bm{x}}^ {\prime},\bm{u}\right\rangle\right|\geq s\right)\leq C_{L}\exp\left(-\frac{s^{ 1/(2L)}}{C_{L}}\right),$ $$

for some constant $C_{L}>0$ that can depend only on $L$. Hence, we have

$$ $\displaystyle\mathop{\mathbb{P\/}}\left(a_{i}^{-1}\left|\left\langle\bm{g}_{i, 1},\bm{u}\right\rangle-\mathop{\mathbb{E\/}}\left\langle\bm{g}_{i,1},\bm{u} \right\rangle\right|\geq s\right)\leq C_{L}\exp\left(-\frac{(s/P)^{1/(2L)}}{C_ {L}}\right).$ $$

In particular, this implies that typical value of $a_{i}^{-1}\bm{g}_{i,1}$ is bounded by $\Theta(P)$.
Similarly, for $\bm{g}_{i,2}$, we have

$$ $\mathop{\mathbb{P\/}}\left(a_{0}^{-2}\left|\left\langle\bm{g}_{i,1},\bm{u} \right\rangle-\mathop{\mathbb{E\/}}\left\langle\bm{g}_{i,1},\bm{u}\right \rangle\right|\geq s\right)\leq C_{L}\exp\left(-\frac{(s/m)^{1/(2L)}}{C_{L}} \right),$ (8) $$

or equivalently,

$$ $\mathop{\mathbb{P\/}}\left(a_{0}^{-1}\left|\left\langle\bm{g}_{i,1},\bm{u} \right\rangle-\mathop{\mathbb{E\/}}\left\langle\bm{g}_{i,1},\bm{u}\right \rangle\right|\geq s\right)\leq C_{L}\exp\left(-\frac{(s/(a_{0}m))^{1/(2L)}}{C _{L}}\right).$ $$

Note that since $a_{0}m=o(1)\ll P$, the RHS of this inequality is much smaller than the RHS of
(8) when we choose the same $s$.
Combine the above bounds together, and we obtain that for each fixed $i\in[m]$,

$$ $\mathop{\mathbb{P\/}}\left(a_{0}^{-1}\left|\left\langle\nabla_{{\bm{v}}_{i}}l( {\bm{x}}),\bm{u}\right\rangle-\left\langle\nabla_{{\bm{v}}}\mathcal{L},\bm{u} \right\rangle\right|\geq s\right)\leq C_{L}\exp\left(-\frac{(s/P)^{1/(2L)}}{C_ {L}}\right),$ $$

for some constant $C_{L}>0$ that can depend only on $L$ and is potentially different from the $C_{L}$
in (8).
As a corollary, we have

$$ $\displaystyle\mathop{\mathbb{P\/}}\left(a_{0}^{-1}\left\|\nabla_{{\bm{v}}_{i}} l({\bm{x}})-\nabla_{{\bm{v}}}\mathcal{L}\right\|\geq s\right)$ $\displaystyle\leq\;$ $\displaystyle C_{L}\sum_{k=1}^{d}\mathop{\mathbb{P\/}}\left(a_{0}^{-1}\left| \left\langle\nabla_{{\bm{v}}_{i}}l({\bm{x}}),\bm{e}_{k}\right\rangle-\left \langle\nabla_{{\bm{v}}}\mathcal{L},\bm{e}_{k}\right\rangle\right|\geq s/\sqrt {d}\right)$ $\displaystyle\leq\;$ $\displaystyle C_{L}\exp\left(\log(d)-\frac{1}{C_{L}}\left(\frac{s}{P\sqrt{d}} \right)^{1/(2L)}\right).$ $$

Similarly, one can show that $\mathop{\mathbb{E\/}}\left\langle\nabla_{{\bm{v}}_{i}}l({\bm{x}}),\bm{u}\right
\rangle^{2}\leq C_{L}a_{0}^{2}P^{2}$
for some constant $C_{L}>0$ that can depend only on $L$ and is potentially different from the $C_{L}$
in (8).
∎

###### Lemma A.1 (Proposition 11.31 of O’Donnell [ 2014 ] ) .

###### Proof.

###### Theorem A.1 (Theorem 1.3 of Adamczak and Wolff [ 2015 ] ) .

###### Remark on the definition of ∥ ⋅ ∥ J \left\|\cdot\right\|_{J} ∥ ⋅ ∥ start_POSTSUBSCRIPT italic_J end_POSTSUBSCRIPT .

###### Remark on the RHS of ( 7 ) .

###### Proof.

## Appendix B Typical structure at initialization

In this section, we use the results in Section F.1 to analyze
the structure of ${\bm{v}}_{1},\dots,{\bm{v}}_{m}$ at initialization. Recall that we initialize ${\bm{v}}_{i}$ with $\mathrm{Unif}(\mathbb{S}^{d-1})$ independently. Meanwhile,
note that for ${\bm{v}}\sim\mathrm{Unif}(\mathbb{S}^{d-1})$, we have ${\bm{v}}\overset{d}{=}\bm{Z}/\left\|\bm{Z}\right\|$ where
$\bm{Z}\sim\mathcal{N}\left(0,\bm{I}_{d}\right)$.

We start with a lemma on the largest coordinate. This lemma ensures that $\left\|{\bm{v}}\right\|_{2L}^{2L}$ is much smaller than
the second-order terms at least at initialization.

###### Lemma B.1 (Largest coordinate) .

Let ${\bm{v}}\sim\mathrm{Unif}(\mathbb{S}^{d-1})$. For any $K\geq 1$, we have

$$ $\max_{i\in[d]}|v_{i}|\leq\frac{4\sqrt{2K\log d}}{\sqrt{d}}\quad\text{with probability at least }1-\frac{4}{d^{K}}.$ $$

As a corollary, for any $\delta_{\mathop{\mathbb{P\/}}}\in(0,1)$, at initialization, we have

$$ $\max_{i\in[m]}\left\|{\bm{v}}_{i}\right\|_{\infty}\leq\frac{4\sqrt{2\log(4m/ \delta_{\mathop{\mathbb{P\/}}})}}{\sqrt{d}}\quad\text{with probability at least $1-\delta_{\mathop{\mathbb{P\/}}}$}.$ $$

In particular, this implies that at initialization, at least with the same probability, for any $L\geq 2$,

$$ $\max_{i\in[m]}\left\|{\bm{v}}_{i}\right\|_{2L}^{2L}\leq d\left(\frac{4\sqrt{2K \log d}}{\sqrt{d}}\right)^{2L}\leq d\left(\frac{32K\log d}{d}\right)^{L}.$ $$

###### Proof.

Let $\bm{Z}\sim\mathcal{N}\left(0,\bm{I}_{d}\right)$. Recall that $\bm{Z}/\left\|\bm{Z}\right\|$ follows the uniform distribution over the sphere.
By Lemma F.1 with $s=\sqrt{d}/3$, we have $\left\|\bm{Z}\right\|\geq\sqrt{d}/2$ with
probability at least $1-2\exp(-d/18)$. Then, by Lemma F.2,
with probability at least $1-2e^{-d/18}-2e^{-s^{2}/2}$, we have

$$ $\frac{\max_{i\in[d]}|Z_{i}|}{\left\|\bm{Z}\right\|}\leq\frac{\sqrt{2\log d}+s} {\sqrt{d}/2}=\frac{2\sqrt{2\log d}}{\sqrt{d}}+\frac{2s}{\sqrt{d}}.$ $$

Let $K\geq 1$ be arbitrary. Choose $s=\sqrt{2K\log d}$ and the above becomes

$$ $\frac{\max_{i\in[d]}|Z_{i}|}{\left\|\bm{Z}\right\|}\leq\frac{4\sqrt{2K\log d}} {\sqrt{d}}\quad\text{with probability at least }1-\frac{4}{d^{K}}.$ $$

For the corollary, use union bound and choose $K=\log(4m/\delta_{\mathop{\mathbb{P\/}}})/\log d$, we have

$$ $\max_{i\in[m]}\left\|{\bm{v}}_{i}\right\|_{\infty}\leq\frac{4\sqrt{2\log(4m/ \delta_{\mathop{\mathbb{P\/}}})}}{\sqrt{d}}\quad\text{with probability at least }1-\frac{4m}{d^{K}}=1-\delta_{\mathop{\mathbb{P\/}}}.$ $$

∎

Suppose that we only have higher-order terms. Then, for a neuron ${\bm{v}}\in\mathbb{S}^{d-1}$ to converge to a ground-truth
direction $\bm{e}_{k}$ in a reasonable amount of time, we need $v_{k}^{2}$ to be the largest among all $v_{i}^{2}$ and there is gap between
it and the second largest $v_{i}^{2}$. The following lemma ensures that when $m$ is large, for every ground-truth direction
$\{\bm{e}_{k}\}_{k\in[P]}$, there will be at least one neuron satisfying the above property. Note that in our case, we
only need to ensure $v_{k}^{2}$ is the largest among all $\{v_{i}^{2}\}_{i\in[P]}$ instead of $\{v_{i}^{2}\}_{i\in[d]}$, as
the second-order term will help us identify the correct subspace.

###### Lemma B.2 (Existence of good neurons) .

Let $\delta_{\mathop{\mathbb{P\/}}}\in(0,1)$ be given and $c\geq 1$ a universal constant.
Suppose that the number of neurons $m$ satisfies

$$ $m\geq 400cP^{8c^{2}}\sqrt{\log P}\log\left(P\vee\frac{1}{\delta_{\mathop{ \mathbb{P\/}}}}\right).$ $$

Then, at initialization, with probability at least $1-\delta_{\mathop{\mathbb{P\/}}}$, we have

$$ $\displaystyle\forall p\in[P]\,\exists i\in[m]\quad\text{such that}\quad\frac{| v_{i,p}|}{\max_{q\in[P]\setminus\{p\}}|v_{i,q}|}\geq\frac{1+2c}{1+c}.$ $$

###### Remark .

In particular, note that the number of neurons we need is $\operatorname{poly}(P)$ instead of $\operatorname{poly}(d)$.
∎

###### Proof.

Let $\bm{Z}\sim\mathcal{N}\left(0,\bm{I}_{d}\right)$. Note that $|v_{p}|/|v_{q}|\overset{d}{=}|Z_{p}|/|Z_{q}|$. Hence, it suffices to
consider the largest and the second largest among $\{|Z_{i}|\}_{i\in[P]}$. Let $|v|_{(1)}$ and $|v|_{(2)}$
denote the largest and second largest among $\{|v_{i}|\}_{i\in[P]}$.
By Lemma F.4 (with $d$ replaced by $P$),
for any $c\geq 1$, we have

$$ $\mathop{\mathbb{P\/}}\left[\frac{|v|_{(1)}}{|v|_{(2)}}\geq\frac{1+2c}{1+c} \right]\geq\frac{1}{5\pi(1+2c)}\frac{1}{P^{8c^{2}}\sqrt{\log P}}.$ $$

Then, for each $p\in[P]$, by symmetry, we have

$$ $\mathop{\mathbb{P\/}}\left[\frac{|v_{p}|}{\max_{q\in[P]\setminus\{p\}}|v_{q}|} \geq\frac{1+2c}{1+c}\right]\geq\frac{1}{5\pi(1+2c)}\frac{1}{P^{8c^{2}}\sqrt{ \log P}}.$ $$

Now, define the event $G_{p}$ as

$$ $G_{p}=\left\{\exists i\in[m],\,\frac{|v_{i,p}|}{\max_{q\in[P]\setminus\{p\}}|v _{i,q}|}\geq\frac{1+2c}{1+c}\right\}.$ $$

Then, we compute

$$ $\displaystyle\mathop{\mathbb{P\/}}[G_{p}]$ $\displaystyle\geq 1-\left(\mathop{\mathbb{P\/}}\left[\frac{|v_{p}|}{\max_{q\in [P]\setminus\{p\}}|v_{q}|}<\frac{1+2c}{1+c}\right]\right)^{m}$ $\displaystyle\geq 1-\left(1-\frac{1}{5\pi(1+2c)}\frac{1}{P^{8c^{2}}\sqrt{\log P }}\right)^{m}$ $\displaystyle\geq 1-\exp\left(-\frac{1}{5\pi(1+2c)}\frac{m}{P^{8c^{2}}\sqrt{ \log P}}\right).$ $$

By union bound, we have

$$ $\mathop{\mathbb{P\/}}\left[\bigwedge_{p=1}^{P}G_{p}\right]\geq 1-\exp\left( \log P-\frac{1}{5\pi(1+2c)}\frac{m}{P^{8c^{2}}\sqrt{\log P}}\right).$ $$

Let $\delta_{\mathop{\mathbb{P\/}}}\in(0,1)$ be given. Choose

$$ $m\geq 400cP^{8c^{2}}\sqrt{\log P}\log\left(P\vee\frac{1}{\delta_{\mathop{ \mathbb{P\/}}}}\right).$ $$

Then, the above becomes $\mathop{\mathbb{P\/}}[\bigwedge_{p=1}^{P}G_{p}]\geq 1-\delta_{\mathop{\mathbb{
P\/}}}$.
∎

###### Lemma B.3 (Typical structure at initialization) .

Let $\delta_{\mathop{\mathbb{P\/}}}\in(e^{-\log^{C}d},1)$ be given. Suppose that $\{{\bm{v}}_{k}\}_{k=1}^{m}\sim\mathrm{Unif}(\mathbb{S}^{d-1})$ independently
with

$$ $m=400P^{8}\log^{1.5}\left(P\vee 1/\delta_{\mathop{\mathbb{P\/}}}\right).$ $$

Then, with probability at least $1-3\delta_{\mathop{\mathbb{P\/}}}$, we have

$$ $\displaystyle\forall p\in[P]\,\exists i\in[m]\quad\text{such that}\quad\frac{| v_{i,p}|}{\max_{q\in[P]\setminus\{p\}}|v_{i,q}|}\geq\frac{3}{2},$ $\displaystyle\forall i\in[m],\,\quad\left\|{\bm{v}}_{i}\right\|_{\infty}\leq \frac{20\sqrt{\log(P/\delta_{\mathop{\mathbb{P\/}}})}}{\sqrt{d}},$ $\displaystyle\forall i\in[m],\,\quad\frac{\sqrt{P}}{3\sqrt{d}}\leq\frac{\left \|{\bm{v}}_{\leq P}\right\|}{\left\|{\bm{v}}\right\|}\leq\frac{3\sqrt{P}}{ \sqrt{d}}.$ $$

###### Proof.

The first two bounds comes directly from Lemma B.1 and
Lemma B.2. By Lemma F.1,
we have

$$ $\displaystyle\mathop{\mathbb{P\/}}\left(\left|\left\|\bm{Z}\right\|-\mathop{ \mathbb{E\/}}\left\|\bm{Z}\right\|\right|\geq\sqrt{d}/2\right)\leq 2e^{-d/8},$ $\displaystyle\mathop{\mathbb{P\/}}\left(\left|\left\|\bm{Z}_{\leq P}\right\|- \mathop{\mathbb{E\/}}\left\|\bm{Z}_{\leq P}\right\|\right|\geq\sqrt{P}/2\right )\leq 2e^{-P/8}.$ $$

As a result, for any ${\bm{v}}\sim\mathrm{Unif}(\mathbb{S}^{d-1})$, we have with probability at least $1-4e^{-P/8}$ that

$$ $\frac{\left\|{\bm{v}}_{\leq P}\right\|}{\left\|{\bm{v}}\right\|}\overset{d}{=} \frac{\left\|\bm{Z}_{\leq P}\right\|}{\left\|\bm{Z}\right\|}=\frac{\mathop{ \mathbb{E\/}}\left\|\bm{Z}_{\leq P}\right\|\pm\sqrt{P}/2}{\mathop{\mathbb{E\/} }\left\|\bm{Z}\right\|\pm\sqrt{d}/2}=[1/3,3]\times\sqrt{\frac{P}{d}}.$ $$

Since we assume $P\geq\log^{C^{\prime}}d$ for a large $C^{\prime}$, we have $4e^{-P/8}\leq\delta_{\mathop{\mathbb{P\/}}}/m$. This gives the third
bound.
∎

###### Lemma B.1 (Largest coordinate) .

###### Proof.

###### Lemma B.2 (Existence of good neurons) .

###### Remark .

###### Proof.

###### Lemma B.3 (Typical structure at initialization) .

###### Proof.

## Appendix C Stage 1: recovery of the subspace and directions

In this section, we consider the stage where the second layer is fixed to be a small value and the first layer is
trained using online spherical SGD. Let ${\bm{v}}$ be an arbitrary first-layer neuron. By Lemma 2.2, we can write its update rule as(^5^55See the remark following Lemma 2.2 for the meaning of an arbitrary first-layer neuron ${\bm{v}}$. Also recall that we assume w.l.o.g. that ${\bm{v}}^{*}_{k}=\bm{e}_{k}$.)

$$ $\hat{{\bm{v}}}_{t+1}={\bm{v}}_{t}+\frac{\eta}{a_{0}}\left(\tilde{\nabla}_{{\bm {v}}}\mathcal{L}+a_{0}\bm{Z}_{t+1}\right),\quad{\bm{v}}_{t+1}=\frac{\hat{\bm{v }}_{t+1}}{\left\|\hat{\bm{v}}_{t+1}\right\|},$ $$

where $\bm{Z}_{t+1}=a_{0}^{-1}(\bm{I}-{\bm{v}}{\bm{v}}^{\top})(\nabla_{{\bm{v}}}l({
\bm{x}})-\nabla_{{\bm{v}}}\mathcal{L})$ and

$$ $\displaystyle-\tilde{\nabla}_{{\bm{v}}}\mathcal{L}$ $\displaystyle=-(\bm{I}-{\bm{v}}{\bm{v}}^{\top})\nabla_{{\bm{v}}}\mathcal{L}$ $\displaystyle=2a_{0}\sum_{k=1}^{P}v_{k}(\bm{I}-{\bm{v}}{\bm{v}}^{\top})\bm{e}_ {k}+2La_{0}\sum_{k=1}^{P}v_{k}^{2L-1}(\bm{I}-{\bm{v}}{\bm{v}}^{\top})\bm{e}_{k }\pm_{2}2Lma_{0}^{2}.$ $$

In particular, for each $k\in[d]$, we have(^6^66We will often drop the subscript $t$ when it is clear from the context.)

$$ $\hat{v}_{t+1,k}=v_{t,k}+\eta\left(\mathbbm{1}\{k\leq P\}\left(2+2Lv_{k}^{2L-2} \right)-\rho\right)v_{k}+\eta Z_{t+1,k}\pm 2\eta Lma_{0},$ $$

where

$$ $\rho:=2\sum_{i=1}^{P}v_{i}^{2}+2L\sum_{i=1}^{P}v_{i}^{2L}=2\left\|{\bm{v}}_{ \leq P}\right\|^{2}+2L\left\|{\bm{v}}_{\leq P}\right\|_{2L}^{2L}.$ (9) $$

In addition, we have the following lemma on the dynamics of $v_{k}^{2}$. The proof is routine calculation and is deferred
to the end of this section.

###### Lemma C.1 (Dynamics of $v_{k}^{2}$ ) .

For any first-layer neuron ${\bm{v}}$ and $k\in[d]$, we have

$$ $\displaystyle\hat{v}_{t+1,k}^{2}$ $\displaystyle=\left(1+2\eta\left(\mathbbm{1}\{k\leq P\}\left(2+2Lv_{k}^{2L-2} \right)-\rho\right)\right)v_{k}^{2}+2\eta v_{k}Z_{k}$ $\displaystyle\qquad\pm 300L^{3}\eta ma_{0}\pm 300L^{3}\eta^{2}\left(1\vee Z_{k }^{2}\right).$ $$

To proceed, we split Stage 1 into two substages. In Stage 1.1, we rely on the second-order terms to learn the relevant
subspace. We will also show that the gap between largest and second-largest coordinates, which can be guaranteed
with certain probability at initialization, is preserved throughout Stage 1.1. These give Stage 1.2 a nice starting
point. Then, we show that in Stage 1.2, online spherical SGD can recover the directions using the $2L$-th order terms.

###### Lemma C.1 (Dynamics of v k 2 superscript subscript 𝑣 𝑘 2 v_{k}^{2} italic_v start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT ) .

### C.1 Stage 1.1: recovery of the subspace and preservation of the gap

In this subsection, first we show that the ratio $\left\|{\bm{v}}_{\leq P}\right\|^{2}/\left\|{\bm{v}}_{>P}\right\|^{2}$ will grow from $\Omega(P/d)$ to
$\Theta(1)$ within $\tilde{O}(dP)$ iterations and during this phase. We will rely on the second-order terms and bound
the influence of higher-order terms. This leads to the desired complexity. The next goal to show the initial randomness
is preserved. In our case, we only to the gap between the largest and the second-largest coordinate to be preserved.
This ensures that the neurons will not collapse to one single direction. Formally, we have the following lemma.

###### Lemma C.2 (Stage 1.1) .

Let ${\bm{v}}\in\mathbb{S}^{d-1}$ be an arbitrary first-layer neuron satisfying
$\left\|{\bm{v}}\right\|_{\infty}\leq\log^{2}d/(2d)$ and $\left\|{\bm{v}}_{\leq P}\right\|^{2}/\left\|{\bm{v}}_{>P}\right\|^{2}\geq 0.1P/d$ at initialization.
Let $\delta_{\mathop{\mathbb{P\/}}}\in(e^{-\log^{C}d},1)$ be given. Suppose that we choose

$$ $ma_{0}\lesssim_{L}\frac{1}{d\log^{3}d}\quad\text{and}\quad\eta\lesssim_{L} \frac{\delta_{\mathop{\mathbb{P\/}}}}{dP^{2}\log^{4L+1}(d/\delta_{\mathop{ \mathbb{P\/}}})}=\tilde{\Theta}_{L}\left(\frac{\delta_{\mathop{\mathbb{P\/}}}} {dP^{2}}\right).$ $$

Then, with probability at least $1-O(\delta_{\mathop{\mathbb{P\/}}})$, we have

$$ $\frac{\left\|{\bm{v}}_{\leq P}\right\|^{2}}{\left\|{\bm{v}}_{<P}\right\|^{2}} \geq 1\quad\text{within $T=\frac{1+o(1)}{4\eta}\log\left(\frac{d}{P}\right)= \tilde{\Theta}(dP^{2})$ iterations}.$ $$

Moreover, if at initialization, $v_{p}^{2}$ is the largest among $\{v_{k}^{2}\}_{k\in[P]}$ and is $1.5$ times larger
than the second-largest $\{v_{k}^{2}\}_{k\in[P]}$, then at the end of Stage 1.1, it is still $1.25$ times larger
than the second-largest $\{v_{k}^{2}\}_{k\in[P]}$.

###### Remark .

To make the above result hold uniformly over all $m=\operatorname{poly}(P)$ neurons, it suffices to replace $\delta_{\mathop{\mathbb{P\/}}}$ with
$\delta_{\mathop{\mathbb{P\/}}}/m$. In addition, by Lemma B.3, the hypotheses of this lemma
hold with high probability at initialization.
∎

###### Proof.

It suffices to combine Lemma C.4, Lemma C.5 and Lemma C.6.
∎

To prove this lemma, we will use stochastic induction (cf. Section F.2), in particular,
Lemma F.6, Lemma F.8, and
Lemma F.10. For example, to analyze the dynamics of
$\left\|{\bm{v}}_{\leq P}\right\|^{2}/\left\|{\bm{v}}_{>P}\right\|^{2}$, it suffices to write down the update rule of
$\left\|{\bm{v}}_{\leq P}\right\|^{2}/\left\|{\bm{v}}_{>P}\right\|^{2}$ and decompose it into a signal growth term, a higher-order error term, and
a martingale difference term as in Lemma F.6. Then, we bound the higher-order
error terms, and estimate the covariance of the martingale difference terms, assuming the induction hypotheses.

The induction hypotheses we will maintain in this substage are the following:

$$ $\displaystyle\frac{\left\|{\bm{v}}_{t,\leq P}\right\|^{2}}{\left\|{\bm{v}}_{t, >P}\right\|^{2}}=\Theta(1)(1+4\eta)^{t}\frac{\left\|{\bm{v}}_{0,\leq P}\right \|^{2}}{\left\|{\bm{v}}_{0,>P}\right\|^{2}},\quad v_{p}^{2}\leq\frac{\log^{2}d }{P}.$ $$

They are established in Lemma C.4, Lemma C.9
and Lemma C.8.

###### Lemma C.2 (Stage 1.1) .

###### Remark .

###### Proof.

#### C.1.1 Learning the subspace

Now, we derive formulas for the dynamics of the ratio $\left\|{\bm{v}}_{\leq P}\right\|^{2}/\left\|{\bm{v}}_{>P}\right\|^{2}$. Since we will use
Lemma F.6 to analyze it, the goal here is separate the signal terms,
martingale difference terms, and higher-order error terms.

###### Lemma C.3 (Dynamics of the norm ratio) .

Assume the induction hypotheses.
Let ${\bm{v}}$ be an arbitrary first-layer neuron. For any $t\leq T$, we have

$$ $\displaystyle\frac{\left\|{\bm{v}}_{t+1,\leq P}\right\|^{2}}{\left\|{\bm{v}}_{ t+1,>P}\right\|^{2}}$ $\displaystyle=\frac{\left\|{\bm{v}}_{\leq P}\right\|^{2}}{\left\|{\bm{v}}_{>P} \right\|^{2}}\left(1+4\eta+\varepsilon_{v}\right)+\xi_{t+1}$ $\displaystyle\quad-\frac{\left(1+4\eta-2\eta\rho+\varepsilon_{v}\right)\left\| {\bm{v}}_{\leq P}\right\|^{2}}{\left(1-2\eta\rho\right)\left\|{\bm{v}}_{>P} \right\|^{2}}\frac{2\eta\left\langle{\bm{v}}_{>P},\bm{Z}_{>P}\right\rangle}{ \left(1-2\eta\rho\right)\left\|{\bm{v}}_{>P}\right\|^{2}}+\frac{2\eta\left \langle{\bm{v}}_{\leq P},\bm{Z}_{\leq P}\right\rangle}{\left(1-2\eta\rho\right )\left\|{\bm{v}}_{>P}\right\|^{2}},$ $$

where $\varepsilon_{v}:=4L\eta\left\|{\bm{v}}_{\leq P}\right\|_{2L}^{2L}/\left\|{\bm{
v}}_{\leq P}\right\|^{2}$ and
for any $\delta_{\mathop{\mathbb{P\/}}}\in(0,1)$, we have with probability at least $1-\delta_{\mathop{\mathbb{P\/}}}$, that

$$ $|\xi_{t+1}|\leq C_{L}(1+4\eta)^{t}\eta P\left(ma_{0}\vee\eta P^{3}\log^{4L} \left(\frac{1}{\delta_{\mathop{\mathbb{P\/}}}}\right)\right),$ $$

where $C_{L}>0$ is a constant that can depend on $L$.

###### Proof.

Recall from Lemma C.1 that

$$ $\displaystyle\hat{v}_{t+1,k}^{2}$ $\displaystyle=\left(1+2\eta\left(\mathbbm{1}\{k\leq P\}\left(2+2Lv_{k}^{2L-2} \right)-\rho\right)\right)v_{k}^{2}+2\eta v_{k}Z_{k}$ $\displaystyle\qquad\pm 300L^{3}\eta ma_{0}\pm 300L^{3}\eta^{2}\left(1\vee Z_{k }^{2}\right).$ $$

Hence, for the norms, we have
(the higher order terms are changed; additional $P$, $d$ factors)

$$ $\displaystyle\left\|\hat{{\bm{v}}}_{\leq P}\right\|^{2}$ $\displaystyle=\left(1+2\eta\left(2-\rho\right)\right)\left\|{\bm{v}}_{\leq P} \right\|^{2}+4L\eta\left\|{\bm{v}}_{\leq P}\right\|_{2L}^{2L}+2\eta\left \langle{\bm{v}}_{\leq P},\bm{Z}_{\leq P}\right\rangle$ $\displaystyle\qquad\underbrace{\pm 300L^{3}P\eta ma_{0}\pm 300L^{3}\eta^{2} \left(P\vee\left\|\bm{Z}_{\leq P}\right\|^{2}\right)}_{=:\;\xi_{\leq P,t}},$ $\displaystyle\hat{v}_{t+1,k}^{2}$ $\displaystyle=\left(1-2\eta\rho\right)\left\|{\bm{v}}_{>P}\right\|^{2}+2\eta \left\langle{\bm{v}}_{>P},\bm{Z}_{>P}\right\rangle$ $\displaystyle\qquad\underbrace{\pm 300L^{3}d\eta ma_{0}\pm 300L^{3}\eta^{2} \left(d\vee\left\|\bm{Z}_{\geq P}^{2}\right\|\right)}_{=:\;\xi_{>P,t}}.$ $$

For notational simplicity, put $\varepsilon_{v}=4L\eta\left\|{\bm{v}}_{\leq P}\right\|_{2L}^{2L}/\left\|{\bm{v
}}_{\leq P}\right\|^{2}$.
Note that $\left\|{\bm{v}}_{\leq P}\right\|/\left\|{\bm{v}}_{>P}\right\|=\left\|\hat{\bm{
v}}_{\leq P}\right\|/\left\|\hat{\bm{v}}_{>P}\right\|$.
Thus, we have

$$ $\displaystyle\frac{\left\|{\bm{v}}_{t+1,\leq P}\right\|^{2}}{\left\|{\bm{v}}_{ t+1,>P}\right\|^{2}}$ $\displaystyle=\frac{\left(1+2\eta\left(2-\rho\right)+\varepsilon_{v}\right) \left\|{\bm{v}}_{\leq P}\right\|^{2}+2\eta\left\langle{\bm{v}}_{\leq P},\bm{Z} _{\leq P}\right\rangle+\xi_{\leq P}}{\left(1-2\eta\rho\right)\left\|{\bm{v}}_{ >P}\right\|^{2}+2\eta\left\langle{\bm{v}}_{>P},\bm{Z}_{>P}\right\rangle+\xi_{> P}}$ $\displaystyle=\frac{\left(1+4\eta-2\eta\rho+\varepsilon_{v}\right)\left\|{\bm{ v}}_{\leq P}\right\|^{2}}{\left(1-2\eta\rho\right)\left\|{\bm{v}}_{>P}\right\| ^{2}}\left(1-\frac{2\eta\left\langle{\bm{v}}_{>P},\bm{Z}_{>P}\right\rangle}{ \left\|\hat{{\bm{v}}}_{t+1,>P}\right\|^{2}}-\frac{\xi_{>P}}{\left\|\hat{{\bm{v }}}_{t+1,>P}\right\|^{2}}\right)$ $\displaystyle\qquad+\frac{2\eta\left\langle{\bm{v}}_{\leq P},\bm{Z}_{\leq P} \right\rangle}{\left\|\hat{{\bm{v}}}_{t+1,>P}\right\|^{2}}+\frac{\xi_{\leq P}} {\left\|\hat{{\bm{v}}}_{t+1,>P}\right\|^{2}}$ $\displaystyle=\frac{\left(1+4\eta-2\eta\rho+\varepsilon_{v}\right)\left\|{\bm{ v}}_{\leq P}\right\|^{2}}{\left(1-2\eta\rho\right)\left\|{\bm{v}}_{>P}\right\| ^{2}}$ $\displaystyle\qquad-\frac{\left(1+4\eta-2\eta\rho+\varepsilon_{v}\right)\left \|{\bm{v}}_{\leq P}\right\|^{2}}{\left(1-2\eta\rho\right)\left\|{\bm{v}}_{>P} \right\|^{2}}\frac{2\eta\left\langle{\bm{v}}_{>P},\bm{Z}_{>P}\right\rangle}{ \left\|\hat{{\bm{v}}}_{t+1,>P}\right\|^{2}}+\frac{2\eta\left\langle{\bm{v}}_{ \leq P},\bm{Z}_{\leq P}\right\rangle}{\left\|\hat{{\bm{v}}}_{t+1,>P}\right\|^{ 2}}$ $\displaystyle\qquad-\frac{\left(1+4\eta-2\eta\rho+\varepsilon_{v}\right)\left \|{\bm{v}}_{\leq P}\right\|^{2}}{\left(1-2\eta\rho\right)\left\|{\bm{v}}_{>P} \right\|^{2}}\frac{\xi_{>P}}{\left\|\hat{{\bm{v}}}_{t+1,>P}\right\|^{2}}+\frac {\xi_{\leq P}}{\left\|\hat{{\bm{v}}}_{t+1,>P}\right\|^{2}}.$ $$

Note that up to some higher-order terms, the first line contains the signal terms and the second line contains the
martingale difference terms. Now, our goal is to factor out those higher-order terms.
For the first line, first recall from (13) that $\rho\leq 4L$, and then we use the fact that

$$ $\frac{1}{1+z}=1-z\pm 2z^{2},\quad\forall|z|\leq 1/2,$ (10) $$

to obtain

$$ $\displaystyle\frac{\left(1+4\eta-2\eta\rho+\varepsilon_{v}\right)\left\|{\bm{v }}_{\leq P}\right\|^{2}}{\left(1-2\eta\rho\right)\left\|{\bm{v}}_{>P}\right\|^ {2}}$ $\displaystyle=\frac{\left\|{\bm{v}}_{\leq P}\right\|^{2}}{\left\|{\bm{v}}_{>P} \right\|^{2}}\left(1+4\eta-2\eta\rho+\varepsilon_{v}\right)\left(1+2\eta\rho \pm 64L^{2}\eta^{2}\right)$ $\displaystyle=\frac{\left\|{\bm{v}}_{\leq P}\right\|^{2}}{\left\|{\bm{v}}_{>P} \right\|^{2}}\left(1+4\eta+\varepsilon_{v}\pm 2000L^{3}\eta^{2}\right).$ $$

Similarly, for the second line, we write

$$ $\frac{1}{\left\|\hat{{\bm{v}}}_{t+1,>P}\right\|^{2}}=\frac{1}{\left(1-2\eta \rho\right)\left\|{\bm{v}}_{>P}\right\|^{2}}\left(1-\frac{2\eta\left\langle{ \bm{v}}_{>P},\bm{Z}_{>P}\right\rangle+\xi_{>P}}{\left\|\hat{{\bm{v}}}_{t+1,>P} \right\|^{2}}\right).$ $$

By the tail bounds in Lemma 2.2 and the union bound, for any $\delta_{\mathop{\mathbb{P\/}}}\in(0,1)$,
we have

$$ $\displaystyle\left|\left\langle\overline{{\bm{v}}_{>P}},\bm{Z}_{>P}\right \rangle\right|\leq C_{L}^{2L}P\log^{2L}\left(\frac{C_{L}}{\delta_{\mathop{ \mathbb{P\/}}}}\right),\quad\left|\left\langle\overline{{\bm{v}}_{\leq P}},\bm {Z}_{\leq P}\right\rangle\right|\leq C_{L}^{2L}P\log^{2L}\left(\frac{C_{L}}{ \delta_{\mathop{\mathbb{P\/}}}}\right),$ $\displaystyle|Z_{k}|\leq C_{L}^{2L}P\log^{2L}\left(\frac{C_{L}d}{\delta_{ \mathop{\mathbb{P\/}}}}\right),\quad\forall k\in[d],$ $$

with probability at least $1-2\delta_{\mathop{\mathbb{P\/}}}$. In particular, note that the second bound also implies, with at
least the same probability, we have

$$ $\displaystyle|\xi_{\leq P}|$ $\displaystyle\leq 600L^{3}\eta P\left(ma_{0}\vee\eta C_{L}^{4L}P^{2}\log^{4L} \left(\frac{C_{L}d}{\delta_{\mathop{\mathbb{P\/}}}}\right)\right),$ $\displaystyle|\xi_{>P}|$ $\displaystyle\leq 600L^{3}\eta d\left(ma_{0}\vee\eta C_{L}^{4L}P^{2}\log^{4L} \left(\frac{C_{L}d}{\delta_{\mathop{\mathbb{P\/}}}}\right)\right).$ $$

By our definition of Stage 1.1, we have $\left\|\hat{{\bm{v}}}_{t+1,>P}\right\|^{2}\geq 1/2$. Therefore, with probability at least
$1-2\delta_{\mathop{\mathbb{P\/}}}$, we have

$$ $\frac{1}{\left\|\hat{{\bm{v}}}_{t+1,>P}\right\|^{2}}=\frac{1}{\left(1-2\eta \rho\right)\left\|{\bm{v}}_{>P}\right\|^{2}}\left(1\pm C_{L}^{\prime}\eta P \log^{2L}\left(\frac{1}{\delta_{\mathop{\mathbb{P\/}}}}\right)\right),$ $$

for some constant $C_{L}^{\prime}>0$ that can depend on $L$.
Thus, for the ratio of the norms, we have

$$ $\displaystyle\frac{\left\|{\bm{v}}_{t+1,\leq P}\right\|^{2}}{\left\|{\bm{v}}_{ t+1,>P}\right\|^{2}}$ $\displaystyle=\frac{\left\|{\bm{v}}_{\leq P}\right\|^{2}}{\left\|{\bm{v}}_{>P} \right\|^{2}}\left(1+4\eta+\varepsilon_{v}\pm 2000L^{3}\eta^{2}\right)$ $\displaystyle\qquad-\frac{\left(1+4\eta-2\eta\rho+\varepsilon_{v}\right)\left \|{\bm{v}}_{\leq P}\right\|^{2}}{\left(1-2\eta\rho\right)\left\|{\bm{v}}_{>P} \right\|^{2}}\frac{2\eta\left\langle{\bm{v}}_{>P},\bm{Z}_{>P}\right\rangle}{ \left(1-2\eta\rho\right)\left\|{\bm{v}}_{>P}\right\|^{2}}\left(1\pm C_{L}^{ \prime}\eta P\log^{2L}\left(\frac{1}{\delta_{\mathop{\mathbb{P\/}}}}\right)\right)$ $\displaystyle\qquad+\frac{2\eta\left\langle{\bm{v}}_{\leq P},\bm{Z}_{\leq P} \right\rangle}{\left(1-2\eta\rho\right)\left\|{\bm{v}}_{>P}\right\|^{2}}\left( 1\pm C_{L}^{\prime}\eta P\log^{2L}\left(\frac{1}{\delta_{\mathop{\mathbb{P\/}} }}\right)\right)$ $\displaystyle\qquad-\frac{\left(1+4\eta-2\eta\rho+\varepsilon_{v}\right)\left \|{\bm{v}}_{\leq P}\right\|^{2}}{\left(1-2\eta\rho\right)\left\|{\bm{v}}_{>P} \right\|^{2}}\frac{\xi_{>P}}{\left\|\hat{{\bm{v}}}_{t+1,>P}\right\|^{2}}+\frac {\xi_{\leq P}}{\left\|\hat{{\bm{v}}}_{t+1,>P}\right\|^{2}}.$ $$

Collect the higher-order terms into $\xi_{t+1}$, so that the above becomes

$$ $\displaystyle\frac{\left\|{\bm{v}}_{t+1,\leq P}\right\|^{2}}{\left\|{\bm{v}}_{ t+1,>P}\right\|^{2}}$ $\displaystyle=\frac{\left\|{\bm{v}}_{\leq P}\right\|^{2}}{\left\|{\bm{v}}_{>P} \right\|^{2}}\left(1+4\eta+\varepsilon_{v}\right)+\xi_{t+1}$ $\displaystyle\quad-\frac{\left(1+4\eta-2\eta\rho+\varepsilon_{v}\right)\left\| {\bm{v}}_{\leq P}\right\|^{2}}{\left(1-2\eta\rho\right)\left\|{\bm{v}}_{>P} \right\|^{2}}\frac{2\eta\left\langle{\bm{v}}_{>P},\bm{Z}_{>P}\right\rangle}{ \left(1-2\eta\rho\right)\left\|{\bm{v}}_{>P}\right\|^{2}}+\frac{2\eta\left \langle{\bm{v}}_{\leq P},\bm{Z}_{\leq P}\right\rangle}{\left(1-2\eta\rho\right )\left\|{\bm{v}}_{>P}\right\|^{2}}.$ $$

For the higher-order terms, we have with probability at least $1-O(\delta_{\mathop{\mathbb{P\/}}})$

$$ $\displaystyle|\xi_{t+1}|$ $\displaystyle\lesssim_{L}\frac{\left\|{\bm{v}}_{\leq P}\right\|^{2}}{\left\|{ \bm{v}}_{>P}\right\|^{2}}\eta^{2}+\frac{\left\|{\bm{v}}_{\leq P}\right\|^{2}}{ \left\|{\bm{v}}_{>P}\right\|^{2}}\frac{\eta|\left\langle{\bm{v}}_{>P},\bm{Z}_{ >P}\right\rangle|}{\left\|{\bm{v}}_{>P}\right\|^{2}}\eta P\log^{2L}\left(\frac {1}{\delta_{\mathop{\mathbb{P\/}}}}\right)$ $\displaystyle\qquad+\frac{\eta|\left\langle{\bm{v}}_{\leq P},\bm{Z}_{\leq P} \right\rangle|}{\left\|{\bm{v}}_{>P}\right\|^{2}}\eta P\log^{2L}\left(\frac{1} {\delta_{\mathop{\mathbb{P\/}}}}\right)+\frac{\left\|{\bm{v}}_{\leq P}\right\| ^{2}}{\left\|{\bm{v}}_{>P}\right\|^{2}}\frac{|\xi_{>P}|}{\left\|{\bm{v}}_{>P} \right\|^{2}}+\frac{|\xi_{\leq P}|}{\left\|{\bm{v}}_{>P}\right\|^{2}}$ $\displaystyle\lesssim_{L}\frac{\left\|{\bm{v}}_{\leq P}\right\|^{2}}{\left\|{ \bm{v}}_{>P}\right\|^{2}}\eta^{2}+\left(\frac{\left\|{\bm{v}}_{\leq P}\right\| ^{2}}{\left\|{\bm{v}}_{>P}\right\|^{3}}+\frac{\left\|{\bm{v}}_{\leq P}\right\| }{\left\|{\bm{v}}_{>P}\right\|^{2}}\right)\eta^{2}P^{2}\log^{4L}\left(\frac{1} {\delta_{\mathop{\mathbb{P\/}}}}\right)$ $\displaystyle\qquad+\left(\frac{d\left\|{\bm{v}}_{\leq P}\right\|^{2}}{\left\| {\bm{v}}_{>P}\right\|^{4}}+\frac{P}{\left\|{\bm{v}}_{>P}\right\|^{2}}\right) \eta\left(ma_{0}\vee\eta P^{2}\log^{4L}\left(\frac{d}{\delta_{\mathop{\mathbb{ P\/}}}}\right)\right)$ $\displaystyle\lesssim_{L}(1+4\eta)^{t}\eta P\left(ma_{0}\vee\eta P^{2}\log^{4L }\left(\frac{d}{\delta_{\mathop{\mathbb{P\/}}}}\right)\right),$ $$

where we use the induction hypothesis $\left\|{\bm{v}}_{\leq P}\right\|^{2}/\left\|{\bm{v}}_{>P}\right\|^{2}=\Theta
\left((1+4\eta)^{t}P/d\right)$
to handle the $d\left\|{\bm{v}}_{\leq P}\right\|^{2}/\left\|{\bm{v}}_{>P}\right\|^{4}$ factor in the last line.
∎

With the above formula, we can now use Lemma F.6 to analyze the dynamics of
ratio of the norms.

###### Lemma C.4 (Learning the subspace) .

Let ${\bm{v}}$ be an arbitrary fixed first-layer neuron. Suppose that

$$ $ma_{0}\lesssim_{L}\frac{1}{d\log d}\quad\text{and}\quad\eta\lesssim_{L}\frac{ \delta_{\mathop{\mathbb{P\/}}}}{dP^{2}\log^{4L+1}\left(d/\delta_{\mathop{ \mathbb{P\/}}}\right)}=\tilde{\Theta}_{L}\left(\frac{\delta_{\mathop{\mathbb{P \/}}}}{dP^{2}}\right),$ $$

Then, throughout Stage 1.1, we have

$$ $\frac{(1+4\eta)^{t}}{2}\frac{\left\|{\bm{v}}_{0,\leq P}\right\|^{2}}{\left\|{ \bm{v}}_{0,>P}\right\|^{2}}\leq\frac{\left\|{\bm{v}}_{\leq P}\right\|^{2}}{ \left\|{\bm{v}}_{>P}\right\|^{2}}\leq\frac{3(1+4\eta)^{t}}{2}\frac{\left\|{\bm {v}}_{0,\leq P}\right\|^{2}}{\left\|{\bm{v}}_{0,>P}\right\|^{2}},$ $$

and Stage 1.1 takes at most $(1+o(1))(4\eta)^{-1}\log\left(d/P\right)=\tilde{O}_{L}\left(dP^{2}/\delta_{
\mathop{\mathbb{P\/}}}\right)$ iterations.
To obtain estimates that uniformly hold for all neurons, it suffices to replace $\delta_{\mathop{\mathbb{P\/}}}$ with $\delta_{\mathop{\mathbb{P\/}}}/m$.

###### Proof.

By Lemma C.3, we have

$$ $\displaystyle\frac{\left\|{\bm{v}}_{t+1,\leq P}\right\|^{2}}{\left\|{\bm{v}}_{ t+1,>P}\right\|^{2}}$ $\displaystyle=\frac{\left\|{\bm{v}}_{\leq P}\right\|^{2}}{\left\|{\bm{v}}_{>P} \right\|^{2}}\left(1+4\eta+\varepsilon_{v}\right)+\xi_{t+1}$ $\displaystyle\quad\underbrace{-\frac{\left(1+4\eta-2\eta\rho+\varepsilon_{v} \right)\left\|{\bm{v}}_{\leq P}\right\|^{2}}{\left(1-2\eta\rho\right)\left\|{ \bm{v}}_{>P}\right\|^{2}}\frac{2\eta\left\langle{\bm{v}}_{>P},\bm{Z}_{>P} \right\rangle}{\left(1-2\eta\rho\right)\left\|{\bm{v}}_{>P}\right\|^{2}}}_{=: \;H^{(1)}_{t+1}}+\underbrace{\frac{2\eta\left\langle{\bm{v}}_{\leq P},\bm{Z}_{ \leq P}\right\rangle}{\left(1-2\eta\rho\right)\left\|{\bm{v}}_{>P}\right\|^{2} }}_{=:\;H^{(2)}_{t+1}},$ $$

where $\varepsilon_{v}:=4L\eta\left\|{\bm{v}}_{\leq P}\right\|_{2L}^{2L}/\left\|{\bm{
v}}_{\leq P}\right\|^{2}$ and
for any $\delta_{\mathop{\mathbb{P\/}}}\in(0,1)$, we have with probability at least $1-\delta_{\mathop{\mathbb{P\/}}}/T$, that

$$ $|\xi_{t+1}|\leq C_{L}(1+4\eta)^{t}\eta P\left(ma_{0}\vee\eta P^{2}\log^{4L} \left(\frac{T}{\delta_{\mathop{\mathbb{P\/}}}}\right)\right),$ $$

where $C_{L}>0$ is a constant that can depend on $L$.
By our induction hypothesis $v_{p}^{2}\leq\log^{2}d/P$, we

$$ $\varepsilon_{v}=\frac{4L\eta}{\left\|{\bm{v}}_{\leq P}\right\|^{2}}\sum_{p=1}^ {P}v_{p}^{2L}\leq\frac{4L\eta}{\left\|{\bm{v}}_{\leq P}\right\|^{2}}\left\|{ \bm{v}}_{\leq P}\right\|_{\infty}^{2L-2}\sum_{p=1}^{P}v_{p}^{2}\leq\eta\frac{4 L\log^{2L-2}(d)}{P^{L-1}}=:\eta\delta_{v}.$ $$

In particular, note that $\delta_{v}$ does not depend on $t$ and is $o(1)$.
For the martingale difference terms, by Lemma 2.2, we have

$$ $\displaystyle\mathop{\mathbb{E\/}}\left[(H^{(1)}_{t+1})^{2}\mid\mathcal{F}_{t}\right]$ $\displaystyle\lesssim_{L}\eta^{2}\frac{\left\|{\bm{v}}_{\leq P}\right\|^{4}}{ \left\|{\bm{v}}_{>P}\right\|^{6}}\mathop{\mathbb{E\/}}\left[\left\langle \overline{{\bm{v}}_{>P}},\bm{Z}_{>P}\right\rangle^{2}\mid\mathcal{F}_{t}\right ]\lesssim_{L}\eta^{2}P^{2}\frac{\left\|{\bm{v}}_{\leq P}\right\|^{4}}{\left\|{ \bm{v}}_{>P}\right\|^{4}},$ $\displaystyle\mathop{\mathbb{E\/}}\left[(H^{(2)}_{t+1})^{2}\mid\mathcal{F}_{t}\right]$ $\displaystyle\lesssim_{L}\eta^{2}\frac{\left\|{\bm{v}}_{\leq P}\right\|^{2}}{ \left\|{\bm{v}}_{>P}\right\|^{4}}\mathop{\mathbb{E\/}}\left[\left\langle \overline{{\bm{v}}_{\leq P}},\bm{Z}_{\leq P}\right\rangle^{2}\mid\mathcal{F}_{ t}\right]\lesssim_{L}\eta^{2}P^{2}\frac{\left\|{\bm{v}}_{\leq P}\right\|^{2}}{ \left\|{\bm{v}}_{>P}\right\|^{2}}.$ $$

Put $H_{t+1}:=H^{(1)}_{t+1}+H^{(2)}_{t+1}$. The above bounds imply that

$$ $\mathop{\mathbb{E\/}}\left[H^{2}_{t+1}\mid\mathcal{F}_{t}\right]\lesssim_{L} \eta^{2}P^{2}\frac{\left\|{\bm{v}}_{\leq P}\right\|^{2}}{\left\|{\bm{v}}_{>P} \right\|^{2}}\lesssim_{L}\eta^{2}P^{2}(1+4\eta)^{t}\frac{\left\|{\bm{v}}_{0, \leq P}\right\|^{2}}{\left\|{\bm{v}}_{0,>P}\right\|^{2}}\lesssim_{L}\frac{\eta ^{2}P^{3}}{d}(1+4\eta)^{t}$ $$

where the second inequality comes from our induction hypothesis.

For notational simplicity, put $X_{t}:=\left\|{\bm{v}}_{\leq P}\right\|^{2}/\left\|{\bm{v}}_{>P}\right\|^{2}$,
$x^{-}_{t}=(1+4\eta)^{t}X_{0}$ and $x^{+}_{t}=(1+4\eta(1+\delta_{v}))^{t}X_{0}$. $x^{\pm}$ will serve as the
lower and upper bounds for the deterministic counterpart of $X$, since

$$ $\displaystyle\left(1+4\eta\right)X_{t}+\xi_{t+1}+H_{t+1}\leq X_{t+1}\leq\left( 1+4\eta(1+\delta_{v})\right)X_{t}+\xi_{t+1}+H_{t+1}.$ $$

Moreover, note that for any $t\leq T$, we have

$$ $\displaystyle\frac{x^{+}_{t}}{x^{-}_{t}}=\left(\frac{1+4\eta(1+\delta_{v})}{1+ 4\eta}\right)^{t}$ $\displaystyle=\left(\left(1+4\eta(1+\delta_{v})\right)\left(1-4\eta\pm 16\eta^ {2}\right)\right)^{t}$ $\displaystyle\leq\left(1+4\eta\delta_{v}\pm 40\eta^{2}\right)^{t}$ $\displaystyle\leq\exp\left(40\eta T\left(\delta_{v}+\eta\right)\right).$ $$

Since $T\leq\log d/\eta$, the above implies

$$ $1\leq\frac{x^{+}_{t}}{x^{-}_{t}}\leq\exp\left(40\log d\left(\delta_{v}+\eta \right)\right)\leq 1+80\log d\left(\delta_{v}+\eta\right)=1+o(1),$ $$

where the last (approximate) identity holds whenever

$$ $\delta_{v}\ll\frac{1}{\log d}\quad\Leftarrow\quad\frac{4L\log^{2L-2}(d)}{P^{L- 1}}\ll\frac{1}{\log d}\quad\Leftarrow\quad P\gg(4L)^{1/(L-1)}\log^{2}d.$ $$

In particular, this implies that the (multiplicative) difference between $x^{+}_{t}$ and $x^{-}_{t}$ is small.

Now, we apply Lemma F.6 to $X_{t}$. In our case, we have

$$ $\Xi\lesssim_{L}\eta P\left(ma_{0}\vee\eta P^{2}\log^{4L}\left(\frac{T}{\delta_ {\mathop{\mathbb{P\/}}}}\right)\right),\quad\sigma_{Z}^{2}\lesssim_{L}\frac{ \eta^{2}P^{3}}{d},$ $$

$\alpha=4(1+o(1))\eta$ and $X_{0}=\Theta(P/d)$. Recall that $T\leq O(\log d/\eta)$.
Hence, to meet the conditions of Lemma F.6, it suffices to choose

$$ $\displaystyle\eta P\left(ma_{0}\vee\eta P^{2}\log^{4L}\left(\frac{T}{\delta_{ \mathop{\mathbb{P\/}}}}\right)\right)\lesssim_{L}\frac{X_{0}}{T}$ $\displaystyle\quad\Leftarrow\quad\left\{\begin{aligned} &ma_{0}\lesssim_{L} \frac{1}{d\log d},\\ &\eta\lesssim_{L}\frac{1}{dP^{2}\log^{4L}\left(T/\delta_{\mathop{\mathbb{P\/}} }\right)\log d}\end{aligned}\right.$ $\displaystyle\frac{\eta^{2}P^{3}}{d}\lesssim_{L}\frac{\delta_{\mathop{\mathbb{ P\/}}}\alpha X_{0}^{2}}{16}$ $\displaystyle\quad\Leftarrow\quad\eta\lesssim_{L}\frac{\delta_{\mathop{\mathbb {P\/}}}}{dP}.$ $$

To satisfy the above conditions, it suffices to choose

$$ $ma_{0}\lesssim_{L}\frac{1}{d\log d}\quad\text{and}\quad\eta\lesssim_{L}\frac{ \delta_{\mathop{\mathbb{P\/}}}}{dP^{2}\log^{4L+1}\left(d/\delta_{\mathop{ \mathbb{P\/}}}\right)}.$ $$

Then, by Lemma F.6,
we have, with probability at least $1-\Theta(\delta_{\mathop{\mathbb{P\/}}})$, $0.5x^{-}_{t}\leq X_{t}\leq 1.5x^{+}_{t}$. Since $x^{+}_{t}=(1+o(1))x^{-}_{t}$, this implies $0.5x_{t}\leq X_{t}\leq 2x_{t}$. To complete the proof, it suffices to note that
for $x_{t}$ to grow from $\Theta(P/d)$ to $1$, the number of iterations needed is bounded by
$(1+o(1))(4\eta)^{-1}\log\left(d/P\right)$.
∎

###### Lemma C.3 (Dynamics of the norm ratio) .

###### Proof.

###### Lemma C.4 (Learning the subspace) .

###### Proof.

#### C.1.2 Preservation of the gap

Now, we show that the gap between the largest coordinate and the second-largest coordinate can be preserved
in Stage 1.1. Let $p=\operatorname*{argmax}_{i\in[P]}v_{i}^{2}(0)$ and consider the ratio $v_{p}^{2}/v_{q}^{2}$, where $q\in[P]$ is
arbitrary. The proof is conceptually very similar to the previous one, except that we will use
Lemma F.8 instead of Lemma F.6.
However, there is still some technical subtlety that is not involved in the previous analysis. When $v_{q}^{2}$ is close
to 0 0, the dynamics of $v_{p}^{2}/v_{q}^{2}$ can be unstable, violating the conditions of Lemma F.8. Intuitively, this should not cause any fundamental issue, since we are only interested
in the square of largest and second-largest coordinates, both of which should be at least $\Omega(1/d)$ throughout
Stage 1.1. To handle this technical issue, we will partition $q\in[P]$ based on the initial value $v_{0,q}^{2}$.
When $v_{0,q}^{2}=\Omega(1/d)$, we consider the dynamics of the ratio $v_{p}^{2}/v_{q}^{2}$ directly. If
$v_{0,q}^{2}$ is small, we will use Lemma C.7 and
Lemma C.8, and bound the ratio in a more direct way.

###### Lemma C.5 (Gap between large and small coordinates) .

Consider $p,q\in[P]$. There exists a universal constant $c_{v}>0$ such that
if $v_{0,p}^{2}\geq 1/d$ and $v_{0,q}^{2}\leq c_{v}/d$, and we choose the hyperparameters according to
Lemma C.7 and Lemma C.8, then we have with probability at least $1-O(\delta_{\mathop{\mathbb{P\/}}})$, that
$v_{p}^{2}\geq 2v_{q}^{2}$ throughout Stage 1.1.

###### Proof.

By Lemma C.7, we have

$$ $v_{t,p}^{2}\geq\frac{1}{2}(1+4\eta)^{t}v_{0,p}^{2}\geq\frac{1}{2}(1+4\eta)^{t} \frac{1}{d},$ $$

with probability at least $1-O(\delta_{\mathop{\mathbb{P\/}}})$. Meanwhile, by Lemma C.8, we have

$$ $v_{t,q}^{2}\leq 2C(1+4\eta)^{t}\frac{c_{v}}{d},$ $$

with probability at least $1-O(\delta_{\mathop{\mathbb{P\/}}})$. Hence, as long as $c_{v}\leq 1/(8C)$, we have
$v_{t,q}^{2}\leq v_{t,p}^{2}/2$ throughout Stage 1 with probability at least
$1-O(\delta_{\mathop{\mathbb{P\/}}})$.
∎

###### Lemma C.6 (Gap between large coordinates) .

Consider $p,q\in[P]$ and let $c_{v}>0$ be the universal constant in the previous lemma. Suppose that
$v_{0,p}^{2}\geq v_{0,q}^{2}\geq c_{v}/d$. Let $\varepsilon_{R}\in(0,1)$ be given.
Suppose that the hyperparameters satisfy the conditions in Lemma C.7 and

$$ $ma_{0}\lesssim_{L}\frac{\varepsilon_{R}}{d\log^{3}d},\quad P\gtrsim_{L}\frac{ \log^{3}d}{\varepsilon_{R}},\quad\eta\lesssim_{L}\frac{\varepsilon_{R}\sqrt{ \delta_{\mathop{\mathbb{P\/}}}}}{dP^{2}\log^{2L+2}\left(d/\delta_{\mathop{ \mathbb{P\/}}}\right)}.$ $$

Then, we have $\left|v_{p}^{2}/v_{q}^{2}-v_{0,p}^{2}/v_{0,q}^{2}\right|\leq\varepsilon_{R}$ throughout Stage 1.1 with
probability at least $1-\Theta(\delta_{\mathop{\mathbb{P\/}}})$.

###### Proof.

First, note that by Lemma C.7, we have $v_{t,q}^{2}\geq c_{v}/(2d)$
throughout Stage 1.1 with probability at least $1-O(\delta_{\mathop{\mathbb{P\/}}})$.
Recall from Lemma C.1 that for any $k\leq P$, we have

$$ $\displaystyle\hat{v}_{t+1,k}^{2}$ $\displaystyle=\left(1+2\eta\left(2Lv_{k}^{2L-2}+2-\rho\right)\right)v_{k}^{2}+ 2\eta v_{k}Z_{k}\underbrace{\pm 300L^{3}\eta ma_{0}\pm 300L^{3}\eta^{2}\left(1 \vee Z_{k}^{2}\right)}_{=:\;\xi_{k}}.$ $$

Hence, for any $p,q\in[P]$, we have

$$ $\displaystyle\frac{v_{p,t+1}^{2}}{v_{q,t+1}^{2}}$ $\displaystyle=\frac{\left(1+2\eta\left(2Lv_{p}^{2L-2}+2-\rho\right)\right)v_{p }^{2}+2\eta v_{p}Z_{p}+\xi_{p}}{\left(1+2\eta\left(2Lv_{q}^{2L-2}+2-\rho\right )\right)v_{q}^{2}+2\eta v_{q}Z_{q}+\xi_{q}}$ $\displaystyle=\frac{v_{p}^{2}}{v_{q}^{2}}-\frac{v_{p}^{2}}{v_{q}^{2}}\frac{2 \eta v_{q}Z_{q}}{\left(1+2\eta\left(2-\rho\right)\right)v_{q}^{2}+4L\eta v_{q} ^{2L}}+\frac{2\eta v_{p}Z_{p}}{\left(1+2\eta\left(2Lv_{q}^{2L-2}+2-\rho\right) \right)v_{q}^{2}}$ $\displaystyle\qquad-\frac{2\eta v_{p}Z_{p}}{\left(1+2\eta\left(2Lv_{q}^{2L-2}+ 2-\rho\right)\right)v_{q}^{2}}\frac{2\eta v_{q}Z_{q}+\xi_{q}}{\hat{v}_{q,t+1}^ {2}}$ $\displaystyle\qquad+\frac{v_{p}^{2}}{v_{q}^{2}}\frac{2\eta v_{q}Z_{q}}{\left(1 +2\eta\left(2-\rho\right)\right)v_{q}^{2}+4L\eta v_{q}^{2L}}\frac{2\eta v_{q}Z _{q}+\xi_{q}}{\hat{v}_{q,t+1}^{2}}$ $\displaystyle\qquad+\frac{\xi_{p}+4L\eta v_{p}^{2L}}{\hat{v}_{q,t+1}^{2}}+ \frac{v_{p}^{2}}{v_{q}^{2}}\frac{4L\eta v_{q}^{2L}+\xi_{q}}{\hat{v}_{q,t+1}^{2 }}.$ $$

The first line contains the signal term and the martingale difference terms. The other three lines contain
the higher-order error terms. First, for the martingale difference terms, by our induction hypotheses
and the variance bound in Lemma 2.2, we have

$$ $\displaystyle\mathop{\mathbb{E\/}}\left[\left(\frac{v_{p}^{2}}{v_{q}^{2}}\frac {2\eta v_{q}Z_{q}}{\left(1+2\eta\left(2-\rho\right)\right)v_{q}^{2}+4L\eta v_{ q}^{2L}}\right)^{2}\;\bigg{|}\;\mathcal{F}_{t}\right]$ $\displaystyle\lesssim_{L}\eta^{2}P^{2}\frac{v_{p}^{4}}{v_{q}^{6}}\lesssim_{L} \eta^{2}dP^{2}\log^{4}d,$ $\displaystyle\mathop{\mathbb{E\/}}\left[\left(\frac{2\eta v_{p}Z_{p}}{\left(1+ 2\eta\left(2Lv_{q}^{2L-2}+2-\rho\right)\right)v_{q}^{2}}\right)^{2}\;\bigg{|} \;\mathcal{F}_{t}\right]$ $\displaystyle\lesssim_{L}\eta^{2}P^{2}\frac{v_{p}^{2}}{v_{q}^{4}}\lesssim_{L} \eta^{2}dP^{2}\log^{2}d$ $$

where we have used the induction hypotheses $v_{q}^{2}\geq\Theta(1/d)$ and $v_{p}^{2}/v_{q}^{2}=\Theta(v_{0,p}^{2}/v_{0,q}^{2})=O(\log^{2}d)$. Using the language of Lemma F.8, these imply

$$ $\sigma_{Z}^{2}\lesssim_{L}\eta^{2}dP^{2}\log^{4}d.$ (11) $$

Then, for the higher-order terms, first by the tail bounds in Lemma 2.2, we have
for any $\delta_{\mathop{\mathbb{P\/}},\xi}\in(0,1)$, that

$$ $|Z_{p}|\vee|Z_{q}|\leq C_{L}^{2L}P\log^{2L}\left(\frac{C_{L}}{\delta_{\mathop{ \mathbb{P\/}},\xi}}\right)\quad\text{with probability at least $1-2\delta_{ \mathop{\mathbb{P\/}},\xi}$.}$ $$

In particular, this implies that with at least the same probability, we have

$$ $|\xi_{p}|\vee|\xi_{q}|\lesssim_{L}\eta ma_{0}\vee\eta^{2}P^{2}\log^{4L}\left( \frac{1}{\delta_{\mathop{\mathbb{P\/}},\xi}}\right).$ $$

Suppose that $\eta\leq 1/d$. Then, we have

$$ $\displaystyle\left|\frac{\xi_{p}+4L\eta v_{p}^{2L}}{\hat{v}_{q,t+1}^{2}}+\frac {v_{p}^{2}}{v_{q}^{2}}\frac{4L\eta v_{q}^{2L}+\xi_{q}}{\hat{v}_{q,t+1}^{2}}\right|$ $\displaystyle\lesssim_{L}\log^{2}d\left(\frac{|\xi_{p}|+|\xi_{q}|}{v_{q}^{2}}+ \eta\left(1+\frac{v_{p}^{2}}{v_{q}^{2}}\right)\left(v_{p}^{2L-2}+v_{q}^{2L-2} \right)\right)$ $\displaystyle\lesssim_{L}\eta ma_{0}d\log^{2}d+\eta\frac{\log^{2L}d}{P^{L-1}}+ \eta^{2}dP^{2}\log^{4L+2}\left(\frac{d}{\delta_{\mathop{\mathbb{P\/}}}}\right),$ $$

and

$$ $\displaystyle\left|\frac{2\eta v_{p}Z_{p}}{\left(1+2\eta\left(2Lv_{q}^{2L-2}+2 -\rho\right)\right)v_{q}^{2}}\frac{2\eta v_{q}Z_{q}+\xi_{q}}{\hat{v}_{q,t+1}^{ 2}}\right|$ $\displaystyle\lesssim_{L}\;$ $\displaystyle\frac{\eta^{2}|v_{p}Z_{p}|}{v_{q}^{3}}|Z_{q}|+\frac{\eta|v_{p}Z_{ p}|}{v_{q}^{4}}|\xi_{q}|$ $\displaystyle\lesssim_{L}\;$ $\displaystyle\eta^{2}dP^{2}\log^{4L+1}\left(\frac{d}{\delta_{\mathop{\mathbb{P \/}}}}\right)+\eta^{3}d^{1.5}P^{3}\log^{6L+1}\left(\frac{d}{\delta_{\mathop{ \mathbb{P\/}}}}\right)+\eta^{2}d^{1.5}P\log^{2L+1}\left(\frac{d}{\delta_{ \mathop{\mathbb{P\/}}}}\right)ma_{0},$ $$

and, similarly,

$$ $\displaystyle\left|\frac{v_{p}^{2}}{v_{q}^{2}}\frac{2\eta v_{q}Z_{q}}{\left(1+ 2\eta\left(2-\rho\right)\right)v_{q}^{2}+4L\eta v_{q}^{2L}}\frac{2\eta v_{q}Z_ {q}+\xi_{q}}{\hat{v}_{q,t+1}^{2}}\right|$ $\displaystyle\lesssim_{L}\;$ $\displaystyle\frac{v_{p}^{2}\eta|Z_{q}|}{|v_{q}|^{5}}\left(\eta|v_{q}Z_{q}|+| \xi_{q}|\right)$ $\displaystyle\lesssim_{L}\;$ $\displaystyle\eta^{2}dP^{2}\log^{4L+2}\left(\frac{d}{\delta_{\mathop{\mathbb{P \/}}}}\right)+\eta^{3}d^{1.5}P^{3}\log^{6L+2}\left(\frac{d}{\delta_{\mathop{ \mathbb{P\/}}}}\right)+\eta^{2}d^{1.5}P\log^{2L+2}\left(\frac{d}{\delta_{ \mathop{\mathbb{P\/}}}}\right)ma_{0}.$ $$

Suppose that $\eta\leq 1/(dP^{2})$, which is implied by the condition of Lemma C.4.
Then, using the language of Lemma F.8, we have

$$ $\Xi\lesssim_{L}\eta ma_{0}d\log^{2}d+\eta\frac{\log^{2L}d}{P^{L-1}}+\eta^{2}dP ^{2}\log^{4L+2}\left(\frac{d}{\delta_{\mathop{\mathbb{P\/}}}}\right).$ (12) $$

Combine this with (11), recall $T\eta=O(\log d)$, apply
Lemma F.8, and we obtain

$$ $\displaystyle\left|\frac{v_{p}^{2}}{v_{q}^{2}}-\frac{v_{0,p}^{2}}{v_{0,q}^{2}}\right|$ $\displaystyle\lesssim_{L}T\eta ma_{0}d\log^{2}d+T\eta\frac{\log^{2L}d}{P^{L-1} }+T\eta^{2}dP^{2}\log^{4L+2}\left(\frac{d}{\delta_{\mathop{\mathbb{P\/}}}} \right)\sqrt{\delta_{\mathop{\mathbb{P\/}}}^{-1}T\eta^{2}dP^{2}\log^{4}d}$ $\displaystyle\lesssim_{L}ma_{0}d\log^{3}d+\frac{\log^{2L+1}d}{P^{L-1}}+\eta dP ^{2}\log^{4L+3}\left(\frac{d}{\delta_{\mathop{\mathbb{P\/}}}}\right)\sqrt{ \delta_{\mathop{\mathbb{P\/}}}^{-1}\eta dP^{2}\log^{5}d},$ $$

throughout Stage 1.1 with probability at least $1-\Theta(\delta_{\mathop{\mathbb{P\/}}})$. For the RHS to be bounded by $\varepsilon_{R}\in(0,1)$,
it suffices to require

$$ $ma_{0}\lesssim_{L}\frac{\varepsilon_{R}}{d\log^{3}d},\quad P\gtrsim_{L}\frac{ \log^{3}d}{\varepsilon_{R}},\quad\eta\lesssim_{L}\frac{\varepsilon_{R}\sqrt{ \delta_{\mathop{\mathbb{P\/}}}}}{dP^{2}\log^{2L+2}\left(d/\delta_{\mathop{ \mathbb{P\/}}}\right)}.$ $$

∎

###### Lemma C.5 (Gap between large and small coordinates) .

###### Proof.

###### Lemma C.6 (Gap between large coordinates) .

###### Proof.

#### C.1.3 Other induction hypotheses

First, we verify the induction hypothesis: $v_{p}^{2}\leq\log^{2}d/P$ for all $p\in[P]$. This condition is used
to ensure the influence of the higher-order term is small compared to the influence of the second-order terms.

###### Lemma C.7 (Bounds for moderately large $v_{p}^{2}$ ) .

Let ${\bm{v}}$ be an arbitrary first-layer neuron. Suppose that $p\in[P]$ and
$c_{v}/d\leq v_{0,p}^{2}\ll c_{v}^{\prime}\log^{2}d/d$ for some small $c_{v},c_{v}^{\prime}>0$.
Then, if we choose

$$ $ma_{0}\lesssim_{L}\frac{c_{v}}{d\log d}\quad\text{and}\quad\eta\lesssim_{L} \frac{c_{v}(1\wedge c_{v})\delta_{\mathop{\mathbb{P\/}}}}{dP^{2}\log^{4L+1} \left(d/\delta_{\mathop{\mathbb{P\/}}}\right)},$ $$

then there exists a universal constant $C\geq 1$ such that
with probability at least $1-O(\delta_{\mathop{\mathbb{P\/}}})$, we have

$$ $\frac{1}{2}(1+4\eta)^{t}v_{0,p}^{2}\leq v_{t,p}^{2}\leq\frac{3C}{2}(1+4\eta)^{ t}v_{0,p}^{2},\quad\forall t\leq T.$ $$

In particular, this implies $v_{t,p}^{2}\leq\log^{2}d/P$ throughout Stage 1.1.

###### Proof.

First, by Lemma C.1, for any $p\leq P$, we have

$$ $\displaystyle\hat{v}_{t+1,p}^{2}$ $\displaystyle\leq\left(1+4\eta+4L\eta v_{p}^{2L-2}\right)v_{p}^{2}+2\eta v_{p} Z_{p}+300L^{3}\eta ma_{0}+300L^{3}\eta^{2}\left(1\vee Z_{p}^{2}\right)$ $\displaystyle\leq\left(1+4\eta\left(1+L\left(\frac{\log^{2}d}{P}\right)^{L-1} \right)\right)v_{p}^{2}+2\eta v_{p}Z_{p}+300L^{3}\eta ma_{0}+300L^{3}\eta^{2} \left(1\vee Z_{p}^{2}\right),$ $$

where the second line comes from the induction hypothesis $v_{p}^{2}\leq\log^{2}d/P$. For notational simplicity,
put $\delta_{v}=L\left(\log^{2}d/P\right)^{L-1}$ (as in the proof of Lemma C.4) and $\xi_{t+1,p}=300L^{3}\eta ma_{0}+300L^{3}\eta^{2}\left(1\vee Z_{p}^{2}\right)$, so that the above can be rewritten as

$$ $v_{t+1,p}^{2}\leq\hat{v}_{t+1,p}^{2}\leq\left(1+4\eta(1+\delta_{v})\right)v_{p }^{2}+2\eta v_{p}Z_{p}+\xi_{p}.$ $$

By the tail bound in Lemma 2.2, there exists some constant $C_{L}>0$ that may depend on
$L$ such that for any $\delta_{\mathop{\mathbb{P\/}},\xi}\in(0,1)$, we have

$$ $|Z_{p}|\leq C_{L}^{2L}P\log^{2L}\left(\frac{C_{L}}{\delta_{\mathop{\mathbb{P\/ }},\xi}}\right)\quad\text{with probability at least $1-\delta_{\mathop{\mathbb {P\/}},\xi}$.}$ $$

Meanwhile, for the martingale difference term, by our induction hypothesis on $v_{p}$ and the variance estimate
in Lemma 2.2, we have

$$ $\displaystyle\mathop{\mathbb{E\/}}\left[(2\eta v_{p}Z_{p})^{2}\mid\mathcal{F}_ {t}\right]\leq 4C_{L}\eta^{2}v_{p}^{2}P^{2}$ $\displaystyle\lesssim_{L}\left(1+4\eta(1+\delta_{v})\right)^{t}\eta^{2}v_{0,P} ^{2}P^{2}$ $\displaystyle\lesssim_{L}\left(1+4\eta(1+\delta_{v})\right)^{t}\eta^{2}\frac{P ^{2}\log^{2}d}{d}.$ $$

Using the language of Lemma F.6, these mean

$$ $\Xi\lesssim_{L}\eta\left(ma_{0}\vee\eta P^{2}\log^{4L}\left(\frac{1}{\delta_{ \mathop{\mathbb{P\/}},\xi}}\right)\right),\quad\sigma_{Z}^{2}\lesssim_{L}\eta^ {2}\frac{P^{2}\log^{2}d}{d}.$ $$

Put $x_{t}=(1+4\eta(1+\delta_{v}))^{t}v_{0,p}^{2}$ where $x_{0}=v_{0,p}^{2}\geq c_{v}/d$.
By the proof of Lemma C.4, we know $(1+4\eta)^{T}=\Theta(d/P)$.
In particular, this implies $\eta T=\frac{1+o(1)}{4}\log(d/P)$. Then, by Lemma F.6,
we have $v_{p}^{2}\leq(1\pm 0.5)x_{t}$ with probability at least $1-2\delta_{\mathop{\mathbb{P\/}}}$, as
long as $ma_{0}$ and $\eta$ are chosen so that

$$ $\displaystyle\eta\left(ma_{0}\vee\eta P^{2}\log^{4L}\left(\frac{T}{\delta_{ \mathop{\mathbb{P\/}}}}\right)\right)\lesssim_{L}\frac{x_{0}}{4T}$ $\displaystyle\quad\Leftarrow\quad\left\{\begin{aligned} &ma_{0}\lesssim_{L} \frac{c_{v}}{d\log d},\\ &\eta\lesssim_{L}\frac{c_{v}}{dP^{2}\log^{4L+1}\left(d/\delta_{\mathop{\mathbb {P\/}}}\right)}\end{aligned}\right.$ $\displaystyle\eta^{2}\frac{P^{2}\log^{2}d}{d}\lesssim_{L}\frac{\delta_{\mathop {\mathbb{P\/}}}\alpha x_{0}^{2}}{16}$ $\displaystyle\quad\Leftarrow\quad\eta\lesssim_{L}\frac{\delta_{\mathop{\mathbb {P\/}}}c_{v}^{2}}{dP^{2}\log^{2}d}.$ $$

To complete the proof, we now estimate $x_{t}$.
Clear that $x_{t}\geq(1+4\eta)^{t}x_{0}$. Meanwhile, we have

$$ $\left(\frac{1+4\eta(1+\delta_{v})}{1+4\eta}\right)^{T}=\left(1+\frac{4\eta \delta_{v}}{1+4\eta}\right)^{T}\leq\left(1+4\eta\delta_{v}\right)^{T}\\ \leq\exp\left(4\eta T\delta_{v}\right)\leq\exp\left((1+o(1))\delta_{v}\log \left(\frac{d}{P}\right)\right)\leq(d/P)^{2\delta_{v}}.$ $$

When $P\geq\log^{3}d$, the last term is bounded by a universal constant $C>0$. As a result, we have

$$ $x_{t}\leq(1+4\eta(1+\delta_{v}))^{t}x_{0}=\left(\frac{1+4\eta(1+\delta_{v})}{1 +4\eta}\right)^{t}(1+4\eta)^{t}x_{0}\leq C(1+4\eta)^{t}x_{0}.$ $$

∎

###### Lemma C.8 (Upper bound for small $v_{q}^{2}$ ) .

Let ${\bm{v}}$ be an arbitrary first-layer neuron. Suppose that $q\in[P]$ and $v_{q}^{2}\leq c_{v}/d$ for some $c_{v}>0$.
Then, if we choose

$$ $ma_{0}\lesssim_{L}\frac{c_{v}}{d\log d}\quad\text{and}\quad\eta\lesssim_{L} \frac{c_{v}(1\wedge c_{v})\delta_{\mathop{\mathbb{P\/}}}}{dP^{2}\log^{4L+1} \left(d/\delta_{\mathop{\mathbb{P\/}}}\right)},$ $$

then there exists a universal constant $C\geq 1$ such that
with probability at least $1-O(\delta_{\mathop{\mathbb{P\/}}})$, we have

$$ $v_{t,q}^{2}\leq 2Cc_{v}\frac{(1+4\eta)^{t}}{d},\quad\forall t\leq T.$ $$

###### Proof.

The proof is essentially the same as the previous one. It suffices to use Lemma F.7 in place of Lemma F.6.
∎

The following lemma is not used in our proof. It serves as an example of using Lemma F.10 to obtain $\operatorname{poly}\log$ dependence on $\delta_{\mathop{\mathbb{P\/}}}$.

###### Lemma C.9 .

There exists a constant $C_{L}>0$ that may depend on $L$ such that if we choose

$$ $\displaystyle ma_{0}\leq\frac{\log d}{C_{L}d}\quad\text{and}\quad\eta\leq\frac {1}{C_{L}dP\log^{2L+3}\left(\frac{Tmd}{\delta_{\mathop{\mathbb{P\/}}}}\right)},$ $$

then with probability at least $1-\delta_{\mathop{\mathbb{P\/}}}$, we have

$$ $\sup_{i\in[m]}\sup_{r>P}\sup_{t\leq T}v_{i,t,r}^{2}\leq\frac{\log^{2}d}{d}.$ $$

###### Proof.

We will use Lemma F.10. Fix a first-layer neuron ${\bm{v}}$ and
$r>P$. Assume the induction hypothesis $v_{r}^{2}\leq K_{v}/d$ where $K_{v}>0$ is a parameter to be determined later.
Recall from Lemma C.1 that

$$ $v_{t+1,r}^{2}\leq\hat{v}_{t+1,r}^{2}=\left(1-2\eta\rho\right)v_{r}^{2}+2\eta v _{r}Z_{r}\pm 300L^{3}\eta ma_{0}\pm 300L^{3}\eta^{2}\left(1\vee Z_{r}^{2} \right).$ $$

Let $\xi_{t+1,r}$ denote the last two terms. Then, we can write

$$ $\hat{v}_{t+1,r}^{2}\leq v_{r}^{2}+2\eta v_{r}Z_{r}+\xi_{r}.$ $$

By the tail bound in Lemma 2.2, for any $\delta_{\mathop{\mathbb{P\/}}}\in(0,1)$,

$$ $|Z_{r}|\leq C_{L}^{2L}P\log^{2L}\left(\frac{T}{C_{L}\delta_{\mathop{\mathbb{P \/}}}}\right)\quad\text{with probability at least $1-\delta_{\mathop{\mathbb{P \/}}}/T$}.$ $$

Hence, with probability at least $1-\delta_{\mathop{\mathbb{P\/}}}/T$, we have

$$ $\displaystyle|\xi_{r}|$ $\displaystyle\leq 300L^{3}\eta ma_{0}+300L^{3}\eta^{2}C_{L}^{2L}P\log^{2L} \left(\frac{T}{C_{L}\delta_{\mathop{\mathbb{P\/}}}}\right)$ $\displaystyle\leq 600L^{3}C_{L}^{2L}\eta\left(ma_{0}\vee\eta P\log^{2L}\left( \frac{T}{C_{L}\delta_{\mathop{\mathbb{P\/}}}}\right)\right)=:\Xi.$ $$

Meanwhile, for the martingale difference terms, $Z_{r}$ satisfies the tail bound (15)
with $a=C_{L}$, $b=P^{-1/(2L)}$, $c=1/(2L)$, and $\sigma_{Z}^{2}=C_{L}P^{2}$. Hence, by Lemma F.10, we have

$$ $\displaystyle\sup_{t\leq T}\left|v_{t,r}^{2}-v_{0,r}^{2}\right|$ $\displaystyle\leq T\Xi+\frac{2K_{v}\eta C_{c}}{d}\sqrt{T\left(\sigma_{Z}^{2}+ \frac{1}{b^{2/c}}+\frac{\log^{1/c}\left(\frac{aT}{b\sigma_{Z}\delta_{\mathop{ \mathbb{P\/}}}}\right)}{b^{1/c}}\right)\log\left(\frac{T}{\delta_{\mathop{ \mathbb{P\/}}}}\right)}$ $\displaystyle\leq C_{L}^{\prime}T\eta\left(ma_{0}\vee\eta P\log^{2L}\left( \frac{T}{C_{L}\delta_{\mathop{\mathbb{P\/}}}}\right)\right)+\frac{K_{v}}{d}C_{ L}^{\prime}\sqrt{\eta^{2}T}P\log^{L+1}\left(\frac{C_{L}PT}{\delta_{\mathop{ \mathbb{P\/}}}}\right),$ $$

with probability at least $1-2\delta_{\mathop{\mathbb{P\/}}}$, for some constant $C_{L}^{\prime}>0$ that may depend on $L$.
Recall that $T\leq\eta^{-1}\log d$. Therefore,

$$ $\displaystyle\sup_{t\leq T}\left|v_{t,r}^{2}-v_{0,r}^{2}\right|$ $\displaystyle\leq C_{L}^{\prime}\log d\left(ma_{0}\vee\eta P\log^{2L}\left( \frac{T}{\delta_{\mathop{\mathbb{P\/}}}}\right)\right)+\frac{K_{v}}{d}C_{L}^{ \prime}\sqrt{\eta\log d}P\log^{L+1}\left(\frac{PT}{\delta_{\mathop{\mathbb{P\/ }}}}\right),$ $$

with probability at least $1-2\delta_{\mathop{\mathbb{P\/}}}$. Thus, apply the union bound over all neurons and all $r>P$,
replace $\delta_{\mathop{\mathbb{P\/}}}$ with $\delta_{\mathop{\mathbb{P\/}}}/(2md)$, and we obtain

$$ $\displaystyle\sup_{i\in[m]}\sup_{r>P}\sup_{t\leq T}\left|v_{i,t,r}^{2}-v_{i,0, r}^{2}\right|$ $\displaystyle\leq C_{L}^{\prime\prime}\log d\left(ma_{0}\vee\eta P\log^{2L} \left(\frac{Tmd}{\delta_{\mathop{\mathbb{P\/}}}}\right)\right)$ $\displaystyle\qquad+\frac{K_{v}}{d}C_{L}^{\prime\prime}\sqrt{\eta\log d}P\log^ {L+1}\left(\frac{PTmd}{\delta_{\mathop{\mathbb{P\/}}}}\right),$ $$

with probability at least $1-\delta_{\mathop{\mathbb{P\/}}}$. Finally, recall that we assume
$\sup_{i\in[m]}\sup_{r>P}\sup_{t\leq T}v_{i,0,r}^{2}\leq\log^{2}/(2d)$. Choose $K_{v}=\log^{2}d$. Then,
we have $\sup_{i\in[m]}\sup_{r>P}\sup_{t\leq T}v_{i,t,r}^{2}\leq\log^{2}d/d$ with probability at least
$1-\delta_{\mathop{\mathbb{P\/}}}$ as long as

$$ $\displaystyle C_{L}^{\prime\prime}\log d\left(ma_{0}\vee\eta P\log^{2L}\left( \frac{Tmd}{\delta_{\mathop{\mathbb{P\/}}}}\right)\right)\leq\frac{\log^{2}d}{2d}$ $\displaystyle\quad\Leftarrow\quad\left\{\begin{aligned} &ma_{0}\leq\frac{\log d }{2C_{L}^{\prime\prime}d}\\ &\eta\leq\frac{\log d}{2C_{L}^{\prime\prime}dP\log^{2L}\left(\frac{Tmd}{\delta _{\mathop{\mathbb{P\/}}}}\right)}\end{aligned}\right.$ $\displaystyle\frac{K_{v}}{d}C_{L}^{\prime\prime}\sqrt{\eta\log d}P\log^{L+1} \left(\frac{PTmd}{\delta_{\mathop{\mathbb{P\/}}}}\right)\leq\frac{\log^{2}d}{2d}$ $\displaystyle\quad\Leftarrow\quad\eta\leq\frac{1}{4(C_{L}^{\prime\prime})^{2}P ^{2}\log^{2L+3}\left(\frac{PTmd}{\delta_{\mathop{\mathbb{P\/}}}}\right)}.$ $$

∎

###### Lemma C.7 (Bounds for moderately large v p 2 superscript subscript 𝑣 𝑝 2 v_{p}^{2} italic_v start_POSTSUBSCRIPT italic_p end_POSTSUBSCRIPT start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT ) .

###### Proof.

###### Lemma C.8 (Upper bound for small v q 2 superscript subscript 𝑣 𝑞 2 v_{q}^{2} italic_v start_POSTSUBSCRIPT italic_q end_POSTSUBSCRIPT start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT ) .

###### Proof.

###### Lemma C.9 .

###### Proof.

### C.2 Stage 1.2: recovery of the directions

Let ${\bm{v}}$ be an arbitrary first-layer neuron. Assume w.l.o.g. that $v_{1}^{2}$ is the largest at initialization and
$v_{0,1}^{2}/\max_{2\leq k\leq P}v_{0,k}^{2}\geq 1+c_{g}$ for some small constant $c_{g}>0$.
By Lemma C.2, we know this gap can be approximately preserved. In other words,
we may assume that $v_{T_{1},1}^{2}/\max_{2\leq k\leq P}v_{T_{1},k}^{2}\geq 1+c_{g}$ for some small constant
$c_{g}>0$ that is potentially smaller than the previous $c_{g}$.
In this subsection, we show that $v_{1}^{2}$ will grow from $\Omega(1/P)$ to $3/4$ and then to close to $1$.
Formally, we prove the following lemma.

###### Lemma C.10 (Stage 1.2) .

Let ${\bm{v}}\in\mathbb{S}^{d-1}$ be an arbitrary first-layer neuron satisfying $v_{T_{1},1}^{2}\geq c/P$ and
$v_{T_{1},1}^{2}/\max_{2\leq k\leq P}v_{T_{1},k}^{2}\geq 1+c$ for some small universal constant $c>0$.
Let $\delta_{\mathop{\mathbb{P\/}}}\in(e^{-\log^{C}d},1)$ and $\varepsilon_{v}>0$ be given. Suppose that we choose

$$ $ma_{0}\lesssim_{L}\frac{\varepsilon_{v}}{dP^{2L}\log(1/\varepsilon_{v})}\quad \text{and}\quad\eta\lesssim_{L}\frac{\varepsilon_{v}^{2}\delta_{\mathop{ \mathbb{P\/}}}}{dP^{L+3}\log^{4L}(d/\delta_{\mathop{\mathbb{P\/}}})}.$ $$

Then, with probability at least $1-O(\delta_{\mathop{\mathbb{P\/}}})$, we have $v_{1}^{2}\geq 1-\varepsilon_{v}$ within
$O_{L}\left(\left(P^{L-1}+\log(1/\varepsilon_{v})\right)/\eta\right)$ iterations.

###### Proof.

It suffices to combine Lemma C.12 and Lemma C.13.
∎

###### Lemma C.11 (Dynamics of $v_{1}^{2}$ ) .

We have

$$ $v_{t+1,1}^{2}=v_{1}^{2}\left(1+4L\eta v_{1}^{2L-2}-4L\eta\left\|{\bm{v}}\right \|_{2L}^{2L}\right)+\frac{2\eta v_{1}Z_{1}-2\eta\left\langle{\bm{v}},\bm{Z} \right\rangle}{1+2\eta\left(2-\rho\right)+4L\eta\left\|{\bm{v}}\right\|_{2L}^{ 2L}}+\xi_{t+1}$ $$

where $\xi_{t}$ satisfies $|\xi_{t}|\leq C_{L}\eta d\left(ma_{0}\vee\eta P^{2}\log^{4L}\left(\frac{d}{
\delta_{\mathop{\mathbb{P\/}},\xi}}\right)\right)$,
with probability least $1-\delta_{\mathop{\mathbb{P\/}},\xi}$ for some constant $C_{L}>0$ that can depend on $L$.

###### Proof.

Recall from Lemma C.1 that

$$ $\displaystyle\hat{v}_{t+1,1}^{2}$ $\displaystyle=\left(1+2\eta\left(2Lv_{1}^{2L-2}+2-\rho\right)\right)v_{1}^{2}+ 2\eta v_{1}Z_{1}$ $\displaystyle\qquad\underbrace{\pm 300L^{3}\eta ma_{0}\pm 300L^{3}\eta^{2} \left(1\vee Z_{k}^{2}\right)}_{=:\;\xi_{1,t+1}}$ $\displaystyle=v_{1}^{2}\left(1+2\eta\left(2-\rho\right)+4L\eta v_{1}^{2L-2} \right)+2\eta v_{1}Z_{1}+\xi_{1,t+1},$ $$

where $\rho:=2\left\|{\bm{v}}_{\leq P}\right\|^{2}+2L\left\|{\bm{v}}_{\leq P}\right\|
_{2L}^{2L}$. Meanwhile, we also have

$$ $\displaystyle\left\|\hat{{\bm{v}}}_{t+1}\right\|^{2}$ $\displaystyle=\sum_{k=1}^{d}\left(1+2\eta\left(2Lv_{k}^{2L-2}+2-\rho\right) \right)v_{k}^{2}+2\eta\left\langle{\bm{v}},\bm{Z}\right\rangle+\left\langle\bm {1},\bm{\xi}\right\rangle$ $\displaystyle=1+2\eta\left(2-\rho\right)+4L\eta\left\|{\bm{v}}\right\|_{2L}^{2 L}+2\eta\left\langle{\bm{v}},\bm{Z}\right\rangle+\left\langle\bm{1},\bm{\xi} \right\rangle.$ $$

Then, we compute

$$ $\displaystyle v_{t+1,1}^{2}$ $\displaystyle=\frac{v_{1}^{2}\left(1+2\eta\left(2-\rho\right)+4L\eta v_{1}^{2L -2}\right)+2\eta v_{1}Z_{1}+\xi_{1,t+1}}{1+2\eta\left(2-\rho\right)+4L\eta \left\|{\bm{v}}\right\|_{2L}^{2L}+2\eta\left\langle{\bm{v}},\bm{Z}\right \rangle+\left\langle\bm{1},\bm{\xi}\right\rangle}$ $\displaystyle=v_{1}^{2}\frac{1+2\eta\left(2-\rho\right)+4L\eta v_{1}^{2L-2}}{1 +2\eta\left(2-\rho\right)+4L\eta\left\|{\bm{v}}\right\|_{2L}^{2L}+2\eta\left \langle{\bm{v}},\bm{Z}\right\rangle+\left\langle\bm{1},\bm{\xi}\right\rangle}$ $\displaystyle\qquad+\frac{2\eta v_{1}Z_{1}}{1+2\eta\left(2-\rho\right)+4L\eta \left\|{\bm{v}}\right\|_{2L}^{2L}+2\eta\left\langle{\bm{v}},\bm{Z}\right \rangle+\left\langle\bm{1},\bm{\xi}\right\rangle}+\frac{\xi_{1,t+1}}{\left\| \hat{{\bm{v}}}_{t+1}\right\|^{2}}$ $\displaystyle=:\textnormal{{Tmp}}_{1}+\textnormal{{Tmp}}_{2}+\textnormal{{Tmp} }_{3}.$ $$

For notational simplicity, we define $N_{v}^{2}:=1+2\eta\left(2-\rho\right)+4L\eta\left\|{\bm{v}}\right\|_{2L}^{2L}$.
Meanwhile, by the tail bound in Lemma 2.2, for each $k\in[d]$ and any
$\delta_{\mathop{\mathbb{P\/}},\xi}\in(0,1)$, we have

$$ $|Z_{k}|\leq C_{2L}^{L}P\log^{2L}\left(\frac{C_{L}}{\delta_{\mathop{\mathbb{P\/ }},\xi}}\right)\quad\text{with probability at least $1-\delta_{\mathop{\mathbb {P\/}},\xi}$}.$ $$

Then, by union bound, with at least the same probability, we have

$$ $|\left\langle{\bm{v}},\bm{Z}\right\rangle|\vee\max_{k\in[d]}|Z_{k}|\leq C_{L}^ {2L}P\log^{2L}\left(\frac{2C_{L}d}{\delta_{\mathop{\mathbb{P\/}},\xi}}\right).$ $$

As a result, with at least the same probability, we have

$$ $\displaystyle|\xi_{1}|$ $\displaystyle\leq 600L^{3}\eta\left(ma_{0}\vee\eta C_{L}^{4L}P^{2}\log^{4L} \left(\frac{2C_{L}d}{\delta_{\mathop{\mathbb{P\/}},\xi}}\right)\right),$ $\displaystyle|\left\langle\bm{1},\bm{\xi}\right\rangle|$ $\displaystyle\leq 600L^{3}\eta d\left(ma_{0}\vee\eta C_{L}^{4L}P^{2}\log^{4L} \left(\frac{2C_{L}d}{\delta_{\mathop{\mathbb{P\/}},\xi}}\right)\right).$ $$

Now, we are ready to analyze each of $\textnormal{{Tmp}}_{i}$ ($i\in[3]$).

First, for the signal term $\textnormal{{Tmp}}_{1}$, we write

$$ $\displaystyle\frac{1+2\eta\left(2-\rho\right)+4L\eta v_{1}^{2L-2}}{1+2\eta \left(2-\rho\right)+4L\eta\left\|{\bm{v}}\right\|_{2L}^{2L}+2\eta\left\langle{ \bm{v}},\bm{Z}\right\rangle+\left\langle\bm{1},\bm{\xi}\right\rangle}$ $\displaystyle=\;$ $\displaystyle\frac{1+2\eta\left(2-\rho\right)+4L\eta v_{1}^{2L-2}}{1+2\eta \left(2-\rho\right)+4L\eta\left\|{\bm{v}}\right\|_{2L}^{2L}}\left(1-\frac{2 \eta\left\langle{\bm{v}},\bm{Z}\right\rangle+\left\langle\bm{1},\bm{\xi}\right \rangle}{N_{v}^{2}+2\eta\left\langle{\bm{v}},\bm{Z}\right\rangle+\left\langle \bm{1},\bm{\xi}\right\rangle}\right)$ $\displaystyle=\;$ $\displaystyle\frac{1+2\eta\left(2-\rho\right)+4L\eta v_{1}^{2L-2}}{1+2\eta \left(2-\rho\right)+4L\eta\left\|{\bm{v}}\right\|_{2L}^{2L}}\left(1-\frac{2 \eta\left\langle{\bm{v}},\bm{Z}\right\rangle}{N_{v}^{2}}\left(1-\frac{2\eta \left\langle{\bm{v}},\bm{Z}\right\rangle+\left\langle\bm{1},\bm{\xi}\right \rangle}{\left\|\hat{\bm{v}}_{t+1}\right\|^{2}}\right)-\frac{\left\langle\bm{1 },\bm{\xi}\right\rangle}{\left\|\hat{{\bm{v}}}_{t+1}\right\|^{2}}\right)$ $\displaystyle=\;$ $\displaystyle\frac{1+2\eta\left(2-\rho\right)+4L\eta v_{1}^{2L-2}}{1+2\eta \left(2-\rho\right)+4L\eta\left\|{\bm{v}}\right\|_{2L}^{2L}}\left(1-\frac{2 \eta\left\langle{\bm{v}},\bm{Z}\right\rangle}{N_{v}^{2}}\pm 4\eta^{2}\left \langle{\bm{v}},\bm{Z}\right\rangle^{2}\pm|\left\langle\bm{1},\bm{\xi}\right \rangle|\right).$ $$

For the first factor, by (10), we have

$$ $\displaystyle\frac{1+2\eta\left(2-\rho\right)+4L\eta v_{1}^{2L-2}}{1+2\eta \left(2-\rho\right)+4L\eta\left\|{\bm{v}}\right\|_{2L}^{2L}}$ $\displaystyle=\;$ $\displaystyle\left(1+2\eta\left(2-\rho\right)+4L\eta v_{1}^{2L-2}\right)\left( 1-2\eta\left(2-\rho\right)-4L\eta\left\|{\bm{v}}\right\|_{2L}^{2L}\pm 160L^{2} \eta^{2}\right)$ $\displaystyle=\;$ $\displaystyle 1+4L\eta v_{1}^{2L-2}-4L\eta\left\|{\bm{v}}\right\|_{2L}^{2L}\pm 3 00L^{2}\eta^{2}.$ $$

As a result, we have

$$ $\displaystyle\frac{\textnormal{{Tmp}}_{1}}{v_{1}^{2}}$ $\displaystyle=\left(1+4L\eta v_{1}^{2L-2}-4L\eta\left\|{\bm{v}}\right\|_{2L}^{ 2L}\pm 300L^{2}\eta^{2}\right)\left(1-\frac{2\eta\left\langle{\bm{v}},\bm{Z} \right\rangle}{N_{v}^{2}}\pm 4\eta^{2}\left\langle{\bm{v}},\bm{Z}\right\rangle ^{2}\pm|\left\langle\bm{1},\bm{\xi}\right\rangle|\right)$ $\displaystyle=1+4L\eta v_{1}^{2L-2}-4L\eta\left\|{\bm{v}}\right\|_{2L}^{2L}- \frac{2\eta\left\langle{\bm{v}},\bm{Z}\right\rangle}{N_{v}^{2}}$ $\displaystyle\qquad\pm O_{L}(1)\eta d\left(ma_{0}\vee\eta P^{2}\log^{4L}\left( \frac{d}{\delta_{\mathop{\mathbb{P\/}},\xi}}\right)\right).$ $$

Then, we consider the (approximate) martingale difference term $\textnormal{{Tmp}}_{2}$. We have

$$ $\displaystyle\textnormal{{Tmp}}_{2}$ $\displaystyle=\frac{2\eta v_{1}Z_{1}}{N_{v}^{2}}\left(1-\frac{2\eta\left \langle{\bm{v}},\bm{Z}\right\rangle+\left\langle\bm{1},\bm{\xi}\right\rangle}{ \left\|\hat{\bm{v}}_{t+1}\right\|^{2}}\right)$ $\displaystyle=\frac{2\eta v_{1}Z_{1}}{N_{v}^{2}}\pm O_{L}(1)\eta d\left(ma_{0} \vee\eta P^{2}\log^{4L}\left(\frac{d}{\delta_{\mathop{\mathbb{P\/}},\xi}} \right)\right).$ $$

Thus, we have

$$ $\displaystyle v_{t+1,1}^{2}$ $\displaystyle=v_{1}^{2}\left(1+4L\eta v_{1}^{2L-2}-4L\eta\left\|{\bm{v}}\right \|_{2L}^{2L}\right)-\frac{2\eta\left\langle{\bm{v}},\bm{Z}\right\rangle}{N_{v} ^{2}}+\frac{2\eta v_{1}Z_{1}}{N_{v}^{2}}$ $\displaystyle\qquad\pm O_{L}(1)\eta d\left(ma_{0}\vee\eta P^{2}\log^{4L}\left( \frac{d}{\delta_{\mathop{\mathbb{P\/}},\xi}}\right)\right).$ $$

∎

###### Lemma C.12 (Weak recovery of directions) .

Suppose that we choose

$$ $ma_{0}\leq\frac{c_{g,L}}{dP^{2L}}\quad\text{and}\quad\eta\leq\frac{c_{g,L} \delta_{\mathop{\mathbb{P\/}}}}{dP^{L+3}\log^{4L}\left(d/\delta_{\mathop{ \mathbb{P\/}}}\right)}.$ $$

Then within $O_{L}(\frac{P^{L-1}}{\eta c_{g,L}})$ iterations, we will have $v_{1}^{2}\geq 3/4$ with probability
at least $1-O(\delta_{\mathop{\mathbb{P\/}}})$.

###### Proof.

By Lemma C.11, we have

$$ $v_{t+1,1}^{2}=v_{1}^{2}\left(1+4L\eta v_{1}^{2L-2}-4L\eta\left\|{\bm{v}}\right \|_{2L}^{2L}\right)+\frac{2\eta v_{1}Z_{1}-2\eta\left\langle{\bm{v}},\bm{Z} \right\rangle}{1+2\eta\left(2-\rho\right)+4L\eta\left\|{\bm{v}}\right\|_{2L}^{ 2L}}+\xi_{t+1}$ $$

where $\xi_{t}$ satisfies $|\xi_{t}|\leq C_{L}\eta d\left(ma_{0}\vee\eta P^{2}\log^{4L}\left(\frac{d}{
\delta_{\mathop{\mathbb{P\/}},\xi}}\right)\right)$,
with probability least $1-\delta_{\mathop{\mathbb{P\/}},\xi}$ for some constant $C_{L}>0$ that can depend on $L$. Meanwhile,
by the variance bound in Lemma 2.2, we have

$$ $\mathop{\mathbb{E\/}}\left[\left(\frac{2\eta v_{1}Z_{1}-2\eta\left\langle{\bm{ v}},\bm{Z}\right\rangle}{1+2\eta\left(2-\rho\right)+4L\eta\left\|{\bm{v}} \right\|_{2L}^{2L}}\right)^{2}\;\bigg{|}\;\mathcal{F}_{t}\right]\lesssim_{L} \eta^{2}P^{2}.$ $$

For the signal term, we write

$$ $\displaystyle v_{1}^{2L-2}-\left\|{\bm{v}}_{\leq P}\right\|_{2L}^{2L}$ $\displaystyle=v_{1}^{2L-2}-v_{1}^{2L}-\sum_{k=2}^{P}v_{k}^{2L}$ $\displaystyle=v_{1}^{2L-2}\left(1-v_{1}^{2}\right)-\left(\left\|{\bm{v}}_{\leq P }\right\|^{2}-v_{1}^{2}\right)\sum_{k=2}^{P}\frac{v_{k}^{2}}{\left\|{\bm{v}}_{ \leq P}\right\|^{2}-v_{1}^{2}}v_{k}^{2L-2}.$ $$

Note that the last summation is a weighted average of $\{v_{k}^{2L-2}\}_{2\leq k\leq P}$.
Similar to the proof in Section C.1.2, we can maintain the
induction hypothesis $v_{1}^{2}/\max_{2\leq k\leq P}v_{k}^{2}\geq 1+c_{g}/2$(^7^77 The only difference is that now the $2L$-th order terms cannot be simply ignored as we no longer have the induction hypothesis $v_{p}^{2}\leq\log^{2}d/P$. To handle them, it suffices to note that if $v_{1}^{2}\geq v_{q}^{2}$, then those $2L$-th order terms are also larger for $v_{1}^{2}$, which will even lead to an amplification of the gap. In fact, this is why we can recover the directions using them.), which gives

$$ $\sum_{k=2}^{P}\frac{v_{k}^{2}}{\left\|{\bm{v}}_{\leq P}\right\|^{2}-v_{1}^{2}} v_{k}^{2L-2}\leq\left(\max_{2\leq k\leq P}v_{k}^{2}\right)^{L-1}\leq\left( \frac{v_{1}^{2}}{1+c_{g}/2}\right)^{L-1}=\frac{v_{1}^{2L-2}}{1+c_{g,L}},$ $$

where $c_{g,L}>0$ is a constant that depend on $L$ and $c_{g}$. Therefore,

$$ $\displaystyle v_{1}^{2L-2}-\left\|{\bm{v}}_{\leq P}\right\|_{2L}^{2L}$ $\displaystyle\geq v_{1}^{2L-2}\left(1-v_{1}^{2}\right)-\left(\left\|{\bm{v}}_{ \leq P}\right\|^{2}-v_{1}^{2}\right)\frac{v_{1}^{2L-2}}{1+c_{g,L}}$ $\displaystyle=\frac{v_{1}^{2L-2}}{1+c_{g,L}}\left(1-\left\|{\bm{v}}_{\leq P} \right\|^{2}+c_{g,L}\left(1-v_{1}^{2}\right)\right)$ $\displaystyle\geq\frac{c_{g,L}}{1+c_{g,L}}v_{1}^{2L-2}\left(1-v_{1}^{2}\right).$ $$

As a result, for the signal term, we have

$$ $\displaystyle v_{1}^{2}\left(1+4L\eta v_{1}^{2L-2}-4L\eta\left\|{\bm{v}}\right \|_{2L}^{2L}\right)$ $\displaystyle\geq v_{1}^{2}\left(1+4L\eta\frac{c_{g,L}}{1+c_{g,L}}v_{1}^{2L-2} \left(1-v_{1}^{2}\right)\right)$ $\displaystyle=v_{1}^{2}+4L\eta\frac{c_{g,L}}{1+c_{g,L}}v_{1}^{2L}\left(1-v_{1} ^{2}\right)$ $\displaystyle\geq v_{1}^{2}+\eta\frac{c_{g,L}L}{1+c_{g,L}}v_{1}^{2L},$ $$

where the last line comes from the induction hypothesis $v_{1}^{2}\leq 3/4$. Thus, using the notations of
Lemma F.11, we have

$$ $\displaystyle\alpha=\eta\frac{c_{g,L}L}{1+c_{g,L}},\quad\Xi=C_{L}\eta d\left( ma_{0}\vee\eta P^{2}\log^{4L}\left(\frac{d}{\delta_{\mathop{\mathbb{P\/}},\xi} }\right)\right),\quad\sigma_{Z}^{2}=C_{L}\eta^{2}P^{2},$ $$

for some large constant $C_{L}>0$ that may differ from the previous one. Meanwhile, by Lemma F.12 and the assumption $x_{0}=v_{1}^{2}\geq\Omega(1/P)$, we have

$$ $T\lesssim\frac{1}{x_{0}^{L-1}\alpha}\leq\frac{P^{L-1}}{\alpha}\lesssim_{L} \frac{P^{L-1}}{\eta c_{g,L}}.$ $$

Thus, to meet the conditions of Lemma F.11, it suffices to choose

$$ $\displaystyle\Xi\leq\frac{x_{0}}{4T}$ $\displaystyle\quad\Leftarrow\quad ma_{0}\leq\frac{c_{g,L}}{dP^{L}},\quad\eta \leq\frac{c_{g,L}}{dP^{L+2}\log^{4L}\left(d/\delta_{\mathop{\mathbb{P\/}}} \right)},$ $\displaystyle\sigma_{Z}^{2}\leq\frac{x_{0}^{2}\delta_{\mathop{\mathbb{P\/}}}}{ 16T}$ $\displaystyle\quad\Leftarrow\quad\eta\lesssim_{L}\frac{\delta_{\mathop{\mathbb {P\/}}}c_{g,L}}{P^{L+3}}.$ $$

∎

###### Lemma C.13 (Strong recovery of directions) .

Let ${\bm{v}}\in\mathbb{S}^{d-1}$ be an arbitrary first-layer neuron. Let $\delta_{\mathop{\mathbb{P\/}}}$ and $\varepsilon_{*}$ be given. Suppose that
we choose

$$ $ma_{0}\lesssim_{L}\frac{\varepsilon_{*}}{d\log(1/\varepsilon_{*})}\quad\text{ and}\quad\eta\lesssim_{L}\frac{\varepsilon_{*}^{2}\delta_{\mathop{\mathbb{P\/} }}}{dP^{2}\log^{4L}\left(d/\delta_{\mathop{\mathbb{P\/}}}\right)}.$ $$

Then, with probability at least $1-O(\delta_{\mathop{\mathbb{P\/}}})$, we have $v_{1}^{2}\geq 1-\varepsilon_{*}$ within $O_{L}(\log(1/\varepsilon_{*})/\eta)$
iterations.

###### Proof.

Again, by Lemma C.11, we have

$$ $v_{t+1,1}^{2}=v_{1}^{2}\left(1+4L\eta v_{1}^{2L-2}-4L\eta\left\|{\bm{v}}\right \|_{2L}^{2L}\right)+\frac{2\eta v_{1}Z_{1}-2\eta\left\langle{\bm{v}},\bm{Z} \right\rangle}{1+2\eta\left(2-\rho\right)+4L\eta\left\|{\bm{v}}\right\|_{2L}^{ 2L}}+\xi_{t+1}$ $$

where $\xi_{t}$ satisfies $|\xi_{t}|\leq C_{L}\eta d\left(ma_{0}\vee\eta P^{2}\log^{4L}\left(\frac{d}{
\delta_{\mathop{\mathbb{P\/}},\xi}}\right)\right)$,
with probability least $1-\delta_{\mathop{\mathbb{P\/}},\xi}$ for some constant $C_{L}>0$ that can depend on $L$.
Meanwhile, by the proof of the previous lemma, we have

$$ $\displaystyle v_{1}^{2}\left(1+4L\eta v_{1}^{2L-2}-4L\eta\left\|{\bm{v}}\right \|_{2L}^{2L}\right)$ $\displaystyle\geq v_{1}^{2}\left(1+4L\eta\frac{c_{g,L}}{1+c_{g,L}}v_{1}^{2L-2} \left(1-v_{1}^{2}\right)\right)$ $\displaystyle=v_{1}^{2}+4L\eta\frac{c_{g,L}}{1+c_{g,L}}v_{1}^{2L}\left(1-v_{1} ^{2}\right)$ $\displaystyle\geq v_{1}^{2}+4L\eta\frac{c_{g,L}}{1+c_{g,L}}\left(\frac{3}{4} \right)^{2L}\left(1-v_{1}^{2}\right).$ $$

This implies

$$ $1-v_{t+1,1}^{2}\leq\left(1-v_{1}^{2}\right)\left(1-4L\eta\frac{c_{g,L}}{1+c_{g ,L}}\left(\frac{3}{4}\right)^{2L}\right)-\frac{2\eta v_{1}Z_{1}-2\eta\left \langle{\bm{v}},\bm{Z}\right\rangle}{1+2\eta\left(2-\rho\right)+4L\eta\left\|{ \bm{v}}\right\|_{2L}^{2L}}-\xi_{t+1}$ $$

For the martingale difference term, also by the previous proof, we have

$$ $\mathop{\mathbb{E\/}}\left[\left(\frac{2\eta v_{1}Z_{1}-2\eta\left\langle{\bm{ v}},\bm{Z}\right\rangle}{1+2\eta\left(2-\rho\right)+4L\eta\left\|{\bm{v}} \right\|_{2L}^{2L}}\right)^{2}\;\bigg{|}\;\mathcal{F}_{t}\right]\lesssim_{L} \eta^{2}P^{2}.$ $$

Let $\varepsilon_{*}>0$ denote our target accuracy.
Hence, in the language of Lemma F.6,(^8^88When $\alpha$ is negative, it suffices to replace $x_{0}$ with our target $\varepsilon_{*}$.) we have

$$ $\displaystyle\alpha=-4L\eta\frac{c_{g,L}}{1+c_{g,L}}\left(\frac{3}{4}\right)^{ 2L},\quad$ $\displaystyle\eta T=O_{L}(\log(1/\varepsilon_{*})),$ $\displaystyle\sigma_{Z}^{2}=O_{L}(1)\eta^{2}P^{2},$ $\displaystyle\Xi=O_{L}(1)\eta d\left(ma_{0}\vee\eta P^{2}\log^{4L}\left(\frac{ Td}{\delta_{\mathop{\mathbb{P\/}}}}\right)\right).$ $$

To meet the conditions of Lemma F.6, it suffices to choose

$$ $\displaystyle\Xi\leq\frac{\varepsilon_{*}}{4T}$ $\displaystyle\quad\Leftarrow\quad ma_{0}\lesssim_{L}\frac{\varepsilon_{*}}{d \log(1/\varepsilon_{*})},\quad\eta\lesssim_{L}\frac{\varepsilon_{*}}{dP^{2} \log(1/\varepsilon_{*})\log^{4L}\left(d/\delta_{\mathop{\mathbb{P\/}}}\right)},$ $\displaystyle\sigma_{Z}^{2}\leq\frac{\delta_{\mathop{\mathbb{P\/}}}|\alpha| \varepsilon_{*}^{2}}{16}$ $\displaystyle\quad\Leftarrow\quad\eta\lesssim_{L}\frac{\delta_{\mathop{\mathbb {P\/}}}c_{g,L}\varepsilon_{*}^{2}}{P^{2}}.$ $$

Then, with probability at least $1-O(\delta_{\mathop{\mathbb{P\/}}})$, we have $v_{1}^{2}\geq 1-\varepsilon$ within $T=O_{L}(\log(1/\varepsilon_{*})/\eta)$
iterations.
∎

###### Lemma C.10 (Stage 1.2) .

###### Proof.

###### Lemma C.11 (Dynamics of v 1 2 superscript subscript 𝑣 1 2 v_{1}^{2} italic_v start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT ) .

###### Proof.

###### Lemma C.12 (Weak recovery of directions) .

###### Proof.

###### Lemma C.13 (Strong recovery of directions) .

###### Proof.

### C.3 Deferred proofs in this section

###### Proof of Lemma C.1 .

Recall that

$$ $\hat{v}_{t+1,k}=v_{t,k}+\eta\left(\mathbbm{1}\{k\leq P\}\left(2+2Lv_{k}^{2L-2} \right)-\rho\right)v_{k}+\eta Z_{t+1,k}+\eta O_{t+1,k},$ $$

where $|O_{t+1,k}|\leq 2Lma_{0}$. Then, we compute

$$ $\displaystyle\hat{v}_{t+1,k}^{2}$ $\displaystyle=\left(\left(1+\eta\left(\mathbbm{1}\{k\leq P\}\left(2+2Lv_{k}^{2 L-2}\right)-\rho\right)\right)v_{k}+\eta O_{k}+\eta Z_{k}\right)^{2}$ $\displaystyle=\left(1+2\eta\left(\mathbbm{1}\{k\leq P\}\left(2+2Lv_{k}^{2L-2} \right)-\rho\right)\right)v_{k}^{2}+2\eta v_{k}Z_{k}+2\eta v_{k}O_{k}$ $\displaystyle\qquad+\eta^{2}\left(\mathbbm{1}\{k\leq P\}\left(2+2Lv_{k}^{2L-2} \right)-\rho\right)^{2}v_{k}^{2}$ $\displaystyle\qquad+2\eta^{2}\left(\mathbbm{1}\{k\leq P\}\left(2+2Lv_{k}^{2L-2 }\right)-\rho\right)v_{k}Z_{k}$ $\displaystyle\qquad+2\eta^{2}\left(\mathbbm{1}\{k\leq P\}\left(2+2Lv_{k}^{2L-2 }\right)-\rho\right)v_{k}O_{k}$ $\displaystyle\qquad+\eta^{2}O_{k}^{2}+\eta^{2}Z_{k}^{2}+2\eta^{2}Z_{k}O_{k}.$ $$

The last four lines, which we denote by $\textnormal{{Tmp}}^{(2)}$ for notational simplicity, contain terms that are quadratic in
$\eta$.
The first term is the second line is the “signal term” that corresponds to the GD update, the second term forms a
martingale difference sequence and the second term captures the influence of other neuron and shrinks with $a_{0}$.

First, we bound the second-order terms.
For $\rho$, we have the following naïve upper bound:

$$ $\rho=2\sum_{i=1}^{P}v_{i}^{2}+2L\sum_{i=1}^{P}v_{i}^{2L}\leq\left(2+2L\max_{j \leq P}v_{j}^{2L-2}\right)\left\|{\bm{v}}_{\leq P}\right\|^{2}\leq 2+2L\max_{j \leq P}v_{j}^{2L-2}\leq 4L,$ (13) $$

where the last inequality comes from the fact $L\geq 2$. Similarly, we also have $2+2Lv_{k}^{2L-2}\leq 4L$.
Hence, we have

$$ $\left|\mathbbm{1}\{k\leq P\}\left(2+2Lv_{k}^{2L-2}\right)-\rho\right|\leq 2+2 Lv_{k}^{2L-2}+\rho\leq 8L.$ $$

Thus, for the second-order terms (last four lines), we have

$$ $\displaystyle|\textnormal{{Tmp}}^{(2)}|$ $\displaystyle\leq 64L^{2}\eta^{2}v_{k}^{2}+16L\eta^{2}|v_{k}Z_{k}|+16L\eta^{2} |v_{k}O_{k}|+\eta^{2}O_{k}^{2}+\eta^{2}Z_{k}^{2}+2\eta^{2}Z_{k}O_{k}$ $\displaystyle\leq 100L^{2}\eta^{2}v_{k}^{2}+10L\eta^{2}Z_{k}^{2}+10L\eta^{2}O_ {k}^{2}$ $\displaystyle\leq 300L^{3}\eta^{2}\left(v_{k}^{2}\vee Z_{k}^{2}\vee m^{2}a_{0} ^{2}\right),$ $$

where we use the inequality $ab\leq a^{2}/2+b^{2}/2$ in the second line to handle the cross terms. In other words,
we have

$$ $\displaystyle\hat{v}_{t+1,k}^{2}$ $\displaystyle=\left(1+2\eta\left(\mathbbm{1}\{k\leq P\}\left(2+2Lv_{k}^{2L-2} \right)-\rho\right)\right)v_{k}^{2}+2\eta v_{k}Z_{k}+2\eta v_{k}O_{k}$ $\displaystyle\qquad\pm 300L^{3}\eta^{2}\left(v_{k}^{2}\vee Z_{k}^{2}\vee m^{2} a_{0}^{2}\right).$ $$

Meanwhile, for the last term in the first line, we have $|2\eta v_{k}O_{k}|\leq 4L\eta v_{k}ma_{0}$. Thus,

$$ $\displaystyle\hat{v}_{t+1,k}^{2}$ $\displaystyle=\left(1+2\eta\left(\mathbbm{1}\{k\leq P\}\left(2+2Lv_{k}^{2L-2} \right)-\rho\right)\right)v_{k}^{2}+2\eta v_{k}Z_{k}$ $\displaystyle\qquad\pm 4L\eta v_{k}ma_{0}\pm 300L^{3}\eta^{2}m^{2}a_{0}^{2}\pm 3 00L^{3}\eta^{2}\left(v_{k}^{2}\vee Z_{k}^{2}\right)$ $\displaystyle=\left(1+2\eta\left(\mathbbm{1}\{k\leq P\}\left(2+2Lv_{k}^{2L-2} \right)-\rho\right)\right)v_{k}^{2}+2\eta v_{k}Z_{k}$ $\displaystyle\qquad\pm 300L^{3}\eta ma_{0}\pm 300L^{3}\eta^{2}\left(1\vee Z_{k }^{2}\right).$ $$

∎

###### Proof of Lemma C.1 .

## Appendix D Stage 2: training the second layer

###### Lemma D.1 .

Suppose that for each $p\in[P]$, there exists a first-layer neuron ${\bm{v}}_{i_{p}}$ with
$v_{i_{p},p}^{2}\geq 1-\varepsilon_{v}$ for some small positive $\varepsilon_{v}=O(1/P)$, then we can choose $\bm{a}_{*}\in\mathbb{R}^{m}$
with $\left\|\bm{a}_{*}\right\|=\sqrt{P}$ such that

$$ $\mathcal{L}(\bm{a}_{*},\bm{V}):=\mathop{\mathbb{E\/}}\left(f_{*}({\bm{x}})-f({ \bm{x}};\bm{a}_{*},\bm{V})\right)^{2}\leq 10LP^{2}\varepsilon_{v}.$ $$

###### Proof.

Choose one ${\bm{v}}_{i_{p}}$ for each $p\in[P]$.
Then, we set the $i_{p}$-th entries of $\bm{a}_{*}$ to be $1$ and all other entries 0 0. Then, we write

$$ $\displaystyle\left(f_{*}({\bm{x}})-f({\bm{x}};\bm{a}_{*},\bm{V})\right)^{2}$ $\displaystyle=\left(\sum_{k=1}^{P}\left(\phi(x_{k})-\phi({\bm{v}}_{i_{k}}\cdot {\bm{x}})\right)\right)^{2}$ $\displaystyle=\sum_{k,l=1}^{P}\left(\phi(x_{k})-\phi({\bm{v}}_{i_{k}}\cdot{\bm {x}})\right)\left(\phi(x_{l})-\phi({\bm{v}}_{i_{l}}\cdot{\bm{x}})\right).$ $$

Recall from the proof of Lemma 2.1
(cf. Section A) that
for any ${\bm{v}},{\bm{v}}^{\prime}\in\mathbb{S}^{d-1}$, we have

$$ $\mathop{\mathbb{E\/}}_{{\bm{x}}\sim\mathcal{N}\left(0,\bm{I}\right)}[\phi({\bm {v}}\cdot{\bm{x}})\phi({\bm{v}}^{\prime}\cdot{\bm{x}})]=\left\langle{\bm{v}},{ \bm{v}}^{\prime}\right\rangle^{2}+\left\langle{\bm{v}},{\bm{v}}^{\prime}\right \rangle^{2L}.$ $$

Hence, for $k=l$, we have

$$ $\displaystyle\mathop{\mathbb{E\/}}\left(\phi(x_{k})-\phi({\bm{v}}_{i_{k}}\cdot {\bm{x}})\right)^{2}$ $\displaystyle=\mathop{\mathbb{E\/}}\phi^{2}(x_{k})+\mathop{\mathbb{E\/}}\phi^{ 2}({\bm{v}}_{i_{k}}\cdot{\bm{x}})-2\mathop{\mathbb{E\/}}\phi(x_{k})\phi({\bm{v }}_{i_{k}}\cdot{\bm{x}})$ $\displaystyle=4-2\left(v_{i_{k},k}^{2}+v_{i_{k},k}^{2L}\right)$ $\displaystyle\leq 4L\varepsilon_{v}.$ $$

Meanwhile, for $k\neq l$, we have

$$ $\displaystyle\mathop{\mathbb{E\/}}\left(\phi(x_{k})-\phi({\bm{v}}_{i_{k}}\cdot {\bm{x}})\right)\left(\phi(x_{l})-\phi({\bm{v}}_{i_{l}}\cdot{\bm{x}})\right)$ $\displaystyle=\;$ $\displaystyle\mathop{\mathbb{E\/}}\phi(x_{k})\phi(x_{l})+\mathop{\mathbb{E\/}} \phi({\bm{v}}_{i_{k}}\cdot{\bm{x}})\phi({\bm{v}}_{i_{l}}\cdot{\bm{x}})-\mathop {\mathbb{E\/}}\phi(x_{k})\phi({\bm{v}}_{i_{l}}\cdot{\bm{x}})-\mathop{\mathbb{E \/}}\phi({\bm{v}}_{i_{k}}\cdot{\bm{x}})\phi(x_{l})$ $\displaystyle\leq\;$ $\displaystyle\left\langle{\bm{v}}_{i_{k}},{\bm{v}}_{i_{l}}\right\rangle^{2}+ \left\langle{\bm{v}}_{i_{k}},{\bm{v}}_{i_{l}}\right\rangle^{2L}.$ $$

Note that

$$ $\left\langle{\bm{v}}_{i_{k}},{\bm{v}}_{i_{l}}\right\rangle^{2}\leq 2v_{i_{l},k }^{2}+2\left\langle{\bm{v}}_{i_{k}}-\bm{e}_{k},{\bm{v}}_{i_{l}}\right\rangle^{ 2}\leq 2\varepsilon_{v}+2\left\|{\bm{v}}_{i_{k}}-\bm{e}_{k}\right\|^{2}=2 \varepsilon_{v}+4\left(1-v_{i_{k},k}\right)\leq 6\varepsilon_{v}.$ $$

As a result, $\left\langle{\bm{v}}_{i_{k}},{\bm{v}}_{i_{l}}\right\rangle^{2}+\left\langle{
\bm{v}}_{i_{k}},{\bm{v}}_{i_{l}}\right\rangle^{2L}\leq 10\varepsilon_{v}$. Combining these two
cases, we obtain

$$ $\mathop{\mathbb{E\/}}\left(f_{*}({\bm{x}})-f({\bm{x}};\bm{a}_{*},\bm{V})\right )^{2}\leq 4PL\varepsilon_{v}+10P^{2}\varepsilon_{v}\leq 10LP^{2}\varepsilon_{v}.$ $$

∎

Now, we are ready to prove the following generalization bound for Stage 2. The proof of it is adapted from
Section B.8 of Oko et al. [], which in turn is based on (Damian et al. [], Abbe et al. [], Ba et al. []).

###### Lemma D.2 .

Suppose that for each $p\in[P]$, there exists a first-layer neuron ${\bm{v}}_{i_{p}}$ with
$v_{i_{p},p}^{2}\geq 1-\varepsilon_{v}$ for some small positive $\varepsilon_{v}=O(1/P)$. Then, there exists some $\lambda>0$
such that the ridge estimator $\hat{\bm{a}}$ we obtain in Stage 2 satisfies

$$ $\left\|f(\cdot;\hat{\bm{a}},\bm{V})-f_{*}\right\|_{L^{1}(D)}\leq\frac{8\left\| \bm{a}_{*}\right\|\sqrt{m}}{\sqrt{N}\delta_{\mathop{\mathbb{P\/}}}}+\sqrt{10LP ^{2}\varepsilon_{v}},$ $$

with probability at least $1-2\delta_{\mathop{\mathbb{P\/}}}$.

###### Proof.

For notational simplicity, let $D=\mathcal{N}(0,1)$ and $\hat{D}=\frac{1}{N}\sum_{n=1}^{N}\delta_{{\bm{x}}_{T+n}}$ denote the empirical
distribution of the samples we use in Stage 2. In addition, we write $f_{\bm{a}}$ for $f(\cdot;\bm{a},\bm{V})$ where $\bm{V}$
is the first-layer weights we have obtained in Stage 1 and $\bm{X}=({\bm{x}}_{T+n})_{n=1}^{N}$.

Let $\bm{a}_{*}\in\mathbb{R}^{m}$ denote the second-layer weights we constructed in
Lemma D.1 and $\hat{\bm{a}}\in\mathbb{R}^{m}$ denote the ridge estimator obtained
via minimizing $\bm{a}\mapsto\left\|f_{*}-f_{\bm{a}}\right\|_{L^{2}(\hat{D})}^{2}+\lambda\left
\|\bm{a}\right\|^{2}$. By the equivalence
between norm-constrained linear regression and ridge regression, there exists $\lambda>0$ such that

$$ $\left\|f_{*}-f_{\hat{\bm{a}}}\right\|_{L^{2}(\hat{D})}^{2}\leq\left\|f_{*}-f_{ \bm{a}_{*}}\right\|_{L^{2}(\hat{D})}^{2}\quad\text{and}\quad\left\|\hat{\bm{a} }\right\|\leq\left\|\bm{a}_{*}\right\|.$ $$

Choose this $\lambda$ and let $\mathcal{F}:=\left\{f(\cdot;\bm{a})\,:\,\left\|\bm{a}\right\|\leq\left\|\bm{a}
_{*}\right\|\right\}$ be our hypothesis
class. Note that $f_{\hat{\bm{a}}}\in\mathcal{F}$. Moreover, we have

$$ $\displaystyle\left\|f_{\hat{\bm{a}}}-f_{*}\right\|_{L^{1}(D)}$ $\displaystyle=\left(\left\|f_{\hat{\bm{a}}}-f_{*}\right\|_{L^{1}(D)}-\left\|f_ {\hat{\bm{a}}}-f_{*}\right\|_{L^{1}(\hat{D})}\right)+\left\|f_{\hat{\bm{a}}}-f _{*}\right\|_{L^{1}(\hat{D})}$ $\displaystyle\leq\sup_{\bm{a}\,:\,\left\|\bm{a}\right\|\leq\left\|\bm{a}_{*} \right\|}\left(\left\|f_{\bm{a}}-f_{*}\right\|_{L^{1}(D)}-\left\|f_{\bm{a}}-f_ {*}\right\|_{L^{1}(\hat{D})}\right)+\left\|f_{\hat{\bm{a}}}-f_{*}\right\|_{L^{ 1}(\hat{D})}$ $\displaystyle\leq\sup_{\bm{a}\,:\,\left\|\bm{a}\right\|\leq\left\|\bm{a}_{*} \right\|}\left(\left\|f_{\bm{a}}-f_{*}\right\|_{L^{1}(D)}-\left\|f_{\bm{a}}-f_ {*}\right\|_{L^{1}(\hat{D})}\right)+\left\|f_{\bm{a}_{*}}-f_{*}\right\|_{L^{2} (\hat{D})},$ $$

where we used the fact that $\left\|f_{\hat{\bm{a}}}-f_{*}\right\|_{L^{1}(\hat{D})}\leq\left\|f_{\hat{\bm{a
}}}-f_{*}\right\|_{L^{2}(\hat{D})}\leq\left\|f_{\bm{a}_{*}}-f_{*}\right\|_{L^{
1}(\hat{D})}$ in the last line.

Now, we bound the first term. Let $\bm{\sigma}:=(\sigma_{n})_{n=1}^{N}$ be i.i.d. Rademacher variables that are also independent of
everything else. By symmetrization and Theorem 7 of Meir and Zhang [], we have

$$ $\displaystyle\mathop{\mathbb{E\/}}_{\bm{X}}\left[\sup_{\bm{a}\,:\,\left\|\bm{a }\right\|\leq\left\|\bm{a}_{*}\right\|}\left(\left\|f_{\bm{a}}-f_{*}\right\|_{ L^{1}(D)}-\left\|f_{\bm{a}}-f_{*}\right\|_{L^{1}(\hat{D})}\right)\right]$ $\displaystyle\leq\;$ $\displaystyle 2\mathop{\mathbb{E\/}}_{\bm{X},\bm{\sigma}}\sup_{\bm{a}\,:\, \left\|\bm{a}\right\|\leq\left\|\bm{a}_{*}\right\|}\frac{1}{N}\sum_{t=1}^{N} \sigma_{t}\left|f_{a}({\bm{x}}_{T+n})-f_{*}({\bm{x}}_{T+n})\right|$ $\displaystyle\leq\;$ $\displaystyle 2\mathop{\mathbb{E\/}}_{\bm{X},\bm{\sigma}}\sup_{\bm{a}\,:\, \left\|\bm{a}\right\|\leq\left\|\bm{a}_{*}\right\|}\frac{1}{N}\sum_{t=1}^{N} \sigma_{t}\left(f_{a}({\bm{x}}_{T+n})-f_{*}({\bm{x}}_{T+n})\right)$ $\displaystyle\leq\;$ $\displaystyle\frac{2}{N}\mathop{\mathbb{E\/}}_{\bm{X},\bm{\sigma}}\sup_{\bm{a} \,:\,\left\|\bm{a}\right\|\leq\left\|\bm{a}_{*}\right\|}\sum_{t=1}^{N}\sigma_{ t}f_{a}({\bm{x}}_{T+n})+\cancelto{0}{2\mathop{\mathbb{E\/}}_{\bm{X},\bm{\sigma }}\frac{1}{N}\sum_{t=1}^{N}\sigma_{t}f_{*}({\bm{x}}_{T+n})}.$ $$

Note that the first term is two times the Rademacher complexity $\mathrm{Rad}_{N}(\mathcal{F})$ of $\mathcal{F}$ (see, for example,
Chapter 4 of Wainwright []). By (the proof of) Lemma 48 of Damian et al. [], we have

$$ $\displaystyle\mathrm{Rad}_{N}(\mathcal{F})\leq\frac{\left\|\bm{a}_{*}\right\|} {\sqrt{N}}\sqrt{\mathop{\mathbb{E\/}}_{{\bm{x}}\sim\mathcal{N}\left(0,\bm{I}_{ d}\right)}\left\|\phi(\bm{V}{\bm{x}})\right\|^{2}}$ $\displaystyle=\frac{\left\|\bm{a}_{*}\right\|}{\sqrt{N}}\sqrt{\sum_{k=1}^{m} \mathop{\mathbb{E\/}}_{{\bm{x}}\sim\mathcal{N}\left(0,\bm{I}_{d}\right)}\phi^{ 2}({\bm{v}}_{k}\cdot{\bm{x}})}$ $\displaystyle=\frac{\left\|\bm{a}_{*}\right\|\sqrt{m}}{\sqrt{N}}\sqrt{\mathop{ \mathbb{E\/}}_{x_{1}\sim\mathcal{N}\left(0,1\right)}\phi^{2}(x_{1})}$ $\displaystyle=\frac{2\left\|\bm{a}_{*}\right\|\sqrt{m}}{\sqrt{N}}.$ $$

In other words, we have

$$ $\mathop{\mathbb{E\/}}\sup_{\bm{a}\,:\,\left\|\bm{a}\right\|\leq\left\|\bm{a}_{ *}\right\|}\left(\left\|f_{\bm{a}}-f_{*}\right\|_{L^{1}(D)}-\left\|f_{\bm{a}}- f_{*}\right\|_{L^{1}(\hat{D})}\right)\leq\frac{4\left\|\bm{a}_{*}\right\|\sqrt {m}}{\sqrt{N}}.$ $$

Hence, for any $\delta_{\mathop{\mathbb{P\/}}}\in(0,1)$, by Markov’s inequality, we have

$$ $\sup_{\bm{a}\,:\,\left\|\bm{a}\right\|\leq\left\|\bm{a}_{*}\right\|}\left( \left\|f_{\bm{a}}-f_{*}\right\|_{L^{1}(D)}-\left\|f_{\bm{a}}-f_{*}\right\|_{L^ {1}(\hat{D})}\right)\leq\frac{4\left\|\bm{a}_{*}\right\|\sqrt{m}}{\sqrt{N} \delta_{\mathop{\mathbb{P\/}}}},$ $$

with probability at least $1-\delta_{\mathop{\mathbb{P\/}}}$.
Apply the same argument to $\left\|f_{\bm{a}_{*}}-f_{*}\right\|_{L^{2}(\hat{D})}$ and recall from
Lemma D.1 that $\left\|f_{\bm{a}_{*}}-f_{*}\right\|_{L^{2}(D)}^{2}\leq 10LP^{2}\varepsilon_{v}$, and
we obtain

$$ $\left\|f_{\hat{\bm{a}}}-f_{*}\right\|_{L^{1}(D)}\leq\frac{8\left\|\bm{a}_{*} \right\|\sqrt{m}}{\sqrt{N}\delta_{\mathop{\mathbb{P\/}}}}+\sqrt{10LP^{2} \varepsilon_{v}},$ $$

with probability at least $1-2\delta_{\mathop{\mathbb{P\/}}}$.
∎

###### Lemma D.1 .

###### Proof.

###### Lemma D.2 .

###### Proof.

## Appendix E Proof of the main theorem

See 2.1

###### Proof.

First, by Lemma B.3, we should choose
$m=400P^{8}\log^{1.5}\left(P\vee 1/\delta_{\mathop{\mathbb{P\/}}}\right)$. Meanwhile, by Lemma D.2, to achieve target $L^{1}$-error $\varepsilon_{*}$ with probability at least $1-O(\delta_{\mathop{\mathbb{P\/}}})$, we need

$$ $N\gtrsim\frac{Pm}{\varepsilon_{*}^{2}\delta_{\mathop{\mathbb{P\/}}}^{2}},\quad \varepsilon_{v}=O_{L}\left(\frac{\varepsilon_{*}^{2}}{P^{2}}\right).$ $$

Then, to meet the conditions of Lemma C.2 and Lemma C.10 (uniformly over those $P$ good neurons), we choose

$$ $a_{0}=O_{L}\left(\frac{\varepsilon_{*}^{2}}{mdP^{2L+2}\log^{3}d\log(1/ \varepsilon_{*})}\right),\quad\eta=O_{L}\left(\frac{\varepsilon_{*}^{4}\delta_ {\mathop{\mathbb{P\/}}}}{dP^{L+8}\log^{4L+1}(d/\delta_{\mathop{\mathbb{P\/}}}) }\right).$ $$

By Lemma C.2 and Lemma C.10,
the numbers of iterations needed for Stage 1.1 and Stage 1.2 are $O_{L}(\log(d/P)/\eta)$ and
$O_{L}\left(\left(P^{L-1}+\log(1/\varepsilon_{v})\right)/\eta\right)$, respectively. Thus, the total number of
iterations is bounded by

$$ $T=O_{L}\left(\frac{\log d+P^{L-1}+\log(P/\varepsilon_{*})}{\eta}\right)=\tilde {O}_{L}\left(\frac{d\operatorname{poly}(P)}{\varepsilon_{*}^{4}\delta_{\mathop {\mathbb{P\/}}}}\right).$ $$

∎

###### Proof.

## Appendix F Technical lemmas

### F.1 Concentration and anti-concentration of Gaussian variables

In this subsection, we first present several concentration and anti-concentration results for Gaussian variables.
While almost all of them have been proved in the past in different papers and textbooks such as
(van Handel [], Wainwright []), we provide proofs of most of them for
easier reference.

###### Lemma F.1 (Concentration of norm) .

Let $\bm{Z}\sim\mathcal{N}\left(0,\bm{I}_{d}\right)$. Then, we have

$$ $\mathop{\mathbb{P\/}}\left(\left|\left\|\bm{Z}\right\|-\mathop{\mathbb{E\/}} \left\|\bm{Z}\right\|\right|\geq s\right)\leq 2e^{-s^{2}/2}.$ $$

###### Remark .

$\left\|\bm{Z}\right\|$ follows the chi distribution $\chi_{d}$, whose expectation is
$\sqrt{2}\Gamma((d+1)/2)/\Gamma(d/2)$. With Stirling’s formula, one can show that for any large $d$,

$$ $\sqrt{d}\geq\mathop{\mathbb{E\/}}\left\|\bm{Z}\right\|=\sqrt{d-1}\left(1-\frac {1}{4d}+\frac{O(1)}{d^{2}}\right)=\sqrt{d}\left(1-\frac{2}{d}\right).$ $$

∎

###### Proof.

We will use without proof the following result: if $\bm{Z}\sim\mathcal{N}\left(0,\bm{I}_{d}\right)$ and $f:\mathbb{R}^{d}\to\mathbb{R}$ is $1$-Lipschitz,
then $f(\bm{Z})$ is $1$-subgaussian. We apply this result to the $1$-Lipschitz function $\left\|\cdot\right\|$.
This gives
$\mathop{\mathbb{P\/}}\left(\left\|\bm{Z}\right\|-\mathop{\mathbb{E\/}}\left\|
\bm{Z}\right\|\geq s\right)\leq e^{-s^{2}/2}.$
Apply the same result to $-\left\|\cdot\right\|$ yields the lower tails.
∎

###### Lemma F.2 (Upper tail for the maximum) .

Let $Z_{1},\dots,Z_{d}\sim\mathcal{N}\left(0,1\right)$ be independent. We have the upper tail

$$ $\mathop{\mathbb{P\/}}\left(\max_{i\in[d]}|Z_{i}|\geq\sqrt{2\log d}+s\right) \leq 2e^{-s^{2}/2},\quad\forall s\geq 0.$ $$

###### Proof.

For notational simplicity, put $Z^{*}=\max_{i\in[d]}Z_{i}$.
By union bound and the Chernoff bound, we have for each $s,\theta>0$,

$$ $\mathop{\mathbb{P\/}}(Z^{*}\geq s)=\mathop{\mathbb{P\/}}\left(\bigvee_{i=1}^{d }Z_{i}\geq s\right)\leq d\mathop{\mathbb{P\/}}(Z_{1}\geq s)\leq d\frac{\mathop {\mathbb{E\/}}e^{\theta Z_{1}}}{e^{\theta s}}=de^{\theta^{2}/2-\theta s}.$ $$

Choose $\theta=s$ to minimize the RHS, and we obtain
$\mathop{\mathbb{P\/}}(Z^{*}\geq s)\leq e^{\log d-s^{2}/2}.$
Replace $s$ with $\sqrt{2\log d+s^{2}}$ and this becomes

$$ $\mathop{\mathbb{P\/}}\left(Z^{*}\geq\sqrt{2\log d}+s\right)\leq\mathop{\mathbb {P\/}}\left(Z^{*}\geq\sqrt{2\log d+s^{2}}\right)\leq e^{-s^{2}/2}.$ $$

Use the fact $-\min_{i\in[d]}Z_{i}\overset{d}{=}\max_{i\in[d]}Z_{i}$ and we complete the proof.
∎

###### Lemma F.3 (Lower tail for the maximum) .

Let $Z_{1},\dots,Z_{d}\sim\mathcal{N}\left(0,1\right)$ be independent. Let $c>0$ be any universal constant. We have

$$ $\mathop{\mathbb{P\/}}\left[\max_{i\in[d]}Z_{i}\geq(1+c)\sqrt{2\log d}\right] \geq\frac{1}{8{\pi}(1+c)}\frac{1}{d^{(1+c)^{2}-1}\sqrt{\log d}}.$ $$

###### Proof.

First, we prove a general result on the integral $I(x)=\int_{x}^{\infty}e^{-y^{2}/2}\mathrm{d}y$. Make the change of variable
$y=x\tau$ to obtain $I(x)=x\int_{1}^{\infty}e^{-x^{2}\tau^{2}/2}\mathrm{d}\tau$. Since the integrand decays very fast as
$\tau$ grows, we expand $\tau^{2}/2$ around as $\tau^{2}/2=1/2+(\tau-1)+(\tau-1)^{2}/2$. This gives

$$ $I(x)=xe^{-x^{2}/2}\int_{1}^{\infty}e^{-x^{2}(\tau-1)}e^{-x^{2}(\tau-1)^{2}/2} \mathrm{d}\tau=xe^{-x^{2}/2}\int_{0}^{\infty}e^{-x^{2}\tau}e^{-x^{2}\tau^{2}/2 }\mathrm{d}\tau$ $$

For the second factor, we have

$$ $\displaystyle\int_{0}^{\infty}e^{-x^{2}\tau}e^{-x^{2}\tau^{2}/2}\mathrm{d}\tau$ $\displaystyle\leq\int_{0}^{\infty}e^{-x^{2}\tau}\mathrm{d}\tau=\frac{1}{x^{2}},$ $\displaystyle\int_{0}^{\infty}e^{-x^{2}\tau}e^{-x^{2}\tau^{2}/2}\mathrm{d}\tau$ $\displaystyle\geq\int_{0}^{\infty}e^{-x^{2}\tau}\left(1-\frac{x^{2}\tau^{2}}{2 }\right)\mathrm{d}\tau=\frac{1}{x^{2}}\left(1-\frac{1}{x^{2}}\right).$ $$

Combining these bounds together, we obtain

$$ $\frac{e^{-x^{2}/2}}{x}\left(1-\frac{1}{x^{2}}\right)\leq I(x)\leq\frac{e^{-x^{ 2}/2}}{x}.$ (14) $$

With this estimation, we are ready to prove this lemma. Let $c>0$ be a constant.
Note that by our previous tail bound, $\max_{i\in[d]}Z_{i}\geq(1+c)\sqrt{2\log d}=:\theta$ is a rare event.
We have

$$ $\displaystyle\mathop{\mathbb{P\/}}\left[\max_{i\in[d]}Z_{i}\geq\theta\right]=1 -\left(1-\frac{I(\theta)}{\sqrt{2\pi}}\right)^{d}$ $\displaystyle\geq\frac{d}{2}\frac{I(\theta)}{\sqrt{2\pi}}$ $\displaystyle\geq\frac{d}{4\sqrt{2\pi}}\frac{e^{-\theta^{2}/2}}{\theta}=\frac{ 1}{8{\pi}(1+c)}\frac{1}{d^{(1+c)^{2}-1}\sqrt{\log d}}.$ $$

∎

###### Lemma F.4 (Gap between the largest and the second largest) .

Let $Z_{1},\dots,Z_{d}\sim\mathcal{N}\left(0,1\right)$ be independent. Consider an arbitrary universal constant
$c\geq 1/\sqrt{2}$. Define the good and bad events as

$$ $\displaystyle G$ $\displaystyle:=\left\{\max_{i\in[d]}|Z_{i}|\geq(1+2c)\sqrt{2\log d}\right\},$ $\displaystyle B$ $\displaystyle:=\left\{\exists i\neq j\in[d],\,\min\{|Z_{i}|,|Z_{j}|\}\geq(1+c) \sqrt{2\log d}\right\}.$ $$

We have

$$ $\frac{\mathop{\mathbb{P\/}}(B)}{\mathop{\mathbb{P\/}}(G)}\leq\frac{8\pi(1+2c) \sqrt{\log d}}{d^{1-2c^{2}}}\to 0\quad\text{as $d\to\infty$}.$ $$

Let $|Z|_{(1)}$ and $|Z|_{(2)}$ be the largest and second-largest among $|Z_{1}|,\dots,|Z_{d}|$. We have

$$ $\mathop{\mathbb{P\/}}\left[\frac{|Z|_{(1)}}{|Z|_{(2)}}\geq\frac{1+2c}{1+c} \right]\geq\mathop{\mathbb{P\/}}\left[G\wedge\neg B\right]\geq(1-o(1))\mathop{ \mathbb{P\/}}(G)\geq\frac{1}{5\pi(1+2c)}\frac{1}{d^{4c+4c^{2}}\sqrt{\log d}}.$ $$

###### Proof.

Let $0<c_{1}<c_{2}$ be two universal constants to be determined later. By Lemma F.3, we have

$$ $\mathop{\mathbb{P\/}}(G):=2\mathop{\mathbb{P\/}}\left[\max_{i\in[d]}Z_{i}\geq( 1+c_{2})\sqrt{2\log d}\right]\geq\frac{1}{4\pi(1+c_{2})}\frac{1}{d^{(1+c_{2})^ {2}-1}\sqrt{\log d}}.$ $$

Meanwhile, we have

$$ $\displaystyle\mathop{\mathbb{P\/}}(B)$ $\displaystyle:=\mathop{\mathbb{P\/}}\left[\exists i\neq j\in[d],\,\min\{|Z_{i} |,|Z_{j}|\}\geq(1+c_{1})\sqrt{2\log d}\right]$ $\displaystyle\leq 2\binom{d}{2}\left(\mathop{\mathbb{P\/}}\left[Z_{1}\geq(1+c_ {1})\sqrt{2\log d}\right]\right)^{2}$ $\displaystyle\leq d^{2}\exp\left(-2(1+c_{1})^{2}\log d\right)$ $\displaystyle=d^{-2(1+c_{1})^{2}+2}.$ $$

Combine these bounds together, we obtain

$$ $\frac{\mathop{\mathbb{P\/}}(B)}{\mathop{\mathbb{P\/}}(G)}\leq\frac{4\pi(1+c_{2 })d^{(1+c_{2})^{2}-1}\sqrt{\log d}}{d^{2(1+c_{1})^{2}-2}}=\frac{4\pi(1+c_{2}) \sqrt{\log d}}{d^{2(1+c_{1})^{2}-1-(1+c_{2})^{2}}}.$ $$

Suppose that $c_{1}^{2}=c^{2}>1/2$ and choose $c_{2}=2c_{1}$. Then, the above becomes

$$ $\frac{\mathop{\mathbb{P\/}}(B)}{\mathop{\mathbb{P\/}}(G)}\leq\frac{4\pi(1+2c) \sqrt{\log d}}{d^{1-2c^{2}}}.$ $$

∎

###### Lemma F.1 (Concentration of norm) .

###### Remark .

###### Proof.

###### Lemma F.2 (Upper tail for the maximum) .

###### Proof.

###### Lemma F.3 (Lower tail for the maximum) .

###### Proof.

###### Lemma F.4 (Gap between the largest and the second largest) .

###### Proof.

### F.2 Stochastic induction

Our proof is essentially a large induction. When certain properties hold, we know how to analyze the dynamics and
can show certain quantities are bounded with high probability. Meanwhile, certain properties hold as long as those
quantities are still well-controlled. In the deterministic setting, this seemingly looped argument can be made
formal by either mathematical induction (in discrete time) or the continuity argument (in continuous time).
In this subsection, we show the same can also be done in the presence of randomness and derive a stochastic
version of Gronwall’s lemma and its generalizations.

We start with an example where Doob’s submartingale inequality can be directly used.
Let $(\Omega,\mathcal{F},(\mathcal{F}_{t})_{t},\mathop{\mathbb{P\/}})$ be our filtered probability space and $(Z_{t})_{t}$ be a martingale difference sequence.
Suppose that $\mathop{\mathbb{E\/}}[Z_{t+1}^{2}\mid\mathcal{F}_{t}]$ is uniformly bounded by $\sigma_{Z}^{2}$. Then, by Doob’s submartingale inequality,
for any $M>0$ and $T>0$, we have

$$ $\mathop{\mathbb{P\/}}\left[\sup_{t\leq T}\left|\sum_{s=1}^{t}Z_{s}\right|\geq M \right]\leq M^{-2}\mathop{\mathbb{E\/}}\left(\sum_{s=1}^{T}Z_{s}\right)^{2}= \frac{T\sigma_{Z}^{2}}{M^{2}}.$ $$

In particular, this implies that when $M=\omega(\sigma_{Z}\sqrt{T})$, we have
$\sup_{t\leq T}\left|\sum_{s=1}^{t}Z_{s}\right|\leq M$ with high probability.

Note that there is no need to any kind of “induction” in the above example. However, things become subtle if
instead of assuming $\mathop{\mathbb{E\/}}[Z_{t+1}^{2}\mid\mathcal{F}_{t}]$ is bounded by $\sigma_{Z}^{2}$, we assume it is bounded by
$\sigma_{Z}^{2}$ as long as $\sup_{s\leq t}\left|\sum_{r=1}^{s}Z_{r}\right|\leq M$. Intuitively, since $M$ is chosen
so that $\sup_{t\leq T}\left|\sum_{s=1}^{t}Z_{s}\right|\leq M$ holds with high probability, the bounds
$\mathop{\mathbb{E\/}}[Z_{t+1}^{2}\mid\mathcal{F}_{t}]\leq\sigma_{Z}^{2}$ should also hold with high probability and we can still use
Doob’s submartingale inequality as before. Now, we formalize this argument.

###### Lemma F.5 .

Let $(Z_{t})_{t}$ be a martingale difference sequence. Suppose that there exists $M,\sigma_{Z}>0$ such that
if $\sup_{s\leq t}\left|\sum_{r=1}^{s}Z_{s}\right|\leq M$, then we have $\mathop{\mathbb{E\/}}[Z_{t+1}^{2}\mid\mathcal{F}_{t}]\leq\sigma_{Z}^{2}$.
Then, we have

$$ $\mathop{\mathbb{P\/}}\left[\sup_{t\leq T}\left|\sum_{s=1}^{t}Z_{s}\right|>M \right]\leq\frac{T\sigma_{Z}^{2}}{M^{2}}.$ $$

Note that this bound is the same as the one we obtained with the assumption that
$\mathop{\mathbb{E\/}}[Z_{t+1}^{2}\mid\mathcal{F}_{t}]\leq\sigma_{Z}^{2}$ always holds.

###### Proof.

Consider the stopping time $\tau:=\inf\{t\geq 0\;:\;\left|\sum_{s=1}^{t}Z_{s}\right|>M\}$.
By definition, we have $\sup_{s\leq t}\left|\sum_{r=1}^{s}Z_{s}\right|\leq M$ for all $t\leq\tau$.
Then, we define $Y_{t+1}=Z_{t+1}\mathbbm{1}\{t<\tau\}$. Note that $(Y_{t})$ is a martingale difference sequence
with $\mathop{\mathbb{E\/}}[Y_{t+1}^{2}\mid\mathcal{F}_{t}]\leq\sigma_{Z}^{2}$. As a result, by Doob’s submartingale inequality, we have
$\mathop{\mathbb{P\/}}\left[\sup_{t\leq T}\left|\sum_{s=1}^{t}Y_{s}\right|>M
\right]\leq T\sigma_{Z}^{2}/M^{2}.$
To relate it to $(Z_{t})_{t}$, we compute

$$ $\displaystyle\mathop{\mathbb{P\/}}\left[\sup_{t\leq T}\left|\sum_{s=1}^{t}Z_{s }\right|>M\right]=\mathop{\mathbb{P\/}}\left[\sup_{t\leq T}\left|\sum_{s=1}^{t }Z_{s}\right|>M\wedge\tau\leq T\right]$ $\displaystyle=\mathop{\mathbb{P\/}}\left[\left|\sum_{s=1}^{\tau}Z_{s}\right|>M \wedge\tau\leq T\right]$ $\displaystyle=\mathop{\mathbb{P\/}}\left[\left|\sum_{s=1}^{\tau}Y_{s}\right|>M \wedge\tau\leq T\right]$ $\displaystyle\leq\frac{T\sigma_{Z}^{2}}{M^{2}},$ $$

where the first and second identities comes from the definition of $\tau$ and the third from the fact $Z_{t}=Y_{t}$
for all $t\leq\tau$.
∎

Now, we consider a more complicated case, where is process of interest is not a pure martingale. Suppose that the process $(X_{t})_{t}$ satisfies

$$ $X_{t+1}=(1+\alpha)X_{t}+\xi_{t+1}+Z_{t+1},\quad X_{0}=x_{0}>0,$ $$

where the signal growth rate $\alpha>0$ and initialization $x_{0}>0$ are given and fixed, $(\xi_{t})_{t}$ is an
adapted process, and $(Z_{t})_{t}$ is a martingale difference sequence. In most cases, $(\xi_{t})_{t}$ will represent the
higher-order error terms.

Our goal is control the difference between $X_{t}$ and its deterministic counterpart
$x_{t}=(1+\alpha)^{t}x_{0}$.
To this end, we recursively expand the RHS to obtain

$$ $\displaystyle X_{t+1}$ $\displaystyle=(1+\alpha)^{2}X_{t-1}+(1+\alpha)\xi_{t}+\xi_{t+1}+(1+\alpha)Z_{t }+Z_{t+1}$ $\displaystyle=(1+\alpha)^{t+1}x_{0}+\sum_{s=1}^{t}(1+\alpha)^{t-s}\xi_{s+1}+ \sum_{s=1}^{t}(1+\alpha)^{t-s}Z_{s+1}.$ $$

Divide both sides with $(1+\alpha)^{t+1}$ and replace $t+1$ with $t$. Then, the above becomes

$$ $X_{t}(1+\alpha)^{-t}=x_{0}+\sum_{s=1}^{t}(1+\alpha)^{-s}\xi_{s}+\sum_{s=1}^{t} (1+\alpha)^{-s}Z_{s}.$ $$

Note that $\left((1+\alpha)^{-t}Z_{t}\right)_{t}$ is still a martingale difference sequence. Ideally, $|\xi_{t}|$ should
be small as it represents the higher-order error terms, and we have bounds on the conditional variance of
$Z_{t}$ so that we can apply Doob’s submartingale inequality to the last term. Unfortunately, in many cases,
since $\xi_{t+1}$ and $Z_{t+1}$, particularly their maximum and (conditional) variance, can potentially depend on
$(X_{s})_{s\leq t}$, we may only be able to assume $|\xi_{t+1}|\leq(1+\alpha)^{t}\Xi$ with probability at least $1-\delta_{\mathop{\mathbb{P\/}},\xi}$
(for each $t$) and $\mathop{\mathbb{E\/}}[Z_{t+1}^{2}\mid\mathcal{F}_{t}]\leq(1+\alpha)^{t}\sigma_
{Z}^{2}$ for some $\xi_{\mathop{\mathbb{P\/}},\xi}$, $\Xi$ and $\sigma_{Z}^{2}$ when, say, $X_{t}=(1\pm 0.5)x_{t}$. Still, we can use the previous argument to
estimate the probability that $X_{t}\notin(1\pm 0.5)x_{t}$ for some $t\leq T$.

Let $\tau:=\inf\{t\geq 0\;:\;X_{t}\notin(1\pm 0.5)x_{t}\}$ and then
$\hat{\xi}_{t+1}=\xi_{t+1}\mathbbm{1}\{t\leq\tau\}$, and $\hat{Z}_{t+1}=Z_{t+1}\mathbbm{1}\{t\leq\tau\}$. Clear that
$\tau$ is a stopping time, $\hat{\xi}$ is adapted, and $\hat{Z}$ is still a martingale difference sequence. Moreover,
we have $|\hat{\xi}_{t}|\leq(1+\alpha)^{t}\Xi$ with probability at least $1-\delta_{\mathop{\mathbb{P\/}},\xi}$ and
$\mathop{\mathbb{E\/}}\left[\hat{Z}_{t+1}^{2}\mid\mathcal{F}_{t}\right]\leq(1+
\alpha)^{t}\sigma_{Z}^{2}$ for all $t\geq 0$.
As a result,

$$ $\displaystyle\left|\sum_{s=1}^{t}(1+\alpha)^{-s}\hat{\xi}_{s}\right|$ $\displaystyle\leq\Xi t\leq T\Xi\quad\text{with probability at least $1-T\delta _{\mathop{\mathbb{P\/}},\xi}$},$ $\displaystyle\mathop{\mathbb{E\/}}\left(\sum_{s=1}^{t}(1+\alpha)^{-s}\hat{Z}_{ s}\right)^{2}$ $\displaystyle=\sum_{s=1}^{t}(1+\alpha)^{-2s}\mathop{\mathbb{E\/}}\mathop{ \mathbb{E\/}}\left[\hat{Z}_{s}^{2}\;\mid\;\mathcal{F}_{s-1}\right]\leq\sum_{s= 1}^{t}(1+\alpha)^{-s}\sigma_{Z}^{2}\leq\frac{\sigma_{Z}^{2}}{\alpha}.$ $$

Then, by Doob’s submartingale inequality, we have

$$ $\mathop{\mathbb{P\/}}\left[\sup_{t\leq T}\left|\sum_{s=1}^{T}(1+\alpha)^{-s} \hat{Z}_{s}\right|\geq\frac{x_{0}}{4}\right]\leq\frac{16\sigma_{Z}^{2}}{\alpha x _{0}^{2}}.$ $$

Hence, for any $\delta_{\mathop{\mathbb{P\/}}}\in(0,1)$, if we assume

$$ $\Xi\leq\frac{x_{0}}{4T}\quad\text{and}\quad\sigma_{Z}^{2}\leq\frac{\delta_{ \mathop{\mathbb{P\/}}}\alpha x_{0}^{2}}{16},$ $$

then with probability at least $1-\delta_{\mathop{\mathbb{P\/}}}-T\delta_{\mathop{\mathbb{P\/}},\xi}$, we have

$$ $\left|\sum_{s=1}^{t}(1+\alpha)^{-s}\hat{\xi}_{s}+\sum_{s=1}^{t}(1+\alpha)^{-s} \hat{Z}_{s}\right|\leq\frac{x_{0}}{2},\quad\forall t\in[T].$ $$

Then, similar to the previous argument, we have

$$ $\displaystyle\mathop{\mathbb{P\/}}\left[\exists t\in[T],\,X_{t}\notin(1\pm 0.5 )x_{t}\right]$ $\displaystyle=\mathop{\mathbb{P\/}}\left[\exists t\in[T],\,X_{t}\notin(1\pm 0. 5)x_{t}\wedge\tau\leq T\right]$ $\displaystyle=\mathop{\mathbb{P\/}}\left[X_{\tau}\notin(1\pm 0.5)x_{\tau} \wedge\tau\leq T\right]$ $\displaystyle=\mathop{\mathbb{P\/}}\left[\left|\sum_{s=1}^{\tau}(1+\alpha)^{-s }\xi_{s}+\sum_{s=1}^{\tau}(1+\alpha)^{-s}Z_{s}\right|\geq\frac{x_{0}}{2}\wedge \tau\leq T\right]$ $\displaystyle=\mathop{\mathbb{P\/}}\left[\left|\sum_{s=1}^{T}(1+\alpha)^{-s} \hat{\xi}_{s}+\sum_{s=1}^{T}(1+\alpha)^{-s}\hat{Z}_{s}\right|\geq\frac{x_{0}}{ 2}\wedge\tau\leq T\right]$ $\displaystyle\leq 1-\delta_{\mathop{\mathbb{P\/}}}-T\delta_{\mathop{\mathbb{P \/}},\xi}.$ $$

Namely, we have proved the following discrete-time stochastic Gronwall’s lemma.

###### Lemma F.6 (Stochastic Gronwall’s lemma) .

Suppose that $(X_{t})_{t}$ satisfies

$$ $X_{t+1}=(1+\alpha)X_{t}+\xi_{t+1}+Z_{t+1},\quad X_{0}=x_{0}>0,$ $$

where the signal growth rate $\alpha>0$ and initialization $x_{0}>0$ are given and fixed, $(\xi_{t})_{t}$ is an
adapted process, and $(Z_{t})_{t}$ is a martingale difference sequence. Define $x_{t}=(1+\alpha)^{t}x_{0}$.

Let $T>0$ and $\delta_{\mathop{\mathbb{P\/}}}\in(0,1)$ be given.
Suppose that there exists some $\delta_{\mathop{\mathbb{P\/}},\xi}\in(0,1)$ and $\Xi,\sigma_{Z}>0$ such that for every $t\geq 0$,
if $X_{t}=(1\pm 0.5)x_{t}$, then we have $|\xi_{t+1}|\leq(1+\alpha)^{t}\Xi$ with probability at least $1-\delta_{\mathop{\mathbb{P\/}},\xi}$
and $\mathop{\mathbb{E\/}}[Z_{t+1}^{2}\mid\mathcal{F}_{t}]\leq(1+\alpha)^{t}\sigma_
{Z}^{2}$. Then, if

$$ $\Xi\leq\frac{x_{0}}{4T}\quad\text{and}\quad\sigma_{Z}^{2}\leq\frac{\delta_{ \mathop{\mathbb{P\/}}}\alpha x_{0}^{2}}{16},$ $$

we have $X_{t}=(1\pm 0.5)x_{t}$ for all $t\in[T]$ with probability at least $1-\delta_{\mathop{\mathbb{P\/}}}-T\delta_{\mathop{\mathbb{P\/}},\xi}$.

###### Remark .

With only the dependence on $\alpha$ and $x_{0}$ kept, then conditions become $\Xi\leq O(\alpha x_{0})$
and $\sigma_{Z}\leq O(\sqrt{\alpha}x_{0})$. When $\alpha$ is small, the second condition is much weaker than the first
one.
∎

###### Remark .

The above argument can be easily generalized to cases where we have multiple induction hypotheses. For example,
if we have another process $X_{t+1}^{\prime}=(1+\alpha^{\prime})X_{t}^{\prime}+\xi_{t+1}^{\prime}+Z_{t+1}^
{\prime}$ and we need both $X_{t}=(1\pm 0.5)x_{t}$ and $X_{t}^{\prime}=(1\pm 0.5)x_{t}^{\prime}$ for the bounds on $|\xi_{t+1}|,|\xi^{\prime}_{t+1}|$, $\mathop{\mathbb{E\/}}[Z_{t+1}^{2}\mid\mathcal{F}_{t}]$,
$\mathop{\mathbb{E\/}}[(Z^{\prime}_{t+1})^{2}\mid\mathcal{F}_{t}]$ to hold. In this case, the final failure probability will be bounded by
$T(\delta_{\mathop{\mathbb{P\/}},\xi}+\delta_{\mathop{\mathbb{P\/}},\xi^{\prime
}})+2\delta_{\mathop{\mathbb{P\/}}}$.
∎

If we are interested only in the upper bound, the above lemma can be used instead. In this lemma, the dependence on the
initial value is more lenient.

###### Lemma F.7 .

Suppose that $(X_{t})_{t}$ satisfies

$$ $X_{t+1}=(1+\alpha)X_{t}+\xi_{t+1}+Z_{t+1},\quad X_{0}=x_{0}>0,$ $$

where the signal growth rate $\alpha>0$ and initialization $x_{0}>0$ are given and fixed, $(\xi_{t})_{t}$ is an
adapted process, and $(Z_{t})_{t}$ is a martingale difference sequence. Define $x_{t}^{+}=(1+\alpha)^{t}x_{0}^{+}$, where
$x_{0}^{+}$ is any value that is at least $x_{0}$.

Let $T>0$ and $\delta_{\mathop{\mathbb{P\/}}}\in(0,1)$ be given.
Suppose that there exists some $\delta_{\mathop{\mathbb{P\/}},\xi}\in(0,1)$ and $\Xi,\sigma_{Z}>0$ such that for every $t\geq 0$,
if $X_{t}=(1\pm 0.5)x_{t}$, then we have $|\xi_{t+1}|\leq(1+\alpha)^{t}\Xi$ with probability at least $1-\delta_{\mathop{\mathbb{P\/}},\xi}$
and $\mathop{\mathbb{E\/}}[Z_{t+1}^{2}\mid\mathcal{F}_{t}]\leq(1+\alpha)^{t}\sigma_
{Z}^{2}$. Then, if

$$ $\Xi\leq\frac{x_{0}^{+}}{4T}\quad\text{and}\quad\sigma_{Z}^{2}\leq\frac{\delta_ {\mathop{\mathbb{P\/}}}\alpha(x_{0}^{+})^{2}}{16},$ $$

we have $X_{t}\leq 2x_{t}^{+}$ for all $t\in[T]$ with probability at least $1-\delta_{\mathop{\mathbb{P\/}}}-T\delta_{\mathop{\mathbb{P\/}},\xi}$.

###### Proof.

Similar to the previous proof, we still have

$$ $X_{t}(1+\alpha)^{-t}=x_{0}+\sum_{s=1}^{t}(1+\alpha)^{-s}\xi_{s}+\sum_{s=1}^{t} (1+\alpha)^{-s}Z_{s}.$ $$

Instead of requiring the last two terms to be bounded by $x_{0}/2$, we can simply require them to be bounded by
$x_{0}^{+}/2$ where $x_{0}^{+}$ is any value that is at least $x_{0}$. Then, to complete the proof, it suffices to repeat
the previous argument.
∎

The above lemmas will be used in Stage 1.1 to estimate the growth rate of the signals. The next lemma considers
the case where $\alpha$ is 0 0 and will be used to show the gap between the largest and the second-largest coordinates
can be preserved during Stage 1.1.

###### Lemma F.8 .

Suppose that $(X_{t})_{t}$ satisfies

$$ $X_{t+1}=X_{t}+\xi_{t+1}+Z_{t+1},\quad X_{0}=x_{0}>0,$ $$

where the signal growth rate $\alpha>0$ and initialization $x_{0}>0$ are given and fixed, $(\xi_{t})_{t}$ is an
adapted process, and $(Z_{t})_{t}$ is a martingale difference sequence.

Let $T>0$ and $\delta_{\mathop{\mathbb{P\/}}}\in(0,1)$ be given.
Suppose that there exists some $\delta_{\mathop{\mathbb{P\/}},\xi}\in(0,1)$ and $\Xi,\sigma_{Z}>0$ such that for every $t\leq T$,
if $|X_{t}-x_{0}|\leq T\Xi+\sqrt{T\sigma_{Z}^{2}/\delta_{\mathop{\mathbb{P\/}}}}$, then $|\xi_{t}|\leq\Xi$ with probability at least
$1-\delta_{\mathop{\mathbb{P\/}},\xi}$ and $\mathop{\mathbb{E\/}}[Z_{t+1}^{2}\mid\mathcal{F}_{t}]\leq\sigma_{Z}^{2}$. Then, we have

$$ $\sup_{t\leq T}\left|X_{t}-x_{0}\right|\leq T\Xi+\sqrt{\frac{T\sigma_{Z}^{2}}{ \delta_{\mathop{\mathbb{P\/}}}}}\quad\text{with probability at least $1-T \delta_{\mathop{\mathbb{P\/}},\xi}-\delta_{\mathop{\mathbb{P\/}}}$}.$ $$

###### Proof.

Recursively expand the RHS, and we obtain

$$ $X_{t}=x_{0}+\sum_{s=1}^{t}\xi_{s}+\sum_{s=1}^{t}Z_{s}.$ $$

Consider the stopping time $\tau:=\inf\left\{t\geq 0\;:\;|X_{t}-x_{0}|>T\Xi+\sqrt{T\sigma_{Z}^{2}/\delta_{
\mathop{\mathbb{P\/}}}}\right\}$.
Define $\hat{\xi}_{t+1}=\mathbbm{1}\{t<\tau\}\xi_{t+1}$ and $\hat{Z}_{t+1}=\mathbbm{1}\{t<\tau\}Z_{t+1}$. Clear that

$$ $\sup_{t\leq T}\left|\sum_{s=1}^{t}\hat{\xi}_{t}\right|\leq T\Xi\quad\text{with probability at least $1-T\delta_{\mathop{\mathbb{P\/}},\xi}$}.$ $$

Meanwhile, by Doob’s submartingale inequality, we have

$$ $\mathop{\mathbb{P\/}}\left[\sup_{t\leq T}\left|\sum_{s=1}^{t}\hat{Z}_{s}\right |\geq M\right]\leq\frac{T\sigma_{Z}^{2}}{M^{2}}.$ $$

In other words,

$$ $\sup_{t\leq T}\left|\sum_{s=1}^{t}\hat{\xi}_{t}+\sum_{s=1}^{t}\hat{Z}_{t} \right|\leq T\Xi+\sqrt{\frac{T\sigma_{Z}^{2}}{\delta_{\mathop{\mathbb{P\/}}}}} \quad\text{with probability at least $1-T\delta_{\mathop{\mathbb{P\/}},\xi}- \delta_{\mathop{\mathbb{P\/}}}$}.$ $$

Finally, we compute

$$ $\displaystyle\mathop{\mathbb{P\/}}\left[\sup_{t\leq T}|X_{t}-x_{0}|>T\Xi+\sqrt {\frac{T\sigma_{Z}^{2}}{\delta_{\mathop{\mathbb{P\/}}}}}\right]$ $\displaystyle=\mathop{\mathbb{P\/}}\left[\sup_{t\leq T}|X_{t}-x_{0}|>T\Xi+ \sqrt{\frac{T\sigma_{Z}^{2}}{\delta_{\mathop{\mathbb{P\/}}}}}\wedge T\geq\tau\right]$ $\displaystyle=\mathop{\mathbb{P\/}}\left[\left|\sum_{s=1}^{\tau}\xi_{t}+\sum_{ s=1}^{\tau}Z_{t}\right|>T\Xi+\sqrt{\frac{T\sigma_{Z}^{2}}{\delta_{\mathop{ \mathbb{P\/}}}}}\wedge T\geq\tau\right]$ $\displaystyle=\mathop{\mathbb{P\/}}\left[\left|\sum_{s=1}^{\tau}\hat{\xi}_{t}+ \sum_{s=1}^{\tau}\hat{Z}_{t}\right|>T\Xi+\sqrt{\frac{T\sigma_{Z}^{2}}{\delta_{ \mathop{\mathbb{P\/}}}}}\wedge T\geq\tau\right]$ $\displaystyle\leq 1-T\delta_{\mathop{\mathbb{P\/}},\xi}-\delta_{\mathop{ \mathbb{P\/}}}.$ $$

∎

The above proofs are all based on Doob’s $L^{2}$-submartingale inequality. In other words, it only uses the information
about the conditional variance, whence the dependence on $\delta_{\mathop{\mathbb{P\/}}}$ is $\sqrt{\delta_{\mathop{\mathbb{P\/}}}}$.
It is possible to get a better dependence (of form $\operatorname{poly}\log(1/\delta_{\mathop{\mathbb{P\/}}})$) if we have a full
tail bound similar to the ones in Lemma 2.2. This can be useful when we need to use the
union bound. To this end, we need the following generalization of Freedman’s inequality. The proof of it is deferred
to the end of this section. In short, we truncate $Z_{t}$ at $M$, apply Freedman’s inequality to the truncated sequence,
and estimate the error introduced by the truncation. This and the next lemmas will not be used in the proof of our
main results. We include them here to explain a possible strategy to improve the dependence on $\delta_{\mathop{\mathbb{P\/}}}$.

###### Lemma F.9 (Freedman’s inequality with unbounded variables) .

Let $(Z_{t})_{t}$ be martingale difference sequence with $\mathop{\mathbb{E\/}}[Z_{t}^{2}\mid\mathcal{F}_{t-1}]\leq\sigma_{Z}^{2}$.
Suppose that $Z_{t}$ satisfies the tail bound

$$ $\mathop{\mathbb{P\/}}\left[|Z_{t}|\geq s\;\mid\;\mathcal{F}_{t-1}\right]\leq a \exp\left(-bs^{c}\right),\quad\forall s>0,$ (15) $$

for some $a\geq 1$ and $b,c\in(0,1]$. Then, there exists a constant $C_{c}$ that may depend on $c$ such that
for any $\delta_{\mathop{\mathbb{P\/}}}\in(0,1)$, we have, with probability at least $1-\delta_{\mathop{\mathbb{P\/}}}$ that

$$ $\left|\sum_{t=1}^{T}Z_{t}\right|\leq C_{c}\sqrt{T\left(\sigma_{Z}^{2}+\frac{1} {b^{2/c}}+\frac{\log^{1/c}\left(\frac{aT}{b\sigma_{Z}\delta_{\mathop{\mathbb{P \/}}}}\right)}{b^{1/c}}\right)\log\left(\frac{1}{\delta_{\mathop{\mathbb{P\/}} }}\right)}.$ $$

###### Remark .

Similar bounds hold for a wider range of parameters. We will only use lemma in the proof of Lemma C.9, where the martingale difference sequence is $(Z_{t})_{t}$ satisfies the tail bound in
Lemma 2.2 (without the $\log m$ introduced by the union bound). In other words,
we have $a=C_{L}$, $b=P^{-1/(2L)}$, $c=1/(2L)$, and $\sigma_{Z}^{2}=C_{L}P^{2}$. In particular, note that both $1/b^{2/c}$ and
$\sigma_{Z}^{2}$ have order $P^{2}$.
∎

With this lemma, we can obtain the following variant of Lemma F.8.
Our goal here is to replace $\sqrt{T\sigma_{Z}^{2}/\delta_{\mathop{\mathbb{P\/}}}}$ with $\sqrt{T\sigma_{Z}^{2}/\operatorname{poly}\log\delta_{\mathop{\mathbb{P\/}}}}$.
The proof is essentially the same as the proof of Lemma F.8, and is therefore
deferred to the end of this section. An example of applying is lemma can be found in the proof of Lemma C.9.

###### Lemma F.10 .

Suppose that $(X_{t})_{t}$ satisfies(^9^99Since we require $b\leq 1$ in (15), we need to “normalize” $Z_{t+1}$ here and use $h_{t}$ to keep its size.)

$$ $X_{t+1}=X_{t}+\xi_{t+1}+h_{t}Z_{t+1},\quad X_{0}=x_{0}>0,$ $$

where the signal growth rate $\alpha>0$ and initialization $x_{0}>0$ are given and fixed, $(\xi_{t})_{t}$, $(h_{t})_{t}$
are adapted processes, and $(Z_{t})_{t}$ is a martingale difference sequence.

Let $T>0$ and $\delta_{\mathop{\mathbb{P\/}}}\in(0,1)$ be given.
Suppose that there exists some $\delta_{\mathop{\mathbb{P\/}},\xi}\in(0,1)$ and $\Xi,\sigma_{Z},h^{*}>0$ such that for every $t\leq T$,
if

$$ $|X_{t}-x_{0}|\leq T\Xi+C_{c}h^{*}\sqrt{T\left(\sigma_{Z}^{2}+\frac{1}{b^{2/c}} +\frac{\log^{1/c}\left(\frac{aT}{b\sigma_{Z}\delta_{\mathop{\mathbb{P\/}}}} \right)}{b^{1/c}}\right)\log\left(\frac{T}{\delta_{\mathop{\mathbb{P\/}}}} \right)},$ (16) $$

then $|\xi_{t}|\leq\Xi$ with probability at least $1-\delta_{\mathop{\mathbb{P\/}},\xi}$, $|h_{t}|\leq h^{*}$,
$\mathop{\mathbb{E\/}}[Z_{t+1}^{2}\mid\mathcal{F}_{t}]\leq\sigma_{Z}^{2}$, and $Z_{t+1}^{2}$ satisfies the tail bound (15).
Then, with probability at least $1-T\delta_{\mathop{\mathbb{P\/}},\xi}-\delta_{\mathop{\mathbb{P\/}}}$, (16) holds for all $t\in[T]$.

Now, we consider the case where the signal grows at a polynomial instead of linear rate. This lemma will be used in
Stage 1.2, where the $2L$-th order terms dominate.

###### Lemma F.11 .

Suppose that $(X_{t})_{t}$ satisfies

$$ $X_{t+1}=X_{t}+\alpha X_{t}^{p}+\xi_{t+1}+Z_{t+1},\quad X_{0}=x_{0}>0,$ (17) $$

where $p>1$, the signal growth rate $\alpha>0$ and initialization $x_{0}>0$ are given and fixed, $(\xi_{t})_{t}$ is
an adapted process, and $(Z_{t})_{t}$ is a martingale difference sequence. Let $\hat{x}_{t}$ be the solution to
the deterministic relationship

$$ $\hat{x}_{t+1}=\hat{x}_{t}+\alpha\hat{x}_{t}^{p},\quad\hat{x}_{0}=x_{0}/2.$ $$

Fix $T>0,\delta_{\mathop{\mathbb{P\/}}}\in(0,1)$. Suppose that there exist $\Xi,\sigma_{Z}>0$ and $\delta_{\mathop{\mathbb{P\/}},\xi}\in(0,1)$
such that when $X_{t}\geq\hat{x}_{t}$, we have $|\xi_{t}|\leq\Xi$ with probability at least $1-\delta_{\mathop{\mathbb{P\/}},\xi}$ and
$\mathop{\mathbb{E\/}}[Z_{t+1}\mid\mathcal{F}_{t}]\leq\sigma_{Z}^{2}$. Then, if

$$ $\Xi\leq\frac{x_{0}}{4T}\quad\text{and}\quad\sigma_{Z}^{2}\leq\frac{x_{0}^{2} \delta_{\mathop{\mathbb{P\/}}}}{16T},$ $$

we have $X_{t}\geq\hat{x}_{t}$ for all $t\leq T$.

###### Proof.

Similar to our previous argument, we can assume w.l.o.g. that the bounds on $|{\bm{x}}_{t}|$ and the conditional variance
of $Z_{t+1}$ always hold.

Note that we can rewrite (17) as $X_{t+1}=X_{t}(1+\alpha X_{t}^{p-1})+\xi_{t}+Z_{t}$ and view it as the linear recurrence relationship in Lemma F.6
with a non-constant growth rate. This suggests defining the counterpart of $(1+\alpha)^{t}$ as

$$ $P_{s,t}:=\begin{cases}\prod_{r=s}^{t-1}(1+\alpha X_{r}^{p-1}),&t>s,\\ 1,&t=s.\end{cases}$ $$

Then, we can inductively write (17) as

$$ $\displaystyle X_{1}$ $\displaystyle=X_{0}\left(1+\alpha X_{0}^{p-1}\right)+\xi_{0}+Z_{0},$ $\displaystyle X_{2}$ $\displaystyle=\left(X_{0}\left(1+\alpha X_{0}^{p-1}\right)+\xi_{0}+Z_{0}\right )\left(1+\alpha X_{1}^{p-1}\right)+\xi_{1}+Z_{1}$ $\displaystyle=X_{0}\left(1+\alpha X_{0}^{p-1}\right)\left(1+\alpha X_{1}^{p-1} \right)+\left(1+\alpha X_{1}^{p-1}\right)\left(\xi_{0}+Z_{0}\right)+\xi_{1}+Z_ {1}$ $\displaystyle=X_{0}P_{0,2}+P_{1,2}\left(\xi_{0}+Z_{0}\right)+\xi_{1}+Z_{1},$ $\displaystyle X_{3}$ $\displaystyle=X_{2}\left(1+\alpha X_{2}^{p-1}\right)+\xi_{2}+Z_{2}$ $\displaystyle=\left(X_{0}P_{0,2}+P_{1,2}\left(\xi_{0}+Z_{0}\right)+\xi_{1}+Z_{ 1}\right)\left(1+\alpha X_{2}^{p-1}\right)+\xi_{2}+Z_{2}$ $\displaystyle=X_{0}P_{0,3}+P_{1,3}\left(\xi_{0}+Z_{0}\right)+P_{2,3}\left(\xi_ {1}+Z_{1}\right)+\xi_{2}+Z_{2}.$ $$

Continue the above expansion, and eventually we obtain

$$ $X_{t}=X_{0}P_{0,t}+\sum_{s=1}^{t}P_{s,t}\left(\xi_{s-1}+Z_{s-1}\right).$ $$

By our induction hypothesis, we have $P_{0,s}\geq 1$. Hence, we can divide both sides with $P_{0,t}$ and then the
above becomes

$$ $P_{0,t}^{-1}X_{t}=X_{0}+\sum_{s=1}^{t}P_{0,t}^{-1}P_{s,t}\left(\xi_{s-1}+Z_{s- 1}\right)=X_{0}+\sum_{s=1}^{t}P_{0,s}^{-1}\xi_{s-1}+\sum_{s=1}^{t}P_{0,s}^{-1} Z_{s-1}.$ $$

For the second term, we have

$$ $\left|\sum_{s=1}^{t}P_{0,s}\xi_{s-1}\right|\leq\sum_{s=1}^{t}P_{0,s}|\xi_{s-1} |\leq T\Xi,$ $$

for all $t\leq T$ with probability at least $1-T\delta_{\mathop{\mathbb{P\/}},\xi}$. By our assumption on $\Xi$, this is
bounded by $x_{0}/4$.
For the last term, by Doob’s submartingale
inequality, for any $M>0$, we have

$$ $\mathop{\mathbb{P\/}}\left[\sup_{r\leq t}\left|\sum_{s=1}^{t}P_{0,s}^{-1}Z_{s- 1}\right|\geq M\right]\leq M^{-2}\sum_{s=1}^{t}\mathop{\mathbb{E\/}}\left[P_{0 ,s}^{-2}Z_{s-1}^{2}\right]\leq\frac{\sigma_{Z}^{2}T}{M^{2}}.$ $$

Choose $M=x_{0}/4$ and the RHS becomes $16\sigma_{Z}^{2}T/x_{0}^{2}$, which is bounded by $\delta_{\mathop{\mathbb{P\/}}}$ by our
assumption on $\sigma_{Z}$. Thus, with probability at least $1-T\delta_{\mathop{\mathbb{P\/}},\xi}-\delta_{\mathop{\mathbb{P\/}}}$, we have
$X_{t}\geq P_{0,t}(x_{0}/2)$ for all $t$. In particular, this implies $X_{t}\geq\hat{x}_{t}$ with at least
the same probability.
∎

The above coupling lemma, when combined with the following estimation on the growth rate of the deterministic
process $\hat{x}_{t}$, gives an upper bound on the time needed for $X_{t}$ to grow from a small value to $\Theta(1)$.

###### Lemma F.12 .

Suppose that $(x_{t})_{t}$ satisfies $x_{t+1}=x_{t}+\alpha x_{t}^{p}$ for some $x_{0}\in(0,1)$ and $p>2$ and
$\alpha\ll 1/p$. Then, we have $x_{t}$ must reach $0.9$ within $O(1/(x_{0}^{p-1}\alpha))$ iterations.

###### Proof.

Consider the continuous-time process $\dot{y}_{\tau}=(1-\delta)y_{\tau}^{p}$ where $y_{0}=x_{0}$ and $\delta>0$
is a parameter to be determined later. For $y$, we have the closed-form formula

$$ $y_{\tau}=\left(\frac{1}{x_{0}^{p-1}}-(p-1)(1-\delta)\tau\right)^{-1/(p-1)}.$ $$

Now, we show by induction that $x_{t}\geq y_{t\alpha}$. Clear that this holds when $t=0$. In addition, we have

$$ $x_{t+1}-y_{(y+1)\alpha}=x_{t}-y_{t}+\int_{0}^{\alpha}\left(x_{t}^{p}-(1-\delta )y_{t\alpha+\beta}^{p}\right)\mathrm{d}\beta.$ $$

Note that since $x_{t}\geq y_{t\alpha}$ and $y_{t\alpha+\beta}\leq y_{(t+1)\alpha}$, it suffices to ensure
$y_{t\alpha}\geq(1-\delta)y_{(t+1)\alpha}$. By our closed-form formula for $y_{\tau}$, we have

$$ $\displaystyle y_{t\alpha}\geq(1-\delta)y_{(t+1)\alpha}$ $\displaystyle\quad\Leftrightarrow\quad$ $\displaystyle\frac{1}{x_{0}^{p-1}}-(p-1)(1-\delta)t\alpha\leq(1-\delta)^{1-p} \left(\frac{1}{x_{0}^{p-1}}-(p-1)(1-\delta)(t+1)\alpha\right)$ $\displaystyle\quad\Leftrightarrow\quad$ $\displaystyle(1-\delta)^{p-1}\leq 1-\frac{(p-1)(1-\delta)\alpha}{\frac{1}{x_{0 }^{p-1}}-(p-1)(1-\delta)t\alpha}.$ $$

We are interested in the regime where $\frac{1}{x_{0}^{p-1}}-(p-1)(1-\delta)t\alpha\geq c_{p}$ for some
small constant $c_{p}>0$ that may depend on $p$. In this regime, we have

$$ $\frac{(p-1)(1-\delta)\alpha}{\frac{1}{x_{0}^{p-1}}-(p-1)(1-\delta)t\alpha}\leq c _{p}p\alpha.$ $$

As a result, if $c_{p}p\alpha\leq 0.1$, then
in order for $y_{t\alpha}\geq(1-\delta)y_{(t+1)\alpha}$ in this regime, it suffices to choose

$$ $\displaystyle(1-\delta)^{p-1}\leq 1-c_{p}p\alpha$ $\displaystyle\quad\Leftarrow\quad(1-\delta)^{p-1}\leq e^{-2c_{p}p\alpha}$ $\displaystyle\quad\Leftarrow\quad 1-\delta\leq e^{-4c_{p}\alpha}\quad \Leftarrow\quad\delta\geq 8c_{p}\alpha.$ $$

Let $1$ be our target value for $x_{t}$. To reach $C_{*}$, we need $\frac{1}{x_{0}^{p-1}}-(p-1)t\alpha\leq 1$.
Choose $c_{p}=1$. Then the above implies that $x_{t}\geq y_{t\alpha}$ with $\dot{y}_{\tau}=(1-8\alpha)y_{\tau}^{p}$
when $x_{t}\leq 1$. Combine this with the closed formula for $y_{\tau}$, and we conclude that $x_{\tau}$ must reach
$1/2$ within $O(1/(x_{0}^{p-1}\alpha))$ iterations.
∎

###### Lemma F.5 .

###### Proof.

###### Lemma F.6 (Stochastic Gronwall’s lemma) .

###### Remark .

###### Remark .

###### Lemma F.7 .

###### Proof.

###### Lemma F.8 .

###### Proof.

###### Lemma F.9 (Freedman’s inequality with unbounded variables) .

###### Remark .

###### Lemma F.10 .

###### Lemma F.11 .

###### Proof.

###### Lemma F.12 .

###### Proof.

### F.3 Deferred proofs of this section

###### Proof of Lemma F.9 .

In this proof, $C_{c}>0$ will be a constant that can depend on $c$ and may change across lines.
Let $M>0$ be a parameter to be determined later. Write

$$ $\displaystyle Z_{t}$ $\displaystyle=Z_{t}\mathbbm{1}\{|Z_{t}|\leq M\}-\mathop{\mathbb{E\/}}\left[Z_{ t}\mathbbm{1}\{|Z_{t}|\leq M\}\mid\mathcal{F}_{t-1}\right]$ $\displaystyle\qquad+\mathop{\mathbb{E\/}}\left[Z_{t}\mathbbm{1}\{|Z_{t}|\leq M \}\mid\mathcal{F}_{t-1}\right]+Z_{t}\mathbbm{1}\{|Z_{t}|>M\}.$ $$

Let $\hat{Z}_{t}$ denote the two terms in RHS of the first line. Note that $(\hat{Z}_{t})_{t}$ is a martingale difference
sequence with conditional variance bounded by $\sigma_{Z}^{2}$. Moreover, every $\hat{Z}_{t}$ is bounded by $2M$. Thus,
by Freedman’s inequality, we have

$$ $\mathop{\mathbb{P\/}}\left[\left|\sum_{t=1}^{T}\hat{Z}_{t}\right|\geq s\right] \leq 2\exp\left(-\frac{s^{2}}{2T(\sigma_{Z}^{2}+M)}\right),\quad\forall s\geq 0.$ (18) $$

Now, we estimate the expectation $\mathop{\mathbb{E\/}}\left[Z_{t}\mathbbm{1}\{|Z_{t}|\leq M\}\mid\mathcal{F}_{t
-1}\right]$.
Since $\mathop{\mathbb{E\/}}\left[Z_{t}\mid\mathcal{F}_{t-1}\right]=0$, it is equal to $\mathop{\mathbb{E\/}}\left[Z_{t}\mathbbm{1}\{|Z_{t}|>M\}\mid\mathcal{F}_{t-1}\right]$,
for which we have

$$ $\displaystyle\left|\mathop{\mathbb{E\/}}\left[Z_{t}\mathbbm{1}\{|Z_{t}|>M\} \mid\mathcal{F}_{t-1}\right]\right|$ $\displaystyle\leq\mathop{\mathbb{E\/}}\left[|Z_{t}|\mathbbm{1}\{|Z_{t}|>M\} \mid\mathcal{F}_{t-1}\right]$ $\displaystyle=\int_{M}^{\infty}\mathop{\mathbb{P\/}}\left[|Z_{t}|\geq s\right] \,\mathrm{d}s\leq a\int_{M}^{\infty}\exp\left(-bs^{c}\right)\,\mathrm{d}s.$ $$

Apply the change-of-variables $y=s/M$ and then $z=y^{c}$. Then, the above becomes

$$ $\displaystyle\left|\mathop{\mathbb{E\/}}\left[Z_{t}\mathbbm{1}\{|Z_{t}|>M\} \mid\mathcal{F}_{t-1}\right]\right|$ $\displaystyle\leq\frac{aM}{c}\int_{1}^{\infty}\exp\left(-bM^{c}z\right)z^{1/c- 1}\,\mathrm{d}z$ $\displaystyle\leq\frac{aM}{c}\int_{1}^{\infty}\exp\left(-bM^{c}z+\left(\frac{1 }{c}-1\right)\log z\right)\,\mathrm{d}z.$ $$

Note that $\log z\leq\sqrt{z}\leq z$ for all $z\geq 1$. Hence, as long as $M^{c}\geq 2(1/c-1)/b$, we will have

$$ $\displaystyle\left|\mathop{\mathbb{E\/}}\left[Z_{t}\mathbbm{1}\{|Z_{t}|>M\} \mid\mathcal{F}_{t-1}\right]\right|$ $\displaystyle\leq\frac{aM}{c}\int_{1}^{\infty}\exp\left(-bM^{c}z/2\right)\, \mathrm{d}z$ $\displaystyle\leq\frac{2a}{bc}\exp\left((1-c)\log M-bM^{c}/2\right).$ $$

Note that there exists some constant $C_{c}>0$ that depends on $c$ such that $\log M\leq M^{c/2}$ for all $M^{c}\geq C_{c}$.
Suppose that $M$ is at least $C_{c}$. Then, as long as $M^{c/2}\geq 4(1-c)/b$, we will have

$$ $\left|\mathop{\mathbb{E\/}}\left[Z_{t}\mathbbm{1}\{|Z_{t}|>M\}\mid\mathcal{F}_ {t-1}\right]\right|\leq\frac{2a}{bc}\exp\left(-bM^{c}/4\right).$ $$

In other words, for any $\varepsilon_{0}>0$, we have $\left|\mathop{\mathbb{E\/}}\left[Z_{t}\mathbbm{1}\{|Z_{t}|>M\}\mid\mathcal{F}_
{t-1}\right]\right|\leq\varepsilon_{0}/T$ if

$$ $\displaystyle M^{c}$ $\displaystyle\geq C_{c}\vee\frac{2(1/c-1)}{b}\vee\frac{16(1-c)^{2}}{b^{2}}\vee \frac{4}{b}\log\left(\frac{2aT}{\varepsilon_{0}bc}\right)$ $\displaystyle=C_{c}\left(\frac{1}{b^{2}}\vee\frac{1}{b}\log\left(\frac{aT}{ \varepsilon_{0}b}\right)\right).$ $$

Meanwhile, by union bound and our tail bound on $Z_{t}$, we have

$$ $\mathop{\mathbb{P\/}}\left[\exists t\in[T],Z_{t}\mathbbm{1}\{|Z_{t}|>M\}\neq 0 \right]\leq\sum_{t=1}^{T}\mathop{\mathbb{P\/}}\left[|Z_{t}|>M\right]\leq Ta \exp\left(-bM^{c}\right).$ $$

Combine the above bounds with (18), and we obtain

$$ $\displaystyle\mathop{\mathbb{P\/}}\left[\left|\sum_{t=1}^{T}Z_{t}\right|\geq \varepsilon_{0}+s\right]$ $\displaystyle\leq\mathop{\mathbb{P\/}}\left[\left|\sum_{t=1}^{T}Z_{t}\right| \geq s\right]+\mathop{\mathbb{P\/}}\left[\exists t\in[T],Z_{t}\mathbbm{1}\{|Z_ {t}|>M\}\neq 0\right]$ $\displaystyle\leq 2\exp\left(-\frac{s^{2}}{2T(\sigma_{Z}^{2}+M)}\right)+Ta\exp \left(-bM^{c}\right),$ $$

where $M>0$ satisfies

$$ $M^{c}\geq C_{c}\left(\frac{1}{b^{2}}\vee\frac{1}{b}\log\left(\frac{aT}{ \varepsilon_{0}b}\right)\right).$ $$

Let $\delta_{\mathop{\mathbb{P\/}}}\in(0,1)$ be our target failure probability. We have

$$ $\displaystyle Ta\exp\left(-bM^{c}\right)\leq\frac{\delta_{\mathop{\mathbb{P\/} }}}{2}$ $\displaystyle\quad\Leftarrow\quad M^{c}\geq\frac{1}{b}\log\left(\frac{2Ta}{ \delta_{\mathop{\mathbb{P\/}}}}\right),$ $\displaystyle 2\exp\left(-\frac{s^{2}}{2T(\sigma_{Z}^{2}+M)}\right)\leq\frac{ \delta_{\mathop{\mathbb{P\/}}}}{2}$ $\displaystyle\quad\Leftarrow\quad s^{2}\geq 2T(\sigma_{Z}^{2}+M)\log\left( \frac{4}{\delta_{\mathop{\mathbb{P\/}}}}\right).$ $$

Thus, for any $\delta_{\mathop{\mathbb{P\/}}}\in(0,1)$, we have with probability at least $1-\delta_{\mathop{\mathbb{P\/}}}$, we have

$$ $\left|\sum_{t=1}^{T}Z_{t}\right|\leq\varepsilon_{0}+\sqrt{2T(\sigma_{Z}^{2}+M) \log\left(\frac{4}{\delta_{\mathop{\mathbb{P\/}}}}\right)}\quad\text{where} \quad M^{c}\geq C_{c}\left(\frac{1}{b^{2}}\vee\frac{1}{b}\log\left(\frac{aT}{ \varepsilon_{0}b\delta_{\mathop{\mathbb{P\/}}}}\right)\right).$ $$

To remove the parameter $\varepsilon_{0}$, we choose $\varepsilon_{0}=\sqrt{2T\sigma_{Z}^{2}\log\left(\frac{4}{\delta_{\mathop{
\mathbb{P\/}}}}\right)}$.
Then, the above becomes, with probability at least $1-\delta_{\mathop{\mathbb{P\/}}}$, we have

$$ $\left|\sum_{t=1}^{T}Z_{t}\right|\leq 2\sqrt{2T(\sigma_{Z}^{2}+M)\log\left( \frac{4}{\delta_{\mathop{\mathbb{P\/}}}}\right)}\quad\text{where}\quad M^{c} \geq C_{c}\left(\frac{1}{b^{2}}\vee\frac{1}{b}\log\left(\frac{aT}{b\sigma_{Z} \delta_{\mathop{\mathbb{P\/}}}}\right)\right).$ $$

∎

###### Proof of Lemma F.10 .

As in the proof of Lemma F.8, we write
$X_{t}=x_{0}+\sum_{s=1}^{t}\xi_{s}+\sum_{s=1}^{t}h_{s-1}Z_{s}$, define

$$ $\tau:=\inf\left\{t\geq 0\,:\,|X_{t}-x_{0}|>T\Xi+C_{c}\sqrt{T\left(\sigma_{Z}^{ 2}+\frac{1}{b^{2/c}}+\frac{\log^{1/c}\left(\frac{aT}{b\sigma_{Z}\delta_{ \mathop{\mathbb{P\/}}}}\right)}{b^{1/c}}\right)\log\left(\frac{T}{\delta_{ \mathop{\mathbb{P\/}}}}\right)}\right\},$ $$

and $\hat{\xi}_{t+1}=\xi_{t+1}\mathbbm{1}\{t<\tau\}$, $\hat{Z}_{t+1}=\mathbbm{1}\{t<\tau\}Z_{t+1}$. By construction,
we have

$$ $\sup_{t\leq T}\left|\sum_{s=1}^{t}\hat{\xi}_{t}\right|\leq T\Xi\quad\text{with probability at least $1-T\delta_{\mathop{\mathbb{P\/}},\xi}$}.$ $$

For the martingale difference term, first note that $h_{t}\hat{Z}_{t+1}/h_{*}$ satisfies (15). Hence, by Lemma F.9, with probability at least $1-\delta_{\mathop{\mathbb{P\/}}}$, we have

$$ $\left|\sum_{s=1}^{t}h_{t}\hat{Z}_{t}\right|\leq C_{c}h_{*}\sqrt{T\left(\sigma_ {Z}^{2}+\frac{1}{b^{2/c}}+\frac{\log^{1/c}\left(\frac{aT}{b\sigma_{Z}\delta_{ \mathop{\mathbb{P\/}}}}\right)}{b^{1/c}}\right)\log\left(\frac{1}{\delta_{ \mathop{\mathbb{P\/}}}}\right)}.$ $$

Replace $\delta_{\mathop{\mathbb{P\/}}}$ with $\delta_{\mathop{\mathbb{P\/}}}/T$, apply the union bound, and we obtain

$$ $\sup_{t\leq T}\left|\sum_{s=1}^{t}h_{t}\hat{Z}_{t}\right|\leq C_{c}h_{*}\sqrt{ T\left(\sigma_{Z}^{2}+\frac{1}{b^{2/c}}+\frac{\log^{1/c}\left(\frac{aT}{b \sigma_{Z}\delta_{\mathop{\mathbb{P\/}}}}\right)}{b^{1/c}}\right)\log\left( \frac{T}{\delta_{\mathop{\mathbb{P\/}}}}\right)},$ $$

with probability at least $1-\delta_{\mathop{\mathbb{P\/}}}$. In other words, we have

$$ $\sup_{t\in[T]}\left|\sum_{s=1}^{t}\hat{\xi}_{t}+\sum_{s=1}^{t}h_{t}\hat{Z}_{t} \right|\leq T\Xi+C_{c}h^{*}\sqrt{T\left(\sigma_{Z}^{2}+\frac{1}{b^{2/c}}+\frac {\log^{1/c}\left(\frac{aT}{b\sigma_{Z}\delta_{\mathop{\mathbb{P\/}}}}\right)}{ b^{1/c}}\right)\log\left(\frac{T}{\delta_{\mathop{\mathbb{P\/}}}}\right)},$ $$

with probability at least $1-T\delta_{\mathop{\mathbb{P\/}},\xi}-\delta_{\mathop{\mathbb{P\/}}}$. To complete the proof, it suffices to repeat the
final part of the proof of Lemma F.8.
∎

###### Proof of Lemma F.9 .

###### Proof of Lemma F.10 .

## Appendix G Simulation

We include simulation results for Stage 1 in this section. The goal here is to provide empirical evidence that
(i) if we have both the second- and $2L$-th order terms, then the sample complexity of online SGD scales
linearly with $d$, (ii) the same also holds for the
absolute function (which is a special case of the setting in Li et al. []) and (iii) without the higher-order terms, online SGD cannot recovery the exact directions.

The setting is the same as the one we have described in Section 2. We choose the hyperparameters
roughly according to Theorem 2.1.
To reduce the needed computational resources, we choose $m=\Theta(P^{2})$ instead of $\tilde{\Omega}(P^{8})$. Note
that by the Coupon Collector problem, we need $m=\Omega(P\log P)$ to ensure that for each $p\in[P]$, there exists
at least one neuron ${\bm{v}}$ with $v_{p}^{2}\geq\max_{q\leq P}v_{q}^{2}$. Since we are mostly interested in the dependence
on $d$, for the learning rate, we choose $\eta=c/d$, where $c$ is a tunable constant that is independent of $d$
but can depend on everything else. $T$ is chosen according to Theorem 2.1 and we early-stop the
training when for all $p\in[P]$, there exists a neuron with $v_{p}^{2}\geq 0.95$ (in the moving average sense).

Figure: Figure 1: Recovery of directions. The above plots show the evolution of the correlation with each of the ground-truth directions in Stage 1. We fix the relevant dimension $P=5$ and vary the ambient dimension $d$. Different colors represent different $d$. For each color, one curve represents $\max_{{\bm{v}}}v_{p}^{2}$ for one $p\in[P]$. In the left plot, the link function is $\phi=h_{2}+h_{4}$, a function that is covered by our theoretical results. We claimed that our theoretical results can be extended to even link function with reasonably regular Hermite coefficients. The right plot, where the link function is the absolute value function, provides an empirical evidence for this. We can see that in both cases, online SGD successfully recover all ground-truth directions, and the number of steps/samples it needs scales approximately linearly with $d$.
Refer to caption: x1.png

Figure: Figure 2: Necessity of the higher order terms. In these two figures, we choose $P=10$ and $d=100$. The left plot shows the maximum correlation each of the ground-truth directions (also see Figure 1). We can see that in the isotropic case, whether online SGD can recover the ground-truth directions is determined by the presence/absence of the higher-order terms. The right plot shows the change of $\max_{{\bm{v}}}v_{p}^{2}/\left\|{\bm{v}}_{\leq P}\right\|^{2}$ for each $p\in[P]$ in Stage 1 when the link function is $h_{2}$. One can observe that they are almost unchanged throughout training. This, together with the left plot, shows that the increase of the correlation is caused by learning the subspace instead of the actual directions.
Refer to caption: x3.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: Some possible typos:

* What is $\epsilon_*$ and how does it relate to $\epsilon$ in Theorem 2.1?

* Is there a normalization step missing from Stage 1 of the algorithm in Section 2.3? Assuming there is a normalization step, then it seems the ODE in Line 278 is only considering the gradient step and not the projection. It is mentioned that the second term on the RHS comes from the projection, but it seems that the second term on the RHS is only coming from the higher degree Hermite term.

* I think $v_k = {v^*}_k^\top v$ has not been formally introduced before the equation in Line 278.

* $\Vert v_{\leq P}\Vert^2$ is missing from the numerator of the last in term in Line 304.

* It seems like the RHS of the equation in Line 352 must be $4L(1 - \Vert v_{\leq P}\Vert^2 + c_L(\Vert v_{\leq P}\Vert^2 - v_1^2))v_1^{2L} \geq 4Lc_L(\Vert v_{\leq P}\Vert^2 - v_1^2)v_1^{2L}$.
soundness: 3
strengths: The majority of theoretical work on learning high-dimensional Gaussian single/multi-index models focus on the importance of the first non-zero term in the Hermite expansion of the link function. It is nice to show that sometimes the higher order terms can be useful in the analysis.
summary: This paper studies the problem of learning an additive Gaussian multi-index model with know link function, where the directions in the multi-index model are orthonormal and the non-linearity only includes the Hermite polynomials of degree 2 and $2L$ for some $L \geq 2$. It is shown that when using online SGD on the squared loss, while looking at the 2nd degree Hermite polynomial is not sufficient for strong recovery of all directions, one can leverage the 2nd degree Hermite polynomial to weakly recover the subspace with $\tilde{O}(d)$ samples. From this weak recovery, once can use the higher-order Hermite term to strongly recover the individual directions in the multi-index model, while maintaining the $\tilde{O}(d)$ sample complexity.
weaknesses: My main concern with the current submission is its scope. Specifically, the picture presented here will never be relevant to functions with information exponent larger than 2, as Oko et al., 2024 show that only looking at the first non-zero term in the Hermite expansion is sufficient. Even if the information exponent is exactly 2, as Oko et al., 2024 argue, after the subspace is recovered with online SGD, it is possible to subtract it with quadratic activation from the multi-index model, thus obtaining a new model that has information exponent larger than 2.

I think while the current analysis is interesting, this submission could be more strengthened if we also had a story for other types of link functions with information exponent 2, e.g. $h_2 + h_3$. Also, moving from exactly orthogonal to approximately orthogonal weights and characterizing the effect of this approximate orthogonality can be another valuable contribution.

---

confidence: 2
contribution: 4
initial_rating: 8
presentation: 4
questions: 1. Is it necessary to keep the weights in the first layer as a collection of orthonormal vectors? Would it be sufficient for them to be extracted from a Normal distribution?
2. Do you have some intuition of what happens if $a_0$ is not so small? As you state the dynamics of different weight vectors in the first layer are no more independent, but do you expect this to impact the qualitative behaviour you describe in your theorems? This could be explored numerically.
3. What would change if we were studying the model $\sum_{k=1}^P {\rm He}_2 (z_k) + {\rm He}_3(z_k)$?
4. What is the role of the regularisation $\lambda$ in learning the second layer? It's common not to regularise when training a two layer network. Could you comment on its role here?
soundness: 3
strengths: The paper is exceptionally well written and easy to follow. The outline of the proof in sections 3 and 4 is incredibly clear, and it gives an insightful and intuitive picture of the different stages of learning.

The problem of studying this class of multi-index models is in my opinion incredibly important as they are instrumental in developing a theory of two layer networks, and studying the interaction between different terms in the Hermite expansion is an important piece that needs to be figured out.
summary: The authors study the problem of learning multi-index models of the form $\sum_{k=1}^P {\rm He}_2 (z_k) + {\rm He}_L(z_k)$ where $L$ is even and greater than $2$ with a two layer neural network of width $m$. The training in the first layer is performed using online SGD and keeping the second layer fixed and small, while the second layer is learned at the end of training using linear regression.

Interestingly, the authors discover that in the training of the first layer there are two phases: first the subspaces are recovered, and then the individual directions are recovered.
weaknesses: I think the paper could greatly benefit from a set of numerical experiments implementing the algorithm in 2.3. The learning algorithm described in the paper is different from a standard gradient descent iteration in which one updates the weights in both the first and second layer together, and it's not clear to me if such a case will have qualitatively the same behaviour as what is presented. Additionally, I believe the assumption that the weights in the first layer evolve independently is not true in practice, and I would expect such interactions to be considered in future works. All of these points could have been addressed by a set of numerical experiments, which could have at least offered a justification for the assumptions taken in this work.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: As the authors discuss in the introduction, the information exponent (IE) was introduced by Ben Arous et al., 2021 specifically in the context of characterizing the sample complexity of one-pass SGD for well-specified single-index models (a.k.a. generalized linear models). While this concept was shown to be also relevant to the sample complexity in non-parametric settings (e.g. learning a single-index target with an infinite width 2-layer neural network, e.g. [Bietti et al. 2022] or [1,2]), there are a couple of works that already highlight that the information exponent is not a relevant concept beyond the rather specific scope for which it was introduced. For instance, Abbe et al., 2022 already points that is not sufficient to characterize strong learnability of multi-index functions, and other works show that the IE is not a robust concept to the change of algorithm even for weak learnability of single- and multi-index models [3, 4, 5, 6, Dandi et al., 2024, Damien et al., 2024, Lee et al., 2024]. In other words, it is already a consensus in the community that the IE is not a fundamental notion of computational complexity. For this reason, I find it is quite a stretch to motivate the result in this manuscript around this concept (e.g. as in the first bullet point in "*1.2 Our Contributions*").

In my reading, one contribution of this work is to provide an explicit example of a non-staircase multi-index target where strong learnability is achieved with one-pass SGD in $\tilde{O}(d)$ steps, showing that the staircase structure of Abbe et al. (2022, 2023) is not a necessary property. This hierarchical "weak-to-strong" learning structure is interesting per se.

Below are some specific questions / suggestions:

- In line 280, is $k\in[d]$ a typo?
- I would suggest the authors to avoid using the same index for quantities which are in different ranges, e.g. $k\in [P]$ or $k\in[m]$ for target and model neurons. This could help readability of the technical parts.
- In this work, you consider a noiseless multi-index target. Would the strong recovery result holds true in the presence of label noise?
- The second sentence in the discussion in L99-103 on the differences with respect to Ben Arous et al., (2024) is confusing, specially in light of Appendix A. If the analysis rely on a reduction to tensor decomposition, the settings do not look so different after all. Maybe be more precise on what you mean by "not comparable".    

Some comments on the related literature:

- Reference [Lee et al., 2024] is concurrent to [4], please acknowledge both when referring to this result.   
- The statistical and computational sample complexity results for single-index models in Damien et al. (2024) generalize the ones from [3]. Please acknowledge this in the discussion in L44-47.
- In L47-48: beyond SGD, computational weak-learnability results for multi-index functions were derived in [5] for polynomials and in [6] general link functions but linear sample complexity.
- Convergence rates for gradient flow on multi-index targets were derived in [7]. The setting studied in this manuscript (which assumes a small learning rate with respect to the initial conditions) should be close to that. Can you comment the relationship to this?

**References**

- [1] Raphaël Berthier, Andrea Montanari, Kangjie Zhou. [Learning time-scales in two-layers neural networks](https://link.springer.com/article/10.1007/s10208-024-09664-9). Foundations of Computational Mathematics. 2024 Aug 22:1-84.
- [2] Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, Ludovic Stephan. [How Two-Layer Neural Networks Learn, One (Giant) Step at a Time](https://arxiv.org/abs/2305.18270). arXiv:2305.18270 [stat.ML]
- [3] Jean Barbier, Florent Krzakala, Nicolas Macris, Léo Miolane, Lenka Zdeborová. [Optimal Errors and Phase Transitions in High-Dimensional Generalized Linear Models](https://www.pnas.org/doi/10.1073/pnas.1802705116). PNAS 2019
- [4] Luca Arnaboldi, Yatin Dandi, Florent Krzakala, Luca Pesce, Ludovic Stephan. [Repetita Iuvant: Data Repetition Allows SGD to Learn High-Dimensional Multi-Index Functions](https://arxiv.org/abs/2405.15459). arXiv:2405.15459 [stat.ML]
- [5] Sitan Chen, Raghu Meka. [Learning Polynomials of Few Relevant Dimensions](http://proceedings.mlr.press/v125/chen20a/chen20a.pdf). ICML 2020
- [6] Emanuele Troiani, Yatin Dandi, Leonardo Defilippis, Lenka Zdeborová, Bruno Loureiro, Florent Krzakala. [Fundamental computational limits of weak learnability in high-dimensional multi-index models](https://arxiv.org/abs/2405.15480). arXiv:2405.15480 [cs.LG]
- [7] Alberto Bietti, Joan Bruna, Loucas Pillaud-Vivien. [On Learning Gaussian Multi-index Models with Gradient Flow](https://arxiv.org/pdf/2310.19793). arXiv:2310.19793 [stat.ML]
soundness: 4
strengths: The description of the time scales for the dynamics is interesting. The informal discussion of how the SGD argument can be derived from the GF analysis and "user-friendly" discussion of the martingale argument from (Ben Arous 2021) is a nice addition to the literature.
summary: This work investigates the sample complexity of learning Gaussian multi-index functions of the type $f_{\star}(x) = \sum\limits_{k=1}^{P}\left(h_{2}(\langle v_{k}^{\star},x\rangle)+h_{2L}(\langle v_{k}^{\star},x\rangle)\right)$ with two-layer neural networks trained under a particular two-step spherical one-pass SGD scheme.

The main result is to show that this target can be strongly learned with this procedure with $\tilde{O}({\rm poly}(P))$ hidden-units and $\tilde{O}(d{\rm poly}(P))$ samples / steps, in contrast with the purely quadratic case where the target can only be weakly learned. This result stems from an asymptotic description of the training dynamics in two phases, the first where only a subspace is recovered, followed by a second phase where the exact directions $\{v_{k}^{\star}\}_{k=1}^{P}$ are retrieved.
weaknesses: - The setting studied is quite specific, and the focus on the information exponent without a broader discussion of the Gaussian multi-index literature can be misleading to the reader (see *Questions*).
- Overall, the manuscript is quite technical, which hinders clarity.

Meta-review: (a) **Summary of the scientific claims and findings**  
The paper studies the learning dynamics of multi-index Gaussian functions through the lens of stochastic gradient descent (SGD). It demonstrates that a two-phase approach—using lower-order Hermite terms for subspace recovery followed by higher-order terms for exact directional recovery—can lead to better sample complexities than existing methods for multi-index models. This hierarchical approach emphasizes structural properties of multi-index models, extending beyond traditional information exponent analyses.

(b) **Strengths**  
- The paper identifies key theoretical insights into the interplay of lower- and higher-order terms in Hermite expansions, advancing understanding of multi-index learning.  
- It highlights a novel hierarchical learning dynamic that could inform future algorithm designs.  
- The theoretical results are supported by clear mathematical arguments and connections to existing literature.  

(c) **Weaknesses**  
- **Lack of generality**: The method applies to specific link functions and relies on exact orthogonality assumptions, limiting broader applicability.  
- **Experimental Validation**: The lack of numerical experiments weakens empirical support for the theoretical claims, leaving important practical questions unanswered.  
- **Clarity in Comparisons**: The relationship between this work and prior results, such as those of Oko et al. (2024), needs more precise articulation.  
- **Unaddressed Scenarios**: Important extensions, such as approximate orthogonality and noisy label scenarios, are not thoroughly explored.  

(d) **Decision Rationale**  
The paper offers valuable insights but falls short of ICLR's standard due to limited scope, incomplete empirical validation, and lack of generalizability. These issues, coupled with specific points raised by reviewers, warrant a rejection. However, the authors are strongly encouraged to address these limitations and resubmit.

Decision: Reject