Stochastic variance-reduced Gaussian variational inference on the Bures-Wasserstein manifold

## Abstract

Abstract Optimization in the Bures-Wasserstein space has been gaining popularity in the machine learning community since it draws connections between variational inference and Wasserstein gradient flows. The variational inference objective function of Kullback–Leibler divergence can be written as the sum of the negative entropy and the potential energy, making forward-backward Euler the method of choice. Notably, the backward step admits a closed-form solution in this case, facilitating the practicality of the scheme. However, the forward step is no longer exact since the Bures-Wasserstein gradient of the potential energy involves “intractable” expectations. Recent approaches propose using the Monte Carlo method – in practice a single-sample estimator – to approximate these terms, resulting in high variance and poor performance. We propose a novel variance-reduced estimator based on the principle of control variates. We theoretically show that this estimator has a smaller variance than the Monte-Carlo estimator in scenarios of interest. We also prove that variance reduction helps improve the optimization bounds of the current analysis. We empirically demonstrate that the proposed estimator gains order-of-magnitude improvements over previous Bures-Wasserstein methods.

## 1 Introduction

Variational inference (VI) (Wainwright et al., ; Blei et al., ) provides a fast and scalable alternative to Markov chain Monte Carlo (MCMC), especially for inference tasks in high dimensions. The main principle of VI is to approximate a complicated distribution $\pi$, e.g., posterior distribution in Bayesian inference, by a simpler tractable family of distributions. The approximation $\mu$ within the family is obtained by solving an optimization problem, providing a closed-form representation and e.g. efficient sampling by construction. The choice of the optimization method is heavily influenced by the assumptions made on the approximation family and the information about $\pi$ that can be obtained, ranging from classical coordinate ascent algorithms for mean-field approximations for targets with conditional conjugacy structure (Blei et al., ) to gradient methods using score-function approximations to avoid assumptions on the target density (Ranganath et al., ) or flexible approximations parameterized with neural networks (Rezende & Mohamed, ).

We focus on Gaussian approximations (Honkela & Valpola, ; Opper & Archambeau, ; Xu & Campbell, ; Quiroz et al., ) but with a particular emphasis on the recent research line in the Wasserstein geometric viewpoint of this family (Lambert et al., ; Diao et al., ). Regarding the target $\pi$, we assume access to second order gradients, typically computed by automatic differentiation, similar to the above works. Gaussian VI offers strong statistical guarantees at the optimal solution (Katsevich & Rigollet, ), offers an easy way of modelling dependencies between the variables and, thanks to the
Bernstein-von Mises theorem (Van der Vaart, ), becomes asymptotically exact for Bayesian inference at the limit of infinite observations.

Recently, there has been emerging interest in Gaussian VI with a new geometric Riemannian optimization perspective (Lambert et al., ; Diao et al., ). The family of non-degenerative Gaussian distributions can be parameterized by its mean and covariance matrix, $\mu_{\theta}$ with $\theta=(m,\Sigma)$, henceforth denoted as $\Theta=\mathbb{R}^{d}\times\mathcal{S}^{d}_{++}$ where $\mathcal{S}^{d}_{++}$ is the set of $d\times d$ symmetric, positive definite matrices. Classical VI employs conventional optimization algorithms (Paisley et al., ; Titsias & Lázaro-Gredilla, ; Kucukelbir et al., ) to minimize the Kullback-Leibler (KL) divergence $D_{\mathrm{KL}}({\mu}_{\theta}\|\pi)$ over the parameter space $\Theta$ equipped with the Euclidean geometry.
Lambert et al. () argue that
because the optimization problem is over the space distributions, it is more natural to use the geometry of this space rather than the geometry of the parameter space. The space of Gaussian distributions has a rich, meaningful and tractable geometry known as Bures-Wasserstein (BW) geometry that benefits optimization. Lambert et al. () subsequently established a theoretical framework for performing VI using the BW geometry, which we adopt in this paper.

Figure: Figure 1: Left: Optimization trajectories of our method compared to SGVI (Diao et al., ) and BWGD (Lambert et al., ). The target is a $50$-dimensional Gaussian distribution, visualized here via the marginal distributions of the first two coordinates. Each ellipse represents a contour of a Gaussian: the black is the initial distribution, the red is the target, and the greys are intermediate steps. Our method is dramatically more stable and finds a more accurate final approximation. Right: the corresponding KL divergence, confirming our method is orders of magnitude more accurate.
Refer to caption: x1.png

Let $\pi(x)\propto\exp(-V(x))$ be the target distribution and consider the VI problem

$$ $\displaystyle\hat{\pi}\in\operatorname*{arg\,min}_{\mu\in\operatorname{BW}{( \mathbb{R}^{d})}}D_{\mathrm{KL}}(\mu\|\pi),$ (1) $$

where $\operatorname{BW}{(\mathbb{R}^{d})}$ is the Bures-Wasserstein space of Gaussian distributions with non-generative covariance matrix. The BW space is a Riemannian manifold whose geodesic distance is the Bures-Wasserstein distance. This setting nicely interplays the theory of optimal transport, Wasserstein gradient flows, and variational inference. The optimization problem (1) can be reformulated as

$$ $\displaystyle\hat{\pi}\in\operatorname*{arg\,min}_{\mu\in\operatorname{BW}( \mathbb{R}^{d})}\mathcal{F}(\mu),\quad\text{where}\quad\mathcal{F}(\mu):= \mathcal{E}_{V}(\mu)+\mathscr{H}(\mu).$ (2) $$

Here, $\mathcal{E}_{V}(\mu)=\int{V(x)}d\mu(x)$ is the potential function and $\mathscr{H}(\mu)=\int\log(\mu(x))d\mu(x)$ is the negative entropy. A conceptual and established idea to minimize a functional $\mathcal{F}$ is to perform gradient flow on $\mathcal{F}$ with respect to the geometry of $\operatorname{BW}(\mathbb{R}^{d})$. To be implementable, the flow must be discretized. Lambert et al. () use forward Euler discretization, resulting in a scheme named Bures-Wasserstein stochastic gradient descent (BWGD).

Diao et al. () remark that forward-backward Euler (Bauschke & Combettes, ) should be used instead due to the objective’s composite nature and the entropy’s non-smoothness. This method iteratively applies a forward step to the potential energy $\mathcal{E}_{V}$ and a backward step (proximal operator) to the negative entropy $\mathscr{H}$. They also observe that the backward step in the BW space has a closed-form solution (Wibisono, ). This is crucial because this step is known to be intractable (or computationally expensive) in the full Wasserspace space (Wibisono, ; Salim et al., ; Mokrov et al., ; Luu et al., ). Although the bottleneck of the FB Euler, which is the backward step, has been resolved in this case, the forward step becomes problematic where one has to compute the Bures-Wasserstein gradient of $\mathcal{E}_{V}$ instead of the “friendly” Wasserstein gradient that is just $\nabla V$. The Bures-Wasserstein gradient is not always available in closed form, i.e., at $\mu\in\operatorname{BW}(\mathbb{R}^{d})$, it is given only implicitly by the map $x\mapsto\mathbb{E}_{\mu}\nabla V+(\mathbb{E}_{\mu}\nabla^{2}V)(x-m_{\mu})$ where $m_{\mu}$ is the mean of $\mu$ (Lambert et al., ). This is the orthogonal projection of the Wasserstein gradient onto a tangent space of the Bures-Wasserstein manifold (Chewi et al., ). For general $V$, these expectations are intractable even though the underlying distribution is a Gaussian. Diao et al. () proposed using the Monte Carlo (MC) method with one sample to estimate these expectations at each iteration: sample $X\sim\mu$ and use $\nabla V(X)$ and $\nabla^{2}V(X)$ as unbiased estimators for $\mathbb{E}_{\mu}\nabla V$ and $\mathbb{E}_{\mu}\nabla^{2}V$, respectively. This scheme is called Stochastic Gaussian VI (SGVI).

The problem with SGVI building on this principle is that the Monte Carlo estimates needed for the BW gradient are typically too noisy, particularly in high dimensions, as shown in our experiments (Sect. 5). In practice, high-variance estimators require small step sizes, leading to slow and inefficient convergence. We resolve this fundamental limitation by proposing a variance-reduced estimator with minimal computational overhead while providing robust theoretical guarantees. Fig. 1 shows the improvement over SGVI and BWGD in practice.
Bures-Wasserstein geometry offers a meaningful transition from the initial distribution to the target distribution, and our method follows the path smoothly and is particularly stable around the optimum.

##### Contributions.

We propose a novel variance-reduced estimator for $\mathbb{E}_{\mu}\nabla V$ that does not use any extra samples, with minimal per-iteration computational overhead, using the control variates approach (Owen, ). Our idea is that the variational distribution $\mu$ should be similar to the target distribution $\pi(x)\propto\exp(-V(x))$ as $\mu$ gets closer and closer to $\pi$, so the density of $\mu$ can be used to construct a correlated control variate for the Monte-Carlo estimator $\nabla V(X)$. Sect. 3 presents the detailed construction and its rationale.

On the theoretical side, we derive the following insights:

- Thm. 1
Under a mild smoothness assumption, we prove that there is a region around the optimal solution $\hat{\pi}$ where our estimator has guaranteed smaller variance than the MC estimator.
- Thm. 2
If $V$ is strongly convex, we prove that the proposed estimator has a smaller variance than the MC estimator at every $\mu\in\operatorname{BW}(\mathbb{R}^{d})$ whenever $\mu$ has sufficiently large (greater than a controllable threshold) principal curvatures.

We further show in Thm. 3 and Thm. 4 that whenever variance reduction happens along the algorithm’s iterates, the effect will enter the convergence analysis and improve the optimization bounds derived in (Diao et al., ). These theorems solidly back our proposed method.

On the practical side, we show that reusing the Cholesky decomposition of the covariance matrix (needed to sample from a multivariate Gaussian) keeps the computational overhead of the control variable negligible.
Despite being only a minimal modification to the Monte Carlo estimator, the proposed estimator achieves significant improvements in our experiments.

## 2 Background

A function $f:\mathbb{R}^{d}\to\mathbb{R}$ is called L-smooth if $\|\nabla f(x)-\nabla f(y)\|\leq L\|x-y\|$ for all $x,y\in\mathbb{R}^{d}$. If $f$ is twice continuously differentiable, we define the Laplacian operator of $f$ as $\Delta f=\sum_{i=1}^{d}{(\partial^{2}/\partial x_{i}^{2})}f$. Note that the Laplacian is the trace of the Hessian, $\Delta=\operatorname{Tr}(\nabla^{2})$.

### 2.1 Bures-Wasserstein geometry

We denote by $\mathcal{P}_{2}(\mathbb{R}^{d})$ the space of probability measures $\mu$ over $\mathbb{R}^{d}$ with finite second-moment, i.e., $\int{\|x\|^{2}}d\mu(x)<+\infty$. Equipped with the Wasserstein distance

$$ $W_{2}^{2}(\mu,\nu)=\inf_{\gamma\in\Gamma(\mu,\nu)}{\int_{X\times X}\|x-y\|^{2} d\gamma(x,y)}$ (3) $$

where $\Gamma(\mu,\nu)$ is the set of probability measures over $X\times X$ whose marginals are $\mu$ and $\nu$, the space $\mathcal{P}_{2}(\mathbb{R}^{d})$ becomes the metric space called the Wasserstein space (Ambrosio et al., ). We call $\gamma\in\Gamma(\mu,\nu)$ a (transport) plan and any $\gamma$ that achieves the optimal value in (3) an optimal plan. A pair of random variables whose joint distribution is an optimal plan is called an optimal coupling (between $\mu$ and $\nu$). When $\mu$ is absolutely continuous with respect to the Lebesgue measure, Briener theorem (Brenier, ) asserts that the optimal plan is unique and is given by $(I,T_{\mu}^{\nu})_{\#}\mu$ where $T_{\mu}^{\nu}=\nabla g$ for some convex function $g$. We call $T_{\mu}^{\nu}$ the optimal transport map from $\mu$ to $\nu$. Apart from being a metric space, the Wasserstein space also enjoys some nice properties of Riemannian geometry. Otto’s calculus (Otto, ) endows the Wasserstein space with a formal Riemannian structure, facilitating gradient flows and optimization.

We denote by $\operatorname{BW}(\mathbb{R}^{d})$ the space of Gaussian distributions with non-generative covariance matrices. The Wasserstein distance between two Gaussian distributions $p_{0}=\mathcal{N}(m_{0},\Sigma_{0})$ and $p_{1}=\mathcal{N}(m_{1},\Sigma_{1})$ is given in the closed-form formula $\mathcal{W}_{2}^{2}(p_{0},p_{1})=\|m_{0}-m_{1}\|^{2}+\mathcal{B}^{2}(\Sigma_{0
},\Sigma_{1})$ where $\mathcal{B}^{2}(\Sigma_{0},\Sigma_{1})=\operatorname{Tr}(\Sigma_{0}+\Sigma_{1}
-2(\Sigma_{0}^{\frac{1}{2}}\Sigma_{1}\Sigma_{0}^{\frac{1}{2}})^{\frac{1}{2}})$ is the Bures metric. The optimal transport map is also given in a closed form in this case:
$T_{p_{0}}^{p_{1}}(x)=m_{1}+{\Sigma_{0}}^{-\frac{1}{2}}\left({\Sigma_{0}}^{
\frac{1}{2}}\Sigma_{1}{\Sigma_{0}}^{\frac{1}{2}}\right)^{\frac{1}{2}}{\Sigma_{
0}}^{-\frac{1}{2}}(x-{m}_{0})$.

The BW space is a geodesically convex subset of the Wasserstein space, meaning that a geodesic curve joining two Gaussians lies entirely inside the BW space. The BW space is a Riemannian manifold in its own right.
Let $\mu=\mathcal{N}(m,\Sigma)\in\operatorname{BW}(\mathbb{R}^{d})$, the tangent space of $\operatorname{BW}(\mathbb{R}^{d})$ at $\mu$ is the space of symmetric affine map denoted as $T_{\mu}\operatorname{BW}(\mathbb{R}^{d})=\{x\mapsto S(x-m)+a~{}|~{}a\in\mathbb
{R}^{d},S\in\mathcal{S}^{d}\}$ where $\mathcal{S}^{d}$ is the space of symmetric $d\times d$ matrices. The Riemannian metric defined using the inner product of elements in this tangent space is identified as the $L^{2}(\mu)$ inner product restricted to this space. Given $U,V\in T_{\mu}\operatorname{BW}(\mathbb{R}^{d})$, the metric is $\langle U,V\rangle_{\mu}:=\int{\langle U(x),V(x)\rangle}d\mu(x)$. This Riemannian metric induces the geodesic distance in $\operatorname{BW}(\mathbb{R}^{d})$ that is given by the Wasserstein distance. We refer to (Altschuler et al., ) for further discussions on BW geometry.

### 2.2 Stochastic Gaussian VI

We refer to (Diao et al., ) for a detailed discussion and relevant terminologies. We briefly explain the stochastic Gaussian VI here to motivate our proposed variance reduction version in Sect. 3. Recall from (2) that we aim to minimize $\mathcal{F}(\mu)=\mathscr{H}(\mu)+\mathcal{E}_{V}(\mu)$ over $\operatorname{BW}(\mathbb{R}^{d})$. At the optimum of $\mathcal{F}$, $\hat{\pi}=\mathcal{N}(\hat{m},\hat{\Sigma})$, first-order optimality condition reads (Opper & Archambeau, ; Lambert et al., ; Diao et al., )

$$ $\displaystyle\mathbb{E}_{\hat{\pi}}\nabla V=0\quad\text{and}\quad\mathbb{E}_{ \hat{\pi}}\nabla^{2}V=\hat{\Sigma}^{-1}$ (4) $$

which is derived by zeroing the Bures-Wasserstein gradient of the objective function.

A natural idea to minimize $\mathcal{F}$ over $\operatorname{BW}(\mathbb{R}^{d})$ is to perform gradient flow on $\mathcal{F}$ using the BW geometry of $\operatorname{BW}(\mathbb{R}^{d})$. When the gradient flow is applied over the entire Wasserstein space $\mathcal{P}_{2}(\mathbb{R}^{d})$, it corresponds to the Langevin diffusion (Jordan et al., ), with one of its discretizations being an MCMC method called the unadjusted Langevin algorithm (Roberts & Tweedie, ). When restricted to $\operatorname{BW}(\mathbb{R}^{d})$, the gradient flow can be formulated using Riemannian geometry (Do Carmo, ), as $\operatorname{BW}(\mathbb{R}^{d})$ forms a true Riemannian manifold. This flow is a curve of Gaussian distributions, characterized by the time-dependent evolution of their mean and covariance matrix. Recently, Lambert et al. () shows that this evolution is governed by Särkkä’s ODEs developed in the context of variational Kalman filtering (Särkkä, ).

The negative entropy $\mathscr{H}$ is convex along generalized geodesics but it is a nonsmooth functional. If $V$ is smooth, it induces the smoothness of $\mathcal{E}_{V}$. Therefore, it is natural to apply forward-backward Euler that alternates between two steps: at iteration $k$,

$$ $\displaystyle\mu_{k+\frac{1}{2}}$ $\displaystyle=(I-\eta\nabla_{\operatorname{BW}}\mathcal{E}_{V}(\mu_{k}))_{\#} \mu_{k}\quad\triangleleft\text{forward step}$ $\displaystyle\mu_{k+1}$ $\displaystyle=\operatorname*{arg\,min}_{\mu\in\operatorname{BW}(\mathbb{R}^{d} )}\left\{\mathscr{H}(\mu)+\dfrac{1}{2\eta}W_{2}^{2}\left(\mu,\mu_{k+\frac{1}{2 }}\right)\right\}\quad\triangleleft\text{backward step}$ $$

where $\nabla_{\operatorname{BW}}$ denotes the Bures-Wasserstein gradient. The backward step is also known as the proximal step in the optimization literature or the JKO (Jordan, Kinderlehrer, and Otto) step (with restriction in $\operatorname{BW}(\mathbb{R}^{d})$) in the context of Wasserstein gradient flow (Jordan et al., ). The backward step is intractable in the full Wasserstein space and hence requires (oftentimes expensive) numerical approximations (Mokrov et al., ; Luu et al., ). On the other hand, if restricted to $\operatorname{BW}(\mathbb{R}^{d})$, this step admits a closed-form solution (Wibisono, ): let $\mu_{k+\frac{1}{2}}=\mathcal{N}(m_{k+\frac{1}{2}},\Sigma_{k+\frac{1}{2}})$, then $\mu_{k+1}$ is a Gaussian distribution with mean $m_{k+1}=m_{k+\frac{1}{2}}$ and variance matrix
$\Sigma_{k+1}=\frac{1}{2}\left(\Sigma_{k+\frac{1}{2}}+2\eta I+[\Sigma_{k+\frac{
1}{2}}(\Sigma_{k+\frac{1}{2}}+4\eta I)]^{\frac{1}{2}}\right).$
This tractability of the backward is the main motivation for (Diao et al., ) to study and develop FB Euler in this scenario. The forward step, however, is not always analytically available since the BW gradient of $\mathcal{E}_{V}$, at iterate $k$,

$$ $\displaystyle\nabla_{\operatorname{BW}}\mathcal{E}_{V}(\mu_{k}):x\mapsto \mathbb{E}_{\mu_{k}}\nabla V+(\mathbb{E}_{\mu_{k}}\nabla^{2}V)(x-m_{k}),$ $$

involves intractable expectations. Diao et al. () propose using Monte Carlo approximation for these expectations: sample ${X}_{k}\sim\mu_{k}$ and use $b_{k}:=\nabla V({X}_{k})$ and $S_{k}:=\nabla^{2}V({X}_{k})$ as unbiased estimators for $\mathbb{E}_{\mu_{k}}\nabla V$ and $\mathbb{E}_{\mu_{k}}\nabla^{2}V$, respectively.

## 3 Stochastic variance-reduced Gaussian VI

We present our ideas on constructing stochastic variance-reduced estimators from first principles.
We recall from Sect. 2.2 that stochastic Gaussian VI approximates, at iteration $k$,

$$ $\displaystyle\mathbb{E}_{\mu_{k}}\nabla V\approx b_{k}:=\nabla V({X}_{k})\quad \text{and}\quad\mathbb{E}_{\mu_{k}}\nabla^{2}V\approx S_{k}:=\nabla^{2}V({X}_{ k})\quad\text{where }{X}_{k}\sim\mu_{k}.$ (5) $$

These estimators are typically noisy. Any number of MC samples can be used, but already one is unbiased and proposed by earlier works; we also focus on the single-sample case for computational efficiency.
We aim to design better unbiased estimators for either $\mathbb{E}_{\mu_{k}}\nabla V$ or $\mathbb{E}_{\mu_{k}}\nabla^{2}V$ in the sense that their variances are smaller than those of $b_{k}$ and $S_{k}$, building on the control variates approach (Owen, ) (also see the discussions in Defazio et al. (); Luu ()).

Let us first describe briefly the core idea of control variates in helping reduce the variance. Let $\theta$ be the quantity of interest and $X$ be an unbiased estimator for $\theta$, i.e., $\mathbb{E}X=\theta$. A *control variate* is a random variable $Y$ with a known mean so that $Y$ is correlated with $X$. The random variable $Z=X+c(\mathbb{E}Y-Y)$ where $c\in\mathbb{R}$ is then a tractable *unbiased estimator* for $\theta$. The variance of $Z$ is

$$ $\displaystyle\mathrm{Var}Z=\mathrm{Var}X+c^{2}\mathrm{Var}Y-2c\mathrm{Cov}(X,Y).$ (6) $$

If $X,Y$ are highly correlated in the sense that $2\mathrm{Cov}(X,Y)>\mathrm{Var}Y$, we immediately get $\mathrm{Var}Z<\mathrm{Var}X$ for any $c\in(0,1]$. So, we achieve a reduction in variance by using $Z$. On the other hand, if $X,Y$ are correlated ($\mathrm{Cov}(X,Y)>0$) but not highly correlated, we can also obtain variance reduction effects whenever $c$ is positive and small enough. Furthermore, given the parabolic form with respect to $c$ in (6), one can pinpoint the optimal value of $c$ is $c^{*}:=\mathrm{Cov}(X,Y)/\mathrm{Var}(Y)$, resulting in the maximal variance reduction $\mathrm{Var}Z=(1-\operatorname{Corr}(X,Y)^{2})\mathrm{Var}X<\mathrm{Var}X$ where $\operatorname{Corr}(X,Y)$ denotes correlation between $X$ and $Y$.

We now return to our problem and seek variance-reduced estimators of the forms

$$ $\displaystyle\tilde{b}_{k}:=\nabla V({X}_{k})+c(\mathbb{E}(Z_{k})-Z_{k})\quad \text{and}\quad\tilde{S}_{k}:=\nabla^{2}V({X}_{k})+d(\mathbb{E}(W_{k})-W_{k})$ $$

where $c,d>0$ and $Z_{k},W_{k}$ are a random vector and a random matrix, respectively. Let us first focus on $\tilde{b}_{k}$. As discussed, $Z_{k}$ should be (element-wise) highly correlated with $\nabla V({X}_{k})$ while $\mathbb{E}(Z_{k})$ remains efficiently computable.
We look for $Z_{k}=\nabla U({X}_{k})$ so that $\nabla U$ is as close to $\nabla V$ as possible. We are in the context of approximating $\pi(x)$ by the VI distribution $\mu_{k}=\mathcal{N}(m_{k},\Sigma_{k})$, so it is natural to expect that

$$ $\displaystyle-\nabla V(x)=\nabla\log\pi(x)\approx\nabla\log f(x;m_{k},\Sigma_{ k})=-\Sigma_{k}^{-1}(x-m_{k}).$ $$

where $f(x;m_{k},\Sigma_{k})\propto\exp\left(-\tfrac{1}{2}(x-m_{k})^{\top}\Sigma_{k}^
{-1}(x-m_{k})\right)$
is the PDF of $\mu_{k}$.
Therefore, we propose using $Z_{k}=\Sigma_{k}^{-1}({X}_{k}-m_{k})$ as a control variate. We have $\mathbb{E}(Z_{k})=0$ since $\mathbb{E}({X}_{k})=m_{k}$. It is worth noting that $Z_{k}$ is known as the Stein/Hyvärinen score (Hyvärinen, ) of $\mu_{k}$. The estimator $\tilde{b}_{k}$ then becomes
$\tilde{b}_{k}:=\nabla V({X}_{k})-c\Sigma_{k}^{-1}({X}_{k}-m_{k})$. By applying the same reasoning to $\tilde{S}_{k}$, we can immediately conclude that $W_{k}$ is deterministic and equals $\Sigma_{k}^{-1}$. Consequently, the control variate does not affect $S_{k}$; we keep the standard estimator.
We derive Stochastic variance-reduced Gaussian VI (SVRGVI) as in Alg. 1. Note that the only difference between Alg. 1 and the SGVI in (Diao et al., ) is the estimator $\tilde{b}_{k}$, where the difference is highlighted in blue.

Fig. 2 (left) demonstrates that our proposed estimator (with $c=0.9$) achieves lower variance compared to the standard MC estimator, while both remain unbiased estimators of $\mathbb{E}_{\mu}\nabla V$. In Fig. 2 (right), we vary $c$ from 0 0 to $2$ and calculate the empirical variance of our estimator, revealing a parabolic pattern. Note that when $c=0$, the estimator reduces to the standard estiator, and for all values of $c\in(0,2)$, our proposed estimator consistently exhibits lower variance, with an optimal value of $c$ around 1. At this optimal $c$, the variance is reduced roughly by a factor of $10$. We provide theoretical justification for these empirical observations in Sect. 4.

Figure: Figure 2: Left: $\pi$ is a Gaussian, VI distribution $\mu$ is in the neighborhood of $\pi$. In this case, the true gradient, i.e., the expectation $\mathbb{E}_{\mu}\nabla V$, can be computed exactly (in navy blue). Our proposed estimator with $c=0.9$ (light blue) has a smaller variance than the Monte Carlo estimator (grey). These are 1,000 samples for each estimator, generated by drawing from $\mu$ and substituting the values into the respective estimator formulas. Right: The empirical variance of our proposed estimator when $c$ varies from 0 0 to $2$. Note that $c=0$ corresponds to the Monte Carlo estimator.
Refer to caption: x3.png

##### Minimal extra computational cost

Despite involving calculating the inverse of the covariance matrix, the computational overhead is small. Sampling from multivariate normal in step 1 in Alg. 1 typically requires obtaining the Cholesky factor of the covariance matrix, which is $O(d^{3})$ (Rasmussen & Williams, )
. With the Cholesky factor, obtaining the solution of the inverse of the matrix times a vector is $O(d^{2})$ (Rasmussen & Williams, ). As such, we can reuse this obtained Cholesky factor in step 1 to compute the inverse in step 2, which implies that the estimator adds an overhead of $O(d^{2})$, which is naturally dominated by the $O(d^{3})$ complexity of the original algorithm.

Figure: Algorithm 1 Stochastic variance-reduced Gaussian Variational Inference (SVRGVI)

## 4 Theory

In Sect. 3, we argued that, in the context of variational inference, as $\mu_{k}$ iteratively gets closer $\pi$, $\nabla V({X}_{k})$ is then (highly) correlated to $\Sigma_{k}^{-1}({X}_{k}-m_{k})$, and hence we obtain a variance reduction effect. This argument leads to the construction of the control variate in Alg. 1. One might question whether this approach remains effective when the target distribution $\pi$ is significantly distant from the BW space. Because we are constrained to the BW space, the best we can do is to get closer to $\hat{\pi}$ which is the optimal solution to the problem (1). However, $\hat{\pi}$ might still look very different from $\pi$. Notably, in Thm. 1, we rigorously show that within a certain neighbourhood of $\hat{\pi}$ (to be defined later), our proposed estimator consistently reduces variance, regardless of how different $\pi$ is to a Gaussian distribution. Let us first introduce Lem. 1 to pave the way for Thm. 1 and also to discuss the optimal $c$ in the control variate. In Lem. 1, we compute the variance of the proposed estimator by leveraging multidimensional Stein’s lemma (Lin et al., ).

###### Lemma 1

Assume that $V$ is continuously differentiable. Let $\mu=\mathcal{N}(m,\Sigma)\in\operatorname{BW}(\mathbb{R}^{d})$. Then,

$$ $\displaystyle\underbrace{\mathbb{E}\|\nabla V(X)-c\Sigma^{-1}(X-m)-\mathbb{E} \nabla V(X)\|^{2}}_{\text{variance of our estimator}}$ $\displaystyle=\underbrace{\mathbb{E}\|\nabla V(X)-\mathbb{E}\nabla V(X)\|^{2}} _{\text{variance of the Monte-Carlo estimator}}+\underbrace{c^{2}\operatorname {Tr}(\Sigma^{-1})-2c\operatorname{Tr}(\mathbb{E}\nabla^{2}V(X))}_{\text{extra term}},~{}\text{where }X\sim\mu.$ $$

Proof of Lem. 1 is given in Appendix A.1. Lem. 1 compares the variance of the proposed estimator and the Monte Carlo estimator at a given $\mu\in\operatorname{BW}(\mathbb{R}^{d})$. Recall that the first-order optimality condition (4) of $\hat{\pi}$ reads $\mathbb{E}_{\hat{\pi}}\nabla^{2}V=\hat{\Sigma}^{-1}$. Consequently, at $\hat{\pi}$, the *extra term* in Lem. 1 is simplified as $c(c-2)\operatorname{Tr}(\hat{\Sigma}^{-1})$ which is negative whenever $c\in(0,2)$ and minimized for $c=1$. Therefore, at $\hat{\pi}$, our estimator is always better than the Monte Carlo estimator for $c\in(0,2)$.

###### Remark 1

A practical merit of Lem. 1 is that it implies the optimal value for $c$ to get maximum variance reduction at $\mu$ is $c^{*}=\operatorname{Tr}(\mathbb{E}_{\mu}\nabla^{2}V)/\operatorname{Tr}(\Sigma^
{-1})$. Applying this to Alg. 1, we can pick the adaptive sequence $\{c_{k}\}$ as

$$ $\displaystyle c^{*}_{k}=\dfrac{\operatorname{Tr}(\mathbb{E}_{\mu_{k}}\nabla^{2 }V)}{\operatorname{Tr}(\Sigma_{k}^{-1})}\approx\dfrac{\operatorname{Tr}(S_{k}) }{\operatorname{Tr}(\Sigma_{k}^{-1})}:=c_{k}.$ (7) $$

Again, this computation of $c_{k}$ incurs a negligible extra cost to Alg. 1. We also remark that around $\hat{\pi}$, optimality condition (4) implies the optimal value $c^{*}$ indeed is around $1$.

In Thm. 1, we further show that when the Laplacian $\Delta V$ is smooth, the proposed estimator has a smaller variance than the Monte Carlo estimator in a region around $\hat{\pi}$.

###### Theorem 1 (Variance reduction around the optimal solution)

Assume that the Laplacian $\Delta V$ is $\ell$-smooth.
For any control variate coefficient $c\in(0,2)$, define the region around $\hat{\pi}=\mathcal{N}(\hat{m},\hat{\Sigma})$:

$$ $\displaystyle\mathcal{V}(\hat{\pi},r)=\{\mu=\mathcal{N}(m,\Sigma):2\ell W_{2}( \mu,\hat{\pi})+c|\operatorname{Tr}(\Sigma^{-1})-\operatorname{Tr}(\hat{\Sigma} ^{-1})|<r\}$ $$

where $r=(2-c)\operatorname{Tr}(\hat{\Sigma}^{-1})>0$ is the region’s radius. For any $\mu\in\mathcal{V}(\hat{\pi},r)$, the proposed estimator has a smaller variance than the Monte Carlo estimator.

Proof of Thm. 1 is given in Appendix A.2 with the main idea being that the smoothness of the Laplacian $\Delta V$ propagates the improvement of the proposed estimator at $\hat{\pi}$ to its neighbourhood. We additionally observe that, for small $c>0$, the region $\mathcal{V}(\hat{\pi},r)$ effectively reduces to the Wasserstein ball $\mathcal{B}(\hat{\pi},\ell^{-1}\operatorname{Tr}(\hat{\Sigma}^{-1}))$.

Thm. 1 applies to arbitrary $\pi$, only requiring a mild smoothness condition of its second derivative. In the next theorem, we show that when $\pi$ is strongly log-concave ($\pi$ is now more similar to a Gaussian), variance reduction happens not only around $\hat{\pi}$ but also in many regions of interest.

###### Theorem 2 (Variance reduction for large-curvature distributions)

If $V$ is $\alpha$-strongly convex for some $\alpha>0$, for any control variate $c>0$, the proposed estimator has a smaller variance than the Monte Carlo estimator at every $\mu=\mathcal{N}(m,\Sigma)$ whenever $\operatorname{Tr}(\Sigma^{-1})<\frac{2\alpha d}{c}$.

Proof of Thm. 2 is given in Appendix A.3.

###### Remark 2

A consequence of Thm. 2 is that we obtain variance reduction at $\mu=\mathcal{N}(m,\Sigma)$ whenever $\lambda_{\min}(\Sigma)>\frac{c}{2\alpha}$ regardless of the mean $m$. Here $\lambda_{\min}(\Sigma)$ is the smallest eigenvalue of $\Sigma$. Note that as $c$ is the user-specified parameter, we can gain control over the region where this effect happens. Thm. 2 indeed provides a strong variance reduction guarantee in the context of strongly log-concave sampling.

We further show in Thm. 3 and Thm. 4 that whenever variance reduction happens along the algorithm’s iterates, the effect will propagate to the convergence analyses of (Diao et al., ) and improve their theoretical bounds. Therefore, combining with Thm. 1 and Thm. 2, the overall theory strongly favours SVRGVI over SGVI.

Let $\mathcal{P}_{k}$ denote the information up to the beginning of iteration $k$, i.e., it is the $\sigma$-algebra given by $\mathcal{P}_{k}=\sigma(X_{0},X_{1},\ldots,X_{k-1})$
for $k\in\{1,2,\ldots,N-1\}$ and $\mathcal{P}_{0}$ is, by convention, the trivial $\sigma$-algebra with no information. Assuming variance reduction occurs along the algorithm’s iterates, i.e., for $k=0,1,\ldots,N-1$, it holds

$$ $\displaystyle\mathbb{E}\left(\|\nabla V(X_{k})-c_{k}\Sigma_{k}^{-1}(X_{k}-m_{k })-\mathbb{E}_{\mu_{k}}\nabla V\|^{2}|\mathcal{P}_{k}\right)\leq\tau_{k} \mathbb{E}\left(\|\nabla V(X_{k})-\mathbb{E}_{\mu_{k}}\nabla V\|^{2}|\mathcal{ P}_{k}\right)$ (8) $$

where $\tau_{k}\in[0,1]$. We also note that conditioning on $\mathcal{P}_{k}$ we discard irrelevant past information and the above conditional expectations are the variances at the current iteration.

We require (8) to hold along the iterates $k=0,1,\ldots,N-1$. We can, in principle, relax this by assuming that (8) holds for all $k\geq K_{0}$ for some $K_{0}$, ensuring that we are within the vicinity of $\hat{\pi}$ (Thm. 1) and can begin analysis after this initial warm-up period.

Under condition 8, we now show the improved bounds. Similar to (Diao et al., ), we consider log-concave and strongly-log-concave sampling, meaning that $V$ is assumed to be convex and strongly convex, respectively.

###### Theorem 3 (Convex case)

Suppose that $V$ is convex and $\beta$-smooth and the step size $0<\eta\leq\frac{1}{2\beta}$. If variance reduction happens, i.e., $\tau_{k}<1$ in (8) for $k=0,1,\ldots,N-1$, then,

$$ $\displaystyle\mathbb{E}\left(\min_{k=1,2,\ldots,N}\mathcal{F}(\mu_{k})\right)- \mathcal{F}(\hat{\pi})\lesssim\dfrac{e}{1+\frac{C\eta^{2}(1-\tau_{\max})}{2}} \left(\dfrac{1}{2\eta N}+\dfrac{C\eta}{2}\right)W_{2}^{2}(\mu_{0},\hat{\pi})+3 \eta\beta d(1+\tau_{\max})$ $$

where $\tau_{\max}:=\max\{\tau_{0},\tau_{1},\ldots,\tau_{N-1}\}<1$, $e\approx 2.718$ is the Euler’s number, $C=24\beta^{3}\lambda_{\max}(\hat{\Sigma})$, and $\lesssim$ is asymptotically at the limit of small $\eta$.

Proof of Thm. 3 is given in Appendix A.4.

###### Theorem 4 (Strongly convex case)

Suppose that $V$ is $\alpha$-strongly-convex with $\alpha>0$, and $0<\eta\leq\frac{\alpha^{2}}{48\beta^{3}}$. If variance reduction happens, i.e., $\tau_{k}<1$ in (8) for $k=0,1,\ldots,N-1$, then

$$ $\displaystyle\mathbb{E}W_{2}^{2}(\mu_{N},\hat{\pi})\lesssim\exp\left(-\dfrac{N (3-\tau_{\max})\eta\alpha}{4}\right)W_{2}^{2}(\mu_{0},\hat{\pi})+\dfrac{24(1+ \tau_{\max})\beta\eta d}{(3-\tau_{\max})\alpha}$ (9) $$

where $\tau_{\max}:=\max\{\tau_{0},\tau_{1},\ldots,\tau_{N-1}\}<1$, and $\lesssim$ is asymptotically at the limit of small $\eta$.

Proof of Thm. 4 is given in Appendix A.5.

###### Remark 3

We recall the corresponding bounds for SGVI in (Diao et al., , Thm 5.7, Thm. 5.8)

- •
Convex. $\mathbb{E}\left(\min_{k=1,2,\ldots,N}\mathcal{F}(\mu_{k})\right)-\mathcal{F}(
\hat{\pi})\lesssim\frac{eW_{2}^{2}(\mu_{0},\hat{\pi})}{2N\eta}+\frac{eC\eta}{2
}W_{2}^{2}(\mu_{0},\hat{\pi})+6\beta\eta d.$(^1^11with a minor correction.)
- •
Strongly convex. $\mathbb{E}W_{2}^{2}(\mu_{N},\hat{\pi})\lesssim\exp\left(-\frac{\alpha N\eta}{2
}\right)W_{2}^{2}(\mu_{0},\hat{\pi})+\frac{24\beta\eta d}{\alpha}$.

Putting side-by-side, we see that Thm. 3 and Thm. 4 improve all coefficients of these bounds. In particular, the scale-down involving $d$ is expected to help in high dimensions. It is also worth noting that even when we set $\tau_{\max}=0$, the noise terms in the bounds of Thm. 3 and Thm. 4 would not disappear because of another source of randomness coming from $S_{k}$.

###### Lemma 1

###### Remark 1

###### Theorem 1 (Variance reduction around the optimal solution)

###### Theorem 2 (Variance reduction for large-curvature distributions)

###### Remark 2

###### Theorem 3 (Convex case)

###### Theorem 4 (Strongly convex case)

###### Remark 3

## 5 Experiments

We demonstrate the method in a collection of controlled problems, comparing it against the recent methods for VI in the BW manifold, namely BWGD (Lambert et al., ) and SGVI (Diao et al., ).

We set the step size to $1$ for all algorithms, fix the covariate coefficient $c=0.9$, and show results for 10 runs, with bold line showing the average performance. The comparisons are shown as convergence curves, as the per-iteration cost of all methods is almost identical.

We also compare against a full-rank Gaussian approximation optimised in the Euclidean geometry (denoted as EVI), using low-variance reparameterization gradients of Roeder et al. () with ADAM optimizer, and Laplace approximation that does not optimize the KL divergence but fits a Gaussian distribution at the target mode; see Appendix B for details. As the per-iteration cost of these methods is different from the BW methods, we only report the final accuracy
for carefully optimized approximations
to show how the BW methods compare against commonly used algorithms. The Laplace approximation is omitted for the Gaussian target as it would be optimal by definition.

##### Gaussian targets

We randomly generate the means and covariances for a multivariate Gaussian target distribution $\pi$, considering dimensions of $\{10,50,200\}$.
Fig. 3
demonstrates consistent significant improvement over SGVI and BWGD.

For example, for $d=200$, the the difference between SVRGVI and SGVI/BWGD is $5$ orders of magnitude, $10^{-2}$ versus $10^{3}$.
Fig. 1 shows visually the marginals for $d=50$, providing an interpretation of the improvement seen in KL-divergence.
We also clearly outperform EVI in higher dimensions, unlike previous BW methods.

Figure: Figure 3: KL divergence for Gaussian targets of varying dimensionality.
Refer to caption: x5.png

##### Student’s t targets

We consider a multivariate Student’s t target with a degree of freedom of $4$ in $200$ dimensions. Fig. 4 (a) shows that our algorithm is again clearly the best. BWGD is not stable and, on average, performs worse than even the Laplace approximation.

##### Bayesian logistic regression

We consider a Bayesian logistic regression with a flat prior as in (Diao et al., ): given a set of covariates $X_{i}\sim\mathcal{N}(0,I_{d})$ for $i=1,2,\ldots,n$, consider

$$ $\displaystyle Y_{i}|X_{i},\theta\sim\textrm{Bernoulli}(\sigma(\langle\theta,X_ {i}\rangle)),\text{ where $\sigma$ is the sigmoid function.}$ $$

The negative log posterior is $V(\theta)=\sum_{i=1}^{n}{\left[\ln(1+e^{\langle\theta,X_{i}\rangle})-Y_{i}
\langle\theta,X_{i}\rangle\right]}$.
The model consists of $n=1000$ data points $(X_{i},Y_{i})$ with dimension $d=200$. In The optimal solution is unknown in this case, so we cannot plot the KL divergence along the iterations. Instead, we estimate the objective function of the problem (2), $\mathcal{F}(\mu_{k})$, by drawing samples from $\mu_{k}$. We denote by $\mu_{\text{best}}$ the distribution that obtains the smallest $\mathcal{F}$ among all iterations of all algorithms, comparing against that.
Fig 4 (b) shows the proposed method is again the most accurate.

Figure: Figure 4: Performance of algorithms for Student’s t target and Bayesian logistic regression.
Refer to caption: x8.png

## 6 Discussion

Various variance reduction techniques have been broadly studied in the VI literature, but mainly for methods operating in the Euclidean parameter space. Our work resembles in nature the seminal work of Roeder et al. () that demonstrated how the variance of gradient estimators for VI can be dramatically reduced by a single-line change in the algorithm: We also propose a minor modification that dramatically improves the accuracy, and should always be used. A high-level similarity lies in the heuristic that the VI distribution resembles the target distribution, allowing it to be used to construct control variates for the quantities of interest: in our case, the BW gradient of the potential energy, and in theirs, the Euclidean gradient of the ELBO (Evidence Lower BOund). This difference leads us to use the Hyvärinen score (Hyvärinen, ) as a control variate, while they rely on the Fisher score (Bishop & Nasrabadi, ). However, unlike their heuristic approach, we theoretically prove that our method effectively reduces variance, rather than relying solely on intuition.

In the BW space, Diao () considers variance reduction for large-sum structures based on the nested-loop idea by Johnson & Zhang () to reduce the stochasticity of the minibatch sampling. In contrast, our method addresses the stochasticity arising from Gaussian sampling from the VI. These two methods are complementary and can be effectively combined to enhance overall performance for a suitable class of problems.

Even though our experiments focused on synthetic targets and did not thoroughly study the effect of the step lengths, they expanded on the previous experimentation of VI optimized in the BW space. We confirm the finding of Diao et al. () that BWGD and SGVI are effectively identical except for the unstability of the former, but now show how their performance degrades in higher dimensions. We also showed how the previous methods do not always reach the accuracy of Euclidean optimization in the parameter space, whereas our improved method was consistently the best.

With the exception of the vastly improved accuracy due to significantly lower variance of the gradient estimators, our method retains all qualitative characteristics of the previous BW methods, both positive and negative. That is, we retain the theoretical convergence guarantees and asymptotic optimality for posterior inference, but also the cubic computational cost due to requiring the Hessian of the log-target and the limitation to Gaussian approximations by construction. As highlighted by Xu & Campbell () and Quiroz et al. (), there are tasks for which Gaussian approximations are highly relevant due to efficiently capturing the correlations.

## 7 Conclusion

Our main result is showing that the methods learning a variational approximation by direct optimization of the approximating distribution in the Bures-Wasserstein space of Gaussians can be made practical. The previous works by Lambert et al. () and Diao et al. () introduced the key idea and the algorithms with strong theoretical guarantees. However, as shown here they do not necessarily find as good approximation as simpler parameter-space methods, limiting the impact. Our variance reduction technique that requires only a minor modification for the SVGI algorithm completely resolves this issue, resulting in extremely stable learning.

We demonstrated substantial variance reduction, quantified to be an order of magnitude in one example task, and showed that this reduction results in orders of magnitude improvement in final approximation accuracy, over both the previous BW methods and examples of parameter-space algorithms. This improvement comes with provable variance reduction in the neighborhood of the optimal solution and for all distributions with sufficiently large curvature in the case of strong log-concave targets, and hence the proposed variance reduction technique should always be used.

#### Acknowledgments

This work is supported by the Research Council of Finland’s Flagship programme: Finnish Center for Artificial Intelligence (FCAI), and additionally by grants 363317, 345811 and 348952. We thank Kai Puolamäki for the discussions. The authors wish to acknowledge CSC – IT Center for Science, Finland, for computational resources.

## Appendix A Theory

### A.1 Proof of Lemma 1

For each $\mu=\mathcal{N}(m,\Sigma)\in\operatorname{BW}(\mathbb{R}^{d})$ and $c>0$, we denote

$$ $\displaystyle\mathcal{Q}(\mu)=\mathbb{E}\|\nabla V(X)-\mathbb{E}\nabla V({X}) \|^{2}-\mathbb{E}\|\nabla V({X})-c\Sigma^{-1}({X}-m)-\mathbb{E}\nabla V({X})\| ^{2},\quad X\sim\mu,$ $$

which is the difference between the variances of the Monte Carlo estimator and our proposed estimator. We want $\mathcal{Q}(\mu)>0$. Simple algebras simplify $\mathcal{Q}$ as

$$ $\displaystyle\mathcal{Q}(\mu)=2c\mathbb{E}\langle\nabla V(X)-\mathbb{E}\nabla V (X),\Sigma^{-1}(X-m)\rangle-c^{2}\mathbb{E}\|\Sigma^{-1}(X-m)\|^{2}.$ $$

Recall a standard result: if $X\sim\mathcal{N}(m,\Sigma)$, then its affine transformation $W=AX+b$ has the distribution $\mathcal{N}(Am+b,A\Sigma A^{\top})$. Applying this result, $W:=\Sigma^{-1}(X-m)\sim\mathcal{N}(0,\Sigma^{-1})$. Therefore

$$ $\displaystyle\mathbb{E}\|W\|^{2}=\sum_{i=1}^{d}{\mathbb{E}W_{i}^{2}}= \operatorname{Tr}(\Sigma^{-1}).$ $$

On the other hand,

$$ $\displaystyle\mathbb{E}\langle\nabla V(X)-\mathbb{E}\nabla V(X),\Sigma^{-1}(X- m)\rangle$ $\displaystyle=\mathbb{E}\langle\nabla V(X),\Sigma^{-1}(X-m)\rangle-\mathbb{E} \langle\mathbb{E}\nabla V(X),\Sigma^{-1}(X-m)\rangle$ $\displaystyle=\mathbb{E}\langle\nabla V(X),\Sigma^{-1}(X-m)\rangle-\langle \mathbb{E}\nabla V(X),\mathbb{E}\Sigma^{-1}(X-m)\rangle$ $\displaystyle=\mathbb{E}\langle\nabla V(X),\Sigma^{-1}(X-m)\rangle-\langle \mathbb{E}\nabla V(X),\Sigma^{-1}(\mathbb{E}X-m)\rangle$ $\displaystyle=\mathbb{E}\langle\nabla V(X),\Sigma^{-1}(X-m)\rangle.$ $$

Let us denote $A=\Sigma^{-1}$ and compute $\mathbb{E}\langle\nabla V(X),A(X-m)\rangle$ as follows

$$ $\displaystyle\mathbb{E}\langle\nabla V(X),A(X-m)\rangle$ $\displaystyle=\mathbb{E}\left(\sum_{i=1}^{d}{\dfrac{\partial V}{\partial x_{i} }(X)[A(X-m)]_{i}}\right)$ $\displaystyle=\mathbb{E}\left(\sum_{i=1}^{d}{\dfrac{\partial V}{\partial x_{i} }(X)\sum_{j=1}^{d}{[A]_{ij}(X_{j}-m_{j})}}\right)$ $\displaystyle=\sum_{i=1}^{d}\sum_{j=1}^{d}{[A]_{ij}\mathbb{E}\left(\dfrac{ \partial V}{\partial x_{i}}(X)(X_{j}-m_{j})\right)}.$ (10) $$

We compute $\mathbb{E}\left(\dfrac{\partial V}{\partial x_{i}}(X)(X_{j}-m_{j})\right)$ by leveraging the following Stein’s lemma (Lin et al., ).

###### Lemma 2 (Stein’s lemma)

Let $X\sim\mathcal{N}(m,\Sigma)$ be an $d$-dimensional Gaussian random variable and $g:\mathbb{R}^{d}\to\mathbb{R}$ be continuously differentiable, then

$$ $\displaystyle\mathbb{E}\left(g(X)(X-m)\right)=\Sigma\mathbb{E}(\nabla g(X)).$ $$

Applying Stein’s lemma with $g=(\partial/\partial x_{i})V$,

$$ $\displaystyle\mathbb{E}\left(\dfrac{\partial V}{\partial x_{i}}(X)(X-m)\right)$ $\displaystyle=\Sigma\mathbb{E}\left(\nabla\dfrac{\partial V}{\partial x_{i}}(X )\right)$ $\displaystyle=\Sigma\mathbb{E}\left(\left[\dfrac{\partial^{2}V}{\partial x_{1} \partial x_{i}}(X),\dfrac{\partial^{2}V}{\partial x_{2}\partial x_{i}}(X), \ldots,\dfrac{\partial^{2}V}{\partial x_{d}\partial x_{i}}(X)\right]^{\top} \right).$ $$

By comparing the $j$-th element of both sides, we get

$$ $\displaystyle\mathbb{E}\left(\dfrac{\partial V}{\partial x_{i}}(X)(X_{j}-m_{j} )\right)=\sum_{k=1}^{d}{\Sigma_{jk}\mathbb{E}\left(\dfrac{\partial^{2}V}{ \partial x_{k}\partial x_{i}}(X)\right)}.$ $$

Plugging this expression into (A.1),

$$ $\displaystyle\mathbb{E}\langle\nabla V(X),A(X-m)\rangle$ $\displaystyle=\sum_{i=1}^{d}\sum_{j=1}^{d}{[A]_{ij}\sum_{k=1}^{d}{\Sigma_{jk} \mathbb{E}\left(\dfrac{\partial^{2}V}{\partial x_{k}\partial x_{i}}(X)\right)}}$ $\displaystyle=\sum_{i=1}^{d}\sum_{j=1}^{d}{\sum_{k=1}^{d}{[A]_{ij}\Sigma_{jk} \mathbb{E}\left(\dfrac{\partial^{2}V}{\partial x_{k}\partial x_{i}}(X)\right)}}$ $\displaystyle=\sum_{i=1}^{d}\sum_{k=1}^{d}{\mathbb{E}\left(\dfrac{\partial^{2} V}{\partial x_{k}\partial x_{i}}(X)\right)\sum_{j=1}^{d}{[A]_{ij}\Sigma_{jk}}}$ $\displaystyle=\sum_{i=1}^{d}\sum_{k=1}^{d}{\mathbb{E}\left(\dfrac{\partial^{2} V}{\partial x_{k}\partial x_{i}}(X)\right)[A\Sigma]_{ik}}$ $\displaystyle=\sum_{i=1}^{d}\sum_{k=1}^{d}{\mathbb{E}\left(\dfrac{\partial^{2} V}{\partial x_{k}\partial x_{i}}(X)\right)[I]_{ik}}$ $\displaystyle=\sum_{i=1}^{d}{\mathbb{E}\left(\dfrac{\partial^{2}V}{\partial x_ {i}^{2}}(X)\right)}$ $\displaystyle=\operatorname{Tr}(\mathbb{E}\nabla^{2}V(X)).$ $$

Therefore,

$$ $\displaystyle\mathcal{Q}(\mu)=2c\operatorname{Tr}(\mathbb{E}\nabla^{2}V(X))-c^ {2}\operatorname{Tr}(\Sigma^{-1}),\quad\text{where }X\sim\mu.$ $$

###### Lemma 2 (Stein’s lemma)

### A.2 Proof of Theorem 1

Recall that Lem. 1 and the optimality condition (4) imply $\mathcal{Q}(\hat{\pi})=c(2-c)\operatorname{Tr}(\hat{\Sigma}^{-1})$.

Now let $\mu=\mathcal{N}(m,\Sigma)\in\operatorname{BW}(\mathbb{R}^{d})$, and let $(X,\hat{X})$ be the optimal coupling between $\mu$ and $\hat{\pi}$,

$$ $\displaystyle|\mathcal{Q}(\mu)-\mathcal{Q}(\hat{\pi})|$ $\displaystyle\leq 2c|\operatorname{Tr}(\mathbb{E}\nabla^{2}V({X}))- \operatorname{Tr}(\hat{\Sigma}^{-1})|+c^{2}|\operatorname{Tr}(\Sigma^{-1})- \operatorname{Tr}(\hat{\Sigma}^{-1})|$ $\displaystyle=2c|\operatorname{Tr}(\mathbb{E}\nabla^{2}V({X}))-\operatorname{ Tr}(\mathbb{E}\nabla^{2}V(\hat{X}))|+c^{2}|\operatorname{Tr}(\Sigma^{-1})- \operatorname{Tr}(\hat{\Sigma}^{-1})|$ $\displaystyle\leq 2c\mathbb{E}|\operatorname{Tr}(\nabla^{2}V({X}))- \operatorname{Tr}(\nabla^{2}{V}(\hat{X}))|+c^{2}|\operatorname{Tr}(\Sigma^{-1} )-\operatorname{Tr}(\hat{\Sigma}^{-1})|$ $\displaystyle=2c\mathbb{E}|\Delta V(X)-\Delta V(\hat{X})|+c^{2}|\operatorname{ Tr}(\Sigma^{-1})-\operatorname{Tr}(\hat{\Sigma}^{-1})|$ $\displaystyle\leq 2c\ell\mathbb{E}\|X-\hat{X}\|+c^{2}|\operatorname{Tr}(\Sigma ^{-1})-\operatorname{Tr}(\hat{\Sigma}^{-1})|$ $\displaystyle\leq 2c\ell(\mathbb{E}\|X-\hat{X}\|^{2})^{\frac{1}{2}}+c^{2}| \operatorname{Tr}(\Sigma^{-1})-\operatorname{Tr}(\hat{\Sigma}^{-1})|$ $\displaystyle=2c\ell W_{2}(\mu,\hat{\pi})+c^{2}|\operatorname{Tr}(\Sigma^{-1}) -\operatorname{Tr}(\hat{\Sigma}^{-1})|.$ $$

Therefore

$$ $\displaystyle\mathcal{Q}(\mu)\geq\mathcal{Q}(\hat{\pi})-2c\ell W_{2}(\mu,\hat{ \pi})-c^{2}|\operatorname{Tr}(\Sigma^{-1})-\operatorname{Tr}(\hat{\Sigma}^{-1} )|.$ $$

So $\mathcal{Q}(\mu)>0$ if

$$ $\displaystyle 2c\ell W_{2}(\mu,\hat{\pi})+c^{2}|\operatorname{Tr}(\Sigma^{-1}) -\operatorname{Tr}(\hat{\Sigma}^{-1})|<\mathcal{Q}(\hat{\pi})$ $$

or

$$ $\displaystyle 2\ell W_{2}(\mu,\hat{\pi})+c|\operatorname{Tr}(\Sigma^{-1})- \operatorname{Tr}(\hat{\Sigma}^{-1})|<(2-c)\operatorname{Tr}(\hat{\Sigma}^{-1}).$ $$

### A.3 Proof of Theorem 2

Recall from Lem. 1: for any $\mu=\mathcal{N}(m,\Sigma)\in\operatorname{BW}(\mathbb{R}^{d})$,

$$ $\displaystyle\mathcal{Q}(\mu)=2c\operatorname{Tr}(\mathbb{E}\nabla^{2}V(X))-c^ {2}\operatorname{Tr}(\Sigma^{-1}),\quad\text{where }X\sim\mu.$ $$

Since $V$ is $\alpha$-strongly convex, $\nabla^{2}V(x)\succcurlyeq\alpha I$ for all $x\in\mathbb{R}^{d}$. Therefore, $\mathbb{E}\nabla^{2}V(X)\succcurlyeq\alpha I$. It follows that $\operatorname{Tr}(\mathbb{E}\nabla^{2}V(X))\geq d\alpha$. Therefore, whenever $\operatorname{Tr}(\Sigma^{-1})<(2d\alpha)/c$, $\mathcal{Q}(\mu)>0$ and we get reduced variance.

### A.4 Proof of Theorem 3

Since Alg. 1 differs from SGVI (Diao et al., ) only at $\tilde{b}_{k}$, we will largely leverage the convergence analysis of (Diao et al., ) but will pay extra attention to the transition of the variance reduction effect to the final bounds.

At the iteration $k$, the (deterministic) Bures-Wasserstein gradient of $\mathcal{E}_{V}$ at $\mu_{k}$ is

$$ $\displaystyle\nabla_{\operatorname{BW}}\mathcal{E}_{V}(\mu_{k}):x\mapsto \mathbb{E}_{\mu_{k}}\nabla V+(\mathbb{E}_{\mu_{k}}\nabla^{2}V)(x-m_{k})$ $$

and in Alg. 1 we approximate this gradient by

$$ $\displaystyle x\mapsto\tilde{b}_{k}+S_{k}(x-m_{k})$ $$

where $\tilde{b}_{k}=\nabla V({X}_{k})-c_{k}\Sigma_{k}^{-1}({X}_{k}-m_{k})$, $S_{k}=\nabla^{2}V({X}_{k})$, and $X_{k}\sim\mu_{k}$.

The error of this approximation is

$$ $\displaystyle\tilde{e}_{k}:x\mapsto(S_{k}-\mathbb{E}_{\mu_{k}}\nabla^{2}V)(x-m _{k})+(\tilde{b}_{k}-\mathbb{E}_{\mu_{k}}\nabla V).$ $$

Let $\mathcal{P}_{k}$ denote $\sigma$-algebra containing the information up to the beginning of iteration $k$, $\mathcal{P}_{k}=\sigma(X_{0},X_{1},\ldots,X_{k-1})$
for $k\in\{1,2,\ldots,N-1\}$ and $\mathcal{P}_{0}$ is, by convention, the trivial $\sigma$-algebra. Let us denote

$$ $\displaystyle\tilde{\sigma}^{2}_{k}:=\mathbb{E}(\|\tilde{e}_{k}\|^{2}_{\mu_{k} }|\mathcal{P}_{k})=\mathbb{E}(\mathbb{E}_{x\sim\mu_{k}}\|(S_{k}-\mathbb{E}_{ \mu_{k}}\nabla^{2}V)(x-m_{k})+(\tilde{b}_{k}-\mathbb{E}_{\mu_{k}}\nabla V)\|^{ 2}|\mathcal{P}_{k}).$ (11) $$

Bounding $\tilde{\sigma}_{k}$: we show that

$$ $\displaystyle\tilde{\sigma}^{2}_{k}\leq 3d\beta(1+\tau_{k})+6(1+\tau_{k})\beta ^{3}\lambda_{\max}(\hat{\Sigma})W_{2}^{2}(\mu_{k},\hat{\pi}),$ (12) $$

The proof of (12) is a direct extension of (Diao et al., , Lem. 5.6), but let us partly include it here for completeness.

First, let $\mu=\mathcal{N}(m,\Sigma)$ and $X\sim\mu$, applying Stein’s lemma we get

$$ $\displaystyle\mathbb{E}\left(\dfrac{\partial V}{\partial x_{i}}(X)(X_{i}-m_{i} )\right)=\sum_{k=1}^{d}{\Sigma_{ik}\mathbb{E}\left(\dfrac{\partial^{2}V}{ \partial x_{k}\partial x_{i}}(X)\right)}.$ $$

Summing up for $i=1,2,\ldots,d$

$$ $\displaystyle\sum_{i=1}^{d}{\mathbb{E}\left(\dfrac{\partial V}{\partial x_{i}} (X)(X_{i}-m_{i})\right)}=\sum_{i=1}^{d}\sum_{k=1}^{d}{\Sigma_{ik}\mathbb{E} \left(\dfrac{\partial^{2}V}{\partial x_{k}\partial x_{i}}(X)\right)},$ $$

which can be rewritten as

$$ $\displaystyle\mathbb{E}\langle\nabla V(X),X-m\rangle=\mathbb{E}\langle\nabla^{ 2}V(X),\Sigma\rangle.$ $$

We now recall the Brascamp-Lieb inequality: let $\mu\propto\exp(-W)$ where $W$ is strictly convex and twice continuously differentiable, then

$$ $\displaystyle\mathrm{Var}_{\mu}(f)\leq\mathbb{E}_{\mu}\langle\nabla f,(\nabla^ {2}W)^{-1}\nabla f\rangle$ $$

for any smooth $f$. By using $f=(\partial/\partial x_{i})V$ and $\mu=\mu_{k}$, we obtain

$$ $\displaystyle\mathrm{Var}_{\mu_{k}}((\partial/\partial_{x_{i}})V)\leq\mathbb{E }_{\mu_{k}}[\nabla^{2}V\Sigma_{k}\nabla^{2}V]_{ii}.$ (13) $$

Summing (13) for $i$ from $1$ to $d$

$$ $\displaystyle\mathbb{E}_{\mu_{k}}\|\nabla V-\mathbb{E}_{\mu_{k}}\nabla V\|^{2} \leq\operatorname{Tr}\left(\mathbb{E}_{\mu_{k}}(\nabla^{2}V\Sigma_{k}\nabla^{2 }V)\right)=\mathbb{E}_{\mu_{k}}\langle\nabla^{2}V,\Sigma_{k}\nabla^{2}V\rangle.$ $$

Since $X_{k}$ is the only source of randomness in $\tilde{e}_{k}$ given $\mathcal{P}_{k}$, the conditional expectation in (11) becomes the expectation over the randomness of $X_{k}$, we can write

$$ $\displaystyle\tilde{\sigma}^{2}_{k}=\mathbb{E}\|(\nabla^{2}V(X_{k})-\mathbb{E} _{\mu_{k}}\nabla^{2}V)(X-m_{k})+\nabla V(X_{k})-c_{k}\Sigma_{k}^{-1}(X_{k}-m_{ k})-\mathbb{E}_{\mu_{k}}\nabla V\|^{2}$ $$

where $X,X_{k}\sim\mu_{k}$ and $X,X_{k}$ are independent. We evaluate

$$ $\displaystyle\dfrac{1}{2}\tilde{\sigma}^{2}_{k}$ $\displaystyle\leq\mathbb{E}\|(\nabla^{2}V(X_{k})-\mathbb{E}_{\mu_{k}}\nabla^{2 }V)(X-m_{k})\|^{2}+\mathbb{E}\|\nabla V(X_{k})-c_{k}\Sigma_{k}^{-1}(X_{k}-m_{k })-\mathbb{E}_{\mu_{k}}\nabla V\|^{2}$ $\displaystyle\leq\mathbb{E}((X-m_{k})^{\top}(\nabla^{2}V(X_{k})-\mathbb{E}_{ \mu_{k}}\nabla^{2}V)^{2}(X-m_{k}))+\tau_{k}\mathbb{E}_{\mu_{k}}\|\nabla V- \mathbb{E}_{\mu_{k}}\nabla V\|^{2}$ $\displaystyle=\mathbb{E}\langle(\nabla^{2}V(X_{k})-\mathbb{E}_{\mu_{k}}\nabla^ {2}V)^{2},(X-m_{k})(X-m_{k})^{\top}\rangle+\tau_{k}\mathbb{E}_{\mu_{k}}\| \nabla V-\mathbb{E}_{\mu_{k}}\nabla V\|^{2}$ $\displaystyle=\langle\mathbb{E}_{\mu_{k}}(\nabla^{2}V-\mathbb{E}_{\mu_{k}} \nabla^{2}V)^{2},\Sigma_{k}\rangle+\tau_{k}\mathbb{E}_{\mu_{k}}\|\nabla V- \mathbb{E}_{\mu_{k}}\nabla V\|^{2}$ $\displaystyle=\mathbb{E}_{\mu_{k}}\langle\nabla^{2}V,\Sigma_{k}\nabla^{2}V \rangle-\langle(\mathbb{E}_{\mu_{k}}\nabla^{2}V)^{2},\Sigma_{k}\rangle+\tau_{k }\mathbb{E}_{\mu_{k}}\|\nabla V-\mathbb{E}_{\mu_{k}}\nabla V\|^{2}$ $\displaystyle\leq\mathbb{E}_{\mu_{k}}\langle\nabla^{2}V,\Sigma_{k}\nabla^{2}V \rangle+\tau_{k}\mathbb{E}_{\mu_{k}}\|\nabla V-\mathbb{E}_{\mu_{k}}\nabla V\|^ {2}$ $\displaystyle\leq(1+\tau_{k})\mathbb{E}_{\mu_{k}}\langle\nabla^{2}V,\Sigma_{k} \nabla^{2}V\rangle$ $\displaystyle\leq\beta(1+\tau_{k})\mathbb{E}_{\mu_{k}}\langle\nabla^{2}V, \Sigma_{k}\rangle$ $\displaystyle=\beta(1+\tau_{k})\mathbb{E}\langle\nabla V(X_{k}),X_{k}-m_{k}\rangle.$ $$

Now by using optimal coupling between $\mu_{k}$ and $\hat{\pi}$, one can obtain (Diao et al., , P.27, P.28)

$$ $\displaystyle\mathbb{E}\langle\nabla V(X_{k}),X_{k}-m\rangle\leq\dfrac{3d}{2}+ \left(2\beta+\dfrac{\beta^{2}\operatorname{Tr}(\hat{\Sigma})}{d}\right)W_{2}^{ 2}(\mu_{k},\hat{\pi})$ $$

Therefore,

$$ $\displaystyle\tilde{\sigma}^{2}_{k}$ $\displaystyle\leq 3d\beta(1+\tau_{k})+(1+\tau_{k})\left(4\beta^{2}+\dfrac{2 \beta^{3}\operatorname{Tr}(\hat{\Sigma})}{d}\right)W_{2}^{2}(\mu_{k},\hat{\pi})$ $\displaystyle\leq 3d\beta(1+\tau_{k})+6(1+\tau_{k})\beta^{3}\lambda_{\max}( \hat{\Sigma})W_{2}^{2}(\mu_{k},\hat{\pi}).$ $$

Bound $\mathbb{E}(\min_{k=\overline{1,N}}\mathcal{F}(\mu_{k}))-\mathcal{F}(\hat{\pi})$:

Lem. 5.1 in (Diao et al., ) implies that

$$ $\displaystyle\mathbb{E}W_{2}^{2}(\mu_{k+1},\hat{\pi})\leq(1-\alpha\eta)\mathbb {E}W_{2}^{2}(\mu_{k},\hat{\pi})-2\eta(\mathbb{E}\mathcal{F}(\mu_{k+1})- \mathcal{F}(\hat{\pi}))+2\eta^{2}\mathbb{E}\tilde{\sigma}_{k}^{2}$ (14) $$

where $\alpha\geq 0$ is the strong convexity modulus of $V$.

Now using the bound (12) for $\tilde{\sigma}_{k}$,

$$ $\displaystyle\mathbb{E}W_{2}^{2}(\mu_{k+1},\hat{\pi})$ $\displaystyle\leq(1-\alpha\eta+12(1+\tau_{k})\eta^{2}\beta^{3}\lambda_{\max}( \hat{\Sigma}))\mathbb{E}W_{2}^{2}(\mu_{k},\hat{\pi})$ $\displaystyle\quad-2\eta(\mathbb{E}\mathcal{F}(\mu_{k+1})-\mathcal{F}(\hat{\pi }))+6(1+\tau_{k})\eta^{2}\beta d$ $\displaystyle\leq\exp\left(-\alpha\eta+12(1+\tau_{k})\eta^{2}\beta^{3}\lambda_ {\max}(\hat{\Sigma})\right)\mathbb{E}W_{2}^{2}(\mu_{k},\hat{\pi})$ $\displaystyle\quad-2\eta(\mathbb{E}\mathcal{F}(\mu_{k+1})-\mathcal{F}(\hat{\pi }))+6(1+\tau_{k})\eta^{2}\beta d.$ $$

Therefore

$$ $\displaystyle 2\eta(\mathbb{E}\mathcal{F}(\mu_{k+1})-\mathcal{F}(\hat{\pi}))\leq$ $\displaystyle\exp\left(-\alpha\eta+12(1+\tau_{k})\eta^{2}\beta^{3}\lambda_{ \max}(\hat{\Sigma})\right)\mathbb{E}W_{2}^{2}(\mu_{k},\hat{\pi})$ $\displaystyle-\mathbb{E}W_{2}^{2}(\mu_{k+1},\hat{\pi})+6(1+\tau_{k})\eta^{2}\beta d$ (15) $$

Since we are considering the convex case, let us set $\alpha=0$ and denote $C_{k}=12(1+\tau_{k})\beta^{3}\lambda_{\max}(\hat{\Sigma})$ and $D_{-1}=0,D_{k}=-C_{0}-C_{1}-\ldots-C_{k}$ for $k=0,1,\ldots,N-1$. By definition, $D_{k}+C_{k}=D_{k-1}$ for all $k=0,1,\ldots,N-1$. Rewrite (A.4) as

$$ $\displaystyle 2\eta(\mathbb{E}\mathcal{F}(\mu_{k+1})-\mathcal{F}(\hat{\pi})) \leq\exp\left(C_{k}\eta^{2}\right)\mathbb{E}W_{2}^{2}(\mu_{k},\hat{\pi})- \mathbb{E}W_{2}^{2}(\mu_{k+1},\hat{\pi})+6(1+\tau_{k})\eta^{2}\beta d.$ $$

Multiply both sides with $\exp(D_{k}\eta^{2})$ we get

$$ $\displaystyle 2\eta\exp(D_{k}\eta^{2})(\mathbb{E}\mathcal{F}(\mu_{k+1})- \mathcal{F}(\hat{\pi}))$ $\displaystyle\leq\exp\left((D_{k}+C_{k})\eta^{2}\right)\mathbb{E}W_{2}^{2}(\mu _{k},\hat{\pi})-\exp(D_{k}\eta^{2})\mathbb{E}W_{2}^{2}(\mu_{k+1},\hat{\pi})+6( 1+\tau_{k})\eta^{2}\beta d\exp(D_{k}\eta^{2})$ $$

and, by using the backward recursion $D_{k}+C_{k}=D_{k-1}$, can be rewritten as

$$ $\displaystyle 2\eta\exp(D_{k}\eta^{2})(\mathbb{E}\mathcal{F}(\mu_{k+1})- \mathcal{F}(\hat{\pi}))$ $\displaystyle\leq\exp\left(D_{k-1}\eta^{2}\right)\mathbb{E}W_{2}^{2}(\mu_{k}, \hat{\pi})-\exp(D_{k}\eta^{2})\mathbb{E}W_{2}^{2}(\mu_{k+1},\hat{\pi})+6(1+ \tau_{k})\eta^{2}\beta d\exp(D_{k}\eta^{2})$ $$

Telescope for $k$ from 0 0 to $N-1$

$$ $\displaystyle 2\eta\sum_{k=0}^{N-1}{\exp(D_{k}\eta^{2})(\mathbb{E}\mathcal{F}( \mu_{k+1})-\mathcal{F}(\hat{\pi}))}$ $\displaystyle\leq W_{2}^{2}(\mu_{0},\hat{\pi})-\exp(D_{N-1}\eta^{2})\mathbb{E} W_{2}^{2}(\mu_{N},\hat{\pi})+6\eta^{2}\beta d\sum_{k=0}^{N-1}{(1+\tau_{k})\exp (D_{k}\eta^{2})}$ $\displaystyle\leq W_{2}^{2}(\mu_{0},\hat{\pi})+6\eta^{2}\beta d\sum_{k=0}^{N-1 }{(1+\tau_{k})\exp(D_{k}\eta^{2})}.$ $$

We see that

$$ $\displaystyle D_{k}=-\dfrac{C}{2}\left(k+1+\sum_{i=0}^{k}{\tau_{i}}\right)$ $$

where $C=24\beta^{3}\lambda_{\max}(\hat{\Sigma})$.

Let us denote $\tilde{S}_{N}(\eta)=\sum_{k=0}^{N-1}{\exp(D_{k}\eta^{2})}$. It holds

$$ $\displaystyle\mathbb{E}\left(\min_{k=1,2,\ldots,N}\mathcal{F}(\mu_{k})\right)- \mathcal{F}(\hat{\pi})\leq\dfrac{W_{2}^{2}(\mu_{0},\hat{\pi})}{2\eta\tilde{S}_ {N}(\eta)}+3\eta\beta d\sum_{k=0}^{N-1}{(1+\tau_{k})\dfrac{\exp(D_{k}\eta^{2}) }{\tilde{S}_{N}(\eta)}}.$ $$

It holds

$$ $\displaystyle\sum_{k=0}^{N-1}{(1+\tau_{k})\dfrac{\exp(D_{k}\eta^{2})}{\tilde{S }_{N}(\eta)}}\leq 1+\tau_{\max}$ (16) $$

and

$$ $\displaystyle\tilde{S}_{N}(\eta)$ $\displaystyle=\sum_{k=0}^{N-1}{\exp(D_{k}\eta^{2})}$ $\displaystyle=\sum_{k=0}^{N-1}{\exp\left(-\dfrac{C}{2}(k+1+\sum_{i=0}^{k}{\tau _{i}})\eta^{2}\right)}$ $\displaystyle\geq\sum_{k=0}^{N-1}{\exp\left(-\dfrac{C}{2}(k+1+(k+1)\tau_{\max} )\eta^{2}\right)}$ $\displaystyle=\sum_{k=0}^{N-1}{\left[\exp\left(-C(k+1)\eta^{2}\right)\right]^{ \frac{\tau_{\max}+1}{2}}}.$ $$

On the other hand, for any $b>0$, the function $f(s)=b^{s}$ is convex. By tangent inequality $f(s)\geq f(1)+f^{\prime}(1)(s-1)$, we get

$$ $\displaystyle b^{s}\geq b+b\ln(b)(s-1).$ (17) $$

Applying the inequality (17) with $b=\exp\left(-C(k+1)\eta^{2}\right)$ and $s=(\tau_{\max}+1)/2$

$$ $\displaystyle\left[\exp\left(-C(k+1)\eta^{2}\right)\right]^{\frac{\tau_{\max}+ 1}{2}}$ $\displaystyle\geq\exp\left(-C(k+1)\eta^{2}\right)+C(k+1)\eta^{2}\exp\left(-C(k +1)\eta^{2}\right)\left(\dfrac{1-\tau_{\max}}{2}\right)$ $\displaystyle=\exp\left(-C(k+1)\eta^{2}\right)\left(1+C\eta^{2}(k+1)\left( \dfrac{1-\tau_{\max}}{2}\right)\right)$ $\displaystyle\geq\exp\left(-C(k+1)\eta^{2}\right)\left(1+C\eta^{2}\left(\dfrac {1-\tau_{\max}}{2}\right)\right).$ $$

Therefore,

$$ $\displaystyle\tilde{S}_{N}(\eta)$ $\displaystyle\geq\left(1+\dfrac{C\eta^{2}(1-\tau_{\max})}{2}\right)\sum_{k=1}^ {N}{\exp(-Ck\eta^{2})}$ $\displaystyle\geq\left(1+\dfrac{C\eta^{2}(1-\tau_{\max})}{2}\right)\sum_{k=1}^ {\min\{N,\lfloor(C\eta^{2})^{-1}\rfloor\}}{\exp(-Ck\eta^{2})}$ $\displaystyle\geq\left(1+\dfrac{C\eta^{2}(1-\tau_{\max})}{2}\right)\sum_{k=1}^ {\min\{N,\lfloor(C\eta^{2})^{-1}\rfloor\}}{\dfrac{1}{e}}$ $\displaystyle=\dfrac{1}{e}\left(1+\dfrac{C\eta^{2}(1-\tau_{\max})}{2}\right) \min\{N,\lfloor(C\eta^{2})^{-1}\rfloor\}.$ $$

By using the basic inequality $1/\min(a,b)\leq 1/a+1/b$, we get

$$ $\displaystyle\dfrac{1}{\tilde{S}_{N}(\eta)}$ $\displaystyle\leq\dfrac{e}{1+\dfrac{C\eta^{2}(1-\tau_{\max})}{2}}\left(\dfrac{ 1}{N}+\dfrac{1}{\lfloor(C\eta^{2})^{-1}\rfloor}\right)$ $\displaystyle\lesssim\dfrac{e}{1+\dfrac{C\eta^{2}(1-\tau_{\max})}{2}}\left( \dfrac{1}{N}+C\eta^{2}\right)$ $$

asymptotically at small $\eta>0$.

Therefore,

$$ $\displaystyle\mathbb{E}\left(\min_{k=1,2,\ldots,N}\mathcal{F}(\mu_{k})\right)- \mathcal{F}(\hat{\pi})\leq\dfrac{e}{1+\dfrac{C\eta^{2}(1-\tau_{\max})}{2}} \left(\dfrac{1}{2\eta N}+\dfrac{C\eta}{2}\right)W_{2}^{2}(\mu_{0},\hat{\pi})+3 \eta\beta d(1+\tau_{\max}).$ $$

### A.5 Proof of Theorem 4

Since $V$ is $\alpha$-strongly convex with $\alpha>0$, $\mathbb{E}_{\hat{\pi}}(\nabla^{2}V)\succcurlyeq\alpha I$, so $\lambda_{\min}(\mathbb{E}_{\hat{\pi}}(\nabla^{2}V))\geq\alpha.$

It follows that

$$ $\displaystyle\lambda_{\max}(\hat{\Sigma})=\dfrac{1}{\lambda_{\min}(\hat{\Sigma }^{-1})}=\dfrac{1}{\lambda_{\min}(\mathbb{E}_{\hat{\pi}}(\nabla^{2}V))}\leq \dfrac{1}{\alpha}.$ $$

Using this inequality in the bound for $\tilde{\sigma}_{k}$ in (12), we get

$$ $\displaystyle\tilde{\sigma}_{k}^{2}\leq 3d\beta(1+\tau_{k})+\dfrac{6(1+\tau_{k })\beta^{3}}{\alpha}W_{2}^{2}(\mu_{k},\hat{\pi}).$ $$

Using this bound for (14),

$$ $\displaystyle\mathbb{E}W_{2}^{2}(\mu_{k+1},\hat{\pi})$ $\displaystyle\leq(1-\alpha\eta)\mathbb{E}W_{2}^{2}(\mu_{k},\hat{\pi})-2\eta( \mathbb{E}\mathcal{F}(\mu_{k+1})-\mathcal{F}(\hat{\pi}))$ $\displaystyle\quad+2\eta^{2}\mathbb{E}\left(3d\beta(1+\tau_{k})+\dfrac{6(1+ \tau_{k})\beta^{3}}{\alpha}W_{2}^{2}(\mu_{k},\hat{\pi})\right)$ $\displaystyle=\left(1-\alpha\eta+\dfrac{12(1+\tau_{k})\eta^{2}\beta^{3}}{ \alpha}\right)\mathbb{E}W_{2}^{2}(\mu_{k},\hat{\pi})+6d\beta\eta^{2}(1+\tau_{k})$ $\displaystyle\leq\exp\left(-\alpha\eta+\dfrac{12(1+\tau_{k})\eta^{2}\beta^{3}} {\alpha}\right)\mathbb{E}W_{2}^{2}(\mu_{k},\hat{\pi})+6d\beta\eta^{2}(1+\tau_{ k}).$ $$

Now with $\eta\leq\alpha^{2}/(48\beta^{3})$,

$$ $\displaystyle\dfrac{12(1+\tau_{k})\eta^{2}\beta^{3}}{\alpha}\leq\dfrac{(1+\tau _{k})\eta\alpha}{4}.$ $$

Therefore

$$ $\displaystyle\mathbb{E}W_{2}^{2}(\mu_{k+1},\hat{\pi})$ $\displaystyle\leq\exp\left(\left(\dfrac{-3+\tau_{k}}{4}\right)\eta\alpha\right )\mathbb{E}W_{2}^{2}(\mu_{k},\hat{\pi})+6d\beta\eta^{2}(1+\tau_{k})$ $\displaystyle\leq\exp\left(\left(\dfrac{-3+\tau_{\max}}{4}\right)\eta\alpha \right)\mathbb{E}W_{2}^{2}(\mu_{k},\hat{\pi})+6d\beta\eta^{2}(1+\tau_{\max}).$ $$

Telescope this inequality, we get

$$ $\displaystyle\mathbb{E}W_{2}^{2}(\mu_{N},\hat{\pi})$ $\displaystyle\leq\exp\left(-N\left(\dfrac{3-\tau_{\max}}{4}\right)\eta\alpha \right)W_{2}^{2}(\mu_{0},\hat{\pi})+\dfrac{6(1+\tau_{\max})\eta^{2}\beta d}{1- \exp\left(-\dfrac{(3-\tau_{\max})\eta\alpha}{4}\right)}$ $\displaystyle\lesssim\exp\left(-\dfrac{N(3-\tau_{\max})\eta\alpha}{4}\right)W_ {2}^{2}(\mu_{0},\hat{\pi})+\dfrac{24(1+\tau_{\max})\beta\eta d}{(3-\tau_{\max} )\alpha}$ $$

asymptotically at small $\eta>0$.

## Appendix B Additional experimental details

### B.1 Laplace Approximation

Laplace approximation fits a Gaussian approximation by finding the mode of the target (MAP estimate for infernece) and forming a second order approximation at that point. The approximation is given by

$$ $\mathcal{N}\left(x_{\text{MAP}},\left(\nabla^{2}V(x_{\text{MAP}})\right)^{-1} \right).$ $$

We use BFGS optimizer (Nocedal & Wright, ) as implemented in SciPy (Virtanen et al., ) to find the (numerical) MAP estimate, and form the approximation according to the local curvature around the point. Convergence of the estimate was validated manually.

### B.2 Variational Inference in the Euclidean Geometry

The baseline method EVI optimizes for the approximation over its parameters $m$ and $\Sigma$ in the Euclidean geometry of the parameter space, using Cholesky factorization for parameterizing the covariance.
This is done by maximizing the Evidence Lower BOund (ELBO)

$$ $\mathcal{L}(m,\Sigma)=\mathbb{E}_{q_{m,\Sigma}(z)}\left[\log p(x,z)-\log q_{m, \Sigma}(z)\right],$ $$

which is equivalent to minimizing the KL divergence. We use single-sample reparameterization estimates for the gradient.
Furthermore, by stopping the gradient after sampling $z$, we remove the Fisher score from the gradient computation, giving an unbiased estimator of the gradient of the ELBO with potentially lower variance (Roeder et al., ).
We use Adam (Kingma & Ba, ) optimizer and the learning rates and number of iterations found in Table 1, found to achieve good convergence. Our implementation is based on the code provided by Modi et al. ().

**Table 1: Optimization details for EVI.**
| Experiment | Dimension | Learning Rate | Iterations |
| --- | --- | --- | --- |
| Gaussian | 10 | 0.01 | 5,000 |
| Gaussian | 50 | 0.01 | 5,000 |
| Gaussian | 200 | 0.001 | 10,000 |
| Student-t | 200 | 0.001 | 8,000 |
| Logistic Regression | 200 | 0.01 | 3,000 |

### B.3 Student’s t distribution

Consider a $d$-dimensional Student-t distribution with location $\mu$, scale matrix $\Sigma$ and $\nu$ degrees of freedom. Its negative log density (up to a constant), gradient and Hessian are given by:

$$ $\displaystyle V(x)$ $\displaystyle=\frac{1}{2}\left(\nu+d\right)\log\left(1+\frac{1}{\nu}(x-\mu)^{ \top}\Sigma^{-1}(x-\mu)\right),$ $\displaystyle\nabla V(x)$ $\displaystyle=\frac{(\nu+d)}{\nu+(x-\mu)^{\top}\Sigma^{-1}(x-\mu)}\Sigma^{-1}( x-\mu),$ $\displaystyle\nabla^{2}V(x)$ $\displaystyle=\frac{\nu+d}{\nu+(x-\mu)^{\top}\Sigma^{-1}(x-\mu)}\Sigma^{-1}- \frac{2(\nu+d)}{(\nu+(x-\mu)^{\top}\Sigma^{-1}(x-\mu))^{2}}\Sigma^{-1}(x-\mu)( x-\mu)^{\top}\Sigma^{-1}.$ $$

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 8
presentation: 4
questions: As mentioned above, in my opinion the paper is very well-written. I therefore don’t have any content-related questions, only a small number of typographical suggestions for the authors:
- Perhaps use the en-dash (--) instead of the simple dash (-) in "Bures–Wasserstein". I believe this is the standard when it comes to joining two names.
- I believe that "FB" abbreviation (forward-backward) is never introduced.
- The term “L-smooth” seemed weird to me, why not “L-Lipschitz”?
- Page 4, line 174, “affine map” -> “affine maps”.
- Page 4, line 198, “shows that” -> “showed that”.
- Page 5, line 233, I would remove parentheses around the “also see ...”, they look ugly immediately after a citep.
- Page 5, line 237, perhaps add commas around “where $c \in \mathbb{R}$”, otherwise the sentence is hard to read.
- Perhaps, when introducing Algorithm 1, you could briefly mention there are recipes for defining the constants $c_k$ that you investigate further in the text. An algorithm without parameters you don’t know how to set would come across as more useful.
- Page 8, line 387, “We also note that conditioning” -> “We also note that by conditioning”.
- Page 8, footnote 1, please expand what is the minor correction. Also, perhaps start with a capital letter.
- Page 9, line 457, missing “e” in “covariancs”.
- Page 9, line 472, “In The” -> “In the”.

Note: no need to respond to the small items above in the rebuttal, save yourself some time.
soundness: 4
strengths: - The paper is written very well. It is a pleasure to read.
- The theoretical analysis is convincing (I didn’t rigorously check the proofs but I don’t doubt the claims and I am pretty sure the methods authors use are quite appropriate).
- Empirical evaluation shows very strong improvement over the baselines, more so than the theoretical analysis (which is not claimed to be tight) suggests.
summary: The paper proposes an improved optimization technique for variational inference with a multivariate Gaussian variational family.

Building on the prior work that proposed optimization on the Bures–Wasserstein manifold and simple one-sample stochastic gradient estimators, this work offers improved stochastic gradient estimators. These are based on a simple idea: as iterates get closer to the optimum, they provide additional information about the next gradient that can be used to reduce variance of its stochastic estimator.

Authors support the superiority of the new optimization technique with both theoretical analysis and empirical evaluation, the latter on a set of synthetic benchmarks.
weaknesses: -I’m not an expert in optimization, so it’s somewhat hard for me to judge the impact. I can speculate that the problem might seem somewhat narrow: variational inference with _Gaussian_ variational family. However, I don’t really think so myself. Furthermore, in my opinion, a good solution to even a somewhat narrow problem definitely warrants publication.

I would actually give this paper a score of 9, but only 8 and 10 are available.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: Based on the concerns summarized in the weaknesses section, I would like to pose the following questions, categorized into “theoretical” and “experimental” aspects:

### Questions on Theoretical Analyses:
1.	Could the authors provide a more detailed discussion on the behavior of each term in the derived bound and how the bound behaves as the iteration $N$ approaches infinity, particularly in terms of convergence and conditions for convergence?
2.	In the convex case, variance reduction is guaranteed only near the optimal solution (Thm. 1), while the convergence analysis in Thm. 3 assumes variance reduction is achieved “in all iterations.” Does this assumption imply that parameters are initially close to the optimal solution, and if so, how realistic is this assumption?
3.	The improvement in the bound relies on the assumption that variance reduction is guaranteed for all iterations. How does this assumption align with the results of Thm. 1? Is the radius $r$ of the region always small enough to ensure consistency between the two theorems? If not, there may be a significant gap between these results.

### Questions on Experimental Analyses:
1.	The numerical experiments are primarily conducted on synthetic data. Could the authors consider conducting experiments on real-world datasets to provide more practical insights into the usefulness of the proposed method?
2.	There are no experiments measuring the reduction in gradient variance. Could the authors include empirical measurements of gradient variance to illustrate how much variance is reduced compared to existing methods? The following references may be relevant for this consideration: [Miller et al. (2017); Buckholz et al. (2018)].
3.	In the proposed method, the degree of variance reduction appears to depend on the parameter $c$. Could the authors provide additional practical discussions and sensitivity analyses on the selection of $c$ and its impact on variance reduction performance, particularly in terms of gradient variance? Although Figure 2 presents the relationship between $c$ and the degree of variance reduction, it would be beneficial to confirm whether this behavior is consistent across various experimental settings.
4.	The comparison with the method of Roeder et al. (2017) in Euclidean SVI is understandable; however, could the authors clarify the rationale for not including comparisons with more recent variance-reduction methods?
soundness: 2
strengths: - The authors extend the control variate method—a well-established variance reduction technique in stochastic variational inference (VI)—to variational inference in Bures-Wasserstein space.
- The variance reduction performance of the proposed method is theoretically guaranteed in (strongly) convex settings.
- The authors demonstrate improved convergence over existing SGVI and Bures-Wasserstein gradient descent (BWGD) methods through several numerical experiments.
summary: - The authors propose a control-variate-based variance reduction technique for gradient estimation in stochastic Gaussian variational inference (SGVI).
- They theoretically demonstrate that, in specific cases such as high-curvature distributions, the gradient estimator with their method achieves lower variance than the standard SGVI.
- Additionally, they show that their variance reduction scheme enhances the convergence of SGVI when $V(x)$ is (strongly) convex.
- Empirically, they confirmed that their method improves the convergence properties of SGVI.
weaknesses: I would like to express my sincere respect for the efforts the authors have invested in this paper. However, I am unable to strongly recommend it for acceptance at ICLR 2025 for the following reasons:

### Concerns Regarding Theoretical Analyses:
- The evaluation of the derived bound is primarily qualitative and limited to a comparison with the existing upper bounds in [Diao et al., 2023]. However, detailed analysis of each term and the behavior of the bound as the iteration count $N$ approaches infinity—specifically in terms of convergence and the conditions required for convergence—is somewhat lacking. This limits a comprehensive understanding of the proposed method’s advantages and challenges the assessment of this paper’s contribution.
- In the convex case, the variance reduction effect is guaranteed only in the neighborhood of the optimal solution (Thm. 1). However, the convergence analysis in Thm. 3 assumes that variance reduction holds “throughout all iterations of the algorithm,” which seems to suggest that parameters are already near the optimal solution from the outset—a somewhat unrealistic assumption (or am I misinterpreting this?).
- Additionally, from the proof, it appears that the improvement in the bound is only valid under the aforementioned assumption: “when variance reduction is assured for all iterations.” This assumption seems to diverge from the result in Thm. 1. Is there a clear link between them? Is the radius $r$ of this region always sufficiently small?

### Concerns Regarding Numerical Experiments:
- One key motivation for variance reduction in stochastic variational inference is to enhance convergence and generalization in real-world data applications (e.g., [Ruiz et al. (2016); Roeder et al. (2017); Miller et al. (2017); Buckholz et al. (2018)]). However, the numerical experiments in this paper are primarily conducted on synthetic data, and the discussion of the practical applicability of the proposed method is somewhat limited.
- Additionally, there are no direct measurements of gradient variance, leaving insufficient empirical evidence regarding the extent of gradient variance reduction compared to existing methods. If the empirically measured gradient variance does not differ substantially from that of other methods, it is possible that other factors are driving the observed improvement in convergence. Furthermore, based on the results presented, the effectiveness of gradient variance reduction appears to depend on the setting of $c$. Practical guidance on choosing $c$, based on an analysis of the sensitivity of variance-reduction performance to $c$, would be beneficial. Currently, such guidance is absent, leaving limited information on how to apply the proposed method in practice.
- When comparing with variance-reduction methods in SVI in Euclidean space (EVI), this study adopts the method of Roeder et al. (2017); however, the rationale for this choice is unclear. While both methods share a similarity in achieving performance improvements through modest algorithmic modifications, concluding that the proposed method outperforms EVI without comparison to more recent approaches seems premature (e.g., [Miller et al. (2017); Buckholz et al. (2018)]).

### Lack of Discussion on Other Variance Reduction Studies in EVI:
- Recent years have seen substantial progress in variance reduction research within EVI, with many theoretical analyses and methods proposed (e.g., [Kim et al. (2023); Domke et al. (2023)]). These approaches extend beyond control variate methods to include enhancements to Monte Carlo methods and techniques such as reparameterization tricks in Gaussian settings, as in Roeder et al. (2017). While the use of control variates as an initial approach to variance reduction in SGVI is reasonable, the motivation for extending only the control variate approach may seem somewhat limited given the variety of variance reduction techniques available. It might be helpful to mention, for example, the specific challenges or limitations that prevent the application of alternative variance reduction techniques in SGVI, if applicable.

### Citation:
- Ruiz et al. (2016): Francisco J. R. Ruiz, Michalis K. Titsias, and David M. Blei. The Generalized Reparameterization Gradient. NeurIPS2016. https://arxiv.org/abs/1610.02287.
- Miller et al. (2017): Andrew C. Miller, Nicholas J. Foti, Alexander D'Amour, and Ryan P. Adams. Reducing Reparameterization Gradient Variance. NeurIPS2017. https://arxiv.org/abs/1705.07880
- Buckholz et al. (2018): A. Buchholz, F. Wenzel, and S. Mandt. Quasi-Monte Carlo Variational Inference. ICML2018. https://arxiv.org/abs/1807.01604.
- Kim et al. (2023): K. Kim, K. Wu, J. Oh, and J. R. Gardner. Practical and Matching Gradient Variance Bounds for Black-Box Variational Bayesian Inference. ICML2023. https://proceedings.mlr.press/v202/kim23w.html.
- Domke et al. (2023): J. Domke, R. Gower, and G. Garrigos. Provable convergence guarantees for black-box variational inference. NeurIPS2023. https://proceedings.neurips.cc/paper_files/paper/2023/hash/d0bcff6425bbf850ec87d5327a965db9-Abstract-Conference.html.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: +  Are there any obstacles in applying this to more realistic statistical inference problems?  I really do appreciate the mathematical elegance  of the proposed method but t would be nice to have a better test problem  to make the case for the method.

+ Any head-to-head comparisons with different inferences methods such that MCM, Hamiltonian MC, or likelihood free inferences (e.g. using diffusions models) on a test problem would be helpful to convince the reader of the efficiency of the method., 

+  The improvement due to control variates looks quite impressive (as it  sometimes does with control variates!).  The convergence analysis  is nice but it is not clear it shed light on this improvement. Any other explanation, even the hand-waving kind?
soundness: 4
strengths: I found the paper to be very overall very well-written with a very clear presentation of the background, literature, and of the technical ideas involved.  The results and the proofs are clearly presented.  Since the measure involved are Gaussian, the analysis of the control variate is quite straightforward and this lead to an improved algorithm at essentially no additional cost.  I think that it is the main contribution of the paper to point out this very cheap improvement to the algorithm. The empirical results also supports this improvement of the stability of the algorithm leading to a substantially better approximation.  At the technical level, control variates for Gaussian measures are quite straightforward and the analysis of the variance reduction relies on standard techniques as well (Stein's lemma).  The convergence of the algorithms is then analyzed by leveraging the results recently obtained Diao et a. (ICML 23) and propagating the improved variance through the convergence analysis of Diao et al.  This leads to improved convergence results.
summary: The paper propose improvements on a gaussian variational inference algorithm in the Bures-Wasserstein manifold.  The idea is to find the best Gaussian approximation (in the KL sense) to a probability distribution of the form $Ce^{-V}$ (where $C$ is unknonw) and to use the Wasserstein geometry on the space of Gaussian distribution (Bures-Wasserstein manifold) to guide us toward the optimal solution.  Several authors have studied this question in the last few year leading to an implicit-explicit algorithm.  In this framework the implicit backward (JKO-style) step is computable explicitly and the explicit forward step involves the computation of the the expectations of the gradient and  the Laplacian of $V$ under gaussian measures.  The central contribution of the paper is to perform a variance reduction (using control variates)  at each step and to study the effect of this on the rate of convergence of the algorithm.  Some synthetic examples are studied.
weaknesses: I think the main weakness is the lack of interesting examples. The algorithms seems to do great on the synthetic examples but it remains to be seen how it will perform in a realistic  Bayesian inference example based on real data.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: Please see the weakness section.
soundness: 2
strengths: The paper is well structured, clearly presenting the problem and the proposed solution. The main idea of using a control variate to reduce the variance of BW gradient is both well-motivated and reasonable. Additionally, it provides strong theoretical foundations for the derived variance-reduced Gaussian VI. The experimental results further validate the effectiveness of the proposed method.
summary: This paper considers the variational inference (VI) problem by minimizing the Kullback-Leibler (KL) divergence over the Bures-Wasserstei space of Gaussian distributions. Building on the work by Diao et al., this paper employs the forward-backward Euler method to tackle the composite nature of the KL objective and non-smoothness of the entropy. The forward step involves computing the Bures-Wasserstein (BW) gradient, which needs to be estimated using Monte Carlo (MC) method with a single sample. However, this approach results in noisy BW gradients, leading to inefficient convergence. The main contribution of this paper is to overcome this limitation, proposing a variance-reduced estimator for the BW gradient, leveraging a control variate approach. In addition, this paper also provides convergence analysis and present experimental results that demonstrate the effectiveness of the proposed method.
weaknesses: I have two major concerns regarding the proposed method in this paper:

1) A straightforward approach to addressing the noisy BW gradient is to sample multiple $X_{k}$ from $\mu_{k}$. This is both effective and efficient, requiring low extra computational cost if estimating $\nabla V$ and $\nabla^{2}V$ is not expensive. I am curious how the performance of the proposed method compares to this trivial approach.

2) The paper only considers the case where $\mu_{k}$ is Gaussian, and the proposed method appears to work well for simple target distributions. Note that both compared methods (BWGD and SGVI) can be extended for Gauss mixtures, while extending the proposed method in this paper to cases of Gaussian mixtures may not trivial. In such cases, the proposed control variate approach might lose its effectiveness and efficiency, as $E[Z_{k}]$ is not zero (the mean) and estimating it is non-trivial.

Meta-review: The proposed work claims that recent variational inference methods based on stochastic implementations of the Bures-Wasserstein gradient flow can be improved by an order-of-magnitude. This is achieved by the principle of control variates. 

The strength of the work is the theoretical analysis which gives an improved convergence bound for BW variational inference. The simplicity of the presented method and its good numerical performance is another strength. All reviewers found the paper well written and appreciated the overall presentation.

The main weaknesses are the lack of larger scale numerical experiments, and lack of comparison and acknowledgment of existing works on control variates in VI or even simpler techniques like momentum. Moreover, there are works on variational inference in different geometries (Fisher-Rao) which are arguably more natural for Bayesian inference than the Wasserstein geometry (which is close to the Euclidean geometry). These weaknesses were mostly addressed by the authors in the rebuttal phase. 

Overall I recommend acceptance of the work. The simplicity of the approach, its promising numerical performance as well as the theoretical results will be of interest to the community. I recommend the authors to take the feedback of the reviewers into account for the final version of the manuscript.

Decision: Accept (Poster)