Unsupervised Learning for Optimal Transport plan prediction between unbalanced graphs
Sonia Mazelet, Rémi Flamary, Bertrand Thirion
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.