Memorization and Optimization in Deep Neural Networks with Minimum Over-parameterization
Anonymous Author(s)
AffiliationAddress email
Abstract
The Neural Tangent Kernel (NTK) has emerged as a powerful tool to providememorization, optimization and generalization guarantees in deep neural networks.
A line of work has studied the NTK spectrum for two-layer and deep networkswith at least a layer with Ω( N ) neurons, N being the number of training samples.
Furthermore, there is increasing evidence suggesting that deep networks with sublinear layer widths are powerful memorizers and optimizers, as long as the numberof parameters exceeds the number of samples.
Thus, a natural open question iswhether the NTK is well conditioned in such a challenging sub-linear setup.
In thispaper, we answer this question in the affirmative.
Our key technical contributionis a lower bound on the smallest NTK eigenvalue for deep networks with theminimum possible over-parameterization : up to logarithmic factors, the number ofparameters is Ω( N ) and, hence, the number of neurons is as little as Ω( √ N ) .
Toshowcase the applicability of our NTK bounds, we provide two results concerningmemorization capacity and optimization guarantees for gradient descent training.
Introduction
Training a neural network is a non-convex problem that exhibits disconnected local minima [8, 57, 70].
Yet, in practice gradient descent (GD) and its variants routinely find solutions with zero training loss[72].
A framework to understand this phenomenon comes from the Neural Tangent Kernel (NTK).
This quantity was introduced in [33] where it was proved that, during GD training, the networkfollows the kernel gradient of the functional cost with respect to the NTK.
Furthermore, as the layerwidths go large, the NTK converges to a deterministic limit which stays constant during training.
Hence, in this infinite-width limit, it suffices that the smallest eigenvalue of the NTK is boundedaway from 0 for gradient descent to reach zero loss.
Going to finite widths, a recipe to prove GDconvergence can be summarized as follows: show that (i) the NTK is well conditioned at initialization,and (ii) the NTK has not changed significantly by the time GD has reached zero loss (see e.g.[49, 17, 11]).
A number of papers have exploited this recipe for networks with progressively smallerover-parameterization: two-layer networks [22, 50, 63, 69, 62], deep networks with polynomiallywide layers [5, 21, 73, 74], and deep networks with a single wide layer [47, 43].
Besides optimization,showing that the NTK is well conditioned directly implies a result on memorization capacity [42, 48],and the smallest eigenvalue of the NTK has also been related to generalization [6].
In [48], it is shown that, given N training samples, a single layer with Ω( N ) neurons suffices for the
NTK to be well conditioned in networks of arbitrary depth.
However, there is increasing evidencethat, in the challenging setup in which the layer widths are sub-linear in N , neural networks stillmemorize the training data [14, 71, 66], reach zero loss under GD training in the two-layer setting[61, 50, 42], and in the deep case, GD explores a nicely behaved region of the loss landscape [44].
This is in agreement with simple back-of-the-envelope calculations on CIFAR-10 and ImageNet:
CIFAR-10 has N = 50000 images and roughly 10 6 parameters suffice to fit random labels [72];furthermore, in order to fit random labels to a subset of 1 . 2 · 10 6 ImageNet data points, 2 . 4 · 10 7 parameters are enough [72].
These numbers suggest that having a number of parameters of the sameorder as the dataset size is much closer to practice than having a number of neurons of that order.
Wealso note that, by counting degrees of freedom or bounding the VC dimension [10], Ω( N ) parameters(and, therefore, Ω( √ N ) neurons) are in general necessary to fit N data points.
This naturally bringsforward the following open question:
Is the NTK well conditioned for deep networks with the minimum possible over-parameterization(i.e.
, containing Ω( N ) parameters corresponding to Ω( √ N ) neurons)?
Main contributions.
In this paper, we settle this open question for a large class of deep networks.
We consider (i) a smooth activation function, (ii) i.i.d.
data satisfying Lipschitz concentration (e.g.,data with a Gaussian distribution, uniform on the sphere/hypercube, or obtained via a Generative
Adversarial Network), (iii) a standard initialization of the weights (e.g., He’s or LeCun’s initialization),and (iv) a loose pyramidal topology in which the layer widths can increase by at most a multiplicativeconstant as the network gets deep.
Then, in Theorem 3.1 we show that the NTK is well conditionedunder the minimum possible over-parameterization requirement: the number of parameters betweenthe last two layers has to be ˜Ω( N ) or, equivalently, the number of neurons has to be ˜Ω( √ N ) , where˜Ω includes extra logarithmic factors.
We achieve this goal by giving a lower bound on the smallesteigenvalue of the NTK.
This lower bound is tight when all the layer widths are of the same order.
Our NTK bounds open the way towards understanding the behavior of deep networks with minimumover-parameterization.
In particular, an immediate consequence of the fact that the NTK is wellconditioned is a result on the memorization capacity (Corollary 4.1).
Furthermore, by suitablychoosing the initialization, we provide convergence guarantees for gradient descent training (Theorem4.2).
Finally, we highlight that, in order to obtain our bounds on the smallest NTK eigenvalue, wegive a number of tight estimates on the ℓ 2 norms of feature vectors and of their centered counterparts,which may be of independent interest.
Proof ideas.
To prove Theorem 3.1, we restrict to the kernel K L − 2 =
J L − 2 J ⊤ L − 2 , where J L − 2 isthe Jacobian of the output w.r.t.
the parameters between the last two hidden layers.
This sufficesas K L − 2 is a lower bound on the NTK in the positive semi-definite (PSD) sense.
We note that thei -th row ( i ∈ { 1 , . . . , N } ) of J L − 2 is given by the Kronecker product ⊗ between the feature vectorat layer L − 2 and the backpropagation term from the same layer.
One key technical hurdle is tocenter J L − 2 , so that its rows have the form u ⊗ v − E [ u ⊗ v ] , where E [ u ] =
E [ v ] = 0 , and allexpectations are taken with respect to the (random) training data.
To do so, we perform three steps ofcentering: (i) we center the feature vectors (corresponding to u ), (ii) we center the backpropagationterms (corresponding to v ), and (iii) we center again the whole row (corresponding to u ⊗ v ).
Thesecentering steps are approximate in the sense that the centered matrix is not necessarily a lower bound(in the PSD sense) on the original one.
However, we are able to control the operator norm of thedifference, and show that it scales slower than the smallest eigenvalue of the centered kernel.
At thispoint, we leverage the structure of the rows of the centered Jacobian to bound their sub-exponentialnorm via a version of the Hanson-Wright inequality for (weakly) correlated random vectors [1].
Finally, after providing also a tight estimate on the ℓ 2 norms of such rows, we can exploit a resultfrom [2] to lower bound the smallest singular value of a matrix whose rows are independent randomvectors with well controlled sub-exponential and ℓ 2 norms.
Existing work bounds the smallest NTK eigenvalue for networks with two layers [61, 42], or deepnetworks with a layer containing Ω( N ) neurons [48].
In particular, [61] also exploits the results [1, 2].
However, the centering of the Jacobian is achieved via a combination of whitening and droppingrows, which appears to be difficult to generalize to deep networks.
In contrast, the 3-step centeringwe described above applies to networks of arbitrary depth L .
A different approach is put forwardin [42], and it is based on a decomposition of the kernel via spherical harmonics.
This techniqueallows to obtain the exact limit of the smallest NTK eigenvalue, it has been used to analyze randomfeature models [26, 41, 25, 40] and to obtain generalization bounds for two-layer networks (see again[42]).
However, understanding how to carry out such a decomposition in the multi-layer setup is anopen problem.
Finally, [48] considers the deep case, and it relates the smallest eigenvalue of the NTKto the smallest singular value of a feature matrix.
As feature matrices are full rank only when thenumber of neurons is Ω( N ) , this approach is inherently limited to networks with a linear-width layer.
The rest of the paper is organized as follows: Section 2 discusses the problem setup and our modelassumptions; Section 3 presents our main result on the smallest NTK eigenvalue and gives a roadmap of the argument; Section 4 provides two applications of our NTK bounds: memorization capacity andgradient descent training; Section 5 discusses additional related work, and Section 6 provides someconcluding remarks.
The details of the proofs are deferred to the appendices.
Preliminaries
Neural network setup.
We consider an L -layer neural network with feature maps f l : R d → R n ldefined for every x ∈ R d as
[Equation]
Here, W l ∈ R n l − 1 × n l is the weight matrix at layer l , ϕ is the activation function and, givenan integer n , we use the shorthand [ n ] = { 1 , . . . , n } .
We assume that the network has a single output, i.e. n L = 1 and W L ∈ R n L − 1 , and for consistency we have n 0 = d .
Letg l : R d → R n l be the pre-activation feature map so that f l ( x ) =
ϕ ( g l ( x )) for l ∈ [ L ] .
We define g 0 ( x ) = f 0 ( x ) = x .
Let X =
[ x 1 , . . .
, x N ] ⊤ ∈
R N ×
d be the data matrix containing Nsamples in R d , θ = [vec( W 1 ) , . . .
, vec( W L )] be the vector of the parameters of the network, and
F L ( θ ) =
[ f L ( x 1 ) , . . .
, f L ( x N )]
⊤ be the network output.
We denote by J the Jacobian of F L withrespect to all the parameters of the network:
[Equation]
Our key object of interest is the empirical Neural Tangent Kernel (NTK)
Gram matrix , denoted by
K ∈ R N × N and defined as:
[Equation]
In [33], it is shown that, as n l → ∞ for all l ∈
[ L − 1] , K converges to a deterministic limit , whichstays constant during gradient descent training.
The focus of this paper is on the finite-width behaviorof the empirical NTK (3).
Quantitative bounds for the NTK convergence rate can be obtained from[7, 15].
However, these bounds lead to a significant over-parameterization requirement, see thediscussion at the end of Section 3 in [48].
Here, our main result consists in showing that the NTK iswell conditioned for a class of neural networks with the minimum possible over-parameterization ,i.e., Ω( N ) parameter or, equivalently, Ω( √ N ) neurons.
Weight and data distribution.
We consider the following initialization of the weight matrices:(
W l ) i,j ∼ i .
i .
d .
N (0 , β 2 l /n l − 1 ) for l ∈
[ L − 1] , i ∈ [ n l − 1 ] , j ∈ [ n l ] , where β l is a numerical constantindependent of the layer widths.
This covers the popular cases of He’s and LeCun’s initialization[27, 30, 36].
For the last layer, we assume that ( W L ) i ∼ i .
i .
d .
N (0 , β 2 L ) for i ∈ [ n L − 1 ] .
Throughoutthe paper, we let ( x 1 , . . .
, x N ) be N i.i.d.
samples from the data distribution P X , which satisfies theconditions below.
Assumption 2.1 (Data scaling) .
The data distribution P X satisfies the following properties:
[Equation]
x ∥ 2 dP X ( x ) =
Θ( √ d ) .
(cid:82)
∥ x ∥ 22 dP X ( x ) =
Θ( d ) .
(cid:82) (cid:13)(cid:13) x − (cid:82) x ′ dP X ( x ′ )
(cid:13)(cid:13) 2 2 dP X ( x ) =
Ω ( d ) .
Assumption 2.2 (Lipschitz concentration) .
The data distribution P X satisfies the Lipschitz concentration property .
Namely, there exists an absolute constant c > 0 such that, for every Lipschitz continuous function φ : R d → R , we have E | φ ( X ) | < + ∞ , and for all t > 0 ,
[Equation]
Assumption 2.1 is simply a scaling of the training data points and their centered counterparts.
Assumption 2.2 covers a number of important cases, e.g., standard Gaussian distribution [65], uniform distributions on the sphere and on the unit (binary or continuous) hypercube [65], dataproduced via a Generative Adversarial Network (GAN) 1 [59], and more generally any distributionsatisfying the log-Sobolev inequality with a dimension-independent constant.
We also remark that
Assumption 2.2 is rather common in the related literature [48], or it is even replaced by a strongerrequirement (e.g., Gaussian distribution or uniform on the sphere)
Assumption 2.3 (Activation function) .
The activation function ϕ satisfies the following properties:
ϕ is a non-linear (and therefore also non-constant) M -Lipschitz function;
its derivative ϕ ′ is a M ′ -Lipschitz function;
These requirements are satisfied by common activations, e.g. smoothed ReLU, sigmoid, or tanh .
Assumption 2.4 (Network topology) .
The network satisfies a loose pyramidal topology condition,i.e. n l = O ( n l − 1 ) , for all l ∈
[ L − 1] .
A strict pyramidal topology (namely, non-increasing layer widths) has been considered in priorwork concerning the loss landscape [45, 46] and gradient descent training [47].
Our Assumption 2.4requires a loose pyramidal topology in the sense that, as we go deep, the layer widths are allowed toincrease by a constant multiplicative factor.
We also note that the widths of neural networks used inpractice are often large in the first layers and then start decreasing [28, 60].
Assumption 2.5 (Over-parameterization) .
We have that
[Equation]
Furthermore, there exists γ > 0 such that
[Equation]
Condition (4) requires the number of parameters between the last two hidden layers to be linear inthe number of samples, up to logarithmic factors.
This represents our key over-parameterizationcondition .
When all the widths have the same scaling, (4) reduces to N = ˜
o
( d 2 ) .
This is satisfiedby several “standard” datasets, such as MNIST ( N = 60000 , d = 784 ), CIFAR-10 ( N = 5 · 10 4 ,d = 3 · 32 2 ), and ImageNet ( N = 1 . 4 · 10 7 , d = 2 · 10 5 ).
We also note that, if N ≫ d 2 , the NTKis low-rank and λ min ( K ) = 0 . Furthermore, Corollary 4.1 and Theorem 4.2 cannot generally holdwhen N ≫ d 2 , as there are more data points to fit than parameters to help with the fitting. A milderrequirement is possible for classification tasks, see the discussion at the end of Section 5. The secondcondition (5) is rather mild, as γ can be taken to be arbitrarily small, and it avoids an exponentialbottleneck in the last hidden layer.
Notation.
The feature matrix at layer l is F l =
[ f l ( x 1 ) , . . .
, f l ( x N )]
⊤ ∈ R N × n l , and Σ l ( x ) =diag([ ϕ ′ ( g l,j ( x ))]
n l j =1 ) for l ∈ [ L − 1] , where g l,j ( x ) is the pre-activation neuron.
Given two matrices
F, B ∈ R m × n , we denote by F ◦ B their Hadamard product, and by F ∗ B =
[( F 1: ⊗ B 1: ) , . .
, ( F m : ⊗
B m : )]
T ∈ R m × n 2 their row-wise Kronecker product (also known as Khatri-Rao product).
Given arandom vector v , let ∥ v ∥ ψ 2 and ∥ v ∥ ψ 1 denote its sub-Gaussian and sub-exponential norm, respectively(see also Appendix A for a detailed definition).
Given a matrix A , let ∥
A ∥ op be its operator norm,∥ A ∥ F its Frobenius norm, λ min ( A ) its smallest eigenvalue, and σ min ( A ) its smallest singular value.
We denote by ∥
φ ∥ Lip the Lipschitz constant of the function φ .
All the complexity notations Ω( · ) ,
O ( · ) , o ( · ) and Θ( · ) are understood for sufficiently large N, d, n 1 , n 2 , . . .
, n L − 1 .
Tildes on suchsymbols are meant to neglect logarithmic factors.
NTK Bounds with Minimum Over-parameterization
Our main technical contribution on the smallest eigenvalue of the empirical NTK (3) is stated below.
Theorem 3.1 (Smallest NTK eigenvalue under minimum over-parameterization) .
Consider an Llayer neural network (1) , where the activation function satisfies Assumption 2.3 and the layer widthssatisfy Assumptions 2.4 and 2.5.
Let { x i } Ni =1 ∼ i .
i .
d .
P X , where P X satisfies the Assumptions 2.1-2.2,and let K be the empirical NTK Gram matrix (3) .
Assume that the weights of the network areinitialized as ( W l ) i,j ∼ i .
i .
d .
N (0 , β 2 l /n l − 1 ) for l ∈
[ L − 1] , i ∈ [ n l − 1 ] , j ∈ [ n l ] and ( W L ) i ∼ i .
i .
d .
N (0 , β 2 L ) for i ∈ [ n L − 1 ] .
Then, we have
[Equation]
with probability at least 1 − C Ne − c log 2 n L − 1 − Ce − c log 2 N over ( x i )
Ni
=1 and ( W k )
Lk =1 , where cand C are numerical constants.
Moreover, we have that
[Equation]
with probability at least 1 − Ce − cn L − 1 , over ( x i ) Ni =1 and ( W k )
Lk =1 .
[Figure]
This result implies that the NTK is well conditioned fora class of networks in which the number of parametersn L − 2 n L − 1 between the last two hidden layers is linear inthe number of data samples N , up to logarithmic factors.
This means that the total number of neurons of the networkcan be as little as ˜Ω( √ N ) , which meets the minimum possible amount of over-parameterization .
The lower bound(6) and the upper bound (7) match when all the widths (upto layer L − 2 ) have the same scaling, i.e., d = Θ( n L − 2 ) .
Throughout the paper, we do not explicitly track the dependence of our bounds on L , in the sense that the numericalconstants c, C may depend on L .
In Figure 1, we consider a 3-layer neural network withd = n 1 = n 2 , and we plot λ min ( K ) as a function of d 2 ,for three different values of N .
The inputs are sampledfrom a standard Gaussian distribution, the activation function is the sigmoid σ ( x ) =
(1+ e − x ) − 1 , andwe set β l = 1 for all l ∈
[ L ] .
We repeat the experiment 10 times, and report average and confidenceinterval at 1 standard deviation.
The linear scaling of λ min ( K ) in d 2 is in agreement with the resultof Theorem 3.1.
3.1 Roadmap of the proof of Theorem 3.1
After an application of the chain rule and some standard manipulations, we have
[Equation]
where B k ∈ R N × n k is a matrix whose i -th row is given by
[Equation]
From the decomposition (8), one readily obtains that K ⪰ K L − 2 , since K k is PSD for all k ∈{ 0 , . . .
, L − 1 } .
Hence, it suffices to prove a lower bound on λ min ( K L − 2 ) .
We denote by J L − 2 the
Jacobian obtained by computing the gradient only over the parameters between layer L − 2 and layer
L − 1 .
Then, K L − 2 =
J L − 2 J ⊤ L − 2 and, by using (8) and (9), the i -th row of J L − 2 can be expressed as
[Equation]
where D L is a diagonal matrix with the entries of W L on its diagonal.
Note that, if we fix the weightmatrices ( W l ) Ll =1 , the rows (( J L − 2 ) i : )
Ni =1 are i.i.d., as the training data ( x i )
Ni =1 are i.i.d.
too.
The proof consists of two main parts.
First, we construct the centered Jacobian ˜ J L − 2 , which isobtained from J L − 2 by iteratively removing expectations with respect to X to its parts, and we showthat its smallest singular value is close to the smallest singular value of J L − 2 .
The details of this partare contained in Appendix D. Second, we bound the smallest eigenvalue of the kernel obtained fromthe centered Jacobian ˜ J L − 2 by providing an accurate estimate of the ℓ 2 and sub-exponential norms ofits rows.
The details of this part are contained in Appendix E.
In order to carry out this program, we exploit a number of concentration results on the ℓ 2 norms of feature and backpropagation vectors, and onthe ℓ 2 norms of their centered counterparts.
These results are contained in Appendix C, and they couldbe of independent interest.
In particular, we provide tight high-probability estimates on (i) ∥
f l ( x ) ∥ 2 ,(ii) E x [ ∥ f l ( x ) ∥ 22 ] , (iii) E x [ ∥ f l ( x ) ∥ 2 ] , (iv) E x [ ∥ f l ( x ) − E x
[ f l ( x )] ∥ 22 ] , E x [ ∥ f l ( x ) − E x
[ f l ( x )] ∥ 2 ] , and∥ f l ( x ) − E x
[ f l ( x )] ∥ 2 , and (v)
∥ D L ϕ ′ ( g L − 1 ( x ))
− E x
[ D L ϕ ′ ( g L − 1 ( x ))]
∥ 2 .
Some preliminary calculations are also contained in Appendix B.
Part 1: Centering.
We consider the centered Jacobian ˜ J L − 2 , whose i -th row is defined as
[Equation] [Equation]
where ˜ f L − 2 ( x i )
=
˜ ϕ
( g L − 2 ( x i )) , with ˜ ϕ
( g L − 2 ( x i ))
=
ϕ ( g L − 2 ( x i ))
− E x i
[ ϕ ( g L − 2 ( x i ))] and˜ ϕ ′ ( g L − 1 ( x i ))
=
ϕ ′
( g L − 1 ( x i ))
− E x i [ ϕ ′ ( g L − 1 ( x i ))] being the centered versions of the activation function and of its derivative, respectively.
Our strategy is to relate λ min ( J L − 2 J ⊤ L − 2 ) to λ min ( ˜ J L −
˜ J ⊤ L − 2 )via the following two steps.
Step (a): Centering F L − 2 and B L − 1 .
We show that λ min ( J L − 2 J ⊤ L − 2 ) ≥ λ min ( ˜ J FB ˜
J ⊤ FB ) −o ( n L − 2 n L − 1 ) ,
where ˜ J FB ˜
J ⊤ FB =
˜ F L − 2 ˜ F ⊤ L − 2 ◦ ˜ B L − 1 ˜ B ⊤ L − 1 and ˜ F L − 2 =
F L − 2 − E X
[ F L − 2 ] and˜ B L − 1 =
B L − 1 − E X
[ B L − 1 ] (Lemma D.1 in Appendix D.1).
Our strategy differs from existingwork (e.g., [61, 48]), where either only the backpropagation term is centered (cf.
Proposition 7.1 of[61]) or only the features are centered (cf.
Lemma 5.4 of [48]).
We remark that, in order to handledeep networks with minimum overparameterization 2 , the centering of F L − 2 and B L − 1 is cruciallyperformed at the same time .
In fact, by doing so, certain terms containing an eigenvalue whichis negative and large in modulus suitably cancel.
As a result, the difference between the originalkernel J L − 2 J ⊤ L − 2 and the centered one ˜ J FB ˜
J ⊤ FB can be bounded by a matrix whose operator normis o ( n L − 2 n L − 1 ) .
Step (b): Centering everything.
We show that λ min ( ˜ J FB
[Equation]
⊤ FB ) ≥ λ min ( ˜ J L − 2 ˜
J ⊤ L − 2 ) − o ( n L − 2 n L − 1 ) ,where the rows of ˜ J L − 2 are given by (11) (Lemma D.2 in Appendix D.2).
The idea is to decomposethe difference
˜ J FB
˜
J ⊤ FB − ˜ J L − 2 ˜ J ⊤ L − 2 into a rank-1 matrix plus a PSD matrix.
Then, in order tobound the operator norm of the rank-1 term, we leverage the general version of the Hanson-Wrightinequality given by Theorem 2.3 in [1] (see the tail bound on quadratic forms that we provide in
Lemma B.5).
This strategy avoids whitening and dropping rows (as done in [61]) and, hence, appearsto be better suited to the deep case.
The main result of this part is stated below, and it is proved by combining Lemmas D.1 and D.2.
Theorem 3.2 (Jacobian centering) .
Consider the setting of Theorem 3.1, and let the rows of J L −
2and ˜ J L − 2 be given by (10) and (11) , respectively.
Then, we have
[Equation] [Equation]
with probability at least 1 − C exp( − c log 2 n L − 1 ) − C exp( − c log 2 N ) over ( x i ) Ni =1 and ( W k )
Lk =1 ,where c, C are numerical constants.
Part 2: Bounding the centered Jacobian.
If we fix the weight matrices ( W l ) Ll =1 , the rows of˜ J L − 2 are i.i.d.
vectors of the form u ⊗ v − E [ u ⊗ v ] , where both u and v are centered.
By exploitingthis structure, in Appendix E.1 we bound the ℓ 2 and sub-exponential norms of these rows.
Theℓ 2 norm is bounded in Lemma E.2, and this result relies on the tight estimates on the ℓ 2 norms ofcentered features and centered backpropagation terms given by Lemma C.4 and C.5, respectively.
The sub-exponential norm is bounded in Lemma E.3, and we use again the tail bound of Lemma
B.5, which exploits the version of the Hanson-Wright inequality in [1].
At this point, the problemis reduced to bounding the smallest eigenvalue of a matrix such that its rows are i.i.d.
and we havea good control on their ℓ 2 and sub-exponential norms.
This goal can be achieved via the followingresult, whose proof follows from Theorem 3.2 in [2] and it is presented for completeness in Appendix
E.2.
Proposition 3.3.
Let ˜ J be a matrix with rows ˜ J i : ∈ R n L − 2 n L − 1 , for i ∈
[ N ] .
Assume that { ˜ J i : } i ∈ [ N ]are independent sub-exponential random vectors with ψ = max i ∥
˜ J i : ∥ ψ 1 .
Let η min =
min i ∥
˜ J i : ∥ 2and η max = max i ∥
˜ J i : ∥ 2 .
Set ξ = ψK + K ′ , and ∆ :=
Cξ 2 N 1 / 4 ( n L − 1 n L − 2 ) 3 / 4 .
Then, we havethat
[Equation]
holds with probability at least
[Equation]
where c and C are numerical constants.
The main result of this part is stated below, and it is proved in Appendix E.3.
Theorem 3.4 (Bounding the centered Jacobian) .
Consider the setting of Theorem 3.1, and let therows of ˜ J L − 2 be given by (11) .
Then, we have
[Equation]
with probability at least 1 − Ce − c √ N − C Ne − c log 2 n L − 1 over ( x i )
Ni
=1 and ( W k )
Lk =1 , where c and
C are numerical constants.
Recall that K ⪰ K L − 2 , hence λ min ( K ) ≥ λ min ( K L − 2 ) =
λ min (cid:0) J L − 2 J ⊤ L − 2 (cid:1) .
Thus, the lowerbound (6) follows by combining the results from Theorem 3.2 and Theorem 3.4.
Finally, the upper bound (7) is rather straightforward.
First, we write
[Equation]
Then, by combining the bound on ∥ ( F l ) 1: ∥ 22 of Lemma C.1 with a direct calculation for ∥ ( B l +1 ) 1: ∥ 22 ,the desired result (7) follows.
The details are contained in Appendix F.
Two Applications: Memorization and Optimization
Memorization capacity.
The fact that the NTK Gram matrix (3) is well conditioned readily implies result on memorization capacity.
This was already observed in [42] for two-layer networks andin [48] for deep networks with a layer containing Ω( N ) neurons.
Here, by using Theorem 3.1, weshow that Ω( N ) parameters between the last pair of hidden layers are enough for the network to fit
N real-valued labels up to an arbitrarily small error.
The result is stated below and, for completeness,we give the proof in Appendix G.
Corollary 4.1 (Memorization) .
Consider an L -layer neural network (1) , where the activation functionsatisfies Assumption 2.3 and the layer widths satisfy Assumptions 2.4 and 2.5.
Let { x i } Ni =1 ∼ i .
i .
d .
P X ,where P X satisfies the Assumptions 2.1-2.2.
Then, it holds that for every Y ∈ R N , and for everyε > 0 , there exists a set of parameters θ such that
[Equation]
with probability at least 1 − C Ne − c log 2 n L [Equation] 1 − Ce − c log 2 N over ( x i ) Ni
=1 , where c and C arenumerical constants.
Gradient descent training.
Theorem 3.1 has implications on the convergence of gradient descentalgorithms.
In particular, by choosing carefully the initialization, we show that gradient descenttrained on the N samples { x i } Ni =1 with labels Y
=
( Y 1 , . . .
, Y N ) converges to zero loss.
Thisis – at the best of our knowledge – the first result of this kind for deep networks with minimumover-parameterization .
To define our initialization, let us assume that the ( L − 1)
-th layer has an even number of neurons.
With a slight abuse of notation, we indicate its width as 2 n L − 1 .
Thus, W L − 1 can be written in theform (cid:104) W (1) L − 1 , W (2)
L − 1 (cid:105) , where W (1) L − 1 , W (2) L − 1 ∈ R n L − 2 × n L − 1 .
Similarly, W L ∈ R 2 n L − 1 is the concatenation of W (1) L , W (2) L ∈ R n L − 1 .
Then, we define the initialization θ
= [vec( W 1 , 0 ) , . .
, vec( W L, 0 )]as follows:
[Equation]
i
[Equation] [Equation]
[Equation] [Equation] [Equation]
[Equation]
i .
[Equation] [Equation] [Equation]
) , i
[Equation]
[ n L
[Equation]
2 ] , j ∈ [ n L
[Equation]
1 ] , and W (2) L {+− 1 , 
+}
= W (1) L − 1 , 0 , ( W (1) L, 0 ) i ∼ i .
i .
d .
N (0 , β 2 L γ ) , i ∈ [ n L − 1 ] , and W (2) L, 0 = − W (1) L [Equation]
[Equation]
As usual, ( β l ) Ll =1 are numerical constants independent of the layer widths.
In contrast, the quantity γwill be set to a suitably large value depending on the number of training samples and on the layerwidths.
In words, the initialization is standard for W 1 , 0 , . . .
, W L − 2 , 0 and W (1) L − 1 , 0 ; the variance of
W (1) L, 0 is boosted up by a factor γ ; and we duplicate the neurons of layer L −
so that W (2) L − 1 , 0 =
W (1) L − 1 , 0 and W (2) L, 0 = − W (1) L, 0 .
At this point, we are ready to state our optimization result.
Theorem 4.2 (Optimization via gradient descent) .
Consider an L -layer neural network (1) , wherethe activation function satisfies Assumption 2.3 and the layer widths satisfy Assumptions 2.4 and 2.5.
Consider training data { x i } Ni =1 ∼ i .
i .
d .
P X , where P X satisfies the Assumptions 2.1-2.2, with labels
Y ∈ R N such that ∥ Y ∥
2 = Θ( √ N ) .
We consider solving the least-squares optimization problem
[Equation]
19) by running gradient descent updates of the form θ t +
=
θ t − η ∇L ( θ t ) , where the initialization θ 0 isdefined in (18) with γ = d 3 N 2 and η ≤ C ( γNdn L − 1 ) − 1 .
Then, for all t ≥ 1 ,
[Equation]
(20) with probability at least 1 − C
[Equation]
over ( x i ) Ni
=1 and the initialization θ 0 ,where c and C are numerical constants.
[Figure]
In words, (20) guarantees that gradient descent linearlyconverges to a point with zero loss.
We also note that (ourupper bound on) the convergence rate can be expressedas a simple function of the learning rate η , the number oftraining samples N , and the layer widths { n i } L − 1 i =0 .
In Figure 2, we give an illustrative example that 4-layernetworks achieve 0 loss when the number of parameters isat least linear in the number of training samples, i.e., underminimum over-parameterization.
To ease the experimentalsetup, we use a ReLU activation, with Adam optimizer.
We initialize the network as in the setting of Theorem 3.1,picking β l = 1 for all l ∈
[ L ] .
The inputs, as well as thetargets, are sampled from a standard Gaussian distribution.
The plot is averaged over 10 independent trials.
Aspredicted by Theorem 4.2, the loss experiences a phasetransition and it goes from 1 to 0 when the layer widthsare of order √ N .
We now provide a sketch of the argument of Theorem 4.2, deferring the detailed proof to Appendix
H. As a consequence of Theorem 3.1, the NTK is well conditioned at the initialization θ 0 , whichimplies that gradient descent makes progress at the beginning.
The idea is that the NTK remainssufficiently well conditioned also during training and, hence, the loss vanishes with the linear rate in(20), since the trained weights remain confined in a ball of sufficiently small radius centered at θ 0 .
More specifically, we exploit the optimization framework developed in [49].
There, a convergenceresult is proved if the following condition holds: there exists α, β such that
[Equation]
and
[Equation]
[Equation] for all θ, θ 1 , θ 2 ∈ D = B
( θ 0 , R ) , where D denotes a ball centered at θ 0 of radius R := 4 ∥
F L ( θ 0 ) −
Y ∥ 2 /α (cf.
Proposition H.1).
In order to control the radius R , we take advantage of the form(18) of the initialization θ 0 .
In particular, the spectrum of the Jacobian J (and, hence, α ) scaleslinearly in √ γ , and the network output at initialization F L ( θ 0 ) is equal to 0 (cf.
Lemma H.2).
Thus, as ∥ Y ∥ 2 = Θ( √ N ) , we have that R = Θ( √ N/α ) , which can be controlled by choosingsuitably γ .
At this point, the argument consists in providing a number of estimates on feature vectors,backpropagation terms, and finally the Jacobian J ( θ ) for all θ ∈ D (cf.
Lemmas H.3-H.9).
Thisallows us to find α, β such that (21) holds and conclude.
Related Work and Discussion
Random matrices in deep learning.
The limiting spectra of several random matrices related toneural networks have been the subject of a recent line of work.
In particular, the Conjugate Kernel(CK) – namely, the Gram matrix of the features from the last hidden layer – has been studied formodels with two layers [54, 37, 39, 51], with a bias term [3, 56], and with multiple layers [13].
The
Hessian matrix has been considered in [52], and the closely related Fisher information matrix in [55].
Using tools from free probability, the input-output Jacobian (as opposed to the parameter-output
Jacobian considered in this work) of deep networks is studied in [53] and the NTK of a two-layermodel in [4].
In [23], the spectrum of NTK and CK for deep networks is characterized via an iterated
Marchenko-Pastur map.
The generalization error has also been studied via the spectrum of suitablerandom matrices, see [29, 41, 38, 42] and Section 6 of the review [11].
Most of the existing resultsfocus on the linear-width asymptotic regime, where the widths of the various layers are linearlyproportional.
An exception is [68], which focuses on the ultra-wide two-layer case ( n 1 ≫ N ).
Smallest eigenvalue of empirical kernels.
In the line of work mentioned above, the limitingspectrum of the random matrix is characterized in terms of weak convergence, which does not lead toa direct implication on the behavior of its smallest eigenvalue.
In fact, lower bounding the smallesteigenvalue often requires understanding the speed of convergence to the limit spectrum.
Quantitativebounds for random Fourier features are obtained in [9].
In the two-layer setting, the smallest NTKeigenvalue is lower bounded in [61, 42, 68], and concentration bounds can also be obtained from[4, 63, 31].
In particular, [42] establishes a tight bound for two-layer networks with roughly as manyparameters as training samples (hence, under minimum over-parameterization), and for a wide classof activations.
However, this result still requires that n 1 = Ω( d ) .
For deep networks, a convergencerate on the NTK can be obtained from [7] and potentially from [15], however these results require allthe layers to be wide.
In [48], it is proved that a single layer of linear width suffices for the NTK tobe well conditioned.
Here, Theorem 3.1 provides the first result for deep networks under minimumover-parameterization, therefore allowing for sub-linear layer widths.
Memorization capacity and gradient descent training.
For binary classification, the memorization capacity of neural networks was first studied by Cover [18], who solved the single-neuroncase.
Later, Baum [12] proved that, for two-layer networks, the memorization capacity is lowerbounded by the number of parameters, and upper bounds of the same order were given in [35, 58].
Recently, tight results for two-layer networks in the regression setting have been provided [14, 42].
In particular, [14] generalizes the construction of [12], while the memorization result of [42] comesfrom a bound on the smallest NTK eigenvalue.
This same route is also followed to analyze deepnetworks in [48] and in this work.
Results for multiple layers have been provided in [10, 71, 66, 24].
As concerns efficient algorithms achieving memorization, [19, 20] study classification with two-layernetworks.
A popular line of research has exploited the NTK view to give optimization guaranteeson gradient descent.
Existing work has focused on the two-layer case [61, 22, 50, 63, 69, 62], andoptimal results in terms of over-parameterization have been obtained [42].
For deep networks, [21]assumes a lower bound on the smallest NTK eigenvalue, [5, 73, 74] require all the layer widths to be(rather large) polynomials in the number of samples, and [47, 43] reduce the over-parameterizationto a single layer of linear width.
In particular, in these last two papers, the analysis of the NTK isreduced to that of a certain feature matrix (the first one in [47] and the last one in [43]).
Hence,these approaches do not seem suitable to tackle networks with sub-linear layer widths 3 .
Finally, weremark that milder over-paramerization requirements are sufficient for classification problems.
Morespecifically, it has been proved that polylogarithmic width suffices for two-layer [34] and deep [16]
ReLU networks to converge and generalize.
In fact, achieving memorization for regression is harderthan for classification, where it suffices to satisfy inequality constraints [42].
Concluding Remarks
In this paper, we show that the NTK is well conditioned for deep networks with sub-linear layerwidths.
As consequences of our NTK bounds, we also prove results on memorization and optimizationfor such a class of networks.
Our approach requires smooth activations and layer widths which do notsignificantly grow in size with the depth.
We believe such assumptions to be purely technical, and weleave as an open question their removal.
Our novel NTK analysis could potentially be applied beyondstandard gradient descent, e.g., for optimization with momentum [67], or federated learning [32]. Checklist For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] See Assumptions 2.1, 2.2, 2.3,
2.4 and 2.5 in Section 2; see also l. 384-385 in Section 6.
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
Given the
theoretical nature of the work, we do not see direct paths to applications with negative
societal impacts.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] All the
assumptions are clearly stated either in the description of the model in Section 2 or in
the statements of the main results (Theorem 3.1, Corollary 4.1, and Theorem 4.2).
(b) Did you include complete proofs of all theoretical results?
[Yes] A sketch of the
proof of Theorem 3.1 is provided in Section 3.1.
The complete proof is contained in
Appendices B-F: Appendix B contains preliminary calculations; Appendix C contains
the high-probability estimates on the ℓ 2 norms of feature and backpropagation vectors;
Appendix D contains the complete proofs for Part 1:
Centering ; Appendix E contains
the complete proofs for Part 2: Bounding the centered Jacobian ; and Appendix F
contains the proof of the upper bound (7).
The proof of Corollary 4.1 is contained in
Appendix G. Finally, the proof of Theorem 4.2 is sketched in Section 4 right after the
statement, and the full proof is contained in Appendix H.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main exper-
imental results (either in the supplemental material or as a URL)?
[Yes] The folder
containing the code and the instructions is in the supplemental material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters,
how they
were chosen)?
[Yes]
(c) Did you report error bars (e
g., with respect to the random seed after running experi-
ments multiple times)?
[Yes] In Figure 1, we repeat the experiment 10 times, and wereport the average together with the confidence interval at 1 standard deviation.
In
Figure 2, we repeat the experiment 10 times and we report the average of the final loss
in the heat-map.
(d) Did you include the total amount of compute and the type of resources used (e
g., type
of GPUs, internal cluster, or cloud provider)?
[N/A]
We did not perform large scale
experiments that require to specify this type of details.
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[N/A] We are not usingexisting assets or curating/releasing new assets.
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
We did not use crowdsourcing or conducted research with humansubjects.
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]