Stochastic Gradient Variational Inference with Price's Gradient Estimator from Bures-Wasserstein to Parameter Space
Organizations: University of Pennsylvania, Philadelphia, U.S. · Yale University, New Haven, U.S. · University of California San Diego, San Diego, U.S. · University of British Columbia, Vancouver, Canada.
Abstract
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.
Figures & tables
| Algorithm | Space | Gradient Estimator | Maximum Step Size | Iterations Complexity | Reference |
| SPGD | Reparam. | d\kappa^{2}\mathrm{tr}\mathopen{}\mathclose{{\left(\mu\Sigma_{*}}}\right)\nicefrac{{1}}{{\epsilon}} | Domke et al. (2023) Kim et al. (2023) | ||
| SPBWGD | Bonnet–Price | d\kappa\frac{1}{\epsilon}\log\frac{1}{\epsilon}+\kappa^{3}\log\mathopen{}\mathclose{{\left(\Delta^{2}\frac{1}{\epsilon}}}\right) | Diao et al. (2023) | ||
| SPGD | Bonnet–Price | d\kappa\frac{1}{\epsilon}+\sqrt{d}\kappa^{3/2}\log\mathopen{}\mathclose{{\left(\kappa\Delta^{2}}}\right)\frac{1}{\sqrt{\epsilon}}+\kappa^{2}\log\frac{1}{\epsilon} | Theorem 3.3 | ||
| SPBWGD | Bonnet–Price | d\kappa\frac{1}{\epsilon}+\sqrt{d}\kappa^{3/2}\log\mathopen{}\mathclose{{\left(\kappa\Delta^{2}}}\right)\frac{1}{\sqrt{\epsilon}}+\kappa^{2}\log\frac{1}{\epsilon} | Theorem 3.2 |
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
| Name | Description | Reference | |
| Dogs | Logistic regression model of the traumatic learning behavior of dogs. The dataset is from the Solomon-Wynne experiment. (model: ; dataset: ) | 3 | Gelman and Hill (2021) Solomon and Wynne (1953) |
| Basketball | Hidden Markov model of NBA basketball player SportVU tracking data during a drive event. (model: ; dataset: ) | 6 | Ali (2019) |
| Lynx | Lotka-Volterra model of a lynx-hare population. The dataset is the number of pelts collected by the Hudson’s Bay Company in the years 1910–1920. (model: ; dataset: ) | 8 | Carpenter (2018) Hewitt (1921) |
| NES2000 | Linear model of political party identification. The data set is from the 2000 National Election Study. (model: ; dataset: ) | 10 | Gelman and Hill (2021) |
| Bones | Latent trait model for multiple ordered categorical responses for quantifying skeletal maturity from radiograph maturity ratings with missing entries. (model: ; dataset: ) | 13 | Spiegelhalter et al. (1996) |
| SISLOB | Loss model of insurance claims. The model is the single line-of-business, single insurer (SISLOB) variant, where the dataset is the “ppauto” line of business, part of the “Schedule P loss data” provided by the Casualty Actuarial Society. (model: ; dataset: ) | 15 | Cooney (2017) |