Optimal Transport Graph Neural Networks

TL;DR

OT-GNN introduces a novel graph embedding method using optimal transport and prototypes, capturing structural information more effectively than traditional aggregation, with theoretical guarantees and improved empirical performance.

Contribution

The paper proposes OT-GNN, combining optimal transport with parametric prototypes for graph embeddings, providing universal approximation guarantees and addressing optimization issues.

Findings

Outperforms existing methods on molecular property prediction tasks.

Produces smoother and more informative graph representations.

Addresses collapse issues with a noise contrastive regularizer.

Abstract

Current graph neural network (GNN) architectures naively average or sum node embeddings into an aggregated graph representation -- potentially losing structural or semantic information. We here introduce OT-GNN, a model that computes graph embeddings using parametric prototypes that highlight key facets of different graph aspects. Towards this goal, we successfully combine optimal transport (OT) with parametric graph models. Graph representations are obtained from Wasserstein distances between the set of GNN node embeddings and ``prototype'' point clouds as free parameters. We theoretically prove that, unlike traditional sum aggregation, our function class on point clouds satisfies a fundamental universal approximation theorem. Empirically, we address an inherent collapse optimization issue by proposing a noise contrastive regularizer to steer the model towards truly exploiting the OT geometry. Finally, we outperform popular methods on several molecular property prediction tasks, while exhibiting smoother graph representations.