[Paper Review] Optimal Transport Graph Neural Networks

Benson Chen, Gary Bécigneul+3 more|arXiv (Cornell University)|Jun 8, 2020

Advanced Graph Neural Networks61 references22 citations

TL;DR

This paper introduces OT-GNN, a graph neural network that computes graph embeddings using Wasserstein distances between node embeddings and learnable prototype point clouds, enabling universal approximation and improved molecular property prediction. By combining optimal transport with parametric prototypes and a noise contrastive regularizer, the model avoids aggregation collapse and achieves smoother, more informative representations than standard GNNs.

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.

Motivation & Objective

• To address the limitation of standard GNNs that lose structural and semantic information through simple aggregation (e.g., sum or mean) of node embeddings.
• To develop a parametric graph model that leverages optimal transport (Wasserstein distance) to compare input graphs with abstract, learnable prototypes.
• To overcome optimization challenges in end-to-end training of OT-based models, particularly prototype collapse where prototypes degenerate to single points.
• To provide theoretical justification for the increased representational capacity of the proposed function class over standard GNN aggregation.
• To empirically demonstrate superior performance on molecular property prediction benchmarks while producing smoother, more interpretable graph representations.

Proposed method

• The model uses a standard GNN (e.g., GIN or GAT) to encode input graphs into node embeddings, forming a point cloud in embedding space.
• It introduces learnable prototype point clouds as free parameters, representing abstract structural or functional features (e.g., functional groups) relevant to the downstream task.
• Graph-level representations are computed as the Wasserstein distance between the input graph’s node embedding point cloud and each prototype point cloud.
• The model is trained end-to-end using a noise contrastive regularizer to prevent prototypes from collapsing to single points, preserving the geometric structure of OT.
• The final graph embedding is formed by concatenating the Wasserstein distances to all prototypes, enabling rich, geometry-aware representations.
• The method is theoretically justified by proving that the resulting function class on point clouds is a universal approximator, unlike standard sum-based aggregation.

Experimental results

Research questions

• RQ1 Can optimal transport-based graph representation learning with learnable prototypes outperform standard GNN aggregation in molecular property prediction?
• RQ2 Does the use of parametric prototypes in combination with OT improve the representational capacity of GNNs compared to sum or mean aggregation?
• RQ3 How can optimization challenges—particularly prototype collapse—be mitigated in end-to-end training of OT-based GNNs?
• RQ4 Do the learned prototypes correspond to interpretable molecular substructures or functional groups?
• RQ5 Is the resulting embedding space smoother and more robust than that of standard GNNs?

Key findings

• OT-GNN outperforms state-of-the-art GNNs, including GIN and GAT, on multiple molecular property prediction benchmarks, demonstrating consistent empirical gains.
• The noise contrastive regularizer effectively prevents prototype collapse, enabling the model to exploit the full geometric structure of optimal transport.
• Theoretical analysis confirms that the function class defined by OT-GNN is a universal approximator on point clouds, unlike standard GNN aggregation methods.
• The model produces smoother graph embedding spaces compared to baselines, indicating improved generalization and robustness.
• Interpretability analysis shows that learned prototypes correspond to meaningful molecular substructures: for example, one prototype is closest to molecules with rings (insoluble), while another favors more soluble molecules.
• Hyperparameter search confirms that performance is stable across different numbers of prototypes (M=10,20) and prototype sizes (N=10), with optimal performance achieved at specific hidden dimensions (d=5,10) and regularization coefficients (nc_coef=0.1).