Implicit Jacobian regularization weighted with impurity of probability output

Gradient descent ( GD ) plays a crucial role in the success of deep learning , but it is still not fully understood how GD finds minima that generalize well . In many studies , GD has been understood as a gradient flow in the limit of vanishing learning rate . However , this approach has a fundamental limitation in explaining the oscillatory behavior with iterative catapult in a practical finite learning rate regime . To address this limitation , we rather start with strong empirical evidence of the plateau of the sharpness ( the top eigenvalue of the Hessian ) of the loss function landscape . With this observation , we investigate the Hessian through simple and much lower-dimensional matrices . In particular , to analyze the sharpness , we instead explore the eigenvalue problem for the low-dimensional matrix which is a rank-one modification of a diagonal matrix . The eigendecomposition provides a simple relation between the eigenvalues of the low-dimensional matrix and the impurity of the probability output . We exploit this connection to derive sharpnessimpurity-Jacobian relation and to explain how the sharpness influences the learning dynamics and the generalization performance . In particular , we show that GD has implicit regularization effects on the Jacobian norm weighted with the impurity of the probability output . 1 INTRODUCTION . Deep learning has shown to be powerful for many learning tasks in various areas . There has been a lot of work to understand how the learning algorithm leads to this successful training of deep neural networks . Especially , it is crucial to understand the geometric properties of the loss landscape of neural networks and their interaction with the gradient-based optimization methods , such as Stochastic Gradient Descent ( SGD ) , along the training trajectory . It has been studied both from the optimization ( Gur-Ari et al. , 2018 ; Jastrzębski et al. , 2019 ; Ghorbani et al. , 2019 ; Liu et al. , 2020 ; Lewkowycz et al. , 2020 ; Cohen et al. , 2021 ) and generalization ( Hochreiter & Schmidhuber , 1997 ; Keskar et al. , 2017 ; Dinh et al. , 2017 ; Jastrzębski et al. , 2017 ; Wang et al. , 2018 ; Chaudhari et al. , 2019 ; Fort et al. , 2019 ; Jiang et al. , 2020 ; Barrett & Dherin , 2021 ; Smith et al. , 2021 ) point of view . We aim at investigating the Hessian of the training loss ( with respect to model parameter ) and its top eigenvalue ( also called sharpness ) . The sharpness characterizes the dynamics of neural network training along the optimization trajectory and is correlated with the generalization capability . For example , the sharpness increases in the beginning , and after reaching a certain value , training dynamics becomes unstable , oscillating along the top eigenvector ( Jastrzębski et al. , 2019 ; Cohen et al. , 2021 ) . Moreover , the rapid increase in the sharpness of the loss landscape in the early phase significantly impacts the final generalization performance ( Achille et al. , 2019 ; Jastrzebski et al. , 2020 ; Lewkowycz et al. , 2020 ; Jastrzebski et al. , 2021 ) . However , the Hessian of a deep neural network is very high-dimensional which makes it difficult to analyze its eigensystem . Recently , some researchers studied the Hessian by exploiting tools in randomized numerical linear algebra ( Sagun et al. , 2017 ; Papyan , 2018 ; 2019 ; Ghorbani et al. , 2019 ; Yao et al. , 2020 ) and decomposition of the Hessian ( Papyan , 2018 ; 2019 ; Fort & Ganguli , 2019 ) . In this paper , we present a new decomposition of the Hessian using eigendecomposition of lowdimensional matrices . From the eigensystem of the low-dimensional matrix , we can provide a simple and intuitive explanation on the relation between its eigenvalue and the probability output . This enables us to explain how the sharpness of the loss landscape influences the learning dynamics and the generalization performance . We summarize the main contributions of the paper as follows : • We decompose the Hessian with low dimensional matrices , the logit Hessian and the logit-weight Jacobian ( defined in Definition 1 ) , and investigate the Hessian by the eigendecomposition of the logit Hessian which is a rank-one modification of a diagonal matrix . • We provide connections between the top eigenvalue of the logit Hessian and the impurity of the probability output . • We derive a relation between the sharpness , the top eigenvalue of the logit Hessian and the Jacobian . We call it sharpness-impurity-Jacobian relation . • We explain how the sharpness of the loss landscape influences the learning dynamics and the generalization performance . In particular , we find that gradient-based optimizations have implicit effects on penalizing the Jacobian norm ( Implicit Jacobian Regularization ) in a certain phase of training ( Active Regularization Period ) . 2 RELATED WORK . We summarize some works on the Hessian , learning dynamics , and generalization of neural networks . In particular , we point out the issue of approximating SGD by a stochastic differential equation ( SDE ) because a continuous flow can not capture the oscillatory behavior of discrete updates with iterative catapult , which plays a key role in limiting the sharpness of the loss landscape . Decomposition of the Hessian Sagun et al . ( 2016 ; 2017 ) empirically found that the eigenvalue spectrum of the Hessian during training is composed of two parts , the bulk which is concentrated around zero and the outliers which are scattered positively away from zero . They showed the bulk depends on the size of the network , and the outliers depend on the data . In particular , the number of outliers matches the number of classes of the data . Further , Papyan ( 2019 ) proposed a three-level hierarchical decomposition of the Hessian matrix according to each class , logit coordinate , and example . However , with different decomposition , we analyze the Hessian from another point of view . SGD as a SDE In many studies , SGD has been understood as a SDE in the limit of vanishing learning rate ( Mandt et al. , 2017 ; Li et al. , 2017b ; a ; Smith & Le , 2018 ; Chaudhari & Soatto , 2018 ; Jastrzębski et al. , 2017 ; Zhu et al. , 2019 ; Park et al. , 2019 ) . However , some theoretical concerns have been raised for such approximations ( Yaida , 2019 ) . Moreover , Barrett & Dherin ( 2021 ) argued that the SDE analysis in the limit of vanishing learning rate can not explain the generalization benefits of finite learning rates and they proposed a modified gradient flow for finite learning rates . However , they still consider a continuous gradient flow and thus it has a fundamental limitation in explaining the oscillatory behavior with iterative catapult in a practical learning rate regime ( Smith et al. , 2021 ) , which will be detailed in the following paragraph . Oscillatory catapult and the plateau of the sharpness Xing et al . ( 2018 ) investigated the roles of learning rate and batch size in SGD dynamics through interpolating the loss landscape between consecutive model parameters during training . They observed SGD explores the parameter space , bouncing between walls of valley-like regions . The large learning rate maintains a high valley height , and a small batch size induces gradient stochasity . They both help exploration through the parameter space with different roles in the training dynamics . Jastrzębski et al . ( 2019 ) empirically investigated the evolution of the sharpness ( the top eigenvalue of the Hessian ) along the whole training trajectory of SGD . They observed initial growth of the sharpness as loss decreases , reaching a maximum sharpness determined by learning rate and batch size , and then it decreases towards the end of training . Due to the initial increase of the sharpness , the SGD step becomes too large compared to the shape of the loss landscape . This is consistent with the valley-like structure shown in Xing et al . ( 2018 ) . Lewkowycz et al . ( 2020 ) investigated simple theoretical models with a solvable training dynamics . They showed that , in their setup with a large learning rate , the loss initially increases while the sharpness decreases , then it converges to a flat minimum . This mechanism is called the catapult mechanism . Recently , Cohen et al . ( 2021 ) found that full-batch GD typically operates in a regime called the Edge of Stability where the sharpness can no longer increase and stays near a certain value , and the training loss behaves nonmonotonically but decreases globally . This behavior of the optimization at the Edge of Stability can be seen as repeated catapult mechanisms . They explicitly marked the limit of the sharpness with 2/η ( η =learning rate ) . To describe the aforementioned evolution of the sharpness , Fort & Ganguli ( 2019 ) developed a theoretical model based on a random matrix modelling . To build a simple random model , they introduced assumptions about gradients and Hessians that they are i.i.d isotropic Gaussian with zero mean with varying variance during training . While they focus on building a random model based on the observation , we rather aim to explain the underlying mechanisms . Implicit bias in SGD There have been many studies on the implicit bias in SGD ( Neyshabur , 2017 ; Zhang et al. , 2021 ; Soudry et al. , 2018 ) . We review the most relevant and recent ones . Jastrzebski et al . ( 2021 ) empirically showed that SGD implicitly penalizes the trace of the Fisher Information Matrix ( FIM ) . They also showed the trace of FIM explodes in the early phase of training when using a small learning rate and called it catastrophic Fisher explosion . Barrett & Dherin ( 2021 ) ; Smith et al . ( 2021 ) demonstrated that SGD implicitly penalizes the norm of the total gradient and the non-uniformity of the minibatch gradients . We demonstrate that the ( logit-weight ) Jacobian plays an important role in the generalization performance in each case . 3 BACKGROUND . In this section , we provide some notations , basic equations and definitions for the following sections . Throughout the paper , we use the denominator layout notation for the vector derivatives , i.e. , ∇vu = ( ∂uj ∂vi ) ij ∈ Rv×u where u : Rv → Ru and v ∈ Rv . It is also generalized to the cases of scalar , u = 1 or v = 1 . We consider a problem of learning a C-class classifier which maps an input x ∈ X ⊂ Rd to a target label y ∈ [ C ] where [ C ] = { 1 , 2 , · · · , C } . To this end , we build a parameterized model fθ : X → Z ⊂ RC with a model parameter θ ∈ Θ ⊂ Rm which outputs a logit vector z ≡ fθ ( x ) ∈ Z ⊂ RC ( we often omit the dependence on the input x and the model parameter θ ) . Then , the logit vector z is given as input to the softmax function to yield a probability vector p = softmax ( z ) ∈ ∆C−1 where ∆C−1 = { p ∈ [ 0 , 1 ] C : 1Tp = 1 , p ≥ 0 } . We want the model to match the most probable class c1 to the true label y , where c ( x ) ≡ arg sort ( p ) in descending order . We exchangeably denote the probability value corresponding to the true label y as p ≡ py ∈ [ 0 , 1 ] . The cross-entropy loss , l = l ( z , y ) ∈ R , is equivalent to the negative log-likelihood l = − log p. We use the notations ‖ · ‖ for the Euclidean ` 2-norm of a vector or for the Euclidean operator norm of a matrix ( equivalently , ‖ · ‖σ for a square matrix ) , ‖ · ‖F for the Frobenius norm , and tr ( · ) for the trace of a ( square ) matrix . Starting with a simple computation of the derivatives of the softmax function , Eq ( 1 ) ( see Appendix A ) , we can easily derive the following equations in order : ∇zp = diag ( p ) − ppT ∈ RC×C ( 1 ) ∇zp = [ ∇zp ] : ,y = p ( ey − p ) ∈ RC ( 2 ) ∇zl = ∇zp ∂l ∂p = p ( ey − p ) · −1 p = p− ey ∈ RC ( 3 ) ∇2zl = ∇z ( ∇zl ) = ∇z ( p− ey ) = diag ( p ) − ppT ∈ RC×C ( 4 ) where diag ( p ) = ( δijpi ) ij ∈ RC×C is a diagonal matrix with p as its diagonal entries , and ei = ( δij ) j ∈ RC is a one-hot vector with i-th element as 1 . Next , the Hessian of the loss function l for given example x with respect to the model parameter can be expressed as follows : ∇2θl = ∇θz∇2zl∇θzT + ∑C j=1 ∇2θzj∇zj l ≈ ∇θz∇2zl∇θzT ∈ Rm×m ( 5 ) using a well-known Gauss-Newton approximation . ( see , for example , Schraudolph ( 2002 ) ) . Now , we are ready to consider the training loss for the training set D. We compute the total training loss over D as L = 〈l〉 which yields∇L = 〈∇l〉 and∇2L = 〈∇2l〉 where 〈·〉 is the expectation over the empirical measure of the training set D , equivalently say ÊD [ · ] . We use the notation 〈·〉S when averaging over a subset S . Following from Eq ( 4 ) and Eq ( 5 ) , we define the Hessian matrixH for the total loss and its Gauss-Newton approximation matrixG with the matricesM and J as follows : Definition 1 . We callM the logit Hessian , J the Jacobian ( of the logit function with respect to the model parameter ) , H the Hessian , andG the Gauss-Newton approximation defined as follows : M ≡ ∇2zl = diag ( p ) − ppT ∈ RC×C ( 6 ) J ( = Jzθ ) ≡ ∇θz ∈ Rm×C ( 7 ) H ≡ 〈∇2θl〉 ≈ 〈JMJT 〉 ≡ G ∈ Rm×m ( 8 ) It is interesting to note that while l is dependent on the true label y , the logit HessianM = ∇2zl is independent of y , and so are J , JMJT , andG . In case of the MSE loss l = 12‖z − e y‖2 , we have M = ∇2zl = I and G = 〈JJT 〉 . We mainly focus on the usual cross-entropy loss and defer the investigation on the MSE loss to Appendix M. From Eq ( 8 ) , we will often use the approximation ‖H‖σ ≈ ‖G‖σ as justified in Sagun et al . ( 2017 ) ; Fort & Ganguli ( 2019 ) , but this approximation sometimes fails in the later phase of training when the top eigenvalues of the Gauss-Newton matrix is not sufficiently isolated from the bulk near 0 ( Papyan , 2018 ) . Thus we mainly focus on the early phase of training .

Summary: In this paper, the relationship between the Jacobian (the gradient of the final activation w.r.t parameters) and the Hessian is analyzed for the softmax cross-entropy loss. As a key tool, the approximation of the Hessian with the probability vector (the softmax output) is used, which suggests connections with several optimization (and generalization) concepts such as sharpness v.s. flatness of the loss landscape and the oscillation on the optimization path at GD/SGD.