Transport Clustering: Solving Low-Rank Optimal Transport via Clustering

Optimal transport (OT) finds a least-cost transport plan between two probability distributions. Unlike standard OT, which infers unstructured pointwise mappings, low-rank optimal transport explicitly constrains the rank of the transport plan to infer latent structure. This improves statistical robustness, yields sharper parametric rates for estimating Wasserstein distances, and generalizes $K$-means to co-clustering. However, these advantages come at the cost of a non-convex and NP-hard optimization problem. We introduce Transport Clustering, an algorithm to compute a low-rank OT plan that reduces low-rank OT to a clustering problem on correspondences obtained from a full-rank transport registration step. We prove that this reduction yields polynomial-time, constant-factor approximation algorithms for low-rank OT: specifically, a $(1+\gamma)$ approximation for negative-type metrics, a $(1+\gamma+\sqrt{2\gamma}\,)$ approximation for kernel costs, and a $(1+\gamma+\rho)$ approximation for general metrics satisfying the triangle inequality. Here, $\gamma \in [0,1]$ is the cost ratio of the optimal full-rank to low-rank solutions, and $\rho\in[0,1]$ is an asymmetry coefficient on the cluster variances. Numerically, Transport Clustering outperforms existing solvers on synthetic benchmarks and large-scale datasets.

Optimal transport (OT) finds the most efficient mapping or transport plan between two sets of points for a given notion of cost, like finding the cheapest way to ship goods from a factory to a consumer. Standard OT maps each individual source point directly to its target. Low-rank OT, instead, assumes this transport happens through a small number of intermediate hubs. For instance, rather than a factory shipping millions of items directly to millions of consumers, Low-Rank OT routes them to a small handful of regional warehouses, which then efficiently redistribute the goods. While powerful, solving this problem is highly difficult. In this work, we prove Low-Rank OT is mathematically equivalent to generalized K-means clustering given the solution to the easier problem of finding the full-rank plan. We introduce the algorithm of Transport Clustering, which breaks the problem into two steps: first, we find the most efficient direct correspondences between two sets of points using standard OT; second, we cluster those correspondences using K-means. This equivalence allows Low-Rank OT to be solved using highly optimized algorithms for clustering with strict theoretical guarantees on their performance.