The training of neural networks (NNs) is typically a non-convex optimization problem, but remarkably, simple algorithms like gradient descent (GD) or its variants can usually succeed in finding solutions with low training losses. To understand this phenomenon, a promising idea is to focus on NNs with large widths (a.k.a. under over-parameterization), for which we can derive infinite-width limits under suitable ways to scale the parameters by the widths. For example, under a "1/ √ width" scaling of the weights, the GD dynamics of wide NNs can be approximated by the linearized dynamics around initialization, and as the widths tend to infinity, we obtain the Neural Tangent Kernel (NTK) limit of NNs, where the solution obtained by GD coincides with a kernel method [37]. Importantly, theoretical guarantees for optimization and generalization can be obtained for wide NNs under this scaling [19,5]. Nonetheless, it was pointed out that this NTK analysis replies on a form of lazy training that excludes the learning of features or representations [16,72], which is a crucial ingredient to the success of deep learning, and is therefore not adequate for explaining the success of NNs [27,41]. Meanwhile, for shallow (i.e., one-hidden-layer) NNs, if we choose a "1 / width" scaling, we can derive an alternative mean-field (MF) limit as the widths tend to infinity. Under this scaling, feature learning occurs even in the infinite-width limit, and the training dynamics can be described by the Wasserstein gradient flow of a probability measure on the space of the parameters, which converges to a global minimizer of the loss function under certain conditions [59,50,63,14]. Generalization guarantees have also been proved for learning with shallow NNs under the MF scaling by identifying a corresponding function space [6,48]. However, currently there are three limitations to this model of over-parameterized NNs. First, the global convergence guarantees for shallow NNs only hold in the infinite-width limit (i.e. they are asymptotic). While [12] studies the deviation between finite-width NNs and their infinite-width limits during training, the analysis is done only asymptotically to the next order in width. Second, a convergence rate has yet to be established except under special assumptions or with modifications to the GD algorithm [13,34,43,38]. Third, while several works have proposed to extend the MF formulation to deep (i.e., multi-layer) NNs [4,62,52,23,57], there is less concensus on what the right model should be than for the shallow case. In summary, we still lack a model for the GD optimization of shallow and multi-layer NNs that goes beyond lazy training while admitting fast global convergence.
In this work, we study the optimization of both shallow NNs under the MF scaling and a type of partially-trained multi-layer NNs, and obtain theoretical guarantees of linear-rate global convergence.
We consider the scenario of training NN models to fit a training set of n data points in dimension d, where the model parameters are optimized by gradient flow (GF, which is the continuous-time limit of GD) with respect to the squared loss. Allowing most choices of the activation function, we prove that:
1. For a shallow NN, if the hidden layer is sufficiently wide and the input data are linearly independent (requiring n ≤ d), then with high probability, the training loss converges to zero at a linear rate.
2. For a multi-layer NN where we only train the second-to-last layer, if the hidden layers are both sufficiently wide, then with high probability, the training loss converges to zero at a linear rate. Unlike for shallow NNs, here we no longer need the requirement on input dimension, demonstrating a benefit of jointly having depth and width.
We also run numerical experiments to demonstrate that our model exhibits feature learning and can achieve better generalization performance than its NTK counterpart.
Over-parameterized NNs, NTK and lazy training. Many recent works have studied the optimization landscape of NNs and the benefits of over-parameterization [24,67,60,66,64,54,65,42,3,40,76,15,75,17,39,9,32,20,31,61,10]. One influential idea is the Neural Tangent Kernel (NTK) [37], which characterizes the behavior of GD on the infinite-width limit of NNs under a particular scaling of the parameters (e.g. for shallow NNs, replacing 1/m with 1/ √ m in (1)). In particular, when the network width is polynomially large in the size of the training set, the training loss converges to a global minimum at a linear rate under GD [19,5,55]. Nonetheless, in the NTK limit, due to a relatively large scaling of the parameters at initialization, the hidden-layer features do not move significantly [37,19,16]. For this reason, the NTK scaling has been called the lazy-training regime, as opposed to a feature-learning or rich regime [16,72,26]. Several works have investigated the differences between the two regimes both in theory [28,29,69,47] and in practice [27,41]. In addition, several works have generalized the NTK analysis by considering higher-order Taylor approximations of the GD dynamics or finite-width corrections to the NTK [2,35,7,33].
Mean-field theory of NNs. An alternative path has been taken to study shallow NNs in the meanfield scaling (as in ( 1)), where the infinite-width limit is analogous to the thermodynamic or hydrodynamic limit of interacting particle systems [59,50,63,14,49,70]. Thanks to the interchangeability of the parameters, the neural network is equivalently characterized by a probability measure on the space of its parameters, and the training can then be described by a Wasserstein gradient flow followed by this probability measure, which, in the infinite-width limit, converges to global mimima under mild conditions. Regarding convergence rate, ref. [71] proves that if we train a shallow NN to fit a Lipschitz target function under population loss, the convergence rate cannot beat the curse of dimensionality. In contrast, we will study the setting of empirical risk minimization, where there are finitely many training data. Ref. [34] shows that mean field Langevin dynamics on shallow NNs can converge exponentially to global minimizers in over-regularized scenarios, but we focus on GF without entropic regularization. Besides the question of optimization, shallow NNs under this scaling represent functions in the Barron space [48] or variation-norm function space [6], which provide theoretical guarantees on generalization as well as fluctuation in training [12]. Several works have proposed different mean-field limits of wide multi-layer NNs and proved convergence guarantees [4,62,51,52,23,57,22], but questions remain. First, due to the presence of different symmetries in a multi-layer network compared to a shallow network [57], the limiting object at the infinite-width limit is often quite complicated. Second, it has been pointed out that under the MF scaling of a multi-layer network, an i.i.d. initialization of the weights would lead to a collapse of the diversity of neurons in the middle layers, diminishing the effect of having large widths [23]. In addition, while another line of work develops MF models of residual models [8,46,21,36], we are interested in multi-layer NN models with a large width in every layer.
Feature learning in deep NNs. Ref. [1] demonstrates the importance of hierarchical learning by proving the existence of concept classes that can be learned efficiently by a deep NN with quadratic activations but not by non-hierarchical models. Ref. [11] studies the optimization landscape and generalization properties of a hierarchical model that is similar to ours in spirit, where an untrained embedding of the input is passed into a trainable shallow model, and prove an improvement in sample complexity in learning polynomials by having neural network outputs as the embedding. However, the trainable models they consider are not shallow NNs but their linearized and quadratic-Taylor approximations, and furthermore the convergence rate of the training is not known. Ref. [73] proposes a novel parameterization under which there exists an infinite-width limit of deep NNs that exhibits feature learning, but properties of its training dynamics is not well-understood. Our multi-layer NN models adopt an equivalent scaling (see Appendix C), and our focus is on proving non-asymptotic convergence guarantees for its partial training under GF.
We summarize our notations in Appendix A. Let Ω ⊆ R d denote the input space, and let x = [x 1 , ..., x d ] ⊺ ∈ Ω denote a generic input data vector. A shallow NN model under the MF scaling can be written as:
where m is the width, W ∈ R m×d and c = [c 1 , ..., c m ] ∈ R m are the first-and second-layer weight parameters of the model, and σ : R → R is the activation function. For simplicity of presentation, we neglect the bias terms. In this paper, we study a more general type of models with the following form:
∀i ∈ [m] :
where W ∈ R m×D and c = [c 1 , ..., c m ] ∈ R m are parameters of the model, and φ 1 , ..., φ D are a set of functions from Ω to R that we call the embedding. Each of h 1 , ..., h m is a function from Ω to R, and we will refer to them as the (hidden-layer) feature map or activations. For simplicity, we write
We consider two types of the embedding, Φ, as described below:
Fixed embedding D is fixed and Φ is deterministic. In the simplest example, we set D = d and φ j (x) = x j , ∀j ∈ [D], and recover the shallow NN model in (1). More generally, our definition includes cases where Φ is a deterministic transformation of an input vector in Ω into an embedding vector in R D . This can be understood as input pre-processing or feature engineering.
High-dimensional random embedding D is large and Φ is random. For instance, we can sample each z j i.i.d. in R d and set φ j (x) = σ 1 √ d z ⊺ j x , equivalent to setting φ 1 , ..., φ m as the hidden-layer activations of a shallow NN with randomly-initialized first-layer weights. Then, the model becomes
∀i ∈ [m] :
Thus, we obtain a 3-layer feed-forward NN whose first-layer weights are random and fixed, and we call it a partially-trained 3-layer (P-3L) NN. Note that the scaling in this model is different from both the NTK scaling (1/ √ m instead of 1/m in (4)) and the MF scaling for multi-layer NNs adopted in [4,62,51,57,23] (1/D instead of 1/ √ D in ( 5)). We show in Appendix B that when σ is homogeneous, this scaling is consistent with the Xavier initialization of neural network parameters up to a reparameterization [30,56]. We also show in Appendix C that in certain cases this scaling is equivalent to the maximum-update parameterization proposed in [73]. Numerical experiments that compare different scalings are described in Section 4.
Consider the scenario of supervised least-squares regression, where we are given a set of n training data points together with their target values, (x 1 , y 1 ), ..., (x n , y n ) ∈ Ω × R. We fit our models by minimizing the empirical squared loss:
To do so, we first initialize the parameters randomly by sampling each c i and W ij i.i.d. from probability measures π c ∈ P(R) and π w ∈ P(R D ), respectively, and then perform GD on W . For simplicity of analysis, we leave c untrained, and further assume that Assumption 1. π c = 1 2 δ ĉ(dc) + 1 2 δ -ĉ (dc) for some ĉ > 0 independent from m, which is the law of a scaled Rademacher random variable.
If σ is Lipschitz, it is differentiable almost everywhere, and we write σ ′ (x) to denote the derivative of σ when it is differentiable at x and 0 otherwise. When σ is differentiable at
and the gradient of the loss function with respect to W ij is given by
Thus, we can perform GD updates on W according to the following rule: ∀i ∈ [m] and ∀j ∈ [D],
where δ > 0 is the step size. As we discuss in Appendix B, this is consistent with the standard GD update rule for Xavier-initialzed NNs. In the limit of infinitesimal step size (δ → 0), the evolution of the parameters during training is described by the GF equation: if we use the superscript t ≥ 0 to denote time elaposed during training, the time-derivative of the parameters is given by
where f t denotes the output function and h t 1 , ..., h t m denote the hidden-layer feature maps determined by the parameters at time t. Then, induced by the evolution of W t , each h t i evolves according to
where we define a kernel function
Accordingly, the output function
Thus, the loss function
where we define the (symmetric
and use λ min (G) to denote its least eigenvalue. We will also use
G aa to denote the minimum and maximum diagonal entries of G, respectively. Since G is positive semi-definite, we see that Lt ≤ 0, which means that the loss value is indeed non-increasing along the GF trajectory.
Feature learning Compared to the NTK scaling of neural networks, the crucial difference is the 1/m factor in (2), instead of 1/ √ m. It is known that under the NTK scaling, due to the 1/ √ m factor, the movement of the feature maps, h 1 , ..., h m , is only of order O(1/ √ m) while the function value changes by an amount of order Ω(1). While this greatly simplifies the convergence analysis, it also implies that the hidden-layer representations are not being learned. In contrast, with the
which implies that the average movement of the feature maps is on the same order as the change in function value, and thus the hidden-layer representations as well as the NTK undergoes nontrivial movement during training. In Appendix C, we further justify the occurrence of feature learning using the framework developed in [73].
To prove that the training loss converges to zero, we need a lower bound on the absolute value of Lt . Indeed, if G is positive definite, which depends on Φ and the training data, we can establish one in the following way. First, as a simple case, if we use an activation function whose derivative's absolute value is uniformly bounded from below by a constant K σ ′ > 0, such as linear, cubic or (smoothed) Leaky ReLU activations, we can derive a Polyak-Lojasiewicz (PL) condition [58,45] from ( 14) directly,
which implies L t ≤ L 0 e -2ĉ 2 λmin(G)(K σ ′ ) 2 t , indicating that the training loss decays to 0 at a linear rate.
For more general choices of the activation function, a challenge is to guarantee that, heuristically speaking, for each a ∈ [n], σ ′ h i (x a ) does not become near zero for too many i ∈ [m] before the loss vanishes. To facilitate a finer-grained analysis, we need the following mild assumption on σ: Assumption 2. σ is Lipschitz with Lipschitz constant L σ , and there exists an open interval I = (I l , I r ) ⊆ R on which σ is differentiable and |σ ′ | is lower-bounded by some K σ ′ > 0.
Intuitively, I is an active region of σ, within which the derivative has a magnitude bounded away from zero. This assumption is satisfied by the majority of activation functions in practice, including smooth ones such as tanh and sigmoid as well as non-smooth ones such as ReLU. Then, under the following initialization scheme, we prove a general result for models with a fixed embedding. Assumption 3. π w is the D-dimensional standard Gaussian distribution, i.e., each W ij is sampled independently from a standard Gaussian distribution. Theorem 1 (Fixed embedding). Suppose that Assumptions 1, 2 and 3 are satisfied, and λ min (G) > 0.
Then ∃ĉ 0 , r and
then with probability at least 1δ, it holds that ∀t ≥ 0,
Here, ĉ0 , r and C depend on I, G min , G max , y , L σ and K σ ′ (but not on m, n, d, D, δ, or λ min (G)).
The result is proved in Appendix E, and below we briefly describe the intuition. A key to the proof is to guarantee that enough neurons remain in the active region throughout training. Specifically, with respect to each training data point (i.e. for each a ∈ [n]), we can keep track of the proportion of neurons (among all i ∈ [m]) for which h t i (x a ) ∈ I. We show that if the proportion is large enough at initialization (shown by Lemma 3 in Appendix E.2 under Assumption 3), then it cannot drop dramatically without a simultaneous decrease of the loss value, as long as the c i 's are not too small in absolute value. This property of the dynamics is formalized in the following lemma: Lemma 1. Consider the dynamics of L t and h t i (x a ) i∈[m],a∈ [n] governed by (11) and (14). Assume that λ min (G) > 0, and ∀i ∈ [m], |c i | = ĉ > 0. Under Assumption 2, define
, ∀t ≥ 0 and η0 = min
) (18) Then ∀t ≥ 0, there is
where κ(λ 1 , λ 2 ) = 9λ2(Ir-I l ) 2λ1K σ ′ .
In the case of shallow NNs under the MF scaling, G = G (0) ∈ R n×n , where
Thus, G (0) is positive definite if and only the training data set x 1 , ..., x n } ⊆ R d consists of linearlyindependent vectors, which is possible (and expected if the training data are sampled independently from some non-degenerate distribution) when n ≤ d. In that case, Theorem 1 implies Corollary 1 (Shallow NN with n ≤ d). Suppose that Assumptions 1, 2 and 3 are satisfied. If the training data are linearly-independent vectors, then under GF (10) on the first-layer weights of the shallow NN, the training loss converges to zero at a linear rate.
While the assumption that n ≤ d is restrictive, we note that existing convergence rate guarantees for the GD-type training of shallow NNs in the MF scaling need strong additional assumptions [38], modifications to the GD algorithm [34,13,53], or restrictions to certain special tasks [43].
A clear limitation of Corollary 1 is that it is only applicable when n ≤ d, since otherwise the Gram matrix G (0) cannot be positive definite. This motivates us to consider the use of a high-dimensional embedding Φ to lift the effective input dimension. In particular, we focus on the scenario where D is large and Φ is random. While the Gram matrix G in this case is also random, we only need that it concentrates around a deterministic and positive definite limit as D tends to infinity: Condition 1 (Concentration of G around a positive definite matrix). There exists a (deterministic) positive definite matrix Ḡ ∈ R n×n with least eigenvalue λ min ( Ḡ) > 0 such that ∀δ, u > 0,
Condition 1 is sufficient for us to apply Lemma 1 and obtain the following global convergence guarantee, which extends Theorem 1 to models with a high-dimensional random embedding. The proof is given in Appendix F. Theorem 2 (High-dimensional random embedding). Under Assumptions 1, 2, 3 and Condition 1, ∃ĉ 0 , r and
, then with probability at least 1δ, it holds that ∀t ≥ 0, L t ≤ L 0 e -rĉ 2 λmin(G)t . (21) Here, ĉ0 , r and C depend on I, Ḡmin , Ḡmax , y , L σ and K σ ′ (but not m, n, d, D, δ, or λ min ( Ḡ)).
Consider the P-3L NN model defined in (4). In this case, the Gram matrix is G (1) , defined by
If z 1 , ..., z m are sampled i.i.d. from a probability measure π z on R d and fixed during training, then the limiting Gram matrix, denoted by Ḡ(1) ∈ R n×n , is given by
Thus, for the convergence result, the assumption we need on the limiting Gram matrix is Assumption 4. π z is sub-Gaussian and the matrix Ḡ(1) , which depends on the choice of σ and the training set, is positive definite with λ min Ḡ(1) > 0 and ( Ḡ(1) ) max < ∞.
This assumption also plays an important role in the NTK analysis, and it is satisfied if, for example, π z is the d-dimensional standard Gaussian distribution, no two data points are parallel, and σ is either the ReLU function [19] or analytic and not a polynomial [18]. When Assumption 4 is satisfied, as long as σ is Lipschitz, we can use standard concentration techniques to verify Condition 1. Thus, Theorem 2 implies that Theorem 3 (P-3L NN). Under Assumptions 1, 2, 3 and 4, ∃ĉ 0 , r, C 1 and
then with probability at least 1δ, it holds that ∀t ≥ 0,
Here, ĉ0 , r, C 1 and C 2 depend on I, Ḡ(1) min , Ḡ(1) max , y , K σ ′ as well as the sub-Gaussian norm of µ z (but not on m, n, d, D, δ or λ min ( Ḡ(1) )).
The proof is given in Appendix G. Compared to Corollary 1 for shallow NNs, a highlight of Theorem 3 is that the requirement of n ≤ d is no longer needed. This demonstrates an advantage of the highdimensional random embedding realized by the first hidden layer in the P-3L NN, thus illustrating a benefit of having both depth and width in NNs from the viewpoint of optimization. Compared to the NTK result [18], our analysis assumes the same level of over-parameterization, but crucially allows feature training to occur, which we discuss in Section 2.2 and support empirically in Section 4.3. Furthermore, by using a multi-layer NN with random and fixed weights as the high-dimensional random embedding, we extend the P-3L NN to a partially-trained L-layer NN model in Appendix H, for which similar convergence results can be proved for training its second-to-last layer via GF.
Additional results and details of the experiments are provided in Appendix I.
We train shallow NNs to fit a randomly labeled data set {(x 1 , y 1 ), ..., (x n , y n )} with d = 20. Specifically, we sample each x a i.i.d. with every entry sampled independently from a standard Gaussian distribution, and each y a i.i.d. uniformly on [-1 2 , 1 2 ] and independently from the x a 's. We see from Figure 1 that the convergence happens at a nearly linear rate when n = 20 and 40, and the rate decreases as n becomes larger. This is coherent with our theoretical result (Corollary 1), and interestingly also echoes a prior result that the convergence rate of optimizing a shallow NN using population loss can suffer from the curse of dimensionality [71], which implies a worsening of the convergence rate as the number of data points increases.
We consider a model defined by ( 2) and ( 3) with d = 30 and Φ(x) = vec(xx ⊺ ) ∈ R d 2 , which we call a shallow NN augmented with quadratic embedding. We compare this model against a plain shallow NN (without the extra embedding), both with m = 8192, to fit a series of training sets with various sizes where the target y is given by another shallow NN augmented with quadratic embedding with m = 5. We see from Figure 2 that the augmented shallow NN achieves lower test error given the same number of training samples, demonstrating the benefit of a good embedding.
We consider the P-3L NN model defined in (4) and ( 5) with D = m (i.e. both hidden layers having the same width), and compare it with 3-layer NN models under NTK and MF scalings, as we define in Table 1 based on prior literature [37,4,52,62,23], which undergo partial training in the same fashion. We adopt the data set used in [69] (more details in Appendix I.3), and train the models by minimizing the unregularized squared loss for varying n's and m's.
First, we see from the top-left plot in Figure 4 that, consistently across different m, the training loss converges at a linear rate for the model under our scaling, which is coherent with Theorem 3. Second, we see from the second row that feature learning occurs in the model under our scaling but negligibly in the model under the NTK scaling, as expected [16]. Note also that under the MF scaling, the feature maps h 1 (x), ..., h m (x) concentrate near 0 at initialization due to the small scaling, but gains diversity during training. Third, we see from Figure 3 that our model yields the smallest test errors out of all three, and in addition, as n grows the test error decreases faster under the MF scaling than under the NTK scaling, both indicating an advantage of feature learning compared to lazy training.
We consider a general type of models that includes shallow and partially-trained multi-layer NNs, which exhibits feature learning when trained via GF, and prove non-asymptotic global convergence guarantees that accommodates a general class of activation functions. For a randomly-initialized shallow NN in the MF scaling that is wide enough, we prove that by performing GF on the input-layer weights, the training loss converges to zero at a linear rate if the number of training data does not exceed the input dimension. For a randomly-initialized multi-layer NN with large widths, we prove Each column corresponds to a different scaling of the P-3L NN model, as described in Table 1. Row 1: Evolution of training loss (solid curve) and test error (dashed curve) during training. Row 2: Distribution of the hidden-layer feature map (pre-activation) associated with two particular input data points. Each dot represents a different i, (i.e., neuron in the second hidden layer,) and the xand y-coordinates equal h i (x 1 ) and h i (x 2 ), respectively, where x 1 is an input from the training set and x 2 is an input from the test set.
that by performing GF on the weights in the second-to-last layer, the same result holds except there is no requirement on the input dimension. We also perform numerical experiments to demonstrate the advantage of feature learning in our partially-trained multi-layer NNs relative to their counterparts under the NTK scaling.
Our work focuses on the optimization rather than the approximation or generalization properties of NNs, which are also crucial to understand. In addition, as our current theoretical results on global convergence neglect the bias terms and assume that the last-layer weights are untrained, a more general version is left for future work.
• For a positive integer n, we let [n] denote the set {1, ..., n}.
• We use i, j (as subscripts) to index the neurons in the hidden layers, a, b (as subscripts or superscripts) to index different training data points, t (as a superscript) to denote the training time / time parameter in gradient flow.
• We write a for 1
∀i ∈ [m] :
with weight parameters θ
] are initialized according to Xavier initialization, which means that we sample each θ
ij i.i.d. from N (0, 1 2m ), and each θ
(3) i i.i.d. from N (0, 1 m+1 ). If m ≫ d, both N (0, 1 m+d ) and N (0, 1 m+1 ) can be approximated by N (0, 1 m ). Then, up to this approximation, by redefining
ij and z jk = √ mθ
jk , we can write
∀i ∈ [m] :
and note that c i , W ij and z jk are all initialized i.i.d. of order O(1). In addition, if σ is homogeneous, this is then equivalent to (4) and ( 5) when D = m.
this is equivalent to updating W ij according to
which justifies the m factor on the right-hand-side of (9).
Consider the partially-trained L-layer NN model defined in Section H in the case where D = m ≫ d.
In the framework of abc-parameterization introduced in [73], our model corresponds to setting
Furthermore, as we explain in Appendix B, the appropriate learning rate scales linearly with m (as in ( 9)), which corresponds to having c = -1
Meanwhile, the maximum-update (µP) parameterization [73] is characterized by setting
Recall the symmetry of abc-parameterization derived in [73], which states that one gets a different but equivalent abc-parameterization by setting
Since our parameterization can be obtained from the maximum-update parameterization by applying the transformation above with θ = 1 2 , they are equivalent in the function space. In particular, for our parameterization, the r parameter defined in [73] can be computed as
Hence, according to [73], our parameterization exhibits feature learning.
Since we assume that G is positive definite and
we can derive from ( 14) that
Since I is an open interval, ∃ξ > 0 such that we can find a subinterval I 0 ⊆ I such that the distance between I 0 and the boundaries of I (if I is bounded on either side) is no less than ξ, i.e.,
In particular, we can choose ξ = 1 3 (I r -I l ) and I 0 = (I l + ξ, I rξ). Then there is
and so
(40) Thus, we have
there is
Therefore,
Since ∀ξ ∈ R, there is
we derive that, ∀a ∈
where we set C 1 = λ max (G) (λ min (G))
-1 2 ξ -1 > 0 for simplicity. As a consequence,
Define
Note that at t = 0, there is ηt = min a∈[n]
. Then, on one hand, we know from ( 40) and ( 47) that ∀t ≥ 0,
Hence, (41) implies that
On the other hand, by the definition of ηt ,
where we set G) for simplicity. Therefore, when η t > 0,
which implies that
and so ∀t ≥ 0, 2
To apply Lemma 1, we need two additional lemmas, which we will prove in Appendix E.1 and E.2. The first one guarantees that the loss value at initialization, L 0 , is upper-bounded with high probability:
Lemma 2. ∀δ > 0, if m ≥ Ω ĉ2 log nδ -1 G max / y 2 , then with probability at least 1δ, there is
The second one proves that η0 is lower-bounded with high probability, which heuristically says that there is indeed a nontrivial proportion of neurons in the central part of the active region of σ, for every a ∈ [n]:
2(K(I,Gmin,Gmax)) 2 , then with probability at least 1δ, there is
where
is a positive number that depends on I, λ 1 and λ 2 .
With these two lemmas, we deduce that ∀δ > 0, if m ≥ Ω (1 + ĉ2 / y 2 ) log nδ -1 , then with probability at least δ, there is ∀t ≥ 0, 2
where K(I, G min , G max ) is defined as in Lemma 3. Therefore, if our choice of ĉ satisfies
then there is ∀t ≥ 0,
in which case (50) gives
which will allow us to finally conclude that
where M SG > 0 is some absolute constant. Thus, by Hoeffding's inequality [68], ∀a ∈ [n], ∀r > 0,
where K is some absolute constant. Hence, by union bound,
then with probability at least 1 δ, there is
E.2 PROOF OF LEMMA 3
Since each W 0 ij are sampled i.i.d. from N (0, 1), we know that ∀a ∈ [n], independently for each i ∈ [m], h 0 i (x a ) follows a Gaussian distribution with mean 0 and variance G aa . Therefore,
Hence, by Hoeffding's inequality, ∀a ∈ [n], ∀r > 0,
∀a ∈ [n], choosing r = 1 2 π (I 0 ; G aa ), we then get
and so by union bound
then we have G
ab -Ḡ(1) ab ≥ u with probability at least 1δ. If we choose u = u ′ n and δ ′ = δ n 2 , then we get, if
Hence, with probability at least 1δ, we have
H GENERALIZATION TO DEEPER MODELS By setting Φ to be the activations of the second-to-last hidden-layer of a multi-layer NN, we can obtain generalizations of the P-3L NN to deeper architectures. For example, in the feed-forward case, we can obtain the following partially-trained L-layer NN:
where W (1) , ..., W (L-3) ∈ R D×D and z 1 , ..., z D ∈ R d are sampled randomly and fixed. This model can be written in the form of ( 4) and ( 5) with φ j (x) = σ h (L-2) j (x) . The corresponding Gram matrix is recursively defined and also appears in the NTK analysis [18]. In particular, the results in [18] imply that if σ is analytic and not a polynomial, then Condition 1 holds, and hence similar global convergence results can be obtained as corollaries of Theorem 2.
In our models, c i i∈[m] is sampled i.i.d. from the Rademacher distribution µ c = 1 2 δ 1 + 1 2 δ -1 , z j j∈[D] is sampled i.i.d. from N (0, I d ), and W ij i∈[m],j∈ [D] is initialized by sampling i.i.d. from N (0, 1). In the model under NTK scaling, we additionally symmetrize the model at initialization according to the strategy used in [16] to ensure that the function value at initialization does not blow up when the width is large. We choose to train the models using 50000 steps of (full-batch) GD with step size δ = 1. When the test error is computed, we use a test set of size 500 generated by sampling i.i.d. from the same distribution as the training set.
The experiments are run with NVIDIA GPUs (1080ti and Titan RTX).