Title: GENERALIZATION GUARANTEES FOR NEURAL NETS VIA HARNESSING THE LOW-RANKNESS OF JACOBIAN Anonymous

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.

1. INTRODUCTION 1.MOTIVATION AND CONTRIBUTIONS:

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. 

1.2. MODEL AND TRAINING:

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 ∈ R k×d and V ∈ R K×k to denote the input-to-hidden and hiddento-output weights. The overall input-output relationship of the network is a function f (⋅; W ) ∶ R d → R K that maps an input x ∈ R d to an output via

x ↦ f (x; W ) ∶= V φ(W x).

(1.1)

Given a dataset consisting of n feature/label pairs (x i , y i ) with x i ∈ R d representing the features and y i ∈ R K the associated labels representing one of K classes with one-hot encoding (i.e. y i ∈ {e 1 , e 2 , . . . , e K } where e ∈ R K 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 over W via1 min

W ∈R k×d L(W ) ∶= 1 2 n i=1 V φ (W x i ) -y i 2 2 .

(1.

2)

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

y = ⎡ ⎢ ⎢ ⎢ ⎢ ⎣ y 1 ⋮ y n ⎤ ⎥ ⎥ ⎥ ⎥ ⎦ ∈ R nK and f (W ) = ⎡ ⎢ ⎢ ⎢ ⎢ ⎣ V f (x 1 ; W ) ⋮ V f (x n ; W ) ⎤ ⎥ ⎥ ⎥ ⎥ ⎦ ∈ R nK . (1.3)

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

W ∈R k×d L(W ) ∶= 1 2 f (W ) -y 2 2 .

(1.4)

To optimize this loss starting from an initialization W 0 we run gradient descent iterations of the form

EQUATION

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.

1.3. INFORMATION AND NUISANCE SPACES:

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 τ ) ∈ R kd , 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) T r(w). Here, J (w) ∈ R nK×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 ∈ R nK×p with singular value decomposition given by J = nK s=1 λ s u s v T s = U diag (λ 1 , λ 2 , . . . , λ nK ) V T , with λ 1 ≥ λ 2 ≥ . . . ≥ λ nK denoting the singular values of J in decreasing order and {u s } nK s=1 ∈ R nK and {v s } nK s=1 ∈ R p the corresponding left and right singular vectors forming the orthonormal basis matrices U ∈ R nK×nK and V ∈ R p×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({u s } r s=1 ) and N ∶= span({u s } Kn s=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 (W 0 )J T (W 0 )]

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.

2. MAIN RESULTS:

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 + e z ). 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 ∈ R n×d we use s min (X) and s max (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 ∈ R n×n with eigenvalue decomposition A = ∑ n i=1 λ i u i u T i , the square root matrix is defined as A 1 2 ∶= ∑ n i=1

√ λ i u i u T i . 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 ⊂ R n and point x ∈ R n , Π 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) ∈ R d × R K from a distribution D, the classification error of the network W with respect to D is defined as

Err D (W ) = P arg max 1≤ ≤K y ≠ arg max 1≤ ≤K f (x; W ) .

(2.1)

2.1. RESULTS FOR RANDOM INITIALIZATION:

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 ∈ R d be a vector with N (0, I d ) distribution. Consider a set of n input data points x 1 , x 2 , . . . , x n ∈ R d aggregated into the rows of a data matrix X ∈ R n×d . Associated to the activation φ and the input data matrix X we define the multiclass kernel matrix as

Σ(X) ∶= I K ⊗ E φ ′ (Xw) φ ′ (Xw) T ⊙ XX T ,

where I K 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 weights W 0 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 (W 0 )J (W 0 ) 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. ∼ 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

Consider an i.i.d. training dataset {(x i , y i )} n i=1 ∈ R d × R K with

k ≳ Γ 4 log n α 8 . (2.2)

Then after T ∝ Γ α 2 gradient iterations of (1.5), with high probability, training loss obeys

f (W T ) -y 2 ≲ Π N (y) 2 + e -Γ √ n, . (2.3)

Furthermore, the classification error obeys

Err D (W T ) ≲ Π 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 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 = R Kn and Π N (y) 2 = 0. For this case, the more general Theorem 6.24 yields

Err D (W T ) ≲ √ K √ n y T Σ -1 (X)y while requiring a width of k ≳ K 4 n 4 log n λ 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 (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)).

2.2. GENERALIZATION GUARANTEES WITH ARBITRARY INITIALIZATION:

Our next result provides generalization guarantees from an arbitrary initialization which applies to pre-trained networks (e.g. those that arise in transfer learning applications) as well as intermediate gradient iterates as the weights evolve. This result has a similar flavor to Theorem 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 {(x i , y i )} n i=1 ∈ R d ×R K 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 W 0 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 (W 0 ) based on spectrum cutoff α √ n. Set the initial residual r 0 = f (W 0 ) -y ∈ R nK and assume r 0 2 ≲ √ n. Suppose k ≳ Γ 4 α 8 . After T ∝ Γ α 2 iterations (1.5) with constant learning rate, training loss obeys: f (W T ) -y 2 ≲ Π N (r 0 ) 2 + e -Γ √ n.

Also with high probability, classification error obeys:

Err D (W T ) ≲ Π 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 W 0 = 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.

2.3. CASE STUDY: GAUSSIAN MIXTURE MODEL:

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 {(x i , y i )} n i=1 ∈ R d × R K . 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 µ , ̃ ∈ R d 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 µ , ̃ , σ 2 I d d with the corresponding label set to the one-hot encoding of the class i.e. e . Note that the cluster indexed by ( , ̃ ) contains ñ , ̃ data points satisfying E[ñ , ̃ ] = ñ = 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

K 2 C = # of output nodes × # of distinct inputs.

as there are (i) only KC distinct input vectors and (ii) K output nodes. We can thus set the information space to be the top K 2 C eigenvectors of the multiclass kernel matrix Σ(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, log n). 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 {(x i ,

y i )} n i=1 ∈ R d × R K 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, I d ). Also let Σ(M ) ∈ R KC×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 ≳ Jacobian spectrum KC n J (W 0 ) 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.

Γ 4 K 8 C 4 λ 4 M . after T = 2ΓK 2 C λ M gradient iterations, with high probability, the model obeys Err D (W T ) ≲ Γ K 2 C nλ M .

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 ≳ K 2 C 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 K 2 C . 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 ≳ K 8 C 4 . This compares favorably with Arora et al. (2019) which concerns the K = 1 case. In particular, Arora et al. (2019) requires k ≳ n 8 λ 6

X 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.

3. PRIOR ART:

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) 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. We present experiments demonstrating our theoretical findings on two popular image classification datasets.

4. NUMERICAL EXPERIMENTS:

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 † I y 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 † I y 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.

Π I (y) 2 y 2 Π N (y) 2 y 2 Π I (r0) 2 r0 2 Π N (

Π I (y) 2 y 2 Π N (y) 2 y 2 Π I (r0) 2 r0 2 Π N (r0) 2 r0 2 J train init 0.

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. 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 on N 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. (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 form X = X + Z where X is low-rank (i.e. rank(X) = r << n) with singular value decomposition X = U ΣV T with U ∈ R n×r , Σ ∈ R r×r , V ∈ R d×r , and Z ∈ R n×d is a matrix with i.i.d. N (0, σ 2

x n) entries. We shall also assume the labels are equal to y = y + z with y = Xw * for some w * ∈ Range(V ) and z ∈ R n a Gaussian random vector with i.i.d. N (0, σ 2 y 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 σ 2 x w 2 2 + 1 2 σ 2 y ,

Now let us consider gradient descent iterations with a step size η which take the form

w τ +1 = w τ -η∇L(w τ ) = I -ηX T X w τ + ηX T y.

(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, Σ = I r 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 misfit Xwy on the column space of the uncorrupted portion of the input data (U ) and its complement. That is,

E L(w) = E L I (w) + E L N (w).

where

E L I (w) ∶= E Π I (Xw -y) 2 2 = Xw -y 2 2 + r 2n σ 2 x w 2 2 + r 2n σ 2 y , E L N (w) ∶= E Π N (Xw -y) 2 2 = 1 2 1 - r n σ 2 x w 2 2 + σ 2 y , with Π I = U U T and Π N = I -U U T .

In Figure 8b we plot these two components. This plot clearly shows that E L I (w) goes down quickly while E L N (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 E L I (w) as the test error on this information subspace. Similarly, one can think of the complement of this subspace as the "nuisance" subspace and E L N (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 -ηXX T r τ = I -ηX X T Xw τ -y + noise

Based on the form of this update when the information subspace is closely aligned with the prominent singular vectors of X the test error on the information subspace (E

L I (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 of X (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 ∈ R d denoting the input features, θ ∈ R p denoting the parameters, and f (x; θ) ∈ R K the K outputs of the model denoted by f 1 (x; θ), f 2 (x; θ), . . . , f K (x; θ). For instance in the case of neural networks θ represents its weights. Given a data set of n input/label pairs {(x i , y i )} n i=1 ⊂ R d × R K , we fit to this data by minimizing a nonlinear least-squares loss of the form

L(θ) = 1 2 n i=1 f (x i ; θ) -y i 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 (x 1 ; θ) ⋮ f (x n ; θ) ⎤ ⎥ ⎥ ⎥ ⎥ ⎦ ∈ R n and y ( ) = ⎡ ⎢ ⎢ ⎢ ⎢ ⎣ (y 1 ) ⋮ (y n ) ⎤ ⎥ ⎥ ⎥ ⎥ ⎦ ∈ R n for = 1, 2, . . . , K.

Concatenating these vectors we arrive at

f (θ) = ⎡ ⎢ ⎢ ⎢ ⎢ ⎣ f 1 (θ) ⋮ f K (θ) ⎤ ⎥ ⎥ ⎥ ⎥ ⎦ ∈ R Kn and y = ⎡ ⎢ ⎢ ⎢ ⎢ ⎢ ⎣ y (1) ⋮ y (K) ⎤ ⎥ ⎥ ⎥ ⎥ ⎥ ⎦ ∈ R Kn .

(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 ∈ R p and the Jacobian mapping J (θ 0 ) ∈ R Kn×p . For ε 0 , β > 0, we call J ∈ R Kn×max(Kn,p) an (ε 0 , β) reference Jacobian matrix if it obeys the following conditions,

J ≤ β, J (θ 0 )J T (θ 0 ) -J J T ≤ ε 2 0 , and J (θ 0 ) -J ≤ ε 0 .

where J (θ 0 ) ∈ R Kn×max(Kn,p) is a matrix obtained by augmenting J (θ 0 ) with max(0, Knp) zero columns. Furthermore, consider the singular value decomposition of J given by

J = U diag(λ)V T = Kn s=1 λ s u s v T s .

(5.7)

where λ ∈ R Kn are the vector of singular values and u s ∈ R Kn and v s ∈ R p 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 Kn,p) is obtained from θ 0 by adding max(Knp, 0) zero entries at the end of θ 0 . The optimization problem for fitting the linearized problem has the form

f lin ( θ) = f (θ 0 ) + J ( θ -θ 0 ), where θ 0 ∈ R max(

L lin (θ) = 1 2 f lin (θ) -y 2 2 .

(5.8)

Thus starting from θ0 = θ 0 the iterates θτ on the linearized problem take the form θτ+1 = θτ -η∇L lin ( θτ ), (5.9)

= θτ -ηJ T (f (θ 0 ) + J ( θτ -θ 0 ) -y), = θτ -ηJ T J ( θτ -θ 0 ) -ηJ T (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τ = f lin ( θτ )y = (I -ηJ J T ) τ r 0 .

(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 = U diag(λ)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({u s } r s=1 ) and N ∶= span({u s } Kn s=r+1 ). We also define the truncated reference Jacobian

J I = [u 1 u 2 . . . u r ] diag (λ 1 , λ 2 , . . . , λ r ) [v 1 v 2 . . . v r ]

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 θ ∈ R p the Jacobian mapping associated with the nonlinearity f ∶ R p ↦ R n has bounded spectrum, i.e. J (θ) ≤ β.

Assumption 2 (Bounded perturbation) Consider a point θ 0 ∈ R p 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 ∶ R p ↦ R nK

the multi-class nonlinear mapping, θ ∈ R p the parameters of the model, and y ∈ R nK the concatenated labels as in (5.4). Let θ be zero-padding of θ till size max(Kn, p). Also, consider a point θ 0 ∈ R p with J an ( 0 , β) reference Jacobian associated with J (θ 0 ) per Definition 5.1 and fitting the linearized problem

f lin ( θ) = f (θ 0 ) + J ( θ -θ 0 ) via the loss L lin (θ) = 1 2 f lin (θ) -y 2 2

. Furthermore, define the information I and nuisance N subspaces and the truncated Jacobian J I associated with the reference Jacobian J based on a cut-off spectrum value of α per Definition 5.2. Furthermore, assume the Jacobian mapping J (θ) ∈ R nK×p associated with f obeys Assumptions 1 and 2 for all θ ∈ R p obeying

θ -θ 0 2 ≤ R ∶= 2 J † I r 0 2 + Γ α Π N (r 0 ) 2 + δ Γ α r 0 2 , (5.12)

around a point θ 0 ∈ R p 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 = θτ -η∇L lin ( θτ ) 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τ ∶= f lin ( θτ )y closely track each other. That is,

r τ -rτ 2 ≤ 3 5 δα β r 0 2 and θ τ -θτ 2 ≤ δ Γ α r 0 2 (5.14)

Furthermore, for all iterates τ obeying 0

≤ τ ≤ T ∶= Γ ηα 2 θ τ -θ 0 2 ≤ R 2 = J † I r 0 2 + Γ α Π N (r 0 ) 2 + δ Γ α r 0 2 .

(5.15)

and after τ = T iteration we have

r T 2 ≤ e -Γ Π I (r 0 ) 2 + Π N (r 0 ) 2 + δα β r 0 2 .

(5.16)

6. PROOFS:

Before we proceed with the proof let us briefly discuss some notation used throughout. For a matrix W ∈ R k×d we use vect(W ) ∈ R kd to denote a vector obtained by concatenating the rows

w 1 , w 2 , . . . , w k ∈ R d of W . That is, vect(W ) = w T 1 w T 2 . . . w T k T

. Similarly, we use mat(w) ∈ R k×d to denote a k × d matrix obtained by reshaping the vector w ∈ R kd across its rows. Throughout, for a differentiable function φ ∶ R ↦ R we use φ ′ and φ ′′ to denote the first and second derivative.

6.1. PROOFS FOR GENERAL THEORY (PROOF OF THEOREM 5.3):

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.

6.1.1. PRELIMINARY DEFINITIONS AND LEMMAS:

Throughout we use r) .

U I = [u 1 u 2 . . . u r ] ∈ R nK×r and U N = [u r+1 u r+2 . . . u nK ] ∈ R nK×(nK-

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 r 0 = 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 λ 2 s (⟨u s , r 0 ⟩) 2 + Γ 2 nK s=r+1 λ 2 s α 2 (⟨u s , r 0 ⟩) 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 (r 0 ) 2 2 + Γ 2 Π N (r 0 ) 2 2 1 2 ≤ Γ r 0 2 (6.2) B α,Γ ≥ α λ 1 Π I (r 0 ) 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 (⟨u s , r 0 ⟩) 2 + Γ 2 nK s=r+1 (⟨u s , r 0 ⟩) 2 1 2 = Π I (r 0 ) 2 2 + Γ 2 Π N (r 0 ) 2 2 1 2 ≤ Γ r 0 2 .

To prove the lower bound, we use the facts that

α 2 λ 2 s ≥ α 2 λ 2 1 to conclude that B α,Γ = r s=1 α 2 λ 2 s (⟨u s , r 0 ⟩) 2 + Γ 2 nK s=r+1 λ 2 s α 2 (⟨u s , r 0 ⟩) 2 1 2 , ≥ r s=1 α 2 λ 2 s (⟨u s , r 0 ⟩) 2 1 2 , ≥ α λ 1 Π I (r 0 ) 2 .

It is of course well known that the mapping I -ηAA T is a contraction for sufficiently small values of η. The next lemma shows that if we replace one of the A matrices with a matrix B which is close to A the resulting matrix I -ηAB T , while may not be contractive, is not too expansive.

Lemma 6.3 (Asymmetric PSD increase) Let A, B ∈ R n×p be matrices obeying

A ≤ β, B ≤ β, and B -A ≤ ε.

Then, for all r ∈ R n and η ≤ 1 β 2 we have

I -ηAB T r 2 ≤ 1 + ηε 2 r 2 .

Proof Note that using η ≤ 1 β 2 and B -A ≤ ε we conclude that

I -ηAB T r 2 2 = I -ηBB T -η(A -B)B T r 2 2 = r -η(A -B + B)B T r 2 2 = r 2 2 -2ηr T (A -B + B)B T r + η 2 AB T r 2 2 ≤ r 2 2 -2η B T r 2 2 + 2η (A -B) T r 2 B T r 2 + η 2 A 2 B T r 2 2 = r 2 2 -η B T r 2 2 + 2η (A -B) T r 2 B T r 2 + η 2 A 2 B T r 2 2 -η B T r 2 2 η≤1 β 2 ≤ r 2 2 -η B T r 2 2 + 2η (A -B) T r 2 B T r 2 A-B ≤ ≤ r 2 2 -η B T r 2 2 + 2ηε B T r 2 r 2 =(1 + ηε 2 ) r 2 2 -η ε r 2 -B T r 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 any X ∈ R n×p with p ≥ n obeying A = XX T , there exists a matrix

Y ∈ R n×p obeying B = Y Y T such that Y -X ≤ 2α

Proof First we note that for any two PSD matrices Schmitt (1992) 

A + , B + ∈ R n×n obeying A + , B + ⪰ α 2 4 I n , Lemma 2.2 of

guarantees that A 1 2 + -B 1 2 + ≤ A + -B + α .

In the above for a PSD matrix A ∈ R n×n with an eigenvalue decomposition A = U ΛU T we use A 1 2 to denote the square root of the matrix given by A = U Λ 1 2 U T . We shall use this result with

A + = A + α 2 4 I n and B + = B + α 2 4 I n to conclude that A 1 2 + -B 1 2 + ≤ A + -B + α = A -B α .

Furthermore, using the fact that the eigenvalues of A + and B + are just shifted versions of the eigenvalues of A and B by α 2 4 we can conclude that

A 1 2 + -A 1 2 ≤ α 2 and B 1 2 + -B 1 2 ≤ α 2 .

Combining the latter two inequalities with the assumption that A -B ≤ α 2 we conclude that

A 1 2 -B 1 2 ≤ A 1 2 + -B 1 2 + + A 1 2 + -A 1 2 + B 1 2 + -B 1 2 ≤ A -B α + α 2 + α 2 ≤2α.

(6.4) Suppose p ≥ n and assume the matrices A and B have eigenvalue decompositions given by A

= U A Λ A U T A and B = U B Λ B U T B . Then, any X ∈ R n×p with p ≥ n has the form X = U A Λ 1 2 A V T A with V A ∈ R p×n an orthonormal matrix. Now pick Y = U B Λ 1 2 B U T B U A V T A . Then clearly Y Y T = B. Furthermore, we have X -Y = U A Λ 1 2 A V T A -U B Λ 1 2 B U T B U A V T A = U A Λ 1 2 A U T A U A V T A -U B Λ 1 2 B U T B U A V T A = U A Λ 1 2 A U T A -U B Λ 1 2 B U T B U A V T A = A 1 2 -B 1 2 U A V T A = A 1 2 -B 1 2 .

Combining the latter with (6.4) completes the proof.

6.1.2. KEY LEMMAS FOR GENERAL NONLINEARITIES:

Throughout this section we assume J is the reference Jacobian per Definition 5.1 with eigenvalue decomposition

J = U ΛV T = ∑ Kn s=1 λ s u s v T s with Λ = diag(λ).

We also define a = U T r 0 = U T r0 ∈ R nK 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ηλ 2 s ) τ a s u s . (6.5)

Furthermore, assuming η ≤ 1 λ 2 1 the linear updates θτ obey θτ -θ0

2 2 ≤ r s=1 a 2 s λ 2 s + τ 2 η 2 nK s=r+1 λ 2 s a 2 s . (6.6) Proof Using the fact that J J T = U Λ 2 U T we have (I -ηJ J T ) τ = U I -ηΛ 2 τ U T

Using the latter combined with (5.11) we thus have rτ =(I -ηJ J T ) τ r 0 ,

=U I -ηΛ 2 τ U T r 0 , =U I -ηΛ 2 τ a, = nK s=1 (1 -ηλ 2 s ) τ a s u s ,

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

J T rt = V ΛU T rt = V Λ I -ηΛ 2 t a.

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

θτ -θ0 = -η τ -1 t=0 ∇L lin ( θt ) = -η τ -1 t=0 J T rt = -ηV τ -1 t=0 Λ I -ηΛ 2 t a.

Thus for any s ∈ {1, 2, . . . , nK}

v T s θτ -θ0 = -ηλ s a s τ -1 t=0 1 -ηλ 2 s t = -ηλ s a s 1 -1 -ηλ 2 s τ ηλ 2 s = -a s 1 -1 -ηλ 2 s τ λ s .

Noting that for η ≤ 1 λ 2 1 ≤ 1 λ 2 s we have 1ηλ 2 s ≥ 0, the latter identity implies that

v T s θτ -θ0 ≤ a s λ s . (6.7)

Furthermore, using the fact that 1ηλ 2 s ≤ 1 we have

v T s θτ -θ0 = ηλ s a s τ -1 t=0 1 -ηλ 2 s t

≤ ηλ s a s τ (6.8)

Combining (6.7) for 1 ≤ s ≤ r and (6.8) for s > r we have θτ -θ0

2 2 = nK s=1 v T s θτ -θ0 2 ≤ r s=1 a 2 s λ 2 s + τ 2 η 2 nK s=r+1 λ 2 s a 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 (r 0 ) 2 + Π N (r 0 ) 2 . (6.9)

Furthermore, after T = Γ ηα 2 iterations we have

rT 2 ≤ e -Γ Π I (r 0 ) 2 + Π N (r 0 ) 2 .

(6.10) and θT -θ0

2 2 ≤ r s=1 a 2 s λ 2 s + Γ 2 nK s=r+1 λ 2 s a 2 s α 4 = B 2 α,Γ α 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

U T I rτ = I -ηΛ 2 I τ U T I r0 and U T N rτ = I -ηΛ 2 N τ U T N r0

Thus, using the fact that for s ≤ r we have λ s ≥ α we have (1ηλ 2 s ) τ ≤ (1ηα 2 ) τ and for s > r we have (1ηλ 2 s ) τ ≤ 1, we can conclude that

U T I rτ 2 ≤ 1 -ηα 2 τ U T I r 0 2 and U T N rτ 2 ≤ U T N r 0 2 .

Combining these with the triangular inequality we have

rτ 2 = U T I rτ U T N rτ 2 ≤ U T I rτ 2 + U T N rτ 2 ≤ 1 -ηα 2 τ U T I r 0 2 + U T N r 0 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 subspace N . 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 (r 0 ) 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 ≤ η(ε 2 0 + εβ) rτ 2 + (1 + ηε 2 ) e τ 2 . (6.11)

Proof For simplicity, denote

B 1 = J (θ τ +1 , θ τ ), B 2 = 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 -ηB 1 B T 2 )r τ .

Similarly, for linearized problem we have rτ+1 = (I -ηJ J T )r τ . Thus,

e τ +1 2 = (I -ηB 1 B T 2 )r τ -(I -ηJ J T )r τ 2 = (I -ηB 1 B T 2 )e τ -η(B 1 B T 2 -J J T )r τ 2 ≤ (I -ηB 1 B T 2 )e τ 2 + η (B 1 B T 2 -J J T )r τ 2 ≤ (I -ηB 1 B T 2 )e τ 2 + η (B 1 B T 2 -J J T ) rτ 2 . (6.12)

We proceed by bounding each of these two terms. For the first term, we apply Lemma 6.3 with

A = B 1 and B = B 2 and use B 1 -B 2 ≤ ε to conclude that (I -ηB 1 B T 2 )e τ 2 ≤ (1 + ηε 2 ) e τ 2 .

(6.13)

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

B 1 B T 2 -J J T = B 1 B T 2 -AA T + AA T -J J T ≤ B 1 B T 2 -AA T + AA T -J J T ≤ (B 1 -A)B T 2 + A(B 2 -A) T + AA T -J J T ≤ B 1 -A B 2 + B 2 -A A + AA T -J J T ≤ β ε 2 + β ε 2 + ε 2 0 = ε 2 0 + εβ. (6.14)

In the last inequality we use the fact that per Assumption 2 we have B 1 -A ≤ ε 2 and B 2 -A ≤ ε 2 as well as the fact that per Definition 5.1 AA T -J J T ≤ ε 2 0 . 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 e 0 = 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

e t+1 ≤(1 + ηε 2 )e t + ηΘr t ≤e t + ηε 2 e t + ηΘ (1 -ηα 2 ) t ρ + + ρ - ≤e t + ηΘ ε 2 Λ + (1 -ηα 2 ) t ρ + + ρ -.

Thus

e t+1 -e t Θ ≤ η ε 2 Λ + (1 -ηα 2 ) t ρ + + ρ -.

(6.17)

Summing up both sides of (6.17) for 0 ≤ t ≤ τ -1 we conclude that

e τ Θ = τ -1 t=0 e t+1 -e t Θ ≤ ητ ε 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.

6.1.3. COMPLETING THE PROOF OF THEOREM 5.3:

With the key lemmas in place in this section we wish to complete the proof of Theorem 5.3. We will use induction to prove the result. Suppose the statement is true for some τ -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) ≤ η J T rτ-1 2 + η J (θ τ -1 ) -J rτ-1 2 + η J (θ τ -1 ) r τ -1 -rτ-1 2 (c) ≤ η J T rτ-1 2 + ε 0 + ε β 2 rτ-1 2 + 1 β r τ -1 -rτ-1 2 (d) ≤ η J T rτ-1 2 + 2δα 5β 2 r 0 2 + 1 β r τ -1 -rτ-1 2 (e) ≤ η J T rτ-1 2 + 2δα 5β 2 r 0 2 + 3δα 5β 2 r 0 2 =η J T rτ-1 2 + δα β 2 r 0 2 (f ) ≤ ηβ 2 B α,Γ α + δα β 2 r 0 2 (g) ≤ B α,Γ α + δα β 2 r 0 2 (h) ≤ B α,Γ α + δΓ α r 0 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

J T rτ-1 2 2 = J T I -ηJ J T τ -1 r 0 2 2 = V Λ I -ηΛ 2 τ -1 U T r 0 2 2 = Λ I -ηΛ 2 τ -1 U T r 0 2 2 = nK s=1 λ 2 s (1 -ηλ 2 s ) 2(τ -1) (⟨u s , r 0 ⟩) 2 ≤ nK s=1 λ 2 s (⟨u s , r 0 ⟩) 2 = r s=1 λ 2 s (⟨u s , r 0 ⟩) 2 + nK s=r+1 λ 2 s (⟨u s , r 0 ⟩) 2 ≤β 4 r s=1 1 λ 2 s (⟨u s , r 0 ⟩) 2 + nK s=r+1 λ 2 s (⟨u s , r 0 ⟩) 2 ≤β 4 r s=1 1 λ 2 s (⟨u s , r 0 ⟩) 2 + Γ 2 nK s=r+1 λ 2 s α 4 (⟨u s , r 0 ⟩) 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 ≤ η(ε 2 0 + εβ) rτ-1 2 + (1 + ηε 2 ) e τ -1 2 . This combined with the induction assumption implies that (6.18) holds for all t ≤ τ ≤ T . Furthermore, using Lemma 6.5 equation (6.9) for all t ≤ τ ≤ T we have

e t 2 ≤ η(ε 2 0 + εβ) rt-1 2 + (1 + ηε 2 ) e t-1 2 ,

rt 2 ≤ 1 -ηα 2 t Π I (r 0 ) 2 + Π N (r 0 ) 2 , (6.19)

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

Θ ∶= ε 2 0 + εβ, ρ + = Π I (r 0 ) 2 , ρ -= Π N (r 0 )

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(ε 2 0 + εβ) (6.22) where in the last inequality we used (5.13). This completes the first part of (5.14) via induction.

( Π I (r 0 ) 2 + Γ Π N (r 0 ) 2 ) α 2 ≤ 2Γ(ε 2 0 + εβ) r 0 2 α 2 (6.21) ≤( 2 25 + 2 5 ) δα β r 0 2 ≤ 3 5 δα β r 0 2 ,

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 where J 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 ) -∇L lin ( θt ) 2 = τ -1 t=0 J (θ t ) T r t -J T rt 2 ≤ τ -1 t=0 J (θ t ) T r t -J T rt 2 ≤ τ -1 t=0 (J (θ t ) -J ) T rt 2 + J (θ t ) T (r t -rt ) 2 = τ -1 t=0 (J (θ t ) -J ) T rt 2 + J (θ t ) T e t 2 ≤ τ -1 t=0 (ε + ε 0 ) rt 2 + β e t 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 ≤ r 0 2 to conclude

τ -1 t=0 rt 2 ≤ τ r 0 2 ≤ T r 0 2 = Γ r 0 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 e t 2 ≤ τ 2(εβ + ε 2 0 ) α 2 Γ r 0 2 ≤ 2Γ 2 (εβ + ε 2 0 ) ηα 4 r 0 2 . (6.26)

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

θ τ -θτ 2 ≤ 2Γ(εβ 2 + ε 2 0 β) α 3 + ε + ε 0 α Γ α r 0 2 = ε 2Γβ 2 α 3 + ε 2 0 2Γβ α 3 + ε + ε 0 α Γ α r 0 2 (a) ≤ 2 5 δ + ε 2 0 2Γβ α 3 + ε + ε 0 α Γ α r 0 2 (b) ≤ 2 5 δ + 2 25 δ + ε + ε 0 α Γ α r 0 2 (c) ≤ 2 5 δ + 2 25 δ + 1 5 δ + ε 0 α Γ α r 0 2 (d) ≤ 2 5 δ + 2 25 δ + 1 5 δ + 1 5 δ Γ α r 0 2 = 22 25 δ α Γ r 0 2 .

Here, (a) follows from ε ≤ δα 3 5Γβ 2 per Assumption (5.13), (b) from ε 0 ≤ 1 5 δα 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 ≤ δ α Γ r 0 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 α,Γ α + δ α Γ r 0 2 ≤ J † I r 0 2 + Γ α Π N (r 0 ) 2 + δ α Γ r 0 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

r T 2 (a) ≤ rT 2 + rT -r T 2 (b) ≤ rT 2 + δα β r 0 2 (c) ≤ e -Γ Π I (r 0 ) 2 + Π N (r 0 ) 2 + δα β r 0 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).

6.2. KEY LEMMAS AND IDENTITIES FOR NEURAL NETWORKS:

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

For two matrices

A = ⎡ ⎢ ⎢ ⎢ ⎢ ⎢ ⎢ ⎣ A 1 A 2 ⋮ A p ⎤ ⎥ ⎥ ⎥ ⎥ ⎥ ⎥ ⎦ ∈ R p×m and B = ⎡ ⎢ ⎢ ⎢ ⎢ ⎢ ⎢ ⎣ B 1 B 2 ⋮ B p ⎤ ⎥ ⎥ ⎥ ⎥ ⎥ ⎥ ⎦ ∈ R p×n ,

we define their Khatri-Rao product as

A * B = [A 1 ⊗ B 1 , . . . , A p ⊗ B p ] ∈ R p×mn

, where ⊗ denotes the Kronecker product.

6.2.1. PRELIMINARY IDENTITIES AND CALCULATIONS:

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 of M . We begin by noting that for a one-hidden layer neural network of the form x ↦ V φ (W x), the Jacobian matrix with respect to vect(W ) ∈ R kd takes the form

J (W ) = ⎡ ⎢ ⎢ ⎢ ⎢ ⎣ J 1 (W ) ⋮ J K (W ) ⎤ ⎥ ⎥ ⎥ ⎥ ⎦ ∈ R Kn×kd (6.27)

where J (W ) is the Jacobian matrix associated with the th class. In particular, J (W ) is given by

J (W ) = [J (w 1 ) . . . J (w k )] ∈ R n×kd with J (w s ) ∶= V ,s diag (φ ′ (Xw s )) X.

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

J (W ) = φ ′ XW T diag(v ) * X.

(6.28) An alternative characterization of the Jacobian is via its matrix representation. Given a vector u ∈ R Kn let us partition it into K size n subvectors so that u

= [u T 1 . . . u T K ] T . We have mat J T (W )u = K =1 diag(v )φ ′ W X T diag(u )X.

(6.29)

6.2.2. FUNDAMENTAL PROPERTIES OF THE JACOBIAN OF THE NEURAL NETWORK:

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

Lemma 6.9 (Properties of Single Output Neural Net Jacobian) Let K = 1 so that V T = v ∈ R n . Suppose φ is an activation obeying φ ′ (z) ≤ B for all z. Then, for any W ∈ R k×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 ( x i 2 = 1). Then the Jacobian mapping is Lipschitz with respect to spectral norm i.e. for all W , W ∈ R k×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)φ ′ W X T diag(u)X 2,∞ ≤ v ∞ max 1≤ ≤k φ ′ w T X T diag(u)X 2 ≤ v ∞ X max 1≤ ≤k φ ′ w T X T 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 any W ∈ R k×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 ( x i 2 = 1). Then the Jacobian mapping is Lipschitz with respect to spectral norm i.e. for all W , W ∈ R k×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 = [A T 1 . . . A T K ] T and B = [B T 1 . . . B T K ] 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 )φ ′ W X T diag(u )X 2,∞ ≤ K =1 diag(v )φ ′ W X T diag(u )X 2,∞ ≤ K =1 B 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.

6.2.3. PROPERTIES OF THE JACOBIAN AT RANDOM INITIALIZATION:

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

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 matrix M , for 1 ≤ , ̃ ≤ K, let M [ , ̃ ] denote the ( , ̃ )th submatrix. For C(W ) = J (W )J (W ) T we have

C(W )[ , ̃ ] = k s=1 (XX T ) ⊙ (V ,s V ̃ ,s φ ′ (Xw s )φ ′ (Xw s ) T ). Suppose W 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

I K ⊗ [(XX T ) ⊙ E[φ ′ (Xw s )φ ′ (Xw s ) T ]].

Proof The ( , ̃ )th submatrix of C(W ) is given by

C(W )[ , ̃ ] = ((diag(v )φ ′ (W X T )) * X T )((diag(v ̃ )φ ′ (W X T )) * X T ) T = k s=1 J (w s )J ̃ (w s ) T = k s=1 V ,s V ̃ ,s (diag(φ ′ (Xw s ))X)(diag(φ ′ (Xw s ))X) T = k s=1 V ,s V ̃ ,s (XX T ) ⊙ (φ ′ (Xw s )φ ′ (Xw s ) T ) = k s=1 XX T ⊙ V ,s V ̃ ,s φ ′ (Xw s )φ ′ (Xw s ) T . (6.33) Setting W 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 (XX T ) ⊙ (E[V ,s V ̃ ,s ] E[φ ′ (Xw s )φ ′ (Xw s ) T ]) = k s=1 ν 2 δ( -̃ )[(XX T ) ⊙ E[φ ′ (Xw s )φ ′ (Xw s ) 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) ) Let A, B ∈ R n×n be two Positive Semi-Definite (PSD) matrices. Then,

λ min (A ⊙ B) ≥ min i B ii λ min (A) , λ max (A ⊙ B) ≤ max i B ii λ 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 φ (W x) where the activation φ obeys φ(0) ≤ B and φ ′ (z) ≤ B for all z. Also assume we have n ≥ K data points x 1 , x 2 , . . . , x n ∈ R d with unit euclidean norm ( x i 2 = 1) aggregated as the rows of a matrix X ∈ R n×d . Furthermore, suppose V has i.i.d. ν-scaled Rademacher entries (i.e. ±ν equally-likely). Then, the Jacobian matrix at a random point W 0 ∈ R k×d with i.i.d. N (0, 1) entries obeys

J (W 0 )J (W 0 ) T -E[J (W 0 )J (W 0 ) T ] ≤ 30K √ kν 2 B 2 X 2 log(n).

with probability at least 1 -1 n 100 . In particular, as long as

k ≥ 1000K 2 B 4 X 4 log(n) δ 2 ,

with the same probability, we have that

1 kν 2 J (W 0 )J (W 0 ) T -Σ(X) ≤ δ.

Proof Define C = J (W 0 )J (W 0 ) 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

A s = φ ′ (Xw s ) φ ′ (Xw s ) T ⊙ XX T . Now consider n × n diagonal blocks of C (denoted by C[ , ]) and note that we have C[ , ] = φ ′ XW T diag(v )diag(v )φ ′ W X T ⊙ XX T = k s=1 V 2 ,s A s =ν 2 k s=1 A s .

Furthermore, using Lemma 6.12

A s ≤ max i φ ′ (x T i w s ) 2 X 2 ≤ B 2 X 2 .

Also, using Jensen's inequality

E[A s ] ≤ E A s ≤ B 2 X 2 .

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

(A s -E[A s ]) 2 = A s -E[A s ] 2 ≤ ( A s + E[A s ] ) 2 ≤ 2B 2 X 2 2 .

(6.34) To proceed, we will bound the weighted sum

S = k s=1 ν 2 (A s -E[A s ])

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

P( S ≥ t) ≤ 2ne -t 2 2∆ 2 , where ∆ 2 is an upper bound on ∑ k s=1 ν 4 (A s -E[A s ]) 2 . Using (6.34) we can pick ∆ 2 = ∑ k s=1 (2ν 2 B 2 X 2 ) 2 = 4kν 4 B 4 X 4 . Setting t = 30 √ kν 2 B 2 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=1 V ,s V ̃ ,s A s .

Note that by construction {V ,s V ̃ ,s } k s=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=1 V 2 ,s V 2 ̃ ,s A 2 s = ν 4 k s=1 A 2 s ≤ ν 4 k s=1 A s 2 ≤ ν 4 kB 4 X 4 ,

so that we can take ∆ 2 = ν 4 kB 4 X 4 and again conclude that for t = 30 √

kν 2 B 2 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 ≤ K 2 n -102 ≤ n -100 .

All that remains is to combine the concentration results for the sub-matrices to arrive at the complete bound. In 

[ , ∶] ≤ √ K sup 1≤ ̃ ≤K D[ , ̃ ] ≤ √ Kt ⇒ D ≤ √ K sup 1≤ ≤K D[ , ∶] ≤ √ K √ Kt = Kt = 30K √ kν 2 B 2 X 2 log(n),

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

6.2.4. UPPER BOUND ON INITIAL RESIDUAL:

In this section we prove a lemma concerning the size of the initial misfit. The proof of this lemma (stated below) follows from a similar argument in the proof of (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 φ (W x) where the activation φ has bounded derivatives obeying φ(0) , φ ′ (z) ≤ B. Also assume we have n data points x 1 , x 2 , . . . , x n ∈ R d with unit euclidean norm ( x i 2 = 1) aggregated as rows of a matrix X ∈ R n×d and the corresponding labels given by y ∈ R Kn . Furthermore, assume the entries of V are i.i.d. Rademacher variables scaled by

ν y 2 50B K log(2K)kn

and the entries of W ∈ R k×d are i.i.d. N (0, 1). Then,

V φ W X T F ≤ ν y 2 ,

holds with probability at least 1 -(2K) -100 .

Proof We begin the proof by noting that

V φ W X T 2 F = K =1 v T φ W X T 2 2 .

We will show that for any row v of V , with probability at least 1

-(2K) -101 , v T φ W X T 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 ∈ R k×d we have

φ X W T v 2 -φ XW T v 2 ≤ φ X W T v -φ XW T v 2 ≤ φ X W T -φ XW T v 2 ≤ φ X W T -φ XW T F v 2 (a) = φ ′ S ⊙ X W T + (1 k×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 ≤ S i,j ≤ 1 and 1 k×n the matrix of all ones. Thus, (6.36) holds with probability at least 1e

φ XW T v 2 is a B X v 2 -Lipschitz function of W . Thus, fixing v, for a matrix W with i.i.d. Gaussian entries φ XW T v 2 ≤ E φ XW T v 2 + t,

- t 2 2B 2 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)) ≤ B 2 . (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 (6.38) holds with probability 1 -(2K) -102 . Conditioned on E v , we now upper bound the expectation via

E v ∶= { 1 T v 2 ≤ 250 log K v 2 2 }

E φ XW T v 2 (a) ≤ E φ (XW T ) v 2 2 = n i=1 E (v T φ(W x i )) 2 (b) = √ n E g∼N (0,I k ) (v T φ(g)) 2 (c) = √ n v 2 2 E g∼N (0,1) (φ(g) -E[φ(g)]) 2 + (1 T v) 2 (E g∼N (0,1) [φ(g)]) 2 (d) ≤ √ n v 2 250 × 4B 2 log(2K) + B 2 ≤32 n log(2K)B v 2 .

Here, (a) follows from Jensen's inequality, (b) from linearity of expectation and the fact that for x i with unit Euclidean norm W x i ∼ N (0, I k ), (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 E v we conclude that (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 E v and the conditional concentration over W , the overall probability of success in (6.39) is at least 1 -(2K) -101 concluding the proof of (6.35) and the Lemma.

φ XW T v 2 ≤50 n log(2K)B v 2 = 50 n log(2K)B √ k ν y 2 50B K log(2K)kn = ν y 2 √ K ,

6.3. RADEMACHER COMPLEXITY AND GENERALIZATION BOUNDS:

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. Let F be a function class. Suppose f ∈ F maps R d to R K . Let {ε i } n i=1 be i.i.d. vectors in R K with i.i.d. Rademacher variables. Given i.i.d. samples S = {(x i , y i )} n i=1 ∼ D, we define the empirical Rademacher complexity to be

R S (F) = 1 n E sup f ∈F n i=1 ε T i f (x i ) .

We begin by stating a vector contraction inequality by Maurer Maurer (2016) . This is obtained by Maurer (2016) . 

setting h i (f (x i )) = h(y i , f (x i )) in Corollary 4 of

Lemma 6.15 Let f (⋅) ∶ R d → R K and let ∶ R K ×R K → R

i )} n i=1 , define R S ( , F) = E sup f ∈F n i=1 ε i (y i , f (x i )) .

We have that R S ( , F) ≤ √ 2R S (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 (⋅, ⋅) ∶ R K × R K → [0, 1] be a 1 Lipschitz loss function. Given i.i.d. samples {(x i , y i )} n i=1 , consider the empirical loss L(f, ) = 1 n n i=1 (y i , f (x i )).

With probability 1δ over the samples, for all f ∈ F, we have that

E[L(f, )] ≤ L(f, ) + 2 √ 2R S (F) + 5 log(2 δ) n

Proof Based on Bartlett & Mendelson (2002) ,

E[L(f, )] ≤ L(f, ) + 2R S ( , 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 φ (W x) with W ∈ R k×d and V ∈ R K×k denoting the input and output weight matrices. Suppose V 0 ∈ R K×k is a matrix obeying V 0 ∞ ≤ ν √ kK. Also let W 0 ∈ R k×d be a reference input weight matrix. Furthermore, we define the neural network function space parameterized by the weights as follows

F V,W = f (x; V , W ) such that V ∈ V and W ∈ W with V = V ∶ V -V 0 F ≤ νM V √ Kk and W = W ∶ W -W 0 F ≤ M W and W -W 0 2,∞ ≤ R √ k . (6.40)

Additionally, assume the training data {(x i , y i )} n i=1 are generated i.i.d. with the input data points of unit Euclidean norm (i.e. x i 2 = 1). Also, define the average energy at W 0 as

E = 1 kn n i=1 φ(W 0 x i ) 2 2 1 2

.

Also let {ξ i } n i=1 ∈ R K be i.i.d. vectors with i.i.d. Rademacher entries and define the empirical Rademacher complexity

R S (F V,W ) ∶= 1 n E ⎡ ⎢ ⎢ ⎢ ⎢ ⎣ sup f ∈F V,W n i=1 ξ T i f (x i ) ⎤ ⎥ ⎥ ⎥ ⎥ ⎦ .

Then,

R S (F V,W ) ≤ νB M W + EM V √ n + R 2 + M W M V √ k . (6.41)

Proof We shall use w to denote the rows of W (same for W 0 , V , V 0 ). We will approximate φ (⟨w , x i ⟩) by its linear approximation φ ⟨w 0 , x i ⟩ +φ ′ (⟨w 0 , x i ⟩) ⟨ww 0 , x i ⟩ via the second order Taylor's mean value theorem. We thus have

R S (F V,W ) ≤ 1 n E ⎡ ⎢ ⎢ ⎢ ⎢ ⎣ n i=1 ξ T i V 0 φ (W 0 x i ) ⎤ ⎥ ⎥ ⎥ ⎥ ⎦ + 1 n E ⎡ ⎢ ⎢ ⎢ ⎢ ⎣ sup W ∈W n i=1 ξ T i V 0 diag (φ ′ (W 0 x i )) (W -W 0 ) x i ⎤ ⎥ ⎥ ⎥ ⎥ ⎦ R1 + 1 2n E ξi,j i.i.d. ∼ ±1 ⎡ ⎢ ⎢ ⎢ ⎢ ⎣ sup W ∈W n i=1 k =1 K j=1 ξ T i,j V 0 j, φ ′′ (1 -t i )⟨w 0 , x i ⟩ + t i ⟨w , x i ⟩ ⟨w -w 0 , x i ⟩ 2 ⎤ ⎥ ⎥ ⎥ ⎥ ⎦ R2 + 1 n E sup V ∈V,W ∈W n i=1 ξ T i (V -V 0 )(φ(W x i ) -φ(W 0 x i )) R3 + 1 n E sup V ∈V n i=1 ξ T i (V -V 0 )φ(W 0 x i ) R4

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

R 1 ≤ 1 n E ⎡ ⎢ ⎢ ⎢ ⎢ ⎣ sup W -W0 F ≤M W n i=1 ξ T i V 0 diag (φ ′ (W 0 x i )) (W -W 0 ) x i ⎤ ⎥ ⎥ ⎥ ⎥ ⎦ ≤ 1 n E ⎡ ⎢ ⎢ ⎢ ⎢ ⎣ sup W -W0 F ≤M W ⟨ n i=1 diag (φ ′ (W 0 x i )) V T 0 ξ i x T i , W -W 0 ⟩ ⎤ ⎥ ⎥ ⎥ ⎥ ⎦ ≤ M W n E ⎡ ⎢ ⎢ ⎢ ⎢ ⎣ n i=1 diag V T 0 ξ i φ ′ (W 0 x i ) x T i F ⎤ ⎥ ⎥ ⎥ ⎥ ⎦ ≤ M W n E ⎡ ⎢ ⎢ ⎢ ⎢ ⎣ n i=1 diag V T 0 ξ i φ ′ (W 0 x i ) x T i 2 F ⎤ ⎥ ⎥ ⎥ ⎥ ⎦ 1 2 = M W n ⎡ ⎢ ⎢ ⎢ ⎢ ⎣ n i=1 E diag V T 0 ξ i φ ′ (W 0 x i ) x T i 2 F ⎤ ⎥ ⎥ ⎥ ⎥ ⎦ 1 2 = M W n ⎡ ⎢ ⎢ ⎢ ⎢ ⎣ n i=1 E diag V T 0 ξ i φ ′ (W 0 x i ) 2 2 ⎤ ⎥ ⎥ ⎥ ⎥ ⎦ 1 2 ≤ BM W n ⎡ ⎢ ⎢ ⎢ ⎢ ⎣ n i=1 E V T 0 ξ i 2 2 ⎤ ⎥ ⎥ ⎥ ⎥ ⎦ 1 2 ≤ BM W n V 0 F ≤ BM W ν √ n ,

where in the last inequality we used the fact that V 0 F ≤ ν. For the second term note that

R 2 ≤ 1 2n E ⎡ ⎢ ⎢ ⎢ ⎢ ⎣ sup W -W0 2,∞≤R n i=1 k =1 K j=1 ξ i,j v 0,j, φ ′′ (1 -t i )⟨w 0 , x i ⟩ + t i ⟨w , x i ⟩ ⟨w -w 0 , x i ⟩ 2 ⎤ ⎥ ⎥ ⎥ ⎥ ⎦ ≤ 1 2n k =1 E ⎡ ⎢ ⎢ ⎢ ⎢ ⎣ sup w -w 0 2 ≤R n i=1 K j=1 ξ i,j v 0,j, φ ′′ (1 -t i )⟨w 0 , x i ⟩ + t i ⟨w , x i ⟩ ⟨w -w 0 , x i ⟩ 2 ⎤ ⎥ ⎥ ⎥ ⎥ ⎦ ≤ 1 2kn k =1 n i=1 E ⎡ ⎢ ⎢ ⎢ ⎢ ⎣ K j=1 ξ i,j v 0,j, ⎤ ⎥ ⎥ ⎥ ⎥ ⎦ R 2 B ≤ BR 2 2k V T 0 2,1 ≤ BR 2 ν 2 √ k .

In the above we used M 2,1 for a matrix M to denote the sum of the Euclidean norm of the rows of M . We also used the fact that V T 0 2,1 ≤ ν √ k. To bound the third term note that

R 3 = 1 n E sup V ∈V,W ∈W n i=1 ξ T i (V -V 0 )(φ(W x i ) -φ(W 0 x i )) ≤ 1 n E sup V ∈V,W ∈W n i=1 V -V 0 F ξ i 2 (φ(W x i ) -φ(W 0 x i )) 2 ≤ νM V n √ kK E sup W ∈W n i=1 ξ i 2 φ(W x i ) -φ(W 0 x i ) 2 ≤ νM V n √ k ⋅ sup W ∈W n i=1 φ(W x i ) -φ(W 0 x i ) 2 ≤ νBM V n √ k ⋅ sup W ∈W n i=1 (W -W 0 )x i 2 ≤ νBM V n √ k ⋅ sup W ∈W n i=1 (W -W 0 ) F = νBM V √ k ⋅ sup W ∈W (W -W 0 ) F = νBM V M W √ k .

Finally, to bound the fourth term note that we have

R 4 = 1 n E sup V ∈V n i=1 ξ T i (V -V 0 )φ(W 0 x i ) = 1 n E sup V -V0 F ≤ νM V √ kK ⟨ n i=1 ξ i φ(W 0 x i ) T , (V -V 0 )⟩ = νM V n √ kK E n i=1 φ(W 0 x i )ξ T i F ≤ νM V n √ kK ⎛ ⎝ E n i=1 φ(W 0 x i )ξ T i 2 F ⎞ ⎠ 1 2 = νM V n √ kK n i=1 E φ(W 0 x i )ξ T i 2 F 1 2 = νM V √ n 1 kn n i=1 E φ(W 0 x i ) 2 F 1 2 = νEM V √ n ≤ νBEM V √ n

Combining these four bounds we conclude that

R S (F V,W ) ≤ νB M W √ n + R 2 √ k + M V M W √ k + EM V √ 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 ∈ R d ↦ f (x; V , W ) = V φ (W x) ∈ R K with W ∈ R k×d

denoting the input-to-hidden weights and V ∈ R k×K the hidden-to-output weights. Suppose V 0 ∈ R K×k is a matrix obeying V 0 ∞ ≤ ν √ kK. Also let W 0 ∈ R k×d be a reference input weight matrix. Also define the empirical losses

L(V , W ) = 1 n n i=1 y i -f (x i ; V , W ) 2 2 ,

and

L(f, ) = 1 n n i=1 (f (x i ; V , W ), y i ), with ∶ R K × R K → [0, 1]

a one Lipschitz loss function obeying (y, y) = 0. Additionally, assume the training data {(x i , y i )} n i=1 are generated i.i.d. according to a distribution D with the input data points of unit Euclidean norm (i.e. x i 2 = 1). Also, define the average energy at W 0 as

E = 1 kn n i=1 φ(W 0 x i ) 2 2 1 2

.

Then for all f in the function class F W given by (6.40)

E[L(f, )] ≤ L(V , W ) + 2 √ 2νB M W √ n + R 2 √ k + νBM V M W √ k + EM V √ 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) ∈ R d × R K generated according to the distribution D, define the population classification error Err D (W ) = P(arg max

1≤ ≤K y i ≠ arg max 1≤ ≤K f i (x, W )).

Then, we also have

Err D (W ) ≤ 2 ⎡ ⎢ ⎢ ⎢ ⎢ ⎣ L(V , W ) + 2 √ 2νB M W √ n + R 2 √ k + νBM V M W √ k + EM V √ n + 5 log(2 δ) n ⎤ ⎥ ⎥ ⎥ ⎥ ⎦ . (6.43)

Proof To begin first note that any 1-Lipschitz with (y, y) = 0 obeys (y, ŷ) ≤ yŷ 2 . Thus, we have

L(f, ) ≤ 1 n n i=1 y i -f (x i ; 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 ⋅ R S (F) + 5 log(2 δ) n , ≤ L(W ) + 2 √ 2νB M W √ n + R 2 √ k + νBM V M W √ k + EM V √ 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, ŷ) = min(1, yŷ 2 ).

Note that, given a sample (x, y) ∈ R d × R K with one-hot encoded labels, if arg max

1≤ ≤K y ≠ arg max 1≤ ≤K f (x; V , W ),

this implies (y, f (x; V , W )) ≥ 0.5.

Combining the latter with Markov inequality we arrive at Err:

D (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.

6.4. PROOFS FOR NEURAL NETS WITH ARBITRARY INITIALIZATION (PROOF OF THEOREM 2.3):

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.

6.4.1. PRELIMINARY OPTIMIZATION RESULT:

Lemma 6.19 (Deterministic convergence guarantee) Consider a one-hidden layer neural net of the form x ↦ f (x; W ) ∶= V φ(W x) with input weights W ∈ R k×d and output weights V ∈ R K×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 over W based on the loss

L(W ) = 1 2 n i=1 f (x i ; W ) -y i 2 2 .

Also, consider a point W 0 ∈ R k×d with J an ( 0 , νB X ) reference Jacobian associated with J (W 0 ) per Definition 5.1. Furthermore, define the information I and nuisance N subspaces and the truncated Jacobian J I associated with the reference Jacobian J based on a cut-off spectrum value of α per Definition 5.2. Let the initial residual vector be r 0 = yf (W 0 ) ∈ R nK . Furthermore, assume

ε 0 ≤ α 5 min δ, δα ΓνB X (6.44) and k ≥ 400 ν 6 B 6 X 6 Γ 2 δ 2 α 8 (B α,Γ + δΓ r 0 2 ) 2 , (6.45)

with 0 ≤ δ ≤ 1 and Γ ≥ 1. We run gradient descent iterations of the form

W τ +1 = W τ -η∇L(W τ ) starting from W 0 with step size η obeying η ≤ 1 ν 2 B 2 X 2 . Then for all iterates τ obeying 0 ≤ τ ≤ T ∶= Γ ηα 2 W τ -W 0 F ≤ B α,Γ α + δ Γ α r 0 2 . (6.46) W τ -W 0 2,∞ ≤ 2νBΓ X √ kα 2 r 0 2 . (6.47)

Furthermore, after τ = T iteration we have

r T 2 ≤ e -Γ Π I (r 0 ) 2 + Π N (r 0 ) 2 + δα νB X r 0 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Γν 2 B 2 X 2 = δα 3 5Γβ 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 all W ∈ R k×d obeying

W -W 0 F ≤ R ∶= 2 B α,Γ α + δ Γ α r 0 2

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

J (W ) -J (W 0 ) ≤B √ K V ∞ X R ≤ ν √ k B X R = δα 3 10Γν 2 B 2 X 2 20Γν 3 B 3 X 3 δα 4 (B α,Γ + δΓ r 0 2 ) √ k ≤ δα 3 10Γν 2 B 2 X 2 = ε 2 . Thus, Assumption 2 holds with W -W 0 F ≤ R ∶= 2 B α,Γ α + δ Γ α r 0 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 of W τ . Also note that ∇L(w ) = th row of mat(J (W ) T r τ ).

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 r t 2 . (6.49)

To bound the right-hand side we use the triangular inequality combined with (6.25) and (6.26) to conclude that (6.50) where in the last inequality we used the fact that ε 2 0 ≤ α 2 25Γ per (6.44) and β = δα 3 5Γβ ≤ α 2 5Γ per our choice of . Combining (6.49) and (6.50), we obtain

η τ -1 t=0 r τ 2 ≤η τ -1 t=0 rτ 2 + η τ -1 t=0 r τ -rτ 2 ≤ Γ α 2 r 0 2 + 2Γ 2 (εβ + ε 2 0 ) α 4 r 0 2 = 2Γ(ε 2 0 + εβ) + α 2 α 4 Γ r 0 2 ≤2 Γ α 2 r 0 2 ,

w (τ ) -w (0) 2 ≤ 2νB X Γ √ kα 2 r 0 2 ,

completing the proof of (6.47) and the theorem.

6.4.2. MAIN OPTIMIZATION RESULT:

Lemma 6.20 (Deterministic optimization guarantee) Consider the setting and assumptions of Lemma 6.19. Also assume Π I (r 0 ) 2 ≥ c r 0 2 for a constant c > 0 if ε 0 > 0. Furthermore, assume

EQUATION

and Γ ≥ 1. We run gradient descent iterations of the form

W τ +1 = W τ -η∇L(W τ ) starting from W 0 with step size η obeying η ≤ 1 ν 2 B 2 X 2 . Then for all iterates τ obeying 0 ≤ τ ≤ T ∶= Γ ηα 2 W τ -W 0 F ≤ 2B α,Γ α . (6.53) W τ -W 0 2,∞ ≤ 2νBΓ X √ kα 2 r 0 2 . (6.54)

Furthermore, after τ = T iteration we have

f (W T ) -y 2 ≤ e -Γ Π I (r 0 ) 2 + Π N (r 0 ) 2 + ζ r 0 2 . (6.55)

Proof To prove this lemma we aim to substitute

EQUATION

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 + Γ r 0 2 B α,Γ 2 ν 6 B 6 X 6 Γ 2 B 2 α,Γ α 8 ≥ max α ζνB X , Γ r 0 2 B α,Γ 2 1600ν 6 B 6 X 6 Γ 2 B 2 α,Γ α 8 ≥ 1 min ζνB X α , B α,Γ Γ r0 2 2 1600ν 6 B 6 X 6 Γ 2 B 2 α,Γ α 8 = 1600Γ 2 ν 6 B 6 X 6 B 2 α,Γ δ 2 α 8 =400 Γ 2 ν 6 B 6 X 6 (B α,Γ + B α,Γ ) 2 δ 2 α 8 ≥400 Γ 2 ν 6 B 6 X 6 (B α,Γ + δΓ r 0 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 (r 0 ) 2 λ 1 ≥ α Π I (r 0 ) 2 νB X ≥ cα r 0 2 νB X ⇒ B α,Γ c r 0 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 r 0 2 X = α 5 √ δ cα Γβ = α 5 √ δ ⋅ min √ δ, cα Γβ = α 5 min ⎛ ⎝ δ, δα Γβ ⎞ ⎠

Combining the latter two cases as long as

ε 2 0 ≤ α 2 25 min c B α,Γ α νBΓ 2 r 0 2 X , ζ 2 ν 2 B 2 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 δ

Γ α r 0 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 ≤ ζ.

6.4.3. MAIN GENERALIZATION RESULT (COMPLETING THE PROOF OF THEOREM 2.3):

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 † I r 0 2 rather than Π I (r 0 ) 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

{(x i , y i )} n i=1 ∈ R d × R K 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 W 0 be an arbitrary initial weight matrix. Also assume the output matrix has bounded entries obeying V ∞ ≤ ν √ kK . Furthermore, set J ∶= J (W 0 ) and define the information I and nuisance N subspaces and the truncated Jacobian J I associated with the reference/initial Jacobian J based on a cut-off spectrum value α = νB ᾱ 4 √ n X . Also define the initial residual

r 0 = f (W 0 ) -y ∈ R nK and pick C r > 0 so that r0 2 √ n ≤ C r . Suppose number of hidden nodes k obeys k ≳ C 2 r Γ 4 ᾱ8 ν 2 ζ 2 , (6.58)

with Γ ≥ 1 and tolerance level ζ. Run gradient descent updates (1.5) with learning rate η ≤

1 ν 2 B 2 X 2 . Then, after T = Γ ηα 2 iterations, training loss obeys f (W T ) -y 2 ≤ Π N (r 0 ) 2 + C r e -Γ + ζ √ n.

and with probability at least 1δ, the generalization error obeys

Err D (W T ) ≤ 2 Π N (r 0 ) 2 √ n bias term + 12νB √ n J † I r 0 2 + Γ α Π N (r 0 ) 2 variance term +5 log(2 δ) n + 2C r (e -Γ + ζ).

Theorem 6.21 immediately follows from Theorem 6.22 below by upper bounding D α,Γ (see Definition 6.1) using Lemma 6.2 equation (6.2).

Theorem 6.22 Consider a training data set {(x i , y i )} n i=1 ∈ R d × R K 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 W 0 be an arbitrary initial weight matrix. Also assume the output matrix has bounded entries obeying V ∞ ≤ ν √ kK . Furthermore, set J ∶= J (W 0 ) and define the information I and nuisance N subspaces and the truncated Jacobian J I associated with the reference/initial Jacobian J based on a cut-off spectrum value α = νB ᾱ 4 √ n X . Also define the initial residual r 0 = f (W 0 )y ∈ R nK and pick C r > 0 so that r0 2

√ n ≤ C r . Also assume, the number of hidden nodes k obeys

k ≥ 25600 C 2 r Γ 4 ᾱ8 ν 2 B 2 ζ 2 , (6.59)

with Γ ≥ 1 and tolerance level ζ ≤ 2. Run gradient descent updates (1.5) with learning rate η ≤ 1 ν 2 B 2 X 2 . Then, after T = Γ ηα 2 iterations, with probability at least 1δ, the generalization error obeys

EQUATION

where D α,Γ is the early stopping distance as in Def. (6.1).

Proof First, note that using α ≤ β = νB X ,

B α,Γ

Γ ≤ r 0 2 per (6.2), and

r 0 2 ≤ C r √ n we have α νB X ≤ 1 ≤ Γ r 0 2 B α,Γ ≤ C r Γ √ n B α,Γ .

This together with ζ ≤ 2 implies that

α ζ 2 νB X + Γ r 0 2 B α,Γ ≤ α ζ 2 νB X + Γ r 0 2 ζ 2 B α,Γ ≤2 Γ r 0 2 ζ 2 B α,Γ ≤2 C r ζ 2 Γ B α,Γ √ n. (6.61) Thus when k ≥25600 C 2 r Γ 4 ᾱ8 ν 2 B 2 ζ 2 =6400 ⎛ ⎝ C r Γ ζ 2 ⎞ ⎠ 2 n 2 Γ 2 ν 6 B 6 X 4 α 8 √ n≥ X ≥ 6400 ⎛ ⎝ C r Γ ζ 2 ⎞ ⎠ 2 nν 6 B 6 X 6 α 8 =1600 ⎛ ⎝ 2 C r ζ 2 Γ B α,Γ √ n ⎞ ⎠ 2 B 2 α,Γ ν 6 B 6 X 6 α 8 (6.61) ≥ 1600 ⎛ ⎝ α ζ 2 νB X + Γ r 0 2 B α,Γ ⎞ ⎠ 2 B 2 α,Γ ν 6 B 6 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 τ -W 0 F ≤ 2B α,Γ α = 2D α,Γ . (6.62) and W τ -W 0 2,∞ ≤ 2νBΓ X √ kα 2 r 0 2 r0 2 ≤Cr √ n ≤ 2 √ nC r νBΓ X √ kα 2 . (6.63) and f (W T ) -y 2 ≤e -Γ Π I (r 0 ) 2 + Π N (r 0 ) 2 + ζ 2 r 0 2 ≤ Π N (r 0 ) 2 + e -Γ + ζ 2 r 0 2 r0 2 ≤Cr √ n ≤ Π N (r 0 ) 2 + C r e -Γ + ζ 2 √ n. (6.64)

Furthermore, using the assumption that V ∞ ≤ ν √ kK Lemma 6.18 applies and hence equation (6.43) with

EQUATION

Also note that using (6.59) we have

3νBR 2 √ k ≤ 12C 2 r Γ 2 ν 3 B 3 n X 2 √ kα 4 ≤ C r ζ 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.

6.5.1. OPTIMIZATION RESULT:

Theorem 6.23 (Optimization guarantee for random initialization) Consider a training data set 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 J I associated with the reference Jacobian J based on a cut-off spectrum value of

{(x i , y i )} n i=1 ∈ R d × R K generated i.

α 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 × 10 7 Γ 4 K 4 B 8 X 6 n log(n) c 4 ζ 4 α 8 0 (6.68)

with Γ ≥ 1 and ζ ≤ c 2 . We run gradient descent iterations of the form (1.5) with a learning rate η ≤ 1 ν 2 B 2 X 2 . Then, after T = ΓK ην 2 α 2 0 iterations, the following identities

f (W τ0 ) -y 2 ≤ Π N (y) 2 + e -Γ Π I (y) 2 + 4ζ √ n, (6.69) W τ -W 0 2 ≤ 2 √ K(B α0,Γ (y) + Γζ √ n) να 0 , (6.70) W τ -W 0 2,∞ ≤ 4ΓBK X να 2 0 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 (W 0 ) obeys

f (W 0 ) 2 ≤ ζ y 2 = ζ √ n ≤ √ n 2 . (6.72)

Hence the initial residual obeys r 0 2 ≤ 2 √ n. Furthermore, using ζ ≤ c 2

r 0 + y 2 ≤ζ y 2 ⇒ Π I (r 0 + y) 2 ≤ ζ y 2 . (6.73)

Thus,

Π I (r 0 ) 2 ≥ Π I (y) 2 -Π I (r 0 + y) 2 ≥ Π I (y) 2 -ζ y 2 ≥ (c -ζ) y 2 ≥ c 2 y 2 (6.74) ≥ c 4 r 0 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,Γ (r 0 ) ≤ B α0,Γ (y) + Γ r 0 + 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 (W 0 ) satisfying J J T = E[J (W 0 )J (W 0 ) 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 (W 0 )J (W 0 ) T ] 1 2 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 ≥ 1000K 2 B 4 X 4 log(n) δ 2 (6.77) then, J (W 0 )J (W 0 ) T -E[J (W 0 )J (W 0 ) T ] ≤ δν 2 K . (6.78)

Let J (W ) be obtained by adding max(Knp, 0) zero columns to J (W ). Then, using (6.78) and Lemma 6.4, there exists J satisfying J J T = E[J (W 0 )J (W 0 ) T ] and

J (W 0 ) -J ≤ 2 δν 2 K .

Therefore, J is an ε 2 0 = 4 δν 2 K reference Jacobian. Now set

Θ = min ⎛ ⎝ c B α0,Γ α 0 BΓ 2 X √ nK , ζB √ K X α 0 2 , ζ Γ ⎞ ⎠

and note that using

α = ν √ K α 0 and r 0 2 ≤ 2 √ n Θ = min ⎛ ⎝ c B α0,Γ α 0 BΓ 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 r 0 2 X , ζ 2 ν 2 B 2 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 (W 0 )J (W 0 ) T ]. Thus, using (6.79) to ensure (6.51) holds it suffices to show

ε 2 0 = 4 δν 2 K ≤ α 2 25 Θ 2 = ν 2 α 2 0 50K Θ.

Hence, to ensure (6.51) holds we need to ensure that δ obeys δ ≤ α 2 0 200 Θ.

Thus using δ = α 2 0 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 × 10 7 Γ 4 K 4 B 8 X 6 n log(n) c 4 ζ 4 α 8 0 (6.80) ≥12 × 10 7 Γ 4 K 4 B 8 X 8 log(n) c 4 ζ 4 α 8 0 ≥4 × 10 7 ⋅ 4K 2 B 4 Γ 4 X 4 c 4 α 4 0 + 1 ζ 4 + Γ 2 ζ 2 K 2 B 4 X 4 log(n) α 4 0 ≥4 × 10 7 ⋅ 4KB 4 Γ 4 X 4 c 4 α 4 0 + 1 ζ 4 + Γ 2 ζ 2 K 2 B 4 X 4 log(n) α 4 0 ≥4 × 10 7 ⋅ 4KB 4 Γ 4 X 4 c 4 α 4 0 + α 4 0 ζ 4 B 4 K 2 X 4 + Γ 2 ζ 2 K 2 B 4 X 4 log(n) α 4 0 (a) ≥ 4 × 10 7 ⋅ ⎛ ⎝ nKB 2 Γ 4 X 2 c 2 B 2 α0,Γ α 2 0 + α 4 0 ζ 4 B 4 K 2 X 4 + Γ 2 ζ 2 ⎞ ⎠ K 2 B 4 X 4 log(n) α 4 0 ≥4 × 10 7 ⋅ max ⎛ ⎝ nKB 2 Γ 4 X 2 c 2 B 2 α0,Γ α 2 0 , α 4 0 ζ 4 B 4 K 2 X 4 , Γ 2 ζ 2 ⎞ ⎠ K 2 B 4 X 4 log(n) α 4 0 =4 × 10 7 K 2 B 4 X 4 log(n) α 4 0 ⋅ min cB α 0 ,Γ α0 BΓ 2 X √ nK 2 , ζB √ K X α0 4 , ζ 2 Γ 2 = 1000K 2 B 4 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 (r 0 ) 2 ≥ c 2 y 2 = c 2 √ n which combined imply B α0,Γ ≥ α λ 1 Π I (r 0 ) 2 ≥ α B X Π I (r 0 ) 2 ≥ α 0 c 2 y 2 B X = α 0 c 2 √ 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 × 10 6 K 4 B 8 X 6 Γ 4 n log(n) ζ 4 α 8 0 (6.83) (a) ≥ 25600 K 4 B 6 X 6 Γ 4 n ζ 2 ν 2 α 8 0 (6.84) k ≥12800 α 2 0 ζ 2 B 2 K X 2 + 1 K 4 B 6 X 6 Γ 4 n ν 2 α 8 0 (b) ≥ 3200 ⎛ ⎝ α 2 0 ζ 2 B 2 K X 2 + 4Γ 2 n B 2 α,Γ ⎞ ⎠ K 4 B 6 X 6 Γ 2 B 2 α,Γ ν 2 α 8 0 ≥3200 ⎛ ⎝ α 2 0 ζ 2 B 2 K X 2 + Γ 2 r 0 2 2 B 2 α,Γ ⎞ ⎠ K 4 B 6 X 6 Γ 2 B 2 α,Γ ν 2 α 8 0 ≥1600 α 0 ζ √ KB X + Γ r 0 2 B α,Γ 2 K 4 B 6 X 6 Γ 2 B 2 α,Γ ν 2 α 8 0 =1600 α ζνB X + Γ r 0 2 B α,Γ 2 ν 6 B 6 X 6 Γ 2 B 2 α,Γ α 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 α,Γ ≤ Γ r 0 2 ≤ 2Γ √ n

Note that (6.80) and (6.84) are implied by

k ≥ 12 × 10 7 Γ 4 K 4 B 8 X 6 n log(n) c 4 ζ 4 α 8 0 , (6.86)

which is the same as (6.68). What remains stating the optimization bounds in terms of the labels y. This follows by substituting (6.72), (6.76), and the fact that r 0 2 ≤ 2 √ n into (6.55), (6.53), and (6.54), respectively.

6.5.2 GENERALIZATION RESULT (COMPLETING THE PROOF OF THEOREM 2.2)

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 † I y 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 

{(x i , y i )} n i=1 ∈ R d × R K generated i.

+ 12B √ K √ n J † I y 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). 

D (W T ) ≤ 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 M W and R dictated by Theorem 6.23 where the output layer V is fixed. Observe that 

EQUATION

3νB R 2 √ k = 48nΓ 2 B 3 K 2 X 2 να 4 0 √ k ≤ c 2 ζ 2 230νB log(n) ≤ζ.

(6.93)

Plugging in (6.91), (6.92), and (6.93) into (6.90) concludes the proof.

A THE JACOBIAN OF THE MIXTURE MODEL IS LOW-RANK (PROOFS FOR SECTION 2.

3)

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

Theorem A.1 Consider a data set of size n consisting of input/label pairs {(x i , y i )} n i=1 ∈ R d × R K generated according to the Gaussian mixture model of Definition 2.4 with K 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 = [x 1 . . . x n ] 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 distribution N (0, I d ). Define the neural tangent kernel matrix associated with the cluster centers as

Σ(M ) = (M M T ) ⊙ E g∼N (0,I d ) [φ ′ (M g)φ ′ (M g) T ] ∈ R KC×KC ,

and assume that Σ(M ) is full rank. Then, the following properties hold with probability 1 -KC exp(-n 8KC )

• I is a K 2 C dimensional subspace.

• The concatenated label vector y = y T 1 y T 2 . . . y T n T lies on I.

• The nonzero eigenvalues (top K 2 C eigenvalues) of Σ(X) are between n 2KC s min (Σ(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.5ñ ≤ ñ , ˜ ≤ 2ñ.

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

Σ(X) = kν 2 I K ⊗ Σ(X).

where Σ(X) = (XX T ) ⊙ E g i.i.d.

∼ N (0,1) [φ ′ (Xg)φ ′ (Xg) T ]. Due to this Kronecker product representation, the range space of Σ(X) is separable. In particular, note that with Ĩ = Range Σ(X)

we have I = I K ⊗ Ĩ which also implies rank(I) = K ⋅ rank Ĩ . 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:

• Ĩ has rank KC.

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 σ √

log nB 2 KC λ 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 σ √ log nB 2 KC λ min ). (A.3)

In return, this implies that

Π N (y) 2 ≲ nσ √ log nB 2 KC λ min

To proceed, we pick α 0 = λ min n 2KC and corresponding ᾱ = α0

4 √ n K X B ≥ λ min 2B 2 K 2 C

and apply (??) to find that, classification error is upper bounded by

Err D (W T ) ≲ σ √ log nB 2 KC λ min + ΓBK √ C √ nλ min + 12ζ + 5 log(2 δ) n + 2e -Γ .

B JOINT INPUT-OUTPUT OPTIMIZATION:

In this section we wish to provide the ingredients necessary to prove a result for the case where both set of input and output weights W 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 φ(W x). (B.1)

Denoting the Jacobian associated with (B.1) by J (V , W ) we have that

J (V , W ) = [J (V ) J (W )] ∈ R Kn×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 (v 1 ) J (v 2 ) . . . J (v K )] . (B.2)

where J (v ) ∈ R Kn×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: 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.

mat J T (V )u ∞ ≤ B W 2,∞ X . • Lipschitzness: Given inputs V , V ′ and outputs W , W ′ J (V , W )-J (V ′ , W ′ ) ≤ B X ( √ Kk V -V ′ ∞ + √ K V ∞ W ′ -W F + W -W ′ F ).

Π I (y) 2 y 2 Π N (y) 2 y 2 J † I y 2 y 2 Π I (r0) 2 r0 2 Π N (r0) 2 r0 2 J † I r0 2 r0 2 J train

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. 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.

Π I (y) 2 y 2 Π N (y) 2 y 2 J † I y 2 y 2 Π I (r0) 2 r0 2 Π N (r0) 2 r0 2 J † I r0 2 r0 2 J train init 0.

I (y) 2 y 2 Π N (y) 2 y 2 J † I y 2 y 2 Π I (r0) 2 r0 2 Π N (r0) 2 r0 2 J † I r0 2 r0 2 J train

Π I (y) 2 y 2 Π N (y) 2 y 2 J † I y 2 y 2 Π I (r0) 2 r0 2 Π N (r0) 2 r0 2 J † I r0 2 r0 2 J train init 0.

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 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. 13 : Depiction of the alignment of the initial label/residual with the information/nuisance space using 100% label corruption on 3-class CIFAR-10 data.

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. 

:

For 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.

Based on our understandingArora et al. (2019) requires the number of hidden units to be on the order of k ≳ n 8 λ 6 . Hence our result reduces the dependence on width by a factor of at least n 4 λ 2 .

e.g. obtained by training with data in a related problem as is common in transfer learning.

Note that this theorem and its conclusions hold for any choice of these parameters in the specified range.

Note that this theorem and its conclusions hold for any choice of these parameters in the specified range.

Figures and Tables:

Figure 1: Illustration of a one-hidden layer neural net with d inputs, k hidden units and K outputs along with a one-hot encoded label.

Figure 2: Depiction of the training and generalization dynamics of gradient methods based on the information and nuisance spaces associated with the neural net Jacobian.

unit length input samples and one-hot encoded labels. Consider the neural net in (1.1) parameterized by W and initialized with W 0 i.i.d.

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 W T -W 0 F .

); Arora et al. (2019); Du et al. (2018b); Allen-Zhu et al. (2018b); Cao & Gu

Figure 3: The singular values of the normalized

Figure 4: Histogram of the top 1000 Jacobian singular values on the CIFAR10 dataset.

Figure 5: Evolution of the residual (r τ = f (W τ )y) along the information/nuisance spaces of the final Jacobian on (a) the training data and (b) the test data and c) misclassification error on training and test. This experiment uses uncorrupted labels.

Figure 6: Evolution of the residual (r τ = f (W τ )y) along the information/nuisance spaces of the Jacobian at 50 epochs on (a) training data and (b) test data and (c) misclassification error on training and test data. 50% of the labels have been corrupted.

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.

Figure 8: Plots of the (a) total test error and (b) the test error components for the model in Section 5.1. The test error decreases rapidly over the information subspace but slowly increases over the nuisance subspace.

mathematical terms we need to bound D ∶= C -E[C] .To this aim define D[ , ∶] to denote the th block row of D. Standard bounds on spectral norm in terms of sub-matrices allow us to conclude that D

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. Let W 0 be the initial weight matrix with i.i.d. N (0, 1) entries. Fix a precision level ζ and set

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 W 0 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 J I 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 log n ζ 4 ᾱ8 (6.87) with Γ ≥ 1. We run gradient descent iterations of the form (1.5) with a learning rate η ≤ 1 ν 2 B 2 X 2 . Then, after T = ΓK ην 2 α 2 0 iterations, training loss obeys (6.69) and classification error Err D (W T ) is upper bounded by 2 Π N (y) 2 √ n bias term

Theorem 6.25 (Neural Net -Generalization) Consider a training data set {(x i , y i )} n i=1 ∈ R d ×R K 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 W 0 be the initial weight matrix with i.i.d. N (0, 1) entries. Fix a precision level ζ ≤ c 2 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 J I 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 × 10 7 Γ 4 K 4 B 8 X 4 n 2 log(n) c 4 ζ 4 α 8 0 (6.89) with Γ ≥ 1. We run gradient descent iterations of the form (1.5) with a learning rate η ≤ 1 ν 2 B 2 X 2 . Then, after T = ΓK ην 2 α 2 0 iterations, classification error Err D (W T ) is upper bounded by Err

Figure 9: Scaled spectral densities of the full train Jacobian at different stages of training

Figure 10: Experiments on the original, uncorrupted CIFAR-10 dataset. * : we discuss these plots in Section 4.



Figure 11: Experiments with 25% label corruption on CIFAR-10 data

Figure12: Experiments with 50% label corruption on CIFAR-10 data. * : we discuss these plots in Section 4.

Figure 13: Experiments with 75% label corruption on CIFAR-10 data

Figure 16: Experiments on 3-class uncorrupted CIFAR-10 dataset

Figure 17: Experiments on 3-class uncorrupted CIFAR-10 data, trained with Adam

Figure 18: Experiments with 25% label corruption on 3-class CIFAR-10 data

Figure 20: Experiments with 75% label corruption on 3-class CIFAR-10 data

Figure 21: Experiments with 100% label corruption on 3-class CIFAR-10 data

Figure 23: Experiments with 25% label corruption on MNIST data

Figure 24: Experiments with 50% label corruption on MNIST data

Figure 25: Experiments with 75% label corruption on MNIST data

Depiction of the alignment of the initial residual with the information/nuisance space using 50% label corrupted data and a Multi-class ResNet20 trained with SGD.

Table content: <table><tr><td/><td>32762</td><td>0.94481</td><td>0.32152</td><td>0.9469</td></tr><tr><td>J train f inal</td><td>0.8956</td><td>0.44487</td><td>0.89597</td><td>0.44412</td></tr></table>

Depiction of the alignment of the initial label/residual with the information/nuisance space using uncorrupted CIFAR-10 data. is spanned by the top 50 singular vectors. We marked figures and table entries also included in Section 4 by asterisk.

Table content: <table><tr><td>init J train f inal J test init J test f inal</td><td>*   *</td><td>0.38081 0.9869 0.38184 0.80926</td><td>0.92465 0.16131 0.92423 0.58746</td><td>0.027224 0.00070893 0.060229 0.0013734</td><td>0.37114 0.98669 0.37227 0.80912</td><td>0.92858 0.1626 0.92812 0.58764</td><td>0.027293 0.00070354 0.06037 0.0013716</td></tr></table>

Depiction of the alignment of the initial label/residual with the information/nuisance space using 25% label corruption on CIFAR-10 data.

Table content: <table><tr><td/><td>35328</td><td>0.93552</td><td>0.021592</td><td>0.36954</td><td>0.92921</td><td>0.021568</td></tr><tr><td>J train f inal</td><td>0.92214</td><td>0.38685</td><td>0.00087246</td><td>0.92324</td><td>0.38423</td><td>0.00087304</td></tr></table>

Depiction of the alignment of the initial label/residual with the information/nuisance space using 50% label corruption on CIFAR-10 data.

Table content: <table><tr><td>init J train f inal J test init J test dip J test f inal</td><td>*   *</td><td>0.32762 0.8956 0.38013 0.7041 0.44774</td><td>0.94481 0.44487 0.92493 0.7101 0.89416</td><td>0.017556 0.00096413 0.080777 0.0040147 0.0012216</td><td>0.32152 0.89597 0.37454 0.70229 0.44409</td><td>0.9469 0.44412 0.92721 0.71189 0.89598</td><td>0.017521 0.00096652 0.080766 0.0040423 0.0012157</td></tr></table>

Depiction of the alignment of the initial label/residual with the information/nuisance space using 75% label corruption on CIFAR-10 data.

Table content: <table><tr><td/><td>31756</td><td>0.94824</td><td>0.0056031</td><td>0.325</td><td>0.94571</td><td>0.005592</td></tr><tr><td>J train f inal</td><td>0.50238</td><td>0.86465</td><td>0.00047518</td><td>0.50718</td><td>0.86184</td><td>0.00047967</td></tr></table>

Depiction of the alignment of the initial label/residual with the information/nuisance space using 100% label corruption on CIFAR-10 data.

Table content: <table><tr><td/><td>Π I (y) 2</td><td>Π N (y) 2</td><td>J  † I y</td><td>2</td><td colspan="3">Π I (r0) 2</td><td>Π N (r0) 2</td><td>J  † I r0</td><td>2</td></tr><tr><td>J train init J train f inal</td><td>y 2 0.31581 0.47747</td><td>y 2 0.94882 0.87865</td><td colspan="2">y 2 0.0094479 0.00045241</td><td colspan="3">r0 2 0.31854 0.47921</td><td>r0 2 0.94791 0.8777</td><td>r0 2 0.0094092 0.00045344</td></tr><tr><td/><td/><td>0.8</td><td/><td/><td/><td/><td/></tr><tr><td/><td/><td>0.6</td><td/><td/><td/><td/><td/></tr><tr><td/><td/><td>0.4</td><td/><td/><td>Π N (y) y 2</td><td>2</td><td>final</td></tr><tr><td/><td/><td/><td/><td/><td>Test error</td><td/><td/></tr><tr><td/><td/><td>0.2</td><td/><td/><td/><td/><td/></tr><tr><td/><td/><td>0</td><td>0.25</td><td>0.5</td><td>0.75</td><td/><td>1</td></tr><tr><td/><td/><td/><td colspan="3">Label corruption</td><td/><td/></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td>1</td></tr><tr><td colspan="9">Figure 15: Test error vs. final projection of labels on nuisance subspace for CIFAR-10</td></tr><tr><td colspan="6">experiments. See discussion of this plot in Section 4</td><td/><td/></tr><tr><td colspan="7">C.3 EXPERIMENTS ON SUBSAMPLED 3-CLASS CIFAR-10</td><td/></tr></table>

Depiction of the alignment of the initial label/residual with the information/nuisance space using 3-class uncorrupted CIFAR-10 data.)

Table content: <table><tr><td/><td>Π I (y) 2</td><td>Π N (y) 2</td><td>J  † I y</td><td>2</td><td>Π I (r0) 2</td><td>Π N (r0) 2</td><td>J  † I r0</td><td>2</td></tr><tr><td>J train init J train 200epochs J train f inal J test init J test f inal</td><td>y 2 0.7239 0.97266 0.98743 0.73366 0.8974</td><td>y 2 0.68991 0.23224 0.15804 0.67951 0.44123</td><td colspan="2">y 2 0.0054426 0.0026234 0.0031639 0.010328 0.0027082</td><td>r0 2 0.88552 0.96849 0.97606 0.88827 0.89772</td><td>r0 2 0.4646 0.24905 0.21749 0.45932 0.44057</td><td colspan="2">r0 2 0.0040958 0.0030069 0.0034312 0.0077364 0.0029383</td></tr></table>

Depiction of the alignment of the initial label/residual with the information/nuisance space using 50% label corruption on 3-class CIFAR-10 data.

Table content: <table><tr><td/><td>Π I (y) 2</td><td>Π N (y) 2</td><td>J  † I y</td><td>2</td><td>Π I (r0) 2</td><td>Π N (r0) 2</td><td>J  † I r0</td><td>2</td></tr><tr><td>J train init J train dip J train f inal J test init J test dip J test f inal</td><td>y 2 0.58664 0.61125 0.75143 0.7236 0.81473 0.75362</td><td>y 2 0.80985 0.79144 0.65981 0.69021 0.57984 0.65731</td><td colspan="2">y 2 0.0017197 0.0026809 0.0018702 0.012765 0.009931 0.0033476</td><td>r0 2 0.64281 0.65671 0.76311 0.7594 0.8322 0.77225</td><td>r0 2 0.76602 0.75415 0.64627 0.65062 0.55447 0.63532</td><td colspan="2">r0 2 0.0019814 0.0031267 0.0012039 0.01235 0.010515 0.0021479</td></tr></table>

Depiction of the alignment of the initial label/residual with the information/nuisance space using 75% label corruption on 3-class CIFAR-10 data.

Table content: <table><tr><td/><td>Π I (y) 2</td><td>Π N (y) 2</td><td>J  † I y</td><td>2</td><td>Π I (r0) 2</td><td>Π N (r0) 2</td><td>J  † I r0</td><td>2</td></tr><tr><td>J train init J train f inal</td><td>y 2 0.58032 0.64704</td><td>y 2 0.81439 0.76246</td><td colspan="2">y 2 0.0014306 0.0015808</td><td>r0 2 0.68898 0.78391</td><td>r0 2 0.72478 0.62088</td><td colspan="2">r0 2 0.0015956 0.00091715</td></tr></table>

Depiction of the alignment of the initial label/residual with the information/nuisance space using uncorrupted MNIST data.

Table content: <table><tr><td/><td>Π I (y) 2</td><td>Π N (y) 2</td><td>J  † I y</td><td>2</td><td>Π I (r0) 2</td><td>Π N (r0) 2</td><td>J  † I r0</td><td>2</td></tr><tr><td>J train init J train f inal J test init J test f inal</td><td>y 2 0.43598 0.9946 0.43964 0.99188</td><td>y 2 0.89996 0.10375 0.89817 0.12716</td><td colspan="2">y 2 0.028665 0.00066394 0.071577 0.0016387</td><td>r0 2 0.48168 0.99321 0.48474 0.99044</td><td>r0 2 0.87635 0.11633 0.87466 0.13796</td><td colspan="2">r0 2 0.027921 0.00076602 0.069724 0.0018745</td></tr></table>

Depiction of the alignment of the initial label/residual with the information/nuisance space using 25% label corruption on MNIST data.

Table content: <table><tr><td/><td>Π I (y) 2</td><td>Π N (y) 2</td><td>J  † I y</td><td>2</td><td>Π I (r0) 2</td><td>Π N (r0) 2</td><td>J  † I r0</td><td>2</td></tr><tr><td>J train init J train f inal</td><td>y 2 0.38906 0.94093</td><td>y 2 0.92121 0.33861</td><td colspan="2">y 2 0.02197 0.0010507</td><td>r0 2 0.38586 0.9391</td><td>r0 2 0.92256 0.34365</td><td colspan="2">r0 2 0.022028 0.0010331</td></tr><tr><td/><td>Π I (y) 2</td><td>Π N (y) 2</td><td>J  † I y</td><td>2</td><td>Π I (r0) 2</td><td>Π N (r0) 2</td><td>J  † I r0</td><td>2</td></tr><tr><td>J train init J train f inal J test init J test dip J test f inal</td><td>y 2 0.34434 0.83136 0.44458 0.97007 0.71536</td><td>y 2 0.93885 0.55573 0.89574 0.24284 0.69876</td><td colspan="2">y 2 0.013931 0.00081502 0.076224 0.0043166 0.0018554</td><td>r0 2 0.35864 0.83235 0.45366 0.9705 0.71632</td><td>r0 2 0.93347 0.55425 0.89117 0.24109 0.69777</td><td colspan="2">r0 2 0.013967 0.00081616 0.076101 0.0043041 0.0018564</td></tr></table>

Depiction of the alignment of the initial label/residual with the information/nuisance space using 50% label corruption on MNIST data.

Table content: <table/>

Depiction of the alignment of the initial label/residual with the information/nuisance space using 75% label corruption on MNIST data.

Table content: <table><tr><td/><td>Π I (y) 2</td><td>Π N (y) 2</td><td>J  † I y</td><td>2</td><td>Π I (r0) 2</td><td>Π N (r0) 2</td><td>J  † I r0</td><td>2</td></tr><tr><td>J train init J train f inal</td><td>y 2 0.31995 0.56007</td><td>y 2 0.94744 0.82844</td><td colspan="2">y 2 0.0061168 0.00063427</td><td>r0 2 0.32072 0.55869</td><td>r0 2 0.94718 0.82938</td><td colspan="2">r0 2 0.0061618 0.00062509</td></tr></table>

Peer review (iclr2020):
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. I enjoyed the paper, however, there were many points where I was unclear on the precise nature of the assumptions used / the strength of the results. Disclaimer: I didn’t manage to read through the proofs in the appendix and cannot therefore vouch for its correctness. -- Point 1 -- Leveraging the data structure I am unclear on how exactly you were modelling the structure of the data. From your proofs, it seems that you have been dealing with the matrix X comprising the concatenated flatted vectors of the raw input features (e.g. pixels) of the input data [x1,x2,...,x_datasetsize]^T. In particular, the only place where I see data explicitly enter is in Definition 2.1, where you look at the X X^T and fi’(w X) fi’(w X)^T. If the data is linearly separable in the raw input space on its own, then I see that the matrix X X^T will be low-rank (related to the number of classes). I also see your point about the connection of the y to the relevant (semantic) clusters. The same argument could by applied to fi’(w X) fi’(w X)^T -- provided that the features produced are again linearly separable, we will observe this object to have a low (number of classes - 1) rank. What is unclear to me is whether these assumptions are warranted. I understand that some simple datasets, e.g. MNIST, are essentially linearly separable in the raw pixels, and therefore the X X^T indeed is low rank. However, I doubt anything like that is true for big datasets, such as ImageNet. For deep networks that are used on these big datasets, such a modelling assumptions would likely not be true. I wonder how this relates to your results, since the low-rank nature of the Hessian is observed even for those networks, which is in turn related to the low rank nature of the Jacobian. -- Point 2 -- Data implicitly present in the Jacobian tensor. I wonder how you modelled the Jacobian tensor that you started using on page 3. Since the Jacobian -- the derivative of the output logits with respect to the weights, has to be evaluated at a particular input X, the assumptions you make on the data are in turn having an effect on the Jacobian, and vice versa. I am unclear on exactly what assumptions you make about the object, and whether you are actually saying that its low rank structure comes from the data, is empirical observed and therefore assumed, or due to the network regardless of the data. I recently saw a new arXiv submission that seems to be looking into this on real networks: https://arxiv.org/pdf/1910.05929.pdf Their model explicitly assumes that logit gradients cluster in the weight space in a particular way. -- Point 3 -- Square loss vs softmax You are using the square loss |f(X) - y|^2 throughout your work. Many of the empirical low-rank observations of the Hessian (related to the JJ^T) are performed on real networks with the cross-entropy loss. While the Hessian with the square loss is of the form JJ^T + terms, the softmax in the cross-entropy loss introduces an additional cross-term (let us call it P for now), which in turn makes it JP(PJ)^T + terms. Do you know how this relates to your results? More generally, does the square loss you use make the results significantly different from what we would get for a softmax? -- Point 4 -- Neural Tangent Kernel (NTK) -- assumptions Under the NTK assumption, you still need to model the derivatives of each logit with respect to each weight on each input in order to obtain the Jacobian matrix and in turn JJ^T. I am therefore very confused by “Based on our simulations the M-NTK indeed haslow-rank structure with a few large eigenvalues and many smaller ones. “ on page 5. What assumptions exactly do you use in your model? -- Point 5 -- Neural Tangent Kernel (NTK) -- validity By assuming the NTK holding, do you limit the validity of your results? I think it is believed that NTK might not, generally speaking, be enough to capture the complexity of DNNs, and therefore assuming it might limit the range of applicability of any results derived assuming it. -- Conclusion -- In general, my points of confusion often stem from being unsure as to what parts of the argument were assumed and based on what empirical / theoretical evidence, and what parts were generically true. While the paper seems interesting, I am not sure what its novel contribution is and how broad the claims made actually are in their applicability. Appendix: I was not able to judge the proofs in the appendix.