# Large-Scale Wasserstein Gradient Flows

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

## Key facts

- Cited by 77 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 72 references that are in the map.
- First posted on arXiv 1 Jun 2021.

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

- [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)
- [OptNet: Differentiable Optimization as a Layer in Neural Networks](https://mlsanity.com/paper/1703.00443/) (2017)
- [Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm](https://mlsanity.com/paper/1608.04471/) (2016)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [Molecular Sets (MOSES): A Benchmarking Platform for Molecular Generation Models](https://mlsanity.com/paper/1811.12823/) (2018)
- [Analyzing Inverse Problems with Invertible Neural Networks](https://mlsanity.com/paper/1808.04730/) (2018)
- [Advances in Variational Inference](https://mlsanity.com/paper/1711.05597/) (2017)
- [Optimal transport mapping via input convex neural networks](https://mlsanity.com/paper/1908.10962/) (2019)
- [Stochastic Optimization for Large-scale Optimal Transport](https://mlsanity.com/paper/1605.08527/) (2016)
- [i-RevNet: Deep Invertible Networks](https://mlsanity.com/paper/1802.07088/) (2018)
- [Maximum Mean Discrepancy Gradient Flow](https://mlsanity.com/paper/1906.04370/) (2019)
- [Large-Scale Optimal Transport and Mapping Estimation](https://mlsanity.com/paper/1711.02283/) (2017)
- [Wasserstein-2 Generative Networks](https://mlsanity.com/paper/1909.13082/) (2019)
- [Convex Potential Flows: Universal Probability Distributions with Optimal Transport and Convex Optimization](https://mlsanity.com/paper/2012.05942/) (2020)
- [Optimal Control Via Neural Networks: A Convex Approach](https://mlsanity.com/paper/1805.11835/) (2018)
- [Sliced-Wasserstein Flows: Nonparametric Generative Modeling via Optimal Transport and Diffusions](https://mlsanity.com/paper/1806.08141/) (2018)
- [Continuous Wasserstein-2 Barycenter Estimation without Minimax Optimization](https://mlsanity.com/paper/2102.01752/) (2021)
- [Scalable Computations of Wasserstein Barycenter via Input Convex Neural Networks](https://mlsanity.com/paper/2007.04462/) (2020)
- [Policy Optimization as Wasserstein Gradient Flows](https://mlsanity.com/paper/1808.03030/) (2018)
- [Deep Generative Learning via Variational Gradient Flow](https://mlsanity.com/paper/1901.08469/) (2019)
- [Large-scale Log-determinant Computation through Stochastic Chebyshev Expansions](https://mlsanity.com/paper/1503.06394/) (2015)
- [Dynamical Optimal Transport on Discrete Surfaces](https://mlsanity.com/paper/1809.07083/) (2018)
- [Approximate inference with Wasserstein gradient flows](https://mlsanity.com/paper/1806.04542/) (2018)
- [A variational approach to path estimation and parameter inference of hidden diffusion processes](https://mlsanity.com/paper/1508.00506/) (2015)
- [On Wasserstein Reinforcement Learning and the Fokker-Planck equation](https://mlsanity.com/paper/1712.07185/) (2017)
- [Proximal Recursion for Solving the Fokker-Planck Equation](https://mlsanity.com/paper/1809.10844/) (2018)

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

- [Proximal Optimal Transport Modeling of Population Dynamics](https://mlsanity.com/paper/2106.06345/) (2021)
- [Supervised Training of Conditional Monge Maps](https://mlsanity.com/paper/2206.14262/) (2022)
- [Variational Wasserstein gradient flow](https://mlsanity.com/paper/2112.02424/) (2021)
- [A Computational Framework for Solving Wasserstein Lagrangian Flows](https://mlsanity.com/paper/2310.10649/) (2023)
- [Understanding DDPM Latent Codes Through Optimal Transport](https://mlsanity.com/paper/2202.07477/) (2022)
- [Neural Monge Map estimation and its applications](https://mlsanity.com/paper/2106.03812/) (2021)
- [Normalizing flow neural networks by JKO scheme](https://mlsanity.com/paper/2212.14424/) (2022)
- [Generative Sliced MMD Flows with Riesz Kernels](https://mlsanity.com/paper/2305.11463/) (2023)
- [Neural Optimal Transport with General Cost Functionals](https://mlsanity.com/paper/2205.15403/) (2022)
- [Convergence of flow-based generative models via proximal gradient descent in Wasserstein space](https://mlsanity.com/paper/2310.17582/) (2023)
- [Kernel Neural Optimal Transport](https://mlsanity.com/paper/2205.15269/) (2022)
- [On amortizing convex conjugates for optimal transport](https://mlsanity.com/paper/2210.12153/) (2022)
- [Meta Optimal Transport](https://mlsanity.com/paper/2206.05262/) (2022)
- [Energy-guided Entropic Neural Optimal Transport](https://mlsanity.com/paper/2304.06094/) (2023)
- [Posterior Sampling Based on Gradient Flows of the MMD with Negative Distance Kernel](https://mlsanity.com/paper/2310.03054/) (2023)
- [Discrete Langevin Sampler via Wasserstein Gradient Flow](https://mlsanity.com/paper/2206.14897/) (2022)
- [One-Step Generative Modeling via Wasserstein Gradient Flows](https://mlsanity.com/paper/2605.11755/) (2026)
- [Wasserstein Iterative Networks for Barycenter Estimation](https://mlsanity.com/paper/2201.12245/) (2022)
- [Scalable Wasserstein Gradient Flow for Generative Modeling through Unbalanced Optimal Transport](https://mlsanity.com/paper/2402.05443/) (2024)
- [Learning diffusion at lightspeed](https://mlsanity.com/paper/2406.12616/) (2024)
- [Efficient Gradient Flows in Sliced-Wasserstein Space](https://mlsanity.com/paper/2110.10972/) (2021)
- [Unifying GANs and Score-Based Diffusion as Generative Particle Models](https://mlsanity.com/paper/2305.16150/) (2023)
- [Particle-based Variational Inference with Preconditioned Functional Gradient Flow](https://mlsanity.com/paper/2211.13954/) (2022)
- [Mean-field Variational Inference via Wasserstein Gradient Flow](https://mlsanity.com/paper/2207.08074/) (2022)
- [Self-Consistency of the Fokker-Planck Equation](https://mlsanity.com/paper/2206.00860/) (2022)
- [Input Convex Gradient Networks](https://mlsanity.com/paper/2111.12187/) (2021)
- [Taming Hyperparameter Tuning in Continuous Normalizing Flows Using the JKO Scheme](https://mlsanity.com/paper/2211.16757/) (2022)
- [Self-Consistent Velocity Matching of Probability Flows](https://mlsanity.com/paper/2301.13737/) (2023)
- [Accelerating Sinkhorn Algorithm with Sparse Newton Iterations](https://mlsanity.com/paper/2401.12253/) (2024)
- [Optimal Neural Network Approximation of Wasserstein Gradient Direction via Convex Optimization](https://mlsanity.com/paper/2205.13098/) (2022)

## Abstract (arXiv)

Wasserstein gradient flows provide a powerful means of understanding and solving many diffusion equations. Specifically, Fokker-Planck equations, which model the diffusion of probability measures, can be understood as gradient descent over entropy functionals in Wasserstein space. This equivalence, introduced by Jordan, Kinderlehrer and Otto, inspired the so-called JKO scheme to approximate these diffusion processes via an implicit discretization of the gradient flow in Wasserstein space. Solving the optimization problem associated to each JKO step, however, presents serious computational challenges. We introduce a scalable method to approximate Wasserstein gradient flows, targeted to machine learning applications. Our approach relies on input-convex neural networks (ICNNs) to discretize the JKO steps, which can be optimized by stochastic gradient descent. Unlike previous work, our method does not require domain discretization or particle simulation. As a result, we can sample from the measure at each time step of the diffusion and compute its probability density. We demonstrate our algorithm's performance by computing diffusions following the Fokker-Planck equation and apply it to unnormalized density sampling as well as nonlinear filtering.