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

1 INTRODUCTION . 1.1 MOTIVATION AND CONTRIBUTIONS . Deep neural networks ( DNN ) are ubiquitous in a growing number of domains ranging from computer vision to healthcare . State-of-the-art DNN models are typically overparameterized and contain more parameters than the size of the training dataset . It is well understood that in this overparameterized regime , DNNs are highly expressive and have the capacity to ( over ) fit arbitrary training datasets including pure noise Zhang et al . ( 2016 ) . Mysteriously however neural network models trained via simple algorithms such as ( stochastic ) gradient descent continue to predict well or generalize on yet unseen test data . In this paper we wish to take a step towards demystifying this phenomenon and help explain why neural nets can overfit to noise yet have the ability to generalize when real data sets are used for training . In particular we explore the generalization dynamics of neural nets trained via gradient descent . Using the Jacobian mapping associated with the neural network we characterize directions where learning is fast and generalizable versus directions where learning is slow and leads to overfitting . The main contributions of this work are as follows.● Leveraging dataset structure : We develop new optimization and generalization results that can harness the low-rank representation of semantically meaningful datasets via the Jacobian mapping of the neural net . This sheds light as to why training and generalization is easier using datasets where the features and labels are semantically linked versus others where there is no meaningful relationship between the features and labels ( even when the same network is used for training ) .● Bias–variance tradeoffs : We develop a generalization theory based on the Jacobian which decouples the learning process into information and nuisance spaces . We show that gradient descent almost perfectly interpolates the data over the information space ( incurring only a small bias ) . In contrast , optimization over the nuisance space is slow and results in overfitting due to higher variance . ● Network size vs prediction bias : We obtain data-dependent tradeoffs between the network size and prediction bias . Specifically , we show that larger networks result in smaller prediction bias , but small networks can still generalize well , especially when the dataset is sufficiently structured , but typically incur a larger bias . This compares favorably with recent literature on optimization and generalization of neural networks Jacot et al . ( 2018 ) ; Arora et al . ( 2019 ) ; Du et al . ( 2018b ) ; Allen-Zhu et al . ( 2018b ) ; Cao & Gu ( 2019 ) ; Ma et al . ( 2019 ) ; Allen-Zhu et al . ( 2018a ) ; Brutzkus et al . ( 2017 ) where guarantees only hold for very wide networks with the width of the network growing inversely proportional to the class margins or related notions . See Section 3 for further detail.● Pretrained models : Our framework does not require random initialization and our results continue to apply even with arbitrary initialization . Therefore , our results may shed light on the generalization capabilities of networks initialized with pre-trained models commonly used in meta/transfer learning . Our extensive experiments strongly suggest Jacobian adapts over time in a favorable and data-dependent fashion shedding light on the properties of ( pre ) trained models . 1.2 MODEL AND TRAINING . Our theoretical analysis will focus on neural networks consisting of one hidden layer with d input features , k hidden neurons and K outputs as depicted in Figure 1 . We use W ∈ Rk×d and V ∈ RK×k to denote the input-to-hidden and hiddento-output weights . The overall input-output relationship of the network is a function f ( ⋅ ; W ) ∶ Rd → RK that maps an input x ∈ Rd to an output via x↦ f ( x ; W ) ∶= V φ ( Wx ) . ( 1.1 ) Given a dataset consisting of n feature/label pairs ( xi , yi ) with xi ∈ Rd representing the features and yi ∈ RK the associated labels representing one of K classes with one-hot encoding ( i.e . yi ∈ { e1 , e2 , . . . , eK } where e ` ∈ RK has all zero entries except for the ` th entry which is equal to one ) . To learn this dataset , we fix the output layer and train overW via1 min W ∈Rk×d L ( W ) ∶= 12 n∑ i=1 ∥V φ ( Wxi ) − yi∥2 ` 2 . ( 1.2 ) It will be convenient to concatenate the labels and prediction vectors as follows y = ⎡⎢⎢⎢⎢⎣ y1⋮ yn ⎤⎥⎥⎥⎥⎦ ∈ R nK and f ( W ) = ⎡⎢⎢⎢⎢⎣ V f ( x1 ; W ) ⋮ V f ( xn ; W ) ⎤⎥⎥⎥⎥⎦ ∈ R nK . ( 1.3 ) Using this shorthand we can rewrite the loss ( 1.2 ) as min W ∈Rk×d L ( W ) ∶= 12 ∥f ( W ) − y∥2 ` 2 . ( 1.4 ) To optimize this loss starting from an initializationW0 we run gradient descent iterations of the form Wτ+1 =Wτ − η∇L ( Wτ ) , ( 1.5 ) with a step size η . In this paper we wish to explore the theoretical properties of the model found by such iterative updates with an emphasis on the generalization ability . 1.3 INFORMATION AND NUISANCE SPACES . In order to understand the generalization capabilities of models trained via gradient descent we need to develop better insights into the form of the gradient updates and how it affects the training dynamics . To this aim let us aggregate the weights at each iteration into one large vector wτ ∶=vect ( Wτ ) ∈ Rkd , define the misfit/residual vector r ( w ) ∶= f ( w ) − y and note that the gradient updates take the form wτ+1 =wτ − η∇L ( wτ ) where ∇L ( w ) = ∇L ( w ) = J ( w ) Tr ( w ) . 1For clarity of exposition , we focus only on optimizing the input layer . However , as shown in the supplementary , the technical approach is quite general and applies to arbitrary multiclass nonlinear least-squares problems . In particular , the proofs are stated so as to apply to one-hidden layer networks where both layers are trained . Here , J ( w ) ∈ RnK×kd denotes the Jacobian mapping associated with f defined as J ( w ) = ∂f ( w ) ∂w . Due to the form of the gradient updates the dynamics of training is dictated by the spectrum of the Jacobian matrix as well as the interaction between the residual vector and the Jacobian . If the residual vector is very well aligned with the singular vectors associated with the top singular values of J ( w ) , the gradient update significantly reduces the misfit allowing substantial reduction in the train error . Thus to provide a more precise understanding of the training dynamics and generalization capabilities of neural networks it is crucial to develop a better understanding of the interaction between the Jacobian and the misfit and label vectors . To capture these interactions we require a few definitions . Definition 1.1 ( Information & Nuisance Spaces ) Consider a matrix J ∈ RnK×p with singular value decomposition given by J = nK∑ s=1λsusv T s = Udiag ( λ1 , λ2 , . . . , λnK ) V T , with λ1 ≥ λ2 ≥ . . . ≥ λnK denoting the singular values of J in decreasing order and { us } nKs=1 ∈ RnK and { vs } nKs=1 ∈ Rp the corresponding left and right singular vectors forming the orthonormal basis matrices U ∈ RnK×nK and V ∈ Rp×nK . For a spectrum cutoff α obeying 0 ≤ α ≤ λ1 let r ∶= r ( α ) denote the index of the smallest singular value above α . We define the information and nuisance spaces associated with J as I ∶= span ( { us } rs=1 ) and N ∶= span ( { us } Kns=r+1 ) . In this paper we shall use either the expected value of the Jacobian at the random initialization or the Jacobian at one of the iterates to define the matrix J and the corresponding information/nuisance spaces . More , specifically we will set J to either J = ( E [ J ( W0 ) J T ( W0 ) ] ) 1/2 or J = J ( Wτ ) . Therefore , one can effectively think of the information space as the span of the prominent singular vectors of the Jacobian and the nuisance space as its complement . In particular , as we demonstrate in Section 4 the Jacobian mapping associated with neural networks exhibit low-rank structure with a few large singular values and many small ones leading to natural choices for the cut-off value α as well as the information and nuisance spaces . Furthermore , we demonstrate both ( empirically and theoretically ) that learning is fast over the information space leading to a significant reduction in both train/test accuracy in the early stages of training . However , after a certain number of iterations learning shifts to the nuisance space and reduction in the training error significantly slows down ( see Fig . 2 ) . Furthermore , subsequent iterations in this stage lead to a slight increase in test error . 2 MAIN RESULTS . Our main results establish multi-class generalization bounds for neural networks trained via gradient descent . First , we will focus on networks where both layers are randomly initialized . Next we will provide guarantees for arbitrary initialization with the goal of characterizing the generalization ability of subsequent iterative updates for a given ( possibly pre-trained ) network in terms of its Jacobian mapping . In this paper we focus on activations φ which are smooth and have bounded first and second order derivatives . This would for instance apply to the softplus activation φ ( z ) = log ( 1 + ez ) . We note that utilizing a proof technique developed in Oymak & Soltanolkotabi ( 2019 ) for going from smooth to ReLU activations it is possible to extend our results to ReLU activations with proper modifications . We avoid doing this in the current paper for clarity of exposition . Before we begin discussing our main results we discuss some notation used throughout the paper . For a matrix X ∈ Rn×d we use smin ( X ) and smax ( X ) = ∥X∥ to denote the minimum and maximum singular value of X . For two matrices A and B we use A ⊙B and A ⊗B to denote their Hadamard and Kronecker products , respectively . For a PSD matrix A ∈ Rn×n with eigenvalue decomposition A = ∑ni=1 λiuiuTi , the square root matrix is defined as A1/2 ∶= ∑ni=1 √λiuiuTi . We also use A† to denote the pseudo-inverse of A . In this paper we mostly focus on label vectors y which are one-hot encoded i.e . all entries are zero except one of them . For a subspace S ⊂ Rn and point x ∈ Rn , ΠS ( x ) denotes the projection of x onto S. Finally , before stating our results we need to provide a quantifiable measure of performance for a trained model . Given a sample ( x , y ) ∈ Rd × RK from a distribution D , the classification error of the networkW with respect to D is defined as ErrD ( W ) = P { arg max 1≤ ` ≤K y ` ≠ arg max1≤ ` ≤K f ` ( x ; W ) } . ( 2.1 )

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