Hierarchical Refinement: Optimal Transport to Infinity and Beyond

Optimal transport (OT) has enjoyed great success in machine learning as a principled way to align datasets via a least-cost correspondence, driven in large part by the runtime efficiency of the Sinkhorn algorithm (Cuturi, 2013). However, Sinkhorn has quadratic space complexity in the number of points, limiting scalability to larger datasets. Low-rank OT achieves linear-space complexity, but by definition, cannot compute a one-to-one correspondence between points. When the optimal transport problem is an assignment problem between datasets then an optimal mapping, known as the Monge map, is guaranteed to be a bijection. In this setting, we show that the factors of an optimal low-rank coupling co-cluster each point with its image under the Monge map. We leverage this invariant to derive an algorithm, Hierarchical Refinement (HiRef), that dynamically constructs a multiscale partition of each dataset using low-rank OT subproblems, culminating in a bijective coupling. Hierarchical Refinement uses linear space and has log-linear runtime, retaining the space advantage of low-rank OT while overcoming its limited resolution. We demonstrate the advantages of Hierarchical Refinement on several datasets, including ones containing over a million points, scaling full-rank OT to problems previously beyond Sinkhorn's reach.

Imagine trying to match one million points from a cube to one million points on a sphere in the most efficient way possible; e.g. by minimizing the distance between matched points. One solution would be to compute distances between every pair of points and then to find the best matching. However, this requires evaluating a trillion possibilities, and becomes challenging even for the most powerful computers.To solve this matching problem at scale, we use the following strategy. First, we find a way to “cut” both shapes together so that two halves of each shape are on the same side of this cut. This is done in a way that minimizes the effort to transform one half to the other, where we discard any connections that cross the cut. We repeat this recursively, refining across a hierarchy of cuts, until we reach a one-to-one mapping between the points.This divide-and-conquer strategy, which we call hierarchical refinement, is guided by the recent technique of low-rank optimal transport which helps us find the most efficient way to make each cut. Importantly, this technique is extensible to any type of space with a notion of “distance” and allows for highly flexible, even non-linear, ways of splitting the space.We show that, given an optimal way to make each “cut,” hierarchical refinement can recover the best possible matching between any two datasets, such as physical shapes, images, or biological data like cells.