Neural Optimal Transport with General Cost Functionals

Arip Asadulaev, Alexander Korotin, Vage Egiazarian, Evgeny Burnaev

Abstract: Neural optimal transport techniques mostly use Euclidean cost functions, such as $\ell^1$ or $\ell^2$. These costs are suitable for translation tasks between related domains, but they are hardly applicable to problems where a specific non-Euclidean optimality of the mapping is required such as dataset transfer. To tackle this issue, we introduce a novel neural network-based algorithm to compute optimal transport plans and maps for general cost functionals. Such functionals provide more flexibility and allow using auxiliary information, such as class labels, to construct the required transport map. Our method is based on a saddle point reformulation of the optimal transport problem and generalizes prior methods for weak and strong transport cost functionals. As an application, we construct a functional to map data distributions with preserving the class-wise structure of data.