Wasserstein Gradient Flows (WGF)

• Wasserstein Gradient Flows are continuous-time dynamical systems that characterize the steepest descent evolution of functionals over probability measures via the 2-Wasserstein metric.
• Discrete-time schemes like the JKO and forward–backward splitting methods provide practical approximations with provable convergence rates under convexity assumptions.
• This framework bridges optimal transport, PDE analysis, and machine learning, enabling rigorous and scalable optimization in infinite-dimensional spaces.

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 (Salim et al., 2020).

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 (Salim et al., 2020).

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 (Salim et al., 2020).

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 (Salim et al., 2020).

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 (Salim et al., 2020).

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 (Salim et al., 2020).

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 (Salim et al., 2020).