Graph matching relaxations speed up supervised graph prediction training

Graph Matching Relaxations and Amortization for Supervised Graph Prediction

Machine Learning

Summary

Predicting graphs from data is tricky because graphs can list their nodes in any order, so you need to compare predicted graphs with target graphs in a way that ignores these orders. This is usually slow because it requires solving a complicated matching problem for every example. The authors studied three ways to relax this task mathematically and found one approach, called Gromov-Wasserstein, works best. To make training faster, they propose teaching another model to predict the matching, so this costly step doesn't have to be done repeatedly. They tested their method on simple and real-world tasks, including a new one that predicts chemical scaffolds from mass spectrometry data.

What this means in practice

  • For chemical informatics teams: Predict molecular scaffold graphs directly from mass spectrometry data for improved chemical analysis workflows.
  • For machine learning engineers: Train graph prediction models efficiently by amortizing graph matching computations, reducing training time on datasets with unordered nodes.

Authors

Federico Méndez, Paul Krzakala, Gabriel Melo, Charlotte Laclau, Rémi Flamary, Florence d'Alché-Buc

Abstract

End-to-end Supervised Graph Prediction (SGP) requires a permutation-invariant loss to compare predicted and target graphs with arbitrary node orderings. Such losses typically involve a costly graph-matching problem. We first study three Optimal Transport relaxations of this problem and show, theoretically and empirically, that the Gromov-Wasserstein (GW) objective is the most suitable for SGP. Then, to avoid solving the resulting inner optimization for every training example, we propose to amortize the graph matching (node alignment) problem. For each training sample, the loss function leverages a transport plan provided by a parametric matcher based on the differentiable Sinkhorn algorithm applied on empirical node distributions. The graph prediction module and the matcher are jointly learned. We showcase the efficiency of this approach on toy and real world SGP problems of increasing complexity including a novel Mass-spectra to Scaffold task that we introduce.
View PDFOpen arXiv