Scale-invariant Bayesian Neural Networks with Connectivity Tangent Kernel
Authors: SungYub Kim
Sihwan Park
Kyung-Su Kim
Eunho Yang
Venue: ICLR 2023 notable top 25% (2023)

Abstract: Studying the loss landscapes of neural networks is critical to identifying generalizations and avoiding overconfident predictions. Flatness, which measures the perturbation resilience of pre-trained parameters for loss values, is widely acknowledged as an essential predictor of generalization. While the concept of flatness has been formalized as a PAC-Bayes bound, it has been observed that the generalization bounds can vary arbitrarily depending on the scale of the model parameters. Despite previous attempts to address this issue, generalization bounds remain vulnerable to function-preserving scaling transformations or are limited to impractical network structures. In this paper, we introduce new PAC-Bayes prior and posterior distributions invariant to scaling transformations, achieved through the \textit{decomposition of perturbations into scale and connectivity components}. In this way, this approach expands the range of networks to which the resulting generalization bound can be applied, including those with practical transformations such as weight decay with batch normalization. Moreover, we demonstrate that scale-dependency issues of flatness can adversely affect the uncertainty calibration of Laplace approximation, and we propose a solution using our invariant posterior. Our proposed invariant posterior allows for effective measurement of flatness and calibration with low complexity while remaining invariant to practical parameter transformations, also applying it as a reliable predictor of neural network generalization.

Published as a conference paper at ICLR 2023
SCALE-INVARIANT BAYESIAN NEURAL NETWORKS WITH CONNECTIVITY TANGENT KERNEL
Sung-Yur Yub Kim', Sihwan Park', Kyungsu ; Eunho Yang!
Korea Advanced Institute of Science and Technology (KAIST)', AITRICS Samsung Medical AI Research Center', Sungkyunkwan University School of Medicine", Massachusetts General Hospital and Harvard Medical School sungyub. kim@kaist ac. .r.kr, sihwan park kaist.aca.. .c.kr kskim. doc0 qmail. com, eunhoy@kaist. ac. kr
ABSTRACT
Studying the loss landscapes neural networks is critical to identifying gener- alizations and avoiding overconfident predictions. Flatness, which measures the perturbation resilience of pre-trained parameters for loss values, is widely acknowl- edged as an essential predictor of generalization. While the concept of flatness has been formalized as a PAC-Bayes bound, it has been observed that the generalization bounds can vary arbitrarily depending on the scale of the model parameters. Despite previous attempts to address this issue, generalization bounds remain vulnerable to scaling transformations or are limited to impractical network structures. In this paper, we introduce new PAC-Bayes prior and posterior distribu- tions invariant to scaling transformations, achieved through the decomposition of perturbations into scale and connectivity components. In this way, this approach expands the range of networks to which the resulting generalization bound can be applied, including those with practical transformations such as weight decay with batch normalization. Moreover, we demonstrate that scale-dependency issues flatness can adversely affect the uncertainty calibration of Laplace approximation, and we propose solution using our invariant posterior. Our proposed invariant posterior allows for effective measurement of flatness and calibration with low complexity while remaining invariant to practical parameter transformations, also applying it as a reliable predictor of neural network generalization.
INTRODUCTION
Neural networks (NNs) have succeeded tremendously, but understanding their generalization mech- anism in real-world scenarios remains challenging (Kendall & Gal, 2017; Ovadia et al., 2019). Although is widely recognized that NNs naturally generalize well and avoid overfitting, the un- derlying reasons are not well understood (Neyshabur al., 2015b; Zhang et al., 2017; Arora et al., 2018). Recent studies on the loss landscapes of NNs attempt to address these issues. For example, Hochreiter & Schmidhuber (1995) proposed the flat minima (FM) hypothesis, which states that loss stability for parameter perturbations positively correlates with network generalizability, as empirically demonstrated by Jiang al. (2020).
However, the FM hypothesis still has limitations. According to Dinh et al. (2017), rescaling two successive layers can arbitrarily degrade flatness measure while maintaining the generalizability of NNs. Meanwhile, Li et al. (2018) argued that weight decay (WD) leads to contradiction of the FM hypothesis in practice: Although WD sharpens pre-trained NNs (i.e., decreased loss resilience), it generally improves the generalization. In short, they suggest that transformations on network parameters (e.g., re-scaling layers and WD) may lead to contradictions to the FM hypothesis. A thorough discussion this can be found in Appendix E.
To resolve this contradiction, we investigate AAC-Bayesian prior and posterior distributions to derive new scale-invariant generalization bound, As result, our bound guarantees invariance for general
"Correspondence

---

Published as a conference paper at ICLR 2023
class of prreseiipp scale transformations with broad class of networks. Specifically, our bound is more general than existing works (Tsuzuku et al., 2020; Kwon et al., 2021), both in terms of transformations (c.g., activation wise rescaling (Neyshabur et al., 2015a) and WD with batch normalization (BN; loffe & Szegedy (2015))) that guarantee invariance and in terms of NN architectures. Therefore, our bound ensures FM contradiction for the first time, which should not occur in practical NNs, including ResNet (He et al., 2016) and Transformer (Vaswani et al., 2017).
Our generalization bound is derived from scale invariances of prior and posterior distributions, guaranteeing not only its scale invariance but also the scale invariance of its substance, the Kullback- Leibler (KL) divergence-basse based kerne.. We call this kernel an empirical Connectivity Tangent Kernel (CTK), as modification of empirical Neural Tangent Kernel (Jacot et al., 2018) with the scale- invariance property. Moreover, we define a new sharpness metric as the trace of CTK, named Connectivity Sharpness (CS). We show via empirical studies that CS predicts NN generalization performance better than existing sharpness measures (Liang et al., 2019; Neyshabur et al., 2017).
In Bayesian NN regimes, we connect the contradictions of the FM hypothesis with the issue of amplifying predictive uncertainty. Then, we alleviate this issue by using a Bayesian NN based on the posterior distribution of our PAC-Bayesian analysis. We name this Bayesian NN as Connectivity Laplace (CL), as it can be seen as a variation of Laplace approximation (LA; MacKay (1992)) using different Jacobian. Specifically, we demonstrate the major pitfalls of WD with BN in LA and show how to remedy this issue using CL.,' We summarize our contributions as follows:
Our novel PAC-Bayes generalization bound guarantees invariance for general function- preserving scale transformations with broad class of networks (Sec. 2.2 and 2.3). We empirically verify this bound gives non-vacuour results for ResNet with 1 IM parameters (Sec. 2.4).
Based on our bound, we propose a low-complexity sharpness metric CS (Sec. 2.5), which empiri- cally shows a stronger correlation with generalization error than other metrics (Sec. 4.1).
To prevent overconfident predictions, we show how our scale-invariant Bayesian NN can be used to solve pitfalls of WD with BNs, proving its practicality (Sec. 3 and 4.2).
2 PAC-BAYES BOUND WITH SCALE-INVARIANCE
This section introduces a data-dependent PAC-Bayes generalization bound without -cale-dependency issues. To this end, we introduce our setup in Sec. 2.1, construct the scale-invariant PAC-Bayes prior and posterior in Sec. 2.2, and present the detailed bound in Sec. 2.3. Then, we demonstrate the effectiveness of this bound for ResNet-18 with CIFAR in Sec. 2.4. An efficient proxy of this bound without complex hyper-parament parameter optimization is provided in Sec. 2.5.
2.1 BACKGROUND
Setup and Definitions. We consider a Neural Network (NN), f(..)) IR x R' + R given input x E Rº and network parameter 0 E R... Hereafter, for simplicity, we consider vectors as single-column matrices unless otherwise stated. We use the output of NN f(r, 8) as prediction for input æ. Let S == {I ),)), denote the independently and identically distributed (i.id.) training data drawn from true data distribution D, where I,, E IR and E I are input and output represen- tations of n-th training instance, respectively. For simplicity, we denote concatenated input and output of all instances as X' E RND and ) E R-NKK respectively, and f(x' 0) E R as a concatenation of {ff (I N Given prior distribution of network parameters p(0) and a likelihood function p(S = I P(,.lf 8)), Bayesian inference defines posterior distribution of network param- eter 0 as S(0 = exp(-L(3 (()))((())))))) (S)) where L(S, 0) ;= log p()) > log p(y, (x., 0) is training loss and Z(S)= := ((#)p(S)0)d is the normalizing factor. For example, the likelihood function for a regression task will be Gaussian: P(y|r, 0) (x, 0), o'Ix) where a is (ho- moscedastic) observation noise scale. For a classification task. we treat it as one-hot regression task following Lee et al. (2019a); He et al. (2020). While we adopt this modification for theoretical tractability, Hui & Belkin (2021) showed this modification offers good performance competitive to the cross-entropy loss. Details this modification are given in Appendix C.
'https://ee thub. com/ sungyubki m/ ccnnectn vi writy-tan tangent 1enggentt--eeel
2

---

Published as a conference paper at ICLR 2023
Laplace approximation. In general, the exact computation for the Bayesian posterior of a net- work parameter is intractable. The Laplace approximation (LA; MacKay (1992)) is proposed to approximate the posterior distribution with Gaussian distribution defined PLA PLA (E|S) ~ ( 0", (V)C(S (V 11) where 0" E Rº is a pre-trained parameter with training loss and V) 8*) E PxPP is a Hessian matrix of loss function w.r.t. parameter at l".
Recent works LA replace the Hessian matrix with (Generalized) Gauss- matrix make computation easier (Khan et al., 2019; Immer et al., 2021). With this approximation, the LA posterior of the regression problem can be represented
PLA (+)() ~ (vb10" (I//a (g/o=+))Jo/00
Damping
Curvature
where a, > 0, E RPPP is identity matrix, and Jo NKXP is concatenation of Jo(x, 0*) RKP (Jacobian of f w.r.t. 0 at input x and parameter 8*) along training input X. Inference with LA requires further sub-curvature approximation for modern NN architectures (e.g., ResNet (He et al., 2016) and Transformer ( (Vaswani et al., 2017)) because of the prohibitively large covariance matrix. This approximation includes diagonal, Kronecker- fectore- approximate curvature (KFAC), last-layer, and subnetwork approximation (Ritter et al., 2018; Kristiadi et al., 2020; Daxberger et al., 2021). Meanwhile, it is well known that proper selection of prior scale a is needed balance the dilemma between overconfidence and underfitting in LA
PAC-Bayes bound with data-dependent prior. We consider PAC-Bayes generalization error bound of classification task used in McAllester (1999); Perez-Ortiz et al. (2021) (especially equation (7) of Perez-Ortiz et al. (2021)). Let P be a PAC-Bayes prior distribution over I independent of training dataset S, and err(.,)) IR x K + [0, 1] be an error function defined separately from the loss function. For any constant E (0, 1] and > O, and any PAC-Bayes posterior distribution Q over
RP, the following holds with probability least - 8: err (Q) S errs
2N (f (2. o). y)). errs s (Q) E -,) (f 6). f).3)) y)}, and
where crrp (Q) := (( (aapiii-ee--er((ffffeeee
N denotes the cardinality of S. That is, erro (Q) and err (Q) generalization error and empirical eиror, respectively. The only restriction on P here is that it cannot depend on the dataset S.
Following the recent discussion in Perez-Ortiz et al. (2021), one can construct data-dependent PAC- Bayes bounds by (i) randomly partitioning dataset S into Sq and Sp so that they are independent, (ii) pre-training a PAC-Bayes prior distribution Pp only dependent of Sp (i.e., PP belongs a PaC-Bayes prior due to the independence of Sq), (iii) fine-tuning PAC-Bayes posterior distribution
Q dependent of entire dataset S, and (iv) computing empirical error errs, (Q) with target subset So
(not entire dataset S). In summary, we modify the aforementioned original PAC-Bayes bound through a data-dependent prior PT as
[L[ + 1og(2/No//4
errp (T) S errs (Q) +
+1
2NC
where Nç is the cardinality of Sq. We denote sets of input and output of partitioned datasets (Sy, Sq) by N3, Ve, Kq. V/2 for simplicity.
2.2 SCALE-INVARIANT PRIOR AND POSTERIOR FROM LINEARIZATION W.R.T. CONNECTIVITY
Our goal is to construct scale-invariant PD and Q To this end, we assume pre-trained parameter 0" E IR. with the prior dataset Sp. This parameter can be attained with standard NN optimization procedures (e.g., stochastic gradient descent (SGD) with momentum). Then, we consider linearized NN at the pre-trained parameter with an auxiliary variable c E R as
Jin g0th (z, c) == f(z, 0*) + Jo(z, O")ddia(0"")
(3)
where diag is a vector-to-matrix diagonal operator. Note that equation 3 is the first-order Taylor approximation of NN with perturbation 0* @ e given input x and parameter 6*: gp (a. = f(x, 0* + 0* o c) x lin (z, c), where @ denotes element-wise multiplication of two vectors. Here we express the perturbation in parameter space as 6* @ c instead a single variable such as 8 E R By decomposing the scale and connectivity of perturbation, this linearization design matches the
pert
3

---

Published as a conference paper at ICLR 2023
scale of perturbation (i.e., 8* C c) to the scale of 8* in component- manner. Note that similar decomposition was proposed in pruning at initialization (Lee et al., 2019c;b) to measure the importance of each connection independently of its weight. However, they only consider this form to predict the effect of each connection before pre-training.
Based equation 3, we define data-dependent prior (IP)) over connectivity
Po. (c) == [V(cOOp,e Ir).
(4)
This distribution can be translated to a distribution over parameter by considering the distribution of perturbed parameter (yé := 0 + 0 @ c): Fe- (v) - a didig((0"))))) )²). We now define the PAC-Bayes posterior over connectivity ((c) as follows:
)²)
Qa- Nc Σg),
(5)
HQ = EdJo (V' - f(x,00)
Endiag (0* (J) (V' -(x,(
(6)
x²
x²
P a
J o²
Ip a
diag(O JJ Jodiag(O")
Eq =
+
(7)
σ²
where Je E NK x P is concatenation of Jo(z, 0p) == Jo( 0* diaag(0") E RxP (i.e., Jacobian
of perturbed NN o0tt (xr e) w.r.t. at input x and connectivity 0p) along training input X. Indeed, Qo- the posterior of Bayesian linear regression w.r.t. connectivity c. We refer to Appendix for detailed derivations. Again, it is equivalent to the posterior distribution over parameter Qe- (v)
N (v=10~ + 0 O μcc (diag(@*)
+ J Jo ) where diag(0") -2 := (diag((*)
by as-
suming that all components of 0 are non-zero. This assumption can be easily satisfied by considering the prior and posterior distributions of non-zero components of NNs only. Although we choose this restriction for theoretical tractability, future works can relax it to achieve diverse predictions by considering the distribution of zero coordinates.
In summary, data-dependent PAC-Bayes bound can be computed with our PAC-Bayes distributions.
The validity of this data-dependent PAC-Bayes bound is ensured as follows: our PAC-Bayes prior depends the Sy through 8", but independent the Sq that measures the errors. Note that here two-phase training (i.e.. pre-training with Sü and fine-tuning with S) explained in Sec. 2.1 is used to attain our PAC-Bayes posterior. Similar ideas of two-phase training with linearization were proposed in the context of transfer learning in Achille et al. (2021); Maddox et al. (2021). In transfer learning, there is a distribution shift between Sp and Sq. Therefore, Sp cannot be used for their fine-tuning phase in contrast to our PaC-Bayes posterior.
Now we provide an invariance property of our prior and posterior distributions w.r.t. function- preserving scale transformations as follows: The main intuition behind this proposition is that Jacobian w.r.t connectivity invariant to the function-preserving scaling transformation, i.e., Jo (d))))ia
Ja (zx 0*) ddiag(O))) Representative cases of T in Proposition 2.1 are presented in Appendix to highlight theoretical implications; these include the case of WD applied to the general network, including BN.
Proposition 2.1 Scale-invariance of PAC-Bayes prior and posterior). Let T R R is
invertible diagonal linear transformation such that faz T())) = f(z, w), Vr Ere E RR.VERR" Then, both PAC-Bayes prior and posterior are invariant under T:
Pr(o-)(c) d Fo- (c),
Qr()(()) d Qc- (c).
Furthermore, generalization and empirical errors are also invariant to T.
2.3 RESULTING PAC-BAYES BOUND
Now we plug in our prior and posterior into the modified PAC-Bayes generalization error bound in equation 2. As a result, we obtain a novel generalization error bound, named PAC-Bayes-CTK, which is guaranteed to be invariant to scale transformations (hence without the contradiction of FM hypothesis mentioned in Sec. 1))
4

---

Published as a conference paper at ICLR 2023
Theorem 2.2 (PAC-Bayes-CCl CTK and its invariance). Let assume pre-trained parameter # with data Sp. By applying Pg. and Qo- to data dependennndnnt PAC- Bayes bound (equation 2). we get
KI. divergonce
Hake 4a No
(Ba) 4NQ
(2 No/0)
crrd(Qo-) & CrrS (Qo-)
+
(8)
2.o
(average) perturbution
sharpness
where { are eigenvalues of (Ip+ 32 J.) and h(x) == x - log(z) - 1. This upper bound scale transformation by Proposition 2.
is invariant to T for the
Note that recent works on FM contradiction focus only on the ccale-invariance of sharpness metrics: Indeed, their generalization bounds are not invariant to scale transformations due to the scale- dependent terms (equation (34) in Tsuzuku et al. (2020) and equation (6) in Kwon et al. (2021)). Specifically, these terms proportional to the norm of pre-trained parameters. On the other hand, the generalization bound in Petzka et al. (2021) (Theorem 11 in their paper) only holds for single-layer NNs, whereas our bound has restrictions for network structure. As a result, our PAC-Bayes bound is the first scale-invariant PAC-Bayes bound to the best of our knowledge.
The following corollary explains why we name PAC-Bayes bound in Theorem 2.2 PAC-Bayes-CTTK.
Corollary 2.3 (Relation between CTK and PAC-Bayes-CTK) Let us define empirical Connectivity Tangent Kernel (CTK) of S as C = JJJ = Jodiag(0"" (0" (JT E RNNKNK by removing below term. Note that empirical CTK has T(= NK) non-zero eigenvalues of (A:)T= then the followings hold for {3) in Theorem 2.2: (i) 3 = =0²/(0² + a² >2),))= < 1 = 1, . T and (ii) 3; = 1 for i = T + 1, ....PP. Since h( = O, this means P T terms of summattio in the sharpness part of PAC- PAC-Bayes- B-ye--CyyeKKKKaa vanish to Furthermore, this sharpness term of FAC-Bayes-CIF CTK is a monotonically increasing function for each eigenvalue of empirical CTK.
Corollary 2.3 clarifies why Ei= (B:)//4N in Theorem 2.2 is called the sharpness term of PAC- Bayes-CTK: A sharp pre-trained parameter would have large CTK eigenvalues (since eigenvalues of CTK measure the sensitivity of output w.r.t. connectivity), increasing the sharpness term and the generalization gap. Finally, Proposition 2.4 shows that empirical CTK is also scale-invariant.
Proposition 2.4 (Scale-invariance of empirical CTK). Let T IR"" + !R! be an cunction-preserving scale transformation in Proposition 2. Then empirical CTK at parameter ab is invariant under T:
(v) C% Yz,y E IR. Vou E R...
(9)
Remark 2.5 (Connections to empirical NTK). The empirical CTK C resembles the existing empir- ical Neural Tangent Kernel (NTK) at parameter (Jacot et al., 2018): 0xy == Jo(a, Joyy, T RKKK Note that the deterministic NTK in Jacot et al. (2018) is the infinitewiddth limiting kernel at initialized parameters, while empirical NTK can be defined on any (finite-width) NNs. We focus on empirical kernels for finite pre-trained parameters throughout the paper, and we leave deterministic kernels defined for future studies. Comparing empirical kernels, the main difference between empiri- cal CTK and the existing empirical NTK is in the definition of Jacobian. In CTK, Jacobian is computed w.r.t. connectivity e while the empirical NTK uses Jacobian w.r.t. parameters 9. Therefore, another PAC-Bayes bound can be derived from the linearization of flim (x. 5) == f(x 0*) - Jo(x, 0*)8. As this bound is related to the eigenvalues of 0%,, call this bound MAC-Baye--NTK and provide derivations in Appendix B. Note that PAC- MAC-BayesNNTK is scale-variant as O + O By in general.
2.4 COMPUTING APPROXIMATE BOUND IN REAL WORLD PROBLEMS
To verify that PAC-Bayes bound in Theorem 2.2 is non-vacuous, compute it for real-world problems. We use CIFAR-10 and 100 datasets (Krizhevsky, 2009), where the 50K training instances are randomly partitioned into Sp of cardinality 45K and Sq of cardinality 5K. We refer to Appendix H for detailed experimental settings.
To compute equation& 8, one needs (i) j perturbation term, (ii) C based sharpness term, and (iii) samples from PAC- posterior Qo- faq in equation 6 can be obtained by minimizing
5

---

Published as a conference paper at ICLR 2023
Table 1: Comparison between PAC-Bayes-CCT CTK and PAC-Bayes-NTTE for ResNet--8
CIFAR-10
PAC Bayes CTK
PAC PAC-BByes-We Nuy
scale
4,0
Trace (x:10 ")
=
1.91
1.91
+
1107.55 + 20.64
Penurbation
5.72
5.77 0.08
5.840005
636,55
564.38 7.56
438.21 10:73
288.24
Sharpness
28.92 0.221
28.96=0..2
28.95 0.22
28.95=0.00
728.80 2.09
602.17 2:77
441.91 + 2.08
KL
a7.5
17.34:0.15
17.39 1..36
682.68 = 4.74
583.27 2.88
470.27 4
365.07 2.96
Test (x)10P)
0.12
0.11
0.18
4.77+0.12
13.00
8.26+0.11
6,34 0.09
6.25
Bound (x:10l)
9:200004
9:21
0.02
9:21
39.00 0.61
32.07 0.25
28.24 0.06
24.84 0.11
CIFAR-100
PAC-Rayes-CTK
PAC-Rayes-NTK,
Parameter scale
0.5
10
2.0
4.0
0.5
1.0
2.0
4.01
Trace (x10--) +>
2.33 0.37
0.37
2
0.37
5830,55 + 532.26
1913-90
244.05
1089.55
104.06
955 05
66.16
Pentrration
14.54
1..32-022
1408
0.20
0.14
6,83
569.16 6.94
329
3.34
Sharpness
5:26
42.53
560.20
1 7.89
501.74
KL.
2:51
28.42 = 2.52
=
28.19
657,31
7.21
$74.68 e 8.95
489.77 3.57
417.02 3.98
Test erz. (0P)
21.7
21882 a..22.00
00.18
0.19
2.21
43.39 a 0.64
37.06 0.26
32.32 a 0.32
28.37
Bound (x:00)
27.34
2..76
27.75
27.75
68.44 0.82
59.96 0.19
$1,90 0.04
44,80 0.25
arg min L(c) = bे ID - -(x J + 2 e by first-order optimality condition. Note that this problem
convex optimization problem w.r.t. c, since is the parameter of the linear regression problem. We use Adam optimizer (Kingma & Ba, 2014) with fixed learning rate le-4 to solve this. For the sharpness term, we apply the Lanczos algorithm to approximate the cigenspectrum of C% following Ghorbani et al. (2019). We use 100 Lanczos iterations based on their setting. Lastly, we estimate empirical and test errors with 8 samples of CL/LL implementation of the Randomize-Thie Then- pseudo-code and computational complexity of RTO implementation can be found in Appendix F.
(RTO) framework (Bardsley et al., 2014; Matthews et al., 2017). The
Table provides the bounds and related terms of PAC-Bayes- CTK (Theorem 2.2) and NTK (Theorem B. First, we found that our estimated PAC-Bayes-CTK and NTK are non-vacuous (i.e., estimated bounds are better than guessing at random) for R8 with 11M parameters. Note that deriving non-vacuous bound is challenging in PAC-Bayes analysis: only few PAC-Bayes works (Dziugaite & Roy, 2017; Zhou et al., 2018; Perez-Ortiz et al., 2021) verified the non-vacuous property of their bounds, and other PAC-Bayes works (Foret al., 2020; Tsuzuku al., 2020) did not. To check the invariance property of AAC-Bayes-CTK, we scale the inver-invariant parameters in ResNet- 18 (i.e.. parameters preceding BN layers) for fixed constants {0.5, 1.0, 2.0, 4.0). Due to BN layers, these transformations do not affect the function represented by NN, and the error bounds should be preserved for scale-invariant bounds. Table shows that PAC-Bayes-CTK bound is stable these transformations. On the other hand, PAC-Baycs-NTK bound is very sensitive to the parameter scale.
2.5 CONNECTIVITY SHARPNESS AND ITS EFFICIENT COMPUTATION
Now, we focus on the fact that the trace of CTK is also invariant to the parameter scale by Proposition 2.4. Unlike MAC-Bayes-CTT and NTK, traces of CTK and NTK do not require onerous hyper- parameter selection of ó,a, a. Therefore, we simply define CS(0") := tr(C%) as a practical sharpness measure at 0*, named Connectivity Sharpness (CS) to detour the complex computation of This metric can be easily applied to find NNs with better generalization, similar to other sharpness metrics (e.g., trace of Hessian), as shown Jiang al. (2020). We evaluate the detecting performance of CS in Sec. 4.1. The following corollary shows how CS can explain the generalization performance of NNs, conceptually.
Corollary 2.6 (Connectivity sharpness, Informal). Let us assume CTK and KL divergence terms of PAC-Bayes- CTK as defined in Theorem 2.2. Then, CS vanishes to zero or infinity, the KL divergence term of CC-BByys-CTKK also does so.
As traces of matrix can be efficiently estimated by Hutchinson's 's method (Hutchinson, 1989), one can compute the CS without explicitly computing the entire CTK. We refer to Appendix F for detailed procedures of computing CS. As CS is invariant to function scale transformations by Proposition 2.4, it does not contradict the FM hypothesis.
3 BAYESIAN NNS WITH SCALE-INVARIANCE
In this section, we discuss a practical implication of our posterior distribution (cquation 5) used in the PAC-Bayes analysis. To this end, we first interpret PAC-Bayes posterior as modified result of

---

Published as a conference paper at ICLR 2023
Table 2: Correlation analysis of sharpness measures with generalization gap. We refer Sec. 4.1 for the details of sharpness measures (row) and correlation metrics for sharpness
sharpness-generaizzttioo generalization relationship (column).
trHH
tr(F)
i(("")
S0
PO
SM
PM
AS
FR
CS
(rank comt)
0.706
0.679
0.703
0.490
0.436
0.473
0.636
0.755
0.649
0.837
network depth
0.764
0.652
0.978
-0.358
-0.719
0.774
0.545
0.756
0.771
0.978
aetwork width
0.687
0,922
0.330
-0.533
-0.575
0.495
0.564
0.827
0.921
0.978
mini-batch size
0.976
0.810
0.988
0.859
0.893
0.909
0.750
0.829
0.685
0.905
learning rate
0.966
0.713
1.000
0.829
0.874
0.057
0.621
0.794
0.565
0.897
weight decay
-0.031
-0.103
0.402
0.647
0.711
0.168
0.211
0.710
0.373
0.742
V (avg.)
0.672
0.599
0.739
0.289
0.237
0.481
0.538
0.783
0.663
0.900
I (cond. MI)
0.320
0.243
0.352
0.039
0.041
0.049
0.376
0.483
0.288
0.539
LA (MacKay, 1992). Then, we demonstrate that this modification improves existing LA when WD is applied to NNs with normalization layers (Proposition 3.1).
One can view the parameter space version of Qa- as a modified version of LA posterior (equation 1) by (i) substituting parameter-dependent damping (iiag(0")=3 for isotropic damping and (ii) adding perturbation H* @ Gig to the mean of Gaussian distribution. Here, we focus on the effect of replacing the damping term of LA in sitth-normalized NNs in the presence of WD. We refer to Antoran et al. (2021; 2022) for the discussion on the effect of adding perturbation to the LA with linearized NNs.
The main difference between the covariance terms of LA in equation I and equation 7 is the dcfinition of Jacobian (i.e., parameter or connectivity) similar to the difference between empirical CTK and NTK in Remark 2.5. Therefore, we name PCL as Connectivity Laplace (CL) approximated posterior.
(CL(v) (15(5) ~ (wbl0", (diag((")) 2/a + JJ Jo/o²) -1)
To compare CL posterions against existing LAs, we explain how WD with BN can produce unexpected side effects of amplifying uncertainty. This side effect can be quantified if we consider linearized NN for LA, called Linearized Laplace (LL; Foong et al. (2019)). Assuming 0² < <a?, the predictive distribution of LL and CL are
a20 07 - f (xz, UPLA (118) N(f((z,)), a20 -
(10)
fli (zx (Pr (v|S) ~ (1 ax C²C C. -' x.)
(11)
for any input x E Rd where X' subscript means concatenation. We refer to Appendix G for the detailed derivations. The following proposition illustrates how WD with BN can increase the prediction uncertainty of equation 10.
Proposition 3.1 (Uncertainty amplifying effect for LL). Let assume that Wy :p - R is WD on scale-invariant parameters (e.g., parameters preceding BN layers) by multiplying < 1 and all the non-scale- invariann parameters are fixed. Then, the predictive uncertainty of LL is amplified by 1/72 > 1 while the predictive uncertainty of CTK is preserved as
Var warxtpp (e)s) ( (W,, (000000000 (Ir,00 1²) = Var (W))) (f (I, $))/7²
( Var Vara-perlen (w)S) (essss((f() (0-3 (Ix, 3) Var (pp) f0i0 v))
where Var(.) is variance of random variable.
Since the primal regularization effect of WD actually occurs when combined with BN as experimen- tally shown Zhang et al. (2019), Proposition 3.1 describes a real-world issue. Recently, Antoran et al. (2021; 2022) observed similar pitfalls in Proposition 3.1. However, their solution requires more complicated hyper- prper-parameter search: independent prior selection for each normalized parameter group. On the other hand, CL does not increase the hyper-parameter to be optimized compared to LL. We believe this difference will make CL more attractive to practitioners.
EXPERIMENTS
Here we describe experiments demonstrating (i) the effectiveness of Connectivity Sharpness (CS) as a generalization measurement metric and (ii) the usefulness of Connectivity Laplace (CL) as a general purpose Bayesian NN: With CS and CL, we can resolve the contradiction in the FM hypothesis concerning the generalization of NNs and attain stable calibration performance for various ranges of prior scales.
7

---

Published as a conference paper at ICLR 2023
Table 3: Test negative log likelihood on two UCI variants Beernnndez-Lobare & Adams, 2015; Foong et al.,
2019). We marked the best method among the four in bold and marked the best method among LL/CL in italics.
Original (Hemlndez-I pibato Adams, 2015)
variants (Fiong al., 2019)
Deep Ensemble
MCDO
L1
CL
Deep Ensemble
MCDD
LL
CL
hoston housing
2.90 0.03
0.01
2,85 + 0.61
2.88 0.02
2.71 0.01
2.68 0..01
2.74 + 0.01
2.75 0:01
concrete strength
3,06 0.01
3.20
1.02
3.11 0.02
4,03 1 0.97
3,42 =: 0.00
347 : 0.01
4.03 1 0.02
energy efficiency
0.74 2 0.01
1.92 0.01
2.12
0.83 2 0.01
0.77 a 0.01
1.78 0.01
2,02 2.02
0.90 a
kin8nm
-1,07 0,00
-0.80 0.01
-1.07 -0.90
-0,94 0.00
-0.71+ +000
-087 0.00
-0.93 + 0,00
naval propulsion
4.83 0.00
-3.85 + 0.00
0.00
.476 0.00
-2.22 0.35
-3.36-+0000
-3.66 0.11
-3.80 0.07
power plant
2.81 =: 0.00
2.91 0.00
2.00
2.81 0.00
2,91 1: 0.00
2.97 2: 0.00
2.98 0.00
2.91 :: 0.00
structure
2.89 0.00
2,96 + 0.00
1.91
2.89 0.00
3.11 0.00
3,07 0.00
3.07 0.00
3.13 0.00
wine
1:21 0.00
0,96
1.24 0.07
0.01
1.48 + 0.01
0.00
1,45 0.01
1.43 + 0.00
yacht So hydrodynamics
1.26 0 0.04
2.17 0.06
1.20 0.04
1.25 a 0.04
1,71 +: 0.03
3.06 2 0.02
1,78 0.02
1.74 1 0.01
4.1 CONNECTIVITY SHARPNESS AS A GENERALIZATION MEASUREMENT METRIC
Based on the CIFAR-10 10 dataset, we evaluate three correlation metrics to determine whether CS is more correlated with generalization performance than existing sharpness measures: (a) Kendall's cank-correlation coefficient (T; Kendall (1938)) (b) granulated Kendall's 's coefficient and their average (L;; Jiang al. (2020)) (c) conditional independence test (AC; Jiang et al. (2020)). In all correlation metrics, a higher value indicates a stronger relationship between sharpness and generalization.
We compare CS to the following baseline sharpness measures: trace of Hessian (tr (tr(H); Keskar et al. (2017)), trace of empirical Fisher (tr(F): Jastrzebski et al. (2021)), trace of empirical NTK at 8*, Fisher-Rao (FR; Liang et al. (2019)) metric, Adaptive Sharpness (AS; Kwon al. (2021) )), and four PAC-Bayes bound based measures: Sharpness-Orig. (SO), PAC-Bayes-Oni
(PO), Sharpness-Mag, (SM), and PAC-Bayes-Mage (PM). which are eq. (52), (49), (62), (61) in Jiang et al. (2020). We compute granulated Kendall's correlation by using five typer-parameter width, learning rate, weight decay, and mini-batch size) and three options for each. Thus, we train models with 3 = 243 different training configurations. We vary the depth and width of NN based on VGG-13 (Simonyan Zisserman, 2015). Further experimental details can be found in Appendix H.
(network depth, network
In Table 2, CS shows the best results for т, M and K. compared to all other sharpness measures. Additionally, granulated Kendall of CS is higher than other sharpness measures for out of hyperparameters and competitive with other sharpness measures for the remaining hyperparameters. The main difference between our CS and other sharpness measures is in the correlation with weight decay/network width: We found that SO and PM can capture the correlation with weight decay, and hypothesize that this is due to the weight norm term of SO/PO. As this weight norm term would interfere in capturing the (i.e., width/depth), SO/PO fail to capture correlation with network width in Table 2. On the other hand, CS/AS do not suffer from such problem. Also, it is notable that FR weakly captures this correlation despite its invariant property. For network width, we found that sharpness measures except for CS, tr (F). AS/FR fail to capture a strong correlation. In summary, only CS/AS detect clear correlations with all hy hyperparameters; among them, CS captures clearer correlations.
generalization correlation related to the number of parameters
4.2 CONNECTIVITY LAPLACE AS AN EFFICIENT GENERAI.-PURPOSS BAYESIAN NN
We evaluate CL's effectiveness as a general- purpose Bayesian NN using the UCI and CIFAR datasets. We refer to Appendix H for detailed experimental settings.
UCI regression We implement full-curvature versions of LL and CL and evaluate these to the UCI regression datasets Bternánddz--obbtee & Adams, 2015) and its GAP-variants (Foong et al., 2019) to compare calibration performance on in-between uncertainty. We measure test NLL for LL/CL and 2 baselines (Deep Ensemble (I (Lakshminarayanal et al., 2017) and Monte-Carlo DropOut (MCDO; Gal Ghahramani (2016)). Eight ensemble members are used in Deep Ensemble, and 32 MC samples are used in LL, CL, and MCDO. Table 3 shows that CL performs better than LL 6 of 9 datasets. Even though LL produces better calibration results on 3 of the datasets for both settings, the performance gaps between LL and CL are not as severe as on the other 6 datasets.
Image classification. We evaluate the uncertainty calibration performance of CL on CIFAR-10 and 100. As baseline methods, we consider Deterministic network, Monte-Carlo Dropout (MCDO; (Gal & Ghahramani, 2016)), Monte-Catt Batch Normalization (MCBN; (Teye et al., 2018)), Deep Ensemble (lakshminarayanaa et al., 2017), Batch Ensemble (Wen et al., 2020), and LLL (Khan et al.,
8

---

Published as a conference paper at ICLR 2023
Table 4: Uncertainty calibration results on CIFAR- 100 (Krizhevsky, 2009) for ResNet- (He et al., 2016)
CIFAR-100
NLL (1)
ECE (L)
Brier. (L)
AUC (1)
Deterministic
1,5370 # 0.0117
0.1115 # 0.0017
0.3889 : 0.0031
MCDO
1.4264 : 0.0110
0.0651 0.0008
0.3925 0.0020
0.6907 = 0.0121
MCBN
1,4689 # 0.0106
0.0998 0.0016
0.3750 0.0028
0.7982 0.0210
Batch Ensemble
1,4029 + 0.0031
0.0842 = 0.0005
0,3582 0.0010
0.7887 0.0115
Deep Ensemble
1,0110
0.0507
0,2740
0.7802
Linearized Laplace
1.1673 + 0.0093
0.0532 = 0.0010
0.3597 + 0.0020
0.8066 0.0120
Connectivity Laplace (Ours)
1.1307 # 0.0042
0.0524 + 0.0019
0.3319 0.0005
0.8423 0.0204
244
6.200
2.2
Det
Det.
0.5
Det:
Brer
1.4
0:125
6.1
14
6.90
16-)
10ml
10-1
10F
101
30
30P
129
10-
10-0
13-1
300
(a) NLL
(b) ECE
(c) Brier Score
Figure 1: Sensitivity Expected calibration error (ECE), Negative Log- likelhooddd (NLL), and Brier score results on corrupted CIFAR-100 ResNet-18. Showing mean (line) and standard deviation (shaded area) across four different seeds.
2019). We use the Randonize-TTenn-pppiiii (RTO) implementation of LL/CL in Appendix F. We measure negative log-likelihood (NLL), expected calibration error (ECE; Guo et al. (2017)), and Brier score (Brier.) for ensemble predictions. We also measure the area under receiver operating curve (AUC) for OOD detection, where we set the SVHN (Netzer et al., 2011) 1) dataset as an OOD,
Table 4 shows uncertainty calibration results on CIFAR-100. We refer to Appendix J for results on other settings, including CIFAR- 10 and VGGNet. Except for Deep Ensemble, CL shows better results than baselines for all uncertainty calibration metrics. Deep Ensemble performs best in 3 out of 4 metrics, but each ensemble member requires full training. LL and CL, however, require only post-hoc training on pre-trained NNs. Particularly noteworthy is that CL presents competitive results with Deep Ensemble, even with much smaller computations.
Robustness to the selection of prior scale. Figure shows the uncertainty calibration results over various a values for LL, CL, and Deterministic (Det.) baseline. As mentioned in previous works (Ritter et al., 2018; Kristiadi et al., 2020), the uncertainty calibration results of LL are extremely sensitive to the selection of a. Especially, LL shows severe under-fitting for large a (i.e., small damping) regime. On the other hand, CL shows stable performance in the various ranges of a.
CONCLUSION
In this work, we proposed new approach to enhance the robustness of generalization bound using PAC-Bayes prior and posterior distributions. By separating scales and connectivities, our approach achieved invariance to function- prrerering scale transformations, which is not addressed by existing generalization error bounds. As a result, our method successfully resolved the contradiction in the FM hypothesis caused by general scale transformation. In addition, our posterior distribution for PAC- PAC-Bayes analysis improved the Laplace approximation without significant drawbacks when dealing with weight decay with BN. To further improve understanding of NN generalization effects, future research could explore extending prior and posterior distributions beyond Gaussian distributions to more task-specific distributions. This could help bridge the gap between theory and practice.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENTS
This work was supported by the Institute of Information Communications Technology & Evaluation (IITP) grant funded by the Korea government (MSIT) (No.20 (No.2019-0-00075) 20.2017- 0-01779) and the National Research Foundation of Korea (NRF) grants (No.201 I55A1059921) funded by the Korea govermment (MSIT). This work was also supported by Samsung Electronics Co., Ltd (No.10201214-08133-0000
Planning
REFERENCES
Alessandro Achille, Aditya Golatkar, Avinash Ravichandran, Marzia Polito, and Stefano Soatto. Lqf: Linear quadratie fine-tuning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattem Recognition, pp. 15729-15739, 2021.
Javier Antoran, James Urquhart Allingham, David Janz, Erik Daxberger, Eric Nalisnick, and José Miguel Hernández-Lobato Linearised laplace inference in networks with normalisation layers and the neural g-prior. In Fourth Symposium Advances in Approximate Bayesian Infer-
ence, 2021.
Javier Antoran, David Janz, James U Allingham, Erik Daxberger, Riccardo Rb Barbano, Eric Nal- isnick, and José Miguel Hernández-Lobace Adapting the linearised laplace model evidence for modern deep learning. In International Conference on Machine Learning, pp. 796-821. PMLR, 2022.
Sanjeev Arora, Rong Ge, Behnam Neyshabur, and Yi Zhang. Stronger generalization bounds for deep nets via a compression approach. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 254-263, PMLR, 10-15 Jul 2018.
Johnathan M Bardsley, Antti Solonen, Heikki Haario, and Marko Laine. Randomize-then-ppiin A method for sampling from posterior distributions in nonlinear inverse problems. SIAM Journal on Scientific Computing, 3684):A 2014.
Christopher M. Bishop. Pattern Recognition and Machine Learning. Springer, 2006.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclau- rin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: composable transformations of Python+NumPy programs. github, 2018.
Erik Daxberger, Eric Nalisnick, James U Allingham, Javier Antoran, and Jose Miguel Hernandez- Lobato. Bayesian deep learning via subnetwork inference. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Learning. volume 139 of Proceedings of Machine Learning Research, pp. 2521. PMLR, 18-24 Jul 2021.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. Ieee, 2009.
Jacob Devlin, Ming- Wei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 1810::111000400 2018.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 1019-1028, PMLR, 06-11 Aug 2017.
Gintare Karolina Dziugaite and Daniel Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In Proceedings of the 33rd Annual Conference on Uncertainty in Artificial Intelligence (UAI), 2017.
Andrew YK Foong, Yingzhen Li, José Miguel Hernández-Lobato and Richard E Turner. 'in- between' uncertainty in bayesian neural networks, arXiv preprint arXiv: 1906. 1137, 2019.
10

---

Published as a conference paper at ICLR 2023
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. arXiv preprint arXiv:2010001412, 2020.
Yarin Gal and Zoubin Ghahramani. Dropout bayesian approximation: Representing model uncertainty in deep learning. In international conference on machine learning, pp. 1050-1059.
PMLR, 2016.
Behrooz Ghorbani, Shankar Krishman, and Ying Xiao. An investigation into neural net optimization via hessian eigenvalue density. In International Conference on Machine Learning, pp. 2232-2241. PMLR, 2019.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q. Weinherger. On calibration of modern neural networks. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 1321-1330. PMLR, 06-11 Aug 2017.
Bobby He, Balaji Lakshminarayanan, and Yee Whye Teh. Bayesian deep ensembles via the neural tangent kernel. H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 1010-1022. Curran Associates, Inc., 2020.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
José Miguel Hermández-Lobard and Ryan Adams. Probabilistic backpropagation for scalable learning of bayesian neural networks. In International conference on machine learning, pp. 1861-1869. PMLR, 2015.
Sepp Hochreiter and Jürgen Schmidhuber. Simplifying neural nets by discovering flat minima. In G. Tesauro, D. Touretzky, and T. Leen (eds.), Advances in Neural Information Processing Systems, volume 7. MIT Press, 1995.
Like Hui and Mikhail Belkin. Evaluation of neural architectures trained with square loss vs cross- entropy in classification tasks. In International Conference on Learning Representations, 2021.
Michael F Hutchinson. A stochastic estimator of the trace the influence matrix for laplacian smoothing splines. Communications in Statisties-Simllation and Computation, 1853):1059-1076 1989.
Alexander Immer, Maciej Korzepa, and Matthias Bauer. Improving predictions of bayesian neural nets via local linearization. In AISTATS, pp. 703-711, 2021.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In Francis Bach and David Blei (eds.). Proceedings of the 32nd International Conference on Machine Learning, volume of Proceedings of Machine Learning Research, pp. 448 456, Lille, France, 07-09 Jul 2015. PMLR.
Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018
Stanislaw Jastrzebski, Devansh Arpit, Oliver Astrand, Giancarlo B Kerg, Huan Wang, Caiming Xiong, Richard Socher, Kyunghyun Cho, and Krzysztof Geras, Catastrophic fisher explosion: Early phase fisher matrix impacts generalization. In International Conference Machine Learning, pp. 4772-4784. PMLR, 2021.
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantas- tic generalization measures and where to find them. International Conference on Learning Representations, 2020.
11

---

Published as a conference paper at ICLR 2023
Alex Kendall and Yarin Gal. What uncertainties do need in bayesian deep learning for com- puter vision? In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Maurice G Kendall. A new measure of rank correlation. Biometrika, 30(1/2): 1938.
Nitish Shirish Keskar, Jorge Nocedal, Ping Tak Peter Tang. Dheevatsa Mudigere, and Mikhail Smelyanskiy. On large-batch training for deep learning: Generalization gap and sharp minima. In 5th International Conference Learning Representations, ICLR 2017, 2017.
Mohammad Emtiyaz. E Khan, Alexander Immer, Ehsan Abedi, and Maciej Korzepa. Approximate inference turns deep networks into gaussian processes. In H. Wallach, H. Larochelle, A. Beygelz- imer, F. d' Alchh-BBue Buch E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 2014.
Agustinus Kristiadi, Matthias Hein, and Philipp Hennig. Being bayesian, even just bit, fixes overconfidence in relu networks. In International conference on machine learning, pp. 5436-5446.
PMLR, 2020.
A Krizhevsky. Learning multiple layers of features from tiny images. Master's 's thesis, University of Toronto, 2009.
Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. Asam: Adaptive sharpness- aware minimization for scale-invariant learning of deep neural networks. arXiv preprint arXiv:2102.11600, 2021.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019a.
Jaehoon Lee, Samuel Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: empirical study. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Process- ing Systems, volume 33, pp. 15156-15172. Curran Associates, Inc., 2020.
Namhoon Lee, Thalaiyasingam Ajanthan, Stephen Gould, and Philip HS Torr. A signal propagation perspective for pruning neural networks at initialization. arXiv preprint arXiv: wew.:/996606677, 2019b.
Namhoon Lee, Thalaiyasingam Ajanthan, and Philip Torr. SNIP: SINGLE-SHOT NETWORK PRUNING BASED ON CONNECTION SENSITIVITY. In International Conference on Learning
Representations, 2019c.
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape neural nets. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018
Tengyuan Liang, Tomaso Poggio, Alexander Rakhlin, and James Stokes. Fisher-rao metric, geometry, and complexity of neural networks. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 888 896. PMLR, 2019.
12

---

Published as a conference paper at ICLR 2023
Ekaterina Lobacheva, Maxim Kodryan, Nadezhda Chirkova, Andrey Malinin, and Dmitry P Vetrov. On the periodic behavior of neural network training with batch normalization and weight decay. Advances in Neural Information Processing Systems, 34:21545-2155 21556, 2021.
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:1608.03999 2016.
lya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2019.
David J. C. MacKay. A practical bayesian framework for backpropagation networks. Neural Comput. 4(3):448 472, may 1992. ISSN 0899-7667. doi: 10.1 10.1162/neco. 4.9.44334...9.....4444..
Wesley Maddox, Shuai Tang, Pablo Moreno, Andrew Gordon Wilson, and Andreas Damianou. Fast adaptation with linearized neural networks. International Conference on Artificial Intelligence and Statistics, pp. 2737-2745. PMLR, 2021.
Alexander G de G Matthews, Jiri Hron, Richard E Turner, and Zoubin Ghahramani. Sample-then- optimize posterior sampling for bayesian linear models. In NIPS Workshop on Advances in Approximate Bayesian Inference, 2017.
David A McAllester. Pac-bayesian model averaging. In Proceedings of the twelfth annual conference on Computational learning theory, pp. 164-170, 1999.
Raphael A Meyer, Cameron Musco, Christopher Musco, and David P Woodruff. Hutch++: Optimal stochastic trace estimation. In Symposium Simplicity in Algorithmss (SOSA), pp. 142-155. SIAM, 2021.
Kevin P Murphy. Machine learning: probabilistic perspective. MIT press, 2012.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y. Ng. Reading digits in natural images with unsupervised feature learning. In NIPS Workshop on Deep Learning and Unsupervised Feature Learning 2011, 2011.
Behnam Neyshabur, Russ R Salakhutdinov, and Nati Srebro. Path-sgd: Path- normalized optimization in deep neural networks. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Gamett (eds.), Advances in Neural Information Processing Systems, volume 28. Curran Associates, Inc., 2015a.
Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization deep learning. ICLR (Workshop). 201 5b
Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nathan Srebro. Exploring general- ization in deep learning. arXiv preprint arXiv: 1706. 07700088777 2017.
Roman Novak, Jascha Sohl-Dickstein, and Samuel S Schoenholz. Fast finite width neural tangent kernel. In International Conference on Machine Learning, pp. 77018-17044. PMLR, 2022.
Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, D. Sculley, Sebastian Nowozin, Joshua Dillon, Balaji Lakshminarayanan, and Jasper Snoek. Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift. H. Wallach, H. Larochelle, Beygelzimer, d' d'Alché- Buc, E. Fox. and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Maria Perez-Ortiz, Omar Risvaplata, John Shawe-Taylor, and Csaba Szepesvári. Tighter risk certifi- cates for neural networks. Journal of Machine Learning Research, 22(227):1-40, 2021.
Henning Petzka, Michael Kamp, Linara Adilova, Cristian Sminchisescu, and Mario Boley. Relative flatness and generalization. Advances in Neural Information Processing Systems, 34, 2021.
Jie Ren, Peter J. Liu, Emily Fertig, Jasper Snoek, Ryan Poplin, Mark Depristo, Joshua Dillon, and Balaji Lakshminarayanan. Likelihood ratios for out-of-distributee detection. In H. Wallach, H. Larochelle, A. Beygelzimer, d'Alché Buc, E. Fox, and R. Gamett (eds.), Advances in Neural Information Processing Systemx, volume 32. Curran Associates, Inc., 2019.
13

---

Published as a conference paper at ICLR 2023
Hippolyt Ritter, Aleksandar Botev, and David Barber. A scalable laplace approximation for neural networks. In International Conference on Learning Representations, 2018.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image
recognition. International Conference on Learning Representations, 2015.
Mattias Teye, Hossein Azizpour, and Kevin Smith. Bayesian uncertainty estimation for batch normalized deep networks. In International Conference on Machine Learning, pp. 4907-4916. PMLR, 2018.
Yusuke Tsuzuku, Issei Sato, and Masashi Sugiyama. Normalized flat minima: Exploring scale invariant definition of flat minima for neural networks using PAC-Bayesian analysis. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp- 9636 066666967. PMLR, 13-18 Jul 2020.
Joost Van Amersfoort, Lewis Smith, Yee Whye Tch, and Yarin Gal. Uncertainty estimation using a single deep deterministic neural network. Hal Daumé and Aarti Singh (eds.). Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 9690-9700. PMLR, 13-18 Jul 2020.
Twan Van Laarhoven. L2 regularization versus batch and weight normalization. arXiv preprint arXiv:1706.0533 05350, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Yeming Wen, Dustin Tran, and Jimmy Ba. Batchensemble: an alternative approach to efficient ensemble and lifelong learning. arXiv preprint arXiv:2002.06715, 2020.
M.A. Woodbury. Inverting Modified Matrices. Memorandum Report / Statistical Research Group, Princeton. Statistical Research Group, 1950.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization, 2017.
Guodong Zhang, Chaoqi Wang, Bowen Xu, and Roger Grosse. Three mechanisms of weight decay regularization. In International Conference on Learning Representations, 2019.
Wenda Zhou, Victor Veitch, Morgane Austern, Ryan P Adams, and Peter Orbanz. Non- vacuous generalization bounds at the imagenet scale: a pac-bayesian compression approach. arXiv preprint /804::8800.08800 2018.
14

---

Published as a conference paper at ICLR 2023
A NOTATIONS
Table 5: Notations used in the main paper
In E CRD Un E R
training inputs/outputs parameter of NNs
0, E R
[ 0)
output of NN given 2.
f0 (x, (I, 6)
f(x, 0-) + Ja(x, 0*) (linearization of NN w.r.t parameter)
90- (2,
f(x,0" + Jæ (ar 0* )diag(0") (linearization of NN w.r.t. connectiiv))) Perturbed NN w.r.. parameterrcooneeciiityy
fg (x, 6), I3.
pert
(z,c)
Oa
zero vector with dimension d
La
identity matrix with dimension d x d trace (i.c., sum of diagonal elements) of matrix A diagonal matrix whose diagonal elements correspond to function funcciinnprreeervigg preserving scale transformation (true) data distribution
(r())
diag(e)
T:Rº + R.
D
S
i.id. sampled training dataset from S randomly partitioned prior/posterior datasets from S
Sq,
Np, No
cardinality of S, Sp, Sq
No. Jo J, 8*) E KxP
concatenated inputs and outputs of all instances of S/Se/Sq-
Jacobian of outputs w.r. parameter/comneettvvgy given 2, 8* concatenated outputs of NN along X' concatenated Jacobian w.r.t. paramettrrconmeecivtty along X our PAC-Bayes prior/posterior
[(X, 0) E RNKP
Jg, Je E NNKP
Po. (c), Qa- (c)
eQT)(Q), e
generalization//emprrca error of PAC-Bayes posterior standard deviation of our PAC-Bayes prior scale (standard deviation) of observational noise mean and covariance of PAC-Bayes posterior empirical NTK/CTK of NN given v/ and input pair x, e RDD empirical NTK/CTK of NN for training inputs (A) given ubb eigenvalues of empirical CTK
NQ E R. En E PXP O" C ER 4F E NNN
DA (3,) 1(x))
eigenvalues of (Ip + J..
x logrr) (non-negative convex function w.r.t. 3, See Fig, 2) non-negative concave function w.r.. de. See Fig. 2)
(h o s)(z)
B PROOFS
B. PROOF OF PROPOSITION 2.1
Proof. Since the prior Po- is independent to the parameter scale, Po- (c) PT(,.)(c) is trivial. For Jacobian w.r.t parameters, have
0f(z T(e)) 0f(r,v)
0f(z, v) 00
(x, T(t))
Jo(a, v)
OT(0) 0T(0) 00 OT(0)
Then, the Jacobian of NN w.r.t. connectivity at T(b) holds
(T(1) = Jo(r, () TT
(12)
-o )(diag(t))
(13)
15

---

Published as a conference paper at ICLR 2023
where the first equality holds from the above and the fact that T is a diagonal linear transformation. Therefore, the covariance of posterior is invariant to T.
P diag diag(T((*))J (x, T(0 (x. ((d)aaaa))))
o²
Ip diag(0")J) (2 0" )diag(0"))
Lp diag(#" JJg Jadiag(f))
x²
Moreover, the mean of posterior is also invariant to T.
Endiag [T (0)) (10
o²
Egdiag (2, (x,T(0")) (Y - ((,
Eodiag(0") (2², (3' f(2,6')
x²
Endiag (0) "T (Y' f(),0))
Therefore, equation and equation are invariant to function cunction-preservigg preserving scale transformations. The re- maining part of the proposition related to the definition of sunction-preseree scale transformation T. For generalization error, the following holds
errp(QT(e-) = E(zz,)) (los))D.ommaai Terr, (err(]((z,)) v0), v)]
B( Qr(e) (err(ge- part (x. c)
E(z.a)-DoomCl [err(g6- pert (x.c),1)
((z.g))D, [err( (f +), v) = crrp (Qo-)
WLOG, this proof can be extended to the empirical emror errs..
B.2 PROOF OF THEOREM 2.2
Proof. (Construction of KL divergence) To construct MAC-Bayes-CTK, we need to arrange KL
divergence between posterior and prior follows:
KLIQ
(tr (2p (Ea + (ne - un)(e (e) + log Σz| log Cel - P)
1
tr(2p' (mQ - -m=)(uQ - ))) T))+
2 (tr(2>p Eo) + log [ip] log 1el - P)
/ (uq m²) T "I''ma (wg ue) (Ir(Sp Eσ)) log 15p' -Eg P)
1 (tr(Ep Ed) log Eاp Egl - p)
2a²
perturbation
sharpness
where the first equality uses the KL divergence between two Gaussian distributions, the third equality uses trace property (tr(AB) = tr( BA) and tr(a) = a for scalar a), and the last equality uses the definition of PAC-Bayes prior (Pa. (c) A((0pp, a Ip)). For sharpness term, first compute the
Ep Eq term
a
Ip
JJJ g2
16

---

Proof. In proof of Theorem 2.2, we showed that eigenvalues of Σp Eq can be represented as Published as a conference paper at ICLR 2023
Since a², o² 0 and J Je is positive semi-definite, the matrix E have non-zero eigenvalues of [3. Since trace is the sum of cigenvalues and the log-determinant is the sum of the log of cigenvalues, we have
Pade 20
P Σ (B, log(3;) 1)
M5Pe 2a
[P
+
h(3,)
2
where h((x) = x log(z)) - 1. By plugging this KL divergence into the equation 2, we get equation 8.
(Eigenvalues Ex² To show the scale-invariance of PAC-Bayes-CTK, is sufficient to show that the KL divergence between posterior and prior is scale-invariant: log(2 /6)/////2,, is independent to KL PAC-Bayes prior/posterior. We already show the invariance property of empiiriaalgeeneraiizatooo error term Proposition 2.1.
To show the invariance property of KL divergence, let us write singular value decomposition of Jacobian w.r.t. connectivity J. E RNK as J. = UEVT, where U RNNK x and V c RPxP orthogonal matrices and Σ E RNNKPP is rectangular diagonal matrix with descending order for singular values. Then, the following holds for 2ip -='Z Ie
20
Ip
o² JoJe
Ip
o VS
VE
ΣVT]
Ip
VT
where ΣT PxP is a diagonal matrix with X; == Au for 2 NK. Therefore,
eigenvalues of > Σq are I+a"X./o" o² Now, we consider Connectivity Tangent Kernel
(CTK) as defined in equation 2.3:
Cº = J.J = J NKXNK
Similar to Jo. CTK can be expressed as follows
C JJJ UEV VVVVVVETT" UΣE = UNU
where A' = M E NKNK As the smallest (P N K) eigenvalues of A = Σ' Σ are just zeros,
A' is just reduced diagonal matrix of A with these eigenvalues removed. As result, {A,) eigenvalues of CTK.
JNK
(Scale invariance of CTK) The scale-invariance property CTK is simple application of equa- tion 13:
CT()) J JT(o) (2,
T()))
Jo(n, "T
J, (1
= Jo(z, / (iddd(())))) (iagg) (o)
Cy Vz, E R Vbb E IR".
Therefore, CTK is invariant to any function-preserving scale transformation T and so do its eigen- values. This guarantees the invariance of Ep Eg and its eigenvalues. In summary, we showed the scale-invariance property of the sharpness term of KL divergence. Now all that remains is to show the invariance of the perturbation term. However, this is already proved in the proof of Proposition 2.1. Therefore, show PAC-Bayes-CTI CTK is invariant any Lunction-rresevvng scale transformation. D
B.3 PROOF OF COROLLARY 2.3
d²
a2
17

---

Published as a conference paper at ICLR 2023
Plotting h0d
Plotting of (h sxxd
N50
12
16
60
20
(a) y = h(z))
(b) (h o (I) where a = =a
Figure 2: Functions used in proofs
where {A:) eigenvalues of A = E. By SVD of Jacobian connectivity, A; Aa
D 2 0. Therefore, eigenvalues CTK are squares of singular values of J. and so X 2 0, Vi. As result, 3; < 1 for all i ,,..., (i.e., NK). Now all that remains to show that the sharpness
P for cigenvalues {3: of - and equality holds for A 0
of PAC-Bayes-CTK is
monotonically increasing function on each eigenvalue of CTK. To show this, we first keep in mind that
o²
s(x)
0² + a²x
monotonically decreasing function for 2 0 and h(x) log(x) is monotonically decreasing function for x E (0, 1]. Since the sharpness term of KL divergence
h(B)) 4No
(ho (AA.)
4No
This is monotonically increasing function for x 2 0 since s(x) S for x 2 0. For your information, we plot h(x)) and y = (h (x) in Figure 2.
B.4 PROOF PROPOSITION 2.4
We refer Scale invariance of CTK part of the proof of Theorem 2.2. This is direct application of the scale-invariance property of Jacobian w.r.t. connectivity.
B.5 PROOF COROLLARY 2.6
Proof. Since CS is trace of CTK, it is of the eigenvalues of CTK. As shown in the proof of Corollary 2.3, eigenvalues of CTK are square of singular values of Jacobian w.r.t. connectivity c. Therefore, the eigenvalues of CTK non-negative and vanish to zero if CS vanishes zero.
P =
Σ
-
A = 0 + 3: = s(A,) = 1 + h(B,) 0, Vi =1 P
1=1
This means the sharpness term of KL divergence vanishes to Furthermore, singular values of Jacobian w.r.t. e also vanish to zero in this case. Therefore, μ4q vanishes to zero, also. Similarly, if CS diverges infinity, this means (at least) one of the eigenvalues of CTK diverges to infinity. In this case, the following holds
- = -
1(3)) + xo, 11 P
Therefore, the KI. divergence term of PAC-Bayes-CTV also diverges infinity.
18

---

Published as a conference paper at ICLR 2023
B.6
PROOF OF PROPOSITION 3.1
Proof. By assumption, we fixed all non-scale invariant parameters. This means we exclude these parameters in the sampling procedure of CL and LL. In terms of predictive distribution, this be translated as
10 MlILI NUT(,, (00(2,00),, ",, 0
100 ((1)) (") a, a C C - Xxx
#f(n.0) #f(a' abl (v:)) for scale-invariant where 6% DiEp 0f(z,)) @5(=(w)) and C Eiep 00, 00;
80, 80;
parameter set P. Thereby, mask the gradient of the non-scale-invariant parameters Therefore, this can be arranged as follows
(2r, (J)ddag(()JJ(z./]
- (odiaa())diag(() (d)))d) e)
where 1p E Rº is a masking vector (i.e., one for included components and zero for excluded components). Then, the WD for scale-invariant parameters can be represented as
Wy(1), = (T v. M P. (s... if 4/ P.
Therefore, we get
(W)(v) 22x
= Jo(x, (()))))da((p)Jd(d(( (,,())
Jo(a, w)W, disg:
J,(n, (1op://))dd(((//)))(((...)))
= 1//2Jn(a, + didaaa [o
16 ma
for empirical NTK and
C (w)) xx'
J0 W, Xiii (
Jo (0 Jo,
Ja (ddii(())ddag(((( (ilpgg(111piigga)))d)e))
for empirical CTK. Therefore, we get
Alin
PLA PLA (+|8)
-2 -1
(o (x, 16) PCr. (vb8) (N(( 0*), a²² -000 O -2'))
This gives us
(LIs) (HV, (zx v)) = Var (W|s) (fili \)//
Varo (/ v))
(ar v)
B.7 DERIVATION OF PAC-BAYES-NTK
Theorem B.1 (PAC-Bayes- Let us assume pre-trained parameter 0" with data Sp. Let us assume PAC-Bayes prior and posterior as
(8) = NN 0p. a²Ip)
(14)
Q/ (8) = N66 )x))
(15)
Egrl () (x) σ²
poy
(16)
Ip
Jo
EV
(17)
?
19

---

Published as a conference paper at ICLR 2023
By applying P'o., Vo- to data depeddeendnnt PAC-Bayes bound (equation 2), get
KL divergence
Hep May
h (3)
log(2/Nor/)
errp(Q'o.) S errs (
(18)
4a No
AN
2N
(average) perturbation
shorpness
where ( are eigenvalues of (Ip S Jo) and h(x) = log(z)) 1. This upper bound not invei-ntv in general.
Proof. The main difference between PAC-Bayes-CTK and PAC-Bayes-NTK is the definition of Jacobian: PAC-Bayes-CTK uses Jacobian w.r.t connectivity and PAC-Bayes-NTK uses Jacobian w.r.t. parameter. Therefore, Construction of KL divergence proof of Theorem 2.2 is preserved except
Ep'Exy = (Ip + x² Jo)
and 8 are eigenvalues Ep. Eg. Note that these eigenvalues satisfy
o²
3 02 a
where X/ are eigenvalues of J Jg.
scale transformation to NTK). In contrast the CTK, the scale invariance property is not applicable to the NTK due the Jacobian w.r.t. parameter:
Remark B.2
Ja(x ((b)) σT [(x, T(+)) अहल [( v) Ja (b) 1b)
we assume all parameters are scale-invariant (or equivalently masking the Jacobian for all non- scale-inver inlee-ivvaiiatt parameters as in the proof of Proposition 3.1), ), the scale of NTK is proportional the inverse scale of parameters.
C DETAILS OF SQUARED LOSS FOR CLASSIFICATION TASKS
For the classification tasks in Sec. 4.2, we use the squared loss instead the cross-entropy loss since theoretical results are built the squared loss. Here, we describe how use the squared loss to mimic the cross-entropy loss. There are several works (Lee et al., 2020; Hui Belkin, 2021) that utilized the squared loss for the classification task instead of the cross-entropy loss. Specifically, Lee et al. (2020) used
C(8,0)
[lf 0) vil12 (xi.u))S
Σ
2NK
where C/ is the number of classes, and Hui & Belkin (2021) used
(((x,)) (), 2K
1
((f((2,0) 2 +
Σ 1=1,4
fl(z, 0)2
for single data loss, where (((2, c), 8) is sample loss given input x, target and parameter f( 0) e R is the i-th component of f(x, 0) IR k and M dataset-specific hyper-parameters.
These works used the mean for reducing the vector-valued loss into a scalar loss. However, this be problematic when the number of classes is large. When the number of classes increases, the denominator of the mean (the number of classes) increases while the target value remains 1 (one-hot label). As result, the scale of gradient for the target class becomes smaller. To avoid such
20

---

Published as a conference paper at ICLR 2023
unfavorable effect, the sum for reducing vectur-valued loss into scalar loss instead of taking the mean, i.e.,
K Σ i=1,
{((x, c),
-
Mxxx0)3
This analysis is consistent with the hyper-parameter selection in Hui & Belkin (2021). They used larger and M as the number of classes increases (e.g., k 1, M for CIFAR-10 10 (Krizhevsky, 2009), but & = 15, . 30 for ImageNet (Deng et al., 2009)) which results in manually compensating the scale of the gradient to the target class label.
D DERIVATION OF PAC-BAYES POSTERIOR
Derivation of Qo-
For Bayesian linear regression, we compute the posterior of [ E IR
z4, + Ets for 1. 1...,M
where ~ (O0 o²) is i..d. sampled and the prior of 3 is given as B ~ a(0p,a²Ip)) By concatenating this, we get
x3 +E
where y E RM, X E R Mxp are concatenation y,s 2,, respectively and æ ~ NoO II It well known (Bishop, 2006; Murphy, 2012) that the posterior of for this problem is
(8)p,
ΣXT y
11 =
x²
Ip x x -1
a
Similarly, we define the Bayesian linear regression problem
V/4 f(=,, 0*) Jo( 0" ""d)dd(a))))) E for NK
where N and the regression coefficient is 3 = c this case. Thus, we treat y: - f(x, 8*) as target and J& (z,, 0* (diag(O") an input of linear regression problem. By concatenating this, get
y f(X, 0*) + J.c + + = + (2) -(X, = J.c
By plugging this into the posterior of the Bayesian linear regression problem, we get
Qo (c) = Ncluue Eg)
pQ = EQJT (y - -f(X,0"))
Endiag 0 JJ (y f(X,
0²
r²
Ip
J Ja
-1
Ip a
diag(O' JJ Jodiag(O")) -1
Σq = a
+
σ²
=
σ²
Derivation of Qa- (x)) We define perturbed parameter s as follows
s = 0" + 0" Φ c.
Since ve is affine to we get the distribution of 1b
Qo (v) = NVvlu v0 )
F = + 3 pa
(diag(0))22
JJ J
-1
5 = diag((") Eediag((0*)
04
+
a
21

---

Published as a conference paper at ICLR 2023
E REPRESENTATIVE CASES OF FUNCTION-PRESERVING SCALING TRANSFORMATIONS
Activation-wiz wise rescaling transformation (Tsuzuku et al., 2020; Neyshabur et al., 2015a) For NNs with positive homogencous (c.g., ReLU) activations, following holds for Væ E IR > (xr,0) (az (o)). where rescale transformation Rr,,,( is defined as
7. O. if 0; E (param. subset connecting as input edges to k-th activation at I-th layer)
(R,,,,x()))
0./7 0g E (param, subset connecting as output edges to k-th activation at I-th layer)
0, in the other cases
(19)
Note that R.,,, (-) is finer-grained rescaling transformation than layer-wise rescaling transformation (i.e., common for all activations in layer 1) discussed in Dinh et al. (2017). Dinh et al. (2017) showed that even layer-wise rescaling transformations can sharpen pre-trained solutions in terms of the trace of Hessian (i.e., contradicting the FM hypothesis). This contradiction also occurs in previous PAC-Bayes bounds (Tsuzuku et al., 2020; Kwon et al., 2021) due to the scale-dependent term.
For example, let us assume linear NN as
(1
f(z, (0,, 02 03, 04)) (0),
=0,0g0 + 020
02
Then, the Jacobian of this NN would be
Aggr 04
Jo(2, 0) 01
0z4
From above discussion, (0.5, 1, 2) = R2,1, ((1, 1)) would change the Jacobian while main-
taining the predictions. To show this,
f(z, 1, 1)) fxz (1,1,1,1) = 2x 1,
for all x E IR. However,
(o (0.5, 1, 1))
(0.5z
2x
x
x
x Jo (IT, (1,1,1,1))
0.5z
in general. Therefore, we verified that the iiliiaationwiiss rescaling transformation R2,1,1 is valid dunction-preserving scale transformation.
WD with BN layers (Toffe Szegedy, 2015) For parameters W E R preceding BN layer, (N((diag())W) = BN((w))
(20)
for an input u E 1== and a positive vector 7 E IRT'. This implies that scaling transformations on these parameters preserve function represented by NNs for Va: E Rd, E TR [xx =
S1,,, where scaling transformation (-),,(()) is defined for = 1, 1,...
(S,.1. = 7k 0, if 0g E [param. subset connecting as input edges to k-th activation at 1-th layer)
0 for 6; in the other cases
(21)
Note that the WD (Loshchilov Hutter, 2019; Zhang et al., 2019) can be implemented as realization of (-). (e.g.. 7 0.9995 for all activations preceding BN layers). Therefore, thanks Proposition 2.1 and Theorem 2.4, our CTK-based bound is invariant to WD applied parameters before BN layers. We also refer to (Van Laarhoven, 2017; Lobacheva et al., 2021) for the optimization perspective of WD with BN.
For example, let us assume a BN layer with single activation as BN(Ox) where x, 0 E R and batch of input is given as 0, 0, 2, 2. Then, the Jacobian of NN would be
J (x, 0
T 0
For simple two layer linear NN f(x) == (2o(W1x) with weight matrix W1, W2. the first case of equation 19 corresponds to k-th row of W1 and the second case of equation 19 corresponds to k-th column of W2.
22

---

Published as a conference paper at ICLR 2023
f(x, 0) 6z and the denominator is detached from backpropagation computation for auto- differentiation packages (e.g., TensorFlow, Pytorch, and JAX). From the above discussion, 0.9995 10.9995,1.141) would change the Jacobian while maintaining the predictions. To show this, f(z 1) = fxz 0.9995) = x
|o]
for all inputs. On the other hand,
J ( 0.9995)
+ Jo 1).
0.9995
Therefore, we showed that the WD with BN layer 10,1995,11 is valid function-preserving scale transformation.
F IMPLEMENTATIONS
F.1 RTO IMPLEMENTATION OF CONNECTIVITY LAPLACE
Algorithm 1 Connectivviyy Laplace RTO
Require:
training data S, pre- gre-trained parameter 6", number of samples M, prior scale observa- tional noise scale
for m = 1, M do Sample parameter noise ej Ne Op, Ip) and label noise eth N(e)ONk I for f ,....,T' do
Randomly draw mini-batch B from S.
Define g|in (Xg, emh) = f (Ax, 0*) Jo(X', 8" (ddia(("))c" for mini-batch input Xg.
Define L(c 2 1D's + =B ge (Kg, cl2 + 2]8|a lle cp l12.
Compute backpropagation L(c")) w.r.t. connectivity cm. Update et' with SGD optimizer.
end for
end for return Samples from Connectivity Laplace {c" M
To estimate the empirical/genneaaizaato bound in Sec. 2.4 and calibrate uncertainty in Sec. 4.2, we need sample e from the posterior Qo- (c). For this, we sample perturbations in connectivity space
Ip
J
N~W 810p.
so that μo 6 for equation 6. To sample this, we provide novel approach to sample from LA/CL without curvature approximation. To this end, we consider the following optimization problem
arg min L(c) arg min
2No2 | - f(x -Jcc E|12 + 2Na ||c collž
where E W(E o² INK) and cu ~ N(ccllpp By first-order optimality condition, we have
NV. J. f(x,0' Jec +6) +
c² c
co =pp.
a²
By arranging this w.r.t. optimizer we get
0²
J
J. ((00-1(200)) +
a -0 JnE
J.J. o² Ip
-1
002
J. - -(xx00)))))))
+ J
Ip
40
+
J.
Ip
+ J.J
-1 J ( - f(x,0))
Ip
J Je
-1
Co
J.
+
+
+
a
o²
o²
a
o²
a
o²
Deterministic
Stochastic
23

---

Published as a conference paper at ICLR 2023
Since both and eu sampled from independent Gaussian distributions, have
2 =
Cu
+ +
J
~
2/0p.
Ip Je
a²
a²
Therefore, optimal solution of randomized optimization problem arg min, L(c)
is
c~ (cl
Ip Je
JT 0*)
Ip J.Je
Wc/ue Ex)
o²
σ²
a
r²
Similarly, sampling from CL can be implemented the following optimization problem.
arg min L(c) == arg min 2No
|e - =|l² + 2Na ||e collé
where (N INKK) and co ap, Since we sample the noise of data/perturbation and optimize the perturbation, this can be interpreted as a Randomize-Theer Optimize implementation of Laplace approximation and Connectivity Laplace (Bardsley al., 2014; Matthews et al., 2017). In Algorithm I, we provide pseudo-code for the RTO implementation of CL. Note that both time and memory complexity of computing linearized NN for mini-batch B is comparable to forward propagation as shown in Novak et al. (2022) with jax.jvp function in JAX (Bradbury et al., 2018). Therefore, the time/memory complexity of mini-batch JVP would be ((1B)L)))) + LW2 + N KK for MLPs with width W and depth L.
F.2
COMPUTING CONNECTIVITY SHARPNESS
Algorithm 2 Hutchison's method for computing Connectivity sharpness
Require: training data S, pre-trained parameter 0", number of samples M 38
for m = m= 1,..,, M do Sample nm ((EeOO I'INK) for t = ,,,,,, do
T
Sequentially draw a mini-batch from S.
Compute vector-Jacobian product UB 2g (")agg((")
Compute Tg = Fg + ||u 1//
end for end for return Estimated Connectivity Sharpness 2B
It well known that empirical NTK Jacobian is intractable in the modern architecture of NNs (e.g., ResNet (He et al., 2016) or BERT (Devlin et al., 2018)). Therefore, one might wonder how Connectivity Sharpness (CS) can be computed for these architectures. One can compute CS with Hutchison's 's method (Hutchinson, 1989; Meyer et al., 2021) as it is defined as trace of empirical CTK. According to Hutchison's 's method, trace of a matrix A E Rmxm is
tr(A) = (AIL)) = tr(AE, [22 1) E. (.rtt T) = E,(tr(2] Az]] T2
where E Rm is a random variable with cov(s) I, (e.g., standard normal distribution or Rademacher distribution). Since A C = J.J R.NK in our case, we further use mini- batch approximation compute A2: (i) Sample 2 E TRME MK from Rademacher distribution for mini-batch M with size M and (ii) compute v/
(X'. 0,,) T 2g IR." with JectbrJJaoobian product of JAX (Bradbury et al., 2018) (or it can simply computed using standard backprop) and (iii) compute z/ = ||: I2. Then, the sum of z for all mini-batch in the training dataset is Monte-Carlo approximation of CS with sample size 1. Empirically, we found that this approximation sufficiently stable to capture the correlation between sharpness and generalization as shown in Sec. 4.1. In Algorithm 2, we provide pseudo-code for the implementation. Note that both time and memory complexity of computing ui is comparable to a backward propagation as shown in Novak et al. (2022) with jax.vjp function in JAX (Bradbury al., 2018). Therefore, the time/memory complexity of width W and depth L..
VJP would ((IBBLW"))/( + LW2 + NKW for MLPs with
24

---

Published as a conference paper at ICLR 2023
G PREDICTIVE UNCERTAINTY OF CONNECTIVITY/LIEAARZZZ LAPLACE
In this section, we derive predictive uncertainty of Linearized Laplace (LLL) and Connectivity Laplace (CL). By matrix inversion lemma (Woodbury, 1950), the weight covariance of LL
(2//a² + J,(K Jo( "'o (=)(2 I'I,- a (2 (x No O xx) (1a))"(o(( (x(
Therefore, o²/ + 0, then the weight covariance of LL converges to
a I - a² J ((, 6*) 0 Ja (*). With this weight covariance and linearized NN, the predictive uncertainty of LL (xr, PLA (0|S) ~ (z 0"), a²0 Zx Similarly, the predictive uncertainty of CL. is f (zc, 0)| (0) ~ NZ NN) C* a CZx CX
a² zx
-6'-1 C
H DETAILED EXPERIMENTAL SETTINGS
H.1 BOUND ESTIMATION
We pre-train ResNet-18 (He et al., 2016) with a mini- size of 1K on Sp with SGD initial learning rate 0.4 and momentum 0.9. We use cosine annealing for learning rate scheduling (Loshchilov & Hutter, 2016) with warmup for the initial 10% training step. We fix = 0.1, a = 0.1, and chosen so that the PAC-Bayes-CTV and NTK
1.0 to compute equation 8. These values bounds fall within the non-vacuous range. We use random seeds to compute the standard errors. Additional results few pre-training data with NÇ- 000 and No 45, 000 presented in Appendix I.
H.2 CONNECTIVITY SHARPNESS
Table 6: Configuration of hyper-
network depth network width
1, 2, 3
32, 64, 128 0.1, 0.032, 0.001 0.0, le-4, 5e-4 256, 1024, 4096
learning rate
WD
mini-batch size
To verify that the CS has a better correlation with generalization performance compared to existing sharpness measures, (Kendall, 1938) which considers the consistency of a sharpness measure with generalization gap (i.e.. if one has higher sharpness, then so has higher generalization gap) (b) granulated Kendall's coefficient (Jiang et al., 2020) which examines Kendall's 's conk-correlation coefficient w.r.t. individual hyper-parameters to separately evaluate the effect of each conditional independence test (Jiang et al., 2020) which captures the causal relationship between measure and generalization.
evaluate the three metrics: (a) Kendall's rank-correlation coefficient
generalization gap (c)
Three metrics are compared with the following baselines: trace of Hessian (Ir(H); (Keskar et al., 2017)), trace of Fisher information matrix (t;;;;; (Jastrzebski et al., 2021)), trace of empirical NTK at 0* (tr (tr(O"" )), and four PAC-Bayes bound based measures, sharpness-orig (SO), PAC-Baycs-orig (PO), 1/a' sharpness mag (SM), and 1/o' PAC-Bayes mag (PM), which eq. (52), (49), (62), (61) in Jiang et al. (2020).
For the granulated Kendall's coefficient, use saper-parameters network depth, network width, learning rate, WD and mini-batch size, along with 3 options for each aaper-parameters as Table 6. We the VGG-13 base model and adjust the depth and width of each conv block. We add
25

---

Published as a conference paper at ICLR 2023
Table Example of network configuration with respect the depth width 128 in (Simonyan Zisserman, 2015)-style.
ConvNet Configuration
input (224 224 RGB image)
Conv3-128 BN
ReLU
MaxPool
Conv3-256
BN
ReLU
MaxPool
Conv3-512
BN
ReLU
MaxPool
Conv3-1024
BN
ReLU
MaxPool
Conv3-1024
BN
ReLU
MaxPool
FC-4096
ReLU
FC-4096
ReLU
FC-1000
BN layers after the convolution layer for each block. Specifically, the number of convolution layers of each conv block is the depth and the number of channels of convolution layers of the first conv block is the width. For the subsequent conv blocks, we follow the original VGG width multipliers (x2, x4, x8). An example with depth and width 128 is depicted in Table 7.
We SGD optimizer with momentum of 0.9. We train each model for 200 epochs and cosine learning rate scheduler (Loshchilov & Hutter, 2016) with 30% of initial epochs as warm-up epochs. The standard data augmentations (padding, random crop, random horizontal flip, and normalization) for CIFAR-10 10 are used for training data. For the analysis, we only use models with above 99% training accuracy following Jiang et al. (2020), As a result, we use 200 out 243 trained models for our correlation analysis. For every experiment, we use 8 NVIDIA RTX 3090 GPUs.
H.3 BNN EXPERIMENTS
UCI regression We train MLP with a single hidden layer. We fix = and choose a from (0.01, O.I, 10, 100) using liglliillooodd of validation dataset since the optimal varies for each regression dataset. We use random seeds to compute the average and standard error of the test negative log- loelliklisooos.
Image classification task We pre-train models for 200 epochs 1IFAR-10/100 dataset (Krizhevsky, 2009) with BesNet--88Hee et al., 2016) as mentioned in Section 2.4. We choose ensemble size M
26

---

Published as a conference paper at ICLR 2023
except Deep Ensemble (Lakshminarayanan et al., 2017) and Batch Ensemble (Wen al., 2020). We use ensemble members for Deep Ensemble and Batch Ensemble due to computational cost. We used 4 random seeds to compute the standard errors except for Deep Ensemble, which ensembles the NNs from 4 different random seeds.
For evaluation, we define prediction of dingle-memer the one-hot representation of network out- put with label smoothing. We select the label smoothing coefficient as 0.01 for CIFAR-10 and 0.1 for CIFAR-100. We define ensemble prediction an averaged prediction of single-member predictions. For OOD detection, we use the variance of prediction in output space, which is competitive to recent OOD detection methods (Ren et al., 2019; Van Amersfoort al., 2020). We use 0.01 for a and select the best a with cross-validation. For every experiment, we used 8 NVIDIA RTX 3090 GPUs.
27

---

Published as a conference paper at ICLR 2023
I ADDITIONAL RESULTS ON BOUND ESTIMATION
Table Results for experiments on PAC-Bayes-CTKnn estimation with N2 5, 000 and No. 45, 000.
We use 4 random seeds to compute error bars.
CIFAR-10
PAC-Bayes-CTK
PAC-Bayes-Ny NTX
Parameter scale
0.5
1.0
2.0
4.0
0.5
1.0
2.0
4.0
Trace (x10
0.10
1.31 +.113
=.31
+.11
541.36 32 52.58
139,50 9.06
+4.50
48.86 0.90
Perturbation
99:50
338.85
394.21
833887
202.53 4.144
315.99
447.97
549.35 7.64
Shorpness KL.
64.70+
42.30
174.25+5.03
116.174490
175.38 4499
174.68
133.61 + 1.48
184.27 1.24
246.11 1.86
295.83 5.82
Test (x10%)
23.75
23877 0.19
23.94 0.22
2
31.33 + 0.30
27.99 0.20
26.15 0.22
25.64 0.20
1x1)F)
CIFAR-100
PAC-Bayes-CTT CTK
PAC-Bayes NTK
Parameter scale
0.5
1.0
2.0
4.0
0.5
1.0
2.0
4.0
Trace (x100*)
0.14
2:22
0.14
0.12
655.76 + 43
167.02 5.78
5.78
57.01 = 3.92
50.57 4,71
Pertorhation
445.76-7.75
Sharpness
+7.11+0..0
66.36
KL.
232.49
232.48
4.116
2 3.774
233.01 5.26
338.15 13:96
417.55 23.06
432.85 27.41
Test (x10P)
65.07 D:33
65.15 0.38
65.21
0..22+0.39
69.92 0.50
69.03 0.38
68.37 0.49
68.63 0.50
0.39
J ADDITIONAL RESULTS ON IMAGE CLASSIFICATION
Table 9: Uncertainty calibration results on CIFAR- (Krizhevsky, 2009) for ResNet-18 (He al., 2016).
CIFAR-10
NLL (L)
ECE (L)
Brier. (4)
AUC (1)
Deterministic
0.3135 0.0088
0.0350 + 0.0014
0.0875 0.0026
MCDO
0.2845 0.0148
0,0157 # 0.0012
0.1056 0.0062
0.9172 + 0.0074
MCBN
0.2922 0.0074
0.0325 + 0.0010
0.0838 0.0022
0.9431 + 0.0033
Batch Ensemble
0.2740 0.0030
0,0286 + 0.0005
0.08 14 = 0.0009
0,9376 + 0.0036
Deep Ensemble
0,1753
0.0067
0.0594
0,8527
Lincarized Laplace
0.2077 0.0032
0.0180 + 0.0009
0.0816 0.0010
0.8991 0.0198
Connectivity Laplace (Ours)
0.2089 0.0023
0.0120 + 0.0019
0.0720 = 0.0011
0.9705 + 0.0098
Table 10: Uncertainty calihration results on CIFAR-10 (Krizhevsky, 2009) for VGG-13 (Simonyan & Zisserman,
2015).
CIFAR-I0
NLL (L)
ECE (L)
Brier. (1)
AUC (f)
Deterministic MCDO
0.4086 0.0018
0.0490 + 0.0003
0.1147 # 0.0005
0.3889 0.0049
0.0465 + 0.0009
0.1106 0.0015
0.7765 0.0221
MCBN
0.3852 0.0012
0.0462 + 0.0002
0.1108 = 0.0003
0.9051 0.0065
Batch Ensemble
0.3544 0.0036
0.0399 + 0.0009
0.1064 = 0.0012
0.9067 + 0.0030
Deep Ensemble
0.2243
0.0121
0.0776
0.7706
Linearized Laplace
0.3366 0.0013
0.0398 + 0.0004
0.1035 = 0.0003
0.8883 + 0.0017
Connectivity Laplace (Ours)
0.2674 0.0028
0.0234 + 0.0011
0.0946 + 0.0010
0,9002 + 0.0033
28

---

Published as a conference paper at ICLR 2023
Table 11: Uncertainty calibration results on CIFAR-I00 (Krizhevsky, 2009) VGG-13 (Simonyan & Zisser-
man, 2015).
CIFAR-100
NLL (L)
ECE (U)
Bricr. (1)
AUC (t)
Deterministic
1.8286 0.0066
0.1544 + 0.0010
0.4661 : 0.0018
MCDO
1.7439 0.0089
0.1363 # 0.0008
0.4456 # 0.0017
0.6424 4 0.0099
MCBN
1.7491 0.0075
0.1399 + 0.0010
0.4488 = 0.0015
0. + 0.0197
Batch Ensemble
1.6142 0.0101
0.1077 + 0.0020
0.4143 0.0027
0.7232 + 0,0021
Deep Ensemble
1.2006
0.0456
0.3228
0.6929
Lincarized Laplace
1.5806 0.0054
0.1036 + 0.0004
0.4127 = 0.0010
0.6893 0.0221
Connectivity Laplace (Ours)
1.4073 0.0039
0.0703 0.0028
0.3827 0.0012
0.7254 0.0136
K TRANSPOSED TABLE
Table 12: Transposed Table
Ieston Fonelts
s7ecery, etrrph
meng micerry
kirdern
Fans grppahine
DOAEr plare
wier
Yotht Buttodyamiiiss
Oriazed
MCDI
2,85 6.01
8101
1.30
1:24.000
1 B63
4190+000 0011
418040001
291 860
291 8.00
290 1120
11:00
4
0.9 0161
1.20 6:84
135 10:00
1100
11080
GAP varooes
MCD0
248 1000
2..
8.4
3/42 1130
3.0 801
0.57 2:00
3.16-2611
297 1109
0.0
0.00
1150
1 801
1.78
Du22
Table 13: First columns of transposed table
boston housing
concrete strength
energy -ffficieny
kin8nm
naval propulsion
Deep Ensemble MCDO
2.50
3.06
0.04
-1.07 0.00
4.00
Original
2.63 0.01
3.20
1.01
+0.80
0.00
LL
2,85 = 0.01
3.00
2.02
-0.90 0.00
4.57
CL
2.88 = 0.02
3.11 0.02
0.83 + 0.01
-1.07 + 0.00
-4.76 0.00
Deep Ensemble
2.71 0.01
4.03 0.07
0.77 1 0.01
-0.94 0.00
2.22
MCDO
GAP variants
LL
2.74
0.01
2,02 0.01
-0.87 0.00
-3.66 0.11
CL
2.01
4.03
0.90 0.02
-0,93 0.00
-3.80 + 0.07
Table 14: Last columns of transposed table
power plant
protein structure
wine
yacht hydrodynamics
Deep Ensemble MCDO LL
2.81 = 0.00
2,89 + 0.00
1.21 # 0.00
1.26 # 0.04
2.91 = 0.00
2.96 + 0.00
0,96 + 0.01
2.17 + 0.06
Original
2.91 =.00
2.91 # 0.00
1.24 + 0.01
1.20 # 0.04
CL
2
2.89 + 0.00
1.27
1.25 : 0.04
Deep Ensemble
2.00
+.000
1.48 1.01
1.71 + 0.03
MCDO
2.97
3.07 0.00
1.000
3.06 0.02
GAP variants
L.I.
2.00
3.07 0.00
1.05
1.78 0.02
CL.
2.91 # 0.00
3.13
1.43 + 0.00
1.74 0.01
29

---

Published as a conference paper at ICLR 2023
L
EIGENSPECTRUMS OF EMPIRICAL CTK AND NTK
Eigenspectrum CTK for CIFAR-
Eigenspectrum CTK for CIFAR-100
0.0035
0.005
0.0030
0.004
0.0025
0.0015
0.002
0.0010
0.001
0.000
0.0000
(a) CTK (CIFAR-10) 10)
(b) CTK (CIFAR-100
Eigenspectrum of NTK for CIFAR-10
Eigenspectrum of NTK for CIFAR-100
1.5
2.0
1.2
1.0
0.6
1.0
0.2
0.5
0.0
0.2
0.4
0.6
0.0
1.0
1.2
1e6
100000 200000 300000 400000 500000 600900
(c) NTK (CIFAR-100)
(d) NTK (CIFAR-100)
Figure 3: Eigenspectrums of CTK and NTK for CIFAR datasets
In this section, we follow Algorithm of Ghorbani al. (2019) to visualize the eigenvalue densities of empirical CTK and NTK. We use 100 Lanczos iterations (Appendix F) with realizations and fix the bandwidth of the RBF kernel as the difference between the maximum and minimum eigenvalue following the implementation of Ghorbani et al. (2019). We present the results in Figure 3.
Figure shows that empirical CTKs have positively skewed (i.e., right-tailed) eigenspectrums with modes close to O. In other words, many of the non-zero eigenvalues of CTK are close to zero. Therefore, the corresponding h(B)) for these eigenvalues are also close to zero, as shown Corollary 2.3 (See Fig. 2b for the visualization). As with empirical CTKs, the empirical NTKs are positively skewed, but their eigenvalue scales are much larger than CTKs. summary, our empirical study demonstrates that although empirical CTK can have up to NK non-zero eigenvalues, only few eigenvalues critical to the scale of E!" h(3)).
30


Reviews:
[{"review_id": "yw-O3ycYa8", "reviewer": "Reviewer_HGSe", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "This paper tackles the characterisation of the local curvature of NN loss landscapes in the face of scale-symmetries that appear in these models’ loss landscapes. Specifically, the authors build upon recent work in the linearised Laplace method for uncertainty estimation with neural networks by extending it to account for scale invariance in the NN parameter space. \n\nThe authors first use their construction, dubbed the connectivity tangent Kernel (CTK), to build a PAC-Bayes bound with a data-dependent prior that is invariant to the scale of the NN weights. The authors show this bound to be non-vacuous empirically on CIFAR10 and CIFAR100 with ResNet-18. The authors then propose the trace of the  CTK as an easy-to-compute surrogate quantity for estimating models’ generalisation. They show this quantity to correlate more strongly with test performance than previously proposed quantities, such as the trace of the Fisher information matrix. \n\nFinally, the authors use their connectivity tangent kernel for uncertainty estimation within the linearised Laplace framework. They apply their approach to UCI regression datasets and to CIFAR10+100 with ResNet18. The proposed method provides a modest but significant performance boost on top of the standard linearised Laplace approach. ", "summary_of_the_review": "\nI think this is a technically strong paper which I would like to see accepted. The paper has some major issues regarding clarity but I am confident that the authors can address these in a revision. When they do, I will be happy to raise my score. ", "strength_and_weaknesses": "**Strengths**\n\n* This paper makes 2 technically strong contributions, each of which could almost be a paper on its own.\n\t* A new non-vacuous PAC-based bound for a linear model built on a basis chosen to be the Jacobian of a pre-trained NN.\n\t* A formulation of the above linear model that avoids pathologies related to scale invariance in neural networks \n\n* The experimental validation is also quite strong. I was impressed that the authors evaluated both their generalisation bound and uncertainty estimation methods using ResNet18 (an 11M parameter model) on the CIFAR100 dataset. \n\n\n**Weaknesses**\n\n* Motivation / Problem exposition. The scale-variance issues this paper addresses are quite subtle and are not explained in enough detail. I was lucky to be familiar with these issues as I have dealt with them in my own work. However, I think that the average reader will be quite confused about what problem the authors are solving. I would recommend including a specific example of the Jacobians of a NN changing when the scale of the weights changes, despite the NN output not changing at all. \n\n* The authors evaluate their bound and show it gives very similar results for different scalings of the NN parameters. However, no other non-invariant bounds are shown. I think that it would be illustrative for the reader to see how the tightness of alternative bounds is lost due to the invariance pathologies addressed by this paper.\n\n* A potential issue with the proposed approach is that the re-scaling of the Jacobian basis that is proposed depends on the linearisation point. This vector contains information about the (potentially arbitrary) scaling of the dimensions of the NN Jacobian and information about how to fit the data well with the NN. It is not clear that the latter will be helpful to scale the Jacobian basis and it could in fact introduce some bias into the linearised model. \n\n\n* I found the experimental details provided lacking, even after reading the appendix. In this regard, I found the ResNet-18 CIFAR100 experiments particularly surprising. The dimensionality of the Jacobian basis expansion for a single input point is NOutputs x Nweights and evaluating this expansion requires as many backward passes as output dimensions. Additionally, a single step of Lanczos iteration requires a pass through the whole dataset and asymptotically the method requires as many steps as the rank of the kernel matrix being decomposed. How did the authors deal with these challenges?\n\n* I would suggest the authors repeat the evaluations from tables 1 and 2 for multiple seeds and add errorbars. \n\n* I found the clarity of writing to be poor, partly due to the text being full of grammatical mistakes.  I provide more details below. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "btXO9o0QRT", "reviewer": "Reviewer_xWGj", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes scaling-invariant prior and posterior distributions for PAC Bayes generalization bounds. It supports these empirically and, based on these, introduces a practical uncertainty estimation method that compares competitively with the linearized Laplace approximation.", "summary_of_the_review": "An in principle interesting paper with glaring weaknesses in its presentation, which all things considered make me lean towards a rejection.", "strength_and_weaknesses": "Strengths:\n* The range of contributions is quite broad, from a PAC Bayes bound to a practical approximate inference method.\n* The paper is overall well-structured.\n* The bounds seem to be empirically reasonably tight.\n* The uncertainty estimation method performs competitively.\n\nWeaknesses:\n* I found it quite difficult to follow what the paper is doing in detail and would not be confident that I would be able to implement the paper just based on the manuscript. In particular, I don't really understand which dataset split is used for calculating which quantities.\n* There are a lot of grammatical mistakes in the manuscript (as far as I can tell as a non-native English speaker), to the extent that it made the paper difficult to read for me (I'd usually just mention this as a minor comment, but in this instance I find that it is to a degree that it negatively impacts the already lacking clarity of the paper).\n* There is no discussion of computational cost. It seems like the method needs to evaluate a lot (100 Lanczos iteration; inside of an optimization problem) of Jacobian-vector products, which is not cheap (especially in case each product requires a pass over the training set, although I'm not sure if this is the case).\n\n**************************************************\nPOST REBUTTAL UPDATE\n\nThe authors have extensively updated the paper, addressing the concerns around presentation. I therefore increase my score.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "6KRqghz_Nch", "reviewer": "Reviewer_Q1Nm", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This work centers around the relationship between sharpness (scaling invariance) and generalization in DNNs under the PAC-Bayesian framework. This work decomposes the scale and connectivity parameters, derive the related connectivity tangent kernel (CTK), and builds the PAC-Bayesian generalization bounds. The trace of CTK can be regarded as a good indicator to show the correlation between sharpness and generalization.\n", "summary_of_the_review": "This work introduces new ideas/insights to analyze the relationship between sharpness and generalization bounds, but leaves some issues that I am concerned about. I'm hesitant to give a 5 or 6, and would like to see the authors’ rebuttal.", "strength_and_weaknesses": "**Pros:** \n\n1. Propose an (empirical) connectivity tagenet kernel (CTK) to demonstrate the sharpness and scaling invariance, which can be regarded an weighted NTK by $\\theta^*$ and Jacobian wr.t. $c$\n2. Build the PAC-Bayesian bound via CTK for generalization and connect it to invariance for flatness\n\n\n**Cons:** \n\nThere are several points on the motivation of Laplace approximation and the extension to over-parameterization that require more refinement. In my view, this work is built on a certain type of Taylor expansion, obtains the corresponding NTK matrix, assumes the posterior distribution to be Gaussian via Laplace approximation, and derives the PAC-Bayesian bound for under-parameterized deep neural networks. The detailed comments are as below.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "2wTLseC1FVw", "reviewer": "Reviewer_RRtQ", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors construct a PAC-Bayes bound that is insensitive to the fact that some parameters of neural nets can be arbitrarily rescaled due to normalization layers rectifying such rescalings. In a second step, they design a variant of the Laplacian approximation and evaluate both on a series of experiments.\n", "summary_of_the_review": "A paper with some potential whose main problems are currently in its presentation. ", "strength_and_weaknesses": "## Strengths\nThe proposed PAC-Bayes bound stays stable and independent wrt rescalings of parameter rescalings compared to the baseline bound. The authors are encouraged to move Table 7 to the main paper as it clearly demonstrates the improvement provided by their approach. Keeping it in the appendix essentially hides this.\n\nA major weakness of the draft is its written structure. There are a lot of grammatical mistakes (primarily missing articles), lots of typos, and partially repeated sentences,... Thorough proofreading is very necessary for any revision.\n\n\n### Questions\n- In the abstract, the authors discuss the flatness hypothesis as being widely accepted, while the first two paragraphs of the introduction then consist of summarizing an ongoing discussion in the literature. Can the authors comment on this?\n- Can the authors explain what they mean by the decomposition of scale and connectivity? The connectivity is modified in the Lee et al. paper you cite c is a binary mask, whereas for you it is a continuous vector. In which sense do you get a decomposition between scale and connectivity by the $\\theta(1+c)$ formulation? I seem to be misunderstanding a basic point of your setup. E.g., a c=0 via your prior changes nothing in any weight connectivity.\n-  Table 1: Why do only some rows contain standard deviations (or standard errors?) and over which number of repetitions are these computed? (Same question for any other table containing stds)\n- What is the relation between Table 4 and Figure 1? From my reading, they show the same experimental setting, but if so then the results in the table are misrepresentative as, e.g., for ECE Figure 1 clearly shows lower values for LL than claimed there. \n\n### Further comments\n- before eq (2) speaks about \"_our_ data-dependent prior\" vs _a_ data-dependent prior making it sound as if the split of training data to get a data-dependent prior was a contribution of the current paper.\n- $P_c$ appears in the line before (2) without any introduction and then disappears again for the remainder of the paper\n- Table 3's formatting makes it look as if CL were to improve also about deep ensembles as well as MC dropout. This needs a major restructuring! (E.g. in the direction of Table 4)", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%