Scalable unsupervised alignment of general metric and non-metric structures

TL;DR

This paper introduces a scalable method for aligning complex data structures across domains by transforming the challenging quadratic assignment problem into a more manageable linear assignment problem, applicable to both metric and non-metric data.

Contribution

It proposes a novel scalable framework that learns a linear assignment problem to approximate the Gromov-Wasserstein alignment, extending to non-metric dissimilarities with differentiable ranks.

Findings

Achieves state-of-the-art performance on synthetic and real datasets.

Extends alignment framework to non-metric dissimilarities.

Demonstrates computational efficiency and simplicity.

Abstract

Aligning data from different domains is a fundamental problem in machine learning with broad applications across very different areas, most notably aligning experimental readouts in single-cell multiomics. Mathematically, this problem can be formulated as the minimization of disagreement of pair-wise quantities such as distances and is related to the Gromov-Hausdorff and Gromov-Wasserstein distances. Computationally, it is a quadratic assignment problem (QAP) that is known to be NP-hard. Prior works attempted to solve the QAP directly with entropic or low-rank regularization on the permutation, which is computationally tractable only for modestly-sized inputs, and encode only limited inductive bias related to the domains being aligned. We consider the alignment of metric structures formulated as a discrete Gromov-Wasserstein problem and instead of solving the QAP directly, we propose to learn a related well-scalable linear assignment problem (LAP) whose solution is also a minimizer of the QAP. We also show a flexible extension of the proposed framework to general non-metric dissimilarities through differentiable ranks. We extensively evaluate our approach on synthetic and real datasets from single-cell multiomics and neural latent spaces, achieving state-of-the-art performance while being conceptually and computationally simple.