B EYOND L INEARIZATION : O N Q UADRATIC AND H IGHER -O RDER A PPROXIMATION OF W IDE N EURAL N ETWORKS
Anonymous authors
Paper under double-blind review
A BSTRACT
Recent theoretical work has established connections between over-parametrized neural networks and linearized models governed by the Neural Tangent Kernels (NTKs).
NTK theory leads to concrete convergence and generalization results, yet the empirical performance of neural networks are observed to exceed their linearized models, suggesting insufﬁciency of this theory.
Towards closing this gap, we investigate the training of over-parametrized neural networks that are beyond the NTK regime yet still governed by the Taylor expansion of the network.
We bring forward the idea of randomizing the neural networks, which allows them to escape their NTK and couple with quadratic models.
We show that the optimization landscape of randomized two-layer networks are nice and amenable to escaping-saddle algorithms.
We prove concrete generalization and expressivity results on these randomized networks, which leads to sample complexity bounds (of learning certain simple functions) that match the
NTK and can in addition be better by a dimension factor when mild distributional assumptions are present.
We demonstrate that our randomization technique can be generalized systematically beyond the quadratic case, by using it to ﬁnd networks that are coupled with higher-order terms in their Taylor series.
I NTRODUCTION
Deep Learning has made remarkable impact on a variety of artiﬁcial intelligence applications such as computer vision, reinforcement learning, and natural language processing.
Though immensely successful, theoretical understanding of deep learning lags behind.
It is not understood how nonlinear neural networks can be efﬁciently trained to approximate complex decision boundaries with a relatively few number of training samples.
There has been a recent surge of research on connecting neural networks trained via gradient descent with the neural tangent kernel (NTK) (Jacot et al., 2018; Du et al., 2018a;b; Chizat & Bach, 2018b;
Allen-Zhu et al., 2018a; Arora et al., 2019a;b).
This line of analysis proceeds by coupling the training dynamics of the nonlinear network with the training dynamics of its linearization in a local neighborhood of the initialization, and then analyzing the expressiveness and generalization of the network via the corresponding properties of its linearized model.
Though powerful, NTK is not yet a completely satisfying theory for explaining the success of deep learning in practice.
In theory, the expressive power of the linearized model is roughly the same as, and thus limited to, that of the corresponding random feature space (Allen-Zhu et al., 2018a;
Wei et al., 2019) or the Reproducing Kernel Hilbert Space (RKHS) (Bietti & Mairal, 2019).
While these spaces can approximate any regular (e.g. bounded Lipschitz) function up to arbitrary accuracy, the norm of the approximators can be exponentially large in the feature dimension for certain non-smooth but very simple functions such as a single ReLU (Yehudai & Shamir, 2019).
Using
NTK analyses, the sample complexity bound for learning these functions can be poor whereas experimental evidence suggests that the sample complexity is mild (Livni et al., 2014).
In practice, kernel machines with the NTK have been experimentally demonstrated to yield competitive results on large-scale tasks such as image classiﬁcation on CIFAR-10; yet there is still a non-neglible performance gap between NTK and full training on the same convolutional architecture (Arora et al.,
2019a
; Lee et al., 2019).
It is an increasingly compelling question whether we can establish theories for training neural networks beyond the NTK regime.
In this paper, we study the optimization and generalization of over-parametrized two-layer neural networks via relating to their higher-order approximations , a principled generalization of the NTK.
Our theory starts from the fact that a two-layer neural network f W 0 + W ( x ) (with smooth activation) can be Taylor expanded with respect to the weight matrix W as
[Equation]
Above, f W 0 does not depend on W , and f (1) corresponds to the NTK model, which is the dominant
W -dependent term when { w r } are small and leads to the coupling between the gradient dynamics for training neural net and its NTK f (1) .
Our key observation is that the dominance of f (1) is deduced from comparing the upper bounds — rather than the actual values —of f ( k ) W 0 , W ( x ) .
It is a priori possible that there exists a subset of W ’s in which the dominating term is not f (1) but some other f ( k ) , k ≥ 2 .
If we were able to train in that set, the gradient dynamics would be coupled with the dynamics on f ( k ) rather than f (1) and thus could be very different.
That learning is coupled with f ( k ) could further offer possibilities for expressing certain functions with parameters of lower complexities, or generalizing better, as f ( k is no longer a linearized model.
In this paper, we build on this perspective and identify concrete regimes in which neural net learning is coupled with higher-order f ( k ) ’s rather than its linearization.
The contribution of this paper can be summarized as follows
• We demonstrate that after randomization , the linear NTK f (1) is no longer the dominant term, and so the gradient dynamics of the neural net is no longer coupled with NTK.
Through a simple sign randomization, the training loss of an over-parametrized two-layer neural network can be coupled with that of a quadratic model (Section 3).
We prove that the randomized neural net loss exhibits a nice optimization landscape in that every second-order stationary point has training loss not much higher than the best quadratic model, making it amenable to efﬁcient minimization (Section 4).
• We establish results on the generalization and expressive power of such randomized neural nets (Section 5).
These results lead to sample complexity bounds for learning certain simple functions that matches the NTK without distributional assumptions and are advantageous when mild isotropic assumptions on the feature are present.
In particular, using randomized networks, the sample complexity bound for learning polynomials (and their linear combination) on (relatively) uniform base distributions is O ( d ) lower than using NTK.
• We show that the randomization technique can be generalized to ﬁnd neural nets that are dominated by the k -th order term in their Taylor series ( k > 2 ) which we term as higher-order NTKs . These models also have expressive power similar as the linear NTK, and potentially even better generalization and sample complexity (Section 6 & Appendix D).
1. P RIOR WORK
We review prior work on the optimization, generalization, and expressivity of neural networks.
Neural Net and Kernel Methods Neal (1996)
ﬁrst proposed the connection between inﬁnitewidth networks and kernel methods.
Later work (Daniely et al., 2016; Williams, 1997; Lee et al., 2018; Novak et al., 2019; Matthews et al., 2018) extended this connection to various settings including deep networks and deep convolutional networks.
These works established that gradient descent on only the output layer weights is well-approximated by a kernel method for large width.
More recently, several groups discovered the connection between gradient descent on all the parameters and the neural tangent kernel (Jacot et al., 2018).
Li & Liang (2018); Du et al.
(2018b) utilized the coupling of the gradient dynamics to prove that gradient descent ﬁnds global minimizers of the training loss of two-layer networks, and Du et al.
(2018a); Allen-Zhu et al.
(2018b)
; Zou et al.
(2018) generalized this to deep residual and convolutional networks.
Using the NTK coupling, Arora et al.
(2019b) ; Cao & Gu (2019a;b) proved generalization error bounds that match the kernel method.
Despite the close theoretical connection between NTK and training deep networks, Arora et al.
(2019a)
; Lee et al.
; Chizat & Bach (2018b) empirically found a signiﬁcant performance gap between NTK and actual training.
This gap has been theoretically studied in Wei et al.
; AllenZhu & Li (2019); Yehudai & Shamir (2019); Ghorbani et al. (2019a) which established that NTK has provably higher generalization error than training the neural net for speciﬁc data distributions and architectures.
The idea of randomization is initiated by Allen-Zhu et al.
(2018a), who use randomization to provably learn a three-layer network; however it is unclear how the sample complexity of their algorithm compares against the NTK.
Inspired by their work, we study the potential gains of coupling with a non-linear approximation over the linear NTK — we compare the performance of a quadratic approximation model with the linear NTK on two-layer networks and ﬁnd that under mild data assumptions the quadratic approximation reduces sample complexity under mild data assumptions.
Outside the NTK Regime
It is believed that the success of SGD is largely due to its algorithmic regularization effects.
A large body of work Li et al.
; Nacson et al.
; Gunasekar et al.
(2018b;a; 2017); Woodworth et al.
(2019) shows that asymptotically gradient descent converges to a max-margin solution with a strong regularization effect, unlike the NTK regularization 1 .
For two-layer networks, a series of works used the mean ﬁeld method to establish the evolution of the network parameters via a Wasserstein gradient ﬂow (Mei et al., 2018b; Chizat & Bach, 2018a;
Wei et al., 2018; Rotskoff & Vanden-Eijnden, 2018; Sirignano & Spiliopoulos, 2018).
In the mean ﬁeld regime, the parameters move signiﬁcantly from their initialization, unlike NTK regime, however it is unclear if the dynamics converge to solutions of low training loss.
Finally, Li et al.
(2019) showed how a combination of large learning rate and injected noise ampliﬁes the regularization from the noise and outperforms the NTK of the corresponding architecture.
Landscape Analysis Many prior works have tried to establish favorable landscape properties such as every local minimum is a global minimum (Ge et al., 2017; Du & Lee, 2018;
Soltanolkotabi et al., 2018; Hardt & Ma, 2016; Freeman & Bruna, 2016; Nguyen & Hein, 2017a;b; Haeffele & Vidal, 2015; Venturi et al., 2018).
Combining with existing advances in gradient descent avoiding saddlepoints (Ge et al., 2015; Lee et al., 2016; Jin et al., 2017), these show that gradient descent ﬁnd the global minimum.
Notably, Du & Lee (2018);
Ge et al.
(2017) show that gradient descent converges to solutions also of low test error, with lower sample complexity than their corresponding NTKs.
Complexity Bounds Recently, researchers have studied norm-based generalization based (Bartlett et al., 2017; Neyshabur et al., 2015; Golowich et al., 2017), tighter compression-based bounds (Arora et al., 2018), and PAC-Bayes bounds (Dziugaite & Roy, 2017; Neyshabur et al., 2017) that identify properties of the parameter that allow for efﬁcient generalization.
P RELIMINARIES
Problem setup We consider the standard supervised learning task, in which we are given a labeled dataset D = { ( x 1 , y 1 ) , . . .
, ( x n , y n ) } , where ( x i , y i ) ∈
X × Y are sampled i.i.d.
from some distribution P , and we wish to ﬁnd a predictor f : X → Y .
Without loss of generality, we assume that X = S d − 1 ( B x ) ⊂ R d for some B x > 0 (so that the features are d -dimensional with norm B x .)
Let (cid:96) :
Y × R → R ≥ 0 be a loss function such that (cid:96) ( y, 0) ≤ 1 , and z (cid:55)→ (cid:96) ( y, z ) is convex, 1Lipschitz, and three-times differentiable with the second and third derivatives bounded by one for all y ∈ Y .
This includes for example the logistic and soft hinge loss for classiﬁcation.
We let
[Equation] [Equation]
denote respectively the empirical risk and population risk for any predictor f : X → Y .
Over-parametrized two-layer neural network We consider learning an over-parametrized twolayer neural network of the form
[Equation]
where W = [ w 1 , . .
w r ] ∈ R d
× m is the ﬁrst layer and a =
[ a 1 , . .
, a m ] (cid:62) ∈
R m is the second layer.
The 1 / √ m factor is chosen to account for the effect of over-parametrization and is consistent with the NTK-type scaling of (Du et al., 2018b; Arora et al., 2019b).
In this paper we ﬁx a and only train W (and thus use f W to denote the network.)
Throughout this paper we assume that the activation is second-order smooth in the following sense.
Assumption A (Smooth activation) .
The activation function σ ∈ C 2 ( R ) , and there exists some absolute constant C > 0 such that | σ (cid:48) ( t ) | ≤ Ct 2 , | σ (cid:48)(cid:48) ( t ) | ≤ C | t | , and σ (cid:48)(cid:48) ( · ) is C -Lipschitz.
An example is the cubic ReLU σ ( t ) =
relu 3 ( t ) =
max { t, 0 } 3 .
The reason for requiring σ to be higher-order smooth (and thus excluding ReLU) will be made clear in the subsequent text 2 .
2. N OTATION
We typically reserve lowercases a, b, α, β, . . .
for scalars, bold lowercases a , b , α , β , . . .
for vectors, and bold uppercases A , B , . . .
for matrices.
For a matrix A
= [ a 1 , . .
. , a m ] ∈ R d × m , its 2 , p norm is deﬁned as (cid:107) A ¯ (cid:107) 2 ,p :=
( (cid:80) mr =1 (cid:107) a r (cid:107) p 2 ) 1 /p for all p ∈ [1 , ∞ ] .
In particular we have (cid:107)·(cid:107) 2 , 2 = (cid:107)·(cid:107)
Fr .
We let B 2 ,p
( R ) :
= { W : (cid:107) W (cid:107) 2 ,p ≤ R } denote a 2 , p -norm ball of radius R .
We use standard Big-Oh notation: a = O ( b ) for stating a ≤ Cb for some absolute constant C > 0 , and a = (cid:101) O ( b ) for a ≤ Cb where C depends at most logarithmically in b and all other problem parameters.
For a twice-differentiable function f : R d → R , x (cid:63) is called a second-order stationary point if ∇ f ( x (cid:63) )
=
and ∇ 2 f ( x (cid:63) ) (cid:23) 0 .
E SCAPING NTK VIA RANDOMIZATION
To motivate our study, we now brieﬂy review the NTK theory for over-parametrized neural nets and provide insights on how to go beyond the NTK regime.
Let W 0 denote the weights in a two-layer neural network at initialization and W denote its movement from W 0 (so that the current weight matrix is W 0 + W .)
The observation in NTK theory, or the theory of lazy training (Chizat & Bach, 2018b), is that for small W the neural network f W 0 + W can be Taylor expanded as
[Equation]
so that the network can be decomposed as the sum of the initial network f W 0 , the linearized model f L W , and higher order terms.
Speciﬁcally (ignoring f W 0 for the moment), when m is large and (cid:107) w r (cid:107)
= O ( m − 1 / 2 ) , we expect f L W =
O (1) and higher order terms to be o m (1) , which is indeed the regime when we train f W 0 + W via gradient descent.
Therefore, the trajectory of training f W 0 + W is coupled with the trajectory of training f W 0 + f L W , which is a convex problem and enjoys convergence guarantees (Du et al., 2018b).
Our goal is to ﬁnd subsets of W so that the dominating term is not f L but something else in the higher order part.
The above expansion makes clear that this cannot be achieved through simple ﬁxes such as tuning the leading scale 1 / √ m or the learning rate — the domination of f L appears to hold so long as the movements w r are small.
Randomized coupling with quadratic model We now explain how the idea of randomization , initiated in (Allen-Zhu et al., 2018a), can help get rid of the domination of f L .
Let W be a ﬁxed weight matrix.
Suppose for each weight vector w r , we sample a random variable Σ rr ∈ R and consider instead the random weight matrix
[Equation]
then the second-order Taylor expansion of f W 0 + WΣ can be written as
[Equation]
where we have deﬁned in addition the quadratic part f Q WΣ .
Due to the existence of { Σ rr } , each original weight w r now gets a weight that is different in f L and f Q .
Speciﬁcally, if we choose
[Equation]
{± 1 }
[Equation]
to be random signs , then we have Σ 2 rr ≡ 1 and thus f Q WΣ ( x ) ≡ f Q W ( x ) , whereas E [Σ rr ] = 0 so that E [ f L WΣ ( x )] ≡ 0 .
Consequently, f Q is not affected by such randomization whereas f L WΣ is now mean zero and thus can have substantially lower magnitude than f L W .
More precisely, when (cid:107) w r (cid:107) 2 (cid:16) m − 1 / 4 , the scalings of f L and f Q compares as follows:
• We have E Σ [ f L WΣ ( x )]
= 0 and E Σ
[Equation]
L
[Equation]
so we expect f L WΣ ( x ) =
O ( m − 1 / 4 ) over a random draw of Σ .
• The quadratic part scales as
[Equation]
Therefore, at the random weight matrix WΣ , f Q dominates f L and thus the network is coupled with its quadratic part rather than the linear NTK.
3. L EARNING RANDOMIZED NEURAL NETS
The randomization technique leads to the following recipe for learning W : train W so that (cid:107) w r (cid:107) 2 =
O ( m − 1 / 4 ) and WΣ has in expectation low loss.
We make this precise by formulating the problem as minimizing a randomized neural net risk.
Randomized risk Let (cid:101)
L : R d × m → R denote the vanilla empirical risk for learning f W :
[Equation]
, where we have reparametrized the weight matrix into W 0 + W so that learning starts at W = 0 .
Following our recipe, we now formulate our problem as minimizing the expected risk
[Equation]
where Σ ∈ R m × m is a diagonal matrix with Σ
rr iid ∼ Unif {± 1 } .
To encourage (cid:107) w r (cid:107)
= O ( m − 1 / 4 ) and improve generalization, we consider a regularized version of (cid:101) L and L with (cid:96) 2 , 4 regularization: (cid:101)
L λ ( W ) : =
(cid:101) L ( W ) + λ (cid:107) W (cid:107) 82 , 4 and L λ ( W )
=
L ( W ) + λ (cid:107) W (cid:107) 82 , 4 = E Σ [ (cid:101)
L λ ( WΣ )] .
Our regularizer penalizes W , i.e. the distance from initialization, similar as in (Hu et al., 2019).
Symmetric initialization We initialize the parameters ( a , W 0 ) randomly in the following way: set
[Equation]
Above, we set half of the a i ’s as +1 and half as − 1 , and the weights w 0 ,r are i.i.d.
in the +1 half and copied exactly into the − 1 half.
Such an initialization is almost equivalent to i.i.d.
random W 0 , but has the additional beneﬁt that f W 0 ( x ) ≡ 0 and also leads to simple expressivity arguments.
Our initialization scale B − 2 x is chosen so that for a random draw of w 0 , we have w (cid:62) 0 x ∼ N (0 , 1) , which is on average O (1) 5 .
For technical convenience, we also assume henceforth that the realized { w 0 ,r } satisﬁes the bound
[Equation]
This happens with probability at least 1 − δ under random initialization (see proof in Appendix A.3), and ensures that max r ∈ [ m ] | w (cid:62) 0 ,r x | ≤ (cid:101) O ( √ d ) simultaneously for all x .
O PTIMIZATION
In this section, we show that L λ enjoys a nice optimization landscape.
4. N ICE LANDSCAPE OF CLEAN RISK
As the randomized loss L induces coupling of the neural net f W 0 + WΣ with the quadratic model f Q W , we expect its behavior to resemble the behavior of gradient descent on the following clean risk :
[Equation]
Above, we have deﬁned diagonal matrices D i
= diag( (cid:8) a r σ (cid:48)(cid:48)
( w (cid:62) 0 ,r x i ) (cid:9) r ∈ [ m ] ) ∈
R m × m which are not trained.
We now show that the clean risk L Q , albeit non-convex, possesses a nice optimization landscape.
Lemma 1 (Landscape of clean risk) .
Suppose there exists
W (cid:63) ∈ R d × m such that L Q ( W (cid:63) ) ≤ OPT .
Let
Σ (cid:48) ∈
R m ×
m be a diagonal matrix with Σ (cid:48)
rr iid ∼ Unif {± 1 } , then we have
[Equation]
This result implies that, for W in a certain ball and large m , every point of higher loss than W (cid:63) will have either a ﬁrst-order or a second-order descent direction.
In other words, every approximate second-order stationary point of L Q is also an approximate global minimum.
Our proof utilizes the fact that L Q is similar to the loss function in matrix sensing / learning quadratic neural networks, and builds on recent understandings that the landscapes of these problems are often nice (Soltanolkotabi et al., 2018; Du & Lee, 2018; Allen-Zhu et al., 2018a).
The proof is deferred to Appendix B.1.
4. N ICE LANDSCAPE OF RANDOMIZED NEURAL NET RISK
With the coupling between f W 0 + WΣ ( x ) and f Q W ( x ) in hand, we expect the risk L ( W ) =
E [ (cid:101) L ( WΣ )] to enjoy similar guarantees as the clean risk does L Q ( W ) in Lemma 1.
We make this precise in the following result.
Theorem 2 (Landscape of L ) .
Suppose there exists W (cid:63) ∈ B 2 , 4 ( B w,(cid:63) ) such that L Q ( W (cid:63) ) ≤ OPT , and that
[Equation]
for some ﬁxed (cid:15) ∈ (0 , 1] and B w ≥ B w,(cid:63) , then for all W ∈ B 2 , 4 ( B w ) , we have
[Equation]
As an immediate corollary, we have a similar characterization of the regularized loss L λ .
Corollary 3 (Landscape of L λ ) .
For any B w ≥ B w,(cid:63) , under the conditions of Theorem 2, we have for all λ > 0 and all W ∈ B 2 , 4 ( B w ) that
[Equation]
where C =
O (1) is an absolute constant.
Theorem 2 follows directly from Lemma 1 through the coupling between L and L Q (as well as their gradients and Hessians).
Corollary 3 then follows by controlling in addition the effect of the regularizer.
The full proof of Theorem 2 and Corollary 3 are deferred to Appendices B.4 and B.5.
We now present our main optimization result, which follows directly from Corollary 3.
Theorem 4 (Optimization of L λ ) .
Suppose there exists W (cid:63) such that
[Equation]
for some OPT > 0 .
For any γ = Θ(1) and (cid:15) > 0 , we can choose λ suitably and m ≥ (cid:101) O
(poly( d, B x B w,(cid:63) , (cid:15) − 1 ))
such that the regularized loss L λ satisﬁes the following: any second order stationary point (cid:99) W has low loss and bounded norm:
[Equation]
Proof sketch.
The proof of Theorem 4 consists of two stages: ﬁrst “localize” any second-order stationary point into a (potentially very big) norm ball using the (cid:107)·(cid:107) 82 , 4 regularizer, then use Corollaryin this ball to further deduce that L λ is low and (cid:13)(cid:13)(cid:13)(cid:99) W (cid:13)(cid:13)(cid:13) 2 , 4 ≤ O ( (cid:107) W (cid:63) (cid:107) 2 , 4 ) .
The full proof is deferred to Appendix B.6.
Efﬁcient optimization through escaping-saddle algorithms Theorem 4 states that when the over-parametrization is enough, any second-order stationary point (SOSP) (cid:99) W of L λ has loss competitive with OPT , the performance of best quadratic model.
Consequently, algorithms that are able to ﬁnd SOSPs (escape saddles) such as noisy SGD (Jin et al., 2019) can efﬁciently minimize L λ to up to a multiplicative / additive factor of OPT . Further, by sampling a fresh Σ at each iteration and using the stochastic gradient ∇ W (cid:101) L λ ( WΣ ) (rather than computing the full ∇ L λ ( W ) ), the noisy SGD iterates can be computed efﬁciently.
G ENERALIZATION AND E XPRESSIVITY
We now shift attention to studying the generalization and expressivity of the (randomized) neural net (cid:99) W learned in Theorem 4.
5. G ENERALIZATION
As (cid:99) W is always coupled (through randomization) with the quadratic model f Q (cid:99) W , we begin by studying the generalization of the quadratic model.
Generalization of quadratic models Let
[Equation]
denote the class of quadratic models for W in a (cid:96) 2 , 4 ball.
We ﬁrst present a lemma that relates the
Rademacher complexity of F Q ( B w ) to the expected operator norm of certain feature maps.
Lemma 5 (Bounding generalization of f Q via feature operator norm) .
For any non-negative loss (cid:96) such that z (cid:55)→ (cid:96) ( y, z ) is 1-Lipschitz and (cid:96) ( y, 0) ≤ 1 for all y ∈ Y , we have the Rademacher complexity bound
[Equation]
where σ i iid ∼ Unif {± 1 } are Rademacher variables.
Operator norm based generalization Lemma 5 suggests a possibility for the quadratic model to generalize better than the NTK model: the Rademacher complexity of F Q ( B w ) depends on the “feature maps”
[Equation]
,r x i ) x i x (cid:62) i through their matrix operator norm.
Compared with the (naive) Frobenius norm based generalization bounds, the operator norm is never worse and can be better when additional structure on x is present.
The proof of Lemma 5 is deferred to Appendix C.1.
We now state our main generalization bound on the (randomized) neural net loss L , which concretizes the above insight.
Theorem 6 (Generalization of randomized neural net loss) .
For any data-dependent (cid:99) W such that (cid:13)(cid:13)(cid:13)(cid:99) W (cid:13)(cid:13)(cid:13) 2 , 4 ≤ B w , we have
[Equation]
where M x, op := (cid:16) B − 2 x E x
(cid:104)(cid:13)(cid:13) 1 n (cid:80) ni
=1 x i x (cid:62) i (cid:13)(cid:13) op (cid:105)(cid:17) 1 / 2 is the (rescaled) operator norm of the empirical covariance matrix.
In particular, M x, op ≤ 1 always holds; if in addition v (cid:62) x is K (cid:112)
Var( v (cid:62) x ) sub-Gaussian for all v ∈ S d − 1 (1) and κ (Cov( x )) ≤ κ , then M x, op ≤ κ/ √ d whenever n ≥ O
( K 4 d ) .
The generalization bound in Theorem 6 features two desirable properties:
(1) For large m (e.g. m (cid:38) n 4 ), the bound scales at most logarithmically with the width m , therefore allowing learning with small samples and extreme over-parametrization; (2) The main term (cid:101) O
( B 2 x B 2 w M x, op / √ n ) automatically adapts to properties of the feature distribution and can lower the generalization error than the naive bound by at most O (1 / √ d ) without requiring us to tune any hyperparameter.
Concretely, we have M x, op ≤ O (1 / √ d ) when x has an isotropic distribution such as Unif( S d − 1 ( B x )) or Unif {± B x / √ d } d .
Theorem 6 follows directly from Lemma 5 and a matrix concentration Lemma.
The proof is deferred to Appendix C.2.
5. E XPRESSIVITY AND S AMPLE C OMPLEXITY THROUGH Q UADRATIC M ODELS
In order to concretize our generalization result, we now study the expressive power of quadratic models through the concrete example of learning functions of the form (cid:80) j ≤ k α j ( β (cid:62) j x ) p j , i.e. sum of “one-directional” polynomials (for consistency and comparability with (Arora et al., 2019b).)
Theorem 7 (Expressing a sum of polynomials through f Q ) .
Suppose { ( a r , w 0 ,r ) } are generated according to the symmetric initialization (3) and we use σ ( t ) =
1 6 relu 3 ( t ) (so that σ (cid:48)(cid:48) ( t ) =
relu( t ) .) If f (cid:63) ( x ) = (cid:80) kj
{+=
+}
α j ( β (cid:62) j x ) p j achieves training loss L ( f (cid:63) ) ≤ (cid:15) 0 , where p j − 2 ∈ { 1 } ∪ { 2 (cid:96) } (cid:96) ≥ 0 ,.
Then so long as the width is sufﬁciently large:
[Equation]
we have with probability at least 1 − δ (over W 0 ) that there exists
W (cid:63) ∈ R d × m such that L Q [Equation] W
[Equation]
) ≤ OPT [Equation] = [Equation]
+ [Equation] and [Equation]
W [Equation] [Equation] The proof of Theorem 7 is based on a reduction from expressing degree p polynomials using quadratic models to expressing degree p − 2 polynomials using random feature models.
The proof can be found in Appendix C.4.
Comparison between quadratic and linearized (NTK) models We now illustrate our results in
Theorem 6 and 7 in three concrete examples, in which we compare the sample complexity bounds of the randomized (quadratic) network and the linear NTK when m is sufﬁciently large.
Learning a single polynomial .
Suppose f (cid:63) ( x ) =
α ( β (cid:62) x ) p satisﬁes L ( f (cid:63) ) ≤ (cid:15) , and we wish to ﬁnd (cid:99) W with O ( (cid:15) ) test loss.
By Theorem 7 we can choose W (cid:63) such that L Q ( W (cid:63) ) ≤ OPT = 2 (cid:15) , and by Theorem 4 we can ﬁnd (cid:99) W such that L ( (cid:99) W )
≤ L λ ( (cid:99) W ) ≤ 3 (cid:15) and (cid:107) (cid:99) W (cid:107) 2 , 4 =
O ( B w,(cid:63) ) .
Take
B x = 1 , and assume x is sufﬁciently isotropic so that M x, op
= O ( 1 √ d ) , the sample complexity from
Theorem 6 is
[Equation]
In contrast, the sample complexity for linear NTK (Arora et al., 2019b; Cao & Gu, 2019a) to reach (cid:15) test loss is
[Equation]
We have n Q /n
L
= (cid:101) O
( p/d ) , a reduction by a dimension factor unless p (cid:16) d .
We note that the above comparison is simply comparing upper bounds, since in general the lower bound on the sample complexity of linear NTK is unknown.
Learning a noisy 2 -XOR .
Wei et al.
(2019) established a sample complexity lower bound of linear
NTK of n ≥ n L = Ω( d 2 ) to achieve constant generalization error, which allows for a rigorous comparison against the quadratic model.
The ground truth function in 2 -XOR is f (cid:63) ( x ) =
x 1 x 2 = ([( e 1 + e 2 ) (cid:62) x ] 2 − [( e 1 − e 2 ) (cid:62) x ] 2 ) / 4 , where x ∈ {± 1 } d , and f (cid:63) attains constant margin on the training distribution constructed in Wei et al.
By Theorem 7, f (cid:63) can be (cid:15) -approximated by f Q W (cid:63) with B 4 w,(cid:63) ≤ O (1) .
Thus by
Theorem 6 the sample complexity for learning noisy 2 -XOR through the randomized net (cid:99) W is
[Equation]
This is (cid:101) O
( d ) better than the sample complexity lower bound of linear NTK and thus provably better. Low-rank matrix sensing .
Suppose we wish to learn a symmetric low-rank matrix A (cid:63) ∈ R d × d with (cid:107) A (cid:63) (cid:107) op ≤ 1 and rank( A (cid:63) ) ≤ r through n rank-one observations of the form y i = f (cid:63) ( x i ) = (cid:10) A (cid:63) , x i x (cid:62) i (cid:11) where x ∼ Unif( S d − 1 ( √ d )) . This ground truth function can be written as f (cid:63) ( x i ) =
(cid:80) rj
=1 α j ( v (cid:62) j x i ) 2 where | α j | ≤ 1 are the eigenvalues of A and v j ∈ R
d are the corresponding eigenvectors.
For any 1-Lipschitz loss such as the absolute loss, by Theorem 7, there exists W (cid:63) such that L Q ( W (cid:63) ) ≤ OPT = (cid:15) and B 4 w,(cid:63) ≤ O ( r (cid:80) rj
=1 (cid:107) v j (cid:107) 42 ) =
O ( r 2 ) .
Thus by Theorem 6, the sample
[Table]
through the randomized net (cid:99) W is
[Equation]
This compares favorably against the sample complexity upper bound for linear NTK, which needs
[Equation] samples.  H IGHER - ORDER NTK S In this section, we demonstrate that our idea of randomization for changing the dynamics of learning neural networks can be generalized systematically — through randomization we are able to obtain over-parametrized neural networks in which the k -th order term dominates the Taylor series.
Consider a two-layer neural network with 2 m neurons and symmetric initialization (cf.
[Equation]
Assuming σ is analytic on R (i.e. it equals its Taylor series at any point), we have f W 0 + W ( x ) =
[Equation]
∞ k =0 f ( k ) W 0 , W ( x ) , where we have deﬁned the k -th order NTK
[Equation]
,r x ) k (cid:1) . Note that f (0) ( x ) ≡ 0 due to the symmetric initialization, and f (1) ( x ) is the standard NTK.
Coupling with f ( k ) via randomization We show in Appendix D.1 that for any k ≥ 2 , there exists a randomized W at which f W 0 + W ( x ) =
f ( k ) W 0 , W ( x ) + o m (1) .
More concretely, our choice of W is such that all of f (1) , . .
f ( k − 1) have on average o m (1) in magnitude, and f ( k )
=
[-E [ z k − ] + 1 . Set ( w + ,r , w − ,r ) = ( z + ,r w (cid:63),r , z − ,r w (cid:63),r ) , and take (cid:107) w (cid:63),r (cid:107) 
 =-] O (1) is the dominating term.
Generalization and expressivity
We establish the generalization of expressivity of f ( k ) in Appendix D.2, which systematically extends our results on the quadratic model.
We show that the sample complexity for learning degree ≥ k polynomials through f ( k ) compared with linear NTK can be better by a factor of d k − 1 for large n , when mild distributional assumptions on x such as approximate isotropy (constant condition number of the k th moment tensor) is present.
C ONCLUSION
In this paper we proposed and studied the optimization and generalization of over-parametrized neural networks through coupling with higher-order terms in their Taylor series.
Through coupling with the quadratic model, we showed that the randomized two-layer neural net has a nice optimization landscape (every second-order stationary point has low loss) and is thus amenable to efﬁcient minimization through escape-saddle style algorithms.
These networks enjoy the same expressivity and generalization guarantees as linearized models but in addition can generalize better by a dimension factor when distributional assumptions are present.
We extended the idea of randomization to show the existence of neural networks whose Taylor series is dominated by the k -th order term.
We believe our work brings in a number of open questions, such as how to better utilize the expressivity of quadratic models, or whether the study of higher-order expansions can lead to a more satisfying theory for explaining the success of full training.
We also note that the Taylor series is only one avenue to obtaining accurate approximations of nonlinear neural networks.
It would be of interest to design other approximation schemes for neural networks that are coupled with the network in larger regions of the parameter space.
A T ECHNICAL TOOLS A. A MATRIX OPERATOR NORM CONCENTRATION BOUND
Lemma 8 (Variant of Theroem 4.6.1, (Tropp et al., 2015)) .
Suppose { A r,i } r ∈ [ m ] ,i ∈ [ n ] are ﬁxed symmetric d × d matrices and { σ i } i ∈ [ n ] iid ∼ Unif {± 1 } are Rademacher variables.
Letting
[Equation]
(cid:88)
i
[Equation]
r
[Equation]
i , then we have
[Equation]
where
[Equation]
Proof.
Applying the high-probability bound in (Tropp et al., 2015, Theorem 4.6.1) and the union bound [Equation] we get
[Equation]
Let V := max r ∈ [ m ] v ( Y r ) [Equation] we have by integrating the above bound over t that
[Equation]
A. E XPRESSING POLYNOMIALS WITH RANDOM FEATURES
Lemma 9.
Let σ ( t ) =
relu( t ) and w 0 ∼ N (0 , B − 2 x I d ) be Gaussian random features.
For any p ∈ { 1 } ∪ { 2 (cid:96) } (cid:96) ≥ 0 and β ∈ R d , there exists a random variable a = a ( w 0 ) such that
[Equation]
and a satisﬁes the (cid:96) 2 norm bound
[Equation]
Proof.
Consider the ReLU random feature kernel
[Equation]
and let H K denote the RKHS associated with this kernel.
By the equivalence of feature maps (Minh et al., 2006, Proposition 1), for any feature map φ : S d − 1 ( B x ) (cid:55)→ H (where H is a Hilbert space) that generates K in the sense that
[Equation]
we have for any function f that
[Equation]
11) and the inﬁmum over a is attainable whenever it is ﬁnite.
For the ReLU random feature kernel K , let u := x (cid:62) x (cid:48) /B
2 x and N 2 ( ρ ) denote a bivariate normal distribution with marginals N (0 , 1) and correlation ρ ∈ [ − 1 , 1] .
We have that
[Equation]
where the constants { c p } satisfy c
[Equation]
[Equation]
[Equation] ≥ [Equation] ) for (cid:96) ≥ 1 , (and thus c p ≥ (2 π ( p ∨ 1) 3 ) − 1 for all p ), and x ⊗ k ∈ R d k denote the k -wise tensor product of x .
Therefore, if we deﬁne feature map
[Equation]
we have K ( x , x (cid:48) )
= (cid:104) φ ( x ) , φ ( x (cid:48) )
(cid:105) .
With this feature map, the function f (cid:63) ( x ) =
α ( β (cid:62) x ) p can be represented as
[Equation]
x )
[Equation]
where [Equation] Thus by the feature map equivalence (11), we have f (cid:63) ∈ H K and
[Equation]
Now apply the feature map equivalence (11) again with the random feature map
[Equation]
(which maps into the inner product space of w 0 ∼ N (0 , B − 2 x I d ) ), we conclude that there exists a = a ( w 0 ) such that f (cid:63) =
E w
[relu( w (cid:62) 0 x ) a ] and
[Equation]
A. P ROOF OF E QUATION (4)
Let N be an 1 / 2 -covering of S d − 1 (1) .
We have |N| ≤ 5 d and for any vector w ∈ R d that (cid:107) w (cid:107) 2 ≤ 2 sup v ∈N ( v (cid:62) w ) (see e
[Equation]
g. (Mei et al., 2018a, Section A).)
We thus have
[Equation]
Setting t =
(cid:112) 8( d log 5 + log( m/δ )) =
O ( (cid:112) d + log( m/δ )) ensures that the above probability does not exceed δ as desired
[Equation]
B P ROOFS FOR S ECTIONB. P ROOF OF L EMMA
Computing the gradient of L Q , we obtain
[Equation]
Further computing the Hessian gives
[Equation]
Taking expectation over Σ (cid:48) , and using that (cid:96) (cid:48)(cid:48) ≤ 1 , term II can be bounded as
[Equation]
[Table]
Term I does not involve Σ (cid:48) and can be deterministically bounded as
[Equation]
where (i) follows directly by computing (cid:10) ∇ L Q ( W ) , W (cid:11) and the convexity of z (cid:55)→ (cid:96) ( y, z ) , and (ii) follows from the assumption that L Q ( W (cid:63) ) ≤ OPT .
Combining the bounds for terms I and II gives the desired result.
B. C OUPLING LEMMAS
Lemma 10 (Bound on f Q ) .
For any W ∈ R d × m , the quadratic model f Q W satisﬁes the bound
[Equation]
for all x ∈ S d − 1 ( B x ) .
Proof.
We have
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
Lemma 11 (Coupling between f and f Q ) .
We have for all x ∈ S
d − 1 ( B x ) that
(a) E Σ [ f L WΣ ( x )]
= 0 and E Σ [( f L WΣ ( x )) 2 ] ≤ (cid:101)
O
( d 2 B 2 x (cid:107) W (cid:107) 22 , 4 m − 1 / 2 ) .
(b) | ∆ Q WΣ ( x ) |
≤ O ( B 3 x (cid:107) W (cid:107) 32 , 4 m − 1 / 4 ) (almost surely for all Σ .)
Proof.
(a) Recall that
[Equation]
As Σ rr has mean zero, we have E Σ [ f L ]
= 0 and
[Equation]
Above, (i) follows from the assumption that | σ (cid:48) ( t ) | ≤ Ct 2 , (ii) is Cauchy-Schwarz, (iii) uses the bound (4), and (iv) uses the power mean inequality on (cid:107) w r (cid:107) 2 .
[Equation]
[Equation] [Equation]
[Equation]
where again (i) uses the power mean inequality on (cid:107) w r (cid:107) 2 .
B. C LOSENESS OF LANDSCAPES
Lemma 12 ( L Q close to L ) .
We have for all W ∈ R d × m that
[Equation]
Proof.
Recall that
[Equation]
W 0 + WΣ ( x ))] and L Q ( W ) =
E D [ (cid:96) ( y, f Q W ( x ))] .
By the 1-Lipschitzness of z (cid:55)→ (cid:96) ( y, z ) we have
[Equation]
where the last step uses Lemma 11.
Lemma 13 (Closeness of directional gradients) .
We have (cid:12)(cid:12)
[Equation]
≤
[Equation]
O
[Equation]
W (cid:107)
[Equation]
[Equation] [Equation]
2 (cid:17)
[Equation]
Proof.
Differentiating L and L Q and taking the inner product with W , we get
[Equation]
and
[Equation]
Therefore, by expanding σ (cid:48) (( w 0 ,r + Σ rr w r ) (cid:62) x ) and noticing that Σ 2 rr ≡ 1 , we have
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
(Σ
[Equation]
(cid:62)
[Equation]
(cid:35) (cid:124) (cid:123)(cid:122) (cid:125) III (cid:12)(cid:12)(cid:12)(cid:12)(cid:12) .
We now bound the three terms separately.
Recall that | (cid:96) (cid:48) | ≤ 1 and (cid:96) (cid:48) ( y, z ) is 1-Lipschitz in z .
For term I we have by Cauchy-Schwarz that
[Equation]
For term II, we have
[Equation]
where (i) uses Cauchy-Schwarz and (ii) uses the bounds in Lemma 10 and 11.
For term III we ﬁrst note by the smoothness of σ (cid:48)
that
[Equation]
Substituting this bound into term III yields
[Equation]
Putting together the bounds for term
I
[Equation]
II, III gives the desired result
[Equation]
Lemma 14 (Closeness of Hessians) .
Let Σ (cid:48) denote a diagonal matrix with diagonal entries drawn i.i.d.
from Unif {± 1 } .
We have for all W , W (cid:63) ∈ R d × m that
[Equation]
Proof.
Differentiating L and L Q twice on the direction W (cid:63) Σ (cid:48) , we get ∇ 2 L ( W ) [ W [Equation] Σ (cid:48) , W
[Equation]
Σ (cid:48) ]
=
[Equation]
m (cid:88) r ≤ m a r σ (cid:48)
(( w [Equation] r [Equation]
w r ) (cid:62) x )(Σ rr Σ (cid:48)
rr w (cid:62) (cid:63)
[Equation]
r x ) 
 2    (cid:124) (cid:123)(cid:122) (cid:125) I
( L ) + E Σ , D   (cid:96) (cid:48)
( y, f W
[Equation]
( x )) ·
[Equation]
√
[Equation]
r
[Equation]
m a [Equation] ( w
[Equation]
(cid:62) x )(Σ
[Equation]
(cid:48)
rr
[Equation]
(cid:62) (cid:63),r x )  (cid:124) (cid:123)(cid:122) (cid:125) II( L ) , and ∇ 2 L Q ( W )[ W (cid:63) Σ (cid:48) , W (cid:63) Σ (cid:48) ] = E Σ (cid:48) , D   (cid:96) (cid:48)(cid:48) ( y, f Q W ( x )) ·   1 √ m (cid:88) r ≤ m a r σ (cid:48)(cid:48) ( w (cid:62) 0 ,r x )( w (cid:62) r x )(Σ (cid:48)
rr w (cid:62) (cid:63),r x ) 
 2 
 (cid:124) (cid:123)(cid:122) (cid:125) I
( L Q ) + E Σ (cid:48) , D   (cid:96) (cid:48)
( y, f Q W ( x )) · 1 √ m (cid:88) r ≤ m a r σ (cid:48)(cid:48) ( w (cid:62) 0 ,r x )(Σ (cid:48)
rr w [Equation] We ﬁrst bound the terms I( L ) and I( L Q ) . We have
[Equation]
Using similar arguments on I( L Q ) gives the bound
[Equation]
We now shift attention to bounding II( L ) − II( L Q ) .
First note that
[Equation]
[Table]
[Equation]
Combining all the bounds gives the desired result.
B. P ROOF OF T HEOREM
We apply Lemma 12, 13, and 14 to connect the neural net loss L to the “clean risk” L Q .
First, by
Lemma 12, we have for all the assumed W that
[Equation]
Therefore we have (cid:12)(cid:12) L ( W ) − L Q ( W ) (cid:12)(cid:12) ≤ (cid:15)/ 6 so long as m ≥
[Equation]
O
[Equation]
B
[Equation]
B 12 w (cid:15) − 4 + d
[Equation]
B
[Equation]
B 4 w
[Equation]
− 2 (cid:1) .
Applying Lemma 1, we obtain that
[Equation]
provided that the error term in Lemma 1 is bounded by (cid:15)/
[Equation]
which happens when
[Equation]
Finally
[Equation]
we choose m sufﬁciently large so that
[Equation]
and
[Equation]
which combined with (14) yields the desired result.
By Lemma 13 and 14, it sufﬁces to choose m such that, to satisfy the closeness of directional gradients,
[Equation]
16) and to satisfy the closeness of Hessian quadratic forms,
[Equation]
Collecting the requirements on m in (13), (15), (16), (17) and merging terms using (cid:15) ≤ 1 and
B w,(cid:63) ≤ B w , the desired result holds whenever
[Equation]
This completes the proof
[Equation]
B. P ROOF OF C OROLLARY
Proof.
For all λ ≥ 0 deﬁne
[Equation]
By Lemma 2, it sufﬁces to show that
[Equation] for some absolute constant C . Recall that L λ ( W ) = L ( W ) + λ (cid:107) W (cid:63) (cid:107) 82 , 4 .
By differentiating A (cid:55)→ (cid:107) A (cid:107) 82 , 4 we get (cid:104)∇ ( L λ − L )( W ) , W (cid:105) =
2 (cid:107) W (cid:107) 42 , 4 · (cid:88) r ≤ m (cid:68) 4 λ (cid:107)
[Equation]
(cid:107) 22 w r , w r (cid:69) = 8 λ (cid:107) W (cid:107) 82 ,and ∇ 2 ( L λ
[Equation]
[Equation]
W
[Equation]
Σ (cid:48) , W (cid:63) Σ (cid:48) ]
= 8 λ (cid:107) W (cid:107) 4 2 , (cid:88) r ≤ m (cid:107) w r (cid:107) 2 2 (cid:107) w (cid:63),r Σ (cid:48)
{+rr (cid:107) 2 
+}
+ 2 (cid:104) w r , w (cid:63),r Σ (cid:48) rr (cid:105)  + 32 λ (cid:107) W (cid:107) 42 , 4 · 
[Equation] 

[Equation]
≤ 56 λ (cid:107) W (cid:107) 42 ,
[Equation]
where (i) used Cauchy-Schwarz and (ii) used the AM-GM inequality p 3 q ≤ αp 4 / 4 + 27 q 4 / (4 α 3 ) for all p, q and α > 0 .
Substituting the above expressions into A λ − A 0 yields
[Equation]
Choosing α = 5 / 14 gives the desired result.
B. P ROOF OF T HEOREM
We begin by choosing the regularization strength as
[Equation]
where λ 0 is a constant to be determined.
Let (cid:15) be an accuracy parameter also to be determined.
Localizing second-order stationary points
We ﬁrst argue that any second order stationary point W has to satisfy (cid:107) W (cid:107) 2 , 4 ≤ B w, 0 for some large but controlled B w, 0 . We ﬁrst note that for the clean risk L Q , we have for any W ∈ R d × m that [Equation]
where (i) uses convexity of (cid:96) and (ii) uses the assumption that (cid:96) ( y, 0) ≤ 1 for all y ∈ Y .
Now, applying the coupling Lemma 13, and combining with the fact that (cid:68) ∇ W ( λ (cid:107) W (cid:107) 82 , 4 ) , W (cid:69) = 8 λ (cid:107) W (cid:107) 82 , 4 , we have simultaneously for all W that
[Equation]
[Equation]
Therefore we see that any stationary point W has to satisfy
[Equation]
By Corollary 3, choosing m ≥ poly( λ − 1 0 , d, B w,(cid:63) B x , (cid:15) ) , the coupling error is bounded by (cid:15) in
B 2 , 4 ( B w, 0 ) , i.e. for all W ∈ B 2 , 4 ( B w, 0 ) we have that
[Equation]
where C =
O (1) is an absolute constant.
Bounding loss and norm Choosing
[Equation]
we get that CλB 8 w,
[Equation]
= 2 γ OPT + (cid:15) , and thus the bound (18) reads
[Equation]
For the second-order stationary point (cid:99) W , the gradient term vanishes and the Hessian term is nonnegative, so we get
[Equation]
and thus
[Equation]
Further, by re-writing (18), we obtain
[Equation]
for any γ = O (1) .
This is the desired result.
C P ROOFS FOR S ECTIONC. P ROOF OF L EMMA
As the loss (cid:96) ( y [Table] z ) is 1-Lipschitz in z for all y , by the Rademacher contraction theorem (Wainwright, 2019, Chapter 5) we have that
[Table] [Equation] [Table]
op
[Equation]
m
[Equation]
r
[Equation]
m
[Equation]
[Table]
w
[Table]
[Equation]
[Table]
[Equation]
[Table]
∗   +
[Equation]
√ n ≤ E σ , x  [Equation]
max [Equation] n n [Equation]
i
[Equation] [Table] [Equation] [Table]
[Equation] [Table] [Equation] [Table]
op 
[Table]
[Equation] [Table] 1 √ m [Table]
[Equation]
[Table] [Equation] (cid:125) [Equation]
[Table] w [Equation] 1 √ n, where the last step used the power mean (or Cauchy-Schwarz) inequality on {(cid:107) w r (cid:107) 2 } .
We ﬁrst relate the generalization of L to that of L Q through L
[Equation]
By Lemma 12, we have simultaneously for all W ∈ B 2 , 4 ( B w ) that
[Equation]
Further, from the proof we see that the argument does not depend on the distribution of x (it holds uniformly for all x ∈ S d − 1 ( B x ) , therefore for the population version we also have the bound
[Equation]
These bounds hold for all W ∈ B 2 , 4 ( B w ) so apply to (cid:99) W . Therefore it remains to bound L QP ( (cid:99) W ) − L Q ( (cid:99) W ) , i
[Equation]
e. the generalization of the quadratic model. Generalization of quadratic model By symmetrization and applying Lemma
[Equation]
we have
[Equation]
[Equation] [Equation] We now focus on bounding the expected max operator norm above.
First, we apply the matrix concentration Lemma 8 to deduce that [Table] [Equation] 
 [Equation]
As | σ (cid:48)(cid:48) ( t ) | ≤ Ct and w (cid:62) 0 ,r x i ∼ N (0 , 1) for all ( r, i ) , by standard expected max bound on subexponential variables we have [Equation] Therefore deﬁning
[Equation]
and substituting the above bound into (21) yields that
[Equation]
Combining the bound with the coupling error (19) and (20), we arrive at the desired result.
For M x, op we have two versions of bounds:
(a) We always have (cid:13)(cid:13)(cid:13)(cid:80) i ≤ n x i x (cid:62) i /n
(cid:13)(cid:13)(cid:13) op ≤ B 2 x , and thus M x, op ≤ 1 .
(b) If, in addition, x is uniformly distributed on the sphere S d − 1 ( B x ) or the hypercube (cid:110)
± B x / √ d (cid:111)
d , then we have by standard covariance concentration (Vershynin, 2018, The-
orem 4.7.1) that E x (cid:20)(cid:13)(cid:13)(cid:13)(cid:80) i ≤ n x
i x (cid:62) i /
[Equation]
(cid:13)(cid:13)(cid:13) op (cid:21)
[Equation] B 2 x /d ) when n ≥ d . More generally, if for all v ∈ S d − 1 (1) we have
[Equation]
and
that κ (Cov( x )) ≤ κ , then we have (cid:107) Cov( x ) ) (cid:107) op ≤ κB 2 x /d . Applying
Vershynin, 2018, Theorem we get M x, op ≤ κ/ √ d whenever n ≥ O ( K 4 d )
E XPRESSIVE POWER OF INFINITELY WIDE QUADRATIC MODELS Lemma 15 (Expressivity of f Q with inﬁnitely many neurons) . Suppose f (cid:63) ( x ) =
α ( β (cid:62) x ) p for some α ∈ R , β ∈ R d , and p ≥ 2 and such that p − 2 ∈ { 1 } ∪ { 2 (cid:96) } (cid:96) ≥ 0 .
Suppose further that we use σ ( t )
=
16 relu 3 ( t ) (so that σ (cid:48)(cid:48) ( t ) =
relu( t ) ), then there exists choices of ( w + , w − ) that depends on w 0 such that
[Equation]
and further satisﬁes the norm bound
[Equation]
Proof. Our proof builds on reducing the problem from representing ( β (cid:62) x ) p via quadratic networks to representing ( β (cid:62) x ) p − 2 through a random feature model. More precisely, we consider choosing
[Equation]
where a is a real-valued random scalar that can depend on w
[Equation]
and β is the ﬁxed coefﬁcient vector in f (cid:63) . With this choice
[Equation]
the quadratic network reduces to
[Equation]
Therefore
[Equation]
to let the above express f (cid:63) ( x )
=
α ( β
[Equation]
x ) p
[Equation]
it sufﬁces to choose a such that [Equation] for all x . By Lemma
[Equation]
there exists a = a ( w 0 ) satisfying (23) and such that
[Equation]
Using this a in (22), the quadratic network induced by ( w + , w − ) has the desired expressivity, and further satisﬁes the expected 4th power norm bound
[Equation]
This is the desired result.
We begin by stating and proving the result for k = 1 in Appendix C.4.1, i.e. when f (cid:63) = α ( β (cid:62) x ) p is a single “one-directional” polynomial.
The main theorem then follows as a straightforward extension of the k = 1 case, which we prove in Appendix C.4.2.
C.4. E XPRESSING A SINGLE “ ONE - DIRECTIONAL ” POLYNOMIAL
Theorem 16 (Expressivity of f Q ) .
Suppose { ( a r , w 0 ,r ) } are generated according to the symmetric initialization (3) , and f (cid:63) ( x ) =
α ( β (cid:62) x ) p where p − 2 ∈ { 1 } ∪ { 2 (cid:96) } (cid:96) ≥ 0 .
Suppose further that we use σ ( t ) = 16 relu 3 ( t ) (so that σ (cid:48)(cid:48) ( t ) = relu( t ) ), then so long as the width is sufﬁciently large:
[Equation]
we have with probability at least 1 − δ (over W 0 ) that there exists
W (cid:63) ∈ R d × m such that
[Equation]
L Q ( W (cid:63) ) − L ( f (cid:63) ) (cid:12)(cid:12) ≤ (cid:15) and (cid:107) W (cid:63) (cid:107) 42 , 4 ≤ B 4 w,(cid:63) = O (cid:16) p 3 α 2 B 2( p − 2) x (cid:107) β (cid:107) 2 p [Equation] − 1 (cid:17)
[Equation]
Proof of Theorem
We build on the inﬁnite-neuron construction in Lemma 15. Given the symmetric initialization { w 0 ,r } mr =1 , for all r ∈ [ m/ 2] , we consider W (cid:63) ∈ R d × m deﬁned through
[Equation] where we recall ( w + ( w 0 ) , w − ( w 0 )) = ( (cid:112) a + (
[Equation]
0 ) β , (cid:112) a − ( w 0 ) β ) . We then have [Equation] Bound on (cid:107) W (cid:63) (cid:107) 2 ,
As f
[Equation]
( x ) =
α ( β (cid:62) x ) p , Lemma 15 guarantees that the coefﬁcient a
[Equation]
w 0 ) involved above satisﬁes that
[Equation]
By Markov inequality, we have with probability at least 1 − δ/ 2 that
[Equation]
which yields the bound
[Equation]
Concentration of function Let f m ( x ) = 1 m (cid:80) r ≤ m/ 2 σ (cid:48)(cid:48) ( w (cid:62) 0 ,r x ) a ( w 0 ,r ) . We now show the concentration of f m to f (cid:63),p − 2 ( x ) : = α ( β (cid:62) x ) p − 2 over the dataset { x 1 , . . . , x n } .
We perform a truncation argument: let R be a large radius (to be chosen) satisfying
[Equation]
On this event we have
[Equation]
{+Letting f R(cid:63),p − 2 ( x ) : = E w 
 [ σ (cid:48)(cid:48) ( w+}
[Equation]
0 x ) a ( w 0 ) 1 (cid:8) (cid:107) w 0 (cid:107) 2 ≤ RB − 1 x (cid:9) ] , we have
[Equation]
Applying Chebyshev inequality and a union bound, we get
[Equation] For any (cid:15) > 0 , by substituting in t = (cid:15)B − 2 x (cid:107) β (cid:107) − 2 2 / 2 , we see that [Equation]
[Equation]
ensures that
[Equation]
[Equation]
Next
[Equation]
for any x we have the bound
[Equation]
Choosing R such that
[Equation]
[Equation]
ensures that
[Equation]
Combining (35) and (37), we see that with probability at least 1 − δ ,
[-rr (cid:107) 2 
-][Equation] and thus [Equation] To satisfy the requirements for m and R in (36) and (34), we ﬁrst set R = (cid:101) O ( √ d ) (with sufﬁciently large log factor) to satisfy (36) by standard Gaussian norm concentration (cf.
Appendix A.3), and by (34) it sufﬁces to set m
as
[Equation]
for (38) to hold.
C.4. P ROOF OF MAIN THEOREM
We apply Theorem 16 k times: let
[Equation]
so that f (cid:63) = (cid:80) j ≤ k f (cid:63),j . Associate each j with an independent set of initialization ( a ( j ) 0 , W ( j ) 0 ) .
By
Theorem 16, there exists
W ( j ) (cid:63) ∈ R d × m j , where
[Equation]
such that with probability at least 1 − δ/k we have
[Equation]
and the norm bound
[Equation]
(Note we have slightly abused notation, so that now (cid:110) f Q W ( j ) (cid:63) (cid:111) j ∈ [ k ] use a disjoint set of initial weights ( a ( j ) 0 , W ( j ) 0 ) .) Concatenating all the ( W ( j ) (cid:63) , a ( j ) 0 , W ( j ) 0 ) and applying a union bound, we have the following: so long as the width m
[Equation]

 [Equation]
with probability at least 1 − δ (over a 0 ∈ R m and W 0 ∈ R d × m ), there exists W (cid:63) ∈ R d × m such that [Equation] which by the 1-Lipschitzness of the loss implies that
[Equation]
Further, as W (cid:63) is the concatenation of (cid:110) W ( j ) (cid:63) (cid:111) j ∈ [ k ] , we have the norm bound [Equation] This is the desired result.
D E XISTENCE , GENERALIZATION , AND EXPRESSIVITY OF HIGHER - ORDER NTK S
In this section we formally study the generalization and expressivity of higher-order NTKs outlined in Section 6.
Let k ≥ 2 be an integer, and recall for any W ∈ R d × 2 m the deﬁnition of the k -th order
NTK
[Equation] D. C OUPLING f AND f ( k ) VIA RANDOMZIATION Recall that for analytic σ we have the expansion [Equation]
For an arbitrary W such that (cid:107) w + ,r (cid:107) 2 ,
[Equation]
w
[Equation]
2 = o m (1) , we expect that f (1) ( x ) is the dominating term in the expansion.
Extracting the k -th order term We now describe an approach to ﬁnding W so that
[Equation]
that is, the neural net is approximately the k -th order NTK plus an error term that goes to zero as m → ∞ , thereby “escaping” the NTK regime.
Our approach builds on the following randomization technique: let z + , z − be two random variables (distributions) such that
E [ z j + ] = E [ z j − ] for j = 0 , 1 , . . , k − 1 and E [ z k + ]
= E [ z k − ] + 1 . Set ( w + ,r , w − ,r ) =
( z + ,r w (cid:63),r , z − ,r w (cid:63),r ) , and take (cid:107) w (cid:63),r (cid:107) {+
 = O ( m − 1 / 2 k )+} , we have f ( j ) W 0 , W ( x ) = 1 √ m (cid:88) r ≤ m j ! σ ( j ) ( w (cid:62) 0 ,r x ) ( z j + ,r − z j − ,r ) (cid:124) (cid:123)(cid:122) (cid:125) mean zero ( w (cid:62) (cid:63),r x ) j (cid:124) (cid:123)(cid:122) (cid:125)
O ( m − j/ 2 k ) =
O p ( m − j/ 2 k ) for all j = 1 , . . . , k − 1 , and f ( k ) W 0 , W ( x ) = 1 √ m (cid:88) r ≤ m k ! σ ( k ) ( w (cid:62) 0 ,r x ) ( z k + ,r − z k − ,r ) (cid:124) (cid:123)(cid:122) (cid:125) mean = ( w (cid:62) (cid:63),r x ) k (cid:124) (cid:123)(cid:122) (cid:125) O
( m − 1 / 2 ) =
O P (1) , and f ( k +1) W 0 , W ( x ) = 1 √ m (cid:88) r ≤ m ( k + 1)! σ ( k +1) ( w (cid:62) 0 ,r x ) ( z k +1+ ,r − z k +1 − ,r ) (cid:124) (cid:123)(cid:122) (cid:125) ( w (cid:62) (cid:63),r x ) k + (cid:124) (cid:123)(cid:122) (cid:125) O
( m − ( k +1) / 2 k ) =
O P ( m − 1 / 2 k ) . Therefore, with high probability, all f (1) , . . . ,
f ( k − 1) as well as the remainder term f − (cid:80) j ≤ k f ( j ) has order O ( m − 1 / 2 k ) , and the k -th order NTK f ( k ) can express an O (1) function.
D. G ENERALIZATION AND EXPRESSIVITY OF f ( k )
We now turn to studying the generalization and expressivity of the k -th order NTK f ( k ) , Throughout this subsection, we assume (for convenience)
that
[Equation]
is the ReLU activation. As we have seen in Section 6, we have f ( k ) = O (1) by choosing w r ∼ O ( m − 1 / 2 k ) , therefore we restrict attention on such W ’s by considering the constraint set { W : (cid:107) W (cid:107) 2 k 2 , 2 k ≤ B 2 kw } for some B w
= O m (1) .
Overview of results This subsection establishes the following results for the k -th order NTK. • We bound the generalization of f ( k ) through the tensor operator norm of a certain k -tensor involving the features (Lemma 17). Consequently, the generalization of the k -th order NTK for (cid:107) W (cid:107) 2 , 2 k ≤ B w , when the base distribution of x is uniform on the sphere, scales as
[Equation]
(Theorem 19).
Compared with the distribution-free bound B kx B kw / √ n , the leading term is better by a factor of (cid:112) min { d k − 1 , n } . In particular, when n ≥ d k − 1 , the generalization is better by a factor of √ d k − 1 than the distribution-free bound. • For the polynomial f (cid:63) ( x ) = α ( β (cid:62) x ) p with p ≥ k (and p − k is even or one), when m is sufﬁciently large, there exists a W (cid:63) expressing f (cid:63) such that
[Equation]
(Theorem 20).
Substituting into the generalization bound yields the following generalization error for learning f (cid:63) :
[Equation]
In particular, the leading multiplicative factor is the same for all k (including the linear NTK with k = 1 ), but the
sample complexity is lower by a factor of d k − 1 when n ≥ d k − 1 .
This shows systematically the beneﬁt of higher-order NTKs when distributional assumptions are present.
Tensor operator and nuclear norm Our result requires the deﬁnition of operator norm and nuclear norm for k -tensors, which we brieﬂy review here. The operator norm of a symmetric k -tensor A ∈ R d k is deﬁned as
[Equation]
The nuclear norm (cid:107)·(cid:107) ∗ is deﬁned as the dual norm of the operator norm:
[Equation]
Speciﬁcally, for any rank-one tensor u ⊗ k , we have
[Equation]
i.e. its nuclear norm equals its operator norm (and also the Frobenius norm).
D.2. G ENERALIZATION
We begin by stating a generalization bound for f ( k ) , which depends on the operator norm of a k -th order tensor feature, generalizing Lemma 5.
Lemma 17 (Bounding generalization of f ( k ) via tensor operator norm) .
For any non-negative loss (cid:96) such that z (cid:55)→ (cid:96) ( y, z ) is 1-Lipschitz and (cid:96) ( y, 0) ≤ 1 for all y ∈ Y , we have the Rademacher complexity bound E σ [Equation] x (cid:34)
[Equation]
where σ i iid ∼ Unif {± 1 } are Rademacher variables. Proof. The proof is analogous to that of Lemma 5. As the loss (cid:96) ( y, z ) is 1-Lipschitz in z for all y , by the Rademacher contraction theorem (Wainwright, 2019, Chapter 5)
we have that
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
where the last step used the power mean (or Cauchy-Schwarz) inequality on {(cid:107) w r (cid:107) 2 } .
Bound on tensor operator norm It is straightforward to see that the expected tensor operator norm can be bounded as
[Equation]
without any distributional assumptions on x . We now provide a bound on the expected tensor operator norm appearing in Lemma 17 in the special case of uniform features, i.e. x ∼ Unif( S d − 1 ( B x )) .
Lemma 18 (Tensor operator norm bound for uniform features) . Suppose x i iid ∼ Unif( S d − 1 ( B x )) .
Then for any k ≥ 3 and k = O (1) , we have (with high probability over W 0 )
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
Substituting the above bound into Lemma 17 directly leads to the following generalization bound for f ( k ) :
Theorem 19 (Generalization for f ( k ) with uniform features) .
Suppose x i iid ∼ Unif( S d − 1 ( B x )) . Then for any k ≥ 3 and k = O (1) , we have (with high probability over W 0 )
[Equation]
The proof of Lemma 18 is deferred to Appendix D.3.
Theorem 20 (Expressivity of f ( k ) ) . Suppose { ( a r , w 0 ,r ) } are generated according to the symmetric initialization (3) , and f (cid:63) ( x ) =
α ( β (cid:62) x ) p where p − k ∈ { 1 } ∪ { 2 (cid:96) } (cid:96) ≥ 0 .
Suppose further that σ is such that σ k ( t ) = relu( t ) , then so long as the width is sufﬁciently large:
[Equation]
we have with probability at least 1 − δ (over W 0 ) that there exists
W (cid:63) ∈ R d × m such that (cid:12)(cid:12) L Q ( W (cid:63) )
− L ( f (cid:63) ) (cid:12)(cid:12) ≤ (cid:15) and (cid:107) W (cid:63) (cid:107) 2 k 2 , 2 k ≤ B 2 kw,(cid:63) = O (cid:16) p 3 α 2 B 2( p − k ) x (cid:107) β (cid:107) 2 p 2 δ − 1 (cid:17) .
The proof of Theorem 20 is deferred to Appendix D.4.
D. P ROOF OF L EMMA
We begin by observing for any symmetric tensor A ∈ R d k that [Equation] where N ( (cid:15) ) is an (cid:15) -covering of unit sphere S d − 1 (1) . (The proof follows by bounding A [ u , . . .
, u ] by A [ v , . . . , v ] + k(cid:15) (cid:107) A (cid:107) op through replacing u by v one at a time). Taking (cid:15) = 1 / (2 k )
[Equation]
we have
[Equation]
We now perform a truncation argument to upper bound the above probability.
Let M > 0 be a truncation level to be determined, we have by the Bernstein inequality that
[Equation]
where the (cid:101) O (1) · B 2 kx d − k comes from computing the variance of
[Equation]
using that x i are uniform on the sphere (see, e.g. (Ghorbani et al., 2019b, Proof of Lemma 4)); MB kx is the bound on the variable Z i , and the (cid:101) O (1) comes from the fact that (cid:107) w 0 ,r (cid:107) 2 ≤ (cid:101) O ( √ dB − 1 x ) with high probability. Now, choosing M = ( nt/B kx )
[Equation]
/ 2 , the above bound reads
[Equation]
{+It remains to bound (cid:82) ∞ t =
 p t to give an expectation bound on the desired tensor operator norm. This follows by adding up the following three bounds: (1) For the main branch “ nt 2 / (cid:101) O+}
( B 2 kx d − k ) ” we have [Equation] This follows by integrating the “1” branch for t ≤ (cid:101) O ( (cid:112) B 2 kx d − ( k − 1) /n )
(which yields the right hand side) and integrating the other branch otherwise (the integral being upper bounded by (cid:101)
O ( (cid:112) B 2 kx d − k /n ) , dominated by the right hand side).
The branch “ ( nt/B kx ) 1 / 2 ” is taken only when
[Equation]
On the other hand, the inequality ( nt/B kx ) 1 / 2 > (cid:101) O
( d ) happens when
[Equation]
, which is implied by the preceding condition so long as k ≥ 3 .
Therefore, when this branch is taken, the (cid:101) O ( d ) can already be absorbed into the main term, so the contribution of this branch can be bounded as
[Equation]
(3) We have
[Equation]
using a similar argument as part (1).
Putting together the above three bounds [Equation] we obtain [Equation]
the desired result.
Our proof is analogous to that of Theorem 16, in which we ﬁrst look at the case of inﬁnitely many neurons and then use concentration to carry the result onto ﬁnitely many neurons. Expressivity with inﬁnitely many neurons We ﬁrst consider expressing f (cid:63) ( x ) = α ( β (cid:62) x ) p with inﬁnite-neuron version of f ( k ) , that is, we wish to ﬁnd random variables ( w + , w − ) such that
[Equation]
Choosing [Equation] for some real-valued random scalar a (that depends on w 0 ) [Equation]
we have
[Equation]
therefore the task reduces to ﬁnding a = a ( w 0 ) such that E w 0 (cid:2) relu( w
[Equation] x ) a [Equation] = α ( β
[Equation]
x ) p − k .
By Lemma [Equation] there exists a = a ( w 0 ) satisfying the above and such that
[Equation]
Using this a
[Equation]
the k -th order NTK deﬁned by ( w + , w − ) expresses f (cid:63) and further satisﬁes the bound
[Equation]
Finite neurons Given the symmetric initialization { w 0 ,r } mr
=1 , for all r ∈ [ m ] [Equation] we consider W (cid:63) ∈
R d × m deﬁned through
[Equation]
( w [Equation] where we recall ( w + ( w 0 ) , w − ( w 0 ))
= ( a + ( w 0 ) 1 /k β , a − ( w 0 ) 1 /k β ) .
We then have
[Equation]
r
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table] [Equation] [Table] Bound on (cid:107) W (cid:63) (cid:107) 2 , 2 k
As f (cid:63) ( x ) = α ( β (cid:62) x ) p , (32) guarantees that the coefﬁcient a ( w 0 ) involved above satisﬁes that
[Equation]
By Markov inequality, we have with probability at least 1 − δ/ 2 that
[Equation]
which yields the bound
[Equation]
Concentration of function Let f m ( x ) = 1 m (cid:80) r ≤ m σ k ( w (cid:62) 0 ,r x ) a ( w 0 ,r ) .
We now show the concentration of f m to f (cid:63),p − k ( x ) :
= α ( β (cid:62) x ) p − k over the dataset { x 1 , . . .
, x n } .
We perform a truncation argument: let R be a large radius (to be chosen) satisfying
[Equation]
On this event we have
[Equation]
Letting f R(cid:63),p − k ( x ) :
=
E w
[ σ k ( w (cid:62) 0 x ) a ( w 0 ) 1 (cid:8) (cid:107) w 0 (cid:107) 2 ≤ RB − 1 x (cid:9) ] , we have
[Equation]
Applying Chebyshev inequality and a union bound, we get
[Equation]
For any (cid:15) > 0 , by substituting in t =
(cid:15)B − k x (cid:107) β (cid:107) − k 2 / 2 , we see that
[Equation]
(34) ensures that
[Equation]
Next
[Equation]
for any x we have the bound
[Equation]
Choosing R such that
[Equation]
ensures that
[Equation]
Combining (35) and (37), we see that with probability at least 1 − δ ,
[Equation]
and thus
[Equation]
To satisfy the requirements for m and R in (36) and (34), we ﬁrst set R = (cid:101) O ( √ d ) (with sufﬁciently large log factor) to satisfy (36) by standard Gaussian norm concentration (cf.
Appendix A.3), and by (34) it sufﬁces to set m as
[Equation]
for (38) to hold.