---
title: Wasserstein Probability Flow Methods
url: https://www.emergentmind.com/topics/wasserstein-probability-flow-method
type: topic
---

# Wasserstein Probability Flow Methods

Wasserstein Probability Flow denotes a family of constructions that evolve, estimate, or match probability measures by exploiting Wasserstein geometry. In the recent literature, the expression is used for sliced-Wasserstein proximal schemes and generative probability flows [2110.10972], Wasserstein-on-Wasserstein flow matching [2605.08424], parameterized Wasserstein gradient descent in diffeomorphism space [2603.13579], nonstationary distribution estimation with a Wasserstein penalty between consecutive distributions [2507.05893], and particle-flow algorithms for free-support Wasserstein barycenters [2509.11435]. Across these usages, probability measures are treated as dynamical objects on $\mathcal P_2(\mathbb R^d)$ or related spaces, and the resulting methods are expressed through optimal-transport distances, continuity equations, transport plans, or push-forward maps.

## 1. Terminological scope and uses in the literature

The literature uses the label for multiple non-identical constructions rather than for a single universally standardized algorithm. Some works study Wasserstein gradient flows in the classical Jordan–Kinderlehrer–Otto sense; others use deterministic transport ODEs, particle-flow updates, or network-flow reformulations. This terminological multiplicity is important because papers with closely related names solve different mathematical problems.

| Usage | Core object | Representative reference |
|---|---|---|
| Sliced-Wasserstein probability flow | $\mu_{k+1}\in\arg\min_\mu \bigl[SW_2^2(\mu,\mu_k)/(2\tau)+\mathcal F(\mu)\bigr]$ | [2110.10972] |
| Wasserstein-on-Wasserstein flow matching | Deterministic transport of measures over measures | [2605.08424] |
| Parameterized Wasserstein gradient flow | Push-forward density $\rho_\theta=(T_\theta)_\sharp \rho_{\rm ref}$ | [2603.13579] |
| Nonstationary distribution estimation | $\sum_t\log(P_t\{\widehat\xi_t\})-\lambda\sum_t W(P_t,P_{t+1})$ | [2507.05893] |
| Particle-flow barycenter method | Barycenter atoms advected by averaged optimal-transport displacements | [2509.11435] |

A recurring source of confusion is the proximity between “Wasserstein gradient flow,” “probability flow ODE,” and “particle flow.” In some papers, the flow is the steepest descent of a free-energy functional in Wasserstein space; in others, it is a deterministic transport dynamics learned from couplings or induced by transport plans. The terminology overlaps, but the mathematical objects being optimized are not identical.

## 2. Geometric and variational foundations

A large part of the literature starts from the 2-Wasserstein metric on $\mathcal P_2(\mathbb R^d)$,
$$
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),
$$
and studies the minimization of a functional on probability measures. In the gradient-flow setting, one seeks a curve $\rho_t$ solving
$$
\partial_t\rho_t=\nabla\!\cdot\!\Bigl(\rho_t\nabla\frac{\delta \mathcal F}{\delta \rho}(\rho_t)\Bigr),
$$
which is the steepest-descent dynamics of $\mathcal F$ in Wasserstein space [2110.10972]. Typical functionals include the Fokker–Planck functional
$$
\mathcal F(\mu)=\int V(x)\,d\mu(x)+\mathcal H(\mu),
$$
the interaction energy
$$
\mathcal F(\mu)=\tfrac12\iint W(x-y)\,d\mu(x)d\mu(y),
$$
and KL-type objectives of the form
$$
F[q]=\mathrm{KL}(q\Vert \pi_k)=\int q(x)\log q(x)\,dx+\int q(x)V(x)\,dx+\text{const},
$$
which recover exact posteriors at the global minimizer in filtering problems [2110.10972], [2303.06398].

The standard discrete variational realization is the Jordan–Kinderlehrer–Otto scheme
$$
\mu_{k+1}\in\arg\min_\mu \frac{1}{2\tau}W_2^2(\mu,\mu_k)+\mathcal F(\mu),
$$
whose continuous-time limit yields PDEs such as the Fokker–Planck equation. The same geometric picture appears in several domains: approximate inference for diffusions, nonlinear filtering, generative modeling, barycenters, and constrained decomposition of measures [1806.04542], [2303.06398], [2406.00914].

A related branch interprets the KL divergence on Wasserstein space as a Riemannian objective whose gradient flow is Langevin dynamics. In that setting, $\mathcal P_2(\mathbb R^d)$ is treated as a formal infinite-dimensional Riemannian manifold, and continuous-time Riemannian SGD and SVRG flows are obtained by introducing stochastic-gradient variance terms into the induced Fokker–Planck equation [2401.13530]. This extends the deterministic gradient-flow picture to stochastic optimization on the space of probability measures.

## 3. Numerical realizations: proximal schemes, sliced distances, and path losses

A central numerical difficulty is that direct JKO steps require solving a nested optimization problem at each iteration and are known for computational challenges, especially in high dimension. One response is to replace $W_2$ by the sliced-Wasserstein distance
$$
SW_2^2(\mu,\nu)=\int_{S^{d-1}}W_2^2(P^\theta_\#\mu,P^\theta_\#\nu)\,d\theta,
$$
leading to the SW-JKO scheme
$$
\mu_{k+1}\in\arg\min_{\mu\in\mathcal P_2}\Bigl[\frac{SW_2^2(\mu,\mu_k)}{2\tau}+\mathcal F(\mu)\Bigr].
$$
Because the one-dimensional $W_2^2$ admits the closed form
$$
\int_0^1|F_{\hat\mu}^{-1}(u)-F_{\hat\nu}^{-1}(u)|^2\,du,
$$
computable by sorting, the Monte-Carlo approximation of $SW_2$ is differentiable in the samples. This allows the density at each step to be parameterized by any generative model, including a normalizing flow or any invertible or parametric generative network, with overall cost
$$
O(K\,N_e\,L\,N\log N)
$$
over $K$ outer steps [2110.10972].

Another line avoids spatial discretization by regularizing transport and solving the proximal problem in a dual function space. In approximate inference with Wasserstein gradient flows, the regularized proximal step is written with $W_{2,\gamma}$, then converted by Fenchel–Rockafellar duality into a stochastic program over dual potentials $g,h$, which are restricted to an RKHS. A representer theorem reduces the optimization to finite-dimensional coefficients, and the resulting scheme is explicitly discretization-free [1806.04542].

A third approach avoids both temporal and spatial discretization by enforcing self-consistency of the velocity field. In self-consistent velocity matching, one solves the continuity equation through a fixed-point condition
$$
v_t(x)=f_t(x;\mu_t),\qquad \mu_t=(\Phi_t)_\#\mu_0,
$$
and trains a parameterized velocity field by minimizing a residual loss evaluated along its own probability flow. The method is grid-free, matches a wider range of mass-conserving PDEs, and has complexity that scales linearly in $T$ and $d$ for the reported implementations [2301.13737].

Recent work on generative path finding for Wasserstein gradient flow shifts the focus from one-step proximal descent to optimization of the full path. GenWGP minimizes either a finite-horizon action
$$
S_T[p]=\frac12\int_0^T\|\partial_t p_t+\nabla_{\mathcal W}\mathcal F(p_t)\|_{-1,p_t}^2\,dt
$$
or a reparameterization-invariant geometric action, while enforcing approximately constant intrinsic speed between adjacent network layers. The method is designed to avoid delicate time stepping constraints and to remain largely independent of temporal or geometric discretization [2604.11519].

## 4. Parameterized transport maps and finite-dimensional reductions

A major strand of Wasserstein probability-flow methods represents the evolving density as the push-forward of a fixed reference measure by a parameterized map. In the Gross–Pitaevskii setting, the density is written as
$$
\rho_\theta=(T_\theta)_\sharp \rho_{\rm ref},
\qquad
\rho_\theta(T_\theta(z))=\frac{\rho_{\rm ref}(z)}{|\det \nabla T_\theta(z)|},
$$
with $T_\theta$ an orientation-preserving diffeomorphism. The Wasserstein gradient flow on densities then lifts to an $L^2(\rho_{\rm ref})$ gradient flow on maps, and restriction to a finite-dimensional family induces the pull-back metric
$$
G(\theta)=\int_\Omega [\partial_\theta T_\theta(z)]^\top[\partial_\theta T_\theta(z)]\,\rho_{\rm ref}(z)\,dz.
$$
The parameter update is the natural gradient flow
$$
\dot\theta(\tau)=-G(\theta)^{-1}\nabla_\theta E(\theta),
$$
or, discretely,
$$
\theta^{k+1}=\theta^k-\alpha\,G(\theta^k)^+\nabla_\theta E(\theta^k).
$$
This parameterized Wasserstein gradient flow is entirely mesh-free and preserves the unit-mass constraint without normalization [2603.13579].

For numerical realization, the one-dimensional transport map may be represented by a boundary-preserving Neural ODE,
$$
\dot w(t)=f_\theta(w(t)),\qquad
f_\theta(w)=\bigl(1-w^2/L^2\bigr)g_\theta(w),
$$
where the boundary factor enforces $T_\theta(\pm L)=\pm L$ and guarantees that $T_\theta$ is a diffeomorphism of the domain. The same work reports that the output of the parameterized Wasserstein gradient flow can be used to initialize the $H^1$ Sobolev gradient flow, reducing the initial energy gap by a factor of $7$ in 2D and $4.5$ in 3D compared to trivial initial conditions [2603.13579].

A Hamiltonian analogue replaces gradient descent by a reduced Hamiltonian system on parameter space. In parameterized Wasserstein Hamiltonian flow, the push-forward manifold $P_\Theta=\{\rho_\theta:T_{\theta\#}\lambda=\rho_\theta\}$ inherits a pull-back metric $G(\theta)$, and the finite-dimensional Hamiltonian is
$$
H(\theta,p)=\tfrac12\,p^\top G(\theta)^{-1}p+F(\theta).
$$
The resulting ODE is solved by a symplectic scheme, which is intended to preserve the symplectic structure and approximately conserve the Hamiltonian over long times; the method is fully deterministic and does not involve direct optimization over network parameters [2306.00191].

## 5. Deterministic transport learning, probability-flow ODEs, and Wasserstein-on-Wasserstein extensions

In score-based generative modeling, the term probability flow refers to a deterministic ODE whose marginals coincide with those of a forward SDE. For the linear SDE
$$
dx_t=-f(t)x_t\,dt+g(t)\,dB_t,
$$
the associated probability-flow ODE is
$$
\frac{d\tilde x_t}{dt}
=
f(T-t)\,\tilde x_t+\tfrac12 g(T-t)^2\nabla\log p_{T-t}(\tilde x_t),
\qquad
\tilde x_0\sim p_T.
$$
When the true score is replaced by a learned estimate, one obtains an ODE-based sampler. A non-asymptotic convergence analysis in 2-Wasserstein distance is available under accurate score estimates and smooth log-concave data distributions, with the proof relying on explicit contraction rates, synchronous coupling, and analysis of discretization and score-matching errors [2401.17958].

Flow matching provides a second deterministic-transport viewpoint. Given any coupling $\pi\in C(\mu_0,\mu_1)$, linear interpolation of particles,
$$
\mathrm{proj}^t(x,y)=(1-t)x+t\,y,
$$
induces intermediate measures $\rho_t=(\mathrm{proj}^t)_\#\pi$ and a velocity field
$$
v_t^\pi(z)=\mathbb E_{(x,y)\sim\pi}[\,y-x\mid (1-t)x+t\,y=z\,].
$$
A neural vector field $v_\theta(t,x)$ can then be trained by regressing $v_\theta(t,z_j)$ to $x_j'-x_j$ on interpolated samples. This provides the basis for flow matching in Euclidean Wasserstein space and its generalization to spaces of measures [2605.08424].

The Wasserstein-on-Wasserstein extension lifts the geometry from $\mathcal P_2(\mathbb R^d)$ to
$$
X=\mathcal P_2(\mathbb R^d),\qquad M=\mathcal P_2(X),
$$
with outer distance
$$
\mathbb W_2^2(M_0,M_1)=\min_{\Pi\in C(M_0,M_1)}\int_{X\times X}W_2^2(\mu,\nu)\,d\Pi(\mu,\nu).
$$
The construction uses an outer plan between distributions over measures and a measurable family of inner plans between individual measures. Direct computation is expensive, with cost $O(B^3+B^2N^3)$ per gradient step in the reported formulation, so scalable surrogates are introduced: independent couplings, exact minibatch OT, sliced-Wasserstein couplings with complexity $O(S\,B^2\log B+S\,N\log N)$, and lazy linear Wasserstein couplings with cost $O(B^2+BN)$ [2605.08424].

## 6. Applications, guarantees, and limitations

The application range of Wasserstein probability-flow methods is unusually broad. In free-support barycenter computation, barycenter atoms evolve as particles advected by averaged optimal-transport displacements, with barycentric projections of optimal transport plans used in place of Monge maps when the latter do not exist. The discrete update
$$
z_i^{(k+1)}\leftarrow (1-2\eta_k)z_i^{(k)}+2\eta_k\sum_{n=1}^N\pi_n T_n^{(k)}(z_i^{(k)})
$$
admits monotone descent, convergence to stationary points, stability with respect to perturbations of the inputs, and resolution consistency as the number of atoms increases [2509.11435].

In nonstationary distribution estimation, the objective is not a transport PDE but a penalized likelihood,
$$
\max_{P_1,\dots,P_T\in\mathfrak P(\Xi)}
\sum_{t=1}^T\log(P_t\{\widehat\xi_t\})
-\lambda\sum_{t=1}^{T-1}W(P_t,P_{t+1}),
$$
with $W$ the first-order Wasserstein distance. A support-reduction argument shows that each $P_t$ can be taken discrete on observed points, and the problem reduces to a convex network-flow problem with $O(T^2)$ variables and $O(T)$ flow-conservation constraints. The reported numerical tests include dairy-commodity price forecasting and risk-averse portfolio optimization [2507.05893].

Filtering and inference provide another important application domain. Variational Gaussian filtering via Wasserstein gradient flows approximates the filtering posterior by minimizing $\mathrm{KL}(q\Vert\pi_k)$ and reducing the induced PDE to ODEs for Gaussian or mixture-of-Gaussians parameters. The method is reported to remain competitive in systems with multiplicative noise and multimodal state distributions, where Gaussian approximations typically fail [2303.06398]. Approximate inference with Wasserstein gradient flows for diffusion processes similarly uses a discretization-free JKO approximation and, in a nonlinear filtering task, is evaluated against fine-grid integration, Gaussian filters, and particle filtering [1806.04542].

Several theoretical and practical limitations recur. Some guarantees require strong log-concavity, accurate score estimates, or smoothness assumptions on intermediate marginals [2401.17958]. Sliced-Wasserstein JKO schemes provide provable descent and a simple algorithmic structure, but continuous-time properties such as the limiting PDE remain an open theoretical question [2110.10972]. GenWGP assumes regularity conditions such as densities bounded away from $0$ and $C^2$ smoothness of the flow [2604.11519]. Taken together, these works suggest that “Wasserstein Probability Flow” is best understood not as one fixed procedure but as a research area centered on transporting, estimating, and optimizing probability measures by directly using Wasserstein geometry.

Source: https://www.emergentmind.com/topics/wasserstein-probability-flow-method