A Constrained-JKO Scheme for Effective and Efficient Wasserstein Gradient Flows
摘要
We present an innovative Wasserstein Gradient Flow algorithm for solving optimization problems over the space of probability densities. Differently form other state of the art methods, the proposed approach does not involve any computationally expensive training of neural networks. The main contribution is a novel parametrization of the transport map that allows to recast the widely adopted JKO schema from a scalarized bi-objective to a constrained optimization problem. This provides three relevant advantages: a better control of the desired quality of approximation of the optimal transport map; convergence to the optimum without affecting the quality of the transport map; and a significantly lower computational cost (i.e., runtime-per-iteration). Experimental results on a set of well diversified test cases provide the empirical evidence of the advantages offered by the proposed approach.