Score-based Generative Neural Networks for Large-Scale Optimal Transport
This addresses the computational challenge of optimal transport for large, high-dimensional datasets, though it is an incremental improvement by applying a novel method to a known bottleneck.
The paper tackles the problem of sampling from the optimal transport coupling between source and target distributions by introducing a score-based generative model that learns the Sinkhorn coupling, demonstrating empirical success on large-scale tasks.
We consider the fundamental problem of sampling the optimal transport coupling between given source and target distributions. In certain cases, the optimal transport plan takes the form of a one-to-one mapping from the source support to the target support, but learning or even approximating such a map is computationally challenging for large and high-dimensional datasets due to the high cost of linear programming routines and an intrinsic curse of dimensionality. We study instead the Sinkhorn problem, a regularized form of optimal transport whose solutions are couplings between the source and the target distribution. We introduce a novel framework for learning the Sinkhorn coupling between two distributions in the form of a score-based generative model. Conditioned on source data, our procedure iterates Langevin Dynamics to sample target data according to the regularized optimal coupling. Key to this approach is a neural network parametrization of the Sinkhorn problem, and we prove convergence of gradient descent with respect to network parameters in this formulation. We demonstrate its empirical success on a variety of large scale optimal transport tasks.