Unbalanced CO-Optimal Transport

TL;DR

This paper introduces unbalanced COOT, an extension of CO-optimal transport, which is robust to outliers and noise, improving alignment tasks in heterogeneous data scenarios.

Contribution

It proposes unbalanced COOT, providing the first theoretical proof of robustness to noise in incomparable spaces, enhancing alignment in complex datasets.

Findings

Unbalanced COOT is robust to outliers and noise.

Empirical results show improved alignment in domain adaptation.

Demonstrates effectiveness in single-cell data alignment.

Abstract

Optimal transport (OT) compares probability distributions by computing a meaningful alignment between their samples. CO-optimal transport (COOT) takes this comparison further by inferring an alignment between features as well. While this approach leads to better alignments and generalizes both OT and Gromov-Wasserstein distances, we provide a theoretical result showing that it is sensitive to outliers that are omnipresent in real-world data. This prompts us to propose unbalanced COOT for which we provably show its robustness to noise in the compared datasets. To the best of our knowledge, this is the first such result for OT methods in incomparable spaces. With this result in hand, we provide empirical evidence of this robustness for the challenging tasks of heterogeneous domain adaptation with and without varying proportions of classes and simultaneous alignment of samples and features across single-cell measurements.