Optimal Transport Graph Neural Networks

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.

References in corpus (10)

• Graph Matching Networks for Learning the Similarity of Graph Structured Objects
• Geom-GCN: Geometric Graph Convolutional Networks
• Covariant Compositional Networks For Learning Graphs
• Hyperbolic Graph Neural Networks
• Learning Generative Models with Sinkhorn Divergences
• Deep Graph Matching Consensus
• The optimal assignment kernel is not positive definite
• Constant Curvature Graph Convolutional Networks
• Towards Interpretable Sparse Graph Representation Learning with Laplacian Pooling
• Gromov-Wasserstein Factorization Models for Graph Clustering

Cited by in corpus (3)

• Wasserstein Embedding for Graph Learning
• Graph Classification by Mixture of Diverse Experts
• Balanced Coarsening for Multilevel Hypergraph Partitioning via Wasserstein Discrepancy