Revisiting invariances and introducing priors in Gromov-Wasserstein distances

TL;DR

This paper introduces Augmented Gromov-Wasserstein, a new distance measure that controls invariance flexibility and incorporates feature alignments, improving data comparison in machine learning and biological applications.

Contribution

It proposes a novel optimal transport distance that balances invariance and feature alignment, with theoretical analysis and practical demonstrations.

Findings

Enhanced alignment in single-cell multi-omic data

Improved transfer learning performance

Theoretical properties of the new metric

Abstract

Gromov-Wasserstein distance has found many applications in machine learning due to its ability to compare measures across metric spaces and its invariance to isometric transformations. However, in certain applications, this invariance property can be too flexible, thus undesirable. Moreover, the Gromov-Wasserstein distance solely considers pairwise sample similarities in input datasets, disregarding the raw feature representations. We propose a new optimal transport-based distance, called Augmented Gromov-Wasserstein, that allows for some control over the level of rigidity to transformations. It also incorporates feature alignments, enabling us to better leverage prior knowledge on the input data for improved performance. We present theoretical insights into the proposed metric. We then demonstrate its usefulness for single-cell multi-omic alignment tasks and a transfer learning scenario in machine learning.