[Paper Review] Entropic Wasserstein Gradient Flows

Gabriel Peyré|arXiv (Cornell University)|Feb 22, 2015

Geometric Analysis and Curvature Flows48 references5 citations

TL;DR

This paper introduces an entropic regularization approach to accelerate the computation of Wasserstein gradient flows, replacing the computationally expensive Wasserstein distance with a Kullback-Leibler divergence via entropic smoothing. The method enables fast, parallelizable numerical schemes using Dykstra’s algorithm and linear filtering (e.g., Gaussian convolution), achieving near-linear time complexity on grids and efficient heat diffusion approximations on complex domains.

ABSTRACT

This article details a novel numerical scheme to approximate gradient flows for optimal transport (i.e. Wasserstein) metrics. These flows have proved useful to tackle theoretically and numerically non-linear diffusion equations that model for instance porous media or crowd evolutions. These gradient flows define a suitable notion of weak solutions for these evolutions and they can be approximated in a stable way using discrete flows. These discrete flows are implicit Euler time stepping according to the Wasserstein metric. A bottleneck of these approaches is the high computational load induced by the resolution of each step. Indeed, this corresponds to the resolution of a convex optimization problem involving a Wasserstein distance to the previous iterate. Following several recent works on the approximation of Wasserstein distances, we consider a discrete flow induced by an entropic regularization of the transportation coupling. This entropic regularization allows one to trade the initial Wasserstein fidelity term for a Kulback-Leibler divergence, which is easier to deal with numerically. We show how KL proximal schemes, and in particular Dykstra's algorithm, can be used to compute each step of the regularized flow. The resulting algorithm is both fast, parallelizable and versatile, because it only requires multiplications by a Gibbs kernel. On Euclidean domains discretized on an uniform grid, this corresponds to a linear filtering (for instance a Gaussian filtering when $c$ is the squared Euclidean distance) which can be computed in nearly linear time. On more general domains, such as (possibly non-convex) shapes or on manifolds discretized by a triangular mesh, following a recently proposed numerical scheme for optimal transport, this Gibbs kernel multiplication is approximated by a short-time heat diffusion.

Motivation & Objective

• Address the high computational cost of solving each step in implicit Euler time discretization of Wasserstein gradient flows.
• Develop a numerically efficient alternative to exact Wasserstein distance computation in gradient flow schemes.
• Enable scalable and stable numerical approximation of non-linear diffusion and crowd motion PDEs on complex domains.
• Leverage entropic regularization to transform the Wasserstein fidelity term into a Kullback-Leibler divergence for improved numerical tractability.
• Design a versatile, parallelizable algorithm that only requires matrix multiplications with the Gibbs kernel $e^{-c/ u}$

Proposed method

• Apply entropic regularization to the optimal transport problem in each step of the JKO (Jordan-Kinderlehrer-Otto) gradient flow scheme, replacing the Wasserstein distance with a Kullback-Leibler divergence.
• Use first-order proximal schemes based on the KL divergence to compute each time step, enabling efficient optimization via Dykstra’s algorithm.
• Represent the Gibbs kernel $e^{-c/ u}$ as a linear filter (e.g., Gaussian) on uniform grids, allowing nearly linear-time computation via fast convolution.
• For non-uniform or manifold-discretized domains, approximate the Gibbs kernel multiplication via short-time heat diffusion, leveraging existing numerical schemes.
• Formulate the proximal operator for composite functions using weighted geometric means of probability measures, enabling modular computation across multiple marginals.
• Derive closed-form expressions for proximal operators under various functional forms (e.g., sum, product, marginalization), ensuring computational efficiency

Experimental results

Research questions

• RQ1 Can entropic regularization be effectively applied to accelerate the solution of Wasserstein gradient flows in a stable and numerically efficient manner?
• RQ2 How can the computational bottleneck of solving Wasserstein distance subproblems be alleviated while preserving convergence to weak solutions of PDEs?
• RQ3 To what extent can the Gibbs kernel multiplication be approximated efficiently on non-Euclidean domains such as triangular meshes or manifolds?
• RQ4 Can the resulting algorithm maintain parallelizability and versatility across diverse applications like crowd motion and non-linear diffusion?
• RQ5 What is the impact of entropic regularization on the accuracy and convergence rate of the gradient flow approximation?

Key findings

• The entropic regularization transforms the Wasserstein distance subproblem into a Kullback-Leibler divergence minimization, which is numerically more tractable than the original linear program.
• The resulting algorithm achieves nearly linear time complexity on uniform grids due to the use of fast linear filtering (e.g., Gaussian convolution) for Gibbs kernel operations.
• On complex domains such as non-convex shapes or triangular meshes, the Gibbs kernel multiplication is approximated via short-time heat diffusion, preserving accuracy and stability.
• The method enables efficient computation of JKO-type gradient flows for non-linear PDEs, including porous media and spatially-varying diffusion equations.
• Numerical experiments demonstrate the method’s effectiveness in simulating crowd motion on intricate geometries and non-linear diffusion with variable coefficients.
• The use of Dykstra’s algorithm for KL-based proximal schemes ensures robust and scalable convergence across all tested configurations.