Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow ReLU networks

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/\delta$ or $1/\epsilon$, where $n$ denotes the size of the training set, $\delta$ denotes the failure probability, and $\epsilon$ denotes the target error. These large widths far exceed what is used empirically, constituting a significant gap between theory and practice.

In this paper, we narrow this gap by showing that a two-layer ReLU network with $\Omega(\ln(n/\delta)+\ln(1/\epsilon)^{2})$ hidden units trained by gradient descent achieves classification error $\epsilon$ on test data, meaning both optimization and generalization occur. Unlike prior work, the width is fully polylogarithmic in $n$, $1/\delta$, and $1/\epsilon$; 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.

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 $\widetilde{\Omega}(1/\epsilon^{2})$ samples and $\widetilde{\Theta}(1/\epsilon)$ iterations, gradient descent finds a solution with $\epsilon$ test error (cf. Theorem 3.2 and Corollary 3.3). (As discussed in Remark 3.4, $\widetilde{\Omega}(1/\epsilon)$ samples also suffice via a smoothness-based generalization bound, at the expense of large constant factors.)

considers stochastic gradient descent (SGD) with access to a standard stochastic online oracle. We prove that with width at least polylogarithmic and $\widetilde{\Theta}(1/\epsilon)$ samples, SGD achieves an arbitrarily small test error (cf. Theorem 4.1).

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/\sqrt{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 $\Omega(1/\ln(n))$, and our SGD sample complexity is tight in the infinite-width case.

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 and 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 $\Omega(1/\epsilon^{2})$ hidden units, $\widetilde{\Omega}(1/\epsilon^{4})$ data samples, and $O(1/\epsilon^{2})$ steps, while our result only needs polylogarithmic hidden units, $\widetilde{\Omega}(1/\epsilon^{2})$ data samples, and $\widetilde{O}(1/\epsilon)$ steps.

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 and Gu, 2019). Allen-Zhu et al. (2018c) show that SGD can minimize the regression loss for recurrent neural networks, and Allen-Zhu and Li (2019b) further prove a low generalization error. Allen-Zhu and 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 and 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 $\epsilon$ test error with $\Omega(1/\epsilon^{4})$ samples and $\Omega(1/\epsilon^{14})$ hidden units. Cao and Gu (2019b) consider SGD with an online oracle and give a general result. Under the same assumption as in (Cao and Gu, 2019a), their result requires $\Omega(1/\epsilon^{14})$ hidden units and sample complexity $\widetilde{O}(1/\epsilon^{2})$. By contrast, with the same online oracle, our result only needs polylogarithmic hidden units and sample complexity $\widetilde{O}(1/\epsilon)$.

2 Notation

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 and Liang, 2018; Du et al., 2018b; Arora et al., 2019; Nitanda and Suzuki, 2019). We consider the ReLU activation $\sigma(z)\mathrel{\mathop{\ordinarycolon}}=\max\mathinner{\left\{0,z\right\}}$, though our analysis can be extended easily to Lipschitz continuous, positively homogeneous activations such as leaky ReLU.

For any $t\geq 0$, the gradient descent step is given by $W_{t+1}\mathrel{\mathop{\ordinarycolon}}=W_{t}-\eta_{t}\nabla\widehat{\mathcal{R}}(W_{t})$. Also define

Note that $f_{i}^{(t)}(W_{t})=f_{i}(W_{t})$. This property generally holds due to homogeneity: for any $W$ and any $1\leq s\leq m$,

and thus $\left\langle\nabla f_{i}(W),W\right\rangle=f_{i}(W)$.

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:

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

and particularly define $\phi_{i}\mathrel{\mathop{\ordinarycolon}}=\phi_{x_{i}}$ for the training input $x_{i}$.

As discussed in Section 5, the space $\mathcal{H}$ is the reproducing kernel Hilbert space (RKHS) induced by the infinite-width NTK with respect to $W$, and $\phi_{x}$ maps $x$ into $\mathcal{H}$. Assumption 2.1 supposes that the induced training set $\{(\phi_{i},y_{i})\}_{i=1}^{n}$ can be separated by some $\bar{v}\in\mathcal{H}$, with an additional bound on $\mathinner{\!\left\lVert\bar{v}(z)\right\rVert}_{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 $\bar{v}$ directly; see Section 5 for some examples.

With Assumption 2.1, we state our main empirical risk result.

Under Assumption 2.1, given any risk target $\epsilon\in(0,1)$ and any $\delta\in(0,1/3)$, let

Then for any $m\geq M$ and any constant step size $\eta\leq 1$, with probability $1-3\delta$ over the random initialization,

Moreover for any $0\leq t<T$ and any $1\leq s\leq m$,

While the number of hidden units required by prior work all have a polynomial dependency on $n$, $1/\delta$ or $1/\epsilon$, Theorem 2.2 only requires $m=\Omega\mathinner{\left(\ln(n/\delta)+\ln(1/\epsilon)^{2}\right)}$. The required width has a polynomial dependency on $1/\gamma$, which is an adaptive quantity: while $1/\gamma$ can be $\operatorname{poly}(n)$ for random labels (cf. Proposition 5.2), it can be $\operatorname{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 \mathinner{\bigl{\{}\mathinner{\left(\nabla f_{i}(W_{0}),y_{i}\right)}\bigr{\}}}_{i=1}^{n} to be separable, which is the starting point of an NTK-style analysis, the width has to depend polynomially on $1/\gamma$.

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

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

Given an initialization $(W_{0},a)$, for any $1\leq s\leq m$, define

Lemma 2.3 ensures that with high probability $\overline{U}$ has a positive margin at initialization.

Under Assumption 2.1, given any $\delta\in(0,1)$ and any $\epsilon_{1}\in(0,\gamma)$, if $m\geq\mathinner{\left(2\ln(n/\delta)\right)}/\epsilon_{1}^{2}$, then with probability $1-\delta$, it holds simultaneously for all $1\leq i\leq n$ that

For any $W$, any $\epsilon_{2}>0$, and any $1\leq i\leq n$, define

Lemma 2.4 controls $\alpha_{i}(W_{0},\epsilon_{2})$. It will help us show that $\overline{U}$ has a good margin during the training process.

Under the condition of Lemma 2.3, for any $\epsilon_{2}>0$, with probability $1-\delta$, it holds simultaneously for all $1\leq i\leq n$ that

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

Given any $\delta\in(0,1)$, if $m\geq 25\ln(2n/\delta)$, then with probability $1-\delta$, it holds simultaneously for all $1\leq i\leq n$ that

2 Convergence analysis of gradient descent

We analyze gradient descent in this subsection. First, define

For any $W$ and any $1\leq s\leq m$, $\mathinner{\!\left\lVert\partial f_{i}/\partial w_{s}\right\rVert}_{2}\leq 1/\sqrt{m}$, and thus $\mathinner{\!\left\lVert\nabla f_{i}(W)\right\rVert}_{F}\leq 1$. Therefore by the triangle inequality, $\mathinner{\!\left\lVert\nabla\widehat{\mathcal{R}}(W)\right\rVert}_{F}\leq\widehat{\mathcal{Q}}(W)$.

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

For any $t\geq 0$ and any $\overline{W}$, if $\eta_{t}\leq 1$, then

Consequently, if we use a constant step size $\eta\leq 1$ for $0\leq\tau<t$, then

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

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

We first show that as long as $\|w_{s,t}-w_{s,0}\|_{2}\leq 4\lambda/(\gamma\sqrt{m})$ for all $1\leq s\leq m$, it holds that \widehat{\mathcal{R}}^{(t)}\mathinner{\bigl{(}W_{0}+\lambda\overline{U}\bigr{)}}\leq\epsilon/4. To see this, let us consider $\widehat{\mathcal{R}}^{(0)}$ first. For any $1\leq i\leq n$, Lemma 2.5 ensures that $|\langle\nabla f_{i}(W_{0}),W_{0}\rangle|$ is bounded, while Lemma 2.3 ensures that \big{\langle}\nabla f_{i}(W_{0}),\overline{U}\big{\rangle} is concentrated around $\gamma$ with a large width. As a result, with the chosen $\lambda$ in Theorem 2.2, we can show that \big{\langle}\nabla f_{i}(W_{0}),W_{0}+\lambda\overline{U}\big{\rangle} is large, and $\widehat{\mathcal{R}}^{(0)}(W_{0}+\lambda\overline{U})$ is small due to the exponential tail of the logistic loss. To further handle $\widehat{\mathcal{R}}^{(t)}$, we use a standard NTK argument to control \big{\langle}\nabla f_{i}(W_{t})-\nabla f_{i}(W_{0}),W_{0}+\lambda\overline{U}\big{\rangle} under the condition that $\|w_{s,t}-w_{s,0}\|_{2}\leq 4\lambda/(\gamma\sqrt{m})$.

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 $\widehat{\mathcal{R}}^{(t)}(W_{0}+\lambda\overline{U})\leq\epsilon/4$, we can use it and Lemma 2.6 to control $\sum_{\tau<t}\widehat{\mathcal{Q}}(W_{\tau})$, and then just invoke $\|w_{s,t}-w_{s,0}\|_{2}\leq\eta\sum_{\tau<t}\widehat{\mathcal{Q}}(W_{\tau})/\sqrt{m}$.

The quantity $\sum_{\tau<t}\widehat{\mathcal{Q}}(W_{\tau})$ has also been considered in prior work (Cao and Gu, 2019a; Nitanda and Suzuki, 2019), where it is bounded by $\sqrt{t}\sqrt{\sum_{\tau<t}\widehat{\mathcal{Q}}(W_{\tau})^{2}}$ using the Cauchy-Schwarz inequality, which introduces a $\sqrt{t}$ factor. To make the required width depend only on $\operatorname{polylog}(1/\epsilon)$, we also need an upper bound on $\sum_{\tau<t}\widehat{\mathcal{Q}}(W_{\tau})$ which depends only on $\operatorname{polylog}(1/\epsilon)$. Since the above analysis results in a $\sqrt{t}$ factor, and in our case $\Omega(1/\epsilon)$ steps are needed, it is unclear how to get a $\operatorname{polylog}(1/\epsilon)$ width using the analysis in (Cao and Gu, 2019a; Nitanda and Suzuki, 2019). By contrast, using Lemma 2.6, we can show that $\sum_{\tau<t}\widehat{\mathcal{Q}}(W_{\tau})\leq 4\lambda/\gamma$, which only depends on $\ln(1/\epsilon)$.

The claims of Theorem 2.2 then follow directly from the above two steps and Lemma 2.6.

Generalization

To get a generalization bound, we naturally extend Assumption 2.1 to the following assumption.

for almost all $(x,y)$ sampled from the data distribution $\mathcal{D}$.

The above assumption is also made in (Nitanda and Suzuki, 2019) for smooth activations. (Cao and 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.

Under Assumption 3.1, given any $\epsilon\in(0,1)$ and any $\delta\in(0,1/4)$, let $\lambda$ and $M$ be given as in Theorem 2.2:

Then for any $m\geq M$ and any constant step size $\eta\leq 1$, with probability $1-4\delta$ over the random initialization and data sampling,

where $k$ denotes the step with the minimum empirical risk before $\lceil\nicefrac{{2\lambda^{2}}}{{\eta\epsilon}}\rceil$.

Below is a direct corollary of Theorem 3.2.

Under Assumption 3.1, given any $\epsilon,\delta\in(0,1)$, using a constant step size no larger than $1$ and let

it holds with probability $1-\delta$ that $P_{(x,y)\sim\mathcal{D}}\mathinner{\left(yf(x;W_{k},a)\leq 0\right)}\leq\epsilon$, where $k$ denotes the step with the minimum empirical risk in the first $\widetilde{\Theta}(\nicefrac{{1}}{{\gamma^{2}\epsilon}})$ steps.

To get Theorem 3.2, we use a Lipschitz-based Rademacher complexity bound. One can also use a smoothness-based Rademacher complexity bound (Srebro et al., 2010, Theorem 1) and get a sample complexity $\widetilde{O}(\nicefrac{{1}}{{\gamma^{4}\epsilon}})$. However, the bound will become complicated and some large constant will be introduced. It is an interesting open question to give a clean analysis based on smoothness.

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)\mathrel{\mathop{\ordinarycolon}}=f(x_{i};W,a)$, and perform the following update

Still with Assumption 3.1, we show the following result.

Under Assumption 3.1, given any $\epsilon,\delta\in(0,1)$, using a constant step size and $m=\Omega\mathinner{\left(\nicefrac{{\mathinner{\left(\ln(1/\delta)+\ln(1/\epsilon)^{2}\right)}}}{{\gamma^{8}}}\right)}$, it holds with probability $1-\delta$ that

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

The first step is an extension of Lemma 2.6 to the SGD setting, with a similar proof.

With a constant step size $\eta\leq 1$, for any $\overline{W}$ and any $i\geq 0$,

With Lemma 4.2, we can also extend Theorem 2.2 to the SGD setting and get a bound on $\sum_{i<n}\mathcal{Q}_{i}(W_{i})$, using a similar proof. To further get a bound on the cumulative population risk $\sum_{i<n}\mathcal{Q}(W_{i})$, the key observation is that $\sum_{i<n}\mathinner{\left(\mathcal{Q}(W_{i})-\mathcal{Q}_{i}(W_{i})\right)}$ is a martingale. Using a martingale Bernstein bound, we prove the following lemma; applying it finishes the proof of Theorem 4.1.

Given any $\delta\in(0,1)$, with probability $1-\delta$,

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 $\mathinner{\left\{(x_{i},y_{i})\right\}}_{i=1}^{n}$, the linear kernel is defined as $K_{0}(x_{i},x_{j})\mathrel{\mathop{\ordinarycolon}}=\left\langle x_{i},x_{j}\right\rangle$. The maximum margin achievable by a linear classifier is given by

where $\Delta_{n}$ denotes the probability simplex and $\odot$ denotes the Hadamard product. In addition to the dual definition eq. 5.1, when $\gamma_{0}>0$ there also exists a maximum margin classifier $\bar{u}$ which gives a primal characterization of $\gamma_{0}$: it holds that $\|\bar{u}\|_{2}=1$ and $y_{i}\left\langle\bar{u},x_{i}\right\rangle\geq\gamma_{0}$ for all $i$.

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

Here $\phi$ and $\mathcal{H}$ are defined at the beginning of Section 2. Similar to the dual definition of $\gamma_{0}$, the margin given by $K_{1}$ is defined as

We can also give a primal characterization of $\gamma_{1}$ when it is positive.

The proof is given in Appendix D, and uses the Fenchel duality theory. Using the upper bound $\mathinner{\!\left\lVert\hat{v}(z)\right\rVert}_{2}\leq 1/\gamma_{1}$, we can see that $\gamma_{1}\hat{v}$ satisfies Assumption 2.1 with $\gamma\geq\gamma_{1}^{2}$. However, such an upper bound $\mathinner{\!\left\lVert\hat{v}(z)\right\rVert}_{2}\leq 1/\gamma_{1}$ might be too loose, which leads to a bad rate. In fact, as shown later, in some cases we can construct $\bar{v}$ directly which satisfies Assumption 2.1 with a large $\gamma$. For this reason, we choose to make Assumption 2.1 instead of assuming a positive $\gamma_{1}$.

However, we can use $\gamma_{1}$ to show that Assumption 2.1 always holds when there are no parallel inputs. Oymak and 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}\geq\theta$ and $\|x_{i}+x_{j}\|_{2}\geq\theta$ for some $\theta>0$, then the minimum eigenvalue of $K_{1}$ is at least $\theta/(100n^{2})$. For arbitrary labels $y\in\{-1,+1\}^{n}$, since $\mathinner{\!\left\lVert q\odot y\right\rVert}_{2}\geq 1/\sqrt{n}$, we have the worst case bound $\gamma_{1}^{2}\geq\nicefrac{{\theta}}{{100n^{3}}}$. A direct improvement of this bound is $\nicefrac{{\theta}}{{100n_{S}^{3}}}$, 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 $\mathinner{\left\{(x_{i},y_{i})\right\}}_{i=1}^{n}$ which may have a large margin, replacing $y$ with random labels would destroy the margin, which is what should be expected.

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.

and thus Assumption 2.1 holds with $\gamma=\gamma_{0}/2$.

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:

The factor $\nicefrac{{1}}{{\sqrt{d-1}}}$ ensures that $\|x\|_{2}=1$, and $\times$ above denotes the Cartesian product. Here the label $y$ only depends on the first two coordinates of the input $x$.

Then $\bar{v}$ can de defined as follows. It only depends on the first two coordinates of $z$.

The following result shows that $\gamma=\Omega(1/d)$. Note that $n$ could be as large as $2^{d}$, in which case $\gamma$ is basically $O\mathinner{\left(1/\ln(n)\right)}$.

For any $(x,y)$ sampled from the noisy 2-XOR distribution and any $d\geq 3$, it holds that

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

The first step of an NTK analysis is to show that \mathinner{\bigl{\{}\mathinner{\left(\nabla f_{i}(W_{0}),y_{i}\right)}\bigr{\}}}_{i=1}^{n} is separable. Proposition 5.4 gives an example where \mathinner{\bigl{\{}\mathinner{\left(\nabla f_{i}(W_{0}),y_{i}\right)}\bigr{\}}}_{i=1}^{n} is nonseparable when the network is narrow.

For the noisy 2-XOR data, the separator $\bar{v}$ given by eq. 5.3 has margin $\gamma=\Omega(1/d)$, and $1/\gamma=O(d)$. As a result, if we want \mathinner{\bigl{\{}\mathinner{\left(\nabla f_{i}(W_{0}),y_{i}\right)}\bigr{\}}}_{i=1}^{n} to be separable, the width has to be $\Omega(1/\sqrt{\gamma})$. 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 $\gamma$ 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 $\mathcal{H}$ and RKHS mapping $\phi$ is defined in Section 2. For the second layer, the infinite width NTK is defined by

The corresponding RKHS $\mathcal{K}$ and inner product $\langle w_{1},w_{2}\rangle_{\mathcal{K}}$ are given by

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.

The authors are grateful for support from the NSF under grant IIS-1750051, and from NVIDIA via a GPU grant.

References

Appendix A Omitted proofs from Section 2

By Assumption 2.1, given any $1\leq i\leq n$,

is the empirical mean of i.i.d. r.v.’s supported on $[-1,+1]$ with mean $\mu$. Therefore by Hoeffding’s inequality, with probability $1-\nicefrac{{\delta}}{{n}}$,

Applying a union bound finishes the proof. ∎

Given any fixed $\epsilon_{2}$ and $1\leq i\leq n$,

because $\left\langle w,x_{i}\right\rangle$ is a standard Gaussian r.v. and the density of standard Gaussian has maximum $1/\sqrt{2\pi}$. Since $\alpha_{i}(W_{0},\epsilon_{2})$ is the empirical mean of Bernoulli r.v.’s, by Hoeffding’s inequality, with probability $1-\nicefrac{{\delta}}{{n}}$,

Applying a union bound finishes the proof. ∎

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

and by further using the $1$-Lipschitz continuity of $\sigma$, we have

Given $1\leq i\leq n$, let $h_{i}=\sigma(W_{0}x_{i})/\sqrt{m}$. By Lemma A.1, $\|h_{i}\|_{2}$ is sub-Gaussian with variance proxy $1/m$, and with probability at least $1-\nicefrac{{\delta}}{{2n}}$ over $W_{0}$,

On the other hand, by Jensen’s inequality,

As a result, with probability $1-\nicefrac{{\delta}}{{2n}}$, it holds that $\|h_{i}\|_{2}\leq 1$. By a union bound, with probability $1-\nicefrac{{\delta}}{{2}}$ over $W_{0}$, for all $1\leq i\leq n$, we have $\|h_{i}\|_{2}\leq 1$.

For any $W_{0}$ such that the above event holds, and for any $1\leq i\leq n$, the r.v. $\left\langle h_{i},a\right\rangle$ is sub-Gaussian with variance proxy $\|h_{i}\|_{2}^{2}\leq 1$. By Hoeffding’s inequality, with probability $1-\nicefrac{{\delta}}{{2n}}$ over $a$,

By a union bound, with probability $1-\nicefrac{{\delta}}{{2}}$ over $a$, for all $1\leq i\leq n$, we have $\mathinner{\!\left\lvert f(x_{i};W_{0},a)\right\rvert}\leq\sqrt{2\ln\mathinner{\left(4n/\delta\right)}}$.

The probability that the above events all happen is at least $(1-\nicefrac{{\delta}}{{2}})(1-\nicefrac{{\delta}}{{2}})\geq 1-\delta$, over $W_{0}$ and $a$. ∎

The second-order term of eq. A.1 can be bounded as follows

because $\mathinner{\!\left\lVert\nabla\widehat{\mathcal{R}}(W_{t})\right\rVert}_{F}\leq\widehat{\mathcal{Q}}(W_{t})$, and $\eta_{t},\widehat{\mathcal{Q}}(W_{t})\leq 1$, and $\widehat{\mathcal{Q}}(W_{t})\leq\widehat{\mathcal{R}}(W_{t})$. Combining eqs. A.1, A.2 and A.3 gives

The required width ensures that with probability $1-3\delta$, Lemmas 2.3, 2.4 and 2.5 hold with $\epsilon_{1}=\gamma^{2}/8$ and $\epsilon_{2}=4\lambda/(\gamma\sqrt{m})$.

Let $t_{1}$ denote the first step such that there exists $1\leq s\leq m$ with $\mathinner{\!\left\lVert w_{s,t_{1}}-w_{s,0}\right\rVert}_{2}>4\lambda/(\gamma\sqrt{m})$. Therefore for any $0\leq t<t_{1}$ and any $1\leq s\leq m$, it holds that $\mathinner{\!\left\lVert w_{s,t}-w_{s,0}\right\rVert}_{2}\leq 4\lambda/(\gamma\sqrt{m})$. In addition, we let $\overline{W}\mathrel{\mathop{\ordinarycolon}}=W_{0}+\lambda\overline{U}$.

We will split the left hand side into three terms and control them individually:

The first term of eq. A.4 can be controlled using Lemma 2.5:

The second term of eq. A.4 can be written as

Let S_{c}\mathrel{\mathop{\ordinarycolon}}=\mathinner{\left\{s\ {}\middle|\ {}\mathds{1}\mathinner{\bigl{[}\left\langle w_{s,t},x_{i}\right\rangle>0\bigr{]}}-\mathds{1}\mathinner{\bigl{[}\left\langle w_{s,0},x_{i}\right\rangle>0\bigr{]}}\neq 0,1\leq s\leq m\right\}}. Note that $s\in S_{c}$ implies

where in the last step we use the condition that $m\geq 4096\lambda^{2}/\gamma^{6}$.

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

where we use $m\geq 4096\lambda^{2}/\gamma^{6}$. Therefore,

Putting eqs. A.5, A.6 and A.7 into eq. A.4, we have

for the $\lambda$ given in the statement of Theorem 2.2. Consequently, for any $0\leq t<t_{1}$, it holds that \widehat{\mathcal{R}}^{(t)}\mathinner{\bigl{(}\overline{W}\bigr{)}}\leq\epsilon/4.

Let $T\mathrel{\mathop{\ordinarycolon}}=\lceil\nicefrac{{2\lambda^{2}}}{{\eta\epsilon}}\rceil$. The next claim is that $t_{1}\geq T$. To see this, note that Lemma 2.6 ensures

Suppose $t_{1}<T$, then we have $t_{1}\leq\nicefrac{{2\lambda^{2}}}{{\eta\epsilon}}$, and thus $\mathinner{\!\left\lVert W_{t_{1}}-\overline{W}\right\rVert}_{F}^{2}\leq 2\lambda^{2}$. As a result, using $\|\overline{U}\|_{F}\leq 1$ and the definition of $\overline{W}$,

Furthermore, by the triangle inequality, for any $1\leq s\leq m$

which contradicts the definition of $t_{1}$. Therefore $t_{1}\geq T$.

Now we are ready to prove the claims of Theorem 2.2. The bound on $\mathinner{\!\left\lVert w_{s,t}-w_{s,0}\right\rVert}_{2}$ follow by repeating the steps in eq. A.8. The risk guarantee follows from Lemma 2.6:

Appendix B Omitted proofs from Section 3

The proof of Theorem 3.2 is based on Rademacher complexity. Given a sample $S=(z_{1},\ldots,z_{n})$ (where $z_{i}=(x_{i},y_{i})$) and a function class $\mathcal{H}$, the Rademacher complexity of $\mathcal{H}$ on $S$ is defined as

We will use the following general result.

(Shalev-Shwartz and Ben-David, 2014, Theorem 26.5) If $h(z)\in[a,b]$, then with probability $1-\delta$,

(Shalev-Shwartz and Ben-David, 2014, Lemma 26.9) $\textup{Rad}\mathinner{\left(g\circ\mathcal{F}\circ X\right)}\leq K\textup{Rad}\mathinner{\left(\mathcal{F}\circ X\right)}$.

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

Given a feature sample $X$, the following Lemma B.3 controls the Rademacher complexity of $\mathcal{F}_{\rho}\circ X$. A similar version was given in (Liang, 2016, Theorem 43), and the proof is similar to the proof of (Bartlett and Mendelson, 2002, Theorem 18) which also pushes the supremum through and handles each hidden unit separately.

$\textup{Rad}\mathinner{\left(\mathcal{F}_{\rho}\circ X\right)}\leq\rho\sqrt{m/n}$.

Note that for any $1\leq s\leq m$, the mapping $z\mapsto a_{s}\sigma(z)$ is $1$-Lipschitz, and thus Lemma B.2 gives

Invoking the Rademacher complexity of linear classifiers (Shalev-Shwartz and Ben-David, 2014, Lemma 26.10) then gives

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

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

As a result, invoking eq. B.1 with $\rho=4\lambda/(\gamma\sqrt{m})$, with probability $1-4\delta$ over the random initialization and data sampling,

Invoking $P_{(x,y)\sim\mathcal{D}}\mathinner{\left(yf(x;W,a)\leq 0\right)}\leq 2\mathcal{Q}(W)$ finishes the proof. ∎

Appendix C Omitted proofs from Section 4

Recall that $\mathinner{\!\left\lVert\nabla f_{t}(W_{t})\right\rVert}_{F}\leq 1$, we have

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

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

Under Assumption 3.1, given any $\epsilon\in(0,1)$, any $\delta\in(0,1/3)$, and any positive integer $n_{0}$, let

For any $m\geq M$ and any constant step size $\eta\leq 1$, if $n_{0}\geq n\mathrel{\mathop{\ordinarycolon}}=\lceil\nicefrac{{2\lambda^{2}}}{{\eta\epsilon}}\rceil$, then with probability $1-3\delta$,

We first sample $n_{0}$ data examples $(x_{0},y_{0}),\ldots,(x_{n_{0}-1},y_{n_{0}-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\leq s\leq m$ with $\mathinner{\!\left\lVert w_{s,n_{1}}-w_{s,0}\right\rVert}_{2}>4\lambda/(\gamma\sqrt{m})$. If such a step does not exist, let $n_{1}=n_{0}$.

Let $\overline{W}\mathrel{\mathop{\ordinarycolon}}=W_{0}+\lambda\overline{U}$, in exactly the same way as in Theorem 2.2, we can show that with probability $1-3\delta$, for any $0\leq i<n_{1}$,

Now consider $n\mathrel{\mathop{\ordinarycolon}}=\lceil\nicefrac{{2\lambda^{2}}}{{\eta\epsilon}}\rceil$. Using Lemma 4.2, in the same way as the proof of Theorem 2.2 (replacing $\widehat{\mathcal{Q}}(W_{\tau})$ with $\mathcal{Q}_{i}(W_{i})$, etc.), we can show that $n\leq n_{1}$. Then invoking Lemma 4.2 again, we get

Next we prove Lemma 4.3. We need the following martingale Bernstein bound.

(Beygelzimer et al., 2011, Theorem 1) Let $(M_{t},\mathcal{F}_{t})_{t\geq 0}$ denote a martingale with $M_{0}=0$ and $\mathcal{F}_{0}$ be the trivial $\sigma$-algebra. Let $(\Delta_{t})_{t\geq 1}$ denote the corresponding martingale difference sequence, and let

denote the sequence of conditional variance. If $\Delta_{t}\leq R$ a.s., then for any $\delta\in(0,1)$, with probability at least $1-\delta$,

For any $i\geq 0$, let $z_{i}$ denote $(x_{i},y_{i})$, and $z_{0,i}$ denote $(z_{0},\ldots,z_{i})$. Note that the quantity $\sum_{t<i}\mathinner{\left(\mathcal{Q}(W_{t})-\mathcal{Q}_{t}(W_{t})\right)}$ is a martingale w.r.t. the filtration $\sigma(z_{0,i-1})$. The martingale difference sequence is given by $\mathcal{Q}(W_{t})-\mathcal{Q}_{t}(W_{t})$, which satisfies

Invoking Lemma C.2 with eqs. C.4 and LABEL:eq:sgd_tmp2 gives that with probability $1-\delta$,

Suppose the condition of Lemma C.1 holds. Then we have for $n=\lceil\nicefrac{{2\lambda^{2}}}{{\eta\epsilon}}\rceil$, with probability $1-3\delta$,

Further invoking Lemma 4.3 gives that with probability $1-4\delta$,

Since $P_{(x,y)\sim\mathcal{D}}\mathinner{\left(yf(x;W,a)\leq 0\right)}\leq 2\mathcal{Q}(W)$, we get

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

Appendix D Omitted proofs from Section 5

with optimal primal-dual solutions $(\bar{w},\bar{q})$. Moreover

By strong duality, the inequality holds with equality. It follows that

Now let us look at the dual optimization problem. It is clear that

and thus $f^{*}(A^{*}\bar{q})=\gamma_{1}^{2}/2$. Since $\bar{w}=A^{*}\bar{q}$, we have that $\mathinner{\!\left\lVert\bar{w}\right\rVert}_{\mathcal{H}}=\gamma_{1}$. In addition,

and thus $-\bar{w}$ has margin $\gamma_{1}^{2}$. Moreover, we have

and thus $\mathinner{\!\left\lVert\bar{w}(z)\right\rVert}_{2}\leq 1$. Therefore, $\hat{v}=-\bar{w}/\gamma_{1}$ satisfies all requirements of Proposition 5.1. ∎

Let $\hat{q}$ denote the uniform probability vector $(\nicefrac{{1}}{{n}},\ldots,\nicefrac{{1}}{{n}})$. Note that

Since $0\leq\mathinner{\left(\hat{q}\odot\epsilon\right)}^{\top}K_{1}\mathinner{\left(\hat{q}\odot\epsilon\right)}\leq 1$ for any $\epsilon$, by Markov’s inequality with probability $0.9$, it holds that $\mathinner{\left(\hat{q}\odot\epsilon\right)}^{\top}K_{1}\mathinner{\left(\hat{q}\odot\epsilon\right)}\leq 1/(20n)$, and thus $\gamma_{1}\leq 1/\sqrt{20n}$. ∎

By symmetry, we only need to consider an $(x,y)$ where $(x_{1},x_{2},y)=(\nicefrac{{1}}{{\sqrt{d-1}}},0,1)$. Let $z_{p,q}$ denote $(z_{p},z_{p+1},\ldots,z_{q})$, and similarly define $x_{p,q}$. We have

For any nonzero $p\in A_{1}$, we have $-p\in A_{3}$, and $\left\langle\bar{v}(p),x_{1,2}\right\rangle=1/\sqrt{d-1}$. Therefore

Let $\varphi$ 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]$:

Since $\left\langle q,x_{3,d}\right\rangle$ is a Gaussian variable with standard deviation $\sqrt{\nicefrac{{(d-2)}}{{(d-1)}}}$, we have

Plugging eqs. D.4 and D.5 into eq. D.3 gives:

For $t\in[-1,+1]$, it holds that $\varphi(t)\geq 1\sqrt{2\pi e}$, and thus

To prove Proposition 5.4, we need the following technical lemma.

Given $z_{1}\sim\mathcal{N}(0,1)$ and $z_{2}\sim\mathcal{N}(0,b^{2})$ that are independent where $b>1$, we have

First note that for $z_{3}\sim\mathcal{N}(0,1)$ which is independent of $z_{1}$,

Still let $\varphi$ denote the density of $\mathcal{N}(0,1)$, and let $U(c)$ denote the probability that $z_{3}\in[-c,c]$. We have

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

By symmetry, we only need to consider the following training set:

The $1/\sqrt{d-1}$ factor is omitted also because we only discuss the $0/1$ loss.

For any $s$, let $A_{s}$ denote the event that

We will show that if $m\leq\sqrt{d-2}/4$, then $A_{s}$ is true for all $1\leq s\leq m$ with probability $1/2$, and Proposition 5.4 follows from the fact that the XOR data is not linearly separable.

Since $\mathinner{\left((x_{i})_{1},(x_{i})_{2}\right)}$ is $(1,0)$ or $(0,1)$ or $(-1,0)$ or $(0,-1)$, event $A_{s}$ will happen as long as

Note that $(w_{s})_{1},(w_{s})_{2}\sim\mathcal{N}(0,1)$ while $\sum_{j=3}^{d}(w_{s})_{j}\sim\mathcal{N}(0,d-2)$. As a result, due to Lemma D.1,