Stochastic Gradient Variational Inference with Price's Gradient Estimator from Bures-Wasserstein to Parameter Space

A new research paper titled 'Stochastic Gradient Variational Inference with Price's Gradient Estimator from Bures-Wasserstein to Parameter Space' has made a significant theoretical breakthrough in machine learning optimization. Authored by Kyurae Kim, Qiang Fu, Yi-An Ma, Jacob R. Gardner, and Trevor Campbell, the work resolves a longstanding question about the relative performance of two major variational inference (VI) approaches: Wasserstein VI (WVI) and Black-Box VI (BBVI).

Previously, WVI—which performs gradient descent in measure space (specifically Bures-Wasserstein space)—had demonstrated superior convergence guarantees compared to BBVI, which operates in parameter space. This suggested fundamental advantages to the measure-space approach. However, this new research proves that both methods can achieve identical state-of-the-art iteration complexity guarantees when using the same gradient estimator. The critical factor is Price's gradient estimator, which leverages second-order Hessian information from the target log-density. The authors demonstrate that BBVI can match WVI's performance by incorporating this estimator with minor modifications.

The practical implication is substantial for machine learning practitioners. Instead of choosing between algorithm families based on perceived theoretical superiority, developers can now select the most convenient implementation (parameter space vs. measure space) while using the optimal gradient estimator for their problem. The research also shows that WVI can be made more widely applicable by using the reparameterization gradient, which requires only first-order information. This work effectively decouples algorithmic framework choices from gradient estimation strategies, providing clearer guidance for implementing efficient variational inference in practice.