Unsupervised Learning for Optimal Transport plan prediction between unbalanced graphs

TL;DR

This paper introduces ULOT, a deep learning approach that rapidly predicts optimal transport plans between graphs, significantly reducing computation time and enabling differentiable optimization for graph comparison tasks.

Contribution

ULOT is a novel neural network architecture that predicts graph transport plans using FUGW loss, improving scalability and enabling differentiable graph comparisons.

Findings

ULOT predicts transport plans with up to 100x faster than classical methods.

Predicted plans serve as warm starts to accelerate classical solvers.

The method is effective on synthetic and real cortical surface graphs.

Abstract

Optimal transport between graphs, based on Gromov-Wasserstein and other extensions, is a powerful tool for comparing and aligning graph structures. However, solving the associated non-convex optimization problems is computationally expensive, which limits the scalability of these methods to large graphs. In this work, we present Unbalanced Learning of Optimal Transport (ULOT), a deep learning method that predicts optimal transport plans between two graphs. Our method is trained by minimizing the fused unbalanced Gromov-Wasserstein (FUGW) loss. We propose a novel neural architecture with cross-attention that is conditioned on the FUGW tradeoff hyperparameters. We evaluate ULOT on synthetic stochastic block model (SBM) graphs and on real cortical surface data obtained from fMRI. ULOT predicts transport plans with competitive loss up to two orders of magnitude faster than classical solvers. Furthermore, the predicted plan can be used as a warm start for classical solvers to accelerate their convergence. Finally, the predicted transport plan is fully differentiable with respect to the graph inputs and FUGW hyperparameters, enabling the optimization of functionals of the ULOT plan.