# Mean-field Variational Inference via Wasserstein Gradient Flow

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

## Key facts

- Cited by 19 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 73 references that are in the map.
- First posted on arXiv 17 Jul 2022.

## Builds on (most cited first, 28 of 28)

- [Variational Inference: A Review for Statisticians](https://mlsanity.com/paper/1601.00670/) (2016)
- [Semi-Supervised Learning with Deep Generative Models](https://mlsanity.com/paper/1406.5298/) (2014)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [Black Box Variational Inference](https://mlsanity.com/paper/1401.0118/) (2013)
- [Non-convex learning via Stochastic Gradient Langevin Dynamics: a nonasymptotic analysis](https://mlsanity.com/paper/1702.03849/) (2017)
- [Coordinate Descent Algorithms](https://mlsanity.com/paper/1502.04759/) (2015)
- [Automatic Differentiation Variational Inference](https://mlsanity.com/paper/1603.00788/) (2016)
- [Statistical guarantees for the EM algorithm: From population to sample-based analysis](https://mlsanity.com/paper/1408.2156/) (2014)
- [Sampling as optimization in the space of measures: The Langevin dynamics as a composite optimization problem](https://mlsanity.com/paper/1802.08089/) (2018)
- [The Landscape of Empirical Risk for Non-convex Losses](https://mlsanity.com/paper/1607.06534/) (2016)
- [Structured Logconcave Sampling with a Restricted Gaussian Oracle](https://mlsanity.com/paper/2010.03106/) (2020)
- [Optimal Control Via Neural Networks: A Convex Approach](https://mlsanity.com/paper/1805.11835/) (2018)
- [Variational inference via Wasserstein gradient flows](https://mlsanity.com/paper/2205.15902/) (2022)
- [Large-Scale Wasserstein Gradient Flows](https://mlsanity.com/paper/2106.00736/) (2021)
- [Sharp convergence rates for Langevin dynamics in the nonconvex setting](https://mlsanity.com/paper/1805.01648/) (2018)
- [Concentration of tempered posteriors and of their variational approximations](https://mlsanity.com/paper/1706.09293/) (2017)
- [Optimal dimension dependence of the Metropolis-Adjusted Langevin Algorithm](https://mlsanity.com/paper/2012.12810/) (2020)
- [Convergence Rates of Variational Posterior Distributions](https://mlsanity.com/paper/1712.02519/) (2017)
- [Neural Monge Map estimation and its applications](https://mlsanity.com/paper/2106.03812/) (2021)
- [Singularity, Misspecification, and the Convergence Rate of EM](https://mlsanity.com/paper/1810.00828/) (2018)
- [Automated Variational Inference in Probabilistic Programming](https://mlsanity.com/paper/1301.1299/) (2013)
- [Neural Estimation of Statistical Divergences](https://mlsanity.com/paper/2110.03652/) (2021)
- [On Statistical Optimality of Variational Bayes](https://mlsanity.com/paper/1712.08983/) (2017)
- [Theoretical and Computational Guarantees of Mean Field Variational Inference for Community Detection](https://mlsanity.com/paper/1710.11268/) (2017)
- [Approximate inference with Wasserstein gradient flows](https://mlsanity.com/paper/1806.04542/) (2018)
- [Statistical Inference in Mean-Field Variational Bayes](https://mlsanity.com/paper/1911.01525/) (2019)
- [Bayesian Repulsive Gaussian Mixture Model](https://mlsanity.com/paper/1703.09061/) (2017)
- [On the Convergence of Extended Variational Inference for Non-Gaussian Statistical Models](https://mlsanity.com/paper/1902.05068/) (2019)

## Built on by (most cited first, 19 of 19)

- [Forward-backward Gaussian variational inference via JKO in the Bures-Wasserstein Space](https://mlsanity.com/paper/2304.05398/) (2023)
- [Algorithms for mean-field variational inference via polyhedral optimization in the Wasserstein space](https://mlsanity.com/paper/2312.02849/) (2023)
- [Towards Understanding the Dynamics of Gaussian-Stein Variational Gradient Descent](https://mlsanity.com/paper/2305.14076/) (2023)
- [Independent projections of diffusions: Gradient flows for variational inference and optimal mean field approximations](https://mlsanity.com/paper/2309.13332/) (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)
- [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)
- [Wasserstein Gradient Flow over Variational Parameter Space for Variational Inference](https://mlsanity.com/paper/2310.16705/) (2023)
- [Particle Mean Field Variational Bayes](https://mlsanity.com/paper/2303.13930/) (2023)
- [A Mean Field Approach to Empirical Bayes Estimation in High-dimensional Linear Regression](https://mlsanity.com/paper/2309.16843/) (2023)
- [A Primer on Variational Inference for Physics-Informed Deep Generative Modelling](https://mlsanity.com/paper/2409.06560/) (2024)
- [A Particle Algorithm for Mean-Field Variational Inference](https://mlsanity.com/paper/2412.20385/) (2024)
- [On Representations of Mean-Field Variational Inference](https://mlsanity.com/paper/2210.11385/) (2022)
- [Convergence Analysis of the Wasserstein Proximal Algorithm beyond Geodesic Convexity](https://mlsanity.com/paper/2501.14993/) (2025)
- [Semi-Implicit Functional Gradient Flow for Efficient Sampling](https://mlsanity.com/paper/2410.17935/) (2024)
- [Scalable Mean-Field Variational Inference via Preconditioned Primal-Dual Optimization](https://mlsanity.com/paper/2602.07632/) (2026)
- [Gaussian mixture layers for neural networks](https://mlsanity.com/paper/2508.04883/) (2025)
- [Stability of Finite-Batch Particle Mean-Field Variational Inference Beyond Strong Convexity](https://mlsanity.com/paper/2608.11486/) (2026)

## Abstract (arXiv)

Variational inference, such as the mean-field (MF) approximation, requires certain conjugacy structures for efficient computation. These can impose unnecessary restrictions on the viable prior distribution family and further constraints on the variational approximation family. In this work, we introduce a general computational framework to implement MF variational inference for Bayesian models, with or without latent variables, using the Wasserstein gradient flow (WGF), a modern mathematical technique for realizing a gradient flow over the space of probability measures. Theoretically, we analyze the algorithmic convergence of the proposed approaches, providing an explicit expression for the contraction factor. We also strengthen existing results on MF variational posterior concentration from a polynomial to an exponential contraction, by utilizing the fixed point equation of the time-discretized WGF. Computationally, we propose a new constraint-free function approximation method using neural networks to numerically realize our algorithm. This method is shown to be more precise and efficient than traditional particle approximation methods based on Langevin dynamics.
