# Scalable Optimal Transport in High Dimensions for Graph Distances, Embedding Alignment, and More

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

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Builds on 41 of its 81 references that are in the map.
- First posted on arXiv 14 Jul 2021.

## Builds on (most cited first, 30 of 41)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Predict then Propagate: Graph Neural Networks meet Personalized PageRank](https://mlsanity.com/paper/1810.05997/) (2018)
- [Computational Optimal Transport](https://mlsanity.com/paper/1803.00567/) (2018)
- [Directional Message Passing for Molecular Graphs](https://mlsanity.com/paper/2003.03123/) (2020)
- [Word Translation Without Parallel Data](https://mlsanity.com/paper/1710.04087/) (2017)
- [Learning Generative Models with Sinkhorn Divergences](https://mlsanity.com/paper/1706.00292/) (2017)
- [Joint Distribution Optimal Transportation for Domain Adaptation](https://mlsanity.com/paper/1705.08848/) (2017)
- [Graph Matching Networks for Learning the Similarity of Graph Structured Objects](https://mlsanity.com/paper/1904.12787/) (2019)
- [Learning with a Wasserstein Loss](https://mlsanity.com/paper/1506.05439/) (2015)
- [Stochastic Optimization for Large-scale Optimal Transport](https://mlsanity.com/paper/1605.08527/) (2016)
- [Optimal Transport for structured data with application on graphs](https://mlsanity.com/paper/1805.09114/) (2018)
- [A Survey Of Cross-lingual Word Embedding Models](https://mlsanity.com/paper/1706.04902/) (2017)
- [Gromov-Wasserstein Alignment of Word Embedding Spaces](https://mlsanity.com/paper/1809.00013/) (2018)
- [Computational Optimal Transport: Complexity by Accelerated Gradient Descent Is Better Than by Sinkhorn's Algorithm](https://mlsanity.com/paper/1802.04367/) (2018)
- [Asymmetric LSH (ALSH) for Sublinear Time Maximum Inner Product Search (MIPS)](https://mlsanity.com/paper/1405.5869/) (2014)
- [Practical and Optimal LSH for Angular Distance](https://mlsanity.com/paper/1509.02897/) (2015)
- [Learning Combinatorial Embedding Networks for Deep Graph Matching](https://mlsanity.com/paper/1904.00597/) (2019)
- [Unsupervised Alignment of Embeddings with Wasserstein Procrustes](https://mlsanity.com/paper/1805.11222/) (2018)
- [SimGNN: A Neural Network Approach to Fast Graph Similarity Computation](https://mlsanity.com/paper/1808.05689/) (2018)
- [Smooth and Sparse Optimal Transport](https://mlsanity.com/paper/1710.06276/) (2017)
- [Stabilized Sparse Scaling Algorithms for Entropy Regularized Transport Problems](https://mlsanity.com/paper/1610.06519/) (2016)
- [Loss in Translation: Learning Bilingual Word Mapping with a Retrieval Criterion](https://mlsanity.com/paper/1804.07745/) (2018)
- [GOT: An Optimal Transport framework for Graph comparison](https://mlsanity.com/paper/1906.02085/) (2019)
- [From optimal transport to generative modeling: the VEGAN cookbook](https://mlsanity.com/paper/1705.07642/) (2017)
- [Massively scalable Sinkhorn distances via the Nystr"om method](https://mlsanity.com/paper/1812.05189/) (2018)

## Built on by (most cited first, 11 of 11)

- [Scatterbrain: Unifying Sparse and Low-rank Attention Approximation](https://mlsanity.com/paper/2110.15343/) (2021)
- [Scalable Optimal Transport Methods in Machine Learning: A Contemporary Survey](https://mlsanity.com/paper/2305.05080/) (2023)
- [Sobolev Transport: A Scalable Metric for Probability Measures with Graph Metrics](https://mlsanity.com/paper/2202.10723/) (2022)
- [A Fused Gromov-Wasserstein Framework for Unsupervised Knowledge Graph Entity Alignment](https://mlsanity.com/paper/2305.06574/) (2023)
- [MHMS: Multimodal Hierarchical Multimedia Summarization](https://mlsanity.com/paper/2204.03734/) (2022)
- [Semantics-Consistent Cross-domain Summarization via Optimal Transport Alignment](https://mlsanity.com/paper/2210.04722/) (2022)
- [Importance Sparsification for Sinkhorn Algorithm](https://mlsanity.com/paper/2306.06581/) (2023)
- [Safe Screening for Unbalanced Optimal Transport](https://mlsanity.com/paper/2307.00247/) (2023)
- [Auto-encoding Molecules: Graph-Matching Capabilities Matter](https://mlsanity.com/paper/2503.00426/) (2025)
- [SinkSLOT: Sinkhorn via Sparse Lifted Optimal Transport](https://mlsanity.com/paper/2608.28262/) (2026)
- [FlashSinkhorn 2: Block-Sparse Entropic Optimal Transport](https://mlsanity.com/paper/2610.02395/) (2026)

## Abstract (arXiv)

The current best practice for computing optimal transport (OT) is via entropy regularization and Sinkhorn iterations. This algorithm runs in quadratic time as it requires the full pairwise cost matrix, which is prohibitively expensive for large sets of objects. In this work we propose two effective log-linear time approximations of the cost matrix: First, a sparse approximation based on locality-sensitive hashing (LSH) and, second, a Nystr"om approximation with LSH-based sparse corrections, which we call locally corrected Nystr"om (LCN). These approximations enable general log-linear time algorithms for entropy-regularized OT that perform well even for the complex, high-dimensional spaces common in deep learning. We analyse these approximations theoretically and evaluate them experimentally both directly and end-to-end as a component for real-world applications. Using our approximations for unsupervised word embedding alignment enables us to speed up a state-of-the-art method by a factor of 3 while also improving the accuracy by 3.1 percentage points without any additional model changes. For graph distance regression we propose the graph transport network (GTN), which combines graph neural networks (GNNs) with enhanced Sinkhorn. GTN outcompetes previous models by 48% and still scales log-linearly in the number of nodes.