Stochastic variance-reduced Gaussian variational inference on the Bures-Wasserstein manifold
Authors: Hoang Phuc Hau Luu
Hanlin Yu
Bernardo Williams
Marcelo Hartmann
Arto Klami
Venue: ICLR 2025 Poster (2025)
Keywords: Variational inference
Bures-Wasserstein space
Riemannian manifold
optimization
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)

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 not 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 demonstrate that the proposed estimator gains order-of-magnitude improvements over the previous Bures-Wasserstein methods.

Published as a conference paper at ICLR 2025
STOCHASTIC VARIANCE-REDUCH GAUSSIAN VARI- ATIONAL INFERENCE ON THE BURES-WASSERSTEN
MANIFOLD
Hoang Phuc Hau Luu, Hanlin Yu, Bernardo Williams, Marcelo Hartmann, Arto Klami
Department of Computer Science
University of Helsinki, Helsinki, Finland (hoang-phuc-hau. lua, hanl in. yu, bernardo. wi 11 ammmorrno, ma marcelo. harhaartr sarttann, arto. klam Jehelsinki.f5
ABSTRACT
Optimization in the warssessssnnttiin space has been gaining popularity in the machine learning community since it draws connections between variational infer- ence 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-backwalt 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 exact since the Warssssnasnrttein gradient of the potential energy involves "intractable" expec- tations. 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 ----nnitude improvements over previous Bures Wassersetn methods.
INTRODUCTION
Variational inference (VI) (Wainwright et al. 2008 Blei et al. 2017) provides a fast and scalable al- ternative to Markov chain Monte Carlo (MCMC), especially for inference tasks in high dimensions. The main principle of VI is to approximate a complicated distribution π,, e.g., posterior distribution in Bayesian inference, by simpler tractable family of distributions. The approximation l 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 in- fluenced by the assumptions made on the approximation family and the information about that can be obtained, ranging from classical coordinate ascent algorithms for mean-field approxima- tions for targets with conditional conjugacy structure (Blei et al.) (017) to gradient methods using score-function approximations to avoid assumptions on the target density (Ranganath et al. 2014)
or flexible approximations parameterized with neural networks (Rezende & Mohamed 2015
We focus on Gaussian approximations (Honkela Valpola 2004 Opper Archambeau 2009 Xu
& Campbell 2022 Quiroz et al. (023) but with a particular emphasis on the recent research line
the Wasserstein geometric viewpoint of this family (Lambert et al. 2022, Diao al. 2023). Re- garding the target
assume access to second order gradients, typically computed by automatic differentiation, similar to the above works. Gaussian VI offers strong statistical guarantces at the optimal solution (Katsevich & Rigollet 2024), offers easy way of modelling dependencies be- tween the variables and, thanks to the Bernstein- von Mises theorem (Van der Vaart 2000), 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 opti- mization perspective (Lambert et al. 2022 Diao et al. 2023 The family of non-degenerate Gaus-

---

Published as a conference paper at ICLR 2025
SVRIGI (ours)
SGMI
BWGD
16
107
SGVI
BWGD
10a
SVRGVI (ours)
10a
10
0
10-1
100
200
300
400
500
600
-10
10
-20
-10
32
30
-20
30
time elapsed (r *iteess
Figure 1: Left: Optimization trajectories of our method compared to SGVI (Diao et al. 2023 and BWGD (Lambert et al.) 2022). The target 50-dimensional Gaussian distribution, visualized here via the marginal distributions of the first two coordinates. Each ellipse represents contour of Gaussian: the black 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 orders of magnitude more accurate.
sian distributions can be parameterized by its mean and covariance matrix, μ4 with 0 (m, E),
henceforth denoted as θ = Rd x S where S is the set of d x symmetric, positive definite matrices. Classical VI employs conventional optimization algorithms (Paisley et al. 2012 Titsias &
Lizaro-Gredilla 2014 Kucukelbir et al. 2017) to minimize the Kullback-Leibler (KL) divergence
(kI( over the parameter space & equipped with the Euclidean geometry. Lambert et al.
(022) argue that because the optimization problem is over the space distributions, it is more natural
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-Waar Wasserstein (BW) geometry that benefits optimization. Lambert al. (2022) subsequently established a theoret- ical framework for performing VI using the BW geometry, which we adopt in this paper.
Let x (x) x exp(-V(z)) be the target distribution and consider the VI problem
â E arg min (KI(UT))T #EBW (R4)
where BW (Rd) is the Bures- Wasserstenn space of Gaussian distributions with non-degenerate variance matrix. The BW space 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 (T) can be reformulated as
arg min F(n), where (EBW(R")
F(u) = Ev - H (H)).
Here, (14 f (x))du((= is the potential energy and (l²) f log(n(z))ddp(s) is the neg- ative entropy. A conceptual and established idea to minimize a functional F to perform gradi- ent flow on F with respect to the geometry of (WR).. To be implementable, the flow must be discretized, Lambert el al. (2022) use forward Euler discretization, resulting in a scheme named Bures-Wassersterr stochastic gradient descent (BWGD).
Diao et al. (2023) remark that forward-backward (FB) Euler (Bauschke & Combettes) 2011) should be used instead due to the objective's composite nature and the entropy's 's non-smoothness. This method iteratively applies forward step to the potential energy εy and a backward step (proximal operator) to the negative entropy H. They also observe that the backward step the BW space has a closed-form solution [Wibisono)] 2018). This is crucial because this step is known to be intractable (or computationally expensive) in the full Wasserstein space [Wibisono] 2018 Salim et al. 2020) Mokrov et al. 2021 Luu et al.) 2024). 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-Waasserser Wasserstein gradient of Ey instead of the "friendly" Wasserstein gradient that is just VV. The Bures- Wasserseen gradient is not always available in closed form, i.e., at j BW(Rd) it is given only implicitly by the map + F. VV (E,, V2V)(x my) where my
2

---

Published as a conference paper at ICLR 2025
is the mean of j (Lambert al. 2022). This is the orthogonal projection of the Wasserstein gradient onto a tangent space of the Bures Wasserstein manifold (Chewi et al. 2024). For general V. these expectations are intractable even though the underlying distribution is a Gaussian. Diao et al. (2023) proposed using the Monte Carlo (MC) method with one sample to estimate these expectations at each iteration: sample X jl and use (V(X) and (XV(X) as unbiased estimators for JE VV and E 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, shown in experi- ments (Sect. (5) In practice, high-variance estimators require small step sizes, leading to slow and inefficient convergence. We resolve this fundamental limitation by proposing variance-reduced timator with minimal computational overhead while providing robust theoretical guarantees. Fig. [] shows the improvement over SGVI and BWGD in practice. meaningful ful transition from the initial distribution to the target distribution, and our method follows the path smoothly and is particularly stable around the optimum.
gares-Wassersteinn geometry offers
Contributions. We propose novel variance- estimator for E, VV that does not use any extra samples, with minimal per-iteration computational overhead, using the control variates ap- proach (Owen) 2013). Our idea is that the variational distribution should be similar to the target distribution т (xx) x exp(-V/ (x)) as μ gets closer and closer to म, so the density of y can be used to construct a correlated control variate for the Monte-Carlo estimator (V(X). Sect. 3 presents the detailed construction and its rationale.
On the theoretical side, we derive the following insights:
Thm. Under a mild smoothness assumption, we prove that there is a region around the optimal solution á where estimator has guaranteed smaller variance than the MC estimator.
Thm. If V is strongly convex. we prove that the proposed estimator has a smaller variance than the MC estimator at every l E BW(R)) whenever y has sufficiently large (greater than a controllable threshold) variance.
We further show in Thm. 3 and Thm. that whenever variance reduction happens along the algo- rithm's 's iterates, the effect will enter the convergence analysis and improve the optimization bounds derived in Diao et al. 2023). These theorems solidly back our proposed method.
On the practical side, 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 minimal modification to the Monte Carlo estimator, the proposed estimator achieves significant improvements in our experiments.
BACKGROUND
function f RRd - called L-smooth (or Lipschitz smooth) for all v, y E Rd. If is twice continuously differentiable, define the Laplacian operator of f as Aʃ Dim (0 (0²/Ox?)f. For a random variable r. || == inf[C: | s C almost surely).
S L| yl
2.1 BURES- WASSERSTEIN GEOMETRY
We denote by Pa d the space of probability measures over Hd with finite [[ < +oo. Equipped with the Wasserstein distance
i.c.,
inf / TEI(a.)) xxx ||-- ||1-yy|?dyaa drl2dy(a,))
(H.
(3)
where ('u,, v) the set of probability measures over X x X whose marginals are u and v, the space P2 (2Rd) becomes the metric space called the Wasserstein space (Ambrosio et al. 2005). We call E *[4,, v) a (transport) plan and any that achieves the optimal value in (3) an optimal plan. A pair of random variables whose joint distribution is optimal plan is called optimal coupling (between μl and v/). When μ is absolutely continuous with respect to the Lebesgue measure, Brenier theorem (Brenier 1991) asserts that the optimal plan is unique and is given by (T, TW) # where
3

---

Published as a conference paper at ICLR 2025
TØ Vg for some convex function g. We call Ty the optimal transport map from μ to v. Apart from being a metric space, the Wasserstein space also enjoys some nice properties of Riemannian geometry. Otto's calculus (Otto, 2001) endows the Wasserstein space with formal Riemannian structure, facilitating gradient flows and optimization.
We denote by BW (R4) the space of Gaussian distributions with non-degenerate covariance ma- trices. The Wasserstein distance between two Gaussian distributions Po = (m.. Σg) and P1 = (Mmi, E1) is given in the closed-form formula W (poo.P) pI) = =I|ma0 m, 112 + B (Co. Σ
where B² Σn) Tr(En + Σ - -2 ) is the Bures metric. The optimal transport map
is also given in closed form in this case: T (ar) m1 + Eo E. E2
En
(a: mo).
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 = (m,, 2) E BW(R)). the tangent space of (W(R)) at the space of symmetric affine maps denoted as Tú (BRd)) R E { where 8d is the space of symmetric d x d matrices. The Riemannian metric defined using the inner product of elements in this tangent space is identified as the (2) inner product restricted this space. Given U, V T, (x."). the metric is (U, V) := (Z),,, V(z))du(zx. This Riemannian metric induces the geodesic distance in (W(R") that given by the Wasserstein distance. We refer to Altschuler et al. (202I) for further discussions on BW geometry.
{a + S(x m) + t a
2.2 STOCHASTIC GAUSSIAN VI
We refer to Diao et al. (023) for a detailed discussion and relevant terminologies. We briefly explain
stochastic Gaussian VI to motivate our proposed variance reduction. Recall from that aim to
minimize F(u)= X (n) + (v) over BW (Rd). At the optimum of F, _ N m. Σ). first-order
optimality reads (Opper & Archambeau 2009 Lambert et al. 2022 Diao et al. 2023
E3 VV and Eg V2v D-1
(4)
which is derived by zeroing the Bures Wasserstein gradient of the objective function.
A natural idea minimize over (W to perform gradient flow on F using the BW ge- ometry of PW(IR)). When the gradient flow is applied over the entire Wasserstein space (2( corresponds to the Langevin diffusion (Jordan et al. 1998), with one of its discretizations being an MCMC method called the unadjusted Langevin algorithm (Roberts & Tweedie, 1996). When restricted to BW(RR)) the gradient flow can be formulated using Riemannian geometry (Do Carmo 1992,. as BW (RR) forms true Riemannian manifold. This flow is curve of Gaussian distribu- tions, characterized by the time-dependent evolution of their mean and covariance matrix. Recently, Lambert et al. (2022) showed that this evolution is governed by Särkkä's ODEs developed in the context of variational Kalman filtering Särkka 2007).
The negative entropy H is convex along generalized geodesics but it is nonsmooth functional. If V is smooth, it induces the smoothness of Ey. Therefore, it is natural to apply Euler that alternates between two steps: iteration
Pk+1 (1 17 BV ( 1
forward step
{ (1))
- W2 (L, (M.Me+:))
1
10++1
arg min
< backward step
(CBW(Rd)
27
where Vnw denotes the mal step in the optimization literature or the JKO (Jordan, Kinderlehrer, and Otto) step (with restric- tion in BW (R))) in the context of Wasserstein gradient flow (Jordan et al. 1998). The backward step is intractable in the full Wasserstein space and hence requires (oftentimes expensive) numer- ical approximations (Mokrov et al.) 2021 Luu et al. 2024). On the other hand, if restricted to BWR).. this step admits closed-form solution (Wibisono) 2018): let Px+ -N(m++ 3. Ex+z), then 14+1 is Gaussian distribution with mean m*+1 m*+4 and variance matrix Σx+1
gradient. The backward step is also known as the proxi-
(> + 2m/ + [2 (E++) + 4n/)]²) This tractability of the backward is the main motiva-
tion Diao al. (2023) to study and develop FB Fuler in this scenario. The forward step, however,
4

---

Published as a conference paper at ICLR 2025
is not always analytically available since the BW gradient of εv, at iterate k,
Bv (pss) + Ey VV + (Bye V²V m
(023) propose using Monte Carlo approximation for involves intractable expectations. Diao et al. these expectations: sample Xx Mi and use by
VV(x)) and S (Xxx)) unbiased
estimators for Epg VV and Eps V2V, respectively.
3 STOCHASTIC VARIANCE-REDUCEE 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,
(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 Eg,, VV or Eyi, V in the sense that their variances are smaller than those of by and Sk. building on the control variates approach (wen 2013;; Also see the discussions in Defazio et al. (2014): Luu 2022).
Bug VV by VV(+)) and Epp V²V = S == (2V())) where X jag-
Let us first describe briefly the core idea of control variates in helping reduce the variance. Let be the quantity of interest and X be an unbiased estimator for 0, i.e., E. 9. A control variate random variable with a known mean so that Y is correlated with X. The random variable Z = X + c(EY Y), where E I. is then unbiased estimator for 0. The variance of Z is
VarZ VarX + VarY - Yx,).
If X, Y are highly correlated in the sense that 2Cov Y) VarY, we immediately get Var < Var for any с E (0.1), 1]. So, we achieve a reduction in variance by using Z. On the other hand, if X, Y are correlated (Cov(X, Y) 0) but not highly correlated,
also obtain variance reduction effects whenever e is positive and small enough. Furthermore, given the parabolic form with respect to in (6), one can pinpoint the optimal value of is c* Cov( (X, Y))Var(Y), resulting in the maximal variance reduction VarZ = (1 Corr(X, Y) < Var where Corr(X, Y) denotes correlation between X and Y.
We now return to our problem and seek variance-reduced estimators of the forms
b (Vx))) + c(E(Zz) - Zx) and S₁ v²V(X,) + d(E(Wx) Wk)
where d > O and Zz W are random vector and a random matrix, respectively. Let us first focus on by- As discussed, Z should be (element- elementtwwsse) highly correlated with VV(X,) while E(Z)) remains efficiently computable. We look for Z₁ (U(X;) so that VU as close to VV as possible. We are in the context of approximating т (xx) by the VI distribution Ak = Am,,, Ex), so it is natural to expect that
-V((x) Vog x
(x) = log fa; Ex) -E ' my).
where f(x; Σx) & exp (-2 mR) Σ (x - (x)m is the PDF of μ4a Therefore, we pro- pose using Z Σx (Xx m2 control variate. We have E(Zx) 0 since E(X²) my. is worth noting that Zz is known as the Stein/Hyvärinen score (Hyvärinen 2005) of (4u- The estimator - m2x) By applying the same reasoning to St, we can immediately conclude that Wù is deterministic and equals >n² Consequently, the control variate does not affect Sμ; we keep the standard estimator. We derive Stochastic seriaaccereduuced Gaussian VI (SVRGVI) as in Alg. (we will discuss more about the choice of Cą in Sect. 4. Note that the only difference between Alg. 1 and the SGVI in Diao et al. (0023) is the estimator by. where the difference highlighted in blue.
by then becomes by == (V((x)
Fig. 2 (left) demonstrates that our proposed estimator (with c = 0.9) achieves lower variance com- pared to the standard MC estimator, while both remain unbiased estimators of E, VV. Fig. 2 (right), we vary from 0 to 2 and calculate the empirical variance of our estimator, revealing parabolic pattern. Note that when e = O, the estimator reduces to the standard estiator, and for all values of e E (0, 2), proposed estimator consistently exhibits lower variance, with an optimal value around 1. At this optimal c, the variance is reduced roughly by factor 10. We provide theoretical justification for these empirical observations in Sect.
5

---

Published as a conference paper at ICLR 2025
Nonte Cario gradients
1.75
Verwnce reduction gradients
1.50
Emorecat vertence
1.25
2.00
6.75
2-
0.25
6.00
0.25 n5o
0.15 1.00 1.25 1.50 2.00
Figure 2: Left: π is a Gaussian, VI distribution il is in the neighborhood of æ. In this case, the true gradient, i.e., the expectation E, VV, 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 j and substituting the values into the respective estimator formulas. Right: The empirical variance of our proposed estimator when e varies from 2. Note that corresponds the Monte Carlo estimator.
Minimal extra computational cost Despite involving calculating the inverse of the covariance matrix, the computational overhead is small. Sampling from multivariate normal step in Alg. typically requires obtaining the Cholesky factor of the covariance matrix, which is (dd) (Ras- mussen & Williams) 2006 With the Cholesky factor, obtaining the solution of the inverse of the matrix times a vector (d2)) (Rasmussen & Williams 2006). As such, we can reuse this obtained Cholesky factor in step I to compute the inverse in step 2. which implies that the estimator adds an overhead of (df), which naturally dominated by the (dd²) complexity of the original algorithm.
Algorithm Stochastic variance- Gaussian Variational Inference (SVRGVI)
Input: Target distribution (2) x exp(-V(z)), initial distribution Hi W[ Eu), step size 7) > 0, number of steps N, sequence of control variate parameters {cu (c=)==) 1N-1 where oN
(0, 1] vk E (0, 1, 1..... 1)
for to N - do 1. Draw one sample X ~M([,, Ex)
'XXx 2. Compute estimators: by VV())) cg Eg (Xx my) and S₁ + (2V(X.)
3. Update mean and covariance matrix:
mk+1 < Mc 17bx
MR+1 1 S/
Dx+3 + MR+1 Mx+1
Σ6④1 Ex+ 1 + 2nd +
Ex++ (Ex+ + 4n/1]
end for
Output: nN = N(mN, En)
THEORY
In Sect. 3. we argued that, in the context of variational inference, as μ4 iteratively gets closer n, (V(X,) then (highly) correlated to I-1(Xx - my), and hence we obtain a variance reduction effect. This argument leads to the construction of the control variate in Alg. [ One might question whether this approach remains effective when the target distribution マ is significantly distant from the BW space. Because we are constrained to the BW space, the best can do is to get closer to which is the optimal solution to the problem [1. However, â might still look very different from rigorously show that within certain neighbourhood of # (to be defined
Notably, in Thm. later), our proposed estimator consistently reduces variance, regardless of how different π is to a Gaussian distribution. Let first introduce Lem. to pave the way for Thm. and also to discuss

---

Published as a conference paper at ICLR 2025
the optimal in the control variate. In Lem. compute the variance of the proposed estimator by leveraging multidimensional Stein's 's lemma (Lin al. 2019).
Lemma Assume that V is continuously differentiable. Lei pa - (mm, 2) BW(R"). Then,
E
- c2' (X - EVV
variance of our estimator
VV)V [XII
Tr(E 2c TT( where X J..
variance of the Monte-Carke estimator
extraa term
Proof of Lem. Tis given in Appendix A. Lem. compares the variance of the proposed estimator BW (IR4). Recall that the first-order optimality condition (4p of स reads E& V Σ-1 Consequently, at ft, the extra term in Lem. | is simplified
and the Monte Carlo estimator at given
a(c- 2) T which negative whenever e (0, 2) and minimized for n. our estimator is always better than the Monte Carlo estimator for e (0,
Therefore,
Remark practical merit of Lem. iis that it implies the optimal value for to get maximum variance reduction at μ c² [r(EE V2V) TT(=-= Applying this Alg. ] we can pick the
Tr(E, v2V) Tr(S2) = Cg. Cg. Again, this computation of incurs implies the
adaptive sequence {ck} } as ct
Tr(s:
Tr(E.=)
negligible extra cost to Alg. [
We also remark that around ât optimality condition 4
optimal value c* indeed is around 1.
In Thm. we further show that when the Laplacian AV is smooth, the proposed estimator has a smaller variance than the Monte Carlo estimator region around म.
Theorem (Variance reduction around the optimal solution) Assume that the Laplacian AV is E-moooth For any control variate coefficient E (0,2), define the region around मे N E.):
V(a, r) Ep (m, Σ) 28W2(u, - <r]
where r = -2-c) Tr(É-1) > 0 is the region's 's radius. For any μ E VVftr)), r), the proposed estimator has smaller variance than the Monte Carlo estimator.
Proof of Thm. I is given in Appendix A.2 with the main idea being that the smoothness of the Laplacian AV propagates the improvement of the proposed estimator at ã to its neighbourhood. We additionally observe that, for small >> the region V(nrr) r) effectively reduces to the Wasserstein ball B(t, (r(S-
Thm. applies to arbitrary x, only requiring mild smoothness condition of its second derivative. In the next theorem, we show that when π is strongly log-coneave is now more similar Gaussian), variance reduction happens not only around á but also in many regions of interest.
Theorem (Variance reduction at large-variance distributions) If V is a-strongly convex for some a for any control variate e >> 0, the proposed estimator has a smaller variance than the Monte Carlo estimator at every Nmm, Σ) whenever Tr(E-1) < 2nd
Proof of Thm. is given in Appendix
Remark A consequence of Thm. is that for strongly convex V we can obtain variance reduction each iteration in Alg. by setting ek (0, 2od/ Tr(E ())() )). We can extend this result to the case of convex but mon-strongly-cowee V, i.e. a = 0. We first show that: for any Gaussian (mm, E), Tr(EN V2V) O. We can then obtain variance reduction each iteration by setting e E (0, (T(E VVV/Tr(Ex))) Note that this upper bound of Ck is less explicit as the strongly convex case. However, the point that variance reduction still happens in this case given that Cs is small enough. See Appendix A.3 for the proof.
We further show in Thm. and Thm. that whenever variance reduction happens along the al- gorithm's 's iterates, the effect will propagate the convergence analyses of Diao et al. (2023) and

---

Published as a conference paper at ICLR 2025
improve their theoretical bounds. Therefore, combining with Thm. and Thm. the overall theory strongly favours SVRGVI over SGVI.
Let Pa denote the information up to the beginning of iteration k, ie., it is the o-algebra given by Pk =
for E {1,2...., - 1} and Po is, by convention, the trivial o-- algebra with information. Assuming variance reduction occurs along the algorithm's iterates, holds
ie., for = 0, 0,1,...,N
(IVV(()) c (X, m²) - Ep VVI & T (Xx) - En vVII
(7)
where 73 1 almost surely. We recall from Remark 2 that for convex almost surely for all & if is well-chosen, and it follows that ET, < and Txlloo S for all k. We
holds such that 7
can further guarantee |x < for all under mild boundedness assumptions (Appendix AA.4.
Under condition 7 we now show the improved bounds. Similar to Diao et al. (2023). we con- sider log-concave and -roonllylog--concae sampling, meaning that is assumed to convex and strongly convex, respectively.
Theorem (Convex case) Suppose that V is convex and B-smooth and the step size O < 29" variance reduction happens for k 0, 1,..... then,
1
Cn
min F(ux) F(t) 2 1=122......
W 22o #) 3n3d(1 Tmax,
1 [PP(1-mmax,es
2
where Tmaxx max - TTll bor Tmax,E ; ET, e 2.718 the Euler's num- ber, C' = 24 Amaxx }), and S is asymptotically the limit of small n7.
Proof of Thm. is given in Appendix A.5
Theorem (Strongly convex case) Suppose that V is a-strooglyycoovve with a and 0 < S 4833 If variance reduction happens for NN 1. then
N(3 Tmax, oo ma
24(1 Tmax, F 3nd
(WN, (UN, R) exp
W (No: +) +
(8)
(3 Tmax oc )x
where Tm and F are defined in Theorem
and S is asymptotically at the limit of small n,.
Proof of Thm. 4 is given in Appendix[A.o
Remark 3 We recall the corresponding bounds for SGVI in Diao et al. 2023 Thm 5.7, Thm. 5.8)
Convex. E (mink=1.2.... F(ux)) F(+) S
(Ho
+ eCn W2 #) + 63nd
2Nn
Strongly convex. EW3 S exp
aNn
W (no. #) 248nd
Putting side-by-side, we see that Thm. and Thm. 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 Tmax.co = 0, the noise terms in the bounds of Thm. 3 and Thm. woolld not disappear because of another source of randomness coming from St.
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. 2022) and SGVI (Diao et al.) 2023
We set the step size for all algorithms (for an experiment on varying step size, see
With a minor correction to the coefficients in SGVT's bound.
Our code is available at intttttpp// / cow hub/vr25
8

---

Published as a conference paper at ICLR 2025
Iii
Pallud
Pylusfii
10%
101
107
10-1
10'
102
10
BINGD
10
BWGD
SCVI
SCVI
10-1
SDVI
10
SVRCVI (ours)
10
SVRGVV lours!
10-2
SVASVI (ours)
0
50
1000
150
200
250
300
50
100
150
200
250
300
50
100
150
260
250
300
time clapsed (nx #iters)
time elapsed (nx #iters)
time elapsed (g #iters)
(a) d 10
(b) d 50
(c) d 200
Figure KL divergence for Gaussian targets varying dimensionality.
Appendix B.7). fix the covariate coefficient = 0.9 (for an experiment on the effect of c, see Appendix B.6), 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 full-rank Gaussian approximation optimised the Euclidean geometry (denoted as EVI), using low-variance reparameterization gradients of Roeder et al. (2017) with ADAM optimizer, and Laplace approximation that does not optimize the KL divergence but fits Gaussian distribution cost of these methods is different from the BW methods, we only report the final accuracy for carefully optimized approximations show how the BW methods compare against commonly used algorithms. The Laplace approximation is omitted for Gaussian targets as it is optimal by definition.
the target mode; see Appendix [ for details. As the per-iteration
Gaussian targets We randomly generate the means and covariances for multivariate Gaussian target distribution `,, considering dimensions of {10, 50, 200). Fig. demonstrates consistent sig- nificant 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 Fig. [ shows visually the marginals for d = 50, providing an interpretation of the improvement seen diveeiver.. We also clearly outperform EVI in higher dimensions, unlike previous BW methods.
Student's targets
We consider a multivariate Student's target with a degree of freedom of in (a) shows that our algorithm is again clearly the best. BWGD is not stable and, on average, performs worse than even the Laplace approximation.
200 dimensions. Fig. 4
Bayesian logistic regression We consider a Bayesian logistic regression with flat prior as in Diao et al. (2023): given set of covariates X; ~ (0, La) for i = 1, 1.2....,74 consider
Y| Bernuuli( ((0 X))), where is the sigmoid function.
The negative log posterior is V(Ø) = =" [In11 (..X..) (0, X+)]- The model consists of 1000 data points (X;, Y2) with dimension d 200. 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 F(ue), by drawing samples from M4k- We denote by Abest the distribution that obtains the smallest F among all iterations of all algorithms, comparing against that. Fig 4 (b) shows the proposed method is again the most accurate.
We also measured the variance along iterations in Appendix B.4 showing that we achieved variance reduction empirically, We showed Appendix B.5 that the proposed method is more efficient than minibatching (whether using i.i.d. or quasi Monte-Carlo Samlle))
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 semi- nal work of Roeder et al. (2017) that demonstrated how the variance of gradient estimators for VI can be dramatically reduced by single-line change in the algorithm: We also propose minor

---

Published as a conference paper at ICLR 2025
Ou lempiricall
F(ue) Milaaa) (emplrical)
107
BWGD
BWGD
109
SGVI
10a
SGVI
SVAGVI lours)
SVAGVI (ours]
104
Laplace
Laplace
10
10a
10a
20a
100
150
200
250
300
0
100
150
200
250
time elapsed En Miters)
time elapsed (r) Witers)
(a) Student's
(b) Bayesian logistic regression
Figure 4: Performance of algorithms for Student's target and Bayesian logistic regression.
modification that dramatically improves the accuracy, and should always be used. A high-level sim- ilarity lies in the heuristic that the VI distribution resembles the target distribution, allowing to be used to construct control variates. Perhaps surprisingly, Yi & Liu (0233 recently showed that the continuous-time flow of Roeder et al. (2017) 7) for Gaussian VI the Euclidean geometry (by repa- rameterizing the covariance matrix as > 5S also follows Särkkä's ODEs. In other words, those BW methods and Roeder et al. (2017) are different discretizations of the same dynamics. Recent works (Kim et al. Domke al. 2024) have advanced understanding of variance reduction in Euclidean VI, demonstrating strong convergence and extending beyond traditional control variate approaches. Adopting these new techniques in the BW setting is promising research direction.
In the BW space, Diao (2023) considers variance reduction for large-sum structures based on the nested-loop idea by Johnson & Zhang (2013) to reduce the stochasticity of the minibatch sampling. In contrast, our method addresses the stochasticity arising from Gaussian sampling from the VI.
Even though our experiments focused on synthetic targets, they expanded the previous exper- imentation of VI optimized in the BW space. We confirm the finding of Diao et al.] (2023) that BWGD and SGVI are effectively identical except for the instability of the former, but now show how their performance degrades in higher dimensions while our algorithm remains effective.
With the exception of the vastly improved accuracy due to the significantly lower variance of the gradient estimators, our method retains all qualitative characteristies 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 (2022) and Quiroz al. (2023), 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 op- timization of the approximating distribution in the Busss-Wasserstein space of Gaussians can be made practical. The previous works by Lambert et al. (2022) and Diao et al. (2023) introduced the key idea and the algorithms with strong theoretical guarantees. However, they are prone to high vari- ance from Monte Carlo approximations, 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 and showed that this reduction results 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 variance in the case of strong log-concave targets, and hence the proposed variance reduction technique should always be used.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work is supported by the Research Council of Finland's Flagship programme: Finnish Cen- ter for Artificial Intelligence (FCAI), and additionally by grants 363317, 345811 and 348952. We thank Kai Puolamäki for the discussions. The authors wish acknowledge CSC IT Center for Sci- ence, Finland, for computational resources. We thank the anonymous reviewers for their insightful comments and suggestions.
REFERENCES
Jason Altschuler, Sinho Chewi, Patrik R Gerber, and Austin Stromme. Averaging on the Bures- Wasserstein manifold: dimension-free convergence of gradient descent. Advances in Neural In- formation Processing Systems, 14-22132-221455 2021.
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2005.
Heinz Bauschke and Patrick Combettes. Convex analysis and monotone operator theory in Hilbert spaces, 2011.
David Blei, Alp Kucukelbir, and Jon D McAuliffe, Variational inference: review for statisti- cians. Journal of the American statistical Association, (12(518):859-877) 2017.
Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Commu- nications on pure and applied mathematics, 44(4):375-417, 1991.
Alexander Buchholz, Florian Wenzel, and Stephan Mandt. Quasi-monte carlo variational inference. In International Conference on Machine Learning, pp. 668-677. PMLR, 2018.
Sinho Chewi, Jonathan Jiles-Weed, and Philippe Rigollet. preprint arX1v:2407.18163, 2024.
Statistical optimal transport. arXiv
Aaron Defazio, Francis Bach, and Simon Lacoste-Julien. SAGA: A fast incremental gradient method with support for non-strongly convex composite objectives. Advances in neural infor- mation processing systems, 27, 2014.
Michael Ziyang Diao. Proximal gradient algorithms for Gaussian variational inference: Optimiza- tion in the wassersnsnnn space. Master's thesis, Massachusetts Institute of Technology, 2023.
Michael Ziyang Diao, Krishna Balasubramanian, Sinho Chewi, and Adil Salim. Gaussian variational inference via jko in the Bures-We Wasserstein space. In International Confer- ence on Machine Learning, pp. 7960-7991. PMLR, 2023.
Manfredo Perdigao Do Carmo. Riemannian geometry, volume 2. Birkhäuser, 1992.
Justin Domke, Robert Gower, and Guillaume Garrigos. Provable convergence guarantees for black-
box variational inference. Advances in neural information processing systems, 36, 2024.
Antti Honkela and Harri Valpola. Unsupervised variational Bayesian learning of nonlinear models. Advances in neural information processing systems, 17, 2004.
Aapo Hyvärinen. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
Rie Johnson and Tong Zhang. Accelerating stochastic gradient descent using predictive variance reduction. Advances in neural information processing systems, 26, 2013.
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.
Anya Katsevich and Philippe Rigollet. On the approximation accuracy of Gaussian variational inference. The Annals of Statisfics,
2024.

---

Published as a conference paper at ICLR 2025
Kyurae Kim, Kaiwen Wu, Jisu Oh, and Jacob R Gardner. Practical and matching gradient variance bounds for black-box variational bayesian inference. International Conference on Machine Learning. pp. 16853-16876. PMLR, 2023.
Diederik P. Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. In 3rd Inter- national Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Truck Proceedings, 2015.
Alp Kucukelbir, Dustin Tran, Rajesh Ranganath, Andrew Gelman, and David M Blei. Automatic differentiation variational inference. Journal of machine learning research, 18(14)11-45, 2017.
Mare Lambert, Sinho Chewi, Francis Bach, Silvère Bonnabel, and Philippe Rigollet. Variational inference via Wasserstein gradient flows. Advances in Neural Information Processing Systems, 111444444447, 2022.
Wu Lin, Mohammad Emtiyaz Khan, and Mark Schmidt. Stein's lemma for the reparameterizatioo trick with exponential family mixtures. arXiv preprint arXiv: /910. 13398, 2019.
Hoang Phuc Hau Luu. Advanced machine learning techniques based on DCA and applications to predictive muintenance. PhD thesis, Université de Lorraine, 2022.
Hoang Phuc Hau Luu, Hanlin Yu, Bernardo Williams, Petrus Mikkola, Marcelo Hartmann, Kai Puolamäki, and Arto Klami. non-geodesicalll--cones optimization in the Wasserstein space. Advances in Neural Information Processing Systems, 37:16772-16809, 2024.
Chirag Modi, Robert Gower, Charles Margossian, Yuling Yao, David Blei, and Lawrence Saul. Variational inference with Gaussian score matching. Advances in Neural Information Processing Systems, 36, 2024.
Petr Mokrov, Alexander Korotin, Lingxiao Li, Aude Genevay, Justin M Solomon, and Evgeny Bur- naev. Large-scale Wasserstein gradient flows. Advances in Neural Information Processing Sys- fems, 34:15243-15256 2021.
Jorge Nocedal and Stephen J. Wright. Numerical Optimization. Spinger, 2 edition, 2006.
Manfred Opper and Cédric Archambeau. The variational Gaussian approximation revisited. Neural computation, 21(3):786-792, 2009.
Felix Otto. The geometry of dissipative evolution equations: The porous medium equation. Com- munications Partial Differential Equations, 16(1-2):101-174, 2001.
Art Owen. Monte Carlo theory, methods and examples, 2013.
John Paisley, David M. Blei, and Michael I. Jordan. Variational Bayesian inference with stochas- tic search. In Proceedings of the 29th International Coference on International Conference on Machine Learning, pp. 3363-1370, 2012.
Matias Quiroz, David Nott, and Robert Kohn. Gaussian variational approximations for high- dimensional state space models. Bayesian Analysis, 18(3):989-1016, 2023.
Rajesh Ranganath, Sean Gerrish, and David Blei. Black box variational inference. In Artificial intelligence and statistics, pp. 814-822. PMLR, 2014.
Carl Edward Rasmussen and Christopher K. I. Williams. Gaussian Processes for Muchine Learning. Adaptive computation and machine learning. The MIT Press, Cambridge, MA, USA, 2006.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In Interna- tional conference on machine learning, pp. 1530-1538, PMLR, 2015.
Gareth O. Roberts and Richard L. Tweedie. Exponential convergence of Langevin distributions and their discrete approximations. Bernoulli, 2(4):341 363, 1996.
Geoffrey Roeder, Yuhuai Wu, and David K Duvenaud. Sticking the landing: Simple, lower-varianee gradient estimators for variational inference. Advances in Neural Information Processing Systems, 30, 2017.
12

---

Published as a conference paper at ICLR 2025
Adil Salim, Anna Korba, and Giulia Luise. The Wasserstein proximal gradient algorithm. Advances in Neural Information Processing Systems, 33:12336-1123335 2020.
Simo Särkkä. On unscented Kalman filtering for state estimation continuous- nonlinear systems. IEEE Transactions automatic control, 52(9):1631-144) 2007.
Michalis Titsias and Miguel Sizaro-Gredilla Doubly stochastic variational Bayes for non-conjugate inference. In International conference on machine learning, pp. 1971-1979, PMLR, 2014.
Aad W Van der Vaart. Asymptotic statistics, volume 3. Cambridge University Press, 2000.
Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Courna- peau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jomathan Bright, Stéfan van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nel- son, Eric Jones, Robert Kern, Eric Larson, Carey, Ilhan Polat, Yu Feng, Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, lan Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, Antônio H. Ribeiro, Fabian Pedregosa, Paul van Mul- bregt, and SciPy 1.0 Contributors. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods, 17:261-272 2020.
Martin Wainwright, Michael I Jordan, et al. Graphical models, exponential families, and variational inference. Foundations and Trends@ in Machine Learning. 1(1-2):1-305, 2008.
Andre Wibisono. Sampling as optimization in the space of measures: The Langevin dynamics as composite optimization problem. In Conference Learning Theory, pp. 2093-3027. PMLR, 2018.
Zuheng Xu and Trevor Campbell. The computational asymptotics of Gaussian variational inference and the laplace approximation. Statistics and Computing, 32(4):63, 2022.
Mingxuan Yi and Song Liu. Bridging the gap between variational inference and Wasserstein gradi- ent flows. arXiv preprint arXiv:2310.20090, 2023.
A THEORY
A.1 PROOF OF LEMMA
For each (m, Σ) E (Rd) and c > 0, we denote
EVV ((w) = E - VV)) - cx'
m) - EVV(XII? X~
which is the difference between the variances of the Monte Carlo estimator and our proposed esti- mator. We want Q(μ)) > O. Simple algebras simplify Q as
(p)) (VV()
E- -m) - CEEE-- m)]l?
Recall standard result: if X ~ N(m, Σ). then its affine transformation W = AX + has the distribution. (Am-+ AΣAT). Applying this result, W == = NO..E-1 Therefore
E| = Σ EW? T)
On the other hand,
E (VV( - - -EV(X))) - m))
-XVV m)) E(EVV 'x - m))
-E(VV(X), m)) (EVV(X), XX, EE- (X - m))
-I(VVV(X),- m)) E -m))
EE(VV(X), -m))
13

---

Published as a conference paper at ICLR 2025
Let us denote A Σ and compute (X), mm) as follows
0v (x)[A(X- -m)ls
(XV - -m)) =
0V (X) Σ
E
(Ali(X) - -m3)
Ozs
j=1
[A
OVV
(X) mm;)
(9)
Or;
(OV (V (x)(X;-1)) - m;)) by leveraging the following Stein's 's lemma (Lin al. 2019
We compute E dz,
Lemma (Stein's lemma) Let (m, Σ) an d-dimensional Gaussian random variable and g R + R be continuously disserentiable, then
E (g (()x)-m)) (V)))()))
Applying Stein's lemma with g = (8/0x;VV,
2v Or. (X) m))
OV
E
EE
V
(x)
Oz:
((
82V
0CV Ox Dr.dx: (x) (x0] )
ΣE
(X):
(X),
Üxzdr
By comparing the j-th element of both sides, we get
E
@V
(X)
-m)))
Σ
82V Orxdda
(x))
Oz,
k=1
Plugging this expression into (9)
d d E(VV(X), m)) Σ
8V DxyQa, (x)
[A] li
E5
j=1
02V axyōxx (x)
[A EynB
OºV
E
drpdx, (x)
>
[A Esn
j=1
02V
Orpdx (x)
[AE]
drx (x) Mi
82V
E
k=1
@PV
Σ E
(x)
Ox?
V
Therefore,
2(m) = 2c TT(VVV((X)) -2 Tr(E-11)
where X ~ μ.
14

---

Published as a conference paper at ICLR 2025
A.2 PROOF OF THEOREM 1
Recall that Lem. and the optimality condition 4
imply @ c(2 - Tr(E-').
)
Now let (m, 2) e BW and let (X, X) be the optimal coupling between μ and n,
[(4) [ < 2c| Tr - -
2c| (x() ++ -
- 2cB (V(VV((X) - -
2cE)AV(() - - -
2c0E - -
S X X =
2ctW2( +)+ + -
Therefore
@(4) 2 2(+) - 2c(W2(H, - -
So Q(u) > Off
+ 2(4)
- (2-
A.3 PROOF OF THEOREM 2 AND REMARK
Proof of Theorem
Recall from Lem.
for any Nmmm E) E (W
2(n) = 2c ((( - - TT(ET)) where ~
Since V a-strongly convex, V²V(x) > al for all æ Rd Therefore, (V²V > al. follows that Tr(VVV((X)) 2 da. Therefore, whenever Tr(E-1) < (2da)/c, Q(H) > 0 and we get reduced variance.
Proof of Remark 2
Assuming that V convex (not necessarily strongly convex) and twice continuously differentiable. We show that: for any Gaussian random variable X, Vx(EV²V(())) > O. By contradiction, suppose that (x(E semidefinite. Therefore, for all z (EVVV((X)= 0 or E(z A (XX)z) O. Let's denote by f the pdf of X, we have: given any fixed z:
follows that EV²V(X) since E E (XX symmetric and positive
o (
Since V is convex, the function under the integral is non-negative. Therefore, the integral being zero implies that the function is zero almost everywhere. The continuity of V further implies that the function under the integral has to be identically zero. Since f(x) > for all x, we deduce V=)(=> 0 for all Now pick 2 e; where e; is the i-th basis vector, i.e., [ex], 1 if
0 otherwise, we get 82, V = O. For # we set that 2; 2y 1, # 0 for all à + i, we x?
have 02, V + 82, V + 20 mnaay V = O, or 02 V = O. We conclude that V = 0, or V is an affine
function. This cannot happen since that (x(EV2v(X)) > O. Therefore, by picking c E (0, ( we guarantee ((n) 0.
e is proper distribution. This contradiction implies
15

---

......d) d) such that Published as a conference paper at ICLR 2025
A.4 DISCUSSION ON CONDITION
For convex V, Remark 2 guarantees that we can pick Cj (that is Pa-measurable) so that т < almost surely. To ensure slightly stronger condition 7i |Tlloo < 1 we proceed as follows. Firstly, from Lemma] we can set
Tk = 20k Tr(E V²V) T
Epc VV - E,z VVI2
On the other hand, it was shown in Diao et al. 2023 Appendix D) that can control the eigenval- ues of 1 along the SGVT's iterations as follows: if V is 8-smooth and Co is initialized such that 3-11 < Σo. and that the stepsize satisfies v/ S f. then 8 Σg for all k. We note that this sult also holds for SVRGVI since their computations only involve the update step of the covariance matrix. Now with Eg ž 3- it holds E₁ ž B-1 almost surely for all k.
Now let us consider a > 0 (strongly convex V) and pick ca Tr(==
where c (0, 1), we get
4a d - -c)
S 1 Tr(E, E
VV En VVI
4a dc?dc(1= c)
$ 1
3E VV Fp VVI2
Therefore, with the assumption that there exists deterministic M₁ such that Ep, VVV Ey VVI2 < Mj almost surely, it holds: almost surely
4a dc c)
Th S1
BM
which implies Th oo < 1.
Now consider a = 0 (convex but non-strongly-convex V) and let us pick cs = 2r Tr(r, V2V) e E (0, 1), we get
where
Tr(E:7)
= 4c(1 (Tr(EE V2V))²
3dMM
If we further assume that [| and the eigenvalues of Eg also bounded above in this case, i.e., there exists a deterministic Bg, Eg such that Amax Ama) S Bg and || S Eş almost surely. Recall that Eg ž B-11.
Since V cannot be affine in the entire domain, there exist ag E RRd and E (1,2,...
@ºv dx² (zo) > 0.
By the continuity of @V there exists r > 0, c > such that
a²v
dx (x) 2 e, Yr E B( r).
By denoting by f(x; m. Σ) the pdf N(m, Σ), we evaluate
Tr(Bp V2V) >
(c))((r mk, Ex)dr
du?
B(xg.r)
2 c(2mT) dct( B(zo.r) c ma) Sx (o-ma) dx
2 <(2iT)
$ A
# tas
c -1|i111 dx.
B(xo.r)
16

---

Published as a conference paper at ICLR 2025
A.5 PROOF OF THEOREM S
Since Alg. I differs from SGVI (Diao et al. 2023) only at by we will largely leverage the con-
vergence analysis of Diao et al. (2023) but will pay attention to the transition of the variance reduction effect to the final bounds.
At the iteration k, the (deterministic) Bures- Mures-Wasserstein gradient of Ev pa is VnwEy(p)) x + Eys, VV + (Bng v2 m2)
and in Alg. we approximate this gradient by
x x by + (=
where by VV(X)) Ix (X2 (X,
m Sk (Xx) and Xx pak-
The enror of this approximation is
ē x (St - Ey, (P + (bx Ep VV).
Let Pk denote o-algebra containing the information up to the beginning of iteration k, Pk o (Xo, X,,..., Let us denote
(k-11) xx-1)
for & (1, 1) and Po is, by convention, the trivial r-algebra.
क.. = B( E( (Px) B(Ezopa Eug V²V)(x m&) + (b. Eyg (PP).
(10)
Bounding हz: we show that
o < 3d3(1 Th) + 6(1 + Tx) Amax W (HR: #),
(11)
The proof of 11 direct extension of Diao et al. (203 Lem. 5.6), but let us partly include here for completeness.
First, let μ = (mm, 2) and X ~ μs applying Stein's lemma we get
E Or;
@V
(X)(X,-m))) -ma) = Σ
d
ExE
8V Fr.dxx
(x)
k=1
Summing up for i = 1,2, 1,2,...,d
d
av
(X)(X, (mxx =
d d
82V Oxp
(x))
E
E
Ox,
k=1
which can be rewritten as
E X -XV"
We now recall the Brascamp-Liet inequality: let o exp(-W) where W is strictly convex and
twice continuously differentiable, then
Var,(f) < E. (V (V
for any smooth f. By using f = (0/0z;) and = Pk. we obtain
((800,() < Ep, VVV.,
(12)
Summing for i from 1 to d
E, VV Eyg 2 < Tr (
Ey, VVV)
Since X₁ is the only source of randomness in é given Pt. the conditional expectation in 10 becomes the expectation over the randomness of Xx. we can write
o. = EV E,,, V²V) -m+) VV(X)) Cx -Xx m&) Ep, V
17

---

Published as a conference paper at ICLR 2025
where X, X2+ Ma and X, X3 are independent. We evaluate
2 50 < BVI(VV((X)) - EnD v2 - B - Co (Xt - -mx) - E, VV
E((X m+) T (V Eyg V2V) mz)) TR VV Byy, VVI2
E((V2V(X)) - -Ep D2V) (X mx)(X mx) + TkE
VV Exp. VVI2
(E5m - En V2V), 22) + Tullip VV - -Eu V
Eyu (VV, Σg V2V) - (B V2)) Σx) + =xp VV Eyp VVI
S E,t, VV2V, Ei V²V) + Bu VV - En V
I (1 + Th) Bm (VV,
B(1 + The BB (V Ex)
= 8(1 + (BL++ - -
Now by using optimal coupling between je and म, one can obtain (Diao et al. 2023, P.27, P.28)
3d + 23 + 8² Tr(EE)
E(VV(X)),x - -m) =
W (1:: #)
2
Therefore,
tR S 3d3(1 + mx) + + Tx) 432 + 233 TM())
W2 #)
S 3d3(1 + Tx) + 6(1 + Tx) Amaxx (2) #).
Bound [Eminn F((x)) - F(#):
Lem. 5.1 in Diao al. (0233 implies that
EW2
(1 - ( F())) 2
(13)
where 2 is the strong convexity modulus of V.
Now using the bound (11) for āk.
EW2 (MK++1, #) & (1 an 12(1 (1 ||xx|==))) Amax (2) +) -2n F 6(1 + ETTh )n² 3d
I exp (-an 12(1 | b0 )n/ 3 Amax (E) (
(HK, #)
22n(BF(uu++) - - + 6(1 + ETA Bd.
Therefore
(nF(F((p - F(+)) S Sexp (a0m) + 12(1 + || bx )r2 Xmnx (5)) (5))
EW2 (mu, #)
EW ( a) 6(1 ET②) )n 3d
(14)
Since are considering the convex case, set 0 and denote C2 = 12(1 + |//x/ m0 )B3 Xmxxx {
and D_ = 0, Dj = - Ck for = 0, 1,..., ......N.. 1. By definition, D₁ + C₁ = DR-1 for all k = 0, 1,......
N-
Rewrite (4
2n (EF((t+t1) - F(*)) exp (Cwn?) EW? (14, #) EW? (M441, #) 6(1 ET②) 27 3d.
Multiply both sides with exp( DEnr²) we get
(pap((Da))(BFF(ueea - F(+))
< exp ((Dx + Cx)m? BW3 #) exp(Dem )EW R) 6(1 + ETT 3dexp((w>>
18

---

Published as a conference paper at ICLR 2025
and, by using the backward recursion D; Cj D can be rewritten
21
I eexp (Dk-1n²) (W2,u, (1x, exp( (ppp (1x+1, R) (M + 6(1 + End Telescope for à: from 0 to -1
2n exp(Dan))(BF((xll F(+))
N-1
-
-xpp(( 1EV3 + 6n/² 3d
Σ
(1 + ET cxp( Dkn²²)
k=
$ ()., + 612 3d
Σ
(1 ET exp(
k=
We see that
C
k
k+ 1 +
> k=0
1
2
where 24/3 Amax (..).
Let us denote (1) = NN exp(Dan²)) It holds
W2 #)
exp(Danr²)
E
min k=1.2...,,
Flums) F(i))
3n3dd
Σ k=0
+ ETr.
2n (1)
Sx(1))
holds
exp(Dwn²)
(1 + ET④)
+ Tmax.E
(15)
Sw(1)
k=0
and
N-1
(T)
exp(Dnn)²)
C
exp
(k
+ +1++
Σ 1=0
2
C
(k: 1 1 (Tr + 1) mmxmmxx )17²)
exp
Tinns
Texp
(-C((k + 1)n?)]
On the other hand, for any > 0, the function f(s) = 6" is convex. By tangent inequality f(s) 2 f 1), get
8" b bln))(s
(16)
Applying the inequality 16 [exp (-C(k 1)
with b = exp (-C(k +1m)) and (Tmax, + 1)/2
exp (-C(((++++))) + C(k: 1)n² exp (-C(k: 1)r?))
1
Tmax.
=exp (-sp(-(k+1)n
+1
1+
Tmaxx
2
exp (-C(k +
|+
)) (1
1 Tmax,
2
19

---

Published as a conference paper at ICLR 2025
Therefore,
Sx(n) 2 1 + Cit - 000
exp(-Okr))
Cr²² )maxx
1
Σ
exp(-Cki))
Cit² o0
1
2
e
C1f²(1 - -mmax.oo)
1+
Min ((C
2
By using the basic inequality 1/ //min(a, b) s 1/a + 1/b, we get
1
e
+ [(Crp])-T]
Sx(1)
C1f((1
oo
N
1+
2
e
1
+ CTP)
Cn²((1 Tmaxx
N
14
2
asymptotically at small 17 > 0.
Therefore,
min N= N F(1E)) - -F(#)= S
1
On
W2 R) + +3n3d(1 + Fma
C1f(11 m
2nyNV
2
1+
2
A.6 PROOF OF THEOREM 4
Since V is a-strongly convex with > 0, E. V2V) b al, so Amin (E+ (BzV) 2 a
follows that
1
Xmax
S
Amin (1-1) Aminn (E* (V2V)
a
Using this inequality in the bound for ōx in 11I we get
6(1 + T33
ah < 3d3(1 + +T)
W (m): 4).
Using this bound for 13)
EW 3W24441, #) $ (1 an)) EW? (Pk, #) -
- F(t))
+ 2 (dd3(1 Tx) + 6(1 + T33 (2( #)
(1 am
12(1 + TA ll 2
EW2 (/4:: +) 6d,33)) ET②)
a
12(1 + Th 2 g3
S exp
-an +
EW? (Mk: #) + +6d + ETG).
a
Now with 7) S a=²
12(1 + 1/ 100 )np 233
(1 + | (oollo))m
4
20

---

Published as a conference paper at ICLR 2025
Therefore
((
-33
Tk
EW (1x-1; S exp
BW2 (mes #) + 6dd31)((1 + (ETT)
na
exp
Tmaxx.0
1jat
#) ( 6d3rn Tmax, E.). EW2 (Jks
Telescope this inequality, we get
exp -N ( 3 x
(20,0, +) + 6(1 + Tmax,E) E) 3d
EW2 #)
na
4
(3 Tmax 00 1n0
exp
N(3
Tmax, 00 )naa
24(1 Tmax, E 3rrdd
exp
W.2 (mo: #)
(3 Timax )x
asymptotically small 17 0.
B
ADDITIONAL EXPERIMENTAL DETAILS
B.I LAPLACE APPROXIMATION
Laplace approximation fits Gaussian approximation by finding the mode of the target (MAP esti- mate for infernece) and forming second order approximation at that point. The approximation is given by
N (MMAP (22V(ZMAP))
We use BFGS optimizer (Nocedal & Wright 2006) implemented in SciPy (Virtanen et al. 2020) 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 and E in the Eu- clidean geometry of the parameter space, using Cholesky factorization for parameterizing the co- variance. This done by maximizing the Evidence Lower BOund (ELBO)
[(m. Σ) = - p(g,(a,, 2) log [=)],
which equivalent to minimizing the KL divergence. We use single-sample reparameterization estimates for the gradient. Furthermore, by stopping the gradient after sampling 2, remove the Fisher score from the gradient computation, giving an unbiased estimator of the gradient of the ELBO with potentially lower variance (Roeder al. (017). We Adam (Kingma & Bal 2015) optimizer and the learning rates and number of iterations found in Table [ found to achieve good convergence. Our implementation is based on the code provided by Modi et al. (2244)
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
Table I: Optimization details for EVI.
21

---

Published as a conference paper at ICLR 2025
Monte Carlo estimator
Our estimator
10a
Variance
10
0
50
100
150
200
250
300
time elapsed (n #iters)
Figure 6: Student's 's experiment: variance along iterations
B.3 STUDENT'S ' DISTRIBUTION
Consider a d-dimensional Student-t distribution with location scale matrix E and v degrees of freedom. negative log density (up to constant), gradient and Hessian are given by:
1 (v + d) log V(x)= 2
1 (x - p1) I 1+ V
(x -m)) (x--))
(v +
VV(x)= = (x - 11) (x 11) 1x -H),
V²V(x) v+ - 4) E 1x -
+ d
E-1
2(v + d)
Σ -- p)(x 4) I-1
3- 4))
B.4 VARIANCE ALONG ITERATIONS
We further report in Figures S and the variance of our proposed estimator and the Monte Carlo estimator along the SVRGVI's iterates. The variance is computed empirically using 5000 i.i.d. sam- ples at each iteration. The results demonstrate that our estimator consistently achieves a significantly smaller variance compared to the Monte Carlo estimator.
10'
10a
1.5
Moete Carlo extimator sstimator
Varfance
Moete Carln extimator uxtimator
Variance
20:
estimator
Our estimator
10-4
50 100
150
200
250
300
0
56
100
150
200
750
100
5n
160
190
250
250
300
tic eapsed (nx Piters)
8mc eapsedd (nx #iters)
tie aapeed lox #iter)
(a) d = 10
(b) d = 50
(c) d = 200
Figure 5: Gaussian experiment: variance along iterations.
B.5 COMPARISONS AGAINST THE MINIBATCHI APPROACII
straightforward approach to reduce the variance to use more MC samples per iteration. In this experiment, samples for SGVI at each iteration, where m E (1, 10, 100). Fig.
illustrates that SGVI requires approximately 100 samples per iteration to achieve a performance level comparable SVRGVI using only one sample per iteration.
22

---

Published as a conference paper at ICLR 2025
Pullullii
Pallud
Pyluslfi
109
102
SGVI sample
SGVL sample
SGW, MC sample
101
SEMI 20 MC samples
163
SEVI 10 MC samples
10
SGV, 10 MC camples
102
SGM,
SVNGVV) MC sample
201
102
SGVI, 100
SMRGVII, MC sample
10
SGV, 100 MC samples
SVitVvi, sample
100
10
10--
10
10--
10
10-1
10-3
10-2
0
50
1:00
150
200
250
300
50
100
150
200
255
300
50
160
1500
200
250
300
time elapsed (n #iters)
time elapsed (nxx #iters)
time elapsed (1 #iters)
(a) d 10
(b) d 50
(c) d 200
Figure 7: Minibatch-SGVi with MC samples versus SVRGVI the Gaussian experiment
In Euclidean VI. Buchholz. et al. 2018)) showed that quasi-MC samples can result in better es- timator (with smaller variances) than standard MC samples. Since the idea is universal, we can use quasi-MC samples improve the performance of SGVI well. Fig. g confirms that using quasi-MC samples indeed leads to better performance in practice SGVI now needs around 50 quasi-MC samples to reach our performance, and with 100 quasi-MC samples, SGVI surpasses our performance.
102
SCM,, 10 qMC
10a
SCM, gMC
169
sGv,,
10
SCM, 50 gMC samples
SCM, 50 gMC samples
10
sGv,, gMC samples
10a
SCM, 100 oMC samples
162
sevi, 100 oMC samples
10'
sGvi, 100 oMC samples
SVRGVI, MC
SVRGVI, sample
SVRGV,, sample
100
10
10
10
10P
10
10-2
10-1
10
0
50
100
150
200
250
300
50 300
150
200
250
300
50
100
150
200
250
300
time elapsed (nx diters)
time elapsed (nx #itens)
time elapsed (g) X diters)
(a) d = 10
(b) d = 50
(c) d = 200
Figure 8: Minibatch-SGGi with quasi-MC samples versus SVRGVI in the Gaussian experiment
B.6 EFFECT OF C
In this experiment, we study the impact of e on the performance of SVRGVI. In Figures 9 10 we report the performance of SVRGVI in the Gaussian, Student't, and Bayesian logistic regression experiments when varies in {0.0, 0.5, 0.8, 1.0, 1.5, 2.0}. The results indicate that performance improves as increases from 0 to 1, peaking around = 1, and then degrades as c continues increasing to 2.0. Furthermore, the performance is somewhat symmetric around e 1, e.g., 0.8 and = 1.2 yield similar results. We therefore confirm that e being around 1 works best in practice.
B.7 EFFECT OF STEP SIZE.
We conduct an experiment to compare the performances of different algorithms with varying step sizes. We consider Gaussian targets with D= 100. We fix the number of steps to 300, and vary the step size between (2,125,0 0.25, 0.5, 1.0]. The results, as shown in Figure 12] indicates that while the previous method requires relatively small step size work relatively well. our algorithm is able to work robustly with large step sizes and achieves the best performances under all step sizes.
We exclude the case of using single quasi-MC sample, as it coincides with the mean of the Gaussian variational distribution when employing Serambled Sobol sequence. In this specific experiment, this leads to an optimal but misleading-resul purely by coincidence.
23

---

Published as a conference paper at ICLR 2025
Optieffi
104
107
10
10a
k-0.0
101
t-0.0
10"
t-0.0
10-1
0-0.5
c-0.5
C-0.8
10
cm1.0
10-3
cm1.0
cm1.0
c21.2
C21.2
10
C-2.0
10
c-2.0
c-2.0
diters)
10a
10'
10a
10-1
10
10-1
10-1
10
10-
10"5
10"5
10-1
0.00 0.25 6.50 1.00
1.25 1.50 2.50
0.00 6.506.55 1.00 1.50 1.75 2.50
0.00 0.25 0,506,5 1.00 3.25 1.75
d = 10
d = 50
== d=200
Figure 9: Gaussian experiments: Upper row. KL divergence along iterations; Lower row: final KL divergence.
Di (ri (empirical)
c=0.0
105
c=0.5
102
10a
c=1.0
c=0.8
(oxiiii (empirical)
c-1.2
10a
c-1.5
2-2.0
102
10a
10a
0
50
100
150
200
250
300
0.25 0.50
1.00 1.25
1.50 1.75 2.00
elapsed (n #iters)
(a) KL divergence along iters
(b) Last KI. divergence
Figure 10: Student's 's t experiment
7lua) - M(Hbas) (empirical)
c=0.0
c=0.5
10a
c=0.8
P(Au) Filcou)) Tempirical)
4x
c=10
c=122
10t
c=1.5
2=2.0
10"
100
10a
50
100
150
200
250
300
0.25 0.50 0.75 1.00
1.25
1.50 1.75 2.00
time elapsed (n Witers)
(a) KT. divergence along iters
Last KT. divergence
Figure 11: Bayesian linear regression experiment
24

---

Published as a conference paper at ICLR 2025
Di (ux pM
Dg (My Fi)
102
10
10a
101
102
10a
10a
BWGD
10
BWGD
10
10
SVRGVI
10
SVRGVI (ours)
10
EVI
EVI
50
100
150
200
250
300
10a
number iterations
step sizes
(a) evolution of results with varying step sizes
(b) final results of the algorithms
Figure 12: Results of different algorithms with varying step sizes in the Gaussian experiment
25


Reviews:
[{"review_id": "m9PHssK675", "reviewer": "Reviewer_Sicz", "summary": "The paper proposes an improved optimization technique for variational inference with a multivariate Gaussian variational family.\n\nBuilding 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.\n\nAuthors support the superiority of the new optimization technique with both theoretical analysis and empirical evaluation, the latter on a set of synthetic benchmarks.", "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:\n- 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.\n- I believe that \"FB\" abbreviation (forward-backward) is never introduced.\n- The term “L-smooth” seemed weird to me, why not “L-Lipschitz”?\n- Page 4, line 174, “affine map” -> “affine maps”.\n- Page 4, line 198, “shows that” -> “showed that”.\n- Page 5, line 233, I would remove parentheses around the “also see ...”, they look ugly immediately after a citep.\n- Page 5, line 237, perhaps add commas around “where $c \\in \\mathbb{R}$”, otherwise the sentence is hard to read.\n- 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.\n- Page 8, line 387, “We also note that conditioning” -> “We also note that by conditioning”.\n- Page 8, footnote 1, please expand what is the minor correction. Also, perhaps start with a capital letter.\n- Page 9, line 457, missing “e” in “covariancs”.\n- Page 9, line 472, “In The” -> “In the”.\n\nNote: no need to respond to the small items above in the rebuttal, save yourself some time.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The paper is written very well. It is a pleasure to read.\n- 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).\n- Empirical evaluation shows very strong improvement over the baselines, more so than the theoretical analysis (which is not claimed to be tight) suggests.", "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.\n\nI would actually give this paper a score of 9, but only 8 and 10 are available.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}, {"review_id": "LiJkSu7IwK", "reviewer": "Reviewer_qVPq", "summary": "- The authors propose a control-variate-based variance reduction technique for gradient estimation in stochastic Gaussian variational inference (SGVI).\n- 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.\n- Additionally, they show that their variance reduction scheme enhances the convergence of SGVI when $V(x)$ is (strongly) convex.\n- Empirically, they confirmed that their method improves the convergence properties of SGVI.", "questions": "Based on the concerns summarized in the weaknesses section, I would like to pose the following questions, categorized into “theoretical” and “experimental” aspects:\n\n### Questions on Theoretical Analyses:\n1.\tCould 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?\n2.\tIn 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?\n3.\tThe 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.\n\n### Questions on Experimental Analyses:\n1.\tThe 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?\n2.\tThere 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)].\n3.\tIn 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.\n4.\tThe 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?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 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.\n- The variance reduction performance of the proposed method is theoretically guaranteed in (strongly) convex settings.\n- The authors demonstrate improved convergence over existing SGVI and Bures-Wasserstein gradient descent (BWGD) methods through several numerical experiments.", "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:\n\n### Concerns Regarding Theoretical Analyses:\n- 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.\n- 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?).\n- 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?\n\n### Concerns Regarding Numerical Experiments:\n- 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.\n- 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.\n- 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)]).\n\n### Lack of Discussion on Other Variance Reduction Studies in EVI:\n- 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.\n\n### Citation:\n- 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.\n- 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\n- Buckholz et al. (2018): A. Buchholz, F. Wenzel, and S. Mandt. Quasi-Monte Carlo Variational Inference. ICML2018. https://arxiv.org/abs/1807.01604.\n- 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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}, {"review_id": "lBqbZloQtA", "reviewer": "Reviewer_7JyP", "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.", "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.\n\n+ 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., \n\n+  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?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "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.", "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}, {"review_id": "8iJTjVRDJA", "reviewer": "Reviewer_NBYM", "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.", "questions": "Please see the weakness section.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 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.", "weaknesses": "I have two major concerns regarding the proposed method in this paper:\n\n1) 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.\n\n2) 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)