Title: Riemannian Laplace approximations for Bayesian neural networks
Venue: NeurIPS 2023
Accepted: yes
Keywords: Riemannian geometry

Laplace approximation

Approximate inference

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

Abstract

Bayesian neural networks often approximate the weight-posterior with a Gaussian distribution. However, practical posteriors are often, even locally, highly nonGaussian, 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 weightregions 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.

1 Introduction

Bayesian deep learning estimates the weight-posterior of a neural network given data, i.e. p ( θ |D ) . Due to the generally high dimensions of the weight-space, the normalization of this posterior is intractable and approximate inference becomes a necessity. The most common parametric choice approximates the posterior with a Gaussian distribution, p ( θ |D ) ≈ q ( θ |D ) = N ( θ | µ, Σ) , which is estimated variationally [Blundell et al., 2015], using Laplace approximations [MacKay, 1992] or with other techniques [Maddox et al., 2019]. Empirical evidence, however, suggests that the log-posterior is not locally concave [Sagun et al., 2016], indicating that the Gaussian approximation is overly crude. Indeed, this approximation is known to be brittle as the associated covariance is typically ill-

Riem. LA (ours)

LA

conditioned implying a suboptimal behavior [Daxberger et al., 2021a, Farquhar et 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 to the many benefits of parametric distributions, over e.g. Monte Carlo sampling [Neal, 1995] or deep ensembles [Lakshminarayanan et al., 2017].

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 a Gaussian approximation can locally adapt to the loss by equipping the weight-space with a 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

Riemannian approximate posterior using a 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.i.d.) data D = { x n , y n } N n =1 , consisting of inputs x ∈ R D and outputs y ∈ R C . To enable probabilistic modeling , we use a likelihood p ( y | f θ ( x )) which is either Gaussian (regression) or categorical (classification). This likelihood is parametrized by a deep neural network f θ : R D → R C , where θ ∈ R K represent the weights for which we specify a Gaussian prior p ( θ ) . The predictive distribution of a new test point x ′ equals p ( y | x ′ ) = ∫ p ( y | x ′ , θ ) p ( θ |D )d θ where p ( θ |D ) is the true weight-posterior given the data D . To ensure tractability, this posterior is approximated. This paper focuses on 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-probability density, thereby yielding a Gaussian approximation. When considering inference of the true posterior p ( θ |D ) , LA constructs an approximate posterior distribution q LA ( θ |D ) = N ( θ | θ ∗ , Σ) that is centered at the maximum a-posteriori (MAP) estimate

A Taylor expansion around θ ∗ of the regularized loss L ( θ ) then yields

where we know that ∇ θ L ( θ ) ∣ ∣ θ = θ ∗ ≈ 0 , and H θ [ L ]( θ ) ∈ R K × K denotes the Hessian of the loss. This expansion suggests that the approximate posterior covariance should be the inverse Hessian Σ = H θ [ L ]( θ ) ∣ ∣ -1 θ = θ ∗ . The marginal likelihood of the data is then approximated as p ( D ) ≈ exp( -L ( θ ∗ ))(2 π ) D / 2 det(Σ) 1 / 2 . This is commonly used for training hyper-parameters of both the likelihood and the prior [Immer et al., 2021a, Antorán et al., 2022]. We refer to appendix A 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, or 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 Gauss-Newton 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][Ritter 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 θ ∗ , and use the sample to use the 'linearized' function f lin θ ( x ) = f θ ∗ ( x ) + ⟨∇ θ f θ ( x ) ∣ ∣ θ = θ ∗ , θ -θ ∗ ⟩ as predictive, where ∇ θ f θ ( x ) ∣ ∣ θ = θ ∗ ∈ R C × K 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 the parameters θ , 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 et al., 2018, Kristiadi et al., 2020, Immer et al., 2021a, 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 a parametric approximate posterior that better reflects the local shape of the true posterior and captures nonlinear correlations between parameters. The basic idea is 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 a suitable Riemannian metric (Sec. 3.1) and then construct a Laplace approximation according to this metric (Sec. 3.2).

3.1 A loss-aware Riemannian geometry

For a given parameter value θ ∈ Θ , we can measure the training loss L ( θ ) of the associated neural network. Assuming that the loss changes smoothly with θ , we can interpret the loss surface M = g ( θ ) = [ θ, L ( θ )] ∈ R K +1 as a K -dimensional manifold in R K +1 . The goal of Riemannian geometry [Lee, 2019, do Carmo, 1992] is to do calculations that are restricted to such manifolds.

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 as J g ( θ ) v ∈ R K +1 , where J g : Θ →

θ

v

Exp

θ

(

v

)

R K +1 × K is the Jacobian of g that spans the tangent space T g ( θ ) M at the point g ( θ ) ∈ M and v ∈ R K is 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 ⟨ J g ( θ ) v 1 , J g ( θ ) v 2 ⟩ = v ⊺ 1 J g ( θ ) ⊺ J g ( θ ) v 2 , which, we note, is now expressed directly in the intrinsic coordinates. From this observation, we define the Riemannian metric M ( θ ) = J g ( θ ) ⊺ J g ( θ ) , which gives us a notion of a local inner product in the intrinsic coordinates of the manifold (see ellipsoids in Fig. 2). The Jacobian of g is particularly simple J g ( θ ) = [ I K , ∇ θ L ] ⊺ , such that the metric takes the form

The exponential map. A local inner product allows us to define the length of a curve c : [0 , 1] → Θ as length [ c ] = ∫ 1 0 √ ⟨ ˙ c ( t ) , M ( c ( t )) ˙ c ( t ) ⟩ d t , where ˙ c ( t ) = ∂ t c ( t ) 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 a given starting point and direction (see Fig. 2). Geometrically, we can also think of this as mapping a tangent vector back to the manifold , and we write the map as Exp : M×T θ M→M .

The tangential coordinates v can be seen as a coordinate system for the neighborhood around θ , and since the exponential map is locally a bijection we can represent any point locally with a unique tangent vector. However, these coordinates correspond to the tangent space that is spanned by J g ( θ ) , which implies that by changing this basis the associated coordinates change as well. By

θ

∗

LA

OUR

(a) True posterior

(b) LA sample

(c) OUR sample

(d) LA model

(e) OUR model

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

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 to the metric, we first express the loss in normal coordinates of the tangent space at θ ∗ , h (¯ v ) = L ( Exp θ ∗ ( M ( θ ∗ ) -1 / 2 ¯ v )) . Following the standard Laplace approximation, we perform a second-order Taylor expansion of h as

Tangential Laplace. Similar to the standard Laplace approximation, we get a Gaussian approximate posterior ¯ q (¯ v ) = N (¯ v | 0 , Σ) on the tangent space in the normal coordinates with covariance Σ = H ¯ v [ h ](¯ v ) ∣ ∣ -1 ¯ v =0 . Note that changing the normal coordinates ¯ v to tangential coordinates v is a linear transformation and hence v ∼ N (0 , A ( θ ∗ ) ΣA ( θ ∗ ) ⊺ ) , which means that this covariance is equal to H θ [ L ]( θ ) ∣ ∣ -1 θ = θ ∗ since A ( θ ∗ ) is a symmetric matrix, and hence, it cancels out. The approximate posterior q T ( v ) = N ( v | 0 , Σ ) in tangential coordinates, thus, matches the covariance of the standard Laplace approximation.

where ∂ ¯ v h (¯ v ) ∣ ∣ ¯ v =0 = A ( θ ∗ ) ⊺ ∇ θ L ( θ ) ∣ ∣ θ = θ ∗ ≈ 0 as θ ∗ minimize the loss and H ¯ v [ h ](¯ v ) ∣ ∣ ¯ v =0 = A ( θ ∗ ) ⊺ H θ [ L ]( θ ) A ( θ ∗ ) ∣ ∣ θ = θ ∗ with H θ [ L ]( θ ) the standard Euclidean Hessian matrix of the loss. Further details about this step can be found in appendix B.

The predictive posterior. We can approximate the predictive posterior distribution using Monte Carlo integration as p ( y | x ′ , D ) = ∫ p ( y | x ′ , D , θ ) q ( θ )d θ = ∫ p ( y | x ′ , D , Exp θ ∗ ( v )) q T ( v )d v ≈ 1 S ∑ S s =1 p ( y | x ′ , D , Exp θ ∗ ( v s )) , v s ∼ q T ( v ) . Intuitively, this generates tangent vectors according to the standard Laplace approximation and maps them back to 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 a general construction that applies to the same Bayesian inference problems as the standard Laplace approximation and is not exclusive to Bayesian neural networks.

The above analysis also applies to the linearized Laplace approximation. In particular, when the f lin θ ( x ) is considered instead of the f θ ( x ) the loss function in (1) changes to L lin ( θ ) . Consequently, our Riemannian metric is computed under this new loss, and ∇ θ L lin ( θ ) appears in the metric (3).

Example. To build intuition, we consider a logistic regressor on a linearly separable dataset (Fig. 3). The likelihood of a point x ∈ R 2 to be in one class is p ( C = 1 | x ) = σ ( x ⊺ θ + b ) , where σ ( · ) is the sigmoid function, θ ∈ R 2 and b ∈ R . After learning the parameters, we fix b ∗ and show the posterior with respect to θ together with the corresponding standard Laplace approximation (Fig. 3a).

We see that the approximation assigns significant probability mass to regions where the true posterior is near-zero, and the result of a corresponding sample is 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 a sample at 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. 3e).

3.3 Efficient implementation

Our approach is a 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-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 to the structure of our Riemannian metric (3), the ODE simplifies to

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 L lin ( θ ) and Gaussian likelihood.

Automatic-differentiation. 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-differentiation frameworks to compute the Hessianvector product between H θ [ L ]( c ( t )) and ˙ c ( t ) directly. This both reduces memory use, increases speed, and simplifies the implementation.

Mini-batching. 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 a batch B of B random data points and use the associated loss in the ODE (6). As usual, we assume that L ( θ ) ≈ ( N / B ) L B ( θ ) . Note that we only mini-batch the metric and not the covariance of our approximate posterior q T ( 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 L ( θ ) controls the geodesics via the associated Riemannian metric, so when we consider the full dataset we expect the samples to behave similarly to f θ ∗ ( x ) . In other words, our approximate posterior generates weights near θ ∗ resulting in models with similar or even better loss. When we consider a batch the geometry of the associated loss surface L B ( θ ) 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 to 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 L lin ( θ ) may over-regularize the geodesics. Due to the linear nature of f lin θ ( θ ) the associated Riemannian metric is small only close to θ ∗ so the generated samples are similar to f θ ∗ ( x ) . We relax this behavior 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

50

100

150

200

250

Batch size

0

.

28

0

.

29

0

.

30

0

.

31

0

.

32

Test NLL

12

14

16

18

20

22

Expmap computation time (s)

Test NLL Performance

Computation Time

LIN-RIEM-LA

50

100

150

200

250

0

.

30

0

.

35

0

.

40

0

.

45

Test NLL

10

15

20

25

30

Expmap computation time (s)

Test NLL Performance

Computation Time

Batch size

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, 2011, 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 geometric 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 a first-order Taylor expansion, also known as linearization around the MAP estimate. For long, Laplace'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 et al., 2018]. More recently, the Laplace approximation has become a go-to tool for turning trained neural networks into BNNs in a post-hoc manner, thanks to easy-to-use software [Daxberger et al., 2021a] and new approaches to scale up computation [Antorán et al., 2022]. In this direction, other works have only considered a subset of the network parameters [Daxberger et al., 2021b, Sharma et 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 a 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 et 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 et 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 to make a local refinement of a mean-field variational approximation.

Differential geometry. Differential geometry is increasingly playing a 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 et al., 2018, Hauberg, 2018b]. A similar Riemannian metric has been used as a surrogate for the Fisher metric for Riemannian Hamiltomian Monte-Carlo sampling [Hartmann et al., 2022]. In addition to 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-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 last-layer 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-RIEM-LA

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 to reproduce the results is publicly available at https://github.com/federicobergamin/ riemannian-laplace-approximation

5.1 Regression problem

We consider the toy-regression problem proposed by Snelson and Ghahramani [2005]. The dataset contains 200 data points, and we randomly pick 150 examples as our training set and the remaining 50 as a test set. As shown by Lawrence [2001], Ritter et al. [2018], using samples from the LA posterior performs poorly in regression even if the Hessian is not particularly ill-conditioned, i.e. when the prior precision is optimized. For this reason, the linearization 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, so 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 consider the linearized version of our approach the result is the same as the standard LA-linearization, which we include in the appendix D, where we also report results for in-between uncertainty as proposed by Foong et al. [2019] and a comparison with Hamiltonian Monte Carlo.

5.2 Classification problems

0

1

2

3

4

5

6

-

4

-

3

-

2

-

1

0

1

2

3

4

Our samples

LA samples

LA mean

Our mean

MAP

0

1

2

3

4

5

6

-

4

-

3

-

2

-

1

0

1

2

3

4

Our samples

LA samples

LA mean

Our mean

MAP

Illustrative example. We consider a 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 a 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 RIEM-LA performs better in terms of NLL and Brier score on a separate test set.

As we discussed in Sec. 3.3, using a 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 D we show the 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 LA.

UCI datasets. We compare our approach against the standard LA on a set of six UCI classification datasets using a fully connected network with a 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 log-likelihood (NLL) in the test set. All other metrics are reported in appendix D. We are considering the setting where we optimize the prior-precision post-hoc, which is the optimal setting for LA and linearized LA. We consider our standard approaches without using batches, which we have seen that specifically for our linearized approach may lead to sub-optimal performance.

From the results in Table 2 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 is not optimal. This implies that the loss surface of the linearized loss potentially over-regularizes the geodesics as we analyzed in Sec. 3.3, and in this case, considering mini-batching could have been beneficial.

Image classification. We consider a small convolutional neural network on MNIST and FashionMNIST. Our network consists of two convolutional layers followed by average pooling layers and three fully connected layers. We consider a 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 D 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 NLL and Brier score, meaning that the models are better calibrated, but it also leads to a more accurate classifier than the MAP. In terms of ECE, it seems that considering the linearized

approach is beneficial in producing better-calibrated models in both datasets. This holds both for our approach linearized 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, as they achieve similar performance no matter if the prior precision is optimized or not.

Note that for the 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 batchsize is a 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-LA 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.

6 Conclusion and future directions

We propose an extension to the standard Laplace approximation, which leverages the natural geometry of the parameter space. Our method is parametric in the sense that a Gaussian distribution is estimated using the standard Laplace approximation, but it adapts to the true posterior through a nonparametric Riemannian metric. This is a general mechanism that, in principle, can also apply to, e.g., variational approximations. In a similar vein, while the focus of our work is on Bayesian neural networks, nothing prevents us from applying our method to other model classes.

Empirically, we find that our Riemannian Laplace approximation is better or on par with alternative Laplace approximations. The standard Laplace approximation crucially relies on both linearization and on a fine-tuned prior to give useful posterior predictions. Interestingly, we find that the Riemannian Laplace approximation requires neither. This could suggest that the standard Laplace approximation has a rather poor posterior fit, which our adaptive approach alleviates.

Limitations. The main downside of our approach is the computational cost involved in integrating the ODE, which is a common problem in computational geometry [Arvanitidis et al., 2019]. The cost of evaluating the ODE scales linearly with the number of observations, and in particular it 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. Indeed, the solver needs all the training data points to compute both gradient and hvp at each step. We refer to appendix C.1 for a 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 our automatic-differentiation based approach runs on the GPU, which is inefficient. We expect significant improvements by using an ODE solver that runs exclusively on GPU, while tailor-made numerical integration methods are also of particular interest.

Future directions. We showed that equipping the weight-space with a simple Riemannian metric is a 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 metrics 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, a 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 Foundation through the Center for Basic Machine Learning Research in Life Science (NNF20OC0062606). It 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.-A. Absil, R. Mahony, and Rodolphe Sepulchre. Optimization Algorithms on Matrix Manifolds . 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 Søren Hauberg. Latent space oddity: on the curvature of deep generative models. In International Conference on Learning Representations (ICLR) , 2018.

Georgios Arvanitidis, Søren 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-Newton 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-effortless Bayesian deep learning. In Neural Information Processing Systems (NeurIPS) , 2021a.

Erik Daxberger, Eric Nalisnick, James U Allingham, Javier Antorán, and José Miguel Hernández-Lobato. Bayesian deep learning via subnetwork inference. In International Conference on Machine Learning (ICML) , 2021b.

M.P. do Carmo. Riemannian Geometry . Mathematics (Boston, Mass.). Birkhäuser, 1992.

John R Dormand and Peter J Prince. A family of embedded runge-kutta formulae. Journal of computational and applied mathematics , 6(1):19-26, 1980.

Runa Eschenhagen, Erik Daxberger, Philipp Hennig, and Agustinus Kristiadi. Mixtures of Laplace approximations for improved post-hoc uncertainty in deep learning. In NeurIPS Workshop of Bayesian Deep Learning , 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 on Artificial Intelligence and Statistics (AISTATS) , 2020.

Andrew YK Foong, Yingzhen Li, José Miguel Hernández-Lobato, and Richard E Turner. 'In-Between' Uncertainty in Bayesian Neural Networks. arXiv preprint arXiv:1906.11537 , 2019.

F Dan Foresee and Martin T Hagan. Gauss-Newton approximation to Bayesian learning. In Proceedings of International Conference on Neural Networks (ICNN) . IEEE, 1997.

Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In International Conference on Machine Learning (ICML) , 2016.

Alex Graves. Practical variational inference for neural networks. In Neural Information Processing Systems (NeurIPS) , 2011.

Marcelo 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 variational posterior through iterative optimization. Entropy , 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 differential equations and their application to riemannian statistics. 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://github. com/pytorch/functorch , 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 on Machine Learning (ICML) , 2021a.

Alexander Immer, Maciej Korzepa, and Matthias Bauer. Improving predictions of Bayesian neural nets via local linearization. In International Conference on Artificial Intelligence and Statistics (AISTATS) , 2021b.

Pavel Izmailov, Dmitrii 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 on Machine Learning (ICML) , 2018.

Mohammad Emtiyaz E Khan, Alexander Immer, Ehsan Abedi, and Maciej Korzepa. Approximate inference turns deep networks into Gaussian Processes. In Neural Information Processing Systems (NeurIPS) , 2019.

Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference on 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 on Machine Learning (ICML) , 2020.

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 (NeurIPS) , 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 J 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 Kronecker-factored approximate curvature. In International Conference on 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 arXiv:1611.07476 , 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 pseudo-inputs. In Neural Information Processing Systems (NeurIPS) , 2005.

Alessandra Tosi, Søren 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, C J 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-272, 2020. doi: 10.1038/s41592-019-0686-2.

Riemannian Laplace approximations for Bayesian neural networks (Appendix)

Contents of the Appendix

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.e.:

This is done by fitting a Gaussian approximation to the unnormalized distribution p ( D | θ ) p ( θ ) at its peak, where p ( D | θ ) is the likelihood distribution and p ( θ ) is the prior over the weights. In standard training of neural networks the mean-squared error loss is usually used for regression. This corresponds to optimizing a 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 a Gaussian prior p ( θ ) distribution. More specifically, training with λ || θ || 2 2 corresponds to a Gaussian prior N ( θ ; 0 , λ -2 2 I )

Therefore a natural peak to choose is given by the θ ∗ , which is the set of weights obtained at the end of the training. Indeed, θ ∗ is usually computed by:

Once we have θ ∗ , the Laplace approximation uses a a second-order Taylor expansion of L ( θ ) around θ ∗ , which yields:

where the first-order term ∇ θ L ( θ ) ∣ ∣ θ = θ ∗ ≈ 0 because the gradient at θ ∗ is 0 . By looking at L ( θ ) , we can notice that the Hessian is composed by two terms: a data fitting term and a prior term. Assuming p ( θ ) = N ( θ ; 0 , γ 2 I ) , then the Hessian can be expressed as

where we defined α = 1 γ 2 , i.e. the prior precision.

Using the fact that ˆ L ( θ ) is the negative log-numerator of (A.1), we can get p ( D | θ ) p ( θ ) by taking the exponential of -ˆ L ( θ ) . By doing so we have:

For simplicity, we can define Σ = ( H + α I ) -1 and rewrite the equation above to obtain:

We can then use this approximation to estimate the normalizing constant of our approximate posterior, which corresponds to the marginal log-likelihood p ( D ) :

This can be rewritten as

and by using the Gaussian integral properties we can write it as:

and by taking the logarithm we get

which is the approximation of the log marginal likelihood that we want to maximize to optimize the parameters that appears in L ( θ ∗ ) . In a regression problem, we are interested in optimizing both the variance of the Gaussian likelihood and the prior precision. In classification, instead, we just have the prior precision as a hyperparameter to tune.

B Riemannian geometry

We rely on Riemannian geometry [Lee, 2019, do Carmo, 1992] in order to construct our approximate posterior. In a nutshell, a d -dimensional Riemannian manifold can be seen intuitively as a smooth d -dimensional surface that lies within a Euclidean space of dimension D > d , which allows one to compute distances between points that respect the geometry of the surface.

Definition B.1. A Riemannian manifold M is a smooth manifold together with a Riemannian metric M ( x ) that acts on the associated tangent space T x M at any point x ∈ M .

Definition B.2. A Riemannian metric M : M→ R dim ( M ) × dim ( M ) is a smoothly changing positive definite metric tensor that defines an inner product on the tangent space T x M at any point x ∈ M .

We focus on the Bayesian neural network framework where Θ = R K is the parameter space of the associated deep network, and we consider the manifold M = g ( θ ) = [ θ, L ( θ )] . This is essentially the loss surface which is K -dimensional and lies within a ( 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 a 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 J g ( θ ) ∈ R K +1 × K of the map g spans the tangent space on the manifold, and a tangent vector can be written as J g ( θ ) v , where v ∈ R K 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

Here the matrix M ( θ ) = J g ( θ ) ⊺ J g ( θ ) = I K + ∇ θ L ( θ ) ∇ θ L ( θ ) ⊺ 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 ( θ ) is transformed into a Riemannian manifold. This can be also seen as the abstract manifold definition where we 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 θ 1 , θ 2 ∈ Θ . Let a curve c : [0 , 1] → Θ with c (0) = θ 1 and c (1) = θ 2 , the its length defined under the Riemannian metric as length [ c ] = ∫ 1 0 √ ⟨ ˙ c ( t ) , M ( c ( t )) ˙ c ( t ) ⟩ dt . This quantity is computed intrinsically but it also corresponds to the length of the associated curve on M . The shortest path then is defined as c ∗ ( t ) = arg min c length [ 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)

where vec [ · ] stacks the columns of a matrix and ⊗ the Kronecker product [Arvanitidis et al., 2018]. A 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) = θ 1 and c (1) = θ 2 , we get the geodesic that connects these two points. Let v = ˙ c (0) be the velocity of this curve at t = 0 . When the system is solved as an Initial Value Problem (IVP) with conditions c (0) = θ 1 and ˙ c (0) = v , we get the geodesic c v ( t ) between c v (0) = θ 1 and c v (1) = θ 2 . 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 to 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 in (3) the general ODE s system in (B.2) becomes

Proof. We consider the general system, and we compute each term individually. To simplify notation we will use M := M ( c ( t )) , ∇ := ∇ θ L ( c ( t )) , H := H θ [ L ]( c ( t )) , H i the i -th column of the Hessian and ∇ i the i -th element of the gradient ∇ .

Using the Sherman-Morrison formula we have that

The first term in the brackets is

and the second term in the brackets is

and their difference is equal to

We compute the matrix-vector product between the matrix (B.7) and the Kronocker product ˙ c ( t ) ⊗ ˙ c ( t ) = [ ˙ c 1 ˙ c, . . . , ˙ c K ˙ c ] ⊺ ∈ R D 2 × 1 which gives

where we used that ∑ K i =1 H i ˙ c i = H ˙ c and ∑ K i =1 ∇ i ˙ c i = ∇ ⊺ ˙ c . As a final step, we plug-in this result and the inverse of the metric in the general system which gives

Taylor expansion. As regards the Taylor expansion we consider the space Θ and the Riemannian metric therein M ( θ ) and an arbitrary smooth function f : Θ → R . If we ignore the Riemannian metric the second-order approximation of the function f around a point x ∈ Θ is known to be where ∇ θ f ( θ ) ∣ ∣ θ = x is the vector with the partial derivatives evaluated at x and H θ [ f ]( θ ) ∣ ∣ θ = x the corresponding Hessian matrix with the partial derivatives ∂ 2 f ( θ ) ∂θ i ∂θ j ∣ ∣ ∣ θ = x . When we take into account the Riemannian metric, then the approximation becomes

where Γ k ij are the Christoffel symbols and the Einstein summation is used. Note that even if the Hessian is different, the approximation again is a quadratic function.

Now we consider the Taylor expansion on the associated tangent space at the point x instead of directly on the parameter space Θ . We define the function h ( v ) = f ( Exp x ( v )) on the tangent space centered at x and we get that

where we apply the chain-rule and we use the fact that ∂ v Exp x ( v ) ∣ ∣ v =0 = I and ∂ 2 v Exp x ( v ) ∣ ∣ v =0 = 0 . So, we get that ∂ v f ( Exp x ( v )) ∣ ∣ v =0 = ∇ θ f ( θ ) ∣ ∣ θ = x and H v [ f ]( Exp x ( v )) ∣ ∣ v =0 = ∂ 2 f ( θ ) ∂θ i ∂θ j ∣ ∣ ∣ θ = x . The difference here is that the quadratic function is defined on the tangent space, and the exponential map is a non-linear mapping. Therefore, the actual approximation on the parameter space ˆ f Tangent ( Exp x ( v )) = ˆ h ( v ) is not a quadratic function as before, but it adapts to 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 ˆ f Tangent ( Exp x ( v )) to f ( x ) . In our problem of interest, this behavior is desirable implying that if a parameter θ ′ is connected through low-loss regions with a continuous curve to θ ∗ , then we will assign to θ high approximate posterior density.

We can easily consider the approximation on the normal coordinates, where we know that the Christoffel symbols vanish, by using the relationship u = A ¯ u with A = M ( x ) -1 / 2 , and thus the Taylor approximation of the function ¯ h (¯ u ) = f ( Exp x ( A ¯ u )) becomes

Linearized manifold. Let us consider a regression problem with likelihood p ( y | x , θ ) = N ( y | f θ ( x ) , σ 2 ) , where f θ is a deep neural network, and prior p ( θ ) = N ( θ | 0 , λ I K ) . The loss of the f lin θ ( x ) is then defined as

The gradient and the Hessian of this loss function can be easily computed as

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 et al., 2021a] for fitting the Laplace approximation and for optimizing the hyperparameters by using the marginal log-likelihood. 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:

We use functorch [Horace He, 2021] to compute both the gradient and the Hessian-vector-product. We then rely on scipy [Virtanen et al., 2020] implementation of the explicit Runge-Kutta method of order 5(4) [Dormand and Prince, 1980] to solve the initial-value problem. We use default tolerances in all our experiments.

Linearized manifold. We define our linearized manifold by considering the 'linearized' function f lin θ ( x ) = f θ ∗ ( x ) + ⟨∇ θ f θ ( x ) ∣ ∣ θ = θ ∗ , θ -θ ∗ ⟩ to compute the loss, where ∇ θ f θ ( x ) ∣ ∣ θ = θ ∗ ∈ R C × K is the Jacobian. To compute ⟨∇ θ f θ ( x ) ∣ ∣ θ = θ ∗ , θ -θ ∗ ⟩ we use functorch to compute a jacobian-vector 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 it comes with a computational overhead. To be able to estimate this overhead, we have to assume that our solver does S steps to solve an ODE, which is 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 a single ODE results in O ( SWN ) , where S is the number of steps of the solver, W is the number of model parameters and N is the

500

1000

5000

10000

50000

100000

Dataset size

0

1000

2000

3000

4000

5000

Time per expamp (s)

Prior optimized, velocities from posterior

Model size

100 params (full)

500 params (full)

1k params (full)

5k params (diag)

10k params (diag)

dataset size. Indeed, at each step, we need to compute the gradient and the hessian-vector product, which is O ( NW ) , 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 we get O ( SNW ) . If we want to use K posterior samples to estimate our predictive distribution, then we have to solve K different ODEs, therefore the overall overhead is O ( KSNW ) . However, we 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 an 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 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.

Limitations of our implementation The cost of obtaining a posterior sample using our RIEM-LA scales linearly in terms of the model and dataset size. In addition to that, our implementation is not highly engineered. Although we can generate these samples in parallel, we use a general purpose Python solver ( scipy.integrate.solve_ivp [Virtanen et al., 2020]) to solve the ODE system, which 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, 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 a suitable automatic-differentiation framework (e.g. JAX) solely running on GPU will dramatically improve performance.

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 ( sklearn.datasets.make_blobs with three classes, see Fig. C.1 ( left )) to analyze the running time of computing a single sample from the posterior with our methods. A two layers neural-network is considered, and we vary the number of parameters and datapoints in the dataset. For the biggest models, i.e. the one with 5 k and 10 k parameters, we use a 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 a 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 a weight decay of 1 e -2 for the larger model and 1 e -3 for the smaller model for 35000 and 700000 epochs respectively. In both cases, we use a learning rate of 1 e -3 . In Sec. 5, we show some samples from the posterior distribution obtained by using vanilla LA and our RIEM-LA approach. In Fig. D.1 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.

Amore 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]. In 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 a small model the predictive distribution we get is comparable to linearized LA, when we 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 a subset of the training set 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 a 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 see 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 our 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 a learning rate of 1 e -3 and weight-decaay of 1 e -2 . Although it is a binary classification problem, we use a 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 non-linearized 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 linearized 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 are able to be confident only in regions where there is data. We generate a dataset considering 200

0

1

2

3

4

5

6

-

4

-

3

-

2

-

1

0

1

2

3

4

Ours

LA

Laplace mean

Our mean

MAP

POSTERIOR - SINGLE LAYER MODEL

0

1

2

3

4

5

6

-

4

-

3

-

2

-

1

0

1

2

3

4

MAP

LA mean

Our mean

Test points

2

√

V

[

y

]

Our

2

√

V

[

y

]

LA

PREDICTIVE - SINGLE LAYER MODEL

0

1

2

3

4

5

6

-

4

-

3

-

2

-

1

0

1

2

3

4

Ours

LA

Laplace mean

Our mean

MAP

0

1

2

3

4

5

6

-

4

-

3

-

2

-

1

0

1

2

3

4

MAP

LA mean

Our mean

Test points

2

√

V

[

y

]

Our

2

√

V

[

y

]

LA

POSTERIOR - TWO LAYERS MODEL

PREDICTIVE - TWO LAYERS MODEL

examples per classes and we use 350 example as training set and the remaining as test set. We consider a two layer fully-connected network with 20 hidden units per layer and tanh activation. As before, we train it using SGD with a learning rate of 1 e -3 and a weight decay of 1 e -2 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, we 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, we consider a neural network with a single fully connected layer consisting of 50 hidden units and tanh activation. We train it using Adam optimizer [Kingma and Ba, 2015] for 10000 epochs using a learning rate of 1 e -3 and a weight decay of 1 e -2 .

From the two tables, we can see that our RIEM-LA is better than all the other methods both in terms of NLL, 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.

0

1

2

3

4

5

6

-

4

-

3

-

2

-

1

0

1

2

3

4

Ours

LA

Laplace mean

Our mean

MAP

POSTERIOR - SINGLE LAYER MODEL

0

1

2

3

4

5

6

-

4

-

3

-

2

-

1

0

1

2

3

4

MAP

LA mean

Our mean

Test points

2

√

V

[

y

]

Our

2

√

V

[

y

]

LA

PREDICTIVE - SINGLE LAYER MODEL

0

1

2

3

4

5

6

-

4

-

3

-

2

-

1

0

1

2

3

4

Ours

LA

Laplace mean

Our mean

MAP

0

1

2

3

4

5

6

-

4

-

3

-

2

-

1

0

1

2

3

4

MAP

LA mean

Our mean

Test points

2

√

V

[

y

]

Our

2

√

V

[

y

]

LA

POSTERIOR - TWO LAYERS MODEL

PREDICTIVE - TWO LAYERS MODEL

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: an initial convolutional layer with 4 channels and 5 × 5 filter followed by tanh activation and an average pooling layer. The we have another convolutional layer still with 4 channels and 5 × 5 kernel also followed by tanh activation and an 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 1 e -3 and weight decay of 5 e -4 for 100 epoch. The learning rate is annealed using the cosine decay method [Loshchilov and Hutter, 2017].

In Table 7 and Table 8, we can see that also when we do not 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, we 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-distribution results

The benefit of having meaningful and robust uncertainty estimation is that our model would then be confident in-data region while being uncertain in region without data. Therefore, if this is happening, then we would expect the model to be able to detect out-of-distribution (OOD) examples more successfully.

We consider classic OOD images detection tasks, where we train a model on MNIST and tested on FashionMNIST, EMNIST, and KMNIST and one trained on FMNIST and tested on the remaining datasets. It's well known that linearized LA is one of the strongest method for OOD detection

0

1

2

3

4

5

6

-

4

-

3

-

2

-

1

0

1

2

3

4

Ours

LA

Laplace mean

Our mean

MAP

POSTERIOR - SINGLE LAYER MODEL

0

1

2

3

4

5

6

-

4

-

3

-

2

-

1

0

1

2

3

4

MAP

LA mean

Our mean

Test points

2

√

V

[

y

]

Our

2

√

V

[

y

]

LA

PREDICTIVE - SINGLE LAYER MODEL

0

1

2

3

4

5

6

-

4

-

3

-

2

-

1

0

1

2

3

4

Ours

LA

Laplace mean

Our mean

MAP

0

1

2

3

4

5

6

-

4

-

3

-

2

-

1

0

1

2

3

4

MAP

LA mean

Our mean

Test points

2

√

V

[

y

]

Our

2

√

V

[

y

]

LA

POSTERIOR - TWO LAYERS MODEL

PREDICTIVE - TWO LAYERS MODEL

[Daxberger et al., 2021a]. We consider the same MAP estimates we used in the previous section and present OOD performance in Table 9 and Table 10.

We can see that our proposed method is consistently working better than linearized and classic LA in all the considered setting apart for models trained on FMNIST where we 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.

0

1

2

3

4

5

6

-

4

-

3

-

2

-

1

0

1

2

3

4

Ours

LA

Laplace mean

Our mean

MAP

POSTERIOR - SINGLE LAYER MODEL

0

1

2

3

4

5

6

-

4

-

3

-

2

-

1

0

1

2

3

4

MAP

LA mean

Our mean

Test points

2

√

V

[

y

]

Our

2

√

V

[

y

]

LA

PREDICTIVE - SINGLE LAYER MODEL

0

1

2

3

4

5

6

-

4

-

3

-

2

-

1

0

1

2

3

4

Ours

LA

Laplace mean

Our mean

MAP

POSTERIOR - TWO LAYERS MODEL

0

1

2

3

4

5

6

-

4

-

3

-

2

-

1

0

1

2

3

4

MAP

LA mean

Our mean

Test points

2

√

V

[

y

]

Our

2

√

V

[

y

]

LA

PREDICTIVE - TWO LAYERS MODEL

0

1

2

3

4

5

6

-

4

-

3

-

2

-

1

0

1

2

3

4

Ours

LA

Laplace mean

Our mean

MAP

POSTERIOR - TWO LAYERS MODEL (BATCH)

0

1

2

3

4

5

6

-

4

-

3

-

2

-

1

0

1

2

3

4

MAP

LA mean

Our mean

Test points

2

√

V

[

y

]

Our

2

√

V

[

y

]

LA

PREDICTIVE - TWO LAYERS MODEL (BATCH)

0

1

2

3

4

5

6

-

4

-

3

-

2

-

1

0

1

2

3

4

Predictive distribution - HMC

Mean

Test points

2

√

V

[

y

]

PREDICTIVE - SINGLE LAYER MODEL

0

1

2

3

4

5

6

-

4

-

3

-

2

-

1

0

1

2

3

4

Predictive distribution - HMC

Mean

Test points

2

√

V

[

y

]

PREDICTIVE - SINGLE LAYER MODEL

0

1

2

3

4

5

6

-

4

-

3

-

2

-

1

0

1

2

3

4

Predictive distribution - HMC

Mean

Test points

2

√

V

[

y

]

PREDICTIVE - TWO LAYERS MODEL

0

1

2

3

4

5

6

-

4

-

3

-

2

-

1

0

1

2

3

4

Predictive distribution - HMC

Mean

Test points

2

√

V

[

y

]

PREDICTIVE - TWO LAYERS MODEL

PRIOR PRECISION OPTIMIZED

↓

Brier

005

009

.

.

↓

0

0

0

.

0

008

.

0

002

.

0

.

0

001

029

.

006

0

0

0

.

.

.

0

.

0

0930

0

±

2426

0956

0886

.

0

0890

.

0

.

0967

0887

±

±

±

±

±

±

.

0015

0

.

0046

0

0023

.

0

.

0

0006

.

0

0000

.

0

0024

.

0011

METHOD

MAP

VANILLA LA

LIN-LA

RIEM-LA

RIEM-LA (BATCHES)

LIN-RIEM-LA

LIN-RIEM-LA (BATCHES)

86

↑

0

.

69

±

59

.

86

.

87

50

99

57

.

87

.

87

30

.

87

.

02

77

±

±

±

±

±

±

.

5

34

.

0

.

0

07

37

.

0

.

0

07

08

.

0

38

.

24

0

NLL

.

333

0

.

±

678

0

325

.

0

.

0

287

.

286

0

415

.

0

.

298

±

±

±

±

±

±

5

.

12

ECE

91

.

6

.

2

.

±

08

59

↓

0

±

±

55

2

.

57

6

.

2

.

97

38

.

2

34

.

0

±

±

±

±

.

0

31

41

.

0

37

.

0

04

.

0

43

.

28

22

.

27

.

11

.

10

MCE

75

±

49

44

.

95

6

08

.

21

±

±

±

±

.

24

8

.

04

↓

2

5

1

1

0

±

±

.

.

.

.

.

2

42

05

85

16

67

.

1

.

28

61

MAP SOLUTION

VANILLA LA

VANILLA LA

RIEM-LA

RIEM-LA

LIN-LA

LIN-LA

LIN-RIEM-LA

LIN-RIEM-LA (BATCH)

LIN-RIEM-LA (BATCH)

MAP SOLUTION

LIN-RIEM-LA

LIN-RIEM-LA

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nEven 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?\n\nQUESTIONS:\nIn 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)?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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\n\nQUESTIONS:\n* 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\n\nLIMITATIONS:\nOnly 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\nOne 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.\n\nQUESTIONS:\nN/A\n\nLIMITATIONS:\nThere is a limitation section at the end of the paper.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\nOverall, 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.\n\nQUESTIONS:\na) 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\nLIMITATIONS:\nAs 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}