Regularizing Deep Neural Networks with Stochastic Estimators of Hessian Trace

1 INTRODUCTION . Deep neural networks ( DNNs ) are developing rapidly and are widely used in many fields such as image classification , machine translation , language modeling and speech recognition . As more and more models are proposed in the literature , deep neural networks have shown remarkable improvements in performance . However , among various learning problems , over-fitting on training data is a great problem that affects the test accruacy . So a certain regularization method is often needed in the training process . In linear models , Ridge Regression ( Hoerl & Kennard , 1970 ) and Lasso ( Tibshirani , 1996 ) are usually used to avoid over-fitting . They are also called L2 and L1 regularization . L2 regularization has the effect of shrinkage while L1 regularization can be conductive to both shrinkage and sparsity . From the Bayesian perspective , L2 and L1 regularization can be interpreted with normal prior distribution and laplace prior distribution respectively . Apart from L2 and L1 regularization , there are many other forms of regularizers in DNNs . The most widely used one is Weight-Decay ( Krogh & Hertz , 1992 ) . Loshchilov & Hutter ( 2019 ) also showed that L2 regularization and Weight-Decay are not identical . Dropout ( Srivastava et al. , 2014 ) is another method to avoid over-fitting by reducing co-adapting between units in neural networks . Dropout has inspired a large body of work studying its effects ( Wager et al . ( 2013 ) ; Helmbold & Long ( 2015 ) ; Wei et al . ( 2020 ) ) . After dropout , various regularization schemes can be applied additionally . In this paper , we propose a new regularization by penalizing the trace of second derivative of loss function . We refer to our regularization method as Stochastic Estimators of Hessian Trace ( SEHT ) . On one hand , our hessian regularization is valuable to guarantee good generalization . On the other hand , from the perspective of dynamical system , it influences the stability of the system , in which parameters move in the parameter space on the basis of training data . In our experiments , Hessian regularization shows competing test performance and low time consumption with our stochastic algorithm . 2 RELATED WORK . There are many regularization methods in previous work . Label Smoothing ( Szegedy et al. , 2016 ) estimates the marginalized effect of label-dropout and reduces over-fitting by preventing a network from assigning full probability to each training example . Confidence Penalty ( Pereyra et al. , 2017 ) prevents peaked distributions , leading to better generalization . A network appears to be overconfident when it places all probability on a single class in the training set , which is often a symptom of over-fitting . DropBlock ( Ghiasi et al. , 2018 ) is a structured form of dropout , it drops contiguous regions from a feature map of a layer instead of dropping out independent random units . Data augmentation methods are also used in practice to improve model ’ s accuracy and robustness when training neural networks . Cutout ( DeVries & Taylor , 2017 ) is a data augmentation method where parts of the input examples are zeroed out , in order to encourage the network to focus more on less prominent features , then generalize to situations like occlusion . Mixup ( Zhang et al. , 2017 ) extends the training distribution by incorporating the prior knowledge that linear interpolations of feature vectors should lead to linear interpolations of the associated targets . Sokolić et al . ( 2017 ) first proposed Jacobian regularization , a method focusing on the norm of Jacobian matrix with respect to input data . It was proved that generalization error can be bounded by the norm of Jacobian matrix . Besides that , Jacobian matrix shows improved stability of the model predictions against input perturbations according to Taylor expansion . Hoffman et al . ( 2019 ) showed that Jacobian regularization enlarges the size of decision cells and is practically effective in improving the generalization error and robustness of the models . To simplify calculation , stochastic algorithm of Jacobian regularization was also proposed . Motivated by Jacobian regularization , we consider the generalization error and stability of the model respect to Hessian matrix . Then we combine Linear Stability Analysis and propose Hessian regularization with corresponding stochastic algorithms . We compare our Hessian regularization with other methods and demonstrate promising performance in experiments . The main idea to estimate the trace of Hessian matrix is Hutchinson Method ( Avron & Toledo , 2011 ) and the algorithm was also discussed by Yao et al . ( 2020a ) . We make an improvement by designing a new probability distribution to dropout parameters which decrease time consumption obviously without losing generalization . Hessian information is powerful tool used on analyzing the property of neural networks . Yao et al . ( 2020b ) designed AdaHessian , a second order stochastic optimization algorithm . Yu et al . ( 2021 ) used Hessian trace to measure sensitivity , developing a Hessian Aware Pruning method to find insensitive parameters in a neural network model and a Neural Implant technique to alleviate accuracy degradation . However , their methods are static in essence , we focus on dynamical motion of parameters in parameter space . Sankar et al . ( 2021 ) also proposed a Hessian regularization . They forcused on the layerwise loss landscape via the eigenspectrum of the Hessian at each layer . We start from different perspectives , generalization error and dynamical system of parameters . Our experiments also shows better results than Sankar et al . ( 2021 ) ’ s method . 3 HESSIAN REGULARIZATION . Here we introduce a Hessian regularization method based on generalization error and corresponding stochastic algorithms in details . Then we discuss linear stability analysis to explain why Hessian regularization can prevent neural networks from over-fitting . 3.1 TRACE OF HESSIAN MATRIX . In this study , we consider a multi-class classification problem . Input x is a N-dimensional vector , where x ∈ X ⊆ RN , with X is the input space . Y = { 1 , 2 , ... , M } is the label space , which means that we have M classes . Each input x has a label y ∈ Y . Sample space is defined as S = X × Y . An element of S is denoted by s = ( x , y ) . We assume that samples s from S are drawn according to a probability distribution P . A training set of n samples drawn from P is denoted by Sn = { si } ni=1 = { ( xi , yi ) } ni=1 . Our goal is to find a classification function f , which takes x ∈ RN as input and outputs z = f ( x ) . z is a M-dimensional score vector , where each element zi is the score x belonging to category i ∈ Y . The highest score indicates the most probable label . So the estimated label is given as ŷ = g ( x ) = arg maxi zi = arg maxi∈Y f ( x ) i . A loss function is used to measure the discrepancy between the true label y and the estimated label g ( x ) . In this paper , we use the cross-entropy loss , ` ( f ( x ) , y ) = − log exp ( f ( x ) y ) ∑ y′ exp ( f ( x ) y′ ) . ( 1 ) The empirical loss of the classifier f ( x ) associated with the training set is defined as ` emp ( f ) = Ê [ ` ( f ( x ) , y ) ] = 1 n ∑ si∈Sn ` ( f ( xi ) , yi ) , ( 2 ) and the expected loss of the classifier f ( x ) is defined as ` exp ( f ) = E [ ` ( f ( x ) , y ) ] = Es∼P [ ` ( f ( x ) , y ) ] , ( 3 ) then the difference between ` emp ( f ) and ` exp ( f ) is called generalization error : GE ( f ) = || ` exp ( f ) − ` emp ( f ) || . ( 4 ) Wei et al . ( 2020 ) showed generalization bound of linear models with cross-entropy loss of M classes . Let W is the weight matrix , µ ( W ) = Ê [ ‖J‖2 ] and v ( W ) = Ê [ tr ( H ) ] . For linear models , the Jacobian matrix J is a vector defined as ∂ ` ∂z and the Hessian matrix H is defined as { ∂ 2 ` ∂zi∂zj } M×M . With probability 1 − δ over the training examples , for all weight matrices W satisfying the norm bound ||W T ||2,1 ≤ A , the following bound holds : E [ ¯̀ ] − 1.01Ê [ ¯̀ ] . ( Aµ ( W ) ) 2 3 ( θB ) 1 3 n 1 3 + A √ Bv ( W ) θ√ n + BA2θ n ( log2 ( BA 2θ v ( W ) n ) + 1 ) + ζ . ( 5 ) Here with some fixed bound B > 0 , ¯̀= min { ` , B } , ‖W ‖2,1 = ∑ j √∑ i ( W 2ij ) , θ = log3 ( nM ) max i ‖xi‖22 , ζ = B ( log ( 1/δ ) + log log n ) n . ( 6 ) So one can guarantee good generalization when the trace of Hessian matrix and norm of Jacobian matrix are small . On one hand , when learning with gradient descent , we want to find a local or global minimum of loss function . Naturally , at minimum the gradient is zero and the norm of Jacobian matrix is small near minimum . So gradient descent helps us to ensure the norm of Jacobian Matrix to be small . On the other hand , the trace of Hessian Matrix is hard to be constrained by gradient descent . From this aspect , we proposed Hessian regularization for linear models as : 1 n tr ( H ` , z ) . ( 7 ) It ’ s the trace of second derivative of loss ` with respect to output of linear model z , which is also the end nodes of a linear model . A DNN is consist of many layers , with each layer being viewed as a linear model ( except the nonlinear activation functions ) . Thus . we generalize the Hessian regularization to every node in a DNN and define it as 1 n tr ( H ` , ω ) . ( 8 ) It ’ s the trace of second derivative of loss ` with respect to parameters ω . Here we define a new loss with our Hessian regularization as Loss = ` emp ( f ) + λ · 1 n tr ( H ` , ω ) , ( 9 ) where λ controls the strength of our Hessian regularization . 3.2 HUTCHINSON METHOD . In a typical DNN , there are more than millions of parameters . So the calculation of Hessian matrix is difficult . Hutchinson Method ( Avron & Toledo , 2011 ) is an unbiased estimator for the trace of a matrix . Let A be an n × n symmetric matrix with tr ( A ) 6= 0 . Let σ be a random vector whose entries are i.i.d Rademacher random variables ( Pr ( σi = ±1 ) = 12 ) , then σ TAσ is an unbiased estimator of tr ( A ) , based on the following equation : tr ( A ) = tr ( AI ) = tr ( AE [ σσT ] ) = E [ tr ( AσσT ) ] = E [ σTAσ ] . ( 10 ) In this paper , we consider the trace of Hessian matrix H , which is the second derivative matrix . Since the Rademacher random vector is irrelevant to network parameters , we expand the expression of Hutchinson estimator as follow : σTHσ = σT d dldω dω σ = σT ( d dldω dω · σ + dl dω · dσ dω ) = σT ( d ( dldω · σ ) dω ) . ( 11 ) Based on Equation 11 and dσdω = 0 , we only have to calculate the gradient g = dl dω and the derivative of dldω ·σ , instead of the whole Hessian matrix . The whole calculation process only include two inner products and two derivations . We refer to the Hutchinson stochastic estimator of Hessian trace as SEHT-H. Algorithm 1 : SEHT-H Input : n-dimensional gradient g Output : Estimation of tr ( H ) 1 trace = 0 ; 2 for i = 1 to maxIter do 3 σ ∼ Rademacher ( n ) ; / * n-dim vector with each element sampled from Rademacher distribution . * / 4 v = g · σ ; / * inner product * / 5 h = dvdω ; / * derivative of v * / 6 t = σT · h ; / * inner product * / 7 trace += t ; 8 end 9 return tracemaxIter Even though it ’ s a stochastic algorithm , it cost much time because of a great number of parameters . So we propose another efficient algorithm below based on the basic idea of Hutchinson Method .

Summary: The authors propose a regularizer for training DNNs - specifically adding the trace of the Hessian w.r.t. the parameters of the model. They present two efficient stochastic estimators of the trace, based on the Hutchinson method and their own extension to it, specific to NNs.  The authors also provide a motivation of the regularizer through a known generalization bound and through the perspective of dynamical systems stability.  In three experiments - on CIFAR10, CIFAR100 and Wiki-Text they present results suggesting that the proposed method is performing better on average than other techniques.