Reparameterization invariance in approximate Bayesian inference
Authors: Hrittik Roy
Marco Miani
Carl Henrik Ek
Philipp Hennig
Marvin Pförtner
Lukas Tatzel
Søren Hauberg
Venue: NeurIPS 2024 spotlight (2024)
Keywords: Approximate Bayesian inference
Laplace approximations
Differential geometry
Reparametrization invariance
Primary area: probabilistic_methods
TL;DR: We provide a mathematical analysis of reparametrizations in the context of Laplace approximations, and devise novel a reparametrization invariant Riemannian diffusion-based approximate posterior.

Abstract: Current approximate posteriors in Bayesian neural networks (BNNs) exhibit a crucial limitation: they fail to maintain invariance under reparameterization, i.e. BNNs assign different posterior densities to different parametrizations of identical functions. This creates a fundamental flaw in the application of Bayesian principles as it breaks the correspondence between uncertainty over the parameters with uncertainty over the parametrized function. In this paper, we investigate this issue in the context of the increasingly popular linearized Laplace approximation. Specifically, it has been observed that linearized predictives alleviate the common underfitting problems of the Laplace approximation. We develop a new geometric view of reparametrizations from which we explain the success of linearization. Moreover, we demonstrate that these reparameterization invariance properties can be extended to the original neural network predictive using a Riemannian diffusion process giving a straightforward algorithm for approximate posterior sampling, which empirically improves posterior fit.

Reparameterization invariance in approximate Bayesian inference
Hrittik Roy Marco Mianit Technical University of Denmark (hroy, @dta dk
Carl Henrik Ek University of Cambridge, Karolinska Institutet che29@cam. ac. uk
Philipp Hennig, Marvin Pförtner, Lukas Tatzel University of Tübingen, Tübingen Al Center [philipp. hennig. lukas tatzel, marvin. pfoernerr Quni Suntirrninnntttngeeeeee
Søren Hauberg Technical University of Denmark sohau@dtu.dk
Abstract
Current approximate posteriors in Bayesian neural networks (BNNS) exhibit a crucial limitation: they fail maintain invariance under reparameterization, i.e. BNNS assign different posterior densities to different parametrizations of identical functions. This creates a fundamental flaw in the application of Bayesian principles as it breaks the correspondence between uncertainty over the parameters with uncertainty over the parametrized function. In this paper, we investigate this issue in the context of the increasingly popular linearized Laplace approximation. Specifically, has been observed that linearized predictives alleviate the common underfitting problems of the Laplace approximation. We develop new geometric view of reparametrizations from which we explain the success of linearization. Moreover, we demonstrate that these reparameterization invariance properties can extended to the original neural network predictive using a Riemannian diffusion process giving straightforward algorithm for approximate posterior sampling, which empirically improves posterior fit.
Introduction
Bayesian deep learning has not seen the same degree of success as deep learning in general. Theo- retically, Bayesian posteriors should be superior to point estimates (Devroye et al., 1996). but the practical benefits of having the posterior are all too often not significant enough justify their additional computational burden. This has raised the question if we even should attempt to estimate full posterions all network parameters (Sharma et al., 2023).
As an example, consider the Laplace approximation (MacKay, 1992), which places Gaussian in weight space through second-order Taylor expansion of the log-posterior. When applied to neural networks, this is known to significantly underfit and assign significant probability mass to functions that fail to fit the training data (Lawrence, 2001; Immer al., 2021a). Fig. 2 (top-left) exemplifies this failure mode for small regression problem. Interestingly. this behavior is rarely observed outside neural network models, and the failure appears linked Bayesian deep learning.
Recently, the linearized Laplace approximation (LLA) has been shown to significantly improve on Laplace's 's approximation through an additional linearization of the neural network (Immer al., 2021b; Khan et al., 2019). We are unaware of any theoretical justification for this rather counterintu- itive result: why would an additional degree of approximation improve the posterior fit
Equal contribution authors listed random order.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Puncteo space
Figure 1: The weight space of neu- ral network (Eq. 1) overparametrizes the associated function space. This induces families (orange) of weights corresponding the same functions. Model linearization (left) linearizes these families. In nonlinear models, Gaussian weight distributions (center) do not adapt to the families, while our geometric diffusion (right) captures the associated invariance with met- ric (gray ellipses).
Biraati Lonnd
leneet Bieettom
kenel divecden
Lerred directon
Nereel devectonn
Weight weightspoee space
Lisea sodel, Laplaces spoon.
Novlioeus model, Laphoe'' spgrox.
Noaliaa undel, Laplace dittuiio
Regeranetiizaton Funily
Metri:
Approsinate posteckor
We will show that the failures of Bayesian deep learning partly be explained by insufficient handling of reparameterizations network weights, while the LLA achieves infinitesimal invariance to reparameterizations. To motivate, consider the simple network (Fig. 1)
((x) wi ReLU (w2 2); R IR.
This can be reparametrized to form the same function realization from different weights f(z) wi ReLU (a W2 r) for any > 0. That is, the weight -airs (W1, w2) and (wn aww2) correspond the same function even if the weights are different (Fig. center).
Thus the approximate posterior cannot reflect the fundamental property of the true posterior, that should assign single unique density function regardless of its parametrization.
In this paper, we analyze the reparameteri- zation group driving deep learning and show that is a pseudo-Riemannian with the generalized Gauss-Newton (GGN) matrix as its pseudo-metric. We prove that the commonly observed underfitting of the Laplace approximation (Fig. 2, top left) caused by high in-distribution uncertainty in directions of reparametrizations (Fig. 2, top center).
distribution
contribution
Non-
Lincurized Laplacs Laplace
Lupiace dirfusion upgrosSimaiin
manifold,
Figure 2: The function space is decomposed into directions of reparameterizations (kernel) and functional change (non-kernel). We improve the posterior fit by concentrating probability mass on directions of functional change.
We develop reparametrization invariant dif- fusion posterior that proveably does not under- fit despite using the neural network predictive (Fig. 2, center row). Figure (right) visualizes how this posterior adapts to the geometry of reparamctrizations, thereby not underfitting. The diffusion can be simulated with a multi-step Euler- Maruyama scheme from which the linearized Laplace approximation (LLA) is single-step. This link implies that the LLA infinitesimally is invariant to reparameterizations, due to the otherwise coun- terintuitive linearization (Fig. 1, left). Experimentally, our diffusion consistently improves posterior fit, suggesting that reparameterizations should be given more attention in Bayesian deep learning.
Background: Laplace approximations
Let fw R' + R denote neural network with weights and define likelihood p(yl fw (x)) and a prior p(w). Laplace's 's approximation (MacKay, 1992) performs second-order Taylor ex- pansion of the log-posterior around a mode w. This results in Gaussian approximate posterior N(w w, - H-1) ) where H. is the Hessian matrix. The linearized Laplace approximation (Im- mer et al., 2021b; Khan et al., 2019) further linearize fw chosen weight w, i.e. Jw (x) Rs fi (w,x) x)) = (x) + Jw (x)(w w). where Jw Jw)x Öw (ww)w(x) w=w RR. x is the Jacobian of fw. Here D dim(w) denotes the number of parameters in the network. Applying the usual Laplace approximation to the linearized model yields an approximate posterior (Immer al., 2021b),
q( (W))) = - (ww w. (GGN + -
GGN = Σ (X.) H(x,, )J (x,).
(2)
n=

---

where H(x) = -0 (xx) log logp(yl,f(xx) E ROxO is the Hessian of the log- ligelike and we have assumed a weight prior N a II). Note that it trivial to extend other prior covariances. This particular covariance is known as the generalized Gauss-Newton (GGN) Hessian approximation, which is commonly used Laplace approximations (Daxberger et al., 2021b).
To reduce the notational load we stack the per-datum Jacobians into Ju = [ (x1); (x1):....; J.. (xN) E RNOOD and similarly for the Hessians, and write the GGN matrix as GGN = J. H.J. w- For Gaussian likelihoods, the Hessian is identity matrix and can be disregarded, and for other
likelihoods simple expressions are generally available (Immer et al., 2021b).
Sampled and linearized Laplace. The Laplace approximation gives Gaussian distribution ((w(D) over the weight space with mean w and covariance Σ. A predictive distribution obtained by integrating the approximate posterior against the model likelihood,
ply 1x = (owwolp(y"II(eee x
[[y""[(w;,x'))) W;
We refer this predictive method as sampled Laplace. Recent works have suggested linearizing the neural network in the likelihood model to obtain the predictive distribution (Immer et al., 2021b),
pxy" D) = (olk /////W:W.. (w, x x))
ply ply"ISm(m (w,,x")) W, ~ q.
(4)
This is referred to as linearised Laplace. Immer et al. (2021b) argues that the common choice of approximating the posterior precision with the GGN implicitly lincarizes the neural network and hence the predictive distribution should be modified for consistency.
Sampled Laplace is known severely underfit, whereas the linearized Laplace approximation does not (Immer et al. 2021b; Fig. 2). It is an open problem why the crude linearization is beneficial (Papamarkou et al., 2024). This paper shows that the benefit is linked the lack of reparameterization invariance.
The lack of reparameterization invariance leads to an additional problem for Laplace approx- imations. The precision of the approximate posterior is given either by the Hessian or the GGN, As shown by Dinh et al. (2017), the Hessian of the loss is not invariant reparameterizations of the neural network, and the same holds for the GGN. Depending on which parametrization of the posterior mode is chosen by the optimizer, we, thus, get different covariances for the approximate posterior. Empirically, this can render Laplace's approximation unstable (Warburg et al., 2023). Figure (center) illustrates the phenomena.
3 Reparameterizations of linear functions
Deep learning models excel when they are highly overparametrized, i.c. when they have significantly more parameters than observations (D >> NO). This introduces many degrees of freedom to the model, which will be reflected in the Bayesian posterior. However, approximate Bayesian inference does not correctly capture this and assigns different probability measures to identical functions. Next, we characterize these degrees of freedom to design suitable approximate posteriors. To develop the theory, we first consider the linear setting and then extend to the general case.
have argued, traditional
The reparameterizations of linear functions can be characterized exactly. Consider f(w) Aw + b and possible reparameterization. 9 RD RR," of this function such that f(g(w) = f(w). It is then evident that A(g(w) w) = 0. This implics that for any reparameterization of linear function, we have g(w) w ker(A), where ker(A) denotes the kernel (nullspace) of A. Hence, the linear function cannot be reparametrized we restrict ourselves to the non-kemel subspace of the input space if A has trivial kernel.
A linearized neural network SW' ww fw (x) xx= + Jw (x) w') is linear function in the parameters, where we have linearized around w'. The above analysis then implies that the kernel of the stacked Jacobian J Jw characterizes the reparameterizations of the lincarized network.

---

We can also characterize the reparameterizations through the GGN and the corresponding neural tangent kernel (NTK; Jacot et al. 2018),
GGNw w
NTK J.
(5)
By construction, these have the same non-zero eigenvalues, and thereby also have identical ranks. We, thus, see that the kernel of the Jacobian coincides with that of the GGN, i.e. ker(Jw) = ker ker(GGN
Two orthogonal subspaces. For any self- adjoint operator (such as positive semi-definite matrices like the GGN), the image and the kernel orthogonally span the whole space, i.e.
lorigan
keriGrN))
im(GGNw) ker(GGNw) RD,
(6)
larigun))
where the kernel is the hyperplane of vectors that are mapped zero and the image is the hyperplane of vectors spanned by the operator (Fig. 3). For a linearized neural network, im(GGNN) spans the effective parameters P IR. i.e. the maximal set of parameters that generate different linear functions R + Ro when evaluated on the training set.
Figure 3: The weight space can be decomposed into directions of reparameterizattios and func- tional changes. For linear models (left) these are linear subspaces given by the kernel and the image, respectively. For nonlinear models, these are the nonlinear manifolds P and Pw,, respectively.
A Laplace covariance decomposes into the same subspaces. Recall that the posterior precision is E GGN + al. Let the cigendecomposition of GGN be U AU, and assume that U, and U2 are the eigenvectors corresponding the non-zero eigenvalues Å, and the zero eigenvalues respec- tively. These form a basis in the kernel and image subspace as discussed above. Then the covariance is,
E
U1 U2
T
A
0 0
U1 U2
+ al
= UA - + UUU (7)
Consequently, we can decompose any sample from the Gaussian N(w, Σ) into a kernel and an image contribution, the kernel of GGN and Wim in the image. Note that all probability mass in ker(GGN&) is due to the prior, i.e. place prior probability on functional reparameterizations even if never observe data in support of such.
w + Wker Wins where Wser is the component of the sample that is in
Underfitting in sampled Laplace now be understood. For the linearized approximation, it holds for training data x E X that,
Train Accuney
0.7
0.6
flm W Wker + Wim x) = f (w++Win x).
0.4
Leplace Dillaizs
Easpic
Hence, the linearized predictive only samples in the image subspace consisting of unique func- tions. This is not true for the sampled Laplace approximation, which also samples in the ker- nel subspace. Since sampled Laplace does not lincarize the neural network, the kernel does not correspond reparameterizations. It hence adds "incorrect" degrees of freedom to the posterior as artifacts the Gaussian approximation.
Map acouaey
Rank
330
370
225
275
Figure 4: Underfitting of sampled Laplace is less pronounced when the rank of the GGN is higher for a fixed number of parameters. This is consistent with our hypothesis as a high GGN rank implies lower dimensional kernel. For experimental details, see appendix E.2.
Empirically, sampled Laplace only observed to underfit in overparametrized models. Fig. illustrates this by increasing the amount of training data to decrease the kernel rank, i.e. reduce the reparametrization issue. We find that as the issue lessened, sampled Laplace reduces its underfitting.
4 Reparameterizations of neural networks
We have seen that the parameters of linear models can be decomposed into two linear subspaces corresponding to reparameterizations and functional changes. We next analyze nonlinear models.

---

Intuitively. reparameterizations of nonlinear neural network form continuous trajectories in the parameter space (c.f. Fig. 1). We define that all points along such trajectory are identical, which changes the weight space geometry to be manifold. Likewise, the parameter changes corresponding to actual function changes reside on a nonlinear manifold. This sketched in Fig. 3. Interestingly, the GGN turns out to induce natural (local) inner product on these nonlinear manifolds, which allows us to both understand and generalize the linearized Laplace approximation.
4.1 The effective-parameters quotient space
For nonlinear neural network TR. TR! TR_ the surfaces in weight space along which the function does not change are generally not linear. Here, we formalize these reparameterization invariant surfaces and show that they are a partition of the weight space.
Definition 4.1. Given datapoint x E RR,, for any RRD define the x-reparameterizatton.
as the set R (w) = {w' such that f(w', x) = ((w.x)} Consistently, given collection of points X' C TR, we call the intersection (w( = NxEx (<(w) X-repurameterization.
Trivially, w E (w) for any choice of X. We next define the subset of X-reparamettrriattons which can be obtained via smooth deformation from w.
Definition 4.2. We say that piecewise differentiable function 7 [0, R. homotopy of (w, w) if (0)) = w and 7(1) = w'. The set of X' (t-smoott-eparrmeettzzaain is defined as,
R)((W)= w' such that By a homotopy of (w, w') 7(1) Rxx(() VI [0,
A homotopy is, thus, smooth path along which all neural networks have identical predictions on X. We consider networks, w and w', similar they can be connected by such homotopy. Formally, we define the relation over RDD as wow' if w' E R
We next use this relation form new view on the weight space RD in which similar weights seen as one point. This can be realized using quolient spaces (Lee, 2012). These are well-studied spaces that constructed by considering a collection of points in space as single point in new space. In our case, we have the following result.
Lemma 4.3. ~ is an equivalence relation, iee. it is transitive, symmetric and reflexive. We can form the quofient space P RDD/ of effective parameters. We denote [w] E P the equivalence class of an element w E R.D
This quotient structure gives a rich mathematical foundation construct reparameterization invariant neural networks. Within the quotient, two effective parameters [w1], [w2] E P are the and only w ~ W2. This means that all parameters w [wi] gives the same function over X.
same point if
4.2 The effective-parameters manifold
Geometry is the mathematical language of invariances. To this end would like endow the weight space with geometric structure such that two weights, wT and w2, corresponding to the same function, have distance of zero, i.c.
dist(w W2) =
W1 W2-
8
Since the weights generate the same function, define a metric that measures differences in function values on the training data. Consider weights w and an infinitesimal displacement we then define,
dist? 2 fist?(w, +)) Σ [lf(w, f(w e, x11 = - TGGNN + +O(e),
where the last step follows from first-order Taylor expansion of around w. This is a standard pullback metric (f*H)w = GGNw commonly used in Riemannian geometry. This implies that the GGN matrix infinitesimally defines an inner product, i.e. a Riemannian metric. By integrating over paths, the distance extend to any pair of points and satisfies Eq. 8 (Lee, 2012).
5

---

Watch out! It's pseudo-metric. We have already seen that in overparametrrized models, the GGN is rank-deficient, which implies that it not positive definite. Consequently, it is not Riemannian metric but rather -veudo-Riemannian metric. A pseudo-metric can be a counterintuitive object: two points W1 and W2 at distance zero may have different pseudo-metries (f* H),, * (f*H)w This is reflected in the Laplace approximation. The covariance prescribed by the Laplace approximation is En = (W. [(w) + al) -1, where [(w) is shorthand for the training log- log-likelihood. The Hessian is exactly the pullback pseudo-metre (ww) reparameterizations the neural network. Specifically, for a reparameterization function g that is also diffeomorphism, the change of variable rules states that,
(f w. which is not invariant to
V W ((o(W)) Vwg(WW) (W)T
((w) Vwg(w).
(10)
This means that, while each parameter w has its well-defined covariance Ew. each equivalence class does not have a unique one, since w and g(w) belong to the same equivalence class and Ew # Σg(w)-
Non-Gaussian likelihoods. The Euclidean distance measure in Eq. corresponds to choosing a Gaussian likelihood. The distance definition readily extends to other likelihoods and the corresponding metric takes the form of the generalized Gauss-Newtor matrix J w. H.J w- where denotes the Hessian of the log-likelihood. For both Gaussian and Bernoulli likelihoods, this Hessian is positive definite, but e.g. the cross entropy has a rank-defi Hessian and, thus, induces a pseudo-metric.
An impractical solution. The unfortunate behavior of approximate posteriors assigning different probabilities the same function could be rectified by marginalizing over the set of reparameteriza- tions of w, i.e. (R) q q(w'lD)dw' While this construction solves the highlighted problem, its
complexity makes impractical and we unaware of any works along these lines.
When restricted integral can be thought of as "collapsing" each reparameterization equivalence class to a single point in P = D ~ formalized in Lemma 4.3. Nontrivially, the pullback metric implicitly performs a similar operation, as shown later in Theorem 4.5. This connection motivates the dive into Riemannian geometry: we get a tructable approach to engaging with neurul network reparumeterizations
smaller class of reparameterization (the ones homotopic to the identity), the
4.3 Topological equivalence of the two views
So far we described two a priori very different objects: the quotient space P = RD/ ~ and the pseudo-Riemannian manifold (RD GG We referred to both of them as effective parameters and this is no coincidence as there is a natural relationship between the points at distance zero according the pseudo-metric and the equivalence classes.
Proposition 4.4. For any Wo. W1 E R it holds
d (Wo, w1) =0 - Wo
[w]] E.P
(11)
Even better, these two spaces share the same topological structure. To state this we need a notion of distance the quotient space and the most natural choice is to inherit the Euclidean distance from RR D This distance is defined as
dr([w), [w']) inf (lip) qill +
+ llp,. gn 1m113.
where the infimum is taken over all finite sequences ,I....,,, Pr and ......, such that [w] = [pr]. [pt+1] [q:] and [qx] [w]].
This distance de induces a topology on the quotient space P which is equivalent to the topology
induced by the pullback distance d и on the beeudo-Riemaniian manifold. Formally
Theorem 4.5. For any wo. W1 P. for any > 0 there exists 8 0 such that
dr d[w[)) [w]]) < 8
d ( Wo, wi) <
(12)
dy- и (Wo, w1) < -
dæ (w1]) < e.
(13)
This result connects an abstract quotient space P with the quotient captures useful intuitions but is difficult to leverage computationally. In contrast, the pseudo- metric has some counterintuitive aspects but we can identify the underlying Riemannian structure which leads to tractable algorithms (Sec. 5).
metric GGNw. The

---

A tale of two manifolds. For any given parameter w E RD" and training set X, we show that there exist two Riemannian manifolds (Pw. m) and (Pw, , m embedded in R.,, illustrated in Fig. 3. They capture the functional change and reparameterizatioon properties respectively, but, differently from the previously studied (RD, GGNN), they are Riemannian manifolds without degenerate directions in their metrics. Formally.
Theorem 4.6. For any parameter w suppose the set of parameters that generate the same predictions
denoted by Pd = (w' TR. such that f(w',z) = f(w, 3)) for all E X). Then this set is smooth manifold embedded in IR.". Furthermore, the set of parameters that locally generates unique predictions, Pw is also a submanifold embedded in R.D.
They are the direct generalization to the nonlinear case of the two spaces involved in Eq. 6, where Pw plays the role of the image and Pw plays the role of the kernel. When f is linear, they are identical.
In general, Pw and P# intersect only in w, and the two respective tangent spaces in w span all directions. They can be thought of as two collections of parameters, and the associated functions have different properties: (1) Pw is entirely contained in the same equivalence class P S [w]. thus all the parametrized functions are identical on the train set; in contrast, (2) Pw never intersects the same equivalence class more than one time, at least locally, thus the parametrized functions always changes when moving in any direction. Thus Pw resembles the effective parameter manifold P, but with the difference of being an actual Riemannian manifold. These two manifolds exist under the assumption that Jacobian is full rank (see proof in appendix C).
The metries and m the tangent space of Pw corresponding to the non-zero eigenvectors, i.e. m = GGN While, for m we can inherit the Euclidean metric, i.c, m = al for a > O.
uniquely defined. A natural choice for is to restrict GGN to
5 Exploring manifolds with random walks
SDEs on manifolds. Given Riemannian manifold (M, G), the simplest choice of distribution that respects the Riemannian metric G is a Riemannian diffusion (or Brownian motion, c.f. Hsu (2002))
stopped at time t. This follows the stochastic differential equation (Girolami & Calderhead, 2011),
D a dw = V2FG(w) Md + +Iid where P((w) = Σ (GG(W) Ow;
(14)
Practically speaking this simple process can be simulated using an Muler-Maruyama (Maruyama, 1955) scheme. The Christoffel symbols, (O),), are commonly disregarded as they have a high computational cost, and Li et al. (2015) showed that the resulting error is bounded.
Using the Marurama integrator with step size hes setting T = corresponding to standard Bayesian inference and disregarding the term involving the Christoffel symbols I,, we obtain the simple update rule W w, + Gw.,) &c. where Riemannian manifold. However, the effective-parameter (R GGN is only pseudo-Riemannial we explore three Riemannian alternatives: (IR", GGNw + all,, (P al) and (Pw: GGN
I). This applies to any
Diffusion on (IR GGN, + al). The Laplace approximation can also be written as diffusion on manifold. As we saw in Sec. 2, the Laplace approximation can be written as w/D NNw,, Σ) with x-a G + aI. This can also be written as a sample at t = 1 of a Riemannian diffusion on manifold with a constant metric, (RD, G). where G = GGN + al. The SDE dw G-1dW have a marginal distribution at f = that exactly match the standard Laplace approximation. Note that this formulation does not rely on the approximation of the SDE that disregards the term involving the Christoffel symbols I' as these are zero for constant metrics. Hence, the above is exactly a Riemannian diffusion on the manifold with constant metric given by the GGN at the MAP parameter. Note that this is only a valid diffusion for a > O in which case is not reparametrization invariant.
Kernel-manifold diffusion. The kernel-manifold (P. aI) consists of parameters that generate the same function over the training set. The effect of diffusion on this manifold and using the neural network predictive is similar to sampling from the kernel subspace while using the linearized predictive. On the training set the predictive variance because only samples reparametrizations of the MAP predictions. On cattof-distribution data, the variance is greater than 0 if at least one of
7

---

Table 1: In-distribution performance across methods trained on MNIST, FMNIST and CIFAR-10.
Conf. (1)
NLL (1)
Acc. (1)
Brier (L)
BCE (L)
MCE (4)
CIFAR 10 FMNIST MNIST
Laplace Diffusion (ours)
0,988
0.042a0.007
0.98740.002
0.02240,003
0.137=0.019
0.77520.043
Sampled Laplace
0.589=0.008
3.812a0.284
0. 0.14620.032
1.17610.046
0.44320.026
0,98510.002
Linearised Laplace
0.968=0.004
0.30620.041
0.926x0.008
0.117+0.012
0.251=0.034
0.855=0.041
Laplace Diffusion (ours)
0.900=0.001
0.001=0.000
0.906#0.007
0.141#0.006
0.108:0.015
0.729=0.092
Sampled Laplace
0.618=0,021
4.507+0.000
0.098+00010
1.295+0.014
0:0.01
0,986+0.001
Linearised Laplace
0.897=0.003
0.423+0.000
0.86240.005
0.207+0.006
0. 147=0.017
0.756=0.048
Laplace Diffusion (ours)
0.952=0.007
0.345=0,062
0.905+0.007
0.155+0.019
0.259=0.00 008
0.870=0.021
Sampled Laplace
0.843=0.004
0,997+0.222
0.717=0.049
0,422+0.081
0.80420.080 0.0244
Linearised Laplace
0.951=0.007
4.0.0200020
0.863+0.001
0.222+0.002
0.337=0.022
0.789=0.035
Table 2: Out-of-distribution AUROC (1) performance for MNIST, FMNIST and CIFAR-10.
Trained on
MNIST
FMNIST
CIRRR 10
Tested on
FMNIST
EMNIST
KMNIST
MNIST
EMNIST
KMNIST
CIEAR 100
SVHN
Laplace Diffusion (ouxs)
0.983900.033
0.625=00.088
0.929:80008
0.759HI.045
0.741=8.010
0.749H0.023
0.851=00002
0.862#1.000
Sampled Laplace
0.500
0,49420.006
0.482s00033
0,495a0.037
0.503200066
0,49340.033
0,687.0003
0.599
Lincurised Laplace
0.758a00070
0.602+0.027
0.790=00188
0.62540.050
0.62820013
0.62440.020
0.83760006
0,85420.024
the reparameterizations on the training set not a global reparameterization. This leads to clear separation in the predictive variance of in-distribution and out-of-distributiin data (Fig. 2) and further implies that this diffusion distribution never underfits. Stated formally,
Theorem 5.1. Varw [f(w, x)] = 0 for train data x E X. For test point x2 & X, if there exists reparameterization w' E R (w) such that w' & R ( then Var [f(w, > O.
tatn-kenel-parameer manifold diffusion. The por-keeeel-arrameer manifold (Pw. GGN consists of parameters that generate unique functions over the training set. Diffusion on this manifold samples functions that are necessarily different from the MAP predictions on the training set. However, the predictive variance in the training set is bounded such that the functional diversity in the predictive samples reflects the intrinsic variance of the training data (Fig. 2).
This the only considered diffusion that acts on Riemannian manifold while being reparametrization invariant, i.c. (w(( = {w]. We call this Laplace diffusion and study empirically in Sec.
6 Related work
Bayesian deep learning techniques are still in their infancy and generally involve poorly under- stood approximations. The arguably most popular tool for uncertainty quantification is ensembles (Lakshminarayanan al., 2017; Hansen & Salamon, 1990). Several approaches make Gaussian approximations to the true posterior, including 'Bayes by backprop' (Blundell et al., 2015), stochastic weight averaging (SWAG) (Maddox al., 2019) and the Laplace approximation (MacKay, 1992; Daxberger et al., 2022; Antorán et al., 2023; Deng et al., 2022; Miani et al., 2022).
The high dimensionality of the weight space gives rise to significant computational challenges when constructing Bayesian approximations. This has motivated various low-rank approximations (review in Daxberger et al., 2021a), c.g. last layer approximations (Kristiadi et al., 2020), subnetwork inference (Daxberger et al., 2021c), subspace inference (Izmailov et al., 2020) or even PCA in weight space (Maddox et al., 2019). Such approaches lessen the computational load, while often improving predictive performance. Our analysis sheds light on why crude approximations perform favorably: smaller models are less affected by reparameterization issues. Our diffusion process, thus, provides an altermative, and less heuristic, path forward.
MacKay (1 (1998) noted the importance of the choice of basis in Laplace approximations; our pseudo- Riemannian view can be seen as having a continuously changing basis. Kristiadi et al. (2023) studied how metric transforms under bijective differrentaable change of variables. They enforce geometric
8

---

Laplace Diffusion
Sampled Laplace
Lineriied Leglaos
Figure 5: Benchmark results for Ro- tated MNIST (similar results for FM- NIST and CIFAR are in appendix E.3.2). Sampled Laplace significantly under- fits even for non-rotated data, Laplace diffusion consistently outperforms the other methods.
3.0
2.0
9.4
1.0
0.0
NLL.
EGR
Acc.
50
100 130
100 150 Shift intensity
50
Indistrrbtto a2trks
Shift intensity
consisteney, highlighting, e.g., the non-invariance of the GGN to change of variables. Petzka et al. (2019); Jang et al. (2022) point to the same inconsistency with an emphasis on flatness measures.
Kim et al. (2022) and Antorán et al. (2022) study global (rather than data-dependant) reparametriza- tions associated with specialized architectures. While analytic expressions can be obtained, the results do not apply to general networks. While not expressed in terms of reparametrizations. Izmailov et al. (2021) show that linearly dependent datasets give rise to hyperplane in the kernel manifold. Kim al. (2024) also study the kemnel of the GGN in the context of influence functions. These works characterize subsets the reparametrization group. We provide the first architecture-agnose characterization of all continuous reparametrizations..
In a closely related work, Bergamin et al. (2024) introduced a Riemannian Laplace approximation (Hauberg, 2018) that improves posterior fit over range of tasks. Furthering this line of research, Yu et al. (2023) explored the use of the Fisher information metric within this framework. While sharing the language of Riemannian geometry, of linearized Laplace within the context of neural network reparametrization, instead of primarily aiming to achieve better posterior approximations. This allows to gain deeper insights into the underlying mechanisms that contribute to the success of this approximation technique.
work focuses on analyzing the effectiveness
Experiments
We benchmark Laplace diffusion with neural network predictive against linearized and sampled Laplace to validate the developed theory. Implementation details are in Appendix 1. We will show that the diffusion posterior slightly outperforms linearized Laplace terms of both in-distribution fit and out-of-distribution detection. For completeness, we include comparisons to other baselines such as SWAG, diagonal Laplace, and last-layer Laplace in Appendix E.3. Laplace diffusion is competitive with the best-performing Bayesian methods despite using the neural network predictive (i.e. no linearization). This contrasts sampled Laplace which severely underfits. This is evidence that the developed theory explains the key challenges of Bayesian deep learning.
Experimental details (appendix E.3). We train 44,000 10.00000prmmtteee LeNet(LeCun et al., 1989) on MNIST and FMNIST as well as 770.000-parameee ResNet(He et al., 2016) on CIFAR- 10(Krizhevsky et al., 2009). We sample from the Laplace approximation of the posterior and our Laplace diffusion. For the samples from the Laplace approximation, we consider both the linearized predictive and the neural network predictive, while for diffusion samples, we only consider the neural network predictive. These baselines were chosen to be as similar as possible our approach to case the comparison. We use the same prior precision for all methods to ensure a fair comparison.
In-distribution performance (Table 1). We measure the in-distribution performance of different posteriors on a held-out test set. We report means standard deviations of several metrics: Confi- dence, Accuracy. Negative Log-Likelihood, Brier Score (Brier, 1950), Expected Calibration Error (Naeini et al., 2015) and Mean Calibration Error. We observe that Laplace diffusion has the best calibration and fit. We also confirm the underfitting of Sampled Laplace across cases. For CIFAR- 10 we had to use a large prior precision to get meaningful samples from sampled Laplace, which explains the less severe underfitting. High prior precision is known to help with underfitting in sampled Laplace, but it also shrinks predictive uncertainty to almost zero.
Robustness to dataset shift (Fig. 5, appendix E.3.2). We use ROTATED-MNIST, ROTATED- FMNIST, and ROTATED-CIFAR asses model-calibration and model fit under distribution shift. Fig. 5 plots negative log- ligelloood (NLL) and expected calibration error (ECE) against the degrees of rotation. Laplace diffusion improves on other Laplace approximations.
Code: dttps://githhn www..coooocce.ccece http://ooooggommetrc--aapappp.aa

---

detection (Table 2). On out- -f--riibtiiittionn data from other benchmarks, we see that Laplace diffusion outperforms the other Laplace approximations.
8 Conclusion
While approximate Bayesian inference excels in many areas, it continues to face challenges in deep learning. Techniques that work well in shallow models struggle with deep ones even they remain computationally tractable. This suggests that overparamettriztioon plays negative role in Bayesian models. Our theoretical analysis shows how overparametrization creates a growing reparameterization issue that conflicts with standard Euclidean approximate posteriors. such as the ever-present Gaussian. For small models this issue is negligible, but as models grow, so does the reparamcterization issue.
Our geometric analysis also suggests a solution: we should consider approximate posteriors that respect the group structure of the reparameterizations. We observe that the generalized Gauss-Newton (GGN) matrix commonly used in Laplace approximations induces Aseudo-Riemannie structure on the parameter space that respects the topology of the reparameterization group. This implies that we can use Riemannnn probability distributions as approximate posteriors, and experimented with the obvious choice of geometric diffusion process. We also showed that the state-of-the-an linearized Laplace approximation can be viewed as naïve (or simple) numerical approximation to our proposed diffusion. This helps explain the success of the linearized approximation.
Our proposed approximate posterior does have issues. While sampling has the same complexity as standard Laplace approximations, increases runtime by constant factor. Common Laplace approximations do not sample according to the GGN but rather approximate this matrix with diagonal or block-diagonal matrix. Mathematically, such approximations break the motivational reparameterization invariance, so it is unclear if such approaches should be applied in our framework. Our work, thus, raises the need for new computational pipelines for engaging with the GGN matrix.
Acknowledgments and Disclosure of Funding
This work was supported by research grant (42062) from VILLUM FONDEN. This project received funding from the European Research Council (ERC) under the European Union's 's Horizon 2020 research and innovation programme (grant agreements 757360 and 101 101123955). The work was partly funded by the Novo Nordisk Foundation through the Center for Basic Machine Learning Research in Life Science (NNF200C0062606), In addition to the ERC (above), PH. MP and LT thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) -SS) for support, and gratefully acknowledge financial support by the DFG Cluster of Excellence "Machine Learning New Perspectives for Science",, EXC 2064/1, project number 390727645; the German Federal Ministry of Education and Research (BMBF) through the Tübingen AI Center (FKZ: 011S18039A and funds from the Ministry of Science, Research and Arts of the State of also grateful to Magnus Waldemar Hoff Harder for alerting imprecisions in early draft of this manuscript.
The authors are
References
Allen-Zhu, Z., Li, Y., and Song, Z. A convergence theory for deep learning via over-parametterizaton.. In International conference machine learning, pp. 242-252 PMLR, 2019.
Antorán, J., Janz, D., Allingham, J. U., Daxberger, E.., Barbano, R. R., Nalisnick, E., and Hernández-
Lobato, J. M. Adapting the linearised laplace model evidence for modern deep learning. In International Conference Machine Learning, pp. 796-821. PMLR, 2022.
Antorán, J., Padhy, S., Barbano, R., Nalisnick, E., Janz, D., and Hernández-Lote J. M. Sampling- based inference for large linear models, with application to linearised laplace, 2023.
Bergamin, F., Moreno-Muñoz, P., Hauberg, S., and Arvanitidis, G. Riemannian laplace approxi- mations for bayesian neural networks. Advances in Neural Information Processing Systems, 36, 2024.
10

---

Blundell, C., Cornebise, J.. Kavukcuoglu, K., and Wierstra, D. Weight uncertainty in neural network. In International conference on machine learning, pp. 1613-1622. PMLR, 2015.
Bombari, S., Amani, M. H., and Mondelli, M. Memorization and optimization deep neural networks with minimum peee-parameterization. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
Brier, G. W. Verification of forecasts expressed in terms of probability, Monthly Weather Review, 78(1):1 - 3, 1950. doi: (01155/5324401110007770000000022 CO:2. URL https://journall ametsoc. xml.
htt.:aas.sattttoooooo//oeaaaaaaa ottp://o////ooomaaa///////////// 1960_077__001_.0feett22220
Daxberger, E., Kristiadi, A., Immer, A., Eschenhagen, R., Bauer, M., and Hennig, P. Laplace redux effortless Bayesian deep learning. In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances Neural Information Processing Systems, volume 34, pp. Curran Associates, Inc., 2021a.
Daxberger, E., Kristiadi, A., Immer, A., Eschenhagen, R., Bauer, M., and Hennig, Laplace redux-effortless Bayesian deep learning. In NeurIPS, 2021b.
Daxberger, E., Nalisnick, E., Allingham, J. U., Antorán, J., and Hernández-Looato, J. M. Bayesian deep learning via subnetwork inference. International Conference on Machine Learning, pp. 2510-2521. PMLR, 2021c
Daxherger, E., Kristiadi, A., Immer, A., Eschenhagen, R., Bauer, M., and Hennig, P. Laplace redux effortless bayesian deep learning, 2022.
Deng, Z., Zhou, F., and Zhu, J. Accelerated linearized laplace approximation for bayesian deep learning, 2022.
Devroye, L., Györfi, L., and Lugosi, G. A probubilistic theory of pattern recognition. Springer Science & Business Media, 1996.
Dinh, L., Pascanu, R., Bengio, S., and Bengio, Y. Sharp minima can generalize for deep nets. In International Conference on Machine Learning, pp. 1019-1028. PMLR, 2017.
Du, S., Lee, J., Li, H., Wang, L., and Zhai, X. Gradient descent finds global minima of deep neural networks. In International conference on machine learning, pp. 1675-1685. PMLR, 2019.
George, T., Laurent, C., Bouthillier, X., Ballas, N., and Vincent, Fast approximate natural gradient descent in kronecker factored eigenbasis. Advances in Neural Information Processing Systems, 31, 2018.
Girolami. M. and Calderhead, B. Riemann manifold langevin and hamiltonian monte carlo methods. Journal of the Royal Statistical Society Series Statistical Methodology, 73(2):123-214, 2011.
Hansen, L. K. and Salamon, Neural network ensembles. IEEE transactions on puttern analysis and machine intelligence, 12(10):993-100) 1990.
Hauberg. S. Directional statistics with the spherical normal distribution. In Proceedings of FUSION 2018, 2018
He, K., Zhang. X., Ren, S., and Sun, J. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Hsu, P. Stochastic analysis on manifolds. American Mathematical Soc., 2002.
Immer, A., Bauer, M., Fortuin, V., Rätsch, G., and Emtiyaz, K. M. Scalable marginal likelihood estimation for model selection in deep learning. In International Conference on Machine Learning (ICML), pp. 4563-4573, 202
Immer, A., Korzepa, M., and Bauer, M. Improving predictions of Bayesian neural nets via local linearization. In International Conference on Artificial Intelligence and Statistics (AISTATS), pp- 703-711, 2021b.
11

---

Izmailov, P., Maddox, W. J., Kirichenko, P., Garipov, T., Vetrov, D., and Wilson, A. G. Subspace inference for bayesian deep learning. In Uncertainty in Artificial Intelligence, pp. 1169-1179. PM 2020.
Izmailov, P., Nicholson, P., Lotfi, S., and Wilson, A. G. Dangers of bayesian model averaging under covariate shift. Advances in Neural Information Processing Systems, 34:3309-3322, 2021.
Jacot, A., Gabriel, F., and Hongler, C. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in neural information processing systems, 2018.
Jang, C., Lee, S., Park, F., and Noh, Y.-K. A repaaranetriaaiionnnnaaaaann sharpness measure based on information geometry. Advances in neural information processing systems, 35:27893-27905. 2022.
Khan, M. E. E., Immer, A., Abedi, E., and Korzepa, M. Approximate inference turns deep networks into Gaussian processes, Advances in Neural Information Processing Systems (NeurIPS), 32, 2019.
Kim, S., Park, S., Kim, K.-S,, and Yang, E. Scale-invariant bayesian neural networks with connectivity tangent kernel. In The Eleventh International Conference on Learning Representations, 2022.
Kim, S., Kim, K., and Yang, E. Gex: flexible method for approximating influence via geometric ensemble. Advances in Neural Information Processing Systems, 36, 2024.
Kristiadi, A., Hein, M., and Hennig, P. Being Bayesian, just bit, fixes overconfidence in ReLU networks. In III, H. D. and Singh, A. (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 5436-5446, PMLR, 13-18 Jul 2020.
Kristiadi, A., Dangel, F., and Hennig. P. The geometry of neural nets' parameter spaces under reparametrization. arXiv preprint arXiv:2302.07388 2023.
Krizhevsky, A., Hinton, G., et al. Learning multiple layers of features from tiny images. 2009.
Lakshminarayanan, B., Pritzel, A., and Blundell, C. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems, 30, 2017.
Lanczos, C. An iteration method for the solution of the eigenvalue problem of linear differential and integral operators. Journal of Research of the National Bureau of Standards, 45(4), October 1950. doi: 10.6028/jres.045.oo 02..00... URL https: /hal scienceeeee.cc wwa.nceehhaa http/hall-11a100
Lawrence, N. D. Variational inference in probabilistic models. PhD thesis, Citeseer, 2001.
LeCun, Y., Boser, B., Denker, J. S., Henderson, D., Howard, R. E., Hubbard, W., and Jackel, L. D.
Backpropagation applied to handwritten zip code recognition. Neural computation, 1(4):541-551, 1989.
Lee, J. M. Smooth manifolds. Springer, 2012.
Li, C., Chen, C., Carlson, D., and Carin, Preconditioned stochastic gradient langevin dynamics for deep neural networks, 2015,
Li, Z., Wang. T., and Arora, S. What happens after sgd reaches zero loss?-a mathematical framework. arXiv preprint arXiv:2110.06914 2021.
Lippe, P. UvA Deep Learning Tutorials. latest/, 2022.
io/en/
Liu, C., Zhu, L., and Belkin, M. On the linearity of large non-linear models: when and why the tangent kernel is constant. Advances in Neural Information Processing Systems, 33:15954-15964, 2020.
MacKay, D. J. Choice of basis for laplace approximation. Machine learning, 33:77-86, 1998.
MacKay, D. C. A practical Bayesian framework for backpropagation networks. Neural Computa- tion, 4(3):448-44 4422, 1992.
12

---

Maddox, W. J., Izmailov, P., Garipov, T., Vetrov, D. P., and Wilson, A. G. A simple baseline for bayesian uncertainty in deep learning. Advances in neural information processing systems, 32. 2019.
Maruyama, G. Continuous markov processes and stochastic equations. Rendiconti del Circolo Matematico di Palermo, 4:48-90, 1955.
Miani, M., Warburg, F., Moreoo-Muñoz, P., Detlefsen, N. S., and Hauberg, S. Laplacian autoencoders for learning stochastic representations. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
Naeini, M. P.. Cooper, G., and Hauskrecht, M. Obtaining well calibrated probabilities using bayesian binning. In Proceedings of the AAAI conference on artificial intelligence, volume 29, 2015.
Nguyen, Q., Mondelli, M., and Montufar, G. Tight bounds on the smallest eigenvalue of the neural tangent kernel for deep relu networks. In International Conference on Machine Learning, pp. 119-8129. PMLR, 2021.
Oymak, and Soltanolkotabi, M. Overparameterized nonlinear learning: Gradient descent takes the shortest path? In International Conference Machine Learning. pp. 4951-4960. PMLR, 2019.
Oymak, S. and Soltanolkotabi, M. Toward moderate overparametterization: Global convergence guarantees for training shallow neural networks. IEEE Journal on Selected Areas in Information Theory, 1(1):44-105, 2020.
Papamarkou, T., Skoularidou, M., Palla, K., Aitchison, L., Arbel, J., Dunson, D., Filippone, M., Fortuin, V., Hennig, P., Hubin, A., et al. Position paper: Bayesian deep learning in the age of large-scale ai. arXiv preprint 24X::2402000009, 2024.
Petzka, H., Adilova, L., Kamp, M., and Sminchisescu, C. eeparameterizaationinnaaae flatness 2019.
measure for deep neural networks. arXiv preprint
Ritter, H., Botev, A., and Barber, D. A scalable laplace approximation for neural networks. In 6th international conference on learning representations, ICLR 2018-conference track proceedings, volume 6. International Conference Representation Learning, 2018.
Sharma, M., Farquhar, S., Nalisnick, E., and Rainforth, T. Do Bayesian Neural Networks Need To Be Fully Stochastic? In International Conference on Artificial Intelligence and Statistics (AISTATS, notable paper award), 2023.
Warburg, F., Miani, M., Brack, S., and Hauberg, S. Bayesian metric learning for uncertainty quantification in image retrieval. In Neural Information Processing Systems (NeurIPS), 2023.
Yu, H., Hartmann, M., Williams, B., Girolami, M., and Klami, A. Riemannian laplace approximation with the fisher metric, 2023.
13

---

A Recap
Consider function ʃ IR R. Ro with Jacobian Jw(x) (x)) w-w RoxD with respect to w evaluated in x and w. For a given log-likelihood we define the Hessian w.r.t (x) log f(y)ff(x)) e RoxO and we assume not be dependent on the output Hx(x) -0
Notation.
y (which is true, for example, for exponential families).
Consider being given a dataset of finite size N, here we do not care about labels and we only refer to the collections of the datapoints X' {xi,.... XN} R!.
Consider the stacking of the per-datum Jacobians Jw = [J (x1)i.... Jw (xx) E ROXD where dropped the dependence on X. Similarly consider Hw diag(H, (x1) i.)i..... Hw (xN)) I NO NO the block diagonal stacking of the Hessians.
Consider the Generalized Gauss-Newton (GGN) and the Neural Tangent Kernel (NTK) matrices
GGN J HwJ. E RDD
NTK = H Jw H E R NOXNO
(15)
Recall also that the pullback pseudo-metts metric is (f*H)w GGNw- Assumptions.
that is We assume uniform upper and lower bound on the eigenvalues of the NTK matrix,
NTKwul ||l||
31 EER
such that
0 0<1S
L
VEERNO Vw VWERD,
(16)
where uniform means uniform over parameters, i.e. the bounds 1, T. holds for every w. Moreover we assume that the Jacobian function w >> Jw(x) is Lipschitz for every x E X.
How unreasonable are the assumptions? The assumption of an upper bound L is equivalent to assuming that w + f(w, x) is Lipschitz for each datapoint activation functions if we restrict the parameter space to a ball of fixed radius.
X. This true with all standard
The assumption of lower bound I is strongly supported by the literature NTK, thanks to its direct implications on memorization capacity and generalization. The general trend that the more overparametrized the network is. the stronger such lower bounds are. In the hardest setting of minimum overparametrization, Bombari et al. (2022) proved bound that holds with high probability for fully connected MLPs at initialization. Similar results hold with bigger overparametrizatioos Nguyen et al. (2021): Allen-Zhu et al. (2019); Du et al. (2019). Building on top of that, other lines of work Liu et al. (2020); Oymak & Soltanolkotabi (2019, 2020) proved that the too much during training thanks to the PL inequality framework, and in particular the proximity of the neural network dynamics the one described by NTK, is supported by spectral bounds on the Hessian of the landscape.
does change
Lastly. the Lipschitzness assumption the Jacobian is potentially the most unrealistic, although would hold, for example, if the derivatives activations are Lipschitz and the parameters are restricted to a finite radius ball. Nonetheless, control that the kernel of the GGN does not "rotate too fast", which is a much weaker assumption, but also much more cluttered to state formally.
emphasize that this assumption is only used to
B Proof for equivalence of the two settings
This section contains the proof of Proposition 4.4 and Theorem 4.5 involving the pseudo-Riemannias manifold (RP, H) and the quotient group (P,dp) with Euclidean-induced metric.
To case the readability, recall the two involved notions of distance and their respective definition:
dr. (Wo, wi) the geodesic distance for any two parameters Wo, W1 D
dp((wo, wi]) the quotient Fuclidean distance for any two equivalence classes [wo], [wi]
p
14

---

dw. wi) inf LEN LEN
(17)
=inf 7 J. (1((())) de
(18)
=inf/ S. Ve(OTT H -
(19)
where the infimum is taken over smoothly differentiable curves y [0, 1] + Rd such that 7(0) Wo and (11) = Wi
dr([wo], [w1]) = inf Σ lipe qill aill.
(20)
where the infimum is taken over all finite sequences {pi (p....... {4 lE=1... C. RR such that [wo]
[Pi], [Pe+1] = [9.] and [q.] = [wi].
Let us state theorem that encapsulate together both the O-distance part in Proposition 4.4 and the c-Ō part in Theorem 4.5 in more unified way.
Theorem B.1. For any Wo. W1 E R holds
d H (Wo w1)
=
dp([wol, (w1l) O,
(21)
and also that, for any 0 there exists 8 >> 0 such that
dp([wo]; (w])) <
dy-H W1) <e dp([wo), (wi]) <e,
(22)
dy =1 (Wo; W1) < =
(23)
We prove the 3 points separately, in Appendix B.I, Appendix B.2 and Appendix B.3 respectively.
B.1
Proof of Eq. 21
The proof logic is
d--nn W1) =
[wo]
[w]]
dp([wo]: [[]]ol/[w.)) 0.
(24)
and we prove the two steps in the two following Propositions, respectivlly Proposition B.2. For any Wo, W1 R it holds
d a (Wo: w1 =
[wo] = [W1]-
(25)
Proof. By definition [wo] = w ] if and only if there exists a piecewise differentiable y [0, 1] - R such that (0) Wo. (11 W, and f(wo, x) (f)(t), for any [0, 1] and X' Then Consider y from the definition of the equivalence relation ~ and define the points w, = Yt))
for ease of notation. Then for any / [].
((1) = lim
Wite
W.
(26)
c-+
For any X it holds that f(W., x) = f(we,x)Vt,
x)VI, l' E [0, 1]. which implies that f [Weee: x)
f(we, x) = 1 E [0, 1JVe E [0, 1-1-- -1]. Thus,
f(Wi+e, x) f(We. Jw, (x) lim lim 50
Wi+e W. Jw. (x) (().
(27)
e-0
This holds to any x X, so the same holds for the per-datum stacked jacobians Jw, (t(() = 0. Thus.
(0)3-,) T f too 7((1) 7()) J Hw, Jw: -1(()= =0
(28)
and we can measure the length of y in the pullback metric as
LEN (7) = f." (t)ll)ls- H di = S/' Odt = O,
(29)
15

---

which gives an upper bound on the geodesic distance
dr- H (Wo wi) = inf LEN)-=(()))) & LENy- = 0,
(30)
7
thus, d1- (Wo, (Wo wi) = 0 and this implication is proven.
3 d H1 Wo, W1)
implies that there exists a O-length differentiable y [o, 1] + TRD such
that 7(0) Wo. ((1 WT. Without loss of generality, we can assume to be non-stationary, i.e.
y1()) # 0. Here O-length means
0 ( (7) S. 1h H,
(31)
which implies that (1)((t) Hy0 = 0 for any t [0, 1] except zero-measure set which we can neglect later. Then
0 = (t((t) = (t) J Jw. (t)
Jw. (t) 0,
(32)
by positive definitess of Hw,, assumed as hypothesis. We highlight that the leftmost 0 in the previous equation is scalar, while the rightmost O is vector in I2 RNO as obtained by the matrix-vector product of J E NOXD with (1(1) R.D. Looking at the equation Jw, 7 = 0 componentwise implies that
Wehighlight
Vx E X, Vo {1, (1...., O), Vt [0. 1]. where [o]. refers to the oth component of a vector v. Thus, for T E [0, 1]. by Fundamental Theorem of Calculus we have
(f ([)]]((())o (())
(33)
[f(wr, - [f(wo: 1/" (Vwlf(w,.x)))
0
V N. Vo E (1....,0),
(34)
Then f(wr,x) = f(wo, x) Vx X and VT' E [0, 1]. So proved that is homotopy of (w, w') such that (t) E (w) Vt [0, 1]. Thus,
w' E (w)
w
[w]] [w].
(35)
and this completes the proof.
L
Proposition B.3. For any Wo. W1 E R it holds
[wol = [wi]
dp([wo, wi]) =
(36)
Proof. This arrow is trivially true by considering the two sequences in the definition of the quotient distance to be of length one and such that p1 = wo and qı = W1-
=dd([woo wil) implies that, by definition of inf, there exists a sequence + 0 such that
Vm E N there exists two finite sequences of points p ,n...pn (m) E R2D and q .....ghh gr (m) E R such that
lpp
(mm)
(mm) = Em:
(37)
where [wo] = 1,(")), (1 lg. " and mm [W1]. Note that the sequence length n may depend on m. Lipschitzness of f in parameters, assumed by hypothesis, means that
1|-- all < e = (1, - f(g, x)l < Ihe
Xx 0,
(38)
thus,
(m)
(m)
Cm
Σ 11 ,x) -s(g:") (m) ,x]1 < 2Lem
X (39)
1=1
Also, by definition of the equivalence class, it holds Vx E X that
[wo] = mm)
[[[(w, x) -(p (m)) ,x]1 = 0
(40)
[mm] = [g. (m))
(p!m) x) -f( (m) x)] 0
(41)
[n( = [w]]
(g!")) x) f(w,, x)] = 0.
(42)
16

---

Thus, by the triangular inequality, the last four equations imply
[[((wo x) f(wi, x)Il - II x) fp1 (m) x)
[1s(p1 (m)
x) f(gh") x)]
(43)
+ ( mm x) f(W1,x)]
[ (m) x) f(ah")))) ,x)l
(44)
m-1
I (m)) x) - f(g (m)) x)]l
| x) x)) -)) -(( xx]] (45)
n-1
20 +
Σ 0 2L
(46)
and this holds for any m N. Taking the limit m -+ oo, 6m -> 0 implies ( x) f(W1, x)] s
Thus, f(wo, x) f(w1. x) Vx E X, and, thus, [wo] [wi]. which completes the proof.
0
B.2
Proof of Eq. 22
In order prove Eq. 22, we first prove a weaker statement
Proposition B.4. For any wo, RD and for any ô > O it holds that
Woo will < =
dg-n(WW, wi) < L8.
(47)
Proof. y [0,1] R be defined as ((t) = (1- + twi, then
y = S. (t)(()) ((())llsHHc)) dt = 1-² || - Ho
(48)
$ / L| - wolldt = L||w] - wolll
and thus,
dy- = inf LEN LENg- L| woll < LD.
(49)
Definition B.5. Let 71 [0, - RR a sequence of piecewise differentiable paths such that
(1) (0) Vi nn..
(50)
Consider the piecewise differentiable path that is the concatenation of the paths one after the other, 3 CAT(T)..... Tx.) defined
it)) 7i i) iigisitt.
(51)
It is straightforward to that LEN (+) - LEN(H))
The choice of e, that prove Eq. 22 trivially follows from the following Proposition B.6. With the assumption of eigenvalue upper bound for any Wo. Wi E R for any > 0 it holds that
any
dr([wo, wil) <
dy- (Wos w1) 3L8.
(52)
Proof. dp dp([wo; wil) < implies that, by the definition of inf, there exists two sequences
{pr) (p:)iml...m fgitiml... C R such that [wo = [pi], [p lasl. [q.] [wi] and such that
Σ 1p: - qill < 25.
(53)
Now the idea define 1 paths the equivalence classes and paths connecting them, then the stacking of the 2n + will give an upper bound on the geodesic distance.
Let us first define the paths 724 for i O....,7 by making use of Proposition B.2
17

---

Woo = [Pi] imply that there exists Yo [0, 1] + R. be such that 2o(0) = wo. 70(1) = P1
and such that LEN N (%)) L6/n+1
[4,] = [w]] imply that there exists 72n [0, 1] + RD be such that Y2n (0) = 9nr 72n (1) Wi and such that LEN 11 (72m) < L6/n+1
for all i 1, , 1, [p++1] [ys] imply that there exists 72i [0, 1] I be such that (0) = pi+1, (1) = ại and such that LEN v (72i) < L6/ n+1
And then for all i = 1....... Proposition B.4 and E Ilp: dill < 28 imply 72:-1 [0, - RD be such that (2) = p.. -2: (1) = and such that Ei=1 L < 2L5. Then concatenating these 2n + 1 paths, we have
(Wo. wi) inf LEN fz (f) LEN y (CAT(TO0......
(54)
2m
LEN LE (()) (%)
LEN (126-1 +
Σ
L- 2L6 + LP. (55)
k=0
B.3 Proof of Eq. 23
Proposition B.7. The pullback metric does not "rotate" too much when moving from to + formally
(1 [[le[DDllll>>H S || S (1
(56)
Proof. Follows from the K -ischiiiz assumption on w J w
The choice of e, that prove Eq. 23 trivially follows from the following
Proposition B.8. With the assumption eigenvalue upper bound for any Wo, Wi RD for any
> 0 it holds that
65
d dr- w1) <
dp([wo, (i]lo.ww])
(57)
Proof. d (Wo, wi) 5 implies that there exist path y: (0, 1] RDD with
26 = LEN,- [ (t'(t))l
(58)
Consider also the Euclidean length LENg(7) which does not depend on 5. Then, for any d there exist (()) E N such that, considering the uniform partition of [0, 1]. t = t/n(6) for = 0, O.... n6) holds the discrete approximation of the integral
(681
36%
ln(t:41) - P: + p: >(t))l F-x
(59)
n(8) Σ
1h(t:+1) Polls Hyw + lp; ((ti)l. ((t.))..HH...
Haee
(60)
where is defined for every i E {0, (0,....n(6)) 1} as follow. Consider the projection PK on the kernel of f HHyttt) and the projection P/¹ on the image of f Hy(t.)) such that the two projections orthogonal and pk + P = I. Define the point P; D2D p; pkk (7(4441) >(t)) + 7tt), which implies that I|:: - 7(44)) Hy = 0. By definition of the projections, it also hold that
p. 7((41) - P/ (x((+11) (t.)), which implies that 7(++1) P. P/ (r(()+1) (tt)) is
18

---

aligned with the non-zero eigenvalues of H',t),, this means that we can resort to the lower bound / on the non-zero eigenvalues and Proposition B.7 to see that
(h((t+1)= Pills Hy(u) Hyyu)
(1 ((.))))))()))))) pill H-
(61)
R
Ki]l)(t+11) - -pill
(62)
[1(((41)) -Pll
Vi {0, ((6) 17,
(63)
where, recalling that (1+(4++1 - (t,)ll = 3/n(6) by construction, there always exist an n6)) big enough such that K 2 1/2, which we can assume without loss of generality. Rearranging the terms, we have
m1
68
Σ
(l((t+z) - Psll-
(64)
1=0
Finally, we can define the points di = 7(441) and we have the two sequences po, .... PP(6)) and
20, 9m(5)-1 whose satisfies the costrains [wo] [po]. [pe+1] [9.] and [1m())-1] [w] We highlight that the dependence on 8 of the length of the sequence is not problematic, as it is sufficient
that (8) is finite for every fixed 6, which Thus,
m(6)-1
65
do((wo. wil) <
Σ
ll qel
(65)
which concludes the proof.
C Proof of existence of the two Riemannian manifolds
This section contains the proof for the existence of the two Riemannian manifolds (Pw. m) and (Pa, m²) embedded in RDD In the rest of the section, consider fixed E R.D.
Fix a training set = {xI,..., x C IR size N and consider the stacked partial evaluation of the network defined as
R:RD - NO
(66)
W W (f(, J(W, xN))
(67)
And define 3(w)= E RNO The differential Vw Jw NOXD equals the stacking of the per-datum Jacobians, and we assume to be full rank thanks to the uniform lower bound on the eigenvalues of the NTK matrix. Full-rankness in the overparametrized setting D NO implies the surjectivity of the differential operator.
03, Ow, i=1.....O:;j=I....
J
(68)
We reorder the w, so that the first NO columns are independent. Then the NO NO matrix
03,
R = Ow, w/ NOI....NO,
(69)
non-singular.
We consider the map
aWw1.... wD) = (3(w) .............................. B(.....()o.o.... S(W)NO.WNOHH.... WNO+I...., WD).
(70)
We obtain
Oos Ow, w i=l.... D.j=ll....D
R
V V
(71)
and this is non-singular. By the inverse function theorem, a is local diffeomorphism. So there an open W C D containing vे such that alw W + a(W) is smooth with smooth inverse.
19

---

Finally, define
P
a 1 : Pr.... .D-NNO))
for p E RD-NO
C
(72)
and similarly
Pw
a a PNO. O.....0) for p RNO D-NO
CRP
(73)
We claim that the two restrictions a are slice charts of Pw and Pw, respectively. Since it is smooth diffeomorphism. it is certainly chart. Moreover, by construction, the points in P are exaclty those whose image under 3 is V. thus, P S [w]. On the other hand, the points Pw parametrize functions that take the values ((w) p in a local neighbourhood of y. Thus, locally never intersects the same equivalence class more than time.
D Proof of Theorem 5.1
By definition of the kernel manifold, we have that if w E Pau, then we have that w E [w/]. Hence for all w E P and for all x E &
f(w,x) f(w', (w',xx x).
(74)
It follows that Varw-Pt, If(w x)] = O for any X X.
For the second statement notice that Varw [f(we if and only if f(w, Xteat) for all w P and for some constant c. Suppose w R₂ and w & "uu(xan)" then f(wr,xtest Xtest) # f(w', xtest). which means that f(w, Xtest) not constant. Hence we have that Var Varw~py [S(w, Xtast] >
E Further results and experimental setup
E.1
Implementation details of the Laplace approximation
Sampling from Laplace's approximation requires computing the inverse square root of a matrix of size D x D, where D is the number of parameters. For most models this problem is intractable. The standard approach to this problem function such as KFAC, Last-Layer, and Diagonal approximations. However, these approximations introduce additional complexity making the task of validating our theoretical analysis much harder. In light of these considerations, choose to sample from Laplace's approximation in a way that is closest to the theoretical ideal, at the cost performing expensive computations.
consider sparse approximations the Hessian of the loss
In small experiments with the toy regression problem, we instantiate the exact GGN and compute vector products with its inverse-square root. For experiments with LeNet and ResNet, we rely on the empirical observation that the spectrum of the GGN is dominated by its leading eigenval- ues (Figure 6)
This makes low-rank approximations of the GGN particu- larly attractive. We choose the Lanezos algorithm (Lanc- zos, 1950) with full reorthogonalization and run it for very high number of iterations, to ensure numerical stabil- ity and very low reconstruction eтror to form our low-rank approximations of GGN. Additionally, Lanczos only requires implicit access to the matrix ory cost and GGN vector products for neural networks can be performed efficiently using Jacobian- Vector products and Vector-Jacobiaa products. If we do
Figure 6: Eigenvalues of the GGN of Convolutional Neural Network trained on MNIST.
avoid the mem-
20

---

Algorithm Laplace diffusion
1: Input: Observed data D, trained MAP point w', number of steps T, number of samples S, rank
k.
2: Initialize samples w....wg as the MAP estimate w/
3: for j in ........ S do
4: for t in I.... ......To T
5: Sample ~ NV Ix)
6:
Compute the top-k eigenvalues(Af) and eigenvectors(U)) of GGNw)
7: w/
+ w JN UJ( a
8: end for
9:
end for
10: Return posterior samples w W
a sufficiently large number of iterations we obtain the non-zero eigenvalues A, and correspond- ing eigenvectors U1. For obtaining samples from the diffusion we can use Algorithm 1 with the eigenvalues and eigenvectors computed using the Lanczos algorithm.
Given the non-zero eigenvalues A, and corresponding cigenvectors U1 we can also form inverse vector products with the square root of GGN + aT. It should be evident from the discussion in section 3 about decomposing the covariance that this vector product with a vector e given by:
(GGN al) U,(A + +
- 22 +
U2 Va
(75)
=U,(A +
-2 (I-
Va --
(76)
=U+((A+o++*
_
1 Va
1
(77)
Va
This allows us to form inverse-square root vector products with the GGN + al given the non-zero eigenspectrum. We run the sampling algorithm H100 GPUs to run the high-order Lanczos decom- position. This approach sampling from Laplace's approximation has O(pk?) time complexity, and O(pk) memory cost, where l is the number of Lanezos iterations and p the number of parameters.
E.2
Experimental details and further results for toy experiments
E.2.1 Toy regression in Figure
In this experiment, we fit small MLP, with hidden layers of width 10 on the sine curve. Due to the small size of the GGN it is possible instantiate and do all the computations explicitly. We sample from the exact Laplace's 's approximation, the non-kernel and kernel subspace of the GGN, and use the neural network and the linearized predictive functions for the top row and the bottom row respectively. For the middle, we simulate a diffusion on the kernel manifold for the center plot, a diffusion in the non-kernel manifold for the right plot and do alternating steps in the manifolds for the left plot which gives us the full distribution.
We also do similar experiment show that the same phenomenon also holds for classification. We use a small convolutional neural network, with convolutional layers with kernel of size 3, to classify a 2-class mixture of Gaussians and look at uncertainties of sampled Laplace, linearized Laplace, and Laplace's 's diffusion. We decompose these uncertainties into their kernel and non non-kernel components respectively. We see the same effect for classification as we did in regression in Figure 7
The main takeaway of these experiments that sampled Laplace underfits in-distribution and this effect is related to the kernel component of the distribution.
E.2.2 Effect of kernel rank on in-distribution in figure
For this experiment, we train small convolutional neural network, with two convolutional layers and kernel of size 3, on MNIST. In this case, the GGN can be instantiated explicitly. We recall that the
21

---

Sampled Full Distnbutton
Sampled Kermel Distribution
Sampled Non-Kernel Distribution
0.00
9/00
Linearized Laplace uncertainly
Linearized Kernel uncertainty
Linearized Non-Kerne uncertainty
8:25
0:45
-5
0.13
Laplace Diffusion uncertainly
uncerfainly
uncertainly
9.8
B.8
-5
Figure 7: Decomposition of uncertainties of Laplace Approximation for the Gaussian mixture classification.
-6
-2
-4,
22 0
GGN is a sum of J Jw (X,) (X (x,) over the dataset where X, are the individual data points. We recall that the rank of the GGN is bounded by NO where N is the number of data points and O is the output dimensions. This suggests considering partial sums by subsampling the data points gives GGN with lower rank. Equivalently, this is GGN with higher dimensional kernel, and hence the usual covariance from Laplace's approximation (GGN + al) has a higher contribution from the kernel subspace.
consider multiple such subsamples and plot the training accuracy for samples from Laplace's approximation against the kernel subspace dimension. Here we see clear trend that the underfitting in sampled Laplace decreases as the rank of GGN increases, or the contribution from the kernel component decreases. This serves to further support suggested hypothesis that the underfitting in sampled Laplace caused by its kernel component and is hence deeply related to reparameterizations.
E.3
Additional benchmarks and results for image classification
We use standard LeNet for the MNIST and FashionMNIST experiments and a smaller version of ResNet (Lippe, 2022), with 272.378 parameters consisting of times group of ResNet blocks. We use those instead of the standard ResNets due to constraints on the computational budget for the CIFAR-10 experiments. We train LeNet with Adam optimizer and a learning rate of 10 10-3. For the ReNet we use SGD with a learning rate of 0.1 with momentum and weight decay.
Training details:
Hyperparameters: We benchmark Laplace diffusion against SWAG, diagonal Laplace, last-layer Laplace, and MAP addition to linearised Laplace and sampled Laplace.
For choosing the prior precision for diagonal Laplace and last-Layer Laplace for each benchmark we do a grid search over the set {0.1, 1.0, 5.0, 10.0, 50.0, 100.03. For Laplace diffusion, sampled Laplace, and Linearised to ensure that the comparison can validate the theory it is preferable to have the same prior precision for all of these methods. So we only do the grid search tune the prior precision for sampled Laplace and use this for all three methods. We keep the hyperparameters for these three methods as similar as possible to have the most informative comparisons.
For Laplace diffusion on MNIST and FMNIST, we simulate the diffusion with a step size of 0.05, with 2000 Lanezos iterations and predict using 20 MC samples. For sampled Laplace and Linearised Laplace, we also use 2000 Lanczos iterations and we predict using 20 MC samples. For the CIFAR-10 experiments, we simulate the diffusion with step size of 0.2, with 5000 Lanezos iterations and we predict using 20 MC samples. For sampled Laplace andlLinearised Laplace use the same number of Lanczos iterations and MC samples.
For SWAG we use a learning rate of 10 with momentum of 0.9 and weight decay of 3e-4 and the low-rank covariance structure in all experiments. For the MNIST and FMNIST experiments we collect 20 models and for the CIFAR- 10 experiments we collect models to sample from the posterior.
Last-layer Laplace is the recommended method by Daxberger et al. (2021c) so should approximate the best performance one get using various possible configurations. For the CIFAR-10 experi-
22

---

ments, the last layer of ResNet is too large to instantiate the full GGN matrix. So we instead use the last 1000 parameters of the model to construct the covariance matrix of the posterior.
Diagonal Laplace requires high prior precision to ensure it does not severely underfit in-distribution (similar to sampled Laplace). It often becomes almost deterministie. So exclude it from the CIFAR results. This has also been observed by Deng et al. (2022) and Ritter et al. (2018).
All additional information about the experimental setup can be found in the submitted code.
E.3.1 In-distribution fit and calibration
We extend Table to benchmark Laplace's diffusion against various other Bayesian methods. Here we see that despite using the neural network predictive is competitive with the best-performing Bayesian methods whereas Sampled Laplace performs significantly worse.
MNIST
Conf. (1)
NIJ.
Acc. (1)
Brier (J)
ECE (J)
MCE (J)
Laplace's diffusion
0.988=0.001
0.040=0.007
0.98610.002
0.02240.00
0.146:0.011
0.773+0.050
Sampled Laplace
0.593=0.002
3.669=0.15 157
0.157+0.031
1.162=0.044
0.436*000027
0.984=0.001
Linearised Laplace
0.967=0.002
0.295=0.041
0.930+0.004
0.111#0.000
0,239+0.032
0.862=000044
SWAG
0.993=0.001
0.032=0.001
0.989#0.002
0.019#0.001
0.187=0.024
0.901=0.041
Last-Layeer Laplace
0.991=0.001
0.047=0.003
0.987+0.002
0.021
0914440.034
0.721=00001
Diagonal Laplace
0.963=0.005
0.078=0.011
0.976*0.003
0.038+0.005
0.095
0.692=0.0229
MAP
0.992=0.002
0.042=0.004
0,988+0.000
0.021#0.000
0.19840.041
0.685=0.091
FMNIST
Conf. (f)
NLL (1)
Acc. (t)
Brier (4)
ECE (J)
MCE (1)
Laplace's 's diffusion
0.900+0.001
0.275+0.016
0.90620.007
0.108 0.141
0.729 0.02
Sumpled Lapluce
0.618:0.021
1.507=0.160
0.098+0.010
1.295+0.014
0.518=0.013
0.986#0.001
Linearised Laplace
0.897=0.003
0.423+0.014
0.86240.005
0.207+0.006
0.147=0.017
0.75610.048
SWAG
0.925+0.002
0.25960.004
0.911+0.006
0.13540.004
0.152
0.75240.067
Last-Layer Laplace
0.91420.001
0.280+0.016
0.901#0.006
0.144+0.004
0.131=0.003
0.67340.083
Diagonal Laplace
0.86240.009
0.323+0.022
0.889#0.010
0.165#0.011
0.102=0.000
0.66040.051
MAP
0.914+0.000
0.279a0.015
0,90410.006
0.143:0.004
0.5500000
0.609#0.033
CIFAR-10
Conf. (1)
NLJ. (L)
Acc. (t)
Brier (1)
ECE (L)
MCE (L)
Laplace's 's diffusion
0,948 0004
0.403+0.007
0,88940.00 005
0.180+0.009
0.26420.048
0,87940.038
Sampled Laplace
0.843+0.004
0.997:0.222
0.71740.049
0.422+0.081
0.221
0,804+0.080
Linearised Laplace
0,951 007
0.61410.020
0.86340.001
0,222+0.008 002
0,337+0.022
0.78940.035
SWAG
0,94240.000 0.03
0.39340.004
0.88420.002
0.17620.001
0.234+0.014
0.91240.016
Last-Layer Laplace
0.953+0.004
0.343+0.033
0.899a0.001
0.155+0.007
0.260=0.000
0.884+0.029
MAP
0.960+0.003
0.333+0.030
0.913+0.007
0.143+0.009
0.282+0.002
0.932+0.006
E.3.2
Robustness to dataset shift
In these experiments (Fig. 8), to measure in-distribution fit and calibration, we report accuracy, negative log-likelihood (NLL), and expected calibration error (ECE) all evaluated on the standard test sets. We measure the robustness of dataset shift of various baselines by plotting the negative log-likelihood and the expected calibration error against shift intensity. The desired behavior is good in-distribution fit, as close as possible to MAP, and stable calibration errors and NLL under distribution shifts. We see that the Laplace's diffusion competitive against other Bayesian methods.
E.3.3 Out-of-distribattion detection
We extend to benchmark Laplace's diffusion against various other Bayesian methods for Out-of- Distribution Detection. Once again, observe that, despite using the neural network predictive,
23

---

Canpart 20P
Dip-tta Calmratio Ereer
Skith insucrty
AES LAT SHAEE
1.5
Figure 8: Model Fit and Calibration of various posterior sampling methods on in-distribution data(first column) and under distribution shift for (tIST(top row), Fashion MNIST row) and CIFAR- 10(bottom row). We use rotated MNIST, rotated FMNIST, and rotated CIFAR in the second and third columns. Shift intensities denote angles of rotation.
it is competitive with the best-performing Bayesian methods in terms of having a higher AUROC,
whereas Sampled Laplace performs significantly worse
MNIST
Tested on
FMNIST
EMNIST
KMNIST
Conf.())
AUROC(1)
Conf.(
AUROC(1)
Conf.())
AUROC(1)
Laplace's diffusion
0.810=0.031
0.911#0.037
0.947=0.007
0.629#0.021
0.817=0.015
0.930#0.009
Sampled Laplace
0.583=0.015
0.51540.021
0.58420.003
0.51540.004
0.58820.013
0.507+0.015
Linearised Laplace
0.895=0.027
0.71540.086
0.934s0.004
0.53640.028
0.863+0.009
0.757+0.018
SWAG
0.827=0.054
0.949#0.018
0.955=0.004
0.627+0.016
0.823+0.013
0.947+0.0007
Last-Layer Laplace
0.84220.046
0.90420.033
0.958=0.005
0.625+0.028
0.837+0.009
0.935+0.004
Diagonal Laplace
0.770=0.026
0.866:0.030
0.907+0.005
0.61740.009
0.747+0.013
0.897+0.005
MAP
0.850=0.035
0.907+0.038
0.959+0.005
0.634+0.019
0.837+0.010
0.938+0.006

---

FMNIST
Tested
MNIST
EMNIST
KMNIST
Conf.())
AUROC(1)
Conf.(.)
AUROC(1)
Conf.())
AUROC(1)
Laplace's diffusion
0.734a0.039
0.759=0.045
0.730+0.020
0.741#0.010
0.746#0.012
0.749=0.023
Sampled Laplace
0.597=0.027
0.495=0.037
0.593+0.026
0.50340.036
0.59840.024
0.493=0.033
Linearised Laplace
0.817=0.022
0.625=0.050
0.813+0.003
0.628+0.013
0.816#0.014
0.624#0.020
SWAG
0.727=0.013
0.817=0.015
0.763+0.028
0.769#0.026
0.777+0.007
0.782=0.010
Last-Layer Laplace
0.757=0.019
0.761=0.031
0.76040.017
0.73540.018
0.77240.008
0.747=0.021
Diagonal Laplace
0.652=0.033
0.767=0.032
0.682+0.019
0.719+0.033
0.696*0.0225
0.728=0.032
MAP
0.759=0.019
0.757=0.032
0.762+0.018
0.730=0.015
0.773+0.010
0.743 024
CIFAR-10
Tested on
CIFAR-100
SVHN
Conf.())
AUROC(1)
Conf.(L)
AUROC(()
Laplace's s diffusion
0.790=0.002
0.851#0.002
0.76440.000
0.862+0.010
Sampled Laplace
0.727=0.026
0.687+0.033
0.79240.022
0.599=0.038
Linearised Laplace
0.818+0.005
0,837+0.006
0.809#0.033
0.85420.024
SWAG
0.776+0.00%
0.845+0.000
0.729 009
0.876+0.003
Layt-Layer Laplace
0.789=0.006
0.864#0.001
0.786#0.029
0.868+0.015
MAP
0.797=0.004
0.873+0.001
0.79240.034
0.878+0.017
E.4
Short discussion on benchmarks
Sparse Approximations of GGN. Our theoretical analysis is mainly concerned with the ideal versions of Laplace's 's approximations, where we consider the full GGN in the covariance without any approximations. However, it can also shed some light on other Bayesian methods.
is common to use sparse approximations of the GGN when doing Laplace's approximations. Interestingly observe that Laplace's 's approximations with sparse GGN such diagonal Laplace, Last Layer, etc do not benefit from linearization to the same degree (Fig. 9).
Shagle Laglon Uecoraally
Figure 9: Predictive uncertainty Laplace's approximation with neural network and linearized predictive (top row) and diagonal Laplace with neural network and linearized predictive (bottom row).
This is perfectly consistent with our analysis as we show that the benefit of linearization primarily comes from the Jacobian term in the linearized predictive and the GGN sharing kernel. Diagonal and other sparse approximations
not share the spectral properties with the Jacobian the linearized predictive. George et al. (2018) note the potential advantages of having the spectrum of the approx- imate curvature more aligned with the true curvature. This suggests future directions to improve various approximations to the GGN by accounting for reparameterizattion..
25

---

SWAG. Another baseline that can be explained using our method is the SWAG, It has been shown that in (Li et al., 2021) SGD steps close the optimum can be decomposed into a normal space component and tangent space component. In our terminology, this can be thought of as diffusion step in the Kernel manifold and diffusion step the Non-kernel manifold. Hence it can be shown that SWAG roughly approximates diffusion-basee basedd posterior. Hence our analysis can provide some theoretical grounding for heuristic methods like SWAG.
26

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The theoretical findings are states in abstract and introduction clear as
possible withouth first introducing all the notation of the paper. The experimental results match the claim and confirm the theoretical findings.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The conclusion raises several points of concern.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model cell-specfiiaatinn. asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
27

---

Answer: [Yes]
Justification:
all theoretical statements are given detailed derivation in the appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: We provide detailed appendix describing the experimental setup. We will further release all code to reproduce the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
28

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We rely established benchmark data. All code will be released under open source software licence upon paper acceptance.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cm
wwwpiiiisssoooooooloooooo onPolicy icy) for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / /ttpss.ccpppliiiiim wwtppccpcccc////cccccccccccc..lt htalllicguiieee http//uudeeee//ee https//oooSSuuiisseeeeeee
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We specify the details about the experimental setup in the appendix, Further- more, we also provide the code supplemental material which contains the full details of the experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification:
We provide errors and/or standard deviations for presented results.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
29

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: these details are presented in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https httpp://neetee //tp:rr//////pp/// cc/ Ceapblic/thhcscuideeee.
Answer: [Yes]
Justification:
The work is predominantly theoretical and does not raise ethical concerns.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction..
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification: The paper is predominantly theoretical. However, uncertainty quantification is potential remedy for several issues with currently deployed machine learning models, so there a potential for positive societal impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
30

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
we do not deem this be relevant.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best. faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: only consider well-established benchmark data, which cite appropri- ately.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
31

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: will release open source code for reproducing experiments upon paper acceptance. This well documented.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: this paper does not involve crowdsourcing other forms of research involving human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: this paper does not involve crowdsourcing other forms of research involving human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
32

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
33


Reviews:
[{"review_id": "ldunpY9eB2", "reviewer": "Reviewer_AZGF", "summary": "This work questions why the linearized Laplace approximation can be effective for uncertainty estimation, whereas the vanilla Laplace approximation can lead to poor performance, such as underfitting. This work argues that this degradation is attributed to the fact that the random function derived from each parameterization of a DNN may not form an invariant function space. \n\nSpecifically, since the null space $ \\text{ker} (\\text{GGN}_{w}) $ is not empty for over-parameterized NNs, sampling the random function containing elements of null space could result in an inconsistent random function. This can lead to under-fitting issues as observed in Laplace approximation. \n\nTherefore, the authors believe that sampling the random function on $ \\text{im} (\\text{GGN}_{w}) $, i.e., the linearized Laplace, could be more effective than the Laplace method. \n\nFurthermore, for a given weight $w$, they explore how to sample the same function with  $ f(w) $, i.e., $ f(g(w)) = f(w) $, by sampling the weight on a specific manifold $ g(w) $. To this end, they employ the concept of quotient space to elaborately define the manifold and diffusion to sample $ w $ from the manifold $ g(w) $.", "questions": "* To confirm my understanding, does the Laplace diffusion denote that for a given $w$, the sample functions $ { w_t }$ are obtained by the update rule using $\\text{GGN}^{+}$ as described in the SDE on the manifold?\n\n* Does $\\text{GGN}^{+}$ mean the positive definite matrix of $\\text{GGN}$, which is obtained by applying SVD on $\\text{GGN}$ and then using the eigenvectors with positive eigenvalues?\n\n* Is the Laplace diffusion the post-hoc method, meaning that $w$ is first obtained by MAP inference, i.e., $w_{\\text{MAP}} = R^{f}_{\\mathcal{X}}(w)$, and then $w_t$ is obtained by Laplace diffusion with $\\text{GGN}^{+}$?\n\n* Compared to the performance of SWAG described in attached appendix, the performance of the Laplace diffusion does not seem significantly improved. \nIn this context, for a given $w$, is it important to sample the weight parameter according to the invariant Riemannian manifold? \n\n* Rather, as considering that SWAG focuses on finding good neighborhood of $w$ in training procedure and obtains comparable performance, isn't it more important to focus on how to find good $w$ (for example, $ w_{\\text{swa}} $ ) and explore the subspace of $w$? I am just curious about your opinion on this.", "limitations": "See above weaknesses and limitations.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* This work reasonably reveals why the linearized Laplace approximation can be more effective than the standard Laplace approximation. This explanation appears to be novel.\n\n* This work attempts to justify the above reasoning using a mathematical elaboration framework.", "weaknesses": "* The lack of explanation of background knowledge, such as quotient space and Riemannian manifold, makes the paper difficult to understand.\n\n* Although it appears effective compared to Laplace and Linearized Laplace, the performance improvement seems marginal when compared to other baselines such as SWAG and Last-layer Laplace.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable review and questions, which will further improve the paper.\n\n- “The lack of explanation of background knowledge, such as quotient space and Riemannian manifold, makes the paper difficult to understand.”\n\nWe acknowledge this issue, which is due to the space constraints of a conference paper. Because of the additional allowed page, we will be happy to extend the background section to cover the necessary knowledge in Riemannian geometry to benefit future readers.\n\n- “Although it appears effective compared to Laplace and Linearized Laplace, the performance improvement seems marginal when compared to other baselines such as SWAG and Last-layer Laplace.”\n\nWhile this is a valid comment we believe it stems from a slight misunderstanding of the main point of our experiments. \n\nThe main goal of this paper is to understand the uncertainties obtained from Laplace approximations and more generally to shed light on some of the problems encountered in Bayesian deep learning. More specifically decomposing the uncertainty into its two fundamental components and how they interact with the choice of the predictive (neural network or the linearized neural network). Our purpose is not to advance the state-of-the-art of Bayesian methods.\n\nWe include the baselines in the appendix and omit them from the main paper for precisely this reason. We want to validate the theory in realistic settings while staying as close as possible to the theoretical ideal (these choices are also discussed in Sec. 7 and Appendix E.1).  \n\nOne way to think about the experiments is as follows:\nObservation: Linearized Laplace underfits less (in-distribution) than Sampled Laplace. \nHypothesis: This difference is due to improper accounting of the reparameterization properties in the Sampled Laplace.\nTheory: The theory presented in the paper provides mathematical justification for this hypothesis.\nExperiment: In addition to the theoretical justification, we build a method that is as close to Sampled Laplace as possible while also accounting for the reparameterization properties of the neural network. This is the Laplace diffusion which turns out to not underfit in-distribution, hence supporting the extensive theory developed in the paper. \n\nOf course, we do not disregard the importance of building practical, performant, and scalable Bayesian methods. In future work, we intend to build on the here-presented theoretical insights to achieve this exact goal. However, understanding why Bayesian deep learning struggles (this paper) is a prerequisite. The main takeaway of this paper is not a particular method but a general principle: Think about reparameterization issues when building Bayesian methods.\n\nTo confirm my understanding, does the Laplace diffusion denote that for a given 𝑤, the sample functions 𝑤𝑡 are obtained by the update rule using GGN+ as described in the SDE on the manifold?\nDiffusion can be performed by simulating the SDE in the kernel manifold and the non-kernel manifold or both (using alternating steps, for example). In Figs. 2 and 7 we perform all 3 diffusions. For the larger experiments (Sec. 7), your description is correct. We only perform the diffusion in the non-kernel manifold. The reasons for this are discussed in Section 5.\n\n- \"Does GGN+ mean the positive definite matrix of GGN, which is obtained by applying SVD on GGN [...]?\"\n\nEssentially yes. However, we use a Lanczos decomposition of the GGN instead of SVD.\n\nIs the Laplace diffusion the post-hoc method, meaning that 𝑤 is first obtained by MAP inference, i.e., 𝑤 MAP=𝑅𝑋𝑓(𝑤), and then 𝑤𝑡 is obtained by Laplace diffusion with GGN+?\n\nYes, it is a post-hoc method. As stated above we want to build a method as close to Sampled Laplace as possible which additionally is reparameterization invariant. Since Sampled Laplace is typically a post-hoc method it makes sense that Laplace diffusion should be as well.\n\n- \"Compared to the performance of SWAG described in attached appendix, the performance of the Laplace diffusion does not seem significantly improved. In this context, for a given 𝑤, is it important to sample the weight parameter according to the invariant Riemannian manifold? Rather, as considering that SWAG focuses on finding good neighborhood of 𝑤 in training procedure and obtains comparable performance, isn't it more important to focus on how to find good 𝑤(for example, 𝑤 swa) and explore the subspace of 𝑤? I am just curious about your opinion on this.\"\n\nThis is an interesting point of discussion, which we will approach in steps.\nFirst, note that while the mode (i.e. the MAP) of the Gaussian approximate posterior is very important, the reparametrization issue we discuss is mainly related to the covariance. SWAG restricts the covariance to a subspace.\nSWAG involves several heuristics that are hard to characterize theoretically, so here we consider a simplified version:\nFix a mode parameter $w_m$, then perform one SGD step to find a new parameter $w_1$. Repeat this (starting from $w_m$) several times.\nCompute the empirical covariance of $w_1,w_2,..$ and use this for the Gaussian posterior approximation.\nIt’s important to clarify that the kernel directions (i.e. 0-eigenvalue directions of the GGN) depend on the dataset. We can, thus, consider both the full-dataset-kernel directions and the batch-kernel directions. Notably, the full-dataset-kernel is always contained in the batch-kernel.\nThe covariance obtained by simplified-SWAG is then guaranteed to be 0 along the full-dataset-kernel direction.\nThis shows that the simplified-SWAG approximate posterior is locally reparametrization invariant, thereby demonstrating how the developed theory can help us understand existing methods and eventually develop new ones.\n\nWe thank the reviewer for their efforts and we encourage them to reconsider their assessment of the experiments in light of the additional context we have provided above.", "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": "", "extra_scores": {}}, {"review_id": "pcBmD74X1A", "reviewer": "Reviewer_12ke", "summary": "The authors provide theoretical analysis of the underfitting problem of Laplace approximation. Specifically, they show that the underfitting of Laplace approximation is due to the approximate posterior covariance is not invariant under reparameterization. Moreover, they propose a reparametrization invariant diffusion posterior to address undefitting. The method are evaluated over standard image classification benchmarks. The quality of the uncertainty for both in-distribution and out-of-distribution tasks are considered.", "questions": "See weakness.", "limitations": "The authors do not address potential negative social impact since the paper is predominantly theoretical.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper provides, to my knowledge, the first theoretical justification of why Laplace approximation sufferes from underfitting: the approximate posterior covariance is not reparameterization invariance.\n2. Given the theoretical analysis, the idea of reparametrization invariant posterior is natural and proves to be effective empirically.\n3. The paper is easy to follow and well organized.", "weaknesses": "1. The proposed method, as mentioned by the authors as well, sufferes from more expensive computation. It would be good to report the time metric for readers to properly assess the practicability of the method.\n2. The datasets considered are a bit outdated and the networks considred seem to be quite small, such that overall the performance is on the lower end (e.g. <90% acc. for CIFAR10). Recent BDL papers typically consider larger datasets (such as ImageNet) and deeper networks (e.g. [1]). Furthermore, last-layer Laplace seems to be competitive with or even outperform the proposed method in some experiments.\n3. The method is tailored to improving Laplace approximation, and there is no discussion related to any other approximate inference techiques (e.g. variational inference).\n\nReferences\n[1] Antoran et al. Sampling-based inference for large linear models with application to linearised Laplace. ICLR 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the excellent review, which we will reply to in parts.\n\n- “The paper provides, to my knowledge, the first theoretical justification of why Laplace approximation sufferes from underfitting: the approximate posterior covariance is not reparameterization invariance.”\n\nThis is correct. Additionally, we are the first to provide a complete characterization of continuous reparameterizations of a neural network. Certain reparameterizations such as scale reparameterizations for ReLUs are already known but we provide a complete characterization of these in terms of the kernel manifold. Furthermore, to the best of our knowledge, we are the first to consider data-dependent reparameterizations rather than only global reparameterization (a subset of the former). \n\n- “The proposed method, as mentioned by the authors as well, suffers from more expensive computation. It would be good to report the time metric for readers to properly assess the practicability of the method.”\n\nFor a $P$ parameters network, computing $S$ samples by approximating the diffusion with $T$ discrete time steps and a low-rank $K$ approximation of the GGN requires $O(KP + SP)$ in memory and $O(K^2PST)$ in time, i.e. it is $ST$ times as expensive as linearized Laplace.\n\nHowever, we want to emphasize that our main goal here is to understand the uncertainties obtained from Laplace approximations and more generally to understand some of the problems encountered in Bayesian deep learning. More specifically decomposing the uncertainty into its two fundamental components and how they interact with the choice of the predictive (neural network or the linearized neural network). Our purpose is not to advance the state-of-the-art of Bayesian methods.\n\nOne way to think about the experiments is as follows:\n\nObservation: Linearized Laplace underfits less (in-distribution) than Sampled Laplace. \nHypothesis: This difference is due to improper accounting of the reparameterization properties in the Sampled Laplace.\nTheory: The theory presented in the paper provides mathematical justification for this hypothesis.\nExperiment: In addition to the theoretical justification, we build a method that is as close to Sampled Laplace as possible while also accounting for the reparameterization properties of the neural network. This is the Laplace diffusion which turns out to not underfit in-distribution, hence supporting the extensive theory developed in the paper. \n\nFor this reason, we voluntarily pay the extra computational factor $T$ to get as close as possible to the theoretical ideal. \nFor example, we could have made the method more scalable by fixing low values for $K$ or $T$. While such approximations might be more practical and scalable, the resulting experiments can no longer validate the theory due to the brutal approximations. This is also why we only include comparison with baselines in the appendix omitting them from the main paper, where we only compare the different predictive functions.\n\nOf course, we do not disregard the importance of building practical, performant, and scalable Bayesian methods. In future work, we intend to build on the here-presented theoretical insights to achieve this exact goal. However, understanding why Bayesian deep learning struggles (this paper) is a prerequisite. The main takeaway of this paper is not a particular method but a general principle: Think about reparameterization issues when building Bayesian methods.\n\nFurther discussion of these issues and time complexity is available in Appendix E.1.\n\n- “The datasets considered are a bit outdated and the networks considred seem to be quite small [...]”\n\nFor the experiments, we have chosen to replicate the key experiments from the Laplace Redux paper [1] which is an important reference for Laplace approximations. Their largest experiments are also on CIFAR-10 with ResNets so we don’t believe the scale of our experiments is too unreasonable. \n\nSecondly, we are certainly aware of the recent works that do Laplace approximation on larger models (see references in the paper). But they all invariably rely on fairly brutal approximations such as KFAC, last-layer, or LoRa. We are not at all opposed to this but hopefully, it is clear from the previous response why such approximations are inappropriate in the context of our experiments: these aim for theory-validation rather than the usual focus on performance and scalability. \n\n- “The method is tailored to improving Laplace approximation, and there is no discussion related to any other approximate inference techiques (e.g. variational inference).”\n\nWhile it’s true that our primary focus is on Laplace approximations, we believe that the tools developed here are more generally applicable. The theory of reparameterizations can find applicability in various topics such as continual learning and in studying the loss landscape. The issues considered here are also relevant to other Bayesian methods. It's not possible to explicate this fully but we have a brief discussion of this direction in appendix E.4. Another possible direction is Variational Inference with the diffusion posterior as a variational family. Secondly, we believe that the underfitting of Sampled Laplace is an important open problem in the community, which should not be understated. For example, this has also been discussed in a recent position paper [2]\n\nWe thank the reviewer for their efforts and we urge them to reconsider their assessment of the experiments in light of the additional context we have provided above.\n\n[1] Daxberger, Erik, et al. \"Laplace redux-effortless Bayesian deep learning.\" Advances in Neural Information Processing Systems 34 (2021): 20089-20103.\n\n[2] Papamarkou, Theodore, et al. \"Position paper: Bayesian deep learning in the age of large-scale AI.\" arXiv preprint arXiv:2402.00809 (2024)", "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": "", "extra_scores": {}}, {"review_id": "HGop9xUyoV", "reviewer": "Reviewer_gWds", "summary": "The paper addresses the problem of maintaining invariance under reparameterization of Bayesian Neural Networks (BNNs). This issue undermines the reliability of BNNs since different parameterizations of identical functions produce different posterior densities. Through theoretical analysis and empirical validation, the authors demonstrate the effectiveness of their approach by extending a geometric view (in Linearized Laplace) and using a Riemannian diffusion process to achieve reparameterization invariance.", "questions": "* Although mentioned in related work, I feel that this paper needs comparative experiments with methods such as Riemannian Laplace and (Bergman et al., 2024) and Connectivity Laplace (Kim et al., 2023).\n* As mentioned in Weakness section, I would recommend the author to evaluate the Riemannian Diffusion in other domain dataset & architectures (e.g., Transformers).", "limitations": "Limitations are mentioned in \"Weakness\" and \"Questions\".", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* Proposed Riemannian diffusion process incorporates the properties of the linearized Laplace approximation into the original neural network prediction, resulting in a better posterior fit.\n* The method leverages advanced mathematical tools from differential geometry, offering a sophisticated approach to maintaining reparameterization invariance.\n* In both in-distribution and out-of-distribution scenarios, the proposed method consistently outperforms other approaches.\n* The paper is well-organized and well-written, with each section building logically on its predecessor.", "weaknesses": "* Although the proposed method is theoretically sound, its computational complexity is higher than simpler approximations. This could limit its practicality for very large-scale neural networks.\n* Similar to previous works in Linearized Laplace, the need for new computational pipelines to handle the GGN matrix and its induced structure is acknowledged but not fully addressed in the paper.\n* Although the experiments are comprehensive, they demonstrate results on standard datasets and relatively small to medium-sized neural networks.Testing the method on more diverse datasets and larger-scale models would strengthen the empirical validation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable review. \n\n- “Although the proposed method is theoretically sound, its computational complexity is higher than simpler approximations. [...]”\n\nWhile this is a valid comment we believe it stems from a slight misunderstanding of the main point that is being made with our experiments. \n\nThe main goal of this paper is to understand the uncertainties obtained from Laplace approximations and more generally to understand some of the problems encountered in Bayesian deep learning. More specifically decomposing the uncertainty into its two fundamental components and how they interact with the choice of the predictive (neural network or the linearized neural network). Our purpose is not to advance the state-of-the-art of Bayesian methods.\n\nOne way to think about the experiments is as follows:\nObservation: Linearized Laplace underfits less (in-distribution) than Sampled Laplace. \nHypothesis: This difference is due to improper accounting of the reparameterization properties in the Sampled Laplace.\nTheory: The theory presented in the paper provides mathematical justification for this hypothesis.\nExperiment: We build a method that is as close to Sampled Laplace as possible while also accounting for the reparameterization properties of the neural network. This is the Laplace diffusion which turns out to not underfit in-distribution, hence supporting the extensive theory developed in the paper. \n\nFor this reason, we voluntarily spend the extra compute to get as close as possible to the theoretical ideal. Specifically, for a $P$ parameters network, computing $S$ samples by approximating the diffusion with $T$ discrete time steps and a low-rank $K$ approximation of the GGN requires $O(KP + SP)$ in memory and $O(K^2PST)$ in time, i.e. it is $ST$ times as expensive as linearized Laplace.\n\nWe could have made the method more scalable by fixing low values for $K$ or $T$. However, the resulting experiments can no longer validate the theory due to the brutal approximations. This is also why we only include comparison with baselines in the appendix omitting them from the main paper, where we only compare the different predictive functions.\n\nOf course, we do not disregard the importance of building practical, performant, and scalable Bayesian methods. In future work, we intend to build on the here-presented theoretical insights to achieve this exact goal. However, understanding why Bayesian deep learning struggles (this paper) is a prerequisite. The main takeaway of this paper is not a particular method but a general principle: Think about reparameterization issues when building Bayesian methods.\n\n- “Similar to previous works in Linearized Laplace, the need for new computational pipelines to handle the GGN matrix and its induced structure is acknowledged but not fully addressed in the paper.”\n\nIndeed, our paper concludes with a brief discussion of this computational challenge. We intended to acknowledge the inherent difficulty and to draw more attention to this problem. That being said, the paper does contribute advances in this direction as our implementation uses efficient matrix-free access to the GGN using automatic differentiation combined with iterative solvers to obtain samples from the Laplace approximation. This is a significant improvement over existing implementations. We deemphasize these computational aspects of our work and only discuss them implicitly because our main focus is on the theoretical explanation.\n\n- “Although the experiments are comprehensive, they demonstrate results on standard datasets and relatively small to medium-sized neural networks. Testing the method on more diverse datasets and larger-scale models would strengthen the empirical validation.”\n\nWe have chosen to replicate the key experiments from the Laplace Redux paper [1] which is an important reference for Laplace approximations. Their largest experiments are also on CIFAR-10 with ResNets so we don’t believe the scale of our experiments is too unreasonable. Furthermore, scaling to larger datasets and models necessarily requires making some fairly brutal approximations, by doing subnetwork inference, KFAC, diagonal approximations, etc. We are not at all opposed to this but hopefully, it is clear from the previous response why such approximations are inappropriate in the context of theory-validation.\n\n- “Although mentioned in related work, I feel that this paper needs comparative experiments with methods such as Riemannian Laplace and (Bergman et al., 2024) and Connectivity Laplace (Kim et al., 2023).”\n\nAs mentioned above, our main experimental goal is to validate the theory. Further note that the public code for the first reference is orders of magnitude slower than our code and it is infeasible to run on as large models as we do. The public repository for the second reference has no documentation and no code for reproducing the experiments in their paper. Consequently, it would be highly non-trivial to run them on our benchmarks.\n\n- “As mentioned in Weakness section, I would recommend the author to evaluate the Riemannian Diffusion in other domain dataset & architectures (e.g., Transformers).”\n\nTo our knowledge, it's not standard practice to evaluate Laplace approximations on transformers. We are only aware of [2] which evaluates Laplace approximations on transformers and they perform Laplace approximations only on the low-rank adaptation weights. Hopefully from the previous response, it is clear why such approximations are not appropriate for our experiments given our purpose. \n\nWe thank the reviewer for their efforts and we urge them to reconsider their assessment of the experiments in light of the additional context we have provided above.\n\n[1] Daxberger, Erik, et al. \"Laplace redux-effortless Bayesian deep learning.\" Advances in Neural Information Processing Systems 34 (2021)\n\n[2] Yang, Adam X., et al. \"Bayesian low-rank adaptation for large language models.\" arXiv:2308.13111 (2023)", "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": "", "extra_scores": {}}, {"review_id": "rIpijWORIU", "reviewer": "Reviewer_v5yb", "summary": "The paper studies the effects of the invariance of Bayesian neural networks under reparametrization on their approximate posterior inference. The primary effect is the ambiguity on the uncertainty of the inferred function as it can be represented by multiple reparameterizations, each of which may be assigned a different uncertainty estimate. The paper takes Laplace approximation as a case and studies its properties when used for Bayesian neural net inference using linear algebraic and differential geometric tools. The paper derives an algorithm from its key theoretical outcomes and shows it to outperform its counterparts in various downstream uncertainty quantification tasks.", "questions": "* How does the argumentation in the first paragraph of Section 3 come together with the main message of the paper? Is the problem not instead that the posterior assigns separate masses to different weights that represent the same functional form? How exactly does this produce a pathology in the posterior besides the need to detect multiple equivalent representations and add their masses? Why do we actually need to design the related degrees of freedom? Overall I buy the argument and see the problem, but my point is that probably this bridge paragraph may be misleading the reader.\n\n * What does f_lin^w precisely mean in Eq 4? Is it the first-order Taylor approximation of f?", "limitations": "The paper briefly discusses its limitations in the final paragraph of the Conclusion section. The potential negative societal impacts of the studied work are not discussed probably because they are thought to be inapplicable.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "* This is a very high-calibre paper, a complete piece of work, that studies an original and significant problem with impressive theoretical rigor. \n\n* The presentation is extremely clear. Figures 1 and 2 are truly helpful to the reader to grasp the visual intuition behind the studied problem and the proposed solution.\n\n* The compilation of theoretical tools such as the a parameter manifold on the Generalized Gauss-Newton (GGN) approximation and a diffusion defined on it are original, advanced, and elegant.\n\n* The reported numerical results are both comprehensive and strongly in favor of the central claims.", "weaknesses": "* Overall this is excellent piece of work with a mature write-up. Just its presentation may be improved tiny bit. For instance, certain terms are used very early on in the paper such as the introduction and Figure 1 in their meanings a bit unusual for the broad ML research readership, such as a \"kernel\" and \"diffusion\". While their allocation in the particular context is appropriate, reading experience may be improved if their meanings are made more explicit, at least verbally, in their first use.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the excellent review and we are very grateful for a thorough engagement with our work and for pinpointing its highlights.\n\nWe acknowledge your point (under Weaknesses) about providing a bit more explanation of unconventional terms such as kernel (that in machine learning often means something entirely different). We will take this feedback to heart and update the manuscript accordingly.\n\nRegarding the opening paragraph of Sec. 3, then we acknowledge that, in particular, the phrasing “pathological representations” is too vague and opens the door to misinterpretations. We suggest to rephrase the sentence\n “However, as we have argued, traditional approximate Bayesian inference does not correctly capture this, leading to pathological representations” \ninto \n“However, as we have argued, traditional approximate Bayesian inference does not correctly capture this and assigns different probability measures to identical functions.”\nWe believe such a change clarifies the text, but are happy to take further feedback.\n\nYes, $f_{lin}^w$ is exactly the mentioned Taylor expansion. The equation is given in line 107 of the submitted manuscript, but it would have been better to already introduce this notation around line 64. We will make this change.\n\nWe hope this answers your questions and concerns, and are happy to discuss further. Once again, we are grateful for the favorable review that exactly captures our intention, and we hope that you will continue to fight for the paper.", "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": "", "extra_scores": {}}]

Author rebuttal:
We are grateful to the four reviewers who all argue in favor of acceptance.

We observe a general agreement that the developed theory is both novel and sheds significant light on the difficulties of Bayesian deep learning induced by reparametrizations. Some concerns have been raised about the experimental part of the paper. Specifically, there were concerns about the scalability and performance of Laplace diffusion. Here we want to clarify that the stated intent of the experiments is to support the developed theory rather than to develop a state-of-the-art and scalable method (despite Laplace diffusion’s favorable performance). For this reason, we do not optimize the method for performance but keep it as close as possible to Sampled Laplace and we also avoid approximations that might make it more scalable. We argue that the work presented provides the theoretical foundation opening up new insights that will allow for the development of new and more efficient methods that will be both computationally efficient and accurate. The main takeaway of this paper is not a particular method but a general principle: “ Think about reparameterization issues when building Bayesian methods”

We reply to the individual reviews below.

Meta-review: This work studies the reparameterization invariance of Laplace approximations in Bayesian neural networks. The study begins by noting that the linearized Laplace suffers less from underfitting than the sample Laplace approach. This observation leads to the hypothesis that the discrepancy arises from a failure to properly account for reparameterization invariance in the nonlinear setting. The authors provide a rigorous theoretical analysis to support this claim and propose a modification to the sample Laplace method to address the issue. The empirical results strongly support the theoretical findings. All reviewers agree that this work presents an original perspective, addresses a significant problem, and offers a complete theory backed by convincing empirical evidence. Therefore, I recommend acceptance.

The authors are encouraged to include a detailed introduction to background concepts, such as quotient spaces and Riemannian manifolds, in the final version.
Decision: Accept (spotlight)