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 method for estimating the trace of the Hessian of the cross-entropy loss with respect to the weights of a neural network classifier. They suggest adding the trace estimator as a regularizing penalty term for training neural networks with improved generalization.  The authors give two theoretical motivations for regularizing the Hessian trace: - For linear models the trace of the Hessian of the cross-entropy loss with respect to the logits appears as a factor in a term for bounding generalization error. The authors draw a connection between the Hessian with respect to the weights and the Hessian with respect to the logits and argue that penalizing the trace of the Hessian with respect to the weights also leads to a tighter bound on the generalization error. - For the continuous gradient dynamics around the optimum the Hessian eigenspectrum is the negative eigenspectrum of the Jacobian of the gradient dynamics which determine the stability of the dynamic system. The authors argue that decreasing Hessian trace increases Jacobian trace and hence reduces data dependend stability which will avoid overfitting.  The authors' suggested method for estimating the trace uses the stochastic Hutchinson trace estimator which samples a random vector sigma with E(sigma sigma^T) = I to get E(sigma^T H sigma) = tr(H). The authors appear to use back-propagation through a back-propagation gradient approach to compute sigma^T H sigma by first computing the gradient g = dl/domega of the loss l with respect to the weights omega, and then computing the gradient of the inner product g^T sigma which is a scalar, with respect to the weights again. This gradient-over-gradient can be inner product multiplied again with sigma.  In order to save cost the authors suggest a drop out method setting a fraction of the terms of sigma to 0 which saves a fraction of the derivative computation and correspondingly drops terms of the trace sum.  The authors then present experimental results evaluating their suggested regularization method in combination with and compared to other regularization methods on - CIFAR10 / Resnet18 - CIFAR100 / Wide Residual Network - WikiText-2 / 2-layer LSTM