Optimal Transport Graph Neural Networks

Gary Bécigneul, Octavian-Eugen Ganea, Benson Chen, Regina Barzilay, Tommi S. Jaakkola

Keywords: graph neural networks, optimal transport, molecular representations, molecular property prediction

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 are (to our knowledge) the first to successfully combine optimal transport 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 optimal transport geometry. Finally, we consistently report better generalization performance on several molecular property prediction tasks, while exhibiting smoother graph representations.

One-sentence Summary: We compute graph representations based on abstract prototypes that leverage optimal transport and graph neural networks.

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics

Community Implementations: [ 2 code implementations](https://www.catalyzex.com/paper/arxiv:2006.04804/code)

Reviewed Version (pdf): https://openreview.net/references/pdf?id=zDE5P6nL77