Statistical Optimal Transport via Factored Couplings

TL;DR

This paper introduces a novel regularization approach for estimating Wasserstein distances in high-dimensional spaces, leveraging low transport rank couplings to improve accuracy and computational efficiency.

Contribution

It proposes a new structural regularization based on low transport rank couplings, enhancing high-dimensional optimal transport estimation beyond traditional methods.

Findings

Significant improvement in high-dimensional domain adaptation tasks.

Theoretical evidence that transport rank mitigates curse of dimensionality.

Effective estimation of Wasserstein distances from sample data.

Abstract

We propose a new method to estimate Wasserstein distances and optimal transport plans between two probability distributions from samples in high dimension. Unlike plug-in rules that simply replace the true distributions by their empirical counterparts, our method promotes couplings with low transport rank, a new structural assumption that is similar to the nonnegative rank of a matrix. Regularizing based on this assumption leads to drastic improvements on high-dimensional data for various tasks, including domain adaptation in single-cell RNA sequencing data. These findings are supported by a theoretical analysis that indicates that the transport rank is key in overcoming the curse of dimensionality inherent to data-driven optimal transport.