Dissecting Hessian: Understanding Common Structure of Hessian in Neural Networks
Authors: Yikai Wu
Xingyu Zhu
Chenwei Wu
Annie N. Wang
Rong Ge
Venue:  (2021)
Keywords: Hessian
neural network
Kronecker factorization
PAC-Bayes bound
eigenspace
eigenvalue

Abstract: Hessian captures important properties of the deep neural network loss landscape. We observe that eigenvectors and eigenspaces of the layer-wise Hessian for neural network objective have several interesting structures -- top eigenspaces for different models have high overlap, and top eigenvectors form low rank matrices when they are reshaped into the same shape as the weight matrix of the corresponding layer. These structures, as well as the low rank structure of the Hessian observed in previous studies, can be explained by approximating the Hessian using Kronecker factorization. Our new understanding can also explain why some of these structures become weaker when the network is trained with batch normalization. Finally, we show that the Kronecker factorization can be combined with PAC-Bayes techniques to get better generalization bounds.

Dissecting Hessian: Understanding Common Structure of Hessian in Neural Networks
Yikai Wu* Department of Computer Science Duke University Durham, NC 27708 yikai. wu@duke. edu
Xingyu Zhu" Department of Computer Science
Duke University Durham, NC 27708 xingyu. wwwduuke..uuuu.....m edu
Chenwei Wu Department of Computer Science
Annie Wang
Department of Computer Science
Duke University
Duke University
Durham, NC 27708
Durham, NC 27708
chenwei.. vu5920duke edu
annie wang0298duke. edu
Rong Ge
Department of Computer Science Duke University
Durham, NC 27708 ronggeecs duke. edu
Abstract
Hessian captures important properties the deep neural network loss landscape. Previous works have observed low rank structure in the Hessians of neural networks. We make several new observations about the top eigenspace of layer-wise Hessian top eigenspaces for different models have surprisingly high overlap, and top eigenvectors form low rank matrices when they are reshaped into the same shape the corresponding weight matrix. Towards formally explaining such structures of the Hessian, we show that the new eigenspace structure can be explained by approximating the Hessian using Kronecker factorization; also prove the low rank structure for random data at random initialization for over- wer-parametrized two-layer neural nets. Our new understanding can explain why some of these structures become weaker when the network is trained with batch normalization. The Kronecker factorization also leads to better explicit generalization bounds.
Introduction
The loss landscape for neural networks is crucial for understanding training and generalization. In this paper we focus on the structure of Hessians, which capture important properties of the loss landscape. For optimization, Hessian information is used explicitly in second order algorithms, and even for gradient-based algorithms properties the Hessian are often leveraged analysis (Sra et al. 2012). For generalization, the Hessian captures the local structure of the loss function near local minimum, which is believed to be related generalization gaps (Keskar et al. 2017).
"Contributed equally, listed in alphabetical order
Preprint, Under review.

---

Several previous results including (Sagun et al. 2018 Papyan 2018) observed interesting structures in Hessians for neural networks it often has around large eigenvalues where is the number of classes. In this paper we ask:
Why does the Hessian of neural networks have special structures in its top eigenspace?
rigorous analysis of the Hessian structure would potentially allow us to understand what the top eigenspace of the Hessian depends (e.g.. the weight matrices or data distribution), well as predicting the behavior of the Hessian when the architecture changes.
Towards this goal, we first focus the structure for the top eigenspace of layer-wise Hessians. We observe that the top eigenspace of Hessians are far from random models trained with different random initializations still have large overlap in their top eigenspace, and the top eigenvectors close to rank when they are reshaped into the same shape as the corresponding weight matrix.
We formalize conjecture that allows us to understand all these structures using Kronecker decomposition. We also analyze the Hessian an over-pore wver-parametrized two-layer neural network for random data, proving that the output Hessian is approximately rank and its top eigenspace can be easily computed based on weight matrices.
1.1 Our Results
0.6
11.25
0.50
0.5
0111
0.5
0.2
0.4
0.2
-...: Oirget Dim 19)
-.::: Girput Din: 61
11.00
0.0
100
200
50
100
10
10
10
10
(a) Overlap between dominate cigenspace of layer-wise Hessian at different minima for fel:LeNet5 (left) with output dimension 120
(b) Top 10 singular values of the top 4 cigenvectors of the layer-wise Hessian of fel:LeeNet5 after reshaped as matrix.
and convl 1 I:ResNet18 8-66 (right) with output dimension 64.
Figure Some interesting observations on the structure of layer-wise Hessians. The eigenspace overlap is defined in Definition 4.1 and the reshape operation is defined in Definition 4.2)
Structure of Top Eigenspace for Hessians:
Consider two neural networks trained with different random initializations and potentially different hyper-parameters: their weights are usually nearly orthogonal. One might expect that the top eigenspace of their layer-wiss wise Hessians different. However, this surprisingly false: the top cigenspace of the layer-wise Hessians have very high overlap, and the overlap peaks at the dimension of the layer's 's output (see Fig. lap. Another interesting phenomenon is that we express the top eigenvectors a layer-wise Hessian as a matrix with the same dimensions as the weight matrix, then the matrix is approximately rank 1. In Fig. lb we show the singular values of several such reshaped eigenvectors.
also very
Understanding Hessian Structure using Kronecker Factorization: We show that both of these new properties laser-wise Hessians can explained by a Kronecker Factorization. Under a decoupling conjecture, we can approximate the layer-wise Hessian using the Kronecker product of the output Hessian and input auto-correlation. This Kronecker approximation directly implics that the cigenvectors of the layer-wiss Hessian should be approximately rank I when viewed matrix. Moreover, under stronger assumptions, we can generalize the approximation for the top cigenvalues and eigenvectors of the full Hessian.
Structure of auto-correlation: The auto-correlation of the input is often very close to rank matrix. We show that when the input auto-correlation component is approximately rank 1, the layer-wise Hessians indeed have high overlap at the dimension the layer's output, and the spectrum of the layer- wise Hessian is similar to the spectrum of the output Hessian. On the contrary, when the model trained with batch normalization, the input auto-correlation matrix is much farther from rank 1 and the layer- wise Hessian often does not have the same low rank structure.

---

Structure of output Hessian: For the output Hessian, we prove that in an over parametrized two- layer neural network on random data, the Hessian is approximately rank c 1. Further, we can compute the top c 1 eigenspace directly from weight matrices. We show that this calculation can be extended to multiple layers and the result has a high overlap with the actual top eigenspace of Hessian in most settings.
Applications: As a direct application of our results, we show that the Hessian structure can be used to improve the PAC-Bayes bound computed in Dziugaite & Roy (017))
2 Related Works
Hessian-based analysis for neural networks (NNs): Hessian matrices for NNs reflect the second order information about the loss landscape, which is important in characterizing SGD dynamics (Jastrzebski et al., 2019) and related to generalization (Li et al., 2020), robustness adversaries (Yo et al.] 2018) and interpretation of NNs (Singla et al., 2019). People have empirically observed several interesting phenomena of the Hessian, e.g., the gradient during training converges to the top eigenspace of Hessian (Gur-Ari et al. 2018 Ghorbani et al.) 2019), and the eigenspectrum of Hessian contains a "spike" which has about
large eigenvalues and continuous "bulk" (Sagun et al..
2016 2018
Papyan 2018).
Understanding structures of Hessian: People have developed different frameworks to explain the low rank structure of the Hessians including hierarchical clustering of logit gradients (Papyan 2019 2020), independent Gaussian model for logit gradients (Fort & Ganguli, 2019), and Neural Tangent Kernel (Jacot et al., 2020). These results work in different settings but are not directly comparable (among thee and paper). A distinguishing feature of this work is that we can characterize the top eigenspace of the Hessian directly by the weight matrices of the network.
Theoretical analysis for large eigenvalues of Fisher Information Matrices (FIM): The FIM can be seen as a component of the Hessian matrix for neural networks. Karakida et al. (20T9b) showed that the largest eigenvalues of the FIM for a randomly initialized neural network are much larger than the others. Their results rely on the cigenvalue spectrum analysis in Karakida et al. (201 90 a). which assumes the weights used during forward propagation are drawn independently from the weights used in back propagation (Schoenholz et al. 2017].
Liyer-wise Kronecker factorization for training NNs: The idea of approximating the FIM using
Kronecker factorizations can be dated back to to/Heskes (2000). More recently Martens & Grosse (2015)
proposed fronecker-factoree approximate curvature (K-FAC) which approximates the inverse of FIM using layer-wise Kronecker product and perform approximated natural gradient descent (NGD) in training NNs. Kronecker factored eigenbasis has also been utilized (George et al., 2018)_ K-FAC has been generalized to convolutional and recurrent NNs (Grosse & Martens [016] Martens et al. 2018). Bayesian deep learning (Zhang et al., 2018). and structured pruning (Wang al. 2019). Unlike these previous works which focus on accelerating computations, in this paper we use Kronecker factorization to explain the structures of Hessians.
PAC-Bayes generalization bounds: People have established generalization bounds for neural net- works under BAC-Bayes framework by McAAlesterr (1999). This bound was further tightened by Langford & Seeger (200T), and Catoni (2007) proposed faster-rate version. For neural networks, Dz] jugaite & Roy (2017) proposed the first non-vacuous generalization bound, which used PAC- PAC-Bayesian approach with optimization to bound the generalization error for a stochastic neural network, Their bound was then extended to ImageNet scale by Zhou et al. (2019) using compression.
(200))
3 Preliminaries and Notations
Basic Notations: In this paper, we generally follow the default notation suggested by Goodfellow et al. 2016). Additionally, for a matrix M, let |M denote its Frobenius norm and [M denote its spectral norm. For two matrices M E RRt xb, N E Iaa xb2 let M @ N E IR (b() be their Kronecker product such that [M @ N - 1)xbb++ Mi N
We consider classification problems with cross-entropy loss. For a c-class classification problem, we are given a collection of training samples S = {(x;, J where Vi [N], (æ,, y:) Rd x R We assume S is i..d sampled from the underlying data distribution D.
Neural Networks:
3

---

Consider an L-layer fully connected ReLU neural network without skip connection fo : - RR With (xx) z1,20 as the Rectified Linear Unit (ReLU) function, the output of this network is series of logits z E R computed recursively as
2 (p) (()) xp() BIP).
xx = (z)(1))
(2)
Here (P)) and 2(P) are called the input and output of the p-th layer, and we set x!1) = I, := fo(x)) = z(4) We denote == (wx(1, ,,(1), u0(2) network. In particular, w is the flattened i-th layer weight coefficient matrix W and 6(0) is its bias vector. For convolutional networks, similar analogue presented in Appendii A.2
( b(8),.... ... ww(L, b(L)) E Rp the parameters of the
For a single input a: E Rd with label y and logit output let n((P) and (p) be the lengths of (P)) and 2(p), For convolutional layers, we consider the number of output channels as (P) and width of unfolded input as n(P). Note that x(1) = z, (L) 2 = fo(æ). We also denote p = softmax(z)
/// e: the output confidence.
With the loss function E y) = --E
log(p;) E R+ being the cross-entrops loss between the softmax of logits = fo (2,) IR and the one-hot label y E Re, the training process the neural network optimizes parameter 0 to minimize the empirical training loss:
1 N Σ [Jo(a:) v.) =
((0)
E (m.y)ss
le (z, y)
Fixing the parameter 0, we use H2 a:) to denote the Hessian of some vector e with respect to scalar loss function & at input x.
Hessians:
(v, V= = (2 y)-
-
(4)
Note here can be any vector. For example, the parameter Hessian is H2 x), where 0. The layer-wise weight Hessian of the p-th layer is Hz(w(r), x))
For simplicity, define E as the empirical expectation stated otherwise. We focus the layer-wise weight Hessians H Hc(w((), = E. (w) æ)) with respect to loss, which are diagonal blocks in the full Hessian Hc(0)= = ( æ:) corresponding
the training sample S unless explicitly
to the cross terms between the weight coefficients of the same layer. We define M Hz(z(p),
as the Hessian of output (p) with respect to empirical loss. With the notations defined above, we have the p-th wise hessian for a single input as
(P)
(P) x)) = V2 Vee (
((, y) = M(() @ (o
(5)
follows that
(c(w(() = =E M @ x x T = E
M
8 xx
(6)
The subscription æ and the superscription (p) will be omitted when there is no confusion, as our analysis primarily focuses the same layer unless otherwise stated.
4 Kronecker Factorization of Hessian
The fact that layer-wise Hessian for single sample can be decomposed into Kronecker product of two components naturally leads to the following conjecture:
Conjecture (Decoupling Conjecture). The layer-wise Hessian can be approximated by Kronecker
product of the expectation of its two components, that is
H (c(w()) E
M & man
~ E[M] ? E(mx=""
In particular, the top eigenvalues and eigenspace of Hc(w(()) close to those of [ } 6 E(aa=?).
Note that this conjecture is certainly true when M and xæ. are approximately statistically indepen- dent. (Section 4.1 and [Section 4.2 we will show that this conjecture is true in practice. We then
analyze the two components separately in Sections 4.3 and F.3.3

---

Assuming the decoupling conjecture, we can analyze the layer-wise Hessian by analyzing the two components separately. Note that E[M] is the Hessian of the layer-wise output with respect to empirical loss, and [xxx== " ] is the auto-correlation matrix of the layer-wise inputs. For simplicity we call E[M] the output Hessian and TP] [[==" the input auto-correlation. For convolutional layers, we define a similar factorization E. M @ IE [[(x=: for the layer-wise Hessian, but with different M motivated by Grosse& & Martens 2016). (SeelAppendis A.2)
In this paper we also note that the off-diagonal blocks of the full Hessian can also be decomposed similarly. We can then approximate each block using the Kronecker factorization, and when the input auto-correlation matrices are close to rank 1, this allows us to approximate the eigenvalues and eigenvectors the full parameter Hessian. The details of this approximation stated in Appendix B
Experiment Setup: We conduct experiments on the CIFAR-10, CIFAR-I 100 (MIT) (Krizhevsky 2009), and MNIST (CC BY-SA 3.0) (LeCun et al., 1998) datasets as well as their random labeled versions, namely MNIST-R and CIFARTO-R. We use PyTorch (Paszke et al. 2019) framework for all experiments. We used several different fully connected (fc) networks fc network with m hidden layers and nt neurons each hidden layer is denoted as F-n""), several variations of LeNet (LeCun et al. 1998), VGGII (Simonyan & Zisserman 2015), and ResNet18 (He et al. 2016). More representative results are included in Appendix EJ The eigenvalues and eigenvectors of the exact layer-wise Hessians are approximated using modified Lanczos algorithm (Golmant et al. 2018) which is described in detail in Appendix CJ We use "layer.ettwrr" to denote a layer of particular network. For example, conv2:LeNet5 refers to the second convolutional layer in LeNet5.
4.1 Approximation of Layer-wise Hessian and Full Hessian
We compare the top eigenvalues and eigenspaces of the approximated Hessian and the true Hessian. We use the standard definition of subspace overlap to measure the similarity between top eigenspaces. Definition 4.1 (Subspace Overlap). For k-dimensional subspaces U, V in Id (d 2 k) where the basis vectors u, 's and u,'s 's column vectors, with & the size i vector of canonical angles between U and V, we define the subspace overlap of U and V as
Overlap(U, V) := -IIUT V 1
cos $12/k.
Note that when k 1, the overlap is equivalent to the squared dot product between the two vectors.
Appoodinated
Appeodinated
Eineet
0.8
leneet
0.8
1.5-
0.6
0.6
0.2
122
0.2
0.0
0.0
10
20
40
20
40
The i-ta eigovalue
Top Rigerepuoc
The i-th rigenvalue
Top & Elgenspuce
(a) Top eigenvalues of layer-wise Hessian of fcl
(b) Top eigenspace of layer-wise Hessians
(c) Top eigenvalues of the full Hessian
(d) Top eigenspace of the full Hessian
Figure 2: Comparison between the approximated and true layer-wise Hessian of F-2002.
As shown in Fig. this approximation works reasonably well for the top eigenvalues and eigenspaces of both layer wise weight Hessians and the full parameter Hessian.
4.2 Eigenvector Correspondence for Layer-wise Hessians
Suppose the i-th eigenvector for Ejæx?] is vr and the j-th eigenvector for E[M] is u,. Then the Kronecker product [[M] @ E[xx²" has an eigenvector uj @ u, Therefore if the decoupling conjecture holds, one would expect that the top eigenvector of the layer-wise Hessian has clear correspondence with the top eigenvectors of its two components. Since u é v is just the flattened matrix uoT, we may naturally define the following reshape operation.
Definition 4.2 (Layer-wise Eigenvector Matricization). Consider a layer with input dimension n and output dimension m. For an eigenvector h E Rm of its layer-wise Hessian, the matricized form of h is Mat(h) e mx where Mat(h) h(i-1m++--
5

---

More concretely, to demonstrate the correspondence between the eigenvectors of the layer-wise Hessian and the cigenvectors of matrix E [ and E(xx ", we introduce "eigenvector correspondence matrices" as shown in Fig.3.
Definition 4.3 (Eigenvector Correspondence Matrices). For layer-wise Hessian matrix H Rmmxmm with eigenvectors h,,''. hmn and its corresponding auto-correlation matrix E E(xx")) E [axx with eigenvectors v1.-.., Un. The correspondence between v, and h2 can be defined as
Corr(w,, h2) == (aa(h))o:
For the output Hessian matrix E E|M] E R=*m with eigenvectors " Um. we can likewise define correspondence between v; and hy as
Corr(u; h2) == Mat(h,)) "u.l??
(10)
We may then define the eigenvector correspondence matrix between H and E[xx' as mn matrix whose j-h entry is Corr (v; h2). and the eigenvector correspondence matrix between H and M[M] as a m mn matrix whose j-th entry is Corr(u:, h2).
Intuitively, if the j-th entry of the correspondence matrix is close I, then the eigenvector h2 is likely to be the Kronecker product of v, (or u.) with some vector. If the decoupling conjecture holds, every eigenvector of the layer-wise Hessian (column of the correspondence matrices) should have perfect correlation of with exactly one of v, and one of u.
(a) H with E/ex" ?1
(b) H with [[M]
(c) H with [ax=")
(d) H with B(xx
Figure 3: Heatmap of Eigenvector Correspondence Matrices for fel:LeNet5, which has 120 output neurons. Here we take the top left corner of the eigenvector correspondence matrices. Similarities between (a)(c) and (b) respectively verify the decoupling conjecture.
In Fig. we can see that the correspondence matrices for the true layer-wise Hessian H approximately satisfies this property for top eigenvectors. The similarity between the correspondence patterns for the true and Kroneckor product approximated Hessian H also verifies the validity of the Kronecker approximation for dominate eigenspace. From Fig. 3a and Fig. 3c the top m eigenvectors of the true layer-wise Hessian and the approximated Hessian cigenvector of E[xx""] From Fig. 3b and diagonal pattern for both the true Hessian and the Kronecker approximation. Thus for small i we have h; R 01 & U,:
all highly correlated with v1. the first Fig. 3d the correspondence with the IE [[]] component has
4.3 Structure of Input Auto-correlation Matrix [[xx"]
To understand the structure of the auto-correlation matrix, a key observation that the input æ for most
layers are outputs of a ReLU, hence it is nonnegative. We can decompose the auto-correlation matrix
as IE [[(xx" ] = E[x T + [x - [[=))(x - E(x))". We denote Σe := E[[ [[x]) (x E [[=])"] the auto-covariance matrix. As every sample
nonnegative, the expectation [B[xx] [ has large norm and usually dominates the covariance matrix Ex
We empirically verified this phenomenon on various networks and datasets. The first eigenvector of E|æx?" has very high overlap with E|æ] (squared dot product mean: 0.997, range: 0.964- 1.000). Meanwhile |E(x E x || significantly larger than |x,]] in our experiments [|[ ||] |lllllll/l/ Σg
mean: 12.08, range: 2.28-30.03). This suggests that E[a] [ar]" is approximately equal to axr" and dominates the covariance Σ,, Similar phenomenon also exists for convolution layers. The
complete experiment results are provided in Appendix E. We also observe the E [aaa"" matrices are all close to rank throughout the training trajectory as shown in Appendix E.4

---

4.4 Low Rank Structure of [
Previous works observed the gap in Hessian eigenspectrun around the number of classes с (where e 10 in our experiments on CIFARIO and MNIST), Since E[æx") is close to rank and the Kronecker factorization is good approximation for top eigenspace. the top eigenvalues of layer-wise Hessian can be approximated as the top eigenvalues of E, M] multiplied by the first eigenvalue of Elax"") Thus, the top eigenvalues of Hessians should have the same relative ratios as the top eigenvalues of their corresponding E[M]'s. Therefore, the outliers should also appear in E MM.
-3
2.5
1.0
1.0
2.0
0.5
2.0
1.5
2-
1.5
(1.0
0
10
20
30
10
20
30
0
10
20
30
10
20
30
0
10
20
30
0
10
20
30
E|M]
Hg
[[M]
E[M]
11,
(a) fel:LeNet5 at initialization (C)
(b) fel:LeNet5 at minimum
(c) fel:LeNet5 at minimum
(CIFARIO).
(CIFARIO-R).
Figure 4: Eigenspectrum of the layer- layer-wise output Hessian E and the layer-wise weight Hessian H (r)w The vertical axes denote the cigenvalues. Similarity between the two eigenspectra is direct consequence of a low rank R[xx"" and the decoupling conjecture.
Fig. shows the similarity eigenvalue spectrum between M and layer-wise Hessians different situations, which agrees with our prediction. However, the eigengap only appear at initialization and at minimum for true labels Fig.. 4a and Fig. 4b), but not at minimum for random labels (Fig. 4c). To understand the structure of E| itself, we consider simplified setting where have random two-layer neural network with random data.
Theorem 4.1 (informal). For two-layer neural network with Gaussian input, at initialization, when network large, the output Hessian the first layer approximately rank (c 1) and the corresponding top eigenspace is R(W R(W(2)) (W((2). 1) and ((W denotes the row space of the weight matrix (2) of the second layer.
The formal statement of this theorem and the full proof is in Appendix H| The closed form cal- culation can be heuristically extended to the case with multiple layers, that the top eigenspace of the output Hessian of the k-layer would be approximately R(((()) {S(()) 1}. where S(A) = ((=-1).. (I++1)
and (+)( is the row space of (AA)
Though result was only proven for random initialization and random data, observe that this subspace also has high overlap with the top eigenspace of output Hessian at the minima of models trained with real datasets. The corresponding empirical results shown in Appendix H.4
5
Understanding Structures of Layer- Hessian
5.1 Eigenspace Overlap of Different Models
Several interesting structures of layer-wise Hessians can be explained using the decoupling conjecture and cigenvalue correspondence. Consider models with the same network structure that are trained on the same dataset using different random initializations, despite no obvious similarity between their parameters, we observe surpisingly high overlap between the dominating eigenspace of their layer-wise Hessians.
Fig. includes 4 different variants of LeNet5 trained on CIFAR10, 3 different variants of ResNet18 trained on CIFAR 100, and 3 different variants of VGG1 I trained on CIFAR100, For each structural variant, models are trained from independent random initializations. We plot the average pairwise overlap between the top eigenspaces of those models' layer-wise Hessians. In each figure, we vary the number of output neuron/channels. For the same structure, the top eigenspaces of different models exhibits a highly non-trivial overlap which peaks near the dimension of the layer's output.
7

---

0.5
0.4
0.1
Ovezlap
0.3
Ovezlpp
Overlap
0.4
0.2
60
48
20 18
0.1
0.1
150
0.0
0.0
0.0
0
50
100
100
100
200
Dimennsion
l
Dimensson &
Dimension &
(a) conv 12:ReeNeee888 (CIFAR100) with 48/64/80 output channels
VGGv61 (CIFAR100) with (b)
(c) fel:L.NNet5 (CIFAR 10) with
48/64/80 output channels
80/100/120/150 output neurons
Figure 5: Overlap between the top & dominating eigenspace of different independently trained models. The overlap peaks at the output dimension m. The eigenspace overlap is defined in Definition 4.1
As we observed in Section 4.3 the auto-correlation matrix E[æx") is approximately [[=]][[a) Thus if the i-th eigenvector of OE [M] is ug the i-th eigenvector of the layer-wise Hessian would be close to u, & [ax)] where R[[x] = B[æz]/ [[]]]] [l. is normalized B[ar). Even though the directions of u, s can be very different for different models, at rank m these vectors always span the entire space, as a result the top-m eigenspace for layer-wise Hessian is close to I, @ E[x]]
Now suppose we have two different models with [[x]] and B[x] respectively. Their top-m eigenspaces are close to Im @ ((x)) and Im @ [[ac] respectively. In this case, it is easy to check that the overlap at m is approximately (B[az]] [B(x]2)² Since B[ær] and [B/z]]2 are the same for the input layer and all non negativee for other layers, the inner-product between them is large and the overlap is expected to be high at dimension
Our explanations of the overlap relies two properties: the auto-correlation matrix needs to be very close to rank I and the eigenspectrum of output Hessian should have heavy-tail. While both are true in many shallow networks and in later layers of deeper networks, they are not satisfied for earlier layers of deeper networks. In Anppendiix F.3 we explain how one can still understand the overlap using correspondence matrices when the above simplified argument does not hold.
5.2 Dominating Eigenvectors Layer-wise Hessian are Low Rank
A natural corollary for the Kronecker factorization approximation of layer-wise Hessians is that the eigenvectors of the layer-wise Hessians are low rank. Let h; be the i-th eigenvector of layer-wise Hessian. The rank of Mat(h;) can be considered as an indicator of the complexity of the eigenvector. Consider the case that h; is one of the top eigenvectors. From Section 5.1 we have h, u, & B[x]] Thus, Mht R u. Raz)", which is approximately rank 1.
Experiments shows that first singular values of Mat(h,) divided by its Frobenius Norm are usually much larger than 0.5, indicating the top eigenvectors of the layer-wise Hessians very close to rank 1. Fig. 22 shows first singular values of Mat(h;) divided by its Frobenius Norm for from to 200. We see that the top eigenvectors of the layer-wise Hessians are very close rank 1.
convo
CODV2
kl
fx2
R3
0.8
0.6
0.2
Figure 6: Ratio between top singular value and Frobenius norm of matricized dominating eigenvectors. (LeNet5 on CIFAR10). The horizontal axes denote the index i of eigenvector h4, and the vertical axes denote || Mat(h.))]/I] Mai(h:)ll::
8

---

5.3 Batch Normalization and Zero-mean Input
According to our explanation, the good approximation and high overlap of top eigenspace both depend on the low rank structure of B[æx" Also, the low rank structure caused by the fact that [|=] dominates Σæ in most cases. Therefore, it's natural to conjecture that models trained using Batch Normalization (BN) (Ioffe Szegedy 2015) will change these phenomena as [[(r] will be zero and [xxx"' Σn for those models. Indeed, as shown in Ghorbani et al. (2019), BN suppresses the outliers in the Hessian eigenspectrum and Papyan (2020) provided an explanation.
We experiment on our networks with BN. The results are shown in Appendix F.4. We found that E[æx=" is no longer close rank 1 for models trained with BN. However, E [/æ=="] still have few large eigenvalues. In this case, all the previous structures (c outliers, high eigenspace overlap, low rank eigenvectors) become weaker. The decoupling conjecture itself also becomes less accurate. However, the approximation still gives meaningful information.
6 Tighter PAC-Bayes Bound with Hessian Information
PAC-Bayes bound is commonly used to derive upper bounds on the generalization error.
Theorem 6.1 (PAC-Bayes Bound). (McAllester 1999. Langford & Seeger 2001 With the hypothesis space H. parametrized by model parameters. For any prior distribution P in H that is chosen
independently from the training set S, and any posterior distribution Q in H whose choice may inference S, with probability 5,
DKL (Q]I)) + log
[S]
DKL (E(Q)]le(Q)) S
(1I)
S
where e(Q) is the expected classification error for the posterior over the underlying data distribution and é(Q) is the classification error for the posterior over the training set.
Intuitively, if one can find a posterior distribution Q that has low loss on the training set, and is close to the prior P. then the generalization error on Q must be small. Dziugaite & Roy (2017) uses optimization techniques find an optimal posterior in the family of Gaussians with diagonal covariance. They showed that the bound is nonvacuous for several neural network models.
We follow Dziugaite& & Roy 2017) set the prior P to be multi- variant Gaussian. The covariance is a multiple of identity. Thus, is invariant with respect to the change of basis. For the posterior, when the variance in one direction is larger, the distance with the prior decreases; however this also has the risk of increasing the empirical loss over the posterior. In general, one would expect the variance to be larger along a flatter direction in the loss landscape. However, since the covariance matrix of Q is fixed to be diagonal in Dziugaite & Roy (2017). the search of optimal deviation happens in standard basis vectors which are not aligned with the local loss landscape.
Using the Kronecker factorization as in Eq. (7) we can approximate the layer-wise Hessian's 's eigenspace. We set Q to be Gaussian whose covariance is diagonal in the eigenbasis of the layer- wise Hessians. We expect the alignment of sharp and flat directions will result in a better optimized posterior Q and thus tighter bound classification error.
We perform the same optimization process as proposed by Dziugaite Roy (017). Our algorithm is called Approx Hessian (APPR) when we fix the layer-wise Hessian eigenbasis to the one at the minimum and Iterative Hessian (ITER) when we update the cigenbasis dynamically with the mean of the Gaussian which is being optimized. To accelerate Iterative Hessian, we used generalization of Theorem 4.1 to directly approximate the output hessian, which is then used to compute the eigenbasis. We call this variant algorithm of Iterative Hessian with approximated output Hessian (ITER.M). The results of this variant are only slightly worse than Iterative Hessian, which also suggests the approximation of the output hessian is reasonable.
We used identical dataset, network structures and experiment settings as in Dziugaite & Roy 2017 with a few adjustments in hyperparameters. We also added T-200 used in Section 4 T-600 in and T-200% are trained on standard MNIST while all others are trained on MNIST-2 (see Appendix D. The results are shown in Table with confidence of 0.965.

---

Table Optimized PAC-Bayes bounds using different methods. T-n and R-n" represents network F-n²" trained with true/random labels. TESTER. gives the empirical generalization gap. BASE represents the bound given by the algorithm proposed by Dziugaite & Roy 2.17). APPR, ITER, and
ITER.M represents the bound given by our algorithms.
Model
T-600
T-1200
T-300
T-600
R-600
T-60010
T-200
TestEr.
0.015
0.016
0.015
0.015
0.493
0.018
0.021
BASE
0.154
0.175
0.169
0.192
0.605
0.287
0.417
APPR (ours)
0.146
0.173
0.142
0.171
0.565
0.242
0.273
ITER (ours)
0.120
0.142
0.125
0.146
0.568
0.213
0.215
ITER.M (ours)
0.126
0.149
0.131
0.150
0.562
0.223
0.273
We also plotted the final posterior variance, 8 for network T-200 in Fig. 7 APPR, ITER, and ITER.M, variance. This agrees with our presumption that top eigenvectors are aligned with sharper directions and should have smaller variance after optimization. Detailed algorithm description and experiment results are shown in Appendix G
For our algorithms, that direction associated with larger eigenvalue has smaller
BASE
Avvn
Irim
Tras.M
10"3
11
10
10
103000
Figure 7: Optimized posterior variance s using different algorithms (fe1:T-2002 trained on MNIST). The horizontal axis denotes the eigenbasis ordered with decreasing eigenvalues. The abbreviation of algorithms are the same as in Table
Conclusions
In this paper identified Hessians for neural networks. Specifically, the eigenspace overlap reveals novel similarity between different models. We showed that under decoupling conjecture, these structures can be explained by a Kronecker factorization. We analyze each component in the factorization, in particular we prove that the output Hessian is approximately rank gives simple heuristic formula to estimate the top-eigenspace of the Hessian. As proof of concept, we showed that these structures can be used to find better explicit generalization bounds. Since the dominating eigenspace of Hessian, which is the sharpest directions on the loss landscape, plays an important role both generalization (Keskar et al. 2017 Jiang al. 2020) and optimization (Gur-Ari et al. 2018), we believe our new understanding can benefit both fields. We hope this work would be a starting point towards formally proving the structures of neural network Hessians.
surprising structures in the dominating eigenspace of layerwise
for random two-layer neural networks. Our proof
Limitations and Open Problems Most of our work focuses on the layerwise Hessian except for (Section 4.1 and Appendix B The eigenspace overlap phenomenon depends on properties of auto-coorelltion and output Hessian, which are weaker for earlier layers of larger networks. Our theoretical results need to assume the parameters are random, and only applies to fully-connected networks. The immediate open problems include why the decoupling conjecture is correct and why the output Hessian is low rank in more general networks.
10

---

References
Catoni, O. Pac-bayesian supervised classification: the thermodynamics of statistical learning. arXiv preprint arXiv:0712.02488 2007.
Dangel, F., Harmeling, S., and Hennig. P. Modular block-diagonal curvature approximations for feedforward architectures. International Conference on Artificial Intelligence and Statistics, pp- 799-801 808, 2020.
Dziugaite, G. K. and Roy, D. M. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In Proceedings of the Thirty-Thine Conference on Uncertainty Artificial Intelligence, UAI, 2017.
Fort, and Ganguli, S. Emergent properties of the local geometry of neural loss landscapes. arXiv preprint arXiv:1910.05929. 2019.
George, T. Laurent, C., Bouthillier, X., Ballas, N., and Vincent, P. Fast approximate natural gradient descent in a kronecker factored eigenbasis. In Advances in Neural Information Processing Systems, pp. -5500955000 2018.
Ghorbani, B., Krishnan, S., and Xiao, Y. An investigation into neural net optimization via hessian cigenvalue density. International Conference on Machine Learning, pp. 2232-2241, 2019.
Glorot, X. and Bengio. Y. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference artificial intelligence and statistics, pp. 249-256, 2010.
Golmant, N., Yao, Z., Gholami, A., Mahoney, M., and Gonzalez, J. sytorch-hessian-cieeninng efficient pytorch hessian cigendecomposition, 2018. /tth:::://ttthtubb -ytorch- hessi an eigenthings
//githbbtttttttt com/noahgolmmat/
Goodfellow, I., Bengio, Y., and Courville, A. Deep Learning, MIT Press, 2016.
Grosse, R. and Martens, J. A kronecker-factored approximate fisher matrix for convolution layers. In International Conference on Machine Learning, pp. 573-582, 2016.
Gur-Ari, G., Roberts, A., and Dyer, E. Gradient descent happens a tiny subspace. arXiv preprint arXiv:1812.04754, 2018.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 770-778, 2016. doi: 10.1 N0.100//CVR.2201..990
Heskes, T. On "natural" learning and pruning in multilayered perceptrons. Neural Computation, 12 (4):881-901, 2000.
loffe, S. and Szegedy, C. Batch normalization: Accelcrating deep network training by reducing internal covariate shift. In International Conference on Machine Learning, pp. 448-456 2015.
Jacot, A., Gabriel, F., and Hongler, C. The asymptotic spectrum of the hessian of DNN throughout training. In 8th International Conference on Learning Representations, ICIR, 2020.
Jastrzebski, S., Kenton, Z., Ballas, N., Fischer, A., Bengio, Y., and Storkey, A. J. On the relation between the sharpest directions of DNN loss and the SGD step length. In 7th International Conference on Learning Representations, ICLR, 2019.
Jiang, Y., Neyshabur, B., Mobahi, H., Krishnan, D., and Bengio, S. Fantastic generalization measures and where to find them. International Conference on Learning Representations, 2020.
Karakida, R., Akaho, S., and Amari, S.-i. The normalization method for alleviating pathological sharpness in wide neural networks. In Advances in Newral Information Processing Systems, volume 32, pp. 6406 6406-6416, 2019a.
Karakida, R., Akaho, S., and Amari, S.-i. Pathological spectra of the fisher information metric and its variants in deep neural networks. arXiv preprint arXiv:1910.0599.. 2019b.
11

---

Karakida. R., Akaho, S., and Amari, S.-i. Universal statistics of fisher information in deep neural networks: Mean field approach. The 22nd International Conference on Artificial Intelligence and Statistics, pp. 1032-1041. PMLR, 2019c.
Keskar, N. S., Mudigere, D., Nocedal, J., Smelyanskiy, M., and Tang, P. T. P. On large-batch training for deep learning: Generalization gap and sharp minima. In 5th International Conference on Learning Representations, ICLR, 2017.
Kleinman, D. and Athans, M. The design of suboptimal linear time-varying systems. IEEE Transac- lions Automalic Control, (3(2):150-115 1968.
Krizhevsky, A. Learning multiple layers of features from tiny images. Technical report, 2009.
Langford, J. and Seeger, M. Bounds for averaging classifiers. Technical report, 2001.
Laurent, B. and Massart, P. Adaptive estimation a quadratic functional by model selection. Annals of Statistics, pp. 1302-1338, 2000.
LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P. Gradient-based learning applied document recognition. Proceedings of the IEEE, 86(11):22782233 -2244, 1998.
Li, X., Gu, Q.. Zhou, Y., Chen, T., and Banerjee, A. Hessian based analysis of sgd for deep nets: Dynamics and generalization. In Proceedings of the 2020 SIAM International Conference on Data Mining, pp. 190-198. SIAM, 2020.
Martens, and Grosse, R. Optimizing neural networks with kronecker-factored approximate curvature, In International conference on machine learning, pp. 2408-2417, 2015.
Martens, J., Ba, J., and Johnson, M. Kronecker- curvature approximations for recurrent neural networks. In International Conference on Learning Representations, 2018.
McAllester, D. A. Some pac-bayesian theorems. Machine Learning, 37(3)3355-36,, 1999.
Papyan, V. The full spectrum of deepnet hessians at scale: Dynamics with sgd training and sample size. arXiv preprint arXiv: 181 f. 1.07062, 2018
Papyan, V. Measurements of three-level hierarchical structure in the outliers in the spectrum of deepnet hessians. In International Conference Machine Learning, pp. 5012-5021, 2019.
Papyan, Traces of class/cross-class structure pervade deep learning spectra. arXiv preprint arXiv:2008.1 1865, 2020.
Paszke, A., Gross, S., Chintala, S., Chanan, G., Yang, E., Vetoo Z., Lin, Z., Desmaison, A., Antiga, L., and Lerer, A. Automatic differentiation in pytorch. Technical report, 2017.
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L.,, Desmaison, A., Kopf, A., Yang. E.,, Vitoo Z.., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S. Pytorch: An imperative style, high-performane deep learning library. Advances in Neural Information Processing Systems 32, pp. 8024-8035.
Curran Associates, Inc., 2019,
Sagun, L., Bottou, L.,, and LeCun, Y. Eigenvalues of the hessian in deep learning: Singularity and beyond. arXiv preprint 1rXiv:1611.074777 2016.
Sagun, I.., Evci, U., Güney, V. U., Dauphin, Y. N., and Bottou, L. Empirical analysis of the hessian of wver-parametrized neural networks. In 6th International Conference on Learning Representations, ICLR 2018, Workshop Track Proceedings, 2018.
Schoenholz, S. S., Gilmer, J., Ganguli, S., and Sohl-Dickstein, Deep information propagation. In International Conference on Learning Representations, ICLR, 2017.
Simonyan, K. and Zisserman, A. Very deep convolutional networks for large-scale image recognition. In Bengio, Y. and LeCun, (eds.), 3rd International Conference on Learning Representations, ICLR, 2015.
12

---

Singla, S., Wallace, E., Feng. S., and Feizi, S. Understanding impacts of high-order loss approx- imations and features in deep learning interpretation. In International Conference on Machine Learning, pp. 5848-8848888 2019.
Skorski, M. Chain rules for hessian and higher derivatives made easy by tensor calculus. arXiv preprint arXiv:1911.13292 2019.
Sra, S., Nowozin, S., and Wright, S. Optimization for machine learning. Mit Press, 2012.
Torralba, A., Fergus, R., and Freeman, W. T. 80 million tiny images: A large data set for nonparametric object and scene recognition. IEEE transactions on pattern analysis and machine intelligence, 30 (11):1958-1970, 2008.
Van der Maaten, L. and Hinton, G. Visualizing data using t-sne. Journal of machine learning research, 9(11), 2008.
Wang, C., Grosse, R., Fidler, S., and Zhang, Eigendamage: Structured pruning in the kronecker- factored eigenbasis. In International Conference on Machine Learning, pp. 1566-6575 2019.
Yao, Z., Gholami, A., Lei, Q., Keutzer, K., and Mahoney, M. W. Hessian- analysis of large batch training and robustness to adversaries. In Advances in Neural Information Processing Systems, pp- 4949-4959, 2018.
Yao, Z., Gholami, A., Keutzer, K., and Mahoney, M. Pyhessian: Neural networks through the lens of the hessian. arXiv preprint
2019.
Zhang, G., Sun, S., Duvenaud, D., and Grosse, R. Noisy natural gradient variational inference. In International Conference on Machine Learning, pp. 5852-5861, 2018.
Zhou, W., Veitch, V., Austern, M., Adams, R. P., and Orbanz, Non-vacuous generalization bounds at the imagenet scale: pac-bayesian compression approach. In International Conference Learning Representations, 2019.
Zhu, S. A short note on the tail bound of wishart distribution. arXiv preprint
2012.
13

---

Appendix Roadmap
Appendix A we provide the detailed derivations of Hessian for fully-connected and convolutional layers.
2. In/Appendix B provide detailed description the approximation of dominating eigen- vectors of the full hessian.
3. In/Appendix C
we explain how we compute the eigenvalues and cigenvectors of full and layer-wise Hessian numerically.
4. In AAppeddx D
give the detailed experiment setups, including the datasets, network structures, and the training settings we use.
5. In Anppeddixx we provide detailed experimental results that are not fully included in the
main text.
6. In|Appendix F found.
we provide additional and more general explanations of the phenomena
7. In Appendix G we give a detailed description of the PAC-Bayes bound that we optimize and the algorithm we use optimize the bound.
8. In Appendix H. we provided the complete statement and proof for Theorem 4.
A Detailed Derivations
A.1 Derivation of Hessian
For an input æ with label y. we define the Hessian of single input loss with respect to vector as
H2 x)) ()((f(())) y) Vil(z, y).
(12)
We define the Hessian of loss with respect to for the entire training sample as
Hc(0)=VV,c(0-- Σ v)l())(z))))))))
Σ
(13)
We now derive the Hessian for a fixed input label pair (I,, y). Following the definition and notations in Section 3. we also denote output as = (xx) We fix layer p for the layer-wise Hessian. Here the layer-wise weight Hessian is He (x) We also have the output for the layer as z 2 Since (x) only appear in the layer but not the subsequent layers, we can consider where only contains the layers after the p-th layer and does not depend on (p) Thus, using the Hessian Chain rule (Skorski) 2019). we have
fo(æ)) go (p()(p) (w, .))
(P( .)) ( Oz((0) T
Hz(2(p), xx)
Oz(p)
mibi
00(2, y) dz!²)
V2 (p)
win
(14)
Owi(p)
where z( is the ith entry of (P) and (P) is the number of neurons in p-th layer (size of (()).
Since (p) (1)) + ((p) and (P) vee(W(()) we have
0z(p)
Ow(() mn @ (p)
(15)
Since 0 Swiss does not depend on wp))) for all we have V2 2 O. Thus.
H₂ (P(w((), x) (Imin @ (p)) Hz(z(p), x)) (Imm @ (
(16)
We define (( He(z(p), w) as in Section that
(r) x)) (Imon 8 v0())) M (Iusio @ x MM(P) (P)T
(17)
14

---

We now look into M H2(2(2), 2).
Again we have (0(=(1)) and can use chain rule here,
Oz Oz(p)
Oz
de y)
Hz(z x)
Hi(z, x)
(18)
0z(p)
Ozi
By letting p := softmax(z) be the output confidence vector, we define the Hessian with respect to output logit z as Az and have
A He x) (2 y) diag(p) ppp".
(19)
according to (Singla al. (2019).
We also define the Jacobian of z with respect to (P)) (informally logit gradient for layer p) as 0=
Glf dz For FC layers with ReLUs,
consider ReLU after the p-th layer multiplying (P)) by an indicator function 200000 To use matrix multiplication, we can turn the indicator function
into a diagonal matrix and define D (p) where
(P)) = diag (120)00)
(20)
Thus, we have the input of the next layer as ((++)) = (P)) ((1), The FC layers can then be considered sequential matrix multiplication and have the final output
W ( ((-1)) (())))
(())
(21)
Thus,
Oz
G(p)
D(r),
(22)
@z(r)
Since G(P) is independent of 2(P), we have
V 2i = 00 Vi.
(23)
Thus,
M H. x)) G Ag
(24)
Moreover, loss Hessian with respect to the bias term ((p) equals to that with respect to the output of that layer 2(P), We thus have
Hy (b) x)) M
G(PTT' G.PTT A.G!P
(25)
The Hessians of loss for the entire training sample simply the empirical expectations of the Hessian for single input. We have the formula as the following:
He(w(P)) = E H.(w ,x) = M
(P)
(P)
x
(p)T)
(26)
Hz(b())) = (z(z(r)) = MP() = E C
A (p)
(27)
Note that we can further decompose A& = = Qx, where
Qx = diag (vP) (I. 1xp²)
(28)
with 1, is all one vector of size c, proved in Papyan (2099).
We can further extend the close form expression to off diagonal blocks and the bias entries to get the full Gauss-Newton term of Hessian. Let
G (1)T ar
G (1)T
G (2)
X (2) G (2)T
FT
(29)
Ge (LIT 8 (n)
G (LYT
The full Hessian given by
c
00 y
Hc(0)= = F
A F
+
(30)
15

---

A.2 Approximating Weight Hessian of Convolutional Layers
The approximation of weight Hessian of convolutional layer is a trivial extension from the approxi- mation of Fisher information matrix of convolutional layer by Grosse & Martens (016).
Consider a two dimensional convolutional layer of neural network with input channels and output channels. Let its input feature map X be of shape (n, X1, X2) and output feature map Z be of shape (m, P1, P2). Let its convolution kernel be of size K1 x K2. Then the weight W is of shape (m, n. ,, 2). and the bias b is of shape (m). Let P be the number patches slide over by the convolution kernel, we have P P.
Follow Dangel et al. (2020), we define Z E IRmXP as the reshaped matrix of and W E Rm Ks as the reshaped matrix of W. Define B E RmxPP by broadcasting b to P dimensions. Let X Rnk K2 x P be the unfolded X with respect to the convolutional layer. The unfold operation (Paszke ce al. 2019) is commonly used in computation to model convolution matrix operations.
After the above transformation, we have the linear expression of the p-th convolutional layer similar to FC layers:
Z (p) W (P) X (P) + (P))
(31)
We still omit superscription of (p) for dimensions for simplicity. We also denote 2(r) as the vector form of (P) and has size mP. Similar to fully connected layer, we have analogue of Eq. (17) for convolutional layer
H (w X)
I s x(x)) M (Im 00 X
(32)
where M H) X) and is a mP x mP matrix. Also, since convolutional layers can also be considered as linear operations (matrix multiplication with reshape) together with FC layers and ReLUs, Eq. (23) still holds. Thus, we still have
He(2()) (P) X)= M
G( A.
(33)
where G
मश and has dimension x mP,, although is cannot be further decomposed as direct multiplication of weight matrices as in the FC layers.
However, for convolutional layers, X (P) is matrix instead of vector. Thus, cannot make Eq. (32) into the form of a Kronecker product as in Eq. (17)
Despite this, it is still possible to have Kronecker factorization of the weight Hessian in the form
H (P) X)
M (P) @ X (
(34)
using further approximation motivated by Grosse & Martens 2016). Note that M (P) need to have
different shape (n m) from M (mP mP), since H, (P X) is mnn 122 mnK1K2
and X (P) x (P]T is nK 1K2 nK K12.
Since we can further decompose Ag = Qx Qa. we then have
M (p) G(P)TT A Go)) ())) (Q.Gz ())
(35)
We define ((6) Ge (p) Here Qu is c x and Ge (p) is e mP so that N is x mP. We can reshape ( into cP x m matrix Nr(p) We then reduce M(1) (mP x mP) into a m m matrix as
1 ( N
M (P)
(36)
The scalar 1. is normalization factor since squeeze dimension of size P into size 1.
Thus, we can have similar Kronecker factorization approximation as
(P) 15
H. (P)
x) E
Lm X x (p)
(()
M(1) (Im 2 X x
(37)
E (m)
(p)
E (D)
& E x (p) (()T
(38)
16

---

B Structure of Dominating Eigenvectors of the Full Hessian.
Although it is not possible apply Kronecker factorization the full Hessian directly, we can construct an approximation of the top eigenvectors and eigenspace using similar ideas and our findings.
In this section, we will always have superscript (p) for all layer-wise matrices and vectors in order to distinguish them from the full versions.
As shown in Appendix A. Eq. (30) we have the full Hessian of fully connected networks as
dl(z, y
H2 = E
F A
+ E
V
(39)
where
G (1)T (1)
G (1)2
Gn (2)2 & (2
FT
G (2)T
(40)
G (L) & x (n)
G (LYT
In order to simplify the formula, we define
x(P))
(41)
be the extended input of the p-th layer. Thus, the terms in the Hessian attributed the bias be included the Kronecker product with the extended input, and F` can be simplified as
Ga (1)T @ a())
G (2)r
@ a(2)
FT
(42)
Gx (L)T
② ((m)
As discussed in several previous works (Sagun et al., 2019), the full Hessian can be decomposed in the G-term and the H-ter. Specifically, the G-term
2016 Papyan 2018 2019: & Ganguli
is E FT Fa]. and the H-term is E Σ Di=1 21
08(2.y)
in Eq. (39)
Empirically, the G-term usually dominates the H-term, and the top eigenvalues and eigenspace of the Hessian are mainly attributed to the G-term. Since we focus the top eigenspace, we can approximate our full Hessian using the G-term, as
Hc(0) E E
(43)
In our approximation of the layer-wise Hessian H ((w()) Eq. (6) the two parts of the Kronecker (2) and the auto-correlation matrix of the input (fx(F) x (p)T), Although we cannot apply Kronecker factorization to [F A Fz]. we can still approximate its eigenspace using the eigenspace of the full output Hessian.
factorization are the layer-wise output Hessian EE E
Note here that the full output Hessian is not common definition. Let vi. = E m be the sum of output dimension of each layer. We define a full output vector ž R by concatenating all the layerwise outputs together,
(2)
(44)
(L)
17

---

We then define the full output Hessian is the Hessian w.r.t 2. Let the full output Hessian for single input æ be M2 E R Similar Eq. (26) be expressed as
M H)(3, A.
(45)
where
G C
GT
(46)
G
similar
Eq. (42) The full output Hessian for the entire training sample is thus
- E[MA] B G.)
(47)
We can then approximate the eigenvectors of the full Hessian Hc (0) using the eigenvectors of E Ma] Let the i-th eigenvector of (0) be and that of M[Ma] [ be u. We may then break up
u, into segments corresponding different layers as in
U1
(48)
u (L)
where for all layer p, ((())
Rm Motivated by the relation between Ge and Fu, the i-th eigenvector of (0) can be approximated as the following. Let
E ze
11)
E
(2)
w;;
(49)
LL)
E (1o(
u
We then have
w;
RV
(50)
We can then use the
process to get the basis vectors of the approximated eigenspace. Another reason for this approximation is that the expectation is the input of each layer [[(x(1)) dominates covariance as shown in Appendix F.T Thus, the approximate is accurate for top eigenvectors and also top eigenspace. For latter eigenvectors, the approximation would not be as accurate since this approximate loses information in the covariance of the inputs.
We also approximated the eigenvalues using this approximation. Let the i-th eigenvalue of Hc(0) be A and that of [M[M)] be o,, We have
A o: 1
(51)
Below we show the approximation of the eigenvalues top eigenspace using this method. The eigenspace overlap is defined as in Definition 4.1 We experimented on several fully connected networks, the results shown below are for F-2002 (same as Fig. 2d in the main text), F-200 F-600 and F-6008, all with dimension 50.
18

---

2.5
Approsimated
2.0
Fsoed:
0.8
1.5
1.0
0.5
0.2
0.0
10
20
30
10
50
20
40
The i-th rigenvalue
Top & Eigenspare
(a) Rigenvalues for F-200
(b) Eigenspace overlap for F-200
1.5
Approxinated
1.0
0.4
0.5
0.0
0.0
10
20
30
40
50
The i-th eigenvalue
Top & Eigensparn
Eigenvalues for F-200*
(d) Eigenspace overlap for F-200
1.25
Apprasimated Exnet
1.00
0.75
0.6
0.50
0.25
0.2
0.00
0.0
20
40
50
20
40
8-lh eigenvalue
& Eigenspace
(e) Eigenvalues for F-600*
(I) Eigenspace overlap for F-600
Approsinstted Exneet
0.8
0.10
0.4
0.2
0.00
10
20
30
10
50
20
30
40
50
The i-th eigenvulue
Top & Eigenspace
(g) Eigenvalues for 8-600
(h) Eigenspace overlap for F-600
Figure 8: Top 50 Eigenvalues and Eigenspace approximation for full Hessian
19

---

C Computation of Hessian Eigenvalues and Eigenvectors
For Hessian approximated using Kronecker factorization, compute [[ and E[xx"] explicitly. Let m and v be an eigenvector of [[M and E[æx²" respectively, with corresponding eigenvalues Am and Ap. Since both matrices are positive semi-definite m&v is an eigenvector of [ & x=x" with eigenvalue Amdo. >g. In this way, since E[M has m eigenvectors and EJ [[xx") has eigenvectors, we can approximate mn eigenvectors for the layer-wise Hessian. All these calculation can be done directly.
However, it is almost prohibitive calculate the true Hessian explicitly. Thus, methods with automatic differentiation (Paszke et al. (217)) to calculate them. The packages we use Golmant et (2018) and use the Lanczos method in most of the calculations. We also use package in Yao et al. 2019) as reference.
numerical
For layer-wise Hessian, modified the Golmant al. (2018) package. In particular, the package relies on the calculation of Hessian-vector product v. where is a vector with the same size as parameter 0. To calculate eigenvalues and eigenvectors for layer-wise Hessian at the p-th layer, cut the v into different layers. Then, we only leave the part corresponding to weights of the p-th layer and set all other entries to 0. Note that the dimension does not change. We let the new vector be ((P) and get the value of Hv(r) using auto differentiation. Then, we do the same operation to u and get (P).
20

---

D
Detailed Experiment Setup
D.1
Datasets
We conduct experiment CIFAR-10, CIFAR-100 (MIT) (Krizhessky 2009) (https:/ /t::::/wwwr www toronto. edu/"kriz porontoo.du//krizz. cifar html), and MNIST (CC BY-SA 3.0) (LeCun et al. 1998) (http:// yann. Lanun cww/ exdb/ wwt/tm/sst///t////m/mmmmmmm The datasets are downloaded through torchvision (Paszke et al. 2019) (https://pytorrd) org/ woggvvisie visiiiiiiiii/// stable/ /nwdexin.... html). We used their default splitting of training and testing set.
To compare our work on PAC-Bayes bound with the work of Dziugaite & Roy (2017) 7), we created custom dataset MNIST-2 by setting the label of images 0-4 to 0 and 5-9 to 1. We also created random- labeled datasets MNIST-R and CIFARIO-P by randomly labeling the images from the training set of MNIST and CIFARIO. The dataset information is summarized in Table 2
Table 2: Datasets
# Data Points
Dataset
Train
Test
Input Size
Classes
Label
CIFAR 10
50000
10000
3 x 32 x 32
10
True
CIFAR 10-R
50000
10000
3 32
32
10
Random
CIFAR 100
50000
10000
3 x 32 x 32
100
True
MNIST
60000
10000
28 28
10
True
MNIST-2
60000
10000
28 x 28
True
MNIST-R
60000
10000
28 28
10
Random
All the datasets (MNIST, CIFAR-10, and CIFAR-100) used publicly available. According to their descriptions on the contents and collection methods, they should not contain any personal information offensive content. MNIST is remix of datasets from the National Institute of Standards and Technology (NIST), which obtained consent for collecting the data. However, we also note that CIFAR-I0 and CIFAR-I00 are subsets of the dataset 80 Million Tiny Image (Torralba et al., 2008) hhtp:://group: groups csail. .mi.. intt.edd/visse sedu/vision/TinyImggoo collection and includes some offensive images.
which used automatic
D.2 Network Structures
Fully Connected Network: We used several different fully connected networks varying in the number of hidden layers and the number neurons for each hidden layer. The output all layers except the last layer passed into ReLU before feeding into the subsequent layer. As described in Section 4.1 we denote fully connected network with m hidden layers and n neurons each hidden layer by F-n For networks without uniform layer width, denote them by sequence of numbers (e.g. for a network with three hidden layers, where the first two layers has 200 neurons each and the third has 100 neurons, we denote it as F-200--100), For example, the structure of F-200 is shown in
Table 3
Table 3: Structure of F-2002 on MNIST
Name
Module
In Shape
Out Shape
Flatten
(28,28)
784
fel
Linear(784, 200)
784
200
ReLU
200
200
4 fc2
Linear(200, 200)
200
200
ReLU
200
200
fc3
Linear(200, 10)
200
10
output
21

---

LeNet5: We adopted the LeNet5 structure proposed by LeCun et al. (9988 for MNIST, and slightly modified the input convolutional layers to adapt the imput of CIFAR- 10 dataset. The standard LeNet5 structure we used in the experiments is shown in Table 4 We further modified the dimension of fel and conv2 to create several variants for the experiment Sectiin 5.1] Take the model whose first fully connected layer is adjusted to have 80 neurons as example, we denote it as LeNet5-(fe1-80)
Table 4: Structure of LeNet5 on CIFAR-10
#
Name
Module
In Shape
Out Shape
1
conv1
Conv2D((3, 6,
(3, 32, 32)
(6, 28, 28)
2
ReLU
(6, 28, 28)
(6, 28, 28)
3
maxpooll
MaX2ooling22(2..2)
(6, 28, 28)
(6, 14, 14)
4
conv2 Conv2D(6, 16, 5)
(6, 14, 14)
(16, 10, 10)
5
ReLU
(16, 10, 10)
(16, 10, 10)
6
maxpool2
M2XPooling22(2.22
(16, 10, 10)
(16, 5)
7
Flatten
(16, 5,
400
8
fel
Linear(400, 120)
400
120
9
ReLU
120
120
10
fc2
Linear( (120, 84)
120
84
11
ReLU
84
84
12
fc3
Linear(84, 10)
84
10
output
Networks with Batch Normalization: In Appendix F.4 conducted several experiments regard- ing the effect batch normalization on our results. For those experiments, we use the existing structures and add batch normalization layer for each intermediate output after it passes the ReLU module. In order for the Hessian to be well-defined, we fix the running statistics of batch normaliza- tion and treat it linear layer during inference. We also turn off the learnable parameters A and [Ioffe & Szegedy 2015) for simplicity. For network structure X, we denote the variant with batch normalization after all hidden layers X-BN, For example, the detailed structure LeNet5-BN is shown in Table
Table 5: Structure of LeNet5-Be -BN on CIFAR-10
#
Name
Module
In Shape
Out Shape
1
convl
Conv2D(3, 6,
32, 32)
(6, 28, 28)
2
ReLU
(6, 28, 28)
(6, 28, 28)
3
BatchNorm2D
(6, 28, 28)
(6, 28, 28)
4
maxpool1
MuxPoolingg2D(2.2)
(6, 28, 28)
(6, 14, 14)
5
conv2
Conv2D(6, 16, 5,
(6, 14, 14)
(16, 10, 10)
6
ReLU
(16, 10, 10)
(16, 10, 10)
7
BatchNorm21D
(16, 10, 10)
(16, 10, 10)
8
maxpool2
MaxPooling2DD2.22)
(16, 10, 10)
(16, 5, 5)
9
Flatten
(16,
400
10
fel
Linear(400, 120)
400
120
11
ReLU BatchNorm Linear(120, 84)
120
120
12
120
120
13
fe2
120
84
14
ReLU
84
84
BatchNorm
84
16
fc3
Linear(84, 10)
84
10
output
22

---

To verify that our results apply to larger networks, we trained number of variant of VGGII (originally named VGG-A the paper, but commonly refered as VGG11) proposed by [Simonyan & Zisserman 2015). For simplicity, we removed the dropout regularization in the original network. To adapt the structure, which is originally designed for the 3 x 224 x 224 input of ImageNet, to 3 x 32 x 32 input of CIFAR-10.
Variants of VGG11:
Since the original VGGI1 network dimensions, we reduce the number of output channels of each convolution layer in the network to 32, 48, 64, 80, and 200. We denote the small size variants as VGGI1-W32, VGG11-W4R VGGI1-W64, VGG1 1-00 and VGGI 1- -W200 respectively. We use convl I conv8 and fel to denote the layers of VGGI1 where convl closest to the imput feature and fel the classification layer.
large for computing the top eigenspace up to hundreds of
Variants of ResNet18: We also trained number of variant of ResNet18 18 proposed by He et al. (2016). As batch normalization will change the low rank structure of the auto correlation matrix and reduce the overlap, we removed all batch normalization operations. Following the adaptation of ResNet to CIFAR dataset as https: //gp::://itte //githubt com com/kuanglli//yytoooeeeooooo cifar) we changed the input size to 3 x 32 32 and added Ixl convolutional layer for each shortcut after the first block.
Similar to VGGI1, I, we reduce the number of output channels each convolution layer in the network 48, 64, 80. We denote the small size variants as ResNet18- B-44, ResNet18-WGo and ResNet18- W80 respectively.
We use convl conv 17 and fel to denote the layers of the ResNet18 backbone where convl I is closest to the imput feature and fel the classification layer. For the 1x1 convolutional layers in the shortcut, denote them by sc-convl sc-conv3. where sc-convl is the convolutional layer the shorteut of the second ResNet block and sc-conv3 is the convolutional layer on the shortcut of the fourth ResNet block.
D.3 Training Process and Hyperparameter Configuration
For all datasets, we used the default splitting of training and testing set. All models (except explicitly stated otherwise) are trained using batched stochastic gradient descent (SGD) with batch-size 128 and fixed learning rate 0.01 for 1000 epochs. No momentum and weight decay regularization were may assume that the final models are
used. The loss objective converges by the end of training, at local minima.
For generality we also used training scheme with fixed learning rate at 0.001, and a training scheme with fixed learning rate 0.01 with momentum of 0.9 and weight-decay factor of 0.0005. Models trained with these settings will be explicitly stated. Otherwise we assume they were trained with the default scheme mentioned above.
Follow the default initialization scheme of Py PyTorch(Paast Paszke et al. 2019), the weights of linear layers and convolutional layers are initialized using the Xavier method (Glorot Bengio 2010). and bias of each layer are initialized to be zero.
23

---

E Additional Experiment Results
E.1 Low Rank Structure of Auto-correlation Matrix E(ax""
We have briefly discussed about the autocorrelation matrix E[æx") being approximately rank in Section 4.3 in the main text. In particular, we claimed that the mean of layer input dominate the covariance, that E [[xx'') se E[x]E[<"") In this section we provide some additional empirical results supporting that claim.
We use two metrics to quantify the quality of this approximation: the squared dot product between normalized EE[x] and the first eigenvector of E[xx:" and the ratio between the first and second eigenvalue of TE [xxx" Intuitively if the first quantity close to I and the second quantity large, then the approximation is accurate.
Formally. for fully connected layers, define B[æ] as the normalized expectation of the layer input æ, namely [[[=]]|[[[<=]]] For convolutional layers, following the notations in Appendix A.2 define [[[x]] the first left singular vector of E. [[ where B(ar)] RnK,K, Abusing notations for simplicity, we use E [[x=") to denote the n. n K2 x nK K₂ matrix EXXxT) In this section we consider the
squared dot product between E[2x] and the first eigenvector of E[xx"), namely (u7" [[]])
For the spectral ratio, let X be the first eigenvalue of [[[===" and A2 be the second. We have
A
1 Ex1
||[ar] [[[]]][[]] E(a)"11
(52)
A2
Σx
where Ex is the covariance of æ. Thus, the spectral norm of E E[a]" divided by that of Ex gives lower bound to A A1/22. In experiments, usually have \ ||<x| |/Elll
Table 6: Squared dot product (uf [[x=))² and spectral ratio \ A2 for fully connected layers in selection of network structures and datasets. We independently trained 5 runs for each instance and compute the mean, minimum, and maximum of the two quantities over all layers (except the first layer which takes the input with mean-zero) in all runs.
(of É B(x))² min
\ / /A2
Dataset
Network
fc
mean
max
mean
min
max
MNIST
F-200
2
1.000
1,000
1.000
12.29
9.65
16.16
F-6002
2
0,999
0.999
0.999
12.00
11.42
13.00
F-600+
4
1,000
0.999
1.000
17.81
7.33
28.00
F-600$
0.991 8
0.965
1.000
6.63
2.28
11.15
CIFARIO
F-6002
2
0.999
0.998
1,000
9.24
4.74
13.74
F-15003
0.999 3
0.997
1.000
13.27
6.10
18.41
LeNet5
0,998 3
0.997
0.999
7.21
5.88
9.02
LeNet5-(fe1-80)
3
0.998
0.996
0.999
7.80
6.77
11.01
LeNet5-(fe1-100)
0.997 3
0.995
0.999
7.42
6.20
9.10
LeNe15-(fc1-150)
3
0.998
0.992
0.999
7.35
5.34
9.62
VGG11-W32
0,990
0.988
0.993
6.02
5.57
6,51
VGG11-W64
0.996 1
0.993
0.999
5.87
5.32
6.26
VGG11-W64
0.995
0.993
0.996
6.24
5.97
6.70
CIFARI00
VGG11-W43
0,999
0.999
0.999
17.861
15.456
20.491
VGG11-W
0.999
0.999
1.000
19.185
18.358
20.410
VGG11- 1-W80
0,999
0.999
1.000
19.455
18.120
21.450
ResNet 8- W48
1,000
1.000
1.000
28.23
27.37
29.27
ResNet18 8- W66
1.000
1.000
1.000
27.07
25.72
29.50
ResNet1 8-W80
1.000
1.000
1.000
28.23
25.98
30.03

---

Table 7: Squared dot product ((ar) and spectral ratio /2 for convolutional layers in the selection of network structures and datasets Table
(ufE(z))²
\1/2o min
Dataset
Network
conv
mean
min
max
mean
max
CIFARIO
LeNet5
1
0.999
0.998
0.999
15.87
11.15
27.20
LeNet5-(fc1-80)
0.998
0.998
0.999
12.36
9.53
13.36
LeNet5-(fe1-100)
1
0.999
0.999
0.999
19.49
16.69
21.92
LeNet5-(fc1-150)
1
0.999
0.998
0.999
12.86
7.65
16.34
VGG11-W32
7
0.995
0.991
0.999
5.31
2.39
9.09
VGG11-W64
7
0.997
0.993
1.000
5.76
2.50
9.98
VGG11-W64
7
0.998
0.995
1.000
5.81
2.53
10.62
CIFARI00
VGG11-W48
7
0.996
0.991
0.999
5.72
2.46
9.90
VGG11-W64
7
0.995
0.991
0.999
5.66
2.50
10.79
VGG11-W80
7
0.994
0.988
0.998
5.18
2.50
8.45
ResNet18-Wei
19
0.981
0.917
0.998
3.79
1.89
7.56
ResNett18-W
19
0.985
0.910
0.998
3.96
1.81
7.53
ResNet18-W
19
0.987
0.954
0.997
4.16
2.11
7.04
As we can see from Table and the autocorrelation matrix B[xx" for tully connected layers. Similar phenomenon also holds for convolutional layers in the modern architectures, but the spectral gap are generally smaller compared that of the fully connected layers.
Table a variety of settings, E[>]E[z]" indeed dominated
E.2 Eigenspace Overlap Between Different Models
The non trivial overlap between top eigenspaces of layer-wise Hessians one of our interesting observations that had been discusses in Section 5.1 Here we provide more related empirical results. Some will further verify our claim in Section 5.1 and some will appear be challenge that. Both results will be explained discussed more extensively in in[ppendixx
E.2.1 Overlap preserved when varying hyper-parameters:
We first verify that the overlap also exists for set of models trained with the different hyper- parameters. Using the LeNet5 (defined in Table 4 as the network structure. We train 6 models using the default training scheme (SGD, 1r=0.01, momentum=O), models using a smaller learning rate (SGD. 1r=0.00,, momentum=O), and 5 models using combination of optimization tricks (SGD, 1r=0.01, momentum-0.i. weight decay=0.0005), With these 16 models. we compute the pairwise eigenspace overlap of their layer- wise Hessians (120 pairs in total) and plot their average in Fig. 9 The shade areas in the figure represents the standard deviation. The pattern of overlap clearly preserved, and the position of the peak roughly agrees with the output dimension m. demonstrating that the phenomenon is caused by common structure instead similarities in training process.
m
0.20
0.5
0.5
0.5
0.4
0.15
0,3-
0.2
0.1
0.2
0.2
0.1
0.0
0.11
0.0
0.000
0.0
10
200
40
1000
200
100
2010
30
Diuomniio
Dimcesioo
Deameioo
Dincrrdon
Figure 9: Eigenspace overlap of different models of LeNet5 trained with different hyper parameters.
Note that for fc3 (the final output layer), we are not observing linear growth starting from 0 like other layers. This can be explained by the lack of neuron permutation. Related details will be discussed along with the reason for the linear growth pattern for other layers in Appendix F.3
25

---

E.2.2 Eigenspace overlap for convolutional layers in large models:
Even though the exact Kroneckor Factorization for layer-wise Hessians is only well-defined for fully connected layers, we also observe similar nontrivial eigenspace overlap for convolutional layers in larger and deeper networks including variants of VGGI I and ResNer18 8 datasets CIFAR 10 and CIFARI00 100. Some representative results are shown Fig. 25 and Fig. 11 For each model on each dataset, we independently train 5 models and compute the average pairwise eigenspace overlap. The shade areas represents the standard deviation.
For most of the convolutional layers, the eigenspace overlap peaks around the dimension which is equal the number of output channels of that layer, which similar the layers LeNet5 in Fig. 9 The eigenspace overlap of the final fully connected-laye also behaves similar to fe3::eNet5, which remains around constant then drops after exceeding the dimension of final output. However, there are also layers whose overlap does not peak around the output dimensions, (e.g. conv5 of Fig. Ilap. We will cluster these "failure cases" in the following paragraph.
Fig. 10b) and conv?
113
04-
044
0.5
044
0.4
0.33
0.2
110
0.35
2203
3.3
01005
0.00
0.1
0.1
8:005
332
0.05
1e0er 24 33
CERVO 33
10004 1. 30
040655
1.8
008
01003
3.10
(100
0.30
3.06
3.12-
3.00
0.50
3.64
0.61-
0.02
3.22
002
04077
330
2.
200
30
200
00
200
400
Dexemiice
(a) VGGI1-W32 (CIFARIO)
(b) VGGII-W200 (CIFARIO)
1.5
115
00
0.0
3.55
123
0.2
5.00
324
0.2
0.5
0.3
004
0.1
3.35
033
101
01
04663 25: 7.
3.00
033
033
0.5
1.33
5.32
5:3
2.23
0.3
02
0.2
03277
30
100
30
304
131
(c) VGG11-WAL (CIFAR100)
(d) VGG1I-W80 (CIFARI00)
Figure 10: Top Eigenspace overlap for varients of VGG11 CIFARIO and CIFAR 100
26

---

(.3)
1.00
310-
E.00
8.755
8.334
L0A22
80
180
16
160
180
36
33
008
314
64-
0.
24-
C4
14
02
12
37
30
199
000
301
Daanka
Dawn
(a) ResNet18- W48 (CIFARI00)
311
0.12
0.5
015
8.30
3000
GC322 6--
D4 64
6005 21
0005
01077.1
100
16
-0
180
120
180
00
10-
112
324
4.10
110
206
114
0.
E1-
00
20
180
30
(b) ResNet18 W66 (CIFARI00)
112.
1/2
E311
300
0002
180
300
010
K.70
0.35
1.30
1.30
115
1.77
C20
305
114
CI0P12
63-
E4-
03-
62
8.2
675
62x
32
200
010
180
101
(c) ResNet18 W80 (CIFARI00)
Figure 11: Top Eigenspace overlap for variants of ResNet18 on CIFAR100
27

---

E.2.3 Failure Cases
As in Fig. and Fig. 11 there is small portion of layers, usually closer to the input, whose eigenspace overlap does peak around the output dimensions. These layers can be clustered into the following two general cases.
a. Early Peak Low Overlap
For layers shown in Fig. 12 The overlap of dominating eigenspaces are significantly lower than the other layers. Also there exists small peak very small dimensions.
0.010
0.10
0.00
0.01
0.02
0.006
0.01
0.12
Outpes Diu:
Oxipes
Iro Outpes Dia:: 8]
100
200
900
0.000
100
300
0.000
0,000
Dinoctaion
0.015
0.01.55
0.00
0.04
0.015
01133
0.010
0.0000
0.055
0.0035
0.01
0.000
lihe.
0.0000
10
20
Dimousion (Zoom
(Zoom
(Zoom
(a) fc2:F-2002 (MNIST)
(b) BIGN::VGGL1-W222 (c) conv 2:GG:G-GGIWWB (d) con (CIFARIO)
(CIFAR 10FA
(CIFAR100)
Figure 12: Top eigenspace overlap for layers with an early low peak. Figures in the second row are the zoomed in versions of the figures in the first row.
b. Delayed Peak / Peak Doesn't Decline
For layers shown in Fig. 13 the top eigenspaces has nontrivial overlap, but the peak dimension is larger than predicted output dimension.
0.3-
0.125
0.15
0.100
0.2-
0.10
0.075
0.1
0.2
0.05
0.1
010.5
I.: Outjec Bin
Dla:
Oatpet Diu:
Ouppes Dio:
0.0
0
100
300
300
100
30
100
100
50
100
150
Dimumiio
Dinamaiu
Dinamaion
Dinamaioo
cony2-VGGiI VGG11--20 cony7-VGG11-Wal (CIFAR100)
(c) com77-VGG11-448 (d) conv Resseeee1
(CIFAR10)
(CIFARI00)
(CIFARI00)
Figure 13: Top eigenspace overlap for layers with delayed peak.
However, the existence of such failure cases does not undermine the theory of Kronecker factorization approximation. In fact, both "failure cases" appear because top hessian eigenspace not completely spanned by E(ær). and can be predicted by computing the auto correlation matrices and the output Hessians. The details will also be elaborated in Appendix F3 with the help of correspondence matrices.
28

---

E.3 Eigenvector Correspondence
Here we present the correspondence matrix for fel, fc2, conv1, and conv2 layer of LeNet5. The top eigenvectors for layers shows a strong correlation with the first eigenvector of R[ax") (which approximately E[z]). However, the diagonal pattern in the correspondence matrix with E[M] for fc2 not clear as the one for fel.
00
40
60
60
25
50
75
100
125
160
175
25
75
100
125
160
176
(a) Correspondence with E((rx""
")-
(b) Correspondence with [[M].
Figure 14: Eigenvector Correspondence for fel:LeNet5. (m=1200
60
100
100
116
21
60
100
160
(a) Correspondence with E[ww")
(b) Correspondence with B[M].
Figure 15; Eigenvector Correspondence for fc2:LeNet5. (m==84)
(a) Correspondence with E((zz ")
(b) Correspondence with E[M].
Figure 16: Eigenvector Correspondence for conv I:LeNet5. (m=6)
(a) Correspondence with E[wx
(b) Correspondence with R[M].
Figure 17: Eigenvector Correspondence for conv2:LeNet5. (m=16)
29

---

E.4 Structure of E[xx") and [[M] During Training
We observed the pattern of "xx " matrix and E [MM] matrix along the training trajectory Fig.
Fig. 19p. It shows that E [xx' ] is always approximately rank 1, and E[M] always have around large eigenvalues. According our analysis, since the nontrivial eigenspace overlap is likely to be
consequence of approximately rank 1 E [x=="" we would conjecture that the overlap phenomenon likely to happen on the training trajectory as well.
Eigenvalues of E(xx") (normalized)
1.0
0.6
0.2
200
400
600
800
1000
Epoch
Figure 18: Top eigenvalues of E [xxx) along training trajectory. (fcl:LeNet5)
Eigenvalues of E (normalized)
1.0
Aga
0.6
ArT
0.4
0.2
0.0
200
400
600
800
1000
Epoch
Figure 19: Top eigenvalues of E| along training trajectory. (fcl:LeNet5)
30

---

F Additional Explanations
F.1
Outliers in Hessian Eigenspectrum
One characteristic of Hessian that has been mentioned by many the outliers the spectrum of eigenvalues. [Sagun et al. (2018) suggests that there is gap in Hessian eigenvalue distribution around the number of classes in most cases, where e = 10 in our case. A popular theory to explain the gap is the class logit clustering of the logit gradients (Fort & Ganguli 2019 Papyan 2019 2020). Note that these explanations can be consistent with our heuristic formula for the top eigenspace of output Hessian initialization- in the two-layer setting we considered the logit gradients are indeed clustered.
In the layer-wise setting, the clustering claim be formalized as follows: For each class k: [c] and logit entry E [c]. with Q be defined as in Fq. (28) and (x, y) as the input, label pair, let
dzx Qx
(53)
Oun (p)
Then at the initialization, for each logit entry J. {As.) hee[ is clustered around the "logit center" at the minima, for each class i, { ( hej is clustered around the "class center" Á A Ejelel A A,ll With the decoupling conjectures, may also consider similar claims for output Hessians, where
Á A E|e[c] A A,5:
Ozx
T J E
Qx
y
(54)
dz (P)
A natural extension of the clustering phenomenon on output Hessians then as follows: At the initialization, for each logit entry j. {T heel is clustered around r & [T ): at the minima, for each class i, J is clustered around r; A Ejelq(FT,)) Note that we have the layer-wise Hessian and wise output Hessian satisfying
())
E (JE(d)
[
A. M
E A.E
(55)
Low-rank Hessian at Random Initialization and Logit Gradient Clustering
We first briefly recapture our explanation on the low-rankness of Hessian at random initialization.
In In/Apenndix H. we have shown that for two layer ReLU network with Gaussian random initial- ization and Gaussian random input, the output hessian of the first layer M is approximately
1 (2)T AW We then heuristically extend this approximation to randomly initialized L-layer network, that with ((r) (L--). (.)t)) the output Hessian of the p-th layer H (p) can be approximated by M (p) where
M A S AS
(56)
4h-p
Since A is strictly rank c c with null space of the all-one vector, (P) is strictly rank c 1. Thus (P)) is approximately rank c - 1, and so is the corresponding layerwise Hessian according to the decoupling conjecture.
Now we discuss the connection between our analysis with the theory of logit gradient clustering. As previously observed by Papyan (2019), for each logit entry 1, (A hreje are clustered around the logit gradients Eee(s] Aijl. Similar clustering effects for (Ti, (ee)) were also empirically observed by our experiments. Moreover, through the approximation above and the decoupling conjecture, for each logit entry j, the cluster centers I' and A can be approximated by
S Q);
(57)
(E|æ) & ST (E)),
Following Papyan (2019). we used t-SNE (Van der Maaten & Hinton, 2008 to visualize the logit gradients. As we see in Fig. 20, the "logit centers" the clustering directly corresponds to the approximated dominating eigenvectors of the Hessian, which is consistent with our analysis.
31

---

Appecsinated light ontirs
onees 5,
X
Appisincatel eattrs 5
clotes uratien Cy
diattz cxciboes A.
20-
20-
-20
-21-
-31
LSNE
USNE
(a) Eigenspectrum of E at initialization.
(b) Clustering of T with logits at initialization.
(c) Clustering of A with logits at initialization.
Figure 20: Logit clustering behavior of A and P at initialization (fc1:T-2002)
Gradient Clustering at Minima Currently our theory does not provide an explanation to the low rank structure of Hessian at the minima. However we have observed that the class clustering of logit gradients does not universally apply all models at the minima, even when the models have around e significant large eigenvalues. As shown in Fig. 21 the class clustering is very weak but there are still around significant large eigenvalues. We conjecture that the class clustering of logit gradients may be sufficient but not necessary condition for the Hessian to be low rank at minima.
clase clutte corree I.
Charts swwatee Fuy
daas cope A.
clutee turzhee Dug
(2-
0.00
40-
3.03
50-
0.003
1000
-00-
0:00
-72
50
L.90E X
LSNNE X
(a) Eigenspectrum of B [MM] at minimum.
(b) Clustering of I with class at minimum. (c) Clustering of A with class at minimum.
Figure 21: Class clustering behavior of A and I at minimum. (fc1:T-2002)
F22 Dominating Eigenvectors of Layer-wise Hessian are Low Rank
natural corollary for the Kronecker factorization approximation of layer-wise Hessians is that the eigenvectors of the layer-wise Hessians are low rank. Let h; be the i-th eigenvector of layer-wise Hessian. The rank of Mat(h;) can be considered as an indicator of the complexity of the eigenvector. Consider the case that h; is one of the top eigenvectors. From Section 5.1 we have h; R u, 69 1E(x)). Thus, Mat(h;) ~ E.x]"T which is approximately rank 1. Experiments shows that first singular values of Mat(h;) divided by its Frobenius Norm are usually much larger than 0.5, indicating the top eigenvectors of the layer-wise Hessians are very close to rank 1. Fig. 22 shows first singular values of Mat (h.) divided by its Frobenius Norm for i from to 200. We can see that the top eigenvectors of the layer-wise Hessians are very close to rank 1.
32

---

1.0
//
0.8
0.4
0.2
30
40
20
10
20
Figure 22: Ratio between top singular value and Frobenius norm of matricized dominating eigenvectors. (LeNet5 on CIFARIO). The horizontal axes denote the index of eigenvector h; and the vertical axes denote || Mat(h+)ll/I) (attt))
Eigenspace Overlap of Different Models
From the experiment results in Appendix together with Fig. 5 we can see that our approximation and explanation stated in Section 5.1 of the main text is approximately correct but may not be so general explanation which addresses why the
accurate for some layers. We present overlap before rank-m grows linearly. We will also explain some exceptional cases as shown in Appendix E.2 and possible discrepancies of our approximation.
Let h; be the i-th cigenvector of the layer-wise Hessian Hc(w(P)), under the assumption that the autocorrelation matrix [B[xx") is approximately rank that E[æx") E[<]E]<>] T" for all S. m, we can approximate the h; u, @ [B[=]/|[[[<]]] where u, is the i-th eigenvector of EE [MM. Formally, the trend of top eigenspace can characterized by the following theorem. For simplicity of notations, we abuse the superscript within parentheses to refer the two models instead of layer number in this section.
Theorem F.1. Consider 2 different models with the same network structure trained on the same dataset. Fix the p-th hidden layer with input dimension n and output dimension m. For the first
model, denote its output Hessian as E [iM with eigenvalues Ji (1) (1) cigenvectors
(1) 2 Tm N 0 and T' denote its autocorrelation matrix as E[xx" with eigenvalues (1)
T2
(1) (1) 2 12 2 ... 7/ 2 0 and eigenvectors { + E R The variables for the second matrices are defined identically by changing 1 in the superscript parenthesis to 2.
Assume the Kronecker factorization approximation is accurate that (z) = [[M] @ E[xa²1))) and Hc(w(P))(2) as [[M](2) @ E(xx""((), Also assume the autocorrelation ma- trices of two models are sufficiently close to rank in the sense that Tm 71 > 7'')'/" (2) (2)
(1) and
Tm 71 > ri (2),,(2) Then for all S m, the overlap of top k eigenspace between their layerwise Hessians Hc (P)w()))(L and (zlw(P))(2) will be approximately l (t((") (())),, Consequently, the top eigenspace overlap will show a linear growth before it reaches dimension m. The peak at m is approximately (t, ?
Proof. Let h (2)
be the i-th eigenvector of the layer-wise Hessian for the first model H ())
and gi be that of the second model H Hu(w(()
(P)
Consider the first model. By the Kronecker
(1)) Ti (1)))())) (1) the top m eigenvalues of the layer-wise (1) (1) Consequently, for all i
factorization approximation, since 7m
r (1)T & ( Thus, for any h < m, we have its top k eigenspace as vi(2) @ ((1), where V (L) Rmmx has column vectors ri ...k
Hessian 1i (1) (1)
(
have h
(1)
Similarly, for the second model we have h
(2)
ri
(2)
1 and the top
...., . eigenspace v((2) @ x(2) where v:(2) has column vectors (2) the 2 models at dimension & is thus
(2) The eigenspace overlap of
Overlap (vi & tt
V (2) 3
(
1 k
V. (1)T (2)
(1X a 1lee
(58)
(1)) (12) O Overlap (v. (2)
33

---

Note that for all i S m. (2) !r! E IR", which is the space corresponding to the neurons. Since for hidden layers, the output neurons (channels for convolutional layers) can be arbitrarily permuted to give equivalent models while changing cigenvectors. For h; T, @ ti, permuting neurons will permute entries in rt. Thus, we can assume that for two models, r3 and ry are not correlated and thus have an expected inner product of V1/m.
Bl(r.") (2) 2 ) ( A
It follows from Definition 4.1 that E[Overlap(V
and thus the eigenspace overlap of dimension & would be approximately (t(1) (())).. explains the peak at dimension and the linear growth before it.
This
From our results on autocorrelation matrices in Section 4.3 and Appendix E.1 we have [1 o
and [2x) se (1 where É is the normalized expectation. Hence when = the overlap is approximately (B[ae) (1) [2))) Since [[ax]()) and E[ax](2 the identical for the input layers, the overlap is expected to be very high at dimension m for input layers. For other hidden layers in
ReLU network, a are output of ReLU and thus non-negative. Two non-negative vectors B[x][() and B[æ] (2)
still have relatively large dot product, which contributes the high overlap peak.
F331
The Decreasing Overlap After Output Dimension
Consider the (m 1)-th eigenvector h pm+ of the first model. Following the Kronecker factorization approximation and assumptions in Theorem F.1 have h h r1 St ( Since top eigenspace of the first model is approximately In @ 1 and t.2 is orthogonal to (22) the ( eigenvector will be orthogonal to the top m eigenspace of the first model. will also have low overlap with Im & Ef?) since (É(z)()) [[zz((2) is large.
Moreover, since the remaining eigenvectors of the autocorrelation matrix no longer has the all positive property as the first eigenvector and structure of the convariance Ex is directly associated with the ordering of the input neurons which are randomly permuted across different models, the overlap between other cigenvectors of the autocorrelation matrix across different models will be close to random, hence the overlap after the top m dimension will decrease until the eigenspaces has sufficiently many basis vectors to make the random overlap large.
F.3.2 The Output Layer
Note that for the last layer satisfying the assumptions in Theorem F.I] the overlap will stay high before dimension and be approximately (t1 +2)² since the output neurons directly correspondence classes, and hence neurons cannot be permuted. In this case, the overlap will be approximately (t1 t² for all dimension & < m. This is consistent with our observations.
0..
0.3-
0.2-
0.2
0.2
0.1-
0.1
0.1
Ourget
1000 Ohapet Jan 10
Chrppt Diu: Bi
H.e Dusjer Din 301
0.0
0.0
DU0
0.0
10
20
10
100
150
Dimonsion
Dimasoian
Dimutnion
(a) 2c3:F-2002 (MNIST)
(b) fe3:5eNee5 (CIFAR10)
(c) fe1:VGG1 BEGG111-W 1-2000 (CIFAR1O)
(d) fel:ResNet18-W6i (CIFARIO0)
Figure 23: Top eigenspace overlap for the final fully connected layer.
34

---

F.3.3 Explaining "Failure Cases" of Eigenspace Overlap
As shown in Fig. 12 and Fig. 13 the nontrivial top eigenspace overlap does not necessarily peak at the output dimension for all layers. Some layers has low peak very small dimensions and others has a peak at larger dimension. With the more complete analysis provided above, we now proceed explain these two phenomenons. The major reason for such phenomenons that the assumption of autocorrelation matrix being sufficiently close to rank I is not always satisfied. In particular, following the notations in Theorem F.I] for these exceptional layers we have Tm?1 < 7172-
We first consider the first phenomenon (early peak of low overlap) and take fc2:F-200" (MNIST) in example. Here Fig. is is identical to Fig. 12a which displays the early peak around 10.
001125
00
60-
25-
2.5
00103
10077
1.5
123
2-
Ourpes 201
150
100
010000
20
10
23
10
17A
17.5
(Zooas
KM
10
(a) Eigenspace overlap (zoomed in)
(b) Eigenspectrum of E|M] and E[æx""
(c) True Hessian with (d) Approximated Hessian B B(xc=")
with B B(xx'
Figure 24: Eigenspace overlap, eigenspectrum, and cropped (upper 20 x 20 block) eigenvector correspondence matrices for fc2:F-2002 (MNIST)
As shown in Fig. 24b, the second eigenvalue of the auto correlation E| is as large approx- imately 1/10 of the first eigenvalue. With the output Hessian have e e- 9 significant large eigenvalues as described in it has 710"1 71722 Thus through the Kronecker factorization ap- proximation, the top dimensional eigenspace is no longer simply Im & B[xx), but subset of top eigenvectors of the output Hessian Kroneckered with a subset of top eigenvectors of E B[x=") as reflected in Fig. 24d This "mixture" of Kronecker product is moreover verified in (ig. 24c
As reflected by the first row of Fig. 24c and Fig. 24d. for i S 9 we have h, R r, @ È R(æ:). which falls in the regime of Theorem F.1 Hence we are seeing an linearly growing pattern of the overlap for dimension less than 10 and reaches mean overlap of around 0.012 by dimension 9, If following this linear trend, the overlap would close to 0.25 by the output dimension of 200. However, since the 10-th eigenvalue of the output Hessian is significantly smaller, little of the 10-19 dimensional eigenspace were contributed by El(z). hence the overlap of dimension larger than 10 falls into the regime discussed in Appendix F.3.1 for which we see a sharp decrease of overlap after dimension 9. be used to predict when conditions
Note that this example shows that Kronecker factorization in Theorem F. fails and also predict the condition can be satisfied up to which dimension. As shown similar explanation also applies convolutional layers in larger networks.
Fig. 25
60
00
355
2.22
2.0
0.010
222
50
125
J..:: Ouipes Daac 20)
160
01000
175
17.5
21
Dinorsion
[[44]]
10
(a) Eigenspace overlap (zoomed in)
(b) Eigenspectrum of E|M] and E[æx"]
(c) True Hessian with (d) Approximated Hessian E[æx""
with B|xx'
Figure 25: Eigenspace overlap, eigenspectrum, and cropped (upper 20 x 20 block) eigenvector correspondence matrices for conv5: VGG1: 1- (CIFARIO)
35

---

We then consider the second phenomenon (delayed peak) and take con2::VGG11-W22 (CIFARIO) in as an example. Here Fig. 26a identical to Fig. 1 which has the overlap peak later than the output dimension 200. In this case, the second eigenvalue of the auto correlation matrix is still not negligible compared to the top eigenvalue. What differentiate this case from the first phenomenon that the eigenvalues of the output Hessian no longer has a significant peak instead it has heavy tail which is necessary for high overlap.
Towards dimension m there gradually exhibits higher correspondence to later eigenvectors of the input autocorrelation matrix and hence less correspondence to B[xx] This eventually results in the delayed and flattened peak.
10-
10-
0.2
23
0.1
31-
0.-
91-
91-
100
200
S00
050
0.3
-0
Dimension
(a) Eigenspace overlap (zoomed in)
(b) Eigenspectrum of
(c) True Hessian with
(d) Approximated Hessian
E [M] and E [axx"
E [ww=")
with E B(æx²)
Comressaece vatt [33
68-
0.4-
6G-
644
00
100
155
1000
100
200
500
400
50
100
150
100
120
200
350
100
(th elgenoeiter (sorrage armie Algoneettoe)
(1th elgeoooeet (somaage aomee elgenoeetoo)
(c) First Row of Correspondence Matrix of True Hessian with B[æx"]
(f) First Row of Correspondence Matrix of Approximated Hessian with B(xx*
Figure 26: Eigenspace overlap, eigenspectrum, and cropped (upper 50 x 50 block) eigenvector correspondence matrices for VGG/: (CIFAR1O)
Since the full correspondence matrices are too large to be visualized, we plotted their first rows up to 400 dimensions in Fig. 26c and Fig. 26f in which each dot represents the average of correlation with B[œx for the 10 eigenvector nearby. From these figures it is straightforward to see the gradual decreasing correlation with [[xx].
F.4 Batch Normalization and Zero-mean Input
In this section, we show the results on networks with using Batch normalization (BN) (Lffe & Szegedy 2015). For layers after BN. we have E[x] z 0 so that E| [|=]" no longer dominates Σn and the low rank structure of [[xx"" should disappear. Thus, we can further expect that the overlap between top eigenspace of layer-wise Hessian among different models will not have a peak.
Table 8 shows the same experiments done in Table 6. The values for each network are the average of 3 different models. is clear that the high inner product and large spectral ratio both do hold here, except for the first layer where there is no normalization applied. Note that we had channel- normalization (zero-mean for each channel but not zero-mean for x) for convl I in LeNet5 so that the spectral ratio is also small.
Fig. 27a shows that E [[>x?]] is no longer close rank 1 when having BN. This is as expected. However, B[<x7] still has a few large eigenvalues.
Fig. 27b shows the eigenvector correspondance matrix of True Hessian with E[ax=" for fel:LeNet5. Because E[[ma=") is no longer close to rank 1, only very few eigenvectors of the layer-wise Hessian will have high correspondance with the top eigenvector of E[æx"), as expected. This directly leads
36

---

Table Structure of E xx T for BN networks
(vf T[(=))2 min
X X2 min
Dataset
Network
fc
max
MNIST
F-200
0.062
0.001
0.260
1.16
1.04
1.30
BB6002-BN
2
0.026
0,000
0.063
1.13
1.02
1.26
F-600*-BN
4
0.027
0.000
0.146
1.11
1.03
1.19
CIFARIO
LeNet5-BR
3
0.210
0.001
0.803
1.54
1.20
1.89
to the disappearance of peak in top eigenspace overlap of different models, as shown in peak still exists in convl because BN is not applied to the input.
Fig. 28 The
fr1
001
00-
223
3.00
1.75
7.3
7.5
1.50
102
162
180
17.3
17.3
10
231
10
[
(b) True Hessian with (fel:LeNet5-BN)
(c) Approx Hessian with
(a) Eigenspectrum for E|æxx
Elax"
Figure 27: Eigenspectrum and Eigenvector correspondence matrices with E (wxx") for LeNe15-BN.
fel
fr/2
0.5
0.4
0.025
0.06
0.3
0.00
0.5-
0.01
0.005
0.0
0.0
0.000
0.0000
0.00
Drmension
Dicontsion
Dimonetoc
Dimendoo
Dimenson
Figure 28: Eigenspace overlap different models of LeNet5-BN.
Comparing Fig, 27b and Fig. 27c, we can see that the Kronecker factorization still gives reasonable approximation for the eigenvector correspondence matrix with R[ax""", although worse than the cases without BN (Fig.
Fig. 29 compare the eigenvalues and top eigenspaces of the approximated Hessian and the true Hessian for LeNet5 with BN. The approximation using Kronecker factorization is also worse than the case without BN (Fig. 2). However, the approximation still gives meaningful information the overlap of top eigenspace is still highly nontrivial.
37

---

0.5
A Approximated
12-
Exact
0.4
0.8
0.3
0.6
0.2
Ec1
0.4
fc2
0.1
0
20
20
40
60
80
The 1-th eigenvalue
Dimension Dominating Eigenspace
(a) Top eigenvalues of approximated and exact layer-wise Hessian for fc2.
(b) Top eigenspace overlap between approximated and true layer-wise Hessian.
Figure 29: Comparison between the true and approximated layer-wise Hessians for LeNet5-BN,
G Computing PAC-Bayes Bounds with Hessian Approximation
Given a model parameterized with 0 and an input-label pair (æ, y) R x R,, the classification error of over the input sample æ is İ(0, x) := aarg max fa(x) arg max y). With the underlying data distribution D and training set i.i.d. sampled from D, we define
E(0) (( x)),
e(0)
1 N
Σ
((0 2;)]
(59)
the expected and empirical classification error of 0, respectively. We define the measurable hypothesis space parameters H := TR.". For any probabilistic measure P in H. let e(P) Ee)re(f), é( P) = rê((), and 8( P) (()r Here ((P) (P) serves differentiable convex surrogate of é(P).
Theorem G.1 (Pac-Bayes Bound), (McAllester 1999 Langford & Seeger 2001) For any prior distribution P in H. that is chosen independently from the training set S, and any posterior distribution Q in H. whose choice may inference S, with probability 1 ō,
DRL (Ql()) log
S
DKL (ε(Q)]le(Q)) I
(60)
|
Fix some constant b, c and On E H random initialization, Dziugaite & Roy (017) shows that when setting Q = N(w, diag(s)). P = (6,.. AIp). where w, s E 7l and À cexp (-j/b) for some j E N, and solve the optimization problem
DKL (kL(Q[ + log [S]
min ě(Q) w.ewÀ
(61)
2([S] - 1)
with initialization w = 0, = 8², one can achieved nonvacous PAC- Bayes bound by Eq. (60)
In order to avoid discrete optimization for N, Dziugaite & Roy (2017) uses the BRE term to replace the bound in Eq. (60) The BRE term defined as
BRE(W, 8, 0) = (kl(P)))) + 2log(b {) log *118|
(62)
[
where Q = N(w, diag(s)), P = (Bb, AI)). The optimization goal actually used in the implemen- tation thus
{(Q) + V BRE(W, 8)
min WER" MER
(63)
38

---

Algorithm 1 shows the algorithm for Iterative Hessian (ITER) PAC- BByes Optimization. If we set ?) T. the algorithm will be come Approximate Hessian (APPR) PAC-Bayes Optimization. It is based on Algorithm in DZiugaite & Roy (2017). The initialization of w is different from Dziugaite Roy (2017) because believe what they wrote, abs( ) is typo and log log(abs((w) (w) is what they actually means. It is more reasonable to initialize the variance 8 as w² instead of exp abs(w)).
Algorithm 1 PAC-Bayes bound optimization using layer-wise Hessian eigenbasis
Input:
Wo E RRP
> Network parameters (Initialization)
E R.P.
D Network parameters (SGD solution)
> Training examples
8 E (0
Confidence parameter
be N, CE 1)
> Precision and bound for
TH (0, 1),
Learning rate; No. of iterations Epoch interval for Hessian calculation
E N
Output w
> Optimized network parameters
8
Optimized posterior variances in Hessian eigenbasis D Optimized prior variancce
À
1: procedure HEAAIIVEE-HSSIA-PPAC-BBBAA
2:
Ç 6- log(abs(w)]
> where s(<) = exp 2<
3:
2 g+ -3
where x(e) exp (2)
4:
R(w, 8,
BRE(W,, X;
BRE term
5:
B(w, s, S. w') =[(w') + R(w, 8, 8)
Optimization goal
6:
for 00 do
D Run SGD for iterations
7:
ift mod == O then HESSIANCAL.C(W) end if
8:
9:
10:
Sample E ~~ NO 1)
11:
w' w'ww, w TOSTANDARD & exp(s))
Generate noisy parameter for SNN
12:
w f- w - u, s, [ww'(w)) < S TOHESSIAN (V w C(w')) @ o exp(s))
13:
14:
0 g o +V vRR((w s. (()))
b Gradient descent
15:
end for
16: return w. s()), A(2)
17:
end procedure
In the algorithm, HFSSTANCALC((U)) is the process to calculate Hessian information with respect the posterior
in order produce the Hessian eigenbasis to perform the change of basis. For very small networks, we can calculate Hessian explicitly but it prohibitive for most common networks. However, efficient approximate change of basis can be performed using our approximated layer-wise Hessians. In this case, we would just need to calculate the full eigenspace of E M.M] and that of [[xx="] for each layer. For pth layer, we denote them as U(e) and (p) respectively with eigenvectors columns. We can also
the corresponding eigenvalues by doing pairwise multiplications between cigenvalues of [[ and B[xx"").
After getting the cigenspaces, of basis on vectors with the same dimensionality as the parameter vector (or the posterior mean). ToHESSIAN(u) is the process to put vector in the standard basis to the Hessian eigenbasis. We first break into different layers and let (P) be the vector for the pth layer. We then define Mat the reshape of vector to the shape of the parameter matrix (p)) of that layer. We have the new vector (p) in Hessian basis
perform the change of basis. Note that perform change
(P) vec [T(TT Mat (at (rul))
(64)
The new vector TOHESSIAN(u) is thus the concatenation of all the (p).
39

---

TOSTANDARD(P) is the process to put a vector in the Hessian eigenbasis to the standard basis. It the reverse process TOHESSIAN. We also break into layers and let the vector for the pth layer be v(p), Then, the new vector (P)) is
u vec (TrST) Matl") ((())
(65)
The new vector TOSTANDARD(E) is thus the concatenation of all ((P))
After getting optimized w, s, X, we compute the final bound using Monte Carlo methods same as in Dziugaite & Roy (2017). Note that the prior P is invariant with respect to the change of basis, since its covariance matrix is multiple of identity p. Thus, the KL divergence can be calculate in the Hessian eigenbasis without changing the value of X. In the Iterative Hessianwith approximated output Hessian (lTER.M), we use M to approximate M as in Eq. (56)
We followed the experiment setting proposed by Dziugaite Roy 2017 in general. In all the results we present, we first trained the models from Gaussian random initialization wo to the initial posterior mean estimate w using SGD (Ir=0.01) with batch-size 128 and epoch number 1000.
We then optimize the posterior mean and variance with layer-wise Hessian information using Algo rithm I, where 8 0.025. b 100, and 0.1. We train for 2000 epochs, with learning rate initialized at 0.001 and decays with ratio 0.1 every 400 epochs. For Approximated Hessian algorithm, set 17 1. For lterative Hessian algorithm, schedule as learning rate (multiply by 10 every time the learning rate is multiplied by 0. and the results are similar to those without decay. We also used the same Monte Carlo method as in Dziugaite & Roy
set 10. We also tried 1) with the same decay
[ (2017) to calculate the final PAC-Bayes bound. Except that we used 50000 iterations instead of 150000 iterations because extra iterations do not further tighten the bound significantly. We use sample frequency 100 and 8' = 0.01 as in that paper.
The complete experiment results are listed in Table 9 We follow the same naming convention in Dziugaite & Roy| (2017) except adding T-200 we introduced in Section 4 T-60010. T-600, and T-200 trained standard MNIST with 10 classes, and others trained on MNIST-2 (see Appendix D. in which we combined class 0-4 and class 5-9.
Prev means the previous results in Dziugaite & Roy (201)) APPR means Approximated Hessian, ITER means Iterative Hessian, ITER (D) means Iterative Hessian with decaying 7), ITER. means Iterative Hessian with approximated oulput Hessian. BASE are Base PAC-Bayes optimization as in the previous paper.
In Table 9
We also plotted the final posterior variance. s. Fig. 30 shown below is for T-200 For posterior variance optimized with our algorithms (APPR, ITER, and ITER.M) we can see that direction associated with larger eigenvalue has a smaller variance. This agrees with our presumption that top cigenvectors are aligned with sharper directions and should have smaller variance after optimization. The effect is more significant and consistent for Iterative Hessian, where the PAC-Bayes bound is also tighter.
ITER
ITER.M M
10-1
10-2
10
10
50000
100000
154000
50000
100000
150000
50000
100000
150000
50000
100000
150000
Figure 30: Optimized posterior variance, 8. (fc1:T-2002. trained on MNIST), the horizontal axis is ordered with decreasing eigenvalues.
40

---

Table 9; Full PAC-Bayes bound optimization results
PAC-Bayer Bound
KL Divergence
SNN loss
Test Error
Network
Method
À (prior)
T-600
PREV
0.161
5144
0.028
0.017
BASE
0.154
4612.6
0.03373
-1.3313
0.0153
APPR
0.1432
3980.6
0.03417
-1.6063
0.0153
ITER
0.1198
3766.1
0.02347
-1.2911
0.0153
(TER(D)
0.1199
3751.1
0.02366
-1.2913
0.0153
ITER.M
0.1255
3929.9
0,02494
-1.3213
0.0153
T-600
PREV
0.186
6534
0.028
0.016
BASE
0.1921
6966.6
0.03262
-1.163
0.0148
APPR
0.1658
5176.1
0.03468
-2.0963
0.0148
ITER
0.1456
5086.5
0.02473
-1.7963
0.0148
(EE(D))
0.1443
4956.8
0.02523
-1.7963
0.0148
ITER.M
0.1502
5024.5
0.02767
-1.8363
0.0148
T-1200
PREV 0.179
5977
0.027
0.016
BASE
0.1754
5917.6
0.03295
-1.4463
0.0161
APPR
0.1725
5318.8
0.03701
-1.8313
0.0161
ITER
0.1417
5071
0.02292
-1.4763
0.0161
ITE((D)
0.1413
5021.1
0.02316
-1.4763
0.0161
ITERM
0.1493
5185.4
0.02576
-15366
0.0161
T-300
PREV
0.17
5791
0.027
0.015
BASE
0.1686
5514.9
0.03329
-.1513
0.015
APPR
0.1434
4105.4
0.03296
-1.8063
0.015
ITER
0.1249
3873.2
0.02514
-1.766
0.015
ITER(D)
0.1244
3833.7
0.02526
-1.763
0.015
ITER.M
0.1308
3987.2
0.02721
-15771
0.015
R-600
PREV
1.352
201131
0,112
0.501
BASE
0.6046
1144.8
0.507
-1.8263
0.4925
APPR
0.5653
390.25
0.5066
-2.7713
0.4925
(EE
0.5681
431.62
0.5066
-2.4513
0.4925
ITER.M
0.5616
340.62
0.5065
-2.5263
0.4925
T-200%
BASE
0,4165
21896
0.04706
-1.1513
0.0208
APPR
0.2621
11068
0.0366
-1.4213
0.0208
ITER
0.2145
9821
0.02229
-1.1513
0.0208
ITER(D)
0.2311
9758.5
0.03071
-1.1513
0.0208
ITER.M
0,2728
13406
0.02605
-1.5513
0.0208
T-600
BASE 0,2879
12674
0.03854
-1.1513
0.018
APPR
0.2424
9095.8
0.04159
-1.013
0.018
ITER
0.2132
8697.9
0.02947
-130663
0.018
ITER.M
0.2227
8870.9
0.03294
-1.4613
0.018
T-600%
BASE 0.3472
17212
0.03884
1.1513
0.0186
APPR
0.2896
11618
0.04723
2.0563
0.0186
ITER
0.2431
10568
0.03057
-1.5713
0.0186
41

---

H Proofs and experiment for the low rank structure of the output Hessian
H.1
Preliminaries
H.1.1 Notations
We use [n]] to denote the set f1, f1,...,,.) and ||M]] to denote the spectral norm of a matrix M. We use (A, B) to denote the Frobenius inner product of two matrices A and B, namely (A, B) - Eig A., Denote tr(M) the trace of matrix M and denote 1. the all-one vector of dimension e (the subseript may be omitted when it's clear from the context). Furthermore, for notation simplicity, will say "with probability over (w event E is true" to denote
lim lim W W
Pr
[E] = 1.
(66)
H.1.2 Problem Setting
Consider two layer fully connected ReLU neural network with input dimension d. hidden layer dimension n and output dimension c. The network is trained with cross-entropy objective L. Let o denote the element-wise ReLU activation function which acts as (z) = x 1=20. Let W E Rn and (2) E RR** denote the weight matrices of the first and second layer respectively.
Let the neural network have standard normal input x ~ NN Ia). Denoting the output of the first and second layer as and respectively, we have
a (W()x) and z W y. Let p softmax(z) denote the softmax output of the network. Let = diag(p) - pp" From the previous analysis of Hessian, we have the output Hessian of the second layer can be written as M == (2) AW()P))) where D := diag(1y20) is the diagonal random matrix representing the activations of ReLU function after the first layer.
In this problem, we look into the state of random Gaussian initialization, in which entries of both matrices are i.i.d. sampled from standard normal distribution, and then re-scaled such that each row of (1)) and (22) has norm 1. When taking n and d to infinity, with the concentration of norm in high-dimensional Gaussian random variables, in this problem we assume that entries of W(1) are iid sampled from zero-mean distribution with variance 1/d, and entries of (2)) are iid sampled from zero-mean distribution with variance 1/n. This initialization is standard in training neural networks.
Since our formula for the top eigenspace is going to depend on W(22, throughout the section when condition on the value of Wil)) and W72), The expectation is only taken over the input x ~ (0, Ia) (due to concentration taking expectation on x is similar to having many samples from the input distribution). In this case, the output Hessian defined as:
take expectation
M A DDW (2)T (2)T ()
(67)
H.2 Main Theorem and Proof Sketch
In this section, we will provide formal statement of our main theorem and its proof sketch. First, we state our main theorem:
Theorem H.1. For all > 0,
(
Ac(M
lim lim 1=+00 d W W.w.
Pr
<E =1
(68)
Besides, for all e > O, if we define S₁ as the top c 1 eigenspace of M, and S2 as R R(W) where R(W) is the row space of W. then
1}
lim lim n-000 d- W.
Pr
Overlap (2, > e =1
(69)
Proof Theorem H.Il
First of all, let us repeat the expression for the output Hessian M:
M A (2) ( D).
(70)
42

---

In the proof, we will first analyze the properties of D, W A separately:
Firstly, D is diagonal matrix with 0/1 entries, and the following lemma shows that its entries independent when the input dimension tends to infinity.
Lemma 1. When d + co, with probability / over W the entries of D are independent.
Secondly, since each entry of ( is sampled i.i.d. from spherical Gaussian distribution, this matrix enjoys some very nice properties when the network width goes to infinity, We have the following lemma for (2).
Lemma 2. (informal) When is large enough, each row of ( has norm very close to 1, and these rows are nearly orthogonal to each other. Besides, the entries (and the average of the entries) cannot be too large.
These properties (along with other useful properties) will be formally stated and proved in Ap- pendix H.3.2
As for matrix A, it's very hard to compute its expectation explicitly because the generation of A involves softmax, but we are able to prove some useful properties of [[A] as shown in the following lemma.
Lemma 3. À 1om [[A] exist and is rank- (0-1] 1) with probability I over W
This lemma is true because A is positive semi-definite (PSD) and is almost always of rank (c - 1). Besides, its null space always contains the all-one vector 1..
Having these properties of these three matrices, we then look at the expression of M again to see how to compute this expectation. This expectation is not easy to compute as we condition on (1) and (2),, With this conditioning, D and A are correlated and hard to decompose. This is when need the most important observation in proof: When the input dimension and the network width tend to infinity, and D can be considered independent when computing M. These two matrices are actually not independent even when take the limit: For example, if D happens to be 1 whenever the first row of (2) is positive, then the first logit output is going be much larger than the rest which significantly skews the distribution of A. However, since the computation of M only contains finite-derree polynomials of A and D, we only need weaker form of independence, i.e,, the distribution of A is approximately invariant condition on finite entries of D, as shown in the following lemma:
Lemma 4. Let Dx denote the distribution of X, and (D1, D2) denote the total variation distance between D1 and D2. Then Vi,j E [ni], Ve >
lim lim n=000 d-000 Wiw.N N... PAIDA.PADD..D.... > <] 0.
Pr
(71)
The intuition behind this is that A is uniquely determined by the output of the second layer z, where z = ( is c-dimensional and y is an n--imensiinal vector. Since D = diag(1y2o). fixing finite number of entries in D is equivalent fixing the signs of finite entries of y. When n is large enough compare to c, only constraining finite entries of y shouldn't 't change the distribution of z by much. The formal proof of this theorem is given in in/ppendix H.3.4
Having Lemma 4 we can equivalently consider A and D independent matrices. To formalize this, we need the following definition:
Definition H.I. Let D' be an independent copy of D and also independent of A. Define M & (2)T (').
In other words, M is the matrix which has the same expression as M except that we assume D is independent of A in M*. Then we know that M and M' are essentially the same. Since D is a diagonal matrix with 0/1 entries, multiplying D at both sides of a matrix is equivalent to independently zero out each row and corresponding column with probability 5. Thus, the probability of each diagonal entry to be kept Y while for off-diagonal ones it's 1. Formally, we have
1 EE M
W
(2 A A2 + diag(E W (2)7 (w
(72)
4
43

---

We have two terms on the right hand side of this equation: T, E [W Aw and T2 - (iag(E[(W(a) A (. We make two observations of these two terms. On the hand, they have the same trace. This is because the diagonal entries of these two matrices are exactly the same. On the other hand, Ti is low rank matrix but T2 is approximately full rank. For Ti, since the expectation is taken over only, know that T W2)TE [A] W Note that E[A] is a rank (c 1) PSD matrix, so T2 is also PSD and has rank at most (c - 1). T2 a diagonal matrix, and each diagonal entry equals quadratic form w. un("TT) [[A]ws where w is the i-th column of ( (i E [n]). This
term is always positive unless w. lies in the span of c. which happens with probability 0. Actually,
due to the random nature of w. " the diagonal terms in T3 do differ too much from one another.
To summarize, Ti and T2 are both PSD matrices with the same trace, while Tj is low-rank but Tg is approximately full rank. This intuitively indicates that the positive eigenvalues of T2 is significantly large than those of T2, making the positive eigenvalues of T₃ the dominating eigenvalues of M and those of T2 the thin but long tail of M eigenvalue spectrum.
Now we have know that Ti is almost the only contributing term to top eigenvalues and eigenspaces of M so we only need analyze these for Ti ( E [A] (2)2, Since the rows of (2) are close to rank and almost mutually orthogonal, the matrix Ti will roughly keeps all the eigenvalues of E[A], and the top (c 1) eigenspace should roughly be the "'ottred version of R (le). i.e, (( ( (2). lc).
Despite the arguments above, we have some technical difficulties, the biggest of which is that the dimensions of M and M will become infinite when n goes to infinity. To tackle this problem, we introduce indirect and more complicated way to do the proof. We first project these matrices onto the row span of (2) and show that this projection roughly keeps all the information of these matrices. Formally, we have the following lemma:
Lemma 5. With probability over (
W WWWWTI
lim
=1
M
After that, we do the analysis in this finite-dimensional span and finish the proof of our main theorem.
H.3
Detailed Proof
H.3.1
Proof Lemma
We first restate Lemma restateLLemaa here:
Lemma 1, When d + 0o, with probability over W(), the entries D
D are independent.
The off-diagonal entries of D are always 0 and independent of anything. For the diagonal entries, each diagonal entry is decided by corresponding entry y. Therefore, we only need to prove that the entries of are independent, and we have the following lemma:
Proof of Lemma Remember that D =
Lemma 6. When d + co, with probability over WL), the entries of y are independent.
Proof of Lemma 6. We will prove this lemma using the multivariate Lindeberg- Feller CLT. For each i [n], let w, R& denote the i-th column vector of (1)), Let u, w,x,, then have
w,x; = ((1
(73)
Note that x;'s are i.i.d standard Gaussian. It has the moments:
[[ O,
B[(x; [[xx])=]=
E[(( E[x.))
(74)
follows that
Varlu:] = Var[wxx.] w.
(75)
44

---

Let V = Ei=1 Varlu,)
V famm w
(76)
As d + 20, from Lemma 9 (we replace probability, therefore lim In.
W with d and W we have w(tr In in
We now verify the Lindeberg condition of independent random vectors {ui...., Uz First observe that the fourth moments the u, are sufficiently small.
27
u
lim
E
Σ
(1) 58
X3)
d
d
lim
E
n²
max
W (2)
x3
(77)
jEIn|
i=1
1
d
lim n²
max ie (E())EE())
(x
js
Σ
(x+
E|x,))
Since [x(x; E = 3 and [dxje(d].je[]] it follows from above that
(2) < 2d with probability 1 from Lemma 11
lim
|
S n² lim
2d
Σ
3 = 48n² lim d' id = 48n² lim d = 0. (78) d=000
i=1
For any
O, since |u:
in the domain of integration,
d
lim
u
1 [|/,,/|
e]]
lim d00
E
uu
|u
[1]
(79)
n
lim
E
|
O.
d=000
As the Lindeberg Condition is satisfied, with limn V
Ie in probability we have
lim == lim d= d000
u
(80)
Thus, since u converges to N (0, In) in distribution with probability over (1),, entries are independent. Since y = a (u) and ReLu is an entry-wise operator, entries of y are independent. D
Since the diagonal entries of D are uniquely determined by the corresponding entries of y, know that when d + with probability 1 over (1),, the entries of D are independent. This finishes the proof of Lemma
H.3.2 Proof of Lemma 2
We first restate Lemma 2 here:
Lemma 2. (informal) When is large enough, each of ( has norm very close to 1, and these rows are nearly orthogonal to each other. Besides, the entries (and the average of the entries) cannot be too large.
This is not a formal lemma and will act as the intuition behind the properties of (2), this section, we will formally state the properties need and prove them. Besides, for simplicity of notations, unless otherwise stated.
use W to denote (2) from
Lemma 7. For all i [c], for all > lim Pr (Σ)= WNy 2 26) = 0.
45

---

Proof of Lemma 7 Since each entry of W is initialized independently from (00 "),, by Central Limit Theorem we have دE=1 Waj (0, ). For any 0, fix By chebyshev's inequality,
lim 1 n-000 ne²
lim P
Σ W 26
0.
(81)
n-+00
Lemma 8. For all 6 >> 0, 7im Pr (1 1/. c 2 6) =0.
Besides, for all i E [c], lim Pr
Will - - 1 6) =0
Proof of Lemma 8.
Since each entry of W is initialized independently from (0, H.), we know
Ei= Σ nW2, follows xen distribbutionn Using the tail bound provided by Lemma 1 in Laurent & Massart 2000), we know that for large enough
that n [
Pr In [[]] - cnl 2 ne)
2 Pr (Iz W cnl 2 2//n2//4 3/4 + +2 S 2exp(-n''
(82)
In other words,
lim Pr W -cl 2(
lim Pr [n | cnl nc = 0.
(83)
n-000
Similarly, Vi E [c]. n W follows a x distribiiut so for large enough n,
Pr [n w wi - n| 2 ne)
s Pr [n | ml 2 3/3 + /2) < 2 exxp((=p (aap(-n11)) 1/3),,
(84)
which indicates that
lim Pr [ = lim Pr 11000
[n | ml 2 nc = 0.
(85)
Lemma 9. For all c >> 0, lim Pr (WWW Ie 2 = 0.
Besides, for all i.j [c], sim Pr ((WWT),, degl 2 c) Here ö the Kronecker delta
function, i.e. 8155 = 1 [i=]]
Proof of Lemma 9, Since each entry of W is initialized independently from M/0, h,, we know that WWT follows Wishart distribution W. ( n). Using the third tail bound in Theorem I of Zhu 20122, for large enough n, we get
Pr (r[] Iell 2 e) = Pr
1 WWT
1
2 S
Pr
-
1 WWW
1
1e
> 2(c+ 1)m 1)n-m)n------- 2cn 2cn-1/)) 1/2)
1/4
2ecxp(-vii
(86)
Therefore,
Ve > 0, lim Pr Pr Wr( I.] 2 2() 2()=0. =0
(87)
n-
Moreover, for all i, E [c], we have
Pr -
Σ (WWT)), -8hs)² -8 2 =²
1+=1
(|WWT - 1ol/22e²
(88)
Pr
(|WWT loll Vc
Pr
46

---

which implies that for all 8,j [c].
lim Pr (IWWWT)s 8 2 e) c)= =0.
(89)
n-000
Lemma 10. Let Pw be the projection matrix onto the row span of W, then for all e > 0,
lim Pr [ - Pw 11 of 0.
(90)
Proof of Lemma 10, Without loss of generality, we assume that 1. Let Wi(i [c)) be the i-th row of W, and we will do the Gram-Schmidt process for the rows of W. Specifically, the Gram-Sctie Schmidt process is as following: Assume that (W.).., are the already normalized basis, we Wi+s Finally, from the definition of ||
set W
W
W+I, W and W k-1
projection matrix, we know that Pw
W W.
Let c1166244 from Lemma 9 we know that Vi.j E lim, Pr (W,W, -,,1
O. Besides, from Lemma 8 we know that Vi [c], lim Pr (11WWl1 1| 2c) 0. Then we use induction bound the difference between W and W. Specifically, we will show that Vi [c], W|: wil I 8'e.. For notation simplicity, in the following proof we will not repeat the probability argument and assume that Vi, E [c], W; W - S é' and E [|] We will only use these inequalities finite times so applying union bound will give the probability result.
é'
W. sc. If our inductive hypothesis holds for i < k. then for i = k + 1, we have
For i 1, we know that W 1
W and Will 1 S c, [
Vj - (W, W M < (W,, W, + KW, W, W,)l
< < + Will W, w,
(91)
$ é' + (1 + é'8'é
s (2 1)e'.
Therefore,
W will Σ (W W 3 < Σ (23i+1 1)e' (23R+2 1)
(92)
EE|K
jEIM
and
[ - 1/ W. - 1 I - W.l S 23k+26
(93)
Thus,
W. W. < W. - Will + W - W.ll
< W| - 1| W - -Wil
(94)
<8t4166
which finishes the induction and implies that Ve > 0, vi E [c] W; - w.|l < 8'e' Thus,
||
w| F
Σ
W. Wi 2 cc.16%..
(95)
EE(c)
This means that
WTW W Pw
= | W w W
F <2| -wl 1 W W
(96)
2c 16% Se.
47

---

Lemma 11. The largest entry of ( is reasonably small with high probability as n goes to infinity, namely,
lim P
w (2)
2n
0
(97)
[Ele)]E[[
Proof of Lemma 11 For i..dd random variables X1; Xn ~ (0,, 1), by concentration inequality on maximum of Gaussian random variables, for any 0, we have
P maxx, (og((2))) + < 2e-
(98)
For any i, since W are i.i.d sampled from (O, "),, with rescaling of 1/ we may substitute Xj with (=(2) ij
follows that
/2log((cn) +
P
max #E|d]JE(m)
W!
2e -c-
(99)
vn
Taking L = nž, with e as constant, for large n we have v log(2cn)
< nz. Thus for large n.
n + n
max [E(c).E(
W(2) (2)
2n
=
max JEle)JE(n)
W (2)
by
72
(100)
// log 2n) + no
max EEIc).EEn
W
2e
With the same argument, have
P
min
w!! (2)
-2n
2e
(101)
[E(e)]E[N)
Passing n to infinity completes the proof.
H.3.3 Proof Lemma 3
We first restate Lemma 3 here:
Lemma 3. À lim, E[A] exist and is rank-(c 1) with probability I over W(2)
Before proving Lemma 3] we need some knowledge about the distribution of A. Since A is determined by the vector suffice know the distribution of
Lemma 12. 4om z 4 AO,
x-1 Ie) with probability over (2). 2x
Proof of Lemma 12 We will prove this lemma using the multivariate Lindeberg-Feller CLT. For
each E [n]. let w. E Re denote the i-th column vector of (2)) Let v, = wily, - E(y,]). then we have
v
E Σ Wi.
(102)
1=1
1=1
1=1
From Lemma 6 we know y, are i.i.d. rectified half standard normal. It has the moments:
[yyy 2m
-
B (ly,-EJy -(y,)))
6T
10T - 3 < 1. (103)
2m
4n²
follows that
-- w.w!
Var[v:] V Varlway,
(104)
2T
Let V En Varlv:].
n
π
w. w.
w (2) (2)T
(105)
2T
2T
48

---

As n + 0o, from Lemma 9 we have W + Le in probability, therefore limm V 2
I.
We now verify the Lindeberg condition of independent random vectors (VI,..., VnJ. First observe that the fourth moments of the v,'s are sufficiently small.
27
lim
E
V+
lim
E
W
(y: -Bly,))
1=1
(y. (y.-E(y,)) }
max wl?)) jelc
(106)
lim
E
c²2
1=1
4
lim
max iE(n)oEld
( (2)
ja
Σ (y, - B
100
Since E - - < and (m] it follows that
< 2n 1 with probability 1 from Lemma
ji
lim
E [[MIll]]]
S c² lim 2n
Σ
lim 16n in = 166. lim n 0. (107)
For any O, since |/s
e in the domain of integration,
lvvll c2
lim
E
[ivll [1lvll > c]
lim 1-0000
E
l|vill
1 11lv
(108)
lim
E
||v: [1]
=
0.
2 n-0
i=1
As the Lindeberg Condition is satisfied, with lom V
I-1 Ie
we have
2
lim
v.
d N 0,
T-
π
Ie
(109)
2T
By Lemma 7. we have lim w, Ö with probability over (2),, therefore plugging Eg. (109)
into (Eq. (102) we have
lim 4 N (0, π --111- 10
(110)
2T
After that, we can proceed to prove Lemma 3
Proof of Lemma 3 Note that each entry of is a quadratic function of p, and p is a continuous function of z. Therefore, we consider A function of and write A(z) when necessary. From Lemma 12 we know lim, z follows a standard normal distribution A (0, aTT)) with probability 1 over W, where a is some absolute constant. Therefore, Á A lim,, [|] exist and it equals (0 2)] [Z [A (z)]. For notation simplicity, will omit the statement "with probability I over W when there is no confusion.
From the definition of A we know that A - (rag())-pp' where p is the vector obtained by applying softmax to so Di= Pi = 1 and Vi E [c],2: E (0, 1). Therefore, for any vector p satisfying the previous conditions, we have
T
Pi
PiP
Σ (p: - p.) =
(111)
where 1 is the all-one vector. Therefore, we know that A has an eigenvalue O with eigenvector 7 This means that [[A] also has an eigenvalue with eigenvector Ve Thus, [[[A] is at most rank (c 1).
49

---

Then we analyze the other (c - 1) eigenvalues of A. Since A = QQ" where Q diag( vP) Ip"). we know that A is always positive semi-definite (PSD) matrix, which indicates that [[A] must also be PSD. Assume the eigenvalues of À are Y 2 X2 2 Ac-1 de 0. Therefore, by definition, we have
x=-1
min v" Au Ez-A(( v =1
min |ES.||e||=1
v Ar
(112)
where S (RRRR is the orthogonal subspace of the span of 1. v E S implies that v 1 1, i.e.,
E 0 = 0.
Direct computation gives us
v" Av Σ
vip; Σ ViPi
(113)
i=1
Define two vectors a, E R2 as Vi E [c], a, U. b Vis then ||6|] Ei= P: = and
x² Au |a|| (a. b)² ||a|@ ||b|| (a, b)?.
(114)
Therefore,
A 2 |||| |bb| sin" 0(a b). where ((a, b) is the angle between and b, i.e. 0(a, b) - arccos then
(115)
(a.b) la T
Define po min Pi.
||a||
po ||/|| Po- vipo
(116)
1=1
i=1
Since ||b]| = 1, we have
l|a
(a, b) bll ||a |
sin? 8(a, b)
(117)
Besides,
laa - (a, b) b11² =
Σ
1i DiyPi
Σ
DiPj
Vpi
i=1
j=1
2
Pi
1t:
(118)
Σ
2
U
PjPs
1=1
j=1
Define arg maxie(]] ii and A arg minie[e] Ui, then
(U,
vz)
V
Σ PyPs
+ - UyPj
> 2
(119)
j=l
From ||l|| = 1 we know that maxsele) |D: 2 Ve Besides, since De - Vi
0, we have > 0> vr.
Therefore, U. - vz > maxxeld
|dy| 2 Ve' As a result,
||a - (a, b) b112 2 po
(U. - v2)²
Po
(120)
2
2c
Moreover,
c
||a 2 Σ
Di <
pi = 1.
(121)
i=1
50

---

Thus,
sin 0(a, b) 2
2c.
Po
(122)
2c'
which means that
Av 2 Pa
1 Pu
P6
(123)
2c
2c
Now we analyze the distribution of Po. Since z follows spherical Gaussian distribution (0, ade). we know that the entries of are totally independent. Besides, for each entry z;(i E [c]). we have |z.l < a with probability 3. where 3 0.68 an absolute constant. Therefore, with probability 8".
forall entries 2i E [c]), we have | < a. In this case,
exp(minse(c) zi)
exp(-cx)
Po Ei=1 exp(zi) c cexp(a)
(124)
In other cases, know that pa >> 0. Thus,
exp(-a) coxp(n)
A--1 E2-A(0,,13)
min |ES.|||||
v Au
(125)
2c
The right hand side is independent of Therefore, A--1 0, which means that A has exactly (c - 1) positive eigenvalues and 0 eigenvalue, and the eigenvalue gap between the smallest positive eigenvalue and ( is independent of n.
H.3.4 Proof of Lemma 4
We first restate Lemma 4 here:
Lemma 4. Dx denote the distribution o, and (D1, D2) denote the total variation distance between Dj and D2. Then Vi,j E [m]. Ve > 0
lim lim n d00 Wi. N..
Pr
PAD.A..AD....D...
0.
(71)
Proof of Lemma The proof of this lemma requires knowledge about the distributions of A condition on two diagonal entries of D. Since A can be uniquely determined by z. is enough for
know the distribution of condition on the entries of D. We use the following lemma to analyze this:
Lemma 13. With probubility 1 over (22) for any i,j E [nr. any (p,g) E {0,
lim ((z) Dn Da p, D3s 9) S z. 120000
(126)
With {wi,... w and {vi...., vn3 defined above, since different summands contributing to z are independent, and every v, is only affected by its corresponding Da.
Proof of Lemma 13
P(z|Da, = p. Djj = q) = z + P(v,DDx = p) - + P(v, = 9) w. (y; P(y, D p)) + w P(y -)).
(127)
For any i E [n]. with the condition of D2 p. when p DI, p) O: when conditioned with p = 1, D. = p) is of a half standard normal distribution truncated at O. In both cases the conditional distribution of P(y, D p) and hence y; Ply p) has bounded mean and variance. For any w, by Lemma we have
2 ij
||w:|| ||wiil
s/c
max (( ( [E(c).E(M))
<
4en
with probability over (2)2,
(128)
Since lim V4cm-5 0, goes to infinity we have
wi(y, - Ply p)) 4 0 with probability I over (
(129)
51

---

Therefore
= p, Djs = =
w. Piy,DD = p)) wy(y, pjy =4) S (130)
From Lemma 13 we conclude that
lim lim 11 11 to Wo, (maa) I.a), W0..
Pr
=
(131)
Since A can be uniquely determined by z, we have
(D; DD... 2 (V(DA,AAD -.I....DDD.. -D,,--)
(132)
Therefore,
Pr R--00 d00 W [ Win.
lim lim
PAp...AA > =0
(133)
This finishes the proof of Lemma 4
H.3.5 Proof of Lemma 5
We first restate Lemma 5 here:
Lemma 5,
With probability over (
WMWW T
lim 70000
M 2
Proof of Lemma 5, To prove the equivalence between other terms, including terms containing M', bridges. To prove the equivalence between M and M' we need the following lemma which explains the reason why we only need the weaker sense of independence (Lemma 4) instead of the total independence between A and D.
W MWT and | f we need some
Lemma 14. Let p(A, D) be homogeneous polynomial of A and D and is degree in A and degree 2 in D, and let the coefficients p are upper bounded E -norm by an absolute constant. Also let A' be independent copy of A. Then
lim lim +-+00 to0 WO L. W.O. [E EIp(A, > >d=0.
Pr
(134)
Proof of Lemma 14,
Assume that p(A,D) = Em (A0(0) (0 (u( P (e()) then from
linearity of expectation we know
(j( = Σ & CE D (1 Dom
(135)
Since the entries of D can only be or have
E (1) D (()) D (()) E(A D Du())())) = D De()).(() = 1].
(136)
Assume that the upper bound of the sum of lcsls is i.., Em |c| 2 Set é' and from Lemma 4 we know that
lim lim 10-000 +oo W W.
Pr
[TV -D,
0. (137)
In other words, with probability we have TV DAID..:D,,.. =1), S c. Besides, since Vi, [c], # j.P., Ps.P: + Py E 10, each entry of A (cither p4 ps or -P:P,) must in (- 1). Therefore, when and goes to infinity, with probability 1 over W(1) and (2)) we have
| [B(AA D Pa().(t)) D Do())e(s) =
E A (00)
(138)
<TV PAI...D,.
5 2
52

---

Thus,
-BppA',
D)II &
Σ
cal |B[A (1),)(1) D Du(t),u()) = D Do(),.0) = 1]
2=1
d'
c
S a
<e.
2
20x
(139)
This finishes the proof of Lemma 14
After this, using Lemma we have the following lemmas: Lemma 15. with probability over W,
| 1.
lim 70000
M
Proof Lemma 15
Let (D', A') be an independent copy of (D, then
[M F ||E|D""WD]]| F
E [(DW' AW', A'WDD')
(140)
= Ir (DW W' A'WW'")
= E " AWDD'W' A')]
Expressing term inside expectation as a polynomial entries of A, D, A' and D', we get
(WD'DW" AWDD'W"
A²),,
(WD'DW*A), A),, (WDD'WTA))
(141)
W W.
ΣΣ W W W.
k=
Σ 1
W.WWW DD.
Now bound the norm of the coefficient of this polynomial follows (note that absolute value of each entry of is bounded by 1):
Σ E
W W. W j.t W
W [aal
Wa W..
k, 11:1
n Σ
n
Wal Wall-Waal
Σ
Waal W
1==1
j,s=
(142)
W2 L.
W
c
n
W2
W2
8.6
W + |
W +
[< )2 = è² |
53

---

When n + co, we know that W = O(c) with probability 1 over W. so the coefficient of this polynomial norm bounded. We know from Lemma that the distribution of A is invariant condition on two entries of D. Furthermore, since A' and D' are independent copies of A and D, we know that the distribution of (A, A') invariant conditioning on two entries of and two entries of D'. Each term in this polynomial is 4-th order term containing two entries from D and two from D'. This combined with Lemma 14 gives us
J M 1F
lim 70000
1.
(143)
Lemma 16. For all i E [c], lom ((W
MW NJs
(W (WM))))..) Thus,
WMWWT
lim
=1
WM-WW
Proof of Lemma 16, This proof is very similar to that of Lemma 15 First, we focus on a single entry of the matrix W MW" and express it as a polynomial of entries of D:
(WMW ).s = B[(WDWT AWDW )isl
= E
WDWT ( OR.S
E
W W DuA
Σ
W W SriD
(144)
E=1
E
A W W. W.W.W.WWW...... W., W..DD.
1,,=1,==1
Then bound the l norm of the coefficients of this polynomial follows:
A & s, Wa. W
k,a=1 1,==1
Wil Waal
WW,l [W j.
k,s=11t=1
n Σ
W W. (Wall)
W Wi
(145)
W
W2
W W
2
2
||| + [ (p))
c W
|
5 (2c W | 4e2 |
Similar to Lemma 15 this coefficient is l norm bounded. Therefore, using Lemma 14, we have with probability over W. for all i,j [c), (WW (W M' )i.)) 0, which indicates that
WMWT
lim
WM-WT =1
Lemma 17.
WM-WT.
lim
M
54

---

Proof of Lemma 17, The proof of this lemma will be divided into two parts. In the first part. we will estimate the Frobenius norm of M and in the second part we do the same thing for W M W
Part I: We know from the definition of M that
1 [E] [E[WT AW) + diag(E(W"AA))
M
(146)
4
Define A E[A], then
S[WTAW] --::WTTAAAY WTAW.
(147)
From Lemma
Ve' > O, with probability we have WWW S é'. Besides, from Kleinman
& Athans 1968) we know that for positive semi-definite matrices A and B we have Xunin (A)tr( (B) 2
Ir(B)) S Amaa (A)tr(B) so
A 2
WTWW
1T(TT AWW - I(A)
WwWWW'AWW'A A) - I(AA)
([]WWT 4ell + ('W À) (r(A)
(148)
([[WWT- -Lll + 1) (WW AA) -
(/WW" -oll - + +1) tr (AA) - tr(AA)
$ WWT - -all A F
+ WWT - Tell 11
For any 6 > 0, set = min x gives us with probability
W Áw
A
lim R=+00
(149)
A
F
i...
W A
W
lim
(150)
+-+000
A
Besides, if we denote the i-th column of W by wi then
[diag([((TTW]]))
Σ ! Aw:)²
[|w: 1|2 |A
(151)
1=1
|m /1
Since Ein² |w>ll'] c² + 2c, by the additive form of Chernoff bound we get
Pr wol||* c² + 3c
Σ n² |wall
Pr
(62 + +2) 2 c
Se
2nc
(152)
n
Therefore, when + with probability we have
(Ed AWD)]
<
|
A
Σ
11
S
A
² 2 3c
(153)
55

---

Thus, with probability 1,
diag (E W AW)
2
lim
F
=
0,
(154)
WTTW
F
i.c.,
A
16
lim n-
(155)
||**| 1I
Part 2: Plug equation Eq. (146) into WM W and we get
1 WMW- = WE AWW"] + diag W AW)
(156)
Similar to Part 1, when n + oo, with probability L, we have
E WWW.WAwwE F =1
lim +20000
(157)
A
Besides, when n - co, with probability we have
Wdi AW])W= S
[|
A
2
w
3c
W ||
(158)
n
As result, with probability
W Wdiag W AW])
lim
-0
(159)
WWTAWW!
i.e.
2
A
16
lim n=000
1.
(160)
W M W F
Combining the results of Part and Part proves this lemma.
Combining Lemma 15 Lemma 16 and Lemma 17 directly finishes the proof of Lemma
H.3.6 Proof of Theorem Theorem H.1
Proof of Theorem Theorem H.1. Now proceed to the proof of our main theorem. In this proof, we will use the bounds for [ which are formalized into the lemma below:
Lemma 18. With probability 1, lim [ p is both lower bounded and upper bounded by constants that are independent n.
Proof of Lemma 18.
From
Lemma 15 we know that lim
1, so we only Since M TT AW] + diag(E\W"T and
[
|M
need bound lom M
[E[WT diag(EIE (W ) N O, we have
| F N B[ AWi],, = w AA
(161)
F
From equation
Eq. (125) know that the first (c- 1) eigenvalues of Á is bounded by some constant
which is independent of n. Then we analyze the eigenvalues of WT ÁW: From Lemma 9 and set 2' we know that the smallest singular value of W is lower bounded by } Theretore, for any unit vector v in the row span W, we have
30 exppiar
of W ÁW (We) A(w)) 2 [el 2 7
(162)
56

---

Thus.
W Aw
2 1. which is some constant that independent n.
F
Besides, since D is a dignonal matrix with 0/1 entries, and the absolute value of each entry of A is bounded by 1, we have
[ F E EDW' AWD|
< E W[ AW = W |AIl F S Sc|
(163)
From Lemma we know that with probability 1, W 1 S 2c, therefore, || is upper bounded by 2c which independent of
Now we are ready to prove our main theorem.
From Lemma 5 we have
W WWMW
lim
(164)
|M
Then we consider W WWW W
Note that
W w
(165)
tr wWW' wWW.WWWW WMW
From Lemma 9 we know that for all > O, lim,,+ Pr Pr( WWT Ia]l 2 = O. For notation simplicity, this proof we will omit the limit and probability arguments which can be dealt with using union bound. Therefore, we will directly state WWW" Ia s é'. From Kleinman Athans (9668) we know that for positive semi cemi-definite matrices and B we have Amia (iin(AAtre (A)tr(B) S tr(AB) S maxx A))
[tWW WWW WW" WMW un WWWWW'w'WWM. W |
WW WW WWWWWMW. ww (WMW'WMW) M
< maxx1 Amiu Amin (WW") Amax Ww:::////////////w/// www
WWT Lell U(W wwW.""""""W""))>>> etr WWW.WWWWWWW WWw. SAMAI
(166)
Similarly.
WWT WMW'WWWW L MW W (WMM
<WWMWW WwWWWWW
(997
WWT I. WWWWWWWWWWWW WMW") & (
Therefore,
W WMW T w WIE
1/1
-trWWTT W WWWT WwW.WWWWM. -WW" w W
<tr(WWT WMW'WW'WMWW" www www. WMWT)
+ +|tr(W WWWWWWW'WW'WWW'' "IWWMMI"""I"WWWWMWWVV))
(168)
<6((1 + (W) WWW
<(2e' + (rW WWW) (2'' (+))) WMW
For all <> select é < minf V 4. we have
2
W WWW w |WW < | <
(169)
In other words,
WWW =
lim n-000
(170)
WMWTI 2
Hence get
WWWMMWTM 1.
lim
(171)
|
57

---

Next, consider the orthogonal projection matrix Pw that projects vectors in R" into the subspace spanned by all rows of W. We will consider the matrix Pw M MPw. Define 8 - W Pw. then from Lemma 10 we we get || I c'. Therefore,
W wwwww P M wwwllle | 1 |P + 16MPW IF + |6
- || (2|1Pw]lr |01| + [101l²²
(172)
< || (2 4626' + (6)2).
For all 0, we choose e' < min LE
T6 } and have
W www.wwg wl, - PP M Mwlll <C,
(173)
M
which means that
WWWTWMW'N F WwwMPwlll
WWW W
P
M
Pw 11
lim
lim 1-000
0.
10000
WWw.MWWWW.
(174)
Thus,
PW PwMPWwld
|P MPw F
lim
lim
1.
(175)
MM|
[WTW W MW
Note that |M i| MP || fore,
MW
PW
There-
Pw W
+ PP M MPMM/w/// +
MW lim
M F
|lw M PM 1/2
lim
0.
M 1F
|M
(176)
In other words,
Pw M Pwll
PW M MP |
Pw M Pwll M 1p
lim
lim
lim
0.
(177)
M I
From Lemma 18
we know that oom || |r (this hasn't proved be exist, so we perhaps need
to say "for large enough n²") is lower bounded by some constant that is independent of so
lim 7-000
PwMPPllF
lim |PWll lim PWWMPWII 11 X-000
(178)
Note that
M PwMPW + MW + PWW + PWMP
(179)
Thus,
lim
M - Pw = =0
(180)
For any e O, set m 2c }. from Lemma 10 know that with probability
IIw WT WTwll < 6. Therefore,
Pw MPw W WWW.WWW//
Pw W w'] M + 2 Pw W w
| PP
(181)
<8² 2c² 28 26
<
In other words,
lim
Pw WWW W =0
(182)
n n0000
Now we conclude that
lim - WMMWWWWWWWW/-C +00
(183)
58

---

From Lemma 16 we know that VŁ,j E [c]. (W ).. -WWM- ).)) = 0, i.e.
lim
WMWTT
WM
(184)
Since
W wwW
W WMW WM F. (185)
from Lemma 8 which bounds the Frobenius norm of W know that
lim
W WWW W - W WM' 0.
(186)
Thus,
lim 100
M W WWW'WW' WTr =0
(187)
Note that M = [E[W7 AW + diag AW])), so
4W WM WMWWW wwww"Aww'' W w W (E W AW])w W. (188)
We will first analyze the second term on the RHS of equation Eq. (188) Ve > 0, set e' = Sa.
and from Lemma 9 we know that WWWT - [e] < with probability
1, which means that cl with probability 1. Set c = c, we know that WWT 2c with probability
WWT
Note that
W (F AW])"AW)WWWT F - W w
diagg(E W AWIII
= WW
(E W AWD!I
(189)
s 4c² [iag(
Combine this with equation
Eq. (154) and we have
W T W di
AW] W
lim
O.
(190)
20000
WTAW
F
From equation Eq. (162)
we know that
W ÅW
with probability 1, so
lim 22 +00
W
- wV'WwWANN Aww" =0
(191)
Similarly, define 8 A WWT Le then
w WWW WWW.AwW'W
W W AW
W &AoW
+
W Ao
(192)
W
|
À
+
W 1r
||
A
Set é' < minf
8.4 / 8ca }, then from Lemma we know that ||01| < d' with probability 1, and from
Lemma 8 we have [W]] 2c with probability 1. We also have Ä
I since each entry of A
bounded by 1 in absolute value. Therefore,
W WWW''W'W'' W W'AW
. c 2cd c
E + = e, (193)
2
which means that
lim n=000
W T Wwww'AWw'''
Aw
= 0.
(194)
From equations Eq. (187)
Eq. (191)
and Eq. (194)
we get
lim
M
W
ÁW
=0
(195)
200
4
59

---

Besides, from equation Eq. (95) in Lemma 10 we know that for any c' > 0,
W W
W.
W; <6
(196)
EE
where Wi is the orthogonal version of W, i... we run the Gram-Schmidt process for the rows of W.
Define f W W, for any E > 0, set = min s2
V
2c
J.
we have with probability 1,
W ÁW W AW S 2 2|
A
W +
|181
A
(197)
c c(c)2 <c.
Therefore,
lim
W AW
w' AW
=0,
(198)
which implics
lim 20000
M
W Aw
=
(199)
F
From Lemma we know that with probability 1, Á is rank (c 1). Since A 1 is always true, the top (c 1) eigenspace of Á [1)(1). Note that the rows W are of unit norm and orthogonal each other, conclude that w" AW is of rank (c 1) and the corresponding eigenspace is
R W W 11. Moreover, the minimum positive eigenvalue of W AW is lower bounded by y.
As for the top c eigenvectors of M, define 8 - M
w ÁW, then M W ÁW +
Define S1 as the top cigenspaces for M, and Sg to be the top 1 eigenspaces for W AW. Then from Davis Kahan Theorem we know that
|
||sin 1(S1, FF S
(200)
2 W AW
Here OSS1, S2) is a (c 1) (c- 1) diagonal matrix whose i-th diagonal entry is the i-th canonical
angle between S1 and S2. Since lim | F= and with probability Ac-1 W AW) 2 ?
which is independent of n, we have with probability 1,
lim n00
||sin ((S), S2ll O,
(201)
which indicates that the top c - eigenspaces for M and W ÁW are the same when n - oo.
Notice that the top c- eigenspaces of W AW are W 1 \ ( 1), so M will also have the
same top c eigenspaces. Besides, from equation Eq. (95) we know that lim W W
0, so R [.W. 1 \ (W 1} are the same as R(W theorem.
W 1}. This completes the proof of this
H.4 Experiment Results
Table 10: Overlap of ((S(1)) (S(K). 1} and the top e different layers at minima.
dimension eigenspace of E M (k)
of
Dataset Network
MNIST
MNIST-R
CIFAR10
CIFAR10-R
F-15003
LeNet5
F-1500
LeNet5
F-1500
LeNet5
F-1500
LeNet5
fcl
0.602
0.890
0.235
0.518
0.880
0.951
0.903
0.213
fc2
0.967
0.931
0.801
0.912
0.943
0.972
0.931
0.701
fc3
0.982
0.999
0.998
0.999
0.993
0.999
0.996
0.999
Note that the overlap can be low for random-labe datasets which do not have clear eigengap (as in Fig. 4). Understanding how the data could change the behavior of the Hessian is an interesting open problem. Other papers have given alternative explanations which are not directly comparable to ours, however ours is the only one that gives closed-fore formula for top eigenspace. In In[Appendix F.T will discuss the other explanations in more details.
60


Reviews:
[{"review_id": "yQDZYq2KPo8", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### 1. Content summary\nThe authors study the layer-wise Hessian of several well-known vision networks on several vision classification datasets. They, to me be very surprisingly, observe that the subspaces formed by top eigenvectors of those layer-wise Hessians have non-trivial overlap between independently initialized and trained models. They explain this by proposing the usage of Kronecker factorization, whose validity they experimentally show. They trace the surprising top eigenspace overlap to the structure of the data and show that it weakens with the application of batch norm. They also show that this overlap can tighten PAC-Bayes bounds.\n\n### 2. Strong points\n* The paper is well written and motivated. I went from the intro to the conclusion in one go and felt I didn't have to jump around to fish out definitions etc.\n* The question is well motivated and interesting.\n* The proposed simplifications (Kronecker factorization) are validated on real networks.\n* The networks and data used are non-trivial and close to practice, which makes this paper better than many others in bridging the gap between theory and experiments.\n* The story fits nicely with existing observations about the effect of batch norm on class-specific structure of the Hessian.\n* The obvious application to PAC-Bayes is nice!\n\n### 3. Weaker points\n* The paper would be strengthened by a broader suite of experiments: add ResNets, WideResNets, let's see what skip connections do, try CIFAR-100 to see the effect of >10 classes. I am not suggesting you work on ImageNet, but more experiments of even the easier variety would make the paper even more convincing. Perhaps FASHION MNIST and SVHN might be quick enough to run?\n* I am missing a summary table or plot showing how well the approximations proposed match the reality for different architectures and datasets. This is a small detail, but I kept wondering how valid many things were. I was convinced by your detailed figures (e.g. Figure 3) and I am sure a figure like  that for each combination of net x architecture would clutter the paper, but a summary plot showing the match between the assumptions and reality would be very nice and convincing.\n\n### 4. Points to clarify\n* In Figure 4 the scales for the E[M] and H_L are different, the difference presumable coming from the scale of E[xx^T]. Is that the case?\n* Am I understanding it correctly that the top eigenvector overlap is driven mainly by the structure of the data that through E[xx^T] \\approx E[x]E[x]^T becomes almost rank 1?\n\n### 5. Potential papers to look at\nStiffness: A New Perspective on Generalization in Neural Networks by Stanislav Fort, Paweł Krzysztof Nowak, Stanislaw Jastrzebski, Srini Narayanan (https://arxiv.org/abs/1901.09491) looks at the gradient tensorproduct gradient structure which is a part of the Hessian and its class dependence and training time dependence. It might be relevant.\n\n### 6. Summary\nThis is a good, well-written and motivated paper with a (for me) surprising observation and well-formed explanation for it. More experiments would make it even stronger, but I already like it as it is! ", "extra_scores": {}, "extra_text": {}}, {"review_id": "fpOECxB50R", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 4, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\nThe paper studies the structure of the Hessian matrix of loss functions by approximating Hessians using Kronecker factorizations. Combining the Kronecker factorization with PAC-Bayes, the authors provide a tighter bound on classification error. \n\nPros:\n- The paper contains many experiments.\n- Interesting observations on the models trained with batch normalization. \n- The improvement of classification error bound.\n\nConcerns:\n- I suggest adding the Conclusion section. This section should highlight the contributions and summarize the study, and the paper will certainly benefit from it. \n- The main paper contribution, optimized PAC-Bayes bound, is not emphasized enough; many important details are moved into the Appendix section, while most part of the paper is devoted to eigenspace overlap. While eigenspace overlap looks quite interesting, it does not give many insights. One of the most interesting properties of the Hessian structure is the gap in eigenvalues distribution around the number of classes, but this is prior work.\n- The figure captions are not detailed enough; many figures lack axis labeling. For example, it is not apparent to me what is illustrated in Figure 3. What do the axes depict? Similarly, for Figure 4.\n- The key concern is an incremental contribution comparing to Martens & Grosse (2015). While Martens & Grosse (2015) use the Kronecker factorization for training acceleration, your paper focuses more on eigenspace overlap, which is supported by many rigorous empirical results and illustrations, but not on the possible interpretations and insights.\n\nUPD:\nThank you for your answer and for addressing the points raised in the review. Still, I agree with the concerns raised by AnonReviewer4 on the small scale of the networks used in the experiments (e.g., LeNet, fc). I am still not convinced by the experiments and the excess of hardly understandable plots given to support the main bulk of the paper's claims. Therefore, I leave my rating the same.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "JRcpdQPu0RL", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 2, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "## Summary\nThe paper studies the empirical properties of the Kronecker Factorization based approximation to the Hessian of the loss. The paper makes the following two empirical observations: \n1. Top eigenspace of the layer-wise hessian has a high overlap for neural networks trained with different initializations and hyper-parameter setting. \n2. Top eigenvectors of the layerwise hessian as a matrix is approximately rank 1. \nPaper makes the decoupling conjecture which states that $\\nabla^2_{z^{(p)}}\\ell(z, y)$, which is Hessian wrt loss evaluated at the output of layer p, is independent of $x_{(p)}x_{(p)}^T$, correlation of the input to the layer $p$. Under this assumption, the paper gives an explanation of why the above-stated phenomenon holds. Finally, the paper uses this approximation to optimize the PAC-Bayes bounds for the stochastic classifier where the eigenbasis of the proposed layer-wise Hessian is used for the covariance of the posterior. \n\n## Evaluation\n\nOverall, the paper is not clearly written. It really needs to improve on the notation and needs to explain the ideas more clearly. Starting from Eq(3) in the paper, we are computing a gradient wrt $v$ for the hessian, but it’s really not clear from the notation how the output of $f_\\theta(x)$ is dependent on $v.$ Writing this precisely would ease the burden on the reader. \n\nHaving said that authors do make new empirical observations with regards to the layer-wise hessian approximation but they don’t do a good job in explaining why these observations are relevant in a more general context. I have some questions for the authors to help me evaluate this work.\n\n1. How close is the layer-wise Hessian approximation to the true Hessian? There is no discussion in the paper about the approximation. \n2. A majority of the paper hinges on the decoupling conjecture. How close are the matrices in Eq(6) in terms of the actual norm on matrices? In the paper, we’re only given visual proof for the closeness in terms of the eigenvector correspondence matrices. Honestly,  I had a hard time following how it was computed.\n3. For the PAC-Bayes bound, how the posterior variance is parameterized? It’ll be useful to give a mathematical characterization of it. Is the prior variance wrt standard basis still? If so, how is the KL computed? I am happy to just see the expression because couldn't find it in the Appendix. \n\nMinor comments: \nPlease review the citations added in references. For many of the papers, only Arxiv posts are cited whereas the papers are published at other venues. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "GT0K-M-85rZ", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 4, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "1.\tI want to congratulate the authors on producing an attractive narrative, with a very nice flow to it and a very clear heuristic content.  I believe if I were seeing these ideas for the first time, I would definitely find this presentation friendlier and clearer than if I learned it from some other presentations.\n2.\tThe topic is important. The analysis of the spectrum of the Hessian of the loss function gives us significant insight, and it is rather amazing and beautiful that the spectrum has definite mathematical properties which also have mathematical explanation. The narrative excels in making us feel this as something clear and natural.\n3.\tThis is my first time ever reviewing for conferences of this nature; I am a bit unclear about the level of contribution we are looking for. My understanding is that researchers in this field are writing numerous very short papers each year and a paper doesn’t need to represent more than an incremental bump over previous work.\n4.\tOne reason I bid to review this Mss. is that the topic is not new to me. In the last few years three separate researchers at my institution have written papers on this general topic; and so I feel more conversant with this topic than some other material I might have been assigned. So …\n5.\tI feel that this work may not know about or may not have fully assimilated the contributions by other authors. I can mention these papers:\n\narXiv:1901.10159 \nAn Investigation into Neural Net Optimization via Hessian Eigenvalue Density\n\tAuthors: Behrooz Ghorbani, Shankar Krishnan, Ying Xiao\nAbstract: To understand the dynamics of optimization in deep neural networks, we develop a tool to study the evolution of the entire Hessian spectrum throughout the optimization process. Using this, we study a number of hypotheses concerning smoothness, curvature, and sharpness in the deep learning literature. We then thoroughly analyze a crucial structural feature of the spectra: in non-batch normalized networks, we observe the rapid appearance of large isolated eigenvalues in the spectrum, along with a surprising concentration of the gradient in the corresponding eigenspaces. In batch normalized networks, these two effects are almost absent. We characterize these effects, and explain how they affect optimization speed through both theory and experiments.\n\n\nSo arXiv:1901.10159 clearly calls out the phenomenon of batch norm changing the hessian. In comparison, the paper under review gives a heuristically very clear explanation of why such an effect should be present, i.e. why batch normalization should be effective at making the objective easier to optimize.  However, the paper under review does not fully discuss the above research contribution of arXiv:1901.10159,  and it would be very easy for a reader of this paper to imagine that the phenomenon being presented here originates with this paper.   I should point out that the authors of arXiv:1901.10159 have made company-wide presentations about that work. Since the company is Google, that means that their work is fairly well known.\n\n\narXiv:1811.07062  \nThe Full Spectrum of Deepnet Hessians at Scale: Dynamics with SGD Training and Sample Size\nAuthors: Vardan Papyan\nAbstract: We apply state-of-the-art tools in modern high-dimensional numerical linear algebra to approximate efficiently the spectrum of the Hessian of modern deepnets, with tens of millions of parameters, trained on real data. Our results corroborate previous findings, based on small-scale networks, that the Hessian exhibits \"spiked\" behavior, with several outliers isolated from a continuous bulk. We decompose the Hessian into different components and study the dynamics with training and sample size of each term individually\n\nTo my knowledge arXiv:1811.07062  is the first paper to show the empirical deep learning community that one can compute the eigenvalue density spectra of modern deepnet classifiers eg those of the kind that are in use on large datasets like imagenet and in real applications. The result of having the technology in arXiv:1811.07062  is that the author is able to show that some of the key features of eigenvalue spectra that were seen in small scale situations are also present at full scale. The examples shown in the paper under review, in contrast, are of quite limited scale, and to my understanding do not represent the current state of the art. Consequently, although tools re available to test out the authors’ ideas on realistic problems, we are left wondering what seen in these small examples might generalize. \n\narXiv:1901.08244 Measurements of Three-Level Hierarchical Structure in the Outliers in the Spectrum of Deepnet Hessians\nAuthors: Vardan Papyan\nAbstract: We consider deep classifying neural networks. We expose a structure in the derivative of the logits with respect to the parameters of the model, which is used to explain the existence of outliers in the spectrum of the Hessian. Previous works decomposed the Hessian into two components, attributing the outliers to one of them, the so-called Covariance of gradients. We show this term is not a Covariance but a second moment matrix, i.e., it is influenced by means of gradients. These means possess an additive two-way structure that is the source of the outliers in the spectrum. This structure can be used to approximate the principal subspace of the Hessian using certain \"averaging\" operations, avoiding the need for high-dimensional eigenanalysis. We corroborate this claim across different datasets, architectures and sample sizes\n\nTo my knowledge arXiv:1901.08244 goes over some of the same territory as the present manuscript, but at full scale and in a much more penetrating way. arXiv:1901.08244  shows that not only do the C class means influence the eigenvalues, but actually there are C(C-1) eigenvalues that have structure deriving from means. The examples shown in the paper under review, in contrast, are of quite limited scale, and in contrast don’t seem able to show the full structure which we now know to be present in deepnet spectra across a very wide range of networks and datasets at full scale.  However, the paper under review does not discuss the research contribution of arXiv:1901.08244,  and in consequence a reader could get the misleading impression that this paper’s heuristic explanations of the mean structure are the only knowledge we have about this phenomenon, when we have actually dramatically more information. \n\n\narXiv:2008.11865  Traces of Class/Cross-Class Structure Pervade Deep Learning Spectra\nAuthors: Vardan Papyan\nAbstract: Numerous researchers recently applied empirical spectral analysis to the study of modern deep learning classifiers. We identify and discuss an important formal class/cross-class structure and show how it lies at the origin of the many visually striking features observed in deepnet spectra, some of which were reported in recent articles, others are unveiled here for the first time. These include spectral outliers, \"spikes\", and small but distinct continuous distributions, \"bumps\", often seen beyond the edge of a \"main bulk\". The significance of the cross-class structure is illustrated in three ways: (i) we prove the ratio of outliers to bulk in the spectrum of the Fisher information matrix is predictive of misclassification, in the context of multinomial logistic regression; (ii) we demonstrate how, gradually with depth, a network is able to separate class-distinctive information from class variability, all while orthogonalizing the class-distinctive information; and (iii) we propose a correction to KFAC, a well-known second-order optimization algorithm for training deepnets.\n\nTo my knowledge arXiv:2008.11865  again goes over some of the same territory as the present manuscript, but at full scale and in a much more penetrating way. arXiv:2008.11865  in particular discusses the KFAC approximation of Martens and Grosse and in fact claims, with substantial evidence gleaned from many realistic examples, that KFAC is not a good approximation. The paper under review doesn’t cite arXiv:2008.11865, but does make heuristic claims about the adequacy of the KFAC approximation.  However, the evidence presented is much weaker and the implications much more sketchy than what is discussed at greater length  in arXiv:2008.11865.\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper studies different properties of the top eigenspace of the Hessian of a deep neural network and their overlap. It raised quite a lot of discussion, which finally went in not very constructive way. The reviewers generally agree that the paper has potential, but the actual contribution is limited.

Pros: 
- The idea that top eigenspaces between different models have high overlap is interesting
- The explanation that these structures can be explained by Kronecker-product approximation of the Hessian.

Cons: 
- The connection to PAC-Bayes is unclear and seems artificial.
- Many of the related work is missing
- The models and datasets are too simple, and general conclusions can not be made on such kind of models. Much more testing is needed to verify the claims, including state-of-the art architectures and datasets.
Decision: Reject