# Riemannian Laplace Approximation with the Fisher Metric

Hanlin Yu  Department of Computer Science, University of Helsinki  Marcelo
Hartmann  Department of Computer Science, University of Helsinki  Bernardo
Williams  Department of Computer Science, University of Helsinki  Mark
Girolami  Department of Engineering, University of Cambridge  Alan Turing
Institute  Arto Klami  Department of Computer Science, University of Helsinki

###### Abstract

The Laplace’s method approximates a target density with a Gaussian
distribution at its mode. It is computationally efficient and asymptotically
exact for Bayesian inference due to the Bernstein-von Mises theorem, but for
complex targets and finite-data posteriors it is often too crude an
approximation. A recent generalization of the Laplace Approximation transforms
the Gaussian approximation according to a chosen Riemannian geometry providing
a richer approximation family, while still retaining computational efficiency.
However, as shown here, its properties heavily depend on the chosen metric,
indeed the metric adopted in previous work results in approximations that are
overly narrow as well as being biased even at the limit of infinite data. We
correct this shortcoming by developing the approximation family further,
deriving two alternative variants that are exact at the limit of infinite
data, extending the theoretical analysis of the method, and demonstrating
practical improvements in a range of experiments.

##  1 Introduction

Functional distributional approximations offer a computationally attractive
alternative to Markov Chain Monte Carlo (MCMC) methods for statistical
(posterior in the Bayesian setting) inference of potentially large models,
building on variational inference (VI) [Blei et al., 2017], expectation
propagation (EP) [Minka, 2001], or Laplace approximation (LA) [Tierney and
Kadane, 1986, Rue et al., 2009]. Typically the methods need to make a
compromise between computational efficiency and accuracy, and often the
flexibility is achieved by non-linear transformations modelled as neural
networks that need to be separately trained for a given model, as in amortized
VI [Margossian and Blei, 2023] or normalizing flows [Papamakarios et al.,
2021]. We focus on perhaps the oldest approximation strategy, the Laplace
Approximation [Tierney and Kadane, 1986], and study how to increase the
flexibility of the approximation without resorting to trainable
transformations.

Figure 1: Left: The Fisher metric (green) captures the local curvature of the
target density (black). Right: Samples from a Gaussian distribution (orange)
are deterministically mapped using the Fisher metric to provide a flexible
approximation (blue).

The Laplace Approximation fits a Gaussian distribution at the mode of the
target density and is used in machine learning from non-conjugate Gaussian
processes [Rasmussen and Williams, 2006] to deep learning [Daxberger et al.,
2021]. It is computationally efficient, only requiring the maximum a
posteriori (MAP) estimate and the Hessian of the log-posterior that determine
the mean and covariance of the approximation. Despite its simplicity, it is
asymptotically exact at the limit of infinite data, due to the Bernstein-von
Mises theorem stating that the posterior converges to a Gaussian when the
probabilistic model has identifiable parameterisation [van der Vaart, 1998].
Obviously, as has been shown many times in the literature, the disadvantage is
that the approximation family is very limited for finite-data posteriors.

We seek to create an approximation that retains the advantages but increases
the flexibility, and do so by relying on the tools of Riemannian geometry
[Manfredo, 1992]. Instead of changing the approximation, we change the metric
of the parameter space in such a manner that after this modification the
approximation more accurately characterises the target. Concepts of Riemannian
geometry have been broadly considered in the approximate inference literature,
with Riemannian extensions of MCMC [Girolami and Calderhead, 2011],
stochastic-gradient MCMC [Patterson and Teh, 2013] and VI [Frank et al.,
2021]. Recently, Bergamin et al. [2023] introduced a Riemannian generalization
of the Laplace Approximation. The core idea (illustrated conceptually in
Figure 1, using the new metric from this paper) is to transform samples from a
Gaussian distribution using numerical integrators to follow geodesic paths
induced by a chosen geometry.

We build on Bergamin et al. [2023], expanding both theoretical and practical
understanding of the approximation. They used the Riemannian metric Hartmann
et al. [2022] proposed for Riemannian MCMC, which is computationally
attractive due to the efficient _exponential map_ for transforming the
samples. However, as shown here, the approximation is not asymptotically exact
and is also biased in practical cases. We resolve this theoretical and
practical limitation in two alternative ways. We first show that the method of
Bergamin et al. [2023] can be fixed by incorporating a _logarithmic map_ to
correct the bias, resulting in asymptotically exact but somewhat
computationally heavy and unstable approximation. A more practical solution is
to replace the metric itself. We explain how a metric based on the Fisher
Information Matrix (FIM) gives an approximation that is exact for targets that
are diffeomorphisms of a Gaussian, requires less function evaluations in
numerical integration, and is superior in how it performs in a range of tasks.

##  2 Preliminaries

###  2.1 Laplace Approximation

Let $\operatorname{\boldsymbol{\theta}}\in\operatorname{\boldsymbol{\Theta}}$
denote the $D$-dimensional vector in the parameter space
$\operatorname{\boldsymbol{\Theta}}$. Denote the log-posterior distribution
given observed data set $\operatorname{\boldsymbol{y}}$ $=$
$\\{y_{n}\\}_{n=1}^{N}$ as
$\ell_{\operatorname{\boldsymbol{y}}}(\operatorname{\boldsymbol{\theta}})$
$:=$
$\log\pi(\operatorname{\boldsymbol{\theta}}|\operatorname{\boldsymbol{y}})$.
The Laplace’s method expands a second-order Taylor series of
$\ell_{\operatorname{\boldsymbol{y}}}(\operatorname{\boldsymbol{\theta}})$ at
the mode (MAP) where the first-order term disappears. This is equivalent to
using a Gaussian approximation

$\pi(\operatorname{\boldsymbol{\theta}}|\operatorname{\boldsymbol{y}})\approx\operatorname{\mathcal{N}}\big{(}\operatorname{\boldsymbol{\theta}}|\hat{\operatorname{\boldsymbol{\theta}}},-\nabla^{2}\ell_{\operatorname{\boldsymbol{y}}}(\hat{\operatorname{\boldsymbol{\theta}}})^{-1}\big{)},$

where $\hat{\operatorname{\boldsymbol{\theta}}}$ is the MAP estimate and
$\nabla^{2}\ell_{\operatorname{\boldsymbol{y}}}(\hat{\operatorname{\boldsymbol{\theta}}})$
is the Hessian matrix of the log-posterior evaluated at
$\hat{\operatorname{\boldsymbol{\theta}}}$. We refer to this as ELA, for
Euclidean LA.

###  2.2 Riemannian geometry

##### Key concepts.

A _manifold_ is a space more general than an Euclidean space but that locally
acts like a Euclidean space. A _differentiable manifold_ is a manifold that
can be represented with a system of coordinates such that it will allow us to
extend notions of derivatives. A _Riemannian manifold_ is a differentiable
manifold endowed with a metric $g$ and this metric will, intrinsically,
generalise the local notions of angles, distances and derivatives on the
manifold.

The metric $g$ is formally defined as a function
$g_{\operatorname{\boldsymbol{\theta}}}:T_{\operatorname{\boldsymbol{\theta}}}\operatorname{\boldsymbol{\Theta}}\times
T_{\operatorname{\boldsymbol{\theta}}}\operatorname{\boldsymbol{\Theta}}\rightarrow\mathbb{R}$
where
$T_{\operatorname{\boldsymbol{\theta}}}\operatorname{\boldsymbol{\Theta}}$
denotes the _tangent space_ at $\operatorname{\boldsymbol{\theta}}$, given by
$T_{\operatorname{\boldsymbol{\theta}}}\operatorname{\boldsymbol{\Theta}}=\\{\tfrac{\mathrm{d}}{\mathrm{d}t}c_{i}(t)|_{t=0}=\operatorname{\boldsymbol{v}}_{i}\textrm{,
and}\ c_{i}(0)=\operatorname{\boldsymbol{\theta}},i\in\mathbb{N}\\}$. The
$\operatorname{\boldsymbol{v}}_{i}$s are known as the _tangent vector_ s. The
function $g$ is like the dot-product in Euclidean spaces, but takes a more
general form as a mapping
$(\operatorname{\boldsymbol{v}},\operatorname{\boldsymbol{u}})\mapsto\langle\operatorname{\boldsymbol{v}},\operatorname{\boldsymbol{G}}(\operatorname{\boldsymbol{\theta}})\operatorname{\boldsymbol{u}}\rangle$
where $\operatorname{\boldsymbol{G}}(\operatorname{\boldsymbol{\theta}})$ is a
symmetric positive-definite matrix called _metric tensor_ that collects the
coefficients of the metric as
$g\big{(}\tfrac{\partial}{\partial_{\theta_{i}}},\tfrac{\partial}{\partial_{\theta_{j}}}\big{)}=\operatorname{\boldsymbol{G}}(\operatorname{\boldsymbol{\theta}})_{ij}$.
We are free to choose the metric $g$ to encode desired properties for the
manifold.

##### Gradients and Hessians in Riemannian geometry.

In a given basis, the _Riemannian gradient_ of
$\ell_{\operatorname{\boldsymbol{y}}}(\operatorname{\boldsymbol{\theta}})$ is
given by $\textrm{grad}\
\ell_{\operatorname{\boldsymbol{y}}}(\operatorname{\boldsymbol{\theta}})=\operatorname{\boldsymbol{G}}(\operatorname{\boldsymbol{\theta}})^{-1}\nabla\ell_{\operatorname{\boldsymbol{y}}}(\operatorname{\boldsymbol{\theta}})$,
and the _Riemannian Hessian_ is $\textrm{Hess}\
\ell_{\operatorname{\boldsymbol{y}}}(\operatorname{\boldsymbol{\theta}})[\operatorname{\boldsymbol{v}}]=\nabla_{\operatorname{\boldsymbol{v}}}\textrm{grad}\
\ell_{\operatorname{\boldsymbol{y}}}(\operatorname{\boldsymbol{\theta}})=(\nabla^{2}\ell_{\operatorname{\boldsymbol{y}}}(\operatorname{\boldsymbol{\theta}})-\sum_{k}\operatorname{\boldsymbol{\Gamma}}(\operatorname{\boldsymbol{\theta}})^{k}\nabla\ell_{\operatorname{\boldsymbol{y}}}(\operatorname{\boldsymbol{\theta}})_{k})\operatorname{\boldsymbol{v}}$.
Here
$\operatorname{\boldsymbol{\Gamma}}(\operatorname{\boldsymbol{\theta}})^{k}=\\{\operatorname{\boldsymbol{\Gamma}}(\operatorname{\boldsymbol{\theta}})^{k}\\}_{i,j}$
are the _Christoffel symbols_ of the second kind in matrix forms. They can
also be expressed using the metric if the connection is chosen to be the Levi-
Civita connection [Manfredo, 1992] as

$\operatorname{\boldsymbol{\Gamma}}^{k}_{ij}=\frac{1}{2}\boldsymbol{G}^{kl}\left(\partial_{i}\boldsymbol{G}_{jl}+\partial_{j}\boldsymbol{G}_{il}-\partial_{l}\boldsymbol{G}_{ij}\right),$

where $\partial_{i}=\tfrac{\partial}{\partial{\theta_{i}}}$. Following
differential geometric convention, we use the Einstein summation. Note that in
this convention $\boldsymbol{G}^{kl}$ refers to elements of the _inverse_ of
the matrix
$\operatorname{\boldsymbol{G}}(\operatorname{\boldsymbol{\theta}})$.

##### Geodesic paths.

Another important notion for the development of this work is the _exponential
map_ and its inverse known as the _logarithmic map_. Formally, the
exponential-map is defined as a function
$\mathrm{Exp}_{\operatorname{\boldsymbol{\theta}}}:T_{\operatorname{\boldsymbol{\theta}}}\operatorname{\boldsymbol{\Theta}}\rightarrow\operatorname{\boldsymbol{\Theta}}$
that takes a vector on the tangent space of
$\operatorname{\boldsymbol{\Theta}}$ and maps it back onto
$\operatorname{\boldsymbol{\Theta}}$ [Boumal, 2023].

Intuitively, the exponential map is a function that returns the final position
after following the shortest path at $\operatorname{\boldsymbol{\theta}}$ in
the direction of $\operatorname{\boldsymbol{v}}$ for unit time. These paths
are called geodesics and they solve the geodesic equation given by [Manfredo,
1992]

$\derivative[2]{\operatorname{\boldsymbol{\theta}}_{k}}{t}+\operatorname{\boldsymbol{v}}^{i}\operatorname{\boldsymbol{v}}^{j}\operatorname{\boldsymbol{\Gamma}}^{k}_{ij}=0,\;k=1,2,\dots,D.$
(1)

The inverse
$\mathrm{Log}_{\operatorname{\boldsymbol{\theta}}}:\operatorname{\boldsymbol{\Theta}}\rightarrow
T_{\operatorname{\boldsymbol{\theta}}}\operatorname{\boldsymbol{\Theta}}$ of
the above is defined as
$\mathrm{Log}_{\operatorname{\boldsymbol{\theta}}}(\tilde{\operatorname{\boldsymbol{\theta}}})=\arg\min_{\operatorname{\boldsymbol{v}}\in
T_{\operatorname{\boldsymbol{\theta}}}\operatorname{\boldsymbol{\Theta}}}g_{\operatorname{\boldsymbol{\theta}}}(\operatorname{\boldsymbol{v}},\operatorname{\boldsymbol{v}})$
subject to
$\mathrm{Exp}_{\operatorname{\boldsymbol{\theta}}}(\operatorname{\boldsymbol{v}})=\tilde{\operatorname{\boldsymbol{\theta}}}$
[see Boumal, 2023]. It gives the $\operatorname{\boldsymbol{v}}$ for which the
exponential map starting from $\operatorname{\boldsymbol{\theta}}$ maps to
$\tilde{\operatorname{\boldsymbol{\theta}}}$.

##  3 Riemannian LA

Here we outline a general recipe for Riemannian Laplace Approximation (RLA).
The formulation matches the one recently proposed by Bergamin et al. [2023],
though here presented in a more general form where e.g. the choice of the
metric is still left open. We will later show how the metric influences both
the theoretical and practical properties of the method.

The Bernstein-von Mises Theorem states that, under certain regularity
conditions, the posterior at the limit of infinite data becomes Gaussian with
inverse of the FIM as the covariance [van der Vaart, 1998]. This makes ELA
exact at the limit. We want a Riemannian generalization of LA to retain this
property but to adapt for the local curvature of the target for non-Gaussian
targets and finite-data posteriors.

###  3.1 Principle and algorithm

The Riemannian extension of LA is based on a Taylor series expansion on the
tangent space
$T_{\operatorname{\boldsymbol{\theta}}}\operatorname{\boldsymbol{\Theta}}$ at
a given point
$\operatorname{\boldsymbol{\theta}}\in\operatorname{\boldsymbol{\Theta}}$.
Bergamin et al. [2023] used here the particular tangent space
$T_{\hat{\operatorname{\boldsymbol{\theta}}}}\operatorname{\boldsymbol{\Theta}}$,
where $\hat{\operatorname{\boldsymbol{\theta}}}$ is the MAP estimate. When the
parameter space of the model $\operatorname{\boldsymbol{\Theta}}$ is endowed
with a Riemannian geometry, the Taylor approximation of the function
$\operatorname{\boldsymbol{\Theta}}\ni\operatorname{\boldsymbol{\theta}}\mapsto\ell_{\operatorname{\boldsymbol{y}}}(\operatorname{\boldsymbol{\theta}})$
at $\hat{\operatorname{\boldsymbol{\theta}}}$ assumes a more general
formulation [see Boumal, 2023, Bergamin et al., 2023]

$\displaystyle\ell_{\operatorname{\boldsymbol{y}}}\big{(}\mathrm{Exp}_{\hat{\operatorname{\boldsymbol{\theta}}}}(\operatorname{\boldsymbol{v}})\big{)}$
$\displaystyle\approx\ell_{\operatorname{\boldsymbol{y}}}(\hat{\operatorname{\boldsymbol{\theta}}})+g_{\hat{\operatorname{\boldsymbol{\theta}}}}\big{(}\textrm{grad}\
\ell_{\operatorname{\boldsymbol{y}}}(\hat{\operatorname{\boldsymbol{\theta}}}),\operatorname{\boldsymbol{v}}\big{)}$
(2)
$\displaystyle+\frac{1}{2}g_{\hat{\operatorname{\boldsymbol{\theta}}}}\big{(}\textrm{Hess}\
\ell_{\operatorname{\boldsymbol{y}}}(\hat{\operatorname{\boldsymbol{\theta}}})[\operatorname{\boldsymbol{v}}],\operatorname{\boldsymbol{v}}\big{)}+\mathcal{O}(g_{\hat{\operatorname{\boldsymbol{\theta}}}}(\operatorname{\boldsymbol{v}},\operatorname{\boldsymbol{v}})^{3}).$

Algorithm 1 RLA-B, producing $N$ samples
$\operatorname{\boldsymbol{\theta}}^{[n]}$.

Obtain MAP estimate $\hat{\operatorname{\boldsymbol{\theta}}}$

Set
$\operatorname{\boldsymbol{\Sigma}}=(-\nabla^{2}\ell_{\operatorname{\boldsymbol{y}}}(\hat{\operatorname{\boldsymbol{\theta}}}))^{-1}$

for $n\leftarrow 1,\dots,N$ do

Obtain velocity
$\operatorname{\boldsymbol{v}}^{[n]}\sim\mathcal{N}(\operatorname{\boldsymbol{0}},\operatorname{\boldsymbol{\Sigma}})$

$\operatorname{\boldsymbol{\theta}}^{[n]}=\mathrm{Exp}_{\hat{\operatorname{\boldsymbol{\theta}}}}(\operatorname{\boldsymbol{v}}^{[n]})$

end for

Generally we cannot evaluate the approximation analytically, but can obtain
samples from it following Algorithm 1 that comprises of three distinct steps,
where the last two are repeated for each sample (and can be computed in
parallel):

  1. 1.

Find $\hat{\operatorname{\boldsymbol{\theta}}}$ where the approximation is
placed;

  2. 2.

Sample initial velocity $\operatorname{\boldsymbol{v}}\in
T_{\hat{\operatorname{\boldsymbol{\theta}}}}\operatorname{\boldsymbol{\Theta}}$
from a Gaussian distribution with suitably chosen covariance
$\operatorname{\boldsymbol{\Sigma}}$;

  3. 3.

Solve exponential map at $\hat{\operatorname{\boldsymbol{\theta}}}$ with
velocity $\operatorname{\boldsymbol{v}}$ to obtain the sample
$\operatorname{\boldsymbol{\theta}}$.

This constructs a warped distribution [De Bortoli et al., 2022] on the
manifold. The primary design choice is the metric $g$ that encodes the
intrinsic geometry of the problem, but as will be explained later there is
freedom also in the choices for the first two steps. Once $g$ is given, the
geodesic equation (1) can be formulated as an initial value problem of an
ordinary system of differential equations (ODE) and solved relatively
efficiently with numerical integration.

##### Euclidean Laplace Approximation.

Consider the special case of $g$ being the Euclidean inner product. Then
$\operatorname{\boldsymbol{G}}(\operatorname{\boldsymbol{\theta}})=\operatorname{\boldsymbol{I}}_{D}$
and the exponential map becomes
$\mathrm{Exp}_{\hat{\operatorname{\boldsymbol{\theta}}}}(\operatorname{\boldsymbol{v}})$
$=$ $\hat{\operatorname{\boldsymbol{\theta}}}+\operatorname{\boldsymbol{v}}$,
a unit step in the direction of the velocity $\operatorname{\boldsymbol{v}}$.
If velocities are sampled using
$\operatorname{\boldsymbol{\Sigma}}=-(\nabla^{2}\ell_{\operatorname{\boldsymbol{y}}}(\hat{\operatorname{\boldsymbol{\theta}}}))^{-1}$
we get the classical ELA.

##### Method of Bergamin et al. [2023].

In Equation (2), Bergamin et al. [2023] sets the metric tensor as
$\operatorname{\boldsymbol{G}}(\operatorname{\boldsymbol{\theta}})=\operatorname{\boldsymbol{I}}_{D}+\nabla\ell_{\operatorname{\boldsymbol{y}}}(\operatorname{\boldsymbol{\theta}})\nabla\ell_{\operatorname{\boldsymbol{y}}}(\operatorname{\boldsymbol{\theta}})^{\top}$,
which is a special case of the metric Hartmann et al. [2022] used for
geometric MCMC and coined as the _Monge metric_. This metric is defined solely
based on gradient information, and is a way to introduce curvature information
using the first order derivative of the target density function. It enables
fast computation of Christoffel symbols for the exponential map. We denote
this specific instance of RLA as RLA-B. Even though RLA-B was shown to work
well in a range of tasks and scales up for large problems, we will later
observe that it does not always improve over ELA.

|
---|---
Figure 2: Euclidean LA (left) is exact for a Gaussian target, but RLA-B
(right) is biased. Lines are the true contours and marginals, with samples and
histograms characterizing the approximation.

More importantly, the method is not exact even for Gaussian targets and hence
not for infinite-data posteriors, but instead underestimates the uncertainty.
This is seen by analysis of the exponential map. The solutions of the geodesic
equation keep the norm
$\norm{\operatorname{\boldsymbol{v}}(t)}^{2}_{\operatorname{\boldsymbol{G}}(\operatorname{\boldsymbol{\theta}}(t))}$
constant [Lee, John M., 2018], which in this metric expands to
$\norm{\operatorname{\boldsymbol{v}}(t)}^{2}+\langle\operatorname{\boldsymbol{v}}(t),\nabla\ell_{\operatorname{\boldsymbol{y}}}(\operatorname{\boldsymbol{\theta}}(t))\rangle^{2}$.
At the MAP estimate we have
$\nabla\ell_{\operatorname{\boldsymbol{y}}}(\hat{\operatorname{\boldsymbol{\theta}}})=0$
and hence
$\operatorname{\boldsymbol{G}}(\hat{\operatorname{\boldsymbol{\theta}}})=\operatorname{\boldsymbol{I}}_{D}$.
Denote the norm of the initial velocity by
$\norm{\operatorname{\boldsymbol{v}}_{0}}$. For every other $t$ we have
$\norm{\operatorname{\boldsymbol{v}}(t)}<\norm{\operatorname{\boldsymbol{v}}_{0}}$
because
$\langle\operatorname{\boldsymbol{v}}(t),\nabla\ell_{\operatorname{\boldsymbol{y}}}(\operatorname{\boldsymbol{\theta}}(t))\rangle^{2}>0$
except in the rare cases when the gradient is null or exactly orthogonal to
velocity. Since the exponential map integrates for unit time, the distance of
each sample to $\hat{\operatorname{\boldsymbol{\theta}}}$ (in the Euclidean
sense) is almost always smaller than for ELA, for which the distance is
$\norm{\operatorname{\boldsymbol{v}}_{0}}$. This is empirically demonstrated
in Figure 2, which shows that RLA-B does not recover even a two dimensional
isotropic Gaussian.

##  4 Improving Riemannian LA

Next we describe two alternative asymptotically valid RLA methods, and discuss
the choice of the starting point and covariance of velocities.

###  4.1 RLA-B with logarithmic map

We can make RLA-B exact for Gaussian targets by employing a logarithmic map to
solve for the initial velocity $\operatorname{\boldsymbol{v}}$ such that the
result of the exponential map retains the correct Euclidean distance to
$\hat{\operatorname{\boldsymbol{\theta}}}$.

Consider two manifolds $\operatorname{\mathcal{P}}$ and
$\operatorname{\mathcal{M}}$.
$\operatorname{\mathcal{P}}=(\operatorname{\boldsymbol{\Theta}},g_{\operatorname{\mathcal{P}}})$
has the metric tensor

$\displaystyle\operatorname{\boldsymbol{G}}_{\operatorname{\mathcal{P}}}(\operatorname{\boldsymbol{\theta}})$
$\displaystyle=I+\nabla\log\mathcal{N}(\operatorname{\boldsymbol{\theta}}|\hat{\operatorname{\boldsymbol{\theta}}},\operatorname{\boldsymbol{\Sigma}})\nabla\log\mathcal{N}(\operatorname{\boldsymbol{\theta}}|\hat{\operatorname{\boldsymbol{\theta}}},\operatorname{\boldsymbol{\Sigma}})^{\top}$

where
$\operatorname{\boldsymbol{\Sigma}}=(-\nabla^{2}\ell_{\operatorname{\boldsymbol{y}}}(\hat{\operatorname{\boldsymbol{\theta}}}))^{-1}$,
and
$\operatorname{\mathcal{M}}=(\operatorname{\boldsymbol{\Theta}},g_{\operatorname{\mathcal{M}}})$
the metric tensor

$\operatorname{\boldsymbol{G}}_{\operatorname{\mathcal{M}}}(\operatorname{\boldsymbol{\theta}})=I+\nabla\ell_{\operatorname{\boldsymbol{y}}}(\operatorname{\boldsymbol{\theta}})\nabla\ell_{\operatorname{\boldsymbol{y}}}(\operatorname{\boldsymbol{\theta}})^{\top}.$

$T_{\hat{\operatorname{\boldsymbol{\theta}}}}\operatorname{\mathcal{P}}$ and
$T_{\hat{\operatorname{\boldsymbol{\theta}}}}\operatorname{\mathcal{M}}$
coincide at $\hat{\operatorname{\boldsymbol{\theta}}}$ as the metric tensors
are identical, and hence we have a diffeomorphism

$\textrm{Exp}_{\operatorname{\mathcal{M}},\hat{\operatorname{\boldsymbol{\theta}}}}\circ\textrm{Log}_{\operatorname{\mathcal{P}},\hat{\operatorname{\boldsymbol{\theta}}}}:\operatorname{\mathcal{P}}\rightarrow\operatorname{\mathcal{M}}.$

Suppose we sample from a Gaussian distribution as in ELA and consider the
samples as points on manifold $\operatorname{\mathcal{P}}$. If
$\operatorname{\mathcal{P}}$ is equal to $\operatorname{\mathcal{M}}$ the
transformation given by
$\textrm{Exp}_{\operatorname{\mathcal{M}},\hat{\operatorname{\boldsymbol{\theta}}}}\circ\textrm{Log}_{\operatorname{\mathcal{P}},\hat{\operatorname{\boldsymbol{\theta}}}}$
results in exact samples from the target associated with
$\operatorname{\mathcal{M}}$. For non-Gaussian targets the method is not
exact, but we observed it can still help reducing the underestimation tendency
of RLA-B.

In Algorithm 2 the exponential map is computed as before on manifold
$\operatorname{\mathcal{M}}$, and the logarithmic map corresponds to a
boundary value problem of the geodesic ODE on manifold
$\operatorname{\mathcal{P}}$. It is solved numerically and benefits again from
efficient computation of Christoffel symbols. The logarithmic map is in
general more costly than the exponential map [Arvanitidis et al., 2016], but
we observed it can help reducing the number of function evaluations during the
exponential map and the overall cost is not necessarily higher. We refer to
this variant as RLA-BLog.

Algorithm 2 RLA-BLog, producing $N$ samples
$\operatorname{\boldsymbol{\theta}}^{[n]}$. The manifold
$\operatorname{\mathcal{P}}$ uses the metric induced by Gaussian
$\mathcal{N}(\tilde{\operatorname{\boldsymbol{\theta}}}|\hat{\operatorname{\boldsymbol{\theta}}},\operatorname{\boldsymbol{\Sigma}})$
and the manifold $\operatorname{\mathcal{M}}$ the metric induced by the target
$\pi(\operatorname{\boldsymbol{\theta}}|\operatorname{\boldsymbol{y}})$.

Obtain MAP estimate $\hat{\operatorname{\boldsymbol{\theta}}}$

$\operatorname{\boldsymbol{\Sigma}}=(-\nabla^{2}\ell_{\operatorname{\boldsymbol{y}}}(\hat{\operatorname{\boldsymbol{\theta}}}))^{-1}$

for $n\leftarrow 1\dots N$ do

Obtain base sample
$\bar{\operatorname{\boldsymbol{\theta}}}^{[n]}\sim\mathcal{N}(\operatorname{\boldsymbol{0}},\operatorname{\boldsymbol{\Sigma}})$

$\operatorname{\boldsymbol{v}}^{[n]}=\text{Log}_{\operatorname{\mathcal{P}},\hat{\operatorname{\boldsymbol{\theta}}}}(\bar{\operatorname{\boldsymbol{\theta}}}^{[n]})$

$\operatorname{\boldsymbol{\theta}}^{[n]}=\text{Exp}_{\operatorname{\mathcal{M}},\hat{\operatorname{\boldsymbol{\theta}}}}(\operatorname{\boldsymbol{v}}^{[n]})$

end for

###  4.2 Fisher Information Matrix as metric

The choices of Bergamin et al. [2023] were motivated by computational
arguments and relied on extrinsic notions in differential geometry, whereas we
turn the attention to intrinsic notions and the natural Riemannian metric
tensor for probabilistic machine learning.

###  4.3 Fisher metric

From a statistical point of view FIM is the lower bound for the variance of
unbiased estimators. At a first glance, it may not tell us much about possible
underlying geometries on $\operatorname{\boldsymbol{\Theta}}$. However, FIM
transforms like a 2-covariant tensor for regular probabilistic models as the
expected score function is null and it is symmetric and positive-definite by
definition [see Lehmann, 2003, Schervish, 2012]. Therefore the FIM can act as
the coefficients
$\operatorname{\boldsymbol{G}}(\operatorname{\boldsymbol{\theta}})$ for the
metric $g$ and the pair $(\operatorname{\boldsymbol{\Theta}},g)$ defines an
abstract Riemannian manifold.

We propose using a metric tensor

$\operatorname{\boldsymbol{G}}(\operatorname{\boldsymbol{\theta}})=\operatorname{\mathbb{E}}_{\operatorname{\boldsymbol{Y}}|\operatorname{\boldsymbol{\theta}}}\left[-\nabla^{2}\log\pi(\operatorname{\boldsymbol{Y}}|\operatorname{\boldsymbol{\theta}})\right]-\nabla^{2}\log\pi(\operatorname{\boldsymbol{\theta}}),$
(3)

which is the FIM plus the negative Hessian of the log-prior, accounting for
both the prior and the likelihood. This formulation has been used by Girolami
and Calderhead [2011], Lan et al. [2015] as the metric for Riemann Manifold
MCMC algorithms and by Hartmann and Vanhatalo [2019] in natural gradient
descent for Gaussian processes with non log-concave likelihoods. For
simplicity, we call it the _Fisher metric_ and the resulting approximation
RLA-F.

The same form can also be used as $\operatorname{\boldsymbol{\Sigma}}$ when
sampling the velocities $\operatorname{\boldsymbol{v}}$, instead of the
negative Hessian as in ELA and RLA-B. The negative Hessian is not guaranteed
to be positive definite (e.g. for a neural network), which causes numerical
problems, whereas (3) is. In addition, we will later show that the Fisher
metric coincides with the negative Hessian for a specific class of targets
(Theorem 3).

####  4.3.1 Hausdorff MAP

The classical MAP estimate is not invariant under reparameterizations in
Euclidean geometry [see Jermyn, 2005]. From the Riemannian perspective, a
natural invariant alternative is the maximum value of the posterior density
under the Hausdorff measure. When $\operatorname{\boldsymbol{\Theta}}$ is
endowed with $g$, the probability density function on the manifold is given by

$\pi^{\operatorname{\boldsymbol{G}}}(\operatorname{\boldsymbol{\theta}})=\frac{\pi(\operatorname{\boldsymbol{\theta}})}{\sqrt{\det\operatorname{\boldsymbol{G}}(\operatorname{\boldsymbol{\theta}})}};$

see Theorem 3.2.5 by Federer [1969].

If sampling the velocities $\operatorname{\boldsymbol{v}}$ using the Fisher
metric, we can justify the use of the Hausdorff MAP as a reparameterization
(from the differential geometric viewpoint), where the Fisher is locally
identity around the MAP. This makes the gradients zero and consequently the
second-order Taylor series still corresponds to a Gaussian. We will later show
that both Hausdorff and Euclidean MAP can provide good approximations.

####  4.3.2 Theoretical basis

This variant is asymptotically exact for Gaussian targets, including posterior
distributions of identifiable models at the limit of large data, as determined
by the following theorems. Proofs for all of these theorems are provided in
Appenfix A. Note that Knollmüller and Enßlin [2020] also showed conclusions
similar to Theorem 1 in context of VI.

######  Theorem 1.

For Gaussian (or uniform) prior and Gaussian likelihood with fixed covariance,
the Fisher metric is constant.

######  Theorem 2.

Under the conditions of Bernstein-von Mises theorem, the Fisher metric
converges to a fixed matrix at the limit of infinite data.

For constant metric (or for one that converges to a constant one), the
geodesics become straight lines in the Euclidean sense. Consequently the
approximation becomes identical to ELA and hence exact for the target that is
(or converges to) a Gaussian.

We can also make a stronger statement for specific types of targets and
priors:

######  Theorem 3.

With an invariant prior, e.g. Jeffreys prior, RLA-F with Hausdorff MAP is
exact for probabilistic models whose target distributions are diffeomorphic
with Gaussians, for which the negative Hessian at the Hausdorff MAP coincides
with the Fisher metric.

Denote a diffeomorphic transformation
$\operatorname{\boldsymbol{\theta}}=\operatorname{\boldsymbol{\phi}}(\operatorname{\boldsymbol{\psi}})$
for a Gaussian distribution
$\mathcal{N}(\operatorname{\boldsymbol{\psi}}|\operatorname{\boldsymbol{\mu}},\operatorname{\boldsymbol{S}})$.
Since Gaussians are symmetric, observing how Riemannian metrics transform,
define a metric for the space of $\operatorname{\boldsymbol{\theta}}$ as

$\operatorname{\boldsymbol{G}}_{\operatorname{\boldsymbol{\Psi}}}=\left(\partialderivative{\operatorname{\boldsymbol{\theta}}}{\operatorname{\boldsymbol{\psi}}}\right)^{\top}\operatorname{\boldsymbol{G}}_{\operatorname{\boldsymbol{\Theta}}}\partialderivative{\operatorname{\boldsymbol{\theta}}}{\operatorname{\boldsymbol{\psi}}}.$

With
$(\operatorname{\boldsymbol{G}}_{\operatorname{\boldsymbol{\Psi}}})_{i,j}$ $=$
$\operatorname{\mathbb{E}}_{\operatorname{\boldsymbol{\psi}}}\big{(}-\partial^{2}_{\mu_{i},\mu_{j}}\log\mathcal{N}(\operatorname{\boldsymbol{\psi}}|\operatorname{\boldsymbol{\mu}},\operatorname{\boldsymbol{S}})\big{)}=\operatorname{\boldsymbol{S}}^{-1}_{i,j}$
inspired by the Fisher metric, we have

$\operatorname{\boldsymbol{G}}_{\operatorname{\boldsymbol{\Theta}}}=\left(\partialderivative{\operatorname{\boldsymbol{\theta}}}{\operatorname{\boldsymbol{\psi}}}\right)^{-\top}\operatorname{\boldsymbol{S}}^{-1}\left(\partialderivative{\operatorname{\boldsymbol{\theta}}}{\operatorname{\boldsymbol{\psi}}}\right)^{-1}.$

##### Example.

Consider the Squiggle distribution, used previously e.g. for evaluation of
MCMC samplers [Hartmann et al., 2022],

$\displaystyle\pi(\theta_{1},\theta_{2}|\operatorname{\boldsymbol{\mu}},\operatorname{\boldsymbol{S}})$
$\displaystyle=\operatorname{\mathcal{N}}(\operatorname{\boldsymbol{\phi}}^{-1}_{a}(\theta_{1},\theta_{2})|\operatorname{\boldsymbol{\mu}},\operatorname{\boldsymbol{S}}),$
$\displaystyle\operatorname{\boldsymbol{\phi}}^{-1}_{a}(\theta_{1},\theta_{2})$
$\displaystyle=(\theta_{1},\theta_{2}+\sin(a\theta_{1})).$

This can be seen as a probability density as

$\displaystyle\pi(\operatorname{\boldsymbol{\mu}}|\operatorname{\boldsymbol{\theta}})$
$\displaystyle=\operatorname{\mathcal{N}}(\operatorname{\boldsymbol{\mu}}|\operatorname{\boldsymbol{\phi}}^{-1}_{a}(\operatorname{\boldsymbol{\theta}}),\operatorname{\boldsymbol{S}})\
\textrm{and}\
\pi_{J}(\operatorname{\boldsymbol{\theta}})=\sqrt{\det\operatorname{\boldsymbol{G}}_{\operatorname{\boldsymbol{\Theta}}}}$

where $\pi_{J}$ denotes the Jeffreys prior.

The Jacobian to form the Fisher metric is given by

$\left(\partialderivative{\operatorname{\boldsymbol{\theta}}}{\operatorname{\boldsymbol{\psi}}_{a}}\right)^{-1}=\begin{bmatrix}1&0\\\
a\cos(a\theta_{1})&1\end{bmatrix}$

and for $\operatorname{\boldsymbol{\mu}}=\operatorname{\boldsymbol{0}}$ used
here the Hausdorff MAP is also $\operatorname{\boldsymbol{0}}$. As stated by
Theorem 3, RLA-F is exact for this target due to the diffeomorphism detailed
above. Figure 3 empirically validates this.

|
---|---
Figure 3: RLA-F (left) is exact (already for finite data) for the squiggle
distribution with complex shape due to diffeomorphism with a Gaussian, whereas
RLA-B (right) is too narrow everywhere and generally biased.

####  4.3.3 Computation of the Fisher metric

##### FIM for neural networks (NNs).

The Fisher metric is known for many probabilistic models of basic form, i.e.
the likelihood part [see Yang and Berger, 1996, George Casella, 2001]. For
instance, for exponential family the FIM always has closed-form expression
[George Casella, 2001]. Building on these, we can easily form the Fisher
metric for more complex models. Let $\phi$ be the parameters of a model
$\pi(y|\phi)$ in its basic form in a $P$-dimensional parameter space $\Phi$.
Denote the FIM for $\Phi$ as
$\operatorname{\boldsymbol{G}}_{\operatorname{\boldsymbol{\Phi}}}(\cdot)$.
Consider inputs $\\{\operatorname{\boldsymbol{x}}_{n}\\}_{n=1}^{N}$ with
associated observations $\operatorname{\boldsymbol{y}}=\\{y_{n}\\}_{n=1}^{N}$
and a NN that maps $\mathcal{X}$ $\ni$ $\operatorname{\boldsymbol{x}}_{n}$
$\overset{f_{\operatorname{\boldsymbol{\theta}}}}{\mapsto}\Phi$, where
$\mathcal{X}$ is the space of inputs and $\operatorname{\boldsymbol{\theta}}$
the vector of parameters. Let
$\ell_{\operatorname{\boldsymbol{y}}}(\operatorname{\boldsymbol{\theta}})=\sum_{n}^{N}\log\pi\big{(}y_{n}|f_{\operatorname{\boldsymbol{\theta}}}(\operatorname{\boldsymbol{x}}_{n}))$
denote the likelihood function of $\operatorname{\boldsymbol{\theta}}$ and
define
$\phi_{n}=f_{\operatorname{\boldsymbol{\theta}}}(\operatorname{\boldsymbol{x}}_{n})\in\operatorname{\boldsymbol{\Phi}}\
\forall n$. Using the chain rule for the Hessian matrix of
$\ell_{\operatorname{\boldsymbol{y}}}(\operatorname{\boldsymbol{\theta}})$
w.r.t $\operatorname{\boldsymbol{\theta}}$ the expression involving the score
function vanishes in expectation. Therefore, the metric tensor on
$\operatorname{\boldsymbol{\theta}}$ reads

$\displaystyle\operatorname{\mathbb{E}}_{\operatorname{\boldsymbol{Y}}}(-$
$\displaystyle\nabla^{2}\ell_{\operatorname{\boldsymbol{Y}}}(\operatorname{\boldsymbol{\theta}})_{i,j})=-\sum_{n}^{N}\operatorname{\mathbb{E}}_{Y_{n}}\bigg{(}\partial^{2}_{i,j}\log\pi\big{(}Y_{n}|\underbrace{f_{\operatorname{\boldsymbol{\theta}}}(\operatorname{\boldsymbol{x}}_{n})}_{\phi_{n}})\bigg{)}$
$\displaystyle=\sum_{n}^{N}\boldsymbol{J}_{n}^{\top}\operatorname{\boldsymbol{G}}_{\operatorname{\boldsymbol{\Phi}}}(f_{\operatorname{\boldsymbol{\theta}}}(\operatorname{\boldsymbol{x}}_{n}))\boldsymbol{J}_{n},$

where $\boldsymbol{J}_{n}$ $=$
$[\nabla(\boldsymbol{\phi}_{n})_{1}\cdots\nabla(\boldsymbol{\phi}_{n})_{P}]^{\top}$
is a $P\times D$ Jacobian matrix of the NN at the $n^{th}$ input. In other
words, we can form the FIM for the whole network by transforming the FIM on
its basic form with Jacobians that can be computed using standard automatic
differentiation. Formally, this is a pullback metric from the parameter space
$\operatorname{\boldsymbol{\Phi}}$ to the NN parameter space
$\operatorname{\boldsymbol{\Theta}}$; see further details in Appendix F.2.

##### Computational cost.

The core computational cost for RLA comes from obtaining the acceleration
given the position and velocity, which is needed in each step of the ODE
integrator. We have

$\partialderivative[2]{\operatorname{\boldsymbol{\theta}}_{k}}{t}=-\frac{1}{2}\operatorname{\boldsymbol{G}}^{kl}\left[\left(\partial_{i}\operatorname{\boldsymbol{G}}^{kl}+\partial_{j}\operatorname{\boldsymbol{G}}_{il}-\partial_{l}\operatorname{\boldsymbol{G}}_{ij}\right)\operatorname{\boldsymbol{v}}^{i}\operatorname{\boldsymbol{v}}^{j}\right],$

which is a product of the inverse of the metric
($\operatorname{\boldsymbol{G}}^{kl}$) and a vector. For RLA-F, the
computational cost is dominated by the inversion. The cost for drawing $N$
samples hence becomes $O(NTD^{3})$, where $T$ is the number of evaluations
during integration. RLA-B has lower cost due to avoiding direct inversion, but
we will later see that with Fisher metric we can often use considerably
smaller $T$ that balances the difference. Also note that already standard ELA
has cost $O(D^{3}+ND^{2})$ due to inversion of Hessian and multivariate
sampling. The sampling parallelizes over $N$ for all methods.

##  5 Experiments

Code for reproducing the experiments is available at
https://github.com/ksnxr/RLAF. Additional details for the experiments, such as
choices for various hyperparameters, are provided in Appendix G, and Appenfix
H provides additional experimental results to complement the main experiments
reported here.

###  5.1 Experimental setup

We evaluate three methods: RLA-F using the Fisher metric (Section 4.2), RLA-B
as proposed by Bergamin et al. [2023], and RLA-BLog using the logarithmic map
(Section 4.1). In addition, we show results for standard ELA and discuss the
choice of the MAP estimate and covariance for sampling initial velocities when
relevant. We repeat the experiments $5$ times and report averaged results over
the repetitions.

For evaluating the approximation accuracy, we generate $20,000$ samples using
the NUTS sampler in Stan [Stan Dev Team, 2023] as the ground truth. For low-
dimensional problems we directly compare the approximation with the posterior
samples by computing the Wasserstein distance $\mathcal{W}$ as in Zhang et al.
[2022]. For the neural networks we compare the model predictions instead, due
to non-identifiability of the posterior, measured using mean squared error
(MSE) and negative log-likelihood (NLL).

Following Bergamin et al. [2023], we use a numerical ODE solver with adaptive
step sizes. We report the actual number of integration steps $T$ (and in some
case running time) for the exponential map as an indicator of the complexity
of the metric. The numbers are for producing one sample since the algorithm
parallelizes trivially over the samples. For the NN experiment we use SciPy
[Virtanen et al., 2020] for integration, and for other experiments Diffrax
[Kidger, 2021].

###  5.2 Banana distribution

|  |
---|---|---
Figure 4: Left: The banana has two Euclidean MAPs (blue) but only one Hausdorff MAP (brown). RLA-F depends on the MAP choice, with Hausdorff (middle) being here superior to Euclidean (right).  | ELA | RLA-B | RLA-BLog | RLA-F
---|---|---|---|---
MAP | $\mathcal{W}$ | $\mathcal{W}$ | $T$ | $\mathcal{W}$ | $T$ | $\mathcal{W}$ | $T$
Euclidean | [1.434, 0.01] | [0.811, 0.009] | 70.4 | [0.788, 0.013] | 70.4 | [0.791, 0.014] | 24.9
Hausdorff | [1.386, 0.009] | [0.341, 0.006] | 51.3 | [0.208, 0.007] | 60.0 | [0.143, 0.009] | 32.7
Table 1: Banana distribution results as $[\text{mean},\text{std}]$. Bold font
indicates the best method.

We first compare the approximations in a 2D problem with difficult geometry,
the banana distribution

$\displaystyle\operatorname{\boldsymbol{\theta}}\overset{\textrm{ind.}}{\sim}\operatorname{\mathcal{N}}(0,\sigma_{\operatorname{\boldsymbol{\theta}}}^{2})\
\textrm{and}\
Y_{n}|\operatorname{\boldsymbol{\theta}}\sim\operatorname{\mathcal{N}}(\theta_{1}+\theta_{2}^{2},\sigma_{Y}^{2}),$

where we set $\sigma_{\operatorname{\boldsymbol{\theta}}}=\sigma_{Y}=2$,
$N=100$. The samples are obtained using $\theta_{1}=0.5$ and
$\theta_{2}^{2}=0.75$. The distribution has closed-form FIM (Section 4.3.3
also applies).

|  | ELA | RLA-B | RLA-BLog | RLA-F
---|---|---|---|---|---
| data | $\mathcal{W}$ | $\mathcal{W}$ | $T$ | $\mathcal{W}$ | $T$ | $\mathcal{W}$ | $T$
stand. | Ripl | [0.106, 0.004] | [0.236, 0.001] | 29.7 | [0.16, 0.065] | 45.2 | [0.064, 0.002] | 12.2
Pima | [0.149, 0.0] | [0.274, 0.0] | 41.7 | [0.21, 0.007] | 78.5 | [0.147, 0.0] | 12.2
Hear | [0.529, 0.001] | [0.649, 0.0] | 40.2 | [1.154, 0.189] | 71.6 | [0.514, 0.0] | 13.2
Aust | [0.441, 0.001] | [0.522, 0.001] | 51.5 | [0.745, 0.203] | 79.1 | [0.417, 0.001] | 18.0
Germ | [0.388, 0.001] | [0.431, 0.0] | 58.4 | [0.676, 0.1] | 104.0 | [0.387, 0.0] | 14.2
raw | Ripl | [0.437, 0.017] | [0.489, 0.012] | 30.2 | [0.554, 0.253] | 45.0 | [0.247, 0.012] | 12.7
Pima | [0.21, 0.01] | [0.294, 0.004] | 5632.9 | [0.216, 0.014] | 1690.4 | [0.112, 0.008] | 15.4
Hear | [0.842, 0.026] | [0.96, 0.012] | 7868.2 | [0.898, 0.029] | 3161.1 | [0.644, 0.012] | 17.9
Aust | [0.454, 0.004] | [0.455, 0.01] | 18513.8 | [0.467, 0.005] | 12298.6 | [0.378, 0.005] | 18.0
Germ | [0.846, 0.002] | [0.92, 0.003] | 3545.9 | [0.898, 0.003] | 2985.5 | [0.823, 0.001] | 17.8
Table 2: Logistic regression results as $[\text{mean},\text{std}]$. Bold font
indicates the best method. Italic fonts indicate the integrator reached
maximum number of steps in at least one run.

The distribution has two modes symmetric across the x-axis with the same log-
posterior. However, the Hausdorff MAP is unique and at the x-axis. Figure 4
compares the resulting approximations with RLA-F for the different MAP
estimates; with Euclidean MAP, we use Hessian precision for sampling the
velocities, whereas with Hausdorff MAP we use the Fisher precision. Table 1
shows that all Riemannian methods clearly outperform classical ELA for this
target, and for all Riemannian methods using the Hausdorff MAP is considerably
better. Both of the newly proposed methods outperform RLA-B.

###  5.3 Bayesian logistic regression

Following Girolami and Calderhead [2011] and Lan et al. [2015], we applied
Bayesian logistic regression on five datasets (details in Appendix G.9) using
the model

$\displaystyle
Y_{n}|\operatorname{\boldsymbol{\theta}}\sim\textrm{Bernoulli}(\sigma(\operatorname{\boldsymbol{\theta}}^{\top}\operatorname{\boldsymbol{x}}_{n}))\
\ \textrm{and}\ \
\operatorname{\boldsymbol{\theta}}\overset{\textrm{ind.}}{\sim}\operatorname{\mathcal{N}}(0,\alpha),$

where $\sigma:\mathcal{X}\rightarrow(0,1)$ is the sigmoid function and
$\alpha=100$. The parameter space of the Bernoulli model in its basic form is
$\operatorname{\boldsymbol{\phi}}\in(0,1)=\Phi$, whose FIM is known. Therefore
the Fisher metric on the parameter space $\operatorname{\boldsymbol{\Theta}}$
becomes $\operatorname{\boldsymbol{G}}$ $=$
$\operatorname{\boldsymbol{X}}^{\top}\operatorname{\boldsymbol{\Lambda}}\operatorname{\boldsymbol{X}}+\alpha^{-1}\operatorname{\boldsymbol{I}}$,
where
$\operatorname{\boldsymbol{X}}=[\operatorname{\boldsymbol{x}}_{1}\cdots\operatorname{\boldsymbol{x}}_{N}]$
is the covariate matrix (or inputs) and $\operatorname{\boldsymbol{\Lambda}}$
is a diagonal matrix with elements
$\operatorname{\boldsymbol{\Lambda}}_{nn}=\sigma\big{(}\boldsymbol{\theta}^{\top}\boldsymbol{x}_{n}\big{)}\big{(}1-\sigma(\boldsymbol{\theta}^{\top}\boldsymbol{x}_{n}\big{)}\big{)}$.
We use Euclidean MAP, and here the negative Hessian coincides with the Fisher
metric [Girolami and Calderhead, 2011] (so both choices for
$\operatorname{\boldsymbol{\Sigma}}$ are the same).

Table 2 reports approximation accuracy for two setups: With standardized
(z-score) inputs to make the geometry of the problem easier, and with the raw
inputs where the metric also needs to handle potentially large scale
differences. For both cases, RLA-F is clearly the best on all data sets, and
the approximations using the Monge metric are worse than ELA. The Fisher
metric results in smoother integration surface, especially when not
standardizing the inputs, with the Monge metric variants needing up to $1000$
times more evaluations and consequently also more time despite lower
complexity (see Appendix H.4 for times).

###  5.4 Neural network regression

Finally, we ran an experiment similar to the one by Bergamin et al. [2023],
with data from Snelson and Ghahramani [2005]. The task is a 1D regression
problem using an NN of size $1$-$10$-$1$ with $\tanh$ activation. The methods
are implemented using Daxberger et al. [2021], with Fisher precision (the
default option for full LA), $\hat{\operatorname{\boldsymbol{\theta}}}$ found
by standard MAP training, and prior precision and noise optimized posthoc. MSE
and NLL are computed using $500$ samples.

|  |  |
---|---|---|---
|  |  |
Figure 5: NN regression with complete (top) and gap (bottom) training data. Methods from left to right: RLA-B; RLA-BLog; RLA-F; NUTS. Gray shading denotes the part of $x$-axis with training data, dark line is the mean prediction, and blue lines are samples. | Complete | Gap
---|---|---
method | MSE | NLL | $T$ | time | MSE | NLL | $T$ | time
ELA | [2.235, 0.617] | [3.02, 0.064] | N/A | N/A | [1.365, 0.407] | [3.171, 0.075] | N/A | N/A
RLA-B | [0.2, 0.237] | [0.193, 0.003] | 6677.8 | 34.1 | [1.646, 2.722] | [1.031, 0.02] | 6573.2 | 34.0
RLA-BLog | [0.107, 0.054] | [0.201, 0.004] | 6674.7 | 37.7 | [0.27, 0.068] | [1.122, 0.029] | 4793.9 | 28.5
RLA-F | [0.072, 0.0] | [0.121, 0.002] | 242.9 | 0.8 | [0.564, 0.066] | [1.063, 0.017] | 727.7 | 2.4
NUTS | [0.073, 0.0] | [0.126, 0.002] | N/A | N/A | [0.394, 0.02] | [0.872, 0.029] | N/A | N/A
Table 3: NN regression results as $[\text{mean},\text{std}]$. Bold font
indicates the best method.

Fisher metric and Christoffel symbols follow from Section 4.3.3 and are

$\displaystyle\operatorname{\boldsymbol{G}}_{kl}$
$\displaystyle=\frac{1}{\sigma^{2}}\sum_{n=1}^{N}\partial_{k}f_{\operatorname{\boldsymbol{\theta}}}(\operatorname{\boldsymbol{x}}_{n})\partial_{l}f_{\operatorname{\boldsymbol{\theta}}}(\operatorname{\boldsymbol{x}}_{n})+\alpha^{-1}\operatorname{\boldsymbol{I}}\;\;\text{and}$
$\displaystyle\operatorname{\boldsymbol{v}}^{i}\operatorname{\boldsymbol{v}}^{j}\operatorname{\boldsymbol{\Gamma}}^{k}_{ij}$
$\displaystyle=\frac{1}{\sigma^{2}}\operatorname{\boldsymbol{G}}^{kl}\sum_{n=1}^{N}\partial_{l}f_{\operatorname{\boldsymbol{\theta}}}(\operatorname{\boldsymbol{x}}_{n})\Big{(}\operatorname{\boldsymbol{v}}^{i}\operatorname{\boldsymbol{v}}^{j}\partial^{2}_{i,j}f_{\operatorname{\boldsymbol{\theta}}}(\operatorname{\boldsymbol{x}}_{n})\Big{)},$

as we use Gaussian prior with precision $\alpha$ on every
$\operatorname{\boldsymbol{\theta}}$ and Gaussian probabilistic model for $Y$;
see Calin and Udrişte [2014] and Song et al. [2018] for derivations.

Following Bergamin et al. [2023], we run experiments with training samples
covering a continuous area of inputs (complete) and with a gap between $1.5$
and $3$ (gap) to measure the in-between uncertainty quantification [Foong et
al., 2019]. Table 3 quantifies the results and shows that RLA-F is in general
the best approximation and for the complete case matches also NUTS. As is well
known, Euclidean LA does not work without linearization of the network
[Daxberger et al., 2021]. Figure 5 illustrates the two scenarios, where RLA-F
is effectively indistinguishable from NUTS, whereas both RLA-B and RLA-BLog
overestimate predictive variance and some of the posterior samples are clearly
off.

##  6 Discussion and conclusions

##### Main observation.

Our key message is that the Riemannian Laplace Approximation proposed by
Bergamin et al. [2023] is promising as practical and flexible approximation,
but their specific variant is biased. The bias can be resolved by either
slightly modifying the algorithm or by considering more suitable metrics, and
the corrected approximation family has interesting theoretical properties. For
instance, with the Fisher metric, it is exact for targets diffeomorphic with a
Gaussian. This opens up new research directions for the study of metrics in
inference tasks.

##### On Monge metric.

One likely reason for the poor empirical performance of RLA-B is the (lack of)
relative scaling of the identity and the outer product in the metric. Hartmann
et al. [2022] and Yu et al. [2023] empirically showed that for MCMC the latter
may need to be heavily down-weighted and even then it is not always better
than Euclidean. RLA-B likely requires some form of scaling as well and
downweighting the latter term would reduce the bias, but selecting the scaling
remains an open problem.

##### Recommendations.

For small-to-medium problems RLA-F is clearly the best choice, based on
consistent best performance and stable computation. For large-scale problems,
for instance in the common use case of NNs [Daxberger et al., 2021], RLA-B
remains currently the best choice, despite the bias.

We showed how to compute the FIM for arbitrary NNs, but only experimented on
tiny ones to facilitate direct comparisons against NUTS. Our preliminary
experiments indicate that for $D$ in the range of $200$-$250$ the running
times of RLA-B and RLA-F are often comparable, with the quality of the MAP
estimate making large differences. We could naturally run RLA-F for still
larger $D$ in reasonable time, but very large models necessarily call for a
more efficient metric. A promising direction for scalable and accurate method
could build on scalable approximations for FIM [Martens and Grosse, 2015,
George et al., 2018] with practical implementations as in Botev and Martens
[2023] and George [2021], or on exact subsampled Fisher [Benzing, 2022].

#### Acknowledgements

The work was supported by the Research Council of Finland Flagship programme:
Finnish Center for Artificial Intelligence FCAI, and additionally by the
grants 336019, 345811, 348952, 324852. The authors acknowledge support from
CSC – IT Center for Science, Finland, for computational resources.

## References

  * Arvanitidis et al. [2016] G. Arvanitidis, L. K. Hansen, and S. r. Hauberg.  A Locally Adaptive Normal Distribution.  In _Advances in Neural Information Processing Systems_ , volume 29. Curran Associates, Inc., 2016.
  * Atkinson et al. [2009] K. Atkinson, W. Han, and D. Stewart.  _Numerical Solutions of Ordinary Differential Equations_.  John Wiley & Sons, 2009.
  * Benzing [2022] F. Benzing.  Gradient Descent on Neurons and its Link to Approximate Second-order Optimization.  In _Proceedings of the 39th International Conference on Machine Learning_ , volume 162 of _Proceedings of Machine Learning Research_ , pages 1817–1853. PMLR, July 2022.
  * Bergamin et al. [2023] F. Bergamin, P. Moreno-Muñoz, S. Hauberg, and G. Arvantidis.  Riemannian Laplace approximations for Bayesian neural networks.  In _Neural Information Processing Systems (NeurIPS)_ , 2023\.
  * Blei et al. [2017] D. M. Blei, A. Kucukelbir, and J. D. McAuliffe.  Variational Inference: A Review for Statisticians.  _Journal of the American Statistical Association_ , 112(518):859–877, 2017.  Publisher: Taylor & Francis.
  * Botev and Martens [2023] A. Botev and J. Martens.  KFAC-JAX, 2023.  URL https://github.com/google-deepmind/kfac-jax.
  * Boumal [2023] N. Boumal.  _An introduction to optimization on smooth manifolds_.  Cambridge University Press, 2023.
  * Bradbury et al. [2018] J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, and Q. Zhang.  JAX: composable transformations of Python+NumPy programs, 2018.  URL http://github.com/google/jax.
  * Brofos and Lederman [2021] J. A. Brofos and R. R. Lederman.  On Numerical Considerations for Riemannian Manifold Hamiltonian Monte Carlo, 2021.  URL https://arxiv.org/abs/2111.09995.
  * Calin and Udrişte [2014] O. Calin and C. Udrişte.  _Geometric Modeling in Probability and Statistics_.  Springer International Publishing, 1 edition, 2014.
  * Daxberger et al. [2021] E. Daxberger, A. Kristiadi, A. Immer, R. Eschenhagen, M. Bauer, and P. Hennig.  Laplace Redux - Effortless Bayesian Deep Learning.  In _Advances in Neural Information Processing Systems_ , volume 34, pages 20089–20103. Curran Associates, Inc., 2021.
  * De Bortoli et al. [2022] V. De Bortoli, E. Mathieu, M. Hutchinson, J. Thornton, Y. W. Teh, and A. Doucet.  Riemannian Score-Based Generative Modelling.  In _Advances in Neural Information Processing Systems_ , volume 35, pages 2406–2422. Curran Associates, Inc., 2022.
  * Dormand and Prince [1980] J. R. Dormand and P. J. Prince.  A family of embedded Runge-Kutta formulae.  _Journal of Computational and Applied Mathematics_ , 6(1):19–26, 1980.
  * Federer [1969] H. Federer.  _Geometric Measure Theory_.  Springer-Verlag, 1969.  ISBN 9780387045054,0387045058.
  * Feng et al. [1992] S. Feng, C., K. A., M. S., S., and R. Henery.  Statlog Project, 1992.  Published: UCI Machine Learning Repository.
  * Flamary et al. [2021] R. Flamary, N. Courty, A. Gramfort, M. Z. Alaya, A. Boisbunon, S. Chambon, L. Chapel, A. Corenflos, K. Fatras, N. Fournier, L. Gautheron, N. T. H. Gayraud, H. Janati, A. Rakotomamonjy, I. Redko, A. Rolet, A. Schutz, V. Seguy, D. J. Sutherland, R. Tavenard, A. Tong, and T. Vayer.  POT: Python Optimal Transport.  _Journal of Machine Learning Research_ , 22(78):1–8, 2021.
  * Foong et al. [2019] A. Y. K. Foong, Y. Li, J. M. Hernández-Lobato, and R. E. Turner.  ’In-Between’ Uncertainty in Bayesian Neural Networks, 2019\.  URL https://arxiv.org/abs/1906.11537.
  * Frank et al. [2021] P. Frank, R. Leike, and T. A. Enßlin.  Geometric Variational Inference.  _Entropy_ , 23(7), 2021.
  * George [2021] T. George.  NNGeometry: Easy and Fast Fisher Information Matrices and Neural Tangent Kernels in PyTorch, Feb. 2021.  URL https://doi.org/10.5281/zenodo.4532597.
  * George et al. [2018] T. George, C. Laurent, X. Bouthillier, N. Ballas, and P. Vincent.  Fast Approximate Natural Gradient Descent in a Kronecker Factored Eigenbasis.  In _Advances in Neural Information Processing Systems_ , volume 31. Curran Associates, Inc., 2018.
  * George Casella [2001] R. L. B. George Casella.  _Statistical Inference_.  Duxbury Press, 2° edition, 2001.  ISBN 0534243126,9780534243128.
  * Girolami and Calderhead [2011] M. Girolami and B. Calderhead.  Riemann manifold Langevin and Hamiltonian Monte Carlo methods.  _Journal of the Royal Statistical Society: Series B (Statistical Methodology)_ , 73(2):123–214, 2011.
  * Grosse [2022] R. Grosse.  Chapter 3: Metrics, 2022.  URL https://www.cs.toronto.edu/~rgrosse/courses/csc2541_2022/readings/L03_metrics.pdf.
  * Harris et al. [2020] C. R. Harris, K. J. Millman, S. J. v. d. Walt, R. Gommers, P. Virtanen, D. Cournapeau, E. Wieser, J. Taylor, S. Berg, N. J. Smith, R. Kern, M. Picus, S. Hoyer, M. H. v. Kerkwijk, M. Brett, A. Haldane, J. F. d. Río, M. Wiebe, P. Peterson, P. Gérard-Marchant, K. Sheppard, T. Reddy, W. Weckesser, H. Abbasi, C. Gohlke, and T. E. Oliphant.  Array programming with NumPy.  _Nature_ , 585(7825):357–362, Sept. 2020.  Publisher: Springer Science and Business Media LLC.
  * Hartmann and Vanhatalo [2019] M. Hartmann and J. Vanhatalo.  Laplace approximation and natural gradient for Gaussian process regression with heteroscedastic student-t model.  _Statistics and Computing_ , 29(4):753–773, July 2019.
  * Hartmann et al. [2022] M. Hartmann, M. Girolami, and A. Klami.  Lagrangian manifold Monte Carlo on Monge patches.  In _Proceedings of The 25th International Conference on Artificial Intelligence and Statistics_ , volume 151 of _Proceedings of Machine Learning Research_ , pages 4764–4781. PMLR, Mar. 2022.
  * Jermyn [2005] I. H. Jermyn.  Invariant Bayesian estimation on manifolds.  _The Annals of Statistics_ , 33:583–605, 2005.  ISSN 0090-5364.
  * Kidger [2021] P. Kidger.  _On Neural Differential Equations_.  PhD Thesis, University of Oxford, 2021.
  * Kingma and Ba [2015] D. P. Kingma and J. Ba.  Adam: A Method for Stochastic Optimization.  In _3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings_ , 2015.
  * Knollmüller and Enßlin [2020] J. Knollmüller and T. A. Enßlin.  Metric Gaussian Variational Inference, 2020.  URL https://arxiv.org/abs/1901.11033.
  * Kristiadi et al. [2023] A. Kristiadi, F. Dangel, and P. Hennig.  The Geometry of Neural Nets’ Parameter Spaces Under Reparametrization, 2023.  URL https://arxiv.org/abs/2302.07384.
  * Kunstner et al. [2019] F. Kunstner, P. Hennig, and L. Balles.  Limitations of the empirical Fisher approximation for natural gradient descent.  In _Advances in Neural Information Processing Systems_ , volume 32. Curran Associates, Inc., 2019.
  * Lan et al. [2015] S. Lan, V. Stathopoulos, B. Shahbaba, and M. Girolami.  Markov Chain Monte Carlo From Lagrangian Dynamics.  _Journal of Computational and Graphical Statistics_ , 24(2):357–378, 2015.
  * Lee, John M. [2018] Lee, John M.  _Introduction to Riemannian Manifolds_.  Graduate Texts in Mathematics. Springer International Publishing AG, Cham, Switzerland, 2nd edition, 2018.
  * Lehmann [2003] G. C. Lehmann.  _Theory of Point Estimation_.  Springer texts in statistics. Springer, 2nd ed edition, 2003.
  * Magnusson et al. [2022] M. Magnusson, P. Bürkner, and A. Vehtari.  posteriordb: a set of posteriors for Bayesian inference and probabilistic programming, Nov. 2022.  URL https://github.com/stan-dev/posteriordb.
  * Manfredo [1992] d. C. Manfredo.  _Riemannian Geometry_.  Birkhäuser, Boston, 1992.
  * Margossian and Blei [2023] C. C. Margossian and D. M. Blei.  Amortized Variational Inference: When and Why?, 2023.  URL https://arxiv.org/abs/2307.11018.
  * Martens [2020] J. Martens.  New Insights and Perspectives on the Natural Gradient Method.  _Journal of Machine Learning Research_ , 21(146):1–76, 2020.
  * Martens and Grosse [2015] J. Martens and R. Grosse.  Optimizing Neural Networks with Kronecker-factored Approximate Curvature.  In _Proceedings of the 32nd International Conference on Machine Learning_ , volume 37 of _Proceedings of Machine Learning Research_ , pages 2408–2417, Lille, France, July 2015. PMLR.
  * Meurer et al. [2017] A. Meurer, C. P. Smith, M. Paprocki, O. Čertík, S. B. Kirpichev, M. Rocklin, A. Kumar, S. Ivanov, J. K. Moore, S. Singh, T. Rathnayake, S. Vig, B. E. Granger, R. P. Muller, F. Bonazzi, H. Gupta, S. Vats, F. Johansson, F. Pedregosa, M. J. Curry, A. R. Terrel, v. Roučka, A. Saboo, I. Fernando, S. Kulal, R. Cimrman, and A. Scopatz.  SymPy: symbolic computing in Python.  _PeerJ Computer Science_ , 3:e103, Jan. 2017.  ISSN 2376-5992.  doi: 10.7717/peerj-cs.103.  URL https://doi.org/10.7717/peerj-cs.103.
  * Minka [2001] T. P. Minka.  Expectation propagation for approximate bayesian inference.  In _Proceedings of the 17th Conference in Uncertainty in Artificial Intelligence_ , UAI ’01, page 362–369. Morgan Kaufmann Publishers Inc., 2001.
  * Neal [2003] R. M. Neal.  Slice sampling.  _The Annals of Statistics_ , 31(3):705 – 767, 2003.  Publisher: Institute of Mathematical Statistics.
  * OpenAI [2023] OpenAI.  ChatGPT, 2023.  URL https://chat.openai.com/.
  * Papamakarios et al. [2021] G. Papamakarios, E. Nalisnick, D. J. Rezende, S. Mohamed, and B. Lakshminarayanan.  Normalizing Flows for Probabilistic Modeling and Inference.  _Journal of Machine Learning Research_ , 22(57):1–64, 2021.
  * Paszke et al. [2019] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala.  Pytorch: An imperative style, high-performance deep learning library.  In _Advances in Neural Information Processing Systems 32_ , pages 8024–8035. Curran Associates, Inc., 2019.
  * Patterson and Teh [2013] S. Patterson and Y. W. Teh.  Stochastic Gradient Riemannian Langevin Dynamics on the Probability Simplex.  In _Advances in Neural Information Processing Systems_ , volume 26. Curran Associates, Inc., 2013.
  * Petersen and Pedersen [2012] K. B. Petersen and M. S. Pedersen.  The Matrix Cookbook, Nov. 2012.  URL http://www2.compute.dtu.dk/pubdb/pubs/3274-full.html.
  * Rasmussen and Williams [2006] C. E. Rasmussen and C. K. I. Williams.  _Gaussian processes for machine learning_.  Adaptive computation and machine learning. MIT Press, 2006.
  * Ripley [1994] B. D. Ripley.  Neural networks and related methods for classification (with discussion).  _Journal of the Royal Statistical Society series B_ , pages 409–456, 1994.
  * Ripley [1996] B. D. Ripley.  _Pattern Recognition and Neural Networks_.  Cambridge University Press, 1996.
  * Rue et al. [2009] H. Rue, S. Martino, and N. Chopin.  Approximate bayesian inference for latent gaussian models by using integrated nested laplace approximations.  _Journal Of The Royal Statistical Society_ , 71:319–392, 2009.
  * Schervish [2012] M. Schervish.  _Theory of Statistics_.  Springer Series in Statistics. Springer New York, 2012.
  * Smith et al. [1988] J. W. Smith, J. E. Everhart, W. C. Dickson, W. C. Knowler, and R. S. Johannes.  Using the ADAP Learning Algorithm to Forecast the Onset of Diabetes Mellitus.  Orlando, 1988.
  * Snelson and Ghahramani [2005] E. Snelson and Z. Ghahramani.  Sparse Gaussian Processes using Pseudo-inputs.  In _Advances in Neural Information Processing Systems_ , volume 18. MIT Press, 2005.
  * Song et al. [2018] Y. Song, J. Song, and S. Ermon.  Accelerating Natural Gradient with Higher-Order Invariance.  In _Proceedings of the 35th International Conference on Machine Learning_ , volume 80 of _Proceedings of Machine Learning Research_ , pages 4713–4722. PMLR, July 2018.
  * Stan Dev Team [2023] Stan Dev Team.  CmdStanPy, 2023.  URL https://pypi.org/project/cmdstanpy/.
  * Stan Development Team [2023] Stan Development Team.  Stan Modeling Language Users Guide and Reference Manual, 2023\.  URL https://mc-stan.org.
  * Tierney and Kadane [1986] L. Tierney and J. B. Kadane.  Accurate approximations for posterior moments and marginal densities.  _Journal of the American Statistical Association_ , 81(393):82–86, 1986.
  * van der Vaart [1998] A. van der Vaart.  _Asymptotic Statistics_.  Cambridge University Press, Cambridge, UK, 1998.
  * Virtanen et al. [2020] P. Virtanen, R. Gommers, T. E. Oliphant, M. Haberland, T. Reddy, D. Cournapeau, E. Burovski, P. Peterson, W. Weckesser, J. Bright, S. J. van der Walt, M. Brett, J. Wilson, K. J. Millman, N. Mayorov, A. R. J. Nelson, E. Jones, R. Kern, E. Larson, C. J. Carey, I. Polat, Y. Feng, E. W. Moore, J. VanderPlas, D. Laxalde, J. Perktold, R. Cimrman, I. Henriksen, E. A. Quintero, C. R. Harris, A. M. Archibald, A. H. Ribeiro, F. Pedregosa, P. van Mulbregt, and SciPy 1.0 Contributors.  SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python.  _Nature Methods_ , 17:261–272, 2020.
  * Yang and Berger [1996] R. Yang and J. Berger.  _A Catalog of Noninformative Priors_.  Discussion papers. Institute of Statistics and Decision Sciences, Duke University, 1996.
  * Yu et al. [2023] H. Yu, M. Hartmann, B. Williams, and A. Klami.  Scalable Stochastic Gradient Riemannian Langevin Dynamics in Non-Diagonal Metrics.  _Transactions on Machine Learning Research_ , 2023.
  * Zhang et al. [2022] L. Zhang, B. Carpenter, A. Gelman, and A. Vehtari.  Pathfinder: Parallel quasi-Newton variational inference.  _Journal of Machine Learning Research_ , 23(306):1–49, 2022.

##  Appendix A Proofs

Here we provide the proofs for the Theorems stated in the main paper.

######  Theorem 1.

For Gaussian (or uniform) prior and Gaussian likelihood with fixed covariance,
the Fisher metric is constant.

###### Proof.

Consider the prior distribution on the parameters of interest. We consider two
cases, the first where the prior is Gaussian with a fixed covariance and the
second where the prior is uniform. For the first case, we have
$\operatorname{\boldsymbol{\theta}}\sim\operatorname{\mathcal{N}}(\operatorname{\boldsymbol{\mu}},\operatorname{\boldsymbol{S}})$
for some $\operatorname{\boldsymbol{\mu}}$ and
$\operatorname{\boldsymbol{S}}$. Therefore,

$\log\pi(\operatorname{\boldsymbol{\theta}})=-\frac{k}{2}\log(2\pi)-\frac{1}{2}\log\det(\operatorname{\boldsymbol{S}})-\frac{1}{2}(\operatorname{\boldsymbol{\theta}}-\operatorname{\boldsymbol{\mu}})^{\top}\operatorname{\boldsymbol{S}}^{-1}(\operatorname{\boldsymbol{\theta}}-\operatorname{\boldsymbol{\mu}}),$

and

$\partialderivative[2]{\log\pi(\operatorname{\boldsymbol{\theta}})}{\operatorname{\boldsymbol{\theta}}}=-\operatorname{\boldsymbol{S}}^{-1}.$

The negative Hessian is therefore constant.

For the second case, where the prior is uniform, the negative Hessian is $0$
everywhere, thus is also constant.

Consider the Gaussian distribution
$\operatorname{\boldsymbol{Y}}|\operatorname{\boldsymbol{\theta}}\sim\operatorname{\mathcal{N}}(\operatorname{\boldsymbol{\theta}},\operatorname{\boldsymbol{\Sigma}})$
for some fixed $\operatorname{\boldsymbol{\Sigma}}$. We have

$\partialderivative[2]{\log\pi(\operatorname{\boldsymbol{y}}|\operatorname{\boldsymbol{\theta}})}{\operatorname{\boldsymbol{\theta}}}=-\operatorname{\boldsymbol{\Sigma}}^{-1},$

which does not depend on $\operatorname{\boldsymbol{y}}$, thus the FIM is
constant.

Therefore, the resulting Fisher metric is constant. ∎

######  Theorem 2.

Under the conditions of Bernstein-von Mises theorem, the Fisher metric
converges to a fixed matrix at the limit of infinite data.

###### Proof.

At the limit of infinite data, up to the regularity conditions characterised
by the Bernstein-von Mises theorem [van der Vaart, 1998], the likelihood
dominates the prior and the posterior converges in distribution to a Gaussian
given by
$\operatorname{\mathcal{N}}(\operatorname{\boldsymbol{\theta}}_{\text{MLE}},(NI(\operatorname{\boldsymbol{\theta}}_{\text{MLE}}))^{-1})$,
where $\operatorname{\boldsymbol{\theta}}_{\text{MLE}}$ denotes the maximum
likelihood estimate. In other words, with the same conditions, the prior’s
contribution converges to $0$, and the likelihood converges to a Gaussian
function, up to a scaling, given by
$\operatorname{\mathcal{N}}(\operatorname{\boldsymbol{\theta}}_{\text{MLE}}|\operatorname{\boldsymbol{\theta}},(NI(\operatorname{\boldsymbol{\theta}}_{\text{MLE}}))^{-1})$.

Since the FIM of the Gaussian model is given by
$NI(\operatorname{\boldsymbol{\theta}}_{\text{MLE}})$, which is constant for
all $\operatorname{\boldsymbol{\theta}}$, the FIM of the limiting model is
also constant, and the resulting convergence limit of the metric of the
current model is also constant. ∎

######  Theorem 3.

With an invariant prior, e.g. Jeffreys prior, RLA-F with Hausdorff MAP is
exact for probabilistic models whose target distributions are diffeomorphic
with Gaussians, for which the negative Hessian at the Hausdorff MAP coincides
with the Fisher metric.

###### Proof.

Consider first the simple case with no transformations, such that the
probabilistic model reduces to the following form

$\displaystyle\pi(\operatorname{\boldsymbol{\mu}}|\operatorname{\boldsymbol{\psi}})$
$\displaystyle=\operatorname{\mathcal{N}}(\operatorname{\boldsymbol{\mu}}|\operatorname{\boldsymbol{\psi}},\operatorname{\boldsymbol{S}}),$
$\displaystyle\pi(\operatorname{\boldsymbol{\psi}})$
$\displaystyle=\text{Jeffreys}(\operatorname{\boldsymbol{\psi}}).$

The FIM for $\operatorname{\boldsymbol{\psi}}$ is constant, given by
$\operatorname{\boldsymbol{S}}^{-1}$. Therefore, The prior corresponds to a
uniform distribution in the Euclidean sense, and the posterior is given by

$\pi(\operatorname{\boldsymbol{\psi}}|\operatorname{\boldsymbol{\mu}})=\operatorname{\mathcal{N}}(\operatorname{\boldsymbol{\psi}}|\operatorname{\boldsymbol{\mu}},\operatorname{\boldsymbol{S}}).$

It is clear that RLA-F is exact for this posterior.

Consider now a diffeomorphic transformation
$\operatorname{\boldsymbol{\phi}}:\operatorname{\boldsymbol{\psi}}\rightarrow\operatorname{\boldsymbol{\theta}}$.
It is well known that Fisher Information Matrix automatically transforms
according to the corrrect transformation rule of a Riemannian metric [Martens,
2020, Kristiadi et al., 2023]. The distribution of the transformed parameters
is given by
$\pi(\operatorname{\boldsymbol{\theta}}|\operatorname{\boldsymbol{\mu}})=\pi(\operatorname{\boldsymbol{\phi}}^{-1}(\operatorname{\boldsymbol{\theta}})|\operatorname{\boldsymbol{\mu}})\left|\det\partialderivative{\operatorname{\boldsymbol{\phi}}^{-1}}{\operatorname{\boldsymbol{\theta}}}\right|$.
It coincides with the posterior of the probabilistic model

$\displaystyle\pi(\operatorname{\boldsymbol{\mu}}|\operatorname{\boldsymbol{\theta}})$
$\displaystyle=\operatorname{\mathcal{N}}(\operatorname{\boldsymbol{\mu}}|\operatorname{\boldsymbol{\phi}}^{-1}(\operatorname{\boldsymbol{\theta}}),\operatorname{\boldsymbol{S}}),$
$\displaystyle\pi_{J}(\operatorname{\boldsymbol{\theta}})$
$\displaystyle=\sqrt{\det\operatorname{\boldsymbol{G}}_{\operatorname{\boldsymbol{\Theta}}}(\operatorname{\boldsymbol{\theta}})},$

since Jeffreys prior accounts for the change of variable of the
transformation,

$\displaystyle\pi_{J}(\operatorname{\boldsymbol{\theta}})=\sqrt{\det\operatorname{\boldsymbol{G}}_{\operatorname{\boldsymbol{\Theta}}}(\operatorname{\boldsymbol{\theta}})}=\left|\det\partialderivative{\operatorname{\boldsymbol{\phi}}^{-1}}{\operatorname{\boldsymbol{\theta}}}\right|\sqrt{\det\operatorname{\boldsymbol{G}}_{\operatorname{\boldsymbol{\Psi}}}(\operatorname{\boldsymbol{\phi}}^{-1}(\operatorname{\boldsymbol{\theta}}))}.$

The parameters that give the maximum value under the Hausdorff measure
${\pi(\operatorname{\boldsymbol{\theta}}|\operatorname{\boldsymbol{\mu}})}(\det\operatorname{\boldsymbol{G}}_{\operatorname{\boldsymbol{\Theta}}}(\operatorname{\boldsymbol{\theta}}))^{-\tfrac{1}{2}}\propto\pi(\operatorname{\boldsymbol{\mu}}|\operatorname{\boldsymbol{\phi}}^{-1}(\operatorname{\boldsymbol{\theta}}))$
are always given by the MLE estimate
$\hat{\operatorname{\boldsymbol{\theta}}}=\operatorname{\boldsymbol{\phi}}(\operatorname{\boldsymbol{\mu}})$.

The argument for exactness is as follows. Consider the two manifolds
$\operatorname{\mathcal{M}}:(\operatorname{\boldsymbol{\Psi}},\operatorname{\boldsymbol{G}}_{\operatorname{\boldsymbol{\Psi}}})$
and
$\operatorname{\mathcal{N}}:(\operatorname{\boldsymbol{\Theta}},\operatorname{\boldsymbol{G}}_{\operatorname{\boldsymbol{\Theta}}})$.
Since the corresponding metrics satisfy the transformation rule

$\operatorname{\boldsymbol{G}}_{\operatorname{\boldsymbol{\Theta}}}=\left(\partialderivative{\operatorname{\boldsymbol{\psi}}}{\operatorname{\boldsymbol{\theta}}}\right)^{\top}\operatorname{\boldsymbol{G}}_{\operatorname{\boldsymbol{\Psi}}}\partialderivative{\operatorname{\boldsymbol{\psi}}}{\operatorname{\boldsymbol{\theta}}},$
(4)

there exists an isomorphism between the tangent spaces of
$\operatorname{\mathcal{M}}$ and $\operatorname{\mathcal{N}}$. As a result,
the exponential map on $\operatorname{\mathcal{M}}$ and
$\operatorname{\mathcal{N}}$ and the tangent vectors at the center of the
distribution transform correctly. Since RLA-F is exact for the distribution
corresponding to $\operatorname{\mathcal{M}}$, it is also exact for the
distribution corresponding to $\operatorname{\mathcal{N}}$.

Moreover, the negative Hessian at the Hausdorff MAP coincides with the Fisher
metric. This can be seen from the transformation rule of Hessian matrix under
a transformation given by
$\operatorname{\boldsymbol{\theta}}=\operatorname{\boldsymbol{\phi}}(\operatorname{\boldsymbol{\psi}})$

$\nabla_{\operatorname{\boldsymbol{\theta}}}^{2}f=\left(\partialderivative{\operatorname{\boldsymbol{\psi}}}{\operatorname{\boldsymbol{\theta}}}\right)^{\top}\nabla_{\operatorname{\boldsymbol{\psi}}}^{2}f\left(\partialderivative{\operatorname{\boldsymbol{\psi}}}{\operatorname{\boldsymbol{\theta}}}\right)+[\nabla_{\operatorname{\boldsymbol{\theta}}}^{2}\operatorname{\boldsymbol{\phi}}^{-1}_{1}(\operatorname{\boldsymbol{\theta}}),\dots,\nabla_{\operatorname{\boldsymbol{\theta}}}^{2}\operatorname{\boldsymbol{\phi}}^{-1}_{n}(\operatorname{\boldsymbol{\theta}})]\left(\nabla_{\operatorname{\boldsymbol{\psi}}}f\otimes\operatorname{\boldsymbol{I}}_{n}\right),$

because for the MAP estimate, $\nabla_{\operatorname{\boldsymbol{\psi}}}f=0$.
∎

##  Appendix B Using Hausdorff MAP and Fisher precision

In Section 4.3.1, we briefly discussed the choices of using Hausdorff MAP
together with Fisher precision. Here we provide additional remarks on this.

We use notations inspired by Kristiadi et al. [2023]. Consider a
reparametrization
$\operatorname{\boldsymbol{\phi}}:\operatorname{\boldsymbol{\psi}}\rightarrow\operatorname{\boldsymbol{\theta}}$.
Denote its Jacobian as $\operatorname{\boldsymbol{J}}$, with

$\operatorname{\boldsymbol{J}}_{ij}=\frac{\partial\operatorname{\boldsymbol{\theta}}_{i}}{\partial\operatorname{\boldsymbol{\psi}}_{j}}.$

As noted by Kristiadi et al. [2023], the transformation rule of vector
components is

$\operatorname{\boldsymbol{v}}_{\operatorname{\boldsymbol{\Theta}}}=\operatorname{\boldsymbol{J}}\operatorname{\boldsymbol{v}}_{\operatorname{\boldsymbol{\Psi}}},$

and the transformation rule of covector components is

$\operatorname{\boldsymbol{w}}_{\operatorname{\boldsymbol{\Theta}}}=\operatorname{\boldsymbol{J}}^{-\top}\operatorname{\boldsymbol{w}}_{\operatorname{\boldsymbol{\Psi}}}.$

Recall that the FIM transforms as a Riemannian metric under one-to-one
mappings [Martens, 2020, Kristiadi et al., 2023]. Therefore, with an invariant
prior or ignoring the effect of the prior, the resulting metric transforms
automatically.

Typical MAP estimates are not invariant under reparametrization. However, we
can define an invariant MAP estimate by taking into account the Riemannian
structure [Kristiadi et al., 2023]. Specifically, instead of calculating the
maximum of the PDF under Lebesgue measure, we calculate it under the Hausdorff
measure. Denote the PDF under Lebesgue measure as
$\pi_{\operatorname{\boldsymbol{\Theta}}}(\operatorname{\boldsymbol{\theta}})$,
and the PDF under Hausdorff measure as
$\pi_{\operatorname{\boldsymbol{\Theta}}}^{\operatorname{\boldsymbol{G}}}(\operatorname{\boldsymbol{\theta}})$.
Then, using the Riemannian volume form [Lee, John M., 2018], we have

$\pi_{\operatorname{\boldsymbol{\Theta}}}^{\operatorname{\boldsymbol{G}}}(\operatorname{\boldsymbol{\theta}})=\frac{\pi_{\operatorname{\boldsymbol{\Theta}}}(\operatorname{\boldsymbol{\theta}})}{\sqrt{\det\operatorname{\boldsymbol{G}}_{\operatorname{\boldsymbol{\Theta}}}(\operatorname{\boldsymbol{\theta}})}}.$

######  Theorem 4.

The PDF under Hausdorff measure is invariant across reparametrization.

###### Proof.

$\displaystyle\pi_{\operatorname{\boldsymbol{\Theta}}}^{\operatorname{\boldsymbol{G}}}(\operatorname{\boldsymbol{\theta}})$
$\displaystyle=\frac{\pi_{\operatorname{\boldsymbol{\Psi}}}(\operatorname{\boldsymbol{\phi}}^{-1}(\operatorname{\boldsymbol{\theta}}))\sqrt{\det\operatorname{\boldsymbol{J}}^{-1}(\operatorname{\boldsymbol{\theta}})}}{\sqrt{\det(\operatorname{\boldsymbol{J}}^{-T}(\operatorname{\boldsymbol{\theta}})\operatorname{\boldsymbol{G}}_{\operatorname{\boldsymbol{\Psi}}}(\operatorname{\boldsymbol{\phi}}^{-1}(\operatorname{\boldsymbol{\theta}}))\operatorname{\boldsymbol{J}}^{-1}(\operatorname{\boldsymbol{\theta}}))}}$
$\displaystyle=\frac{\pi_{\operatorname{\boldsymbol{\Psi}}}(\operatorname{\boldsymbol{\phi}}^{-1}(\operatorname{\boldsymbol{\theta}}))}{\sqrt{\det\operatorname{\boldsymbol{G}}_{\Psi}(\operatorname{\boldsymbol{\phi}}^{-1}(\operatorname{\boldsymbol{\theta}}))}}=\frac{\pi_{\operatorname{\boldsymbol{\Theta}}}(\operatorname{\boldsymbol{\theta}})}{\sqrt{\det\operatorname{\boldsymbol{G}}_{\operatorname{\boldsymbol{\Theta}}}(\operatorname{\boldsymbol{\theta}})}}.$

∎

This naturally leads to the following corollary

######  Corollary 1.

The Hausdorff MAP is invariant across reparameterization.

The following theorem shows that we can ignore the first order gradients in
the Taylor series expansions as in Equation 2 of the main paper

######  Theorem 5.

Under normal coordinates, the Hausdorff MAP is a critical point under Lebesgue
measure.

###### Proof.

It is a standard result in Riemannian geometry that for a Riemannian manifold,
at each tangent space there exists a normal coordinate [Lee, John M., 2018],
such that the components of the metric at the point form an identity matrix,
and all first partial derivatives of the metric vanishes at the point [see
Proposition 5.24 from Lee, John M., 2018]. Therefore, consider such a
transformation to the normal coordinate
$\operatorname{\boldsymbol{\phi}}:\operatorname{\boldsymbol{\psi}}\rightarrow\operatorname{\boldsymbol{\theta}}$.
$\hat{\operatorname{\boldsymbol{\theta}}}=\operatorname{\boldsymbol{\phi}}(\hat{\operatorname{\boldsymbol{\psi}}})$
is the Hausdorff MAP due to Theorem 1. Recall the definition of Hausdorff MAP;
it is clear that
$\partialderivative{\pi^{\operatorname{\boldsymbol{G}}}(\hat{\operatorname{\boldsymbol{\theta}}})}{\theta_{i}}=0$.

Moreover, since the first partial derivatives of the metric vanishes under a
normal coordinate, recall
$\partial(\det(\operatorname{\boldsymbol{X}}))=\det(\operatorname{\boldsymbol{X}})\Tr(\operatorname{\boldsymbol{X}}^{-1}\partial\operatorname{\boldsymbol{X}})$
[Petersen and Pedersen, 2012], we have
$\partialderivative{\sqrt{\det\operatorname{\boldsymbol{G}}(\hat{\operatorname{\boldsymbol{\theta}}})}}{\theta_{i}}=0.$
Therefore,

$\partialderivative{\pi(\hat{\operatorname{\boldsymbol{\theta}}})}{\theta_{i}}=\partialderivative{\pi^{\operatorname{\boldsymbol{G}}}(\hat{\operatorname{\boldsymbol{\theta}}})\sqrt{\det\operatorname{\boldsymbol{G}}(\hat{\operatorname{\boldsymbol{\theta}}})}}{\theta_{i}}=\partialderivative{\pi^{\operatorname{\boldsymbol{G}}}(\hat{\operatorname{\boldsymbol{\theta}}})}{\theta_{i}}\sqrt{\det\operatorname{\boldsymbol{G}}(\hat{\operatorname{\boldsymbol{\theta}}})}+\pi^{\operatorname{\boldsymbol{G}}}(\hat{\operatorname{\boldsymbol{\theta}}})\partialderivative{\sqrt{\det\operatorname{\boldsymbol{G}}(\hat{\operatorname{\boldsymbol{\theta}}})}}{\theta_{i}}=0.$

∎

Under an invariant prior or ignoring the effect of prior, the Fisher metric
transforms as a Riemannian metric, and RLA-F with Hausdorff MAP and Fisher
precision can be interpreted as forming the approximate Gaussian with Fisher
metric in place of the negative Hessian at
$\hat{\operatorname{\boldsymbol{\theta}}}$ under normal coordinates and
transforming the approximation back to the current approximation.

The following theorem justifies using the Fisher metric as the precision of
the Gaussian approximation from a differential geometry viewpoint.

######  Theorem 6.

When the precision of a Gaussian distribution is a Riemannian metric, the
samples from the Gaussian follow the transformation rule of the tangent
vectors.

###### Proof.

We analyze the transformation property of samples from the resulting Gaussian
distribution.

The transformation rule of the components of a Riemannian metric, which is a
$(0,2)$ tensor, follow

$\operatorname{\boldsymbol{G}}_{\operatorname{\boldsymbol{\Theta}}}=\operatorname{\boldsymbol{J}}^{-\top}\operatorname{\boldsymbol{G}}_{\operatorname{\boldsymbol{\Psi}}}\operatorname{\boldsymbol{J}}^{-1}.$

The transformation rule of the inverse of the Riemannian metric follows

$\operatorname{\boldsymbol{G}}^{-1}_{\operatorname{\boldsymbol{\Theta}}}=\operatorname{\boldsymbol{J}}\operatorname{\boldsymbol{G}}^{-1}_{\operatorname{\boldsymbol{\Psi}}}\operatorname{\boldsymbol{J}}^{\top}=(\operatorname{\boldsymbol{J}}\operatorname{\boldsymbol{L}}_{\operatorname{\boldsymbol{\Psi}}})(\operatorname{\boldsymbol{J}}\operatorname{\boldsymbol{L}}_{\operatorname{\boldsymbol{\Psi}}})^{\top},$

where $\operatorname{\boldsymbol{L}}$ indicates the decomposition of the
inverse of the metric. We therefore have

$\operatorname{\boldsymbol{L}}_{\operatorname{\boldsymbol{\Theta}}}=\operatorname{\boldsymbol{J}}\operatorname{\boldsymbol{L}}_{\operatorname{\boldsymbol{\Psi}}},$

and the samples from the resulting Gaussian distribution transform in the same
way as vector components. ∎

##  Appendix C Coverage of RLA samples

An interesting insight is that, under mild assumptions, we can obtain samples
over the entire $\operatorname{\mathbb{R}}^{D}$ space for any metric and any
starting point.

######  Theorem 7.

For target distributions of unconstrained parameters, if the induced
Riemannian manifold $(\operatorname{\boldsymbol{\Theta}},g)$ is connected,
there always exists one length-minimizing geodesic starting from the MAP and
passing through an arbitrary point $\operatorname{\boldsymbol{\theta}}$ on the
induced manifold.

###### Proof.

Under unconstrained parametrization, for a point $p$ on the connected
Riemannian manifold $\operatorname{\boldsymbol{\Theta}}$, the geodesics are
defined on the entire tangent space. Therefore, using Lemma 6.18 from [Lee,
John M., 2018], for any two points
$\operatorname{\boldsymbol{p}},\operatorname{\boldsymbol{q}}\in\operatorname{\boldsymbol{\Theta}}$,
there exists one length-minimizing geodesic. ∎

##  Appendix D Geodesic ODE

In this section, we provide some background on numerical solutions of ODE,
while also demonstrating the formulations of exponential map and logarithmic
map as ODE problems. Unless otherwise stated, the technical contents are based
on Atkinson et al. [2009].

Generally, an ODE system can be written in the form of

$\derivative{\operatorname{\boldsymbol{x}}(t)}{t}=\operatorname{\boldsymbol{f}}(t,\operatorname{\boldsymbol{x}}(t)),$

where $\operatorname{\boldsymbol{x}}$ is an unknown function dependent on time
$t$, and $\operatorname{\boldsymbol{f}}$ is a function describing the change
in $\operatorname{\boldsymbol{x}}$.

The geodesic equation induces the following ODE system

$\derivative{[\operatorname{\boldsymbol{\theta}}(t),\operatorname{\boldsymbol{v}}(t)]}{t}=[\operatorname{\boldsymbol{v}}(t),\operatorname{\boldsymbol{a}}(t)],$

where $\operatorname{\boldsymbol{a}}(t)$ is given by the geodesic equation,
such that

$\operatorname{\boldsymbol{a}}^{k}(t)=-\operatorname{\boldsymbol{v}}^{i}(t)\operatorname{\boldsymbol{v}}^{j}(t)\operatorname{\boldsymbol{\Gamma}}^{k}_{ij}(t).$

It is a nonlinear second order ODE.

An initial value problem (IVP) solves for $\operatorname{\boldsymbol{x}}(b)$
given the following conditions

$\displaystyle\derivative{\operatorname{\boldsymbol{x}}(t)}{t}$
$\displaystyle=\operatorname{\boldsymbol{f}}(t,\operatorname{\boldsymbol{x}}(t)),\quad
t_{0}\leq t\leq b,$ $\displaystyle\operatorname{\boldsymbol{x}}(t_{0})$
$\displaystyle=\operatorname{\boldsymbol{x}}_{0}.$

Probably the simplest solver is Euler’s method, which divides the times into a
discrete set of nodes $t_{0}<t_{1}<t_{2}<\dots<t_{N}\leq b$, and iteratively
solves it as

$\operatorname{\boldsymbol{x}}(n+1)=\operatorname{\boldsymbol{x}}(n)+h\operatorname{\boldsymbol{f}}(t_{n},\operatorname{\boldsymbol{x}}_{n}).$

However, this is a naive method, and often leads to large integration errors.
In practice, one commonly used solver is Dormand-Prince’s $5(4)$ method
[Dormand and Prince, 1980]; this is the default option in Virtanen et al.
[2020]. It generally requires $6$ function evaluations per step.

A general nonlinear two-point boundary value problem (BVP) can be formulated
as

$\displaystyle\operatorname{\boldsymbol{x}}^{\prime\prime}(t)$
$\displaystyle=\operatorname{\boldsymbol{f}}(t,\operatorname{\boldsymbol{x}}(t),\operatorname{\boldsymbol{x}}^{\prime}(t)),\quad
a<t<b,$
$\displaystyle\operatorname{\boldsymbol{A}}\begin{bmatrix}\operatorname{\boldsymbol{x}}(a)\\\
\operatorname{\boldsymbol{x}}^{\prime}(a)\end{bmatrix}$
$\displaystyle+\operatorname{\boldsymbol{B}}\begin{bmatrix}\operatorname{\boldsymbol{x}}(b)\\\
\operatorname{\boldsymbol{x}}^{\prime}(b)\end{bmatrix}=\begin{bmatrix}\operatorname{\boldsymbol{\gamma}}_{1}\\\
\operatorname{\boldsymbol{\gamma}}_{2}\end{bmatrix},$

where $\operatorname{\boldsymbol{A}}$ and $\operatorname{\boldsymbol{B}}$ are
square matrices, and the conditions on the second line are known as the
boundary conditions.

The logarithmic map of the geodesic ODE can be formulated as a BVP, by
observing that we can make the following choices

$\operatorname{\boldsymbol{A}}=\begin{bmatrix}\operatorname{\boldsymbol{I}}&0\\\
0&0\end{bmatrix},\;\operatorname{\boldsymbol{B}}=\begin{bmatrix}0&0\\\
\operatorname{\boldsymbol{I}}&0\end{bmatrix},\;\operatorname{\boldsymbol{\gamma}}_{1}=\operatorname{\boldsymbol{x}}_{a},\;\operatorname{\boldsymbol{\gamma}}_{2}=\operatorname{\boldsymbol{x}}_{b},$

and can thus be solved by a BVP solver.

##  Appendix E Christoffel symbols based on expression for Riemannian metric

With modern autodiff frameworks, it is possible to directly calculate the
Christoffel symbols for a given $\operatorname{\boldsymbol{\theta}}$, given
that we have a differentiable expression of the Riemannian metric. As noted by
Song et al. [2018], we only need the quantities in the form of

$\sum_{i,j}\operatorname{\boldsymbol{\Gamma}}^{k}_{ij}\operatorname{\boldsymbol{v}}^{i}\operatorname{\boldsymbol{v}}^{j}.$

We provide a JAX [Bradbury et al., 2018] implementation to obtain the above
quantities as below. 111In all presented experiments jax.numpy.linalg.solve is
used instead of jax.scipy.linalg.solve. However, the implementation in
jax.scipy could make use of the positive definite structure of the matrix,
theoretically leading to faster inversions, and we thus present this
implementation.

⬇

import jax

import jax.numpy as jnp

import jax.scipy as jsp

def christoffel_fn(g, theta, v):

# adapted based on ChatGPT

d_g = jax.jacfwd(g)(theta)

# Compute the Christoffel symbols

partial_1 = jnp.einsum(”jli,i,j->l”, d_g, v, v)

partial_2 = jnp.einsum(”ilj,i,j->l”, d_g, v, v)

partial_3 = jnp.einsum(”ijl,i,j->l”, d_g, v, v)

result = jsp.linalg.solve(g(theta), 0.5 * (partial_1 + partial_2 - partial_3),

assume_a=”pos”)

return result

In the above code, theta and v are the current position and velocity, and g is
a differentiable function that returns the Riemannian metric given the
position. It is possible to analytically derive the Christoffel symbols given
the expression of the metric. We observe that using analytical expressions can
lead to more efficient and numerically stable results. We show some analytical
derivations in Section F. However, we observed that for small scale problems,
using the numerical results are often fast and accurate enough.

##  Appendix F Fisher metric and Christoffel symbols for specific models

###  F.1 Logistic regression

In Section 5.3, we considered Bayesian logistic regression. Here we present
some derivations on the Christoffel symbols of Fisher metric for that problem.

For logistic regression, the partial derivatives of the Fisher metric are
given by [Girolami and Calderhead, 2011]

$\partialderivative{\operatorname{\boldsymbol{G}}(\operatorname{\boldsymbol{\theta}})}{\theta_{i}}=\operatorname{\boldsymbol{X}}^{\top}\operatorname{\boldsymbol{\Lambda}}\operatorname{\boldsymbol{V}}^{i}\operatorname{\boldsymbol{X}},$

where, using $s_{k}$ to denote
$s(\operatorname{\boldsymbol{X}}_{k}\operatorname{\boldsymbol{\theta}})$, we
have $\operatorname{\boldsymbol{\Lambda}}_{kk}=s_{k}(1-s_{k})$ and
$\operatorname{\boldsymbol{V}}^{i}_{kk}=(1-2s_{k})\operatorname{\boldsymbol{X}}_{ni}$.

One interesting property is the following [Lan et al., 2015]

$\partialderivative{(\operatorname{\boldsymbol{G}}(\operatorname{\boldsymbol{\theta}}))_{jl}}{\theta_{i}}=\partialderivative{(\operatorname{\boldsymbol{G}}(\operatorname{\boldsymbol{\theta}}))_{il}}{\theta_{j}}=\partialderivative{(\operatorname{\boldsymbol{G}}(\operatorname{\boldsymbol{\theta}}))_{ij}}{\theta_{l}}.$

In order to see that, note that

$\displaystyle\partialderivative{(\operatorname{\boldsymbol{G}}(\operatorname{\boldsymbol{\theta}}))_{jl}}{\theta_{i}}$
$\displaystyle=(\operatorname{\boldsymbol{X}}^{\top}\operatorname{\boldsymbol{\Lambda}}\operatorname{\boldsymbol{V}}^{i}\operatorname{\boldsymbol{X}})_{jl}=\sum_{k}\operatorname{\boldsymbol{X}}_{kj}(\operatorname{\boldsymbol{\Lambda}}\operatorname{\boldsymbol{V}}^{i})_{kk}\operatorname{\boldsymbol{X}}_{kl}=\sum_{k}\operatorname{\boldsymbol{X}}_{kj}s_{k}(1-s_{k})(1-2s_{k})\operatorname{\boldsymbol{X}}_{ki}\operatorname{\boldsymbol{X}}_{kl}$
$\displaystyle=\partialderivative{(\operatorname{\boldsymbol{G}}(\operatorname{\boldsymbol{\theta}}))_{il}}{\theta_{j}}=\partialderivative{(\operatorname{\boldsymbol{G}}(\operatorname{\boldsymbol{\theta}}))_{ij}}{\theta_{l}}.$

Therefore, the resulting Christoffel symbols have a relatively simple form,
given by

$\operatorname{\boldsymbol{\Gamma}}^{m}_{ij}=\frac{1}{2}\operatorname{\boldsymbol{G}}^{ml}\left(\partial_{i}\operatorname{\boldsymbol{G}}_{jl}+\partial_{j}\operatorname{\boldsymbol{G}}_{il}-\partial_{l}\operatorname{\boldsymbol{G}}_{ij}\right)=\frac{1}{2}\operatorname{\boldsymbol{G}}^{ml}\partial_{i}\operatorname{\boldsymbol{G}}_{jl}.$

###  F.2 Neural Networks (NN)

Here we expand the discussions on FIM for NNs in Section 4.3.3.

Recall the probabilistic model
$\pi(\operatorname{\boldsymbol{y}}|\operatorname{\boldsymbol{\phi}})$ on its
basic form. Given an NN that maps $\mathcal{X}$
$\overset{f_{\operatorname{\boldsymbol{\theta}}}}{\mapsto}\operatorname{\boldsymbol{\Phi}}$
where $\operatorname{\boldsymbol{\Phi}}$ is a $P$-dimensional parameter space
with given Fisher metric, we can define a Fisher metric on the parameter space
of the NN as follows. Using the chain rule for the Hessian matrix of
$\ell_{\operatorname{\boldsymbol{y}}}(\operatorname{\boldsymbol{\theta}})$
w.r.t $\operatorname{\boldsymbol{\theta}}$, observe that the extra expression
involving the score function will be zero in expectation [Schervish, 2012,
Definition 2.78, p. 111]. Therefore, the metric tensor w.r.t
$\operatorname{\boldsymbol{\theta}}$ reads

$\displaystyle\operatorname{\mathbb{E}}(-\nabla^{2}\ell_{\operatorname{\boldsymbol{Y}}}(\operatorname{\boldsymbol{\theta}})_{i,j})$
$\displaystyle=\operatorname{\mathbb{E}}\left(-\sum_{n}^{N}\partial^{2}_{i,j}\log\pi\big{(}Y_{n}|\operatorname{\boldsymbol{\phi}}_{n})\right)$
$\displaystyle=\sum_{n}^{N}\sum_{p}^{P}\sum_{q}^{P}\operatorname{\mathbb{E}}\left(-\partial^{2}_{q,p}\log\pi\big{(}Y_{n}|\operatorname{\boldsymbol{\phi}}_{n})\right)\partial_{i}(\operatorname{\boldsymbol{\phi}}_{n})_{p}\partial_{j}(\operatorname{\boldsymbol{\phi}}_{n})_{q}.$
$\displaystyle=\sum_{n}^{N}\sum_{p}^{P}\sum_{q}^{P}\operatorname{\boldsymbol{G}}_{q,r}(\operatorname{\boldsymbol{\phi}}_{n})\partial_{i}(\operatorname{\boldsymbol{\phi}}_{n})_{p}\partial_{j}(\operatorname{\boldsymbol{\phi}}_{n})_{q}$
$\displaystyle=\sum_{n}^{N}\sum_{p}^{P}\sum_{q}^{P}\operatorname{\boldsymbol{G}}_{q,r}(f_{\operatorname{\boldsymbol{\theta}}}(\operatorname{\boldsymbol{x}}_{n}))\partial_{i}(f_{\operatorname{\boldsymbol{\theta}}}(\operatorname{\boldsymbol{x}}_{n}))_{p}\partial_{j}(f_{\operatorname{\boldsymbol{\theta}}}(\operatorname{\boldsymbol{x}}_{n}))_{q},$

where the expectation taken w.r.t $\operatorname{\boldsymbol{Y}}$ is over the
Hessian of
$\ell_{\operatorname{\boldsymbol{Y}}}(\operatorname{\boldsymbol{\theta}})$ on
$\operatorname{\boldsymbol{\phi}}$ (in between 2nd and 3rd step) and that
comes for granted once the Fisher metric is known in the basic form of the
probabilistic model
$\pi(\operatorname{\boldsymbol{y}}|\operatorname{\boldsymbol{\phi}})$. In the
last passage we see that
$\partial_{i}(f_{\operatorname{\boldsymbol{\theta}}}(\operatorname{\boldsymbol{x}}_{n}))_{p}$
is the Euclidean derivative with respect to
$\operatorname{\boldsymbol{\theta}}_{i}$ of the $p^{th}$ output of the neural
network $f$ evaluated at the $n^{th}$ input. Then using matrix notations we
write,

$\displaystyle\operatorname{\boldsymbol{G}}(\operatorname{\boldsymbol{\theta}})=\sum_{n}^{N}\boldsymbol{J}_{n}^{\top}\operatorname{\boldsymbol{G}}_{\operatorname{\boldsymbol{\Phi}}}(f_{\operatorname{\boldsymbol{\theta}}}(\operatorname{\boldsymbol{x}}_{n}))\boldsymbol{J}_{n},$

where $\boldsymbol{J}_{n}$ $=$
$[\nabla(\boldsymbol{\operatorname{\boldsymbol{\phi}}}_{n})_{1}\cdots\nabla(\boldsymbol{\operatorname{\boldsymbol{\phi}}}_{n})_{P}]^{\top}$
denotes the $P\times D$ Jacobian matrix of the NN at the $n^{th}$ input.

Now, it is clear from the above form that its structure resembles that of a
pullback metric in Riemannian manifolds. To see this fix the inputs (or
covariates) in the input space and look instead to the output of the NN as a
map
$\operatorname{\boldsymbol{\Theta}}\overset{h_{n}}{\rightarrow}\operatorname{\boldsymbol{\Phi}}$
with
$h_{n}(\operatorname{\boldsymbol{\theta}})=f_{\operatorname{\boldsymbol{\theta}}}(\operatorname{\boldsymbol{x}}_{n})\
\forall n$. Once we have endowed a metric
$g_{\operatorname{\boldsymbol{\Phi}}}$ on $\operatorname{\boldsymbol{\Phi}}$
and defined a map for every $n$, the pullback metric on
$\operatorname{\boldsymbol{\Theta}}$ given all the inputs can be formally
written as the sum of all input information related to
$\operatorname{\boldsymbol{\theta}}$. That is,

$\displaystyle
g_{\operatorname{\boldsymbol{\Theta}}}(\operatorname{\boldsymbol{u}},\operatorname{\boldsymbol{v}})=\sum_{n}^{N}g_{\operatorname{\boldsymbol{\Phi}}}\big{(}\mathrm{d}h_{n}(\operatorname{\boldsymbol{\theta}})[\operatorname{\boldsymbol{u}}],\mathrm{d}h_{n}(\operatorname{\boldsymbol{\theta}})[\operatorname{\boldsymbol{v}}]\big{)},$

where
$\mathrm{d}h_{n}(\operatorname{\boldsymbol{\theta}})[\cdot]:T_{\operatorname{\boldsymbol{\theta}}}\operatorname{\boldsymbol{\Theta}}\rightarrow
T_{\operatorname{\boldsymbol{\phi}}_{n}}\operatorname{\boldsymbol{\Phi}}$ is
the differential at a point $\operatorname{\boldsymbol{\theta}}\ \forall n$.
Since
$\operatorname{\boldsymbol{\phi}}_{n}=f_{\operatorname{\boldsymbol{\theta}}}(\operatorname{\boldsymbol{x}}_{n})$
can be seen as a function of $\operatorname{\boldsymbol{\theta}}$, the
differential takes the form
$\mathrm{d}h_{n}(\operatorname{\boldsymbol{\theta}})[\operatorname{\boldsymbol{v}}]=\boldsymbol{J}_{n}\operatorname{\boldsymbol{v}}$.
Plugging this into the expression above and using that the matrix of
coefficients of the metric $g_{\operatorname{\boldsymbol{\Phi}}}$ is
$\operatorname{\boldsymbol{G}}_{\operatorname{\boldsymbol{\phi}}}$, we obtain
the metric tensor above; similar observation can be found in e.g. Grosse
[2022]. Also note that because the function $h_{n}$ may not be one-to-one for
all $n$, $\operatorname{\boldsymbol{x}}_{n}\in\mathcal{X}$, the metric tensor
$\operatorname{\boldsymbol{G}}(\operatorname{\boldsymbol{\theta}})$ may not
satisfy the definition of a pullback [see Manfredo, 1992].

##  Appendix G Experimental details

In this section, we describe the computation environment, some implementation
details, the employed methods to find the MAP estimate, to obtain the NUTS
samples, to integrate the ODEs and to calculate the Wasserstein distances, and
describe the basic setups for several experiments.

###  G.1 Computation environment

Apart from some preliminary ones, all experiments are carried out on the
computer cluster using CPUs with type Intel(R) Xeon(R) Gold 6230 CPU @
2.10GHz. When we benchmark the running times of the algorithms, we always use
two cores. The detailed configurations differ based on the nature of the jobs.

###  G.2 Implementation details

Sacred [Greff2017] is used to manage the experiments. Code for experiments
other than neural network ones are implemented mainly using JAX [Bradbury et
al., 2018], benefiting from fast computations of gradients and other
quantities and easy vectorizations. Code for neural network experiments are
written mainly in PyTorch [Paszke et al., 2019], using torch.func to obtain
the gradients and hessian vector products, etc, because we build on the
existing library provided by Daxberger et al. [2021]. We always use the
default data types of the libraries (float64 in NumPy [Harris et al., 2020]
and float32 in PyTorch and JAX) when initializing the variables. Note that in
experiments other than neural network ones, the random numbers are generated
by NumPy; In neural network experiments, they are generated by PyTorch.

Bergamin et al. [2023] formed the metric as identity plus the outer product of
gradients of negative log-posterior, while Hartmann et al. [2022] formed the
metric as identity plus a scaled version of the outer product of gradients of
log-posterior. We note that, when the scaling of Monge metric is $1$, these
are strictly mathematically equivalent formulations. In experiments other than
neural network ones, similar to Monge metric [Hartmann et al., 2022], we form
the metric as identity plus the outer product of gradients of log-posterior;
for neural network experiments, we form the metric similar to Bergamin et al.
[2023].

We used SymPy [Meurer et al., 2017] to verify some derivations, and ChatGPT
[OpenAI, 2023] as an assistance for writing code, scientific discussions etc.

###  G.3 Finding the MAP estimate

For experiments other than neural network experiments, unless there is an
analytical solution of the MAP estimate, we find the MAP estimate by
collecting $20$ runs of BFGS optimization with random initialization using the
implementation from Virtanen et al. [2020] which yield negative log-posterior
values other than plus infinity or NaN, with maximum number of iterations
$1e6$. We choose the run that yields the smallest negative log-posterior with
numpy.argmin.

For neural network experiments, we find the MAP estimate using $lr=1e-2$,
$weight\\_decay=1e-5$ for $20000$ epochs with the Adam [Kingma and Ba, 2015]
optimizer. The prior precision and noise std are optimized posthoc following
standard procedure from Daxberger et al. [2021]. During the optimization
process, we fix the random seeds, in order to have consistent prior precision
and noise std for NUTS.

###  G.4 NUTS samples

We use CmdStanPy [Stan Dev Team, 2023] to generate the NUTS samples. Following
the default options for reference posterior in posteriordb [Magnusson et al.,
2022], we run the NUTS sampler using $10$ chains, a thinning of $10$, a warm-
up of $10000$ iterations and $20000$ iterations per chain. We fix the random
seeds for reproducibility, generating a total of $20000$ samples.

For sampling from banana distribution, based on preliminary runs, we set
adapt_delta to be $0.95$. We otherwise keep the hyper parameters to be their
default values.

###  G.5 ODE integration

For the exponential maps, other than neural network experiments, ODE
integrations are carried out using Diffrax [Kidger, 2021]. Inspired by the
default settings of the function solve_ivp in SciPy [Virtanen et al., 2020],
we use Dormand-Prince’s $5(4)$ method [Dormand and Prince, 1980]. Diffrax
reports the number of steps num_steps used by the integrator. In order to make
the results approximately comparable to the number of function evaluations
nfev reported by SciPy, we multiply num_steps by $6$, the number of function
evaluations per step of Dormand-Prince’s $5(4)$ method. We use adaptive step
size with initial step size chosen by the algorithm, with $rtol=1e-3$ and
$atol=1e-6$, matching the default option of SciPy. We set the maximum number
of steps allowed to be $4096$ (approximately equivalent to $24576$ function
evaluations), following the default configuration of Diffrax. For neural
network experiments, following Bergamin et al. [2023], we use SciPy’s
solve_ivp and use the default hyper parameters.

For the logarithmic maps, we implement it using analytical gradients and
hessian vector products in NumPy.

In all experiments apart from logistic regression (in which we use analytical
gradients), we use the gradients and hessian vector products calculated by the
autodiff framework.

We always use analytical expressions for the Fisher metric. For logistic
regression and NN experiments, we use analytical expressions of the
Christoffel symbols, while in other experiments the Christoffel symbols are
solved numerically based on the metric.

When reporting the numbers of function evaluations, we report the numbers of
function evaluations required to solve the exponential maps. We observe that
in certain scenarios the number of integration steps can exceed the upper
limit. In which case, we treat the number of steps for those samples exceeding
the limit as the limit $4096$ while calculating $T$, and calculate the
Wasserstein distance of the samples not exceeding the limit.

###  G.6 Calculating the Wasserstein distance

We use Python Optimal Transport [Flamary et al., 2021] to calculate the
Wasserstein distance, which can be formulated as [Flamary et al., 2021]

$\displaystyle\min_{\operatorname{\boldsymbol{\gamma}}}\langle\operatorname{\boldsymbol{\gamma}},\operatorname{\boldsymbol{M}}\rangle_{F},$
$\displaystyle
s.t.\quad\operatorname{\boldsymbol{\gamma}}\operatorname{\boldsymbol{1}}=\operatorname{\boldsymbol{a}},$
$\displaystyle\operatorname{\boldsymbol{\gamma}}^{\top}\operatorname{\boldsymbol{1}}=\operatorname{\boldsymbol{b}},$
$\displaystyle\operatorname{\boldsymbol{\gamma}}\geq 0,$

where $\operatorname{\boldsymbol{\gamma}}$ is the transportation plan,
$\operatorname{\boldsymbol{M}}$ is the distance matrix induced by the
Euclidean metric for the two sets of samples, and
$\operatorname{\boldsymbol{a}}$ and $\operatorname{\boldsymbol{b}}$ are
uniform numbers summing up to $1$ with the lengths the same as the numbers of
respective samples. The maximum number of iterations is set to be $1e10$.

Wasserstein distance is a proper distance metric, and it becomes zero when the
two distributions are the same.

###  G.7 Squiggle

The experiment presented in Section 4.3.2 was carried out using the following
hyper parameters on the squiggle distribution

$a=1.5,\;\operatorname{\boldsymbol{S}}=\begin{bmatrix}5.0&0.0\\\
0.0&0.05\end{bmatrix}.$

###  G.8 Banana distribution

Results for sampling from banana distribution were presented in Section 5.2.
The Fisher metric for banana distribution can be found in Brofos and Lederman
[2021], and is given by

$\operatorname{\boldsymbol{G}}(\theta_{1},\theta_{2})=\begin{bmatrix}\frac{1}{\sigma_{\operatorname{\boldsymbol{\theta}}}^{2}}+\frac{N}{\sigma_{\operatorname{\boldsymbol{y}}}^{2}}&&\frac{2N\theta_{2}}{\sigma_{\operatorname{\boldsymbol{y}}}^{2}}\\\
\frac{2N\theta_{2}}{\sigma_{\operatorname{\boldsymbol{y}}}^{2}}&&\frac{1}{\sigma_{\operatorname{\boldsymbol{\theta}}}^{2}}+\frac{4N\theta_{2}^{2}}{\sigma_{\operatorname{\boldsymbol{y}}}^{2}}\end{bmatrix}.$

###  G.9 Bayesian logistic regression

We report the details on the $5$ logistic regression datasets, covering the
dimensionality of parameter of interest $D$ and the number of data points $N$.

Ripley [Ripley, 1994, 1996], abbreviated as Ripl: $D=3$, $N=250$.

Pima [Smith et al., 1988]: $D=8$, $N=532$.

Heart [Feng et al., 1992], abbreviated as Hear: $D=14$, $N=270$.

Australian [Feng et al., 1992], abbreviated as Aust: $D=15$, $N=690$.

German [Feng et al., 1992], abbreviated as Germ: $D=25$, $N=1000$.

###  G.10 Neural networks

Similar to Bergamin et al. [2023], we use the dataset from Snelson and
Ghahramani [2005], which has $200$ data points with labels, with the $x$
coordinates distributed between $0$ and $6$.

For complete training set, a test set of size $50$ is selected from the $200$
data points with labels with random seed $1$. As mentioned in the main paper,
the gap training set is constructed by using data points with labels with $x$
coordinates not between $1.5$ and $3.0$.

##  Appendix H Additional experimental results

###  H.1 Samples from banana distribution

In Section 5.2, we consider the problem of sampling from banana distribution.
We reported the numerical qualities of the results in Table 1, and the
different MAP estimates and the resulting approximations using RLA-F in Figure
4. Here we provide additional plots that show how all other variants work for
the same problem.

In the Banana distribution, there are two Euclidean MAPs, with approximately
the same log-posterior and approximately symmetric on the two sides of axis.
The results of starting from one Euclidean MAP and the Hausdorff MAP are shown
in Figure 6. We report the results of sampling from one particular Euclidean
MAP; the two MAPs are symmetric across the x-axis and starting from the other
Euclidean MAP would simply (approximately) mirror the approximation.

|  |  |
---|---|---|---
|  |  |
Figure 6: The first row shows approximations starting from one of the
Euclidean MAPs, and the second row shows approximations starting from the
Hausdorff MAP. Each row from left to right: ELA; RLA-B; RLA-BLog; RLA-F

Visually, RLA-B, RLA-BLog and RLA-F all adapt to the local curvature. RLA-B is
generally narrower, and is slightly in the wrong direction. RLA-BLog
alleviates these issues similar to RLA-F, and for all methods the Hausdorff
MAP is clearly better as already indicated by the numerical results in the
main paper.

###  H.2 Squiggle

|
---|---
Figure 7: Samples from RLA-B (left) and RLA-F (right) on another version of
squiggle.

In Section 4.3.2, RLA-B was demonstrated to struggle sampling from a squiggle
distribution. Moreover, in the particular runs shown in Figure 3 which
generated $10000$ samples respectively, RLA-B took on average $75$ function
evaluations, while Fisher only took $32$.

We note that its properties can vary. Figure 7 shows the results of RLA-B and
RLA-F on another version of squiggle, given by

$a=1.5,\;\operatorname{\boldsymbol{S}}=\begin{bmatrix}10.0&0.0\\\
0.0&0.001\end{bmatrix}.$

The high density area is clearly narrower. While in this version of squiggle
RLA-B seems to yield better results than the version as shown in the main
paper, there still exists a non-negligible mismatch. Moreover, the resulting
integration problem seems more difficult than using RLA-F: For this particular
run, RLA-F only took $41$ function evaluations on average compared to $300$
for RLA-B; Moreover, there were at least one sample where RLA-B exceeded the
maximum steps, while there were no such issues with RLA-F.

###  H.3 Funnel

Here we demonstrate how the approximations work in one more case of a
distribution of complex geometry that is diffeomorphic with a Gaussian. We
consider the problem of sampling from the funnel distribution [Neal, 2003]. It
can be formulated as

$p(\operatorname{\boldsymbol{\theta}})=\operatorname{\mathcal{N}}(\theta_{D}|0,\sigma)\prod_{d=1}^{D-1}\operatorname{\mathcal{N}}\left(\theta_{d}|0,\exp(\tfrac{1}{2}\theta_{D})\right).$

Consider the two dimensional funnel. It admits a natural reparameterization,
and can be formulated as a transformed version of multivariate Gaussian [Stan
Development Team, 2023], given by

$\displaystyle\operatorname{\boldsymbol{\psi}}$
$\displaystyle\sim\operatorname{\mathcal{N}}(\operatorname{\boldsymbol{0}},\operatorname{\boldsymbol{I}}),$
$\displaystyle\operatorname{\boldsymbol{\theta}}$
$\displaystyle=\operatorname{\boldsymbol{\phi}}(\operatorname{\boldsymbol{\psi}})=\begin{bmatrix}\exp(\sigma\psi_{2}/2)\psi_{1}\\\
\sigma\psi_{2}\end{bmatrix}.$

As such, it is possible to obtain a Riemannian metric, following the general
procedure in Section 4.3.2 of the main paper. We have

$\displaystyle\operatorname{\boldsymbol{\phi}}^{-1}(\operatorname{\boldsymbol{0}})$
$\displaystyle=\operatorname{\boldsymbol{0}},$
$\displaystyle\left(\partialderivative{\operatorname{\boldsymbol{\theta}}}{\operatorname{\boldsymbol{\psi}}}\right)^{-1}$
$\displaystyle=\begin{bmatrix}\exp(-\frac{1}{2}\theta_{2})&-\frac{1}{2}\theta_{1}\exp(-\frac{1}{2}\theta_{2})\\\
0&\frac{1}{\sigma}\end{bmatrix},$
$\displaystyle\operatorname{\boldsymbol{G}}_{\operatorname{\boldsymbol{\theta}}}$
$\displaystyle=\left(\partialderivative{\operatorname{\boldsymbol{\theta}}}{\operatorname{\boldsymbol{\psi}}}\right)^{-\top}\left(\partialderivative{\operatorname{\boldsymbol{\theta}}}{\operatorname{\boldsymbol{\psi}}}\right)^{-1}.$

The samples obtained using RLA-F on a two dimensional funnel with $\sigma=3$
are shown in Figure 8. Interestingly, all variants capture the marginal
distribution of $\theta_{2}$ well even though many MCMC methods struggle to
reach the narrow funnel. However, ELA and RLA-B fail to expand in the
direction of $\theta_{1}$ for larger values of $\theta_{2}$. RLA-BLog
alleviates the issue but is still not perfect, but RLA-F is exact.

|  |  |
---|---|---|---
Figure 8: Results on funnel. From left to right: ELA; RLA-B; RLA-BLog; RLA-F.

###  H.4 Bayesian logistic regression

In Section 5.3, we considered Bayesian logistic regression. We observed that
for RLA-B and RLA-BLog, there can be huge differences in terms of number of
function evaluations between standardized and raw inputs. We hypothesize that
this is in part caused by the large differences in the scales of the raw
inputs, and hence the gradient term in Monge metric, which is of quadratic
nature, has very different scales across dimensions. This induces strong
curvature and leads to challenging numerical integration problems.

We benchmark the time to obtain a single sample based on a sample from
Gaussian distribution in seconds of different algorithms and report them in
Table 4. We use a slightly optimized version of function for logarithmic map
implemented in NumPy and functions for exponential maps are implemented in
JAX. We use SciPy for integrations, averaging the obtained times over $2500$
samples. RLA-F consistently takes less time to obtain one sample, despite
requiring direct matrix inversions.

| stand. | raw
---|---|---
data | RLA-B | RLA-BLog | RLA-F | RLA-B | RLA-BLog | RLA-F
Ripl | 0.011 | 0.048 | 0.006 | 0.011 | 0.045 | 0.006
Pima | 0.016 | 0.26 | 0.006 | 1.989 | 1.183 | 0.006
Hear | 0.015 | 0.638 | 0.007 | 2.33 | 2.683 | 0.006
Aust | 0.02 | 0.734 | 0.01 | 6.183 | 4.632 | 0.008
Germ | 0.026 | 2.215 | 0.014 | 1.435 | 3.4 | 0.015
Table 4: Time per sample in seconds for logistic regression. Bold font
indicates the fastest method.

###  H.5 Neural networks

In Section 5.4, we presented NN regression results. Here we present further
visualizations.

For completeness, Figure 9 explicitly shows how ELA fails to generate
meaningful samples for this problem.

|
---|---
Figure 9: Neural network regression with ELA. Left: complete data; right: gap
data.

####  H.5.1 Numerical stability

|
---|---
Figure 10: Bad runs with RLA-B on neural networks. The left shows the complete case, and the right shows the gap case. |
---|---
Figure 11: Bad runs with RLA-BLog on neural networks. The left shows the
complete case, and the right shows the gap case. Figure 12: Bad run with RLA-F
on neural networks, for the gap case.

In Figures 5 (in main paper) and 5, we purposefully reported runs which
resulted in relatively well-behaved predictive distributions for the methods
under comparison. RLA-F is numerically stable but the other methods
occasionally fail. This is demonstrated in Figures 10 and 11 that show
(examples of) bad runs where both RLA-B and RLA-BLog have extremely bad and
isolated samples that influence even the mean prediction significantly. For
completeness, Figure 12 shows the only run among the $10$ independent runs
where RLA-F has any issues, in form of two minor outlier samples that do not
have notable influence on the whole predictive distribution.

###  H.6 Empirical Fisher

| data | $\mathcal{W}$ | $T$
---|---|---|---
stand. | Ripl | [0.106, 0.004] | 12.0
Pima | [0.148, 0.0] | 12.1
Hear | [0.53, 0.0] | 13.8
raw | Ripl | [0.426, 0.011] | 12.1
Pima | [0.199, 0.001] | 12.5
Hear | [0.942, 0.02] | 17.0
Table 5: Logistic regression results with Empirical Fisher metric, reported as
$[\text{mean},\text{std}]$.

Some works have considered an alternative Riemannian metric known as the
Empirical Fisher Information Matrix, given by

$\sum_{n=1}^{N}\nabla\log\pi(y_{n}|\operatorname{\boldsymbol{\theta}})\nabla\log\pi(y_{n}|\operatorname{\boldsymbol{\theta}})^{\top}.$

There have been arguments that it may not lead to desirable behaviors
[Kunstner et al., 2019]. Lan et al. [2015] also considers another Empirical
Fisher Information Matrix whose estimates are more accurate than the previous
ones. We tried replacing the FIM with the above empirical version in the
Fisher metric and performed a range of preliminary experiments to evaluate
whether the Empirical Fisher could also be used as metric in RLA, always using
numerical Christoffel symbols.

While it works reasonably well for sampling from banana distribution, yielding
a Wasserstein distance of $[0.799,0.008]$ with Euclidean MAP and
$[0.141,0.003]$ with Hausdorff MAP, it becomes much worse for Bayesian
logistic regression, as shown in Table 5. Especially, for the experiment on
Heart dataset without standardization, it is much worse than Euclidean. We
therefore did not explore it further.