# Algorithms for mean-field variational inference via polyhedral optimization in the Wasserstein space

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

## Key facts

- Cited by 20 papers in the map (data through 14 Oct 2026).
- Builds on 43 of its 92 references that are in the map.
- First posted on arXiv 5 Dec 2023.

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

- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Variational Inference: A Review for Statisticians](https://mlsanity.com/paper/1601.00670/) (2016)
- [Computational Optimal Transport](https://mlsanity.com/paper/1803.00567/) (2018)
- [Convex Optimization: Algorithms and Complexity](https://mlsanity.com/paper/1405.4980/) (2014)
- [How to train your neural ODE: the world of Jacobian and kinetic regularization](https://mlsanity.com/paper/2002.02798/) (2020)
- [Fast Computation of Wasserstein Barycenters](https://mlsanity.com/paper/1310.4375/) (2013)
- [Wasserstein Distributionally Robust Optimization: Theory and Applications in Machine Learning](https://mlsanity.com/paper/1908.08729/) (2019)
- [Optimal Entropy-Transport problems and a new Hellinger-Kantorovich distance between positive measures](https://mlsanity.com/paper/1508.07941/) (2015)
- [Sampling as optimization in the space of measures: The Langevin dynamics as a composite optimization problem](https://mlsanity.com/paper/1802.08089/) (2018)
- [Convex Potential Flows: Universal Probability Distributions with Optimal Transport and Convex Optimization](https://mlsanity.com/paper/2012.05942/) (2020)
- [Entropic estimation of optimal transport maps](https://mlsanity.com/paper/2109.12004/) (2021)
- [Variational inference via Wasserstein gradient flows](https://mlsanity.com/paper/2205.15902/) (2022)
- [A Wasserstein-type distance in the space of Gaussian Mixture Models](https://mlsanity.com/paper/1907.05254/) (2019)
- [Plugin Estimation of Smooth Optimal Transport Maps](https://mlsanity.com/paper/2107.12364/) (2021)
- [Rates of Estimation of Optimal Transport Maps using Plug-in Estimators via Barycentric Projections](https://mlsanity.com/paper/2107.01718/) (2021)
- [Optimal transport for Gaussian mixture models](https://mlsanity.com/paper/1710.07876/) (2017)
- [Accelerating Langevin Sampling with Birth-death](https://mlsanity.com/paper/1905.09863/) (2019)
- [Forward-backward Gaussian variational inference via JKO in the Bures-Wasserstein Space](https://mlsanity.com/paper/2304.05398/) (2023)
- [Averaging on the Bures-Wasserstein manifold: dimension-free convergence of gradient descent](https://mlsanity.com/paper/2106.08502/) (2021)
- [Provable Smoothness Guarantees for Black-Box Variational Inference](https://mlsanity.com/paper/1901.08431/) (2019)
- [Provable convergence guarantees for black-box variational inference](https://mlsanity.com/paper/2306.03638/) (2023)
- [On Linear Optimization over Wasserstein Balls](https://mlsanity.com/paper/2004.07162/) (2020)
- [Fast and Smooth Interpolation on Wasserstein Space](https://mlsanity.com/paper/2010.12101/) (2020)
- [On the Convergence of Black-Box Variational Inference](https://mlsanity.com/paper/2305.15349/) (2023)
- [Optimal transport map estimation in general function spaces](https://mlsanity.com/paper/2212.03722/) (2022)
- [Theoretical and Computational Guarantees of Mean Field Variational Inference for Community Detection](https://mlsanity.com/paper/1710.11268/) (2017)
- [Mean-field Variational Inference via Wasserstein Gradient Flow](https://mlsanity.com/paper/2207.08074/) (2022)
- [Fr'echet Means and Procrustes Analysis in Wasserstein Space](https://mlsanity.com/paper/1701.06876/) (2017)
- [Learning normalizing flows from Entropy-Kantorovich potentials](https://mlsanity.com/paper/2006.06033/) (2020)

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

- [Measure-to-measure interpolation using Transformers](https://mlsanity.com/paper/2411.04551/) (2024)
- [Convergence of coordinate ascent variational inference for log-concave measures via optimal transport](https://mlsanity.com/paper/2404.08792/) (2024)
- [Mirror and Preconditioned Gradient Descent in Wasserstein Space](https://mlsanity.com/paper/2406.08938/) (2024)
- [Theoretical Guarantees for Variational Inference with Fixed-Variance Mixture of Gaussians](https://mlsanity.com/paper/2406.04012/) (2024)
- [Convergence rate of random scan Coordinate Ascent Variational Inference under log-concavity](https://mlsanity.com/paper/2406.07292/) (2024)
- [Variational Inference in Location-Scale Families: Exact Recovery of the Mean and Correlation Matrix](https://mlsanity.com/paper/2410.11067/) (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)
- [Constructive approximate transport maps with normalizing flows](https://mlsanity.com/paper/2412.19366/) (2024)
- [A Particle Algorithm for Mean-Field Variational Inference](https://mlsanity.com/paper/2412.20385/) (2024)
- [Stability of Mean-Field Variational Inference](https://mlsanity.com/paper/2506.07856/) (2025)
- [Linearized Wasserstein Barycenters: Synthesis, Analysis, Representational Capacity, and Applications](https://mlsanity.com/paper/2410.23602/) (2024)
- [Nearly Dimension-Independent Convergence of Mean-Field Black-Box Variational Inference](https://mlsanity.com/paper/2505.21721/) (2025)
- [Rotated Mean-Field Variational Inference and Iterative Gaussianization](https://mlsanity.com/paper/2510.07732/) (2025)
- [Radon--Wasserstein Gradient Flows for Interacting-Particle Sampling in High Dimensions](https://mlsanity.com/paper/2602.05227/) (2026)
- [Static and Dynamic Approaches to Computing Barycenters of Probability Measures on Graphs](https://mlsanity.com/paper/2603.26940/) (2026)
- [Wasserstein Contraction of Coordinate Ascent Variational Inference](https://mlsanity.com/paper/2605.30253/) (2026)
- [Gaussian mixture layers for neural networks](https://mlsanity.com/paper/2508.04883/) (2025)
- [Mode Collapse of Mean-Field Variational Inference](https://mlsanity.com/paper/2510.17063/) (2025)
- [Stability of Finite-Batch Particle Mean-Field Variational Inference Beyond Strong Convexity](https://mlsanity.com/paper/2608.11486/) (2026)

## Abstract (arXiv)

We develop a theory of finite-dimensional polyhedral subsets over the Wasserstein space and optimization of functionals over them via first-order methods. Our main application is to the problem of mean-field variational inference, which seeks to approximate a distribution π over ℝd by a product measure π^star. When π is strongly log-concave and log-smooth, we provide (1) approximation rates certifying that π^star is close to the minimizer π^star_diamond of the KL divergence over a polyhedral set P_diamond, and (2) an algorithm for minimizing KL(·|π) over P_diamond based on accelerated gradient descent over Rd. As a byproduct of our analysis, we obtain the first end-to-end analysis for gradient-based algorithms for MFVI.