Unbalanced Low-rank Optimal Transport Solvers

TL;DR

This paper introduces scalable, unbalanced low-rank optimal transport algorithms that combine computational efficiency with modeling flexibility, demonstrated on spatial transcriptomics data.

Contribution

It merges low-rank and unbalanced OT methods, creating versatile solvers that are both scalable and less rigid, applicable to complex real-world problems.

Findings

Achieves linear-time complexity for low-rank OT

Effectively handles outliers with unbalanced OT

Demonstrates practical relevance on spatial transcriptomics

Abstract

The relevance of optimal transport methods to machine learning has long been hindered by two salient limitations. First, the $O(n^3)$ computational cost of standard sample-based solvers (when used on batches of $n$ samples) is prohibitive. Second, the mass conservation constraint makes OT solvers too rigid in practice: because they must match \textit{all} points from both measures, their output can be heavily influenced by outliers. A flurry of recent works in OT has addressed these computational and modelling limitations, but has resulted in two separate strains of methods: While the computational outlook was much improved by entropic regularization, more recent $O(n)$ linear-time \textit{low-rank} solvers hold the promise to scale up OT further. On the other hand, modelling rigidities have been eased owing to unbalanced variants of OT, that rely on penalization terms to promote, rather than impose, mass conservation. The goal of this paper is to merge these two strains, to achieve the promise of \textit{both} versatile/scalable unbalanced/low-rank OT solvers. We propose custom algorithms to implement these extensions for the linear OT problem and its Fused-Gromov-Wasserstein generalization, and demonstrate their practical relevance to challenging spatial transcriptomics matching problems.