Wasserstein Probability Flow Methods

• Wasserstein Probability Flow is a collection of techniques that evolve and optimize probability measures using the geometry of optimal transport, including sliced-Wasserstein and gradient flow formulations.
• These methods enable applications in generative modeling, barycenter computation, nonlinear filtering, and nonstationary distribution estimation, offering versatile tools for modern inference and optimization tasks.
• They combine variational principles with parameterized transport maps and proximal schemes to address high-dimensional computational challenges while preserving essential mass-conservation properties.

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 (Bonet et al., 2021), Wasserstein-on-Wasserstein flow matching (Piening et al., 8 May 2026), parameterized Wasserstein gradient descent in diffeomorphism space (Zhang et al., 13 Mar 2026), nonstationary distribution estimation with a Wasserstein penalty between consecutive distributions (Anderson et al., 8 Jul 2025), and particle-flow algorithms for free-support Wasserstein barycenters (You, 14 Sep 2025). 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]$$	(Bonet et al., 2021)
Wasserstein-on-Wasserstein flow matching	Deterministic transport of measures over measures	(Piening et al., 8 May 2026)
Parameterized Wasserstein gradient flow	Push-forward density $\rho_\theta=(T_\theta)_\sharp \rho_{\rm ref}$	(Zhang et al., 13 Mar 2026)
Nonstationary distribution estimation	$$\sum_t\log(P_t\{\widehat\xi_t\})-\lambda\sum_t W(P_t,P_{t+1})$$	(Anderson et al., 8 Jul 2025)
Particle-flow barycenter method	Barycenter atoms advected by averaged optimal-transport displacements	(You, 14 Sep 2025)

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 (Bonet et al., 2021). Typical functionals include the Fokker–Planck functional

$$\mathcal F(\mu)=\int V(x)\,d\mu(x)+\mathcal H(\mu),$$

the interaction energy

$$\mu_{k+1}\in\arg\min_\mu \bigl[SW_2^2(\mu,\mu_k)/(2\tau)+\mathcal F(\mu)\bigr]$$0

and KL-type objectives of the form

$$\mu_{k+1}\in\arg\min_\mu \bigl[SW_2^2(\mu,\mu_k)/(2\tau)+\mathcal F(\mu)\bigr]$$1

which recover exact posteriors at the global minimizer in filtering problems (Bonet et al., 2021, Corenflos et al., 2023).

The standard discrete variational realization is the Jordan–Kinderlehrer–Otto scheme

$$\mu_{k+1}\in\arg\min_\mu \bigl[SW_2^2(\mu,\mu_k)/(2\tau)+\mathcal F(\mu)\bigr]$$2

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 (Frogner et al., 2018, Corenflos et al., 2023, Han et al., 2024).

A related branch interprets the KL divergence on Wasserstein space as a Riemannian objective whose gradient flow is Langevin dynamics. In that setting, $$\mu_{k+1}\in\arg\min_\mu \bigl[SW_2^2(\mu,\mu_k)/(2\tau)+\mathcal F(\mu)\bigr]$$3 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 (Yi et al., 2024). 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 $$\mu_{k+1}\in\arg\min_\mu \bigl[SW_2^2(\mu,\mu_k)/(2\tau)+\mathcal F(\mu)\bigr]$$4 by the sliced-Wasserstein distance

$$\mu_{k+1}\in\arg\min_\mu \bigl[SW_2^2(\mu,\mu_k)/(2\tau)+\mathcal F(\mu)\bigr]$$5

leading to the SW-JKO scheme

$$\mu_{k+1}\in\arg\min_\mu \bigl[SW_2^2(\mu,\mu_k)/(2\tau)+\mathcal F(\mu)\bigr]$$6

Because the one-dimensional $$\mu_{k+1}\in\arg\min_\mu \bigl[SW_2^2(\mu,\mu_k)/(2\tau)+\mathcal F(\mu)\bigr]$$7 admits the closed form

$$\mu_{k+1}\in\arg\min_\mu \bigl[SW_2^2(\mu,\mu_k)/(2\tau)+\mathcal F(\mu)\bigr]$$8

computable by sorting, the Monte-Carlo approximation of $$\mu_{k+1}\in\arg\min_\mu \bigl[SW_2^2(\mu,\mu_k)/(2\tau)+\mathcal F(\mu)\bigr]$$9 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

$\rho_\theta=(T_\theta)_\sharp \rho_{\rm ref}$0

over $\rho_\theta=(T_\theta)_\sharp \rho_{\rm ref}$1 outer steps (Bonet et al., 2021).

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 $\rho_\theta=(T_\theta)_\sharp \rho_{\rm ref}$2, then converted by Fenchel–Rockafellar duality into a stochastic program over dual potentials $\rho_\theta=(T_\theta)_\sharp \rho_{\rm ref}$3, which are restricted to an RKHS. A representer theorem reduces the optimization to finite-dimensional coefficients, and the resulting scheme is explicitly discretization-free (Frogner et al., 2018).

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

$\rho_\theta=(T_\theta)_\sharp \rho_{\rm ref}$4

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 $\rho_\theta=(T_\theta)_\sharp \rho_{\rm ref}$5 and $\rho_\theta=(T_\theta)_\sharp \rho_{\rm ref}$6 for the reported implementations (Li et al., 2023).

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

$\rho_\theta=(T_\theta)_\sharp \rho_{\rm ref}$7

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 (Liu et al., 13 Apr 2026).

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}$8

with $\rho_\theta=(T_\theta)_\sharp \rho_{\rm ref}$9 an orientation-preserving diffeomorphism. The Wasserstein gradient flow on densities then lifts to an $$\sum_t\log(P_t\{\widehat\xi_t\})-\lambda\sum_t W(P_t,P_{t+1})$$0 gradient flow on maps, and restriction to a finite-dimensional family induces the pull-back metric

$$\sum_t\log(P_t\{\widehat\xi_t\})-\lambda\sum_t W(P_t,P_{t+1})$$1

The parameter update is the natural gradient flow

$$\sum_t\log(P_t\{\widehat\xi_t\})-\lambda\sum_t W(P_t,P_{t+1})$$2

or, discretely,

$$\sum_t\log(P_t\{\widehat\xi_t\})-\lambda\sum_t W(P_t,P_{t+1})$$3

This parameterized Wasserstein gradient flow is entirely mesh-free and preserves the unit-mass constraint without normalization (Zhang et al., 13 Mar 2026).

For numerical realization, the one-dimensional transport map may be represented by a boundary-preserving Neural ODE,

$$\sum_t\log(P_t\{\widehat\xi_t\})-\lambda\sum_t W(P_t,P_{t+1})$$4

where the boundary factor enforces $$\sum_t\log(P_t\{\widehat\xi_t\})-\lambda\sum_t W(P_t,P_{t+1})$$5 and guarantees that $$\sum_t\log(P_t\{\widehat\xi_t\})-\lambda\sum_t W(P_t,P_{t+1})$$6 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 $$\sum_t\log(P_t\{\widehat\xi_t\})-\lambda\sum_t W(P_t,P_{t+1})$$7 Sobolev gradient flow, reducing the initial energy gap by a factor of $$\sum_t\log(P_t\{\widehat\xi_t\})-\lambda\sum_t W(P_t,P_{t+1})$$8 in 2D and $$\sum_t\log(P_t\{\widehat\xi_t\})-\lambda\sum_t W(P_t,P_{t+1})$$9 in 3D compared to trivial initial conditions (Zhang et al., 13 Mar 2026).

A Hamiltonian analogue replaces gradient descent by a reduced Hamiltonian system on parameter space. In parameterized Wasserstein Hamiltonian flow, the push-forward manifold $\mathcal P_2(\mathbb R^d)$0 inherits a pull-back metric $\mathcal P_2(\mathbb R^d)$1, and the finite-dimensional Hamiltonian is

$\mathcal P_2(\mathbb R^d)$2

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 (Wu et al., 2023).

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

$\mathcal P_2(\mathbb R^d)$3

the associated probability-flow ODE is

$\mathcal P_2(\mathbb R^d)$4

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 (Gao et al., 2024).

Flow matching provides a second deterministic-transport viewpoint. Given any coupling $\mathcal P_2(\mathbb R^d)$5, linear interpolation of particles,

$\mathcal P_2(\mathbb R^d)$6

induces intermediate measures $\mathcal P_2(\mathbb R^d)$7 and a velocity field

$\mathcal P_2(\mathbb R^d)$8

A neural vector field $\mathcal P_2(\mathbb R^d)$9 can then be trained by regressing $$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),$$0 to $$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),$$1 on interpolated samples. This provides the basis for flow matching in Euclidean Wasserstein space and its generalization to spaces of measures (Piening et al., 8 May 2026).

The Wasserstein-on-Wasserstein extension lifts the geometry from $$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),$$2 to

$$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),$$3

with outer 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),$$4

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 $$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),$$5 per gradient step in the reported formulation, so scalable surrogates are introduced: independent couplings, exact minibatch OT, sliced-Wasserstein couplings with complexity $$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),$$6, and lazy linear Wasserstein couplings with cost $$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),$$7 (Piening et al., 8 May 2026).

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

$$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),$$8

admits monotone descent, convergence to stationary points, stability with respect to perturbations of the inputs, and resolution consistency as the number of atoms increases (You, 14 Sep 2025).

In nonstationary distribution estimation, the objective is not a transport PDE but a penalized likelihood,

$$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),$$9

with $\rho_t$0 the first-order Wasserstein distance. A support-reduction argument shows that each $\rho_t$1 can be taken discrete on observed points, and the problem reduces to a convex network-flow problem with $\rho_t$2 variables and $\rho_t$3 flow-conservation constraints. The reported numerical tests include dairy-commodity price forecasting and risk-averse portfolio optimization (Anderson et al., 8 Jul 2025).

Filtering and inference provide another important application domain. Variational Gaussian filtering via Wasserstein gradient flows approximates the filtering posterior by minimizing $\rho_t$4 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 (Corenflos et al., 2023). 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 (Frogner et al., 2018).

Several theoretical and practical limitations recur. Some guarantees require strong log-concavity, accurate score estimates, or smoothness assumptions on intermediate marginals (Gao et al., 2024). 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 (Bonet et al., 2021). GenWGP assumes regularity conditions such as densities bounded away from $\rho_t$5 and $\rho_t$6 smoothness of the flow (Liu et al., 13 Apr 2026). 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.