Learning Over-Parametrized Two-Layer ReLU Neural Networks beyond NTK

Introduction

Gradient-based optimization methods are the method of choice for learning neural networks. However, it has been challenging to understand their working on non-convex functions. Prior works prove that stochastic gradient descent provably convergences to an approximate local optimum Ge et al. (2015); Sun et al. (2015); Lee et al. (2017); Kleinberg et al. (2018). Remarkably, for many highly complex neural net models, gradient-based methods can also find high-quality solutions Sun (2019) and interpretable features Zeiler and Fergus (2014).

Recent studies made the connection between training wide neural networks and Neural Tangent Kernels (NTK) Jacot et al. (2018); Arora et al. (2019b); Cao and Gu (2019); Du et al. (2018c). The main idea is that training neural networks with gradient descent with a particular initialization is equivalent to using kernel methods. However, the NTK approach has not yet provided a fully satisfactory theory for explaining the success of neural networks. Empirically, there seems to be a non-negligible gap between the test performance of neural networks trained by SGD and that of the NTK Arora et al. (2019a); Li et al. (2019b). Recent works have suspected that the gap stems from that the NTK approach has difficulty dealing with non-trivial explicit regularizers or does not sufficiently leverage the implicit regularization of the algorithm Wei et al. (2019); Chizat and Bach (2018b); Li et al. (2019a); HaoChen et al. (2020).

In this work, we provide a new convergence analysis of the gradient descent dynamic on an over-parametrized two-layer ReLU neural network. We prove that for learning a certain two-layer target network with orthonormal ground truth weights, gradient descent is provably more accurate than any kernel method that uses polynomially large feature maps.

where $a_{i}$ is in $[\frac{1}{\kappa d},\frac{\kappa}{d}]$ for an absolute constant $\kappa\geq 1$ and satisfies $\sum_{i\in[d]}a_{i}=1$, and $\{w_{i}^{\star}\}_{i=1}^{d}$ forms an orthonormal basis. Equation (1.1) can also be written as the sum of $2d$ neurons with ReLU activation:

Let $\mathcal{Z}=\{(x_{j},y_{j})\}_{j=1}^{N}$ be a training dataset of $N$ i.i.d. samples from the Gaussian distribution with identity covariance and $y_{j}=f^{\star}(x_{j})$ for any $1\leq j\leq N$.

We learn the target network $f^{\star}$ using an over-parametrized two-layer ReLU network with $m\geq 2d$ neurons $W=\{w_{i}\}_{i=1}^{m}$, given by:

Note that we have re-parametrized the output layer with the norm of the corresponding neuron, so that we only have one set of parameters $W$. This is without loss of generality for learning $f^{\star}$ because when $a_{i}\geq 0$, $a_{i}\cdot\operatorname*{ReLU}(w_{i}^{\top}x)$ is equal to $\|w_{i}^{\prime}\|\cdot\operatorname*{ReLU}({{w_{i}^{\prime}}^{\top}x})$ where $w_{i}^{\prime}=\sqrt{a_{i}/\|w_{i}\|}\cdot w_{i}$. Given a training dataset $\mathcal{Z}=\{(x_{i},y_{i})\}_{i=1}^{N}$, we learn the target network by minimizing the following empirical loss:

Let $L(W)$ denote the population loss given by the expectation of $\hat{L}(W)$ over $\mathcal{Z}$.

Algorithm. We focus on truncated gradient descent with random initialization. Algorithm 1 describes the procedure.An interesting feature is that when a neuron becomes larger than a certain threshold, we no longer update the neuron. This is a variant of gradient clipping often used in training recurrent neural networks (e.g. Merity et al. (2017); Gehring et al. (2017); Peters et al. (2018)) — here we drop the gradients of the large weights instead of re-scaling them. The truncation allows us to upper bound the norm of every neuron. Our main result is to show that Algorithm 1 learns the target network accurately in polynomially many iterations.

Let $\mathcal{Z}$ be a training dataset with $N=\operatorname{poly}_{\kappa}(d)$ samples generated by the model described above.Let $\operatorname{poly}(d)$ denote a polynomial of $d$ and $\operatorname{poly}_{\kappa}(d)$ denote a polynomial whose degree may depend on $\kappa$. Let $C(\kappa)$ be a sufficiently large constant that only depends on $\kappa$. Let $0<Q<1/100$ be a sufficiently small absolute constant that does not depend on $\kappa$. Let $\lambda_{0}$ be a sufficiently small value on the order of $1/\operatorname{poly}(d)$ and $\lambda_{1}\leq\lambda_{0}/O(\operatorname{poly}_{\kappa}(d))$ be a sufficiently small value on the order of $1/\operatorname{poly}_{\kappa}(d)$. For a learning rate $\eta<\min\left(\lambda_{0}^{2},O(\frac{1}{\operatorname{poly}_{\kappa}(d)})\right)$, a network width $m\geq\Omega(\operatorname{poly}(d)/\operatorname{poly}(\lambda_{1}))$, and truncation parameters $\lambda_{0},\lambda_{1}$, let $\hat{W}$ be the final network learned by Algorithm 1. With probability $1-\frac{1}{\operatorname{poly}(d)}$ over the choice of the random initialization, we have that the population loss of $\hat{W}$ satisfies

The intuition behind our main result is as follows. We build on a connection between the popluation $L(W)$ and tensor decomposition for Gaussian inputs Ge et al. (2017, 2018). By expanding the population loss in the Hermite polynomial basis, the optimization problem becomes an infinite sum of tensor decompositions problems (cf. equation (2) in Section 2) To analyze the gradient descent dynamic on the infinite sum tensor decomposition objective, we first analyze the infinite-width case – when $m$ goes to infinity. We establish a conditional-symmetry condition on the population of neurons, which greatly simplifies the analysis. This is established using the fact that our input distribution and labeling function (the absolute value activation) are both symmetric. Our analysis uncovers a stage-wise convergence of the gradient descent dynamic as follows, which matches our observations in simulations.

First, Algorithm 1 minimizes the 0th and 2nd order tensor decompositions. Informally, the distribution of neurons is fitting to the 0th moment and the 2nd moment of $\left\{w_{i}^{\star}\right\}_{i=1}^{d}$.

Second, Algorithm 1 minimizes the 4th and higher order tensor decompositions. Initially, there is a long plateau where the evolution is slow, but after a certain point gets faster. As a remark, this behavior has been observed for randomly initialized tensor power method Anandkumar et al. (2017). Because the solution to the 4th and higher order orthogonal tensor decomposition problems is unique, we can learn the ground truth weights $\left\{w_{i}^{\star}\right\}_{i=1}^{d}$.

Then we show that the sampling error between the infinite-width case and the finite-width case is small. The finite-width case can be thought of as a finite sample of the infinite-width case. As the network width increases, the sampling error reduces. In Section 3 and 4, we will first present a proof overview. The full proof is given in Section A and B.

As a complement, we show that the generalization error bound of Theorem 1.1 cannot be achieved by kernel functions with polynomially large feature map. Hence, by minimizing the higher order tensor decomposition terms, the learned neural network is provably more accurate than kernel functions that simply fit the lower order terms. Our result is stated as follows.

Under either of the following two situations,

Comparing the above result with Theorem 1.1, we conclude that provided with polynomially many samples, Algorithm 1 can recover the target two-layer neural network more accurately than the feature map and kernel method described above. Section C shows how to prove Theorem 1.2.

2 Related Work

Neural tangent kernel (NTK). A sequence of recent work shows that the learning process of gradient descent on over-parametrized neural networks, under certain initializations, reduces to the learning process of the associated neural tangent kernel. See Jacot et al. (2018); Arora et al. (2019b); Cao and Gu (2019); Du et al. (2018c); Arora et al. (2019a); Allen-Zhu and Li (2019b); Allen-Zhu et al. (2019c, b); Li and Liang (2018); Zou et al. (2018); Du et al. (2018a); Daniely et al. (2016); Ghorbani et al. (2019); Li et al. (2019a); Hanin and Nica (2019); Yang (2019) and the references therein. For NTK based results, the learning process of gradient descent can be viewed as solving convex kernel regression. Our work analyzes a non-convex objective that involves an infinite sum of tensor decomposition problems. By analyzing the higher order tensor decompositions, we can achieve a smaller generalization error than kernel methods.

Allen-Zhu and Li (2019a, 2020a) show that over-parametrized neural networks can learn certain concept class more efficient than any kernel method. Their work assumes the target network satisfies a certain “information gap” assumption between the first and second layer, while our target network does not require such gaps. Allen-Zhu et al. (2019a); Bai and Lee (2019) go beyond NTK by studying quadratic approximations of neural networks. Our work further analyzes higher-order tensor decompositions that are present in the Taylor expansion of the loss objective.

Two-layer neural networks given Gaussian inputs. There is a large body of work on learning two-layer neural networks over the last few years, such as Kawaguchi (2016); Soudry and Carmon (2016); Xie et al. (2016); Soltanolkotabi et al. (2017); Tian (2017); Brutzkus and Globerson (2017); Boob and Lan (2017); Vempala and Wilmes (2018); Oymak and Soltanolkotabi (2019); Bakshi et al. (2018); Yehudai and Shamir (2019); Zhang et al. (2018); Li and Liang (2017); Li and Dou (2020); Allen-Zhu and Li (2020b). Our work is particularly related to those that learn a two-layer neural network given Gaussian inputs. Li and Yuan (2017); Zhong et al. (2017) consider learning two-layer networks with ReLU activations with a warm start tensor initialization, as opposed to from a random initialization. Du et al. (2017) consider learning a target function consisting of a single ReLU activation. Brutzkus and Globerson (2017); Tian (2017) study the case where the weight vector for each neuron has disjoint support. Apart from the gradient descent algorithm, the method of moments has also been shown to be an effective strategy with provable guarantees (e.g. Bakshi et al. (2018); Ge et al. (2018)).

The closest work to ours is Ge et al. (2017) that consider a similar concept class. However, their work requires designing a complicated loss function, which is different from the mean squared loss. The learner network also uses a low-degree activation function as opposed to the ReLU activation. These are introduced to address the challenge of analyzing non-convex optimization for tensor decomposition with multiple components as variables, because prior works mostly focus on the non-convex formulation that optimizes over a single component (e.g., see Ge and Ma (2017)). Ge et al. (2017) have stated the question of analyzing the gradient descent dynamic for minimizing the sum of second and fourth order tensor decompositions as a challenging open question. Our analysis not only applies to this setting, but also allows for more even order tensor decompositions. Apart from ReLU activations, quadratic activations have been studied in Li et al. (2018); Oymak and Soltanolkotabi (2019); Soltanolkotabi et al. (2017).

Infinite-width neural networks. Previous work such as Mei et al. (2018); Chizat and Bach (2018a) show that as the hidden layer width goes to infinity, gradient descent approaches the Wasserstein gradient flow. Mei et al. (2018) use tools from partial differential equations to prove the global convergence of the gradient descent. Both of these results do not provide explicit convergence rates. Wei et al. (2018) show that under a certain regularity assumption on the activation function, the Wasserstein gradient flow converges in polynomial iterations for infinite-width neural networks..

Organizations. The rest of the paper is organized as follows. In Section 2, we reduce our setting to learning a sum of tensor decomposition problems. In Section 3, we describe an overview of the analysis for the infinite-width case. In Section 4, we show how to connect the above case to the gradient descent dynamic on the empirical loss for polynomially-wide networks. Finally we validate our theoretical insight on simulations in Section 5. In Section A, we provide the proof of the infinite-width case. In Section B, we provide an error analysis of the infinite-width case and complete the proof of Theorem 1.1. In Section C, we present the proof of Theorem 1.2.

Preliminaries

Recall that the ground-truth weights $\left\{w_{i}^{*}\right\}_{i=1}^{d}$ forms an orthonormal basis. Since the input distribution $x\sim\mathcal{N}(0,I_{d\times d})$ and the initialization $\left\{w_{i}\right\}_{i=1}^{m}\sim\mathcal{N}\left(0,\frac{1}{d}\cdot\operatorname{Id}_{d\times d}\right)$ are both rotation invariant, without loss of generality we can assume that $w_{i}^{*}=e_{i}$, for all $1\leq i\leq d$.

We can average out the randomness in $x$ by applying Theorem 2.1 of Ge et al. (2017) on the loss function $L(W)$, by expanding the activations function in the Hermite basis O’Donnell (2014).

where $c_{k}=\frac{2[(k-3)!!]^{2}}{\pi\cdot k!}$ is the Hermite coefficients of the absolute value function for any $k\geq 0$. We remark that the population loss is a infinite sum of orthogonal tensor decomposition problems! For example, the -th order tensor decomposition concerns the $l_{2}$-norm of the weights. More generally, the $k$-order tensor decomposition concerns the $k$-th moment of the weights.

Correspondingly, the population loss of $f_{\mathcal{P}}(x)$ is given as

Gradient descent update. It has been shown in prior works that gradient descent in the (natural) parameter space corresponds to Wasserstein gradient descent in the distributional space. However, we found that the Wasserstein gradient perspective is not particularly helpful for us to analyze our algorithms and therefore we work with the update in the parameter space. The distribution $\mathcal{P}$ can be viewed as a collection of infinitesimal neurons. The gradient of each neuron $v$ is given by computing the gradient of the objective $L(W)$ w.r.t a particle $v$ assuming the rest of the particles follow the distribution $\mathcal{P}$. Let $\nabla_{v}L_{\infty}(\mathcal{P})$ denote the gradient of $v$. We have that

where $b_{0}=4c_{0},b_{1}=2c_{1}$, and for any $j\geq 2$, $b_{2j}=(4j)\times c_{2j}=\Theta\left(\frac{1}{j^{2}}\right)$ and $b_{2j^{\prime}}=(4j-4)\times c_{2j}$. We use $\nabla_{v}L_{\infty}$ and $\nabla_{v}$ as a shorthand for $\nabla_{v}L_{\infty}(\mathcal{P})$. Based on equation (2.5), we can further decompose $\nabla_{v}L_{\infty}(\mathcal{P})$ into the sum of $\nabla_{2j,v}L_{\infty}(\mathcal{P})$ for $j\geq 0$, where the $2j$-th gradient refers to the gradient of the $2j$-th tensor decomposition. As a result, given a neural network with neuron distribution $\mathcal{P}^{(t)}$, the neuron distribution after a truncated gradient descent step, denoted by $\mathcal{P}^{(t+1)}$, satisfies that

Finite-width case. We briefly describe the connection between the above infinite-width case and the finite-width case. Intuitively, we can think of the finite-width case as sampling $m$ neurons randomly from the neuron population $\mathcal{P}$ in the infinite-width case. There are two sources of sampling error that arise from the above process: (i) the error of the gradients between the finite neuron distribution and the infinite neuron distribution; (ii) the error between the empirical loss and the population loss. Because of gradient truncation, the norm of every neuron is bounded by $1/\lambda$. Therefore, the sampling error reduces as $m$ and $N$ increases, as shown in the following claim.

With probability at least $1-\delta$ over the randomness of $\{w_{i}\}_{i=1}^{m}$ and the training dataset $\mathcal{Z}$, for every $w\in W$, we have that:

Claim 2.1 can be proved by standard concentration inequalities such as the Chernoff bound.

Overview of the Infinite-Width Case

We begin by studying Algorithm 1 for minimizing the population loss using an infinite-width neural network. The infinite-width case plays a central role in our analysis. First, the infinite-width case allows us to simplify the gradient update rule through a conditional-symmetry condition that we describe below. Second, the finite-width case can be reduced to the infinite-width case by bounding the sampling error of the two cases — we describe the reduction in the next section.

A natural starting point for the infinite-width case is to simply set the network width $m$ to infinity in Theorem 1.1. However, this will include negligible outliers such as those with large norms in the Gaussian distribution. Therefore, we focus on a truncated probability measure $\mathcal{P}^{(0)}$ of $\mathcal{N}(0,\operatorname{Id}_{d\times d})$ by enforcing a certain bounded condition. The precise definition of $\mathcal{P}^{(0)}$ is presented in Definition A.1 of Appendix A. For the purpose of providing an overview of the analysis, it suffices to think of $\mathcal{P}^{(0)}$ as a Gaussian-like distribution that satisfies the following property.

Provided with $\mathcal{P}^{(0)}$ as initialization, we are ready to state the main result of the infinite-width case as follows.

In the setting of Theorem 1.1, let the number of samples $N$ go to infinity. Starting from the initialization $W^{(0)}$ as the neuron distribution $\mathcal{P}^{(0)}$, let $\hat{W}$ be the final output network by Algorithm 1. The population loss of $\hat{W}$ satisfies $L(\hat{W})\leq O(1/d^{1+Q})$.

In the rest of this section, we present an overview of the proof of Theorem 3.1 and provide pointers to the proof details to be found in Section A. First, we provide a simplifying formula for the gradient of $L_{\infty}(\mathcal{P})$. We describe an overview of the two stages of Algorithm 1 in Section 3.1 and 3.2, respectively.

Suppose the update rule of $\mathcal{P}^{(t)}$ is given in equation (2.6). If $\mathcal{P}^{(t)}$ is conditionally-symmetric, then $\mathcal{P}^{(t+1)}$ is also conditionally-symmetric.

To see that Claim 3.1 is true, we first observe that the 1st order tensor decomposition is always zero when $\mathcal{P}^{(t)}$ is conditionally symmetric. For the even order tensor decompositions, we observe that for every neuron $v$ in $\mathcal{P}^{(t)}$ and every $1\leq j\leq d$, subject to $v_{-j}$ being fixed, $\nabla_{v}L_{\infty}(\mathcal{P})$ is a polynomial of $v_{j}$ that only involves odd degree monomials. Therefore, as long as $\mathcal{P}^{(t)}$ is conditionally-symmetric, then $\mathcal{P}^{(t+1)}$ is still conditionally-symmetric. Since $\mathcal{P}^{(0)}$ is conditionally-symmetric by definition, we conclude that the neuron distribution is conditionally-symmetric throughout Algorithm 1. Based on this claim, we simplify equation (2.5) as follows.

Suppose that $\mathcal{P}=\mathcal{P}^{(t)}$ is conditionally-symmetric. For any $j\geq 0$, let $\nabla_{2j,v}$ be a shorthand for the gradient of the 2j-th tensor $\nabla_{2j,v}L_{\infty}(\mathcal{P})$. For any $1\leq i\leq d$, let $[\nabla_{2j,v}]_{i}$ be the $i$-th coordinate of $\nabla_{2j,v}$. We have that $[\nabla_{2j,v}]_{i}$ is equal to the following for each value of $j$:

The proof of Claim 3.2 is by applying Claim 3.1 to equation (2.5), which zeroes out the coordinates in $w$ that has an odd order before taking the expectation of $w$ in $\mathcal{P}$. For the 2nd order gradient $[\nabla_{2,v}]_{i}$, we have that

Similar arguments apply to the gradient of higher order tensor decompositions. Claim 3.1 and 3.2 together implies that for the infinite-width case, the gradient descent update is given by equation (3.2) and (3.3).

We show that Algorithm 1 minimizes the 0th and 2nd order tensor decompositions of the objective $L_{\infty}$ to zero first.

First, we show that the gradient of the 4th and higher order tensor decompositions is dominated by $\nabla_{0,v}$ and $\nabla_{2,v}$. We observe that for $v\sim\mathcal{P}^{(0)}$, the $i$-th coordinate of $\nabla_{0,v}$ and $\nabla_{2,v}$ satisfies that

This is because $\mathcal{P}^{(0)}$ is a suitable truncation of $\mathcal{N}(0,\operatorname{Id}_{d\times d}/d)$. We further have that

After the 0th and 2nd order tensor decompositions are minimized to a small enough value, the gradient of higher order tensor decompositions begins to dominate the update. In Lemma A.3, we show that for a small fraction of neurons, their norms become much larger than an average neuron — a phenomenon that we term as “winning the lottery ticket”. The main intuition is as follows.

In Proposition A.10, we show that the gradient of most neurons $v$ except a small fraction can be approximated by a signal term from the 4th order gradient plus an $O(1/d^{2})$ error term:

where $C_{t}(\kappa)$ is a function that only depends on $\kappa$ but grows slowly with $t$. To see that equation (3.5) is true, except a small set of neurons with probability mass at most $1/d^{\alpha}$ where $\alpha$ will be specified later, any other neuron $w$ satisfies $\|w\|_{\infty}^{2}\leq{\alpha\log d}/{d}$. For the small set of neurons, since we stop updating a neuron when its norm grows larger than $1/\lambda_{0}$, the norm of any of these neurons is less than ${1}/{\lambda_{0}}$. Thus, provided with a sufficiently large $\alpha$, the contribution of these neurons to the gradient is negligible. Combined together, we prove equation (3.5) in Proposition A.10.

Next, we reduce the dynamic to tensor power method. Based on equation (3.5), we observe that the update of $v_{i}$ is approximately $v_{i}^{(t+1)}\approx v_{i}^{(t)}+\eta\cdot a_{i}\cdot(v_{i}^{(t)})^{3}$, which is analogous to performing power method over a fourth order tensor decomposition problem. Hence, for larger initializations of $v_{i}$, $v_{i}$ also grows faster. Based on the intuition, we introduce the set of “basis-like” neurons $\mathcal{S}_{i,good}$ in the population $\mathcal{P}$, which are defined more precisely in Lemma A.3. Intuitively, $\mathcal{S}_{i,good}$ includes any neuron $v$ that satisfies $[v_{i}^{(0)}]^{2}\geq{C^{2}\log d}/{d}$, which has probability measure at least ${1}/{d^{C^{2}}}$ by standard anti-concentration inequalities. Following equation (3.5), we show that the neurons in $\mathcal{S}_{i,good}$ keeps growing until they become roughly equal to $e_{i}/(\lambda_{0}\operatorname{poly}(d))$.

As shown in Lemma A.3, Algorithm 1 goes through a long plateau of $O_{\kappa}({d^{2}}/({\eta}\operatorname{poly}\log(d))$ iterations, until the neurons of $\mathcal{S}_{i,good}$ are sufficiently large. Intuitively, the scaling of $d^{2}$ in the number of iterations arises from the $1/d^{2}$ increment in equation (3.5). This concludes Stage 1. The update of these basis-like neurons will be the focus of Stage 2.

2 Dynamic during Stage 2

In the second stage, we reduce the gradient truncation parameter in Algorithm 1 from $\lambda_{0}=\Theta(1/\operatorname{poly}(d))$ to a smaller value $\lambda_{1}=\Theta(1/\operatorname{poly}_{\kappa}(d))$. This allows the neurons that are close to basis vectors to fit the target network more accurately.

In Lemma A.5, we show that after $\Theta(d\log d/\eta)$ iterations, the population loss reduces to less than $o(1/(d\log^{0.01}d))$. The proof of Lemma A.5 involves analyzing the 0th and 2nd order tensor decompositions, similar to Stage 1.1.

At the end of Stage 2.1, the weights of the learner neural network form a “warm start” initialization, meaning that its population loss is less than $o(1/d)$ Li and Yuan (2017); Zhong et al. (2017). The final substage will show that the population loss can be further reduced from $o(1/(d\log^{0.01}d))$ to $O(1/d^{1+Q})$, where $Q$ is a fixed constant defined in Theorem 1.1.

In Lemma A.6, we show that the population loss further reduces to $O({1}/d^{1+Q})$ after $\Theta(d^{1+10Q}/\eta)$ iterations. We describe an informal argument by contrasting the gradient update of neurons in $\mathcal{S}_{i,good}$ and the rest of the neurons for a particular coordinate $i\in[d]$.

For any neuron $v\in\mathcal{S}_{i,good}$, in Claim A.10, we show that the $i$-th coordinate of $v$ approximately follows the following update (cf. equation (A.45)):

where $c_{t}$ is a function that grows with $t$ but bounded above by $O(d^{Q})$ and $C(\kappa)$ is a function that only depends on $\kappa$. For any neuron $v\notin\mathcal{S}_{i,good}$, in Claim A.10, we show that $v_{i}$ follows a similar update but its corresponding value of $c_{t}$ is much smaller than that of neurons in $\mathcal{S}_{i,good}$. Thus, basis-like neurons grow faster than the rest of neurons by an additive factor that scales with $c_{t}/d^{2}$.

Once Lemma A.6 is finished, Algorithm 1 has learned an accurate approximation of $f^{\star}(\cdot)$ and we can conclude the proof of Theorem 3.1. We show that the population loss has also become less than $O(d^{1+Q})$ (cf. equation (A.10)). Thus, we have finished the analysis of Algorithm 1 for $L_{\infty}(\mathcal{P})$. We provide the proof details of Theorem 3.1 in Section A.

Overview of the Finite-Width Case

Based on the analysis of the infinite-width case, we reduce the finite-width case to the infinite-width case. By applying Claim 2.1 with $\mathcal{P}=\mathcal{P}^{(t)}$, when $\{w_{i}^{(t)}\}_{i=1}^{m}$ are i.i.d. samples from $\mathcal{P}^{(t)}$, the empirical loss and its gradient are tightly concentrated around the population loss and its gradient. Furthermore, as we increase the number of neurons $m$ and the number of samples $N$, the sampling error reduces. Therefore, the goal of our reduction is to show that the sampling error remains small throughout the iterations of Algorithm 1. We describe our reduction informally and leave the details to Section B.

Combined together, we show in Lemma B.1 that $\xi_{w}^{(t)}$ indeed remains polynomially small. For Stage 2, we analyze the propagation of $\xi_{w}^{(t)}$ in Lemma B.2 and B.3 using similar arguments.

Combining the above three lemmas on error propagation and Theorem 3.1, we complete the proof of Theorem 1.1 in Section B.

Simulations

We provide simulations to complement our theoretical result. We consider a setting where $w_{i}^{\star}=e_{i}$ and $a_{i}=1/d$, for $1\leq i\leq d$. The input is drawn from the Gaussian distribution. For the $i$th order tensor, we measure the corresponding tensor decomposition loss from the population loss $L(W)$.

We validate the insight of our analysis, which shows that the convergence of gradient descent has several stages. We use the labeling function of equation (1.1) and a learner network with absolute value activation functions as in Section 3 and Section A. First, the 0th and 2nd order tensor decomposition losses converge to zero quickly. Second, the 4th and higher order tensor decomposition losses converge to zero followed by a long plateau. Figure 2 shows the result. Here we use $d=30$ and $m=100>2d$. The number of samples is $10^{4}$.

We can see that initially, the 0th and 2nd order tensor decompositions have higher loss than the 4th and higher order tensor decompositions. Then, both the 0th and the 2nd order losses decrease significantly from the initial value and converge to below $10^{-1}$ very quickly. Moreover, after a quick warm up period, the 0th order loss always stays smaller than the 2nd order loss, as our theory predicts. This is followed by a long plateau, which corresponds to Stage 1.2 of our analysis. During this stage, the 4th and higher order losses dominate dynamic, where a small fraction of neurons converge to basis-like neurons. Eventually, the learner neural network accumulates enough basis-like neurons from the 4th and higher tensors in the network. The 4th and higher order losses become less than $10^{-2}$. The 0th and 2nd order losses further reduce to closer to zero. Our theory provides an in-depth explanation of these phenomena.

It has been observed that for properly parametrized gradient descent, gradient descent can get stuck starting from a random initialization Ge et al. (2017); Du et al. (2018b). We show that this is because the higher order losses remain large even though the 0th order loss has become small. We consider the same setting as the previous experiment but use $m=2d$. Figure 2 shows the result. We can see that the 0th order loss still reduces to less than $10^{-2}$. However, the 2nd, 4th and 6th order losses are still larger than $10^{-1}$ even after $10^{5}$ iterations.

Conclusions and Discussions

In this work, we have shown that for learning a certain target network with absolute value activation, a truncated gradient descent algorithm can provably converge in polynomially many iterations starting from a random initialization. The learned network is more accurate compared to any kernel method that uses polynomially large feature mappings.

We describe several interesting questions for future work. First, it would be interesting to extend our result to a setting where the target network uses ReLU activation, i.e. $f^{\star}(x)=a^{\top}\operatorname*{ReLU}(Wx)$. We note that there is a straightforward reduction from the above setting to our setting by simply solving a linear regression. After applying the reduction, we could then apply our result. The challenge of directly analyzing gradient descent for learning $f^{\star}(x)=a^{\top}\operatorname*{ReLU}(Wx)$ is that the 1st order tensor decomposition in the Hermite expansion of $f^{\star}(x)$ breaks the conditionally-symmetric property. Second, it would be interesting to extend our result to settings where $W^{\star}$ is not necessarily orthonormal. The challenge is to analyze the gradient descent dynamic beyond orthogonal tensors. We leave this question for future research.

The work is in part supported by SDSI and SAIL. T. M is also supported in part by Lam Research and Google Faculty Award.

References

The appendix provides complete proofs to Theorem 1.1 and 1.2.

In Section A, we describe the proof of Theorem 3.1 for the infinite-width case. This section comprises the bulk of the appendix.

In Section B, we describe the proof of Theorem 1.1 by reducing the finite-width case to the infinite-width case.

In Section C, we prove Theorem 1.2 using ideas from the work of Allen-Zhu and Li [2019a].

Appendix A Proof of the Infinite-Width Case

We provide the proof of Theorem 3.1, which shows that running truncated gradient descent on an infinite-width network can recover the target network with population loss at most $O(d^{1+Q})$, where $Q$ is a sufficiently small constant defined in Theorem 3.1. Recall from Section 3 that our analysis begins by setting up the random initialization and then proceeds in two stages. We fill in the proof details left from Section 3. The rest of this section is organized as follows.

Initialization: We set up the random initialization used by Algorithm 1.

Stage 1: We fill in the proof details of the dynamic during Stage 1, which subsumes Stage 1.1 and Stage 1.2 described in Section 3.1. This stage runs for $\Theta(\frac{d^{2}}{\eta C(\kappa)\log d})$ iterations.

Stage 2: We fill in the proof details of the dynamic during Stage 2, which subsumes Stage 2.1 and Stage 2.2 described in Section 3.2. This stage runs for $\Theta(\frac{d^{1+10Q}}{\eta})$ iterations.

Recall that for the infinite-width case, our initialization of the neuron distribution is a probability measure truncated from a Gaussian distribution with identity covariance. We formally define the truncation and the initialization, denoted by $\mathcal{P}^{(0)}$, as follows.

The maximum entry of $w$ is bounded: $\|w\|_{\infty}\leq\frac{\operatorname{poly}\log(d)}{\sqrt{d}}$.

Both $\|w\|_{2}^{2}$ and $\sum_{i=1}^{d}a_{i}d\cdot w_{i}^{2}$ are in the range

There are at most $O(\log^{0.01}(d))$ coordinates $i\in[d]$ of $w$ such that $w_{i}^{2}\geq\frac{\log d}{d}$.

We define $\mathcal{P}^{(0)}$ as the probability measure of $\mathcal{N}(0,\operatorname{Id}_{d\times d}/d)$ conditional on the support set $\mathcal{S}_{g}$.

Remark. For our purpose of proving the finite-width case later in Section B, it suffices to consider $\mathcal{P}^{(0)}$ as the initialization as opposed to $\mathcal{N}(0,\operatorname{Id}_{d\times d}/d)$. This is because when Algorithm 1 samples $m=\operatorname{poly}_{\kappa}(d)$ neurons from $\mathcal{N}(0,\operatorname{Id}_{d\times d}/d)$, with high probability all the $m$ samples are in the set $\mathcal{S}_{g}$. To see this, by standard concentration inequalities for the Gaussian distribution, we can show that the set $\mathcal{S}_{g}$ has probability measure at least $\mu(\mathcal{S}_{g})\geq 1-\frac{1}{d^{\Omega(1)}}$. Thus by union bound, with high probability all $m$ samples are in $\mathcal{S}_{g}$.

As stated in Section 3, we are going to heavily use the conditionally-symmetric property (cf. Definition 3.1). We observe that the initialization $\mathcal{P}^{(0)}$ is indeed conditionally-symmetric. This is because $\mathcal{N}\left(0,\operatorname{Id}_{d\times d}/d\right)$ satisfies the conditionally-symmetric property and our truncation in Definition A.1 only involves conditions on the square of the coordinates of $w$. Hence the truncation of $\mathcal{N}(0,\operatorname{Id}_{d\times d}/d)$ to $\mathcal{S}_{g}$ preserves the conditionally-symmetric condition.

Notations for gradients. Before describing the analysis, we introduce several notations first. Recall from Claim 3.2 that the gradient of a neuron $v$ in the distribution $\mathcal{P}$ can be simplified given the conditionally-symmetric property. For each coordinate $1\leq i\leq d$, the gradient of neuron $v$ satisfies that $[\nabla_{v}]_{i}=\sum_{j\geq 0}\left[\nabla_{2j,v}\right]_{i}$, where $\nabla_{v}=\nabla_{v}L_{\infty}(\mathcal{P})$, $\nabla_{2j,v}=\nabla_{2j,v}L_{\infty}(\mathcal{P})$ denotes the gradient of $v$ for the $2j$-th loss, and $[\nabla_{v}]_{i}$ denotes the $i$-th coordinate of $\nabla_{v}$. Let $B_{1,2j}=b_{2j}+b_{2j}^{\prime}$ and $B_{2,2j}=b_{2j}^{\prime}$, where $b_{2j}$ and $b_{2j}^{\prime}$ are the Hermite coefficients of the $2j$-th loss given in Section 2. For a vector $w\in\mathcal{S}_{g}$, let $w^{(0)}$ denote a neuron with initialization $w$ in the initialization $\mathcal{P}^{(0)}$. Let $\mathcal{P}^{(t)}$ denote the $t$-th iterate of $\mathcal{P}^{(0)}$ following the update rule of equation (2.6).

Recall from Section 3.1 that the goal of Stage 1 is to show that a small fraction of neurons becomes basis-like, i.e. close to a basis $e_{i}$ times a scaling factor of $\operatorname{poly}(d)$ at the end of $\Theta_{\kappa}(d^{2}/\eta\log d)$ iterations for some $i\in[d]$. To facilitate the analysis, we maintain an inductive hypothesis throughout Stage 1 that provides an upper bound on the norm of a typical neuron during the update. We first introduce the set of neurons that will not become basis-like by the end of Stage 1.

Let $C_{0}$ be a large enough constant. Let $c_{0}=C_{0}\log d$ and $\mathcal{S}$ be the set of all vectors $w$ in $\mathcal{S}_{g}$ such that

where $\bar{w}=w/\|w\|$ denotes $w$ being normalized to norm $1$.

Based on the above definition, we introduce the following inductive hypothesis that shows the neurons in $\mathcal{S}$ remain “small and dense” (i.e. not basis-like) throughout Stage 1. This stage runs for $\Theta(\frac{d^{2}}{\eta\log d})$ iterations. We use $\kappa_{1}$ to denote a value that is less than $O(\exp(\operatorname{poly}(\kappa)))$.

In the setting of Theorem 3.1, let $T_{2}=\Theta(\frac{d^{2}}{\eta c_{0}\exp(\operatorname{poly}(\kappa))})$. There exists an increasing sequence $\left\{c_{t}\right\}_{t=1}^{T_{2}}$ where $c_{t}\leq\exp(\operatorname{poly}(\kappa))\log d$ such that for every $w\in\mathcal{S}$ and every $t\leq T_{2}$, the $t$-th iterate of the neuron $w^{(t)}$ with initialization $w^{(0)}=w$ satisfies that

Furthermore, for every coordinate $i\in[d]$, we have that in expectation,

Equation (A.2) and (A.3), which we also refer to as inductive hypothesis $\mathcal{H}_{1}$, show that the norm of any neuron in $\mathcal{S}$ will not grow beyond $O_{\kappa}(\log d/d)$. Hence they will not become basis-like during Stage 1.

The set $\mathcal{S}$ contains most neurons in $\mathcal{P}^{(0)}$ because by standard anti-concentration inequalities, the measure of the set $\mathcal{S}$ is at least $1-d^{-O(C_{0})}$. Hence, $1-\mu(\mathcal{S})$ is at most $d^{-O(C_{0})}$. Based on this fact, we state a simple claim on the norm of neurons that are not in $\mathcal{S}$ that will be used later:

To see that equation (A.4) is true, recall that the truncation of Algorithm 1 ensures that $\|w\|^{2}\leq 1/\lambda_{0}$. Combined with the fact that $1-\mu(\mathcal{S})\leq d^{-O(C)}$ and $\lambda_{0}=\Theta(1/\operatorname{poly}(d))$, we have that equation (A.4) holds for a sufficiently large constant $C_{0}$. This finishes our introduction of the inductive hypothesis $\mathcal{H}_{1}$. The proof of Proposition A.1 can be found in Section A.2.2.

Given the inductive hypothesis $\mathcal{H}_{1}$, we can state the formal result that corresponds to Stage 1.1 in Section 3.1. For a neuron distribution $\mathcal{P}$, let us first introduce the following notations, which corresponds to the population loss of the 0th and 2nd order tensor decompositions.

Based on the above notations, we show the following convergence result at the end of Stage 1.1.

In the setting of Theorem 3.1, suppose that Proposition A.1 holds. Let $T_{1}=\Theta\left(\frac{\operatorname{poly}(\kappa_{1})d\log d}{\eta}\right)$. Then, for every $t\geq T_{1}$, we have that $\Delta^{(t)},\delta_{+}^{(t)},\delta_{-}^{(t)}$ are all less than $\frac{c_{t}\operatorname{poly}(\kappa_{1})}{d^{2}}$, where $c_{t}$ is given in Proposition A.1.

The above result implies that after $T_{1}$ iterations, the 0th and 2nd order losses remain smaller than $c_{t}\operatorname{poly}(\kappa_{1})/d^{2}$. The proof of Lemma A.2 can be found in Section A.2.1.

Once Stage 1.1 is finished, recall from Section 3 that the higher order gradients begin to dominate the dynamic. Hence Algorithm 1 enters Stage 1.2. We introduce the following notations in order to state the formal result. Let $T_{2}^{\prime}=T_{2}-\frac{d^{2}}{\eta\operatorname{poly}\log(d)}$. For every $1\leq i\leq d$, let $\Gamma_{i}=\frac{1}{2B_{1,4}(a_{i}^{2}d)(\eta T_{2}^{\prime})}$. Let $\rho=\frac{\operatorname{poly}(\kappa_{1})\cdot\log d}{d}$. Here, by our assumption, we know that $a_{i}^{2}=\Theta({1}/{d^{2}})$. Since $T_{2}^{\prime}=\Theta({d^{2}}/(\eta\log d))$, we can see that $\Gamma_{i}=\Theta({\log d}/{d})$. Consider a coordinate $i\in[d]$. We define the set of good neurons whose $i$-th coordinate is larger than $\Gamma_{i}+\rho$ as

Then we define the set of bad neurons that have two large coordinates as

The following lemma shows that, among other statements, the neurons in $\mathcal{S}_{i,good}$ will win the lottery and become basis-like at the end of Stage 1.2 in the sense described below.

In the setting of Theorem 3.1, suppose that Proposition A.1 holds. At iteration $T_{2}$ (recall that $T_{2}$ is defined in Proposition A.1), the following holds for $\mathcal{S}_{i,good}$ and $\mathcal{S}_{i,bad}$:

For every $i\in[d]$ and every $v\in\mathcal{S}_{i,good}$, we have that

For every $1\leq i\leq d$ and every $v\in\mathcal{S}_{g}$, if there exists $j\neq i$ such that $|v_{i}^{(T_{2})}|$ and $|v_{j}^{(T_{2})}|$ are both greater than $\frac{2(\log d)^{2}}{\sqrt{d}}$, then the neuron $v$ is in the union of $\mathcal{S}_{i,bad}$ and $\mathcal{S}_{j,bad}$.

For every $i\in[d]$, the probability measure of $\mathcal{S}_{i,good}$ and $\mathcal{S}_{i,bad}$ satisfies that

In the above result, the set $\mathcal{S}_{i,good}$ contains neurons that become approximately a large scaling of the basis $e_{i}$ after $T_{2}$ iterations, a phenomenon that we term as winning the lottery ticket. The norm of these neurons become much larger than those in $\mathcal{S}$, whose norm is bounded by $O_{\kappa}(\log d/d)$. The set $\mathcal{S}_{i,bad}$ contains neurons whose coordinate $i$ might be large in the end, but not close to a basis. The final statement in this lemma shows that the probability measure of bad neurons is small compared to good neurons. Lemma A.3 is proved in Section A.3. This concludes Stage 1.

The second stage begins by reducing the gradient truncation parameter from $\lambda_{0}=\Theta(\frac{1}{\operatorname{poly}(d)})$ to $\lambda_{1}=\Theta(\frac{1}{\operatorname{poly}_{\kappa}(d)})$.As a remark, the rational for this technical twist is that the neurons do not grow too large Stage 1. This is useful for the error analysis later in the finite-width case. Recall from Section 3.2 that the goal of Stage 2 is to allow basis-like neurons to grow until they fit the target network with population loss at most $O(d^{1+Q})$.

The first substage of the analysis shows that the population loss reduces below $O(\frac{1}{d\log^{0.01}d})$, after $T_{3}=\Theta({d\log d}/{\eta})$ many iterations.

The second substage of the analysis shows that the population loss further reduces below $O\left({1}/{d^{1+Q}}\right)$, after $T_{4}=\Theta({d^{1+10Q}}/{\eta})$ many iterations.

To facilitate the analysis, we introduce an inductive hypothesis throughout Stage 2 that describes the behavior of the good and bad neurons. Let us introduce several notations first. Let the union of the bad neurons for all coordinates be given by

The set of potential neurons for coordinate $i\in[d]$ is given by

We remark that these are the set of neurons whose coordinate $i$ can become larger than $O(\frac{\operatorname{poly}\log(d)}{\sqrt{d}})$ at the end of Stage 1 (cf. Section A.2.1). The set of good neurons $\mathcal{S}_{i,good}$ is a subset of $\mathcal{S}_{i,pot}$. Let the union of the potential neurons for all coordinates be given by

We maintain the following running hypothesis that, among other things, specifies the behavior of the potential, good, and bad neurons in detail.

In the setting of Theorem 3.1, there exists a monotonically increasing sequence $\mathcal{c_{t}}_{t=T_{2}}^{T_{4}}$ such that $c_{T_{2}}=\operatorname{poly}(\log d)\leq c_{t}\leq d^{O(Q)}\leq d^{1/10}$ and for every $T_{2}<t\leq T_{4}$, the following list of properties holds for the neuron distribution $\mathcal{P}^{(t)}$:

For every $v\in\mathcal{S}_{g}$, we have that $\|v^{(t)}\|_{2}^{2}\leq 1/\lambda_{1}$. As a result, gradient truncation never happens during this stage.

For every $v\notin\mathcal{S}_{pot}$, we have that

For every $i\in[d]$, every $v\in\mathcal{S}_{i,pot}\backslash\mathcal{S}_{bad}$, and $j\not=i$, we have that

The probability mass of the set of bad neurons satisfies that

For every $i\in[d]$ and every $v\in\mathcal{S}_{i,good}$, we have that $\|v_{i}^{(t)}\|_{2}^{2}\geq\frac{1}{\lambda_{0}\operatorname{poly}(d)}$.

For every $i\in[d]$, the following claims regarding the set of potential neurons and bad neurons hold:

where $\kappa_{2}$ denotes $\exp(\operatorname{poly}(\kappa_{1}))$ and $\kappa_{1}$ denotes $\exp(\operatorname{poly}(\kappa))$.

We remark that in the above inductive hypothesis, equation (A.5) and (A.6) show similar conditions as equation (A.2) provided in Proposition A.1. For the rest of the section, we refer to the conclusion of Proposition A.4 as inductive hypothesis $\mathcal{H}_{2}$. The proof of Proposition A.4 can be found in Section A.4.1.

Given the inductive hypothesis, we can state the formal result that corresponds to Stage 2.1 in Section 3.2. We introduce the notation $\Delta^{(t)}=2b_{0}\left(\sum_{i=1}^{d}(\gamma_{i}^{(t)}+\beta_{i}^{(t)})-\sum_{i=1}^{d}a_{i}\right)$ that measures the average error of the neurons across all coordinates at iteration $t$. We show that by the end of $T_{3}=T_{2}+\Theta(d\log d/\eta)$ iterations, we have obtained a warm start neuron distribution for $\Delta^{(t)}$, $\left\{\beta_{1}^{(t)},\dots,\beta_{d}^{(t)}\right\}$, and $\left\{\gamma_{1}^{(t)},\dots,\gamma_{d}^{(t)}\right\}$. We state the result below.

In the setting of Theorem 3.1, suppose Proposition A.4 holds. There exists an iteration $T_{3}=T_{2}+\Theta(d\log d/\eta)$ such that at iteration $T_{3}$, the following holds:

The above result implies that the set of potential neurons has fit the $i$-th coordinate of the target network with error less than $o(1/d)$. The 0th order loss has also been reduced below $o(1/d)$. The proof of Lemma A.5 can be found in Appendix A.3.

In the end, we describe the formal result that corresponds to Stage 2.2 in Section 3.2. We construct a potential function to show that $\beta_{i}^{(t)}+\gamma_{i}^{(t)}$ converges to $a_{i}$ when $t\geq T_{3}$. After running for $T_{4}=T_{3}+\Theta(\frac{d^{1+10Q}}{\eta})$ many iterations, we show that a certain set of potential neurons has converged to $a_{i}$ with error at most $O(1/d^{2+Q})$, for every $1\leq i\leq d$.

The result is shown in Lemma A.6 below. We introduce the following notations for defining the potential function at iteration $t$:

where $C_{1},C_{2}$ denote two sufficiently large constants. Consider the following functions (recall that $\Delta_{+}$ and $\Delta_{-}$ have been defined in Stage 1):

Let $\beta_{+}^{(t)}=\frac{1}{C}\max_{i\in[d]}\{\beta_{i}^{(t)}\}$. Let $\Phi^{(t)}=\max\{\Phi_{+}^{(t)},\Phi_{-}^{(t)},\beta_{+}^{(t)}\}$ be our potential function. Lemma A.5 implies that by the end of $t=T_{3}$ iterations, we have that $\delta_{-}^{(t)},\delta_{+}^{(t)},\beta_{+}^{(t)},\Delta_{+}^{(t)},\Delta_{-}^{(t)}$ are all less than $O\left({1}/{d\log^{0.01}d}\right)$. Hence $\Phi^{(T_{3})}\leq O(1/(d\log^{0.01}d))$. The result below shows that after iteration $T_{3}$, $\Phi^{(t)}$ further decreases whenever $\Phi^{(t)}$ is at least $O(\frac{\operatorname{poly}(\kappa_{2})c_{t})}{d^{2}})$.

In the setting of Theorem 3.1, suppose that Proposition A.4 holds. Let $C_{1}$ be a fixed constant. For any $T_{3}<t\leq T_{4}$, as long as $\Phi^{(t)}\geq\frac{\operatorname{poly}(\kappa_{2})c_{t}}{d^{2}}$ (recalling that $c_{t}$ is defined in Proposition A.4) we have that

By combining the results of Stage 1 and Stage 2, we are ready to prove Theorem 3.1.

When Proposition A.1 and A.4 hold, using the induction hypothesis in equation (A.7), we have that for the infinite-width case, the population loss $L_{\infty}(\mathcal{P}^{(t)})$ satisfies:

where the first term comes from the 0th order loss and the second term comes from 2nd and higher order losses. This claim also implies that

At the beginning of Stage 2.2, by Lemma A.5, we know that $\Phi^{(T_{3})}\leq 1/(d\log^{0.01}d)$. During Stage 2.2, by Lemma A.6, as long as $\Phi^{(t)}\geq O_{\kappa}(c_{t}/d^{2})$, $\Phi^{t+1}\leq\Phi^{(t)}\leq\Phi^{(t)}(1-O(\Phi^{(t)}))$. Hence, after at most $d^{1+O(Q)}/\eta$ iterations (or $T_{4}-T_{3}$ more precisely), $\Phi^{(T_{4})}$ reduces to below $O(d^{1+Q})$. Applying this result to equation (A.11), we conclude that $L_{\infty}(\mathcal{P}^{(T_{4})})\leq O(1/d^{1+Q})$.

A.1 Stage 1.1: Proof of Convergence for 0th and 2nd Order Tensors

This section provides the proof of Lemma A.2 is organized as follows.

In Proposition A.7, we first show that the gradients from 4th and higher order tensor decompositions are small compared to that of the 0th and 2nd order tensor decompositions.

The above shows that the dynamic is mainly dominated by the 0th and 2nd losses initially. In Proposition A.8 and Proposition A.9, we show the gradient update of the 0th and 2nd order. Based on these, we show the proof Lemma A.2 at the end of this subsection.

We first show that the 4th and higher order tensor gradients do not have much contribution to the gradient, for all the neurons in $\mathcal{S}$. We introduce the following notations for convenience. For a neuron distribution $\mathcal{P}$, let the following denote the gradient of $v$ involving only other neurons $w$.

Recall that $\nabla_{2j,v}$ is the gradient of $v$ for the 2j-th tensor (cf. equation (3.3)). Let

The following result provides an upper bound on the higher order gradients.

In the setting of Lemma A.2, suppose Proposition A.1 holds. Then there exists an absolute constant $C>0$ such that for every $i\in[d]$ and $v\in\mathcal{S}_{g}$, at the $t$-th iteration for $t\leq T_{2}$, the neuron $v$ from distribution $\mathcal{P}^{(t)}$ satisfies that

Moreover, the gradient from the network satisfies

As a corollary, for every $v\in\mathcal{S}\subseteq\mathcal{S}_{g}$, we have that

Let us focus on $\nabla_{4,v}$ first. We now bound each term in $[\nabla_{4,v}]_{i}$ in equation (3.3) separately.

For the signal term in the gradient, $a_{i}\langle e_{i},v\rangle\langle e_{i},\bar{v}\rangle^{2}$, because $a_{i}\leq\kappa/d$, we have

Another term in the gradient is (again, using the fact that for $w$ with $w^{(0)}\in\mathcal{S}$, $\|w\|_{\infty}^{2}\leq\frac{c_{t}}{d}$):

The last term in the gradient is given by:

Combining Eq (A.13), Eq (A.14), Eq (A.15) and Eq (A.16), we obtain that

For $\Delta_{2j,v}$, with $j\geq 3$, we can apply the same calculation as above, and show that

Since $\nabla_{\geq 4,v}=\sum_{j\geq 2}\nabla_{2j,v}$ and $\sum_{j}b_{2j}=O(1)$ we complete the proof. ∎

Based on the above result, we describe the dynamic of the 0th order tensor in the following proposition.

$|\Delta^{(t)}|\leq\frac{c_{t}\operatorname{poly}(\kappa_{1})}{d^{2}}$.

If $\Delta^{(t)}>0$, then $\Delta^{(t)}\leq\delta_{-}^{(t)}\left(1-\frac{1}{16\kappa^{2}d}\right)$. If $\Delta^{(t)}<0$, then $|\Delta^{(t)}|\leq\left(1-\frac{1}{5\kappa_{1}^{5}d}\right)\delta_{+}^{(t)}$.

Moreover, when $\Delta^{(t)}\geq\max\left\{8\kappa^{2}\delta_{+}^{(t)},\frac{c_{t}\operatorname{poly}(\kappa_{1})}{d^{2}}\right\}$ , it holds that

where $\mathcal{S}^{(t)}_{+}$ is the set of all $i\in[d]$ with $\delta_{i}^{(t)}\geq 0$ and $\left|\mathcal{S}^{(t)}_{+}\right|$ denotes its cardinality.

Consider the iteration $t$, we have that for every $i$ and every $v\in\mathcal{S}_{g}$, the update of $v_{i}^{(t)}$ is given as:

Hence, using Proposition A.7 and inequality (A.4), it holds that

This implies that for every sufficiently small $\eta\leq\lambda_{0}^{2}$ and $\Lambda\leq\lambda_{0}^{2}$, it holds:

Let us consider two cases when $\operatorname*{abs}{\Delta^{(t)}}=\Omega(\kappa_{1}^{8}c_{t}/d^{2})$.

Now, consider a value $\rho=\frac{1}{4\kappa^{2}d}$, when $\Delta^{(t)}\geq\delta_{-}^{(t)}(1-\frac{\rho}{2})$, it holds that $(1-\frac{\rho}{2})\Delta^{(t)}\geq\delta_{-}^{(t)}(1-\rho)$. Therefore, when $\Delta^{(t)}\geq\delta_{-}^{(t)}(1-\frac{\rho}{2})$, Eq (A.21) implies that

Summing up all $\delta_{i}^{(t)}$ with $\delta_{i}^{(t)}\leq 0$, this implies that

Combine the above inequality with inequality (A.20), we have that (using $\Delta^{(t)}\geq 0$ so that $\Delta_{+}^{(t)}\geq\Delta_{-}^{(t)}$):

Therefore we conclude that when $\Delta^{(t)}\geq\Omega\left(\frac{c_{t}\kappa^{8}}{d^{2}}\right)$ and $\Delta^{(t)}\geq\delta_{-}^{(t)}(1-\frac{\rho}{2})$, it must holds that

Here the second inequality comes from Eq (A.23). This implies that $\Delta^{(t+1)}$ will decrease faster than $\delta_{-}^{(t+1)}$ at the next iteration. Hence, when $\Delta^{(t)}\geq\Omega\left(\frac{c_{t}\kappa^{8}}{d^{2}}\right)$, then $\Delta^{(t)}\geq\delta_{-}^{(t)}(1-\frac{\rho}{2})$ can never happen. Hence, by our choice of $\rho$, we conclude that as long as $\Delta^{(t)}\geq\Omega\left(\frac{c_{t}\kappa^{8}}{d^{2}}\right)$, then

On the other hand, even when $\Delta^{(t)}\leq\delta_{-}^{(t)}(1-\frac{\rho}{2})$ but $\Delta^{(t)}=\Omega\left(\frac{\kappa_{1}^{8}c_{t}}{d^{2}}\right)$ , we still have that for every $i$ with $\delta_{i}^{(t)}\leq 0$, by Eq (A.22):

Hence, as long as $\Delta^{(t)}=\Omega\left(\frac{\kappa_{1}^{8}c_{t}}{d^{2}}\right)$, we will always have

Combining the above with equation (A.20), we have that

Here, we are using the fact that $\Delta^{(t)}_{-}\leq\Delta^{(t)}_{+}\leq\delta_{+}^{(t)}|\mathcal{S}^{(t)}_{+}|$. Now, this implies that when $\Delta^{(t)}\geq\max\left\{8\kappa^{2}\delta_{+}^{(t)},\Omega\left(\frac{\kappa_{1}^{8}c_{t}}{d^{2}}\right)\right\}$ , it also holds that

The proof follows by a similar argument to Case 1. ∎

Based on the above result, next we describe the dynamic of the 2nd order tensor.

Moreover, when $\Delta^{(t)}>0$, we have the following improved bound for $\delta_{+}^{(t+1)}$:

By the update rule, we can obtain (in Eq (A.19)) that

which proves the condition. On the other hand when $\delta_{-}^{(t)}\leq\frac{c_{t}\operatorname{poly}(\kappa_{1})}{d^{2}}$, we directly completes the proof by choosing a larger poly in $\frac{c_{t}\operatorname{poly}(\kappa_{1})}{d^{3}}$. We can apply the same argument for $\delta_{+}$, and the improved bound for the case when $\Delta^{(t)}>0$. ∎

A.1.1 Proof of the Main Lemma

Now we are ready to show the final convergence lemma. We first provide the following claim that shows on average, each coordinate of the neuron distribution lies in a bounded range. This also proves the first equation of (A.3) in the inductive hypothesis $\mathcal{H}_{1}$.

Hence, we only need to consider $t\in\mathcal{T}$, for these iterations, by Proposition A.9 we know that

On the other hand by Proposition A.8, we have that when $\Delta^{(t)}\geq\max\left\{8\kappa^{2}\delta_{+}^{(t)},\frac{c_{t}\operatorname{poly}(\kappa_{1})}{d^{2}}\right\}$ , it holds that:

Now, let us define $\gamma_{0}=\delta_{+}^{0}\leq\frac{2\kappa}{d}$, with $\gamma_{t+1}=\gamma_{t}\left(1-\eta\frac{1}{4\kappa d}\right)+\eta\frac{c_{t}\operatorname{poly}(\kappa_{1})}{d^{3}}$ for every $t\in\mathcal{T}$ and $\gamma_{t+1}=\gamma_{t}+\eta\frac{c_{t}\operatorname{poly}(\kappa_{1})}{d^{3}}$ otherwise. We know that as long as $\Delta^{(t)}\geq\frac{c_{t}\operatorname{poly}(\kappa_{1})}{d^{2}}$, we have:

Clearly, by Proposition A.9 once $\delta_{+}^{(t)}$ or $\delta_{-}^{(t)}\leq\frac{c_{t}\operatorname{poly}(\kappa_{1})}{d^{2}}$, they will stay within the interval for the next iterations. By Proposition A.8, after both $\delta_{+}^{(t)}$ and $\delta_{-}^{(t)}\leq\frac{c_{t}\operatorname{poly}(\kappa_{1})}{d^{2}}$, we know $\Delta^{(t)}$ will be within the interval as well.

Hence, we just need to consider the first time that $\delta_{+}^{(t)}$ and $\delta_{-}^{(t)}$ goes outside the interval. Following Proposition A.9, we know that when $\delta_{+}^{(t)}\geq\frac{c_{t}\operatorname{poly}(\kappa_{1})}{d^{2}}$, it holds:

which gives the convergence error rate of $\delta_{+}$ after $T_{1}$ iterations. The same holds for $\delta_{-}^{(t)}$. ∎

Finally, we have an estimate of how big each coordinate is for the neurons at the end of Stage 1.1, which can be given by the output layer weights $\{a_{i}\}_{i=1}^{d}$. We show the following claim, which will be used in the proof of Stage 2.

In the setting of Lemma A.2, at iteration $T_{1}$ (recalling that $T_{1}=\Theta(\frac{\operatorname{poly}(\kappa_{1})d\log d}{\eta})$), for every $v\in\mathcal{S}_{g}$ and every $i\in[d]$, we have that

Let us first show the upper bound. For every $v^{(0)}\in\mathcal{S}_{g}$. By the update rule, we have that

On the other hand, we have that by Eq (A.18), it holds: