Unsupervised Learning for Optimal Transport plan prediction between unbalanced graphs

Sonia Mazelet, Rémi Flamary, Bertrand Thirion

Keywords: Non-Convex Optimization, Unsupervised Learning, optimal transport, gnn

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.

Primary Area: Deep learning (e.g., architectures, generative models, optimization for deep networks, foundation models, LLMs)

Submission Number: 13169