Generalization Bounds of Stochastic Gradient Descent for Wide and Deep Neural Networks

Introduction

Deep learning has achieved great success in a wide range of applications including image processing (Krizhevsky et al., 2012), natural language processing (Hinton et al., 2012) and reinforcement learning (Silver et al., 2016). Most of the deep neural networks used in practice are highly over-parameterized, such that the number of parameters is much larger than the number of training data. One of the mysteries in deep learning is that, even in an over-parameterized regime, neural networks trained with stochastic gradient descent can still give small test error and do not overfit. In fact, a famous empirical study by Zhang et al. (2017) shows the following phenomena:

Even if one replaces the real labels of a training data set with purely random labels, an over-parameterized neural network can still fit the training data perfectly. However since the labels are independent of the input, the resulting neural network does not generalize to the test dataset.

If the same over-parameterized network is trained with real labels, it not only achieves small training loss, but also generalizes well to the test dataset.

While a series of recent work has theoretically shown that a sufficiently over-parameterized (i.e., sufficiently wide) neural network can fit random labels (Du et al., 2019b; Allen-Zhu et al., 2019b; Du et al., 2019a; Zou et al., 2019), the reason why it can generalize well when trained with real labels is less understood. Existing generalization bounds for deep neural networks (Neyshabur et al., 2015; Bartlett et al., 2017; Neyshabur et al., 2018; Golowich et al., 2018; Dziugaite and Roy, 2017; Arora et al., 2018; Li et al., 2018; Wei et al., 2019; Neyshabur et al., 2019) based on uniform convergence usually cannot provide non-vacuous bounds (Langford and Caruana, 2002; Dziugaite and Roy, 2017) in the over-parameterized regime. In fact, the empirical observation by Zhang et al. (2017) indicates that in order to understand deep learning, it is important to distinguish the true data labels from random labels when studying generalization. In other words, it is essential to quantify the “classifiability” of the underlying data distribution, i.e., how difficult it can be classified.

Certain effort has been made to take the “classifiability” of the data distribution into account for generalization analysis of neural networks. Brutzkus et al. (2018) showed that stochastic gradient descent (SGD) can learn an over-parameterized two-layer neural network with good generalization for linearly separable data. Li and Liang (2018) proved that, if the data satisfy certain structural assumption, SGD can learn an over-parameterized two-layer network with fixed second layer weights and achieve a small generalization error. Allen-Zhu et al. (2019a) studied the generalization performance of SGD and its variants for learning two-layer and three-layer networks, and used the risk of smaller two-layer or three-layer networks with smooth activation functions to characterize the classifiability of the data distribution. There is another line of studies on the algorithm-dependent generalization bounds of neural networks in the over-parameterized regime (Daniely, 2017; Arora et al., 2019a; Cao and Gu, 2020; Yehudai and Shamir, 2019; E et al., 2019), which quantifies the classifiability of the data with a reference function class defined by random features (Rahimi and Recht, 2008, 2009) or kernelsSince random feature models and kernel methods are highly related (Rahimi and Recht, 2008, 2009), we group them into the same category. More details are discussed in Section 3.2.. Specifically, Daniely (2017) showed that a neural network of large enough size is competitive with the best function in the conjugate kernel class of the network. Arora et al. (2019a) gave a generalization error bound for two-layer ReLU networks with fixed second layer weights based on a ReLU kernel function. Cao and Gu (2020) showed that deep ReLU networks trained with gradient descent can achieve small generalization error if the data can be separated by certain random feature model (Rahimi and Recht, 2009) with a margin. Yehudai and Shamir (2019) used the expected loss of a similar random feature model to quantify the generalization error of two-layer neural networks with smooth activation functions. A similar generalization error bound was also given by E et al. (2019), where the authors studied the optimization and generalization of two-layer networks trained with gradient descent. However, all the aforementioned results are still far from satisfactory: they are either limited to two-layer networks, or restricted to very simple and special reference function classes.

In this paper, we aim at providing a sharper and generic analysis on the generalization of deep ReLU networks trained by SGD. In detail, we base our analysis upon the key observations that near random initialization, the neural network function is almost a linear function of its parameters and the loss function is locally almost convex. This enables us to prove a cumulative loss bound of SGD, which further leads to a generalization bound by online-to-batch conversion (Cesa-Bianchi et al., 2004). The main contributions of our work are summarized as follows:

We give a bound on the expected -$1$ error of deep ReLU networks trained by SGD with random initialization. Our result relates the generalization bound of an over-parameterized ReLU network with a random feature model defined by the network gradients, which we call neural tangent random feature (NTRF) model. It also suggests an algorithm-dependent generalization error bound of order $\widetilde{\mathcal{O}}(n^{-1/2})$, which is independent of network width, if the data can be classified by the NTRF model with small enough error.

Our analysis is general enough to cover recent generalization error bounds for neural networks with random feature based reference function classes, and provides better bounds. Our expected -$1$ error bound directly covers the result by Cao and Gu (2020), and gives a tighter sample complexity when reduced to their setting, i.e., $\widetilde{\mathcal{O}}(1/\epsilon^{2})$ versus $\widetilde{\mathcal{O}}(1/\epsilon^{4})$ where $\epsilon$ is the target generalization error. Compared with recent results by Yehudai and Shamir (2019); E et al. (2019) who only studied two-layer networks, our bound not only works for deep networks, but also uses a larger reference function class when reduced to the two-layer setting, and therefore is sharper.

Our result has a direct connection to the neural tangent kernel studied in Jacot et al. (2018). When interpreted in the language of kernel method, our result gives a generalization bound in the form of $\widetilde{\mathcal{O}}(L\cdot\sqrt{\mathbf{y}^{\top}(\bm{\Theta}^{(L)})^{-1}\mathbf{y}/n})$, where $\mathbf{y}$ is the training label vector, and $\bm{\Theta}^{(L)}$ is the neural tangent kernel matrix defined on the training input data. This form of generalization bound is similar to, but more general and tighter than the bound given by Arora et al. (2019a).

Problem Setup

where $\sigma(\cdot)$ is the entry-wise activation function. In this paper, we only consider the ReLU activation function $\sigma(z)=\max\{0,z\}$, which is the most commonly used activation function in applications. It is also arguably one of the most difficult activation functions to analyze, due to its non-smoothess. We remark that our result can be generalized to many other Lipschitz continuous and smooth activation functions. For simplicity, we follow Allen-Zhu et al. (2019b); Du et al. (2019a) and assume that the widths of each hidden layer are the same. Our result can be easily extended to the setting that the widths of each layer are not equal but in the same order, as discussed in Zou et al. (2019); Cao and Gu (2020).

When $L=1$, the neural network reduces to a linear function, which has been well-studied. Therefore, for notational simplicity we focus on the case $L\geq 2$, where the parameter space is defined as

The goal of neural network learning is to minimize the expected risk, i.e.,

The initialization scheme for $\mathbf{W}^{(1)}$ given in Algorithm 1 generates each entry of the weight matrices from a zero-mean independent Gaussian distribution, whose variance is determined by the rule that the expected length of the output vector in each hidden layer is equal to the length of the input. This initialization method is also known as He initialization (He et al., 2015). Here the last layer parameter is initialized with variance $1/m$ instead of $2/m$ since the last layer is not associated with the ReLU activation function.

Main Results

In this section we present the main results of this paper. In Section 3.1 we give an expected -$1$ error bound against a neural tangent random feature reference function class. In Section 3.2, we discuss the connection between our result and the neural tangent kernel proposed in Jacot et al. (2018).

For any $\mathbf{W}\in\mathcal{W}$, we define its $\omega$-neighborhood as

Below we introduce the neural tangent random feature function class, which serves as a reference function class to measure the “classifiability” of the data, i.e., how easy it can be classified.

Let $\mathbf{W}^{(1)}$ be generated via the initialization scheme in Algorithm 1. The neural tangent random feature (NTRF) function class is defined as

where $R>0$ measures the size of the function class, and $m$ is the width of the neural network.

The name “neural tangent random feature” is inspired by the neural tangent kernel proposed by Jacot et al. (2018), because the random features are the gradients of the neural network with random weights. Connections between the neural tangent random features and the neural tangent kernel will be discussed in Section 3.2.

We are ready to present our main result on the expected -$1$ error bound of Algorithm 1.

For any $\delta\in(0,e^{-1}]$ and $R>0$, there exists

such that if $m\geq m^{*}(\delta,R,L,n)$, then with probability at least $1-\delta$ over the randomness of $\mathbf{W}^{(1)}$, the output of Algorithm 1 with step size $\eta=\kappa\cdot R/(m\sqrt{n})$ for some small enough absolute constant $\kappa$ satisfies

where the expectation is taken over the uniform draw of $\widehat{\mathbf{W}}$ from $\{\mathbf{W}^{(1)},\ldots,\mathbf{W}^{(n)}\}$.

The expected -$1$ error bound given by Theorem 3.3 consists of two terms: The first term in (3.1) relates the expected -$1$ error achieved by Algorithm 1 with a reference function class–the NTRF function class in Definition 3.2. The second term in (3.1) is a standard large-deviation error term. As long as $R=\widetilde{\mathcal{O}}(1)$, this term matches the standard $\widetilde{\mathcal{O}}(n^{-1/2})$ rate in PAC learning bounds (Shalev-Shwartz and Ben-David, 2014).

The parameter $R$ in Theorem 3.3 is from the NTRF class and introduces a trade-off in the bound: when $R$ is small, the corresponding NTRF class $\mathcal{F}(\mathbf{W}^{(1)},R)$ is small, making the first term in (3.1) large, and the second term in (3.1) is small. When $R$ is large, the corresponding function class $\mathcal{F}(\mathbf{W}^{(1)},R)$ is large, so the first term in (3.1) is small, and the second term will be large. In particular, if we set $R=\widetilde{\mathcal{O}}(1)$, the second term in (3.1) will be $\widetilde{\mathcal{O}}(n^{-1/2})$. In this case, the “classifiability” of the underlying data distribution $\mathcal{D}$ is determined by how well its i.i.d. samples can be classified by $\mathcal{F}(\mathbf{W}^{(1)},\widetilde{\mathcal{O}}(1))$. In other words, Theorem 3.3 suggests that if the data can be classified by a function in the NTRF function class $\mathcal{F}(\mathbf{W}^{(1)},\widetilde{\mathcal{O}}(1))$ with a small training error, the over-parameterized ReLU network learnt by Algorithm 1 will have a small generalization error.

The expected -$1$ error bound given by Theorem 3.3 is in a very general form. It directly covers the result given by Cao and Gu (2020). In Appendix A.1, we show that under the same assumptions made in Cao and Gu (2020), to achieve $\epsilon$ expected -$1$ error, our result requires a sample complexity of order $\widetilde{\mathcal{O}}(\epsilon^{-2})$, which outperforms the result in Cao and Gu (2020) by a factor of $\epsilon^{-2}$.

Our generalization bound can also be compared with two recent results (Yehudai and Shamir, 2019; E et al., 2019) for two-layer neural networks. When $L=2$, the NTRF function class $\mathcal{F}(\mathbf{W}^{(1)},\widetilde{\mathcal{O}}(1))$ can be written as

In contrast, the reference function classes studied by Yehudai and Shamir (2019) and E et al. (2019) are contained in the following random feature class:

2 Connection to Neural Tangent Kernel

Besides quantifying the classifiability of the data with the NTRF function class $\mathcal{F}(\mathbf{W}^{(1)},\widetilde{\mathcal{O}}(1))$, an alternative way to apply Theorem 3.3 is to check how large the parameter $R$ needs to be in order to make the first term in (3.1) small enough (e.g., smaller than $n^{-1/2}$). In this subsection, we show that this type of analysis connects Theorem 3.3 to the neural tangent kernel proposed in Jacot et al. (2018) and later studied by Yang (2019); Lee et al. (2019); Arora et al. (2019b). Specifically, we provide an expected -$1$ error bound in terms of the neural tangent kernel matrix defined over the training data. We first define the neural tangent kernel matrix for the neural network function in (2.1).

Then we call $\bm{\Theta}^{(L)}=[(\widetilde{\bm{\Theta}}_{i,j}^{(L)}+\bm{\Sigma}_{i,j}^{(L)})/2]_{n\times n}$ the neural tangent kernel matrix of an $L$-layer ReLU network on training inputs $\mathbf{x}_{1},\ldots,\mathbf{x}_{n}$.

Definition 3.7 is the same as the original definition in Jacot et al. (2018) when restricting the kernel function on $\{\mathbf{x}_{1},\ldots,\mathbf{x}_{n}\}$, except that there is an extra coefficient $2$ in the second and third lines. This extra factor is due to the difference in initialization schemes–in our paper the entries of hidden layer matrices are randomly generated with variance $2/m$, while in Jacot et al. (2018) the variance of the random initialization is $1/m$. We remark that this extra factor $2$ in Definition 3.7 will remove the exponential dependence on the network depth $L$ in the kernel matrix, which is appealing. In fact, it is easy to check that under our scaling, the diagonal entries of $\bm{\Sigma}^{(L)}$ are all $1$’s, and the diagonal entries of $\widetilde{\bm{\Theta}}^{(L)}$ are all $L$’s.

The following lemma is a summary of Theorem 1 and Proposition 2 in Jacot et al. (2018), which ensures that $\bm{\Theta}^{(L)}$ is the infinite-width limit of the Gram matrix $(m^{-1}\langle\nabla_{\mathbf{W}}f_{\mathbf{W}^{(1)}}(\mathbf{x}_{i}),\nabla_{\mathbf{W}}f_{\mathbf{W}^{(1)}}(\mathbf{x}_{j})\rangle)_{n\times n}$, and is positive-definite as long as no two training inputs are parallel.

For an $L$ layer ReLU network with parameter set $\mathbf{W}^{(1)}$ initialized in Algorithm 1, as the network width $m\rightarrow\infty$The original result by Jacot et al. (2018) requires that the widths of different layers go to infinity sequentially. Their result was later improved by Yang (2019) such that the widths of different layers can go to infinity simultaneously., it holds that

where the expectation is taken over the randomness of $\mathbf{W}^{(1)}$. Moreover, as long as each pair of inputs among $\mathbf{x}_{1},\ldots,\mathbf{x}_{n}\in S^{d-1}$ are not parallel, $\bm{\Theta}^{(L)}$ is positive-definite.

Lemmas 3.8 clearly shows the difference between our neural tangent kernel matrix $\bm{\Theta}^{(L)}$ in Definition 3.7 and the Gram matrix $\mathbf{K}^{(L)}$ defined in Definition 5.1 in Du et al. (2019a). For any $i,j\in[n]$, by Lemma 3.8 we have

In contrast, the corresponding entry in $\mathbf{K}^{(L)}$ is

It can be seen that our definition of kernel matrix takes all layers into consideration, while Du et al. (2019a) only considered the last hidden layer (i.e., second last layer). Moreover, it is clear that $\bm{\Theta}^{(L)}\succeq\mathbf{K}^{(L)}$. Since the smallest eigenvalue of the kernel matrix plays a key role in the analysis of optimization and generalization of over-parameterized neural networks (Du et al., 2019b, a; Arora et al., 2019a), our neural tangent kernel matrix can potentially lead to better bounds than the Gram matrix studied in Du et al. (2019a).

Let $\mathbf{y}=(y_{1},\ldots,y_{n})^{\top}$ and $\lambda_{0}=\lambda_{\min}(\bm{\Theta}^{(L)})$. For any $\delta\in(0,e^{-1}]$, there exists $\widetilde{m}^{*}(\delta,L,n,\lambda_{0})$ that only depends on $\delta,L,n$ and $\lambda_{0}$ such that if $m\geq\widetilde{m}^{*}(\delta,L,n,\lambda_{0})$, then with probability at least $1-\delta$ over the randomness of $\mathbf{W}^{(1)}$, the output of Algorithm 1 with step size $\eta=\kappa\cdot\inf_{\widetilde{y}_{i}y_{i}\geq 1}\sqrt{\widetilde{\mathbf{y}}^{\top}(\bm{\Theta}^{(L)})^{-1}\widetilde{\mathbf{y}}}/(m\sqrt{n})$ for some small enough absolute constant $\kappa$ satisfies

where the expectation is taken over the uniform draw of $\widehat{\mathbf{W}}$ from $\{\mathbf{W}^{(1)},\ldots,\mathbf{W}^{(n)}\}$.

Apparently, by choosing $\widetilde{\mathbf{y}}=\mathbf{y}$ in Corollary 3.10, one can also obtain a bound on the expected $0-1$ error of the form \widetilde{\mathcal{O}}\big{[}L\cdot\sqrt{\mathbf{y}^{\top}(\bm{\Theta}^{(L)})^{-1}\mathbf{y}/n}\big{]}+\mathcal{O}\big{[}\sqrt{\log(1/\delta)/n}\big{]}.

Corollary 3.10 gives an algorithm-dependent generalization error bound of over-parameterized $L$-layer neural networks trained with SGD. It is worth noting that recently Arora et al. (2019a) gives a generalization bound \widetilde{\mathcal{O}}\big{(}\sqrt{\mathbf{y}^{\top}(\mathbf{H}^{\infty})^{-1}\mathbf{y}/n}\big{)} for two-layer networks with fixed second layer weights, where $\mathbf{H}^{\infty}$ is defined as

Our result in Corollary 3.10 can be specialized to two-layer neural networks by choosing $L=2$, and yields a bound \widetilde{\mathcal{O}}\big{(}\sqrt{\mathbf{y}^{\top}(\bm{\Theta}^{(2)})^{-1}\mathbf{y}/n}\big{)}, where

Corollary 3.10 is based on the asymptotic convergence result in Lemma 3.8, which does not show how wide the network need to be in order to make the Gram matrix close enough to the NTK matrix. Very recently, Arora et al. (2019b) provided a non-asymptotic convergence result for the Gram matrix, and showed the equivalence between an infinitely wide network trained by gradient flow and a kernel regression predictor using neural tangent kernel, which suggests that the generalization of deep neural networks trained by gradient flow can potentially be measured by the corresponding NTK. Utilizing this non-asymptotic convergence result, one can potentially specify the detailed dependency of $\widetilde{m}^{*}(\delta,L,n,\lambda_{0})$ on $\delta$, $L$, $n$ and $\lambda_{0}$ in Corollary 3.10.

Corollary 3.10 demonstrates that the generalization bound given by Theorem 3.3 does not increase with network width $m$, as long as $m$ is large enough. Moreover, it provides a clear characterization of the classifiability of data. In fact, the $\sqrt{\widetilde{\mathbf{y}}^{\top}(\bm{\Theta}^{(L)})^{-1}\widetilde{\mathbf{y}}}$ factor in the generalization bound given in Corollary 3.10 is exactly the NTK-induced RKHS norm of the kernel regression classifier on data $\{(\mathbf{x}_{i},\widetilde{y}_{i})\}_{i=1}^{n}$. Therefore, if $y=f^{*}(\mathbf{x})$ for some $f^{*}(\cdot)$ with bounded norm in the NTK-induced reproducing kernel Hilbert space (RKHS), then over-parameterized neural networks trained with SGD generalize well. In Appendix E, we provide some numerical evaluation of the leading terms in the generalization bounds in Theorem 3.3 and Corollary 3.10 to demonstrate that they are very informative on real-world datasets.

Proof of Main Theory

In this section we provide the proof of Theorem 3.3 and Corollary 3.10, and explain the intuition behind the proof. For notational simplicity, for $i\in[n]$ we denote $L_{i}(\mathbf{W})=L_{(\mathbf{x}_{i},y_{i})}(\mathbf{W})$.

Before giving the proof of Theorem 3.3, we first introduce several lemmas. The following lemma states that near initialization, the neural network function is almost linear in terms of its weights.

There exists an absolute constant $\kappa$ such that, with probability at least $1-\mathcal{O}(nL^{2})\cdot\exp[-\Omega(m\omega^{2/3}L)]$ over the randomness of $\mathbf{W}^{(1)}$, for all $i\in[n]$ and $\mathbf{W},\mathbf{W}^{\prime}\in\mathcal{B}(\mathbf{W}^{(1)},\omega)$ with $\omega\leq\kappa L^{-6}[\log(m)]^{-3/2}$, it holds uniformly that

There exists an absolute constant $\kappa$ such that, with probability at least $1-\mathcal{O}(nL^{2})\cdot\exp[-\Omega(m\omega^{2/3}L)]$ over the randomness of $\mathbf{W}^{(1)}$, for any $\epsilon>0$, $i\in[n]$ and $\mathbf{W},\mathbf{W}^{\prime}\in\mathcal{B}(\mathbf{W}^{(1)},\omega)$ with $\omega\leq\kappa L^{-6}m^{-3/8}[\log(m)]^{-3/2}\epsilon^{3/4}$, it holds uniformly that

The locally almost convex property of the loss function given by Lemma 4.2 implies that the dynamics of Algorithm 1 is similar to the dynamics of convex optimization. We can therefore derive a bound of the cumulative loss. The result is given in the following lemma.

For any $\epsilon,\delta,R>0$, there exists

such that if $m\geq m^{*}(\epsilon,\delta,R,L)$, then with probability at least $1-\delta$ over the randomness of $\mathbf{W}^{(1)}$, for any $\mathbf{W}^{*}\in\mathcal{B}(\mathbf{W}^{(1)},Rm^{-1/2})$, Algorithm 1 with $\eta=\nu\epsilon/(Lm)$, $n=L^{2}R^{2}/(2\nu\epsilon^{2})$ for some small enough absolute constant $\nu$ has the following cumulative loss bound:

We now finalize the proof by applying an online-to-batch conversion argument (Cesa-Bianchi et al., 2004), and use Lemma 4.1 to relate the neural network function with a function in the NTRF function class.

Note that for any $i\in[n]$, $\mathbf{W}^{(i)}$ only depends on $(\mathbf{x}_{1},y_{1}),\ldots,(\mathbf{x}_{i-1},y_{i-1})$ and is independent of $(\mathbf{x}_{i},y_{i})$. Therefore by Proposition 1 in Cesa-Bianchi et al. (2004), with probability at least $1-\delta$ we have

for all $\mathbf{W}^{*}\in\mathcal{B}(\mathbf{W}^{(1)},Rm^{-1/2})$. We now compare the neural network function $f_{\mathbf{W}^{*}}(\mathbf{x}_{i})$ with the function $F_{\mathbf{W}^{(1)},\mathbf{W}^{*}}(\mathbf{x}_{i}):=f_{\mathbf{W}^{(1)}}(\mathbf{x}_{i})+\langle\nabla f_{\mathbf{W}^{(1)}}(\mathbf{x}_{i}),\mathbf{W}^{*}-\mathbf{W}^{(1)}\rangle\in\mathcal{F}(\mathbf{W}^{(1)},R)$. We have

Taking infimum over $\mathbf{W}^{*}\in\mathcal{B}(\mathbf{W}^{(1)},Rm^{-1/2})$ and rescaling $\delta$ finishes the proof. ∎

2 Proof of Corollary 3.10

In this subsection we prove Corollary 3.10. The following lemma shows that at initialization, with high probability, the neural network function value at all the training inputs are of order $\widetilde{\mathcal{O}}(1)$.

For any $\delta>0$, if $m\geq KL\log(nL/\delta)$ for a large enough absolute constant $K$, then with probability at least $1-\delta$, $|f_{\mathbf{W}^{(1)}}(\bm{x}_{i})|\leq\mathcal{O}(\sqrt{\log(n/\delta)})$ for all $i\in[n]$.

We now present the proof of Corollary 3.10. The idea is to construct suitable target values $\widehat{y}_{1},\ldots,\widehat{y}_{n}$, and then bound the norm of the solution of the linear equations $\widehat{y}_{i}=\langle\nabla f_{\mathbf{W}^{(1)}}(\mathbf{x}_{i}),\mathbf{W}\rangle$, $i\in[n]$. In specific, for any $\widetilde{\mathbf{y}}$ with $\widetilde{y}_{i}y_{i}\geq 1$, we examine the minimum distance solution to $\mathbf{W}^{(1)}$ that fit the data $\{(\mathbf{x}_{i},\widetilde{y}_{i})\}_{i=1}^{n}$ well and use it to construct a specific function in \mathcal{F}\big{(}\mathbf{W}^{(1)},\widetilde{\mathcal{O}}\big{(}\sqrt{\widetilde{\mathbf{y}}^{\top}(\bm{\Theta}^{(L)})^{-1}\widetilde{\mathbf{y}}}\big{)}\big{)}.

Therefore by (4.5) and the fact that $\|\widehat{\mathbf{y}}\|_{2}^{2}=\overline{B}^{2}n$, we have

and therefore \mathbf{W}\in\mathcal{B}\big{(}\mathbf{0},\mathcal{O}\big{(}\sqrt{\widetilde{\mathbf{y}}^{\top}(\bm{\Theta}^{(L)})^{-1}\widetilde{\mathbf{y}}}\cdot m^{-1/2}\big{)}\big{)}. Moreover, by (4.6), we have $\widehat{y}_{i}=\langle\nabla_{\mathbf{W}}f_{\mathbf{W}^{(1)}}(\mathbf{x}_{i}),\mathbf{W}\rangle$. Plugging this into (4.4) then gives

Since \widehat{f}(\cdot)=f_{\mathbf{W}^{(1)}}(\cdot)+\langle\nabla_{\mathbf{W}}f_{\mathbf{W}^{(1)}}(\cdot),\mathbf{W}\rangle\in\mathcal{F}\big{(}\mathbf{W}^{(1)},\widetilde{\mathcal{O}}\big{(}\sqrt{\widetilde{\mathbf{y}}^{\top}(\bm{\Theta}^{(L)})^{-1}\widetilde{\mathbf{y}}}\big{)}\big{)}, applying Theorem 3.3 and taking infimum over $\widetilde{\mathbf{y}}$ completes the proof. ∎

Conclusions and Future Work

In this paper we provide an expected -$1$ error bound for wide and deep ReLU networks trained with SGD. This generalization error bound is measured by the NTRF function class. The connection to the neural tangent kernel function studied in Jacot et al. (2018) is also discussed. Our result covers a series of recent generalization bounds for wide enough neural networks, and provides better bounds.

An important future work is to improve the over-parameterization condition in Theorem 3.3 and Corollary 3.10. Other future directions include proving sample complexity lower bounds in the over-parameterized regime, implementing the results in Jain et al. (2019) to obtain last iterate bound of SGD, and establishing uniform convergence based generalization bounds for over-parameterized neural networks with methods developped in Bartlett et al. (2017); Neyshabur et al. (2018); Long and Sedghi (2019).

Acknowledgement

We would like to thank Peter Bartlett for a valuable discussion, and Simon S. Du for pointing out a related work (Arora et al., 2019b). We also thank the anonymous reviewers and area chair for their helpful comments. This research was sponsored in part by the National Science Foundation CAREER Award IIS-1906169, IIS-1903202, and Salesforce Deep Learning Research Award. The views and conclusions contained in this paper are those of the authors and should not be interpreted as representing any funding agencies.

Appendix A Comparison with Recent Results

In this section we compare our result in Theorem 3.3 with recent generalization error bounds for over-paramerized neural networks by Cao and Gu (2020); Yehudai and Shamir (2019); E et al. (2019), and backup our discussions in Remark 3.5 and Remark 3.6.

In this section we provide direct comparison between our result in Theorem 3.3 and Theorem 4.4 in Cao and Gu (2020). To concretely compare these two results, we apply our result to the setting studied in Cao and Gu (2020), which is based on the following assumption.

Under Assumption 3.1 and Assumption A.1, for any $\delta\in(0,e^{-1}]$, there exists

such that if $m\geq m^{*}(\delta,R,L,n)$, then with probability at least $1-\delta$ over the randomness of $\mathbf{W}^{(1)}$, the parameters given by Algorithm 1 with $\eta=\kappa\cdot R/(m\sqrt{n})$ for some small enough absolute constant $\kappa$ satisfies

where the expectation is taken over the draws of training examples $\{(\mathbf{x}_{i},y_{i})\}_{i=1}^{n}$ as well as the uniform draw of $\widehat{\mathbf{W}}$ from $\{\mathbf{W}^{(1)},\ldots,\mathbf{W}^{(n)}\}$.

By setting the expected -$1$ loss bound to $\epsilon$, we obtain a sample complexity of order $\widetilde{\mathcal{O}}(4^{L}\cdot\gamma^{-2}\epsilon^{-2})$, which is better than the sample complexity given in Cao and Gu (2020) by a factor of $\epsilon^{-2}$.

A.2 Comparison with Yehudai and Shamir (2019); E et al. (2019)

Here we give a detailed explanation to Remark 3.6, where we compare our result with Yehudai and Shamir (2019); E et al. (2019). The reference function classes studied in these two papers share the same general form:

and $\sigma(\cdot)$ is the activation function of interest.

We compare our result with the bounds given by Yehudai and Shamir (2019); E et al. (2019) by comparing the reference function classes we use. Apparently, a larger reference function class in general gives a better generalization error bound. Such a comparison requires us to adjust the scaling of initialized parameters. Based on our previous discussion, it is easy to see that the initialized second layer weights in our work and Yehudai and Shamir (2019); E et al. (2019) are all of the same scaling. However, the $\|\cdot\|_{2}$ of first layer weight matrix in Yehudai and Shamir (2019); E et al. (2019) is larger than ours by a factor of $\sqrt{m}$. Adjusting this scaling difference will give an extra factor $\sqrt{m}$, which matches the $\sqrt{m}$ factor in the definition of our neural network function. Note that even after adjusting the scaling of parameters, these random feature function classes are not directly comparable, since the activation functions and the distributions of random weights are different. However, an informal comparison can already clearly show the advantage of our result. Moreover, we remark that at least for two-layer networks, our analysis can be easily generalized to other activation functions and initialization methods, and the resulting NTRF class should be strictly larger than the random feature function classes used in Yehudai and Shamir (2019); E et al. (2019). This justifies our discussion in Remark 3.6.

Appendix B Proofs of Technical Lemmas in Section 4

In this section we provide the proofs of the technical lemmas in Section 4. We first introduce some extra notations. Following Allen-Zhu et al. (2019b), for a parameter collection $\mathbf{W}$ and $i\in[n]$, we denote

as the hidden layer outputs of the network. We also define binary diagonal matrices

For $i\in[n]$ and $l\in[L-1]$, we use $\mathbf{h}_{i,l}^{\prime}$, $\mathbf{D}_{i,l}^{\prime}$ and $\mathbf{h}_{i,l}^{(1)}$, $\mathbf{D}_{i,l}^{(1)}$ to denote the hidden layer outputs and binary diagonal matrices with parameter collections $\mathbf{W}^{\prime}$ and $\mathbf{W}^{(1)}$ respectively. We also implement the following matrix product notation which is also used in Zou et al. (2019); Cao and Gu (2020):

With this notation, we have the following matrix product representation of the neural network gradients:

The following two lemmas are proved based on several results given by Allen-Zhu et al. (2019b). Note that in their paper, both the first and the last layers of the network are fixed, which is slightly different from our setting. We remark that this difference does not affect the result.

If $\omega\leq\mathcal{O}(L^{-9/2}[\log(m)]^{-3})$, then with probability at least $1-\mathcal{O}(nL)\cdot\exp[-\Omega(m\omega^{2/3}L)]$, $1/2\leq\|\mathbf{h}_{i,l}\|_{2}\leq 3/2$ for all $\mathbf{W}\in\mathcal{B}(\mathbf{W}^{(1)},\omega)$, $i\in[n]$ and $l\in[L-1]$.

If $\omega\leq\mathcal{O}(L^{-6}[\log(m)]^{-3})$, then with probability at least $1-\mathcal{O}(nL^{2})\cdot\exp[-\Omega(m\omega^{2/3}L)]$, uniformly over:

any $i\in[n]$, $1\leq l_{1}<l_{2}\leq L-1$

any diagonal matrices $\mathbf{D}_{i,1}^{\prime\prime},\ldots,\mathbf{D}_{i,L-1}^{\prime\prime}\in^{m\times m}$ with at most $\mathcal{O}(m\omega^{2/3}L)$ non-zero entries,

For all $\mathbf{W}\in\mathcal{B}(\mathbf{W}^{(1)},\omega)$, $\|\prod_{r=l_{1}}^{l_{2}}(\mathbf{D}_{i,r}+\mathbf{D}_{i,r}^{\prime\prime})\mathbf{W}_{r}\|_{2}\leq\mathcal{O}(\sqrt{L})$.

For all $\mathbf{W}\in\mathcal{B}(\mathbf{W}^{(1)},\omega)$, $\|\mathbf{W}_{L}\prod_{r=l_{1}}^{L-1}(\mathbf{D}_{i,r}+\mathbf{D}_{i,r}^{\prime\prime})\mathbf{W}_{r}\|_{2}\leq\mathcal{O}(1)$.

For all $\mathbf{W},\mathbf{W}^{\prime}\in\mathcal{B}(\mathbf{W}^{(1)},\omega)$,

Since $f_{\mathbf{W}^{\prime}}(\mathbf{x}_{i})=\sqrt{m}\cdot\mathbf{W}_{L}^{\prime}\mathbf{h}_{i,L-1}^{\prime}$, $f_{\mathbf{W}}(\mathbf{x}_{i})=\sqrt{m}\cdot\mathbf{W}_{L}\mathbf{h}_{i,L-1}$, by direct calculation, we have

By (iii) in Lemma B.2, with probability at least $1-\mathcal{O}(nL^{2})\cdot\exp[-\Omega(m\omega^{2/3}L)]$, we have

where the last inequality follows by Lemma B.1. This inequality finishes the proof. ∎

B.2 Proof of Lemma 4.2

Intuitively, Lemma 4.2 follows by the fact that the composition of a convex function and an almost linear function is almost convex. The detailed proof is as follows.

Therefore by triangle inequality, we have

where the last inequality again follows by \omega\leq\mathcal{O}\big{(}L^{-9/4}m^{-3/8}[\log(m)]^{-3/8}\epsilon^{3/4}\big{)}. ∎

B.3 Proof of Lemma 4.3

To prove Lemma 4.3, we first introduce the following lemma which provides an upper bound for the gradient of the neural network function near initialization.

There exists an absolute constant $\kappa$ such that, with probability at least $1-\mathcal{O}(nL^{2})\cdot\exp[-\Omega(m\omega^{2/3}L)]$, for all $i\in[n]$, $l\in[L]$ and $\mathbf{W}\in\mathcal{B}(\mathbf{W}^{(1)},\omega)$ with $\omega\leq\kappa L^{-6}[\log(m)]^{-3}$, it holds uniformly that

We now provide the final proof of Lemma 4.3.

Let $\omega=C_{1}L^{-6}m^{-3/8}[\log(m)]^{-3}\epsilon^{3/4}$, where $C_{1}$ is a small enough absolute constant such that the conditions on $\omega$ given in Lemmas 4.2 and B.3 hold. It is easy to see that as long as $m\geq C_{1}^{-8}R^{8}L^{48}[\log(m)]^{12}\epsilon^{-6}$, we have $\mathbf{W}^{*}\in\mathcal{B}(\mathbf{W}^{(1)},\omega)$. We now show that under our parameter choice, $\mathbf{W}^{(1)},\ldots,\mathbf{W}^{(n)}$ are inside $\mathcal{B}(\mathbf{W}^{(1)},\omega)$ as well.

This result follows by simple induction. Clearly we have $\mathbf{W}^{(1)}\in\mathcal{B}(\mathbf{W}^{(1)},\omega)$. Suppose that $\mathbf{W}^{(1)},\ldots,\mathbf{W}^{(i)}\in\mathcal{B}(\mathbf{W}^{(1)},\omega)$. Then by Lemma B.3, for $l\in[L]$ we have $\|\nabla_{\mathbf{W}_{l}}L_{i}(\mathbf{W}^{(i)})\|_{F}\leq\mathcal{O}(\sqrt{m})$. Therefore

Plugging in our parameter choice $\eta=\nu\epsilon/(Lm)$, $n=L^{2}R^{2}/(2\nu\epsilon^{2})$ for some small enough absolute constant $\nu$ gives

where the last inequality holds as long as $m\geq C_{2}R^{16}L^{56}[\log(m)]^{12}\epsilon^{-14}$ for some large enough constant $C_{2}$. Therefore by induction we see that $\mathbf{W}^{(1)},\ldots,\mathbf{W}^{(n)}\in\mathcal{B}(\mathbf{W}^{(1)},\omega)$. As a result, the conditions of Lemmas 4.2 and B.3 are satisfied for $\mathbf{W}^{*}$ and $\mathbf{W}^{(1)},\ldots,\mathbf{W}^{(n)}$.

In the following, we utilize the results of Lemmas 4.2 and B.3 to prove the bound of cumulative loss. First of all, by Lemma 4.2, we have

Note that for the matrix inner product we have the equality $2\langle\mathbf{A},\mathbf{B}\rangle=\|\mathbf{A}\|_{F}^{2}+\|\mathbf{B}\|_{F}^{2}-\|\mathbf{A}-\mathbf{B}\|_{F}^{2}$. Applying this equality to the right hand side above gives

By Lemma B.3, for $l\in[L]$ we have $\|\mathbf{W}_{l}^{(i)}-\mathbf{W}_{l}^{(i+1)}\|_{F}\leq\eta\|\nabla_{\mathbf{W}_{l}}L_{i}(\mathbf{W}^{(i)})\|_{F}\leq\mathcal{O}(\eta\sqrt{m})$. Therefore

Telescoping over $i=1,\ldots,n$, we obtain

where in the first inequality we simply remove the term $-\|\mathbf{W}_{l}^{(n+1)}-\mathbf{W}_{l}^{*}\|_{F}^{2}/(2\eta)$ to obtain an upper bound, and the second inequality follows by the assumption that $\mathbf{W}^{*}\in\mathcal{B}(\mathbf{W}^{(1)},Rm^{-1/2})$. Plugging in the parameter choice $\eta=\nu\epsilon/(Lm)$, $n=L^{2}R^{2}/(2\nu\epsilon^{2})$ for some small enough absolute constant $\nu$ gives

B.4 Proof of Lemma 4.4

Here we prove Lemma 4.4. The proof essentially follows by standard Gaussian tail bound and a bound on the length of last hidden layer output vector.

By Lemma 4.1 in Allen-Zhu et al. (2019b), with probability at least $1-\mathcal{O}(nL)\cdot\exp[-\Omega(m/L)]>1-\delta/2$ over the randomness of $\mathbf{W}^{(1)}_{1},\ldots,\mathbf{W}^{(1)}_{L-1}$, $\|\mathbf{h}_{i,L-1}^{(0)}\|_{2}\in[1/2,3/2]$ for all $i\in[n]$. Condition on $\mathbf{W}^{(1)}_{1},\ldots,\mathbf{W}^{(1)}_{L-1}$, $f_{\mathbf{W}^{(1)}}(\mathbf{x}_{i})=\sqrt{m}\cdot\mathbf{W}^{(1)}_{L}\mathbf{h}_{i,L-1}$ is a Gaussian random variable with variance $\|\mathbf{h}_{i,L-1}\|_{2}^{2}$. Therefore by standard Gaussian tail bound and union bound, with probability at least $1-\delta$, $|f_{\mathbf{W}^{(1)}}(\bm{x}_{i})|\leq\mathcal{O}(\sqrt{\log(n/\delta)})$ for all $i\in[n]$. ∎

Appendix C Proofs of Results in Section A

In this section we provide the proofs of Corollary A.2 and Lemma A.3.

The following lemma is a simplified version of Lemma C.2 in Cao and Gu (2020). Since the proof is almost the same as the proof of Lemma C.2 in Cao and Gu (2020), except replacing the $\epsilon$-net argument with a simple union bound over $n$ training examples, we omit the proof detail here.

By Lemma C.1, with probability at least $1-\delta$, there exists $\bm{\alpha}_{L-1}\in S^{m-1}$ such that $y_{i}\cdot\langle\bm{\alpha}_{L-1},\mathbf{h}_{i,L-1}\rangle\geq 2^{-L}\gamma$ for all $i\in[n]$. Therefore, setting $R=(B+B^{\prime})\cdot 2^{L}\gamma^{-1}=\widetilde{\mathcal{O}}(2^{L}\gamma^{-1})$, we have

Moreover, $f^{*}(\cdot):=f_{\mathbf{W}^{(1)}}(\cdot)+\langle\nabla_{\mathbf{W}}f_{\mathbf{W}^{(1)}}(\cdot),\mathbf{W}\rangle$ satisfies $f^{*}\in\mathcal{F}(\mathbf{W}^{(1)},R)$, and

C.2 Proof of Lemma A.3

Here we give the proof of Lemma A.3. It is based on a simple construction.

For any $f(x)=\mathbf{W}_{2}\sigma(\mathbf{W}_{1}^{(1)}\mathbf{x})$ with $\|\mathbf{W}_{2}\|_{F}\leq Cm^{-1/2}$, by the assumption that $\|\mathbf{W}_{2}^{(1)}\|_{F}\leq Km^{-1/2}$ for some $K=\widetilde{\mathcal{O}}(1)$, we have $\mathbf{W}_{2}^{\prime}:=\mathbf{W}_{2}-\mathbf{W}_{2}^{(1)}$ satisfies $\|\mathbf{W}_{2}^{\prime}\|_{F}\leq(C+K)\cdot m^{-1/2}$. Therefore

Appendix D Proofs of Lemmas in Section B

In this section we give the proofs of lemma B.1, Lemma B.2 and Lemma B.3 in Section B.

By Lemma 4.1 in Allen-Zhu et al. (2019b), with probability at least $1-\mathcal{O}(nL)\cdot\exp[-\Omega(m/L)]$, $\|\mathbf{h}_{i,l}^{(1)}\|_{2}\in[3/4,5/4]$ for all $i\in[n]$ and $l\in[L-1]$. Moreover, by Lemma 5.2 in Allen-Zhu et al. (2019b) and the $1$-Lipschitz continuity of $\sigma(\cdot)$, with probability at least $1-\mathcal{O}(nL)\cdot\exp[-\Omega(m\omega^{2/3}L)]$, $\|\mathbf{h}_{i,l}-\mathbf{h}_{i,l}^{(1)}\|_{2}\leq\mathcal{O}(\omega L^{5/2}\sqrt{\log(m)})$. Therefore by the assumption that $\omega\leq\mathcal{O}(L^{-9/2}[\log(m)]^{-3})$, we have $\|\mathbf{h}_{i,l}\|_{2}\in[1/2,3/2]$ for all $i\in[n]$ and $l\in[L-1]$. ∎

D.2 Proof of Lemma B.2

We first introduce the following lemma characterizing the activation changes between networks with two close enough parameter sets $\mathbf{W}$ and $\mathbf{W}^{\prime}$. This lemma directly follows by Lemma 8.2 in Allen-Zhu et al. (2019b) and triangle inequality.

If $\omega\leq\mathcal{O}(L^{-9/2}[\log(m)]^{-3/2})$, then with probability at least $1-\mathcal{O}(nL)\cdot\exp[-\Omega(m\omega^{2/3}L)]$,

for all $\mathbf{W},\mathbf{W}^{\prime}\in\mathcal{B}(\mathbf{W}^{(1)},\omega)$, $i\in[n]$ and $l\in[L-1]$.

We first prove (i) and (iii), and then use (iii) to prove (ii).

By Lemma D.1, with probability at least $1-\mathcal{O}(nL)\cdot\exp(-\Omega(L\omega^{2/3}m))$, $\|\mathbf{D}_{i,l}-\mathbf{D}_{i,l}^{(1)}\|_{0}\leq\mathcal{O}(L\omega^{2/3}m)$ for all $i\in[n]$ and $l\in[L-1]$. Therefore we have $\|\mathbf{D}_{i,r}+\mathbf{D}_{i,r}^{\prime\prime}-\mathbf{D}_{i,l}^{(1)}\|_{0}\leq\mathcal{O}(L\omega^{2/3}m)$ for all $i\in[n]$ and $l\in[L-1]$. Therefore by Lemma 5.6 in Allen-Zhu et al. (2019b), with probability at least $1-\mathcal{O}(nL^{2})\cdot\exp[-\Omega(m\omega^{2/3}L)]$ we have \big{\|}\prod_{r=l_{1}}^{l_{2}}(\mathbf{D}_{i,r}+\mathbf{D}_{i,r}^{\prime\prime})\mathbf{W}_{r}\big{\|}_{2}\leq\mathcal{O}(\sqrt{L}). This completes the proof of (i) in Lemma B.2.

Similarly, to prove (iii), applying Lemma D.1 to $\mathbf{W}^{\prime}$ gives that with probability at least $1-\mathcal{O}(nL)\cdot\exp(-\Omega(L\omega^{2/3}m))$, $\|\mathbf{D}_{i,l}^{\prime}+\mathbf{D}_{i,r}^{\prime\prime}-\mathbf{D}_{i,l}^{(1)}\|_{0}\leq\mathcal{O}(L\omega^{2/3}m)$ for all $i\in[n]$ and $l\in[L-1]$. Now by Lemma 5.7 in Allen-Zhu et al. (2019b)Note that $\sqrt{m}\cdot\mathbf{W}_{L}^{(1)}$ is a random vector following the Gaussian distribution $N(\mathbf{0},\mathbf{I})$, which matches the distribution of the last layer parameters in Allen-Zhu et al. (2019b) for the binary classification case, where the output dimension of the network is $1$. with $s=\mathcal{O}(m\omega^{2/3}L)$ to $\mathbf{W}$ and $\mathbf{W}^{\prime}$, we have

Combining equations (D.1), (D.2), (D.3), (D.4) and applying triangle inequality gives the desired final result (iii).

Applying (iii) and (b) in Lemma 4.4 in Allen-Zhu et al. (2019b), with probability at least $1-\mathcal{O}(nL)\cdot\exp[-\Omega(m/L)]$, we obtain

D.3 Proof of Lemma B.3

for all $\mathbf{W}\in\mathcal{B}(\mathbf{W}^{(1)},\omega)$ and $i\in[n]$. For $l\in[L-1]$, by direct calculation we have

Therefore by Lemma B.1 and (ii) in Lemma B.2, we have

Finally, for $\|\nabla_{\mathbf{W}_{l}}L_{i}(\mathbf{W}^{(i)})\|_{F}$ we have

Appendix E Experimental Results

In this section we provide numerical calculations of the generalization bounds given by Theorem 3.3 and Corollary 3.10 on the MNIST dataset (LeCun et al., 1998). The main goal of these calculations is to demonstrate that the bounds given in our results are informative and can provide practical insight.

We have done experiments of a five-layer fully connected NN on MNIST dataset (3 versus 8), and calculated the first terms in the bounds given by Theorem 3.3 and Corollary 3.10.

To demonstrate the scaling of the bound in Corollary 3.10, we calculate the value of $\sqrt{\mathbf{y}^{\top}(\bm{\Theta}^{(L)})^{-1}\mathbf{y}/n}$, where $\mathbf{y}$ is the true label vector with random flips. We plot $\sqrt{\mathbf{y}^{\top}(\bm{\Theta}^{(L)})^{-1}\mathbf{y}/n}$ in Figure 1 by varying the level of label noise, i.e., ratio of the labels that are flipped. Note that to simplify calculation, we do not consider the $\widetilde{\mathbf{y}}$ introduced in Corollary 3.10. Clearly, our calculation here gives an upper bound of the generalization bound in Corollary 3.10.

References