A Combinatorial Algorithm for the Semi-Discrete Optimal Transport Problem

Optimal Transport (OT, also known as the Wasserstein distance) is a popular metric for comparing probability distributions and has been successfully used in many machine-learning applications. In the semi-discrete 2-Wasserstein problem, we wish to compute the cheapest way to transport all the mass from a continuous distribution µ to a discrete distribution ν in ℝ^d for d ≥ 1, where the cost of transporting unit mass between points a and b is d(a, b) = ∥a - b∥². When both distributions are discrete, a simple combinatorial framework has been used to find the exact solution (see e.g. [Orlin, STOC 1988]). In this paper, we propose a combinatorial framework for the semi-discrete OT, which can be viewed as an extension of the combinatorial framework for the discrete OT but requires several new ideas. We present a new algorithm that given µ and ν in ℝ² and a parameter ε > 0, computes an ε-additive approximate semi-discrete transport plan in O(n⁴ log n log 1/ε) time (in the worst case), where n is the support-size of the discrete distribution ν and we assume that the mass of µ inside a triangle can be computed in O(1) time. Our algorithm is significantly faster than the known algorithms, and unlike many numerical algorithms, it does not make any assumptions on the smoothness of µ. As an application of our algorithm, we describe a data structure to store a large discrete distribution µ (with support size N) using O(N) space so that, given a query discrete distribution ν (with support size k), an ε-additive approximate transport plan can be computed in O(k³√N log 1/ε) time in 2 dimensions. Our algorithm and data structure extend to higher dimensions as well as to p-Wasserstein problem for any p ≥ 1.

Duke Scholars

Author Pankaj K. Agarwal Computer Science