Understanding Why Neural Networks Generalize Well Through GSNR of Parameters

1 INTRODUCTION . Deep neural networks typically contain far more trainable parameters than training samples , which seems to easily cause a poor generalization performance . However , in fact they usually exhibit remarkably small generalization gaps . Traditional generalization theories such as VC dimension ( Vapnik & Chervonenkis , 1991 ) or Rademacher complexity ( Bartlett & Mendelson , 2002 ) can not explain its mechanism . Extensive research focuses on the generalization ability of DNNs ( Neyshabur et al. , 2017 ; Arora et al. , 2018 ; Keskar et al. , 2016 ; Dinh et al. , 2017 ; Hoffer et al. , 2017 ; Novak et al. , 2018 ; Dziugaite & Roy , 2017 ; Jakubovitz et al. , 2019 ; Kawaguchi et al. , 2017 ; Advani & Saxe , 2017 ) . Unlike that of shallow models such as logistic regression or support vector machines , the global minimum of high-dimensional and non-convex DNNs can not be found analytically , but can only be approximated by gradient descent and its variants ( Zeiler , 2012 ; Kingma & Ba , 2014 ; Graves , 2013 ) . Previous work ( Zhang et al. , 2016 ; Hardt et al. , 2015 ; Dziugaite & Roy , 2017 ) suggests that the generalization ability of DNNs is closely related to gradient descent optimization . For example , Hardt et al . ( 2015 ) claims that any model trained with stochastic gradient descent ( SGD ) for reasonable epochs would exhibit small generalization error . Their analysis is based on the smoothness of loss function . In this work , we attempt to understand the generalization behavior of DNNs through GSNR and reveal how GSNR affects the training dynamics of gradient descent . Stanislav Fort ( 2019 ) studied a new gradient alignment measure called stiffness in order to understand generalization better and stiffness is related to our work . The GSNR of a parameter is defined as the ratio between its gradient ’ s squared mean and variance over the data distribution . Previous work tried to use GSNR to conduct theoretical analysis on deep learning . For example , Rainforth et al . ( 2018 ) used GSNR to analyze variational bounds in ∗corresponding author unsupervised DNNs such as variational auto-encoder ( VAE ) . Here we focus on analyzing the relation between GSNR and the generalization gap . Intuitively , GSNR measures the similarity of a parameter ’ s gradients among different training samples . Large GSNR implies that most training samples agree on the optimization direction of this parameter , thus the parameter is more likely to be associated with a meaningful “ pattern ” and we assume its update could lead to a better generalization . In this work , we prove that the GSNR is strongly related to the generalization performance , and larger GSNR means a better generalization . To reveal the mechanism of DNNs ’ good generalization ability , we show that the gradient descent optimization dynamics of DNN naturally leads to large GSNR of model parameters and therefore good generalization . Furthermore , we give a complete analysis and a detailed interpretation to this phenomenon . We believe this is probably the key to DNNs remarkable generalization ability . In the remainder of this paper we first analyze the relation between GSNR and generalization ( Section 2 ) . We then show how the training dynamics lead to large GSNR of model parameters experimentally and analytically in Section 3 . 2 LARGER GSNR LEADS TO BETTER GENERALIZATION . In this section , we establish a quantitative relation between the GSNR of model parameters and generalization gap , showing that larger GSNR during training leads to better generalization . 2.1 GRADIENTS SIGNAL TO NOISE RATIO . Consider a data distribution Z = X × Y , from which each sample ( x , y ) is drawn ; a model ŷ = f ( x , θ ) parameterized by θ ; and a loss function L. The parameters ’ gradient w.r.t . L and sample ( xi , yi ) is denoted by g ( xi , yi , θ ) or gi ( θ ) : = ∂L ( yi , f ( xi , θ ) ) ∂θ ( 1 ) whose j-th element is gi ( θj ) . Note that throughout this paper we always use i to index data examples and j to index model parameters . Given the data distribution Z , we have the ( sample-wise ) mean and variance of gi ( θ ) . We denote them as g̃ ( θ ) = E ( x , y ) ∼Z ( g ( x , y , θ ) ) and ρ2 ( θ ) = Var ( x , y ) ∼Z ( g ( x , y , θ ) ) , respectively . The gradient signal to noise ratio ( GSNR ) of one model parameter θj is defined as : r ( θj ) : = g̃2 ( θj ) ρ2 ( θj ) ( 2 ) At a particular point of the parameter space , GSNR measures the consistency of a parameter ’ s gradients across different data samples . Figure 1 intuitively shows that if GSNR is large , the parameter gradient space tends to be distributed in the similar direction and if GSNR is small , the gradient vectors are then scatteredly distributed . 2.2 ONE-STEP GENERALIZATION RATIO . In this section we introduce a new concept to help measure the generalization performance during gradient descent optimization , which we call one-step generalization ratio ( OSGR ) . Consider training set D = { ( x1 , y1 ) , ... , ( xn , yn ) } ∼ Zn with n samples drawn from Z , and a test set D′ = { ( x′1 , y′1 ) , ... , ( x′n′ , y′n′ ) } ∼ Zn ′ . In practice we use the loss on D′ to measure generalization . For simplicity , we assume the sizes of training and test datasets are equal , i.e . n = n′ . We denote the empirical training and test loss as : L [ D ] = 1 n n∑ i=1 L ( yi , f ( xi , θ ) ) , L [ D ′ ] = 1 n n∑ i=1 L ( y′i , f ( x ′ i , θ ) ) , ( 3 ) respectively . Then the empirical generalization gap is given by L [ D′ ] − L [ D ] . In gradient descent optimization , both the training and test loss would decrease step by step . We use ∆L [ D ] and ∆L [ D′ ] to denote the one-step training and test loss decrease during training , respectively . Let ’ s consider the ratio between the expectations of ∆L [ D′ ] and ∆L [ D ] of one single training step , which we denote as R ( Z , n ) . R ( Z , n ) : = ED , D ′∼Zn ( ∆L [ D ′ ] ) ED∼Zn ( ∆L [ D ] ) ( 4 ) Note that this ratio also depends on current model parameters θ and learning rate λ . We are not including them in the above notation as we will not explicitly model these dependencies , but rather try to quantitatively characterize R for very small λ and for θ at the early stage of training ( satisfying Assumption 2.3.1 ) . Also note that the expectation of ∆L [ D′ ] is over D and D′ . This is because the optimization step is performed on D. We refer to R ( Z , n ) as OSGR of gradient descent optimization . Statistically the training loss decreases faster than the test loss and 0 < OSGR ( t ) < 1 ( Middle panel of Figure 2 ) , which usually results in a non-zero generalization gap at the end of training . If OSGR ( t ) is large ( ≈ 1 ) in the whole training process ( Right panel of Figure 2 ) , generalization gap would be small when training completes , implying good generalization ability of the model . If OSGR ( t ) is small ( = 0 ) , the test loss will not decrease while the training loss normally drops ( Left panel of Figure 2 ) , corresponding to a large generalization gap . 2.3 RELATION BETWEEN GSNR AND OSGR . In this section , we derive a relation between the OSGR during training and the GSNR of model parameters . This relation indicates that , for the first time as far as we know , the sample-wise gradient distribution of parameters is related to the generalization performance of gradient descent optimization . In gradient descent optimization , we take the average gradient over training set D , which we denote as gD ( θ ) . Note that we have used gi ( θ ) to denote gradient evaluated on one data sample and g̃ ( θ ) to denote its expectation over the entire data distribution . Similarly we define gD′ ( θ ) to be the average gradient over test set D′ . gD ( θ ) = 1 n n∑ i=1 g ( xi , yi , θ ) = ∂L [ D ] ∂θ , gD′ ( θ ) = 1 n n∑ i=1 g ( x′i , y ′ i , θ ) = ∂L [ D′ ] ∂θ ( 5 ) Both the training and test dataset are randomly generated from the same distribution Zn , so we can treat gD ( θ ) and gD′ ( θ ) as random variables . At the beginning of the optimization process , θ is randomly initialized thus independent of D , so gD ( θ ) and gD′ ( θ ) would obey the same distribution . After a period of training , the model parameters begin to fit the training dataset and become a function of D , i.e . θ = θ ( D ) , therefore distributions of gD ( θ ( D ) ) and gD′ ( θ ( D ) ) become different . However we choose not to model this dependency and make the following assumption for our analysis : Assumption 2.3.1 ( Non-overfitting limit approximation ) The average gradient over the training dataset and test dataset gD ( θ ) and gD′ ( θ ) obey the same distribution . Obviously the mean of gD ( θ ) and gD′ ( θ ) is just the mean gradient over the data distribution g̃ ( θ ) . ED∼Zn [ gD ( θ ) ] = ED , D′∼Zn [ gD′ ( θ ) ] = g̃ ( θ ) ( 6 ) We denote their variance as σ2 ( θ ) , i.e . VarD∼Zn [ gD ( θ ) ] = VarD , D′∼Zn [ gD′ ( θ ) ] = σ 2 ( θ ) ( 7 ) It is straightforward to show that : σ2 ( θ ) = VarD∼Zn [ 1 n n∑ i=1 gi ( θ ) ] = 1 n ρ2 ( θ ) ( 8 ) where σ2 ( θ ) is the variance of the average gradient over the dataset of size n , and ρ2 ( θ ) is the variance of the gradient of a single data sample . In one gradient descent step , the model parameter is updated by ∆θ = θt+1 − θt = −λgD ( θ ) where λ is the learning rate . If λ is small enough , the one-step training and test loss decrease can be approximated by ∆L [ D ] ≈ −∆θ · ∂L [ D ] ∂θ +O ( λ2 ) = λgD ( θ ) · gD ( θ ) +O ( λ2 ) ( 9 ) ∆L [ D′ ] ≈ −∆θ · ∂L [ D ′ ] ∂θ +O ( λ2 ) = λgD ( θ ) · gD′ ( θ ) +O ( λ2 ) ( 10 ) Usually there are some differences between the directions of gD ( θ ) and gD′ ( θ ) , so statistically ∆L [ D ] tends to be larger than ∆L [ D′ ] and the generalization gap would increase during training . When λ → 0 , in one single training step the empirical generalization gap increases by ∆L [ D ] − ∆L [ D′ ] , for simplicity we denote this quantity as ▽ : ▽ : = ∆L [ D ] −∆L [ D′ ] ≈ λgD ( θ ) · gD ( θ ) − λgD ( θ ) · gD′ ( θ ) ( 11 ) = λ ( g̃ ( θ ) + ϵ ) ( g̃ ( θ ) + ϵ− g̃ ( θ ) − ϵ′ ) ( 12 ) = λ ( g̃ ( θ ) + ϵ ) ( ϵ− ϵ′ ) ( 13 ) Here we replaced the random variables by gD ( θ ) = g̃ ( θ ) + ϵ and gD′ ( θ ) = g̃ ( θ ) + ϵ′ , where ϵ and ϵ′ are random variables with zero mean and variance σ2 ( θ ) . Since E ( ϵ′ ) = E ( ϵ ) = 0 , ϵ and ϵ′ are independent , the expectation of ▽ is ED , D′∼Zn ( ▽ ) = E ( λϵ · ϵ ) +O ( λ2 ) = λ ∑ j σ2 ( θj ) +O ( λ 2 ) ( 14 ) where σ2 ( θj ) is the variance the of average gradient of the parameter θj . For simplicity , when it involves a single model parameter θj , we will use only a subscript j instead of the full notation . For example , we use σ2j , rj , and gD , j to denote σ 2 ( θj ) , r ( θj ) , and gD ( θj ) respectively . Consider the expectation of ∆L [ D ] and ∆L [ D′ ] when λ → 0 ED∼Zn ( ∆L [ D ] ) ≈ λED∼Zn ( gD ( θ ) · gD ( θ ) ) = λ ∑ j ED∼Zn ( g 2 D , j ) ( 15 ) ED , D′∼Zn ( ∆L [ D ′ ] ) = ED , D′∼Zn ( ∆L [ D ] −▽ ) ( 16 ) ≈ λ ∑ j ( ED∼Zn ( g 2 D , j ) − σ2j ) ( 17 ) = λ ∑ j ( ED∼Zn ( g 2 D , j ) − ρ2j/n ) ( 18 ) Substituting ( 18 ) and ( 15 ) into ( 4 ) we have : R ( Z , n ) = 1− ∑ j ρ 2 j n ∑ j ED∼Zn ( g 2 D , j ) ( 19 ) Although we derived eq . ( 19 ) from simplified assumptions , we can empirically verify it by estimating two sides of the equation on real data . We will elaborate on this estimation method in section 2.4 . We can rewrite eq . ( 19 ) as : R ( Z , n ) = 1− 1 n ∑ j ED∼Zn ( g 2 D , j ) ∑ j′ ED∼Zn ( g 2 D , j′ ) ρ2j ED∼Zn ( g2D , j ) ( 20 ) = 1− 1 n ∑ j ED∼Zn ( g 2 D , j ) ∑ j′ ED∼Zn ( g 2 D , j′ ) 1 rj + 1 n ( 21 ) where ED∼Zn ( g2D , j ) = V arD∼Zn ( gD , j ) + E 2 D∼Zn ( gD , j ) = 1 nρ 2 j + g̃ 2 j . We define ∆Lj [ D ] to be the training loss decrease caused by updating θj . We can show that when λ is very small ∆Lj [ D ] = λg2D , j +O ( λ 2 ) . Therefore when λ → 0 , we have R ( Z , n ) = 1− 1 n ∑ j Wj 1 rj + 1 n , where Wj : = ED∼Zn ( ∆Lj [ D ] ) ED∼Zn ( ∆L [ D ] ) with ∑ j Wj = 1 ( 22 ) Eq . ( 22 ) shows that the GSNR rj plays a crucial role in the model ’ s generalization ability—the onestep generalization ratio in gradient descent equals one minus the weighted average of 1 rj+ 1 n over all model parameters divided by n. The weight is proportional to the expectation of the training loss decrease resulted from updating that parameter . This implies that larger GSNR of model parameters during training leads to smaller generalization gap growth thus better generalization performance of the trained model . Also note when n → ∞ , we have R ( Z , n ) → 1 , meaning that training on more data helps generalization .

Summary: In this work, the authors suggest a new point of view on generalization through the lens of the distribution of the per-sample gradients. The authors consider the variance and mean of the per-sample gradients for each parameter of the model and define for each parameter the Gradient Signal to Noise ratio (GSNR). The GSNR of a parameter is the ratio between the mean squared of the gradient per parameter per sample (computed over the samples) and the variance of the gradient per parameter per sample (also computed over the samples). The GSNR is promising as a measure of generalization and the authors provide a nice leading order derivation of the GSNR as a proxy for the measure of the generalization gap in the model. After the derivation, experimental results on MNIST are presented and suggest that empirically there is a relation between the generalization gap of neural network trained by gradient descent and the GSNR quantity given in the paper. Next the author analyze the GNSR of DNNs as opposed with shallow models or other learning techniques and observe that the GSNR differs when using random labels (lower GSNR) as compared with true labels and exhibits different behavior along training for DNNs and gradient descent.