Year: 2019

Abstract:
Modern neural network architectures often generalize well despite containing many more parameters than the size of the training dataset. This paper explores the generalization capabilities of neural networks trained via gradient descent. We develop a data-dependent optimization and generalization theory which leverages the low-rank structure of the Jacobian matrix associated with the network. Our results help demystify why training and generalization is easier on clean and structured datasets and harder on noisy and unstructured datasets as well as how the network size affects the evolution of the train and test errors during training. Specifically, we use a control knob to split the Jacobian spectum into “information" and “nuisance" spaces associated with the large and small singular values. We show that over the information space learning is fast and one can quickly train a model with zero training loss that can also generalize well. Over the nuisance space training is slower and early stopping can help with generalization at the expense of some bias. We also show that the overall generalization capability of the network is controlled by how well the labels are aligned with the information space. A key feature of our results is that even constant width neural nets can provably generalize for sufficiently nice datasets. We conduct various numerical experiments on deep networks that corroborate our theoretical findings and demonstrate that: (i) the Jacobian of typical neural networks exhibit low-rank structure with a few large singular values and many small ones leading to a low-dimensional information space, (ii) over the information space learning is fast and most of the labels falls on this space, and (iii) label noise falls on the nuisance space and impedes optimization/generalization.

heading: 1 INTRODUCTION

heading: 1.1 MOTIVATION AND CONTRIBUTIONS
text: 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 Zhang et al. (2016). 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 with 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.● 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).● Bias–variance tradeoffs: We develop a generalization 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.
● 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 compares favorably with recent literature on optimization and generalization of neural networks Jacot et al. (2018); Arora et al. (2019); Du et al. (2018b); Allen-Zhu et al. (2018b); Cao & Gu (2019); Ma et al. (2019); Allen-Zhu et al. (2018a); Brutzkus et al. (2017) where guarantees only hold for very wide networks with the width of the network growing inversely proportional to the class margins or related notions. See Section 3 for further detail.● Pretrained models: Our framework does not require random initialization 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 commonly used in meta/transfer learning. Our extensive experiments strongly suggest Jacobian adapts over time in a favorable and data-dependent fashion shedding light on the properties of (pre)trained models.

heading: 1.2 MODEL AND TRAINING
text: Our theoretical analysis will focus on neural networks consisting of one hidden layer with d input features, k hidden neurons and K outputs as depicted in Figure 1. We use W ∈ Rk×d and V ∈ RK×k to denote the input-to-hidden and hiddento-output weights. The overall input-output relationship of the network is a function f(⋅;W ) ∶ Rd → RK that maps an input x ∈ Rd to an output via x↦ f(x;W ) ∶= V φ(Wx). (1.1) Given a dataset consisting of n feature/label pairs(xi,yi) with xi ∈ Rd representing the features and yi ∈ RK the associated labels representing one of K classes with one-hot encoding (i.e. yi ∈{e1,e2, . . . ,eK} where e` ∈ RK has all zero entries except for the `th entry which is equal to one).
To learn this dataset, we fix the output layer and train overW via1
min W ∈Rk×d L(W ) ∶= 12 n∑ i=1 ∥V φ (Wxi) − yi∥2`2 . (1.2)
It will be convenient to concatenate the labels and prediction vectors as follows
y = ⎡⎢⎢⎢⎢⎣ y1⋮ yn ⎤⎥⎥⎥⎥⎦ ∈ R nK and f(W ) = ⎡⎢⎢⎢⎢⎣ V f(x1;W )⋮ V f(xn;W ) ⎤⎥⎥⎥⎥⎦ ∈ R nK . (1.3)
Using this shorthand we can rewrite the loss (1.2) as
min W ∈Rk×d L(W ) ∶= 12 ∥f(W ) − y∥2`2 . (1.4)
To optimize this loss starting from an initializationW0 we run gradient descent iterations of the form Wτ+1 =Wτ − η∇L(Wτ), (1.5)
with a step size η. 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.

heading: 1.3 INFORMATION AND NUISANCE SPACES
text: In order to understand the generalization capabilities of models trained via gradient descent we need to develop better insights into the form of the gradient updates and how it affects the training dynamics. To this aim let us aggregate the weights at each iteration into one large vector wτ ∶=vect(Wτ) ∈ Rkd, define the misfit/residual vector r(w) ∶= f(w) − y and note that the gradient updates take the form wτ+1 =wτ − η∇L(wτ) where ∇L(w) = ∇L(w) = J (w)Tr(w). 1For clarity of exposition, we focus only on optimizing the input layer. However, as shown in the supplementary, the technical approach is quite general and applies to arbitrary multiclass nonlinear least-squares problems. In particular, the proofs are stated so as to apply to one-hidden layer networks where both layers are trained.
Here, J (w) ∈ RnK×kd denotes the Jacobian mapping associated with f defined as J (w) = ∂f(w) ∂w
. Due to the form of the gradient updates the dynamics of training is dictated by the spectrum of the Jacobian matrix as well as the interaction between the residual vector and the Jacobian. If the residual vector is very well aligned with the singular vectors associated with the top singular values of J (w), the gradient update significantly reduces the misfit allowing substantial reduction in the train error. Thus to provide a more precise understanding of the training dynamics and generalization capabilities of neural networks it is crucial to develop a better understanding of the interaction between the Jacobian and the misfit and label vectors. To capture these interactions we require a few definitions. Definition 1.1 (Information & Nuisance Spaces) Consider a matrix J ∈ RnK×p with singular value decomposition given by
J = nK∑ s=1λsusv T s = Udiag (λ1, λ2, . . . , λnK)V T ,
with λ1 ≥ λ2 ≥ . . . ≥ λnK denoting the singular values of J in decreasing order and {us}nKs=1 ∈ RnK and {vs}nKs=1 ∈ Rp the corresponding left and right singular vectors forming the orthonormal basis matrices U ∈ RnK×nK and V ∈ Rp×nK . For a spectrum cutoff α obeying 0 ≤ α ≤ λ1 let r ∶= r(α) denote the index of the smallest singular value above α. We define the information and nuisance spaces associated with J as I ∶= span({us}rs=1) and N ∶= span({us}Kns=r+1). In this paper we shall use either the expected value of the Jacobian at the random initialization or the Jacobian at one of the iterates to define the matrix J and the corresponding information/nuisance spaces. More, specifically we will set J to either J = (E[J (W0)J T (W0)])1/2 or J = J (Wτ). Therefore, one can effectively think of the information space as the span of the prominent singular vectors of the Jacobian and the nuisance space as its complement. In particular, as we demonstrate in Section 4 the Jacobian mapping associated with neural networks exhibit low-rank structure with a few large singular values and many small ones leading to natural choices for the cut-off value α as well as the information and nuisance spaces. Furthermore, we demonstrate both (empirically and theoretically) that learning is fast over the information space leading to a significant reduction in both train/test accuracy in the early stages of training. However, after a certain number of iterations learning shifts to the nuisance space and reduction in the training error significantly slows down (see Fig. 2). Furthermore, subsequent iterations in this stage lead to a slight increase in test error.

heading: 2 MAIN RESULTS
text: Our main results establish multi-class generalization bounds for neural networks trained via gradient descent. First, we will focus on networks where both layers are randomly initialized. Next we will provide guarantees for arbitrary initialization with the goal of characterizing the generalization ability of subsequent iterative updates for a given (possibly pre-trained) network in terms of its Jacobian mapping. In this paper we focus on activations φ which are smooth and have bounded first and
second order derivatives. This would for instance apply to the softplus activation φ(z) = log (1 + ez). We note that utilizing a proof technique developed in Oymak & Soltanolkotabi (2019) for going from smooth to ReLU activations it is possible to extend our results to ReLU activations with proper modifications. We avoid doing this in the current paper for clarity of exposition. Before we begin discussing our main results we discuss some notation used throughout the paper. For a matrix X ∈ Rn×d we use smin(X) and smax(X) = ∥X∥ to denote the minimum and maximum singular value of X . For two matrices A and B we use A ⊙B and A ⊗B to denote their Hadamard and Kronecker products, respectively. For a PSD matrix A ∈ Rn×n with eigenvalue decomposition A = ∑ni=1 λiuiuTi , the square root matrix is defined as A1/2 ∶= ∑ni=1 √λiuiuTi . We also use A† to denote the pseudo-inverse of A. In this paper we mostly focus on label vectors y which are one-hot encoded i.e. all entries are zero except one of them. For a subspace S ⊂ Rn and point x ∈ Rn, ΠS(x) denotes the projection of x onto S. Finally, before stating our results we need to provide a quantifiable measure of performance for a trained model. Given a sample (x,y) ∈ Rd × RK from a distribution D, the classification error of the networkW with respect to D is defined as
ErrD(W ) = P{arg max 1≤`≤K y` ≠ arg max1≤`≤K f`(x;W )}. (2.1)

heading: 2.1 RESULTS FOR RANDOM INITIALIZATION
text: To explore the generalization of randomly initialized networks, we utilize the neural tangent kernel.
Definition 2.1 (Multiclass Neural Tangent Kernel (M-NTK) Jacot et al. (2018)) Let w ∈ Rd be a vector with N (0,Id) distribution. Consider a set of n input data points x1,x2, . . . ,xn ∈ Rd aggregated into the rows of a data matrix X ∈ Rn×d. Associated to the activation φ and the input data matrixX we define the multiclass kernel matrix as
Σ(X) ∶= IK ⊗ E [ (φ′ (Xw)φ′ (Xw)T )⊙ (XXT ) ], where IK is the identity matrix of size K. Here, the ` th diagonal block of Σ(X) corresponds to the kernel matrix associated with the ` th network output for 1 ≤ ` ≤ K. This kernel is intimately related to the multiclass Jacobian mapping. In particular, suppose the initial input weightsW0 are distributed i.i.d. N (0,1) and the output layer V has i.i.d. zero-mean entries with ν2/K variance. Then E[J (W0)J (W0)T ] = ν2Σ(X). We use the square root of this multiclass kernel matrix (i.e. Σ(X)1/2) to define the information and nuisance spaces for our random initialization result. The following theorem is a (non-rigorous) simplification of our main result Theorem 6.24 where we ignore constants and log factors, and state a weaker but simpler generalization bound. Theorem 2.2 Fix numbers Γ ≥ 1 and α > 0. Consider an i.i.d. training dataset {(xi,yi)}ni=1 ∈ Rd × RK with unit length input samples and one-hot encoded labels. Consider the neural net in (1.1) parameterized by W and initialized with W0
i.i.d.∼ N (0,1) entries. Set V with i.i.d. Rademacher entries (properly scaled). Define the information I and nuisance N spaces with respect to Σ(X)1/2 with spectrum cutoff α √ nK per Definition 1.1. Furthermore, assume
k ≳ Γ4 logn α8 . (2.2)
Then after T ∝ Γ/α2 gradient iterations of (1.5), with high probability, training loss obeys ∥f(WT ) − y∥`2 ≲ ∥ΠN (y)∥`2 + e−Γ√n, . (2.3)
Furthermore, the classification error obeys
ErrD(WT ) ≲ ∥ΠN (y)∥`2√ n + e−Γ + Γ α √ n .
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 y mostly lies on the information space.
● Generalization bound: The generalization error has two core components: bias and variance. The bias component ∥ΠN (y)∥`2/√n + e−Γ arises from the training loss and corresponds to the portion of the labels that falls over the nuisance space. The variance component Γ/α√n corresponds to the Rademacher complexity of the model space which connects to the distance ∥WT −W0∥F . If y is aligned with the information space, the bias term ΠN (y) will be small. Additionally, if the kernel matrix is low-rank, we can pick a large α to ensure small variance as well as small network width. In particular with a constant α the required network width is logarithmic in n.
We note however that our results continue to apply even when the kernel is not approximately low-rank. In particular, consider the extreme case where we select α √ nK = √λ ∶= √λmin (Σ(X)).
This sets I = RKn and ∥ΠN (y)∥`2 = 0. For this case, the more general Theorem 6.24 yields ErrD(WT ) ≲ √ K√ n √ yTΣ−1(X)y while requiring a width of k ≳ K4n4 logn λ4 . (2.4)
We note that in this special case our results improve upon the required width in recent literature Arora et al. (2019)2 that focuses on K = 1 and a conclusion of the form (2.4). However, as we demonstrate in our numerical experiments in practice λ is very small or even zero (e.g. see the toy model in Section 2.3) so that requirements of the form (2.4) may require unrealistically (or even infinitely) wide networks. In contrast, our results apply to all Jacobian spectrums, however can further harness the low-rank structure of the Jacobian to give even stronger bounds.● Small width is sufficient for generalization: Based on our simulations the M-NTK indeed has low-rank structure with a few large eigenvalues and many smaller ones. As a result a reasonable scaling choice of α is constant. In that case our result states that as soon as the number of hidden nodes are logarithmic in n, good generalization can be achieved. This favorably compares to related works Jacot et al. (2018); Arora et al. (2019); Du et al. (2018b); Allen-Zhu et al. (2018b); Cao & Gu (2019) where network size is required to grow polynomial with n and inversely with the distance between the inputs or other notions of margin.● Network size–Bias tradeoff: Based on the requirement (2.2) if the network is large (in terms of # of hidden units k), we can choose a small cut-off α. This in turn allows us to enlargen the information space and reduce the training bias further. In summary, as the network capacity grows, we can gradually interpolate finer detail and reduce bias.● Fast convergence: Note that the number of gradient iterations is upper bounded by Γ/α2. Hence, the training speed is dictated by and is inversely proportional to the the smallest singular value over the information space. Specifically, picking α to be a constant, convergence on the information space will be fast requiring only a constant number of iterations to reach any fixed accuracy (see (2.3)).

heading: 2.2 GENERALIZATION GUARANTEES WITH ARBITRARY INITIALIZATION
text: 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 2.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 model3 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. The following theorem is a (non-rigorous) simplification of Theorem 6.21. Theorem 2.3 Let Γ ≥ 1, α be arbitrary scalars. Consider i.i.d. training data {(xi,yi)}ni=1 ∈ Rd×RK with unit length inputs and one-hot encoded labels. Also consider a neural net with k hidden nodes as in (1.1) parameterized by W . Let W0 be an arbitrary initial weight matrix and assume the output matrix has bounded entries obeying ∥V ∥`∞ ≤ 1√kK . Define the nuisance space N associated with J (W0) based on spectrum cutoff α√n. Set the initial residual r0 = f(W0) − y ∈ RnK and assume ∥r0∥`2 ≲ √n. Suppose k ≳ Γ4/α8. After T ∝ Γ/α2 iterations (1.5) with constant learning rate, training loss obeys: ∥f(WT ) − y∥`2 ≲ ∥ΠN (r0)∥`2 + e−Γ√n. 2Based on our understanding Arora et al. (2019) requires the number of hidden units to be on the order of k ≳ n8/λ6. Hence our result reduces the dependence on width by a factor of at least n4/λ2. 3e.g. obtained by training with data in a related problem as is common in transfer learning.
Also with high probability, classification error obeys: ErrD(WT ) ≲ ∥ΠN (r0)∥`2√n + e−Γ + Γα√n . 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 reasonable size. As with its randomized counterpart 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 develops more favorable properties with the labels/initial residuals becoming more correlated with the information space of the Jacobian. As mentioned earlier, due to 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 via its Jacobian, it is more likely to generalize well. Furthermore, given its deterministic nature the theorem can be applied at any iteration by setting W0 =Wτ , implying that if the Jacobians of any of the iterates provides a better low-rank representation then one can provide better generalization guarantees. Our numerical experiments demonstrate that the Jacobian of the neural network adapts to the dataset over time with a more substantial amount of the labels lying on the information space. While we defer the rigorous theory of this adaptation to future, Section D provides a proof sketch of evolution of Jacobian rank for a simple dataset model. Such a result when combined with our result above can potentially provide significantly tighter bounds. This is particularly important in light of recent literature Chizat & Bach (2018b); Ghorbani et al. (2019c); Yehudai & Shamir (2019) suggesting a significant generalization gap between kernel methods/linearized neural nets when compared with neural nets operating beyond the linearized regime (e.g. mean field regime). As a result we view our deterministic result as a first step towards moving beyond the NTK regime.

heading: 2.3 CASE STUDY: GAUSSIAN MIXTURE MODEL
text: To illustrate a concrete example, we consider a distribution based on multiclass mixture models. Definition 2.4 (Gaussian mixture model) Consider a size n dataset {(xi,yi)}ni=1 ∈ Rd × RK . We assume this dataset consists of K classes each comprising of C clusters with a total of KC clusters. We index each cluster with (`, ̃̀) denoting the ̃̀th cluster from the `th class. We assume the dataset in cluster (`, ̃̀) is centered around a cluster center µ`,̃̀ ∈ Rd with unit Euclidian norm. We assume the dataset is generated i.i.d. with the cluster membership assigned uniformly of the clusters with probability 1 KC and the input samples associated with the cluster (`, ̃̀) are generated i.i.d. according
to N (µ`,̃̀, σ2Id/d) with the corresponding label set to the one-hot encoding of the class ` i.e. e`. Note that the cluster indexed by (`, ̃̀) contains ñ`,̃̀data points satisfying E[ñ`,̃̀] = ñ = n/KC. 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 σ = 0 where we have a discrete input distribution over the cluster centers. In this scenario, the multi-class Jacobian matrix is at most rank K2C = # of output nodes × # of distinct inputs. as there are (i) onlyKC distinct input vectors and (ii)K output nodes. We can thus set the information space to be the top K2C eigenvectors of the multiclass kernel matrix Σ(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 y perfectly lies on the information space.
In Figure 3 we numerically verify that the approximate rank and singular values of the Jacobian indeed scale as above even when σ > 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. ζ,B, logn). In this theorem we use ≳ to denote inequality up to constant/logarithmic factors. Theorem 2.5 (Generalization for Gaussian Mixture Models-simplified) Consider a data set of size n consisting of input/label pairs {(xi,yi)}ni=1 ∈ Rd × RK generated according to Def. 2.4 with the standard deviation obeying σ ≲ K n . Let M = [µ1,1 . . . µK,C]T be the matrix obtained by
aggregating all the cluster centers and let g ∼ N (0,Id). Also let Σ(M) ∈ RKC×KC be the M-NTK
associated with the cluster centers M per Def. 2.1. Furthermore, set λM = λmin(Σ(M)), and assume λM > 0. If the number of hidden nodes obeys k ≳ Γ4K8C4λ4
M
. after T = 2ΓK2C λM gradient
iterations, with high probability, the model obeys ErrD(WT ) ≲ Γ√K2CnλM .
90 180 360 450
100
101
102
Singular value index (x-axis marks K2C)
C = 10 C = 20 C = 25 C = 40 C = 50
1Figure 3: The singular values of the normalized Jacobian spectrum √ KC n J (W0) of a neural
network with K = 3. Here, the data is generated according to the Def. 2.4 with K classes and σ = 0.1. The cluster centers are picked so that the distance between any two is at least 0.5. We consider two cases: n = 30C (solid line) and n = 60C (dashed line). These plots demonstrate that the top KC singular values grow proportional to √ n.
We note that λM captures how diverse the cluster centers are. In this sense λ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 λM scales like a constant Oymak & Soltanolkotabi (2019). This theorem focuses on the regime where the noise level σ is small. In this case one can achieve good generalization as soon as the sample size scales as n ≳ K2C which is the effective rank of the M-NTK matrix. This result follows from our main result with random initialization by setting the cutoff at α2 ∼ λM
K2C . This
demonstrates that in this model α does indeed scale as a constant. Finally, the required network width is independent of n and only depends on K and C specifically k ≳ K8C4. This compares favorably with Arora et al. (2019) which concerns the K = 1 case. In particular, Arora et al. (2019) requires k ≳ n8/λ6X which depends on n (in lieu of K and C) and the minimum eigenvalue λX of the NTK matrix Σ(X) (rather than λM ). Furthermore, as σ → 0, Σ(X) becomes rank deficient and λX → 0 so that Arora et al. (2019) requires infinite width.

heading: 3 PRIOR ART
text: Neural networks have impressive generalization abilities even when they are trained with more parameters than the data Zhang et al. (2016). Thus, optimization/generalization properties of neural nets have been the topic of recent literature Zhang et al. (2016). Below we discuss the works on statistical learning, optimization, and implicit bias. Statistical learning theory: Statistical properties of neural networks have been studied since 1990’s Anthony & Bartlett (2009); Bartlett et al. (1999); Bartlett (1998). 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 Dziugaite & Roy (2017); Arora et al. (2018); Neyshabur et al. (2017b); Golowich et al. (2017). Bartlett et al. (2017); Neyshabur et al. (2017a) established tight sample complexity results for deep networks based on spectral norms. See also Nagarajan & Kolter (2019) for improvements via leveraging various properties of the inter-layer Jacobian and Long & Sedghi (2019) for results with convolutional networks. Related, Arora et al. (2018) leverages compression techniques for constructing tighter bounds. Yin et al. (2018) jointly studies statistical learning and adversarial robustness. These interesting results, provide generalization guarantees for the optimal solution to the empirical risk minimizer. 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 Keskar et al. (2016); Hardt et al. (2015); Sagun et al. (2017); Chaudhari et al. (2016); Hoffer et al. (2017); Goel & Klivans (2017); Goel et al. (2018). For linearly separable datasets, Soudry et al. (2018); Gunasekar et al. (2018); Brutzkus et al. (2017); Ji & Telgarsky (2018a;b) 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 Dou & Liang (2019); Belkin et al. (2018a;b; 2019); Liang & Rakhlin (2018); Belkin et al. (2018c). Chizat & Bach (2018a); Song et al. (2018); Mei et al. (2018); Sirignano & Spiliopoulos (2018); Rotskoff & Vanden-Eijnden (2018) relate the distribution of the network weights to Wasserstein gradient flows using mean field analysis.
Global convergence and generalization of neural nets: Closest to our work, recent literature Cao & Gu (2019); Arora et al. (2019); Ma et al. (2019); Allen-Zhu et al. (2018a) provides generalization bounds for overparameterized networks trained via gradient descent. Also see Li et al. (2018); Huang et al. (2019) for interesting visualization of the optimization and generalization landscape. Jacot et al. (2018) introduced NTK and observed that principal directions of NTK is optimized faster than smaller eigendirections for infinitely wide networks. In connection to this, our Def 1.1 helps quantify the low-rankness and bimodality of the Jacobian spectrum (same as NTK for random initialization). Similar to Thm 2.2, Arora et al. (2019) uses the NTK to provide generalization guarantees in a similar framework to Jacot et al. (2018) (see (2.4) for comparison). Li et al. (2019a) leverages low-rank Jacobian structure to establish robustness to label noise. Very recent work Su & Yang (2019) uses low-rankness to better capture approximation power of neural nets. These works build on global convergence results of randomly initialized networks Du et al. (2018b;a); Allen-Zhu et al. (2018b); Chizat & Bach (2018b); Zhang et al. (2019); Nitanda & Suzuki (2019); Oymak & Soltanolkotabi (2018); Zou et al. (2018) which study the gradient descent trajectory via comparisons to a NTK linearization. These results however typically require unrealistically wide networks for optimization where the width grows poly in n and poly-inversely proportional to the distance between the input samples. Example distance measures are class margin for logistic loss and minimum eigenvalue of the NTK matrix for least-squares. Our work circumvents this issue by allowing a capacity-dependent interpolation. We prove that even 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 2.3).
Finally, low-rankness of the Jacobian plays a central role in this work. Hessian and Jacobian of neural nets are investigated by multiple papers which contain related findings on the bimodal (approximately low-rank) spectrum Papyan (2018); Ghorbani et al. (2019b); Papyan (2019b); Sagun et al. (2017); Li et al. (2019b); Javadi et al. (2019). Our key empirical contribution is establishing (in great detail) that multiclass Jacobian adapts over time to align its information space with the labels to better represent the data. This alignment leads to tighter generalization bounds in our analysis shedding light on representation learning and gradient dynamics beyond NTK.
4 NUMERICAL EXPERIMENTS
We present experiments demonstrating our theoretical findings on two popular image classification datasets. In this section we focus on a set of CIFAR-10 experiments and discuss how our theory is strongly supported by what we observe in practice. To provide more detail and show that our theory holds across different datasets, in addition to the experiments discussed in this section we perform additional experiments on a modified 3-class version of CIFAR-10 and MNIST in Appendix C.
Experimental setup. The CIFAR-10 dataset consists of 50k training images and 10k test images in 10 classes. 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) and relatively few parameters (0.27M ). 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 Kn ≈ 500000) of the neural network Jacobian. In order to be consistent with our theoretical formulation we make the following modifications to the default ResNet20 architecture: (1) we scale the output of the final fully connected layer to ensure that the output is small, consistent with Theorem 2.2 (2) we turn off batch normalization and (3) we do not pass the network output through a soft-max function. We train the network with SGD on least-squares loss with batch size 128 and without any form of data augmentation. We set the initial learning rate to 0.1 and adjust 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 conduct two sets of experiments to illustrate the results on
uncorrupted and corrupted data. In this section we highlight some of these results and relate them to our theoretical framework. For the complete set of experiments we refer the reader to Appendix C. Jacobian eigenstructure. Calculating the exact full singular value decomposition of the Jacobian at this scale (500k × 270k) is not tractable due to computation/memory limitations. In order to verify the bimodal structure of the Jacobian with exact singular values we plot the histogram of the top 1000 singular values of the Jacobian mapping at initalization and after training in Figure 4. This figure clearly demonstrates that the Jacobian has low-rank structure. In both cases we observe that singular values are concentrated around zero with a relatively small density distributed over higher singular values. This observation serves as a natural basis for decomposition of the label space into information I (large singular values, low-dimensional) and nuisance space N (small singular values, high-dimensional). We note that while calculating all the eigenvalues is not possible, we verify the bimodal structure of the entire Jacobian spectrum by approximating its spectral density in App. C. Experiments without label corruption. First, we present experiments on the original training data described above with no label corruption. We train the network for 400 epochs to achieve a good fit to the training data. Our theory predicts that the sum of ∥J†Iy∥`2 and ∥ΠN (y)∥`2 determines the classification error (Theorems 2.2 and 6.24). Table 1 collects these values for the initial and final Jacobian. These values demonstrate that the label vector is indeed correlated with the top eigenvectors of both the initial and final Jacobians. An interesting aspect of these results is that this correlation increases from the initial to the final Jacobian so that more of the label energy lies on the information space of the final Jacobian in comparison with the initial Jacobian. Stated differently, we observe a significant adaptation of the Jacobian to the labels after training compared to the initial Jacobian so that our predictions become more and more accurate as the iterates progress. In particular, the first column of Table 1 shows that the fraction of label energy lying on the information subspace of the Jacobian drastically increases after training (from 0.38 to 0.99). Consequentially, less energy falls on the nuisance space (decreases from 0.92 to 0.16 after training), while ∥J†Iy∥`2 remains relatively small resulting in better generalization. Towards explaining these, Section D provides a preliminary analysis showing Jacobian spectrum indeed adapts to data. We also track the projection of the residual rτ on the information and nuisance subspaces throughout training on both training and test data and depict the results in Figures 5a and 5b. In agreement with our theory, these plots show that learning on I is fast and the residual energy decreases rapidly on this space. On the other hand, residual energy on N goes down rather slowly and the decrease in total residual energy is overwhelmingly governed by I, suggesting that most information relevant to learning lies in this space. We also plot the training and test error in Figure 5c. We observe that as learning progresses, the residual on both spaces decrease in tandem with training and test error.
Experiments with label corruption. Our next experiments 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 epoch 760 to fit to the training data.
0 0.25 0.5 0.75 1
0.2
0.4
0.6
0.8
Label corruption
∥ΠN (y)∥`2∥y∥`2 final Test error
1Figure 7: Fraction of the energy of the label vector that lies on the nuisance space of the initial Jacobian and final Jacobian as well as the test error as a function of the amount of label corruption. We again track the projection of the residual rτ on the information/nuisance spaces throughout training on both training and test data and depict the results in Figs 6a and 6b. We also track the train and test errors in Figure 6c. From Figure 6c it is evident that while the training error steadily decreases, test error exhibits a very different behavior compared to 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 6a). The lowest test error is observed at epoch 50 after which a steady increase follows. In the second phase, the network overfits to the corrupted data resulting in more test error on the uncorrupted test data (Figure 6b). 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. In particular the residual onN slowly increases while residual on I drops sharply creating a dip in both test error and total residual energy around epoch 50. This phenomenon is further explained in the appendix (see Sec. 5.1) via a linear model.
In Table 2 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 predicted by our theory).
To demonstrate the connection between generalization and information/nuisance spaces, we repeat the experiment with 25%, 75% and 100% label corruption and depict the results after 800 epochs in
Fig. 7. As expected, the test error increases with the corruption. 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.

heading: 5 TECHNICAL APPROACH AND GENERAL THEORY

heading: 5.1 PRELUDE: FITTING A LINEAR MODEL
text: To gain better insights into what governs the generalization capability of gradient based iterations let us consider the simple problem of fitting a linear model via gradient descent. This model maps an input/feature vector x ∈ Rd into a one-dimensional output/label via x↦ f(x,w) ∶=wTx. We wish to fit a model of this form to n training data consisting of input/label pairs {(xi, yi)}ni=1 ∈ Rd × R. Aggregating this training data as rows of a feature matrix X ∈ Rn×d and label vector y ∈ Rn, the training problem takes the form
L(w) = 1 2 ∥Xw − y∥2`2 . (5.1)
We focus on an overparameterized model where there are fewer training data than the number of parameters i.e. n ≤ d. We assume the feature matrix can be decomposed into the formX =X +Z whereX is low-rank (i.e. rank(X) = r << n) with singular value decompositionX = UΣV T with U ∈ Rn×r, Σ ∈ Rr×r, V ∈ Rd×r, and Z ∈ Rn×d is a matrix with i.i.d. N (0, σ2x/n) entries. We shall also assume the labels are equal to y = y + z with y =Xw∗ for somew∗ ∈ Range(V ) and z ∈ Rn a Gaussian random vector with i.i.d. N (0, σ2y/n) entries. One can think of this as a linear regression model where the features and labels are corrupted with Gaussian noise. The goal of course is to learn a model which fits to the clean uncorrupted data and not the corruption. In this case the population loss (i.e. test error) takes the form
E [L(w)] = 1 2 ∥Xw − y∥2 `2 + 1 2 σ2x ∥w∥2`2 + 12σ2y,
Now let us consider gradient descent iterations with a step size η which take the form wτ+1 =wτ − η∇L(wτ) = (I − ηXTX)wτ + ηXTy. (5.2) To gain further insight into the generalization capabilities of the gradient descent iterations we shall consider an instance of this problem where the subspaces U and V are chosen uniformly at random, Σ = Ir with n = 200, d = 500, r = 5, and σx = 0.2, σy = 2. In Figure 8a 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 misfitXw − y on the column space of the uncorrupted portion of the input data (U) and its complement. That is,
EL(w) = ELI(w) + ELN (w).
where
ELI(w) ∶=E [ ∥ΠI (Xw − y)∥2`2 ] = ∥Xw − y∥2`2 + r2nσ2x ∥w∥2`2 + r2nσ2y, ELN (w) ∶=E [ ∥ΠN (Xw − y)∥2`2 ] = 12 (1 − rn)(σ2x ∥w∥2`2 + σ2y) ,
with ΠI = UUT and ΠN = I −UUT . In Figure 8b we plot these two components. This plot clearly shows that ELI(w) goes down quickly while ELN (w) slowly increases with their sum creating the dip in the test error. Since U is a basis for the range of the uncorrupted portion of the features (X) one can think of span(U) as the “information" subspace and ELI(w) as the test error on this information subspace. Similarly, one can think of the complement of this subspace as the “nuisance" subspace and ELN (w) as the test error on this nuisance subspace. Therefore, one can interpret Figure 8a as the test error decreasing rapidly in the first few iterations over the information subspace but slowly increasing due to the contributions of the nuisance subspace.
To help demystify this behavior note that using the gradient descent updates from (5.2) the update in terms of the misfit/residual rτ =Xwτ − y takes the form rτ+1 = (I − ηXXT )rτ = (I − ηX XT ) (Xwτ − y) + noise Based on the form of this update when the information subspace is closely aligned with the prominent singular vectors ofX the test error on the information subspace (ELI(w) ≈ ∥Xwτ − y∥2`2 ) quickly decreases in the first few iterations. However, the further we iterate the parts of the residual aligned with the less prominent eigen-directions ofX (which correspond to the nuisance subspace) slowly pick up more energy contributing to a larger total test error. In this section, we outline our approach to proving robustness of over-parameterized neural networks. Towards this goal, we consider a general formulation where we aim to fit a general nonlinear model of the form x ↦ f(x;θ) with x ∈ Rd denoting the input features, θ ∈ Rp denoting the parameters, and f(x;θ) ∈ RK the K outputs of the model denoted by f1(x;θ), f2(x;θ), . . . , fK(x;θ). For instance in the case of neural networks θ represents its weights. Given a data set of n input/label pairs {(xi,yi)}ni=1 ⊂ Rd × RK , we fit to this data by minimizing a nonlinear least-squares loss of the form
L(θ) = 1 2 n∑ i=1 ∥f(xi;θ) − yi∥2`2 . (5.3)
To continue let us first aggregate the predictions and labels into larger vectors based on class. In particular define
f`(θ) = ⎡⎢⎢⎢⎢⎣ f`(x1;θ)⋮ f`(xn;θ) ⎤⎥⎥⎥⎥⎦ ∈ R n and y(`) = ⎡⎢⎢⎢⎢⎣ (y1)`⋮(yn)` ⎤⎥⎥⎥⎥⎦ ∈ R
n for ` = 1,2, . . . ,K. Concatenating these vectors we arrive at
f(θ) = ⎡⎢⎢⎢⎢⎣ f1(θ)⋮ fK(θ) ⎤⎥⎥⎥⎥⎦ ∈ R Kn and y = ⎡⎢⎢⎢⎢⎢⎣ y(1)⋮ y(K) ⎤⎥⎥⎥⎥⎥⎦ ∈ RKn. (5.4)
Using the latter we can rewrite the optimization problem (5.3) into the more compact form
L(θ) = 1 2 ∥f(θ) − y∥2`2 . (5.5)
To solve this problem we run gradient descent iterations with a learning rate η starting from an initial point θ0. These iterations take the form θτ+1 = θτ − η∇L(θτ) with ∇L(θ) = J T (θ) (f(θ) − y) . (5.6) 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 J (formally defined below) that is close to the Jacobian at initialization J (θ0).
Definition 5.1 (Reference Jacobian and its SVD) Consider an initial point θ0 ∈ Rp and the Jacobian mapping J (θ0) ∈ RKn×p. For ε0, β > 0, we call J ∈ RKn×max(Kn,p) an (ε0, β) reference Jacobian matrix if it obeys the following conditions, ∥J∥ ≤ β, ∥J (θ0)J T (θ0) − JJT ∥ ≤ ε20, and ∥J (θ0) − J∥ ≤ ε0. where J (θ0) ∈ RKn×max(Kn,p) is a matrix obtained by augmenting J (θ0) with max(0,Kn − p) zero columns. Furthermore, consider the singular value decomposition of J given by
J = Udiag(λ)V T = Kn∑ s=1λsusv T s . (5.7)
where λ ∈ RKn are the vector of singular values and us ∈ RKn and vs ∈ Rp are the left/right singular vectors.
One natural choice for this reference Jacobian is J = J (θ0). However, we shall also use other reference Jacobians in our results. We will compare the gradient iterations (5.6) to the iterations associated with fitting a linearized model around θ0 defined as flin(θ̃) = f(θ0) + J(θ̃ − θ0), where θ0 ∈ Rmax(Kn,p) is obtained from θ0 by adding max(Kn − p,0) zero entries at the end of θ0. The optimization problem for fitting the linearized problem has the form
Llin(θ) = 1 2 ∥flin(θ) − y∥2`2 . (5.8)
Thus starting from θ̃0 = θ0 the iterates θ̃τ on the linearized problem take the form θ̃τ+1 = θ̃τ − η∇Llin(θ̃τ), (5.9)= θ̃τ − ηJT (f(θ0) + J(θ̃τ − θ0) − y),= θ̃τ − ηJTJ(θ̃τ − θ0) − ηJT (f(θ0) − y) .
The iterates based on the linearized problem will provide a useful reference to keep track of the evolution of the original iterates (5.6). Specifically we study the evolution of misfit/residuals associated with the two problems
Original residual: rτ = f(θτ) − y. (5.10) Linearized residual: r̃τ = flin(θ̃τ) − y = (I − ηJJT )τr0. (5.11)
To better understand the dynamics of convergence of the linearized iterates next we define two subspaces associated with the reference Jacobian and its spectrum.
Definition 5.2 (Information/Nuisance Subspaces) Let J denote the reference Jacobian per Definition 5.1 with eigenvalue decomposition J = Udiag(λ)V T per (5.7). For a spectrum cutoff α obeying 0 ≤ α ≤ λ1 let r(α) denote the index of the smallest singular value above the threshold α, that is, r(α) = min ({s ∈ {1,2, . . . , nK} such that λs ≥ α}) . We define the information and nuisance subspaces associated with J as I ∶= span({us}rs=1) andN ∶= span({us}Kns=r+1). We also define the truncated reference Jacobian JI = [u1 u2 . . . ur]diag (λ1, λ2, . . . , λr) [v1 v2 . . . vr]T which is the part of the reference Jacobian that acts on the information subspace 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.
Assumption 1 (Bounded spectrum) For any θ ∈ Rp the Jacobian mapping associated with the nonlinearity f ∶ Rp ↦ Rn has bounded spectrum, i.e. ∥J (θ)∥ ≤ β.
Assumption 2 (Bounded perturbation) Consider a point θ0 ∈ Rp and positive scalars ε,R > 0. Assume that for any θ obeying ∥θ − θ0∥`2 ≤ R, we have∥J (θ) −J (θ0)∥ ≤ ε
2 .
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. Theorem 5.3 (Meta Theorem) Consider a nonlinear least squares problem of the form L(θ) = 1 2 ∥f(θ) − y∥2`2 with f ∶ Rp ↦ RnK the multi-class nonlinear mapping, θ ∈ Rp the parameters of
the model, and y ∈ RnK the concatenated labels as in (5.4). Let θ be zero-padding of θ till size max(Kn,p). Also, consider a point θ0 ∈ Rp with J an ( 0, β) reference Jacobian associated withJ (θ0) per Definition 5.1 and fitting the linearized problem flin(θ̃) = f(θ0) + J(θ̃ − θ0) via the lossLlin(θ) = 12 ∥flin(θ) − y∥2`2 . Furthermore, define the information I and nuisance N subspaces and the truncated Jacobian JI associated with the reference Jacobian J based on a cut-off spectrum value of α per Definition 5.2. Furthermore, assume the Jacobian mapping J (θ) ∈ RnK×p associated with f obeys Assumptions 1 and 2 for all θ ∈ Rp obeying ∥θ − θ0∥`2 ≤ R ∶= 2(∥J†Ir0∥`2 + Γα ∥ΠN (r0)∥`2 + δΓα ∥r0∥`2) , (5.12) around a point θ0 ∈ Rp for a tolerance level δ obeying 0 < δ ≤ 1 and stopping time Γ obeying Γ ≥ 1. Finally, assume the following inequalities hold
ε0 ≤ min(δα, √ δα3/Γβ)
5 and ε ≤ δα3 5Γβ2 . (5.13)
We run gradient descent iterations of the form θτ+1 = θτ − η∇L(θτ) and θ̃τ+1 = θ̃τ − η∇Llin(θ̃τ) on the original and linearized problems starting from θ0 with step size η obeying η ≤ 1/β2. Then for all iterates τ obeying 0 ≤ τ ≤ T ∶= Γ
ηα2 the iterates of the original (θτ ) and linearized (θ̃τ ) problems
and the corresponding residuals rτ ∶= f(θτ) − y and r̃τ ∶= flin(θ̃τ) − y closely track each other. That is, ∥rτ − r̃τ∥`2 ≤ 35 δαβ ∥r0∥`2 and ∥θτ − θ̃τ∥`2 ≤ δΓα∥r0∥`2 (5.14) Furthermore, for all iterates τ obeying 0 ≤ τ ≤ T ∶= Γ
ηα2
∥θτ − θ0∥`2 ≤ R2 = ∥J†Ir0∥`2 + Γα ∥ΠN (r0)∥`2 + δΓα ∥r0∥`2 . (5.15) and after τ = T iteration we have ∥rT ∥`2 ≤ e−Γ ∥ΠI(r0)∥`2 + ∥ΠN (r0)∥`2 + δαβ ∥r0∥`2 . (5.16)

heading: 6 PROOFS
text: Before we proceed with the proof let us briefly discuss some notation used throughout. For a matrix W ∈ Rk×d we use vect(W ) ∈ Rkd to denote a vector obtained by concatenating the rows w1,w2, . . . ,wk ∈ Rd of W . That is, vect(W ) = [wT1 wT2 . . . wTk ]T . Similarly, we use mat(w) ∈ Rk×d to denote a k × d matrix obtained by reshaping the vector w ∈ Rkd across its rows. Throughout, for a differentiable function φ ∶ R↦ R we use φ′ and φ′′ to denote the first and second derivative.

heading: 6.1 PROOFS FOR GENERAL THEORY (PROOF OF THEOREM 5.3)
text: 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 r̃τ . 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.

heading: 6.1.1 PRELIMINARY DEFINITIONS AND LEMMAS
text: Throughout we use
UI = [u1 u2 . . . ur] ∈ RnK×r and UN = [ur+1 ur+2 . . . unK] ∈ RnK×(nK−r). to denote the basis matrices for the information and nuisance subspaces from Definition 5.2. Similarly, we define the information and nuisance spectrum as λI = [λ1 λ2 . . . λr]T and λN = [λr+1 λr+2 . . . λnK]T . We also define the diagonal matrices Λ = diag(λ), ΛI = diag(λI), and ΛN = diag(λN ). Definition 6.1 (early stopping value and distance) Consider Definition 5.2 and let Γ > 0 be a positive scalar. Associated with the initial residual r0 = f(θ0) − y and the information/nuisance subspaces of the reference Jacobian J (with a cut-off level α) we define the (α,Γ) early stopping value as
Bα,Γ =( r∑ s=1 α2 λ2s (⟨us,r0⟩)2 + Γ2 nK∑ s=r+1 λ2s α2 (⟨us,r0⟩)2)1/2 . (6.1) We also define the early stopping distance as
Dα,Γ = Bα,Γ α .
The goal of early stopping value/distance is understanding the behavior of the algorithm at a particular stopping time that depends on Γ and the spectrum cutoff α. 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. Lemma 6.2 (Bounds on Early-Stopping Value) The early stopping value Bα,Γ from Definition 6.1 obeys
Bα,Γ ≤ (∥ΠI(r0)∥2`2 + Γ2∥ΠN (r0)∥2`2)1/2 ≤ Γ ∥r0∥`2 (6.2)Bα,Γ ≥ α λ1
∥ΠI(r0)∥`2 . (6.3) Proof To prove the upper bound we use the fact that α ≤ λs for s ≤ r and α ≥ λs for s ≥ r to conclude that
Bα,Γ ≤ ( r∑ s=1 (⟨us,r0⟩)2 + Γ2 nK∑ s=r+1 (⟨us,r0⟩)2)
1/2
= (∥ΠI(r0)∥2`2 + Γ2 ∥ΠN (r0)∥2`2)1/2≤ Γ∥r0∥`2 . To prove the lower bound, we use the facts that α2/λ2s ≥ α2/λ21 to conclude that
Bα,Γ =( r∑ s=1 α2 λ2s (⟨us,r0⟩)2 + Γ2 nK∑ s=r+1 λ2s α2 (⟨us,r0⟩)2)1/2 , ≥( r∑
s=1 α2 λ2s (⟨us,r0⟩)2)1/2 ,
≥ α λ1 ∥ΠI(r0)∥`2 .
It is of course well known that the mapping (I − ηAAT ) is a contraction for sufficiently small values of η. The next lemma shows that if we replace one of theA matrices with a matrixB which is close toA the resulting matrix (I − ηABT ), while may not be contractive, is not too expansive.
Lemma 6.3 (Asymmetric PSD increase) LetA,B ∈ Rn×p be matrices obeying ∥A∥ ≤ β, ∥B∥ ≤ β, and ∥B −A∥ ≤ ε. Then, for all r ∈ Rn and η ≤ 1/β2 we have ∥(I − ηABT )r∥
`2 ≤ (1 + ηε2) ∥r∥`2 .
Proof Note that using η ≤ 1/β2 and ∥B −A∥ ≤ ε we conclude that ∥(I − ηABT )r∥2
`2 = ∥(I − ηBBT − η(A −B)BT )r∥2
`2= ∥r − η(A −B +B)BTr∥2 `2= ∥r∥2`2 − 2ηrT (A −B +B)BTr + η2 ∥ABTr∥2`2≤ ∥r∥2`2 − 2η∥BTr∥2`2 + 2η ∥(A −B)Tr∥`2 ∥BTr∥`2 + η2∥A∥2∥BTr∥2`2= ∥r∥2`2 − η∥BTr∥2`2 + 2η ∥(A −B)Tr∥`2 ∥BTr∥`2 + (η2∥A∥2∥BTr∥2`2 − η∥BTr∥2`2)
η≤1/β2≤ ∥r∥2`2 − η∥BTr∥2`2 + 2η ∥(A −B)Tr∥`2 ∥BTr∥`2∥A−B∥≤ ≤ ∥r∥2`2 − η∥BTr∥2`2 + 2ηε∥BTr∥`2∥r∥`2 =(1 + ηε2) ∥r∥2`2 − η (ε ∥r∥`2 − ∥BTr∥`2)2≤(1 + ηε2) ∥r∥2`2 ,
completing the proof.
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. Lemma 6.4 Let A and B be n × n positive semi-definite matrices satisfying ∥A −B∥ ≤ α2 for a scalar α ≥ 0. Then for anyX ∈ Rn×p with p ≥ n obeyingA =XXT , there exists a matrix Y ∈ Rn×p obeyingB = Y Y T such that ∥Y −X∥ ≤ 2α Proof First we note that for any two PSD matricesA+,B+ ∈ Rn×n obeyingA+,B+ ⪰ α24 In, Lemma 2.2 of Schmitt (1992) guarantees that
∥A1/2+ −B1/2+ ∥ ≤ ∥A+ −B+∥ α .
In the above for a PSD matrix A ∈ Rn×n with an eigenvalue decomposition A = UΛUT we use A1/2 to denote the square root of the matrix given byA = UΛ1/2UT . We shall use this result with A+ =A + α24 In andB+ =B + α24 In to conclude that
∥A1/2+ −B1/2+ ∥ ≤ ∥A+ −B+∥ α = ∥A −B∥ α .
Furthermore, using the fact that the eigenvalues of A+ and B+ are just shifted versions of the eigenvalues ofA andB by α2/4 we can conclude that
∥A1/2+ −A1/2∥ ≤ α 2 and ∥B1/2+ −B1/2∥ ≤ α 2 .
Combining the latter two inequalities with the assumption that ∥A −B∥ ≤ α2 we conclude that ∥A1/2 −B1/2∥ ≤ ∥A1/2+ −B1/2+ ∥ + ∥A1/2+ −A1/2∥ + ∥B1/2+ −B1/2∥
≤∥A −B∥ α + α 2 + α
2≤2α. (6.4)
Suppose p ≥ n and assume the matrices A and B have eigenvalue decompositions given by A = UAΛAU T A and B = UBΛBUTB . Then, any X ∈ Rn×p with p ≥ n has the form X = UAΛ1/2A V TA
with VA ∈ Rp×n an orthonormal matrix. Now pick Y = UBΛ1/2B UTBUAV TA . Then clearly Y Y T =B. Furthermore, we have ∥X −Y ∥ = ∥UAΛ1/2A V TA −UBΛ1/2B UTBUAV TA ∥
= ∥UAΛ1/2A UTAUAV TA −UBΛ1/2B UTBUAV TA ∥ = ∥(UAΛ1/2A UTA −UBΛ1/2B UTB)UAV TA ∥ = ∥(A1/2 −B1/2)UAV TA ∥= ∥A1/2 −B1/2∥ .
Combining the latter with (6.4) completes the proof.

heading: 6.1.2 KEY LEMMAS FOR GENERAL NONLINEARITIES
text: Throughout this section we assume J is the reference Jacobian per Definition 5.1 with eigenvalue decomposition J = UΛV T = ∑Kns=1 λsusvTs with Λ = diag(λ). We also define a = UTr0 = UT r̃0 ∈ RnK be the coefficients of the initial residual in the span of the column space of this reference Jacobian.
We shall first characterize the evolution of the linearized parameter θ̃τ and residual r̃τ vectors from (5.11) in the following lemma.
Lemma 6.5 The linearized residual vector r̃τ can be written in the form
r̃τ = U (I − ηΛ2)τ a = nK∑ s=1(1 − ηλ2s)τasus. (6.5)
Furthermore, assuming η ≤ 1/λ21 the linear updates θ̃τ obey ∥θ̃τ − θ̃0∥2`2 ≤ r∑
s=1 a2s λ2s + τ2η2 nK∑ s=r+1λ 2 sa 2 s. (6.6)
Proof Using the fact that JJT = UΛ2UT we have (I − ηJJT )τ = U (I − ηΛ2)τ UT
Using the latter combined with (5.11) we thus have
r̃τ =(I − ηJJT )τr0,=U (I − ηΛ2)τ UTr0, =U (I − ηΛ2)τ a, =nK∑ s=1(1 − ηλ2s)τasus,
completing the proof of (6.5).
We now turn our attention to proving (6.6) by tracking the representation of θ̃τ in terms of the right singular vectors of J . To do this note that using (6.5) we have JT r̃t = V ΛUT r̃t = V Λ (I − ηΛ2)t a. Using the latter together with the gradient update on the linearized problem we have
θ̃τ − θ̃0 = − η (τ−1∑ t=0 ∇Llin(θ̃t)) = −η ( τ−1∑ t=0 J T r̃t) = −ηV (τ−1∑ t=0 Λ (I − ηΛ2) t)a.
Thus for any s ∈ {1,2, . . . , nK} vTs (θ̃τ − θ̃0) = −ηλsas (τ−1∑
t=0 (1 − ηλ2s) t) = −ηλsas 1 − (1 − ηλ2s) τ ηλ2s = −as 1 − (1 − ηλ2s) τ λs .
Noting that for η ≤ 1/λ21 ≤ 1/λ2s we have 1 − ηλ2s ≥ 0, the latter identity implies that ∣vTs (θ̃τ − θ̃0)∣ ≤ ∣as∣λs . (6.7) Furthermore, using the fact that 1 − ηλ2s ≤ 1 we have ∣vTs (θ̃τ − θ̃0)∣ = ηλs ∣as∣ (τ−1∑
t=0 (1 − ηλ2s) t) ≤ ηλs ∣as∣ τ (6.8)
Combining (6.7) for 1 ≤ s ≤ r and (6.8) for s > r we have ∥θ̃τ − θ̃0∥2`2 = nK∑ s=1 ∣vTs (θ̃τ − θ̃0)∣ 2 ≤ r∑ s=1 a2s λ2s + τ2η2 nK∑ s=r+1λ 2 sa 2 s,
completing the proof of (6.6).
For future use we also state a simple corollary of the above Lemma below.
Corollary 6.6 Consider the setting and assumptions of Lemma 6.5. Then, after τ iterations we have ∥r̃τ∥`2 ≤ (1 − ηα2)τ ∥ΠI(r0)∥`2 + ∥ΠN (r0)∥`2 . (6.9) Furthermore, after T = Γ
ηα2 iterations we have
∥r̃T ∥`2 ≤ e−Γ∥ΠI(r0)∥`2 + ∥ΠN (r0)∥`2 . (6.10) and
∥θ̃T − θ̃0∥2`2 ≤ r∑ s=1 a2s λ2s + Γ2 nK∑ s=r+1 λ2sa 2 s α4 = B2α,Γ α2 .
with Bα,Γ 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 UTI r̃τ = (I − ηΛ2I)τ UTI r̃0 and UTN r̃τ = (I − ηΛ2N )τ UTN r̃0 Thus, using the fact that for s ≤ r we have λs ≥ α we have (1 − ηλ2s)τ ≤ (1 − ηα2)τ and for s > r we have (1 − ηλ2s)τ ≤ 1, we can conclude that ∥UTI r̃τ∥`2 ≤ (1 − ηα2)τ ∥UTI r0∥`2 and ∥UTN r̃τ∥`2 ≤ ∥UTNr0∥`2 . Combining these with the triangular inequality we have
∥r̃τ∥`2 = ∥[UTI r̃τUTN r̃τ]∥`2 ≤ ∥U TI r̃τ∥`2 + ∥UTN r̃τ∥`2 ≤ (1 − ηα2)τ ∥UTI r0∥`2 + ∥UTNr0∥`2 ,
concluding the proof of (6.9). The second bound on the residual simply follows from the fact that(1 − ηα2)T ≤ e−Γ. The bound on ∥θ̃T − θ̃0∥2`2 is trivially obtained by using T 2 = Γ2η2α4 in (6.6). The lemma above shows that with enough iterations, gradient descent on the linearized problem fits the residual over the information space and the residual is (in the worst case) unchanged over the nuisance subspaceN . Our hypothesis is that, when the model is generalizable the residual mostly lies on the information space I which contains the directions aligned with the top singular vectors. Hence, the smaller term ∥ΠN (r0)∥`2 over the nuisance space will not affect generalization significantly. To make this intuition precise however we need to connect the residual of the original problem to that of the linearized problem. The following lemma sheds light on the evolution of the original problem (5.6) by characterizing the evolution of the difference between the residuals of the original and linearized problems from one iteration to the next.
Lemma 6.7 (Keeping track of perturbation - one step) Assume Assumptions 1 and 2 hold and θτ and θτ+1 are within an R neighborhood of θ0, that is,∥θτ − θ0∥`2 ≤ R and ∥θτ+1 − θ0∥`2 ≤ R. Then with a learning rate obeying η ≤ 1/β2, the deviation in the residuals of the original and linearized problems eτ+1 = rτ+1 − r̃τ+1 obey∥eτ+1∥`2 ≤ η(ε20 + εβ)∥r̃τ∥`2 + (1 + ηε2)∥eτ∥`2 . (6.11) Proof For simplicity, denoteB1 = J (θτ+1,θτ),B2 = J (θτ),A = J (θ0) where
J (b,a) = ∫ 1 0 J (tb + (1 − t)a)dt.
We can write the predictions due to θτ+1 as f(θτ+1) = f(θτ − η∇L(θτ)) = f(θτ) + ηJ (θτ+1,θτ)∇L(θτ)= f(θτ) + ηJ (θτ+1,θτ)J T (θτ)(f(θτ) − y).
This implies that rτ+1 = f(θτ+1) − y = (I − ηB1BT2 )rτ . Similarly, for linearized problem we have r̃τ+1 = (I − ηJJT )r̃τ . Thus, ∥eτ+1∥`2 = ∥(I − ηB1BT2 )rτ − (I − ηJJT )r̃τ∥`2= ∥(I − ηB1BT2 )eτ − η(B1BT2 − JJT )r̃τ∥`2≤ ∥(I − ηB1BT2 )eτ∥`2 + η∥(B1BT2 − JJT )r̃τ∥`2≤ ∥(I − ηB1BT2 )eτ∥`2 + η ∥(B1BT2 − JJT )∥ ∥r̃τ∥`2 . (6.12)
We proceed by bounding each of these two terms. For the first term, we apply Lemma 6.3 with A =B1 andB =B2 and use ∥B1 −B2∥ ≤ ε to conclude that ∥(I − ηB1BT2 )eτ∥`2 ≤ (1 + ηε2)∥eτ∥`2 . (6.13) Next we turn our attention to bounding the second term. To this aim note that
∥B1BT2 − JJT ∥ = ∥B1BT2 −AAT +AAT − JJT ∥≤ ∥B1BT2 −AAT ∥ + ∥AAT − JJT ∥≤ ∥(B1 −A)BT2 ∥ + ∥A(B2 −A)T ∥ + ∥AAT − JJT ∥≤ ∥B1 −A∥∥B2∥ + ∥B2 −A∥∥A∥ + ∥AAT − JJT ∥ ≤ β ε
2 + β ε 2 + ε20
= ε20 + εβ. (6.14) In the last inequality we use the fact that per Assumption 2 we have ∥B1 −A∥ ≤ ε/2 and ∥B2 −A∥ ≤ ε/2 as well as the fact that per Definition 5.1 ∥AAT − JJT ∥ ≤ ε20. Plugging (6.13) and (6.14) in (6.12) completes the proof.
Next we prove a result about the growth of sequences obeying certain assumptions. As we will see later on in the proofs this lemma allows us to control the growth of the perturbation between the original and linearized residuals (eτ = ∥eτ∥`2 ). Lemma 6.8 (Bounding residual perturbation growth for general nonlinearities) Consider positive scalars Γ, α, ε, η > 0. Also assume η ≤ 1/α2 and α ≥ √2Γε and set T = Γ
ηα2 . Assume the scalar
sequences eτ (with e0 = 0) and r̃τ obey the following identities r̃τ ≤(1 − ηα2)τρ+ + ρ−, eτ ≤(1 + ηε2)eτ−1 + ηΘr̃τ−1, (6.15) for all 0 ≤ τ ≤ T and non-negative values ρ−, ρ+ ≥ 0. Then, for all 0 ≤ τ ≤ T , eτ ≤ ΘΛ holds with Λ = 2(Γρ− + ρ+)
α2 . (6.16)
Proof We shall prove the result inductively. Suppose (6.16) holds for all t ≤ τ − 1. Consequently, we have
et+1 ≤(1 + ηε2)et + ηΘr̃t≤et + ηε2et + ηΘ ((1 − ηα2)tρ+ + ρ−)≤et + ηΘ (ε2Λ + (1 − ηα2)tρ+ + ρ−) . Thus
et+1 − et Θ
≤ η (ε2Λ + (1 − ηα2)tρ+ + ρ−) . (6.17) Summing up both sides of (6.17) for 0 ≤ t ≤ τ − 1 we conclude that
eτ Θ = τ−1∑ t=0 et+1 − et Θ
≤ ητ (ε2Λ + ρ−) + ηρ+ τ−1∑ t=0(1 − ηα2)t = ητ (ε2Λ + ρ−) + ηρ+ 1 − (1 − ηα2) τ
ηα2
≤ η (τε2Λ + ρ+ ηα2 + τρ−) = ητ(ε2Λ + ρ−) + ρ+ α2≤ ηT (ε2Λ + ρ−) + ρ+ α2 = Γε2Λ + Γρ− + ρ+ α2 = Γε2Λ α2
+ Λ 2≤ Λ,
where in the last inequality we used the fact that α2 ≥ 2Γε2. This completes the proof of the induction step and the proof of the lemma.

heading: 6.1.3 COMPLETING THE PROOF OF THEOREM 5.3
text: 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 τ − 1 ≤ T − 1. In particular, we assume the identities (5.14) and (5.15) hold for all 0 ≤ t ≤ τ − 1. We aim to prove these identities continue to hold for iteration τ . We will prove this result in multiple steps. Step I: Next iterate obeys ∥θτ − θ0∥`2 ≤ R. We first argue that θτ lies in the domain of interest as dictated by (5.12), i.e. ∥θτ − θ0∥`2 ≤ R. To do this note that per the induction assumption (5.15) holds for iteration τ − 1 and thus ∥θτ−1 − θ0∥`2 ≤ R/2. As a result using the triangular inequality to show ∥θτ − θ0∥`2 ≤ R holds it suffices to show
that ∥θτ − θτ−1∥`2 ≤ R/2 holds. To do this note that∥θτ − θτ−1∥`2 =η∥∇L(θτ−1)∥`2=η ∥J T (θτ−1)rτ−1∥`2 =η ∥J T (θτ−1)rτ−1∥
`2(a)≤ η∥J T (θτ−1)r̃τ−1∥`2 + η∥J T (θτ−1)(rτ−1 − r̃τ−1)∥`2 (b)≤ η ∥JT r̃τ−1∥`2 + η ∥J (θτ−1) − J∥ ∥r̃τ−1∥`2 + η ∥J (θτ−1)∥ ∥rτ−1 − r̃τ−1∥`2(c)≤ η ∥JT r̃τ−1∥`2 + ε0 + εβ2 ∥r̃τ−1∥`2 + 1β ∥rτ−1 − r̃τ−1∥`2 (d)≤ η∥JT r̃τ−1∥`2 + 2δα5β2 ∥r0∥`2 + 1β ∥rτ−1 − r̃τ−1∥`2 (e)≤ η∥JT r̃τ−1∥`2 + 2δα5β2 ∥r0∥`2 + 3δα5β2 ∥r0∥`2 =η∥JT r̃τ−1∥`2 + δαβ2 ∥r0∥`2 (f)≤ ηβ2Bα,Γ
α + δα β2 ∥r0∥`2 (g)≤ Bα,Γ
α + δα β2 ∥r0∥`2 (h)≤ Bα,Γ
α + δΓ α ∥r0∥`2 =R
2 .
Here, (a) and (b) follow from a simple application of the triangular inequality, (c) from the fact that∥J (θτ−1) − J∥ ≤ ∥J (θτ−1) −J (θ0)∥ + ∥J (θ0) − J∥ ≤ ε + ε0, (d) from combining the bounds in (5.13), (e) from the induction hypothesis that postulates (5.14) holds for iteration τ − 1, (f) from considering the SVD J = UΛV T which implies that
∥JT r̃τ−1∥2`2 = ∥JT (I − ηJJT )τ−1 r0∥2`2 = ∥V Λ (I − ηΛ2)τ−1UTr0∥ 2 `2
= ∥Λ (I − ηΛ2)τ−1UTr0∥2 `2 =nK∑ s=1λ 2 s(1 − ηλ2s)2(τ−1) (⟨us,r0⟩)2
≤nK∑ s=1λ 2 s (⟨us,r0⟩)2
= r∑ s=1λ 2 s (⟨us,r0⟩)2 + nK∑ s=r+1λ 2 s (⟨us,r0⟩)2
≤β4 r∑ s=1 1 λ2s (⟨us,r0⟩)2 + nK∑ s=r+1λ 2 s (⟨us,r0⟩)2
≤β4 ( r∑ s=1 1 λ2s (⟨us,r0⟩)2 + Γ2 nK∑ s=r+1 λ2s α4 (⟨us,r0⟩)2) =β4 (Bα,Γ
α )2
(g) from the fact that η ≤ 1 β2 , and (h) from the fact that α ≤ β and Γ ≥ 1. Step II: Original and linearized residuals are close (first part of (5.14)). In this step we wish to show that the first part of (5.14) holds for iteration τ . Since we established in
the previous step that ∥θτ − θ0∥`2 ≤ R the assumption of Lemma 6.7 holds for iterations τ − 1 and τ . Hence, using Lemma 6.7 equation (6.11) we conclude that ∥eτ∥`2 ≤ η(ε20 + εβ)∥r̃τ−1∥`2 + (1 + ηε2)∥eτ−1∥`2 . This combined with the induction assumption implies that ∥et∥`2 ≤ η(ε20 + εβ)∥r̃t−1∥`2 + (1 + ηε2)∥et−1∥`2 , (6.18) holds for all t ≤ τ ≤ T . Furthermore, using Lemma 6.5 equation (6.9) for all t ≤ τ ≤ T we have ∥r̃t∥`2 ≤ (1 − ηα2)t ∥ΠI(r0)∥`2 + ∥ΠN (r0)∥`2 , (6.19) To proceed, we shall apply Lemma 6.8 with the following variable substitutions
Θ ∶= ε20 + εβ, ρ+ = ∥ΠI(r0)∥`2 , ρ− = ∥ΠN (r0)∥`2 , eτ ∶= ∥eτ∥`2 , r̃τ ∶= ∥r̃τ∥`2 . (6.20) We note that Lemma 6.8 is applicable since (i) η ≤ 1/β2 ≤ 1/α2, (ii) based on (5.13) we have α ε ≥ 5Γ δ β2 α2 ≥ √2Γ, (iii) τ obeys τ ≤ T = Γ ηα2 , and (iv) (6.15) holds based on (6.18) and (6.19). Thus using Lemma 6.8 we can conclude that
∥eτ∥`2 ≤2(ε20 + εβ)(∥ΠI(r0)∥`2 + Γ∥ΠN (r0)∥`2)α2 ≤2Γ(ε20 + εβ)∥r0∥`2
α2 (6.21)
≤( 2 25 + 2 5 )δα β ∥r0∥`2 ≤ 35 δαβ ∥r0∥`2 , (6.22) where in the last inequality we used (5.13). This completes the first part of (5.14) via induction.
Step III: Original and linearized parameters are close (second part of (5.14)). In this step we wish to show that the second part of (5.14) holds for iteration τ . To do this we begin by noting that by the fact that J is a reference Jacobian we have ∥J (θ0) − J∥ ≤ ε0 whereJ augments J (θ0) by padding zero columns to match size of J . Also by Assumption 2 we have∥J (θ) −J (θ0)∥ ≤ ε2 . Combining the latter two via the triangular inequality we conclude that∥J (θτ) − J∥ ≤ ε0 + ε. (6.23) Let θ and ∇L̄(θ) be vectors augmented by zero padding θ,∇L(θ) so that they have dimension max(Kn,p). Now, we track the difference between θ and linearized θ̃ as follows
∥θτ − θ̃τ∥`2 η = ∥τ−1∑ t=0 ∇L̄(θt) −∇Llin(θ̃t)∥`2
= ∥τ−1∑ t=0 J (θt)Trt − JT r̃t∥`2 ≤ τ−1∑ t=0 ∥J (θt)Trt − JT r̃t∥`2 ≤ τ−1∑ t=0 ∥(J (θt) − J)T r̃t∥`2 + ∥J (θt)T (rt − r̃t)∥`2 = τ−1∑ t=0 ∥(J (θt) − J)T r̃t∥`2 + ∥J (θt)Tet∥`2 ≤ τ−1∑ t=0(ε + ε0)∥r̃t∥`2 + β∥et∥`2 . (6.24)
In the last inequality we used the fact that ∥J (θt) − J∥ ≤ ε + ε0 and ∥J∥ ≤ β. We proceed by bounding each of the two terms in (6.24) above. For the first term we use the fact that ∥r̃τ∥`2 ≤ ∥r0∥`2 to conclude
τ−1∑ t=0 ∥r̃t∥`2 ≤ τ∥r0∥`2 ≤ T ∥r0∥`2 = Γ∥r0∥`2 ηα2 . (6.25)
To bound the second term in (6.24) we use (6.21) together with τ ≤ T ≤ Γ ηα2 to conclude that
τ−1∑ t=0 ∥et∥`2 ≤ τ 2(εβ + ε20) α2 Γ∥r0∥`2 ≤ 2Γ2(εβ + ε20)ηα4 ∥r0∥`2 . (6.26) Combining (6.25) and (6.26) in (6.24), we conclude that
∥θτ − θ̃τ∥`2 ≤(2Γ(εβ2 + ε20β)α3 + ε + ε0α ) Γα∥r0∥`2 =(ε2Γβ2
α3 + ε20 2Γβα3 + ε + ε0α ) Γα∥r0∥`2
(a)≤ (2 5 δ + ε20 2Γβα3 + ε + ε0α ) Γα∥r0∥`2 (b)≤ (2 5 δ + 2 25 δ + ε + ε0 α ) Γ α ∥r0∥`2 (c)≤ (2 5 δ + 2 25 δ + 1 5 δ + ε0 α ) Γ α ∥r0∥`2 (d)≤ (2 5 δ + 2 25 δ + 1 5 δ + 1 5 δ) Γ α ∥r0∥`2
=22 25 δ α Γ∥r0∥`2 .
Here, (a) follows from ε ≤ δα3 5Γβ2 per Assumption (5.13), (b) from ε0 ≤ 15√ δα3Γβ per Assumption (5.13), (c) from ε ≤ δα3 5Γβ2 ≤ δα 5Γ ≤ δα 5 per Assumption (5.13), and (d) from ε0 ≤ δα5 per Assumption (5.13). Thus,
∥θτ − θ̃τ∥`2 ≤ δαΓ∥r0∥`2 . Combining the latter with the fact that ∥θ̃τ − θ0∥`2 ≤ Bα,Γα (which follows from Lemma 6.5 equation (6.6)) we conclude that ∥θτ − θ0∥`2 = ∥θτ − θ0∥`2 ≤ ∥θ̃τ − θ0∥`2 + ∥θτ − θ̃τ∥`2 ≤ Bα,Γα + δαΓ∥r0∥`2 ≤ ∥J†Ir0∥`2 + Γα ∥ΠN (r0)∥`2 + δαΓ∥r0∥`2 The completes the proof of the bound (5.15).
Step V: Bound on residual with early stopping. In this step we wish to prove (5.16). To this aim note that
∥rT ∥`2 (a)≤ ∥r̃T ∥`2 + ∥r̃T − rT ∥`2(b)≤ ∥r̃T ∥`2 + δαβ ∥r0∥`2 (c)≤ e−Γ∥ΠI(r0)∥`2 + ∥ΠN (r0)∥`2 + δαβ ∥r0∥`2
where (a) follows from the triangular inequality, (b) from the conclusion of Step II (first part of (5.14)), and (c) from Corollary 6.6 equation (6.10). This completes the proof of (5.16).

heading: 6.2 KEY LEMMAS AND IDENTITIES FOR NEURAL NETWORKS
text: 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.
For two matrices
A = ⎡⎢⎢⎢⎢⎢⎢⎣ A1 A2⋮ Ap ⎤⎥⎥⎥⎥⎥⎥⎦ ∈ Rp×m and B = ⎡⎢⎢⎢⎢⎢⎢⎣ B1 B2⋮ Bp ⎤⎥⎥⎥⎥⎥⎥⎦ ∈ Rp×n,
we define their Khatri-Rao product asA ∗B = [A1 ⊗B1, . . . ,Ap ⊗Bp] ∈ Rp×mn, where ⊗ denotes the Kronecker product.

heading: 6.2.1 PRELIMINARY IDENTITIES AND CALCULATIONS
text: We begin by discussing a few notations. Throughout we use w` and v` to denote the `th row of input and output weight matrices W and V . Given a matrix M we use ∥M∥2,∞ to denote the largest Euclidean norm of the rows ofM . We begin by noting that for a one-hidden layer neural network of the form x↦ V φ (Wx), the Jacobian matrix with respect to vect(W ) ∈ Rkd takes the form
J (W ) = ⎡⎢⎢⎢⎢⎣ J1(W )⋮JK(W ) ⎤⎥⎥⎥⎥⎦ ∈ R Kn×kd (6.27)
where J`(W ) is the Jacobian matrix associated with the `th class. In particular, J`(W ) is given byJ`(W ) = [J`(w1) . . . J`(wk)] ∈ Rn×kd with J`(ws) ∶= V`,sdiag (φ′(Xws))X. Alternatively using Khatri-Rao products this can be rewritten in the more compact formJ`(W ) = (φ′ (XW T )diag(v`)) ∗X. (6.28) An alternative characterization of the Jacobian is via its matrix representation. Given a vector u ∈ RKn let us partition it into K size n subvectors so that u = [uT1 . . . uTK]T . We have mat (J T (W )u) = K∑̀=1 diag(v`)φ′ (WXT )diag(u`)X. (6.29)

heading: 6.2.2 FUNDAMENTAL PROPERTIES OF THE JACOBIAN OF THE NEURAL NETWORK
text: In this section we prove a few key properties of the Jacobian mapping of a one-hidden layer neural network. Lemma 6.9 (Properties of Single Output Neural Net Jacobian) Let K = 1 so that V T = v ∈ Rn. Suppose φ is an activation obeying ∣φ′(z)∣ ≤ B for all z. Then, for anyW ∈ Rk×d and any unit length vector u, we have ∥J (W )∥ ≤B√k ∥v∥`∞ ∥X∥ and ∥mat (J T (W )u)∥2,∞ ≤B∥v∥`∞∥X∥ (6.30) Furthermore, suppose φ is twice differentiable and ∣φ′′(z)∣ ≤ B for all z. Also assume all data points have unit Euclidean norm (∥xi∥`2 = 1). Then the Jacobian mapping is Lipschitz with respect to spectral norm i.e. for all W̃ ,W ∈ Rk×d we have∥J (W̃ ) −J (W )∥ ≤ B ∥v∥`∞ ∥X∥ ∥W̃ −W ∥F . Proof The result on spectral norm and Lipschitzness of J (W ) have been proven in Oymak & Soltanolkotabi (2019). To show the row-wise bound (6.30), we use (6.29) to conclude that∥mat (J T (W )u)∥2,∞ = ∥diag(v)φ′ (WXT )diag(u)X∥2,∞≤ ∥v∥`∞ max
1≤`≤k ∥φ′ (wT` XT )diag(u)X∥`2≤ ∥v∥`∞∥X∥ max 1≤`≤k ∥φ′ (wT` XT )diag(u)∥`2≤ B∥v∥`∞∥X∥∥u∥`2= B∥v∥`∞∥X∥.
Next we extend the lemma above to the multi-class setting.
Lemma 6.10 (Properties of Multiclass Neural Net Jacobian) Suppose φ is an activation obeying∣φ′(z)∣ ≤ B for all z. Then, for anyW ∈ Rk×d and any unit length vector u, we have ∥J (W )∥ ≤ B√Kk ∥V ∥`∞ ∥X∥ (6.31)
and ∥mat (J T (W )u)∥2,∞ ≤ B√K∥V ∥`∞∥X∥. (6.32) Furthermore, suppose φ is twice differentiable and ∣φ′′(z)∣ ≤ B for all z. Also assume all data points have unit Euclidean norm (∥xi∥`2 = 1). Then the Jacobian mapping is Lipschitz with respect to spectral norm i.e. for all W̃ ,W ∈ Rk×d we have
∥J (W̃ ) −J (W )∥ ≤ B√K ∥V ∥`∞ ∥X∥ ∥W̃ −W ∥F . Proof The proof will follow from Lemma 6.9. First, given A = [AT1 . . . ATK]T and B =[BT1 . . . BTK]T , observe that
∥A∥ ≤ √K sup 1≤`≤K ∥A`∥ and ∥A −B∥ ≤ √ K sup
1≤`≤K ∥A` −B`∥. These two identities applied to the components J`(W ) and J`(W̃ ) −J`(W ) completes the proof of the bound on the spectral norm and the perturbation. To prove the bound in (6.32) we use the identity (6.29) to conclude that
∥mat (J T (W )u)∥2,∞ = ∥ K∑̀=1 diag(v`)φ′ (WXT )diag(u`)X∥2,∞ ≤ K∑̀=1 ∥diag(v`)φ′ (WXT )diag(u`)X∥2,∞ ≤ K∑̀=1B∥V ∥`∞∥X∥∥u`∥`2 = B∥V ∥`∞∥X∥( K∑̀=1 ∥u`∥`2) ≤ B∥V ∥`∞∥X∥√K ( K∑̀=1 ∥u`∥2`2) 1/2
= B∥V ∥`∞∥X∥√K, where the penultimate inequality follows from Cauchy Schwarz, completing the proof.

heading: 6.2.3 PROPERTIES OF THE JACOBIAN AT RANDOM INITIALIZATION
text: In this section we prove a few lemmas characterizing the properties of the Jacobian at the random initialization.
Lemma 6.11 (Multiclass covariance) Given input and output layer weights V and W , consider the Jacobian described in (6.27). Given an Kn×Kn matrixM , for 1 ≤ `, ̃̀≤K, letM[`, ̃̀] denote the (`, ̃̀)th submatrix. For C(W ) = J (W )J (W )T we have
C(W )[`, ̃̀] = k∑ s=1(XXT )⊙ (V`,sṼ̀,sφ′(Xws)φ′(Xws)T ).
SupposeW i.i.d.∼ N (0,1) and V has i.i.d. zero-mean entries with ν2 variance. Then E[C(W )] is a block diagonal matrix given by the Kronecker product E[C(W )] = kν2Σ(X). where Σ(X) is equal to IK ⊗ [(XXT )⊙ E[φ′(Xws)φ′(Xws)T ]].
Proof The (`, ̃̀)th submatrix of C(W ) is given by C(W )[`, ̃̀] = ((diag(v`)φ′(WXT )) ∗XT )((diag(ṽ̀)φ′(WXT )) ∗XT )T
= k∑ s=1J`(ws)J̃̀(ws)T
= k∑ s=1V`,sṼ̀,s(diag(φ′(Xws))X)(diag(φ′(Xws))X)T
= k∑ s=1V`,sṼ̀,s(XXT )⊙ (φ′(Xws)φ′(Xws)T )
= k∑ s=1 (XXT )⊙ (V`,sṼ̀,sφ′(Xws)φ′(Xws)T ) . (6.33)
SettingW i.i.d.∼ N (0,1) and V with i.i.d. zero-mean and ν2-variance entries, we conclude that E[C(W )[`, ̃̀]] = k∑
s=1(XXT )⊙ (E[V`,sṼ̀,s]E[φ′(Xws)φ′(Xws)T ]) = k∑ s=1ν
2δ(` − ̃̀)[(XXT )⊙ E[φ′(Xws)φ′(Xws)T ]] = kδ(` − ̃̀)ν2Σ̃(X),
where δ(x) is the discrete δ function which is 0 for x ≠ 0 and 1 for x = 0 and Σ̃(X) is single output kernel matrix which concludes the proof.
Next we state a useful lemma from Schur (1911) which allows us to bound the eigenvalues of the Hadamard product of the two PSD matrices. Lemma 6.12 (Schur (1911)) LetA,B ∈ Rn×n be two Positive Semi-Definite (PSD) matrices. Then, λmin (A⊙B) ≥(min
i Bii)λmin (A) ,
λmax (A⊙B) ≤(max i Bii)λmax (A) .
Next we state a lemma regarding concentration of the Jacobian matrix at initialization.
Lemma 6.13 (Concentration of the Jacobian at initialization) Consider a one-hidden layer neural network model of the form x ↦ V φ (Wx) where the activation φ obeys ∣φ(0)∣ ≤ B and∣φ′(z)∣ ≤ B for all z. Also assume we have n ≥ K data points x1,x2, . . . ,xn ∈ Rd with unit euclidean norm (∥xi∥`2 = 1) aggregated as the rows of a matrix X ∈ Rn×d. Furthermore, suppose V has i.i.d. ν-scaled Rademacher entries (i.e. ±ν equally-likely). Then, the Jacobian matrix at a random pointW0 ∈ Rk×d with i.i.d. N (0,1) entries obeys
∥J (W0)J (W0)T − E[J (W0)J (W0)T ]∥ ≤ 30K√kν2B2∥X∥2 log(n). with probability at least 1 − 1/n100. In particular, as long as
k ≥ 1000K2B4∥X∥4 log(n) δ2 ,
with the same probability, we have that
∥ 1 kν2 J (W0)J (W0)T −Σ(X)∥ ≤ δ.
Proof Define C = J (W0)J (W0)T . We begin by showing that the diagonal blocks of C are concentrated. To do this first for 1 ≤ s ≤ k define the random matrices
As = (φ′ (Xws)φ′ (Xws)T )⊙ (XXT ) .
Now consider n × n diagonal blocks of C (denoted by C[`, `]) and note that we have C[`, `] = (φ′ (XW T )diag(v`)diag(v`)φ′ (WXT ))⊙ (XXT )
= k∑ s=1V 2 `,sAs
=ν2 k∑ s=1As.
Furthermore, using Lemma 6.12
∥As∥ ≤ (max i (φ′(xTi ws))2) ∥X∥2 ≤ B2 ∥X∥2 . Also, using Jensen’s inequality ∥E[As]∥ ≤ E ∥As∥ ≤ B2 ∥X∥2 . Combining the latter two identities via the triangular inequality we conclude that ∥(As − E[As])2∥ = ∥As − E[As]∥2 ≤ (∥As∥ + ∥E[As]∥)2 ≤ (2B2∥X∥2)2 . (6.34) To proceed, we will bound the weighted sum
S = k∑ s=1ν 2(As − E[As]) in spectral norm. To this aim we utilize the Matrix Hoeffding inequality which states that P(∥S∥ ≥ t) ≤ 2ne− t22∆2 , where ∆2 is an upper bound on ∥∑ks=1 ν4(As − E[As])2∥. Using (6.34) we can pick ∆2 =∑ks=1(2ν2B2∥X∥2)2 = 4kν4B4∥X∥4. Setting t = 30√kν2B2∥X∥2√log(n), we conclude that P{ ∥C[`, `] − E[C[`, `]]∥ ≥ t} = P(∥S∥ ≥ t) ≤ n−102 concluding the proof of concentration of the diagonal blocks of C. For the off-diagonal blocks C[`, ̃̀] using (6.33) from the proof of Lemma 6.11 we have that C[`, ̃̀] = k∑
s=1V`,sṼ̀,sAs. Note that by construction {V`,sṼ̀,s}ks=1 are i.i.d. ±ν2 Rademacher variables and thus C[`, ̃̀] is sum of zero-mean i.i.d. matrices and we are again in the position to apply Hoeffding’s inequality. To this aim note that
∥ k∑ s=1V 2 `,sV 2̃̀,sA2s∥ = ν4 ∥ k∑ s=1A 2 s∥ ≤ ν4 k∑ s=1 ∥As∥2 ≤ ν4kB4 ∥X∥4 , so that we can take ∆2 = ν4kB4 ∥X∥4 and again conclude that for t = 30√kν2B2∥X∥2 log(n) we have
P{∥C[`, ̃̀]∥ ≥ t} ≤ n−102 Using the fact that E[C[`, ̃̀]] = 0 and K ≤ n, combined with a union bound over all sub-matrices 1 ≤ `, ̃̀≤K we conclude that P{∥C[`, ̃̀] − E [C[`, ̃̀]]∥ ≥ t} ≤K2n−102 ≤ n−100. All that remains is to combine the concentration results for the sub-matrices to arrive at the complete bound. In mathematical terms we need to bound D ∶= ∥C − E[C]∥.To this aim define D[`, ∶] to denote the `th block row ofD. Standard bounds on spectral norm in terms of sub-matrices allow us to conclude that∥D[`, ∶]∥ ≤√K sup
1≤̃̀≤K ∥D[`, ̃̀]∥ ≤ √ Kt ⇒
∥D∥ ≤√K sup 1≤`≤K ∥D[`, ∶]∥ ≤ √ K √ Kt =Kt = 30K√kν2B2∥X∥2 log(n),
concluding the proof. The result in terms of δ is obtained by using the population covariance Lemma 6.11.

heading: 6.2.4 UPPER BOUND ON INITIAL RESIDUAL
text: 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 (Oymak & Soltanolkotabi, 2019, Lemma 6.12).
Lemma 6.14 (Upper bound on initial residual) Consider a one-hidden layer neural network model of the form x ↦ V φ (Wx) where the activation φ has bounded derivatives obeying∣φ(0)∣, ∣φ′(z)∣ ≤ B. Also assume we have n data points x1,x2, . . . ,xn ∈ Rd with unit euclidean norm (∥xi∥`2 = 1) aggregated as rows of a matrix X ∈ Rn×d and the corresponding labels given by y ∈ RKn. Furthermore, assume the entries of V are i.i.d. Rademacher variables scaled by ν∥y∥`2 50B √ K log(2K)kn and the entries ofW ∈ Rk×d are i.i.d. N (0,1). Then,
∥V φ (WXT )∥ F ≤ ν ∥y∥`2 ,
holds with probability at least 1 − (2K)−100. Proof We begin the proof by noting that
∥V φ (WXT )∥2 F = K∑̀=1 ∥vT` φ (WXT )∥ 2 `2 .
We will show that for any row v of V , with probability at least 1 − (2K)−101, ∥vT` φ (WXT )∥`2 ≤ ν√K ∥y∥`2 . (6.35)
so that a simple union bound can conclude the proof. Therefore, all that remains is to show (6.35) holds. To prove the latter, note that for any two matrices W̃ ,W ∈ Rk×d we have
∣∥φ (XW̃ T )v∥ `2 − ∥φ (XW T )v∥ `2 ∣ ≤ ∥φ (XW̃ T )v − φ (XW T )v∥
`2≤ ∥φ (XW̃ T ) − φ (XW T )∥ ∥v∥`2≤ ∥φ (XW̃ T ) − φ (XW T )∥ F ∥v∥`2(a)= ∥(φ′ (S ⊙XW̃ T + (1k×n −S)⊙XW T ))⊙ (X(W̃ −W )T )∥F ∥v∥`2≤B ∥X(W̃ −W )T ∥
F ∥v∥`2≤B ∥X∥ ∥v∥`2 ∥W̃ −W ∥F ,
where in (a) we used the mean value theorem with S a matrix with entries obeying 0 ≤ Si,j ≤ 1 and 1k×n the matrix of all ones. Thus, ∥φ (XW T )v∥`2 is a B ∥X∥ ∥v∥`2 -Lipschitz function ofW . Thus, fixing v, for a matrixW with i.i.d. Gaussian entries
∥φ (XW T )v∥ `2 ≤ E [ ∥φ (XW T )v∥ `2 ] + t, (6.36)
holds with probability at least 1 − e− t22B2∥v∥2`2 ∥X∥2 . Next given g ∼ N (0,1), we have ∣E[φ(g)]∣ ≤ ∣E[φ(0)]∣ + ∣E[φ(g) − φ(0)]∣ ≤ B +B E[∣g∣] ≤ 2B and Var(φ(g)) ≤ B2. (6.37)
where the latter follows from Poincare inequality (e.g. see (Ledoux, 2001, p. 49)). Furthermore, since v has i.i.d. Rademacher entries, applying Bernstein bound, event
Ev ∶= {∣1Tv∣2 ≤ 250 logK∥v∥2`2} (6.38)
holds with probability 1 − (2K)−102. Conditioned on Ev , we now upper bound the expectation via E [ ∥φ (XW T )v∥ `2 ] (a)≤ √E [ ∥φ (XW T )v∥2`2 ]
= ¿ÁÁÀ n∑
i=1 E [ (vTφ(Wxi))2 ] (b)= √n√Eg∼N (0,Ik) [ (vTφ(g))2 ] (c)= √n√∥v∥2`2 Eg∼N (0,1) [ (φ(g) − E[φ(g)])2 ] + (1Tv)2(Eg∼N (0,1)[φ(g)])2(d)≤ √n ∥v∥`2 √250 × 4B2 log(2K) +B2≤32√n log(2K)B ∥v∥`2 .
Here, (a) follows from Jensen’s inequality, (b) from linearity of expectation and the fact that for xi with unit Euclidean normWxi ∼ N (0,Ik), (c) from simple algebraic manipulations, (d) from the inequalities (6.38) and (6.37). Thus using t = 18√n log(2K)B ∥v∥`2 in (6.36), conditioned on Ev we conclude that
∥φ (XW T )v∥ `2 ≤50√n log(2K)B ∥v∥`2 = 50√n log(2K)B√k ν ∥y∥`250B√K log(2K)kn = ν∥y∥`2√ K ,
(6.39)
holds with probability at least 1−exp(−102 log(2K) n∥X∥2 ) ≥ 1−(2K)−102 where we used n ≥ ∥X∥2. Using a union bound over Ev and the conditional concentration overW , the overall probability of success in (6.39) is at least 1 − (2K)−101 concluding the proof of (6.35) and the Lemma.

heading: 6.3 RADEMACHER COMPLEXITY AND GENERALIZATION BOUNDS
text: In this section we state and prove some Rademacher complexity results that will be used in our generalization bounds. We begin with some basic notation regarding Rademacher complexity. LetF be a function class. Suppose f ∈ F maps Rd to RK . Let {εi}ni=1 be i.i.d. vectors in RK with i.i.d. Rademacher variables. Given i.i.d. samples S = {(xi,yi)}ni=1 ∼ D, we define the empirical Rademacher complexity to be
RS(F) = 1 n E [ sup f∈F n∑ i=1ε T i f(xi)].
We begin by stating a vector contraction inequality by Maurer Maurer (2016). This is obtained by setting hi(f(xi)) = h(yi, f(xi)) in Corollary 4 of Maurer (2016). Lemma 6.15 Let f(⋅) ∶ Rd → RK and let ` ∶ RK ×RK → R be a 1 Lipschitz loss function with respect to second variable. Let {εi}ni=1 be i.i.d. Rademacher variables. Given i.i.d. samples {(xi,yi)}ni=1, define
RS(`,F) = E [ sup f∈F n∑ i=1 εi`(yi, f(xi))].
We have that RS(`,F) ≤ √2RS(F). Combining the above result with standard generalization bounds based on Rademacher complexity Bartlett & Mendelson (2002) allows us to prove the following result. Lemma 6.16 Let `(⋅, ⋅) ∶ RK × RK → [0,1] be a 1 Lipschitz loss function. Given i.i.d. samples{(xi,yi)}ni=1, consider the empirical loss L(f, `) = 1
n n∑ i=1 `(yi, f(xi)).
With probability 1 − δ over the samples, for all f ∈ F , we have that E[L(f, `)] ≤ L(f, `) + 2√2RS(F) + √ 5 log(2/δ)
n
Proof Based on Bartlett & Mendelson (2002),
E[L(f, `)] ≤ L(f, `) + 2RS(`,F) + √ 5 log(2/δ) n
holds with 1 − δ probability. Combining the latter with Lemma 6.15 completes the proof. Lemma 6.17 Consider a neural network model of the form x ↦ f(x;V ,W ) = V φ (Wx) with W ∈ Rk×d and V ∈ RK×k denoting the input and output weight matrices. Suppose V0 ∈ RK×k is a matrix obeying ∥V0∥`∞ ≤ ν/√kK. Also let W0 ∈ Rk×d be a reference input weight matrix. Furthermore, we define the neural network function space parameterized by the weights as follows
FV,W = {f(x;V ,W ) such that V ∈ V and W ∈W} with V = {V ∶ ∥V −V0∥F ≤ νMV√ Kk } and W = {W ∶ ∥W −W0∥F ≤MW and ∥W −W0∥2,∞ ≤ R√
k }. (6.40)
Additionally, assume the training data {(xi,yi)}ni=1are generated i.i.d. with the input data points of unit Euclidean norm (i.e. ∥xi∥`2 = 1). Also, define the average energy atW0 as
E = ( 1 kn n∑ i=1 ∥φ(W0xi)∥2`2)
1/2 .
Also let {ξi}ni=1 ∈ RK be i.i.d. vectors with i.i.d. Rademacher entries and define the empirical Rademacher complexity
RS (FV,W) ∶= 1 n E ⎡⎢⎢⎢⎢⎣ supf∈FV,W n∑ i=1ξ T i f(xi) ⎤⎥⎥⎥⎥⎦. Then,
RS (FV,W) ≤ νB (MW +EMV√ n + R2 +MWMV√ k
) . (6.41) Proof We shall use w` to denote the rows of W (same for W0,V ,V0). We will approximate φ (⟨w`,xi⟩) by its linear approximation φ (⟨w0` ,xi⟩)+φ′(⟨w0` ,xi⟩) (⟨w` −w0` ,xi⟩) via the second order Taylor’s mean value theorem. We thus have
RS (FV,W) ≤ 1 n E ⎡⎢⎢⎢⎢⎣ n∑ i=1ξ T i V0φ (W0xi) ⎤⎥⎥⎥⎥⎦ + 1 n E ⎡⎢⎢⎢⎢⎣ supW ∈W n∑ i=1ξ T i V0diag (φ′ (W0xi)) (W −W0)xi
⎤⎥⎥⎥⎥⎦´¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¸¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¶R1 + 1
2n E ξi,j i.i.d.∼ ±1 ⎡⎢⎢⎢⎢⎣ supW ∈W n∑ i=1 k∑̀=1 K∑ j=1 ξ T i,jV 0 j,`φ ′′ ((1 − ti`)⟨w0` ,xi⟩ + ti`⟨w`,xi⟩) (⟨w` −w0` ,xi⟩)2 ⎤⎥⎥⎥⎥⎦´¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¸¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¶R2
+ 1 n E [ sup V ∈V,W ∈W n∑ i=1ξ T i (V −V0)(φ(Wxi) − φ(W0xi))]
´¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¸¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¶R3 + 1 n E [ sup V ∈V n∑ i=1ξ T i (V −V0)φ(W0xi)]
´¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¸¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¶R4
We proceed by bounding each of these four terms. For the first term note that
R1 ≤ 1 n E ⎡⎢⎢⎢⎢⎣ sup∥W−W0∥F ≤MW n∑ i=1ξ T i V0diag (φ′ (W0xi)) (W −W0)xi ⎤⎥⎥⎥⎥⎦ ≤ 1 n E ⎡⎢⎢⎢⎢⎣ sup∥W−W0∥F ≤MW ⟨ n∑ i=1 diag (φ′ (W0xi))V T0 ξixTi ,W −W0⟩
⎤⎥⎥⎥⎥⎦ ≤MW n E ⎡⎢⎢⎢⎢⎣∥( n∑ i=1 diag (V T0 ξi)φ′ (W0xi)xTi )∥F
⎤⎥⎥⎥⎥⎦ ≤MW n E ⎡⎢⎢⎢⎢⎣∥( n∑ i=1 diag (V T0 ξi)φ′ (W0xi)xTi )∥ 2 F ⎤⎥⎥⎥⎥⎦ 1/2
=MW n ⎡⎢⎢⎢⎢⎣ n∑ i=1 E ∥diag (V T0 ξi)φ′ (W0xi)xTi ∥ 2 F ⎤⎥⎥⎥⎥⎦ 1/2
=MW n ⎡⎢⎢⎢⎢⎣ n∑ i=1 E ∥diag (V T0 ξi)φ′ (W0xi)∥2`2 ⎤⎥⎥⎥⎥⎦ 1/2
≤BMW n ⎡⎢⎢⎢⎢⎣ n∑ i=1 E ∥V T0 ξi∥2`2 ⎤⎥⎥⎥⎥⎦ 1/2 ≤BMW n
∥V0∥F ≤BMWν√
n ,
where in the last inequality we used the fact that ∥V0∥F ≤ ν. For the second term note that
R2 ≤ 1 2n E ⎡⎢⎢⎢⎢⎣ sup∥W−W0∥2,∞≤R n∑ i=1 k∑̀=1 K∑ j=1 ξi,jv0,j,`φ ′′ ((1 − ti`)⟨w0` ,xi⟩ + ti`⟨w`,xi⟩) (⟨w` −w0` ,xi⟩)2 ⎤⎥⎥⎥⎥⎦
≤ 1 2n k∑̀=1 E ⎡⎢⎢⎢⎢⎣ sup∥w`−w0` ∥`2≤R n∑ i=1 RRRRRRRRRRR K∑ j=1 ξi,jv0,j,` RRRRRRRRRRR ∣φ ′′ ((1 − ti`)⟨w0` ,xi⟩ + ti`⟨w`,xi⟩)∣ (⟨w` −w0` ,xi⟩)2 ⎤⎥⎥⎥⎥⎦ ≤ 1
2kn k∑̀=1 n∑ i=1 E ⎡⎢⎢⎢⎢⎣ RRRRRRRRRRR K∑ j=1 ξi,jv0,j,` RRRRRRRRRRR ⎤⎥⎥⎥⎥⎦R 2B
≤BR2 2k ∥V T0 ∥2,1 ≤BR2ν
2 √ k .
In the above we used ∥M∥2,1 for a matrixM to denote the sum of the Euclidean norm of the rows ofM . We also used the fact that ∥V T0 ∥2,1 ≤ ν√k. To bound the third term note that
R3 = 1 n E [ sup V ∈V,W ∈W n∑ i=1ξ T i (V −V0)(φ(Wxi) − φ(W0xi))]
≤ 1 n E [ sup V ∈V,W ∈W n∑ i=1 ∥V −V0∥F ∥ξi∥`2 ∥(φ(Wxi) − φ(W0xi))∥`2 ] ≤ νMV n √ kK E [ sup W ∈W n∑ i=1 ∥ξi∥`2∥φ(Wxi) − φ(W0xi)∥`2] ≤ νMV n √ k ⋅ sup W ∈W n∑ i=1 ∥φ(Wxi) − φ(W0xi)∥`2 ≤ νBMV n √ k ⋅ sup W ∈W n∑ i=1 ∥(W −W0)xi∥`2 ≤ νBMV n √ k ⋅ sup W ∈W n∑ i=1 ∥(W −W0)∥F = νBMV√ k ⋅ sup W ∈W ∥(W −W0)∥F = νBMVMW√ k .
Finally, to bound the fourth term note that we have
R4 = 1 n E [ sup V ∈V n∑ i=1ξ T i (V −V0)φ(W0xi)]
= 1 n E [ sup ∥V −V0∥F ≤ νMV√kK ⟨ n∑ i=1ξiφ(W0xi)T , (V −V0)⟩] = νMV n √ kK E [∥ n∑ i=1φ(W0xi)ξTi ∥F ]
≤ νMV n √ kK ⎛⎝E [∥ n∑ i=1φ(W0xi)ξTi ∥ 2
F
]⎞⎠ 1/2
= νMV n √ kK ( n∑ i=1 E [ ∥φ(W0xi)ξTi ∥ 2
F ])1/2
= νMV√ n ( 1 kn n∑ i=1 E [ ∥φ(W0xi)∥2F ])
1/2
= νEMV√ n ≤ νBEMV√ n
Combining these four bounds we conclude that
RS (FV,W) ≤ νB (MW√ n + R2√ k + MVMW√ k + EMV√ n ) ,
concluding the proof of (6.41).
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.
Lemma 6.18 Consider a one-hidden layer neural network with input to output mapping of the form x ∈ Rd ↦ f(x;V ,W ) = V φ (Wx) ∈ RK withW ∈ Rk×d denoting the input-to-hidden weights and
V ∈ Rk×K the hidden-to-output weights. Suppose V0 ∈ RK×k is a matrix obeying ∥V0∥`∞ ≤ ν/√kK. Also letW0 ∈ Rk×d be a reference input weight matrix. Also define the empirical losses
L(V ,W ) = 1 n n∑ i=1 ∥yi − f(xi;V ,W )∥2`2 ,
and
L(f, `) = 1 n n∑ i=1 `(f(xi;V ,W ),yi),
with ` ∶ RK × RK → [0,1] a one Lipschitz loss function obeying `(y,y) = 0. Additionally, assume the training data {(xi,yi)}ni=1are generated i.i.d. according to a distribution D with the input data points of unit Euclidean norm (i.e. ∥xi∥`2 = 1). Also, define the average energy atW0 as
E = ( 1 kn n∑ i=1 ∥φ(W0xi)∥2`2)
1/2 .
Then for all f in the function class FW given by (6.40) E[L(f, `)] ≤ √L(V ,W ) + 2√2νB (MW√
n + R2√ k + νBMVMW√ k + EMV√ n ) +
√ 5 log(2/δ)
n ,
(6.42)
holds with probability at least 1 − δ. Furthermore, Suppose labels are one-hot encoded and thus unit Euclidian norm. Given a sample (x,y) ∈ Rd × RK generated according to the distribution D, define the population classification error
ErrD(W ) = P(arg max 1≤`≤K yi ≠ arg max1≤`≤K fi(x,W )).
Then, we also have
ErrD(W ) ≤ 2⎡⎢⎢⎢⎢⎣ √L(V ,W ) + 2√2νB (MW√ n + R2√ k + νBMVMW√ k + EMV√ n ) + √ 5 log(2/δ) n ⎤⎥⎥⎥⎥⎦ . (6.43)
Proof To begin first note that any 1-Lipschitz ` with `(y,y) = 0 obeys `(y, ŷ) ≤ ∥y − ŷ∥`2 . Thus, we have
L(f, `) ≤ 1 n n∑ i=1 ∥yi − f(xi;V ,W )∥`2 ≤ √L(V ,W ), where the last inequality follows from Cauchy-Schwarz. Consequently, applying Lemmas 6.16 and 6.17 we conclude that
E[L(f, `)] ≤ L(f, `) + 2√2 ⋅RS(F) + √ 5 log(2/δ) n ,
≤ √L(W ) + 2√2νB (MW√ n + R2√ k + νBMVMW√ k + EMV√ n ) +
√ 5 log(2/δ)
n ,
which yields the first statement.
To prove the second statement on classification accuracy, we pick the ` function as follows `(y, ŷ) = min(1, ∥y − ŷ∥`2). Note that, given a sample (x,y) ∈ Rd × RK with one-hot encoded labels, if arg max
1≤`≤K y` ≠ arg max1≤`≤K f`(x;V ,W ), this implies
` (y, f(x;V ,W )) ≥ 0.5. Combining the latter with Markov inequality we arrive at ErrD(W ) ≤ 2 E(x,y)∼D[`(y, f(x;V ,W ))] = 2 E[L(`,W )]. Now since ` is 1 Lipschitz and bounded, it obeys (6.42), which combined with the above identity yields (6.43), completing the proof.

heading: 6.4 PROOFS FOR NEURAL NETS WITH ARBITRARY INITIALIZATION (PROOF OF THEOREM 2.3)
text: In this section we prove Theorem 2.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 2.3.

heading: 6.4.1 PRELIMINARY OPTIMIZATION RESULT
text: Lemma 6.19 (Deterministic convergence guarantee) Consider a one-hidden layer neural net of the form x ↦ f(x;W ) ∶= V φ(Wx) with input weights W ∈ Rk×d and output weights V ∈ RK×k and an activation φ obeying ∣φ(0)∣ ≤ B, ∣φ′(z)∣ ≤ B, and ∣φ′′(z)∣ ≤ B for all z. Also assume V is fixed with all entries bounded by ∥V ∥`∞ ≤ ν√kK and we train overW based on the loss
L(W ) = 1 2 n∑ i=1 ∥f(xi;W ) − yi∥2`2 .
Also, consider a point W0 ∈ Rk×d with J an ( 0, νB∥X∥) reference Jacobian associated withJ (W0) per Definition 5.1. Furthermore, define the information I and nuisance N subspaces and the truncated Jacobian JI associated with the reference Jacobian J based on a cut-off spectrum value of α per Definition 5.2. Let the initial residual vector be r0 = y − f(W0) ∈ RnK . Furthermore, assume
ε0 ≤ α 5 min(δ,√ δα ΓνB∥X∥) (6.44)
and
k ≥ 400ν6B6 ∥X∥6 Γ2 δ2α8
(Bα,Γ + δΓ∥r0∥`2)2 , (6.45) with 0 ≤ δ ≤ 1 and Γ ≥ 1. We run gradient descent iterations of the formWτ+1 =Wτ − η∇L(Wτ) starting fromW0 with step size η obeying η ≤ 1ν2B2∥X∥2 . Then for all iterates τ obeying 0 ≤ τ ≤ T ∶=
Γ ηα2
∥Wτ −W0∥F ≤Bα,Γ α + δΓ α ∥r0∥`2 . (6.46)
∥Wτ −W0∥2,∞ ≤2νBΓ ∥X∥√ kα2
∥r0∥`2 . (6.47) Furthermore, after τ = T iteration we have
∥rT ∥`2 ≤ e−Γ∥ΠI(r0)∥`2 + ∥ΠN (r0)∥`2 + δανB ∥X∥∥r0∥`2 . (6.48) Proof To prove this lemma we wish to apply Theorem 5.3. We thus need to ensure that the assumptions of this theorem are satisfied. To do this note that by Lemma 6.10 Assumption 1 holds with β = νB ∥X∥. Furthermore, we pick ε = δα3
5Γν2B2∥X∥2 = δα35Γβ2 which together with (6.44) guarantees (5.13) holds. We now turn our attention to verifying Assumption 2. To this aim note that for allW ∈ Rk×d obeying
∥W −W0∥F ≤ R ∶= 2(Bα,Γα + δΓα∥r0∥`2)
as long as (6.45) holds by Lemma 6.10 we have ∥J (W ) −J (W0)∥ ≤B√K ∥V ∥`∞ ∥X∥R≤ ν√ k B ∥X∥R
= δα3 10Γν2B2 ∥X∥2
20Γν3B3∥X∥3 δα4 (Bα,Γ + δΓ∥r0∥`2)√ k
≤ δα3 10Γν2B2 ∥X∥2 =ε 2 .
Thus, Assumption 2 holds with ∥W −W0∥F ≤ R ∶= 2 (Bα,Γα + δ Γα∥r0∥`2). Now that we have verified that the assumptions of Theorem 5.3 hold so do its conclusions and thus (6.46) and (6.48) hold. We now turn our attention to proving the row-wise bound (6.47). To this aim let w(τ)` denote the `th row ofWτ . Also note that ∇L(w`) = `th row of mat(J (W )Trτ). Hence, using Lemma 6.10 equation (6.32) we conclude that
∥∇L(w(τ)` )∥`2 ≤ B √ K ∥V ∥`∞ ∥X∥ ∥rτ∥`2 ≤ νB ∥X∥√k ∥rτ∥`2 .
Consequently, for any row 1 ≤ ` ≤ k, we have ∥w(τ)` −w(0)` ∥`2 ≤ η νB ∥X∥√k τ−1∑ t=0 ∥rt∥`2 . (6.49)
To bound the right-hand side we use the triangular inequality combined with (6.25) and (6.26) to conclude that
η τ−1∑ t=0 ∥rτ∥`2 ≤η τ−1∑ t=0 ∥r̃τ∥`2 + η τ−1∑ t=0 ∥rτ − r̃τ∥`2
≤ Γ α2 ∥r0∥`2 + 2Γ2(εβ + ε20)α4 ∥r0∥`2 =2Γ(ε20 + εβ) + α2
α4 Γ∥r0∥`2
≤2 Γ α2 ∥r0∥`2 , (6.50) where in the last inequality we used the fact that ε20 ≤ α225Γ per (6.44) and β = δα35Γβ ≤ α25Γ per our choice of . Combining (6.49) and (6.50), we obtain ∥w(τ)` −w(0)` ∥`2 ≤ 2νB∥X∥Γ√kα2 ∥r0∥`2 , completing the proof of (6.47) and the theorem.

heading: 6.4.2 MAIN OPTIMIZATION RESULT
text: Lemma 6.20 (Deterministic optimization guarantee) Consider the setting and assumptions of Lemma 6.19. Also assume ∥ΠI(r0)∥`2 ≥ c∥r0∥`2 for a constant c > 0 if ε0 > 0. Furthermore, assume
ε20 ≤ α225 min(c Bα,ΓανBΓ2∥r0∥`2 ∥X∥ , ζ2ν2B2 ∥X∥2 α2 , ζ Γ ) , (6.51)
and
k ≥ 1600( α ζνB ∥X∥ + Γ∥r0∥`2Bα,Γ ) 2 ν6B6 ∥X∥6 Γ2B2α,Γ α8 , (6.52)
and Γ ≥ 1. We run gradient descent iterations of the form Wτ+1 =Wτ − η∇L(Wτ) starting from W0 with step size η obeying η ≤ 1ν2B2∥X∥2 . Then for all iterates τ obeying 0 ≤ τ ≤ T ∶= Γηα2
∥Wτ −W0∥F ≤ 2Bα,Γ α . (6.53) ∥Wτ −W0∥2,∞ ≤ 2νBΓ ∥X∥√ kα2
∥r0∥`2 . (6.54) Furthermore, after τ = T iteration we have
∥f(WT ) − y∥`2 ≤ e−Γ∥ΠI(r0)∥`2 + ∥ΠN (r0)∥`2 + ζ∥r0∥`2 . (6.55) Proof To prove this lemma we aim to substitute
δ = min(ζνB ∥X∥ α , Bα,Γ
Γ∥r0∥`2 ) ≤ 1, (6.56) 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 δ from (6.56) combined with (6.52) ensures that
k ≥1600( α ζνB ∥X∥ + Γ∥r0∥`2Bα,Γ ) 2 ν6B6 ∥X∥6 Γ2B2α,Γ α8
≥max( α ζνB ∥X∥ , Γ∥r0∥`2Bα,Γ ) 2 1600ν6B6 ∥X∥6 Γ2B2α,Γ α8
≥ 1 min ( ζνB∥X∥ α , Bα,Γ Γ∥r0∥`2 )2 1600ν6B6 ∥X∥6 Γ2B2α,Γ α8 =1600Γ2ν6B6 ∥X∥6 B2α,Γ δ2α8 =400Γ2ν6B6 ∥X∥6 (Bα,Γ + Bα,Γ)2 δ2α8 ≥400Γ2ν6B6 ∥X∥6 (Bα,Γ + δΓ∥r0∥`2)2 δ2α8 ,
so that (6.45) holds. We thus turn our attention to proving (6.44). If ε0 = 0, the statement already holds. Otherwise, note that based on Lemma 6.2 equation (6.3) we have
Bα,Γ ≥ α∥ΠI(r0)∥`2 λ1 ≥ α∥ΠI(r0)∥`2 νB∥X∥ ≥ cα∥r0∥`2νB∥X∥ ⇒ Bα,Γc∥r0∥`2 ≥
α
νB ∥X∥ . (6.57) Recall that α = να0√
K , which implies that
• If δ = ζνB∥X∥ α : For (6.44) to hold it suffices to have ε0 ≤ α5 min( ζνB∥X∥α ,√ ζΓ).
• If δ = Bα,Γ Γ∥r0∥`2 : For (6.44) to hold it suffices to have ε0 ≤ α5
√ c
Bα,Γα νBΓ2∥r0∥`2∥X∥ as based on
(6.57) we have √
cα βΓ = √ cα νBΓ∥X∥ ≤ √δ
ε0 ≤α 5
√ cBα,Γα
νBΓ2∥r0∥`2 ∥X∥ =α
5
√ δ
√ cα
Γβ
=α 5
√ δ ⋅min(√δ,√ cα
Γβ )
=α 5 min ⎛⎝δ,
√ δα
Γβ ⎞⎠ Combining the latter two cases as long as
ε20 ≤ α225 min(c Bα,ΓανBΓ2∥r0∥`2 ∥X∥ , ζ2ν2B2 ∥X∥2 α2 , ζ Γ ) ⇔ (6.51),
then (6.44) holds. As a result when (6.51) and (6.52) hold with δ = min ( ζνB∥X∥ α , Bα,Γ Γ∥r0∥`2 ) both assumptions of Theorem 6.19 also hold and so do its conclusions. In particular, (6.53) follows from (6.46) by noting that based on our choice of δ we have δ Γ α ∥r0∥`2 ≤ Bα,Γα , (6.54) follows immediately from (6.47), and (6.55) follows from (6.48) by noting that based on our choice of δ we have δα
νB∥X∥ ≤ ζ.

heading: 6.4.3 MAIN GENERALIZATION RESULT (COMPLETING THE PROOF OF THEOREM 2.3)
text: We state the following rigorous and stronger version of Theorem 2.3. The differences (besides constant terms) are the use of ∥X∥ rather than √n (we always have ∥X∥ ≤ √n) and use of the tighter bound ∥J†Ir0∥`2 rather than ∥ΠI(r0)∥`2 /α. Theorem 6.21 Let ζ,Γ, ᾱ be scalars obeying ζ ≤ 1/2, Γ ≥ 1, and ᾱ ≥ 0 which determine the overall precision, cut-off and learning duration, respectively.4 Consider a training data set{(xi,yi)}ni=1 ∈ Rd × RK generated i.i.d. according to a distribution D where the input samples have unit Euclidean norm. Also consider a neural net with k hidden nodes as described in (1.1) parameterized by W where the activation function φ obeys ∣φ′(z)∣ , ∣φ′′(z)∣ ≤ B. Let W0 be an arbitrary initial weight matrix. Also assume the output matrix has bounded entries obeying ∥V ∥`∞ ≤ ν√kK . Furthermore, set J ∶= J (W0) and define the information I and nuisance N subspaces and the truncated Jacobian JI associated with the reference/initial Jacobian J based on a cut-off spectrum value α = νBᾱ 4√n√∥X∥. Also define the initial residual r0 = f(W0) − y ∈ RnK and pick Cr > 0 so that ∥r0∥`2√
n ≤ Cr. Suppose number of hidden nodes k obeys
k ≳ C2rΓ4 ᾱ8ν2ζ2 , (6.58)
with Γ ≥ 1 and tolerance level ζ . Run gradient descent updates (1.5) with learning rate η ≤ 1 ν2B2∥X∥2 . Then, after T = Γ ηα2 iterations, training loss obeys
∥f(WT ) − y∥`2 ≤ ∥ΠN (r0)∥`2 +Cr (e−Γ + ζ)√n. and with probability at least 1 − δ, the generalization error obeys ErrD(WT ) ≤ 2∥ΠN (r0)∥`2√
n´¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¸¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¶ bias term
+ 12νB√ n (∥J†Ir0∥`2 + Γα∥ΠN (r0)∥`2)´¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¸¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¶ variance term
+5 √
log(2/δ) n
+ 2Cr(e−Γ + ζ). 4Note that this theorem and its conclusions hold for any choice of these parameters in the specified range.
Theorem 6.21 immediately follows from Theorem 6.22 below by upper boundingDα,Γ (see Definition 6.1) using Lemma 6.2 equation (6.2).
Theorem 6.22 Consider a training data set {(xi,yi)}ni=1 ∈ Rd × RK generated i.i.d. according to a distribution D where the input samples have unit Euclidean norm. Also consider a neural net with k hidden nodes as described in (1.1) parameterized by W where the activation function φ obeys ∣φ′(z)∣ , ∣φ′′(z)∣ ≤ B. Let W0 be an arbitrary initial weight matrix. Also assume the output matrix has bounded entries obeying ∥V ∥`∞ ≤ ν√kK . Furthermore, set J ∶= J (W0) and define the information I and nuisance N subspaces and the truncated Jacobian JI associated with the reference/initial Jacobian J based on a cut-off spectrum value α = νBᾱ 4√n√∥X∥. Also define the initial residual r0 = f(W0) − y ∈ RnK and pick Cr > 0 so that ∥r0∥`2√n ≤ Cr. Also assume, the number of hidden nodes k obeys
k ≥ 25600 C2rΓ4 ᾱ8ν2B2ζ2 , (6.59)
with Γ ≥ 1 and tolerance level ζ ≤ 2. Run gradient descent updates (1.5) with learning rate η ≤ 1
ν2B2∥X∥2 . Then, after T = Γηα2 iterations, with probability at least 1− δ, the generalization error obeys
ErrD(WT ) ≤ 2∥ΠN (r0)∥`2√ n´¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¸¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¶
bias term
+ 12νBDα,Γ√ n´¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¸¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¶
variance term
+5 √
log(2/δ) n + 2Cr(e−Γ + ζ), (6.60)
where Dα,Γ is the early stopping distance as in Def. (6.1). Proof First, note that using α ≤ β = νB ∥X∥, Bα,Γ Γ ≤ ∥r0∥`2 per (6.2), and ∥r0∥`2 ≤ Cr√n we have
α νB∥X∥ ≤ 1 ≤ Γ∥r0∥`2Bα,Γ ≤ Cr Γ √ nBα,Γ .
This together with ζ ≤ 2 implies that α
ζ 2 νB ∥X∥ + Γ∥r0∥`2Bα,Γ ≤ αζ2νB ∥X∥ + Γ∥r0∥`2 ζ 2 Bα,Γ
≤2Γ∥r0∥`2 ζ 2 Bα,Γ ≤2Cr ζ 2 ΓBα,Γ √ n. (6.61)
Thus when
k ≥25600 C2rΓ4 ᾱ8ν2B2ζ2
=6400⎛⎝CrΓζ 2 ⎞⎠ 2 n2Γ2ν6B6 ∥X∥4 α8
√ n≥∥X∥≥ 6400⎛⎝CrΓζ
2
⎞⎠ 2 nν6B6 ∥X∥6 α8
=1600⎛⎝2Crζ 2 ΓBα,Γ √ n ⎞⎠ 2 B2α,Γν6B6 ∥X∥6 α8
(6.61)≥ 1600⎛⎝ αζ 2 νB ∥X∥ + Γ∥r0∥`2Bα,Γ ⎞⎠ 2 B2α,Γν6B6 ∥X∥6 α8
Thus, (6.52) holds. Also (6.51) trivially holds for ε0. Thus applying Theorem 6.20 with ε0 = 0 the following three conclusions hold
∥Wτ −W0∥F ≤ 2Bα,Γ α
= 2Dα,Γ. (6.62) and
∥Wτ −W0∥2,∞ ≤2νBΓ ∥X∥√ kα2
∥r0∥`2 ∥r0∥`2≤Cr√n≤ 2√nCrνBΓ ∥X∥√
kα2 . (6.63)
and
∥f(WT ) − y∥`2 ≤e−Γ∥ΠI(r0)∥`2 + ∥ΠN (r0)∥`2 + ζ2∥r0∥`2 ≤∥ΠN (r0)∥`2 + (e−Γ + ζ2) ∥r0∥`2
∥r0∥`2≤Cr√n≤ ∥ΠN (r0)∥`2 +Cr (e−Γ + ζ2)√n. (6.64) Furthermore, using the assumption that ∥V ∥`∞ ≤ ν√kK Lemma 6.18 applies and hence equation (6.43) with W =WT , √L(WT ) = ∥f(WT )−y∥`2√n , MW = 2Dα,Γ, MV = 0, and R = 2√nCrνBΓ∥X∥α2 implies that
ErrD(WT ) ≤ 2⎡⎢⎢⎢⎢⎣ ∥f(WT ) − y∥`2√ n + 3νB (2Dα,Γ√ n + R2√ k ) + √ 5 log(2/δ) n ⎤⎥⎥⎥⎥⎦ (6.65) Also note that using (6.59) we have
3νBR2√ k ≤ 12C2rΓ2ν3B3n∥X∥2√ kα4 ≤ Cr ζ 2
(6.66)
Plugging (6.64) and (6.66) into (6.65) completes the proof.
6.5 PROOFS FOR NEURAL NETWORK WITH RANDOM INITIALIZATION (PROOF OF THEOREM 2.2)
In this section we prove Theorem 2.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 2.2.

heading: 6.5.1 OPTIMIZATION RESULT
text: Theorem 6.23 (Optimization guarantee for random initialization) Consider a training data set{(xi,yi)}ni=1 ∈ Rd × RK generated i.i.d. according to a distribution D where the input samples have unit Euclidean norm and the concatenated label vector obeys ∥y∥`2 = √n (e.g. one-hot encoding). Consider a neural net with k hidden layers as described in (1.1) parameterized by W where the activation function φ obeys ∣φ′(z)∣ , ∣φ′′(z)∣ ≤ B. LetW0 be the initial weight matrix with i.i.d. N (0,1) entries. Fix a precision level ζ and set
ν = ζ 50B √ log(2K) . (6.67)
Also assume the output layer V has i.i.d. Rademacher entries scaled by ν√ kK . Furthermore,
set J ∶= Σ(X)1/2 and define the information I and nuisance N spaces and the truncated Jacobian JI associated with the reference Jacobian J based on a cut-off spectrum value of α0 = ᾱ 4√n√K ∥X∥B ≤ B√K ∥X∥ per Definition 1.1 so as to ensure ∥ΠI(y)∥`2 ≥ c ∥y∥`2 for some constant c. Assume
k ≥ 12 × 107 Γ4K4B8∥X∥6n log(n) c4ζ4α80
(6.68)
with Γ ≥ 1 and ζ ≤ c 2
. We run gradient descent iterations of the form (1.5) with a learning rate η ≤ 1
ν2B2∥X∥2 . Then, after T = ΓKην2α20 iterations, the following identities ∥f(Wτ0) − y∥`2 ≤ ∥ΠN (y)∥`2 + e−Γ∥ΠI(y)∥`2 + 4ζ√n, (6.69) ∥Wτ −W0∥`2 ≤ 2 √ K(Bα0,Γ(y) + Γζ√n)
να0 , (6.70)
∥Wτ −W0∥2,∞ ≤ 4ΓBK∥X∥ να20
√ n
k , (6.71)
hold with probability at least 1 − (2K)−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 − (2K)−100, the initial prediction vector f(W0) obeys
∥f(W0)∥`2 ≤ ζ∥y∥`2 = ζ√n ≤ √ n
2 . (6.72)
Hence the initial residual obeys ∥r0∥`2 ≤ 2√n. Furthermore, using ζ ≤ c/2 ∥r0 + y∥`2 ≤ζ∥y∥`2 Ô⇒ ∥ΠI(r0 + y)∥`2 ≤ ζ∥y∥`2 . (6.73)
Thus,
∥ΠI(r0)∥`2 ≥∥ΠI(y)∥`2 − ∥ΠI(r0 + y)∥`2≥∥ΠI(y)∥`2 − ζ∥y∥`2≥ (c − ζ) ∥y∥`2 ≥ c
2 ∥y∥`2 (6.74) ≥ c 4 ∥r0∥`2 . (6.75)
Thus the assumption on the ratio of information to total energy of residual holds and we can replace c with c 4 in Theorem 6.20. Furthermore, since Bα0,Γ(⋅) is Γ-Lipschitz function of its input vector in `2 norm hence we also have Bα0,Γ(r0) ≤ Bα0,Γ(y) + Γ∥r0 + y∥`2 ≤ Bα0,Γ(y) + Γζ∥y∥`2 . (6.76) Next we wish to show that (6.51) holds. In particular we will show that there exists an ε0reference Jacobian J for J (W0) satisfying JJT = E[J (W0)J (W0)T ]. Note that, such a J will have exactly same information/nuisance spaces as the square-root of the multiclass kernel matrix i.e. (E[J (W0)J (W0)T ]) 12 since these subspaces are governed by the left eigenvectors. Applying Lemmas 6.13 (with a scaling of the Jacobian by 1/√kK due to the different scaling of V ), we find that if
k ≥ 1000K2B4∥X∥4 log(n) δ2
(6.77)
then,
∥J (W0)J (W0)T − E[J (W0)J (W0)T ]∥ ≤ δν2 K . (6.78)
Let J (W ) be obtained by adding max(Kn − p,0) zero columns to J (W ). Then, using (6.78) and Lemma 6.4, there exists J satisfying JJT = E[J (W0)J (W0)T ] and
∥J (W0) − J∥ ≤ 2 √ δν2
K .
Therefore, J is an ε20 = 4 δν2K reference Jacobian. Now set Θ = min⎛⎝c Bα0,Γα0BΓ2 ∥X∥√nK ,(ζB √ K ∥X∥ α0 )2 , ζ Γ
⎞⎠ and note that using α = ν√ K α0 and ∥r0∥`2 ≤ 2√n
Θ =min⎛⎝c Bα0,Γα0BΓ2 ∥X∥√nK ,(ζB √ K ∥X∥ α0 )2 , ζ Γ ⎞⎠ =min⎛⎝c Bα0,ΓανBΓ2 ∥X∥√n,(νζB ∥X∥α ) 2 , ζ Γ
⎞⎠ ≤2 ⋅min(c Bα,Γα
νBΓ2∥r0∥`2 ∥X∥ , ζ2ν2B2 ∥X∥2 α2 , ζ Γ ) (6.79)
To continue further, note that Bα0,Γ calculated with respect to Σ(X)1/2 with cutoff α0 is exactly same as Bα,Γ calculated with respect to J with cutoff α = να0√K which is a square-root of E[J (W0)J (W0)T ]. Thus, using (6.79) to ensure (6.51) holds it suffices to show
ε20 = 4δν2K ≤ α 2 25 Θ 2 = ν2α20 50K Θ.
Hence, to ensure (6.51) holds we need to ensure that δ obeys
δ ≤ α20 200 Θ.
Thus using δ = α20 200
Θ to ensure (6.51) we need to make sure k is sufficiently large so that (6.77) holds with this value of δ. Thus it suffices to have
k ≥12 × 107 Γ4K4B8∥X∥6n log(n) c4ζ4α80
(6.80)
≥12 × 107 Γ4K4B8∥X∥8 log(n) c4ζ4α80 ≥4 × 107 ⋅ (4K2B4Γ4 ∥X∥4 c4α40 + 1 ζ4 + Γ2 ζ2 ) K2B4∥X∥4 log(n) α40 ≥4 × 107 ⋅ (4KB4Γ4 ∥X∥4 c4α40 + 1 ζ4 + Γ2 ζ2 ) K2B4∥X∥4 log(n) α40 ≥4 × 107 ⋅ (4KB4Γ4 ∥X∥4 c4α40 + α40 ζ4B4K2 ∥X∥4 + Γ2 ζ2 ) K2B4∥X∥4 log(n) α40
(a)≥ 4 × 107 ⋅ ⎛⎝nKB 2Γ4 ∥X∥2 c2B2α0,Γα20 + α40 ζ4B4K2 ∥X∥4 + Γ2 ζ2 ⎞⎠ K 2B4∥X∥4 log(n) α40
≥4 × 107 ⋅max⎛⎝nKB 2Γ4 ∥X∥2 c2B2α0,Γα20 , α40 ζ4B4K2 ∥X∥4 , Γ2 ζ2 ⎞⎠ K 2B4∥X∥4 log(n) α40 =4 × 107 K2B4∥X∥4 log(n) α40 ⋅min(( cBα0,Γα0BΓ2∥X∥√nK ) 2 ,( ζB√K∥X∥ α0 )4 , ζ2 Γ2
) =1000K2B4∥X∥4 log(n)
δ2 (6.81)
Here, (a) follows from the fact that ∥Σ(X)1/2∥ ∶= λ1 ≤ B∥X∥, equation (6.3), and ∥ΠI(r0)∥`2 ≥ c 2 ∥y∥`2 = c2√n which combined imply
Bα0,Γ ≥ αλ1 ∥ΠI(r0)∥`2 ≥ αB ∥X∥ ∥ΠI(r0)∥`2 ≥ α0c2 ∥y∥`2 B∥X∥ = α0c2
√ n
B∥X∥ . (6.82)
To be able to apply Theorem 6.20 we must also ensure (6.52) holds. Therefore, it suffices to have
k ≥64 × 106K4B8 ∥X∥6 Γ4n log(n) ζ4α80
(6.83)
(a)≥ 25600K4B6 ∥X∥6 Γ4n ζ2ν2α80
(6.84)
k ≥12800( α20 ζ2B2K ∥X∥2 + 1) K4B6 ∥X∥6 Γ4n ν2α80
(b)≥ 3200⎛⎝ α 2 0 ζ2B2K ∥X∥2 + 4Γ2nB2α,Γ ⎞⎠ K4B6 ∥X∥6 Γ2B2α,Γ ν2α80
≥3200⎛⎝ α 2 0 ζ2B2K ∥X∥2 + Γ2∥r0∥2`2B2α,Γ ⎞⎠ K4B6 ∥X∥6 Γ2B2α,Γ ν2α80 ≥1600( α0 ζ √ KB ∥X∥ + Γ∥r0∥`2Bα,Γ ) 2 K4B6 ∥X∥6 Γ2B2α,Γ ν2α80 =1600( α ζνB ∥X∥ + Γ∥r0∥`2Bα,Γ ) 2 ν6B6 ∥X∥6 Γ2B2α,Γ α8
(6.85)
Here, (a) follows from the fact that n ≥ K and the relationship between ζ and ν per (6.88) and (b) follows from the fact that per equation (6.2) we have
Bα,Γ ≤ Γ ∥r0∥`2 ≤ 2Γ√n Note that (6.80) and (6.84) are implied by
k ≥ 12 × 107 Γ4K4B8∥X∥6n log(n) c4ζ4α80 , (6.86)
which is the same as (6.68). What remains is stating the optimization bounds in terms of the labels y. This follows by substituting (6.72), (6.76), and the fact that ∥r0∥`2 ≤ 2√n into (6.55), (6.53), and (6.54), respectively.

heading: 6.5.2 GENERALIZATION RESULT (COMPLETING THE PROOF OF THEOREM 2.2)
text: The theorem below is the formal statement and a more general version of Theorem 2.2. The differences (besides constant terms) are the use of ∥X∥ rather than √n (we always have ∥X∥ ≤ √n) and use of the tighter bound ∥J†Iy∥`2 rather than ∥ΠI(y)∥`2 /α0. Theorem 6.24 Let ζ,Γ, ᾱ be scalars obeying ζ ≤ 1/2, Γ ≥ 1, and ᾱ ≥ 0 which determine the overall precision, cut-off and learning duration, respectively.5 Consider a training data set {(xi,yi)}ni=1 ∈ Rd × RK generated i.i.d. according to a distribution D where the input samples have unit Euclidean norm and the concatenated label vector obeys ∥y∥`2 = √n (e.g. one-hot encoding). Consider a neural net with k hidden nodes as described in (1.1) parameterized by W where the activation function φ obeys ∣φ′(z)∣ , ∣φ′′(z)∣ ≤ B. Let W0 be the initial weight matrix with i.i.d. N (0,1) entries. Fix a precision level ζ and set ν = ζ/(50B√log(2K)). Also assume the output layer V has i.i.d. Rademacher entries scaled by ν√
kK . Furthermore, set J ∶= (Σ(X))1/2 and define the
information I and nuisance N spaces and the truncated Jacobian JI associated with the Jacobian J based on a cut-off spectrum value of α0 = ᾱ 4√n√K ∥X∥B per Definition 1.1. Assume
k ≳ Γ4 logn ζ4ᾱ8
(6.87)
5Note that this theorem and its conclusions hold for any choice of these parameters in the specified range.
with Γ ≥ 1. We run gradient descent iterations of the form (1.5) with a learning rate η ≤ 1 ν2B2∥X∥2 . Then, after T = ΓK ην2α20
iterations, training loss obeys (6.69) and classification error ErrD(WT ) is upper bounded by
2∥ΠN (y)∥`2√ n´¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¸¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¶
bias term
+ 12B√K√ n
(∥J†Iy∥`2 + Γα0 ∥ΠN (y)∥`2)´¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¸¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¶ variance term +12(1+ Γ ᾱ 4 √ n∥X∥2 )ζ+5
√ log(2/δ)
n +2e−Γ,
holds with probability at least 1 − (2K)−100 − δ. The theorem below is a restatement of Theorem 6.24 after substituting the upper bound on the early stopping distance Dα0,Γ of Def. (6.1). Theorem 6.25 (Neural Net – Generalization) Consider a training data set {(xi,yi)}ni=1 ∈ Rd×RK generated i.i.d. according to a distribution D where the input samples have unit Euclidean norm and the concatenated label vector obeys ∥y∥`2 = √n (e.g. one-hot encoding). Consider a neural net with k hidden nodes as described in (1.1) parameterized by W where the activation function φ obeys∣φ′(z)∣ , ∣φ′′(z)∣ ≤ B. LetW0 be the initial weight matrix with i.i.d. N (0,1) entries. Fix a precision level ζ ≤ c
2 and set
ν = ζ 50B √ log(2K) . (6.88)
Also assume the output layer V has i.i.d. Rademacher entries scaled by ν√ kK . Furthermore, set
J ∶= Σ(X)1/2 and define the information I and nuisance N spaces and the truncated Jacobian JI associated with the Jacobian J based on a cut-off spectrum value of α0 = ᾱ 4√n√K ∥X∥B ≤ B∥X∥ per Definition 1.1 chosen to ensure ∥ΠI(y)∥`2 ≥ c∥y∥`2 for some constant c > 0. Assume
k ≥ 12 × 107 Γ4K4B8∥X∥4n2 log(n) c4ζ4α80
(6.89)
with Γ ≥ 1. We run gradient descent iterations of the form (1.5) with a learning rate η ≤ 1 ν2B2∥X∥2 . Then, after T = ΓK ην2α20
iterations, classification error ErrD(WT ) is upper bounded by ErrD(WT ) ≤ 2∥ΠN (y)∥`2 + e−Γ ∥ΠI(y)∥`2√ n +12B√K√ n Dα0,Γ+12(1+ Γ ᾱ 4 √ n∥X∥2 )ζ+10 √ log(2/δ) n , holds with probability at least 1 − (2K)−100 − δ. Proof Under the stated assumptions, Theorem 6.23 holds with probability 1 − (2K)−100. The proof will condition on outcomes of the Theorem 6.23. Specifically, we shall apply (6.43) of Lemma 6.18 with MW and R dictated by Theorem 6.23 where the output layer V is fixed. Observe that∥V ∥F = √Kk∥V ∥`∞ = ν, we have
ErrD(WT ) ≤ 2⎡⎢⎢⎢⎢⎣ ∥f(WT ) − y∥`2√ n + 3νB (MW√ n + R2√ k ) + √ 5 log(2/δ) n ⎤⎥⎥⎥⎥⎦ . (6.90) Theorem 6.23 yields
∥f(WT ) − y∥`2√ n ≤ ∥ΠN (y)∥`2 + e−Γ∥ΠI(y)∥`2√ n
+ 4ζ. (6.91) Using (6.70) for MW
νBMW√ n
≤ 2B √ KDα0,Γ(y)√ n + 2B√KΓζ α0 . (6.92)
Using (6.71) on row bound R and lower bound on k
3νB R2√ k =48nΓ2B3K2 ∥X∥2 να40 √ k
≤ c2ζ2 230νB log(n)≤ζ. (6.93)
Plugging in (6.91), (6.92), and (6.93) into (6.90) concludes the proof.

heading: A THE JACOBIAN OF THE MIXTURE MODEL IS LOW-RANK (PROOFS FOR SECTION 2.3)
text: 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 K2C information space associated with this Jacobian.
Theorem A.1 Consider a data set of size n consisting of input/label pairs {(xi,yi)}ni=1 ∈ Rd × RK generated according to the Gaussian mixture model of Definition 2.4 withK classes each consisting of C clusters with the cluster centers given by {µ`,̃̀}(K,C)(`,̃̀)=(1,1) and σ = 0. Let Σ(X) be the multiclass neural tangent kernel matrix associated with input matrix X = [x1 . . . xn]T with the standard deviation of the output layer set to ν = 1√
k . Also define the information space I to be the range space
of Σ(X). Also let M = [µ1,1 . . . µK,C]T be the matrix obtained by aggregating all the cluster centers as rows and let g be a Gaussian random vector with distributionN (0,Id). Define the neural tangent kernel matrix associated with the cluster centers as Σ̃(M) = (MMT )⊙ Eg∼N (0,Id)[φ′(Mg)φ′(Mg)T ] ∈ RKC×KC , and assume that Σ̃(M) is full rank. Then, the following properties hold with probability 1 − KC exp(− n 8KC )
• I is a K2C dimensional subspace. • The concatenated label vector y = [yT1 yT2 . . . yTn ]T lies on I. • The nonzero eigenvalues (top K2C eigenvalues) of Σ(X) are between n 2KC smin(Σ(X))
and 2n KC ∥Σ(X)∥. Hence the eigenvalues of the information space grow with n KC .
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(− n 8KC ) 0.5ñ ≤ ñ`,˜̀ ≤ 2ñ. Note that based on Lemma 6.11, the multiclass covariance is given by Σ(X) = kν2IK ⊗ Σ̃(X). where Σ̃(X) = (XXT )⊙E
g i.i.d.∼ N (0,1)[φ′(Xg)φ′(Xg)T ]. Due to this Kronecker product represen-
tation, the range space of Σ(X) is separable. In particular, note that with Ĩ = Range (Σ̃(X)) we have I = IK ⊗ Ĩ which also implies rank(I) = K ⋅ rank (Ĩ). 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:
• Ĩ has rank KC.
• The nonzero eigenvalues of Σ̃ (X) are between 0.5ñsmin(Σ̃(M)) to 2ñ∥Σ̃(M)∥. • The portion of the label vector associated with class ` i.e. y(`) ∈ Rn (see (5.4)) lies on I.
Hence, the concatenated vector y lies on I = IK ⊗ Ĩ. To prove these statements let J`(X;W0) and J`(M ;W0) be the Jacobian associated with the `th output of the neural net (see (5.4)) for data matrices X and M . Observe that the columns ofJ`(X;W0) are chosen from J`(M ;W0) and in particular each column of J`(M ;W0) is repeated between 0.5ñ to 2ñ times. To mathematically relate this, define the KC dimensional subspace S of Rn where for any v ∈ S , entries vi and vj of v are equal iff data point xi and xj are equal (i.e. belong to the same class/cluster pair). Now, we define the orthonormal matrixUS ∈ Rn×KC as the 0-1 matrix with orthogonal rows that map RKC to S as follows. Assume the ith data point xi belongs to the class/cluster pair (`i, ̃̀i). We then set the ith row of US as vect(e`ieT̃̀ i ). Using US we have USJ`(M ;W0) = J`(X;W0). Now note that using the above identity we have USΣ̃(M)UTS = Σ̃(X). Since US is tall and orthogonal, the range of Σ̃(X) is exactly the range of US hence Ĩ = S which is KC dimensional. Furthermore, nonzero eigenvectors of Σ̃(X) lie on S and any eigenvector v satisfies
vT Σ̃(X)v ≥ smin(US)2smin(Σ̃(M)) ≥ 0.5n̄smin(Σ̃(M)) and similarly vT Σ̃(X)v ≤ 2n̄∥Σ̃(M)∥ which follows from the fact that `2-norm-squared of columns of U are between 0.5n̄ to 2n̄. Finally, we will argue that label vector y(`) lies on S. Note that for all samples i that belong to the same cluster y(`)i will be the same (either zero or one), thus y(`) ∈ S. Next lemma provides a perturbation analysis when there is noise.
Lemma A.2 Consider the single-output NTK kernel given by Σ̃(X) = E [φ′(Xw)φ′(Xw)T ]⊙ (XXT ) , and assume that this matrix has rank r so that λr+1 (Σ̃(X)) = λr+2 (Σ̃(X)) = . . . = λn (Σ̃(X)) = 0. Also assume a noise corrupted version ofX given by
X̃ =X + σ√ d Z
with Z a matrix consisting of i.i.d. N (0,1) entries. Then, ∥Σ̃(X̃) − Σ̃(X)∥ ≲ ∆ where ∆ ∶= σ2B2 logn ∥X∥2 + σ2B2(n/d + 1) +√logn ⋅ σB2 ∥X∥2 + σB2√n/d + 1 ∥X∥ (A.1) holds with probability at least 1 − 2ne− d2 . Whenever σ ≤ 1√ logn , ∆ is upper bounded as
∆ n ≲ B2σ√logn. (A.2)
Furthermore, let Ṽ ,V ∈ Rn×r be orthonormal matrices corresponding to the top r eigenvalues of Σ̃(X̃) and Σ̃(X). Then,
∥Ṽ Ṽ T −V V T ∥ ≤ ∆ λr(Σ̃(X)) −∆
Proof Note that diag (φ′ (X̃w)) X̃ − diag (φ′ (Xw))X =diag (φ′ (X̃w)) X̃ − diag (φ′ (Xw))X
=diag (φ′ (X̃w) − φ′ (Xw))X + diag (φ′ (X̃w)) (X̃ −X)
Now define M̃ = diag (φ′(X̃w)) X̃ andM = diag (φ′(Xw))X and note that using the above we can conclude that
∥M̃ −M∥ ≤ ∥diag (φ′ (X̃w) − φ′ (Xw))X∥ + ∥diag (φ′ (X̃w)) (X̃ −X)∥
≤B ∥(X̃ −X)w∥ `∞ ∥X∥ +B ∥X̃ −X∥
Now using the fact that ∥M̃M̃T −MMT ∥ ≤ ∥M̃ −M∥2 + 2 ∥M̃ −M∥ ∥M∥ , we conclude that
∥Σ̃(X̃) − Σ̃(X)∥ = ∥E [M̃M̃T −MMT ]∥ ≤E [ (B ∥(X̃ −X)w∥
`∞ ∥X∥ +B ∥X̃ −X∥)2 ] + 2B ∥X∥ (B ∥X∥E [ ∥(X̃ −X)w∥
`∞ ] +B ∥X̃ −X∥) ≤2B2 ∥X∥2 E [ ∥(X̃ −X)w∥2
`∞ ] + 2B2 ∥X̃ −X∥2+ 2B2 ∥X∥2 E [ ∥(X̃ −X)w∥ `∞ ] + 2B2 ∥X̃ −X∥ ∥X∥
To proceed further, with probability 1 − n exp(−d/2), each row of X̃ −X is upper bounded by 2σ. Hence, using a standard tail bound over supremum of n Gaussian random variables (which follows by union bounding) we have
E[∥(X̃ −X)w∥2 `∞]1/2 ≤ 2σ√2 logn
holds with the same probability. Furthermore, spectral norm bound on Gaussian random matrix implies that
∥X̃ −X∥2 ≤ (2(√n +√d))2 σ2 d
≤ 8(n/d + 1)σ2, holds with probability at least 1 − e− 12 (n+d). Plugging these two probabilistic bounds into the chain of inequalities we conclude that
∥Σ̃(X̃) − Σ̃(X)∥ ≲ σ2B2 logn ∥X∥2 + σ2B2(n/d + 1) +√logn ⋅ σB2 ∥X∥2 + σB2√n/d + 1 ∥X∥ To establish (A.2), observe that B2σ √ logn∥X∥2 dominates over other terms in the regime σ√logn is small. The final bound is a standard application of Davis-Kahan Theorem Yu et al. (2014) when we use the fact that Σ̃(X) is low-rank. The following lemma plugs in the critical quantities of Theorem 2.2 for our mixture model to obtain a generalization bound. Theorem A.3 (Generalization for Mixture Model) Consider a dataset {xi,yi}ni=1 generated i.i.d. from the Gaussian mixture model in Definition 2.4. Let λM = λmin(Σ(M)) whereM ∈ RKC×d is the matrix of cluster centers. Suppose input noise level σ obeys
σ ≲ λmin B2KC √ logn
Consider the setup of Theorem 2.2 with quantities ζ and Γ. Suppose network width obeys
k ≳ Γ4B8K8C4 logn ζ4λ4min .
With probability 1 − ne−d/2 −KC exp(− n 8KC ) − (2K)−100 − δ, running gradient descent for T = 2ΓK2C ην2nλmin with learning rate η ≤ 1 ν2B2∥X∥2 , we have that
ErrD(WT ) ≲ ¿ÁÁÀσ√lognB2KC λmin + ΓBK√C√ nλmin + 12ζ + 5 √ log(2/δ) n + 2e−Γ.
Proof The proof is an application of Lemma A.2 and Theorem A.1. Let I ′ be the information space corresponding to noiseless dataset where input samples are identical to cluster centers. Let P ′,P correspond to the projection matrices to I and I ′. First, using Lemma A.2 and the bound on σ, we have
∥P ′ −P ∥ ≤ cσ√lognB2KC λmin
for some constant c > 0. Next we quantify ΠI(y) using the fact that (i) ΠI′(y) = y via Theorem A.1 as follows
∥ΠI(y)∥`2 ≥ ∥ΠI′(y)∥`2 − ∥ΠI(y) −ΠI′(y)∥`2 ≥ √n(1 − cσ √ lognB2KC
λmin ). (A.3)
In return, this implies that
∥ΠN (y)∥`2 ≲ ¿ÁÁÀnσ√lognB2KC
λmin To proceed, we pick α0 = √λminn2KC and corresponding ᾱ = α04√n√K∥X∥B ≥ √ λmin 2B2K2C
and apply (??) to find that, classification error is upper bounded by
ErrD(WT ) ≲ ¿ÁÁÀσ√lognB2KC λmin + ΓBK√C√ nλmin + 12ζ + 5 √ log(2/δ) n + 2e−Γ.

heading: B JOINT INPUT-OUTPUT OPTIMIZATION
text: In this section we wish to provide the ingredients necessary to prove a result for the case where both set of input and output weightsW and V are trained. To this aim, we consider the combined neural net Jacobian associated with input and output layers given by x↦ f(x;V ,W ) ∶= V φ(Wx). (B.1) Denoting the Jacobian associated with (B.1) by J (V ,W ) we have that J (V ,W ) = [J (V ) J (W )] ∈ RKn×k(K+d) Here, J (W ) is as before whereas J (V ) is the Jacobian with respect to V and is given by J (V ) = [J (v1) J (v2) . . . J (vK)] . (B.2) where J (v`) ∈ RKn×k is so that its `’th block rows of size n × k is nonzero for 1 ≤ ` ≤K i.e.
̃̀th block row of J (v`) = {0 if ` ≠ ̃̀ φ(XW T ) else .
Hence, J (V ) is K ×K block diagonal with blocks equal to φ(XW T ). The following theorem summarizes the properties of the joint Jacobian. Theorem B.1 (Properties of the Combined Input/Output Jacobian) J (V ,W ) satisfies the following properties.
• Upper bound: ∥J (V ,W )∥ ≤ B∥X∥(∥W ∥F +√Kk ∥V ∥`∞). • Row-bound: For unit length u: ∥mat (J T (W )u)∥2,∞ ≤ B√K∥V ∥`∞∥X∥. • Entry-bound: For unit length u: ∥mat (J T (V )u)∥`∞ ≤ B∥W ∥2,∞∥X∥. • Lipschitzness: Given inputs V ,V ′ and outputsW ,W ′
∥J (V ,W )−J (V ′,W ′)∥ ≤ B ∥X∥ (√Kk ∥V −V ′∥`∞+√K ∥V ∥`∞ ∥W ′ −W ∥F+∥W −W ′∥F ).
Proof First, we prove results concerning J (V ). First, note that ∥J (V )∥ ≤ ∥φ(XW T )∥ ≤ B∥X∥∥W ∥F . Next, note that for u = [u1 . . . uK] ∈ RKn we have ∥J T (V )u∥`∞ = max
1≤`≤K ∥φ(WXT )u`∥`∞= max 1≤s≤k ∣φ(wsXT )u`∣= B∥W ∥2,∞∥X∥.
Let J1,J2 be the Jacobian matrices restricted to V andW of J (V ,W ). To prove Lipschitzness, first observe that
∥J (V ,W ) −J (V ′,W ′)∥ ≤ ∥J1(V ,W ) −J1(V ′,W ′)∥ + ∥J2(V ,W ) −J2(V ′,W ′)∥. Next, observe that ∥J1(V ,W ) −J1(V ′,W ′)∥ ≤ ∥φ(XW T ) − φ(XW ′T )∥ ≤ B∥X∥∥W −W ′∥F . We decompose J2 via ∥J2(V ,W ) −J2(V ′,W ′)∥ ≤ ∥J2(V ,W ) −J2(V ,W ′)∥ + ∥J2(V ,W ′) −J2(V ′,W ′)∥
≤ B√K ∥V ∥`∞ ∥X∥ ∥W ′ −W ∥F + ∥J2(V ,W ′) −J2(V ′,W ′)∥. To address the second term, note that, Jacobian is linear with respect to output layer hence ∥J2(V ,W ′) −J2(V ′,W ′)∥ = ∥J2(V −V ′,W ′)∥ ≤ B√Kk ∥V −V ′∥`∞ ∥X∥ . Combining the latter two identities we arrive at ∥J2(V ,W ) −J2(V ′,W ′)∥ ≤ B ∥X∥ (√Kk ∥V −V ′∥`∞ +√K ∥V ∥`∞ ∥W ′ −W ∥F ), completing the proof.

heading: C FURTHER NUMERICAL EXPERIMENTS
text: We depict the approximate eigenstructure in Section C.1. Complementary to Section 4, in Appendix C.2 we provide the complete set of experiments we performed on the original 10-class CIFAR-10 dataset covering various levels of label corruption. In Appendix C.3 we show numerical results on a sub-sampled 3-class version of CIFAR-10. Moreover, we demonstrate that our theory holds across different datasets by providing experimental results on MNIST in Appendix C.4.
C.1 EXPERIMENTS ON FULL (10-CLASS) CIFAR-10
As mentioned earlier while calculating all the eigenvalues is not possible, we also verify the bimodal structure of the Jacobian using the entire spectrum by approximating its spectral density using the method recently described in Ghorbani et al. (2019a) and Papyan (2019a) for Hessian eigenvalue density estimation. We detail the parameters of the algorithm in Appendix C.5. Figure 9 depicts the estimated spectrum before and after training. We observe a similar bimodal spectrum with a few outliers. Moreover, in both depictions of the Jacobian eigenstructure we observe increasing separation between the low and high end of the spectrum after training.
C.2 EXPERIMENTS ON FULL (10-CLASS) CIFAR-10
In addition to what has been described in Section 4 here we disclose the complete set of experiments performed on the CIFAR-10 dataset. We trained the modified ResNet20 model described in Section 4 with SGD as longs as it was necessary to achieve a good fit to the training data. Information subspace
is spanned by the top 50 singular vectors. We marked figures and table entries also included in Section 4 by asterisk.
Experiments without label corruption. We trained the modified ResNet20 model described in Section 4 with SGD for 400 epochs with learning rate 0.1 on the original dataset without any form of data augmentation. The network output after the last layer has been scaled by s = 0.025 in all experiments to follow.
Experiments with label corruption. We corrupt the labels in the training data by switching each label to a strictly different (incorrect) class. We train the network on the corrupted dataset for 800 epochs with initial step size 0.1 decayed to 0.01 after 760 epochs.
0 200 400 600 800 10−6 10−5 10−4 10−3 10−2 10−1 100
Epochs
R es id u a l en
er gy
1 (a) Residual projection on initial train Jacobian
0 200 400 600 800
10−6 10−5 10−4 10−3 10−2 10−1 100
Epochs
R es id u a l en
er gy
∥ rτ ∥2`2∥r0∥2`2∥ΠI(rτ )∥2`2∥r0∥2`2∥ΠN (rτ )∥2`2∥r0∥2`2
1 (b) Residual projection on train Jacobian at 50 epochs (dip)∗
0 200 400 600 800
10−6 10−5 10−4 10−3 10−2 10−1 100
Epochs
R es id u al
en er gy
1 (c) Residual projection on final train Jacobian
0 200 400 600 800 10−2
10−1
100
Epochs
R es id u al en er gy ∥ rτ ∥2`2∥r0∥2`2∥ΠI(rτ )∥2`2∥r0∥2`2∥ΠN (rτ )∥2`2∥r0∥2`2
1 (d) Residual projection on initial test Jacobian
0 200 400 600 800
10−1
100
Epochs
R es id u al
en er gy
1 (e) Residual projection on test Jacobian at 50 epochs (dip)∗
0 200 400 600 800
10−2
10−1
100
Epochs
R es id u al en er gy ∥ rτ ∥2`2∥r0∥2`2∥ΠI(rτ )∥2`2∥r0∥2`2∥ΠN (rτ )∥2`2∥r0∥2`2
1 (f) Residual projection on final test Jacobian
C.3 EXPERIMENTS ON SUBSAMPLED 3-CLASS CIFAR-10
We created a new dataset of 3 classes by sub-sampling the original CIFAR-10 dataset. To do this, we discarded all training and test data except those belonging to classes 0 (airplane), 1 (automobile) and 2 (bird) and sampled 3333 examples of each classes for a total of 9999 training images. We trained the neural network model described in Section 4 with no output scaling using SGD and Adam. We applied standard data augmentation (random crop and flip) to increase generalization. Information subspace is spanned by the top 50 singular vectors.
Experiments without label corruption. First we trained the network on the sub-sampled 3-class dataset keeping the original labels. For the SGD experiments, we used initial learning rate 0.01 decreased by a factor of 10 at epochs 260 and 360 for a total of 400 epochs with batch size 128. For the Adam experiment learning rate 0.01 for 400 epochs was sufficient for a good fit to the training data. We observed better Jacobian adaptation using Adam compared to SGD on this dataset (0.98743 from Table 8 vs. 0.9969 from Table 9).
Experiments with label corruption. We corrupt the training labels by the corruption model described in C.2. We train the network on the corrupted dataset for 800 epochs with initial step size 0.01 and batch size 128. We decrease the learning rate by a factor of 10 at the following epochs: at 500 epochs for 25% corruption, at 700 epochs for 50% and 75% corruption and at 500 and 700 for 100% corruption.
0 200 400 600 800
10−4
10−3
10−2
10−1
100
Epochs
R es id u a l en
er gy
1 (a) Residual projection on initial train Jacobian
0 200 400 600 800
10−4
10−3
10−2
10−1
100
Epochs
R es id u a l en
er gy
∥ rτ ∥2`2∥r0∥2`2∥ΠI(rτ )∥2`2∥r0∥2`2∥ΠN (rτ )∥2`2∥r0∥2`2
1 (b) Residual projection on train Jacobian at 100 epochs (dip)
0 200 400 600 800
10−3
10−2
10−1
100
Epochs
R es id u al
en er gy
1 (c) Residual projection on final train Jacobian
0 200 400 600 800
10−1
100
Epochs
R es id u al en er gy ∥ rτ ∥2`2∥r0∥2`2∥ΠI(rτ )∥2`2∥r0∥2`2∥ΠN (rτ )∥2`2∥r0∥2`2
1 (d) Residual projection on initial test Jacobian
0 200 400 600 800
10−1
100
Epochs
R es id u al
en er gy
1 (e) Residual projection on test Jacobian at 100 epochs (dip)
0 200 400 600 800
10−1
100
Epochs
R es id u al en er gy ∥ rτ ∥2`2∥r0∥2`2∥ΠI(rτ )∥2`2∥r0∥2`2∥ΠN (rτ )∥2`2∥r0∥2`2
1 (f) Residual projection on final test Jacobian
C.4 EXPERIMENTS ON MNIST
The MNIST dataset contains handwritten digits in 10 classes divided into 60k training and 10k test images. To demonstrate our theoretical findings on a dataset different from CIFAR-10 we repeat all experiments on MNIST under various levels of label corruption. We run SGD with batch size 128 on least-squares loss using the modified ResNet20 model from Section 4. In all of the following MNIST experiments the model output after the last layer has been scaled by 0.1. Information subspace is spanned by the top 50 singular vectors
Experiments without label corruption. We train the network on the original dataset for 100 epochs with initial step size 0.1 decayed to 0.01 after 60 epochs.
Experiments with label corruption. We corrupt various portions of the labels in the training data by switching the labels to a random incorrect class. We train the network on the corrupted dataset for 400 epochs with initial step size 0.1 decayed to 0.01 after 360 epochs.
C.5 SINGULAR VALUE DECOMPOSITION OF THE JACOBIAN
The matrix representation of the Jacobian arising in the numerical experiments is in general very large and pose a computational challenge. In particular, the Jacobian matrix for a ResNet20 model is 500k×270k (or 540 GB memory with float32 entries, more than 1 TB with float64) on the CIFAR-10 training data. Storing these matrices in memory for singular value decomposition is not feasible and therefore other methods are necessary to approximate the singular values and singular vectors.
The Lanczos algorithm Lanczos (1950) is an adaptation of power methods, which guarantees fast convergence to the largest eigenvalues of a hermitian matrix. More importantly, it only requires matrix-vector products to estimate the top eigenvectors and therefore resolves the memory issues with large Jacobians. Using the ARPACK implementation of the implicitly restarted Lanczos method Lehoucq et al. (1998) it takes about 1 hour on 64 cores to find the top 50 eigenvalues and eigenvectors of the 500k × 500k symmetric matrix JJT for CIFAR-10 training data. The same algorithm has been used to calculate the histogram of top 1000 singular values.
The required computation time for Lanczos method is dominated by large Jacobian-vector products, which increases linearly with the number of calculated eigenvalues. Therefore it is necessary to depart from deterministic power iteration type algorithms and take a different approach. Stochastic Lanczos quadrature algorithm has been recently introduced to the machine learning community in Ghorbani et al. (2019a) and Papyan (2019a) to approximate the eigenvalue density of large neural network Hessians. Finding the exact spectral density might be difficult therefore first the problem is relaxed by approximating the spectrum convolved with a Gaussian density of variance σ2 = 10−5. As described in Ghorbani et al. (2019a), we draw k = 10 i.i.d. Gaussian random vectors and run m = 80 steps of the Lanczos algorithm starting from each random vector to estimate the parameters of the spectral density. We modified the algorithm by replacing Hessian-vector products by JJT -vector products. Finally, we normalize the density and scale it by the dimension of the Jacobian output space.

heading: D JACOBIAN ADAPTATION FOR LINEAR DATA
text: In this section, we will (non-rigorously) argue that Jacobian and its singular value spectrum indeed adapts to the data for simple linear datasets. Specifically, we will show that the very first gradient iteration from a random initialization is already powerful enough to encourage an approximately low-rank Jacobian. Dataset: Let (xi)ni=1 i.i.d.∼ N (0,Id) and fix θ ∈ Rd. Fix labels yi = θTxi and the dataset (xi, yi)ni=1. Neural net: Let W0 ∈ Rk×d i.i.d.∼ N (0,1) and fix v ∈ Rk with half +σ and half −σ entries. Set activation to be ReLU. (Non-rigorous) Statement: With small initialization (σ ≈ 0) and a single large gradient step, W1 = W0 − η∇L(W0) has a (properly scaled) Jacobian matrix with approximately rank 2d. Instead, at random initializationW0, Jacobian has full rank n (assuming kd≫ n). Specifically, atW1, we have
σ2d+1(J (X,W1)) ≲√d/n and σ2d(J (X,W1)) ≳ 1. (D.1) where σi is the ith largest singular value. In words, J (X,W1) is approximately rank 2d when n ≳ d. This is in stark contrast to the properties of Jacobian at random initialization. For instance Li & Liang (2018); Zou et al. (2018); Oymak & Soltanolkotabi (2019) actually show that Jacobian is well conditioned with minimum singular value lower bounded by the minimum separation between the samples (implying full rank of n). Such separation (independent of n) holds for Gaussian data even when n is polynomially larger than d via standard concentration / packing bounds (also see Oymak & Soltanolkotabi (2019); Li & Liang (2018)). Hence, as further formalized below, we argue that even a single step of gradient descent can lead to an approximately low-rank Jacobian.
D.1 APPROACH
Fixing weights ofW0, we study the gradient of the random data. Let wr be rth row ofW0. Fixing hidden node r, let x′i = xi if xTi wr ≥ 0 (i.e. the rth node is active) and 0 else. The gradient of the node r atW0 is given by
∇L(W ) ∂wr = vr n n∑ i=1(f(xi,W0) − yi)x′i.
If initialization is small enough (σ ≈ 0), we can approximate f(xi,W0) − yi ≈ yi hence we study ∇L(W ) ∂wk ≈ vr n n∑ i=1yix ′ i. We will show that gradients of each row are ±σθ/2 in expectation. Decompose a samplex ∼ N (0,Id) as x = g wr∥wr∥`2 + g where g is standard normal vector over the orthogonal complement of vector wr and g ∼ N (0,1). Denoting label associated to x by y, consequently, we find
E[yx′] = E[yx ∣ xTwr ≥ 0]P(xTwr ≥ 0) (D.2) = 1
2 E[(g wr∥wr∥`2 + g)(gwr + g)Tθ ∣ g ≥ 0] (D.3)
= 1 2 E[ggT ]θ + E[g2 ∣ g ≥ 0]wrwTr θ (D.4) = θ
2 . (D.5)
Hence
E[∇L(W ) ∂wk ] ≈ viθ/2.
This implies that, in population (n→∞), gradient vector has an extremely simple rank-one form: EX[∇L(W ) ∣ W0] = 12vθT .
Standard results on non-asymptotic statistics then yield that with exponentially high-probability (i.e. 1 − exp(−cn)),
∥∇L(W ) ∂wk − EX[∇L(W ) ∂wk ]∥`2 ≲√d/n As long as k ≲ poly(n), union bounding over all rows, we find the spectral norm bound
∥∇L(W ) ∣ W0 − EX[∇L(W ) ∣ W0]∥ ≤ σ√kd/n On the other hand, EX[∇L(W ) ∣ W0] is rank 1 with nonzero singular value σ√k/2. Consequently, a large gradient step will ensure that W1 =W0 − η∇L(W ) ∣ W0 is approximately rank one as well. Specifically, set the rank one matrix W̄1 = EX[η∇L(W ) ∣ W0]. Then the tail of the spectrum is at most ∥W1 −W̄1∥ ≲√d/n∥W̄1∥ (i.e. √d/n as large as the spectral norm).
Finally, we need to move from showing low-rankness of the weight matrix to low-rankness of Jacobian. Define the Jacobians associated with data at W̄1 andW1 by
J̄ = J (X,W̄1) = diag(v)φ′(W̄1XT ) ∗XT , J = J (X,W1) = diag(v)φ′(W1XT ) ∗XT The Jacobian associated with J̄ is exactly rank 2d since matrix diag(v)φ′(W̄1XT ) has exactly two distinct rows and kronecker producting withXT will yield d rank for each (adding up to 2d). Since X is well-conditioned with high probability (as it is N (0,1) and fat with n ≳ d), top 2d singular values and σ2d(J̄ ) are strictly positive with a lower bound independent of n and σ (after scaling by 1/σ∥X∥). Finally, using smoothness of φ′ (which excludes ReLU but includes smoother versions of ReLU such as softplus), using Jacobian perturbation bound (6.31), we have (after scaling by 1/σ∥X∥).
∥J − J̄ ∥ ≲ ∥W1 − W̄1∥ ≲√d/n. Together, this leads to the desired conclusion (D.1) since singular value perturbation inequality yields σ2d+1(J ) ≤ ∥J − J̄ ∥ ≲√d/n (as σ2d+1(J̄ ) = 0).

Review summaries:
The authors use the empirically supported assumption of the low-rank nature of the neural network Jacobian to provide new elements of data-dependent optimization and generalization theory. By modelling the data as a low-rank object itself, they analytically study the evolution of the train and test errors. The paper divides the space of weights and biases into the “information” and “nuisance” subspaces, spanned by the top largest and the remaining singular vectors of the Jacobian respectively. They use this division and its alignment with the low-rank structure of the data to talk about convergence speed. Finally, they provide numerical experiments to back their claims.