Generalization Guarantees for Neural Networks via Harnessing the Low-rank Structure of the Jacobian

Introduction

Deep neural networks (DNN) are ubiquitous in a growing number of domains ranging from computer vision to healthcare. State-of-the-art DNN models are typically overparameterized and contain more parameters than the size of the training dataset. It is well understood that in this overparameterized regime, DNNs are highly expressive and have the capacity to (over)fit arbitrary training datasets including pure noise . Mysteriously however neural network models trained via simple algorithms such as (stochastic) gradient descent continue to predict well or generalize on yet unseen test data. In this paper we wish to take a step towards demystifying this phenomenon and help explain why neural nets can overfit to noise yet have the ability to generalize when real data sets are used for training. In particular we explore the generalization dynamics of neural nets trained via gradient descent. Using the Jacobian mapping associated to the neural network we characterize directions where learning is fast and generalizable versus directions where learning is slow and leads to overfitting. The main contributions of this work are as follows.

$\bullet$ Leveraging dataset structure: We develop new optimization and generalization results that can harness the low-rank representation of semantically meaningful datasets via the Jacobian mapping of the neural net. This sheds light as to why training and generalization is easier using datasets where the features and labels are semantically linked versus others where there is no meaningful relationship between the features and labels (even when the same network is used for training). $\bullet$ Bias–variance tradeoffs: We develop a bias–variance theory based on the Jacobian which decouples the learning process into information and nuisance spaces. We show that gradient descent almost perfectly interpolates the data over the information space (incurring only a small bias). In contrast, optimization over the nuisance space is slow and results in overfitting due to higher variance. $\bullet$ Network size vs prediction bias: We obtain data-dependent tradeoffs between the network size and prediction bias. Specifically, we show that larger networks result in smaller prediction bias, but small networks can still generalize well, especially when the dataset is sufficiently structured, but typically incur a larger bias. This is in stark contrast to recent literature on optimization and generalization of neural networks where guarantees only hold for very wide networks with the width of the network growing inversely proportional to the distance between the input samples or class margins or related notions. See Section 3.4 for further detail. $\bullet$ Pretrained models: In our framework we do not require the initialization to be random and our results continue to apply even with arbitrary initialization. Therefore, our results may shed light on the generalization capabilities of networks initialized with pre-trained models such as those commonly used in meta/transfer learning.

2 Model and training

It will be convenient to concatenate the labels and prediction vectors as follows

Using this shorthand we can rewrite the loss (1.2) as

To optimize this loss starting from an initialization $\bm{W}_{0}$ we run gradient descent iterations of the form

with a step size $\eta$. In this paper we wish to explore the theoretical properties of the model found by such iterative updates with an emphasis on the generalization ability.

Components of a Jacobian-based theory of generalization

Now let us consider gradient descent iterations with a step size $\eta$ which take the form

To gain further insight into the generalization capabilities of the gradient descent iterations we shall consider an instance of this problem where the subspaces $\bm{U}$ and $\bm{V}$ are chosen uniformly at random, $\bm{\Sigma}=\bm{I}_{r}$ with $n=200$, $d=500$, $r=5$, and $\sigma_{x}=0.2,~{}\sigma_{y}=2$. In Figure 2(a) we plot the population loss evaluated at different iterations. We observe an interesting phenomenon, in the first few iterations the test error goes down quickly but it then slowly increases. To better understand this behavior we decompose the population loss into two parts by tracking the projection of the misfit ${\bm{X}}\bm{w}-\bm{y}$ on the column space of the uncorrupted portion of the input data $(\bm{U})$ and its complement. That is,

To help demystify this behavior note that using the gradient descent updates from (2.2) the update in terms of the misfit/residual $\bm{r}_{\tau}={\bm{X}}\bm{w}_{\tau}-\bm{y}$ takes the form

2 Information and nuisance spaces of the Jacobian

We define the information and nuisance spaces associated with $\bm{J}$ as $\mathcal{I}:=\text{span}(\{{\bm{u}}_{s}\}_{s=1}^{r})$ and $\mathcal{N}:=\text{span}(\{{\bm{u}}_{s}\}_{s=r+1}^{{K}n})$. We also define the truncated Jacobian

which is the part of the reference Jacobian that acts on the information space $\mathcal{I}$.

Main results

To explore the generalization of randomly initialized networks, we utilize the neural tangent kernel.

with $\Gamma\geq 1$. We run gradient descent iterations of the form (1.5) with a learning rate $\eta\leq\frac{1}{\nu^{2}B^{2}\left\|{\bm{X}}\right\|^{2}}$. Then, after $T=\frac{\Gamma K}{\eta\nu^{2}\alpha_{0}^{2}}$ iterations, classification error $\text{Err}_{{\cal{D}}}(\bm{W}_{T})$ is upper bounded by

holds with probability at least $1-(2{K})^{-100}-\delta$.

This theorem shows that even networks of moderate width can achieve a small generalization error if (1) the data has low-dimensional representation i.e. the kernel is approximately low-rank and (2) the inputs and labels are semantically-linked i.e. the label vector $\bm{y}$ mostly lies on the information space.

holds as long as ${\bm{{\Sigma}}}({\bm{X}})$ is invertible and the width of the network obeys

We note that in this special case our results improve upon the required width in recent literature Based on our understanding requires the number of hidden units to be at least on the order of $k\gtrsim\frac{n^{8}}{\lambda^{6}}$. Note that using the fact that $\left\|{\bm{X}}\right\|\leq\sqrt{n}$ our result reduces the dependence on width by a factor of at least $\frac{n^{4}}{\lambda^{2}}$. We note that $\left\|{\bm{X}}\right\|$ often scales with $\sqrt{\frac{n}{d}}$ so that the improvement in width is even more pronounced in typical instances. that focuses on ${K}=1$ and a conclusion of the form (3.4). However, as we demonstrate in our numerical experiments in practice $\lambda$ can be rather small or even zero (e.g. see the toy model in Section 3.3) so that requirements of the form (3.5) may require unrealistically (or even infinitely) wide networks. In contrast, as discussed above by harnessing the low-rank structure of the Jacobian our results show that neural networks generalize well as soon as the width grows at most logarithmically in the size of the training data (even when $\lambda=0$).

$\bullet$ Small width is sufficient for generalization: Based on our simulations the M-NTK (or more specifically Jacobian at random initialization) indeed has low-rank structure with a few large eigenvalues and many smaller ones. As a result a typical scaling of the cut-off $\alpha_{0}$ is so that $\bar{\alpha}$ scales like a constant. In that case our result states that as soon as the number of hidden nodes are moderately large (e.g. logarithmic in $n$) then good generalization can be achieved. Specifically we can achieve good generalization by using width on the order of $\log n$ and picking small values for $\zeta$ and $\bar{\alpha}$ and large values for $\Gamma$.

$\bullet$ Network size–Bias tradeoff: Based on the requirement (6.86) if the network is large (in terms of # of hidden units $k$), we can choose a small cut-off $\alpha_{0}$. This in turn allows us to enlargen the information space and reduce the training bias. In summary, as network capacity grows, we can gradually interpolate finer detail and reduce bias. On the other hand, choosing a properly large $\alpha_{0}$, we can obtain good bounds for even small network sizes $k$ as long as the portion of the labels that fall on the nuisance space is small. This is in stark contrast to related works where network size grows inversely proportional to the distance between the input samples or other notions of margin.

$\bullet$ Fast convergence: We note that by setting learning rate to $\eta=\frac{1}{\nu^{2}B^{2}\left\|{\bm{X}}\right\|^{2}}$, the number of gradient iterations is upper bounded by $\frac{\Gamma}{\bar{\alpha}^{2}}$. Hence, the training speed is dictated by and is inversely proportional to the the smallest singular value over the information space. Specifically, when the Jacobian is sufficiently low-rank so that we can pick $\bar{\alpha}$ to be a constant, convergence on the information space is rather fast requiring only a constant number of iterations to converge to any fixed constant accuracy. See the proofs for further detail on the optimization dynamics of the training problem (e.g. results/proofs for linear convergence of the empirical loss).

2 Generalization guarantees with arbitrary initialization

Our next result provides generalization guarantees from an arbitrary initialization which applies to pre-trained networks (e.g. those that arise in transfer learning applications) as well as intermediate gradient iterates as the weights evolve. This result has a similar flavor to Theorem 3.2 with the key difference that the information and nuisance spaces are defined with respect to any arbitrary initial Jacobian. This shows that if a pre-trained modele.g. obtained by training with data in a related problem as is common in transfer learning. provides a better low-rank representation of the data in terms of its Jacobian, it is more likely to generalize well. Furthermore, given its deterministic nature the theorem can be applied at any iteration, implying that if the Jacobians of any of the iterates provides a better low-rank representation of the data then one can provide sharper generalization guarantees.

with $\Gamma\geq 1$ and tolerance level $\zeta$. Run gradient descent updates (1.5) with learning rate $\eta\leq\frac{1}{\nu^{2}B^{2}\left\|{\bm{X}}\right\|^{2}}$. Then, after $T=\frac{\Gamma}{\eta\alpha^{2}}$ iterations, with probability at least $1-\delta$, the generalization error obeys

As with the random initialization result, this theorem shows that as long as the initial residual is sufficiently correlated with the information space, then high accuracy can be achieved for neural networks with moderate width. As with its randomized counter part this result also allows us to study various tradeoffs between bias-variance and network size-bias. Crucially however this result does not rely on random initialization. The reason this is particularly important is two fold. First, in many scenarios neural networks are not initialized at random. For instance, in transfer learning the network is pre-trained via data from a different domain. Second, as we demonstrate in Section 4 as the iterates progress the Jacobian mapping seems to develop more favorable properties with the labels/initial residuals becoming more correlated with the information space of the Jacobian. As mentioned earlier, due its deterministic nature the theorem above applies in both of these scenarios. In particular, if a pre-trained model provides a better low-rank representation of the data in terms of its Jacobian, it is more likely to generalize well. Furthermore, given its deterministic nature the theorem can be applied at any iteration by setting $\bm{\theta}_{0}=\bm{\theta}_{\tau}$, implying that if the Jacobians of any of the iterates provides a better low-rank representation of the data then one can provide sharper generalization guarantees. Our numerical experiments demonstrate that the Jacobian of the neural network seems to adapt to the dataset over time with a more substantial amount of the labels lying on the information space. While we have not formally proven such an adaptation behavior in this paper, we hope to develop rigorous theory demonstrating this adaptation in our future work.Such a result when combined with our arbitrary initialization guarantee above can potentially provide significantly tighter generalization bounds. This is particularly important in light of a few recent literature suggesting a significant gap between generalization capabilities of kernel methods/linearized neural nets when compared with neural nets operating beyond a linear or NTK learning regime (e.g. mean field regime). As a result we view our deterministic result as a first step towards moving beyond the NTK regime.

3 Case Study: Gaussian mixture model

To illustrate a concrete example, we consider a distribution based on a Gaussian mixture model consisting of $K$ classes where each class consists of $C$ clusters.

This distribution is an ideal candidate to demonstrate why the Jacobian of the network exhibits low-rank or bimodal structure. Let us consider the extreme case $\sigma=0$ where we have a discrete input distribution over the cluster centers. In this scenario, we can show that the multi-class Jacobian matrix is at most rank

as there are (i) only ${K}C$ distinct input vectors and (ii) ${K}$ output nodes. We can thus set the information space to be the top ${K}^{2}C$ eigenvectors of the multiclass kernel matrix ${\bm{{\Sigma}}}({\bm{X}})$. As formalized in the appendix, it can be shown that

The singular values of the information space grow proportionally with $n/KC$.

The concatenated label vector $\bm{y}$ perfectly lies on the information space.

In Figure 4 we numerically verify that the approximate rank and singular values of the Jacobian indeed scale as above even when $\sigma>0$. The following informal theorem leverages these observations to establish a generalization bound for this mixture model. This informal statement is for exposition purposes. See Theorem A.3 in Appendix A for a more detailed result capturing the exact dependencies (e.g. $\zeta,B,\log n$). In this theorem we use $\gtrsim$ to denote inequality up to constant/logarithmic factors.

Then, after running gradient descent for $T=\frac{2\Gamma K^{2}C}{\lambda_{\bm{M}}}$ iterations, the model obeys

We note that $\lambda_{\bm{M}}$ captures how diverse the cluster centers are. In this sense $\lambda_{\bm{M}}>0$ intuitively means that neural network, specifically the neural tangent kernel, is sufficiently expressive to interpolate the cluster centers. In fact when the cluster centers are in generic position $\lambda_{\bm{M}}$ scales like a constant . This theorem focuses on the regime where the noise level $\sigma$ is small. In this case we show that one can achieve good generalization as soon as the number of data points scale with the square of the number classes times the total number of cluster (i.e. $n\gtrsim K^{2}C$) which is the effective rank of the M-NTK matrix. We note that this result follows from our main result with random initialization by setting the cutoff level at $\alpha_{0}^{2}\sim\frac{\lambda_{{\bm{M}}}n}{KC}$. This demonstrates that in this model $\bar{\alpha}$ does indeed scale as a constant. Finally, the required network width is independent of $n$ and only depends on $K$ and $C$ specifically we require $k\gtrsim K^{8}C^{4}$. This is in stark contrast with in the binary case. To the best of understanding requires $k\gtrsim\frac{n^{8}}{\lambda_{{\bm{X}}}^{6}}$ which depends on $n$ (in lieu of $K$ and $C$) and the minimum eigenvalue $\lambda_{{\bm{X}}}$ of the NTK matrix ${\bm{{\Sigma}}}({\bm{X}})$ (rather than $\lambda_{{\bm{M}}}$). Furthermore, in this case as $\sigma\rightarrow 0$, ${\bm{{\Sigma}}}({\bm{X}})$ becomes rank deficient and $\lambda_{{\bm{X}}}\rightarrow 0$ so that the required width of grows to infinity.

4 Prior Art

Neural networks have impressive generalization abilities even when they are trained with more parameters than the size of the dataset . Thus, optimization and generalization properties of neural networks have been the topic of many recent works . Below we discuss related work on classical learning theory as well as optimization and implicit bias.

Statistical learning theory: Statistical properties of neural networks have been studied since 1990’s . With the success of deep networks, there is a renewed interest in understanding capacity of the neural networks under different norm constraints or network architectures . established tight sample complexity results for deep networks based on the product of appropriately normalized spectral norms. See also for improvements via leveraging various properties of the inter-layer Jacobian and for results with convolutional networks. Related, leverages compression techniques for constructing tighter bounds. jointly studies statistical learning and adversarial robustness. These interesting results, provide generalization guarantees for the optimal solution to the empirical risk minimizer. In contrast, we focus on analyzing the generalization dynamics of gradient descent iterations.

Properties of gradient descent: There is a growing understanding that solutions found by first-order methods such as gradient descent have often favorable properties. Generalization properties of stochastic gradient descent is extensively studied empirically . For linearly separable datasets, show that first-order methods find solutions that generalize well without an explicit regularization for logistic regression. An interesting line of work establish connection between kernel methods and neural networks and study the generalization abilities of kernel methods when the model interpolates the training data . relate the distribution of the network weights to Wasserstein gradient flows using mean field analysis. This literature is focused on asymptotic characterizations rather than finite-size networks.

Global convergence and generalization of neural nets: Closer to this work, recent literature provides generalization bounds for overparameterized networks trained via gradient descent. Also see for interesting visualization of the optimization and generalization landscape. Similar to Theorem 3.2, uses the NTK to provide generalization gurantees. leverages low-rank Jacobian structure to establish robustness to label noise. These works build on global convergence results of randomly initialized neural networks which study the gradient descent trajectory via comparisons to a a linearized Neural Tangent Kernel (NTK) learning problem. These results however typically require unrealistically wide networks for optimization where the width grows poly-inversely proportional to the distance between the input samples. Example distance measures are class margin for logistic loss and minimum eigenvalue of the kernel matrix for least-squares. Our work circumvents this by allowing a capacity-dependent interpolation. We prove that even rather small networks (e.g. of constant width) can interpolate the data over a low-dimensional information space without making restrictive assumptions on the input. This approach also leads to faster convergence rates. In terms of generalization, our work has three distinguishing features: (a) bias-variance tradeoffs by identifying information/nuisance spaces, (b) no margin/distance/minimum eigenvalue assumptions on data, (c) the bounds apply to multiclass classification as well as pre-trained networks (Theorem 3.3).

Numerical experiments

Experimental setup. We present experiments supporting our theoretical findings on the CIFAR-10 dataset, which consists of $50k$ training images and $10k$ test images in $10$ classes. For our experiments, we reduced the number of classes to $3$ (automobile, airplane, bird) and subsampled the training data such that each class is represented by $3333$ images ($9999$ in total). This is due to the fact that calculating the full spectrum of the Jacobian matrix over the entire data set is computationally intensiveWe plan to perform more comprehensive set of experiments by calculating the Jacobian spectrum in a distributed manner.. For testing, we used all examples of the $3$ classes ($3000$ in total). In all of our experiments we set the information space to be the span of the top 50 singular vectors (out of total dimension of ${K}n\approx 30000$).

We demonstrate our results on ResNet20, a state-of-the-art architecture with a fairly low test error on this dataset ($8.75\%$ test error reported on $10$ classes) and relatively few parameters ($0.27M$). In order to be consistent with our theoretical formulation we made the following modifications to the default architecture: (1) we turned off batch normalization and (2) we did not pass the network output through a soft-max function. We trained the network using a least-squares loss with SGD with batch size $128$ and standard data augmentation (e.g. random crop and flip). We set the initial learning rate to $0.01$ and adjusted the learning rate schedule and number of epochs depending on the particular experiment so as to achieve a good fit to the training data quickly. The figures in this section depict the minimum error over a window consisting of the last 10 epochs for visual clarity. We also conducted two sets of experiments to illustrate the results on uncorrupted and corrupted data.

Experiments without label corruption. First, we present experiments on the original training data described above with no label corruption. We train the network to fit to the training data by using $400$ epochs and decreasing the learning rate at $260$ and $360$ epochs by a factor of $10$.

In Figure 5 we plot the histogram of the eigenvalues of the Jacobian calculated on the training data at initialization and after training. This figure clearly demonstrates that the Jacobian has low-rank structure as there are tens of large singular values with the remaining majority of the spectrum consisting of small singular values. This observation serves as a natural basis for decomposition of the label space into the information space $\mathcal{I}$ (large singular values, low-dimensional) and nuisance space $\mathcal{N}$ (small singular values, high-dimensional).

We also track the projection of the residual $\bm{r}_{\tau}$ on the information and nuisance subspaces throughout training on both training and test data and depict the results in Figures 6(a) and 6(b). In agreement with our theory, these plots show that learning on $\mathcal{I}$ is fast and the residual energy decreases rapidly on this space. On the other hand, residual energy on $\mathcal{N}$ goes down rather slowly and the decrease in total residual energy is overwhelmingly governed by $\mathcal{I}$, suggesting that most information relevant to learning lies in this space. We also plot the training and test error in Figure 6(c). We observe that as learning progresses, the residual on both spaces decrease in tandem with training and test error.

Experiments with 50% label corruption. In our next series of experiments we study the effect of corruption. Specifically, we corrupt $50\%$ of the labels by randomly picking a label from a (strictly) different class. We train the network for $800$ epochs and divide the learning rate by $10$ at $700$ epochs to fit to the training data.

Similar to the uncorrupted case, we track the projection of the residual $\bm{r}_{\tau}$ on the information and nuisance spaces throughout training on both training and test data and depict the results in Figures 8(a) and 8(b). We also track the train and test misclassification error in Figure 8(c). From Figure 8(c) it is evident that while the training error steadily decreases, test error exhibits a very different behavior from the uncorrupted experiment. In the first phase, test error drops rapidly as the network learns from information contained in the uncorrupted data, accompanied by a corresponding decrease in residual energy on the information subspace on the training data (Figure 8(a)). The lowest test error is observed at $100$ epochs after which a steady increase follows. In the second phase, the network overfits to the corrupted data resulting in larger test error on the uncorrupted test data (Figure 8(b)). More importantly, the increase of the test error is due to the nuisance space as the error over information space is stable while it increases over the nuisance space. In particular the residual on $\mathcal{N}$ slowly increases while residual on $\mathcal{I}$ drops sharply creating a dip in both test error and total residual energy at approximately $100$ epochs. This phenomenon closely resembles the population loss decomposition of the linear model discussed in Section 2.1 (see Figure 2), where we observe a dip in total test error caused by an increasing component along the nuisance space and a simultaneously decreasing component along information space.

In Table 3 we again depict the fraction of the energy of the labels and the initial residual that lies on the information/nuisance spaces. The Jacobian continues to adapt to the labels/initial residual even in the presence of label corruption, albeit to a smaller degree. We note that due to corruption, labels are less correlated with the information space of the Jacobian and the fraction of the energy on the nuisance space is higher which results in worse generalization (as also predicted by our theory).

In order to demonstrate the connection between generalization error and information/nuisance spaces of the Jacobian, we repeat the experiment with $25\%$, $75\%$ and $100\%$ label corruption and depict the results after $800$ epochs in Figure 9. As expected, the test error increases with the corruption level. Furthermore, the corrupted labels become less correlated with the information space, with more of the label energy falling onto the nuisance space. This is consistent with our theory which predicts worse generalization in this case.

Technical approach and General Theory

To continue let us first aggregate the predictions and labels into larger vectors based on class. In particular define

Using the latter we can rewrite the optimization problem (5.1) into the more compact form

To solve this problem we run gradient descent iterations with a learning rate $\eta$ starting from an initial point $\bm{\theta}_{0}$. These iterations take the form

As mentioned earlier due to the form of the gradient the convergence/generalization of gradient descent naturally depends on the spectral properties of the Jacobian. To capture these spectral properties we will use a reference Jacobian $\bm{J}$ (formally defined below) that is close to the Jacobian at initialization $\mathcal{J}(\bm{\theta}_{0})$.

Thus starting from $\widetilde{\bm{\theta}}_{0}=\overline{\bm{\theta}}_{0}$ the iterates $\widetilde{\bm{\theta}}_{\tau}$ on the linearized problem take the form

The iterates based on the linearized problem will provide a useful reference to keep track of the evolution of the original iterates (5.4). Specifically we study the evolution of misfit/residuals associated with the two problems

To better understand the dynamics of convergence of the linearized iterates next we define two subspaces associated with the reference Jacobian and its spectrum.

Let $\bm{J}$ denote the reference Jacobian per Definition 5.1 with eigenvalue decomposition $\bm{J}=\bm{U}$diag$(\bm{\lambda})\bm{V}^{T}$ per (5.5). For a spectrum cutoff $\alpha$ obeying $0\leq\alpha\leq\lambda_{1}$ let $r(\alpha)$ denote the index of the smallest singular value above the threshold $\alpha$, that is,

We define the information and nuisance subspaces associated with $\bm{J}$ as $\mathcal{I}:=\text{span}(\{{\bm{u}}_{s}\}_{s=1}^{r})$ and $\mathcal{N}:=\text{span}(\{{\bm{u}}_{s}\}_{s=r+1}^{{K}n})$. We also define the truncated reference Jacobian

which is the part of the reference Jacobian that acts on the information subspace $\mathcal{I}$.

We will show rigorously that the information and nuisance subspaces associated with the reference Jacobian dictate the directions where learning is fast and generalizable versus the directions where learning is slow and overfitting occurs. Before we make this precise we list two assumptions that will be utilized in our result.

With these assumptions in place we are now ready to discuss our meta theorem that demonstrates that the misfit/residuals associated to the original and linearized iterates do in fact track each other rather closely.

We run gradient descent iterations of the form $\bm{\theta}_{\tau+1}=\bm{\theta}_{\tau}-\eta\nabla\mathcal{L}(\bm{\theta}_{\tau})$ and $\widetilde{\bm{\theta}}_{\tau+1}=\widetilde{\bm{\theta}}_{\tau}-\eta\nabla\mathcal{L}_{lin}(\widetilde{\bm{\theta}}_{\tau})$ on the original and linearized problems starting from $\bm{\theta}_{0}$ with step size $\eta$ obeying $\eta\leq 1/\beta^{2}$. Then for all iterates $\tau$ obeying $0\leq\tau\leq T:=\frac{\Gamma}{\eta\alpha^{2}}$ the iterates of the original ($\bm{\theta}_{\tau}$) and linearized ($\widetilde{\bm{\theta}}_{\tau}$) problems and the corresponding residuals $\bm{r}_{\tau}:=f(\bm{\theta}_{\tau})-\bm{y}$ and $\bm{\widetilde{r}}_{\tau}:=f_{\text{lin}}(\widetilde{\bm{\theta}}_{\tau})-\bm{y}$ closely track each other. That is,

Furthermore, for all iterates $\tau$ obeying $0\leq\tau\leq T:=\frac{\Gamma}{\eta\alpha^{2}}$

Proofs

In this section we prove our result for general nonlinearities. We begin with a few notations and definitions and preliminary lemmas in Section 6.1.1. Next in Section 6.1.2 we prove some key lemmas regarding the evolution of the linearized residuals $\bm{\widetilde{r}}_{\tau}$. In Section 6.3 we establish some key Rademacher complexity results used in our generalization bounds. Finally, in Section 6.1.3 we use these results to complete the proof of Theorem 5.3.

to denote the basis matrices for the information and nuisance subspaces from Definition 5.2. Similarly, we define the information and nuisance spectrum as

Consider Definition 5.2 and let $\Gamma>0$ be a positive scalar. Associated with the initial residual $\bm{r}_{0}=f(\bm{\theta}_{0})-\bm{y}$ and the information/nuisance subspaces of the reference Jacobian $\bm{J}$ (with a cut-off level $\alpha$) we define the $(\alpha,\Gamma)$ early stopping value as

We also define the early stopping distance as

The goal of early stopping value/distance is understanding the behavior of the algorithm at a particular stopping time that depends on $\Gamma$ and the spectrum cutoff $\alpha$. In particular, as we will see later on the early stopping distance characterizes the distance from initialization at an appropriate early stopping time. We continue by stating and proving a few simple lemmas. The first Lemma provides upper/lower bounds on the early stopping value.

The early stopping value ${\cal{B}}_{\alpha,\Gamma}$ from Definition 6.1 obeys

Proof To prove the upper bound we use the fact that $\alpha\leq\lambda_{s}$ for $s\leq r$ and $\alpha\geq\lambda_{s}$ for $s\geq r$ to conclude that

To prove the lower bound, we use the facts that $\alpha^{2}/\lambda_{s}^{2}\geq\alpha^{2}/\lambda_{1}^{2}$ to conclude that

It is of course well known that the mapping $\left(\bm{I}-\eta{\bm{A}}{\bm{A}}^{T}\right)$ is a contraction for sufficiently small values of $\eta$. The next lemma shows that if we replace one of the ${\bm{A}}$ matrices with a matrix ${{\bm{B}}}$ which is close to ${\bm{A}}$ the resulting matrix $\left(\bm{I}-\eta{\bm{A}}{{\bm{B}}}^{T}\right)$, while may not be contractive, is not too expansive.

Proof Note that using $\eta\leq 1/\beta^{2}$ and $\left\|{{\bm{B}}}-{\bm{A}}\right\|\leq\varepsilon$ we conclude that

The next lemma shows that if two PSD matrices are close to each other then an appropriate square root of these matrices will also be close.

Furthermore, using the fact that the eigenvalues of ${\bm{A}}_{+}$ and ${{\bm{B}}}_{+}$ are just shifted versions of the eigenvalues of ${\bm{A}}$ and ${{\bm{B}}}$ by $\alpha^{2}/4$ we can conclude that

Combining the latter two inequalities with the assumption that $\left\|{\bm{A}}-{{\bm{B}}}\right\|\leq\alpha^{2}$ we conclude that

Then clearly ${\bm{Y}}{\bm{Y}}^{T}={{\bm{B}}}$. Furthermore, we have

Combining the latter with (6.1.1) completes the proof.

1.2 Key lemmas for general nonlinearities

We shall first characterize the evolution of the linearized parameter $\widetilde{\bm{\theta}}_{\tau}$ and residual $\bm{\widetilde{r}}_{\tau}$ vectors from (5.9) in the following lemma.

The linearized residual vector $\bm{\widetilde{r}}_{\tau}$ can be written in the form

Furthermore, assuming $\eta\leq 1/\lambda_{1}^{2}$ the linear updates $\widetilde{\bm{\theta}}_{\tau}$ obey

Proof Using the fact that $\bm{J}\bm{J}^{T}=\bm{U}\bm{\Lambda}^{2}\bm{U}^{T}$ we have

Using the latter combined with (5.9) we thus have

We now turn our attention to proving (6.6) by tracking the representation of $\widetilde{\bm{\theta}}_{\tau}$ in terms of the right singular vectors of $\bm{J}$. To do this note that using (6.5) we have

Using the latter together with the gradient update on the linearized problem we have

Noting that for $\eta\leq 1/\lambda_{1}^{2}\leq 1/\lambda_{s}^{2}$ we have $1-\eta\lambda_{s}^{2}\geq 0$, the latter identity implies that

Furthermore, using the fact that $1-\eta\lambda_{s}^{2}\leq 1$ we have

Combining (6.7) for $1\leq s\leq r$ and (6.8) for $s>r$ we have

For future use we also state a simple corollary of the above Lemma below.

Consider the setting and assumptions of Lemma 6.5. Then, after $\tau$ iterations we have

Furthermore, after $T=\frac{\Gamma}{\eta\alpha^{2}}$ iterations we have

with ${\cal{B}}_{\alpha,\Gamma}$ given by (6.1) per Definition 6.2.

Proof To prove the first bound on the residual ((6.9))note that using (6.5) we have

Thus, using the fact that for $s\leq r$ we have $\lambda_{s}\geq\alpha$ we have $(1-\eta\lambda_{s}^{2})^{\tau}\leq(1-\eta\alpha^{2})^{\tau}$ and for $s>r$ we have $(1-\eta\lambda_{s}^{2})^{\tau}\leq 1$, we can conclude that

Combining these with the triangular inequality we have

Assume Assumptions 1 and 2 hold and $\bm{\theta}_{\tau}$ and $\bm{\theta}_{\tau+1}$ are within an $R$ neighborhood of $\bm{\theta}_{0}$, that is,

Then with a learning rate obeying $\eta\leq 1/\beta^{2}$, the deviation in the residuals of the original and linearized problems $\bm{e}_{\tau+1}=\bm{r}_{\tau+1}-\bm{\widetilde{r}}_{\tau+1}$ obey

Proof For simplicity, denote ${{\bm{B}}}_{1}={\cal{J}}(\bm{\theta}_{\tau+1},\bm{\theta}_{\tau})$, ${{\bm{B}}}_{2}={\cal{J}}(\bm{\theta}_{\tau})$, ${\bm{A}}={\cal{J}}(\bm{\theta}_{0})$ where

We can write the predictions due to $\bm{\theta}_{\tau+1}$ as

Similarly, for linearized problem we have $\bm{\widetilde{r}}_{\tau+1}=({\bm{I}}-\eta\bm{J}\bm{J}^{T})\bm{\widetilde{r}}_{\tau}$. Thus,

We proceed by bounding each of these two terms. For the first term, we apply Lemma 6.3 with ${\bm{A}}={{\bm{B}}}_{1}$ and ${{\bm{B}}}={{\bm{B}}}_{2}$ and use $\|{{\bm{B}}}_{1}-{{\bm{B}}}_{2}\|\leq\varepsilon$ to conclude that

Next we turn our attention to bounding the second term. To this aim note that

In the last inequality we use the fact that per Assumption 2 we have $\|{{\bm{B}}}_{1}-{\bm{A}}\|\leq\varepsilon/2$ and $\|{{\bm{B}}}_{2}-{\bm{A}}\|\leq\varepsilon/2$ as well as the fact that per Definition 5.1 $\|{\bm{A}}{\bm{A}}^{T}-\bm{J}\bm{J}^{T}\|\leq\varepsilon_{0}^{2}$. Plugging (6.13) and (6.1.2) in (6.1.2) completes the proof.

Consider positive scalars $\Gamma,\alpha,\varepsilon,\eta>0$. Also assume $\eta\leq 1/\alpha^{2}$ and $\alpha\geq\sqrt{2\Gamma}\varepsilon$ and set $T=\frac{\Gamma}{\eta\alpha^{2}}$. Assume the scalar sequences $e_{\tau}$ (with $e_{0}=0$) and $\widetilde{r}_{\tau}$ obey the following identities

for all $0\leq\tau\leq T$ and non-negative values $\rho_{-},\rho_{+}\geq 0$. Then, for all $0\leq\tau\leq T$,

Proof We shall prove the result inductively. Suppose (6.16) holds for all $t\leq\tau-1$. Consequently, we have

Summing up both sides of (6.17) for $0\leq t\leq\tau-1$ we conclude that

where in the last inequality we used the fact that $\alpha^{2}\geq 2\Gamma\varepsilon^{2}$. This completes the proof of the induction step and the proof of the lemma.

1.3 Completing the proof of Theorem 5.3

With the key lemmas in place in this section we wish to complete the proof of Theorem 5.3. We will use induction to prove the result. Suppose the statement is true for some $\tau-1\leq T-1$. In particular, we assume the identities (5.12) and (5.13) hold for all $0\leq t\leq\tau-1$. We aim to prove these identities continue to hold for iteration $\tau$. We will prove this result in multiple steps.

Here, (a) and (b) follow from a simple application of the triangular inequality, (c) from the fact that $\left\|{\cal{J}}(\bm{\theta}_{\tau-1})-\bm{J}\right\|\leq\left\|{\cal{J}}(\bm{\theta}_{\tau-1})-{\cal{J}}(\bm{\theta}_{0})\right\|+\left\|{\cal{J}}(\bm{\theta}_{0})-\bm{J}\right\|\leq\varepsilon+\varepsilon_{0}$, (d) from combining the bounds in (5.11), (e) from the induction hypothesis that postulates (5.12) holds for iteration $\tau-1$, (f) from considering the SVD $\bm{J}=\bm{U}\bm{\Lambda}\bm{V}^{T}$ which implies that

(g) from the fact that $\eta\leq\frac{1}{\beta^{2}}$, and (h) from the fact that $\alpha\leq\beta$ and $\Gamma\geq 1$.

This combined with the induction assumption implies that

holds for all $t\leq\tau\leq T$. Furthermore, using Lemma 6.5 equation (6.9) for all $t\leq\tau\leq T$ we have

To proceed, we shall apply Lemma 6.8 with the following variable substitutions

We note that Lemma 6.8 is applicable since (i) $\eta\leq 1/\beta^{2}\leq 1/\alpha^{2}$, (ii) based on (5.11) we have $\frac{\alpha}{\varepsilon}\geq\frac{5\Gamma}{\delta}\frac{\beta^{2}}{\alpha^{2}}\geq\sqrt{2\Gamma}$, (iii) $\tau$ obeys $\tau\leq T=\frac{\Gamma}{\eta\alpha^{2}}$, and (iv) (6.15) holds based on (6.18) and (6.19). Thus using Lemma 6.8 we can conclude that

where in the last inequality we used (5.11). This completes the first part of (5.12) via induction.

Step III: Original and linearized parameters are close (second part of (5.12)). In this step we wish to show that the second part of (5.12) holds for iteration $\tau$. To do this we begin by noting that by the fact that $\bm{J}$ is a reference Jacobian we have $\|{\overline{\cal{J}}}(\bm{\theta}_{0})-\bm{J}\|\leq\varepsilon_{0}$ where ${\overline{\cal{J}}}$ augments ${\cal{J}}(\bm{\theta}_{0})$ by padding zero columns to match size of $\bm{J}$. Also by Assumption 2 we have $\|{\cal{J}}(\bm{\theta})-{\cal{J}}(\bm{\theta}_{0})\|\leq\frac{\varepsilon}{2}$. Combining the latter two via the triangular inequality we conclude that

To bound the second term in (6.24) we use (6.21) together with $\tau\leq T\leq\frac{\Gamma}{\eta\alpha^{2}}$ to conclude that

Combining (6.25) and (6.26) in (6.24), we conclude that

Here, (a) follows from $\varepsilon\leq\frac{\delta\alpha^{3}}{5\Gamma\beta^{2}}$ per Assumption (5.11), (b) from $\varepsilon_{0}\leq\frac{1}{5}\sqrt{\frac{\delta\alpha^{3}}{\Gamma\beta}}$ per Assumption (5.11), (c) from $\varepsilon\leq\frac{\delta\alpha^{3}}{5\Gamma\beta^{2}}\leq\frac{\delta\alpha}{5\Gamma}\leq\frac{\delta\alpha}{5}$ per Assumption (5.11), and (d) from $\varepsilon_{0}\leq\frac{\delta\alpha}{5}$ per Assumption (5.11). Thus,

The completes the proof of the bound (5.13).

Step V: Bound on residual with early stopping. In this step we wish to prove (5.14). To this aim note that

where (a) follows from the triangular inequality, (b) from the conclusion of Step II (first part of (5.12)), and (c) from Corollary 6.6 equation (6.10). This completes the proof of (5.14).

2 Key lemmas and identities for neural networks

In this section we prove some key lemmas and identities regarding the Jacobian of one-hidden layer networks as well as the size of the initial residual that when combined with Theorem 5.3 allows us to prove theorems involving neural networks. We begin with some preliminary identities and calculations in Section 6.2.1. Next, in Section 6.2.2 we prove a few key properties of the Jacobian mapping of a one-hidden layer neural network. Section 6.2.3 focuses on a few further properties of the Jacobian at a random initialization. Finally, in Section 6.2.4 we provide bounds on the initial misfit.

Alternatively using Khatri-Rao products this can be rewritten in the more compact form

2.2 Fundamental properties of the Jacobian of the neural network

In this section we prove a few key properties of the Jacobian mapping of a one-hidden layer neural network.

Proof The result on spectral norm and Lipschitzness of ${\cal{J}}(\bm{W})$ have been proven in . To show the row-wise bound (6.30), we use (6.29) to conclude that

Next we extend the lemma above to the multi-class setting.

Proof The proof will follow from Lemma 6.9. First, given ${\bm{A}}=[{\bm{A}}_{1}^{T}~{}\dots~{}{\bm{A}}_{K}^{T}]^{T}$ and ${{\bm{B}}}=[{{\bm{B}}}_{1}^{T}~{}\dots~{}{{\bm{B}}}_{K}^{T}]^{T}$, observe that

where the penultimate inequality follows from Cauchy Schwarz, completing the proof.

2.3 Properties of the Jacobian at random initialization

In this section we prove a few lemmas characterizing the properties of the Jacobian at the random initialization.

Setting $\bm{W}\overset{\text{i.i.d.}}{\sim}\mathcal{N}(0,1)$ and $\bm{V}$ with i.i.d. zero-mean and $\nu^{2}$-variance entries, we conclude that

Next we state a useful lemma from which allows us to bound the eigenvalues of the Hadamard product of the two PSD matrices.

Next we state a lemma regarding concentration of the Jacobian matrix at initialization.

with probability at least $1-1/n^{100}$. In particular, as long as

Proof Define ${\bm{C}}={\cal{J}}(\bm{W}_{0}){\cal{J}}(\bm{W}_{0})^{T}$. We begin by showing that the diagonal blocks of ${\bm{C}}$ are concentrated. To do this first for $1\leq s\leq k$ define the random matrices

Combining the latter two identities via the triangular inequality we conclude that

To proceed, we will bound the weighted sum

in spectral norm. To this aim we utilize the Matrix Hoeffding inequality which states that

concluding the proof of concentration of the diagonal blocks of ${\bm{C}}$.

so that we can take $\Delta^{2}=\nu^{4}kB^{4}\left\|{\bm{X}}\right\|^{4}$ and again conclude that for $t=30\sqrt{k}\nu^{2}B^{2}\|{\bm{X}}\|^{2}\log(n)$ we have

concluding the proof. The result in terms of $\delta$ is obtained by using the population covariance Lemma 6.11.

2.4 Upper bound on initial residual

In this section we prove a lemma concerning the size of the initial misfit. The proof of this lemma (stated below) follows from a similar argument in the proof of [48, Lemma 6.12].

holds with probability at least $1-(2{K})^{-100}$.

We will show that for any row $\bm{v}$ of ${\bm{V}}$, with probability at least $1-(2{K})^{-101}$,

where the latter follows from Poincare inequality (e.g. see [35, p. 49]). Furthermore, since $\bm{v}$ has i.i.d. Rademacher entries, applying Bernstein bound, event

holds with probability $1-(2{K})^{-102}$. Conditioned on $E_{\bm{v}}$, we now upper bound the expectation via

holds with probability at least $1-\exp(-102\log(2{K})\frac{n}{\left\|{\bm{X}}\right\|^{2}})\geq 1-(2{K})^{-102}$ where we used $n\geq\|{\bm{X}}\|^{2}$. Using a union bound over $E_{\bm{v}}$ and the conditional concentration over $\bm{W}$, the overall probability of success in (6.38) is at least $1-(2{K})^{-101}$ concluding the proof of (6.34) and the Lemma.

3 Rademacher complexity and generalization bounds

We begin by stating a vector contraction inequality by Maurer . This is obtained by setting $h_{i}(f(\bm{x}_{i}))=h(\bm{y}_{i},f(\bm{x}_{i}))$ in Corollary 4 of .

Combining the above result with standard generalization bounds based on Rademacher complexity allows us to prove the following result.

With probability $1-\delta$ over the samples, for all $f\in\mathcal{F}$, we have that

holds with $1-\delta$ probability. Combining the latter with Lemma 6.15 completes the proof.

We proceed by bounding each of these four terms. For the first term note that

where in the last inequality we used the fact that $\|{{\bm{V}}_{0}}\|_{F}\leq\nu$. For the second term note that

In the above we used $\|\bm{M}\|_{2,1}$ for a matrix $\bm{M}$ to denote the sum of the Euclidean norm of the rows of $\bm{M}$. We also used the fact that $\|{\bm{V}}_{0}^{T}\|_{2,1}\leq\nu\sqrt{k}$. To bound the third term note that

Finally, to bound the fourth term note that we have

Combining these four bounds we conclude that

Next we state a crucial lemma that connects the test error measured by any Lipschitz loss to that of the quadratic loss on the training data.

Then for all $f$ in the function class $\mathcal{F}_{{\cal{W}}}$ given by (6.39)

where the last inequality follows from Cauchy-Schwarz. Consequently, applying Lemmas 6.16 and 6.17 we conclude that

Combining the latter with Markov inequality we arrive at

4 Proofs for neural nets with arbitrary initialization (Proof of Theorem 3.3)

In this section we prove Theorem 3.3. We first discuss a preliminary optimization result in Section 6.4.1. Next, in Section 6.4.2 we build upon this result to prove our main optimization result. Finally, in Section 6.4.3 we use these optimization results to prove our main generalization result, completing the proof of Theorem 3.3.

with $0\leq\delta\leq 1$ and $\Gamma\geq 1$. We run gradient descent iterations of the form $\bm{W}_{\tau+1}=\bm{W}_{\tau}-\eta\nabla\mathcal{L}(\bm{W}_{\tau})$ starting from $\bm{W}_{0}$ with step size $\eta$ obeying $\eta\leq\frac{1}{\nu^{2}B^{2}\left\|{\bm{X}}\right\|^{2}}$. Then for all iterates $\tau$ obeying $0\leq\tau\leq T:=\frac{\Gamma}{\eta\alpha^{2}}$

Furthermore, after $\tau=T$ iteration we have

as long as (6.44) holds by Lemma 6.10 we have

Hence, using Lemma 6.10 equation (6.31) we conclude that

To bound the right-hand side we use the triangular inequality combined with (6.25) and (6.26) to conclude that

where in the last inequality we used the fact that $\varepsilon_{0}^{2}\leq\frac{\alpha^{2}}{25\Gamma}$ per (6.43) and $\epsilon\beta=\frac{\delta\alpha^{3}}{5\Gamma\beta}\leq\frac{\alpha^{2}}{5\Gamma}$ per our choice of $\epsilon$. Combining (6.48) and (6.4.1), we obtain

completing the proof of (6.46) and the theorem.

4.2 Main Optimization Result

and $\Gamma\geq 1$. We run gradient descent iterations of the form $\bm{W}_{\tau+1}=\bm{W}_{\tau}-\eta\nabla\mathcal{L}(\bm{W}_{\tau})$ starting from $\bm{W}_{0}$ with step size $\eta$ obeying $\eta\leq\frac{1}{\nu^{2}B^{2}\left\|{\bm{X}}\right\|^{2}}$. Then for all iterates $\tau$ obeying $0\leq\tau\leq T:=\frac{\Gamma}{\eta\alpha^{2}}$

Furthermore, after $\tau=T$ iteration we have

Proof To prove this lemma we aim to substitute

in Theorem 6.19. To do this we need to verify the assumptions of Theorem 6.19. To this aim note that the choice of $\delta$ from (6.55) combined with (6.51) ensures that

so that (6.44) holds. We thus turn our attention to proving (6.43). If $\varepsilon_{0}=0$, the statement already holds. Otherwise, note that based on Lemma 6.2 equation (6.3) we have

Recall that $\alpha=\frac{\nu\alpha_{0}}{\sqrt{K}}$, which implies that

If $\delta=\frac{\zeta\nu B\left\|{\bm{X}}\right\|}{\alpha}$: For (6.43) to hold it suffices to have $\varepsilon_{0}\leq\frac{\alpha}{5}\min\left(\frac{\zeta\nu B\left\|{\bm{X}}\right\|}{\alpha},\sqrt{\frac{\zeta}{\Gamma}}\right)$.

Combining the latter two cases as long as

4.3 Main generalization result (completing the proof of Theorem 3.3)

Theorem 3.3 immediately follows from Theorem 6.21 below by upper bounding ${\cal{D}}_{\alpha,\Gamma}$ (see Definition 6.1) using Lemma 6.2 equation (6.2).

with $\Gamma\geq 1$ and tolerance level $\zeta\leq 2$. Run gradient descent updates (1.5) with learning rate $\eta\leq\frac{1}{\nu^{2}B^{2}\left\|{\bm{X}}\right\|^{2}}$. Then, after $T=\frac{\Gamma}{\eta\alpha^{2}}$ iterations, with probability at least $1-\delta$, the generalization error obeys

where ${\cal{D}}_{\alpha,\Gamma}$ is the early stopping distance as in Def. (6.1).

This together with $\zeta\leq 2$ implies that

Thus, (6.51) holds. Also (6.50) trivially holds for $\varepsilon_{0}$. Thus applying Theorem 6.20 with $\varepsilon_{0}=0$ the following three conclusions hold

Plugging (6.4.3) and (6.64) into (6.63) completes the proof.

5 Proofs for neural network with random initialization (proof of Theorem 3.2)

In this section we prove Theorem 3.2. We first discuss and prove an optimization result in Section 6.5.1. Next, in Section 6.5.2 we build upon this result to complete the proof of Theorem 3.2.

with $\Gamma\geq 1$ and $\zeta\leq\frac{c}{2}$. We run gradient descent iterations of the form (1.5) with a learning rate $\eta\leq\frac{1}{\nu^{2}B^{2}\left\|{\bm{X}}\right\|^{2}}$. Then, after $T=\frac{\Gamma K}{\eta\nu^{2}\alpha_{0}^{2}}$ iterations, the following identities

hold with probability at least $1-(2{K})^{-100}$.

Proof To prove this result we wish to apply Theorem 6.20. To do this we need to verify the assumptions of this theorem. To start with, using Lemma 6.14 with probability at least $1-(2{K})^{-100}$, the initial prediction vector $f(\bm{W}_{0})$ obeys

Therefore, $\bm{J}$ is an $\varepsilon_{0}^{2}=4\frac{\delta\nu^{2}}{K}$ reference Jacobian. Now set

Hence, to ensure (6.50) holds we need to ensure that $\delta$ obeys

Thus using $\delta=\frac{\alpha_{0}^{2}}{200}\Theta$ to ensure (6.50) we need to make sure $k$ is sufficiently large so that (6.75) holds with this value of $\delta$. Thus it suffices to have

To be able to apply Theorem 6.20 we must also ensure (6.51) holds. Therefore, it suffices to have

Here, (a) follows from the fact that $n\geq K$ and the relationship between $\zeta$ and $\nu$ per (6.85) and (b) follows from the fact that per equation (6.2) we have

Note that (6.78) and (6.82) are implied by

5.2 Generalization result (completing the proof of Theorem 3.2)

Theorem below is a restatement of Theorem 3.2 after substituting the upper bound on the early stopping distance ${\cal{D}}_{\alpha_{0},\Gamma}$ of Def. (6.1).

with $\Gamma\geq 1$. We run gradient descent iterations of the form (1.5) with a learning rate $\eta\leq\frac{1}{\nu^{2}B^{2}\left\|{\bm{X}}\right\|^{2}}$. Then, after $T=\frac{\Gamma K}{\eta\nu^{2}\alpha_{0}^{2}}$ iterations, classification error $\text{Err}_{{\cal{D}}}(\bm{W}_{T})$ is upper bounded by

holds with probability at least $1-(2{K})^{-100}-\delta$.

Using (6.69) on row bound $R$ and lower bound on $k$

Plugging in (6.88), (6.89), and (6.5.2) into (6.87) concludes the proof.

Acknowledgements

M. Soltanolkotabi is supported by the Packard Fellowship in Science and Engineering, a Sloan Research Fellowship in Mathematics, an NSF-CAREER under award #1846369, the Air Force Office of Scientific Research Young Investigator Program (AFOSR-YIP) under award #FA9550-18-1-0078, an NSF-CIF award #1813877, and a Google faculty research award.

References

Appendix A The Jacobian of the Mixture Model is low-rank (Proofs for Section 3.3)

The following theorem considers a simple noiseless mixture model and proves that its Jacobian is low-rank and the concatenated multiclass label vectors lie on a rank $K^{2}C$ information space associated with this Jacobian.

and assume that $\widetilde{{\bm{{\Sigma}}}}({\bm{M}})$ is full rank. Then, the following properties hold with probability $1-KC\exp(-\frac{n}{8KC})$

$\mathcal{I}$ is a $K^{2}C$ dimensional subspace.

The concatenated label vector $\bm{y}=\begin{bmatrix}\bm{y}_{1}^{T}&\bm{y}_{2}^{T}&\ldots&\bm{y}_{n}^{T}\end{bmatrix}^{T}$ lies on $\mathcal{I}$.

The nonzero eigenvalues (top $K^{2}C$ eigenvalues) of ${\bm{{\Sigma}}}({\bm{X}})$ are between $\frac{n}{2KC}s_{\min}({\bm{{\Sigma}}}({\bm{X}}))$ and $\frac{2n}{KC}\|{\bm{{\Sigma}}}({\bm{X}})\|$. Hence the eigenvalues of the information space grow with $\frac{n}{{K}C}$.

Proof First, we establish that each cluster has around the same size. Applying Chernoff bound and a union bound, we find that with probability $1-KC\exp(-\frac{n}{8KC})$

Note that based on Lemma 6.11, the multiclass covariance is given by

we have ${\mathcal{I}}={\bm{I}}_{{K}}\otimes\widetilde{{\mathcal{I}}}$ which also implies $\text{rank}({\mathcal{I}})={K}\cdot\text{rank}\left(\widetilde{{\mathcal{I}}}\right)$. Hence, this identity allows us to reduce the problem to a single output network. To complete the proof we will prove the following three identities:

$\widetilde{{\mathcal{I}}}$ has rank ${K}C$.

The nonzero eigenvalues of $\widetilde{{\bm{{\Sigma}}}}\left({\bm{X}}\right)$ are between $0.5\widetilde{n}s_{\min}(\widetilde{{\bm{{\Sigma}}}}({\bm{M}}))$ to $2\widetilde{n}\|\widetilde{{\bm{{\Sigma}}}}({\bm{M}})\|$.

Now note that using the above identity we have

Since ${{\bm{U}}}_{\mathcal{S}}$ is tall and orthogonal, the range of $\widetilde{{\bm{{\Sigma}}}}({\bm{X}})$ is exactly the range of ${\bm{U}}_{\mathcal{S}}$ hence $\widetilde{{\mathcal{I}}}=\mathcal{S}$ which is ${K}C$ dimensional. Furthermore, nonzero eigenvectors of $\widetilde{{\bm{{\Sigma}}}}({\bm{X}})$ lie on $\mathcal{S}$ and any eigenvector $\bm{v}$ satisfies

Next lemma provides a perturbation analysis when there is noise.

Consider the single-output NTK kernel given by

and assume that this matrix has rank $r$ so that $\lambda_{r+1}\left(\widetilde{\bm{\Sigma}}({\bm{X}})\right)=\lambda_{r+2}\left(\widetilde{\bm{\Sigma}}({\bm{X}})\right)=\ldots=\lambda_{n}\left(\widetilde{\bm{\Sigma}}({\bm{X}})\right)=0$. Also assume a noise corrupted version of ${\bm{X}}$ given by

with $\bm{Z}$ a matrix consisting of i.i.d. $\mathcal{N}(0,1)$ entries. Then, $\left\|\widetilde{\bm{\Sigma}}(\widetilde{{\bm{X}}})-\widetilde{\bm{\Sigma}}({\bm{X}})\right\|\lesssim\Delta$ where

Now define $\widetilde{\bm{M}}=\text{diag}\left(\phi^{\prime}(\widetilde{{\bm{X}}}\bm{w})\right)\widetilde{{\bm{X}}}$ and $\bm{M}=\text{diag}\left(\phi^{\prime}({\bm{X}}\bm{w})\right){\bm{X}}$ and note that using the above we can conclude that

To proceed further, with probability $1-n\exp(-d/2)$, each row of $\widetilde{{\bm{X}}}-{\bm{X}}$ is upper bounded by $2\sigma$. Hence, using a standard tail bound over supremum of $n$ Gaussian random variables (which follows by union bounding) we have

holds with the same probability. Furthermore, spectral norm bound on Gaussian random matrix implies that

The following lemma plugs in the critical quantities of Theorem 3.2 for our mixture model to obtain a generalization bound.

Consider the setup of Theorem 3.2 with quantities $\zeta$ and $\Gamma$. Suppose network width obeys

Proof The proof is an application of Lemma A.2 and Theorem A.1. Let $\mathcal{I}^{\prime}$ be the information space corresponding to noiseless dataset where input samples are identical to cluster centers. Let ${\bm{P}}^{\prime},{\bm{P}}$ correspond to the projection matrices to $\mathcal{I}$ and $\mathcal{I}^{\prime}$. First, using Lemma A.2 and the bound on $\sigma$, we have

for some constant $c>0$. Next we quantify $\Pi_{\mathcal{I}}(\bm{y})$ using the fact that (i) $\Pi_{\mathcal{I}^{\prime}}(\bm{y})=\bm{y}$ via Theorem A.1 as follows

To proceed, we pick $\alpha_{0}=\sqrt{\frac{\lambda_{\min}n}{2KC}}$ and corresponding $\bar{\alpha}=\frac{\alpha_{0}}{\sqrt{n}\sqrt{K\left\|{\bm{X}}\right\|}B}\geq\sqrt{\frac{\lambda_{\min}}{2B^{2}K^{2}C}}$ and apply (3.3) to find that, classification error is upper bounded by

Appendix B Joint input-output optimization

In this section we wish to provide the ingredients necessary to prove a result for the case where both set of input and output weights $\bm{W}$ and ${\bm{V}}$ are trained. To this aim, we consider the combined neural net Jacobian associated with input and output layers given by

Denoting the Jacobian associated with (B.1) by ${\cal{J}}({\bm{V}},\bm{W})$ we have that

Here, ${\cal{J}}(\bm{W})$ is as before whereas ${\cal{J}}({\bm{V}})$ is the Jacobian with respect to ${\bm{V}}$ and is given by

Hence, ${\cal{J}}({\bm{V}})$ is ${K}\times{K}$ block diagonal with blocks equal to $\phi({\bm{X}}\bm{W}^{T})$. The following theorem summarizes the properties of the joint Jacobian.

${\cal{J}}({\bm{V}},\bm{W})$ satisfies the following properties.

Lipschitzness: Given inputs ${\bm{V}},{\bm{V}}^{\prime}$ and outputs $\bm{W},\bm{W}^{\prime}$

Proof First, we prove results concerning ${\cal{J}}({\bm{V}})$. First, note that

Let ${\cal{J}}_{1},{\cal{J}}_{2}$ be the Jacobian matrices restricted to ${\bm{V}}$ and $\bm{W}$ of ${\cal{J}}({\bm{V}},\bm{W})$. To prove Lipschitzness, first observe that

To address the second term, note that, Jacobian is linear with respect to output layer hence

Combining the latter two identities we arrive at