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.
Candelieri, A., Ponti, A., Archetti, F. (2025). A Constrained-JKO Scheme for Effective and Efficient Wasserstein Gradient Flows. In Learning and Intelligent Optimization 18th International Conference, LION 18, Ischia Island, Italy, June 9–13, 2024, Revised Selected Papers (pp.66-80). Springer Science and Business Media Deutschland GmbH [10.1007/978-3-031-75623-8_6].
A Constrained-JKO Scheme for Effective and Efficient Wasserstein Gradient Flows
Candelieri A.;Ponti A.;Archetti F.
2025
Abstract
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.I documenti in IRIS sono protetti da copyright e tutti i diritti sono riservati, salvo diversa indicazione.