Warpspeed Computation of Optimal Transport, Graph Distances, and Embedding Alignment

1 Jan 2021  ·  Johannes Klicpera, Marten Lienen, Stephan Günnemann ·

Optimal transport (OT) is a cornerstone of many machine learning tasks. The current best practice for computing OT is via entropy regularization and Sinkhorn iterations. This algorithm runs in quadratic time and requires calculating the full pairwise similarity matrix, which is prohibitively expensive for large sets of objects. To alleviate this limitation we propose to use a sparse approximation of the cost matrix based on locality sensitive hashing (LSH). Moreover, we fuse this sparse approximation with the Nyström method, resulting in the locally corrected Nyström method (LCN). These approximations enable the first log-linear time algorithms for entropy-regularized OT that work for complex real-world tasks using high-dimensional spaces with little to no loss in accuracy. We thoroughly demonstrate this by evaluating different Sinkhorn approximations both directly and as a component of two real-world models. Using approximate Sinkhorn for unsupervised word embedding alignment enables us to train the model full-batch in a fraction of the time while improving upon the original on average by 3.1 percentage points without any model changes. For graph distance regression we propose the graph transport network (GTN), which combines graph neural networks (GNNs) with enhanced Sinkhorn and outcompetes previous models by 52%. LCN-Sinkhorn enables GTN to achieve this while still scaling log-linearly in the number of nodes.

PDF Abstract
No code implementations yet. Submit your code now

Datasets


Results from the Paper


  Submit results from this paper to get state-of-the-art GitHub badges and help the community compare results to other papers.

Methods