Original: W ARPSPEED C OMPUTATION OF O PTIMAL T RANSPORT , G RAPH D ISTANCES , AND E MBEDDING A LIGNMENT
Revised: W ARPSPEED C OMPUTATION OF O PTIMAL T RANSPORT , G RAPH D ISTANCES , AND E MBEDDING A LIGNMENT

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Optimal transport (OT) is a cornerstone of many machine learning tasks.
Revised: Optimal transport (OT) is a cornerstone of many machine learning tasks.

Original: The current best practice for computing OT is via entropy regularization and Sinkhorn iterations.
Revised: The current best practice for computing OT is via entropy regularization and Sinkhorn iterations.

Original: This algorithm runs in quadratic time and requires calculating the full pairwise cost matrix, which is prohibitively expensive for large sets of objects.
Revised: This algorithm runs in quadratic time and requires calculating the full pairwise cost matrix, which is prohibitively expensive for large sets of objects.

Original: To alleviate this limitation we propose to instead use a sparse approximation of the cost matrix based on locality sensitive hashing (LSH).
Revised: To alleviate this limitation we propose to instead use a sparse approximation of the cost matrix based on locality sensitive hashing (LSH).

Original: Moreover, we fuse this sparse approximation with the Nyström method, resulting in the locally corrected Nyström method (LCN).
Revised: Moreover, we fuse this sparse approximation with the Nyström method, resulting in the locally corrected Nyström method (LCN).

Original: These approximations enable general log-linear time algorithms for entropy-regularized OT that perform well even in complex, high-dimensional spaces.
Revised: These approximations enable general log-linear time algorithms for entropy-regularized OT that perform well even in complex, high-dimensional spaces.

Original: We thoroughly demonstrate these advantages via a theoretical analysis and by evaluating multiple approximations both directly and as a component of two real-world models.
Revised: We thoroughly demonstrate these advantages via a theoretical analysis and by evaluating multiple approximations both directly and as a component of two real-world models.

Original: 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.
Revised: 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.

Original: 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 48 % .
Revised: 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 48 % .

Original: LCN-Sinkhorn enables GTN to achieve this while still scaling log-linearly in the number of nodes.
Revised: LCN-Sinkhorn enables GTN to achieve this while still scaling log-linearly in the number of nodes.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Measuring the distance between two distributions or sets of objects is a central problem in machine learning.
Revised: Measuring the distance between two distributions or sets of objects is a central problem in machine learning.

Original: One common method of solving this is optimal transport (OT).
Revised: One common method of solving this is optimal transport (OT).

Original: OT is concerned with the problem of ﬁnding 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.
Revised: OT is concerned with the problem of ﬁnding 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.

Original: some pointwise cost function (e.g. the Euclidean distance).
Revised: some pointwise cost function (e.g. the Euclidean distance).

Original: 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).
Revised: 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).

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Sparse
Revised: Sparse

Original: 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 inﬂuence the result.
Revised: 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 inﬂuence the result.

Original: 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.
Revised: 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.

Original: LCN-Sinkhorn thus fuses sparse Sinkhorn and Nyström-Sinkhorn (Altschuler et al., 2019).
Revised: LCN-Sinkhorn thus fuses sparse Sinkhorn and Nyström-Sinkhorn (Altschuler et al., 2019).

Original: Both sparse Sinkhorn and LCN-Sinkhorn run in log-linear time.
Revised: Both sparse Sinkhorn and LCN-Sinkhorn run in log-linear time.

Original: We theoretically analyze these approximations and show that sparse corrections can lead to signiﬁcant improvements over the Nyström approximation.
Revised: We theoretically analyze these approximations and show that sparse corrections can lead to signiﬁcant improvements over the Nyström approximation.

Original: We furthermore validate these approximations by showing that they are able to reproduce both the Sinkhorn distance and transport plan signiﬁcantly better than previous methods across a wide range of regularization parameters and computational
Revised: We furthermore validate these approximations by showing that they are able to reproduce both the Sinkhorn distance and transport plan signiﬁcantly better than previous methods across a wide range of regularization parameters and computational

Original: 1 ¯ P ij ¯ P M
Revised: 1 ¯ P ij ¯ P M

Original: S i j
Revised: S i j

Original: Multiscale OT 1 ¯ P ij ¯ P N y s , i j
Revised: Multiscale OT 1 ¯ P ij ¯ P N y s , i j

Original: D i a gon a l ( p e rf ec t c o rr e l a ti on )
Revised: D i a gon a l ( p e rf ec t c o rr e l a ti on )

Original: Nyström Sinkhorn 1 ¯ P ij ¯ P s p
Revised: Nyström Sinkhorn 1 ¯ P ij ¯ P s p

Original: i j
Revised: i j

Original: Sparse Sinkhorn 1 ¯ P ij ¯ P L C
Revised: Sparse Sinkhorn 1 ¯ P ij ¯ P L C

Original: N , i j
Revised: N , i j

Original: LCN-Sinkhorn budgets (as e.g. demonstrated in Fig. 1).
Revised: LCN-Sinkhorn budgets (as e.g. demonstrated in Fig. 1).

Original: We then show the impact of these improvements by employing Sinkhorn approximations end-to-end in two high-impact machine learning tasks.
Revised: We then show the impact of these improvements by employing Sinkhorn approximations end-to-end in two high-impact machine learning tasks.

Original: First, we incorporate them into Wasserstein Procrustes for word embedding alignment (Grave et al., 2019).
Revised: First, we incorporate them into Wasserstein Procrustes for word embedding alignment (Grave et al., 2019).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: GTN with LCN-Sinkhorn is the ﬁrst model that both overcomes the bottleneck of using a single embedding per graph and scales log-linearly in the number of nodes.
Revised: GTN with LCN-Sinkhorn is the ﬁrst model that both overcomes the bottleneck of using a single embedding per graph and scales log-linearly in the number of nodes.

Original: In summary, our paper’s main contributions are:
Revised: In summary, our paper’s main contributions are:

Original: • Locally Corrected Nyström (LCN), a ﬂexible, log-linear time approximation for similarity matrices, leveraging both local (sparse) and global (low-rank) approximations.
Revised: • Locally Corrected Nyström (LCN), a ﬂexible, log-linear time approximation for similarity matrices, leveraging both local (sparse) and global (low-rank) approximations.

Original: • Entropy-regularized optimal transport (a.k.a.
Revised: • Entropy-regularized optimal transport (a.k.a.

Original: Sinkhorn distance) with log-linear runtime via sparse Sinkhorn and LCN-Sinkhorn.
Revised: Sinkhorn distance) with log-linear runtime via sparse Sinkhorn and LCN-Sinkhorn.

Original: These are the ﬁrst log-linear approximations that are stable enough to substitute full entropy-regularized OT in models that leverage high-dimensional spaces.
Revised: These are the ﬁrst log-linear approximations that are stable enough to substitute full entropy-regularized OT in models that leverage high-dimensional spaces.

Original: • The graph transport network (GTN), which combines a graph neural network (GNN) with multihead unbalanced LCN-Sinkhorn.
Revised: • The graph transport network (GTN), which combines a graph neural network (GNN) with multihead unbalanced LCN-Sinkhorn.

Original: GTN both sets the state of the art on graph distance regression and still scales log-linearly in the number of nodes.
Revised: GTN both sets the state of the art on graph distance regression and still scales log-linearly in the number of nodes.

Original:  S PARSE S INKHORN
Revised:  S PARSE S INKHORN

Original: Entropy-regularized optimal transport.
Revised: Entropy-regularized optimal transport.

Original: In this work we focus on optimal transport between two discrete sets of points.
Revised: In this work we focus on optimal transport between two discrete sets of points.

Original: We furthermore add entropy regularization, which enables fast computation and often performs better than regular OT (Cuturi, 2013).
Revised: We furthermore add entropy regularization, which enables fast computation and often performs better than regular OT (Cuturi, 2013).

Original: Formally, given two categorical distributions modelled via the vectors p ∈ R n and q ∈ R m supported on two sets of points X p = { x p 1 , . . .
Revised: Formally, given two categorical distributions modelled via the vectors p ∈ R n and q ∈ R m supported on two sets of points X p = { x p 1 , . . .

Original: , x p n } and X q = { x q 1 , . . .
Revised: , x p n } and X q = { x q 1 , . . .

Original: , x q m } in R d and the cost function c : R d × R d → R (e.g. the squared Ldistance) giving rise to the cost matrix C ij =
Revised: , x q m } in R d and the cost function c : R d × R d → R (e.g. the squared Ldistance) giving rise to the cost matrix C ij =

Original: c ( x p i , x q i ) we aim to ﬁnd the Sinkhorn distance d λc and the associated optimal transport plan ¯ P (Cuturi, 2013)
Revised: c ( x p i , x q i ) we aim to ﬁnd the Sinkhorn distance d λc and the associated optimal transport plan ¯ P (Cuturi, 2013)

Original: [Equation]
Revised: [Equation]

Original: with the Frobenius inner product (cid:104) .
Revised: with the Frobenius inner product (cid:104) .

Original: (cid:105) F and the entropy H ( P )
Revised: (cid:105) F and the entropy H ( P )

Original: = − (cid:80)
Revised: = − (cid:80)

Original: ni
Revised: ni

Original: =1 (cid:80) mj
Revised: =1 (cid:80) mj

Original: =1 P ij log P ij .
Revised: =1 P ij log P ij .

Original: Note that d λc includes the entropy and can thus be negative, while Cuturi (2013) originally used d 1 /λ Cuturi ,
Revised: Note that d λc includes the entropy and can thus be negative, while Cuturi (2013) originally used d 1 /λ Cuturi ,

Original: c = (cid:104)
Revised: c = (cid:104)

Original: ¯ P , C (cid:105) F .
Revised: ¯ P , C (cid:105) F .

Original: This optimization problem can be solved by ﬁnding the vectors ¯ s and ¯ t that normalize the columns and rows of the matrix ¯ P = diag(¯ s ) K diag(¯ t ) with the similarity matrix K ij =
Revised: This optimization problem can be solved by ﬁnding the vectors ¯ s and ¯ t that normalize the columns and rows of the matrix ¯ P = diag(¯ s ) K diag(¯ t ) with the similarity matrix K ij =

Original: e − C ijλ , so that ¯ P 1 m = p and ¯ P T 1 n = q .
Revised: e − C ijλ , so that ¯ P 1 m = p and ¯ P T 1 n = q .

Original: This is usually achieved via the Sinkhorn algorithm, which initializes the normalization vectors as s (1) = 1 n and t (1) = 1 m and then updates them alternatingly via
Revised: This is usually achieved via the Sinkhorn algorithm, which initializes the normalization vectors as s (1) = 1 n and t (1) = 1 m and then updates them alternatingly via

Original: [Equation]
Revised: [Equation]

Original: until convergence, where (cid:11) denotes elementwise division.
Revised: until convergence, where (cid:11) denotes elementwise division.

Original: Sparse Sinkhorn.
Revised: Sparse Sinkhorn.

Original: The Sinkhorn algorithm is faster than non-regularized EMD algorithms, which run in O ( n 2 m log n log( n max( C )))
Revised: The Sinkhorn algorithm is faster than non-regularized EMD algorithms, which run in O ( n 2 m log n log( n max( C )))

Original: (Tarjan, 1997).
Revised: (Tarjan, 1997).

Original: However, its computational cost is still quadratic in time, i.e. O ( nm ) , which is prohibitively expensive for large n and m .
Revised: However, its computational cost is still quadratic in time, i.e. O ( nm ) , which is prohibitively expensive for large n and m .

Original: 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 in K ’s computation.
Revised: 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 in K ’s computation.

Original: We propose to leverage this by approximating C via the sparse matrix C sp , where
Revised: We propose to leverage this by approximating C via the sparse matrix C sp , where

Original: [Equation]
Revised: [Equation]

Original: if
Revised: if

Original: [Equation]
Revised: [Equation]

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: K sp and ¯ P sp follow according to the deﬁnitions of K and ¯ P .
Revised: K sp and ¯ P sp follow according to the deﬁnitions of K and ¯ P .

Original: In this work we primarily consider neighbors with distance lower than r 1 as “near”.
Revised: In this work we primarily consider neighbors with distance lower than r 1 as “near”.

Original: Finding such neighbors can be efﬁciently solved via locality sensitive hashing (LSH) on X p ∪ X q .
Revised: Finding such neighbors can be efﬁciently solved via locality sensitive hashing (LSH) on X p ∪ X q .

Original: Locality sensitive hashing.
Revised: Locality sensitive hashing.

Original: LSH tries to ﬁlter “near” from “far” data points by putting them into different hash buckets.
Revised: LSH tries to ﬁlter “near” from “far” data points by putting them into different hash buckets.

Original: Points closer than a certain distance r 1 are put into the same bucket with probability at least p 1 , while those beyond some distance r 2 = c · r 1 with c > 1 are put into the same bucket with probability at most p 2 (cid:28) p 1 .
Revised: Points closer than a certain distance r 1 are put into the same bucket with probability at least p 1 , while those beyond some distance r 2 = c · r 1 with c > 1 are put into the same bucket with probability at most p 2 (cid:28) p 1 .

Original: 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.
Revised: 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.

Original: In this work we focus on cross-polytope LSH
Revised: In this work we focus on cross-polytope LSH

Original: (Andoni et al., 2015) and k-means LSH
Revised: (Andoni et al., 2015) and k-means LSH

Original: (Paulevé et al., 2010), depending on the cost function (see App.
Revised: (Paulevé et al., 2010), depending on the cost function (see App.

Original: H).
Revised: H).

Original: Sparse Sinkhorn with LSH scales log-linearly with the number of points, i.e. O ( n log n ) for n ≈ m (see App.
Revised: Sparse Sinkhorn with LSH scales log-linearly with the number of points, i.e. O ( n log n ) for n ≈ m (see App.

Original: A and
Revised: A and

Original: App. K for details).
Revised: App. K for details).

Original: Unfortunately, LSH can fail when e.g. the cost between pairs is very similar (see
Revised: Unfortunately, LSH can fail when e.g. the cost between pairs is very similar (see

Original: App. B).
Revised: App. B).

Original: However, we can alleviate these limitations by fusing K sp with the Nyström approximation.
Revised: However, we can alleviate these limitations by fusing K sp with the Nyström approximation.

Original:  L OCALLY C ORRECTED N YSTRÖM AND LCN-S INKHORN
Revised:  L OCALLY C ORRECTED N YSTRÖM AND LCN-S INKHORN

Original: Nyström method.
Revised: Nyström method.

Original: The Nyström method is a popular way of approximating similarity matrices that provides performance guarantees for many important tasks (Williams & Seeger, 2001; Musco &
Revised: The Nyström method is a popular way of approximating similarity matrices that provides performance guarantees for many important tasks (Williams & Seeger, 2001; Musco &

Original: Musco, 2017).
Revised: Musco, 2017).

Original: It approximates a positive semi-deﬁnite (PSD) similarity matrix K via its low-rank decomposition K Nys = UA − 1 V .
Revised: It approximates a positive semi-deﬁnite (PSD) similarity matrix K via its low-rank decomposition K Nys = UA − 1 V .

Original: Since the optimal decomposition via SVD is too expensive to compute, Nyström instead chooses a set of l landmarks L = { x l 1 , . . .
Revised: Since the optimal decomposition via SVD is too expensive to compute, Nyström instead chooses a set of l landmarks L = { x l 1 , . . .

Original: , x l l } and obtains the matrices via U ij = k ( x p i , x l j ) , A ij = k
Revised: , x l l } and obtains the matrices via U ij = k ( x p i , x l j ) , A ij = k

Original: ( x l i , x l j ) , and V ij = k ( x l i , x q j ) , where k ( x 1 , x 2 ) is an arbitrary
Revised: ( x l i , x l j ) , and V ij = k ( x l i , x q j ) , where k ( x 1 , x 2 ) is an arbitrary

Original: PSD kernel, e.g. k ( x 1 , x 2 ) =
Revised: PSD kernel, e.g. k ( x 1 , x 2 ) =

Original: e − c ( x 1 , x 2) λ for Sinkhorn.
Revised: e − c ( x 1 , x 2) λ for Sinkhorn.

Original: Common methods of choosing landmarks from X p ∪ X q are uniform and ridge leverage score (RLS) sampling.
Revised: Common methods of choosing landmarks from X p ∪ X q are uniform and ridge leverage score (RLS) sampling.

Original: We instead focus on k-means
Revised: We instead focus on k-means

Original: Nyström and sampling via k-means++, which we found to be signiﬁcantly faster than recursive RLS sampling (Zhang et al., 2008) and perform better than both uniform and RLS sampling (see App. H).
Revised: Nyström and sampling via k-means++, which we found to be signiﬁcantly faster than recursive RLS sampling (Zhang et al., 2008) and perform better than both uniform and RLS sampling (see App. H).

Original: Sparse vs. Nyström.
Revised: Sparse vs. Nyström.

Original: Exponential kernels like the one used for K (e.g. the Gaussian kernel) typically have a reproducing kernel Hilbert space that is inﬁnitely dimensional.
Revised: Exponential kernels like the one used for K (e.g. the Gaussian kernel) typically have a reproducing kernel Hilbert space that is inﬁnitely dimensional.

Original: The resulting Gram matrix K thus always has full rank.
Revised: The resulting Gram matrix K thus always has full rank.

Original: 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 .
Revised: 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 .

Original: As such, it is ill-suited for any moderately low entropy regularization parameter, where the transport matrix ¯ P resembles a permutation matrix.
Revised: As such, it is ill-suited for any moderately low entropy regularization parameter, where the transport matrix ¯ P resembles a permutation matrix.

Original: Sparse Sinkhorn, on the other hand, cannot account for global structure and instead approximates all non-selected distances as inﬁnity.
Revised: Sparse Sinkhorn, on the other hand, cannot account for global structure and instead approximates all non-selected distances as inﬁnity.

Original: It will hence fail if more than a handful of neighbors are required per point.
Revised: It will hence fail if more than a handful of neighbors are required per point.

Original: These approximations are thus opposites of each other, and as such not competing but rather complementary approaches.
Revised: These approximations are thus opposites of each other, and as such not competing but rather complementary approaches.

Original: Locally corrected Nyström.
Revised: Locally corrected Nyström.

Original: Since we know that the entries in our sparse approximation are exact, fusing this matrix with the Nyström method is rather straightforward.
Revised: Since we know that the entries in our sparse approximation are exact, fusing this matrix with the Nyström method is rather straightforward.

Original: For all non-zero values in the sparse approximation K sp we ﬁrst calculate the corresponding Nyström approximations, obtaining the sparse matrix K spNys .
Revised: For all non-zero values in the sparse approximation K sp we ﬁrst calculate the corresponding Nyström approximations, obtaining the sparse matrix K spNys .

Original: To obtain the locally corrected Nyström (LCN) approximation we remove these entries from K Nys and replace them with their exact values, i.e.
Revised: To obtain the locally corrected Nyström (LCN) approximation we remove these entries from K Nys and replace them with their exact values, i.e.

Original: [Equation]
Revised: [Equation]

Original: LCN-Sinkhorn.
Revised: LCN-Sinkhorn.

Original: To obtain the approximate transport plan ¯ P LCN we run the Sinkhorn algorithm with K LCN instead of K .
Revised: To obtain the approximate transport plan ¯ P LCN we run the Sinkhorn algorithm with K LCN instead of K .

Original: However, we never fully instantiate K LCN .
Revised: However, we never fully instantiate K LCN .

Original: Instead, we only save the decomposition and directly use these parts in Eq.
Revised: Instead, we only save the decomposition and directly use these parts in Eq.

Original: (2) via K LCN t =
Revised: (2) via K LCN t =

Original: U ( A − 1 V t )
Revised: U ( A − 1 V t )

Original: + K sp ∆ t , similarly to Altschuler et al.
Revised: + K sp ∆ t , similarly to Altschuler et al.

Original: As a result we obtain the decomposition of ¯ P LCN =
Revised: As a result we obtain the decomposition of ¯ P LCN =

Original: ¯ P Nys + ¯ P sp ∆ = ¯
Revised: ¯ P Nys + ¯ P sp ∆ = ¯

Original: P U ¯ P W + ¯ P sp − ¯ P spNys and the approximate distance (using Lemma A from Altschuler et al.
Revised: P U ¯ P W + ¯ P sp − ¯ P spNys and the approximate distance (using Lemma A from Altschuler et al.

Original: [Equation]
Revised: [Equation]

Original: This approximation scales log-linearly with dataset size (see App.
Revised: This approximation scales log-linearly with dataset size (see App.

Original: A and App.
Revised: A and App.

Original: K for details).
Revised: K for details).

Original: It allows us to smoothly move from Nyström-Sinkhorn to sparse Sinkhorn by varying the number of neighbors and landmarks.
Revised: It allows us to smoothly move from Nyström-Sinkhorn to sparse Sinkhorn by varying the number of neighbors and landmarks.

Original: We can thus freely choose the optimal “operating point” based on the underlying problem and regularization parameter.
Revised: We can thus freely choose the optimal “operating point” based on the underlying problem and regularization parameter.

Original: We discuss the limitations of LCN-Sinkhorn in App.
Revised: We discuss the limitations of LCN-Sinkhorn in App.

Original: B.
Revised: B.

Original:  T HEORETICAL A NALYSIS
Revised:  T HEORETICAL A NALYSIS

Original: Approximation error.
Revised: Approximation error.

Original: The main question we aim to answer in our theoretical analysis is what improvements to expect from adding sparse corrections to Nyström Sinkhorn.
Revised: The main question we aim to answer in our theoretical analysis is what improvements to expect from adding sparse corrections to Nyström Sinkhorn.

Original: To do so, we ﬁrst analyse approximations of K in a uniform and a clustered data model.
Revised: To do so, we ﬁrst analyse approximations of K in a uniform and a clustered data model.

Original: In these we use Nyström and
Revised: In these we use Nyström and

Original: LSH schemes that largely resemble k -means, as used in most of our experiments.
Revised: LSH schemes that largely resemble k -means, as used in most of our experiments.

Original: Relevant proofs and notes for this section can be found in App.
Revised: Relevant proofs and notes for this section can be found in App.

Original: C to G.
Revised: C to G.

Original: Theorem 1.
Revised: Theorem 1.

Original: Let X p and X q have n samples that are uniformly distributed in a d -dimensional closed, locally Euclidean manifold with unit volume.
Revised: Let X p and X q have n samples that are uniformly distributed in a d -dimensional closed, locally Euclidean manifold with unit volume.

Original: Let furthermore C ij = (cid:107) x
Revised: Let furthermore C ij = (cid:107) x

Original: p
Revised: p

Original: i − x q j (cid:107) 2 and
Revised: i − x q j (cid:107) 2 and

Original: K ij =
Revised: K ij =

Original: e − C ij /λ .
Revised: e − C ij /λ .

Original: Let the l landmarks L be arranged optimally and regularly so that the expected Ldistance to the closest landmark is minimized.
Revised: Let the l landmarks L be arranged optimally and regularly so that the expected Ldistance to the closest landmark is minimized.

Original: Denote R = 12 min x , y ∈ L, x (cid:54) =
Revised: Denote R = 12 min x , y ∈ L, x (cid:54) =

Original: y (cid:107) x − y (cid:107) 2 .
Revised: y (cid:107) x − y (cid:107) 2 .

Original: Assume that the sparse correction K sp ij = K ij
Revised: Assume that the sparse correction K sp ij = K ij

Original: if
Revised: if

Original: and only if x q j is one of the k − 1 nearest neighbors of x p i , and that the distance to x p
Revised: and only if x q j is one of the k − 1 nearest neighbors of x p i , and that the distance to x p

Original: i ’s k -nearest neighbor δ k (cid:28) R .
Revised: i ’s k -nearest neighbor δ k (cid:28) R .

Original: Then the expected maximum error in row i of the LCN approximation K LCN is
Revised: Then the expected maximum error in row i of the LCN approximation K LCN is

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: with j denoting the index of x p i ’s k -nearest neighbor.
Revised: with j denoting the index of x p i ’s k -nearest neighbor.

Original: Using the upper incomplete Gamma function Γ( ., . )
Revised: Using the upper incomplete Gamma function Γ( ., . )

Original: we can furthermore bound the second term by
Revised: we can furthermore bound the second term by

Original: [Equation]
Revised: [Equation]

Original: The error in Eq.
Revised: The error in Eq.

Original: (6) is dominated by the ﬁrst term since δ k (cid:28) R .
Revised: (6) is dominated by the ﬁrst term since δ k (cid:28) R .

Original: Note that R only decreases slowly with the number of landmarks since R ≥ ( ( d/ 2)!
Revised: Note that R only decreases slowly with the number of landmarks since R ≥ ( ( d/ 2)!

Original: l ) 1 /d 1 2 √ π (Cohn, 2017).
Revised: l ) 1 /d 1 2 √ π (Cohn, 2017).

Original: Moving from pure
Revised: Moving from pure

Original: Nyström to LCN by correcting the nearest neighbors’ entries thus provides signiﬁcant beneﬁts, even for uniform data.
Revised: Nyström to LCN by correcting the nearest neighbors’ entries thus provides signiﬁcant beneﬁts, even for uniform data.

Original: For example, by just correcting the ﬁrst neighbor we obtain a 68 % improvement in the ﬁrst term ( d =32 , λ =0 .
Revised: For example, by just correcting the ﬁrst neighbor we obtain a 68 % improvement in the ﬁrst term ( d =32 , λ =0 .

Original: 05 , n =1000 ).
Revised: 05 , n =1000 ).

Original: This is even more pronounced in clustered data.
Revised: This is even more pronounced in clustered data.

Original: Theorem 2.
Revised: Theorem 2.

Original: Let X p , X q ⊆ R
Revised: Let X p , X q ⊆ R

Original: d be distributed inside the same c clusters with cluster centers x c .
Revised: d be distributed inside the same c clusters with cluster centers x c .

Original: Let r be the maximum L 2 distance of a point to its cluster center and D the minimum distance between two points from different clusters, with r (cid:28) D .
Revised: Let r be the maximum L 2 distance of a point to its cluster center and D the minimum distance between two points from different clusters, with r (cid:28) D .

Original: Let each LSH bucket used for the sparse approximation
Revised: Let each LSH bucket used for the sparse approximation

Original: K sp cover at least one cluster.
Revised: K sp cover at least one cluster.

Original: Let K Nys use 1 ≤ l ≤ d and K LCN use l = 1 optimally distributed landmarks per cluster.
Revised: Let K Nys use 1 ≤ l ≤ d and K LCN use l = 1 optimally distributed landmarks per cluster.

Original: Then the maximum error is
Revised: Then the maximum error is

Original: [Equation]
Revised: [Equation]

Original: Since we can lower bound Eq.
Revised: Since we can lower bound Eq.

Original: (8) by 1 − le − 2 r/λ − O ( e − D/λ ) we can conclude that the error in
Revised: (8) by 1 − le − 2 r/λ − O ( e − D/λ ) we can conclude that the error in

Original: K Nys is close to 1 for any reasonably large rλ (which is the maximum error possible).
Revised: K Nys is close to 1 for any reasonably large rλ (which is the maximum error possible).

Original: The errors in
Revised: The errors in

Original: K sp and K LCN on the other hand are vanishingly small, since r (cid:28) D .
Revised: K sp and K LCN on the other hand are vanishingly small, since r (cid:28) D .

Original: Moreover, these maximum approximation error improvements directly translate to improvements in the Sinkhorn approximation.
Revised: Moreover, these maximum approximation error improvements directly translate to improvements in the Sinkhorn approximation.

Original: We can show this by slightly adapting the error bounds for an approximate Sinkhorn transport plan and distance due to Altschuler et al.
Revised: We can show this by slightly adapting the error bounds for an approximate Sinkhorn transport plan and distance due to Altschuler et al.

Original: Theorem 3 (Altschuler et al.
Revised: Theorem 3 (Altschuler et al.

Original: Let X p , X q ⊆ R
Revised: Let X p , X q ⊆ R

Original: d
Revised: d

Original: have n samples.
Revised: have n samples.

Original: Denote ρ as the maximum distance between two samples.
Revised: Denote ρ as the maximum distance between two samples.

Original: Let ˜ K be an approximation of the similarity matrix K with K ij = e −(cid:107)
Revised: Let ˜ K be an approximation of the similarity matrix K with K ij = e −(cid:107)

Original: x p i − x q j (cid:107) 2 /λ and (cid:107)
Revised: x p i − x q j (cid:107) 2 /λ and (cid:107)

Original: ˜
Revised: ˜

Original: K − K (cid:107) ∞ ≤ ε (cid:48) 

Revised: K − K (cid:107) ∞ ≤ ε (cid:48)

Original: e − ρ/λ , where ε (cid:48) = min(1 , ε 50( ρ + λ log λnε ) ) .
Revised: e − ρ/λ , where ε (cid:48) = min(1 , ε 50( ρ + λ log λnε ) ) .

Original: When performing the Sinkhorn algorithm until (cid:107) ˜ P
Revised: When performing the Sinkhorn algorithm until (cid:107) ˜ P

Original: [Equation]
Revised: [Equation]

Original: N
Revised: N

Original: [Equation]
Revised: [Equation]

Original: p (cid:107) 1 + (cid:107) ˜
Revised: p (cid:107) 1 + (cid:107) ˜

Original: P T 1 N − q
Revised: P T 1 N − q

Original: [Equation]
Revised: [Equation]

Original: 1 ≤ ε (cid:48) / 2 , the resulting approximate transport plan ˜ P and distance ˜
Revised: 1 ≤ ε (cid:48) / 2 , the resulting approximate transport plan ˜ P and distance ˜

Original: d λc are bounded by
Revised: d λc are bounded by

Original: [Equation]
Revised: [Equation]

Original: Convergence rate.
Revised: Convergence rate.

Original: We next show that approximate Sinkhorn converges as fast as regular Sinkhorn by slightly adapting the convergence bound by Dvurechensky et al.
Revised: We next show that approximate Sinkhorn converges as fast as regular Sinkhorn by slightly adapting the convergence bound by Dvurechensky et al.

Original: (2018) to account for sparsity.
Revised: (2018) to account for sparsity.

Original: Theorem 4 (Dvurechensky et al.
Revised: Theorem 4 (Dvurechensky et al.

Original: Given the matrix ˜ K ∈ R n × n and p , q the Sinkhorn algorithm gives a transport plan satisfying (cid:107)
Revised: Given the matrix ˜ K ∈ R n × n and p , q the Sinkhorn algorithm gives a transport plan satisfying (cid:107)

Original: ˜
Revised: ˜

Original: P 1 N − p (cid:107) 

Revised: P 1 N − p (cid:107)

Original: + (cid:107)
Revised: + (cid:107)

Original: [Equation]
Revised: [Equation]

Original: −
Revised: −

Original: [Equation]
Revised: [Equation]

Original: ε in iterations
Revised: ε in iterations

Original: [Equation]
Revised: [Equation]

Original: Backpropagation.
Revised: Backpropagation.

Original: Efﬁcient gradient computation is almost as important for modern deep learning models as the algorithm itself.
Revised: Efﬁcient gradient computation is almost as important for modern deep learning models as the algorithm itself.

Original: These models usually aim at learning the embeddings in X p and X q and therefore need gradients w.r.t.
Revised: These models usually aim at learning the embeddings in X p and X q and therefore need gradients w.r.t.

Original: the cost matrix C .
Revised: the cost matrix C .

Original: We can estimate these either via automatic differentiation of the unrolled Sinkhorn iterations or via the analytic solution that assumes exact convergence.
Revised: We can estimate these either via automatic differentiation of the unrolled Sinkhorn iterations or via the analytic solution that assumes exact convergence.

Original: Depending on the problem at hand, either the automatic or the analytic estimator will lead to faster overall convergence (Ablin et al., 2020).
Revised: Depending on the problem at hand, either the automatic or the analytic estimator will lead to faster overall convergence (Ablin et al., 2020).

Original: LCN-Sinkhorn works ﬂawlessly with automatic backpropagation since it only relies on basic linear algebra (except for choosing Nyström landmarks and LSH neighbors, for which we use a simple straight-through estimator (Bengio et al., 2013)).
Revised: LCN-Sinkhorn works ﬂawlessly with automatic backpropagation since it only relies on basic linear algebra (except for choosing Nyström landmarks and LSH neighbors, for which we use a simple straight-through estimator (Bengio et al., 2013)).

Original: To enable fast analytic backpropagation we provide analytic gradients in Proposition 1.
Revised: To enable fast analytic backpropagation we provide analytic gradients in Proposition 1.

Original: Note that both backpropagation methods have runtime linear in the number of points n and m .
Revised: Note that both backpropagation methods have runtime linear in the number of points n and m .

Original: Proposition 1.
Revised: Proposition 1.

Original: The derivatives of the distances d λc and d λ LCN ,c (Eqs.
Revised: The derivatives of the distances d λc and d λ LCN ,c (Eqs.

Original: (1) and (5) )
Revised: (1) and (5) )

Original: and the optimal transport plan ¯ P ∈ R n × m w.r.t.
Revised: and the optimal transport plan ¯ P ∈ R n × m w.r.t.

Original: the (decomposed) cost matrix C ∈ R n × m in entropy-regularized
Revised: the (decomposed) cost matrix C ∈ R n × m in entropy-regularized

Original: OT and LCN-Sinkhorn are
Revised: OT and LCN-Sinkhorn are

Original: [Equation]
Revised: [Equation]

Original: with δ ij denoting the Kronecker delta and † the Moore-Penrose pseudoinverse.
Revised: with δ ij denoting the Kronecker delta and † the Moore-Penrose pseudoinverse.

Original: Using these decompositions we can backpropagate through LCN-Sinkhorn in time O (( n + m ) l 2 + l 3 ) .
Revised: Using these decompositions we can backpropagate through LCN-Sinkhorn in time O (( n + m ) l 2 + l 3 ) .

Original:  G RAPH T RANSPORT N ETWORK
Revised:  G RAPH T RANSPORT N ETWORK

Original: Graph distance learning.
Revised: Graph distance learning.

Original: The ability to predict similarities or distances between graph-structured objects is useful across a wide range of applications.
Revised: The ability to predict similarities or distances between graph-structured objects is useful across a wide range of applications.

Original: It can e.g. be used to predict the reaction rate between molecules (Houston et al., 2019), search for similar images (Johnson et al., 2015), similar molecules for drug discovery (Birchall et al., 2006), or similar code for vulnerability detection (Li et al., 2019).
Revised: It can e.g. be used to predict the reaction rate between molecules (Houston et al., 2019), search for similar images (Johnson et al., 2015), similar molecules for drug discovery (Birchall et al., 2006), or similar code for vulnerability detection (Li et al., 2019).

Original: We propose the graph transport network (GTN) to evaluate approximate Sinkhorn and advance the state of the art on this task.
Revised: We propose the graph transport network (GTN) to evaluate approximate Sinkhorn and advance the state of the art on this task.

Original: Graph transport network.
Revised: Graph transport network.

Original: GTN ﬁrst uses a Siamese graph neural network (GNN) to embed two graphs independently as sets of node embeddings.
Revised: GTN ﬁrst uses a Siamese graph neural network (GNN) to embed two graphs independently as sets of node embeddings.

Original: These embeddings are then matched using enhanced entropy-regularized optimal transport.
Revised: These embeddings are then matched using enhanced entropy-regularized optimal transport.

Original: Given an undirected graph G =
Revised: Given an undirected graph G =

Original: ( V , E ) , with node set V and edge set E , node attributes x i ∈
Revised: ( V , E ) , with node set V and edge set E , node attributes x i ∈

Original: R H x and (optional) edge attributes e i,j ∈ R H e , with i, j ∈ V , we update the node embeddings in each GNN layer via
Revised: R H x and (optional) edge attributes e i,j ∈ R H e , with i, j ∈ V , we update the node embeddings in each GNN layer via

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: with N i denoting the neighborhood of node i , h (0) i
Revised: with N i denoting the neighborhood of node i , h (0) i

Original: = x i , h ( l ) i ∈
Revised: = x i , h ( l ) i ∈

Original: R H N for l ≥ 1 , the bilinear layer
Revised: R H N for l ≥ 1 , the bilinear layer

Original: W edge ∈ R H N × H N × H e , and the degree normalization η (1) i,j = 1 and η ( l ) i,j = 1 / (cid:112) deg
Revised: W edge ∈ R H N × H N × H e , and the degree normalization η (1) i,j = 1 and η ( l ) i,j = 1 / (cid:112) deg

Original: i deg j for l > 1 .
Revised: i deg j for l > 1 .

Original: This choice of η i,j allows our model to handle highly skewed degree distributions while still being able to represent node degrees.
Revised: This choice of η i,j allows our model to handle highly skewed degree distributions while still being able to represent node degrees.

Original: We found the choice of non-linearity σ not to be critical and chose a
Revised: We found the choice of non-linearity σ not to be critical and chose a

Original: LeakyReLU.
Revised: LeakyReLU.

Original: We do not use the bilinear layer W edge e i,j if there are no edge attributes.
Revised: We do not use the bilinear layer W edge e i,j if there are no edge attributes.

Original: We aggregate each layer’s node embeddings to obtain the ﬁnal embedding of node
Revised: We aggregate each layer’s node embeddings to obtain the ﬁnal embedding of node

Original: i
Revised: i

Original: [Equation]
Revised: [Equation]

Original: Having obtained the embedding sets H ﬁnaland H ﬁnalof both graphs we use the L 2 distance as a cost function and then calculate the Sinkhorn distance, which is symmetric and permutation invariant w.r.t.
Revised: Having obtained the embedding sets H ﬁnaland H ﬁnalof both graphs we use the L 2 distance as a cost function and then calculate the Sinkhorn distance, which is symmetric and permutation invariant w.r.t.

Original: the sets H ﬁnaland H ﬁnal 2 .
Revised: the sets H ﬁnaland H ﬁnal 2 .

Original: We obtain the embeddings for matching via h (0) i =
Revised: We obtain the embeddings for matching via h (0) i =

Original: MLP( h ﬁnal i ) and obtain the ﬁnal prediction via d = d λc w out + b out , with learnable w out and b out .
Revised: MLP( h ﬁnal i ) and obtain the ﬁnal prediction via d = d λc w out + b out , with learnable w out and b out .

Original: All weights in GTN are trained end-to-end via backpropagation.
Revised: All weights in GTN are trained end-to-end via backpropagation.

Original: For small graphs we use the full Sinkhorn distance and scale to large graphs by leveraging LCN-Sinkhorn.
Revised: For small graphs we use the full Sinkhorn distance and scale to large graphs by leveraging LCN-Sinkhorn.

Original: GTN is more expressive than models that aggegrate node embeddings to a single ﬁxed-size embedding for the entire graph but still scales log-linearly in the number of nodes, as opposed to previous approaches that scale quadratically.
Revised: GTN is more expressive than models that aggegrate node embeddings to a single ﬁxed-size embedding for the entire graph but still scales log-linearly in the number of nodes, as opposed to previous approaches that scale quadratically.

Original: Note that GTN inherently performs graph matching and can therefore also be applied to this task.
Revised: Note that GTN inherently performs graph matching and can therefore also be applied to this task.

Original: Learnable unbalanced OT.
Revised: Learnable unbalanced OT.

Original: Since GTN regularly encounters graphs with disagreeing numbers of nodes it needs to be able to handle cases where (cid:107) p (cid:107) 1 (cid:54)
Revised: Since GTN regularly encounters graphs with disagreeing numbers of nodes it needs to be able to handle cases where (cid:107) p (cid:107) 1 (cid:54)

Original: = (cid:107) q (cid:107) 1 or where not all nodes in one graph have a corresponding node in the other and thus P 1 m < p or P T 1 n < q .
Revised: = (cid:107) q (cid:107) 1 or where not all nodes in one graph have a corresponding node in the other and thus P 1 m < p or P T 1 n < q .

Original: Unbalanced OT allows us to handle both of these cases (Peyré & Cuturi, 2019).
Revised: Unbalanced OT allows us to handle both of these cases (Peyré & Cuturi, 2019).

Original: Previous methods did so by swapping these requirements with a uniform divergence loss term on p and q (Frogner et al., 2015; Chizat et al., 2018).
Revised: Previous methods did so by swapping these requirements with a uniform divergence loss term on p and q (Frogner et al., 2015; Chizat et al., 2018).

Original: However, these approaches uniformly penalize deviations from balanced OT and therefore cannot adapt to only ignore parts of the distribution.
Revised: However, these approaches uniformly penalize deviations from balanced OT and therefore cannot adapt to only ignore parts of the distribution.

Original: We propose to alleviate this limitation by swapping the cost matrix C with the bipartite matching (BP) matrix (Riesen & Bunke, 2009)
Revised: We propose to alleviate this limitation by swapping the cost matrix C with the bipartite matching (BP) matrix (Riesen & Bunke, 2009)

Original: C BP
Revised: C BP

Original: [Equation]
Revised: [Equation]

Original: and adaptively computing the costs
Revised: and adaptively computing the costs

Original: c i,ε , c ε,j
Revised: c i,ε , c ε,j

Original: and c ε,ε based on the input sets X p and X q .
Revised: and c ε,ε based on the input sets X p and X q .

Original: Using the
Revised: Using the

Original: BP matrix adds minor computational overhead since we only need to save the diagonals c p ,ε and c ε, q of C p ,ε and C ε, q .
Revised: BP matrix adds minor computational overhead since we only need to save the diagonals c p ,ε and c ε, q of C p ,ε and C ε, q .

Original: We can then include the additional parts of C BP in the Sinkhorn algorithm (Eq.
Revised: We can then include the additional parts of C BP in the Sinkhorn algorithm (Eq.

Original: (2)) via
Revised: (2)) via

Original: [Equation]
Revised: [Equation]

Original: where ˆ t denotes the upper and ˇ t the lower part of the vector t .
Revised: where ˆ t denotes the upper and ˇ t the lower part of the vector t .

Original: To calculate d λc we can decompose the transport plan P BP in the same way as C BP , with a single scalar for P ε,ε .
Revised: To calculate d λc we can decompose the transport plan P BP in the same way as C BP , with a single scalar for P ε,ε .

Original: For GTN we obtain the deletion cost via c
Revised: For GTN we obtain the deletion cost via c

Original: i,ε = (cid:107) α (cid:12) x p i (cid:107) 2 , with a learnable vector α ∈ R d .
Revised: i,ε = (cid:107) α (cid:12) x p i (cid:107) 2 , with a learnable vector α ∈ R d .

Original: Multi-head OT.
Revised: Multi-head OT.

Original: Inspired by attention models (Vaswani et al., 2017) we further improve GTN by using multiple OT heads.
Revised: Inspired by attention models (Vaswani et al., 2017) we further improve GTN by using multiple OT heads.

Original: Using K heads means that we calculate OT in parallel for K separate sets of embeddings representing the same pair of objects and obtain a set of distances d λc ∈ R K .
Revised: Using K heads means that we calculate OT in parallel for K separate sets of embeddings representing the same pair of objects and obtain a set of distances d λc ∈ R K .

Original: We can then transform these distances to a ﬁnal distance prediction using a set of linear layers h (k) i =
Revised: We can then transform these distances to a ﬁnal distance prediction using a set of linear layers h (k) i =

Original: W ( k ) h ﬁnal i for head k and obtain the ﬁnal prediction via d = MLP( d λc ) .
Revised: W ( k ) h ﬁnal i for head k and obtain the ﬁnal prediction via d = MLP( d λc ) .

Original: Note that both learnable unbalanced
Revised: Note that both learnable unbalanced

Original: OT and multi-head OT might be of independent interest.
Revised: OT and multi-head OT might be of independent interest.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Log-linear optimal transport.
Revised: Log-linear optimal transport.

Original: For an overview of optimal transport and its foundations see Peyré & Cuturi (2019).
Revised: For an overview of optimal transport and its foundations see Peyré & Cuturi (2019).

Original: On low-dimensional grids and surfaces OT can be solved using dynamical OT (Papadakis et al., 2014; Solomon et al., 2014), convolutions (Solomon et al., 2015), or embedding/hashing schemes (Indyk & Thaper, 2003; Andoni et al., 2008).
Revised: On low-dimensional grids and surfaces OT can be solved using dynamical OT (Papadakis et al., 2014; Solomon et al., 2014), convolutions (Solomon et al., 2015), or embedding/hashing schemes (Indyk & Thaper, 2003; Andoni et al., 2008).

Original: In higher dimensions we can use tree-based algorithms (Backurs et al., 2020) or hashing schemes (Charikar, 2002), which are however limited to a previously ﬁxed set of points X p , X q , on which only the distributions p and q change.
Revised: In higher dimensions we can use tree-based algorithms (Backurs et al., 2020) or hashing schemes (Charikar, 2002), which are however limited to a previously ﬁxed set of points X p , X q , on which only the distributions p and q change.

Original: For sets that change dynamically (e.g. during training) one common method of achieving log-linear runtime is a multiscale approximation of entropy-regularized OT (Schmitzer, 2019; Gerber & Maggioni, 2017).
Revised: For sets that change dynamically (e.g. during training) one common method of achieving log-linear runtime is a multiscale approximation of entropy-regularized OT (Schmitzer, 2019; Gerber & Maggioni, 2017).

Original: Tenetov et al.
Revised: Tenetov et al.

Original: (2018) recently proposed using a low-rank approximation of the Sinkhorn similarity
Revised: (2018) recently proposed using a low-rank approximation of the Sinkhorn similarity

Original: [Table]
Revised: [Table]

Original: matrix obtained via a semidiscrete approximation of the Euclidean distance.
Revised: matrix obtained via a semidiscrete approximation of the Euclidean distance.

Original: Altschuler et al.
Revised: Altschuler et al.

Original: (2019) improved upon this approach by using the Nyström method for the approximation.
Revised: (2019) improved upon this approach by using the Nyström method for the approximation.

Original: These approaches still struggle with high-dimensional real-world problems, as we will show in Sec. 7.
Revised: These approaches still struggle with high-dimensional real-world problems, as we will show in Sec. 7.

Original: Sliced Wasserstein distance.
Revised: Sliced Wasserstein distance.

Original: Another approach to reduce the computational complexity of optimal transport (without entropy regularization) are sliced Wasserstein distances (Rabin et al., 2011).
Revised: Another approach to reduce the computational complexity of optimal transport (without entropy regularization) are sliced Wasserstein distances (Rabin et al., 2011).

Original: However, they require the L 2 distance as a cost function and are either unstable in convergence or prohibitively expensive for high-dimensional problems ( O ( nd 3 ) )
Revised: However, they require the L 2 distance as a cost function and are either unstable in convergence or prohibitively expensive for high-dimensional problems ( O ( nd 3 ) )

Original: (Meng et al., 2019).
Revised: (Meng et al., 2019).

Original: Fast Sinkhorn.
Revised: Fast Sinkhorn.

Original: Another line of work has been pursuing accelerating entropy-regularized OT without changing its computational complexity w.r.t.
Revised: Another line of work has been pursuing accelerating entropy-regularized OT without changing its computational complexity w.r.t.

Original: the number of points.
Revised: the number of points.

Original: Original Sinkhorn requires
Revised: Original Sinkhorn requires

Original: O (1 /ε 2 ) iterations (Dvurechensky et al., 2018) and Jambulapati et al.
Revised: O (1 /ε 2 ) iterations (Dvurechensky et al., 2018) and Jambulapati et al.

Original: (2019) recently proposed an algorithm that reduces them to O (1 /ε ) .
Revised: (2019) recently proposed an algorithm that reduces them to O (1 /ε ) .

Original: Alaya et al.
Revised: Alaya et al.

Original: (2019) proposed to reduce the size of the Sinkhorn problem by screening out neglectable components, which allows for approximation guarantees.
Revised: (2019) proposed to reduce the size of the Sinkhorn problem by screening out neglectable components, which allows for approximation guarantees.

Original: Genevay et al.
Revised: Genevay et al.

Original: (2016) proposed using a stochastic optimization scheme instead of Sinkhorn iterations.
Revised: (2016) proposed using a stochastic optimization scheme instead of Sinkhorn iterations.

Original: Essid & Solomon (2018) and Blondel et al.
Revised: Essid & Solomon (2018) and Blondel et al.

Original: (2018) proposed alternative regularizations to obtain OT problems with similar runtimes as the Sinkhorn algorithm.
Revised: (2018) proposed alternative regularizations to obtain OT problems with similar runtimes as the Sinkhorn algorithm.

Original: This work is largely orthogonal to ours.
Revised: This work is largely orthogonal to ours.

Original: Embedding alignment.
Revised: Embedding alignment.

Original: For an overview of cross-lingual word embedding models see Ruder et al.
Revised: For an overview of cross-lingual word embedding models see Ruder et al.

Original: Unsupervised word embedding alignment was proposed by Conneau et al.
Revised: Unsupervised word embedding alignment was proposed by Conneau et al.

Original: (2018), with subsequent advances by Alvarez-Melis & Jaakkola (2018)
Revised: (2018), with subsequent advances by Alvarez-Melis & Jaakkola (2018)

Original: ; Grave et al.
Revised: ; Grave et al.

Original: ; Joulin et al.
Revised: ; Joulin et al.

Original: Graph matching and distance learning.
Revised: Graph matching and distance learning.

Original: Most recent approaches for graph matching and graph distance learning either rely on a single ﬁxed-dimensional graph embedding (Bai et al., 2019; Li et al., 2019), or only use attention or some other strongly simpliﬁed variant of optimal transport (Bai et al., 2019;
Revised: Most recent approaches for graph matching and graph distance learning either rely on a single ﬁxed-dimensional graph embedding (Bai et al., 2019; Li et al., 2019), or only use attention or some other strongly simpliﬁed variant of optimal transport (Bai et al., 2019;

Original: Riba et al., 2018; Li et al., 2019).
Revised: Riba et al., 2018; Li et al., 2019).

Original: Others break permutation invariance and are thus ill-suited for this task (Ktena et al., 2017; Bai et al., 2018).
Revised: Others break permutation invariance and are thus ill-suited for this task (Ktena et al., 2017; Bai et al., 2018).

Original: So far only approaches using a single graph embedding allow faster than quadratic scaling in the number of nodes.
Revised: So far only approaches using a single graph embedding allow faster than quadratic scaling in the number of nodes.

Original: Compared to the
Revised: Compared to the

Original: Sinkhorn-based image model concurrently proposed by Wang et al.
Revised: Sinkhorn-based image model concurrently proposed by Wang et al.

Original: GTN uses no CNN or cross-graph attention, but an enhanced GNN and embedding aggregation scheme.
Revised: GTN uses no CNN or cross-graph attention, but an enhanced GNN and embedding aggregation scheme.

Original: OT has recently been proposed for graph kernels (Maretic et al., 2019; Vayer et al., 2019), which can (to a limited extent) be used for graph matching, but not for distance learning.
Revised: OT has recently been proposed for graph kernels (Maretic et al., 2019; Vayer et al., 2019), which can (to a limited extent) be used for graph matching, but not for distance learning.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: Approximating Sinkhorn.
Revised: Approximating Sinkhorn.

Original: We start by directly investigating different Sinkhorn approximations.
Revised: We start by directly investigating different Sinkhorn approximations.

Original: To do so we compute entropy-regularized OT on pairs of 10 000 word embeddings from Conneau et al.
Revised: To do so we compute entropy-regularized OT on pairs of 10 000 word embeddings from Conneau et al.

Original: (2018), which we preprocess with Wasserstein Procrustes alignment in order to obtain both close and distant neighbors.
Revised: (2018), which we preprocess with Wasserstein Procrustes alignment in order to obtain both close and distant neighbors.

Original: We let every method use the same total number of 40 neighbors and landmarks (LCN uses 20 each) and set λ = 0 .
Revised: We let every method use the same total number of 40 neighbors and landmarks (LCN uses 20 each) and set λ = 0 .

Original: 05 (as in Grave et al.
Revised: 05 (as in Grave et al.

Original: We measure transport plan approximation quality by (a) calculating the Pearson correlation coefﬁcient (PCC) between all entries in the approximated plan and the true ¯ P and (b) comparing the sets of 0 .
Revised: We measure transport plan approximation quality by (a) calculating the Pearson correlation coefﬁcient (PCC) between all entries in the approximated plan and the true ¯ P and (b) comparing the sets of 0 .

Original: 1 % largest entries in the approximated and true ¯ P using the Jaccard similarity (intersection over union, IoU).
Revised: 1 % largest entries in the approximated and true ¯ P using the Jaccard similarity (intersection over union, IoU).

Original: In all ﬁgures the error bars denote standard deviation across 5 runs, which is often too small to be visible.
Revised: In all ﬁgures the error bars denote standard deviation across 5 runs, which is often too small to be visible.

Original: Table 1 shows that both sparse Sinkhorn, LCN-Sinkhorn and factored OT (Forrow et al., 2019) obtain distances that are signiﬁcantly closer to the true d λc than Multiscale OT and Nyström-Sinkhorn.
Revised: Table 1 shows that both sparse Sinkhorn, LCN-Sinkhorn and factored OT (Forrow et al., 2019) obtain distances that are signiﬁcantly closer to the true d λc than Multiscale OT and Nyström-Sinkhorn.

Original: Furthermore, the transport plan computed by sparse Sinkhorn and LCN-Sinkhorn show both a PCC and IoU that are around twice as high as Multiscale OT, while Nyström-Sinkhorn and factored OT exhibit almost no correlation.
Revised: Furthermore, the transport plan computed by sparse Sinkhorn and LCN-Sinkhorn show both a PCC and IoU that are around twice as high as Multiscale OT, while Nyström-Sinkhorn and factored OT exhibit almost no correlation.

Original: LCN-Sinkhorn performs especially well in this regard.
Revised: LCN-Sinkhorn performs especially well in this regard.

Original: This is also
Revised: This is also

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Table]
Revised: [Table]

Original: evident in Fig.
Revised: evident in Fig.

Original: [Table]
Revised: [Table]

Original: which shows how the 10 4 × 10 4 approximated OT plan entries compared to the true Sinkhorn values
Revised: which shows how the 10 4 × 10 4 approximated OT plan entries compared to the true Sinkhorn values

Original: [Table]
Revised: [Table]

Original: Fig.
Revised: Fig.

Original: 2 shows that sparse Sinkhorn offers the best trade-off between runtime and OT plan quality.
Revised: 2 shows that sparse Sinkhorn offers the best trade-off between runtime and OT plan quality.

Original: Factored OT exhibits a runtime 2 to 10 times longer than the competition due to its iterative reﬁnement scheme.
Revised: Factored OT exhibits a runtime 2 to 10 times longer than the competition due to its iterative reﬁnement scheme.

Original: LCN-Sinkhorn performs best for use cases with constrained memory (few neighbors/landmarks), as shown in Fig.
Revised: LCN-Sinkhorn performs best for use cases with constrained memory (few neighbors/landmarks), as shown in Fig.

Original: The number of neighbors and landmarks directly determines memory usage and is linearly proportional to the runtime (see App. K).
Revised: The number of neighbors and landmarks directly determines memory usage and is linearly proportional to the runtime (see App. K).

Original: Fig.
Revised: Fig.

Original: 9 shows that sparse
Revised: 9 shows that sparse

Original: Sinkhorn performs best for low regularizations, where LCN-Sinkhorn fails due to the Nyström part going out of bounds.
Revised: Sinkhorn performs best for low regularizations, where LCN-Sinkhorn fails due to the Nyström part going out of bounds.

Original: Nyström Sinkhorn performs best at high values and LCN-Sinkhorn always performs better than both (as long as it can be calculated).
Revised: Nyström Sinkhorn performs best at high values and LCN-Sinkhorn always performs better than both (as long as it can be calculated).

Original: Interestingly, all approximations except factored OT seem to fail at high λ .
Revised: Interestingly, all approximations except factored OT seem to fail at high λ .

Original: We defer analogously discussing the distance approximation to
Revised: We defer analogously discussing the distance approximation to

Original: App.
Revised: App.

Original: L. All approximations scale linearly both in the number of neighbors/landmarks and dataset size, as shown in App.
Revised: L. All approximations scale linearly both in the number of neighbors/landmarks and dataset size, as shown in App.

Original: K. Overall, we see that sparse Sinkhorn and LCN-Sinkhorn yield signiﬁcant improvements over previous approximations.
Revised: K. Overall, we see that sparse Sinkhorn and LCN-Sinkhorn yield signiﬁcant improvements over previous approximations.

Original: However, do these improvements also translate to better performance on downstream tasks?
Revised: However, do these improvements also translate to better performance on downstream tasks?

Original: Embedding alignment.
Revised: Embedding alignment.

Original: Embedding alignment is the task of ﬁnding the orthogonal matrix
Revised: Embedding alignment is the task of ﬁnding the orthogonal matrix

Original: R ∈ R d × d that best aligns the vectors from two different embedding spaces, which is e.g. useful for unsupervised word translation.
Revised: R ∈ R d × d that best aligns the vectors from two different embedding spaces, which is e.g. useful for unsupervised word translation.

Original: We use the experimental setup established by Conneau et al.
Revised: We use the experimental setup established by Conneau et al.

Original: by migrating
Revised: by migrating

Original: Grave et al.
Revised: Grave et al.

Original: (2019)’s implementation to PyTorch.
Revised: (2019)’s implementation to PyTorch.

Original: The only change we make is using the full set of 20 000 word embeddings and training for 300 steps, while reducing the learning rate by half every 100 steps.
Revised: The only change we make is using the full set of 20 000 word embeddings and training for 300 steps, while reducing the learning rate by half every 100 steps.

Original: We do not change any other hyperparameters and do not use unbalanced OT.
Revised: We do not change any other hyperparameters and do not use unbalanced OT.

Original: After training we match pairs via cross-domain similarity local scaling (CSLS) (Conneau et al., 2018).
Revised: After training we match pairs via cross-domain similarity local scaling (CSLS) (Conneau et al., 2018).

Original: We use 10 Sinkhorn iterations, 40 neighbors for sparse Sinkhorn, and 20 neighbors and landmarks for
Revised: We use 10 Sinkhorn iterations, 40 neighbors for sparse Sinkhorn, and 20 neighbors and landmarks for

Original: LCN-Sinkhorn (for details see App. H).
Revised: LCN-Sinkhorn (for details see App. H).

Original: We allow both multiscale OT and Nyström Sinkhorn to use as many landmarks and neighbors as can ﬁt into GPU memory and ﬁnetune both methods.
Revised: We allow both multiscale OT and Nyström Sinkhorn to use as many landmarks and neighbors as can ﬁt into GPU memory and ﬁnetune both methods.

Original: Table 2 shows that using full Sinkhorn yields a signiﬁcant improvement in accuracy on this task
Revised: Table 2 shows that using full Sinkhorn yields a signiﬁcant improvement in accuracy on this task

Original: [Table]
Revised: [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: compared to the original approach of performing Sinkhorn on randomly sampled subsets of embeddings (Grave et al., 2019).
Revised: compared to the original approach of performing Sinkhorn on randomly sampled subsets of embeddings (Grave et al., 2019).

Original: LCN-Sinkhorn even outperforms the full version in most cases
Revised: LCN-Sinkhorn even outperforms the full version in most cases

Original: [Table]
Revised: [Table]

Original: which is likely due to regularization effects from the approximation.
Revised: which is likely due to regularization effects from the approximation.

Original: It also runs 4.6x faster than full Sinkhorn and 3.1x faster than the cannot match its accuracy.
Revised: It also runs 4.6x faster than full Sinkhorn and 3.1x faster than the cannot match its accuracy.

Original: LCN-Sinkhorn still outcompetes the original method after reﬁning the embeddings with iterative local CSLS (Conneau et al., 2018).
Revised: LCN-Sinkhorn still outcompetes the original method after reﬁning the embeddings with iterative local CSLS (Conneau et al., 2018).

Original: Both multiscale OT and Nyström
Revised: Both multiscale OT and Nyström

Original: Sinkhorn fail at this task, despite their larger computational budget.
Revised: Sinkhorn fail at this task, despite their larger computational budget.

Original: This shows that the improvements achieved by sparse Sinkhorn and LCN-Sinkhorn have an even larger impact in practice.
Revised: This shows that the improvements achieved by sparse Sinkhorn and LCN-Sinkhorn have an even larger impact in practice.

Original: Graph distance regression.
Revised: Graph distance regression.

Original: The graph edit distance (GED) is useful for various tasks, such as image retrieval (Xiao et al., 2008) or ﬁngerprint matching (Neuhaus & Bunke, 2004), but its computation is NP-complete (Bunke & Shearer, 1998).
Revised: The graph edit distance (GED) is useful for various tasks, such as image retrieval (Xiao et al., 2008) or ﬁngerprint matching (Neuhaus & Bunke, 2004), but its computation is NP-complete (Bunke & Shearer, 1998).

Original: Therefore, to use it on larger graphs we need to learn an approximation.
Revised: Therefore, to use it on larger graphs we need to learn an approximation.

Original: We use the Linux dataset by Bai et al.
Revised: We use the Linux dataset by Bai et al.

Original: and generate 2 new datasets by computing the exact GED using the method by Lerouge et al.
Revised: and generate 2 new datasets by computing the exact GED using the method by Lerouge et al.

Original: (2017) on small graphs ( ≤ 30 nodes) from the AIDS dataset (Riesen & Bunke, 2008) and a set of preferential attachment graphs.
Revised: (2017) on small graphs ( ≤ 30 nodes) from the AIDS dataset (Riesen & Bunke, 2008) and a set of preferential attachment graphs.

Original: We compare GTN to 3 state-of-the-art baselines: SiameseMPNN (Riba et al., 2018), SimGNN (Bai et al., 2019), and the Graph Matching Network (GMN)
Revised: We compare GTN to 3 state-of-the-art baselines: SiameseMPNN (Riba et al., 2018), SimGNN (Bai et al., 2019), and the Graph Matching Network (GMN)

Original: (Li et al., 2019).
Revised: (Li et al., 2019).

Original: We tune the hyperparameters of all baselines and GTN on the validation set via a grid search.
Revised: We tune the hyperparameters of all baselines and GTN on the validation set via a grid search.

Original: For more details see App.
Revised: For more details see App.

Original: H to J.
Revised: H to J.

Original: We ﬁrst test both GTN and the proposed OT enhancements.
Revised: We ﬁrst test both GTN and the proposed OT enhancements.

Original: Table 3 shows that GTN improves upon competing models by 20 % with a single head and by 48 % with 8 OT heads.
Revised: Table 3 shows that GTN improves upon competing models by 20 % with a single head and by 48 % with 8 OT heads.

Original: These improvements break down when using regular balanced OT, showing the importance of learnable unbalanced OT.
Revised: These improvements break down when using regular balanced OT, showing the importance of learnable unbalanced OT.

Original: Having established GTN as a state-of-the-art model we next ask whether we can sustain its performance when using approximate OT.
Revised: Having established GTN as a state-of-the-art model we next ask whether we can sustain its performance when using approximate OT.

Original: To test this we additionally generate a set of larger graphs with around 200 nodes and use the Pyramid matching (PM) kernel (Nikolentzos et al., 2017) as the prediction target, since these graphs are too large to compute the GED.
Revised: To test this we additionally generate a set of larger graphs with around 200 nodes and use the Pyramid matching (PM) kernel (Nikolentzos et al., 2017) as the prediction target, since these graphs are too large to compute the GED.

Original: See App.
Revised: See App.

Original: J for hyperparameter details.
Revised: J for hyperparameter details.

Original: Table 4 shows that both sparse Sinkhorn and the multiscale method using 4 (expected) neighbors fail at this task, demonstrating that the low-rank approximation in LCN has a crucial stabilizing effect during training.
Revised: Table 4 shows that both sparse Sinkhorn and the multiscale method using 4 (expected) neighbors fail at this task, demonstrating that the low-rank approximation in LCN has a crucial stabilizing effect during training.

Original: Nyström Sinkhorn with 4 landmarks performs surprisingly well on the AIDS dataset, suggesting an overall low-rank structure with Nyström acting as regularization.
Revised: Nyström Sinkhorn with 4 landmarks performs surprisingly well on the AIDS dataset, suggesting an overall low-rank structure with Nyström acting as regularization.

Original: However, it does not perform as well on the other two datasets.
Revised: However, it does not perform as well on the other two datasets.

Original: Using LCN-Sinkhorn with 2 neighbors and landmarks works well on all three datasets, with an RMSE increased by only 10 % compared to full
Revised: Using LCN-Sinkhorn with 2 neighbors and landmarks works well on all three datasets, with an RMSE increased by only 10 % compared to full

Original: GTN.
Revised: GTN.

Original: App.
Revised: App.

Original: K furthermore shows that GTN with LCN-Sinkhorn indeed scales linearly in the number of nodes across multiple orders of magnitude.
Revised: K furthermore shows that GTN with LCN-Sinkhorn indeed scales linearly in the number of nodes across multiple orders of magnitude.

Original: This model thus allows to perform graph matching and distance learning on graphs that are considered large even for simple node-level tasks ( 20 000 nodes).
Revised: This model thus allows to perform graph matching and distance learning on graphs that are considered large even for simple node-level tasks ( 20 000 nodes).

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: Locality sensitive hashing (LSH) and the novel locally corrected Nyström (LCN) method enable fast and accurate approximations of entropy-regularized OT with log-linear runtime: Sparse Sinkhorn and LCN-Sinkhorn.
Revised: Locality sensitive hashing (LSH) and the novel locally corrected Nyström (LCN) method enable fast and accurate approximations of entropy-regularized OT with log-linear runtime: Sparse Sinkhorn and LCN-Sinkhorn.

Original: The graph transport network (GTN) is one example for such a model, which can be substantially improved with learnable unbalanced OT and multi-head OT.
Revised: The graph transport network (GTN) is one example for such a model, which can be substantially improved with learnable unbalanced OT and multi-head OT.

Original: It sets the new state of the art for graph distance learning while still scaling log-linearly with graph size.
Revised: It sets the new state of the art for graph distance learning while still scaling log-linearly with graph size.

Original: These contributions enable new applications and models that are both faster and more accurate, since they can sidestep workarounds such as pooling.
Revised: These contributions enable new applications and models that are both faster and more accurate, since they can sidestep workarounds such as pooling.

Original: A C OMPLEXITY A NALYSIS
Revised: A C OMPLEXITY A NALYSIS

Original: Sparse Sinkhorn.
Revised: Sparse Sinkhorn.

Original: A common way of achieving a high p 1 and low p 2 in LSH is via the AND-OR construction.
Revised: A common way of achieving a high p 1 and low p 2 in LSH is via the AND-OR construction.

Original: In this scheme we calculate B · r hash functions, divided into B sets (hash bands) of r hash functions each.
Revised: In this scheme we calculate B · r hash functions, divided into B sets (hash bands) of r hash functions each.

Original: A pair of points is considered as neighbors if any hash band matches completely.
Revised: A pair of points is considered as neighbors if any hash band matches completely.

Original: Calculating the hash buckets for all points with b hash buckets per function scales as O (( n + m ) dBbr ) for the hash functions we consider.
Revised: Calculating the hash buckets for all points with b hash buckets per function scales as O (( n + m ) dBbr ) for the hash functions we consider.

Original: As expected, for the tasks and hash functions we investigated we obtain approximately m/b r and n/b r neighbors, with b r hash buckets per band.
Revised: As expected, for the tasks and hash functions we investigated we obtain approximately m/b r and n/b r neighbors, with b r hash buckets per band.

Original: Using this we can ﬁx the number of neighbors to a small, constant β in expectation with b r = min( n, m ) /β .
Revised: Using this we can ﬁx the number of neighbors to a small, constant β in expectation with b r = min( n, m ) /β .

Original: We thus obtain a sparse cost matrix C sp with O (max( n, m ) β ) non-inﬁnite values and can calculate s and t in linear time
Revised: We thus obtain a sparse cost matrix C sp with O (max( n, m ) β ) non-inﬁnite values and can calculate s and t in linear time

Original: O
Revised: O

Original: ( N sink max( n, m ) β ) , where N sink ≤ 2 + − 4ln(min i,j { ˜ K ij | ˜ K ij > 0 } min i,j { p i , q j } )
Revised: ( N sink max( n, m ) β ) , where N sink ≤ 2 + − 4ln(min i,j { ˜ K ij | ˜ K ij > 0 } min i,j { p i , q j } )

Original: ε (see Theorem 4) denotes the number of Sinkhorn iterations.
Revised: ε (see Theorem 4) denotes the number of Sinkhorn iterations.

Original: Calculating the hash buckets with r = logmin( n,m ) − log β log b takes O (( n + m ) dBb (log min( n, m )
Revised: Calculating the hash buckets with r = logmin( n,m ) − log β log b takes O (( n + m ) dBb (log min( n, m )

Original: − log β ) / log b ) .
Revised: − log β ) / log b ) .

Original: Since B , b , and β are small, we obtain roughly log-linear scaling with the number of points overall, i.e. O ( n log n ) for n ≈ m .
Revised: Since B , b , and β are small, we obtain roughly log-linear scaling with the number of points overall, i.e. O ( n log n ) for n ≈ m .

Original: LCN-Sinkhorn.
Revised: LCN-Sinkhorn.

Original: Both choosing landmarks via k-means++ sampling and via k-means with a ﬁxed number of iterations have the same runtime complexity of O (( n + m ) ld ) .
Revised: Both choosing landmarks via k-means++ sampling and via k-means with a ﬁxed number of iterations have the same runtime complexity of O (( n + m ) ld ) .

Original: Precomputing W can be done in time
Revised: Precomputing W can be done in time

Original: O
Revised: O

Original: ( nl 2 + l 3 ) .
Revised: ( nl 2 + l 3 ) .

Original: The low-rank part of updating the vectors s and t can be computed in
Revised: The low-rank part of updating the vectors s and t can be computed in

Original: O ( nl + l 2 + lm ) , with l chosen constant, i.e. independently of n and m .
Revised: O ( nl + l 2 + lm ) , with l chosen constant, i.e. independently of n and m .

Original: Since sparse Sinkhorn with
Revised: Since sparse Sinkhorn with

Original: LSH has a log-linear runtime we again obtain log-linear overall runtime for LCN-Sinkhorn.
Revised: LSH has a log-linear runtime we again obtain log-linear overall runtime for LCN-Sinkhorn.

Original: B L IMITATIONS
Revised: B L IMITATIONS

Original: Sparse Sinkhorn.
Revised: Sparse Sinkhorn.

Original: Using a sparse approximation for K works well in the common case when the regularization parameter λ is low and the cost function varies enough between data pairs, such that the transport plan P resembles a sparse matrix.
Revised: Using a sparse approximation for K works well in the common case when the regularization parameter λ is low and the cost function varies enough between data pairs, such that the transport plan P resembles a sparse matrix.

Original: However, it can fail if the cost between pairs is very similar or the regularization is very high, if the dataset contains many hubs, i.e. points with a large number of neighbors, or if the distributions p or q are spread very unevenly.
Revised: However, it can fail if the cost between pairs is very similar or the regularization is very high, if the dataset contains many hubs, i.e. points with a large number of neighbors, or if the distributions p or q are spread very unevenly.

Original: Furthermore, sparse
Revised: Furthermore, sparse

Original: Sinkhorn can be too unstable to train a model from scratch, since randomly initialized embeddings often have no close neighbors (see Sec. 7).
Revised: Sinkhorn can be too unstable to train a model from scratch, since randomly initialized embeddings often have no close neighbors (see Sec. 7).

Original: LCN-Sinkhorn largely alleviates these limitations.
Revised: LCN-Sinkhorn largely alleviates these limitations.

Original: LCN-Sinkhorn.
Revised: LCN-Sinkhorn.

Original: Since we cannot calculate the full cost matrix, LCN-Sinkhorn cannot provide accuracy guarantees in general.
Revised: Since we cannot calculate the full cost matrix, LCN-Sinkhorn cannot provide accuracy guarantees in general.

Original: Highly concentrated distributions p and q might have adverse effects on LCN-Sinkhorn.
Revised: Highly concentrated distributions p and q might have adverse effects on LCN-Sinkhorn.

Original: However, we can compensate for these by sampling landmarks or neighbors proportional to each point’s probability mass.
Revised: However, we can compensate for these by sampling landmarks or neighbors proportional to each point’s probability mass.

Original: We therefore do not expect LCN-Sinkhorn to break down in this scenario.
Revised: We therefore do not expect LCN-Sinkhorn to break down in this scenario.

Original: If the regularization parameter is low or the cost function varies greatly, we sometimes observed stability issues (over- and underﬂows) with the Nyström approximation because of the inverse
Revised: If the regularization parameter is low or the cost function varies greatly, we sometimes observed stability issues (over- and underﬂows) with the Nyström approximation because of the inverse

Original: A − 1 , which cannot be calculated in log-space.
Revised: A − 1 , which cannot be calculated in log-space.

Original: Due to its linearity the Nyström method furthermore sometimes approximates similarities as negative values, which leads to a failure if the result of the matrix product in Eq.
Revised: Due to its linearity the Nyström method furthermore sometimes approximates similarities as negative values, which leads to a failure if the result of the matrix product in Eq.

Original: (2) becomes negative.
Revised: (2) becomes negative.

Original: In these extreme cases we also observed catastrophic elimination caused by the correction K sp ∆ .
Revised: In these extreme cases we also observed catastrophic elimination caused by the correction K sp ∆ .

Original: Since this essentially means that optimal transport will be very local, we recommend using sparse Sinkhorn in these scenarios.
Revised: Since this essentially means that optimal transport will be very local, we recommend using sparse Sinkhorn in these scenarios.

Original: This again demonstrates the complementarity of the sparse approximation and Nyström:
Revised: This again demonstrates the complementarity of the sparse approximation and Nyström:

Original: In cases where one fails we can often resort to the other.
Revised: In cases where one fails we can often resort to the other.

Original: C P ROOF OF T HEOREM
Revised: C P ROOF OF T HEOREM

Original: We ﬁrst prove a lemma that will be useful later on.
Revised: We ﬁrst prove a lemma that will be useful later on.

Original: Lemma A. Let ˜ K be the Nyström approximation of the similarity matrix K ij = e −(cid:107) x i − x j (cid:107) 2 /λ .
Revised: Lemma A. Let ˜ K be the Nyström approximation of the similarity matrix K ij = e −(cid:107) x i − x j (cid:107) 2 /λ .

Original: Let x i and x j be data points with equal L 2 distance r i and r j to all l landmarks, which have the same distance ∆ to each other.
Revised: Let x i and x j be data points with equal L 2 distance r i and r j to all l landmarks, which have the same distance ∆ to each other.

Original: Then
Revised: Then

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: The inter-landmark distance matrix is
Revised: The inter-landmark distance matrix is

Original: [Equation]
Revised: [Equation]

Original: where 1 l × l denotes the constant 1 matrix.
Revised: where 1 l × l denotes the constant 1 matrix.

Original: Using the identity
Revised: Using the identity

Original: [Equation]
Revised: [Equation]

Original: we compute
Revised: we compute

Original: [Equation]
Revised: [Equation]

Original: Now consider the error (cid:107) K
Revised: Now consider the error (cid:107) K

Original: i, : − K LCN ,i, : (cid:107) ∞ .
Revised: i, : − K LCN ,i, : (cid:107) ∞ .

Original: The k − 1 nearest neighbors are covered by the sparse correction and therefore the next nearest neighbor has distance δ k .
Revised: The k − 1 nearest neighbors are covered by the sparse correction and therefore the next nearest neighbor has distance δ k .

Original: The expected distance from a landmark is greater than the expected distance inside the surrounding d -ball of radius R , i.e. E [ r ] ≥ E V ( R )
Revised: The expected distance from the closest landmark is greater than the expected distance inside the surrounding d -ball of radius R , i.e. E [ r ] ≥ E V ( R )

Original: [ r ] = dd +1 R .
Revised: [ r ] = dd +1 R .

Original: Because furthermore δ k (cid:28) R , the error is dominated by the ﬁrst term and the maximum error in row i is given by the k -nearest neighbor of i , denoted by j .
Revised: Because furthermore δ k (cid:28) R , the error is dominated by the ﬁrst term and the maximum error in row i is given by the k -nearest neighbor of i , denoted by j .

Original: Thus
Revised: Thus

Original: E [ (cid:107) K i, : − K LCN ,
Revised: E [ (cid:107) K i, : − K LCN ,

Original: i
Revised: i

Original: [Equation]
Revised: [Equation]

Original: : (cid:107)
Revised: : (cid:107)

Original: [Equation]
Revised: [Equation]

Original: Note that we can lower bound the ﬁrst term using Jensen’s inequality.
Revised: Note that we can lower bound the ﬁrst term using Jensen’s inequality.

Original: However, we were unable to ﬁnd a reasonably tight upper bound and the resulting integral (ignoring exponentially small boundary effects, see Percus & Martin (1998))
Revised: However, we were unable to ﬁnd a reasonably tight upper bound and the resulting integral (ignoring exponentially small boundary effects, see Percus & Martin (1998))

Original: [Equation]
Revised: [Equation]

Original: with the volume of the d -ball
Revised: with the volume of the d -ball

Original: [Equation]
Revised: [Equation]

Original: does not have an analytical solution.
Revised: does not have an analytical solution.

Original: We thus have to resort to calculating this expectation numerically.
Revised: We thus have to resort to calculating this expectation numerically.

Original: We lower bound the second term by (1) ignoring every landmark except the closest one, since additional landmarks can only increase the estimate K LCN ,i,j .
Revised: We lower bound the second term by (1) ignoring every landmark except the closest one, since additional landmarks can only increase the estimate K LCN ,i,j .

Original: We then (2) upper bound the Ldistance to the closest landmark r by √ dR/ 2 , since this would be the furthest distance to the closest point in a d -dimensional grid.
Revised: We then (2) upper bound the Ldistance to the closest landmark r by √ dR/ 2 , since this would be the furthest distance to the closest point in a d -dimensional grid.

Original: Any optimal arrangement minimizing E [min y ∈ L (cid:107) x − y (cid:107) 2 | x ∈ X p ] would be at least as good as a grid and thus have furthest distances as small or smaller than those in a grid.
Revised: Any optimal arrangement minimizing E [min y ∈ L (cid:107) x − y (cid:107) 2 | x ∈ X p ] would be at least as good as a grid and thus have furthest distances as small or smaller than those in a grid.

Original: Thus,
Revised: Thus,

Original: [Equation]
Revised: [Equation]

Original: We upper bound this expectation by considering that any point outside the inscribed sphere of the space closest to a landmark (which has radius R ) would be further away from the landmarks and thus have a lower value e − d/λ .
Revised: We upper bound this expectation by considering that any point outside the inscribed sphere of the space closest to a landmark (which has radius R ) would be further away from the landmarks and thus have a lower value e − d/λ .

Original: We can therefore reduce the space over which the expectation is taken to the ball with radius R , i.e.
Revised: We can therefore reduce the space over which the expectation is taken to the ball with radius R , i.e.

Original: [Equation]
Revised: [Equation]

Original: Next we (1) ignore the contributions of all landmarks except for the closest 2, since a third landmark must be further away from the data point than √ 3 R , adding an error of O ( e − 2 √ 3 R/λ ) .
Revised: Next we (1) ignore the contributions of all landmarks except for the closest 2, since a third landmark must be further away from the data point than √ 3 R , adding an error of O ( e − 2 √ 3 R/λ ) .

Original: We then (2) lower bound the distances of both points to both landmarks by the closest distance to a landmark r =
Revised: We then (2) lower bound the distances of both points to both landmarks by the closest distance to a landmark r =

Original: [Equation]
Revised: [Equation]

Original: {(cid:107)
Revised: {(cid:107)

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: x
Revised: x

Original: [Equation]
Revised: [Equation]

Original: − x l 2 (cid:107) 2 } and use Lemma A to obtain
Revised: − x l 2 (cid:107) 2 } and use Lemma A to obtain

Original: [Equation]
Revised: [Equation]

Original: Assuming Euclideanness in V ( R ) we obtain
Revised: Assuming Euclideanness in V ( R ) we obtain

Original: [Equation]
Revised: [Equation]

Original: D P ROOF OF T HEOREM
Revised: D P ROOF OF T HEOREM

Original: Note that this theorem does not use probabilistic arguments but rather geometrically analyzes the maximum possible error.
Revised: Note that this theorem does not use probabilistic arguments but rather geometrically analyzes the maximum possible error.

Original: K sp is correct for all pairs inside a cluster and 0 otherwise.
Revised: K sp is correct for all pairs inside a cluster and 0 otherwise.

Original: We therefore obtain the maximum error by considering the closest possible pair between clusters.
Revised: We therefore obtain the maximum error by considering the closest possible pair between clusters.

Original: By deﬁnition, this pair has distance D and thus
Revised: By deﬁnition, this pair has distance D and thus

Original: [Equation]
Revised: [Equation]

Original: LCN is also correct for all pairs inside a cluster, so we again consider the closest possible pair x i , x j between clusters.
Revised: LCN is also correct for all pairs inside a cluster, so we again consider the closest possible pair x i , x j between clusters.

Original: We furthermore only consider the landmarks of the two concerned clusters, adding an error of O ( e − D/λ ) .
Revised: We furthermore only consider the landmarks of the two concerned clusters, adding an error of O ( e − D/λ ) .

Original: Hence,
Revised: Hence,

Original: K LCN,
Revised: K LCN,

Original: [Equation]
Revised: [Equation]

Original: landmarks ,
Revised: landmarks ,

Original: [Equation]
Revised: [Equation]

Original: 33) and thus
Revised: 33) and thus

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: For pure Nyström we need to consider the distances inside a cluster.
Revised: For pure Nyström we need to consider the distances inside a cluster.

Original: In the worst case two points overlap, i.e. K ij = 1 , and lie at the boundary of the cluster.
Revised: In the worst case two points overlap, i.e. K ij = 1 , and lie at the boundary of the cluster.

Original: Since r (cid:28) D we again only consider the landmarks in the concerned cluster, adding an error of O ( e − D/λ ) .
Revised: Since r (cid:28) D we again only consider the landmarks in the concerned cluster, adding an error of O ( e − D/λ ) .

Original: Because of symmetry we can optimize the worst-case distance from all landmarks by putting them on a ( l − 1) -simplex centered on the cluster center.
Revised: Because of symmetry we can optimize the worst-case distance from all landmarks by putting them on a ( l − 1) -simplex centered on the cluster center.

Original: Since there are at most d landmarks in each cluster there is always one direction in which the worst-case points are r away from all landmarks.
Revised: Since there are at most d landmarks in each cluster there is always one direction in which the worst-case points are r away from all landmarks.

Original: The circumradius of an ( l − 1) -simplex with side length ∆ is (cid:113) l − 1 2 l ∆ .
Revised: The circumradius of an ( l − 1) -simplex with side length ∆ is (cid:113) l − 1 2 l ∆ .

Original: Thus, the maximum distance to all landmarks is (cid:113) r 2 + l −
Revised: Thus, the maximum distance to all landmarks is (cid:113) r 2 + l −

Original: [Equation]
Revised: [Equation]

Original: 2 l ∆ 2 .
Revised: 2 l ∆ 2 .

Original: Using Lemma A we therefore obtain the Nyström approximation
Revised: Using Lemma A we therefore obtain the Nyström approximation

Original: [Equation]
Revised: [Equation]

Original: E N OTE ON T HEOREM
Revised: E N OTE ON T HEOREM

Original: Lemmas C-F and and thus Theorem 1 by Altschuler et al.
Revised: Lemmas C-F and and thus Theorem 1 by Altschuler et al.

Original: (2019) are also valid for Q outside the simplex so long as (cid:107) Q (cid:107) 1 = n
Revised: (2019) are also valid for Q outside the simplex so long as (cid:107) Q (cid:107) 1 = n

Original: and it only has non-negative entries.
Revised: and it only has non-negative entries.

Original: Any ˜ P returned by Sinkhorn fulﬁlls these conditions.
Revised: Any ˜ P returned by Sinkhorn fulﬁlls these conditions.

Original: Therefore the rounding procedure given by their Algorithm 4 is not necessary for this result.
Revised: Therefore the rounding procedure given by their Algorithm 4 is not necessary for this result.

Original: Furthermore, to be more consistent with Theorems 1 and 2 we use the L 2 distance instead of L 22 in this theorem, which only changes the dependence on ρ .
Revised: Furthermore, to be more consistent with Theorems 1 and 2 we use the L 2 distance instead of L 22 in this theorem, which only changes the dependence on ρ .

Original: F N OTES ON T HEOREM
Revised: F N OTES ON T HEOREM

Original: To adapt Theorem 1 by Dvurechensky et al.
Revised: To adapt Theorem 1 by Dvurechensky et al.

Original: (2018) to sparse matrices (i.e. matrices with some
Revised: (2018) to sparse matrices (i.e. matrices with some

Original: K ij = 0 ) we need to redeﬁne
Revised: K ij = 0 ) we need to redeﬁne

Original: [Equation]
Revised: [Equation]

Original: (36) i.e. take the minimum only w.r.t.
Revised: (36) i.e. take the minimum only w.r.t.

Original: non-zero elements in their Lemma 1.
Revised: non-zero elements in their Lemma 1.

Original: G P ROOF OF P ROPOSITION
Revised: G P ROOF OF P ROPOSITION

Original: Theorem A (Danskin’s theorem) .
Revised: Theorem A (Danskin’s theorem) .

Original: Consider a continuous function φ : R k × Z → R , with the compact set Z ⊂ R j .
Revised: Consider a continuous function φ : R k × Z → R , with the compact set Z ⊂ R j .

Original: If φ ( x , z ) is convex in x for every z ∈ Z and φ ( x , z ) has a unique maximizer ¯ z , the derivative of
Revised: If φ ( x , z ) is convex in x for every z ∈ Z and φ ( x , z ) has a unique maximizer ¯ z , the derivative of

Original: [Equation]
Revised: [Equation]

Original: (37) is given by the derivative at the maximizer, i.e.
Revised: (37) is given by the derivative at the maximizer, i.e.

Original: ∂f ∂
Revised: ∂f ∂

Original: [Equation]
Revised: [Equation]

Original: ∂φ ( x
Revised: ∂φ ( x

Original: [Equation]
Revised: [Equation]

Original: ¯ z ) ∂ x .
Revised: ¯ z ) ∂ x .

Original: We start by deriving the derivatives of the distances.
Revised: We start by deriving the derivatives of the distances.

Original: To show that the Sinkhorn distance fulﬁlls the conditions for Danskin’s theorem we ﬁrst identify x = C , z = P , and φ ( C , P )
Revised: To show that the Sinkhorn distance fulﬁlls the conditions for Danskin’s theorem we ﬁrst identify x = C , z = P , and φ ( C , P )

Original: = −(cid:104) P , C (cid:105) F + λH ( P ) .
Revised: = −(cid:104) P , C (cid:105) F + λH ( P ) .

Original: We next observe that the restrictions P 1 m = p and P T 1 n = q deﬁne a compact, convex set for P .
Revised: We next observe that the restrictions P 1 m = p and P T 1 n = q deﬁne a compact, convex set for P .

Original: Furthermore, φ is a continuous function and linear in C , i.e. both convex and concave for any ﬁnite P .
Revised: Furthermore, φ is a continuous function and linear in C , i.e. both convex and concave for any ﬁnite P .

Original: Finally, φ ( C , P ) is concave in P since (cid:104) P , C (cid:105) F is linear and λH ( P ) is concave.
Revised: Finally, φ ( C , P ) is concave in P since (cid:104) P , C (cid:105) F is linear and λH ( P ) is concave.

Original: Therefore the maximizer ¯ P is unique and Danskin’s theorem applies to the Sinkhorn distance.
Revised: Therefore the maximizer ¯ P is unique and Danskin’s theorem applies to the Sinkhorn distance.

Original: Using
Revised: Using

Original: [Equation]
Revised: [Equation]

Original: and the chain rule we can calculate the derivative w.r.t.
Revised: and the chain rule we can calculate the derivative w.r.t.

Original: the cost matrix as
Revised: the cost matrix as

Original: [Equation]
Revised: [Equation]

Original: and ∂d λ LCN ,c ∂ log K sp and ∂d λ LCN ,c ∂ log K sp Nys directly follow from ∂d λc ∂ C
Revised: and ∂d λ LCN ,c ∂ log K sp and ∂d λ LCN ,c ∂ log K sp Nys directly follow from ∂d λc ∂ C

Original: [Equation]
Revised: [Equation]

Original: Theorem B (Implicit function theorem) .
Revised: Theorem B (Implicit function theorem) .

Original: Let f : R n (cid:48)
Revised: Let f : R n (cid:48)

Original: × R m (cid:48)
Revised: × R m (cid:48)

Original: → R m (cid:48) be a continuously differentiable function with f ( a , b ) = 0 .
Revised: → R m (cid:48) be a continuously differentiable function with f ( a , b ) = 0 .

Original: If its Jacobian matrix J f i ( x , y ) , y j = ∂f i ∂ y j ( a , b ) is invertible, then there exists an open set a ∈ U ⊂ R n (cid:48) on which there exists a unique continuously differentiable function g : U → R m (cid:48) with g ( a )
Revised: If its Jacobian matrix J f i ( x , y ) , y j = ∂f i ∂ y j ( a , b ) is invertible, then there exists an open set a ∈ U ⊂ R n (cid:48) on which there exists a unique continuously differentiable function g : U → R m (cid:48) with g ( a )

Original: = b and ∀ x ∈ U : f
Revised: = b and ∀ x ∈ U : f

Original: ( x , g ( x )) = 0 .
Revised: ( x , g ( x )) = 0 .

Original: Moreover,
Revised: Moreover,

Original: [Equation]
Revised: [Equation]

Original: Next we derive the transport plan derivatives.
Revised: Next we derive the transport plan derivatives.

Original: To apply the implicit function theorem we identify x = C and y
Revised: To apply the implicit function theorem we identify x = C and y

Original: = P as ﬂattened matrices.
Revised: = P as ﬂattened matrices.

Original: We will index these ﬂat matrices via index pairs to simplify interpretation.
Revised: We will index these ﬂat matrices via index pairs to simplify interpretation.

Original: We furthermore identify
Revised: We furthermore identify

Original: [Equation]
Revised: [Equation]

Original: The minimizer ¯ P cannot lie on the boundary of P ’s valid region since lim p → 0 ∂∂p p log p =
Revised: The minimizer ¯ P cannot lie on the boundary of P ’s valid region since lim p → 0 ∂∂p p log p =

Original: lim p → 0 log p + 1 =
Revised: lim p → 0 log p + 1 =

Original: −∞ and therefore ¯ P ij > 0 .
Revised: −∞ and therefore ¯ P ij > 0 .

Original: Hence, f ( C , ¯ P ) = 0 with ¯ P ( C ) =
Revised: Hence, f ( C , ¯ P ) = 0 with ¯ P ( C ) =

Original: arg min P (cid:104) P , C (cid:105) F − λH ( P ) and we ﬁnd that g ( x ) =
Revised: arg min P (cid:104) P , C (cid:105) F − λH ( P ) and we ﬁnd that g ( x ) =

Original: ¯ P ( C ) .
Revised: ¯ P ( C ) .

Original: We furthermore obtain
Revised: We furthermore obtain

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: J f ij ( C , P ) , P ab ( C , ¯ P ( C )) is hence a diagonal matrix and invertible since λ/ ¯ P ij > 0 .
Revised: J f ij ( C , P ) , P ab ( C , ¯ P ( C )) is hence a diagonal matrix and invertible since λ/ ¯ P ij > 0 .

Original: We can thus use the implicit function theorem and obtain
Revised: We can thus use the implicit function theorem and obtain

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: To extend this result to LCN-OT we use
Revised: To extend this result to LCN-OT we use

Original: [Equation]
Revised: [Equation]

Original: and the chain rule to obtain
Revised: and the chain rule to obtain

Original: [Equation]
Revised: [Equation]

Original: ∂ ¯ P sp ∂ log K sp and ∂ ¯ P spNys ∂ log K sp Nys directly follow from ∂ ¯ P ∂ C .
Revised: ∂ ¯ P sp ∂ log K sp and ∂ ¯ P spNys ∂ log K sp Nys directly follow from ∂ ¯ P ∂ C .

Original: We can calculate the pseudoinverses ¯ P † U
Revised: We can calculate the pseudoinverses ¯ P † U

Original: =
Revised: =

Original: ( ¯ P TU ¯ P U ) − 1 ¯ P U and ¯ P † W = ¯ P TW ( ¯ P W ¯ P TW ) − 1 in time
Revised: ( ¯ P TU ¯ P U ) − 1 ¯ P U and ¯ P † W = ¯ P TW ( ¯ P W ¯ P TW ) − 1 in time

Original: O (( n + m ) l 2 + l 3 ) since ¯ P U ∈ R n × l and ¯ P W ∈ R l × m .
Revised: O (( n + m ) l 2 + l 3 ) since ¯ P U ∈ R n × l and ¯ P W ∈ R l × m .

Original: We do not fully instantiate the matrices required for backpropagation but instead save their decompositions, similar to the transport plan ¯
Revised: We do not fully instantiate the matrices required for backpropagation but instead save their decompositions, similar to the transport plan ¯

Original: P Nys = ¯ P
Revised: P Nys = ¯ P

Original: U ¯ P W .
Revised: U ¯ P W .

Original: We can then compute backpropagation in time O (( n + m ) l 2 ) by applying the sums over i and j in the right order.
Revised: We can then compute backpropagation in time O (( n + m ) l 2 ) by applying the sums over i and j in the right order.

Original: We thus obtain O (( n + m ) l 2 + l 3 ) overall runtime for backpropagation.
Revised: We thus obtain O (( n + m ) l 2 + l 3 ) overall runtime for backpropagation.

Original: H C HOOSING LSH N EIGHBORS AND N YSTRÖM L ANDMARKS
Revised: H C HOOSING LSH N EIGHBORS AND N YSTRÖM L ANDMARKS

Original: We focus on two LSH methods for obtaining near neighbors.
Revised: We focus on two LSH methods for obtaining near neighbors.

Original: Cross-polytope LSH (Andoni et al., 2015) uses a random projection matrix
Revised: Cross-polytope LSH (Andoni et al., 2015) uses a random projection matrix

Original: R ∈ R d × b/ 2 with the number of hash buckets b , and then decides on the hash bucket via h ( x ) = arg max([ x T R (cid:107) − x T R ]) , where (cid:107) denotes concatenation.
Revised: R ∈ R d × b/ 2 with the number of hash buckets b , and then decides on the hash bucket via h ( x ) = arg max([ x T R (cid:107) − x T R ]) , where (cid:107) denotes concatenation.

Original: K-means LSH computes k-means and uses the clusters as hash buckets.
Revised: K-means LSH computes k-means and uses the clusters as hash buckets.

Original: We further improve the sampling probabilities of cross-polytope LSH via the AND-OR construction.
Revised: We further improve the sampling probabilities of cross-polytope LSH via the AND-OR construction.

Original: In this scheme we calculate B · r hash functions, divided into B sets (hash bands) of r hash functions
Revised: In this scheme we calculate B · r hash functions, divided into B sets (hash bands) of r hash functions

Original: [Table]
Revised: [Table]

Original: each.
Revised: each.

Original: A pair of points is considered as neighbors if any hash band matches completely.
Revised: A pair of points is considered as neighbors if any hash band matches completely.

Original: K-means
Revised: K-means

Original: LSH does not work well with the AND-OR construction since its samples are highly correlated.
Revised: LSH does not work well with the AND-OR construction since its samples are highly correlated.

Original: For large datasets we use hierarchical k-means instead (Paulevé et al., 2010; Nistér & Stewénius, 2006)
Revised: For large datasets we use hierarchical k-means instead (Paulevé et al., 2010; Nistér & Stewénius, 2006)

Original: [Table]
Revised: [Table]

Original: Since the graph transport network (GTN) uses the L 2 distance between embeddings as a cost function we use (hierarchical) k-means LSH and k-means Nyström in both sparse OT and LCNOT.
Revised: Since the graph transport network (GTN) uses the L 2 distance between embeddings as a cost function we use (hierarchical) k-means LSH and k-means Nyström in both sparse OT and LCNOT.

Original: For embedding alignment we use cross-polytope LSH for sparse OT since similarities are measured via the dot product.
Revised: For embedding alignment we use cross-polytope LSH for sparse OT since similarities are measured via the dot product.

Original: For LCN-OT we found that using k-means LSH works better with
Revised: For LCN-OT we found that using k-means LSH works better with

Original: Nyström using k-means++ sampling than cross-polytope LSH.
Revised: Nyström using k-means++ sampling than cross-polytope LSH.

Original: This is most likely due to a better alignment between LSH samples and Nyström.
Revised: This is most likely due to a better alignment between LSH samples and Nyström.

Original: We convert the cosine similarity to a distance via d cos = (cid:114) 1 − x T p x q (cid:107) x p (cid:107) 2 (cid:107) x q (cid:107) 2 (Berg et al., 1984) to use k-means with dot product similarity.
Revised: We convert the cosine similarity to a distance via d cos = (cid:114) 1 − x T p x q (cid:107) x p (cid:107) 2 (cid:107) x q (cid:107) 2 (Berg et al., 1984) to use k-means with dot product similarity.

Original: Note that this is actually based on cosine similarity, not the dot product.
Revised: Note that this is actually based on cosine similarity, not the dot product.

Original: Due to the balanced nature of OT we found this more sensible than maximum inner product search (MIPS).
Revised: Due to the balanced nature of OT we found this more sensible than maximum inner product search (MIPS).

Original: For both experiments we also experimented with uniform and recursive RLS sampling but found that the above mentioned methods work better.
Revised: For both experiments we also experimented with uniform and recursive RLS sampling but found that the above mentioned methods work better.

Original: I I MPLEMENTATIONAL D ETAILS
Revised: I I MPLEMENTATIONAL D ETAILS

Original: Our implementation runs in batches on a GPU via PyTorch (Paszke et al., 2019) and PyTorch Scatter (Fey & Lenssen, 2019).
Revised: Our implementation runs in batches on a GPU via PyTorch (Paszke et al., 2019) and PyTorch Scatter (Fey & Lenssen, 2019).

Original: To avoid over- and underﬂows we use log-stabilization throughout, i.e. we save all values in log-space and compute all matrix-vector products and additions via the log-sum-exp trick log (cid:80)
Revised: To avoid over- and underﬂows we use log-stabilization throughout, i.e. we save all values in log-space and compute all matrix-vector products and additions via the log-sum-exp trick log (cid:80)

Original: i e x i = max j x j + log( (cid:80) i e x i − max j x j ) .
Revised: i e x i = max j x j + log( (cid:80) i e x i − max j x j ) .

Original: Since the matrix A is small we compute its inverse using double precision to improve stability.
Revised: Since the matrix A is small we compute its inverse using double precision to improve stability.

Original: Surprisingly, we did not observe any beneﬁt from using the Cholesky decomposition or not calculating A − 1 and instead solving the equation
Revised: Surprisingly, we did not observe any beneﬁt from using the Cholesky decomposition or not calculating A − 1 and instead solving the equation

Original: B = AX for X .
Revised: B = AX for X .

Original: We furthermore precompute W =
Revised: We furthermore precompute W =

Original: A − 1 V to avoid unnecessary operations.
Revised: A − 1 V to avoid unnecessary operations.

Original: We use 3 layers and an embedding size H N = 32 for GTN.
Revised: We use 3 layers and an embedding size H N = 32 for GTN.

Original: The MLPs use a single hidden layer, biases and LeakyReLU non-linearities.
Revised: The MLPs use a single hidden layer, biases and LeakyReLU non-linearities.

Original: The single-head MLP uses an output size of H N, match =
Revised: The single-head MLP uses an output size of H N, match =

Original: H N and a hidden embedding size of 4 H N , i.e. the same as the concatenated node embedding, and the multi-head MLP uses a hidden embedding size of H N .
Revised: H N and a hidden embedding size of 4 H N , i.e. the same as the concatenated node embedding, and the multi-head MLP uses a hidden embedding size of H N .

Original: To stabilize initial training we scale the node embeddings by ¯ d ¯ n √ H N,match directly before calculating OT.
Revised: To stabilize initial training we scale the node embeddings by ¯ d ¯ n √ H N,match directly before calculating OT.

Original: ¯ d denotes the average graph distance in the training set, ¯ n the average number of nodes per graph, and H N, match the matching embedding size, i.e.
Revised: ¯ d denotes the average graph distance in the training set, ¯ n the average number of nodes per graph, and H N, match the matching embedding size, i.e.

Original: 32 for single
Revised: 32 for single

Original: head
Revised: head

Original: and 128 for
Revised: and 128 for

Original: OT.
Revised: OT.

Original: J G RAPH D ATASET G ENERATION AND E XPERIMENTAL D ETAILS
Revised: J G RAPH D ATASET G ENERATION AND E XPERIMENTAL D ETAILS

Original: The dataset statistics are summarized in Table 5.
Revised: The dataset statistics are summarized in Table 5.

Original: Each dataset contains the distances between all graph pairs in each split, i.e. 10 296 and 1128 distances for preferential attachment.
Revised: Each dataset contains the distances between all graph pairs in each split, i.e. 10 296 and 1128 distances for preferential attachment.

Original: The AIDS dataset was generated by randomly sampling graphs with at most 30 nodes from the original AIDS dataset (Riesen & Bunke, 2008).
Revised: The AIDS dataset was generated by randomly sampling graphs with at most 30 nodes from the original AIDS dataset (Riesen & Bunke, 2008).

Original: Since not all node types are present in the training set and our choice of
Revised: Since not all node types are present in the training set and our choice of

Original: GED is permutation-invariant w.r.t.
Revised: GED is permutation-invariant w.r.t.

Original: types, we permuted the node types so that there are no previously unseen types in the validation and test sets.
Revised: types, we permuted the node types so that there are no previously unseen types in the validation and test sets.

Original: For the preferential attachment datasets we ﬁrst generated 12, 4, and 4 undirected “seed” graphs (for train, val, and test) via the initial attractiveness model with randomly chosen parameters: 1 to 5 initial nodes, initial attractiveness of 0 to 4 and 1 / 2¯ n and 3 / 2¯ n total nodes, where ¯ n is the average number of nodes ( 20 , 200 , 2000 , and 20 000 ).
Revised: For the preferential attachment datasets we ﬁrst generated 12, 4, and 4 undirected “seed” graphs (for train, val, and test) via the initial attractiveness model with randomly chosen parameters: 1 to 5 initial nodes, initial attractiveness of 0 to 4 and 1 / 2¯ n and 3 / 2¯ n total nodes, where ¯ n is the average number of nodes ( 20 , 200 , 2000 , and 20 000 ).

Original: We then randomly label every node (and edge) in these graphs uniformly.
Revised: We then randomly label every node (and edge) in these graphs uniformly.

Original: To obtain the remaining graphs we edit the “seed” graphs between ¯ n/ 40 and ¯ n/ 20 times by randomly adding, type editing, or removing nodes
Revised: To obtain the remaining graphs we edit the “seed” graphs between ¯ n/ 40 and ¯ n/ 20 times by randomly adding, type editing, or removing nodes

Original: [Table] [Table]
Revised: [Table] [Table]

Original: and edges.
Revised: and edges.

Original: Editing nodes and edges is 4x and adding/deleting edges 3x as likely as adding/deleting nodes.
Revised: Editing nodes and edges is 4x and adding/deleting edges 3x as likely as adding/deleting nodes.

Original: Most of these numbers were chosen arbitrarily, aiming to achieve a somewhat reasonable dataset and process.
Revised: Most of these numbers were chosen arbitrarily, aiming to achieve a somewhat reasonable dataset and process.

Original: We found that the process of ﬁrst generating seed graphs and subsequently editing these is crucial for obtaining meaningfully structured data to learn from.
Revised: We found that the process of ﬁrst generating seed graphs and subsequently editing these is crucial for obtaining meaningfully structured data to learn from.

Original: For the GED we choose an edit cost of 1 for changing a node or edge type and 2 for adding or deleting a node or an edge.
Revised: For the GED we choose an edit cost of 1 for changing a node or edge type and 2 for adding or deleting a node or an edge.

Original: We represent node and edge types as one-hot vectors.
Revised: We represent node and edge types as one-hot vectors.

Original: We train all models except SiamMPNN (which uses SGD) and GTN on Linux with the Adam optimizer and mean squared error (MSE) loss for up to 300 epochs and reduce the learning rate by a factor of 10 every 100 steps.
Revised: We train all models except SiamMPNN (which uses SGD) and GTN on Linux with the Adam optimizer and mean squared error (MSE) loss for up to 300 epochs and reduce the learning rate by a factor of 10 every 100 steps.

Original: On Linux we train for up to 1000 epochs and reduce the learning rate by a factor of 2 every 100 steps.
Revised: On Linux we train for up to 1000 epochs and reduce the learning rate by a factor of 2 every 100 steps.

Original: We use the parameters from the best epoch based on the validation set.
Revised: We use the parameters from the best epoch based on the validation set.

Original: We choose hyperparameters for all models using multiple steps of grid search on the validation set, see Tables 6 to 8 for the ﬁnal values.
Revised: We choose hyperparameters for all models using multiple steps of grid search on the validation set, see Tables 6 to 8 for the ﬁnal values.

Original: We use the originally published result of SimGNN on Linux and thus don’t provide its hyperparameters.
Revised: We use the originally published result of SimGNN on Linux and thus don’t provide its hyperparameters.

Original: GTN uses 500 Sinkhorn iterations.
Revised: GTN uses 500 Sinkhorn iterations.

Original: We obtain the ﬁnal entropy regularization parameter from λ base via λ
Revised: We obtain the ﬁnal entropy regularization parameter from λ base via λ

Original: = λ base ¯ d ¯ n 1 log n , where ¯ d denotes the average graph distance and ¯ n the average number of nodes per graph in the training set.
Revised: = λ base ¯ d ¯ n 1 log n , where ¯ d denotes the average graph distance and ¯ n the average number of nodes per graph in the training set.

Original: The factor ¯ d/ ¯ n serves to estimate the embedding distance scale and 1 / log n counteracts the entropy scaling with n log n .
Revised: The factor ¯ d/ ¯ n serves to estimate the embedding distance scale and 1 / log n counteracts the entropy scaling with n log n .

Original: Note that the entropy regularization parameter was small, but always far from 0, which shows that entropy regularization actually has a positive
Revised: Note that the entropy regularization parameter was small, but always far from 0, which shows that entropy regularization actually has a positive

Original: [Table]
Revised: [Table]

Original: Full Sinkhorn scales quadratically, while all approximationes scale at most linearly with the size.
Revised: Full Sinkhorn scales quadratically, while all approximationes scale at most linearly with the size.

Original: Sparse approximations are 2-4x faster than low-rank approximations, and factored OT is multiple times slower due to its iterative reﬁnement scheme.
Revised: Sparse approximations are 2-4x faster than low-rank approximations, and factored OT is multiple times slower due to its iterative reﬁnement scheme.

Original: Note that similarity matrix computation time ( K ) primarily depends on the LSH/Nyström method, not the OT approximation.
Revised: Note that similarity matrix computation time ( K ) primarily depends on the LSH/Nyström method, not the OT approximation.

Original: [Table]
Revised: [Table]

Original: OT
Revised: OT

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table]
Revised: [Table]

Original: Sparse Skh.
Revised: Sparse Skh.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

Original: sustaining accuracy.
Revised: sustaining accuracy.

Original: effect on learning.
Revised: effect on learning.

Original: On the pref.
Revised: On the pref.

Original: att.
Revised: att.

Original: 200 dataset we use no L 2 regularization, λ base = 0 . 5 , and a batch size of 200.
Revised: 200 dataset we use no L 2 regularization, λ base = 0 . 5 , and a batch size of 200.

Original: For pref.
Revised: For pref.

Original: att.
Revised: att.

Original: 2k we use λ base = 2 and a batch size of 20 for full Sinkhorn and 100 for
Revised: 2k we use λ base = 2 and a batch size of 20 for full Sinkhorn and 100 for

Original: LCN-OT.
Revised: LCN-OT.

Original: For pref.
Revised: For pref.

Original: att.
Revised: att.

Original: 20k we use λ base = 50 and a batch size of 4.
Revised: 20k we use λ base = 50 and a batch size of 4.

Original: λ base scales with graph size due to normalization of the PM kernel.
Revised: λ base scales with graph size due to normalization of the PM kernel.

Original: For LCN-OT we use roughly 10 neighbors for LSH (20 k-means clusters) and 10 k-means landmarks for Nyström on pref.
Revised: For LCN-OT we use roughly 10 neighbors for LSH (20 k-means clusters) and 10 k-means landmarks for Nyström on pref.

Original: att.
Revised: att.

Original: We double these numbers for pure Nyström Sinkhorn, sparse OT, and multiscale OT.
Revised: We double these numbers for pure Nyström Sinkhorn, sparse OT, and multiscale OT.

Original: For pref.
Revised: For pref.

Original: att.
Revised: att.

Original: 2k we use around 15 neighbors ( 10 · 20 hierarchical clusters) andlandmarks and for pref.
Revised: 2k we use around 15 neighbors ( 10 · 20 hierarchical clusters) andlandmarks and for pref.

Original: att.
Revised: att.

Original: 20k we use roughly 30 neighbors ( 10 · 10 · 10 hierarchical clusters) and 20 landmarks.
Revised: 20k we use roughly 30 neighbors ( 10 · 10 · 10 hierarchical clusters) and 20 landmarks.

Original: The number of neighbors for the 20k dataset is higher and strongly varies per iteration due to the unbalanced nature of hierarchical k-means.
Revised: The number of neighbors for the 20k dataset is higher and strongly varies per iteration due to the unbalanced nature of hierarchical k-means.

Original: This increase in neighbors and landmarks and
Revised: This increase in neighbors and landmarks and

Original: PyTorch’s missing support for ragged tensors largely explains LCN-OT’s deviation from perfectly linear runtime scaling.
Revised: PyTorch’s missing support for ragged tensors largely explains LCN-OT’s deviation from perfectly linear runtime scaling.

Original: We perform all runtime measurements on a compute node using one Nvidia GeForce GTX 1080 Ti, two Intel Xeon E5-2630 v4, and 256GB RAM.
Revised: We perform all runtime measurements on a compute node using one Nvidia GeForce GTX 1080 Ti, two Intel Xeon E5-2630 v4, and 256GB RAM.

Original: K R UNTIMES
Revised: K R UNTIMES

Original: Table 9 compares the runtime of the full Sinkhorn distance with different approximation methods using 40 neighbors/landmarks.
Revised: Table 9 compares the runtime of the full Sinkhorn distance with different approximation methods using 40 neighbors/landmarks.

Original: We separate the computation of approximate K from the optimal transport computation (Sinkhorn iterations), since the former primarily depends on the LSH and
Revised: We separate the computation of approximate K from the optimal transport computation (Sinkhorn iterations), since the former primarily depends on the LSH and

Original: Nyström methods we choose.
Revised: Nyström methods we choose.

Original: We observe a 2-4x speed difference between sparse (multiscale OT and sparse Sinkhorn) and low-rank approximations (Nyström Sinkhorn and LCN-Sinkhorn), while factored OT is multiple times slower due to its iterative reﬁnement scheme.
Revised: We observe a 2-4x speed difference between sparse (multiscale OT and sparse Sinkhorn) and low-rank approximations (Nyström Sinkhorn and LCN-Sinkhorn), while factored OT is multiple times slower due to its iterative reﬁnement scheme.

Original: In Fig.
Revised: In Fig.

Original: 5 we observe that this runtime gap stays constant independent of the number of neighbors/landmarks, i.e. the relative
Revised: 5 we observe that this runtime gap stays constant independent of the number of neighbors/landmarks, i.e. the relative

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: Sparse Sinkhorn mostly performs best, with LCNSinkhorn coming in second.
Revised: Sparse Sinkhorn mostly performs best, with LCNSinkhorn coming in second.

Original: Factored OT performs well with very few landmarks
Revised: Factored OT performs well with very few landmarks

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: 10λ
Revised: 10λ

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

Original: R e l .
Revised: R e l .

Original: e rr .
Revised: e rr .

Original: d λ c difference decreases as we increase the number of neighbors/landmarks.
Revised: d λ c difference decreases as we increase the number of neighbors/landmarks.

Original: This gap could either be due to details in low-level CUDA implementations and hardware or the fact that low-rank approximations require 2x as many multiplications for the same number of neighbors/landmarks.
Revised: This gap could either be due to details in low-level CUDA implementations and hardware or the fact that low-rank approximations require 2x as many multiplications for the same number of neighbors/landmarks.

Original: In either case, both
Revised: In either case, both

Original: Table 9 and Fig.
Revised: Table 9 and Fig.

Original: 5 show that the runtimes of all approximations scale linearly both in the dataset size and the number of neighbors and landmarks, while full Sinkhorn scales quadratically.
Revised: 5 show that the runtimes of all approximations scale linearly both in the dataset size and the number of neighbors and landmarks, while full Sinkhorn scales quadratically.

Original: We furthermore investigate whether GTN with approximate Sinkhorn indeed scales log-linearly with the graph size by generating preferential attachment graphs with 200 , 2000 , and 20 000 nodes ( ± 50 % ).
Revised: We furthermore investigate whether GTN with approximate Sinkhorn indeed scales log-linearly with the graph size by generating preferential attachment graphs with 200 , 2000 , and 20 000 nodes ( ± 50 % ).

Original: We use the Pyramid matching (PM) kernel (Nikolentzos et al., 2017) as prediction target.
Revised: We use the Pyramid matching (PM) kernel (Nikolentzos et al., 2017) as prediction target.

Original: Fig.
Revised: Fig.

Original: 6 shows that the runtime of LCN-Sinkhorn scales almost linearly (dashed line) and regular full
Revised: 6 shows that the runtime of LCN-Sinkhorn scales almost linearly (dashed line) and regular full

Original: Sinkhorn quadraticly (dash-dotted line) with the number of nodes, despite both achieving similar accuracy and LCN using slightly more neighbors and landmarks on larger graphs to sustain good accuracy.
Revised: Sinkhorn quadraticly (dash-dotted line) with the number of nodes, despite both achieving similar accuracy and LCN using slightly more neighbors and landmarks on larger graphs to sustain good accuracy.

Original: Full Sinkhorn went out of memory for the largest graphs.
Revised: Full Sinkhorn went out of memory for the largest graphs.

Original: L D ISTANCE A PPROXIMATION
Revised: L D ISTANCE A PPROXIMATION

Original: Figs.
Revised: Figs.

Original: 7 and 8 show that for the chosen λ = 0 .
Revised: 7 and 8 show that for the chosen λ = 0 .

Original: sparse Sinkhorn offers the best trade-off between computational budget and distance approximation, with LCN-Sinkhorn and multiscale OT coming in second.
Revised: sparse Sinkhorn offers the best trade-off between computational budget and distance approximation, with LCN-Sinkhorn and multiscale OT coming in second.

Original: Factored OT is again multiple times slower than the other methods and thus not included in
Revised: Factored OT is again multiple times slower than the other methods and thus not included in

Original: Fig.
Revised: Fig.

Original: Note that d λc can be negative due to the entropy offset.
Revised: Note that d λc can be negative due to the entropy offset.

Original: This picture changes as we increase the regularization.
Revised: This picture changes as we increase the regularization.

Original: For higher regularizations LCN-Sinkhorn is the most precise at constant computational budget (number of neighbors/landmarks), as shown in Fig. 9.
Revised: For higher regularizations LCN-Sinkhorn is the most precise at constant computational budget (number of neighbors/landmarks), as shown in Fig. 9.

Original: Note that the crossover points in this ﬁgure roughly coincide with those in Fig.
Revised: Note that the crossover points in this ﬁgure roughly coincide with those in Fig.

Original: Keep in mind that in most cases the OT plan is more important than the raw distance approximation, since it determines the training gradient and tasks like embedding alignment don’t use the distance at all.
Revised: Keep in mind that in most cases the OT plan is more important than the raw distance approximation, since it determines the training gradient and tasks like embedding alignment don’t use the distance at all.

Original: This becomes evident in the fact that sparse
Revised: This becomes evident in the fact that sparse

Original: Sinkhorn achieves a better distance approximation than LCN-Sinkhorn but performs worse in both downstream tasks investigated in Sec. 7.
Revised: Sinkhorn achieves a better distance approximation than LCN-Sinkhorn but performs worse in both downstream tasks investigated in Sec. 7.