Riemannian Laplace approximations for Bayesian neural networks
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.
Federico Bergamin, Pablo Moreno-Muñoz, Søren Hauberg, Georgios Arvanitidis
# 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 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.

## 1 Introduction

_Bayesian deep learning_ estimates the weight-posterior of a neural network given data, i.e. \(p(\theta|\mathcal{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(\theta|\mathcal{D})\approx q(\theta|\mathcal{D})=\mathcal{N}(\theta|\mu,\Sigma)\), 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-conditioned implying a suboptimal behavior (Daxberger et al., 2021; 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

Figure 1: Our Riemannian Laplace approximation is a simple parametric distribution, which is shaped according to the local loss landscape through a Riemannian metric.
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 \(\mathcal{D}=\{\mathbf{x}_{n},\mathbf{y}_{n}\}_{n=1}^{N}\), consisting of inputs \(\mathbf{x}\in\mathbb{R}^{D}\) and outputs \(\mathbf{y}\in\mathbb{R}^{C}\). To enable _probabilistic modeling_, we use a likelihood \(p(\mathbf{y}|f_{\theta}(\mathbf{x}))\) which is either Gaussian (regression) or categorical (classification). This likelihood is parametrized by a deep neural network \(f_{\theta}:\mathbb{R}^{D}\rightarrow\mathbb{R}^{C}\), where \(\theta\in\mathbb{R}^{K}\) represent the weights for which we specify a Gaussian prior \(p(\theta)\). The predictive distribution of a new test point \(\mathbf{x}^{\prime}\) equals \(p(\mathbf{y}|\mathbf{x}^{\prime})=\int p(\mathbf{y}|\mathbf{x}^{\prime},\theta )p(\theta|\mathcal{D})\mathrm{d}\theta\) where \(p(\theta|\mathcal{D})\) is the true weight-posterior given the data \(\mathcal{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.

### 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(\theta|\mathcal{D})\), la constructs an approximate posterior distribution \(q_{\textsc{la}}(\theta|\mathcal{D})=\mathcal{N}(\theta|\theta_{*},\Sigma)\) that is centered at the _maximum a-posteriori_ (map) estimate

\[\theta_{*}=\operatorname*{arg\,max}_{\theta}\left\{\log p(\theta|\mathcal{D}) \right\}=\operatorname*{arg\,min}_{\theta}\underbrace{\left\{-\sum_{n=1}^{N} \log p(\mathbf{y}_{n}\mid\mathbf{x}_{n},\theta)-\log p(\theta)\right\}}_{ \mathcal{L}(\theta)}. \tag{1}\]

A Taylor expansion around \(\theta_{*}\) of the regularized loss \(\mathcal{L}(\theta)\) then yields

\[\hat{\mathcal{L}}(\theta)\approx\mathcal{L}(\theta_{*})+\langle\nabla_{\theta} \mathcal{L}(\theta)\big{|}_{\theta=\theta_{*}},(\theta-\theta_{*})\rangle+ \frac{1}{2}\langle(\theta-\theta_{*}),\mathrm{H}_{\theta}[\mathcal{L}](\theta )\big{|}_{\theta=\theta_{*}}(\theta-\theta_{*})\rangle, \tag{2}\]

where we know that \(\nabla_{\theta}\mathcal{L}(\theta)\big{|}_{\theta=\theta_{*}}\approx 0\), and \(\mathrm{H}_{\theta}[\mathcal{L}](\theta)\in\mathbb{R}^{K\times K}\) denotes the Hessian of the loss. This expansion suggests that the approximate posterior covariance should be the inverse Hessian \(\Sigma=\mathrm{H}_{\theta}[\mathcal{L}](\theta)\big{|}_{\theta=\theta_{*}}^{-1}\). The marginal likelihood of the data is then approximated as \(p(\mathcal{D})\approx\exp(-\mathcal{L}(\theta_{*}))(2\pi)^{D/2}\det(\Sigma)^{ \nicefrac{{1}}{{2}}}\). This is commonly used for training hyper-parameters of both the likelihood and the prior (Immer et al., 2021a; Antoran 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. (2021) 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 \(\theta_{*}\), and use the sample to use the "linearized" function \(f_{\theta}^{\text{lin}}(\mathbf{x})=f_{\theta_{*}}(\mathbf{x})+\langle\nabla_{ \theta}f_{\theta}(\mathbf{x})\big{|}_{\theta=\theta_{*}},\theta-\theta_{*}\rangle\) as predictive, where \(\nabla_{\theta}f_{\theta}(\mathbf{x})\big{|}_{\theta=\theta_{*}}\in\mathbb{R}^ {C\times K}\) is the Jacobian. Recently, this approach has been justified by Khan et al. (2019), Immer et al. (2021), 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 \(\theta\), 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., 2021; Daxberger et al., 2021). 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 \(\Theta\) 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).

### A loss-aware Riemannian geometry

For a given parameter value \(\theta\in\Theta\), we can measure the training loss \(\mathcal{L}(\theta)\) of the associated neural network. Assuming that the loss changes smoothly with \(\theta\), we can interpret the loss surface \(\mathcal{M}=g(\theta)=[\theta,\mathcal{L}(\theta)]\in\mathbb{R}^{K+1}\) as a \(K\)-dimensional manifold in \(\mathbb{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 \(\Theta\) as being the _intrinsic coordinates_ of the manifold \(\mathcal{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 \(\mathbf{J}_{g}(\theta)\mathbf{v}\in\mathbb{R}^{K+1}\), where \(\mathbf{J}_{g}:\Theta\rightarrow\mathbb{R}^{K+1\times K}\) is the Jacobian of \(g\) that spans the tangent space \(\mathcal{T}_{g(\theta)}\mathcal{M}\) at the point \(g(\theta)\in\mathcal{M}\) and \(\mathbf{v}\in\mathbb{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 which, we note, is now expressed directly in the intrinsic coordinates. From this observation, we define the _Riemannian metric_\(\mathbf{M}(\theta)=\mathbf{J}_{g}(\theta)^{\mathsf{T}}\mathbf{J}_{g}(\theta)\), 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 \(\mathbf{J}_{g}(\theta)=[\mathbb{I}_{K},\nabla_{\theta}\mathcal{L}]\)! such that the metric takes the form

\[\mathbf{M}(\theta)=\mathbb{I}_{K}+\nabla_{\theta}\mathcal{L}(\theta)\nabla_{ \theta}\mathcal{L}(\theta)^{\mathsf{T}}. \tag{3}\]

The exponential map.A local inner product allows us to define the length of a curve \(c:[0,1]\rightarrow\Theta\) as \(\texttt{length}[c]=\int_{0}^{1}\sqrt{\langle\hat{c}(t),\mathbf{M}(c(t))\hat{ c}(t)\rangle}\mathrm{d}t\), where \(\hat{c}(t)=\partial_{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 \(\texttt{Exp}:\mathcal{M}\times T_{\theta}\mathcal{M}\rightarrow\mathcal{M}\).

The tangential coordinates \(\mathbf{v}\) can be seen as a coordinate system for the neighborhood around \(\theta\), 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 \(\mathbf{J}_{g}(\theta)\), which implies that by changing this basis the associated coordinates change as well. By

Figure 2: The parameter space \(\Theta\) of the bnn together with examples of the Riemannian metric and the exponential map. Note that the Riemannian metric adapts to the shape of the loss which causes the geodesic to follow its shape.
orthonormalizing this basis we get the _normal coordinates_ where the metric vanishes. Let \(\mathbf{v}\) the tangential coordinates and \(\bar{\mathbf{v}}\) the corresponding normal coordinates, then it holds that

\[\langle\mathbf{v},\mathbf{M}(\theta)\mathbf{v}\rangle=\langle\bar{\mathbf{v}}, \bar{\mathbf{v}}\rangle\Rightarrow\mathbf{v}=\mathbf{A}(\theta)\bar{\mathbf{v }}\quad\text{with}\quad\mathbf{A}(\theta)=\mathbf{M}(\theta)^{-\nicefrac{{1}}{ {2}}}. \tag{4}\]

We will use the normal coordinates when doing Taylor expansions of the log-posterior, akin to standard Laplace approximations.

### 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 \(\theta_{*}\), \(h(\bar{\mathbf{v}})=\mathcal{L}(\text{Exp}_{\theta_{*}}(\mathbf{M}(\theta_{* })^{-\nicefrac{{1}}{{2}}}\bar{\mathbf{v}}))\). Following the standard Laplace approximation, we perform a second-order Taylor expansion of \(h\) as

\[\hat{h}(\bar{\mathbf{v}})\approx h(0)+\langle\partial_{\bar{\mathbf{v}}}h( \bar{\mathbf{v}})\big{|}_{\bar{\mathbf{v}}=0},\bar{\mathbf{v}}\rangle+\frac{1 }{2}\langle\bar{\mathbf{v}},\text{H}_{\bar{\mathbf{v}}}[h](\bar{\mathbf{v}}) \big{|}_{\bar{\mathbf{v}}=0}\bar{\mathbf{v}}\rangle, \tag{5}\]

where \(\partial_{\bar{\mathbf{v}}}h(\bar{\mathbf{v}})\big{|}_{\bar{\mathbf{v}}=0}= \mathbf{A}(\theta_{*})\bar{\mathbf{v}}\nabla_{\theta}\mathcal{L}(\theta)\big{|} _{\theta=\theta_{*}}\approx 0\) as \(\theta_{*}\) minimize the loss and \(\text{H}_{\bar{\mathbf{v}}}[h](\bar{\mathbf{v}})\big{|}_{\bar{\mathbf{v}}=0}= \mathbf{A}(\theta_{*})\bar{\mathbf{v}}\text{H}_{\theta}[\mathcal{L}](\theta) \mathbf{A}(\theta_{*})\big{|}_{\theta=\theta_{*}}\) with \(\text{H}_{\theta}[\mathcal{L}](\theta)\) the standard Euclidean Hessian matrix of the loss. Further details about this step can be found in appendix B.

**Tangential Laplace.** Similar to the standard Laplace approximation, we get a Gaussian approximate posterior \(\bar{q}(\bar{\mathbf{v}})=\mathcal{N}(\bar{\mathbf{v}}\mid 0,\ \overline{ \Sigma})\) on the tangent space in the normal coordinates with covariance \(\overline{\Sigma}=\text{H}_{\bar{\mathbf{v}}}[h](\mathbf{v})\big{|}_{\bar{ \mathbf{v}}=0}^{-1}\). Note that changing the normal coordinates \(\bar{\mathbf{v}}\) to tangential coordinates \(\mathbf{v}\) is a linear transformation and hence \(\mathbf{v}\sim\mathcal{N}(0,\mathbf{A}(\theta_{*})\overline{\mathbf{\Sigma}} \mathbf{A}(\theta_{*})\bar{\mathbf{v}})\), which means that this covariance is equal to \(\text{H}_{\theta}[\mathcal{L}](\theta)\big{|}_{\theta=\theta_{*}}^{-1}\) since \(\mathbf{A}(\theta_{*})\) is a symmetric matrix, and hence, it cancels out. The approximate posterior \(q_{\mathcal{T}}(\mathbf{v})=\mathcal{N}(\mathbf{v}\mid 0,\Sigma)\) 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|\mathbf{x}^{\prime},\mathcal{D})=\int p(y|\mathbf{x}^{\prime},\mathcal{D},\theta)q(\theta)\mathrm{d}\theta=\int p(y|\mathbf{x}^{\prime},\mathcal{D}, \text{Exp}_{\theta_{*}}(\mathbf{v}))q_{\mathcal{T}}(\mathbf{v})\mathrm{d} \mathbf{v}\approx\frac{1}{3}\sum_{s=1}^{S}p(y|\mathbf{x}^{\prime},\mathcal{D},\text{Exp}_{\theta_{*}}(\mathbf{v}_{s})),\ \mathbf{v}_{s}\sim q_{\mathcal{T}}(\mathbf{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_{\theta}^{\text{lin}}(\mathbf{x})\) is considered instead of the \(f_{\theta}(\mathbf{x})\) the loss function in (1) changes to \(\mathcal{L}^{\text{lin}}(\theta)\). Consequently, our Riemannian metric is computed under this new loss, and \(\nabla_{\theta}\mathcal{L}^{\text{lin}}(\theta)\) 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 \(\mathbf{x}\in\mathbb{R}^{2}\) to be in one class is \(p(C=1|\mathbf{x})=\sigma(\mathbf{x}^{\prime}\theta+b)\), where \(\sigma(\cdot)\) is the sigmoid function, \(\theta\in\mathbb{R}^{2}\) and \(b\in\mathbb{R}\). After learning the parameters, we fix \(b_{*}\) and show the posterior with respect to \(\theta\) together with the corresponding standard Laplace approximation (Fig. 2(a)).

Figure 3: The la assigns probability mass to regions where the true posterior is nearly zero, and a sample from this region corresponds to a poor classifier. Considering this sample as the initial velocity for the exponential map, the generated sample falls within the true posterior and the associated classifier performs well. As a result, our model quantifies better the uncertainty.
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. 2(b)). 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. 2(c)). We also show the predictive distribution for both approaches and even if both solve easily the classification problem, our model better quantifies uncertainty (Fig. 2(e)).

### 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

\[\ddot{c}(t)=-\nabla_{\theta}\mathcal{L}(c(t))\left(1+\nabla_{\theta}\mathcal{ L}(c(t))^{\intercal}\nabla_{\theta}\mathcal{L}(c(t))\right)^{-1}\langle\dot{c}(t ),\mathrm{H}_{\theta}[\mathcal{L}](c(t))\dot{c}(t)\rangle, \tag{6}\]

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 \(\mathcal{L}^{\text{lin}}(\theta)\) 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 Hessian-vector product between \(\mathrm{H}_{\theta}[\mathcal{L}](c(t))\) and \(\dot{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 \(\mathcal{B}\) of \(B\) random data points and use the associated loss in the ode (6). As usual, we assume that \(\mathcal{L}(\theta)\approx(\nicefrac{{N}}{{B}})\mathcal{L}_{\mathcal{B}}( \theta)\). Note that we only mini-batch the metric and not the covariance of our approximate posterior \(q_{\mathcal{T}}(\mathbf{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 \(\mathcal{L}(\theta)\) controls the geodesics via the associated Riemannian metric, so when we consider the full dataset we expect the samples to behave similarly to \(f_{\theta_{*}}(\mathbf{x})\). In other words, our approximate posterior generates weights near \(\theta_{*}\) resulting in models with similar or even better loss. When we consider a batch the geometry of the associated loss surface \(\mathcal{L}_{\mathcal{B}}(\theta)\) 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 \(\mathcal{L}^{\text{lin}}(\theta)\) may over-regularize the geodesics. Due to the linear nature of \(f_{\theta}^{\text{lin}}(\theta)\) the associated Riemannian metric is small only close to \(\theta_{*}\) so the generated samples are similar to \(f_{\theta_{*}}(\mathbf{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.

Figure 4: Analysis of mini-batching
Related work

**Bayesian neural networks.** Exact inference for bnn 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 bnn 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 bnn in a _post-hoc_ manner, thanks to easy-to-use software (Daxberger et al., 2021) and new approaches to scale up computation (Antoran et al., 2022). In this direction, other works have only considered a subset of the network parameters (Daxberger et al., 2021; 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. (2021), 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, 2018). 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 (2018) 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., 2021). 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 available in appendix D together with the complete training and modeling details.
### 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).

### Classification problems

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. 5(a)) with high uncertainty both within and away from the data support. Instead, the other three methods (Fig. 5(b)-5(d)) 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. 5(d) 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

Figure 5: Posterior samples under a simple _(top)_ and an overparametrized model _(bottom)_. Vanilla la is known to generate bad models, while our samples from riem-la quantify well the uncertainty.

Figure 6: Binary classification confidence estimate using \(100\) Monte-Carlo 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 of the data. Confidence plots of all different methods can be found in the supplementary material. Black lines are the decision boundaries.
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.

\begin{table}
\begin{tabular}{c c c c c c c} \hline \hline  & \multicolumn{6}{c}{brain dataset} \\ \cline{2-7}  & \multicolumn{3}{c}{prior optimized} & \multicolumn{3}{c}{pero not optimized} \\ \cline{2-7} method & Accuracy \(\uparrow\) & NLL\(\downarrow\) & Brier\(\downarrow\) & Accuracy \(\uparrow\) & NLL\(\downarrow\) & Brier\(\downarrow\) \\ \hline MAP & \(86.69\pm 0.34\) & \(0.333\pm 0.005\) & \(0.0930\pm 0.0015\) & \(86.69\pm 0.34\) & \(0.333\pm 0.005\) & \(0.0930\pm 0.0015\) \\ VAEIL-LA & \(95.50\pm 5.07\) & \(0.678\pm 0.009\) & \(0.2426\pm 0.0046\) & \(48.58\pm 2.32\) & \(0.700\pm 0.003\) & \(0.2534\pm 0.0017\) \\ LIN-LA & \(86.59\pm 0.37\) & \(0.325\pm 0.008\) & \(0.0956\pm 0.0023\) & \(86.92\pm 0.40\) & \(0.403\pm 0.012\) & \(0.1196\pm 0.0044\) \\ REM-LA & \(\mathbf{87.57\pm 0.07}\) & \(\mathbf{0.287\pm 0.002}\) & \(\mathbf{0.0886\pm 0.0006}\) & \(\mathbf{87.14\pm 0.20}\) & \(\mathbf{0.285\pm 0.001}\) & \(\mathbf{0.0878\pm 0.0006}\) \\ REM-LA (ACRES) & \(87.30\pm 0.08\) & \(\mathbf{0.286\pm 0.001}\) & \(\mathbf{0.0890\pm 0.0000}\) & \(\mathbf{87.32\pm 0.17}\) & \(\mathbf{0.294\pm 0.002}\) & \(\mathbf{0.0895\pm 0.0004}\) \\ LIN-BIBI-LA & \(87.02\pm 0.38\) & \(0.415\pm 0.029\) & \(0.0067\pm 0.0024\) & \(85.33\pm 0.31\) & \(0.884\pm 0.037\) & \(0.1252\pm 0.0022\) \\ LIN-BIBI-LA (AGRE) & \(\mathbf{87.72\pm 0.24}\) & \(0.298\pm 0.006\) & \(\mathbf{0.0887\pm 0.0012}\) & \(86.16\pm 0.21\) & \(0.352\pm 0.002\) & \(0.0934\pm 0.0011\) \\ \hline \hline \end{tabular}
\end{table}
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 5 different seeds.

\begin{table}
\begin{tabular}{c c c c c c} \hline \hline  & \multicolumn{6}{c}{prior precision optimized} \\ \cline{2-7} dataset & map & vanilla la & riem-la & linearized la & linearized rim-LA \\ \hline vehicle & \(0.975\pm 0.081\) & \(1.209\pm 0.020\) & \(\mathbf{0.454\pm 0.024}\) & \(0.875\pm 0.020\) & \(\mathbf{0.494\pm 0.044}\) \\ glass & \(2.084\pm 0.323\) & \(1.737\pm 0.037\) & \(\mathbf{1.047\pm 0.242}\) & \(1.365\pm 0.058\) & \(1.359\pm 0.299\) \\ ionosphere & \(1.032\pm 0.175\) & \(0.673\pm 0.013\) & \(\mathbf{0.344\pm 0.068}\) & \(0.497\pm 0.015\) & \(0.625\pm 0.110\) \\ waveform & \(1.076\pm 0.110\) & \(0.888\pm 0.030\) & \(\mathbf{0.459\pm 0.057}\) & \(0.640\pm 0.002\) & \(0.575\pm 0.065\) \\ australian & \(1.306\pm 0.146\) & \(0.684\pm 0.011\) & \(\mathbf{0.541\pm 0.053}\) & \(\mathbf{0.570\pm 0.016}\) & \(0.833\pm 0.108\) \\ breast cancer & \(\mathbf{0.225\pm 0.076}\) & \(0.594\pm 0.030\) & \(\mathbf{0.176\pm 0.092}\) & \(0.327\pm 0.022\) & \(\mathbf{0.202\pm 0.073}\) \\ \hline \hline \end{tabular}
\end{table}
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 5 different seeds.
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 batch-size 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

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 we have considered the 'obvious' mini-batching solution. Empirically, we find that this introduces some stochasticity in the sampling, which can actually be helpful in the posterior exploration. The computational cost also grows with the dimensionality of the parameter space, predominantly because the number of necessary solver steps increases as well. Our implementation relies on an off-the-shelf ode solver, and we expect that significant improvements can be obtained using a tailor-made numerical integration method.

\begin{table}
\begin{tabular}{c c c c c c} \hline \hline  & \multicolumn{5}{c}{CNN on MNIST - prior precision optimized} \\ \cline{2-6} method & Accuracy \(\uparrow\) & NLL\(\downarrow\) & Brier\(\downarrow\) & EC\(\downarrow\) & MCE\(\downarrow\) \\ \hline map & \(95.02\pm 0.17\) & \(0.167\pm 0.005\) & \(0.0075\pm 0.0002\) & \(1.05\pm 0.14\) & \(39.94\pm 14.27\) \\ vanilla la & \(88.69\pm 1.84\) & \(0.871\pm 0.026\) & \(0.0393\pm 0.0013\) & \(42.11\pm 1.22\) & \(50.52\pm 1.45\) \\ lin-la & \(94.91\pm 0.26\) & \(0.204\pm 0.006\) & \(0.0087\pm 0.0003\) & \(6.30\pm 0.8\) & \(39.30\pm 16.77\) \\ riem-la & \(\mathbf{96.74\pm 0.12}\) & \(\mathbf{0.115\pm 0.003}\) & \(\mathbf{0.0052\pm 0.0002}\) & \(2.48\pm 0.06\) & \(38.03\pm 15.02\) \\ riem-la (batches) & \(95.67\pm 0.19\) & \(0.170\pm 0.005\) & \(0.0072\pm 0.0002\) & \(5.40\pm 0.06\) & \(22.40\pm 0.51\) \\ lin-riem-la & \(95.44\pm 0.18\) & \(0.149\pm 0.004\) & \(0.0068\pm 0.0003\) & \(\mathbf{0.66\pm 0.03}\) & \(39.40\pm 14.75\) \\ lin-riem-la (batches) & \(95.14\pm 0.20\) & \(0.167\pm 0.004\) & \(0.0076\pm 0.0002\) & \(3.23\pm 0.04\) & \(\mathbf{18.10\pm 2.50}\) \\ \hline \hline \end{tabular} 
\begin{tabular}{c c c c c c} \hline \hline  & \multicolumn{5}{c}{CNN on FashionMNIST - prior precision optimized} \\ \cline{2-6} Method & Accuracy \(\uparrow\) & NLL\(\downarrow\) & Brier\(\downarrow\) & EC\(\downarrow\) & MCE\(\downarrow\) \\ \hline map & \(79.88\pm 0.09\) & \(0.541\pm 0.002\) & \(0.0276\pm 0.0000\) & \(\mathbf{1.66\pm 0.07}\) & \(24.07\pm 1.50\) \\ vanilla la & \(74.88\pm 0.83\) & \(1.026\pm 0.046\) & \(0.0482\pm 0.0019\) & \(31.63\pm 1.28\) & \(43.61\pm 2.95\) \\ lin-la & \(79.85\pm 0.13\) & \(0.549\pm 0.001\) & \(0.0278\pm 0.0000\) & \(3.23\pm 0.44\) & \(37.88\pm 17.98\) \\ riem-la & \(\mathbf{83.33\pm 0.17}\) & \(0.472\pm 0.001\) & \(\mathbf{0.0237\pm 0.0001}\) & \(3.13\pm 0.48\) & \(10.94\pm 2.11\) \\ riem-la (batches) & \(81.65\pm 0.18\) & \(0.525\pm 0.004\) & \(0.0263\pm 0.0002\) & \(5.80\pm 0.73\) & \(35.30\pm 18.40\) \\ lin-riem-la & \(81.33\pm 0.10\) & \(0.521\pm 0.004\) & \(0.0261\pm 0.0002\) & \(\mathbf{1.59\pm 0.40}\) & \(25.53\pm 0.10\) \\ lin-riem-la (batches) & \(80.49\pm 0.13\) & \(0.529\pm 0.003\) & \(0.0269\pm 0.0002\) & \(2.10\pm 0.42\) & \(\mathbf{6.14\pm 1.42}\) \\ \hline \hline \end{tabular}
\end{table}
Table 3: Image classification results using a CNN on MNIST and FashionMNIST. The network is trained on \(5000\) examples and we test the in-distribution 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 3 different seeds.
## 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.
