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

1 INTRODUCTION . Measuring the distance between two distributions or sets of objects is a central problem in machine learning . One common method of solving this is optimal transport ( OT ) . OT is concerned with the problem of finding the transport plan for moving a source distribution ( e.g . a pile of earth ) to a sink distribution ( e.g . a construction pit ) with the cheapest cost w.r.t . some pointwise cost function ( e.g . the Euclidean distance ) . The advantages of this method have been shown numerous times , e.g . in generative modelling ( Arjovsky et al. , 2017 ; Bousquet et al. , 2017 ; Genevay et al. , 2018 ) , loss functions ( Frogner et al. , 2015 ) , set matching ( Wang et al. , 2019 ) , or domain adaptation ( Courty et al. , 2017 ) . Motivated by this , many different methods for accelerating OT have been proposed in recent years ( Indyk & Thaper , 2003 ; Papadakis et al. , 2014 ; Backurs et al. , 2020 ) . However , most of these approaches are specialized methods that do not generalize to modern deep learning models , which rely on dynamically changing high-dimensional embeddings . In this work we aim to make OT computation for point sets more scalable by proposing two fast and accurate approximations of entropy-regularized optimal transport : Sparse Sinkhorn and LCNSinkhorn , the latter relying on our newly proposed locally corrected Nyström ( LCN ) method . Sparse Sinkhorn uses a sparse cost matrix to leverage the fact that in entropy-regularized OT ( also known as the Sinkhorn distance ) ( Cuturi , 2013 ) often only each point ’ s nearest neighbors influence the result . LCN-Sinkhorn extends this approach by leveraging LCN , a general similarity matrix approximation that fuses local ( sparse ) and global ( low-rank ) approximations , allowing us to simultaneously capture both kinds of behavior . LCN-Sinkhorn thus fuses sparse Sinkhorn and Nyström-Sinkhorn ( Altschuler et al. , 2019 ) . Both sparse Sinkhorn and LCN-Sinkhorn run in log-linear time . We theoretically analyze these approximations and show that sparse corrections can lead to significant improvements over the Nyström approximation . We furthermore validate these approximations by showing that they are able to reproduce both the Sinkhorn distance and transport plan significantly better than previous methods across a wide range of regularization parameters and computational budgets ( as e.g . demonstrated in Fig . 1 ) . We then show the impact of these improvements by employing Sinkhorn approximations end-to-end in two high-impact machine learning tasks . First , we incorporate them into Wasserstein Procrustes for word embedding alignment ( Grave et al. , 2019 ) . LCN-Sinkhorn improves upon the original method ’ s accuracy by 3.1 percentage points using a third of the training time without any further model changes . Second , we develop the graph transport network ( GTN ) , which combines graph neural networks ( GNNs ) with optimal transport , and further improve it via learnable unbalanced OT and multi-head OT . GTN with LCN-Sinkhorn is the first model that both overcomes the bottleneck of using a single embedding per graph and scales log-linearly in the number of nodes . In summary , our paper ’ s main contributions are : • Locally Corrected Nyström ( LCN ) , a flexible , log-linear time approximation for similarity matrices , leveraging both local ( sparse ) and global ( low-rank ) approximations . • Entropy-regularized optimal transport ( a.k.a . Sinkhorn distance ) with log-linear runtime via sparse Sinkhorn and LCN-Sinkhorn . These are the first log-linear approximations that are stable enough to substitute full entropy-regularized OT in models that leverage high-dimensional spaces . • The graph transport network ( GTN ) , which combines a graph neural network ( GNN ) with multihead unbalanced LCN-Sinkhorn . GTN both sets the state of the art on graph distance regression and still scales log-linearly in the number of nodes . 2 SPARSE SINKHORN . Entropy-regularized optimal transport . In this work we focus on optimal transport between two discrete sets of points . We furthermore add entropy regularization , which enables fast computation and often performs better than regular OT ( Cuturi , 2013 ) . Formally , given two categorical distributions modelled via the vectors p ∈ Rn and q ∈ Rm supported on two sets of points Xp = { xp1 , . . . , xpn } and Xq = { xq1 , . . . , xqm } in Rd and the cost function c : Rd × Rd → R ( e.g . the squared L2 distance ) giving rise to the cost matrix Cij = c ( xpi , xqi ) we aim to find the Sinkhorn distance dλc and the associated optimal transport plan P̄ ( Cuturi , 2013 ) dλc = min P 〈P , C〉F − λH ( P ) , s.t . P1m = p , P T1n = q , ( 1 ) with the Frobenius inner product 〈. , .〉F and the entropy H ( P ) = − ∑n i=1 ∑m j=1Pij logPij . Note that dλc includes the entropy and can thus be negative , while Cuturi ( 2013 ) originally used d 1/λ Cuturi , c = 〈P̄ , C〉F . This optimization problem can be solved by finding the vectors s̄ and t̄ that normalize the columns and rows of the matrix P̄ = diag ( s̄ ) K diag ( t̄ ) with the similarity matrixKij = e− Cij λ , so that P̄1m = p and P̄ T1n = q . This is usually achieved via the Sinkhorn algorithm , which initializes the normalization vectors as s ( 1 ) = 1n and t ( 1 ) = 1m and then updates them alternatingly via s ( i ) = p ( Kt ( i−1 ) ) , t ( i ) = q ( KTs ( i ) ) ( 2 ) until convergence , where denotes elementwise division . Sparse Sinkhorn . The Sinkhorn algorithm is faster than non-regularized EMD algorithms , which run in O ( n2m log n log ( nmax ( C ) ) ) ( Tarjan , 1997 ) . However , its computational cost is still quadratic in time , i.e . O ( nm ) , which is prohibitively expensive for large n and m. We propose to overcome this by observing that the matrix K , and hence also P̄ , is negligibly small everywhere except at each point ’ s closest neighbors because of the exponential used inK ’ s computation . We propose to leverage this by approximating C via the sparse matrix Csp , where Cspij = { Cij if xpi and xqj are “ near ” , ∞ otherwise . ( 3 ) Ksp and P̄ sp follow according to the definitions of K and P̄ . In this work we primarily consider neighbors with distance lower than r1 as “ near ” . Finding such neighbors can be efficiently solved via locality sensitive hashing ( LSH ) on Xp ∪Xq . Locality sensitive hashing . LSH tries to filter “ near ” from “ far ” data points by putting them into different hash buckets . Points closer than a certain distance r1 are put into the same bucket with probability at least p1 , while those beyond some distance r2 = c · r1 with c > 1 are put into the same bucket with probability at most p2 p1 . There is a plethora of LSH methods for different cost functions ( Wang et al. , 2014 ; Shrivastava & Li , 2014 ) , so we do not have to restrict our approach to a limited set of functions . In this work we focus on cross-polytope LSH ( Andoni et al. , 2015 ) and k-means LSH ( Paulevé et al. , 2010 ) , depending on the cost function ( see App . H ) . Sparse Sinkhorn with LSH scales log-linearly with the number of points , i.e . O ( n log n ) for n ≈ m ( see App . A and App . K for details ) . Unfortunately , LSH can fail when e.g . the cost between pairs is very similar ( see App . B ) . However , we can alleviate these limitations by fusingKsp with the Nyström approximation . 3 LOCALLY CORRECTED NYSTRÖM AND LCN-SINKHORN . Nyström method . The Nyström method is a popular way of approximating similarity matrices that provides performance guarantees for many important tasks ( Williams & Seeger , 2001 ; Musco & Musco , 2017 ) . It approximates a positive semi-definite ( PSD ) similarity matrixK via its low-rank decomposition KNys = UA−1V . Since the optimal decomposition via SVD is too expensive to compute , Nyström instead chooses a set of l landmarks L = { xl1 , . . . , xll } and obtains the matrices via Uij = k ( xpi , xlj ) , Aij = k ( xli , xlj ) , and Vij = k ( xli , xqj ) , where k ( x1 , x2 ) is an arbitrary PSD kernel , e.g . k ( x1 , x2 ) = e− c ( x1 , x2 ) λ for Sinkhorn . Common methods of choosing landmarks from Xp ∪Xq are uniform and ridge leverage score ( RLS ) sampling . We instead focus on k-means Nyström and sampling via k-means++ , which we found to be significantly faster than recursive RLS sampling ( Zhang et al. , 2008 ) and perform better than both uniform and RLS sampling ( see App . H ) . Sparse vs. Nyström . Exponential kernels like the one used forK ( e.g . the Gaussian kernel ) typically have a reproducing kernel Hilbert space that is infinitely dimensional . The resulting Gram matrixK thus always has full rank . A low-rank approximation like the Nyström method can therefore only account for its global structure and not the local structure around each point x . As such , it is ill-suited for any moderately low entropy regularization parameter , where the transport matrix P̄ resembles a permutation matrix . Sparse Sinkhorn , on the other hand , can not account for global structure and instead approximates all non-selected distances as infinity . It will hence fail if more than a handful of neighbors are required per point . These approximations are thus opposites of each other , and as such not competing but rather complementary approaches . Locally corrected Nyström . Since we know that the entries in our sparse approximation are exact , fusing this matrix with the Nyström method is rather straightforward . For all non-zero values in the sparse approximationKsp we first calculate the corresponding Nyström approximations , obtaining the sparse matrixKspNys . To obtain the locally corrected Nyström ( LCN ) approximation we remove these entries fromKNys and replace them with their exact values , i.e . KLCN = KNys +K sp ∆ = KNys −K sp Nys +K sp . ( 4 ) LCN-Sinkhorn . To obtain the approximate transport plan P̄LCN we run the Sinkhorn algorithm with KLCN instead of K. However , we never fully instantiate KLCN . Instead , we only save the decomposition and directly use these parts in Eq . ( 2 ) viaKLCNt = U ( A−1V t ) +K sp ∆t , similarly to Altschuler et al . ( 2019 ) . As a result we obtain the decomposition of P̄LCN = P̄Nys + P̄ sp ∆ = P̄U P̄W + P̄ sp − P̄ spNys and the approximate distance ( using Lemma A from Altschuler et al . ( 2019 ) ) dλLCN , c = λ ( sT P̄U P̄W1m + 1 T n P̄U P̄W t+ s T P̄ sp∆ 1m + 1 T n P̄ sp ∆ t ) . ( 5 ) This approximation scales log-linearly with dataset size ( see App . A and App . K for details ) . It allows us to smoothly move from Nyström-Sinkhorn to sparse Sinkhorn by varying the number of neighbors and landmarks . We can thus freely choose the optimal “ operating point ” based on the underlying problem and regularization parameter . We discuss the limitations of LCN-Sinkhorn in App . B .

Summary: The paper considers the problem of approximating Sinkhorn divergence and corresponding transportation plan by combining low-rank and sparse approximation for the Sinkhorn kernel and using Nystrom iterations as a substitute for Sinkhorn's iterations. The corresponding approach is amenable to differentiation and can be used as a building block in different architectures. Numerical experiments in several settings are performed to compare  the proposed approach with existing ones and demonstrate its scalability.