Title: POLYLOGARITHMIC WIDTH SUFFICES FOR GRADIENT DESCENT TO ACHIEVE ARBITRARILY SMALL TEST ER-ROR WITH SHALLOW RELU NETWORKS

Abstract: Recent theoretical work has guaranteed that overparameterized networks trained by gradient descent achieve arbitrarily low training error, and sometimes even low test error. The required width, however, is always polynomial in at least one of the sample size n, the (inverse) target error 1 / , and the (inverse) failure probability 1 /δ. This work shows that Θ( 1 / ) iterations of gradient descent with Ω( 1 / 2 ) training examples on two-layer ReLU networks of any width exceeding polylog(n, 1 / , 1 /δ) suffice to achieve a test misclassification error of . We also prove that stochastic gradient descent can achieve test error with polylogarithmic width and Θ( 1 / ) samples. The analysis relies upon the separation margin of the limiting kernel, which is guaranteed positive, can distinguish between true labels and random labels, and can give a tight sample-complexity analysis in the infinitewidth setting.

1. INTRODUCTION:

Despite the extensive empirical success of deep networks, their optimization and generalization properties are still not fully understood. Recently, the neural tangent kernel (NTK) has provided the following insight into the problem. In the infinite-width limit, the NTK converges to a limiting kernel which stays constant during training; on the other hand, when the width is large enough, the function learned by gradient descent follows the NTK (Jacot et al., 2018) . This motivates the study of overparameterized networks trained by gradient descent, using properties of the NTK. In fact, parameters related to the NTK, such as the minimum eigenvalue of the limiting kernel, appear to affect optimization and generalization (Arora et al., 2019) .

However, in addition to such NTK-dependent parameters, prior work also requires the width to depend polynomially on n, 1/δ or 1/ , where n denotes the size of the training set, δ denotes the failure probability, and denotes the target error. These large widths far exceed what is used empirically, constituting a significant gap between theory and practice.

Our contributions. In this paper, we narrow this gap by showing that a two-layer ReLU network with Ω(ln(n/δ)+ln(1/ ) 2 ) hidden units trained by gradient descent achieves classification error on test data, meaning both optimization and generalization occur. Unlike prior work, the width is fully polylogarithmic in n, 1/δ, and 1/ ; the width will additionally depend on the separation margin of the limiting kernel, a quantity which is guaranteed positive (assuming no inputs are parallel), can distinguish between true labels and random labels, and can give a tight sample-complexity analysis in the infinite-width setting. The paper organization together with some details are described below.

Section 2 studies gradient descent on the training set. Using the 1 geometry inherent in classification tasks, we prove that with any width at least polylogarithmic and any constant step size no larger than 1, gradient descent achieves training error in Θ(1/ ) iterations (cf. Theorem 2.2). As is common in the NTK literature (Chizat & Bach, 2019) , we also show the parameters hardly change, which will be essential to our generalization analysis.

Section 3 gives a test error bound. Concretely, using the preceding gradient descent analysis, and standard Rademacher tools and exploiting how little the weights moved, we show that with Ω(1/ 2 ) samples and Θ(1/ ) iterations, gradient descent finds a solution with test error (cf. Theorem 3.2 and Corollary 3.3). (As discussed in Remark 3.4, Ω(1/ ) samples also suffice via a smoothness-based generalization bound, at the expense of large constant factors.)

Section 4 considers stochastic gradient descent (SGD) with access to a standard stochastic online oracle. We prove that with width at least polylogarithmic and Θ(1/ ) samples, SGD achieves an arbitrarily small test error (cf. Theorem 4.1).

Section 5 discusses the separation margin, which is in general a positive number, but reflects the difficulty of the classification problem in the infinite-width limit. While this margin can degrade all the way down to O(1/ √ n) for random labels, it can be much larger when there is a strong relationship between features and labels: for example, on the noisy 2-XOR data introduced in (Wei et al., 2018) , we show that the margin is Ω(1/ ln(n)), and our SGD sample complexity is tight in the infinite-width case.

Section 6 concludes with some open problems.

1.1. RELATED WORK:

There has been a large literature studying gradient descent on overparameterized networks via the NTK. The most closely related work is (Nitanda & Suzuki, 2019) , which shows that a two-layer network trained by gradient descent with the logistic loss can achieve a small test error, under the same assumption that the NTK with respect to the first layer can separate the data distribution. However, they analyze smooth activations, while we handle the ReLU. They require Ω(1/ 2 ) hidden units, Ω(1/ 4 ) data samples, and O(1/ 2 ) steps, while our result only needs polylogarithmic hidden units, Ω(1/ 2 ) data samples, and O(1/ ) steps.

Additionally on shallow networks, Du et al. (2018b) prove that on an overparameterized two-layer network, gradient descent can globally minimize the empirical risk with the squared loss. Their result requires Ω(n 6 /δ 3 ) hidden units. Oymak & Soltanolkotabi (2019) ; Song & Yang (2019) further reduce the required overparameterization, but there is still a poly(n) dependency. Using the same amount of overparameterization as (Du et al., 2018b) , Arora et al. (2019) further show that the twolayer network learned by gradient descent can achieve a small test error, assuming that on the data distribution the smallest eigenvalue of the limiting kernel is at least some positive constant. They also give a fine-grained characterization of the predictions made by gradient descent iterates; such a characterization makes use of a special property of the squared loss and cannot be applied to the logistic regression setting. Li & Liang (2018) show that stochastic gradient descent (SGD) with the cross entropy loss can learn a two-layer network with small test error, using poly( , 1/ ) hidden units, where is at least the covering number of the support of the feature distribution using balls whose radii are no larger than the smallest distance between two data points with different labels. Allen-Zhu et al. (2018a) consider SGD on a two-layer network, and a variant of SGD on a three-layer network. The three-layer analysis further exhibits some properties not captured by the NTK. They assume a ground truth network with infinite-order smooth activations, and they require the width to depend polynomially on 1/ and some constants related to the smoothness of the activations of the ground truth network.

On deep networks, a variety of works have established low training error (Allen-Zhu et al., 2018b; Du et al., 2018a; Zou et al., 2018; Zou & Gu, 2019) . Allen-Zhu et al. (2018c) show that SGD can minimize the regression loss for recurrent neural networks, and Allen-Zhu & Li (2019b) further prove a low generalization error. Allen-Zhu & Li (2019a) show that using the same number of training examples, a three-layer ResNet can learn a function class with a much lower test error than any kernel method. Cao & Gu (2019a) assume that the NTK with respect to the second layer of a two-layer network can separate the data distribution, and prove that gradient descent on a deep network can achieve test error with Ω(1/ 4 ) samples and Ω(1/ 14 ) hidden units. Cao & Gu (2019b) consider SGD with an online oracle and give a general result. Under the same assumption as in (Cao & Gu, 2019a) , their result requires Ω(1/ 14 ) hidden units and sample complexity O(1/ 2 ).

By contrast, with the same online oracle, our result only needs polylogarithmic hidden units and sample complexity O(1/ ).

1.2. NOTATION:

The dataset is denoted by {(x i , y i )} n i=1 where x i ∈ R d and y i ∈ {-1, +1}. For simplicity, we assume that x i 2 = 1 for any 1 ≤ i ≤ n, which is standard in the NTK literature.

The two-layer network has weight matrices W ∈ R m×d and a ∈ R m . We use the following parameterization, which is also used in (Du et al., 2018b; Arora et al., 2019) :

f (x; W, a) := 1 √ m m s=1 a s σ w s , x , with initialization w s,0 ∼ N (0, I d ), and a s ∼ unif {-1, +1} .

Note that in this paper, w s,t denotes the s-th row of W at step t. We fix a and only train W , as in (Li & Liang, 2018; Du et al., 2018b; Arora et al., 2019; Nitanda & Suzuki, 2019) . We consider the ReLU activation σ(z) := max {0, z}, though our analysis can be extended easily to Lipschitz continuous, positively homogeneous activations such as leaky ReLU.

We use the logistic (binary cross entropy) loss (z) := ln 1 + exp(-z) and gradient descent. For any 1 ≤ i ≤ n and any W , let f i (W ) := f (x i ; W, a). The empirical risk and its gradient are given by

R(W ) := 1 n n i=1 y i f i (W ) , and ∇ R(W ) = 1 n n i=1 y i f i (W ) y i ∇f i (W ).

For any t ≥ 0, the gradient descent step is given by

W t+1 := W t -η t ∇ R(W t ). Also define f (t) i (W ) := ∇f i (W t ), W , and R (t) (W ) := 1 n n i=1 y i f (t) i (W ) . Note that f (t) i (W t ) = f i (W t ).

This property generally holds due to homogeneity: for any W and any 1 ≤ s ≤ m,

∂f i ∂w s = 1 √ m a s 1 w s , x i > 0 x i , and ∂f i ∂w s , w s = 1 √ m a s σ w s , x i ,

and thus ∇f i (W ), W = f i (W ).

2. EMPIRICAL RISK MINIMIZATION:

In this section, we consider a fixed training set and empirical risk minimization. We first state our assumption on the separability of the NTK, and then give our main result and a proof sketch.

The key idea of the NTK is to do the first-order Taylor approximation:

f (x; W, a) ≈ f (x; W 0 , a) + ∇ W f (x; W 0 , a), W -W 0 .

In other words, we want to do learning using the features given by ∇f i (W 0 ) ∈ R m×d . A natural assumption is that there exists U ∈ R m×d which can separate ∇f i (W 0 ), y i n i=1 with a positive margin:

min 1≤i≤n y i U , ∇f i (W 0 ) = min 1≤i≤n   y i 1 √ m m s=1 a s ūs , x i 1 w s,0 , x i > 0   > 0.

(2.1)

The infinite-width limit of eq. (2.1) is formalized as Assumption 2.1, with an additional bound on the (2, ∞) norm of the separator. A concrete construction of U using Assumption 2.1 is given in eq. (2.2).

Let µ N denote the Gaussian measure on R d , given by the Gaussian density with respect to the Lebesgue measure on R d . We consider the following Hilbert space

H := w : R d → R d w(z) 2 2 dµ N (z) < ∞ . For any x ∈ R d , define φ x ∈ H by φ x (z) := x1 z, x > 0 ,

and particularly define φ i := φ xi for the training input x i . Assumption 2.1. There exists v ∈ H and γ > 0, such that v(z) 2 ≤ 1 for any z ∈ R d , and for any

1 ≤ i ≤ n, y i v, φ i H := y i v(z), φ i (z) dµ N (z) ≥ γ.

♦:

As discussed in Section 5, the space H is the reproducing kernel Hilbert space (RKHS) induced by the infinite-width NTK with respect to W , and φ x maps x into H. Assumption 2.1 supposes that the induced training set {(φ i , y i )} n i=1 can be separated by some v ∈ H, with an additional bound on v(z) 2 which is crucial in our analysis. It is also possible to give a dual characterization of the separation margin (cf. eq. (5.2)), which also allows us to show that Assumption 2.1 always holds when there are no parallel inputs (cf. Proposition 5.1). However, it is often more convenient to construct v directly; see Section 5 for some examples.

With Assumption 2.1, we state our main empirical risk result. Theorem 2.2. Under Assumption 2.1, given any risk target ∈ (0, 1) and any δ ∈ (0, 1/3), let λ := 2 ln(4n/δ) + ln(4/ ) γ/4 , and M := 4096λ 2 γ 6 . Then for any m ≥ M and any constant step size η ≤ 1, with probability 1 -3δ over the random initialization, 1 T t<T R(W t ) ≤ , where T := 2λ 2 /η .

Moreover for any 0 ≤ t < T and any 1 ≤ s ≤ m,

w s,t -w s,0 2 ≤ 4λ γ √ m .

While the number of hidden units required by prior work all have a polynomial dependency on n, 1/δ or 1/ , Theorem 2.2 only requires m = Ω ln(n/δ) + ln(1/ ) 2 . The required width has a polynomial dependency on 1/γ, which is an adaptive quantity: while 1/γ can be poly(n) for random labels (cf. Proposition 5.2), it can be polylog(n) when there is a strong feature-label relationship, for example on the noisy 2-XOR data introduced in (Wei et al., 2018) (cf. Proposition 5.3). Moreover, we show in Proposition 5.4 that if we want ∇f i (W 0 ), y i n i=1 to be separable, which is the starting point of an NTK-style analysis, the width has to depend polynomially on 1/γ.

In the rest of Section 2, we give a proof sketch of Theorem 2.2. The full proof is given in Appendix A.

2.1. PROPERTIES AT INITIALIZATION:

In this subsection, we give some nice properties of random initialization.

Given an initialization (W 0 , a), for any 1 ≤ s ≤ m, define

ūs := 1 √ m a s v(w s,0 ), (2.2)

where v is given by Assumption 2.1. Collect ūs into a matrix U ∈ R m×d . It holds that ūs 2 ≤ 1/ √ m, and U F ≤ 1.

Lemma 2.3 ensures that with high probability U has a positive margin at initialization.

Lemma 2.3. Under Assumption 2.1, given any δ ∈ (0, 1) and any 1 ∈ (0, γ), if m ≥ 2 ln(n/δ) / 2 1 , then with probability 1 -δ, it holds simultaneously for all 1 ≤ i ≤ n that

y i f (0) i U = y i ∇f i (W 0 ), U ≥ γ - 2 ln(n/δ) m ≥ γ -1 .

For any W , any 2 > 0, and any

1 ≤ i ≤ n, define α i (W, 2 ) = 1 m m s=1 1 w s , x i ≤ 2 .

Lemma 2.4 controls α i (W 0 , 2 ). It will help us show that U has a good margin during the training process.

Lemma 2.4. Under the condition of Lemma 2.3, for any 2 > 0, with probability 1 -δ, it holds simultaneously for all 1 ≤ i ≤ n that

α i (W 0 , 2 ) ≤ 2 π 2 + ln(n/δ) 2m ≤ 2 + 1 2 .

Finally, Lemma 2.5 controls the output of the network at initialization.

Lemma 2.5. Given any δ ∈ (0, 1), if m ≥ 25 ln(2n/δ), then with probability 1 -δ, it holds simultaneously for all

1 ≤ i ≤ n that f (x i ; W 0 , a) ≤ 2 ln 4n/δ .

2.2. CONVERGENCE ANALYSIS OF GRADIENT DESCENT:

We analyze gradient descent in this subsection. First, define

Q(W ) := 1 n n i=1 -y i f i (W ) .

We have the following observations.

• For any W and any 1 ≤ s ≤ m, ∂f i /∂w s 2 ≤ 1/ √ m, and thus ∇f i (W ) F ≤ 1.

Therefore by the triangle inequality, ∇ R(W )

F ≤ Q(W ).

• The logistic loss satisfies 0 ≤ -≤ 1, and thus 0 ≤ Q(W ) ≤ 1.

• The logistic loss satisfies -≤ , and thus

Q(W ) ≤ R(W ).

The quantity Q first appeared in the perceptron analysis (Novikoff, 1962) for the ReLU loss, and has also been analyzed in prior work (Ji & Telgarsky, 2018; Cao & Gu, 2019a; Nitanda & Suzuki, 2019) . In this work, Q specifically helps us prove the following result, which plays an important role in obtaining a width which only depends on polylog(1/ ).

Lemma 2.6. For any t ≥ 0 and any W , if η t ≤ 1, then

η t R(W t ) ≤ W t -W 2 F -W t+1 -W 2 F + 2η t R (t) W .

Consequently, if we use a constant step size η ≤ 1 for 0 ≤ τ < t, then

η   τ <t R(W τ )   + W t -W 2 F ≤ W 0 -W 2 F + 2η   τ <t R (τ ) W   .

The proof of Lemma 2.6 starts from the standard iteration guarantee:

W t+1 -W 2 F = W t -W 2 F -2η t ∇ R(W t ), W t -W + η 2 t ∇ R(W t ) 2 F .

We can then handle the inner product term using the convexity of and homogeneity of ReLU, and control ∇ R(W t ) 2 F by R(W t ) using the above properties of Q(W t ). Lemma 2.6 is similar to (Allen-Zhu & Li, 2019a, Fact D.4 and Claim D.5) , where the squared loss is considered.

Using Lemmas 2.3 to 2.6, we can prove Theorem 2.2. Below is a proof sketch; the full proof is given in Appendix A.

1. We first show that as long as

w s,t -w s,0 2 ≤ 4λ/(γ √ m) for all 1 ≤ s ≤ m, it holds that R (t) W 0 + λU ≤ /4. To see this, let us consider R (0) first. For any 1 ≤ i ≤ n, Lemma 2.5 ensures that | ∇f i (W 0 ), W 0 | is bounded, while Lemma 2.3 ensures that ∇f i (W 0 ), U is concentrated

around γ with a large width. As a result, with the chosen λ in Theorem 2.2, we can show that ∇f i (W 0 ), W 0 + λU is large, and R (0) (W 0 + λU ) is small due to the exponential tail of the logistic loss. To further handle R (t) , we use a standard NTK argument to control ∇f

i (W t ) -∇f i (W 0 ), W 0 + λU under the condition that w s,t -w s,0 2 ≤ 4λ/(γ √ m).

2. We then prove by contradiction that the above bound on w s,t -w s,0 2 holds for at least the first T iterations. The key observation is that as long as R (t) (W 0 + λU ) ≤ /4, we can use it and Lemma 2.6 to control τ <t Q(W τ ), and then just invoke

w s,t -w s,0 2 ≤ η τ <t Q(W τ )/ √ m.

The quantity τ <t Q(W τ ) has also been considered in prior work (Cao & Gu, 2019a; Nitanda & Suzuki, 2019) , where it is bounded by √ t τ <t Q(W τ ) 2 using the Cauchy-Schwarz inequality, which introduces a √ t factor. To make the required width depend only on polylog(1/ ), we also need an upper bound on τ <t Q(W τ ) which depends only on polylog(1/ ). Since the above analysis results in a √ t factor, and in our case Ω(1/ ) steps are needed, it is unclear how to get a polylog(1/ ) width using the analysis in (Cao & Gu, 2019a; Nitanda & Suzuki, 2019) . By contrast, using Lemma 2.6, we can show that τ <t Q(W τ ) ≤ 4λ/γ, which only depends on ln(1/ ). 3. The claims of Theorem 2.2 then follow directly from the above two steps and Lemma 2.6.

3. GENERALIZATION:

To get a generalization bound, we naturally extend Assumption 2.1 to the following assumption. Assumption 3.1. There exists v ∈ H and γ > 0, such that v(z) 2 ≤ 1 for any z ∈ R d , and

y v(z), x 1 z, x > 0 dµ N (z) ≥ γ

for almost all (x, y) sampled from the data distribution D. ♦

The above assumption is also made in (Nitanda & Suzuki, 2019) for smooth activations. (Cao & Gu, 2019a ) make a similar separability assumption, but in the RKHS induced by the second layer a; by contrast, Assumption 3.1 is on separability in the RKHS induced by the first layer W .

Here is our test error bound with Assumption 3.1.

Theorem 3.2. Under Assumption 3.1, given any ∈ (0, 1) and any δ ∈ (0, 1/4), let λ and M be given as in Theorem 2.2:

λ := 2 ln(4n/δ) + ln(4/ ) γ/4

, and M := 4096λ 2 γ 6 .

Then for any m ≥ M and any constant step size η ≤ 1, with probability 1 -4δ over the random initialization and data sampling,

P (x,y)∼D yf (x; W k , a) ≤ 0 ≤ 2 + 16 2 ln(4n/δ) + ln(4/ ) γ 2 √ n + 6 ln(2/δ) 2n ,

where k denotes the step with the minimum empirical risk before 2λ 2 /η .

Below is a direct corollary of Theorem 3.2. Corollary 3.3. Under Assumption 3.1, given any , δ ∈ (0, 1), using a constant step size no larger than 1 and let

n = Ω 1 γ 4 2 , and m = Ω ln(n/δ) + ln(1/ ) 2 γ 8 ,

it holds with probability 1 -δ that P (x,y)∼D yf (x; W k , a) ≤ 0 ≤ , where k denotes the step with the minimum empirical risk in the first Θ( 1 /γ 2 ) steps.

The proof of Theorem 3.2 uses the sigmoid mapping -(z) = e -z /(1+e -z ), the empirical average Q(W k ), and the corresponding population average 

Q(W k ) := E (x,y)∼D -yf (x; W k , a) . As noted in (Cao & Gu, 2019a), because P (x,y)∼D yf (x; W k , a) ≤ 0 ≤ 2Q(W k ), it is enough to control Q(W k ). As Q(W k ) is controlled by Theorem 2.2, it is enough to control the generalization error Q(W k ) -Q(W k ).

4. STOCHASTIC GRADIENT DESCENT:

There are some different formulations of SGD. In this section, we consider SGD with an online oracle. We randomly sample W 0 and a, and fix a during training. At step i, a data example (x i , y i ) is sampled from the data distribution. We still let f i (W ) := f (x i ; W, a), and perform the following update

W i+1 := W i -η i y i f i (W i ) y i ∇f i (W i ).

Note that here i starts from 0.

Still with Assumption 3.1, we show the following result. Theorem 4.1. Under Assumption 3.1, given any , δ ∈ (0, 1), using a constant step size and m = Ω (ln(1/δ)+ln(1/ ) 2 ) /γ 8 , it holds with probability 1 -δ that

1 n n i=1 P (x,y)∼D yf (x; W i , a) ≤ 0 ≤ , for n = Θ( 1 /γ 2 ).

Below is a proof sketch of Theorem 4.1; the complete proof is given in Appendix C. For any i and W , define

R i (W ) := y i ∇f i (W i ), W , and Q i (W ) := -y i ∇f i (W i ), W .

Due to homogeneity, it holds that

R i (W i ) = y i f i (W i ) and Q i (W i ) = -y i f i (W i ) .

The first step is an extension of Lemma 2.6 to the SGD setting, with a similar proof. Lemma 4.2. With a constant step size η ≤ 1, for any W and any i ≥ 0,

η   t<i R t (W t )   + W i -W 2 F ≤ W 0 -W 2 F + 2η   t<i R t W   .

With Lemma 4.2, we can also extend Theorem 2.2 to the SGD setting and get a bound on i<n Q i (W i ), using a similar proof. To further get a bound on the cumulative population risk i<n Q(W i ), the key observation is that i<n Q(W i ) -Q i (W i ) is a martingale. Using a martingale Bernstein bound, we prove the following lemma; applying it finishes the proof of Theorem 4.1. Lemma 4.3. Given any δ ∈ (0, 1), with probability

1 -δ, t<i Q(W t ) ≤ 4 t<i Q t (W t ) + 4 ln 1 δ .

5. ON SEPARABILITY:

In this section we give some discussion on Assumption 2.1, the separability of the NTK. The proofs are all given in Appendix D.

Given a training set (x i , y i ) n i=1 , the linear kernel is defined as K 0 (x i , x j ) := x i , x j . The maximum margin achievable by a linear classifier is given by γ 0 := min q∈∆n (q y) K 0 (q y).

(5.1)

where ∆ n denotes the probability simplex and denotes the Hadamard product. In addition to the dual definition eq. ( 5.1), when γ 0 > 0 there also exists a maximum margin classifier ū which gives a primal characterization of γ 0 : it holds that ū 2 = 1 and y i ū, x i ≥ γ 0 for all i.

In this paper we consider another kernel, the infinite-width NTK with respect to the first layer:

K 1 x i , x j := E ∂f (x i ; W 0 , a) ∂W 0 , ∂f (x j ; W 0 , a) ∂W 0 = E w∼N (0,I d ) x i 1 x i , w > 0 , x j 1 x j , w > 0 = φ i , φ j H .

Here φ and H are defined at the beginning of Section 2. Similar to the dual definition of γ 0 , the margin given by K 1 is defined as γ 1 := min q∈∆n (q y) K 1 (q y).

(5.2)

We can also give a primal characterization of γ 1 when it is positive. Proposition 5.1. If γ 1 > 0, then there exists v ∈ H such that v H = 1, and

y i v, φ i H ≥ γ 1 for any 1 ≤ i ≤ n. Additionally v(z) 2 ≤ 1/γ 1 for any z ∈ R d .

The proof is given in Appendix D, and uses the Fenchel duality theory. Using the upper bound v(z) 2 ≤ 1/γ 1 , we can see that γ 1 v satisfies Assumption 2.1 with γ ≥ γ 2 1 . However, such an upper bound v(z) 2 ≤ 1/γ 1 might be too loose, which leads to a bad rate. In fact, as shown later, in some cases we can construct v directly which satisfies Assumption 2.1 with a large γ. For this reason, we choose to make Assumption 2.1 instead of assuming a positive γ 1 .

However, we can use γ 1 to show that Assumption 2.1 always holds when there are no parallel inputs. Oymak & Soltanolkotabi (2019, Corollary I. 2) prove that if for any two feature vectors x i and x j , we have x i -x j 2 ≥ θ and x i + x j 2 ≥ θ for some θ > 0, then the minimum eigenvalue of K 1 is at least θ/(100n 2 ). For arbitrary labels y ∈ {-1, +1} n , since q y 2 ≥ 1/ √ n, we have the worst case bound γ 2 1 ≥ θ /100n 3 . A direct improvement of this bound is θ /100n 3 S , where n S denotes the number of support vectors, which could be much smaller than n with real world data.

On the other hand, given any training set (x i , y i ) n i=1 which may have a large margin, replacing y with random labels would destroy the margin, which is what should be expected.

Proposition 5.2. Given any training set (x i , y i ) n i=1 , if the true labels y are replaced with random labels ∼ unif {-1, +1} n , then with probability 0.9 over the random labels, it holds that γ 1 ≤ 1/ √ 20n.

Although the above bounds all have a polynomial dependency on n, they hold for arbitrary or random labels, and thus do not assume any relationship between the features and labels. Next we give some examples where there is a strong feature-label relationship, and thus a much larger margin can be proved.

5.1. THE LINEARLY SEPARABLE CASE:

Suppose the data distribution is linearly separable with margin γ 0 : there exists a unit vector ū such that y ū, x ≥ γ 0 almost surely. Then we can define v(z) := ū for any z ∈ R d . For almost all (x, y), we have

y v(z), x 1 z, x > 0 dµ N (z) = y ū, x 1 z, x > 0 dµ N (z) ≥ γ 1 z, x > 0 dµ N (z) = γ 0 2 ,

and thus Assumption 2.1 holds with γ = γ 0 /2.

5.2. THE NOISY 2-XOR DISTRIBUTION:

We consider the noisy 2-XOR distribution introduced in (Wei et al., 2018) . It is the uniform distribution over the following 2 d points:

(x 1 , x 2 , y, x 3 , . . . , x d ) ∈ 1 √ d -1 , 0, 1 , 0, 1 √ d -1 , -1 , -1 √ d -1 , 0, 1 , 0, -1 √ d -1 , -1 × -1 √ d -1 , 1 √ d -1 d-2

.

The factor 1 / √ d-1 ensures that x 2 = 1, and × above denotes the Cartesian product. Here the label y only depends on the first two coordinates of the input x.

To construct v, we first decompose R 2 into four regions:

A 1 := (z 1 , z 2 ) z 1 ≥ 0, |z 1 | ≥ |z 2 | , A 2 := (z 1 , z 2 ) z 2 > 0, |z 1 | < |z 2 | , A 3 := (z 1 , z 2 ) z 1 ≤ 0, |z 1 | ≥ |z 2 | \ {(0, 0)}, A 4 := (z 1 , z 2 ) z 2 < 0, |z 1 | < |z 2 | .

Then v can de defined as follows. It only depends on the first two coordinates of z.

v(z) :=          (1, 0, 0, . . . , 0) if (z 1 , z 2 ) ∈ A 1 , (0, -1, 0, . . . , 0) if (z 1 , z 2 ) ∈ A 2 , (-1, 0, 0, . . . , 0) if (z 1 , z 2 ) ∈ A 3 , (0, 1, 0, . . . , 0) if (z 1 , z 2 ) ∈ A 4 .

(5.

3)

The following result shows that γ = Ω(1/d). Note that n could be as large as 2 d , in which case γ is basically O 1/ ln(n) . Proposition 5.3. For any (x, y) sampled from the noisy 2-XOR distribution and any d ≥ 3, it holds that

y v(z), x 1 z, x > 0 dµ N (z) ≥ 1 60d

.

We can prove two other interesting results for the noisy 2-XOR data.

The width needs a poly(1/γ) dependency for initial separability. The first step of an NTK analysis is to show that ∇f i (W 0 ), y i n i=1 is separable. Proposition 5.4 gives an example where ∇f i (W 0 ), y i n i=1 is nonseparable when the network is narrow. Proposition 5.4. Let D = {(x i , y i )} 4 i=1 denote an arbitrary subset of the noisy 2-XOR dataset such that x i 's have the same last (d -2) coordinates. For any d ≥ 20, if m ≤ √ d -2/4, then with probability 1/2 over the random initialization of W 0 , for any weights V ∈ R m×d , it holds that y i V, ∇f i (W 0 ) ≤ 0 for at least one i ∈ {1, 2, 3, 4}.

For the noisy 2-XOR data, the separator v given by eq. ( 5.3) has margin γ = Ω(1/d), and 1/γ = O(d). As a result, if we want ∇f i (W 0 ), y i n i=1 to be separable, the width has to be Ω(1/ √ γ).

For a smaller width, gradient descent might still be able to solve the problem, but a beyond-NTK analysis would be needed.

A tight sample complexity upper bound for the infinite-width NTK. (Wei et al., 2018) give a d 2 sample complexity lower bound for any NTK classifier on the noisy 2-XOR data. It turns out that γ could give a matching sample complexity upper bound for the NTK and SGD. (Wei et al., 2018) consider the infinite-width NTK with respect to both layers. For the first layer, the infinite-width NTK K 1 is defined in Section 5, and the corresponding RKHS H and RKHS mapping φ is defined in Section 2. For the second layer, the infinite width NTK is defined by

K 2 x i , x j := E ∂f (x i ; W 0 , a) ∂a , ∂f (x j ; W 0 , a) ∂a = E w∼N (0,I d ) σ w, x i σ w, x j .

The corresponding RKHS K and inner product w 1 , w 2 K are given by

K := w : R d → R w(z) 2 dµ N (z) < ∞ , and w 1 , w 2 K = w 1 (z)w 2 (z) dµ N (z).

Given any x ∈ R d , it is mapped into ψ x ∈ K, where ψ x (z) := σ z, x . It holds that K 2 (x i , x j ) = ψ xi , ψ xj K . The infinite-width NTK with respect to both layers is just K 1 +K 2 . The corresponding RHKS is just H × K with the inner product

(v 1 , w 1 ), (v 2 , w 2 ) H×K = v 1 , v 2 H + w 1 , w 2 K .

The classifier v considered in eq. ( 5.3) has a unit norm (i.e., v H = 1) and margin γ on the space H. On H × K, it is enough to consider (v, 0), which also has a unit norm and margin γ. Since the infinite-width NTK model is a linear model in H × K, (Ji & Telgarsky, 2018, Lemma 2.5 ) can be used to show that SGD on the RKHS H×K could obtain a test error of with a sample complexity of O( 1 /γ 2 ). (The analysis in (Ji & Telgarsky, 2018) is done in R d , but it still works with a well-defined inner product.) Since γ = Ω(1/d), to achieve a constant test accuracy we need O(d 2 ) samples. This mathces (up to logarithmic factors) the sample complexity lower bound of d 2 given by Wei et al. (2018) .

6. OPEN PROBLEMS:

In this paper, we analyze gradient descent on a two-layer network in the NTK regime, where the weights stay close to the initialization. It is an interesting open question if gradient descent learns something beyond the NTK, after the iterates move far enough from the initial weights. It is also interesting to extend our analysis to other architectures, such as multi-layer networks, convolutional networks, and residual networks. Finally, in this paper we only discuss binary classification; it is interesting to see if it is possible to get similar results for other tasks, such as regression.

A OMITTED PROOFS FROM SECTION 2

Proof of Lemma 2.3. By Assumption 2.1, given any

1 ≤ i ≤ n, µ := E w∼N (0,I d ) y i v(w), x i 1 w, x i > 0 ≥ γ.

On the other hand,

y i f (0) i U = 1 m m s=1 y i v(w s,0 ), x i 1 w s,0 , x i > 0

is the empirical mean of i.i.d. r.v.'s supported on [-1, +1] with mean µ. Therefore by Hoeffding's inequality, with probability 1 -δ /n,

y i f (0) i U -γ ≥ y i f (0) i U -µ ≥ - 2 ln(n/δ) m .

Applying a union bound finishes the proof.

Proof of Lemma 2.4. Given any fixed 2 and 1 ≤ i ≤ n,

E α i (W 0 , 2 ) = P w, x i ≤ 2 ≤ 2 2 √ 2π = 2 π 2 ,

because w, x i is a standard Gaussian r.v. and the density of standard Gaussian has maximum 1/ √ 2π. Since α i (W 0 , 2 ) is the empirical mean of Bernoulli r.v.'s, by Hoeffding's inequality, with probability 1 -δ /n,

α i (W 0 , 2 ) ≤ E α i (W 0 , 2 ) + ln(n/δ) 2m ≤ 2 π 2 + ln(n/δ) 2m .

Applying a union bound finishes the proof.

To prove Lemma 2.5, we need the following technical result.

Lemma A.1. Consider the random vector X = (X 1 , . . . , X m ), where X i = σ(Z i ) for some σ : R → R that is 1-Lipschitz, and Z i are i.i.d. standard Gaussian r.v.'s. Then the r.v. X 2 is 1-sub-Gaussian, and thus with probability 1 -δ,

X 2 -E X 2 ≤ 2 ln(1/δ). Proof. Given a ∈ R m , define f (a) = m i=1 σ(a i ) 2 = σ(a) 2 ,

where σ(a) is obtained by applying σ coordinate-wisely to a. For any a, b ∈ R m , by the triangle inequality, we have

f (a) -f (b) = σ(a) 2 -σ(b) 2 ≤ σ(a) -σ(b) 2 = m i=1 σ(a i ) -σ(b i ) 2 ,

and by further using the 1-Lipschitz continuity of σ, we have

f (a) -f (b) ≤ m i=1 σ(a i ) -σ(b i ) 2 ≤ m i=1 (a i -b i ) 2 = a -b 2 .

As a result, f is a 1-Lipschitz continuous function w.r.t. the 2 norm, indeed f (X) is 1-sub-Gaussian and the bound follows by Gaussian concentration (Wainwright, 2015, Theorem 2.4 ).

Proof of Lemma 2.5. Given 1 ≤ i ≤ n, let h i = σ(W 0 x i )/ √ m. By Lemma A.1, h i 2 is sub-Gaussian with variance proxy 1/m, and with probability at least 1 -δ /2n over W 0 ,

h i 2 -E h i 2 ≤ 2 ln(2n/δ) m ≤ 2 ln(2n/δ) 25 ln(2n/δ) ≤ 1 - √ 2 2 .

On the other hand, by Jensen's inequality,

E h i 2 ≤ E h i 2 2 = √ 2 2 .

As a result, with probability 1 -δ /2n, it holds that h i 2 ≤ 1. By a union bound, with probability 1 -δ /2 over W 0 , for all 1 ≤ i ≤ n, we have h i 2 ≤ 1.

For any W 0 such that the above event holds, and for any 1 ≤ i ≤ n, the r.v. h i , a is sub-Gaussian with variance proxy h i 2 2 ≤ 1. By Hoeffding's inequality, with probability 1 -δ /2n over a,

h i , a = f (x i ; W 0 , a) ≤ 2 ln 4n/δ .

By a union bound, with probability 1 -δ /2 over a, for all 1 ≤ i ≤ n, we have f (x i ; W 0 , a) ≤ 2 ln 4n/δ .

The probability that the above events all happen is at least (1

-δ /2)(1 -δ /2) ≥ 1 -δ, over W 0 and a.

Proof of Lemma 2.6. We have

W t+1 -W 2 F = W t -W 2 F -2η t ∇ R(W t ), W t -W + η 2 t ∇ R(W t ) 2 F . (A.1) Therefore Lemma 2.4 ensures that |S c | ≤ s | w s,0 , x i | ≤ 2 ≤ m 4λ γ √ m + 1 2 = m 4λ γ √ m + γ 2 16 .

and thus

y i ∇f i (W t ) -∇f i (W 0 ), W 0 ≤ 1 √ m • |S c | • 4λ γ √ m ≤ 16λ 2 γ 2 √ m + λγ 4 ≤ λγ 2 , (A.6)

where in the last step we use the condition that m ≥ 4096λ 2 /γ 6 .

• The third term of eq. (A.4) can be bounded as follows: by Lemma 2.3,

y i ∇f i (W t ), U = y i ∇f i (W 0 ), U + y i ∇f i (W t ) -∇f i (W 0 ), U ≥ γ -1 + y i ∇f i (W t ) -∇f i (W 0 ), U .

In addition,

y i ∇f i (W t ) -∇f i (W 0 ), U = y i 1 m m i=1 1 w s,t , x i > 0 -1 w s,0 , x i > 0 v(w s,0 ), x i ≥ - 1 m • |S c | ≥ - 4λ γ √ m - 1 2 ≥ - γ 2 16 - 1 2 ,

where we use m ≥ 4096λ 2 /γ 6 . Therefore,

y i ∇f i (W t ), U ≥ γ -1 - γ 2 16 - 1 2 = γ - γ 2 4 ≥ 3γ 4 . (A.7)

Putting eqs. (A.5) to (A.7) into eq. (A.4), we have

y i ∇f i (W t ), W ≥ -2 ln 4n δ - λγ 2 + 3λγ 4 = λγ 4 -2 ln 4n δ = ln 4 ,

for the λ given in the statement of Theorem 2.2. Consequently, for any 0 ≤ t < t 1 , it holds that R (t) W ≤ /4.

Let T := 2λ 2 /η . The next claim is that t 1 ≥ T . To see this, note that Lemma 2.6 ensures

W t1 -W 2 F ≤ W 0 -W 2 F + 2η   t<t1 R (t) W   ≤ λ 2 + 2 ηt 1 .

Suppose t 1 < T , then we have t 1 ≤ 2λ 2 /η , and thus

W t1 -W 2 F ≤ 2λ 2 .

As a result, using

U F ≤ 1 and the definition of W , √ 2λ ≥ W t1 -W F ≥ W t1 -W , U = W t1 -W 0 , U -W -W 0 , U ≥ W t1 -W 0 , U -λ.

Moreover, due to eq. (A.7),

W t1 -W 0 , U = -η τ <t1 ∇ R(W τ ), U = η τ <t1 1 n n i=1 -y i f i (W τ ) y i ∇f i (W τ ), U ≥ η τ <t1 Q(W τ ) 3γ 4 .

As a result,

η τ <t1 Q(W τ ) ≤ 4( √ 2 + 1)λ 3γ ≤ 4λ γ .

Furthermore, by the triangle inequality, for any

1 ≤ s ≤ m w s,t -w s,0 2 ≤ η τ <t 1 n n i=1 y i f i (W τ ) y i ∂f i ∂w s,τ 2 ≤ η τ <t 1 n n i=1 y i f i (W τ ) • ∂f i ∂w s,τ 2 ≤ η τ <t Q(W τ ) 1 √ m ≤ η τ <t1 Q(W τ ) 1 √ m ≤ 4λ γ √ m , (A.8)

which contradicts the definition of t 1 . Therefore t 1 ≥ T . Now we are ready to prove the claims of Theorem 2.2. The bound on w s,t -w s,0 2 follow by repeating the steps in eq. (A.8). The risk guarantee follows from Lemma 2.6:

1 T t<T R(W t ) ≤ W 0 -W 2 F ηT + 2 T t<T R (t) W ≤ 2 + 2 = .

B OMITTED PROOFS FROM SECTION 3

The proof of Theorem 3.2 is based on Rademacher complexity. Given a sample S = (z 1 , . . . , z n ) (where z i = (x i , y i )) and a function class H, the Rademacher complexity of H on S is defined as

Rad (H • S) := 1 n E ∼{-1,+1} n   sup h∈H n i=1 i h(z i )   .

We will use the following general result.

Lemma B.1. (Shalev-Shwartz & Ben-David, 2014, Theorem 26.5) If h(z) ∈ [a, b], then with prob- ability 1 -δ, sup h∈H   E z∼D h(z) - 1 n n i=1 h(z i )   ≤ 2Rad (H • S) + 3(b -a) ln(2/δ) 2n .

We also need the following contraction lemma. Consider a feature sample X = (x 1 , . . . , x n ) and a function class F on X. For each 1 ≤ i ≤ n, let g i : R → R denote a K-Lipschitz function. Let g • F denote the class of functions which map x i to g i (f (x i )) for some f ∈ F.

Lemma B.2. (Shalev-Shwartz & Ben-David, 2014, Lemma 26.9 )

Rad (g • F • X) ≤ KRad (F • X).

To prove Theorem 3.2, we need one more Rademacher complexity bound. Given a fixed initialization (W 0 , a), consider the following classes:

EQUATION

F ρ := x → f (x; W, a) W ∈ W ρ .

Given a feature sample X, the following Lemma B.3 controls the Rademacher complexity of F ρ •X.

A similar version was given in (Liang, 2016, Theorem 43) , and the proof is similar to the proof of (Bartlett & Mendelson, 2002, Theorem 18) which also pushes the supremum through and handles each hidden unit separately.

Lemma B.3. Rad F ρ • X ≤ ρ m/n.

Proof of Lemma B.3. We have Invoking the Rademacher complexity of linear classifiers (Shalev-Shwartz & Ben-David, 2014, Lemma 26.10 ) then gives

E   sup W ∈Wρ n i=1 i f (x i ; W, a)   = E   sup W ∈Wρ n i=1 i m s=1 1 √ m a s σ w s , x i   = E   1 √ m sup W ∈Wρ m s=1 n i=1 i a s σ w s , x i   = E    1 √ m m s=1   sup ws-ws,0 2 ≤ρ n i=1 i a s σ w s , x i      = 1 √ m m i=1 E   sup

Rad F ρ • X = 1 n E   sup W ∈Wρ n i=1 i f (x i ; W, a)   ≤ ρ √ m √ n .

Now we are ready to prove the main generalization result Theorem 3.2.

Proof. Fix an initialization (W 0 , a), and let H := (x, y) → -yf (x) f ∈ F ρ . Since for any h ∈ H and any z, h(z) ∈ [0, 1], Lemma B.1 ensures that with probability 1 -δ over the data sampling,

sup h∈H   E z∼D h(z) - 1 n n i=1 h(z i )   = sup W ∈Wρ Q(W ) -Q(W ) ≤ 2Rad (H • S) + 3 ln(2/δ) 2n .

Since for each 1 ≤ i ≤ n, the mapping z → -(y i z) is (1/4)-Lipschitz, Lemma B.2 further ensures that Rad (H • S) ≤ Rad F ρ • X /4, and thus

sup W ∈Wρ Q(W ) -Q(W ) ≤ ρ √ m 2 √ n + 3 ln(2/δ) 2n . (B.1)

On the other hand, Theorem 2.2 ensures that under the conditions of Theorem 3.2, for any fixed dataset, with probability 1 -3δ over the random initialization, we have

Q(W k ) ≤ R(W k ) ≤ , and w s,k -w s,0 2 ≤ 4λ γ √ m .

As a result, invoking eq. (B.1) with ρ = 4λ/(γ √ m), with probability 1 -4δ over the random initialization and data sampling,

Q(W k ) ≤ Q(W k ) + 2λ γ √ n + 3 ln(2/δ) 2n ≤ + 8 2 ln(4n/δ) + ln(4/ ) γ 2 √ n + 3 ln(2/δ) 2n .

Invoking P (x,y)∼D yf (x; W, a) ≤ 0 ≤ 2Q(W ) finishes the proof.

C OMITTED PROOFS FROM SECTION 4:

Proof of Lemma 4.2. Recall that ∇f t (W t ) F ≤ 1, we have

W t+1 -W 2 F ≤ W t -W 2 F -2η y t f t (W t ) y t ∇f t (W t ), W t -W + η 2 y t f t (W t ) 2 . (C.1)

Similar to the proof of Lemma 2.6, the first order term of eq. (C.1) can be handled using the convexity of and homogeneity of ReLU as follows

y t f t (W t ) y t ∇f t (W t ), W t -W ≥ R t (W t ) -R t W , (C.2)

and the second-order term of eq. (C.1) can be bounded as follows 

EQUATION

ηR t (W t ) ≤ W t -W 2 F -W t+1 -W 2 F + 2ηR t W .

Telescoping gives the claim.

With Lemma 4.2, we give the following result, which is an extension of Theorem 2.2 to the SGD setting.

Lemma C.1. Under Assumption 3.1, given any ∈ (0, 1), any δ ∈ (0, 1/3), and any positive integer n 0 , let λ := 2 ln(4n 0 /δ) + ln(4/ ) γ/4 , and M := 4096λ 2 γ 6 .

For any m ≥ M and any constant step size η ≤ 1, if n 0 ≥ n := 2λ 2 /η , then with probability

1 -3δ, 1 n i<n Q i (W i ) ≤ .

Proof. We first sample n 0 data examples (x 0 , y 0 ), . . . , (x n0-1 , y n0-1 ), and then feed (x i , y i ) to SGD at step i. We only consider the first n 0 steps.

The proof is similar to the proof of Theorem 2.2. Let n 1 denote the first step before n 0 such that there exists some 1 ≤ s ≤ m with w s,n1 -w s,0 2 > 4λ/(γ √ m). If such a step does not exist, let n 1 = n 0 .

Let W := W 0 + λU , in exactly the same way as in Theorem 2.2, we can show that with probability 1 -3δ, for any 0 ≤ i < n 1 , y i ∇f i (W i ), W ≥ ln 4 , and thus R i W ≤ /4. Now consider n := 2λ 2 /η . Using Lemma 4.2, in the same way as the proof of Theorem 2.2 (replacing Q(W τ ) with Q i (W i ), etc.), we can show that n ≤ n 1 . Then invoking Lemma 4.2 again, we get

1 n i<n Q i (W i ) ≤ 1 n i<n R i (W i ) ≤ W 0 -W 2 F ηn + 2 n i<n R i W ≤ 2 + 2 = .

Next we prove Lemma 4.3. We need the following martingale Bernstein bound. Lemma C.2. (Beygelzimer et al., 2011 , Theorem 1) Let (M t , F t ) t≥0 denote a martingale with M 0 = 0 and F 0 be the trivial σ-algebra. Let (∆ t ) t≥1 denote the corresponding martingale difference sequence, and let

V t := t j=1 E ∆ 2 j F j-1

denote the sequence of conditional variance. If ∆ t ≤ R a.s., then for any δ ∈ (0, 1), with probability at least 1 -δ,

M t ≤ V t R (e -2) + R ln 1 δ .

Proof of Lemma 4.3. For any i ≥ 0, let z i denote (x i , y i ), and z 0,i denote (z 0 , . . . , z i ). Note that the quantity t<i Q(W t ) -Q t (W t ) is a martingale w.r.t. the filtration σ(z 0,i-1 ). The martingale difference sequence is given by Q 

(W t ) -Q t (W t ), which satisfies Q(W t ) -Q t (W t ) = E (x,y)∼D -yf (x; W t , a) + y t f (x t ; W t , a) ≤ 1, (C.4) since -1 ≤ ≤ 0. Moreover, we have E Q(W t ) -Q t (W t ) 2 σ(z 0,t-1 ) = Q(W t ) 2 -2Q(W t )E Q t (W t ) σ(z 0,t-1 ) + E Q t (W t ) 2 σ(z 0,t-1 ) = -Q(W t ) 2 + E Q t (W t ) 2 σ(z 0,t-1 ) ≤ E Q t (W t ) 2 σ(z 0,t-1 ) ≤ E Q t (W t ) σ(z 0,t-1 ) = Q(W t ).

-δ, t<i Q(W t ) -Q t (W t ) ≤ (e -2) t<i Q(W t ) + ln 1 δ . Consequently, t<i Q(W t ) ≤ 4 t<i Q t (W t ) + 4 ln 1 δ .

Finally, we prove Theorem 4.1.

Proof of Theorem 4.1. Suppose the condition of Lemma C.1 holds. Then we have for n = 2λ 2 /η , with probability 1 -3δ,

1 n i<n Q i (W i ) ≤ .

Further invoking Lemma 4.3 gives that with probability 1 -4δ,

1 n i<n Q(W i ) ≤ 4 n i<n Q i (W i ) + 4 n ln 1 δ ≤ 5 .

Since P (x,y)∼D yf (x; W, a) ≤ 0 ≤ 2Q(W ), we get

1 n n i=1 P (x,y)∼D yf (x; W i , a) ≤ 0 ≤ 10 .

For the condition of Lemma C.1 to hold, it is enough to let

n 0 = Θ ln(1/δ) ηγ 2 2 , which gives M = Θ ln(1/δ) + ln(1/ ) 2 γ 8 and n = Θ ln(1/δ) + ln(1/ ) 2 γ 2 . D OMITTED PROOFS FROM SECTION 5 Proof of Proposition 5.1. Define f : H → R by f (w) := 1 2 w(z) 2 2 dµ N (z) = 1 2 w 2 H .

It holds that f is continuous, and f * has the same form. Define g : R n → R by

g(p) := max 1≤i≤n p i , with conjugate g * (q) = 0, if q ∈ ∆ n , +∞, o.w.

Finally, define the linear mapping A : H → R n by (Aw) i = y i w, φ i H .

Since f , f * , g and g * are lower semi-continuous, and dom g -Adom f = R n , and dom f * -A * dom g * = H, Fenchel duality may be applied in each direction (Borwein & Zhu, 2005, Theorem 4.4.3) , and ensures that

inf w∈H f (w) + g(Aw) = sup q∈R n -f * (A * q) -g * (-q) .

with optimal primal-dual solutions ( w, q). Moreover

inf w∈H f (w) + g(Aw) = inf w∈H,u∈R n sup q∈R n f (w) + g(Aw + u) + q, u ≥ sup q∈R n inf w∈H,u∈R n f (w) + g(Aw + u) + q, u = sup q∈R n inf w∈H,u∈R n f (w) -A * q, w H + g(Aw + u) --q, Aw + u = sup q∈R n -f * (A * q) -g * (-q) .

By strong duality, the inequality holds with equality. It follows that w = A * q, and supp(-q) ⊂ arg max 1≤i≤n (A w) i .

where eq. (D.1) is due to the independence between z 1,2 and z 3,d , and in eq. (D.2) we use the fact that v(z) 1,2 only depends on z 1,2 and v(z) 3,d are all zero. Since v(z) 1,2 , x 1,2 = 0 for z 1,2 ∈ A 2 ∪ A 4 , we only need to consider A 1 and A 3 in eq. (D.3). For simplicity, we will denote z 1,2 by p ∈ R 2 , and v(z) 1,2 by v(p), and z 3,d by q ∈ R d-2 . For any nonzero p ∈ A 1 , we have -p ∈ A 3 , and v(p), x 1,2 = 1/ √ d -1. Therefore y v(p), x 1,2 1 p, x 1,2 + q, x 3,d > 0 dµ N (q)

+ y v(-p), x 1,2 1 -p, x 1,2 + q, x 3,d > 0 dµ Since q, x 3,d is a Gaussian variable with standard deviation (d-2) /(d-1), we have 

EQUATION

EQUATION

(z) = 1 √ d -1 U p 1 √ d -2 1 [p ∈ A 1 ] dµ N (p) = 1 √ d -1 ∞ 0 U p 1 √ d -2 p1 -p1

ϕ(p 2 ) dp 2 ϕ(p 1 ) dp

1 = 1 √ d -1 ∞ 0 U p 1 √ d -2 U (p 1 )ϕ(p 1 ) dp 1 ≥ 1 √ d -1 1 0 U p 1 √ d -2 U (p 1 )ϕ(p 1 ) dp 1 .

For t ∈ [-1, +1], it holds that ϕ(t) ≥ 1 √ 2πe, and thus

U (a) = a -a ϕ(t) dt ≥ 2a √ 2πe .

Therefore eq. (D.3) is lower bounded by

1 √ d -1 1 0 U p 1 √ d -2 U (p 1 )ϕ(p 1 ) dp 1 ≥ 1 √ d -1 1 0 2 √ 2πe • p 1 √ d -2 • 2p 1 √ 2πe • 1 √ 2πe dp 1 ≥ 1 20 (d -1)(d -2) 1 0 p 2 1 dp 1 = 1 60 (d -1)(d -2) ≥ 1 60d

.

To prove Proposition 5.4, we need the following technical lemma. Lemma D.1. Given z 1 ∼ N (0, 1) and z 2 ∼ N (0, b 2 ) that are independent where b > 1, we have

P |z 1 | < |z 2 | > 1 - 1 b .

Proof. First note that for z 3 ∼ N (0, 1) which is independent of z 1 ,

P |z 1 | < |z 2 | = P |z 1 | < b|z 3 | = 1 -P |z 3 | < 1 b |z 1 | .

Still let ϕ denote the density of N (0, 1), and let U (c) denote the probability that z 3 ∈ [-c, c]. We have

P |z 3 | < 1 b |z 1 | = 1 |z 3 | < 1 b |z 1 | ϕ(z 3 )ϕ(z 1 ) dz 3 dz 1 = U 1 b |z 1 | ϕ(z 1 ) dz 1 ≤ 2 √ 2πb |z 1 |ϕ(z 1 ) dz 1 = 2 πb < 1 b ,

where we use the facts that U (c) ≤ 2c/ √ 2π and E[|z 1 |] = 2/π.

We now give the proof of Proposition 5.4 using Lemma D.1.

Proof of Proposition 5.4. By symmetry, we only need to consider the following training set:

x 1 = (1, 0, 1, . . . , 1), y 1 = 1, x 2 = (0, 1, 1, . . . , 1), y 2 = -1, x 3 = (-1, 0, 1, . . . , 1), y 3 = 1, x 4 = (0, -1, 1, . . . , 1), y 4 = -1.

The 1/ √ d -1 factor is omitted also because we only discuss the 0/1 loss.

For any s, let A s denote the event that 1 w s , x 1 > 0 = 1 w s , x 2 > 0 = 1 w s , x 3 > 0 = 1 w s , x 4 > 0 .

We will show that if m ≤ √ d -2/4, then A s is true for all 1 ≤ s ≤ m with probability 1/2, and Proposition 5.4 follows from the fact that the XOR data is not linearly separable.

For any s and i, w s , x i = (w s ) 1 (x i ) 1 + (w s ) 2 (x i ) 2 + d j=3 (w s ) j .

Since (x i ) 1 , (x i ) 2 is (1, 0) or (0, 1) or (-1, 0) or (0, -1), event A s will happen as long as 

2 √ d -2 m ≥ 1 - 2 √ d -2 √ d -2 4 = 1 2 .

:



Figures and Tables:

Note that for any 1 ≤ s ≤ m, the mapping z → a s σ(z) is 1-Lipschitz, and thus Lemma B.

C.3) since η, -≤ 1 and -≤ . Combining eqs. (C.1) to (C.3) gives

eqs. (C.4) and (C.5) gives that with probability 1

ϕ denote the density function of the standard Gaussian distribution, and for c > 0, let U (c) denote the probability that a standard Gaussian random variable lies in the interval [-c, c]: U (c) := c -c ϕ(t) dt.

.5) Plugging eqs. (D.4) and (D.5) into eq. (D.3) gives: y v(z), x 1 z, x > 0 dµ N

) j , and(w s ) 2 < d s=3 (w s ) j . Note that (w s ) 1 , (w s ) 2 ∼ N (0, 1) while d j=3 (w s ) j ∼ N (0, d-2). As a result, due to Lemma D.1, Using a union bound, P(A s ) > 1 -2 / √ d-2. If m ≤ √ d -2/4, then by a union bound again,

Peer review (iclr2020):
Summary and Decision The authors of this paper studied the optimization and generalization properties of shallow ReLU networks. In particular, the authors were able to show a width dependence that is polylogarithmic on the number of samples, probability or failure, error tolerance, and a margin parameter. This work is unique in that the authors showed how to bound many key quantities in terms of the margin parameter, and drew a connection with the neural tangent kernel*s maximum margin. Furthermore, the overall reading experience was very smooth, although I do have some minor comments later. The main concern from me is on the implicit dependence of the margin parameter gamma, most notably this can lead to the width dependence to be polynomial in terms of the number of samples and the minimum separation distance. While this concern warrants a careful discussion (below), I believe the paper still offers a nice analysis of shallow networks. Overall, I would recommend accept for this paper. Background There has been a large number of works studying very wide networks, showing both optimization and generalization results. While there has been great progress, most of these existing results require the width of both deep (and shallow) networks to be very large. For example, even by being polynomial in the number of samples, the networks are already unrealizable in practice. Therefore, guarantees with much better dependence is highly desirable. Discussion of Contributions As the title may suggest, it is a bit surprising that we can show that polylog width is sufficient. Intuitively, we can imagine that the classification margin can grow exponentially more complex as the number of samples increase. I believe the nice result can be attributed to a careful analysis of key quantities such as | W_t - W_0 |_F in terms of the margin parameter from Assumption 2.1. Some nice examples of the analysis in this paper include the introduction of the weight matrix overline{U} and overline{W} as an intermediate between W_0 and W_t, and observing that the activations of the ReLU xi_{i,t} do not change very much during training. The tricks together led to a very tight bound on the change in weights | W_t - W_0 |_F in terms of the margin parameter gamma. As the authors mentioned on page 6, this tight control was used to bound the Rademacher complexity later. The connection drawn between the margin assumption and neural tangent kernel (Proposition 5.1) is also interesting on its own. The authors intended this result to serve as a justification of the margin assumption (2.1). Discussion of the Margin Parameter Let me start by saying I*m not completely certain on how to interpret this margin parameter gamma in Assumption 2.1. Perhaps I*m missing some obvious ideas here, but I would still like the authors respond with some more details. At the same time, I don*t believe this is a sufficient criticism to reject this paper, as I believe the analysis in terms of gamma is still valuable. On one hand, if we were to assume the margin condition holds for all possible data points (i.e. Assumption 3.1), then there is no concern about polynomial dependence on the number of samples, and this is certainly a reasonable assumption in some applications. On the other hand, many of the previous analysis on wide networks were in terms of a minimum separation distance, i.e. assume there exists a delta > 0 such that for all i 
eq j, we have | x_i - x_j | geq delta . The authors have provided a discussion in section 5, including both a worst case bound of gamma geq delta / (100 n^2), by Oymak and Soltanokotabi (2019) and an example where the margin is O( n^{-1/2} ) with high probability. Using either bounds on gamma, we will have a width with polynomial dependence in terms of the number of samples and minimum separation distance. Therefore if we were to compare against previous works in the same benchmark, i.e. using a minimum separation assumption instead, then arguably this work did not achieve a width that is only polylog in terms of the number of samples. That being said, I don*t believe the authors were intentionally trying to hide sample dependence inside an assumption. The paper is presented in a very transparent way, and the authors were being honest in chapter 5 about the worst case dependence on the number of samples. To summarize, it is unclear to me whether the paper truly achieved a width dependence that is polylog in terms of the number of samples, but the analysis in terms of gamma remains a valuable contribution. I welcome the authors and the other reviewers to provide additional comments on whether the title and claims of this paper is appropriate. Minor Comments For the sake of improving readability, I also have some minor comments that do not contribute towards the decision. - On page 5, the first observation bullet point in section 2.2, it is written here that by triangle inequality | 
abla widehat{R}(W) |_F leq widehat{Q}(W) , I thought you needed an absolute value on the right hand side, perhaps you should mention that ell is strictly decreasing. - On page 6, in the statement of Corollary 3.3, you are missing eta leq 1, and perhaps the 	ilde on Omega and O should be defined. - On the same note, while it is obvious that Theorem 3.2 implies this corollary, it is still worth writing up a proof to compute the constants. - On page 7, the notation Delta_n and odot are not defined, I had to infer definition from the proof. - On page 12, just below the second equation, it wasn*t immediately clear how widehat{R}( overline{W} ) leq epsilon / 4. I believe it*s worth expanding the definitions a bit, and explicitly plug in Lemma 2.5. - On page 12, in the second last equation, I*m actually not sure where this inequality comes from | W_t - overline{W} |_F geq langle W_t - overline {W} , overline{U} angle Perhaps it*s obvious, but I currently don*t see it.