Reparameterization invariance in approximate
Bayesian inference

Hrittik Roy† , Marco Miani†
Technical University of Denmark{hroy, mmia}@dtu.dk

Philipp Hennig, Marvin Pförtner, Lukas Tatzel
University of Tübingen, Tübingen AI Center{philipp.hennig, lukas.tatzel,
marvin.pfoertner}@uni-tuebingen.deCarl Henrik Ek
University of Cambridge,
Karolinska Institutet
che29@cam.ac.uk
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 to maintain invariance under reparameterization, i.e.
BNN s 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 withuncertainty over the parametrized function. In this paper, we investigate thisissue 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 geometricview of reparametrizations from which we explain the success of linearization.
Moreover, we demonstrate that these reparameterization invariance properties canbe 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.

1

Introduction

Bayesian deep learning has not seen the same degree of success as deep learning in general. Theoretically, Bayesian posteriors should be superior to point estimates (Devroye et al., 1996), but thepractical benefits of having the posterior are all too often not significant enough to justify theiradditional computational burden. This has raised the question if we even should attempt to estimate
full posteriors of all network parameters (Sharma et al., 2023).
As an example, consider the Laplace approximation (MacKay, 1992), which places a Gaussian inweight space through a second-order Taylor expansion of the log-posterior. When applied to neuralnetworks, this is known to significantly underfit and assign significant probability mass to functions
that fail to fit the training data (Lawrence, 2001; Immer et al., 2021a). Fig. 2 (top-left) exemplifies thisfailure mode for a small regression problem. Interestingly, this behavior is rarely observed outsideneural network models, and the failure appears linked to Bayesian deep learning.
Recently, the linearized Laplace approximation (LLA) has been shown to significantly improve on
Laplace’s approximation through an additional linearization of the neural network (Immer et al.,2021b; Khan et al., 2019). We are unaware of any theoretical justification for this rather counterintuitive result: why would an additional degree of approximation improve
improve the posterior fit?
†

Equal contribution authors listed in random order.

Preprint. Under review.

Figure 1: The weight space of a neural network (Eq. 1) overparametrizes
the associated function space. This
induces families (orange) of weights
corresponding to the same functions.
Model linearization (left) linearizes
these families. In nonlinear models,
Gaussian weight distributions (center)do not adapt to the families, while ourgeometric diffusion (right) captures
the associated invariance with a metric (gray ellipses).
We will show that the failures of Bayesian deep learning can partly be explained by insufficient
handling of reparameterizations of network weights, while the LLA achieves infinitesimal invarianceto reparameterizations. To motivate, consider the simple network (Fig. 1)
f (x) = w1 ReLU (w2 x) ;f : R → R.

(1)

This can be reparametrized to form the same function realization from different weights as f (x) =w1/α ReLU (α w x) for any α > 0. That is, the weight-pairs (w , w ) and (w1/α, αw ) correspond22
1
2
to the same function even if the weights are different (Fig. 1, center).
Thus the approximate posterior cannot reflect the fundamental property of the true posterior, that itshould assign a single unique density to a function regardless of its parametrization.
In this paper, we analyze the reparameterization group driving deep learning and showthat it is a pseudo-Riemannian manifold,with the generalized Gauss-Newton (GGN)
matrix as its pseudo-metric. We prove that thecommonly observed underfitting of the Laplaceapproximation (Fig. 2, top left) is caused byhigh in-distribution uncertainty in directions of
reparametrizations (Fig. 2, top center).

Figure 2: The function space is decomposed
We develop a reparametrization invariant difinto directions of reparameterizations (kernel) andfusion posterior that proveably does not underfunctional change (non-kernel). We improve thefit despite using the neural network predictive
posterior fit by concentrating probability mass on(Fig. 2, center row). Figure 1 (right) visualizes
directions of functional change.
how this posterior adapts to the geometry of
reparametrizations, thereby not underfitting. The diffusion can be simulated with a multi-step EulerMaruyama scheme from which the linearized Laplace approximation (LLA) is a single-step. This
link implies that the LLA infinitesimally is invariant to reparameterizations, due to the otherwise counterintuitive linearization (Fig. 1, left). Experimentally, our diffusion consistently improves posterior
fit, suggesting that reparameterizations should be given more attention in Bayesian deep learning.

2

Background: Laplace approximations

Let fw : RI → RO denote a neural network with weights w, and define a likelihood p(y|fw (x))
and a prior p(w). Laplace’s approximation (MacKay, 1992) performs a second-order Taylor
expansion of the log-posterior around a mode ŵ. This results in a Gaussian approximate pos−1
terior N (w|ŵ, −Hŵ), where Hw is the Hessian matrix. The linearized Laplace approximation (Immer et al., 2021b; Khan et al., 2019) further linearize fw at a chosen weight ŵ, i.e.
fw (x) ≈ fŵ (x) + Jŵ (x)(w − ŵ), where Jŵ (x) = ∂w fw (x)|w=ŵ ∈ RO×D is the Jacobianof 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 et al., 2021b),q(w|D) = N w | ŵ, (GGNŵ + αI)−1

GGN ŵ =

N
X
n=1

2

Jŵ (xn )⊤ H(xn )Jŵ (xn ),(2)

where H(x) = −∂f2ŵ (x) log p(y|fŵ (x)) ∈ RO×O is the Hessian of the log-likelihood and we haveassumed a weight prior N (0, α−1 I). Note that it is trivial to extend to other prior covariances. Thisparticular covariance is known as the generalized Gauss-Newton (GGN) Hessian approximation,
which is commonly used in Laplace approximations (Daxberger et al., 2021b).
To reduce the notational load we stack the per-datum Jacobians into Jŵ = [Jŵ (x1 ); . . . ; Jŵ (xN )] ∈RN O×D and similarly for the Hessians, and write the GGN matrix as GGNŵ = J⊤ŵ HJŵ . For
Gaussian likelihoods, the Hessian is an identity matrix and can be disregarded, and for otherlikelihoods simple expressions are generally available (Immer et al., 2021b).
Sampled and linearized Laplace. The Laplace approximation gives a Gaussian distribution
q(w|D) over the weight space with mean ŵ and covariance Σ. A predictive distribution is obtainedby integrating the approximate posterior against the model likelihood,S
1X
p(y∗ |x∗, D) = Ew∼q [p(y∗ |f (w, x∗ ))] ≈p(y∗ |f (wi , x∗ )), wi ∼ q.
(3)
S i=1
We refer to this predictive method as sampled Laplace. Recent works have suggested linearizing theneural network in the likelihood model to obtain the predictive distribution (Immer et al., 2021b),
S
1X
ŵ
ŵ
p(y∗ |x∗, D) = Ew∼q [p(y∗ |flin(w, x∗ ))] ≈p(y∗ |flin(wi , x∗ )), wi ∼ q.
(4)
S i=1
This is referred to as linearised Laplace. Immer et al. (2021b) argues that the common choice ofapproximating the posterior precision with the GGN implicitly linearizes the neural network andhence the predictive distribution should be modified for consistency.
Sampled Laplace is known to 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). In this paper, we show that the benefit is linked to the lack of
reparameterization invariance.
The lack of reparameterization invariance leads to an additional problem for Laplace approximations. 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 to reparameterizations of
the neural network, and the same holds for the GGN. Depending on which parametrization of theposterior 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 1 (center) illustrates the phenomena.

3

Reparameterizations of linear functions

Deep Learning models excel when they are highly overparametrized, i.e. when they have significantly
more parameters than observations (D ≫ N O). This introduces many degrees of freedom to themodel, which will be reflected in the Bayesian posterior. However, as we have argued, traditional approximate Bayesian inference does not correctly capture this, leading to pathological representations.
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 it to the general case.
The reparameterizations of linear functions can be characterized exactly.
Considerf (w) = Aw + b and a possible reparameterization, g : RD → RD , of this function suchthat f (g(w)) = f (w). It is then evident that A(g(w) − w) = 0. This implies that for any
reparameterization of a linear function, we have g(w) − w ∈ ker(A), where ker(A) denotes thekernel (nullspace) of A. Hence, the linear function cannot be reparametrized if we restrict ourselvesto the non-kernel subspace of the input space or if A has a trivial kernel.
′w
A linearized neural network flin
: w, x 7→ fw′ (x) + Jw′ (x)(w − w′ ) is a linear function in theparameters, where we have linearized around w′ . The above analysis then implies that the kernelof the stacked Jacobian Jw′ characterizes the reparameterizations of the linearized network.

We can also characterize the reparameterizations through the GGN and the corresponding neuraltangent kernel (NTK; Jacot et al. 2018),⊤⊤
GGN w = Jw Jw ,NTK w = Jw Jw .
(5)3

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(GGNw ).
Two orthogonal subspaces.
For any selfadjoint operator (such as positive semi-definite
matrices like the GGN), the image and the kernel
orthogonally span the whole space, i.e.

112

im(GGNw ) ⊕ ker(GGNw ) = RD ,(6)
where the kernel is the hyperplane of vectorsthat are mapped to zero and the image is thehyperplane of vectors spanned by the operator(Fig. 3). For a linearized neural network,
im(GGNw ) spans the effective parameters
P ⊂ RD, i.e. the maximal set of parameters thatgenerate different linear functions RI → RO
when evaluated on the training set.

2

3

3

Figure 3: The weight space can be decomposed
into directions of reparameterizations and functional changes. For linear models (left) these arelinear subspaces given by the kernel and the image,respectively. For nonlinear models, these are the⊥nonlinear manifolds Pw
and Pwi , respectively.
i

A Laplace covariance decomposes into the same subspaces. Recall that the posterior precision isΣ−1 = GGNŵ + αI. Let the eigendecomposition of GGNŵ be UT ΛU, and assume that U1 and U2
are the eigenvectors corresponding to the non-zero eigenvalues Λ̃, and the zero eigenvalues respectively. These form a basis in the kernel and image subspace as discussed above. Then the covariance is,
!−1

T 


U1
U1
Λ̃ 0
−1
Σ=
+ αI
= UT
U1 + α−1 UT(7)
1 (Λ̃ + αIk )2 U2 .
U2
U2
0 0
Consequently, we can decompose any sample from the Gaussian N (ŵ, Σ) into a kernel and an
image contribution, w = ŵ + wker + wim , where wker is the component of the sample that is in
the kernel of GGNŵ and wim is in the image. Note that all probability mass in ker(GGNŵ ) is dueto the prior, i.e. we place prior probability on functional reparameterizations even if we can neverobserve data in support of such.
Underfitting in sampled Laplace can now be
understood. For the linearized approximation, it
holds for training data x ∈ X that,
ŵ
ŵ
flin
(ŵ+wker +wim , x) = flin(ŵ+wim , x).
Hence, the linearized predictive only samples in
the image subspace consisting of unique functions. This is not true for the sampled Laplaceapproximation, which also samples in the kernel subspace. Since sampled Laplace does not
linearize the neural network, the kernel does notcorrespond to reparameterizations. It hence adds
“incorrect” degrees of freedom to the posterior
as artifacts of the Gaussian approximation.

Figure 4: Underfitting of sampled Laplace is lesspronounced when the rank of the GGN is higher fora fixed number of parameters. This is consistent
with our hypothesis as a high GGN rank implies alower dimensional kernel. For experimental details,
see appendix E.2.
Empirically, sampled Laplace is only observed to underfit in overparametrized models. Fig. 4
investigates this by increasing the amount of training data to decrease the kernel rank, i.e. reduce thereparametrization issue. We find that as the issue is 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 a nonlinear neural network form continuous trajectories in theparameter space (c.f. Fig. 1). We define that all points along such a trajectory are identical, whichchanges the weight space geometry to be a manifold. Likewise, the parameter changes corresponding
to actual function changes reside on a nonlinear manifold. This is sketched in Fig. 3. Interestingly, theGGN turns out to induce a natural (local) inner product on these nonlinear manifolds, which allows us
to both understand and generalize the linearized Laplace approximation.
44.1

The effective-parameters quotient group

For a nonlinear neural network f : RD × RI → RO , 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 a datapoint x ∈ RI , for any w ∈ RD we define the x-reparameterizations
as the set Rfx (w) = {w′ such that f (w′ , x) = f (w, x)}. Consistently, given a collection of points
T
X ⊆ RI , we call the intersection RfX (w) = x∈X Rfx (w) X -reparameterizations.
Trivially, w ∈ RfX (w) for any choice of X . We next define the subset of X -reparameterizations
which can be obtained via a smooth deformation from w.
Definition 4.2. We say that a piecewise differentiable function γ : [0, 1] → RD is a homotopy of(w, ŵ) if γ(0) = w and γ(1) = w′ . The set of X -smooth-reparameterizations is defined as,∃γ a homotopy of (w, w′ ).
R̄fX (w) = w′ such that
γ(t) ∈ RfX (w) ∀t ∈ [0, 1]
We define the relation ∼ over RD as w∼w′ if w′ ∈ R̄fX (w).
Lemma 4.3. ∼ is an equivalence relation and we can form the quotient group P = RD / ∼ ofeffective parameters. We denote [w] ∈ P the equivalence class of an element w ∈ RD .
This group structure gives a rich mathematical foundation to construct reparameterization invariant
neural networks. Within the group, two effective parameters [w1 ], [w2 ] ∈ P are the same point if
and only if w1 ∼ w2 . This means that all parameters w ∈ [w1 ] gives the same function over X .
4.2

The effective-parameters manifold

Geometry is the mathematical language of invariances. To this end would like to endow the weightspace with a geometric structure such that two weights, w1 and w2 , corresponding to the samefunction, have a distance of zero, i.e.
dist(w1 , w2 ) = 0⇔

w1 ∼ w2 .

(8)

Since the weights generate the same function, we define a metric that measures differences in function
values on the training data. Consider weights w and an infinitesimal displacement ϵ, we then define,dist2 (w, w + ϵ) =NX

∥f (w, xn ) − f (w + ϵ, xn )∥2 = ϵ⊤ GGNw ϵ + O(ϵ3 ),

(9)

n=1

where the last step follows from a first-order Taylor expansion of f 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. it is a Riemannian metric. By integratingover paths, the distance extend to any pair of points and satisfies Eq. 8 (Lee & Lee, 2012).
Watch out! It’s a pseudo-metric. We have already seen that in overparametrized models, the
GGN is rank-deficient, which implies that it is not positive definite. Consequently, it is not a
Riemannian metric but rather a pseudo-Riemannian metric. A pseudo-metric can be a counterintuitive
object: two points w1 and w2 at distance zero may have different pseudo-metrics (f ∗ H)w1 ̸=(f ∗ H)w2 . This is reflected in the Laplace approximation. The covariance prescribed by the Laplaceapproximation is Σŵ = (∇2w L(ŵ) + αI)−1 , where L(w) is shorthand for the training log-likelihood.
The Hessian is exactly the pullback pseudo-metric ∇2w L(ŵ) = (f ∗ H)ŵ , which is not invariant toreparameterizations of the neural network. Specifically, for a reparameterization function g that is
also a diffeomorphism, the change of variable rules states that,∇2w L(g(ŵ)) = ∇w g(ŵ)⊤ ∇2w L(ŵ) ∇w g(ŵ).
|{z}| {z }
Σ−1
−αI
g(ŵ)

(10)

Σ−1
ŵ −αIThis means that, while each parameter w has its well-defined covariance Σw , each equivalence classdoes not have a unique one, since ŵ and g(ŵ) belong to the same equivalence class and Σŵ ̸= Σg(ŵ) .
5

Non-Gaussian likelihoods. The Euclidean distance measure in Eq. 9 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-Newton matrix J⊤w HJw , where H denotes the Hessianof the log-likelihood. For both Gaussian and Bernoulli likelihoods, this Hessian is positive definite,
but e.g. the cross entropy has a rank-deficient Hessian and, thus, induces a pseudo-metric.
An impractical solution. The unfortunate behavior of approximate posteriors assigning different
probabilities to Rthe same function could be rectified by marginalizing over the set of reparameterizations of w, i.e. w′ ∈R(w) q(w′ |D)dw′ . While this construction solves the highlighted problem, its
complexity makes it impractical and we are unaware of any works along these lines.
When restricted to a smaller class of reparameterization (the ones homotopic to the identity), theintegral can be thought of as “collapsing” each reparameterization equivalence class to a single point
in P = RD / ∼ formalized in Lemma 4.3. Nontrivially, the pullback metric implicitly performs asimilar operation, as shown later in Theorem 4.5. This connection motivates the dive into Riemannian
geometry: we get a tractable approach to engaging with neural network reparameterizations.
4.3

Topological equivalence of the two views

So far we described two a priori very different objects: the quotient group P = RD / ∼ and thepseudo-Riemannian manifold (RD , GGNw ). We referred to both of them as effective parameters andthis is no coincidence as there is a one-to-one pairing between the two sets.
Proposition 4.4. There exists a bijection between (RD , GGNw ) and P.
Even better, this bijection is a homeomorphism, thus equating the topological properties of the two
spaces. To state that we first need to define a notion of distance on the quotient group and the mostnatural choice is to inherit the Euclidean distance ∥ · ∥ from RD . This distance is defined asdP ([w], [w′ ]) = inf {∥p1 − q1 ∥ + . . . + ∥pn − qn ∥}where the infimum is taken over all finite sequences p1 , . . . , pn and q1 , . . . , qn such that [w] = [p1 ],[pi+1 ] = [qi ] and [qn ] = [w′ ].
Theorem 4.5. The pseudo-Riemannian manifold obtained with the pullback pseudo-metric(RD , GGNw ) is homeomorphic to the quotient group (P, dP ) of reparameterization.
This result connects an abstract quotient group P with the pseudo-Riemannian metric GGNw . The
quotient group captures useful intuitions but is difficult to leverage computationally. In contrast, thepseudo-metric has some counterintuitive aspects but we can identify the underlying Riemannian
structure which leads to tractable algorithms (Sec. 5).
A tale of two manifolds. For any given parameter w ∈ RD and training set X , we show that there⊥
exist two Riemannian manifolds (Pw , m) and (Pw, m⊥ ) embedded in RD , illustrated in Fig. 3. They
capture the functional change and reparameterization properties respectively, but, differently from thepreviously studied (RD , GGNw ), they are Riemannian manifolds without degenerate directions intheir metrics. 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 Pw
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 Pw⊆ [w], thus
all the parametrized functions are identical on the train set; in contrast, (2) Pw never intersects thesame equivalence class more than one time, at least locally, thus the parametrized functions always
change when moving in any direction. Thus Pw resembles the effective parameter manifold P, butwith the difference of being an actual Riemannian manifold. These two manifolds exist under theassumption that Jacobian is full rank (see proof in appendix C).

The two metrics m and m⊥ are not uniquely defined. A natural choice for m is to restrict GGNw to⊥
the tangent space of Pw corresponding to the non-zero eigenvectors, i.e. m = GGN+
w . While, for m⊥
we can inherit the Euclidean metric, i.e. m = αI for α > 0, since the restricted GGN is zero.
6

MNIST

Laplace Diffusion (ours)
Sampled Laplace
Linearised Laplace

Conf. (↑)
0.988±0.001
0.589±0.008
0.968±0.004

NLL (↓)
0.042±0.007
3.812±0.284
0.306±0.041

Acc. (↑)
0.987±0.002
0.146±0.032
0.926±0.008

Brier (↓)
0.022±0.003
1.176±0.046
0.117±0.012

ECE (↓)
0.137±0.019
0.443±0.026
0.251±0.034

MCE (↓)
0.775±0.043
0.985±0.002
0.855±0.041

CIFAR-10 FMNISTTable 1: In-distribution performance across methods trained on MNIST, FMNIST and CIFAR-10.

Laplace Diffusion (ours)
Sampled Laplace
Linearised Laplace

0.900±0.001
0.618±0.021
0.897±0.003

0.001±0.000
4.507±0.000
0.423±0.000

0.906±0.007
0.098±0.010
0.862±0.005

0.141±0.006
1.295±0.014
0.207±0.006

0.108±0.015
0.518±0.013
0.147±0.017

0.729±0.092
0.986±0.001
0.756±0.048

Laplace Diffusion (ours)
Sampled Laplace
Linearised Laplace

0.952±0.007
0.843±0.004
0.951±0.007

0.345±0.062
0.997±0.222
0.614±0.020

0.905±0.007
0.717±0.049
0.863±0.001

0.155±0.019
0.422±0.081
0.222±0.002

0.259±0.008
0.221±0.047
0.337±0.022

0.870±0.021
0.804±0.080
0.789±0.035

Table 2: Out-of-distribution AUROC (↑) performance for MNIST, FMNIST and CIFAR-10.
Trained on
Tested on
Laplace Diffusion (ours)
Sampled Laplace
Linearised Laplace

5

MNIST

FMNISTCIFAR-10

FMNIST

EMNISTKMNIST

MNISTEMNISTKMNISTCIFAR-100

SVHN

0.909±0.033
0.500±0.026
0.758±0.070

0.625±0.018
0.494±0.006
0.602±0.027

0.929±0.008
0.482±0.013
0.790±0.018

0.759±0.045
0.495±0.037
0.625±0.050

0.741±0.010
0.503±0.036
0.628±0.013

0.749±0.023
0.493±0.033
0.624±0.020

0.851±0.002
0.687±0.033
0.837±0.006

0.862±0.010
0.599±0.038
0.854±0.024

Exploring manifolds with random walks

SDEs on manifolds. Given a 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),
dw =

√

1

2τ G(w)− 2 dW + τ Γdt

where Γi (w) =D
X∂(G(w)−1 )ij .
∂w
j
j=1

(11)

Practically speaking this simple process can be simulated using an Euler–Maruyama (Maruyama,1955) scheme. The Christoffel symbols, Γi (θ), are commonly disregarded as they have a highcomputational cost, and Li et al. (2015) showed that the resulting error is bounded.
Using the Euler–Maruyama integrator with step size ht , setting τ = 1 corresponding to standard
Bayesian inference and disregarding√the term involving the Christoffel symbols Γ, we obtain the1
simple update rule wt+1 = wt + 2ht G(wt )− 2 ϵ, where ϵ ∼ N (0, I). This applies to any
Riemannian manifold. However, the effective-parameter (RD , GGNw ) is only pseudo-Riemannian,⊥
we explore three Riemannian alternatives: (RD , GGNw + αI), (Pw, αI) and (Pw , GGN+ω)
Diffusion on (RD , GGNw + αI). The Laplace approximation can also be written as a diffusion ona manifold. As we saw in Sec. 2, the Laplace approximation can be written as w|D ∼ N (ŵ, Σ) with
Σ−1 = GGNŵ + αI. This can also be written as a sample at t = 1 of a Riemannian diffusion on a1
manifold with a constant metric, (RD , G), where G = GGNŵ + αI. The SDE dw = G− 2 dW havea marginal distribution at t = 1 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 Γ as these are zero for constant metrics. Hence, the above is exactly a
Riemannian diffusion on the manifold with a constant metric given by the GGN at the MAP parameter.
Note that this is only a valid diffusion for α > 0 in which case it is not reparametrization invariant.
⊥
Kernel-manifold diffusion. The kernel-manifold (Pw, αI) consists of parameters that generate
the same function over the training set. The effect of diffusion on this manifold and using theneural network predictive is similar to sampling from the kernel subspace while using the linearized
predictive. On the training set the predictive variance is 0 because it only samples reparametrizationsof the MAP predictions. On out-of-distribution data, the variance is greater than 0 if at least one of
the reparameterizations on the training set is not a global reparameterization. This leads to a clear
separation in the predictive variance of in-distribution and out-of-distribution data (Fig. 2) and further
implies that this diffusion distribution never underfits. Stated formally,

7

Theorem 5.1. Varw∼Pŵ⊥ [f (w, x)] = 0 for train data x ∈ X . For a test point xt ̸∈ X , if there exists
a reparameterization w′ ∈ R̄fX (ŵ) such that w′ ̸∈ R̄fX ∪{xt }(ŵ), then Varw∼Pŵ⊥ [f (w, xt )] > 0.
Non-kernel-parameter manifold diffusion. The non-kernel-parameter 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 is the only considered diffusion that acts on a Riemannian manifold while being reparametrization
invariant, i.e. R̄fX (w) = {w}. We call this Laplace diffusion and study it empirically in Sec. 7.

6

Related work

Bayesian deep learning techniques are still in their infancy and generally involve poorly understood approximations. The arguably most popular tool for uncertainty quantification is ensembles(Lakshminarayanan et 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 et 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), e.g. last layer approximations (Kristiadi et al., 2020), subnetwork
inference (Daxberger et al., 2021c), subspace inference (Izmailov et al., 2020) or even PCA in weightspace (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 alternative, and less heuristic, path forward.
MacKay (1998) noted the importance of the choice of basis in Laplace approximations; our pseudoRiemannian view can be seen as having a continuously changing basis. Kristiadi et al. (2023) studied
how a metric transforms under a bijective differentiable change of variables. They enforce geometricconsistency, highlighting, e.g., the non-invariance of the GGN to a 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) reparametrizations 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 a hyperplane in the kernel manifold.
Kim et al. (2024) also study the kernel of the GGN in the context of influence functions. These
works characterize subsets of the reparametrization group. We provide the first architecture-agnostic
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 a 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, our work focuses on analyzing the effectivenessof linearized Laplace within the context of neural network reparametrization, instead of primarily
aiming to achieve better posterior approximations. This allows us to gain deeper insights into theunderlying mechanisms that contribute to the success of this approximation technique.

7

Experiments

We benchmark Laplace diffusion with neural network predictive against linearized and sampled
Laplace to validate the developed theory. Implementation details are in Appendix E.1. We will showthat the diffusion posterior slightly outperforms linearized Laplace in 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
Code: https://github.com/h-roy/geometric-laplace.

8

Figure 5: Benchmark results for Rotated MNIST (similar results for FM NIST and CIFAR are in appendix E.3.2).
Sampled Laplace significantly underfits even for non-rotated data. Laplace
diffusion consistently outperforms theother methods.
In -d istr ib u tion m e tr ics

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 thatthe developed theory explains the key challenges of Bayesian deep learning.
Experimental details (appendix E.3). We train a 44,000-parameter LeNet(LeCun et al., 1989) on
MNIST and FMNIST as well as a 270,000-parameter 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 andthe neural network predictive, while for diffusion samples, we only consider the neural network
predictive. These baselines were chosen to be as similar as possible to our approach to ease thecomparison. 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: Confidence, 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 to asses model-calibration and model fit under distribution shift.
Fig. 5 plots negative log-likelihood (NLL) and expected calibration error (ECE) against the degreesof rotation. Laplace diffusion improves on other Laplace approximations.
Out-of-distribution detection (Table 2). On out-of-distribution 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 deeplearning. Techniques that work well in shallow models struggle with deep ones even if they remaincomputationally tractable. This suggests that overparametrization plays a negative role in Bayesianmodels. 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 reparameterization 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 a pseudo-Riemannian structure on
the parameter space that respects the topology of the reparameterization group. This implies that we
can use pseudo-Riemannian probability distributions as approximate posteriors, and we experimented
with the obvious choice of a geometric diffusion process. We also showed that the state-of-the-art
linearized Laplace approximation can be viewed as a 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, it increases runtime by a constant factor. Common Laplaceapproximations do not sample according to the GGN but rather approximate this matrix with adiagonal or block-diagonal matrix. Mathematically, such approximations break the motivational9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 a research grant (42062) from VILLUM FONDEN. This project received
funding from the European Research Council (ERC) under the European Union’s Horizon 2020
research and innovation programme (grant agreements 757360 and 101123955). The work was partly
funded by the Novo Nordisk Foundation through the Center for Basic Machine Learning Researchin Life Science (NNF20OC0062606). In addition to the ERC (above), PH, MP and LT thank theInternational Max Planck Research School for Intelligent Systems (IMPRS-IS) for support, andgratefully 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: 01IS18039A); and funds
from the Ministry of Science, Research and Arts of the State of Baden-Württemberg.

References
Allen-Zhu, Z., Li, Y., and Song, Z. A convergence theory for deep learning via over-parameterization.
In International conference on machine learning, pp. 242–252. PMLR, 2019.
Antorán, J., Janz, D., Allingham, J. U., Daxberger, E., Barbano, R. R., Nalisnick, E., and HernándezLobato, J. M. Adapting the linearised laplace model evidence for modern deep learning. In
International Conference on Machine Learning, pp. 796–821. PMLR, 2022.
Antorán, J., Padhy, S., Barbano, R., Nalisnick, E., Janz, D., and Hernández-Lobato, J. M. Samplingbased inference for large linear models, with application to linearised laplace, 2023.
Bergamin, F., Moreno-Muñoz, P., Hauberg, S., and Arvanitidis, G. Riemannian laplace approximations for bayesian neural networks. Advances in Neural Information Processing Systems, 36,
2024.
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 in deep neuralnetworks with minimum over-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: 10.1175/1520-0493(1950)078<0001:VOFEIT>2.0.
CO;2. URL https://journals.ametsoc.org/view/journals/mwre/78/1/1520-0493_
1950_078_0001_vofeit_2_0_co_2.xml.
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 in Neural Information Processing Systems, volume 34, pp.
20089–20103. Curran Associates, Inc., 2021a.
Daxberger, E., Kristiadi, A., Immer, A., Eschenhagen, R., Bauer, M., and Hennig, P. Laplace
redux–effortless Bayesian deep learning. In NeurIPS, 2021b.
Daxberger, E., Nalisnick, E., Allingham, J. U., Antorán, J., and Hernández-Lobato, J. M. Bayesian
deep learning via subnetwork inference. In International Conference on Machine Learning, pp.
2510–2521. PMLR, 2021c.
Daxberger, 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 deeplearning, 2022.
10

Devroye, L., Györfi, L., and Lugosi, G. A probabilistic 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 neuralnetworks. In International conference on machine learning, pp. 1675–1685. PMLR, 2019.
George, T., Laurent, C., Bouthillier, X., Ballas, N., and Vincent, P. Fast approximate natural gradientdescent in a 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 B: Statistical Methodology, 73(2):123–214, 2011.
Hansen, L. K. and Salamon, P. Neural network ensembles. IEEE transactions on pattern analysis
and machine intelligence, 12(10):993–1001, 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 Proceedingsof the IEEE conference on computer vision and pattern recognition, pp. 770–778, 2016.
Hsu, E. 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, 2021a.
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.
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.
PMLR, 2020.
Izmailov, P., Nicholson, P., Lotfi, S., and Wilson, A. G. Dangers of bayesian model averaging undercovariate shift. Advances in Neural Information Processing Systems, 34:3309–3322, 2021.
Jacot, A., Gabriel, F., and Hongler, C. Neural tangent kernel: Convergence and generalization inneural networks. In Advances in neural information processing systems, 2018.
Jang, C., Lee, S., Park, F., and Noh, Y.-K. A reparametrization-invariant sharpness measure basedon 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: A flexible method for approximating influence via geometricensemble. Advances in Neural Information Processing Systems, 36, 2024.
Kristiadi, A., Hein, M., and Hennig, P. Being Bayesian, even just a 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.07384, 2023.
11

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 andintegral operators. Journal of Research of the National Bureau of Standards, 45(4), October 1950.
doi: 10.6028/jres.045.026. URL https://hal.science/hal-01712947.
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. and Lee, J. M. Smooth manifolds. Springer, 2012.
Li, C., Chen, C., Carlson, D., and Carin, L. Preconditioned stochastic gradient langevin dynamics fordeep 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. https://uvadlc-notebooks.readthedocs.io/en/
latest/, 2022.
Liu, C., Zhu, L., and Belkin, M. On the linearity of large non-linear models: when and why thetangent 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. J. C. A practical Bayesian framework for backpropagation networks. Neural Computation, 4(3):448–472, 1992.
Maddox, W. J., Izmailov, P., Garipov, T., Vetrov, D. P., and Wilson, A. G. A simple baseline forbayesian 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., Moreno-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 bayesianbinning. In Proceedings of the AAAI conference on artificial intelligence, volume 29, 2015.
Nguyen, Q., Mondelli, M., and Montufar, G. F. Tight bounds on the smallest eigenvalue of the neuraltangent kernel for deep relu networks. In International Conference on Machine Learning, pp.
8119–8129. PMLR, 2021.
Oymak, S. and Soltanolkotabi, M. Overparameterized nonlinear learning: Gradient descent takes theshortest path? In International Conference on Machine Learning, pp. 4951–4960. PMLR, 2019.
Oymak, S. and Soltanolkotabi, M. Toward moderate overparameterization: Global convergenceguarantees for training shallow neural networks. IEEE Journal on Selected Areas in Information
Theory, 1(1):84–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 oflarge-scale ai. arXiv preprint arXiv:2402.00809, 2024.
12

Petzka, H., Adilova, L., Kamp, M., and Sminchisescu, C. A reparameterization-invariant flatness
measure for deep neural networks. arXiv preprint arXiv:1912.00058, 2019.
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 on 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

Notation. Consider a function f : RD × RI → RO with Jacobian Jw (x) = ∂w fw (x)|w=w ∈RO×D with respect to w evaluated in x and w. For a given log-likelihood we define the Hessian w.r.t.
to the output Hw (x) = −∂f2w (x) log p(y|fw (x)) ∈ RO×O and we assume it not to be dependent ony (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 tothe collections of the datapoints X = {x1 , . . . , xN } ⊂ RI .
Consider the stacking of the per-datum Jacobians Jw = [Jw (x1 ); . . . ; Jw (xN )] ∈ RN O×D , where
we dropped the dependence on X . Similarly consider Hw = diag(Hw (x1 ); . . . ; Hw (xN )) ∈RN O×N O the block diagonal stacking of the Hessians.
Consider the Generalized Gauss-Newton (GGN) and the Neural Tangent Kernel (NTK) matrices⊤

GGN w = Jw Hw Jw ∈ R

D×D

1/2

⊤

1/2

NTK w = Hw Jw Jw Hw∈ RN O×N O .

(12)Recall also that the pullback pseudo-metric is (f ∗ H)w = GGNw .
Assumptions. We assume uniform upper and lower bound on the eigenvalues of the NTK matrix,
that is∃l, L ∈ R

such that0<l≤∥ NTKw v∥≤ L ∀v ∈ RN O , ∀w ∈ RD ,
∥v∥

(13)

where uniform means uniform over parameters, i.e. the bounds l, L holds for every w. Moreover weassume that the Jacobian function w 7→ Jw (x) is Lipschitz for every x ∈ X .
How unreasonable are the assumptions? The assumption of an upper bound L is equivalent toassuming that w 7→ f (w, x) is Lipschitz for each datapoint x ∈ X . This is true with all standard
activation functions if we restrict the parameter space to a ball of fixed radius.
The assumption of a lower bound l is strongly supported by the literature on NTK, thanks to its
direct implications on memorization capacity and generalization. The general trend is that the more
overparametrized the network is, the stronger such lower bounds are. In the hardest setting ofminimum overparametrization, Bombari et al. (2022) proved a bound that holds with high probability
for fully connected MLPs at initialization. Similar results hold with bigger overparametrizations
Nguyen et al. (2021); Allen-Zhu et al. (2019); Du et al. (2019). Building on top of that, other lines ofwork Liu et al. (2020); Oymak & Soltanolkotabi (2019, 2020) proved that the NTK does not change
too much during training thanks to the PL inequality framework, and in particular the proximity of
the neural network dynamics to the one described by NTK, is supported by spectral bounds on the
Hessian of the landscape.
Lastly, the Lipschitzness assumption on the Jacobian is potentially the most unrealistic, although
it would hold, for example, if the derivatives of activations are Lipschitz and the parameters are
restricted to a finite radius ball. Nonetheless, we emphasize that this assumption is only used tocontrol that the kernel of the GGN does not “rotate too fast”, which is a much weaker assumption, butalso much more cluttered to state formally.

B

Proof for equivalence of the two settings

This section contains the proof of Proposition 4.4 and Theorem 4.5 involving the pseudo-Riemannian
manifold (RD , f ∗ H) and the quotient group (P, dP ) with Euclidean-induced metric.
To ease the readability, we recall the two involved notions of distance and their respective definition:• df ∗ H (w0 , w1 ) the geodesic distance for any two parameters w0 , w1 ∈ RD
• dP ([w0 , w1 ]) the quotient Euclidean distance for any two equivalence classes [w0 ], [w1 ] ∈
P
14

(14)

df ∗ H (w0 , w1 ) = inf LENf ∗ H (γ)γ
Z 1∥γ ′ (t)∥f ∗ Hγ(t) dt
= inf
γ

= inf
γ

(15)

0

Z 1q
γ ′ (t)⊤ · f ∗ Hγ(t) · γ ′ (t)dt,

(16)

0

where the infimum is taken over smoothly differentiable curves γ : [0, 1] → RD such that γ(0) = w0
and γ(1) = w1 .
( n)
XdP ([w0 ], [w1 ]) = inf∥pi − qi ∥ ,
(17)
i=1

where the infimum is taken over all finite sequences {pi }i=1...n , {qi }i=1...n ⊂ RD such that [w0 ] =[p1 ], [pi+1 ] = [qi ] and [qn ] = [w1 ].
Let us state a more quantitative theorem
Theorem B.1. For any w0 , w1 ∈ RD it holds⇐⇒

df ∗ H (w0 , w1 ) = 0

dP ([w0 , w1 ]) = 0,(18)

and also that, for any ϵ > 0 there exists δ > 0 such thatdP ([w0 , w1 ]) < δ
df ∗ H (w0 , w1 ) < δ=⇒=⇒

df ∗ H (w0 , w1 ) < ϵdP ([w0 , w1 ]) < ϵ,

(19)
(20)

and note that by definition of bijection Eq. 18 is equivalent to Proposition 4.4 and, by definition of
homeomorphism, Eq. 19 and Eq. 20 are equivalent to Theorem 4.5.
We prove the 3 points separately, in Appendix B.1, Appendix B.2 and Appendix B.3 respectively.
B.1

Proof of Eq. 18

The proof logic is
df ∗ H (w0 , w1 ) = 0⇐⇒[w0 ] = [w1 ]⇐⇒dP ([w0 , w1 ]) = 0.

(21)

and we prove the two steps in the two following Propositions, respectively.
Proposition B.2. For any w0 , w1 ∈ RD it holds⇐⇒

df ∗ H (w0 , w1 ) = 0[w0 ] = [w1 ].

(22)

Proof. By definition [w0 ] = [w1 ] if and only if there exists a piecewise differentiable γ : [0, 1] → RDsuch that γ(0) = w0 , γ(1) = w1 and f (w0 , x) = f (γ(t), x) for any t ∈ [0, 1] and x ∈ X . Then⇐= Consider a γ from the definition of the equivalence relation ∼ and define the points wt = γ(t)for ease of notation. Then for any t ∈ [0, 1]
wt+ϵ − wt.
(23)
ϵFor any x ∈ X it holds that f (wt , x) = f (wt′ , x)∀t, t′ ∈ [0, 1], which implies that f (wt+ϵ , x) −f (wt , x) = 0 ∀t ∈ [0, 1]∀ϵ ∈ [0, 1 − t]. Thus,
γ ′ (t) = lim

ϵ→0f (wt+ϵ , x) − f (wt , x)wt+ϵ − wt= Jwt (x) · lim= Jwt (x) · γ ′ (t).
(24)
ϵ→0
ϵ
ϵ
This holds to any x ∈ X , so the same holds for the per-datum stacked jacobians Jwt · γ ′ (t) = 0.
Thus,′∥γ ′ (t)∥2f ∗ Hγ(t) = γ ′ (t)⊤ · f ∗ Hγ(t) · γ ′ (t) = γ ′ (t)⊤ · J⊤(25)
wt Hwt Jwt · γ (t) = 0
0 = limϵ→0and we can measure the length of γ in the pullback metric as
Z 1
Z 1
LEN f ∗ H (γ) =∥γ ′ (t)∥f ∗ Hγ(t) dt =0 dt = 0,
0

0

15

(26)

which gives an upper bound on the geodesic distance(27)

df ∗ H (w0 , w1 ) = inf LENf ∗ H (γ̂) ≤ LENf ∗ H (γ) = 0,γ̂thus, df ∗ H (w0 , w1 ) = 0 and this implication is proven.
=⇒ df ∗ H (w0 , w1 ) = 0 implies that there exists a 0-length differentiable γ : [0, 1] → RD suchthat γ(0) = w0 , γ(1) = w1 . Without loss of generality, we can assume γ to be non-stationary, i.e.
γ ′ (t) ̸= 0. Here 0-length means
Z 1
0 = LENf ∗ H (γ) =∥γ ′ (t)∥f ∗ Hγ(t) dt,(28)0

which implies that ∥γ ′ (t)∥f ∗ Hγ(t) = 0 for any t ∈ [0, 1] except a zero-measure set which we canneglect later. Then′0 = ∥γ ′ (t)∥2f ∗ Hγ(t) = γ ′ (t)⊤ · J⊤
wt Hwt Jwt · γ (t)=⇒Jwt · γ ′ (t) = 0,(29)

by positive definitess of Hwt , assumed as hypothesis. We highlight that the leftmost 0 in the previous
equation is a scalar, while the rightmost 0 is a vector in RN O as obtained by the matrix-vector productof Jwt ∈ RN O×D with γ ′ (t) ∈ RD . Looking at the equation Jwt · γ ′ (t) = 0 componentwise implies
that⟨∇w [f (wt , x)]o , γ ′ (t)⟩ = 0 ∀x ∈ X , ∀o ∈ {1, . . . , O}, ∀t ∈ [0, 1],(30)
where [v]o refers to the oth component of a vector v. Thus, for T ∈ [0, 1], by Fundamental Theoremof Calculus we have
Z T
[f (wT , x)]o − [f (w0 , x)]o =⟨∇w [f (wt , x)]o , γ ′ (t)⟩dt = 0∀x ∈ X , ∀o ∈ {1, . . . , O}.
0(31)
Then f (wT , x) = f (w0 , x) ∀x ∈ X and ∀T ∈ [0, 1]. So we proved that γ is an homotopy of (w, w′ )such that γ(t) ∈ RfX (w) ∀t ∈ [0, 1]. Thus,w′ ∈ R̄fX (w)=⇒w′ ∼ w[w′ ] = [w],

=⇒

(32)

and this completes the proof.
Proposition B.3. For any w0 , w1 ∈ RD it holds[w0 ] = [w1 ]

⇐⇒

(33)

dP ([w0 , w1 ]) = 0.

Proof. =⇒ This arrow is trivially true by considering the two sequences in the definition of thequotient distance to be of length one and such that p1 = w0 and q1 = w1 .
⇐= dP ([w0 , w1 ]) = 0 implies that, by definition of inf, there exists a sequence ϵm → 0 such that
(m)
(m)
(m)
(m)
∀m ∈ N there exists two finite sequences of points p1 , . . . , pn ∈ RD and q1 , . . . , qn ∈ RD
such that
n
X
(m)
(m)
∥pi − qi ∥ = ϵm ,
(34)
(m)

(m)

∥p − q∥ < ϵ

=⇒

i=1
(m)

(m)

where [w0 ] = [p1 ], [pi+1 ] = [qi ] and [qn ] = [w1 ]. Note that the sequence length n may
depend on m.
Lipschitzness of f in parameters, assumed by hypothesis, means that
thus,
n
X

(m)

∥pi

(m)

− qi

∥ = ϵm∥f (p, x) − f (q, x)∥ < Lϵ

=⇒

i=1

n
X

(m)

∥f (pi∀p, q ∈ RD , ∀x ∈ X , ∀ϵ > 0,

(35)

(m)

(36)

, x) − f (qi, x)∥ < 2Lϵm

∀x ∈ X .

i=1

Also, by definition of the equivalence class, it holds ∀x ∈ X that(m)[w0 ] = [p1 ](m)(m)[pi+1 ] = [qi ][qn(m) ] = [w1 ]

=⇒
=⇒
=⇒

(m)

∥f (w0 , x) − f (p1 , x)∥ = 0

(37)

(m)
(m)
∥f (pi+1 , x) − f (qi , x)∥ = 0∥f (qn(m) , x) − f (w1 , x)∥ = 0.

(38)

16

(39)

Thus, by the triangular inequality, the last four equations imply
(m)

(m)

∥f (w0 , x) − f (w1 , x)∥ ≤ ∥f (w0 , x) − f (p1 , x)∥ + ∥f (p1 , x) − f (qn(m) , x)∥(40)+ ∥f (qn(m) , x) − f (w1 , x)∥(m)= ∥f (p1 , x) − f (qn(m) , x)∥
≤

n
X

(m)

∥f (pi(m), x) − f (qi(41), x)∥ +

i=1

n−1
X

(m)

(m)

∥f (pi+1 , x) − f (qi, x)∥ (42)i=1

< 2Lϵm +

n−1
X

(43)

0 = 2Lϵm ,i=1

and this holds for any m ∈ N. Taking the limit m → ∞, ϵm → 0 implies ∥f (w0 , x)−f (w1 , x)∥ ≤ 0.
Thus, f (w0 , x) = f (w1 , x) ∀x ∈ X , and, thus, [w0 ] = [w1 ], which completes the proof.
B.2

Proof of Eq. 19

In order to prove Eq. 19, we first prove a weaker statement
Proposition B.4. For any w0 , w1 ∈ RD and for any δ > 0 it holds that∥w0 − w1 ∥ < δ=⇒

df ∗ H (w0 , w1 ) < Lδ.

(44)

Proof. Let γ : [0, 1] → RD be defined as γ(t) = (1 − t)w0 + tw1 , then
Z 1
Z 1
LEN f ∗ H (γ) =∥γ ′ (t)∥f ∗ Hγ(t) dt =∥w1 − w0 ∥f ∗ Hγ(t) dt
0

0

Z 1
≤

(45)

L∥w1 − w0 ∥dt = L∥w1 − w0 ∥,0and thus,
df ∗ H (w0 , w1 ) = inf LENf ∗ H (γ̂) ≤ LENf ∗ H (γ) = L∥w1 − w0 ∥ < Lδ.
γ̂

(46)

Definition B.5. Let γ1 , . . . , γn : [0, 1] → RD a sequence of piecewise differentiable paths such that∀i ∈ {0, . . . , n}.

γi−1 (1) = γi (0)(47)

Consider the piecewise differentiable path that is the concatenation of the paths one after the other,γ̂ = CAT(γ1 , . . . , γn ) defined asγ̂(t) = γi (nt − i)
It is straightforward to see that LEN(γ̂) =if i ≤ t ≤ i + 1.

(48)

Pn

i=1 LEN (γi ).

The choice of ϵ, δ that prove Eq. 19 trivially follows from the following
Proposition B.6. With the assumption of eigenvalue upper bound L, for any w0 , w1 ∈ RD for any
δ > 0 it holds thatdP ([w0 , w1 ]) < δ =⇒ df ∗ H (w0 , w1 ) < 3Lδ.
(49)
Proof. dP ([w0 , w1 ]) < δ implies that, by the definition of inf, there exists two sequences{pi }i=1...n , {qi }i=1...n ⊂ RD such that [w0 ] = [p1 ], [pi+1 ] = [qi ], [qn ] = [w1 ] and such thatnX∥pi − qi ∥ < 2δ.

(50)

i=1

Now the idea is to define n + 1 paths on the equivalence classes and n paths connecting them, thenthe stacking of the 2n + 1 will give an upper bound on the geodesic distance.
Let us first define the paths γ2i for i = 0, . . . , n by making use of Proposition B.2
17• [w0 ] = [p1 ] imply that there exists γ0 : [0, 1] → RD be such that γ0 (0) = w0 , γ0 (1) = p1
and such that LENf ∗ H (γ0 ) < Lδ/n+1• [qn ] = [w1 ] imply that there exists γ2n : [0, 1] → RD be such that γ2n (0) = qn , γ2n (1) =w1 and such that LENf ∗ H (γ2n ) < Lδ/n+1• for all i = 1, . . . , n − 1, [pi+1 ] = [qi ] imply that there exists γ2i : [0, 1] → RD be such that
γ2i (0) = pi+1 , γ2i (1) = qi and such that LENf ∗ H (γ2i ) < Lδ/n+1
Pn
D
And then for all i = 1, . . . , n, Proposition B.4 and i=1 ∥pPin− qi ∥ < 2δ imply γ2i−1 : [0, 1] → R
be such that γ2i−1 (0) = pi , γ2i−1 (1) = qi and such that i=1 LENf ∗ H (γ2i−1 ) < 2Lδ.
Then concatenating these 2n + 1 paths, we have(51)

df ∗ H (w0 , w1 ) = inf LENf ∗ H (γ̂) ≤ LENf ∗ H (CAT(γ0 , . . . , γ2n ))
γ̂

=

2n
X
i=0

B.3

LEN f ∗ H (γi ) =n
X

LEN f ∗ H (γ2i−1 ) +

i=1

n
X

LEN f ∗ H (γ2i ) < 2Lδ + Lδ. (52)i=0Proof of Eq. 20

Proposition B.7. The pullback metric does not “rotate” too much when moving from w to w + ϵ,formally(1 − K∥ϵ∥)∥v∥f ∗ Hw+ϵ ≤ ∥v∥f ∗ Hw ≤ (1 + K∥ϵ∥)∥v∥f ∗ Hw+ϵ .
(53)
Proof. Follows from the K-Lischitz assumption on w 7→ Jw
The choice of ϵ, δ that prove Eq. 20 trivially follows from the following
Proposition B.8. With the assumption of eigenvalue upper bound L, for any w0 , w1 ∈ RD for any
δ > 0 it holds that6δ
df ∗ H (w0 , w1 ) < δ =⇒ dP ([w0 , w1 ]) <.
(54)
lProof. df ∗ H (w0 , w1 ) < δ implies that there exist a path γ : [0, 1] → RD withZ 1
2δ = LENf ∗ H (γ) =∥γ ′ (t)∥f ∗ Hγ(t) dt.

(55)

0

Consider also the Euclidean length LENE (γ) which does not depend on δ. Then, for any δ there
exist n(δ) ∈ N such that, considering the uniform partition of [0, 1], ti = i/n(δ) for i = 0, . . . , n(δ) it
holds the discrete approximation of the integraln(δ)−13δ >Xn(δ)−1∥γ(ti+1 ) − γ(ti )∥f ∗ Hγ(ti ) =i=0X∥γ(ti+1 ) − pi + pi − γ(ti )∥f ∗ Hγ(ti )

(56)

i=0
n(δ)−1=X
i=0∥γ(ti+1 ) − pi ∥f ∗ Hγ(ti ) + ∥pi − γ(ti )∥f ∗ Hγ(ti ) ,
|
{z
}
=0

(57)

where pi is defined for every i ∈ {0, . . . , n(δ) − 1} as follow. Consider the projection PiK on the
kernel of f ∗ Hγ(ti ) , and the projection PiI on the image of f ∗ Hγ(ti ) , such that the two projections
are orthogonal and PiK + PiI = I. Define the point pi ∈ RD as pi = PiK (γ(ti+1 ) − γ(ti )) + γ(ti ),which implies that ∥pi − γ(ti )∥f ∗ Hγ(ti ) = 0. By definition of the projections, it also hold thatpi = γ(ti+1 ) − PiI (γ(ti+1 ) − γ(ti )), which implies that γ(ti+1 ) − pi = PiI (γ(ti+1 ) − γ(ti )) is18

aligned with the non-zero eigenvalues of f ∗ Hγ(ti ) , this means that we can resort to the lower bound lon the non-zero eigenvalues and Proposition B.7 to see that∥γ(ti+1 ) − pi ∥f ∗ Hγ(ti ) ≥ (1 − K∥γ(ti+1 ) − γ(ti )∥)∥γ(ti+1 ) − pi ∥f ∗ Hγ(ti+1 )
{z
}
|

(58)

K̃

≥ K̃l∥γ(ti+1 ) − pi ∥l≥ ∥γ(ti+1 ) − pi ∥2(59)
∀i ∈ {0, . . . , n(δ) − 1},(60)

where, recalling that ∥γ(ti+1 ) − γ(ti )∥ = 1/n(δ) by construction, there always exist an n(δ) big
enough such that K̃ ≥ 1/2, which we can assume without loss of generality. Rearranging the terms,
we haven(δ)−1
X
6δ
≥∥γ(ti+1 ) − pi ∥.
(61)
l
i=0
Finally, we can define the points qi = γ(ti+1 ) and we have the two sequences p0 , . . . , pn(δ)−1 and
q0 , . . . , qn(δ)−1 whose satisfies the costrains [w0 ] = [p0 ], [pi+1 ] = [qi ] and [qn(δ)−1 ] = [w1 ]. We
highlight that the dependence on δ of the length of the sequence is not problematic, as it is sufficientthat n(δ) is finite for every fixed δ, which it is. Thus,n(δ)−1dP ([w0 , w1 ]) ≤X

∥pi − qi ∥ <

i=0

6δ
,
l

(62)

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⊥(Pw̄, m⊥ ) embedded in RD . In the rest of the section, we consider a fixed w̄ ∈ RD .
Fix a training set X = {x1 , . . . , xN } ⊂ RI of size N and consider the stacked partial evaluation of
the network defined as
(63)
(64)

F : RD −→ RN O
w 7−→ (f (w, x1 ), . . . , f (w, xN )) .

And define ȳ = F(w̄) ∈ RN O . The differential ∇w F w̄ = Jw̄ ∈ RN O×D equals the stacking of
the per-datum Jacobians, and we assume it to be full rank thanks to the uniform lower bound on the
eigenvalues of the NTK matrix. Full-rankness in the overparametrized setting D > N O implies thesurjectivity of the differential operator.


∂Fi
Jw̄ =
.
(65)
∂wj w̄ i=1,...,N O; j=1,...,D
We reorder the wi so that the first N O columns are independent. Then the N O × N O matrix


∂Fi
R=
(66)
∂wj w̄ i=1,...,N O; j=1,...,N O
is non-singular. We consider the map
α(w1 , . . . , wD ) = (F(w)1 , . . . , F(w)N O , wN O+1 , . . . , wD ) .
We obtain∇w α w̄ =∂αi∂wj w̄


=i=1,...,D; j=1,...,D

R
0


∗
,
I

(67)
(68)

and this is non-singular. By the inverse function theorem, α is a local diffeomorphism. So there is an
open W ⊆ RD containing w̄ such that α|W : W → α(W ) is smooth with smooth inverse.
19Finally, define
⊥
Pw̄
=





α−1 (ȳ1 , . . . , ȳN O , p1 , . . . , pD−N O )|{z}for p ∈ RD−N O




(69)

⊆ RD ,

NOand similarly
Pw̄ =





α−1 (p1 , . . . , pN O , 0, . . . , 0)
| {z }

for p ∈ RN O




(70)

⊆ RD .



D−N O⊥
We claim that the two restrictions of α are slice charts of Pw̄
and Pw̄ , respectively. Since it is a⊥smooth diffeomorphism, it is certainly a chart. Moreover, by construction, the points in Pw̄
are⊥exaclty those whose image under F is ȳ, thus, Pw̄ ⊆ [w̄]. On the other hand, the points w ∈ Pw̄
parametrize functions that take the values F(w) = p in a local neighbourhood of ȳ. Thus, locally it
never intersects the same equivalence class more than one time.

D

Proof of Theorem 5.1⊥′
By definition of the kernel manifold, we have that if w ∈ Pw
′ then we have that w ∈ [w ]. Hence⊥for all w ∈ Pw′ and for all x ∈ X
f (w, x) = f (w′ , x).
(71)
It follows that Varw∼P ⊥′ [f (w, x)] = 0 for any x ∈ X .
w

For the second statement notice that Varw∼P ⊥′ [f (w, xtest ] = 0 if and only if f (w, xtest ) = c for⊥
all w ∼ Pw
′ and for some constant c.

w

Suppose ŵ ∈ R̄fX and ŵ ̸∈ R̄fX ∪{xtest } , then f (ŵ, xtest ) ̸= f (w′ , xtest ), which means thatf (w, xtest ) is not constant. Hence we have that Varw∼P ⊥′ [f (w, xtest ] > 0w

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 ofsize D × D, where D is the number of parameters. For most models this problem is intractable. The
standard approach to this problem is to consider sparse approximations to the Hessian of the loss
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, we choose to sample from Laplace’s approximation in a way that isclosest to the theoretical ideal, at the cost of performing expensive computations.

4000

3000

Eigenvalues

In small experiments with the toy regression problem, we
instantiate the exact GGN and compute vector productswith 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 eigenvalues (Figure 6)

2000

This makes low-rank approximations of the GGN particularly attractive. We choose the Lanczos algorithm (Lanczos, 1950) with full reorthogonalization and run it for avery high number of iterations, to ensure numerical stability and very low reconstruction error to form our low-rank Figure 6: Eigenvalues of the GGN of aapproximations of the GGN. Additionally, Lanczos only Convolutional Neural Network trained
requires implicit access to the matrix so we avoid the mem- on MNIST.
ory cost and GGN vector products for neural networks canbe performed efficiently using Jacobian-Vector products and Vector-Jacobian products. If we do
1000

0

0

50

100

150

200

Index

20

250

300

350

Algorithm 1 Laplace diffusion
1: Input: Observed data D, trained MAP point w′ , number of steps T , number of samples S, rankk.

0
2: Initialize samples w10 . . . wS
as the MAP estimate w′
3: for j in 1, . . . S do
4:
for t in 1, . . . T do
5:
Sample ϵ ∼ N (0, Ik )6:
Compute the top-k eigenvalues(Λtj ) and eigenvectors(Ujt ) of GGNwjt
1

7:
wjt ← wjt + √1N Ujt (Λtj + αI)− 2 ϵ8:
end for
9: end forT
10: Return posterior samples w1T . . . wS.

a sufficiently large number of iterations we obtain the non-zero eigenvalues Λ̃, and corresponding 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 Λ̃, and corresponding eigenvectors U1 we can also form inverse
vector products with the square root of GGN + αI. It should be evident from the discussion insection 3 about decomposing the covariance that this vector product with a vector v is given by:
1
1
1(GGN + αI)− 2 v = U1 (Λ̃ + αIk )− 2 v + √ U2 v
α1
1
= U1 (Λ̃ + αIk )− 2 v + √ (I − U1 )vα
1
1
1= U1 ((Λ̃ + αIk )− 2 − √ Ik )v + √ v
α
α

(72)
(73)
(74)

This allows us to form inverse-square root vector products with the GGN + αI given the non-zero
eigenspectrum. We run the sampling algorithm on H100 GPUs to run the high-order Lanczos decomposition. This approach of sampling from Laplace’s approximation has O(pk 2 ) time complexity, andO(pk) memory cost, where k is the number of Lanczos iterations and p is the number of parameters.
E.2
E.2.1

Experimental details and further results for toy experiments
Toy regression in Figure 2

In this experiment, we fit a small MLP, with 2 hidden layers of width 10 on the sine curve. Due to thesmall size of the GGN it is possible to instantiate and do all the computations explicitly. We sample
from the exact Laplace’s approximation, the non-kernel and kernel subspace of the GGN, and use theneural 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 we do alternating steps in the two manifolds for the left
plot which gives us the full distribution.
We also do a similar experiment to show that the same phenomenon also holds for classification. We
use a small convolutional neural network, with 2 convolutional layers with kernel of size 3, to classify
a 2-class mixture of Gaussians and look at uncertainties of sampled Laplace, linearized Laplace, andLaplace’s diffusion. We decompose these uncertainties into their kernel and 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 is that sampled Laplace underfits in-distribution and thiseffect is related to the kernel component of the distribution.
E.2.2

Effect of kernel rank on in-distribution fit in figure 4

For this experiment, we train a small convolutional neural network, with two convolutional layers anda kernel of size 3, on MNIST. In this case, the GGN can be instantiated explicitly. We recall that the21Sampled Full Distribution

Sampled Kernel Distribution

Sampled Non-Kernel Distribution

Linearized Laplace uncertainty

Linearized Kernel uncertainty

Linearized Non-Kernel uncertainty

Laplace Diffusion uncertainty

Diffusion Kenrel uncertainty

Diffusion Non-Kernel uncertainty

5
0
5
5
0
5
5
0
5
6

4

2

0

2

4

6

6

4

2

0

2

4

6

6

4

2

0

2

4

6

0.90
0.84
0.78
0.72
0.66
0.60
0.54
0.48
0.42
1.05
0.90
0.75
0.60
0.45
0.30
0.15
0.00
0.8
0.7
0.6
0.5
0.4
0.3
0.2
0.1
0.0

Figure 7: Decomposition of uncertainties of Laplace Approximation for the Gaussian mixture
classification.
GGN is a sum of Jw′ (xn )⊤ H(xn )Jw′ (xn ) over the dataset where xn are the individual data points.
We recall that the rank of the GGN is bounded by N O 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
us a GGN with a lower rank. Equivalently, this is GGN with a higher dimensional kernel, and hence
the usual covariance from Laplace’s approximation (GGN + αI)−1 has a higher contribution fromthe kernel subspace.

We consider multiple such subsamples and plot the training accuracy for samples from Laplace’s
approximation against the kernel subspace dimension. Here we see a 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 our suggested hypothesis that the underfitting insampled Laplace is caused by its kernel component and is hence deeply related to reparameterizations.
E.3

Additional benchmarks and results for image classification

Training details: We use a standard LeNet for the MNIST and FashionMNIST experiments and asmaller version of ResNet (Lippe, 2022), with 272,378 parameters consisting of 3 times a group of 3
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−3 . For the ReNet we use SGD with a learning rate of 0.1 with momentum and weight decay.
Hyperparameters: We benchmark Laplace diffusion against SWAG, diagonal Laplace, last-layer
Laplace, and MAP in 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.0}. For Laplace diffusion, sampled
Laplace, and Linearised to ensure that the comparison can validate the theory it is preferable to havethe same prior precision for all of these methods. So we only do the grid search to tune the priorprecision for sampled Laplace and use this for all three methods. We keep the hyperparameters forthese 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 Lanczos iterations and we 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 a step size of 0.2, with 5000 Lanczos iterations and we
predict using 20 MC samples. For sampled Laplace andlLinearised Laplace use the same number ofLanczos iterations and MC samples.
For SWAG we use a learning rate of 10−2 with momentum of 0.9 and weight decay of 3e−4 andthe 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 3 models to sample from theposterior.
Last-layer Laplace is the recommended method by Daxberger et al. (2021c) so it should approximate
the best performance one can get using various possible configurations. For the CIFAR-10 experi22

ments, the last layer of ResNet is too large to instantiate the full GGN matrix. So we instead use thelast 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 deterministic. So we 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 1 to benchmark Laplace’s diffusion against various other Bayesian methods. Here
we see that despite using the neural network predictive it is competitive with the best-performing
Bayesian methods whereas Sampled Laplace performs significantly worse.
MNIST
Laplace’s diffusion
Sampled Laplace
Linearised Laplace
SWAG
Last-Layer Laplace
Diagonal Laplace
MAP

Conf. (↑)
0.988±0.001
0.593±0.002
0.967±0.002
0.993±0.001
0.991±0.001
0.963±0.005
0.992±0.002

NLL (↓)
0.040±0.007
3.669±0.157
0.295±0.041
0.032±0.001
0.047±0.003
0.078±0.011
0.042±0.004

Acc. (↑)
0.986±0.002
0.157±0.031
0.930±0.004
0.989±0.002
0.987±0.002
0.976±0.003
0.988±0.000

Brier (↓)
0.022±0.003
1.162±0.044
0.111±0.006
0.019±0.001
0.021±0.001
0.038±0.005
0.021±0.000

ECE (↓)
0.146±0.011
0.436±0.027
0.239±0.032
0.187±0.024
0.198±0.034
0.095±0.004
0.198±0.041

MCE (↓)
0.773±0.050
0.984±0.001
0.862±0.044
0.901±0.041
0.721±0.091
0.692±0.029
0.685±0.091

Conf. (↑)
0.900±0.001
0.618±0.021
0.897±0.003
0.925±0.002
0.914±0.001
0.862±0.009
0.914±0.000

NLL (↓)
0.275±0.016
4.507±0.160
0.423±0.014
0.259±0.004
0.280±0.016
0.323±0.022
0.279±0.015

Acc. (↑)
0.906±0.007
0.098±0.010
0.862±0.005
0.911±0.006
0.901±0.006
0.889±0.010
0.904±0.006

Brier (↓)
0.141±0.006
1.295±0.014
0.207±0.006
0.135±0.004
0.144±0.004
0.165±0.011
0.143±0.004

ECE (↓)
0.108±0.015
0.518±0.013
0.147±0.017
0.152±0.008
0.131±0.003
0.102±0.003
0.125±0.010

MCE (↓)
0.729±0.092
0.986±0.001
0.756±0.048
0.752±0.067
0.673±0.083
0.660±0.051
0.609±0.033

Conf. (↑)
0.948±0.004
0.843±0.004
0.951±0.007
0.942±0.003
0.953±0.004
0.960±0.003

NLL (↓)
0.403±0.007
0.997±0.222
0.614±0.020
0.393±0.004
0.343±0.033
0.333±0.030

Acc. (↑)
0.889±0.005
0.717±0.049
0.863±0.001
0.884±0.002
0.899±0.001
0.913±0.007

Brier (↓)
0.180±0.009
0.422±0.081
0.222±0.002
0.176±0.001
0.155±0.007
0.143±0.009

ECE (↓)
0.264±0.048
0.221±0.047
0.337±0.022
0.234±0.014
0.260±0.000
0.282±0.002

MCE (↓)
0.879±0.038
0.804±0.080
0.789±0.035
0.912±0.016
0.884±0.029
0.932±0.006

FMNIST
Laplace’s diffusion
Sampled Laplace
Linearised Laplace
SWAG
Last-Layer Laplace
Diagonal Laplace
MAP

CIFAR-10Laplace’s diffusion
Sampled Laplace
Linearised Laplace
SWAG
Last-Layer Laplace
MAP

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 isgood 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 is competitive against other Bayesian methods.
E.3.3

Out-of-distribution detection

We extend 2 to benchmark Laplace’s diffusion against various other Bayesian methods for Out-ofDistribution Detection. Once again, we observe that, despite using the neural network predictive,23

12

Laplace's Diffusion
Sampled Laplace
Linearised Laplace
SWAG
Last-Layer Laplace
MAP

3.0

10

2.0

1.5

0.7

0.6

8

Negative log-likelihood

2.5

Laplace's Diffusion posteriors
Sampled Laplace posteriors
Linearised Laplace posteriors
SWAG
Last-Layer Laplace posteriors
MAP

Expected Calibration Error

3.5

6

4

0.5

0.4

0.3

1.0
2

Laplace's Diffusion posteriors
Sampled Laplace posteriors
Linearised Diffusion posteriors
SWAG
Last-Layer Laplace posteriors
MAP

0.2

0.5
0
0.0

NLL

ECE

In-distribution metrics

Acc.

0

Laplace's Diffusion
Sampled Laplace
Linearised Laplace
SWAG
Last-Layer Laplace
MAP

4

25

50

75

100

Shift intensity

125

150

175

0

Laplace's Diffusion posteriors
Sampled Laplace posteriors
Linearised Laplace posteriors
SWAG
Last-Layer Laplace posteriors
MAP

8

25

50

75

100

Shift intensity

125

150

175

0.7

3

2

6

4

2

1

ECE

In-distribution metrics

Acc.

0

Laplace's Diffusion
Sampled Laplace
Linearised Laplace
SWAG
Last-Layer Laplace
MAP

1.0

25

50

75

100

Shift intensity

125

150

Laplace's Diffusion posteriors
Sampled Laplace posteriors
Linearised Diffusion posteriors
SWAG
Last-Layer Laplace posteriors
MAP
0

25

50

75

100

Shift intensity

125

150

175

0.6

Negative log-likelihood

6

0.6

0.4

5

4

3

2
0.2

Laplace's Diffusion posteriors
Sampled Laplace posteriors
Linearised Laplace posteriors
SWAG
Last-Layer Laplace posteriors
MAP

1

0.0

0.3

175

7

0.8

0.4

0.1

0
NLL

0.5

0.2

Expected Calibration Error0

Expected Calibration Error

Negative log-likelihood

0.6

NLL

ECE

In-distribution metrics

Acc.

0

0

25

50

75

100

Shift intensity

125

150

175

0.5

0.4

0.3
Laplace's Diffusion posteriors
Sampled Laplace posteriors
Linearised Diffusion posteriors
SWAG
Last-Layer Laplace posteriors
MAP

0.2

0

25

50

75

100

Shift intensity

125

150

175

Figure 8: Model Fit and Calibration of various posterior sampling methods on in-distribution data(firstcolumn) and under distribution shift for MNIST(top row), Fashion MNIST(middle row) and CIFAR10(bottom row). We use rotated MNIST, rotated FMNIST, and rotated CIFAR in the second and thirdcolumns. 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
Laplace’s diffusion
Sampled Laplace
Linearised Laplace
SWAG
Last-Layer Laplace
Diagonal Laplace
MAP

Conf.(↓)
0.810±0.031
0.583±0.015
0.895±0.027
0.827±0.054
0.842±0.046
0.770±0.026
0.850±0.035

FMNIST
AUROC(↑)
0.911±0.037
0.515±0.021
0.715±0.086
0.949±0.018
0.904±0.033
0.866±0.030
0.907±0.038

Conf.(↓)
0.947±0.007
0.584±0.003
0.934±0.004
0.955±0.004
0.958±0.005
0.907±0.005
0.959±0.005

24

EMNIST
AUROC(↑)
0.629±0.021
0.515±0.004
0.536±0.028
0.627±0.016
0.625±0.028
0.617±0.009
0.634±0.019

Conf.(↓)
0.817±0.015
0.588±0.013
0.863±0.009
0.823±0.013
0.837±0.009
0.747±0.013
0.837±0.010

KMNIST
AUROC(↑)
0.930±0.009
0.507±0.015
0.757±0.018
0.947±0.007
0.935±0.004
0.897±0.005
0.938±0.006

FMNIST
Tested on

MNIST
AUROC(↑)
0.759±0.045
0.495±0.037
0.625±0.050
0.817±0.015
0.761±0.031
0.767±0.032
0.757±0.032

Conf.(↓)
0.734±0.039
0.597±0.027
0.817±0.022
0.727±0.013
0.757±0.019
0.652±0.033
0.759±0.019

Laplace’s diffusion
Sampled Laplace
Linearised Laplace
SWAG
Last-Layer Laplace
Diagonal Laplace
MAP

EMNIST
AUROC(↑)
0.741±0.010
0.503±0.036
0.628±0.013
0.769±0.026
0.735±0.018
0.719±0.033
0.730±0.015

Conf.(↓)
0.730±0.020
0.593±0.026
0.813±0.003
0.763±0.028
0.760±0.017
0.682±0.019
0.762±0.018

Conf.(↓)
0.746±0.012
0.598±0.024
0.816±0.014
0.777±0.007
0.772±0.008
0.696±0.025
0.773±0.010

KMNIST
AUROC(↑)
0.749±0.023
0.493±0.033
0.624±0.020
0.782±0.010
0.747±0.021
0.728±0.032
0.743±0.024

CIFAR-10
Tested on

CIFAR-100
Conf.(↓)
AUROC(↑)
0.790±0.002 0.851±0.002
0.727±0.026 0.687±0.033
0.818±0.005 0.837±0.006
0.776±0.009 0.845±0.008
0.789±0.006 0.864±0.001
0.797±0.004 0.873±0.001

Laplace’s diffusion
Sampled Laplace
Linearised Laplace
SWAG
Last-Layer Laplace
MAP
E.4

SVHN
Conf.(↓)
AUROC(↑)
0.764±0.008 0.862±0.010
0.792±0.022 0.599±0.038
0.809±0.033 0.854±0.024
0.729±0.009 0.876±0.003
0.786±0.029 0.868±0.015
0.792±0.034 0.878±0.017

Short discussion on benchmarks

Sparse Approximations of GGN. Our theoretical analysis is mainly concerned with the ideal
versions of Laplace’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.
It is common to use sparse approximations of the GGN when doing Laplace’s approximations.
Interestingly we observe that Laplace’s approximations with sparse GGN such as diagonal Laplace,
Last Layer, etc do not benefit from linearization to the same degree (Fig. 9).
Sampled Laplace Uncertainty

Linearised Laplace Uncertainty
Laplace Mean predictions
Predictions
Targets

3
2

2

1

1

Laplace Mean Linearised predictions
Predictions
Targets

0
0
1
1

2
3

2
1.0

0.5

1.5

0.0

0.5

1.0

Diagonal Laplace Uncertainty

1.5

2.0

1.0

Diagonal Laplace Mean predictions
Predictions
Targets

1.0

0.5

1.5

0.0

0.5

0.5

0.0

0.0

0.5

0.5

1.5

2.0

Diagonal Laplace Mean Linearised predictions
Predictions
Targets

1.0

0.5

1.0

Linearised Diagonal Laplace Uncertainty

1.0

1.0

1.5
1.0

0.5

0.0

0.5

1.0

1.5

2.0

1.0

0.5

0.0

0.5

1.0

1.5

2.0

Figure 9: Predictive uncertainty of Laplace’s approximation with neural network and linearized
predictive (top row) and diagonal Laplace with neural network and linearized predictive (bottomrow).
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 a kernel. Diagonal andother sparse approximations do not share the spectral properties with the Jacobian in the linearized
predictive. George et al. (2018) note the potential advantages of having the spectrum of the approximate curvature more aligned with the true curvature. This suggests future directions to improvevarious approximations to the GGN by accounting for reparameterizations.
25

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

26