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 networks can escape the “lazy training,” or Neural Tangent Kernel (NTK) regime, where the network is coupled with its ﬁrst order Taylor expansion at initialization.
While the NTK is minimax optimal for learning dense polynomials [25], it cannot learn features, and hence has poor sample complexity for learning many classes of functions including sparse polynomials.
Recent works have thus aimed to identify settings where gradient based algorithms provably generalize better than the NTK.
One such example is the “QuadNTK” approach of Bai and Lee [7], which analyzes the second-order term in the Taylor expansion.
Bai and Lee [7] show that the second-order term can learn sparse polynomials efﬁciently; however, it sacriﬁces the ability to learn general dense polynomials.
In this paper, we analyze how gradient descent on a two-layer neural network can escape the NTK regime by utilizing a spectral characterization of the NTK [39] and building on the QuadNTK approach.
We ﬁrst expand upon the spectral analysis to identify “good” directions in parameter space in which we can move without harming generalization.
Next, we show that a wide two-layer neural network can jointly use the NTK and QuadNTK to ﬁt target functions consisting of a dense low-degree term and a sparse high-degree term – something neither the NTK nor the QuadNTK can do on their own.
Finally, we construct a regularizer which encourages the parameter vector to move in the “good" directions, and show that gradient descent on the regularized loss will converge to a global minimizer, which also has low test error.
This yields an end to end convergence and generalization guarantee 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 spanning computer vision, natural language processing, reinforcement learning, and the sciences.
Despite these 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 gradient descent are able to interpolate training data despite the nonconvexity of the loss landscape – and generalization – why these solutions found by gradient descent require relatively few samples to generalize – are still not well understood.
One successful approach for understanding optimization has been the Neural Tangent Kernel (NTK) theory [41, 29, 15, 20].
The NTK approach couples the gradient descent dynamics of a wide neural network under 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 kernel converges almost surely to a deterministic kernel, also referred to as the NTK, and properties of this kernel and its corresponding Reproducing Kernel Hilbert Space can be studied.
However, recent work has shown that the NTK theory fails to explain the generalization capabilities of neural networks.
While the equivalence between neural networks in the NTK regime and kernel methods implies that such models perform no better than kernels, in practice, neural networks have been shown to outperform kernel methods on a number of tasks [5, 33].
Theoretically, a recent line of work [25, 39, 38] has provided a precise statistical analysis of the generalization properties of rotationally invariant kernels on the unit sphere, which includes the NTK.
[25] proves a sample complexity lower bound for such kernels, showing that d k samples are needed to learn any degree k polynomial in d dimensions.
As a result, the NTK is no better than a polynomial kernel, and cannot adapt to low-dimensional structure.
The limitations of the NTK can be further understood from the linearization perspective.
Consider a two-layer neural network f ( x ; W ) with input x , width m , ﬁrst layer weights initialized as W 0 ∈ R d × m , second layer weights a ∈ R m , activation function σ , and displacement from initialization
W
[Equation] [Equation]
For simplicity we assume the second layer weights are held ﬁxed, and so W are the trainable parameters.
The NTK theory states that when W has small norm, the gradient descent dynamics can be well approximated by replacing the model with its ﬁrst-order Taylor expansion about the initialization:
[Equation]
Here, { w r } r ∈ [ m ] and { w 0 ,r } r ∈ [ m ] are the columns of W and W 0 respectively, ϕ ( x ) := vec ( ∇ W f
( x ; W )
| W = 0 ) is a random feature vector with norm independent of m , and
[Equation]
is hereafter referred to 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 due to local convexity gradient descent will stay in this small norm ball around the initialization while interpolating the training data.
This small movement of each individual neuron gives rise to the name lazy training for networks in the NTK regime [15].
The equivalence to kernel methods and the poor generalization of neural networks in NTK regime, along with their failure to describe the dynamics of neural networks in practice, motivate our goal to understand how neural networks 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 NTK developed in [25, 39, 38] to understand the mechansism by which it overﬁts to the training data.
For d k (cid:28) n (cid:28) d k +1 , consider a dataset of n training samples ( x i , y i ) , where the x i are sampled i.i.d from S d − 1 ( √ d ) (the d -dimensional sphere of radius √ d ) and
y i =
f ∗ ( x ) for an unknown function f ∗ .
Montanari and Zhong [39] decompose the SVD of the empirical feature matrix Φ ∈ R n × md into the block matrix form
[Equation]
where U 1 , ˜ Λ 1 , V 1 are the top r singular values/vectors and U 2 , ˜ Λ 2 , V 2 are the bottom n − r .
Here, r = O ( d k ) is chosen speciﬁcally so that V 1 is the “high-variance" subspace which can express polynomials of degree ≤ k and generalizes well [9], while V 2 is used to interpolate the training data while not affecting generalization.
[25, 39] show that the NTK will learn the projection of f ∗ onto degree k polynomials.
Furthermore, [9, 26] 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 not affecting the test predictions.
This ﬁrst stage is still desirable for our goal – if the network can ﬁt and generalize from part of the signal using the NTK, then it should, and previous work [27] has shown that for general networks gradient descent will learn the optimal degree-1 polynomial in the early stages of training.
The V 2 directions, however, are “bad" directions the parameters should avoid moving 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 bounded by O (1) .
This heuristic argument yields our ﬁrst criterion for escaping the NTK regime.
[Equation] [Equation]
minimally in the V 2 directions.
Generalizing to Test Data.
While the previous criterion prevents the network from overﬁtting with the NTK, we must also prevent a movement of (cid:107) W (cid:107) F (cid:29) 1 from causing the test predictions to explode.
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)on 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 r eigenvectors of Σ be Q 1 , the next s eigenvectors be Q 2 , and the bottom md − r − s eigenvectors be Q 3 , where Λ 1 , Λ 2 , Λ 3 are the corresponding diagonal matrices of eigenvalues, and r = Θ( d k ) , s =
d Θ( k ) .
We use the results of [39, 38] 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 arbitrary degree ≤ k polynomials.
This partitioning of the eigenvectors tells us that Q 1 are informative, large eigenvalue directions which help the NTK to learn a low degree signal, Q 2 are the medium directions which will cause the test predictions to grow too large if W moves too far from initialization, and Q 3 are the “good" directions which W is free 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.
One challenge is that we cannot distinguish the Q 2 directions from the Q 3 directions with d k samples.
Nevertheless, the existence of an eigenvalue gap will allow us to constrain movement in Q 2 .
The Quadratic NTK.
Once we have moved (cid:29) 1 from the initialization, we can no longer couple to the network’s linearization.
The network is still, however, in a local regime, and we instead can couple the training dynamics to the second-order Taylor expansion of our model, where the second-order term is denoted the Quadratic-NTK [7] f Q ( x ; W
[Equation] [Equation]
Bai and Lee [7] showed that f Q ( x ; W ) can effectively learn low-rank polynomials with better sample complexity than the NTK.
In particular, they show that d k samples are needed to ﬁt a target function of the form f ∗ ( x ) =
( β T x ) k +1 , an improvement over the d k +1 samples needed by the NTK.
In doing so, however, the QuadNTK sacriﬁces its ability to learn general dense polynomials; furthermore, Bai and Lee [7] require a randomization trick to artiﬁcially delete the f L term.
A later followup work [8] showed that in a number of standard experimental settings, the second-order Taylor expansion of the network better tracks the true gradient descent dynamics and acheives lower test loss than the network’s linearization (i.e NTK) does.
However, there is no existing result which shows that both the linear term and the quadratic term can provably learn a component of the signal.
Based on the preceeding motivation, we thus aim to show that we can jointly utilize both the NTK and the QuadNTK to learn a larger class of functions than either the NTK or QuadNTK can learn on their own.
1.2 Our Contributions
With the previous intuition in hand, we outline the main contributions of our work.
We ﬁrst prove a technical result on the eigendecomposition of Σ , and show that the eigenvectors can indeed be partitioned into 3 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, that enforces 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 signal f ∗ consisting of an arbitrary degree ≤ k component and sparse degree k +component.
The key technical challenge is to construct a solution with large enough movement so that the QuadNTK can ﬁt the high degree term and hence improve generalization, while simultaneously preventing this large movement from interfering with the NTK training predictions (Goal 1) or greatly increasing test loss (Goal 2).
Our main result, Theorem 1, is that gradient descent on a polynomially wide two-layer neural network converges to an approximate global minimizer of the regularized loss function, which generalizes well to the test distribution.
As a result, we show d k samples are needed to learn f ∗ up to vanishingly small test loss.
This ultimately gives us the “best of both worlds”, as we leverage both the linear and quadratic term to learn the target f ∗ with sample complexity better than either the NTK or QuadNTK alone.
Overall, our work identiﬁes which directions weights can move further from initialization in and provably generalize better than the NTK.
The outline of our paper is as follows.
In Section 2 we formally deﬁne the problem setup.
In Sectionwe 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, showing the optimization landscape has favorable geometry, a gradient descent convergence result, and a generalization bound.
We conclude with experiments supporting our main theorem and demonstrating the relevance of the low-degree plus sparse task to standard neural networks.
1.3 Related Work
The NTK approach [41, 29, 15, 32, 20], which couples a neural network to its linearization at initialization, has been utilized to show global convergence of gradient descent on neural networks [19, 35, 47, 4].
The equivalence to kernel methods has also been used to prove generalization bounds, based on generalization bounds for kernels [6, 11, 3].
However, neural networks have been shown to perform far better than their NTK in practice [5, 33].
Further, [25] shows that kernels cannot adapt to low-dimensional structure, and proves a sample complexity lower bound of d k samples needed to learn a degree k polynomial in d dimensions.
A number of recent works [44, 18, 46, 1, 23, 2, 24, 16, 45, 36, 14, 37] have thus aimed to provide examples of learning problems where a neural network trained with a gradient-based algorithm has a provable sample complexity improvement over any kernel method.
One such approach has been to understand higher-order approximations of the training dynamics [7, 8, 12, 28].
Here, the network is no longer coupled to its linearization, but rather higher order terms in the Taylor expansion.
On the empirical side, [8] shows that these higher order Taylor expansions better track the optimization dynamics and can obtain lower test loss.
Theoretically, [7, 12] prove that the second-order term, the QuadNTK, can be used to obtain sample complexity improvements.
However, the QuadNTK has poor sample complexity for learning dense polynomials, and [7, 12] do not consider training on the original network, but rather only the second-order term after the linear term has been deleted.
This paper, on the other hand, provides an end-to-end convergence and generalization result for training on the full two-layer neural network.
We leverage the NTK to efﬁciently learn polynomials with both a dense and sparse component, and are thus the ﬁrst work showing that both the linear and quadratic term can learn part of the signal.
The technical results in our paper rely on the statistical characterization of the NTK developed in the series of works [25, 39, 38].
Furthermore, our optimization results rely on a line of work showing that quadratically parameterized models have nice landscape properties such as all secondorder saddle points are global minima [21, 22, 41, 18]; the fact that gradient descent avoids saddle points [21, 34, 30, 31] 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 ] , has n samples, where ( x i , y i ) ∈
X × Y are sampled i.i.d from a distribution µ on X × Y .
µ is deﬁned so that ( x , y ) ∼ µ satisﬁes x ∼ Unif ( S d − 1 ( √ d )) , the uniform distribution on the d -dimensional sphere of 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 following low-degree plus sparse structure:
Assumption 1 (Low-degree plus sparse signal) .
Let f ∗ ( x ) =
f k ( x ) + f sp ( x ) , where
• f k ( x ) is an arbitrary degree ≤ k polynomial with E x
∼
µ [ f k ( x ) 2 ] = 1 .
(Low Degree) • f sp ( x )
=
(cid:80) Ri =
α i ( β Ti x ) k +1 where |
α
i | ≤ 1 , (cid:107) β
i (cid:107) 2 = 1 .
(Sparse)
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 and a = [ 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) ∞ , (cid:107) σ (cid:48) (cid:107) ∞ , (cid:107) σ (cid:48)(cid:48) (cid:107) ∞ < 1 .
We also require σ (cid:48) , σ (cid:48)(cid:48) to satisfy Assumption 3, a particular technical condition on their harmonic expansions.
These two assumptions are satisﬁed by commonly used activations, such as the sigmoid with 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 in z , 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ﬁned as
[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 .
We use big O notation to ignore absolute constants that do not depend on n, d, m , as well as polynomial dependencies on the rank R .
We write a d (cid:46) b
d if a 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 depend logarithmically on d .
We also treat k
= O (1) .
Finally, all our results hold for d > C , where C is a 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 the operator 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]
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 the target function f ∗ .
In practice, Σ can be approximated to arbitrary precision by using a large dataset of unlabeled data, or by computing the harmonic expansion of σ (cid:48) , as detailed in Appendix A.
Let Σ admit the eigendecomposition Σ
= (cid:80) md
i =
λ i ( Σ ) v i v T i , where the λ i ( Σ ) are nonnegative and nonincreasing.
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 Σ ≤ r := r (cid:88)
[Equation]
We deﬁne our regularizers as follows
[Equation]
[Equation]
Intuitively, R 3 constrains movement in the V 2 directions to enforce Goal #1, R 1 constrains movement in the Q 2 directions to enforce Goal #2, and R 2 , R 4 are weight-decay like terms necessary for generalization.
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 [31] with learning rate η and noise level σ 2 .
That is, if W t denotes the weights at time step t , the update is given by
[Equation]
(21) 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) d k · poly ( R ) · max( ε − 2 , log d ) , and the width be m = poly ( n, d, R, ε − 1 ) .
Let the sequence of iterates { W t } t ≥follow the update in (21) 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 that the 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 sparse function 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 over the upper bound for the sample complexity of the quadratic NTK given in [7], in which Ω( K 2 d k ) samples are needed to learn a rank K polynomial of degree k + 1 .
This is polynomially worse than our bound since the dense, low-degree term f k can have rank (cid:29) d .
Proof Sketch
The proof of Theorem 1 follows similar high-level steps to [7].
We ﬁrst construct a W ∗ ∈ R d × m which ﬁts f ∗ and is small on the regularizers.
Next, we show that the optimization landscape of the regularized loss has a favorable geometry, and as a result that gradient descent converges to a global minimum.
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 NTK regime, 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 [25, 39], 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 (QuadNTK can ﬁt high degree component) .
Let m ≥ d 2 k .
With high probability over the initialization, there exists W Q such that
[Equation]
This generalizes the corresponding result in [7] to more activations.
The proof of this lemma is presented 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 ) , where n k is deﬁned in Appendix A, to be the dimension of the subspace which can express degree ≤ k polynomials.
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.4, and relies on key lemmas from [39] relating the 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 result for random features was shown in [38].
Unlike [38], 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 the deﬁnition of the regularizers, after setting r = n k :
[Equation]
Also, deﬁne the empirical loss of the quadratic model as:
[Equation] [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 ) .
With probability 1 − 1 / poly ( d ) , there exists W ∗ such that
[Equation]
29) and:
[Equation]
The proof of this Theorem is presented in Appendix B.3, and again relies on the eigendecomposition of Σ .
As outlined in the introduction, we show Σ
= Q 1 Λ 1 Q T 1 + Q 2 Λ 2 Q T 2 + Q 3 Λ 3 Q 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 on average.
This tells us that the medium directions Q 2 are undesirable, as any Ω m (1) movement in these directions will case f L ( x ; W ) to grow large.
On the other hand, movement in a “sufﬁciently random” 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 with this direct construction is that (cid:107) W Q (cid:107) F (cid:29) 1 , so a priori W Q can have a large effect on the linear term.
We thus require f L ( x ; W Q ) to be small, both on the sample and over the population.
The key insight is the following: since dim ( V 1 ) (cid:28) dim ( V 2 ) , any sufﬁciently random direction lies almost entirely in V 2 .
If u is this random direction, then E n (cid:2) ( f L ( x ; u )) 2 (cid:3) is small.
Similarly, the 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.
By deﬁnition f Q ( x ; SW Q ) =
f Q ( x ; W Q ) , furthermore, we show SW Q is now sufﬁciently random in that 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 expectation over S , a solution of the form W L + SW Q acheives small regularized loss, which implies the existence of such a desirable solution via the probabilistic method.
4.2 Landscape
Deﬁne ﬁrst and second-order stationary points [31] 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 4 d 26( k +1)
ε
[Equation]
22 / 3 min , and ν ≤ m − 14 .
Let W be a ν -ﬁrst order stationary point, and let W ∗ be the solution 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 [30, 31], which is that perturbed gradient descent will ﬁnd a ( ν, γ ) -SOSP in poly (1 /ν, 1 /γ ) time.
The challenge with applying these results directly is that L λ is no longer smooth or Hessian-Lipschitz due to the R 4 regularizer.
To circumvent this, we ﬁrst prove that the iterates in perturbed gradient descent are bounded in a Frobenius norm ball.
Then, it sufﬁces to use the bound on smoothness and Hessian-Lipschitzness in this ball to prove convergence.
Our main optimization result, with proof in Appendix D, is the following:
Theorem 3.
There exists a choice of learning rate η and perturbation radius σ such that with probability 1 − 1 / poly ( d ) , perturbed gradient descent (c.f [31, Algorithm 1]) reaches a ( ν, γ ) -SOSP within T = poly ( m ) timesteps.
2 log d t     3 ( W ) 3 =0.0 3 =1e3 =0.0001 3 =0.
3 =0.01 3 =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
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 12 d − k − 12 ε min , λ 2 = ε min , λ
= 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 ( ˆ W ) (cid:46)
[Equation]
+ (cid:114) d
[Equation]
n .
The proof of this theorem is presented in
Appendix E.
The key is to use the small values of the regularizers at ˆ W to bound the test loss.
R 2 and R 4 are used to bound the Rademacher complexities of the linear and quadratic terms respectively, while R 1 and R 3 control the inﬂuence of the high degree 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 linear and quadratic model f L ( x ; W ) + f Q ( x ; W ) via gradient descent on the square loss, with signal f ∗ ( x ) = f 1 ( x ) + ( β T x ) 2 where f 1 ( x ) = x 1 − 1 .
Our covariates are dimension d = 100 , we have n = 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 .
Rather than computing Σ , we use the top n k right singular vectors of Φ to estimate R 3 .
We observe that for all λ 3 we reach near zero training error.
However, when λ 3 is zero or very small, the model struggles to learn the entire signal, and test error plateaus near 0.6.
In the leftmost pane, we plot the value of R 3 ( W ) over the course of training.
We observe that R 3 ( W ) grows large for small λ 3 , which implies that the model is using the “bad" directions in V 2 to overﬁt to the training data.
For large values of λ 3 , however, Rprevents W from moving in the “bad" directions to overﬁt the data.
This is seen in the leftmost pane, where the value of R 3 ( W ) is small.
Instead, the parameter moves in the “good” Q 3 directions, and the Quad-NTK term kicks in to ﬁt the remaining component of the signal while generalizing.
This leads to the consistently lower test loss (around 0.2) as shown in the rightmost 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 we generalize well without explicitly regularizing R 1 .
We hypothesize that the noise in the perturbed gradient 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 of the perturbations, one might expect that the solutions P ≤ k ˆ W + P k ˆ WS , over all choices of random signs S , can be reached with roughly equal probability.
Since this set of solutions is “sufﬁciently random,” in expectation over S it generalizes well, and thus gradient descent is likely to converge to a solution which generalizes well.
In Appendix G we conduct experiments with both R 1 and R 3 to support 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.
[Figure]
On the Low-Degree Plus Sparse Task.
Theorem 1 shows that two-layer neural networks trained via noisy gradient descent with a speciﬁc regularizer can more efﬁciently learn low-degree plus sparse target functions.
In Figure 2 we show empirically that neural networks with standard initialization and trained via vanilla gradient descent efﬁciently learn a “dense quadratic plus sparse cubic.”
For varying values of d , we train a two-layer neural network to convergence and plot the smallest n such that the test loss was < 0 .
We observe this minimal n roughly scales with d 2 , the optimal (in d ) sample complexity.
This provides evidence that standard networks trained with vanilla GD can effectively learn the low-degree plus sparse task, and hence that this is a sensible task to study and that our work presents a step towards understanding why standard neural networks perform better in practice more generally.
See Appendix G for more details.
On Second-Order Taylor Expansions of Standard Networks.
Throughout this paper we studied the quadratic Taylor expansion of a two-layer neural network.
In Appendix G, we complement the results of [8] and show that for a standard architecture and data distribution, the quadratic Taylor expansion better approximates the optimization dynamics and acheives lower test loss than the linearization (NTK) does.
Discussion
The goal of this work is to better understand how neural networks can escape the NTK regime.
By analyzing 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 test predictions to explode, and “good” directions in which the parameters can move a large distance.
We then constructed a regularizer which encourages movement in these good directions, and showed how a network can jointly use the linear and quadratic terms in its Taylor expansion to ﬁt a low-degree plus sparse signal.
Altogether, we provided an end-to-end convergence and generalization guarantee with a provable sample complexity improvement over the NTK and QuadNTK.
As discussed above, one interesting direction of future work is to understand the role of R 1 .
Other directions of future work include understanding whether our analysis can be used to leverage higherorder terms in the Taylor expansion, understanding the connection between the QuadNTK and feature learning, and investigating whether increasing the depth of the network can allow the NTK and QuadNTK to jointly learn a hierarchical representation [12].
Acknowledgements
EN acknowledges support from a National Defense Science & Engineering Graduate Fellowship.
JDL and EN acknowledge support of the ARO under MURI Award W911NF-11-1-0304, the Sloan
Research Fellowship, NSF CCF 2002272, NSF IIS 2107304, NSF CIF 2212262, ONR Young
Investigator Award, and NSF-CAREER under award
The authors would like to thank
Alex Damian for helpful discussions throughout the project.
Limitations
Our main theorem relies on a few assumptions, which are detailed in Section 2.
One assumption is that σ satisﬁes Assumptions 2 and 3, which require σ to be sufﬁciently differentiable and have nonzero Hermite/Gegenbauer coefﬁcients.
While this allows for common activations such as sigmoid or tanh (with a generic bias term), it does not include activations such as ReLU .
We believe that either considering a smoothed version of ReLU or including a randomized bias term can help us work around this non-differentiability issue.
Additionally, we believe the boundedness assumption is not essential and can be overcome with an appropriate truncation argument; however, making the boundedness assumption is useful for simplifying the proof.
Another assumption is that the covariates are drawn from the uniform distribution on the unit sphere
S d − 1 ( √ d ) .
This assumption is necessary in order to invoke the statistical characterization of the NTK developed in [25, 39], and we note that either the uniform-on-sphere or Gaussian data assumption has been made in a number of prior works [25, 23, 39, 26].
In practice, data can be normalized to be isotropic.
Furthermore, neural networks have been observed to generalize better than kernel methods in a wide variety of settings, and thus we believe our results are indicative of a more general phenomenon.
A Spherical Harmonics: Technical Background
Below we present relevant results on spherical harmonics, Gegenbauer polynomials, and Hermite polynomials.
These results are from [25, 39], and a more in depth discussion of the technical background can be found in those references.
For (cid:96) ∈
Z ≥ 0 , deﬁne B ( d, (cid:96) ) as
[Equation]
and deﬁne
[Equation] [Equation]
We let L 2 ( S d − 1 ( √ d ) , µ ) denote the space of square-integrable functions over the sphere of radius √ d , with respect to the uniform probability measure µ .
We use the shorthand (cid:104)· , ·(cid:105) L 2 = (cid:104)· , ·(cid:105)
L 2 ( S d − 1 ( √ d ) ,µ ) , and likewise for (cid:107) f (cid:107) L 2 .
The normalized spherical harmonics { Y ( d ) k,i } 0 ≤ i ≤ B ( d,k ) ,k ≥ 0 are a sequence of polynomials such that
Y ( d ) k,i is degree k , and the Y ( d ) k,i form an orthonormal basis of L 2 ( S d − 1 ( √ d ) , µ ) , i.e:
[Equation]
Y
[Equation]
For x ∼ S ( d − 1) ( √ d ) , let ˜ τ 1 d − 1 be the probability measure of √ d (cid:104) x , e (cid:105) and τ 1 d − 1 be the measure of (cid:104) x , e (cid:105) , where e is an arbitrary unit vector.
The Gegenbauer polynomials { Q ( d ) } k ≥ 0 are a basis of L 2 ([ − d, d ] , ˜ τ
1 d − 1 ) such that Q ( d ) k is a degree k polynomial with
Q ( d ) k ( d ) = 1 and
[Equation]
The following identity relates spherical harmonics and Gegenbauer polynomials:
[Equation]
We also use the following fact about Gegenbauer polynomials:
[Equation]
Furthermore
[Equation]
f ∈ L 2 ([ −√ d, √ d ]
[Equation]
τ 1 d − 1 ) can be decomposed via Gegenbauer polynomials as
[Equation]
[Equation]
Thus (cid:107) f (cid:107) L
[Equation]
([ −√
[Equation] [Equation]
=
[Equation]
Let γ be the measure of a standard Gaussian on R .
The normalized Hermite polynomials { h k } h ≥are an orthonormal basis of L 2 ( R , γ ) such that h k is degree k .
For f ∈ L 2 ( R , γ ) , let µ k ( f ) :
= (cid:104) f, h k (cid:105)
L 2 ( R ,γ ) be the k th hermite coefﬁcient.
One observes that τ 1 d − 1 converges weakly to γ .
As a result we have the following connection between Hermite and Gegenbauer coefﬁcients:
[Equation]
A.1 Assumptions
We additionally require σ (cid:48) , σ (cid:48)(cid:48) to satisfy the following assumptions, which are that the Hermite/Gegenbauer coefﬁcients are nonzero and well behaved:
Assumption 3.
We assume σ, σ (cid:48) satisfy the following: (a) Let σ (cid:48) satisfy µ (cid:96) ( σ (cid:48) )
(cid:54) = 0 for (cid:96) ≤ 4 k
and (cid:80) (cid:96)> 4 k µ 2 (cid:96) ( σ (cid:48) )
> 0 .
As a result, we can let µ 2 (cid:96) ( σ (cid:48) )
=
Θ d (1) for (cid:96) < 4 k .
(b) Let σ (cid:48)(cid:48) satisfy
[Equation]
42) where λ d,(cid:96) ( σ (cid:48)(cid:48) )
= (cid:104) σ (cid:48)(cid:48) , Q ( d ) (cid:96) ( √ d · ) (cid:105)
L 2 ([ −√ d, √ d ] ,τ 1 d − 1 ) .
A.2 Computing Σ
Below we discuss how to express Σ in terms of the Gegenbauer coefﬁcients of σ (cid:48) , λ d,(cid:96) ( σ (cid:48) ) .
First, observe that that Σ is a matrix of d × d blocks, where the i, j th block is equal to
[Equation]
where u : S d − 1 ( √ d ) × S d − 1 ( √ d ) → R d × d is the function
[Equation]
[25, Lemma 7] shows that there exist scalar valued functions u 1 , u 2 , u 3 such that
[Equation]
where u 1 , u 2 , u 3 can be computed in terms of the quantities
[Equation]
It thus sufﬁces to compute these quantities.
We assume that we can compute arbitrarily many
Gegenbauer coefﬁcients of σ (cid:48) .
Note that
[Equation]
All these expressions are of the form E µ [ f ( θ T 1 x ) g
( θ T 2 x )] .
For arbitrary f, g , let their decompositions into Gegenbauer polynomials be:
[Equation]
Then, by Equation 38,
[Equation]
which can be computed to desired precision by truncating this inﬁnite sum accordingly.
This only requires knowledge of the Gegenbauer coefﬁcients of f, g .
Given the Gegenbauer coefﬁcients of a function ψ ( z ) , [25, Lemma 6] gives a formula for the Gegenbauer coefﬁcients of zψ ( z ) .
We can therefore write the Gegenbauer coefﬁcients of zσ (cid:48)
( z ) , z 2 σ (cid:48) ( z ) in terms of those of σ (cid:48) , and thus we can approximate this sum to the desired precision.
This procedure allows us to express Σ in terms of the Gegenbauer coefﬁcients of σ (cid:48) .
B Expressivity Proofs B.1 Quad-NTK Proofs B.1.1 Preliminaries
Lemma 4 (Expressing polynomials with random features) .
Let p ≥ 0 , and let σ satisfy the following two assumptions:
(a)
[Equation]
∈
[Equation] [Equation]
[Equation] [Equation]
) (b) d p · min k ≤ p λ d,k ( σ ) =
Ω d (1) , where λ d,k ( σ ) 2 = (cid:104) σ, Q ( d ) k ( √ d · ) (cid:105)
L 2 ([ −√ d, √ d ] ,τ 1 d − 1 ) .
For | α | ≤ 1 , (cid:107) β (cid:107) =
1 , there exists a function a ∈ L 2 ( S d − 1 (1)) such that
[Equation]
and a satisﬁes the norm bound
[Equation]
Proof.
We can decompose σ into a sum over Gegenbauer polynomials
[Equation]
By Equation 37,
[Equation]
Let a be decomposed into spherical harmonics as
[Equation]
for some coefﬁcients c k,i with (cid:80) k ≥ 0 (cid:80) B ( d,k ) i =
c
[Equation]
,i < ∞ .
Since the spherical harmonics form an orthonormal basis of S d − 1 ( √ d ) , we have
[Equation]
Next, for an arbitrary function f ∈ L 2 ([ − √ d, √ d ] , τ 1 d − 1 ) , we can decompose
[Equation]
For f ( t ) =
αt p , λ d,k ( f ) = 0 for k > p , and thus
[Equation]
Deﬁne the sequence of coefﬁcients { c k,i } 0 ≤ k ≤ p, 1 ≤ i ≤ B ( d,k ) by c k,i
[Equation]
Y
[Equation]
d ) k,i ( √ dβ )
λ d
[Equation]
k
[Equation]
f ) λ − 1 d,k ( σ ) , which are well deﬁned for sufﬁciently large d by assumption (b) of the lemma.
Since there are only ﬁnitely many nonzero c k,i , the function a ( w 0 ) is in L 2 ( S d − 1 (1)) .
Also,
[Equation]
as desired.
To obtain a norm bound on a , we can write
[Equation]
since
[Equation]
and thus (cid:107) f (cid:107) L 2 ([ −√ d, √ d ] ,τ 1 d − 1 ) =
Θ d (1) .
Lemma 5 (Expressivity via inﬁnitely many neurons) .
Let k ≥ 1 , and let σ be a twice-differentiable activation such that σ (cid:48)(cid:48) satisﬁes Assumption 3.
Then, there exist functions w + , w − : S d − 1 (1) → R such that
[Equation]
and
[Equation]
Proof.
Note that since σ (cid:48)(cid:48) is continuous and bounded, σ
(cid:48)(cid:48) ∈
L 2 ([ −√ d, √ d ] , τ 1 d − 1 ) .
Therefore applying Lemma 4 with activation σ (cid:48)(cid:48) and degree k − 1 , there exists a function a satisfying
[Equation]
and
[Equation]
Deﬁne
[Equation]
Then
[Equation]
so
[Equation]
Finally, we have the norm bound
[Equation]
B.1.2 Proof of Lemma
Proof.
We show this Lemma holds with probability 1 − d − 10 .
Deﬁne M := (cid:98) m 2 R (cid:99) .
Deﬁne for i ∈
[ R ] , deﬁne the subnetwork f iQ ( x , W Q ) by
[Equation]
We will now construct W Q ∈ R d × m .
For RM <
r ≤ m/ 2 , set w Q,r = w Q,r + m/ 2 = 0 .
As a result, we have that
[Equation]
Our construction will proceed by expressing α i ( β Ti x ) k +1 with f iQ ( x , W Q ) .
For ﬁxed i ∈
[ R ] , let w i + , w i − be from the inﬁnite width construction in Lemma 5, so
that
[Equation]
For integers ( i − 1) M + 1 ≤ r ≤ iM , deﬁne
[Equation]
Then,
[Equation]
Note that we can bound
[Equation]

[Equation]

[Equation]
(cid:32) k −
[Equation]
(cid:88) k (cid:48)
=
B ( d, k ) (cid:33) 1 /(cid:46)
d k − 1 .
Therefore letting Z r = σ (cid:48)(cid:48) ( w T 0 ,r x ) a i ( w 0 ,r )( β T x ) 2 , we have | Z r | (cid:46) d k .
Also, the Z r are i.i.d and satisfy E [ Z r ] =
( β T x ) k +1 .
Therefore by Hoeffding’s inequality, with probability 1 − 12 d − 11 n − 1 we have
[Equation]
where we omit poly ( R ) dependencies inside the big O notation.
Union bounding over j ∈ [ n ] , with probability 1 − 12 d − 11 over the initialization,
[Equation]
Since the above holds for all i ∈
[ R ] , union bounding over i yields that with probability 1 − 12 d − 11 R ≥ 1 − 12 d − 10 ,
[Equation]
To bound the norm of W Q , observe that
[Equation]
Since we can upper bound
[Equation]
−
[Equation]
for ﬁxed i ∈
[ R ] by Hoeffding we have that with probability 1 − 12 d − 11 over the initialization
[Equation]
M iM
[Equation]
Union bounding over each i ∈
[ R ] and using M = Θ( m/R ) =
Θ( m ) , with probability 1 − 12 d −over the initialization we have that
[Equation]
as desired.
Corollary 2.
The solution W Q constructed in Lemma 1 satisﬁes
[Equation]
Proof.
From the proof of Lemma 1, either (cid:107) w r (cid:107) 2 = 0 , or ( i − 1) M + 1 ≤ r ≤ iM or m/ 2 + ( i − 1) M + 1 ≤ r ≤ m/ 2 + iM for some i in which case
[Equation]
B.2 NTK Proofs B.2.1 Symmetric Initialization
Recall the deﬁnition of the NTK featurization map
[Equation]
The symmetric initialization makes this different from the NTK features in [25, 39], which for width ˜ m = m/ 2 is given by
[Equation]
These two features are related by
[Equation]
For the bulk of this section we consider the features ˜ ϕ
( x ) in order to invoke the results from [25, 39].
B.2.2 Preliminaries
For arbitary N (cid:28) md , let D N
=
{ x i } i ∈
[ N ] be a dummy dataset of size N , where each x i is sampled i.i.d from S d − 1 ( √ d ) .
We deﬁne the following random matrices which depend on D N .
Denote by
[Equation]
the feature matrix, and let
[Equation]
be the empirical kernel matrix.
The inﬁnite-width kernel matrix K ∞ N ∈ R N × N has entries
[Equation]
Also, deﬁne Σ N ∈ R ˜
md × ˜ md
to be the empirical covariance matrix, so that
[Equation]
and let
[Equation]
be the population covariance matrix.
We let σ (cid:48) satisfy assumption 3.
Along with the boundedness of σ (cid:48) , this allows us to invoke the following lemmas from [39]:
Lemma 6 ([39], Theorem 3.2) .
With probability 1 − d − 11 ,
[Equation]
As in [39], let Ψ ≤ (cid:96) ∈ R N × n (cid:96) be the evaluations of the degree ≤ (cid:96) spherical harmonics on the N data points.
Lemma 7 ([39], Lemma 2) .
Let d (cid:96) log 2 d (cid:28) N (cid:28) d (cid:96) +1 / log d C .
With probability 1 − d − 11 the inﬁnite-width kernel matrix can be decomposed as
[Equation]
where γ k ≥ 0 is a sequence satisfying γ
= d − 1 (1 + o d (1)) µ 21 ( σ (cid:48) )
[Equation]
γ k = µ 2 k −
[Equation]
σ (cid:48)
[Equation]
+ o d (1) for k
[Equation] [Equation]
and Λ 2 ≤ (cid:96) is a diagonal matrix where B ( d, k ) − 1 γ k has multiplicity B ( d, k ) , for k ≤ (cid:96) .
Furthermore, the remainder ∆ satisﬁes
[Equation]
and the spherical harmonic features Ψ ≤ (cid:96) satisfy
[Equation]
By assumption 3, γ
=
Θ( d − 1 ) and γ (cid:96) (cid:48) =
Θ(1) for (cid:96) (cid:48) ≥ 1 .
The following Lemma gives the eigenvalues of the empirical kernel matrix K N and the inﬁnite-width kernel matrix K ∞ N :
Lemma 8 (Follows from [39], Lemma 6) .
With probability 1 − d − 11 , the following all hold:
For 1 ≤
[Equation]
≤
[Equation]
, n
[Equation]
Additionally,
[Equation]
Finally, for i > n (cid:96) ,
[Equation]
We also use the following classical results throughout the proofs in this section.
Lemma 9 (Weyl’s Inequality) .
For two psd matrices A 1 , A 2 ∈ R p × p ,
[Equation]
for all i ∈ [ p ] .
Deﬁnition 3 ([13]) .
For r ≤ p , let U 1 , U 2 ∈ R p ×
r
both have orthonormal columns.
Then the distance between the subspaces spanned by U 1 , U 2 is
[Equation]
where O r × r is the space of r × r orthogonal matrices.
Lemma 10 (Davis-Kahan sin-
θ theorem [17, 13]) .
For two psd matrices A 1 , A 2 ∈ R p × p , let
A 1 =
U 1 Λ 1 U T 1 and A
= U 2 Λ 2 U T 2 be their eigendecompositions (sorted by decreasing eigenvalues), and let U
=
[ U 1 , ≤ r
U 1 ,>r ] , U 2 =
[ U 2 , ≤ r
U 2 ,>r ] be their eigenvectors, where
U 1 , ≤ r , U 2 , ≤ r ∈ R p × r .
Furthermore, assume (cid:107) A 1 − A 2 (cid:107) op <
(1 − 1 / √ 2)( λ r ( A 1 ) − λ r +
A 1 )) .
Then dist
[Equation] [Equation]
B.2.3 Eigenvector Lemmas
Throughout this section, we condition on the event where Lemmas 6, 7, 8 are all true.
Lemma 11.
Deﬁne
[Equation]
and let ˜ K have eigendecomposition ˜
U ˜ Λ 2 ˜ U T .
Deﬁne Ψ ≤ k ∈ R N × n k to be the ﬁrst n k columns of Ψ ≤ k (where k < (cid:96) ), and also let ˜ U ≤ k be the ﬁrst n k columns of ˜ U , with ˜ U >k the remaining columns.
Then
[Equation]
Proof.
Let Ψ k : (cid:96) ∈ R N × n (cid:96) − n k be the ( n k + 1) -th to n (cid:96) -th columns.
We can then decompose
[Equation]
and
[Equation]
where ˜ U ≤ k ∈ R N × n k .
Let u ∈ R N − n k be a vector such that (cid:107) u (cid:107)
=
and (cid:107) Ψ T ≤ k ˜ U >k u (cid:107)
2 = (cid:107) Ψ T ≤ k ˜ U >k (cid:107) op .
Finally, deﬁne ˜ u
= Ψ T ≤ k ˜ U >k u .
We then have:
[Equation]
since λ min ( Λ 2 ≤ k ) =
Θ( d − k ) .
Deﬁne
[Equation]
and deﬁne ˜ K >k =
˜ K − ˜ K ≤ k .
By Weyl’s inequality,
[Equation]
For i > n k , λ i ( K ≤ k ) = 0 , in which case
[Equation]
Therefore we can upper bound
[Equation]
Altogether, we have
[Equation]
which yields
[Equation]
By deﬁnition,
[Equation]
as desired.
Lemma 12.
Let K ∞ N have eigendecomposition K ∞ N = UΛ 2 U T , where U ≤ k , U >k are the ﬁrst n k , remaining N − n k columns of U respectively.
Then
[Equation]
N
[Equation]
Proof.
By Lemma 8, λ n k ( K ∞ N ) − λ n k +
( K ∞ N ) (cid:38) dN · d − k .
Since (cid:107) K ∞ N − ˜ K
(cid:107) op =
d (cid:107) ∆ (cid:107) op ≤ ˜ O
[Equation]
N d (cid:96) −(cid:33) (cid:28) λ n k ( K ∞ N ) − λ n k +
( K ∞ N ) by Davis-Kahan we have dist
[Equation]
Lemma 13.
Let K N have eigendecomposition K N =
U N Λ 2 N U TN .
Then
[Equation]
Proof.
By Lemma 6, we can write
[Equation]
where ∆ N ≤ ˜ O (cid:16)(cid:113) N ˜ md (cid:17) .
Rearranging, we get
[Equation]
, where we can bound
[Equation]
since (cid:107) K ∞ N (cid:107) op
(cid:46) N .
Since K ∞ N has eigengap Θ( N · d 1 − k ) we can again apply Davis-Kahan to get dist (cid:0) U ≤ k , U N, ≤ k
[Equation]
(cid:46) (cid:107) K
[Equation]
op
[Equation]
Nd
[Equation]
Lemma 14.
Let Σ N have eigendecomposition
Σ N = V N Λ 2 N N V TN , where V N ∈ R
md
× N .
Let ˜ Σ have eigendecomposition ˜ Σ
= VΛ 2 V T .
Let V N
=
[ V N, ≤ k
V N,>k ] , V =
[ V ≤ k
V >k ] where
V N, ≤ k , V ≤ k ∈ R md × n k .
Then, with probability 1 − n − 11 , dist ( V N, ≤ k , V ≤ k
[Equation]
Proof.
By [42, 5.6.4], with probability 1 − n − 11 we have
[Equation]
where r =
Tr ( ˜ Σ ) / (cid:107) ˜ Σ (cid:107) op is the effective rank.
We can upper bound Tr ( ˜ Σ ) ≤ d , and thus
[Equation]
since (cid:107) ˜ Σ (cid:107) op ≤ d .
By Lemma 8,
[Equation]
Therefore by Davis-Kahan, we can bound
[Equation]
The following lemma is a consequence of the preceeding eigenstructure and matrix perturbation lemmas, and partitions the eigenvectors of Σ into groups corresponding to large, medium, and small eigenvalues.
Lemma 15.
Let N ≥ d 4 k , m ≥ N 5 / 2 .
For 1 ≤ k (cid:48) ≤ 2 k , n k (cid:48) − 1 < i ≤ n k (cid:48) , we have λ i ( Σ ) =
Θ( d 1 − k (cid:48) ) .
Also, for i > n 2 k , λ i ( Σ ) (cid:28) λ
[Table]
Proof.
Since
[Equation]
by Weyl’s inequality we have
[Equation]
Since for n k (cid:48) − 1 < i ≤ n k (cid:48)
we have λ
i ( Σ N ) =
N − 1 λ i
( K N ) =
Θ( d 1 − k (cid:48) )
(cid:29) (cid:113) d 2 N , since N ≥ d 4 k ; therefore λ i ( Σ ) =
Θ( d 1 − k (cid:48) ) .
Furthermore, for i > n 2 k ,
[Equation]
As a consequence, we can write the eigendecomposition of Σ as
[Equation]
[Equation] [Equation]

[Equation]
Q
[Equation]
Q
[Equation]


[Equation]
where Q 1 ∈ R md × n k are the large eigenvectors, Q 2 ∈ R md × n 2 k − n k are the medium eigenvectors, and Q 3 ∈ R md × md − n 2 k are the small eigenvectors; concretely,
[Equation]
) and
[Equation]
B.2.4 Proof of Lemma
Proof.
We show this Lemma holds with probability 1 − d − 10 .
Condition on the events where
Lemmas 6, 7, 8, 14 hold.
Pick N so that d (cid:96) (cid:28) N (cid:28) d (cid:96) +1 for (cid:96) = 4 k , and choose ˜ m
=
N 5 / 2 .
We form a dataset of N samples by adding another N − n samples i.i.d from S d − 1 ( √ d ) .
Let y N ∈ R
N be the vector where y N,i =
f k ( x i ) where i ∈ [ N ] .
Recall that Ψ ≤ k ∈ R N × n k denotes the evaluations of the degree ≤ k spherical harmonics on the N data points.
Since f k ( x ) is a degree ≤ k polynomial, we can orthogonally decompose it in terms of the degree ≤ k spherical harmonics, i.e
[Equation]
where (cid:80) k k (cid:48)
=
(cid:80) B ( d,k ) t = w 2 k (cid:48) t
= 1 .
Therefore we have y N = Ψ ≤ k w ∗ for w ∗ ∈ R n k where (cid:107) w ∗ (cid:107)
= 1 .
Observe that Φ N has SVD Φ N =
U N Λ N V TN .
Deﬁne ˜ w ∗ =
˜ U T ≤ k Ψ ≤ k w ∗ .
Then
[Equation]
, and thus
[Equation]
by Lemma 11.
Also,
[Equation]
By Lemmas 12 and13, we have
[Equation]
Thus there exists an orthogonal O 1 such that (cid:13)(cid:13)(cid:13) ˜
U ≤ k − U N, ≤ k O 1 (cid:13)(cid:13)(cid:13) op = dist (cid:16) ˜
U ≤ k , U N, ≤ k (cid:17) and hence
[Equation]
Since U N, ≤ k =
Φ N V N, ≤ k Λ − 1 N, ≤ k , we have (cid:107) y
[Equation]
Φ
[Equation]
(cid:46)
[Equation]
d
Finally, by Lemma 14, we have dist( V N, ≤ k , V ≤ k ) (cid:46) (cid:114) d 2 k
N ≤ d − k , and thus there exists an orthogonal O 2 such that (cid:107) V N, ≤ k − V ≤ k O 2 (cid:107) op ≤ d − k and hence (cid:107) y N − Φ N V ≤ k O 2 Λ − 1 N, ≤ k O 1 ˜ w ∗ (cid:107) ≤ (cid:107) y N
[Equation]
Φ N V N, ≤ k Λ − 1 N, ≤ k O 1 ˜
[Equation]
Therefore, letting v ∗ =
V ≤ k O 2 Λ − 1 N, ≤ k O 1 ˜ w ∗
[Equation]
Note that v ∗ satisﬁes
[Equation]
Since the { x i } i ∈ [ N ] are i.i.d, we can treat the dataset D as a subsample of n data points.
Since | f k ( x ) | 2 ≤ n k = Θ( d k ) , and | ˜ ϕ
( x ) T v ∗ | 2 ≤ (cid:107) ˜ ϕ
( x ) (cid:107) 2 (cid:107) v ∗ (cid:107) 2 =
Θ( d k ) , we can bound ( f k ( x ) − ˜
ϕ
( x ) T v ∗ ) 2 ≤ d k .
Also,
[Equation]
Therefore by Bernstein’s Inequality
[Equation]
with probability 1 − d −
[Equation]
we can bound
[Equation]
To conclude we must relate ˜ ϕ to ϕ .
Observe that
[Equation]
Therefore λ i ( Σ ) =
λ i ( ˜ Σ ) for i ≤ ˜ m , and λ i ( Σ ) = 0 otherwise.
Furthermore, if v is an eigenvector of ˜ Σ , then (cid:20) v − v (cid:21) is an eigenvector of Σ .
As a result, if u ∈ span ( { v i ( ˜ Σ ) } i ∈ [ n k ] ) , then (cid:20) u − u (cid:21) ∈ span ( P ≤ k ) .
Letting z ∗ =(cid:20) v ∗ − v ∗ (cid:21) , we get that z ∗ ∈ span ( P ≤ k ) , and also ϕ ( x ) T z ∗ =
˜ ϕ
( x ) T v ∗ .
Therefore
E
[Equation]
and also (cid:107) z ∗ (cid:107)
(cid:46)
d k − 1 .
Union bounding over all high probability events, this holds with probability 1 − 4 d − 11 ≥ 1 − d − 10 , as desired.
B.3 Proof of Theorem
Proof.
We condition on the events of Lemmas 1, 2, 16, 17, 18 holding, which occurs with probability 1 − 5 d − 10 ≥ 1 − d − 9 .
We proceed by the probabilistic method.
For r ∈ [ m ] , let the σ r be random variables with σ r ∼
Unif ( {± 1 } ) i.i.d, and let S = diag ( σ 1 , . . .
, σ m ) be the diagonal matrix of random signs.
Let
W ∗ S =
W L + W Q S be a (random) solution.
We will show that in expectation over S the training error and all the regularizers are small, which implies the existence of a S which makes all these quantities small.
Bounding the Empirical Loss.
First, observe that we can write
[Equation] [Equation]
Since f Q ( x ; W Q ) =
f Q ( x ; W Q S )
[Equation]
the second term can be deterministically bounded as
[Equation]
We next bound the ﬁrst term in expectation using Lemma 19:
[Equation]
Since the loss is Lipschitz, we can bound the empirical loss as
[Equation]
where we used the bounds in Lemmas 1, 2, along with the lower bounds on m in the assumption of the theorem.
Bounding the Regularizers.
We begin with R 1 :
[Equation]
since our construction guarantees that P >k W L = 0 .
By Lemma
[Equation]
Therefore
[Equation]
R 2 can be bounded as
[Equation]
W
[Equation]
By Lemma 20,
[Equation]
where the last step uses Lemma 16.
Therefore, again applying Lemma 19,
[Equation]
R 3 can be bounded as
[Equation]
since our construction guarantees P >k W L = 0 .
By Lemma 19,
[Equation]
By Lemma 20,
[Equation]
which can be upper bounded by 2 (cid:107) f L ( x ; W Q S ) (cid:107) 2 L 2 .
Therefore
[Equation]
Altogether,
[Equation]
Finally, to bound R 4 , observe that
[Equation]
By the construction (cid:107) W Q (cid:107) 2 , 4 ≤ d k − 14 .
Also, since W L ∈ span ( P ≤ k ) , by Lemma 21 we have
[Equation]
Therefore (cid:107) W ∗ S (cid:107) 2 , 4 (cid:46) d
k − 14 , and thus
[Equation]
Corollary 3.
The solution W ∗ satisﬁes
[Equation]
Proof.
We have
[Equation]
By construction (cid:107) W ∗ L (cid:107) F (cid:46) d k − 12 , and also
[Equation]
The desired claim follows since we assume m (cid:29) d k − 1 .
Corollary 4.
[Equation]
Proof.
We can writte
[Equation]
By Corollary 2, (cid:107) W ∗ Q (cid:107) 2 , ∞ (cid:46) m − 1 / 4 d k − 12 .
Also, since W ∗ L ∈ span ( P ≤ k ) , we can apply Lemmato get (cid:107) W ∗ L (cid:107) 2 , ∞ ≤ m − 12 d k 2 (cid:107) W ∗ L (cid:107) F ≤ m − 12 d k − 12 .
Altogether, since m ≥ d 2 k , we get (cid:107) W ∗
[Equation] [Equation]
∞
[Equation]
B.4 Expressivity Helper Lemmas
Lemma 16.
With probability 1 − d − 10 ,
E n [ f k ( x ) 2 ] (cid:46) 1 .
Proof.
Since (cid:107) f k (cid:107) L 2 = 1 and f is degree k , | f k ( x ) | 2 ≤ n k for all x .
Furthermore,
E µ [( | f k ( x ) | 2 − 1) 2 ] ≤ E µ [ f k ( x ) 4 ] ≤ n k .
Therefore by Bernstein’s Inequality, with probability 1 − d − 10 we have
[Equation]
since
[Table]
Lemma 17 ([42, Exercise 4.7.3]) .
With probability 1 − d − 10 ,
[Equation]
Lemma 18.
Recall
[Equation]
With probability 1 − d − 10 ,
[Equation]
Proof.
Observe that
[Equation]
Therefore by [42, 5.6.4]
[Equation]
with probability 1 − d −
[Equation]
since n (cid:38) d k log d .
Lemma 19.
On the event where Lemma 17 holds, for all W we have
[Equation]
Furthermore, we have
[Equation]
Proof.
We have
[Equation]
The proof in the population case is identical
[Equation]
Lemma 20.
On the event where Lemma 18 holds,
[Equation]
for all W .
Proof.
We can write
[Equation]
Therefore
[Equation]
by Lemma
[Equation]
as desired.
Lemma 21.
For any W ∈ span ( P ≤ k ) , we have
[Equation]
and
[Equation]
Proof.
For ( r, s ) ∈ [ m ] × [ d ] , let e ( r,s ) denote the (( d − 1) r + s ) th canonical basis vector in R md , so that e T ( r,s ) vec ( W ) = { w r } s .
Let c 1 , . .
c n k be scalars such that (cid:80) n k i =
1 c 2 i = (cid:107) W (cid:107) 2 F and vec ( W
[Equation]
= n k (cid:88) i =
c i v i .
By Cauchy, we can bound
[Equation]
Observe that
[Equation]
Furthermore, since e T ( r,s ) ϕ ( x ) = 1 √ m σ (cid:48) ( w T 0 ,r x ) x s
[Equation]
e
[Equation]
r
[Equation]
s
[Equation]
(cid:107) 2Σ =
E µ (cid:104) ( e T ( r,s ) ϕ (
[Equation]
)) 2 (cid:105) ≤
[Equation]
µ [ x 2 s ] ≤ 1 m.
By Lemma 15 λ n k = Θ( d − k +1 ) , and thus we can bound
[Equation]
Thus every row w r satisﬁes
[Equation]
so
[Equation]
C Landscape Proofs C.1 Coupling Lemmas
Recall that ˆ L Q ( W ) =
E n [ (cid:96) ( y, f L ( x ; W ) + f Q ( x ; W ))] denotes the empirical loss of the quadratic model.
As in [7], we begin by showing the losses, gradients, and Hessians for ˆ L Q ( W ) and ˆ L ( W ) are close for W satisfying a norm bound.
This is given by the following 3 coupling lemmas.
Lemma 22 (Coupling of Losses) .
[Equation]
Lemma 23 (Coupling of Gradients) .
[Equation]
Lemma 24 (Coupling of Hessians) .
[Equation]
While these Lemmas are similar in structure to those in [7], extra care must be taken to properly deal with the effect of the f L terms.
These proofs are presented in Appendix C.1.2.
C.1.1 Auxiliary Results
We ﬁrst prove some intermediate results which are used in the proofs of the coupling lemmas:
Lemma 25 (Coupling of Function Values) .
[Equation]
Proof.
[Equation]
Lemma 26 (Coupling of Function Gradients)
[Equation]
Proof.
Taylor expanding σ (cid:48) , we have
[Equation]
Therefore
[Equation]
Lemma 27.
[Equation]
[Figure]
[Equation]
Proof.
We can decompose
[Equation]
We can bound
[Equation]
So max i
[Equation] [Equation]
C.1.2 Proof of Coupling Lemmas
Proof of Lemma 22.
By Lipschitzness of the loss,
[Equation]
Proof of Lemma 23.
The gradients are
[Equation]
Therefore
[Equation]
(cid:105)
[Equation]
(cid:104)∇
[Equation] [Equation]
W
[Equation]
˜
[Equation]
(cid:105) (cid:12)(cid:12)(cid:12)(cid:105) +
[Equation]
(cid:12)(cid:12)(cid:12) (cid:104)∇
[Equation]
[Equation]
(cid:105) (cid:12)(cid:12)(cid:12)(cid:105)
[Equation]
(cid:12)(cid:12)(cid:12)
[Equation]
(cid:105) (cid:12)(cid:12)(cid:12)(cid:105)
[Equation]
by Lemma 26.
The ﬁrst term is
[Equation]
The second term can be bounded as
[Equation]
by Lemma
[Equation]
Proof of Lemma 24.
The Hessians are
[Equation]
and
[Equation]
The difference between the ﬁrst terms can be bounded by
[Equation]
[Table]
[Equation]
The difference between the second terms is upper bounded by
[Equation]
The ﬁrst term can be bounded by
[Equation]
Also, by Lemma 26 we can bound
[Table]
C.2 Proof of Lemma
Proof.
We would ﬁrst like to show that the quadratic model has good landscape properties.
To prove this, we begin by showing any approximate stationary point must be “localized," in that the values of the regularizers at these stationary points must not be too large.
Lemma 28 (Localization) .
Let λ 2 = ε min , λ
=
m 12 d − k − 12 ε min , λ
4 = d − 2( k − 1) ε min .
Assume
[Equation]
min
[Equation]
n
[Equation]
4 min (cid:17) , and ν ≤ m − 14 .
Then, for any ν -ﬁrst-order stationary point W of L λ , we have
[Equation]
Next, we show that for these localized points, the landscape of ˆ L Q is “good.”
Lemma 29.
Let W ∗ L = P
[Equation]
k W ∗ , W L = P ≤ k W .
There exists a universal constant C such that
[Equation]
As a corollary, we obtain that for localized W , this error term can be made arbitrarily small for sufﬁciently large m .
Corollary 5.
Assume m (cid:38) d 14 k + ε − 22 / 3 min .
Then, additionally under the assumptions of Lemma 3, any ν -ﬁrst-order stationary point W satisﬁes
[Equation]
Our ﬁnal step is to show that the landscape of L λ is good.
To do this, we ﬁrst show that the landscapes of the quadratic model and original model couple for localized points.
This coupling is given by the following lemma.
Lemma 30.
Assume the conditions of Lemma 3, Corollary 5, and let W be a ν -ﬁrst order stationary point.
Furthermore, assume m (cid:38) n 4 d 26( k +1)
ε − 22 / 3 min .
Then,
[Equation]
An immediate corollary is that the landscape of ˆ L must be good.
Corollary 6.
Let W be a ν -ﬁrst-order stationary point of L λ .
Under the same conditions of Lemma 3,
Corollary
[Equation]
Lemma
[Equation]
we have
[Equation]
To conclude, we must show that adding the regularizers has a benign effect on the landscape
Lemma 31.
Deﬁne
[Equation]
to be the total regularization term.
Under the conditions of Lemma 3, we have
[Equation]
Lemma 3 now directly follows by adding the results of Corollary 6 and Lemma 31.
C.3 Proof of Corollary
Proof.
Let W be an ( ν, γ ) -SOSP of L λ ( W ) .
Then
[Equation]
since we have chosen ν ≤ m − 1 / 2 , m ≥ d 4 k − 2 3 ε − 14 / 3 .
Also,
[Equation]
since we have chosen γ ≤ m − 3 / 4 , m ≥ d 2( k − 1) ε − 4 .
Altogether, we can bound
L λ ( W ) (cid:46)
L λ ( W ∗ )
+ (cid:104)∇ L λ ( W ) , W − 2 W ∗ L + W L (cid:105) − E S [ ∇ 2 L λ ( W )[ W ∗ S , W ∗ S ]] + ε min (cid:46) ε min , as desired.
C.4 Proofs of Intermediate Results C.4.1 Proof of Lemma
Proof.
Let W be an ν -ﬁrst-order stationary point of L λ .
Then
[Equation]
We have that
[Equation]
First, by convexity we can bound
[Equation]
Secondly, we can bound
[Equation]
Finally
[Equation]
by Lemma
[Equation]
Altogether,
[Equation]
We next turn to the regularizers.
R i for i = 1 , 2 , 3 are all quadratics, so (cid:104)∇R i ( W ) , W (cid:105)
= 2 R i ( W ) .
Also it is true that (cid:104)∇R 4 ( W ) , W (cid:105)
=
8 R 4 ( W ) .
Plugging into (53), we get ν
[Equation]
Therefore (using ν ≤ m − 1 / 4 ),
[Equation]
F
[Equation]
where the last step follows from Lemma
[Equation]
We can bound
[Equation]
Therefore by AM-GM, λ
[Equation]
where the last step uses
[Equation]
W
[Equation]
Since m − 14 n ≤ ε min / 2 ≤ λ 2 / 2 , λ 3 / 2 , we have
[Equation]
Therefore, plugging in λ 4 = d
[Equation] [Equation]
[Equation]
Since ε min < 1 we trivially have d k − 14 ε − 1 / 8 min <
d 2 k − 1 6 ε − 1 / 6 min .
Also, since m ≥ d 4 k +4 n 2 ε − 2 min , we have
[Equation]
Additionally, assuming m ≥ d 16 / 3( k +1) ε − 8 / 3 min , we have
[Equation]
Therefore we can bound
[Equation]
and thus
[Equation]
In this case, the RHS of (57) can be upper bounded by d (cid:107) W (cid:107) 22 , 4 .
Plugging back into (56), we get
[Equation]
which yields the bounds
[Equation]
C.4.2 Proof of Lemma
Proof.
Observe that
[Equation]
and therefore for any diagonal matrix of random signs S = diag ( σ 1 , . . .
, σ m ) , ∇ ˆ L Q ( W )[ W ∗ S , W ∗ S ]
[Equation] [Equation]
[ (cid:96)
[Equation]
( y,
[Equation]
L ( x ; W ) + f Q ( x ; W )
[Equation]
f Q ( x ; W ∗ )]
[Equation]

[Equation]
[Table]
The expectation of the second term over the random signs S can be upper bounded by
[Equation]
Therefore
[Equation] [Equation]
Next, deﬁne ∆ =
W − 2 W ∗ L + W L .
We have
[Equation]
ˆ
[Equation]
Q
[Equation]
W ) , ∆
[Equation]
=
E n
[Equation] [Equation]
The second term can be bounded in magnitude as (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
Also, we have that | E n [
[Equation]
, f
[Equation] [Equation]
[Equation]
; W ))
[Equation]
W
[Equation] [Equation]
) |
[Equation] [Equation]
W
[Equation] [Equation]
and similarly | E n [ (cid:96) (cid:48) ( y, f L ( x ; W ) + f Q ( x ; W ))
( f L ( x ; W ∗ ) −
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
Altogether, we have for some constant C
[Equation]
> 0 ,
[Equation] [Equation]
Finally, by convexity of (cid:96) we have
[Equation]
Combining (63), (65)
[Equation]
and (67), we get that
[Equation]
as desired.
C.4.3 Proof of Corollary
Proof.
First, note that we have the bounds
[Equation]
Also, by Lemma
[Equation]
for ν -ﬁrst order stationary point W we can bound
[Equation]
Furthermore, since vec ( W L ) ∈ span ( P ≤ k )
[Equation]
we can write
[Equation]
Therefore
[Equation]
Altogether
[Equation]
by Lemma
[Equation]
we can bound
[Equation]
· d 5 k − 1 6 ε − 2 / 3 min + m − 14 d 7 k +112 ε − 2 / 3 min
[Equation]
≤
[Equation]
, since we have assumed m (cid:38) d 14 k + ε − 22 / 3 min .
C.4.4 Proof of Lemma
Loss Coupling.
By Lemma 22, we can bound
[Equation]
for m ≥ d 4 k +4 ε − 6 min .
Similarly,
[Equation]
since m ≥ d 3( k +1) ε − 4 min .
Gradient Coupling.
Next, by Lemma 23, we have
[Equation]
where ∆ := W − 2 W ∗ L + W L .
First, observe that
[Equation]
where we applied Lemma 21.
Plugging in the bounds for (cid:107) W (cid:107) 2 , 4 , (cid:107) W L (cid:107) 2 , 4 , we get (cid:107) ∆ (cid:107)
[Equation]
since m ≥ d 4 k ε −
[Equation]
Also, by Lemma 27,
[Equation]
Observe that
[Equation]
where we ﬁrst decomposed W = P >k W + W L and then used the deﬁnition of R 3 .
Since
W L ∈ span ( P ≤ k ) , by Lemma 20 we have
[Equation]
and similarly
[Equation]
( f L ( x ; W ∗ L )) 2 (cid:3) (cid:46) E n (cid:2) ( f k ( x ) 2 (cid:3) + E n (cid:2) ( f L ( x ; W ∗ L ) − f k ( x )) 2 (cid:3) (cid:46) 1 .
Altogether
E
[Equation]
Plugging this back in,
[Equation]
since d − k (cid:28) ε min .
Therefore (cid:12)(cid:12) (cid:12) (cid:104)∇ ˆ L ( W ) , W −
[Equation]
W ∗
[Equation]
+ W L (cid:105) −
[Equation]
ˆ
[Equation]
W )
[Equation]
W
[Equation]
W
[Equation]
W
[Equation]
(cid:46) d
[Equation]
· d k − 12 ε
[Equation] [Equation]
min · n 12 d 2 k +23 ε − 1 /
[Equation]
min = m − 1 / 4 n 12 d
[Equation]
+
[Equation]
≤
[Equation]
since m ≥ n 2 d 20( k +
[Equation] [Equation]
− 22 / 3 min .
Hessian Coupling.
By Lemma 24, we have
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
By Lemma
[Equation]
[Table]
[Equation]
where we used Lemma 19 to bound E S E n (cid:2) ( f L ( x ; W ∗ S )) 2 (cid:3) , and then plugged in the bound for (cid:107) W (cid:107) 2 , 4 from Lemma 28 and the bound for (cid:107) W ∗ (cid:107) 2 , 4 from Theorem 2.
By Corollary 4, we can bound (cid:107) W ∗ (cid:107) 2 , ∞ (cid:46) m
− 14 d k − 12 .
Plugging these two bounds in, along with the bound for (cid:107) W (cid:107) 2 , 4 , we get
[Equation]
since we’ve assumed m ≥ n 4 d 26( k +
[Equation]
ε − 22 / 3 min .
C.4.5 Proof of Lemma
Proof.
Let us ﬁrst consider the regularizer R 4 ( W ) .
From the proof of [7, Corollary 3], we have
E S ∇ 2 R 4 ( W )[ W ∗ S , W ∗ S ] − (cid:104)∇R 4 ( W ) , W (cid:105) + 2 R 4 ( W ) − 2 R 4 ( W ∗ ) ≤ −R 4 ( W ) + C R 4 ( W ∗ ) for an absolute constant C .
Thus
E
[Equation]
By Lemma
[Equation]
we can bound
[Equation]
Plugging this in, along with the bound for (cid:107) W (cid:107) 2 , 4 , yields
[Equation]
and thus
[Equation]
where we used λ 4 R 4 ( W ∗ ) ≤ ε min and λ
4 = d − 2( k − 1) ε min , and then used the assumption m ≥ d (20 k +
[Equation]
/ 3 ε − 22 / 3 min .
We next deal with the other 3 regularizers.
Observe that we can write
[Equation]
tot
[Equation]
for some psd
A ∈ R md × md .
We get that
[Equation]
Since W L , W ∗ L ∈ span ( P ≤ k ) , we get that
[Equation]
, and
[Equation]
Therefore
[Equation]
Also,
[Equation]
Therefore
[Equation]
Summing this with the effect of R 4 on the landscape
[Equation]
we obtain
[Equation]
as desired
[Equation]
D Optimization Proofs D.1 Geometric Properties
In this section we show that the regularized loss L λ is (cid:96) -smooth and ρ -Hessian-Lipschitz inside a norm ball.
Lemma 32 (Loss Hessians are Lipschitz) .
[Equation]
Proof.
Recall that
[Equation]
To bound the ﬁrst term
[Equation]
since (cid:96) (cid:48)
[Equation]
σ (cid:48)(cid:48) are both Lipschitz and can be upper bounded by
[Equation]
we get
[Equation]
+ 1 √ m m
[Equation]
=
[Equation]
w
[Equation] [Equation]
w
[Equation]
w
[Equation]
x )
[Equation]
≤ (cid:12)(cid:12) ( w
[Equation]
r
[Equation]
w
[Equation]
r )
[Equation]
(cid:12)(cid:12)
[Equation]
m (cid:88) s =
[Equation]
(cid:12)(cid:12)
[Equation]
w
[Equation]
s
[Equation] [Equation]
s
[Equation]
(cid:12)(cid:12)
Therefore
[Equation]
m
[Equation]
To bound the second term
[Equation]
we have
[Equation]
Since σ (cid:48) is bounded by
[Equation]
we can bound (cid:107)∇ W f ( x ; W 2 ) (cid:107) ≤ √ d .
Next, we have
[Equation]
Finally
[Equation]
Altogether,
[Equation]
Therefore
[Equation]
so
[Equation]
Lemma 33 (Regularized loss is Hessian-Lipschitz) .
The regularized loss L λ is O ( λ 4 Γ 5 ) -HessianLipschitz inside the region { W | (cid:107) W (cid:107) F ≤ Γ }
Proof.
For i = 1 , 2 , 3 , the regularizer R i is a convex quadratic, so ∇ 2 R
i ( W 1 ) = ∇ 2 R
i ( W 2 ) .
As for the regularizer R 4 , we have ∇
[Equation]
R 4 ( W )[ ˜ W , ˜ W ] =
[Equation]
Therefore, for (cid:107) ˜ W (cid:107) F = 1 , (cid:12)(cid:12)(cid:12) ∇ 2 R 4 ( W 1 )[ ˜ W , ˜ W ] − ∇ 2 R 4 ( W 2 )
[ ˜ W , ˜ W ] (cid:12) (cid:12)(cid:12) ≤(cid:32) m (cid:88) r
= ( (cid:107)
w 1 ,r (cid:107) 2 w T 1 ,r
+ (cid:107)
[Equation]
We can bound the ﬁrst term using
[Equation]
and
[Equation]
For the second term, we bound
[Equation]
[Table]
F .
Finally, we bound m (cid:88) r =
[Equation]
≤
[Equation]
≤
[Equation]
r =
[Equation]
(cid:46) (cid:107) W 1 (cid:107) 2 F (cid:107) W 1 − W 2 (cid:107) 2 , 4 ( (cid:107) W 1 (cid:107) 32 , 4 + (cid:107) W 2 (cid:107) 32 , 4 ) + (cid:107) W 2 (cid:107) 42 , 4 ( (cid:107) W 1 (cid:107) F + (cid:107) W 2 (cid:107) F ) (cid:107) W 1 − W 2 (cid:107) F ≤ (cid:0) (cid:107) W 1 (cid:107) 2 F ( (cid:107) W 1 (cid:107) 32 , 4 + (cid:107) W 2 (cid:107) 32 , 4 ) + (cid:107) W 2 (cid:107) 42 , 4 ( (cid:107) W 1 (cid:107) F + (cid:107) W 2 (cid:107) F ) (cid:1) (cid:107) W 1 − W 2 (cid:107) F .
Altogether, when (cid:107) W 1 (cid:107) F , (cid:107) W 2 (cid:107) F ≤ Γ , and using (cid:107) W (cid:107) 2 , 2 k ≤ (cid:107) W (cid:107) F for k ≥ 1 , we get that (cid:12)(cid:12)(cid:12) ∇ 2 R 4 ( W 1 )
[ ˜ W , ˜ W ] − ∇ 2 R 4 ( W 2 )
[ ˜ W , ˜ W ] (cid:12)(cid:12)(cid:12) (cid:46) Γ 5 (cid:107) W 1 − W F (cid:107) .
Therefore L λ is O
(Γ 5 ) -Hessian-Lipschitz.
Lemma 34 (Regularized loss is smooth) .
The regularized loss L λ is O ( λ 4 Γ 6 + m 1 / 2 ) -smooth.
Proof.
Recall that
[Equation]
and thus, for (cid:107) ˜ W (cid:107) F = 1 ,
[Equation]
Therefore ˆ L is O ( d ) -smooth.
Next, consider the regularizers.
R 1 , R 2 , R 3 are all convex quadratics, and since (cid:107) ϕ
( x ) (cid:107) ≤ √ d , we can upper bound the smoothness of R 1 , R 2 , R 3 by d .
Therefore the smoothness of λ 1 R 1 + λ 2 R 2 + λ 3 R 3 is at most m 1 / 2 d − k − 12 ε min · d (cid:46) m 1 / 2 .
Finally, we consider R 4 .
For (cid:107) ˜ W (cid:107) F = 1 , we can bound ∇ 2 R 4 (
[Equation]
=
[Equation] [Equation]
[Equation]
24 (cid:107) W (cid:107)
[Equation]
m (cid:88) r = (cid:107) w
[Equation]
(cid:107)
[Equation]
(cid:107)
[Table]
w
[Equation]
(cid:107)≤ 32 (cid:107) W
[Equation] [Equation]
(cid:107)
[Table]
W
[Equation]
F
[Equation]
24 (cid:107) W
[Equation]
(cid:107) W (cid:107)
[Equation]
F ≤ 32 (cid:107) W (cid:107)
[Equation]
24 (cid:107) W (cid:107)
[Equation]
(cid:107) W (cid:107)
[Equation]
F ≤ 56 (cid:107) W (cid:107) 6 F .
Therefore L λ is O ( λ 4 Γ 6 + m 1 / 2 ) -smooth.
D.2 Proof of Theorem
We prove the following formal version of Theorem 3.
Theorem 5.
For ν, γ , choose η = c
λ 4 m 3 for sufﬁciently small constant c and deﬁne ˜ ε
= min( ν, γ 2 m − 5 / 2 ) , σ = ˜Θ(˜ ε ) .
Then with probability 1 − d − 8 , perturbed gradient descent reaches a ( ν, γ ) -SOSP within T =
˜ O ( m 3 ˜ ε − 2 ) timesteps.
Proof.
We follow the same strategy as [30, Theorem 8].
We ﬁrst show that perturbed gradient descent stays in a bounded region.
Then, we can use the smoothness and Hessian-Lipschitz parameters in this region for the purpose of a convergence result.
Throughout, we condition on the high probability event where the construction in Theorem 2 holds.
Let Γ =
m
I claim that perturbed gradient descent stays in the region {(cid:107) W (cid:107) | (cid:107) W (cid:107) ≤ Γ } .
We prove the claim by induction.
(cid:107) W 0 (cid:107) =
so clearly the base case holds.
Assume that (cid:107) W t (cid:107) ≤ Γ .
The gradient descent update on W is
W
[Equation]
←
[Equation]
ˆ
[Equation]
W
[Equation]
W
[Equation]
Observe that the learning rate is η = c
λ 4 m 3 =
c λ 4 Γ 6 for a sufﬁciently small constant c .
We can bound (cid:107)∇ ˆ L ( W t ) (cid:107) ≤ √ d and, as in the proof of the previous lemma,
[Equation]
Therefore
[Equation]
Finally, since the w r component of ∇R 4 ( W ) is 8 (cid:107) W
[Equation] [Equation]
w r
[Equation]
2 w r , we have
[Equation]
Note that
[Equation]
for c small enough.
We then have
[Equation]
We split into two cases.
If (cid:107) W t (cid:107) F ≤ Γ / 2 , then
[Equation]
[Equation]
≤ Γ /and (cid:13)(cid:13)(cid:13) η ∇ ˆ L ( W t ) + η ∇ (cid:0) λ
[Equation]
W
[Equation]
) + λ
[Equation]
2 ( W
[Equation]
) + λ 3 R 3 ( W t ) (cid:1)(cid:13)(cid:13)(cid:13) (cid:46) d 2( k − 1) ε − 1 min Γ − 5 m 1 /≤ Γ / 4 , so (cid:107) W t +1 (cid:107) F ≤ 3Γ / 4 .
Otherwise, Γ ≥ (cid:107) W t (cid:107) F ≥ Γ / 2 , so
[Equation]
Then
[Equation]
since Γ 6 = m 3 (cid:38) m 5 / 2 d 2( k − 1) ε − 1 min .
Finally, the perturbation moves at most η (cid:107) Ξ t (cid:107) F , and since E (cid:107) Ξ t (cid:107) F = σ , with probability 1 − d −each of the T perturbations is bounded by Γ m − 3 η − 1 (cid:29) σ
Therefore even after the perturbation we have (cid:107) W t +1 (cid:107) F ≤ Γ , completing the induction step.
Lemmas 33, 34 tell us L λ is O ( λ 4 m 5 / 2 ) Hessian Lipschitz and O ( λ 4 m 3 ) smooth throughout the entire gradient descent trajectory.
Our goal is to converge to a ( ν, γ ) -SOSP; this is equivalent to converging to a ˜ ε -SOSP as deﬁned in [30, 31], with ˜ ε :=
min( ν, γ 2 λ 4 m 5 / 2 ) ≥ min( ν, γ 2 m − 5 / 2 ) .
Therefore by [30, 31], with probability 1 − d − 9 perturbed gradient descent on the regularized loss with learning rate η = c
λ 4 m 3 and perturbation radius σ = ˜Θ(˜ ε ) will encounter an ˜ ε -SOSP in at most T =
˜ O (cid:0) λ
4 m 3 ˜ ε − 2 (cid:1) ≤ ˜ O
(cid:0) m 3 ˜ ε − 2 (cid:1) timesteps.
Union bounding over the high probability events, this occurs with probability 1 − 3 d − 9 ≥ 1 − d − 8 .
E Generalization Proofs E.1 Proof of Theorem
Recall the deﬁnition of the empirical Rademacher complexity:
Deﬁnition 4.
Let F be a class of functions from R d to R .
Given a dataset D = { x 1 , . . .
, x n } , the empirical Rademacher complexity of F is deﬁned as
[Equation]
We next show that the Rademacher complexities of the linear term and the quadratic term can be bounded.
Lemma 35 (Rademacher complexity of linear term) .
Let W ⊂ R m × d , and deﬁne the function class
F
[Equation]
P
[Equation]
: W ∈ W} .
Then, with probability 1 − d − 9 over the draw of D ,
[Equation]
Lemma 36 (Rademacher complexity of quadratic term) .
Let W ⊂ R m × d , and deﬁne the function class F Q ( W ) :
= { x (cid:55)→ f Q ( x ; W ) : W ∈ W} .
Then, with probability 1 − d − 9 over the draw of D ,
[Equation]
[Equation]
Lemma 35 is presented in Appendix E.2; Lemma 36 follows directly from [7, Lemma 5, Theorem 6].
Equipped with these Rademacher complexity lemmas, we can now prove the main generalization result.
Proof of Theorem 4.
By Lipschitzness of the loss and Lemma 25, we can bound
L ( W ) =
E µ [ (cid:96) ( y, f ( x ; W ))] ≤ E µ [ (cid:96) ( y, f Q ( x ; W ) + f L ( x ; W ))]
+ E µ | f ( x ; W ) − f Q ( x ; W ) − f L ( x ; W ) | ≤ E µ [ (cid:96) ( y, f Q ( x ; W ) + f L ( x ; W ))]
+ m −m (cid:88) r =
E µ | w Tr x |≤ E µ [ (cid:96) ( y, f Q ( x ; W ) + f L ( x ; W ))]
+ Cm − 14 (cid:107) W (cid:107) 32 , 4 .
Again by Lipschitzness, we can bound
E µ [ (cid:96) ( y, f Q ( x ; W ) + f L ( x ; W ))] ≤ E µ [ (cid:96) ( y, f Q ( x ; W ) + f L ( x ; P ≤ k W ))]
+ E µ | f L ( x ; P >k W ) | ≤ E µ [ (cid:96) ( y, f Q ( x ; W ) + f L ( x ; P ≤ k W ))]
+ (cid:107) f L ( x ; P >k W ) (cid:107) L 2 .
Similarly, we can lower bound ˆ L ( W ) ≥ E n [ (cid:96) ( y, f Q ( x ; ˆ W ) + f L ( x ; P ≤ k ˆ W ))] − Cm − 14 (cid:107) W (cid:107) 32 , 4 − (cid:0) E n [( f L ( x ; P >k )) 2 ] (cid:1) 12 .
Since L λ ( ˆ W ) ≤ Cε min , the value of each regularizer satisﬁes R i ( ˆ W ) ≤ Cλ − 1 i ε min .
By our choice of ( λ 1 , λ 2 , λ 3 , λ 4 ) , we have
[Equation]
L 2 (cid:46) m − 12 d k − 12 ,
R 2 ( ˆ W ) =
(cid:107) f L ( · ; P ≤ k W ) (cid:107) 2 L 2 (cid:46) 1 ,
R 3 ( ˆ W ) =
E n (cid:2) ( f L ( x ; P >k W )) 2 (cid:3) (cid:46) m − 12 d k − 12 ,
R 4 ( ˆ W ) =
(cid:107) W (cid:107) 82 , 4 (cid:46)
d 2( k − 1) .
Therefore
[Equation]
and similarly
[Equation]
Since ˆ L ( ˆ W ) ≤ Cε min , we thus have that
[Equation]
Next, deﬁne the set
[Equation]
By construction, we have ˆ W ∈ W .
Furthermore, deﬁne the function class
[Equation]
By the standard empirical Rademacher complexity bound
[Equation]
with probability 1 − d − 9 over the draw of
D
[Equation]
E µ [ (cid:96) ( y, f Q ( x ; ˆ W ) + f L ( x ; P ≤ k ˆ W ))] − E n [ (cid:96) ( y, f Q ( x ; ˆ W ) + f L ( x ; P ≤ k ˆ W ))]] ≤ 2 R D ( L ) ,
By the Rademacher contraction Lemma [43], since (cid:96) is 1-Lipschitz we can bound (conditioning on
Lemmas 35, 36)
[Equation]
n
[Equation]
Union bounding over the high probability events, with probability 1 − 3 d − 9 ≥ 1 − d − 8 , we have
[Equation]
n
[Equation]
since m (cid:38) d 3( k − 1) ε − 4 min .
E.2 Proof of Lemma
Proof.
Recall that we can write f L ( x ; W ) =
ϕ ( x ) T vec ( W ) , where ϕ ( x ) is the NTK featurization map.
Also, recall
[Equation]
We can then bound
[Equation]
where the last step follows by Cauchy-Schwarz.
By deﬁnition,
[Equation]
Also, we can bound
E σ
[Equation]
By Lemma 18, with probability 1 − d − 9 we have
[Equation]
Furthermore, E n (cid:104) ( Σ †≤ n k ) 1 / 2 P ≤ k ϕ ( x ) ϕ ( x ) T P ≤ k ( Σ †≤ n k ) 1 / 2 (cid:105) and P ≤ k have the same span
[Equation]
which is dimension
n k = Θ
( d k ) .
Therefore
[Equation]
Altogether, we get the bound
E
[Equation]
so
[Equation]
F Proof of Theorem
Proof.
Choose m (cid:38) n 4 d 26( k +1)
ε − 22 / 3 .
Set the regularization parameters as λ= m 1 / 2 d − k − 12 ε min , λ 2 = ε min , λ
=
m 1 / 2 d − k − 12 ε min η, λ 4 =
d − 2( k − 1) ε min .
Also, set r = n k , η = cλ − 1 4 m − 3 for small constant c , and σ = ˜ O ( m − 4 ) .
For ν = m − 1 / 2 , γ = m − 3 / 4 , ˜ ε
= min( ν, γ 2 m − 5 / 2 ) = m − 4 .
Therefore by Theorem 5, with probability 1 − d − 8 we reach a ( ν, γ ) -SOSP within T =
˜ O ( m 3 ε − 2 ) =
˜ O ( m 11 ) timesteps.
Call this point ˆ W .
By Corollary 1, L λ ( ˆ W ) ≤ Cε min .
Finally, by Theorem 4, with probability 1 − d − 8 , L ( ˆ W ) ≤ Cε min .
Setting ε min = ε/C , we get that L ( ˆ W ) ≤ (cid:15) with probability 1 − d − 7 , as desired.
G Additional Experiments G.1 Additional Simulations
[Figure]
[Figure] [Figure]
1 =0.
[Figure] [Figure]
1 =0.
2.0 log d t 0.
T r a i n
L o ss 0.
2.0 log d t 0.
T e s t L o ss
In Figure 3, we conduct the same experiment as in Figure 1, while additionally using the Rregularizer.
We ﬁx λ 3 = 0 .
01 while varying λ 1 .
Since we cannot compute R 1 exactly, we use an unbiased estimate at every timestep by sampling a new set of x ’s and computing E [ f L ( x ; P >k W )] on this set.
In the leftmost pane, we plot a moving average of our estimate of R 1 .
First, we observe that even when λ 1 = 0 , the regularizer R 1 is an order of magnitude smaller than
R 3 was when we set λ 3 =
(50 versus 500).
Furthermore, in the rightmost pane, we see that the test loss of the model is small regardless of which value of λ 1 was chosen.
This provides additional evidence that R 1 is kept small throughout the training process.
G.2 CIFAR10 Experiments
To demonstrate the signiﬁcance of our approach on “realistic" datasets/models, we consider experiments with CNNs on CIFAR10.
[8] showed that, in practice, training the second-order Taylor expansion of the network tracks the true gradient descent dynamics far better than the network’s linearization does.
This is further demonstrated in Figure 4.
Here, we train a 4-layer CNN with width 512, ReLU activation, average pooling between each layer, and the standard PyTorch initialization, on the cats vs. horses
CIFAR10 classiﬁcation task.
We train via SGD with batch size 128.
For both the train loss and test loss, the dynamics from training f L + f Q tracks the true network dynamics better than just the linearization f L .
Furthermore, f L + f Q acheives a lower test loss than f L (the true network beats both Taylorizations).
[Figure]
In Table 1, we additionally measure the test loss of the linear and quadratic terms after 100 epochs of training the full model .
We observe that both the linear and quadratic term have a nontrivial loss ( < 1 ), and thus learned a nonzero component of the signal.
This provides evidence that in real neural networks, both the linear and quadratic terms learn a nontrivial component of the signal.
[Table]
G.3 Standard MLP Experimental Details
In Figure 2, we demonstrated that “standard” neural networks can effectively learn low-degree dense and high-degree sparse polynomials.
We trained a 2-layer neural network with standard PyTorch initialization and width 100 to learn the target function f ∗ ( x ) =
x T Ax + h 3 ( β T x ) , where A is a high-rank matrix chosen so that x → x T Ax has an L 2 norm of 1.
Here, h 3 is the 3rd Hermite polynomial, and thus h 3 ( β T x ) is a sparse cubic only depending on the random direction β (the
Hermite polynomial is chosen for this task so that it is orthogonal to the quadratic term, making it the “hardest” low-rank cubic to learn).
For varying values of dimension d from 10 to 100 and number of samples n , we train our network via vanilla gradient descent with ﬁxed learning rate 0 .
The initialization, small width, ﬁxed learning rate, and lack of regularization are designed to mimic a standard deep learning setup.
For each value of d , we compute the minimum n required such that the test loss is < 0 .
1 (note that the test loss of the zero predictor is 1 . 0 ).
Figure 2 is a log − log plot of d versus this optimal n .
In Figure 2, we observe that the number of samples needed to obtain 0 .
1 test loss roughly scales with d 2 .
We convincingly see that much fewer than d 3 samples (the red dashed line) are needed.
The NTK, on the otherhand, requires Ω( d 3 ) samples to learn any cubic function.
The minimax sample complexity to learn arbitrary quadratics is Θ( d 2 ) , and therefore this experiment shows that standard neural networks learn “dense qudratic plus sparse cubic" functions with optimal sample complexity.
This provides further evidence that the low-degree plus sparse task is worthy of theoretical study.
Experimental Details.
All experiments were run on an NVIDIA RTX A6000 GPU.
We use the JAX framework [10] along with the Neural Tangents API [40].
Code for all experiments can be found at https://github.com/eshnich/escape_NTK .