remarkRemark hypothesisHypothesis claimClaim Affine Invariant Gradient FlowsY.
Chen, D.Z.Huang, J. Huang, S. Reich and A. M. Stuart

# Gradient Flows for Sampling: Mean-Field Models, Gaussian Approximations and
Affine Invariance††thanks: Submitted to the editors DATE. The authors are in
alphabetical order.

Y. Chen California Institute of Technology, Pasadena, CA (, , )
<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>D.Z.
Huang22footnotemark: 2 J. Huang University of Pennsylvania, Philadelphia, PA
<EMAIL_ADDRESS>S. Reich Universität Potsdam, Potsdam, Germany ()
<EMAIL_ADDRESS>A.M. Stuart22footnotemark: 2

###### Abstract

Sampling a probability distribution with an unknown normalization constant is
a fundamental problem in computational science and engineering. This task may
be cast as an optimization problem over all probability measures, and an
initial distribution can be evolved to the desired minimizer (the target
distribution) dynamically via gradient flows. Mean-field models, whose law is
governed by the gradient flow in the space of probability measures, may also
be identified; particle approximations of these mean-field models form the
basis of algorithms. The gradient flow approach is also the basis of
algorithms for variational inference, in which the optimization is performed
over a parameterized family of probability distributions such as Gaussians,
and the underlying gradient flow is restricted to the parameterized family.

By choosing different energy functionals and metrics for the gradient flow,
different algorithms with different convergence properties arise. In this
paper, we concentrate on the Kullback–Leibler divergence as the energy
functional after showing that, up to scaling, it has the unique property
(among all $f$-divergences) that the gradient flows resulting from this choice
of energy do not depend on the normalization constant of the target
distribution. For the metrics, we focus on variants of the Fisher-Rao,
Wasserstein, and Stein metrics; we introduce the affine invariance property
for gradient flows, and their corresponding mean-field models, determine
whether a given metric leads to affine invariance, and modify it to make it
affine invariant if it does not.

We study the resulting gradient flows in both the space of all probability
density functions and in the subset of all Gaussian densities. The flow in the
Gaussian space may be understood as a Gaussian approximation of the flow in
the density space. We demonstrate that, under mild assumptions, the Gaussian
approximation based on the metric and through moment closure coincide; the
moment closure approach is more convenient for calculations. We establish
connections between these approximate gradient flows, discuss their relation
to natural gradient methods in parametric variational inference, and study
their long-time convergence properties showing, for some classes of problems
and metrics, the advantages of affine invariance. Furthermore, numerical
experiments are included which demonstrate that affine invariant gradient
flows have desirable convergence properties for a wide range of highly
anisotropic target distributions.

###### keywords:

Bayesian inference, sampling, gradient flow, mean-field dynamics, Gaussian
approximation, variational inference, affine invariance.

68Q25, 68R10, 68U05

##  1 Introduction

###  1.1 Context

This paper is concerned with the problem of sampling a probability
distribution (the target) known up to normalization. This problem is
fundamental in many applications arising in computational science and
engineering and is widely studied in the applied mathematics, machine learning
and statistics communities. A particular application is Bayesian inference for
large-scale inverse problems; such problems are ubiquitous, arising in
applications from climate science [65, 118, 63, 91], through numerous problems
in engineering [136, 38, 21] to machine learning [111, 100, 28, 32]. These
applications have fueled the need for efficient and scalable algorithms which
employ noisy data to learn about unknown parameters $\theta$ appearing in
models and perform uncertainty quantification for predictions then made by
those models.

Mathematically, the objective is to sample the target probability distribution
with density $\rho_{\rm post}(\cdot)$, for the parameter
$\theta\in\mathbb{R}^{N_{\theta}}$, given by

(1) $\displaystyle\rho_{\rm post}(\theta)\propto\exp(-\Phi_{R}(\theta)),$

where $\Phi_{R}:\mathbb{R}^{N_{\theta}}\to\mathbb{R}_{+}$ is a known function.
We use the notation $\rho_{\rm post}$ because of the potential application to
Bayesian inference; however we do not explicitly use the Bayesian structure in
this paper, and our analysis applies to arbitrary target distributions.

We study the use of gradient flows in the space of probability distributions
in order to create algorithms to sample the target distribution. By studying
gradient flows with respect to different metrics, and by studying mean-field
based particle models and Gaussian approximations, both related to these
underlying gradient flows, we provide a unifying approach to the construction
of a wide family of algorithms. The choice of metric plays a key role in the
behavior of the resulting methods and we highlight the importance of affine
invariance in this regard. In Section 1.2 we provide a literature review
pertinent to our contributions; the contributions we make are described in
Section 1.3 and in Section 1.4 we describe the organization of the remainder
of the paper.

###  1.2 Literature Review

In this subsection, we describe the research landscape in which our work sits.
We start by discussing the general background, and describing our work in this
context, and then we give more detailed literature reviews relating to the
topics of gradient flows, mean-field models, Gaussian approximations, and
affine invariance.

####  1.2.1 Background

Numerous approaches to the sampling problem have been proposed in the
literature. One way of classifying them is into: a) methods which deform a
given source measure (for example the prior in Bayesian inference) into the
target measure, in a fixed time of fixed finite number of steps or in a finite
continuous time interval; and b) methods which transform any initial measure
into the target measure after an infinite number of steps, or at time infinity
in continuous time. Continuous time formulations are used for insight into the
algorithms; discrete time must be used in practice. Typical methods in
category a) are sequential Monte Carlo (SMC) approaches [44], with (typically
not optimal) transport being the underpinning continuous time concept [127];
typical methods in category b) are Markov chain Monte Carlo (MCMC) approaches
[18], with stochastic differential equations (SDEs) which are ergodic with
respect to the target, such as Langevin equations [106], being the
underpinning continuous time concept. Making practical algorithms out of these
ideas, for large scale problems in science and engineering, often requires
invocation of further reduction of the space in which solutions are sought,
for example by variational inference [15] or by ensemble Kalman approximation
[20].

In this paper we focus primarily on methods in category b) and describe a
general methodology for the derivation of a wide class of sampling algorithms;
however the methods can be interpreted as being partially motivated by
dynamics of transports arising in the methods of type a). Specifically, we
focus on methods created by studying the gradient flow, in various metrics,
induced by an energy which measures divergence of the current estimate of the
target from the true target. With this perspective, we provide a unifying
viewpoint on a number of sampling methods appearing in the literature and a
methodology for deriving new methods. We focus on the Fisher-Rao, Wasserstein,
and Stein metrics, and variants thereof. Creating useful algorithms out of
this picture requires further simplifications; we study mean-field models,
which lead to particle methods, and methods based on confining the gradient
descent of the energy to the space of Gaussians. In both settings we precisely
define the concept of being affine invariant; roughly speaking this concept
requires that any invertible affine transformation of $\theta$ makes no
difference to the gradient flow. We include numerical experiments which
demonstrate the advantage of affine invariant methods for anisotropic targets.
Because the analysis is cleaner we work in continuous time; but time-
discretization is employed to make implementable algorithms. Furthermore we
emphasize that our statements about the existence of gradient flows are purely
formal. For the rigorous underpinnings of gradient flows see [3]; and for a
recent extension of this rigorous analysis to a sub-class of gradient flows
with respect to an affine invariant metric see [19].

####  1.2.2 Gradient Flows

There is existing literature on the use of gradient flows in the probability
density space, employing a variety of different metric tensors, to minimize an
energy defined as the Kullback–Leibler (KL) divergence between the current
density and the target distribution. Particle realizations of these flows then
lead to sampling schemes. For example, the Wasserstein gradient flow [68, 105]
and Stein variational gradient flow [89, 88] have led to sampling algorithms
based on Langevin dynamics and Stein variational gradient descent
respectively; in [94], the Fisher-Rao gradient flow, using kernel-based
density approximations, has been proposed for sampling. Furthermore, the paper
[93] proposed the Wasserstein-Fisher-Rao gradient flow to sample multi-modal
distributions. In [54, 55], the Kalman-Wasserstein metric was introduced and
gradient flows with respect to this metric were advocated. Accelerated
gradient flows in the probability space have been studied in [134]. A recent
overview of the use of gradient flows in optimization and sampling can be
found in [126].

The Wasserstein gradient flow was identified in the seminal work [68]. The
authors showed that the Fokker-Planck equation is the Wasserstein gradient
flow of the KL divergence of the current density estimate from the target.
Since then, Wasserstein gradient flow has played a significant role in optimal
transport [115], sampling [30, 80], machine learning [34, 114], partial
differential equations [105, 22] and many other areas. The Fisher-Rao metric
was introduced by C.R. Rao [110] via the Fisher information matrix. The
original definition is in parametric density spaces, and the corresponding
Fisher-Rao gradient flow in the parameter space leads to natural gradient
descent [1]. The Fisher-Rao metric in infinite dimensional probability spaces
was discussed in [51, 122]. The concept underpins information geometry [2, 6].
The gradient flow of the KL divergence under the Fisher-Rao metric is induced
by a mean-field model of birth-death type. The birth-death process has been
used in sequential Monte Carlo samplers to reduce the variance of particle
weights [40] and to accelerate Langevin sampling [93, 94]. The discovery of
the Stein metric [88] follows the introduction of the Stein variational
gradient descent algorithm [89]. The study of the Stein gradient flow [88, 92,
45] sheds light on the analysis and improvements of the algorithm [41, 133,
134].

####  1.2.3 Mean-Field Models

It is natural to ask which evolution equations in state space
$\mathbb{R}^{N_{\theta}}$ give rise to a given gradient flow in the space of
probability measures. Continuous time linear Markov processes with continuous
sample paths are limited to Itô SDEs (or equivalent models written in terms of
Stratonovich or other stochastic integals) [106]. It is thus natural to seek
mean-field models in the form of Itô SDEs which depend on their own density,
which is therefore governed by a nonlinear Fokker-Planck equation. Examples of
particle models giving rise to linear and nonlinear Fokker-Planck equations
with gradient structure include Langevin dynamics [68, 105] and Stein
variational gradient descent [89, 88] for sampling the Wasserstein gradient
flow and the Stein variational gradient flow respectively. It is also of
potential interest to go beyond Itô SDEs and include Levy (jump) processes
[13, 4], as well as birth-death models [73]. Finally, we mention mean-field
models for ensemble Kalman type algorithms [20]; these typically do not have a
law which is a gradient flow in the space of probability measures, except in
the linear-Gaussian setting. These mean-field models combine gradient flow,
Gaussian approximations, and mean-field equations [42, 16, 20].

In practice, mean-field models are approximated by interacting particle
systems [66], in which integration against the density is replaced by
integration against the empirical measure of the particle system. At the level
of the nonlinear Markov process for the density on $\mathbb{R}^{N_{\theta}}$
defined by the mean-field model, this corresponds to approximation by a linear
Markov process for the density on $\mathbb{R}^{JN_{\theta}}$, where $J$ is the
number of particles; the concepts of exchangeability and propagation of chaos
may be used to relate the two Markov processes. See [99, 124, 26] and the
references therein.

####  1.2.4 Gaussian Approximations

There is substantial work on the use of gradient flows in the space of
Gaussian, or other parametric density spaces, to minimize the Kullback–Leibler
(KL) divergence [129, 15]. These methods, in the Gaussian setting, aim to
solve the problem

(2)
$\displaystyle(m^{\star},C^{\star})=\operatorname*{arg\,min}_{m,C}~{}\mathrm{KL}[\mathcal{N}(m,C)\|\rho_{\rm
post}(\theta)].$

Again, but now restricted to variations in the space of Gaussian densities,
different metric tensors lead to different gradient flows to identify
$(m^{\star},C^{\star}).$ Recently, the work [80] proved the global convergence
of the Wasserstein natural gradient descent algorithm when the posterior is
log-concave. Other work on the use of Gaussian variational inference methods
includes the papers [104, 109, 74, 86, 52, 137].

In addition to their role in parametric variational inference, Gaussian
approximations have been widely deployed in various generalizations of Kalman
filtering [71, 121, 70, 130, 48]. For Bayesian inverse problems, iterative
ensemble Kalman samplers have been proposed which are in category a) defined
in subsection 1.2.1 [47, 31, 130]. The paper [62] introduced an ensemble
Kalman methodology falling in category b), defined in subsection 1.2.1, based
on a novel mean-field dynamical system that depends on its own filtering
distribution. For all these algorithms based on a Gaussian ansatz, the
accuracy depends on some measure of being close to Gaussian. Regarding the use
of Gaussian approximations in Kalman inversion we highlight, in addition to
the approximate Bayesian methods already cited, the use of ensemble Kalman
methods for optimization: see [64, 24, 76, 63, 135]. Kalman filtering has also
been used in combination with variational inference [78]. The relation between
iterative Kalman filtering and Gauss-Newton or Levenberg Marquardt algorithms
are studied in [11, 10, 63, 25], and leads to ensemble Kalman based
optimization methods which are affine invariant.

####  1.2.5 Affine Invariance

The idea of affine invariance was introduced for MCMC methods in [58, 49],
motivated by the empirical success of the Nelder-Mead simplex algorithm [101]
in optimization. Sampling methods with the affine invariance property can be
effective for highly anisotropic distributions; this is because they behave
identically in all coordinate systems related through an affine
transformation; in particular, they can be understood by studying the best
possible coordinate system, which reduces anisotropy to the maximum extent
possible within the class of affine transformations. The numerical studies
presented in [58] demonstrate that affine-invariant MCMC methods offer
significant performance improvements over standard MCMC methods. This idea has
been further developed to enhance sampling algorithms in more general
contexts. Preconditioning strategies for Langevin dynamics to achieve affine-
invariance were discussed in [83]. And in [54], the Kalman-Wasserstein metric
was introduced, gradient flows in this metric were advocated and in [55] the
methodology was shown to achieve affine invariance. Moreover, the authors in
[54, 55, 107] used the empirical covariance of an interacting particle
approximation of the mean-field limit, leading to a family of derivative-free
sampling approaches in continuous time. Similarly, the work [90] employed the
empirical covariance to precondition second order Langevin dynamics. Affine
invariant samplers can also be combined with the pCN (preconditioned
Crank–Nicolson) MCMC method [36], to boost the performance of MCMC in function
space [37, 46]. Another family of affine-invariant sampling algorithms is
based on Newton or Gauss-Newton, since the use of the Hessian matrix as the
preconditioner in Newton’s method induces the affine invariance property. Such
methods include stochastic Newton MCMC [98] and the Newton flow with different
metrics [41, 133].

###  1.3 Our Contributions

The primary contributions of the work are as follows:

  * •

we highlight a general methodology for the design of algorithms to sample a
target probability distribution known up to normalization, unifying and
generalizing an emerging scattered literature;

  * •

the methodology is based on the introduction of gradient flows of the KL
divergence between the time-dependent density which is solution of the
gradient flow, and the target density;

  * •

we justify the choice of the KL divergence as the energy functional by showing
that, among all $f$-divergences, it is the unique choice (up to scaling) for
which the resulting gradient flow is independent of the normalization constant
of the target distribution;

  * •

the notion of gradient flow requires a metric and we employ the Fisher-Rao,
Wasserstein and Stein metrics to provide concrete instantiations of the
methodology;

  * •

to design implementable algorithms from the gradient flows we discuss the use
of particle approximations of mean-field models, whose law is governed by the
gradient flow, and restriction of the gradient flow to a parameterized
Gaussian family, which we show to be equivalent to a moment closure approach;

  * •

we define the concept of affine invariant metrics, demonstrate links to affine
invariant mean-field models and variational methods restricted to the set of
Gaussians, and describe numerical results highlighting the benefits of affine
invariant methods;

  * •

we prove results concerning the long time behavior of the underlying gradient
flows, in both the full and Gaussian density spaces, further highlighting the
benefits of affine invariance in some cases.

Our code is accessible online:

https://github.com/Zhengyu-Huang/InverseProblems.jl

###  1.4 Organization

The remainder of the paper is organized as follows. In Section 2, we introduce
energy functionals in the density space. In Section 3, we review the basics of
gradient flows in the space of all probability density functions; we make
links to mean-field models and we propose a definition of affine invariant
metrics, leading to affine invariant gradient flows and mean-field models.
Examples of affine invariant Fisher-Rao, Wasserstein and Stein gradient flows
are discussed and their convergence properties are studied theoretically; some
of these results highlight the benefits of affine invariance. In Section 4, we
review the basics of Gaussian approximate gradient flows. We define different
Gaussian approximate gradient flows under the aforementioned metrics,
computing the dynamics governing the evolution of the mean and covariance,
studying their convergence properties, and again identifying the effects of
affine invariance. A by-product of our computations is to show that the
evolution equations for mean and covariance can be computed by simple use of
moment closure. In Section 5, numerical experiments are provided to
empirically confirm the theory and in particular to demonstrate the
effectiveness of the affine invariance property in designing algorithms for
certain classes of problems. We make concluding remarks in Section 6. Four
appendices contain details of the proofs of the results stated in the main
body of the paper.

##  2 Energy Functional

Consider the space of all strictly positive probability density functions,
defined by111Here, and in what follows, we omit the domain
$\mathbb{R}^{N_{\theta}}$ of the integral, unless including it is needed for
purposes of disambiguation. Furthermore, extension to non-negative densities
is possible and is relevant in some applications; but we work in the simpler
setting of strict positivity throughout this paper.

(3) $\mathcal{P}=\Bigl{\\{}\rho\in
L_{1}(\mathbb{R}^{N_{\theta}}):\int\rho\mathrm{d}\theta=1,\,\rho>0\Bigr{\\}}.$

Now consider any $\rho\in\mathcal{P}$ which is absolutely continuous with
respect to $\rho_{\rm post}$, noting that we may then define the energy
$\mathcal{E}(\rho)$ by

(4) $\displaystyle\mathcal{E}(\rho)=\mathrm{KL}[\rho\|\rho_{\rm
post}]=\int\rho\log\Bigl{(}\frac{\rho}{\rho_{\rm
post}}\Bigr{)}\,\mathrm{d}\theta;$

we can extend $\mathcal{E}(\cdot)$ to the whole of $\mathcal{P}$ by setting it
to $\infty$ when $\rho$ is not absolutely continuous with respect to
$\rho_{\rm post}.$ The minimization problem

(5) $\displaystyle\min_{\rho\in\mathcal{P}}\mathrm{KL}[\rho\|\rho_{\rm post}]$

has as unique global minimizer $\rho=\rho_{\rm post}$. This suggests the
derivation of algorithms to identify $\rho_{\rm post}$ based on minimization
of $\mathcal{E}(\rho)$ over $\mathcal{P}$; we refer to this as nonparametric
variational inference.

When minimizing $\mathcal{E}(\rho)$ the first variation plays a central role.
For the specific choice of KL divergence as the energy functional, the first
variation is given by

(6) $\displaystyle\frac{\delta\mathcal{E}}{\delta\rho}=\log\rho-\log\rho_{\rm
post}-\int(\log\rho-\log\rho_{\rm post}){\rm d}\theta,$

where we have used the fact that $(\rho\log\rho)^{\prime}=\log\rho+1$, and we
impose $\int\frac{\delta\mathcal{E}}{\delta\rho}(\theta){\rm d}\theta=0.$ From
the formula Eq. 6 we see that, for the KL divergence,
$\frac{\delta\mathcal{E}}{\delta\rho}$ remains unchanged if we scale
$\rho_{\rm post}$ by any positive constant $c>0$, i.e. if we change $\rho_{\rm
post}$ to $c\rho_{\rm post}$. This property eliminates the need to know the
normalization constant of $\rho_{\rm post}$ in order to calculate the first
variation. It is common in Bayesian inference for the normalization to be
unknown and indeed the fact that MCMC methods do not need the normalization
constant is central to their widespread use; it is desirable that the
methodology presented here has the same property.

The following Proposition 2.1 shows that this property of the KL divergence is
special: among all $f$-divergences with continuously differentiable $f$
defined on the positive reals it is the only one to have this property. Here
the $f$-divergence between two continuous density functions $\rho$ and
$\rho_{\rm post}$, positive everywhere, is defined as

$D_{f}[\rho\|\rho_{\rm post}]=\int\rho_{\rm
post}f\Bigl{(}\frac{\rho}{\rho_{\rm post}}\Bigr{)}{\rm d}\theta.$

For convex $f$ with $f(1)=0$, Jensen’s inequality implies that
$D_{f}[\rho\|\rho_{\rm post}]\geq 0$. The KL divergence used in (4)
corresponds to the choice $f(x)=x\log x.$ In what follow we view this
$f$-divergence as a function of probability measure $\rho$, parameterized by
$\rho_{\rm post}$; in particular we observe that this parameter-dependent
function of probability density $\rho$ makes sense if $\rho_{\rm post}$ is
simply a positive function: it does not need to be a probability density; we
may thus scale $\rho_{\rm post}$ by any positive real.

######  Proposition 2.1.

Assume that $f:(0,\infty)\to\mathbb{R}$ is continuously differentiable and
$f(1)=0$. Then the KL divergence is the only $f$-divergence (up to scalar
factors) whose first variation with respect to $\rho$ is invariant with
respect to $\rho_{\rm post}\mapsto c\rho_{\rm post}$, for any $c\in(0,\infty)$
and for any $\rho_{\rm post}\in\mathcal{P}$.

The proof of the proposition can be found in Section A.1.

a consequence of Proposition 2.1, the gradient flows defined by the energy (4)
do not depend on the normalization constant of the posterior, as we will see
in the next section. Hence the numerical implementation is more
straightforward in comparison with the use of other divergences or metrics to
define the energy $\mathcal{E}(\rho)$. This justifies the choice of KL
divergence as energy functional for sampling. and our developments in most of
this paper are specific to the energy Eq. 4. However, other energy functionals
can be, and are, used for constructing gradient flows; for example, the chi-
squared divergence [33, 87]:

(7) $\chi^{2}(\rho\|\rho_{\rm post})=\int\rho_{\rm
post}\Bigl{(}\frac{\rho}{\rho_{\rm
post}}-1\Bigr{)}^{2}\mathrm{d}\theta=\int\frac{\rho^{2}}{\rho_{\rm
post}}\mathrm{d}\theta-1.$

The normalization constant can appear explicitly in the gradient flow equation
for general energy functionals. Additional structures need to be explored to
simulate such flows. For example, when the energy functional is the chi-
squared divergence, in [33], kernelization is used to avoid the normalization
constant in the Wasserstein gradient flow. Moreover, in [87] where a
modification of the Fisher-Rao metric is used, ensemble methods with birth-
death type dynamics are adopted to derive numerical methods; the normalization
constant can be absorbed into the birth-death rate.

In the context of algorithms, it is also of interest to consider minimization
of $\mathcal{E}(\cdot)$ given by Eq. 4 over parameterized manifolds in
$\mathcal{P}$: parametric variational inference. To illustrate this, we
consider the manifold of Gaussian densities222The extension to $C\succeq 0$
may be relevant for some applications but we work in the simpler, strictly
positive covariance, setting here. $\mathcal{P}^{G}\subset\mathcal{P}$

(8a) $\displaystyle\mathcal{P}^{G}$
$\displaystyle:=\Bigl{\\{}\rho_{a}:\rho_{a}(\theta)=\frac{\exp\bigl{(}{-\frac{1}{2}(\theta-m)^{T}C^{-1}(\theta-m)}\bigr{)}}{\sqrt{|2\pi
C|}}\textrm{ with }a=(m,C)\in\mathcal{A}\Bigr{\\}},$ (8b)
$\displaystyle\mathcal{A}$
$\displaystyle=(m\in\mathbb{R}^{N_{\theta}},\,C\succ
0\in\mathbb{R}^{N_{\theta}\times N_{\theta}});$

here $|\cdot|$ denotes the determinant when the argument is a matrix. This
definition leads to Gaussian variational inference:

(9) $\min_{\rho\in\mathcal{P}^{G}}\mathrm{KL}[\rho\|\rho_{\rm post}].$

Any minimizer $\rho_{a_{\star}}=\mathcal{N}(m_{\star},C_{\star})$ satisfies
[79]333We use $\nabla_{\theta}\nabla_{\theta}f(\theta)$ to denote the Hessian
matrix associated with scalar field $f(\theta)$. In doing so we follow the
convention in the continuum mechanics literature, noticing that the Hessian
operator is formed as the composition of the gradient acting on a scalar
field, followed by the gradient acting on a vector field [57]. The notation
$\nabla_{\theta}^{2}f(\theta)$ is used by some authors to denote the Hessian;
we avoid this because of potential confusion with its useage in some fields to
denote the Laplacian (trace of the Hessian).

(10)
$\displaystyle\mathbb{E}_{\rho_{a_{\star}}}\bigl{[}\nabla_{\theta}\log\rho_{\rm
post}(\theta)\bigr{]}=0\quad\textrm{and}\quad
C_{\star}^{-1}=-\mathbb{E}_{\rho_{a_{\star}}}\bigl{[}\nabla_{\theta}\nabla_{\theta}\log\rho_{\rm
post}(\theta)\bigr{]}.$

##  3 Gradient Flow

In this section, we start by introducing the concept of metric, the gradient
flow of energy Eq. 4 that it induces, the related mean-field dynamics in the
state space $\mathbb{R}^{N_{\theta}}$, and the concept of affine invariance,
all in Section 3.1. Then, in subsequent subsections, we introduce the Fisher-
Rao (Section 3.2), Wasserstein (Section 3.3) and Stein gradient flows (Section
3.4), together with affine invariant modifications when relevant. In fact the
Fisher-Rao metric has a stronger invariance property: it is invariant under
any diffeomorphism of the parameter space. Finally, we discuss the convergence
properties of these gradient flows in Section 3.5.

###  3.1 Basics of Gradient Flows

In this subsection, we introduce gradient flows in the probability space and
affine invariance in this context. Our focus is on formal calculations to
derive these flows. We do not focus on the rigorous analytical underpinnings
of gradient flows in a metric space; the reader interested in further details
should consult [3].

####  3.1.1 Metric

For simplicity, the probability space we consider in this paper is the
manifold $\mathcal{P}$ of smooth strictly positive densities given by Eq. 3.
Then, at any $\rho\in\mathcal{P}$, the tangent space of $\mathcal{P}$ is

(11) $\displaystyle T_{\rho}\mathcal{P}:=\Bigl{\\{}\sigma\in
C^{\infty}(\mathbb{R}^{N_{\theta}}):\int\sigma\mathrm{d}\theta=0\Bigr{\\}}.$

We can identify the cotangent space $T_{\rho}^{*}\mathcal{P}\simeq
T_{\rho}\mathcal{P}$. This isomorphism relates the tangent vector $\sigma\in
T_{\rho}\mathcal{P}$ to the covector in $T_{\rho}^{*}\mathcal{P}$ via the
identification $\langle\sigma,\cdot\rangle$, and where
$\langle\sigma_{1},\sigma_{2}\rangle=\int\sigma_{1}\sigma_{2}$ denotes the
inner product in $T_{\rho}\mathcal{P}$. From any cotangent vector in
$T_{\rho}^{*}\mathcal{P}$ we can use the Riesz representation theorem to
represent it in $T_{\rho}\mathcal{P}.$

Given a metric tensor at $\rho$, defined by
$M(\rho):T_{\rho}\mathcal{P}\rightarrow T_{\rho}^{*}\mathcal{P}$, we define
Riemannian metric $g_{\rho}:T_{\rho}\mathcal{P}\times
T_{\rho}\mathcal{P}\to\mathbb{R}$ via $g_{\rho}(\sigma_{1},\sigma_{2})=\langle
M(\rho)\sigma_{1},\sigma_{2}\rangle$. The inverse of $M(\rho)$ is sometimes
referred to as the Onsager operator.

In the definition of $M(\rho)$, we identify $M(\rho)\sigma_{1}\in
T^{*}_{\rho}\mathcal{P}$ as its representer in $T_{\rho}\mathcal{P}$ and thus
$\langle M(\rho)\sigma_{1},\sigma_{2}\rangle$ is well-defined, given that
$\langle\cdot,\cdot\rangle$ is the $L^{2}$ inner product in the tangent space.
In fact, for ease of notation, in this paper we sometimes use the same
notation for a dual element and its representer in the primal space. One may
also understand $\langle\cdot,\cdot\rangle$ in $\langle
M(\rho)\sigma_{1},\sigma_{2}\rangle$ as the primal-dual pairing between
$T^{*}_{\rho}\mathcal{P}$ and $T_{\rho}\mathcal{P}$; however we employ the
$L^{2}$ inner product interpretation throughout this paper for convenience.

The geodesic distance
$\mathcal{D}:\mathcal{P}\times\mathcal{P}\to\mathbb{R}^{+}$ under metric $g$
is defined by the formula

(12)
$\mathcal{D}(\rho_{A},\rho_{B})^{2}=\inf_{\rho_{t}}\Bigl{\\{}\int_{0}^{1}g_{\rho_{t}}(\partial_{t}{\rho}_{t},\partial_{t}{\rho}_{t})\mathrm{d}t:\,\rho_{0}=\rho_{A},\rho_{1}=\rho_{B}\Bigr{\\}}.$

The distance $\mathcal{D}$ defines a metric on probability measures; however,
to avoid confusion with the Riemannian metric $g$, in this paper we always
refer to $\mathcal{D}$ as a distance.

We also recall that the geodesic distance $\mathcal{D}$ has the following
property [43]:

(13) $\lim_{\epsilon\to
0}\frac{1}{\epsilon^{2}}\mathcal{D}(\rho+\epsilon\sigma,\rho)^{2}=g_{\rho}(\sigma,\sigma)=\langle
M(\rho)\sigma,\sigma\rangle.$

####  3.1.2 Flow Equation

Recall that the dual element of the first variation of $\mathcal{E}(\rho)$,
denoted $\frac{\delta\mathcal{E}}{\delta\rho}\in T_{\rho}^{*}\mathcal{P}$, is
defined by

$\displaystyle\Bigl{\langle}\frac{\delta\mathcal{E}}{\delta\rho},\sigma\Bigr{\rangle}=\lim_{\epsilon\rightarrow
0}\frac{\mathcal{E}(\rho+\epsilon\sigma)-\mathcal{E}(\rho)}{\epsilon},$

for any $\sigma\in T_{\rho}\mathcal{P}$. In the preceding, we identify
$\frac{\delta\mathcal{E}}{\delta\rho}\in T_{\rho}^{*}\mathcal{P}$ via its
representer in $T_{\rho}\mathcal{P}$ so that the inner product notation is
well-defined and consistent; see LABEL:remark:_primal_dual_identification. The
gradient of $\mathcal{E}$ under the Riemannian metric, denoted by
$\nabla\mathcal{E}$, is defined via the condition

$\forall\sigma\in T_{\rho}\mathcal{P}\qquad
g_{\rho}(\nabla\mathcal{E},\sigma)=\Bigl{\langle}\frac{\delta\mathcal{E}}{\delta\rho},\sigma\Bigr{\rangle}.$

Using the metric tensor, we can write
$\nabla\mathcal{E}(\rho)=M(\rho)^{-1}\frac{\delta\mathcal{E}}{\delta\rho}$.

The gradient flow of $\mathcal{E}$ with respect to this metric is thus defined
by

(14) $\frac{\partial\rho_{t}}{\partial
t}=-\nabla\mathcal{E}(\rho_{t})=-M(\rho_{t})^{-1}\frac{\delta\mathcal{E}}{\delta\rho}\Bigr{|}_{\rho=\rho_{t}},$

in which the right hand side is an element in $T_{\rho_{t}}\mathcal{P}.$

e gradient flow can also be interpreted from the proximal perspective. Given
the metric $g$ and the geodesic distance function under this metric,
$\mathcal{D}$, the proximal point method uses the following iteration

(15)
$\displaystyle\mathfrak{r}_{{n+1}}=\operatorname*{arg\,min}_{\rho\in\mathcal{P}}\Bigl{(}\mathcal{E}(\rho)+\frac{1}{2\Delta
t}\mathcal{D}(\rho,\mathfrak{r}_{n})^{2}\Bigr{)}$

to minimize the energy functional $\mathcal{E}$ in density space
$\mathcal{P}$. When $\Delta t$ is small it is natural to seek
$\mathfrak{r}_{n+1}=\mathfrak{r}_{n}+\Delta t\sigma_{n}$ and note that,
invoking the approximation implied by Eq. 13,

$\displaystyle\sigma_{{n}}\approx\operatorname*{arg\,min}_{\sigma\in
T_{\mathfrak{r}_{n}}\mathcal{P}}\Bigl{(}\mathcal{E}(\mathfrak{r}_{n}+\Delta
t\sigma)+\frac{1}{2}{\Delta
t}\langle\sigma,M(\mathfrak{r}_{n})\sigma\rangle\Bigr{)}.$

To leading order in $\Delta t$, this expression is minimized by choosing

$\sigma_{n}\approx-M(\mathfrak{r}_{n})^{-1}\frac{\delta\mathcal{E}}{\delta\rho}\Bigr{|}_{\rho=\mathfrak{r}_{n}}.$

Letting $\mathfrak{r}_{n}\approx\rho_{n\Delta t}$, the formal continuous time
limit of the proximal algorithm leads to the corresponding gradient flow Eq.
14 [68].

For the specific choice (4) of KL divergence functional as the energy
functional, the first variation is given by (6). Thus the resulting gradient
flow has the form

(16) $\frac{\partial\rho_{t}}{\partial
t}=-M(\rho_{t})^{-1}\bigl{(}\log\rho_{t}-\log\rho_{\rm
post}-\int(\log\rho-\log\rho_{\rm post}){\rm d}\theta\bigr{)}.$

Most of the paper is focused on this gradient flow, for specific choices of
metric, mean-field models that have law governed by this gradient flow, or
restrictions of this gradient flow to variations in the manifold of Gaussians.

####  3.1.3 Affine Invariance

We now introduce the concept of affine invariance. Roughly speaking, affine
invariant gradient flows are invariant under any invertible affine
transformations of the density variables; as a consequence, the convergence
rate is independent of the affine transformation. It is thus natural to expect
that algorithms with this property have an advantage for sampling highly
anisotropic posteriors.

Let $\varphi:\theta\rightarrow\tilde{\theta}$ denote a diffeomorphism in
$\mathbb{R}^{N_{\theta}}$. When $\varphi(\theta)=A\theta+b$,
$A\in\mathbb{R}^{N_{\theta}\times N_{\theta}},b\in\mathbb{R}^{N_{\theta}}$ and
$A$ is invertible, then the diffeomorphism is an affine transformation.

######  Definition 3.1.

We define the pushforward operation $\\#$ for various objects as follows:

  * •

for density $\rho$, we write $\tilde{\rho}=\varphi\\#\rho$, which satisfies
$\tilde{\rho}(\tilde{\theta})=\rho(\varphi^{-1}(\tilde{\theta}))|\nabla_{\tilde{\theta}}\varphi^{-1}(\tilde{\theta})|$;

  * •

for tangent vector $\sigma\in T_{\rho}\mathcal{P}$, we have
$\tilde{\sigma}=\varphi\\#\sigma\in T_{\tilde{\rho}}\mathcal{P}$ which
satisfies

$\tilde{\sigma}(\tilde{\theta})=\sigma(\varphi^{-1}(\tilde{\theta}))|\nabla_{\tilde{\theta}}\varphi^{-1}(\tilde{\theta})|;$

  * •

for functional $\mathcal{E}$ on $\mathcal{P}$, we define
$\tilde{\mathcal{E}}=\varphi\\#\mathcal{E}$ via
$\tilde{\mathcal{E}}(\tilde{\rho})=\mathcal{E}(\varphi^{-1}\\#\tilde{\rho})$.

In the above definition, the forms of $\tilde{\rho}$ and $\tilde{\sigma}$ can
be derived by the standard change-of-variable formula [128].

Then we can define affine invariant gradient flow, metric, and mean-field
dynamics.

######  Definition 3.2 (Affine Invariant Gradient Flow).

Fix a Riemannian metric $g$ and the gradient operation $\nabla$ with respect
to this metric. Consider the gradient flow

$\frac{\partial\rho_{t}}{\partial t}=-\nabla\mathcal{E}(\rho_{t})$

and the affine transformation $\tilde{\theta}=\varphi(\theta)=A\theta+b$. Let
$\tilde{\rho}_{t}:=\varphi\\#\rho_{t}$ denote the distribution of
$\tilde{\theta}$ at time $t$ and set
$\tilde{\mathcal{E}}=\varphi\\#\mathcal{E}.$ The gradient flow is affine
invariant if

$\frac{\partial\tilde{\rho}_{t}}{\partial
t}=-\nabla\tilde{\mathcal{E}}(\tilde{\rho}_{t}).$

The key idea in the preceding definition is that, after the change of
variables, the dynamics of $\tilde{\rho}_{t}$ is itself a gradient flow, in
the same metric as the gradient flow in the original variables.

######  Definition 3.3 (Affine Invariant Metric).

Define the pull-back operator on Riemannian metric $g$ by

$(\varphi^{\\#}g)_{\rho}(\sigma_{1},\sigma_{2})=g_{\varphi\\#\rho}(\varphi\\#\sigma_{1},\varphi\\#\sigma_{2}),$

for any $\rho\in\mathcal{P}$ and $\sigma_{1},\sigma_{2}\in
T_{\rho}\mathcal{P}$. We say that Riemannian metric $g$ is affine invariant if
$\varphi^{\\#}g=g$ for any affine transformation $\varphi$.

The affine invariance of gradient flows is closely related to that of the
Riemannian metric:

######  Proposition 3.4.

The following two conditions are equivalent:

  1. 1.

the gradient flow under Riemannian metric $g$ is affine invariant for any
$\mathcal{E}$;

  2. 2.

the Riemannian metric $g$ is affine invariant.

We provide a proof for this proposition in Section B.2. Given this, it
suffices to focus on the affine invariance of the Riemannian metrics that we
consider in this paper; furthermore, we may modify them where needed to make
them affine invariant. Proposition 3.4, we consider the affine invariance
property to hold for any $\mathcal{E}$; the metric is independent of
$\mathcal{E}$. However, it is possible to choose a metric that depends on the
energy functional $\mathcal{E}$. An example of this is Newton’s method where
the Riemannian metric is given by the Hessian of the energy functional,
assuming it is positive definite; see the discussion of Newton’s flow on
probability space in [133]. Recall that our motivation for introducing the
affine invariance property is that algorithms with this property will, in
settings where an affine transformation removes anisotropy, have favorable
performance when sampling highly anisotropic posteriors. Our current
definition of affine invariance is tied to the energy functional $\mathcal{E}$
without direct reference to $\rho_{\rm post}$. Given $\mathcal{E}$, an affine
invariant gradient flow has the same convergence property when the energy
functional changes to $\mathcal{E}(\varphi^{-1}\\#\rho)$ where $\varphi$ is an
invertible affine transformation. To connect the transformation of the energy
functional to that of $\rho_{\rm post}$, we note that the KL divergence
satisfies the property

(17) $\mathcal{E}(\varphi^{-1}\\#\rho)={\rm KL}[\varphi^{-1}\\#\rho\|\rho_{\rm
post}]={\rm KL}[\rho\|\varphi\\#\rho_{\rm post}].$

Therefore, affine invariant gradient flows of the KL divergence have the same
convergence property when $\rho_{\rm post}$ changes to $\varphi\\#\rho_{\rm
post}$, for any invertible affine transformation $\varphi$. This suggests that
the flow will have favorable behaviour for sampling highly anisotropic
posteriors provided, under at least one affine transformation, the anisotropy
is removed.

####  3.1.4 Mean-Field Dynamics

Approximating the dynamics implied by Eq. 14 is often a substantial task. One
approach is to identify a mean-field stochastic dynamical system, with state
space $\mathbb{R}^{N_{\theta}}$, defined so that its law is given by Eq. 14.
For example, we may introduce the Itô SDE

(18) $\displaystyle\mathrm{d}\theta_{t}=f(\theta_{t};\rho_{t},\rho_{\rm
post})\mathrm{d}t+h(\theta_{t};\rho_{t},\rho_{\rm post}){\rm d}W_{t},$

where $W_{t}\in\mathbb{R}^{N_{\theta}}$ is a standard Brownian motion. Because
the drift
$f:\mathbb{R}^{N_{\theta}}\times\mathbb{R}\times\mathbb{R}\rightarrow\mathbb{R}^{N_{\theta}}$
and diffusion coefficient
$h:\mathbb{R}^{N_{\theta}}\times\mathbb{R}\times\mathbb{R}\rightarrow\mathbb{R}^{N_{\theta}\times
N_{\theta}}$ are evaluated at $\rho_{t}$, the density of $\theta_{t}$ itself,
this is a mean-field model.

The density is governed by a nonlinear Fokker-Planck equation

(19) $\frac{\partial\rho_{t}}{\partial
t}=-\nabla_{\theta}\cdot(\rho_{t}f)+\frac{1}{2}\nabla_{\theta}\cdot\bigl{(}\nabla_{\theta}\cdot(hh^{T}\rho_{t})\bigr{)}.$

By choice of $f,h$ it may be possible to ensure that Eq. 19 coincides with Eq.
14. Then an interacting particle system can be used to approximate Eq. 18,
generating an empirical measure which approximates $\rho_{t}.$

As the affine invariance property is important for gradient flows, we also
need to study this property for mean-field dynamics that are used to
approximate these flows.

######  Definition 3.5 (Affine Invariant Mean-Field Dynamics).

Consider the mean-field dynamics Eq. 18 and the affine transformation
$\tilde{\theta}=\varphi(\theta)=A\theta+b$. The mean-field dynamics is affine
invariant, when

(20a) $\displaystyle\begin{split}&Af(\theta;\rho,\rho_{\rm
post})=f(\varphi(\theta);\varphi\\#\rho,\varphi\\#\rho_{\rm
post}),\end{split}$ (20b) $\displaystyle\begin{split}&Ah(\theta;\rho,\rho_{\rm
post})=h(\varphi(\theta);\varphi\\#\rho,\varphi\\#\rho_{\rm
post}),\end{split}$

for any affine transformation $\varphi$. This implies that
$\tilde{\theta}_{t}=\varphi(\theta_{t})$ satisfies a SDE of the same form as
Eq. 18:

(21)
$\displaystyle\mathrm{d}\tilde{\theta}_{t}=f(\tilde{\theta}_{t};\tilde{\rho}_{t},\tilde{\rho}_{\rm
post})\mathrm{d}t+h(\tilde{\theta}_{t};\tilde{\rho}_{t},\tilde{\rho}_{\rm
post}){\rm d}W_{t},$

where $\tilde{\rho}_{t}=\varphi\\#\rho_{t}$ and $\tilde{\rho}_{\rm
post}=\varphi\\#\rho_{\rm post}$ by definition.

If we use this definition, then mean-field dynamics of affine invariant
gradient flows need not be affine invariant, since there may be different
$f,h$ giving rise to the same flow – equivalence classes. For the affine
invariance of the corresponding mean-field dynamics, we have the following
proposition, noting that the condition on the energy is satisfied for Eq. 4 by
Eq. 17.

######  Proposition 3.6.

Consider the energy functional $\mathcal{E}(\rho;\rho_{\rm post})$, making
explicit the dependence on $\rho_{\rm post}$, and assume that
$\mathcal{E}(\varphi^{-1}\\#\rho;\rho_{\rm
post})=\mathcal{E}(\rho;\varphi\\#\rho_{\rm post})$ holds. Then, corresponding
to any any affine invariant gradient flow of $\mathcal{E}$ there is a mean-
field dynamics of the form (18) which is affine invariant.

The proof of this proposition may be found in Section B.3.

As a consequence, Proposition 3.6 unifies the affine invariance property of
the gradient flow in probability space and the corresponding mean-field
dynamics. We note, however, that the mean-field dynamics is not unique and we
only prove the existence of one choice (amongst many) which is affine
invariant. In our later discussions, we will give some specific construction
of the mean-field dynamics for several gradient flows, and show that they are
indeed affine invariant.

e condition assumed in Proposition 3.6 indicates that the pushforward of the
functional $\mathcal{E}$ (See Definition 3.1) satisfies

(22)
$\tilde{\mathcal{E}}(\tilde{\rho})=\tilde{\mathcal{E}}(\tilde{\rho};\rho_{\rm
post})=\mathcal{E}(\varphi^{-1}\\#\tilde{\rho};\rho_{\rm
post})=\mathcal{E}(\tilde{\rho};\varphi\\#\rho_{\rm
post})=\mathcal{E}(\tilde{\rho};\tilde{\rho}_{\rm post}).$

Thus, this condition allows to connect the affine invariance defined via the
transformation of energy functional and via the transformation of the target
posterior distributions, as explained in
LABEL:remark:_transform_E_and_transform_posterior. Beyond the KL divergence
(see Eq. 17), the condition is also satisfied by various widely used energy
functionals, such as the Hellinger distance and the chi-squared divergence.

###  3.2 Fisher-Rao Gradient Flow

####  3.2.1 Metric

The Fisher-Rao Riemannian metric is

$g_{\rho}^{\mathrm{FR}}(\sigma_{1},\sigma_{2})=\int\frac{\sigma_{1}\sigma_{2}}{\rho}\mathrm{d}\theta.$

Writing tangent vectors on a multiplicative scale, by setting
$\sigma=\rho\psi_{\sigma}$, we see that this metric may be written as

$g_{\rho}^{\mathrm{FR}}(\sigma_{1},\sigma_{2})=\int\psi_{\sigma_{1}}\psi_{\sigma_{2}}\rho\mathrm{d}\theta,$

and hence that in the $\psi_{\sigma}$ variable the metric is described via the
$L_{\rho}^{2}$ inner-product.

The Fisher-Rao metric tensor $M^{\mathrm{FR}}(\rho)$ associated to
$g_{\rho}^{\mathrm{FR}}$ satisfies

(23a) $\displaystyle
M^{\mathrm{FR}}(\rho)\sigma=\psi_{\sigma}-\int\psi_{\sigma}{\rm
d}\theta,\quad\forall~{}\sigma\in T_{\rho}\mathcal{P}$ (23b) $\displaystyle
M^{\mathrm{FR}}(\rho)^{-1}\psi=\rho(\psi-\mathbb{E}_{\rho}[\psi]),\quad\forall~{}\psi\in
T_{\rho}^{*}\mathcal{P}.$

The corresponding geodesic distance $\mathcal{D}^{\rm
FR}:\mathcal{P}\times\mathcal{P}\to\mathbb{R}^{+}$ is

(24) $\mathcal{D}^{\rm
FR}(\rho_{A},\rho_{B})^{2}=\inf_{\rho_{t}}\Bigl{\\{}\int_{0}^{1}\mathrm{d}t\int\frac{|\partial_{t}{\rho}_{t}|^{2}}{\rho_{t}}\mathrm{d}\theta:\,\rho_{0}=\rho_{A},\rho_{1}=\rho_{B}\Bigr{\\}}.$

If we do not restrict the distributions $\rho_{t}$ to be on the probability
space (i.e., we allow they to have any positive mass), then by using the
relation

$\frac{|\partial_{t}{\rho}_{t}|^{2}}{\rho_{t}}=4\Bigl{|}\frac{\mathrm{d}}{\mathrm{d}t}\sqrt{\rho_{t}}\Bigr{|}^{2}$

and the Cauchy-Schwarz inequality, we can solve the optimization problem in
Eq. 24 explicitly. The optimal objective value will be
$4\int|\sqrt{\rho_{A}}-\sqrt{\rho_{B}}|^{2}\mathrm{d}\theta.$ This is (up to a
constant scaling) the Hellinger distance [56].

On the other hand, if we constrain $\rho_{t}$ to be on the probability space,
then the geodesic distance will be (up to a constant scaling) the spherical
Hellinger distance:

$\mathcal{D}^{\rm
FR}(\rho_{A},\rho_{B})^{2}=4\operatorname{arccos}^{2}\left(\int\sqrt{\rho_{A}}\sqrt{\rho_{B}}{\rm
d}\theta\right).$

For more discussions, see [61, 81, 94]. In view of this relation, Fisher-Rao
gradient flows are sometimes referred to as spherical Hellinger gradient flows
in the literature [87, 94].

####  3.2.2 Flow Equation

From (6) and (23b) we see that the Fisher-Rao gradient flow of the KL
divergence is

(25) $\displaystyle\frac{\partial\rho_{t}}{\partial t}=$
$\displaystyle-M^{\mathrm{FR}}(\rho_{t})^{-1}\frac{\delta\mathcal{E}}{\delta\rho}\Bigr{|}_{\rho=\rho_{t}},$
$\displaystyle=$ $\displaystyle\rho_{t}\bigl{(}\log\rho_{\rm
post}-\log\rho_{t}\bigr{)}-\rho_{t}\mathbb{E}_{\rho_{t}}[\log\rho_{\rm
post}-\log\rho_{t}].$

e gradient flow Eq. 25 in probability space has the form typical of a mean-
field model which is a birth-death process – it is possible to create and kill
particles to sample this process. However, the support of the empirical
distribution using this algorithm never increases during evolution. To address
this issue, the work [93] added Langevin diffusion to the birth-death process,
resulting in what they term the Wasserstein-Fisher-Rao gradient flow.
Alternatively, the authors in [87] utilized a Markov chain kernel and MCMC to
sample the birth death dynamics arising from the Fisher-Rao gradient flow,
using the chi-squared divergence [85] instead of Eq. 4.

When the target distribution (1) arises from a Bayesian inverse problem it may
be written in the form

(26) $\displaystyle\rho_{\rm
post}(\theta)\propto\exp(-\Phi(\theta))\rho_{0}(\theta);$

function $\Phi:\mathbb{R}^{N_{\theta}}\to\mathbb{R}_{+}$ is the negative log
likelihood and $\rho_{0}$ is the prior. In this context it is interesting to
consider

(27)
$\displaystyle\mathcal{E}(\rho)=\int\rho(\theta)\,\Phi(\theta)\,\mathrm{d}\theta$

with associated Fisher-Rao gradient flow

(28) $\displaystyle\frac{\partial\rho_{t}}{\partial
t}=-\rho_{t}\left(\Phi-\mathbb{E}_{\rho_{t}}[\Phi]\right).$

It may be shown that the density $\rho_{t}$ is explicitly given by

(29)
$\displaystyle\rho_{t}(\theta)=\frac{\exp(-t\Phi(\theta))\rho_{0}(\theta)}{\mathbb{E}_{\rho_{0}}[\exp(-t\Phi)]}.$

Hence we recover (26) at $t=1$. This observation is at the heart of homotopy-
based approaches to Bayesian inference [40], leading to methods based on
particle filters; the link to an evolution equation for $\rho_{t}$ is employed
and made explicit in various other approaches to filtering [39, 112]. See [35,
20] for overviews. Such Fisher-Rao gradient flow structure for Bayes updates
has also been identified in the context of filtering in [82, 61, 60].

We also note that, by letting $t\to\infty$ one finds that

(30) $\displaystyle\lim_{t\to\infty}\rho_{t}=\delta_{\theta^{\ast}}$

where $\theta^{\ast}$ denotes the (assumed unique) minimizer of $\Phi$, in the
support of $\rho_{0}$, and $\delta_{\theta^{\ast}}$ denotes the Dirac delta
function centred at $\theta^{\ast}$.

####  3.2.3 Affine Invariance

The Fisher-Rao metric is affine invariant. One may understand this property
through the affine invariance property of Newton’s method when the energy
functional is the KL divergence. To see this note that, from Eq. 6, the
Hessian of $\mathcal{E}$ given by Eq. 4 has the form

(31)
$\displaystyle\frac{\delta^{2}\mathcal{E}(\rho)}{\delta\rho^{2}}=\frac{\delta^{2}\mathrm{KL}[\rho\|\rho_{\rm
post}]}{\delta\rho^{2}}=\frac{1}{\rho}=M^{\text{FR}}(\rho).$

Therefore, the Fisher-Rao gradient flow of the KL divergence behaves like
Newton’s method, which is affine invariant. In fact the Fisher-Rao metric is
affine invariant under any diffeomorphism of the parameter space, not just
invertible affine transformations. Indeed, it is the only metric, up to
constant, that satisfies this strong invariance property [23, 5, 9].

####  3.2.4 Mean-Field Dynamics

The Fisher-Rao gradient flow (25) in $\rho_{t}$ can be realized as the law of
a mean-field ordinary differential equation in $\theta_{t}$

(32) $\frac{\mathrm{d}\theta_{t}}{\mathrm{d}t}=f(\theta_{t};\rho_{t},\rho_{\rm
post}).$

Writing the nonlinear Liouville equation associated with this model and
equating it to Eq. 25 shows that drift $f$ satisfies

(33) $-\nabla_{\theta}\cdot(\rho_{t}f)=\rho_{t}\bigl{(}\log\rho_{\rm
post}-\log\rho_{t}\bigr{)}-\rho_{t}\mathbb{E}_{\rho_{t}}[\log\rho_{\rm
post}-\log\rho_{t}].$

Note that $f$ is not uniquely determined by (33). Writing $f$ as a gradient of
a potential, with respect to $\theta$, shows that the potential satisfies a
linear elliptic PDE, and under some conditions this will have a unique
solution; but there will be other choices of $f$ which are not a pure
gradient, leading to nonuniqueness.

By Proposition 3.6, for affine invariant gradient flows, there exist mean-
field dynamics (i.e., via choosing certain $f$ in (33)) that are affine
invariant. Here, we construct a specific class of $f$ that leads to affine
invariant mean-field dynamics for the Fisher-Rao gradient flow.

First, we introduce a matrix valued function:
$P:\mathbb{R}^{N_{\theta}}\times\mathcal{P}\to\mathbb{R}^{N_{\theta}\times
N_{\theta}}_{\succ 0},$ where the output space is the cone of positive-
definite symmetric matrices; we refer to matrices such as $P$ as
_preconditioners_ throughout this paper. Then, the following proposition shows
that the choice of $f=P(\theta,\rho_{t})\nabla\phi$ leads to affine invariance
of the dynamics, under certain condition on $P$. The proof can be found in
Section B.4.

######  Proposition 3.7.

Consider the invertible affine transformation
$\tilde{\theta}=\varphi(\theta)=A\theta+b$ and correspondingly
$\tilde{\rho}=\varphi\\#\rho$. Assume that the preconditioning matrix
satisfies

(34) $\displaystyle P(\tilde{\theta},\tilde{\rho})=AP(\theta,\rho)A^{T}.$

Assume, furthermore, that the solution $\phi(\theta;\rho,\rho_{\rm post})$ of
the equation

(35) $-\nabla_{\theta}\cdot(\rho
P\nabla_{\theta}\phi)=\rho\bigl{(}\log\rho_{\rm
post}-\log\rho\bigr{)}-\rho\mathbb{E}_{\rho}[\log\rho_{\rm post}-\log\rho]$

exists, is unique (up to constants) and belongs to
$C^{2}(\mathbb{R}^{N_{\theta}})$, for any $\rho\in\mathcal{P}$. Then, the
corresponding mean-field equation Eq. 32 with $f=P\nabla_{\theta}\phi$ is
affine invariant.

re generally, given any alternative functional $\mathcal{E}$, such as (27),
one can define affine invariant mean-field ordinary differential equations of
the form (32) with drift $f=P\nabla_{\theta}\phi$ and potential $\phi$
satisfying the equation

(36) $\displaystyle-\nabla_{\theta}\cdot(\rho
P\nabla_{\theta}\phi)=\rho\Bigl{(}\frac{\delta\mathcal{E}}{\delta\rho}-\mathbb{E}_{\rho}\Bigl{[}\frac{\delta\mathcal{E}}{\delta\rho}\Bigr{]}\Bigr{)}.$

In addition to the above choice of mean field models, birth-death type mean
field dynamics have also been used to simulate Fisher-Rao gradient flows for
sampling; see [93, 94].

###  3.3 Wasserstein Gradient Flow

####  3.3.1 Metric

Generalizing the relationship between $\sigma$ and $\psi_{\sigma}$ introduced
in the Fisher-Rao context, we define $\psi_{\sigma}$ to be solution of the PDE

(37) $-\nabla_{\theta}\cdot(\rho\nabla_{\theta}\psi_{\sigma})=\sigma.$

This definition requires specification of function spaces to ensure unique
invertibility of the divergence form elliptic operator. One then defines the
Wasserstein metric tensor and its inverse by

Elementary manipulations show that the corresponding Riemannian metric is
given by

(39a) $\displaystyle g_{\rho}^{\mathrm{W}}(\sigma_{1},\sigma_{2})$
$\displaystyle=\langle\sigma_{1},$ (39b) $\displaystyle=\langle$ (39c)
$\displaystyle=\int\rho(\theta)\nabla_{\theta}\psi_{\sigma_{1}}(\theta)^{T}\nabla_{\theta}\psi_{\sigma_{2}}(\theta)\mathrm{d}\theta.$

Here $g_{\rho}^{\mathrm{W}}$ is positive-definite and hence a valid metric. It
is termed the Wasserstein Riemannian metric throughout this paper.

e Wasserstein Riemannian metric has a transport interpretation. To understand
this fix $\sigma\in T_{\rho}\mathcal{P}$ and consider the family of _velocity
fields_ $v$ related to $\sigma$ via the constraint
$\sigma=-\nabla_{\theta}\cdot(\rho v)$. Then define
$v_{\sigma}=\operatorname*{arg\,min}_{v}\int\rho|v|^{2}$ in which the
minimization is over all $v$ satisfying the constraint. A formal Lagrange
multiplier argument can be used to deduce that
$v_{\sigma}=\nabla_{\theta}\psi_{\sigma}$ for some $\psi_{\sigma}$. This
motivates the relationship appearing in Eq. 37 as well as the form of the
Wasserstein Riemannian metric appearing in Eq. 39 which may then be viewed as
measuring the kinetic energy $\int\rho|v_{\sigma}|^{2}\mathrm{d}\theta$. We
emphasize that, for ease of understanding, our discussion on the Riemannian
structure of the Wasserstein metric is purely formal; for rigorous treatments,
the reader can consult [3].

To further develop the preceding discussion, consider the Liouville equation
for dynamical system in $\mathbb{R}^{N_{\theta}}$ driven by vector field
$v_{\sigma}:=\nabla_{\theta}\psi_{\sigma}$. Let $\rho_{A},\rho_{B}$ be two
elements in $\mathcal{P}$ and let $\rho_{t}$ be a path in time governed by
this Liouville equation, and satisfying the boundary conditions
$\rho_{0}=\rho_{A},\rho_{1}=\rho_{B}.$ Then

(40) $\frac{\partial\rho_{t}}{\partial
t}+\nabla_{\theta}\cdot(\rho_{t}\nabla_{\theta}\psi_{t})=0,\rho_{0}=\rho_{A},\rho_{1}=\rho_{B}.$

With this equation, we can write the geodesic distance $\mathcal{D}^{\rm
W}:\mathcal{P}\times\mathcal{P}\to\mathbb{R}^{+}$ as:

(41) $\begin{split}\mathcal{D}^{\rm
W}(\rho_{A},\rho_{B})^{2}&=\inf_{\rho_{t}}\Bigl{\\{}\int_{0}^{1}g_{\rho_{t}}^{\rm
W}(\partial_{t}{\rho}_{t},\partial_{t}{\rho}_{t})\mathrm{d}t:\,\rho_{0}=\rho_{A},\rho_{1}=\rho_{B}\Bigr{\\}}\\\
&=\inf_{\psi_{t}\in\mathsf{L}}\Bigl{\\{}\int_{0}^{1}\mathrm{d}t\int\rho_{t}|\nabla_{\theta}\psi_{t}|^{2}\mathrm{d}\theta\Bigr{\\}},\end{split}$

where $\mathsf{L}$ is the set of time-dependent potentials $\psi_{t}$ such
that equation Eq. 40 holds. This is the celebrated Benamou-Brenier formula for
the 2-Wasserstein distance [12].

####  3.3.2 Flow Equation

The Wasserstein gradient flow of the KL divergence is

(42) $\displaystyle\frac{\partial\rho_{t}}{\partial t}$ $\displaystyle=-$
$\displaystyle=\nabla_{\theta}\cdot\big{(}\rho_{t}(\nabla_{\theta}\log\rho_{t}-\nabla_{\theta}\log\rho_{\rm
post})\big{)}$
$\displaystyle=-\nabla_{\theta}\cdot(\rho_{t}\nabla_{\theta}\log\rho_{\mathrm{post}})+\Delta_{\theta}\rho_{t}.$

This is simply the Fokker-Planck equation for the Langevin dynamics

(43)
$\mathrm{d}\theta_{t}=\nabla_{\theta}\log\rho_{\mathrm{post}}(\theta)\mathrm{d}t+\sqrt{2}\mathrm{d}W_{t},$

where $W_{t}\in\mathbb{R}^{N_{\theta}}$ is a standard Brownian motion. This is
a trivial mean-field model of the form Eq. 18 in the sense that there is no
dependence on the density $\rho_{t}$ associated with the law of $\theta.$

note that elliptic equations defining certain potentials arise in the context
of both the Fisher-Rao as well as the Wasserstein metric. However, while (37)
appears in the definition of the Wasserstein metric only, solving (36) is
required for obtaining the mean-field equations (32) in the Fisher-Rao
setting. Returning to the cost functional (27), we find that the associated
Wasserstein gradient mean-field dynamics simply reduces to gradient descent

(44) $\displaystyle{\rm d}\theta_{t}=-\nabla_{\theta}\Phi(\theta){\rm d}t$

while the associated Fisher-Rao mean-field equations are more complex and
linked to Bayesian inference as discussed earlier in LABEL:rem:homotopy.

####  3.3.3 Affine Invariance

The Wasserstein Riemannian metric Eq. 39 is not affine invariant. Hence, in
this subsection, we introduce an affine invariant modification to the
Wasserstein metric. To this end, we consider preconditioner
$P:\mathbb{R}^{N_{\theta}}\times\mathcal{P}\to\mathbb{R}^{N_{\theta}\times
N_{\theta}}_{\succ 0},$ where the output space is the cone of positive-
definite symmetric matrices.

We generalize Eq. 37 and let $\psi_{\sigma}$ solve the PDE

(45) $-\nabla_{\theta}\cdot(\rho
P(\theta,\rho)\nabla_{\theta}\psi_{\sigma})=\sigma,$

again noting that specification of function spaces is needed to ensure unique
invertibility of the divergence form elliptic operator (see Proposition 3.7
where similar considerations arise). We may then generalize the metric tensor
in Eq. 38 to obtain and inverse given by

Manipulations similar to use in Eq. 39, but using , show that

$\displaystyle g_{\rho}^{\mathrm{AIW}}(\sigma_{1},\sigma_{2})$
$\displaystyle=\langle\sigma_{1},$ $\displaystyle=\langle$
$\displaystyle=\int\rho(\theta)\nabla_{\theta}\psi_{\sigma_{1}}(\theta)^{T}P(\theta,\rho)\nabla_{\theta}\psi_{\sigma_{2}}(\theta)\mathrm{d}\theta.$

It follows that $g_{\rho}^{\mathrm{AIW}}$ is positive-definite and hence a
valid metric tensor. We have the following proposition to guarantee this
metric tensor is affine invariant:

######  Proposition 3.8.

Under the assumption on $P$ given in Proposition 3.7, leading to (34), the
metric corresponding to $M^{\mathrm{AIW}}$ is affine invariant. Consequently,
the associated gradient flow of the KL divergence, namely

(47) $\displaystyle\frac{\partial\rho_{t}(\theta)}{\partial t}$
$\displaystyle=\nabla_{\theta}\cdot\Bigl{(}\rho_{t}P(\theta,\rho_{t})(\nabla_{\theta}\log\rho_{t}-\nabla_{\theta}\log\rho_{\rm
post})\Bigr{)},$

is affine invariant.

The proof of this proposition is provided in Section B.5. Henceforth we refer
to $M^{\mathrm{AIW}}$ satisfying the condition of the preceding proposition as
an affine invariant Wasserstein metric tensor.

####  3.3.4 Mean-Field Dynamics

As discussed in relation to the topic of affine invariance in Section 3.1.4,
mean-field models with a given law are not unique. In the specific context of
the Wasserstein gradient flow this suggests looking beyond Eq. 43 for a mean-
field model with governing law given by Eq. 42. This can be achieved as
follows [123]. Fix arbitrary
$h:\mathbb{R}^{N_{\theta}}\times\mathbb{R}\rightarrow\mathbb{R}^{N_{\theta}\times
N_{\theta}}$, define
$D(\theta,\rho)=\frac{1}{2}h(\theta,\rho)h(\theta,\rho)^{T}$ and choose
$d(\theta,\rho)=\nabla_{\theta}\cdot D(\theta,\rho)$. Then, for any $h$,
consider the SDE

(48) $\displaystyle\mathrm{d}\theta_{t}=\Bigl{(}\nabla_{\theta}\log\rho_{\rm
post}(\theta_{t})+\bigl{(}D(\theta_{t},\rho_{t})-I\bigr{)}\nabla_{\theta}\log\rho_{t}(\theta_{t})-d(\theta_{t},\rho_{t})\Bigr{)}\mathrm{d}t+h(\theta_{t},\rho_{t}){\rm
d}W_{t},$

When $h=\sqrt{2I}$ we recover Eq. 43. When this condition does not hold, so
that $D(\theta,\rho_{t})\neq I$, the equation requires knowledge of the score
function $\nabla_{\theta}\log\rho_{t}(\theta_{t})$; and particle methods to
approximate Eq. 48 will require estimates of the score; see [96, 132, 119,
17]. See also [120] and references therein for discussion of score estimation.

We now apply similar considerations to the preconditioned Wasserstein gradient
flow Eq. 47. Employing the same choices of $d$ and $D$ from $h$ as in the
unpreconditioned case we obtain the following mean-field evolution equation:

(49) $\displaystyle\mathrm{d}\theta_{t}$
$\displaystyle=P(\theta_{t},\rho_{t})\nabla_{\theta}\log\rho_{\rm
post}(\theta_{t})\mathrm{d}t$
$\displaystyle\quad\quad\quad+\Bigl{(}\bigl{(}D(\theta_{t},\rho_{t})-P(\theta_{t},\rho_{t})\bigr{)}\nabla_{\theta}\log\rho_{t}(\theta_{t})-d(\theta_{t},\rho_{t})\Bigr{)}\mathrm{d}t$
$\displaystyle\qquad\qquad\qquad\qquad\qquad\qquad+h(\theta_{t},\rho_{t}){\rm
d}W_{t}.$

For this specific mean-field equation Eq. 49, we can also establish affine
invariance; see the following proposition and its proof in Section B.6.

######  Proposition 3.9.

The mean-field equation Eq. 49 is affine invariant under the assumption on the
preconditioner $P$ given in Proposition 3.7, leading to (34), and the
assumptions on $h$ given in LABEL:eq:aMFD-AI-sigma.

In particular, let $C(\rho)$ denote the covariance matrix of $\rho.$ If we
take $P(\theta,\rho)=C(\rho)$ then we recover the affine invariant Kalman-
Wasserstein metric introduced in [54, 55]. Furthermore, then making the choice
of $h(\theta,\rho)=\sqrt{2C(\rho)}$ leads to the following affine invariant
overdamped Langevin equation, also introduced in [54, 55]:

(50)
$\displaystyle\mathrm{d}\theta_{t}=C(\rho_{t})\nabla_{\theta_{t}}\log\rho_{\rm
post}(\theta_{t})\mathrm{d}t+\sqrt{2C(\rho_{t})}{\rm d}W_{t}.$

Comparison with Eq. 43 demonstrates that it is a preconditioned version of the
standard overdamped Langevin equation.

###  3.4 Stein Gradient Flow

####  3.4.1 Metric

Generalizing Eq. 37 we let $\psi_{\sigma}$ solve the integro-partial
differential equation

(51)
$-\nabla_{\theta}\cdot\Bigl{(}\rho(\theta)\int\kappa(\theta,\theta^{\prime},\rho)\rho(\theta^{\prime})\nabla_{\theta^{\prime}}\psi_{\sigma}(\theta^{\prime})\mathrm{d}\theta^{\prime}\Bigr{)}=\sigma(\theta).$

As before definition of function space setting is required to ensure that this
equation is uniquely solvable. Now define the Stein metric tensor , and its
inverse, as follows:

Computations analogous to those shown in Eq. 39 show that the Stein Riemannian
metric implied by metric tensor is given by

(53a) $\displaystyle g_{\rho}^{\mathrm{S}}(\sigma_{1},\sigma_{2})$
$\displaystyle=\langle\sigma_{1},$ (53b)
$\displaystyle=\int\int\kappa(\theta,\theta^{\prime},\rho)\rho(\theta)\nabla_{\theta}\psi_{\sigma_{1}}(\theta)^{T}\nabla_{\theta^{\prime}}\psi_{\sigma_{2}}(\theta^{\prime})\rho(\theta^{\prime})\mathrm{d}\theta\mathrm{d}\theta^{\prime}.$

in the Wasserstein setting, the Stein Riemannian metric [88] also has a
transport interpretation. The Stein metric identifies, for each $\sigma\in
T_{\rho}\mathcal{P}$, the set of velocity fields $v$ satisfying the constraint
$\sigma=-\nabla_{\theta}\cdot(\rho v)$. Then
$v_{\sigma}=\operatorname*{arg\,min}_{v}\|v\|_{\mathcal{H}_{\kappa}}^{2}$,
with minimization over all $v$ satisfying the constraint, and where
$\mathcal{H}_{\kappa}$ is a Reproducing Kernel Hilbert Space (RKHS) with
kernel $\kappa$. A formal Lagrangian multiplier argument shows that

$v_{\sigma}=\int\kappa(\theta,\theta^{\prime},\rho)\rho(\theta^{\prime})\nabla_{\theta^{\prime}}\psi_{\sigma}(\theta^{\prime})\mathrm{d}\theta^{\prime}$

for some $\psi_{\sigma}.$ The Stein metric measures this transport change via
the RKHS norm $\|v_{\sigma}\|_{\mathcal{H}_{\kappa}}^{2}$, leading to the
interpretation that the Stein Riemannian metric can be written in the form

$g_{\rho}^{\mathrm{S}}(\sigma_{1},\sigma_{2})=\langle
v_{\sigma_{1}},v_{\sigma_{2}}\rangle_{\mathcal{H}_{\kappa}}.$

Analogously to Eq. 40, for any $\rho_{A},\rho_{B}\in\mathcal{P}$, we may write
a path to connect these end points and defined by

(54) $\frac{\partial\rho_{t}}{\partial
t}+\nabla_{\theta}\cdot\Bigl{(}\rho_{t}\int\kappa(\theta,\theta^{\prime},\rho_{t})\rho_{t}(\theta^{\prime})\nabla_{\theta^{\prime}}\psi_{t}(\theta^{\prime})\mathrm{d}\theta^{\prime}\Bigr{)}=0,\rho_{0}=\rho_{A},\rho_{1}=\rho_{B}.$

The corresponding geodesic distance $\mathcal{D}^{\rm
S}:\mathcal{P}\times\mathcal{P}\to\mathbb{R}^{+}$ is

(55) $\begin{split}\mathcal{D}^{\rm
S}(\rho_{A},\rho_{B})^{2}&=\inf_{\rho_{t}}\Bigl{\\{}\int_{0}^{1}g_{\rho_{t}}^{\rm
S}(\partial_{t}{\rho}_{t},\partial_{t}{\rho}_{t})\mathrm{d}t:\,\rho_{0}=\rho_{A},\rho_{1}=\rho_{B}\Bigr{\\}}\\\
&=\inf_{\psi_{t}\in\mathsf{L}}\Bigl{\\{}\int_{0}^{1}\mathrm{d}t\int\int\kappa(\theta,\theta^{\prime},\rho_{t})\rho_{t}(\theta)\nabla_{\theta}\psi_{t}(\theta)\cdot\nabla_{\theta^{\prime}}\psi_{t}(\theta^{\prime})\rho_{t}(\theta^{\prime})\mathrm{d}\theta\mathrm{d}\theta^{\prime}\Bigr{\\}},\end{split}$

where $\mathsf{L}$ is the set of time-dependent potentials $\psi_{t}$ such
that equation Eq. 54 holds.

####  3.4.2 Flow Equation

The Stein variational gradient flow is

(56) $\displaystyle\frac{\partial\rho_{t}(\theta)}{\partial t}$
$\displaystyle=-\Bigl{(}$
$\displaystyle=\nabla_{\theta}\cdot\Bigl{(}\rho_{t}(\theta)\int\kappa(\theta,\theta^{\prime},\rho_{t})\rho_{t}(\theta^{\prime})\nabla_{\theta^{\prime}}\bigl{(}\log\rho_{t}(\theta^{\prime})-\log\rho_{\rm
post}(\theta^{\prime})\bigr{)}\mathrm{d}\theta^{\prime}\Bigr{)}.$

####  3.4.3 Affine Invariance

The Stein metric Eq. 52 is not affine invariant. To address this, in this
subsection we introduce an affine invariant modification. The generalization
is similar to that undertaken to obtain an affine invariant version of the
Wasserstein metric and so we will make the presentation brief. We define

so that for any $\psi\in T_{\rho}^{*}\mathcal{P}$, it holds that Here
$\kappa:\mathbb{R}^{N_{\theta}}\times\mathbb{R}^{N_{\theta}}\rightarrow\mathbb{R}$
is a positive definite kernel and we factorize the preconditioner
$P:\mathbb{R}^{N_{\theta}}\times\mathbb{R}^{N_{\theta}}\times\mathcal{P}\to\mathbb{R}^{N_{\theta}\times
N_{\theta}}$, which can be written in the form
$P(\theta,\theta^{\prime},\rho)=L(\theta,\rho)L(\theta^{\prime},\rho)^{T}.$
With this is hand it follows that

$\displaystyle\langle\psi,$ $\displaystyle=$
$\displaystyle\int\int\kappa(\theta,\theta^{\prime},\rho)\rho(\theta)\left(L(\theta,\rho)^{T}\nabla_{\theta}\psi(\theta)\right)^{T}\left(L(\theta^{\prime},\rho)^{T}\nabla_{\theta^{\prime}}\psi(\theta^{\prime})\right)\rho(\theta^{\prime})\mathrm{d}\theta\mathrm{d}\theta^{\prime}\geq
0$

and the resulting metric is well-defined. We have the following proposition to
guarantee this metric tensor is affine invariant:

######  Proposition 3.10.

Consider the invertible affine transformation
$\tilde{\theta}=\varphi(\theta)=A\theta+b$ and correspondingly
$\tilde{\rho}=\varphi\\#\rho$; moreover
$\tilde{\theta}^{\prime}=\varphi(\theta^{\prime})$. Assume that the
preconditioning matrix satisfies

$\kappa(\tilde{\theta},\tilde{\theta}^{\prime},\tilde{\rho})P(\tilde{\theta},\tilde{\theta}^{\prime},\tilde{\rho})=\kappa(\theta,\theta^{\prime},\rho)AP(\theta,\theta^{\prime},\rho)A^{T}.$

Then the metric corresponding to $M^{\mathrm{AIS}}$ is affine invariant.
Consequently, the associate gradient flow of the KL divergence, namely

(58) $\displaystyle\frac{\partial\rho_{t}(\theta)}{\partial t}$
$\displaystyle=\nabla_{\theta}\cdot(\mathsf{f})$ $\displaystyle\mathsf{f}$
$\displaystyle=\Bigl{(}\rho_{t}(\theta)\int\kappa(\theta,\theta^{\prime},\rho_{t})\rho_{t}(\theta^{\prime})P(\theta,\theta^{\prime},\rho_{t})\nabla_{\theta^{\prime}}\bigl{(}\log\rho_{t}(\theta^{\prime})-\log\rho_{\rm
post}(\theta^{\prime})\bigr{)}\mathrm{d}\theta^{\prime}\Bigr{)}$

is affine invariant.

The proof of this proposition is in Section B.7. Henceforth we refer to
$M^{\mathrm{AIS}}$ satisfying the condition of the preceding proposition as an
affine invariant Stein metric tensor. As an example, we can obtain an affine
invariant Stein metric by making the choices $P=C(\rho)$ and
$\kappa(\theta,\theta^{\prime},\rho)\propto\exp\bigl{\\{}-\frac{1}{2}(\theta-\theta^{\prime})^{T}C(\rho)^{-1}(\theta-\theta^{\prime})\bigr{\\}}$;
this set-up is considered in our numerical experiments; see Section 5.

####  3.4.4 Mean-Field Dynamics

The Stein gradient flow (56) has the following as a mean-field counterpart
[89, 88] in $\theta_{t}$ with the same law $\rho_{t}$:

(59) $\displaystyle\frac{\mathrm{d}\theta_{t}}{\mathrm{d}t}$
$\displaystyle=\int\kappa(\theta_{t},\theta^{\prime},\rho_{t})\rho_{t}(\theta^{\prime})\nabla_{\theta^{\prime}}\bigl{(}\log\rho_{\rm
post}(\theta^{\prime})-\log\rho_{t}(\theta^{\prime})\bigr{)}\mathrm{d}\theta^{\prime}$
$\displaystyle=\int\kappa(\theta_{t},\theta^{\prime},\rho_{t})\rho_{t}(\theta^{\prime})\nabla_{\theta^{\prime}}\log\rho_{\rm
post}(\theta^{\prime})+\rho_{t}(\theta^{\prime})\nabla_{\theta^{\prime}}\kappa(\theta_{t},\theta^{\prime},\rho_{t})\mathrm{d}\theta^{\prime}.$

Here, the second equality is obtained using integration by parts; it
facilitates an expression which avoids the score (gradient of log density
function of $\rho_{t}$). This is useful because, when implementing particle
methods, the resulting integral can then be approximated directly by Monte
Carlo methods.

Similarly, for the preconditioned Stein gradient flow Eq. 58, we can construct
the following mean-field equation:

(60) $\displaystyle\frac{\mathrm{d}\theta_{t}}{\mathrm{d}t}$
$\displaystyle=\int\Bigl{(}\kappa(\theta_{t},\theta^{\prime},\rho_{t})\rho_{t}(\theta^{\prime})P(\theta_{t},\theta^{\prime},\rho_{t})\nabla_{\theta^{\prime}}\log\rho_{\rm
post}(\theta^{\prime})$
$\displaystyle\qquad\qquad\qquad\qquad\qquad\qquad+\nabla_{\theta^{\prime}}\cdot(\kappa(\theta_{t},\theta^{\prime},\rho_{t})P(\theta_{t},\theta^{\prime},\rho_{t}))\rho_{t}(\theta^{\prime})\Bigr{)}\mathrm{d}\theta^{\prime}.$

The mean-field equation Eq. 60 is affine invariant; see the following
proposition and its proof in Section B.8.

######  Proposition 3.11.

The mean-field equation Eq. 60 is affine invariant under the assumption on the
preconditioner in Proposition 3.10.

###  3.5 Large-Time Asymptotic Convergence

In the three preceding subsections we studied gradient flows, under various
different metrics, of the energy $\mathcal{E}$ given in Eq. 4. In this
subsection, we study the convergence of these gradient flows both surveying
known, and adding new, results. In short, the convergence of the Fisher-Rao
gradient flow occurs at rate $\mathcal{O}(\exp(-t))$ and is hence problem
independent; this reflects the invariance of the metric under any
diffeomorphism. In contrast, the proven results for Wasserstein and Stein
gradient flows have convergence rates that depend on the problem, even after
being modified to be affine invariant. We note, however, that when $\rho_{\rm
post}$ is Gaussian, the affine invariant Wasserstein gradient flows also
achieve $\mathcal{O}(\exp(-t))$ [54, 55]. Numerical results illustrating and
complementing the analysis in this section may be found in Section 5.

####  3.5.1 Fisher-Rao Gradient Flow

We have the following proposition concerning large-time convergence of the
gradient flow:

######  Proposition 3.12.

Let $\rho_{t}$ solve the Fisher-Rao gradient flow Eq. 25. Assume also that
there exist constants $K,B>0$ such that the initial density $\rho_{0}$
satisfies

(61) $\displaystyle
e^{-K(1+|\theta|^{2})}\leq\frac{\rho_{0}(\theta)}{\rho_{\rm post}(\theta)}\leq
e^{K(1+|\theta|^{2})},$

and both $\rho_{0},\rho$ have bounded second moment

(62) $\displaystyle\int|\theta|^{2}\rho_{0}(\theta)\mathrm{d}\theta\leq
B,\quad\int|\theta|^{2}\rho(\theta)\mathrm{d}\theta\leq B.$

Then, for any $t\geq\log\bigl{(}(1+B)K\bigr{)}$,

(63) $\displaystyle{\rm KL}[\rho_{t}(\theta)\|\rho_{\rm
post}(\theta)]\leq(2+B+eB)Ke^{-t}.$

It is notable that the exponential convergence rate is independent of the
properties of the target distribution $\rho_{\rm post};$ this reflects
invariance of the flow under any diffeomorphism. The proof of this proposition
is in Section C.1. Similar propositions are in [93, Theorem 3.3] and [94,
Theorem 2.3]; our results relax the assumptions required on the initial
condition.

####  3.5.2 Wasserstein Gradient Flow

The convergence of the Wasserstein gradient flow Eq. 42 is widely studied
[128]. A variety of different conditions on $\rho_{\rm post}$ lead to the
exponential convergence of the Wasserstein gradient flow to $\rho_{\rm post}$
with convergence rate $e^{-2\alpha t}$ [8]. They include that $\rho_{\rm
post}$ is $\alpha$-strongly logconcave (Definition 3.13) [7] or that
$\rho_{\rm post}$ satisfies the log-Sobolev inequality [59] or Poincaré
inequality [108] with constant $1/\alpha$. We have the following proposition
concerning the convergence of the affine-invariant Wasserstein gradient flow
Eq. 47:

######  Definition 3.13.

The distribution $\rho_{\rm post}(\theta)$ is called $\alpha$-strongly
logconcave, if

(64) $\displaystyle-\nabla_{\theta}\nabla_{\theta}\log\rho_{\rm
post}(\theta)\succeq\alpha I.$

######  Proposition 3.14.

Assume $\rho_{\rm post}(\theta)$ is $\alpha$-strongly logconcave and there
exists $\lambda>0$ such that $P(\theta,\rho)\succeq\lambda I$. Then the
solution of the affine-invariant Wasserstein gradient flow Eq. 47 $\rho_{t}$
saisfies

$\frac{1}{2}\lVert\rho_{t}-\rho_{\rm
post}\rVert^{2}_{1}\leq\mathrm{KL}[\rho_{t}\|\rho_{\rm
post}]e^{-2\alpha\lambda t},$

where $\lVert\cdot\rVert_{1}$ denotes the $L_{1}$ norm.

The proof of the proposition is in Section C.2. It is a generalization of [54,
Proposition 3.1] which concerns the specific preconditioner
$P_{t}:=P(\theta_{t},\rho_{t})$ chosen to equal $C_{t}$, the covariance at
time $t.$ A key point to appreciate is that, in contrast to the exponential
rates reported for Fisher-Rao gradient descent, the exponential rates reported
here depend on the problem. When $\rho_{\rm post}$ is Gaussian, the affine
invariant Wasserstein gradient flows, however, provably achieves convergence
rate $\mathcal{O}(\exp(-t))$ [54, 55]; it would be of interest to identify
classes of non-Gaussian problems where this rate is also achievable for the
affine invariant Wasserstein gradient flow.

The result shows that a key determining factor in the exponential rates is
lower-bounding the preconditioner. Regarding the choice of the covariance as
preconditioner we have the following lower bounds:

######  Lemma 3.15.

Let $\rho_{t}$ solve the Wasserstein gradient flow Eq. 47 with preconditioner
$P_{t}=C_{t}$. Assume that there exist $\lambda_{0,\min},\beta>0$ such that
$C_{0}\succeq\lambda_{0,\min}$ and
$-\nabla_{\theta}\nabla_{\theta}\log\rho_{\rm post}\preceq\beta$. Then we
have:

  * •

if $N_{\theta}=1$, then $C_{t}\geq\min\\{\frac{1}{\beta},\lambda_{0,\min}\\};$

  * •

if $N_{\theta}>1$ and ${\rm tr}\,C_{t}\leq K$, then
$C_{t}\succeq\min\\{\frac{1}{\beta^{2}K},\lambda_{0,\min}\\}$

The proof of this lemma is in Section C.3. A similar result, providing a non-
degeneracy guarantee for the empirical covariance in particle approximations,
can be found in [55, Proposition 4.4].

####  3.5.3 Stein Gradient Flow

For the Stein gradient flow Eq. 56 the solution $\rho_{t}$ converges weakly to
$\rho_{\rm post}$ as $t\rightarrow\infty$, under certain assumptions [92,
Theorem 2.8][77, Proposition 2]; the exponential rates are problem-dependent,
similar to those for Wasserstein gradient flows in the preceding subsection,
and in contrast to those for the Fisher-Rao gradient flow which give a
universal rate across wide problem classes. Quantitative rates and necessary
functional inequalities for the exponential convergence near the equilibrium
in terms of the decay of the KL divergence are discussed in [45]. However, the
speed of convergence for initial distributions far from equilibrium remains an
open and challenging problem.

##  4 Gaussian Approximate Gradient Flow

In this section, we revisit the gradient flows of the energy Eq. 4 under the
Fisher-Rao, Wasserstein and Stein metrics. We confine variations to the
manifold of Gaussian densities $\mathcal{P}^{G}$ defined in Eq. 8, in contrast
to the previous Section 3, in which we consider variations in the whole of
$\mathcal{P}$ defined in Eq. 3. The corresponding Gaussian approximate
gradient flows underpin Gaussian variational inference, which aims to identify
the minimizers of Eq. 2. We first introduce the basics of metrics and gradient
flow in the Gaussian density space, identify the ways that Gaussian
approximations can be made and develop the concept of affine invariance for
them in Section 4.1. Then we introduce the Gaussian approximate Fisher-Rao
gradient flow in Section 4.2, Gaussian approximate Wasserstein graident flow
in Section 4.3 and Gaussian approximate Stein gradient flow in Section 4.4; in
all cases we also discuss affine invariance and introduce affine invariant
modifications where appropriate. We find that different affine invariant
metrics lead to very similar gradient flows; and in particular to flows with
very similar large time behaviour. We discuss the large time convergence
properties of these Gaussian approximate gradient flows in Section 4.5.

###  4.1 Basics of Gaussian Approximate Gradient Flows

In this subsection, we introduce gradient flows in the Gaussian density space;
we follow the structure of Section 3.1. We study the problem from perspective
of the metric in Section 4.1.1, the perspective of the flow equations in
Section 4.1.2, the perspective of affine invariance in Section 4.1.3, and the
perspective of mean-field equations in Section 4.1.4. For Gaussian evolutions
the mean-field models are evolution equations for the state defined by affine
(in the state) tangent vector field; the affine map is defined by mean-field
expectations with respect to the Gaussian with mean and covariance of the
state.

####  4.1.1 Metric

Recall the manifold of Gaussian densities in Eq. 8, which has dimension
$N_{a}$. We assume we are given a metric $g_{\rho}$ and metric tensor
$M(\rho)$, depending on $\rho\in\mathcal{P}$, and we now wish to find
corresponding objects defined for parametric variations within the family of
Gaussian densities $\mathcal{P}^{G}$.444In fact our development is readily
generalized to the determination of the corresponding objects for any
parametrically dependent manifold of densities, not just Gaussians. To this
end we introduce $\rho_{a}$, with $a\in\mathbb{R}^{N_{a}}$, denoting the
parametric family. We aim to find reduced metric $\mathfrak{g}_{a}$ and metric
tensor $\mathfrak{M}(a)$ in the parameter space $\mathbb{R}^{N_{a}}$ rather
than in $\mathcal{P}.$

Noting that

(65) $\displaystyle\lim_{\epsilon\rightarrow
0}\frac{\rho_{a+\epsilon\sigma}-\rho_{a}}{\epsilon}=\nabla_{a}\rho_{a}\cdot\sigma,$

we see that any element in the tangent space $T_{\rho_{a}}\mathcal{P}^{G}$ can
be identified with a vector $\sigma\in\mathbb{R}^{N_{a}}$. We denote the
Riemannian metric tensor restricted to $\mathcal{P}^{G}$ at $\rho_{a}$ as
$\mathfrak{g}_{a}$. Then

(66)
$\displaystyle\mathfrak{g}_{a}(\sigma_{1},\sigma_{2}):=g_{\rho_{a}}(\nabla_{a}\rho_{a}\cdot\sigma_{1},\nabla_{a}\rho_{a}\cdot\sigma_{2})=\langle\mathfrak{M}(a)\sigma_{1},\sigma_{2}\rangle_{\mathbb{R}^{N_{a}}},$

where $\sigma_{1},\sigma_{2}\in T_{\rho_{a}}(\mathcal{P}^{G})$, and the
induced metric tensor is given by

(67)
$\mathfrak{M}(a):=\int\nabla_{a}\rho_{a}(\theta)\bigl{(}M(\rho_{a})\nabla_{a}\rho_{a}^{T}\bigr{)}(\theta)\mathrm{d}\theta.$

####  4.1.2 Flow Equation

Given Eq. 67 it is intuitive that the gradient flow in the parameter space
implied by the gradient flow in the manifold of Gaussians is given by

(68) $\displaystyle\frac{\partial a_{t}}{\partial
t}=-\mathfrak{M}(a_{t})^{-1}\left.\frac{\partial\mathcal{E}(\rho_{a})}{\partial
a}\right|_{a=a_{t}}.$

We refer to Eq. 68 as the Gaussian approximate gradient flow; it is formulated
as an evolution equation in the parameter space. It is also possible to write
an evolution equation for $\rho_{a_{t}}$ in the space of Gaussian probability
densities $\mathcal{P}^{G}$.

Our goal now is to show that Eq. 68 may be derived by using any one of the
following proximal, Riemannian, and moment closure perspectives. In
particular, these perspectives justify that the gradient flow in the space of
Gaussian densities $\mathcal{P}^{G}$ is a Gaussian approximation of the
gradient flow on the whole probability space. Such approximation can be
interpreted either by constraining the minimization underlying the proximal
perspective, by the projection of the flow field based on the Riemannian
metric, or through a moment closure reduction of probability densities. The
latter moment closure approach is particularly expedient for determination of
the form of equation Eq. 68.

##### Proximal Perspective

Given the metric $g$ and the corresponding distance function $\mathcal{D}$,
the proximal point method Eq. 15 can be restricted to the space of Gaussian
desnities, leading to the iteration

(69)
$\displaystyle\mathfrak{r}_{{n+1}}=\operatorname*{arg\,min}_{\rho\in\mathcal{P}^{G}}\Bigl{(}\mathcal{E}(\rho)+\frac{1}{2\Delta
t}\mathcal{D}(\rho,\mathfrak{r}_{{n}})^{2}\Bigr{)}$

to minimize the energy functional $\mathcal{E}$ in Gaussian density function
space $\mathcal{P}^{G}$. Since elements in $\mathcal{P}^{G}$ are uniquely
defined via a point $a\in\mathbb{R}^{N_{a}}$, the map
$\mathfrak{r}_{{n}}\mapsto\mathfrak{r}_{{n+1}}$ implicitly defines a map
$a_{n}\mapsto a_{n+1}.$ Thus we write $\mathfrak{r}_{{n}}=\rho_{a_{n}}$ and
determine the update equation for $a_{n}$. When $\Delta t$ is small it is
natural to seek $a_{n+1}=a_{n}+\Delta t\sigma_{n}$ and note that, invoking the
approximations implied by Eq. 65 and Eq. 13,

$\displaystyle\sigma_{{n}}$
$\displaystyle\approx\operatorname*{arg\,min}_{\sigma\in\mathbb{R}^{N_{a}}}\Bigl{(}\mathcal{E}(\rho_{a_{n}}+\Delta
t\nabla_{a}\rho_{a_{n}}\cdot\sigma)+\frac{1}{2}{\Delta
t}\langle\nabla_{a}\rho_{a_{n}}\cdot\sigma,M(\rho_{a_{n}})\nabla_{a}\rho_{a_{n}}\cdot\sigma\rangle\Bigr{)},$
$\displaystyle\approx\operatorname*{arg\,min}_{\sigma\in\mathbb{R}^{N_{a}}}\Bigl{(}\mathcal{E}(\rho_{a_{n}}+\Delta
t\nabla_{a}\rho_{a_{n}}\cdot\sigma)+\frac{1}{2}{\Delta
t}\langle\sigma,\mathfrak{M}(a_{n})\sigma\rangle_{\mathbb{R}^{N_{a}}}\Bigr{)}.$

To leading order in $\Delta t$, this expression is minimized by choosing

$\sigma_{n}=-\mathfrak{M}(a_{n})^{-1}\Bigl{\langle}\frac{\delta\mathcal{E}}{\delta\rho}\Bigr{|}_{\rho=\rho_{a_{n}}},\nabla_{a}\rho_{a_{n}}\Bigr{\rangle}=-\mathfrak{M}(a_{n})^{-1}\frac{\partial\mathcal{E}(\rho_{a})}{\partial
a}\Bigr{|}_{a=a_{n}}.$

Letting $a_{n}\approx a_{n\Delta t}$ shows that the formal continuous time
limit of the proximal algorithm leads to the corresponding gradient flow Eq.
68.

##### Riemannian Perspective

We start by defining the projection $P^{G}:T_{\rho}\mathcal{P}\to
T_{\rho}\mathcal{P}^{G}$ as follows: for any $\psi\in T_{\rho}\mathcal{P}$ we
define $P^{G}$ by requiring that

(70) $g_{\rho}(\psi,\sigma)=g_{\rho}(P^{G}\psi,\sigma),\quad\forall\sigma\in
T_{\rho}\mathcal{P}^{G}.$

Now consider the gradient flow

(71a) $\displaystyle\frac{\partial\rho_{t}}{\partial t}$
$\displaystyle=\sigma_{t}\in T_{\rho_{t}}\mathcal{P},$ (71b)
$\displaystyle\sigma_{t}$
$\displaystyle=-M(\rho_{t})^{-1}\frac{\delta\mathcal{E}}{\delta\rho}\Bigr{|}_{\rho=\rho_{t}}$

designed to decrease the functional $\mathcal{E}$ under the metric $g$. We
note that, by virtue of Eq. 6, $\sigma_{t}=\sigma_{t}(\theta,\rho_{t}).$

We may now consider restriction of the gradient flow to variations in the
manifold of Gaussian densities, leading to equation for
$\rho_{a_{t}}\in\mathcal{P}^{G}\subseteq\mathcal{P}$, defined through the
corresponding gradient flow

(72) $\displaystyle\frac{\partial\rho_{a_{t}}}{\partial t}=P^{G}\sigma_{t}\in
T_{\rho_{a_{t}}}\mathcal{P}^{G}.$

The proof of the following proposition may be found in Section D.2.

######  Proposition 4.1.

The flow of the parameter $a_{t}$ implied by the evolution equation Eq. 72 for
the density $\rho_{a_{t}}$ is the Gaussian approximate gradient flow Eq. 68.

##### Moment Closure Perspective

For any gradient flow Eq. 71 designed to decrease the functional $\mathcal{E}$
under the metric $g$, we consider the following moment closure approach to
obtain a Gaussian approximation. First, we write evolution equations for the
mean and covariance under Eq. 71 noting that they satisfy the following
identities:

(73)
$\displaystyle\frac{\mathrm{d}m_{t}}{\mathrm{d}t}=\frac{\mathrm{d}}{\mathrm{d}t}\int\rho_{t}(\theta)\theta\mathrm{d}\theta=\int\sigma_{t}(\theta,\rho_{t})\theta\mathrm{d}\theta,$
$\displaystyle\frac{\mathrm{d}C_{t}}{\mathrm{d}t}=\frac{\mathrm{d}}{\mathrm{d}t}\int\rho_{t}(\theta)(\theta-
m_{t})(\theta-
m_{t})^{T}\mathrm{d}\theta=\int\sigma_{t}(\theta,\rho_{t})(\theta-
m_{t})(\theta-m_{t})^{T}\mathrm{d}\theta.$

This is not, in general, a closed system for the mean and covariance; this is
because $\rho_{t}$ is not, in general, determined by only first and second
moments. To close the system, we replace $\sigma_{t}(\theta,\rho_{t})$ by
$\sigma_{t}(\theta,\rho_{a_{t}})$, where
$\rho_{a_{t}}=\mathcal{N}(m_{t},C_{t})$. We obtain the following closed system
for the evolution of $(m_{t},C_{t})$:

(74)
$\displaystyle\frac{\mathrm{d}m_{t}}{\mathrm{d}t}=\int\sigma_{t}(\theta,\rho_{a_{t}})\theta\mathrm{d}\theta,$
$\displaystyle\frac{\mathrm{d}C_{t}}{\mathrm{d}t}=\int\sigma_{t}(\theta,\rho_{a_{t}})(\theta-
m_{t})(\theta-m_{t})^{T}\mathrm{d}\theta.$

The proof of the following proposition, which shows that this moment closure
approach delivers the mean and covariance evolution equation of the Gaussian
approximate gradient flow Eq. 68, may be found in Section D.2.

######  Proposition 4.2.

Suppose that the cotangent space in the Gaussian submanifold corresponding to
the metric $g$ is a collection of linear and quadratic functions of $\theta$:

(75)

Here is the metric tensor and $T_{\rho_{a}}\mathcal{P}^{G}$ is the tangent
space. Then the mean and covariance evolution equations LABEL:eq:mC-Momentum2
are equivalent to the Gaussian approximate gradient flow Eq. 68.

Furthermore, Section D.2 also contains proof of the following Lemma 4.3
indicating that several of the metrics considered later in this Section 4 do
indeed satisfy the assumption Eq. 75 sufficient for Proposition 4.2 to hold.

######  Lemma 4.3.

Assumption Eq. 75 holds for the Fisher-Rao metric, the affine invariant
Wasserstein metric with preconditioner $P$ independent of $\theta$, and affine
invariant Stein metric with preconditioner $P$ independent of $\theta$ and
with a bilinear kernel
$\kappa(\theta,\theta^{\prime},\rho)=(\theta-m)^{T}A(\rho)(\theta^{\prime}-m)+b(\rho)$
($b\neq 0$, and $A$ nonsingular).

e moment closure perspective was used in [117] as a heuristic approach to
state estimation in the context of the unscented Kalman filter. A connection
between the heuristics and gradient flow on the Bures–Wasserstein space of
Gaussian distributions was established in [80]. The latter is equivalent to
the Gaussian approximate gradient flow under the Wasserstein metric, also
called the Gaussian approximate Wasserstein gradient flow in this paper; see
also the discussion in Section 4.3.1.

####  4.1.3 Affine Invariance

We now study the affine invariance concept in the setting of Gaussian
approximate gradient flows. Let $\varphi:\theta\rightarrow\tilde{\theta}$
denote an invertible affine transformation in $\mathbb{R}^{N_{\theta}}$, where
$\varphi(\theta)=A\theta+b$ with
$A\in\mathbb{R}^{N_{\theta}}\times\mathbb{R}^{N_{\theta}}$,
$b\in\mathbb{R}^{N_{\theta}}$, and $A$ invertible.

We define the push forward operator for various objects.

  * •

For a parametrically-defined density $\rho_{a}$, we write
$\rho_{\tilde{a}}=\varphi\\#\rho_{a}$, so that
$\rho_{\tilde{a}}(\tilde{\theta})=\rho_{a}(\varphi^{-1}(\tilde{\theta}))|A^{-1}|$.
Specifically, for Gaussian density space where $a=(m,C)$, we have an
invertible affine transformation in $\mathbb{R}^{N_{a}}$, such that
$\tilde{a}=A^{G}a+b^{G}$, where $A^{G}$ and $b^{G}$ depend only on $A$ and $b$
and are defined by the identities $\tilde{m}=Am+b$ and $\tilde{C}=ACA^{T}.$

  * •

For a tangent vector $\sigma\in\mathbb{R}^{N_{a}}$ corresponding to
$\nabla_{a}\rho_{a}(\theta)\cdot\sigma$ in $T_{\rho_{a}}\mathcal{P}^{G}$, we
have $\tilde{\sigma}=A^{G}\sigma\in\mathbb{R}^{N_{a}}$ corresponding to
$\nabla_{\tilde{a}}\rho_{\tilde{a}}(\tilde{\theta})\cdot\tilde{\sigma}$ in
$T_{\rho_{\tilde{a}}}\mathcal{P}^{G}$, and note that this satisfies
$\nabla_{\tilde{a}}\rho_{\tilde{a}}(\tilde{\theta})\cdot\tilde{\sigma}=\varphi\\#(\nabla_{a}\rho_{a}(\theta)\cdot\sigma)=\nabla_{a}\rho_{a}(\varphi^{-1}(\tilde{\theta}))\cdot\sigma|A^{-1}|$.

  * •

For a functional $\mathcal{E}$ on $\mathcal{P}^{G}$, we define
$\tilde{\mathcal{E}}=\varphi\\#\mathcal{E}$ via
$\tilde{\mathcal{E}}(\rho_{\tilde{a}})=\mathcal{E}(\varphi^{-1}\\#\rho_{\tilde{a}})$.

With the above, we can make precise the definition of affine invariance for
the Gaussian approximate gradient flows. The definition is similar to
Definition 3.2.

######  Definition 4.4 (Affine Invariant Gaussian Approximate Gradient Flow).

The Gaussian approximate gradient flow Eq. 68 is called _affine invariant_ if,
under any invertible affine transformation $\tilde{a}_{t}=\varphi(a_{t})$, the
dynamics of $\tilde{a}_{t}$ is itself a gradient flow of
$\tilde{\mathcal{E}}$, in the sense that

(76) $\displaystyle\frac{\partial\tilde{a}_{t}}{\partial
t}=-\mathfrak{M}(\tilde{a}_{t})^{-1}\frac{\partial\tilde{\mathcal{E}}(\rho_{\tilde{a}})}{\partial\tilde{a}}\Big{|}_{\tilde{a}=\tilde{a}_{t}}.$

Naturally, if the gradient flow in probability space is affine invariant, then
the Gaussian approximate flow has the same property; see the following
proposition.

######  Proposition 4.5.

For any affine invariant metric $g$ defined via Definition 3.3, the Gaussian
approximate gradient flow under the corresponding metric $\mathfrak{g}_{a}$ is
affine invariant for any $\mathcal{E}$.

We provide a proof for this proposition in Section D.3.

####  4.1.4 Mean-Field Dynamics

The Gaussian approximate gradient flow can also be realized as a mean-field
ordinary differential equation. Recall that for Gaussian evolutions the mean-
field models are evolution equations for mean and covariance defined via mean-
field expectations with respect to the Gaussian with this mean and covariance.
We have the following lemma.

######  Lemma 4.6.

Consider the mean-field equation

(77) $\frac{\mathrm{d}\theta_{t}}{\mathrm{d}t}=\mathsf{A}(\rho_{t},\rho_{\rm
post})(\theta_{t}-m_{t})+\mathsf{b}(\rho_{t},\rho_{\rm post}),$

where
$\mathsf{A}:\mathcal{P}\times\mathcal{P}\rightarrow\mathbb{R}^{N_{\theta}\times
N_{\theta}}$ and
$\mathsf{b}:\mathcal{P}\times\mathcal{P}\rightarrow\mathbb{R}$, $\rho_{t}$ is
the law of $\theta_{t}$ and $m_{t}$ is the mean under $\rho_{t}$. If the law
of $\theta_{0}$ is a Gaussian distribution, then $\theta_{t}$ solving (77) is
also Gaussian distributed for any $t>0$; thus we can write
$\rho_{t}=\rho_{a_{t}}$ where $a_{t}=(m_{t},C_{t})$ denotes the mean and
covariance of the distribution. The evolution of $m_{t}$ and $C_{t}$ is given
by

(78)
$\begin{split}\frac{\mathrm{d}m_{t}}{\mathrm{d}t}&=\mathsf{b}(\rho_{a_{t}},\rho_{\rm
post}),\\\
\frac{\mathrm{d}C_{t}}{\mathrm{d}t}&=\mathsf{A}(\rho_{a_{t}},\rho_{\rm
post})C_{t}+C_{t}\mathsf{A}(\rho_{a_{t}},\rho_{\rm post})^{T}.\end{split}$

We provide a proof of this lemma in Section D.4. The lemma allows us to
identify the corresponding mean-field dynamics Eq. 77 of the Gaussian
approximate gradient flow Eq. 76. Furthermore the evolution equation (78) for
the mean and covariance is defined via vector field for the evolution defined
by expectations under the Gaussian with this mean and covariance. We will
elaborate on this identification in detail for specific metric tensors in
later subsections. Regarding the affine invariance property of the mean-field
equation, we have the following proposition:

######  Proposition 4.7.

Suppose the mean and covariance evolution equation Eq. 78 is affine invariant,
in the sense that under the invertible affine transformation
$\tilde{\theta}=\varphi(\theta)=A\theta+b$ and correspondingly
$\tilde{m}_{t}=Am_{t},\tilde{C}_{t}=AC_{t}A^{T},\tilde{\rho}_{a_{t}}=\varphi\\#\rho_{a_{t}},\tilde{\rho}_{\rm
post}=\varphi\\#\rho_{\rm post}$, it holds that

(79)
$\begin{split}\frac{\mathrm{d}\tilde{m}_{t}}{\mathrm{d}t}&=\mathsf{b}(\tilde{\rho}_{a_{t}},\tilde{\rho}_{\rm
post}),\\\
\frac{\mathrm{d}\tilde{C}_{t}}{\mathrm{d}t}&=\mathsf{A}(\tilde{\rho}_{a_{t}},\tilde{\rho}_{\rm
post})\tilde{C}_{t}+\tilde{C}_{t}\mathsf{A}(\tilde{\rho}_{a_{t}},\tilde{\rho}_{\rm
post})^{T}.\end{split}$

Then, the corresponding mean-field equation Eq. 77 is also affine invariant.

We provide a proof of this proposition in Section D.5.

###  4.2 Gaussian Approximate Fisher-Rao Gradient Flow

####  4.2.1 Metric

In the Gaussian density space, where $\rho$ is parameterized by
$a=[m,C]\in\mathbb{R}^{N_{a}}$, the induced Fisher-Rao metric tensor
$\mathfrak{M}(a)\in\mathbb{R}^{N_{a}\times N_{a}}$ has entries

(80) $\displaystyle\mathfrak{M}(a)_{jk}$
$\displaystyle=\int\frac{\partial\log\rho_{a}(\theta)}{\partial
a_{j}}\frac{\partial\log\rho_{a}(\theta)}{\partial
a_{k}}\rho_{a}(\theta)\mathrm{d}\theta.$

This is also the Fisher information matrix, which has explicit formula in the
Gaussian space:

(81) $\displaystyle\mathfrak{M}(a)_{jk}$ $\displaystyle=\frac{\partial
m}{\partial a_{j}}^{T}C^{-1}\frac{\partial m}{\partial a_{k}}+\frac{1}{2}{\rm
tr}\Bigl{(}C^{-1}\frac{\partial C}{\partial a_{j}}C^{-1}\frac{\partial
C}{\partial a_{k}}\Bigr{)}.$

####  4.2.2 Flow Equation

The moment closure approach in Section 4.1.2 delivers the evolution of the
mean and covariance by virtue of Lemma 4.3. Using this leads to the following
equations:

$\displaystyle\frac{\mathrm{d}m_{t}}{dt}$
$\displaystyle=\mathrm{Cov}_{\rho_{a_{t}}}[\theta,\,\log\rho_{\rm
post}-\log\rho_{a_{t}}]$
$\displaystyle=C_{t}\mathbb{E}_{\rho_{a_{t}}}[\nabla_{\theta}(\log\rho_{\rm
post}-\log\rho_{a_{t}})],$ $\displaystyle\frac{\mathrm{d}C_{t}}{dt}$
$\displaystyle=\mathbb{E}_{\rho_{a_{t}}}[(\theta-m_{t})(\theta-
m_{t})^{T}(\log\rho_{\rm
post}-\log\rho_{a_{t}}-\mathbb{E}_{\rho_{a_{t}}}[\log\rho_{\rm
post}-\log\rho_{a_{t}}])]$
$\displaystyle=C_{t}\mathbb{E}_{\rho_{a_{t}}}[\nabla_{\theta}\nabla_{\theta}(\log\rho_{\rm
post}-\log\rho_{a_{t}})]C_{t},$

where $\rho_{a_{t}}\sim\mathcal{N}(m_{t},C_{t})$. Thus we obtain

(82) $\displaystyle\frac{\mathrm{d}m_{t}}{dt}$
$\displaystyle=C_{t}\mathbb{E}_{\rho_{a_{t}}}[\nabla_{\theta}\log\rho_{\rm
post}],$ $\displaystyle\frac{\mathrm{d}C_{t}}{dt}$
$\displaystyle=C_{t}+C_{t}\mathbb{E}_{\rho_{a_{t}}}[\nabla_{\theta}\nabla_{\theta}\log\rho_{\rm
post}]C_{t}.$

turning to LABEL:rem:homotopy and assuming a quadratic negative log likelihood

(83) $\displaystyle\Phi(\theta)=\frac{1}{2}(H\theta-y)^{\rm
T}R^{-1}(H\theta-y)$

and a Gaussian prior distribution, the functional (27) leads to the following
Fisher-Rao gradient flow equations

(84) $\displaystyle\frac{\mathrm{d}m_{t}}{dt}$ $\displaystyle=-C_{t}H^{\rm
T}R^{-1}(Hm_{t}-y),$ $\displaystyle\frac{\mathrm{d}C_{t}}{dt}$
$\displaystyle=-C_{t}H^{\rm T}R^{-1}HC_{t},$

for the mean $m_{t}$ and covariance matrix $C_{t}$. These define the well-
known Kalman-Bucy filter [72] from linear state estimation; see the text [67]
for further details. Their extension to general log-likelihood functions
$\Phi$ under Gaussian approximation is discussed in [107].

Equation 82 also corresponds to the gradient flow under the finite dimensional
Fisher-Rao metric in the parameter space [104, 75]; in this context it goes by
the nomenclature _natural gradient flow_ [1, 97, 138]. The connection between
Fisher-Rao natural gradient methods and Kalman filters has been studied in
[102, 103].

####  4.2.3 Affine Invariance

Equation 82 is affine invariant.

####  4.2.4 Mean-Field Dynamics

Using Lemma 4.6 we can read off a choice of the pair $\mathsf{A},\mathsf{b}$
defining the mean-field equation for the Gaussian approximate Fisher-Rao
gradient flow (82); we obtain

(85) $\frac{\mathrm{d}\theta_{t}}{\mathrm{d}t}=\frac{1}{2}\Bigl{[}$

Following Proposition 4.7, the mean-field equation Eq. 85 is also affine
invariant.

###  4.3 Gaussian Approximate Wasserstein Gradient Flow

####  4.3.1 Metric

Recall the preconditioner $P:$ where the output space is the cone of positive-
definite symmetric matrices. In the Gaussian density space, where $\rho$ is
parameterized by $a\in$, the preconditioned Wasserstein metric tensor
$\mathfrak{M}(a)\in$ has entries

(86)
$\displaystyle\mathfrak{M}(a)_{jk}=\int\psi_{j}(\theta)\frac{\partial\rho_{a}(\theta)}{\partial
a_{k}}\mathrm{d}\theta\quad\textrm{where}\quad-\nabla_{\theta}\cdot\Bigl{(}\rho_{a}(\theta)P(\theta,\rho_{a})\nabla_{\theta}\psi_{j}(\theta)\Bigr{)}=\frac{\partial\rho_{a}}{\partial
a_{j}}(\theta).$

When $P(\theta,\rho_{a})$ is the identity operator, the metric tensor
$\mathfrak{M}(a)$ has an explicit formula [30, 125, 95, 14, 84], and the
corresponding Gaussian density space is called the Bures–Wasserstein space
[128].

####  4.3.2 Flow Equation

Again using the moment closure approach from Section 4.1.2, which is shown to
apply here in Lemma 4.3, the mean and covariance evolution equations for the
Gaussian approximate Wasserstein gradient flow with $\theta-$independent $P$
are

(87)
$\begin{split}\frac{\mathrm{d}m_{t}}{\mathrm{d}t}&=\int\Bigl{[}\rho_{a_{t}}P(\rho_{a_{t}})\nabla_{\theta}(\log\rho_{\rm
post}-\log\rho_{a_{t}})\Bigr{]}\mathrm{d}\theta=P(\rho_{a_{t}})\mathbb{E}_{\rho_{a_{t}}}[\nabla_{\theta}\log\rho_{\rm
post}],\\\
\frac{\mathrm{d}C_{t}}{\mathrm{d}t}&=\int\nabla_{\theta}\cdot\Bigl{[}\rho_{a_{t}}P(\rho_{a_{t}})\nabla_{\theta}(\log\rho_{a_{t}}-\log\rho_{\rm
post})\Bigr{]}(\theta-m_{t})(\theta-m_{t})^{T}\mathrm{d}\theta\\\
&=2P(\rho_{a_{t}})+P(\rho_{a_{t}})\mathbb{E}_{\rho_{a_{t}}}[\nabla_{\theta}\nabla_{\theta}\log\rho_{\rm
post}]C_{t}+C_{t}\mathbb{E}_{\rho_{a_{t}}}[\nabla_{\theta}\nabla_{\theta}\log\rho_{\rm
post}]P(\rho_{a_{t}}),\end{split}$

where $\rho_{a_{t}}\sim\mathcal{N}(m_{t},C_{t})$.

When we set $P(\rho)\equiv$, the evolution equation becomes

(88) $\displaystyle\frac{\mathrm{d}m_{t}}{\mathrm{d}t}$
$\displaystyle=\mathbb{E}_{\rho_{a_{t}}}\bigl{[}\nabla_{\theta}\log\rho_{\rm
post}(\theta_{t})\bigr{]},$ $\displaystyle\frac{\mathrm{d}C_{t}}{\mathrm{d}t}$
$\displaystyle=2$

This corresponds to the gradient flow under the constrained Wasserstein metric
in Gaussian density space [117, 70, 80].

####  4.3.3 Affine Invariance

We now allow the preconditioner to depend on $\rho_{a_{t}}$ and set
$P(\rho_{a_{t}})=C_{t}$. This choice satisfies the affine invariant condition
Proposition 3.8. The resulting evolution equations for the corresponding mean
and covariance are

(89) $\displaystyle\frac{\mathrm{d}m_{t}}{\mathrm{d}t}$
$\displaystyle=C_{t}\mathbb{E}_{\rho_{a_{t}}}\bigl{[}\nabla_{\theta}\log\rho_{\rm
post}(\theta_{t})\bigr{]},$ $\displaystyle\frac{\mathrm{d}C_{t}}{\mathrm{d}t}$
$\displaystyle=2C_{t}+2C_{t}\mathbb{E}_{\rho_{a_{t}}}\bigl{[}\nabla_{\theta}\nabla_{\theta}\log\rho_{\rm
post}(\theta_{t})\bigr{]}C_{t}.$

Equation Eq. 89 is similar to the Gaussian approximate Fisher-Rao gradient
flow Eq. 82, but with scaling factor $2$ in the covariance evolution.

####  4.3.4 Mean-Field Dynamics

Employing Lemma 4.6 we can again identify a pair $\mathsf{A},\mathsf{b}$
leading to a mean-field equation for the Gaussian approximate Wasserstein
gradient flow (87) with $\theta$-independent $P$:

(90)
$\frac{\mathrm{d}\theta_{t}}{\mathrm{d}t}=\Bigl{[}C_{t}^{-1}+P(\rho_{a_{t}})\mathbb{E}_{\rho_{a_{t}}}[\nabla_{\theta}\nabla_{\theta}\log\rho_{\rm
post}]\Bigr{]}(\theta_{t}-m_{t})+P(\rho_{a_{t}})\mathbb{E}_{\rho_{a_{t}}}[\nabla_{\theta}\log\rho_{\rm
post}].$

From Proposition 4.7, its corresponding mean-field equation Eq. 90 with
$P(\rho_{a_{t}})=C_{t}$ is also affine invariant.

###  4.4 Gaussian Approximate Stein Gradient Flow

####  4.4.1 Metric

We work in the general preconditioned setting, as in the Wasserstein case. In
the Gaussian density space, where $\rho$ is parameterized by $a\in$, the Stein
metric tensor $\mathfrak{M}(a)\in$ is

(91)
$\displaystyle\mathfrak{M}(a)_{jk}=\int\psi_{j}(\theta)\frac{\partial\rho_{a}(\theta)}{\partial
a_{k}}\mathrm{d}\theta,\quad\textrm{where}$
$\displaystyle-\nabla_{\theta}\cdot\Bigl{(}\rho_{a}(\theta)\int\kappa(\theta,\theta^{\prime},\rho_{a})\rho_{a}(\theta^{\prime})P(\theta,\theta^{\prime},\rho_{a}(\theta),\rho_{a}(\theta^{\prime}))\nabla_{\theta}\psi_{j}(\theta^{\prime})\mathrm{d}\theta^{\prime}\Bigr{)}=\frac{\partial\rho_{a}}{\partial
a_{j}}(\theta).$

####  4.4.2 Flow Equation

We consider the setting in which $P=P(\rho_{a})$ is independent of $\theta$,
and we choose bilinear kernel

(92)
$\kappa(\theta,\theta^{\prime},\rho)=(\theta-m)^{T}A(\rho)(\theta^{\prime}-m)+b(\rho).$

where $A:\mathcal{P}\rightarrow$, $b:\mathcal{P}\rightarrow$ and $m$ is the
mean under $\rho.$

In the following let $P_{t}:=P(\rho_{a_{t}})$, evaluate $A$ and $b$ at
$\rho=\rho_{a_{t}}$, writing the resulting time-dependent matrix- and vector-
valued functions as $A_{t}$ and $b_{t}$ where $A_{\cdot}:$ and $b_{\cdot}:$,
and let $m_{t}$ denote mean under $\rho_{a_{t}}$ so that $m_{\cdot}:$ Using
the moment closure approach, from Section 4.1.2, which is shown to apply here
in Lemma 4.3, the mean and covariance evolution equations of the
preconditioned Stein gradient flow with bilinear kernel Eq. 92 can be shown to
be

(93)
$\displaystyle\frac{\mathrm{d}m_{t}}{\mathrm{d}t}=-\int\Bigl{(}\rho_{a_{t}}(\theta)\int\kappa(\theta,\theta^{\prime},\rho_{a_{t}})\rho_{a_{t}}(\theta^{\prime})P_{t}\nabla_{\theta^{\prime}}\bigl{(}\log\rho_{a_{t}}(\theta^{\prime})-\log\rho_{\rm
post}(\theta^{\prime})\bigr{)}\mathrm{d}\theta^{\prime}\Bigr{)}\mathrm{d}\theta,$
$\displaystyle\frac{\mathrm{d}C_{t}}{\mathrm{d}t}=$
$\displaystyle-\int\Bigl{(}\rho_{a_{t}}(\theta)\int\kappa(\theta,\theta^{\prime},\rho_{a_{t}})\rho_{a_{t}}(\theta^{\prime})P_{t}\nabla_{\theta^{\prime}}\bigl{(}\log\rho_{a_{t}}(\theta^{\prime})-\log\rho_{\rm
post}(\theta^{\prime})\bigr{)}\mathrm{d}\theta^{\prime}\Bigr{)}(\theta-
m_{t})^{T}$ $\displaystyle+(\theta-
m_{t})\Bigl{(}\rho_{a_{t}}(\theta)\int\kappa(\theta,\theta^{\prime},\rho_{a_{t}})\rho_{a_{t}}(\theta^{\prime})P_{t}\nabla_{\theta^{\prime}}\bigl{(}\log\rho_{a_{t}}(\theta^{\prime})-\log\rho_{\rm
post}(\theta^{\prime})\bigr{)}\mathrm{d}\theta^{\prime}\Bigr{)}^{T}\mathrm{d}\theta.$

where $\rho_{a_{t}}\sim\mathcal{N}(m_{t},C_{t})$. Imposing the form of the
bilinear kernel Eq. 92 we obtain

(94) $\displaystyle\frac{\mathrm{d}m_{t}}{\mathrm{d}t}=$ $\displaystyle
b_{t}P_{t}\mathbb{E}_{\rho_{a_{t}}}[\nabla_{\theta}\log\rho_{\rm post}],$
$\displaystyle\frac{\mathrm{d}C_{t}}{\mathrm{d}t}=$ $\displaystyle
P_{t}A_{t}C_{t}+C_{t}A_{t}P_{t}$
$\displaystyle\quad\quad\quad+P_{t}\mathbb{E}_{\rho_{a_{t}}}[\nabla_{\theta}\nabla_{\theta}\log\rho_{\rm
post}]C_{t}A_{t}C_{t}+C_{t}A_{t}C_{t}\mathbb{E}_{\rho_{a_{t}}}[\nabla_{\theta}\nabla_{\theta}\log\rho_{\rm
post}]P_{t}.$

fferent choices of the preconditioner $P$ and the bilinear kernel $\kappa$
allows us to recover different Gaussian variational inference methods
appearing in the literature. Choosing the preconditioner $P_{t}=$ and bilinear
kernel Eq. 92 with $A_{t}=C_{t}^{-1}$ and $b_{t}=1$ recovers the Gaussian
approximate Wasserstein gradient flow LABEL:eq:Gaussian-Wasserstein. Setting
the preconditioner $P_{t}=$ and bilinear kernel Eq. 92 with $A_{t}=I$ and
$b_{t}=1$ recovers the Gaussian sampling approach introduced in [53]:

(95) $\displaystyle\frac{\mathrm{d}m_{t}}{\mathrm{d}t}=$
$\displaystyle\mathbb{E}_{\rho_{a_{t}}}[\nabla_{\theta}\log\rho_{\rm post}],$
$\displaystyle\frac{\mathrm{d}C_{t}}{\mathrm{d}t}=$ $\displaystyle
2C_{t}+\mathbb{E}_{\rho_{a_{t}}}[\nabla_{\theta}\nabla_{\theta}\log\rho_{\rm
post}]C_{t}^{2}+C_{t}^{2}\mathbb{E}_{\rho_{a_{t}}}[\nabla_{\theta}\nabla_{\theta}\log\rho_{\rm
post}].$

####  4.4.3 Affine Invariance

Recalling that $\rho_{a_{t}}\sim\mathcal{N}(m_{t},C_{t})$ and setting
$P_{t}=P(\rho_{a_{t}})=C_{t}$ and choosing bilinear kernel Eq. 92 with
$A_{t}=\frac{1}{2}C_{t}^{-1}$ and $b_{t}=1$, which satisfies the affine
invariant condition Proposition 3.10, leads to the Gaussian approximate
Fisher-Rao gradient flow Eq. 82.

####  4.4.4 Mean-Field Dynamics

Using Lemma 4.6 we can deduce that the Gaussian approximate Stein gradient
flow (LABEL:eq:Gaussian-Stein-CTL0) with $\theta$-independent $P$ has the
following mean-field equation:

$\frac{\mathrm{d}\theta_{t}}{\mathrm{d}t}=\Bigl{[}P_{t}A_{t}+P_{t}\mathbb{E}_{\rho_{a_{t}}}[\nabla_{\theta}\nabla_{\theta}\log\rho_{\rm
post}]C_{t}A_{t}\Bigr{]}(\theta_{t}-m_{t})+b_{t}P_{t}\mathbb{E}_{\rho_{a_{t}}}[\nabla_{\theta}\log\rho_{\rm
post}].$

Here $P_{t}=P(\rho_{a_{t}})$. From Proposition 4.7, we know that this
corresponding mean-field equation is also affine invariant.

###  4.5 Convergence to Steady State

Recall that the objective of Gaussian variational inference is to solve the
minimization problem Eq. 9. Furthermore all critical points satisfy Eq. 10.
Regular gradient descent, in metric defined via the Euclidean inner-product in
gives rise to the dynamical system

(96) $\displaystyle\frac{\mathrm{d}m_{t}}{\mathrm{d}t}$
$\displaystyle=\mathbb{E}_{\rho_{a_{t}}}\bigl{[}\nabla_{\theta}\log\rho_{\rm
post}(\theta_{t})\bigr{]},$ $\displaystyle\frac{\mathrm{d}C_{t}}{\mathrm{d}t}$
$\displaystyle=\frac{1}{2}C_{t}^{-1}+\frac{1}{2}\mathbb{E}_{\rho_{a_{t}}}\bigl{[}\nabla_{\theta}\nabla_{\theta}\log\rho_{\rm
post}(\theta_{t})\bigr{]}.$

Note that steady states of this dynamical system necessarily satisfy Eq. 10.
In the preceding subsections we have derived a number of different gradient
flows in the manifold of Gaussian densities, including the Gaussian
approximate Fisher-Rao gradient flow Eq. 82 and the Gaussian approximate
Wasserstein gradient flow LABEL:eq:Gaussian-Wasserstein; note that both of
these dynamical systems also necessarily satisfy Eq. 10 in steady state. The
convergence properties of Eq. 89 obtained by the affine-invariant Wasserstein
gradient flow are similar to those of the Gaussian approximate Fisher-Rao
gradient flow Eq. 82. We omit detailed discussion from the paper to avoid
redundant discussions. In this subsection, we survey and study the convergence
of these aforementioned Gaussian approximate gradient flows in three settings:
the Gaussian posterior case; logconcave posterior case; and general posterior
case.

####  4.5.1 Gaussian Posterior Case

Assume the posterior distribution Eq. 1 is Gaussian so that $\rho_{\rm
post}(\theta)\sim\mathcal{N}(m_{\star},C_{\star})$ where

(97)
$\displaystyle\Phi_{R}(\theta)=\frac{1}{2}(x-m_{\star})^{T}C_{\star}^{-1}(x-m_{\star}).$

######  Proposition 4.8.

Consider the posterior distribution Eq. 1 under assumption Eq. 97 so that the
posterior is Gaussian. Then the Gaussian approximate Fisher-Rao gradient flow
Eq. 82 has the analytical solution

(98a) $\displaystyle
m_{t}=m_{\star}+e^{-t}\Bigl{(}(1-e^{-t})C_{\star}^{-1}+e^{-t}C_{0}^{-1}\Bigr{)}^{-1}C_{0}^{-1}\bigl{(}m_{0}-m_{\star}\bigr{)},$
(98b) $\displaystyle
C_{t}^{-1}=C_{\star}^{-1}+e^{-t}\bigl{(}C_{0}^{-1}-C_{\star}^{-1}\bigr{)}.$

The proof is in Section E.1. We remark that both mean and covariance converge
exponentially fast to $m_{\star}$ and $C_{\star}$ with convergence rate
$\mathcal{O}(e^{-t})$. This rate is independent of $C_{\star}$.

For the Gaussian approximate gradient flow Eq. 96 and the Gaussian approximate
Wasserstein gradient flow LABEL:eq:Gaussian-Wasserstein, if the norm of
$C_{\star}$ is large, their convergence rate is much slower than the Gaussian
approximate Fisher-Rao gradient flow. Indeed, we have the following
convergence result:

######  Proposition 4.9.

Consider the posterior distribution Eq. 1 under assumption Eq. 97 so that the
posterior is Gaussian; this posterior is the unique minimizer of the Gaussian
variational inference problem Eq. 9. Denote the largest eigenvalue of
$C_{\star}$ by $\lambda_{\star,\max}$. For gradient flows with initialization
$C_{0}=\lambda_{0}$, the following hold:

  1. 1.

for the Gaussian approximate gradient flow Eq. 96:

$\|m_{t}-m_{\star}\|_{2}=\mathcal{O}(e^{-t/\lambda_{\star,\max}}),\|C_{t}-C_{\star}\|_{2}=\mathcal{O}(e^{-t/(2\lambda_{\star,\max}^{2})});$

  2. 2.

for the Gaussian approximate Fisher-Rao gradient flow Eq. 82:

$\|m_{t}-m_{\star}\|_{2}=\mathcal{O}(e^{-t}),\|C_{t}-C_{\star}\|_{2}=\mathcal{O}(e^{-t});$

  3. 3.

for the Gaussian approximate Wasserstein gradient flow LABEL:eq:Gaussian-
Wasserstein:

$\|m_{t}-m_{\star}\|_{2}=\mathcal{O}(e^{-t/\lambda_{\star,\max}}),\|C_{t}-C_{\star}\|_{2}=\mathcal{O}(e^{-2t/\lambda_{\star,\max}}),$

where the implicit constants depend on $m_{\star}$, $C_{\star}$ and
$\lambda_{0}$.

The proof is in Section E.2.

####  4.5.2 Logconcave Posterior Case

In this subsection, we consider the case that the posterior distribution
$\rho_{\rm post}(\theta)$ given by Eq. 1 is strongly log-concave.

######  Proposition 4.10.

Assume that the posterior distribution $\rho_{\rm post}(\theta)$ is
$\alpha$-strongly logconcave (Definition 3.13); moreover assume that
$-\nabla_{\theta}\nabla_{\theta}\log\rho_{\rm post}\preceq\beta$. Assume
further that the initial covariance matrix satisfies $\lambda_{0,\min}$. Then
for the dynamics Eq. 96, the Gaussian approximate Fisher-Rao gradient flow Eq.
82, and the Gaussian approximate Wasserstein gradient flow LABEL:eq:Gaussian-
Wasserstein, we have

(99) $\displaystyle{\mathrm{KL}}\Bigl{[}\rho_{a_{t}}(\theta)\Big{\|}\rho_{\rm
post}(\theta)\Bigr{]}\leq
e^{-Kt}{\mathrm{KL}}\Bigl{[}\rho_{a_{0}}(\theta)\Big{\|}\rho_{\rm
post}(\theta)\Bigr{]}+(1-e^{-Kt}){\mathrm{KL}}\Bigl{[}\rho_{a_{\star}}\Big{\|}\rho_{\rm
post}(\theta)\Bigr{]},$

where $\rho_{a_{0}}\sim\mathcal{N}(m_{0},C_{0})$ is the initial condition and
$\rho_{a_{\star}}\sim\mathcal{N}(m_{\star},C_{\star})$ is the unique global
minimizer of Eq. 9. The rate constant $K$ depends on
$\alpha,\beta,\lambda_{0,\min},\lambda_{0,\max}$. Specifically, we have
$K=2\alpha/\max\\{1,4/\alpha,4\lambda_{0,\max}\\}$ for the Gaussian
approximate gradient flow Eq. 96; $K=\alpha\min\\{1/\beta,\lambda_{0,\min}\\}$
for the Gaussian approximate Fisher-Rao gradient flow Eq. 82; and $K=2\alpha$
for the Gaussian approximate Wasserstein gradient flow LABEL:eq:Gaussian-
Wasserstein. In addition, we also have that $\rho_{a_{t}}$ converges to
$\rho_{a_{\star}}$ exponentially fast in terms of the Wasserstein metric:

(100) $\displaystyle
W^{2}_{2}(\rho_{a_{t}},\rho_{a_{\star}})\leq\frac{2e^{-Kt}}{\alpha}\left({\mathrm{KL}}[\rho_{a_{0}}\|\rho_{\rm
post}(\theta)]-{\mathrm{KL}}[\rho_{a_{\star}}\|\rho_{\rm
post}(\theta)]\right).$

The proof is in Section E.3, and is inspired by the work [80].

On the other hand, we have the following proposition about the lower bound of
the local convergence rate of the Gaussian approximate Fisher-Rao gradient
flow Eq. 82.

######  Proposition 4.11.

Assume the posterior distribution $\rho_{\rm post}(\theta)$ is
$\alpha$-strongly logconcave (Definition 3.13); moreover assume that
$-\nabla_{\theta}\nabla_{\theta}\log\rho_{\rm post}\preceq\beta$. Denote the
unique minimizer of the Gaussian variational inference porblem Eq. 9 by
$\rho_{a_{\star}}:=\mathcal{N}(m_{\star},C_{\star})$. For $N_{\theta}=1$, let
$\lambda_{\star,\max}<0$ denote the largest eigenvalue of the linearized
Jacobian matrix of the Gaussian approximate Fisher-Rao gradient flow Eq. 82
around $m_{\star}$ and $C_{\star}$; this number determines the local
convergence rate of the Gaussian approximate Fisher-Rao gradient flow Eq. 82.
Then we have

$\lambda_{\star,\max}\leq-\frac{1}{(7+\frac{4}{\sqrt{\pi}})\bigl{(}1+\log(\frac{\beta}{\alpha})\bigr{)}}.$

Moreover, the bound is sharp: it is possible to construct a sequence of
$\rho_{\rm post}$, where the infimum of the negative eigenvalues is
$\mathcal{O}(1/(1+\log\frac{\beta}{\alpha}))$.

The proof can be found in Section E.4.

####  4.5.3 General Posterior Case

The previous sections consider Gaussian and then logconcave posteriors; for
all three Gaussian approximate gradient flows we demonstrate exponential
convergence, and for the Fisher-Rao based methodology we have some invariance
of the rates of convergence with respect to conditioning of the problem. In
this subsection, however, we construct counterexamples showing that, for
general posteriors, the convergence of all three Gaussian approximate gradient
flows to a stationary point can be arbitrarily slow.

######  Proposition 4.12.

For any $K>0$ there exist targets $\rho_{\rm post}$ such that, for the three
Gaussian approximate gradient flows Eq. 82, LABEL:eq:Gaussian-Wasserstein and
Eq. 96, the convergence to their stationary points can be as slow as
$\mathcal{O}(t^{-\frac{1}{2K}})$.

The proof is in Section E.5.

##  5 Numerical Experiments

In this section, we perform numerical experiments to study the behaviour of
the aforementioned gradient flows for sampling, which complements our
theoretical study. We observe the following:

  * •

In the probability density space, affine invariant gradient flows outperform
their non-affine invariant counterparts, for the Gaussian posterior case (Fig.
1), the logconcave posterior case (Fig. 3), and the general posterior case
(Fig. 6).

  * •

In the restricted Gaussian density space, affine invariant gradient flows
outperform their non-affine invariant counterparts, for the Gaussian posterior
case (Fig. 2), the logconcave posterior case (Fig. 4), and the general
posterior case (Fig. 8).

  * •

For general non-Gaussian posteriors, the convergence rates of all gradient
flows deteriorate when the posterior becomes more anisotropic (Fig. 6);
consequently accurately estimating the summary statistics of these posteriors
is challenging.

  * •

The convergence curves from the use of affine invariant Wasserstein gradient
flow, implemented with Langevin dynamics, oscillate slightly due to the added
noise; those obtained from affine invariant Stein gradient flow, implemented
by Stein variational gradient descent, are smooth (See Figs. 1, 3, and 6).
However, the added noise helps for sampling non-Gaussian, highly anisotropic
posteriors in comparison with the affine invariant Stein variational gradient
descent (See Figs. 5 and 6).

In the following subsections, we first introduce all test problems in Section
5.1 and the setup for the numerical methods in Section 5.2. Then we present
numerical results for the Gaussian posterior case in Section 5.3, the
logconcave posterior case in Section 5.4, and the general posterior case in
Section 5.5.

###  5.1 Overview of Test Problems

We focus our experiments on three two-dimensional posteriors. In defining them
we use the notation $\theta=[\theta^{(1)},\theta^{(2)}]^{T}$.

  1. 1.

_Gaussian Posterior._

$\Phi_{R}(\theta)=\frac{1}{2}\theta^{T}\begin{bmatrix}1&0\\\
0&\lambda\end{bmatrix}\theta\quad\textrm{with}\quad\lambda=0.01,\,0.1,\,1.$

We initialize the gradient flows from

$\theta_{0}\sim\mathcal{N}\Bigl{(}\begin{bmatrix}10\\\
10\end{bmatrix},\begin{bmatrix}\frac{1}{2}&0\\\ 0&2\end{bmatrix}\Bigr{)}.$

  2. 2.

_Logconcave Posterior_.

$\Phi_{R}(\theta)=\frac{(\sqrt{\lambda}\theta^{(1)}-\theta^{(2)})^{2}}{20}+\frac{(\theta^{(2)})^{4}}{20}\quad\textrm{with}\quad\lambda=0.01,\,0.1,\,1.$

We initialize the gradient flows from
$\theta_{0}\sim\mathcal{N}\Bigl{(}\begin{bmatrix}10\\\
10\end{bmatrix},\begin{bmatrix}4&0\\\ 0&4\end{bmatrix}\Bigr{)}$.

  3. 3.

_General Posterior._

$\Phi_{R}(\theta)=\frac{\lambda(\theta^{(2)}-(\theta^{(1)})^{2})^{2}}{20}+\frac{(1-\theta^{(1)})^{2}}{20}\quad\textrm{with}\quad\lambda=0.01,\,0.1,\,1.$

This example is known as the Rosenbrock function [58]. We initialize the
gradient flows from

$\theta_{0}\sim\mathcal{N}\Bigl{(}\begin{bmatrix}0\\\
0\end{bmatrix},\begin{bmatrix}4&0\\\ 0&4\end{bmatrix}\Bigr{)}.$

The summary statistics that we use to compare the resulting solution with the
ground truth are the expectation $\mathbb{E}[\theta]$, the covariance
$\mathrm{Cov}[\theta]$, and $\mathbb{E}[\cos(\omega^{T}\theta+b)]$; in the
latter case we randomly draw $\omega\sim\mathcal{N}(0,I)$ and
$b\sim\textrm{Uniform}(0,2\pi)$ and report the average MSE over 20 random
draws of $\omega$ and $b$. The ground truths of these summary statistics are
evaluated by integrating $\rho_{\rm post}$ numerically (See Appendix F for
details).

###  5.2 Numerical Method Setup

The gradient flows in the nonparametric density space studied in Section 3 are
implemented by interacting particle systems with $J=1000$ particles. Since we
are not aware of a straightforward way of implementing the Fisher-Rao gradient
flow that does not suffer from the immobility of the support, we do not
consider it here. We compare the following nonparametric gradient flows (GFs)
in our experiments.

  * •

Wasserstein GF: The Wasserstein gradient flow with $P=I$, which is implemented
as stochastic particle dynamics Eq. 48 with $\sigma_{t}=\sqrt{2}$.

  * •

Affine-invariant Wasserstein GF: The affine-invariant Wasserstein gradient
flow with $P=C_{t}$, which is implemented as stochastic particle dynamics Eq.
50.

  * •

Stein GF: The Stein gradient flow with

$P=I,\quad\kappa(\theta,\theta^{\prime},\rho)=(1+4\log(J+1)/N_{\theta})^{N_{\theta}/2}\exp(-\frac{1}{h}\lVert\theta-\theta^{\prime}\rVert^{2}),$

which is implemented as deterministic particle dynamics Eq. 59. Here
$h=\textrm{med}^{2}/\log(J+1)$ and $\textrm{med}^{2}$ is the squared median of
the pairwise Euclidean distance between the current particles, following [89].

  * •

Affine-invariant Stein GF: The affine-invariant Stein gradient flow with

$P=C,\quad\kappa(\theta,\theta^{\prime},\rho)=(1+2)^{N_{\theta}/2}\exp(-\frac{1}{2}(\theta-\theta^{\prime})^{T}C^{-1}(\theta-\theta^{\prime})),$

which is implemented as deterministic particle dynamics Eq. 60.

Here the scaling constants in the definition of kernels are chosen such that

(101)
$\displaystyle\int\int\kappa(\theta,\theta^{\prime},\rho)\mathcal{N}(\theta,m,C)\mathcal{N}(\theta^{\prime},m,C)\mathrm{d}\theta\mathrm{d}\theta^{\prime}=1.$

This choice makes the Stein gradient flows comparable with the Wasserstein
gradient flows in terms of implementation cost per step. Since we cannot
analytically compute the integral Eq. 101 for the kernel of Stein GF, we
estimate its scaling constant by replacing $\textrm{med}^{2}$ with
$N_{\theta}C$.

For the gradient flows in the Gaussian density space we consider the three
mean and covariant dynamics given in equations Eq. 82, LABEL:eq:Gaussian-
Wasserstein and Eq. 96. The expectations in the evolution equations are
calculated by the unscented transform [69] with $J=2N_{\theta}+1=5$ quadrature
points. Therefore, Gaussian approximation has considerable speedup in
comparison with the previously mentioned particle-based sampling approaches,
where $J=1000$.

###  5.3 Gaussian Posterior Case

The convergence of different gradient flows, according to the three summary
statistics, are presented in Fig. 1 (nonparameteric density space) and in Fig.
2 (Gaussian density space). In both nonparametric and Gaussian density spaces,
the imposition of the affine invariance property makes the convergence rate
independent of the anisotropy $\lambda$ and accelerates the sampling for badly
scaled Gaussian ($\lambda=0.01$). However, all these gradient flows in the
nonparametric density space do not converge within machine precision because
of the limited number of particles. The convergence rates of Gaussian
approximate gradient flows match well with the predictions of Proposition 4.9.

Figure 1: Gaussian posterior case: convergence of different gradient flows in
terms of the $L_{2}$ error of $\mathbb{E}[\theta_{t}]$, the relative Frobenius
norm error of the covariance
$\frac{\lVert\mathrm{Cov}[\theta_{t}]-\mathrm{Cov}[\theta_{{\rm
true}}]\rVert_{F}}{\lVert\mathrm{Cov}[\theta_{\rm true}]\rVert_{F}}$, and the
error of $\mathbb{E}[\cos(\omega^{T}\theta_{t}+b)]$.  Figure 2: Gaussian
posterior case: convergence of different dynamics in terms of $L_{2}$ error of
$\mathbb{E}[\theta_{t}]$, the relative Frobenius norm error of the covariance
$\frac{\lVert\mathrm{Cov}[\theta_{t}]-\mathrm{Cov}[\theta_{{\rm
true}}]\rVert_{F}}{\lVert\mathrm{Cov}[\theta_{\rm true}]\rVert_{F}}$, and the
error of $\mathbb{E}[\cos(\omega^{T}\theta_{t}+b)]$.

###  5.4 Logconcave Posterior Case

The convergence of different gradient flows, according to the three summary
statistics, are presented in Fig. 3 (nonparameteric density space) and in Fig.
4 (Gaussian density space). In both nonparametric and Gaussian density spaces,
the imposition of the affine invariance property makes the convergence rate
independent of the anisotropy $\lambda$ and accelerates the sampling in the
highly anisotropic case ($\lambda=0.01$). We observe that the corresponding
Gaussian approximate gradient flows can reach lower errors for this case with
the present numerical method setup defined in Section 5.2. We also observe
that the convergence rate of the Gaussian approximate Fisher-Rao gradient flow
does not deteriorate with increased anisotropy constant $\lambda$; this
indicates that the convergence rate in Proposition 4.10, for this gradient
flow, may not be tight.

Figure 3: Logconcave posterior case: convergence of different gradient flows
in terms of the $L_{2}$ error of $\mathbb{E}[\theta_{t}]$, the relative
Frobenius norm error of the covariance
$\frac{\lVert\mathrm{Cov}[\theta_{t}]-\mathrm{Cov}[\theta_{{\rm
true}}]\rVert_{F}}{\lVert\mathrm{Cov}[\theta_{\rm true}]\rVert_{F}}$, and the
error of $\mathbb{E}[\cos(\omega^{T}\theta_{t}+b)]$. Figure 4:  Logconcave
posterior case: convergence of different dynamics in terms of $L_{2}$ error of
$\mathbb{E}[\theta_{t}]$, the relative Frobenius norm error of the covariance
$\frac{\lVert\mathrm{Cov}[\theta_{t}]-\mathrm{Cov}[\theta_{{\rm
true}}]\rVert_{F}}{\lVert\mathrm{Cov}[\theta_{\rm true}]\rVert_{F}}$, and the
error of $\mathbb{E}[\cos(\omega^{T}\theta_{t}+b)]$

###  5.5 General Posterior Case

We note that the Rosenbrock function is a non-convex function. Although its
minimizer is at $[1,1]$, the expectation and covariance of the posterior
density function is (See Appendix F)

$\mathbb{E}[\theta]=\begin{bmatrix}1\\\
11\end{bmatrix}\qquad\mathrm{Cov}[\theta]=\begin{bmatrix}10&20\\\
20&\frac{10}{\lambda}+240\end{bmatrix}.$

The particles obtained by different nonparameteric gradient flows at $t=15$
are depicted in Fig. 5, and their convergences according to the three summary
statistics are depicted in Fig. 6. Estimated posterior densities (3 standard
deviations) obtained by different Gaussian approximate gradient flows are
presented in Fig. 7, and their convergences according to the three summary
statistics are depicted in Fig. 8. For small $\lambda$ (e.g., $\lambda=0.01$),
${\theta^{(2)}}$ is the stretch direction, and therefore the imposition of the
affine invariance property makes the convergence faster. However, when
$\lambda$ increases, the posterior density concentrates on a manifold with
significant curvature (See Fig. 5). Although the particle positions match well
with the density contours, the convergence of different gradient flows
significantly deteriorates; the imposition of affine invariance does not
relieve the situation. Furthermore, the Gaussian approximation cannot
represent the posterior distribution at all well because the posterior is far
from Gaussian.

Figure 5: General posterior case: particles obtained by different gradient
flows at $t=15$. Grey lines represent the contour of the true posterior.
Figure 6: General posterior case: convergence of different gradient flows in
terms of the $L_{2}$ error of $\mathbb{E}[\theta_{t}]$, the relative Frobenius
norm error of the covariance
$\frac{\lVert\mathrm{Cov}[\theta_{t}]-\mathrm{Cov}[\theta_{{\rm
true}}]\rVert_{F}}{\lVert\mathrm{Cov}[\theta_{\rm true}]\rVert_{F}}$, and the
error of $\mathbb{E}[\cos(\omega^{T}\theta_{t}+b)]$. Figure 7: General
posterior case: density functions (3 standard deviations) obtained by
different dynamics at $t=15$. Grey lines represent the contour of the true
posterior.  Figure 8: General posterior case: convergence of different
dynamics in terms of $L_{2}$ error of $\mathbb{E}[\theta_{t}]$, the relative
Frobenius norm error of the covariance
$\frac{\lVert\mathrm{Cov}[\theta_{t}]-\mathrm{Cov}[\theta_{{\rm
true}}]\rVert_{F}}{\lVert\mathrm{Cov}[\theta_{\rm true}]\rVert_{F}}$, and the
error of $\mathbb{E}[\cos(\omega^{T}\theta_{t}+b)]$.

##  6 Conclusions

###  6.1 Summary

In this work, we have studied various gradient flows in both nonparametric and
Gaussian density spaces for sampling distributions with unknown normalization
constants, focusing on the affine invariance property. We introduce the
concept of affine invariant metric and use it to develop general affine
invariant Wasserstein, Stein, and Fisher-Rao gradient flows in both
nonparametric and Gaussian density spaces. We provide theoretical analysis of
these gradient flows and demonstrate that affine-invariance effectively
improves the efficiency of sampling log-concave distributions. Numerically, we
demonstrate that the affine-invariance property can accelerate the convergence
of gradient flows for highly anisotropic distributions.

Nevertheless, these strategies may still not perform well for general
posterior distributions. In particular for multimodal distributions, or
distributions which concentrate on manifolds with significant curvature such
as the Rosenbrock function used here. In a companion work [29], we are
exploring the direction of approximating the Fisher-Rao Gradient flow using
Gaussian mixtures; this has the potential to capture multiple dominant modes
efficiently. We are also interested in exploring other invariant properties
and approximations that could deal with these more complex distributions. In
addition, for high dimensional problems, it is of interest to develop a
systematic study of model reduction of these gradient flows; see for example
the projected Stein and Wasserstein gradient flows studied in [27, 131].

Finally we highlight that, for Bayesian inverse problems [116] the methods
developed here do not exploit the structure of the forward problem. This is in
contrast to ensemble Kalman based methods which have demonstrable performance
advantages for some problems in this class [20]; notably these ensemble
methods have derivative free implementations which are favorable in large