Title: BEYOND LINEARIZATION: ON QUADRATIC AND HIGHER-ORDER APPROXIMATION OF WIDE NEURAL NETWORKS

Abstract: 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 insufficiency 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 lead 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 find networks that are coupled with higher-order terms in their Taylor series.

1. INTRODUCTION:

Deep Learning has made remarkable impact on a variety of artificial 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 efficiently 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 classification on CIFAR-10; yet there is still a non-neglible per-formance 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 W0+W (x) (with smooth activation) can be Taylor expanded with respect to the weight matrix W as

f W0+W (x) = m r=1

a r σ((w 0,r + w r ) x) = m r=1 a r σ(w 0,r x)

EQUATION

.

Above, f W0 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 boundsrather than the actual values-of f (k) W0,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 efficient 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 find 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.1. PRIOR WORK:

We review prior work on the optimization, generalization, and expressivity of neural networks.

Neural Net and Kernel Methods Neal (1996) first proposed the connection between infinitewidth 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) 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 find 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. (2017); Nacson et al. (2019) ; 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 regularization1 .

For two-layer networks, a series of works used the mean field method to establish the evolution of the network parameters via a Wasserstein gradient flow (Mei et al., 2018b; Chizat & Bach, 2018a; Wei et al., 2018; Rotskoff & Vanden-Eijnden, 2018; Sirignano & Spiliopoulos, 2018) . In the mean field regime, the parameters move significantly 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 amplifies 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 find 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 efficient generalization.

2. PRELIMINARIES:

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 find 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 : Y × R → R ≥0 be a loss function such that (y, 0) ≤ 1, and z → (y, z) is convex, 1-Lipschitz, 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 classification. We let

L(f ) := E D [ (y, f (x))] := 1 n n i=1 (y i , f (x i )) and L P (f ) := E (x,y)∼P [ (y, f (x))]

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 first layer and a = [a 1 , . . . , a m ] ∈ 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 fix 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 σ ∈ C2 (R), and there exists some absolute constant

C > 0 such that |σ (t)| ≤ Ct 2 , |σ (t)| ≤ C|t|, and σ (•) 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.1. NOTATION:

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 defined as A ¯ 2,p := ( 

3. ESCAPING NTK VIA RANDOMIZATION:

To motivate our study, we now briefly 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 W0+W can be Taylor expanded as

f W0+W (x) = 1 √ m r≤m a r σ((w 0,r + w r ) x) = 1 √ m r≤m a r σ(w 0,r x) f W 0 (x) + 1 √ m r≤m a r σ (w 0,r x)(w r x) :=f L W (x) +O   1 √ m r≤m (w r x) 2   ,

so that the network can be decomposed as the sum of the initial network f W0 , the linearized model f L W , and higher order terms. Specifically (ignoring f W0 for the moment), when m is large and w r 2 = 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 W0+W via gradient descent. Therefore, the trajectory of training f W0+W is coupled with the trajectory of training f W0 + f L W , which is a convex problem and enjoys convergence guarantees (Du et al., 2018b) .

Our goal is to find 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 fixes 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 fixed weight matrix. Suppose for each weight vector w r , we sample a random variable Σ rr ∈ R and consider instead the random weight matrix

WΣ := Wdiag({Σ rr } m r=1 ) = [Σ 11 w 1 , . . . , Σ rr w r ],

then the second-order Taylor expansion of f W0+WΣ can be written as

f W0+WΣ (x) = 1 √ m m r=1 a r σ (w 0,r + w r Σ rr ) x = f W0 (x) + 1 √ m m r=1 a r σ (w 0,r x)(Σ rr w r x) =f L WΣ (x) + 1 2 √ m m r=1 a r σ (w 0,r x)Σ 2 rr (w r x) 2 :=f Q WΣ (x) + . . . ,

where we have defined in addition the quadratic part f Q WΣ . Due to the existence of {Σ rr }, each original weight w r now has an additional a scalar that is different in f L and f Q . Specifically, if we choose Σ rr iid ∼ Unif{±1}

(2) 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 w r 2 m -1/4 , the scalings of f L and f Q compare as follows:

• We have E Σ [f L WΣ (x)] = 0 and E Σ (f L WΣ (x)) 2 = 1 m m r=1 a 2 r σ (w 0,r x) 2 (w r x) 2 = O 1 m m r=1 w r 2 2 = O(m -1/2 ), so we expect f L WΣ (x) = O(m -1/4

) over a random draw of Σ. • The quadratic part scales as

f Q WΣ (x) = f Q W (x) = 1 2 √ m m r=1 a r σ (w 0,r x)(w r x) 2 = O 1 √ m m r=1 w r 2 2 = O(1).

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.1. LEARNING RANDOMIZED NEURAL NETS:

The randomization technique leads to the following recipe for learning W: train W so that w r 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 L : R d×m → R denote the vanilla empirical risk for learning f W :

L(W) = E D [ (y, f W0+W (x))] ,

where we have reparametrized the weight matrix into W 0 + W so that learning starts at W = 0.

Following our randomization recipe, we now formulate our problem as minimizing the expected risk

L(W) := E Σ [ L(WΣ)] = E Σ,D [ (y, f W0+WΣ (x))] ,

where Σ ∈ R m×m is a diagonal matrix with Σ rr iid ∼ Unif{±1}. To encourage w r 2 = O(m -1/4 ) and improve generalization, we consider a regularized version of L and L with 2,4 regularization:

L λ (W) := L(W) + λ W 8 2,4 and L λ (W) = L(W) + λ W 8 2,4 = E Σ [ L λ (WΣ)].

Our regularizer penalizes W, i.e. the distance from initialization, similar as in (Hu et al., 2019). 3 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 benefit that f W0 (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 0 x ∼ N(0, 1), which is on average O(1) 4 . For technical convenience, we also assume henceforth that the realized {w 0,r } satisfies the bound

EQUATION

This happens with probability at least 1 -δ under random initialization (see proof in Appendix A. 

4. OPTIMIZATION:

In this section, we show that L λ enjoys a nice optimization landscape.

4.1. NICE LANDSCAPE OF CLEAN RISK:

As the randomized loss L induces coupling of the neural net f W0+WΣ with the quadratic model f Q W , we expect its behavior to resemble the behavior of gradient descent on the following clean risk:

L Q (W) := 1 n n i=1 (y i , f Q W (x i )) = 1 n n i=1 y i , 1 2 √ m x i x i , WD i W .

Above, we have defined diagonal matrices

D i = diag( a r σ (w 0,r x i ) 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 ∈ R d×m such that L Q (W ) ≤ OPT.

Let Σ ∈ R m×m be a diagonal matrix with Σ rr iid ∼ Unif{±1}, then we have

E Σ ∇ 2 L Q (W)[W Σ , W Σ ] ≤ ∇L Q (W), W -2(L Q (W) -OPT) + O dB 4 x W 2 2,4 W 2 2,4 m -1 .

(5)

This result implies that, for W in a certain ball and large m, every point of higher loss than W will have either a first-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.2. NICE LANDSCAPE OF RANDOMIZED NEURAL NET RISK:

With the coupling between f W0+WΣ (x) and f Q W (x) in hand, we expect the risk L(W) = E[ 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 ∈ B 2,4 (B w, ) such that L Q (W ) ≤ OPT, and that m ≥ O B 12 x B 12 w + d 4 B 4 x B 4 w + d 2 B 20 x B 20 w -4 + d 5 B 8 x B 8 w -2 .

(6) for some fixed ∈ (0, 1] and B w ≥ B w, , 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, , 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 such that

EQUATION

for some OPT > 0. For any γ = Θ(1) and > 0, we can choose λ suitably and m ≥ O(poly(d, B x B w, , -1 )) such that the regularized loss L λ satisfies the following: any second order stationary point W has low loss and bounded norm:

EQUATION

Proof sketch. The proof of Theorem 4 consists of two stages: first "localize" any second-order stationary point into a (potentially very big) norm ball using the • 8 2,4 regularizer, then use Corollary 3 in this ball to further deduce that L λ is low and

W 2,4 ≤ O( W 2,4 ). The full proof is deferred to Appendix B.6.

Efficient optimization & allowing large learning rate Theorem 4 states that when the overparametrization is enough, any second-order stationary point (SOSP) W of L λ has loss competitive with OPT, the performance of best quadratic model. Consequently, algorithms that are able to find SOSPs (escape saddles) such as noisy SGD (Jin et al., 2019) can efficiently 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 L λ (WΣ) (rather than computing the full ∇L λ (W)), the noisy SGD iterates can be computed efficiently. We note in passing that our coupling results work in any 2,4 ball of O(1) size further allows the use of a large learning rate: as soon as the learning rate is bounded by O(m 1/4 ), we would have W t 2,4 ≤ O(1) in a constant number of iterations, and thus our coupling and landscape results would hold. This is in contrast with the NTK regime which requires the learning rate to be bounded by O(1) (Du et al., 2018b) .

5. GENERALIZATION AND EXPRESSIVITY:

We now shift attention to studying the generalization and expressivity of the (randomized) neural net W learned in Theorem 4.

5.1. GENERALIZATION:

As W is always coupled (through randomization) with the quadratic model f Q W , we begin by studying the generalization of the quadratic model.

Generalization of quadratic models Let:

F Q (B w ) := x → f Q W (x) : W 2,4 ≤ B w

denote the class of quadratic models for W in a 2,4 ball. We first 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 such that z → (y, z) is 1-Lipschitz and (y, 0) ≤ 1 for all y ∈ Y, we have the Rademacher complexity bound

E σ,x sup W 2,4 ≤Bw 1 n n i=1 σ i (y i , f Q W (x i )) ≤ B 2 w E σ,x   max r∈[m] 1 n n i=1 σ i σ (w 0,r x i )x i x i op   + 1 √ n ,

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" 1 n n i=1 σ i σ (w 0,r x i )x i x 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 W such that W 2,4 ≤ B w , we have

E W0,D L( W) -L P ( W) ≤ O B 2 x B 2 w M x,op √ n + 1 √ n + O B 3 x B 3 w m -1/4 + d 2 B 2 x B 2 w m -1/2 , where M x,op := B -2 x E x 1 n n i=1 x i x i op 1/2

is the (rescaled) operator norm of the empiri-

cal covariance matrix. In particular, M x,op ≤ 1 always holds; if in addition v x is K Var(v 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 n 4 ), the bound scales at most logarithmically with the width m, therefore allowing learning with small samples and extreme over-parametrization; Theorem 6 follows directly from Lemma 5 and a matrix concentration Lemma. The proof is deferred to Appendix C.2.

5.2. EXPRESSIVITY AND SAMPLE COMPLEXITY THROUGH QUADRATIC MODELS:

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 j≤k α j (β j x) pj , 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 σ (t) = relu(t).)

If f (x) = k j=1 α j (β j x) pj achieves training loss L(f ) ≤ 0 , where p j -2 ∈ {1} ∪ {2 } ≥0 ,.

Then so long as the width is sufficiently large:

m ≥ O   ndk 2 k j=1 p 3 j α 2 j (B x β j 2 ) 2pj -2   ,

we have with probability at least 1 -δ (over W 0 ) that there exists W ∈ R d×m such that

L Q (W ) ≤ OPT := 0 + and W 4 2,4 ≤ B 4 w, = O   k k j=1 p 3 j α 2 j B 2(pj -2) x β 2pj 2 δ -1   .

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 sufficiently large.

Learning a single polynomial. Suppose f (x) = α(β x) p satisfies L(f ) ≤ , and we wish to find W with O( ) test loss. By Theorem 7 we can choose W such that L Q (W ) ≤ OPT = 2 , and by Theorem 4 we can find W such that L( W) ≤ L λ ( W) ≤ 3 and W 2,4 = O(B w, ). Take B x = 1, and assume x is sufficiently isotropic so that

M x,op = O( 1 √ d ), the sample complexity from Theorem 6 is n ≥ O B 4 x B 4 w M 2 x,op 2 = O p 3 α 2 β 2p 2 d 2 := n Q .

In contrast, the sample complexity for linear NTK (Arora et al., 2019b; Cao & Gu, 2019a) to reach test loss is

n ≥ O p 2 α 2 β 2p 2 2 := n L .

We have n Q /n L = O(p/d), a reduction by a dimension factor unless p 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 on the noisy 2-XOR problem, which allows for a rigorous comparison against the quadratic model.

The ground truth function in:

2-XOR is f (x) = x 1 x 2 = ([(e 1 + e 2 ) x] 2 -[(e 1 -e 2 ) x] 2 )/4

, where x ∈ {±1} d , and f attains constant margin on the training distribution constructed in Wei et al. (2019) . By Theorem 7, f can be -approximated by f Q W with B 4 w, ≤ O(1). Thus by Theorem 6 the sample complexity for learning noisy 2-XOR through the randomized net W is

n ≥ n Q = O B 4 x B 4 w, M 2 x,op 2 = O d 2 .

This is 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 ∈ R d×d with A op ≤ 1 and rank(A ) ≤ r through n rank-one observations of the form

y i = f (x i ) = A , x i x i where x ∼ Unif(S d-1 ( √ d))

. This ground truth function can be written as f (x i ) = r j=1 α j (v 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

EQUATION

2 ) = O(r 2 ). Thus by Theorem 6, the sample complexity of reaching 2 test loss

E x A , xx -f W (x) ≤ 2 through the randomized net W is n ≥ n Q = O B 4 x B 4 w, M 2 x,op 2 = O dr 2 2 .

This compares favorably against the sample complexity upper bound for linear NTK, which needs

n ≥ n L = O B 4 x • ( r j=1 α j v j 2 2 ) 2 2 = O d 2 r 2 2 samples.

6. HIGHER-ORDER NTKS:

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 overparametrized neural networks in which the k-th order term dominates the Taylor series. Due to space constraints, we provide our concrete results on the coupling, generalization, and expressivity of higher-order models in Appendix D.

7. CONCLUSION:

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 efficient 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 

σ i A r,i ,

then we have

E σ max r∈[m] Y r op ≤ 4 max r∈[m] v(Y r ) log(2md), where v(Y) := E σ [Y 2 ] op .

Proof. Applying the high-probability bound in (Tropp et al., 2015, Theorem 4.6 .1) and the union bound, we get

P max r∈[m] Y r op ≥ t ≤ 2d r∈[m] exp(-t 2 /2v(Y r )) ≤ 2dm exp(-t 2 /2 max r∈[m] v(Y r )) = exp - t 2 2 max r∈[m] v(Y r ) + log(2dm) . Let V := max r∈[m] v(Y r )

, we have by integrating the above bound over t that

E max r∈[m] Y r op ≤ ∞ 0 min exp - t 2 2V + log(2dm) , 1 dt ≤ 4V log(2dm) + ∞ √ 4V log(2dm) exp(-t 2 /2V + log(2dm))dt ≤ 4V log(2dm) + ∞ √ 4V log(2dm) exp(-t 2 /4V )dt ≤ 4V log(2dm) + √ 4πV 2dm ≤ 4 V log(2dm).

A.2 EXPRESSING 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 } ≥0 and β ∈ R d , there exists a random variable a = a(w 0 ) such that E w0 [σ(w 0 x)a] = α(β x) p and a satisfies the 2 norm bound

E w0 [a 2 ] ≤ 2π(p ∨ 1) 3 α 2 B 2(p-1) x d β 2p 2 .

Proof. Consider the ReLU random feature kernel

K(x, x ) = E w0∼N(0,B -2 x I d ) [relu(w 0 x)relu(w 0 x )],

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 ) → H (where H is a Hilbert space) that generates K in the sense that K(x, x ) = φ(x), φ(x ) H , we have for any function f that

EQUATION

and the infimum over a is attainable whenever it is finite.

For the ReLU random feature kernel K, let u := x x /B 2 x and N 2 (ρ) denote a bivariate normal distribution with marginals N(0, 1) and correlation ρ ∈

[-1, 1]. We have that K(x, x ) = E w0∼N(0,B -2 x I d ) relu(w 0 x)relu(w 0 x ) = E (Z1,Z2)∼N2(u) [relu(Z 1 )relu(Z 2 )] = 1 2π u(π -arccos u) + 1 -u 2 = 1 2π 1 + π 2 u + ∞ =1 (2 -3)!! (2 -2)!!(2 -1)(2 ) u 2 = p∈{0,1}∪{2 } ≥1 c p (x x ) p B -2p x = p∈{0,1}∪{2 } ≥1 √ c p B -p x x ⊗p , √ c p B -p x (x ) ⊗p ,

where the constants {c p } satisfy

c 0 = 1/(2π), c 1 = 1/4, c 2 ≥ 1 2π(2 -1) 2 (2 )

for ≥ 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 define feature map φ(x) := √ c p B -p x x ⊗p p∈{0,1}∪{2 } ≥1 , we have K(x, x ) = φ(x), φ(x )

. With this feature map, the function f (x) = α(β x) p can be represented as

f (x) ≡ c , φ(x) where c = 0, . . . , 0, α • c -1/2 p B p

x β ⊗p , 0, . . . . Thus by the feature map equivalence (11), we have f ∈ H K and f 2

EQUATION

2 . Now apply the feature map equivalence (11) again with the random feature map

x → relu(w 0 x) w0

(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 = E w0 [relu(w 0 x)a] and

E w0 [a 2 ] ≤ f 2 H K ≤ 2π(p ∨ 1) 3 α 2 B 2p x β 2p 2 .

A.3 PROOF OF EQUATION (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 w 2 ≤ 2 sup v∈N (v w) (see e.g. (Mei et al., 2018a , Section A).) We thus have Computing the gradient of L Q , we obtain

P max r∈[m] B x w 0,r 2 ≥ t ≤ max v∈N B x (v w 0,r ) ≥ t/2 ≤ exp(-t 2 /8 + log |N | + log m) ≤ exp(-t 2 /8 + d log 5 + log m).

∇L Q (W) = 2 n n i=1 (y i , f Q W (x i )) 1 2 √ m x i x i WD i .

Further computing the Hessian gives

∇ 2 L Q (W)[W Σ , W Σ ] = 2 n n i=1 (y i , f Q W (x i )) • 1 2 √ m x i x i , W Σ D i Σ W f Q W (xi) + 4 n n i=1 (y i , f Q W (x i )) • 1 2 √ m x i x i , WD i W Σ := yi 2 = 2 n n i=1 (y i , f Q W (x i ))f Q W (x i ) I + 4 n n i=1 (y i , f Q W (x i )) y 2 i II .

Taking expectation over Σ , and using that ≤ 1, term II can be bounded as

E Σ [II] ≤ E Σ 4 n n i=1 y 2 i = E Σ ,D   2 m r≤m σ (w 0,r x) 2 (w r x) 2 (Σ rr w ,r x) 2   ≤ C • E D   1 m r≤m (w 0,r x) 2 (w r x) 2 (w ,r x) 2   ≤ CB 4 x max r∈[m],i∈[n] (w 0,r x i ) 2 • 1 m r≤m w r 2 2 w ,r 2 2 ≤ O dB 4 x W 2 2,4 W 2 2,4 m -1 ,

where the last step used Cauchy-Schwarz on { w r 2 } and w ,r 2 .

Published as a conference paper at ICLR 2020 Term I does not involve Σ and can be deterministically bounded as

I = 2E D [ (y, f Q W (x))f Q W (x)] = 2E D [ (y, f Q W (x))f Q W (x)] + 2E D [ (y, f Q W (x))(f Q W (x) -f Q W (x))] (i) ≤ ∇L Q (W), W + 2E D [ (y, f Q W (x)) -(y, f Q W (x))] (ii) = ∇L Q (W), W -2(L Q (W) -OPT).

where (i) follows directly by computing ∇L Q (W), W and the convexity of z → (y, z), and (ii) follows from the assumption that L Q (W ) ≤ OPT. Combining the bounds for terms I and II gives the desired result.

B.2 COUPLING LEMMAS:

Lemma 10 (Bound on f Q ). For any W ∈ R d×m , the quadratic model f Q W satisfies the bound

f Q W (x) ≤ O √ dB 2 x W 2 2,4

for all x ∈ S d-1 (B x ).

Proof. We have

f Q W (x) = 1 √ m r≤m a r σ (w 0,r x)(w r x) 2 ≤ 1 √ m r≤m C w 0,r x • (w r x) 2 ≤ C √ mB 2 x max r∈[m] w 0,r x • 1 m r≤m w r 2 2 ≤ C √ mB 2 x O( √ d) •   1 m r≤m w r 4 2   1/2 = O √ dB 2 x R w,0 W 2 2,4 .

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 ] ≤ O(d 2 B 2 x W 2 2,4 m -1/2 ). (b) |∆ Q WΣ (x)| ≤ O(B 3 x W 3 2,4 m -1/4 ) (almost surely for all Σ.)

Proof.

(a) Recall that

f L WΣ (x) = 1 √ m r≤m a r σ (w 0,r x)(Σ rr w r x).

As Σ rr has mean zero, we have

E Σ [f L ] = 0 and E Σ [(f L ) 2 ] = 1 m r≤m a 2 r σ (w 0,r x) 2 (w r x) 2 (i) ≤ 1 m r C(w 0,r x) 4 (w r x) 2 (ii) ≤ C max r∈[m] (w 0,r x) 4 • 1 m r B 2 x w r 2 2 (iii) ≤ O(d 2 B 2 x ) • 1 m r≤m w r 2 2 (iv) ≤ CB 6 x R 4 w,0 •   1 m r≤m w r 4 2   1/2 = O d 2 B 2 x W 2 2,4 m -1/2 .

Above, (i) follows from the assumption that |σ (t)| ≤ Ct 2 , (ii) is Cauchy-Schwarz, (iii) uses the bound (4), and (iv) uses the power mean inequality on w r 2 .

(b) We have by the Lipschitzness of σ that

|∆ Q W (x)| = 1 √ m r≤m a r σ((w 0,r + Σ rr w r ) x) -σ(w 0,r x) -σ (w 0,r x)(Σ rr w r x) -σ (w 0,r x)(Σ rr w r x) 2 ≤ 1 √ m r≤m C|Σ rr w r x| 3 (i) ≤ C √ mB 3 x • 1 m r≤m w r 3 2 ≤ C √ mB 3 x •   1 m r≤m w r 4 2   3/4 = O B 3 x W 3 2,4 m -1/4 ,

where again (i) uses the power mean inequality on w r 2 .

B.3 CLOSENESS OF LANDSCAPES:

Lemma 12 (L Q close to L). We have for all W ∈ R d×m that

|L(W) -L Q (W)| ≤ O B 3 x W 3 2,4 m -1/4 + d 2 B 2 x W 2 2,4 m -1/2 .

Proof. Recall that

L(W) = E Σ,D [ (y, f W0+WΣ (x))] and L Q (W) = E D [ (y, f Q W (x))]

. By the 1-Lipschitzness of z → (y, z) we have

L(W) -L Q (W) ≤ E Σ,D f W0+WΣ (x) -f Q W (x) ≤ E Σ,D (f L WΣ (x) + ∆ Q WΣ (x)) 2 1/2 ≤ 2E Σ,D (f L WΣ (x)) 2 + 2E Σ,D (∆ Q WΣ (x)) 2 1/2 = O B 3 x W 3 2,4 m -1/4 + d 2 B 2 x W 2 2,4 m -1/2 ,

where the last step uses Lemma 11.

Lemma 13 (Closeness of directional gradients). We have

∇L(W), W -∇L Q (W), W ≤ O dB x W 2,4 + √ dB 5 x W 5 2,4 + B 3 x W 3 2,4 m -1/4 + d 2.5 B 4 x W 4 2,4 m -1/2 .

Proof. Differentiating L and L Q and taking the inner product with W, we get

∇L(W), W = E Σ,D   (y, f W0+WΣ (x)) • 1 √ m r≤m a r σ ((w 0,r + Σ rr w r ) x)(Σ rr w r x)   ,

and

∇L Q (W), W = E D   (y, f Q W (x)) • 1 √ m r≤m a r σ (w 0,r x) • (w r x) 2   .

Therefore, by expanding σ ((w 0,r + Σ rr w r ) x) and noticing that Σ 2 rr ≡ 1, we have

∇L(W) -∇L Q (W), W = E Σ,D   (y, f W0+WΣ (x)) • 1 √ m r≤m a r σ (w 0,r x)(Σ rr w r x)   I + E Σ,D   (y, f W0+WΣ (x)) -(y, f Q W (x)) • 1 √ m r≤m a r σ (w 0,r x) • (w r x) 2   II + E Σ,D (y, f W0+WΣ (x)) • 1 √ m r≤m a r σ ((w 0,r + Σ rr w r ) x) -σ (w 0,r x) -σ (w 0,r x)(Σ rr w r x) (Σ rr w r x) III .

We now bound the three terms separately. Recall that | | ≤ 1 and (y, z) is 1-Lipschitz in z. For term I we have by Cauchy-Schwarz that

|I| ≤   E D   1 m r≤m a 2 r σ (w 0,r x) 2 (w r x) 2     1/2 ≤   C max r∈[m],i∈[n] (w 0,r x i ) 4 • 1 m r≤m w r 2 2 B 2 x   1/2 ≤ O(dB x ) •   1 m r≤m w r 4 2   1/4 = O dB x W 2,4 m -1/4 .

For term II, we have

|II| (i) ≤ E Σ,D (f W0+WΣ (x) -f Q W (x)) 2 1/2 • E D (f Q W (x)) 2 1/2 (ii) ≤ O B 3 x W 3 2,4 m -1/4 + d 2 B 2 x W 2 2,4 m -1/2 • O( √ dB 2 x W 2 2,4 ) = O √ dB 5 x W 5 2,4 m -1/4 + d 2.5 B 4 x W 4 2,4 m -1/2 .

where (i) uses Cauchy-Schwarz and (ii) uses the bounds in Lemma 10 and 11. For term III we first note by the smoothness of σ that a r σ ((w 0,r + Σ rr w r ) x) -σ (w 0,r x) -σ (w 0,r x)(Σ rr w r x) (Σ rr w r x)

≤ C Σ rr w r x 3 ≤ CB 3 x w r 3 2 . Substituting this bound into term III yields |III| ≤ 1 √ m r≤m CB 3 x w r 3 2 ≤ C √ mB 3 x • 1 m r≤m w r 3 2 ≤ C √ mB 3 x •   1 m r≤m w r 4 2   3/4 = O B 3 x W 3 2,4 m -1/4 .

Putting together the bounds for term I, II, III gives the desired result.

Lemma 14 (Closeness of Hessians). Let Σ denote a diagonal matrix with diagonal entries drawn i.i.d. from Unif{±1}. We have for all W, W ∈ R d×m that

E Σ ∇ 2 L(W) -∇ 2 L Q (W) [W Σ , W Σ ] ≤ O B 3 x W 2,4 + √ dB 5 x W 3 2,4 W 2 2,4 m -1/4 + d 2.5 B 4 x W 2 2,4 W 2 2,4 + B 2 x (d 2 + W 4 2,∞ B 4 x ) W 2 2,4 m -1/2 + dB 4 x W 2 2,4 W 2 2,4 m -1 .

Proof. Differentiating L and L Q twice on the direction W Σ , we get

∇ 2 L(W)[W Σ , W Σ ] = E Σ,Σ ,D    (y, f W0+WΣ (x)) •   1 √ m r≤m a r σ ((w 0,r + Σ rr w r ) x)(Σ rr Σ rr w ,r x)   2    I(L) + E Σ,D   (y, f W0+WΣ (x)) • 1 √ m r≤m a r σ ((w 0,r + Σ rr w r ) x)(Σ rr Σ rr w ,r x) 2   II(L)

, and

∇ 2 L Q (W)[W Σ , W Σ ] = E Σ ,D    (y, f Q W (x)) •   1 √ m r≤m a r σ (w 0,r x)(w r x)(Σ rr w ,r x)   2    I(L Q ) + E Σ ,D   (y, f Q W (x)) • 1 √ m r≤m a r σ (w 0,r x)(Σ rr w ,r x) 2   II(L Q )

.

We first bound the terms I(L) and I(L Q ). We have

EQUATION

2,4 m -1/2 . 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

2 .

Then we have, by applying the bounds in Lemma 10 and 11,

II(L) -II(L Q ) = E Σ,D   (y, f W0+WΣ (x)) • 1 √ m r≤m ∆ σ r (x)   + E Σ,D (y, f Q W0+WΣ (x)) -(y, f Q W (x)) • 2f Q W (x) ≤ C • 1 √ m r≤m B 3 x w r 2 w ,r 2 2 + C E f L WΣ (x) + ∆ Q WΣ (x) 2 • E f Q W (x) 2 1/2 ≤ O B 3 x W 2,4 W 2 2,4 m -1/4 + O B 3 x W 3 2,4 m -1/4 + d 2 B 2 x W 2 2,4 m -1/2 • O √ dB 2 x W 2 2,4 . = O B 3 x W 2,4 + √ dB 5 x W 3 2,4 W 2 2,4 m -1/4 + d 2.5 B 4 x W 2 2,4 W 2 2,4 m -1/2 .

Combining all the bounds gives the desired result.

B.4 PROOF OF THEOREM 2:

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

)

Applying Lemma 1, we obtain that

E Σ ∇ 2 L Q (W)[W Σ , W Σ ] -∇L Q (W), W ≤ 2(L Q (W) -OPT) + /3 ≤ 2(L(W) -OPT) + 2 /3 (14)

provided that the error term in Lemma 1 is bounded by /3, which happens when

EQUATION

)

Finally, we choose m sufficiently large so that

E Σ (∇ 2 L(W) -∇ 2 L Q (W))[W Σ , W Σ ] ≤ /6

and

∇L(W) -∇L Q (W), W ≤ /6,

which combined with ( 14) yields the desired result. By Lemma 13 and 14, it suffices to choose m such that, to satisfy the closeness of directional gradients,

EQUATION

)

and to satisfy the closeness of Hessian quadratic forms,

EQUATION

)

Collecting the requirements on m in ( 13), ( 15), ( 16), ( 17) and merging terms using ≤ 1 and B w, ≤ B w , the desired result holds whenever

m ≥ O B 12 x B 12 w + d 4 B 4 x B 4 w + d 2 B 20 x B 20 w -4 + d 5 B 8 x B 8 w -2 .

This completes the proof.

B.5 PROOF OF COROLLARY 3:

Proof. For all λ ≥ 0 define

A λ := E Σ ∇ 2 L λ (W)[W Σ , W Σ ] -∇L λ (W), W + 2(L λ (W) -OPT),

By Lemma 2, it suffices to show that

A λ -A 0 ≤ Cλ W 8 2,4 -λ W 8 2,4

for some absolute constant C.

Recall that L λ (W) = L(W) + λ W 8 2,4 . By differentiating A → A 8 2,4 we get ∇(L λ -L)(W), W = 2 W 4 2,4 • r≤m 4λ w r 2 2 w r , w r = 8λ W 8 2,4

and

∇ 2 (L λ -L)(W)[W Σ , W Σ ] = 8λ W 4 2,4   r≤m w r 2 2 w ,r Σ rr 2 2 + 2 w r , w ,r Σ rr 2   + 32λ W 4 2,4 •   r≤m w r 2 2 w r , w ,r Σ rr   2 ≤ 56λ W 4 2,4 r≤m w r 2 2 w ,r 2 2 (i) ≤ 56λ W 6 2,4 W 2 2,4 (ii) ≤ 14λα W 8 2,4 + 378λ α 3 W 8 2,4 ,

where (i) used Cauchy-Schwarz and (ii) used the AM-GM inequality p 3 q ≤ αp 4 /4 + 27q 4 /(4α 3 ) for all p, q and α > 0. Substituting the above expressions into A λ -A 0 yields

A λ -A 0 ≤ 14λα W 8 2,4 + 378λ α 3 W 8 2,4 -8λ W 8 2,4 + 2λ W 8 2,4 = (14λα -6λ) W 8 2,4 + 378λ α 3 W 8 2,4 .

Choosing α = 5/14 gives the desired result.

B.6 PROOF OF THEOREM 4:

We begin by choosing the regularization strength as

λ = λ 0 B -8 w,

, where λ 0 is a constant to be determined. Let be an accuracy parameter also to be determined.

Localizing second-order stationary points We first argue that any second order stationary point W has to satisfy W 2,4 ≤ B w,0 for some large but controlled B w,0 . We first note that for the clean risk L Q , we have for any

W ∈ R d×m that ∇L Q (W), W = E D (y, f Q W (x)) • 2f Q W (x) = 2E D (y, f Q W (x)) • (f Q W (x) -f Q 0 (x)) (i) ≥ 2(L Q (W) -L Q (0)) (ii) ≥ -2,

where (i) uses convexity of and (ii) uses the assumption that (y, 0) ≤ 1 for all y ∈ Y.

Now, applying the coupling Lemma 13, and combining with the fact that

∇ W (λ W 8 2,4 ), W = 8λ W 8 2,4 , we have simultaneously for all W that ∇L λ (W), W ≥ ∇ W (λ W 8 2,4 ), W + ∇L Q (W), W -∇(L -L Q )(W), W ≥ 8λ W 8 2,4 -2 -O dB x W 2,4 + √ dB 5 x W 5 2,4 + B 3 x W 3 2,4 m -1/4 + d 2.5 B 4 x W 4 2,4 m -1/2 .

Therefore we see that any stationary point W has to satisfy

W 2,4 ≤ B w,0 := O λ -1/8 + (λ -1 dB x m -1/4 ) 1/7 + (λ -1 √ dB 5 x m -1/4 ) 1/3 + (λ -1 B 3 x ) 1/5 + (λ -1 d 2.5 B 4 x m -1/2 ) 1/4 .

By Corollary 3, choosing m ≥ poly(λ -1 0 , d, B w, B x , ), the coupling error is bounded by 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:

λ 0 = 1 C • (2γOPT + ),

we get that CλB 8 w, = 2γOPT + , and thus the bound ( 18) reads

E Σ ∇ 2 L λ (W)[W Σ , W Σ ] ≤ ∇L λ (W), W -2(L λ (W) -OPT) -λ W 8 2,4 + 2γOPT + 2 .

For the second-order stationary point W, the gradient term vanishes and the Hessian term is nonnegative, so we get

2(L λ ( W) -OPT) ≤ 2(γOPT + ) -λ W 8 2,4 ≤ 2(γOPT + )

and thus

L λ ( W) ≤ (1 + γ)OPT + .

Further, by re-writing (18), we obtain

λ W 8 2,4 ≤ CλB 8 w, + 2(OPT -L λ ( W)) + ≤ CλB 8 w, + 2OPT + ≤ CλB 8 w, • 1 + 2OPT + 2γOPT + = O(1) • λB 8 w, ,

for any γ = O(1). This is the desired result.

C PROOFS FOR SECTION 5:

C.1 PROOF OF LEMMA 5

As the loss (y, z) is 1-Lipschitz in z for all y, by the Rademacher contraction theorem (Wainwright, 2019, Chapter 5) we have that

E σ,x sup W 2,4 ≤Bw 1 n n i=1 σ i (y i , f Q W (x i )) ≤ 2E σ,x sup W 2,4 ≤Bw 1 n n i=1 σ i f Q W (x i ) + E σ,x 1 n n i=1 σ i (y i , 0) ≤ E σ,x   sup W 2,4 ≤Bw 1 √ m r≤m 1 n n i=1 σ i a r σ (w 0,r x i )x i x i , w r w r   + 1 √ n ≤ E σ,x   sup W 2,4 ≤Bw max r∈[m] 1 n n i=1 a r σ i σ (w 0,r x i )x i x i op • 1 √ m r≤m w r w r *   + 1 √ n ≤ E σ,x   max r∈[m] 1 n n i=1 σ i σ (w 0,r x i )x i x i op   • sup W 2,4 ≤Bw 1 √ m r≤m w r 2 2 ≤B 2 w + 1 √ n ,

where the last step used the power mean (or Cauchy-Schwarz) inequality on { w r 2 }.

C.2 PROOF OF THEOREM 6:

We first relate the generalization of L to that of L Q through

L P ( W) -L( W) ≤ L P ( W) -L Q P ( W) + L Q P ( W) -L Q ( W) + L Q ( W) -L( W).

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 W. Therefore it remains to bound

L Q P ( W) - L Q ( W), i.e.

the generalization of the quadratic model.:

Generalization of quadratic model By symmetrization and applying Lemma 5, we have

EQUATION

We now focus on bounding the expected max operator norm above. First, we apply the matrix concentration Lemma 8 to deduce that

E W0,σ,x   max r∈[m] 1 n n i=1 σ i σ (w 0,r x i )x i x i op   ≤ 4 log(2dm) • E W0,x   max r∈[m] 1 n 2 n i=1 σ (w 0,r x i ) 2 x i 2 2 x i x i op   ≤ 4B x log(2dm) n • E W0,x   max r,i σ (w 0,r x i ) 2 • 1 n n i=1 x i x i op   ≤ 4B x log(2dm) n   E W0,x max r,i σ (w 0,r x i ) 2 • E x   1 n n i=1 x i x i op     1/2

As |σ (t)| ≤ Ct and w 0,r x i ∼ N(0, 1) for all (r, i), by standard expected max bound on subexponential variables we have

E W0,x max σ (w 0,r x i ) 2 ≤ O(log(mn)) = O(1). Therefore defining M x,op := B -2 x • E x 1 n n i=1 x i x i 1/2

, and substituting the above bound into (21) yields that

E W0,D L Q P ( W) -L Q ( W) ≤ O B 2 x B 2 w M x,op √ n + 1 √ n .

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

i≤n x i x i /n 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

±B x / √ d d

, then we have by standard covariance concentration (Vershynin, 2018, Theorem 4.7 .1) that

E x i≤n x i x i /n op ≤ B 2 x /d • O(1 + d/n + d/n) = O(B 2 x /d) when n ≥ d. More generally, if for all v ∈ S d-1 (1) we have v x ψ2 ≤ K v Cov(x)v,

and that κ(Cov(x)) ≤ κ, then we have Cov(x)) op ≤ κB 2

x /d. Applying (Vershynin, 2018, Theorem 4.7 .1), we get

M x,op ≤ κ/ √ d whenever n ≥ O(K 4 d).

C.3 EXPRESSIVE POWER OF INFINITELY WIDE QUADRATIC MODELS:

Lemma 15 (Expressivity of f Q with infinitely many neurons). Suppose f (x) = α(β x) p for some α ∈ R, β ∈ R d , and p ≥ 2 and such that p -2 ∈ {1} ∪ {2 } ≥0 . Suppose further that we use σ(t) = 1 6 relu 3 (t) (so that σ (t) = relu(t)), then there exists choices of (w + , w -) that depends on w 0 such that E w0 σ (w 0 x) (w + x) 2 -(w -x) 2 = f (x) and further satisfies the norm bound

E w0 w + 4 2 + w - 4 2 ≤ 2π((p -2) ∨ 1) 3 α 2 B 2(p-2) x β 2p 2 .

Proof. Our proof builds on reducing the problem from representing (β x) p via quadratic networks to representing (β 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 0 , and β is the fixed coefficient vector in f . With this choice, the quadratic network reduces to

E w0 σ (w 0 x) (w + x) 2 -(w -x) 2 = E w0 σ (w 0 x) a + (β x) 2 -a -(β x) 2 = (β x) 2 E w0 σ (w 0 x)a .

Therefore, to let the above express f (x) = α(β x) p , it suffices to choose a such that

EQUATION

for all x. By Lemma 9, there exists a = a(w 0 ) satisfying ( 23) and such that

E w0 [a 2 ] ≤ 2π((p -2) ∨ 1) 3 α 2 B 2(p-2) x β 2(p-2) 2

.

Using this a in ( 22), the quadratic network induced by (w + , w -) has the desired expressivity, and further satisfies the expected 4th power norm bound

EQUATION

2 . This is the desired result.

C.4 PROOF OF THEOREM 7:

We begin by stating and proving the result for k = 1 in Appendix C.4.1, i.e. when f = α(β 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.1 EXPRESSING 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 (x) = α(β x) p where p -2 ∈ {1} ∪ {2 } ≥0 . Suppose further that we use σ(t) = 1 6 relu 3 (t) (so that σ (t) = relu(t)), then so long as the width is sufficiently large:

m ≥ O ndp 3 α 2 (B x β 2 ) 2p -2 ,

we have with probability at least 1 -δ (over W 0 ) that there exists W ∈ R d×m such that

L Q (W ) -L(f ) ≤ and W 4 2,4 ≤ B 4 w, = O p 3 α 2 B 2(p-2) x β 2p 2 δ -1 .

Proof of Theorem 16:

We build on the infinite-neuron construction in Lemma 15. Given the symmetric initialization {w 0,r } m r=1 , for all r ∈ [m/2], we consider W ∈ R d×m defined through (w ,r , w ,r+m/2 ) = 2m -1/4 w + (w 0,r ), 2m -1/4 w -(w 0,r ) , where we recall (w + (w 0 ), w -(w 0 )) = ( a + (w 0 )β, a -(w 0 )β). We then have

f Q W (x) = 1 2 √ m r≤m/2 σ (w 0,r x) (w ,r x) 2 -(w ,r+m/2 x) 2 = 2 m r≤m/2 σ (w 0,r x) (w + (w 0,r ) x) 2 -(w -(w 0,r ) x) 2 =   1 m/2 r≤m/2 σ (w 0,r x)a(w 0,r )   • (β x) 2 .

Bound on W 2,4 As f (x) = α(β x) p , Lemma 15 guarantees that the coefficient a(w 0 ) involved above satisfies that

R 2 a := E w0 [a(w 0 ) 2 ] ≤ 2π((p -2) ∨ 1) 3 α 2 B 2(p-2) x β 2(p-2) 2

.

By Markov inequality, we have with probability at least 1 -δ/2 that

1 m/2 r≤m a(w 0,r ) 2 ≤ 4π((p -2) ∨ 1) 3 α 2 B 2(p-2) x β 2(p-2) 2 δ -1 , which yields the bound W 4 2,4 = r≤m w ,r 4 2 ≤ β 4 2 • r≤m/2 16m -1 a(w 0,r ) 2 = 8 β 4 2 • 1 m/2 r≤m/2 a(w 0,r ) 2 ≤ 32π[(p -2) 3 ∨ 1]α 2 B 2(p-2) x β 2p 2 δ -1 .

Concentration of function Let f m (x) = 1 m r≤m/2 σ (w 0,r x)a(w 0,r ). We now show the concentration of f m to f ,p-2 (x) := α(β 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

f m (x) = 1 m r≤m σ (w 0,r x)a(w 0,r )1 w 0,r 2 ≤ RB -1 x := f R m (x). Letting f R ,p-2 (x) := E w0 [σ (w 0 x)a(w 0 )1 w 0 2 ≤ RB -1 x ],

we have

E W0 f m (x) -f R ,p-2 (x) 2 = 1 m E w0 σ (w 0 x)a 2 (w 0 )1 { w 0 2 ≤ R} ≤ C R 2 R 2 a m .

Applying Chebyshev inequality and a union bound, we get

P max i |f m (x i ) -f ,p-2 (x i )| ≥ t ≤ C nR 2 R 2 a mt 2 .

For any > 0, by substituting in t = B -2

EQUATION

ensures that max

EQUATION

)

Next, for any x we have the bound

f R ,p-2 (x) -f ,p-2 (x) = E w0 [σ (w 0 x)a(w 0 )1 { w 0 2 > R}] ≤ E[a(w 0 ) 2 ] 1/2 • E[σ (w 0 x) 4 ] 1/4 • P( w 0 2 > R) 1/4 ≤ R a • C/ √ d • P( w 0 2 > R) 1/4 .

Choosing R such that

EQUATION

)

Published as a conference paper at ICLR 2020 Combining ( 35) and ( 37), we see that with probability at least 1 -δ,

max i∈[n] f Q W (x i ) -f (x i ) = max i∈[n] |f m (x i ) -f ,p-2 (x i )| • (β x i ) 2 ≤ 2 • B -2 x β -2 2 2 • B 2 x β 2 2 =

and thus

EQUATION

)

To satisfy the requirements for m and R in ( 36) and ( 34 

m ≥ O nd(p -2) 3 α 2 (B x β 2 ) 2p -2 .

for (38) to hold.

C.4.2 PROOF OF MAIN THEOREM:

We apply Theorem 16 k times: let f ,j (x) := α j (β j x) pj , so that f = j≤k f ,j . Associate each j with an independent set of initialization (a (j) 0 , W (j) 0 ). By Theorem 16, there exists W (j) ∈ R d×mj , where

m j = O ndk 2 p 3 j α 2 j (B x β j 2 ) 2pj

-2 such that with probability at least 1 -δ/k we have

max i∈[n] f Q W (j) (x i ) -f (x i ) ≤ /k

and the norm bound

W (j) 4 2,4 ≤ O kp 3 j α 2 j B 2(pj -2) x β j 2pj 2 δ -1 .

(Note we have slightly abused notation, so that now f Q

W (j) j∈[k]

use a disjoint set of initial weights (a

EQUATION

).) Concatenating all the (W (j) , a

EQUATION

0 ) and applying a union bound, we have the following: so long as the width

m ≥ k j=1 m j = O   ndk 2 k j=1 p 3 j α 2 j (B x β j 2 ) 2pj -2   , with probability at least 1 -δ (over a 0 ∈ R m and W 0 ∈ R d×m ), there exists W ∈ R d×m such that max i∈[n] f Q W (x i ) -f (x i ) ≤ ,

which by the 1-Lipschitzness of the loss implies that

L Q (W ) ≤ L(f ) + = 0 + .

Further, as W is the concatenation of W (j)

j∈[k]

, we have the norm bound

W 4 2,4 = k j=1 W (j) 4 2,4 = O   k k j=1 p 3 j α 2 j B 2(pj -2) x β j 2pj 2 δ -1   .

This is the desired result.

D EXISTENCE, GENERALIZATION, AND EXPRESSIVITY OF HIGHER-ORDER NTKS:

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×2m the definition of the k-th order NTK

EQUATION

D.1 COUPLING f AND f (k) VIA RANDOMZIATION

Recall that for analytic σ we have the expansion

f W0+W (x) = 1 √ m r≤m σ((w 0,r + w +,r ) x) -σ((w 0,r + w -,r ) x) = ∞ k=0 f (k) W0,W (x),

For an arbitrary W such that w +,r 2 , w -,r 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 finding W so that

f W0+W (x) = f (k) W0,W (x) + o m (1)

, 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 ,r , z -,r w ,r ), and take w ,r 2 = O(m -1/2k ), we have

f (j) W0,W (x) = 1 √ m r≤m 1 j! σ (j) (w 0,r x) (z j +,r -z j -,r ) mean zero (w ,r x) j O(m -j/2k ) = O p (m -j/2k

) for all j = 1, . . . , k -1, and

EQUATION

f (k+1) W0,W (x) = 1 √ m r≤m 1 (k + 1)! σ (k+1) (w 0,r x) (z k+1 +,r -z k+1 -,r ) (w ,r x) k+1 O(m -(k+1)/2k ) = O P (m -1/2k ).

Therefore, with high probability, all f (1) , . . . , f (k-1) as well as the remainder term f -j≤k f (j) has order O(m -1/2k ), and the k-th order NTK f (k) can express an O(1) function.

D.2 GENERALIZATION 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

σ k (t) := 1 k! σ (k) (t) ≡ relu(t)

is the ReLU activation.

As we have seen in Section 6, we have f (k) = O(1) by choosing w r ∼ O(m -1/2k ), therefore we restrict attention on such W's by considering the constraint set {W :

W 2k 2,2k ≤ B 2k w } 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 W 2,2k ≤ B w , when the base distribution of x is uniform on the sphere, scales as

O B k x B k w 1 √ nd k-1 + 1 n + 1 √ n .

(Theorem 19). Compared with the distribution-free bound B k x B k w / √ n, the leading term is better by a factor of 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 (x) = α(β x) p with p ≥ k (and p -k is even or one), when m is sufficiently large, there exists a W expressing f such that 20) . Substituting into the generalization bound yields the following generalization error for learning f :

EQUATION

O p 3 α 2 (B x β 2 ) p 1 √ nd k-1 + 1 n .

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 benefit of higher-order NTKs when distributional assumptions are present.

Tensor operator and nuclear norm Our result requires the definition of operator norm and nuclear norm for k-tensors, which we briefly review here. The operator norm of a symmetric k-tensor A ∈ R d k is defined as

A op := sup v 2 =1 A, v ⊗k = sup v 2 =1 A[v, . . . , v].

The nuclear norm • * is defined as the dual norm of the operator norm:

A * := sup B op =1 A, B .

Specifically, for any rank-one tensor u ⊗k , we have

u ⊗k * = sup B op =1 u ⊗k , B = u k 2 ,

i.e. its nuclear norm equals its operator norm (and also the Frobenius norm).

D.2.1 GENERALIZATION:

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 such that z → (y, z) is 1-Lipschitz and (y, 0) ≤ 1 for all y ∈ Y, we have the Rademacher complexity bound

E σ,x sup W 2,2k ≤Bw 1 n n i=1 σ i (y i , f (k) W0,W (x i )) ≤ 2B k w E σ,x   max r∈[m] 1 n n i=1 σ i σ k (w 0,r x i )x ⊗k i op   + 1 √ n ,

where σ i iid ∼ Unif{±1} are Rademacher variables.

Proof. The proof is analogous to that of Lemma 5. As the loss (y, z) is 1-Lipschitz in z for all y, by the Rademacher contraction theorem (Wainwright, 2019, Chapter 5) we have that

E σ,x sup W 2,2k ≤Bw 1 n n i=1 σ i (y i , f (k) W0,W (x i )) ≤ 2E σ,x sup W 2,2k ≤Bw 1 n n i=1 σ i f (k) W0,W (x i ) + E σ,x 1 n n i=1 σ i (y i , 0) ≤ 2E σ,x   sup W 2,2k ≤Bw 1 √ m r≤m 1 n n i=1 σ i a r σ k (w 0,r x i )x ⊗k i , w ⊗k r   + 1 √ n ≤ 2E σ,x   sup W 2,2k ≤Bw max r∈[m] 1 n n i=1 a r σ i σ k (w 0,r x i )x ⊗k i op • 1 √ m r≤m w ⊗k r *   + 1 √ n ≤ 2E σ,x   max r∈[m] 1 n n i=1 σ i σ k (w 0,r x i )x ⊗k i op   • sup W 2,2k ≤Bw 1 √ m w r k 2 ≤B k w + 1 √ n ,

where the last step used the power mean (or Cauchy-Schwarz) inequality on { w r 2 }.

Bound on tensor operator norm It is straightforward to see that the expected tensor operator norm can be bounded as O B k x / √ n 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

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 )

E D sup W 2,2k ≤Bw L P (f (k) W0,W ) -L(f (k) W0,W ) ≤ O B k x B k w 1 √ nd k-1 + 1 n + 1 √ n .

The proof of Lemma 18 is deferred to Appendix D.3.

D.2.2 EXPRESSIVITY:

Theorem 20 (Expressivity of f (k) ). Suppose {(a r , w 0,r )} are generated according to the symmetric initialization (3), and f (x) = α(β x) p where p -k ∈ {1} ∪ {2 } ≥0 . Suppose further that σ is such that σ k (t) = relu(t), then so long as the width is sufficiently large:

m ≥ O ndp 3 α 2 (B x β 2 ) 2p -2 ,

we have with probability at least 1 -δ (over W 0 ) that there exists W ∈ R d×m such that

L Q (W ) -L(f ) ≤ and W 2k 2,2k ≤ B 2k w, = O p 3 α 2 B 2(p-k) x β 2p 2 δ -1 .

The proof of Theorem 20 is deferred to Appendix D.4.

D.3 PROOF OF LEMMA 18:

We begin by observing for any symmetric tensor

A ∈ R d k that A op ≤ 1 1 -k sup v∈N ( ) A[v, . . . , v],

where N ( ) is an -covering of unit sphere S d-1 (1). (The proof follows by bounding A[u, . . . , u] by A[v, . . . , v] + k A op through replacing u by v one at a time). Taking = 1/(2k), we have

P σ,x   max r∈[m] 1 n n i=1 σ i σ k (w 0,r x i )x ⊗k i op ≥ t   ≤ P σ,x max r∈[m],v∈N (1/(2k)) 1 n n i=1 σ i σ k (w 0,r x i )(v x i ) k ≥ t/2 .

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 On the other hand, the inequality (nt/B k x ) 1/2 > O(d) happens when t > O d 2 B k x /n , which is implied by the preceding condition so long as k ≥ 3. Therefore, when this branch is taken, the O(d) can already be absorbed into the main term, so the contribution of this branch can be bounded as (3) We have 

P σ,x   max r∈[m] 1 n n i=1 σ i σ k (w 0,r x i )x ⊗k i op ≥ t   ≤ P σ,x max r∈[m],v∈N (1/(2k)) 1 n n i=1 σ i σ k (w 0,r x i )(v x i ) k 1 σ k (w 0,r x i ) ≤ M ≥ t/2 + P x max

D.4 PROOF OF THEOREM 20:

Our proof is analogous to that of Theorem 16, in which we first look at the case of infinitely many neurons and then use concentration to carry the result onto finitely many neurons.

Expressivity with infinitely many neurons We first consider expressing f (x) = α(β x) p with infinite-neuron version of f (k) , that is, we wish to find random variables (w + , w -) such that E w0 relu(w 0 x) (w + x) k -(w -x) k = f (x). Choosing (w + , w -) = ([a] + ) 1/k β, ([a] -) 1/k β for some real-valued random scalar a (that depends on w 0 ), we have E w0 relu(w 0 x) (w + x) k -(w -x) k = (β x) k • E w0 relu(w 0 x)a , therefore the task reduces to finding a = a(w 0 ) such that E w0 relu(w 0 x)a = α(β x) p-k . By Lemma 9, there exists a = a(w 0 ) satisfying the above and such that

E w0 [a 2 ] ≤ 2π((p -k) ∨ 1) 3 α 2 B 2(p-k) x β 2(p-k) 2

.

(32) Using this a, the k-th order NTK defined by (w + , w -) expresses f and further satisfies the bound

E w0 w + 2k 2 + w - 2k 2 = Ew 0 [a 2 ] • β 2k 2 ≤ 2π((p -k) ∨ 1) 3 α 2 B 2(p-k) x β 2p 2 .

Finite neurons Given the symmetric initialization {w 0,r } m r=1 , for all r ∈ [m], we consider W ∈ R d×m defined through (w ,r , w ,r+m ) = m -1/2k w + (w 0,r ), m -1/2k w -(w 0,r ) , where we recall (w + (w 0 ), w -(w 0 )) = (a + (w 0 ) 1/k β, a -(w 0 ) 1/k β). We then have Concentration of function Let f m (x) = 1 m r≤m σ k (w 0,r x)a(w 0,r ). We now show the concentration of f m to f ,p-k (x) := α(β 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 f m (x) = 1 m r≤m σ k (w 0,r x)a(w 0,r )1 w 0,r 2 ≤ RB -1

x := f R m (x).

Letting f R ,p-k (x) := E w0 [σ k (w 0 x)a(w 0 )1 w 0 2 ≤ RB -1

x ], we have

E W0 f m (x) -f R ,p-k (x) 2 = 1 m E w0 σ k (w 0 x)a 2 (w 0 )1 { w 0 2 ≤ R} ≤ C R 2 R 2 a m .

Applying Chebyshev inequality and a union bound, we get

P max i |f m (x i ) -f ,p-k (x i )| ≥ t ≤ C nR 2 R 2 a mt 2 .

For any > 0, by substituting in t = B -k

x β -k

2 /2, we see that

EQUATION

ensures that max

EQUATION

)

Next, for any x we have the bound

f R ,p-k (x) -f ,p-k (x) = E w0 [σ k (w 0 x)a(w 0 )1 { w 0 2 > R}] ≤ E[a(w 0 ) 2 ] 1/2 • E[σ k (w 0 x) 4 ] 1/4 • P( w 0 2 > R) 1/4 ≤ R a • C/ √ d • P( w 0 2 > R) 1/4 .

Choosing R such that

EQUATION

ensures that

EQUATION

) max i∈[n] f Q W (x i ) -f (x i ) = max i∈[n] |f m (x i ) -f ,p-k (x i )| • (β x i ) k ≤ 2 • B -k x β -k 2 2 • B k x β k 2 =

and thus

EQUATION

)

To satisfy the requirements for m and R in ( 36) and ( 34 for (38) to hold.

:

As a concrete example,Woodworth et al. (2019) showed that for matrix completion the NTK solution estimates zero on all unobserved entries and the max-margin solution corresponds to the minimum nuclear norm solution.

We note that the only restrictive requirement in Assumption A is the Lipschitzness of σ , which guarantees second-order smoothness of the objectives. The bounds on derivatives (and specifically their bound near zero) are merely for technical convenience and can be weakened without hurting the results.

Our specific choice of • 2,4 norm is needed for measuring the average magnitude of f Q W , whereas the high (8-th) power is not essential and can be replaced by any (4 + )-th power without affecting the result.

Our choice covers two commonly used scales in neural net analyses: Bx = 1, w0,r ∼ N(0, I d ) in e.g.(Arora et al., 2019b; Allen-Zhu et al., 2018a); Bx = √ d, w0,r ∼ N(0, I d /d) in e.g.(Ghorbani et al., 2019b).

Figures and Tables:

; Du et al. (2018b) utilized the coupling of the gradient dynamics to prove that gradient descent finds 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. (2019); Chizat & Bach (2018b) empirically found a significant performance gap between NTK and actual training. This gap has been theoretically studied in Wei et al. (2019); Allen-Zhu & Li (2019); Yehudai & Shamir (2019); Ghorbani et al. (2019a) which established that NTK has provably higher generalization error than training the neural net for specific data distributions and architectures.

for all p ∈ [1, ∞]. In particular we have • 2,2 = • Fr . We let B 2,p (R) := {W : W 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 = 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 is called a second-order stationary point if ∇f (x ) = 0 and ∇ 2 f (x ) 0.

3), and ensures that max r∈[m] |w 0,r x| ≤ O( √ d) simultaneously for all x.

The main term 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 .

Setting t = 8(d log 5 + log(m/δ)) = O( d + log(m/δ)) ensures that the above probability does not exceed δ as desired. B PROOFS FOR SECTION 4 B.1 PROOF OF LEMMA 1

), we first set R = O( √ d) (with sufficiently large log factor) to satisfy (36) by standard Gaussian norm concentration (cf. Appendix A.3), and by (34) it suffices to set m as

where the O(1) • B 2kx d -k comes from computing the variance ofZ i := σ i σ k (w 0,r x i )(v x i ) kusing that x i are uniform on the sphere (see, e.g.(Ghorbani et al., 2019b, Proof of Lemma 4)); M B k x is the bound on the variable Z i , and the O(1) comes from the fact that w 0,r 2 ≤ O( √ dB -1 x ) with high probability. Now, choosing M = (nt/B k x ) 1/2 , the above bound reads exp -c min nt 2 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 / O(B 2k x d -k )" we have This follows by integrating the "1" branch for t ≤ O( B 2k x d -(k-1) /n) (which yields the right hand side) and integrating the other branch otherwise (the integral being upper bounded by O( B 2k x d -k /n), dominated by the right hand side). (2) The branch "(nt/B k x ) 1/2 " is taken only when nt B 2k x d -k ) i.e. t > O n -1/3 B k x d 2k/3 .

k (w 0,r x i )x ⊗k i

f

w 0,r x) (w ,r x) k -(w ,r+m x) k = 1 m r≤m σ k (w 0,r x) (w + (w 0,r ) x) k -(w -(w 0,r ) x) k =   1 m r≤m σ k (w 0,r x)a(w 0,r )   • (β x) k .Bound on W 2,2k As f (x) = α(β x) p , (32) guarantees that the coefficient a(w 0 ) involved above satisfies thatR 2 a := E w0 [a(w 0 ) 2 ] ≤ 2π((p -k) ∨ 1) 3 α 2 B 2(p-k)By Markov inequality, we have with probability at least 1 -δ/2 that1 m r≤m a(w 0,r ) 2 ≤ 4π((p -k) ∨ 1) 3 α 2 B 2(p-k) a(w 0,r ) 2 ≤ 4π[(p -k) 3 ∨ 1]α 2 B 2(p-k)

), we first set R = O( √ d) (with sufficiently large log factor) to satisfy (36) by standard Gaussian norm concentration (cf. Appendix A.3), and by (34) it suffices to set m asm ≥ O nd[(p -k) 3 ∨ 1]α 2 (B x β 2 ) 2p -2 .

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. Colin Wei, Jason D Lee, Qiang Liu, and Tengyu Ma. On the margin theory of feedforward neural networks. arXiv preprint arXiv:1810.05369, 2018. Colin Wei, D Lee Jason, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization of neural nets vs their induced kernel. arXiv preprint arXiv:1810.05369, 2019. Christopher KI Williams. Computing with infinite networks. In Advances in neural information processing systems, pp. 295-301, 1997. Blake Woodworth, Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Kernel and deep regimes in overparametrized models. arXiv preprint arXiv:1906.05827, 2019. Gilad Yehudai and Ohad Shamir. On the power and limitations of random features for understanding neural networks. arXiv preprint arXiv:1904.00687, 2019. Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Stochastic gradient descent optimizes over-parameterized deep ReLU networks. arXiv preprint arXiv:1811.08888, 2018.

Table content: <table/>

Peer review (iclr2020):
This paper presents an approach for going beyond NTK regime, namely the linear Taylor approximation for network function. By employing the idea of randomization this paper manages to reduce the magnitude of the linear part while the quadratic part is unaffected. This technique enables further analysis of both the optimization and generalization based on the nice property of quadratic approximation. I believe this paper should be accepted because of its novel approach to go beyond the linear part, which I view as a main contribution. Also, this paper is well written and presents its optimization result and proof in a clear manner. Although it is an innovation in the theoretical perspective, I still want to raise two questions about the object this paper trying to analyze: 1. The activation function is designed so that it has really nice property: it has second-order derivative which is lipshitz. Actually I believe Assumption A is first motivated by cubic ReLU. Why is cubic ReLU so favorable in this paper? Is it possible to use quadratic ReLU in the proof? Also I wonder what is the key property of activation function which is allowed here. 2. The network model considered here, is modified to a randomized neural network. So maybe this paper just circumvents the difficulty in going beyond NTK regime? I have this concern because in reality optimizing this network model seems intractable. To evaluate the loss or we take expectation over ; when doing gradient descent, apparently the gradient also takes exponential time to compute. Overall, I believe this paper has high quality and I enjoyed reading it. However, I do expect response from authors which could address my concerns raised above. This can help me achieve a better understanding and a more precise evaluation on this paper. *** I have read the author*s response and decide to remain the rating as weak accept.