Variational inference via Wasserstein gradient flows
Authors: Marc Lambert
Sinho Chewi
Francis Bach
Silvère Bonnabel
Philippe Rigollet
Venue: NeurIPS 2022 Accept (2022)
Keywords: Bures-Wasserstein
Kalman filter
mixture of Gaussians
variational inference
Wasserstein gradient flow

Abstract: Along with Markov chain Monte Carlo (MCMC) methods, variational inference (VI) has emerged as a central computational approach to large-scale Bayesian inference. Rather than sampling from the true posterior $\pi$, VI aims at producing a simple but effective approximation $\hat \pi$ to $\pi$ for which summary statistics are easy to compute. However, unlike the well-studied MCMC methodology, algorithmic guarantees for VI are still relatively less well-understood. In this work, we propose principled methods for VI, in which $\hat \pi$ is taken to be a Gaussian or a mixture of Gaussians, which rest upon the theory of gradient flows on the Bures--Wasserstein space of Gaussian measures. Akin to MCMC, it comes with strong theoretical guarantees when $\pi$ is log-concave.

Variational inference via Wasserstein gradient flows
Mare Lambert DGA, INRIA, Ecole Normale Supérieure, PSI. Research University marc. lambert@inria.fr
Sinho Chewi MIT schewi edu
Francis Bach INRIA, Ecole Normale Supérieure, PSL Research University francis www.baaria.haa....... fr
Silvère Bonnabel MINES Paris PSL,
Philippe Rigollet
MIT rigollet@math.. mit. edu
Université de la Nouéelle-Calédonie silvere. bonnabel maminesparis.psl...
Abstract
Along with Markov chain Monte Carlo (MCMC) methods, variational inference (v1) has emerged as central computational approach to large-scale Bayesian inference. Rather than sampling from the true posterior æ, VI aims producing a simple but effective approximation स to for which summary statistics are easy to compute. However, unlike the well-studied MCMC methodology, algorithmic guarantees for VI are still relatively less well-understood. In this work, we propose principled methods for VI, which is taken to be a Gaussian or a mixture of Gaussians, which rest upon the theory of gradient flows on the Gaussian measures. Akin to MCMC, it comes with strong theoretical guarantees when is log-concave.
space of
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 method- ology typically generates 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 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, - (X V2dB,
where (B.) denotes standard Brownian motion on Rd, admits as a stationary distribution. Crucially, the Langevin diffusion can be discretized and implemented without knowledge of the normalizing constant of π, leading practical algorithms for Bayesian inference. Recent theoretical efforts have produced sharp asnmtttmitoicc convergence guarantees for algorithms based on the Langevin diffusion (or variants thereof), with many results known when strongly log-concave or satisfies isoperimetric assumptions [see, e.g., Durmus et al., 2019, Shen and Lee, 2019, Vempala and Wibisono, 2019, Chen et al., 2020, Dalalyan and Riou-Durane 2020, Chewi et al., 2021, Lee et al., 2021, Ma et al., 2021, Wu et al., 2022].
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

More recently, Variational Inference (V1) has emerged as viable alternative to MCMC [Jordan et al., 1999. Wainwright and Jordan, 2008, Blei al., 2017]. The goal of VI is approximate the posterior by more tractable distribution + E P such that
A arg arg min KL(p x).
PEP
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 x, that is, we take P to be the class of non-degenerate Gaussian distributions on IRd [Barber and Bishop, 1997, Seeger, 1999, Honkela and Valpola, 2004, Opper and Archambeau, 2009, Zhang et al., 2018]. The expressive power of the variational model may then be further increased by considering mixture distributions [Lin et al., 2019, Daudel and Douc, 2021, Daudel et al., 202 l].
Although the solution R 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-suce such as the potentially long "burn-in" period of samplers and the lack of effective stopping criteria for the algorithn-whin highlight the fact that obtaining an approximation for the covariance matrix of x 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 V1 problem (2).
still producing informative summary statistics. In this regard, we
However, in contrast with MCMC methods, to date there have not been many theoretical guar- antees for VI, even when т strongly log- concave and P is taken to be the class of Gaus- sians (m,, E.) The problem stems from the fact that the objective in (2) is typically non- convex in the pair (m, E). Obtaining such guar- antees remains 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 frame- work where (TT)), represents the marginal laws of stochastic process indexed by time t, which can be discrete or continuous. The goal to recursively build Gaussian approximation to (Tz)
Figure I: Left: randomly initialized mixture of 20 Gaussians (the initial covariances are depicted as red circles) and contour plot of a logistic target m. Right: contour lines of a mixture of Gaussians approximation obtained from the gradient flow in Section 5.
As concrete example, suppose that (mz) denotes the marginal law of the solution to the Langevin diffusion (1). In the context of Bayesian optimal filtering and smoothing, Särkkä [2007] proposed the following heuristic. Let (m,, Σx) denote the mean and covariance matrix of ग.. Then, it can be
checked (see Section B.4) that
m E (X))
(3)
S = 21 - (X,) @ (Xx-my) m.) + (X, - m.) & VV(X.)
where X, ~ TG. These ordinary differential equations (ODEs) are intractable because they involve expectations under the law of X, ~ Te, which is not available to the practitioner. However, if we replace X, Te with a Gaussian Y, ~ = Amm,, Sim Ex) with the same mean and covariance as X then the system of ODEs
m = - (V)
Σ = 21 (X) & (Y₁ m;) (Y m.) & VV(Y))
yields well-defined evolution of Gaussian distributions (Pt):20- which we may optimistically believe to be a good approximation of (TT))20- Moreover, the system of ODEs can numerically approxi- mated efficiently in practice using Gaussian quadrature rules to compute the above expectations. This is the principle behind the unscented Kalman filter [Julier et al., 2000].

---

In the context of the Langevin diffusion, Särkkä's heuristic (4) provides promising avenue towards computational VI. Indeed, since π x exp(-V) is the unique stationary distribution of the Langevin diffusion (1), an algorithm to approximate (ITe 120 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 approxi- mation (pe) in Särkkä's heuristic (4) arises precisely as the gradient flow of the Kullback-Leibler (KL) divergence (L(- x) on the wassersnsnn space of Gaussian distributions on Rd endowed with the 2- 2-Wasserstein distance from optimal transport [Villani, 2003]. This perspective allows 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 Jordan et al. [1998], 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 gradient flow of KL(- IT) on the Wasserstein space P2 (I(")) of probability measures with finite second moments. Subsequently, in order to emphasize the Riemannian geometry underlying this result, Otto [2001] developed his eponymous calculus P2 (TR*), framework which has had tremendous impact in analysis, geometry, PDE, probability, and statistics.
Inspired by this perspective, we show in Theorem that Särkkä's approximation (pt) is also gra- dient flow of KL(- || IT), with the main difference being that it is constrained to lie on the submanifold (W of P2 consisting of Gaussian distributions, known as the wasserrennnn manifold. In turn, our result paves the way for new theoretical understanding via the powerful theory of gradient flows. As first step, using well-known results about convex functionals on the Wasserstein space, we show in Corollary that (pt) converges rapidly to the solution of the VI problem (2) with P = (Rd) 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-as -on-asymptoiic discretization guarantees as pioneered by Dalalyan [2017] for MCMC. To further close the theoretical gap between VI and the state of the art for MCMC. we propose in Section 4.2 stochastic gradient descent (SGD) algorithm as a time discretization of the Mures-Wassersse
gradient flow. This algorithm comes with convergence guarantees that establish VI as a solid competitor to MCMC not only from 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 the substantially more flexible class of mixtures of Gaussians. Namely, the space of mixtures of Gaussians can be identified as a Wasserstein space over (WRR and hence inherits Otto's 's differential calculus. Leveraging this viewpoint, in Theorem 3 we derive the gradient flow of KL(- || T) over the space of mixtures of Gaussians and propose to implement via system of interacting partieles. Unlike typical particle-based algorithms. here our particles correspond to Gaussian distributions, and the collection thereof to Gaussian mixture which better equipped to approximate 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 broadly useful extension of classical particle methods for PDEs.
Related work. Classical VI methods define a parametric family P {po 0 E Q} and minimize 0 ++ KL |l T) over 8 θ using off-the-shelf optimization algorithms [Paisley et al., 2012, Ran- ganath et al., 2014]. 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. this regard, previous approaches to VT using natural gradients implicitly employ different geometry [Wu et al., 2019, Huang et al., 2022, Khan and Havard, 2022], namely the agparameterizationnnnaar invarrann Fisher-Rao geometry [Amari and Nagaoka, 2000J. The application of Wasserstein gradient flows to VI was
3

---

introduced earlier in work on normalizing flows and Stein Variational Gradient Descent (SVGD) [Liu and Wang, 2016, Liu, 2017].
Our work falls in line with a number of recent papers aiming to place VI on a solid theoretical footing [Alquier al., 2016, Wang and Blei, 2019, Domke, 2020, Knoblauch al., 2022]. Some of these works in particular have obtained non-asymptotic algorithmic guarantees for specific examples, see, e.g., Challis and Barber [2013], Alquier and Ridgway [2020].
The connection between VI and Kalman filtering was studied in the static case by Lambert et al. [2021, 2022a], and extended to the dynamical case by Lambert et al. [2022b], providing first justification of Särkkä's heuristic in terms local variational Gaussian approximation. In particular, the closest linear process to the Langevin diffusion (1) is Gaussian process governed by a McKean -lasovv 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 [Carlen and Gangbo, 2003, Caglioti et al., 2009, Tudorascu and Wunsch, 2011, Eberle al., 2017) 7], although our interpretation of Särkkä's heuristic is, to the best of our knowledge, new.
2 Background
In order to define gradient flows on the space of probability measures, we must first endow this space with geometry; see Appendix for more details. Given probability measures μ and on R&, define the 2- Wasserstein distance
[ inf (yeeu,.)) / /Ike y/12 dr(x, v) 1/2
V2(,,
where C(p, v/) is the set of couplings of μ and v, that is, joint distributions on Rd x Id whose marginals are j and respectively. This quantity is finite as long as μ and v belong to the space P2(R of probability measures over Id with finite second moments, The 2- 2-Wasserstein distance has the interpretation of measuring the smallest possible mean squared displacement of mass required to transport μ to v; we refer to Villani [2003, 2009], Santambrogio [2015] for textbook treatments on optimal transport. Unlike other notions of distance between probability measures, such as the total variation distance, the 2- 2-Wasserstein distance respects the geometry of the underlying space Rd leading to numerous applications in modern data science [see, e.g., Peyré and Cuturi, 2019].
The space (IR)))). W2) is a metric space [Villani, 2003, Theorem 7.3], and we refer to it as the Wasserstein space. However, as shown by Otto [Otto, 2001]. it has far richer geometric structure: formally, (P2 (IE"), W2) can be viewed as a Riemannian manifold, a fact which allows for considering gradient flows of functionals on P2 (I. ). A fundamental example of such a functional is the KL divergence KL(- iT) to a target density π x exp(-V) on RV, for which Jordan et al. [1998] showed that the Wasserstein gradient flow is the same as the evolution of the marginal law of the Langevin diffusion (I). This optimization perspective has had tremendous impact on our understanding and development of MCMC algorithms [Wibisono, 2018].
3 Variational inference with Gaussians
In this section we describe our problem using two equivalent approaches: variational approach based on a modified version of the JKO scheme of Jordan et al. [1998] (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).
3.1 Variational approach: the Bures- Bures-JKO scheme
The space of non-degeneraee Gaussian distributions on Id equipped with the W2 distance forms the Bures Wasserstesr space BW(R C P2(R On ( the Wasserstein distance W (po. P1) between two Gaussians Pa = [m Σn) and pi = A/ (m,, Σ,) admits the following closed form:
W (po. PI) = ||mo - -m,1l² 2 + 132 (Eo, E
(5)

---

where B² (E,, Σn)
tr(E. + ΣΙ - 2 m)
is the squared Bures metric [Bures, 1969].
Given target density cxp(-V) R and with step size h >> O, may define the iterates of the proximal point algorithm
1 2h
Pk+1, = arg min (EBW(R4)
[L(p
W (2.
(6)
Using (5), this is explicit optimization problem involving the mean and covariance matrix of Although (6) not solvable in closed form, by letting / Y 0 we obtain a limiting curve (pr (20) via pt. limas Ple/Aj.h- which can be interpreted as the Bures-Wasserstein gradient flow of the KL divergence KL(- т). This procedure mimics the JKO scheme [Jordan et al., 1998] with the additional constraint that the iterates lie in (W)) and therefore call it the Bures- scheme.
3.2 Geometric approach: the Bures- Wasserstein gradient flow of the KL divergence
In the formal sense of Otto descrihed above, (W( submanifold of P2 (Rd). Moreover, since Gaussians can be parameterized by their mean and covariance, (W(R)) can be identified with the manifold I x + where $ the cone of symmetrie positive definite BW( is genuine Riemannian manifold in its own right [see Modin, 2017, Malagò et al., 2018, Bhatia et al., 2019], and gradient flows can be defined using Riemannian geometry [do Carmo, 1992]. See Section B.3 for more details. Since the functional μ + ((4) = KL(μ || Tт) defined over (2 restricts to a functional over (W we can also consider the gradient flow of I over the Bures-War Srees-Wassersteinn space; note that this latter gradient flow is necessarily a curve (p:) such that each is Gaussian measure.
d matrices. Hence,
3.3 Variational inference via the Bures-Wasserster gradient flow
Using either approach, we can prove the following theorem.
Theorem 1. Let x exp(-V) be the target density on R.d. Then, the limittin curve (Pt) where Pt (m,, Σx) is obtained via the Bures-JKO scheme (6), or equivalently, the Bures-Wassersen gradient flow (Pt) of the KL divergence KL(- || T). satisfies Sürkkä's 's system of ODEs (4).
Proof. The proof using the Bures-JKO scheme is given in Section A.l 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 the submanifold (W(. Equipped with this interpretation, are able to obtain information about the asymptotic behavior of the approximation (p:)-2o- Namely, we can hope that it converges to constrained minimizer म arg solution to the VI problem (2). In the next section, we show that this convergence in fact holds as soon V is convex, and moreover with quantitative rates.
KL(p || T), i.c., precisely the
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., Opper and Archambeau, 2009], and we recall standard facts about स here for completeness. It is known that á satisfies the equations
Ex VV 0 and
= --
where E 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
zero). particular, follows from (7) that if V2V enjoys the bounds al 3 V²V 31 for some -oo I S B S co,
then any solution á to the constrained problem also satisfies 8 1 E (0)
3.4 Continuous-time convergence
Besides providing an intuitive interpretation Särkkä's heuristic, Theorem readily yields conver- gence criteria for the system (4) which rest upon general principles for gradient flows. We begin with
5

---

VI
10-- KL divergence
10:
KL diver divergence
Laplace
10a
10*
300
100
step
step
Figure 2: Two left plots: approximation bimodal target and 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 1.4 for details).
key observation. For functional F BW((R R U (0o) and R, we say that F is a-convex if for all constant-speed geodesics (Pr)rE(o.1) in (W())"),
F(p:) s (1 11 FD)) tp()) at (1
(2o.
te e 1].
2
Lemma 1. For any R. if VVV ž al, then KL(- π) is a-convex on BW(R.).
Proof. The assumption that V2V 2 al entails that the functional KL(- T) is a-convex on the entire Wasserstein space (P2 (R W2) [see, Villani, 2009, Theorem 17.15]. Since (W(R)) is geodesically convex subset of P2 (IR) (see Section B.3), then the geodesics in BW(R") agree with geodesics (2()))) from which it follows that KL(- π) is a-convex on BW(Rd).
Consequently, we obtain the following corollary. Its proof is postponed to Section D. Corollary 1. Suppose that v²V Z al for some a E IR Then, for any po BW(R)), there is unique solution the BW (R) gradient of KL(- T) started at po- Moreover:
1. a > 0, then all 2 0, (2 +) < exp(- -oot) W (po:
2. 1f 0, then for all I 2 0, KL(pr π) KL(## )) S exp( 2od) (KL(po
|| (] KL(# m))
3. If 0, then for all 0, KL(p, TT) KL π 22
W (P..
The assumption that 22 > a for assumption the MCMC literature. Under this same assumption, Corollary yields convergence for the wassersesnnn gradient flow KL(- || т): however, the flow must first be discretized in time for implementation. If we assume additionally that the smoothness condition V²V 3 31 holds, then 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.
i.e., that is strongly log-concave, is standard
Time discretization of the Bures-Wassersee gradient flow
We are now equipped with dual perspectives a dynamical solution to Gaussian ODE and gradient flow. Each perspective leads to a different implementation. On the one hand, we discretize the system of ODEs defined (4) using numerical integration. On the other, we discretize the gradient flow using stochastic gradient descent in the Bures wassersesnn space.
4.1 Numerical integration of the ODEs
The system of ODEs (4) can be integrated in time using classical Runge-Kutte scheme. The expectations under a Gaussian support are approximated by cubature rules used in Kalman filter- ing [Arasaratnam and Haykin, 2009].
Moreover, square root version of the ODE also considered to ensure that covariance matrices remain symmetric and positive. See Appendix 1.2 for more details. We have tested our method on bimodal distribution and on a posterior distribution arising from logistic regression problem. We observe fast convergence shown in Figure 2.

---

4.2 Bures warss-sasssrrtiiin SGD and theoretical guarantees for VI
Although the ODE discretization proposed in the preceding section enjoys strong empirical perfor- mance, it is unclear how to quantify its impact on the convergence rates established in Corollary 1. Therefore, we now propose stochastic gradient descent algorithm over the Bures Wassersesnn space, for which useful analysis tools have been developed [Chewi et al., 2020, Altschuler et al., 2022]. 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 maintains sequence of Gaussian distributions (p% ken) here (mx, Στ) 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 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.
Algorithm 1 Bures Wasserstein SGD
Require: strong convexity parameter a >> 0; step size 0; mean my and covariance matrix Eg for k = 1, 1.....N .N do
draw sample X ~ Pk set m*++1 < mh - h (V())) set Mr + 1 - (V -E
set Σx x- M&
set Dx+1 f clip >
end for
Theorem 2. Assume that 0 al 3 V 1.
a² and that we initialize Algorithm I matrix satisfying 9 I II 3 Also, assume that h S 60 Then, for all k E N,
36dh
E E A) exp(-akh) (2 a) a²
a² In particular, we obtain E W? (Ps, +) e² provided we set h x and the number of iterations to be k 2 - Log(W2(po, #)/E)
The upper bound V < 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 is similar in flavor to modern results for MCMC, both in terms of the assumptions (Hessian bounds and query access to the derivatives' 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 closing the gap between the two approaches.
5 Variational inference with mixtures of Gaussians
Thus far, we have shown that the tractability Gaussians can be readily exploited in the context of Busse-Wassnnttein 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 natural and powerful alternative; indeed, universal approximation of arbitrary probability measures by mixtures of Gaussians is well-known [see, e.g., Delon and Desolneux, 2020]. As we show next, the space of mixtures of Gaussians can also be equipped with Wasserstein structure which gives rise to implementable gradient flows.
5.1 Geometry of the space of mixtures of Gaussians
We begin with the key observation already made by Chen et al. [2019], that any mixture of Gaussians can be canonically identified with a probability distribution (the mixing distribution) over the param- eter space θ Rd x dd (the space of means and covariance matrices). Explicitly probability measure μ E P()) corresponds to a Gaussian mixture as follows:
Pp / du((B).
(8)
A notable downside of Algorithm I is the requirement of a Hessian oracle for V. which results in higher per-iteration cost than typical MCMC samplers.
7

---

where is the Gaussian distribution with parameters & E θ Equivalently, can be thought of as probability measure over BW(R") and hence the space of Gaussian mixtures on I2 can be identified with the Wasserstein space (B(BB over the Bures-Wassersser the distance (5) between Gaussian measures. Indeed, the theory of optimal transport can be developed with any Riemannian manifold (rather than IRd) as the base space [Villani, 2009]. As before, the space (B(B)( is endowed with a formal Riemannian structure, which respects the geometry of the base space (W4)),, and we can consider Wasserstein gradient flows over P2(BW((R)),
space which is endowed with
Note that this framework encompasses both discrete mixtures of Gaussians (when μ is discrete measure) and continuous mixtures of Gaussians. In the case when the mixing distribution μ is discrete, the geometry of P2 (BWBBR)) was studied by Chen et al. [2019], Delon and Desolneux [2020]. An important insight of our work, however, is that it is fruitful to consider the full space (B(B)(I)) for deriving gradient flows, even if we eventually develop algorithms which propagate a finite number of mixture components.
5.2 Gradient flow of the KL divergence and particle discretization
We consider the gradient flow of the KL divergence functional
+ I()) = (L(P T)
over the space P2 (BBBB(())) The proof of the following theorem is given in Section F.
Theorem 3. The gradient flow (He) of the functional 3 defined in (9) over (2( can be described as follows. Let 0g (mo, E.0) ~ No, and let Og (m,, Σ.)) evolve according to the ODE
PHc (Y))
my =
E
V
In
(10)
E = E 2
PM (Y
E
Eg
E
2
In
PK (Y2) π
where Y, ~ N(m,, Sieee,,,,, Ex). Then 0₂ ~ 14+.
The gradient flow in Theorem 3 describes the evolution of a particle 0g which describes the parameters of 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 μlg- By evolving all those particles through (10), which interact with each other via the term PP. 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 μr, which, in turn, corresponds to Gaussian mixture. Since an infinite number of Gaussian particles is impractical, consider initializing this evolution finitely supported distribution foo, thus corresponding more familiar Gaussian mixture model with a finite number of components:
N
1
1 N
No
N
6
(
N
8
(mg
(
Ppo =
Σ
P
A=1
Interestingly, it can be readily checked that the system of ODEs (10) thus initialized maintains finite mixture distribution:
N
N
Σ 500
= Σ (x))))
H4 =
x(())
N
N
1=1
where the parameters 0(0) = (mc (8)
1, evolve according to the following interacting particle
system, for [N]
Pиee (y.())) π
m v In
(11)
(1))
v in π
Pus (r
E
E
2
In
Pas (Y, π
(12)
where (1) Pact. This finite system of particles can now be implemented using the same numerical tools as for Gaussian see Section J. Note that due to this property of the dynamics, can hope

---

best to converge to the best mixture of N Gaussians approximating т, but this approximation error expected vanish as N - oc. Also, similarly (4), it is possible to write down Hessian-free updates using integration by parts, see Appendix A.2.
The ahove system of particles may also be derived using proximal point method similar to the Bures-JKO scheme, see Section A.2. Indeed, infinitesimally, it has the variational interpretation
(4h1
(N) L+h
arg min (00 EE
KL
1 N
Σ i=1
Pows
σ + 2
W2 (pocs, Paf Popo)
Reassuringly, Equations (11)-(12 reduce to (4) when μ4 imm, Snn) is a point mass, indicating that the theorem provides natural extension of our previous results. However, although the model (8) substantially more expressive than the Gaussian 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 I 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 f, the distribution Pu, is continuous. This extension provides considerably more flexibility-from mixture of point masses particle-based algorithms hitherto considered for either sampling [Liu and Wang, 2016, Liu, 2017, Duncan et al., 2019, Chewi et al., 2020], or solving partial differential equations [Carrillo et al., 2011, 2012, Bonaschi et al., 2015, Craig and Bertozzi, 2016, Carrillo et al., 2019, Craig et al., 2022].
mixture of Gaussians-
to interacting
step No (Approximated density)
step O
N*
N
10
10-
10-
10-
5-
-10
-10-
-10-
-10
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 The right figure shows the final step with the approximated density in contour-lines. More figures available in Appendix J.
Conclusion
Using the powerful theory of Wasserstein gradient flows, derived algorithms for VI using either Gaussians or mixtures 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 for time discretization which, unlike classical particle methods, maintains a contimuous probability distribution at each time.
Gaussian particle method
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 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 KI. divergence with respect to the Sintersttin-ishherr geometry [Liero et al., 2016, Chizat et al., 2018, Liero et al., 2018], which yields 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
We (m;
)
Let v,( N(m, (1) ). and let (0)
w Then, the system of ODEs is given by
Pur (8)
mi E v
In
(()
E
2 In
Prse (Y,
(1)
E
())
E
2
In
Pho (Y,
A(0
Pa E
1 N
Σ
Eln π
P (v( (()))
π
We have implemented these equations and their empirical performance is encouraging. However, fuller investigation of algorithms for VI with changing weights is beyond the scope of this work and we leave it for future research.
Code for the experiments is available at httph::t//// htttpp://gichabb.commarrrree http:///aaarrellabeee/eeeeeee
Acknowledgments and Disclosure of Funding
We thank Yian Ma for helpful discussions, as well anonymous reviewers for useful references and suggestions. ML acknowledges support from the French Defence procurement agency (DGA). SC is supported by the Department of Defense (DoD) through the National Defense Science & Engineering Graduate Fellowship (NDSEG) Program. French government under the management of the Agence Nationale de la Recherche part of the "Investissements d' "'avenni" program, reference ANR- 19-P31A-0001 (PRAIRIE 31A Institute), as well as from the European Research Council (grant SEQUOIA 724063), PR is supported by NSF grants I1S-1838071, 1, DMS-2022448, and CCF-2106377.
and MI. acknowledge support from the
References
Pierre Alquier and James Ridgway. Concentration of tempered posteriors and of their variational approximations. Ann. Statist., 48(3):1475-1497, 2020.
Pierre Alquier, James Ridgway, and Nicolas Chopin. On the properties of variational approximations of Gibbs posteriors. J. Mach. Learn. Res., 17:Paper No. 239, 41, 2016.
Jason Altschuler, Sinho Chewi, Patrik Gerber, and Austin J. Stromme. Averaging on the Bures- Wasserstein manifold: dimension-free convergence of gradient descent. arXiv e-prints, art. arXiv:2106.08502. 2022.
Shun-ichi Amari and Hiroshi Nagaoka. Methods of information geometry, volume 191 of Translations of Mathematical Monographs. American Mathematical Society, Providence, RI, 2000.
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows in metric spaces and in the space of probubility measures. Lectures in Mathematics ETH Züiicc,, Birkhäuser Verlag, Basel, second edition, 2008
Ienkaran Arasaratnam and Simon Haykin. Cubature Kalman filters. IEEE Trans. Automat. Control, 54(6):1254 1269, 2009.
Nihat Ay, Jürgen Jost, Hông Vân Lč, and Lorenz Schwachhöfer. Information geometry, volume 64 of Ergebnisse der Mathematik und ihrer Grenzgebiete. 3. Folge. A Series of Modern Surveys in Mathematics [Results Mathematics and Related Areas. 3rd Series. Series of Modern Surveys in Mathematics). Springer, Cham, 2017.
Dominique Bakry, Ivan Gentil, and Michel I. Ledoux. Analysis and geometry of Markov diffusion oper- ators, volume 348 of Grundlehren der Mathematischen Wissenschaften [Fundamental Principles of Mathematical Sciences]. Springer, Cham, 2014.
10

---

David Barber and Christopher Bishop. Ensemble learning for multi-layer networks. In Advances in Neural Information Processing Systems, volume 10, 1997.
Jean-David Benamou and Yann Brenier. A numerical method for the optimal time-continuous mass transport problem and related problems. In Monge Ampère equation: applications to geometry and optimization (Deerfield Beach, FL, 1997), volume 226 of Contemp. Math., pages 1-11. Amer. Math. Soc., Providence, RI, 1999.
Rajendra Bhatia, Tanvi Jain, and Yongdo Lim. On the Bures-Wasserstii distance between positive definite matrices. Expo. Math., 161-:196,1 2019.
Christopher M. Bishop. Pattern recognition and machine learning. Information Science and Statistics. Springer, New York, 2006.
David M. Blei, Alp Kucukelbir, and Jon D. McAuliffe.
Variational inference: A review for statisticians. Journal of the American Statistical Association, 1 112(518):859-8777 2017.
Giovanni Bonaschi, José Carrillo, Marco Di Francesco, and Mark A. Peletier. Equivalence of gradient flows and entropy solutions for singular nonlocal interaction equations in 1D. ESAIM Control Optim. Cale. Var., 21(2):414-441, 2015.
Donald Bures. An extension Kakutani's theorem on infinite product measures to the tensor product of semifinite w- -algebbas.. Trans. Amer. Math. Soc., 135:199-22 1969.
Emanuele Caglioti, Mario Pulvirenti, and Frédéric Rousset. On constrained 2-D Navier-Stokes equation. Comm. Math. Phys., 290(2):651-61) 2009.
Eric A. Carlen and Wilffri Gangbo. Constrained steepest descent in the 2-Wasserstein metric. Ann. of Math. (2), 157(3):807-846 2003.
José A. Carrillo, Marco Di Francesco, Alessio Figalli, Thomas Laurent, and Dejan Slepčev. Global- in-time weak measure solutions and finite-time aggregation for nonlocal interaction equations. Duke Math. J., 156()):229-25)) 271, 2011.
José A. Carrillo, Marco Di Francesco, Alessio Figalli, Thomas Laurent, and Dejan Slepëev. Confine- ment in nonlocal interaction equations. Nonlinear Anal., 75(2):550-5588 2012.
José A. Carrillo, Katy Craig, and Francesco S. Patacchini. blob method for diffusion. Calc. Var. Partial Differential Equations, 58(2):Paper No. 53, 53, 2019.
Edward Challis and David Barber. Gaussian Kullback-Leible approximate inference. Mach. Learn. Res., 14:2239-222 2013.
Yongxin Chen, Tryphon T. Georgiou, and Allen Tannenbaum. Optimal transport for Gaussian mixture models. IEEE Access, 7:6269-6278, 2019.
Yuansi Chen, Raaz Dwivedi, Martin Wainwright, and Bin Yu. Fast mixing of Metropolized Hamiltonian Monte Carlo: benefits of multi-step gradients. J. Mach. Learn. Res., 21:Paper No. 92, 71, 2020.
Sinho Chewi, Thibaut Le Gouic, Chen Lu, Tyler Maunu, and Philippe Rigollet. SVGD as a kernelized Wasserstein gradient flow of the chi-squared divergence. In Advances in Neural Information Processing Systems, volume 33, pages 2098-2109 2209, 2020.
Sinho Chewi, Tyler Maunu, Philippe Rigollet, and Austin J. Stromme. Gradient descent algorithms for Bures-Wasserstech barycenters. In Proceedings of the Conference on Learning Theory, volume 125, pages 1276-1304 PMLR, 09-12 Jul 2020.
Sinho Chewi, Murat Erdogdu, Mufan B. Li, Ruoqi Shen, and Matthew Zhang. Analysis of Langevin Monte Carlo from Poincaré to log-Sobolev. arXiv e-prints, art. arXiv:2112.12662 2021.
I Lénaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and Trançois-Xavier Vialard. An interpolating distance between optimal transport and Fisher-Rao metrics. Found. Comput. Math., 18(1):1-44, 2018.
11

---

Katy Craig and Andrea L. Bertozzi. A blob method for the aggregation equation. Math. Comp., 85 1300):16111177, 2016.
Katy Craig, Karthik Elamvazhuthi, Matt Haberland, and Olga Turanova. A blob method for inho- mogeneous diffusion with applications to multi-agent control and sampling. arXiv e-prints, art. N.ii:2202.11927,, March 2022.
Arnak Dalalyan. Theoretical guarantees for approximate sampling from smooth and log-concave densities. Journal of the Royal Statistical Society. Series B (Statistical Methodology), 79(3): 651-676, 2017.
Arnak S. Dalalyan and Lionel Riou-Durand. On sampling from log-concave density using kinetic Langevin diffusions. Bernoulli, 26(3):1956-198) 2020.
Kamélia Daudel and Randal Douc. Mixture weights optimisation for alpha-divergence variational inference. In Advances in Neural Information Processing Systems, volume 34, pages 4397-4408, 2021.
Kamélia Daudel, Randal Douc, and François Portier. Infinite-dimensional gradient- based descent for alpha-divergence minimisation. Ann. Statist., 49(44)2250-233 2270, 2021.
Julie Delon and Agnès Desolneux. A Wasserstein-type distance in the space of Gaussian mixture models. SIAM J. Imaging Sci., 13(2):936 970, 2020.
Manfredo P. do Carmo. Riemannian geometry. Mathematics: Theory & Applications. Birkhäuser Boston, Inc., Boston, MA, 1992. Translated from the second Portuguese edition by Francis Flaherty.
Justin Domke. Provable smoothness guarantees for blxck-box variational inference. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 2587-2596. PMLR, 13-18 Jul 2020.
Andrew Duncan, Nikolas Nuesken, and Lukasz Szpruch. On the geometry of Stein variational gradient descent. arXiv e-prints, art. 191.0099,,,,0944, December 2019.
Alain Durmus, Szymon Majewski, and Btažej Miasojedow. Analysis of Langevin Monte Carlo via convex optimization. J. Mach. Learn. Res., 20:Paper No. 73, 46, 2019.
Simon Eberle, Barbara Niethammer, and André Schlichting. Gradient flow formulation and longtime behaviour of constrained Folker-Planck equation. Nonlineur Anal., 158: 142-167, 2017.
Antti Honkela and Harri Valpola. Unsupervised variational Bayesian learning of nonlinear models. In Advances Neural Information Processing Systems, volume 17, 2004.
Daniel Z. Huang, Jiaoyang Huang, Sebastian Reich, and Andrew M. Stuart. Efficient derivative-free Bayesian inference for large-scale inverse problems. arXiv e-prints, art. MtXiv:2204.04386, 2022.
Michael Jordan, Zoubin Ghahramani, Tommi S. Jaakkola, and Lawrence K. Saul. An introduction to variational methods for graphical models. Mach. Learn., 37(2):183-23) 1999.
Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the Fokker-Planck equation. SIAM Journal on Mathematical Analysis, 29(1):1-17, 1998.
Simon Julier and Jeffrey K. Uhlmann. Unscented filtering and nonlinear estimation. Proceedings of the IEEE, 92(3):401-422 2004.
Barrant-Whyte. Simon Julier, Jeffrey Uhlmann, and Hugh F. Durrant- Whyte. A new method for the nonlinear transformation of means and covariances in filters and estimators. IEEE Trans. Automat. Control, 45(3):477-482, 2000.
Mohammad Emtiyaz Khan and Rue Håvard. The Bayesian learning rule.
2022.
12

---

Jeremias Knoblauch, Jack Jewson, and Theodoros Damoulas. An optimization view on Bayes' rule: reviewing and generalizing variational inference. Journal of Machine Learning Research, 23 (132):1-109, 2022.
Mare Lambert, Silvère Bonnabel, and Francis Bach. The limited-memory recursive variational Gaussian approximation (L-RVGA). hal-03501920, 2021.
Marc Lambert, Silvère Bonnabel, and Francis Bach. The recursive variational Gaussian approximation (R-VGA). Statistics and Computing, 32(1):10, 2022a.
Mare Lambert, Silvère Bonnabel, and Francis Bach. The continuous-discrete variational Kalman filter (CD-VKF). In 2022 61 IEEE Conference Decision and Control (CDC), 2022b.
Yin Tat Lee, Ruoqi Shen, and Kevin Tian. Structured logconcave sampling with restricted Gaussian oracle. In Proceedings of the Conference on Learning Theory, volume 134, pages 2993-3050, 15-19 Aug 2021.
Matthias Liero, Alexander Mielke, and Giuseppe Savaré. Optimal transport in competition with reaction: the Hellinger-Kantorovvia distance and geodesic curves. SIAM J. Math. Anal., 48(4): 2869-2911 2016.
Matthias Liero, Alexander Mielke, and Giuseppe Savaré. Optimal entropy-transport problems and a new Hellinger- Kantorovcc distance between positive measures. Invent. Math., 11117) 2018.
Wu Lin, Mohammad E. Khan, and Mark Schmidt. Fast and simple natural gradienn variational inference with mixture of exponential-famie approximations. In Proceedings of the International Conference on Machine Learning, volume 97, pages 3992-4002, 09-15 Jun 2019.
Dong C. Liu and Jorge Nocedal. On the limited memory BFGS method for large scale optimization. Math. Programming, 45(3, (Ser. 5)):503-528, 1989.
Qiang Liu. Stein variational gradient descent as gradient flow. In Advances in Neural Information Processing Systems, volume 30, 2017.
Qiang Liu and Dilin Wang. Stein variational gradient descent: general purpose Bayesian inference
algorithm. In Advances in Neural Information Processing Systems, volume 29, 2016.
Yulong Lu, Jianfeng Lu, and James Nolen. Accelerating Langevin sampling with birth-death. arXiv
e-prints, art, urXiv:190S, 19w.s99965...,,,. May 2019.
Yi-An Ma, Niladri S. Chatterji, Xiang Cheng, Nicolas Flammarion, Peter L. Bartlett, and Michael I. Jordan. Is there an analog of Nesterov acceleration for gradient-based MCMC? Bernoulli, 27(3): 1942 1992, 2021.
Luigi Malagò, Luigi Montrucchio, and Giovanni Pistone. Wasserstein Riemannian geometry of Gaussian densities. Inf. Geom., 1(2):137-179, 2018.
Klas Modin. Geometry of matrix decompositions seen through optimal transport and information geometry. Geom. Mech,, 933:335-3 2017.
Martin Morf, Bernard Levy, and Thomas Kailath. Square-root algorithms for the continuous time linear least squares estimation problem. In 1977 IEEE Conference on Decision and Control including the 16th Symposium on Adaptive Processes and A Special Symposium on Fuzzy Set Theory and Applications, pages 944-947, 1977.
Manfred Opper and Cédric Archambeau. The variational Gaussian approximation revisited. Neural Comput., (3)3:7767779, 2009.
Felix Otto. Dynamics of labyrinthine pattern formation in magnetic fluids: mean-field theory, Arch. Rational Mech. Anal., 141(1):63-103, 1998.
Felix Otto. The geometry of dissipative evolution equations: the porous medium equation. Comm. Partial Differential Equations, 26(1-2);101-177 2001.
13

---

John Paisley, David M. Blei, and Michael I. Jordan. Variational Bayesian inference with stochastic search. In Proceedings of the International Conference on Machine Learning, pages 1363-1377 2012.
Gabriel Peyré and Marco Cuturi. Computational optimal transport: with applications data science. Now, 2019.
Rajesh Ranganath. Sean Gerrish, and David M. Blei. Black box variational inference. In Proceedings of International Conference on Artificial Intelligence and Statistics, volume 33, pages 814-822, Reykjavik, Iceland, 22-25 Apr 2014.
Filippo Santambrogio. Optimal transport for applied mathematicians, volume 87 of Progress in Nonlineur Differential Equations and their Applications. Birkhäuser/Springer, Cham, 2015. Calculus of variations, PDEs, and modeling.
Simo Särkkä. On unscented Kalman filtering for state estimation of continuous-time nonlinear systems. IEEE Trans. Automat. Control, 52(9):663111411 2007.
Matthias Seeger. Bayesian model selection for support vector machines, Gaussian processes and other kernel classifiers. In Advances in Neural Information Processing Systems, volume 12, 1999.
Ruoqi Shen and Yin Tat Lee. The randomized midpoint method for log-concave sampling. In Advances in Neural Information Processing Systems, volume 32, 2019.
Adrian Tudorascu and Marcus Wunsch. On nonlinear, nonlocal parabolic problem with conservation of mass, mean and variance. Comm. Partial Differential Equations, 36(8): 1454, 2011.
Santosh Vempala and Andre Wibisono. Rapid convergence of the unadjusted Langevin algorithm: isoperimetry suffices. In Advances in Neural Information Processing Systems 32, pages 8094-88 8106. 2019.
Cédric Villani. Topics in optimal transportation, volume 58 of Graduate Studies in Mathematics. American Mathematical Society, Providence, RI, 2003.
Cédric Villani. Optimal transport, volume 338 of Grundlehren der Mathematiscchn Wissenschaften [Fundamental Principles of Mathematical Sciences]. Springer-Ver Verlag, Berlin, 2009. Old and new.
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., (14657))1147-1116) 2019.
Andre Wibisono. Sampling as optimization in the space of measures: the Langevin dynamics as 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,220 2022.
Lin Wu, Emtiyaz K. Mohammad, and Schmidt Mark. Stein's 's lemma for the reparameterization trick with exponential family mixtures. 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.
14


Reviews:
[{"review_id": "ex0_EqVR9Nx", "reviewer": "Reviewer_8QPd", "summary": "This paper proposes a variational inference method to find the optimal approximation of the target distribution, which has the smallest KL divergence to the target distribution in the family of Gaussian distribution or Gaussian mixture. For the family of Gaussian distribution, the convergence analysis of the continuous-time dynamics and the discrete-time stochastic gradient algorithms are established under the strong log-concavity assumption of the target density. For the Gaussian mixture, they propose the gradient flow of the KL divergence and the particle discretization by parametrizing the Gaussian mixture model as a probability measure over the Bures-Wasserstein space. \n", "questions": "The interpretation of the Gaussian mixture model as the measure on the Bures-Wasserstein space is interesting. I’m wondering if the target density itself is a Gaussian mixture, can the proposed algorithm exactly fit the posterior distribution?\n\nFor the Gaussian mixture, are there objective functionals other than the KL divergence, for instance, maximum mean discrepancy, to make the resulting problem on the measure of the Bures-Wasserstein space convex? \n\nFor the Gaussian family, the gradient flow in equ. (4) does not require the calculation of the Hessian of the log posterior density. However, the gradient flow of Gaussian mixture in equ. (10) does require the Hessian. As the calculation of the Hessian for high dimensional inference problem can be problematic and the proposed method itself is a first-order method, is it possible to only utilize the gradient information of the log-posterior density?\n\nSome literature related to the Bures-Wasserstein space are missing, see [1, 2].\n\n[1] Malagò, Luigi, Luigi Montrucchio, and Giovanni Pistone. \"Wasserstein Riemannian geometry of Gaussian densities.\" Information Geometry 1.2 (2018): 137-179.\n\n[2] Modin, Klas. \"Geometry of matrix decompositions seen through optimal transport and information geometry.\" arXiv preprint arXiv:1601.01875 (2016).\n", "limitations": "Yes.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n- The paper is well-motivated\n- It provides an interesting interpretation of the variational Kalman filtering from the perspective of Wasserstein gradient flow.\n- The convergence analysis of the proposed Bures-Wasserstein SGD is novel.\n- The numerical results of the proposed method for Gaussian mixture seem promising.\n\nWeakness:\n- The theoretical contribution is not enough. For the Gaussian family, the analysis of the continuous-time dynamics seems to be a direct application of the convergence results of Wasserstein gradient flow. For the Gaussian mixture, although the author acknowledge that the problem is non-convex, some theoretical analysis of the proposed method can strengthen the paper. See details in the Questions. \n- The structure of the paper can be improved. The algorithm for Gaussian mixture is of great interest in practice but the detailed description of that algorithm is not presented in the paper.\n- The comparison of the proposed method for Gaussian mixture with other variational inference methods using Gaussian mixture is missing. Some experiments on Bayesian inference can strengthen the paper. \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "TnJPSTWDwyK", "reviewer": "Reviewer_qdtG", "summary": "The paper provides an analysis of variational inference algorithms based on Gaussian approximation of variational distributions. The paper exploits the theory around Bures-Wasserstein spaces of Gaussian measures to come construct a sequence of Gaussian distributions (or mixtures) that converge to the target measure, i.e., the variational approximation of the target distribution. \nThe paper uses Wasserstein gradient flows for their theoretical analysis and establishes that Saerkkae’s heuristic coincides with the sequence of distributions arising in the gradient flow of the KL divergence. \nThis analysis is extended to the use Gaussian mixture distributions. \nTheoretical guarantees are derived and discussed in detail. The theoretical results rely on the assumption of log concave distribution, an assumption that is also made in the study of Langevin type algorithms. ", "questions": "see above", "limitations": "Yes", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The theoretical contributions of the paper look sound and the work is well situated among other contributions in the field. \n\nTo add: \nIn the related work section it might be worthwhile to add reference to the work on sequential Monte Carlo algorithms. \n\nMinor: \nLine 272: generality? \n\nContributions:\n- Add more numerical results in the main body of the paper? This would make the work stronger on spark practical applications\n\nLine 335 : “ …. yields a new Gaussian particle method which appears to be significantly more powerful than classical particle methods.” \nThis is an overly strong statement that requires more justification. You are not explicitly discussing how your method is more powerful. This should happen both from an empirical and theoretical point of view (which you both don’t do). \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "CbhN-frIqd1", "reviewer": "Reviewer_E7dv", "summary": "This paper produces an algorithm for computing variational posteriors that are normals or mixtures of normals. As the title suggests, this is achieved via Wasserstein gradient flows. Remarkably, the system of the Langevin-type stochastic process this gives rise to was proposed before---but as a heuristic in the setting of sequential bayesian inference/kalman filtering. Relative to existing variational methods, the benefits of the proposed approach are a number of theoretical guarantees that are usually not available, and hold under reasonable (even if somewhat limiting) assumptions so long as one uses a single normal (rather than a mixture of normals) as the variational family. ", "questions": "I would advise the authors to tone down the claims about VI being a mostly heuristic field without significant theoretical advances; and to embedd their paper into the context of [1-4] and related work provided above.", "limitations": "I think that the paper does not point to the limitations of its results enough, but I don't think this is such a big problem---there is only so much space a NeurIPS submission is allowed to cover, and the paper does a better job than most others. The three most important limitations are\n(1) V has to be convex\n(2) guarantees only hold for the single Gaussian measure case; and not even for the mixture Gaussian case [and to be fair, this one is pointed out explicitly!]\n(3) it is unclear if the algorithm is actually faster/better than standard VI techniques/optimisation methods in practice", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Originality: \nThe work is highly original from a technical point of view, & to the best of my knowledge the first attempt at using Wasserstein Gradient flows in a computationally attractive way for the purposes of variational inference. Its originality lies not so much in the end goal of the algorithm---there are plenty other variational methods using Gaussian families/mixtures of Gaussians---but in using an area of mathematics that has hitherto seen little use in the context of variational inference to derive a new algorithm for achieving an (old/well-known) end goal and guaranteeing a number of attractive properties for the resulting algorithm. \n\nQuality:\nThis paper is of the highest technical standard. Beyond that, the authors have made a very serious effort to show that the method is not just a mathematical curiosity, but practically meaningful. It has been a while since a paper has inspired me to dive into a different literature and read around 5-10 other papers to fully appreciate all its details, but this paper has managed it effortlessly.\n\nClarity: \nThe paper articulates its (complicated) subject clearly, and leaves the reader satisfied that she or he can understand every component part of the paper. The supplementary material is extensive, and maintains an exceptionally high standard.\n\nSignificance:\nI think the paper is significant in showing how to use Wasserstein gradient flows for the design of attractive variational methods. If I have one point of criticism towards the paper, it would be that the research within the field of variational inference is misrepresented. For example, the paper's abstract states that \"VI is still poorly understood and dominated by heuristics. In this work, we propose principled methods for VI [...]\". I think this sentiment does severe injustice to a lot of good work that has happened on VI in recent times. I understand that the authors are focused specifically on optimisation algorithms with theoretical guarantees that are common in optimisation, but other theoretical work from other fields and specifically statistics has recently put VI on a much more solid foundation than it used to be. To name but a few examples, there are PAC-Bayesian bounds for variational objectives [1] as well as asymptotic studies [2], work motivating variational methods as constrained optimisation [3], and even some prior work on convergence guarantees (under conditions that are in some sense milder than those of the current paper) [4]. I think the paper would do good toning down the claims of the abstract & the other claims in the paper aligned with the spirit of the abstract---or at least make them more specific to the kind of guarantees/theory that the paper is focused on. It certainly isn't accurate to describe all of the research on VI as 'dominated by heuristics'. I think the paper would also do well positioning itself in the context of a recent line of work reaching across various disciplines and dedicated to making VI a more principled set of tools as in [1-4].\n\n[1] https://jmlr.org/papers/volume17/15-290/15-290.pdf\n\n[2] https://www.tandfonline.com/doi/full/10.1080/01621459.2018.1473776\n\n[3] https://www.jmlr.org/papers/v23/19-1047.html\n\n[4] https://proceedings.mlr.press/v119/domke20a.html\n\n\nTypos:\n- p.7, l.270 'variation' -> 'variational'", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept