Marc Lambert
email: marc.lambert@inria.frsinhochewimit

Sinho Chewi
email: schewi@mit.edu

Francis S Bach
email: francis.bach@inria.fr

Silvère Bonnabel
email: silvere.bonnabel@minesparis.psl.eu

Philippe Rigollet
email: rigollet@math.mit.edu

Variational

Marc Lambert Dga

Variational inference via Wasserstein gradient flows

teaching and research institutions in France or abroad, or from public or private research centers.

Introduction

This work brings together three active research areas: variational inference, variational Kalman filtering, and gradient flows on the Wasserstein space.

Variational inference. The development of large-scale Bayesian methods has fueled the need for fast and scalable methods to approximate complex distributions. More specifically, Bayesian methodology typically generates a high-dimensional posterior distribution π ∝ exp(-V ) that is known only up to normalizing constants, making the computation even of simple summary statistics such as the mean and covariance a major computational hurdle. To overcome this limitation, two distinct computational approaches are largely favored. The first approach consists of Markov chain Monte Carlo (mcmc) methods that rely on carefully constructed Markov chains which (approximately) converge to π. For example, the Langevin diffusion

dX t = -∇V (X t ) dt + √ 2 dB t , (1)
where (B t ) t≥0 denotes standard Brownian motion on R d , admits π as a stationary distribution.

Crucially, the Langevin diffusion can be discretized and implemented without knowledge of the normalizing constant of π, leading to practical algorithms for Bayesian inference. Recent theoretical efforts have produced sharp non-asymptotic convergence guarantees for algorithms based on the Langevin diffusion (or variants thereof), with many results known when π is strongly log-concave or satisfies isoperimetric assumptions [see, e.g., [START_REF] Durmus | Analysis of Langevin Monte Carlo via convex optimization[END_REF][START_REF] Shen | The randomized midpoint method for log-concave sampling[END_REF][START_REF] Vempala | Rapid convergence of the unadjusted Langevin algorithm: isoperimetry suffices[END_REF], Chen et al., 2020[START_REF] Dalalyan | On sampling from a log-concave density using kinetic Langevin diffusions[END_REF][START_REF] Chewi | Analysis of Langevin Monte Carlo from Poincaré to log-Sobolev[END_REF][START_REF] Tat | Structured logconcave sampling with a restricted Gaussian oracle[END_REF][START_REF] Ma | Is there an analog of Nesterov acceleration for gradient-based MCMC?[END_REF], Wu et al., 2022].

More recently, Variational Inference (vi) has emerged as a viable alternative to mcmc [START_REF] Michael | An introduction to variational methods for graphical models[END_REF], Wainwright and Jordan, 2008[START_REF] Blei | Variational inference: A review for statisticians[END_REF]. The goal of vi is to approximate the posterior π by a more tractable distribution π ∈ P such that π ∈ arg min p∈P KL(p π) .

(2)

A common example arises when P is the class of product distributions, in which case π is called the mean-field approximation of P. Unfortunately, by definition, mean-field approximations fail to capture important correlations present in the posterior π, and various remedies have been proposed, with varied levels of success. In this paper, we largely focus on obtaining a Gaussian approximation to π, that is, we take P to be the class of non-degenerate Gaussian distributions on R d [START_REF] Barber | Ensemble learning for multi-layer networks[END_REF][START_REF] Seeger | Bayesian model selection for support vector machines, Gaussian processes and other kernel classifiers[END_REF][START_REF] Honkela | Unsupervised variational Bayesian learning of nonlinear models[END_REF][START_REF] Opper | The variational Gaussian approximation revisited[END_REF], Zhang et al., 2018]. The expressive power of the variational model may then be further increased by considering mixture distributions [START_REF] Lin | Fast and simple natural-gradient variational inference with mixture of exponential-family approximations[END_REF], Daudel and Douc, 2021, Daudel et al., 2021].

Although the solution π of ( 2) is no longer equal to the true posterior, variational inference remains heavily used in practice because the problem (2) can be solved for simple models P via scalable optimization algorithms. In particular, vi avoids many of the practical hurdles associated with mcmc methods-such as the potentially long "burn-in" period of samplers and the lack of effective stopping criteria for the algorithm-while still producing informative summary statistics. In this regard, we highlight the fact that obtaining an approximation for the covariance matrix of π via mcmc methods requires drawing potentially many samples, whereas for many choices of P (e.g., the Gaussian approximation) the covariance matrix of π can be directly obtained from the solution to the vi problem (2).

Figure 1: Left: randomly initialized mixture of 20 Gaussians (the initial covariances are depicted as red circles) and contour plot of a logistic target π. Right: contour lines of a mixture of Gaussians approximation π obtained from the gradient flow in Section 5.

However, in contrast with mcmc methods, to date there have not been many theoretical guarantees for vi, even when π is strongly logconcave and P is taken to be the class of Gaussians N (m, Σ). The problem stems from the fact that the objective in ( 2) is typically non-convex in the pair (m, Σ). Obtaining such guarantees remains a pressing challenge for the field.

Variational Kalman filtering. There is also considerable interest in extending ideas behind variational inference to dynamical settings of Bayesian inference. Consider a general framework where (π t ) t represents the marginal laws of a stochastic process indexed by time t, which can be discrete or continuous. The goal is to recursively build a Gaussian approximation to (π t ) t .

As a concrete example, suppose that (π t ) t≥0 denotes the marginal law of the solution to the Langevin diffusion (1). In the context of Bayesian optimal filtering and smoothing, [START_REF] Särkkä | On unscented Kalman filtering for state estimation of continuous-time nonlinear systems[END_REF] proposed the following heuristic. Let (m t , Σ t ) denote the mean and covariance matrix of π t . Then, it can be checked (see Section B.4) that ṁt = -E ∇V (X t ) Σt = 2I -E[∇V (X t ) ⊗ (X t -m t ) + (X t -m t ) ⊗ ∇V (X t )]

(3) where X t ∼ π t . These ordinary differential equations (ODEs) are intractable because they involve expectations under the law of X t ∼ π t , which is not available to the practitioner. However, if we replace X t ∼ π t with a Gaussian Y t ∼ p t = N (m t , Σ t ) with the same mean and covariance as X t , then the system of ODEs

ṁt = -E ∇V (Y t ) Σt = 2I -E[∇V (Y t ) ⊗ (Y t -m t ) + (Y t -m t ) ⊗ ∇V (Y t )] (4)
yields a well-defined evolution of Gaussian distributions (p t ) t≥0 , which we may optimistically believe to be a good approximation of (π t ) t≥0 . Moreover, the system of ODEs can be numerically approximated efficiently in practice using Gaussian quadrature rules to compute the above expectations. This is the principle behind the unscented Kalman filter [START_REF] Julier | A new method for the nonlinear transformation of means and covariances in filters and estimators[END_REF].

In the context of the Langevin diffusion, Särkkä's heuristic (4) provides a promising avenue towards computational vi. Indeed, since π ∝ exp(-V ) is the unique stationary distribution of the Langevin diffusion (1), an algorithm to approximate (π t ) t≥0 is expected to furnish an algorithm to solve the VI problem (2). However, at present there is little theoretical understanding of how the system (4) approximates (3); moreover, Särkkä's heuristic only provides Gaussian approximations, and it is unclear how to extend the system (4) to more complex models (e.g., mixtures of Gaussians).

Our contributions: bridging the gap via Wasserstein gradient flows. We show that the approximation (p t ) t≥0 in Särkkä's heuristic (4) arises precisely as the gradient flow of the Kullback-Leibler (KL) divergence KL(• π) on the Bures-Wasserstein space of Gaussian distributions on R d endowed with the 2-Wasserstein distance from optimal transport [START_REF] Villani | Topics in optimal transportation[END_REF]. This perspective allows us to not only understand its convergence but also to extend it to the richer space of mixtures of Gaussian distributions, and propose an implementation as a novel system of interacting "Gaussian particles". Below, we proceed to describe our contributions in greater detail.

Our framework builds upon the seminal work of [START_REF] Jordan | The variational formulation of the Fokker-Planck equation[END_REF], which introduced the celebrated JKO scheme in order to give meaning to the idea that the evolving marginal law of the Langevin diffusion (1) is a gradient flow of KL(• π) on the Wasserstein space P 2 (R d ) of probability measures with finite second moments. Subsequently, in order to emphasize the Riemannian geometry underlying this result, [START_REF] Otto | The geometry of dissipative evolution equations: the porous medium equation[END_REF] developed his eponymous calculus on P 2 (R d ), a framework which has had tremendous impact in analysis, geometry, PDE, probability, and statistics.

Inspired by this perspective, we show in Theorem 1 that Särkkä's approximation (p t ) t≥0 is also a gradient flow of KL(• π), with the main difference being that it is constrained to lie on the submanifold BW(R d ) of P 2 (R d ) consisting of Gaussian distributions, known as the Bures-Wasserstein manifold. In turn, our result paves the way for new theoretical understanding via the powerful theory of gradient flows. As a first step, using well-known results about convex functionals on the Wasserstein space, we show in Corollary 1 that (p t ) t≥0 converges rapidly to the solution of the vi problem (2) with P = BW(R d ) as soon as V is convex. Moreover, in Section 4.1, we apply numerical integration based on cubature rules for Gaussian integrals to the system of ODEs (4), thus arriving at a fast method with robust empirical performance (details in Sections I and J). This combination of results brings vi closer to Langevin-based mcmc both on the practical and theoretical fronts, but still falls short of achieving non-asymptotic discretization guarantees as pioneered by [START_REF] Arnak | Theoretical guarantees for approximate sampling from smooth and log-concave densities[END_REF] for mcmc. To further close the theoretical gap between vi and the state of the art for mcmc, we propose in Section 4.2 a stochastic gradient descent (SGD) algorithm as a time discretization of the Bures-Wasserstein gradient flow. This algorithm comes with convergence guarantees that establish vi as a solid competitor to mcmc not only from a practical standpoint but also from a theoretical one. Both have their relative merits; whereas mcmc targets the true posterior, vi leads to fast computation of summary statistics of the approximation π to π.

In Section 5, we consider an extension of these ideas to the substantially more flexible class of mixtures of Gaussians. Namely, the space of mixtures of Gaussians can be identified as a Wasserstein space over BW(R d ) and hence inherits Otto's differential calculus. Leveraging this viewpoint, in Theorem 3 we derive the gradient flow of KL(• π) over the space of mixtures of Gaussians and propose to implement it via a system of interacting particles. Unlike typical particle-based algorithms, here our particles correspond to Gaussian distributions, and the collection thereof to a Gaussian mixture which is better equipped to approximate a continuous measure. We validate the empirical performance of our method with promising experimental results (see Section J). Although we focus on the vi problem in this work, we anticipate that our notion of "Gaussian particles" may be a broadly useful extension of classical particle methods for PDEs.

Related work.

Classical vi methods define a parametric family P = {p θ : θ ∈ Θ} and minimize θ → KL(p θ π) over θ ∈ Θ using off-the-shelf optimization algorithms [START_REF] Paisley | Variational Bayesian inference with stochastic search[END_REF][START_REF] Ranganath | Black box variational inference[END_REF]. Since ( 2) is an optimization problem over the space of probability distributions, we argue for methods that respect a natural geometric structure on this space. In this regard, previous approaches to vi using natural gradients implicitly employ a different geometry [START_REF] Lin | Fast and simple natural-gradient variational inference with mixture of exponential-family approximations[END_REF][START_REF] Daniel | Efficient derivative-free Bayesian inference for large-scale inverse problems[END_REF][START_REF] Emtiyaz | The Bayesian learning rule[END_REF], namely the reparameterization-invariant Fisher-Rao geometry [START_REF] Amari | Methods of information geometry[END_REF]. The application of Wasserstein gradient flows to vi was introduced earlier in work on normalizing flows and Stein Variational Gradient Descent (SVGD) [START_REF] Liu | Stein variational gradient descent as gradient flow[END_REF]Wang, 2016, Liu, 2017].

Our work falls in line with a number of recent papers aiming to place vi on a solid theoretical footing [START_REF] Alquier | On the properties of variational approximations of Gibbs posteriors[END_REF], Wang and Blei, 2019[START_REF] Domke | Provable smoothness guarantees for black-box variational inference[END_REF][START_REF] Knoblauch | An optimization-centric view on Bayes' rule: reviewing and generalizing variational inference[END_REF]. Some of these works in particular have obtained non-asymptotic algorithmic guarantees for specific examples, see, e.g., [START_REF] Challis | Gaussian Kullback-Leibler approximate inference[END_REF], [START_REF] Alquier | Concentration of tempered posteriors and of their variational approximations[END_REF].

The connection between vi and Kalman filtering was studied in the static case by [START_REF] Lambert | The limited-memory recursive variational Gaussian approximation[END_REF]Lambert et al. [ , 2022a]], and extended to the dynamical case by Lambert et al. [2022b], providing a first justification of Särkkä's heuristic in terms of local variational Gaussian approximation. In particular, the closest linear process to the Langevin diffusion (1) is a Gaussian process governed by a McKean-Vlasov equation whose Gaussian marginals have parameters evolving according to Särkkä's ODEs.

Constrained gradient flows on the Wasserstein space have also been extensively studied [START_REF] Carlen | Constrained steepest descent in the 2-Wasserstein metric[END_REF][START_REF] Caglioti | On a constrained 2-D Navier-Stokes equation[END_REF][START_REF] Tudorascu | On a nonlinear, nonlocal parabolic problem with conservation of mass, mean and variance[END_REF][START_REF] Eberle | Gradient flow formulation and longtime behaviour of a constrained Fokker-Planck equation[END_REF], although our interpretation of Särkkä's heuristic is, to the best of our knowledge, new.

Background

In order to define gradient flows on the space of probability measures, we must first endow this space with a geometry; see Appendix B for more details. Given probability measures µ and ν on R d , define the 2-Wasserstein distance

W 2 (µ, ν) = inf γ∈C(µ,ν)
x -y 2 dγ(x, y)

1/2
, where C(µ, ν) is the set of couplings of µ and ν, that is, joint distributions on R d × R d whose marginals are µ and ν respectively. This quantity is finite as long as µ and ν belong to the space P 2 (R d ) of probability measures over R d with finite second moments. The 2-Wasserstein distance has the interpretation of measuring the smallest possible mean squared displacement of mass required to transport µ to ν; we refer to [START_REF] Villani | Topics in optimal transportation[END_REF][START_REF] Villani | Optimal transport, volume 338 of Grundlehren der Mathematischen Wissenschaften[END_REF], [START_REF] Santambrogio | Optimal transport for applied mathematicians[END_REF] for textbook treatments on optimal transport. Unlike other notions of distance between probability measures, such as the total variation distance, the 2-Wasserstein distance respects the geometry of the underlying space R d , leading to numerous applications in modern data science [see, e.g., [START_REF] Peyré | Computational optimal transport: with applications to data science[END_REF].

The space (P 2 (R d ), W 2 ) is a metric space [Villani, 2003, Theorem 7.3], and we refer to it as the Wasserstein space. However, as shown by Otto [START_REF] Otto | The geometry of dissipative evolution equations: the porous medium equation[END_REF], it has a far richer geometric structure: formally, (P 2 (R d ), W 2 ) can be viewed as a Riemannian manifold, a fact which allows for considering gradient flows of functionals on P 2 (R d ). A fundamental example of such a functional is the KL divergence KL(• π) to a target density π ∝ exp(-V ) on R d , for which [START_REF] Jordan | The variational formulation of the Fokker-Planck equation[END_REF] showed that the Wasserstein gradient flow is the same as the evolution of the marginal law of the Langevin diffusion (1). This optimization perspective has had tremendous impact on our understanding and development of mcmc algorithms [Wibisono, 2018].

Variational inference with Gaussians

In this section we describe our problem using two equivalent approaches: a variational approach based on a modified version of the JKO scheme of [START_REF] Jordan | The variational formulation of the Fokker-Planck equation[END_REF] (Section 3.1), and a Wasserstein gradient flow approach based on Otto calculus (Section 3.2). Both lead to the same result (Section 3.3). While the former is more accessible to readers who are unfamiliar with gradient flows on the Wasserstein space, the latter leads to strong convergence guarantees (Section 3.4).

Variational approach: the Bures-JKO scheme

The space of non-degenerate Gaussian distributions on R d equipped with the W 2 distance forms the Bures-Wasserstein space BW(R d ) ⊆ P 2 (R d ). On BW(R d ), the Wasserstein distance W 2 2 (p 0 , p 1 ) between two Gaussians p 0 = N (m 0 , Σ 0 ) and p 1 = N (m 1 , Σ 1 ) admits the following closed form:

W 2 2 (p 0 , p 1 ) = m 0 -m 1 2 + B 2 (Σ 0 , Σ 1 ) , (5)
where

B 2 (Σ 0 , Σ 1 ) = tr(Σ 0 + Σ 1 -2 (Σ 1 2 0 Σ 1 Σ 1 2 0 ) 1
2 ) is the squared Bures metric [START_REF] Bures | An extension of Kakutani's theorem on infinite product measures to the tensor product of semifinite w * -algebras[END_REF]. Given a target density π ∝ exp(-V ) on R d , and with a step size h > 0, we may define the iterates of the proximal point algorithm

p k+1,h := arg min p∈BW(R d ) KL(p π) + 1 2h W 2 2 (p, p k,h ) . (6)
Using (5), this is an explicit optimization problem involving the mean and covariance matrix of p. Although (6) is not solvable in closed form, by letting h 0 we obtain a limiting curve (p t ) t≥0 via p t = lim h 0 p t/h ,h , which can be interpreted as the Bures-Wasserstein gradient flow of the KL divergence KL(• π). This procedure mimics the JKO scheme [START_REF] Jordan | The variational formulation of the Fokker-Planck equation[END_REF] with the additional constraint that the iterates lie in BW(R d ), and we therefore call it the Bures-JKO scheme.

Geometric approach: the Bures-Wasserstein gradient flow of the KL divergence

In the formal sense of Otto described above, BW(R d ) is a submanifold of P 2 (R d ). Moreover, since Gaussians can be parameterized by their mean and covariance, BW(R d ) can be identified with the manifold R d × S d ++ , where S d ++ is the cone of symmetric positive definite d × d matrices. Hence, BW(R d ) is a genuine Riemannian manifold in its own right [see [START_REF] Modin | Geometry of matrix decompositions seen through optimal transport and information geometry[END_REF][START_REF] Malagò | Wasserstein Riemannian geometry of Gaussian densities[END_REF][START_REF] Bhatia | On the Bures-Wasserstein distance between positive definite matrices[END_REF], and gradient flows can be defined using Riemannian geometry [START_REF] Manfredo | Riemannian geometry. Mathematics: Theory & Applications[END_REF]. See Section B.3 for more details. Since the functional µ → F(µ) = KL(µ π) defined over P 2 (R d ) restricts to a functional over BW(R d ), we can also consider the gradient flow of F over the Bures-Wasserstein space; note that this latter gradient flow is necessarily a curve (p t ) t≥0 such that each p t is a Gaussian measure.

Variational inference via the Bures-Wasserstein gradient flow

Using either approach, we can prove the following theorem.

Theorem 1. Let π ∝ exp(-V ) be the target density on R d . Then, the limiting curve (p t ) t≥0 where p t = N (m t , Σ t ) is obtained via the Bures-JKO scheme (6), or equivalently, the Bures-Wasserstein gradient flow (p t ) t≥0 of the KL divergence KL(• π), satisfies Särkkä's system of ODEs (4).

Proof. The proof using the Bures-JKO scheme is given in Section A.1 and the proof using Otto calculus is presented in Section C. This theorem shows that Särkkä's heuristic (4) precisely yields the Wasserstein gradient flow of the KL divergence over the submanifold BW(R d ). Equipped with this interpretation, we are now able to obtain information about the asymptotic behavior of the approximation (p t ) t≥0 . Namely, we can hope that it converges to constrained minimizer π = arg min p∈BW(R d ) KL(p π), i.e., precisely the solution to the vi problem (2). In the next section, we show that this convergence in fact holds as soon as V is convex, and moreover with quantitative rates.

The solution π to (2), and consequently the limit point of Särkkä's approximation, is well-studied in the variational inference literature [see, e.g., [START_REF] Opper | The variational Gaussian approximation revisited[END_REF], and we recall standard facts about π here for completeness. It is known that π satisfies the equations

E π ∇V = 0 and E π ∇ 2 V = Σ-1 , ( 7
)
where Σ is the covariance matrix of π (these equations can also be derived as first-order necessary conditions by setting the Bures-Wasserstein gradient derived in Section C to zero). In particular, it follows from (7) that if ∇ 2 V enjoys the bounds αI ∇ 2 V βI for some -∞ ≤ α ≤ β ≤ ∞, then any solution π to the constrained problem also satisfies β -1 I Σ (α ∨ 0) -1 I.

Continuous-time convergence

Besides providing an intuitive interpretation of Särkkä's heuristic, Theorem 1 readily yields convergence criteria for the system (4) which rest upon general principles for gradient flows. We begin with a key observation. For a functional F : Lemma 1.

BW(R d ) → R ∪ {∞} and α ∈ R, we say that F is α-convex if for all constant-speed geodesics (p t ) t∈[0,1] in BW(R d ), F(p t ) ≤ (1 -t) F(p 0 ) + t F(p 1 ) - α t (1 -t) 2 W 2 2 (p 0 , p 1 ) , t ∈ [0, 1] .
For any α ∈ R, if ∇ 2 V αI, then KL(• π) is α-convex on BW(R d ).
Proof.

The assumption that ∇ 2 V αI entails that the functional KL(• π) is α-convex on the entire Wasserstein space (P 2 (R d ), W 2 ) [see, e.g., Villani, 2009, Theorem 17.15]. Since BW(R d ) is a geodesically convex subset of P 2 (R d ) (see Section B.3), then the geodesics in BW(R d ) agree with the geodesics in P 2 (R d ), from which it follows that KL(• π) is α-convex on BW(R d ).

Consequently, we obtain the following corollary. Its proof is postponed to Section D.

Corollary 1. Suppose that ∇ 2 V αI for some α ∈ R. Then, for any p 0 ∈ BW(R d ), there is a unique solution to the BW(R d ) gradient flow of KL(• π) started at p 0 . Moreover:

1. If α > 0, then for all t ≥ 0, W 2 2 (p t , π) ≤ exp(-2αt) W 2 2 (p 0 , π). 2. If α > 0, then for all t ≥ 0, KL(p t π) -KL(π π) ≤ exp(-2αt) {KL(p 0 π) -KL(π π)}. 3. If α = 0, then for all t > 0, KL(p t π) -KL(π π) ≤ 1 2t W 2 2 (p 0 , π).
The assumption that ∇ 2 V αI for some α > 0, i.e., that π is strongly log-concave, is a standard assumption in the mcmc literature. Under this same assumption, Corollary 1 yields convergence for the Bures-Wasserstein gradient flow of KL(• π); however, the flow must first be discretized in time for implementation. If we assume additionally that the smoothness condition ∇ 2 V βI holds, then a surge of recent research has succeeded in obtaining precise non-asymptotic guarantees for discretized mcmc algorithms. In Section 4.2 below, we will show how to do the same for vi.

Time discretization of the Bures-Wasserstein gradient flow

We are now equipped with dual perspectives on a dynamical solution to Gaussian vi: ODE and gradient flow. Each perspective leads to a different implementation. On the one hand, we discretize the system of ODEs defined in (4) using numerical integration. On the other, we discretize the gradient flow using stochastic gradient descent in the Bures-Wasserstein space.

Numerical integration of the ODEs

The system of ODEs (4) can be integrated in time using a classical Runge-Kutta scheme. The expectations under a Gaussian support are approximated by cubature rules used in Kalman filtering [START_REF] Arasaratnam | Cubature Kalman filters[END_REF].

Moreover, a square root version of the ODE is also considered to ensure that covariance matrices remain symmetric and positive. See Appendix I.2 for more details. We have tested our method on a bimodal distribution and on a posterior distribution arising from a logistic regression problem. We observe fast convergence as shown in Figure 2.

Bures-Wasserstein SGD and theoretical guarantees for VI

Although the ODE discretization proposed in the preceding section enjoys strong empirical performance, it is unclear how to quantify its impact on the convergence rates established in Corollary 1. Therefore, we now propose a stochastic gradient descent algorithm over the Bures-Wasserstein space, for which useful analysis tools have been developed [Chewi et al., 2020[START_REF] Altschuler | Averaging on the Bures-Wasserstein manifold: dimension-free convergence of gradient descent[END_REF]. This approach bypasses the use of the system of ODEs (4), and instead discretizes the Bures-Wasserstein gradient flow directly. Under the standard assumption of strong log-concavity and log-smoothness, it leads to an algorithm (Algorithm 1) for approximating π with provable convergence guarantees.

Algorithm 1 Bures-Wasserstein SGD Require: strong convexity parameter α > 0; step size h > 0; mean m 0 and covariance matrix

Σ 0 for k = 1, . . . , N do draw a sample Xk ∼ p k set m k+1 ← m k -h ∇V ( Xk ) set M k ← I -h (∇ 2 V ( Xk ) -Σ -1 k ) set Σ + k ← M k Σ k M k set Σ k+1 ← clip 1/α Σ + k end for
Algorithm 1 maintains a sequence of Gaussian distributions (p k ) k∈N ; here (m k , Σ k ) denote the mean vector and covariance matrix at iteration k (see Section E for a derivation of the algorithm as SGD in the Bures-Wasserstein space). The clipping operator clip τ , which is introduced purely for the purpose of theoretical analysis, simply truncates the eigenvalues from above; see Section E. Our theoretical result for vi is given as the following theorem, whose proof is deferred to Section E.

Theorem 2. Assume that 0 ≺ αI ∇ 2 V I. Also, assume that h ≤ α 2
60 and that we initialize Algorithm 1 at a matrix satisfying α 9 I Σ µ 0

1 α I. Then, for all k ∈ N, E W 2 2 (p k , π) ≤ exp(-αkh) W 2 2 (p 0 , π) + 36dh α 2 .
In particular, we obtain E W 2 2 (p k , π) ≤ ε 2 provided we set h α 2 ε 2 d and the number of iterations to be k

d α 3 ε 2 log(W 2 (p 0 , π)/ε).
The upper bound ∇ 2 V I is notationally convenient for our proof but not necessary; in any case, any strongly log-concave and log-smooth density π can be rescaled so that the assumption holds.

Theorem 2 is similar in flavor to modern results for mcmc, both in terms of the assumptions (Hessian bounds and query access to the derivatives1 of V ) and the conclusion (a non-asymptotic polynomial-time algorithmic guarantee). We hope that such an encouraging result for vi will prompt more theoretical studies aimed at closing the gap between the two approaches.

Variational inference with mixtures of Gaussians

Thus far, we have shown that the tractability of Gaussians can be readily exploited in the context of Bures-Wasserstein gradient flows and translated into useful results for variational inference. Nevertheless, these results are limited by the lack of expressivity of Gaussians, namely their inability to capture complex features such as multimodality and, more generally, heterogeneity. To overcome this limitation, mixtures of Gaussians arise as a natural and powerful alternative; indeed, universal approximation of arbitrary probability measures by mixtures of Gaussians is well-known [see, e.g., [START_REF] Delon | A Wasserstein-type distance in the space of Gaussian mixture models[END_REF]. As we show next, the space of mixtures of Gaussians can also be equipped with a Wasserstein structure which gives rise to implementable gradient flows.

Geometry of the space of mixtures of Gaussians

We begin with the key observation already made by [START_REF] Chen | Optimal transport for Gaussian mixture models[END_REF], that any mixture of Gaussians can be canonically identified with a probability distribution (the mixing distribution) over the parameter space Θ = R d × S d ++ (the space of means and covariance matrices). Explicitly a probability measure µ ∈ P(Θ) corresponds to a Gaussian mixture as follows:

µ ↔ p µ := p θ dµ(θ) , (8)
where p θ is the Gaussian distribution with parameters θ ∈ Θ. Equivalently, µ can be thought of as a probability measure over BW(R d ), and hence the space of Gaussian mixtures on R d can be identified with the Wasserstein space P 2 (BW(R d )) over the Bures-Wasserstein space which is endowed with the distance (5) between Gaussian measures. Indeed, the theory of optimal transport can be developed with any Riemannian manifold (rather than R d ) as the base space [START_REF] Villani | Optimal transport, volume 338 of Grundlehren der Mathematischen Wissenschaften[END_REF]. As before, the space P 2 (BW(R d )) is endowed with a formal Riemannian structure, which respects the geometry of the base space BW(R d ), and we can consider Wasserstein gradient flows over P 2 (BW(R d )).

Note that this framework encompasses both discrete mixtures of Gaussians (when µ is a discrete measure) and continuous mixtures of Gaussians. In the case when the mixing distribution µ is discrete, the geometry of P 2 (BW(R d )) was studied by [START_REF] Chen | Optimal transport for Gaussian mixture models[END_REF], [START_REF] Delon | A Wasserstein-type distance in the space of Gaussian mixture models[END_REF]. An important insight of our work, however, is that it is fruitful to consider the full space P 2 (BW(R d )) for deriving gradient flows, even if we eventually develop algorithms which propagate a finite number of mixture components.

Gradient flow of the KL divergence and particle discretization

We consider the gradient flow of the KL divergence functional

µ → F(µ) := KL(p µ π) (9)
over the space P 2 (BW(R d )). The proof of the following theorem is given in Section F.

Theorem 3. The gradient flow (µ t ) t≥0 of the functional F defined in (9) over P 2 (BW(R d )) can be described as follows. Let θ 0 = (m 0 , Σ 0 ) ∼ µ 0 , and let θ t = (m t , Σ t ) evolve according to the ODE

ṁt = -E ∇ ln p µt π (Y t ) Σt = -E ∇ 2 ln p µt π (Y t ) Σ t -Σ t E ∇ 2 ln p µt π (Y t ) (10)
where

Y t ∼ N (m t , Σ t ). Then θ t ∼ µ t .
The gradient flow in Theorem 3 describes the evolution of a particle θ t which describes the parameters of a Gaussian measure, hence the name Gaussian particle. The intuition behind this evolution is as follows. Suppose we draw infinitely many initial particles (each being a Gaussian) from µ 0 . By evolving all those particles through (10), which interact with each other via the term p µt , they tend to aggregate in some parts of the space of Gaussian parameters and spread out in others. This distribution of Gaussian particles is precisely the mixing measure µ t , which, in turn, corresponds to a Gaussian mixture. Since an infinite number of Gaussian particles is impractical, consider initializing this evolution at a finitely supported distribution µ 0 , thus corresponding to a more familiar Gaussian mixture model with a finite number of components:

µ 0 = 1 N N i=1 δ θ (i) 0 = 1 N N i=1 δ (m (i) 0 ,Σ (i) 0 ) ↔ p µ 0 := 1 N N i=1 p (m (i) 0 ,Σ (i) 0 ) .
Interestingly, it can be readily checked that the system of ODEs (10) thus initialized maintains a finite mixture distribution:

µ t = 1 N N i=1 δ θ (i) t = 1 N N i=1 δ (m (i) t ,Σ (i) t ) ,
where the parameters θ

(i) t = (m (i) t , Σ (i)
t ) evolve according to the following interacting particle system, for

i ∈ [N ] ṁ(i) t = -E ∇ ln p µt π (Y (i) t ) , (11)
Σ(i) t = -E ∇ 2 ln p µt π (Y (i) t ) Σ (i) t -Σ (i) t E ∇ 2 ln p µt π (Y (i) t ) , (12)
where

Y (i) t ∼ p θ (i) t
. This finite system of particles can now be implemented using the same numerical tools as for Gaussian vi, see Section J. Note that due to this property of the dynamics, we can hope at best to converge to the best mixture of N Gaussians approximating π, but this approximation error is expected to vanish as N → ∞. Also, similarly to (4), it is possible to write down Hessian-free updates using integration by parts, see Appendix A.2.

The above system of particles may also be derived using a proximal point method similar to the Bures-JKO scheme, see Section A.2. Indeed, infinitesimally, it has the variational interpretation

(θ (1) t+h , . . . , θ (N ) t+h ) ≈ arg min θ (1) ,...,θ (N ) ∈Θ KL 1 N N i=1 p θ (i) π + 1 2N h N i=1 W 2 2 (p θ (i) , p θ (i) t ) .
Reassuringly, Equations ( 11)-( 12) reduce to (4) when µ 0 = δ (m 0 ,Σ 0 ) is a point mass, indicating that the theorem provides a natural extension of our previous results. However, although the model ( 8) is substantially more expressive than the Gaussian vi considered in Section 3, it has the downside that we lose many of the theoretical guarantees. For example, even when V is convex, the objective functional F considered here need not be convex; see Section G. We nevertheless validate the practical utility of our approach in experiments (see Figure 3 and Section J).

Unlike typical interacting particle systems which arise from discretizations of Wasserstein gradient flows, at each time t, the distribution p µt is continuous. This extension provides considerably more flexibility-from a mixture of point masses to a mixture of Gaussians-compared to interacting particle-based algorithms hitherto considered for either sampling [START_REF] Liu | Stein variational gradient descent: a general purpose Bayesian inference algorithm[END_REF][START_REF] Liu | Stein variational gradient descent as gradient flow[END_REF][START_REF] Duncan | On the geometry of Stein variational gradient descent[END_REF], Chewi et al., 2020], or solving partial differential equations [START_REF] José | Globalin-time weak measure solutions and finite-time aggregation for nonlocal interaction equations[END_REF], 2012[START_REF] Bonaschi | Equivalence of gradient flows and entropy solutions for singular nonlocal interaction equations in 1D[END_REF][START_REF] Craig | A blob method for the aggregation equation[END_REF][START_REF] José | A blob method for diffusion[END_REF][START_REF] Craig | A blob method for inhomogeneous diffusion with applications to multi-agent control and sampling[END_REF].

Conclusion

Using the powerful theory of Wasserstein gradient flows, we derived new algorithms for vi using either Gaussians or mixtures of Gaussians as approximating distributions. The consequences are twofold. On the one hand, strong convergence guarantees under classical conditions contribute markedly to closing the theoretical gap between mcmc and Gaussian vi. On the other hand, discretization of the Wasserstein gradient flow for mixtures of Gaussians yields a new Gaussian particle method for time discretization which, unlike classical particle methods, maintains a continuous probability distribution at each time.

We conclude by briefly listing some possible directions for future study. For Gaussian variational inference, our theoretical result (Theorem 2) can be strengthened by weakening the assumption that π is strongly log-concave, or by developing algorithms which do not require Hessian information for V . For mixtures of Gaussians, it is desirable to design a principled algorithm which also allows for the mixture weights to be updated.

Towards the latter question, in Section H we derive the gradient flow of the KL divergence with respect to the Wasserstein-Fisher-Rao geometry [START_REF] Liero | Optimal transport in competition with reaction: the Hellinger-Kantorovich distance and geodesic curves[END_REF][START_REF] Chizat | An interpolating distance between optimal transport and Fisher-Rao metrics[END_REF][START_REF] Liero | Optimal entropy-transport problems and a new Hellinger-Kantorovich distance between positive measures[END_REF], which yields an interacting system of Gaussian particles with changing weights. The equations are given as follows: at each time t, the mixing measure is the discrete measure

µ t = N i=1 w (i) t δ (m (i) t ,Σ (i) t ) . Let Y (i) t ∼ N (m (i) t , Σ (i) t ), and let r (i) t = w (i)
t . Then, the system of ODEs is given by

ṁ(i) t = -E ∇ ln p µt π (Y (i) t ) , Σ(i) t = -E ∇ 2 ln p µt π (Y (i) t ) Σ (i) t -Σ (i) t E ∇ 2 ln p µt π (Y (i) t ) , ṙ(i) t = -E ln p µt π (Y (i) t ) - 1 N N j=1 E ln p µt π (Y (j) t ) r (i) t .
We have implemented these equations and their empirical performance is encouraging. However, a fuller investigation of algorithms for vi with changing weights is beyond the scope of this work and we leave it for future research. A Proofs via the Bures-JKO scheme

A.1 Proof of Theorem 1

Given a Gaussian distribution at time t written p t = N (m t , Σ t ) and a target distribution π, we seek the solution p at time t + h of the following JKO scheme, where p = N (m, Σ) is constrained to lie on the space BW(R d ) of Gaussians equipped with the Wasserstein distance:

min p∈BW(R d ) L(p) = KL(p π) + 1 2h W 2 2 (p, p t ) . ( 13
)
Using the expression for the Wasserstein distance W 2 2 (p t , p) given in (5) it is equivalent to finding the Gaussian parameters which solve:

min m,Σ L(m, Σ) = KL(N (m, Σ) π) + 1 2h m t -m 2 + 1 2h B 2 (Σ t , Σ) . ( 14
)
We first compute the critical points of L and then take the limit as h 0 to get the desired differential equations (ODEs) for the parameters. This boils down to computing the Wasserstein gradient flow of KL(• π) over the Bures-Wasserstein manifold BW(R d ).

The left KL divergence is a sum of two terms

KL(p π) = H(p) -E p [ln π], where H(p) is the negative entropy of a Gaussian. It satisfies ∇ m H(p) = 0 and ∇ Σ H(p) = ∇ Σ (-1 2 ln det Σ) = -1 2 Σ -1 .
To alleviate notation, for any function f we both let x denote its argument and E p [f (x)] denote expectation over x ∼ p = N (m, Σ) throughout the present proof, depending on the context.

The gradient of the left KL divergence with respect to m is given by:

∇ m KL(p π) = -∇ m E p [ln π(x)] = -E p [∇ x ln π(x)],
where we have used integration by parts (assuming π is continuously differentiable) and the property of Gaussian densities

∇ m N (x | m, Σ) = -∇ x N (x | m, Σ)
to get a derivative with respect to x. The critical point of L given by ( 14) w.r.t. the mean parameter m thus writes:

∇ m L(m, Σ) = 1 h (m -m t ) -E p [∇ x ln π(x)] = 0 . ( 15
)
Taking the limit as h 0, we find that m t must satisfy the following ODE:

ṁt = E pt [∇ x ln π(x)] = -E pt [∇ x V (x)] ,
where we recall that π ∝ exp(-V ). This recovers the first line of (4).

The gradient of the left KL divergence with respect to Σ is given by:

∇ Σ KL(p π) = - 1 2 Σ -1 -∇ Σ E p [ln π(x)] = - 1 2 Σ -1 - 1 2 E p [∇ 2 x ln π(x)] ,
where we have used two integrations by parts (supposing π is twice continuously differentiable) and the property of Gaussian densities

∇ Σ N (x | m, Σ) = 1 2 ∇ 2 x N (x | m, Σ)
to let a Hessian w.r.t. x appear. The Bures derivative is given by [see [START_REF] Bhatia | On the Bures-Wasserstein distance between positive definite matrices[END_REF]:

∇ Σ B 2 (Σ t , Σ) = I -T Σ,Σt ,
where T A,B is the optimal transport map from N (0, A) to N (0, B), with the explicit expression

T A,B = A -1 2 (A 1 2 BA 1 2 ) 1 2 A -1 2 = (T B,A ) -1
. The gradient of the variational loss L in ( 14) is thus:

∇ Σ L(m, Σ) = 1 2h I - 1 2h T Σ,Σt - 1 2 Σ -1 - 1 2 E p [∇ 2 x ln π(x)] .
Zeroing this equation gives:

I = T Σ,Σt + h Σ -1 + h E p [∇ 2 x ln π(x)] . (16)
Multiplying by Σ on the left, as well as on the right, yields the two following equations:

Σ = ΣT Σ,Σt + hI + h Σ E p [∇ 2 x ln π(x)] , (17)
Σ = T Σ,Σt Σ + hI + h E p [∇ 2 x ln π(x)] Σ . ( 18
)
Averaging the equations we obtain the symmetrized form:

Σ = 1 2 T Σ,Σt Σ + 1 2 ΣT Σ,Σt + hI + 1 2 h Σ E p [∇ 2 x ln π(x)] + 1 2 h E p [∇ 2 x ln π(x)] Σ . (19)
Let us denote T Σ,Σt = T (Σ). Since T (Σ) pushes forward Σ to Σ t , it follows that T (Σ) Σ T (Σ) = Σ t (which can be checked directly from the expression for T Σ,Σt ). The first variation of this equality w.r.t. Σ at I gives

dT Σ t + dΣ + Σ t dT = 0 . ( 20
)
Let us now term Σ = Σ t+h the solution to (19). Up to the first order in h we have Σ t+h = Σ t +dΣ = Σ t +h Σt . Let dT denote the corresponding first variation of T , that is, T (Σ t+h ) = T (Σ t )+dT = I +dT up to the first order in h. Substituting into (19), using the previously found relation (20), dividing by h and letting h 0 , we finally obtain the desired ODE:

Σt = 2I + Σ t E pt [∇ 2 x ln π(x)] + E pt [∇ 2 x ln π(x)] Σ t (21) = 2I + E pt [∇ x ln π(x) ⊗ (x -m t )] + E pt [(x -m t ) ⊗ ∇ x ln π(x)] , (22)
where the relation

E p [∇ 2 x ln π(x)] Σ = E p [∇ x ln π(x) ⊗ (x -m)
] comes from Gaussian integration by parts and yields a Hessian-free form. Letting π ∝ exp(-V ) yields the second line of (4).

Interpretation in terms of Wasserstein gradient flows. Let T t+h→t denote the optimal transport map from p t+h to p t , so that T t+h→t = m t + T Σ t+h ,Σt (x -m t+h ). Combining the equations ( 15) and ( 16), it reads

T t+h→t (x) -x h = 1 h {m t -m t+h + (T Σ t+h ,Σt -I) (x -m t+h )} = E p t+h ∇V + (E p t+h ∇ 2 V -Σ -1 t+h ) (x -m t+h ) .
In Section C, this equality will be written

T t+h→t -id h = [∇ BW KL(• π)](p t+h ) , (23)
where 1 h (T t+h→t -id) and [∇ BW KL(• π)](p t+h ) are the Bures-Wasserstein gradients of the functionals -1 2h W 2 2 (•, p t ) and KL(• π) at p t+h respectively. The equation ( 23) is a first-order optimality condition for the Bures-JKO scheme (6) and mimics the known optimality condition for the original JKO scheme, see [Santambrogio, 2015, equation (8.4)].

The quantity 1 h (T t+h→t -id) is a difference quotient which measures the infinitesimal displacement of a particle traveling along the gradient flow. As h 0, we will interpret this quantity as -v t , the negative of the tangent vector to the curve at time t (the negative sign appears because T t+h→t is the transport map backwards in time). Hence, the equation ( 23) states that as h 0, the tangent vector to the curve (p t ) t≥0 is the negative Bures-Wasserstein gradient of the KL divergence, which is the definition of a gradient flow.

From this perspective, the computation of the linearization in ( 20) is equivalent to computing the tangent vector to the Wasserstein geodesic, which is given in (32).

A.2 Extension to mixtures of Gaussians

We now consider a finite Gaussian mixture model p = 1 N N i=1 p θ (i) where θ (i) = (m (i) , Σ (i) ). We consider the following variational problem:

min θ (1) ,...,θ (N ) ∈Θ 1 2N h N i=1 W 2 2 (p θ (i) , p θ (i) t ) + KL 1 N N i=1 p θ (i) π ,
where, as before, W 2 is the Wasserstein distance between two Gaussians distribution:

W 2 2 (p θ (i) , p θ (i) t ) = m (i) -m (i) t 2 + B 2 (Σ (i) , Σ (i) t ) .
The KL divergence is now written

KL 1 N N i=1 p θ (i) π = 1 N N i=1 p θ (i) ln p - 1 N N i=1 p θ (i) ln π .
For k ∈ [N ], the derivative of this divergence with respect to m (k) gives:

∇ m (k) KL(p π) = 1 N ∇ m (k) p θ (k) ln p + p ∇ m (k) ln p - 1 N ∇ m (k) p θ (k) ln π = 1 N p θ (k) ∇ x ln p π ,
where we have used the same integration by parts as in the Section A.1, i.e., p θ (k) ∇ m (k) ln p = p θ (k) ∇ x ln p, and the Fisher score property p ∇ m (k) ln p = 0. Mimicking Section A.1, see ( 15), we obtain in the limit h 0

ṁ(k) = -E p θ (k) ∇ x ln p π ,
which is the desired equation ( 11).

The derivative of the KL divergence with respect to Σ (k) gives:

∇ Σ (k) KL(p π) = 1 N ∇ Σ (k) p θ (k) ln p + p ∇ Σ (k) ln p - 1 N ∇ Σ (k) p θ (k) ln π = 1 N 1 2 p θ (k) ∇ 2 x ln p - 1 2 p θ (k) ∇ 2 x ln π = 1 2N E p θ (k) ∇ 2 x ln p π ,
where we have used a double integration by parts

∇ Σ (k) p θ (k) ln p = 1 2 p θ (k) ∇ 2
x ln p as in Section A.1 and the Fisher score property p ∇ Σ (k) ln p = 0.

Using the Bures derivative, the critical points of the variational loss with respect to Σ k satisfy:

1 2N h (I -T Σ (k) ,Σ (k) t ) + 1 2N E p θ (k) ∇ 2 x ln p π = 0 .
Multiplying on the left and on the right by Σ (k) and taking the average as in Section A.1, we find:

1 h Σ (k) - 1 2 (Σ (k) T Σ (k) ,Σ (k) t + T Σ (k) ,Σ (k) t Σ (k) ) = - 1 2 E p θ (k) ∇ 2 x ln p π Σ (k) + Σ (k) E p θ (k) ∇ 2 x ln p π .
We can now use the first-order approximation 1

2 (Σ (k) T Σ (k) ,Σ (k) t + T Σ (k) ,Σ (k) t Σ (k) ) ≈ Σ (k) -h 2 Σ(k)
shown in Section A.1 to obtain:

Σ(k) = -E p θ (k) ∇ 2 x ln p π Σ (k) -Σ (k) E p θ (k) ∇ 2 x ln p π .
This yields the desired ODE (12), which can be rewritten in a Hessian-free form:

Σ(k) = -E p θ (k) ∇ x ln p π ⊗ (x -m k ) -E p θ (k) (x -m k ) ⊗ ∇ x ln p π .

B Background on Otto calculus

B.1 Overview and history

Historically, the connection between dissipative evolution equations and the theory of gradient flows on the Wasserstein space was discovered in Otto [1998]. Subsequently, this link was further developed and strengthened in the seminal works [START_REF] Jordan | The variational formulation of the Fokker-Planck equation[END_REF], [START_REF] Otto | The geometry of dissipative evolution equations: the porous medium equation[END_REF]. Although the paper [START_REF] Jordan | The variational formulation of the Fokker-Planck equation[END_REF] chronologically precedes [START_REF] Otto | The geometry of dissipative evolution equations: the porous medium equation[END_REF], the intuition of the former is based heavily on the work of Otto in the latter paper, in which he develops the formal2 rules governing the calculus which now bears his name.

Otto calculus endows the space P 2 (R d ) of probability measures over R d with finite second moment with a formal Riemannian structure inspired by fluid dynamics. To describe the idea, suppose that (µ t ) t≥0 is a curve of probability measures, with µ t representing the fluid density at time t. Also, let (v t ) t≥0 denote the velocity vector fields governing the dynamics of the particles; this means that the trajectory t → x t of an individual particle evolves according to the ODE ẋt = v t (x t ) .

(24)

In probabilistic language, if x 0 is a random variable drawn from the density µ 0 and it evolves according to (24), then x t ∼ µ t for all t ≥ 0. From this, we can derive a partial differential equation (PDE) governing the evolution of (µ t ) t≥0 as follows: fix a test function ϕ : R d → R (which is bounded, smooth, etc.). Formally, if the integration by parts is justified, then

ϕ ∂ t µ t = ∂ t ϕ dµ t = ∂ t E ϕ(x t ) = ∇ϕ, v t dµ t = -ϕ div(µ t v t )
from which we deduce the continuity equation of fluid dynamics:

∂ t µ t + div(µ t v t ) = 0 . ( 25
)
Conversely, if (µ t ) t≥0 is a sufficiently nice curve, then it is always possible to find a family of vector fields (v t ) t≥0 such that the equation ( 25) holds, i.e., we can interpret (µ t ) t≥0 as the evolution of a fluid density. However, the choice of vector fields is not unique, since we may always replace v t with another vector field ṽt such that div(µ t (v t -ṽt )) = 0. This motivates the search for a distinguished choice of vector fields to describe the evolution of the curve of measures.

To do so, we pick v t to minimize the kinetic energy,

v t = arg min w t 2 dµ t w t : R d → R satisfies div(µ t w t ) = -∂ t µ t .
If µ t is regular (admits a density w.r.t. Lebesgue measure), then the minimum is attained at a gradient vector field: v t = ∇ψ t for a function ψ t : R d → R. We are led to define the tangent space

T µ P 2 (R d ) = {∇ψ | ψ : R d → R}
and endow it with the inner product v, w µ = v, w dµ .

This yields a formal Riemannian structure on P 2 (R d ). Moreover, the choice of picking the vector field with minimal kinetic energy is closely related to the idea of optimal transport of mass [see [START_REF] Villani | Topics in optimal transportation[END_REF], and in fact [START_REF] Benamou | A numerical method for the optimal time-continuous mass transport problem and related problems[END_REF] showed that

W 2 2 (µ 0 , µ 1 ) = inf v t 2 µt dt (µ t , v t ) t∈[0,1] solves the continuity equation (25) . ( 26
)
From the lens of Riemannian geometry, this says that the notion of distance induced by the Riemannian structure is precisely the quadratic Wasserstein distance, and hence we refer to the space P 2 (R d ) equipped with this Riemannian structure as the Wasserstein space. This formal picture already allows one to compute gradients of functionals defined over P 2 (R d ) and hence to consider gradient flows, as well as to derive criteria which imply quantitative rates of convergence for these flows. However, it is a considerable technical undertaking to make the preceding formal considerations fully rigorous, and this was only accomplished later in the comprehensive monograph [START_REF] Ambrosio | Gradient flows in metric spaces and in the space of probability measures[END_REF]. Instead, in [START_REF] Jordan | The variational formulation of the Fokker-Planck equation[END_REF], the authors sidestep this difficulty by considering an implicit time-discretization scheme which only requires the metric structure of (P 2 (R d ), W 2 ). For a step size h > 0, define the discrete updates µ h,k+1 := arg min

µ∈P 2 (R d ) F(µ) + 1 2h W 2 2 (µ, µ h,k ) , (27)
where F : P 2 (R d ) → R ∪ {∞} is the functional of interest defined over the Wasserstein space. Note that in optimization, this is known as the "proximal point method" for minimizing F.

As h 0, one hopes that we have convergence µ h, t/h → µ t in a suitable sense, and then the limiting curve (µ t ) t≥0 can be interpreted as the Wasserstein gradient flow of F. This is indeed what [START_REF] Jordan | The variational formulation of the Fokker-Planck equation[END_REF] showed in a particular, but important case. Namely, if π ∝ exp(-V ) is a density on R d obeying mild regularity conditions, and we take the functional to be the KL divergence, F(µ) = KL(µ π), then the sequence of discrete approximations converges to the solution of the Fokker-Planck equation

∂ t µ t = div µ t ∇ ln µ t π . ( 28
)
It is well-known that the Fokker-Planck equation governs the evolution of the marginal law of the Langevin diffusion

dX t = -∇V (X t ) dt + √ 2 dB t ,
where (B t ) t≥0 is a standard Brownian motion on R d . Hence, this celebrated result says that the Langevin diffusion can be interpreted as the Wasserstein gradient flow of the KL divergence. The implicit discretization ( 27) is now commonly known as the "JKO scheme" after the authors Jordan, Kinderlehrer, and Otto.

Although the Wasserstein space is not truly a Riemannian manifold, many of the formal calculations of [START_REF] Otto | The geometry of dissipative evolution equations: the porous medium equation[END_REF] can now be justified rigorously, under appropriate technical conditions, due to the extensive theory developed in [START_REF] Ambrosio | Gradient flows in metric spaces and in the space of probability measures[END_REF], [START_REF] Villani | Optimal transport, volume 338 of Grundlehren der Mathematischen Wissenschaften[END_REF]. This perspective leads to intuitive derivations of gradient flows, as explained in Section C, and much more.

B.2 Geometry of the Wasserstein space

In this section, we provide further details about the geometry of (P 2 (R d ), W 2 ).

Let µ 0 , µ 1 ∈ P 2 (R d ), and for simplicity assume that µ 0 admits a density with respect to Lebesgue measure. Then, Brenier's theorem [Villani, 2003, Theorem 2.12] says that there exists a proper, convex, lower semicontinuous ϕ : R d → R ∪ {∞} such that ∇ϕ solves the optimal transport problem from µ 0 to µ 1 : namely, (∇ϕ) # µ 0 = µ 1 and W 2 2 (µ 0 , µ 1 ) = ∇ϕ(x) -x 2 dµ 0 (x). We refer to ∇ϕ as the optimal transport map from µ 0 to µ 1 .

The (unique) constant-speed geodesic (µ t ) t∈[0,1] joining µ 0 to µ 1 is then described via

µ t = (∇ϕ t ) # µ 0 , ∇ϕ t := (1 -t) id + t ∇ϕ . ( 29
)
In view of the fluid dynamical perspective, the constant-speed geodesics in the Wasserstein space correspond to particle trajectories t → x t which are straight lines traversed at constant speed: indeed, x t = ∇ϕ t (x 0 ) = (1 -t) x 0 + t ∇ϕ(x 0 ). Since ẋt = ∇ϕ(x 0 ) -x 0 = (∇ϕ -id) • (∇ϕ t ) -1 (x t ), then along the geodesic we see that (µ t , v t ) t∈[0,1] solves the continuity equation ( 25), where the vector field is v t = (∇ϕ -id) • (∇ϕ t ) -1 . This solution achieves the minimum in (26).

Recall that on a Riemannian manifold M, the Riemannian exponential map at p is defined on a subset of the tangent space T p M, and it maps v to the endpoint of the constant-speed geodesic at time 1 which emanates from p with velocity v (at time 0). The Riemannian logarithmic map log p is the inverse mapping: it maps an element q ∈ M to the element v ∈ T p M such that the constant-speed geodesic joining p to q in one unit of time has velocity v at time 0. In the previous paragraph, we have identified the logarithmic map: log µ ν = ∇ϕ µ→ν -id, where ∇ϕ µ→ν is the optimal transport map from µ to ν. Thus, the Riemannian exponential map is exp µ v = (id + v) # µ.

B.3 The Bures-Wasserstein space

The space of non-degenerate Gaussian distributions equipped with the W 2 metric is known as the Bures-Wasserstein space, after [START_REF] Bures | An extension of Kakutani's theorem on infinite product measures to the tensor product of semifinite w * -algebras[END_REF]. We denote this space as BW(R d ).

Given m ∈ R d and Σ 0, we denote by p m,Σ the Gaussian on R d with mean m and covariance Σ. Conversely, for a non-degenerate Gaussian p we write (m p , Σ p ) for its mean and covariance. Via this correspondence, we can therefore identify the space of non-degenerate Gaussians with the manifold

R d × S d
++ , where S d ++ denotes the cone of positive definite matrices. Abusing notation, we will do so whenever there is no danger of confusion.

Suppose that p m 0 ,Σ 0 , p m 1 ,Σ 1 ∈ BW(R d ). Then, the optimal transport map from p 0 := p m 0 ,Σ 0 to

p 1 := p m 1 ,Σ 1 is ∇ϕ(x) = m 1 + Σ -1/2 0 (Σ 1/2 0 Σ 1 Σ 1/2 0 ) 1/2 Σ -1/2 0 (x -m 0 ) .
Observe that ∇ϕ is an affine map. Since the pushforward of a Gaussian via an affine map is also Gaussian, it follows from (29) that the constant speed geodesic (p t ) t∈[0,1] joining p 0 to p 1 also lies in BW(R d ). In other words, BW(R d ) is a geodesically convex subset of P 2 (R d ).

The tangent vector to the geodesic at time 0 is always an affine map of the form x → a+S (x-m p 0 ), where a ∈ R d and S is a symmetric matrix. The tangent space is

T p BW(R d ) = {x → a + S (x -m p ) | a ∈ R d , S ∈ S d } ,
which can therefore be identified with pairs (a, S) ∈ R d × S d . With this abuse of notation, if (a, S), (a , S ) ∈ T p BW(R d ), then (a, S), (a , S ) p = a + S (x -m p ), a + S (x -m p ) dp(x) = a, a + S, Σ p S .

Specializing the notions from the previous section, we obtain

log p (q) = m q -m p , Σ -1/2 p (Σ 1/2 p Σ q Σ 1/2 p ) 1/2 Σ -1/2 p -I , exp p (a, S) = m p + a + (S + I) (• -m p ) # p = N m p + a, (S + I) Σ p (S + I) .
Here, exp p (a, S) is defined if S -I. This definition of the tangent space is consistent with the Wasserstein space, in that we have the inclusion

T p BW(R d ) → T p P 2 (R d ), but the abuse of notation T p BW(R d ) = R d × S d can sometimes cause confusion. Indeed, if (p t = p mt,Σt ) t∈[0,1] is a constant-speed geodesic in BW(R d
), and the tangent vector at time 0 is (a, S), then p t = exp p 0 t (a, S) = N m p + ta, (tS + I) Σ p (tS + I) .

In particular, Σ t = Σ 0 + t (S -I), and ṁ0 = a , (31) Σ0 = SΣ 0 + Σ 0 S .

(32)

Although we derived the equations ( 31) and (32) for geodesic curves, they also hold for any curve (p t ) t≥0 with tangent vector equal to (a, S) at time 0. Using this, we can derive an expression for the Bures-Wasserstein gradient ∇ BW f of a function f : R d × S d ++ → R. By definition, this satisfies, for any curve (m t , Σ t ) t≥0 with tangent vector (a, S) at time 0,

∇ BW f (m 0 , Σ 0 ), (a, S) p m 0 ,Σ 0 = ∂ t t=0 f (m t , Σ t ) . Write (ā, S) = ∇ BW f (m 0 , Σ 0 ). Then, we want ā, a + S, Σ 0 S = ∇ m f (m 0 , Σ 0 ), ṁ0 + ∇ Σ f (m 0 , Σ 0 ), Σ0 = ∇ m f (m 0 , Σ 0 ), a + 2 ∇ Σ f (m 0 , Σ 0 ), Σ 0 S ,
where ∇ m , ∇ Σ denote the usual Euclidean gradients. Hence, by identification, we conclude that the Bures-Wasserstein gradient of f is related to the Euclidean gradient of f via

∇ BW f (m, Σ) = ∇ m f (m, Σ), 2 ∇ Σ f (m, Σ) . ( 33
)
See [Altschuler et al., 2022, Appendix A] for further discussion.

B.4 Evolution of the mean and covariance along the Fokker-Planck equation

It is known that the Wasserstein gradient of

F := KL(• π) is ∇ W 2 F(µ) = ∇ ln µ π . ( 34
)
[See, e.g., Ambrosio et al., 2008, Theorem 10.4.13.] Also, as shown by [START_REF] Jordan | The variational formulation of the Fokker-Planck equation[END_REF], the Langevin diffusion is the gradient flow of KL(• π). In Otto calculus, this means that the law (π t ) t≥0 of the Langevin diffusion obeys the continuity equation ( 25) with velocity vector field v t = -∇ W 2 F(π t ) = -∇ ln(π t /π), which is consistent with the Fokker-Planck equation ( 28).

According to the particle interpretation (24) of dynamics in the Wasserstein space, if x 0 ∼ π 0 and

ẋt = v t (x t ) = -∇ ln π t π (x t ) ,
then x t ∼ π t . Note that (x t ) t≥0 is not the Langevin diffusion (1) as it is the solution to a deterministic ODE (albeit with random initial condition), but the marginal law of (x t ) t≥0 agrees with that of the Langevin diffusion. This provides a convenient tool for calculating the evolution of the mean and covariance along the Fokker-Planck equation, as we now demonstrate. The evolution of the mean is

ṁt = ∂ t E x t = E ẋt = -E ∇ ln π t π (x t ) .
Since E ∇ ln π t (x t ) = 0 (which is verified via integration by parts), and π ∝ e -V , this can also be written as

ṁt = -E πt ∇V .
Next, for the evolution of the covariance,

∂ t E(x t ⊗ x t ) = E(x t ⊗ ẋt + ẋt ⊗ x t ) = -E x t ⊗ ∇ ln π t π (x t ) + ∇ ln π t π (x t ) ⊗ x t ∂ t E(x t ) ⊗ E(x t ) = m t ⊗ E( ẋt ) + E( ẋt ) ⊗ m t = -E m t ⊗ ∇ ln π t π (x t ) + ∇ ln π t π (x t ) ⊗ m t which yields Σt = -E (x t -m t ) ⊗ ∇ ln π t π (x t ) + ∇ ln π t π (x t ) ⊗ (x t -m t ) .
Integration by parts yields

(• -m t ) ⊗ ∇ ln π t dπ t + ∇ ln π t ⊗ (• -m t ) dπ t = (• -m t ) ⊗ ∇π t + ∇π t ⊗ (• -m t ) = -2I .
Hence,

Σt = 2I -E πt [∇V ⊗ (• -m t ) + (• -m t ) ⊗ ∇V ] .
This verifies equation ( 3). The equations in this section can also be derived using Itô calculus.

C Proofs via Otto calculus

Our aim in this section is to derive the Wasserstein gradient flow of the KL divergence KL(• π) constrained to lie in the Bures-Wasserstein space of non-degenerate Gaussian measures. Since the Bures-Wasserstein space can be formally viewed as a submanifold of the Wasserstein space, it leads to two natural approaches for computing the constrained gradient flow. In the first approach, we take the Wasserstein gradient of KL(• π) and we compute the orthogonal projection onto the tangent space of the Bures-Wasserstein space. In the second approach, we note that the geometry of the Bures-Wasserstein space has been studied in its own right [see, e.g., [START_REF] Bhatia | On the Bures-Wasserstein distance between positive definite matrices[END_REF] and in particular, the explicit expression (33) for the Bures-Wasserstein gradient is known. We can therefore view KL(• π) as a functional over BW(R d ) and compute its gradient directly using (33).

C.1 Orthogonal projection approach

First, we justify why computing the orthogonal projection of the P 2 (R d ) gradient gives the same result as computing the intrinsic gradient on BW(R d ). Let F be any functional on P 2 (R d ). By definition, the Bures-Wasserstein gradient ∇ BW F satisfies

∂ t F(p t ) = ∇ BW F(p t ), v t pt (35)
for any curve (p t ) t∈R in BW(R d ) with tangent vectors (v t ) t∈R . Here, ∇ BW F(p t ) ∈ T pt BW(R d ). On the other hand, since (p t ) t≥0 is also a curve in P 2 (R d ) and the Riemannian structure of BW(R d ) is consistent with that of P 2 (R d ), the definition of the gradient in P 2 (R d ) yields

∂ t F(p t ) = ∇ W 2 F(p t ), v t pt .
Note that the orthogonal projection

proj Tp t BW(R d ) ∇ W 2 F(p t ) = arg min w∈Tp t BW(R d ) w -∇ W 2 F(p t ) 2 pt is characterized as the unique element of T pt BW(R d ) satisfying proj Tp t BW(R d ) ∇ W 2 F(p t ), v pt = ∇ W 2 F(p t ), v pt
for all v ∈ T pt BW(R d ). Thus, (35) holds with

∇ BW F(p) = proj TpBW(R d ) ∇ W 2 F(p) .
This argument clearly works for arbitrary Riemannian submanifolds.

Next, we compute the projection of the P 2 (R d ) gradient of the KL divergence.

Using the formula (34) for the P 2 (R d ) gradient of the KL divergence and the description of the tangent space to BW(R d ) in Section B.3 and ( 30), the projected gradient

(ā, S) ∈ R d × S d is such that for all (a, S) ∈ R d × S d , ∇ ln p π (x), + S (x -m p ) dp(x) = (ā, S), (a, S) p = ā, a + S, Σ p S .
Using ∇p(x) = -Σ -1 p (x -m p ) p(x) and integration by parts,

∇ ln p π (x), a + S (x -m p ) dp(x) = E p ∇ ln p π , a + Σ p S ∇ ln p π (x), Σ -1 p (x -m p ) dp(x) = E p ∇ ln p π , a - Σ p S ∇ ln p π (x), ∇p(x) dx = E p ∇ ln p π , a + div Σ p S ∇ ln p π (x) dp(x) = E p ∇ ln p π , a + E p ∇ 2 ln p π , Σ p S .
Hence,

(ā, S) = E p ∇ ln p π , E p ∇ 2 ln p π . ( 36
)
Using the fact that E p ∇ ln p = 0, this can also be written

(ā, S) = (E p ∇V, E p ∇ 2 V -Σ -1 p )
which corresponds to the affine map

x → E p ∇V + (E p ∇ 2 V -Σ -1 p ) (x -m p ) . ( 37
)
If (p t = p mt,Σt ) t≥0 evolves according to the constrained gradient flow, then using the expression for the projected Wasserstein gradient together with ( 31) and ( 32),

ṁt = -E pt ∇V , Σt = 2I -Σ t E pt ∇ 2 V -E pt ∇ 2 V Σ t .
The sign in the above equations comes from the fact that we perform steepest descent in Bures-Wasserstein descent, i.e., the tangent vector to the curve at time t is -proj Tp

t BW(R d ) ∇ W 2 F(p t ).
The system of equations we have derived here differs from the system (4), but we can check that they agree using integration by parts. Indeed,

Σt = 2I -Σ t ∇ 2 V dp t -∇ 2 V dp t = 2I + Σ t ∇p t ⊗ ∇V + ∇V ⊗ ∇p t Σ t = 2I + Σ t ∇ ln p t ⊗ ∇V dp t + ∇V ⊗ ∇ ln p t dp t Σ t = 2I -E pt [(• -m t ) ⊗ ∇V + ∇V ⊗ (• -m t )] .

C.2 Alternate proof using direct Bures-Wasserstein calculation

In the second approach, we view F as a functional on the Bures-Wasserstein space. Explicitly,

F(m, Σ) = p m,Σ ln p m,Σ π .
Using (33),

∇ BW F(m, Σ) = ∇ m F(m, Σ), 2 ∇ Σ F(m, Σ) = ∇ m p m,Σ ln p m,Σ π , 2 ∇ Σ p m,Σ ln p m,Σ π . ( 38
)
Furthermore, using the identities

∇ m p m,Σ (x) = -∇ x p m,Σ (x) and ∇ Σ p m,Σ (x) = 1 2 ∇ 2 x p m,Σ (x) (39)
for the Gaussian distribution, integration by parts verifies that (38) agrees with (36).

D Proof of Corollary 1

Corollary 1 is a consequence of general and well-known principles for gradient flows. To emphasize this generality, we will consider an abstract α-convex differentiable functional F defined over a geodesically convex subset of a Riemannian manifold; this ensures that the logarithmic map is well-defined in the following calculations. We assume that F is minimized at p ; by adding a constant to F, we can assume inf F = 0. Let d denote the distance function on the manifold. If (p t ) t≥0 , (q t ) t≥0 are two solutions to the gradient flow for F, then

∂ t d 2 (p t , q t ) = 2 log pt (q t ), ∇F(p t ) pt + 2 log qt (p t ), ∇F(q t ) qt .
(The reader who is unfamiliar with Riemannian geometry should keep in mind that in Euclidean space, log p (q) = q -p.) Next, the α-convexity of F implies

F(p t ) ≥ F(q t ) + ∇F(q t ), log qt (p t ) qt + α 2 d 2 (p t , q t ) , F(q t ) ≥ F(p t ) + ∇F(p t ), log pt (q t ) pt + α 2 d 2 (p t , q t ) .
Adding these equations and rearranging yields

∂ t d 2 (p t , q t ) ≤ -2α d 2 (p t , q t ) .
By Grönwall's inequality, it implies

d 2 (p t , q t ) ≤ exp(-2αt) d 2 (p 0 , q 0 ) .
This inequality has two consequences. First, for any α ∈ R, p 0 = q 0 implies p t = q t : the solution to the gradient flow is unique. Second, if α > 0, then we can set q t = p for all t ≥ 0 to deduce exponential contraction of the gradient flow to the minimizer p , which is the first statement of Corollary 1.

To obtain convergence in functional values, observe that by definition of the gradient flow, we have on the one hand that

∂ t F(p t ) = -∇F(p t ) 2 pt . (40)
On the other hand, if α > 0, the convexity inequality and Young's inequality respectively, yield

0 = F(p ) ≥ F(p) + ∇F(p), log p (p ) p + α 2 d 2 (p, p ) (41)
≥ F(p) - 1 2α ∇F(p) 2 p - α 2 log p (p ) 2 p =d 2 (p,p ) + α 2 d 2 (p, p )
and hence ∇F(p) 2 ≥ 2α F(p). Substituting this into (40) and applying Grönwall's inequality again, we deduce

F(p t ) ≤ exp(-2αt) F(p 0 ) .
Finally, suppose α = 0. We consider the Lyapunov functional

L t := t F(p t ) + 1 2 d 2 (p t , p ) .
Differentiating in time,

∂ t L t = F(p t ) -t ∇F(p t ) 2 pt + log pt (p ), ∇F(p t ) pt .
On the other hand, applying the convexity inequality in (41) with α = 0 yields ∂ t L t ≤ 0. Hence, L t ≤ L 0 , and

F(p t ) ≤ d 2 (p 0 , p ) 2t .

E Proof of Theorem 2

In this section, we use the Riemannian exponential and logarithmic maps, as discussed in Section B.3. Also, let F := KL(• π) denote the KL divergence. For τ > 0, the eigenvalue clipping operation is defined as

clip τ : Σ = d i=1 λ i u i u T i → clip τ Σ := d i=1 (λ i ∧ τ ) u i u T i . (42)
In the proof of Theorem 1 in Section C, we showed that the Bures-Wasserstein gradient is

g p := ∇ BW F(p) = E p ∇V, E p ∇ 2 V -Σ -1 ( 43
)
where Σ is the covariance matrix of p. Here, the first component of the gradient governs the evolution of the mean, whereas the second component governs the evolution of the covariance; see Section B.3. We propose to estimate the gradient in (43) via a sample,

ĝp := ∇V ( X), ∇ 2 V ( X) -Σ -1 , X ∼ p .
By comparing Algorithm 1 and the definition of the exponential map in Section B.3, one can check that for p

+ k := p m k+1 ,Σ + k and 3 h ≤ 1 p + k = exp p k (-hĝ k ) , where ĝk ∈ T p k BW(R d ) is the stochastic gradient ĝk (x) = ∇V ( Xk ) + (∇ 2 V ( Xk ) -Σ -1 k ) (x -m k ) .
Thus, aside from the eigenvalue clipping operation (which is harmless, due to Lemma 3 below), Algorithm 1 is exactly a stochastic gradient descent scheme on BW(R d ). Note also that from the definition of the exponential map in Section B.2, the update can also be written at the particle level:

if X k ∼ p k is independent of ĝk , then X + k := X k -h ĝk (X k ) ∼ p + k . (44)
In the next lemma, we obtain a uniform control on the smallest eigenvalues of the covariance matrices of the iterates.

Lemma 2. Assume that 0 ≺ αI ∇ 2 V I holds and h ≤ α 2 /60. Also, in Algorithm 1, assume that Σ k α 9 I. Then, Σ + k α 9 I.

Proof. Since the statement of the lemma only involves the covariance matrices, we can suppose that all of the mean vectors are zero.

The key is to write Σ + k as a generalized Bures-Wasserstein barycenter at Σ k for an appropriate distribution. Recall that

Σ + k = I + h Σ -1 k -h ∇ 2 V ( Xk ) Σ k I + h Σ -1 k -h ∇ 2 V ( Xk ) . (45)
Note that Σ -1 k is the optimal transport map from the Gaussian p 0,Σ k to p 0,Σ -1

k . 4 Hence, h Σ -1 k -h ∇ 2 V ( Xk ) = h (Σ -1 k -I) + h (I -∇ 2 V ( Xk )) = h log Σ k (Σ -1 k ) + h log Σ k ( Σ)
where we defined the matrix

Σ = (2I -∇ 2 V ( Xk )) Σ k (2I -∇ 2 V ( Xk )).
To check that this is valid, we need 2I -∇ 2 V ( Xk ) 0, i.e., ∇ 2 V ( Xk ) 2I, which follows from ∇ 2 V I. We have shown that

Σ + k = exp Σ k log Σ k (Σ) dP (Σ)
where

P = (1 -2h) δ Σ k + h δ Σ -1 k + h δ Σ = (1 -2h) δ Σ k + 2h 1 2 δ Σ -1 k + 1 2 δ Σ .
This is precisely the definition of a generalized Bures-Wasserstein barycenter. Next, suppose that Σ k λI for some λ > 0. Since Σ k α -1 I, and

I 2I -∇ 2 V ( Xk ) 2I, α I Σ -1 k 1 λ I , and λ I Σ 4 α I .
Then, [START_REF] Altschuler | Averaging on the Bures-Wasserstein manifold: dimension-free convergence of gradient descent[END_REF], Theorem 1]5 implies the following. If we define the quantities

λ -:= 1 2 √ α + 1 2 √ λ 2 , λ + := 1 2 1 λ + 1 2 4 α , then for step sizes 2h ≤ λ - 2λ + and if Σ k λ - 4 I, we also have Σ + k λ - 4 I.
To use this result, let us choose λ such that λ - 4 = λ; it can be seen that this holds with λ = α 9 . Since λ + = 13 2α , the step size condition then translates into h ≤ 2α 2 117 , for which it suffices to have h ≤ α 2 60 . We also recall an important fact about the eigenvalue clipping operation.

Lemma 3 ( [Altschuler et al., 2022, Proposition 3]). For any m ∈ R d , τ > 0, and

Σ, Σ ∈ S d ++ , W 2 (p m,clip τ Σ , p m,clip τ Σ ) ≤ W 2 (p m,Σ , p m,Σ ) .
We now turn towards the proof of Theorem 2. In the proof, we let F k := σ( X0 , X1 , X2 , . . . , Xk-1 ) be the σ-algebra generated by the random samples up until iteration k.

Proof of Theorem 2. Conditioned on F k , and independently of Xk , let X k ∼ p k and Z ∼ π be optimally coupled; let Ē denote the expectation taken w.r.t. (X k , Z). Using Lemma 3, the fact that Σ 1 α I (see discussion in Section 3.3), and (44), we have

E[W 2 2 (p k+1 , π) | F k ] ≤ E[W 2 2 (p + k , π) | F k ] ≤ E Ē[ X k -h ĝk (X k ) -Z 2 ] F k = E Ē[ X k -Z 2 -2h ĝk (X k ), X k -Z + h 2 ĝk (X k ) 2 ] F k = W 2 2 (p k , π) -2h Ē g k (X k ), X k -Z + h 2 E Ē[ ĝk (X k ) 2 ] F k ,
where we abbreviated

g k := g p k . From strong convexity of KL(• π) on BW(R d ) (Lemma 1), Ē g k (X k ), X k -Z ≥ KL(p k π) -KL(π π) + α 2 W 2 2 (p k , π) ≥ α W 2 2 (p k , π) .
Thus,

E[W 2 2 (p k+1 , π) | F k ] ≤ (1 -2αh) W 2 2 (p k , π) + h 2 E Ē[ ĝk (X k ) 2 ] F k =:err .
It remains to bound the error term.

Recall that

ĝk (X k ) = (∇ 2 V ( Xk ) -Σ -1 k ) (X k -m k ) + ∇V ( Xk ) .
We bound the terms one by one. First,

Ē[ Σ -1 k (X k -m k ) 2 ] = tr(Σ -1 k ) ≤ 9d α
where we used Lemma 2. Next, since ∇ 2 V I by assumption,

Ē[ ∇ 2 V ( Xk ) (X k -m k ) 2 ] ≤ Ē[ X k -m k 2 ] = tr(Σ k ) ≤ d α .
Lastly, let Ẑ ∼ π be optimally coupled with Xk . By the optimality condition for π (Section 3.3), we know that E ∇V ( Ẑ) = 0. Applying the Poincaré inequality for π (which holds because π is strongly log-concave, see [Bakry et al., 2014, Theorem 4.8.4])

Ē[ ∇V ( Xk ) 2 ] ≤ 2 Ē[ ∇V ( Ẑ) 2 ] + 2 Ē[ Xk -Ẑ 2 ] ≤ 2 α E π[ ∇ 2 V 2 HS ] + 2 W 2 2 (p k , π) ≤ 2d α + 2 W 2 2 (p k , π) .
Collecting the terms,

err ≤ 36d α + 6 W 2 2 (p k , π) .
From the assumption h ≤ α 2 60 .

E[W 2 2 (p k+1 , π) | F k ] ≤ (1 -αh) W 2 2 (p k , π) + 36dh 2 α .
Iterating this bound proves the result.

F Proof of Theorem 3

In order to present the proof of Theorem 3, we first review relevant facts about the Wasserstein space over a Riemannian manifold (M, g). We refer readers to [START_REF] Villani | Optimal transport, volume 338 of Grundlehren der Mathematischen Wissenschaften[END_REF] for an in-depth treatment.

Similarly to the Euclidean setting, we can define the space of probability measures over M with finite second moment,

P 2 (M) := µ ∈ P(M) d 2 (p 0 , •) dµ < ∞ for some p 0 ∈ M ,
where d denotes the induced distance on M. We equip P 2 (M) with the 2-Wasserstein metric

W 2 2 (µ, ν) := inf γ∈C(µ,ν) d 2 (x, y) dγ(x, y) 1/2
, which makes (P 2 (M), W 2 ) into a metric space. Moreover, at each regular measure µ ∈ P 2 (M), we can define the tangent space

T µ P 2 (M) := {∇ψ | ψ ∈ C ∞ c (M)} L 2 (µ)
equipped with the inner product v, w µ := g p v(p), w(p) dµ(p) , which endows (P 2 (M), W 2 ) with the structure of a formal Riemannian manifold. Curves (µ t ) t≥0 in P 2 (M) are still described the continuity equation

∂ t µ t = div(µ t v t ) (46)
where now v t is an element of the tangent bundle T M and div denotes the divergence operator on the Riemannian manifold. Equation ( 46) is to be interpreted in the weak sense, i.e., for any test function ϕ : M → R,

∂ t ϕ dµ t = g(∇ϕ, v t ) dµ t . (47)
If (µ t ) t≥0 is a smooth curve such that µ t admits a density ρ t w.r.t. the Riemannian volume measure, then this is equivalent to the partial differential equation (PDE)

∂ t ρ t = div(ρ t v t ) .
As before, the continuity equation admits a particle interpretation: if p 0 ∼ µ 0 and (p t ) t≥0 evolves via the ODE

ṗt = v t (p t ) , (48)
then p t ∼ µ t for all t ≥ 0. Given a functional F : P 2 (M) → R ∪ {∞} defined over the Wasserstein space, its gradient at µ is, by definition, the element ∇ W 2 F(µ) ∈ T µ P 2 (M) such that: for all curves (µ t ) t∈R satisfying the continuity equation ( 46) with µ 0 = µ, it holds that

∂ t t=0 F(µ t ) = ∇ W 2 F(µ), v 0 µ = g ∇ W 2 F(µ), v 0 dµ .
Using the continuity equation ( 47), it follows by direct identification that

∇ W 2 F(µ) = ∇δF(µ) ,
where δF(µ) : M → R, the first variation of F at µ, is defined up to an additive constant and satisfies

∂ t t=0 F(µ) = δF(µ) ∂ t t=0 µ t .
A gradient flow of F is a curve (µ t ) t≥0 which satisfies the continuity equation ( 46) with velocity vector field v t = -∇ W 2 F(µ t ), which in turn admits the particle interpretation (48).

We now consider the functional

F(µ) := KL(p µ π)
and compute its first variation. Let m denote the Riemannian volume measure; let (ρ t ) t∈R be a smooth curve of densities ρ t = dµt dm . Since

F(µ) = V dp µ + p µ ln p µ = V dp θ ρ(θ) dm(θ) + ln p θ ρ(θ ) dm(θ ) dp θ ρ(θ) dm(θ) then ∂ t F(µ t ) = V dp θ ρt (θ) dm(θ) + p θ ρt (θ ) dm(θ ) p θ ρ t (θ ) dm(θ ) dp θ ρ t (θ) dm(θ) + ln p θ ρ(θ ) dm(θ ) dp θ ρt (θ) dm(θ) = (V + ln p µt + 1) dp θ ρt (θ) dm(θ) .
From this, δF(µ) : θ → (V + ln p µ + 1) dp θ = ln p µ π dp θ + 1 .

Next, we compute the Bures-Wasserstein gradient using ( 33) and ( 39):

∇ BW δF(µ)(m, Σ) = ln p µ π ∇ m p m,Σ , 2 ln p µ π ∇ Σ p m,Σ = ∇ ln p µ π dp m,Σ , ∇ 2 ln p µ π dp m,Σ .
Finally, to derive the system of ODEs (10), we combine the above expression for the Wasserstein gradient of F together with the particle interpretation (48) and the equations ( 31) and ( 32) for dynamics on the Bures-Wasserstein space.

Hence,

H(p µt ) = p µt ln p µt = - 1 2 ln(2πe) - 1 2 ln 1 + (1 -t + tτ ) 2 .
Then, the convexity of t → H(p µt ) is equivalent to the convexity of t → -ln(1 + (1 -t + tτ ) 2 ), which fails when, e.g., τ = 1/2; in that case, the function is, in fact, concave on the interval [0, 1].

H The Wasserstein-Fisher-Rao gradient flow

Similarly to the setting in Section 5, here we identify probability measures µ over the Bures-Wasserstein space with the corresponding Gaussian mixture p µ . The aim of this section is to derive the gradient flow of the KL divergence µ → KL(p µ π), except we now equip the space P 2 (BW(R d ))

with the Wasserstein-Fisher-Rao geometry [START_REF] Liero | Optimal transport in competition with reaction: the Hellinger-Kantorovich distance and geodesic curves[END_REF][START_REF] Chizat | An interpolating distance between optimal transport and Fisher-Rao metrics[END_REF][START_REF] Liero | Optimal entropy-transport problems and a new Hellinger-Kantorovich distance between positive measures[END_REF].

Deriving the gradient flow with respect to this geometry leads to dynamics for a system of interacting Gaussian particles in which the weight of each particle is also updated at each iteration.

H.1 Background on Wasserstein-Fisher-Rao geometry

Here we briefly summarize the relevant background on the Wasserstein-Fisher-Rao (WFR) geometry.

The WFR metric is also called the Hellinger-Kantorovich metric by some authors.

The Fisher-Rao metric. The Fisher-Rao metric is a metric on the space M + (R d ) of positive measures (not necessarily probability measures). It is the induced metric on M + (R d ) if we enforce that the mapping µ → √ µ (defined for smooth probability densities µ) is an isometry into L 2 (R d ).

This means that

d 2 FR (µ 0 , µ 1 ) = ( √ µ 0 - √ µ 1 ) 2 ,
and if µ 0 and µ 1 are probability measures then this is known to statisticians (up to a constant factor) as the squared Hellinger distance. (If we apply the analogous procedure to discrete probability measures, then this amounts to identifying the simplex with a subset of the unit sphere.) The Fisher-Rao metric is well-studied in the field of information geometry [START_REF] Amari | Methods of information geometry[END_REF]Nagaoka, 2000, Ay et al., 2017]. Next, we describe the Riemannian geometry underlying the Fisher-Rao metric. Consider a curve t → µ t of positive measures with time derivative μ. Since the Fisher-Rao metric endows the square root of the density with a Hilbert metric, we place endow the time derivative of the square root, √µ = μ/(2 √ µ), with the Hilbert norm μ/(2 √ µ) L 2 (R d ) . Thus, the norm at the tangent space

T µ M + (R d ) is given by μ 2 µ = μ2 4µ .
Actually, because we are working with positive measures (called unbalanced measures to distinguish from the usual optimal transport problem which requires the measures to have the same total mass), this kind of geometry is useful for studying problems in which the total mass changes over time. For example, PDEs of the form ∂ t µ t = α t µ t are called reaction equations because they describe, e.g., how the concentration of a chemical changes over time in reaction to the environment. Motivated by this application, we parameterize μ via μ = αµ, in which case the norm is

α 2 µ = 1 4 α 2 dµ . (49)
R d × R + and identify all of the points with zero mass which sit at the "tip of the cone"). The cone space is the natural setting for WFR geometry; for example, one can introduce a metric on C(R d ) and show that the WFR distance is an optimal transport problem w.r.t. this metric. We will not go into such detail, but nevertheless we introduce the cone space because is important for the particle interpretation of WFR dynamics. Curves of measures (µ t ) t∈[0,1] in the WFR geometry admit a particle interpretation in terms of trajectories on C(R d ). Namely, the equation ( 50) can be interpreted as follows. There exists a curve of measures t → µ t over the cone space C(R d ), such that if r : C(R d ) → R + denotes the mapping (x, r) → r, and x : C(R d ) → R d maps (x, r) → x, then

µ t = x # (r 2 µ t ) .
Moreover, if we draw (x 0 , r 0 ) ∼ µ 0 and follow the ODEs

ẋt = v t (x t ) , ṙt = α t (x t ) -α t dµ t r t , then (x t , r t ) ∼ µ t .
Here the notation ∼ is an (egregious) abuse of notation because µ t is not a probability measure; by (x, r) ∼ µ more precisely we mean that µ t = (ODE t ) # µ 0 where ODE t is the solution mapping (x 0 , r 0 ) → (x t , r t ) to the above system of ODEs at time t.

To make this interpretation more concrete, we specialize to the case of discrete measures. Suppose that we start at a probability measure

µ 0 = N i=1 w (i) 0 δ x (i) 0 .
Then, we lift to the cone space:

µ 0 = N i=1 δ (x (i) 0 , w (i) 0 ) = N i=1 δ (x (i) 0 ,r (i) 0 )
where we set r

(i) t = w (i)
t . Next, we follow the ODEs

ẋ(i) t = v t (x (i) t ) , ṙ(i) t = α t (x (i) t ) - N j=1 w (j) t α t (x (j) t ) r (i) t .
Upon projecting back to the base space, we obtain another discrete measure

µ t = N i=1 w (i) t δ x (i) t = N i=1 (r (i) t ) 2 δ x (i) t .
As a sanity check, we check that these dynamics ensure that µ t is a probability measure for all t.

The time derivative of the sum of the weights is

∂ t N i=1 w (i) t = 2 N i=1 r (i) t ∂ t r (i) t = 2 N i=1 (r (i) t ) 2 α t (x (t) i ) -E µt α t = 2 N i=1 w (i) t α t (x (t) i ) -E µt α t = 0 .

H.2 Derivation of the gradient flow

Next, we derive the Wasserstein-Fisher-Rao gradient flow of the functional µ → F(µ) := KL(p µ π) on the space (P 2 (BW(R d )), WFR) of Gaussian mixtures equipped with the Wasserstein-Fisher-Rao metric (over the Bures-Wasserstein space). The WFR gradient of F, ∇ WFR F(µ), is the pair

∇ WFR F(µ) = ∇ BW δF(µ), δF(µ) -δF(µ) dµ .
This result is essentially stated as Lu et al. [2019, Proposition A.1], although we have generalized the formula to hold when the base space is no longer R d . Note also that we have already calculated the first variation of F, as well as the BW gradient, in Section F. The interpretation of the formula is that in the gradient flow of F, we have a particle (m, Σ) associated with some mass w evolving according to

ṁ = -E p m,Σ ∇ ln p µ π , Σ = -Σ E p m,Σ ∇ 2 ln p µ π -E p m,Σ ∇ 2 ln p µ π Σ , ṙ = -E p m,Σ ln p µ π -E pµ ln p µ π r ,
where r = √ w. The interpretation may be clearer in the discrete case, so suppose that we initialize the dynamics at a discrete measure

µ 0 = N i=1 w (i) 0 δ (m (i) 0 ,Σ (i) 0 ) .
Next we solve the coupled system of ODEs, for i ∈ i) ,

[N ], ṁ(i) = -E p m (i) ,Σ (i) ∇ ln p µ π , Σ(i) = -Σ (i) E p m (i) ,Σ (i) ∇ 2 ln p µ π -E p m (i) ,Σ (i) ∇ 2 ln p µ π Σ (i) , ṙ(i) = -E p m (i) ,Σ (i) ln p µ π -E pµ ln p µ π r (
where r (i) = √ w (i) and

µ t = N i=1 w (i) t δ (m (i) t ,Σ (i) t ) .
Since the normalization constant of π cancels out in the above equations, they are implementable without this knowledge.

I Experiments for Gaussian VI

The goal of the present section is to conduct numerical experiments that illustrate the convergence of the Gaussian distribution corresponding to the ODE (4) to an approximation of the target distribution. We consider two kinds of targets: a mixture of two Gaussians, and a log-concave target that corresponds to the likelihood function in logistic regression.

• Covariance matrices in square root form: To numerically enforce that the covariance matrix Σ remains symmetric and positive at each step, as is customary in the Kalman filtering literature, we consider a continuous-time "square-root" form of the covariance as developed in [START_REF] Morf | Square-root algorithms for the continuous-time linear least squares estimation problem[END_REF] and applied in [START_REF] Särkkä | On unscented Kalman filtering for state estimation of continuous-time nonlinear systems[END_REF]. Let R be a lower triangular matrix such that Σ = RR T . An ODE for R is obtained as follows.

Σ = ṘR T + R ṘT (58)
Multiplying by R -1 on the left and R -T on the right yields:

R -1 Ṙ + ṘT R -T = R -1 Σ R -T . ( 59
) As R -1 Ṙ + ṘT R -T = R -1 Ṙ + (R -1 Ṙ)
T , the solution is given by:

R -1 Ṙ = Tria(R -1 Σ R -T ) , (60)
Ṙ = R Tria(R -1 Σ R -T ) , (61)
where Tria(A) gives the lower triangular matrix L corresponding to A such that A = L + L T where L i,i = 1 2 A i,i , L i,j = A i,j if i > j, and L i,j = 0 otherwise. Letting Σ be as in ( 4), this yields an ODE in terms of the square root factor R.

• Computing expectations: We compute Gaussian expectations using a quadrature rule based on 2d sigma points x 1 , . . . , x 2d [START_REF] Julier | Unscented filtering and nonlinear estimation[END_REF]]:

E p m,Σ [f (x)] ≈ 2d n=1 α n f (x n ) ,
where the sigma points are distributed according to x n = m + c n Re n , where RR T = Σ, e n | n=1,...,d is a basis, and e n | n=d+1,...,2d is its negative. Many variants exist to choose α n and c n ; here, we consider the cubature points of [START_REF] Arasaratnam | Cubature Kalman filters[END_REF] defined by α n = 1 2d and c n = √ d which are well-adapted for Gaussian integration.

I.3 Results in dimension 2

We first conduct experiments in dimension 2 to easily visualize the true posterior (normalization is performed using a discrete grid of size 100 × 100).

I.3.1 Trajectories generated by numerical integration of the ODEs

In Figure 5, we see that Gaussian vi converges quickly to one mode of the bimodal target, and to the unique mode of the logistic target. As shown in Figure 6, the results still hold if we choose a larger step size for the Runge-Kutta scheme.

I.3.2 Comparison with the Laplace approximation

We compare Gaussian vi with the Laplace approximation on the logistic target in dimension 2 for the setting described in Section I.1 with an arbitrary Σ * and N = 10. We plot the convergence speed of our algorithm for Gaussian vi in Figure 7 for separation parameters s = 1.5 and s = 2, the latter corresponding to a sharper density. Gaussian vi converges very fast and produces a better approximation of the target in terms of KL divergence than the Laplace approximation.

Figure 5: Approximation of a bimodal target (left) and a logistic target (right). We use a Runge-Kutta scheme with step size 0.1 and a time duration of T = 30 (i.e., 300 steps). The ellipsoids represent the Gaussian computed at successive steps.

Figure 6: Same as Figure 5 with a larger Runge-Kutta step size 1 (i.e., 30 steps). In both cases, the algorithm converges to the same approximation as in Figure 5.

I.4 Results in higher dimensions

We now compare Gaussian vi with the Laplace approximation on the logistic target in dimension d = 10 and d = 100. We consider the setting described in Section I.1 where we let Σ * = 1 d I, to have consistent norms of the inputs accross dimensions.

For Gaussian vi in high dimension, we find that a step size 1 for the Runge-Kutta integration method is too large and leads to singular covariance matrices. We thus take the step size equal to 0.1. The initial Gaussian is taken to be N (0, 100I), to better cover regions of low density initially.

Results are shown in Figures 8 and9 in dimension d = 10 and 100 respectively. Gaussian vi converges very fast and always produces a better approximation of the target in terms of KL divergence than the Laplace approximation. Note that the Laplace approximation can have a very high left KL divergence when the target distribution is sharp (i.e., when the two classes are well-separated). This is because the Gaussian approximation computed with the Laplace method tends to spill out of the target distribution in region of very low densities. 8 but with dimension d = 100, N = 500, with separation factor s = 0.05 (upper row) and s = 0.3 (lower row). The unnormalized KL is computed letting Z = 1 (upper row) and Z = 10 100 (lower row). The unnormalized KL divergence for the Laplace method is not shown in the lower plot because it is too large to be visualized.

J.2.2 Initialization of the Gaussian particles

We start by illustrating the sensitivity of the algorithm to the initialization on a simple example with one Gaussian particle and a bimodal target (Figure 10). When the initial particle is close to one of the two modes and has same covariance as each mode, then it moves towards that mode and its covariance remains constant. When the particle is equidistant from the two modes, then the mean of the particle converges to the average of the two modes, and its covariance increases. Perturbing the initial condition slightly leads the particle to be attracted to one of the two modes. Figure 10: Trajectory of a Gaussian particle for different initial conditions. In the three left plots, we initialize the particle with the same covariance as each mode, and in the right plot we initialize the particle with a large covariance.

To avoid bad initialization, the idea is to generate instead more particles than the number of modes of the target. Finally, we initialize our Gaussian particles with means randomly chosen from a Euclidean ball which covers most of the mass of the target density.

J.3 Experimental results

We show qualitative fits by plotting the contour lines of the approximated density (compared to the true density), as well as quantitative evaluation of the KL divergence to the target.

The true posterior is computed using a discrete grid of size 100 × 100. The KL divergences are evaluated using Monte Carlo sampling.

J.3.1 Simple targets

We consider a mixture of 20 Gaussians to approximate the targets defined in Section I.1. We see in Figure 11 that the algorithm captures both modes of the bimodal distribution, and approximates well the logistic target also, see Figure 12.

Figure 11: Approximation of the bimodal target using 20 Gaussian particles at initialization (left) and at final step (right). We use Runge-Kutta integration with step size 0.1 and integrattion time T = 30 (i.e., 300 steps).

Figure 12: Approximation of the logistic target with 20 Gaussian particles. We use Runge-Kutta integration with step size 0.1 and integration time T = 30 (i.e., 300 steps).

J.3.2 More complex targets

We assess the sensitivity to the number of particles in Figures 13,14,[START_REF] Modin | Geometry of matrix decompositions seen through optimal transport and information geometry[END_REF] When the number of particles increases, better KL divergence is achieved and the distribution is better approximated. We also note that when the samples initially cover a low density mode as in Figure 14, they tend to overestimate the local density before they escape the mode.

Figure 2 :

 2 Figure 2: Two left plots: approximation of a bimodal target and a logistic target. Two right plots: convergence of the KL in dimension 2 and 100 for the logistic target. Our algorithm yields better approximation in KL than the Laplace approximation (see Appendix I.4 for details).

Figure 3 :

 3 Figure 3: Approximation of a Gaussian mixture target π with 40 Gaussian particles. The particles are represented by their covariance ellipsoids shown at Steps 0, 1, and 2. The right figure shows the final step with the approximated density in contour-lines. More figures are available in Appendix J.

Figure 7 :Figure 8 :

 78 Figure7: Results in dimension d = 2, N = 10 for a separation factor s = 1.5 (upper row) and s = 2 (lower row). The left column shows the true density via contour lines, the true mean (black dot) and covariance (black ellipsoid), and the results of the Laplace and Wasserstein vi approximations as blue and red ellipsoids respectively. The right column shows the evolution of the left KL divergence for Gaussian vi on a logarithmic scale. The corresponding KL divergence obtained with Laplace approximation is shown as a blue straight line.

Figure 9 :

 9 Figure9: Same as Figure8but with dimension d = 100, N = 500, with separation factor s = 0.05 (upper row) and s = 0.3 (lower row). The unnormalized KL is computed letting Z = 1 (upper row) and Z = 10 100 (lower row). The unnormalized KL divergence for the Laplace method is not shown in the lower plot because it is too large to be visualized.

Figure 13 :

 13 Figure 13: A target with 4 equally weighted modes and isotropic covariances.

Figure 14 :

 14 Figure 14: A target with 4 non-equally weighted modes and isotropic covariances.

Figure 15 :

 15 Figure 15: A target with 6 non-equally weighted modes and isotropic covariances.

  Code for the experiments is available at https://github.com/marc-h-lambert/W-VI. SC is supported by the Department of Defense (DoD) through the National Defense Science & Engineering Graduate Fellowship (NDSEG) Program. FB and ML acknowledge support from the French government under the management of the Agence Nationale de la Recherche as part of the "Investissements d'avenir" program, reference ANR-19-P3IA-0001 (PRAIRIE 3IA Institute), as well as from the European Research Council (grant SEQUOIA 724063). PR is supported by NSF grants IIS-1838071, DMS-2022448, and CCF-2106377. Martin J. Wainwright and Michael I. Jordan. Graphical models, exponential families, and variational inference. Foundations and Trends in Machine Learning, 1(1-2):1-305, 2008. Yixin Wang and David M. Blei. Frequentist consistency of variational Bayes. J. Amer. Statist. Assoc., 114(527):1147-1161, 2019.

	Andre Wibisono. Sampling as optimization in the space of measures: the Langevin dynamics as
	a composite optimization problem. In Proceedings of the 31st Conference On Learning Theory,
	volume 75, pages 2093-3027, 2018.
	Keru Wu, Scott Schmidler, and Yuansi Chen. Minimax mixing time of the Metropolis-adjusted
	Langevin algorithm for log-concave sampling. Journal of Machine Learning Research, 23(270):
	1-63, 2022.
	Lin Wu, Emtiyaz K. Mohammad, and Schmidt Mark. Stein's lemma for the reparameterization trick
	with exponential family mixtures. arXiv:1910.13398, 2019.

Guodong Zhang, Shengyang Sun, David Duvenaud, and Roger Grosse. Noisy natural gradient as variational inference. In Proceedings of the International Conference on Machine Learning, volume 80, pages 5852-5861, 2018.

A notable downside of Algorithm 1 is the requirement of a Hessian oracle for V , which results in a higher per-iteration cost than typical mcmc samplers.

Here, "formal" is not a synonym for "rigorous".

This latter requirement is needed because BW(R d ) has a finite injectivity radius.

This observation was also used in the analysis of Bures-Wasserstein gradient descent for entropically regularized barycenters in[START_REF] Altschuler | Averaging on the Bures-Wasserstein manifold: dimension-free convergence of gradient descent[END_REF].

See the latest revision.

Acknowledgements

We thank Yian Ma for helpful discussions, as well as anonymous reviewers for useful references and suggestions. ML acknowledges support from the French Defence procurement agency (DGA).

G Lack of convexity of the KL divergence for mixtures of Gaussians

In this section, we provide counterexamples for the lack of convexity of the objective functional µ → F(µ) = KL(p µ π) on the space P 2 (BW(R d )).

First, we point out that even when π is strongly log-concave, the functional F can be badly behaved. For example, if π = p 0,1 = N (0, 1) is a Gaussian of variance 1, then we can write it as a Gaussian mixture in many ways: π = N (m, a) dν 1-a (m) for any a ∈ [0, 1], where ν a = N (0, a). In particular, the set of minimizers of F is not a singleton, and includes all of the measures ν a ⊗ δ a ((m, σ 2 ) is a random pair with independent components, where m ∼ N (0, 1 -a) and σ 2 = a almost surely) for a ∈ [0, 1] (as well as all convex combinations-i.e., mixtures-thereof).

Next, we give an explicit example which demonstrates the lack of convexity of the entropy functional µ → H(p µ ) := p µ ln p µ . This can be understood as the KL divergence with zero potential (V = 0). Note that the entropy functional H is convex on P 2 (R d ) [Ambrosio et al., 2008, Section 9.4], but our claim is that its composition with the map µ → p µ is not convex on P 2 (BW(R d )).

In one dimension let µ 0 = N (0, 1) ⊗ δ 1 and µ 1 = N (0, τ 2 ) ⊗ δ 1 . In words, a random pair (m 0 , σ 2 0 ) drawn from µ 0 satisfies m 0 ∼ N (0, 1) and σ 2 0 = 1, and similarly for µ 1 . What is the optimal coupling of µ 0 and µ 1 ? Clearly σ 2 0 = σ 2 1 = 1 is the trivial coupling, and since the Bures-Wasserstein distance over the means is the same as the Euclidean distance between the means, we want the usual W 2 optimal coupling between N (0, 1) and N (0, τ 2 ); it follows that m 1 = τ m 0 . Hence, the Bures geodesic between is {(m t , σ 2 t ) = ((1 -t + tτ ) m 0 , 1)} t∈[0,1] ; equivalently the (Bures-)Wasserstein geodesic between µ 0 and µ 1 is

Next, recall that the Gaussian mixture p µt is the law of X drawn in the two-stage procedure: first we draw (m t , σ 2 t ) ∼ µ t , and given (m t , σ 2 t ) we draw X ∼ p mt,σ 2 t . Thus,

Wasserstein geometry. We recall from Section B that Wasserstein geometry is motivated by a completely different class of PDEs, namely transport equations encoded by the continuity equation

which describe the evolving law of a particle x t tracing out an integral curve of the family of vector fields: ẋt = v t (x t ). The Riemannian structure is obtained by equipping the tangent space T µ P 2 (R d ) with the norm

Wasserstein-Fisher-Rao geometry. Next we combine the two geometric structures, which can model transport-reaction equations such as

The tangent space norm is then given by the combination

(At this point some authors add a factor 1 4 in front of the α 2 , which is natural in view of (49). This is convenient for studying geometric properties of the space, but it is not necessary for our purposes.) As in the pure Fisher-Rao case, this is a metric on the space of positive measures M + (R d ).

It induces the distance

One can show that the tangent space to M + (R d ) consists of pairs (α, v) for which α = u and v = ∇u for some function u : R d → R. Thus, compared to the Wasserstein metric in which the tangent space norm is the Ḣ1 (µ) norm u Ḣ1 (u) = ∇u L 2 (µ) , the Wasserstein-Fisher-Rao metric has the interpretation of completing the tangent space norm to the full Sobolev norm H 1 (µ).

Constraining the dynamics to lie within probability measures. In order to have our dynamics stay on the space of probability measures, we follow [START_REF] Lu | Accelerating Langevin sampling with birth-death[END_REF] and consider instead the equation

which now conserves mass. The tangent space norm is modified to read

Particle interpretation. The particle interpretation of the WFR geometry is more complicated to state than for the Wasserstein geometry, but it can be done. Instead of considering a particle x, we consider a pair (x, r) consisting of a particle x ∈ R d and a number r > 0 (this number is actually interpreted as the square root of the mass of the particle). The pair (x, r) should be thought of as an element of the cone space We define a bimodal target as a mixture of two Gaussians π = 1 2 N (µ 1 , Σ 1 ) + 1 2 N (µ 2 , Σ 2 ) where Σ 1 and Σ 2 have non isotropic covariances with a ratio of 3 between the largest and the smallest eigenvalues.

Log-concave target: Bayesian logistic regression The proposed log-concave target is generated in the context of the Bayesian treatment of logistic regression associated with a two-class synthetic dataset D = {(x i , y i ) : i = 1, . . . , N }. The probability of the binary label y i ∈ {0, 1} given the corresponding covariate x i and parameter z ∈ R d is defined by the following Bernoulli distribution:

where σ(x) = 1/(1 + exp(-x)) is the logistic function. We define the target distribution as the posterior associated to data D starting from an uninformative (flat) prior on z, that is,

with Z the normalization constant. The Langevin dynamics are associated with the gradient of V then defined by:

To generate the synthetic data D, we randomly draw labels y i ∈ {0, 1} and for the problem to be well-specified we have drawn the class-conditional covariates x i from Gaussian distributions N (m * y i , Σ * ) with m * 1 = -m * 0 = m * . We call s the separation factor defined by m * 1 -m * 0 = 2m * = s. For illustrative purposes we also plot Fisher's linear discriminant vector defined by z * = 2Σ * -1 m * [see Bishop, 2006, chapter 4]. An example of the generated data is displayed in Figure 4.

I.1.2 Evaluation of the KL divergence for the proposed log-concave targets

The target distribution (52) may be written π(z|D) = 1 Z π(z|D) where π(z|D) = N i=1 π(y i |x i , z) is the unnormalized distribution. The divergence between any Gaussian distribution p = N (m, Σ) and the target then writes

where H(p) is the negative entropy of a Gaussian distribution for which a closed-form expression is known. We will see shortly we can approximate the expectation under the Gaussian p as follows

using K = 2d sigma points with cubature rules

for all i, and where R is defined via the Cholesky decomposition RR T = Σ (see Section I.2 for details).

I.1.3 The Laplace approximation as a baseline

We use the widespread Laplace approximation [see Bishop, 2006, chapter 4] as a baseline for comparisons. In dimension 2, we compute the normalization constant Z of (52) using a grid. When we turn to high dimension, normalization becomes intractable. However, we may still compare our algorithm with Laplace approximation as follows. Since our goal is mainly to illustrate the convergence of our algorithm using Laplace approximation as a baseline, we may choose an arbitrary value for the normalization constant Z when evaluating the divergence to the target in equation ( 55). This allows for comparison of the KL divergence between the approximating distribution-given by either Gaussian vi or Laplace approximation-and the target π up to the same additive constant for both methods. By default, we let Z = 1, but we sometimes use larger values of Z in order to avoid plotting negative values for the unnormalized KL (albeit an arbitrary choice).

To obtain the Laplace approximation, we first compute a mode of the target distribution π. Once the mode z 0 has been found, we consider the following Taylor approximation around the mode:

where H is the Hessian of the negative log-likelihood around z 0 defined by H = ∇ 2 log 1 π (z 0 ). Renormalizing, this yields the approximation π ≈ πLaplace = N (z 0 , H -1 ) .

(57)

In our experiments, we use the L-BFGS algorithm [START_REF] Liu | On the limited memory BFGS method for large scale optimization[END_REF]] to find the mode z 0 .

I.2 Implementation

We follow [START_REF] Särkkä | On unscented Kalman filtering for state estimation of continuous-time nonlinear systems[END_REF], Lambert et al. [2022b] to compute the expectations involved in equation ( 4) using quadrature rules. We then numerically integrate the set of coupled ODEs in equation ( 4) using a fourth-order Runge-Kutta method. As a first step, we introduce a method to enforce that the covariance matrix Σ remains symmetric and positive at all times.

J Experiments for mixture of Gaussians VI

In this section, we consider a mixture of Gaussians model to approximate a target distribution in the simple two-dimensional case. The goal is to illustrate the convergence of the approximating particles system ( 11)-( 12) to an approximation of the target in the form of a finite mixture of Gaussians.

J.1 Setup

We consider the bimodal and logistic targets defined in Section I, as well as more complex targets defined as finite mixtures of Gaussians:

The gradient ∇ x log π(x) then writes:

We consider K Gaussian samples equally weighted such that our mixture model is p

Even if we are using an approximation with equal weights, contrary to the target (which can be arbitrary in practice), we can hope from Theorem 3 convergence to a good approximation of π when letting K M .

J.2 Implementation details

J.2.1 Integration of the ODEs

Following equations ( 11)-( 12), we implement the system of ODEs

We recall that these equations arise from applying Theorem 3 to a discrete mixing measure and applying integration by parts to obtain Hessian-free updates. To constrain the covariance matrix to remain definite positive along the numerical integration process, we use the same method as in the Gaussian vi case (Section I.2): we replace each ODE for a covariance matrix Σ by an ODE for its lower triangular matrix factor R where Σ = RR T . To compute the expectations, we use the sigma points with cubature rules as described in Section I.2. Finally, to solve the ODEs we consider a classical Runge-Kutta scheme of 4 th order. The coupling between the ODEs is taken into account by applying the Runge-Kutta algorithm on the joint ODE Ẋ = F (X) where the Gaussian parameters are stacked as follows: X = [m 1 , . . . , m K , vec(R 1 ), . . . , vec(R K )].

For our problem, setting the Runge-Kutta step size to 0.1 is sufficient. We observe that asymptotic convergence, i.e., complete stability of the ODE system, may require many iterations when we propagate a large number of coupled Gaussian particles. On the other hand, the KL divergence is roughly stable after 30 steps.