Scalable unsupervised alignment of metric and nonmetric structures

Sanketh Vedula, Valentino Maiorca, Lorenzo Basile, Francesco Locatello, Alexander Bronstein

Keywords: unpaired alignment, single-cell multiomics, Gromov-Wasserstein, optimal transport

TL;DR: We propose a new framework for solving the Gromov-Wasserstein problem; the proposed is inductive, thus scable, and it produces high-quality assignments on single-cell multiomic alignment.

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.

Submission Number: 122