ON FEATURE LEARNING IN NEURAL NETWORKS WITH GLOBAL CONVERGENCE GUARANTEES
Year: 2022

Abstract:
We study the gradient flow optimization of over-parameterized neural networks (NNs) in a setup that allows feature learning while admitting non-asymptotic global convergence guarantees. First, we prove that for wide shallow NNs under the mean-field (MF) scaling and with a general class of activation functions, when the input dimension is at least the size of the training set, the training loss converges to zero at a linear rate under gradient flow. Building upon this analysis, we study a model of wide multi-layer NNs with random and untrained weights in earlier layers, and also prove a linear-rate convergence of the training loss to zero, regardless of the input dimension. We also show empirically that, unlike in the Neural Tangent Kernel (NTK) regime, our multi-layer model exhibits feature learning and can achieve better generalization performance than its NTK counterpart.

heading: 1 INTRODUCTION
text: The training of neural networks (NNs) can often succeed despite being a non-convex optimization problem and involving a large number of parameters. To understand this mystery, one recent direction is to focus on the over-parameterized regime, where the NNs have large widths, by studying their infinite-width limits under certain scalings of the parameters. For example, under the NTK or “1/ √ width” scaling, the gradient descent (GD) dynamics of the training of a wide NN can be approximated by the linearized dynamics around initialization, and in the infinite-width limit, we obtain a kernel method equipped with the Neural Tangent Kernel (NTK) [37]. Theoretical guarantees for optimization and generalization can be obtained for wide NNs under the NTK scaling [19, 5]. Nonetheless, it has been pointed out that the 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].
Alternatively, for shallow (i.e. one-hidden-layer) NNs, if we initialize the parameters according to a scaling of “1 / width”, we can derive the mean-field (MF) limit of shallow NNs. In this limit, feature learning occurs, 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 via a corresponding function space [6, 48].
Currently, however, three limitations exist to the mean-field analysis of NNs. First, the global convergence guarantees for shallow NNs only hold in the infinite-width limit, and the deviation between finite-width NNs and their infinite-width limits has only been studied asymptotically [12]. Second, a convergence rate is 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. More generally, we still lack a model for the optimization of shallow and multi-layer NNs that goes beyond lazy training while admitting fast global convergence.
The aim of this work is to obtain theoretical guarantees of fast global convergence of feature learning in both shallow and a type of multi-layer NNs.

heading: 1.1 SUMMARY OF MAIN CONTRIBUTIONS
text: We consider the scenario of training an NN model to fit a set of n training data with input dimension d under least-squares regression, where the model parameters are optimized by gradient flow. With 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-layer widths are sufficiently large, 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 that n ≤ d, demonstrating a benefit of jointly having depth and width.
To our knowledge, we are presenting the first theoretical guarantee of linear-rate convergence of feature learning in multi-layer NNs. We also run numerical experiments that complement the theory.

heading: 1.2 RELATED WORKS
text: 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 NNs by its linearization around initialization, and becomes exact if the width tends to infinity under a particular scaling of the parameters (e.g. replacing 1/m with 1/ √ m in (1)). In particular, when the network width is polynomially large in the number of training set size, the training loss converges to global minima 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 gradient flow 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 than 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 quadraticTaylor approximations, and furthermore the convergence rate of the training is not known. Ref. [73] proposes an alternative scaling under which there exists an infinite-width limit of deep NNs that exhibits feature learning, but properties of its training dynamics is still not well-understood. Our pseudo-multi-layer NN models adopt the same scaling, and our focus is on proving non-asymptotic convergence guarantees.

heading: 2 PROBLEM SETUP

heading: 2.1 MODEL
text: We summarize our notations in Appendix A. Let Ω ⊆ Rd denote the input space, and let x = [x1, ..., xd]
ᵀ ∈ Ω denote a generic input data vector. A shallow NN model under the MF scaling can be written as:
f(x) = 1
m m∑ i=1 ciσ ( 1√ d d∑ j=1 Wijxj ) , (1)
where m is the width, W ∈ Rm×d and c = [c1, ..., cm] ∈ Rm 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:
f(x) = 1
m m∑ i=1 ciσ ( hi(x) ) , (2)
∀i ∈ [m] : hi(x) = 1√ D D∑ j=1 Wijφj(x) , (3)
where W ∈ Rm×D and c = [c1, ..., cm] ∈ Rm are parameters of the model, and φ1, ..., φD are a set of functions from Ω to R that we call the embedding. Each of h1, ..., hm is a function from Ω to R, and we will refer to them as the (hidden-layer) feature map or activations. For simplicity, we write Φ(x) = [φi(x), ...φD(x)] ᵀ ∈ RD. We consider two types of the embedding, Φ:
Fixed embedding D is fixed and Φ is deterministic. In the simplest example, we set D = d and φj(x) = xj , ∀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 RD. 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 zj i.i.d. in Rd 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
f(x) = 1
m m∑ i=1 ciσ ( hi(x) ) , (4)
∀i ∈ [m] : hi(x) = 1√ D D∑ j=1 Wijσ ( 1√ d zᵀj x ) . (5)
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 a special case of this scaling is equivalent to the maximum-update scaling scheme proposed in [73]. Numerical experiments that compare different scalings are described in Section 4.

heading: 2.2 TRAINING WITH GRADIENT FLOW
text: Consider the scenario of supervised least-squares regression, where we are given a set of n training data points together with their target values, (x1, y1), ..., (xn, yn) ∈ Ω× R. We fit our models by minimizing the empirical least-squares loss:
L[f ] = 1 2 n∑ a=1 ( f(xa)− ya )2 (6)
To do so, we first initialize the parameters randomly by sampling each ci and Wij i.i.d. from probability measures πc ∈ P(R) and πw ∈ P(RD), respectively, and then perform GD on W . For simplicity of analysis, we leave c untrained, and further assume that Assumption 1. πc = 12δĉ(dc) + 1 2δ−ĉ(dc) for some ĉ > 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) for the value of its derivative when σ is differentiable at x and 0 otherwise. When σ is differentiable at hi(x), there is
∂f
∂Wij (x) =
1 m √ D ciσ ′(hi(x))φj(x) . (7)
Thus, the gradient of the loss function with respect to Wij is given by
∂L[f ] ∂Wij = 1 m √ D ci n∑ a=1 (f(xa)− ya)σ′ ( hi(xa) ) φj(xa) . (8)
Thus, we can perform GD updates on W according to the following rule (consistent with standard GD for Xavier-initialzed NNs, as discussed in Appendix B): ∀i ∈ [m] and ∀j ∈ [D],
Wij ← Wij −mδ ∂L[f ] ∂Wij = Wij − δ√ D ci n∑ a=1 (f(xa)− ya)σ′ ( hi(xa) ) φj(xa) , (9)
where δ > 0 is the step size. In the limit of infinitesimal step size (δ → 0), the evolution of the parameters during training is described by the gradient flow equation: if we use the superscript t ≥ 0 to denote time elaposed during training, the time-derivative of the parameters is given by
Ẇ tij = − ci√ D n∑ a=1 ( f t(xa)− ya ) σ′ ( hti(xa) ) φj(xa) , (10)
where f t denotes the output function and ht1, ..., 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 hti evolves according to
∀x ∈ Ω : ḣti(x) = 1√ D D∑ j=1 Ẇ tijφj(x) = −ci n∑ a=1 G(x,xa) ( f t(xa)− ya ) σ′ ( hti(xa) ) , (11)
where we define a kernel function G : Ω× Ω→ R as
G(x,x′) = 1 D D∑ j=1 φj(x)φj(x ′) = 1 D (Φ(x))ᵀΦ(x′) . (12)
Accordingly, the output function f t satisfies
∀x ∈ Ω : ḟ t(x) = 1 m m∑ i=1 ciσ ′(hti(x))ḣti(x)
=− ĉ 2
m m∑ i=1 σ′ ( hti(x) ) n∑ a=1 G(x,xa) ( f t(xa)− ya ) σ′ ( hti(xa) ) ,
(13)
Thus, the loss function Lt := L[f t] evolves according to
L̇t = n∑ a=1 ( f t(xa)− ya ) ḟ t(xa)
=− ĉ 2
m m∑ i=1 n∑ a,b=1 Gab ( f t(xa)− ya ) ( f t(xb)− yb ) σ′ ( hi(xa) ) σ′ ( hi(xb) ) ≤− ĉ 2λmin(G)
m
m∑ i=1 n∑ a=1 ( f t(xa)− ya )2 ( σ′ ( hi(xa) ))2 ,
(14)
where we define the symmetric Gram matrix G ∈ Rn×n with entries Gab = Gba = G(xa,xb), and use λmin(G) to denote its least eigenvalue. We will also use Gmin = mina∈[n]Gaa and Gmax = maxa∈[n]Gaa to denote the minimum and maximum diagonal entries of G, respectively. Since G is positive semidefinite, we see that L̇t ≤ 0, which means that the loss value is indeed non-increasing along the gradient flow trajectory.
Feature learning Compared to the NTK scaling of neural networks, the crucial difference is the 1/m instead of 1/ √ m factor in (2). It is known that under the NTK scaling, due to the 1/ √ m factor, the movement of the feature maps, h1, ..., hm, is only of order O(1/ √ m) before the function value changes by an amount of order Ω(1). While this greatly simplifies the convergence analysis, it also implies a lack of learning of the hidden representations. By contrast, with the 1/m factor in (2), if σ is Lipschitz with Lipschitz constant Lσ , there is |f t2(x)− f t1(x)| ≤ Lσ ĉm ∑m i=1 |ht2i (x)− ht1i (x)|, ∀t1, t2 ≥ 0. Therefore, regardless of m and D,
1
m m∑ i=1 |ht1i (x)− ht2i (x)| ≥ (ĉ)−1(Lσ)−1|f t1(x)− f t2(x)| , (15)
which implies that the average movement of the feature maps is at least on the same order as the change in function value, and thus the NTK undergoes nontrivial movement. In Appendix C, we further justify the occurrence of feature learning using the framework developed in [73].

heading: 3 CONVERGENCE ANALYSIS

heading: 3.1 MODELS WITH A FIXED EMBEDDING
text: Based on (14), to prove that the training loss converges to zero, we need a tighter upper bound on the value of L̇t than 0. Indeed, if G is positive definite, which depends on Φ and the training data, we can achieve this as follows. As a warm-up, if we use activation functions whose derivative has a magnitude uniformly bounded from below by a constant Kσ′ > 0, such as linear, cubic or (smoothed) Leaky ReLU activations, we can directly establish a Polyak-Lojasiewicz (PL) condition [58, 45],
L̇t ≤− 2ĉ2λmin(G) (Kσ′)2 Lt , (16)
and hence Lt ≤ L0e−2ĉ 2λmin(G)(Kσ′ ) 2t, implying 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], σ′ ( hi(xa) ) does not become near zero for too many i ∈ [m] before the loss vanishes. To facilitate a fine-grained analysis, we need the following mild assumption on σ:
Assumption 2. σ is Lipschitz with Lipschitz constant Lσ, and there exists an open interval I = (Il, Ir) ⊆ 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 law of N (0, IdD), i.e., the product distribution of D independent standard Gaussian random variables.
Theorem 1 (Fixed embedding). Suppose that Assumptions 1, 2 and 3 are satisfied, and λmin(G) > 0. Then ∃ĉ0, r and C > 0 such that ∀δ > 0, if ĉ ≥ ĉ0λmax(G)/λmin(G) and m ≥ C(1 + ĉ2/‖y‖2) log (n/δ), then with probability at least 1− δ, it holds that ∀t ≥ 0,
Lt ≤ L0e−rĉ 2λmin(G)t . (17)
The values of ĉ0, r and C depend on I,Gmin, Gmax, Lσ and Kσ′ (but not m, n, d, D, δ, λmin(G)).
The result is proved in Appendix E, and below we briefly describe the intuition. For each a ∈ [n], among all neurons i ∈ [m], we can keep track of the proportion of neurons for which hti(xa) ∈ I . In particular, 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 ci’s are not too small in magnitude. This property of the dynamics is formalized in the following lemma:
Lemma 1. Consider the dynamics ofLt and { hti(xa) } i∈[m],a∈[n] governed by (11) and (14). Assume that λmin(G) > 0, and ∀i ∈ [m], |ci| = ĉ > 0. Under Assumption 2, define
ηt = min a∈[n]
{ 1
m m∑ i=1 1hti(xa)∈I
} ,∀t ≥ 0 and η̃0 = min
a∈[n]
{ 1
m m∑ i=1 1 h0i (xa)∈( 2Il+Ir 3 , Il+2Ir 3 ) } (18)
Then ∀t ≥ 0, there is( ηt ) 3 2 ≥ ( η̃0 ) 3 2 − κ ( λmin (G) , λmax (G) )( (L0) 12 − (Lt) 12 ) /ĉ , (19)
where κ(λ1, λ2) = 9λ2(Ir−Il)
2λ1Kσ′ .
3.1.1 EXAMPLE: SHALLOW NEURAL NETWORKS WHEN n ≤ d
In the case of shallow NNs, G = G(0) ∈ Rn×n, where
G (0) ab :=
1 d (xa) ᵀxb (20)
Thus, G(0) is positive definite if and only the training data set { x1, ...,xn} ⊆ Rd 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 the gradient flow (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].

heading: 3.2 MODELS WITH A HIGH-DIMENSIONAL RANDOM EMBEDDING
text: 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. With this in mind, we consider the type of models where D is large and Φ is a high-dimensional random embedding. While the Gram matrix G in this case is also random, we are interested in its limit as D tends to infinity. Inspired by the analysis above, we desire the following condition on the random embedding: Condition 1 (Concentration of G around a positive definite matrix). There exists a (deterministic) positive definite matrix Ḡ ∈ Rn×n with least eigenvalue λmin(Ḡ) > 0 such that ∀δ, u > 0, ∃Dmin(δ, u) > 0 such that if D ≥ Dmin(δ, u), then P(‖G− Ḡ‖2 > u) < δ. Since the sampling of c and W is independent from the randomness of Φ, Condition 1 is sufficient for us to apply Lemma 1 and obtain the following global convergence guarantee that 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, ∃ĉ0, r and C > 0 such that ∀δ > 0, if ĉ ≥ ĉ0λmax(Ḡ)/λmin(Ḡ), m ≥ C(1 + ĉ2/‖y‖2) log (n/δ) and D ≥ Dmin( 12δ, 12λmin(Ḡ)), then with probability at least 1− δ, it holds that ∀t ≥ 0,
Lt ≤ L0e−rĉ 2λmin(G)t . (21)
The values of ĉ0, r and C depend on I, Ḡmin, Ḡmax, Lσ and Kσ′ (but not m, n, d, D, δ, λmin(Ḡ)).
3.2.1 EXAMPLE: PARTIALLY-TRAINED THREE-LAYER NEURAL NETWORKS (P-3L NNS)
Consider the P-3L NN model defined in (4). In this case, the Gram matrix is G(1), defined by
G (1) ab =
1
D D∑ j=1 σ ( 1√ d zᵀj xa ) σ ( 1√ d zᵀj xb ) (22)
If z1, ...,zm are sampled i.i.d. from a probability measure πz on Rd and fixed during training, then the limiting Gram matrix, denoted by Ḡ(1) ∈ Rn×n, is given by
Ḡ (1) ab = Ez∼πz
[ σ ( 1√
d zᵀxa
) σ ( 1√
d zᵀxb
)] (23)
Thus, for the convergence result, the assumption we need on the limiting Gram matrix is Assumption 4. πz is sub-Gaussian and the matrix Ḡ(1), which depends on the choice of σ and the training set, is positive definite with λmin ( Ḡ(1) ) > 0 and (Ḡ(1))max <∞.
This assumption also plays an important role in the NTK analysis, and it is satisfied if, for example, πz = N (0, Id), 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 and πz is sub-Gaussian (including Gaussian, uniform and Rademacher distributions, for example), 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, ∃ĉ0, r, C1 and C2 > 0 such that ∀δ > 0, if ĉ ≥ ĉ0λmax(Ḡ(1))/λmin(Ḡ(1)), m ≥ C1(1 + ĉ2/‖y‖2) log (n/δ) and D ≥ C2n 2 log(n/δ)/λmin(Ḡ (1))2, then with probability at least 1− δ, it holds that ∀t ≥ 0,
Lt ≤ L0e−rĉ 2λmin(Ḡ (1))t . (24)
The values of ĉ0, r, C1 and C2 depend on I, Ḡ (1) min, Ḡ (1) max, Kσ′ as well as the sub-Gaussian norm of µz (but not m, n, d, D, δ, λmin(Ḡ(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 as the activations of the first hidden layer in the P-3L NN, which illustrates a benefit of jointly having depth and width in NNs. Compared to the NTK analysis [18], our result assumes the same level of over-parameterization, but crucially allows feature training to occur, as discussed in Section 2.2.
Furthermore, by using a multi-layer NN with random and fixed weights as the high-dimensional random embedding, we extend the P-3L NN model to the partially-trained L-layer NN model in Appendix H, under which similar results of global convergence can be proved.

heading: 4 NUMERICAL EXPERIMENTS
text: Additional results and details of the experiments are provided in Appendix I.

heading: 4.1 EXPERIMENT 1: CONVERGENCE OF TRAINING WHEN FITTING RANDOM DATA
text: We train a shallow NN to fit a randomly labeled data set {(x1, y1), ..., (xn, yn)} with d = 20, where we sample each xa i.i.d. fromN (0, I20), with I20 being the 20×20 identity matrix, and each ya i.i.d. uniformly on [− 12 , 12 ] and independently from the xa’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 also echoes the negative result by [71] on the convergence rate of optimizing a shallow NN using population loss, which implies a worsening of the convergence rate as the number of data points increases to infinity.

heading: 4.2 EXPERIMENT 2: BENEFIT OF INPUT EMBEDDING
text: We consider a model defined by (2) and (3) with Φ(x) = vec(xxᵀ) ∈ Rd2 with d = 30, and call this model a shallow NN augmented with quadratic embedding. We compare this model against a plain shallow NN (without extra embedding), both with m = 8192, on a regression task where the target y is given by another shallow NN augmented with quadratic embedding with m = 5 and varying n’s. We see from Figure 2 that the augmented shallow NN indeed achieves lower test error given the same number of training samples, demonstrating the benefit of receiving a good embedding.

heading: 4.3 EXPERIMENT 3: FEATURE LEARNING V.S. LAZY TRAINING
text: We consider the P-3L NN model defined in (4) and (5) with D set to be m (i.e. both hidden layers having the same width, m), and compare it with the same model under NTK and MF scalings as described in Table 1. We adopt the data set used in [69] (more details in Appendix I.3), and train the models by unregularized least-squares regression 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 very little in the model under NTK scaling, as expected [16]. Note also that in the MF scaling, the feature maps h1(x), ..., hm(x) concentrate near 0 at initialization due to the small scaling, but gains diversity during training. Third, we see from Figure 3 that our scaling yields the smallest test errors out of the three, and as n grows the test error also decreases faster in the MF scaling than in the NTK scaling, which indicates the advantage of feature learning compared to lazy training.

heading: 5 CONCLUSIONS AND LIMITATIONS
text: We consider a general type of models that includes shallow and partially-trained multi-layer NNs, which exhibits feature learning when trained via gradient flow, and prove non-asymptotic global convergence guarantees that accommodates a general class of activation functions. For a randomlyinitialized shallow NN, we prove that by performing gradient flow on the input-layer weights, the training loss converges to zero in a linear rate if the number of training data does not exceed the input dimension. For a randomly-initialized multi-layer NN, we prove that by performing gradient flow on the weights in the second-to-last layer, the same result holds except there is no more 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 only the optimization and not the approximation or generalization of NNs. In addition, as the current theoretical results on global convergence neglect the bias terms and assume that the last-layer weights are untrained, and a more general version is left for future work.

heading: A ADDITIONAL NOTATIONS
text: • 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 n ∑n a=1.
• We use bold letters (e.g. x, z, c, y) to denote vectors.
• We use W and {Wij}i∈[m],j∈[D] interchangeably to refer to the same set of parameters.

heading: B CONSISTENCY OF THE SCALING AND GD UPDATE RULE WITH XAVIER INITIALIZATION
text: Consider a three-layer network defined by
f(x) = m∑ i=1 θ (3) i σ ( hi(x) ) (25)
∀i ∈ [m] : hi(x) = m∑ j=1 θ (2) ij σ ( 1√ d d∑ k=1 θ (1) jk xk ) (26)
with weight parameters { θ
(1) jk } j,k∈[m], { θ (2) ij } i,j∈[m] and { θ (3) i } i∈[m] are initialized according to
Xavier initialization, which means that we sample each θ(1)jk i.i.d. from N (0, 1m+d ), each θ (2) ij i.i.d. from N (0, 12m ), and each θ (3) i i.i.d. from N (0, 1m+1 ). If m d, both N (0, 1m+d ) and N (0, 1m+1 ) can be approximated by N (0, 1m ). Then, up to this approximation, by redefining ci = √ mθ (3) i , Wij = √ mθ (2) ij and zjk = √ mθ (1) jk , we can write
f(x) = 1√ m m∑ i=1 ciσ ( hi(x) ) , (27)
∀i ∈ [m] : hi(x) = 1√ m D∑ j=1 Wijσ ( 1√ md zᵀj x ) , (28)
and note that ci,Wij and zjk 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.
Moreover, there is ∂f∂Wij (x) = 1√ m ∂f
∂θ (2) ij
(x). Then, since performing GD on θ(2)ij with step size δ
means updating θ(2)ij according to
θ (2) ij ← θ (2) ij − δ ∂L[f ] ∂θ
(2) ij
, (29)
this is equivalent to updating Wij according to
Wij ← √ m ( θ (2) ij − δ
∂L[f ] ∂θ
(2) ij ) =Wij −mδ
∂L[f ] ∂Wij (x) ,
(30)
which justifies the m factor on the right-hand-side of (9).

heading: C RELATIONSHIP TO THE ABC-PARAMETERIZATION AND FEATURE LEARNING
text: Consider the pseudo-L-layer NN model defined in Section H in the case where D = m >> d. In the language of abc-parameterization defined in [73], our parameterization corresponds to having
a1 =0, a2 = ... = aL = 1
2 , aL+1 = 1
bl =0,∀l ∈ [L+ 1] Furthermore, as we explain in Appendix B, the appropriate learning rate scales linearly with m (as in (9)), which corresponds to having c = −1 (31) Meanwhile, the maximum-update (µP) scaling [73] is characterized by setting
a1 =− 1
2 , a2 = ... = aL = 0, aL+1 =
1 2 (32)
bl = 1
2 , ∀l ∈ [L+ 1] (33)
c =0 (34) Recall the symmetry of abc-parameterization derived in [73], which states that one gets a different but equivalent abc-parameterization by setting
al ← al + θ, bl ← bl − θ, c← c− 2θ (35) Since our parameterization can be obtained from the maximum-update parameterization by applying the transformation above with θ = 12 , they are equivalent in the function space. In particular, for our parameterization, the r parameter defined in [73] can be computed as
r = min{bL+1, aL+1 + c}+ aL+1 + c+ min l=1,...,L {2al − 1l 6=1}
= min{0, 1 + (−1)}+ 1 + (−1) + min{2 · 0− 0, 2 · 1 2 − 1}
=0
(36)
Hence, according to [73], our parameterization exhibits feature learning.

heading: D PROOF OF LEMMA 1
text: Since we assume that G is positive definite and |ci| = ĉ > 0, ∀i ∈ [m], we can derive from (14) that
L̇t =− ĉ 2
m m∑ i=1 n∑ a,b=1 ( f t(xa)− ya ) ( f t(xb)− yb ) σ′ ( hti(xa) ) σ′ ( hti(xa) ) Gab
≤− ĉ2λmin ( G ) 1 m m∑ i=1 n∑ a=1 ( f t(xa)− ya )2 ( σ′ ( hi(xa) ))2 =− ĉ2λmin
( G ) n∑ a=1 ( f t(xa)− ya )2 1 m m∑ i=1 ( σ′ ( hi(xa) ))2 ≤− ĉ2λmin(G)
n∑ a=1 ( f t(xa)− ya )2 1 m m∑ i=1 1hti(xa)∈I (Kσ′) 2
≤− ĉ2λmin(G) (Kσ′)2 n∑ a=1 ( f t(xa)− ya )2 min b∈[n]
{ 1
m m∑ i=1 1hti(xb)∈I
}
≤− 2ĉ2λmin(G) (Kσ′)2Lt min a∈[n]
{ 1
m m∑ i=1 1hti(xa)∈I
}
(37)
Since I is an open interval, ∃ξ > 0 such that we can find a subinterval I0 ⊆ I such that the distance between I0 and the boundaries of I (if I is bounded on either side) is no less than ξ, i.e.,
inf u∈I0,u′∈R\I
|u− u′| ≥ ξ (38)
In particular, we can choose ξ = 13 (Ir − Il) and I0 = (Il + ξ, Ir − ξ). Then there is 1hti(xa)∈I ≥1h0i (xa)∈I0 , hti(xa)∈I
≥1h0i (xa)∈I0 , |hti(xa)−h0i (xa)|<ξ ≥1h0i (xa)∈I0 − 1|hti(xa)−h0i (xa)|≥ξ
(39)
and so
min a∈[n]
{ 1
m m∑ i=1 1hti(xa)∈I } ≥ min a∈[n] { 1 m m∑ i=1 1h0i (xa)∈I0 } − max a∈[n] { 1 m m∑ i=1 1|hti(xa)−h0i (xa)|≥ξ } (40)
Thus, we have L̇t ≤ −2ĉ2λmin(G) (Kσ′)2 Lt (
min a∈[n]
{ 1
m m∑ i=1 1h0i (xa)∈I0 } − max a∈[n] { 1 m m∑ i=1 1|hti(xa)−h0i (xa)|≥ξ }) (41)
Meanwhile, since
ḣti(xa) = −ci n∑ b=1 ( f t(xb)− yb ) σ′ ( hti(xb) ) Gab , (42)
there is
1
m m∑ i=1 ∣∣∣ḣti(xa)∣∣∣ ≤ĉ 1m m∑ i=1 ∣∣∣∣∣ n∑ b=1 ( f t(xb)− yb ) σ′ ( hti(xb) ) Gab ∣∣∣∣∣ ≤ĉ  1 m m∑ i=1 ∣∣∣∣∣ n∑ b=1 ( f t(xb)− yb ) σ′ ( hti(xb) ) Gab ∣∣∣∣∣ 2  12
≤ĉ ( 1
m m∑ i=1 n∑ b=1 ∣∣(f t(xb)− yb)σ′(hti(xb))Gab∣∣2 ) 1 2
≤ĉλmax(G) ( 1
m m∑ i=1 n∑ b=1 ( f t(xb)− yb )2 ( σ′ ( hti(xb)
))2) 12
≤ĉλmax(G)
 ∣∣∣L̇t∣∣∣
(ĉ)2λmin (G)
 1 2
≤λmax(G) (λmin (G))− 1 2 ∣∣∣L̇t∣∣∣ 12
(43)
Therefore,
1
m m∑ i=1 ∣∣hti(xa)− h0i (xa)∣∣ ≤∫ t 0 1 m m∑ i=1 ∣∣∣ḣsi (xa)∣∣∣ ds ≤λmax(G) (λmin (G))− 1 2 ∫ t 0
∣∣∣L̇s∣∣∣ 12 ds (44) Since ∀ξ ∈ R, there is
ξ · 1|hti(xa)−h0i (xa)|≥ξ ≤ ∣∣hti(xa)− h0i (xa)∣∣ , (45)
we derive that, ∀a ∈ [n], 1
m m∑ i=1 1|hti(xa)−h0i (xa)|≥ξ ≤ξ −1 1 m m∑ i=1 ∣∣hti(xa)− h0i (xa)∣∣ ≤C1 ∫ t 0 ∣∣∣L̇s∣∣∣ 12 ds , (46)
where we set C1 = λmax(G) (λmin (G)) − 12 ξ−1 > 0 for simplicity. As a consequence,
max a∈[n]
{ 1
m m∑ i=1 1|hti(xa)−h0i (xa)|≥ξ
} ≤ C1 ∫ t 0 ∣∣∣L̇s∣∣∣ 12 ds (47) Define
η̃t = min a∈[n]
{ 1
m m∑ i=1 1h0i (xa)∈I0
} − C1 ∫ t 0 ∣∣∣L̇s∣∣∣ 12 ds (48) Note that at t = 0, there is η̃t = mina∈[n] { 1 m ∑m i=1 1h0i (xa)∈I0 } . Then, on one hand, we know from (40) and (47) that ∀t ≥ 0,
ηt ≥ min a∈[n]
{ 1
m m∑ i=1 1h0i (xa)∈I0 } − max a∈[n] { 1 m m∑ i=1 1|hti(xa)−h0i (xa)|≥ξ } ≥ η̃t , (49)
Hence, (41) implies that
L̇t ≤− 2ĉ2λmin(G) (Kσ′)2 Ltηt
≤− 2ĉ2λmin(G) (Kσ′)2 Ltη̃t (50)
On the other hand, by the definition of η̃t,
˙̃ηt =− C1 ∣∣∣L̇t∣∣∣ 12
≥− C1 ∣∣∣L̇t∣∣∣ · ∣∣∣L̇t∣∣∣− 12
≥C1L̇t ( 2ĉ2λmin(G) (Kσ′) 2 Ltη̃t )− 12 ≥C2(ĉ)−1L̇t ( Lt )− 12 (η̃t)− 12 ,
(51)
where we set C2 = 2− 1 2C1 (λmin(G)) − 12 (Kσ′)−1 = λmax(G)√
2ξKσ′λmin(G) ≤ nGmax√ 2ξKσ′λmin(G) for simplic-
ity. Therefore, when ηt > 0,
d
dt
( 2
3
( η̃t ) 3 2 ) = ( η̃t ) 1 2 ˙̃ηt ≥ C2(ĉ)−1L̇t ( Lt )− 12 ≥ C2(ĉ)−1 d
dt
( 2Lt ) 1 2 , (52)
which implies that
2
3
( η̃t ) 3
2 − 2 3
( η̃0 ) 3 2 ≥ C2(ĉ)−1 (( 2Lt ) 1 2 − ( 2L0 ) 1 2 ) ≥ −C2(ĉ)−1 ( 2L0 ) 1 2 (53)
and so ∀t ≥ 0, 2
3
( ηt ) 3
2 ≥ 2 3
( η̃t ) 3
2 ≥ 2 3
( η̃0 ) 3 2 − C2(ĉ)−1 ( 2L0 ) 1 2 (54)

heading: E PROOF OF THEOREM 1
text: To apply Lemma 1, we need two additional lemmas to be proved in Appendix E.1 and E.2. The first one guarantees that the loss value at initialization, L0, is upper-bounded with high probability: Lemma 2. ∀δ > 0, if m ≥ Ω ( ĉ2 log ( nδ−1 ) Gmax ) , then with probability at least 1− δ, there is
L0 ≤ ‖y‖2 (55)
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]:
Lemma 3. ∀δ > 0, if m ≥ log(nδ −1)
2(K(I,Gmin,Gmax)) 2 , then with probability at least 1− δ, there is
η0 > K(I,Gmin, Gmax) , (56) where K(I, λ1, λ2) = 16√2πλ2 (Ir − Il) exp { −max{|Il|,|Ir|}(λ1)2 } is a positive number that depends on

heading: I , λ1 and λ2.
text: With these two concentration lemmas, we deduce that, ∀δ > 0, if m ≥ Ω ( (1 + ĉ2) log ( 2n δ )) , then with probability at least δ, there is ∀t ≥ 0, 2
3
( ηt ) 3
2 ≥ 2 3
(K(I,Gmin, Gmax)) 3 2 − √ 2C2(ĉ) −1‖y‖ , (57)
where K(I,Gmin, Gmax) is defined as in Lemma 3. Therefore, if our choice of ĉ satisfies
ĉ ≥ 3 √
2C2‖y‖ (K(I,Gmin, Gmax)) 3 2 , (58)
then there is ∀t ≥ 0, ηt ≥ 2− 23K(I,Gmin, Gmax) > 0 , (59)
in which case (50) gives
L̇t ≤ −2 13 ĉ2λmin(G) (Kσ′)2 LtK(I,Gmin, Gmax) , (60) which will allow us to finally conclude that
Lt ≤ L0 exp { −2 13λmin(G) (Kσ′)2K(I,Gmin, Gmax)ĉ2t } (61)
Note that (60) establishes a PL condition. Several other convergence analyses of NNs have also relied on variants of the PL condition [25, 44, 74].
E.1 PROOF OF LEMMA 2 Proof. Since at initialization, { ci } i∈[m] and { W 0ij } i∈[m],j∈[D] are both sampled i.i.d. and { ci } i∈[m]
have mean zero, we know that ∀a ∈ [n], f0(xa) = 1m ∑m i=1 ciσ ( hti(xa) ) is the sample mean of i.i.d.
random variables with zero-mean. Moreover, since { W 0ij } i∈[m],j∈[D] are sampled from N (0, 1), we
know that ∀i ∈ [m], the random variable ciσ ( hti(xa) ) is sub-Gaussian, with sub-Gaussian norm [68]
‖ciσ ( hti(xa) ) ‖ψ2 ≤ĉ‖σ ( hti(xa) ) ‖ψ2
≤ĉLσ(Gaa) 1 2MSG
≤ĉLσ(Gmax) 1 2MSG ,
(62)
where MSG > 0 is some absolute constant. Thus, by Hoeffding’s inequality [68], ∀a ∈ [n], ∀r > 0, P (∣∣f0(xa)∣∣ ≥ u) =P( ∣∣∣∣∣ 1m m∑ i=1 ciσ ( hti(xa) )∣∣∣∣∣ ≥ u )
≤2 exp { − Ku 2m
‖ciσ ( hti(xa) ) ‖2ψ2
}
≤2 exp { − Ku 2m
ĉ2(Lσ)2Gmax(MSG)2
} ,
(63)
where K is some absolute constant. Hence, by union bound,
P (∑ a ∣∣f0(xa)∣∣2 ≥ ‖y‖2) ≤ n∑ a=1 P (∣∣f0(xa)∣∣ ≥ ‖y‖)
≤2n exp { − K‖y‖ 2m
ĉ2(Lσ)2Gmax(MSG)2
} (64)
Thus, ∀δ > 0, if
m ≥ ĉ2(Lσ)2GmaxK−1(MSG)2‖y‖−2 log ( 2n
δ
) (65)
then with probability at least 1− δ, there is
L0 =1 2 ∑ a ∣∣f0(xa)− ya∣∣2 ≤1
2 ∑ a ∣∣f0(xa)∣∣2 + 1 2 ‖y‖2
≤‖y‖2
(66)
E.2 PROOF OF LEMMA 3
Since each W 0ij are sampled i.i.d. from N (0, 1), we know that ∀a ∈ [n], independently for each i ∈ [m], h0i (xa) follows a Gaussian distribution with mean 0 and variance Gaa. Therefore,
m∑ i=1 1h0i (xa)/∈I0 ∼ Binomial (m, 1− π (I0;Gaa)) , (67)
Hence, by Hoeffding’s inequality, ∀a ∈ [n], ∀r > 0,
P
( 1
m m∑ i=1 1h0i (xa)/∈I0 ≥ 1− π(I0;Gaa) + r ) ≤ exp { −2mr2 } (68)
∀a ∈ [n], choosing r = 12π (I0;Gaa), we then get
P
( 1
m m∑ i=1 1h0i (xa)∈I0 ≤ 1 2 π (I0;Gaa)
) =P ( 1
m m∑ i=1 1h0i (xa)/∈I0 ≥ 1− 1 2 π (I0;Gaa)
)
≤ exp { −1
2 m (π (I0;Gaa))
2 } ≤ exp { −1
2 m ( min b∈[n] {π (I0;Gbb)} )2}
(69)
and so by union bound
P ( η̃0 ≤ 1
2 min b∈[n]
{π (I0;Gbb)} ) =P ( min a∈[n] { 1 m m∑ i=1 1h0i (xa)∈I0 } < 1 2 min b∈[n] {π (I0;Gbb)} )
≤ n∑ a=1 P
( 1
m m∑ i=1 1h0i (xa)∈I0 < 1 2 min b∈[n]
{π (I0;Gbb)} )
≤ n∑ a=1 P
( 1
m m∑ i=1 1h0i (xa)∈I0 < 1 2 π (I0;Gaa)
)
≤n exp { −1
2 m ( min b∈[n] {π (I0;Gbb)} )2}
(70)
Since ∀b ∈ [n], there is Gmin ≤ Gbb ≤ Gmax,
π (I0;Gbb) = 1√
2πGbb ∫ Ir−ξ Il+ξ e − u2Gbb du
≥ 1√ 2πGbb (Ir − Il − 2ξ) inf Il+ξ≤u≤Ir−ξ exp
{ − u 2
Gbb } ≥ 1√
2πGmax (Ir − Il − 2ξ) exp
{ −max{|Il|, |Ir|}
(Gmin)2 } ≥ 1
3 √ 2πGmax (Ir − Il) exp
{ −max{|Il|, |Ir|}
(Gmin)2
} (71)
Letting K(I, λ1, λ2) = 16√2πλ2 (Ir − Il) exp { −max{|Il|,|Ir|}(λ1)2 } > 0, we can then write
P ( η̃0 ≤ K(I,Gmin, Gmax) ) ≤P ( η̃0 ≤ 1
2 min b∈[n]
{π (I0;Gbb)} )
≤n exp { −2m (K(I,Gmin, Gmax))2 } (72) Thus, ∀δ > 0, if m ≥ log(nδ −1)
2(K(I,Gmin,Gmax)) 2 , then with probability at least 1 − δ, it holds that
η̃0 > K(I,Gmin, Gmax) > 0.

heading: F PROOF OF THEOREM 2
text: By Condition 1, we know that ∀δ > 0, if D ≥ Dmin( 12δ, 12λmin(Ḡ)), then with probability at least 1 − 12δ, there is ‖G − Ḡ‖2 ≤ 12λmin(Ḡ), and hence λmin(G) ≥ 12λmin(Ḡ), Gmin ≥ 12 Ḡmin, λmax(G) ≤ λmax(Ḡ) + 12λmin(Ḡ) ≤ 2λmax(Ḡ), and Gmax ≤ 2Ḡmax. We then perform the following analysis conditioned on the event that ‖G− Ḡ‖2 ≤ 12λmin(Ḡ). Since the sampling of { ci } i∈[m] and { W 0ij } i∈[m],j∈[D] is independent from the realization of G, we know from Lemma 3 that if m ≥ log(4nδ −1)
2(K(I, 12λmin(Ḡ),2λmax(Ḡ))) 2 ≥ log(4nδ
−1)
2(K(I,Gmin,Gmax)) 2 , then with
probability at least 1− 14δ, there is
η̃0 > K(I,Gmin, Gmax) ≥ K(I, 1
2 Ḡmin, 2Ḡmax) (73) From Lemma 2, we also know that if m ≥ Ω ( ĉ2 log ( n δ ) λmax(Ḡ) ) ≥ Ω ( ĉ2 log ( n δ ) λmax(Ḡ) ) , then with probability at least 1 − 14δ, there is L0 ≤ ‖y‖2. Therefore, in total, we know that with probability at least 1− δ, the following conditions all hold:
‖G− Ḡ‖2 ≤ 1
2 λmin(Ḡ) , (74)
η̃0 ≥K(I, 1 2 Ḡmin, 2Ḡmax) , (75)
L0 ≤‖y‖2 , (76) in which case, by applying Lemma 1 with G = G(1), we get(
ηt ) 3 2 ≥ ( η̃0 ) 3 2 −K1(λmin (G) , λmax (G))(ĉ)−1 (( L0 ) 1 2 − ( Lt ) 1 2 ) , (77)
where K1(λ1, λ2) = 92λ −1 1 λ2K −1 σ′ (Ir − Il) > 0. Thus, by the definition of K1(·, ·), we know that K1(λmin (G) , λmax (G)) ≤ 4K1(λmin ( Ḡ ) , λmax ( Ḡ ) ) , (78)
and so ∀t ≥ 0,( ηt ) 3 2 ≥ ( η̃0 ) 3 2 − 4K1(λmin ( Ḡ ) , λmax ( Ḡ ) )(ĉ)−1 (( L0 ) 1 2 − ( Lt ) 1 2 ) ≥ ( K(I, 1
2 Ḡmin, 2Ḡmax)
) 3 2 − 4K1(λmin ( Ḡ ) , λmax ( Ḡ ) )(ĉ)−1‖y‖
(79)
Therefore, if our choice of ĉ satisfies ĉ ≥ 8K1(λmin ( Ḡ ) , λmax ( Ḡ ) )‖y‖(
K(I, 12 Ḡmin, 2Ḡmax) ) 3 2
(80)
then there is ∀t ≥ 0, ηt ≥ 2− 23K(I, 1
2 Ḡmin, 2Ḡmax) (81)
Hence, (50) implies that ∀t ≥ 0,
L̇t ≤− 2ĉ2λmin(G)(Kσ′)2Lt2− 2 3K(I,
1 2 Ḡmin, 2Ḡmax)
≤− 2− 23 ĉ2λmin(Ḡ)(Kσ′)2LtK(I, 1
2 Ḡmin, 2Ḡmax)
(82)
and therefore
Lt ≤ L0 exp { −2− 23 ĉ2λmin(Ḡ)(Kσ′)2K(I, 1
2 Ḡmin, 2Ḡmax)
} (83)

heading: G PROOF OF THEOREM 3
text: In view of Theorem 2, it is sufficient to verify that Condition 1 holds for Dmin(δ, u) = Ω ( n2u−2 log(nδ−1) ) , which is given by the following lemma:
Lemma 4. ∀δ ≥ 0, if D ≥ Ω ( n2u−2 log(nδ−1) ) , then with probability at least 1− δ,
‖G(1) − Ḡ(1)‖2 ≤ u (84)
G.1 PROOF OF LEMMA 4
Let Z be a random vector on Rd with law given by πz , and then we can write Ḡ(1)ab = E [σ(xᵀaZ)σ(x ᵀ bZ)] for a, b ∈ [n]. By assumption, Z is sub-gaussian with sub-gaussian norm
‖Z‖ψ2 := sup x∈Sd−1 ‖xᵀZ‖ψ2 <∞ (85)
Thus, ∀a ∈ [n], we have ‖σ(xᵀaZ)‖ψ2 ≤ Lσ‖xᵀZ‖ψ2 ≤ Lσ‖Z‖ψ2 (86)
Hence, by Lemma 2.7.7 in [68], we know that ∀a, b ∈ [n], σ(xᵀaZ)σ(xᵀbZ) is a sub-exponential random variable with sub-exponential norm
‖σ(xᵀaZ)σ(xᵀbZ)‖ψ1 ≤ ‖σ(xᵀaZ)‖ψ2‖σ(x ᵀ bZ)‖ψ2 ≤ (Lσ)2‖Z‖2ψ2 (87)
Then, by Bernstein’s inequality (Theorem 2.8.1 in [68]), since each zj is sampled i.i.d. from πz , we have that ∀a, b ∈ [n] and ∀u > 0,
P (∣∣∣G(1)ab − Ḡ(1)ab ∣∣∣ ≥ u) =P ∣∣∣∣∣∣ 1D D∑ j=1 σ(xᵀaz j)σ(xᵀb z j)− E [σ(xᵀaZ)σ(xᵀbZ)] ∣∣∣∣∣∣ ≥ u 
≤2 exp { −K min { u2D
‖σ(xᵀaZ)σ(xᵀbZ)‖2ψ1 ,
uD
‖σ(xᵀaZ)σ(xᵀbZ)‖ψ1
}}
≤2 exp { −K min { u2D
(Lσ)4‖Z‖4ψ2 ,
uD
(Lσ)2‖Z‖2ψ2
}} ,
(88)
where K > 0 is some absolute constant. In other words, for any δ′ > 0, if D ≥ max { (Lσ) 4‖Z‖4ψ2 log(2(δ′)−1)
u2K ,
(Lσ) 2‖Z‖2ψ2 log(2(δ′)−1)
uK
} , (89)
then we have ∣∣∣G(1)ab − Ḡ(1)ab ∣∣∣ ≥ u with probability at least 1− δ. If we choose u = u′n and δ′ = δn2 , then we get, if
D ≥ max { n2(Lσ)
4‖Z‖4ψ2 log(2n2δ−1) K(u′)2 , n(Lσ) 2‖Z‖2ψ2 log(2n2δ−1) Ku′
} , (90)
then
P ( ‖G(1) − Ḡ(1)‖2F ≥ (u′) 2 ) ≤ n∑ a,b=1 P (∣∣∣G(1)ab − Ḡ(1)ab ∣∣∣ ≥ u′n ) ≤n2 δ
n2
≤δ
(91)
Hence, with probability at least 1− δ, we have
‖G(1) − Ḡ(1)‖22 ≤ ‖G(1) − Ḡ(1)‖2F ≤ (u′) 2 , (92)

heading: H GENERALIZATION TO DEEPER MODELS
text: 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:
f(x) = 1
m m∑ i=1 ciσ ( h (L−1) i (x) ) ,
∀i ∈ [m] : h(L−1)i (x) = 1√ D D∑ j=1 Wijσ ( h (L−2) j (x) ) ,
∀l ∈ [L− 3],∀i ∈ [D] : h(l+1)i (x) = 1√ D D∑ j=1 W̄ (l) ij σ ( h (l) j (x) ) ,
∀j ∈ [D] : h(1)j (x) = 1√ d zᵀj x ,
where W̄ (1), ..., W̄ (L−3) ∈ RD×D and z1, ...,zD ∈ Rd 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.

heading: I FURTHER DETAILS OF THE NUMERICAL EXPERIMENTS
text: In our models, { ci } i∈[m] are sampled i.i.d. from the Rademacher distribution µc = 1 2δ1 + 1 2δ−1,{
zj } j∈[D] are sampled i.i.d. from N (0, Id), and { Wij } i∈[m],j∈[D] are initialized by sampling i.i.d. fromN (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. We choose train the models using 50000 steps of (full-batch) GD with step size δ = 1. When a 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 titanrtx).
I.1 EXPERIMENT 1
We choose σ to be tanh. For each choice of n, we run the experiment with 5 different random seeds, and Figure 1 plots the evolution of the training loss during GD averaged over the 5 runs with m = 8192.
Figure 5 is the same plot as Figure 1 except for with m = 4096. We see that the two two plots agree well.
I.2 EXPERIMENT 2
We choose σ to be ReLU. For each choice of n and each of the two models, we experiment with 5 different random seeds, and Figure 2 plots the test error at the 50000 GD step averaged over the 5 runs ± its standard deviation. In Figure 6, we plot the evolution of the training loss and test error over GD for the two models with m = 2048 and 8192 and different choices of n, averaged over 5 runs with different random seeds. We see in particular that the difference between the two choices of m is negligible, suggesting that it is unlikely to obtain performance improvements with further over-parameterization.
I.3 EXPERIMENT 3
We choose σ to be ReLU and input dimension d = 50. We use a training set of size n = 600 for the results reported in Figure 4. The data set is inspired by [69]: We sample both the training and the test set i.i.d. from the distribution (x, y) ∼ D on Rd+1, under which the joint distribution of (x1, x2, y) is
P(x1 = 1, x2 = 0, y = 1) = 1
4 (93)
P(x1 = −1, x2 = 0, y = 1) = 1
4 (94)
P(x1 = 0, x2 = 1, y = −1) = 1
4 (95)
P(x1 = 0, x2 = −1, y = −1) = 1
4 (96)
(97)
and x3, ..., xd each follow the uniform distribution in [−1, 1], independently from each other as well as x1, x2 and y.
Figures 7 and 8 are the same plot as Figure 4 except for setting n = 400 and 800, respectively. We see that as n increases, test error improves for all three models, and the features learned by the pseudo-3L NN (in our scaling) become “sharper” visually.

Review summaries:
In this paper, the authors studied the optimization problem of shallow and deep neural network. They showed that under the non-degeneracy condition on certain Gram matrix, gradient descent (GD) can converge to 0 training loss efficiently. One important difference with the existing NTK and mean-field literatures is that a different scaling factor was used in this paper. Experiment results show that neural network with this scaling is different from NTK and mean-field scale, while it is still able to do feature learning.

The authors consider the training dynamics of neural networks under a modification of the mean field parameterization in the high dimensional (d >= n) setting.  They provide an especially simple global convergence result under very mild assumptions (beyond the more restrictive d>=n one).  The proof is particularly clean and simple in the setting of leaky-relu type activations, and somewhat more technical arguments are required to get the result to hold for more general activations.  The d>=n assumption is key, as their results depend upon the minimal eigenvalue of the empirical kernel matrix defined by the inner product K(x,x') = <x, x'>, which is strictly positive when d>=n but not necessarily otherwise.  They extend the results to the feature map/embedding setting, assuming the data come from feature maps \phi(x) where \phi(x) lies in a sufficiently high dimensional space so that the kernel defined over embeddings is positive definite provided the dimension of the embedding is large enough. 

This paper studies the optimization of a pseudo-three-layer network. The first (output) layer is scaled by a law of large number (LNN) scaling, while the second layer is scaled by a central limit theorem (CLT)-scaling. The inputs are fixed random feature embeddings. Only the weights of the second layer are learnable. They proved that when the second layer is sufficiently over-parameterized (wrt the number of samples), GD flow converges to a zero-loss solution exponentially fast. In particular, the authors claim that, unlike the NTK regime, their setting exhibits feature learnings.