Learning High Dimensional Wasserstein Geodesics
This work addresses the problem of computing Wasserstein geodesics in high dimensions, which is a challenge for researchers working with optimal transport and generative models.
This paper proposes a new formulation and learning strategy to compute the Wasserstein geodesic between two high-dimensional probability distributions. The method leverages deep neural networks and a sample-based bidirectional learning algorithm to find the geodesic, and as a result, it can also sample from the geodesic and compute the Wasserstein distance and optimal transport map.
We propose a new formulation and learning strategy for computing the Wasserstein geodesic between two probability distributions in high dimensions. By applying the method of Lagrange multipliers to the dynamic formulation of the optimal transport (OT) problem, we derive a minimax problem whose saddle point is the Wasserstein geodesic. We then parametrize the functions by deep neural networks and design a sample based bidirectional learning algorithm for training. The trained networks enable sampling from the Wasserstein geodesic. As by-products, the algorithm also computes the Wasserstein distance and OT map between the marginal distributions. We demonstrate the performance of our algorithms through a series of experiments with both synthetic and realistic data.