[论文解读] Entropic Wasserstein Gradient Flows

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

Geometric Analysis and Curvature Flows参考文献 48被引用 5

一句话总结

本文提出了一种基于熵正则化的加速方法，用于计算Wasserstein梯度流，通过熵平滑将计算成本高昂的Wasserstein距离替换为Kullback-Leibler散度。该方法利用Dykstra算法和线性滤波（例如高斯卷积）实现快速、可并行化的数值格式，在网格上实现近线性时间复杂度，并在复杂域上实现高效的热扩散近似。

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.

研究动机与目标

• 解决隐式欧拉时间离散化Wasserstein梯度流过程中每一步计算成本过高的问题。
• 为梯度流方案中精确Wasserstein距离计算开发一种数值高效的替代方法。
• 实现在复杂域上对非线性扩散和人群运动PDE的可扩展且稳定的数值近似。
• 利用熵正则化将Wasserstein保真项转化为Kullback-Leibler散度，以提升数值可处理性。
• 设计一种通用且可并行化的算法，仅需与Gibbs核 $e^{-c/ u}$ 的矩阵乘法操作。

提出的方法

• 在JKO（Jordan-Kinderlehrer-Otto）梯度流方案的每一步中对最优传输问题应用熵正则化，用Kullback-Leibler散度替代Wasserstein距离。
• 基于KL散度使用一阶邻近算法计算每一步时间推进，通过Dykstra算法实现高效优化。
• 在均匀网格上将Gibbs核 $e^{-c/ u}$ 表示为线性滤波器（例如高斯核），通过快速卷积实现近似线性时间计算。
• 对于非均匀或流形离散化的域，通过短时间热扩散近似Gibbs核乘法，利用现有数值格式。
• 通过概率测度的加权几何平均形式化复合函数的邻近算子，实现在多个局部变量上的模块化计算。
• 推导出在各种函数形式（例如和、积、边缘化）下的邻近算子的闭式表达式，确保计算效率。

实验结果

研究问题

• RQ1 熵正则化能否以稳定且数值高效的方式有效应用于加速Wasserstein梯度流的求解？
• RQ2 如何缓解求解Wasserstein距离子问题的计算瓶颈，同时保持对PDE弱解的收敛性？
• RQ3 在非欧几里得域（如三角形网格或流形）上，Gibbs核乘法在多大程度上可被高效近似？
• RQ4 所得到的算法能否在人群运动和非线性扩散等多样化应用中保持可并行性和通用性？
• RQ5 熵正则化对梯度流近似方法的精度和收敛速度有何影响？

主要发现

• 熵正则化将Wasserstein距离子问题转化为Kullback-Leibler散度最小化问题，相比原始线性规划更具数值可处理性。
• 由于利用快速线性滤波（如高斯卷积）实现Gibbs核操作，该算法在均匀网格上实现了近似线性时间复杂度。
• 在复杂域（如非凸形状或三角形网格）上，通过短时间热扩散近似Gibbs核乘法，保持了精度和稳定性。
• 该方法实现了对非线性PDE（包括多孔介质方程和空间变系数扩散方程）的JKO型梯度流的高效计算。
• 数值实验表明，该方法在复杂几何结构上模拟人群运动以及变系数非线性扩散方面均表现出高效性。
• 使用Dykstra算法求解基于KL的邻近方案，确保了在所有测试配置下鲁棒且可扩展的收敛性。