Identifying good directions to escape the NTK regime and efﬁciently learn low-degree plus sparse polynomials
Abstract
A recent goal in the theory of deep learning is to identify how neural networkscan escape the “lazy training,” or Neural Tangent Kernel (NTK) regime, wherethe network is coupled with its ﬁrst order Taylor expansion at initialization.
Whilethe NTK is minimax optimal for learning dense polynomials [22], it cannot learnfeatures, and hence has poor sample complexity for learning many classes offunctions including sparse polynomials.
Recent works have thus aimed to identifysettings where gradient based algorithms provably generalize better than the NTK.
One such example is the “QuadNTK” approach of Bai and Lee [7], which analyzesthe second-order term in the Taylor expansion.
Bai and Lee [7] show that thesecond-order term can learn sparse polynomials efﬁciently; however, it sacriﬁcesthe ability to learn general dense polynomials.
In this paper, we analyze how gradient descent on a two-layer neural network canescape the NTK regime by utilizing a spectral characterization of the NTK [36] andbuilding on the QuadNTK approach.
We ﬁrst expand upon the spectral analysisto identify “good” directions in parameter space in which we can move withoutharming generalization.
Next, we show that a wide two-layer neural network canjointly use the NTK and QuadNTK to ﬁt target functions consisting of a denselow-degree term and a sparse high-degree term – something neither the NTK northe QuadNTK can do on their own.
Finally, we construct a regularizer whichencourages our parameter vector to move in the “good" directions, and show thatgradient descent on the regularized loss will converge to a global minimizer, whichalso has low test error.
This yields an end to end convergence and generalizationguarantee with provable sample complexity improvement over both the NTK and
QuadNTK on their own.
Introduction
In recent years, Deep Learning has acheived a number of practical successes, in domains spanningcomputer vision, natural language processing, reinforcement learning, and the sciences.
Despitethese impressive empirical results, the theory underlying deep learning is far from complete.
In fact,the dual questions of optimization – the mechanism by which neural networks trained with gradientdescent are able to interpolate training data despite the nonconvexity of the loss landscape – andgeneralization – why these solutions found by gradient descent require relatively few samples togeneralize – are still not well understood.
One successful approach for understanding optimization has been the Neural Tangent Kernel (NTK)theory [37, 26, 12, 17].
The NTK approach couples the gradient descent dynamics of a wide neuralnetwork with a speciﬁc initialization to the gradient descent dynamics of a particular kernel regression problem, with a random, initialization dependent kernel.
In the limit of inﬁnite width, this kernelconverges almost surely to a deterministic kernel, also referred to as the NTK, and properties of thiskernel and its corresponding Reproducing Kernel Hilbert Space can be studied.
However, recent work has shown that the NTK theory fails to explain the generalization capabilitiesof neural networks.
While the equivalence between neural networks in the NTK regime and kernelmethods implies that such models perform no better than kernels, in practice, neural networks havebeen shown to outperform kernel methods on a number of tasks [5, 30].
Theoretically, a recentline of work [22, 36, 35] has provided a precise statistical analysis of the generalization propertiesof rotationally invariant kernels on the unit sphere, which includes the NTK.
[22] proves a samplecomplexity lower bound for such kernels, showing that d k samples are needed to learn any degree kpolynomial in d dimensions.
As a result, the NTK is no better than a polynomial kernel, and cannotadapt to low-dimensional structure.
The limitations of the NTK can be further understood from the linearization perspective.
Consider atwo-layer neural network f ( x ; W ) of width m with ﬁrst layer weights initialized as W 0 ∈ R d × m ,second layer weights a ∈ R m , and activation function σ , and let W ∈ R
d
× m be the displacementfrom initialization, i.e
[Equation]
The NTK theory states that when W has small norm, the gradient descent dynamics can be wellapproximated by replacing the model with its ﬁrst-order Taylor expansion about the initialization:
[Equation]
where ϕ ( x ) := vec ( ∇ W f ( x ; W ) | W = 0 ) is a random feature vector with norm independent of m ,and
[Equation]
is hereafter denoted as the linear term .
Ignoring d dependence, there exists a global minimizer with(cid:107) W (cid:107) F (cid:39) 1 and (cid:107) w r (cid:107) 2 (cid:39) m − 1 / 2 , and thus gradient descent stays in this small norm ball around theinitialization while interpolating the training data.
This small movement of each individual neurongives rise to the name lazy training for networks in the NTK regime [12].
This equivalence to kernelmethods and the poor generalization of neural networks in NTK regime, along with their failure todescribe the dynamics of neural networks in practice, motivate our goal to understand how neuralnetworks can escape the NTK regime.
We concretely ask the following question:
Q: How can we encourage each neuron to move (cid:29) m − 1 / 2 , thus escaping the NTK regime?And does this allow us to break the NTK sample complexity lower bounds?
1.1 Motivation
Escaping the NTK Regime.
To answer this, we invoke the statistical characterization of the NTKdeveloped in [22, 36, 35] to understand the mechansism by which it overﬁts to the training data.
Assume that for d k (cid:28) n (cid:28) d k +1 we have n training samples ( x i , y i ) , where the x i are sampled i.i.dfrom S d − 1 ( √ d ) (the d -dimensional sphere of radius √ d ) and y i
=
f ∗ ( x ) for an unknown functionf ∗ .
Montanari and Zhong [36] show that the SVD of the empirical feature matrix Φ ∈ R n × md canbe written in a block matrix form asΦ =   ϕ ( x 1 ) T · · · ϕ ( x n ) T 
 =
[ U
U 2 ] (cid:20) ˜ Λ  ˜ Λ(cid:21)
(cid:20) V T
V T(cid:21) , where U 1 , ˜ Λ 1 , V 1 are the top r singular values/vectors and U 2 , ˜ Λ
2 , V 2 are the bottom n − r .r
=
O ( d k ) is chosen speciﬁcally so that V 1 is a “high-variance" subspace which can expresspolynomials of degree ≤ k and generalizes well [8], and V 2 interpolates the training data while notaffecting generalization.
Montanari and Zhong [36] show that the NTK will learn the projection of f ∗ onto degree k polynomials.
Furthermore, [8, 23] show that gradient descent will ﬁrst move in the subspace spanned by
V 1 to learn this projection, then move in the V 2 directions to interpolate the training data, while notmodifying the test predictions.
This ﬁrst stage is still desirable for our goal – if the network can ﬁtand generalize from part of the signal using the NTK, then it should, and previous work [24] hasshown that for general networks gradient descent will learn the optimal degree-1 polynomial in theearly stages of training.
The V 2 directions, however, are “bad" directions the parameters should avoidmoving in, as they are only used for the NTK to overﬁt.
Instead, W should move in the null space ofΦ , where (cid:107) W (cid:107) F can be Ω(1) while keeping the evaluation of f L ( · ; W ) on the training data boundedby O (1) .
This heuristic argument yields our ﬁrst criterion for escaping the NTK regime.
Goal #1: Move minimally in the V 2 directions.
Generalizing to Test Data.
While the previous criterion prevents the network from overﬁtting withthe NTK, we must also prevent a movement of (cid:107) W (cid:107) F (cid:29) 1 from causing the test predictions toexplode.
Since (cid:107) ϕ
( x ) (cid:107)
=
Θ m (1) , it is a priori possible for f L ( x ; W ) =
vec ( W ) T ϕ ( x ) to be (cid:29) 1on
the population, which would necessarily cause a large test loss.
To identify a set of good directions, deﬁne the feature covariance matrix Σ ∈ R md × md by
[Equation]
Our ﬁrst technical contribution is a characterization of the eigendecomposition of Σ .
Let the top reigenvectors of Σ be Q 1 , the next s eigenvectors be Q 2 , and the bottom md − r − s be Q 3 , whereΛ 1 , Λ 2 , Λ 3 are the corresponding diagonal matrices of eigenvalues, and r = Θ( d k ) , s = d Θ( k ) .
Weuse the results of [36, 35] to show that Σ has an eigenvalue gap in that λ min ( Λ 1 ) (cid:29) λ max ( Λ 2 ) andλ min ( Λ 2 ) (cid:29) λ max ( Λ 3 ) , and furthermore that Q 1 can ﬁt degree ≤ k polynomials.
This partitioningof the eigenvectors tells us that Q 1 are informative, large eigenvalue directions which help the NTKto learn a low degree signal, Q 2 are the medium directions which will cause the test predictions togrow too large if W moves too far from initialization, and Q 3 are the “good" directions which W isfree to move a distance of (cid:29) 1 in.
This yields the following criterion for generalizing well:
Goal #2: Move in the Q 3 directions, but minimally in the Q 2 directions.
A challenge, however, is that given d k samples we are unable to distinguish the Q 2 directions fromthe Q 3 directions.
Nevertheless, the existence of an eigenvalue gap allows us to constrain movementin Q 2 .
The Quadratic NTK.
Once we have moved (cid:29) 1 from the initialization, we can no longer couple tothe network’s linearization.
The network is still, however, in a local regime, and we instead can couplethe training dynamics to the second-order Taylor expansion of our model, where the second-orderterm is denoted the Quadratic-NTK [7] f Q ( x ; W ) :
[Equation]
Bai and Lee [7] showed that f Q ( x ; W ) can effectively learn low-rank polynomials with better samplecomplexity than the NTK.
In particular, they show that d k samples are needed to ﬁt a target functionof the form f ∗ ( x ) =
( β T x ) k +1 , an improvement over the d k +1 samples needed by the NTK.
Indoing so, however, the QuadNTK sacriﬁces its ability to learn general dense polynomials.
Basedon the preceeding motivation, we thus aim to show that we can jointly utilize both the NTK and
QuadNTK to learn a larger class of functions than either the NTK or QuadNTK can learn on theirown.
1.2 Our Contributions
With this intuition in hand, we outline the main contributions of this work.
We ﬁrst prove a technicalresult on the eigendecomposition of Σ , and show that the eigenvectors can indeed be partitioned into3 categories corresponding to large (degree ≤ k ), medium ("bad") and small ("good") eigenvalues.
We then construct a regularizer, depending only on the covariate distribution and initialization, thatenforces goals 1 and 2 by preventing the parameters from moving in either of the bad sets of directions( V 2 and Q 2 ) .
Furthermore, we show how to jointly use the NTK and QuadNTK to ﬁt a target signalf ∗ consisting of an arbitrary degree ≤ k component and sparse degree k + 1 component.
The key technical challenge is to construct a solution with large enough movement so that the QuadNTKcan ﬁt the high degree term and hence improve generalization, while simultaneously preventing thislarge movement from interfering with the NTK training predictions (Goal 1) or greatly increasingtest loss (Goal 2).
Our main result, 1, is that gradient descent on a polynomially wide two-layerneural network converges to an approximate global minimizer of the regularized loss function, whichgeneralizes well to the test distribution.
As a result, we show d k samples are needed to learn f ∗ up tovanishingly small test loss.
This ultimately gives us the “best of both worlds," where we can learn thetarget f ∗ with sample complexity better than either the NTK or Quad-NTK alone.
Overall, our workidentiﬁes which directions weights can move further from initialization in and provably generalizebetter than the NTK.
The outline of our paper is as follows.
In Section 2 we formally deﬁne the problem setup.
In Section 3we deﬁne our regularizers, and present Theorem 1.
Section 4 is an outline of the proof of Theorem 1,which we split into four components – expressing f ∗ with the linear-plus-quad model, showingthe optimization landscape has favorable geometry, a gradient descent convergence result, and ageneralization bound.
We conclude with experiments supporting our main results.
1.3 Related Work
The NTK approach [37, 26, 12, 29, 17], which couples a neural network to its linearization atinitialization, has been utilized to show global convergence of gradient descent on neural networks [16,32 , 43, 4].
The equivalence to kernel methods has also been used to prove generalization bounds,based on generalization bounds for kernels [6, 10, 3].
However, neural networks have been shown toperform far better than their NTK in practice [5, 30].
Further, [22] shows that kernels cannot adaptto low-dimensional structure, and proves a sample complexity lower bound of d k samples neededto learn a degree k polynomial in d dimensions.
A number of recent works [40, 15, 42, 1, 20, 2, 21,13, 41, 33, 34] have thus aimed to provide examples of learning problems where a neural networktrained with a gradient-based algorithm has a provable sample complexity improvement over anykernel method.
One such approach has been to understand higher-order approximations of the training dynamics [7,11, 25].
Here, the network is no longer coupled to its linearization, but rather higher order termsin the Taylor expansion.
[7, 11] show that the second-order term, the QuadNTK, can be used toobtain sample complexity improvements.
However, the QuadNTK has poor sample complexity forlearning dense polynomials; our work shows how polynomials with a dense and sparse componentcan be learned jointly.
Furthermore, these works do not consider training on the original network,but rather only the second-order term after the linear term has been deleted.
We consider training onthe full two-layer neural network, and leverage the NTK to provide an end-to-end convergence andgeneralization result.
The technical results in our paper rely on the statistical characterization of the NTK developedin the series of works [22, 36, 35].
Furthermore, our optimization results rely on a line of workshowing that quadratically parameterized models have nice landscape properties such as all secondorder saddle points are global minima [18, 19, 37, 15]; the fact that gradient descent avoids saddlepoints [18, 31, 27, 28] can then be used to show convergence.
Preliminaries
2.1 Problem Setup
Our problem setup is the standard supervised learning setting.
Our dataset D n = { ( x i , y i ) } i ∈ [ n ] , hasn samples, where ( x i , y i ) ∈
X × Y are sampled i.i.d from a distribution µ on X × Y .
µ is deﬁned sothat ( x , y ) ∼ µ satisﬁes x ∼ Unif ( S d − 1 ( √ d )) , the uniform distribution on the d -dimensional sphereof radius √ d , and y = f ∗ ( x ) for some deterministic, unknown function f ∗ : S d − 1 ( √ d ) → R .
We assume that d k (cid:28) n (cid:28) d k +1 for some integer k , and that the target f ∗ has the followinglow-degree plus sparse structure:
Assumption 1 (Low-degree plus sparse signal) .
Let f ∗ ( x ) =
( β T x ) k +
+ f k ( x ) , where (cid:107) β (cid:107) ≤ 1and f k is an arbitrary degree ≤ k polynomial with E x ∼ µ [ f k ( x ) 2 ] =
We aim to ﬁt f ∗ with f ( x ; W ) , a two-layer neural network as deﬁned in (1).
Here, W =[ w 1 , . . .
, w r ] ∈ R d × m is the ﬁrst layer weight’s distance from initialization and is the trainable parameter.
W 0 =
[ w 0 , 1 , . . .
, w 0 ,r ] denotes the ﬁrst layer weight at initialization, and anda =
[ a 1 , . .
. , a m ] T ∈ R m is the second layer weight, which is held ﬁxed throughout training.
We consider the following symmetric initialization of a , W 0 , which ensures that f ( · ; 0 ) = 0 identically.
[Equation]
σ ∈ C 2 ( R ) is our nonlinear activation function.
We make the following assumption on σ :
Assumption 2.
The activation σ satisﬁes (cid:107) σ (cid:107) ∞ ,
[Equation]
[Equation]
We also require σ (cid:48) , σ (cid:48)(cid:48) to satisfy Assumption 3, a particular technical condition on their harmonicexpansions.
These two assumptions are satisﬁed by commonly used activations, such as the sigmoidwith generic shift b : σ ( z )
=
1 1+exp( b − z ) .
We assume the loss function (cid:96) : R × R → R ≥ 0 satisﬁes (cid:96) ( y, z ) ≤ 1 , (cid:96) ( y, y )
= 0 , (cid:96) ( y, z ) convex inz , and (cid:107) ∂∂z (cid:96) (cid:107) ∞ , (cid:107) ∂ 2 ∂z 2 (cid:96) (cid:107) ∞ , (cid:107) ∂ 3 ∂z 3 (cid:96) (cid:107) ∞ ≤ 1 .
The empirical loss ˆ L and population loss L are deﬁnedas
[Equation]
where for a function g ( x , y ) , E n [ g ( x , y )] := 1 n (cid:80)
ni
=1 g ( x i , y i ) denotes the empirical expectation,while E µ [ g ( x , y )] denotes the population expectation over ( x , y ) ∼ µ .
Notation.
For f ∈ L 2 ( S d ( √ d ) , µ ) , deﬁne (cid:107) f (cid:107) L 2 :
= (cid:107) f (cid:107) L 2 ( S d ( √ d ) ,µ ) =
(cid:0) E x ∼ µ [( f ( x )) 2 ] (cid:1) 1 / 2 .
Weuse big O notation to ignore absolute constants that do not depend on n, d, m , and write a d (cid:46) b
d ifa d
= O ( b d ) , a d (cid:28) b d if lim d →∞ a d /b d = 0 .
We also use ˜ O notation to ignore terms that dependlogarithmically on d .
We also treat k
= O (1) .
Finally, all our results hold for d > C , where C isa universal constant.
For a matrix A , we let (cid:107) A (cid:107) F be its Frobenius norm, (cid:107) A (cid:107) =
(cid:107) A (cid:107) op be theoperator norm, and (cid:107) A (cid:107) 2 ,p :=
( (cid:80) i (cid:107) a i (cid:107) p 2 ) 1 /p be the 2 , p norm.
2.2 Linear and Quadratic Expansion
For W small, f ( · ; W ) can be approximated by its second order Taylor expansion about W 0 :
[Equation]
We deﬁne f L ( x ; W ) , f Q ( x ; W ) to be the linear and quadratic terms of the network:
[Equation]
Main Theorem
Deﬁne the NTK featurization map ϕ : S d − 1 ( √ d ) → R md as
[Equation]
11) and the feature covariance matrix Σ ∈ R
md × md as
[Equation]
Note that Σ depends only on the network at initialization and the input distribution, and not on thetarget function f ∗ .
In practice, Σ can be approximated to arbitrary precision by using a large datasetof unlabeled data, or by computing the harmonic expansion of σ (cid:48) , as detailed in Appendix A.
Let Σ admit the eigendecomposition Σ
= (cid:80) mdi
=1 λ i ( Σ )
v i v Ti , where the λ i ( Σ ) are nonnegative andnonincreasing.
For r ∈ [ md ] , we let Π ≤ r be the projection operator onto span ( v 1 , . . .
, v r ) , and letΠ >r
=
I
md − Π ≤ r .
Furthermore, deﬁne
[Equation]
We deﬁne our regularizers as follows
[Equation]
Intuitively, R 3 constrains movement in the V 2 directions to enforce Goal #1, R 1 constrains movementin the Q 2 directions to enforce Goal #2, and R 2 , R 4 are weight-decay like terms necessary forgeneralization.
Although R 1 does not know the directions Q 2 , the eigenvalue gap between Q 2 and
Q 3 ensures that whenever R 1 is small, movement in Q 2 must be small as well.
Given regularization parameters λ
= ( λ 1 , λ 2 , λ 3 , λ 4 ) , deﬁne the regularized loss L λ ( W ) as
[Equation]
Finally, we train our model trained via perturbed gradient descent [28] with learning rate η and noiselevel σ 2 .
That is, if W t denotes the weights at time step t , the update is given by
[Equation] [Equation]
where Ξ t ∈ R d × m are i.i.d random matrices with each entry i.i.d N ( σ 2 md ) .
Given these deﬁnitions, we now present our main theorem:
Theorem 1.
Let ε > 0 be a target test accuracy, the number of samples be n (cid:38) max( d k ε − 2 , d k log d ) ,and the width be m = poly ( n, d, ε − 1 ) .
Let the sequence of iterates { W t } t ≥ 0 follow the updatein (19) with initialization W 0 = 0 .
Then, there exists a choice of parameters ( λ 1 , λ 2 , λ 3 , λ 4 , r, σ 2 , η )such that with high probability over W 0 , D n , and { Ξ t } t ≥ 0 , there exists a T = poly ( m ) such thatthe predictor ˆ W :=
W T satisﬁes L ( ˆ W ) ≤ ε .
Remark 1.
Theorem 1 tells us n = ˜Θ( d k ) samples are needed to learn the low-degree plus sparsefunction f ∗ .
This is an improvement over the sample complexity needed to learn f ∗ via the NTK,which is Ω( d k +1 ) samples as f ∗ is a degree ( k + 1) -polynomial
This also improves overthe upper bound for the sample complexity of the quadratic NTK given in [7], in which Ω( R 2 d k − 1 )samples are needed to learn a rank R polynomial of degree k
+ 1 .
This is polynomially worse thanour bound since the dense, low-degree term f k can have rank (cid:29) d .
Proof Sketch
The proof of Theorem 1 follows high-level similar steps to [7].
We ﬁrst construct a W ∗ ∈ R d × mwhich ﬁts f ∗ and is small on the regularizers.
Next, we show that the optimization landscape of theregularized loss has a favorable geometry, and as a result that gradient descent converges to a globalminimum.
We conclude with a generalization bound to show that global minima have low test loss.
Throughout the proof sketch, we emphasize how the regularizers encourage us to escape the NTKregime, and discuss the challenges posed by the existence of the f L term in the dynamics.
4.1 Expressivity
We begin by showing that the f L and f Q terms can ﬁt the low degree and sparse components,respectively, of the signal.
As in [22, 36], the derivations in this section rely on spherical harmonics;an overview of the technical results used are presented in Appendix A.
The following lemma shows that f Q ( x ; W ) can ﬁt the sparse, high degree term in f ∗ .
Lemma 1.
Let m ≥ d 2 k .
With high probability over the initialization, there exists W Q such that
[Equation]
and (cid:107) W Q
[Equation] [Equation]
[Equation]
This generalizes the corresponding result in [7] to more activations.
The proof of this lemma ispresented in Appendix B.1.2.
Next, we show that f L ( x ; W ) can ﬁt the low degree term.
Here, we choose r = n k = Θ( d k ) , wheren k is deﬁned in Appendix A, to be the dimension of the subspace which can express degree ≤ kpolynomials.
We deﬁne P ≤ k = Π ≤ r to be the projection on the top n k eigenvectors of Σ .
Lemma 2 (NTK can ﬁt low degree component) .
Let m ≥ d 10 k .
With high probability, there exists
W L with vec ( W L ) ∈ span ( P ≤ k ) such that
[Equation]
The proof of this Lemma is presented in Appendix B.2.3, and relies on key lemmas from [36] relatingthe spherical harmonics of degree ≤ k to the eigenstructure of the kernel.
A key intermediate result is
Lemma 15, which characterizes the spectrum of the population covariance matrix Σ .
A similar resultfor random features was shown in [35].
Unlike [35], we do not characterize all the eigenvalues of Σ ;however, simply partitioning them into the three categories is sufﬁcient for our purposes.
Finally, we use the W L , W Q to construct a W ∗ which has small regularized loss.
Recall thedeﬁnition of the regularizers, after setting r = n k :
[Equation]
Also, deﬁne the empirical loss of the quadratic model as:
[Equation]
The following theorem is the central expressivity result:
Theorem 2.
For ε min > 0 , let m (cid:38) max
( d 3( k − 1) ε − 4 min , d 10 k ) , n (cid:38) max( d k ε − 2 min , d k log d ) .
Withprobability 1 − 1 / poly ( d ) , there exists W ∗ such that
[Equation] [Equation]
and:
[Equation]
[Equation]
The proof of this Theorem is presented in Appendix B. [Equation] and again relies on the eigendecompositionof Σ .
As outlined in the introduction, we show Σ
= Q 1 Λ 1 Q [Equation] [Equation]
[Equation]
T 3 , whereΛ 2 are the medium eigenvalues and Λ 3 are the small eigenvalues.
Formally, Λ 2 contain Θ( d − i +1 )with multiplicity Θ( d i ) , for integers
i ∈
[ k + 1 , 2 k ] , and the entires of Λ 3 are equal to 1 /m onaverage.
This tells us that the medium directions Q 2 are undesirable, as any Ω m (1) movement inthese directions will case f L ( x ; W ) to grow large.
On the other hand, movement in a “sufﬁcientlyrandom" Q 3 direction will minimally affect the population value of f L ( x ; W ) .
To prove Theorem 2, we will construct W ∗ to be of the form W L + W Q , where W L ﬁts the lowdegree term and W Q ﬁts the sparse term.
The issue this direct construction is that (cid:107) W Q (cid:107) F (cid:29) 1 , soa priori W Q can have a large effect on the linear term.
We thus require f L ( x ; W Q ) to be small, bothon the sample and over the population.
The key insight is the following: since dim ( V 1 ) (cid:29) dim ( V 2 ) ,any sufﬁciently random direction lies almost entirely in V 2 .
If u is this random direction, then 273 E n (cid:2) ( f L ( x ; u )) 2 (cid:3) is small.
Similarly, the sufﬁciently random direction u lies almost entirely in Q 3 .
Since the Q 3 directions have eigenvalues 1 /m on average, E µ (cid:2) ( f L ( x ; u )) 2 (cid:3) ≈ 1 m (cid:107) u
(cid:107) 2 (cid:28) 1 .
We thus consider a “sufﬁciently random" version of W Q so that W Q minimally affects f L .
Speciﬁcally, we consider the weight SW Q , where S ∈ R m × m is a diagonal matrix of random signs.
Bydeﬁnition f Q ( x ; SW Q ) =
f Q ( x ; W Q ) , furthermore, we show SW Q is now sufﬁciently random inthat E n (cid:2) ( f L ( x ; u )) 2 (cid:3) , E µ (cid:2) ( f L ( x ; u )) 2 (cid:3) are both small.
This allows us to prove that in expectationover S , a solution of the form W L + SW Q acheives small regularized loss, which implies theexistence of such a desirable solution via the probabilistic method.
4.2 Landscape
Deﬁne ﬁrst and second-order stationary points [28] as follows:
Deﬁnition 1.
W is a ν -ﬁrst-order stationary point of f if (cid:107)∇ f ( W ) (cid:107) ≤ ν .
Deﬁnition 2.
W is a ( ν, γ ) -second-order stationary point (SOSP) of f if ∇ 2 f ( W ) (cid:23) − γ I .
The following lemma is our central landscape result:
Lemma 3.
Let r =
d k , λ 2 = ε min , λ
= m 12 d − k − 12 ε min , λ
4 = d − 2( k − 1) ε min .
Assume m ≥n 8 d 26( k +1)
ε − 22 / 3 min , and ν ≤ m − 14 .
Let W be a ν -ﬁrst order stationary point, and let W ∗ be thesolution constructed in 2.
Then,
[Equation]
As a corollary, we show that any ( ν, γ ) -SOSP of L λ ( W ) has small loss.
Corollary 1.
Set r, λ as in Lemma 3.
Let ˆ
W be a ( ν, γ ) -second-order stationary point of L λ ( W ) ,with ν ≤ m − 1 / 2 , γ ≤ m − 3 / 4 .
Then L λ ( ˆ W ) (cid:46)
ε min .
The proofs of these results are deferred to Appendix C
4.3 Optimization
We next invoke the main theorem of [27, 28], that perturbed gradient descent will ﬁnd a ( ν, γ ) -SOSPin poly (1 /ν, 1 /γ ) time.
The challenge with applying these results directly is that L λ is no longersmooth or Hessian-Lipschitz due to the R 4 regularizer.
To circumvent this, we ﬁrst prove that theiterates in perturbed gradient descent are bounded in a Frobenius norm ball.
Then, it sufﬁces to usethe bound on smoothness and Hessian-Lipschitzness in this ball to prove convergence.
Our mainoptimization result, with proof in Appendix D, is the following:
Theorem 3.
There exists a choice of learning rate η and perturbation radius σ such that withprobability 1 − 1 / poly ( d ) , perturbed gradient descent (c.f [28][Algorithm 1]) reaches a ( ν, γ ) -SOSPwithin T = poly ( m ) timesteps.
4.4 Generalization
Finally, we conclude by showing that any ˆ W with small L λ ( ˆ W ) also has small test loss:
Theorem 4 (Main generalization theorem) .
Let r = d k , λ
=
m 1 2 d − k − 1 2 ε min , λ 2 = ε min , λ 3 =m 12 d − k − 12 ε min , λ
4 = d − 2( k − 1) ε min .
Assume m (cid:38) ε − 4 min d 3( k − 1) .
With probability 1 − 1 / poly ( d )over the draw of D , any data dependent ˆ W with L λ ( ˆ W ) ≤
Cε min has population loss
L
[Equation]
The proof of this theorem is presented in
Appendix E.
The key is to use the small values of theregularizers at ˆ W to bound the test loss.
R 2 and R 4 are used to bound the Rademacher complexitiesof the linear and quadratic terms respectively, while R 1 and R 3 control the inﬂuence of the highdegree component of the linear term on the population loss and empirical loss respectively.
Experiments
We conclude with experiments to support our main theorem.
In Figure 1 we train the joint linearand quadratic model f L ( x ; W ) + f Q ( x ; W ) via gradient descent on the square loss, with signal
2 log d t     3 ( W ) {+=0.
 =1e-+} =0.0001 =0.
=0.01 =0.
2 log d t 0.
T r a i n
L o ss  2 log d t 0.
T e s t
L o ss f ∗ ( x ) = f 1 ( x ) + ( β T x ) 2 where
f 1 ( x ) = x 1 − 1 .
Our covariates are dimension d = 100 , we haven =
d 1 . 5 = 1000 samples, and our network has width m = 10000 .
We train our model with the regularizer R 3 ( W ) for varying values of λ 3 .
For all λ 3 we reach nearzero training error.
However, we observe that when λ 3 is zero or very small, the model strugglesto learn the entire signal, and test error plateaus near 0.6.
In the leftmost pane, we plot the valueof R 3 ( W ) over the course of training.
We observe that R 3 ( W ) grows large for small λ 3 , whichimplies that the model is using the “bad" directions in V 2 to overﬁt to the training data.
For largevalues of λ 3 , however, R 3 prevents W from moving in the "bad" directions to overﬁt the data.
Thisis seen in the leftmost pane, where the value of R 3 ( W ) is small.
Instead, the parameter moves inthe “good” Q 3 directions, and the Quad-NTK term kicks in to ﬁt the remaining component of thesignal while generalizing.
This leads to the consistently lower test loss (around 0.2) as shown in therightmost pane.
On the Regularizer R 1 .
Our proof required the existence of R 1 to prevent movement in the bad
Q 2 directions.
While R 1 being small is necessary for generalization, in Figure 1 we observe that wegeneralize well without explicitly regularizing R 1 .
We hypothesize that the noise in the perturbedgradient descent update may be implicitly regularizing R 1 .
A heuristic justiﬁcation is as follows:recall that any sufﬁciently random direction u lies mostly in Q 3 .
Due to the rotational symmetry ofthe perturbations, one might expect that the solutions P ≤ k ˆ W + SP k ˆ W , over all choices of randomsigns S , can be reached with roughly equal probability.
Since this set of solutions is “sufﬁcientlyrandom," in expectation over S it generalizes well, and thus gradient descent is likely to converge toa solution which generalizes well.
In Appendix G we conduct experiments with both R 1 and R 3 tosupport the claim that explicit regularization of R 1 is not needed to acheive low test error; however,proving this claim is left for future work.
Discussion
The goal of this work is to better understand how neural networks can escape the NTK regime.
Byanalyzing the eigendecomposition of the feature covariance matrix, we identiﬁed 3 sets of directions –ones that can ﬁt low degree signal, “bad" directions which either cause the NTK to overﬁt or the testpredictions to explode, and “good" directions in which the parameters can move a large distance.
Wethen constructed a regularizer which encourages movement in these good directions, and showed howa network can jointly use the linear and quadratic terms in its Taylor expansion to ﬁt a low-degreeplus sparse signal.
Altogether, we provided an end-to-end convergence and generalization guaranteewith a provable sample complexity improvement over the NTK and Quad-NTK.
As discussed above, one interesting direction of future work is to understand the role of R 1 .
Otherdirections of future work include understanding whether our analysis can be used to leverage higherorder terms in the Taylor expansion, and investigating whether increasing the depth of the networkcan allow the NTK and Quad-NTK to jointly learn a hierarchical representation [11].
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’scontributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] We discuss the assumptions ourresults rely on in Section 2.
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
This work
is a theoretical work, and does not have any direct 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] Assumptionson our model are given in Section 2.
(b) Did you include complete proofs of all theoretical results?
[Yes] We give a proof sketchof our main result, Theorem 1, in Section 4, and present the full proof in the Appendix.3.
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] We specify
instructions in Appendix G, and will submit code with the supplementary material.
(b) Did you specify all the training details (e g., data splits, hyperparameters, how theywere chosen)?
[Yes] We specify training details in Appendix G.
(c) Did you report error bars (e.g., with respect to the random seed after running experi-
ments multiple times)?
[Yes] We report error bars in all plots.
(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)?
[Yes] Computational resources are
detailed in Appendix G
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?
[Yes] We cite JAX [9], the
framework used for our experiments.
Our experiments use synthetic data.
(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 identiﬁable
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]
(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]