🤖 AI Summary
To address the high computational cost and curse-of-dimensionality challenges in sampling optimal transport (OT) couplings for large-scale, high-dimensional data, this paper proposes an efficient learning framework based on score-based generative models. Specifically, conditioned on source samples, it iteratively generates target samples following the Sinkhorn-regularized OT coupling via Langevin dynamics. Crucially, it jointly parameterizes the score function and Sinkhorn potential functions—enabling, for the first time, end-to-end co-learning of score-based generation and OT coupling. We theoretically establish the convergence of gradient descent on the network parameters under mild assumptions. Experiments demonstrate that our method significantly improves both accuracy and speed of coupling estimation across diverse large-scale OT tasks, while maintaining scalability and practical applicability.
📝 Abstract
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.
