Implicit Jacobian regularization weighted with impurity of probability output
Authors: Sungyoon Lee
Jinseong Park
Jaewook Lee
Venue: ICLR 2022 Submitted (2022)
Keywords: deep learning
gradient descent
implicit bias
implicit regularization
Hessian
sharpness

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 sharpness-impurity-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.

Under review as a conference paper at ICLR 2022
IMPLICIT JACOBIAN REGULARIZATION WEIGHTED WITH IMPURITY OF PROBABILITY OUTPUT
Anonymous authors Paper under double-bling review
ABSTRACT
Gradient descent (GD) plays a crucial role in the success 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 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 dower-dimensional matrices. In particular, to analyze the sharpness, we instead explore the eigenvalue problem for the low-dimensional matrix which is rank-one modification of diagonal matrix. The eigendecomposition provides a simple relation between the eigenvalues the low-dimensional matrix and the impurity of the probability output. We exploit this connection to derive sharpness- impurity-Jacobiar relation and to explain how the sharpness influences the learning dynamics and the generalization performance. In particular, show that GD has implicit regularization effects on the Jacobian norm weighted with the impurity of the probability output.
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 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-bast 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; Jastrzebski et al.) 2019 Ghorbani et al. 2019 Liu et al. 2020
Lewkowyez al. 2020 Cohen al. (021)) and generalization (Hochreiter & Schmidhuber 1997
Keskar et al. [011 Dinh et al. 2017 Jastrzębski er al., 2017 Wang et al., 2018 Chaudhari et al. 2019
Fort et al. 2019 Jiang al. [2020] Barrett & Dherin 202 Smith 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 the beginning, and after reaching a certain value, training dynamics becomes unstable, oscillating along the top eigenvector (Jastrzębski et al. 2019 Cohen 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 Dastrzebski et al. 2020 Lewkowyez et al., 2020 Jastrzebski al., (02T). However, the Hessian of a deep neural network is very sigh-dimensiona which makes it difficult to analyze its eigensystem. Recently, some researchers studied the Hessian by exploiting tools randomized numerical linear algebra (Sagun et al. 2017 Papyan 2018 2019 Ghorbani al. 2019 Yao et al. 2020) and decomposition of the Fort & Ganguli 2019).
Hessian (Papyan
In this paper, we present a decomposition of the Hessian using eigendecomposition of low- dimensional 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

---

Under review as a conference paper at ICLR 2022
enables 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 follows:
We decompose the Hessian with low dimensional matrices, the logit Hessian and the logit- weight Jacobian (defined in Definition D. and investigate the Hessian by the eigende- composition 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 relation.
Jacobian. We call it
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 penalizing the Jacobian norm (Implicit Jacobian Regularization) in 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 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 the limit vanishing learning rate (Mandr et al. 2017 Li et al. 201 7b a [mith & Lej 2018 Chaudhari & Soatto 2018
Jastrzębski et al.) 2017] Zhu et al.) 2019 Park et al. 2019). However, some theoretical concerns have been raised for such approximations (aida 2019). Moreover, Barrett & Dherin (021) argued that the SDE analysis in the limit of vanishing learning rate cannot explain the generalizatio benefits of finite learning rates and they proposed modified gradient flow for finite learning rates. However, they still consider a continuous gradient flow and thus it has fundamental limitation in explaining the oscillatory behavior with iterative catapult in 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. (0188) 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 high valley height, and small batch size induces gradient stochasity. They both help exploration through the parameter space with different roles in the training dynamics. Jasttzębski 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) 8). Lewkowyez et al. (2020) investigated simple theoretical models with a solvable training dynamics. They showed that, in their setup with 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 regime
2

---

Under review as a conference paper at ICLR 2022
called the Edge of Stability where the sharpness can no longer increase and stays near certain value, and the training loss behaves nonmonotonically but decreases globally. This behavior of the optimization at the Edge of Stability can be seen repeated catapult mechanisms. They explicitly marked the limit of the sharpness with 2/1 (1) =learning rate).
To describe the aforementioned evolution of the sharpness, Fort & Ganguli (2019) developed 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 explain the underlying mechanisms.
Implicit bias in SGD There have been many studies on the implicit bias in SGD (Neyshabur) 2017;
Zhang al. 202T; Soudry et al. 2018). We review the most relevant and recent ones. Dastrzebski 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); [mmth 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 each case.
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., V,
Ou.
-> T and E TR.. It is also generalized to the cases of scalar, Rexu where u T
ij
M
1.
We consider problem of learning a C-class classifier which maps an input æ E X' C IRdd to target label y E [C] where [C] 21, .C). To this end, we build a parameterized model fo : X Z C RC with a model parameter 0 E θ C RT which outputs a logit vector = = fa( E Z C R (we often omit the dependence on the input x and the model parameter Ø). Then, the logit vector is given input to the softmax function yield a probability vector p softmax(
C
where A C-1 = {p E [0, 1 = 1, 2 0}. We want the model to match the most probable class e₁ the true label y. where c(x>) arg (p)) in descending order. We exchangeably denote the probability value corresponding to the true label y as p = Py E [0, 1]. The cross cross-entropy loss, 1= ((z.)) y) IR, is equivalent to the negative log-likelihood for the Euclidean ε--norm of vector or for the Euclidean operator norm of a matrix (equivalently, ||- I-ll for a square matrix), || F| for the Frobenius norm, and tr(.) for the trace of a (square) matrix.
log p. We use the notations ||
Starting with simple computation of the derivatives of the softmax function, Eq (see Appendix we can easily derive the following equations in order:
V=P = diag(p) - pp" E CCC
VzP = [V=P). = p(e" - pp) E C
(2)
V2l = V=P Op p(e" e² - -p) - = --)---p--- ed E C
Ø1
1
(3)
p
V2I = (V (p --")
e²) = diag(p) - pp" E CCC
(4)
where diag(p) = (b.p E CCC is a diagonal matrix with p as its diagonal entries, and e (des)s E RC is one-hot vector with i-th element as 1.
Next, the Hessian of the loss function for given example with respect to the model parameter can be expressed as follows:
Vole VozV + +
j=1 Vo VozV R mmxm
(5)
using a well- Gauss-Newtor approximation. (see, for example, Schhaudollhh 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. (Z) which yields VL. (VI) and 2 = (V21) where (+) is the expectation over
3

---

Under review as a conference paper at ICLR 2022
the empirical measure of the training set D, equivalently say Epl.]. We use the notation ()s when averaging over subset S. Following from Eq ④ and Eq (5). we define the Hessian matrix for the total loss and its Newton approximation matrix G with the matrices M and J as follows:
Definition 1. We call M the logit Hessian. J the Jacobian (of the logit function with respect to the model parameter), H the Hessian, and G the Newstnn approximation defined as follows:
M = 122 diag(p) -pp" E R
(6)
J = Jă) Voz E g
(7)
H (Va() (JMJ") G R
mm
It is interesting note that while 1 is dependent the true label y, the logit Hessian M V is independent of y, and so are J, JJMJ" and G. In case of the MSE loss - - e²1 we have M = V = I and G = (JJ"). We mainly focus on the usual cross-entropy loss and defer the investigation on the MSE loss to Appendix M From Eq ( we will often use the approximation
|H R |Gl|l as justified in (Saunn et al. (017);; Fort & Ganguli 2019), but this approximation
sometimes fails in the later phase training when the top eigenvalues of the Gauss-Newton matrix
is not sufficiently isolated from the bulk near 0 (Papyan 2018 Thus mainly focus on the carly
phase of training.
DECOMPOSITION OF THE HESSIAN H
In the previous section, we introduced the Gauss-Newtor approximation G of the Hessian H, and decomposition of G with the Jacobian J and the logit Hessian M, i.e., G = (JMJ")) Now, we focus on the logit Hessian matrix M and eigendecomposition, estimate the top eigenvalue of M with upper/lower bounds, and explore the evolution of the top eigenvalue during training.
4.1 COSANECCOMPOSTTION OF THE LOGIT HESSIAN
The lower-dimensional matrix M E CxC is simple and fully characterized by only the probability vector p M diag(p) pp" in Eq (6). but it turns out be important for understanding the much higher-dimensional matrix G E Ru (C < m). Since M = diag(p) - pp" is rank-one modification of 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 Hessian M can be eigendecomposed as M = QAQ" = E ( (3)2 where ( the i-th largest eigenvalue of M and ( is its corresponding normalized eigenvector. For simplicity, we also use the same ordered index of (i) [C] with parentheses for the probability output p A i.e., C (i) and P(1) P(2) 2 ... 2 P(c) 0, because this ordering related the eigenvalues [ demonstrated in the following theorem. We defer the proof to Appendix
J==1 as
0.5
0.5
0.5
1
2
10
-0.01
0.00
0.02
0.10
0.20
0.30
0.35
5.0
2.5
-5.0
-0.1
0.0
0.1
0.2
0.3
0.4
Figure 1: [Eigenvalues of the logit Hessian] Graph of the secular function v(X) in Eq which has zeros at the eigenvalues [ P(s) with red vertical lines. It illustrates Theorem (a) and (c).
/r of M = / (blue curves). We highlighted the singularities
4

---

Under review as a conference paper at ICLR 2022
Theorem 1 (ef. Golub (173); : Bunch al. 1978 ). The eigenvalues X (A(1 2 2 2 ... 1(0))
and the corresponding normalized eigenvectors q of the logit Hessian M V22l diag pp satisfy the following properties:
(a) The eigenvalue ( is the i-th largest solution of the following equation:
u(A) =1 p
(9)
(b) The eigenvector (8) is aligned with the direction of (diag(p) - 1)
(c) P(++1) A00 for SisC. and ((C) =0
(d) ((n)( S 10) I Cini(p(1)) where Gini(q) = - 9² (1 1)- 20(1 q) is the Gini impurity for the binary case (4, 1)-
Figure illustrates the secular function vAA defined as Eq ( in Theorem (a). The function v(A)
has singularities at the probability values (P() 1e and zeros at the eigenvalues {x satisfying
Theorem (c). Moreover, Theorem (c) and (d) provide the upper/lower bounds the eigenvalues
of M. Especially, the top eigenvalue (1) is bounded by Gini((pa) S A < (i))( and thus we call it impurity.
4.2
EVOLUTION OF IMPURITY
0.5
5000
0.3-
0.4
0.3
3000
0.2
0.2
2000
0.1
0.1
upper bound
1000
lower bound
0.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
4000
Step
[Evolution of Impurity] The impurity (2) increases and then decreases as P increases during training. Left: The impurity of typical example is plotted against P(1)- We together plot the upper bound min Gini(P(1)) and the lower bound 1 (i)i(P(1)) from Theorem (c) and (d). Right: The impurity is plotted against the training step. Blue curve indicates its mean value (1) and sky-blue area shows the 25-75% quantile range of the impurity for the training data.
Figure 2:
2
We explore the top eigenvalue (1) of M (also referred as impurity) during training. Figure demonstrates the evolution of the impurity during training, which increases the beginning and then decreases in the later phase. We trained model to zero training loss, and thus, for most examples, the probability Py for the true class eventually becomes the highest probability P(1). As the top probability P(1) increases from 1/C to during training, the impurity starts from (() x J
Il Gini àl []= 1 }] (Theorem I (c) and (d)), increases at the initial phase of the training
because lower bounded by | ((()) = P(1) (1 P(1)) which increases for P(1) E [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(Pa)) 2P(1) (1 P(1)) which decreases for P(1) E [0.5, 1]. Note that Cohen et (2021) tried estimate similar value, but they use Py, not P(1)- We again emphasize that M Is independent of the label y but dependent only on the probability output p.
5 IMPLICIT JACOBIAN REGULARIZATION
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
5

---

Under review as a conference paper at ICLR 2022
influences the learning dynamics and the generalization performance. Detailed experimental settings for each Figure are described in Appendix
||M||
||
||Í||
||Í|]
100
800
100
100
700
2000
600
60
500
60
60
1500
1000
400
300
750
1000
200
500
20-
20
20-
500
100
250
40
80
120
160
200
1000
1040
1080
1120
1160
1200
1500
1540
1580
1620
1660
1700
Step
Step
Figure 3: The sharpness H and the Jacobian norm []] show similar oscillating behavior up
to factor X which is locally constant and slowly changes during training (CIFAR-10. 17 0.04,
[3 = 128; left: 0-200, middle: 1000-1200, right: 1500-1700 steps). We highlighted | 1- = 2/n
with the dashed horizontal line. Note that they have different right y-axes.
5.1 SHARRNESSIMPPRITYYACCOOM
RELATION
We first take closer look at the sharpness of the loss landscape during training and build a relation between the sharpness | ||. the impurity (2), and the Jacobian J. Since the Gauss-Newton matrix G is known to approximate the true Hessian H well, especially for the top eigenspace (Sagun et al.
2017
Fort& & Ganguli, 2019 Papyan 2019), we write the sharpness H ll as follows:
H
||G|| - IKJMJ*))l
(10)
It implies the impurity || and the squared norm of the Jacobian J are highly correlated with the sharpness H |lr as demonstrated in the following theorem. We defer the proof to Appendix
Theorem 2. For some 0 & X" & x for each x: E 'D we can bound
JJMJ (A 1 < (11131113)
(11)
In other words, >+||J[ 12 has the expected value of approximately | oo. and we briefly state this
relation with > |- ~ || Ia- Here, X in X J acts as an adaptive regularization weight for the regularization on the Jacobian norm || as detailed in the following section. Since it is computationally inefficient track [| for every æ E D, we instead investigate [I? We expect we observe that H l| and [ show similar
|| lo {+) for some X= In Figure 3
oscillating behavior up to factor X which is locally constant and slowly changes during training.
5.2 IMPLICIT JACOBIAN REGULARIZZTTON
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 | 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 [| 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 begin- ning, the width of the valley of the loss landscape becomes narrower than the discrete step size the cradient-based optimization. After the sharpness reaching this threshold, the iterate starts to bounce

---

Under review as a conference paper at ICLR 2022
530
100
123
1n4-
430
200
25-
12-
-10
30
120 125
15.0
175
20.
225
25.
Figure 4: Oscillatory catapult in the optimization trajectory {0(0} (from blue to red) of full- batch GD (MeInnes et al.) 2018). Left: UMAP of the model parameters trained on CIFAR-10 for the first 500 steps. Right: Zoom-in for the oscillatory steps [100, 300]. After few steps (~100), the sharpness reaches a threshold and the iterate shows the oscillatory behavior with iterative catapult.
off from one side of the valley to the other, and repeats this (Xing et al. 2018 Jastrzębski et al. 2019). Figure shows this oscillatory behavior with iterative catapult after the sharpness reaching the threshold, using UMAP (MeInnes 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 (Lewkowyzz et al. 2020). Figure 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 the threshold. Therefore,
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/n y (n 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 [ (Right). This oscillatory catapult and the plateau of the sharpness are attributed the discrete dynamics of the gradient-basse based optimization with finite learning rate and cannot be described by continuous gradient flow.
Implicit Jacobian Regularization (IJR) Due to this catapult effect, the Jacobian norm cannot continue increase. In other words, the gradient- based optimization implicitly penalizes the Jacobian norm since | H| タンー This IJR effect begins after the sharpness reaches the
пA* loosened) with increasing P(1) 0.5 in the later phase (see Figure 2 and 6) because X A 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 (11 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).
threshold. And the regularization effect diminishes as (!) decreases (the upper limit 2
test acc
train
|||||
test loss
500
15000
0.30
-0.6
400
Active
111
12500
jacolblan Norm
0.25
-0.5
Sharpress
Aegularization
10000
0.20
1.5
300
Peried
7500
0.15
Impartt
200
1.0
5000
0.10
0.2
100
2500
0.1
0.5
0
12000
16000
20000
0.000
0.0
1.00
Stop
Figure 5: Three phases of Implicit Jacobian Regularization (LJR), It shows the IJR effects in the Active Regularization Period (II).
7

---

Under review as a conference paper at ICLR 2022
Table Explicit Jacobian Regularization (EJR) enhances the test accuracy in various settings. We report improvement (AAcc.) and Error Reduction Rate (ERR) on CIFAR-IOCCIFAR-I0 when trained with EJR (+EJR), compared the standard training (Baseline).
Network Architecture
Regularization
Test Accuracy
ÅAcc.
ERR
Dataset
Batch Size
Coeflicient
Baseline
+EJR
(%p)
(%)
0.003
66.71
75.40
+8.69
26.10
0.01
67.88
75.62
+7.74
24.10
0.03
Areg 0.01
69.83
75.53
+5.70
18.89
0.1
70.33
74.29
+3,96
13.35
SimpleCNN
0.3
69.34
73.47
+4.13
13.47
CIFAR-I0
0.01
66.81
74.43
+7.62
22.96
50000
0,03
67.72
74.31
+6.59
20.42
Area 0.001
(full-hatch)
0.1
67.53
73.69
+6.16
18.97
0.3
61.08
72.15
+11.07
28.44
WRN-28-10
128
0.1
Preg
96. 10=ous
97.07
+0,97
24.87
Preg 2, Mreg 0.03
97.38
+1.28
32.82
CIFAR-100 WRN-28-10
0.1
Preg
80.69
83.73
+3.04
15.74
Figure 5 shows that the gradient-bace based optimization has implicit regularization effects on the Jacobian
norm a certain period. To be specific, there are three phases that the Jacobian norm (I) initially rapidly increasing before the sharpness reaches near the threshold, (II) actively regularized with gentle slope, and (III) again exponentially increasing as the regularization effect diminishes (as the regularization weight X S X decreases) with the slope being gradually steeper. We call the second phase Active Regularization Period.
||M||
11001
|
11011
[4/1)
300
10000
0.30
500
0.30
14000
8000
400
12000
Shargnns
200
6000
1..00
Stargness
300
10000
Lombien Nori
0.20
150
0.15
0.15
4000
200
6000
4000
50
2000
0.05
100
2000
0.05
0
1.00
0.00
5000
10000
15000
Step
Step
Figure 6: The IJR effects vary depending on the hyperparameters used in the training. The plateau of the sharpness and the n-shaped evolution of the impurity are clearly observed as expected. We used SGD (Left) with fixed batch size 128 and different learning rates 1/ = 0.01/0.02/0.0 and (Right) with fixed 7 = 0.01 and different batch sizes 128/64/32 (solid/dashed/ddoted lines) on CIFAR- 10. Training with a large learning rate and a small batch size (dotted line) penalizes the Jacobian norm more strongly with lower limits of the sharpness. Curves are smoothed for visual clarity.
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 (Jastrzębski al., 2019: Lewkowyez al. 2020 Cohen
let al. 202T). ). As expected, GD with a large learning rate limits the sharpness with a lower value of 2/7. The large learning rate encourages the stronger implicit regularization on the Jacobian norm. Figure shows, comparing the three red lines (solid/dashed/dootee) with different learning rates (0.01/0.02/0.033 and batch sizes (128/64/32), training with a larger learning rate and 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 IJ, we aim to explicitly regularize the Jacobian norm. However, it is computationally hard to back pook-prop through the computation graph of the operator norm of [| for a practical neural network even with simple iterative method (see Algorithm
2 in Appendix E Thus, we instead penalize an upper bound, the Frobenius norm [| (J) ( (J) with the regularization coefficient
Areg/ C, i.e., minimize I. + Areg (J) / C. See Appendix for its variants with other regular-
8

---

Under review as a conference paper at ICLR 2022
train loss
train loss
test loss
test acc
test loss
test acc
1.0
0.8
2.0
0.9
0.7
1.5
0.6
Loss
ACC
Lo55
1.0-
0.8
0.5
0.4
0.5
0.7
0.3
0.0
0.6
0.2
40000
80000 120000
40000 80000 120000
Step
Step
Figure 7: The effectiveness of EJR (dashed lines) compared to Baseline (solid lines) during training (Left: CIFAR-10, Right: CIFAR-100) WRN-28-10. With EJR, it mitigates the overfitting, especially after each learning rate decay (undesirable decrease/increase of test accuracy/loss).
ization coefficients preg and pree. The Frobenius regularization term can be efficiently computed with an unbiased estimator ||? (Bw=urgo-r)(l()))iu
CEw-urgo- V (u² 2) where u is randomly drawn from the unit hypersphere SC-1 Since the batch-size is large enough, we efficiently use single sample u U(SC-1) for each batch suggested in Hoffman al. (2019). We expect improvements in the generalization performance when introducing EJR. This would support the effectiveness IJR that it efficiently controls the capacity of the model. Table 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 al. 2021 with G instead of the Fisher Information Matrix (FIM). The trace of can be written as follows:
tr(G) (I(JMM*)) E ==
X [01 (tr(M)
(12)
where we assume 11 ~ || since Ei= [| I1 J (see Figure (Left) in Appendix for the empirical evidence). Here, the trace of the logit Hessian M can be equivalently written C-class Gini impurity: tr(M) > P,(1 p.) - p = Gini (p) which C-1 for the initial uniform distribution and 0 for a one-hot probability. Thus, penalizing (r( induces the effects of penalizing [J][ especially in the early phase training with large Gini (p). Thus, as Jastrzzebbkk al. (2021) argued that Fisher Penalty the trace of the FIM improves the generalization performance by limiting the memorization, the Jacobian regularization may have similar effects. Moreover, since Vol(z,9) Vo=V 2 9) (p e²), the trace of the FIM they approximately used is simply 1 E V z, 5)]
[[ e )]I with [ (p elI]ll and the explicit Jacobian regularizer | ClE 1) (B [Ju]ll In each case, we emphasize that the Jacobian J plays important role in the generalization performance.
a single sample ÿ. the gradient norm penalty Barrett & Dherin 2021 is [[
6 CONCLUSION
We have investigated the eigensystem of the Hessian through a new decomposition using eigen- decomposition of the di-eenssiional logit Hessian. By doing so, we could provide simple and intuitive explanation on the relation between the gradient-based optimization, the learning dynamics and the generalization performance neural networks. We hope this research could help answer other intriguing questions on the learning dynamics and generalization.

---

Under review as a conference paper at ICLR 2022
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).
REFERENCES
Alessandro Achille, Matteo Rovere, and Stefano Soatto. Critical learning periods in deep networks. In International Conference Learning Representations, 2019. URL U httphtt:/p/ppeereeeee net fett/oormms dd=
David Barrett and Benoit Dherin. Implicit gradient regularization. In International Confer-
ence on Learning Representations, 2021. URL ULL/tttss///peneeeeew...ttt iit.rrrriiww.neettooumiiiiii
3q5IqUzkeF
James Bunch, Christopher P Nielsen, and Danny C Sorensen. Rank-one modification of symmetric eigenproblem. Numerische Mathematik, 31(1):31-48, 1978.
Pratik Chaudhari and Stefano Soatto. Stochastic gradient descent performs variational inference, con- verges to limit cycles for deep networks, In International Conference on Learning Representations, 2018.
Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-sgd: Biasing gradient descent into wide valleys. Journal of Statistical Mechanics: Theory and Experiment, (11) 88, 2019.
Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference Learning Representations, 2021. httthhttts////peeeeee fot
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In International Conference on Machine Learning, pp. 1019-1028. PMLR, 2017.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalizaatinnn International Conference on Learning Representations, 2021. URL https: httpp:://peeneeieee.
Stanislav Fort and Surya Ganguli. Emergent properties of the local geometry of neural loss landscapes. arXiv preprint arXiv: 1910,0599, 2019.
Stanislav Fort, Pawel Krzysztof Nowak, Stanislaw Jastrzebski, and Srini Narayanan. Stiffness: A new perspective on generalization neural networks. arXiv preprint wwXiv:1901.094/ 2019.
William Fulton. Eigenvalues, invariant factors, highest weights, and schubert calculus. Bulletin of the American Mathematical Sociery, 37(3):209 2299. 2000.
Semyon Aranovich Gershgorin. Uber die abgrenzung der eigenwerte einer matrix. Иаnecтия Роесийсской академии наук. Cерна математическая. (6):749-754, 1931.
Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao. An investigation into neural net optimization via hessian eigenvalue density. In International Conference Machine Learning, pp. 2232-22411 PMLR, 2019.
Gene H Golub. Some modified matrix eigenvalue problems. Siam Review, 15(2):318-33 334, 1973.
Gene H Golub and Charles F Van Loan. Malrix computations, volume 3. JHU press, 2013.
10

---

Under review as a conference paper at ICLR 2022
Guy Gur-Ari, Daniel A Roberts, and Ethan Dyer. Gradient descent happens in tiny subspace. arXiv preprint arXiv:1812.0473 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition,
pp. 770-778, 2016.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural computation, 91):1144, 1997.
Judy Hoffman, Daniel Roberts. and Sho Yaida. Robust learning with jacobian regularization. arXiv preprint arXiv:1908.022 027229, 2019.
Stanislaw Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in sgd. arXiv preprint arXiv: A.KV:1711..0623. 2017.
Stanislaw Jastrzebski, Maciej Szymczak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho and Krzysztof Geras The break-even point on optimization trajectories of deep neural networks. In International Conference on Learning Representations, 2020. URL https:// openreview net wtte:/oromnnni1g111110
Stanislaw Jastrzebski. Devansh Arpit, Oliver Astrand, Giancarlo B Kerg, Huan Wang, Caiming Xiong, Richard Socher, Kyunghyun Cho, and Krzysztof Geras. Catastrophic fisher explosion: Early phase fisher matrix impacts generalization. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 4772-4784. PMLR, 18-24 Jul 2021. URL hhttps://proceeeeeeeeee wtt.peess//1333/asttttss
html
Stanislaw Jastrzebski, Zachary Kenton, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amost Storkey. On the relation between the sharpest directions of DNN loss and the SGD step length. In International Conference on Learning Representations, 2019. URL http:s:://ppeereeiiiii
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantas- tic generalization measures and where find them. In International Conference on Learning Representations, 2020. URL
net www.ooooumnndniieeeee
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In 5th International Conference Learning Representations, ICLR 2017, Toulon, France, April 24-26. 2017, Conference Track Proceedings. OpenReview.net, 2017. URL
htt///ommmddillllytlo 1y Ygg
Alex Krizhevsky, Geoffrey Hinton, al. Learning multiple layers of features from tiny images. 2009.
Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.022/18 18, 2020.
Chris Junchi Li. Lei Li, Junyang Qian, and Jian-Gue Liu. Batch size matters: A diffusion approxima- tion framework on nonconvex stochastic gradient descent. stat, 1050:22, 2017a.
Qianxiao Li, Cheng Tai, and E Weinan. Stochastic modified equations and adaptive stochastic gradient algorithms. In International Conference on Machine Learning, pp. 2101-2110. PMLR. 2017b.
Chaoyue Liu, Libin Zhu, and Mikhail Belkin. Toward theory of optimization for wwer-parameterized systems of non-lineat equations: the lessons of deep learning. arXiv preprint arXiv:2003.00300) 2020.
Hanxiao Liu, Karen Simonyan, and Yiming Yang. DARTS: Differentiable architecture search. In International Conference on Learning Representations, 2019. bttp:stss::/// /ttp:::://ppeeviieeee

---

Under review as a conference paper at ICLR 2022
Stephan Mandt, Matthew D Hoffman, and David M Blei. Stochastic gradient descent approximate bayesian inference. Journal of Machine Learning Research, 18:1-35, 2017.
Leland McInnes, John Healy, Nathaniel Saul, and Lukas Großberger. Umap: Uniform manifold approximation and projection. Journal of Open Source Software, 3(29):861, 2018. doi: 10.21 105/
joss.00861. https:/ /ttp::://do..ooo htt//oi.orrg ooriiiorg//00 10. 211 05/ 20ss 00861
Behnam Neyshabur. Implicit regularization in deep learning. arXiv preprint arXiv:1709.01955 2017.
Vardan Papyan. The full spectrum of deepnet hessians at scale: Dynamics with sgd training and sample size. arXiv preprint 1 2018.
Vardan Papyan. Measurements of three-level hierarchical structure the outliers in the spectrum of deepnet hessians. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.). Proceedings of the 36th International Conference on Machine Learning, volume 97 Proceedings of Machine Learning Research, pp. 5012-5021 PMLR, 09-15 Jun 2019. URL http::://///tccct httttp:://proceedingg mlr.
Daniel Park, Jascha Sohl-Dickstein, Quoe Le, and Samuel Smith. The effect of network width on stochastic gradient descent and generalization: an empirical study. In International Conference on Machine Learning, pp. 5042-5051, PMLR, 2019.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematical statistics, pp. 400-407, 1951.
Levent Sagun, Leon Bottou, and Yann LeCun. Eigenvalues of the hessian in deep learning: Singularity and beyond. arXiv preprint 1 2016.
Levent Sagun, Utku Evci, V Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis of the hessian of over-parametrized neural networks. arXiv preprint arXiv: 2017.
Nicol N Schraudolph. Fast curvature matrix-vector products for second-order gradient descent. Neural computation, 14(7):173-11738, 2002.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference Learning Representations, ICLR 2015, San Diego CA, USA, May 7-9, 2015, Conference Track Proceedings. 2015. bttp:t //ap::////xiiiii www.ooogaa..... abs 1499. 1556
Samuel I. Smith and Quoe V Le. bayesian perspective on generalization and stochastic gradient descent. In International Conference on Learning Representations, 2018.
Samuel Smith, Benoit Dherin, David Barrett, and Soham De. On the origin implicit regularization in stochastic gradient descent. In International Conference on Learning Representations, 2021. URL htt.pss://ppereeiieeee wtwp:erveveieeeeeeoorrrr.om www.///oorriddrrggoooooo
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent separable data. The Journal of Machine Learning Research, 19(1): 2822-2878, 2018
Huan Wang, Nitish Shirish Keskar, Caiming Xiong, and Richard Socher. Identifying generalization properties in neural networks. arXiv preprint arXiv: 07402. 2018.
Hermann Weyl. Das asymptotische verteilungsgesetz der eigenwerte linearer partieller differentialgle- ichungen (mit einer anwendung auf die theorie der hohlraumstrahlung). Mathematische Annalen, 71(4):441-49, 1912.
Chen Xing, Devansh Arpit, Christos Tsirigotis, and Yoshua Bengio. A walk with sgd. arXiv preprint arXiv:1802.08877, 2018
Sho Yaida. Pluctuation-disiipatton ence on Learning Representations, 2019. URL http::://ppezeviieeeeee/ wwt/ffeutttttteeeeeeeeeeeeeeteet Forum?id
dissipation relations for stochastic gradient descent. In International Confer-
SkNksoRctQ
12

---

Under review as a conference paper at ICLR 2022
Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael W Mahoney. Pyhessian: Neural networks through the lens of the hessian. 2020 IEEE International Conference on Big Data (Big Data), pp. 581-590. IEEE, 2020.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. arXiv preprint arXiv:1605.07 146. 2016.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):107-115, 2021.
Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects. In International Conference Machine Learning, pp. 7654-7663, PMLR, 2019.
13

---

Under review as a conference paper at ICLR 2022
A PROOF OF EQ
= diag(p) pp" Roxe
Proof. By definition of the softmax function,
exp(zz)
Dj [softmax(2)) = Ee exp(z*) = -ep(2z))s-
(13)
where s EE exp(=*), have
Vx,P, { (exp(z;)s~exp(e -PiPj:
i
(14)
-exp(zz)) eexp(z;) (ex((2) + exp(2;)) (= -p? Pp, =f
which leads to V=P (V2,Pi)y = -pp" + diag(p).
B PROOF OF THEOREM
Theorem (restated). The eigenvalues ( (A(1 2 ((2 2 ... 2 (()) and the corresponding
normalized eigenvectors (0) of the logit Hessian M = -V1 = diag(p) - pp" satisfy the following properties:
(a) The eigenvalue X is the i-th largest solution of the following equation:
u(A) =1 De
(b) The eigenvector (8)) is aligned with the direction of (diag (p) 1)-
(c) P(++)) (1) PG) for <i and (C) =
(d) Gini(((a) S (L) S Gini(P(1)) where Ginilg) = 1 - 9² - -1- 22 9) is the
Gini impurity for the binary case (q, 4).
Proof. The eigenvalues (+) of M = diag(p) pp" are the zeros of the following characteristic
polynomial:
(M det(diag(p) pp" XI)
(15)
del(diag(p) AT) del(I (diag(p) - AT)-'pp")
(16)
p
(p. X)
(17)
Pj
j=1
where the second equality follows from A pp" diag(p) XI, and the third inequality holds because det I uv
A(I
--ppp") with the matrix A
= + u ar for vectors u and
Then equivalent to solving the following equation:
p P
c(A) =1
(18)
which implies (a). Note that this result also implies (c) as shown in Figure
Next. to prove (b), put A diag(p) Al and q A-'p. Then it is required to show that
(M >f)q (A- pp" )q for the eigenvalues (0) We have
(A pp' ) (A pp A
pp A A
(19)
14

---

Under review as a conference paper at ICLR 2022
Here, (pp" A A) p);
Ej,xPiP/ P.P, A ΣS PiPjO 1)-1pk pk Σx PiPk (Pk
A) Ps = pi E Pk//-PP A) = ps. The last equality holds for the eigenvalues À = ( which
follows from (a).
Now, we want to prove the statement (c). Since
(C) & (A) + (B (C) 2 (A)
(20)
(A) + X (B)
(21)
for C A+ E RCCC where (D) is the i-th largest eigenvalue of a matrix D (eyl 1912.
Fulton 2000), we can get (C) (A) + (B) and ( (C) ((++)
(A) X (B).
Thus, for A diag(p) and B -pp" we can get
P(++)) S X (8) M P()) for 1 sisc-1
(22)
since (A) P(j), 1 (A) P(++1) and
(B) (B) 0. Moreover, since
M1 = - pp" p p E, P: = 0, the smallest eigenvalue ( =c)=
Lastly, we prove the statement (d). From the Gershgorin circle theorem (Gershgorin 1931 we have
A U Σ ( (( (PPp)) P())).P1((LL 1100011-P1111 P())) [0. 21) P())] (23)
i7i
which implies ( < 2p (1-P())) Note that (1)(11P(1)) 2 P(o) (1- P()) since gt) = t(1 t)
is increasing for 0 S 1 S 0.5. In detail, if P(1) 2 0.5, since P(s) - 1 P(1) S 0.5, we have g (P(a)) I g(1 - p(n) = g(P()) Otherwise (P(1) < 0.5), since Po) S P(1), it leads to the same inequality (P()) S (P())) With the Rayleigh principle, we can express the largest eigenvalue as (1) = max
u Mu, and thus (1)T ( = M(ja = Pt) - P())) S A(),
1 |
C
PROOF OF THEOREM 2
Theorem (restated), For some 0 5 X S ( for each x E D, we can bound
= = s (111)
Proof. We start with the Rayleigh principle:
IKMMJ")ll max q (JMJ" q -= max T JMJJ ④) 1-
(24)
=1
Since M = = 1 (0) (T,, we can continue by putting v = J q. and then
Eq 1
max (0" MD)
max ||||| 11
(iT 0)3)
(25)
Then by putting Ā E. 7 (
with (0 = (q))=((())) (T)) E, (1) T)) 2 0 OEE 1===11 1),
Eq 25
max |-|=1
(q (1) w)
max (Ā
(q (i)T JT 9)
(26)
||||=1
max (Ā llgl==1
(q (T 9))
(27)
Since [ an orthonormal basis of C (eigenvectors symmetric matrix M), we have the || 1 í with q² arg [X
following by putting X
|
Eq
27
max ÁLlj g|l (A][)" (111)]11))
(28)
[|
15

---

Under review as a conference paper at ICLR 2022
Because of the definition of X* and with ||q"|| 1 and Σ, (7(1) 2 0), we have
X S A (
(29)
which leads to the final inequality in Eq and 7(() 0 (i + 1) for all DD
The equality holds when || q"1 || ((0 =
16

---

Under review as a conference paper at ICLR 2022
D GRADIENT DESCENT IN THE TOP HESSIAN SUBSPACE
Gur-Ari et al. 2018) showed that the gradient of the loss quickly converges to a tiny subspace spanned by few top eigenvectors the Hessian after short training. Then, the top Hessian subspace does not evolve much, which implies gradient descent happens in tiny subspace. However, the underlying mechanism has not been fully understood.
Direction of q and two salient elements 1 of M. The eigenvector
We investigate the direction of the eigenvector q
(30)
Pj X()
can be obtained from Theorem | (b) for some a > 0. Here, the magnitude of the denominator Ip; X is small for the two indices j = (t), (7+ 1), and is large for the others. This is because the eigenvalue ( lies between P(++)) and P(a) (Theorem) (c)). Therefore, the eigenvector q has
(8) and large negative value in ( compared the other relatively large positive value in q
components.
(1)
0.44
0.10
0.05
0.03
0.02
0.01
0.01
0.01
0.00
0.00
0.4
0.27
0.25
0.07
0.03
0.02
0.01
0.01
0.01
0.00
0.00
0.3
0.17
0.2
0.00
0.05
0.18
0.13
0.04
0.02
0.01
0.01
0.00
0.00
0.02
0.03
0.05
0.16
0.09
0.03
0.01
0.01
0.01
0.00
0..01
00.01
0.02
0.04
-0113
0.07
0.02
0.01
0.01
0.00
-0.
-0.01
-0.01
0002
-0.02
-0.04
-0.12
0.05
0.01
0.01
0.00
0.0--00.01
00.01
-0.01
-0.01
-0.01
-0.02
-0.09
0.04
0.01
0.00
-02
(C-1) 00000
-0.00
-0.01
-0.01
-0.01
-0.01
-0.02
-0.08
0.02
0.00
-03
(C) 00.00
-0.00
-0.00
-0.00
0..00
-0.00
-0.06
0.00
-0.4
VAnl g(1)
Vi (12)
VA 33
(V q (D) ..... :...; (C-1) averaged over the training set 'D where M = QQQ Each column of Q/1/2 visualizes the color-encoded direction of (0) multiplied by Vx(1) We highlighted the elements q and qitt1) with the dashed boxes for Osic C (see Theorem (b)).
Figure 8: Heatmap of the matrix Q/
(1)
Figure 8 shows the directions of (9) with the heatmap of the matrix QA where QA1/2 [VX(T] P;... VVCC-I (C-1) E CCxC for M = QAQ As expected, considering each column of 1/2/2, the eigenvector (8) is colored in red (+) at 9 and in blue (-) ( for
(4)
(6)
isisC- 1. The two salient elements are highlighted with the dashed boxes.
Direction (t) and margin maximization
light of the previous discussion, the direction of lo E R m
Jg (1) =
Vo (q (0(+))
(31)
approximately a direction maximizing ( z(²) ((+1) ()+ at the current parameter ((t) e because the other terms are relatively small. In other words, it tends to maximize the margin (()) ((++)) in the logit space Z between the classes (i) and (i+ 1) (see Figure In particular,
Jq(1) is approximately a direction that maximizes the margin between the most likely class and the second most likely class.
17

---

Under review as a conference paper at ICLR 2022
airplane 1.00
0.07
0.02
0.06
automobile
-0.07
0.05
00.0
0.75
bind -0.12
0.05
1.00
0.00
0.49
0.11
0.10
0.01
0. 0.02
0.50
0-..3
cat
0.005.0.08
0.06
0.25
deer 0.00 0.02
0.45
0:02
1.00
0.00
0.37
0.19
0.05 0.01
dog
0.06 -0.05-0.11
-0.32
0.00
1.00
0.15
0.06
0.00
0.08
p 0.5
frag
0.02 0.02
00-000000
0.11
-025
horse
-0.02 0.00
-0.01
-0.03
1119-00006
0.08
1.00
0.06
0.04
-0.50
0.40
0.06
-0.05
0.00
0.06
1.00
-0.75
truck 0.04
...00.022
0.04 -0.01 0.08
.0.04 0.04
0.119
airplane
automobile
bird-
tak-
dear-
doa
horse
ship
Figure 9: There are C' clusters of V Jq according to the most likely class (not the true class). Left:: (Within-class similarity) Directional data of m from D, and Dj for the two classes, dog (blue) and automobile (red). They projected onto the 2D-plane spanned by the two mean vectors indicated with the arrows. We highlight the MRL for each class. The directional data of m concentrated within the class but separated from each other. Right: (Between-class dissimilarity) Cosine similarities between each pair of {m'² J. They are mostly orthogonal, but some pairs even negatively aligned, for example, automobile and truck. This is because the examples predicted to be automobile mostly have the second most probable class as truck.
Clustering of Jq and the most probable class We first define following subsets of the training
set according to the most probable class (and the second most onc): D; = {æ D C1 (x) = i} c D
and Dij = ( E D ej (a: = i, c2(x)) = j} C DD for + j E [C]. Note that D; = Ujzi i Dij- Given two examples from Dij, their (1) are expected to be highly aligned to each other. This is because the direction of Jq (1) is approximately
direction of maximizing the margin and of learning the features to discriminate the class from the class j. Moreover, two examples from D, also have highly-aligned (1) Figure 2 (Left) shows the concentration the directional data of J from D, We also compute the mean resultant length (MRL) measure the concentration. The MRL p of the directional variable E gm- {v= E I || | = 1} defined as p = [E[]]] E 10 indicates how V is distributed (the higher, the more concentrated).
Now, we focus on vA(I) as the other (O-terms are dominated by the (1) after few epochs (see Appendix [ for details). Then, we follow a similar approach from (Papyan 2019) and provide the following equation:
(mm") = Σ 7mmmm E
mmmm m(m'm + ((m mm )D))
(32)
where 7 = DI and mn' = (m) Here, the covariance term ((m m (m m
D. is
weak as m is concentrated within D, and thus we can roughly approximate G with Σ 7: m IT This implies that the top eigensubspace
the Hessian highly overlaps with the at most C-dimensional subspace spanned by { (m') ] /==1- Figure 9 (Right) demonstrates that the mean vectors {m C well separated from each other. This also implies the outliers in the Hessian spectrum 2016 2017).
Tee1 are (Sgunn et al.
Why gradient descent happens mostly in the top Hessian subspace? Given input æ. after the model becomes correctly predict the true label y, the gradient descent dircction - [ee p) used in the training tends to be highly aligned with ( This is because ex p and (1) both have similar direction. They have positive values 1 Py and (
(1) in y(= (1))-th element, negative value
-P, and q: (1) in the others, and especially large negative value for the second most probable class i= (2). Figure[o0 (Middle) (Middle) shows the cosine similarity between the gradient descent direction -g and ()) As expected, they are highly aligned with the cosine similarity near the two vectors e p and (1) become more aligned to each other.
18

---

Under review as a conference paper at ICLR 2022
Trai Lows
Adl
AuD
20
0.5
8.6
0.6
1.0
0.4
0.4
0.5
1000
2000
2000
4000
1000
2000
3000
4000
1000
2000
3000
4000
Stop
Stop
Stop
Figure 10: Left: Total gradient - (g) is aligned with the top eigenvector r of the Hessian H at each step during training (Jastrzebski et al., 2019; Gur-Ari et al., 2018). They have large cosine similarities considering that they are very high-dimensional, We highlighted the cosine value for random m-dimensional vectors in & with the dashed horizontal line (about le-3). Middle: Jq (or m) is highly aligned with the gradient g for given example each step during training. They have cosine similarities near 1 as the model becomes to correctly predict the true label. See Figure 2 (Right) together. Right: Total gradient g² and the top eigenvector of Hessian H mostly lie in the at most C-dimensional subspace S spanned by {m') J'i=1 The subspace alignment measure As is defined in Eq 33
Next, we move on to the subspace S span( ( spanned by {m') JEE As each ( is highly aligned with -g. it is reasonably expected that the total gradient g" lies the subspace S. To measure how much the vector * 0 is aligned with the subspace S, we define the cosine similarity between the vector v and its projection Ps(v) onto the subspace as follows:
As(v) = cos(u, Ps(v)) = ||Ps())]|[|o|]
(33)
In particular, As(v) means S and As(v) 1 means v E S. Figure (Right) shows the high alignment of the total gradient gº in the subspace S span m although the subspace S is of dimension at most C < m. Moreover, since G can be roughly approximated by Ee m. m the top eigenvector of the Hessian H is also highly aligned with the subspace S.
160
140
1200
120
1000
800
80
600
60
40
400
20
200
10
10
10a
1
0.0
0.2
0.4
0.6
0.8
1.0
cos cos? "T
Figure 11: Histograms of
[]
and cos (Jq (g())))))))) 5))
Σ, 1/
19

---

Under review as a conference paper at ICLR 2022
E EXPERIMENTAL SETTINGS
E. DATA
We the CIFAR-10 dataset (Krizhevsky al. (2009) https: www. torento edu/ seriz/cifat 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) allow the training loss converge to small value. case of the WRN-28-110 (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 batch- normalization, VGG for CIFAR-100 wihtout butch-normalization, 6-layer CNN (6CNN), SimpleCNN used in Jastrzebski et al. (222I) (SimpleCNN) two 3-layer connecco networks (3FCN-CIFAR and 3FCN-MNIST), and WRN-28-10 for CIFAR-IOCCIFAR-10 with the number of model parameters, m 9750922, 9797092, 268346, 511926, 3617661770 706, 656810, 199210, 36539199,,,6336888,,,,,,,,,,4,,4 respectively.
ResNet-20 (ResNet) (He al.) 2016)
We use a modified version of the implementation of VGG-11 from without the dropout Locus abaaab//eeee edge of- of-stability wttp///tttubssccc/oooooooctttctt wtwthhub//xc/eeeeeeee stt_ciiaaar_. We change the last linear layer for the CIFAR- 100 dataset. The 6CNN model can be expressed in the Pytorch code as follows:
com/
layers and ResNet-20 from http:///..aaa///////////////////
breegyyangfu/yytooee vgg
nn. Sequential (
nn.Conv2d (3, 32, 3, stride=1, padding=1, bias=False)
nn.ReLU ()
nn.Conv2d (32, 32, stride-2, padding-1, bias-False) nn.ReLU
nn.Conv2d (32, 64, stride=l, padding=1, bias=False) nn.ReLU
mn.Conv2d (64, 64, stride=2, padding=1, bias=False) nn.ReLU(),
nn.Flatten() (),
Ln. (4096, 100, bias-True), nn.ReLU Lnneir (100, 10, bias=True),
and the 3FCN architecture is as follows: nn Sequential nn. nn.Flatten() (), Ln.Line (n, 200, bias-True), nn.ReLU(). Ln.einear (200, 200, blas=True), nn.ReLU
Ln. (200, 10, bias=True),
where 78784 for 3FCN-MNIST, and n=3072 for 3FCN-CIFAR (the same one used in Cohen et al. (021))
E.3 HYPERPARAMETERS
SGD Robbins & Monro 1951 with the learning rate can be expressed as follows:
0 (++1) ((t) "11" (1() (0(t)
(34)
where (( is the model parameter, B(() C D is the training batch at t-th step, = Vol, and g =g) We mostly use the simplest form of (S)GD described in Eq (34) without momentum,
20

---

Under review as a conference paper at ICLR 2022
weight decay, and learning rate decay. The only exception is when train WRN-28-1U 10, 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 EACH FIGURE
For Figure and 10 we ran the experiments on the CIFAR-IB 10 dataset using 6CNN and trained the model using GD with the learning rate 7) 0.04. See Appendix H 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) Appendix D
E.5 HESSIAN
When computing the top eigenvalue H and the corresponding eigenvector of the Hessian H, use the tool developed in PyHessian (ao et al. (2020), phtpss: / //g1
thub. com/ ammi righlamm Pylees iissiaa an MIT License) based on power iteration method using a small subset (5-25%) of training dataset D.
E.6 POWER ITERATION ALGORITHH
Even though we have the secular function v(A) in Eq 9 and the algorithms for computing the eigenvectors (Bunch et al. 1978),
the power iteration in Algorithm to get the top eigenvalue X(1 of the logit Hessian M E RCCC since we can run the algorithm for a mini-batch in parallel.
Then, we can compute the corresponding top eigenvector q from Theorem (b). To compute the
second largest eigenvalue, we apply the power iteration to M' = M x (1)T instead of M after computing (1) and (2)
Algorithm 1 Power iteration
Input: matrix M, maximum iteration Mmmax tolerance bound Output: the spectral norm ( = || la of the matrix M Initialize u IRC with random vector.
i i 0
repeat
v + Mu/||Mul u M u/ i iii+1
until it converges within the tolerance bound e or 2 7mmax
return v- Mu
We also compute the operator of the Jacobian [| with the power iteration in It requires(C + (ajuiresC++))--imees scalar function differentiations with respect to @ for each iteration.
21

---

Under review as a conference paper at ICLR 2022
Algorithm 2 Power iteration for Jacobian
Input: logit function ze (not matrix (J)), maximum iteration 7lmax, tolerance bound Output: the operator norm [|]]]] of the Jacobian
Initialize u RRC with random vector.
i 0
repeat
Compute (J)u Volu" 2) (1 scalar function differentiation)
V + (J)u/ JJ
Compute (J) 0 (D) = .o (o(z)).... o'' [z()))" (Ox scalar function differen-
tiations)
u u (J) T [
i+i+1
until it converges within the tolerance bound return of (()))
2 1lmax
22

---

Under review as a conference paper at ICLR 2022
F THE TENDENCY OF THE JACOBIAN NORM TO INCREASE
The weight norm |000] increases (Figure 12 (the third one)) in order to increase the logit norm
|2(0() (Figure 12 (the leftmost figure)) and to minimize the cross-entropy loss during training (Soudry al. 2018). This also leads to the increase the layerwise weight norms (Figure 14] and the Jacobian norm (Figure 6). While the Jacobian norm tends to increase, it stops increasing at certain point, stays for few steps, and increases again. We ran the experiments on the CIFAR-10 10 dataset and 6CNN using GD with learning rate 0.04.
12
21
20
0
Step
Figure 12: (Left to Right): the logit norm (0()() the absolute value of the logit (0(())) the weight norm g(t) the distance from the initial weight |[(t)) 0001| during
training (every 10 steps). See together with Figure (Left, solid red line) and Figure 5
3.0
1.0
17.34
0.8
0.6
11812 atti)
2.0
17.33
F-111
0.4
17.32
1.0
0.2
Step
Stop
Step
Figure 13: (Left Right): the logit norm (I)= the absolute value of the logit sum [| (0)")))) the weight norm | ||, the distance from the initial weight ||0(0) ((0) in the early phase of training (every step).
85
70
70
2000
4000
2000
4000
2000
4000
Step
Step
Step
25
150
20
145
140
15
135
50-
10:
40
2000
4000
2000
4000
2000
4000
Step
Step
Step
Figure 14: The layerwise weight norms (6 layers from left right and from top bottom) of the 6CNN model in the early phase of training (every 10 step).
23

---

Under review as a conference paper at ICLR 2022
G H (J)
Surprisingly, we empirically observed that H & (11)) during training where 1/vC E "... Since it requires further theoretical grounding, we left it as future work. We observe this relation for variety of learning rates (Figure 15), network architectures (Figure T6 and T7p.
batch sizes (Figure 18 and 19). and datasets (Figure 20 and 21 At least, H| and {> b) (J)
they increase and decrease together. We emphasize that Figure 21 shows the case when the sharpness
did not reach the limit 2/1 This because the learning rate is relatively low and it easy to train
model for the MNIST dataset, and thus (A(1) decreases before the sharpness reaches the limit.
||H||
200
175
125
100
75
25
1000
2000
3000
4000
Step
||]|
||||
A PP1i
||]|
175
175
175
125
125
125
100
75
50
50
50
100
400
1000
500
Step
Step
Figure 15: The relation | & (J) on the CIFAR-10 dataset and the 6CNN model trained using GD with different learning rates 7 20200.04//..0 (from left to right). Bottom figures are plotted for the early phase of training. Top figure is plotted for = 0.02. Curves are plotted for every step.
24

---

Under review as a conference paper at ICLR 2022
|||]
||1||
120
25
120
25
100
20
100
20
60
60
10
10
40
20
1000
200
1000
Step
Step
||M|
||i||
30
30
140
140
120
100
20
100
20
60
10
10
40
40
300
700
900
800
700
900
Stop
Stop
Figure 16: The relation H (J)il] on the CIFAR-10 dataset and the VGG model trained using GD with different learning rates 17 = 0.04/0....0. (left/right) in the early phase of training. Curves are plotted for every step (Top: 0-1000 stpes and Bottom: 500-100 1000 stpes).
||M||a
200
50
150
125
30
100
75
10
25
500
1000
Step
Figure 17: The relation ||]] x (J)) on the CIFAR-10 dataset and the ResNet model trained using SGD with learning rate = 0.04 and batch sizes [B(()] 128 during training (0-1000 steps). Curves are plotted for every step.
25

---

Under review as a conference paper at ICLR 2022
|
||Í|]
||M|
200
200
200
200
100
140
150
150
150
150
60
100
100
100
80
100
75
75
75
50
50
50
40
20
25
25
25
20
1.00
200
300
400
500
1.00
200
300
400
500
100
2:00 300
400
500
Stap
Stop
Stop
(I) on the CIFAR-10 dataset and the 6CNN model trained using SGD with fixed learning rate 1] 0.04 and different batch sizes [B( 512/128 11/2/22/2222 (from left to right) in the initial phase (0-500 steps). Note that the proportionality con- stant may change according the batch size (the smaller the batch size, the larger the proportionality constant). Curves are plotted for every step.
Figure 18: The relation H
||M||
||i|
||i|
A A
100
20.0
100
20.0
100
20.0
15.0
15.0
15.0
60
12.5
60
12.5
60
12.5
7.5
7.5
7.5
20
20
20
2.5
2.5
2.5
5000
10000
10000 20000 30000 40000
40000 80000 220000160000
Stap
Step
Step
Figure 19: The relation Hl x (J) on the CIFAR-10 dataset and the VGG model trained using SGD with fixed learning rate 7) 0.04 and different batch sizes [B(O] 512/128/32 (from left right) during training (0-100 epochs). Curves are plotted for every steps (n 97/388/ 1552) where 97 /5000//52]..
26

---

Under review as a conference paper at ICLR 2022
||ile
|Mile
|Mile
80
80
80
60
40
70
60
60
60
30
50
50
40
40
30
30
30
20
15
10
10
10
10
10
5
500
1000
500
1000
5.00
1000
Step
Step
Step
||i|ee
||iie
||ille
80
80
80
60
40
70
70
50
60
60
60
30
50
40
40
40
30
20
30
20
20
20
10
10
10
10
300
1000
800
1000
300
1006
Stop
Step
Step
the CIFAR-100 dataset and the VGG model trained using SGD with fixed learning rate 7) 0.1 and different batch sizes [ 128/6e 225/6/32 (from left to right) in the early phase of training. Note that the proportionality constant may change according the batch size (the smaller the batch size, the larger the proportionality constant), Curves are plotted for every step (Top: 0- 0-1000 stpes and Bottom: 500-1000 stpes).
Figure 20: The relation | (J)
[H]
[H]
50
80
50
80
50-
80
40
40
40
20
20
20-
500
1000
1500
1000
1500
1000
1500
Figure 21: The relation | x (J) on the MNIST dataset and the 3FNST model trained using GD with different learning rates 0.02/0.00 00.04/0.00 (from left to right) in the early phase of training, Note that the sharpness did not reach the limit 2/n (the dashed horizontal line). Curves plotted for every step.
27

---

Under review as a conference paper at ICLR 2022
H IMPLICIT REGULARIZATION ON (J)il²
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 || | x (]) in the previous section and its effect on i For the early phase of training,
hard to see the initial rapid growth of the sharpness in this smoothed curves and when exactly the regularization begins 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, observe that training with larger learning rate and smaller batch size limits ||) with smaller value (dotted red lines) in the Active Regularization Period. Curves are smoothed for visual clarity.
||i|
200
140
120
0.25
Spectral nomm
100
0.20
Norm
60
0.10
40
20
0.05
0.00
1000
2000
3000
4000
5000
Step
Figure 22: The evolution of | [or, 1 and (A on the CIFAR-10 dataset and the VGG model trained using GD with the learning rates 0.04/0.08 (solid/dotted lines). See the Figure caption together.
||]|
|
A
1000
0.25
Spectral norm
800
0.20
Norm
0.15
400
0.10
200
0.05
0.00
2000
4000
6000
Step
Figure 23: The evolution of И .. ||) and (A(1) on the CIFAR-10 dataset and the 6CNN model trained using GD with the learning rates 77 = 0.02//000////// 002/0.04/0.0 (solid/ddshee/ddotted lines). See the Figure caption together.
28

---

Under review as a conference paper at ICLR 2022
|||||
W
A
175
175
0.25
150
Spectral norm
0.20
125
125
100
100
Nomm
0.15
75
75
0.10
50
50
25
25
0.00
3000
9000
Step
Figure 24: The evolution of | Il. [I il and (X() on the CIFAR-10 dataset and the VGG model trained using SGD with the fixed learning rate 7) 0.04 and the batch sizes [B(0] (D)00//D)/512//// (solid/dashed/dotted lines). Training with a batch size 512 shows similar evolutions to the GD training. See the Figure 6 caption together.
||H||>
w/
x
50
0.30
40
800
0.25
Spectral norm
0.20
30
600
Eigenvalue
Norm
0.15
20
400
0.10
10:
200
0.05
0.00
500
1000
1500 Step
2000
2500
3000
Figure 25: The evolution of |.,, | and (A(1) the CIFAR-100 dataset and the VGG model trained using SGD with the fixed learning rate 1) 0.1 and the batch sizes [B(t)) 128 /6//34/22222 (solid/dashed/dotted lines). See the 6 together.
29

---

Under review as a conference paper at ICLR 2022
I FIGURE 4 VISUALIZATION OF THE OPTIMIZATION TRAJECTORY)
We use UMAP McInnes et al. 2018) to visualize the optimization trajectory (0(3) heer C + C m 2D space. As GD enters into the Edge of Stability (Cohen et al. 2021), oscillates in 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).
12.5
450
400
10.0
350
7.5
300
5.0
250
2.5
200
150
0.0
100
-2.5
0
-10
-5
20
800
700
600
10
500
400
300
0-
100
-10
10
20
25
900
800
700
10
500
5-
400
0-
300
200
100
-20 -10
10
Figure 26: Visualization of the optimization trajectory using UMAP. (Top) UMAP on the CIFAR- 10 dataset trained with 6CNN for the first 500 steps. (Middle) on the CIFAR- 10 dataset trained with VGG for the first 1000 steps, and (Bottom) on the MNIST dataset trained with 3FCN-MNIST for the first 1000 steps (from red to blue).
30

---

Under review as a conference paper at ICLR 2022
J FIGURE 8 (2A
Figure 27 shows the matrix Q/ (Figure 8) for different training steps. Figure 8 and Figure 27 (Top Right) plotted at the equivalent step (t 1000). Figure 27 demonstrates that X becomes more dominant than the others as training progresses. The argument that there are two salient elements in each qd ()) in its (i)- and (i + 1) elements is empirically shown to be valid throughout the training.
2
1 2
Figure 27: Evolution of QA 1/2 for some training steps during the training. They visualized for 100/50001100000000 00/00/0000020000/400//0000 steps from left to right and from top to bottom.
31

---

Under review as a conference paper at ICLR 2022
3.80
6:29
401
057 4.00
+000
4:00
0.-83 5:0
40.07
1.00
8.0 USE 340
0000 080
Figure 28: Cosine similarities of (m²)
1=1 during training. They are visualized for 3tc==50001100//55//000000000 from left right and from top to bottom. See the Figure caption.
1.0
0.5-
0.2
0.0
500
1000 1500 2000 2500 3000 3500
Figure 29: The Mean Resultant Lengths (MRL) p for each D,. Note that p not defined for first few steps because some D; are empty.
32

---

Under review as a conference paper at ICLR 2022
K FIGURE 9 (CLUSTERS OF m AROUND EACH m
Figure 30 shows Figure (Left) for some different class pairs (i, with negative cosine similarity. ie., cos(mm ,m²) < O. We use the model trained for = 1000 steps. Figure 28 shows Figure
(Right) for different training steps. Figure 29 shows the evolution of the Mean Resultant Length (MRL) p in Figure ( (Left) during training.
pattg
Figure 30: Directional data of m from D; and D,. They are visualized for (i = (airplane, ship)/(atomobilee truck)/(dog, cat)/(deer, bird) from left right. See the Figure caption.
L
FIGURE 10 SUBSPACE)
(GRADIENT DESCENT HAPPENS MOSTLY IN THE TOP HESSIAN
Figure 31
shows similar results with different settings with VGG and learning rate r 0.08.
Train Loss
coosi o M|
Astg" P)
Astri
1.0
1.0
0.8
1.5
0.6
Subspace
1.0
0.5
0.2
0.2
0.0
1000
2000
3000
1000
2000
3000
Step
Step
Figure 31: (Left) Alignment between the vectors gº and r, and (Right) alignment of g in the subspace S span((m' are plotted for every 25 steps. Note that As not defined for first few steps (about 0-800 steps) because some D. are empty.
1 (iz)) using VGG with 0.08. See the Figure 10 caption. They
33

---

Under review as a conference paper at ICLR 2022
M ANALYSIS OF THE MSE LOSS
In the main text, focus the cross-entropy loss. Here, we briefly analyze the MSE loss, / I| ed| 2 Then, we have M V2l I, ( = |M and G (JJ"). It leads to the
same conclusion as in Theorem
|| [| M|ا K||JJ' (1
(|
(35)
We empirically observed that H |l
|| 2 as shown in Figure 32
test
||H||a
11/)112
acc
300
250
0.6
250
200
0.5
Jacobian Norm*
200
Sharpness
M
150
0.4
150
0.3 <
100
100
0.2
50
50
0.1
0
0.0
20000
40000 Step
60000
80000
Figure 32: The sharpness H a and Jacobian norm || 12 during training with the MSE loss
34

---

Under review as a conference paper at ICLR 2022
N DETAILS OF EXPLICIT JACOBIAN REGULARIZATION (EJR)
We first propose simple form of EJR with the regularized loss as follows:
[(0) L( + Xrea /
(36)
and update the model parameter
e g(t) 77 (VoL((())) + Xreg Voll())
(37)
However, this requires to build a computational graph for [| which 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:
0 (+++1) g(0) v)Vo)(0())
(38)
where 0- - -0- - Preg (J) and Preg Preg/ (J) in Foret et al. 2021); ; Liu et al. (019))
Second, we propose another variant using L instead of Eq (38) follows:
L= L + prrg u (2)
(39)
We approximate L as follows:
((0) ((0) (VoL(0)) (Prey (J)u)
(40)
L(() (VoL(0))) (J)u) + (Volreg u (z))² (Preg(J)u)
(41)
L(0) (Pre (J)u) + (Preg (J)u) ( (Preg(J)e)
(42)
= -(()) + (VeL(0))) (Jyu) + PregPregh (J)
(43)
~ ((()) (VeL(0)) (Prea (J'r + Preg Preg || kJ)/
(44)
We used the first-order Taylor expansion of i + ( (J) in Eq 40 of minimizing Mrogut PregPreg ((T] ll//) /VC ze PregPreg (J)
We expect additional effect
C/CC compared to the first
variant.
Ávag
Arog
Ave
Avg 0.01
Ang 0.03
76
74
72
70
68
66
10-2
10-1
learning rate n
Figure 33: [Explicit Jacobian Regularization] The explicit Jacobian regularization enhances the test accuracy. We plot the test accuracy for different learning rates y and regularization coefficients Aveg- The models are trained with batch size of [ = 128 on CIFAR-10.
35

---

Under review as a conference paper at ICLR 2022
Table Effectiveness of EJR. We report improvement (AAcc.) and Error Reduction Rate (ERR) on CIFAR-10 10 when trained with EJR (+EJR), compared to the standard training (Baseline).
Network
Test Accuracy
AAcc.
ERR
Dataset
Batch Size
Reg. param.
Architecture
Bascline
+EJR
(%p)
(%)
0.003
Aveg 0.01
66.71
75.40
+8.69
26.10
0.01
Areg 0.03
67.88
75.66
+7.78
24.22
0.03
Areg
69.83
75.53
+5.70
18.89
0.1
Aveg 0.03
70.33
74.59
+4,26
14.36
SimpleCNN
0.3
Aveg 0.01
69.34
73.47
+4.13
13.47
0.01
66,81
74.43
+7.62
22.96
50000
0.03
67.72
74.31
+6.59
20.42
Aveg 0.001
(full-batch)
0.1
67.53
73.69
+6.16
18.97
CIFAR-10
0.3
61.08
72.15
+ +11.07
28.44
Prog 0.5
96.44
+0.51
12.72
WRN-28-10 (200 epuchs)
0.1
Preg
95.93
96.57
+0.64
15.72
Preg
96.62
+0.69
16.95
Prog =
96.30
+0.37
9.09
Oreg 0.5
96.65
+0.55
14.10
WRN-28-10 (400 epochs)
128
0.1
Preg
96. 10
96.78
+0.68
17.44
Prog 2
97.07
+0.97
24.87
Prog
96.79
+4.669
17.69
Prog = 0.5
80.42
+0.13
0.66
Preg
81.11
+0.82
4.16
WRN-28-UE (200 epochs)
Preg
81,50
+1.21
6.14
128
0.1
Prog
20.29
82.51
+2.22
11.26
Preg
82.65
+2.36
11.97
Preg 5
82.31
+2.02
10.25
Preg
82.03
+1.74
8.83
CIFAR-100
Preg
82.55
+1.86
9.63
Preg
82.51
+1.82
9.42
Preg
82.84
+2.15
11.13
WRN-28-10 (400 epochs)
Preg 4
83.35
+2.66
13.78
128
0.1
80.691021
Prog
83.73
+3.04
15.74
Preg
83.16
+2,47
12.79
Preg
83.12
+2.43
12.58
Preg
81.16
+0.47
2.43
Table 3: Effectiveness of EJR.v2. We report improvement (AAcc.) and Error Reduction Rate (ERR) on CIFAR-10 when trained with the second variant of EJR (+EJR.v2), compared to the standard training (Baseline).
Network
Test Accuracy
AAcc.
ERR
Dataset
Batch
Reg, param.
Architecture
Baseline
+E.RRv2
(%p)
(%)
Prog Prng 0.001
97.28
+1.18
30.26
Preg 2, Preg 0.003
97.32
+1.23
31.28
Preg 2, freg 0.01
97.33
+1.23
31.54
Preg 2, phrey 0.03
97.38
+1.28
32.82
WRN-28-10
CIFAR-10
128
0.1
Prog Plrey 0.1
10
97.17
+1.07
27.44
(400 epochs)
Preg 1,Preg 0.01
97.07
+0.97
24.87
Prgg 1,prey 0.02
97.34
+1.24
31.79
Preg 1,Preg 0.06
97.33
+1.23
31.54
Preg 1, 1, 0.1
97.26
+1.16
29.74
36


Reviews:
[{"review_id": "Ri8xeH_gVu4", "reviewer": "Reviewer_4LVi", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths \n\n1. This study links three concepts in optimization and generalization (flat minima, oscillation, gradient regularization). It is interesting to see how the diversity (i.e. entropy) of the softmax outputs is connected to the sharpness of the loss landscape. \n2. The eigenvalues of the approximated Hessian (logit Hessian) are theoretically analyzed (Theorem 1).\n3. There are experiments that are intended to support the theoretical analysis.\n\n\nWeaknesses\n\n1. This paper often lacks clear descriptions of technical details that are necessary to be contained. I read this paper three times from the beginning, but I'm still puzzled. My questions are:\n- What is the experimental setting of Figure 3? Appendix C4 says CNN is trained on CIFAR10, but the caption says MNIST and Resnet are also used. \n- What is the model architecture of Figure 4?\n- In Eq. (11), the expectation \\hat{E}_D is separately applied to \\lambda and J. Why you can do that? I concern this operation yields a significant approximation error. Indeed, in Figure 5, at the point around 3500 steps ||J1|| reaches 1000 and \\lambda^{(1)} takes around 0.1, which means ||H|| should be around 100 according to (11). However, the actual value is 50. This gap looks large. \n- There is a sentence \"This explains why the behavior of the sharpness in the early phase of training seems to impact the final generalization.\" in Section 5.2. Could you elaborate on this? Why the sharpness in the early phase influences the final performance?\n- The same paragraph says \"λ acts as a regularization coefficient\". Please clarify this. Do you mean the objective is approximated as `original_loss + λ ||H||_σ`? If so, please derive this mathematically. \n2. Jacobian regularization is not significantly evaluated. This study suggests using the Frobenius norm of the gradient as a regularizer. However, its effect is not comprehensively examined in an empirical manner. How is it better than popular regularizers (e.g. weight decay, l1, gradient penalty, Jacobian regularizers described in [*])? How is the effect consistent for different architectures and datasets?\n3. The motivation is not solid. As is described in Section 2, there are several studies working on the decomposition of the Hessian (e.g. Papyan (2019)). Why the decomposition H ~= JMJ should be studied in this paper?\n\n\n[*] https://www.ejournals.eu/pliki/art/13926/", "summary_of_the_paper": "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. ", "summary_of_the_review": "This paper provides a new view that connects three notions in optimization and generalization (flat minima, oscillation, gradient regularization) with theoretical and empirical supports. However, the current manuscript lacks clear descriptions; without these, I cannot evaluate the correctness. Furthermore, the proposed method (Jacobian regularization) is not significantly evaluated. Also, the motivation of using the proposed decomposition is not grounded. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "CwWKuJ6a43", "reviewer": "Reviewer_6UXZ", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "- Theorem 1: What is novel in the claim and proof compared to prior work? In particular, I am referring here to the contributions from linear algebra by Golub (1973), Golub & van Loan, and Bunch et al. (1978). For example, Eq. (9) is known as the secular equation of the eigen-decomposition problem for rank-one modification of a diagonal matrix. Results (a-c) are for sure quite well-known in that literature (this might be a misrepresentation of the theoretical contributions). So, the novel contribution here might only be (d), which should then be formulated as a separate proposition and Theorem 1 should be cited with Golub and others.\n\n- Section 4.2: In the introduction, there has been a lot of stress on eigenvalues of the Hessian, which should be matrix H. This section, on the other hand, states that the evolution of top eigenvalues of matrix M are tracked. This is the matrix of an exponential family model that completely ignores all the layers of a neural network apart from the last one where the softmax operates. The relation to prior work by Cohen et al (2021) is also unclear -- what would be a novel finding compared to that work? In what sense, the current work extends the estimation procedure from that work?\n\n- Figure 3 and opening of Section 5 is unclear and too informal for a conference publication. The appendix is optional and findings should be described clearly in the main part of the paper.\n\n- Section 5.1: Impurity here refers to the top eigenvalue of M? This notion is not defined, there is just a bound of the top eigenvalue in terms of Gini impurity.\n\n- Eq. (11): I did not follow the derivation for this?\n\n- Section 5.2: Could you please clarify the implication that weight norm increase leads to increase in the Jacobian norm? The Jacobian norm should be about the geometry of the loss landscape which can be independent of the actual parameter values? Alternatively, if you start with all-zero parameters then this is to be expected.\n\n- The reasoning on the Jacobian regularization is unclear. Would not Jacobian norm need to decrease for GD to be able to converge?", "summary_of_the_paper": "The paper investigates the dynamics of leading eigenvalues in the so called logit-Hessian matrix and tries to correlate that with the convergence to a good local optima. The work also tracks the behavior of the Jacobian matrix for the pre-softmax outputs and introduces the notion of implicit Jacobian regularization.", "summary_of_the_review": "The paper is poorly written and too informal in a number of places. There are many claims which are presented as facts but are in reality hypotheses. I have expected an analysis based on the Hessian of a neural network and it turns out that the work focuses only on the Hessian matrix of the exponential family model, completely ignoring the network parameters up to the sofmax operator. I also fail to grasp the relevance of the Jacobian regularization and do not see why the drop in the magnitude of Jacobian norm would be of interest. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "BEkp9ZYg7wS", "reviewer": "Reviewer_utPV", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is overall well written with clear mathematics. \n\nCons:  Theorem 1 is enlightening. \n\nWeakness:  Can the author provide some analytical examples. E.g., how do the neural network structures and activation functions take effect in the logit Hessian matrix?  A  two-layer neural network example is needed.\n\nIn addition, there is essential literature on related information matrices, which are missed.\n\nLi, Zhao, Wasserstein information matrix, arXiv:1910.11248", "summary_of_the_paper": "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. ", "summary_of_the_review": "The paper proposes to study an approximate Hessian operator of the loss function. This is helpful in understanding the fundamental non-convexity questions in learning. \n\nA more analytical example is needed to improve the current paper. The example should address how the simplified Hessian matrix depends on the neural network functions and architectures. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "1zdUFWFEAy", "reviewer": "Reviewer_Z9ex", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Strengths:**\n- While not totally novel, it is still a nice, relatively unexplored, approach to look at the low-dimensional logit Hessian to derive some insights. \n- In particular, Theorem 1 provides an interesting relation between eigenvalues of (per-sample) logit Hessian and corresponding probability output. This also provides a way to get some bounds on the top eigenvalue of the logit-Hessian (impurity) in terms of the largest probability for a class. The corresponding figures are also quite nice and convey the point. \n- The sharpness-impurity-jacobian behaviour and its connection to oscillatory, catapult aspects and others are interesting. (But unfortunately, unclear how confidently we can take away the message, see below. Also, the fact that such a relation between sharpness-impurity-jacobian exists is not too far-fetched, given the form of the Gauss-Newton term.)\n\n\n**Weaknesses:**\n1. Theoretically unmotivated usage of the particular 'Jacobian norm': In Section 5, the authors propose an approximate relation between the Hessian spectral norm, 'impurity', and the 'Jacobian norm'. However, the Jacobian norm is not the usual Jacobian Frobenius norm -- as one would expect or as used in related work (Hoffmann et. al. 2019) -- but is the 'norm' with the particular vector all-one vector (normalized to unit length). This usage is far from being innocuous. It seems that, ideally, the authors would have liked to have the eigenvectors of the logit Hessian, and in particular, q^(1), for the spectral-norm, as also evident from the discussion in the appendix --- but this is perhaps hard to come by practically. Hence the issue is that, in turn, the particular normalized all-one vector used is nothing but the eigenvector corresponding to the zero eigenvalue of the logit Hessian. This is very strange given the particular relationship they want to derive. \nThus, it seems to be an unmotivated heuristic which \"surprisingly\" seems to hold in practice as the figure 3 shows. (On that note, in Figure 3, the ||H||_\\sigma is the spectral norm of the entire Hessian or the Gauss-Newton term G alone? Over how many samples is it computed?)\n\n2. Explicit regularization: Here, the authors do not use their version of the 'Jacobian norm', but rather just the Frobenius norm. Yes, both are closely related and perhaps ultimately regularizing either might have similar effects. But, why is such a choice made in the first place, while not following their own derivation? What happens if you explicitly regularize by your particular Jacobian norm ||J 1||^2? Otherwise, as such, it currently makes the whole 'implicit Jacobian regularization' slightly fishy. \n\n3. Behaviour of the actual Jacobian Frobenius norm: In Figure 3 and 5, it would be very relevant to see how does the Jacobian (Frobenius) norm behaves, not just the particular version of the 'Jacobian norm' used by the authors. Can you please add these to the figures?\n\n4. What does the active regularization period correspond to in terms of the progress in loss or performance? This implicit Jacobian regularization claimed by the paper is only active with a particular (early) phase of training but does not indicate the values of the loss there (please provide them) --- thus, it is hard to concretely understand how 'early' is this phase as well as the fact that why or if it suffices to be present only in this phase?\n\n5. Applicability is restricted to only cross-entropy (CE) loss: An important downside of their approach is that this only applies for CE loss, not something like MSE --- where, the logit-Hessian-equivalent is just the identity matrix. Thus, the offered explanations -- if correct -- will not describe the overall behaviour of generalization in neural networks, only the specific case of CE. I think the authors should very explicitly mention this downside. \n\n6. Other puzzling empirical aspects:\n- Regarding the connection between the Jacobian and Fisher penalty, at what point are get Figures in appendix 11 computed? When is this valid in general? Is this valid, at least in the entire 'active regularization period'?\n- Figure 7: Here most of the performance gains, with explicit regularization, are 5-10% which, in my opinion, are pretty large to be true. So, i would like to understand this better: what are the particular networks considered (especially the ones on the bottom-left side)? are the networks trained long enough to a similar level of loss (not just fixed # of epochs) so as to render them comparable? I believe the \"real\" boost, if so, is represented by the point in the upper-right of Figure 7 (right) --- but I would love if my belief is proved wrong.\n- Figure 6: If i read the plot correctly, the trend of impurity, as well as the 'Jacobian-norm', for increasing LR and decreasing batch sizes seems to be reversed --- which is not clear why?\n\n7. Presentation of the Sharpness-impurity-jacobian relation: This relation is also vaguely/poorly expressed.\n- In Fig 3, when you say \\lambda^(1) and J are averaged over the training set, what do you precisely mean in the context of the 'Jacobian-norm'? Is it that the Jacobian J is averaged over the training set and then || E(J) 1||^2 is computed or is it the expectation of the norm over the training set, i.e., E(||J_x 1||^2)? Then why is one preferred over the other? Do they have similar behaviour?\n- Given that you refer to \\lambda^(1) as per-sample impurity in the previous section, and here implicitly the average over training set is considered. Thus, I ask the authors to properly delineate these two usages, say via the horizontal bar at the top. \n\n8. Unclear/incorrect statements:\n- The authors say that Gauss-Newton approximation sometimes fails in the later phase of training. However, I think they are being imprecise here, and strictly speaking, Gauss-Newton approximation (for the Hessian as a whole) can be justified only towards this late phase or end of training. See Sagun et. al., 2017 page 5 \"At a point close to a local minimum ....\". I guess what the authors mean to say this, only in the context of the top-eigenvalue, that in the late-phase the top-eigenvalue is not clearly isolated from the bulk (which makes sense). So, the authors should really clarify that top-eigenvalue approximation from the Gauss-Newton is valid only during early phase, rather than (incorrectly) conflate explanations as presently done. \n- \" However, the generalization gap keeps increasing, which provides another counter-example that flat minima generalizes poorly\": Where are the corresponding train and test loss curves based on which you make this inference?\n- Figure 3 caption is also a bit weird. Can you specify what is the default optimizer, dataset, network? otherwise only some mixed and partial information is present.\n", "summary_of_the_paper": "The paper investigates properties of the logit Hessian spectrum, how it relates to the top-eigenvalue of the Hessian (which they refer to as 'sharpness' though it can be debated), and implies a particular implicit regularization of the Jacobian during an early phase of training (which is shown to yield better generalization). They use the derived relations to 'explain' previous observations on sharpness and generalization.", "summary_of_the_review": "The analysis of the logit Hessian spectrum, in itself, is very interesting. But other than that, most of the other contributions are not so far-fetched or novel, the claim about implicit regularization is not motivated theoretically, many of the empirical aspects are unexplained and questionable, the approach is very specific to cross-entropy loss and thus likely not a general explanation if so. As a result, overall it is unclear if their offered explanations are conclusive and not speculative.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper studies an interesting question of the relationship between the eigenvalues of the Hessian matrix with the probability of the output in the logistic loss, and use this to propose a regularization that can improve the performance of the neural networks.

 All the reviewers agree that although the question is interesting, the paper lacks significantly in terms of representation and would benefit from another round of revision.
Decision: Reject