We are thankful to the reviewer for the positive comments on our work.

Weaknesses

1. Using more samples indeed reduces the variances, but is far less efficient. Our method is effectively as easy to implement and has virtually no computational overhead. We now added a new experiment where we empirically compare against the naive approach to make the difference clear. Figure 7 in Appendix B.5 shows that to reach the same convergence rate one needs to use around 100 samples (and hence gradient oracles) per iteration, compared to just one sample for our method.

2. To clarify, while BWGD can be extended for Gaussian mixture approximations, to our knowledge, SGVI cannot. We build on SGVI and retain the same limitations. The reason, we think, is quite understandable: the whole point of SGVI is the tractability of the backward step (JKO step), which is known to have a closed-formed solution for the family of single Gaussians. It is still unclear whether this step is still tractable for Gaussian mixtures.

Thank you for the reply.

We further clarify that the difference between BWGD and SGVI is: that BWGD is forward Euler, while SGVI is forward-backward Euler. Although we focus on SGVI (because the entropy is non-smooth, making SGVI the method of choice in general), we can use the same idea to improve the Bures-Wasserstein gradient used in BWGD for a single Gaussian case (because the BW gradient used in BWGD is the BW gradient used in the forward step of SGVI plus a computable deterministic term).

Additionally, we clarify that: different from the single Gaussian case, BWGD for mixture of Gaussians is a deterministic method -- using quadrature based on sigma points to compute expectations. Therefore, developing variance reduction for a deterministic method is unnecessary.

We are thankful to the reviewer for the positive comments on our work.

Questions

• In principle, there are no obstacles to running for realistic statistical inference problems because the method is universal -- as long as we can access the gradient and the Hessian of the log-target. Our experimentation focused on synthetic problems as in previous BW-VI works to best illustrate and quantify the effect of the variance reduction. We leave empirical comparisons on more challenging tasks as future work; they need more work on careful tuning of learning rates etc. that could not be fit in this paper.

• While comparisons across different families of inference methods (here VI vs MCMC) are interesting, they are best left for papers where that is the main objective contribution (for example, Yao et al. "Yes, but Did It Work?: Evaluating Variational Inference", ICML 2018). Already defining the metrics for such a comparison would be outside of our scope. Both our theoretical analysis and empirical results are on convergence rates and gradient variances, which are meaningful quantities for VI formulated as an optimization problem, but not for MCMC. Accuracy comparisons, in turn, would mostly depend on the experimental setup (e.g. whether the approximation family contains the true posterior or not) and not so much on the specific VI algorithm.

• We now empirically quantify the variance reduction for the gradients in Appendix B.5, to show that the improved convergence is because of the variance reduction. One possible explanation is that the deterministic forward-backward (FB) Euler using BW geometry works very well, hence closely following the trajectory of FB Euler would be optimal. We also note that the improvement is more pronounced in high dimensions. Comparing the bounds of SGVI and SVRGVI, variance reduction helps scale down the term relating to the dimension $d$, where a small scaling factor can make a big difference when $d$ is large.

Thank you for the detailed comments and the effort put into the review.

Concerns Regarding Theoretical Analyses:

• Our convergence analysis is in two stages for clarity; We first show our estimator has a smaller variance than the Monte Carlo estimator in several scenarios, and then show that the benefit of variance reduction will propagate through the proof of SGVI [Diao2023] to improve their bounds. We can discuss the behaviour of the bounds as $N$ approaches infinity as follows. As $N \to \infty$, the term associated with $W_2^2(\mu_0,\hat{\pi})$ goes to zero (sublinearly for the convex case and linearly for the strongly convex case), while the constant term associated with the noise behaves like $O(\eta)$ where $\eta$ is the stepsize. For a fixed $\eta>0$, the constant terms specify a "convergence neighbourhood" where further convergence is not possible, especially due to the noise coming from the Hessian $\nabla^2 V(X_k)$ that does not diminish. However, our convergence neighbourhood is smaller than the one in SGVI of Diao et al., meaning that the variance reduction helps our algorithm reach closer to the optimal solution than SGVI.

• First, we clarify that Theorem 1 holds also without the convexity assumption. For Theorem 3, the analysis can easily be relaxed to the later iterations when we are within the vicinity of the optimal solution (see our comment in lines 389 -- 391), similar to warm-up optimization: initially run the vanilla version of SGVI (simply set $c=0$) and switch to using SVRGVI (set, e.g., $c=0.9$) at the convergence phase. We also note that the convergence phase is where variance reduction is needed the most for stable convergence. The method always helps (in theory) for stable convergence around the optimum, and, in practice, provides benefits right from the start. We further note that for strongly log-concave targets, variance reduction happens also further from the optimal solution (Theorem 2).

• For the smooth case the radius $r$ depends on the smoothness parameter $\ell$ of the Laplacian of the log target as well as the variance of the optimal distribution $\hat{\pi}$. The smoother the log-target and the larger the variance of the optimal distribution, the larger the radius. However, we currently do not have means of quantifying it or determining what is sufficiently large (and relative to what).

Concerns Regarding Numerical Experiments:

• We build on SGVI and BWGD, both of which only considered synthetic data as well. We fully agree there is clear value in real-world examples, but for a still developing family of algorithms, the effects can be best demonstrated and quantified in controlled experiments. Note that we considered more targets and in higher dimensions, compared to the previous works.

• This is a very good point. We now added new result figures (Appendix B.4) showing the variance reduction along iterations in practice -- we obtain substantially smaller variance compared to the standard single-sample Monte-Carlo estimator. This explains the variance reduction as the driving force for the improved convergence. Regarding $c$, we have added the sensitivity experiment in Appendix B.6, showing that $c$ being around $1$ works best in practice.

• Our primary objective is to improve Bures-Wasserstein Variational Inference (BW-VI) methods and explain when and how they work. Note that we did not claim those methods (including ours) would beat all Euclidean Variational Inference (EVI) methods. Instead, the comparison against EVI was included to show that in some cases the previous BW-VI methods can be worse than rather standard Euclidean VI methods -- an important observation that is not apparent based on the previous literature -- and that this is at least partly due to the high variance of the gradient estimators that can be easily resolved. We chose the method of Roeder et al. (2017) as a representative reference method because it is also a variance-reduction method using a score and shares similar characteristics to our variance-reduction method in the Bures-Wasserstein space. In the revised version we will discuss newer EVI methods and point out thorough empirical comparison as a valuable future work.

Citations: [Diao2023] Diao, M. Z., Balasubramanian, K., Chewi, S., & Salim, A. Forward-backward Gaussian variational inference via JKO in the Bures-Wasserstein space. ICML 2023.

Lack of Discussion on Other Variance Reduction Studies in Euclidean VI:

More advanced techniques beyond control variates in Euclidean VI are very appreciated and it would be a promising future work to integrate these techniques to further improve BW VI. We extended the discussion on Euclidean VI as suggested in lines 503-506. However, since we have not yet studied these methods in our setting, we are not able to discuss what challenges they may cause.

Thank you for your detailed and thoughtful rebuttal. While your responses have addressed several of my concerns, I still have some remaining questions and would appreciate further clarification on the following points:

Additional Questions for Theory:

• Applicability of Theorem 2 to Convex Functions:
  • Based on Theorem 2, variance reduction is claimed to be achievable for all iterations. However, this result seems to hold only under the $\alpha$-strongly convex assumption. Does this imply that Theorem 2 cannot be directly applied to the general convex case covered in Theorem 3? If so, could Theorem 2 be extended to address the general convex scenario?
• Condition $\mathrm{Tr}(\Sigma^{-1}) < \frac{2\alpha d}{c}$ in Theorem 2:
  • Is this condition generally satisfied in practical scenarios? Remark 2 suggests that achieving variance reduction requires aligning the target distribution properties with $\alpha$ and calculating $\lambda_{\mathrm{min}}(\Sigma)$ at each iteration. However, this step is not included in Algorithm 1. How should practitioners use the proposed method for approximating target distributions based on convex functions in such cases?
• Practical Intuition for Theorem Conditions:
  • The assumptions and results in theorems, such as Theorem 2’s variance reduction or Theorem 1’s smaller variance, are not intuitively clear in terms of practical implications (this reflect my lack of understanding, maybe?). For instance, according to the rebuttal, variance reduction in Theorem 1 depends on $r$, but identifying $r$ is challenging, leaving it unclear whether the reduced variance is sufficiently small. Does this property align with the goal of ensuring small variance at each iteration? How does the proposed method compare to variance reduction approaches based on sampling techniques, such as Quasi-Monte Carlo?

Additional Questions for Experiments

• Importance of Hyperparameter Settings:

  • As indicated by Theorem 3, settings such as learning rates are crucial for achieving the theoretical results. Would the authors agree that further experimentation and discussions on this aspect are necessary prior to publication?

• Visualization and Sensitivity Analyses:

  • Could the authors use an anonymous image-sharing service to share results showing measured variance and sensitivity analyses? These results are critical for verifying whether the claimed variance reduction is genuinely achieved by the proposed method.

• Regarding the comparison to EVI with variance reduction

  • I believe the current demonstration of the advantages of the proposed method over EVI-based approaches still appears insufficient. For example, Figure 3 shows that EVI with variance reduction outperforms both SGVI and BWGD. This raises the concern that more advanced variance reduction techniques developed for EVI might surpass the proposed method in performance. If feasible, could you provide comparisons with recent state-of-the-art variance reduction techniques in EVI to clarify how the proposed method stands in relation to them?

We thank the reviewer for the reply and the thoughtful comments.

First, we address that we updated the paper, so please check out our new PDF file in OpenReview.

Additional Questions for Theory:

It is more natural for us to address your three bullet points in the following order: second - first - third.

• Condition $Tr(\Sigma^{-1})<\frac{2\alpha d}{c}$: Yes, we can verify and check this condition in practical scenarios. Theorem 2 says that: for $\alpha$-strongly convex $V$, we obtain variance reduction whenever $c<(2\alpha d)/Tr(\Sigma^{-1})$. We also note that, at the current iteration, we already have the Cholesky decomposition of $\Sigma$ as $\Sigma = L L^{\top}$, available from the sampling step. Therefore, we can compute $Tr(\Sigma^{-1})$ based on the triangular matrix $L$ with the complexity of $O(d^2)$, which is neglectable. Therefore, we can always guarantee variance reduction if we, at the current iteration, pick $c$ smaller than the computable term $(2\alpha d)/Tr(\Sigma^{-1})$. This is also our recommendation for the practitioner in the worst case -- if they want to be safe; However, in practice, $c$ can be chosen more flexibly. For example, as stated in Theorem 1, in a neighbourhood of the optimal solution, any $c \in (0,2)$ should guarantee variance reduction (even without requiring the target to be log-concave!).

• Applicability of Theorem 2 to Convex Functions: Yes, we can extend to convex case ($\alpha=0$). We assume that $V$ is twice continuously differentiable. We show that: for any Gaussian random variable $X$, $Tr(\mathbb{E} \nabla^2 V(X)) > 0$. By contradiction, suppose that $Tr(\mathbb{E} \nabla^2 V(X)) = 0$, it follows that $\mathbb{E} \nabla^2 V(X) = 0$ since $\mathbb{E} \nabla^2 V(X)$ is symmetric and positive semidefinite. Therefore, for all $z$: $z^{\top} (\mathbb{E} \nabla^2 V(X)) z = 0$ or $\mathbb{E}(z^{\top} \nabla^2 V(X) z) = 0$. Let's denote by $f$ the pdf of $X$, we have: given any fixed $z$:

$

\int{(z^{\top} \nabla^2 V(x) z) f(x)dx} = 0

$

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 $\nabla^2 V$ further implies that the function under the integral has to be identically zero. Since $f(x)>0$ for all $x$, we deduce $z^{\top} \nabla^2 V(x) z = 0$ for all $x$. Now pick $z = e_i$ where $e_i$ is the i-th basis vector, i.e., the i-th component is $1$, and $0$ otherwise, we get $\partial^2_{x_i^2} V \equiv 0$. For $i \neq j$, we set $z$ so that $z_i = z_j = 1$, $z_k=0$ for all $k \neq i,j$, we have $\partial^2_{x_i^2} V + \partial^2_{x_j^2} V + 2 \partial^2_{x_i x_j} V \equiv 0$, or $\partial^2_{x_i x_j} V \equiv 0$. We conclude that $\nabla^2 V \equiv 0$, or $V$ is an affine function. This cannot happen since $V = -\log(\pi)$. This contradiction implies that $Tr(\mathbb{E} \nabla^2 V(X)) > 0$. From Lemma 1, we deduce that variance reduction happens whenever $0< c< 2 Tr(\mathbb{E} \nabla^2 V(X)) / Tr(\Sigma^{-1})$. In this case, the upper bound of $c$ is not given as explicitly as in the strongly convex case, but the point here is that variance reduction still happens in this context.

• Practical Intuition for Theorem Conditions: First, as answered above, the practical implication for Theorem 2 is now clear. We empirically verify in Appendix B.4 that we obtained strong variance reduction effectively from the beginning. About the comparison to the trivial approach such as Monte Carlo with multiple samples (see also our answer to Reviewer NBYM), we added an experiment in Appendix B.5, showing that the Monte Carlo method needs roughly 100 samples (hence gradient oracles) per iteration to reach our performance using just one sample per iteration.

Additional Questions for Experiments

• Importance of Hyperparameter Settings: Yes, we now added a study on the stepsize in Appendix B.7, showing that our method strongly outperforms previous BW methods under different stepsizes, our method is also less sensitive and allows using larger step sizes.

• Visualization and Sensitivity Analyses: Yes, we measured the variance along iterations in Appendix B.4 and performed sensitivity analyses regarding $c$ in Appendix B.6.

• Regarding the comparison to EVI with variance reduction: we are trying to incorporate one more EVI as a baseline if time allows.

Given the above rebuttal, we would like to additionally make this comment to resolve the main concern of the reviewer which is about the condition that "variance reduction holds for all iterations":

In the strongly convex case (Theorem 4), yes we can guarantee variance reduction for all iterations if we set $c_k$ such that $0< c_k < 2\alpha d / Tr(\Sigma_k^{-1})$. In the convex (but not strongly convex) case (Theorem 3), variance reduction also happens at all iterations if we set $c_k$ small enough.

In response to the additional experiment regarding EVI, we remark that Miller et al. (2017) consider only Gaussian with diagonal covariance, it is misleading and besides the point to compare it against our method that considers full covariance matrices. About Buckholz et al. (2018), the core idea of the method is to replace Monte Carlo samples with quasi-Monte Carlo samples to improve the estimator. While their idea was used in EVI, it is universal and it is straightforward to incorporate the idea into Bures-Wasserstein VI methods. Consequently, instead of plotting the results of Buckholz et al. (2018) directly, we adopted their quasi-Monte Carlo sampling approach to improve SGVI, enabling a more direct and meaningful comparison to our method. To this end, we have added a new experiment in Appendix B.5, which demonstrates that using quasi-Monte Carlo samples outperforms traditional Monte Carlo samples. However, to achieve comparable performance to our method SVRGVI (with one MC sample per iteration), SGVI requires around 50 quasi-Monte Carlo samples per iteration. With 100 quasi-Monte Carlo samples, SGVI surpasses our method (which still uses one sample per iteration).

Citations:

• 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.

Thank you for the thoughtful and positive comments on our work and the acceptance recommendation. We have addressed the minor points raised in the question section in the revised version.