---
title: Wasserstein Gradient Flows (WGF)
url: https://www.emergentmind.com/topics/wasserstein-gradient-flows-wgf
type: topic
---

# Wasserstein Gradient Flows (WGF)

Wasserstein Gradient Flows (WGF) are continuous-time dynamical systems that characterize the steepest descent evolution of a functional over the space of probability measures endowed with the 2-Wasserstein metric. The WGF framework provides a rigorous, geometrically-intrinsic generalization of gradient descent to infinite-dimensional spaces, with foundational relevance across optimal transport, partial differential equations, and probabilistic machine learning.

## 1. The 2-Wasserstein Space: Metric, Geometry, and Geodesics

The space of Borel probability measures on ℝᵈ with finite second moments,
$$\mathcal{P}_2(\mathbb{R}^d) := \left\{ \mu \text{ probability on } \mathbb{R}^d : \int \|x\|^2\,d\mu(x) < \infty \right\},$$
equipped with the 2-Wasserstein distance,
$$W_2^2(\mu,\nu) := \inf_{\pi\in\Gamma(\mu,\nu)} \int_{\mathbb{R}^d\times \mathbb{R}^d} \|x-y\|^2\,d\pi(x,y),$$
becomes a geodesic metric space, where $\Gamma(\mu,\nu)$ is the set of couplings of μ and ν. When μ is absolutely continuous, the optimal transport map is given by the gradient of a convex function (Brenier's theorem), and constant-speed geodesics can be constructed as pushforwards via interpolated maps: $\mu_t = ((1-t)\mathrm{Id} + t\,T)_{\#}\mu$ for t∈[0,1], where $T$ is the optimal transport map. The geodesic structure is central for defining “steepest descent” in this space [2002.03035].

## 2. Continuous-Time Formulation: Evolution Equation and Variational Characterization

For a given functional $\mathcal{F} : \mathcal{P}_2 \to (-\infty, +\infty]$, the curve $\mu(t)$ solving the *Wasserstein gradient flow* is characterized by the Evolution Variational Inequality (EVI):
$$
\forall \nu \in \mathcal{P}_2, \quad \frac{d}{dt} W_2^2(\mu(t), \nu) \leq -2 \big[ \mathcal{F}(\mu(t)) - \mathcal{F}(\nu) \big].
$$
Under regularity conditions, this is equivalent to a PDE for the density $\rho(t,x)$:
$$
\partial_t \rho + \nabla \cdot (\rho v) = 0, \quad v = -\nabla_x \left(\frac{\delta\mathcal{F}}{\delta \rho} \right),
$$
where $\frac{\delta\mathcal{F}}{\delta \rho}$ denotes the first variation of $\mathcal{F}$. For example, if $\mathcal{F}(\mu)=\int V\,d\mu + \int \rho\log\rho$, the gradient flow yields the Fokker-Planck equation, a prototypical diffusive evolution [2002.03035].

## 3. Discrete-Time Schemes: JKO and Forward-Backward Splitting

The canonical time-discretization of WGF is the Jordan–Kinderlehrer–Otto (JKO) implicit Euler scheme:
$$
\mu_{n+1} \in \arg\min_{\mu \in \mathcal{P}_2} \left\{ \mathcal{F}(\mu) + \frac{1}{2\gamma} W_2^2(\mu, \mu_n) \right\}.
$$
This yields a sequence whose piecewise-constant interpolation converges to the continuous WGF as $\gamma \to 0$.

When the objective function decomposes as $\mathcal{F} = \mathcal{U} + \mathcal{G}$ with $\mathcal{U}$ smooth and $\mathcal{G}$ possibly nonsmooth but geodesically convex, the Forward–Backward (FB) proximal-gradient algorithm over $\mathcal{P}_2$ is defined as:
- **Forward (gradient) step** for $\mathcal{U}$: $\nu_{n+1} := (\mathrm{Id} - \gamma \nabla F)_{\#} \mu_n$,
- **Backward (proximal) step** for $\mathcal{G}$: $\mu_{n+1} \in \arg\min_\mu \left\{ \mathcal{G}(\mu) + (1/2\gamma) W_2^2(\mu, \nu_{n+1}) \right\} =: \mathrm{Prox}_{\gamma \mathcal{G}}(\nu_{n+1})$,

mirroring the classical Euclidean proximal-point framework. Here, $\mathrm{Prox}_{\gamma \mathcal{G}}$ is a JKO step for $\mathcal{G}$ only [2002.03035].

## 4. Convergence Theory for Proximal Splitting and Rates

Suppose $F$ is $L$-smooth and $\lambda$-strongly convex, and $\mathcal{G}$ is proper, lower semicontinuous, and convex along generalized geodesics. If $\gamma < 1/L$, the FB scheme satisfies a discrete EVI:
$$
W_2^2(\mu_{n+1}, \mu_*) \leq (1-\gamma\lambda) W_2^2(\mu_n, \mu_*) - 2\gamma [\mathcal{F}(\mu_{n+1}) - \mathcal{F}(\mu_*)].
$$
- If $\lambda=0$, $\mathcal{F}(\mu_n)-\mathcal{F}(\mu_*) = O(1/(\gamma n))$.
- If $\lambda>0$, $W_2^2(\mu_n, \mu_*) \leq (1-\gamma\lambda)^n W_2^2(\mu_0, \mu_*)$ (linear convergence).

This result establishes WGF-FB as an infinite-dimensional analog of the proximal gradient method, retaining convergence guarantees familiar from convex Euclidean optimization [2002.03035].

## 5. Practical Implementation, Computational Aspects, and Examples

Continuous-time WGF enjoys exact decay rates, while discrete-time schemes (JKO, FB) match these rates up to step-size constraints. The main numerical challenge is evaluating the proximal map (JKO subproblem), which, depending on $\mathcal{G}$, may admit:
- Closed-form solutions (e.g., negative entropy/heat flow),
- PDE-based solvers (for more complex energies),
- Entropic regularization or Sinkhorn algorithms for approximation.

FB splitting reduces the implicit computation to the $\mathcal{G}$ part only, with the $\mathcal{U}$ part handled by a simple push-forward. In the canonical quadratic-plus-entropy example (sampling from a Gaussian), each FB step maintains Gaussianity, and closed-form recursions for mean and covariance yield linear W₂-convergence. Particle-based (sample-wise) push-forward strategies with optional heat flow accurately reflect continuous-time contraction, even in high dimensions [2002.03035].

## 6. Extensions, Applications, and Open Directions

The FB splitting framework for Wasserstein gradient flows enables:
- Handling composite objectives with both smooth and nonsmooth contributions,
- Direct generalization from Euclidean optimization,
- Provable convergence under geodesic convexity,
- Scalability to high dimensions when approximate or closed-form JKO operators are available.

Ongoing research targets efficient algorithms for more general energy landscapes (including non-convex energies, non-Euclidean underlying domains), adaptive schemes, high-dimensional and large-scale applications, and connections to stochastic optimization and sampling [2002.03035].

## Table: Summary of Classical vs. Proximal-Splitting WGF Schemes

| Method        | Iteration Definition                                                    | Complexity per Step    |
|---------------|------------------------------------------------------------------------|-----------------------|
| JKO           | $\mu_{n+1} \leftarrow \arg\min_{\mu} \mathcal{F}(\mu) + \frac{1}{2\gamma}W_2^2(\mu, \mu_n)$ | Full proximal (often hard/expensive) |
| FB-Splitting  | First pushforward by $-\gamma\nabla F$, then Prox$_{\gamma\mathcal{G}}$ | Cheaper: only Prox$_{\gamma\mathcal{G}}$ |


The Wasserstein Proximal Gradient framework thus defines and analyzes an efficient and theoretically well-founded approach to composite optimization over the space of measures, with direct applicability to variational inference, sampling, and PDE evolution models [2002.03035].

Source: https://www.emergentmind.com/topics/wasserstein-gradient-flows-wgf