Diffusion-Regularized Wasserstein Gradient Flow

• DWGF are gradient flow models on Wasserstein space that incorporate entropic diffusion regularization for enhanced stability and convergence.
• They originate as scaling limits of entropic optimal transport algorithms, such as the Sinkhorn method, linking discrete iterations to continuous dynamics.
• DWGF underpin novel computational strategies in optimization and sampling, leveraging mirror descent geometry in high-dimensional probability measures.

Diffusion-regularized Wasserstein Gradient Flow (DWGF) refers to a family of gradient flow formulations and computational methods that combine the geometry of optimal transport (Wasserstein space) with diffusion (entropic or stochastic) regularization. These flows arise naturally in the scaling limits of entropic optimal transport algorithms (such as Sinkhorn), in stochastic particle models converging to macroscopic evolution equations, and in theoretical formulations bridging discrete iterative procedures and continuous-time stochastic or deterministic dynamics. DWGF provides the foundation for understanding, analyzing, and simulating entropic regularization phenomena in high-dimensional transport and for designing algorithms with robust convergence and stability properties, especially in the context of optimization and sampling over probability measures.

1. Definition and Conceptual Framework

A DWGF is a gradient flow on the Wasserstein space of probability measures where the steepest descent trajectory is regularized (typically by an entropy or diffusion term). The flow is not the classical Wasserstein gradient flow of a functional $\mathcal{F}(\rho)$ but instead is often derived as the scaling limit of entropic-regularized transport problems or as the macroscopic evolution of stochastic particle systems.

In the canonical setting, consider the evolution of a probability density $\rho_t$ governed by the continuity equation

$\partial_t \rho_t + \mathrm{div}(\rho_t v_t) = 0,$

where the velocity field $v_t$ is not the classical Wasserstein gradient, but is often modified either by dual coordinates (mirror geometry) or by regularization. Notably, in the limit of small regularization parameter, the discrete Sinkhorn algorithm gives rise to a continuous DWGF, called the Sinkhorn flow, which is a mirror gradient flow with respect to the relative entropy and the squared Wasserstein distance as mirror map ("Wasserstein Mirror Gradient Flow as the limit of the Sinkhorn Algorithm" (Deb et al., 2023)).

DWGF is thus characterized by:

• Regularization (e.g., entropy, noise, or smoothing) applied to the transport, yielding flows with improved stability, regularity, and convergence;
• A microscopic perspective: DWFGs arise naturally when considering the large-deviations rate functionals for empirical measures under independent diffusions, resulting in a variational structure involving both entropy and Wasserstein cost (Adams et al., 2010);
• Algorithmic significance: The continuous-time limit of regularized transport algorithms (e.g., Sinkhorn, IPFP) and projected Langevin methods for entropically regularized couplings are DWGFs.

2. Mathematical Formulation: Sinkhorn Flow and Mirror Gradient Structure

A central example of DWGF is the Sinkhorn flow—a Wasserstein mirror gradient flow emerging as the scaling limit of Sinkhorn iterations when the entropic regularization parameter $\varepsilon \to 0$ and the number of iterations is scaled as $1/\varepsilon$. Define the squared Wasserstein distance to a reference measure $\nu = e^{-g}$ as the mirror function: $U(\rho) = \frac{1}{2} W_2^2(\rho, \nu).$ Given a convex (energy) functional $F(\rho)$, the mirror gradient flow in Wasserstein space is governed by

$$\partial_t \rho_t + \mathrm{div}\left(\rho_t v_t\right) = 0,$$

where the velocity is given by

$$v_t(x) = - \left(\nabla^2 u_t(x)\right)^{-1} \nabla_x \left( f(x) + \log \rho_t(x) \right),$$

with $u_t$ the Brenier potential (the convex function whose gradient pushes $\rho_t$ to $\nu$), and $f$ the potential corresponding to the energy functional.

Equivalently, the flow can be written for the potential $u_t$ as a parabolic Monge-Ampère equation: $$\partial_t u_t(x) = f(x) - g(x^{(u_t)}) + \frac{\partial}{\partial x^{(u_t)}} u_t(x),$$ where $x^{(u_t)}$ denotes the "mirror variable" defined by the optimal transport map.

A crucial property of the Sinkhorn flow is that the (metric) speed of the flow measured in the linearized optimal transport (LOT) metric is given by

$$\lim_{\delta \to 0} \frac{1}{\delta} \mathrm{LOT}(\nu, \rho_{t+\delta}, \rho_t) = \| v_t \|_{L^2(\rho_t)},$$

indicating that the flow's regularity and convergence are captured in this mirror-geometric sense.

3. Diffusion-Regularized Dynamics and McKean–Vlasov Processes

DWGFs admit stochastic process representations, notably as McKean–Vlasov diffusions whose time marginals follow the mirror flow. The stochastic differential equation for the Sinkhorn diffusion is

$$dX_t = \left[ -\frac{\partial f}{\partial u_t}(X_t) - \frac{\partial g}{\partial u_t}(X_t^{(u_t)}) + \frac{\partial h_t}{\partial u_t}(X_t) \right] dt + \sqrt{2 \frac{\partial X_t}{\partial X_t^{(u_t)}}} dB_t,$$

where $h_t = -\log \rho_t$ and $B_t$ is Brownian motion. This SDE has drift components that are "mirrored" via the Brenier map and a noise term regularized by the geometry induced by $u_t$.

In contrast to the classical Langevin diffusion (gradient flow of the relative entropy in Wasserstein space), the Sinkhorn diffusion corresponds to a mirror descent in probability space, with stochastic dynamics regularizing the movement along the mirrored geometry.

Under sufficient regularity conditions, the time-marginal distribution of the process $X_t$ solves the Sinkhorn PDE, thus establishing a direct connection between stochastic microscopic dynamics and the macroscopic DWGF.

4. Analytical and Computational Implications

DWGF delivers several analytical and algorithmic benefits:

• Entropic/Noise Regularization: The $\varepsilon$-regularized optimal transport (Sinkhorn) problem introduces smoothing and strong convexity, overcoming singularities and instability of classical transport maps;
• Convergence Behavior: In the scaling limit, the mirror flow may converge exponentially, or in certain cases with finite termination, depending on the choice of energies and geometry;
• Numerical Schemes: The underlying mirror structure of DWGF legitimizes energy-dissipation algorithms and JKO-type (minimizing movement) schemes with entropic regularization for faster convergence and better numerical stability;
• Sampling and Optimization: The construction of McKean–Vlasov diffusions governed by DWGF pathwise evolution enables the development of accelerated sampling methods and novel algorithms for variational inference by incorporating mirror/Wasserstein geometry.

Crucially, DWGF provides the rigorous limiting structure tying together discrete regularized transport algorithms (e.g., Sinkhorn, IPFP), continuous stochastic processes, and variational flows in Wasserstein or mirror-geometry-modified spaces.

5. Connection to Classical and Modern Optimal Transport

DWGFs generalize classical Wasserstein gradient flows by allowing the geometry (the "mirror map") to be induced by entropic or stochastic regularization. While the canonical Fokker–Planck equation is the Wasserstein gradient flow of the negative entropy,

$$\partial_t \rho_t = \nabla \cdot \left( \rho_t \nabla ( f + \log \rho_t ) \right ),$$

the DWGF (e.g., Sinkhorn flow) modifies the velocity field by composing with the inverse Hessian of the potential $u_t$, yielding

$$\partial_t \rho_t(x) + \mathrm{div} \left( \rho_t(x) (\nabla^2 u_t(x))^{-1} \nabla_x(f(x) + \log \rho_t(x)) \right ) = 0.$$

This difference alters both the convergence behavior and the geometry of the flow. For instance, in some cases, the mirror flow converges exponentially or even in finite time, contrary to the typical exponential or sublinear decay of entropy in classical gradient flows.

By identifying DWGF as the scaling limit of entropic optimal transport algorithms ("Sinkhorn algorithm"), the theory provides a geometric and probabilistic justification for the widespread empirical utility of Sinkhorn regularization and related computational methods in high-dimensional optimal transport and machine learning (Deb et al., 2023).

6. Applications and Future Directions

DWGF plays a central role in bridging discrete entropically regularized algorithms and continuous transport PDEs, offering:

• Rigorous geometric understanding of the continuous limit of Sinkhorn and IPFP iterates as gradient flows in mirrored Wasserstein geometry;
• A framework for designing sampling and optimization schemes in high-dimensional spaces using stochastic differential equations structured by mirror geometry;
• A foundation for developing and analyzing extensions of variational inference and learning via regularized transport geometry, with direct algorithmic implications for large-scale optimal transport, generative modeling, and Bayesian computation.

A promising future direction is exploiting the mirror geometry underlying DWGF for constructing accelerated or robust optimization and sampling algorithms that—unlike classical Langevin or Wasserstein gradient flows—harness the smoothing and geometric features of the mirror flow, particularly in the context of high-dimensional generative models and large-scale computational optimal transport.

---

In summary, diffusion-regularized Wasserstein gradient flows (DWGFs) are a critical linking concept in modern optimal transport theory, uniting entropic regularization in computation, stochastic process theory, and mirror descent geometry in probability spaces. The Sinkhorn flow, as a paradigmatic DWGF, captures the scaling limit of entropic OT algorithms and unlocks a new class of stochastic processes and algorithmic strategies exploiting mirror-regularized Wasserstein geometry for advanced optimization and sampling over probability measures.