G ENERALIZATION G UARANTEES FOR N EURAL N ETS VIA H ARNESSING THE L OW - RANKNESS OF J ACOBIAN
Anonymous authors
Paper under double-blind review
A BSTRACT
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.
Speciﬁcally, 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 sufﬁciently nice datasets.
We conduct various numerical experiments on deep networks that corroborate our theoretical ﬁndings 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.
I NTRODUCTION 1. M OTIVATION 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)ﬁt 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 overﬁt 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 overﬁtting.
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 overﬁtting due to higher variance.
●
Network size vs prediction bias: We obtain data-dependent tradeoffs between the network size and prediction bias.
Speciﬁcally, we show that larger networks result in smaller prediction bias, but small networks can still generalize well, especially when the dataset is sufﬁciently structured, but typically incur a larger bias.
This compares favorably with recent literature on optimization and generalization of neural networks Jacot et al. Arora et al.
Du et al. (2018b); Allen-Zhu et al.
(2018b); Cao & Gu (2019) ; Ma et al.
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.
[Figure]
1. M ODEL 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
[Equation]
and
[Equation]
× k to denote the input-to-hidden and hiddento-output weights.
The overall input-output relationship of the network is a function f [Equation] ⋅
[Equation] → R K that maps an input x ∈
R d to an output via
[Equation]
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 (cid:96) ∈
R K has all zero entries except for the (cid:96) th entry which is equal to one).
To learn this dataset, we ﬁx the output layer and train over W via
[Equation]
It will be convenient to concatenate the labels and prediction vectors as follows y
[Equation] ⎡⎢⎢⎢⎢⎣ [Equation]
nK and f (
W
[Equation]
⎡⎢⎢⎢⎢⎣
[Equation]
nK
[Equation] Using this shorthand we can rewrite the loss (1.2) as
[Equation]
To optimize this loss starting from an initialization W 0 we run gradient descent iterations of the form
W τ + =
W τ − η ∇L(
W τ
[Equation]
(1.5) with a step size η .
In this paper we wish to explore the theoretical properties of the model found by such iterative updates with an emphasis on the generalization ability.
1. I NFORMATION AND N UISANCE S PACES
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 [Equation]
kd [Equation]
deﬁne the misﬁt/residual vector r [Equation] − y and note that the gradient updates take the form
[Equation] where
[Equation] [Figure]
Here,
[Equation]
kd denotes the Jacobian mapping associated with f deﬁned as
[Equation]
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
[Equation] the gradient update signiﬁcantly reduces the misﬁt 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 misﬁt and label vectors.
To capture these interactions we require a few deﬁnitions.
Deﬁnition 1.1 (Information & Nuisance Spaces) Consider a matrix J ∈
R nK [Equation] p with singular value decomposition given by
[Equation]
with λ≥ λ≥ ...
≥ λ nK denoting the singular values of J in decreasing order and { u s } [Equation]
∈ R nK and { v s } nKs = ∈
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≤ α ≤ λ 1 let r ∶= r ( α ) denote the index of the smallest singular value above α .
We deﬁne the information and nuisance spaces associated with J as
I ∶= span ({ u s } rs = ) and
N ∶= span ({ [Equation]
Kns = [Equation] + ) .
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 deﬁne the matrix J and the corresponding information/nuisance spaces.
More, speciﬁcally we will set J to either J
[Equation]
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 signiﬁcant 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 signiﬁcantly slows down (see
Fig. 2).
Furthermore, subsequent iterations in this stage lead to a slight increase in test error.
M AIN 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 ﬁrst and second order derivatives.
This would for instance apply to the softplus activation φ ( z ) = log ( + 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 modiﬁcations.
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 =
∑ ni = 1 λ i u
i u Ti , the square root matrix is deﬁned as A/ ∶= ∑ ni =
√
λ
i u i u
Ti .
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 quantiﬁable measure of performance for a trained model.
Given a sample ( x , y ) ∈
R d ×
R K from a distribution
D , the classiﬁcation error of the network W with respect to
D is deﬁned as
Err
[Equation]
W
[Equation]
2. R ESULTS FOR RANDOM INITIALIZATION
To explore the generalization of randomly initialized networks, we utilize the neural tangent kernel.
Deﬁnition 2.1 (Multiclass Neural Tangent Kernel (M-NTK)
Jacot et al.
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 deﬁne the multiclass kernel matrix as
[Equation]
where I K is the identity matrix of size K .
Here, the (cid:96) th diagonal block of Σ (
X ) corresponds to the kernel matrix associated with the (cid:96) th network output for≤ (cid:96) ≤
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 ,) and the output layer V has i.i.d.
zero-mean entries with ν/
K variance.
Then E [J (
W)J (
W)
T ] =
ν 2 Σ (
X ) .
We use the square root of this multiclass kernel matrix (i.e. Σ (
X ) / 2 ) to deﬁne the information and nuisance spaces for our random initialization result.
The following theorem is a (non-rigorous) simpliﬁcation of our main result Theorem 6.24 where we ignore constants and log factors, and state a weaker but simpler generalization bound.
Theorem 2.2 Fix numbers Γ ≥ 1 and α > 0 .
Consider an i.i.d.
training dataset {( x i , y i )} ni
= ∈ R d ×
R K with 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.
∼ N( 0 ,) entries.
Set V with i.i.d.
Rademacher entries (properly scaled).
Deﬁne the information
I and nuisance
N spaces with respect to Σ (
X ) / with spectrum cutoff α √ nK per Deﬁnition 1.1.
Furthermore, assume
[Equation]
[Equation]
Then after T ∝ Γ / α 2 gradient iterations of (1.5) , with high probability, training loss obeys
[Equation]
[Equation]
Furthermore, the classiﬁcation error obeys
[Equation]
This theorem shows that even networks of moderate width can achieve a small generalization error if (1) the data has low-dimensional representation i.e. the kernel is approximately low-rank and (2) the inputs and labels are semantically-linked i.e. the label vector y mostly lies on the information space.
●
Generalization bound: The generalization error has two core components: bias and variance.
The bias component
[Equation]
Γ arises from the training loss and corresponds to the portion of the labels that falls over the nuisance space.
The variance component Γ
[Equation]
n corresponds to the
Rademacher complexity of the model space which connects to the distance
[Equation]
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 α
[Equation]
This sets
I
[Equation]
Kn and
[Equation]
For this case, the more general Theorem 6.24 yields
[Equation]
while requiring a width of
[Equation]
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 inﬁnitely) 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 sufﬁcient 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. Arora et al.
Du et al. (2018b); Allen-Zhu et al.
(2018b)
; Cao & Gu (2019) where network size is required to grow polynomial with n and inversely with the distance between the inputs or other notions of margin.
● Network size–Bias tradeoff: Based on the requirement (2.2) if the network is large (in terms of # of hidden units k ), we can choose a small cut-off α .
This in turn allows us to enlargen the information space and reduce the training bias further.
In summary, as the network capacity grows, we can gradually interpolate ﬁner 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.
Speciﬁcally, picking α to be a constant, convergence on the information space will be fast requiring only a constant number of iterations to reach any ﬁxed accuracy (see (2.3)).
2. G ENERALIZATION 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 ﬂavor to Theorem 2.2 with the key difference that the information and nuisance spaces are deﬁned with respect to any arbitrary initial
Jacobian.
This shows that if a pre-trained model 3 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) simpliﬁcation of Theorem 6.21.
Theorem 2.3 Let Γ ≥ 1 ,α be arbitrary scalars.
Consider i.i.d.
training data { [Equation]
x [Equation]
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 ∥
[Equation]
Deﬁne the nuisance space
N associated with
[Equation]
based on spectrum cutoff α √ n .
Set the initial residual r= f (
W
[Equation]
− y ∈
R nK and assume [Equation] n .
Suppose k ≳ Γ/
α 8 .
After T ∝ Γ / α 2 iterations (1.5) with constant learning rate, training loss obeys:
[Equation]
Also with high probability, classiﬁcation error obeys: Err
D (
W T )
≲ ∥ Π
N ( r)∥ (cid:96)√ n + e − Γ + Γ α √ n .
As with the random initialization result, this theorem shows that as long as the initial residual is sufﬁciently 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=
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 signiﬁcantly tighter bounds.
This is particularly important in light of recent literature Chizat & Bach (2018b); Ghorbani et al. (2019c); Yehudai & Shamir (2019) suggesting a signiﬁcant generalization gap between kernel methods/linearized neural nets when compared with neural nets operating beyond the linearized regime (e.g. mean ﬁeld regime).
As a result we view our deterministic result as a ﬁrst step towards moving beyond the NTK regime.
2. C ASE S TUDY : G AUSSIAN MIXTURE MODEL
To illustrate a concrete example, we consider a distribution based on multiclass mixture models.
Deﬁnition 2.4 (Gaussian mixture model) Consider a size n dataset {( x i , y i )} ni = ∈
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 ( (cid:96), ̃ (cid:96) ) denoting the ̃ (cid:96) th cluster from the (cid:96) th class.
We assume the dataset in cluster ( (cid:96), ̃ (cid:96) ) is centered around a cluster center µ (cid:96), ̃ (cid:96) ∈
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 ( (cid:96), ̃ (cid:96) ) are generated i.i.d.
according to
N ( µ (cid:96), ̃ (cid:96) ,σ 2 I d / d ) with the corresponding label set to the one-hot encoding of the class (cid:96) i.e. e (cid:96) .
Note that the cluster indexed by ( (cid:96), ̃ (cid:96) ) contains ̃ n (cid:96), ̃ (cid:96) data points satisfying E
[̃ n (cid:96), ̃ (cid:96) ]
=
̃ n = n /
KC .
This distribution is an ideal candidate to demonstrate why the Jacobian of the network exhibits low-rank or bimodal structure.
Let us consider the extreme case σ = 0 where we have a discrete input distribution over the cluster centers.
In this scenario, the multi-class Jacobian matrix is at most rank
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-simpliﬁed) Consider a data set of size n consisting of input/label pairs {( x i , y i )} ni = ∈
R d ×
R K generated according to Def. 2.4 with the standard deviation obeying σ ≲
Kn .
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 ≳ Γ 4 K 8 C 4 λ
M .
after T = 2Γ K 2 C λ M gradient iterations, with high probability, the model obeys Err D (
W T )
≲ Γ
[Figure]
√
KCn
J (
W) of a neural network with K = 3 .
Here, the data is generated according to the Def. 2.4 with K classes and σ = 0 .
The cluster centers are picked so that the distance between any two is at least 0 .
We consider two cases: n = 30 C (solid line) and n = 60 C (dashed line).
These plots demonstrate that the top KC singular values grow proportional to √ n .
We note that λ M captures how diverse the cluster centers are.
In this sense λ M > 0 intuitively means that neural network, speciﬁcally the neural tangent kernel, is sufﬁciently 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 α∼ λ 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 speciﬁcally k ≳
K 8 C 4 .
This compares favorably with Arora et al.
(2019) which concerns the K = case.
In particular, Arora et al.
(2019) requires k ≳ n/ λ 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 deﬁcient and λ X → so that Arora et al.
(2019) requires inﬁnite width.
P RIOR A RT
Neural networks have impressive generalization abilities even when they are trained with more parameters than the data Zhang et al.
Thus, optimization/generalization properties of neural nets have been the topic of recent literature Zhang et al.
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.
; 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.
; Neyshabur et al.
(2017b)
; Golowich et al.
Bartlett et al.
; 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.
leverages compression techniques for constructing tighter bounds.
Yin et al.
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 ﬁrst-order methods such as gradient descent have often favorable properties.
Generalization properties of stochastic gradient descent is extensively studied empirically Keskar et al.
; Hardt et al.
Sagun et al.
; Chaudhari et al.
; Hoffer et al.
; Goel & Klivans (2017); Goel et al.
For linearly separable datasets, Soudry et al.
; Gunasekar et al.
; Brutzkus et al.
(2017); Ji & Telgarsky (2018a;b) show that ﬁrst-order methods ﬁnd 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.
Mei et al.
; Sirignano & Spiliopoulos (2018); Rotskoff & Vanden-Eijnden (2018) relate the distribution of the network weights to Wasserstein gradient ﬂows using mean ﬁeld analysis.
Global convergence and generalization of neural nets: Closest to our work, recent literature Cao & Gu (2019); Arora et al.
; Ma et al.
Allen-Zhu et al.
(2018a) provides generalization bounds for overparameterized networks trained via gradient descent.
Also see Li et al.
; Huang et al.
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 inﬁnitely wide networks.
In connection to this, our Def 1.1 helps quantify the low-rankness and bimodality of the Jacobian spectrum (same as NTK for random initialization).
Similar to Thm 2.2, Arora et al.
(2019) uses the NTK to provide generalization guarantees in a similar framework to Jacot et al. (2018) (see (2.4) for comparison).
Li et al.
(2019a)
leverages low-rank Jacobian structure to establish robustness to label noise.
Very recent work Su & Yang (2019) uses low-rankness to better capture approximation power of neural nets.
These works build on global convergence results of randomly initialized networks Du et al.
(2018b;a); Allen-Zhu et al.
(2018b);
Chizat & Bach (2018b); Zhang et al.
; Nitanda & Suzuki (2019); Oymak & Soltanolkotabi (2018); Zou et al.
(2018) which study the gradient descent trajectory via comparisons to a NTK linearization.
These results however typically require unrealistically wide networks for optimization where the width grows poly in n and poly-inversely proportional to the distance between the input samples.
Example distance measures are class margin for logistic loss and minimum eigenvalue of the NTK matrix for least-squares.
Our work circumvents this issue by allowing a capacity-dependent interpolation.
We prove that even small networks (e.g. of constant width) can interpolate the data over a low-dimensional information space without making restrictive assumptions on the input.
This approach also leads to faster convergence rates.
In terms of generalization, our work has three distinguishing features: (a) bias-variance tradeoffs by identifying information/nuisance spaces, (b) no margin/distance/minimum eigenvalue assumptions on data, (c) the bounds apply to multiclass classiﬁcation 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 ﬁndings on the bimodal (approximately low-rank) spectrum Papyan (2018); Ghorbani et al.
(2019b); Papyan (2019b); Sagun et al.
; Li et al.
(2019b)
; Javadi et al.
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.
N UMERICAL EXPERIMENTS
[Figure]
We present experiments demonstrating our theoretical ﬁndings on two popular image classiﬁcation datasets.
In this section we focus on a set of CIFAR-10 experiments and discuss how our theory is strongly supported by what we observe in practice.
To provide more detail and show that our theory holds across different datasets, in addition to the experiments discussed in this section we perform additional experiments on a modiﬁed 3-class version of CIFAR-10 and
MNIST in Appendix C.
Experimental setup.
The CIFAR-10 dataset consists of 50 k training images and 10 k test images inclasses.
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 .
27 M ).
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 modiﬁcations to the default ResNet20 architecture: (1) we scale the output of the ﬁnal 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 ﬁt to the training data quickly.
The ﬁgures 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
∥ Π
I ( y )∥ (cid:96)∥ y ∥ (cid:96)∥ Π
N ( y )∥
(cid:96)∥ y ∥ (cid:96)∥ Π
I ( r)∥ (cid:96)∥ r∥ (cid:96)∥ Π
N ( r)∥ (cid:96)∥ r∥ (cid:96) 2 J traininit 0.
J trainfinal 0.
[Figure]
∥
[Figure]
∥
[Figure]
residual ( r τ = f
W τ ) − y
) along the information/nuisance spaces of the ﬁnal Jacobian on (a) the training data and (b) the test data and c) misclassiﬁcation error on training and test.
This experiment uses uncorrupted labels.
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 ( 500 k × 270 k ) 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 ﬁgure 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 ﬁt to the training data.
Our theory predicts that the sum of ∥
J [Equation]
I y [Equation] 2 and ∥ Π
N ( y )∥ (cid:96) 2 determines the classiﬁcation error (Theorems 2.2 and 6.24).
Table 1 collects these values for the initial and ﬁnal
Jacobian.
These values demonstrate that the label vector is indeed correlated with the top eigenvectors of both the initial and ﬁnal Jacobians.
An interesting aspect of these results is that this correlation increases from the initial to the ﬁnal Jacobian so that more of the label energy lies on the information space of the ﬁnal Jacobian in comparison with the initial Jacobian.
Stated differently, we observe a signiﬁcant 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 ﬁrst 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 ∥ (cid:96) 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.
[Figure] [Figure]
[Table]
[Figure] [Figure]
information/nuisance space using 50% label corrupted data and a Multi-class ResNet20 trained with SGD.
Experiments with label corruption.
Our next experiments study the effect of corruption.
Speciﬁcally, 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 ﬁt to the training data.
[Figure]
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 ﬁrst 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 overﬁts 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 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.
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.
T ECHNICAL APPROACH AND G ENERAL T HEORY 5. P RELUDE : FITTING A LINEAR MODEL
To gain better insights into what governs the generalization capability of gradient based iterations let us consider the simple problem of ﬁtting a linear model via gradient descent.
This model maps an input/feature vector x ∈
R d into a one-dimensional output/label via x ↦ f ( x , w ) ∶= w T x .
We wish to ﬁt a model of this form to n training data consisting of input/label pairs {( x i ,y i )} ni = ∈
R d ×
R .
Aggregating this training data as rows of a feature matrix X ∈
R n × d and label vector y ∈
R n , the training problem takes the form
L( w [Equation] =  ∥
[Equation]
− y ∥
[Equation]
(cid:96)
[Equation]
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 ﬁts to the clean uncorrupted data and not the corruption.
In this case the population loss (i.e. test error) takes the form
[Equation]
Now let us consider gradient descent iterations with a step size η which take the form
[Equation]
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 ﬁrst 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 misﬁt Xw − y on the column space of the uncorrupted portion of the input data (
U ) and its complement.
That is,
[Equation]
where
E
L I ( w ) ∶=
E [∥ Π
I (
Xw − y )
∥ 2 (cid:96)] =
∥
Xw − y ∥
(cid:96)+ r 2 nσ 2 x ∥ w ∥ 2 (cid:96)+ r 2 nσ 2 y ,
E
L N ( w ) ∶=
E [∥ Π
N (
Xw − y )
∥ 2 (cid:96)]
=  ( − r n )( σ 2 x ∥ w ∥ 2 (cid:96)+ σ 2 y ) , with Π
I =
UU T and Π
N =
I −
UU 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 ﬁrst 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 misﬁt/residual r τ =
Xw τ − y takes the form
[Equation]
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
[Equation]
w )
[Equation]
y
[Equation]
2 ) quickly decreases in the ﬁrst 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 ﬁt 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( x ; θ ) ,f( 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 )} ni = ⊂
R d ×
R K , we ﬁt to this data by minimizing a nonlinear least-squares loss of the form
[Equation]
To continue let us ﬁrst aggregate the predictions and labels into larger vectors based on class.
In particular deﬁne
[Equation]
and y
[Equation] [Equation]
y n
[Equation]
(cid:96) ⎤⎥⎥⎥⎥⎦ ∈
R n for (cid:96) = 1 , 2 ,...,K.
Concatenating these vectors we arrive at
[Equation]
and y
[Equation] [Equation]
Using the latter we can rewrite the optimization problem (5.3) into the more compact form
[Equation]
To solve this problem we run gradient descent iterations with a learning rate η starting from an initial point θ 0 .
These iterations take the form
[Equation]
∇L
[Equation] [Equation]
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 deﬁned below) that is close to the Jacobian at initialization
[Equation]
Deﬁnition 5.1 (Reference Jacobian and its SVD) Consider an initial point θ∈
R p and the Jacobian mapping
J
[Equation]
p .
For ε 0 ,β > 0 , we call J ∈
R Kn
[Equation]
max
[Equation]
Kn,p
[Equation]
an (
[Equation] [Equation]
β ) reference
Jacobian matrix if it obeys the following conditions,
[Equation]
and
[Equation]
where
[Equation]
is a matrix obtained by augmenting
[Equation]
zero columns.
Furthermore, consider the singular value decomposition of J given by
[Equation]
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 ( θ) .
However, we shall also use other reference Jacobians in our results.
We will compare the gradient iterations (5.6) to the iterations associated with ﬁtting a linearized model around θ 0 deﬁned as f lin
[Equation]
̃
[Equation]
) , where
[Equation]
is obtained from θ 0 by adding max (
[Equation]
) zero entries at the end of θ 0 .
The optimization problem for ﬁtting the linearized problem has the form
[Equation]
Thus starting from ̃ θ= θ
the iterates ̃ θ τ on the linearized problem take the form
[Equation]
[Equation]
The iterates based on the linearized problem will provide a useful reference to keep track of the evolution of the original iterates (5.6).
Speciﬁcally we study the evolution of misﬁt/residuals associated with the two problems
[Equation]
To better understand the dynamics of convergence of the linearized iterates next we deﬁne two subspaces associated with the reference Jacobian and its spectrum.
Deﬁnition 5.2 (Information/Nuisance Subspaces) Let J denote the reference Jacobian per Deﬁnition 5.1 with eigenvalue decomposition J =
U diag ( λ )
V T per (5.7) .
For a spectrum cutoff α obeying ≤ α ≤ λ 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 deﬁne the information and nuisance subspaces associated with J as
I ∶= span ({ u s } rs = ) and
N ∶= span ({ u s }
Kns = r + ) .
We also deﬁne the truncated reference Jacobian
J
I = [ uu... u r ] diag ( λ 1 ,λ 2 ,...,λ r )[ vv... 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 overﬁtting 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 θ∈
R p and positive scalars ε,R > 0 .
Assume that for any θ obeying ∥
[Equation]
R , we have ∥J
[Equation]
With these assumptions in place we are now ready to discuss our meta theorem that demonstrates that the misﬁt/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( θ ) =
∥ f ( θ ) − y ∥ 2 (cid:96) 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 θ∈
R p with J an ( (cid:15) 0 ,β ) reference Jacobian associated with
J ( θ) per Deﬁnition 5.1 and ﬁtting the linearized problem f lin (̃ θ ) =
f ( θ)
+
J (̃ θ − θ) via the loss
L lin ( θ ) =
∥ f lin ( θ ) − y ∥
2 (cid:96) 2 .
Furthermore, deﬁne 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 Deﬁnition 5.2.
Furthermore, assume the Jacobian mapping
J ( θ ) ∈
R nK × p associated with f obeys Assumptions 1 and 2 for all θ ∈
R p obeying ∥ θ − θ∥
[Equation]
.12) around a point θ∈
R p for a tolerance level δ obeying< δ ≤ 1 and stopping time Γ obeying Γ ≥ 1 .
Finally, assume the following inequalities hold
[Equation]
and
[Equation] [Equation]
We run gradient descent iterations of the form θ τ
[Equation]
and
[Equation]
θ
[Equation]
) on the original and linearized problems starting from θ 0 with step size η obeying η ≤ / β 2 .
Then for all iterates τ obeying
[Equation]
the iterates of the original ( θ τ ) and linearized ( ̃ θ τ ) problems and the corresponding residuals r τ
[Equation]
y closely track each other.
That is,
[Equation]
and
[Equation]
Furthermore, for all iterates τ obeying
[Equation]
Γ ηα∥ θ τ − θ∥ (cid:96)≤
R = ∥
J †
I r∥ (cid:96)+ Γ α ∥ Π
N ( r)∥ (cid:96)+ δ Γ α ∥ r∥ (cid:96) 2 .
(5.15) and after τ =
T iteration we have ∥ r T ∥ (cid:96)≤ e − Γ ∥
Π
I ( r)∥ (cid:96)+ ∥
Π
N ( r)∥ (cid:96)+ δα β ∥
r∥ (cid:96) 2 .
P ROOFS
Before we proceed with the proof let us brieﬂy 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 Tw T... w Tk ]
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 ﬁrst and second derivative.
6. P ROOFS FOR G ENERAL T HEORY (P ROOF OF T HEOREM 5.3)
In this section we prove our result for general nonlinearities.
We begin with a few notations and deﬁnitions 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. P RELIMINARY DEFINITIONS AND LEMMAS
Throughout we use
U
[Equation]
U
[Equation]
to denote the basis matrices for the information and nuisance subspaces from Deﬁnition 5.2.
Similarly, we deﬁne the information and nuisance spectrum as
[Equation]
We also deﬁne the diagonal matrices
[Equation]
Deﬁnition 6.1 (early stopping value and distance) Consider Deﬁnition 5.2 and let Γ > 0 be a positive scalar.
Associated with the initial residual r
[Equation]
y and the information/nuisance subspaces of the reference Jacobian J (with a cut-off level α ) we deﬁne the ( α, Γ ) early stopping value as
[Equation]
We also deﬁne the early stopping distance as
[Equation]
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 ﬁrst Lemma provides upper/lower bounds on the early stopping value.
Lemma 6.2 (Bounds on Early-Stopping Value)
The early stopping value
B α, Γ from Deﬁnition 6.1 obeys
[Equation]
[Equation]
Proof To prove the upper bound we use the fact that α ≤ λ s for s ≤ r and α ≥ λ s for s ≥ r to conclude that
[Equation]
To prove the lower bound, we use the facts that α/ λ
[Equation]
21 to conclude that
[Equation]
(⟨ u s , r⟩) ) / , ≥( r ∑ s = αλ 2 s (⟨ u s , r⟩) ) / , ≥ α λ∥ Π
I ( r)∥ (cid:96) 2 .
It is of course well known that the mapping (
I − η AA T ) is a contraction for sufﬁciently 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 η ≤ / β 2 we have ∥(
I − η AB T ) r ∥
(cid:96)≤ ( + ηε)∥ r ∥ (cid:96) 2 .
Proof Note that using η ≤ / β 2 and ∥
B −
A ∥ ≤ ε we conclude that ∥(
I − η AB T ) r ∥ (cid:96)=∥(
I − η BB T − η (
A −
B )
B T ) r ∥ (cid:96)=∥ r − η (
A −
B +
B )
B T r ∥ (cid:96)=∥ r ∥ 2 (cid:96)− 2 η r T (
A −
B +
B )
B T r + η∥
AB T r ∥ (cid:96)≤∥ r ∥ 2 (cid:96)− 2 η ∥
B T r ∥ 2 (cid:96)+ 2 η ∥(
A −
B )
T r ∥ (cid:96)∥
B T r ∥ (cid:96)+ η∥
A ∥ ∥
B T r ∥ 2 (cid:96)=∥ r ∥ 2 (cid:96)− η ∥
B T r ∥ 2 (cid:96)+ 2 η ∥(
A −
B )
T r ∥ (cid:96)∥
B T r ∥ (cid:96)+ ( η∥
A ∥ ∥
B T r ∥ 2 (cid:96)− η ∥
B T r ∥ 2 (cid:96))
η ≤ / β≤ ∥
r ∥ 2 (cid:96)− η ∥
B T r ∥ 2 (cid:96)+ 2 η ∥(
A −
B )
T r ∥ (cid:96)∥
B T r ∥ (cid:96)∥
A −
B ∥≤ (cid:15) ≤ ∥ r ∥ 2 (cid:96)− η ∥
B T r ∥ 2 (cid:96)+ 2 ηε ∥
B T r ∥ (cid:96)∥ r ∥ (cid:96)=( + ηε)∥ r ∥ 2 (cid:96)− η ( ε ∥ r ∥ (cid:96)− ∥
B T r ∥ (cid:96)) ≤( + ηε)∥ r ∥ 2 (cid:96) 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-deﬁnite 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 A + , B + ∈
R n × n obeying A + , B + ⪰ α 2
I n , Lemma 2.2 of Schmitt (1992) guarantees that ∥
A/ + −
B/ + ∥ ≤ ∥
A + −
B + ∥ α .
In the above for a PSD matrix A ∈
R n × n with an eigenvalue decomposition A =
U Λ U T we use
A/ 2 to denote the square root of the matrix given by A =
U Λ/ 2 U T .
We shall use this result with
A + =
A + α 2
I n and B + =
B + α 2
I n to conclude that ∥
A
[Equation] [Equation]
+ −
B/
[Equation]
+
[Equation]
A + −
B +
[Equation]
=
∥
A
[Equation]
B ∥ α .
Furthermore, using the fact that the eigenvalues of A + and B + are just shifted versions of the eigenvalues of A and B by α/ 4 we can conclude that ∥
A/ + −
A/ ∥ ≤ α and ∥
B/ + −
B/ ∥ ≤ α 2 .
Combining the latter two inequalities with the assumption that ∥
A −
B ∥ ≤ α 2 we conclude that ∥
A/ −
B/ ∥ ≤∥
A/ + −
B/ + ∥ + ∥
A/ + −
A/ ∥ + ∥
B/ + −
B/ ∥ ≤∥
A −
B ∥
α + α + α ≤ 2 α.
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 Λ/ 2 A V T A with V A ∈
R p × n an orthonormal matrix.
Now pick
Y
[Equation]
Then clearly Y Y T =
B .
Furthermore, we have
[Equation]
Combining the latter with (6.4) completes the proof.
6.1. K EY LEMMAS FOR GENERAL NONLINEARITIES
Throughout this section we assume J is the reference Jacobian per Deﬁnition 5.1 with eigenvalue decomposition J =
[Equation]
Λ
[Equation]
∑
Kns = 1 λ s u s v Ts with Λ = diag ( λ ) .
We also deﬁne a =
U T r=
U T ̃ r∈
R nK be the coefﬁcients of the initial residual in the span of the column space of this reference Jacobian.
We shall ﬁrst characterize the evolution of the linearized parameter ̃ θ τ and residual ̃ r τ vectors from (5.11) in the following lemma.
Lemma 6.
The linearized residual vector ̃ r τ can be written in the form
[Equation]
Furthermore, assuming η ≤ / λ 21 the linear updates
[Equation]
obey
[Equation]
Proof Using the fact that JJ T =
U Λ 2 U T we have
[Equation]
Using the latter combined with (5.11) we thus have
[Equation]
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
[Equation]
Using the latter together with the gradient update on the linearized problem we have
[Equation]
= ∇L lin (̃ θ t ))
= − η ( τ − ∑ t =
J T ̃ r t )
= − η V ( τ − ∑ t = Λ (
I − η Λ) t )
a .
Thus for any s ∈ { 1 , 2 ,...,nK }
[Equation]
Noting that for η ≤
[Equation]
s we have
[Equation]
, the latter identity implies that
[Equation]
Furthermore, using the fact that
[Equation]
Combining (6.7) for≤ s
[Equation]
r and (6.8) for s
[Equation]
r we have
[Equation]
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 ∥̃
[Equation]
Furthermore, after T =
Γ ηα 2 iterations
we have ∥̃
[Equation] [Equation]
.10) and ∥̃
[Equation]
with
B α, Γ given by (6.1) per Deﬁnition 6.2.
Proof To prove the ﬁrst bound on the residual ((6.9))note that using (6.5) we have
[Equation]
and U T
N ̃
r τ
[Equation]
Thus, using the fact that for s ≤ r we have λ s ≥ α we have ( − ηλ 2 s ) τ ≤ ( − ηα)
τ and for s > r
we have ( − ηλ 2 s )
τ ≤ 1 , we can conclude that ∥
[Equation]
ηα
[Equation] [Equation]
∥ (cid:96)and ∥
U T
N ̃
r τ ∥ (cid:96)
[Equation]
U
[Equation]
N
[Equation]
Combining these with the triangular inequality we have
[Equation]
concluding the proof of (6.9).
The second bound on the residual simply follows from the fact that (
[Equation]
Γ .
The bound on ∥̃
[Equation]
2 is trivially obtained by using T
[Equation] [Equation]
4 in (6.6).
The lemma above shows that with enough iterations, gradient descent on the linearized problem ﬁts 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
[Equation]
2 over the nuisance space will not affect generalization signiﬁcantly.
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, ∥ θ τ − θ∥ (cid:96)≤
R and ∥
θ τ + − θ∥ (cid:96)≤
R.
Then with a learning rate obeying η ≤ / β 2 , the deviation in the residuals of the original and linearized problems e τ
[Equation]
∥ (cid:96)
[Equation]
η (
[Equation]
+ εβ )∥̃ r τ ∥ (cid:96)+ ( + ηε)∥ e τ ∥ (cid:96) 2 .
Proof For simplicity, denote B= J (
[Equation]
, B
[Equation]
J ( θ
[Equation]
) , A = J ( θ) where
J ( b , a ) = ∫
J ( t b + ( − t ) a ) dt.
We can write the predictions due to θ τ + 1 as f
[Equation]
η
[Equation]
= f
[Equation]
= f
[Equation]
η
[Equation]
f
[Equation]
− y ) .
This implies that
[Equation]
Similarly, for linearized problem we have
[Equation]
[Equation] [Equation]
We proceed by bounding each of these two terms.
For the ﬁrst term, we apply Lemma 6.3 with
A =
B 1 and B =
B 2 and use ∥
B
[Equation]
ε to conclude that
[Equation] [Equation]
Next we turn our attention to bounding the second term.
To this aim note that
[Equation]
∥ ≤ ∥(
B−
A )
B T∥ + ∥
A (
B−
A )
T ∥ + ∥
AA T −
JJ T ∥ ≤ ∥
B−
A ∥∥
B∥ + ∥
B−
A ∥∥
A ∥ + ∥
AA T −
JJ T ∥ ≤ β ε + β ε + ε= ε+ εβ.
In the last inequality we use the fact that per Assumption 2 we have ∥
B−
A ∥ ≤ ε / 2 and ∥
B−
A ∥ ≤ ε / 2 as well as the fact that per Deﬁnition 5.
∥
AA T −
JJ T ∥ ≤ ε 20 .
Plugging (6.13) and (6.14) in (6.12) completes the proof.
Next we prove a result about the growth of sequences obeying certain assumptions.
As we will see later on in the proofs this lemma allows us to control the growth of the perturbation between the original and linearized residuals ( e τ = ∥
e τ ∥ (cid:96) 2 ).
Lemma 6.8 (Bounding residual perturbation growth for general nonlinearities) Consider positive scalars Γ ,α,ε,η > 0 .
Also assume η ≤ / α 2 and α ≥ √ 2Γ ε and set T =
Γ ηα 2 .
Assume the scalar sequences e τ (with e= 0 ) and ̃ r τ obey the following identities ̃ r τ ≤( − ηα)
τ ρ + + ρ − , e τ ≤( + ηε) e τ − + η Θ ̃
r
τ − 1 , (6.15) for all≤ τ ≤
T and non-negative values ρ − ,ρ + ≥ 0 .
Then, for all≤ τ ≤
T , e τ ≤ ΘΛ holds with Λ =
( Γ ρ − + ρ + ) α.
Proof We shall prove the result inductively.
Suppose (6.16) holds for all t ≤ τ − 1 .
Consequently, we have
[Equation]
Thus e
[Equation]
Summing up both sides of (6.17) for≤ t ≤ τ − 1 we conclude that e
[Equation]
[Equation]
≤ Λ
[Equation]
where in the last inequality we used the fact that α≥ 2Γ ε 2 .
This completes the proof of the induction step and the proof of the lemma.
6.1. C OMPLETING THE PROOF OF T HEOREM 5.
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 τ − ≤
T − 1 .
In particular, we assume the identities (5.14) and (5.15) hold for all≤ 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
[Equation]
R .
We ﬁrst argue that θ τ lies in the domain of interest as dictated by (5.12), i.e.
∥ θ τ − θ∥ (cid:96)≤
R .
To do this note that per the induction assumption (5.15) holds for iteration τ − 1 and thus ∥ θ
τ − − θ∥ (cid:96)≤
R / 2 .
As a result using the triangular inequality to show ∥ θ τ − θ∥ (cid:96)≤
R holds it sufﬁces to show
that ∥ θ τ − θ τ − ∥ (cid:96)≤
R / 2 holds.
To do this note
that ∥ θ τ − θ τ − ∥ (cid:96)= η ∥∇L( θ τ − )∥ (cid:96)= η ∥J
T ( θ τ − ) r τ − ∥ (cid:96)= η ∥J
T ( θ τ − ) r τ − ∥ (cid:96)( a ) ≤ η ∥J
T ( θ τ − )̃
r τ − ∥ (cid:96)+ η ∥J
T ( θ τ − )( r τ − − ̃ r τ − )∥ (cid:96)( b ) ≤ η ∥
J T ̃ r τ − ∥ (
cid:96)+ η ∥J ( θ τ − ) −
J ∥∥̃ r τ − ∥ (cid:96)+ η ∥J ( θ τ − )∥∥
r τ − − ̃ r τ − ∥ (cid:96)( c ) ≤ η ∥
J T ̃ r τ − ∥ (cid:96)+ ε+ ε β∥̃ r τ − ∥ (
cid:96)+ β
∥
r τ − − ̃ r τ − ∥ (cid:96)( d ) ≤ η ∥
J T ̃ r τ − ∥ (cid:96)+ 2 δα 5 β∥ r∥
(cid:96)+ β ∥ r τ − − ̃ r τ − ∥ (cid:96)( e ) ≤ η ∥
J T ̃ r τ − ∥ (cid:96)+ 2 δα 5 β∥ r∥
(cid:96)+ 3 δα 5 β∥ r∥ (cid:96)= η ∥
J T ̃ r τ − ∥ (cid:96)+ δα β∥ r∥ (cid:96)( f ) ≤ ηβ
B α, Γ α + δα β∥ r∥ (cid:96)( g ) ≤ B α, Γ α + δα β∥ r∥ (cid:96)( h ) ≤ B α, Γ α + δ Γ α ∥ r∥ (cid:96)=
R 2 .
Here, (a) and (b) follow from a simple application of the triangular inequality, (c) from the fact that ∥J ( θ τ − ) −
[Equation]
≤
[Equation] [Equation]
[Equation] [Equation]
J ∥
[Equation]
ε + ε 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
[Equation] [Equation]
V Λ (
I − η Λ)
τ − 1 U T r∥ (cid:96)=∥ Λ (
I − η Λ)
τ − 1 U T r∥ (cid:96)= nK ∑ s = λ 2 s ( − ηλ 2 s ) ( τ − ) (⟨ u s , r⟩)
≤ nK ∑ s = λ 2 s (⟨ u s , r⟩) =
r ∑ s = λ 2 s (⟨ u s , r⟩)
+ nK
∑ s = r + λ 2 s (⟨ u s , r⟩) ≤ β 4 r ∑ s =  λ 2 s (⟨ u s , r⟩)
+ nK
∑ s = r + λ 2 s (⟨ u s , r⟩) ≤ β( r ∑ s =  λ 2 s (⟨ u s , r⟩)
+ Γ 2 nK ∑ s = r + λ 2 s
α(⟨ u s , r⟩) )
=
β(B α, Γ α )
(g) from the fact that η ≤ 1 β 2 , and (h) from the fact that α
≤ β and Γ ≥ 1 .
Step
II: Original and linearized residuals are close
ﬁrst
part of (5.14) ).
In this step we wish to show that the ﬁrst part of (5.14) holds for iteration τ .
Since we established in the previous step that
[Equation]
R the assumption of Lemma 6.7 holds for iterations τ − 1 and τ .
Hence, using Lemma 6.7 equation (6.11) we conclude that
[Equation]
This combined with the induction assumption implies that
[Equation]
holds for all t ≤ τ ≤
T .
Furthermore, using Lemma 6.5 equation (6.9) for all t ≤ τ ≤
T we have ∥̃
[Equation]
To proceed, we shall apply Lemma 6.8 with the following variable substitutions
[Equation]
We note that Lemma 6.8 is applicable since (i) η ≤ / β≤ / α 2 , (ii) based on (5.13) we have αε ≥ 5Γ δ β 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
[Equation] [Equation]
where in the last inequality we used (5.13).
This completes the ﬁrst part of (5.14) via induction.
Step III: Original and linearized parameters are close (second part of (5.14) ).
In this step we wish to show that the second part of (5.14) holds for iteration τ .
To do this we begin by noting that by the fact that J is a reference Jacobian we have ∥J ( θ) −
J ∥ ≤ ε 0 where
J augments
J ( θ) by padding zero columns to match size of J .
Also by Assumption 2 we have ∥J ( θ ) − J ( θ)∥ ≤ ε 2 .
Combining the latter two via the triangular inequality we conclude that
[Equation]
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 ∥ θ τ − ̃ θ τ ∥ (cid:96)η = ∥
τ − ∑ t = ∇ ¯
L( θ t ) − ∇L lin
(̃ θ t )∥
(cid:96)= ∥ τ − ∑ t =
J ( θ t )
T r t −
J T ̃ r t ∥
(cid:96)≤ τ − ∑ t = ∥J ( θ t )
T r t −
J T ̃ r t ∥
(cid:96)≤ τ − ∑ t =
∥(J ( θ t ) −
J )
T ̃ r t ∥
cid:96)+ ∥J ( θ t )
T ( r t − ̃ r t )∥ (cid:96)= τ − ∑ t =
∥(J ( θ t ) −
J )
T ̃ r t ∥
cid:96)+ ∥J ( θ t )
T e t ∥ (cid:96)≤ τ − ∑ t = ( ε + ε)∥̃ r t ∥ (cid:96)+ β ∥ e t ∥ (cid:96) 2 .
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 ﬁrst term we use the fact that ∥̃ r τ ∥ (cid:96)
[Equation]
to conclude
[Equation]
To bound the second term in (6.24) we use (6.21) together with τ
[Equation]
≤
[Equation]
2 to conclude that
[Equation] [Equation]
Combining (6.25) and (6.26) in (6.24), we conclude that
[Equation]
∥ (cid:96)=  δ
[Equation]
∥ r
[Equation]
∥ (cid:96) 2 .
Here, (a) follows from ε ≤ δα 3 5Γ β 2 per Assumption (5.13), (b) from ε≤ √ δα 3 Γ β per Assumption (5.13), (c) from ε ≤ δα 3 5Γ β≤ δα 5Γ ≤ δα 5 per Assumption (5.13), and (d) from ε≤ δα 5 per Assumption (5.13).
Thus,
[Equation]
Combining the latter with the fact that
[Equation]
α (which follows from Lemma 6.5 equation (6.6))
we conclude that
[Equation]
I
[Equation]
∥
[Equation] [Equation]
)∥ (cid:96)
[Equation]
δ
[Equation]
∥ r∥ (cid:96)
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 ∥ (cid:96)
[Equation]
( a ) ≤ ∥̃ r T ∥ (cid:96)+ ∥̃ r T − r T ∥ (cid:96)( b ) ≤ ∥̃ r T ∥ (cid:96)+ δα β ∥ r∥ (cid:96)( c )
≤ e − Γ ∥
Π
I ( r)∥ (cid:96)+ ∥
Π
N ( r)∥ (cid:96)+ δα β ∥
r∥ (cid:96)where (a) follows from the triangular inequality, (b) from the conclusion of Step II (ﬁrst part of (5.14)), and (c) from Corollary 6.6 equation (6.10).
This completes the proof of (5.16).
6. K EY 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 misﬁt.
For two matrices
[Equation]
we deﬁne their Khatri-Rao product as A
[Equation]
mn , where ⊗ denotes the Kronecker product.
6.2. P RELIMINARY IDENTITIES AND CALCULATIONS
We begin by discussing a few notations.
Throughout we use w (cid:96) and v (cid:96) to denote the (cid:96) 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 φ (
[Equation]
) , the Jacobian matrix with respect to vect (
[Equation]
∈
R kd takes the form
[Equation] [Equation]
27) where
J (cid:96) (
W ) is the Jacobian matrix associated with the (cid:96) th class.
In particular,
J (cid:96) (
W ) is given by
[Equation]
with
[Equation]
Alternatively using Khatri-Rao products this can be rewritten in the more compact form
J
[Equation] [Equation]
[Equation]
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 TK ]
T .
We have mat (J
T (
W ) u )
[Equation]
K ∑
[Equation]
=
[Equation] [Equation]
[Equation] [Equation]
[Equation]
X
[Equation]
6.2. F UNDAMENTAL PROPERTIES OF THE J ACOBIAN 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 ∥ (cid:96) ∞
∥
X ∥ and ∥ mat (J
T (
W ) u )
∥ 2 , ∞ ≤
B ∥ v ∥ (cid:96) ∞
∥
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 ∥ (cid:96)= 1 ).
Then the Jacobian mapping is Lipschitz with respect to spectral norm i.e. for all ̃
W , W
[Equation]
we have
[Equation]
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
[Equation]
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
[Equation]
.31) and ∥ mat
[Equation]
Furthermore, suppose φ is twice differentiable and ∣
[Equation]
)∣ ≤
B for all z .
Also assume all data points have unit Euclidean norm (
[Equation]
Then the Jacobian mapping is Lipschitz with respect to spectral norm i.e. for all ̃
[Equation]
W
[Equation]
d we have
[Equation]
Proof The proof will follow from Lemma 6.9.
First, given A =
A T 1 ...
A TK ]
T and B =
B T 1 ...
B TK ]
T , observe that
[Equation]
K
[Equation]
∥ and
[Equation]
These two identities applied to the components
J (cid:96) (
W ) and
[Equation]
W )
[Equation]
J
[Equation]
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
[Equation]
where the penultimate inequality follows from Cauchy Schwarz, completing the proof.
6.2. P ROPERTIES OF THE J ACOBIAN 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≤ (cid:96), ̃ (cid:96) ≤
K , let M [ (cid:96), ̃ (cid:96) ] denote the ( (cid:96), ̃ (cid:96) )
th submatrix.
For C (
W )
=
J (
W )J (
W )
T we have
C (
W )
[ (cid:96), ̃ (cid:96) ] =
[Equation]
∑
[Equation]
Suppose W i.i.d.
∼ N( 0 ,) 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
( (cid:96), ̃ (cid:96) )
th submatrix of C (
W ) is given by
[Equation]
[Equation]
Setting W i.i.d.
∼ N( 0 ,) and V with i.i.d.
zero-mean and ν 2 -variance entries, we conclude that
E
[Equation]
= k
∑
[Equation]
) ⊙ (
E [
V (cid:96),s V ̃ (cid:96),s ]
E [ φ ′ (
Xw s ) φ ′ (
Xw s )
T ]) =
k ∑ s = ν 2 δ ( (cid:96) − ̃ (cid:96) )
XX T ) ⊙
E [ φ ′ (
Xw s ) φ ′ (
Xw s )
T ]] =
kδ ( (cid:96) − ̃ (cid:96) ) ν
[Equation]
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-Deﬁnite (PSD) matrices.
Then,
[Equation]
min
[Equation]
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 ∣ φ ( )∣ ≤
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 ∥ (cid:96)= 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∈
R k × d with i.i.d.
N( 0 ,) entries obeys
[Equation]
with probability at least− / n 100 .
In particular, as long as
[Equation]
with the same probability, we have that
[Equation]
Proof Deﬁne C =
J (
W)J (
W)
T .
We begin by showing that the diagonal blocks of C are concentrated.
To do this ﬁrst for≤ s ≤ k deﬁne the random matrices
[Equation]
Now consider n × n diagonal blocks of C (denoted by C [ (cid:96),(cid:96) ] ) and note that we have
[Equation]
Furthermore, using Lemma 6.
[Equation]
Also, using Jensen’s inequality
[Equation]
Combining the latter two identities via the triangular inequality we conclude that
[Equation]
To proceed, we will bound the weighted sum
[Equation]
]) in spectral norm.
To this aim we utilize the Matrix Hoeffding inequality which states that
[Equation]
where ∆ 2 is an upper bound on
[Equation]
Using (6.34) we can pick ∆= ∑ ks =
( 2 ν 2 B∥
X ∥ )
= 4 kν 4 B∥
X ∥ 4 .
Setting t = √ kν 2 B∥
X ∥ √ log ( n ) , we conclude that
P {∥
C [ (cid:96),(cid:96) ] −
E [
C [ (cid:96),(cid:96) ]]∥ ≥ t } =
P (∥
S ∥ ≥ t ) ≤ n − concluding the proof of concentration of the diagonal blocks of C .
For the off-diagonal blocks C [ (cid:96), ̃ (cid:96) ] using (6.33) from the proof of Lemma 6.11 we have that
C [ (cid:96), ̃ (cid:96) ]
[Equation]
k ∑ s =
[Equation]
V
[Equation]
V ̃ (cid:96),s A s .
Note that by construction {
V (cid:96),s V ̃ (cid:96),s } ks
= 1 are i.i.d.
± ν 2 Rademacher variables and thus C [ (cid:96), ̃ (cid:96) ] 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 ∥
[Equation]
∥
X ∥
[Equation]
so that we can take ∆=
[Equation]
4 kB∥
X ∥ 4 and again conclude that for t = √ kν 2 B∥
X ∥ 2 log ( n ) we have
P {∥
C [ (cid:96), ̃ (cid:96) ]
∥ ≥ t } ≤ n −
Using the fact that E [
C [ (cid:96), ̃ (cid:96) ]]
= 0 and K ≤ n , combined with a union bound over all sub-matrices ≤ (cid:96), ̃ (cid:96) ≤
K we conclude that
P
[Equation]
C
[Equation]
E
[Equation]
All that remains is to combine the concentration results for the sub-matrices to arrive at the complete bound.
In mathematical terms we need to bound D ∶= ∥
C −
E [
C ]∥ .To
this aim deﬁne D [ (cid:96), ∶] to denote the (cid:96) th block row of D .
Standard bounds on spectral norm in terms of sub-matrices allow us to conclude that
[Equation]
concluding the proof.
The result in terms of δ is obtained by using the population covariance Lemma 6.11.
6.2. U PPER BOUND ON INITIAL RESIDUAL
In this section we prove a lemma concerning the size of the initial misﬁt.
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 ∣ φ ( )∣ , ∣ φ ′ ( z )∣ ≤
B .
Also assume we have n data points x 1 , x 2 ,..., x n ∈
R d with unit euclidean norm ( ∥ x i ∥ (cid:96)= 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 ν
[Equation]
kn and the entries of W ∈
R k × d are i.i.d.
N( 0 ,) .
Then, ∥
V φ (
W X T )∥
F ≤ ν ∥ y ∥ (cid:96) 2 , holds with probability at least− ( 2 K ) − 100 .
Proof We begin the proof by noting that
[Equation]
We will show that for any row v of V , with probability at least− ( 2 K ) − 101 , ∥ v T(cid:96)
φ
[Equation]
W
[Equation]
∥
[Equation]
≤
[Equation]
(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 ∥ (cid:96)
[Equation]
∥ φ (
[Equation]
) v
[Equation]
≤∥ φ (
[Equation]
̃
[Equation]
) v − φ (
[Equation]
) v ∥ (cid:96)≤∥ φ (
X ̃
[Equation]
) − φ (
[Equation]
T )
∥∥ v ∥ (cid:96)≤∥ φ (
X ̃
W T ) − φ (
XW T )∥
F ∥ v ∥ (cid:96)( a ) =
∥( φ ′ (
S ⊙
X ̃
W
[Equation] [Equation]
k
[Equation]
n −
S ) ⊙
XW T )) ⊙ (
X (̃
W −
W )
T )∥
F ∥ v ∥ (cid:96)≤
B ∥
X (̃
W −
W )
T ∥
F ∥ v ∥ (cid:96)≤
B ∥
X ∥∥ v ∥ (cid:96)∥̃
W −
W ∥
F , where in (a) we used the mean value theorem with S a matrix with entries obeying≤
S i,j ≤ and 1 k × n the matrix of all ones.
Thus, ∥
[Equation]
v
[Equation]
X
[Equation]
Lipschitz function of W .
Thus, ﬁxing v , for a matrix W with i.i.d.
Gaussian entries
[Equation] [Equation]
36) holds with probability at least− e −
[Equation]
v
[Equation]
2 ∥ X ∥ 2 .
Next given g ∼ N( 0 ,) , we have ∣
E
[Equation]
and Var
[Equation]
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
[Equation]
holds with probability− ( 2 K ) − 102 .
Conditioned on E v , we now upper bound the expectation via
E [∥ φ (
[Equation]
) v ∥
[Equation]
] ( a )
[Equation]
√ [Equation] [Equation] φ (
[Equation]
v ∥
[Equation]
] = ¿``(cid:192) n
∑ i =
[Equation]
E [( v
[Equation]
W
[Equation] [Equation]
] ( b ) =
[Equation] [Equation] [Equation]
( g )) ]
( c ) = √ n √∥ v ∥ 2 (cid:96) 2 E g ∼N( 0 ,)
[( φ ( g ) −
E [ φ ( g )]) ]
+ ( 1 T v ) (
E g ∼N( 0 ,) [ φ ( g )]) ( d ) ≤ √ n ∥ v ∥ (cid:96) 2 √× 4 B 2 log ( 2 K ) +
B≤ √ n log ( 2 K )
B ∥ v ∥ (cid:96) 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 = √ n log ( 2 K )
B ∥ v ∥ (cid:96) 2 in (6.36), conditioned on E v we conclude that
[Equation]
∥ (cid:96)≤ √ n log ( 2 K )
B ∥ v ∥ (cid:96)= √ n log ( 2 K )
B √ k ν ∥ y ∥ (cid:96)50 B √
K log ( 2 K ) kn = ν ∥ y ∥ (cid:96)√
K , (6.39) holds with probability at least− exp (− 102log ( 2 K ) n ∥
X ∥ ) ≥ −( 2 K ) − 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− ( 2 K ) − 101 concluding the proof of (6.35) and the Lemma.
6. R ADEMACHER 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 } ni = 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 )} ni = ∼ D , we deﬁne the empirical
Rademacher complexity to be
R S (F) = n E [ sup f ∈F n ∑ i = ε
Ti f ( x i )] .
We begin by stating a vector contraction inequality by Maurer Maurer (2016).
This is obtained by setting h i ( f ( x i ))
=
h ( y i ,f ( x i )) in Corollary 4 of Maurer (2016).
Lemma 6.15 Let f (⋅) ∶
R d →
R K and let (cid:96) ∶
R K ×
R K →
R be a 1 Lipschitz loss function with respect to second variable.
Let { ε i } n i = 1 be i.i.d.
Rademacher variables.
Given i.i.d.
samples {( x i , y i )} n i = 1 , deﬁne
[Equation]
We have that
[Equation]
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 (cid:96) (⋅ , ⋅)
∶
R K ×
R K → [ 0 ,] be a 1 Lipschitz loss function.
Given i.i.d.
samples {( x i , y i )} ni = 1 , consider the empirical loss
L(
[Equation]
(cid:96)
[Equation]
n n ∑
[Equation]
( x
[Equation]
With probability− δ over the samples, for all f ∈ F , we have that
[Equation]
Proof Based on Bartlett & Mendelson (2002),
[Equation]
holds with− δ 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
[Equation]
V , W ) =
V φ (
W x ) with
W ∈
R k × d and V ∈
R K × k denoting the input and output weight matrices.
Suppose V∈
R K × k is a matrix obeying
[Equation]
ν /√ kK .
Also let W∈
R k × d be a reference input weight matrix.
Furthermore, we deﬁne 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
[Equation]
V
[Equation]
V
[Equation]
∥
[Equation]
νM
[Equation]
√
Kk } and
W = {
W ∶ ∥
W −
W∥
F ≤
M
W and ∥
W −
W∥ 2 , ∞ ≤
R √ k } .
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 ∥ (cid:96)= 1 ).
Also, deﬁne the average energy at W 0 as
[Equation]
∥ 2 (cid:96)) / .
Also let { ξ i } ni = ∈
R K be i.i.d.
vectors with i.i.d.
Rademacher entries and deﬁne the empirical
Rademacher complexity
[Equation]
⎤⎥⎥⎥⎥⎦ .
Then,
[Equation]
Proof We shall use w (cid:96) to denote the rows of W (same for W 0 , V , V 0 ).
We will approximate φ (⟨ w (cid:96) , x i ⟩) by its linear approximation φ (⟨ w 0 (cid:96) , x i ⟩)+ φ ′ (⟨ w 0 (cid:96) , x i ⟩)(⟨ w (cid:96) − w 0 (cid:96) , x i ⟩) via the second order Taylor’s mean value theorem.
We thus have
R
[Equation]
⎡⎢⎢⎢⎢⎣
[Equation]
⎡⎢⎢⎢⎢⎣
[Equation]
⎤⎥⎥⎥⎥⎦ ·„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„‚„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„¶
[Equation]
[Equation]
⎡⎢⎢⎢⎢⎣
[Equation]
−
[Equation] [Equation]
⟨ w
[Equation]
x
[Equation]
⟩)(⟨ w
[Equation]
−
[Equation] [Equation]
⎤⎥⎥⎥⎥⎦ ·„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„‚„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„¶ R + n E [ sup V ∈V , W ∈W n ∑ i =
ξ Ti (
V −
V)( φ (
W x i ) − φ (
W 0 x i ))] ·„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„‚„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„¶ R + n E [ sup V ∈V n ∑ i =
ξ Ti (
V −
V) φ (
W 0 x i )] ·„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„‚„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„¶ R
We proceed by bounding each of these four terms.
For the ﬁrst term note that
R ≤ n E ⎡⎢⎢⎢⎢⎣ sup ∥
W −
W∥
F ≤
M
W n ∑ i = ξ Ti V 0 diag ( φ ′ (
W 0 x i ))
W −
W) x i ⎤⎥⎥⎥⎥⎦ ≤ n E ⎡⎢⎢⎢⎢⎣ sup ∥
W −
W∥
F ≤
M
W ⟨ n ∑ i = diag ( φ ′ (
W 0 x i ))
V T 0 ξ i x Ti , W −
W⟩ ⎤⎥⎥⎥⎥⎦ ≤
M
W n E ⎡⎢⎢⎢⎢⎣∥( n ∑ i = diag
V T 0 ξ i ) φ ′ (
W 0 x i ) x Ti )∥
F ⎤⎥⎥⎥⎥⎦ ≤
M
W n E ⎡⎢⎢⎢⎢⎣∥( n ∑ i = diag
V T 0 ξ i ) φ ′ (
W 0 x i ) x T i )∥
F ⎤⎥⎥⎥⎥⎦ / =
M
W n ⎡⎢⎢⎢⎢⎣ n ∑ i = 1 E ∥ diag (
V T 0 ξ i ) φ ′ (
W 0 x i ) x
Ti ∥
F ⎤⎥⎥⎥⎥⎦ / =
M
W n ⎡⎢⎢⎢⎢⎣ n ∑ i = 1 E ∥ diag (
V T 0 ξ i ) φ ′ (
W 0 x i )∥ 2 (cid:96)⎤⎥⎥⎥⎥⎦ / ≤
BM
W n ⎡⎢⎢⎢⎢⎣ n ∑ i = 1 E ∥
V T 0 ξ i ∥ 2 (cid:96)⎤⎥⎥⎥⎥⎦ / ≤
BM
W n ∥
V∥
F ≤
BM
W ν √ n , where in the last inequality we used the fact that ∥
V∥
F ≤ ν .
For the second term note that
R ≤ 2 n E ⎡⎢⎢⎢⎢⎣ sup ∥
W −
W∥ 2 , ∞ ≤
R n ∑ i =
k ∑ (cid:96) =
K ∑ j = ξ i,j v 0 ,j,(cid:96) φ ′′ (( − t i(cid:96) )
⟨ w 0 (cid:96) , x i ⟩ + t i(cid:96) ⟨ w (cid:96) , x i ⟩)(⟨ w (cid:96) − w 0 (cid:96) , x i ⟩) ⎤⎥⎥⎥⎥⎦ ≤ 2 n k
∑ (cid:96) =
1 E ⎡⎢⎢⎢⎢ ⎣ sup ∥ w (cid:96) − w 0 (cid:96) ∥
(cid:96)≤
R n ∑ i = 1 RRRRRRRRR
RR
K ∑ j = ξ i,j v 0 ,j,(cid:96)
RRRRRRRRR
RR ∣ φ ′′ (( − t i(cid:96) )
⟨ w 0 (cid:96) , x i ⟩ + t i(cid:96) ⟨ w (cid:96) , x i ⟩)∣(⟨ w (cid:96) − w 0 (cid:96) , x i ⟩) ⎤⎥⎥⎥⎥ ⎦ ≤ 2 kn k
∑ (cid:96) =
n ∑ i = 1 E ⎡⎢⎢⎢⎢⎣RRRRRRRRRRR
K ∑ j = ξ i,j v 0 ,j,(cid:96)
RRRRRRRRRRR⎤⎥⎥⎥⎥⎦
R 2 B ≤
BR2 k ∥
V T∥ 2 ,≤
BR 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∥ 2 ,≤ ν √ k .
To bound the third term note that
R = n E [ sup V ∈V , W ∈W n ∑ i =
ξ Ti (
V −
V)( φ (
W x i ) − φ (
W 0 x i ))] ≤ n E [ sup V ∈V , W ∈W n ∑ i = ∥
V −
V∥
F ∥ ξ i ∥ (cid:96)∥( φ (
W x i ) − φ (
W 0 x i ))
∥ (cid:96)] ≤ νM
V n √ kK E [ sup W ∈W n ∑ i = ∥
ξ i ∥ (cid:96)∥ φ (
W x i ) − φ (
W 0 x i )∥
(cid:96)] ≤ νM
V n √ k
⋅ sup W ∈W n ∑ i = ∥
φ (
W x i ) − φ (
W 0 x i )∥
(cid:96)≤ νBM
V n √ k
⋅ sup W ∈W n ∑ i = ∥(
W −
W)
x i ∥ (cid:96)≤ νBM
V n √ k
⋅ sup W ∈W n ∑ i = ∥(
W −
W)∥
F = νBM
V √ k ⋅ sup W
[Equation]
k .
Finally, to bound the fourth term note that we have
R
[Equation] [Equation]
n ∑ i
[Equation]
V
[Equation] [Equation]
[Equation] [Equation]
∥
[Equation]
F
[Equation]
∑ i
[Equation] [Equation]
Ti
[Equation] [Equation]
n √ kK
[Equation]
E
[Equation]
n
[Equation]
i = φ
[Equation]
W
[Equation]
Ti ∥
F ]
[Equation]
/ = νM
[Equation]
n √ kK
[Equation]
n ∑ i =
[Equation]
E
[Equation]
Ti
[Equation]
F ]
[Equation] [Equation]
= νM
[Equation]
√ n
[Equation]
kn
[Equation] [Equation]
V
[Equation]
νBEM
V √ n
Combining these four bounds we conclude that
[Equation]
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∈
R K × k is a matrix obeying ∥
V∥ (cid:96) ∞ ≤ ν /√ kK .
Also let W∈
R k × d be a reference input weight matrix.
Also deﬁne the empirical losses
L(
V , W )
=
n n ∑ i =
[Equation]
∥
y
[Equation]
− f (
[Equation]
V , W )
[Equation]
(cid:96)
[Equation]
, and
L( f,(cid:96) )
=
n n ∑ i = (cid:96) ( f ( x
[Equation]
; V , W )
[Equation]
y
[Equation]
) , with (cid:96) ∶
R K ×
R K → [ 0 ,] a one Lipschitz loss function obeying (cid:96) ( y , y )
= 0 .
Additionally, assume the training data {( x i , y i )} ni = 1 are generated i.i.d.
according to a distribution
D with the input data points of unit Euclidean norm (i.e.
∥ x i ∥ (cid:96)= 1 ).
Also, deﬁne the average energy at W 0 as
E
[Equation] [Equation]
∑
[Equation] [Equation]
[Equation]
Then for all f in the function class
F W given by (6.40)
E [L( f,(cid:96) )] ≤
[Equation]
V , W )
+
[Equation]
M
W √ n
[Equation]
V
M
W √
[Equation]
EM
V √ n
[Equation]
√ 5log
[Equation]
6.42) holds with probability at least− δ .
Furthermore, Suppose labels are one-hot encoded and thus unit
Euclidian norm.
Given a sample
[Equation]
x , y
[Equation]
K generated according to the distribution
D , deﬁne the population classiﬁcation error
[Equation]
Then, we also have
Err
[Equation]
⎤⎥⎥⎥⎥⎦
[Equation]
Proof To begin ﬁrst note that any 1-Lipschitz (cid:96) with (cid:96) ( y , y )
=
obeys (cid:96) ( y
[Equation]
ˆ y ) ≤ ∥ y − ˆ y ∥ (cid:96) 2 .
Thus, we have
L( f
[Equation]
where the last inequality follows from Cauchy-Schwarz.
Consequently, applying Lemmas 6.16 and 6.17 we conclude that
E
[Equation]
V
[Equation]
W
[Equation]
k
[Equation]
V √ n
[Equation]
√ 5log (
[Equation]
/ δ ) n , which yields the ﬁrst statement.
To prove the second statement on classiﬁcation accuracy, we pick the (cid:96) function as follows
[Equation]
Note that, given a sample
[Equation]
K with one-hot encoded labels, if
[Equation]
this implies
[Equation]
Combining the latter with Markov inequality we arrive at
[Equation]
Now since (cid:96) is 1 Lipschitz and bounded, it obeys (6.42), which combined with the above identity yields (6.43), completing the proof.
6. P ROOFS FOR NEURAL NETS WITH ARBITRARY INITIALIZATION (P ROOF OF T HEOREM 2.3)
In this section we prove Theorem 2.3.
We ﬁrst 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. P RELIMINARY O PTIMIZATION R ESULT
Lemma 6.19 (Deterministic convergence guarantee) Consider a one-hidden layer neural net of the form x ↦
[Equation]
) with input weights W
[Equation]
and output weights V
[Equation]
K
[Equation]
and an activation φ obeying ∣
[Equation] [Equation]
z )∣ ≤
B , and ∣ φ
[Equation]
( z )∣
[Equation]
B for all z .
Also assume V is ﬁxed with all entries bounded by
[Equation]
and we train over W based on the loss
[Equation]
Also, consider a point W∈
R k × d with J an ( (cid:15)
[Equation]
,νB ∥
X ∥) reference Jacobian associated with
J (
W) per Deﬁnition 5.1.
Furthermore, deﬁne 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 Deﬁnition 5.2.
Let the initial residual vector be r [Equation] y
[Equation]
R nK .
Furthermore, assume
[Equation]
6.44) and
[Equation]
6.45) with≤ δ ≤ 1 and Γ ≥ 1 .
We run gradient descent iterations of the form W τ +
[Equation]
starting from W 0 with step size η obeying η ≤ 1 ν 2 B∥
X ∥ 2 .
Then for all iterates τ obeying
[Equation]
∶
[Equation]
Γ
[Equation] [Equation]
Furthermore, after τ =
T iteration we have
[Equation]
Proof To prove this lemma we wish to apply Theorem 5.3.
We thus need to ensure that the assumptions of this theorem are satisﬁed.
To do this note that by Lemma 6.10 Assumption 1 holds with β = νB ∥
X ∥ .
Furthermore, we pick ε
[Equation]
X ∥ = δα 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 ∈
[Equation]
obeying
[Equation]
(cid:96)
[Equation]
as long as (6.45) holds by Lemma 6.10 we have
[Equation]
Thus, Assumption 2 holds with
[Equation]
Now that we have veriﬁed 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 ( τ ) (cid:96) denote the (cid:96) th row of W τ .
Also note that
[Equation]
Hence, using Lemma 6.10 equation (6.32) we conclude that
[Equation]
Consequently, for any row
[Equation]
k , we have
[Equation]
To bound the right-hand side we use the triangular inequality combined with (6.25) and (6.26) to conclude that
[Equation]
[Equation]
where in the last inequality we used the fact that ε≤
[Equation]
2 25Γ per (6.44) and (cid:15)β = δα 3 5Γ β ≤ α 2 5Γ per our choice of (cid:15) .
Combining (6.49) and (6.50), we obtain ∥ w
[Equation]
completing the proof of (6.47) and the theorem.
6.4. M AIN O PTIMIZATION R ESULT
Lemma 6.20 (Deterministic optimization guarantee) Consider the setting and assumptions of
Lemma 6.19.
Also assume
[Equation]
for a constant c
[Equation]
if ε> 0 .
Furthermore, assume
[Equation]
and
[Equation]
and Γ ≥ 1 .
We run gradient descent iterations of the form W τ
[Equation]
) starting from
W 0 with step size η obeying η ≤
[Equation]
X ∥ 2 .
Then for all iterates τ obeying
[Equation] [Equation]
Furthermore, after τ =
T iteration we have
[Equation]
[Equation]
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
[Equation]
ζνB ∥
X ∥
+ Γ ∥ r∥ (cid:96)
B α, Γ ) 2 ν 6 B∥
X ∥ 6 Γ
B 2 α, Γ α≥ max ( α ζνB ∥
X ∥ , Γ ∥
r∥ (cid:96)
B α, Γ ) 2 1600 ν 6 B∥
X ∥ 6 Γ
B 2 α, Γ α≥ min ( ζνB ∥
X ∥
α ,
B α, Γ Γ ∥ r∥ (cid:96))
ν 6 B∥
X ∥ 6 Γ
B 2 α, Γ α= 1600Γ 2 ν 6 B∥
X ∥
B 2 α, Γ δ 2 α= 400Γ 2 ν 6 B∥
X ∥
(B α, Γ + B α, Γ ) δ 2 α≥ 400Γ 2 ν 6 B∥
X ∥
(B α, Γ + δ Γ ∥ r∥ (cid:96))
δ 2 α, so that (6.45) holds.
We thus turn our attention to proving (6.44).
If ε= 0 , the statement already holds.
Otherwise, note that based on Lemma 6.2 equation (6.3) we have
B α, Γ ≥ α ∥ Π
I ( r)∥ (cid:96)λ≥ α ∥ Π
I ( r)∥ (cid:96)νB ∥
X ∥ ≥ cα ∥ r∥ (cid:96)νB ∥
X ∥ ⇒
B α, Γ c
∥ r∥ (cid:96)≥ α νB ∥
X ∥ .
Recall that α = να√ K , which implies that
• If δ = ζνB ∥ X ∥ α : For (6.44) to hold it sufﬁces to have ε≤ α 5 min (
[Equation]
X ∥
[Equation]
• If δ =
[Equation]
2 : For (6.44) to hold it sufﬁces to have ε
[Equation]
c
[Equation]
r∥ (cid:96)∥
X ∥ as based on (6.57) we have √
[Equation]
⎠
Combining the latter two cases as long as
[Equation]
then (6.44) holds.
As a result when (6.51) and (6.52) hold with δ
[Equation]
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 δ Γ α ∥
[Equation]
(6.54) follows immediately from (6.47), and (6.55) follows from (6.48) by noting that based on our choice of δ we have
[Equation]
6.4. M AIN GENERALIZATION RESULT ( COMPLETING THE PROOF OF T HEOREM 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 ∥
[Equation]
2 rather than
[Equation]
Theorem 6.21 Let ζ, Γ , ¯ α be scalars obeying ζ ≤ / 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 )}
[Equation]
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 ∣
[Equation]
B .
Let W 0 be an arbitrary initial weight matrix.
Also assume the output matrix has bounded entries obeying ∥
[Equation]
Furthermore, set J ∶= J (
W) and deﬁne 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 α =
[Equation]
√ n √∥
X ∥ .
Also deﬁne the initial residual r= f (
W) − y ∈
R nK and pick C r > so that
[Equation]
≤
C r .
Suppose number of hidden nodes k obeys k ≳
[Equation]
¯
[Equation]
2 , (6.58) with Γ ≥ 1 and tolerance level ζ .
Run gradient descent updates (1.5) with learning rate η ≤ 1 ν 2 B∥
X ∥ 2 .
Then, after T =
Γ ηα 2 iterations, training loss obeys ∥
f (
W T )
− y ∥ (cid:96)≤ ∥ Π
N ( r)∥ (cid:96)+
C r ( e − Γ + ζ )√ n. and with probability at least− δ , the generalization error obeys
Err
D (
W T ) ≤ ∥ Π
N ( r)∥ (cid:96)√ n ·„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„‚„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„¶ bias term + 12 νB √ n (∥
J †
I r∥ (cid:96)+ Γ α ∥ Π
N ( r)∥ (cid:96)) ·„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„‚„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„¶
[Table]
+ 5 √ log ( / δ ) n + 2 C r ( e − Γ + ζ ) .
Theorem 6.21 immediately follows from Theorem 6.22 below by upper bounding
D α, Γ (see Deﬁnition 6.1) using Lemma 6.2 equation (6.2).
Theorem 6.22 Consider a training data set
[Equation]
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
[Equation]
B .
Let W 0 be an arbitrary initial weight matrix.
Also assume the output matrix has bounded entries obeying
[Equation]
Furthermore, set J ∶= J (
W) and deﬁne 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 ¯ α√ n √∥
X ∥ .
Also deﬁne the initial residual r= f (
W
[Equation]
) − y ∈
R nK and pick C r >
so that
[Equation]
≤
C r .
Also assume, the number of hidden nodes k obeys
[Equation]
.59) with Γ ≥ 1 and tolerance level ζ ≤ 2 .
Run gradient descent updates (1.5) with learning rate η ≤
[Equation]
B∥
X ∥ 2 .
Then, after T =
Γ ηα 2 iterations, with probability at least− δ , the generalization error obeys
Err
[Equation]
where
D α, Γ is the early stopping distance as in Def.
[Equation]
Proof First, note that using α
[Equation]
per (6.2), and
[Equation]
n we have
[Equation]
This together with ζ ≤ 2 implies that
[Equation]
n
[Equation]
Thus when k ≥
[Equation] [Equation]
√
[Equation] [Equation]
[Equation] [Equation]
[Equation]
Thus, (6.52) holds.
Also (6.51) trivially holds for ε 0 .
Thus applying Theorem 6.20 with ε=
the following three conclusions hold
[Equation]
(6.62) and
[Equation]
(6.63) and ∥
[Equation]
Furthermore, using the assumption that
[Equation]
Lemma 6.18 applies and hence equation (6.43) with W
[Equation]
D α, Γ , M
V = 0 , and R = √ nC r νB Γ ∥
X ∥
α 2 implies that
Err
[Equation]
Also note that using (6.59) we have
[Equation]
Plugging (6.64) and (6.66) into (6.65) completes the proof.
6. P ROOFS FOR NEURAL NETWORK WITH RANDOM INITIALIZATION ( PROOF OF T HEOREM 2.2)
In this section we prove Theorem 2.2.
We ﬁrst 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. O PTIMIZATION RESULT
Theorem 6.23 (Optimization guarantee for random initialization) Consider a training data set {( x i , y i )} ni = ∈
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 ∥
[Equation]
= √ 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 ,) entries.
Fix a precision level ζ and set
[Equation]
Also assume the output layer V has i.i.d.
Rademacher entries scaled by ν √ kK .
Furthermore, set J ∶= Σ
[Equation]
X ) / 2 and deﬁne 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 α= ¯ α
[Equation]
∥ per Deﬁnition 1.1 so as to ensure ∥ Π
I ( y )∥ (cid:96)≥ c ∥ y ∥ (cid:96) 2 for some constant c .
Assume
[Equation] [Equation]
with Γ ≥ 1 and ζ ≤ c 2 .
We run gradient descent iterations of the form (1.5) with a learning rate η ≤
[Equation]
ν
[Equation]
∥
X ∥ 2 .
Then, after T =
Γ K ην 2 α 20 iterations, the following identities
[Equation] [Equation]
hold with probability at least− ( 2 K ) − 100 .
Proof To prove this result we wish to apply Theorem 6.20.
To do this we need to verify the assumptions of this theorem.
To start with, using Lemma 6.14 with probability at least− ( 2 K ) − 100 , the initial prediction vector f (
[Equation]
obeys
[Equation]
Hence the initial residual obeys
[Equation]
√ n .
Furthermore, using ζ ≤ c
[Equation]
Thus, ∥
[Equation] [Equation]
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 (cid:96)norm hence we also have
[Equation]
Next we wish to show that (6.51) holds.
In particular we will show that there exists an ε 0reference Jacobian J for
J (
W) satisfying JJ T =
E [J (
W
[Equation]
)J (
W
[Equation]
T ] .
Note that, such a J will have exactly same information/nuisance spaces as the square-root of the multiclass kernel matrix i.e.
[Equation]
12 since these subspaces are governed by the left eigenvectors.
Applying
Lemmas 6.13 (with a scaling of the Jacobian by/√ kK due to the different scaling of V ), we ﬁnd that if
[Equation]
then, ∥J
[Equation]
Let
J (
W ) be obtained by adding max (
Kn − p,) zero columns to
J (
W ) .
Then, using (6.78) and
Lemma 6.4, there exists J satisfying JJ T
[Equation]
and
[Equation]
Therefore, J is an ε= 4 δν 2 K reference Jacobian.
Now set
[Equation]
⎠ and note
that using α = ν √ K α 0 and ∥ r
[Equation] [Equation]
To continue further, note that
B α 0 , Γ calculated with respect to Σ (
X )
[Equation]
2 with cutoff α 0 is exactly same as
B α, Γ calculated with respect to J with cutoff α = να√ K which is a square-root of
E [J (
[Equation]
T ] .
Thus, using (6.79) to ensure (6.51) holds it sufﬁces to show
[Equation]
Hence, to ensure (6.51) holds we need to ensure that δ obeys
[Equation]
α200Θ .
Thus using δ = α 20 200 Θ to ensure (6.51) we need to make sure k is sufﬁciently large so that (6.77) holds with this value of δ .
Thus it sufﬁces to have
[Equation] [Equation]
≥
[Equation]
⎛ ⎝
[Equation]
Here, (a) follows from the fact that
[Equation]
, equation (6.3)
[Equation]
and ∥
[Equation]
y
[Equation]
n which combined imply
[Equation]
To be able to apply Theorem 6.20 we must also ensure (6.52) holds.
Therefore, it sufﬁces to have
[Equation]
[Equation] [Equation]
≥
[Equation]
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
[Equation]
Note that (6.80) and (6.84) are implied by
[Equation]
which is the same as (6.68).
What remains is stating the optimization bounds in terms of the labels y .
This follows by substituting (6.72), (6.76), and the fact that ∥ r∥
[Equation]
≤ √ n into (6.55), (6.53), and (6.54), respectively.
6.5. G ENERALIZATION RESULT ( COMPLETING THE PROOF OF T HEOREM 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
[Equation]
rather than
[Equation]
Theorem 6.24 Let ζ, Γ , ¯ α be scalars obeying ζ ≤ / 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 )} ni
= ∈ 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 ∥
[Equation]
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 ∣
[Equation]
z )∣
[Equation]
B .
Let W 0 be the initial weight matrix with i.i.d.
N( 0 ,) entries.
Fix a precision level ζ and set ν = ζ /
[Equation]
√
[Equation]
Also assume the output layer V has i.i.d.
Rademacher entries scaled by ν
[Equation]
kK .
Furthermore, set J ∶= (
[Equation]
X
[Equation]
) / 2 and deﬁne 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 α
[Equation]
per Deﬁnition 1.1.
Assume
[Equation]
with Γ ≥ 1 .
We run gradient descent iterations of the form (1.5) with a learning rate η ≤ 1 ν 2 B∥
X ∥ 2 .
Then, after T =
Γ K ην 2 α 20 iterations, training loss obeys (6.69) and classiﬁcation error Err
D (
W T ) is upper bounded by ∥ Π
N ( y )∥ (cid:96)√ n ·„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„‚„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„¶
[Table]
+ 12 B √
K √ n (∥
J †
I y ∥ (cid:96)+ Γ α∥ Π
N ( y )∥ (cid:96)) ·„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„‚„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„¶ [Table] +
[Equation]
holds with probability at least
[Equation]
The theorem below is a restatement of Theorem 6.24 after substituting the upper bound on the early stopping distance
[Equation]
Theorem 6.25 (Neural Net – Generalization) Consider a training data set {( x i , y i )} ni
[Equation]
∈
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
[Equation]
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 )∣ , ∣ φ
[Equation]
( z )∣ ≤
B .
Let W 0 be the initial weight matrix with i.i.d.
N( 0 ,) entries.
Fix a precision level ζ ≤ c 2 and set
[Equation]
Also assume the output layer V has i.i.d.
Rademacher entries scaled by ν √ kK .
Furthermore, set
J ∶= Σ (
X ) / 2 and deﬁne 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 α= ¯ α√ n √
K ∥
X ∥
B ≤
B ∥
X ∥ per Deﬁnition 1.1 chosen to ensure
[Equation]
≥ c ∥ y ∥ (cid:96) 2 for some constant c > 0 .
Assume
[Equation]
6.89) with Γ ≥ 1 .
We run gradient descent iterations of the form (1.5) with a learning rate η ≤ 1 ν 2 B∥
X ∥ 2 .
Then, after T =
[Equation]
, classiﬁcation error Err
D (
W T ) is upper bounded by
Err
[Equation]
holds with probability at least− (
[Equation]
K
[Equation]
− − δ .
Proof Under the stated assumptions, Theorem 6.23 holds with probability− ( 2 K ) − 100 .
The proof will condition on outcomes of the Theorem 6.23.
Speciﬁcally, 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 ﬁxed.
Observe that ∥
V ∥
F
[Equation]
Kk ∥
V ∥ (cid:96) ∞ =
ν , we have
Err
[Equation]
⎤⎥⎥⎥⎥⎦ .
[Equation]
Theorem 6.23 yields
[Equation]
Using (6.70) for M
[Equation]
3 νB R√ k = 48 n Γ 2 B 3 K∥
X ∥ να√ k ≤ c 2 ζ230 νB log ( n ) ≤ ζ.
Plugging in (6.91), (6.92), and (6.93) into (6.90) concludes the proof.
A T HE J ACOBIAN OF THE M IXTURE M ODEL IS LOW - RANK (P ROOFS FOR S ECTION 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 )} ni = ∈
R d ×
R K generated according to the Gaussian mixture model of Deﬁnition 2.4 with K classes each consisting of
C clusters with the cluster centers given by { µ (cid:96), ̃ (cid:96) } (
K,C ) ( (cid:96), ̃ (cid:96) )=( 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 ν = √ k .
Also deﬁne 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 ) .
Deﬁne the neural tangent kernel matrix associated with the cluster centers as
[Equation]
and assume that ̃
[Equation]
) is full rank.
Then, the following properties hold with probability−
KC exp (− n 8 KC ) •
I is a K 2 C dimensional subspace.
• The concatenated label vector
y =
[Equation]
y
Tn
T lies on I .
• The nonzero eigenvalues (top K 2 C eigenvalues) of Σ ( X ) are between n 2 KC s min (
Σ (
X )) and 2 nKC ∥ Σ ( X )∥ .
Hence the eigenvalues of the information space grow with nKC .
Proof First, we establish that each cluster has around the same size.
Applying Chernoff bound and a union bound, we ﬁnd that with probability
[Equation]
Note that based on Lemma 6.11, the multiclass covariance is given by
[Equation]
where ̃ Σ
[Equation]
) =
[Equation]
[ φ
[Equation]
Xg ) φ ′ (
Xg )
T ] .
Due to this Kronecker product representation, the range space of Σ (
X ) is separable.
In particular, note that with ̃I =
Range (̃ Σ (
X ))
we have
I =
I K ⊗ ̃I which also implies rank (I) =
K ⋅ rank (̃I) .
Hence, this identity allows us to reduce the problem to a single output network.
To complete the proof we will prove the following three identities:
• ̃ I has rank KC .
• The nonzero eigenvalues of
̃ Σ (
X
are between 0 .
̃ ns min (̃ Σ (
M ))
to
̃ n ∥̃ Σ (
M )∥
• The portion of the label vector associated with class (cid:96) i.e. y
( (cid:96) ) ∈
R n (see (5.4)) lies on I .
Hence, the
concatenated vector y lies on
I =
I K ⊗ ̃I .
To prove these statements let
J (cid:96) (
X ; W) and
J (cid:96) (
M ; W) be the Jacobian associated with the (cid:96) th output of the neural net (see (5.4)) for data matrices X and M .
Observe that the columns of
J (cid:96) (
X ; W) are chosen from
J (cid:96) (
M ; W) and in particular each column of
J (cid:96) (
M ; W) is repeated between 0 .̃ n tõ n times.
To mathematically relate this, deﬁne the KC dimensional subspace
S of
R n where for any v ∈ S , entries v i and v j of v are equal iff data point x i and x j are equal (i.e. belong to the same class/cluster pair).
Now, we deﬁne the orthonormal matrix U
S ∈
R n ×
KC as the 0-1 matrix with orthogonal rows that map R KC
to
S as follows.
Assume the i th data point x i belongs to the class/cluster pair ( (cid:96) i , ̃ (cid:96) i ) .
We then set the i th row of U
S as vect ( e (cid:96) i e T ̃ (cid:96) i ) .
Using U
S we have
U
[Equation]
Now note that using the above identity we have
[Equation]
Since U
S is tall and orthogonal, the range of ̃ Σ (
X ) is exactly the range of U
S hence ̃I
=
S which is KC dimensional.
Furthermore, nonzero eigenvectors of ̃ Σ (
X ) lie on
S and any eigenvector v satisﬁes
[Equation]
and similarly
[Equation]
∥ which follows from the fact that (cid:96) 2 -norm-squared of columns of U are between 0 .
5¯ n to 2¯ n .
Finally, we will argue that label vector y ( (cid:96) ) lies on
S .
Note that for all samples i that belong to the same cluster y (
[Equation]
) i will be the same (either zero or one), thus y
[Equation]
Next lemma provides a perturbation analysis when there is noise.
Lemma A.2 Consider the single-output NTK kernel given by
[Equation]
and assume that this matrix has rank r so that λ r +
[Equation]
X ))
= 0 .
Also assume a noise corrupted version of X given by ̃
X =
X + σ √ d Z with Z a matrix consisting of i.i.d.
N( 0 ,) entries.
Then, ∥̃ Σ ( ̃
X ) − ̃ Σ (
X )∥ ≲ ∆ where ∆
[Equation]
log n ∥
[Equation]
X
[Equation]
(A.1) holds with probability at least− 2 n e − d 2 .
Whenever σ ≤ √ log n , ∆ is upper bounded as
[Equation] [Equation]
Furthermore, let ̃
V , V ∈
R n × r be orthonormal matrices corresponding to the top r eigenvalues of ̃ Σ ( ̃
X ) and ̃ Σ (
X ) .
Then, ∥ ̃
V ̃
V T −
V V T ∥ ≤ ∆ λ r (̃ Σ (
X ))
− ∆
Proof Note that diag ( φ ′ ( ̃
Xw ))
̃
X − diag ( φ ′ (
Xw ))
X = diag ( φ ′ ( ̃
Xw ))
̃
X − diag ( φ ′ (
Xw ))
X = diag ( φ ′ ( ̃
Xw ) − φ
Xw ))
X + diag ( φ ′ ( ̃
Xw ))
( ̃
X −
X )
Now deﬁne ̃
M = diag ( φ ′ ( ̃
Xw ))
̃
X and M = diag ( φ ′ (
Xw ))
X and note that using the above we can conclude that
[Equation]
∥ Now using the fact that
[Equation]
we conclude that
[Equation]
w
[Equation]
X
[Equation]
X
[Equation]
X ∥
To proceed further, with probability− n exp (− d / ) , each row of ̃
X −
X is upper bounded by 2 σ .
Hence, using a standard tail bound over supremum of n Gaussian random variables (which follows by union bounding) we have
E
[Equation]
holds with the same probability.
Furthermore, spectral norm bound on Gaussian random matrix implies that
[Equation]
holds with probability at least− e
[Equation] [Equation]
Plugging these two probabilistic bounds into the chain of inequalities we conclude that ∥̃
[Equation]
To establish (A.2), observe that B 2 σ √ log n ∥
X ∥ 2 dominates over other terms in the regime σ √ log n is small.
The ﬁnal bound is a standard application of Davis-Kahan Theorem Yu et al.
(2014) when we use the fact that ˜ Σ (
X ) is low-rank.
The following lemma plugs in the critical quantities of Theorem 2.2 for our mixture model to obtain a generalization bound.
Theorem A.3 (Generalization for Mixture Model) Consider a dataset { x i , y i } ni
= 1 generated i.i.d.
from the Gaussian mixture model in Deﬁnition 2.4.
Let λ M = λ min ( Σ (
M ))
where M ∈
R KC ×
d is the matrix of cluster centers.
Suppose input noise level σ obeys σ ≲
[Equation]
min
[Equation]
log n
Consider the setup of Theorem 2.2 with quantities ζ and Γ .
Suppose network width obeys
[Equation]
4min .
With probability
[Equation]
, running gradient descent for T = 2Γ K 2 C ην 2 nλ min with learning rate η
[Equation]
2 , we have that
Err
[Equation]
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
[Equation]
for some constant c > 0 .
Next we quantify Π
I ( y ) using the fact that (i) Π
I ′ ( y
[Equation]
= y via Theorem A. as follows
[Equation]
In return, this implies that
[Equation]
To proceed, we pick α= √
[Equation]
KC and corresponding ¯ α
[Equation]
B 2 K 2 C and apply ( ?? )
to ﬁnd that, classiﬁcation error is upper bounded by
Err
[Equation]
B J OINT 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
[Equation] [Equation]
Denoting the Jacobian associated with (B.1) by
J
[Equation]
, W ) we have that
[Equation]
Here,
J (
W ) is as before whereas
J (
V ) is the Jacobian with respect to V and is given by
[Equation]
where
J (
[Equation]
∈
R Kn × k is so that its (cid:96) ’
th block rows of size n × k is nonzero for≤ (cid:96) ≤
K i.e.
̃ (cid:96) th block row of
J ( v (cid:96) )
=
{ 0 if (cid:96) ≠ ̃ (cid:96) φ (
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 ) satisﬁes the following properties.
• Upper bound:
[Equation]
• Row-bound:
For unit length u :
[Equation]
) u )∥
[Equation]
B √
K ∥
V ∥
[Equation]
∥ .
• Entry-bound: For unit length u :
∥ mat (J
T
[Equation]
V
[Equation]
u
[Equation]
∥
[Equation]
B ∥
W ∥
[Equation]
X ∥ .
• Lipschitzness:
Given inputs V , V ′ and outputs W , W
′ ∥J
[Equation]
W
[Equation]
[Equation] [Equation]
B
[Equation]
K ∥
V ∥
[Equation]
W ′ −
W ∥
F +∥
W −
W ′ ∥
F ) .
Proof First, we prove results concerning
J (
V ) .
First, note that
[Equation]
Next, note that for u
= [ u 1 ... u K
[Equation]
R Kn we have
[Equation]
Let
J 1 ,
J 2 be the Jacobian matrices restricted to V and W of
J (
V , W ) .
To prove Lipschitzness, ﬁrst observe that ∥J
[Equation]
Next, observe that ∥J
[Equation]
We decompose
J 2 via ∥J
[Equation]
− J (
V ′ , W ′ )
∥ ≤
B √
K ∥
V ∥ (cid:96) ∞
∥
X ∥∥
W ′ −
W ∥
F + ∥J (
V , W ′ )
− J (
V ′ , W ′ )
∥ .
To address the second term, note that, Jacobian is linear with respect to output layer hence ∥J
[Equation] [Equation]
W
[Equation]
B √
Kk ∥
V −
V ′ ∥ (cid:96) ∞
∥
X ∥ .
Combining the latter two identities we arrive at ∥J
[Equation]
V ∥
[Equation]
W ′ −
W ∥
F ) , completing the proof.
C F URTHER NUMERICAL EXPERIMENTS
We depict the approximate eigenstructure in Section C.1.
Complementary to Section 4, in Appendix
C.2 we provide the complete set of experiments we performed on the original 10-class CIFARdataset covering various levels of label corruption.
In Appendix C.3 we show numerical results on a sub-sampled 3-class version of CIFAR-10.
Moreover, we demonstrate that our theory holds across different datasets by providing experimental results on MNIST in Appendix C.4.
C. E XPERIMENTS ON FULL (10- CLASS ) CIFAR-
As mentioned earlier while calculating all the eigenvalues is not possible, we also verify the bimodal structure of the Jacobian using the entire spectrum by approximating its spectral density using the method recently described in Ghorbani et al. (2019a) and Papyan (2019a) for Hessian eigenvalue density estimation.
We detail the parameters of the algorithm in Appendix C.5.
Figure 9 depicts the estimated spectrum before and after training.
We observe a similar bimodal spectrum with a few outliers.
Moreover, in both depictions of the Jacobian eigenstructure we observe increasing separation between the low and high end of the spectrum after training.
C. E XPERIMENTS ON FULL (10- CLASS ) CIFAR-
In addition to what has been described in Section 4 here we disclose the complete set of experiments performed on the CIFAR-10 dataset.
We trained the modiﬁed ResNet20 model described in Sectionwith SGD as longs as it was necessary to achieve a good ﬁt to the training data.
Information subspace
[Figure]
S c a l e d d e n s i t y
[Figure]
[Figure]
J traininit ∗ 0.
[Figure] [Figure]
[Figure]
J trainfinal ∗ 0.
J testinit 0.
J testfinal 0.
[Table]
is spanned by the top 50 singular vectors.
We marked ﬁgures and table entries also included in
Section 4 by asterisk.
Experiments without label corruption.
We trained the modiﬁed ResNet20 model described in
Section 4 with SGD for 400 epochs with learning rate 0 .
1 on the original dataset without any form of data augmentation.
The network output after the last layer has been scaled by s = 0 .
025 in all experiments to follow.
Experiments with label corruption.
We corrupt the labels in the training data by switching each label to a strictly different (incorrect) class.
We train the network on the corrupted dataset forepochs with initial step size 0 .
1 decayed to 0 .
01 after 760 epochs.
[Table]
of the alignment of the initial label/residual with the information/nuisance space using 25% label corruption on CIFAR-10 data.
[Figure]
[Figure] [Figure]
y )
∥ (cid:96)∥ y ∥ (cid:96)
[Figure]
∥ Π
N ( y )∥ (cid:96)
[Figure]
∥ y ∥ (cid:96)
[Table]
[Figure]
∥ Π
I ( y )∥ (cid:96)∥ y ∥ (cid:96)∥ Π
N ( y )∥
(cid:96)∥ y ∥ (cid:96)∥
J †
I y ∥ (cid:96)∥ y ∥ (cid:96)∥ Π I ( r)∥ (cid:96)∥ r∥ (cid:96)∥ Π
N ( r)∥ (cid:96)∥ r∥ (cid:96)∥
J †
I r∥ (cid:96)∥ r∥ (cid:96) 2 J traininit 0.
J trainfinal
[Table]
[Figure] [Figure]
[Figure] [Figure] [Figure]
[Figure] [Figure]
∗
∗ : we discuss in
[Figure] [Figure]
∥ Π
I ( y )∥ (cid:96)∥ y ∥ (cid:96)∥ Π
N ( y )∥
(cid:96)∥ y ∥ (cid:96)∥
J †
I y ∥ (cid:96)∥ y ∥ (cid:96)∥ Π I ( r)∥ (cid:96)∥ r∥ (cid:96)∥ Π
N ( r)∥ (cid:96)∥ r∥ (cid:96)∥
J †
I r∥ (cid:96)∥ r∥ (cid:96) 2 J traininit 0.
J trainfinal 0.
[Figure]
experiments.
See discussion of this plot in Section
C. E XPERIMENTS ON SUBSAMPLED 3- CLASS CIFAR-
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 ﬂip) to increase generalization.
Information subspace is spanned by the top 50 singular vectors.
Experiments without label corruption.
First we trained the network on the sub-sampled 3 -class dataset keeping the original labels.
For the SGD experiments, we used initial learning rate 0 .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 sufﬁcient for a good ﬁt to the training data.
We observed better Jacobian adaptation using Adam compared to SGD on this dataset ( 0 .from Table 8 vs. 0 .
9969 from Table 9).
Experiments with label corruption.
We corrupt the training labels by the corruption model described in C.2.
We train the network on the corrupted dataset for 800 epochs with initial step size 0 .
01 and batch size 128 .
We decrease the learning rate by a factor of 10 at the following epochs: at 500 epochs for 25% corruption, at 700 epochs for 50% and 75% corruption and at 500 and 700 for 100% corruption.
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
∥ Π
I ( y )∥ (cid:96)∥ y ∥ (cid:96)∥ Π
N ( y )∥
(cid:96)∥ y ∥ (cid:96)∥
J †
I y ∥ (cid:96)∥ y ∥ (cid:96)∥ Π I ( r)∥ (cid:96)∥ r∥ (cid:96)∥ Π
N ( r)∥ (cid:96)∥ r∥ (cid:96)∥
J †
I r∥ (cid:96)∥ r∥ (cid:96) 2 J traininit 0.
J train 200 epochs 0.
J trainfinal 0.
J test init 0.
J testfinal 0.
[Figure] [Figure]
J traininit 0.
J trainfinal 0.
[Figure]
i du a l e n e r
g
y
[Figure] [Figure]
Π
I ( y )∥ (cid:96)∥ y ∥ (cid:96)∥ Π
N ( y )∥
(cid:96)∥ y ∥ (cid:96)∥
J †
I y ∥ (cid:96)∥ y ∥ (cid:96)∥ Π I ( r)∥ (cid:96)∥ r∥ (cid:96)∥ Π
N ( r)∥ (cid:96)∥ r∥ (cid:96)∥
J †
I r∥ (cid:96)∥ r∥ (cid:96) 2 J train init 0.
J trainfinal 0.
[Table]
[Figure] [Figure]
[Figure]
a l
[Figure]
n
[Figure] [Figure]
[Figure] [Figure]
∥ Π
I ( y )∥ (cid:96)∥ y ∥ (cid:96)∥ Π
N ( y )∥
(cid:96)∥ y ∥ (cid:96)∥
J †
I y ∥ (cid:96)∥ y ∥ (cid:96)∥ Π I ( r)∥ (cid:96)∥ r∥ (cid:96)∥ Π
N ( r)∥ (cid:96)∥ r∥ (cid:96)∥
J †
I r∥ (cid:96)∥ r∥ (cid:96) 2 J traininit 0.
J traindip 0.
J trainfinal 0.
J testinit 0.
J testdip 0. J testfinal 0.
[Figure]
a l
[Figure]
n e r g y
[Figure]
∥ Π
[Figure]
( y )∥
(cid:96)∥ y ∥ (cid:96)
[Figure]
∥ Π
N ( y )∥ (cid:96)
[Figure]
∥ y ∥ (cid:96)
[Figure]
∥
J
[Figure]
I y ∥ (cid:96)
[Figure]
∥ y ∥
[Figure]
∥ Π
[Figure] [Figure]
)∥ (cid:96)
[Figure]
∥
[Figure]
∥
[Figure]
∥
[Figure]
∥ (cid:96)
[Figure]
∥
[Figure]
∥ (cid:96)
[Figure]
J traininit 0.
J trainfinal 0.
[Figure] [Figure]
J traininit 0.
[Figure]
J trainfinal 0.
corruption on 3-class CIFAR-10 data.
C. E XPERIMENTS ON MNIST
The MNIST dataset contains handwritten digits in 10 classes divided into 60k training and 10k test images.
To demonstrate our theoretical ﬁndings 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 modiﬁed ResNet20 model from Section 4.
In all of the following MNIST experiments the model output after the last layer has been scaled by 0 .
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.
∥ Π
I ( y )∥ (cid:96)∥ y ∥ (cid:96)∥ Π
N ( y )∥
(cid:96)∥ y ∥ (cid:96)∥
J †
I y ∥ (cid:96)∥ y ∥ (cid:96)∥ Π I ( r)∥ (cid:96)∥ r∥ (cid:96)∥ Π
N ( r)∥ (cid:96)∥ r∥ (cid:96)∥
J †
I r∥ (cid:96)∥ r∥ (cid:96) 2 J traininit 0.
J train final 0.
[Table]
[Figure] [Figure]
[Figure]
J traininit 0.
[Figure]
[Figure]
[Figure]
J trainfinal 0.
[Table] [Table]
[Figure]
MNIST data ∥
Π
I ( y )∥ (cid:96)∥ y ∥ (cid:96)∥ Π
N ( y )∥
(cid:96)∥ y ∥ (cid:96)∥
J †
I y ∥ (cid:96)∥ y ∥ (cid:96)∥ Π I ( r)∥ (cid:96)∥ r∥ (cid:96)∥ Π
N ( r)∥ (cid:96)∥ r∥ (cid:96)∥
J †
I r∥ (cid:96)∥ r∥ (cid:96)
[Table]
[Table] [Table]
[Table] [Table]
[Figure] [Figure]
∥ Π
I ( y )∥ (cid:96)∥ y ∥ (cid:96)∥ Π
N ( y )∥
(cid:96)∥ y ∥ (cid:96)∥
J †
I y ∥ (cid:96)∥ y ∥ (cid:96)∥ Π I ( r)∥ (cid:96)∥ r∥ (cid:96)∥ Π
N ( r)∥ (cid:96)∥ r∥ (cid:96)∥
J †
I r∥ (cid:96)∥ r∥ (cid:96) 2 J traininit 0.
J trainfinal 0.
[Figure] [Figure]
[Figure]
C. S INGULAR VALUE DECOMPOSITION OF THE J ACOBIAN
The matrix representation of the Jacobian arising in the numerical experiments is in general very large and pose a computational challenge.
In particular, the Jacobian matrix for a ResNet20 model is 500 k × 270 k
or 540 GB memory with ﬂoat32 entries, more than 1 TB with ﬂoat64) on the CIFARtraining data.
Storing these matrices in memory for singular value decomposition is not feasible and therefore other methods are necessary to approximate the singular values and singular vectors.
The Lanczos algorithm Lanczos (1950) is an adaptation of power methods, which guarantees fast convergence to the largest eigenvalues of a hermitian matrix.
More importantly, it only requires matrix-vector products to estimate the top eigenvectors and therefore resolves the memory issues with large Jacobians.
Using the ARPACK implementation of the implicitly restarted Lanczos method
Lehoucq et al.
(1998) it takes about 1 hour on 64 cores to ﬁnd the top 50 eigenvalues and eigenvectors of the 500 k × 500 k symmetric matrix JJ T for CIFAR-10 training data.
The same algorithm has been used to calculate the histogram of top 1000 singular values.
The required computation time for Lanczos method is dominated by large Jacobian-vector products, which increases linearly with the number of calculated eigenvalues.
Therefore it is necessary to depart from deterministic power iteration type algorithms and take a different approach.
Stochastic
Lanczos quadrature algorithm has been recently introduced to the machine learning community in
Ghorbani et al.
(2019a)
and Papyan (2019a) to approximate the eigenvalue density of large neural network Hessians.
Finding the exact spectral density might be difﬁcult therefore ﬁrst the problem is relaxed by approximating the spectrum convolved with a Gaussian density of variance σ= − 5 .
As described in Ghorbani et al. (2019a), we draw k = 10 i.i.d.
Gaussian random vectors and run m = steps of the Lanczos algorithm starting from each random vector to estimate the parameters of the spectral density.
We modiﬁed the algorithm by replacing Hessian-vector products by JJ T -vector products.
Finally, we normalize the density and scale it by the dimension of the Jacobian output space.
D J ACOBIAN A DAPTATION FOR L INEAR D ATA
In this section, we will (non-rigorously) argue that Jacobian and its singular value spectrum indeed adapts to the data for simple linear datasets.
Speciﬁcally, we will show that the very ﬁrst gradient iteration from a random initialization is already powerful enough to encourage an approximately low-rank Jacobian.
Dataset: Let ( x i ) ni
= 1 i.i.d.
∼ N( 0 , I d ) and ﬁx θ ∈
R d .
Fix labels y
i
=
θ
T x i and the dataset ( x i ,y i ) ni
= 1 .
Neural net: Let W∈
R k ×
d i.i.d.
∼ N( 0 ,) and ﬁx v ∈
R k with half + σ and half − σ entries.
Set activation to be ReLU.
(Non-rigorous) Statement: With small initialization ( σ ≈ 0 ) and a single large gradient step,
W=
W− η ∇L(
W) has a (properly scaled)
Jacobian matrix with approximately rank 2 d .
Instead, at random initialization W 0 , Jacobian has full rank n (assuming kd ≫ n ).
Speciﬁcally, at W 1 , we have σ 2 d + (J (
X , W))
≲ √ d / n and σ 2 d (J (
X , W)) ≳ 1 .
(D.1) where σ i is the i th largest singular value.
In words,
J (
X , W) is approximately rank 2 d when n ≳ d .
This is in stark contrast to the properties of Jacobian at random initialization.
For instance Li &
Liang (2018); Zou et al.
; Oymak & Soltanolkotabi (2019) actually show that Jacobian is well conditioned with minimum singular value lower bounded by the minimum separation between the samples (implying full rank of n ).
Such separation (independent of n ) holds for Gaussian data even when n is polynomially larger than d via standard concentration / packing bounds (also see Oymak &
Soltanolkotabi (2019); Li & Liang (2018)).
Hence, as further formalized below, we argue that even a single step of gradient descent can lead to an approximately low-rank Jacobian.
D. A PPROACH
Fixing weights of W 0 , we study the gradient of the random data.
Let w r be r th row of W 0 .
Fixing hidden node r , let x ′ i =
[Equation]
i
if x Ti w r ≥
(i.e. the r th node is active) and 0 else.
The gradient of the node r at W 0 is given by
[Equation]
If initialization is small enough ( σ ≈ 0 ), we can approximate f
[Equation]
y i hence we study
[Equation]
We will show that gradients of each row are ± σ θ / 2 in expectation.
Decompose a sample x ∼ N( 0 , I d ) as
[Equation]
where g is standard normal vector over the orthogonal complement of vector w r and g ∼ N( 0 ,) .
Denoting label associated to x by y , consequently, we ﬁnd
E [
[Equation]
D.
[Equation]
g w r + g )
T θ ∣ g ≥ ] (D.3) =
[Equation]
E [
[Equation]
T ] θ +
E [ g
[Equation]
∣ g ≥ ] w r
[Equation]
T r θ (D.4)
=
θ 2 .
(D.5)
Hence
E [
[Equation]
W
[Equation] [Equation]
This implies that, in population ( n → ∞ ), gradient vector has an extremely simple rank-one form:
[Equation]
Standard results on non-asymptotic statistics then yield that with exponentially high-probability (i.e.
[Equation]
As long as k ≲ poly ( n ) , union bounding over all rows, we ﬁnd the spectral norm bound
[Equation]
On the other hand, E X [∇L
[Equation]
] is rank 1 with nonzero singular value σ √
[Equation]
Consequently, a large gradient step will ensure that
[Equation]
is approximately rank one as well.
Speciﬁcally, set the rank one matrix ¯
[Equation]
Then the tail of the spectrum is at most
[Equation]
.e.
[Equation]
n as large as the spectral norm).
Finally, we need to move from showing low-rankness of the weight matrix to low-rankness of
Jacobian.
Deﬁne the Jacobians associated with data at ¯ W 1 and W 1 by
[Equation]
The Jacobian associated with ¯
J is exactly rank
d
since matrix diag ( v ) φ ′ ( ¯ W 1 X T ) has exactly two distinct rows and kronecker producting with X T will yield d rank for each (adding up to 2 d ).
Since
X is well-conditioned with high probability (as it is
N( 0 ,) and fat with n ≳ d ), top 2 d singular values and σ 2 d ( ¯
J ) are strictly positive with a lower bound independent of n and σ (after scaling by / σ ∥
X ∥ ).
Finally, using smoothness of φ ′ (which excludes ReLU but includes smoother versions of ReLU such as softplus), using Jacobian perturbation bound (6.31), we have (after scaling by
[Equation]
Together, this leads to the desired conclusion (D.1) since singular value perturbation inequality yields σ
[Equation]