# Wasserstein Embedding for Graph Learning

arXiv: https://arxiv.org/abs/2006.09430 · HTML: https://mlsanity.com/paper/2006.09430/

## Key facts

- Cited by 56 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 16 Jun 2020.

## Built on by (most cited first, 30 of 56)

- [Walking Out of the Weisfeiler Leman Hierarchy: Graph Learning Beyond Message Passing](https://mlsanity.com/paper/2102.08786/) (2021)
- [A Survey on Optimal Transport for Machine Learning: Theory and Applications](https://mlsanity.com/paper/2106.01963/) (2021)
- [Wasserstein Task Embedding for Measuring Task Similarities](https://mlsanity.com/paper/2208.11726/) (2022)
- [Hierarchical Multi-Marginal Optimal Transport for Network Alignment](https://mlsanity.com/paper/2310.04470/) (2023)
- [Molecular Representation Learning via Heterogeneous Motif Graph Neural Networks](https://mlsanity.com/paper/2202.00529/) (2022)
- [Template based Graph Neural Network with Optimal Transport Distances](https://mlsanity.com/paper/2205.15733/) (2022)
- [Learning Prototype-oriented Set Representations for Meta-Learning](https://mlsanity.com/paper/2110.09140/) (2021)
- [Expected Sliced Transport Plans](https://mlsanity.com/paper/2410.12176/) (2024)
- [Pave Your Own Path: Graph Gradual Domain Adaptation on Fused Gromov-Wasserstein Geodesics](https://mlsanity.com/paper/2505.12709/) (2025)
- [Graph Classification by Mixture of Diverse Experts](https://mlsanity.com/paper/2103.15622/) (2021)
- [Weisfeiler and Leman Go Infinite: Spectral and Combinatorial Pre-Colorings](https://mlsanity.com/paper/2201.13410/) (2022)
- [Going Deeper into Permutation-Sensitive Graph Neural Networks](https://mlsanity.com/paper/2205.14368/) (2022)
- [Flow Factorized Representation Learning](https://mlsanity.com/paper/2309.13167/) (2023)
- [Graph Classification via Reference Distribution Learning: Theory and Practice](https://mlsanity.com/paper/2408.11370/) (2024)
- [Lightspeed Geometric Dataset Distance via Sliced Optimal Transport](https://mlsanity.com/paper/2501.18901/) (2025)
- [SLOSH: Set LOcality Sensitive Hashing via Sliced-Wasserstein Embeddings](https://mlsanity.com/paper/2112.05872/) (2021)
- [On a linear fused Gromov-Wasserstein distance for graph structured data](https://mlsanity.com/paper/2203.04711/) (2022)
- [Linear Optimal Partial Transport Embedding](https://mlsanity.com/paper/2302.03232/) (2023)
- [FSW-GNN: A Bi-Lipschitz WL-Equivalent Graph Neural Network](https://mlsanity.com/paper/2410.09118/) (2024)
- [Finding Heterophilic Neighbors via Confidence-based Subgraph Matching for Semi-supervised Node Classification](https://mlsanity.com/paper/2302.09755/) (2023)
- [Gaussian process regression with Sliced Wasserstein Weisfeiler-Lehman graph kernels](https://mlsanity.com/paper/2402.03838/) (2024)
- [Multiscale Graph Comparison via the Embedded Laplacian Discrepancy](https://mlsanity.com/paper/2201.12064/) (2022)
- [Fast Estimation of Wasserstein Distances via Regression on Sliced Wasserstein Distances](https://mlsanity.com/paper/2509.20508/) (2025)
- [Set Representation Learning with Generalized Sliced-Wasserstein Embeddings](https://mlsanity.com/paper/2103.03892/) (2021)
- [Supervised Contrastive Learning with Structure Inference for Graph Classification](https://mlsanity.com/paper/2203.07691/) (2022)
- [Non-isotropic Persistent Homology: Leveraging the Metric Dependency of PH](https://mlsanity.com/paper/2310.16437/) (2023)
- [SeTformer is What You Need for Vision and Language](https://mlsanity.com/paper/2401.03540/) (2024)
- [Embedding Graphs on Grassmann Manifold](https://mlsanity.com/paper/2205.15068/) (2022)
- [Regularized Optimal Transport Layers for Generalized Global Pooling Operations](https://mlsanity.com/paper/2212.06339/) (2022)
- [Equivariant vs. Invariant Layers: A Comparison of Backbone and Pooling for Point Cloud Classification](https://mlsanity.com/paper/2306.05553/) (2023)

## Abstract (arXiv)

We present Wasserstein Embedding for Graph Learning (WEGL), a novel and fast framework for embedding entire graphs in a vector space, in which various machine learning models are applicable for graph-level prediction tasks. We leverage new insights on defining similarity between graphs as a function of the similarity between their node embedding distributions. Specifically, we use the Wasserstein distance to measure the dissimilarity between node embeddings of different graphs. Unlike prior work, we avoid pairwise calculation of distances between graphs and reduce the computational complexity from quadratic to linear in the number of graphs. WEGL calculates Monge maps from a reference distribution to each node embedding and, based on these maps, creates a fixed-sized vector representation of the graph. We evaluate our new graph embedding approach on various benchmark graph-property prediction tasks, showing state-of-the-art classification performance while having superior computational efficiency. The code is available at https://github.com/navid-naderi/WEGL.