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

For approximating a target distribution given only its unnormalized log-density, stochastic gradient-based variational inference (VI) algorithms are a popular approach. For example, Wasserstein VI (WVI) and black-box VI (BBVI) perform gradient descent in measure space (Bures-Wasserstein space) and parameter space, respectively. Previously, for the Gaussian variational family, convergence guarantees for WVI have shown superiority over existing results for black-box VI with the reparametrization gradient, suggesting the measure space approach might provide some unique benefits. In this work, however, we close this gap by obtaining identical state-of-the-art iteration complexity guarantees for both. In particular, we identify that WVI's superiority stems from the specific gradient estimator it uses, which BBVI can also leverage with minor modifications. The estimator in question is usually associated with Price's theorem and utilizes second-order information (Hessians) of the target log-density. We will refer to this as Price's gradient. On the flip side, WVI can be made more widely applicable by using the reparametrization gradient, which requires only gradients of the log-density. We empirically demonstrate that the use of Price's gradient is the major source of performance improvement.

Variational inference (VI) is a collection of algorithms for approximating a distribution, for which we may have access only to its unnormalized log densities, with some other distribution that is easier to work with. For example, Gaussians. This problem setup commonly arises in fields such as Bayesian statistics and generative modeling. While many algorithms for VI exist, we study two particular algorithms based on stochastic gradient descent, called Wasserstein VI (WVI) and black-box VI (BBVI). Unlike BBVI, WVI utilizes the distributional nature of the VI problem through the Wasserstein geometry. Intuitively, one would expect WVI to perform better than BBVI, since it leverages more of the VI problem's inherent properties. Indeed, previous theoretical and empirical results suggested that WVI converges faster than the typical implementations of BBVI. In this work, we theoretically and empirically demonstrate that this superiority was due to a particular implementation detail---the gradient estimator used in WVI and BBVI---rather than to the algorithm's geometric nature. Hence, there is unfortunately no further evidence to support the claim that WVI is inherently superior to BBVI in the considered setup.