# Forward-backward Gaussian variational inference via JKO in the Bures-Wasserstein Space

arXiv: https://arxiv.org/abs/2304.05398 · HTML: https://mlsanity.com/paper/2304.05398/

## Key facts

- Cited by 36 papers in the map (data through 14 Oct 2026).
- Builds on 48 of its 88 references that are in the map.
- First posted on arXiv 10 Apr 2023.

## Builds on (most cited first, 30 of 48)

- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Variational Inference: A Review for Statisticians](https://mlsanity.com/paper/1601.00670/) (2016)
- [Optimization Methods for Large-Scale Machine Learning](https://mlsanity.com/paper/1606.04838/) (2016)
- [Improving Variational Inference with Inverse Autoregressive Flow](https://mlsanity.com/paper/1606.04934/) (2016)
- [Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm](https://mlsanity.com/paper/1608.04471/) (2016)
- [Black Box Variational Inference](https://mlsanity.com/paper/1401.0118/) (2013)
- [Theoretical guarantees for approximate sampling from smooth and log-concave densities](https://mlsanity.com/paper/1412.7392/) (2014)
- [Stein Variational Gradient Descent as Gradient Flow](https://mlsanity.com/paper/1704.07520/) (2017)
- [Analysis of Langevin Monte Carlo via convex optimization](https://mlsanity.com/paper/1802.09188/) (2018)
- [Sampling as optimization in the space of measures: The Langevin dynamics as a composite optimization problem](https://mlsanity.com/paper/1802.08089/) (2018)
- [Noisy Natural Gradient as Variational Inference](https://mlsanity.com/paper/1712.02390/) (2017)
- [Variational inference via Wasserstein gradient flows](https://mlsanity.com/paper/2205.15902/) (2022)
- [A Unified Theory of SGD: Variance Reduction, Sampling, Quantization and Coordinate Descent](https://mlsanity.com/paper/1905.11261/) (2019)
- [On the geometry of Stein variational gradient descent](https://mlsanity.com/paper/1912.00894/) (2019)
- [Improved analysis for a proximal algorithm for sampling](https://mlsanity.com/paper/2202.06386/) (2022)
- [Concentration of tempered posteriors and of their variational approximations](https://mlsanity.com/paper/1706.09293/) (2017)
- [SVGD as a kernelized Wasserstein gradient flow of the chi-squared divergence](https://mlsanity.com/paper/2006.02509/) (2020)
- [Towards a Theory of Non-Log-Concave Sampling: First-Order Stationarity Guarantees for Langevin Monte Carlo](https://mlsanity.com/paper/2202.05214/) (2022)
- [A Non-Asymptotic Analysis for Stein Variational Gradient Descent](https://mlsanity.com/paper/2006.09797/) (2020)
- [Efficient constrained sampling via the mirror-Langevin algorithm](https://mlsanity.com/paper/2010.16212/) (2020)
- [Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networks](https://mlsanity.com/paper/2106.00774/) (2021)
- [Fast and Simple Natural-Gradient Variational Inference with Mixture of Exponential-family Approximations](https://mlsanity.com/paper/1906.02914/) (2019)
- [Scaling limit of the Stein variational gradient descent: the mean field regime](https://mlsanity.com/paper/1805.04035/) (2018)
- [On Riemannian Optimization over Positive Definite Matrices with the Bures-Wasserstein Geometry](https://mlsanity.com/paper/2106.00286/) (2021)
- [Averaging on the Bures-Wasserstein manifold: dimension-free convergence of gradient descent](https://mlsanity.com/paper/2106.08502/) (2021)
- [Handling the Positive-Definite Constraint in the Bayesian Learning Rule](https://mlsanity.com/paper/2002.10060/) (2020)
- [Provable Smoothness Guarantees for Black-Box Variational Inference](https://mlsanity.com/paper/1901.08431/) (2019)
- [Fast Matrix Square Roots with Applications to Gaussian Processes and Bayesian Optimization](https://mlsanity.com/paper/2006.11267/) (2020)
- [Projected Wasserstein gradient descent for high-dimensional Bayesian inference](https://mlsanity.com/paper/2102.06350/) (2021)
- [A Finite-Particle Convergence Rate for Stein Variational Gradient Descent](https://mlsanity.com/paper/2211.09721/) (2022)

## Built on by (most cited first, 30 of 36)

- [Convergence of flow-based generative models via proximal gradient descent in Wasserstein space](https://mlsanity.com/paper/2310.17582/) (2023)
- [Provable convergence guarantees for black-box variational inference](https://mlsanity.com/paper/2306.03638/) (2023)
- [On the Convergence of Black-Box Variational Inference](https://mlsanity.com/paper/2305.15349/) (2023)
- [Batch and match: black-box variational inference with a score-based divergence](https://mlsanity.com/paper/2402.14758/) (2024)
- [Algorithms for mean-field variational inference via polyhedral optimization in the Wasserstein space](https://mlsanity.com/paper/2312.02849/) (2023)
- [Solving Inverse Problems via Diffusion-Based Priors: An Approximation-Free Ensemble Sampling Approach](https://mlsanity.com/paper/2506.03979/) (2025)
- [Towards Understanding the Dynamics of Gaussian-Stein Variational Gradient Descent](https://mlsanity.com/paper/2305.14076/) (2023)
- [Momentum Particle Maximum Likelihood](https://mlsanity.com/paper/2312.07335/) (2023)
- [Theoretical Guarantees for Variational Inference with Fixed-Variance Mixture of Gaussians](https://mlsanity.com/paper/2406.04012/) (2024)
- [Variational Inference in Location-Scale Families: Exact Recovery of the Mean and Correlation Matrix](https://mlsanity.com/paper/2410.11067/) (2024)
- [Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme](https://mlsanity.com/paper/2506.01502/) (2025)
- [Bridging the Gap Between Variational Inference and Wasserstein Gradient Flows](https://mlsanity.com/paper/2310.20090/) (2023)
- [Proximal Interacting Particle Langevin Algorithms](https://mlsanity.com/paper/2406.14292/) (2024)
- [Sampling from the Mean-Field Stationary Distribution](https://mlsanity.com/paper/2402.07355/) (2024)
- [Variational Inference with Mixtures of Isotropic Gaussians](https://mlsanity.com/paper/2506.13613/) (2025)
- [Non-geodesically-convex optimization in the Wasserstein space](https://mlsanity.com/paper/2406.00502/) (2024)
- [A Mean Field Approach to Empirical Bayes Estimation in High-dimensional Linear Regression](https://mlsanity.com/paper/2309.16843/) (2023)
- [Provably Scalable Black-Box Variational Inference with Structured Variational Families](https://mlsanity.com/paper/2401.10989/) (2024)
- [Training Latent Diffusion Models with Interacting Particle Algorithms](https://mlsanity.com/paper/2505.12412/) (2025)
- [A Particle Algorithm for Mean-Field Variational Inference](https://mlsanity.com/paper/2412.20385/) (2024)
- [Optimization Guarantees for Square-Root Natural-Gradient Variational Inference](https://mlsanity.com/paper/2507.07853/) (2025)
- [Stochastic variance-reduced Gaussian variational inference on the Bures-Wasserstein manifold](https://mlsanity.com/paper/2410.02490/) (2024)
- [Variational Online Mirror Descent for Robust Learning in Schr"odinger Bridge](https://mlsanity.com/paper/2504.02618/) (2025)
- [Dense associative memory for Gaussian distributions](https://mlsanity.com/paper/2509.23162/) (2025)
- [DC-LA: Difference-of-Convex Langevin Algorithm](https://mlsanity.com/paper/2601.22932/) (2026)
- [Nearly Dimension-Independent Convergence of Mean-Field Black-Box Variational Inference](https://mlsanity.com/paper/2505.21721/) (2025)
- [Busemann Functions in the Wasserstein Space: Existence, Closed-Forms, and Applications to Slicing](https://mlsanity.com/paper/2510.04579/) (2025)
- [Large-scale Score-based Variational Posterior Inference for Bayesian Deep Neural Networks](https://mlsanity.com/paper/2602.05873/) (2026)
- [Convergence of projected stochastic natural gradient variational inference for various step size and sample or batch size schedules](https://mlsanity.com/paper/2604.00683/) (2026)
- [You Only Accept Samples Once: Fast, Self-Correcting Stochastic Variational Inference](https://mlsanity.com/paper/2406.02838/) (2024)

## Abstract (arXiv)

Variational inference (VI) seeks to approximate a target distribution π by an element of a tractable family of distributions. Of key interest in statistics and machine learning is Gaussian VI, which approximates π by minimizing the Kullback-Leibler (KL) divergence to π over the space of Gaussians. In this work, we develop the (Stochastic) Forward-Backward Gaussian Variational Inference (FB-GVI) algorithm to solve Gaussian VI. Our approach exploits the composite structure of the KL divergence, which can be written as the sum of a smooth term (the potential) and a non-smooth term (the entropy) over the Bures-Wasserstein (BW) space of Gaussians endowed with the Wasserstein distance. For our proposed algorithm, we obtain state-of-the-art convergence guarantees when π is log-smooth and log-concave, as well as the first convergence guarantees to first-order stationary solutions when π is only log-smooth.