This reposority contains a sample implementation code pertaining to our recent work, Wasserstein Embedding for Graph Learning (WEGL), in which we leverage the linear optimal transport (LOT) theory to introduce a novel and fast framework for embedding entire graphs in a vector space that can then be used for graph classification tasks.
The Jupyter Notebook included in this repository runs WEGL on the OGBG-molhiv dataset from the Open Graph Benchmark, achieving state-of-the-art results on molecular property prediction using a downstream random forest classifier with much-reduced training complexity compared to the existing graph neural network (GNN) approaches.
For further details on the approach and more comprehensive evaluation results, please visit our paper webpage.
- OGB: The Open Graph Benchmark.
- POT: A Python library for Optimal Transport.
- PyTorch
- PyTorch Geometric
- scikit-learn
- NumPy
- Matplotlib
- tqdm
- PrettyTable
If you use WEGL in your work, please cite our paper using the BibTeX citation below:
@inproceedings{
kolouri2021wasserstein,
title={Wasserstein Embedding for Graph Learning},
author={Soheil Kolouri and Navid Naderializadeh and Gustavo K. Rohde and Heiko Hoffmann},
booktitle={International Conference on Learning Representations},
year={2021},
url={https://openreview.net/forum?id=AAes_3W-2z}
}