Optimal Transport Graph Neural Networks

8 Jun 2020  ·  Benson Chen, Gary Bécigneul, Octavian-Eugen Ganea, Regina Barzilay, Tommi Jaakkola ·

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.

PDF Abstract


Results from the Paper

 Ranked #1 on Graph Regression on Lipophilicity (using extra training data)

     Get a GitHub badge
Task Dataset Model Metric Name Metric Value Global Rank Uses Extra
Training Data
Result Benchmark
Drug Discovery BACE ProtoW-L2 AUC 0.873 # 3
Drug Discovery BBBP ProtoW-L2 AUC 0.92 # 1
Graph Regression ESOL ProtoW-dot RMSE .594 # 1
Graph Regression Lipophilicity ProtoS-L2 RMSE 0.580 # 1