The gradient flow of a function over the space of probability densities with respect to the Wasserstein metric often exhibits nice properties and has been utilized in several machine learning applications. The standard approach to compute the Wasserstein gradient flow is the finite difference which discretizes the underlying space over a grid, and is not scalable. In this work, we propose a scalable proximal gradient type algorithm for Wasserstein gradient flow. The key of our method is a variational formulation of the objective function, which makes it possible to realize the JKO proximal map through a primal-dual optimization. This primal-dual problem can be efficiently solved by alternatively updating the parameters in the inner and outer loops. Our framework covers all the classical Wasserstein gradient flows including the heat equation and the porous medium equation. We demonstrate the performance and scalability of our algorithm with several numerical examples.
翻译:瓦西斯坦度量值的概率密度空间上的一个函数的梯度流,在瓦西斯坦度值的概率密度空间上,其特性往往很好,并被若干机器学习应用程序所利用。计算瓦西斯坦度梯度流的标准方法是有限的差异,它将基底空间分散在一个网格上,无法缩放。在这项工作中,我们为瓦西斯坦度梯度流提出了一个可缩放的准梯度型算法。我们方法的关键是目标函数的变式配方,它使得通过原始的双优化实现 JKO 准十进制地图成为可能。通过更新内环和外环的参数,可以有效地解决这一原始的双重问题。我们的框架涵盖了所有古典瓦西斯坦度梯度流,包括热方程和多孔介质方程。我们用几个数字实例来展示我们的算法的性能和可缩放性。