On feature learning in neural networks with global convergence guarantees
Authors: Zhengdao Chen
Eric Vanden-Eijnden
Joan Bruna
Venue: ICLR 2022 Poster (2022)
Keywords: neural networks
feature learning
gradient descent
global convergence

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.

ON FEATURE LEARNING IN NEURAL NETWORKS WITH GLOBAL CONVERGENCE GUARANTEES
Zhengdao Chen
Eric Vanden-Eijnden
Courant Institute of Mathematical Sciences New York University New York, NY 10012, USA zc121 60nyu.
Courant Institute of Mathematical Sciences
New York University
New York, NY 10012, USA
eve2@cino nyu. edu
Joan Bruna
Courant Institute of Mathematical Sciences and Center for Data Science
New York University
New York, NY 10012, USA
brunaš cims. nyu. edu
ABSTRACT
We study the optimization of wide neural networks (NNs) via gradient flow (GF) in setups that allow feature learning while admitting non-asympotiic global con- vergence guarantees. First, for wide shallow NNs under the mean-field scaling and with general class of activation functions, dimension is no less than the size of the training set, the training loss converges to zero at linear rate under GF. Building upon this analysis, we study model of wide multi-layer NNs whose second-to-last layer is trained via GF, for which we also prove a linear-rate convergence of the training loss to zero, but 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.
prove that when the input
INTRODUCTION
The training of neural networks (NNs) is typically a non-convex optimization problem, but remarkably, simple algorithms like gradient descent (GD) its variants 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- vneerparammttrizztionn), for which suitable ways to scale the parameters by the widths. For example, under "1/ Vwidth" scaling of the weights, the GD dynamics of wide NNs can be approximated by the lincarized dynamics around initialization, and as the widths tend to infinity,
derive ffnite-width limits under
obtain the Newral 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 crucial ingredient the success of deep learning, and is therefore not adequate for explaining the success of NNs 27
Meanwhile, for shallow (i.e.. he--hhdde--lyeee NNs, if we choose "I / width" scaling, we can derive an alternative mean-field (MF) limit as the widths tend to infinity. Under this scaling, feature learning occurs even the infinite-width limit, and the training dynamics can described by the Wasserstein gradient flow of probability measure on the space of the parameters, which converges to global minimizer the loss function under certain conditions [59 50 63 [4]. 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-paree wver-parameterized NNs. First, the global convergence guarantees for shallow NNs only hold in the infinite- infinite-width width limit (i.e. they are asymptotic). While [12] studies the deviation between finite-width NNs and their infinite-wide limits during training, the analysis is done only asymptotically to the

---

next order in width. Second, 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., layer NNs [4 62 52 23 [7 there is less concensus on what the right model should be than for the shallow case. In summary, we still lack 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 type of partially-trained multi-layer NNs, and obtain theoretical guarantees of linear-rate global convergence.
1.1 SUMMARY OF MAIN CONTRIBUTIONS
We consider the scenario of training NN models to fit a training set of n data points in dimension d. where the model parameters 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 S 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 numerical experiments to demonstrate that model exhibits feature learning and can achieve better generalization performance than its NTK counterpart.
1.2 RELATED WORKS
Over-parameterized NNs, NTK and lazy training. Many recent works have studied the opti-
mization landscape of NNs and the benefits of premettraettiiiooooonn 124 67 60, 66 64 54 65 42]
9 32 20, 2T 61 10). One influential idea is the Neural Tangent Kernel (NTK) (77, which characterizes the behavior of GD the infinite-widh width limit of NNs under particular scaling of the parameters (e.g. for shallow NNs, replacing 1 /m with 1/vm m in ①.. 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 relatively large scaling of the parameters at initialization, the hidden-layer features do not move significantly [6). For this reason, the NTK scaling has been called the lazy-training regime, as opposed to a feature-learning or rich regime [16 the differences between the two regimes both in theory 18
40 76
13
75 17
39
72
26 Several works have investigated
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 mean-field scaling (as in (I), where the infinite-wied width limit is analogous the thermodynamic
49 70). Thanks to the interchangeability of the parameters, the neural network is equivalently characterized by 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. [7I] 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) variation-na 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
or hydrodynamic limit of interacting particle systems (59 50 63
2

---

guarantees 4 62 51
but questions remain. First, due to the presence of different symmetries in a multi -ayer network compared to a shallow network [57], the limiting object at the infinite-widhh limit is often quite complicated. Second, it has been pointed out that under the MF
52
23
57 22],
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 [3]. In
addition, while another line of work develops MF models of residual models [ 46 21 36
we are
interested in multi-layer NN models with a large width in every layer.
Feature learning deep NNs. Ref. [ 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. [[T]] studies the optimization landscape and generalization properties of hierarchical model that is similar to ours in spirit, where an untrained embedding the input is passed into a trainable shallow model, and prove an improvement sample complexity in learning polynomials by having neural network outputs 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 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 not well-understood. Our multi-layer NN models adopt an equivalent scaling (see Appendix CC. and our focus is on proving non-asymptotic convergence guarantees for its partial training under GF.
2
PROBLEM SETUP
2.1 MODEL
We summarize our notations in Appendix A [ri,..., _ E Ω denote a generic input data vector. A shallow NN model under the MF scaling can be written as:
Let Ω S. TRd denote the input space, and let æ
m
f(x)
Σ
GO
W W
where is the width, W E Rex and c [c1 Cml E Rm the first- and second-layer weight parameters of the model, and a R + R is the activation function. For simplicity of presentation, we neglect the bias terms. In this paper, study more general type of models with the following form:
1 m
f(x)
Σ Go c (h,(z))
(2)
Vi E [m] : h1(x)
(x)
(3)
where W E R and c = [c.. Cm] 1Tmh are parameters of the model, and .i..... &D are a set of functions from Ω to IR that we call the embedding. Each of hi,..., fy is function from Ω to R, and we will refer to them as the (hidden-lay layer) feature map or activations. For simplicity, we write Φ (0,(x), ........0111 (D()) RR D. We consider two types of the embedding, Φ as described below:
Fixed embedding D is fixed and Φ is deterministic. In the simplest example, we D d and 6,(x) = zj. Vjj E [D]. and recover the shallow NN model in I More generally, our definition includes cases where vector in D. This can be understood as input pre- pre-processing or feature engineering.
deterministic transformation of an input vector in Ω into an embedding
D is large and & is random. For instance, can sample each zj i.i.d. in IRd and set (5 = o ( equivalent to setting di,,.., em as the hidden-layer
High-dimennsionll random embedding
3

---

activations of a shallow NN with mondomly-initializn first-layer weights. Then, the model becomes
1 Σ Go (ha(a))
f(x) =
(4)
Vi [m]
(x) VD Σ
W 1j0
1 [22
(5)
Thus, we obtain a 3-layer feed-forward NN whose first-layer weights are random and fixed, and we call it 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
23 (1) D1/D instead of 1 D in (5). We show in Appendix B that when is homogeneous, this scaling consistent with the Xavier initialization of neural network parameters up to reparameterization (30] (5)). We also show in Appendix C that in certain cases this scaling is equivalent to the maximum-update parametcrization proposed in [73]. Numerical experiments that compare different scalings are described in Section
adopted in 4 62 51 57
2.2 TRAINING WITH GRADIENT FLOW
Consider the scenario of supervised least-squares regression, where we are given a set n training data points together with their target values, (I11 y1),
-.. (In; Un) Q R. We fit our models by
minimizing the empirical squared loss:
cifl 2 Σ (f(xa) ya)²
(6)
To do so, we first initialize the parameters randomly by sampling each c and Wi iid. from probability measures ne E (RR) and Tw P(RP), respectively, and then perform GD on W. For simplicity of analysis,
leave untrained, and further assume that
Assumption 1. Te = ()d( + (dd)(dc) for some > O independent from m, which is the law of scaled Rademacher random variable.
If a is Lipschitz, it is differentiable almost everywhere, and we write (xx to denote the derivative of when it is differentiable and 0 otherwise. When is differentiable at h; (2x), there is
0S (æ) mvD caa (h)(
1
ow
and the gradient of the loss function with respect Wij given by
ac[s] 1
co (((aa)) ya) o' (h:(
(8)
aW
mvD
Thus, we can perform GD updates on W according to the following rule: vi E [m] and Vj E [D),
dC]f Wy + W3y - md ow; = Wij
C.
(f(x)) - yo) o' (((h( (xa)
(9)
where 5 >> 0 is the step size. As we discuss in Appendix 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 2 0 to denote time elaposed during training, the time-derivative of the parameters is given by
B this is consistent with the standard GD)
W
Σ (f))))) - va) o (h!(#a))
(10)
where f denotes the output function and hf...., hi, denote the hidden-layer feature maps determined by the parameters at time t. Then, induced by the evolution of W each h evolves according to
Wa E Ω : h((x)= = VD
We (x) = -c Σ (x, x.) (f*(x)) ( ya) o' ((hh())))) (11)
4

---

where we define a kernel function 9 1: x x - R
D
1 ( (0(2))T0(0')
21x = D j=1
_
(12)
Accordingly, the output function f satisfies
1 m
Ve E 2
(x)
Σ
oo (M)(
m
(13)
62
o' (h((a))
Σ g(x, (f'(z))) yo) o' o'(x((z)))
m
Thus, the loss function =² := f evolves according to
-= Σ (fa())) - ya) (xx
a=1
11
Geb (f)( - ya) (f)())) yb) o' ((h((za))d((h(z)) (14)
m
Amin (G) m
(M(a) - yya)² ya)² (o'(h((z)))))
m
where we define the (symmetric) Gram matrix G E Rux* with entries Gab Gba G(x x.). and use (G) to denote its least cigenvalue. We will also Gmin min Gan and Gmax = max Gaa to denote the minimum and maximum diagonal entries of G, respectively. Since G' is positive semi-definite, non-increasing along the GF trajectory.
that L < 0, which means that the loss value is indeed
Feature learning Compared the NTK scaling of neural networks, the crucial difference the 1/m factor in (2) instead of 1/vmm. It is known that under the NTK scaling, due to the 1/vm factor, the movement of the feature maps, luy....,hm hm is only of order ((I/vm) while the function value changes by an amount of order (2(1). While this greatly simplifies the convergence analysis, it also implies that the hidden-layer representations are not being learned. In contrast, with the
1 1/m factor in (2). if is Lipschitz with Lipschitz constant T. there is [f*x - fx)) <
Loc m | (x: h/2 (x)), Vt,, to 2 0. Therefore, regardless of m and D,
I (z) 1: [x) 2 (2) (L.) (x)) (z)
(15)
which implies that the average movement 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. Appendix C we further justify the occurrence of feature learning using the framework developed in [73
3 CONVERGENCE ANALYSIS
3.1 MODELS WITH A FIXED EMBEDDING
To prove that the training loss converges to zero, we need a lower bound on the absolute value of Ê. Indeed, if G is positive definite, which depends on & and training data, we can establish one in the following way. First, as simple case, if we use activation function whose derívative's absolute value is uniformly bounded from below by a constant K2 > O, such as linear, cubic or (smoothed) Leaky ReLU activations, can derive Dolyak-Lojasiewice (PL) condition [5 45 from 14) directly,
Č < 282 Amin (G) (Ka))² ct
(16)
5

---

which implies L, S Lue 28 (0))(())) indicating that the training loss decays to 0 at linear rate.
For more general choices of the activation function, challenge is to guarantee that, heuristically speaking, for each a E [n], o' (h;((x)) does not become near zero for too many E [m] before the loss vanishes. To facilitate finer-grained analysis, we need the following mild assumption on a:
Assumption 2. a is Lipschitz with Lipschitz constant Lo. and there exists an open interval T = (J,, I,) S IR. on which a is differentiable and |a') is lower-bounded by some K > O.
Intuitively, I is active region of a, within which the derivative has magnitude bounded away from zero. This assumption is satisfied by the majority of activation functions in practice, including smooth ones such 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.
Tw is the D-dimensional standard Gaussian distribution, i.e., each Way is sampled independently from a standard Gaussian distribution.
Assumption 3.
Theorem 1 (Fixed embedding). Suppose that Assumptions Then Beg r and > 0 such that 48 > O, if 2 &gdmax (G] then with probubility at least 1 - 6, it holds that VI 2 O,
2 and 3 are satisfied, and Xmin (G) > 0. Amin and m 2 C(1 + log (n/8).
L, S Loe -o Ammn(G))
(17)
Here, Co, r and C depend I, Gmin, Gmax 11 Lo and Ko (but not on m, d, D. 8, or Amin (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 E [n)), we can keep track of the proportion of neurons (among all i [ml) for which h (axa) E I. We show that the proportion 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 cy'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 and (1(((a)) / [e(m],ae[a] governed by and 14 Assume
that Amin (G) > O, and Vi E [m]] |csl = >> 0. Under Assumption
2
define
m
m 1 (e )E( 2 (422r)
1/' min [E(n]
V1 and AiGNJET
min aE(n)
m
m
(18)
Then VL 2 O, there is
(n/) 2
(77)
xi (Amin (G), Amax -((()))))(()))))
(19)
where k(A1;>2) )2) = I12(T,-1))
2A, Ka
3.1.1 EXAMPLE: SHALLOW NEURAL NETWORKS WHEN n S d
In the case of shallow NNs under the MF scaling, G = -(()) E Raxx, where
(0)
(20)
Thus, (( is positive definite if and only the training data set { 2. S Rdd consists of linearly- independent vectors, which is possible (and expected if the training data are sampled independently from some non-degenerate distribution) when n S In that case, Theorem iimplies
Corollary 1 (Shallow NN with < d). Suppose that Assumptions 7 and Bare satisfied. If the training data are linearly-indepent shallow NN, the training loss converges to zero at a linear rate.
vectors, then under GF (10) on the first-layer weights of the
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 [88. modifications to the GD algorithm 34
13 53
restrictions to certain special tasks (43

---

3.2 MODELS WITH HIGH-DIMENSIONA, RANDOM EMBEDDING
A clear limitation of Corollary is that it is only applicable when n. S d, since otherwise the Gram matrix ( cannot be positive definite. This motivates us to consider the use high-dimensional embedding & to lift the effective imput dimension. In particular, we focus 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 deterministic and positive definite limit as D tends infinity:
Condition (Concentration of G around a positive definite matrix). There exists a (determinis- tic) positive definite matrix G R= with least eigenvalue Amin (G) 0 such that Võ, > 0, mii (8, u) > such that if D 2 Dmin (ò, u), then P[|| - Gll2 > u) < <..
Condition 1 is sufficient for us to apply Lemma 1 and obtain the following global convergence guarantee, which extends Theorem to models with a nigh-dimensional random embedding. The proof is given in Appendix
Theorem 2 (High-dimensional random embedding). Under Assumptions 2 3 and Condition 7 3dg and C > O such that 4g >> 0, if à 2 Codmax /min (ain(G),,, m 2 C e log (n/8) and
D 2 minn (2,, 2 Amin (G)), then with probability at least 1 holds that Vt 2 0.
Ls S Lge- Ammin (G)n
(21)
Here, Cg r and C depend on I, Gmin: Gmaxs || Lo and K (but not m, n, d, D. or Amin (G)).
3.2.1 EXAMPLE: PARTIALLY-TRAINEI TIREE-LAYER NEURAL NETWORKS (P-3L NNs)
Consider the P-3L NN model defined in In this case, the Gram matrix is ( defined by
D E
( =
1 (
(22)
If 21,..., Z. are sampled i.i.d. from a probability measure т= on Rdd and fixed during training, then the limiting Gram matrix, denoted by G(1) E RRnn, is given by
G
Eg
Ta( (
1 =T To (
(23)
Thus, for the convergence result, the assumption we need on the limiting Gram matrix is Assumption 4. Ty is sub- Gaussian and the matrix G(1), which depends on the choice of a and the training set, is positive definite with Amin (G
and (G1))) < 00.
This assumption also plays an important role in the NTK analysis, and it is satisfied if, for example, Tx is the d-dimensional standard Gaussian distribution, no two data points are parallel, and a is either
the ReI.U function [19 or analytic and not a polynomial [8] When Assumption 4 is satisfied, as
long as is Lipschitz, we can use standard concentration techniques to verify Condition [ Thus, Theorem 2 jimplies that
Theorem 3 (N), Under Assumptions 7 3 and 4 2čo C₁ and C2 > 0 such that V8 > 0,
è 2 čo (D0m))(())) (/ /mmn (G(1), m 2 C1 + log (n/ and D 2 Czn² log(n/))) Amin (Q
then with probability at least 1 - 8, it holds that Vt 2 O,
L. S Lae - Amni (G
(24)
Here, Cg r, C1 and C2 depend on I, (()) ((1 ||y|], Ko as well as the sub-Gaussian norm of μ
min
(but not on m, n. d, D, or Amin \minn(())) (G())).
The proof is given in Appendix G Compared to Corollary I for shallow NNs, a highlight of Theorem that the requirement of n S is no longer needed. This demonstrates an advantage of the high- dimensional random embedding realized by the first hidden layer in the P-3L NN, thus illustrating 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- weer-praamtterizatoon but crucially allows feature training to occur, which we discuss in Section 2.2 and support empirically in Section 4
Furthermore, by using layer NN with random and fixed weights as the high-dimensionall 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-lass layer via GF.
7

---

Table 1: Three different scalings of the partially-trainet 3L NN model considered in Experiment 3.
Model
Ours
NTK
MF [4 62]
23
f(x)
2iz1 0:0 (hi(x))
Co (h(()))
Ca (1,(z) =i=1
lu(#)
2= Wejo
m
j=1 Wijor va
m
ž Wi V
m
Vm W
Wig
m8
0c
W
Bck aw
W
m²²8 oc
aw
aw
NUMERICAL EXPERIMENTS
Additional results and details of the experiments are provided in Appendix !|
4.1 EXPERIMENT 1: CONVERGENCE OF TRAINING WHEN FITTING RANDOM DATA
We train shallow NNs to fit a randomly labeled data set {(21, in)..., (æn, yn)} with d = 20. Specifically, we sample each a i.i.d. with every entry sampled independently from standard Gaussian distribution, and each Ya i.i.d. uniformly on 1- ğ and independently from the a, 's. We see from Figure I that the convergence happens at nearly linear rate when 20 and 40, and the rate decreases as n becomes larger. This is coherent with our theoretical result (Corollary [ and interestingly also echoes a prior result that the convergence rate of optimizing 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.
4.2 EXPERIMENT 2: BENEFIT OF INPUT EMBEDDING
We consider a model defined by ( and (3) with d = 30 and
= =vc=== E Rdd, which we call shallow NN augmented with quadratic embedding. We compare this model against 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 given by another shallow NN augmented with quadratic embedding with m = 5. We see from Figure thaat the augmented shallow NN achieves lower test error given the same number of training samples, demonstrating the benefit of good embedding.
4.3 EXPERIMENT 3: FEATURE LEARNING V.S. LAZY TRAINING
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-yayer NN models under NTK and MF scalings, as we define
in Table based on prior literature [37 4 52 62 23), which undergo partial training in the same fashion. We adopt the data set used in 669 (more details in Appendix [.3] and train the models by minimizing the unregularized squared loss for varying rt's and m's.
First, we from the top-left plot in TigureAthare consistently across different m. the training loss converges at 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 h2 (x), hm (x)) concentrate near 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 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.
CONCLUSIONS AND LIMITATIONS
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 mnndomly-iiiiiliizedd shallow NN in the MF scaling that is wide enough, we prove that by performing GF on the imput-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 randomly-initialized multi-layer NN with large widths, prove
8

---

101
P-3L NN
10-3-
1.0
train loss
test
test
10
0.6
0.0
0.4
0.4
0.2
0.2
0.0
0.0
100
150
700
250
200
400
6000
800
GD step
11
Figure/ 1: Training loss v.S. number of GD steps for dif- ferent in Experiment I.
Figure 2: Test error v.s. n in Experiment 2 by the two models with 8192.
Figure 3: Test error v.8. in Ex- periment by the three models with 8192.
P-3L NN
P-3L NN (NTK)
P-3L NN (MF)
00-2
10
20"h
03"1
10
10
00
2182
30031
0000
3100
6C201
20030
0000
GD step
GD stop
GD step
-0-
h1(x
Figure 4: Results of Experiment 3 when n = 600. Each column corresponds to different scaling of the P-3L NN model, as described in Table [ 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.c.. neuron in the second hidden layer,) and the and y-coordinates equal hy (xx) and h,(x2), respectively, where Φi is an input from the training set and Ø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 requirement 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, 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.
ACKNOWLEDGMENTS
The authors acknowledge support from the Henry MacCracken Fellowship. NSF RI-1816753, NSF CAREER CIF 1845360, NSF CHS-1901091 and NSF DMS-MoDL. 2134216.

---

REFERENCES
[1] Zeyuan Allen- -he and Yuanzhi Li. Backward feature correction: How deep learning performs deep learning. arXiv preprint arXiv:2001.044133 2020.
[2] Zeyuan Ahu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparame- terized neural networks, going beyond two layers. In H. Wallach, H. Larochelle, A. Beygelzimer, d'Alché-Buee E. Fox, and R. Garnett, editors, Advances Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[3] Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via pver-parameterizaaion. In International Conference on Machine Learning, pages 242-252, 2019.
[4] Dyego Araújo, Roberto Oliveira, and Daniel Yukimura. A mean-field limit for certain deep neural networks. arXiv preprint arXiv:1906.0011 2019.
[5] 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 MXXiv:9901.08888 2019.
[6] Francis Bach. Breaking the curse of dimensionality with convex neural networks. The Journal of Machine Learning Research, 18()):629-681, 2017.
[7] Yu Bai and Jason D. Lee. Beyond linearization: On quadratic and higher-order approximation of wide neural networks. In International Conference on Leurning Representations, 2020.
[8] Peter Bartlett, Dave Helmbold, and Philip Long- Gradient descent with identity initializa- tion efficiently learns positive definite linear transformations by deep residual networks. In International conference on machine learning, pages 521-530. PMLR, 2018.
[9] Niladri S Chatterji and Philip M Long. Finite-sample analysis of interpolating linear classifiers in the overparameterizeed regime. arXiv preprint arXiv:2004. /2019, 2020.
[10] Niladri Chatterji, Philip M Long, and Peter L Bartlett. When does gradient descent with logistic loss interpolate using deep networks with smoothed relu activations? arXiv preprint MxXiv:2102.04998, 2021.
[11] Minshuo Chen, Yu Bai, Jason D Lee, Tuo Zhao, Huan Wang, Caiming Xiong, and Richard Socher. Towards understanding hierarchical learning: Benefits of neural representations. arXiv preprint arXiv:2006. 2020.
[12] Zhengdao Chen, Grant Rotskoff, Joan Bruna, and Eric Vanden Eijjnen A dynamical central limit theorem for two-layer neural networks. Advances in Neural Information Processing Systems, 33, 2020.
[13] Lenaic Chizat. Sparse optimization measures with over-parameterized gradient descent. arXiv preprint arXiv: 1907. 10300, 2019.
[14] Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. In Advances in Neural Information Processing Systems, pages 3036-3046, 2018.
[15] Lénaic Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. arXiv preprint urXiv:2002.04488, 2020.
[16] Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable program- ming. In Advances in Neural Information Processing Systems, pages 2937-2947, 2019.
[17] Simon Du and Wei Hu. Width provably matters in optimization for deep linear neural networks. In International Conference Machine Learning, pages 1655-1664. PMLR, 2019.
10

---

[18] Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 1675-1685. PMLR, 09-15 Jun 2019.
[19] 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: 1./....0 1.03804, 2018
[20] Ethan Dyer and Guy Gur-Ari. Asymptotics of wide networks from feynman diagrams. In International Conference on Learning Representations, 2020.
[21] Weinan E, Chao Ma, and Lei Wu. Machine learning from a continuous viewpoint, i. Science China Mathematics, 63(1 1):2233-226,, 2020.
[22] Weinan E and Stephan Wojtowytsch. On the banach spaces associated with multi-layer relu net- works: Function representation, approximation theory and gradient descent dynamics. CSIAM Transactions on Applied Mathematics, ((3):377-44)) 2020.
[23] Cong Fang, Jason D Lee, Pengkun Yang, and Tong Zhang. Modeling from features: mean-field framework for over-parame wver-parameterized deep neural networks. arXiv preprint arXiv:2007.01452, 2020.
[24] Daniel Freeman and Joan Bruna. Topology and geometry of half-rectified network optimiza- 2016.
tion. arXiv preprint
[25] Spencer Frei and Quanquan Gu. Proxy convexity: A unified framework for the analysis of neural networks trained by gradient descent. arXiv preprint arXiv:2106.1377 2021.
[26] Mario Geiger, Leonardo Petrini, and Matthieu Wyart. Landscape and training regimes in deep learning. Physics Reports, 2021.
[27] Mario Geiger, Stefano Spigler, Arthur Jacot, and Matthieu Wyart. Disentangling feature and lazy learning deep neural networks: an empirical study. arXiv preprint arXiv:1906.080%) 2019.
[28] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Limitations of lazy training of two-layers neural network. In Advances in Neural Information Processing Systems, pages 911 0111-9121, 2019.
[29] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. When do neural networks outperform kernel methods? arXiv preprint Stti::2006.144009 2020.
[30] Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Yee Whye Teh and Mike Titterington, editors, Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statisties, volume 9 of Proceedings of Machine Learning Research, pages
Chia Laguna Resort, Sardinia, Italy, 13-15 May
2010. PMLR.
[31] Sebastian Goldt, Madhu Advani, Andrew M Saxe, Florent Krzakala, and Lenka Zdeborová. Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup. In Advances in Neural Information Processing Systems, pages 9979-6989, 2019.
[32] Anna Golubeva, Guy Gur-Ari, and Behnam Neyshabur. Are wider nets better given the same number parameters? In International Conference on Learning Representations, 2021.
[33] Boris Hanin and Mihai Nica. Finite depth and width comrectins to the neural tangent kemne. In International Conference on Learning Representations, 2020.
[34] Kaitong Hu, Zhenjie Ren, David Siska, and Lukasz Szpruch. Reaanfiiel langevin dynamics and energy landscape of neural networks. arXiv preprint arXiv:1905.099 077669 2019.
[35] Jiaoyang Huang and Horng- Tzer Yau. Dynamics deep neural networks and neural tangent hierarchy. arXiv preprint arXiv: 0009000855000 2019.

---

[36] Jean-François Jabir, David Siška, and Lukasz Szpruch. Mean-fiell neural odes via relaxed optimal control. arXiv preprint arXiv:1912.0547,, 2019.
[37] 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.
[38] Adel Javanmard, Marco Mondelli, and Andrea Montanari. Analysis of two-layer neural network via displacement convexity. The Annals of Statistics, 48(6):3619-3442, 2020.
[39] Kenji Kawaguchi. Deep learning without poor local minima. arXiv preprint arXiv:1605.07110, 2016.
[40] Jachoon Lec, Yasaman Bahri, Roman Novak, Samuel S Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks as gaussian processes. arXiv preprint arXiv:1711.00165. 2017.
[41] Jachoon Lee, Samuel S. Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein, Finite versus infinite neural networks: an empirical study. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Anmual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020.
[42] Yuanzhi Li and Yingyu Liang. Learning overparameterrized neural networks via stochastic gradient descent on structured data, arXiv preprint arXiv: 01204 2018.
[43] Yuanzhi Li, Tengyu Ma, and Hongyang R. Zhang. Learning over-parametrized two-layer neural networks beyond ntk. In Jacob Abernethy and Shivani Agarwal, editors, Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pages 2613-2682. PMLR, 09-12 Jul 2020.
[44] Chaoyue Liu, Libin Zhu, and Mikhail Belkin. Loss landscapes and optimization in over- parameterized non-linear systems and neural networks. arXiv preprint arXiv:2003.0330, 2020.
[45] Stanislaw Lojasiewicz. A topological property of real analytic subsets. Coll. du CNRS, Les équations aux dérivées partielles, 117(87-89):2, 1963.
[46] Yiping Lu, Chao Ma. Yulong Lu, Jianfeng Lu, and Lexing Ying. A mean-field analysis of deep resnet and beyond: Towards provable optimization via overparameterization from depth. arXiv preprint 20Xi::2003.05508, 2020.
[47] Tao Luo, Zhi-Qin John Xu, Zheng Ma, and Yaoyu Zhang. Phase diagram for two-layer relu neural networks at infinite-width limit. Journal of Machine Learning Research, 22(71):1 1-77,, 2021.
[48] Chao Ma, Lei Wu, et al. The barron space and the flow-induced function spaces for neural network models. Constructive Approximation, 2021.
[49] Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. arXiv preprint arXiv:1902.00 2019.
[50] Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115633):E7665- E7671, 2018.
[51] Phan-Minh Nguyen. Mean field limit of the learning dynamics of multilayer neural networks. arXiv preprint arXiv: 1902002880 2019.
[52] Phan-Minh Nguyen and Huy Tuan Pham. A rigorous framework for the mean field limit of multilayer neural networks. arXiv preprint arXiv:2001. 11443, 2020.
12

---

[53] Atsushi Nitanda, Denny Wu, and Taiji Suzuki. Particle dual averaging: Optimization of mean field neural networks with global convergence rate analysis. arXiv preprint arXiv:2012.15477, 2020.
[54] Samet Oymak and Mahdi Soltanolkooabi. Overparameterized nonlinear learning: Gradient descent takes the shortest path? In International Conference Machine Learning, pages 1951-4960. PMLR, 2019.
[55] 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):84-105, 2020.
[56] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, al. Pytorch: An imperative style, perf-oaefo deep learning library. arXiv preprint arXiv:1912.01703, 2019.
[57] Huy Tuan Pham and Phan-Minh Nguyen. Global convergence of three-layer neural networks in the mean field regime. ICLR, 2021.
[58] Boris Polyak. Gradient methods for the minimisation of functionals. Ussr Computational Mathematics and Mathematical Physics, 3:864-878, 1963.
[59] Grant Rotskoff and Eric Handen-Eijnden Parameters as interacting particles: long time convergence and asymptotic error scaling of neural networks. In Advances in Neural Information Processing Systems, pages 7 7146-7155, 2018.
[60] Itay Safran and Ohad Shamir. Spurious local minima are common in two-layer ReLU neural networks. International Conference on Machine Learning, pages 4433-4441, 2018.
[61] Stefano Sarao Mannelli, Eric Handen-Eijnden and Lenka Zdeborová. Optimization and gen- eralization of shallow neural networks with quadratic activation functions. In Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 13445-13455 Curran Associates, Inc., 2020.
[62] Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of deep neural networks. arXiv preprint arXiv: 1903.04440, 2019.
[63] Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: A law of large numbers. SIAM Journal on Applied Mathematics, 80(2):225-752, 2020.
[64] Mahdi Soltanolkotabi, Adel Javanmard, and Jason D Lee. Theoretical insights into the op- timization landscape of over-parameterizee shallow neural networks. IEEE Transactions on Information Theory, 65(2):742-769, 2018.
[65] Daniel Soudry and Yair Carmon. No bad local minima: Data independent training error guarantees for multilayer neural networks. arXiv preprint 1605vv1605.0836) 2016.
[66] Yuandong Tian. An analytical formula of population gradient for two-layered ReLU network and its applications in convergence and critical point analysis. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 3404-3413, PMLR, 06-11 Aug 2017.
[67] Luca Venturi, Afonso Bandeira, and Joan Bruna. Spurious valleys in one-hidden-layer neural network optimization landscapes. Journal of Machine Learning Research, 20(133): 2019.
[68] Roman Vershynin. High- -igeeniimeii Probability: An Introduction with Applications in Data Science. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2018
[69] Colin Wei, Jason D. Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization of neural nets vs their induced kernel. Advances in Neural Information Processing Systems, pages 9709-9721, 2019.
13

---

[70] Stephan Wojtowytsch. On the convergence of gradient descent training for two-layer relu- networks in the mean field regime. arXiv preprint arXiv:2005. 2020.
[71] Stephan Wojtowytsch and E Weinan. Can shallow neural networks beat the curse of dimen- sionality? mean field training perspective. IEEE Transactions on Artificial Intelligence, 12):1221-12, 2020.
[72] Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrizee models. arXiv preprini arXiv:2002.09277, 2020.
[73] Greg Yang and Edward J. Hu. Tensor programs iv: Feature learning in infinite-widd neural networks. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 11727-11737. PMLR, 18-24 Jul 2021.
[74] Mo Zhou, Rong Ge, and Chi Jin. A local convergence theory for mildly over-parameterized two-layer neural network. arXiv preprint arXiv:2102.02410, 2021.
[75] Zhihui Zhu, Daniel Soudry, Yonina C Eldar, and Michael Wakin. The global optimization geometry of shallow lincar neural networks. Journal of Mathematical Imaging and Vision, 22(3) 2020.
[76] Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Gradient descent optimizes over- parameterized deep relu networks. Machine Learning, 109(3)4477-92, 2020.
14

---

A
ADDITIONAL NOTATIONS
For a positive integer let [nx] denote the set f1...., n).
We use i.j (as subscripts) to index the neurons in the hidden layers, a, (as subscripts or superscripts) to index different training data points, t (as superseript) to denote the training time / time parameter in gradient flow.
We write Σ. for Σx=1
We use bold letters (e.g. 2, y) to denote vectors.
We use W and [E() interchangeably refer to the same set of parameters.
B CONSISTENCY OF THE SCALING AND GD UPDATE RULE WITH XAVIER INITIALIZATION
Consider three-layer network defined by
122
f(x)
Σ
8
(3)
a (h)(2))
(25)
m
d
Vi [m] h1(2) = >
o
2k 2%
(26)
j=1
) se(m) are initialized according to Xavier initialization, which means that we sample each 022) i.i.d. from N(0, m+a): ) each 0! i.i.d.
with weight parameters (0 J.EEm (0 } JajEtm and (0!30)
from (O, ). and each 8(0) i...d. from ((0, m If m > d, both N(O, mt and (0(0, m+11 can be approximated by (O, 1) Then, up to this approximation, by redefining & /mb (3) write
Wi / and Zjk /moll) (L)
1
f(x)
co (h((x))
(27)
m
1
vi E [m]]
h,(x:)
Wijo
Vmd 23xz)
(28)
m
j=1
and note that c22 Way and Ei are all initialized i.i.d. order 0(1). addition, if homogeneous, this is then equivalent and (5) when D
8f (x)) Then, since performing GD on 0(2) with step size 8f (x) vm 2
Moreover, there aW;
means updating 8(2) according to
JC
(29)
00
this is equivalent to updating Wiy according to
aCls]
Wi < mm
(30)
0C (x)
W md
oW,3
which justifies the factor on the right-hand-side of (9)
15

---

C RELATIONSHIP TO THE MAXIMUM-UPDATE PARAMETERIZATION AND FEATURE LEARNING
Consider the partially-tranned I-layer NN model defined in Section H in the where D m >> d. In the framework of abe- abe-parameterization introduced in (73), our model corresponds to setting
Q1
02 _.. To 11+1 =1 0,
0, A E [1+
Furthermore, as we explain in Appendix ()), which corresponds to having
the appropriate learning rate scales linearly with m (as in
c=-1
(31)
Meanwhile, the muximum-update (pP) parameterization
is characterized by setting
1
1 2
d1
2'
12
=01 = 0, ar4l
(32)
1
bi
41 E []
(33)
2
c c
(34)
Recall the symmetry of abc- abe--arametteization derived in 773, which states that one gets a different but equivalent parameterizzaion by setting
aj 0, by + b - 0,
20
(35)
Since our parameterization can be obtained from the maximum-upee uupdatee parameterization by applying the transformation above with 2. they are equivalent in the function space. In particular, for our parameterization, the r parameter defined in 73 can be computed as
r = mimmmbbbiia. aL+1 + cJ + aL+I + min {2a; 1 I=.....
1 min (0, 1 (-1)) +1 (-1) min min 0 0,2 1}
(36)
Hence, according to 73 our parameterization exhibits feature learning.
D PROOF OF LEMMA
Since assume that G is positive definite and |cal = 0, Vi E [m], we derive from 14 that
(fM(z)) vo) (J²(2))) - vn) o' (h:( (h:(ar))) Gab
2=1
a,0=1
< 8d Amin (G)
(f((())) - ya) (o' (h(((aa))))
m
1
8d Amin (G)
(f((xa)
yya)
(o' (ha(
a=1
1=1
(37)
m
S Amin (G) a=1
(f((x))) - ya)
1 (eweee (Kor)²
m
m
S 8 Amin (min(G) 2
(f((z)) - -ya)
min be(n)
1 JaN
m
a=l
1
< 22 X (G) (Kor)²
min aE(n
1 (a((a)E
m
1=1
16

---

Since I is open interval, 35 such that find subinterval Iu S I such that the distance between Iu and the boundaries of I (if is bounded on either side) is no less than &, i.e.
inf
1 - ul125
(38)
WEYO
In particular, we can choose { = (I+ I1) and (II + E. - E). Then there is
Dns(ee)er 2 2 M((o))r
> [ , (M((N)-A((m))
(39)
> ( [[((=.))A((<=)))))
and so
m
1
min <Ein]
min (n(n)
1 [a((a))E
max aE(m)
1 (k((er) -n (20
m
m
m
(40)
Thus, we have
m
}
m
L' < 22 Amiu (G) (Ko,)) C³
min aE(ni
max aE(n]
1 (m( ( 110
m
m
(41)
Meanwhile, since
M(xa) =
(f'())) Wb) of (((((mb))
(42)
b=1
there is
1
(((a))
<6
(fr - vo) of (h)((x)))
1=1
b=1
144
<8
Σ
(f((z)) 20) (hb((E))))
m
15-
<8
(fr v6) o' (h)( Gab
m
b=14=1 b=1
(43)
<c)max
(f 20)² (0 (d((((zz))))
m
<C)max
(2)2 Amiu (G)
<dmax (G) (Amin (G)) |
Therefore,
rt m <
(((x))) 11 (x)(
ha (xxa) ds
Jo i=1
(44)
< <mmax (G) (Amin (G)) f-' co ds
Since Ag E R, there is
E 1n(x,))) S (M(())) 11(2)a))
(45)
17

---

derive that, Va [n],
1
1
SS
(((aa - (xx))
i=1
(46)
S01 / L.
ds
where we C₁ Amax (G) (Amin (G)) g- for simplicity, As a consequence.
m
(0(( 12(0n)12 sc C L-' ( ds
(47)
aEin]
m
i=1
Define
m
1 (o. } }-c, Ci / [20
7
min aE(n)
ds
(48)
m
Note that at = 0, there is f7 min
mm
1 E IApGeNEL
Then, on one hand, we know
from (40) and (47 that Vt 2 0,
n min aEIn]
1
Σ
A (e.))l
max aE/n]
1
2 1t
(49)
m i=1
m
Hence,
41 implies that
L 20 Amin (G) (3 Ctr
(50)
S - 26 Amin (G) (K.,) ('' C'
On the other hand, by the definition of 1j'
f = C1 [
2 C1 [0 1.20
(51)
20,2' (20 Amin (G) (Kor) ('''''')))
202(0)-' (c') - (77)-²
where we set C2 = 2- (Amin (G)) 1 (Ko ity. Therefore, when 17² > 0, d
Amsna V2 (miin(())
n6 V2eK Amia
for simplic-
dt (1)) 8 ()))))=- (17) , 2t (c)) C2(0) d dt (2C9)
(52)
which implies that
(77)
#
2
(79)
#
2 -2( ((2c)) (2c9) -C2( (269))
(53)
3
and Vt 2 0,
2
(1)))
2
(1t))
2
(7))
v C2(6) 1(29))
(54)
3
3
3
E PROOF OF THEOREM
To apply Lemma] 1 need two additional lemmas, which we will prove in Appendix E. and
E.2 The first one guarantees that the loss value at initialization, L is upper-bounded with high probability:
18

---

a positive number that depends on Lemma 2. Ag > 0. if 52 (ce log (no-1) /max/ /lbas///ll1//)...)) then with probability least 1 6, there
_ y
(55)
The second one proves that 77 is lower-bounded with high probability, which heuristically says that there is indeed nontrivial proportion of neurons in the central part of the active region of o, for every a E [n]]
Lemma 3. Vg > O, m 2 2(K(I,Gmin Gmax))= then with probability at least 1 - 6, there is
log
170 > K(I, Gmin: Gmax)
(56)
6v27Ta (I, - I1) exp { max Mel
where KII, A1, 22) I, X and x2.
(2,)²
With these two lemmas, we deduce that AA > 0f 2 9 ((1 //1)) log (nō-1)), then with probability least 6, there is 14 0,
2
2 (K Gmins Gmax)) V2C2(c)
(57)
3
3
where K (1, Guins Gmax) is defined in Lemma Therefore, choice of é satisfies
3v 3v2C2
22 (K (T, Gmin, Gmax))
(58)
then there is Vt 2 0,
1/¹ 2 2 KL,, Gmin; Gmax) 0,
(59)
which case
(50) gives
C < - Amia (G) (I, Guinn Gmax))
(60)
which will allow us finally conclude that
L < LD exp
{ -22 Amin (G) (G) (Ko K(1, Gmin Gmax)
(61)
Note that (60) establishes a PL condition. Several other convergence analyses of NNs have also relied variants of the PL condition 25 44 74.
E.1 PROOF LEMMA
is ie jE(D] are both sampled i.i.d. and {ci se(m] has mean zero, we know that Va E [n)] (xa. = 1 E Ga (N.)).) is the sample mean of i.i.d. random variables with zero-mean. Moreover, since {wo jE(D is sampled from N/O0, 1), we
Proof. Since at initialization, {cr rermi and {w-
m
know that Vi E [m]], the random variable c: (x)()) is sub-Gaussian (68) with sub-Gaussian norm
|ccar (b)())) 1l, Sāllo (a:(
SèLo (Gao) Msc
(62)
(G MsG
where MsG > is some some absolute constant. Thus, by Hoeffding's inequality 68 Va E [m], Vr > 0.
m
P (If((x)) 2 u) =P
Go (h)(
2
m
Ku2m
<2exp
(63)
|c:,o (h;
Ku²m m
<2exp
2L ).u.) (M
19

---

where K is some absolute constant. Hence, by union bound.
P
r ( 12 2 |
(1
2 11
a=1
(64)
<2nexp {
K 1 11ll111 1m
(I,) ) (Msc)
Thus, VS > 0, if
2n
m 2 (L))) Gmax (Msc\= -2 ~logg
(65)
8
then with probability least 1 8, there is
(((()) (2aa - -Yal
2 a=
1 7
f (x) ) 1² + 1/11² |
1
(66)
2
E.2 PROOF OF LEMMA
Since each W are sampled i..d. from (0, 1), we know that Va E [m]. independently for each i E [m], (x,)) follows a Gaussian distribution with mean 0 and variance Goa- Therefore,
di Binomial (m, n (Io; Gaa))
(67)
Hence, by Hoeffding's 's inequality, Va E [n], Vr > 0,
m 1 (e. 2 1 T(Io: Goo) tr I sexp ( (2mr²)
(68)
Va E [n]], choosing r b (To: Gaa), we then get
m
1 (g())e 2 - 1
P
(ew)ee S n (Tu Gaa)
=F
2 (Iu: Gao)
m
2
m
i=1
< exp
- ( (lo; Gan))
(69)
{ (To Guoll )²}
1 m
S exp
min bE|n)
2
and by union bound
min
{ (Io: Gna)]
=P
min eEin
min { (To: Gnadl bE(m)
2
bE(m)
m
2
im1
P
min {o (To; Gaa) bE(n)
m
2
0=1
1=1
(70)
<
P
Ag(G.) < < (To: Gaa)
6=1
1=1
min { (la; G(a)) 2) bE(n]
Snexp
-
20

---

Since Ab [n], there is Guin I Gbb Gmax.
pla-s e Gbb du
π (Lus Gbb)
V2TGub lit5
u²
11
inf
exp
Gb
11 1+tESuST-
Gbb
(71)
1
max 4+
(1r -11-1---
exp exp
12
(Gmin)²
max
( |
(1)1))exp exp.
3v 2n 2Taxx
(Gmin)
Letting 21, = 6v/2 (L, 11) exp
{
max((411,11111 (A1)²
> 0, we can then write
1
IP (ip S (T, Gmins Gmax)) SP
1j0 S min f (Tos Gab)) bE(n)
2
(72)
Sn Exexp
{
{-2m 2m (K(I, Gmin Gmax))
Thus, V8 0, if m
( 2) then with probability at least 1 - 6, it holds that
log(no
70 > K[I, Gmin; Gmax)
> 0.
F PROOF OF THEOREM 2
By Condition we know that > 0, Df mmi (mmin then with probability at least 1 - 5, there is ||G - Gl|2 < 2 Amin (G), and hence Amin (G) 2 Amii (G), Gmin 2 G Xmax(()) < GG) + Amin (G) I 22max (G). and Giax I 2G We then perform the following analysis conditioned the event that ||G - Gll2 I 1 Ammin(( (G).
Since the sampling of {cr (e)ieem and (W:) jE(m) is independent from the realization of G, log(4n8-1)
log(And- then with we know from Lemma 3 that if m 2 2(K((, Amin (G),2Amma (O)))) 2[K(I.Gain Gmax))
probability at least 1 oo, there is
7" K(I, Gmin; Gmax) 2 K(1, G 2Gmax)
(73)
From Lemma
2 we also know that m 2 Ω (6² log (no-1) Amax AG)x 12)
Ω (22 log (nō-1) Amax Amax(G)/lll (G)))lyl??, then with probability at least 1 - 18. there is L < |g/| There- fore, total, we know that with probability least 1 - 5, the following conditions all hold:
| - Amin(G))
(74)
7 2 (T,, Gmin, 26max)
(75)
" < 1
(76)
in which case, by applying Lemma 1 with G = (1),, we get
(n²) $ 2 (170 Ki (dmin (G) Amax (G)
((C))+
(77)
where (((A,,)2))
: 'eK (I, 11) > 0. Thus, by the definition of (K(())) )) know that
Ki (Amin (G), Amax (G)) < 4K 4K1(Amm (Amin (G) Amax (G)),
(78)
and so Vt 2 0,
(1²))
(70) 4Kg (Amin (G) Amax (G) ((C))
(79)
2 (K(I,
G Guin, 2G 2mmax)) (KKg(Ani (G) Amax (G))
2
21

---

Therefore, if our choice of é satisfies
8K1 (Amin (G) Amax (G) )))ll
(80)
2 2Gmax)
then there is Vt 2 0,
1 G Griny 2G 2Gmax)
o/ 2 122-3K(
(81)
Hence,
500 implies that Vt 2 0,
L v - 26 Xmin (G) K(I, 1
Guain 2Gmax)
2
(82)
S - 2 2 Amin (G) ( K(1, 2 Gminn 20 max)
and therefore
L' S' S - exp { - ( K(1, 1 2
2miin 20 20mxx))
(83)
G PROOF OF THEOREM
In view of Theorem 2 it is sufficient to verify that Condition holds for Dunin (6, u)
Ω (n"u== log((n ")). which is given by the following lemma:
Lemma 4. Ag 2 0, D 2 (n²u* Plo((88-))) then with probability at least 1 - 5,
[1 (22) Su
(84)
G.1 PROOF OF LEMMA 4
Let Z be a random vector on Rd with law given by T2. and then we can write G()) E |o (x)] for a, b E [n]. By assumption, Z is sub-gaussian with sub-gaussian norm
ab
2 1l == sup ||x Z 11iz < 00
(85)
xESé-1
Thus, Va E [n], we have
| S Lo | 11 S La 111
(86)
Hence, by Lemma 2.7.7 in 68. know that Va, b E [n], (II is sub-exponential
random variable with sub-exponential norm
(a(= S [|(+
(L,)
(87)
Then, by Bernstein's 's inequality (Theorem 2.8.1 in (88) since each zj is sampled i.d. from u we have that Va, E [m] and Vu > 0,
P (Ga
G 2u) =P
1 D
(xT
Eo
2)
2u
j=l
2
uD
<2exp
K min
|oo (2) 2
[o ( Z
u²D
uD
<2exp
K min
(T.g |
(I, ) 1
(88)
where K > is some absolute constant. other words, for any 8' > if
(Teg) [1 1 1og(2(6")-1)
(T.g) 2 7 2 log ( ((
D > max
(89)
u² K
uk
22

---

(=j} is sampled i.idd from (0, La), and (Wi)), vE(m)_E[D) is initialized by sampling i.i.d. from (0, 1). 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 test set of size 500 generated by sampling G G.
then we have then get, if
2 u with probability at least 6. we choose u
and 8 5
n2.
n² (L())
11 log(2n28-1)
nlLo)
log(2n28
2 max
(90)
(22²)²
Ku'
then
P - G - 2 (u')²)
1 Σ
G(l
4
2
ab
ab
0,b=1
(91)
Sn²²
112
<ó
Hence, with probability at least 6, we have
||G(]] - 1 = - & (ul)²
(92)
H GENERALIZATION TO DEEPER MODELS
By setting & to be the activations of the second-to-last hidden-layer multi-layer NN, we can obtain generalizations of the P-3I. NN to deeper architectures. For example, in the feed-forward case, we can obtain the following partially-trained L-layer NN:
((x)
hh (L- (4-))(()))
E [m]
h (L- (x)
Wijo (
VD
D
W W Vi [D]
h2
(x)
(1+1)
(x)
W σ (x))
vD
==1
Vj [D]
(x)
1
Vd
where (()), .... W(L-3) E DxD and Zi,..., E Rd are sampled randomly and fixed. This model can be written the form matrix is recursively defined and also appears in the NTK analysis [18] In particular, the results
and (5) with dy (x) = (1) (æ)). The corresponding Gram
in (18] imply that a is analytic and not polynomial, then Condition holds, and hence similar
global convergence results can be obtained as corollaries of Theorem 2
I FURTHER DETAILS OF THE NUMERICAL EXPERIMENTS
In our models, {c] is sampled i.i.d. from the Rademacher distribution p. 101 + 1
step size 8 = 1. When the test error is computed, i..d. from the same distribution as the training set.
The experiments
with NVIDIA GPUs (1080ti and Titan RTX).
23

---

109
10"3
train loss
10
300
2000
10000
GD step
Figure 5: Training loss v.s. number of GD steps for different n in Experiment with m 4096.
I.1 EXPERIMENT
We choose to be tanh. For each choice of n, we run the experiment with 5 different random seeds, and Figure D plots the evolution of the training loss during GD averaged over the runs with m 8192.
Figure is the same as Figure except for having m 4096. We see that the two two plots agree well.
1.2 EXPERIMENT 2
We choose a to be ReLU. For each choice of n and each the two models, we experiment with different random seeds, and Figure plots the test error the 50000 GD step averaged over the 5 runs + its standard deviation.
In Figure 6 plot the evolution of the training loss and test error during GD for the two different models, with m = 2048 or 8192 and different choices of n, averaged over 5 runs with different random seeds. We in particular that the difference between the two choices of is negligible, suggesting that it is unlikely to obtain performance improvements with further over-parameterizztion..
I.3 EXPERIMENT
We choose to be ReLU and input dimension d = 50. We use training set of size 600 for the results reported in Figure The data set inspired by [69]: We sample both the training and the test set i.i.d. from the distribution (z, D on Id+ under which the joint distribution of (æ1, 22, y)
P(#1 = 1, = 0,y 4
(93)
P(ri -1,42 0, 1)
(94)
P(zı 0, 1, 1)
(95)
P( = 0, = - -1--=
(96)
(97)
and #3 ...,.... Fa each follow the uniform distribution in [-1,1], independently from each other as well T1, 22 and y.
Figures 7 and g are the same as Figure exxcpp for having n = 400 and 800, respectively. We see that increases, test error improves for all three models, while our P-3L NN model remains the one achieving the lowest test error.
24

---

0=50; Shallow acc
n=50; Shallow- Suallow-Quad
0.4
0.4
0.2
0.2
0.0
0.0
-0.2
-0.2
2000
3000
100, Shallow acc
100: Shallow Quad
1.2
1.0
1.0
2000
3000
T500
10000
2000
5000)
7500
10000
n= 150; Shallow acc
n= 150; Shallow Quad
2008
1000
254
-02
-0.2
1000
3000
T000
2000
3000
2000 Shallow
n=200; Shallow Quad
1.5
1.2
20-0
1.0
1.0
0.8
0.5
0.0
0.0
0.4
2000
3000
3000
5000
n- 250; Shallow acc
n
250;
Shallow + Quad
2010
2040
1000
-0.2
-0.2
2000
3000
7000
10000
2000
5000
7900
10020
Figure 6: Test error GD steps in Experiment 2 for the two models and difference choices of n.
25

---

P-3L NN
P-3L NN (NTK)
P-3L NN (MF)
10
10
10
10-3
10
10-1
10
2243
27-2340
mud455
4000
20000
4200
GD step
GD step
GD step
teker
beier
abe
h1(x
-6-
b;(x²)
1;
Figure Same as Figure 4 except for setting 400.
P-3L NN
P-3L NN (NTK)
P-3L NN (MF)
199
107
19a4
10"1
10"1
10-7-
10
10-7-
10
10
10
10
10
10
1392
10"4
10a4
10
20000
1000
2000
4000)
20000
42000
GD step
step
GD step
0-
hi(x
0-
-0-
22
h;(x²)
1x
Figure 8: Same as Figure except for setting n = 800.
26


Reviews:
[{"review_id": "i2ZlzShbI_N", "reviewer": "Reviewer_XJCF", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "\nAlthough the n <= d setting is somewhat limited, I quite liked this paper, as all previous mean field neural network analyses were quite complicated.  This one also appears novel in that there is the connection to kernels which was unexpected.  I should note I am not an expert on mean field analysis of neural networks, but to my knowledge the paper's contributions are novel and commendable.\n\nThere were a number of points of clarification I'm hoping the authors could comment on.\n\n(1) The choice of the scaling in the network is non-standard and deserves more discussion.  The authors claim at top of pg. 4 that the scaling agrees with the scaling of Yang & Hu '20---I assume the authors are referring to the \\mu P scaling in Table 1. But it seems to me that this is quite a different scaling than the one there, and also the standard mean field scalings.  The \\mu P scaling requires the scaling of sqrt(D), not 1/sqrt(D), for the inner layer weights, since a_1 = -1/2.  Moreover, the learning rate scaling in \\mu P is order 1, while the learning rate scaling in this paper is order m.  I understand the need for learning rate scaling of order m; this is standard in mean field and is explained well in Appendix B.  \n\nWhat happens when the scaling is 1/D vs. 1/sqrt(D)?  Or even unit scaling (as is standard in mean field, see MFP column in Yang & Hu '20)?  My understanding is that it only changes the scaling of lambda_min(G), and thus doesn't really affect the rates (at least in the instance of leaky relu activations).  \n\n(2) Regarding Experiment 1: Are different d chosen for the different n?  The results only hold when d <= n, so it's odd to consider d=20 for each of them but this seems to be what the authors have done.  Regarding the negative result of (65; Wojtowytsch & E '20), please give more details on what precisely the negative result is, and how your experiments relate to it.  My understanding is those results show that the risk can only decrease at rate t^{-c/d}, so that as d increases the rate should get worse.  But here you are increasing n.  It was also unclear to me why such low dimensions and sample sizes were used here.\n\n\nTypos and minor comments\n\n\"feature training\" on page 2, bottom of Sec 1.1\n\nreferences for claim at bottom of page 2 about \"weights would lease to collapse of the diversity\"?\n\nThe usage of summation terms everywhere rather than vectorized forms was distracting and atypical.  If accepted, for the camera ready I would recommend re-writing everything in the form of e.g. f(x) = (1/m) c^T \\sigma(Wx)\n\nPresumably \\hat c = \\Theta(1) in assumption 1; if c depends on m things could go badly?\n\ni\\in [m] on top of pg.6, not i\\in [n]\n\n(30) should be \\partial \\mathcal L, not \\partial f, presumably\n\n\"Appendix ??\" on pg 17\n\nI should note that  I did not check the details of all the proofs in the appendix.", "summary_of_the_paper": "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. ", "summary_of_the_review": "The paper makes a novel contribution for convergence of NNs in the mean field regime with a clean and simple proof.  ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "XMT7SzWKq3T", "reviewer": "Reviewer_jNHj", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "[Contributions]\n\nRecently, the mean field setting of neural networks has become an important topic in the context of global convergence analysis of neural networks because of the presence of feature learning whereas the kernel (lazy) regime basically describes the local behavior of the dynamics of training. However, an optimization theory in the mean field regime is significantly more challenging and usually requires involved mathematical tools. \n\nIn this sense, this study makes a certain contribution in finding a simple condition. The proof can be considered as a sort of extension of the NTK-theory. Indeed, this theory basically builds upon showing the positivity of (finite-width) NTK as well as NTK-theory but does not require fixation of NTK, unlike NTK-theory. Specifically, the positivity of NTK can be reduced to the positivity of the Gram-matrix of input (random) features and $d \\geq n$ by decoupling the parameter-dependent part from the NTK.\n\n[Improvements]\n\n- It is misleading to say that this theory covers multi-layer neural networks because the trainable layer is limited to the second-to-last layer. Therefore, the model in this paper is essentially two-layer neural networks with random feature inputs.\n- The explanation of feature learning in the last paragraph of Section 2 is insufficient. In particular, the large stepsize in proportion to the width $m$ is also needed to exhibit feature learning as well as network scaling $1/m$. (See feature learning paper [67] for the detail. Although this submission considers the continuous-time dynamics, this point should be mentioned.)\n- Some important references are missing. For instance, (modified) PL-condition based analysis of overparameterized neural networks is relevant to this submission in the sense that the theory essentially relies on the positivity of NTK.\n\n    - Spencer Frei and Quanquan Gu. Proxy Convexity: A Unified Framework for the Analysis of Neural Networks Trained by Gradient Descent. NeurIPS, 2021.\n\n    - Chaoyue Liu, Libin Zhu, and Mikhail Belkin. Loss landscapes and optimization in over-parameterized non-linear systems and neural networks. 2020.\n\n    - Mo Zhou, Rong Ge, and Chi Jin. A Local Convergence Theory for Mildly Over-Parameterized Two-Layer Neural Network. COLT, 2021.\n\n\n[Additional references]\n\nMore recently, global convergence rate analyses in the mean field regime under KL-regularization were established by several papers.\n\n- Kaitong Hu, Zhenjie Ren, David Siska, and Lukasz Szpruch. Mean-field Langevin dynamics and energy landscape of neural networks. 2019.\n\n- Jean-François Jabir, David Šiška, and Łukasz Szpruch. Mean-Field Neural ODEs via Relaxed Optimal Control. 2019.\n\n- Atsushi Nitanda, Denny Wu, and Taiji Suzuki. Particledual averaging: Optimization of mean field neural networks with global convergence rate analysis. NeurIPS, 2021.\n\n[Minor comments]\n\n- Typo (line2, page 6): $i \\in [n]$ → $i \\in [m]$.\n- There are typos in the definition of pseudo-L-layer NN in Section 3.2.2:\nThe description of $\\phi_j$ seems redundant. \nFor $h_i^{L-1}$, the index should be $i \\in [m]$.\n", "summary_of_the_paper": "This paper establishes the global convergence analysis, with a linear convergence rate, of gradient flow for the neural networks in the mean field regime which possesses a feature learning aspect. The contribution of this paper is to show that the positivity of the Gram-matrix of input (random) features is sufficient for guaranteeing global convergence.", "summary_of_the_review": "I think this paper makes a certain contribution in finding out a simple condition for the global convergence analysis in the mean-field regime. However, the quality of the manuscript could be improved. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "VctadgOpby", "reviewer": "Reviewer_zizE", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength\n-\tUnderstanding the optimization of neural network beyond NTK (lazy training) is an important direction.\n-\tUnder certain condition on the Gram matrix, it can be shown that GD converges to 0 training loss. Empirically, the neural network indeed shows the ability of feature learning.\n\nWeakness:\n-\tMajor concern: I was wondering if authors could clarify the convergence rate dependency on \\lambda_min (G) and \\lambda_max (G) in Theorem 1, Theorem 2 and Theorem 3, i.e., the relation between r,\\hat{c}_min,C and \\lambda_min (G), \\lambda_max(G). There seems no discussion about it. It seems to me that \\hat{c}_min exponentially depends on 1/\\lambda_min (G) (based on Lemma 3, especially the definition of K(I,\\lambda_1,\\lambda_2)), which implies the width m exponentially depends on 1/\\lambda_min (G). If this is true, I feel the results in this paper are less interesting, since it would require exponentially number of neurons for the global convergence.\n\n-\tFor multi-layer neural network, only the second-to-last layer is trained while all other layers are random sampled and fixed. This is different from the setting in practice. It is understandable that there might be technical challenges when analyzing the case of training all layers, so I would not view this as a major limitation.\n\nMinor comments:\n-\tAt the end if page 3 (fixed embedding): \\phi_j(x) = x_j instead of \\bm{x}_j\n", "summary_of_the_paper": "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.", "summary_of_the_review": "Based on my above comments, I currently would not vote for accept because of the major concern above.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "IpN7gUMxwyM", "reviewer": "Reviewer_48yw", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "As far as I can tell, the setting of this paper is still limited to the lazy training regime. In other words, the essential model of fitting the $n$ data points is following linearized model:\n$$\nf(x;W+\\Delta) \\approx f(x;W) + \\langle \\nabla_W f(x;W), \\Delta \\rangle.\n$$\n\nThis can be seen as follows.\n\n- Notice that for the output of ``feature function'' $h_i(x)=\\frac{1}{\\sqrt{D}}\\sum_{j=1}^D W_{i,j}\\phi_j(x)$ changing an $O(1)$ value, $\\{W_{i,j}\\}$ only needs to change $O(1/\\sqrt{D})$. This change shrinks as increasing $D$. Also it is indeed required that $D\\geq poly(n)$ in the proof. In contrast, in Theorem 1, they only require m (the width of the output layer which is in LNN scaling) to be larger than $\\log(n)$.\n\n- On the other hand, the movement of W can roughly be estimated as follows,\n  $$\n  ||\\Delta|| \\leq \\int_0^\\infty ||\\nabla_W L(W_t)|| d t \\leq \\int_0^\\infty \\sqrt{\\lambda_\\max(G) L(W_t)} dt\\leq \\sqrt{\\lambda_\\max(G)}\\int_0^\\infty e^{-r\\lambda_\\min(G) t/2}dt \\leq c\\sqrt{\\frac{\\lambda_\\max(G)}{\\lambda_\\min(G)}} = O(poly(n)).\n  $$\n  Hence, on average, for each $i,j$, $W_{i,j}$ roughtly only moves $poly(n)/\\sqrt{D}\\ll 1$ when $D$ is sufficiently large.\n\nPlease feel free to correct me if you think I misunderstood the results. ", "summary_of_the_paper": "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. \n\n\n\n\n", "summary_of_the_review": "The studied setting is not too much different from previous work of non-convex optimization in the lazy regime. Therefore, I cannot recommend accepting this paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "1: The main claims of the paper are incorrect or not at all supported by theory or empirical results.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper studies optimization of over-parametrized neural networks in the mean-field scaling. Specifically, when the input dimension in larger than the number of training samples, the paper shows that the training loss converges to 0 at a linear rate under gradient flow. It's possible to extend the result by random feature layers to handle the case when input dimension is low. Empirically the dynamics in this paper seems to achieve better generalization performance than the NTK counterpart, but no theoretical result is known. Overall this is a solid contribution to the hard problem of analyzing the training dynamics of mean-field regime. There was some debate between reviewers on what is the definition of "feature learning" and I recommend the authors to give an explicit definition of what they mean (and potentially use a different term).
Decision: Accept (Poster)