IMPLICIT JACOBIAN REGULARIZATION WEIGHTED WITH IMPURITY OF PROBABILITY OUTPUT
Year: 2021

Abstract:
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.

text: 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.

heading: 1 INTRODUCTION
text: 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).

heading: 2 RELATED WORK
text: 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 cannot 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 cannot 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.

heading: 3 BACKGROUND
text: 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.

heading: 4 DECOMPOSITION OF THE HESSIAN H
text: In the previous section, we introduced the Gauss-Newton approximationG of the HessianH , and decomposition of G with the Jacobian J and the logit Hessian M , i.e., G = 〈JMJT 〉. Now, we focus on the logit Hessian matrixM and its eigendecomposition, estimate the top eigenvalue ofM with upper/lower bounds, and explore the evolution of the top eigenvalue during training.

heading: 4.1 EIGENDECOMPOSITION OF THE LOGIT HESSIAN
text: The lower-dimensional matrixM ∈ RC×C is simple and fully characterized by only the probability vector p as M = diag(p) − ppT in Eq (6), but it turns out to be important for understanding the much higher-dimensional matrixG ∈ Rm×m (C m). SinceM = diag(p)− ppT is a rank-one modification of a simple diagonal matrix diag(p), we can obtain the eigenvalues and the eigenvectors from the theory of the rank-one modification of the eigenproblem (see, for example, Bunch et al. (1978); Golub (1973) and (Golub & Van Loan, 2013, Section 8.4.3)). Then, the logit HessianM can be eigendecomposed asM = QΛQT = ∑C i=1 λ
(i)q(i)q(i)T where λ(i) is the i-th largest eigenvalue of M and q(i) is its corresponding normalized eigenvector. For simplicity, we also use the same ordered index of (i) ∈ [C] with parentheses for the probability output p ∈ ∆C−1, i.e., ci = (i) and p(1) ≥ p(2) ≥ · · · ≥ p(C) ≥ 0, because this ordering is related to the eigenvalues {λ(i)}Ci=1 as demonstrated in the following theorem. We defer the proof to Appendix B.
Theorem 1 (cf. Golub (1973); Bunch et al. (1978)). The eigenvalues λ(i) (λ(1) ≥ λ(2) ≥ · · · ≥ λ(C)) and the corresponding normalized eigenvectors q(i) of the logit HessianM = ∇2zl = diag(p)−ppT satisfy the following properties:
(a) The eigenvalue λ(i) is the i-th largest solution of the following equation:
v(λ) = 1− ∑C
i=1 p2i pi − λ = 0 (9)
(b) The eigenvector q(i) is aligned with the direction of (diag(p)− λ(i)I)−1p
(c) p(i+1) ≤ λ(i) ≤ p(i) for 1 ≤ i ≤ C − 1, and λ(C) = 0
(d) 12Gini(p(1)) ≤ λ (1) ≤ Gini(p(1)) where Gini(q) = 1− q2 − (1− q)2 = 2q(1− q) is the
Gini impurity for the binary case (q, 1− q).
Figure 1 illustrates the secular function v(λ) defined as Eq (9) in Theorem 1 (a). The function v(λ) has singularities at the probability values {p(i)}Ci=1 and zeros at the eigenvalues {λ(i)}Ci=1, satisfying Theorem 1 (c). Moreover, Theorem 1 (c) and (d) provide the upper/lower bounds on the eigenvalues ofM . Especially, the top eigenvalue λ(1) is bounded by 12Gini(p(1)) ≤ λ
(1) ≤ Gini(p(1)), and thus we call it impurity.

heading: 4.2 EVOLUTION OF IMPURITY
text: We explore the top eigenvalue λ(1) of M (also referred to as impurity) during training. Figure 2 demonstrates the evolution of the impurity during training, which increases in the beginning and then decreases in the later phase. We trained a model to zero training loss, and thus, for most examples, the probability py for the true class y eventually becomes the highest probability p(1). As the top probability p(1) increases from 1/C to 1 during training, the impurity starts from λ(1) ≈ 1C ∈ [ 12Gini( 1 C ), 1 C ] = [ C−1 C2 , 1 C ] (Theorem 1 (c) and (d)), increases at the initial phase of the training because it is lower bounded by 12Gini(p(1)) = p(1)(1 − p(1)) which increases for p(1) ∈ [0, 0.5]. Then λ(1) decreases as p(1) becomes larger than 0.5 which leads λ(1) to almost 0 at the later phase because it is upper bounded by Gini(p(1)) = 2p(1)(1 − p(1)) which decreases for p(1) ∈ [0.5, 1]. Note that Cohen et al. (2021) tried to estimate a similar value, but they use py, not p(1). We again emphasize thatM is independent of the label y but dependent only on the probability output p.

heading: 5 IMPLICIT JACOBIAN REGULARIZATION
text: In this section, from the results of the previous sections, we aim to derive a relation between the sharpness, the impurity and the Jacobian, and answer how the sharpness of the loss landscape
influences the learning dynamics and the generalization performance. Detailed experimental settings for each Figure are described in Appendix E.

heading: 5.1 SHARPNESS-IMPURITY-JACOBIAN RELATION
text: We first take a closer look at the sharpness of the loss landscape during training and build a relation between the sharpness ‖H‖σ , the impurity λ(1), and the Jacobian J . Since the Gauss-Newton matrix G is known to approximate the true HessianH well, especially for the top eigenspace (Sagun et al., 2017; Fort & Ganguli, 2019; Papyan, 2019), we can write the sharpness ‖H‖σ as follows:
‖H‖σ ≈ ‖G‖σ = ‖〈JMJT 〉‖σ (10)
It implies the impurity ‖M‖σ and the squared norm of the Jacobian J are highly correlated with the sharpness ‖H‖σ as demonstrated in the following theorem. We defer the proof to Appendix C. Theorem 2. For some 0 ≤ λ∗ ≤ λ(1) for each x ∈ D, we can bound
‖〈JMJT 〉‖σ = 〈λ∗‖J‖2〉 ≤ 〈λ(1)‖J‖2〉 (11)
In other words, λ∗‖J‖2 has the expected value of approximately ‖H‖σ, and we briefly state this relation with λ∗‖J‖2 ∼ ‖H‖σ. Here, λ∗ in λ∗‖J‖2 acts as an adaptive regularization weight for the regularization on the Jacobian norm ‖J‖2 as detailed in the following section. Since it is computationally inefficient to track ‖J‖2 for every x ∈ D, we instead investigate ‖〈J〉‖2. We expect ‖H‖σ = λ̂∗‖〈J〉‖2 for some λ̂∗. In Figure 3, we observe that ‖H‖σ and ‖〈J〉‖2 show similar oscillating behavior up to a factor λ̂∗ which is locally constant and slowly changes during training.

heading: 5.2 IMPLICIT JACOBIAN REGULARIZATION
text: Now, we are ready to answer how the sharpness of the loss landscape influences the learning dynamics and the generalization performance.
Growing Jacobian and sharpness in the early phase of training The weight norm ‖θ‖ increases in order to increase the logit norm ‖z‖ and to minimize the cross-entropy loss during training (Soudry et al., 2018) (see Appendix F for details). This also leads to the increase in the layerwise weight norms and the Jacobian norm. Thus, the increase of the sharpness in the early phase of training can be mainly attributed to the increase of the Jacobian norm ‖J‖. However, as will be discussed in the following paragraphs, the Jacobian norm does not continuously increase throughout training.
Oscillatory catapult and the plateau of the sharpness As the sharpness increases in the beginning, the width of the valley of the loss landscape becomes narrower than the discrete step size of the gradient-based optimization. After the sharpness reaching this threshold, the iterate starts to bounce
off from one side of the valley to the other, and repeats this (Xing et al., 2018; Jastrzębski et al., 2019). Figure 4 shows this oscillatory behavior with iterative catapult after the sharpness reaching the threshold, using UMAP (McInnes et al., 2018). Due to the catapult, the iterate cannot stay in a sharper area and is catapulted to another area. This mechanism is called the catapult mechanism (Lewkowycz et al., 2020). Figure 3 shows fine-grained patterns that the sharpness oscillates around the threshold by the two conflicting effects: the Jacobian norm tends to increase the sharpness and the catapult mechanism reduces it again when the sharpness is over the threshold. Therefore, we can observe the plateau of the sharpness in a coarser scale (see Figure 5). For GD with the MSE loss, the threshold of the sharpness can be easily obtained to be 2/η (η = learing rate), and the threshold value is similar for the cross-entropy loss as shown in Figure 3 (Cohen et al., 2021). The plateau of the sharpness is also observed in the case of SGD, but with a different threshold value depending on the batch size (the smaller the batch size, the lower the plateau) as shown in Figure 6 (Right). This oscillatory catapult and the plateau of the sharpness are attributed to the discrete dynamics of the gradient-based optimization with a finite learning rate and cannot be described by a continuous gradient flow.
Implicit Jacobian Regularization (IJR) Due to this catapult effect, the Jacobian norm cannot continue to increase. In other words, the gradient-based optimization implicitly penalizes the Jacobian norm since ‖J‖2 ∼ ‖H‖σ/λ∗ / 2ηλ∗ . This IJR effect begins after the sharpness reaches the threshold. And the regularization effect diminishes as λ(1) decreases (the upper limit 2ηλ∗ is loosened) with increasing p(1) ≥ 0.5 in the later phase (see Figure 2, 5 and 6) because λ∗ ≤ λ(1) acts as a regularization coefficient. This explains why the behavior of the sharpness in the early phase of training seems to impact the final generalization. Moreover, as λ(1) decreases to a very small value, we can observe the decrease in the sharpness. However, the test accuracy is saturated, which provides another counter-example that flat minima generalizes poorly (Dinh et al., 2017).
Figure 5 shows that the gradient-based optimization has implicit regularization effects on the Jacobian norm in a certain period. To be specific, there are three phases that the Jacobian norm is (I) initially rapidly increasing before the sharpness reaches near the threshold, (II) actively regularized with a gentle slope, and (III) again exponentially increasing as the regularization effect diminishes (as the regularization weight λ∗ ≤ λ(1) decreases) with the slope being gradually steeper. We call the second phase Active Regularization Period.
The evolution of the sharpness of the loss landscape is highly affected by the hyperparameters used in training such as learning rate and batch size (Jastrzębski et al., 2019; Lewkowycz et al., 2020; Cohen et al., 2021). As expected, GD with a large learning rate η limits the sharpness with a lower value of 2/η. The large learning rate encourages the stronger implicit regularization on the Jacobian norm. Figure 6 shows, comparing the three red lines (solid/dashed/dotted) with different learning rates (0.01/0.02/0.03) and batch sizes (128/64/32), training with a larger learning rate and a smaller batch size limits the Jacobian norm with a smaller value in the Active Regularization Period. This could explain why training with a large learning rate and a small batch size yields better generalization.
Explicit Jacobian Regularization (EJR) To further investigate and boost the effectiveness of IJR, we aim to explicitly regularize the Jacobian norm. However, it is computationally hard to back-propagate through the computation graph of the operator norm of ‖〈J〉‖2 for a practical neural network even with a simple iterative method (see Algorithm 2 in Appendix E). Thus, we instead penalize an upper bound, the Frobenius norm ‖〈J〉‖2F (≥ ‖〈J〉‖2), with the regularization coefficient λreg/C, i.e., we minimize L+ λreg‖〈J〉‖2F /C. See Appendix L for its variants with other regular-
ization coefficients ρreg and µreg. The Frobenius regularization term can be efficiently computed with an unbiased estimator ‖〈J〉‖2F = CEu∼U(SC−1)[‖〈J〉u‖2] = CEu∼U(SC−1)[‖∇θ〈uTz〉‖2] where u is randomly drawn from the unit hypersphere SC−1. Since the batch-size is large enough, we efficiently use a single sample u ∼ U(SC−1) for each batch as suggested in Hoffman et al. (2019). We expect improvements in the generalization performance when introducing EJR. This would support the effectiveness of IJR that it efficiently controls the capacity of the model. Table 1 shows clear improvements in the test accuracy when introducing EJR.
Connections between the Jacobian and Fisher/Gradient Penalty Our explanation of the implicit bias in SGD may extend to the catastrophic Fisher explosion (Jastrzebski et al., 2021) withG instead of the Fisher Information Matrix (FIM). The trace ofG can be written as follows:
tr(G) = 〈tr(JMJT )〉 = 〈 ∑C−1
i=1 λ(i)‖Jq(i)‖2〉 ≈ 〈tr(M)‖J‖2F 〉/C (12)
where we assume ‖Jq(i)‖2 ≈ ‖J‖2F /C since ∑C i=1 ‖Jq(i)‖2 = ‖J‖2F (see Figure 11 (Left) in Appendix D for the empirical evidence). Here, the trace of the logit HessianM can be equivalently written as a C-class Gini impurity: tr(M) = ∑C i=1 pi(1− pi) = 1− ∑C i=1 p 2 i ≡ Gini
C(p) which is C−1C for the initial uniform distribution and 0 for a one-hot probability. Thus, penalizing tr(G) induces the effects of penalizing ‖J‖F , especially in the early phase of training with large GiniC(p). Thus, as Jastrzebski et al. (2021) argued that Fisher Penalty on the trace of the FIM improves the generalization performance by limiting the memorization, the Jacobian regularization may have similar effects. Moreover, since ∇θl(z, ŷ) = ∇θz∇zl(z, ŷ) = J(p − eŷ), the trace of the FIM they approximately used is simply ‖Êx∼BEŷ∼p[∇θl(z, ŷ)]‖2 = ‖Êx∼BEŷ∼p[J(p − eŷ)]‖2 with a single sample ŷ, the gradient norm penalty (Barrett & Dherin, 2021) is ‖Ê(x,y)∼B[J(p − ey)]‖ and the explicit Jacobian regularizer is ‖J‖2F = CEu∼U(SC−1)‖Ê(x,y)∼B[Ju]‖2. In each case, we emphasize that the Jacobian J plays an important role in the generalization performance.

heading: 6 CONCLUSION
text: We have investigated the eigensystem of the Hessian through a new decomposition using eigendecomposition of the low-dimensional logit Hessian. By doing so, we could provide a simple and intuitive explanation on the relation between the gradient-based optimization, the learning dynamics and the generalization performance of neural networks. We hope this research could help answer other intriguing questions on the learning dynamics and generalization.
ETHICS STATEMENT
No ethics statement is needed for this study.
REPRODUCIBILITY STATEMENT
We provide code to reproduce our experiments. We refer the readers to the supplementary material (README.md).

heading: A PROOF OF EQ (1)
text: ∇zp = diag(p)− ppT ∈ RC×C (1)
Proof. By definition of the softmax function,
pj = [softmax(z)]j = exp(zj)∑ k exp(zk) = exp(zj)s −1 (13)
where s = ∑ k exp(zk), we have
∇zipj = { −exp(zj)s−2exp(zi) = −pipj , if i 6= j −exp(zi)s−2exp(zi) + exp(zi)s−1 = −p2i + pi, if i = j
(14)
which leads to∇zp = (∇zipj)ij = −ppT + diag(p).

heading: B PROOF OF THEOREM 1
text: Theorem (restated). The eigenvalues λ(i) (λ(1) ≥ λ(2) ≥ · · · ≥ λ(C)) and the corresponding normalized eigenvectors q(i) of the logit HessianM = ∇2zl = diag(p)− ppT satisfy the following properties:
(a) The eigenvalue λ(i) is the i-th largest solution of the following equation:
v(λ) = 1− ∑C
i=1 p2i pi − λ = 0 (9)
(b) The eigenvector q(i) is aligned with the direction of (diag(p)− λ(i)I)−1p
(c) p(i+1) ≤ λ(i) ≤ p(i) for 1 ≤ i ≤ C − 1, and λ(C) = 0
(d) 12Gini(p(1)) ≤ λ (1) ≤ Gini(p(1)) where Gini(q) = 1− q2 − (1− q)2 = 2q(1− q) is the
Gini impurity for the binary case (q, 1− q).
Proof. The eigenvalues λ(i) of M = diag(p) − ppT are the zeros of the following characteristic polynomial:
φM (λ) = det(diag(p)− ppT − λI) (15) = det(diag(p)− λI) det(I − (diag(p)− λI)−1ppT ) (16)
= C∏ i=1 (pi − λ) 1− C∑ j=1 p2j pj − λ  (17) where the second equality follows from A − ppT = A(I − A−1ppT ) with the matrix A = diag(p)− λI , and the third inequality holds because det(I +uvT ) = 1 +uTv for vectors u and v. Then it is equivalent to solving the following equation:
v(λ) = 1− C∑ i=1 p2i pi − λ = 0 (18)
which implies (a). Note that this result also implies (c) as shown in Figure 1.
Next, to prove (b), put A ≡ diag(p) − λI and q ≡ A−1p. Then it is required to show that (M − λI)q = (A− ppT )q = 0 for the eigenvalues λ = λ(i). We have
(A− ppT )q = (A− ppT )A−1p = p− ppTA−1p (19)
Here, (ppTA−1p)i = ∑ j,k pipjA −1 jk pk = ∑ j,k pipjδjk(pj − λ)−1pk = ∑ k pipk(pk −
λ)−1pk = pi ∑ k p 2 k/(pk − λ) = pi. The last equality holds for the eigenvalues λ = λ(i) which follows from (a).
Now, we want to prove the statement (c). Since
λ(i)(C) ≤ λ(j)(A) + λ(k)(B) if k + j − i = 1 (20) λ(i)(C) ≥ λ(j)(A) + λ(k)(B) if k + j − i = C (21)
for C = A+B ∈ RC×C where λ(i)(D) is the i-th largest eigenvalue of a matrixD (Weyl, 1912; Fulton, 2000), we can get λ(i)(C) ≤ λ(i)(A) + λ(1)(B) and λ(i)(C) ≥ λ(i+1)(A) + λ(C−1)(B). Thus, forA = diag(p) andB = −ppT , we can get
p(i+1) ≤ λ(i)(M) ≤ p(i) for 1 ≤ i ≤ C − 1 (22)
since λ(i)(A) = p(i), λ(i+1)(A) = p(i+1) and λ(1)(B) = λ(C−1)(B) = 0. Moreover, since M1 = p− ppT1 = p− p ∑ i pi = 0, the smallest eigenvalue is λ (C) = 0.
Lastly, we prove the statement (d). From the Gershgorin circle theorem (Gershgorin, 1931), we have λ(1) ∈ ⋃ i B(Mii, ∑ j 6=i |Mij |) = B(p(1)(1− p(1)),p(1)(1− p(1))) = [0, 2p(1)(1− p(1))] (23)
which implies λ(1) ≤ 2p(1)(1−p(1)). Note that p(1)(1−p(1)) ≥ p(i)(1−p(i)) since g(t) = t(1− t) is increasing for 0 ≤ t ≤ 0.5. In detail, if p(1) ≥ 0.5, since p(i) ≤ 1 − p(1) ≤ 0.5, we have g(p(i)) ≤ g(1 − p(1)) = g(p(1)). Otherwise (p(1) < 0.5), since p(i) ≤ p(1), it leads to the same inequality g(p(i)) ≤ g(p(1)). With the Rayleigh principle, we can express the largest eigenvalue as λ(1) = max‖u‖2=1 u TMu, and thus e(1)TMe(1) = M(1)(1) = p(1)(1− p(1)) ≤ λ(1).

heading: C PROOF OF THEOREM 2
text: Theorem (restated). For some 0 ≤ λ∗ ≤ λ(1) for each x ∈ D, we can bound
‖〈JMJT 〉‖σ = 〈λ∗‖J‖2〉 ≤ 〈λ(1)‖J‖2〉 (11)
Proof. We start with the Rayleigh principle:
‖〈JMJT 〉‖σ = max ‖q‖=1 qT 〈JMJT 〉q = max ‖q‖=1 〈qTJMJTq〉 (24)
SinceM = ∑ i λ (i)q(i)q(i)T , we can continue by putting v = JTq, and then
Eq(24) = max ‖q‖=1 〈vTMv〉 = max ‖q‖=1 〈 ∑ i λ(i)(q(i)Tv)2〉 (25)
Then by putting λ̃ = ∑ i γ (i)λ(i) with γ(i) = (q(i)Tv)2/ ∑ i(q (i)Tv)2 ≥ 0 ( ∑ i γi = 1),
Eq(25) = max ‖q‖=1 〈λ̃ ∑ i (q(i)Tv)2〉 = max ‖q‖=1 〈λ̃ ∑ i (q(i)TJTq)2〉 (26)
= max ‖q‖=1 〈λ̃ ∑ i (q(i)T (JTq))2〉 (27)
Since {q(i)}Ci=1 is an orthonormal basis of RC (eigenvectors of a symmetric matrixM ), we have the following by putting λ∗ = ‖J
T q∗‖2 ‖J‖2 λ̃ with q ∗ = arg max‖q‖=1〈λ̃‖JTq‖2〉,
Eq(27) = max ‖q‖=1
〈λ̃‖JTq‖2〉 = 〈λ̃‖JTq∗‖2〉 = 〈λ∗‖J‖2〉 (28)
Because of the definition of λ∗ and λ̃ with ‖q∗‖ = 1 and ∑ i γ (i) = 1 (γ(i) ≥ 0), we have
λ∗ ≤ λ̃ ≤ λ(1) (29)
which leads to the final inequality in Eq (11). The equality holds when ‖JTq∗‖ = ‖J‖, γ(1) = 1 and γ(i) = 0 (i 6= 1) for all x ∈ D.

heading: D GRADIENT DESCENT IN THE TOP HESSIAN SUBSPACE
text: Gur-Ari et al. (2018) showed that the gradient of the loss quickly converges to a tiny subspace spanned by a few top eigenvectors of the Hessian after a short training. Then, the top Hessian subspace does not evolve much, which implies gradient descent happens in a tiny subspace. However, the underlying mechanism has not been fully understood.
Direction of q(i) and two salient elements We investigate the direction of the eigenvector q(i) (1 ≤ i ≤ C − 1) ofM . The eigenvector
q(i) = α
( pj
pj − λ(i) ) j
(30)
can be obtained from Theorem 1 (b) for some α > 0. Here, the magnitude of the denominator |pj − λ(i)| is small for the two indices j = (i), (i+ 1), and is large for the others. This is because the eigenvalue λ(i) lies between p(i+1) and p(i) (Theorem 1 (c)). Therefore, the eigenvector q(i) has a relatively large positive value in q(i)(i) and a large negative value in q (i) (i+1) compared to the other components.
Figure 8 shows the directions of q(i) with the heatmap of the matrix QΛ1/2 where QΛ1/2 = [ √ λ(1)q(1); · · · ; √ λ(C−1)q(C−1); 0] ∈ RC×C for M = QΛQT . As expected, considering each column of QΛ1/2, the eigenvector q(i) is colored in red (+) at q(i)(i) and in blue (−) at q (i) (i+1) for 1 ≤ i ≤ C − 1. The two salient elements are highlighted with the dashed boxes.
Direction of Jq(i) and margin maximization In light of the previous discussion, the direction of Jq(i) = (Jq(i))|θ=θ(t) = ∇θ ( q(i)(θ(t))Tz(θ) ) ∣∣ θ=θ(t) ∈ Rm (31)
is approximately a direction maximizing q(i)(i)z(i) + q (i) (i+1)z(i+1) at the current parameter θ (t) ∈ Θ because the other terms are relatively small. In other words, it tends to maximize the margin z(i)− z(i+1) in the logit space Z between the two classes (i) and (i+ 1) (see Figure 8). In particular, Jq(1) is approximately a direction that maximizes the margin between the most likely class and the second most likely class.
Clustering of Jq(1) and the most probable class We first define following subsets of the training set according to the most probable class (and the second most one): Di = {x ∈ D : c1(x) = i} ⊂ D and Dij = {x ∈ D : c1(x) = i, c2(x) = j} ⊂ D for i 6= j ∈ [C]. Note that Di = ⋃ j 6=iDij . Given two examples from Dij , their Jq(1) are expected to be highly aligned to each other. This is because the direction of Jq(1) is approximately a direction of maximizing the margin and of learning the features to discriminate the class i from the class j. Moreover, two examples from Di also have highly-aligned Jq(1). Figure 9 (Left) shows the concentration of the directional data of Jq(1) from Di. We also compute the mean resultant length (MRL) to measure the concentration. The MRL ρ of the directional variable V ∈ Sm−1 ≡ {v ∈ Rm : ‖v‖ = 1} defined as ρ ≡ ‖E[V ]‖ ∈ [0, 1] indicates how V is distributed (the higher, the more concentrated). Now, we focus onm ≡ √ λ(1)Jq(1) as the other λ(i)-terms are dominated by the λ(1)-term after a few epochs (see Appendix J for details). Then, we follow a similar approach from (Papyan, 2019) and provide the following equation:
〈mmT 〉 = ∑C
i=1 γi〈mmT 〉Di = ∑C i=1 γi(m imiT + 〈(m−mi)(m−mi)T 〉Di) (32)
where γi = |Di|/|D| and mi = 〈m〉Di . Here, the covariance term 〈(m −mi)(m −mi)T 〉Di is weak asm is concentrated withinDi, and thus we can roughly approximateG with ∑C i=1 γim
imiT . This implies that the top eigensubspace of the Hessian highly overlaps with the at mostC-dimensional subspace spanned by {mi}Ci=1. Figure 9 (Right) demonstrates that the mean vectors {mi}Ci=1 are well separated from each other. This also implies the outliers in the Hessian spectrum (Sagun et al., 2016; 2017).
Why gradient descent happens mostly in the top Hessian subspace? Given input x, after the model becomes to correctly predict the true label y, the gradient descent direction −g = J(ey − p) used in the training tends to be highly aligned with Jq(1). This is because ey − p and q(1) both have similar direction. They have positive values 1− py and q(1)(1) in y(= (1))-th element, negative value −pi and q(1)i in the others, and especially large negative value for the second most probable class i = (2). Figure 10 (Middle) shows the cosine similarity between the gradient descent direction −g and Jq(1). As expected, they are highly aligned with the cosine similarity near 1 as the two vectors ey − p and q(1) become more aligned to each other.
Next, we move on to the subspace S ≡ span({mi}Ci=1) spanned by {mi}Ci=1. As each m =√ λ(1)Jq(1) is highly aligned with −g, it is reasonably expected that the total gradient gD lies in the subspace S. To measure how much the vector v 6= 0 is aligned with the subspace S, we define the cosine similarity between the vector v and its projection PS(v) onto the subspace S as follows:
AS(v) ≡ cos(v, PS(v)) = ‖PS(v)‖/‖v‖ (33)
In particular, AS(v) = 0 means v ∈ S⊥ and AS(v) = 1 means v ∈ S. Figure 10 (Right) shows the high alignment of the total gradient gD in the subspace S ≡ span({mi}Ci=1) although the subspace S is of dimension at most C m. Moreover, since G can be roughly approximated by∑C i=1 γ imimiT , the top eigenvector r of the HessianH is also highly aligned with the subspace S .

heading: E EXPERIMENTAL SETTINGS
text: E.1 DATA
We use the CIFAR-10 dataset (Krizhevsky et al. (2009), https://www.cs.toronto.edu/ ~kriz/cifar.html) and the MNIST dataset which have C = 10 number of classes. We also conduct some experiments on the CIFAR-100 dataset with the number of classes C = 100. We mostly do not use the data augmentation for training (1) not to introduce the randomness in the training loss and (2) to allow the training loss to converge to a small value. In case of the WRN-28-10 (Zagoruyko & Komodakis, 2016), we use the data augmentation to improve the performance.
E.2 NETWORK ARCHITECTURES
We use the following models: VGG-11 (VGG) (Simonyan & Zisserman, 2015) without batchnormalization, VGG for CIFAR-100 (VGG-CIFAR-100), ResNet-20 (ResNet) (He et al., 2016) wihtout batch-normalization, a 6-layer CNN (6CNN), SimpleCNN used in Jastrzebski et al. (2021) (SimpleCNN) two 3-layer fully-connected networks (3FCN-CIFAR and 3FCN-MNIST), and WRN-28-10 for CIFAR-10/CIFAR-100 with the number of model parameters, m = 9750922, 9797092, 268346, 511926, 361706, 656810, 199210, 36479194, 36536884, respectively.
We use a modified version of the implementation of VGG-11 from https://github.com/ chengyangfu/pytorch-vgg-cifar10/blob/master/vgg.py without the dropout layers and ResNet-20 from https://github.com/locuslab/edge-of-stability/ blob/github/src/resnet_cifar.py. We change the last linear layer for the CIFAR-100 dataset. The 6CNN model can be expressed in the Pytorch code as follows: nn.Sequential(
nn.Conv2d(3, 32, 3, stride=1, padding=1, bias=False) nn.ReLU(), nn.Conv2d(32, 32, 4, stride=2, padding=1, bias=False) nn.ReLU(), nn.Conv2d(32, 64, 3, stride=1, padding=1, bias=False) nn.ReLU(), nn.Conv2d(64, 64, 4, stride=2, padding=1, bias=False) nn.ReLU(), nn.Flatten(), nn.Linear(4096, 100, bias=True), nn.ReLU(), nn.Linear(100, 10, bias=True),
)
and the 3FCN architecture is as follows: nn.Sequential(
nn.Flatten(), nn.Linear(n, 200, bias=True), nn.ReLU(), nn.Linear(200, 200, bias=True), nn.ReLU(), nn.Linear(200, 10, bias=True),
)
where n=784 for 3FCN-MNIST, and n=3072 for 3FCN-CIFAR (the same one used in Cohen et al. (2021)).
E.3 HYPERPARAMETERS
SGD (Robbins & Monro, 1951) with the learning rate η can be expressed as follows:
θ(t+1) = θ(t) − ηgB (t) (θ(t)) (34)
where θ(t) is the model parameter, B(t) ⊂ D is the training batch at t-th step, g = ∇θl, and gB = 〈g〉B. We mostly use the simplest form of (S)GD as described in Eq (34) without momentum,
weight decay, and learning rate decay. The only exception is when we train WRN-28-10, we used the momentum of 0.9, the weight decay of 0.0005, and the learning rate decay of 0.2 in [30%, 60%, 80%] of the whole training epochs.
E.4 DETAILED SETTINGS FOR EACH FIGURE
For Figure 2, 9, and 10, we ran the experiments on the CIFAR-10 dataset using 6CNN and trained the model using GD with the learning rate η = 0.04. See Appendix H for the setting used in Figure 6. We also plot variants of the figures in the main paper for other settings (training steps, data, network architectures, and hyperparameters) in Appendix J-L.
E.5 HESSIAN
When computing the top eigenvalue ‖H‖σ and the corresponding eigenvector r of the HessianH , we use the tool developed in PyHessian (Yao et al. (2020), https://github.com/amirgholami/ PyHessian, MIT License) based on power iteration method using a small subset (5-25%) of training dataset D.
E.6 POWER ITERATION ALGORITHM
Even though we have the secular function v(λ) in Eq (9) and the algorithms for computing the eigenvectors (Bunch et al., 1978), we use the power iteration in Algorithm 1 to get the top eigenvalue λ(1) of the logit Hessian M ∈ RC×C since we can run the algorithm for a mini-batch in parallel. Then, we can compute the corresponding top eigenvector q(1) from Theorem 1 (b). To compute the second largest eigenvalue, we apply the power iteration toM ′ ≡M − λ(1)q(1)q(1)T instead ofM after computing λ(1) and q(1).
Algorithm 1 Power iteration Input: matrixM , maximum iteration nmax, tolerance bound Output: the spectral norm λ(1) = ‖M‖σ of the matrixM Initialize u ∈ RC with a random vector. i← 0 repeat v ←Mu/‖Mu‖ u←MTv/‖MTv‖ i← i+ 1
until it converges within the tolerance bound or i ≥ nmax return vTMu
We also compute the operator norm of the Jacobian ‖〈J〉‖ with the power iteration as in 2. It requires(C + 1)-times scalar function differentiations with respect to θ for each iteration.
Algorithm 2 Power iteration for Jacobian Input: logit function zθ (not matrix 〈J〉), maximum iteration nmax, tolerance bound Output: the operator norm ‖〈J〉‖ of the Jacobian Initialize u ∈ RC with a random vector. i← 0 repeat
Compute 〈J〉u = ∇θ〈uTz〉 (1× scalar function differentiation) v ← 〈J〉u/‖〈J〉u‖ Compute 〈J〉Tv = (vT 〈J〉)T = [vT∇θ〈z1〉, · · · ,vT∇θ〈zC〉]T (C× scalar function differentiations) u← 〈J〉Tv/‖〈J〉Tv‖ i← i+ 1
until it converges within the tolerance bound or i ≥ nmax return vT (〈J〉u)

heading: F THE TENDENCY OF THE JACOBIAN NORM TO INCREASE
text: The weight norm ‖θ(t)‖ increases (Figure 12 (the third one)) in order to increase the logit norm ‖z(θ(t))‖ (Figure 12 (the leftmost figure)) and to minimize the cross-entropy loss during training (Soudry et al., 2018). This also leads to the increase in the layerwise weight norms (Figure 14) and the Jacobian norm (Figure 6). While the Jacobian norm tends to increase, it stops increasing at a certain point, stays for a few steps, and increases again. We ran the experiments on the CIFAR-10 dataset and 6CNN using GD with learning rate η = 0.04.
G ‖H‖σ ∝ 〈λ(1)〉‖〈J〉1̂‖2
Surprisingly, we empirically observed that ‖H‖σ ∝ 〈λ(1)〉‖〈J〉1̂‖2 during training where 1̂ = 1/ √ C ∈ RC . Since it requires further theoretical grounding, we left it as future work. We observe this relation for a variety of learning rates (Figure 15), network architectures (Figure 16 and 17), batch sizes (Figure 18 and 19), and datasets (Figure 20 and 21). At least, ‖H‖σ and 〈λ(1)〉‖〈J〉1̂‖2, they increase and decrease together. We emphasize that Figure 21 shows the case when the sharpness did not reach the limit 2/η. This is because the learning rate is relatively low and it is easy to train a model for the MNIST dataset, and thus 〈λ(1)〉 decreases before the sharpness reaches the limit.
H IMPLICIT REGULARIZATION ON ‖〈J〉1̂‖2
We ran the experiments for Figure 2 on the CIFAR-10 dataset. We used 6CNN and VGG for Figure 2 (Left) and Figure 2 (Right), respectively.
We further investigate the relation ‖H‖σ ∝ 〈λ(1)〉‖〈J〉1̂‖2 in the previous section and its effect on ‖〈J〉1̂‖2. For the early phase of training, it is hard to see the initial rapid growth of the sharpness in this smoothed curves and when exactly the regularization begins to activate. We refer the readers to the previous section, Appendix G, for the fine-grained analysis of the early phase of training. We provide plots with different settings. Again, we observe that training with a larger learning rate and a smaller batch size limits ‖〈J〉1̂‖2 with a smaller value (dotted red lines) in the Active Regularization Period. Curves are smoothed for visual clarity.

heading: I FIGURE 4 (VISUALIZATION OF THE OPTIMIZATION TRAJECTORY)
text: We use UMAP (McInnes et al., 2018) to visualize the optimization trajectory {θ(t)}t∈[T ] ⊂ Θ ⊂ Rm in a 2D space. As GD enters into the Edge of Stability (Cohen et al., 2021), it oscillates in a direction nearly orthogonal to its global descent direction (Xing et al., 2018). Note that GD may not enter the Edge of Stability as shown in Figure 26 (Bottom).

heading: J FIGURE 8 (QΛ1/2)
text: Figure 27 shows the matrixQΛ1/2 (Figure 8) for different training steps. Figure 8 and Figure 27 (Top Right) are plotted at the equivalent step (t = 1000). Figure 27 demonstrates that λ(1) becomes more dominant than the others as training progresses. The argument that there are two salient elements in each q(i) in its (i)- and (i+ 1)-th elements is empirically shown to be valid throughout the training.
K FIGURE 9 (CLUSTERS OF m AROUND EACH mi)
Figure 30 shows Figure 9 (Left) for some different class pairs (i, j) with negative cosine similarity, i.e., cos(mi,mj) < 0. We use the model trained for t = 1000 steps. Figure 28 shows Figure 9 (Right) for different training steps. Figure 29 shows the evolution of the Mean Resultant Length (MRL) ρ in Figure 9 (Left) during training.

heading: L FIGURE 10 (GRADIENT DESCENT HAPPENS MOSTLY IN THE TOP HESSIAN SUBSPACE)
text: Figure 31 shows similar results with different settings with VGG and learning rate η = 0.08.

heading: M ANALYSIS OF THE MSE LOSS
text: In the main text, we focus on the cross-entropy loss. Here, we briefly analyze the MSE loss, l = 1 2‖z − e
y‖2. Then, we have M = ∇2zl = I , λ(1) = ‖M‖σ = 1 and G = 〈JJT 〉. It leads to the same conclusion as in Theorem 2:
‖G‖σ = ‖〈JJT 〉‖σ ≤ 〈‖JJT ‖σ〉 = 〈‖J‖2〉 (35)
We empirically observed that ‖H‖σ ∝ ‖〈J〉‖2 as shown in Figure 32.

heading: N DETAILS OF EXPLICIT JACOBIAN REGULARIZATION (EJR)
text: We first propose a simple form of EJR with the regularized loss as follows:
L̃(θ) = L(θ) + λreg‖〈J〉‖2F /C (36)
and update the model parameter as
θ(t+1) = θ(t) − η(∇θL(θ(t)) + λreg∇θ‖〈J〉‖2F /C) (37)
However, this requires to build a computational graph for ‖〈J〉‖2F which is inefficient for a large network (e.g. WRN-28-10).
To this end, we propose two efficient variants of EJR. First, we propose to update the model parameter as follows:
θ(t+1) = θ(t) − η∇θL(θ̂(t)) (38)
where θ̂ = θ + ρ̃reg〈J〉u and ρ̃reg = ρreg/‖〈J〉u‖ as in Foret et al. (2021); Liu et al. (2019). Second, we propose another variant using L̂ instead of L in Eq (38) as follows:
L̂ = L+ µregu T 〈z〉 (39)
We can approximate L̂ as follows:
L̂(θ̂) ≈ L̂(θ) + (∇θL̂(θ))T (ρ̃reg〈J〉u) (40) = L̂(θ) + (∇θL(θ))T (ρ̃reg〈J〉u) + (∇θµreguT 〈z〉)T (ρ̃reg〈J〉u) (41) = L̂(θ) + (∇θL(θ))T (ρ̃reg〈J〉u) + (µreg〈J〉u)T (ρ̃reg〈J〉u) (42) = L̂(θ) + (∇θL(θ))T (ρ̃reg〈J〉u) + µregρreg‖〈J〉u‖ (43) ≈ L̂(θ) + (∇θL(θ))T (ρ̃reg〈J〉u) + µregρreg‖〈J〉‖F / √ C (44)
We used the first-order Taylor expansion of L̂(θ + ρ̂reg〈J〉u) in Eq (40). We expect additional effect of minimizing µreguTz + µregρreg‖〈J〉‖F / √ C ≈ µregρreg‖〈J〉‖F / √ C compared to the first variant.

Review summaries:
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. 

The authors study the largest eigenvalue and eigenvector of the Hessian of the loss function. The authors approximate the Hessian matrix by a low-dimensional matrix, which is a rank-one modification of a diagonal matrix. The eigendecomposition helps to explain how the sharpness influences the gradient descent method and the generalization properties. They show that GD has implicit regularization effects on the Jacobian norm weighted with the impurity of the probability output, which is also related to the Fisher information matrix.