Learning and Generalization in Overparameterized Neural Networks, Going Beyond Two Layers

Introduction

Neural network learning has become a key machine learning approach and has achieved remarkable success in a wide range of real-world domains, such as computer vision, speech recognition, and game playing . In contrast to the widely accepted empirical success, much less theory is known. Despite a recent boost of theoretical studies, many questions remain largely open, including fundamental ones about the optimization and generalization in learning neural networks.

One key challenge in analyzing neural networks is that the corresponding optimization is non-convex and is theoretically hard in the general case . This is in sharp contrast to the fact that simple optimization algorithms like stochastic gradient descent (SGD) and its variants usually produce good solutions in practice even on both training and test data. Therefore,

what functions can neural networks provably learn?

Another key challenge is that, in practice, neural networks are heavily overparameterized (e.g., ): the number of learnable parameters is much larger than the number of the training samples. It is observed that overparameterization empirically improves both optimization and generalization, appearing to contradict traditional learning theory.For example, Livni et al. observed that on synthetic data generated from a target network, SGD converges faster when the learned network has more parameters than the target. Perhaps more interestingly, Arora et al. found that overparameterized networks learned in practice can often be compressed to simpler ones with much fewer parameters, without hurting their ability to generalize; however, directly learning such simpler networks runs into worse results due to the optimization difficulty. We also have experiments in Figure 1(a). Therefore,

why do overparameterized networks (found by those training algorithms) generalize?

Most existing works analyzing the learnability of neural networks make unrealistic assumptions about the data distribution (such as being random Gaussian), and/or make strong assumptions about the network (such as using linear activations). Other works focusing on structured data such as sparse coding model and Li and Liang shows that two-layer ReLU networks can learn classification tasks when the data come from mixtures of arbitrary but well-separated distributions.

A theorem without distributional assumptions on data is often more desirable. Indeed, how to obtain a result that does not depend on the data distribution, but only on the concept class itself, lies in the center of PAC-learning which is one of the foundations of machine learning theory . Also, studying non-linear activations is critical because otherwise one can only learn linear functions, which can also be easily learned via linear models without neural networks.

Brutzkus et al. prove that two-layer networks with ReLU activations can learn linearly-separable data (and thus the class of linear functions) using just SGD. This is an (improper) PAC-learning type of result because it makes no assumption on the data distribution. Andoni et al. proves that two-layer networks can learn polynomial functions of degree $r$ over $d$-dimensional inputs in sample complexity $O(d^{r})$. Their learner networks use exponential activation functions, where in practice the rectified linear unit (ReLU) activation has been the dominant choice across vastly different domains.

On a separate note, if one treats all but the last layer of neural networks as generating a random feature map, then training only the last layer is a convex task, so one can learn the class of linear functions in this implicit feature space . This result implies low-degree polynomials and compositional kernels can be learned by neural networks in polynomial time. Empirically, training last layer greatly weakens the power of neural networks (see Figure 1).

Our Result. We prove that an important concept class that contains three-layer (resp. two-layer) neural networks equipped with smooth activations can be efficiently learned by three-layer (resp. two-layer) ReLU neural networks via SGD or its variants.

Specifically, suppose in aforementioned class the best network (called the target function or target network) achieves a population risk $\mathsf{OPT}$ with respect to some convex loss function. We show that one can learn up to population risk $\mathsf{OPT}+\varepsilon$, using three-layer (resp. two-layer) ReLU networks of size greater than a fixed polynomial in the size of the target network, in $1/\varepsilon$, and in the “complexity” of the activation function used in the target network. Furthermore, the sample complexity is also polynomial in these parameters, and only poly-logarithmic in the size of the learner ReLU network.

We stress here that this is agnostic PAC-learning because we allow the target function to have error (e.g., $\mathsf{OPT}$ can be positive for regression), and is improper learning because the concept class consists of smaller neural networks comparing to the networks being trained.

Our Contributions. We believe our result gives further insights to the fundamental questions about the learning theory of neural networks.

To the best of our knowledge, this is the first result showing that using hidden layers of neural networks one can provably learn the concept class containing two (or even three) layer neural networks with non-trivial activation functions.In contrast, Daniely focuses on training essentially only the last layer (and the hidden-layer movement is negligible). After this paper has appeared online, Arora et al. showed that neural networks can provably learn two-layer networks with a slightly weaker class of smooth activation functions. Namely, the activation functions that are either linear functions or even functions.

Our three-layer result gives the first theoretical proof that learning neural networks, even with non-convex interactions across layers, can still be plausible. In contrast, in the two-layer case the optimization landscape with overparameterization is almost convex ; and in previous studies on the multi-layer case, researchers have weakened the network by applying the so-called NTK (neural tangent kernel) linearization to remove all non-convex interactions .

In fact, the techniques introduced in this work is later used in to show how multi-layer neural networks can provably learn concept classes where no kernel methods can learn efficiently, and used in to show how the initial large learning rate in training a neural network can improve generalization, also not achievable by the convex kernel methods.

To some extent we explain the reason why overparameterization improves testing accuracy: with larger overparameterization, one can hope to learn better target functions with possibly larger size, more complex activations, smaller risk $\mathsf{OPT}$, and to a smaller error $\varepsilon$.

We establish new tools to tackle the learning process of neural networks in general, which can be useful for studying other network architectures and learning tasks. (E.g., the new tools here have allowed researchers to study also the learning of recurrent neural networks .)

Other Related Works. We acknowledge a different line of research using kernels as improper learners to learn the concept class of neural networks . This is very different from us because we use “neural networks” as learners. In other words, we study the question of “what can neural networks learn” but they study “what alternative methods can replace neural networks.”

There is also a line of work studying the relationship between neural networks and NTKs (neural tangent kernels) . These works study neural networks by considering their “linearized approximations.” There is a known performance gap between the power of real neural networks and the power of their linearized approximations. For instance, ResNet achieves 96% test error on the CIFAR-10 data set but NTK (even with infinite width) achieves 77% . We also illustrate this in Figure 1.

With sufficient overparameterization, for two-layer neural networks Li and Liang and Du et al. show that gradient descent can perfectly fit the training samples when the data is not degenerated. Allen-Zhu et al. show that deep neural networks can also provably fit the training samples in polynomial time. These results do not cover generalization so are not relevant in studying the learnability of neural networks.

2 Why Do Overparameterized Networks Generalize?

Our result above assumes that the learner network is sufficiently overparameterized. So, why does it generalize to the population risk and give small test error? More importantly, why does it generalize with a number of samples that is (almost) independent of the number of parameters?

This question cannot be studied under the traditional VC-dimension learning theory since the VC dimension grows with the number of parameters. Several works explain generalization by studying some other “complexity” of the learned networks. Most related to the discussion here is where the authors prove a generalization bound in the norms (of weight matrices) of each layer, as opposed to the number of parameters. There are two main concerns with those results.

Learnability $=$ Trainability $+$ Generalization. It is not clear from those results how a network with both low “complexity” and small training loss can be found by the training method. Therefore, they do not directly imply PAC-learnability for non-trivial concept classes (at least for those concept classes studied by this paper).

Their norms are “sparsity induced norms”: for the norm not to scale with the number of hidden neurons $m$, essentially, it requires the number of neurons with non-zero weights not to scale with $m$. This more or less reduces the problem to the non-overparameterized case.

At a high level, our generalization is made possible with the following sequence of conceptual steps.

Good networks with small risks are plentiful: thanks to overparameterization, with high probability over random initialization, there exists a good network in the close neighborhood of any point on the SGD training trajectory. (This corresponds to Section 6.2 and 6.3.)

The optimization in overparameterized neural networks has benign properties: essentially along the training trajectory, there is no second-order critical points for learning three-layer networks, and no first-order critical points for two-layer. (This corresponds to Section 6.4.)

In the learned networks, information is also evenly distributed among neurons, by utilizing either implicit or explicit regularization. This structure allows a new generalization bound that is (almost) independent of the number of neurons. (This corresponds to Section 6.5 and 6.6, and we also empirically verify it in Section 7.1.)

Since practical neural networks are typically overparameterized, we genuinely hope that our results can provide theoretical insights to networks used in various applications.

3 Roadmap

In the main body of this paper, we introduce notations in Section 2, present our main results and contributions for two and three-layer networks in Section 3 and 4, and conclude in Section 5.

For readers interested in our novel techniques, we present in Section 6 an 8-paged proof sketch of our three-layer result. For readers more interested in the practical relevance, we give more experiments in Section 7. In the appendix, we begin with mathematical preliminaries in Appendix A. Our full three-layer proof is in Appendix C. Our two-layer proof is much easier and in Appendix B.

Notations

For notation simplicity, with high probability (or w.h.p.) means with probability $1-e^{-c\log^{2}m}$ for a sufficiently large constant $c$ for two-layer network where $m$ is the number of hidden neurons, and $1-e^{-c\log^{2}(m_{1}m_{2})}$ for three-layer network where $m_{1},m_{2}$ are the number of hidden neurons for the first and second layer (to be precisely defined later). In this paper, $\widetilde{O}$ hides factors of $\polylog(m)$ for two-layer networks, or $\polylog(m_{1},m_{2})$ for three-layer networks.

Function complexity. The following notion measures the complexity of any smooth activation function $\phi(z)$. Suppose $\phi(z)=\sum_{i=0}^{\infty}c_{i}z^{i}$. Given a non-negative $R$, the complexity

where $C^{*}$ is a sufficiently large constant (e.g., $10^{4}$). Intuitively, $\mathfrak{C}_{\mathfrak{s}}$ measures the sample complexity: how many samples are required to learn $\phi$ correctly; while $\mathfrak{C}_{\varepsilon}$ bounds the network size: how much over-parameterization is needed for the algorithm to efficiently learn $\phi$ up to $\varepsilon$ error. It is always true that $\mathfrak{C}_{\mathfrak{s}}(\phi,R)\leq\mathfrak{C}_{\varepsilon}(\phi,R)\leq\mathfrak{C}_{\mathfrak{s}}(\phi,O(R))\times\poly(1/\varepsilon)$.Recall \big{(}\frac{\sqrt{\log(1/\varepsilon)}}{\sqrt{i}}C^{*}\big{)}^{i}\leq e^{O(\log(1/\varepsilon))}=\frac{1}{\poly(\varepsilon)} for every $i\geq 1$. While for $\sin z,\exp(z)$ or low degree polynomials, $\mathfrak{C}_{\mathfrak{s}}(\phi,O(R))$ and $\mathfrak{C}_{\varepsilon}(\phi,R)$ only differ by $o(1/\varepsilon)$.

Result for Two-Layer Networks

the complexity of $F^{*}$ and assume they are bounded.

Standard two-layer networks $f_{r}^{*}(x)=\sum_{i=1}^{p}a^{*}_{r,i}\phi(\langle w^{*}_{1,i},x\rangle)$ are special cases of (3.1) (by setting $w_{2,i}^{*}=(0,\dots,0,1)$ and $\phi_{i}=\phi$). Our formulation (3.1) additionally captures combinations of correlations between non-linear and linear measurements of different directions of $x$.

Learning Process. Let $W^{(0)}$ denote the initial value of the hidden weight matrix, and let $W^{(0)}+W_{t}$ denote the value at time $t$. (Note that $W_{t}$ is the matrix of increments.) The weights are initialized with Gaussians and then $W$ is updated by the vanilla SGD. More precisely,

entries of $W^{(0)}$ and $b^{(0)}$ are i.i.d. random Gaussians from $\mathcal{N}(0,1/m)$,

entries of each $a^{(0)}_{r}$ are i.i.d. random Gaussians from $\mathcal{N}(0,\varepsilon_{a}^{2})$ for some fixed $\varepsilon_{a}\in(0,1]$.We shall choose $\varepsilon_{a}=\widetilde{\Theta}(\varepsilon)$ in the proof due to technical reason. As we shall see in the three-layer case, if weight decay is used, one can relax this to $\varepsilon_{a}=1$.

For notation simplicity, with high probability (or w.h.p.) means with probability $1-e^{-c\log^{2}m}$ for a sufficiently large constant $c$, and $\widetilde{O}$ hides factors of $\polylog(m)$.

For every \varepsilon\in\big{(}0,\frac{1}{pk\mathfrak{C}_{\mathfrak{s}}(\phi,1)}\big{)}, there exists

such that for every $m\geq M_{0}$ and every $N\geq\widetilde{\Omega}(N_{0})$, choosing $\varepsilon_{a}=\varepsilon/\widetilde{\Theta}(1)$ for the initialization, choosing learning rate \eta=\widetilde{\Theta}\big{(}\frac{1}{\varepsilon km}\big{)} and

with high probability over the random initialization, SGD after $T$ iteration satisfies

SGD only takes one example per iteration. Thus, sample complexity is also at most $T$.

We note sample complexity $T$ is (almost) independent of $m$, the amount of overparametrization.

2 Our Interpretations

Overparameterization improves generalization. By increasing $m$, Theorem 1 supports more target functions with possibly larger size, more complex activations, and smaller population risk $\mathsf{OPT}$. In other words, when $m$ is fixed, among the class of target functions whose complexities are captured by $m$, SGD can learn the best function approximator of the data, with the smallest population risk. This gives intuition how overparameterization improves test error, see Figure 1(a).

Large margin non-linear classifier. Theorem 1 is a nonlinear analogue of the margin theory for linear classifiers. The target function with a small population risk (and of bounded norm) can be viewed as a “large margin non-linear classifier.” In this view, Theorem 1 shows that assuming the existence of such large-margin classifier, SGD finds a good solution with sample complexity mostly determined by the margin, instead of the dimension of the data.

Inductive bias. Recent works (e.g., ) show that when the network is heavily overparameterized (that is, $m$ is polynomial in the number of training samples) and no two training samples are identical, then SGD can find a global optimum with classification error (or find a solution with $\varepsilon$ training loss) in polynomial time. This does not come with generalization, since it can even fit random labels. Our theorem, combined with , confirms the inductive bias of SGD for two-layer networks: when the labels are random, SGD finds a network that memorizes the training data; when the labels are (even only approximately) realizable by some target network, then SGD learns and generalizes. This gives an explanation towards the well-known empirical observations of such inductive bias (e.g., ) in the two-layer setting, and is more general than Brutzkus et al. in which the target network is only linear.

Result for Three-Layer Networks

Concept class and target function $F^{*}(x)$. This time we consider more powerful target functions $F^{*}=(f_{1}^{*},\cdots,f_{k}^{*})$ of the form

to denote the complexity of the two layers, and assume they are bounded.

Our concept class contains measures of correlations between composite non-linear functions and non-linear functions of the input, there are plenty of functions in this new concept class that may not necessarily have small-complexity representation in the previous formulation (3.1), and as we shall see in Figure 1(a), this is the critical advantage of using three-layer networks compared to two-layer ones or their NTKs. The learnability of this correlation is due to the non-convex interactions between hidden layers. As a comparison, studies the regime where the changes in hidden layers are negligible thus can not show how to learn this concept class with a three-layer network.

are only special cases of (4.1). Also, even in the special case of $\Phi_{i}(z)=z$, the target

captures combinations of correlations of combinations of non-linear measurements in different directions of $x$. This we do not know how to compute using two-layer networks.

In fact, our results of this paper even apply to the following general form:

with the mild requirement that $\langle w^{*}_{1,j},w^{*}_{3,j}\rangle=0$ for each $j\in[p_{2}]$. We choose to present the slightly weaker formulation (4.1) for cleaner proofs.

Learner network $F(x;W,V)$. Our learners are three-layer networks $F=(f_{1},\ldots,f_{k})$ with

Again for simplicity, we only update $W$ and $V$. The weights are randomly initialized as:

entries of $W^{(0)}$ and $b_{1}=b^{(0)}_{1}$ are i.i.d. from $\mathcal{N}(0,1/m_{1})$,

entries of $V^{(0)}$ and $b_{2}=b^{(0)}_{2}$ are i.i.d. from $\mathcal{N}(0,1/m_{2})$,

entries of each $a_{r}=a^{(0)}_{r}$ are i.i.d. from $\mathcal{N}(0,\varepsilon_{a}^{2})$ for $\varepsilon_{a}=1$. Recall in our two-layer result we have chosen $\varepsilon_{a}<1$ due to technical reasons; thanks to weight decay, we can simply select $\varepsilon_{a}=1$ in our three-layer case.

As for the optimization algorithm, we use SGD with weight decay and an explicit regularizer.

For some $\lambda\in(0,1]$, we will use $\lambda F(x;W,V)$ as the learner network, i.e., linearly scale $F$ down by $\lambda$. This is equivalent to replacing $W$, $V$ with $\sqrt{\lambda}W$, $\sqrt{\lambda}V$, since a ReLU network is positive homogenous. The SGD will start with $\lambda=1$ and slowly decrease it, similar to weight decay.We illustrate the technical necessity of adding weight decay. During training, it is easy to add new information to the current network, but hard to forget “false” information that is already in the network. Such false information can be accumulated from randomness of SGD, non-convex landscapes, and so on. Thus, by scaling down the network we can effectively forget false information.

We also use an explicit regularizer for some $\lambda_{w},\lambda_{v}>0$ with This $\|\cdot\|_{2,4}$ norm on $W$ encourages weights to be more evenly distributed across neurons. It can be replaced with $\|\sqrt{\lambda_{t-1}}W_{t-1}\|_{2,2+\alpha}^{2+\alpha}$ for any constant $\alpha>0$ for our theoretical purpose. We choose $\alpha=2$ for simplicity, and observe that in practice, weights are automatically spread out due to data randomness, so this explicit regularization may not be needed. See Section 7.1 for an experiment.

Now, in each round $t=1,2,\dots,T$, we use (noisy) SGD to minimize the following stochastic objective for some fixed $\lambda_{t-1}$:

2 Main Theorems

For notation simplicity, with high probability (or w.h.p.) means with probability $1-e^{-c\log^{2}(m_{1}m_{2})}$ and $\widetilde{O}$ hides factors of $\polylog(m_{1},m_{2})$.

Consider Algorithm 2. For every constant $\gamma\in(0,1/4]$, every $\varepsilon_{0}\in(0,1/100]$, every $\varepsilon=\frac{\varepsilon_{0}}{kp_{1}p_{2}^{2}\mathfrak{C}_{\mathfrak{s}}(\Phi,p_{2}\mathfrak{C}_{\mathfrak{s}}(\phi,1))\mathfrak{C}_{\mathfrak{s}}(\phi,1)^{2}}$, there exists

such that for every $m_{2}=m_{1}=m\geq M$, and properly set $\lambda_{w},\lambda_{v},\sigma_{w},\sigma_{v}$ in Table 1, as long as

there is a choice $\eta=1/\poly(m_{1},m_{2})$ and $T=\poly(m_{1},m_{2})$ such that with probability $\geq 99/100$,

For general $\Phi$, ignoring $p_{2}$, the complexity of three-layer networks is essentially $\mathfrak{C}_{\varepsilon}(\Phi,\mathfrak{C}_{\varepsilon}(\phi,1))$. This is necessary in some sense: consider the case when $\phi(x)=Cx$ for a very large parameter $C$, then $\Phi(\phi(z))$ is just a function $\Phi^{\prime}(z)=\Phi(Cz)$, and we have $\mathfrak{C}_{\varepsilon}(\Phi^{\prime},1)=\mathfrak{C}_{\varepsilon}(\Phi,C)\approx\mathfrak{C}_{\varepsilon}(\Phi,\Theta(\mathfrak{C}_{\varepsilon}(\phi,1)))$.

3 Our Contributions

Our sample complexity $N$ scales polynomially with the complexity of the target network, and is (almost) independent of $m$, the amount of overparameterization. This itself can be quite surprising, because recent results on neural network generalization require $N$ to be polynomial in $m$. Furthermore, Theorem 2 shows three-layer networks can efficiently learn a bigger concept class (4.1) comparing to what we know about two-layer networks (3.1).

From a practical standpoint, one can construct target functions of the form (4.1) that cannot be (efficiently) approximated by any two-layer target function in (3.1). If data is generated according to such functions, then it may be necessary to use three-layer networks as learners (see Figure 1).

From a theoretical standpoint, even in the special case of $\Phi(z)=z$, our target function can capture correlations between non-linear measurements of the data (recall Remark 4.1). This means $\mathfrak{C}_{\varepsilon}(\Phi,\mathfrak{C}_{\varepsilon}(\phi,1)\sqrt{p_{2}})\approx O(\sqrt{p_{2}}\mathfrak{C}_{\varepsilon}(\phi,1))$, so learning it is essentially in the same complexity as learning each $\phi_{s,j}$. For example, a three-layer network can learn $\cos(100\langle w^{*}_{1},x\rangle)\cdot e^{100\langle w^{*}_{2},x\rangle}$ up to accuracy $\varepsilon$ in complexity $\poly(1/\varepsilon)$, while it is unclear how to do so using two-layer networks.

Technical Contributions. We highlight some technical contributions in the proof of Theorem 2.

In recent results on the training convergence of neural networks for more than two layers , the optimization process stays in a close neighborhood of the initialization so that, with heavy overparameterization, the network becomes “linearized” and the interactions across layers are negligible. In our three-layer case, this means that the matrix $W$ never interacts with $V$. They then argue that SGD simulates a neural tangent kernel (NTK) so the learning process is almost convex .

In our analysis, we directly tackle non-convex interactions between $W$ and $V$, by studying a “quadratic approximation” of the network. See Remark 6.1 for a mathematical comparison. In other words, we introduced a second-order version of NTK, and our new proof techniques could be useful for future theoretical applications.

Comparison to Daniely . Daniely studies the learnability of multi-layer networks when (essentially) only the output layer is trained, which reduces to a convex task. He shows that multi-layer networks can learn a compositional kernel space, which implies two/three-layer networks can efficiently learn low-degree polynomials. He did not derive the general sample/time complexity bounds for more complex functions such as those in our concept classes (3.1) and (4.1), but showed that they are finite.

In contrast, our learnability result of concept class (4.1) is due to the non-convex interaction between hidden layers. Since Daniely studies the regime when the changes in hidden layers are negligible, if three layer networks are used, to the best of our knowledge, their theorem cannot lead to similar sample complexity bounds comparing to Theorem 2 by only training the last layer of a three-layer network. Empirically, one can also observe that training hidden layers is better than training the last layer (see Figure 1).

Conclusion and Discussion

We show by training the hidden layers of two-layer (resp. three-layer) overparameterized neural networks, one can efficiently learn some important concept classes including two-layer (resp. three-layer) networks equipped with smooth activation functions. Our result is in the agnostic PAC-learning language thus is distribution-free. We believe our work opens up a new direction in both algorithmic and generalization perspectives of overparameterized neural networks, and pushing forward can possibly lead to more understanding about deep learning.

One can also combine this paper with properties of recurrent neural networks (RNNs) to derive PAC-learning results for RNNs , or use the existential tools of this paper to derive PAC-learning results for three-layer residual networks (ResNet) . The latter gives a provable separation between neural networks and kernels in the efficient PAC-learning regime.

There are plenty of open directions following our work, especially how to extend our result to larger number of layers. Even for three layers, our algorithm currently uses an explicit regularizer $R(W^{\prime},V^{\prime})$ on the weights to control the generalization. In contrast, in practice, it is observed that neural networks actually implicitly regularizes: even without any restriction to weights, the network learned by an unconstraint overparameterized neural network still generalizes. It is an interesting direction to explain this inductive bias for three layers and beyond. (This present paper only explains inductive bias for two-layer networks.) As for a third direction, currently our result does not directly apply to target networks with ReLU activations. While there are functions $\phi$ with $\mathfrak{C}_{\varepsilon}(\phi,1)=2^{O(1/\varepsilon)}$ that approximate ReLU function in $$with$\varepsilon$ error, obtaining an polynomial or sub-exponential bound on the sampling complexity would be of great interest. On the other hand, the result by indicates that learning ReLU could also be hard for algorithms such as SGD.

We present key technical lemmas we used for proving the three-layer network theorems in Section 6 so that interested readers do not need to go through the appendix. Many of them can be of independent interests and have found further applications beyond this paper (such as for residual networks and for recurrent networks ). The full proof is in Appendix C.

The two-layer result is based on similar ideas but simpler, and the full proof is in Appendix B.

We give more experiments in Section 7 on Page 7. Our appendix starts on page A.

Main Lemmas for Three Layer Networks

In Section 6.1, we state the main theorem for the first variant of the SGD, which we excluded from the main body due to space limitation. In Section 6.2, we show the existence of some good “pseudo network” that can approximate the target. In Section 6.3, we present our coupling technique between a real network and a pseudo network. In Section 6.4, we present the key lemma about the optimization procedure. In Section 6.5, we state a simple generalization bound that is compatible with our algorithm. These techniques together give rise to the proof of Theorem 3. In Section 6.6, we present additional techniques needed to show Theorem 2.

In the first variant of SGD, in each round $t=1,2,\dots,T$, we use (noisy) SGD to minimize the following stochastic objective for some fixed $\lambda_{t-1}$:

Below is the main theorem for using the first variant of SGD to train three-layer networks.

Consider Algorithm 3. In the same setting as Theorem 2, for every $m_{2}=m_{1}=m\geq M$, as long as

there is choice $\eta=1/\poly(m_{1},m_{2})$, $T=\poly(m_{1},m_{2})$ such that with probability at least $99/100$,

As $m$ goes large, this sample complexity $N\approx m^{3/2}$ polynomially scales with $m$ so may not be very efficient (we did not try hard to improve the exponent $3/2$). Perhaps interesting enough, this is already non-trivial, because $N$ can be much smaller than $m^{2}$, the number of parameters of the network or equivalently the naive VC dimension bound. Recall in our second variant of SGD, the sample complexity $N$ only grows polylogarithmically in $m$.

2 Existence

$D_{v,x}\in\{0,1\}^{m_{2}\times m_{2}}$ denote the diagonal sign matrix of the second layer at random initialization.

Consider the output of a three-layer network at randomly initialized sign without bias as

The above pseudo network can be reminiscent of the simpler linearized NTK approximation of a network used in prior work , which in our language means

In such linearization it is clear that the weight matrices $W$ and $V$ do not interact with each other. In contrast, in our quadratic formula $a_{r}D_{v,x}VD_{w,x}Wx$, the matrices $W$ and $V$ are multiplied together, resulting in a non-convex interaction after putting into the loss function. This can be viewed as a second-order version of NTK (but is itself not NTK). We shall see in Section 6.3 that the pseudo network can be made close to the real network in some sense.

For every \varepsilon\in\big{(}0,\frac{1}{kp_{1}p_{2}^{2}\mathfrak{C}_{\mathfrak{s}}(\Phi,p_{2}\mathfrak{C}_{\mathfrak{s}}(\phi,1))\mathfrak{C}_{\mathfrak{s}}(\phi,1)^{2}}\big{)}, there exists

such that if $m_{1},m_{2}\geq M$, then with high probability, there exists weights $W^{\divideontimes},V^{\divideontimes}$ with

In other words, at randomly initialized signs, there exist choices of $W^{\divideontimes}$ and $V^{\divideontimes}$ with small norms so that $G^{(0)}(x;W^{\divideontimes},V^{\divideontimes})$ approximates the target. Later, we will combine this with the coupling lemma Section 6.3 to show a main structural property of overparameterized networks: solutions with good population risks are dense in the parameter space, in the sense that with high probability over the random initialization, there exists a good solution in the “close ” neighborhood of the initialized weights.

where $\alpha_{1},\beta_{1},b_{0}\sim\mathcal{N}(0,1)$ are independent random variables.

The reason Lemma lem:fit_fun_mainb is equivalent to Lemma lem:fit_fun_maina consists of a few thinking steps. Without loss of generality, one can assume $w^{*}=(1,0,0,\dots,0)$ and write $\phi(\langle w^{*},x\rangle)=\phi(x_{1})$ and write $h(\langle w,w^{*}\rangle,b_{0})=h(w_{1},b_{0})$ so it only depends on the first coordinate of $w$ and $b_{0}$. Under such simplifications, the second through last coordinates do not make any difference, so we can assume without loss of generality that $w^{*},w,x$ are only in 2 dimensions, and write $w=(\alpha_{1},\beta_{1})$ and $x=(x_{1},\sqrt{1-x_{1}^{2}})$. In sum, proving Lemma lem:fit_fun_maina suffices in establishing Lemma lem:fit_fun_mainb.

Given Lemma 6.3, we can directly apply it to the two-layer case (see Appendix B.1) to show the existence of good pseudo networks. As for the three-layer case, we need to apply Lemma 6.3 twice: once for (each neuron of) the second hidden layer and once for the output.

Consider the the input (without bias) to a single neuron of the second hidden layer at random initialization. Without loss of generality, say the first neuron, given as:

Even though $n_{1}(x)$ is completely random, using Lemma 6.3, we can derive the following lemma which rewrites $n_{1}(x)$ in the direction of an arbitrary function $\phi$.

Moreover, letting $C=\mathfrak{C}_{\varepsilon}(\phi,1)$ be the complexity of $\phi$, and if $v_{1,i}^{(0)}\sim\mathcal{N}(0,\frac{1}{m_{2}})$ and $w_{i,j}^{(0)},b^{(0)}_{1,i}\sim\mathcal{N}(0,\frac{1}{m_{1}})$ are at random initialization, then we have

For every fixed $x$, $\rho\left(v_{1}^{(0)},W^{(0)},b_{1}^{(0)}\right)$ is independent of $B\left(x,v_{1}^{(0)},W^{(0)},b_{1}^{(0)}\right)$.

$\rho\left(v_{1}^{(0)},W^{(0)},b_{1}^{(0)}\right)\sim\mathcal{N}\left(0,\frac{1}{100C^{2}m_{2}}\right)$.

For every $x$ with $\|x\|_{2}=1$, $|\phi_{\varepsilon}(x)-\phi(\langle w^{*},x\rangle)|\leq\varepsilon$.

For every fixed $x$ with $\|x\|_{2}=1$, with high probability $\left|R\left(x,v_{1}^{(0)},W^{(0)},b_{1}^{(0)}\right)\right|\leq\widetilde{O}\left(\frac{1}{{\sqrt{m_{1}m_{2}}}}\right)$ and $\left|B\left(x,v_{1}^{(0)},W^{(0)},b_{1}^{(0)}\right)\right|\leq\widetilde{O}\left(\frac{1}{{\sqrt{m_{2}}}}\right)$.

Furthermore, there exists real-valued function $\widetilde{\rho}(v_{1}^{(0)})$ satisfying with high probability:

Lemma 6.4 shows that, up to some small noise $B$ and $R$, we can “view” the input to any neuron of the second layer essentially as “a Gaussian variable $\rho$” times “the target function $\phi(\langle w^{*},x\rangle)$” in the first layer of the target network. This allows us to apply Lemma 6.3 again for the output layer, so as to construct for instance a composite function $\Phi(\cdot)$ with $\phi_{i}$ as its inputs.One technical issue is the following. When considering multiple neurons of the second layer, those Gaussian variables $\rho$ are not independent because they all depend on $W^{(0)}$. The notion of $\widetilde{\rho}$ in Lemma 6.4 removes such dependency across multiple neurons, at the expense of small Wasserstein distance.

Lemma 6.4 may sound weird at first look because random initialization cannot carry any information about the target. There is no contradiction here, because we will show, $B$ is essentially another Gaussian (with the same distribution as $\rho$) times $\sqrt{\text{constant}-\phi^{2}(\langle w^{*},x\rangle)}$, thus $n_{1}(x)$ can still be independent of the value of $\phi(\langle w^{*},x\rangle)$. Nevertheless, the decomposition in Lemma 6.4 shall enable us to show that, when we start to modify the hidden weights $W,V$, the learning process will start to discover this structure and make the weight of the term relating to $\phi(\langle w^{*},x\rangle)$ stand out from other terms.

Using Lemma 6.4 and applying Lemma 6.3 once more, we can prove Lemma 6.2.

3 Coupling Between Real and Pseudo Networks

Suppose we are currently at weights $W^{(0)}+W^{\prime}+W^{\rho}$, $V^{(0)}+V^{\prime}+V^{\rho}$, where matrices $W^{\rho},V^{\rho}$ are random Gaussian matrices such that:

for some $\sigma_{v},\sigma_{w}\in[1/(m_{1}m_{2}),1]$ to be specified later, and $W^{\prime},V^{\prime}$ are matrices with bounded norms that can depend on the randomness of $W^{(0)},b_{1}^{(0)},V^{(0)},b_{2}^{(0)}$. Intuitively, $W^{\prime}$ and $V^{\prime}$ capture how much the algorithm has moved away from the initialization, while $W^{\rho},V^{\rho}$ are introduced for adding smoothness in the optimization, see Section 6.4.

Let us introduce the notion of pseudo networks at the current weights. Let

$D_{w,x}$ denote the diagonal sign matrix of the first layer at random initialization $W^{(0)}$,

$D_{v,x}$ denote the diagonal sign matrix of the second layer at random initialization $W^{(0)},V^{(0)}$,

$D_{v,x}+D_{v,x}^{\prime}\in\{0,1\}^{m_{2}\times m_{2}}$ denote the diagonal sign matrix of the second layer at these weights.

For a fixed $r\in[k]$, let us denote row vector $a_{r}=(a_{r,i})_{i\in[m_{2}]}$. Define the pseudo network (and its semi-bias, bias-free versions) as

As a sanity check, at $W^{(0)}+W^{\prime}+W^{\rho}$, $V^{(0)}+V^{\prime}+V^{\rho}$ the pseudo network equals the true one:

Suppose \tau_{v}\in\big{[}0,1\big{]}, \tau_{w}\in\big{(}\frac{1}{m_{1}^{3/2}},\frac{1}{m_{1}^{1/2}}\big{]}, \sigma_{w}\in\big{(}\frac{1}{m_{1}^{3/2}},\frac{\tau_{w}}{m_{1}^{1/4}}\big{]}, \sigma_{v}\in\big{(}0,\frac{1}{m_{2}^{1/2}}\big{]} and $\eta>0$. Given fixed unit vector $x$, and perturbation matrices $W^{\prime},V^{\prime},W^{\prime\prime},V^{\prime\prime}$ (that may depend on the randomness of $W^{(0)},b_{1}^{(0)},V^{(0)},b_{2}^{(0)}$ and $x$) satisfying

(Sparse sign change). $\|D_{w,x}^{\prime}\|_{0}\leq\widetilde{O}(\tau_{w}^{4/5}m_{1}^{6/5})$, $\|D_{v,x}^{\prime}\|_{0}\leq\widetilde{O}\left(\sigma_{v}m_{2}^{3/2}+\tau_{v}^{2/3}m_{2}+\tau_{w}^{2/3}m_{1}^{1/6}m_{2}\right)$.

The first statement “sparse sign change” of Lemma 6.5 says that, if we move from random initialization $W^{(0)},V^{(0)}$ to $W^{(0)}+W^{\prime}+W^{\rho},V^{(0)}+V^{\prime}+V^{\rho}$, then how many signs of the ReLUs (in each layer) will change, as a function of the norms of $W^{\prime}$ and $V^{\prime}$. This calculation is similar to but slightly more involved due to the $\|\cdot\|_{2,4}$ norm that we use here.

We quickly point out a corollary by applying the coupling and existential lemmas together.

Recall in the existential Lemma 6.2, we have studied a pseudo network $G^{(0)}$ where the signs are determined at the random initialization $W^{(0)},V^{(0)}$. Now, the coupling Lemma 6.5 says that the amount of sign change from $G^{(0))}$ to $G^{(b,b)}$ can be controlled. Therefore, if parameters are chosen appropriately, the existential Lemma 6.2 should also apply to $G^{(b,b)}$. Formally,

In the same setting as Lemma 6.2, given perturbation matrices $W^{\prime},V^{\prime}$ (that may depend on the randomness of the initialization and the data distribution $\mathcal{D}$) with

Using parameter choices from Table 1, w.h.p. there exist $W^{\divideontimes}$ and $V^{\divideontimes}$ (independent of the randomness of $W^{\rho},V^{\rho}$) satisfying

As we shall see later, Corollary 6.6 gives rise to $W^{\divideontimes}$ and $V^{\divideontimes}$ that shall be used as a descent direction for the objective. (Corollary 6.6 did not use all the parameters inside Table 1, and some of those parameters shall be used in later subsections.)

4 Optimization

use the property that solutions with good risks are dense in the parameter space (i.e., Corollary 6.6) to show that the optimization landscape of the overparameterized three-layer neural network is benign: it has no spurious local minimal or (more or less) even any second-order critical points; and

use existing theorems on escaping saddle points (such as ) to show that SGD will not be stuck in saddle points and thus converges.

Key issue. Unfortunately, before digging into the details, there is already a big hole in this approach. ReLU networks are not second-order-differentiable: a ReLU activation does not have a well-defined Hessian/sub-Hessian at zero. One may na ïvelythink that since a ReLU network is infinite-order differentiable everywhere except a measure zero set, so we can safely ignore the Hessian issue and proceed by pretending that the Hessian of ReLU is always zero. This intuition is very wrong. Following it, we could have run into the absurd conclusion that any piece-wise linear function is convex, since the Hessian of it is zero almost everywhere. In other words, the only non-smooth point of ReLU has a Hessian value equal to the Dirac $\delta$-function, but such non-smooth points, albeit being measure zero, are actually turning points in the landscape. If we want a meaningful second-order statement of the ReLU network, we must not na ïvelyignore the “Hessian” of ReLU at zeros.

In practice, since the solution $W_{t},V_{t}$ are found by stochastic gradient descent starting from random initialization, they will have a non-negligible amount of intrinsic noise. Thus, the additional smoothing in the algorithm might not be needed by an observation in . Smooth analysis might also be used for analyzing the effect of such noise, but this is beyond the scope of this paper.

Actual algorithm. Let us consider the following smoothed, and regularized objective:

where $W^{\rho},V^{\rho}$ are Gaussian random matrices with each entry i.i.d. from $\mathcal{N}(0,\sigma_{w}^{2})$ and $\mathcal{N}(0,\sigma_{v}^{2})$, respectively. $R(\sqrt{\lambda_{t}}W_{t},\sqrt{\lambda_{t}}V_{t})=\lambda_{v}\|\sqrt{\lambda_{t}}V_{t}\|_{F}^{2}+\lambda_{w}\|\sqrt{\lambda_{t}}W_{t}\|_{2,4}^{4}$ and $\lambda_{v},\lambda_{w}$ are set such that $\lambda_{v}\|\sqrt{\lambda_{t}}V^{\divideontimes}\|_{F}^{2}\leq\varepsilon_{0}$ and $\lambda_{w}\|\sqrt{\lambda_{t}}W^{\divideontimes}\|_{2,4}^{4}\leq\varepsilon_{0}$ for every $W^{\divideontimes}$ and $V^{\divideontimes}$ coming from Lemma 6.2. See Algorithm 3 (first SGD variant) for the details. We prove the following lemma.

For every $\varepsilon_{0}\in(0,1)$ and $\varepsilon=\frac{\varepsilon_{0}}{kp_{1}p_{2}^{2}\mathfrak{C}_{\mathfrak{s}}(\Phi,p_{2}\mathfrak{C}_{\mathfrak{s}}(\phi,1))\mathfrak{C}_{\mathfrak{s}}(\phi,1)^{2}}$, for every constant $\gamma\in(0,1/4]$, consider the parameter choices in Table 1, and consider any $\lambda_{t},W_{t},V_{t}$ (that may depend on the randomness of $W^{(0)},b^{(0)},V^{(0)},b^{(1)}$ and $\mathcal{Z}$) with

With high probability over the random initialization, there exists $W^{\divideontimes},V^{\divideontimes}$ with $\|W^{\divideontimes}\|_{F},\|V^{\divideontimes}\|_{F}\leq 1$ such that for every $\eta\in\left[0,\frac{1}{\poly(m_{1},m_{2})}\right]$:

More interestingly, since (noisy) stochastic gradient descent is capable of finding approximate second-order critical points with a global convergence rate, one can show the following final convergence rate for minimizing $L^{\prime}$ for Algorithm 3.

In the setting of Theorem 3, with probability at least $99/100$, Algorithm 3 (the first SGD variant) converges in $TT_{w}=\poly\left(m_{1},m_{2}\right)$ iterations to a point

$D_{w,x,\rho}$ denote the diagonal matrix with diagonals being 0-1 signs of the first layer at $W+W^{\rho}$;

$D_{v,x,\rho}$ denote that of the second layer at weights $W+W^{\rho}$ and $V+V^{\rho}$; and

For a fixed $r\in[k]$, consider the real network $P_{\rho,\eta}$ and the pseudo network $P^{\prime}_{\rho,\eta}$:

There exists $\eta_{0}=\frac{1}{\poly(m_{1},m_{2})}$ such that, for every $\eta\leq\eta_{0}$, for every fixed $x$ with $\|x\|_{2}=1$, for every $W^{\prime},V^{\prime},W^{\prime\prime},V^{\prime\prime}$ that may depend on the randomness of the initialization and

where $O_{p}$ hides polynomial factor of $m_{1},m_{2}$.

In other words, when working with a Gaussian smoothed network (with output $P_{\rho,\eta}$), it suffices to study a pseudo network (with output $P_{\rho,\eta}^{\prime}$). In the proof of Lemma 6.7, this allows us to go from the real (smoothed) network to the pseudo (smoothed) network, and then apply Lemma 6.5.

5 Generalization

In our first variant of SGD Algorithm 3, we show a very crude Rademacher complexity bound that can be derived from the contraction lemma (a.k.a. Talagrand’s concentration inequality).

For every $\tau^{\prime}_{v},\tau^{\prime}_{w}\geq 0$, every $\sigma_{v}\in(0,1/\sqrt{m_{2}}]$, w.h.p. for every $r\in[k]$ and every $N\geq 1$, the empirical Rademacher complexity is bounded by

Since the population risk is bounded by the Rademacher complexity, combining this with Lemma 6.8, one can easily prove Theorem 3.

6 Second Variant of SGD

To show that this gives a better sampling complexity bound, we need the following stronger coupling lemma. It gives a somewhat better bound on the “cross term vanish” part comparing to the old coupling Lemma 6.5.

Under parameter choices Table 1, the last error term is at most $\varepsilon/k$.

For this SGD variant, we also have the following the stronger Rademacher complexity bound. It relies on Lemma 6.11 to reduce the function class to pseudo networks, which are only linear functions in $W^{\prime}$ and $V^{\prime}$, and then computes its Rademacher complexity.

For every $\tau^{\prime}_{v}\in$, \tau^{\prime}_{w}\in\big{[}\frac{1}{m_{1}^{3/4}},\frac{1}{m_{1}^{9/16}}\big{]}, every $\sigma_{w}\in[0,1/\sqrt{m_{1}}]$ and $\sigma_{v}\in[0,1/\sqrt{m_{2}}]$, w.h.p. for every $r\in[k]$ and every $N\geq 1$, we have by our choice of parameters in Lemma 6.7, the empirical Rademacher complexity is bounded by

Under parameter choices in Table 1, this is at most \widetilde{O}\big{(}\frac{\tau^{\prime}_{w}\tau^{\prime}_{v}m_{1}^{1/4}\sqrt{m_{2}}}{\sqrt{N}}\big{)}+\varepsilon/k.

After plugging Lemma 6.11 and Lemma 6.12 into the final proof, we can show Theorem 2.

Empirical Evaluations

To be consistent with our theorems, we choose random initialization as follows. Entries of $a_{r}$ are i.i.d. from $\mathcal{N}(0,1)$, and entries of $W,V,b_{1},b_{2}$ are i.i.d. from $\mathcal{N}(0,\frac{1}{m})$. This ensures that the output at random initialization is $\Theta(1)$.

In Figure 2, we provide an additional experiment for target function $y=F^{*}(x)=(\tanh(8x_{1})+\tanh(8x_{2})+\tanh(8x_{3})-2)^{2}\cdot\tanh(8x_{4})$.

Recall our three-layer theorem requires a slightly unconventional regularizer, namely, the $\|W\|_{2,4}$ norm of the first layer to encourage weights to be more evenly distributed across neurons. Is this really necessary? To get some preliminary idea we run our aforementioned three-layer experiment on the first data set,

once with the traditional weight decay on $W$ (which corresponds to minimizing $\|W\|_{F}$), and

once with the $\|W\|_{2,4}$ norm without weight decay.

In both cases we best tune the learning rates as well as the weight decay parameter (or the weight of the $\|W\|_{2,4}^{4}$ regularizer). We present our findings in Figure 3.

In Figure 3(a), we observe that there is no real difference between the two regularizers in terms of test error. In one case ($m=200$) using $\|W\|_{2,4}$ even gives slightly better test accuracy (but we do not want to claim this as a general phenomenon).

More importantly, since by Cauchy-Schwarz we have $\|W\|_{F}\geq\|W\|_{2,4}\geq m^{-1/4}\|W\|_{F}$. Let us consider the following norm ratio:

In Figure 3(b), we see that $ratio$ is roughly the same between two types of regularizers. This means, even if we regularize only the Frobenius norm, weights are still sufficiently distributed across neurons comparing to what we can do by regularizing $\|W\|_{2,4}$. We leave it a future work to study why SGD encourages such implicit regularization.

We provide technical preliminaries in Appendix A, give our two-layer proofs in Appendix B, and three-layer proofs in Appendix C.

Appendix A Technical Preliminaries

where the infimum is taken over all possible joint distributions over $(X,Y)$ where the marginal on $X$ (resp. $Y$) is distributed in the same way as $A$ (resp. $B$).

Slightly abusing notation, in this paper, we say a random variable $X$ satisfies $|X|\leq B$ with high probability if (1) $|X|\leq B$ w.h.p. and (2) $\mathcal{W}_{2}(X,0)\leq B$. For instance, if $g=\mathcal{N}(0,1/m)$, then $|g|\leq\widetilde{O}(1/\sqrt{m})$ with high probability.

Let $(n_{1},\alpha_{1},a_{1,1},a_{2,1}),\cdots,(n_{m},\alpha_{m},a_{1,m},a_{2,m})$ be $m$ i.i.d. samples from some distribution, where within a 4-tuples:

the marginal distribution of $a_{1,i}$ and $a_{2,i}$ is standard Gaussian $\mathcal{N}(0,1)$;

$n_{i}$ and $\alpha_{i}$ are not necessarily independent;

$a_{1,i}$ and $a_{2,i}$ are independent; and

$n_{i}$ and $\alpha_{i}$ are independent of $a_{1,i}$ and $a_{2,i}$.

Since this holds for every choice of $\{n_{i},\alpha_{i}\}_{i\in[m]}$ we can complete the proof. The second inequality follows from sub-exponential concentration bounds. ∎

If $X_{1},X_{2}$ are independent, and $X_{1},X_{3}$ are independent conditional on $X_{2}$, then $X_{1}$ and $X_{3}$ are independent.

Multiplying $\operatornamewithlimits{\mathbf{Pr}}[X_{2}=x_{2}]$ on both side leads to:

Marginalizing away $X_{2}$ gives $\operatornamewithlimits{\mathbf{Pr}}[X_{1}=x_{1},X_{3}=x_{3}]=\operatornamewithlimits{\mathbf{Pr}}[X_{1}=x_{1}]\operatornamewithlimits{\mathbf{Pr}}[X_{3}=x_{3}]$, so $X_{1}$ and $X_{3}$ are independent. ∎

A.2 Central Limit Theorem

The following Wasserstein distance bound of central limit theorem is easy to derive from known results:

In other words—by choosing $t=R_{i}/\Sigma_{i}$— we have for every $R_{i}\geq 5C^{2}$, letting $Z_{i}\sim\mathcal{N}(0,R_{i})$, $Z_{i+1}\sim\mathcal{N}(0,R_{i}+\Sigma_{i})$ be independent gaussian (also independent of $X_{i}$), it satisfies

Repeatedly applying the above inequality and starting with $Z_{1}\sim\mathcal{N}(0,5C^{2})$ and choosing $R_{i}=5C^{2}+\sum_{j=1}^{i-1}\Sigma_{j}$, we have $Z_{i}\sim\mathcal{N}(0,5C^{2}+\sum_{j=1}^{i-1}\Sigma_{j})$ and $Z_{m+1}\sim\mathcal{N}(0,5C^{2}+V)$. Using $\Sigma_{1}\geq\Sigma_{2}\geq\cdots\geq\Sigma_{m}$, we know that $\frac{5C\Sigma_{i}}{R_{i}}\leq\frac{5C}{i-1}$ for $i\geq 2$. This implies that

Finally, since $|X_{1}|\leq C$ we have $\mathcal{W}_{2}(X_{1},0)\leq C$ and $\Sigma_{1}^{2}\leq C^{2}$. By triangle inequality we have

There exists $Z\sim\mathcal{N}(0,V)$ such that $\mathcal{W}_{2}(Z_{m+1},Z)=O(C)$. All of these together imply

A.3 Interval Partition

(Indicator). $\mathfrak{s}(y,g)=0$ if $g\notin I(y)$, and $\mathfrak{s}(y,g)\in\{-1,1\}$ otherwise.

(Balanced). $\operatornamewithlimits{\mathbf{Pr}}_{g\sim\mathcal{N}(0,1)}[g\in I(y)]=\tau$ for every $y\in$.

(Symmetric). $\operatornamewithlimits{\mathbf{Pr}}_{g\sim\mathcal{N}(0,1)}\left[\mathfrak{s}(y,g)=1\right]=\operatornamewithlimits{\mathbf{Pr}}_{g\sim\mathcal{N}(0,1)}\left[\mathfrak{s}(y,g)=-1\right]$.

(Bounded). $\max_{x\in I(y)}\{\mathfrak{s}(y,x)x\}-\min_{x\in I(y)}\{\mathfrak{s}(y,x)x\}\leq 10\tau$.

We refer to $I(y)$ as an “interval” although it may actually consist of two disjoint closed intervals.

Let us just prove the case when $y\geq 0$ and the other case is by symmetry. It is clear that, since there are only two degrees of freedom, there is a unique interval $I_{1}(y)=[y-a(y),y+b(y)]$ with $a(y),b(y)\geq 0$ such that

(Half probability). $\operatornamewithlimits{\mathbf{Pr}}_{g\sim\mathcal{N}(0,1)}[g\in I_{1}(y)]=\frac{\tau}{2}$.

$[y-a(y),y+b(y)]$ and $[-y-b(y),-y+a(y)]$ intersect. In this case, consider the unique interval

It must satisfy $\operatornamewithlimits{\mathbf{Pr}}_{g\sim\mathcal{N}(0,1)}[g\in I_{2}(y)\,\wedge\,g>0]<\tau/2$, because otherwise we must have $y-a(y)\geq 0$ and the two intervals should not have intersected.

Define $\tau^{\prime}(y)=\operatornamewithlimits{\mathbf{Pr}}_{g\sim\mathcal{N}(0,1)}[g\in I_{2}(y)]<\tau$. Let $c(y)>e(y)$ be the unique positive real such that

Let $d(y)\in[e(y),c(y)]$ be the unique real such that

Finally, we define $I(y)=[-c(y),c(y)]$ and

In both cases, one can carefully verify that properties 1, 2, 3, 4 hold. Property 5 follows from the standard property of Gaussian random variable under condition $\tau\leq 1/100$ and $y\in$.

To check the final Lipschitz continuity property, recall for a standard Gaussian distribution, inside interval $[-\frac{1}{10},\frac{1}{10}]$ it behaves, up to multiplicative constant factor, similar to a uniform distribution. Therefore, the above defined functions $a(y)$ and $b(y)$ are $O(1)$-Lipschitz continuous in $y$. Let $y_{0}\geq 0$ be the unique constant such that $y-a(y)=0$ (it is unique because $y-a(y)$ monotonically decreases as $y\to 0+$. It is clear that for $y_{0}\leq y_{1}\leq y_{2}$ it satisfies

As for the turning point of $y=y_{0}$, it is clear that

so the function $I(\cdot)$ is continuous at point $y=y_{0}$. Finally, consider $y\in[-y_{0},y_{0}]$. One can verify that $e(y)$ is $O(1)$-Lipschitz continuous in $y$, and therefore the above defined $\tau^{\prime}(y)$, $d(y)$ and $c(y)$ are also $O(1)$-Lipschitz in $y$. This means, for $-y_{0}\leq y_{1}\leq y_{2}\leq y_{0}$, it also satisfies

This proves the Lipschitz continuity of $I(y)$. ∎

A.4 Hermite polynomials

Let $h_{i}(i\geq 0)$ denote the degree-$i$ (probabilists’) Hermite polynomial

where $\delta_{i,j}=1$ if $i=j$ and $\delta_{i,j}=0$ otherwise. They have the following summation and multiplication formulas.

For even $i>0$, for any $x_{1}\in$ and $b$,

For odd $i>0$, for any $x_{1}\in$ and $b$,

(Throughout this paper, the binomial $\binom{n}{m}$ is defined as $\frac{\Gamma(n+1)}{\Gamma(m+1)\Gamma(n+1-m)}$, and this allows us to write for instance $\binom{5/2}{-1/2}$ without notation change.)

Using the summation formula of Hermite polynomial, we have:

Using the multiplication formula of Hermite polynomial, we have:

Consider even $i>0$. By Lemma A.7, we have for even $i\geq 0$:

Consider odd $i>0$. By Lemma A.7, we have for odd $i>0$:

by a similar calculation as in the even $i$ case.

Then $L_{i,b}$’s are given by the recursive formula:

As a result (with the convention that $0!!=1$ and $(-1)!!=1$)

The base cases $L_{0,b}$ and $L_{1,b}$ are easy to verify. Then the lemma comes from induction. ∎

A.5 Optimization

Then, $\lambda_{\min}(\nabla^{2}f(x))\leq-{\varepsilon}$, where $\lambda_{\min}$ is the minimal eigenvalue.

We also recall the following convergence theorem of SGD for escaping saddle point.The original proof of was for constant probability but it is easy to change it to $1-p$ at the expense of paying a polynomial factor in $1/p$. The original proof did not state the objective “non-increasing” guarantee but it is easy to show it for mini-batch SGD. Recall $\eta=\frac{1}{\poly(d,B,1/\varepsilon,1/\delta)}$ in so in each iteration, the original SGD may increase the objective by no more than $\frac{\eta^{2}B}{b}$ if a batch size $b$ is used. If $b$ is polynomially large in $1/\eta$, this goal is easily achievable. In practice, however, using a very small batch size suffices.

A.6 Rademacher Complexity

Suppose $\mathcal{X}=(x_{1},\dots,x_{N})$ where each $x_{i}$ is generated i.i.d. from a distribution $\mathcal{D}$. If every $f\in\mathcal{F}$ satisfies $|f|\leq b$, for every $\delta\in(0,1)$ with probability at least $1-\delta$ over the randomness of $\mathcal{Z}$, it satisfies

Let $\mathcal{F}^{\prime}$ be the class of functions by composing $L$ with $\mathcal{F}_{1},\dots,\mathcal{F}_{k}$, that is, $\mathcal{F}^{\prime}=\{L_{x}\circ(f_{1},\dots,f_{k})\mid f_{1}\in\mathcal{F}_{1}\cdots f_{k}\in\mathcal{F}_{k}\}$. By the (vector version) of the contraction lemma of Rademacher complexity There are slightly different versions of the contraction lemma in the literature. For the scalar case without absolute value, see [43, Section 3.8]; for the scalar case with absolute value, see [14, Theorem 12]; and for the vector case without absolute value, see . it satisfies $\widehat{\mathfrak{R}}(\mathcal{Z};\mathcal{F}^{\prime})\leq O(1)\cdot\sum_{r=1}^{k}\widehat{\mathfrak{R}}(\mathcal{Z};\mathcal{F}_{r})$. ∎

(addition) $\widehat{R}(\mathcal{X};\mathcal{F}_{1}+\mathcal{F}_{2})=\widehat{R}(\mathcal{X};\mathcal{F}_{1})+\widehat{R}(\mathcal{X};\mathcal{F}_{2})$.

(contraction) $\widehat{R}(\mathcal{X};\sigma\circ\mathcal{F})\leq\widehat{R}(\mathcal{X};\mathcal{F})$.

satisfies $\widehat{R}(\mathcal{X};\mathcal{F}^{\prime})\leq 2\|w\|_{1}\max_{j\in[m]}\widehat{R}(\mathcal{X};\mathcal{F}_{j})$.

satisfies \widehat{R}(\mathcal{X};\mathcal{F}^{\prime})\leq 2D\sum_{j\in[m]}\widehat{R}(\mathcal{X};\mathcal{F}_{j})+O\big{(}\frac{BR\log m}{\sqrt{N}}\big{)}.

The first three are trivial, and the contraction lemma is classical (see for instance ). The derivations of Lemma prop:rade and Lemma prop:radf are less standard.

For Lemma prop:rade, let us choose an arbitrary $f_{j}^{(0)}$ from each $\mathcal{F}_{j}$. We write $f\in\mathcal{F}$ to denote $(f_{1},\dots,f_{m})\in\mathcal{F}_{1}\times\cdots\times\mathcal{F}_{m}$.

Above, ① is because $\xi_{i}w_{j}\sigma(f_{j}^{(0)}(x_{i}))$ is independent of $f_{j}$ and thus zero in expectation; ② uses the non-negativity of \sup_{f_{j}\in\mathcal{F}_{j}}\sum_{i\in[N]}\xi_{i}\big{(}\sigma(f_{j}(x_{i}))-\sigma(f_{j}^{(0)}(x_{i}))\big{)}; ③ is for the same reason as ①; and ④ is by Proposition prop:radd.

Above, ① is from the same derivation as Proposition prop:rade; and ② uses Proposition prop:radb by viewing function class \big{\{}x\mapsto\sum_{j=1}^{m}v_{j}\cdot\sigma(f_{j}^{(0)}(x))\big{\}} as linear in $v$. ∎

Appendix B Proofs for Two-Layer Networks

Recall from (3.1) the target $F^{*}=(f_{1}^{*},\cdots,f_{k}^{*})$ for our two-layer case is

We consider another function defined as $G(x;W)=(g_{1}(x;W),\ldots,g_{k}(x;W))$ for the weight matrix $W$, where

where $w_{i}$ is the $r$-th row of $W$ and $w_{i}^{(0)}$ is the $r$-th row of $W^{(0)}$ (our random initialization). We call this $G$ a pseudo network. For convenience, we also define a pseudo network $G^{(b)}(x;W)=(g^{(b)}_{1}(x;W),\ldots,g^{(b)}_{k}(x;W))$ without bias

Roadmap. Our analysis begins with showing that w.h.p. over the random initialization, there exists a pseudo network in the neighborhood of the initialization that can approximate the target function (see Section B.1). We then show that, near the initialization, the pseudo network approximates the actual ReLU network $F$ (see Section B.2). Therefore, there exists a ReLU network near the initialization approximating the target. Furthermore, it means that the loss surface of the ReLU network is close to that of the pseudo network, which is convex. This then allows us to show the training converges (see Section B.3). Combined with a generalization bound localized to the initialization, we can prove the final theorem that SGD learns a network with small risk.

The proof is much simpler than the three-layer case. First, the optimization landscape is almost convex, so a standard argument for convex optimization applies. While for three-layer case, the optimization landscape is no longer this nice and needs an escaping-from-saddle-point argument, which in turn requires several technicalities (smoothing, explicit regularization, weight decay, and Dropout-like noise). Second, the main technical part of the analysis, the proof of the existential result, only needs to deal with approximating one layer of neurons in the target, while in the three-layer case, a composition of the neurons needs to be approximating, requiring additional delicate arguments. It is also similar with the generalization bounds. However, we believe that the analysis in the two-layer case already shows some of the key ideas, and suggest the readers to read it before reading that for the three-layer case.

The main focus of this subsection is to show that there exists a good pseudo network near the initialization. (Combining with the coupling result of the next subsection, this translates to the real network near the initialization.)

For every $\varepsilon\in(0,\frac{1}{pk\mathfrak{C}_{\mathfrak{s}}(\phi,1)})$, letting $\varepsilon_{a}=\varepsilon/\widetilde{\Theta}(1)$, there exists

such that if $m\geq M$, then with high probability there exists $W^{\divideontimes}=(w^{\divideontimes}_{1},\ldots,w^{\divideontimes}_{m})$ with $\|W^{\divideontimes}\|_{2,\infty}\leq\frac{kpC_{0}}{\varepsilon_{a}m}$ and $\|W^{\divideontimes}\|_{F}\leq\widetilde{O}(\frac{kp\mathfrak{C}_{\mathfrak{s}}(\phi,1)}{\varepsilon_{a}\sqrt{m}})$ and

In the same setting as Lemma B.1, we have that w.h.p.

Recall the pseudo network without bias is given by

where $\alpha_{1},\beta_{1},b_{0}\sim\mathcal{N}(0,1)$ are independent random Gaussians.

where $\sqrt{m}\langle w_{j}^{(0)},w_{1,i}^{*}\rangle,\sqrt{m}b_{j}^{(0)}$ has the same distribution with $\alpha_{1},b_{0}$ in Lemma 6.3. By Lemma 6.3, we have that

Fit a combination $\sum_{i\in[p]}a^{*}_{r,i}\phi_{i}(\langle w_{1,i}^{*},x\rangle)\langle w_{2,i}^{*},x\rangle$. We can re-define (the norm grows by a maximum factor of $p$)

Fit multiple outputs. If there are $k$ outputs let us re-define (the norm grows by a maximum factor of $k$)

Now, re-scaling each $w^{\divideontimes}_{j}$ by a factor of $\frac{1}{m}$ and re-scaling $\varepsilon$ by $\frac{1}{2pk}$, we can write

Now, we apply the concentration from Lemma A.1, which implies for our parameter choice of $m$, with high probability

The above concentration holds for every fixed $x$ with high probability, and thus also holds in expectation with respect to $(x,y)\sim\mathcal{D}$. This proves the first statement. As for the second statement on $L(G^{(b)}(x;W^{\divideontimes}),y)$, it follows from the Lipschitz continuity of $L$.