Learning High Dimensional Wasserstein Geodesics

TL;DR

This paper introduces a novel deep learning approach to compute Wasserstein geodesics in high-dimensional spaces, enabling efficient sampling, distance computation, and optimal transport map estimation.

Contribution

It formulates a minimax problem using Lagrange multipliers for high-dimensional Wasserstein geodesics and proposes a neural network-based bidirectional learning algorithm.

Findings

Successfully computes Wasserstein geodesics in high dimensions

Enables sampling from the Wasserstein geodesic

Accurately estimates Wasserstein distance and OT map

Abstract

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.