Riemannian Laplace approximations for Bayesian neural networks
Authors: Federico Bergamin
Pablo Moreno-Muñoz
Søren Hauberg
Georgios Arvanitidis
Venue: NeurIPS 2023 poster (2023)
Keywords: Riemannian geometry
Laplace approximation
Approximate inference
Bayesian neural networks
TL;DR: We propose a Riemannian extension of the Laplace approximation that can be used for Bayesian neural networks.

Abstract: Bayesian neural networks often approximate the weight-posterior with a Gaussian distribution. However, practical posteriors are often, even locally, highly non-Gaussian, and empirical performance deteriorates. We propose a simple parametric approximate posterior that adapts to the shape of the true posterior through a Riemannian metric that is determined by the log-posterior gradient. We develop a Riemannian Laplace approximation where samples naturally fall into weight-regions with low negative log-posterior. We show that these samples can be drawn by solving a system of ordinary differential equations, which can be done efficiently by leveraging the structure of the Riemannian metric and automatic differentiation. Empirically, we demonstrate that our approach consistently improves over the conventional Laplace approximation across tasks. We further show that, unlike the conventional Laplace approximation, our method is not overly sensitive to the choice of prior, which alleviates a practical pitfall of current approaches.

Riemannian Laplace approximations for Bayesian neural networks
Federico Bergamin, Pablo Moreno-Muñoz, Søren Hauberg, Georgios Arvanitidis Section for Cognitive Systems, DTU Compute, Technical University of Denmark {fedbe, pabmo, sohau, gear) dk
Abstract
Bayesian neural networks often approximate the weight-posterior with Gaussian distribution. However, practical posteriors are often, even locally, highly non- Gaussian, and empirical performance deteriorates. We propose simple parametric approximate posterior that adapts to the shape of the true posterior through a Riemannian metric that is determined by the log-posterior gradient. We develop a Riemannian Laplace approximation where samples naturally fall into weight- regions with low negative tog-posterion We show that these samples can be drawn by solving a system of ordinary differential equations, which can be done efficiently by leveraging the structure of the Riemannian metric and automatic differentiation. Empirically, we demonstrate that our approach consistently improves over the conventional Laplace approximation across tasks. We further show that, unlike the conventional Laplace approximation, our method is not overly sensitive to the choice of prior, which alleviates practical pitfall of current approaches.
Introduction
Bayesian deep learning estimates the weight-posterion of a neural network given data, i.c. p(0 Due to the generally high dimensions of the weight-space, the nor- malization of this posterior intractable and approximate inference becomes necessity. The most common para- metric choice approximates the posterior with a Gaussian distribution, p( D) x D) _N(0)ua Σ), which is timated variationally [Blundell et al., 2015), using Laplace approximations [MacKay, 1992] with other techniques [Maddox et al., 2019]. Empirical evidence, however, sug- gests that the log-posterior is not locally concave [Sagun et al., 2016], indicating that the Gaussian approximation overly crude. Indeed, this approximation is known to be brittle as the associated covariance is typically ill- conditioned implying suboptimal behavior [Daxberger et al., 2021a, Farquhar al., 2020], and for this reason, alternative approaches have been proposed to fix this issue [Mackay. 1992]. Nonetheless, the Gaussian approximation is widely used due the many benefits of parametric distributions, over e.g. Monte Carlo sampling [Neal, 1995] or deep ensembles [Lakssmiinrrayanan et al., 2017].
Riem. LA
Figure 1: Our Riemannian Laplace ap- proximation is a simple parametric dis- tribution, which is shaped according to the local loss landscape through Rie- mannian metric.
In this paper we argue that the underlying issue is not with the Gaussian approximation, but rather with the weight-space over which the approximation is applied. We show that Gaussian approxima- tion can locally adapt the loss by equipping the weight-space with simple Riemannian metric and performing the approximation tangentially to the associated manifold. Practically, this ensures that samples from the Riemannian approximate posterior land in regions of weight-space yielding low training loss, which significantly improves over the usual Gaussian approximation. We obtain our
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Riemannian approximate posterior using generalization of the Laplace approximation [MacKay, 1992] to general Riemannian manifolds. Sampling from this distribution requires solving a system of ordinary differential equations, which we show can be performed efficiently by leveraging the structure of the used Riemannian metric and automatic differentiation. Empirically, we demonstrate that this significantly improves upon conventional Laplace approximations across tasks.
2 Background
Notation & assumptions. We consider independent and identically distributed (i.id.) data D (Xa,Yn Yn mn=1. consisting of inputs RD and outputs y R To enable probabilistic modeling, we use a likelihood (x)) which is either Gaussian (regression) or categorical (classification). This likelihood is parametrized by deep neural network fo IR" IRC, where 0 E R represent the weights for which we specify a Gaussian prior p(0). The predictive distribution of a new test point x' equals p(y|x') = p(y|x', (D)( where p(0 is the true weight-posterior given the data D. To ensure tractability, this posterior is approximated. This paper focuses the Laplace approximation, though the bulk of the methodology applies to other approximation techniques as well.
2.1 The Laplace approximation
The Laplace approximation (LA) is widely considered in probabilistic models for approximating intractable densities [Bishop, 2007]. The idea is to perform a second-order Taylor expansion of an unnormalized log- probabbbbbby density, thereby yielding Gaussian approximation. When considering inference of the true posterior (0]D), LA constructs an approximate posterior distribution qLA qL. (0)D) = (0 2) that is centered at the maximum a-posteriori (MAP) estimate
0. arg max (log p(0 arg min
log pogp(yn Xn, 0) p()
(I)
(())
A Taylor expansion around 0, of the regularized loss L(0) then yields
1
(()) = [(0.) + (VeC())) (0 0.)) + 2 ((0 0w), Ho Ho( (0 0,)),
(2)
where we know that VoL(0)) & O, and Hg (())) E R denotes the Hessian of the loss. This expansion suggests that the approximate posterior covariance should be the inverse Hessian Σ (0) The marginal likelihood of the data is then approximated as p(D) s exp(-2(0.) (2T) P/2 dE) This is commonly used for training hyper-parameters of both the likelihood and the prior [Immer et al., 2021a, Antorán al., 2022]. We refer appendix for further details.
Tricks of the trade. Despite the simplicity of the Laplace approximation, its application to modern neural networks is not trivial. The first issue is that the Hessian matrix is too large to be stored in memory, which is commonly handled by approximately reducing the Hessian to being diagonal, low-rank, Kronecker factored, only considered for a subset of parameters (see Daxberger et al. [2021a] for a review). Secondly, the Hessian is generally not positive definite [Sagun et al., 2016), which is commonly handled by approximating the Hessian with the generalized approximation [Foresee and Hagan, 1997, Schraudolph, 2002]. Furthermore, estimating the predictive distribution using Monte Carlo samples from the Laplace approximated posterior usually performs poorly [Lawrence, 2001, Chapter 5][Riteer et al., 2018] even for small models. Indeed, the Laplace approximation can place probability mass in low regions of the posterior. A solution, already proposed by [Mackay, 1992, Chapter 4], is to consider a first-order Taylor expansion around f., and use the sample to the "lincarized" function fiin()) fo. (x) + offo 0= A 0.) predictive, where Vefe(x)) p=0, E IR. CxC is the Jacobian. Recently, this approach has been justified by Khan et al. [2019], Immer et al. [2021b], who proved that the generalized Gauss-Newton approximation is the exact Hessian of this new linearized model. Even if this is a linear function with respect to parameters 0, empirically it achieves better performance than the classic Laplace approximation.

---

Although not theoretically justified, optimizing the prior precision post-hoc has been shown to play a crucial role in the Laplace approximation [Ritter al., 2018, Kristiadi et al., 2020, Immer et al.. 2021a, l, Daxberger et al., 2021a]. This is usually done either using cross-validation or by maximizing the log-marginal likelihood. In principle, this regularizes the Hessian, and the associated approximate posterior concentrates around the MAP estimate.
Strengths & weaknesses. The main strength of the Laplace approximation is its simplicity in implementation due to the popularization of automatic differentiation. The Gaussian approximate posterior is, however, quite crude and often does not capture the shape locally of the true posterior [Sagun et al., 2016]. Furthermore, the common reduction of the Hessian to not correlate all model parameters limit the expressive power of the approximate posterior.
3 Riemannian Laplace approximations
We aim to construct parametric approximate posterior that better reflects the local shape of the true posterior and captures nonlinear correlations between parameters. The basic idea to retain the Laplace approximation but change the parameter space θ to locally encode the training loss. To realize this idea, we will first endow the parameter space with suitable Riemannian metric (Sec. 3.1) and then construct Laplace approximation according to this metric (Sec. 3.2).
3.1 A loss-aware Riemannian geometry
0
Expa (v)
For a given parameter value 0 E O, we can measure the training loss (()) of the associated neural network. As- suming that the loss changes smoothly with 0, we can interpret the loss surface .M g(0)
0
00, [()) E R-1
as a K -d-iimensional manifold in IR k+1 The goal of Rie- mannian geometry [Lee, 2019, do Carmo, 1992] do calculations that are restricted to such manifolds.
The parameter space θ of the BNN together with examples of the Rie- mannian metric and the exponential map. Note that the Riemannian metric adapts to the shape of the loss which causes the geodesic follow its shape.
Figure 2:
The metric. We can think of the parameter space θ as being the intrinsic coordinates of the manifold M, and it is beneficial to do all calculations directly in these coordinates. Note that a vector tangential to the manifold can be written J, (0)v E RR+1 where J, θ
KK+Ix is the Jacobian of g that spans the tangent space Tate) M at the point g0)) E .M and IR the vector of tangential coordinates for this basis of the tangent space. We can take inner products between two tangent vectors in the same tangent space as (1)
= (OT,, which, we note, is now expressed directly in the intrinsic coordinates. From this observation, we define the Riemannian metric M(0)) J. (O)T J, (0), which gives notion of a local inner product in the intrinsic coordinates of the manifold (see ellipsoids in Fig. 2). The Jacobian of is particularly simple (0) [4x, V VoL.. such that the metric takes the form
M(0) Ix + V VoL(O))
(3)
The exponential map. A local inner product allows us to define the length of a curve e: [0, 1] θ as length [c] = fo / ((()) M(c(I))(())d where é(l) = 0,c(l) is the velocity. From this, the distance between two points can be defined as the length of the shortest connecting curve, where the latter is known as the geodesic curve. Such geodesics can be expressed as solutions to a system of second-order non-linear ordinary differential equations (ODES), which is given in appendix B alongside further details on geometry. Of particular interest to us is the exponential map, which solves these ODES subject to an initial position and velocity. This traces out a geodesic curve with given starting point and direction (see Fig. 2). Geometrically, we can also think of this as mapping a tangent vector back the manifold, and we write the map Exp: M x TeM + M.
The tangential coordinates v can be seen as a coordinate system for the neighborhood around 0, and since the exponential map is locally bijection we can represent any point locally with unique tangent vector. However, these coordinates correspond to the tangent space that is spanned by J. ( (0), which implies that by changing this basis the associated coordinates change as well. By
3

---

LA
OUR
(a) True posterior
(b) sample
(c) OUR sample
(d) LA model
(c) OUR model
Figure 3:
The LA assigns probability mass to regions where the true posterior is nearly zero, and sample from this region corresponds
poor classifier. Considering this sample the initial velocity for the exponential map, the generated sample falls within the true posterior and the associated classifier performs well. As result, our model quantifies better the uncertainty.
orthonormalizing this basis we get the normal coordinates where the metric vanishes. Let v the tangential coordinates and v the corresponding normal coordinates, then it holds that
(v, M()vv) (v, v). =
A(0)v with A()) M(0)-
(4)
We will use the normal coordinates when doing Taylor expansions of the log-posterior, akin to standard Laplace approximations.
3.2 The proposed approximate posterior
In order to Taylor-expand the loss according the metric, we first express the loss in normal coordinates of the tangent space 0., h(v) L(Exp,, (M(0.)- Following the standard Laplace approximation, we perform a second-order Taylor expansion of as
1 (v. Ho Hclhl(() h(v)= = h(0) + (0ch(v) -= v)
(5)
2
where Bch(v)lo=o A AOO VO..VV
R 0 as 0, minimize the loss and Hy [h][h[( v= = (()))(((( with H. (6) the standard Euclidean Hessian matrix of the loss. Further details about this step can be found in appendix B.
Tangential Laplace. Similar the standard Laplace approximation, we get Gaussian approximate posterior ((v) = N(V O, Σ) on the tangent space in the normal coordinates with covariance E = H. [V] v- Note that changing the normal coordinates to tangential coordinates is linear transformation and hence ~ No A
which means that this covariance is equal to H, since ( a symmetric matrix, and hence, it cancels out. The approximate posterior qr (v) | 0) in tangential coordinates, thus, matches the covariance of the standard Laplace approximation.
The predictive posterior. We can approximate the predictive posterior distribution using Monte Carlo integration as p(y|x', D) p(y|x', (. ())()) = ply|x', D, Exp,, (v))gr(v)ds s
"ESip(p xx' D, Exp., (v=)), V. (v).. Intuitively, this generates tangent vectors according to the standard Laplace approximation and maps them back the manifold by solving the geodesic ODE, This lets the Riemannian approximate posterior take shape from the loss landscape, which is largely ignored by the standard Laplace approximation. We emphasize that this is general construction that applies to the same Bayesian inference problems as the standard Laplace approximation and is not exclusive Bayesian neural networks.
The above analysis also applies to the linearized Laplace approximation. In particular, when the flim (x) is considered instead of the fo(x) the loss function in (1) changes (l) Consequently, Riemannian metric is computed under this new loss, and V (i appears in the metric (3).
Example. build intuition, we consider a logistic regressor on linearly separable dataset (Fig. 3). The likelihood of point E R22 to be in one class is p(C = 1/x) = (xTO + b), where a(-) is the sigmoid function, E IR² and b E After learning the parameters, we and show the posterior with respect to 8 together with the corresponding standard Laplace approximation (Fig. 3a).

---

We see that the approximation assigns significant probability mass to regions where the true posterior near-zero, and the result of a corresponding sample a poor classifier (Fig. 3b). Instead, when we consider this sample as the initial velocity and compute the associated geodesic with the exponential map, we generate sample the tails of the true posterior which corresponds to a well-behaved model (Fig. 3c). We also show the predictive distribution for both approaches and even if both solve easily the classification problem, our model better quantifies uncertainty (Fig. 3c).
3.3 Efficient implementation
Our approach is natural extension of the standard Laplace approximation, which locally adapts the approximate posterior to the true posterior. The caveat is that computational cost increases since we need to integrate an ODE for every sample. We now discuss partial alleviations.
Integrating the ODE. In general, the system of second-order nonlinear ODES (see appendix B for the general form) is non-trivial as it depends on the geometry of the loss surface, which is complicated in the over-varannn parametrized regime [Li et al., 2018]. In addition, the dimensionality of the parameter space is high, which makes the solution of the system even harder. Nevertheless, due the structure of our Riemannian metric (3), the ODE simplifies to
{()) V (1 Vellc())VV (())
which can be integrated reasonably efficiently with standard solvers. In certain cases, this ODE can be further simplified, for example when we consider the linearized loss C (li(()) and Gaussian likelihood.
The ODE (6) requires computing both gradient and Hessian, which are high-dimensional objects for modern neural networks. While we need to compute the gradient explicitly, we do not need to compute and store the Hessian matrix, which is infeasible for large networks. Instead, we rely on modern automatic-differeniaatoo frameworks to compute the Hessian- vector product between Hg (c(()) and é(t) directly. This both reduces memory use, increases speed, and simplifies the implementation.
iitoomati-difffernntttioo
batiing The cost of computing the metric, and hence the ODE scales linearly with the number of training data, which can be expensive for large datasets. A reasonable approximation is to mini-batch the estimation of the metric when generating samples, i.e. construct batch B of B random data points and use the associated loss in the ODE (6). As usual, we assume that (() x (N/B) Note that we only mini-batch the metric and not the covariance of our approximate posterior qr (v).
We analyze the influence of mini-batching in our methods and provide empirical evidence in Fig. 4. In principle, the geometry of the loss surface C(0) controls the geodesies via the associated Riemannian metric, when we con- sider the full dataset we expect the samples to behave similarly to fo. (x). In other words, our approximate pos- terior generates weights near 0, resulting in models with similar even better loss. When we consider batch the geometry of the associated loss surface L&(0), controls the generated geodesic. So if the batch represents well the structure of the full dataset, then the resulting model will be meaningful with respect the original problem, and in addition, it may exhibit some variation that is beneficial from the Bayesian perspective for the quantification of the uncertainty. The same concept applies in the linearized version, with the difference that when the full dataset is considered the geometry of (0) may over-regularize the geodesics. Due to the linear nature of flin the asso- ciated Riemannian metric is small only close to f., so the generated samples are similar to so_ (x). We relax this be- havior and potentially introduce variations in the resulting models when we consider a different batch whenever we generate a sample. Find more details in appendix D.
RIEM-LA
0:32
Trst NLL Performance Compuiaaiion Time
130 13
Teat NLL
0.79
0.28
d0
200
Botch size
LIN-RIEM-LA
Test NLL Performance
Computation Time
0.42
Tixt NLL
0:40
0:35
+00
Botth
Figure 4: Analysis of mini-batching
5

---

4 Related work
Bayesian neural networks. Exact inference for BNNS is generally infeasible when the number of parameters is large. Several methods rely on approximate inference, which differs in their trade-off between computational cost and the goodness of the approximation. These techniques are usually based on the Laplace approximation [MacKay. 1992], variational inference [Graves, 201 I, Blundell et al., 2015, Khan et al., 2018], dropout [Gal and Ghahramani, 2016]. stochastic weight averaging [Izmailov et al.. 2018, Maddox et al., 2019] or Monte Carlo based methods [Neal, 1995], where the latter is often more expensive.
Laplace approximations. In this work, we are primarily focused on Laplace approximations, although the general geometrie idea can be used in combination with any other inference approach listed above. Particularly, Laplace's method for BNNs was first proposed by Mackay [1992] in his evidence framework, where a closed-form approximation of predictive probabilities was also derived. This one uses first-order Taylor expansion, also known as linearization around the MAP estimate. For long, Laplace's 's method was infeasible for modern architectures with large networks due to the exact computation of the Hessian. The seminal works of Martens and Grosse [2015] and Botev et al. [2017] made it possible to approximate the Hessian of large networks, which made Laplace approximations feasible once more [Ritter al., 2018]. More recently, the Laplace approximation has become go-to tool for turning trained neural networks into BNNS in a post-hoc manner, thanks to easy-to-use software [Daxberger al., 2021a] and new approaches scale up computation [Antorán et al., 2022]. In this direction, other works have only considered a subset of the network parameters [Daxberger al., 2021b, Sharma al., 2023], especially the last-layer. This is de facto the only current method competitive with ensembles [Lakshminarayanan et al., 2017].
Posterior refinement. Much work has gone into building more expressive approximate posteriors. Recently, Kristiadi et al. [2022] proposed to use normalizing flows to get a non-Gaussian approximate distribution using the Laplace approximation as base distribution. Although this requires training an additional model, they showed that few bijective transformations are enough to improve the last-layer posterior approximation. Immer al. [2021b). instead, propose to refine the Laplace approximation by using Gaussian variational Bayes or a Gaussian process. This still results in a Gaussian distribution, but it has proven beneficial for linearized Laplace approximations. Other approaches rely on a mixture of distributions to improve the goodness of the approximation. Miller al. [2017] expand a variational approximation iteratively adding components to a mixture, while Eschenhagen et al. [2021] use a weighted sum of posthoc Laplace approximations generated from different pre-trained networks. Havasi et al. [2021], instead, introduces auxiliary variables make local refinement of a mean-field variational approximation.
Differential geometry. Differential geometry is increasingly playing role in inference. Arvanitidis et al. [2016] make a Riemannian normal distribution locally adapt to data by learning a suitable Riemannian metric from data. In contrast, our metric is derived from the model. This is similar in spirit to work that investigates pull-back metrics in latent variable models [Tosi et al., 2014, Arvanitidis al., 2018, Hauberg, 2018b]. A similar Riemannian metric has been used as surrogate for the Fisher metric for Riemannian Hamiltomian Monte-Carlo sampling [Hartmann et al., 2022]. In addition that, the geometry of the latent parameter space of neural networks was recently analyzed by Kristiadi et al. [2023] focusing on the invariance of flatness measures with respect to re- re-parametrizations. Finally, we note that Hauberg [2018a] considers Laplace approximations on the sphere as part of constructing a recursive Kalman-like filter.
5 Experiments
We evaluate our Riemannian LA (RIEM-LA) using illustrative examples, image datasets where we use a convolutional architecture, and real-world classification problems. We compare our method and its linearized version to standard and linearized LA. All predictive distributions are approximated using Monte Carlo (MC) samples. Although laser LA is widely used lately, we focus on approximating the posterior of all the weights of the network. In all experiments, we maximize the marginal log-likelihood to tune the hyperparameters of the prior and the likelihood as proposed in [Daxberger et al., 2021a]. To evaluate the performance in terms of uncertainty estimation we considered the standard metrics in the literature: negative log-likelihood (NLL), the Brier score (BRIER), the expected calibration error (ECE), and the maximum calibration error (MCE). More experiments are

---

(a) VANILLA LA
(b) RIEM-LA
(c) LIN-LA
(d) LIN- RINRRIEMLA
Figure 6: Binary classification confidence estimate using 100 Monte-Carle samples on the banana dataset. Vanilla LA underfit, while the other three methods are able to be certain within the data and uncertain far away. Note, for linearized RIEM-LA we solve the expmap using a different subset the data. Confidence plots of all different methods can be found in the supplementary material. Black lines are the decision boundaries.
available in appendix D together with the complete training and modeling details. In appendix C.2 we analyze the runtime to compute the exponential map over different dataset and model sizes. Code reproduce the results publicly available at bttpp://gittub com com//federicobbergamin ri riemannian. -aplaacea-eee-eeeaaeaaaaaaaaaela approximation
5.1 Regression problem
samppes
LA stmplis
We consider the toy-regression problem proposed by Snel- and Ghahramani [2005]. The dataset contains 200 data points, and we randomly pick 150 examples as our training set and the remaining 50 test set. As shown by Lawrence [2001], Ritter et al. [2018], using samples from if the Hessian is not particularly ill-conditioned, i.c. when the prior precision is optimized. For this reason, the lineariza- tion approach is necessary for regression with standard LA. Instead, we show that even our basic approach fixes this problem when the prior is optimized. We tested our approach by considering two fully connected networks, one with one hidden layer with 15 units and one with two layers with 10 units each, both with tanh activations. Our approach approximates well the true posterior locally, the resulting function samples follow the data. Of course, if the Hessian is extremely degenerate our approach also suffers, as the initial velocities are huge. When we con- sider the linearized version of our approach the result is the same as the standard L.-llinarization, which we in- clude in the appendix D, where also report results for in- between uncertainty as proposed by Foong et al. [2019] and comparison with Hamiltonian Monte Carlo.
Our moan
the LA posterior performs poorly in regression
Ouz somposs
LA semples
men
Our
Figure Posterior samples under a sim- plc (top) and overparametrized model (bottom). Vanilla LA is known to gener- ate bad models, while our samples from RIEM-LA quantify well the uncertainty.
5.2 Classification problems
Illustrative example. We consider 2-dimensional binary classification problem using the banana dataset which is shown in Fig. 6. We train a 2-layer fully connected neural net with 16 hidden units per layer and tanh activation. For all methods, we use 100 MC samples for the predictive distribution.
As in regression, direct samples from the vanilla LA lead to really poor model (Fig. 6a) with high uncertainty both within and away from the data support. Instead, the other three methods (Fig. 6b-6d) show a better-behaved confidence that decreases outside of the data support. This is also supported by the metrics in Table 1, where remarkably RIFM-L. performs better in terms of NLL and Brier score separate test set.

---

Table 1: In-distribution results in the banana dataset. We use 100 MC samples both for the LA and our variants. ECE and MCE are computed using M = 10 bins. We report mean and standard error over different seeds.
BANANA DATASET
PRIOR OPTIMIZED
PRIOR NOT OPTEMIZED
METHOD
Accuracy
NLL4
Briers
Accurary 1
NLLL
Brier]
86.69 0.34
0.333 0.0015
0.0320
0.0015
86.69 0.334
0.333 0.005
0.0560
0.0015
0.009
0.2426
2.332
0.003
0.2534
0.0017
LIN-LA
86.99 0.33
0.325
0.008
0.0356
2.0023
86.92=0.40
0.403 + 0.012
0.1196
0.0044
RIES--LA
87.57
0.07
0:287
0.002
0.0886
0006
87:14 0:220
0:285
0.001
0.0878
0.0006
RIEM-LA (BATCHES)
87.30 - 0.08
0.286
= 0.001
0. 0.0890
+.0 0.000.000 0000
87.32220.0
0.294 + 00000
0 0.0895
0.0004
LIN REEM LA
87.02
0.38
0.415
0.029
0.0967
O. 0,0024
85,33
165311000 0.33
0.884 1: 0.037
0.1252
0.00/2
LIN-IEM-LLA (BATCHES)
87.77 0.24
0,298 0.006
0 0.0687
0.00011
86.16 0.211
0.352
0.002
0.0594
0.0011
As we discussed in Sec. 3.3, using subset of the dataset for computing the exponential map can be beneficial for our linearized manifold in addition to speeding up computation. In Fig. 6d we plot the confidence for our linearized approach using batches while in appendix confidence of the same approach using the full data for solving the ODES. We can see that our linearized RIEM-LA tends to be overconfident outside the data region and also close to the decision boundary. This behaviour can be found in the high NLL that linearized RIEM-LA gets compared to our vanilla approach and linearized I.A.
we show the
UCI datasets. We compare our approach against the standard LA on a set of six UCI classification datasets using fully connected network with single layer, 50 hidden units and tanh activation. The predictive distribution is estimated using MC with 30 samples from the approximate posterior of each approach. In Table 2 we compare the methods in terms of their negative liglliieldooo (NLL) in the test set. All other metries are reported in appendix D. We are considering the setting where optimize the consider our standard approaches without using batches, which we have seen that specifically for our linearized approach may lead to sub-optimal performance.
post-hoc, which is the optimal setting for LA and linearized LA. We
From the results in Table we see that our RIEM-LA consistently performs better in terms of negative log-likelihood than vanilla and linearized LA. We also observe that in two datasets the performance of our linearized RIEM-LA not optimal. This implies that the loss surface of the linearized loss potentially over-regularizes the geodesics we analyzed in Sec. 3.3, and in this case, considering mini- hini-batching could have been beneficial.
Table 2: Negative log-likelihood (lower is better) on UCI datasets for classification. Predictive distribution is estimated using 30 MC samples. Mean and standard error over different seeds.
PRIOR PRECISION OPTIMIZED
MAP
VANILLA LA
RIEM-LA
LINEARIZED LA
LINEARIZED RIEM LA
VEHICLE R
0.975 1: 0.081
1.209 1: 0.020
0.454 1 0.024
0.875 + 0.020
0.494 + 0.044
2.084 0.323
1.737 0.037
1.047
0.224
1.365 + 0.058
1.359 = 0.299
JONOSPHERE
1.032 +: 0.175
0.673
1: 0.013
0.344 1: 0.068
0.497 1 0.015
0.625 +:2 0.110
WAVEFORM
1.076 - 0.110
0.888 0.030
0.459
0.057
0.640 + 0.002
0.575 = 0.065
AUSTRALIAN
1.306 = 0.146
0,081 188100 0.011
0.541 + 0.053
0.570 = 0.016
0.833 = 0.108
BREAST CANCER
0.225 4 0.076
0.594 + 0.030
0.176 + 0.092
0.327 = 0.022
0.202 = 0.073
Image classification. We consider a small convolutional neural network on MNIST and Fash- ionMNIST. Our network consists of two convolutional layers followed by average pooling layers and three fully connected layers. We consider model of this size as the high dimensionality of the parameter space is one of the main limitations of the ODE solver. For the training of the model, we subsample each dataset and we consider 5000 observations by keeping the proportionality of labels, and we test in the full test set containing 8000 examples. In Table 3 we compare the different methods with the prior precision optimized as this is the ideal setting for the linearized LA. We refer to appendix for the setting with the prior precision not optimized.
From the results we observe that our standard RIEM-LA performs better than all the other methods in terms of NIL and Brier score, meaning that the models are better calibrated, but also leads to more accurate classifier than the MAP. In terms of ECE, it seems that considering the linearized
8

---

Table 3: Image classification results using a CNN on MNIST and FashionMNIST. The network is trained on 5000 examples and we test the in on-distribttion performance on the test set, which contains 8000 examples. We use 25 Monte Carlo samples to approximate the predictive distribution and 1000 datapoints per batch in our batched manifolds. Calibration metrics are computed using M = 15 bins. We report mean and standard error for each metric over different seeds.
CNN ON MNIST PRIOR PRECISION OFTINIZED
Accuracy 1
NLL4
Brier]
BCE]
MCEJ
MAP
95.02 +: 0.17
0.167 + 0.005
U.0075
0.0002
1.03 + 0.14
39.94
14.27
VANILLA
88.69 1: 1.84
0.871 1. 0.026
0.0393
0.0013
42.11 1.22
50.52 1: 1.45
LIN-LA
94.91 + 0.26
0.204 + 0.006
0.0087 + 0.0003
6.30 0.08
39.30 + 16.77
RIEM-LA
96.74 0.12
0.115 1:
0.0052
0,0002
2.48
38.03
15.02
RIEM-LA (BATCHES)
95.67 1 0.19
0.170 1 0.005
0.0072 1 0.0002
5.40 0.06
22.40 + 0.51
LIN-RIEM-LA
95.44
+ 0.18
0.149 + 0.004
0.0068
0.0003
39.40
14.75
LIN-RIEM-LA (BATCHES)
95.11 = 0.20
0.167 + 0.001
0.0076 0.0002
3.23 0.01
18,10 = 2.50
CNN FASHIONMNIST
OPTIMIZED
METHOD
Accuracy f
NLL4
Brier]
ECF
MCF4
MAP
79.88 + 0.09
0.541 + 0.002
0.0276 + 0.0000
1.66 + 0.07
24.07 = 1.50
VANILLA
0.83
1.026 + 0.016
0.0019
31.63
2.95
LIN-LA
79.85 + 0.13
0.549 + 0.001
0.0278 + 0.0000
3.23 0.44
37.88 + 17.98
RIEM-LA
83.33 = 0.17
0.472 + 0.001
0.0237
0.0001
3.13
10.94 - 2.11
RIEM-LA (BATCHES)
81.65 + 0.18
0.325 + 0.004
0,0263
0.0002
5.80 0.73
35.30 + 18.40
LIN RIEM LA
81.33 1: 0.10
0.521 1. 0.004
0.0261 1 0.0002
1.59 1 0.40
25.53
1 0.10
LIN-RIEM-LA (BATCHES)
80.49
+ 0.13
0.529 + 0.003
0.0269 0.0002
2.10 0.42
6.14 1.42
approach is beneficial in producing better-callbrated models both datasets. This holds both for our approach lincarized RIEM-LA and the standard LA. Optimizing the prior precision post-hoc is crucial for the vanilla LA and associated results can be seen in appendix D. Instead, both our methods appear to be robust and consistent, they achieve similar performance matter if the prior precision is optimized or not.
Note that for the mini mini-batches for our approaches, we consider 20% of the data by randomly selecting 1000 observations while we respect the label frequency based on the full dataset. Clearly, the batch- size is hyperparameter for our methods and can be estimated systematically using cross-validation. Even if we do not optimize this hyperparameter, we see that our batched version of RIEM-LA and LIN-RIEM-LA perform better than the standard LA and on-par with our LIN-RIEM-LE without batches, implying that a well-tuned batch-size can potentially further improve the performance. Nevertheless, this also shows that our method is robust with respect to the batch-size.
Conclusion and future directions
We propose extension to the standard Laplace approximation, which leverages the natural geometry of the parameter space. Our method is parametric the sense that a Gaussian distribution estimated using the standard Laplace approximation, but it adapts to the true posterior through nonparametric Riemannian metric. This is general mechanism that, in principle, can also apply to, e.g., variational approximations. In a similar vein, while the focus of our work is Bayesian neural networks, nothing prevents us from applying our method to other model classes
Empirically, we find that our Riemannian Laplace approximation is better on par with alternative Laplace approximations. The standard Laplace approximation crucially relies on both lineariza- tion and on fine-tuned prior to give useful posterior predictions. Interestingly, Riemannian Laplace approximation requires neither. This could suggest that the standard Laplace approximation has rather poor posterior fit, which our adaptive approach alleviates.
find that the
Limitations. The main downside of our approach is the computational cost involved in integrating the ODE, which is common problem in computational geometry [Arvanitidis et al., 2019]. The cost of evaluating the ODE scales linearly with the number of observations, and particular it OSSW N). where S is the number of steps of the solver, W is the number of model parameters and N is the dataset size. Indeed, the solver needs all the training data points to compute both gradient and hvp at each step. We refer to appendix C.I for detailed explanation. For big datasets, we have considered

---

the "obvious" trick of using a random (small) batch of the data when solving the ODE to reduce the complexity. Empirically, we find that this introduces some stochasticity in the sampling, which sometimes is beneficial to explore the posterior distribution better and eventually boost performance, motivating further research. The computational cost also grows with the dimensionality of the parameter space, as the number of necessary solver steps increases, as well as the cost of the Hessian vector products. Our implementation relies on an off-the-shelf ODE solver which runs on the CPU while cioomatic-iiffereniattoo based approach runs on the GPU, which is inefficient. We expect significant improvements by using an ODE solver that runs exclusively on GPU, while taieor-made numerical integration methods also of particular interest.
Future directions. We showed that equipping the weight-space with simple Riemannian metric is promising approach that solves some of the classic LA issues. We believe that this opens up different interesting research directions that either aim to use different metries instead of the one used in this work or to make this method scale to bigger dataset. Regarding the metric, a possible extension of this work would be to consider the Fisher information matrix in the parameter space. Potential ideas to improve efficiency is to consider approximations of the Riemannian metric leading to simpler ODE systems, for example by using the KFAC instead of the full-Hessian in (6). Another directions, instead, would be to focus on developing a better solver. Indeed, we know the structure and behavior of our ODE system, i.e. geodesics start from low loss which increases along the curve, therefore, potential direction would be to develop solvers that exploit this information. Usually, general purpose solvers aim for accuracy, while in our case even inexact solutions could be potentially useful if computed fast.
Acknowledgments and Disclosure of Funding
This work was funded by the Innovation Fund Denmark (0175-00014B) and the Novo Nordisk Foun- dation through the Center for Basic Machine Learning Research in Life Science (NNF200C0062666). also received funding from the European Research Council (ERC) under the European Union's Horizon 2020 research, innovation programme (757360). SH was supported in part by a research grant (42062) from VILLUM FONDEN.
References
P.-. Absil, Mahony, and Rodolphe Sepulchre. Optimization Algorithms on Matrix Manifoods. Princeton University Press, 2008
Javier Antorán, David Janz, James U Allingham, Erik Daxberger, Riccardo Rb Barbano, Eric Nalisnick, and José Miguel Hernández-Lobato. Adapting the linearised laplace model evidence for modern deep learning. In International Conference on Machine Learning (ICML), 2022.
Georgios Arvanitidis, Lars K Hansen, and Søren Hauberg. A locally adaptive normal distribution. In Neural Information Processing Systems (NeurIPS), 2016.
Georgios Arvanitidis, Lars Kai Hansen, and Soren Hauherg. Latent space addity: on the curvature of deep generative models. In International Conference on Learning Representations (ICI.), 2018
Georgios Arvanitidis, Spren Hauberg, Philipp Hennig, and Michael Schober. Fast and robust shortest paths on manifolds learned from data. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2019.
Christopher M. Bishop. Pattern Recognition and Machine Learning. Springer, 2007.
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In International conference on machine learning (ICML), 2015.
Aleksandar Botev, Hippolyt Ritter, and David Barber. Practical Gauss newss- optimisation for deep learning. In International Conference on Machine Learning (ICML), 2017.
Erik Daxberger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hennig. Laplace redux-efforteen effortless Bayesian deep learning. In Neural Information Processing Systems (NeurIPS). 10
10

---

Erik Daxberger, Eric Nalisnick, James U Allingham, Javier Antorán, and José Miguel Hernández-Loban Bayesian deep learning via subnetwork inference. In International Conference Machine Learning (ICML). 2021b.
M.P. do Carmo. Riemannian Geometry. Mathematics (Boston, Mass.). Birkhäuser, 1992.
John R Dormand and Peter Prince. A family of embedded runge-kutta formulac. Journal of computational and applied mathematics, 6(1):19-26, 1980.
Runa Eschenhagen, Erik Daxherger, Philipp Hennig, and Agustinus Kristiadi. Mixtures of Laplace approxima-
tions for improved post-hoe uncertainty in deep learning, In NeurlPS Workshop of Bayesian Deep Leurning.
2021.
Sebastian Farquhar, Michael A Osborne, and Yarin Gal. Radial bayesian neural networks: Beyond discrete support in large-scale bayesian deep learning. In International Conference Artificial Intelligence and Statistics (AISTATS), 2020.
Andrew YK Foong, Yingzhen Li, José Miguel Hernández-Lobato, and Richard E Turner. 'In-Between' Uncer- tainty in Bayesian Neural Networks. arXiv preprint arXiv: 10.9906/15533 /537,, 2019.
F Dan Foresee and Martin T Hagan. Gauss-Newton approximation to Bayesian learning. In Proceedings of International Conference on Neural Networks (ICNN). IECE, 1997.
Yarin Gal and Zoubin Chahramani. Dropout a bayesian approximation: Representing model uncertainty in deep learning. In International Conference Machine Learning (ICML), 2016.
Alex Graves, Practical variational inference for neural networks. In Neural Information Processing Systems (NeurIPS), 201
Marcclo Hartmann, Mark Girolami, and Arto Klami. Lagrangian manifold monte carlo on monge patches. In International Conference on Artificial Intelligence and Statistics, 2022.
Søren Hauberg. Directional statistics with the spherical normal distribution. In FUSION 2018, 2018a.
Søren Hauberg. Only bayes should learn a manifold. arXiv preprint, 2018b.
Marton Havasi, Jasper Snoek, Dustin Tran, Jonathan Gordon, and José Miguel Hernández-Lobato. Refining the variationaa posterior through iterative optimization. Entrpp, 2021.
E Heirer, SP Nørsett, and G Wanner. Solving ordinary differential equations i: Nonstiff problems, 1987.
Philipp Hennig and Søren Hauberg. Probabilistic solutions to diffcrential equations and their application to riemannian statisties. In Proceedings of the 17th international Conference on Artificial Intelligence and Statistics (AISTATS), 2014.
Richard Zou Horace He. functorch: Jax-like composable function transforms for pytorch. https://githur http//ppoorcc/fuuorrcchhhh 2021.
Alexander Immer, Matthias Bauer, Vincent Fortuin, Gunnar Rätsch, and Khan Mohammad Emtiyaz. Scalable marginal likelihood estimation for model selection in deep learning. In International Conference Machine Learning (ICMI.), 2021a.
Alexander Immer, Maciej Korzepa, and Matthias Bauer. Improving predictions of Bayesian neural nets via local linearization. In International Conference Artificial Intelligence and Statistics (AISTATS), 2021b.
Pavel Izmailov, Dmitri Podoprikhin. Timur Garipov, Dmitry P. Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. In Uncertainty in Artificial Intelligence (UAI). 2018.
Mohammad Khan, Didrik Nielsen, Voot Tangkaratt, Wu Lin, Yarin Gal, and Akash Srivastava. Fast and scalable bayesian deep learning by weight-perturbation in adam. In International Conference Machine Learning (ICMI.), 2018.
2018.
Mohammad Emtiyaz E Khan, Alexander Immer, Ehsan Abedi, and Maciej Korzepa. Approximate inference deep networks into Gaussian Processes. Neural Information Processing Systems (NeurIPS), 2019.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic opúmizztiion In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference Learning Representations, ICLR, 2015.
Agustinus Kristiadi, Matthias Hein, and Philipp Hennig. Being bayesian, even just a bit, fixes overconfidence in ReLU networks. In International Conference Machine Learning (ICML), 2020.
11

---

Agustinus Kristiadi, Runa Eschenhagen, and Philipp Hennig, Posterior refinement improves sample efficiency in bayesian neural networks. In Neural Information Processing Systems (NeurIPS), 2022.
Agustinus Kristiadi, Felix Dangel, and Philipp Hennig. The geometry of neural nets' parameter spaces under reparametrization. arXiv preprint, 2023.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles. In Neural Information Processing Systems (NeurlPS), 2017.
Neil David Lawrence. Variational inference in probabilistic models, PhD thesis, Citeseer, 2001.
J.M. Lee. Introduction to Riemannian Manifolds. Graduate Texts in Mathematics. Springer International Publishing, 2019.
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. In Neural Information Processing Systems (NeurIPS), 2018.
Ilya Loshchilov and Frank Hutter. SGDR: stochastic gradient descent with warm restarts. In 5th International Conference on Learning Representations, ICLR, 2017.
David J. C. MacKay, A practical Bayesian framework for backpropagation networks. Neural Computation, 1992.
David John Cameron Mackay. Bayesian methods for adaptive models. California Institute of Technology, 1992.
Wesley Maddox, Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson. A simple baseline for Bayesian uncertainty in deep learning. In Neural Information Processing Systems (NeurIPS), 2019.
James Martens and Roger Grosse, Optimizing neural networks with Kronceker-factored approximatc curvature. In International Conference Machine Learning (ICML), 2015.
Andrew C Miller. Nicholas J Foti, and Ryan P Adams. Variational boosting: Iteratively refining posterior
approximations. In International Conference on Machine Learning (ICML), 2017.
Radford M Neal. Bayesian learning for neural networks. PhD thesis, University of Toronto, 1995.
Xavier Pennec. Intrinsic Statistics on Riemannian Manifolds: Basic Tools for Geometric Measurements. Journal of Mathematical Imaging and Vision, 2006.
Hippolyt Ritter, Aleksandar Botev, and David Barber. A scalable laplace approximation for neural networks. In International Conference on Learning Representations (ICLR), 2018.
Levent Sagun, Leon Bottou, and Yann LeCun. Eigenvalues of the hessian in deep learning: Singularity and beyond. arXiv preprint 111.0:661. 2016.
Nicol N Schraudolph. Fast curvature matrix-vector products for second-order gradient descent. Neural computation, 2002.
Mrinank Sharma, Sebastian Farquhar, Eric Nalisnick, and Tom Rainforth. Do Bayesian Neural Networks Need To Be Fully Stochastic? In International Conference on Artificial Intelligence and Statistics (AISTATS), 2023.
Edward Snelson and Zoubin Ghahramani, Sparse Gaussian processes using pscudo-inputs. In Neural Information Processing Systems (NeurIPS), 2005.
Alessandra Tosi, Spren Hauberg, Alfredo Vellido, and Neil D. Lawrence. Metrics for probabilistic geometries. In Uncertainty in Artificial Intelligence (UAI), 2014.
Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Cournapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, Carey, ilhan Polat, Yu Feng. Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, Antônio H. Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SciPy 1.0 Contributors. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods, 17:261-277 2020. doi: 10. 1033//411992-011 -01--006------22
12

---

Riemannian Laplace approximations for Bayesian neural networks
(Appendix)
Contents of the Appendix
A Laplace background.
B Riemannian geometry
Implementation details
C.I Cost of solving the ODE system
C.2 Analysis of exponential map
D Experiments details and additional results
D.I Regression example
D.2 Illustrative 2D classification example
D.3 An additional illustrative example
D.4 Additional results in the UCI classification tasks
D.5 Complete results on MNIST and FMNIST
D.6
Out-of-distribution results
A
Laplace background
We provide a more detailed introduction to Laplace approximation and discuss the optimiziation of the prior and likelihood hyperparameters by maximizing the marginal likelihood. In BNNs, Laplace approximation is use to approximate the posterior distribution, i.c.:
p(0 D) p(D 0)p(0) p(D ()p(()
p(D) PoP(D 0)p(0) z P(D ())(())
(A.1)
This is done by fitting a Gaussian approximation the unnormalized distribution p(D p)p(0) at its peak, where ptD 0) is the likelihood distribution and p(0) is the prior over the weights. In standard training neural networks the mean-squared error loss is usually used for regression. This corresponds to optimizing Gaussian log-likelihood up to a scaling factor while using the cross-entropy loss in classification correspond to minimizing the negative log-likelihood. The usual weight decay, or L2 regularization, corresponds instead to Gaussian prior p(0) distribution. More specifically, training with \|101| corresponds to Gaussian prior N/ 0, A
Therefore natural peak to choose given by the d., which is the set of weights obtained at the end of the training. Indeed, 0, is usually computed by:
0, = arg min
log plyn Xn, 0) - log p(0)
(A.2)
(())

---

Once we have d., the Laplace approximation uses second-order Taylor expansion of [()) around 0.. which yields:
(()) = C(0.) (Vo[(0)) 0.)) ((0 0.). Ho]((()))))))))))))))))))))))))) b-o. (0- 0.)), (A.3)
where the first-order term VoL(0) 2 O because the gradient at 0, is O.
By looking L(0), notice that the Hessian is composed by two terms: data fitting term and a prior term. Assuming p(0) = (0; 0, I), then the Hessian can be expressed as
Ho| 0=0.- =y 1
Vi log x4) la.= (H + al),
(A.4)
where defined 72z i.e. the prior precision.
Using the fact that (() is the negative log- ong-nmmratorr of (A.1), we can get p(D ()p(O) by taking the exponential of (()). By doing
have:
p(D p)) = exp( -())) exp(-_((0.) 2 (0- 0.) (H+aaI(A 0.)).
(A.5)
For simplicity, we can define = (H + aI) and rewrite the equation above to obtain:
p(D 0)p(0) ~ exp( exp( - -0- 0.)))
(A.6)
We can then use this approximation to estimate the normalizing constant our approximate posterior, which corresponds to the marginal log- log-likeldhood p(D):
(D) Z x [ exp(- exp( -00 -.) ' -..)))do.
(A.7)
This can be rewritten
p(D) exp(-L(0.)) exp( 1 0.) (0-00)))0--0.))) 0.))),
(A.8)
and by using the Gaussian integral properties can write as:
(DD ~ exp( exp(( (det 1)
(A.9)
and by taking the logarithm we get
log log)(D) ~ C(O.) + log((2))(22) + log(det 2/1/2
(A.10)
which is the approximation of the log marginal likelihood that we want maximize optimize the parameters that appears in L(O.). In a regression problem, we are interested in optimizing both the variance of the Gaussian likelihood and the prior precision. classification, instead, we just have the prior precision hyperparameter to tune.
B Riemannian geometry
We rely Riemannian geometry [Lee, 2019, do Carmo, 1992] in order to construct approximate posterior. In a nutshell, a d-dimensional Riemannian manifold can be seen intuitively as a smooth d-dimensional surface that lies within Euclidean space of dimension D d, which allows one to compute distances between points that respect the geometry of the surface.
Definition B.I. Riemannian manifold M a smooth manifold together with a Riemannian metric M(x) that acts the associated tangent space T at any point M.
Definition B.2. A Riemannian metric M M - Rdim(M)xdi((MM) is smoothly changing positive definite metric tensor that defines inner product the tangent space T&M at any point x e M.

---

We focus on the Bayesian neural network framework where θ = R is the parameter space of the associated deep network, and we consider the manifold .M g(0) = [8. This is essentially the loss surface which is K --dimensional and lies within (K + 1) )-dimensional Euclidean space. In order to satisfy the smoothness condition for M we restrict to activation functions as the tanh, while common loss function as the mean squared error and softmax are also smooth.
The parameter space θ is parametrization of this surface and technically represents the intrinsic coordinates of the manifold, which is known as the global chart in the literature. The Jacobian _0(0) E I K K+1xk of the map 9 spans the tangent space on the manifold, and a tangent vector can be written as J g (0)v, where IR are the tangential coordinates. We can thus compute the inner product between two tangent vectors in the ambient space using the Euclidean metric therein as
JJ( = (v, (v)
(B.1)
Here the matrix (()) = (0 TK + VoL(O)VoL(O) is a Riemannian metric that is known in the literature as the pull-back metric. Note that the flat space θ is technically a smooth manifold and together with M(0) transformed into a Riemannian manifold. This can be also seen as the abstract manifold definition where only need the intrinsic coordinates and the Riemannian metric to compute geometric quantities, and not the actual embedded manifold in the ambient space.
One example of a geometric quantity is the shortest path between two points 02, 02 E O. Let a curve c: [0. 1] + with c(0) 01 and c(1) 82, the its length defined under the Riemannian metric as length(c] [c] = To V(c(t), M(c(t))(tt)d. This quantity is computed intrinsically but it also corresponds the length of the associated curve M. The shortest path then is defined as e² (t) arg min. length[c] [c]] but as the length is invariant under re-parametrizations of time, we consider the energy functional instead, which we optimize using the Euler-Lagrange equations. This gives the following system of second-order nonlinear ordinary differential equations (ODES)
a(t) = M (c(t((
(M(((()) (cg(t)
(((()))) Ôvec ((c(t)(((t))] dc(t)
2
(é(t) @ c(t)), (B.2)
2
ack(1) (t)
where vec[-] stacks the columns of matrix and @ the Kronecker product [Arvanitidis et al., 2018). curve that satisfies this system is known as geodesic and is potentially the shortest path. When the system is solved as a Boundary Value Problem (BVP) with initial condition
c(0) = 0₁ and c(1) 02. we get the geodesic that connects these two points. Let v = E(0) be the velocity of this curve at t = O. When the system is solved as an Initial Value Problem (IVP) with conditions c(0) 81 and c(0) = v, we get the geodesic cv (t) between cv (0) = 02 and cv (1) = 02. This operation is known as the exponential map and we use it for our approximate posterior.
In general, an analytic solution for this system of ODES does not exist [Hennig and Hauberg, 2014, Arvanitidis et al., 2019], and hence, we rely on an approximate numerical off-the-shelf ODE solver. Note that this is a highly complicated system, especially when the Riemannian metric depends on a finite data set, while it is computationally expensive evaluate it. As we show below, the structure of the Riemannian metric that we consider, allows us to simplify significantly this system.
Lemma B.3, For the Riemannian metric (3) the general ODES system in (B.2) becomes
vol(c(t)) Vol(c(t))) (c(1), (((((((((()))((
a()==- 1+
1+ 1
(B.3)
Proof. We consider the general system, and we compute each term individually. To simplify notation we will use M = M(c(t)), and V, the i-th element of the gradient V.
VoL(c(t)). H := Ho (c)((())) H, the i-th column of the Hessian
Using the Sherman-Morrieo formula have that
V Volle(t)VV&l((tt)
M(c(t))-1 =I 1K
(B.4)
(VoCic()
The first term in the brackets is = 2 [H, VT + VH. V......, H + VH DXD=
2 aM((())
aM(c(I))
(B.5)
(ty(t)
(cx(t)
3

---

and the second term in the brackets is
Ôvec (M((tt))] ac(t)
[ HKH+H E
(B.6)
and their difference is equal
H* V VKH. [2 H V,H, V.H,...,,VVI, 2VH
(B.7)
We compute the matrix-vector product between the matrix (B.7) and the Kronocker product e(t) x
c(t)
RpDxx which gives
K
K Σ
K V
K
2VH + H HćH = 2V
Σ H
- Hè
V = (é,
1=1
==1
(B.8)
where we used that E H24: = Hé and E V = VT As a final step, we plug-in this result and the inverse of the metric in the general system which gives
1
V
T
2V
(c, Hc).
(B.9)
(V
Taylor expansion. As regards the Taylor expansion we consider the space θ and the Riemannian metric therein M(()) and arbitrary smooth function f θ + R. If ignore the Riemannian metric the second-order approximation of the function f around a point x E known to be
f (x + vv = f(xx (Vof(0)) v) 2 (v. He 10)
(B.10)
where Vof(0)) is the vector with the partial derivatives evaluated at x and Hş [f) d=x the corresponding Hessian matrix with the partial derivatives 8 f(0) 80,00, 0==x
When we take into account the
Riemannian metric, then the approximation becomes
fRimm + = f(x) (Vof()) v) + (H - 1 - [of(0)
lo=xv), v), (B.11)
where r are the Christoffel symbols and the Einstein summation is used. Note that even the Hessian different, the approximation again is quadratic function.
Now we consider the Taylor expansion on the associated tangent space at the point x instead of directly on the parameter space O. We define the function h(v) = [[Expx on the tangent space centered at x and we get that
huu) = h(0) + (a. (Expx u)
(B.
1 + (u, (u, (Expx,))))))) - T, (Expx v>=
u),
(B.13)
where we apply the chain-rule and we use the fact that O,Exp, (v) Iand oe Exp, (v)
v= 0. The
So, we get that 8 (Exp = Vol(O) and H [[ v=0 26, 38, difference here is that the quadratic function is defined on the tangent space, and the exponen- tial map is non-linear mapping. Therefore, the actual approximation Frangeen (Exp, (v)) h(v) is not quadratic function as before, but adapts the structure of the Riemannian metric. Intuitively, the closer a point is to the base point x with respect to the Riemannian distance, the more similar is the associated value fTangem (Exgeem(Exxxx (v)) to f(x). In our problem of interest, this behavior desirable implying that a parameter 0' connected through low-loss regions with continuous curve l., then will assign to 0 high approximate posterior density.
5²
the parameter space

---

We can easily consider the approximation on the normal coordinates, where we know that the Christoffel symbols vanish, by using the relationship u Aü with A = M(x)-V; and thus the Taylor approximation the function h(u) (Exp, (Au)) becomes
(()) h(0) (AT p-x: Au) + (i, (x)
(B.14)
Further details can be found in related textbooks [Absil et al., 2008] and articles [Pennec, 2006].
Linearized manifold. Let consider regression problem with likelihood p(y|x, 0) ((y (x) where fe is a deep neural network, and prior p(0) [N000, Alx). The loss of the f (x) then defined
_ (0)
(y (x.) (Vo[(xa) 0=0.- 0.))² + 1112.
(B.15)
20
The gradient and the Hessian of this loss function can be easily computed as
N
Voc 1i (0)
δ² (y
Σ
(x -- +
(B.16)
n=l
N
1
Hg (p)
Vof(xx))o-b Vof(xn) 2
(B.17)
σ²
which can be used to evaluate the ODES system in (6). A similar result can be derived for the binary cross entropy loss and the Bernoulli likelihood.
C Implementation details
In this section we present the implementation details of our work. The code will be released upon acceptance.
Gradient, Hessian, and Jacobian computations. The initial velocities used for our methods are samples from the Laplace approximation. We rely in the Laplace library [Daxberger al., 202 la] for fitting the Laplace approximation and for optimizing the hyperparameters by using the marginal lig-likeihood. We also used the same library to implement all the baselines consider in this work. As we have seen from Sec. 3.3, to integrate the ODE we have to compute (6), which we report also here for clarity:
č(t) = -VooC(((((() (1 D (600, (d0,
(C.1)
We use functorch [Horace He, 2021] to compute both the gradient and the Hessian-vector -ector--rooo...uu.t...... We then rely on scipy [Virtanen et al., 2020] implementation of the explicit Runge-Kutta method of order (4) [Dormand and Prince, 1980] to solve the initial-value problem. We use default tolerances in all our experiments.
Linearized manifold. We define linearized manifold by considering the "linearized" function fl [x. + (VoJo(x) 8=0, 0.) to compute the loss, where VaJo(x) ø=0_ E RCCK K is the Jacobian. To compute (Vofo(x) 0=a_- 9 0_) we use functorch to compute a jacobian-vectoe product. This way, we avoid having to compute and store the Jacobian, which for big networks and large dataset is infeasible to store.
C.1 Cost of solving the ODE system
Our proposed approach builds on top of Laplace approximation, and therefore comes with computational overhead. To be able to estimate this overhead, we have to assume that our solver does steps to solve an ODE, which not entirely correct, since the explicit Runge-Kutta method of order 5(4) (RK-45) [Dormand and Prince, 1980] uses an adaptive step-size and therefore the value of S differs for every ODE. With this assumption, the cost of evaluating single ODE results O(SWN) N). where S is the number of steps of the solver, W is the number of model parameters and N is the
5

---

5000
pazzza
100-
parere (611)
36 percee (llg]
parzzn (dee)
5000
1000
1000
0000
1000
50.000
10000
Figure C.I: Left: Average runtime per posterior sample in terms of model size and dataset size. Right: Average runtime and standard eтror per posterior sample in terms of model size and dataset size computed over different samples. For the models with 5k and 10k parameters, we used the diagonal Hessian over all the weights to sample the initial velocities.
dataset size. Indeed, at each step, need compute the gradient and the hessian- vector product, which is ONNW), and perform a Runge- Kutta step, which scales linearly with the dimension of the problem, i.e. O(W) [Heirer et al., 1987]. Since the solver would take S steps to solve the system get OSSNW). If we want to use K posterior samples to estimate our predictive distribution, then we have solve K different ODES, therefore the overall overhead is O(K SNW). However, want to highlight that all these computations happen before deployment of the model. At test time, the complexity is the same as LA when using MC sample to approximate the predictive distribution.
In the derivation of the complexity of our method, we assumed that the ODE solver did exactly S steps. As mentioned before, RK-45 uses adaptive step-size, therefore the number of steps the solver needs to converge mainly depends on the complexity of the associated ODE problem, which is defined by the geometry of the loss landscape and the initial velocity. To show this, in the next section, we include a benchmark using a synthetic example to analyze how the runtime changes with respect the size of the model and dataset. While increasing the model parameters and dataset size affect W and N, may be the case that ODE systems do not necessarily get harder, therefore the value of S would be small. In other words, evaluating the ODE becomes more expensive dimensions increase, but perhaps the actual system gets easier to solve.
Limitations of our implementation The cost of obtaining a posterior sample using RIEM-LA scales linearly in terms of the model and dataset size. addition to that, our implementation is not highly engineered. Although Python solver (scipy. integrate solve_ivy [Virtanen et al., 2020J) to solve the ODE system, which runs the CPU. When the solver needs evaluate the ODE, based approach runs on the GPU, and the result is moved to the CPU (. (-detach() cpu () numpy () causing significant overhead, which is inefficient. Especially when dimensions increase, both the transfer of the data and the computations on the CPU are sub-optimal. Implementing an ODE solver on suitable iifomati--iifferentattio framework (e.g. JAX) solely running GPU will dramatically improve performance.
generate these samples in parallel,
general purpose
nutommtii-diiffrenttatio
C.2 Analysis of exponential map
We mentioned scalability as one of the main limitation of our proposed method. We use a simple toy-example (skl earn datasets. make_ make_blobs with three classes, see Fig. C.I (leff)) to analyze the running time of computing a single sample from the posterior with our methods. A two layers neural-networf is considered, and we vary the number of parameters and datapoints in the dataset. For the biggest models, i.e. the one with 5k and 10k parameters, we use diagonal approximation of the Hessian to sample the initial velocities, while for all the others we use the full Hessian. In all cases, the ODE system always relies on the hvp which uses the full Hessian. In Fig. C.1 (right), we present runtime for five different models and six dataset sizes. We report the mean and standard deviation computed over five different exponential maps. The results show that the runtime increases

---

for all models as the dataset size increases. We also see that using diagonal approximation of the Hessian leads to faster exponential maps.
D Experiments details and additional results
D.1 Regression example
For the regression example we consider two fully connected networks, one with one hidden layer with 15 units and one with two layers with 10 units each, both with tanh activations. We train both model using full-dataset GD, using weight decay of l 2 for the larger model and le 3 for the smaller model for 35000 and 700000 epochs respectively. In both cases, we use learning rate of le Sec. 5, we show some samples from the posterior distribution obtained by using vanilla LA and our RIEM-LA approach. In Fig. D.I we report the predictive distribution for our classic approach while in Fig. D.2 we show both the posterior and the predictive for our linearized manifold and linearized LA. We can see that our linearized approach perform similarly to linearized LA in terms of posterior samples and predictive distribution.
In
more interesting experiment is to consider a gap in our dataset to measure the in-between uncertainty. A good behaviour would be to give calibrated uncertainty estimates in between separated regions of observations [Foong et al., 2019].
our regression example, we consider the points between 1.5 and 3 as test set, and report posterior samples and predictive distribution for both our methods, vanilla and linearized, and LA. From Fig. D.3, we can notice that our RIEM-LA is able to generate uncertainty estimates that are reliable both in the small and the overparametrized model. It gets more interesting when we consider our linearized manifold approach as it can be seen in Fig. D.4. While for small model the predictive distribution we get is comparable to linearized LA, when consider the overparametrized model, it tends to overfit to a solution that it is better from the perspective of the linearized loss but very different from the MAP. Therefore, this results in uncertainty estimates that are not able to capture the true data set in this case. This behaviour is related to our discussion in Sec. 3.3. By using subset of the training to solve the ODES system we alleviate this behavior by giving more reliable uncertainty estimates.
For each all the regression setting described above, we also report the mean and the uncertainty estimates by sampling from a closer approximation of the true posterior obtained using Hamiltonian Monte Carlo (HMC) and No-U-Turn sampler. We use one chain, samples 1000 weights, discard the first 500 as warm-up and use the last 500 to compute the predictive distribution. Results can be seen in Fig. D.5. We can that in this simple example the uncertainty produced by our method is really close to the one produced by HMC.
D.2 Illustrative 2D classification example
We consider the banana dataset as an illustrative example to study the confidence of proposed method against Laplace and linearized Laplace. We train a 2-layer fully connected neural net with 16 hidden units per layer and tanh activation using SGD for 2500 epochs. We use learning rate of le - 3 and weight-decaay of le 2. Although it is a binary classification problem, we use network with two outputs and use the cross-entropy loss to train the model because the Laplace library does not support binary cross-entropy at the moment. For all methods, we use 100 MC samples for the predictive distribution.
As we have mentioned in Sec. 5, our linearized manifold when we solve the ODES system by using the entire training set tends to be overconfident compared to the our classic loneaineer approach. This can be easily seen form Fig. D.6 where we can compare the last two rows and see that solving the ODES system using batches is beneficial in terms of uncertainty quantification. We also plot the confidence of all different methods when we do not optimize the prior precision. We can see that our proposed approaches are robust to it, while línearized LA highly depends on it.
D.3 An additional illustrative example
We consider the pinwheel dataset with five different classes as an additional illustrative classification example. Given the way these classes clusters, it is interesting to see if some of the approaches able to be confident only in regions where there is data. We generate a dataset considering 200
7

---

Outs
MAP LA mann
LA
Lapiorn
Out moren
Test points
SLAP
P D'gg
BVV lakk
POSTERIOR SINGLE LAYER MODEL
PREDICTIVE SINGLE LAYER MODEL
Outs
MAP
Lapbon meam
Our meee
Tres points
PV/ W/00
Pv/
1-
2-
-2
POSTERIOR TWO LAYERS MODEL
PREDICTIVE TWO LAYERS MODEL
Figure D.I: Regression results in terms of posterior and predictive distribution using vanilla LA and our RIEM-LA approach. Our proposed approach is able to get better samples and to give a reliable uncertainty estimate compared to vanilla LA. If we consider overparametrized model.
examples per classes and we use 350 example as training set and the remaining as test set. We consider two layer fully-connected network with 20 hidden units per layer and tanh activation. As before, train it using SGD with learning rate of 1e 3 and weight decay of 1e for 5000 epochs.
From Fig. D.7, we can see not only vanilla LA but also linearized LA is failing in being confident also in-data region when we use 50 posterior samples. Our proposed approaches instead give meaningful uncertainty estimates, but we can see that by optimizing the prior precision, our methods get slightly more confident far from the data.
D.4 Additional results in the UCI classification tasks
In the main paper we present results on five different UCI classification tasks in terms of test negative log-likelihood. Here, report also results in terms of test accuracy, Brier score, and expected calibration error. Results with prior precision optimized are shown in Table 5, while for prior precision not optimized we refer to Table 6. In both cases, consider neural network with single fully connected layer consisting of 50 hidden units and tanh activation. We train using Adam optimizer [Kingma and Ba, 2015] for 10000 epochs using learning rate of le 3 and weight decay of le
From the two tables, we can see that our RIEM-LA is better than all the other methods both in terms of NIL,, brier, and ECE. It is also surprising that our method, both the classic and the linearized approach, are able to improve the accuracy of the MAP estimate in most datasets.

---

Ous
MAP LA mann
LA
Laplorn
Our tiren
Test points
SEAP
P g/gi
BVV lakk
POSTERIOR SINGLE LAYER MODEL
PREDICTIVE SINGLE LAYER MODEL
Outs
MAP
Lapbon mean
Out meee
TTes points
PV/ 1/00
2-
2-
22
POSTERIOR TWO LAYERS MODEL
PREDICTIVE TWO LAYERS MODEL
Figure D.2: Regression results in terms of posterior and predictive distribution using linearized LA and our linearized manifold approach. Our linearized manifold perform similarly to linearized LA no matter what model architecture are considering.
D.5 Complete results on MNIST and FMNIST
For these two image classification tasks we consider a small convolutional neural network. Our network consists of the following layers: initial convolutional layer with channels and 5 x filter followed by tanh activation and an average pooling layer. The we have another convolutional layer still with channels and 5 x 5 kernel also followed by tanh activation and average pooling layer. Then we have three fully connected layer with 16, 10, and 10 hidden units respectively and tanh activation.
We train both models using SGD with a learning rate of le 3 and weight decay 5e for 100 epoch. The learning rate is annealed using the cosine decay method [Loshchilov and Hutter, 2017].
In Table and Table 8, we can that also when we do optimize the prior precision our RIEM-LA is mostly performing better than all the alternatives. In particular, for the CNN trained on MNIST and no prior precision optimization,
have that the MAP is also performing well in terms of NLL and Brier score. On FashionMNIST.. instead, if we do not optimize the prior we have that both our approaches are better than all the other methods.
D.6 Out-of-distributtio results
The benefit of having meaningful and robust uncertainty estimation is that model would then be confident in-data region while being uncertain in region without data. Therefore, if this is happening. then would expect the model to be able to detect -iitt--iistiibuutinn (OOD) examples more successfully.
We consider classic OOD images detection tasks, where we train a model on MNIST and tested on FashionMNISTT, EMNIST, and KMNIST and one trained on FMNIST and tested on remaining datascts. It's well known that linearized LA is one of the strongest method for OOD detection

---

Out
MAP
LA
LA main
Out moren
Test points
SLAP
PVT g'gg PVV BiLk
POSTERIOR SINGLE LAYER MODEL
PREDICTIVE SINGLE LAYER MODEL
MAP
Trsn points
PV/ D'O
PV/
Chane
Laplane meen
-2
Oue noan
MAP
POSTERIOR TWO LAYERS MODEL
PREDICTIVE TWO LAYERS MODEL
Figure D.3: In-between uncertainty example using vanilla LA and riem-la la approach. Also in this setting, our vanilla approach is able to overcome difficulties vanilla LA in this problem setting and giving both meaningful samples from the posterior and reliable predictive distribution.
[Daxberger et al., 2021a]. We consider the same MAP estimates we used in the previous section and present OOD performance Table and Table 10.
We can see that our proposed method consistently working better than linearized and classic LA in all the considered setting apart for models trained FMNIST where do not optimize the prior precision to compute our initial velocities. In that setting, however, our linearized approach using batches is getting similar performance than linearized LA.
10

---

Out
MAP
Lagiaro meam
Our mnan
Out ticsa
Tesk points
Pv/V GYo
By/V
11
POSTERIOR SINGLE LAYER MODEL
PREDICTIVE SINGLE LAYER MODEL
moun Our mnann
Laplorn mean
Our
Test points
Pv/V Gi
By/V TIAIA
11
2-
-2-
POSTERIOR TWO LAYERS MODEL
PREDICTIVE TWO LAYERS MODEL
mam. Our
Lapleon
Tesh points
XEAP
PyOog BulV Tiki
-3
POSTERIOR- TWO LAYERS MODEL (BATCH)
PREDICTIVE: TWO LAYERS MODEL (BATCH)
Figure D.4: In-between uncertainty example using linearized LA and our linearized manifold. For linearized approach both consider solving the ODES system using the entire training set and using subsets of it. We already highlight how our linearized approach tend to overfit on the linearized loss, and we clearly it here from the plots in the center. Using batches to solve the ODES system gives more reliable uncertainty estimates.
11

---

Productive
Toxt poato 31
Toi paaale 21/
PREDICTIVE SINGLE LAYER MODEL
PREDICTIVE SINGLE LAYER MODEL
Prodictvv dirribraiinn HNC
Paslictive ditribution HMC
potza 2/
poiate
2y P
PREDICTIVE TWO LAYERS MODEL
PREDICTIVE TWO LAYERS MODEL
Figure D.5: Predictive distribution obtained using HMC all the different regression examples considered using a one-layer and a two-layers model.
Table 4: Full in-distribution results in the banana dataset. We used 100 samples both for the Laplace approximation and our method. ECE and MCE are computed using M 10 bins.
BANANA DATASET
PRIOR PRECISION NOT OPTIMIZED
МETHOD
Accumcy
NI.IJJ
Brier]
BCE
MCFJ
0.34
0.333 0.005
0.0930
5.91 + 0.34
22.75
VANILLA
48.85
2.32
0.700 0.003
0.2534
0.0017
7.10 1.96
23.10 6.69
LIN-LA
86.92 + 0.40
0.403 0.012
0.1196 0.0044
13.04 0.60
17.11
0.54
RIEM-LA (RATCHES)
87.32
0.17
0.294 + 0.002
0.0004 0.0895
4.79 0.31
8.28 + 0.84
LIN-RIEM-LA
85.33
0.31
0.884 0.037
0.1252
11.50 0.31
36.27 2.84
LIN-RIEM-LA (RATCHES)
0.21 86.16
0.352 0.002
0.0994 0.0011
4.06 0.11
13.67 0.69
PRIOR PRECISION OPTIMIZED
METHOD
Accumay
NL.1
Brier]
ECE.
MCFJ
MAP
86.69 + 0.34
0.333 + 0.005
0.0930
0.0015
5.91 4 0.34
22.75
2.42
VANILLA LA
59.50 1: 5.07
0.678 0.009
0.2426
0.0046
12.08 1 2.31
27.49
5.05
86.99 0.37
0.325 0.008
0.0956
0.0023
6.59 + 0.41
11.44
1.85
RIEM-LA
87.57 1: 0.07
0.287 0.002
0.0886
0.0006
2.55 1 0.37
10.95
1.16
RIEM-LA (BATCHES)
0.08
2.57 0.004
0.67
LIN-RIEM-LA
87.02 + 0.38
0.415 0.029
0.0967
0.0021
6.97 0.43
21.21
1001:20 2.28
LIN-RIEM-LA (BATCHES)
87.77
0.24
0.298 0.006
0.0887 0.0011
2.38 4 0.28
8.04 1 1.61
12

---

MAP SOLUTION
PRIOR OPTIMIZED
PRIOR NOT OPTIMIZED
VANILLA LA
VANILLA LA
RIEM-LA
RIEM-LA
LIN-LA
LIN-LA
I.N-RIEM-LA
LIN-RIEM-LA
LIN-RIEM-LA (BATCH)
LIN-RIEM-LA (BATCH)
Figure D.6: Confidence for all the considered approach on the banana dataset. All plots were computed using 50 MC samples from the posterior. We can notice that our RIEM-LA and LIN-RIEM- LA with batches are giving the best confidence of all methods followed by linearized LA. We can also that our approaches are robust prior optimization, while linearized LA gets really better if we optimize the prior precision.
13

---

MAP SOLUTION
92
PRIOR OPTIMIZED
PRIOR NOT OPTIMIZED
VANILLA LA
VANILLA LA
RIEM-LA
RIEM-LA
LIN-LA
LIN-LA
LIN-RIEM-LA
LIN-RIEM-LA
Figure D.7: Confidence for all the considered approach on the pinwheel dataset. All plots were computed using 50 MC samples from the posterior. We can notice that in this setting RIEM-LA gives the better confidence followed by LIN-RIEM-LA. However, priorprecision optimization is making our approaches more confident outside the data region. Linearized LA struggle in producing meaningful uncertainty estimates.
14

---

Table 5: Results for all the different techniques on UCI datasets for classification. Predictive distribution is estimated using 30 MC samples when prior precision is optimized. Mean and standard emror over different seeds.
PRECISION OPTIMIZED
ACCURACY
DATASET
MAP
VANILLA LA
RIEM-LA
LINEARIZED LA
LINEARIZED RIEM-LA
VEHICLE
79.52 + 0.99
49.61 + 3.05
80.47
1.12
73.86 = 1.47
78.58 + 1.31
IONOSPHERE
56.04 + 2.30
58.11 4.04
90.19
1.95
80.38
1.57
88.30 + 1.12
WAVEFORM
82.27
1.66
64.27
+ 3.19
84.93
1.42
17.07
1.52
93
1.08
$2.30
1.00
56.92
- 2.33
84.81
1.03
75.00
1.90
82.69 + 0.82
BREAST CANCER
96.08
0.73
71.96
1 8.58
96.86
0.85
94.90
0.89
96.08
1.30
NLL
DATASET
VANILLA LA
RIEM-LA
LINEARIZED
LINEARIZED RIEM-LA
1:
0.454
GLASS
2.084 1: 0.323
1.737 0.037
1.047
1 0.224
1.365 1: 0.058
1.359 1: 0.299
IONOSPHERE
1.032 + 0.175
0.015
0.
WAVETORM
1.076 + 0.110
0.888 0.030
0.459
+ 0.057
0.640 = 0.002
0.575 0.065
AUSTRALIAN
1.306
+ 0.146
0.684 0.011
0.541 + 0.053
0.570 : 0.016
0.833 0.108
CANCER
0.225
0.076
0.594 0.030
0.176
0.092
0.327 1 0.022
0.202 0.073
DATASET
MAP
VANILLA LA
RIEM-LA
LINEARIZED LA
LINEARIZED RIEM-LA
VEHICLE
0.0877 +: 0.0052
0.1654
0.0028
0.0871 1: O 0.0035
0.1210 + 0.0028
0.0720 +: 0.0048
KONOSPHERE
0.1308 + 0.0182
0.0064 0.2308
0.0840 1 0.0127
0.1596 + 0.0069
0.0868
0.0095
WAVEFORM
0.1043 + 0.0091
0.1766
0. 0.0825
+ 0.07
0.1260 = 0.0011
0.0871
0. 0.0053
AUSTRALIAN
0.1642 4 0.0105
0.2152
0.0052
0.1909
1 0.0113
0.1901 + 0.0069
0.1340
0.0082
BREAST CANCER
0.0351 1 0.0077
0.2020 0.0148
0.0269 0.0075
0.0866 1 0.0082
0.0310 0.0094
BCE
DATASET
MAP
VANILLA LA
RIEM-LA
LINEARIZED
LINEARIZED RIEM-LA
2.41
00
GLASS
31.77
1: 2.67
11.77 1 1.08
17.06
2.96
23.17
1.47
20.84 1 2.57
11.22 4:
9.05 +
WAVEDORM
15.73
1.29
18.73 + 2/11
10.72
1.31
19.12
1.60
9.10 0.94
AUSTRALIAN
16.25
0.96
5.55
1.41
8.49
+ 1.62
12.42
1.43
10.20 - 0.90
CANCER
3.85 1 0.89
20.39
1 3.28
3.20
0.72
20.24
0.99
3.16 0.92
15

---

Table 6: Results for all the different techniques on UCI datasets for classification. Predictive distribution is estimated using 30 MC samples when prior precision is not optimized. Mcan and standard error over different seeds.
PRIOR PRECISION NOT OPTIMIZED
ACCURACY
DATASET
MAP
VANILLA LA
RIEM-LA
LINEARIZED LA
LINEARIZED RIEM-LA
VEHICLE
79.52
0.99
23.94 + 1.62
82.05 = 1.10
44.88
1.74
80.16 + 1.45
1:
IONOSPHERE
86.04
2.30
47.55 -0 1.72
91.32 = 1.26
69.81 = 4.10
89.81 + 0.68
WAVEFORM
82.27
1.66
24.13 - 4.30
83.60
= 1.34
54.0
2111
83.47 + 1.28
82.30
1.00
40.19 + 1.77
86.73
1.10
61.92
2.57
81.04
1 0.80
BREAST CANCER
96.08
0.73
51.37
11.68
97.06
0.73
84.51
3.72
95.29
1.22
NLL J
DATASET
MAP
VANILLA
RIEM-LA
LINEARIZED LA
LINEARIZED RIEM-LA
+
1.197 1197:00000
GLASS
2.084 1 0.323
2.057 1 0.084
0.981 1. 0.202
1.697 0.044
1.116 1 0.270
IONOSPHERE
+
0.726 0.011
0.014
WAVETORM
1.076
0.110
1.185 = 0.031
0.465
0.051
0.962 = 0.015
0.456 + 0.059
AUSTRALIAN
1.306 U.146
0.750 0.009
0.518
0.070
0.658 + 0.012
0.593
0.030
CANCER
0.225
0.076
0.690 0.051
0.197
0.093
0.503 0.0222
0.171 0.064
DATASET
MAP
VANILLA LA
RIEM-LA
LINEARIZED LA
LINEARIZED RIEM-LA
VEHICLE
0.0877 + 0.0052
0.1917 + 0.0020
0.0604
0.0024
0.1645 1: 0.0023
0.0657 + 0.0027
KONOSPHERE
0.1308 = 0.0182
0.2654 = 0.0053
0.0689
0.0076
0. 1596 + 0.0069
0.0868 + 0.0095
WAVEFORM
0.1043 0.0091
0.2403 0.0068
0.0789
0056
0.1929 + 0.0033
0.0821 = 0.0060
AUSTRALIAN
0.1642 = 0.0105
0.2774
0.0042
0.1081
0. 0098
0.2328
+ 0.0059
0.1205 + 0.0072
BREAST CANCER
0.0351
0.0077
0.2486
0.0248
0.0283 0. 0088
0.1597 1 0.0106
0.0316 0.0080
ECE 1
DATASET
MAP
VANILLA
RIEM-LA
LINEARIZED LA
LINEARIZED RIEM-LA
16.75
5.43
GLASS
31.77 2.67
2.67
15.27 1: 1.11
14.84 1 2.23
8.18 1 2.30
19.63 1 1.74
JONOSPHERE
13.92
4:
13.25
+
WAVEDORM
15.73
1.29
18.50
- 1-11
6.34 + 1.01
10.17 1.74
5.32 = 0.94
AUSTRALIAN
16.25
0.96
16.98
+ 1.70
5.92
1.50
7.17 - 0.79
6.36 0.97
3.85
0.89
23.00
1 5.45
3.27
0.68
21.63 3.01
3.09 1 0.55
16

---

Table 7: Results simple CNN architecture MNIST dataset. We train the network 5000 examples and test the in-distribution performance on the test set, which contains 8000 examples. We used 25 Monte Carlo samples to approximate the predictive distribution and in cases used a subset of the data to solve the ODEs system we rely on 1000 samples. Calibration metrics are computed using 15 bins.
MNIST DATASET
PRIOR PRECISION NOT OPTMIZED
METHOD
Accuracy 1
NLL.
Brier
ECEL
MCEJ
MAP
95.02 + 0.17
0.167 1: 0.005
0.0075
0.0002
1.05 0.14
39.94
14.27
VANILLA LA
8.10 = 0.74
2.521 + 0.006
0.0937 + 0.0000
11.79 0.83
35.09 = 6.47
LIN-LA
0.0143
0.09
RIEM-LA
96.18 0.23
0.177 + 0.007
0.0073 = 0.0003
7.10
25.41
1.96
RIEM-LA (BATCHES)
94.98 + 0.20
0,28-1 0.008
0.0111 0,0001
13.64
0.20
29.23
0.50
LIN-RIEM-LA
95.12 1 0.27
0,180 1: 0.006
0.0080
0.0003
4.19 1 0.07
22.02
3.97
LIN AI--RIM--L (BATCHES)
94.63 1: 0.19
0.229 1 0.007
0.0097 0.0004
7.84 1 0.18
28.74 1 3.67
PRIOR PRECISION OPTIMIZED
METHOD
Accuracy t
NIL.J.
Brier
BCEL
MCEJ
0.005
39.94
VANILLA
88.69 + 1.84
0.871 + 0.026
0.0393 0.0013
42.11
50.52
1.45
LIN-LA
94.91 + 0.26
0.204 +: 0.006
0.0087
0.0003
6.30 0.00
39.30 1: 16.77
RIEM-LA
96.74 0.12
0.115 1. 0.003
0.0052
0.0002
2.48 0.06
38.03
15.02
RIEM-LA (BATCHES)
95.67 + 0.19
0.170 + 0.005
0.0072
0.0002
6.40
22.40 = 0.51
LIN-RIEM-LA (BATCHES)
95.14 + 0.20
0.167 + 0.004
0.0076 + 0.0002
3.23
18.10 : 2.50
Table 8: Results simple CNN architecture FMNIST dataset. We train the network 5000 examples and test the in-distribution performance on the test set, which contains 8000 examples. We used 25 Monte Carlo samples approximate the predictive distribution and in cases used a subset of the data to solve the ODEs system we rely on 1000 samples. Calibration metrics are computed using 15 bins.
FMNIST DATASET
PRIOR PRECISION NOT OPTIMIZED
METHOD
Accuracy f
NLL.
Brier]
BCEL
MCE.
MAP
79.88 1: 0.09
0.541 1. 0.002
0.0276 0.0000
1.66 1 0.07
24.07 1: 1.50
VANILLA LA
10.16 + 0.59
2.548 4 0.050
0.0936 1: 0.0007
10.79
1.05
27.19 + 6.64
RIEM-LA
82.70 = 0.23
0.522
0.004
0.0262 1 0.0002
9.95 0.38
19.11 - 0.43
RIEM-LA (BATCHES)
80.77
+ 0.10
0.582 + 0.004
0.0285
0.0001
10.34 0.00
18.52
0.64
LIN-RIEM-LA
80.91
+ 0.20
0.528 + 0.004
0.0265
0.0002
1.59 = 0.118
17.58
= 3.60
LIN-RIEM-LA (BATCHES)
79.95
+ 0.22
0.567 + 0.005
0.0281 + 0.0002
4.79 - 0.58
16.30 + 1.88
PRIOR PRECISION OPTIMIZED
METHOD
Accuracy 1
NLL.
Brier]
ECEL
MCEJ
+:
1.66
1.50
VANILLA
74.88 + 0.83
1.026 + 0.046
0.0482
0.0019
31.63
1.28
43.61 - 2.95
LIN-LA
79.85 + 0.13
0.349 + 0.001
0.0278
0.0000
3.23
37.88
17.98
RIEM-LA
83.33
0.17
0.472 1 0.001
0.0237
0.0001
3.13 0.48
10.94 1 2.11
RIEM-LA (BATCHES)
81.65
1 0.18
0.525 1 0.004
0.0263
0.0002
5.80
0.73
35.30 4: 18.40
LIN RIEM-LA (BATCHES)
80.49
1 0.13
0.529 1 0.003
0.0269 1: 0.0002
2.10
6.14
1.43
17

---

Table 9: AUROC T for OOD tasks for models trained MNIST and tested against FashionMNIST, EMNIST, KMNIST. (B) indicates exponential maps solved using a subset of the training set. We can notice that our proposed RIEM-LA is consistently performing better than all the other approaches for OOD detection.
PRIOR PRECISION NOT OPTIMIZEO
VANILLA
RIEM-LA
RIEM-LA
RIEM-LA
LIN. RIEM-LA (B)
FMNIST
0.777 0.057
0.523 + 0.020
0.911 + 0.011
0.876 2876-+0.001 0.018
0.573 + 0.023
0.801 +: 0:0.013
0.862 + 0.026
EMNIST
0.851
0.008
0.490 0.015
0.900 0.004
0.897 0.005
0.905 1. 0.003
0.872 0.005
0.901 0.004
XMNIST
0.877
0.005
0.511 0.013
0.949 0.002
0,930 0.002
0.941 1 0.002
0.938 0.003
0.941 0.002
PRIOR PRECISION OPTIMIZED
0O0 DATA
MAP
VANILLA LA
KIEM-LA
LIN,
LIN. RIEM-LA
RIEM-LA (8)
LIN. RIBM-LA (B)
FMNIST
0.777 = 0.057
0.681 :: 0,029
0.917 1 0.017
0.854 4 0.024
0.809 + 0.027
0.897 0.012
0,841 +: 0.022
EMNIST
0,851 0.008
0.708 0.018
0.917 0.001
0.006
0.871 = 0.003
0.899 0.003
0.858 0.003
KMNIST
0.877 0.005
0.813 + 0.012
0.953 + 0.001
0,921 0.003
0.906 + 0.002
0.948 0.002
0.920 + 0.002
Table 10: AUROC for OOD tasks for models trained on FashionMNIST and tested against MNIST, EMNIST, KMNIST. (B) indicates exponential maps solved using subset of the training set. It is interesting to notice that if we do not optimize the prior precision, then LIN-LA is performing the best together with our linearized approaches. we optimize the prior, RIEM-LA performs the best of all approaches.
PRIOR PRECISION NOT OPTIMIZED
VANILLA
LIN. RIEM-LA
RIEM-LA
RIEM-LA (R)
MNIST
0.715 0.022
0.494 = 0.014
0.895 0.010
0.937 0.012
0.929 0.015
0.917 1. 0.008
0.940 0.012
HMNIST
0.649 0.009
0.4% = 0.013
0.794 0.011
0.907
0.007
0.881 + 0.009
0.814 + 0.016
0.801 0.008
KMNIST
0.718
0.495
0.898
PRIOR PRECISION OPTIMIZED
0O0 DATA
MAP
VANILLA
RIRM-IA
LIN. 1.4
LIN. RIEM-1A
RIFM-LA (8)
LIN. RIFM-LA (R)
MNIST
0.715:1:000 0.022
0.861
0.037
0.921 0.011
0.861 = 0.020
0.807 0.016
0.931 1: 0.005
0.569 0.016
EMNIST
0.649 0.009
0.765 0.050
0.857 0.009
0.806 3 0.004
0.750 0.012
0.851 1. 0.010
0.795 0.001
0.718-+0.00
0.827
0.546
+.0000
0.907
0.009
18


Reviews:
[{"review_id": "ovsMglyjDk", "reviewer": "Reviewer_CqAb", "summary": "The Laplace approximation offers a practical posterior but is limited due to the symmetry of the weight space it is parameterised in. The method proposed to improve posterior quality by adapting the posterior shape through a Riemannian metric that is determined by the log-posterior gradient. ", "questions": "a) What are the method's memory and computational budget limitations?\n\nb) How accurate is the ODE solver when the number of dimensions grows?\nIf the quality of the solutions degrades in higher dimensions, can we expect the method to remain functional for larger models?\n\nc) What are current limitations to scale the method to larger models and datasets (e.g. resnets/transformers models at cifar/imagenet scale)?\n\nd) MacKay also considered notes that the choice of basis for a Laplace approximation is important in [1], which might be a relevant reference. What are the most important reasons to consider the proposed adaptation using the Riemannian metric over other potential changes of the basis?\n\n[1] MacKay, David JC. \"Choice of basis for Laplace approximation.\" Machine learning 33 (1998): 77-86.\n\n", "limitations": "As mentioned, I think being able to adapt the quality of Laplace approximations by considering geometrical aspects of the loss landscape is very interesting. My concerns are mainly in practicality of the approach.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* Practically\n\nThe community is interested in using Laplace approximations for Bayesian posterior estimation due to its simplicity and wide applicability. The proposed paper deals with improving the quality of such approximations, which would benefit many.\n\n* Methodologically\n\nSince the quality of the Laplace approximation is inherent to the space it is parameterised in, it is very interesting to see how Riemmanian geometry of the loss landscapes provides a better understanding of the approximations and provides avenues to improve posterior quality.\n\n* Results\n\nIt is very promising to see that the Riemannian Laplace approximation allows for good posterior fits even when no linearization or prior tuning is being used.", "weaknesses": "Overall, I think the promise of the paper of using the Riemmannian geometry of the loss surface to improve the posterior quality of Laplace approximation is very promising. My concerns mainly lie in the practicality of the approach in terms of scaling the method to larger models (deeper networks and models with more parameters) and larger datasets. Since the paper presents the method in the context of Bayesian deep learning, the scalability of the method is important.\n\n- Scalability to deeper networks\nThe posterior samples from 'POSTERIOR - TWO LAYERS MODEL' of Fig. D.4 seem off. It would be good if the authors addresses how the method would perform for more complex model classes. Does the method break down in this case, or could this potentially be mitigated?\n\n- Scalability to models with more parameters\n\nAs mentioned in the paper, there is a computational cost associated with the growing dimensionality of the parameter space because the number of necessary solver steps increases. I am worried that the method can not be applied to larger deeper NNs as the solutions that are found by ODE solver in practice for larger dimensions will not be of sufficient quality. Since deep neural networks typically consist of many more parameters than the models, this seems like a very big limitation.\n\n- Scalability to larger datasets\n\nThe linear scaling in the number of data points.\n\n- Quantitative results and comparisons\n\nThe method only considers very small models (e.g. single or 2 layer NNs) and small toy problem datasets. This small data regime would allow computing of (close to) exact posteriors, which would allow better quantitative assessment of the posterior found by Riemmanian Laplace approximation. Furthermore, it would be interesting to see larger model and data regimes.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Response to Reviewer CqAb\nWe thank the reviewer for the positive consideration on our work. We also appreciate the time taken to review our paper. We addressed all concerns you highlight under the Weaknesses and Questions sections. \n\n\n**Point 1 / Weaknesses**\n> *... Fig. D.4 seem off. It would be good if the authors addresses how the method would perform for more complex model classes. Does the method break down in this case, or could this potentially be mitigated?*\n\nIn this example, we consider the linearized version of our approach. The data that we use imply a specific loss landscape structure that our samples respect. In the regions where data exist our samples agree with the MAP, while in the other regions the behavior of the samples cannot really be \"judged\" but potentially the uncertainty. Interestingly, if the prior is not optimized, then our method behaves \"as expected\" respecting the MAP even in the region without the data (see attached PDF for a plot in that setting). The reason is that with the prior optimized the loss landscape of our linearzied approach has a particular \"biased\" behavior, which batching seems to help alleviating.\n\n**Point 2 / Weaknesses**\n> *...This small data regime would allow computing of (close to) exact posteriors, which would allow better quantitative assessment...*\n\nThank you for the suggestion. In the attached PDF we added a comparison with the predictive distribution obtained by using HMC on the regression examples. \n\n\n**Point 3 / Weaknesses**\n> *The method only considers very small models (e.g. single or 2 layer NNs) and small toy problem datasets.* \n\nFor the experiments we use an off-the-shelf solver for solving ODEs in ~5000 dimensions. We also use it for a LeNet with ~44000 parameters (see attached PDF). We would like to remark that from a differential geometry viewpoint this is already a surprising result. We expect to be able to push further the dimensions by developing suitable ODE solvers.\n\n\n**Point 4 / Questions**\n> *What are the method's memory and computational budget limitations?*\n\nAs we explained in the general comment the overhead our method has on top of Laplace approximation is given by the need to solve an ODE to get a sample. If we define $N$ to be the number of datapoints, $W$ to be the number of parameters, then both computing the gradient and the hvp is $O(NW)$. A single step of Runge-Kutta method of order $5(4)$ is $O(W)$ [1]. At each step we have to evaluate the ODE which requires the computation of gradient and hvp. Therefore, if we assume that the solver performs $S$ steps, the computational cost of the solver is $O(SNW)$.\n\n1. Hairer, E., Nørsett, S. P., and Wanner, G. Solving Ordinary Differential Equations I, Nonstiff Problems. Springer, 1993.\n\n\n**Point 5 / Questions**\n> *How accurate is the ODE solver when the number of dimensions grows?*\n\nThe accuracy of the solution does not depend on the dimensions of the problem per se, but on the complexity of the ODE problem and the solver. We solve ODEs using a general purpose Numpy ODE solver that is based on a high-accuracy algorithm (Runge-Kutta), where accuracy here means that the solution satisfies perfectly the ODE system for all time steps. Therefore, our solutions are highly accurate, with the price of increased computational complexity. We refer to the general comment for potential improvements in efficiency. We also conjecture that “this type of accuracy” may not be critical for our method. Instead, it may be sufficient and \"accurate\" for our method the generated geodesic to travel within low loss region and stop when the loss increases. A specialized solver with this characteristics is of particular interest.\n\n**Point 6 / Questions**\n> *What are the most important reasons to consider the proposed adaptation using the Riemannian metric over MacKay's choice of basis for a Laplace approximation?*\n\nThank you for the reference, this is indeed related to our work and we will include it the associated section. MacKay proposed to reparametrize the model/parameter space such that to be as near as possible to a Gaussian. This way, LA will be a good approximation. However, this is not always straightforward, especially in the deep networks regime. Our method is “similar” in spirit, but instead of reparametrizing the parameter space we make the approximation adapt to it by finding several local bases instead of a general one.\n"}, {"review_id": "Qpwk6PXZyS", "reviewer": "Reviewer_am4X", "summary": "This paper presents a novel Laplace Approximation for Bayesian Neural Networks. A key insight is to examine the local loss landscape with a Riemannian metric, which is determined by the gradient of the log posterior. Using this metric and an exponential map, a Laplace Approximation technique is developed to draw posterior samples that fall into regimes with low negative log posterior. The paper also develops a sampling method, which relies on the 2nd order ODE solver. Several experiments are conducted. When compared to the standard Laplace Approximation, evidences are provided to illustrate the improvements. ", "questions": "N/A", "limitations": "There is a limitation section at the end of the paper.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- the contribution provided by this work is original and novel to the best of my knowledge.\n\n- the paper is polished well generally. Despite that the materials are developed on differential geometry, intuitions are relatively provided well.\n\n- Laplace Approximation has been increasingly popular in recent years. Such extensions to incorporate Riemannian geometry could be relevant to the Bayesian Deep Learning community.", "weaknesses": "One complaint about the paper is that, without referring to the appendix, it is difficult to comprehend the material fully. For example, in section 3.3, I wish that the connection between an ODE and Riemannian metric is difficult to understand directly. Differential geometry is not often thought in engineering courses at many universities. It may make sense to recap the essential concept in the main paper. How the method could be used for linearized Laplace Approximation is made very short.\n\nAnother point for improvement is the choice of the baselines. It would make sense to include a deep ensemble and MC-dropout as a minimum. This could show how far the proposed Laplace Approximation can compete with popular methods in practice.\n\nIn the experiments, the paper could improve on analyzing the computational complexity, in comparison to the standard Laplace Approximation. DaxBerger et al 2021 claim that the major benefit of Laplace Approximation is on simplicity, it would be great to project this paper's method more on the plateau of quality of uncertainty Vs computational complexity. While it is great for research, I think the methods based on differential geometry may have certain drawbacks among practitioners. The paper could be more convincing by providing the gains in uncertainty, but additional overhead due to the added complexity of the pipeline.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Response to Reviewer am4X\n\nWe would like to thank the reviewer for their thoughtful consideration of our work. We appreciate the time you took to review our paper. We have taken the time to address all the points raised under Weaknesses.\n\n**Point 1/ Weaknesses**\n> *\"...without referring to the appendix, it is difficult to comprehend the material fully.\"*\n> *\"...How the method could be used for linearized Laplace Approximation is made very short.\"*\n\nThank you for the feedback. The paper is based on two unrelated fields Bayesian NNs and differential geometry. Due to space limitations thorough discussion and further information have been moved in the appendix, as well as a concrete example for the linearized version of our method. As our standard approach usually works better than the linearized version and is interpretable, we mainly focus on that in the paper. However, we will take into account the feedback and update the camera-ready version accordingly.\n\n**Point 2/ Weaknesses**\n> *\"Another point for improvement is the choice of the baselines. …  compete with popular methods in practice\"*\n\nIn related works where extensions of LA are proposed, linearized LA is considered as the main baseline for comparison, as this is the LA approach that competes well with Deep-Ensemble for Bayesian NNs. Therefore, it is already a quite strong baseline. We are happy to include more comparisons (ensemble methods and dropout) to the camera-ready version.\n\n**Point 3/ Weaknesses**\n\n> *\"In the experiments, the paper could improve on analyzing the computational complexity, in comparison to the standard Laplace Approximation. ... due to the added complexity of the pipeline.\"*\n\nLA is already simple and cheap. Our method is more flexible and expressive, but it comes with the price of increased computational cost as all differential geometric techniques. However, implementing our method is simple as in practice only an ODE (initial value problem) has to be solved when generating a sample (for scalability and complexity see the general comment). Another benefit of our method is interpretability. Even if linearized LA works well, it is not yet understood theoretically why it performs as such. Instead, our method is interpretable, which can be easily seen from the regression example where the LA is known to behave poorly. \n\n\nWe mentioned in the general comment that the computational cost on top of Laplace of our proposed method is $O(SNW)$ where $S$ is the number of step of the ODE solver, $N$ is the number of datapoints and $W$ is the number of parameters in the model. Getting $K$ samples from the posterior is therefore $O(KSNW)$. To briefly explain the complexity cost, at each step of the solver, we need to compute the gradient and the hessian-vector product, which are $O(NW)$, and perform a Runge-Kutta step, which scales linearly with the dimension of the problem, i.e. $O(W)$ [1]. Since the solver would take $S$ steps (RK-45 uses an adaptive step-size therefore this values changes in every ODE), we get $O(SNW)$.\n\n1. Hairer, E., Nørsett, S. P., and Wanner, G. Solving Ordinary Differential Equations I, Nonstiff Problems. Springer, 1993."}, {"review_id": "kXaRlrY88W", "reviewer": "Reviewer_BssD", "summary": "The paper presents a Laplace approximation for Bayesian neural networks that adapts the covariance to the local geometry of the loss, effectively overcoming the quadratic approximation of the loss. The authors report competitive performance with the standard Laplace approximation (both Monte Carlo sampled and linearized) on regression and MNIST-scale classification problems as well as a reduced reliance on tuning the precision of the prior.\n\nThe approach is explained clearly and makes a lot of sense (at least with my rather superficial understanding of Riemannian geometry), it is applicable in more general probabilistic models and I would expect it to lead to various follow-up works. While methodologically this is a very nice paper, I feel like it is let down by the empirical evaluation. The method is only tested on UCI and MNIST-scale datasets, which are hardly relevant for deep learning these days. The authors mention the computational cost of their method, but only discuss the reasons superficially without providing exact benchmark figures to give a sense of where the main bottlenecks arise. Given the apparent computational limitations of the method, an experiment with a non-NN model could have strengthened the paper.\n\nAll things considered, the strong methodological contribution outweighs the unconvincing empirical evaluation for me, so I would lean towards acceptance, although I wish could have given the paper a much higher rating.", "questions": "* What are the absolute runtimes of each method per sampled prediction in the expeirments? It would be great to get a better sense of this as a function of dataset/minibatch size and number of parameters/network size (I could see a synthetic experiment be illuminating here).\n* How much of the cost of the method lies in calculating the Hessian-vector products? These might take up a significant chunk of the total ODE compute time, so I’m wondering if explicit Hessian approximations (last-layer, KFAC, subset, ...) might improve scalability?\n* Could you elaborate on the discussion at the end of section 3.3? In particular I don’t follow how using the full dataset would over-regularize the geodesic as judging from the $N/B$ factor in the inline equation you seem to correctly rescale the mini-batch loss to match the full loss in expectation (note: I assume by ‘over-regularize’ you mean concentrating the samples around the mean, i.e. effectively reduce the entropy of the Gaussian approximate posterior). Any ideas for overcoming this over-regularization?\n* Could you comment on the results for the Riemannian Laplace approximation being noticeably better on the UCI datasets with a non-optimized precision?\n\n\n**Minor note**: the x axis labels in Fig 4 are cut off at the bottom", "limitations": "Only discussed superficially and in a quite hand-wavy manner, I would have wanted to see concrete benchmarks and an analysis of how the compute time evolves with increasing dataset size and number of parameters respectively as well as some evidence from the literature that a tailor-made ODE solver could indeed allow the method to take the step from small conv nets to more modern architectures.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "* The core idea makes a lot of sense, is applicable beyond inference in neural networks and seems to work well for the experiments that are considered. It has the potential to address the rather restrictive approximation of a quadratic loss in the Laplace approximation.\n* I am confident that the paper will inspire various pieces of follow-up work.\n* The paper is well-structured and -written.\n* Effective use of illustrative examples throughout.", "weaknesses": "* Only small scale problems are considered in the experiments\n* This is exacerbated by lack of analysis of the computational cost. It is not really clear to me what specifically is preventing the method from being applied to larger networks and datasets (even something like CIFAR with ResNets would have been great). The discussion mentions scaling issues w.r.t. number of datapoints and parameters, but the relative behavior is not benchmarked at all and it is also not clear to me how much time the ODE solver actually spends e.g. calculating the Hessian-vector products vs computations independent of that. Given that this paper opens up a new direction for research, I think there should be much clearer pointers as to where exactly the current bottlenecks and limitations are and where improvements can realistically be achieved.\n* There is no experiment demonstrating the efficacy of the approach on a non-NN probabilistic model. Given the apparent computational cost of the method, this would have seemed like a rather natural experiment to include and neural networks don’t seem like a particularly good fit for the approach.\n* The explanation of why the method would work better with a mini-batch rather than the full dataset (section 3.3/fig 4) is not exactly clear and seems rather hand-wavy.\n* Lack of HMC ground-truth baselines for the regression problems", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Response to Reviewer BssD\n\nWe thank the reviewer for the positive consideration of our work and constructive feedback. We appreciate the time spent to review our paper. We address all your concerns below and we refer to the general comment for the questions regarding scalability.\n\n\n**Point 1/ Weaknesses**\n> *“There is no experiment demonstrating the efficacy of the approach on a non-NN probabilistic model. … good fit for the approach.”*\n\nIndeed, our model can be also used to approximate posteriors other than BNNs. We have already in the paper two constructive examples: the Rosenbrock function (Fig. 1) and the logistic regression (Fig. 3). We are happy to include more examples in the camera-ready appendix. The reasons why we decided to focus on the BNN problem are due to the challenges, the potential impact and the research questions that arise. Even with the current basic setting, our approach performs on-par/better than linearized LA, which is considered among the strongest approximations for turning NNs into BNNs. We also note that our method is interpretable compared to linearized LA, which performance is not yet theoretically understood.\n\n\n**Point 2/ Weaknesses**\n> *\"The explanation of why the method would work better with a mini-batch rather than the full dataset seems rather hand-wavy\"*\n\nWe agree that we have not yet fully analyzed the influence of batching on the result. The terminology can lead to misunderstanding, but for “batching” in this setting we refer to solving each ODE only using a subset of the data. In the context of this paper we propose this as an “obvious and simple” way to scale the method. However, further research should be conducted to properly analyze the behavior. We believe that this is closely related to the generalization concept in deep learning, when stochasticity is induced in the training algorithm.\n\n**Point 3/ Weaknesses**\n> *\"Lack of HMC ground-truth baselines for the regression problems\"*\n\nThank you for the suggestion. We added the predictive distribution obtained using HMC on the attached pdf. We will add them in the appendix too.\n\n\n**Point 4 / Questions**\n> *\"What are the absolute runtimes of each method per sampled prediction in the expeirments? ...(I could see a synthetic experiment be illuminating here)\"*\n\nThank you for the suggestion. We provide some initial results in the attached PDF, and we will include further analysis in the camera-ready version. We also briefly mention the challenges that influence this benchmark on the general comment.\n\n\n**Point 5 / Questions**\n>*“How much of the cost of the method lies in calculating the Hessian-vector products? … might improve scalability?*\n\nThis is a great suggestion for future research. Indeed approximations to the metric and/or the ODE are of particular interest to reduce complexity, as we also mentioned in the general comment.\n\n**Point 6 / Questions**\n> *\"Could you elaborate on the discussion at the end of section 3.3? ... Any ideas for overcoming this over-regularization?\"*\n\nIndeed, the over-regularization means that due to the linearization and the prior precision optimization, especially if precision gets a high value, the low loss region concentrates closely around the MAP. Therefore, our samples are generated only near the MAP, and this bias limits the variability of the sampled functions. This is for example not happening in our standard approach. We believe that the batching is an interesting way to alleviate this issue as it seems empirically to be beneficial. However, it poses challenging questions and perhaps future insights. For example, there might be a correlation between the quality of the MAP with respect to generalization, and the sampled models associated with the loss surfaces implied by each batch. \n\n\n\n**Point 7 / Questions**\n> *“Could you comment on the results for the Riemannian Laplace approximation being noticeably better on the UCI datasets with a non-optimized precision?”*\n\nIf the optimized prior precision is high, this corresponds to a stronger L2 regularization, which implies that more models will be generated near the MAP, and this is true even for our model. When the prior is not optimized, our model is capable of generating models with higher variability, which helps to calibrate the uncertainty better. Therefore, the test NLL is expected to be better in our case, as points near the boundary that are missclassified do not get high confidence. Instead, this is the case when the prior is optimized, since the sampled functions are closer to the MAP. \n"}, {"review_id": "eFyVjKbpao", "reviewer": "Reviewer_dyBq", "summary": "This paper develops a Riemannian Laplace approximation, which is a Laplace approximation that takes into account the Riemannian geometry of the loss surface. The contributions of this paper are as follows: i) showing that such a loss-aware Laplace approximation is better able to capture the true posterior (and predictive); ii) presenting the Riemannian geometry framework for the Laplace approximation, with the Hessian in both the normal and tangential space; iii) a practical algorithm for efficiently integrating the required ODE; iv) experimental evidence on several commonly used datasets.", "questions": "In Figure 4, why does the NLL have a minimum and not strictly decrease with larger mini-batch size? Shouldn't the estimate become more and more precise? Is the wall-clock time really on the order of tens of seconds? What is the model size? \n\nFor the CNN, what is the kernel size? \n\nWhich approximation to the Hessian do you use for the Riemannian Laplace approximation (e.g. Diagonal or Kronecker-Factorized)? ", "limitations": "The method has been applied to very small neural networks and the authors argue that this is because the method is computationally very expensive. While the authors acknowledge this in the limitations, it potentially limits the applicability to realistic neural network model sizes. It would be important to know more precisely what the computational complexity is (also as a function of the model size). Furthermore, I would like to see a comparison in wall clock time for the entire approximation compared to standard Laplace approximation, not only for different mini-batch sizes, but also for different model sizes, including large model sizes.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is for the most part very well written, it is for the most part easy to follow (main idea and background). And the authors give several examples (including figures) to help the readers further.\nThe experimental section is decent, having both a toy example that helps gain intuition and a quantitative evaluation on several standard datasets, where the method is benchmarked against vanilla Laplace approximation and the MAP estimate. I think for the purpose of this paper, the various alternative approximate inference algorithms would not be necessary to compare to, as this is a direct extension of Laplace.", "weaknesses": "Even though the paper is generally very well written, I did, however, struggle with some parts of the main section. As a reviewer who is very familiar with Bayesian neural networks and the Laplace approximation but less so with Riemannian geometry, I would prefer to have a more in-depth background section on Riemannian geometry. I think the Laplace approximation section can be shortened if space is needed (e.g. tricks of the trade and strengths and weaknesses could be shortened).\n\nThe main weakness is the scalability of the method (see limitations). Solving the ODE takes a very long time. According to Fig. (4), this is in the order of tens of seconds for a mini-batch. Even for very small neural networks. The models used in the experimental section are tiny, e.g. single hidden layer networks with only 50 hidden units. For the CNN, the authors mention that they use 2 conv layers, but I didn't find the kernel size.\n\nIt is also not clear to me if and which approximation the authors use for the Hessian. Can the complexity be reduced e.g. with a diagonal Hessian approximation? How would this compare to a diagonal or Kronecker-Factorised approximation of the standard LA?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Response to Reviewer dyBq\n\nWe would like to thank the reviewer for their positive consideration. We appreciate the time and effort spent in reviewing our paper. We addressed all the remaining concerns below.\n\n**Point 1 / Weaknesses**\n> *“Even though the paper is generally very well written, I did, however, struggle with some parts of the main section. … I think the Laplace approximation section can be shortened if space is needed (e.g. tricks of the trade and strengths and weaknesses could be shortened).”*\n\nThank you for the feedback. The paper is based on two unrelated fields Bayesian NN and differential geometry. Due to space limitations we included in the main paper the most important information needed for the understanding of the idea, while thorough discussion and further information have been moved in the appendix. We will take into account the feedback and update the camera-ready version accordingly. \n\n\n**Point 2 / Weaknesses and Questions**\n> *\"It is also not clear to me if and which approximation the authors use for the Hessian. … How would this compare to a diagonal or Kronecker-Factorised approximation of the standard LA?\"*\n\nIn all the experiments, we consider the full Hessian for computing the covariance of the Gaussian approximation both for classic and linearized LA and our method. This is computed using the Laplace library which computes the GGN approximation. Therefore, for our method we are sampling the velocities from a Gaussian approximation obtained by computing the full Hessian. In the ODE solver, since we rely on “hpv” of functorch, which also uses the full Hessian in a computationally efficient way, we avoid materializing it. \n\nWhile for complexity we can refer to the general comment, the influence of the choice of the hessian approximation should be futher investigated, but empirically we have seen that diagonal covariances imply faster geodesics (see benchmark in the attached PDF). We will add also a comparison in the camera-ready appendix, where we will analyse how sampling the initial velocities from less accurate Hessian approximations will affect all models.\n\n\n**Point 3 / Questions**\n> *In Figure 4, why does the NLL have a minimum and not strictly decrease with larger mini-batch size? Shouldn't the estimate become more and more precise? Is the wall-clock time really on the order of tens of seconds? What is the model size?*\n\nThank you for pointing out this issue, we agree that it can be confusing. In this plot we report the test NLL i.e., the negative log-likelihood on test unseen data. Indeed, the train NLL should follow the behavior you described as the more points in the batch, the closer the sampled functions to the MAP estimate. However, the MAP may not be the optimal model for the test distribution. Instead, using a (small) batch when solving an ODE system allows our method to generate functions similar to the MAP, but which exhibit some variability and potentially generalize better. Indeed, if the batch is still representative of the whole dataset, the sample will differ from the MAP mostly close to the decision boundary and away from data.\nRegarding the implementation, the result corresponds to the Banana experiment in Sec 5.2. We will update the plot and the description accordingly.\n\n\n**Point 4  / Questions**\n> *“For the CNN, what is the kernel size?”*\n\nWe used a $5\\times5$ kernel, as reported in Appendix D.5\n"}]

Author rebuttal:
## General Comment to all reviewers
We would like to thank the reviewers for their thoughtful comments, positive considerations and suggestions for improving the paper. We appreciate that you found our work novel, well-written, with potential impact for the community, and inspiring for follow-up works. We make a first general comment about the scalability of our approach and we also reply individually to each reviewer. We are happy to clarify further during the discussion phase if some concerns remain.

**Scalability**

The scalability of our proposed approach and the applicability to big models was a common topic of discussion among the reviewers, which they acknowledge that we also highlight in the paper as the main limitation of our approach. We will elaborate more about this issue in the camera-ready version based on the following discussion.

**Evaluation of the ODE**:
- Based on the structure of the Riemannian metric we simplified the original ODE, and its final form allows to apply automatic-differentiation to evaluate it. Otherwise, evaluating the original ODE is rather prohibitive as it needs to compute the Hessian and the Gradient of the loss individually.
- Evaluating the ODE needs all the training data points which is prohibited with big data. We proposed the “obvious” trick of using a random (small) batch when solving an ODE, which empirically in some cases even boosts performance motivating further research. Another idea is to evaluate in parallel the ODE using batching, and then, collect all sub-results for the final ODE result.
- Other potential ideas is to approximate the Riemannian metric with surrogates leading to simpler ODE systems or providing approximations for the current ODE e.g., considering a diagonal Hessian approximation in the ODE as reviewers BssD and dyBq suggested.

**ODE solver**:
- For solving the ODE system we use a general purpose Python solver (``scipy.integrate.solve_ivp``) that runs on the CPU. When the solver needs to evaluate the ODE, our automatic-differentiation based approach runs on the GPU, and the result is moved to the CPU (``.detach().cpu().numpy()``) causing a significant overhead. Especially when dimensions increase both the transfer of the data and the computations on the CPU are sub-optimal. Implementing an ODE solver on a suitable automatic-differentiation framework (e.g. JAX) solely running on GPU will dramatically improve performance.
- As we know the structure and behavior of our ODE system (geodesics start from low loss which increases along the curve), a potential future work would be to develop solvers that exploit this information. Usually, general purpose solvers aim for accuracy, while in our case even inexact solutions could be potentially useful if computed fast [1].
- The benchmark result (see attached PDF) shows that, in general, increasing data and dimensionality makes the ODE system potentially more expensive to solve. However, there are cases where the solution of the ODE is fast in big models. Related research shows that the loss landscape of overparametrized models exhibits behaviors that may be beneficial to our method e.g., easily connected minima with (nonlinear) continuous paths. Therefore, we believe that even in high dimensions, our method has the potential to scale. It may also inspire new ways to study generalization via the loss landscape.

The cost of evaluating a single ODE is $O(SWN)$, where $S$ is the number of steps of the solver, $W$ is the number of model parameters and $N$ is the dataset size. This is the complexity overhead on top of usual Laplace approximation. However, the number of steps the solver needs to converge mainly depends on the complexity of the associated ODE problem, which is defined by the geometry of the loss landscape and the initial velocity. In the attached PDF we include a benchmark synthetic example that shows how the runtime changes with respect to the size of the model and dataset. While increasing the model parameters and dataset size affect $W$ and $N$, it may be the case that ODE systems do not necessarily get harder, therefore the value of $S$ would be small. In other words, evaluating the ODE becomes more expensive as dimensions increase, but perhaps the actual system gets easier to solve.
With the current implementation we manage to solve ODEs up to ~44000 dimensions (LeNet), which is surprising from the differential geometry perspective, and show that the method performs well in this regime (see attached PDF), but more sophisticated implementations will speed this up. Moreover, in a practical application samples are generated offline and not during test time. 

Overall, in this paper we propose a Riemannian extension to Laplace approximation and empirically verify the claims, by considering the complete formulation and relying on pre-existing tools for the implementation. As reviewers acknowledge there is a spectrum of potential research ideas in between for either improving the efficiency of our approach or for new techniques based on the same differential geometric principles. This is in spirit related to NeuralODEs where the original paper proposed the main concept and follow-up works improve parts of it, for example, specialized ODE integrators [2,3,4,5,6].

**References**
1. "Fast and robust shortest paths on manifolds learned from data". G. Arvanitidis et al., AISTATS 2019
2. "Opening the blackbox: accelerating Neural Differential Equations by regularizing internal solver heuristics". A. Pal et al., ICML 2021
3. "On numerical integration in Neural Ordinary Differential Equations". A. Zhu et al., ICML 2022
3. "On robustness of Neural Ordinary Differential Equations". H. Yan et al., arXiv 2020
5. "MALI: A memory efficient and reverse accurate integrator for Neural ODEs". J. Zhuang et al., arXiv 2021
6. "STEER: Simple temporal regularization for Neural ODEs". A. Ghosh et al., arXiv 2020

Meta-review: All reviewers were positive about the methodological contribution as the combination of Laplace methods and Riemannian geometry applying to Bayesian neural networks is interesting and should be of interest to the community. One significant drawback is the high complexity of the proposed method (the run-time required per example is very high) and this explains the small scale of the networks/datasets considered. Would be great if the authors could (partly) address these issues, improve the exposition (see Reviewer dyBq's comment), and add more detailed descriptions to all figures.
Decision: Accept (poster)