Title: Beyond NTK with Vanilla Gradient Descent: A Mean-Field Analysis of Neural Networks with Polynomial Width, Samples, and Time
Venue: NeurIPS 2023
Accepted: yes
Keywords: Nonconvex Optimization

Mean-Field Analysis

Beyond NTK

Deep Learning Theory
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 trained with projected gradient flow converges in polynomial 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.

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 trained with projected gradient flow converges in poly ( d ) time to a non-trivial error that is not achievable by kernel methods using n ≪ 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.

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 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 [29, 36, 51, 35, 65, 38, 48, 39, 60] 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. [30], Jacot et al. [43], Li and Liang [50], Arora et al. [10], Daniely et al. [26]). However, subsequent research shows that neural networks trained with practical hyperparameter settings typically outperform their corresponding NTK kernels [9]. Furthermore, the initialization and learning rate under the NTK regime does not yield optimal generalization guarantees [76, 22, 77, 34].

Many recent works study modified versions of stochastic gradient descent (SGD) and prove sample complexity and runtime guarantees beyond NTK [23, 2, 58, 52, 1, 7, 5, 79, 78, 25, 6, 67, 70, 20, 59]. 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. [52] 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. [23] use one step of gradient descent with a large learning rate and then optimize only the last layer of the neural net. Nichani et al. [59] construct non-standard regularizers based on the NTK feature covariance matrix, in order to prevent the movement of weights in certain directions which hinder generalization. Additionally, Abbe et al. [2] only train the hidden layer for O (1) time in the first stage of their algorithm, and then train the second layer to convergence in the second stage. They do not study how the population loss decreases during the first stage. 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 is the mean-field approach [21, 56, 57, 61, 28, 2, 44, 62, 76], which views the collection of weight vectors as a (discrete) distribution over 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 [21] and Mei et al. [56] do not provide a concrete bound on the number of iterations needed for convergence. Mei et al. [57] 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. [76] 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. We also discuss additional related works in Appendix A.

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. As a corollary, we show that projected gradient descent with polynomially small step size and polynomially many iterations can converge to a low error, with sample complexity that is superior to NTK. Our proof is similar to the standard bound on the error of Euler's method.

Concretely, the neural network is assumed to have unit-norm weight vectors and no bias, and the second-layer weights are all 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 ⊤ ⋆ 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 ≪ 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 wellstudied 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 [58, 2, 23, 16]. Many works on singleindex models study (stochastic) gradient descent in the setting where only a single vector (or 'neuron') in R d is trained. This includes earlier works where the link function is monotonic, or the convergence is analyzed with quasi-convexity [46, 40, 55, 64], along with more recent work [69, 72, 11] 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 ⋆ , so that even without bias terms, their mixture can represent the link function h . Our analysis demonstrates that gradient flow, and gradient descent with a consistent inverse-polynomial learning rate, can simultaneously learn the feature q ⋆ 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 [58, 1, 2, 23, 14].

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. [52], 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 [54] would result in a d O (1) sample complexity, which is not sufficient to outperform NTK.

2 Preliminaries and Notations

We use O ( · ) , ≲ , ≳ 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 ) | ≤ C | x | , ∀ x ∈ R for some universal constant C > 0 (each occurrence can have a different universal constant C and f ). We use a ≲ b as a shorthand for a ≤ O ( b ) . Similarly, Ω( x ) is a placeholder for some g ( x ) where | g ( x ) | ≥ | x | /C, ∀ x ∈ R for some universal constant C > 0 . We use a ≳ b as a shorthand for a ≥ Ω( b ) and a ≍ b as a shorthand to indicate that a ≳ b and a ≲ b simultaneously hold.

Legendre Polynomials. We summarize the necessary facts about Legendre polynomials below, and present related background more comprehensively in Appendix C. Let P k,d : [ -1 , 1] → R be the degreek un-normalized Legendre polynomial [12], and P k,d ( t ) = √ N k,d P k,d ( t ) be the normalized Legendre polynomial, where N k,d ≜ ( d + k -1 d -1 ) -( d + k -3 d -1 ) is the normalizing factor. The polynomials P k,d ( t ) form an orthonormal basis for the set of square-integrable functions over [ -1 , 1] with respect to the measure µ d ( t ) ≜ (1 -t 2 ) d -3 2 Γ( d/ 2) Γ(( d -1) / 2) 1 √ π , i.e., the density of u 1 when u = ( u 1 , · · · , u d ) is uniformly drawn from sphere S d -1 . Hence, for every function h : [ -1 , 1] → R such that E t ∼ µ d [ h ( t ) 2 ] < ∞ , we can define ˆ h k,d ≜ E t ∼ µ d [ h ( t ) P k,d ( t )] and consequently, we have h ( t ) = ∑ ∞ k =0 ˆ h k,d 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 ) : R d → R that we aim to learn has the form y ( x ) = h ( q ⊤ ⋆ x ) , where h : R → R is an unknown one-dimensional even quartic polynomial (which is called a link function), and q ⋆ is an unknown unit vector in R d . Note that h ( s ) has the Legendre expansion h ( s ) = ˆ h 0 ,d + ˆ h 2 ,d P 2 ,d ( s ) + ˆ h 4 ,d 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 σ ( · ) be the activation function, which can be different from h ( · ) . Using the mean-field formulation, we describe the neural network using the distribution of first-layer weight vectors, denoted by ρ :

For example, when ρ = unif ( { u 1 , . . . , u m } ) , is a discrete, uniform distribution supported on m vectors { u 1 , . . . , u m } , then f ρ ( x ) = 1 m ∑ m i =1 σ ( u ⊤ i x ) , i.e. f ρ corresponds to a finite-width neural network whose first-layer weights are u 1 , . . . , u m . For a continuous distribution ρ , the function f ρ ( · ) can be viewed as an infinite-width neural network where the weight vectors are distributed according to ρ (and can be viewed as taking the limit as m →∞ of the finite-width neural network). We assume that the weight vectors have unit norms, i.e. the support of ρ is contained in S d -1 . The activation σ : R → R is assumed to be a fourth-degree polynomial with Legendre expansion σ ( s ) = ∑ 4 k =0 ˆ σ k,d P k,d ( s ) .

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

Lemma 3.1 (Expressivity) . Let γ 2 = ˆ h 2 ,d / ˆ σ 2 ,d and γ 4 = ˆ h 4 ,d / ˆ σ 4 ,d . Suppose for some d and q ⋆ , there exists a network ρ such that f ρ ( x ) = h ( q ⊤ ⋆ x ) on S d -1 . Then we have ˆ σ 0 ,d = ˆ h 0 ,d , and 0 ≤ γ 2 2 ≤ γ 4 ≤ γ 2 ≤ 1 . Moreover, if this condition holds with strict inequalities, then for sufficiently large d , there exists a network ρ such that f ρ ( x ) = h ( q ⊤ ⋆ x ) on S d -1 . (A more explicit version is stated in Appendix G.1.)

Informally, an almost sufficient and necessary condition to have f ρ ( x ) = h ( q ⊤ ∗ x ) for some ρ is that there exists a random variable w supported on [0 , 1] such that E [ w 2 ] ≈ γ 2 and E [ w 4 ] ≈ γ 4 , which is equivalent to 0 ≤ γ 2 2 ≤ γ 4 ≤ γ 2 ≤ 1 . In particular, assuming the existence of such a random variable w , the perfectly-fit network ρ that fits the target function has the form

Motivated by this lemma, we will assume that γ 2 , γ 4 are universal constants that satisfy 0 ≤ γ 2 2 ≤ γ 4 ≤ γ 2 ≤ 1 , and d is chosen to be sufficiently large (depending on the choice of γ 2 and γ 4 ). In addition, we also assume that γ 4 ≤ O ( γ 2 2 ) , that is, the equality γ 2 2 ≤ γ 4 is somewhat tight this ensures that the distribution of w = q ⊤ ⋆ u under the perfectly-fitted neural network is not too spread-out. We also assume that γ 2 is smaller than a sufficiently small universal constant. This ensures that the distribution of q ⊤ ⋆ 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 ⋆ . 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 ˆ σ 0 ,d = ˆ 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. If ˆ σ 0 ,d = ˆ h 0 ,d , then we can assume without loss of generality that ˆ σ 0 ,d = ˆ h 0 ,d = 0 , since ˆ σ 0 ,d and ˆ h 0 ,d will cancel with each other in the population and empirical mean-squared losses (defined below). In summary, we make the following formal assumptions on the Legendre coefficients of the link function and the activation function.

Assumption 3.2. Let γ 2 = ˆ h 2 ,d / ˆ σ 2 ,d and γ 4 = ˆ h 4 ,d / ˆ σ 4 ,d . We first assume γ 4 ≥ 1 . 1 · γ 2 2 . For any universal constant c 1 > 1 , we assume that ˆ σ 2 2 ,d /c 1 ≤ ˆ σ 2 4 ,d ≤ c 1 · ˆ σ 2 2 ,d , and γ 4 ≤ c 1 γ 2 2 . For a sufficiently small universal constant c 2 > 0 (which is chosen after c 1 is determined), we assume 0 ≤ γ 2 ≤ 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 ˆ h 0 ,d = ˆ σ 0 ,d = 0 , and ˆ h 1 ,d = ˆ h 3 ,d = 0 .

Our intention is to replicate the ReLU activation as well as possible with quartic polynomials; our assumption that ˆ σ 2 2 ,d ≍ ˆ σ 2 4 ,d is indeed satisfied by the quartic expansion of ReLU because ̂ relu 2 ,d ≍ d -1 / 2 and ̂ relu 2 ,d ≍ d -1 / 2 (see Proposition C.3). Following the convention defined in Section 2, we will simply write ˆ σ 2 4 ,d ≍ ˆ σ 2 2 ,d , γ 4 ≥ 1 . 1 · γ 2 2 , and γ 4 ≲ γ 2 2 .

̸

Our assumptions rule out the case that γ 4 = γ 2 2 , for the sake of simplicity. We believe that our analysis could be extended to this case with some modifications. Our analysis also rules out the case where γ 2 = 0 and γ 4 = 0 , due to the restriction that γ 4 ≤ c 1 γ 2 2 . The case γ 2 = 0 and γ 4 = 0 would have a significantly different analysis, and potentially a different sample complexity, since our analysis in Section 4.2 and Section 5 makes use of the fact that the initial phase of the population and empirical dynamics behaves similarly to a power method, which follows from γ 2 being nonzero.

̸

Data distribution, losses, and projected gradient flow. The population data distribution is assumed to be S d -1 . We draw n training examples x 1 , . . . , x n i.i.d ∼ S d -1 . Thus, the population and empirical mean-squared losses are:

To ensure that the weight vectors remain on 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 ρ and

the corresponding Riemannian gradient (which is simply the projection of the gradient to the tangent space of S d -1 ):

Here we interpret L ( ρ ) as a function of a collection of particles (denoted by ρ ) and ∇ u L ( ρ ) as the partial derivative with respect to a single particle u evaluated at ρ . Similarly, the (Riemannian) gradient of the empirical loss ̂ L with respect to the particle u is defined as

Population, Infinite-Width Dynamics. Let the initial distribution ρ 0 of the infinite-width neural network be the uniform distribution over S d -1 . We use χ to denote a particle sampled uniformly at random from the initial distribution ρ 0 . A particle initialized at χ follows a deterministic trajectory afterwards - we use u t ( χ ) to denote the location, at time t , of the particle that was initialized at χ . Because u t ( · ) is a deterministic function, we can use χ 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

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 widthm neural network. Let χ 1 , . . . , χ m i.i.d ∼ S d -1 be the initial weight vectors of the widthm neural network, and let ˆ ρ 0 = unif ( { χ 1 , . . . , χ m } ) be the uniform distribution over these initial neurons. We use χ ∈ { χ 1 , . . . , χ m } to index neurons and denote a single initial neuron as ˆ u 0 ( χ ) = χ . Then, we can describe the projected gradient flow on the empirical loss ̂ L with initialization { χ 1 , . . . , χ m } by:

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 ϵ ∈ (0 , 1) be the target error. Let ρ t be the result of projected gradient flow on the population loss, initialized with the uniform distribution on S d -1 , as defined in Eq. (3.7) . Let T ∗ ,ϵ = inf { t > 0 | L ( ρ t ) ≤ 1 2 (ˆ σ 2 2 ,d +ˆ σ 2 4 ,d ) ϵ 2 } be the earliest time t such that a loss of at most 1 2 (ˆ σ 2 2 ,d +ˆ σ 2 4 ,d ) ϵ 2 is reached. Then, we have

The proof is given in Appendix D. The first term on the right-hand side of Eq. (3.11) 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 ⊤ ⋆ 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 [21, 56, 2] using Wasserstein gradient flow techniques. For instance, while the population dynamics of Abbe et al. [2] only trains the hidden layer for O (1) time, our projected gradient flow updates the hidden layer for O (log d ) time. The main challenge in the proof is to deal with the saddle points that are not strict-saddle [32] in the loss landscape which cannot be escaped simply by adding noise in the parameter space [45, 49, 24]. 3 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. [76] escape the saddles by randomly replacing an exponentially small fraction of neurons, which makes the network require exponential width.

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 ϵ = 1 log log d is the target error and T ∗ ,ϵ is the running time defined in Theorem 3.3. Suppose n ≥ d µ (log d ) Ω(1) for any constant µ > 3 , and the network width m satisfies m ≳ d 2 . 5+ µ/ 2 (log d ) Ω(1) , and m ≤ d C for some sufficiently large universal constant C . Let ˆ ρ t be the projected gradient flow on the empirical loss, initialized with m uniformly sampled weights, defined in Eq. (3.10) . Then, with probability at least 1 -1 d Ω(log d ) over the randomness of the data and the initialization of the finite-width network, we have that L (ˆ ρ T ∗ ,ϵ ) ≲ (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) ϵ 2 .

Plugging in µ = 3 . 1 , Theorem 3.4 suggests that, when the network width is at least d 4 . 05 (up to logarithmic factors), the empirical dynamics of gradient descent could achieve ˆ σ 2 max (1 / log log d ) 2 population loss with d 3 . 1 samples (up to logarithmic factors).

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 can be upper bounded in terms of the growth of the signal q T ⋆ u in the population dynamics. Such a delicate relationship between the growth of the error and that of the signal is the main challenge to proving a sample complexity bound better than NTK. Even an additional constant factor in the growth rate of the coupling error would lead to an additional poly ( d ) factor in the sample complexity. Prior work (e.g. Mei et al. [57], Abbe et al. [2]) also establishes a coupling between the population and empirical dynamics. However, these works obtain bounds on the coupling error which are exponential in the running time - specifically, their bounds have a factor of e KT where K is a universal constant and T is the time. In many settings, including ours, Ω(log d ) time is necessary to achieve a non-trivial error, (as we further discuss in Section 4.2) and thus this would lead to a d O (1) bound on the sample complexity, where O (1) is a unspecified and likely loose constant, which cannot outperform NTK. Our work addresses this challenge by comparing the growth of the coupling error with the growth of the signal, which enables us to control the constant factor in the growth rate of the coupling error.

Projected Gradient Descent with 1 / poly ( d ) Step Size. As a corollary of Theorem 3.4, we show that projected gradient descent with a small learning rate can also achieve low population loss in poly ( d ) iterations. We first define the dynamics of projected gradient descent as follows. As before, we let χ 1 , . . . , χ m i.i.d ∼ S d -1 be the initial weight vectors, and we let ˜ ρ = unif ( { χ 1 , . . . , χ m } ) denote the uniform distribution over these vectors. Thus, we can write a single neuron at initialization as ˜ u 0 ( χ ) = χ for χ ∈ { χ 1 , . . . , χ m } . The dynamics of projected gradient descent, on a finite-width neural network and using the empirical loss ̂ L , can be then described as

where η > 0 is the learning rate. 4 We show that projected gradient descent with a 1 / poly ( d ) learning rate can achieve a low population loss in poly ( d ) iterations:

Theorem 3.5 (Projected Gradient Descent) . Suppose we are in the setting of Theorem 3.4. Let ˜ ρ t be the discrete-time projected gradient descent on the empirical loss, initialized with m weight vectors sampled uniformly from S d -1 , defined in Eq. (3.13) . Finally, assume that η ≤ 1 (ˆ σ 2 2 ,d +ˆ σ 2 4 ,d ) d B for a sufficiently large universal constant B . Then, L (˜ ρ T ∗ ,ϵ /η ) ≲ (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) ϵ 2 .

Theorem 3.5 follows from Theorem 3.4 together with a standard inductive argument to bound the discretization error. The full proof is in Appendix F.

The following theorem states that in the setting of Theorem 3.4, kernel methods with any inner product kernel require Ω( d 4 (ln d ) -6 ) samples to achieve a non-trivial population loss. (Note that the zero function has loss E x ∼ S d -1 [ y ( x ) 2 ] ≥ ( ˆ h 4 ,d ) 2 .)

Theorem 3.6 (Sample complexity lower bound for kernel methods) . Let K be an inner product kernel. Suppose d is larger than a universal constant and n ≲ d 4 (ln d ) -6 . Then, with probability at least 1 / 2 over the randomness of n i.i.d data points { x i } n i =1 drawn from S d -1 , any estimator of y of the form f ( x ) = ∑ n i =1 β i K ( x i , x ) of { x i } n i =1 must have a large error: E x ∼ S d -1 ( y ( x ) -f ( x )) 2 ≥ 3 4 ( ˆ h 4 ,d ) 2 .

Theorem 3.4 and Theorem 3.6 together prove a clear sample complexity separation between gradient flow and NTK. When ˆ h 4 ,d ≍ ˆ σ max (i.e., γ 2 , γ 4 ≍ 1 ), gradient flow with finite-width neural networks can achieve ( ˆ 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 ( ˆ h 4 ,d ) 2 / 2 with d 3 . 9 samples.

On a high level, Abbe et al. [2, 3] 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. [47], Hsu et al. [42], Hsu [41]). 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.6 by proving a Ω( ˆ h 2 k,d ) loss lower bound for any universal constant k ≥ 0 (Theorem H.2). We defer the proof of Theorem 3.6 to Appendix H.

4 Analysis of Population Dynamics

4.1 Symmetry of the Population Dynamics

Akey observation is that due to the symmetry in the data, the population dynamics ρ t has a symmetric distribution in the subspace orthogonal to the vector q ⋆ . As a result, the dynamics ρ t can be precisely characterized by the dynamics in the direction of q ⋆ .

Recall that w = q ⊤ ⋆ u denotes the projection of a weight vector u in the direction of q ⋆ . Let z = ( I -q ⋆ q ⊤ ⋆ ) u be the remaining component. For notational convenience, without loss of generality, we can assume that q ⋆ = e 1 and write u = ( w,z ) , where w ∈ [ -1 , 1] and z ∈ √ 1 -w 2 · S d -2 . We will use this convention throughout the rest of the paper . We will use w t ( χ ) to refer to the first coordinate of the particle u t ( χ ) , and z t ( χ ) to refer to the last ( d -1) coordinates.

Definition 4.1 (Rotational invariance and symmetry) . We say a neural network ρ is rotationally invariant if for u = ( w,z ) ∼ ρ , the distribution of z | w is uniform over √ 1 -w 2 S d -2 almost surely. We also say ρ 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., ˆ ρ ) is very far from rotationally invariant, and therefore the definition is specifically used for population dynamics. If ρ is rotationally invariant and symmetric, then L ( ρ ) has a simpler form that only depends on the marginal distribution of w .

Lemma 4.2. Let ρ be a rotationally invariant neural network. Then, for any target function h and any activation σ ,

In addition, suppose h and σ satisfy Assumption 3.2 and ρ is symmetric. Then

The proof of Lemma 4.2 is deferred to Appendix D.1. Eq. (4.1) says that if ρ is rotationally invariant, then L ( ρ ) only depends on the marginal distribution of w . Eq. (4.2) says that if the distribution of w is additionally symmetric and σ 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 ) ≈ s 2 and P 4 ,d ( s ) ≈ s 4 by Eq. (C.4). Thus, L ( ρ ) essentially corresponds to matching the second and fourth moments of w to some desired values γ 2 and γ 4 . Inspired by this lemma, we define the following key quantities: for any time t ≥ 0 , we define D 2 ,t = E u ∼ ρ t [ P 2 ,d ( w )] -γ 2 and D 4 ,t = E u ∼ ρ t [ P 4 ,d ( w )] -γ 4 , where ρ t is defined according to the population, infinite-width dynamics (Eq. (3.6), Eq. (3.7) and Eq. (3.8)).

We next show that the rotational invariance and symmetry properties of ρ t are indeed maintained:

Lemma 4.3. Suppose we are in the setting of Theorem 3.3. At any time t ∈ [0 , ∞ ) , ρ t is symmetric and rotationally invariant.

Rotational invariance follows from the rotational invariance of the data. To show symmetry, we use Eq. (4.1), and the facts that P k,d is an odd polynomial for odd k and ρ t is symmetric at initialization. Using 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 χ ∈ S d -1 , writing w t := w t ( χ ) , we have

where for any w ∈ [ -1 , 1] , we have P t ( w ) = 2ˆ σ 2 2 ,d D 2 ,t w +4ˆ σ 2 4 ,d D 4 ,t w 3 , and Q t ( w ) = λ d (1) w + λ d (3) w 3 , where | λ d (1) | , | λ d (3) | ≲ ˆ σ 2 2 ,d | D 2 ,t | +ˆ σ 2 4 ,d | D 4 ,t | d . More specifically, λ d (1) = 2ˆ σ 2 2 ,d D 2 ,t · 1 d -1 -2ˆ σ 2 4 ,d D 4 ,t · 6 d +12 d 2 -1 and λ d (3) = 4ˆ σ 2 4 ,d D 4 ,t · 6 d +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 . We henceforth refer to the w t ( χ ) as particles - this is well-defined by Lemma 4.4.

4.2 Analysis of One-Dimensional Population Dynamics

We also use ι ∈ [ -1 , 1] to refer to the first coordinate of χ , the initialization of a particle under the population dynamics. We note that ι = ⟨ χ, e 1 ⟩ and therefore, the distribution of ι is µ d . For any time t ≥ 0 , we use w t ( ι ) to refer to w t ( χ ) for any χ ∈ S d -1 . This notation is also well-defined by Lemma 4.4. We divide our proof of Theorem 3.3 into three phases, which are defined as follows. For ease of presentation, we will only consider particles w t ( ι ) for ι > 0 in the following discussion. Our argument also applies for ι < 0 by the symmetry of ρ t (Lemma 4.3).

Definition 4.5 (Phase 1) . Let w max = 1 log d and ι U = log d √ d . Let T 1 > 0 be the minimum time such that w T 1 ( ι U ) = w max := 1 log d . We refer to the time interval [0 , T 1 ] as Phase 1.

Note that essentially all of the particles are less than ι U = log d √ d at initialization by tail bounds for µ d . During Phase 1, the term corresponding to D 2 ,t in the velocity dominates, and all particles grow by a √ d (log d ) 2 factor. During Phase 1, the loss does not decrease much, but a large portion of the particles have grown by a large factor. During Phase 2, the particles and their velocity will become large.

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 D.8.

Definition 4.7 (Phase 3) . Phase 3 is defined as the time interval [ T 2 , ∞ ) .

We divide this phase into two cases, which are both more challenging to analyze than Phases 1 and 2: (i) D 2 ,T 2 = 0 and D 4 ,T 2 < 0 , and (ii) D 2 ,T 2 < 0 and D 4 ,T 2 = 0 . Here we discuss Case 1 the analyses of Cases 1 and 2 are in Appendix D.5 and Appendix D.6 respectively. Suppose Case 1 holds, i.e. D 2 ,T 2 = 0 and D 4 ,T 2 < 0 . Then, for all t ≥ T 2 , we will have D 2 ,t ≥ 0 and D 4 ,t ≤ 0 (Lemma D.19 and Lemma D.20), meaning that for t ≥ T 2 a poly ( γ 2 ) fraction of particles will be far from 0 and 1 (Lemma D.17). However, this does not guarantee a large average velocity - unlike

in Phases 1 and 2, the velocity v ( w ) = P t ( w ) + Q t ( w ) defined in Lemma 4.4 may have a positive root r ∈ (0 , 1) , and the root r can give rise to bad stationary points because for certain values of r , if ρ degenerates to the singleton distribution on the root r , the velocity at r would be exactly 0 . Indeed, we can construct such an r , intuitively because D 2 ,t and D 4 ,t have different signs in P t ( w ) and Q t ( w ) ≈ O (1 /d ) is a lower order term (Lemma D.35).

Thus, we must leverage some information about the trajectory to show that the average velocity is large when D 2 ,t , D 4 ,t have different signs. To this end, we prove that E u ∼ ρ t [( w -r ) 2 ] ≳ E w,w ′ ∼ ρ t ( w -w ′ ) 2 is large by designing a novel potential function Φ( w ) := log( w √ 1 -w 2 ) . We show that | Φ( w ) -Φ( w ′ ) | is always increasing for any two particles w,w ′ (Lemma D.13). Because Φ is Lipschitz on an interval bounded away from 0 and 1 , a lower bound on | Φ( w ) -Φ( w ′ ) | also leads to a lower bound on ( w -w ′ ) 2 , and thus the particles away from 0 and 1 will have a large variance. Recall that when D 2 ,t > 0 and D 4 ,t < 0 , a large portion of particles are away from 0 and 1 , and hence, the average velocity is large. In Appendix B, we include simulations to illustrate the effects of Phases 1, 2 and 3.

5 Analysis of Empirical, Finite-Width Dynamics

Coupling Between Empirical and Population Dynamics. To analyze the difference between the empirical and population dynamics, we define an intermediate process ¯ ρ t , where the initial particles are from ˆ ρ 0 , but the dynamics of the particles then follow the population trajectories:

For χ ∼ unif ( { χ 1 , . . . , χ m } ) , let ¯ u t ( χ ) = u t ( χ ) . Let Γ t be the joint distribution of (¯ u t ( χ ) , ˆ u t ( χ )) . Then, Γ t forms a natural coupling between ¯ ρ t and ˆ ρ t . We will use ¯ u t and ˆ u t as shorthands for the random variables ¯ u t ( χ ) , ˆ u t ( χ ) respectively in the rest of this section. We define the average distance ∆ 2 t := E (ˆ u t , ¯ u t ) ∼ Γ t [ ∥ ˆ u t -¯ u t ∥ 2 2 ] . Intuitively, f ρ t ( x ) and f ˆ ρ t ( x ) are close when ˆ u t and ¯ u t are close, which is formalized by the following lemma:

Lemma 5.1. In the setting of Theorem 3.4, let T ≤ d O (1) ˆ σ 2 2 ,d +ˆ σ 2 4 ,d . Then, with probability at least 1 -exp( -d 2 ) over the initialization, we have for all t ∈ [0 , T ] that E x ∼ S d -1 [( f ρ t ( x ) -f ¯ ρ t ( x )) 2 ] ≲ (ˆ σ 2 2 ,d +ˆ σ 2 4 ,d ) d 2 (log d ) O (1) m and E x ∼ S d -1 [( f ˆ ρ t ( x ) -f ¯ ρ t ( x )) 2 ] ≲ (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d )∆ 2 t .

The proof of Lemma 5.1 is deferred to Appendix E. As a simple corollary of Lemma 5.1, we can upper bound E x ∼ S d -1 [( f ρ t ( x ) -f ˆ ρ t ( x )) 2 ] by the triangle inequality. Thus, so long as ∆ T ∗ ,ϵ is small, we can show that the empirical and population dynamics achieve similar test error at time T ∗ ,ϵ .

Upper Bound for ∆ T ∗ ,ϵ . The following lemma gives our upper bound on ∆ T ∗ ,ϵ :

Lemma 5.2 (Final Bound on ∆ T ∗ ,ϵ ) . In the setting of Theorem 3.4, we have ∆ T ∗ ,ϵ ≤ d -µ -1 4 .

We prove Lemma 5.2 by induction over the time t ≤ T ∗ ,ϵ . Let us define the maximal distance ∆ max ,t := max (ˆ u t , ¯ u t ) ∈ supp (Γ t ) ∥ ˆ u t -¯ u t ∥ 2 where the max is over the finite support of the coupling Γ t . We will control ∆ t and ∆ max ,t simultaneously using induction. The inductive hypothesis is:

Assumption 5.3 (Inductive Hypothesis) . For some universal constant 1 > ϕ > 1 / 2 and ψ ∈ (0 , ϕ -1 2 ) , we say that the inductive hypothesis holds at time T if, for all 0 ≤ t ≤ T , we have ∆ t ≤ d -ϕ and ∆ max ,t ≤ d -ψ .

Showing that the inductive hypothesis holds for all t requires studying the growth of the error ∥ ˆ u t -¯ u t ∥ . We analyze it using the decomposition 5 d dt ∥ ˆ u t -¯ u t ∥ 2 2 = A t + B t + C t , where A t := -2 ⟨ grad ˆ u L ( ρ t ) -grad ¯ u L ( ρ t ) , ˆ u t -¯ u t ⟩ , B t := -2 ⟨ grad ˆ u L (ˆ ρ t ) -grad ˆ u L ( ρ t ) , ˆ u t -¯ u t ⟩ , and C t := -2 ⟨ grad ˆ u ̂ L (ˆ ρ t ) -grad ˆ u L (ˆ ρ t ) , ˆ u t -¯ u t ⟩ . Intuitively, A t captures the growth of the coupling error due to the population dynamics, B t the growth due to the discrepancy between the gradients from the

finite-width and infinite-width networks, and C t the growth due to the discrepancy between finite samples and infinite samples. The following lemma gives upper bounds on these 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 ≤ T ∗ ,ϵ . Let δ t := ˆ u t -¯ u t . Let T 1 be the runtime of Phase 1 (where Phase 1 is defined in Definition 4.5). Then, we have

Additionally, with probability at least 1 -1 d Ω(log d ) over the initialization and dataset, we have B t ≲ (ˆ σ 2 2 ,d +ˆ σ 2 4 ,d ) · ( d (log d ) O (1) √ m +∆ t ) ∥ δ t ∥ 2 , E [ B t ] ≲ (ˆ σ 2 2 ,d +ˆ σ 2 4 ,d ) d (log d ) O (1) √ m ∆ t +(ˆ σ 2 2 ,d +ˆ σ 2 4 ,d )∆ 3 t and C t ≲ (ˆ σ 2 2 ,d +ˆ σ 2 4 ,d )(log d ) O (1) · (√ d n ∥ δ t ∥ 2 + d 2 -2 ψ n ∆ ∥ δ t ∥ 2 2 + d 2 . 5 -2 ψ n ∆ 2 ∥ δ t ∥ 2 + d 4 -4 ψ n ∆ 2 ∥ δ t ∥ 2 2 ) .

The proof is in Appendix E. We now give an explanation for each of these bounds. 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ˆ σ 2 2 ,d | D 2 ,t | during Phase 1 (which follows from Lemma 4.4) and in Eq. (5.2) we show that the growth rate of the coupling error due to the growth in the signal is also at most 4ˆ σ 2 2 ,d | 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 δ t is mostly in the direction of e 1 , the factor by which δ 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 ( ˆ w ) for any two particles w, ˆ 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ˆ σ 2 2 ,d | D 2 ,t | w is important - a larger constant would lead to δ T 1 being larger by a poly ( d ) factor, thus increasing the sample complexity by 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 ( ˆ w ) than to the growth of w , or vice versa. Thus we use a looser bound O (1) · (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) γ 2 ∥ δ t ∥ 2 in Eq. (5.2). However, since the remaining running time after Phase 1, T ∗ ,ϵ -T 1 , is very short (only 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( 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 E [ B t ] than for B t . In our proof, we use the bounds for B t and E [ B t ] to prove the inductive hypothesis for ∆ t and ∆ 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 I.7).

All of these bounds together control the growth of ∥ ˆ u t -¯ u t ∥ at time t . Using Lemma 5.4 we can show that for some properly chosen ϕ and ψ , the inductive hypothesis in Assumption 5.3 holds until T ∗ ,ϵ (see Lemma E.3 for the statement), which naturally leads to the bound in Lemma 5.2. The full proof can be found in Appendix E.

6 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. The work of Arous et al. [11] shows that if the population loss has information exponent 2 , then a sample complexity of d (up to logarithmic factors) can be achieved with a single-neuron student network it is an open question if a similar sample complexity could be obtained in our setting. A final open question is whether two-layer neural networks can be shown to attain arbitrarily small generalization error ϵ -one limitation of our analysis is that we require ϵ ≥ 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.

References

Emmanuel Abbe, Enric Boix-Adsera, Matthew S Brennan, Guy Bresler, and Dheeraj Nagaraj. The staircase property: How hierarchical structure can guide deep learning. Advances in Neural Information Processing Systems , 34:26989-27002, 2021.

Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory , pages 4782-4887. PMLR, 2022.

Emmanuel Abbe, Enric Boix-Adsera, and Theodor Misiakiewicz. Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics. arXiv preprint arXiv:2302.11055 , 2023.

Radosław Adamczak, Alexander E. Litvak, Alain Pajor, and Nicole Tomczak-Jaegermann. Quantitative estimates of the convergence of the empirical covariance matrix in log-concave ensembles. Journal of the American Mathematical Society , 23(2):535-561, oct 2009. doi: 10.1090/ s0894-0347-09-00650-x. URL https://doi.org/10.1090%2Fs0894-0347-09-00650-x .

Zeyuan Allen-Zhu and Yuanzhi Li. What can resnet learn efficiently, going beyond kernels? arXiv preprint arXiv:1905.10337 , 2019.

Zeyuan Allen-Zhu and Yuanzhi Li. Backward feature correction: How deep learning performs deep learning. arXiv preprint arXiv:2001.04413 , 2020.

Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. Advances in neural information processing systems , 32, 2019.

Luca Arnaboldi, Ludovic Stephan, Florent Krzakala, and Bruno Loureiro. From highdimensional & mean-field dynamics to dimensionless odes: A unifying approach to sgd in two-layers networks, 2023.

Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. arXiv preprint arXiv:1904.11955 , 32, 2019. URL https://proceedings.neurips.cc/paper/2019/ file/dbc4d84bfcfe2284ba11beffb853a8c4-Paper.pdf .

Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. arXiv preprint arXiv:1901.08584 , pages 322-332, 2019.

Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from high-dimensional inference. The Journal of Machine Learning Research , 22(1):4788-4838, 2021.

Kendall Atkinson and Weimin Han. Spherical harmonics and approximations on the unit sphere: an introduction , volume 2044. Springer Science & Business Media, 2012.

Drumi D Bainov and Pavel S Simeonov. Integral inequalities and applications , volume 57. Springer Science & Business Media, 1992.

Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit. Advances in Neural Information Processing Systems , 35:21750-21764, 2022.

Gérard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Algorithmic thresholds for tensor pca. 2020.

Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-index models with shallow neural networks. arXiv preprint arXiv:2210.15651 , 2022.

Nicolas Boumal. An introduction to optimization on smooth manifolds . Cambridge University Press, 2023. doi: 10.1017/9781009166164. URL https://www.nicolasboumal.net/book .

Alon Brutzkus and Amir Globerson. Globally optimal gradient descent for a convnet with gaussian inputs. arXiv preprint arXiv:1702.07966 , 2017.

Fan Chen, Zhenjie Ren, and Songbo Wang. Uniform-in-time propagation of chaos for mean field langevin dynamics, 2023.

Minshuo Chen, Yu Bai, Jason D Lee, Tuo Zhao, Huan Wang, Caiming Xiong, and Richard Socher. Towards understanding hierarchical learning: Benefits of neural representations. Neural Information Processing Systems (NeurIPS) , 2020.

Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for overparameterized models using optimal transport. Advances in neural information processing systems , 31, 2018.

Lenaic Chizat and Francis Bach. A note on lazy training in supervised differentiable programming. arXiv preprint arXiv:1812.07956 , 8, 2018.

Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference on Learning Theory , pages 5413-5452. PMLR, 2022.

Hadi Daneshmand, Jonas Kohler, Aurelien Lucchi, and Thomas Hofmann. Escaping saddles with stochastic gradients. In International Conference on Machine Learning , pages 1155-1164. PMLR, 2018.

Amit Daniely and Eran Malach. Learning parities with neural networks. Advances in Neural Information Processing Systems , 33:20356-20365, 2020.

Amit Daniely, Roy Frostig, and Yoram Singer. Toward deeper understanding of neural networks: The power of initialization and a dual view on expressivity. In Advances In Neural Information Processing Systems , pages 2253-2261, 2016.

Kefan Dong and Tengyu Ma. Toward l ∞ -recovery of nonlinear functions: A polynomial sample complexity bound for gaussian random fields. arXiv preprint arXiv:2305.00322 , 2023.

Xialiang Dou and Tengyuan Liang. Training neural networks as learning data-adaptive kernels: Provable representation and approximation benefits. Journal of the American Statistical Association , 116(535):1507-1520, 2021.

Simon S Du and Jason D Lee. On the power of over-parametrization in neural networks with quadratic activation. International Conference on Machine Learning (ICML) , 2018.

Simon S Du, Jason D Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. arXiv preprint arXiv:1811.03804 , pages 1675-1685, November 2018.

Wolfgang Gabcke. New derivation and explicit remainder estimation of the Riemann-Siegel formula . PhD thesis, Georg-August-Universityät Göttingen, 1979.

Rong Ge, Furong Huang, Chi Jin, and Yang Yuan. Escaping from saddle points-online stochastic gradient for tensor decomposition. In Conference on Learning Theory , COLT 2015, pages 797-842, 2015.

Rong Ge, Jason D Lee, and Tengyu Ma. Learning one-hidden-layer neural networks with landscape design. arXiv preprint arXiv:1711.00501 , 2017.

Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. When do neural networks outperform kernel methods? Advances in Neural Information Processing Systems , 33: 14820-14830, 2020.

Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. arXiv preprint arXiv:1802.08246 , 2018.

Suriya Gunasekar, Jason D Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. In Advances in Neural Information Processing Systems , pages 9461-9471, 2018.

Karl Hajjar and Lenaic Chizat. On the symmetries in the dynamics of wide two-layer neural networks, 2023.

Moritz Hardt and Tengyu Ma. Identity matters in deep learning. arXiv preprint arXiv:1611.04231 , 2016.

Moritz Hardt, Tengyu Ma, and Benjamin Recht. Gradient descent learns linear dynamical systems. Journal of Machine Learning Research , 19:1-44, 2018.

Elad Hazan, Kfir Levy, and Shai Shalev-Shwartz. Beyond convexity: Stochastic quasi-convex optimization. In Advances in Neural Information Processing Systems , pages 1594-1602, July 2015.

Daniel Hsu. Dimension lower bounds for linear approaches to function approximation. Daneil Hsu's homepage , 2021.

Daniel Hsu, Clayton H Sanford, Rocco Servedio, and Emmanouil Vasileios VlatakisGkaragkounis. On the approximation power of two-layer networks of random relus. In Conference on Learning Theory , pages 2423-2461. PMLR, 2021.

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in neural information processing systems , pages 8571-8580, 2018.

Adel Javanmard, Marco Mondelli, and Andrea Montanari. Analysis of a two-layer neural network via displacement convexity. 2020.

Chi Jin, Praneeth Netrapalli, Rong Ge, Sham M Kakade, and Michael I Jordan. On nonconvex optimization for machine learning: Gradients, stochasticity, and saddle points. Journal of the ACM (JACM) , 68(2):1-29, 2021.

Sham M Kakade, Varun Kanade, Ohad Shamir, and Adam Kalai. Efficient learning of generalized linear and single index models with isotonic regression. In Advances in Neural Information Processing Systems , pages 927-935, 2011.

Pritish Kamath, Omar Montasser, and Nathan Srebro. Approximate is good enough: Probabilistic variants of dimensional and margin complexity. In Conference on Learning Theory , pages 2236-2262. PMLR, 2020.

Kenji Kawaguchi. Deep learning without poor local minima. Advances in neural information processing systems , 29, 2016.

Jason D Lee, Max Simchowitz, Michael I Jordan, and Benjamin Recht. Gradient descent converges to minimizers. arXiv preprint arXiv:1602.04915 , 1050:16, 2016.

Yuanzhi Li and Yingyu Liang. Learning overparameterized neural networks via stochastic gradient descent on structured data. In Advances in Neural Information Processing Systems , pages 8157-8166, 2018. URL https://www.aaai.org/ocs/index.php/AAAI/AAAI18/ paper/view/16067 .

Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularization in over-parameterized matrix sensing and neural networks with quadratic activations. In Conference On Learning Theory , pages 2-47. PMLR, 2018.

Yuanzhi Li, Tengyu Ma, and Hongyang R Zhang. Learning over-parametrized two-layer neural networks beyond ntk. In Conference on Learning Theory , pages 2613-2682. PMLR, 2020.

Tengyu Ma. Lecture notes for machine learning theory (cs229m/stats214), June 2022.

Song Mei and Andrea Montanari. The generalization error of random features regression: Precise asymptotics and double descent curve. arXiv preprint arXiv:1908.05355 , 2019.

Song Mei, Yu Bai, and Andrea Montanari. The landscape of empirical risk for nonconvex losses. The Annals of Statistics , 46(6A):2747-2774, 2018.

Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layers neural networks. Proceedings of the National Academy of Sciences , pages E7665-E7671, 2018.

Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. In Conference on Learning Theory , pages 2388-2464. PMLR, 2019.

Alireza Mousavi-Hosseini, Sejun Park, Manuela Girotti, Ioannis Mitliagkas, and Murat A Erdogdu. Neural networks efficiently learn low-dimensional representations with sgd. arXiv preprint arXiv:2209.14863 , 2022.

Eshaan Nichani, Yu Bai, and Jason D. Lee. Identifying good directions to escape the NTK regime and efficiently learn low-degree plus sparse polynomials. In NeurIPS , 2022. URL http://papers.nips.cc/paper_files/paper/2022/hash/ 5d6ae8ba43ecb378030753c4408ef9bd-Abstract-Conference.html .

Samet Oymak and Mahdi Soltanolkotabi. Toward moderate overparameterization: Global convergence guarantees for training shallow neural networks. IEEE Journal on Selected Areas in Information Theory , 1(1):84-105, 2020.

Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks. arXiv preprint arXiv:1805.01053 , 2018.

Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: A central limit theorem. Stochastic Processes and their Applications , 130(3):1820-1852, 2020.

Maciej Skorski. Bernstein-type bounds for beta distribution. Modern Stochastics: Theory and Applications , 10(2):211-228, 2023.

Mahdi Soltanolkotabi. Learning ReLUs via gradient descent. In Advances in Neural Information Processing Systems , pages 2007-2017, 2017.

Mahdi Soltanolkotabi, Adel Javanmard, and Jason D Lee. Theoretical insights into the optimization landscape of over-parameterized shallow neural networks. IEEE Transactions on Information Theory , 65(2):742-769, 2018.

Daniel Soudry and Yair Carmon. No bad local minima: Data independent training error guarantees for multilayer neural networks. arXiv preprint arXiv:1605.08361 , 2016.

Taiji Suzuki and Shunta Akiyama. Benefit of deep learning with non-convex noisy gradient descent: Provable excess risk bound and superiority to kernel methods. In International Conference on Learning Representations .

Taiji Suzuki, Denny Wu, and Atsushi Nitanda. Convergence of mean-field langevin dynamics: Time and space discretization, stochastic gradient, and variance reduction, 2023.

Yan Shuo Tan and Roman Vershynin. Online stochastic gradient descent with arbitrary initialization solves non-smooth, non-convex phase retrieval. arXiv preprint arXiv:1910.12837 , 2019.

Matus Telgarsky. Feature selection with gradient descent on two-layer networks in low-rotation regimes. arXiv preprint arXiv:2208.02789 , 2022.

Yuandong Tian. An analytical formula of population gradient for two-layered relu network and its applications in convergence and critical point analysis. arXiv preprint arXiv:1703.00560 , 2017.

Gal Vardi, Gilad Yehudai, and Ohad Shamir. Learning a single neuron with bias using gradient descent. Advances in Neural Information Processing Systems , 34:28690-28700, 2021.

Roman Vershynin. Introduction to the non-asymptotic analysis of random matrices. arXiv preprint arXiv:1011.3027 , 2010.

Roman Vershynin. High-dimensional probability: An introduction with applications in data science , volume 47. Cambridge university press, 2018.

Colin Wei, Jason D. Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization of neural nets v.s. their induced kernel. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d'Alché-Buc, Emily B. Fox, and Roman Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada , pages 9709-9721, 2019. URL https://proceedings.neurips.cc/paper/2019/hash/ 8744cf92c88433f8cb04a02e6db69a0d-Abstract.html .

Colin Wei, Jason D Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization of neural nets vs their induced kernel. In Advances in Neural Information Processing Systems , pages 9709-9721, 2019.

Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. arXiv preprint arXiv:2002.09277 , 2020.

Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Sham M Kakade. Learning high-dimensional single-neuron relu networks with finite samples. arXiv preprint arXiv:2303.02255 , 2023.

Weihang Xu and Simon S Du. Over-parameterization exponentially slows down gradient descent for learning a single neuron. arXiv preprint arXiv:2302.10034 , 2023.

Xiao Zhang, Yaodong Yu, Lingxiao Wang, and Quanquan Gu. Learning one-hidden-layer relu networks via gradient descent. In The 22nd international conference on artificial intelligence and statistics , pages 1524-1534. PMLR, 2019.

List of Appendices

A Additional Related Works

In addition to the related works on the NTK approach, mean-field analyses, and other works analyzing the training dynamics of neural networks that we discussed in Section 1, we also discuss the following additional related works.

There are several works which study mean-field Langevin dynamics [19, 68]. These works show 'uniform-in-time propagation of chaos' estimates, i.e. their results imply bounds on the coupling error between finite-width and infinite-width trajectories which do not grow exponentially in the time T . However, it is likely non-trivial to apply these analyses to directly obtain good test error and sample complexity in the setting of two-layer neural networks, since mean-field Langevin dynamics may require many iterations to obtain good population loss. More concretely, Theorem 4 of Mei et al. [56] suggests that the inverse temperature λ for mean-field Langevin dynamics has to be at least proportional to the dimension d in order for Langevin dynamics to achieve low population loss. This would cause the log-Sobolev constant in Suzuki et al. [68] to be on the order of e -d , leading to a running time of e -d . In comparison, we show in Theorem 3.5 that projected gradient descent with poly ( d ) iterations can achieve a low population loss. We also note that Chen et al. [19] do not study the impact of finite samples on the coupling error.

We also note that there are several works which reduce the population dynamics to a lowerdimensional dynamics [2, 37, 8], as we do in Section 4.1. The focus of our work is on the analysis of the population dynamics (Section 4.2) and coupling error between the empirical and population dynamics (Section 5), rather than the dimension-free dynamics.

Finally, we note that our setting goes beyond that of Abbe et al. [2]: our target function does not satisfy the merged-staircase property required by Abbe et al. [2], since the lowest-degree term in our target function has degree 2 . The work of Abbe et al. [3] studies functions which have higher 'leap complexity,' meaning that terms of the target function can introduce more than one new coordinate at once, or introduce a new coordinate using a Hermite polynomial with degree greater than 1 (in the case where the inputs are Gaussian). However, Abbe et al. [3] use a non-standard algorithm which separates the coordinates based on whether they are large or small, and applying different projections to each subset of coordinates.

B Simulations for Population Dynamics

Here we include some simulations in a setting which is a simplified version of our 1-dimensional dynamics, to illustrate the different phases. Here, the goal is to find a weight vector w ∈ R d , where d = 1000 , that minimizes the following loss:

where γ 2 = 0 . 8 and γ 4 = 0 . 7 . Intuitively, the w i play the role of the particles w in the 1 -dimensional dynamics of Section 4.2. We initialize w as follows - each coordinate is initialized according to a standard Gaussian distribution with mean 0 and standard deviation 0 . 0001 . We then train w using gradient descent on L ( w ) with learning rate 0 . 1 , for 100 , 000 iterations. We define D 2 = 1 d ∑ d i =1 w 2 i -γ 2 and D 4 = 1 d ∑ d i =1 w 4 i -γ 4 .

The dynamics are illustrated by Figure 1, Figure 2, Figure 3 and Figure 4. At initialization, since the standard deviation is 0 . 0001 , all of the w i are very small, as shown in Figure 1. Now, as shown in Figure 4, while both D 2 and D 4 grow rapidly after a certain point, D 4 reaches 0 first. At the point where D 4 reaches 0 , the distribution of the coordinates of w is shown in Figure 2. At this point, the coordinates w i have grown substantially, with a substantial fraction of coordinates around 0 . 5 or larger. This illustrates the intuition that all of the coordinates grow by a large factor during Phases 1 and 2. The weight histogram after many iterations have passed during Phase 3 is shown in Figure 3. In our infinite-width 1-dimensional population dynamics, when D 4 > 0 and D 2 < 0 , the particles w which are larger than the root r of the velocity will move towards 0 , while the particles w which are smaller than the root r of the velocity will move away from 0 . This intuition is reflected in Figure 3, where not as many particles are concentrated around 0 compared to at the end of Phase 2, but the particles have not grown larger than 1 or -1 in absolute value.

WeightHistogram atIteration0

250

200

150

100

50

-0.0003-0.0002-0.0001

0.0000

0.0001

0.0002

0.0003

End of Phases 1 and 2-D4>0-Time 24349

250

200

150

100

50

0

-1.5

1.0

-0.5

0.0

0.5

1.0

1.5

Final Weight Histogram

350

250

200

150

100

50

0

1.0

-0.5

0.0

0.5

1.0

D2 and D4

0.0

-0.2

0.4

-0.6

D2

D4

-0.8

0

0

20000

40000

60000

80000

100000

C Background on Spherical Harmonics

In the following, we summarize the background needed for this paper based on Atkinson and Han [12, Section 2].

Spherical Harmonics. Spherical harmonics are the eigenfunctions of any inner product kernels on the unit sphere S d -1 . The eigenfunctions corresponding to the k -th eigenvalue are degreek polynomials, and form a Hilbert space Y d k with inner product ⟨ f, g ⟩ = E x ∼ S d -1 [ f ( x ) g ( x )] . The dimension of Y d k is N k,d = ( d + k -1 d -1 ) -( d + k -3 d -1 ) . For any universal constant k , N k,d ≍ d k . Spherical harmonics with different degrees are orthogonal to each other, and their linear combinations can represent all square-integrable functions over S d -1 .

The following theorem states that Y d k is the eigenspace corresponding to the k -th eigenvalue for any inner product kernel on the sphere.

Theorem C.1 (Funk-Hecke formula [12]) . Let h : [ -1 , 1] → R be any one-dimensional function with ∫ 1 -1 | h ( t ) | µ d ( t )d t < ∞ , and λ k = N -1 / 2 k,d E t ∼ µ d [ h ( t ) P k,d ( t )] . Then for any function Y k ∈ Y k,d ,

Legendre Polynomials. The un-normalized Legendre polynomial is defined recursively by

In particular, we have

For any u ∈ S d -1 , P k,d ( ⟨ u, x ⟩ ) is a spherical harmonic of degree k and, for every k, k ′ ≥ 0 and u 1 , u 2 ∈ S d -1 , E x ∼ S d -1 [ P k,d ( ⟨ u 1 , x ⟩ ) P k ′ ,d ( ⟨ u 2 , x ⟩ )] = 1 [ k = k ′ ] P k,d ( ⟨ u 1 , u 2 ⟩ ) (Lemma C.5).

Projection Operator. Let Π k be the projection operator to Y d k . For any function f : S d -1 → R , the projection operator Π k is given by

In addition, if f is a function of the form f ( x ) = h ( ⟨ u, x ⟩ ) for some fixed vector u ∈ S d -1 and one-dimensional function h : [ -1 , 1] → R , the Funk-Hecke formula (Theorem C.1) implies

where ˆ h k,d = E t ∼ µ d [ h ( t ) P k,d ( t )] is the k -th coefficient in the Legendre decomposition of h .

Note that by the orthogonality of Legendre polynomials, we have

The following lemma computes the expectation of P k,d ( ⟨ u, x ⟩ ) with respect to a distribution on u which is rotationally invariant on the last ( d -1) coordinates.

Lemma C.2. Let ρ be a distribution on S d -1 that is rotationally invariant as in Definition 4.1. Then, for any x ∈ S d -1 we have E u ∼ ρ [ P k,d ( ⟨ u, x ⟩ )] = E u ∼ ρ [ P k,d ( ⟨ e 1 , u ⟩ )] P k,d ( ⟨ e 1 , x ⟩ ) .

Proof. We prove this lemma by invoking Eq. (2.167) of Atkinson and Han [12], which states that for every s, t ∈ [ -1 , 1] and k ≥ 0 , d ≥ 3 ,

In the following, for x ∈ S d -1 , we write x = ( ⟨ e 1 , x ⟩ , ζ ) , and for u ∼ ρ we write u = ( w,z ) . Note that when ρ is symmetric and u ∼ ρ , conditioned on u 1 we have ⟨ z, ζ ⟩ ∼ (1 -u 2 1 ) 1 / 2 (1 -x 2 1 ) 1 / 2 µ d -1 for any fixed x . As a result, plugging in t = w,s = ⟨ e 1 , x ⟩ and ξ = (1 -u 2 1 ) -1 / 2 (1 -x 2 1 ) -1 / 2 ⟨ z, ζ ⟩ we get

Finally, taking expectation over w we prove the desired result.

Additional Useful Lemmas. In the following, we present some useful lemmas about spherical harmonics and Legendre polynomials.

The following proposition computes the coefficient in the Legendre polynomial decomposition of ReLU activation.

Proposition C.3 (Lemma C.2 of Dong and Ma [27]) . Let σ ( t ) = max { t, 0 } be the ReLU activation. For every d ≥ 3 and even k , we have

As a corollary, when k is an absolute constant we have

Theorem C.4 (Addition Theorem [12]) . Let { Y k,j : 1 ≤ j ≤ N k,d } be an orthonormal basis of Y k,d , that is,

Lemma C.5. For every k, k ′ ≥ 0 , d ≥ 3 and u, v ∈ S d -1 , we have

Proof. Let { Y k,j : 1 ≤ j ≤ N k,d } be an orthonormal basis of Y k,d with

Then for all x, z ∈ S d -1

Similarly define { Y k ′ ,j : 1 ≤ j ≤ N k ′ ,d } for Y k ′ ,d . The addition theorem (Theorem C.4) implies that

Recall that for k = k ′ , Y k,d and Y k ′ ,d are orthononal subspaces. Consequently,

̸

Lemma C.6. For any k ≥ 0 , d ≥ 3 , there exists a feature mapping ϕ : S d -1 → R N k,d such that for every u, v ∈ S d -1 ,

Proof. Let Y 1 , · · · , Y N k,d : S d -1 → R be an orthonormal basis for the degreek spherical harmonics Y k,d . We let ϕ ( x ) = ( Y l ( x )) l ∈ [ N k,d ] be the feature mapping, and in the following we verify Eqs. (C.21)-(C.23). By the addition theorem (Theorem C.4), for every u, v ∈ S d -1 we have

which proves Eq. (C.21). Note that for every u ∈ S d -1 , P k,d ( ⟨ u, u ⟩ ) = P k,d (1) = 1 . Hence Eq. (C.22) follows directly. To prove Eq. (C.23), note that Y 1 , · · · , Y N k,d are orthonormal. Therefore for every i, j ∈ [ N k,d ]

Hence, Eq. (C.23) follows directly.

Lemma C.7. Let µ d be the distribution of ⟨ u, v ⟩ when v is drawn uniformly at random from the unit sphere S d -1 and u ∈ S d -1 is a fixed vector. For fixed n ≥ 1 , let x 1 , · · · , x n ∈ S d -1 be (not necessarily independent) random variables drawn from the distribution µ d . Then there exists a universal constant c > 0 such that for any integer k ≥ 1 ,

In addition, we also have

Proof. To prove the first part of this lemma, we first invoke Proposition C.8, which states that

When x ∼ µ d , Proposition I.5 states that

Therefore for every t ≥ 4 1+ k ( k ln d ) k we have

By union bound we get

which proves the first part of this lemma.

As a corollary, for any fixed ϵ > 4 1+ k ( k ln d ) k we have

Now we upper bound the integral by changing variables. Letting u = t 1 /k / 32 , we get

(C.40)

where the last inequality comes from an upper bound for the incomplete gamma functions [31, Theorem 4.4.3]. Therefore we get

Finally taking ϵ = (32 k ln( dn )) k , we prove the desired result.

Proposition C.8. For any k ≥ 0 , d ≥ 2 , we have

Proof. Let µ d ( t ) ≜ (1 -t 2 ) d -3 2 Γ( d/ 2) Γ(( d -1) / 2) 1 √ π be the density of u 1 when u = ( u 1 , · · · , u d ) is drawn uniformly from S d -1 . By Atkinson and Han [12, Theorem 2.24] we have

As a result,

By Proposition I.5 we have,

Consequently,

Hence, when t ≥ √ k ln d d we get

D Proofs for Population Dynamics

D.1 Proof of Lemma 4.2

In the following, we prove Lemma 4.2, which provides a way to simplify L ( ρ ) when ρ is rotationally invariant and symmetric.

Proof of Lemma 4.2. Recall that the population loss is

and the activation function σ : R → R and target function h : R → R has the following decomposition:

Therefore we have

When the neural network ρ is rotationally invariant, we can simplify the above equation by invoking Lemma C.2, which states that

Continuing Eq. (D.4) above we have,

Now we expand the square by invoking Lemma C.5. In particular, Lemma C.5 states that

Consequently, as desired.

To prove the second part of this lemma, in the following, we assume that (1) ˆ h k,d = 0 , ∀ k ̸∈ { 0 , 2 , 4 } , (2) ˆ h 0 ,d = ˆ σ 0 ,d , and (3) that σ is a degree-4 polynomial. Then we get

Since ρ is symmetric and P k,d ( w ) is an odd function when k is odd, we get

It follows directly that

D.2 Proof of Lemma 4.3 and Lemma 4.4

The goal of this subsection is to show Lemma 4.3, which states that ρ t remains symmetric and rotationally invariant under the infinite-width population dynamics, and to show Lemma 4.4 which describes the dynamics of the first coordinates of the particles in ρ t .

The following lemma is a first step. Recall that when ρ is rotationally invariant, then Lemma 4.2 allows us to rewrite L ( ρ ) = 1 2 E x ∼ S d -1 ( f ρ ( x ) -y ) 2 as 1 2 ∑ ∞ k =0 ( ˆ σ k,d E u ∼ ρ [ P k,d ( w )] -ˆ h k,d ) 2 . The following lemma states that the gradient from the two formulas for L ( ρ ) are the same when ρ is rotationally invariant.

Lemma D.1. Let ρ be a distribution on S d -1 and let ν be a distribution on [ -1 , 1] . Define

If ρ is rotationally invariant as in Definition 4.1 and ν is the marginal distribution of w for u = ( w,z ) ∼ ρ , then for any particle u = ( w,z ) , we have

where grad w L ( ρ ) denotes the first coordinate of grad u L ( ρ ) .

Proof of Lemma D.1. For any particle u = ( w,z ) , we write the first coordinate of ∇ u L ( ρ ) as ∇ w L ( ρ ) , and the vector consisting of the last ( d -1) coordinates as ∇ z L ( ρ ) . Observe that

On the other hand, suppose ρ satisfies the rotational invariance property, and that ν is the marginal distribution of the first coordinate under ρ . Then, we can write ρ in terms of ν -for any w ∈ [ -1 , 1] , the distribution of z conditioned on w is √ 1 -w 2 S d -1 . Thus, since L ( ρ ) = F ( ν ) by Lemma 4.2, we have by the chain rule that

Moreover, for any fixed w ∈ [ -1 , 1] , the quantity ⟨∇ z L ( ρ ) , z ⟩ does not depend on z , since by Lemma D.2, for any rotation matrix A and z ′ = Az , we have ⟨∇ z ′ L ( ρ ) , z ′ ⟩ = ⟨ A · ∇ z L ( ρ ) , Az ⟩ = ⟨∇ z L ( ρ ) , z ⟩ . Thus, for any w ∈ [ -1 , 1] and z ∈ √ 1 -w 2 S d -2 , we can write

as desired.

In the proof of the above lemma, we used the following fact which we now prove: when ρ is rotationally invariant, ∇ u L ( ρ ) is rotationally equivariant as a function of u , i.e. if u is rotated by a rotation matrix A which only modifies the last ( d -1) coordinates, then ∇ u L ( ρ ) is rotated by A as well.

Lemma D.2 (Gradient is Rotationally Equivariant) . Let ρ be a rotationally invariant distribution as in Definition 4.1. Let A ∈ R d × d be a rotation matrix with Ae 1 = e 1 , i.e. A only modifies the last ( d -1) coordinates. Let u ∈ S d -1 and u ′ = Au . Then, ∇ u ′ L ( ρ ) = A · ∇ u L ( ρ ) , and grad u ′ L ( ρ ) = A · grad u L ( ρ ) .

Proof. The proof is by a straightforward calculation:

(By rotational invariance of uniform distribution on S d -1 )

(B.c. A is a rotation matrix so A ⊤ A = I d × d )

By the rotational invariance of ρ , we have f ρ ( Ax ) = f ρ ( x ) , meaning

This proves the first statement of the lemma. The second statement also follows from a similar calculation:

as desired.

As a corollary, we obtain our first formula for the dynamics of the 1-dimensional particles:

Lemma D.3 (One-Dimensional Velocity) . Suppose we are in the setting of Theorem 3.3 and that ρ t is rotationally invariant. Then, for any particle u t (where we omit the initialization χ for convenience), if w t is the first coordinate of u t , then

Proof of Lemma D.3. This follows directly from Lemma D.1, and because the activation σ and the target h only have nonzero coefficients ˆ σ k,d and ˆ h k,d for k = 0 , . . . , 4 .

We now complete the proof of Lemma 4.3, i.e. that ρ t remains rotationally invariant and symmetric if it is defined according to the population projected gradient flow as in Eq. (3.6) and Eq. (3.7). To prove that ρ t remains symmetric, we make use of the above formula for the 1-dimensional dynamics.

Proof of Lemma 4.3. First, we show the rotational invariance property. Assume inductively that at a time t , ρ t satisfies the desired rotational invariance property - it holds at initialization since ρ 0 is the uniform distribution on S d -1 . Consider a particle u t and let A ∈ R d × d be a rotation matrix which only modifies the last ( d -1) coordinates. Let u ′ t = Au t be another particle. Then, it follows from Lemma D.2 that

Thus, for a fixed w ∈ [ -1 , 1] , if z ∼ √ 1 -w 2 S d -2 , then the distribution of dz dt is uniform on c S d -2 for a fixed scalar c that depends on w . Therefore, ρ will remain rotationally invariant.

Next, we prove that ρ t is symmetric. First observe that ρ t is symmetric at initialization, since it is the uniform distribution on S d -1 . Now, suppose that ρ t is symmetric at time t . The polynomial P k,d is odd if k is odd and even if k is even - this can be seen from induction on Eq. (C.2) and Eq. (C.3). Thus, E u ∼ ρ t [ P k,d ( w )] = 0 for odd k by our assumption that ρ t is symmetric. Additionally, recall that we defined h so that ˆ σ 0 ,d = ˆ h 0 ,d and ˆ h k,d = 0 for odd k . Thus, by our definition of D 2 ,t and D 4 ,t in Section 4,

by Lemma D.3. Since P 2 ,d and P 4 ,d are even polynomials, we know P 2 ,d ′ and P 4 ,d ′ are odd polynomials, and thus dw t dt is an odd polynomial in w t . In other words, if w t and w ′ t are the first

coordinates of two particles u t and u ′ t respectively, such that w ′ t = -w t , then dw ′ t dt = -dw t dt . Thus, the distribution of dw t dt is symmetric, meaning that ρ t will remain symmetric.

In the rest of the proofs for the infinite-width population dynamics, we make use of the symmetry property. In particular, many of the lemmas are stated for w > 0 or ι > 0 - however, the generalizations to w < 0 or ι < 0 also hold. We use this symmetry implicitly throughout the proofs in this section.

In Lemma 4.4, we are simply rewriting our formula for the 1-dimensional dynamics in a more convenient form. The following proof shows the detailed calculation.

Proof of Lemma 4.4. By Lemma D.3 and the fact that only the second and fourth order terms are nonzero (which follows from Lemma 4.3, the fact that P k,d is an odd polynomial for odd k (Eq. (C.2) and Eq. (C.3)) and because ˆ h 1 ,d = ˆ h 3 ,d = 0 ), the velocity of w is equal to the following:

where we use grad w L ( ρ ) to denote the first coordinate of grad u L ( ρ ) . For convenience, during the rest of the proof of this lemma, let η d (1) = d d -1 , η d (2) = d 2 +6 d +8 d 2 -1 and η d (3) = 6 d +12 d 2 -1 . Then, by Eq. (C.4),

Here we have chosen λ d (1) = 2ˆ σ 2 2 ,d ( η d (1) -1) D 2 ( t ) -2ˆ σ 2 4 ,d D 4 ( t ) η d (3) and λ d (3) = 4ˆ σ 2 4 ,d D 4 ( t )( η d (2) -1) . Observe that | λ d (1) | , | λ d (3) | ≤ O ( ˆ σ 2 2 ,d | D 2 ,t | +ˆ σ 2 4 ,d | D 4 ,t | d ) , as desired.

D.3 Analysis of Phase 1

Recall the formal definition of Phase 1 in Definition 4.5. Intuitively, during Phase 1, the particles are all growing at a similar rate - because the particles w are all less than w max = 1 log d in magnitude, the cubic term involving w 3 in the velocity of w does not have a significant effect on the growth of w . We formalize this in the following lemma.

Lemma D.4 (Uniform Growth Lemma) . Suppose we are in the setting of Theorem 3.3. Consider a time interval [ t 0 , t 1 ] ∈ [0 , T 2 ] in Phase 1 or 2. We are interested in quantifying the (relative) growth rate of the neurons indexed (or initialized) at ι 1 and ι 2 (where 0 < ι 1 < ι 2 ), defined as

Let δ ∈ (0 , 1) such that 1 √ d ≲ δ ≤ 1 log F 2 -here, we implicitly assume that F 2 ≤ e d . We also assume that δ ≤ γ 2 2 16 . Further assume that w t 1 ( ι 2 ) = δ . Finally assume that P ( | ι | ≥ | ι 2 | ) ≤ γ 2 2 16 . Then,

and additionally,

Finally,

We note that F 1 and F 2 are both at least 1 .

Note that this lemma also holds for ι 2 < ι 1 < 0 , and with w t 1 ( ι 2 ) = -δ , by the symmetry of ρ t .

Proof of Lemma D.4. In this proof, we define τ = √ γ 2 -then, by Assumption 3.2, we can write γ 4 = β · τ 4 where 1 . 1 ≤ β and β is less than some universal constant c 1 . Suppose at time t 1 , w t 1 ( ι 2 ) = δ . One useful observation, which we use in the rest of this proof, is that for all t ≤ t 1 , w t ( ι 2 ) ≤ δ , since for w ≳ 1 √ d , we have dw dt ≥ 0 by Lemma 4.4 (here we used the assumption that δ ≳ 1 √ d ). Thus,

Thus, writing γ 2 = τ 2 , we have

where the last inequality is by Assumption 3.2 b.c. d ≥ c 3 and γ 4 ≲ γ 2 . By a similar argument,

Thus, for any time t ∈ [ t 0 , t 1 ] , for any ι ∈ [ ι 1 , ι 2 ] , by Lemma 4.4 (and because w t ( ι ) ≤ δ for any ι ∈ [ ι 1 , ι 2 ] by Proposition D.37), we have

In summary,

Thus,

and w t ( ι 2 ) grows by a factor

Rearranging gives

Finally, observe that F O ( δ ) 2 ≤ O (1) , since by the assumption that δ ≤ 1 log F 2 ,

and for any x > 0 , x 1 ln x = e , meaning that F δ 2 ≤ O (1) . This implies that F 1 ≳ F 2 . From the above calculations, we can also obtain

as desired. Similarly, we can obtain

Finally,

as desired. Finally, we show the running time bound. Since δ ≤ γ 2 2 16 and γ 2 is less than a sufficiently small universal constant by Assumption 3.2, for any time t during this interval, v ( w t ( ι 2 )) ≳ ˆ σ 2 2 ,d D 2 ( t ) w t ( ι 2 ) ≳ ˆ σ 2 2 ,d τ 2 w t ( ι 2 ) (where the second inequality is by Eq. (D.50)). Thus, by Gronwall's inequality (Fact I.13), the time elapsed when w t ( ι 2 ) grows by a factor of F 2 is at most 1 ˆ σ 2 2 ,d τ 2 log F 2 .

D.3.1 Summary of Phase 1

As a direct consequence of the above lemma, we can show Lemma D.5, which characterizes how fast the particles w t grow during Phase 1.

Lemma D.5 (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 ι ∈ (0 , ι U ) , we have w T 1 ( ι ) ≍ √ d (log d ) 2 ι . Furthermore, T 1 ≲ 1 ˆ σ 2 2 ,d γ 2 log d . Additionally, we have exp ( ∫ T 1 0 2ˆ σ 2 2 ,d | D 2 ,t | dt ) ≍ √ d (log d ) 2 .

Proof of Lemma D.5. To show the lower bound on w T 1 ( ι ) , we apply Lemma D.4, with t 0 = 0 and t 1 = T 1 , with δ = 1 log d , and with ι 1 = ι and ι 2 = ι U. Let us verify that the assumptions of Lemma D.4 hold:

Observe that w T 1 ( ι U ) ι U = 1 / log d log d/ √ d = √ d (log d ) 2 , meaning that log F 2 ≤ log d and 1 log F 2 ≥ δ .

By Assumption 3.2, δ = 1 log d ≤ γ 2 2 16 , since d ≥ c 3 .

The assumption that P ( | ι | ≥ ι U ) ≤ γ 2 2 16 holds by a standard bound on the tail of µ d .

Using the notation of Lemma D.4,

where w max is as defined in Definition 4.5. By Eq. (D.46), we therefore know that

For the bound on the running time, observe that by Eq. (D.48), T 1 ≲ 1 ˆ σ 2 2 ,d γ 2 log √ d (log d ) 2 . The final statement holds because F 2 = √ d (log d ) 2 and from the second statement of Lemma D.4. This completes the proof.

D.4 Analysis of Phase 2

Recall the formal definition of Phase 2 from Definition 4.6. Our analysis in this section has two main goals: (1) to show that the running time of Phase 2 is small, and (2) to show that a majority of the particles become at least 1 log log d in magnitude, up to a constant factor. We will achieve goal (2) using Lemma D.4 - then, we show that goal (1) holds using goal (2). Additionally, goal (2) will be useful for the subsequent phases as well.

First, the following proposition defines a range of particles which will be relevant during Phase 2 and Phase 3. This range of particles is a strict subset of the interval (0 , ι U ) which we considered during Phase 1 - we will show that this new range of particles grows uniformly during the beginning of Phase 2, until the largest of them becomes about 1 log log d in magnitude. Specifically, we define ι L ≜ κ √ d and ι R ≜ 1 κ √ d where κ ≍ 1 log log d . The following proposition states that at initialization, a very large fraction of the ι 's have magnitude between ι L and ι R.

Proposition D.6. Let κ ∈ (0 , 1) , and suppose ι is drawn uniformly at random from µ d . Define ι L = κ √ d and ι R = 1 κ √ d . Then, P ( | ι | ̸∈ [ ι L , ι R ]) ≲ κ .

Proof of Proposition D.6. First, we upper bound the probability that | ι | ≥ ι R. By Markov's inequality,

Now, we upper bound the probability that | ι | ≤ ι L. By Eqs. (1.16) and (1.18) of Atkinson and Han [12], the probability density function of ι is p ( ι ) = Γ( d/ 2) √ π Γ(( d -1) / 2) (1 -ι 2 ) d -3 2 . Thus, we can simply upper bound the density by its maximum value:

as desired.

Next, as a corollary of Lemma D.4, we show that for a portion of Phase 2, the particles initialized at ι ∈ [ ι L , ι R ] grow by a similar factor. Specifically, this is true until the largest of these particles is ξ = 1 2 log log d in magnitude.

Lemma D.7 (Corollary of Lemma D.4) . Suppose we are in the setting of Theorem 3.3. Let T R > T 1 be the first time such that w T R ( ι R ) = ξ = 1 2 log log d . Then, for all ι with | ι | ∈ [ ι L , ι R ] , we have | w T R ( ι ) | ≳ ξκ 2 , and additionally, w T R ( ι ) w T 1 ( ι ) ≍ ξκ (log d ) 2 . Furthermore, T R -T 1 ≲ 1 ˆ σ 2 2 ,d γ 2 log log d .

Proof of Lemma D.7. By Lemma D.5, w T 1 ( ι R ) ≍ √ d (log d ) 2 ι R = 1 κ (log d ) 2 . Let ι L = κ √ d and ι R = 1 κ √ d , where κ = 1 log log d . Now, we apply Lemma D.4, with t 0 = T 1 and t 1 = T R, and with ι 1 = ι for some ι ∈ [ ι L , ι R ] and ι 2 = ι R, and with δ = 1 2 log log d . Note that we assume without loss of generality that ι > 0 -the statement for ι < 0 follows by the symmetry of ρ t . Let us verify that the assumptions hold:

First, δ = 1 2 log log d ≳ 1 √ d . Addditionally, observe that by Lemma D.5, w T 1 ( ι R ) ι R ≍ √ d (log d ) 2 , meaning that w T 1 ( ι R ) ≍ 1 κ (log d ) 2 . Thus, using the notation of Lemma D.4, we have F 2 ≲ 1 / log log d 1 / ( κ (log d ) 2 ) = κ (log d ) 2 log log d , and log F 2 ≤ 2 log log d -2 log log log d + C < 1 δ , as desired. (Here C is a universal constant, and the last inequality holds by Assumption 3.2 since d ≥ c 3 .)

By Assumption 3.2, δ = 1 2 log log d ≲ γ 2 2 16 since d ≥ c 3 .

Additionally, the probability that | ι | ≥ | ι R | is at most κ = 1 log log d up to constant factors by Proposition D.6, and this is clearly at most γ 2 2 16 by Assumption 3.2.

Thus, using the conclusion of Lemma D.5 that w T 1 ( ι ) ι ≳ √ d (log d ) 2 , we obtain

√

as desired. Thus, we obtain the first statement of this lemma. In addition, by Eq. (D.48), and since w T R ( ι R ) = ξ and w T 1 ( ι R ) ≳ 1 κ (log d ) 2 , the running time bound we obtain for this part of Phase 2 is T R -T 1 ≲ 1 ˆ σ 2 2 ,d τ 2 log ξ 1 / ( κ (log d ) 2 ) ≲ 1 ˆ σ 2 2 ,d τ 2 (log log d ) , as desired.

We now prove some auxiliary lemmas. First, we show an upper bound on L ( ρ 0 ) , which is needed in order to later show an upper bound on the time required to have L ( ρ t ) ≤ (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) ϵ 2 . In the process, we also show a reasonably tight bound on the magnitude of D 2 ,t and D 4 ,t during the first part of Phase 2.

Lemma D.8. Suppose we are in the setting of Theorem 3.3 and Lemma D.7, and suppose 0 ≤ t ≤ T R . Then, -3 γ 2 2 ≤ D 2 ,t ≤ -γ 2 2 and -3 γ 4 2 ≤ D 4 ,t ≤ -γ 4 2 . As a corollary, for all t ≥ 0 , L ( ρ t ) ≲ (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) γ 2 2 .

Proof. Since t ≤ T R, w t ( ι R ) ≤ ξ = 1 2 log log d (because for w ≳ 1 √ d , we have dw dt ≥ 0 by Lemma 4.4). By Proposition D.37, for all ι with | ι | ≤ ι R, | w t ( ι ) | ≤ ξ . Thus, by Proposition D.6,

Thus, D 2 ,t = O ( ξ 2 + κ ) -γ 2 , meaning that -3 γ 2 2 ≤ D 2 ≤ -γ 2 2 , and similarly, -3 γ 4 2 ≤ D 4 ≤ -γ 4 2 . The corollary follows from the fact that γ 4 ≲ γ 2 2 by Assumption 3.2.

Next, we show that the particles are in fact increasing in magnitude during Phases 1 and 2. This makes use of our bound on the magnitudes of D 2 ,t and D 4 ,t up to time T R.

Lemma D.9. Suppose we are in the setting of Theorem 3.3, and suppose [ t 0 , t 1 ] is a time interval with 0 ≤ t 0 < t 1 ≤ T 2 . Then, for all ι with | ι | ∈ [ ι L , ι R ] , | w t ( ι ) | is increasing on [ t 0 , t 1 ] .

Proof. Recall that by Lemma 4.4,

and for w ≥ 0 , it suffices to show that P t ( w ) + Q t ( w ) ≤ 0 . For t ≤ T R,

For t ≥ T R, suppose w ≳ ξκ 2 (which holds at time T R for ι ∈ [ ι L , ι R ] by Lemma D.7). Then,

for some universal constant C > 0 . The final expression on the right-hand side is negative - to see this, note that 2 w -w d > 0 if w ≳ ξκ 2 , and 4 w 3 -w d > 0 if w ≳ ξκ 2 since d ≥ c 3 by Assumption 3.2. This completes the proof of the lemma when ι ≥ 0 . For ι < 0 , the lemma holds by the symmetry of ρ t .

D.4.1 Summary of Phase 2

Finally, we show Lemma D.10 which concludes the analysis of Phase 2. For times t ≥ T R within Phase 2, the particles with | ι | ∈ [ ι L , ι R ] satisfy | w t ( ι ) | ≳ ξκ 2 . Using this observation, and the fact that D 2 ,t and D 4 ,t have the same sign (meaning that the two terms of P t ( w ) = 2ˆ σ 2 2 ,d D 2 ,t w +4ˆ σ 2 4 ,d D 4 ,t w 3 do not cancel) we show that the velocity v ( w ) is large on average. By Lemma D.39, this allows us to show that the loss decreases quickly at any time t ≥ T R in Phase 2, meaning that either the loss goes below (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) ϵ 2 quickly, or we exit Phase 2 quickly.

Lemma D.10 (Phase 2 Summary) . 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 ≲ (log log d ) 18 (ˆ σ 2 2 ,d +ˆ σ 2 4 ,d ) log ( γ 2 ϵ ) or T ∗ ,ϵ -T 1 ≲ (log log d ) 18 (ˆ σ 2 2 ,d +ˆ σ 2 4 ,d ) log ( γ 2 ϵ ) .

Proof of Lemma D.10. First, let T R be as defined in Lemma D.7. By Lemma D.7, T R -T 1 ≲ 1 ˆ σ 2 2 ,d γ 2 log log d . Furthermore, at time T R, for all ι such that | ι | ∈ [ ι L , ι R ] , | w t ( ι ) | ≳ ξκ 2 .

Our proof strategy is now to show that v ( w ) is large on average, and then apply Lemma D.39. At any time t ∈ [ T R , T 2 ] , since D 2 ,t ≤ 0 , Proposition D.38 implies that E w ∼ ρ t [ w 2 ] ≤ γ 2 + O ( 1 d ) . Thus, by Markov's inequality, and because d ≥ c 3 (where c 3 is defined in Assumption 3.2),

For all ι with | ι | ∈ [ ι L , ι R ] , | w t ( ι ) | is increasing as a function of t on [0 , T 2 ] , by Lemma D.9. Thus, by Proposition D.6 and a union bound, P w ∼ ρ t ( ξκ 2 ≤ | w t ( ι ) | ≤ 1 / 2) ≥ 1 -O ( κ ) -5 γ 2 ≥ 1 2 for all t ≥ T R and all ι with | ι | ∈ [ ι L , ι R ] . (Note that 1 -O ( κ ) -5 γ 2 ≥ 1 2 holds because γ 2 is less than a

sufficiently small constant by Assumption 3.2.) Finally, for all w satisfying ξκ 2 ≤ | w | ≤ 1 2 , using Lemma 4.4 we can compute that

(B.c. d is sufficiently large (Assumption 3.2))

Thus, since L ( ρ T R ) ≲ (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) γ 2 (by Lemma D.8), by Gronwall's inequality (Fact I.13), we know that for t ∈ [ T R , T 2 ] ,

Thus, if L ( ρ t ) ≳ (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) ϵ 2 , this implies that

and rearranging gives t -T R ≤ 1 (ˆ σ 2 2 ,d +ˆ σ 2 4 ,d ) ξ 6 κ 12 log( γ 2 ϵ 2 ) . Thus, either T 2 -T R ≤ 1 (ˆ σ 2 2 ,d +ˆ σ 2 4 ,d ) ξ 6 κ 12 log( γ 2 ϵ 2 ) , or there exists a time t ∈ [ T R , T 2 ] such that L ( ρ t ) ≲ (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) ϵ 2 , and such that t -T R ≤ 1 (ˆ σ 2 2 ,d +ˆ σ 2 4 ,d ) ξ 6 κ 12 log( γ 2 ϵ ) . This completes the proof of the lemma.

D.5 Phase 3, Case 1

Phase 3, Case 1 is the case where D 2 ,T 2 = 0 and D 4 ,T 2 < 0 , i.e. the case where D 2 ,t reaches 0 first. As we later show in Lemma D.19 and Lemma D.20, if this occurs, then for all t ≥ T 2 , it will be the case that D 2 ,t ≥ 0 and D 4 ,t ≤ 0 . We now outline our analysis of this case - the goal of our analysis is to show that either | D 2 ,t | ≤ ϵ or | D 4 ,t | ≤ ϵ holds within 1 ˆ σ 2 2 ,d +ˆ σ 2 4 ,d · poly ( log log d ϵ ) time after the start of Phase 3, Case 1.

Proof Strategy for Phase 3, Case 1 Our overall proof strategy in this section is to show that | v ( w ) | is large for a large portion of particles w , and then to apply Lemma D.39 to show that the loss decreases quickly. Recall that

and intuitively, we can ignore Q t ( w ) due to the factor of 1 d in its coefficients. Thus, we can write

Since D 4 ,t < 0 and D 2 ,t > 0 , we can further rewrite the above as

where r = √ -ˆ σ 2 2 ,d D 2 ,t 2ˆ σ 2 4 ,d D 4 ,t is the positive root of the last factor (2ˆ σ 2 2 ,d D 2 ,t + 4ˆ σ 2 4 ,d D 4 ,t w 2 ) . Thus, to show that the velocity v T ( w ) is large, we want to show that there is a large fraction of particles w ∈ [0 , 1] simultaneously satisfying the following: (1) w is far from 1 , (2) w is far from 0 , and (3) w is far from r . Indeed, in Lemma D.17, we show that out of the particles w ∈ [0 , 1] , at least a γ 2 2 fraction of these particles is in [ γ 3 / 4 2 , 1 -γ 1 / 2 2 ] , enabling us to bound the distance of these particles w from 0 and 1 .

Thus, to show that v t ( w ) is large for these particles, it suffices to show that | w -r | is large on average for the particles w ∈ [ γ 3 / 4 2 , 1 -γ 1 / 2 2 ] , i.e. the conditional expectation E w ∼ ρ t [( w -r ) 2 | w ∈ [ γ 3 / 4 2 , 1 -γ 1 / 2 2 ]] is large. In turn, this is at least the conditional variance of w conditioned on w being in [ γ 3 / 4 2 , 1 -γ 1 / 2 2 ] , so it suffices to show that Var( w | w ∈ [ γ 3 / 4 2 , 1 -γ 1 / 2 2 ]) is large, and by Proposition D.40 it suffices to show that E [( w -w ′ ) 2 | w,w ′ ∈ [ γ 3 / 4 2 , 1 -γ 1 / 2 2 ]] is large.

To show that the quantity E [( w -w ′ ) 2 | w,w ′ ∈ [ γ 3 / 4 2 , 1 -γ 1 / 2 2 ]] is large, intuitively we would like to show that during Phase 3, Case 1, w and w ′ are getting farther apart. However, this is not true. If we write the velocity as v t ( w ) = (1 -w 2 ) · w · (4ˆ σ 2 4 ,d | D 4 ,t | w 2 -2ˆ σ 2 2 ,d | D 2 ,t | ) , then the last factor is indeed increasing as a function of w ∈ [0 , 1] , and thus 'pushes' the different particles w,w ′ further apart. This does not work as a formal argument, because as two particles w,w ′ become closer to 0 or 1 , their distance | w -w ′ | will decrease again due to the factors w and (1 -w 2 ) in the velocity. However, we are only concerned with particles w ∈ [ γ 3 / 4 2 , 1 -γ 1 / 2 2 ] -wemust make precise the intuition that the behavior of the particles very close to 0 or 1 does not matter.

To make this proof strategy precise, we define the potential function Φ( w ) = log ( w √ 1 -w 2 ) in Definition D.11 and show that | Φ( w ) -Φ( w ′ ) | is increasing in Phase 3, Case 1. As shown in Lemma D.12, we have d dt Φ( w ) ≈ -2ˆ σ 2 2 ,d D 2 ,t -4ˆ σ 2 4 ,d D 4 ,t w 2 -this is simply because d Φ( w ) dw = 1 w (1 -w 2 ) . In other words, d dt Φ( w ) is exactly v ( w ) , but without the factors w and (1 -w 2 ) -we have selected this potential function Φ( w ) specifically to eliminate these factors. Using this observation, it is easy to show that | Φ( w ) -Φ( w ′ ) | is increasing and thus obtain a lower bound on | Φ( w ) -Φ( w ′ ) | for a large fraction of pairs w,w ′ . Finally, to convert this to a lower bound on | w -w ′ | , we use the fact that Φ is Lipschitz on the interval [ γ 3 / 4 2 , 1 -γ 1 / 2 2 ] (Lemma D.14), which gives us | w -w ′ | ≥ 1 γ O (1) 2 | Φ( w ) -Φ( w ′ ) | .

We now repeat the definition of the potential function:

Definition D.11 (Potential Function Φ ) . We define Φ : [0 , 1] → ( -∞ , ∞ ) by Φ( w ) = log ( w √ 1 -w 2 ) .

We note that d Φ( w ) dw = 1 w (1 -w 2 ) . Next, we note the following useful fact about the potential function (we note that this does not require any assumption that we are in Phase 3, Case 1, and just follows from algebraic manipulations):

Lemma D.12. In the setting of Theorem 3.3,

where λ d (1) and λ d (3) are as in the statement of Lemma 4.4.

Proof of Lemma D.12. For any w t ,

Thus, writing dw t dt = -(1 -w 2 ) · ( P t ( w ) + Q t ( w )) where P t and Q t are defined in Lemma 4.4, we obtain

where λ d (1) and λ d (3) are as in the statement of Lemma 4.4.

Intuitively, when D 2 > 0 and D 4 < 0 , the main part of the velocity is P t ( w ) = -2ˆ σ 2 2 ,d D 2 ,t w -4ˆ σ 2 4 ,d D 4 ,t w 3 which is an upward sloping cubic polynomial - thus, the particles w move away from the positive root of this polynomial, and the distance between any pair w,w ′ of particles is increasing provided that w and w ′ are not too close to 0 or 1 . In the following few lemmas, we make this intuition formal using the potential function Φ . We first show that for two particles w,w ′ , the distance between Φ( w ) and Φ( w ′ ) is increasing as long as D 4 ,t ≤ 0 . Note that this condition holds not only during Phase 3, Case 1, but also during Phases 1 and 2.

Lemma D.13 ( Φ( w ) and Φ( w ′ ) Moving Apart) . In the setting of Theorem 3.3, suppose D 4 ,t ≤ 0 for all t in some time interval [ t 0 , t 1 ] , and let ι 1 , ι 2 > 0 . Then,

is increasing on the interval [ t 0 , t 1 ] .

Proof of Lemma D.13. Suppose ι 1 , ι 2 ∈ [0 , 1] such that ι 1 < ι 2 -by Proposition D.37, we will always have w t ( ι 1 ) ≤ w t ( ι 2 ) . For convenience, in the rest of this proof, we will write w 1 ,t = w t ( ι 1 ) and w 2 ,t = w t ( ι 2 ) . Since Φ is an increasing function of w , it will also always be the case that Φ( w 2 ,t ) ≥ Φ( w 1 ,t ) . Thus, to show that | Φ( w 2 ,t ) -Φ( w 1 ,t ) | is increasing, it suffices to show that Φ( w 2 ,t ) -Φ( w 1 ,t ) is increasing. By Lemma D.12,

where λ d (1) and λ d (3) are as in Lemma 4.4. Recall from the statement of Lemma 4.4 that | λ d (3) | ≲ ˆ σ 2 4 ,d | D 4 ,t | · 1 d . Therefore,

and the right hand side is nonnegative since w 2 ,t ≥ w 1 ,t ≥ 0 . Thus, Φ( w 2 ,t ) -Φ( w 1 ,t ) is increasing, as desired.

Next, we prove a helpful lemma showing that the potential function is bi-Lipschitz. This is useful both when we show that | Φ( w ) -Φ( w ′ ) | is initially large (as it allows us to leverage an initial lower bound on | w -w ′ | ) and when we show that | w -w ′ | is large later during Phase 3, Case 1 (as it allows us to leverage the lower bound we obtain on | Φ( w ) -Φ( w ′ ) | ).

Lemma D.14. For all w,w ′ ∈ [0 , 1] , | Φ( w ) -Φ( w ′ ) | ≥ | w -w ′ | . Additionally, for η < 1 2 and w,w ′ ∈ [ η, 1 -η ] , | Φ( w ) -Φ( w ′ ) | ≤ 1 η 2 | w -w ′ | .

Proof of Lemma D.14. First, for all w ∈ [0 , 1] , d Φ( w ) dw = 1 w (1 -w 2 ) > 1 , and the first statement of the lemma follows. For the second statement of the lemma, observe that for w ∈ [ η, 1 -η ] ,

where the last equality is because η > η 2 (so 2 η -η 2 > η ). Thus, Φ is 1 η 2 -Lipschitz on the interval [ η, 1 -η ] .

We next need to show that the potential initially has a large value, at the beginning of Phase 3, Case 1. We show this in the next two lemmas, by showing that for any two particles w,w ′ , the distance between them grows by a large factor during Phase 2 (specifically, by time T R) and then using this to obtain a lower bound on | Φ( w ) -Φ( w ′ ) | by the bi-Lipschitzness of Φ .

Lemma D.15. In the setting of Theorem 3.3, let T R be as defined in Lemma D.7. Then, for ι, ι ′ ∈ [ ι L , ι R ] ,

Proof of Lemma D.15. Let ι, ι ′ ∈ [ ι L , ι R ] , with ι < ι ′ . For convenience, we will write w t = w t ( ι ) and w ′ t = w t ( ι ′ ) . Our goal is to obtain a lower bound on the factor by which w ′ t -w t grows by time T R. First, define P t and Q t as in Lemma 4.4, and observe that

We first obtain a lower bound on the first term -( P t ( w ′ t ) -P t ( w t ) ) , and then show that the other terms are lower-order terms as long as t ≤ T R. For t ∈ [0 , T R ] ,

(B.c. omitted term is nonnegative w t > w t by Proposition D.37)

Next, let us show that the remaining terms in Eq. (D.107) are lower-order terms which will not decrease the growth rate of w ′ t -w t too much. First let us deal with the absolute value of the second

term in Eq. (D.107). For convenience, let w R,t = w t ( ι R ) . Then,

where the last inequality is also because ˆ σ 2 4 ,d ≲ ˆ σ 2 2 ,d , and γ 4 ≲ γ 2 2 (Assumption 3.2) and | D 2 | ≥ γ 2 2 and | D 4 | ≤ 3 γ 4 2 for t ≤ T R (Lemma D.8), meaning that the first term in Eq. (D.113) dominates up to universal constant factors. Next, let us deal with the absolute value of the third term in Eq. (D.107). By Lemma I.16,

Using the same argument with Lemma I.16, we can bound the fourth term in Eq. (D.107):

Combining the bounds we obtained for all the terms of Eq. (D.107), we obtain

and therefore,

First, let us consider how much w t -w t grows during Phase 1. At this point, we apply Lemma D.4 with ι 2 = ι U and 0 < ι 1 < ι U, with δ = 1 log d . Let us first verify that the assumptions of Lemma D.4 are satisfied. By Lemma D.5, F 2 ≍ √ d 2 , using the notation of Lemma D.4. Thus, log F 2 ≤ log d , meaning δ ≤ 1 log F 2 . By Assumption 3.2, δ ≤ 2 16 , and by Proposition I.5, P ( | ι | > ι U ) ≤ 1 log d ≤ 2 16 clearly holds. Thus, the assumptions of Lemma D.4 are satisfied. As a consequence of Lemma D.4 (specifically Eq. (D.47)) we obtain

′ (log d ) γ 2 γ 2

For t ≤ T 1 , we have w 2 R,t ≤ δ 2 ≤ O ( 1 log d ) by Definition 4.5 (and because w t ( ι R ) ≤ w t ( ι U ) by Proposition D.37). Thus,

Additionally, we apply Lemma D.4 with ι 2 = ι R (and an arbitrary ι 1 ∈ [ ι L , ι R ] ) with δ = ξ = 1 2 log log d , and t 0 = T 1 and t 1 being the time T R such that w T R ( ι R ) = ξ . We verify that the conditions of Lemma D.4 hold. Using the notation of Lemma D.4, F 2 = w T R ( ι R ) w T 1 ( ι R ) = ξκ (log d ) 2 by Lemma D.7, and thus log F 2 ≤ 1 δ . Again by Assumption 3.2, δ ≤ γ 2 2 16 , and by Proposition D.6, P ( | ι | > ι R ) ≤ O ( κ ) ≤ γ 2 2 16 . Thus, we apply Lemma D.4 to obtain

For t ≤ T R, w 2 R,t ≤ ξ 2 ≤ δ 2 , meaning that

In summary, since w ′ t -w t grows by √ d (log d ) 2 (up to a constant factor) from time 0 to time T 1 , and by ξκ (log d ) 2 (up to a constant factor) from time T 1 to time T R, this means

from time 0 to time T R, as desired.

We now use the previous lemma to obtain a lower bound on | Φ( w ) -Φ( w ′ ) | by using the biLipschitzness of Φ .

Lemma D.16 (Initial Large Distance Between w and w ′ ) . Suppose we are in the setting of Theorem 3.3. Let T R be as in the statement of Lemma D.7. If ι, ι ′ are sampled independently from ρ 0 , then with probability at least 1 -O ( ϵ ) ,

Proof of Lemma D.16. By Eqs. (1.16) and (1.18) of Atkinson and Han [12], the probability density function of ι is p ( ι ) = Γ( d/ 2) √ π Γ(( d -1) / 2) (1 -ι 2 ) d -3 2 . Thus, for ι ∈ [ ι L , ι R ] , the conditional density function p ( ι | ι ∈ [ ι L , ι R ]) can be upper bounded as

In particular, for any interval I of length at most κ 2 √ d ,

and from this it follows that

Next, suppose ι, ι ′ ∈ [ ι L , ι R ] with ι ′ > ι , and | ι -ι ′ | ≥ κ 2 √ d . Then, by Lemma D.15, if T R is the time T mentioned in the statement of Lemma D.7, then

Suppose we sample ι, ι ′ independently from ρ 0 . By Proposition D.6, with probability at least 1 -O ( κ ) , ι, ι ′ ∈ [ ι L , ι R ] . Thus, by Eq. (D.125), the probability that ι, ι ′ ∈ [ ι L , ι R ] and | ι -ι ′ | ≥ κ 2 √ d is at least 1 -O ( κ ) . In summary, if we sample ι, ι ′ independently from ρ 0 , then with probability 1 -O ( κ ) ,

Next, we show that during Phase 3, Case 1, at any time, there is a significant fraction of particles whose distance from 0 or 1 can be bounded below.

Lemma D.17. In the setting of Theorem 3.3, suppose D 4 ,t ≤ 0 and D 2 ,t ≥ 0 . Then,

Proof of Lemma D.17. For convenience, we define τ = √ γ 2 and β ≥ 1 . 1 such that γ 4 = βτ 4 . (Here, β and τ are well-defined by Assumption 3.2.) First, since D 2 ,t ≥ 0 during Phase 3 Case 1, we know that

and since D 4 ,t ≤ 0 during Phase 3 Case 1, we know that

Rearranging using P 2 ,d ( t ) = d d -1 t 2 -1 d -1 , we obtain

Similarly, rearranging using P 4 ,d ( t ) = d 2 +6 d +8 d 2 -1 t 4 -6 d +12 d 2 -1 t 2 + 3 d 2 -1 , we obtain

and rearranging gives or

because w ≤ 1 and we can assume βτ 4 ≤ 1 due to Assumption 3.2. Suppose that with probability more than 1 -τ 2 2 under ρ t , w ̸∈ [ τ 3 / 2 , 1 -τ ] . Then,

which by Eq. (D.130) implies that

where the last inequality is because τ ≤ 1 . On the other hand,

By Eq. (D.132), this is a contradiction, since it implies that

where the second inequality is because β ≥ 1 . 1 and d is sufficiently large (see Assumption 3.2). This implies that 128 βτ 2 ≥ 1 , which contradicts Assumption 3.2 (since γ 2 is chosen to be sufficiently small). Thus, our original assumption that P w ∼ ρ t ( w ∈ [ τ 3 / 2 , 1 -τ ]) ≤ τ 2 2 is incorrect.

The purpose of the next three lemmas is to show that if Phase 3, Case 1 occurs, i.e. if D 2 ,T 2 = 0 and D 4 ,T 2 < 0 , then for all t ≥ T 2 , we have D 2 ,t ≥ 0 and D 4 ,t ≤ 0 . First, the following lemma states that if D 4 ,t is much smaller than D 2 ,t in absolute value, and D 4 ,t < 0 and D 2 ,t > 0 , then the velocity of the particles w ≥ 0 will be significantly negative.

Lemma D.18. Suppose we are in the setting of Theorem 3.3, and suppose D 4 ,T 2 = 0 and t ≥ T 2 such that D 4 ,t ≤ 0 and D 2 ,t ≥ 0 . If | D 4 ,t | ≤ ξ | D 2 ,t | where ξ = 1 log log d , then for all w ∈ [0 , 1] ,

In particular, for w ∈ [ γ 3 / 4 2 , 1 -γ 1 / 2 2 ] ,

(B.c. τ ≤ 1 / 2 (see Assumption 3.2))

(B.c. τ ≤ 1 / 4 (by Assumption 3.2, τ is sufficiently small))

Proof. Suppose | D 4 ,t | ≤ ξ | D 2 ,t | where ξ = 1 log log d . Let P t and Q t be as in Lemma 4.4. Then, for any w ∈ [0 , 1] ,

Thus, for w ∈ [0 , 1] ,

In particular, for w ∈ [ γ 3 / 4 2 , 1 -γ 1 / 2 2 ] ,

as desired.

The next lemma essentially implies that if at some point D 4 ,t ≤ 0 and D 2 ,t > 0 , then it will never be the case that D 4 ,t > 0 and D 2 ,t > 0 .

Lemma D.19. Suppose we are in the setting of Theorem 3.3, and suppose D 4 ,t = 0 and D 2 ,t ≥ 0 for some t ≥ 0 . Then, d ds D 4 ,s ∣ ∣ ∣ s = t ≤ 0 .

Proof. We can calculate that

Recall from Eq. (C.4) that

meaning that

Observe that for t ≳ 1 √ d , P 4 ,d ′ ( t ) > 0 and for t ≲ 1 √ d , it may be the case that P 4 ,d ′ ( t ) < 0 . In particular,

meaning that d dt D 4 ,t < 0 , as desired.

The next lemma essentially implies that if at some point D 4 ,t ≤ 0 and D 2 ,t > 0 , it will never be the case that D 4 ,t ≤ 0 and D 2 ,t < 0 .

Lemma D.20. Suppose we are in the setting of Theorem 3.3, and suppose D 4 ,t ≤ 0 and D 2 ,t = 0 , for some t ≥ 0 . Then, d ds D 2 ,s ∣ ∣ ∣ s = t ≥ 0 .

Proof. Recall from Eq. (C.4) that P 2 ,d ( t ) = d d -1 t 2 -1 d -1 meaning that P 2 ,d ′ ( t ) = 2 d d -1 t . Thus,

If D 4 ,t ≤ 0 and D 2 ,t = 0 , then for any particle w ∈ [0 , 1] , we can write

where the last equality is because D 4 ,t ≤ 0 . Observe that v t ( w ) ≥ 0 for w ≳ 1 √ d , while it may be the case that v t ( w ) ≤ 0 for w ≲ 1 √ d . Thus,

as desired.

Putting all of the previous lemmas together, we obtain the following invariant: a lower bound for | Φ( w ) -Φ( w ′ ) | , for a large fraction of w,w ′ , which holds throughout Phase 3, Case 1. Note that in the proof of this invariant, we need to make use of the fact that once we enter the case where D 2 ,t > 0 and D 4 ,t < 0 , we cannot leave this case. If this is not true, then we cannot use the fact that | Φ( w t ) -Φ( w ′ t ) | is increasing as a function of t (Lemma D.13) since if D 4 ,t > 0 at any point, then | Φ( w t ) -Φ( w ′ t ) | could decrease, and we would not have control over this decrease.

Lemma D.21 (Phase 3, Case 1 Invariant) . Suppose we are in the setting of Theorem 3.3. Let T R be as defined in the statement of Lemma D.7, and assume that D 2 ,T 2 = 0 and D 4 ,T 2 < 0 . Then, for all t ≥ T R (in particular, for all times t during Phase 3 Case 1), if ι, ι ′ are sampled independently from ρ 0 ,

with probability at least 1 -O ( κ ) .

Proof of Lemma D.21. By Lemma D.19 and Lemma D.20, if D 2 ,T 2 = 0 and D 4 ,T 2 < 0 , then for all t ≥ T 2 , we will have D 2 ,t ≥ 0 and D 4 ,t ≤ 0 . Thus, the lemma follows from Lemma D.13 and Lemma D.16, as well as the fact that D 4 ,t ≤ 0 for t ∈ [ T R , T 2 ] (meaning that the potential difference is increasing in the time interval [ T R , T 2 ] ).

Finally, we show Lemma D.22 which gives a running time bound for Phase 3, Case 1.

Lemma D.22 (Phase 3, Case 1 Summary) . 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 ≥ T 2 such that L ( ρ t ) ≥ 1 2 (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) ϵ 2 is at most O ( 1 ˆ σ 2 4 ,d γ 8 2 ξ 4 κ 6 log( γ 2 ϵ )) .

Proof of Lemma D.22. Suppose t ≥ T 2 and assume that D 2 ,T 2 = 0 and D 4 ,t < 0 (i.e. Phase 3, Case 1 occurs). First, suppose | D 4 ,t | ≤ ξ | D 2 ,t | where ξ = 1 log log d . Then, by Lemma D.18, if | w | ∈ [ γ 3 / 4 2 , 1 -γ 1 / 2 2 ] , we have

(note that the lemma is stated for w ∈ [ γ 3 / 4 2 , 1 -γ 1 / 4 2 ] , but holds for | w | ∈ [ γ 3 / 4 2 , 1 -γ 1 / 4 2 ] since v t ( w ) is an odd function). Therefore,

Thus, by Lemma D.39, for any time t during Phase 3, Case 1 such that | D 4 ,t | ≤ ξ | D 2 ,t | , we have

Let T ∗ ,ϵ be as in Theorem 3.3, and define

Then, by Eq. (D.168), and we can write

where r = √ -ˆ σ 2 2 ,d D 2 ,t 2ˆ σ 2 4 ,d D 4 ,t is the positive root of P t . For | w | ∈ [ γ 3 / 4 2 , 1 -γ 1 / 2 2 ] , we have

Thus, by Gronwall's inequality (Fact I.13),

Since L ( ρ 0 ) ≲ (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) γ 2 2 by Lemma D.8 and L ( ρ T ∗ ,ϵ ) ≳ (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) ϵ 2 , rearranging gives

Now, suppose t is such that | D 4 ,t | ≥ ξ | D 2 ,t | . Let us lower bound the average velocity for w ∈ [ γ 3 / 4 2 , 1 -γ 1 / 4 2 ] first (then, we can conclude the average velocity is large by Lemma D.17). By Lemma 4.4,

where the last inequality is by the discussion above Eq. (D.177). Now, in order to take the expectation over w ∼ ρ t (conditioned on w ∈ [ γ 3 / 4 2 , 1 -γ 1 / 2 2 ] ), we first compute E w ∼ ρ t [( w -r ) 2 | w ∈ [ γ 3 / 4 2 , 1 -γ 1 / 2 2 ]] :

where the last inequality is by Proposition D.40. By Lemma D.21,

By Lemma D.14, if w,w ′ ∈ [ γ 3 / 4 2 , 1 -γ 1 / 2 2 ] , then

Thus, combining the previous two equations gives

and therefore, by Assumption 3.2, since κ/γ 2 ≲ 1 γ 2 log log d , this is at most 1 2 , since d is sufficiently large. Thus,

where the last inequality is because P w,w ′ ∼ ρ t ( | w -w ′ | < γ 3 / 2 2 ξκ 3 | w,w ′ ∈ [ γ 3 / 4 2 , 1 -γ 1 / 2 2 ]) ≤ 1 2 . By Eq. (D.178), we have

Since P w ∼ ρ t ( w ∈ [ γ 3 / 4 2 , 1 -γ 1 / 2 2 ]) ≥ γ 2 4 by Lemma D.17 and the symmetry of ρ t , this implies that

Thus, by Lemma D.39, dL ( ρ t ) dt ≲ -ˆ σ 2 4 ,d γ 8 2 ξ 4 κ 6 L ( ρ t ) . By a similar argument as for the subcase where | D 4 ,t | ≤ ξ | D 2 ,t | , since L ( ρ 0 ) ≲ ˆ σ 2 2 ,d γ 2 2 (by Lemma D.8), if we define

then

since L ( ρ T ∗ ,ϵ ) ≳ ˆ σ 2 2 ,d ϵ 2 . Combining Eq. (D.173) and Eq. (D.194), we find that the total time spent in Phase 3, Case 1 is at most O ( 1 ˆ σ 2 4 ,d γ 8 2 ξ 4 κ 6 log( γ 2 ϵ )) , as desired.

D.6 Phase 3, Case 2

Next, we analyze Phase 3, Case 2, where D 4 ,T 2 = 0 and D 2 ,T 2 < 0 . In this case, as we will show, for all t ≥ T 2 , we will have D 4 ,t ≥ 0 and D 2 ,t ≤ 0 (Lemma D.28). We will first give an outline of our analysis in this case - our goal is again to show that either | D 2 ,t | ≤ ϵ or | D 4 ,t | ≤ ϵ within 1 ˆ σ 2 2 ,d +ˆ σ 2 4 ,d · poly ( log log d ϵ ) time after the start of Phase 3, Case 2.

Proof Strategy for Phase 3, Case 2 We again show that | v ( w ) is large on average and then use Lemma D.39 to show that the loss decreases quickly. However, our argument to show that | v ( w ) is large on average is significantly different from Phase 3, Case 1. We can approximately write

Letting r = √ ˆ σ 2 2 ,d | D 2 ,t | 2ˆ σ 2 4 ,d | D 4 ,t | be the positive root of P t ( w ) , we obtain

This time, v t ( w ) has a negative slope at r , meaning that the particles are attracted towards r , and we cannot use the same argument as in Phase 3, Case 1.

Instead, the key point of our argument in this case is the following. Recall that D 4 ,t ≥ 0 and D 2 ,t ≤ 0 for all t ≥ T 2 -these roughly imply that E w ∼ ρ t [ w 4 ] ≥ γ 4 and E w ∼ ρ t [ w 2 ] ≤ γ 2 . By Assumption 3.2, we have γ 4 ≥ 1 . 1 γ 2 2 -in other words, E w ∼ ρ t [ w 4 ] ≥ 1 . 1( E w ∼ ρ t [ w 2 ]) 2 . From this observation, we know that a significant fraction of the particles must be far from r -otherwise, if all of the particles are concentrated at r , then we would have E w ∼ ρ t [ w 4 ] ≈ ( E w ∼ ρ t [ w 2 ]) 2 , which would give a contradiction.

One complication in showing that the velocity is large is that, a priori, nearly all of the particles could be close to 0 or 1 -their velocities could be small even if they are far from r . We make use of the potential function Φ defined in the analysis of Phase 3, Case 1 to avoid this issue. In this case, it turns out that for any two particles w,w ′ , their potential difference | Φ( w ) -Φ( w ′ ) | is decreasing (Lemma D.23). We can use this to show that a 1 -O ( γ 4 ) fraction of particles w with initializations in [ ι L , ι M ] (where ι M is defined later) cannot become too close to 0 or 1 in Lemma D.24 and Lemma D.25. Intuitively, this is because if | Φ( w ) -Φ( w ′ ) | is bounded by log B for two particles w,w ′ > 0 , and C is a sufficiently large constant, then w ≤ ( 1 log log d ) C roughly implies that w ′ ≲ B ( 1 log log d ) C , simply by the definition Φ( w ) = log( w √ 1 -w 2 ) together with some algebraic manipulations. Moreover, if all of the particles w ′ in this group are at most B ( 1 log log d ) C , then this implies that D 2 has a very large negative value, and this implies that all particles w have a positive velocity (as argued in Lemma D.25), meaning that the particles w in this group cannot become smaller than B ( 1 log log d ) C . We can show by a similar argument that the particles having initializations in [ ι L , ι M ] cannot become too close to 1 .

As long as the positive root r of P t ( w ) is at least 1 2 , this group of particles will have a large velocity as argued in Lemma D.29, since a significant fraction of these particles must be at most 1 3 . We encounter a potential issue when the positive root r is at most 1 2 : it is possible that the particles initialized in [ ι L , ι M ] get stuck at r . We mentioned above that by the constraints D 2 ,t < 0 and D 4 ,t > 0 , there must always be some mass away from r . However, this mass may only be a γ 5 / 4 4 fraction of the total mass, while the particles initialized at [ ι L , ι M ] only have a 1 -O ( γ 4 ) fraction of the total mass. Thus, once r ≤ 1 2 , it is possible for the mass located far away from r to be disjoint from the set of particles initialized at [ ι L , ι M ] . To resolve this, we argue in Lemma D.31 that the particles larger than w t ( ι M ) (which may be close to 1 ) will move away from 1 at an exponentially fast rate while r ≤ 1 2 . Once these particles, which are larger than w t ( ι M ) , are sufficiently far away from 1 , they will also contribute a significant amount to the velocity. Thus, we can argue as discussed above that there must be a significant fraction of particles away from r due to the constraints D 2 ,t ≤ 0 and D 4 ,t ≥ 0 (Lemma D.32).

We now begin the formal proof. We first show that as long as D 4 ,t ≥ 0 , the potential difference between any two particles w,w ′ is increasing. In particular, this holds during Phase 3, Case 2 (but note that this lemma is no longer applicable during Phase 1 or Phase 2).

Lemma D.23 ( Φ( w ) and Φ( w ′ ) Move Closer in Phase 3 Case 2) . Suppose we are in the setting of Theorem 3.3, and suppose D 4 ,t ≥ 0 for all t in some time interval [ t 0 , t 1 ] . Then, for all ι, ι ′ ,

is decreasing as a function of t on the interval [ t 0 , t 1 ] .

Proof of Lemma D.23. We mostly follow the proof of Lemma D.13. Suppose ι, ι ′ ∈ [0 , 1] such that ι ≤ ι ′ . Then, w t ( ι ′ ) ≥ w t ( ι ) for all times t (by Proposition D.37), and since Φ is an increasing function, Φ( w t ( ι ′ )) ≥ Φ( w t ( ι )) for all times t . Thus, it suffices to show that Φ( w t ( ι ′ )) -Φ( w t ( ι )) is decreasing. For convenience, we write w ′ t = w t ( ι ′ ) and w t = w t ( ι ) . We do a computation similar

to that in the proof of Lemma D.13:

According to the statement of Lemma 4.4, λ d (3) = 4ˆ σ 2 4 ,d D 4 ,t · 6 d +9 d 2 -1 , meaning that

Thus, on any interval [ t 0 , t 1 ] where D 4 ,t ≥ 0 , Φ( w ′ t ) -Φ( w t ) is decreasing, as desired.

In the next lemma, we show an upper bound on | Φ( w ) -Φ( w ′ ) | , at the beginning of Phase 3, Case 2, for particles w,w ′ which are respectively initialized at ι, ι ′ ∈ [ ι L , ι M ] . Here ι M is defined in the statement/proof of the following lemma.

Lemma D.24. Suppose we are in the setting of Theorem 3.3. Assume that D 4 ,T 2 = 0 and D 2 ,T 2 < 0 , i.e. we are in Phase 3, Case 2. Then, there exists ι M ∈ [ ι L , ι R ] such that P ι ∼ ρ 0 ( | ι | > ι M ) ≲ γ 4 and for all ι, ι ′ ∈ [ ι L , ι M ] ,

for some universal constant C .

Proof of Lemma D.24. Since D 4 ,T 2 = 0 , we have E w ∼ ρ T 2 [ P 4 ,d ( w )] = γ 4 . Thus, by Eq. (C.4),

Since | w t ( ι ) | ≤ 1 for all t ≥ 0 and all ι ∈ [ -1 , 1] , the left-hand side of the above equation is E w ∼ ρ t [ w 4 ] ± O ( 1 d ) , and rearranging gives

where the second inequality is because d ≥ c 3 by Assumption 3.2. Let ι 0 be such that w T 2 ( ι 0 ) = 1 2 . Let ι M = min( ι 0 , ι R ) . Then, ι M ∈ [ ι L , ι R ] , and P ι ∼ ρ 0 ( | ι | > ι M ) ≲ γ 4 , because P ι ∼ ρ 0 ( | ι | > ι 0 ) ≲ γ 4 and P ι ∼ ρ 0 ( | ι | > ι R ) ≲ O ( κ ) by Proposition D.6, meaning that P ι ∼ ρ 0 ( | ι | > ι M ) ≲ max( γ 4 , κ ) ≲ γ 4 (where the last inequality is because d ≥ c 3 by Assumption 3.2).

Now, let ι ∈ [ ι L , ι M ] . Observe that on the interval [0 , T 2 ] , for all ι , w t ( ι ) is strictly increasing as a function of t , for all ι (see Lemma D.9). Thus, if T R is as defined in Lemma D.7, then for all ι ∈ [ ι L , ι R ] , | w T R ( ι ) | ≳ ξκ 2 . In particular, for all ι ∈ [ ι L , ι M ] , we can write w T 2 ( ι ) ≥ Cξκ 2 for some universal constant C . Additionally, for ι ≤ ι M, we have w T 2 ( ι ) ≤ 1 2 by Proposition D.37. Thus, since Φ( w ) is increasing in w , for ι ∈ [ ι L , ι M ] we have

By Markov's inequality,

and

Thus, for any ι, ι ′ ∈ [ ι L , ι M ] ,

where we have defined the constant C ′ = -1 2 log 3 -log C .

In the next lemma, we show that for ι ∈ [ ι L , ι M ] , during Phase 3, Case 2, w t ( ι ) can be bounded away from 0 and 1 , provided that the conclusion of Lemma D.24 holds at time t .

Lemma D.25 (Lower and Upper Bounds on Particles Between ι L and ι M) . Suppose that we are in the setting of Theorem 3.3. Let s ≥ T 2 , and suppose that for all t ∈ [ T 2 , s ] , D 4 ,t ≥ 0 and D 2 ,t ≤ 0 . Further suppose that for all t ∈ [ T 2 , s ] , the conclusion of Lemma D.24 holds, i.e.

for all ι ∈ [ ι L , ι M ] , where ι M is such that P ι ∼ ρ 0 ( | ι | > ι M ) ≲ γ 4 , and for some universal constant C . Then, for any t ∈ [ T 2 , s ] , for all ι ≥ ι L , w t ( ι ) ≥ ξ 2 κ 2 . Furthermore, for all t ∈ [ T 2 , s ] , we have w t ( ι M ) ≤ 1 -κ 4 ξ 3 .

Proof of Lemma D.25. Consider some time t ∈ [ T 2 , s ] . Suppose that for some ι ∈ [ ι L , ι M ] , w t ( ι ) ≤ ξ 2 κ 2 . Also, suppose ι ′ ∈ [ ι L , ι M ] such that ι ′ > ι . Then,

(By Eq. (D.212), and b.c. Φ is increasing, ι > ι and Proposition D.37)

Thus, w t ( ι ′ ) ≲ ξ and for all ι ∈ [ ι L , ι M ] , we have shown that w t ( ι ) ≲ ξ . In addition, P ι ∼ ρ 0 ( | ι | > ι M ) ≲ γ 4 . We can use this to calculate D 2 ,t and D 4 ,t :

and

To complete the proof, we will show that if these bounds on D 2 ,t and D 4 ,t hold, then for any w , v t ( w ) ≥ 0 . By Lemma 4.4, it suffices to show that P t ( w ) + Q t ( w ) ≤ 0 for w ≥ 0 , where P t and Q t are as defined in Lemma 4.4. If C 1 and C 2 are two appropriately chosen positive universal constants, then by the above upper bounds on D 2 ,t and D 4 ,t ,

In summary, if w t ( ι ) = ξ 2 κ 2 for some ι ∈ [ ι L , ι M ] , then v t ( w ) > 0 for all w ∈ [0 , 1] . In addition, at time T 2 , we have w t ( ι ) ≳ ξκ 2 for all ι > ι L by Lemma D.9 and Lemma D.7. Thus, for all ι ≥ ι L, w t ( ι ) will never go below ξ 2 κ 2 for t ∈ [ T 2 , s ] , as desired.

To finish the proof, we will show the upper bound on w t ( ι M ) . Suppose 1 -w t ( ι M ) ≤ κ 4 ξ 3 . Then,

Thus, for any ι ∈ [ ι L , ι M ] ,

for some universal constant C . In other words, letting w t := w t ( ι ) and rearranging using the definition of Φ , we obtain

which finally gives 1 -w 2 t ≲ ξ , or w 2 t ≥ 1 -O ( ξ ) . In other words, for all ι ≥ ι L, we have w t ( ι ) ≥ 1 -O ( ξ ) , which implies that

and this contradicts the assumption that for all t ∈ [ T 2 , s ] , D 4 ,t ≥ 0 and D 2 ,t ≤ 0 . Thus, for all t ∈ [ T 2 , s ] , 1 -w t ( ι M ) ≥ κ 4 ξ 3 , i.e. w t ( ι M ) ≤ 1 -κ 4 ξ 3 .

The purpose of the next two lemmas is to show that if Phase 3, Case 2 occurs, then for all t ≥ T 2 , we will have D 2 ,t ≤ 0 and D 4 ,t ≥ 0 , assuming the conclusion of Lemma D.25 holds. We will later put these next two lemmas with the above lemma, Lemma D.25, in order to inductively show that the hypothesis of Lemma D.25 holds for t ≥ T 2 (note that Lemma D.25 assumes that the upper bound on the potential difference holds on an interval [ T 2 , s ] , and more importantly assumes that D 2 ,t ≤ 0 and D 4 ,t ≥ 0 for all t ∈ [ T 2 , s ] ).

Lemma D.26. Suppose we are in the setting of Theorem 3.3, and suppose D 4 ,t = 0 and D 2 ,t < 0 for some t ≥ 0 . Additionally, suppose that the conclusion of Lemma D.25 holds at time t , i.e. for all ι ≥ ι L , w t ( ι ) ≥ ξ 2 κ 2 , and w t ( ι M ) ≤ 1 -κ 4 ξ 3 . Then, d ds D 4 ,s ∣ ∣ s = t > 0 .

Proof. Recall from Eq. (C.4) that

Next we compute the derivative of D 4 ,t :

For positive particles w ≳ 1 √ d , we have P 4 ,d ′ ( w ) ≥ 0 , and for w ≲ 1 √ d , it may be the case that P 4 ,d ′ ( w ) ≤ 0 . Furthermore, when D 4 ,t = 0 and D 2 ,t ≤ 0 , we have

or meaning that

Finally, to lower bound d dt D 4 ,t , we must show that there is a significant fraction of particles for which w ≳ 1 √ d and 1 -w is large. Because we have w t ( ι ) ≥ ξ 2 κ 2 for all ι ≥ ι L, and P ( | ι | < ι L ) ≲ κ by Proposition D.6, using a union bound we therefore know that

where the second inequality is by Assumption 3.2 since d is sufficiently large. Thus,

where the last inequality is by Eq. (D.232) and applying the above bound on v t ( w ) from Eq. (D.234) for | w | ∈ [ ξ 2 κ 2 , 1 -κ 4 ξ 3 ] , and because d ≥ c 3 by Assumption 3.2. Expanding the last line gives

and the right hand side is strictly positive since d is sufficiently large by Assumption 3.2.

Lemma D.27. Suppose we are in the setting of Theorem 3.3, and suppose D 4 ,t > 0 and D 2 ,t = 0 for some t ≥ 0 . Additionally, suppose that the conclusion of Lemma D.25 holds at time t , i.e. for all ι ≥ ι L , w t ( ι ) ≥ ξ 2 κ 2 , and w t ( ι M ) ≤ 1 -κ 4 ξ 3 . Then, d ds D 2 ,s ∣ ∣ s = t < 0 .

Proof. Recall from Eq. (C.4) that

meaning that

Suppose D 2 ,t = 0 and D 4 ,t > 0 . Then, the velocity is

Thus, for positive particles w , if w ≳ 1 √ d , v t ( w ) ≤ 0 , and for w ≲ 1 √ d , the velocity is potentially positive. Using this, we can upper bound d dt D 2 ,t by upper bounding E w ∼ ρ t [ w · v t ( w )] :

where the last inequality is because P ( | ι | > ι M ) ≤ O ( γ 4 ) (by the definition of ι M, see statement of Lemma D.25) and P ( | ι | < ι L ) ≤ O ( κ ) by Proposition D.6. Simplifying further, we obtain,

and the last expression is strictly negative, since d ≥ c 3 by Assumption 3.2 and D 4 ,t ≥ 0 . This proves the lemma.

Now, we combine Lemma D.25, Lemma D.26 and Lemma D.27 to inductively show that for all t ≥ T 2 , we have D 2 ,t ≤ 0 , D 4 ,t ≥ 0 and most importantly, w t ( ι ) can be bounded away from 0 and 1 for ι ∈ [ ι L , ι M ] .

Lemma D.28 (Phase 3 Case 2 Invariants) . Suppose we are in the setting of Theorem 3.3. Assume D 2 ,T 2 < 0 and D 4 ,T 2 = 0 , i.e. Phase 3 Case 2 holds. Then, for all t ≥ T 2 , the following will hold:

D 2 ,t ≤ 0 and D 4 ,t ≥ 0

Let ι M be defined as in Lemma D.25. Then, w t ( ι ) ≥ ξ 2 κ 2 for all ι ∈ [ ι L , ι M ] and w t ( ι M ) ≤ 1 -κ 4 ξ 3 .

Proof of Lemma D.28. For the purpose of this proof, define

and let t 0 = sup I , meaning that for t ∈ [ T 2 , t 0 ] , D 2 ,t ≤ 0 and D 4 ,t ≥ 0 . Assume for the sake of contradiction that t 0 < ∞ . Then, by the continuity of D 2 ,t and D 4 ,t as functions of t , D 2 ,t ≤ 0 and D 4 ,t ≥ 0 . Since t 0 = sup I , either D 2 ,t 0 = 0 or D 4 ,t 0 = 0 (and one of these has to be nonzero since otherwise t 0 = ∞ , which is a contradiction). By Lemma D.24 and Lemma D.23, for all t ∈ [ T 2 , t 0 ] , for all ι, ι ′ ∈ [ ι L , ι M ] ,

where C is a universal constant. Thus, applying Lemma D.25 with s = t 0 , we find that for all t ∈ [ T 2 , t 0 ] , for all ι ≥ ι L, w t ( ι ) ≥ ξ 2 κ 2 , and furthermore, for all t ∈ [ T 2 , t 0 ] , w t ( ι M ) ≤ 1 -κ 4 ξ 3 .

To finish the proof, there are two cases to address: either D 2 ,t 0 < 0 and D 4 ,t 0 = 0 , or D 2 ,t 0 = 0 and D 4 ,t 0 > 0 . In the first case, by Lemma D.26, d dt D 4 ,t ∣ ∣ ∣ t = t 0 > 0 , meaning that for a sufficiently small open interval of time J containing t 0 , D 4 ,t is increasing on this interval. This gives us a contradiction, since it implies that for some t ′ 0 > t 0 (with t ′ 0 -t 0 sufficiently small) D 4 ,t is nonnegative on [ t 0 , t ′ 0 ] , and since D 2 ,t 0 < 0 , by continuity, D 2 ,t < 0 for t ∈ [ t 0 , t ′ 0 ] as long as t ′ 0 -t 0 > 0 is sufficiently small. Thus, sup I > t 0 , which is a contradiction. In the second case, we can apply the same argument using Lemma D.27 in place of Lemma D.26.

In summary, if D 4 ,T 2 = 0 and D 2 ,T 2 < 0 , then for all t ≥ T 2 , D 2 ,t ≤ 0 and D 4 ,t ≥ 0 . As a consequence of Lemma D.23, Lemma D.24, and Lemma D.25, for all ι ∈ [ ι L , ι M ] , w t ( ι ) ≥ ξ 2 κ 2 , and furthermore, w t ( ι M ) ≤ 1 -κ 4 ξ 3 , for all t ≥ T 2 .

We now proceed to bound the running time of Phase 3, Case 2, by dividing it into a few subcases, and showing that the running time that each of these subcases contribute is within our desired bound. The first lemma below bounds the running time spent in the subcase where the positive root r of P t is larger than 1 2 (and where D 4 ,t is large relative to D 2 ,t ).

Lemma D.29. Suppose we are in the setting of Theorem 3.3, and assume that D 4 ,T 2 = 0 and D 2 ,T 2 < 0 , i.e. we are in Phase 3, Case 2. Then, at most 1 ˆ σ 2 2 ,d ξ 4 κ 6 log ( γ 2 ϵ ) time t ≥ T 2 elapses such that the following hold:

If r is the positive root of P t (as defined in Lemma 4.4) then r ≥ 1 2 .

| D 4 ,t | ≥ ξ | D 2 ,t |

L ( ρ t ) ≥ (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) ϵ 2

Proof of Lemma D.29. Let t ≥ T 2 such that | D 4 ,t | ≥ ξ | D 2 ,t | . We can write

where r is the positive root of P t ( w ) . We are now going to lower bound the expected velocity while r ≥ 1 2 by lower bounding P t ( w ) for a large fraction of w .

For all ι ≥ ι L, and t ≥ T 2 , we have w t ( ι ) ≥ ξ 2 κ 2 , by Lemma D.28. Thus, P w ∼ ρ t ( w ≥ ξ 2 κ 2 ) ≤ 1 -O ( κ ) by Proposition D.6. Additionally, because D 2 ,t ≤ 0 , we have E w ∼ ρ t [ w 2 ] ≤ γ 2 + O (1 /d ) by Proposition D.38, meaning that by Markov's inequality, P w ∼ ρ t ( | w | ≥ 1 3 ) ≲ γ 2 + 1 /d ≤ γ 2 , where the last inequality is because d ≥ c 3 by Assumption 3.2. By a union bound,

For particles w such that | w | ∈ [ ξ 2 κ 2 , 1 / 3] , we can lower bound the velocity if r ≥ 1 2 :

where the last inequality is because | D 4 ,t | ≥ ξ | D 2 ,t | and d is sufficiently large by Assumption 3.2. Thus,

and

Since L ( ρ 0 ) ≍ ˆ σ 2 4 ,d γ 2 2 by Lemma D.8 and Assumption 3.2, the amount of time spent in the case where r ( t ) ≥ 1 2 , | D 4 ,t | ≥ ξ | D 4 ,t | and L ( ρ t ) ≥ (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) ϵ 2 after t ≥ T 2 is at most 1 ˆ σ 2 2 ,d ξ 4 κ 6 log ( γ 2 ϵ ) by Gronwall's inequality (Fact I.13).

Next, we deal with the subcase where D 4 ,t is small relative to D 2 ,t .

Lemma D.30. Suppose we are in the setting of Theorem 3.3, and D 4 ,T 2 = 0 and D 2 ,T 2 < 0 , i.e. we are in Phase 3, Case 2. Then, at most 1 ˆ σ 2 2 ,d ξ 4 κ 4 log ( γ 2 ϵ ) time t ≥ T 2 elapses such that the following hold:

Proof of Lemma D.30. Suppose | D 4 ,t | ≤ ξ | D 2 ,t | at some time t ≥ T 2 . Then, we can lower bound the velocity for a large fraction of w 's, as follows. First, we provide a bound on P t ( w ) (defined in Lemma 4.4):

Next, we provide a bound on Q t ( w ) :

Thus,

By Lemma D.28, we have D 2 ,t ≤ 0 , meaning that by Proposition D.38,

Thus, by Markov's inequality and Assumption 3.2,

Furthermore, by Lemma D.28, for all ι ≥ ι L, w t ( ι ) ≥ ξ 2 κ 2 for all times t ≥ T 2 , meaning that w t ≥ ξ 2 κ 2 with probability at least 1 -O ( κ ) under ρ t , by Proposition D.6. Thus, by a union bound,

where the last inequality is by Assumption 3.2 since γ 2 and d are sufficiently small and large respectively. Thus the average velocity is at least

Thus, at any time t where | D 4 ,t | ≤ ξ | D 2 ,t | , we have

Since the initial loss is L ( ρ 0 ) ≲ ˆ σ 2 2 ,d γ 2 2 by Lemma D.8, by an argument similar to that used in Lemma D.29, using Gronwall's inequality (Fact I.13), we find that

i.e. the time spent in the subcase in the statement of this lemma is at most 1 ˆ σ 2 2 ,d ξ 4 κ 4 log ( γ 2 ϵ ) , as desired.

We have already considered the case where | D 4 ,t | ≥ ξ | D 2 ,t | and r ≥ 1 2 . Next, we must deal with the case where | D 4 ,t | ≥ ξ | D 2 ,t | , but where r ≤ 1 2 . We further split this case into two subcases. In the next lemma, we deal with the subcase where w t ( ι R ) ≥ 1 -ξ , meaning a non-negligible portion of the particles are close to 1 .

Lemma D.31. Suppose we are in the setting of Theorem 3.3, and D 4 ,T 2 = 0 and D 2 ,T 2 < 0 , i.e. we are in Phase 3, Case 2. Then, the amount of time t ≥ T 2 that elapses while the following conditions simultaneously hold:

2. r ≤ 1 2 where r is the positive root of P t

is at most 1 ˆ σ 2 4 ,d ϵξ 7 κ 12 log ( γ 2 ϵ ) .

Proof of Lemma D.31. If w ( ι R ) ≥ 1 -ξ at any time, then the subsequent times when it could possibly be moving to towards 1 are:

During phase 2 - note that before phase 2, w t ( ι R ) ≤ 1 log d , so it is not possible for w t ( ι R ) to be larger than 1 -ξ since, by Lemma D.9, w t ( ι R ) is increasing during Phase 1.

The times t ≥ T 2 when | D 4 ,t | ≤ ξ | D 2 ,t |

The times t ≥ T 2 when | D 4 ,t | ≥ ξ | D 2 ,t | and r ≥ 1 2 , where r is the positive root of P t ( w )

The only case not mentioned above consists of the times t ≥ T 2 when | D 4 ,t | ≥ ξ | D 2 ,t | and the positive root r of P t ( w ) is less than 1 2 . In this case, w t ( ι R ) will be moving away from 1 . For this reason, we also do not need to consider the case discussed in Lemma D.32.

Thus, we will use upper bounds on the velocity and time during the first 3 cases listed above, and use a lower bound on the velocity in the last. First note that the velocity can always be upper bounded as follows:

Thus,

During Phase 2. By Lemma D.10, the running time during Phase 2 is T 2 -T 1 ≲ 1 (ˆ σ 2 2 ,d +ˆ σ 2 4 ,d ) ξ 6 κ 12 log ( γ 2 ϵ ) . At the beginning of Phase 2, w t ( ι R ) ≤ 1 log d , by the definition of Phase 1. Thus, by Gronwall's inequality (Fact I.13), at time T 2 ,

Times t ≥ T 2 when | D 4 ,t | ≤ ξ | D 2 ,t | and L ( ρ t ) ≥ (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) ϵ 2 . By Lemma D.30, this case contributes at most 1 ˆ σ 2 2 ,d ξ 4 κ 4 log ( γ 2 ϵ ) time to the overall running time. Thus, by Gronwall's inequality (Fact I.13), the additional factor that this case contributes to 1 -w t ( ι R ) is larger than or equal to

for some universal constant C .

Times t ≥ T 2 when | D 4 ,t | ≥ ξ | D 2 ,t | , r ( t ) ≥ 1 2 and L ( ρ t ) ≥ (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) ϵ 2 . By Lemma D.29, this case contributes at most 1 ˆ σ 2 2 ,d ξ 4 κ 6 log ( γ 2 ϵ ) to the overall running time. Thus, by Gronwall's inequality (Fact I.13), the additional factor that this case contributes to 1 -w t ( ι R ) is larger than or equal to

for some universal constant C .

Times t ≥ T 2 when | D 4 ,t | ≥ ξ | D 2 ,t | , r ( t ) ≤ 1 2 and w R ≥ 1 -ξ . In this case, we obtain a lower bound on d dt (1 -w R ) , where we have written w R = w t ( ι R ) for convenience. First let us obtain a lower bound on P t ( w R ) :

Next let us obtain an upper bound on Q t ( w R ) :

Thus, because 1 ξd ≍ log log d d , and when w R ≥ 1 -ξ ,

Furthermore, since | D 4 ,t | ≥ ξ | D 2 ,t | , this means that | D 4 ,t | ≳ ξϵ , since otherwise, | D 2 ,t | ≲ ϵ and | D 4 ,t | ≲ ξϵ , and the loss would be less than (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) ϵ 2 . Thus, if time T esc is spent in this case, then by Gronwall's inequality (Fact I.13), this case contributes a factor of at least

to 1 -w R.

Summary of 4 Cases. In summary, by Eq. (D.274), Eq. (D.275), Eq. (D.276) and Eq. (D.280),

where T esc is the amount of time spent so far in the case where | D 4 ,t | ≥ ξ | D 2 ,t | , r ( t ) ≥ 1 2 and w R ≥ 1 -ξ . This is in fact a lower bound on 1 -w R when L ( ρ t ) ≥ (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) ϵ 2 , since at all other times when L ( ρ t ) ≥ (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) ϵ 2 , we have that 1 -w R is increasing (i.e. w R is moving away from 1 ). In particular, if

then 1 -w R will be at least ξ .

Finally, the following lemma deals with the last case, where the positive root r of P t is at most 1 2 , and additionally, w t ( ι R ) ≤ 1 -ξ .

Lemma D.32. 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, the amount of time t ≥ T 2 that elapses while the following conditions simultaneously hold:

Then,

<

17

0

(B.c.

d

is sufficiently large by Assumption 3.2)

(B.c.

(

β

-

1)

>

17

τ

2

by Assumption 3.2)

which contradicts the assumption that we are in Phase 3, Case 2 since D 4 ,t > 0 by Lemma D.28. Thus, if r ≤ τ 3 / 2 , then

τ

6

-

(

β

-

1)

τ

4

Thus,

Then,

Since D 2 ,t < 0 by Lemma D.28, this implies that E w ∼ ρ t [ w 2 ] ≤ τ 2 + O (1 /d ) by Proposition D.38. Thus,

Proof of Lemma D.32. For convenience and ease of presentation, in this proof we define τ := √ γ 2 , and β ≥ 1 . 1 such that γ 4 = βτ 4 . (Note that by Assumption 3.2, β is at most a universal constant.) We first show that there is always a large fraction of particles which are far away from the positive root r of P t ( w ) . We consider two cases: (i) r ≤ τ 3 / 2 and (ii) r ≥ τ 3 / 2 .

Case 1: r ≤ τ 3 / 2 . Assume for the sake of contradiction that

Case 2: r ≥ τ 3 / 2 . Suppose that

and taking square roots gives

Thus,

On the other hand,

and since D 4 ,t ≥ 0 by Lemma D.28, E w ∼ ρ t [ w 4 ] ≥ βτ 4 -O (1 /d ) by Eq. (C.4) and because w ∼ ρ t are bounded by 1 in absolute value. Thus, by the above equation,

Rearranging gives and taking fourth roots gives

Combining Eq. (D.297) and Eq. (D.303) gives

Rearranging this equation gives

(B.c. τ ≤ 1 , we have τ 3 / 2 ≤ τ 5 / 4 )

(B.c. d is sufficiently large by Assumption 3.2)

(B.c. d is sufficiently large by Assumption 3.2)

This contradicts Assumption 3.2 - since β ≥ 1 . 1 , the left-hand side is larger than an absolute constant, while the right-hand side can be made sufficiently small (since γ 2 ≤ c 2 and c 2 is chosen to be sufficiently small). Thus, in the case that r ≥ τ 3 / 2 , we obtain

Further rearranging gives

In summary, whether r ≥ τ 3 / 2 or r ≤ τ 3 / 2 , we can conclude that

Further assume that w t ( ι R ) ≤ 1 -ξ , and note that by Lemma D.28, for all ι > ι L, we have w t ( ι ) ≥ ξ 2 κ 2 . Thus, by Proposition D.6, with probability 1 -O ( κ ) under ρ t , we have ξ 2 κ 2 ≤ | w | ≤ 1 -ξ . By a union bound,

Suppose w satisfies || w | -r | ≥ τ 3 / 2 and ξ 2 κ 2 ≤ | w | ≤ 1 -ξ , i.e. the conditions mentioned in Eq. (D.309), and write

Then,

and the average velocity is at least

Since the initial loss is L ( ρ 0 ) ≲ (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) γ 2 2 by Lemma D.8, we can calculate using Gronwall's inequality (Fact I.13) that the time spent when the conditions in the statement of this lemma simultaneously hold is at most 1 ˆ σ 2 4 ,d τ 11 ξ 8 κ 4 log ( γ 2 ϵ ) , as desired.

To conclude, we show Lemma D.33 which gives a bound on the running time during Phase 3, Case 2. Lemma D.33 (Phase 3, Case 2 Summary) . 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 ( 1 ˆ σ 2 2 ,d ϵγ 11 / 2 2 ξ 8 κ 12 log ( γ 2 ϵ )) time t ≥ T 2 elapses while L ( ρ t ) ≳ (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) ϵ 2 .

Proof of Lemma D.33. By combining Lemma D.29, Lemma D.30, Lemma D.31, and Lemma D.32, we find that the overall running time during Phase 3 Case 2 is at most

and this can be upper bounded by O ( 1 ˆ σ 2 2 ,d ϵγ 11 / 2 2 ξ 8 κ 12 log ( γ 2 ϵ )) , where we have used the fact that ˆ σ 2 2 ,d ≍ ˆ σ 2 4 ,d (see Assumption 3.2).

D.7 Summary of Phase 3

The following lemma provides a bound on the running time after T 2 that is required for L ( ρ t ) ≲ (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) ϵ 2 .

Lemma D.34 (Phase 3) . Suppose we are in the setting of Theorem 3.3. Then, T ∗ ,ϵ -T 2 ≲ (log log d ) 20 ˆ σ 2 2 ,d ϵγ 8 2 log ( γ 2 ϵ ) .

Proof of Lemma D.34. By combining Lemma D.22 and Lemma D.33, we find that the overall running time during Phase 3 is at most

and this is at most O ( 1 ˆ σ 2 2 ,d ϵγ 8 2 ξ 8 κ 12 log ( γ 2 ϵ )) , where we have used the fact that ˆ σ 2 2 ,d ≍ ˆ σ 2 4 ,d (see Assumption 3.2)

D.8 Proof of Main Theorem for Population Case

Finally, we prove Theorem 3.3:

Proof of Theorem 3.3. We combine Lemma D.5, Lemma D.10 and Lemma D.34 to find that the total runtime is at most

and we can further simplify the above bound to obtain:

where we have used the fact that ˆ σ 2 2 ,d ≍ ˆ σ 2 4 ,d by Assumption 3.2.

D.9 Additional Lemmas

Lemma D.35. Suppose Assumption 3.2 holds. There exists some r , with r ≳ 1 log log d and r ≤ 1 , such that if ρ is the singleton distribution under which r has probability 1 , then the velocity at r is 0 .

Proof. Let ρ be such that r has probability 1 under ρ , for some r ∈ (0 , 1) . Applying the definitions of D 2 and D 4 to ρ , we have D 2 = P 2 ,d ( r ) -γ 2 and D 4 = P 4 ,d ( r ) -γ 4 . By Lemma D.3, the velocity of r is

The second factor is equal to

and we denote this polynomial of r by F ( r ) . Note that P 2 ,d (1) = P 4 ,d (1) = 1 (see Equation (2.20) of Atkinson and Han [12]). Thus, F (1) > 0 by Assumption 3.2, since P 2 ,d ′ (1) > 0 and P 4 ,d ′ (1) > 0 by Eq. (C.4). On the other hand, suppose r ≍ 1 log log d . Then,

Similarly, | P 4 ,d ( r ) | ≲ 1 (log log d ) 4 . Thus, P 2 ,d ( r ) -γ 2 < 0 and P 4 ,d ( r ) -γ 4 < 0 since d can be chosen to be sufficiently large according to Assumption 3.2. However, by Eq. (C.4), one can show that P 2 ,d ′ ( r ) ≳ 1 log log d and P 4 ,d ′ ( r ) ≳ 1 (log log d ) 3 as long as d is sufficiently large. Thus, F ( r ) < 0 for r ≍ 1 log log d . In summary, F has a root between C log log d and 1 for some universal constant C , and for this value of r , the velocity of r is 0 .

Proof. This follows from Lemma D.10 and Lemma D.34.

D.10 Helper Lemmas

In this subsection, we give some convenient facts that we use in the following subsections.

Proposition D.37. If 0 ≤ ι 1 < ι 2 , then for all t ≥ 0 , w t ( ι 1 ) < w t ( ι 2 ) .

̸

Proof of Proposition D.37. Assume this is not the case, and let t be the infimum of all times s such that w s ( ι 1 ) ≥ w s ( ι 2 ) . Then, w t ( ι 1 ) = w t ( ι 2 ) , meaning that for all t ′ ≥ t , w t ′ ( ι 1 ) = w t ′ ( ι 2 ) . For convenience, let w 0 = w t ( ι 1 ) = w t ( ι 2 ) . Note that t > 0 since ι 1 = ι 2 . However, this implies that for t ′ ∈ (0 , t ) , w t ′ ( ι 1 ) < w t ′ ( ι 2 ) , meaning that for any δ > 0 , w s ( ι 1 ) and w s ( ι 2 ) are non-identical solutions to the initial value problem on the time interval ( t -δ, t + δ ) given by the conditions f ( t ) = w 0 and df dt = v t ( f ( t )) . This contradicts the Picard-Lindelof theorem, and thus there cannot exist a time t such that w t ( ι 1 ) ≥ w t ( ι 2 ) .

Proposition D.38. Suppose we are in the setting of Theorem 3.3, and let t ≥ 0 such that D 2 ,t < 0 . Then, E w ∼ ρ t [ w 2 ] ≤ γ 2 + O ( 1 d ) .

Proof of Proposition D.37. Suppose D 2 ,t ≤ 0 . Then, E w ∼ ρ t [ P 2 ,d ( w )] ≤ γ 2 , which implies that

and rearranging gives as desired.

Lemma D.39 (Decrease in Loss and Average Velocity) . Suppose we are in the setting of Theorem 3.3. Then, dL dt = -E u ∼ ρ t [ ∥ grad u L ( ρ t ) ∥ 2 2 ] .

We note that a similar result was shown in Lemma E.11 of Wei et al. [75], but for gradient flow rather than projected gradient flow.

Proof of Lemma D.39. First, we can calculate

Furthermore, we can expand d dt f ρ t ( x ) as

Therefore,

This completes the proof.

The above lemma implies that L ( ρ t ) is always decreasing. Note that frequently, when we apply the above lemma, we will only consider the first coordinate of grad u L ( ρ t ) .

Proposition D.40. For any distribution ρ on R ,

Proof of Proposition D.40. This follows from a short calculation:

as desired.

E Proofs for Finite-Samples and Finite-Width Dynamics

In this section we give the proof for results related to the empirical dynamics, where the number of samples and network width are both polynomial in d . The end goal is to prove Theorem 3.4, the proof of which is given below.

Proof of Theorem 3.4. We have

where the first inequality is by Cauchy-Schwarz inequality, and the second inequality is by Lemma 5.1, Lemma 5.2 and Theorem 3.3. Since ϵ = 1 log log d , the last term dominates because the conditions of Theorem 3.4 imply that m ≳ d 3 , and because µ -1 2 > 1 , meaning that the first term is at most ˆ σ 2 2 ,d +ˆ σ 2 4 ,d d .

To complete the proof, we give proofs for results in Section 5. We first give a proof for Lemma 5.1:

Proof of Lemma 5.1. For the first statement, for any fixed t ≤ T , by Lemma E.13, we have

with probability 1 -δ -wewill choose δ shortly. In addition, for any t, t ′ ∈ [0 , T ] , by Lemma E.14, we have that ∥ u t ( χ ) -u t ′ ( χ ) ∥ 2 ≲ (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) d 4 | t -t ′ | , which by Lemma E.12 implies that for any x ∈ S d -1 ,

In particular, if | t -t ′ | ≤ 1 (ˆ σ 2 2 ,d +ˆ σ 2 4 ,d ) d Ω(log d ) , then for all x ∈ S d -1 we have

By the same argument, we have that

Now, we perform a union bound over all times t such that t ≤ T and t is an integer multiple of 1 (ˆ σ 2 2 ,d +ˆ σ 2 4 ,d ) d H log d for some sufficiently large absolute constant H > 0 . Since T ≤ d C ˆ σ 2 2 ,d +ˆ σ 2 4 ,d , the number of such times is at most d C + H log d . Thus, with probability at least 1 -δ · d C + H log d , we have

for all t ∈ [0 , T ] . In particular, in order to have exponentially small failure probability, we set δ = 1 e d 2 d C + H log d , and by our assumption that m ≤ d O (log d ) , we obtain

for all t ≤ T , with probability at least 1 -e -d 2 , since log 1 δ = d 2 + O ((log d ) 2 ) . This completes the proof of the first statement of the lemma. Additionally, by Lemma E.16, we have

which proves the second statement of the lemma.

The rest of the section is dedicated to prove Lemma 5.2. To begin with, we first introduce the following lemma which formalizes the decomposition of d dt ∥ ˆ u t -¯ u t ∥ 2 2 discussed in Section 5.

Lemma E.1 (Decomposition of d dt ∥ ˆ u t -¯ u t ∥ 2 2 ) . Suppose we are in the setting of Theorem 3.4. Then, d dt ∥ ˆ u t -¯ u t ∥ 2 2 = A t + B t + C t , where A t = -2 ⟨ grad ˆ u L ( ρ t ) -grad ¯ u L ( ρ t ) , ˆ u t -¯ u t ⟩ , B t = -2 ⟨ grad ˆ u L (ˆ ρ t ) -grad ˆ u L ( ρ t ) , ˆ u t -¯ u t ⟩ , and C t = -2 ⟨ grad ˆ u ̂ L (ˆ ρ t ) -grad ˆ u L (ˆ ρ t ) , ˆ u t -¯ u t ⟩ .

Proof of Lemma E.1. By the chain rule,

meaning that

We can expand the difference in gradients as

Thus,

as desired.

Next, we give a proof of Lemma 5.4, which gives an upper bound for each of A t , B t and C t .

Proof of Lemma 5.4. The proof directly follows from Lemma E.8, Lemma E.7, Lemma E.9, Lemma E.10 and Lemma E.11.

Combining the upper bounds in Lemma 5.4 and the decomposition in Lemma E.1, we have the following Lemma E.2 which upper bounds the growth of ∥ ˆ u t -¯ u t ∥ . Recall that T ∗ ,ϵ is the total runtime for the algorithm (see Theorem 3.3), and T 1 is the time such that w T 1 ( ι U ) = 1 log d (see Definition 4.5).

Lemma E.2 (Growth Rate of ∥ δ t ∥ 2 2 With Dependency on n ) . In the setting of Theorem 3.4, suppose that the inductive hypothesis in Assumption 5.3 holds for some ϕ and ψ up until t for some t ≤ T ∗ ,ϵ . Further assume that ψ < 1 2 and ϕ > 1 2 . Then, when t ≤ T 1 we have

When t ∈ [ T 1 , T ∗ ,ϵ ] , we have

and

and

Proof of Lemma E.2. This result follows directly from Lemma 5.4 and the inductive hypothesis. By the inductive hypothesis, we have ∆ t ≤ d -ϕ . We apply this to the bound for C t in Lemma 5.4 to obtain the desired bounds above. Note that in our bounds for d dt ∆ 2 t , it is important that we use the upper bound for E [ B t ] in Lemma 5.4 instead of the upper bound for B t .

Using this bound on the growth rate, we can show that the inductive hypothesis is maintained.

Lemma E.3 (IH is Maintained) . Suppose we are in the setting of Theorem 3.4. Suppose ψ > 0 and 1 > ϕ > 1 2 + ψ . Let n = d µ (log d ) C where µ ≥ max { 2 + 2 ϕ, 4 -2 ϕ -4 ψ } and C is a sufficiently large universal constant. Suppose the inductive hypothesis in Assumption 5.3 holds for all t ≤ T where T < T ∗ ,ϵ . Then, when m ≳ d 3+2 ϕ (log d ) Ω(1) and m ≤ d C for a sufficiently large universal constant C , we have that the inductive hypothesis holds for all t ≤ T + c IH where c IH > 0 is a constant that only depends on d , ϕ , ψ , ˆ σ 2 ,d , ˆ σ 4 ,d .

Proof. Suppose the inductive hypothesis holds for all t ≤ T where T ≤ T 1 . Then, by Lemma E.2 we have

Here we have used the fact that n = d µ (log d ) O (1) where µ ≥ max(4 -4 ψ -2 ϕ, 2 + 2 ϕ ) as well as the fact that ϕ > 1 2 , as follows:

In the third term in the bound for d dt ∥ ˆ u t -¯ u t ∥ 2 2 from Lemma E.2, only the 1 d ϕ term is important - this term dominates the √ d n and d 2 . 5 -2 ψ -2 ϕ n terms, using the fact that µ > 2 + 2 ϕ .

In the fourth term in the bound for d dt ∥ ˆ u t -¯ u t ∥ 2 2 , the d 2 -2 ψ -ϕ n ∥ ˆ u t -¯ u t ∥ 2 2 term is dominated by the 1 d ϕ ∥ ˆ u t -¯ u t ∥ 2 term - this is because µ ≥ 2 + 2 ϕ .

Since µ ≥ 4 -4 ψ -2 ϕ , if we increase n by a (log d ) O (1) factor, then the d 4 -4 ψ -2 ϕ n ∥ ˆ u t -¯ u t ∥ 2 2 term can be absorbed into the first term, as it is dominated by the 4ˆ σ 2 2 ,d | D 2 ,t | O ( 1 log d ) ∥ ˆ u t -¯ u t ∥ 2 2 term. (Recall that by Lemma D.8, | D 2 ,t | ≳ γ 2 during Phase 1).

In the second inequality above, we have used the fact that m ≳ d 3+2 ϕ (log d ) Ω(1) -this implies that the term involving m is dominated by the (log d ) O (1) d ϕ ∥ ˆ u t -¯ u t ∥ 2 term. Additionally, by Lemma E.2, we have

Here, we apply the bound for d dt ∆ 2 t from Lemma E.2. In the third term in this bound, the √ d n term dominates the d 2 . 5 -2 ψ -2 ϕ n term, using the fact that ϕ > 1 2 and µ ≥ 2 + 2 ϕ > 3 , meaning d 2 . 5 -2 ψ -2 ϕ n ≤ d 1 . 5 n ≤ 1 √ n . In the fourth term, we deal with the d 4 -4 ψ -2 ϕ n term similarly to our calculations for the bound for d dt ∥ ˆ u t -¯ u t ∥ 2 2 . We can also deal with the d 2 -2 ψ -ϕ n term similarly, since by our definition of µ , this will be at most 1 d 3 ϕ ≤ 1 (log d ) O (1) , and will thus be absorbed into the first term involving | D 2 ,t | . Using µ ≥ 2 + 2 ϕ , we then obtain √ d n ≤ 1 d 0 . 5+ ϕ as the coefficient of the ∆ t term (i.e. the first-order term involving ∆ t ). Note that here we need to have m ≳ d 3+2 ϕ (log d ) Ω(1) in order to absorb the term involving m into the 1 d 0 . 5+ ϕ ∆ t term.

Now using Lemma I.12 we have

and

and by Lemma D.8 we know that | D 2 ,t | ≳ γ 2 during Phase 1. Thus,

Since 2 ϕ -1 > 2 ψ and d is chosen to be sufficiently large after γ 2 is chosen (Assumption 3.2), we have that ∆ max ,T < √ 1 d ψ + ϕ -1 / 2 = 1 d 1 2 ϕ + 1 2 ψ -1 4 < 1 d ψ . Note that ∆ max ,t is continuous in t and we have

χ

∥

grad

ˆ

u

t

̂

L

(ˆ

u

t

)

-

grad

¯

u

t

L

(¯

u

t

)

∥

2

(By Cauchy-Schwarz Inequality)

Thus, for µ 1 = ( 1 d ψ -1 d 1 2 ϕ + 1 2 ψ -1 4 ) · 1 (ˆ σ 2 2 ,d +ˆ σ 2 4 ,d ) d 4 we still have the inductive hypothesis for ∆ max until time T + µ 1 , that is, ∆ max ,T + µ 1 ≤ 1 d ψ . Similarly, our calculation above using Lemma I.12 gives us

Thus, we have ∆ T < 1 d ϕ (log d ) O (1) < 1 d ϕ (where in the middle expression, we choose a smaller power of (log d ) O (1) so that the inequality holds). Observe that ∆ t is continuous in t , and

Thus, for µ 2 = ( 1 d ϕ -1 d ϕ (log d ) O (1) ) · 1 (ˆ σ 2 2 ,d +ˆ σ 2 4 ,d ) d 4 we still have the inductive hypothesis for ∆ until time T + µ 2 , i.e. ∆ T + µ 2 ≤ 1 d ϕ . As a result, assuming the inductive hypothesis holds up to time T , it holds up to time T +min( µ 1 , µ 2 ) .

Our above proof strategy applies to all times T ≤ T 1 . For the induction when T ∈ [ T 1 , T ∗ ,ϵ ] we follow a similar argument but instead use the bound on A t obtained from Lemma E.8. Since we know from Lemma D.36 and the assumption that ϵ = O ( 1 log log d ) that

we know that this additional phase will at most add another exp( poly (log log d )) factor to the growth of ∆ max ,t which is smaller than any polynomial in d . Thus, Eq. (E.47) and Eq. (E.51) still hold for this period, which finishes the proof.

Proof of Lemma 5.2. The proof directly follows Lemma E.3 with ϕ = µ -1 4 and ψ = µ -3 8 .

E.1 Upper bound on A t

The goal of this section is to obtain an upper bound on

When there is no risk of confusion, we simply use A t , ˆ u t , and ¯ u t , omitting the dependence on χ .

We first obtain a strong upper bound on this inner product during Phase 1 of the population projected gradient flow. To obtain an upper bound on this inner product, we separately consider the contributions from the first coordinates of the vectors, and the remaining ( d -1) coordinates - in other words, the gradients with respect to w and the gradients with respect to z . For convenience, let grad w L ( ρ ) denote the gradient with respect to w and grad z L ( ρ ) denote the gradient with respect to z . The following lemma is about the contribution of the first coordinates of the particles to A t , and is useful during Phase 1.

Lemma E.4 (Contribution of w 's) . Suppose we are in the setting of Theorem 3.4, and let ρ t be the projected gradient flow on population loss, initialized with the uniform distribution on S d -1 , defined

in Eq. (3.7) . Suppose D 2 ,t < 0 and D 4 ,t < 0 , where D 2 ,t and D 4 ,t are as defined in Section 4.1. Then, the contribution of the first coordinates to A t can be bounded by

where w max is as defined in Definition 4.5.

The constant 4 at the beginning of the first term is particularly important since it determines the growth rate of ∥ δ t ∥ 2 2 , and in turn affects the number of samples we will need. At least during Phase 1, the contribution of the first coordinates will be the dominant term in A t , as we will see in the next few lemmas.

Proof of Lemma E.4. The contribution of the first coordinate to the dot product defining A t is

where v ( w ) is as defined in Lemma 4.4. We can further simplify v ( ˆ w t ) -v ( ¯ w t ) using Lemma 4.4. Let P t and Q t be defined as in Lemma 4.4. Then,

Let us upper bound the term corresponding to Q t first. Since each coefficient of Q t is bounded above by O ( ˆ σ 2 2 ,d | D 2 ,t | +ˆ σ 2 4 ,d | D 4 ,t | d ) (by the conclusion of Lemma 4.4) and Q t only has linear and cubic terms, we have

by Lemma I.16. The terms corresponding to P t ( ¯ w t ) in ( v ( ˆ w t ) -v ( ¯ w t ))( ˆ w t -¯ w t ) can be bounded as follows:

≤ 2ˆ σ 2 2 ,d | D 2 ,t | ( ˆ w t -¯ w t ) 2 +4ˆ σ 2 4 ,d | D 4 ,t | ( ˆ w t -¯ w t ) 2 ( ˆ w 2 t + ˆ w t ¯ w t + ¯ w 2 t ) (Since last two terms are nonnegative) ≤ 2ˆ σ 2 2 ,d | D 2 ,t | ( ˆ w t -¯ w t ) 2 + O (ˆ σ 2 4 ,d | D 4 ,t || ˆ w t + ¯ w t | ( ˆ w t -¯ w t ) 2 ) (Since ¯ w t , ˆ w t ≤ 1 )

Thus, the overall bound on the contribution of the first coordinates to A t is

We can simplify this bound further in terms of w max and ∆ max ,t . First observe that

Thus, by the triangle inequality, we have | ¯ w t | ≤ w max and | ˆ w t | ≤ w max +∆ max ,t meaning

(E.78)

as desired.

Next, we study the contribution of the last ( d -1) coordinates of the particles to A t ( χ ) during Phase 1. We first introduce the following formula for the projected gradient of the population loss with respect to the last ( d -1) coordinates of a particle.

Lemma E.5. Let ρ be a distribution which is rotationally invariant, as in Definition 4.1. Let grad u L ( ρ ) be defined as in Eq. (3.4) . In addition, let grad z L ( ρ ) denote the last ( d -1) coordinates of grad u L ( ρ ) , for any particle u = ( w,z ) . Then, for any particle u = ( w,z ) , we have grad z L ( ρ ) = -( w 1 -w 2 · grad w L ( ρ ) ) z .

Proof of Lemma E.5. Notice that w 2 + z 2 = 1 for any u = ( w,z ) ∈ S d -1 , and this equality is maintained by projected gradient descent on the population loss, as in Eq. (3.7). Also, by Lemma I.15, we have grad z L ( ρ ) = cz for some scalar c , since ρ is rotationally invariant. Thus,

By directly applying this formula for the last ( d -1) coordinates of the projected gradient, we have the following upper bound on the contribution of the last ( d -1) coordinates to A t during Phase 1.

Lemma E.6 (Contribution of z 's During Phase 1.) . Suppose we are in the setting of Theorem 3.4. Let t ≤ T 1 (where T 1 is as defined in Definition 4.5). Let ρ t be defined according to the population projected gradient flow as in Eq. (3.7) . Let grad u L ( ρ ) be as in Eq. (3.4) , and let grad z L ( ρ ) denote the last ( d -1) coordinates of grad u L ( ρ ) for any particle u = ( w,z ) . Then, we have the following upper bound:

Rearranging gives us

Proof of Lemma E.6. Using the formula for projected gradient in Lemma E.5, we have that

The second term is always negative for t ≤ T 1 , because Lemma D.9 implies that v ( ¯ w t ) has the same sign as ¯ w t , meaning that grad ¯ w L ( ρ t ) has the opposite sign. We bound the first term as

as desired.

We can combine Lemma E.4 and Lemma E.6 and get the following bound on A t during Phase 1. Lemma E.7. Suppose we are in the setting of Theorem 3.4. Let t ≤ T 1 (where T 1 is as defined in Definition 4.5). Then,

Proof of Lemma E.7. This directly follows from Lemma E.4 and Lemma E.6.

Next, we give a potentially looser bound for A t which holds even after Phase 1. The running time after Phase 1 is at most (log log d ) O (1 ) (aside from the dependency on ˆ σ 2 ,d , ˆ σ 4 ,d , etc.), meaning that even the following growth rate is enough to ensure that the coupling error grows by only a sub-polynomial factor after Phase 1.

Lemma E.8. Suppose we are in the setting of Theorem 3.4. Then, for all t ≥ 0 , we have

Proof of Lemma E.8. We first consider the part of A t that originates from the non-projected gradients, which is defined as

By the Cauchy-Schwarz inequality, we have

We bound the first factor in the equation above by the loss at initialization, which is

by Lemma D.8. Thus, plugging Eq. (E.101) and Eq. (E.210) (from the statement of Lemma E.15) into Eq. (E.100) gives us

Next, we bound the difference between A t and A ′ t . We have

For the first term of Eq. (E.104), we have

and by an argument similar to the one used to bound A ′ t (here using Eq. (E.211) from Lemma E.15 instead of Eq. (E.210)), we can thus show that

Similarly, for the second term we have

Combining Eq. (E.110), Eq. (E.111) and Eq. (E.102) finishes the proof.

E.2 Upper Bound on B t

The goal of this section is to obtain an upper bound on

When there is no risk of confusion, we simply use B t , ˆ u t , and ¯ u t , omitting the χ dependency. In the next lemma, we first obtain a bound on B t ( χ ) , and after that, we obtain a bound on E χ [ B t ( χ )] . The proof of the following lemma is mostly a straightforward calculation, and follows mostly from Lemma 5.1 and the definition of ∆ .

Lemma E.9. Suppose we are in the setting of Theorem 3.4. Suppose T ≲ d C ˆ σ 2 2 ,d +ˆ σ 2 4 ,d for any absolute constant C . Assume the width m of ¯ ρ, ˆ ρ is at most d O (log d ) . Then, with probability at least 1 -e -d 2 , we have

for all times t ∈ [0 , T ] . Additionally, with probability at least 1 -e -d 2 , for all t ∈ [0 , T ] and χ ∈ ˆ ρ 0 , we have

Proof of Lemma E.9. We first decompose B t into two terms:

where the last equality is because ∥ ˆ u t -¯ u t ∥ 2 2 2 = 1 -⟨ ˆ u t , ¯ u t ⟩ since ∥ ˆ u t ∥ 2 = ∥ ¯ u t ∥ 2 = 1 . For convenience, define δ t := ˆ u t -¯ u t . Then, we have

We first focus on the first term in Eq. (E.121), which is the main contributor to B t .

Analysis of First Term in Eq. (E.121). For convenience, we refer to the first term in Eq. (E.121) as B ′ t . Observe that we can expand B ′ t as

Let us bound the first term of Eq. (E.123). By Lemma 5.1, with probability e -d 2 , for all t ≤ T we have

Thus, we can bound the first term of Eq. (E.123) by

Next, we bound the second term in Eq. (E.123). By applying the Cauchy-Schwarz inequality, we have

Combining Eq. (E.126) and Eq. (E.129) with Eq. (E.123), we find that

Analysis of Second Term in Eq. (E.121). For convenience, we write this term as B ′′ t , i.e. we define B ′′ t := ⟨∇ ˆ u L (ˆ ρ t ) -∇ ˆ u L ( ρ t ) , ˆ u t ⟩·∥ δ t ∥ 2 2 . It suffices to show that the first dot product in the definition of B ′′ t is at most (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) up to constant factors. To start, we have

Both terms in Eq. (E.133) can be bounded as before. For the first term we use Eq. (E.124):

(By Cauchy-Schwarz Inequality)

For the second term, we use Lemma E.16:

(By Cauchy-Schwarz Inequality)

(By Lemma E.16 and Lemma E.15, Eq. (E.211))

Thus, we can conclude that

Overall Bound on B t . Combining our bounds on B ′ t and B ′′ t , and using the fact that ∥ δ t ∥ 2 ≤ 2 to absorb the bound for B ′′ t into the bound for B ′ t , we find that

with probability 1 -e -d 2 (recall that Eq. (E.124) held with probability 1 -e -d 2 ). This completes the proof.

Next, we show a tighter bound when we take the expectation of B t over the neurons, i.e. over χ . Taking the expectation allows us to obtain a tighter bound because of the term

For a single ˆ u t , the best bound we can obtain is by using the Cauchy-Schwarz inequality. However, the key point in the following lemma is that when we take the expectation over (ˆ u t , ¯ u t ) ∼ Γ t , this term becomes

together with a third-order term which is at most (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d )∆ 3 t , which is a significantly better bound than (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d )∆ t ∥ δ t ∥ 2 , which we obtained in the previous lemma. This is because to simplify Eq. (E.138), we can use

meaning that we can essentially substitute σ ′ ( ⟨ ˆ u t , x ⟩ ) ⟨ δ t , x ⟩ ≈ σ ( ⟨ ˆ u t , x ⟩ ) -σ ( ⟨ ¯ u t , x ⟩ ) in Eq. (E.138). Note that we need this better bound on the growth rate of ∆ t since our inductive hypothesis imposes a stricter condition on ∆ t .

Lemma E.10. Suppose we are in the setting of Lemma E.9, and in particular, suppose T satisfies the same assumptions as in the statement of Lemma E.9. Then, for all t ≤ T , we have

Proof of Lemma E.10. We again make use of Eq. (E.121), and we recall from the proof of the previous lemma that B t = B ′ t + B ′′ t , where

and

The main difference between Lemma E.9 is that we now bound the second term of B ′ t in expectation. By Taylor's theorem, we can write

for some λ which is a convex combination of ⟨ ˆ u, x ⟩ and ⟨ ¯ u, x ⟩ . Rearranging gives

Thus, taking the expectation of the second term of B ′ t over (ˆ u t , ¯ u t ) ∼ Γ t gives

Next, let us bound the inner expectation in Eq. (E.149). By the Cauchy-Schwarz inequality, we have

Here to obtain the last inequality, we observe that σ ′′ is convex because P 2 ,d ′′ is a constant function and P 4 ,d ′′ is a quadratic function by Eq. (C.4) - additionally, we used the inequality ( a + b ) 2 ≲ a 2 + b 2 and the fact that λ is a convex combination of ⟨ ˆ u t , x ⟩ and ⟨ ¯ u t , x ⟩ .

By an argument similar to the proof of Lemma E.15, we can show that E x ∼ S d -1 [ σ ′′ ( ⟨ ˆ u t , x ⟩ ) 2 ⟨ δ t , x ⟩ 4 ] 1 / 2 ≲ (ˆ σ 2 2 ,d +ˆ σ 2 4 ,d ) 1 / 2 ∥ δ t ∥ 2 2 and E x ∼ S d -1 [ σ ′′ ( ⟨ ¯ u t , x ⟩ ) 2 ⟨ δ t , x ⟩ 4 ] 1 / 2 ≲ (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) 1 / 2 ∥ δ t ∥ 2 2 . Combining this with Eq. (E.149), we have that

and thus,

by the definition of ∆ 3 t . This gives a bound on the second term of B ′ t . Combining this with Eq. (E.126) to bound the first term of B ′ t , we have

Combining this with Eq. (E.136) (for which we also take the expectation over (ˆ u t , ¯ u t ) ∼ Γ t ), we have that

which is at most the bound that we have for B ′ t , and thus,

as desired.

E.3 Upper Bound on C t

The goal of this section is to obtain an upper bound on

This quantity is essentially the part of the growth of ∥ δ t ∥ 2 2 which is due to the difference between the gradient using finite samples and the gradient using infinitely many samples. When there is no risk of confusion, we simply use C t , ˆ u t , and ¯ u t , omitting the χ dependency. We bound this term with the following lemma:

Lemma E.11. Suppose we are in the setting of Theorem 3.4. Assume that the number of samples x i is n ≤ d C , for any universal constant C > 0 . Assume that the width m ≤ d C for any universal constant C > 0 . Let T > 0 . Assume for all t ∈ [0 , T ] we have a bound B 1 on ∆ t with B 1 ≤ 1 √ d . Additionally, assume for all t ∈ [0 , T ] that we have a bound B 2 on ∆ max ,t such that B 2 ≥ 1 √ d . Then, with probability 1 -1 d Ω(log d ) , for all t ∈ [0 , T ] and χ ∈ S d -1 , we have

Proof. Webound C t by expanding grad ˆ u L (ˆ ρ t ) as the sum of several monomials, as well as expanding grad ˆ u ̂ L (ˆ ρ t ) , and then applying Lemma I.7 to each term from grad ˆ u L (ˆ ρ t ) and the corresponding term

from grad ˆ u ̂ L (ˆ ρ t ) . For convenience, we use ℓ ( x ) (ˆ ρ ) to denote the loss of f ˆ ρ on a single data point x ∈ S d -1 , and grad ˆ u ℓ ( x ) (ˆ ρ ) to denote the projected gradient of ℓ ( x ) (ˆ ρ ) with respect to ˆ u . We have

For convenience, define δ t := ˆ u t -¯ u t . Then we can rewrite the above as

i.e. how close the empirical mean of the g ( x i ) is to its expected value. We will show that C ( g ) is small when g corresponds to each of the terms on the right-hand side of Eq. (E.165), even when we take the supremum over the ˆ u t , which may depend on the x i .

Concentration for Mean of First Term in Eq. (E.165). As a first step, we list all of the terms we obtain when expanding the first term on the right-hand side of Eq. (E.165). The σ ′ ( ⟨ ˆ u t , x ⟩ ) factor consists of the following terms: (1) a ˆ σ 2 ,d √ N 2 ,d ⟨ ˆ u t , x ⟩ term, (2) a ˆ σ 4 ,d √ N 4 ,d ⟨ ˆ u t , x ⟩ 3 term, a ˆ σ 4 ,d √ N 4 ,d ⟨ ˆ u ′ t , x ⟩ 4 term, and (3) a ˆ σ 4 ,d N 4 ,d d ⟨ ˆ u ′ t , x ⟩ 2 term. Here note that ˆ u ′ t is not a single vector, but ranges over the entire support of ˆ ρ t . Thus we implicitly perform a union bound over the m vectors in the support of ˆ ρ t -this does not affect the proof since the failure probability of Lemma I.7 is 1 d Ω(log d ) , while we assume the width m is at most d C where C can be any absolute constant. Additionally, y ( x ) consists of the following terms: (1) a ˆ σ 2 ,d γ 2 √ N 2 ,d ⟨ e 1 , x ⟩ 2 term, (2) a √

and (3) a ˆ σ 4 ,d √ N 4 ,d d ⟨ ˆ u t , x ⟩ term - here, we have stated these terms up to absolute constant factors for convenience. In the f ˆ ρ ( x ) -y ( x ) factor, f ˆ ρ ( x ) contributes (1) a ˆ σ 2 ,d √ N 2 ,d ⟨ ˆ u ′ t , x ⟩ 2 term, (2) √

For convenience, throughout this proof, let ¯ δ t = δ t ∥ δ t ∥ 2 . In each of the above factors, for every occurrence of ⟨ ˆ u t , x ⟩ , we substitute ˆ u t = ¯ u t + δ t and further expand. This is useful since originally ˆ u t depends on the x i and thus, when we applied Lemma I.7 we would have to take the supremum over ˆ u t -however, now ¯ u t can be considered a fixed vector with respect to the x i , and while we need to take the supremum over δ t , in each term where δ t occurs, we obtain an additional factor of ∥ δ t ∥ 2 , which reduces the error from uniform convergence. We now completely expand the first term on the right-hand side of Eq. (E.165). First we consider the contribution to C ( g ) of f ˆ ρ ( x ) and later study y ( x ) . Up to constant factors, it suffices to obtain an upper bound on the concentration of

where we have let ˆ σ max = max( | ˆ σ 2 ,d | , | ˆ σ 4 ,d | ) . Here, the sum ranges over p + q = 1 , 3 and q + s = 2 , 4 since the δ t and ¯ u t terms are due to the σ ′ ( ⟨ ˆ u t , x ⟩ ) factor, while the δ ′ t and ¯ u ′ t terms are due to the f ˆ ρ ( x ) factors.

We now obtain concentration bounds for each of the terms in Eq. (E.168). First, we consider the case where p + q ≥ 1 , as in this case, Lemma I.7 can be applied (since we can choose p 1 , p 2 , p 3 , p 4

For illustration, define C ( g ) as

in Lemma I.7 so that u 1 = ¯ δ t and p 1 = p +1 , and u 2 = ¯ δ ′ t and p 2 = q , and p 3 = p 4 = 0 -the hypotheses are then satisfied since p 1 + p 2 + p 3 + p 4 ≥ 2 ). From Lemma I.7, we obtain the bound

Here in applying Lemma I.7, we have assumed that the number of samples n is at most d C for some absolute constant C , meaning that the 1 d Ω(log d ) term in the conclusion is a lower-order term. Lastly, we consider the case where p + q = 0 . To apply Lemma I.7, we can let u 1 = ¯ u t , p 1 = 1 , u 2 = δ t ∥ δ t ∥ 2 , p 2 = 1 , and let p 3 = p 4 = 0 , with w 1 = ¯ u t , q 1 = r -1 and w 2 = ¯ u ′ t , q 2 = s -since p 1 + p 2 = 2 , the hypotheses of Lemma I.7 are still satisfied. This leads to us obtaining only one power of ∥ δ t ∥ 2 in the bound for this term: if p = q = 0 , then we have the bound

Summing over the possible values of p, q, r, s , we find that the overall contribution to the concentration error of the term with f ˆ ρ is, up to logarithmic factors and a ˆ σ 2 max factor,

Finally, in order to simplify this bound, we isolate the dominant terms among the terms of the form d p + q +1 2 n ∥ δ t ∥ p +1 2 E δ ′ ∼ ˆ ρ ∥ δ ′ t ∥ q 2 , with the following cases. First, consider all the terms where p = 0 . For q ≤ 2 we can bound these terms by d p + q +1 2 n ∆ q ∥ δ t ∥ 2 , and for q > 2 we can bound these terms by d p + q +1 2 n · ∆ 2 · B q -2 2 ∥ δ t ∥ 2 . Thus, out of these terms, the term with q = 4 is dominant, and it contributes

Next, consider the case where p ≥ 1 . In this case, since the bound B 2 on ∆ max ,t is greater than 1 d 1 / 2 , the term with p = 3 and q = 4 dominates, and it contributes

However, we must lastly consider the case where p ≥ 1 and q < 2 -the bound on the concentration error may be larger for q = 1 than q = 2 since ∆ < 1 √ d . In this case, since B 2 > 1 √ d , we must

consider the p = 3 , q = 1 and p = 3 , q = 0 terms, which respectively contribute

Note that the p = 3 , q = 0 term dominates since our bound on ∆ is less than 1 √ d , meaning we can ignore the p = 3 , q = 1 term. Thus our overall bound on the concentration error originating from the first term in Eq. (E.165) is

To deal with the terms originating from y ( x ) , we observe that we can simply treat ⟨ e 1 , x ⟩ in the same way as we deal with ⟨ ¯ u t , x ⟩ , and thus the terms originating from y ( x ) do not make any new contributions to the concentration error.

Concentration for Mean of Second Term in Eq. (E.165). Finally, to deal with the second term on the right-hand side of Eq. (E.165), we note that we can follow the same proof, but instead obtaining an upper bound on the concentration of

and as before, the contribution of the terms originating from y ( x ) is at most the contribution of the terms in Eq. (E.183). Due to the additional ∥ δ t ∥ 2 2 factor, we can in fact ignore the contribution of the terms in Eq. (E.183), as the dominant terms will be less than the dominant terms from Eq. (E.168).

Concluding to Bound C t . In summary, with probability at least 1 -1 d Ω(log d ) , for all choices of δ t , δ ′ t , we have

where g ( x ) = ⟨ grad ˆ u ℓ ( x ) (ˆ ρ t ) , δ t ⟩ . Since C t is exactly equal to

up to signs and constant factors, we therefore have

E.4 Additional Lemmas

Lemma E.12 (Uniform Bound on σ and σ ′ ) . For any t ∈ [ -1 , 1] , we have | σ ( t ) | ≤ | ˆ σ 2 ,d | √ N 2 ,d + | ˆ σ 4 ,d | √ N 4 ,d , and additionally, | σ ′ ( t ) | ≲ | ˆ σ 2 ,d | √ N 2 ,d + | ˆ σ 4 ,d | √ N 4 ,d . Finally, we have | σ ′′ ( t ) | ≲ | ˆ σ 2 ,d | √ N 2 ,d + | ˆ σ 4 ,d | √ N 4 ,d .

Proof of Lemma E.12. The first statement is by the definition of σ , and because | P 2 ,d ( t ) | , | P 4 ,d ( t ) | ≤ 1 for any t ∈ [ -1 , 1] , by Atkinson and Han [12], Eq. (2.116). The second and third statements are by the definition of σ and using Eq. (C.4).

and

Lemma E.13 (Error from Finite Width) . Let ρ be a distribution on S d -1 , and let ¯ ρ be a distribution on S d -1 of support at most m , where the elements of the support of ¯ ρ are sampled i.i.d. from ρ . Additionally, let δ > 0 . Then, for some absolute constant C > 0 ,

with probability at least 1 -δ , where f ρ and f ¯ ρ are defined as in Eq. (3.1) .

Proof of Lemma E.13. We can write f ρ ( x ) = E u ∼ ρ [ σ ( ⟨ u, x ⟩ )] , and f ¯ ρ ( x ) = 1 m ∑ m i =1 σ ( ⟨ u i , x ⟩ ) . Now, let B ∈ (0 , 1) , which we will choose appropriately later, and define the function φ B : [ -1 , 1] → [ -1 , 1] by

Additionally, define ˜ f ρ ( x ) = E u ∼ ρ [ σ ( φ B ( ⟨ u, x ⟩ ))] , and define ˜ f ¯ ρ ( x ) = 1 m ∑ m i =1 σ ( φ B ( ⟨ u i , x ⟩ )) . Let us analyze the error incurred from replacing f ρ be ˜ f ρ and f ¯ ρ by ˜ f ¯ ρ . Observe that

We can simplify the second term on the right-hand side by expanding f ρ and ˜ f ρ :

To simplify the inner expectation, we use Proposition I.5 to find that φ B ( ⟨ u 2 , x ⟩ ) = ⟨ u 2 , x ⟩ with probability at least 1 -2 exp( -B 2 d 2 ) . Thus,

Similarly, we can show that

and therefore,

By the same argument, we can show that

Therefore, it suffices to show that E x ∼ S d -1 [( ˜ f ρ ( x ) -˜ f ¯ ρ ( x )) 2 ] is small with high probability over u 1 , . . . , u m . Observe that for | t | ≤ B , we have | P 2 ,d ( t ) | ≤ B 2 + O (1 /d ) , and | P 4 ,d ( t ) | ≤ B 4 + O ( B 2 /d ) + O (1 /d 2 ) . Thus, if we choose B ≥ 1 / √ d , then for | t | ≤ B we have | P 2 ,d ( t ) | ≲ B 2 and | P 4 ,d ( t ) | ≲ B 4 . In particular, we have

Now, letting x ∈ S d -1 , and using Hoeffding's inequality, we find that if u 1 , . . . , u m are i.i.d. samples from ρ , then

with probability at least 1 -2 exp ( -Cmt 2 ˆ σ 2 2 ,d N 2 ,d B 4 +ˆ σ 2 4 ,d N 4 ,d B 8 ) for some universal constant C . Now, by Corollary 4.2.13 of Vershynin [74], for any γ > 0 , there exists N ⊂ S d -1 of size at most ( C γ ) d such that for any x ∈ S d -1 , there exists y ∈ S d -1 such that ∥ x -y ∥ 2 ≤ γ . Thus, with probability at least 1 -2 ( C γ ) d exp ( -Cmt 2 ˆ σ 2 2 ,d N 2 ,d B 4 +ˆ σ 2 4 ,d N 4 ,d B 8 ) , for all y ∈ N , we have

If Eq. (E.199) holds for any y ∈ N , then for any x ∈ S d -1 , since | σ ′ ( t ) | ≲ | ˆ σ 2 ,d | √ N 2 ,d + | ˆ σ 4 ,d | √ N 4 ,d for any t ∈ [ -1 , 1] (see Lemma E.12), if we let y ∈ N such that ∥ x -y ∥ 2 ≤ γ , then we have

Combining Eq. (E.201), Eq. (E.195) and Eq. (E.196), we have that

with probability at least 1 -2 ( C γ ) d exp ( -Cmt 2 ˆ σ 2 2 ,d N 2 ,d B 4 +ˆ σ 2 4 ,d N 4 ,d B 8 ) . Choosing B = C ′ (log d ) 2 √ d for a sufficiently large universal constant C ′ , and using the fact that N 2 ,d ≲ d 2 and N 4 ,d ≲ d 4 (by Eq. (2.10) of Atkinson and Han [12]) we have that

with probability at least 1 -2 ( C γ ) d exp ( -Ω ( mt 2 (ˆ σ 2 2 ,d +ˆ σ 2 4 ,d )(log d ) 16 )) . We can rearrange the failure probability to find that, up to constant factors, it is at most exp ( -Ω ( mt 2 (ˆ σ 2 2 ,d +ˆ σ 2 4 ,d )(log d ) 16 ) + d log C γ ) . Finally, choosing γ = 1 d Ω(log d ) , and choosing t so that the failure probability is at most δ , we obtain

with probability at least 1 -δ , as desired.

Lemma E.14 (Bound on Gradient) . Let u ∈ S d -1 and ρ be a distribution with support on S d -1 . In addition, let ∇ u L ( ρ ) and grad u L ( ρ ) be defined as in Eq. (3.4) . Then, we have ∥∇ u L ( ρ ) ∥ 2 ≲ (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) d 4 and ∥ grad u L ( ρ ) ∥ 2 ≲ (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) d 4 .

Similarly, let ∇ u ̂ L ( ρ ) and grad u ̂ L ( ρ ) be defined as in Eq. (3.5) . Then, we have ∥∇ u ̂ L ( ρ ) ∥ 2 ≲ (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) d 4 and ∥ grad u ̂ L ( ρ ) ∥ 2 ≲ (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) d 4 .

Proof of Lemma E.14. First, we will bound ∇ u L ( ρ ) for any u ∈ S d -1 and distribution ρ with support in S d -1 :

By Lemma E.12, for any u, x ∈ S d -1 , we have | σ ′ ( ⟨ u, x ⟩ ) | ≲ | ˆ σ 2 ,d | √ N 2 ,d + | ˆ σ 4 ,d | √ N 4 ,d . Additionally, by Eq. (3.1) and by the definition of y ( x ) , we have | f ρ ( x ) -y ( x ) | ≤ | ˆ σ 2 ,d | √ N 2 ,d + | ˆ σ 4 ,d | √ N 4 ,d by Lemma E.12. Thus,

because N 2 ,d ≲ d 2 and N 4 ,d ≲ d 4 by Eq. (2.10) of Atkinson and Han [12]. By the same argument, for any u ∈ S d -1 ,

because I -uu ⊤ is an orthogonal projection matrix. This completes the proof - the proof of the analogous statements for ∇ u ̂ L ( ρ ) and grad u ̂ L ( ρ ) follow from the same arguments.

Lemma E.15. Let u, u ′ ∈ S d -1 . Then, we have

and

Proof of Lemma E.15. For convenience, let δ := u ′ -u . Then, we have

(By definition of σ and P k,d and b.c. ( a + b ) 2 ≲ a 2 + b 2 )

(By Eq. (C.4) and expanding the term with P 4 ,d ′ )

Here the last inequality is due to the identity a 3 -b 3 = ( a -b )( a 2 + ab + b 2 ) , which gives (( u ′⊤ x ) 3 -( u ⊤ x ) 3 ) 2 = (( u ′⊤ x ) 2 + ( u ′⊤ x )( u ⊤ x ) + ( u ⊤ x ) 2 ) 2 ( u ′⊤ x -u ⊤ x ) 2 -expanding the first factor fully gives terms of the form ( u ′⊤ x ) p ( u ⊤ x ) q where p and q are both even and p + q = 4 . Thus, applying Lemma I.17 to the final expression above gives

(By Lemma I.17)

≲

ˆ

σ

2

2

,d

∥

δ

∥

4

2

+ ˆ

σ

2

4

,d

∥

δ

∥

4

2

(B.c. N 2 ,d ≲ d 2 and N 4 ,d ≲ d 4 (see Equation 2.10 of Atkinson and Han [12]))

which completes the proof of the first statement. The proof of the second statement is similar - we expand using the definition of σ and then apply Lemma I.17 to each of the terms. We omit the details of the calculation.

Lemma E.16. Let ρ and ρ ′ be two distributions on S d -1 which both have support size m< ∞ , and consider a coupling ( u, u ′ ) ∼ ( ρ, ρ ′ ) of the two distributions. Additionally, let f ρ and f ρ ′ be defined as in Eq. (3.1) . Then,

Proof of Lemma E.16. We can write the left-hand side as

For any u, u ′ ∈ S d -1 and x ∈ S d -1 , by the mean-value theorem, there exists λ which is a convex combination of ⟨ u, x ⟩ and ⟨ u ′ , x ⟩ such that

Using a similar argument as in the proof of Lemma E.15, we can show that

The proof proceeds by expanding σ ′ ( λ ) into each of its terms and using the inequality | a + b | 3 ≲ | a | 3 + | b | 3 (which holds for any two real numbers a, b ) in order to consider only terms which involve one of u or u ′ -the resulting terms are of the same form as those which appear in Eq. (E.216). We omit the details of the calculation. Thus,

as desired.

F Proof of Theorem 3.5

In this section, we build on our analysis from Appendix E to show that projected gradient descent with 1 poly ( d ) learning rate can achieve low population loss. In the rest of the section, for convenience, we define ˜ δ t = ˜ u t -ˆ u ηt , and omit χ when it is clear from context.

Lemma F.1. Let ρ, ρ ′ be two distributions on S d -1 , and let Γ be a coupling between ρ and ρ ′ . Additionally, let u 1 , u 2 , v ∈ S d -1 . Then, we have

Proof. We can expand the left-hand side as

We first obtain a bound on

By the definition of f ρ , we can rewrite Eq. (F.3) as

By Eq. (C.4), σ is an even fourth-degree polynomial and σ ′ is an odd third-degree polynomial. Thus, the inner expectation can be written as a linear combination of a constant number of terms of the form

where i = 0 , 2 , 4 and j = 1 , 3 . We can also estimate the coefficient which accompanies each individual term Eq. (F.5). We first recall the facts that σ ( t ) = ˆ σ 2 ,d ¯ P 2 ,d ( t ) + ˆ σ 4 ,d ¯ P 4 ,d ( t ) , that N 2 ,d = O ( d 2 ) and N 4 ,d = O ( d 4 ) by Equation (2.10) of Atkinson and Han [12], and that the i th degree term of P 2 ,d has a coefficient which is O ( d i/ 2 -1 ) and the i th degree term of P 4 ,d has a coefficient which is O ( d i/ 2 -2 ) . From these facts it follows that the i th degree term of ¯ P 2 ,d = √ N 2 ,d P 2 ,d has a coefficient which is O ( d i/ 2 ) and the i th degree term of ¯ P 4 ,d = √ N 4 ,d P 4 ,d has a coefficient which is O ( d i/ 2 ) . Thus, each term of the form given in Eq. (F.5) has a coefficient which is at most (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) d ( i + j +1) / 2 in absolute value. We can now upper bound each of the terms of the form given in Eq. (F.5) as follows:

Since the coefficient of this term is at most (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) d ( i + j +1) / 2 in absolute value, this term contributes at most

Thus, taking the expectation over u ′ 1 , u ′ 2 ∼ Γ , we have

A similar but simpler argument can be used to bound

by expanding y ( x ) and σ ′ ( x ) . Thus, we obtain

as desired.

Lemma F.2. Suppose we are in the setting of Theorem 3.5. Assume that n ≳ d 3 (log d ) C for a sufficiently large universal constant C > 0 . Additionally, assume that max χ ∥ ˜ δ t ( χ ) ∥ 2 ≤ 1 √ d . Then, we have

Proof. We can write the left-hand side of Eq. (F.18) as

We proceed by first showing using Lemma I.7 that we do not incur much error when we replace the empirical average over the x i by an expectation over x ∼ S d -1 . We can separate the right hand side of Eq. (F.19) into two parts:

and

We wish to show that S 1 and S 2 are very close to

and

respectively, with high probability - then, we will obtain bounds on M 1 and M 2 . We will show that S 1 and M 1 are close - the proof that S 2 and M 2 are close is similar.

Expanding into Monomials Our proof strategy is to expand the following quantity, for x ∈ S d -1 :

and obtain concentration bounds for each of the terms. In this proof, we let Γ denote the natural coupling between ¯ ρ, ˆ ρ, ˜ ρ which is the distribution over triples (¯ u ηt ( χ ) , ˆ u ηt ( χ ) , ˜ u ηt ( χ )) , where χ is sampled from { χ 1 , . . . , χ m } . Then, we can write G ( x ) as

For convenience, let us define

Clearly we have G ( x ) = E ¯ u ′ ηt , ˆ u ′ ηt , ˜ u ′ t ∼ Γ [ g ( x )] . Now, we consider the effect of expanding g ( x ) by expanding σ and σ ′ into monomials that involve (˜ u ′ t · x ) , (˜ u t · x ) , (ˆ u ′ ηt · x ) and (ˆ u ηt · x ) . Since σ

is an even polynomial of degree 4 and σ ′ is an odd polynomial of degree 3 , we can write g ( x ) as a linear combination of terms of the following form:

where i 1 = 0 , 2 , 4 and i 2 = 1 , 3 - each such term also has a constant factor of O (ˆ σ i 1 ,d ˆ σ i 2 ,d √ N i 1 ,d √ N i 2 ,d ) . We now consider the contribution of Eq. (F.28) for each value of i 1 , i 2 . In the case where i 1 = 0 , we can rewrite Eq. (F.28) as ( ˜ δ t · x ) 2 when i 2 = 1 , and

when i 2 = 3 . In the case where i 1 and i 2 are both nonzero, letting ˜ δ ′ t = ˜ u ′ t -ˆ u ′ ηt we can write Eq. (F.28) as

In summary, Eq. (F.28) can be written as a linear combination of terms of the form

where i 3 + i 4 = 1 , 3 and i 2 = 1 , 3 , and terms of the form

where i 1 = 0 , 2 , 4 and i 3 + i 4 = 0 , 2 .

Contribution of Eq. (F.38) to Concentration Error In Eq. (F.38), we can write ˜ u t = ˆ u ηt + ˜ δ t , and thus write Eq. (F.38) as a linear combination of terms of the form

where i 1 = 0 , 2 , 4 and i 2 + i 3 = 0 , 2 . Next, in order to use Lemma I.7, we write ˆ u ηt = ¯ u ηt + δ t and expand (recalling the definition of δ t from Lemma 5.4). Thus, Eq. (F.38) is a linear combination of terms of the form

where i 1 + i 2 = 0 , 2 , 4 , and i 3 + i 4 + i 5 = 0 , 2 . Finally, we obtain a concentration bound for Eq. (F.40) using Lemma I.7:

Here, when applying Lemma I.7, ¯ u ηt and ¯ u ′ ηt take the role of w 1 and w 2 , while δ ′ ηt / ∥ δ ′ ηt ∥ 2 , δ ηt / ∥ δ ηt ∥ 2 and ˜ δ t / ∥ ˜ δ t ∥ 2 take the role of u 1 , u 2 , u 3 and u 4 (we let u 3 = u 4 = ˜ δ t / ∥ ˜ δ t ∥ 2 and choose p 3 and p 4 so that p 3 + p 4 = 2 + i 5 ). Note that we implicitly perform a union bound over ¯ u ′ ηt and δ ′ ηt drawn from the support of f ˆ ρ ηt -this does not affect the failure probability, since the failure probability for the result of Lemma I.7 is 1 d Ω(log d ) , and it is assumed in Theorem 3.5 that m ≤ d C for some universal constant C . When applying Lemma I.7 at later points in this proof, we implicitly perform this union bound.

In order to compute the final contribution of each such term to the concentration error between S 1 and M 1 , we must first (1) incorporate the coefficients of σ and σ ′ with respect to the Legendre polynomials P 2 ,d and P 4 ,d , (2) take the average of ∥ δ ′ ηt ∥ b 2 with respect to χ (which is implicit throughout this proof), and (3) use the fact that n ≳ d 3 (log d ) O (1) .

Recall that σ = ˆ σ 2 ,d √ N 2 ,d P 2 ,d + ˆ σ 4 ,d √ N 4 ,d P 4 ,d . Using Eq. (C.4) and the fact that N 2 ,d ≍ d 2 and N 4 ,d ≍ d 4 (by Equation (2.10) of Atkinson and Han [12]), we find that, up to constant factors, the coefficient of the 0 th -order term in σ ( t ) can be bounded above by max(ˆ σ 2 ,d , ˆ σ 4 ,d ) , the coefficient of the 2 nd -order term in σ ( t ) can be bounded above by max(ˆ σ 2 ,d , ˆ σ 4 ,d ) d , and the coefficient of the 4 th order term can be bounded above by ˆ σ 4 ,d d 2 . Similarly, up to constant factors, the coefficient of the 1 st -order term in σ ′ ( t ) can be bounded above by max(ˆ σ 2 ,d , ˆ σ 4 ,d ) d and the coefficient of the 3 rd -order term in σ ′ ( t ) can be bounded above by max(ˆ σ 2 ,d , ˆ σ 4 ,d ) d 2 . Thus, each term of degree k originating from σ contributes a factor of max(ˆ σ 2 ,d , ˆ σ 4 ,d ) d k/ 2 to the coefficient of Eq. (F.38), and each term of degree k -1 originating from σ ′ contributes a factor of max(ˆ σ 2 ,d , ˆ σ 4 ,d ) d k/ 2 to the coefficient of Eq. (F.38). Finally, for each term in Eq. (F.40) of total degree i 1 + i 2 + i 3 + i 4 + i 5 +2 , note that σ and σ ′ together contribute a factor of degree i 1 + i 2 + i 3 + i 4 + i 5 +1 , since in Eq. (F.27), aside from the factors contributed by σ and σ ′ , only a single factor of ˜ δ t · x is present.

Thus, to obtain the final contribution of Eq. (F.41) to the overall error between M 1 and S 1 , we multiply by an additional coefficient of max(ˆ σ 2 2 ,d , ˆ σ 2 4 ,d ) d ( i 1 + i 2 + i 3 + i 4 + i 5 ) / 2+1 . Taking this into account, the contribution of Eq. (F.41) is

Next, we simplify this using casework on i 2 . For the case i 2 ≤ 1 , since i 4 + i 5 ≤ 2 , the bound is at most

since d ( i 2 + i 4 + i 5 ) / 2+1 ≤ d 2 . 5 , and because n ≳ d 3 (log d ) O (1) , and we eliminate a factor of ∥ δ ′ ηt ∥ i 2 2 ∥ δ ηt ∥ i 4 2 ∥ ˜ δ t ∥ i 5 2 because all of the particles have ℓ 2 norm at most 1 . Additionally, for the case i 2 ≥ 2 , after taking the average of ∥ δ ′ ηt ∥ i 2 2 ≲ ∥ δ ′ ηt ∥ 2 2 across all χ , the bound is at most

where we have ∆ ηt ≤ 1 √ d because by Lemma E.3, we have that Assumption 5.3 holds as long as ηt ≤ T ∗ ,ϵ . In summary, we have shown that terms of the form Eq. (F.38) contribute at most (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) ∥ ˜ δ t ∥ 2 2 , up to constant factors, to the concentration error between S 1 and M 2 .

Contribution of Eq. (F.37) to Concentration Error We expand Eq. (F.37) by using ˜ u t = ˜ δ t + δ ηt + ¯ u ηt , ˆ u ηt = δ ηt + ¯ u ηt and ˜ u ′ t = ˜ δ ′ t + δ ′ ηt + ¯ u ′ ηt . Thus, we can write Eq. (F.37) as a linear combination of terms of the form

where i 1 + i 2 + i 3 + i 4 + i 5 + i 6 = 6 (since the degrees of each of the monomials must be 8 ) and i 3 ≤ 3 (since each of these terms originates from (˜ u ′ t · x ) i 3 (ˆ u ′ ηt · x ) i 4 (˜ u t · x ) i 2 ( ˜ δ t · x )( ˜ δ ′ t · x ) where i 2 ≤ 3 ). Using Lemma I.7, we obtain a concentration bound for Eq. (F.47):

As discussed in the previous case, after we consider the coefficients of σ ( t ) and σ ′ ( t ) , we find that the term in Eq. (F.47) contributes

Now, suppose i 5 + i 6 ≥ 2 . Then, i 1 + i 2 + i 3 + i 4 ≤ 4 , meaning the contribution to the concentration error is at most

using the fact that n ≳ d 3 (log d ) C for a sufficiently large constant C , and ∥ ˜ δ t ∥ 2 ≲ 1 and ∥ δ ηt ∥ 2 ≲ 1 . On the other hand, if i 5 + i 6 ≤ 1 , then i 1 + i 2 + i 3 + i 4 ≥ 5 , and since i 3 ≤ 3 , this implies that i 1 + i 2 + i 4 ≥ 2 . In this case, the bound is at most

(B.c. i 1 + i 2 + i 3 + i 4 ≤ 6 and n ≳ d 3 (log d ) C )

1

+

i

∥

/

2

δ

′

ηt

∥

i

4

2

d.

Thus, if i 1 + i 2 ≥ 2 , then the contribution of this term to the overall concentration error is at most (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) max χ ∥ ˜ δ t ( χ ) ∥ 2 2 . On the other hand, if i 4 = 1 , then we can take the average of ∥ δ ′ ηt ∥ 2 across χ , and by Lemma E.3, we have that Assumption 5.3 holds for ηt ≤ T ∗ ,ϵ , meaning that E χ ∥ δ ′ ηt ∥ 2 ≤ ∆ ηt ≤ 1 √ d . In this case, the contribution of this term to the overall concentration error is at most (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) max χ ∥ ˜ δ t ( χ ) ∥ 2 2 · 1 d ( i 1 + i 2 + i 4 ) / 2 d = (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) max χ ∥ ˜ δ t ( χ ) ∥ 2 2 . Finally, when i 4 ≥ 2 , we can take the average of ∥ δ ′ ηt ∥ 2 2 over χ -wehave E χ ∥ δ ′ ηt ∥ 2 2 ≤ 1 d , and therefore the contribution of this term to the overall concentration error is at most

In all cases, the contribution of Eq. (F.37) to the overall concentration error between S 1 and M 1 is at most

≲

(ˆ

σ

2

2

,d

+ ˆ

σ

2

4

,d

Overall Concentration Error

In summary, the overall concentration error between S 1 and M 1 is that of G ( x ) . We have expanded G ( x ) into various monomials of dot products involving x , and shown that each monomial contributes at most (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) max χ ∥ ˜ δ t ( χ ) ∥ 2 2 to the concentration error. Thus, with high probability,

χ

∥

˜

δ

t

(

χ

)

∥

2

2

·

d

(

i

1

2

)

In order to bound | M 2 -S 2 | , an almost identical, but simpler, argument can be used, with y ( x ) in the place of f ˜ ρ t ( x ) or f ˆ ρ ηt ( x ) . In the modified argument, ˜ u ′ t , ˆ u ′ ηt and ¯ u ′ ηt will all be replaced with e 1 . Thus, in the analysis of Eq. (F.40), we can replace ¯ u ′ ηt with e 1 and assume that i 2 = 0 (since otherwise this term becomes 0 due to the δ ′ ηt · x factor) - we can then proceed similarly to the subcase where i 2 ≤ 1 in the analysis of Eq. (F.40). Additionally, the term analogous to Eq. (F.47) does not arise in the case where f ˜ ρ t ( x ) and f ˆ ρ ηt ( x ) are replaced by y ( x ) , since ˜ u ′ t and ˆ u ′ ηt are both replaced by e 1 , so the corresponding term is 0 . We also note that the Legendre coefficients of y ( x ) have absolute values which are less than those of σ . Thus, we have

Bounding the Expectation To complete the proof, it suffices to bound | M 1 -M 2 | . We have that

and by Lemma F.1, we have

Combining this with our bound on | M 1 -S 1 | + | M 2 -S 2 | , we find that

as desired.

Lemma F.3. Suppose we are in the setting of Theorem 3.5. Suppose s, t ≥ 0 . Then,

Proof. We have

We bound the right-hand side by several applications of the triangle inequality. First, observe that for any x ∈ S d -1 , we have

| f ˆ ρ t ( x ) -f ˆ ρ s ( x ) | ≤ 1 m ∑ j =1 | σ (ˆ u t ( χ i ) · x ) -σ (ˆ u s ( χ i ) · x ) | (F.69) ≲ 1 m m ∑ j =1 ( | ˆ σ 2 ,d | √ N 2 ,d + | ˆ σ 4 ,d | √ N 4 ,d ) | ˆ u t ( χ i ) · x -ˆ u s ( χ i ) · x | (By Lemma E.12) ≲ ( | ˆ σ 2 ,d | √ N 2 ,d + | ˆ σ 4 ,d | √ N 4 ,d ) m m ∑ j =1 ∥ ˆ u t ( χ i ) -ˆ u s ( χ i ) ∥ 2 (F.70) ≲ ( | ˆ σ 2 ,d | √ N 2 ,d + | ˆ σ 4 ,d | √ N 4 ,d ) max i ∈ [ m ] ∥ ˆ u t ( χ i ) -ˆ u s ( χ i ) ∥ 2 . (F.71)

Additionally, for ˆ u t and ˆ u s with the same initialization χ , we have

Also, for ˆ u t and ˆ u s with the same initialization χ , we have

Finally, for x ∈ S d -1 and any distribution ρ on S d -1 , we have

since we can bound | f ρ ( x ) | using Lemma E.12, and because (ignoring ˆ h 0 ,d since it is equal to ˆ σ 0 ,d ) y ( x ) = γ 2 ˆ σ 2 ,d P 2 ,d + γ 4 ˆ σ 4 ,d P 4 ,d , and using Assumption 3.2.

Thus, we can simplify the right-hand side of Eq. (F.68) using the triangle inequality, and combining Eq. (F.71), Eq. (F.73), Eq. (F.76) and Eq. (F.77):

Therefore, Eq. (F.68) gives

as desired.

Proof of Theorem 3.5. To show that projected gradient descent achieves low population loss, we show that it does not diverge far from projected gradient descent. Suppose ˆ ρ 0 and ˜ ρ 0 are both equal to unif ( χ 1 , . . . , χ m } ) . Then, by induction over t , we bound the difference between ˜ u t ( χ ) and ˆ u ηt ( χ ) for all χ ∈ { χ 1 , . . . , χ m } . At t = 0 , we have ˜ u 0 ( χ ) = ˆ u 0 ( χ ) . Let t be a nonnegative integer, and assume that for all s ≤ t we have max χ ∥ ˜ δ t ( χ ) ∥ 2 ≤ 1 √ d .

By the definition of projected gradient descent, we have

Note that ∥ ˜ u t -η · grad ˜ u t ̂ L (˜ ρ t ) ∥ 2 ≥ ∥ ˜ u t ∥ 2 , since grad ˜ u t ̂ L (˜ ρ t ) is orthogonal to ˜ u t . Thus, letting Π denote the projection onto the ℓ 2 unit ball, we have

Also, we have

Thus,

Further simplifying using the Cauchy-Schwarz inequality, we have

Thus, we have max χ ∥ ˜ δ t +1 ( χ ) ∥ 2 2 ≤ ( 1 + O ( η (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) )) max χ ∥ ˜ δ t ( χ ) ∥ 2 2 + O ( η 2 (ˆ σ 2 2 ,d + ˆ σ 2 4 ,d ) 2 d 8 ) (F.111)

which means that for some universal constant C ,

Thus, as long as η = 1 (ˆ σ 2 2 ,d +ˆ σ 2 4 ,d ) d B for a sufficiently large universal constant B , and choosing η so that T ∗ ,ϵ is an integer multiple of η , we have max χ ∥ ˜ δ t +1 ( χ ) ∥ 2 ≤ 1 √ d , because if t ≤ T ∗ ,ϵ η , then

This completes the induction, and therefore for all t ≤ T ∗ ,ϵ η , we have max χ ∥ ˜ δ t ( χ ) ∥ 2 ≤ 1 √ d . Additionally, by Lemma E.16, and using the coupling between ˜ ρ t and ˆ ρ ηt which is the joint distribution of (˜ u t ( χ ) , ˆ u ηt ( χ )) , we have

which completes the proof.

G Missing Proofs in Section 3

G.1 Proof of Lemma 3.1

In this section, we prove sufficient and necessary conditions for γ 2 and γ 4 when the population loss is 0.

̸

Proposition G.1 (Necessary condition) . Let f ρ ( x ) = E u ∼ ρ [ σ ( u ⊤ x )] be a (not necessarily symmetric nor rotational invariant) two-layer neural network and y ( x ) = h ( e ⊤ 1 x ) the target function. Define γ k = ˆ h k,d ˆ σ k,d , ∀ k ≥ 0 . Suppose E x ∼ S d -1 ( f ρ ( x ) -y ( x )) 2 = 0 and ˆ σ 2 ,d = 0 , ˆ σ 4 ,d = 0 , we have

̸

Proof. Similar to the proof of Lemma 4.2, we have

Continuing the equation by invoking Lemma C.5, we get

Therefore, E x ∼ S d -1 [( f ρ ( x ) -h ( x ⊤ e 1 )) 2 ] = 0 implies that

First we prove Eq. (G.3). By Lemma C.6, there is a feature mapping ϕ : S d -1 → R N 2 ,d such that for every u, v ∈ S d -1 ,

Rewrite Eq. (G.10) using ϕ 2 we get

Or equivalently,

As a result,

Similarly, we also have γ 4 ≤ 1 .

Now we prove Eq. (G.1). Recall that

Combining Eq. (G.10) and Eq. (G.17) we have

Consequently,

Similarly, combining Eq. (G.11) and Eq. (G.18) we get

Consequently,

Now for any fixed ξ ∈ S d -1 , our key observation is that by Jensen's inequality,

On the one hand, by Eq. (G.22) we have

On the other hand, by Eq. (G.26) we have

Finally, combining Eqs (G.27)-(G.31) and set ξ = e 1 we get

It follows directly that

Recall that

Finally we prove Eq. (G.2). For any ξ, u ∈ S d -1 we have | u ⊤ ξ | ≤ 1 . As a result,

Combining with Eq. (G.29) and Eq. (G.31) and setting ξ = e 1 we get

Proposition G.2 (Sufficient condition) . Let f ρ ( x ) = E u ∼ ρ [ σ ( u ⊤ x )] be a two-layer neural network and y ( x ) = h ( e ⊤ 1 x ) the target function. Let γ k = ˆ h k,d ˆ σ k,d , ∀ k ≥ 0 . Suppose ˆ h k,d = 0 , ∀ k ̸∈ { 0 , 2 , 4 } , ˆ h 0 ,d = ˆ σ 0 ,d , σ is a degree-4 polynomial, and

Then there exists a symmetric and rotational invariant neural network ρ such that E x ∼ S d -1 ( f ρ ( x ) -y ( x )) 2 = 0 .

Proof. Recall that by Lemma 4.2, for every symmetric and rotational invariant neural network ρ we have

where u = ( w,z ) , w ∈ [ -1 , 1] , and z ∈ √ 1 -w 2 S d -2 . Therefore, we only need to show that there exists a one-dimensional distribution µ such that

Consequently, Eq. (G.40) and Eq. (G.41) is equivalent to

Let β 2 = d -1 d γ 2 + 1 d and β 4 = d 2 -1 ( d +2)( d +4) γ 4 + 6 d +12 ( d +2)( d +4) ( d -1 d γ 2 + 1 d ) -3 ( d +2)( d +4) . Then we have | β 2 -γ 2 | ≤ O (1 /d ) and | β 4 -γ 4 | ≤ O (1 /d ) . Rewriting Eqs. (G.36)-(G.38) using β 2 and β 4 , we have

By Proposition I.14, there exists a distribution µ ⋆ such that

Hence, by setting the marginal distribution of w equal to µ ⋆ , we have

which proves the desired result.

H Proofs for Sample Complexity Lower Bounds for Kernel Method

In this section, we present a sample complexity lower bound for kernel methods with any inner product kernel. That is, the kernel K : S d -1 × S d -1 → R can be written as K ( x, z ) = κ ( x ⊤ z ) for some one-dimensional function κ : [ -1 , 1] → R . Inner product kernels are rotational invariant and do not specialize to any coordinate. In particular, NTK of two-layer neural networks is an inner product kernel (e.g., Du et al. [30], Wei et al. [76]).

To prove Theorem 3.6, we only need to work on the degreek spherical harmonics component of the target function and its estimation because, by the orthogonality of spherical harmonics, any lower bound for the degreek components directly translates to the original setting. The following lemma states the lower bound for general k ≥ 0 , while Theorem 3.6 only requires the case k = 4 .

Lemma H.1. Fixed any k ≥ 1 , let g ⋆ ( x ) = P k,d ( u ⊤ x ) for some fixed u ∈ S d -1 . When d is larger than a sufficiently large universal constant and n < d k (8 k ) -(3 k +1) (ln d ) -( k +2) , with probability at least 1 / 2 over the randomness of n i.i.d. data { x 1 , · · · , x n } drawn uniformly from S d -1 , any estimation g of the form g ( x ) = ∑ n i =1 β i P k,d ( x ⊤ x i ) must have a constant error:

Proof of Lemma H.1 is deferred to Appendix H.2, where our key observation is that P k,d ( ⟨ x i , ·⟩ ) cannot correlate too much with g ⋆ when x i is uniformly randomly drawn from the unit sphere. Indeed, 〈 P k,d ( ⟨ x i , ·⟩ ) , g ⋆ 〉 = P k,d ( x ⊤ i u ) concentrates within ± 1 /N k,d (Lemma C.7). Hence, we can upperbound ⟨ g, g ⋆ ⟩ = ∑ n i =1 β i P k,d ( x ⊤ i u ) ≲ ∥ β ∥ 2 √ n/N k,d . However, we also have ∥ g ∥ 2 2 = ∑ i,j ∈ [ n ] β i β j P k,d ( x ⊤ i x j ) ≳ ∥ β ∥ 2 2 because the matrix ( P k,d ( x ⊤ i x j )) i,j ∈ [ n ] concentrates around I (Lemma H.3). Since ∥ g -g ⋆ ∥ 2 2 < 3 / 4 requires ⟨ g, g ⋆ ⟩ > ∥ g ∥ 2 / 2 , we must have n ≳ N k,d ≈ d -k .

H.1 Proof of Theorem 3.6

The following theorem states that for every k ≥ 0 , to achieve a constant population error, we need at least C k d -k (ln d ) -( k +2) samples where C k = (8 k ) -(3 k +1) is a constant that only depends on k . Hence, Theorem 3.6 is a direct corollary of Theorem H.2 by taking k = 4 .

Theorem H.2. Let y : S d -1 → R be a function of the form y ( x ) = h ( q ⊤ ⋆ x ) where q ⋆ ∈ S d -1 is a fixed vector and h : [ -1 , 1] → R is a one-dimensional function. Let K : S d -1 × S d -1 → R be any inner product kernel. When d is larger than a universal constant and n < d k (8 k ) -(3 k +1) (ln d ) -( k +2) , with probability at least 1 / 2 over the randomness of n i.i.d. data points { x i } n i =1 , any estimation f of the form f ( x ) = ∑ n i =1 β i K ( x i , x ) must have a significant error:

Proof of Theorem H.2. In the following we first fix a degree k ≥ 1 . For every l > 0 , let Y l,d be the space of degreel spherical harmonics and Π l the projection to Y l,d . By the orthogonality of Y l,d , we have

By Atkinson and Han [12, Section 2.3], the projection operator Π k is given by

Consequently, we have (Π k y )( x ) = √ N k,d E ξ ∼ S d -1 [ P k,d ( x ⊤ ξ ) h ( q ⊤ ⋆ ξ )] . Recall that

is the density of u 1 when u = ( u 1 , · · · , u d ) is drawn uniformly from S d -1 , and ˆ h k,d = E t ∼ µ d [ P k,d ( t ) h ( t )] . By Funk-Hecke formula (Theorem C.1) we get

Similarly, we can write the inner product kernel K as K ( x, z ) = κ ( x ⊤ z ) . Define ˆ κ k,d = E t ∼ µ d [ P k,d ( t ) κ ( t )] . Then we have

Now we apply Lemma H.1 to the function ( ˆ h k,d ) -1 (Π k y ) = P k,d ( ⟨ q ⋆ , ·⟩ ) . As a result, with probability at least 1 / 2 over the randomness of { x i } n i =1 , for any function f of the form f ( · ) = ∑ n i =1 β i K ( x i , · ) , we have ∥ ( ˆ h k,d ) -1 Π 4 ( y -f ) ∥ 2 2 ≥ 3 / 4 , which implies ∥ y -f ∥ 2 2 ≥ 3 4 ( ˆ h k,d ) 2 .

H.2 Proof of Lemma H.1

Proof of Lemma H.1. First we prove that when ∥ g ⋆ ∥ 2 = 1 , ∥ g -g ⋆ ∥ 2 2 < 3 / 4 implies ⟨ g, g ⋆ ⟩ ≥ ∥ g ∥ 2 / 2 . To this end, by basic algebra we get

where the last inequality follows from AM-GM. In the following, we prove that with probability at least 1 / 2 , Eq. (H.8) is impossible when n < d k (8 k ) -(3 k +1) (ln d ) -( k +2) .

Recall that Lemma C.5 states that for every x, z ∈ S d -1 and k ≥ 0 ,

Consequently,

As a result,

Invoking Lemma H.3, with probability at least 3 / 4 over the randomness of { x i } i ∈ [ n ] , we have

For the LHS of Eq. (H.8),

By Lemma C.7, when n < d k (8 k ) -(3 k +1) (ln d ) -( k +2) , with probability at least 3 / 4 we have

Combining Eq. (H.11) and Eq. (H.15), there exists a universal constant d 0 > 0 such that when d > d 0 , with probability at least 1 / 2 over the randomness of { x 1 , · · · , x n } ,

Recall that ∥ g -g ⋆ ∥ 2 2 < 3 / 4 implies ⟨ g, g ⋆ ⟩ ≥ ∥ g ∥ 2 / 2 . Consequently, with probability at least 1 / 2 , when n < d k (8 k ) -(3 k +1) (ln d ) -( k +2) and d > d 0 we have ∥ g -g ⋆ ∥ 2 2 ≥ 3 / 4 .

Lemma H.3. Let x 1 , · · · , x n be i.i.d. random variables drawn uniformly from S d -1 , and P k,d the degreek Legendre polynomial. For any fixed k ≥ 1 , define the matrix M k ∈ R n × n where [ M k ] i,j = P k,d ( ⟨ x i , x j ⟩ ) . There exists universal constants d 0 such that, when n < d k (8 k ) -(3 k +1) (ln d ) -( k +2) and d > d 0 , with probability at least 3 / 4 ,

Proof. In the following, we fix k ≥ 1 . First, we prove that

To this end, we invoke Theorem I.4 by constructing a matrix A such that 1 N k,d A ⊤ A = M k .

By Lemma C.6, there exists a feature mapping ϕ : S d -1 → R N k,d such that for every u, v ∈ S d -1

We set A i = ϕ ( x i ) , ∀ 1 ≤ i ≤ n . Consequently,

In the following, we upper bound the incoherence parameter p defined in Theorem I.4:

̸

By Eq. (H.19) we get

̸

It follows that p ≲ n (32 k ln( dn )) k .

Now that we have established the conditions, by Theorem I.4 we have

Note that N k,d ≥ ( k + d -2 k ) ≥ ( d/k ) k when k ≥ 2 , and clearly N 1 ,d ≥ d . Hence, there exists a universal constant d 0 such that, when n < d k (8 k ) -(3 k +1) (ln d ) -( k +2) .

Hence, there exists a universal constant d 0 such that when d ≥ d 0 ,

Now we prove Eq. (H.17). Let t = 1 -λ min ( M ) be a random variable. Note that by Eq. (H.18) we have E [ t ] ≤ 1 / 8 . By basic algebra, we also have

which implies t ≥ 0 . Therefore, by Markov inequality we get

I Toolbox

I.1 Rademacher Complexity and Generalization Bounds

Lemma I.1 (Special Case of Contraction Principle (Lemma 4.6 of Adamczak et al. [4])) . Let F be a family of functions from D to R , and let φ be an L -Lipschitz function with φ (0) = 0 . Then, for any S = { x 1 , . . . , x N } ⊂ D ,

where the expectations are taken over i.i.d. Rademacher random variables ϵ 1 , . . . , ϵ N .

In the following, for a function family H , we let R N ( H ) denote its Rademacher complexity and R S ( H ) denote its empirical Rademacher complexity (given a fixed dataset S ) - see Section 4.4.2 of Ma [53] for a definition of these concepts.

Lemma I.2 (Empirical Rademacher Complexity of Products) . Suppose F and G are two families of functions from D to R , with D ⊂ R d , which are uniformly bounded by B > 0 . Then, for any set S = { x 1 , . . . , x N } ⊂ D ,

where F · G = { fg | f ∈ F , g ∈ G} . As a corollary,

̸

̸

̸

Proof of Lemma I.2. We can write

Using Lemma I.1 and the fact that the function x ↦→ x 2 is 2 B Lipschitz on the interval [ -B,B ] , we have

as desired.

LemmaI.3 (Symmetrization - Analogous to Corollary 4.7 of Adamczak et al. [4]) . Let X 1 , . . . , X N be independent random variables. Let F be a family of functions from D ⊂ R d to R that is uniformly bounded by B 1 ≥ 1 , and let G be a family of functions from D ⊂ R d to R . Then, for nonnegative integers p 1 , p 2 , p 3 , p 4 and a fixed g ∈ G that is bounded by B 2 ≥ 1 , we have

Proof of Lemma I.3. By Theorem 4.13 of Ma [53],

Finally, we simplify R N ( F p i ) for each i . For nonnegative integers p , since F is uniformly bounded by B 1 , the function x ↦→ x p is pB p -1 1 -Lipschitz on [ -B 1 , B 1 ] for p ≥ 1 and is 1 -Lipschitz for p = 0 . Thus, by Lemma I.1, we have R N ( F p ) ≤ min(1 , pB p -1 1 ) R N ( F ) , and

as desired.

I.2 Concentration Lemmas

Theorem I.4 (Theorem 5.62 of Vershynin [73]) . Let A be an m × n matrix with independent columns { A j } n j =1 and m ≥ n . Suppose A j is isotropic, i.e., E [ A j A ⊤ j ] = I and ∥ A j ∥ 2 = √ m for every j ∈ [ n ] . Let p be the incoherence parameter defined by

̸

Then there is a universal constant C 0 such that

The following proposition gives the tail bound for inner products of random vectors in S d -1 (c.f. Theorem 3.4.6 of Vershynin [74]).

Proposition I.5. Let u be a vector drawn uniformly at random from the unit sphere S d -1 . For any fixed vector v ∈ S d -1 , we have

Proof of Proposition I.5. By the symmetricity of u and v , we can assume that v = (1 , 0 , · · · , 0) without loss of generality. Let x = 1 2 ( u ⊤ v +1) . Then we have x ∼ Beta ( d -1 2 , d -1 2 ) . Hence, the desired result follows directly from Skorski [63, Theorem 1].

Lemma I.6 (Modification of Proposition 4.4 of Adamczak et al. [4]) . Let X 1 , . . . , X N be i.i.d. random vectors drawn uniformly from √ d S d -1 , and let p 1 , p 2 , p 3 , p 4 , q 1 , q 2 be nonnegative integers. Suppose p 1 + p 2 + p 3 + p 4 ≥ 2 . Then, for fixed w 1 , w 2 ∈ S d -1 , as long as N ≤ d C for any universal constant C > 0 , we have

with probability at least 1 -1 d Ω(log d ) . Here, C p 1 ,p 2 ,p 3 ,p 4 ,q 1 ,q 2 is a sufficiently large constant which depends on p 1 , p 2 , p 3 , p 4 , q 1 , q 2 and C p 1 ,p 2 ,p 3 ,p 4 is a constant which depends on p 1 , p 2 , p 3 , p 4 .

Proof of Lemma I.6. We largely follow the proof of Proposition 4.4 of Adamczak et al. [4], with some minor modifications. First, let E 1 be the event that |⟨ X i , w 1 ⟩| ≤ (log d ) 2 for all i = 1 , . . . , N , and define E 2 analogously for w 2 . By Lemma I.10, the sub-exponential norm of ⟨ X i , w 1 ⟩ is at most an absolute constant K > 0 . Thus, by Proposition 2.7.1 of Vershynin [74], we have

Thus, by a union bound, with probability at least 1 -N d Ω(log d ) , for all i = 1 , . . . , N , we have |⟨ X i , w 1 ⟩| ≤ (log d ) 2 , and |⟨ X i , w 2 ⟩| ≤ (log d ) 2 (by performing a similar union bound using the sub-exponential norm of ⟨ X i , w 2 ⟩ ). In summary, E 1 and E 2 simultaneously hold with probability at least 1 -1 d Ω(log d ) since N ≤ d C for some universal constant C > 0 .

Now, as in Adamczak et al. [4], define B > 1 , which we will specify later - this is the amount by which we will truncate ⟨ X i , u 1 ⟩ , ⟨ X i , u 2 ⟩ , ⟨ X i , u 3 ⟩ and ⟨ X i , u 4 ⟩ . For convenience, define the function φ B : R → R given by

Define B ′ = (log d ) 2 , and define φ B ′ similarly to φ B , but with B replaced by B ′ . If E 1 and E 2 hold, then for all u 1 , u 2 , u 3 , u 4 ∈ S d -1 , we have

Thus, the rest of the proof will focus on obtaining an upper bound on

First define

We will first obtain an upper bound on U ′ , then show that U ′ and U are close. We apply Lemma I.3 with F = { x ↦→ φ B ( ⟨ x, u ⟩ ) | u ∈ S d -1 } , and with f i corresponding to φ B ( ⟨ x, u i ⟩ ) , and finally with g ( x ) = | φ B ′ ( ⟨ x, w 1 ⟩ ) | q 1 | φ B ′ ( ⟨ x, w 2 ⟩ ) | q 2 , to find that

(By Lemma I.1 since φ B is 1 -Lipschitz)

Now, to show a high-probability upper bound on U ′ , we apply Lemma 4.8 of Adamczak et al. [4] (Talagrand's concentration inequality), and we ensure that all of the conditions are satisfied:

We set X 1 , . . . , X N in Lemma 4.8 of Adamczak et al. [4] to be as defined in the statement of Lemma I.6.

For convenience, define

Then, we let

where w 1 and w 2 are fixed (i.e. as defined in the statement of Lemma I.6).

We let a = 2 B p 1 + p 2 + p 3 + p 4 ( B ′ ) q 1 + q 2 , which is a uniform bound on all the functions in F .

Observe that for all f ∈ F , E X ∼ √ d S d -1 f ( X ) = 0 .

Thus, we apply Lemma 4.8 of Adamczak et al. [4] with

to find that

with failure probability at most

Observe that Talagrand's concentration inequality only gives us a bound on sup f ∈F ∑ N i =1 f ( X i ) , while we wish to obtain a high-probability bound on sup f ∈F | ∑ N i =1 f ( X i ) | . However, using the same argument, we can obtain a similar high-probability bound on sup f ∈F ∑ N i =1 -f ( X i ) = -inf f ∈F ∑ N i =1 f ( X i ) . Thus, we obtain a lower bound on inf f ∈F ∑ N i =1 f ( X i ) which is the negative of the upper bound that we obtained for sup f ∈F ∑ N i =1 f ( X i ) , meaning that the upper bound we obtained for sup f ∈F ∑ N i =1 f ( X i ) also holds for sup f ∈F | ∑ N i =1 f ( X i ) | . In summary,

with probability at least 1 -exp( -Ω( √ d )) .

Now that we have obtained a high-probability bound on U ′ , let us analyze the difference between U ′ and U . Recall the definition of g B,B ′ ,u 1 ,u 2 ,u 3 ,u 4 ,w 1 ,w 2 from above. Additionally, for convenience, let g B ′ ,u 1 ,u 2 ,u 3 ,u 4 ,w 1 ,w 2 ( x ) = |⟨ x, u 1 ⟩| p 1 |⟨ x, u 2 ⟩| p 2 |⟨ x, u 3 ⟩| p 3 |⟨ x, u 4 ⟩| p 4 | φ B ′ ( ⟨ x, w 1 ⟩ ) | q 1 | φ B ′ ( ⟨ x, w 2 ⟩ ) | q 2 , i.e. we are simply removing the truncation by B . Then, following the proof of Proposition 4.4 in Adamczak et al. [4], we have

where { C } denotes the event that the condition C holds. Here the inequality holds because the difference g B ′ ,u 1 ,u 2 ,u 3 ,u 4 ,w 1 ,w 2 ( X i ) -g B,B ′ ,u 1 ,u 2 ,u 3 ,u 4 ,w 1 ,w 2 ( X i ) is nonzero if and only if 1 ( ⋃ 4 j =1 {|⟨ X i , u j ⟩| ≥ B } ) is nonzero. We have obtained a high probability bound on U ′ . Let us now obtain a bound on the third term on the right-hand side of Eq. (I.62) (here, C p 1 ,p 2 ,p 3 ,p 4 ,q 1 ,q 2 is a constant that depends on p 1 , p 2 , p 3 , p 4 , q 1 , q 2 ):

(By Cauchy-Schwarz inequality)

To obtain the last inequality, we apply the Cauchy-Schwarz inequality multiple times to write

as a product of moments of |⟨ X i , u j ⟩| and |⟨ X i , w j ⟩| , and by Lemma I.10, each of these moments can be bounded above by a constant that depends on p 1 , p 2 , p 3 , p 4 , q 1 , q 2 . Additionally, by Lemma I.10 and Proposition 2.7.1, part (a) of Vershynin [74], we have P ( |⟨ X i , u j ⟩| ≥ B ) ≲ e -Ω( B ) , since |⟨ X i , u j ⟩| has constant sub-exponential norm. In summary,

To complete the proof, we bound the second term on the right-hand side of Eq. (I.62). Following the proof of Proposition 4.4 of Adamczak et al. [4], we apply Lemma I.11 (a variant of Theorem 3.6 of Adamczak et al. [4]) to find that, with failure probability at most e -Ω( √ d ) , for all m -sparse vectors z ∈ S N -1 , we have

Let E sparse be the event that Eq. (I.72) holds. For convenience, let A ∈ R d × N be the matrix whose i th column is X i . Additionally, for any subset S ⊂ [ N ] , let A S be the matrix whose columns have indices in S . If S ⊂ [ N ] with | S | = m , then assuming E sparse holds, the operator norm of A S is at

where the last inequality is because E sparse holds.

Following the proof of Proposition 4.4 of Adamczak et al. [4], we now define S B = { i ∈ [ N ] | |⟨ X i , u j ⟩| ≥ B for some j ∈ [4] } . Note that S B depends on u 1 , u 2 , u 3 , u 4 , and we will now show that if we select B to be a certain value which is not too large, this will still ensure that S B is small. Observe that

and

Combining these we obtain

Therefore, if we select B ≳ log ( N | S B | ) , we can eliminate the second term on the right-hand side of Eq. (I.79), finding that | S B | ≲ d B 2 . Combining this with Eq. (I.73), if B ≳ log N and E sparse holds we obtain

Raising both sides to the ( p 1 + p 2 + p 3 + p 4 ) th power, we finally obtain

with probability 1 -e -Ω( √ d ) , where C ′ p 1 ,p 2 ,p 3 ,p 4 is a constant which depends on p 1 , p 2 , p 3 , p 4 . In summary, assuming E sparse holds and the high-probability bound on U ′ holds (and these both hold simultaneously with probability at least 1 -exp( -Ω( √ d )) ), we obtain

by Eq. (I.62). Finally, by Eq. (I.33), we have

as long as the events E 1 and E 2 hold, and these events simultaneously hold with probability at least 1 -1 d Ω(log d ) . Thus, by our above bound on U we have

Setting B ′ = (log d ) 2 and B ≳ max(log N, (log d ) 2 ) ≍ (log d ) 2 (recall that N ≤ d C for a universal constant C ), we have

with probability at least 1 -1 d Ω(log d )

, as desired.

Lemma I.7 (Corollary for Uniform Samples from S d -1 ) . Let x 1 , . . . , x N be i.i.d. random vectors sampled uniformly from S d -1 , and let p 1 , p 2 , p 3 , p 4 , q 1 , q 2 be nonnegative integers. Suppose p 1 + p 2 + p 3 + p 4 ≥ 2 . Then, for fixed w 1 , w 2 ∈ S d -1 , as long as N ≤ d C for any universal constant C > 0 , we have

with probability at least 1 -1 d Ω(log d ) . Here, C p 1 ,p 2 ,p 3 ,p 4 ,q 1 ,q 2 is a constant which depends on p 1 , p 2 , p 3 , p 4 , q 1 , q 2 , and C p 1 ,p 2 ,p 3 ,p 4 is a constant which depends on p 1 , p 2 , p 3 , p 4 .

Proof of Lemma I.7. This is a direct corollary of Lemma I.6, obtained by dividing both sides of Lemma I.6 by 1 d ( p 1 + p 2 + p 3 + p 4 + q 1 + q 2 ) / 2 .

Lemma I.8 (Moments for Truncated Dot Products with Random Spherical Vectors) . Let P be a positive integer, and for i = 1 , . . . , P , let u i ∈ S d -1 and p i be a nonnegative integer. Additionally, let B ≥ 1 . Finally, let X be a random vector drawn uniformly from √ d S d -1 . Then, for some absolute constant c > 0 we have

where φ B is defined as

and C p 1 ,...,p P is a constant which depends only on p 1 , . . . , p P .

Proof of Lemma I.8. We use an argument similar to one used in part of the proof of Proposition 4.4 in Adamczak et al. [4]. First, for i ∈ [ P ] , we can write

Thus, in the expression

if we expand the first product using Eq. (I.104), then we obtain terms of the form ∏ P i =1 F i , where F i is either | φ B ( ⟨ X,u i ⟩ ) | p i or ( |⟨ X,u i ⟩| p i -B p i )1( |⟨ X,u i ⟩| ≥ B ) . Observe that the term where F i = | φ B ( ⟨ X,u i ⟩ ) | p i for all i ∈ [ P ] cancels with the second term in Eq. (I.105). Thus, in order to obtain an upper bound on

it suffices to upper bound the terms ∏ P i =1 F i where at least one of the F i is ( |⟨ X,u i ⟩| p i -B p i )1( |⟨ X,u i ⟩| ≥ B ) . Observe however that

Thus, because | φ B ( ⟨ X,u i ⟩ ) | ≤ |⟨ X,u i ⟩| , in order to upper bound the terms ∏ P i =1 F i where at least one of the F i is ( |⟨ X,u i ⟩| p i -B p i )1( |⟨ X,u i ⟩ ≥ B ) , it suffices to obtain an upper bound on the expected value of

Without loss of generality, let us bound the expected value of the first term in the sum above:

where the last inequality is by Lemma I.10 and Proposition 2.7.1 of Vershynin [74]. Finally, since the sub-exponential norm of ⟨ X,u 1 ⟩ is at most an absolute constant by Lemma I.10, we have that

for some absolute constant c > 0 . Thus, by Eq. (I.109), we obtain

for some absolute constant c > 0 . This completes the proof.

Lemma I.9. Let N,d ∈ N . Let X 1 , . . . , X N i.i.d. random vectors drawn uniformly from √ d S d -1 . Additionally, let ϵ 1 , . . . , ϵ N be i.i.d. Rademacher random variables. Then,

Proof of Lemma I.9. We first bound the second moment of the norm:

̸

Lemma I.10 (Sub-Exponential Norm of Uniform Distribution on √ d S d -1 ) . Assume d ≥ C for some sufficiently large universal constant C . Then there exists an absolute constant K > 0 such that, if X is a random vector drawn uniformly from √ d S d -1 and v ∈ S d -1 is a fixed vector, then ∥⟨ X,v ⟩∥ ψ 1 ≤ K , where ∥ · ∥ ψ 1 denotes the sub-exponential norm of a random variable. (See Definition 2.7.5 of Vershynin [74].) As a corollary, for any p ≥ 1 , there exists a constant C p > 0 depending on p such that for any v ∈ S d -1 , E X ∼ √ d S d -1 [ |⟨ X,v ⟩| p ] ≤ C p .

Proof of Lemma I.10. Let X be drawn uniformly at random from the uniform distribution on √ d S d -1 , and let v ∈ S d -1 . Observe that the distribution of 1 √ d ⟨ X,v ⟩ is µ d . Thus, the density of 1 √ d ⟨ X,v ⟩ is p ( t ) = Γ( d/ 2) √ π Γ(( d -1) / 2) (1 -t 2 ) d -3 2 by Eqs. (1.16) and (1.18) of Atkinson and Han [12]. By a change of variables, the density of ⟨ X,v ⟩ is

Thus, for any t ≥ 0 ,

Thus, for some absolute constant C > 0 and all t ≥ 0 ,

Thus, by the definition of sub-Gaussian norm (Definition 2.5.6 and Proposition 2.5.2 of Vershynin [74]), the sub-Gaussian norm of ⟨ X,v ⟩ is bounded by a constant, i.e. for any v ∈ S d -1 , we have ∥⟨ X,v ⟩∥ ψ 2 ≤ K for some absolute constant K > 0 (here ∥ · ∥ ψ 2 denotes the sub-Gaussian norm of a random variable). By Propositions 2.5.2 and 2.7.1, and Definition 2.7.5, of Vershynin [74], this implies that ∥⟨ X,v ⟩∥ ψ 1 ≲ 1 , since sub-Gaussian random variables are always sub-exponential (specifically, using item (ii) of Proposition 2.5.2 and item (b) of Proposition 2.7.1 of Vershynin [74]). The last statement of the lemma follows directly from Proposition 2.7.1, item (b) of Vershynin [74].

Lemma I.11 (Modification of Theorem 3.6 of Adamczak et al. [4]) . Let d ≥ C for some absolute constant C , and 1 ≤ N ≤ e √ d . Let X 1 , . . . , X N be i.i.d. random vectors sampled from the uniform distribution on √ d S d -1 . Finally, let t ≥ 1 . Then, for some absolute constant c > 0 , with probability at least 1 -e -ct √ d , for all m ∈ [ N ] ,

Here, we use supp z to denote the set of nonzero coordinates of z .

Proof of Lemma I.11. The proof is the same as that of Theorem 3.6 of Adamczak et al. [4] - while Theorem 3.6 of Adamczak et al. [4] is shown for the case where X 1 , . . . , X N are sampled from an isotropic, log-concave probability distribution, the proof can be generalized to the case where the X i are sampled from √ d S d -1 . To see why, observe that the proof only uses the fact that the distribution of the X i is log-concave in the following places: (i) when applying Lemma 3.1 (also of Adamczak et al. [4]), of which the conclusion is that max i ≤ N ∥ X i ∥ 2 ≤ C 0 K √ d with probability at least 1 -e -K √ d , for some absolute constant C 0 , (ii) when applying Lemmas 3.3 and 3.4 (also of Adamczak et al. [4]) the proof requires that ∥⟨ X i , y ⟩∥ ψ 1 ≤ ψ for some absolute constant ψ > 0 and any y ∈ S d -1 . However, property (i) trivially holds when the X i are on √ d S d -1 since ∥ X i ∥ 2 = √ d , and property (ii) holds by Lemma I.10.

I.3 Helper Lemmas

Lemma I.12. Let f : [0 , ∞ ) → R ≥ 0 such that

where p t , q t > 0 for t ∈ [0 , ∞ ) . Then, for T > 0 we have

Thus, we have

As a result, by Fact I.13,

Proof.

Taking the square of both sides finishes the proof.

Fact I.13 (Gronwall's Inequality [13]) . Let a < b , and let f : [ a, b ] → R ≥ 0 and g : [ a, b ] → R be differentiable functions, such that

for all t ∈ [ a, b ] . Then, for all t ∈ [ a, b ] , f ( t ) ≤ f ( a ) exp( ∫ s a g ( s ) ds ) . Additionally, suppose f, g satisfy

for all t ∈ [ a, b ] . Then, for all t ∈ [ a, b ] , f ( t ) ≥ f ( a ) exp( ∫ s a g ( s ) ds ) .

Proposition I.14. For any β 2 , β 4 such that 0 ≤ β 2 2 ≤ β 4 ≤ β 2 ≤ 1 , there exists a one-dimensional distribution µ that satisfies

Proof. Let p = β 2 2 /β 4 . By the assumptions we have 0 ≤ p ≤ 1 . Consider the distribution

where δ x = t denotes the Dirac measure at x = t . First of all, since β 2 2 ≤ β 4 , the distribution µ ( x ) is well-defined. And we can verify that

Note that 0 ≤ β 4 ≤ β 2 implies β 1 / 2 4 β -1 / 2 2 ∈ [0 , 1] . Consequently,

Lemma I.15. Let ρ be a distribution which is rotationally invariant as in Definition 4.1. Let ∇ u L ( ρ ) be defined as in Eq. (3.4) , and let ∇ z L ( ρ ) denote the last ( d -1) coordinates of ∇ u L ( ρ ) . Then, for any particle u = ( w,z ) , we have that ∇ z L ( ρ ) is a scalar multiple of z . In particular, if grad u L ( ρ ) is defined as in Eq. (3.4) and grad z L ( ρ ) denotes the last ( d -1) coordinates of grad u L ( ρ ) , then grad z L ( ρ ) is a scalar multiple of z .

Proof. From Eq. (3.4), we have

Since ρ is rotationally invariant, if x, x ′ ∈ S d -1 such that ⟨ e 1 , x ⟩ = ⟨ e 1 , x ′ ⟩ , then f ρ ( x ) -y ( x ) = f ρ ( x ′ ) -y ( x ′ ) . Thus, ( f ρ ( x ) -y ( x )) σ ′ ( u ⊤ x ) only depends on ⟨ e 1 , x ⟩ and ⟨ u, x ⟩ .

Let P e 1 ,u ∈ R d × d denote the projection matrix into the subspace spanned by e 1 and u . Then,

because if v ∈ S d -1 is orthogonal to e 1 and u , then the distribution of ⟨ v, x ⟩ conditioned on ⟨ e 1 , x ⟩ and ⟨ u, x ⟩ is symmetric around 0 . Thus,

and we have that ∇ u L ( ρ ) is in the subspace spanned by e 1 and u . In other words, there are scalars c 1 , c 2 ∈ R such that

and taking the last ( d -1) coordinates of both sides of this equation gives ∇ u L ( ρ ) = c 2 z , as desired. The last statement of the lemma holds because grad u L ( ρ ) = ( I -uu ⊤ ) ∇ u L ( ρ ) , which is also a linear combination of e 1 and u .

Lemma I.16. Let P be a polynomial of degree at most k > 0 , where k is bounded by a universal constant, and let w 1 , w 2 ∈ [ -1 , 1] . Suppose the coefficients of P are bounded by B . Then, | P ( w 1 ) -P ( w 2 ) | ≲ B | w 1 -w 2 | .

Proof. Let c i be the coefficient of the i th degree term in P . Then, the lemma follows from expanding P ( w 1 ) -P ( w 2 ) as

as desired.

Lemma I.17 (Powers of Dot Products on S d -1 ) . Let u 1 , u 2 , u 3 , u 4 , u 5 ∈ S d -1 , and p, q, r, s, t ≥ 1 . Then,

Proof of Lemma I.17. By repeated applications of the Cauchy-Schwarz inequality, we have

E S d -1 [ |⟨ u 1 , X ⟩| p |⟨ u 2 , X ⟩| q |⟨ u 3 , X ⟩| r |⟨ u 4 , X ⟩| s |⟨ u 5 , X ⟩| t ] (I.145)

where each line results from a single application of the Cauchy-Schwarz inequality. Thus, by Lemma I.10, we have

Dividing both sides by d ( p + q + r + s + t ) / 2 , we obtain

Peer reviews:
comments: SUMMARY:
This paper studies the global convergence of gradient descent for training two-layer networks in learning a high dimensional quartic function. The authors show that GD converges when the sample size $n = O(d^{3.1})$ and the width of the neural network grows at most polynomially in d. The authors also showed that any kernel method with $n \ll d^4$ cannot achieve the same order of accuracy.

STRENGTHS:
First of all, this is an excellent paper on the non-asymptotic mean field analysis of GD training for two-layer neural networks and its separation from kernel method. Albeit the simplicity of the target function, the convergence analysis requests new ideas and tools. One crucial ingredient of the proof lies in showing that the projected gradient flow (population) dynamics can escape saddle points. The other key ingredient is bounding the coupling error between the empirical particle dynamics and the population dynamics. For the later, the authors tame the usual exponential growth rate in the coupling error with a delicate investigation on the relationship between the growth of the error and that of the signal.

WEAKNESSES:
I do not have many comments on the weakness, but have few questions, which are mostly relevant to future considerations. See below.

QUESTIONS:
The authors established the results for learning the quartic function. Does the proof generalize to a more general class of functions? If so, can the authors highlight the further work to be done? If not, it would be good if the authors can comment on the major challenges.

LIMITATIONS:
None
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studied the projected gradient flow on two-layer neural networks in the mean-field regime with polynomial width and quartic activation function. With data sampled uniformly on the sphere, the authors proved that to learn a single-index model with an even quartic link function, this neural network needs $n=O(d^{3.1})$ training samples. However, we know for any inner product kernel ridge regression requires $n\\gg d^4$ to learn a quartic target function.  This provides a concrete example of practical feature learning of neural networks outperforming the kernel methods and NTK regime when the target function has a low-rank structure.\n\nSTRENGTHS:\nOverall, I found the paper well-written and easy to follow. The paper presents interesting insights into more accurate models of neural network training. The results obtained by the authors are as sharp as can be reasonably expected given the problem in certain cases. Unlike previous analyses of feature learning with two-stage training procedures, this work shows that neural networks can learn both feature direction and nonlinear link function when training in a more practical way with the projected gradient flow. The methods of the proof may provide further insights to obtain more general results on feature learning.\n\nWEAKNESSES:\n1. The main limitation is that the activation function and the target link function are special, only quartic functions. It would be nice to check if the proof techniques can handle more general activation and target functions.  Another thing that can be improved is the precise order of the width: how wide the neural network is sufficient for this feature learning.\n2. There should be a comparison between the results of the current submission and [10]. [10] uses information exponent to determine the sample complexity for online SGD to learn a single-index target. When $\\gamma_2=0$ in Assumption 3.2, [10] implies that sample size should be $n\\gg d^{3}\\log n$ which is similar to the results in the current submission. Further explanations are needed.\n3.  And it should be more convincing to visualize Phases 1-3 of the training dynamic in Section 5 if there is some simulation of synthetic data or real-world data for this finite-width training dynamic.\n\nQUESTIONS:\n1. [Ba et al. 2022] also study the feature learning and single-index model when $n\\asymp d \\asymp m$ with two-stage training processes.  And [Arnaboldi et al. 2023] and [Berthier et al. 2023] also studied the gradient flow dynamics of two-layer neural networks in high-dimension to a single-index model. Some comparisons with the current paper could be made.\n2. In line 88, you mentioned that a similar coupling error was shown in [49] but [49] is for random feature regression with proportional limit. Can you explain more about it?\n3. Line 147 typo.\n4. In Theorem 3.4, the sample complexity is $n\\ge d^\\mu$. Is there an explicit formula or upper bound for the power $\\mu$?\n5. In Theorem 3.5, you considered the limitation of the inner-product kernel with data points drawn from the unit sphere. Does the data distribution have to be uniform unit sphere distribution? Can Theorem 3.5 be extended to a general rotational invariant kernel or Euclidean distance kernel?\n6. Line 233 double \"Appendix\"\n7. Line 274 $w$ should be $w_t$?\n8.  In the analysis, Phases 1 and 2 are controlled by the quadratic components in the activation and target functions, and the dynamics are analogous to a power method update. Is there any relation between this and the PCA warmup proposed by [Chen and Meka 2020] before SGD training?\n9. In Assumption 3.2, we need $\\gamma_4\\ge 1.1\\gamma_2^2$. Why do we need this assumption and coefficient $1.1$? What if $\\gamma_2=\\gamma_4=1$?\n\n===============================================================================================\n- Chen, S. and Meka, R., 2020, July. Learning polynomials in few relevant dimensions. In Conference on Learning Theory (pp. 1161-1227). PMLR.\n- Ba, J., Erdogdu, M.A., Suzuki, T., Wang, Z., Wu, D. and Yang, G., 2022. High-dimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems, 35, pp.37932-37946.\n- Arnaboldi, L., Stephan, L., Krzakala, F. and Loureiro, B., 2023. From high-dimensional & mean-field dynamics to dimensionless ODEs: A unifying approach to SGD in two-layers networks. arXiv preprint arXiv:2302.05882.\n- Berthier, R., Montanari, A. and Zhou, K., 2023. Learning time-scales in two-layers neural networks. arXiv preprint arXiv:2303.00055.\n\nLIMITATIONS:\nI think the authors have adequately addressed the potential negative social impact of their work. The conclusion lays out some suggested next steps and in doing so highlights certain current limitations.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

comments: SUMMARY:
This paper studies the statistical efficiency of the projected gradient dynamics on the sphere for (polynomial-width) two-layer neural networks under the mean-field regime. In particular, this work proves the sample complexity of $O(d^{3.1})$ for learning the single-index model with an unknown quartic link function.

STRENGTHS:
The results can be of interest to the readers and are technically sound. The remarkable contributions compared to most related works are that this work

- studies not unnatural optimization dynamics but the standard projected gradient flow on the sphere,

- shows the separation between two-layer neural networks in the mean field regime and kernel methods by providing the sample complexities.

Technically, the improvement in particle complexity is also worth noting. In a naive way, the approximation accuracy to the distribution by using finite particles exponentially deteriorates in time, but this study prevents such deterioration by making use of the problem and model structures. Hence, I think this work certainly contributes to the context.

WEAKNESSES:
- This theory deals with not standard activation such as ReLU, sigmoid, and tanh but a special type of activation function (i.e., fourth-order polynomial).

- In my understanding, Abbe et al. (2023) [2] also studies the standard optimization dynamics and shows the superiority of the mean-field regime against linear models. A detailed discussion of the relationship to [2] would help clarify the position of the work.

QUESTIONS:
Could you specify the dependence of $m$ stated in the next sentence of Lemma 5.1: ``network width $m$ is sufficiently large polynomial of $d$?'' What does the highest degree of this polynomial depend on?

LIMITATIONS:
The limitation of the paper is well addressed.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-08

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nAnalyzed the (projected) gradient flow dynamics of a two-layer neural network in the mean-field regime in learning a specific degree-4 single-index target function. The main contribution is a polynomial-time convergence guarantee and a sample complexity that outperforms kernel methods. This differs from the naive mean-field analysis where the network width can grow exponentially.\n\nSTRENGTHS:\nMost existing mean-field analyses of two-layer neural network focus on the optimization aspect, with the goal of showing global convergence of the training loss, but the generalization properties and sample complexity in learning certain class of target functions are largely unknown. The only exception that I know of is (Abbe et al. 2022). So this submission definitely tackles a challenging and interesting problem. On the technical side, the analysis differs from many previous works on learning single-index model with two-layer neural network, where the link function mismatch is typically handled by training the second-layer parameters.\n\nWEAKNESSES:\nI have the following concerns. \n\n1. The problem setting is restrictive and convoluted, which is rather underwhelming given the promise of a *clean mean-field analysis* with no unnatural modifications. (i) The training algorithm is the unmodified projected gradient flow, but time discretization is not studied. (ii) The target function is assumed to be an even degree-4 function which is quite restrictive. Moreover, the trained neural network also uses a somewhat unnatural degree-4 activation function, and the Legendre coefficients are restricted by Assumption 3.2. The fact that this submission is lengthy and technically demanding despite these strong assumptions raises the question of whether the same analysis can be extended to more general problem settings. \n\n2. The comparison against prior works is not sufficient, and consequently, the significance of the results cannot be easily evaluated.  \n* In the context of learning single-index model, the sample complexity of gradient descent is decided by the information exponent from (Ben Arous et al. 2021). This definition is originally defined in the analysis of online SGD, but the authors should discuss whether the same mechanism is present in the ERM setting. Specifically, my current reading is that this submission assumes an information exponent of 2, due to the factor of $\\sigma_2^2\\gamma_2$ in the denominator of the stopping time $T_*$ in Theorem 3.3.  \nPlease clarify if this is the case.   \n(Ben Arous et al. 2021) *Online stochastic gradient descent on non-convex losses from high-dimensional inference* \n\n* Related to the previous point, it is also well-known that trained two-layer neural network can outperform kernel models when the link function is a high-degree polynomial with low information exponent. For example, the staircase functions in (Abbe et al. 2022) can be learned in linear sample complexity. While these prior results typically employs a layer-wise training procedure, the similarity in the main message may undermine the significance of the theoretical results in this submission. Therefore, I feel that the authors should use more space to highlight the technical differences and challenges (no bias units and retraining of second layer, etc.).   \n(Abbe et al. 2022) *The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks* \n\n* The simplification of population dynamics into some low-dimensional object via symmetry has appeared in many prior works, such as the dimension-free PDE in (Abbe et al. 2022) and (Hajjar and Chizat 2022).  \nThe authors need to explain the difference and new ingredients in Section 4.1.  \n(Hajjar and Chizat 2022) *Symmetries in the dynamics of wide two-layer neural networks* \n\n* The claim that prior mean-field analysis only proved \"generic exponential growth rate in the coupling error\" is not entirely accurate. For the noisy gradient descent setting, recent papers have provided uniform-in-time propagation of chaos estimates, for example see (Chen et al. 2022) (Suzuki et al. 2023). These results hold for finite-width network and discrete-time algorithm, but the current submission only handles the finite-width error. It would be a good idea to comment on the possibility/difficulty of  handling the discrete projected gradient descent algorithm.   \n(Chen et al. 2022) *Uniform-in-time propagation of chaos for mean field Langevin dynamics*   \n(Suzuki et al. 2023) *Convergence of mean-field Langevin dynamics: Time and space discretization, stochastic gradient, and variance reduction*\n\nQUESTIONS:\nI would be happy to update my evaluation if the authors could address the concerns and questions in the Weaknesses section.  \n\n#####################**Post-rebuttal Update**#####################\n\nThe authors addressed some of my concerns; I have therefore increased my score to 5.  \nIn the revised manuscript, please include a detailed comparison against relevant prior works, as done in the rebuttal.\n\nLIMITATIONS:\nSee weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-11"}