[Paper Review] Wasserstein Embedding for Graph Learning

Soheil Kolouri, Navid Naderializadeh+2 more|arXiv (Cornell University)|Jun 16, 2020

Advanced Graph Neural Networks64 references4 citations

TL;DR

WEGL proposes a linear, computationally efficient framework for graph embedding by using Wasserstein distance between node embeddings to create fixed-size vector representations, enabling fast graph-level prediction with state-of-the-art accuracy and linear complexity in the number of graphs, unlike prior methods requiring quadratic pairwise comparisons.

ABSTRACT

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.

Motivation & Objective

• To address the high computational cost of graph kernel methods that require pairwise similarity calculations between all graphs.
• To overcome the scalability limitations of graph neural networks (GNNs) and kernel-based methods on large-scale graph datasets.
• To develop a graph embedding method that supports any downstream ML model while maintaining metric fidelity via Wasserstein distance.
• To reduce computational complexity from quadratic to linear in the number of graphs for Wasserstein-based graph similarity estimation.
• To enable efficient training and inference for graph property prediction tasks on large-scale datasets.

Proposed method

• WEGL computes Monge maps from a reference node embedding distribution to each graph’s node embeddings to derive transport plans.
• It embeds each graph into a Hilbert space using the optimal transport maps, resulting in fixed-size vector representations.
• The Euclidean distance between embedded graphs approximates their 2-Wasserstein distance, preserving metric structure.
• The method leverages linear optimal transport to avoid pairwise distance calculations between all graphs, reducing complexity from O(M²) to O(M).
• Node embeddings are used as input distributions for Wasserstein distance computation, with the reference distribution derived from a pooled or learned distribution across graphs.
• The framework supports any downstream classifier, including SVM, random forest, and gradient boosted trees, due to its explicit Hilbert space embedding.

Experimental results

Research questions

• RQ1 Can we achieve state-of-the-art graph classification performance while reducing computational complexity in graph embedding?
• RQ2 Can we embed graphs into a Hilbert space such that Euclidean distance approximates 2-Wasserstein distance between node embedding distributions?
• RQ3 Does a linear-complexity approach to Wasserstein-based graph similarity outperform quadratic-complexity methods in scalability and performance?
• RQ4 Can the proposed embedding framework be used with any standard ML classifier, not restricted to kernel methods?
• RQ5 How does WEGL compare in efficiency and accuracy to GNNs and graph kernel baselines on large-scale graph datasets?

Key findings

• WEGL achieves state-of-the-art or competitive classification accuracy on multiple benchmark datasets, including ogbg-molhiv and TUD graph classification tasks.
• On the IMDB-BINARY dataset, WEGL with GBDT achieved 75.2% accuracy, outperforming GIN and other baselines.
• On the PROTEINS dataset, WEGL with GBDT reached 92.9% accuracy, significantly outperforming the next best method.
• WEGL's training time is several orders of magnitude faster than WWL and GIN on datasets with large numbers of graphs, especially due to linear complexity.
• Inference with WEGL is faster than CPU-based GIN and significantly faster than WWL, even when GIN is GPU-accelerated.
• The method maintains high performance across diverse datasets, including molecular property prediction and social network analysis, with minimal computational overhead.