Beyond NTK with Vanilla Gradient Descent: A Mean-Field Analysis of
  Neural Networks with Polynomial Width, Samples, and Time
Despite recent theoretical progress on the non-convex optimization of
two-layer neural networks, it is still an open question whether gradient
descent on neural networks without unnatural modifications can achieve better
sample complexity than kernel methods. This paper provides a clean mean-field
analysis of projected gradient flow on polynomial-width two-layer neural
networks. Different from prior works, our analysis does not require unnatural
modifications of the optimization algorithm. We prove that with sample size $n
= O(d^{3.1})$ where $d$ is the dimension of the inputs, the network trained
with projected gradient flow converges in $\text{poly}(d)$ time to a
non-trivial error that is not achievable by kernel methods using $n \ll d^4$
samples, hence demonstrating a clear separation between unmodified gradient
descent and NTK. As a corollary, we show that projected gradient descent with a
positive learning rate and a polynomial number of iterations converges to low
error with the same sample complexity.
Arvind Mahankali, Jeff Z. Haochen, Kefan Dong, Margalit Glasgow, Tengyu Ma
Beyond NTK with Vanilla Gradient Descent: A Mean-Field Analysis of Neural Networks with Polynomial Width, Samples, and Time

###### Abstract

Despite recent theoretical progress on the non-convex optimization of two-layer neural networks, it is still an open question whether gradient descent on neural networks without unnatural modifications can achieve better sample complexity than kernel methods. This paper provides a clean mean-field analysis of projected gradient flow on polynomial-width two-layer neural networks. Different from prior works, our analysis does not require unnatural modifications of the optimization algorithm. We prove that with sample size \(n=O(d^{3.1})\) where \(d\) is the dimension of the inputs, the network converges in polynomially many iterations to a non-trivial error that is not achievable by kernel methods using \(n\ll d^{4}\) samples, hence demonstrating a clear separation between unmodified gradient descent and NTK.

## 1 Introduction

Training neural networks requires optimizing non-convex losses, which is often practically feasible but still not theoretically understood. The lack of understanding of non-convex optimization also limits the design of new principled optimizers for training neural networks that use theoretical insights.

Early analysis on optimizing neural networks with linear or quadratic activations (Du and Lee, 2018; Gunasekar et al., 2018; Li et al., 2018; Gunasekar et al., 2018; Soltanolkotabi et al., 2018; Hardt and Ma, 2016; Kawaguchi, 2016; Hardt et al., 2018; Oymak and Soltanolkotabi, 2020) relies on linear algebraic tools that do not extend to nonlinear and non-quadratic activations. The neural tangent kernel (NTK) approach analyzes nonconvex optimization under certain hyperparameter settings, e.g., when the initialization scale is large and the learning rate is small (see, e.g., Du et al. (2018); Jacot et al. (2018); Li and Liang (2018); Arora et al. (2019); Daniely et al. (2016)). However, subsequent research shows that neural networks trained with practical hyperparameter settings typically outperform their corresponding NTK kernels (Arora et al., 2019). Furthermore, the initialization and learning rate under the NTK regime does not yield optimal generalization guarantees (Wei et al., 2019, Chizat and Bach, 2018; Woodworth et al., 2020; Ghorbani et al., 2020).
Many recent works study modified versions of stochastic gradient descent (SGD) and prove sample complexity and runtime guarantees beyond NTK (Damian et al., 2022; Abbe et al., 2022; Mousavi-Hosseini et al., 2022; Li et al., 2020; Abbe et al., 2021; Allen-Zhu et al., 2019; Allen-Zhu and Li, 2019; Xu and Du, 2023; Wu et al., 2023; Daniely and Malach, 2020; Allen-Zhu and Li, 2020; Suzuki and Akiyama, Telgarsky, 2022; Chen et al., 2020). These modified algorithms often contain multiple stages that optimize different blocks of parameters and/or use different learning rates or regularization strengths. For example, the work of Li et al. (2020) uses a non-standard parameterization for two-layer neural networks and runs a two-stage algorithm with sharply changing gradient clipping strength; Damian et al. (2022) use one step of gradient descent with a large learning rate and then optimize only the last layer of the neural net. However, oftentimes vanilla (stochastic) gradient descent with a constant learning rate empirically converges to a minimum with good generalization error. Thus, the modifications are arguably artifacts tailored to the analysis, and to some extent, over-using the modification may obscure the true power of gradient descent.

Another technique to analyze optimization dynamics uses the mean-field approach (Chizat and Bach, 2018; Mei et al., 2018; Mei et al., 2019; Sirignano and Spiliopoulos, 2018; Dou and Liang, 2021; Abbe et al., 2022; Javanmard et al., 2020; Sirignano and Spiliopoulos, 2020; Wei et al., 2019), which views the collection of weight vectors as a (discrete) distribution over \(\mathbb{R}^{d}\) (where \(d\) is the input dimension) and approximates its evolution by an infinite-width neural network, where the weight vector distribution is continuous and evolves according to a partial differential equation. However, these works do not provide an end-to-end polynomial runtime bound on the convergence to a global minimum in the concrete setting of two-layer neural networks (without modifying gradient descent). For example, the works of Chizat and Bach (2018) and Mei et al. (2018) do not provide a concrete bound on the number of iterations needed for convergence. Mei et al. (2019) provide a coupling between the trajectories of finite and infinite-width networks with exponential growth of coupling error but do not apply it to a concrete setting to obtain a global convergence result with sample complexity guarantees. (See more discussion below and in Section 3.) Wei et al. (2019) achieve a polynomial number of iterations but require exponential width and an artificially added noise. In other words, these works, without modifying gradient descent, cannot prove convergence to a global minimizer with polynomial width and iterations.

In this paper, we provide a mean-field analysis of projected gradient flow on two-layer neural networks with _polynomial width_ and quartic activations. Under a simple data distribution, we demonstrate that the network converges to a non-trivial error in _polynomial time_ with _polynomial samples_. Notably, our results show a sample complexity that is superior to the NTK approach.

Concretely, the neural network is assumed to have unit-norm weight vectors and no bias, and the second-layer weights are all \(\frac{1}{m}\) where \(m\) is the width of the neural network. The data distribution is uniform over the sphere. The target function is of the form \(y(x)=h(q_{*}^{\top}x)\) where \(h\) is an _unknown_ quartic link function and \(q_{*}\) is an unknown unit vector. Our main result (Theorem 3.4) states that with \(n=O(d^{3.1})\) samples, a polynomial-width neural network with random initialization converges in polynomial time to a non-trivial error, which is statistically not achievable by any kernel method with an inner product kernel using \(n\ll d^{4}\) samples. To the best of our knowledge, our result is the first to demonstrate the advantage of _unmodified_ gradient descent on neural networks over kernel methods.

The rank-one structure in the target function, also known as the single-index model, has been well-studied in the context of neural networks as a simplified case to demonstrate that neural networks can learn a latent feature direction \(q_{*}\) better than kernel methods (Mousavi-Hosseini et al., 2022; Abbe et al., 2022; Damian et al., 2022; Bietti et al., 2022). Many works on single-index models study (stochastic) gradient descent in the setting where only a single vector (or "neuron") in \(\mathbb{R}^{d}\) is trained. This includes earlier works where the link function is monotonic, or the convergence is analyzed with quasi-convexity (Kakade et al., 2011; Hazan et al., 2015; Mei et al., 2018; Soltanolkotabi, 2017), along with more recent work (Tan and Vershynin, 2019; Vardi et al., 2021; Arous et al., 2021) for more general link functions. Since these works only train a single neuron, they have limited expressivity in comparison to a neural network, and can only achieve zero loss when the link function equals the activation. We stress that in our setting, the link function \(h\) is unknown, not necessarily monotonic, and does not need to be equal to the activation function. We show that in this setting, the first layer weights will converge to a _distribution_ of neurons that are correlated with but
not exactly equal to \(q_{\star}\), so that even without bias terms, their mixture can represent the link function \(h\). Our analysis demonstrates that gradient descent with a consistent infinitesimal learning rate can _simultaneously_ learn the feature \(q_{\star}\) and the link function \(h\), which is a key challenge that is side-stepped in previous works on neural-networks which use two-stage algorithms (Mousavi-Hosseini et al., 2022; Abbe et al., 2021, 2022; Damian et al., 2022; Abbe et al., 2023; Barak et al., 2022). The main novelty of our population dynamics analysis is designing a potential function that shows that the iterate stays away from the saddle points.

Our sample complexity results leverage a coupling between the dynamics on the empirical loss for a finite-width neural network and the dynamics on the population loss for an infinite-width neural network. The main challenge stems from the fact that some exponential coupling error growth is inevitable over a certain period of time when the dynamics resemble a power method update. Heavily inspired by Li et al. (2020), we address this challenge by using a direct and sharp comparison between the growth of the coupling error and the growth of the signal. In contrast, a simple exponential growth bound on the coupling error similar to the bound of Mei and Montanari (2019) would result in a \(d^{O(1)}\) sample complexity, which is not sufficient to outperform NTK.

## 2 Preliminaries and Notations

We use \(O(\cdot),\lesssim,\gtrsim\) to hide only absolute constants. Formally, every occurrence of \(O(x)\) in this paper can be simultaneously replaced by a function \(f(x)\) where \(|f(x)|\leq C|x|,\forall x\in\mathbb{R}\) for some universal constant \(C>0\) (each occurrence can have a different universal constant \(C\) and \(f\)). We use \(a\lesssim b\) as a shorthand for \(a\leq O(b)\). Similarly, \(\Omega(x)\) is a placeholder for some \(g(x)\) where \(|g(x)|\geq|x|/C,\forall x\in\mathbb{R}\) for some universal constant \(C>0\). We use \(a\gtrsim b\) as a shorthand for \(a\geq\Omega(b)\) and \(a\asymp b\) as a shorthand to indicate that \(a\gtrsim b\) and \(a\lesssim b\) simultaneously hold.

**Legendre Polynomials.** We summarize the necessary facts about Legendre polynomials below, and present related background more comprehensively in Appendix A. Let \(P_{k,d}:[-1,1]\to\mathbb{R}\) be the degree-\(k\)_unnormalized_ Legendre polynomial (Atkinson and Han, 2012), and \(\overline{P}_{k,d}(t)=\sqrt{N_{k,d}}P_{k,d}(t)\) be the _normalized_ Legendre polynomial, where \(N_{k,d}\triangleq\binom{d+k-1}{d-1}-\binom{d+k-3}{d-1}\) is the normalizing factor. The polynomials \(\overline{P}_{k,d}(t)\) form an orthonormal basis for the set of square-integrable functions over \([-1,1]\) with respect to the measure \(\mu_{d}(t)\triangleq(1-t^{2})^{\frac{d-3}{2}}\frac{\Gamma(d/2)}{\Gamma((d-1)/2 )}\frac{1}{\sqrt{\pi}}\), i.e., the density of \(u_{1}\) when \(u=(u_{1},\cdots,u_{d})\) is uniformly drawn from sphere \(\mathbb{S}^{d-1}\). Hence, for every function \(h:[-1,1]\to\mathbb{R}\) such that \(\mathbb{E}_{t\sim\mu_{d}}[h(t)^{2}]<\infty\), we can define \(\hat{h}_{k,d}\triangleq\mathbb{E}_{t\sim\mu_{d}}[h(t)\overline{P}_{k,d}(t)]\) and consequently, we have \(h(t)=\sum_{k=0}^{\infty}\hat{h}_{k,d}\overline{P}_{k,d}(t)\).

## 3 Main Results

We will formally define the data distribution, neural networks, projected gradient flow, and assumptions on the problem-dependent quantities and then state our main theorems.

_Target function._ The ground-truth function \(y(x):\mathbb{R}^{d}\to\mathbb{R}\) that we aim to learn has the form \(y(x)=h(q_{\star}^{\top}x)\), where \(h:\mathbb{R}\to\mathbb{R}\) is an _unknown_ one-dimensional _even_ quartic polynomial (which is called a link function), and \(q_{\star}\) is an _unknown_ unit vector in \(\mathbb{R}^{d}\). Note that \(h(s)\) has the Legendre expansion \(h(s)=\hat{h}_{0,d}+\hat{h}_{2,d}\overline{P}_{2,d}(s)+\hat{h}_{4,d}\overline{ P}_{4,d}(s)\).

_Two-layer neural networks._ We consider a two-layer neural network where the first-layer weights are all unit vectors and the second-layer weights are fixed and all the same. Let \(\sigma(\cdot)\) be the activation function, which can be different from \(h(\cdot)\). Using the mean-field formulation, we describe the neural network using the distribution of first-layer weight vectors, denoted by \(\rho\):

\[f_{\rho}(x)\triangleq\mathbb{E}_{u\sim\rho}[\sigma(u^{\top}x)]\,. \tag{3.1}\]

For example, when \(\rho=\text{unif}(\{u_{1},\ldots,u_{m}\})\), is a discrete, uniform distribution supported on \(m\) vectors \(\{u_{1},\ldots,u_{m}\}\), then \(f_{\rho}(x)=\frac{1}{m}\sum_{i=1}^{m}\sigma(u_{i}^{\top}x)\), i.e. \(f_{\rho}\) corresponds to a finite-width neural network whose first-layer weights are \(u_{1},\ldots,u_{m}\). For a continuous distribution \(\rho\), the function \(f_{\rho}(\cdot)\) can be viewed as an
infinite-width neural network where the weight vectors are distributed according to \(\rho\) (and can be viewed as taking the limit as \(m\to\infty\) of the finite-width neural network). We assume that the weight vectors have unit norms, i.e. the support of \(\rho\) is contained in \(\mathbb{S}^{d-1}\). The activation \(\sigma:\mathbb{R}\to\mathbb{R}\) is assumed to be a fourth-degree polynomial with Legendre expansion \(\sigma(s)=\sum_{k=0}^{4}\hat{\sigma}_{k,d}\overline{P}_{k,d}(s)\).

The simplified neural network defined in Eq. (3.1), even with infinite width (corresponding to a continuous distribution \(\rho\)), has a limited expressivity due to the lack of biases and trainable second-layer weights. We characterize the expressivity by the following lemma:

**Lemma 3.1** (Expressivity).: _Let \(\gamma_{2}=\hat{h}_{2,d}/\hat{\sigma}_{2,d}\) and \(\gamma_{4}=\hat{h}_{4,d}/\hat{\sigma}_{4,d}\). Suppose for some \(d\) and \(q_{\star}\), there exists a network \(\rho\) such that \(f_{\rho}(x)=h(q_{\star}^{\top}x)\) on \(\mathbb{S}^{d-1}\). Then we have \(\hat{\sigma}_{0,d}=\hat{h}_{0,d}\), and \(0\leq\gamma_{2}^{2}\leq\gamma_{4}\leq\gamma_{2}\leq 1\). Moreover, if this condition holds with strict inequalities, then for sufficiently large \(d\), there exists a network \(\rho\) such that \(f_{\rho}(x)=h(q_{\star}^{\top}x)\) on \(\mathbb{S}^{d-1}\). (A more explicit version is stated in Appendix D.1.)_

Informally, an almost sufficient and necessary condition to have \(f_{\rho}(x)=h(q_{\star}^{\top}x)\) for some \(\rho\) is that there exists a random variable \(w\) supported on \([0,1]\) such that \(\mathbb{E}[w^{2}]\approx\gamma_{2}\) and \(\mathbb{E}[w^{4}]\approx\gamma_{4}\), which is equivalent to \(0\leq\gamma_{2}^{2}\leq\gamma_{4}\leq\gamma_{2}\leq 1.\) In particular, assuming the existence of such a random variable \(w\), the perfectly-fit network \(\rho\) that fits the target function has the form

\[q_{\star}^{\top}u\stackrel{{ d}}{{=}}w,\text{ and }u-q_{\star}q_{ \star}^{\top}u\mid q_{\star}^{\top}u\text{ is uniformly distributed in the subspace orthogonal to }q_{\star}\,. \tag{3.2}\]

Motivated by this lemma, we will assume that \(\gamma_{2},\gamma_{4}\) are universal constants that satisfy \(0\leq\gamma_{2}^{2}\leq\gamma_{4}\leq\gamma_{2}\leq 1\), and \(d\) is chosen to be sufficiently large (depending on the choice of \(\gamma_{2}\) and \(\gamma_{4}\)). In addition, we also assume that \(\gamma_{4}\leq O(\gamma_{2}^{2})\), that is, the equality \(\gamma_{2}^{2}\leq\gamma_{4}\) is somewhat tight -- this ensures that the distribution of \(w=q_{\star}^{\top}u\) under the perfectly-fitted neural network is not too spread-out. We also assume that \(\gamma_{2}\) is smaller than a sufficiently small universal constant. This ensures that the distribution of \(q_{\star}^{\top}u\) under the perfectly-fitted network does not concentrate on \(1\), i.e. the distribution of \(u\) is not merely a point mass around \(q_{\star}\). In other words, this assumption restricts our setting to the most interesting case where the landscape has bad saddle points (and thus is fundamentally more challenging to analyze). We also assume for simplicity that \(\hat{\sigma}_{0,d}=\hat{h}_{0,d}\) (because otherwise, the activation introduces a constant bias that prohibits perfect fitting), even though adding a trainable scalar to the neural network formulation can remove the assumption. In summary, we make the following formal assumptions on the Legendre coefficients of the link function and the activation function.

**Assumption 3.2**.: _Let \(\gamma_{2}=\hat{h}_{2,d}/\hat{\sigma}_{2,d}\) and \(\gamma_{4}=\hat{h}_{4,d}/\hat{\sigma}_{4,d}\). We first assume \(\gamma_{4}\geq 1.1\cdot\gamma_{2}^{2}\). For any universal constant \(c_{1}>1\), we assume that \(\hat{\sigma}_{2,d}^{2}/c_{1}\leq\hat{\sigma}_{4,d}^{2}\leq c_{1}\cdot\hat{ \sigma}_{2,d}^{2}\), and \(\gamma_{4}\leq c_{1}\gamma_{2}^{2}\). For a sufficiently small universal constant \(c_{2}>0\) (which is chosen after \(c_{1}\) is determined), we assume \(0\leq\gamma_{2}\leq c_{2}\). We also assume that \(d\) is larger than a sufficiently large constant \(c_{3}\) (which is chosen after \(c_{1}\) and \(c_{2}\).) We also assume \(\hat{h}_{0,d}=\hat{\sigma}_{0,d}\), and \(\hat{h}_{1,d}=\hat{h}_{3,d}=0\)._

Our intention is to replicate the ReLU activation as well as possible with quartic polynomials; our assumption that \(\hat{\sigma}_{2,d}^{2}\asymp\hat{\sigma}_{4,d}^{2}\) is indeed satisfied by the quartic expansion of ReLU because \(\widehat{\operatorname{relu}_{2,d}}\asymp d^{-1/2}\) and \(\widehat{\operatorname{relu}_{2,d}}\asymp d^{-1/2}\) (see Proposition A.3). Following the convention defined in Section 2, we will simply write \(\hat{\sigma}_{4,d}^{2}\asymp\hat{\sigma}_{2,d}^{2}\), \(\gamma_{4}\geq 1.1\cdot\gamma_{2}^{2}\), and \(\gamma_{4}\lesssim\gamma_{2}^{2}\).

_Data distribution, losses, and projected gradient flow._ The population data distribution is assumed to be \(\mathbb{S}^{d-1}\). We draw \(n\) training examples \(x_{1},\ldots,x_{n}\stackrel{{\mathrm{i.i.d}}}{{\sim}}\mathbb{S}^{d-1}\). Thus, the population and empirical mean-squared losses are:

\[L(\rho)=\frac{1}{2}\cdot\operatorname*{\mathbb{E}}_{x\sim\mathbb{S}^{d-1}} \left(f_{\rho}(x)-y(x)\right)^{2},\qquad\qquad\text{ and }\quad\widehat{L}(\rho)=\frac{1}{2n}\sum_{i=1}^{n}\left(f_{\rho}(x_{i})-y(x_{i}) \right)^{2}\,. \tag{3.3}\]

To ensure that the weight vectors remain on \(\mathbb{S}^{d-1}\), we perform projected gradient flow on the empirical loss. We start by defining the gradient of the population loss \(L\) with respect to a particle \(u\) at \(\rho\) and the
corresponding Riemannian gradient (which is simply the projection of the gradient to the tangent space of \(\mathbb{S}^{d-1}\)):

\[\nabla_{u}L(\rho)=\operatorname*{\mathbb{E}}_{x\sim\mathbb{S}^{d-1}}\left[(f_{ \rho}(x)-y(x))\sigma^{\prime}(u^{\top}x)x\right]\,,\qquad\qquad\text{and}\quad \operatorname{grad}_{u}L(\rho)=(I-uu^{\top})\nabla_{u}L(\rho)\,. \tag{3.4}\]

Here we interpret \(L(\rho)\) as a function of a collection of particles (denoted by \(\rho\)) and \(\nabla_{u}L(\rho)\) as the partial derivative with respect to a single particle \(u\) evaluated at \(\rho\). Similarly, the (Riemannian) gradient of the empirical loss \(\widehat{L}\) with respect to the particle \(u\) is defined as

\[\nabla_{u}\widehat{L}(\rho)=\frac{1}{n}\sum_{i=1}^{n}(f_{\rho}(x_{i})-y(x_{i}) )\sigma^{\prime}(u^{\top}x_{i})x_{i}\,,\qquad\qquad\text{and}\quad\operatorname {grad}_{u}\widehat{L}(\rho)=(I-uu^{\top})\text{grad}_{u}\widehat{L}(\rho)\,.\]

**Population, Infinite-Width Dynamics.** Let the initial distribution \(\rho_{0}\) of the infinite-width neural network be the uniform distribution over \(\mathbb{S}^{d-1}\). We use \(\chi\) to denote a particle sampled uniformly at random from the initial distribution \(\rho_{0}\). A particle initialized at \(\chi\) follows a deterministic trajectory afterwards -- we use \(u_{t}(\chi)\) to denote the location, at time \(t\), of the particle that was initialized at \(\chi\). Because \(u_{t}(\cdot)\) is a deterministic function, we can use \(\chi\) to index the particles at any time based on their initialization. The projected gradient flow on an infinite-width neural network and using population loss \(L\) can be described as

\[\forall\chi\in\mathbb{S}^{d-1},u_{0}(\chi) =\chi\,, \tag{3.5}\] \[\frac{du_{t}(\chi)}{dt} =-\text{grad}_{u_{t}}L(\rho_{t})\,,\] (3.6) \[\text{and }\rho_{t} =\text{distribution of }u_{t}(\chi)\text{ (where }\chi\sim\rho_{0})\,. \tag{3.7}\]

**Empirical, Finite-Width Dynamics.** The training dynamics of a neural network with width \(m\) can be described in this language by setting the initial distribution to be a discrete distribution uniformly supported on \(m\) initial weight vectors. The update rule will maintain that at any time, the distribution of neurons is uniformly supported over \(m\) items and thus still corresponds to a width-\(m\) neural network. Let \(\chi_{1},\dots,\chi_{m}\overset{\text{i.i.d}}{\sim}\mathbb{S}^{d-1}\) be the initial weight vectors of the width-\(m\) neural network, and let \(\hat{\rho}_{0}=\text{unif}(\{\chi_{1},\dots,\chi_{m}\})\) be the uniform distribution over these initial neurons. We use \(\chi\in\{\chi_{1},\dots,\chi_{m}\}\) to index neurons and denote a single initial neuron as \(\hat{u}_{0}(\chi)=\chi\). Then, we can describe the projected gradient descent on the empirical loss \(\widehat{L}\) with initialization \(\{\chi_{1},\dots,\chi_{m}\}\) by:

\[\frac{d\hat{u}_{t}(\chi)}{dt} =-\text{grad}_{\hat{u}_{t}(\chi)}\widehat{L}(\hat{\rho}_{t})\,, \tag{3.8}\] \[\text{and }\hat{\rho}_{t} =\text{distribution of }\hat{u}_{t}(\chi)\text{ (where }\chi\sim\hat{\rho}_{0})\,. \tag{3.9}\]

We first state our result on the population, infinite-width dynamics.

**Theorem 3.3** (Population, infinite-width dynamics).: _Suppose Assumption 3.2 holds, and let \(\epsilon\in(0,1)\) be the target error. Let \(\rho_{t}\) be the result of projected gradient flow on the population loss, initialized with the uniform distribution on \(\mathbb{S}^{d-1}\), as defined in Eq. (3.6). Let \(T_{*,\epsilon}=\inf\{t>0\mid L(\rho_{t})\leq\frac{1}{2}(\hat{\sigma}_{2,d}^{2 }+\hat{\sigma}_{4,d}^{2})\epsilon^{2}\}\) be the earliest time \(t\) such that a loss of at most \(\frac{1}{2}(\hat{\sigma}_{2,d}^{2}+\hat{\sigma}_{4,d}^{2})\epsilon^{2}\) is reached. Then, we have_

\[T_{*,\epsilon}\lesssim\frac{1}{\hat{\sigma}_{2,d}^{2}\gamma_{2}}\log d+\frac{( \log\log d)^{20}}{\hat{\sigma}_{2,d}^{2}\epsilon\gamma_{2}^{8}}\log\Big{(} \frac{\gamma_{2}}{\epsilon}\Big{)}\,. \tag{3.10}\]

The first term on the right-hand side of Eq. (3.10) corresponds to the burn-in time for the network to reach a region around where the Polyak-Lojasiewicz condition holds. We divide our analysis of this burn-in phase into two phases, Phase 1 and Phase 2, and we obtain tight control on the factor by which the signal component, \(q_{*}^{\top}u\), grows during Phase 1, while Phase 2 takes place for a comparatively short period of time. (This tight control is critical for our sample complexity bounds where we must show that the coupling error
does not blow up too much -- see more discussion below Lemma 5.4.) Phases 1 and 2 are mostly governed by the quadratic components in the activation and target functions, and the dynamics behave similarly to a power method update. After the burn-in phase, the dynamics operate for a short period of time (Phase 3) in a regime where the Polyak-Lojasiewicz condition holds. We explicitly prove the dynamics stay away from saddle points during this phase, as further discussed in Section 4.2.

We note that Theorem 3.3 provides a concrete polynomial runtime bound for projected gradient flow which is not achievable by prior mean-field analyses (Chizat and Bach, 2018a; Mei et al., 2018) using Wasserstein gradient flow techniques. The main challenge in the proof is to deal with the saddle points that are not strict-saddle (Ge et al., 2015) in the loss landscape which cannot be escaped simply by adding noise in the parameter space (Jin et al., 2021; Lee et al., 2016; Daneshmand et al., 2018).1 Our analysis develops a fine-grained analysis of the dynamics that shows the iterates stay away from saddle points, which allows us to obtain the running time bound in Theorem 3.3 which can be translated to a polynomial-width guarantee in Theorem 3.4. In contrast, Wei et al. (2019) escape the saddles by randomly replacing an exponentially small fraction of neurons, which makes the network require exponential width.

Footnote 1: There are a long list of prior works on the loss landscape of neural networks (Mei et al., 2018; Soltanolkotabi, 2017; Vardi et al., 2021; Kakade et al., 2011; Bietti et al., 2022; Ben Arous et al., 2020; Arous et al., 2021; Soudry and Carmon, 2016; Tian, 2017; Ge et al., 2017; Zhang et al., 2019; Brutzkus and Globerson, 2017).

Next, we state the main theorem on projected gradient flow on empirical, finite-width dynamics.

**Theorem 3.4** (Empirical, finite-width dynamics).: _Suppose Assumption 3.2 holds. Suppose \(\epsilon=\frac{1}{\log\log d}\) is the target error and \(T_{*,\epsilon}\) is the running time defined in Theorem 3.3. Suppose \(n\geq d^{\mu}(\log d)^{O(1)}\) for any constant \(\mu>3\), and the network width \(m\) is equal to \(d^{C}\) for some sufficiently large universal constant \(C\). Let \(\hat{\rho}_{t}\) be the projected gradient flow on the empirical loss, initialized with \(m\) uniformly sampled weights, defined in Eq. (3.9). Then, with probability at least \(1-\frac{1}{d^{2(\log d)}}\) over the randomness of the data and the initialization of the finite-width network, we have that_

\[L(\hat{\rho}_{T_{*,\epsilon}})\lesssim(\hat{\sigma}_{2,d}^{2}+\hat{\sigma}_{4,d}^{2})\epsilon^{2}\,. \tag{3.11}\]

Plugging in \(\mu=3.1\), Eq. (3.11) suggests that, when the network width is at least \(d^{O(1)}\), the empirical dynamics of gradient descent could achieve \(\hat{\sigma}_{\max}^{2}(1/\log\log d)^{2}\) population loss with \(O(d^{3.1})\) samples.

In our analysis, we will establish a coupling between neurons in the empirical dynamics and neurons in the population dynamics. The main challenge is to bound the coupling error during Phase 1 where the population dynamics are similar to the power method. During this phase, we show that the coupling error (i.e., the distance between coupled neurons) remains small by showing that it grows at most as fast as the growth of the signal \(q_{\star}^{T}u\) in the population dynamics. Such a delicate relationship between the growth of the error and that of the signal is essential to proving our sample complexity bound; even an additional constant factor in the growth rate of the coupling error would lead to an additional \(\operatorname{poly}(d)\) factor in the sample complexity. Prior work (Mei et al., 2019) also establishes a coupling between the population and empirical dynamics. However, without the comparison with the growth of the signal, they only show a generic exponential growth rate for the coupling error. As a result, their technique cannot achieve a better sample complexity than the NTK.

The following theorem states that in the setting of Theorem 3.4, kernel methods with any inner product kernel require \(\Omega(d^{4}(\ln d)^{-6})\) samples to achieve a non-trivial population loss. (Note that the zero function has loss \(\mathbb{E}_{x\sim\mathbb{S}^{d-1}}[y(x)^{2}]\geq(\hat{h}_{4,d})^{2}\).)

**Theorem 3.5** (Sample complexity lower bound for kernel methods).: _Let \(K\) be an inner product kernel. Suppose \(d\) is larger than a universal constant and \(n\lesssim d^{4}(\ln d)^{-6}\). Then, with probability at least \(1/2\) over the randomness of \(n\) i.i.d data points \(\{x_{i}\}_{i=1}^{n}\) drawn from \(\mathbb{S}^{d-1}\), any estimator of \(y\) of the form \(f(x)=\sum_{i=1}^{n}\beta_{i}K(x_{i},x)\) of \(\{x_{i}\}_{i=1}^{n}\) must have a large error:_

\[\mathbb{E}_{x\sim\mathbb{S}^{d-1}}(y(x)-f(x))^{2}\geq\tfrac{3}{4}(\hat{h}_{4, d})^{2}. \tag{3.12}\]
Theorem 3.4 and Theorem 3.5 together prove a clear sample complexity separation between gradient flow and NTK. When \(\hat{h}_{4,d}\asymp\hat{\sigma}_{\max}\) (i.e., \(\gamma_{2},\gamma_{4}\asymp 1\)), gradient flow with finite-width neural networks can achieve \((\hat{h}_{4,d})^{2}(\log\log d)^{-2}\) population error with \(d^{3.1}\) samples, while kernel methods with any inner product kernel (including NTK) must have an error at least \((\hat{h}_{4,d})^{2}/2\) with \(d^{3.9}\) samples.

On a high level, Abbe et al. (2022, 2023) prove similar lower bounds in a different setting where the target function is drawn randomly and the data points can be arbitrary (also see Kamath et al. (2020); Hsu et al. (2021); Hsu (2021)). In comparison, our lower bound works for a fixed target function by exploiting the randomness of the data points. In fact, we can strengthen Theorem 3.5 by proving a \(\Omega(\hat{h}_{k,d}^{2})\) loss lower bound for any universal constant \(k\geq 0\) (Theorem E.2). We defer the proof of Theorem 3.5 to Appendix E.

## 4 Analysis of Population Dynamics

In this section, we give an overview of the analysis for the population infinite-width dynamics \(\rho_{t}\) (Theorem 3.3). The full proof is given in Appendix B.

### Symmetry of the Population Dynamics

A key observation is that due to the symmetry in the data, the population dynamics \(\rho_{t}\) has a symmetric distribution in the subspace orthogonal to the vector \(q_{\star}\). As a result, the dynamics \(\rho_{t}\) can be precisely characterized by the dynamics in the direction of \(q_{\star}\).

Recall that \(w=q_{\star}^{\top}u\) denotes the projection of a weight vector \(u\) in the direction of \(q_{\star}\). Let \(z=(I-q_{\star}q_{\star}^{\top})u\) be the remaining component. For notational convenience, without loss of generality, we can assume that \(q_{\star}=e_{1}\) and write \(u=(w,z)\), where \(w\in[-1,1]\) and \(z\in\sqrt{1-w^{2}}\cdot\mathbb{S}^{d-2}\). _We will use this convention throughout the rest of the paper_. We will use \(w_{t}(\chi)\) to refer to the first coordinate of the particle \(u_{t}(\chi)\), and \(z_{t}(\chi)\) to refer to the last \((d-1)\) coordinates.

**Definition 4.1** (Rotational invariance and symmetry).: _We say a neural network \(\rho\) is rotationally invariant if for \(u=(w,z)\sim\rho\), the distribution of \(z\mid w\) is uniform over \(\sqrt{1-w^{2}}\mathbb{S}^{d-2}\) almost surely. We also say the network is symmetric w.r.t a variable \(w\) if the density of \(w\) is an even function._

We note that any polynomial-width neural network (e.g., \(\hat{\rho}\)) is very far from rotationally invariant, and therefore the definition is specifically used for population dynamics.

If \(\rho\) is rotationally invariant and symmetric, then \(L(\rho)\) has a simpler form that only depends on the marginal distribution of \(w\).

**Lemma 4.2**.: _Let \(\rho\) be a rotationally invariant neural network. Then, for any target function \(h\) and any activation \(\sigma\),_

\[L(\rho)=\tfrac{1}{2}\sum_{k=0}^{\infty}\Big{(}\hat{\sigma}_{k,d} \operatorname{\mathbb{E}}_{u\sim\rho}[P_{k,d}(w)]-\hat{h}_{k,d}\Big{)}^{2}\,. \tag{4.1}\]

_In addition, suppose \(h\) and \(\sigma\) satisfy Assumption 3.2 and \(\rho\) is symmetric. Then_

\[L(\rho)=\tfrac{\hat{\sigma}_{2,d}^{2}}{2}\Big{(}\operatorname{ \mathbb{E}}_{u\sim\rho}[P_{2,d}(w)]-\gamma_{2}\Big{)}^{2}+\tfrac{\hat{\sigma} _{4,d}^{2}}{2}\Big{(}\operatorname{\mathbb{E}}_{u\sim\rho}[P_{4,d}(w)]-\gamma _{4}\Big{)}^{2}\,. \tag{4.2}\]

The proof of Lemma 4.2 is deferred to Appendix B.1. Eq. (4.1) says that if \(\rho\) is rotationally invariant, then \(L(\rho)\) only depends on the marginal distribution of \(w\). Eq. (4.2) says that if the distribution of \(w\) is additionally symmetric and \(\sigma\) and \(h\) are quartic, then the terms corresponding to odd \(k\) and the higher order terms for \(k>4\) in Eq. (4.1) vanish. Note that \(P_{2,d}(s)\approx s^{2}\) and \(P_{4,d}(s)\approx s^{4}\) by Eq. (A.4). Thus, \(L(\rho)\) essentially corresponds to matching the second and fourth moments of \(w\) to some desired values \(\gamma_{2}\) and \(\gamma_{4}\). Inspired by this lemma, we define the following key quantities: for any time \(t\geq 0\), we define \(D_{2,t}=\operatorname{\mathbb{E}}_{u\sim\rho_{t}}[P_{2,d}(w)]-\gamma_{2}\) and \(D_{4,t}=\operatorname{\mathbb{E}}_{u\sim\rho_{t}}[P_{4,d}(w)]-\gamma_{4}\), where \(\rho_{t}\) is defined according to the population, infinite-width dynamics (Eq. (3.5), Eq. (3.6) and Eq. (3.7)).

We next show that the rotational invariance and symmetry properties of \(\rho_{t}\) are indeed maintained:
**Lemma 4.3**.: _Suppose we are in the setting of Theorem 3.3. At any time \(t\in[0,\infty)\), \(\rho_{t}\) is symmetric and rotationally invariant._

The proof of Lemma 4.3 is deferred to Appendix B.2. To show rotational invariance, we use Eq. (3.4) and the rotational invariance of the data in the last \((d-1)\) coordinates. To show symmetry, we use Eq. (4.1), and the facts that \(P_{k,d}\) is an odd polynomial for odd \(k\) and \(\rho_{t}\) is symmetric at initialization.

As a consequence of Lemma 4.2 and Lemma 4.3, we obtain a simple formula for the dynamics of \(w_{t}\):

**Lemma 4.4** (\(1\)-dimensional dynamics).: _Suppose we are in the setting of Theorem 3.3. Then, for any \(\chi\in\mathbb{S}^{d-1}\), writing \(w_{t}:=w_{t}(\chi)\), we have_

\[\frac{dw_{t}}{dt}=\underbrace{-(1-w_{t}^{2})\cdot(P_{t}(w_{t})+Q_{t}(w_{t}))}_ {\triangleq_{v}(w_{t})}, \tag{4.3}\]

_where for any \(w\in[-1,1]\), we have \(P_{t}(w)=2\hat{\sigma}_{2,d}^{2}D_{2,t}w+4\hat{\sigma}_{4,d}^{2}D_{4,t}w^{3}\), and \(Q_{t}(w)={\lambda_{d}}^{(1)}w+{\lambda_{d}}^{(3)}w^{3}\), where \(|{\lambda_{d}}^{(1)}|,|{\lambda_{d}}^{(3)}|\lesssim\frac{\hat{\sigma}_{2,d}^{2 }|D_{2,t}|+\hat{\sigma}_{4,d}^{2}|D_{4,t}|}{d}\). More specifically, \({\lambda_{d}}^{(1)}=2\hat{\sigma}_{2,d}^{2}D_{2,t}\cdot\frac{1}{d-1}-2\hat{ \sigma}_{4,d}^{2}D_{4,t}\cdot\frac{6d+12}{d^{2}-1}\) and \({\lambda_{d}}^{(3)}=4\hat{\sigma}_{4,d}^{2}D_{4,t}\cdot\frac{6d+9}{d^{2}-1}\)._

Eq. (4.3) is a properly defined dynamics for \(w\) because the update rule for \(w_{t}\) only depends on \(w_{t}\) and the quantities \(D_{2,t}\) and \(D_{4,t}\) -- additionally, \(D_{2,t}\) and \(D_{4,t}\) only depend on the distribution of \(w\). As a slight abuse of notation, in the rest of this section, and in Appendix B, we will refer to the \(w_{t}(\chi)\) as particles -- this is well-defined by Lemma 4.4. This does not lead to ambiguity because we no longer need to consider the \(u_{t}(\chi)\) when analyzing the population dynamics. We also use \(\iota\in[-1,1]\) to refer to the first coordinate of \(\chi\), the initialization of a particle under the population dynamics. We note that \(\iota=\langle\chi,e_{1}\rangle\) and therefore, the distribution of \(\iota\) is \(\mu_{d}\). For any time \(t\geq 0\), we use \(w_{t}(\iota)\) to refer to \(w_{t}(\chi)\) for any \(\chi\in\mathbb{S}^{d-1}\). This notation is also well-defined by Lemma 4.4.

### Analysis of One-Dimensional Population Dynamics

We divide our proof of Theorem 3.3 into three phases, which are defined as follows. The first phase corresponds to the period of time under which for most initializations \(\iota\), \(w_{t}(\iota)\) is still small. For ease of presentation, we will only consider particles \(w_{t}(\iota)\) for \(\iota>0\) in the following discussion. Our argument also applies for \(\iota<0\) by the symmetry of \(\rho_{t}\) (Lemma 4.3).

**Definition 4.5** (Phase 1).: _Let \(w_{\max}=\frac{1}{\log d}\) and \(\iota_{\mathrm{U}}=\frac{\log d}{\sqrt{d}}\). Let \(T_{1}>0\) be the minimum time such that \(w_{T_{1}}(\iota_{\mathrm{U}})=w_{\max}:=\frac{1}{\log d}\). We refer to the time interval \([0,T_{1}]\) as Phase 1._

By tail bounds for \(\mu_{d}\), at initialization, essentially all of the particles are less than \(\iota_{\mathrm{U}}=\frac{\log d}{\sqrt{d}}\) in magnitude. Since the magnitudes of the \(w_{t}(\iota)\) maintain their respective orders under the population dynamics (Proposition B.33), Phase 1 corresponds to the duration under which all but a negligible portion of the particles have \(|w_{t}(\iota)|\leq w_{\max}\). During Phase 1, we have \(D_{2,t},D_{4,t}<0\), and the term corresponding to \(D_{2,t}\) in the velocity (Eq. (4.3)) dominates, so all particles \(w\) grow by essentially the same \(\frac{\sqrt{d}}{(\log d)^{2}}\) factor. Note that the particles do not have a large velocity in absolute terms, and thus the loss does not decrease much during this phase.

**Definition 4.6** (Phase 2).: _Let \(T_{2}>0\) be the minimum time such that either \(D_{2,T_{2}}=0\) or \(D_{4,T_{2}}=0\). We refer to the time interval \([T_{1},T_{2}]\) as Phase 2. Note that \(T_{2}>T_{1}\) by Lemma B.7._

In other words, Phase 2 corresponds to the time period after Phase 1, during which \(D_{2,t}\) and \(D_{4,t}\) are still negative. Afterwards, at least one of \(D_{2,t}\) or \(D_{4,t}\) becomes nonnegative. During Phase 2, a large fraction of particles grows to \(\mathrm{poly}(\frac{1}{\log\log d})\). Henceforth, for the rest of Phase 2, all of these particles will have a large velocity, and the loss will decrease quickly for the rest of Phase 2.

**Definition 4.7** (Phase 3).: _Phase 3 is defined as the time interval \([T_{2},\infty)\)._
This phase presents two cases: (i) \(D_{2,T_{2}}=0\) and \(D_{4,T_{2}}<0\), or (ii) \(D_{2,T_{2}}<0\) and \(D_{4,T_{2}}=0\). The analyses of the two cases are in Appendix B.5 and Appendix B.6 respectively. We give an overview of both cases below. Our main goal in Phase 3 is to show that if \((D_{2,t},D_{4,t})\) is far from \((0,0)\), then the velocity is large enough to allow a significant decrease in the loss function. Recall that the velocity \(v(w)=-(1-w^{2})(P_{t}(w)+Q_{t}(w))\) defined in Lemma 4.4 is approximately given by

\[\frac{dw}{dt}\approx-(1-w^{2})P_{t}(w)=-(1-w^{2})\left(2\hat{\sigma}_{2,d}^{2} D_{2,t}w+4\hat{\sigma}_{4,d}^{2}D_{4,t}w^{3}\right)\,. \tag{4.4}\]

**Phase 3, First Case** If \(D_{2,T_{2}}=0\) and \(D_{4,T_{2}}<0\), then for all \(t\geq T_{2}\), we will have \(D_{2,t}\geq 0\) and \(D_{4,t}\leq 0\). (See Lemma B.17 and Lemma B.18.) Lemma B.15 implies that when \(D_{2,t}\geq 0\) and \(D_{4,t}\leq 0\), a \(\text{poly}(\gamma_{2})\) fraction of particles will be far from \(0\) and \(1\), which are two of the roots of the velocity function. However, this does not guarantee a large average velocity -- unlike in Phases 1 and 2, the velocity may have a positive root \(r\in(0,1)\), which can give rise to bad stationary points. As a concrete example, if Assumption 3.2 holds, there exists some \(r\) with \(\frac{1}{\log\log d}\lesssim r\leq 1\) such that if \(\rho\) is the singleton distribution which assigns all of its probability to \(r\), then the velocity under \(\rho\) at \(r\) is exactly \(0\) (Lemma B.31). This is because the velocity of \(r\) under \(\rho\) has a factor \(\hat{\sigma}_{2,d}^{2}(P_{2,d}(r)-\gamma_{2})P_{2,d}{}^{\prime}(r)+\hat{ \sigma}_{4,d}^{2}(P_{4,d}(r)-\gamma_{4})P_{4,d}{}^{\prime}(r)\), and this sixth-degree polynomial in \(r\) has a root between \(\frac{1}{\log\log d}\) and \(1\), by the definitions of \(P_{2,d}\) and \(P_{4,d}\) (Eq. (A.4)).

Thus, we must leverage some information about the trajectory to show that the average velocity is large when \(D_{2,t}\) and \(D_{4,t}\) have different signs. To show that the average velocity is large, it suffices to show that the particles are not concentrated around the root, that is, that \(\mathbb{E}_{u\sim\rho_{t}}[(w-r)^{2}]\gtrsim\mathbb{E}_{w,w^{\prime}\sim\rho_ {t}}(w-w^{\prime})^{2}\) is large. However, it is quite possible that for two particles \(w\) and \(w^{\prime}\), the distance \(|w-w^{\prime}|\) decreases over time. Our main technique to control the distance between pairs of particles is as follows. We define a potential function \(\Phi(w):=\log(\frac{w}{\sqrt{1-w^{2}}})\), and we show that \(|\Phi(w)-\Phi(w^{\prime})|\) is always increasing for any two particles \(w,w^{\prime}\) with the same sign, during Phases 1 and 2, as well as during Phase 3 if \(D_{2,T_{2}}=0\) and \(D_{4,T_{2}}<0\) (Lemma B.11). Using the fact that there is a large portion of particles away from \(0\) and \(1\) as long as \(D_{2,t}>0\) and \(D_{4,t}<0\), and because \(\Phi\) is Lispchitz on an interval bounded away from \(0\) and \(1\), we show that the particles which are away from \(0\) and \(1\) also have large variance, implying that the average velocity is large.

**Phase 3, Second Case** For the case \(D_{4,T_{2}}=0\) and \(D_{2,T_{2}}<0\), we use a somewhat different argument which also involves \(\Phi\). We again have to deal with the fact that the velocity \(v(w)\) may have a positive root \(r\). The key point is that, due to the constraints \(D_{2,t}\leq 0\) and \(D_{4,t}\geq 0\) which hold during Phase 3, Case 2 (Lemma B.25), it is not possible for all the particles \(w\) to be stuck at the root \(r\). This is because, by the definition of \(D_{2,t}\) and \(D_{4,t}\), we will approximately have \(\mathbb{E}[w^{2}]\leq\gamma_{2}\) and \(\mathbb{E}[w^{4}]\geq\gamma_{4}\), and if nearly all the particles \(w\) are stuck at \(r\), then this will imply that \(\gamma_{4}\approx\gamma_{2}^{2}\). This would lead to a contradiction since Assumption 3.2 states that \(\gamma_{4}\geq 1.1\gamma_{2}^{2}\).

However, it is still possible for a large portion of particles \(w\) to be stuck at \(0\) or \(1\), meaning that these particles could have a very small velocity due to the factors \(w\) and \((1-w^{2})\) in \(v(w)\). To deal with such particles, we use the following argument. First, by time \(T_{2}\), nearly all of the particles have grown to at least \(\text{poly}(\frac{1}{\log\log d})\) in absolute value, as mentioned above in the discussion of Phase 2. Additionally, since we approximately have \(\mathbb{E}[w^{4}]\leq\gamma_{4}\) at time \(T_{2}\) since \(D_{4,T_{2}}=0\), we can argue by Markov's inequality that an \(O(\gamma_{4})\) fraction of particles is at most \(\frac{1}{2}\). For the purposes of this overview, we can refer to the set of particles which are larger than \(\text{poly}(\frac{1}{\log\log d})\) and smaller than \(O(\gamma_{4})\) as \(\mathcal{S}_{1}\). We also use \(w_{t}(\iota_{\text{M}})\) to refer to the largest particle in \(\mathcal{S}_{1}\).

Our key observation about \(\Phi\) in Phase 3, Case 2 is that for any two particles \(w,w^{\prime}\) with the same sign, \(|\Phi(w)-\Phi(w^{\prime})|\) is in fact _decreasing_ after the time \(T_{2}\) (Lemma B.20). Using this fact, along with our bound on the initial potential difference between any two particles in \(\mathcal{S}_{1}\) (Lemma B.21), we argue that if one of the particles in \(\mathcal{S}_{1}\) is extremely close to \(0\), then all of the particles in \(\mathcal{S}_{1}\) are extremely close to \(0\), and if one of the particles in \(\mathcal{S}_{1}\) is extremely close to \(1\), then all of the particles in \(\mathcal{S}_{1}\) is extremely close to \(1\) (see Lemma B.22 and its proof). By our definition of \(\mathcal{S}_{1}\), at most an \(O(\gamma_{4})\) fraction of particles are not in \(\mathcal{S}_{1}\), meaning that the constraints \(D_{2,t}\leq 0\) and \(D_{4,t}\geq 0\) are violated if either all particles in \(\mathcal{S}_{1}\) are close to \(0\) or all particles in \(\mathcal{S}_{1}\) are close to \(1\) (we make this argument precise in Lemma B.22). Thus, we can show that all the particles in \(\mathcal{S}_{2}\) are far from \(0\) and \(1\) during Phase 3, Case 2.
However, it may be the case that all of the particles in \(\mathcal{S}_{1}\) are stuck at \(r\), since \(\mathcal{S}_{1}\) contains at least a \(1-O(\gamma_{4})\) fraction of the particles, but using the constraints \(D_{2,t}\leq 0\) and \(D_{4,t}\geq 0\), the best we can show is that an \(\Omega(\gamma_{4}^{5/4})\) fraction of particles is away from \(r\) (see the proof of Lemma B.29). Thus, in principle, it is possible that \(\mathcal{S}_{1}\), and the set of particles \(w\) which are far from \(r\), are disjoint. We can circumvent this issue by finally considering the particles \(w\) which are larger than \(w_{t}(\iota_{\text{M}})\). In principle, these particles can all be stuck at or near 1. However, when the positive root \(r\) of \(v_{t}(w)\) is less than \(\frac{1}{2}\), then all the particles \(w\) which are close to 1 will be drawn towards \(r\), and will leave 1 at an exponentially fast rate (see the proof of Lemma B.28). Additionally, we show that the positive root \(r\) cannot be larger than \(\frac{1}{2}\) for too long, leveraging our observations about \(\mathcal{S}_{1}\) (see the proof of Lemma B.26). Thus, eventually, a \(1-\operatorname{poly}(\frac{1}{\log\log d})\) fraction of the particles will be away from 0 and 1. At this point, the average velocity of the particles in \(\rho_{t}\) will be large since an \(\Omega(\gamma_{4}^{5/4})\) fraction of the particles is away from \(r\), as mentioned above.

We now state our conclusions for each of the phases. Recall that \(T_{*,\epsilon}\) is the amount of time \(t\) needed until the loss \(L(\rho_{t})\) becomes less than \(\frac{1}{2}(\hat{\sigma}_{2,d}^{2}+\hat{\sigma}_{4,d}^{2})\epsilon^{2}\).

**Lemma 4.8** (Phase 1).: _Suppose we are in the setting of Theorem 3.3. At the end of phase 1, i.e. at time \(T_{1}\), for all \(\iota\in(0,\iota_{\text{U}})\), we have \(w_{T_{1}}(\iota)\asymp\frac{\sqrt{d}}{(\log d)^{2}}\iota\). Furthermore, \(T_{1}\lesssim\frac{1}{\hat{\sigma}_{2,d}^{2}\gamma^{2}}\log d\). Additionally, we have \(\exp\Big{(}\int_{0}^{T_{1}}2\hat{\sigma}_{2,d}^{2}|D_{2,t}|dt\Big{)}\asymp\frac {\sqrt{d}}{(\log d)^{2}}\)._

We do not make use of the explicit runtime \(T_{1}\) of Phase 1, but we state it for reference. The quantity \(\exp\Big{(}\int_{0}^{T_{1}}2\hat{\sigma}_{2,d}^{2}|D_{2,t}|dt\Big{)}\) is relevant for the analysis of the empirical, finite-width dynamics, as discussed in Section 5. It is approximately the factor by which the \(w_{t}(\iota)\) grow during Phase 1.

**Lemma 4.9** (Phase 2).: _Suppose we are in the setting of Theorem 3.3. Let \(T_{2}\) be the minimum time such that either \(D_{2,T_{2}}=0\) or \(D_{4,T_{2}}=0\), i.e. \(T_{2}\) is the end of phase 2. Then, either \(T_{2}-T_{1}\lesssim\frac{(\log\log d)^{18}}{(\hat{\sigma}_{2,d}^{2}+\hat{ \sigma}_{4,d}^{2})}\log\Big{(}\frac{\gamma_{2}}{\epsilon}\Big{)}\) or \(T_{*,\epsilon}-T_{1}\lesssim\frac{(\log\log d)^{18}}{(\hat{\sigma}_{2,d}^{2}+ \hat{\sigma}_{4,d}^{2})}\log\Big{(}\frac{\gamma_{2}}{\epsilon}\Big{)}\)._

This lemma, which summarizes Phase 2, states that within \(\frac{(\log\log d)^{18}}{(\hat{\sigma}_{2,d}^{2}+\hat{\sigma}_{4,d}^{2})}\log \Big{(}\frac{\gamma_{2}}{\epsilon}\Big{)}\) time, either Phase 2 ends, or the loss is less than the desired value \((\frac{1}{2}(\hat{\sigma}_{2,d}^{2}+\hat{\sigma}_{4,d}^{2})\epsilon^{2}\).

**Lemma 4.10** (Phase 3, Case 1).: _In the setting of Theorem 3.3, suppose that \(D_{2,T_{2}}=0\) and \(D_{4,T_{2}}<0\) (i.e. Phase 3, Case 1 holds). Then, the total amount of time \(t\geq T_{2}\) such that \(L(\rho_{t})\geq\frac{1}{2}(\hat{\sigma}_{2,d}^{2}+\hat{\sigma}_{4,d}^{2}) \epsilon^{2}\) is at most \(O(\frac{\gamma_{2}}{\hat{\sigma}_{4,d}^{2}\gamma_{2}^{2}\xi^{4}\kappa^{6}} \log(\frac{\gamma_{2}}{\epsilon}))\)._

**Lemma 4.11** (Phase 3, Case 2).: _Suppose we are in the setting of Theorem 3.3. Assume \(D_{4,T_{2}}=0\) and \(D_{2,T_{2}}<0\), i.e. Phase 3 Case 2 holds. Then, at most \(O\Big{(}\frac{1}{\hat{\sigma}_{2,d}^{2}\sigma_{2}^{11/2}\xi^{6}\kappa^{12}} \log\Big{(}\frac{\gamma_{2}}{\epsilon}\Big{)}\Big{)}\) time \(t\geq T_{2}\) elapses while \(L(\rho_{t})\gtrsim(\hat{\sigma}_{2,d}^{2}+\hat{\sigma}_{4,d}^{2})\epsilon^{2}\)._

Combining these lemmas gives Theorem 3.3.

## 5 Analysis of Empirical, Finite-Width Dynamics

The first step is to set up a coupling between the finite-width empirical dynamics and the infinite-width population dynamics -- we will eventually show that these two dynamics do not diverge very far, even up to time \(T_{*,\epsilon}\) (defined in Theorem 3.3).

Coupling Between Empirical and Population Dynamics.To analyze the difference between the empirical dynamics (resulting in \(\hat{\rho}_{t}\)) and the population dynamics (resulting in \(\rho_{t}\)) we define an intermediate process \(\bar{\rho}_{t}\), which has finite support. Specifically, \(\bar{\rho}_{0}=\hat{\rho}_{0}\), and \(\bar{\rho}_{t}\) will always have support size \(m\), but the
particles in \(\bar{\rho}\) are then updated according to the infinite-width population dynamics. We formally define \(\bar{\rho}_{t}\) as follows:

\[\bar{\rho}_{0} =\hat{\rho}_{0}=\text{unif}(\{\chi_{1},\dots,\chi_{m}\})\,,\] \[\text{and }\bar{\rho}_{t} =\text{distribution of }u_{t}(\chi)\text{ (where }\chi\sim\bar{\rho}_{0})\,. \tag{5.1}\]

In other words, \(\bar{\rho}_{t}\) consists of the particles in \(\rho_{t}\) whose initialization is in \(\{\chi_{1},\dots,\chi_{m}\}\).

Now, we define a coupling between \(\bar{\rho}_{t}\) and \(\hat{\rho}_{t}\). For \(\chi\sim\text{unif}(\{\chi_{1},\dots,\chi_{m}\})\), let \(\bar{u}_{t}(\chi)=u_{t}(\chi)\). Let \(\Gamma_{t}\) be the joint distribution of \((\bar{u}_{t}(\chi),\hat{u}_{t}(\chi))\). Then, \(\Gamma_{t}\) forms a natural coupling between \(\bar{\rho}_{t}\) and \(\hat{\rho}_{t}\). We will use \(\bar{u}_{t}\) and \(\hat{u}_{t}\) as shorthands for the random variables \(\bar{u}_{t}(\chi),\hat{u}_{t}(\chi)\) respectively in the rest of this section. Additionally, we define the average distance \(\overline{\Delta}_{t}^{2}:=\mathbb{E}_{(\bar{u}_{t},\bar{u}_{t})\sim\Gamma_{t} }[\|\hat{u}_{t}-\bar{u}_{t}\|_{2}^{2}]\). Intuitively, \(f_{\rho_{t}}(x)\) and \(f_{\bar{\rho}_{t}}(x)\) are close when \(\hat{u}_{t}\) and \(\bar{u}_{t}\) are close, which is formalized by the following lemma:

**Lemma 5.1**.: _In the setting of Theorem 3.4, suppose the network width \(m\) satisfies \(m\leq d^{O(\log d)}\), and let \(T\leq\frac{d^{C}}{\sigma_{2,d}^{2}+\sigma_{4,d}^{2}}\) for some universal constant \(C>0\). Then, with probability at least \(1-\exp(-d^{2})\) over the randomness of \(\{\chi_{1},\dots,\chi_{m}\}\), we have for all \(t\leq T\) that_

\[\mathbb{E}_{x\sim\mathbb{S}^{d-1}}[(f_{\rho_{t}}(x)-f_{\bar{\rho}_{t}}(x))^{2} ]\lesssim\frac{(\hat{\sigma}_{2,d}^{2}+\hat{\sigma}_{4,d}^{2})d^{2}(\log d)^{ O(1)}}{m}\,. \tag{5.2}\]

_Additionally, for all \(t\geq 0\), we have_

\[\mathbb{E}_{x\sim\mathbb{S}^{d-1}}[(f_{\hat{\rho}_{t}}(x)-f_{\hat{\rho}_{t}}( x))^{2}]\lesssim(\hat{\sigma}_{2,d}^{2}+\hat{\sigma}_{4,d}^{2})\overline{ \Delta}_{t}^{2}\,. \tag{5.3}\]

The proof of Lemma 5.1 is deferred to Appendix C. As a simple corollary of Lemma 5.1, we can upper bound \(\mathbb{E}_{x\sim\mathbb{S}^{d-1}}[(f_{\rho_{t}}(x)-f_{\hat{\rho}_{t}}(x))^{2}]\) by the triangle inequality. Thus, so long as \(\overline{\Delta}_{T_{*,\epsilon}}\) is small, we can show that the empirical and population dynamics achieve similar test error at time \(T_{*,\epsilon}\).

Upper Bound for \(\overline{\Delta}_{T_{*,\epsilon}}\).The following lemma gives our upper bound on \(\overline{\Delta}_{T_{*,\epsilon}}\):

**Lemma 5.2** (Final Bound on \(\overline{\Delta}_{T_{*,\epsilon}}\)).: _In the setting of Theorem 3.4, we have_

\[\overline{\Delta}_{T_{*,\epsilon}}\leq d^{-\frac{\mu-1}{4}}\,. \tag{5.4}\]

We prove Lemma 5.2 by induction over the time \(t\leq T_{*,\epsilon}\). Let us define the maximal distance \(\Delta_{\max,t}:=\max_{(\hat{u}_{t},\bar{u}_{t})\in\text{supp}(\Gamma_{t})}\| \hat{u}_{t}-\bar{u}_{t}\|_{2}\) where the max is over the _finite_ support of the coupling \(\Gamma_{t}\). We will control \(\overline{\Delta}_{t}\) and \(\Delta_{\max,t}\) simultaneously using induction. The inductive hypothesis is:

**Assumption 5.3** (Inductive Hypothesis).: _For some universal constant \(1>\phi>1/2\) and \(\psi\in(0,\phi-\frac{1}{2})\), we say that the inductive hypothesis holds at time \(T\) if, for all \(0\leq t\leq T\), we have_

\[\overline{\Delta}_{t}\leq d^{-\phi}\quad\text{and}\quad\Delta_{\max,t}\leq d^{ -\psi}\,. \tag{5.5}\]

Showing that the inductive hypothesis holds for all \(t\) requires studying the growth of the error \(\|\hat{u}_{t}-\bar{u}_{t}\|\). We analyze it using the following decomposition: 2

Footnote 2: Note that \(\hat{u}_{t}-\bar{u}_{t}\) and \(A_{t},B_{t}\) and \(C_{t}\) implicitly depend on the initialization \(\chi\), and could be written as \(A_{t}(\chi),B_{t}(\chi),C_{t}(\chi),\) and \(\delta_{t}(\chi)\), but we omit this dependence for ease of presentation.

\[\frac{d}{dt}\|\hat{u}_{t}-\bar{u}_{t}\|_{2}^{2}= A_{t}+B_{t}+C_{t}\,, \tag{5.6}\]

where \(A_{t}:=-2\langle\text{grad}_{\hat{u}}L(\rho_{t})-\text{grad}_{\bar{u}}L(\rho_{t} ),\hat{u}_{t}-\bar{u}_{t}\rangle\), \(B_{t}:=-2\langle\text{grad}_{\hat{u}}L(\hat{\rho}_{t})-\text{grad}_{\hat{u}}L( \rho_{t}),\hat{u}_{t}-\bar{u}_{t}\rangle\), and \(C_{t}=-2\langle\text{grad}_{\hat{u}}\widehat{L}(\hat{\rho}_{t})-\text{grad}_{ \hat{u}}L(\hat{\rho}_{t}),\hat{u}_{t}-\bar{u}_{t}\rangle\). Intuitively, \(A_{t}\) captures the growth of the coupling error due to the population dynamics, \(B_{t}\) captures the growth of the coupling error due to the discrepancy between the gradients from the finite-width and infinite-width networks, and \(C_{t}\) captures the growth of the coupling error due to the discrepancy between finite samples and infinite samples.

We will use the following lemma to give an upper bound on each of these three terms.
**Lemma 5.4** (Bounds on \(A_{t},B_{t},C_{t}\)).: _In the setting of Theorem 3.4, suppose the inductive hypothesis (Assumption 5.3) holds up to time \(t\), for some \(t\leq T_{*,\epsilon}\). Let \(\delta_{t}:=\hat{u}_{t}-\bar{u}_{t}\). Let \(T_{1}\) be the runtime of Phase 1 (where Phase 1 is defined in Definition 4.5). Then, we have_

\[A_{t}\leq\begin{cases}4\hat{\sigma}_{2,d}^{2}|D_{2,t}|\left\|\delta_{t}\right\| ^{2}+O\big{(}\tfrac{\delta_{2,d}^{2}}{\log d}|D_{2,t}|\left\|\delta_{t}\right\| ^{2}\big{)}&\text{if }0\leq t\leq T_{1}\\ O(1)\cdot(\hat{\sigma}_{2,d}^{2}+\hat{\sigma}_{4,d}^{2})\gamma_{2}\left\| \delta_{t}\right\|^{2}&\text{if }T_{1}\leq t\leq T_{*,\epsilon}\end{cases}. \tag{5.7}\]

_Additionally, assuming that the width \(m\) is at most \(d^{O(\log d)}\), we have with probability at least \(1-e^{-d^{2}}\) over the randomness of the initialization \(\hat{\rho}_{0}\) that \(B_{t}\lesssim(\hat{\sigma}_{2,d}^{2}+\hat{\sigma}_{4,d}^{2})\cdot\Big{(}\tfrac {d(\log d)^{O(1)}}{\sqrt{m}}+\overline{\Delta}_{t}\Big{)}\|\delta_{t}\|_{2}\) and \(\mathbb{E}[B_{t}]\lesssim\frac{(\hat{\sigma}_{2,d}^{2}+\hat{\sigma}_{4,d}^{2})d (\log d)^{O(1)}}{\sqrt{m}}\overline{\Delta}_{t}+(\hat{\sigma}_{2,d}^{2}+\hat{ \sigma}_{4,d}^{2})\overline{\Delta}_{t}^{3}\)._

_Finally, assume that the number of samples \(n\leq d^{C}\) and the width \(m\) is equal to \(d^{C}\) for any universal constant \(C>0\). Assume that \(\overline{\Delta}_{t}\leq\frac{1}{\sqrt{d}}\) (i.e. \(\phi>\frac{1}{2}\)) and that \(\psi<\frac{1}{2}\). Then, with probability \(1-\frac{1}{d^{2(\log d)}}\) over the randomness of the dataset, we have \(C_{t}\lesssim(\hat{\sigma}_{2,d}^{2}+\hat{\sigma}_{4,d}^{2})(\log d)^{O(1)} \cdot\Big{(}\sqrt{\frac{d}{n}}\|\delta_{t}\|_{2}+\tfrac{d^{2-2\psi}}{n} \overline{\Delta}\|\delta_{t}\|_{2}^{2}+\tfrac{d^{2.5-2\psi}}{n}\overline{ \Delta}^{2}\|\delta_{t}\|_{2}+\tfrac{d^{2.5-2\psi}}{n}\overline{\Delta}^{2}\| \delta_{t}\|_{2}+\tfrac{d^{2.5-2\psi}}{n}\overline{\Delta}^{2}\|\delta_{t}\|_ {2}+\tfrac{d^{2.5-2\psi}}{n}\overline{\Delta}^{2}\|\delta_{t}\|_{2}+\tfrac{d^{ 2.5-2\psi}}{n}\overline{\Delta}^{2}\|\delta_{t}\|_{2}^{2}\Big{)}\)._

The proof can be found in Appendix C. We give an explanation for each of these bounds below.

For \(A_{t}\), we establish two separate bounds, one which holds during Phase 1, and one which holds during Phases 2 and 3. Intuitively, the signal part in a neuron (i.e. \(w^{2}\)) grows at a rate \(4\hat{\sigma}_{2,d}^{2}|D_{2,t}|\) during Phase 1 (which follows from Lemma 4.4) and in Eq. (5.7) we show that the growth rate of the coupling error due to the growth in the signal is also at most \(4\hat{\sigma}_{2,d}^{2}|D_{2,t}|\). Intuitively, the growth rates match because the signal parts of the neurons follow similar dynamics to a power method during Phase 1. For example, in the worst case when \(\delta_{t}\) is mostly in the direction of \(e_{1}\), the factor by which \(\delta_{t}\) grows is the same as that of \(w\). More mathematically, this is due to the fact that the first-order term in Lemma 4.4 is the dominant term in \(v(w)-v(\hat{w})\) for any two particles \(w,\hat{w}\) (where \(v(w)\) is the one-dimensional velocity defined in the previous section). To get a sample complexity better than NTK, the precise constant factor \(4\) in the growth rate \(4\hat{\sigma}_{2,d}^{2}|D_{2,t}|w\) is important -- a larger constant would lead to \(\delta_{T_{1}}\) being larger by a \(\operatorname{poly}(d)\) factor, thus increasing the sample complexity by \(\operatorname{poly}(d)\).

The same bound for \(A_{t}\) no longer applies during Phases 2 and 3. This is because the dynamics of \(w\) are no longer similar to a power method, and the higher-order terms may make a larger contribution to \(v(w)-v(\hat{w})\) than to the growth of \(w\), or vice versa. Thus we use a looser bound \(O(1)\cdot(\hat{\sigma}_{2,d}^{2}+\hat{\sigma}_{4,d}^{2})\gamma_{2}\left\|\delta _{t}\right\|^{2}\) in Eq. (5.7). However, since the remaining running time after Phase 1, \(T_{*,\epsilon}-T_{1}\), is very short (only \(\operatorname{poly}(\log\log d)\) -- this corresponds to the second term of the running time in Theorem 3.3), the total growth of the coupling error is at most \(\exp(\operatorname{poly}(\log\log d))\), which is sub-polynomial and does not contribute to the sample complexity. Note that if the running time during Phases 2 and 3 is longer (e.g. \(O(\log d)\)) then the coupling error would grow by an additional \(d^{O(1)}\) factor during Phases 2 and 3, which would make our sample complexity worse than NTK.

Our upper bounds on \(B_{t}\) and \(C_{t}\) capture the growth of coupling error due to the finite width and samples. We establish a stronger bound for \(\mathbb{E}[B_{t}]\) than for \(B_{t}\). In our proof, we use the bounds for \(B_{t}\) and \(\mathbb{E}[B_{t}]\) to prove the inductive hypothesis for \(\overline{\Delta}_{t}\) and \(\Delta_{\max,t}\) respectively. We prove the bound for \(C_{t}\) by expanding the error due to finite samples into second-order and fourth-order polynomials and applying concentration inequalities for higher moments (Lemma F.7).

All of these bounds together control the growth of \(\|\hat{u}_{t}-\bar{u}_{t}\|\) at time \(t\). Using Lemma 5.4 we can show that for some properly chosen \(\phi\) and \(\psi\), the inductive hypothesis in Assumption 5.3 holds until \(T_{*,\epsilon}\) (see Lemma C.3 for the statement), which naturally leads to the bound in Lemma 5.2. The full proof can be found in Appendix C.
Conclusion

In this paper, we prove a clear sample complexity separation between vanilla gradient flow and kernel methods with any inner product kernel, including NTK. Our work leads to several directions for future research. The first question is to generalize our results to the ReLU activation. Another question is whether gradient descent can achieve less than \(d^{3}\) sample complexity in our setting. A final open question is whether two-layer neural networks can be shown to attain arbitrarily small generalization error \(\epsilon\) -- one limitation of our analysis is that we require \(\epsilon\geq\text{poly}(1/\log\log d)\).

## Acknowledgments

The authors would like to thank Zhiyuan Li for helpful discussions. The authors would like to thank the support of NSF IIS 2045685, NSF Grant CCF-1844628 and a Sloan Research Fellowship.
