Neural Monge Map estimation and its applications
This work addresses the need for efficient and flexible optimal transport map estimation in machine learning applications, though it is incremental as it builds on existing neural network-based methods.
The authors tackled the problem of estimating optimal transport maps between probability distributions using neural networks, presenting a scalable algorithm that works with samples from marginals, accommodates different dimensions, and supports general cost functions, achieving competitive performance in text-to-image generation and image inpainting tasks.
Monge map refers to the optimal transport map between two probability distributions and provides a principled approach to transform one distribution to another. Neural network based optimal transport map solver has gained great attention in recent years. Along this line, we present a scalable algorithm for computing the neural Monge map between two probability distributions. Our algorithm is based on a weak form of the optimal transport problem, thus it only requires samples from the marginals instead of their analytic expressions, and can accommodate optimal transport between two distributions with different dimensions. Our algorithm is suitable for general cost functions, compared with other existing methods for estimating Monge maps using samples, which are usually for quadratic costs. The performance of our algorithms is demonstrated through a series of experiments with both synthetic and realistic data, including text-to-image generation and image inpainting tasks.