Linear-Time Gromov Wasserstein Distances using Low Rank Couplings and Costs

The ability to compare and align related datasets living in heterogeneous spaces1 plays an increasingly important role in machine learning . The Gromov-Wasserstein2 ( GW ) formalism can help tackle this problem . Its main goal is to seek an assign-3 ment ( more generally a coupling matrix ) that can register points across otherwise4 incomparable datasets . As a non-convex and quadratic generalization of optimal5 transport ( OT ) , GW is NP-hard . Yet , heuristics are known to work reasonably6 well in practice , the state of the art approach being to solve a sequence of nested7 regularized OT problems . While popular , that heuristic remains too costly to scale,8 with cubic complexity in the number of samples n. We show in this paper how a9 recent variant of the Sinkhorn algorithm can substantially speed up the resolution of10 GW . That variant restricts the set of admissible couplings to those admitting a low11 rank factorization as the product of two sub-couplings . By updating alternatively12 each sub-coupling , our algorithm computes a stationary point of the problem in13 quadratic time with respect to the number of samples . When cost matrices have14 themselves low rank , our algorithm has time complexity O ( n ) . We demonstrate15 the efficiency of our method on simulated and real data.16 1 Introduction17 The ever increasing interest for Gromov-Wasserstein ... Several problems in machine learning18 involve comparing families of points that live in heterogeneous spaces . This situation arises typically19 when realigning two distinct sets of feature representations obtained from the similar source . Recent20 applications to single-cell genomics [ 15 ] and NLP [ 12 , 1 ] provide two cases in point : Thousands21 of cells taken from the same tissue are split in two groups , each group is processed with a different22 experimental protocol , resulting in two distinct sets of heterogeneous feature vectors ; Thousands of23 word embeddings for two languages are learned independently . In both cases , one expects to find24 a meaningful way to register points across sets living in heteregeneous spaces , since they contain25 similar overall information . That realignment is usually carried out using the Gromov-Wasserstein26 ( GW ) machinery proposed by Mémoli [ 26 ] and Sturm [ 36 ] , which seeks a relaxed assignment matrix27 that is as “ close ” to an isometry as possible , using a quadratic score to quantify that closeness . GW28 has a lot of practical appeal : It has been used in supervised learning [ 41 ] , generative modeling [ 7 ] ,29 domain adaptation [ 9 ] , structured prediction [ 37 ] , quantum chemistry [ 27 ] and alignment layers [ 17 ] .30 ... despite its cubic cost . Because it is an NP-hard problem , these applications rely on approximating31 GW , typically by solving a sequence of OT problems using entropic regularization . This heuristic is32 efficient yet costly , since it requires O ( n3 ) operations to register two sets of n samples , a price that is33 paid when re-instantiating each OT problem . Our goal is to reduce substantially that complexity by34 exploiting low-factorization of both parameters ( data ) and variable ( relaxed assignment ) matrices in35 the GW problem , while maintaining state of the art performance in applications.36 Submitted to 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Do not distribute . Wasserstein : from cubic to linear complexity . A comparatively simpler problem is the registration37 of two populations embedded in the same space . This corresponds to the classic optimal transport38 ( OT ) problem , which has received considerable attention in ML [ 28 ] . OT has found applications39 in computer vision [ 29 ] , NLP [ 24 ] , single cell tracking [ 33 ] or multi-task regression in neuro-40 imaging [ 22 ] . While the OT problem is originally cast as a linear program , with a O ( n3 log ( n ) ) cost,41 many of these works rely on solving instead a penalized OT problem using Sinkhorn ’ s algorithm [ 34,42 13 ] . In its most naive implementation , the Sinkhorn has quadratic complexity [ 2 ] . Recent works43 achieve O ( n ) complexity by targeting the matrix-vector updates in Sinkhorn ’ s algorithm using44 low-rank approximations of the data kernel matrix [ 4 , 3 , 31 ] . This idea can be further improved by45 imposing the low-rank constraint on the optimization variables of the original OT problem [ 19 ] , to46 modify Sinkorn ’ s steps by enforcing a low rank factorization of the coupling variable [ 32 ] .47 Gromov-Wasserstein : from NP-hard to linear approximations . The GW problem replaces48 the linear objective function in OT by a non-convex quadratic objective . Much like OT is a re-49 laxation of the optimal assignment problem , GW can be seen as a relaxation of the quadratic50 assignment problem ( QAP ) . Both GW and QAP are NP-hard to solve [ 8 ] . In practice , iteratively51 minimizing a linearization of that quadratic objective using Sinkhorn works surprisingly well [ 20 , 35 ] .52 This method corresponds to a mirror-descent scheme [ 27 ] ,53 and in the special case of Euclidean distance matrices , the54 loss is concave and it can be also interpreted as a bi-linear55 relaxation [ 23 ] . In the most general case , this results in an56 O ( n4 ) algorithm ( the objective is a quadratic function of a57 n⇥n relaxed assignment matrix ) , that is reduced to O ( n3 ) 58 when using separable losses [ 27 ] , a price that remains too59 high for several ML applications . It is possible to replace60 the GW distance by cheaper yet only distantly related prox-61 ies , such as lower bounds based on OT [ 26 ] ( see also [ 30 ] ) 62 or sliced projections [ 38 ] . Whether GW can be efficiently63 sped up remains an open question . We propose in this64 work a novel approach that leverages , as done recently for65 OT , low-rank methods . A very recent line of works attacks66 this problem by quantizing first the two input spaces to67 solve a GW problem of reduced size , thus effectively pro-68 ducing an ad-hoc low-rank coupling [ 11 ] . A nice feature69 of this approach is that it maintains the triangular inequal-70 ity and provides a valid upper-bound on the GW distance.71 Related approaches which also approximate GW distance72 using clustering methods ( possibly in a recursive way ) 73 are [ 6 ] and [ 40 ] . We take in this paper a direct approach:74 instead of separating clustering and GW resolution in 275 independent steps , we propose do address them simulta-76 neously : our method seeks the least-costly ( in GW sense ) 77 let us also denote A : = ( dqX ( xi , xi0 ) ) 1i , i0n 2 Rn⇥n and B : = ( d q Y ( xj , xj0 ) ) 1i , i0m 2 Rm⇥m92 two pairwise cost matrices between the points in the respective supports of µ and ⌫ . The Gromov-93 Wasserstein ( GW ) discrepancy between two discrete metric measure spaces ( µ , dX ) and ( ⌫ , dY ) is94 the solution of the following non-convex quadratic problem , instantiated here for simplicity as a95 function of ( a , A ) and ( b , B ) , which contain all the information that is needed:96 GW ( ( a , A ) , ( b , B ) ) = min P2⇧a , b EA , B ( P ) , where ⇧a , b : = { P 2 Rn⇥m+ |P1m = a , PT1n = b } , ( 1 ) and the energy EA , B is a quadratic function parameterized by a loss L : R⇥ R ! R:97 EA , B ( P ) : = X i , j , i0 , j0 L ( Ai , i0 , Bj , j0 ) Pi , jPi0 , j0 . ( 2 ) A typical choice of the loss is the Lp distance L ( a , b ) = |a b|p with p 1 . In that case , [ 26 ] 98 proves that GW1/p defines a distance on the space of metric measure spaces quotiented by measure-99 preserving isometries . When p = 2 , as we consider from now on , the GW objective can be evaluated100 efficiently using the marginal constraints imposed on P , as follows [ 27 ] :101 EA , B ( P ) = hA 2a , ai+ hB 2b , bi 2hAPB , P i . ( 3 ) Indeed , ( 3 ) can be computed efficiently in O ( n2m + nm2 ) operations , using only matrix/matrix102 multiplications , instead of the O ( n2m2 ) complexity of the naive evaluation of ( 2 ) .103 Entropic Gromov-Wasserstein . The original GW problem ( 1 ) can be regularized using an entropic104 term [ 20 , 35 , 27 ] , leading to the following problem:105 GW '' ( ( a , A ) , ( b , B ) ) = min P2⇧a , b EA , B ( P ) `` H ( P ) , ( 4 ) where H ( P ) : = P i , j Pi , j ( log ( Pi , j ) 1 ) is the entropy of P . By applying a Mirror106 Descent ( MD ) scheme with respect to the KL divergence and by choosing the step-size to107 be = 1/ '' , Peyré et al . [ 27 ] provide a simple algorithm which consists in solving a se-108 quence of regularized OT problem as presented in Algorithm 1 . Indeed , each KL pro-109 jection in Algorithm 1 can be computed efficiently thanks to the Sinkhorn algorithm [ 13 ] .110 Algorithm 1 Entropic-GW Inputs : A , B , a , b , `` P = abT nm for ` = 0 , . . . do C 4APB nm ( n+m ) K '' exp ( C/ '' ) nm P argmin P2⇧ ( a , b ) KL ( P , K '' ) O ( nm ) end Result : EA , B ( P ) nm ( n+m ) 111 Computational complexity . Given a cost matrix C , the112 KL projection of K '' onto the polytope ⇧ ( a , b ) , where113 KL ( P , Q ) = hP , log ( P/Q ) 1i , is carried out in the inner114 loop of Algo . 1 using the Sinkhorn algorithm , through115 matrix-vector products . This quadratic complexity ( in116 red ) is dominated by the cost of updating matrix C at each117 iteration in Algorithm 1 , which requires O ( n2m+ nm2 ) 118 algebraic operations ( cubic , in violet ) . As noted above,119 evaluating the objective EA , B ( P ) has the same order . In120 the following we show that by considering a low rank121 exact decomposition ( or approximation ) of the distance122 matrices , the cubic cost of reupdating C and subsequently123 evaluating EA , B can be brought down to quadratic.124 3 Exploiting a Low-Rank Factorization for Cost Matrices125 Exact factorization of cost matrices . In this section we consider the case where the cost matrices126 A and B admit a low-rank factorization . More precisely , we make the following assumption.127 Assumption 1 . Assume that A and B admit a low-rank factorization , that is there exists A1 , A2 2128 Rn⇥d and B1 , B2 2 Rm⇥d 0 such that A = A1AT2 and B = B1BT2 , where d⌧ n , d0 ⌧ m.129 A case in point is when both A and B are squared Euclidean distance matrices , with a sample size130 that is larger than ambient dimension . This case is highly relevant , covering many applications of OT131 to ML . The d⌧ n assumption is also likely to hold for most applications , since cases where d n132 are known to pose challenges to the estimation of OT [ 16 , 39 ] . Writing X = [ x1 , . . . , xn ] 2 Rd⇥n , if133 A = ⇥ kxi xjk22 ⇤ i , j , then one has , writing z = ( X 2 ) T1d 2 Rn that A = z1Tn + 1nzT 2XTX.134 Therefore by denoting A1 = [ z,1n , p 2XT ] 2 Rn⇥ ( d+2 ) and A2 = [ 1n , z , p 2XT ] 2 Rn⇥ ( d+2 ) 135 we obtain the factorization above.136 Under Assumption 1 , the complexity of Algo . 1 is downgraded to quadratic in sample size : the two137 operations that make Algo . 1 cubic lie in the updates of the cost and the computation of the objective.138 Observe that for any given P 2 Rn⇥m , one can compute at each iteration139 C = 4A1A T 2 PB1B T 2 in nm ( d + d0 ) + dd0 ( n + m ) algebraic operations . Moreover thanks to the reformulation of140 EA , B ( P ) given in ( 3 ) , one can compute it in quadratic time as well . Indeed writing G1 : =141 AT1 PB2 and G2 : = AT2 PB1 , both in Rd⇥d 0 , one has hAPB , P i = 1Td ( G1 G2 ) 1d0 . Com-142 puting G1 , G2 given P requires only 2 ( nmd + mdd0 ) , and computing their dot product adds143 dd0 algebraic operations . The overall complexity to compute EA , B ( P ) is O ( nmd + mdd0 ) .144 Algorithm 2 Quadratic Entropic-GW Inputs : A1 , A2 , B1 , B2 , a , b , `` P = abT nm for ` = 0 , . . . do G2 AT2 PB1 nmd + mdd ’ C 4A1G2BT2 nmd ’ + ndd ’ K '' exp ( C/ '' ) nm P argmin P2⇧ ( a , b ) KL ( P , K '' ) O ( nm ) end c1 hA 2a , ai+ hB 2b , bi O ( nm ) G2 AT2 PB1 nmd + mdd ’ G1 AT1 PB2 nmd + mdd ’ c2 21Td ( G1 G2 ) 1d0 O ( dd ’ ) EA , B ( P ) c1 + c2 Return : EA , B ( P ) 145 General distance matrices . When the original146 cost matrices A , are not low-rank but describe147 distances , we propose to use a recent body of148 work that output their low-rank approximation149 in linear time [ 5 , 21 ] . These algorithms produce,150 for any distance matrix D 2 Rn⇥m and > 0,151 matrices D1 2 Rn⇥d , D2 2 Rm⇥d in O ( ( m +152 n ) poly ( d ) ) algebraic operations such that with153 probability at least 0.99 one has154 kD D1D T 2 k 2 F kD Cdk 2 F + kDk 2 F where Cd denotes the best rank-d approximation155 to D. We fall back on this approach to obtain a156 low-rank factorization of a distance matrix in lin-157 ear time whenever needed , aware that this incurs158 an additional approximation . See Appendix B159 for more details.160 4 Imposing a Low Nonnegative Low-Rank for the Coupling161 In this section , we shift our attention to a different opportunity for speed-ups , without assuming that162 Assumption 1 holds : we regularize the GW problem problem by decomposing the coupling as a163 product of two low-rank couplings , in the footsteps of [ 18 , 32 ] , using the following definition:164 Definition 1 . Given M 2 Rn⇥m , the nonnegative ( NN ) rank of M is the smallest number of165 nonnegative rank-one matrices into which the matrix can be decomposed additively:166 rk+ ( M ) : = min ( q|M = qX i=1 Ri , 8i , rk ( Ri ) = 1 , Ri 0 ) . Following [ 18 , 32 ] , we propose to constrain GW , enforcing a rank r on the coupling:167 GW-LR ( r ) ( ( a , A ) , ( b , B ) ) : = min P2⇧a , b ( r ) EA , B ( P ) , where ⇧a , b ( r ) : = { P 2 ⇧a , b , rk+ ( P ) r } . ( 5 ) Note that the minimum is always attained as ⇧a , b ( r ) is compact and the objective is continuous.168 In [ 32 ] , the authors show that one can parameterize any coupling in ⇧a , b ( r ) as a product of two169 low-rank couplings linked by a common marginal . For any g 2 ⇤r , the interior of r , writing170 ⇧a , g , b : = n P 2 Rn⇥m+ , P = Q diag ( 1/g ) RT , Q 2 ⇧a , g , and R 2 ⇧b , g o . one has that S g2 ⇤r ⇧a , g , b = ⇧a , b ( r ) . Therefore GW-LR introduced in ( 5 ) can be reformulated as171 the following optimization problem172 GW-LR ( r ) ( ( a , A ) , ( b , B ) ) = min ( Q , R , g ) 2C ( a , b , r ) EA , B ( Q diag ( 1/g ) R T ) ( 6 ) where C ( a , b , r ) : = C1 ( a , b , r ) \ C2 ( r ) , with173 C1 ( a , b , r ) : = n ( Q , R , g ) 2 Rn⇥r+ ⇥ Rm⇥r+ ⇥ ( R⇤+ ) r s.t . Q1r = a , R1r = b o , C2 ( r ) : = n ( Q , R , g ) 2 Rn⇥r+ ⇥ Rm⇥r+ ⇥ Rr+ s.t . QT1n = RT1m = g o . Stabilization of the Method . [ 32 ] propose to stabilize the objective defined in ( 6 ) by adding to the174 constraints a lower bound ↵ on the weight vector g such that g ↵ coordinate-wise . Indeed , as175 a solution of ( 6 ) must satisfies g > 0 coordinate-wise , then for ↵ sufficiently small , the solution176 of the same problem where one adds the constraint g ↵ will remain the same . Therefore let us177 introduce our new set of constraints C ( a , b , r , ↵ ) : = C1 ( a , b , r , ↵ ) \ C2 ( r ) where C1 ( a , b , r , ↵ ) : =178 C1 ( a , b , r ) \ { ( Q , R , g ) | g ↵ } . Another way to stabilize the method is by considering a double179 regularization scheme as proposed in [ 32 ] where in addition of constraining the nonnegative rank180 of the coupling , we regularize the objective by adding an entropic term in ( Q , R , g ) , which is to be181 understood as that of the values of the three respective entropies evaluated for each term.182 GW-LR ( r ) '' , ↵ ( ( a , A ) , ( b , B ) ) : = min ( Q , R , g ) 2C ( a , b , r , ↵ ) EA , B ( Q diag ( 1/g ) R T ) `` H ( ( Q , R , g ) ) . ( 7 ) Mirror Descent Scheme . As in [ 27 ] , we propose to use a MD scheme with respect to the KL183 divergence to approximate GW-LR ( r ) '' , ↵ in ( 7 ) . More precisely , for any `` 0 , the MD scheme leads184 for all k 0 to the following updates which require solving a convex barycenter problem per step:185 ( Qk+1 , Rk+1 , gk+1 ) : = argmin ⇣2C ( a , b , r , ↵ ) KL ( ⇣ , Kk ) ( 8 ) where ( Q0 , R0 , g0 ) 2 C ( a , b , r ) is an initial point such that Q0 > 0 and R0 > 0,186 Pk : = Qk diag ( 1/gk ) RTk , Kk : = ( K ( 1 ) k , K ( 2 ) k , K ( 3 ) k ) , K ( 1 ) k : = exp ( 4 APkBRk diag ( 1/gk ) 187 ( `` 1 ) log ( Qk ) ) , K ( 2 ) k : = exp ( 4 BP T k DQk diag ( 1/gk ) ( `` 1 ) log ( Rk ) ) , K ( 3 ) k : =188 exp ( 4 ! k/g2k ( `` 1 ) log ( gk ) ) with [ ! k ] i : = [ Q T kAPkBRk ] i , i for all i 2 { 1 , . . . , r } and189 is a positive step size . Solving ( 8 ) can be done efficiently thanks to the Dykstra ’ s Algorithm as190 showed in [ 32 ] . See Appendix C for more details.191 Initialization . To initialize our algorithm , we adapt the First Lower Bound of [ 26 ] to our case of192 interest . More precisely , we show the following Proposition . See appendix A for the proof.193 Proposition 1 . Let us denote x̃ = A 2a 2 Rn , ỹ = B 2b 2 Rm and C̃ = ( |x̃i ỹj |2 ) i , j 2 Rn⇥m.194 Then for all `` 0 and r 1 we have,195 GW-LR ( r ) ( ( a , A ) , ( b , B ) ) min ( Q , R , g ) 2C ( a , b , r , ↵ ) hC̃ , Q diag ( 1/g ) RT i `` H ( ( Q , R , g ) ) . ( 9 ) Note that the RHS of the inequality ( 9 ) is exactly the problem studied in [ 32 ] for which an algorithm196 was proposed . Therefore to initialize our algorithm , we propose to use their approach . Note that here197 the cost C̃ is the squared Euclidean distance between two families { x̃1 , . . . , x̃n } and { ỹ1 , . . . , ỹm } 198 in 1-D which admits a low-rank factorization . Therefore we can apply the linear-time version of the199 algorithm presented in [ 32 ] to compute the solution . Algorithm 3 summarizes our approach , where200 D ( · ) denotes the operator extracting the diagonal of a square matrix.201 Computational Cost . Computing the initialization goes through the computations of x̃ and ỹ which202 requires O ( n2 + m2 ) algebraic operations . Moreover , applying the algorithm proposed in [ 32 ] 203 when the underlying cost is the squared Euclidean distances between two families in 1-D needs204 only O ( ( n +m ) r ) algebraic operations . Solving the barycenter problem as defined in ( 8 ) can be205 done efficiently thanks to Dykstra ’ s Algorithm . Indeed in [ 32 , Algorithm 2 ] the authors show that206 given ( K ( 1 ) k , K ( 2 ) k , K ( 3 ) k ) , each iteration of their algorithm requires only O ( ( n + m ) r ) algebraic207 operations since it involves only matrix/vector multiplications . However computing the kernel208 matrices ( K ( 1 ) k , K ( 2 ) k , K ( 3 ) k ) at each iteration of Algorithm 3 requires a quadratic complexity with209 respect to the number of samples . Overall the proposed algorithm , while faster than the cubic210 implementation proposed in [ 27 ] , still needs O ( ( n2+m2 ) r ) operations per iteration . In the following211 we will see that by combining both nonnegative low-rank constraints on the coupling and low-rank212 approximations of the distance matrices , we can obtain a linear time algorithm with respect to the213 number of samples which computes an approximation of the GW distance.214 Algorithm 3 Low-Rank GW , GW-LR ( r ) '' , ↵ ( ( a , A ) , ( b , B ) ) Inputs : A , B , a , b , r , `` , ↵ x̃ A 2a , ỹ B 2b O ( m2 + n2 ) Step ( ? ) z1 x̃ 2 , z2 ỹ 2 O ( m+ n ) C̃1 [ z1,1n , p 2x̃ ] , C̃2 [ 1m , z2 , p 2ỹ ] T O ( n+ m ) ( Q , R , g ) argmin ( Q , R , g ) 2C ( a , b , r , ↵ ) hC̃1C̃2 , Q diag ( 1/g ) RT i `` H ( ( Q , R , g ) ) O ( ( n+ m ) r ) for k = 1 , . . . do C1 AQ diag ( 1/g ) , C2 RTB O ( ( n2 + m2 ) r ) Step ( ? ? ) K ( 1 ) exp ( 4 C1C2R diag ( 1/g ) ( `` 1 ) log ( Q ) ) O ( ( m+ n ) r2 ) K ( 2 ) exp ( 4 CT2 C T 1 Q diag ( 1/g ) ( `` 1 ) log ( R ) ) O ( ( m+ n ) r 2 ) ! D ( QTC1C2R ) , K ( 3 ) exp ( 4 ! /g2 ( `` 1 ) log ( g ) ) O ( nr2 ) Q , R , g argmin ⇣2C ( a , b , r , ↵ ) KL ( ⇣ , ( K ( 1 ) , K ( 2 ) , K ( 3 ) ) ) O ( ( m+ n ) r ) end c1 hx̃ , ai+ hỹ , bi n+ m C1 AQ diag ( 1/g ) , C2 RTB O ( ( n2 + m2 ) r ) Step ( ? ? ) G C2R , G C1G , G QTG diag ( 1/g ) O ( ( m+ n ) r2 ) c2 2Tr ( G ) r E c1 + c2 Return : E Convergence of the mirror descent . Even if the objective ( 7 ) is not convex in ( Q , R , g ) , we obtain215 the non-asymptotic stationary convergence of the MD algorithm in this setting . For that purpose216 we consider the same convergence criterion as the one proposed in [ 32 ] to obtain non-asymptotic217 stationary convergence of the MD scheme defined as218 '' , ↵ ( ⇠ , ) : = 1 2 ( KL ( ⇠ , G '' , ↵ ( ⇠ , ) ) + KL ( G '' , ↵ ( ⇠ , ) , ⇠ ) ) where G '' , ↵ ( ⇠ , ) : = argmin⇣2C ( a , b , r , ↵ ) { hrEA , B ( ⇠ ) , ⇣i + 1 KL ( ⇣ , ⇠ ) } . For any 1/r ↵ > 0 , we219 show in the following proposition the non-asymptotic stationary convergence of the MD scheme220 applied to the problem ( 7 ) . See Appendix A for the proof.221 Proposition 2 . Let `` 0 , 1r ↵ > 0 and N 1 . By denoting L '' , ↵ : = 27 ( kAk2kBk2/↵ 4 + `` ) 222 and by considering a constant stepsize in the MD scheme ( 8 ) = 12L '' , ↵ , we obtain that223 min 1kN '' , ↵ ( ( Qk , Rk , gk ) , ) 4L '' , ↵D0 N . where D0 : = EA , B ( Q0 diag ( 1/g0RT0 ) GW-LR ( r ) ( ( a , A ) , ( b , B ) ) is the distance of the initial value224 to the optimal one.225 Recall that for ↵ sufficiently small , we have GW-LR ( r ) '' , ↵ ( ( a , A ) , ( b , B ) ) = GW-LR ( r ) `` ( ( a , A ) , ( b , B ) ) .226 Thus Proposition 2 show that our algorithm reach a stationary point of ( 7 ) . In particular , if `` = 0 , the227 proposed algorithm converges towards a stationary point of ( 5 ) .228 5 Double Low-rank Approach for Linear Time GW229 Almost all operations in Algorithm 3 are linear time , except for the three updates highlighted in230 red , involving C1 and C2 , and the computations of x̃ = A 2a and ỹ = B 2b as they still require a231 quadratic number of algebraic operations . When adding Assumption 1 from §3 to the rank constrained232 approach from §4 , we notice that the strengths of both approaches can work hand in hand , both in233 easier initial evaluations of x̃ , ỹ , but , most importantly , at each new recomputation of a factorized234 linearization of the quadratic objective:235 Linear time outer norms . Because A admits a low-rank factorization , one can obtain a low-rank factorization for A 2 . Indeed , remark that for x , y 2 Rd , hx , yi2 = Pd i , j=1 xixjyiyj . Therefore by studying the rows of matrices A1 : = [ a ( 1 ) 1 ; ... ; a ( 1 ) n ] and A2 : = [ a ( 2 ) 1 ; ... ; a ( 2 ) n ] , if one writes ( x ) : = Vect ( xxT ) 2 Rd2 where Vect ( · ) is the vectorization operation , we obtain that A 2 = Ã1Ã2 T where Ã1 = [ ( a ( 1 ) 1 ) , . . . , ( a ( 1 ) n ) ] T , Ã2 = [ ( a ( 2 ) 1 ) , . . . , ( a ( 2 ) n ) ] T . In Algorithm 3 , the line “ Step ( ? ) ” can thus be replaced by x̃ Ã1Ã2 T a and ỹ B̃1B̃2 T b Note236 that computing Ã1 given A1 requires only O ( nd2 ) operations , so that this alternate code only takes237 O ( nd2 ) +O ( m ( d0 ) 2 ) operations.238 Linear time linearization of the GW objective . The linearization step , the critical step in Algo.1239 that consists in updating C at each iteration , consumes a substantial portion of the computational240 budget of GW . Introducing the low-rank Sinkhorn approach makes this step quadratic in Algo.3 ; the241 complexity of that step is also quadratic using the low-rank assumption on costs A and B , in Algo.2.242 There is therefore an opportunity to marry both to speed-up that important step . We argue that this is243 indeed what happens , in the sense that combining the two yields indeed linear time complexities in244 sample sizes , by replacing in Algorithm 3 , the lines “ Step ( ? ? ) ” by245 C1 A1A T 2 Q diag ( 1/g ) and C2 R TB2B T 1 . Note that this speed-up would not be achieved using other approaches that output a low rank246 approximation of the transport plan [ 4 , 3 , 31 ] . The crucial obstacle to using these methods here is that247 the cost matrix C in GW is “ synthetic “ , in the sense that it is the output of a matrix product APB248 involving the very last transport P . This stands in stark contrast with the requirements in [ 4 , 3 , 31 ] 249 that the kernel matrix corresponding to K '' = e C/ '' admits favorable properties , such as being p.s.d250 or admitting an explicit ( random or not ) finite dimensional feature approximation . Since C changes251 at each iteration in Algo.1 , they are not directly applicable.252 Combining the results in §4 with those from §B results in updates for C1 and C2 that only require253 O ( nrd ) and O ( mrd0 ) operations.254 Linear time GW . Finally all the quadratic operations appearing in Algorithm ( 3 ) can be replaced255 by linear counterparts . The iterations that have not been modified had an overall complexity of256 O ( mr ( r + d0 ) + nr ( r + d ) ) at each iteration . The initialization and linearization steps can now be257 performed in linear time , with respective complexity of respectively O ( n ( r + d2 ) +m ( ( d0 ) 2 + r ) ) 258 and O ( ( nr ( r + d ) +mr ( r + d0 ) ) .259 6 Experiments260 Our goal in this section is to demonstrate that , for a far smaller computational budget , the GW-LR261 approach is competitive with the direct entropic approach on datasets that are either synthesized262 to exhibit local clusters , or directly validated on a real high-dimensional dataset as well . Because263 both approaches have different hyperparameters , our goal is to stick to a realistic evaluation that264 stresses both optimality of solutions as a function of computational effort , as well as performance265 in real life applications . We start by investigating the sensitivity of hyperparamaters `` and on our266 method . Since GW is not convex , these may interact in unexpected ways . Experiments were run on a267 personal MacBook Pro 2019 laptop . We reused code from github.com/meyerscetbon/LOT , and268 downloaded genomics data from github.com/rsinghlab/SCOT.269 Benchmarks . We consider three synthetic problems and one real world problem to evaluate time-270 accuracy trade-offs , and also compare the couplings obtained by our method and that of the entropic271 version [ 27 ] . More precisely , we compare the quadratic approach in GW-LR computed with272 algorithm ( 3 ) ( and its linear time counterpat , Lin GW-LR as presented in §5 ) , with Entropic-273 GW , the cubic implementation of [ 27 ] ( as well as its quadratic counterpart , Quad Entropic-GW274 presented in Algo . 2 ) . For GW-LR and Lin GW-LR , and in all experiments , we set the lower bound275 on entries of g to ↵ = 10 10.276 Initialization To initialize all algorithms with a common strategy , we adapted the first lower bound277 of [ 26 , Def . 6.1 ] to the entropic case . In all experiments showing time-accuracy tradeoffs , we choose278 to use number of operations to provide platform independent quantities . Accuracy is measured279 by evaluating the ground-truth energy E , B ( even in scenarios when the method uses a low rank280 approximation for A , B at optimization time ) .281 Sensitivity to and `` Here we aim at showing the282 dependence in both and `` of our proposed method.283 In Figure 2 , we compare the GW loss obtained by our284 algorithm when varying `` and on two mixtures . We285 show that when `` = 0 , the proposed method manage286 to consistently obtain small GW loss whatever is.287 By allowing `` > 0 , the algorithm is able to reach even288 smaller GW loss , however , the choice of `` depends289 highly on . Therefore in the following experiments,290 we fix `` = 0 for our method . We also show the291 dependence in and `` of our method in other settings292 and observe similar behaviors . See Appendix D.2 for293 more details.294 Remark 1 . As shown in Figure 8 in Appendix D.2,295 allowing `` > 0 may also increase the speed of con-296 vergence of the algorithm . However choosing well `` 297 for a given must be done carefully and we prefer in298 the following experiments to present the performance299 of our method in the simplest setting where `` = 0.300 Synthetic low-rank problem In this experiment301 we aim at comparing the time-accuracy tradeoff of302 the different methods when the underlying distribu-303 tions has a low-rank structure . For that purpose , we304 consider two distributions in respectively 10-D and305 15-D , where the support of each distributions is the306 concatenation of clusters of points , and where the eu-307 clidean distance between the centroids of the clusters308 is bigger than a threshold . Here we set = 10.309 Both distributions are uniform , have the same number310 of clusters and the same number of points in each cluster . Some illustrations of the simulated data311 is provided in Appendix D.3 . In Figure 3 , when the underlying cost is the ( not squared ) Euclidean312 distance , our methods manage to consistently obtain similar accuracy that the ones obtained by313 entropic methods , with very low rank r = n/500 , while being orders of magnitude faster . In Figure 4,314 we also compare the time-accuracy tradeoffs in the more favorable case where the underlying cost315 is the squared Euclidean distance and obtain similar results . We also show more experiments for316 different number of clusters in Appendix D.3 , leading to similar conclusions.317 Experiments on Single Cell Genomics Data . We reproduce the single-cell alignment experiment318 introduced in [ 14 ] . The dataset consists in single-cell multi-omics data generated by co-assays . In319 that setup , the ground truth one–to-one correspondence information between cells is known , and can320 therefore be used to benchmark GW strategies . The dataset considered is the SNAREseq [ 10 ] , with321 n = m = 1047 . We apply the exact same pre-processing steps as proposed in [ 14 ] by computing322 intra-domain distance matrices A and B with a k-NN graphs based on correlations , to compute323 shortest path distance matrices . Note that in that case , one can not obtain directly in linear time a324 low-rank factorization of A and B using [ 5 , 21 ] , since the shortest path distances need to be computed325 first . Therefore we only consider the quadratic GW-LR and the cubic Entropic-GW . In Figure 6,326 we compare the alignment performance through the “ fraction of samples closer than the true match ” 327 ( FOSCTTM ) introduced in [ 25 ] . We see that both algorithm obtain similar performance . However , in328 Figure 5 , we show that whatever the chosen , GW-LR reaches better accuracy while being order of329 magnitude faster than Entropic-GW for a very small rank r = 10.330 Conclusion . While the factorization introduced in [ 32 ] held the promise to speed up classic OT , we331 have shown in this work that it delivers an even larger impact when applied to the GW problem:332 Indeed , the combination of low-rank Sinkhorn factorization with-low rank cost matrices is the only333 one , to our knowledge , that ensures that the linearization step of the GW objective can be carried out334 with a linear complexity , throughout outer iterations . This linear complexity is comparable to that of335 the most recent OT solvers , yet still retains the appealing properties of the Entropic approach , such as336 stability and convergence to meaningful solutions.337 References338 [ 1 ] Jean Alaux , Edouard Grave , Marco Cuturi , and Armand Joulin . Unsupervised hyperalignment339 for multilingual word embeddings . arXiv preprint arXiv:1811.01124 , 2018.340 [ 2 ] Jason Altschuler , Jonathan Weed , and Philippe Rigollet . Near-linear time approximation341 algorithms for optimal transport via sinkhorn iteration . arXiv preprint arXiv:1705.09634 , 2017.342 [ 3 ] Jason Altschuler , Francis Bach , Alessandro Rudi , and Jonathan Niles-Weed . Massively scalable343 sinkhorn distances via the nyström method , 2018.344 [ 4 ] Jason Altschuler , Francis Bach , Alessandro Rudi , and Jonathan Weed . Approximating the345 quadratic transportation metric in near-linear time . arXiv preprint arXiv:1810.10046 , 2018.346 [ 5 ] Ainesh Bakshi and David P. Woodruff . Sublinear time low-rank approximation of distance347 matrices , 2018.348 [ 6 ] Andrew J Blumberg , Mathieu Carriere , Michael A Mandell , Raul Rabadan , and Soledad Villar.349 Mrec : a fast and versatile framework for aligning and matching point clouds with applications350 to single cell molecular data . arXiv preprint arXiv:2001.01666 , 2020.351 [ 7 ] Charlotte Bunne , David Alvarez-Melis , Andreas Krause , and Stefanie Jegelka . Learning352 generative models across incomparable spaces . arXiv preprint arXiv:1905.05461 , 2019.353 [ 8 ] Rainer E Burkard , Eranda Cela , Panos M Pardalos , and Leonidas S Pitsoulis . The quadratic354 assignment problem . In Handbook of combinatorial optimization , pages 1713–1809 . Springer,355 1998.356 [ 9 ] Laetitia Chapel , Mokhtar Alaya , and Gilles Gasso . Partial optimal transport with applications357 on positive-unlabeled learning . In Advances in Neural Information Processing Systems 33358 ( NeurIPS 2020 ) , 2020.359 [ 10 ] Song Chen , Blue B Lake , and Kun Zhang . High-throughput sequencing of the transcriptome360 and chromatin accessibility in the same cell . Nature biotechnology , 37 ( 12 ) :1452–1457 , 2019.361 [ 11 ] Samir Chowdhury , David Miller , and Tom Needham . Quantized gromov-wasserstein . arXiv362 preprint arXiv:2104.02013 , 2021.363 [ 12 ] Alexis Conneau , Guillaume Lample , Marc ’ Aurelio Ranzato , Ludovic Denoyer , and Hervé364 Jégou . Word translation without parallel data . arXiv preprint arXiv:1710.04087 , 2017.365 [ 13 ] Marco Cuturi . Sinkhorn distances : Lightspeed computation of optimal transport . In Advances366 in neural information processing systems , pages 2292–2300 , 2013.367 [ 14 ] Pinar Demetci , Rebecca Santorella , Björn Sandstede , William Stafford Noble , and Ritambhara368 Singh . Gromov-wasserstein optimal transport to align single-cell multi-omics data . bioRxiv,369 2020. doi : 10.1101/2020.04.28.066787.370 [ 15 ] Pinar Demetci , Rebecca Santorella , Bjorn Sandstede , William Stafford Noble , and Ritambhara371 Singh . Gromov-wasserstein optimal transport to align single-cell multi-omics data . BioRxiv,372 2020.373 [ 16 ] Richard Mansfield Dudley et al . Weak convergence of probabilities on nonseparable metric374 spaces and empirical measures on euclidean spaces . Illinois Journal of Mathematics , 10 ( 1 ) :375 109–126 , 1966.376 [ 17 ] Danielle Ezuz , Justin Solomon , Vladimir G Kim , and Mirela Ben-Chen . Gwcnn : A metric377 alignment layer for deep shape analysis . In Computer Graphics Forum , volume 36 , pages 49–57.378 Wiley Online Library , 2017.379 [ 18 ] Aden Forrow , Jan-Christian Hütter , Mor Nitzan , Philippe Rigollet , Geoffrey Schiebinger , and380 Jonathan Weed . Statistical optimal transport via factored couplings , 2018.381 [ 19 ] Aden Forrow , Jan-Christian Hütter , Mor Nitzan , Philippe Rigollet , Geoffrey Schiebinger , and382 Jonathan Weed . Statistical optimal transport via factored couplings . In The 22nd International383 Conference on Artificial Intelligence and Statistics , pages 2454–2465 . PMLR , 2019.384 [ 20 ] Steven Gold and Anand Rangarajan . Softassign versus softmax : Benchmarks in combinatorial385 optimization . Advances in neural information processing systems , pages 626–632 , 1996.386 [ 21 ] Piotr Indyk , Ali Vakilian , Tal Wagner , and David Woodruff . Sample-optimal low-rank approxi-387 mation of distance matrices , 2019.388 [ 22 ] Hicham Janati , Thomas Bazeille , Bertrand Thirion , Marco Cuturi , and Alexandre Gramfort.389 Multi-subject meg/eeg source imaging with sparse multi-task regression . NeuroImage , page390 116847 , 2020.391 [ 23 ] Hiroshi Konno . Maximization of a convex quadratic function under linear constraints . Mathe-392 matical programming , 11 ( 1 ) :117–127 , 1976.393 [ 24 ] Matt Kusner , Yu Sun , Nicholas Kolkin , and Kilian Q Weinberger . From word embeddings to394 document distances . In Proc . of the 32nd Intern . Conf . on Machine Learning , pages 957–966,395 2015.396 [ 25 ] Jie Liu , Yuanhao Huang , Ritambhara Singh , Jean-Philippe Vert , and William Stafford Noble.397 Jointly embedding multiple single-cell omics measurements . BioRxiv , page 644310 , 2019.398 [ 26 ] Facundo Mémoli . Gromov–wasserstein distances and the metric approach to object matching.399 Foundations of computational mathematics , 11 ( 4 ) :417–487 , 2011.400 [ 27 ] Gabriel Peyré , Marco Cuturi , and Justin Solomon . Gromov-wasserstein averaging of kernel and401 distance matrices . In International Conference on Machine Learning , pages 2664–2672 , 2016.402 [ 28 ] Gabriel Peyré and Marco Cuturi . Computational optimal transport . Foundations and Trends in403 Machine Learning , 11 ( 5-6 ) , 2019 . ISSN 1935-8245.404 [ 29 ] Yossi Rubner , Carlo Tomasi , and Leonidas J. Guibas . The earth mover ’ s distance as a metric for405 image retrieval . International Journal of Computer Vision , 40 ( 2 ) :99–121 , November 2000.406 [ 30 ] Ryoma Sato , Marco Cuturi , Makoto Yamada , and Hisashi Kashima . Fast and robust comparison407 of probability measures in heterogeneous spaces . arXiv preprint arXiv:2002.01615 , 2020.408 [ 31 ] Meyer Scetbon and Marco Cuturi . Linear time sinkhorn divergences using positive features,409 2020.410 [ 32 ] Meyer Scetbon , Marco Cuturi , and Gabriel Peyré . Low-rank sinkhorn factorization , 2021.411 [ 33 ] Geoffrey Schiebinger , Jian Shu , Marcin Tabaka , Brian Cleary , Vidya Subramanian , Aryeh412 Solomon , Joshua Gould , Siyan Liu , Stacie Lin , Peter Berube , et al . Optimal-transport analysis413 of single-cell gene expression identifies developmental trajectories in reprogramming . Cell , 176414 ( 4 ) :928–943 , 2019.415 [ 34 ] Richard Sinkhorn . A relationship between arbitrary positive matrices and doubly stochastic416 matrices . Ann . Math . Statist. , 35:876–879 , 1964.417 [ 35 ] Justin Solomon , Gabriel Peyré , Vladimir G Kim , and Suvrit Sra . Entropic metric alignment for418 correspondence problems . ACM Transactions on Graphics ( TOG ) , 35 ( 4 ) :1–13 , 2016.419 [ 36 ] Karl-Theodor Sturm . The space of spaces : curvature bounds and gradient flows on the space of420 metric measure spaces . arXiv preprint arXiv:1208.0434 , 2012.421 [ 37 ] Titouan Vayer , Laetita Chapel , Rémi Flamary , Romain Tavenard , and Nicolas Courty . Fused422 gromov-wasserstein distance for structured objects : theoretical foundations and mathematical423 properties . arXiv preprint arXiv:1811.02834 , 2018.424 [ 38 ] Titouan Vayer , Rémi Flamary , Romain Tavenard , Laetitia Chapel , and Nicolas Courty . Sliced425 gromov-wasserstein . arXiv preprint arXiv:1905.10124 , 2019.426 [ 39 ] Jonathan Weed and Francis Bach . Sharp asymptotic and finite-sample rates of convergence of427 empirical measures in wasserstein distance . Bernoulli , 25 ( 4A ) :2620–2648 , 2019.428 [ 40 ] Hongteng Xu , Dixin Luo , and Lawrence Carin . Scalable gromov-wasserstein learning for graph429 partitioning and matching . arXiv preprint arXiv:1905.07645 , 2019.430 [ 41 ] Hongteng Xu , Dixin Luo , Hongyuan Zha , and Lawrence Carin Duke . Gromov-wasserstein431 learning for graph matching and node embedding . In International conference on machine432 learning , pages 6932–6941 . PMLR , 2019.433 Checklist434 1 . For all authors ... 435 ( a ) Do the main claims made in the abstract and introduction accurately reflect the paper ’ s436 contributions and scope ? [ Yes ] The main claim of the paper is a linear time ( w.r.t.437 sample size , as commonly understood in OT ) computation for GW . Sections §3,4 and 5438 build up that answer . This is experimentally validated across several experiments , both439 synthetic , to help the reader form intuitions , and on real data where GW was deemed440 useful.441 ( b ) Did you describe the limitations of your work ? [ Yes ] Because the GW problem is non-442 convex , these limitations are naturally discussed in the experimental section , Section §6.443 We discuss the effects of and `` on the method , which are not easy to parse due to the444 non-convexity of the method.445 ( c ) Did you discuss any potential negative societal impacts of your work ? [ N/A ] As a purely446 methodological paper , we do not envision potentially negative impact of this work on447 its own.448 ( d ) Have you read the ethics review guidelines and ensured that your paper conforms to449 them ? [ Yes ] We have read these guidelines and confirm our paper does conform to450 them.451 2 . If you are including theoretical results ... 452 ( a ) Did you state the full set of assumptions of all theoretical results ? [ Yes ] Our theoretical453 results are of two nature : a bound in Proposition 1 , and a guarantee on convergence in454 Proposition 2 , which has no direct consequence on our empirical findings , yet remain455 useful . Theory is not our main contribution , but rather algorithms.456 ( b ) Did you include complete proofs of all theoretical results ? [ Yes ] all proofs are in the457 appendix . Space in the main body of the paper was prioritized to include experimental458 validation , which , for this non-convex problem , we believe to be equally important.459 3 . If you ran experiments ... 460 ( a ) Did you include the code , data , and instructions needed to reproduce the main ex-461 perimental results ( either in the supplemental material or as a URL ) ? [ Yes ] We have462 included portions of the code that we have used . We pledge to make the entire code463 available later in the reviewing process.464 ( b ) Did you specify all the training details ( e.g. , data splits , hyperparameters , how they were465 chosen ) ? [ Yes ] The main contribution is compute efficiency , that we have considered466 across several parameter choices.467 ( c ) Did you report error bars ( e.g. , with respect to the random seed after running ex-468 periments multiple times ) ? [ N/A ] Our experiments are deterministic , since we use a469 predefined initialization.470 ( d ) Did you include the total amount of compute and the type of resources used ( e.g. , type471 of GPUs , internal cluster , or cloud provider ) ? [ Yes ] The experiments were run with472 basic computational means , a macbook pro.473 4 . If you are using existing assets ( e.g. , code , data , models ) or curating/releasing new assets ... 474 ( a ) If your work uses existing assets , did you cite the creators ? [ Yes ] We have reused a475 single toolbox , and accessed data available publicly.476 ( b ) Did you mention the license of the assets ? [ Yes ] All licenses are open source , see477 supplementary.478 ( c ) Did you include any new assets either in the supplemental material or as a URL ? [ Yes ] 479 Yes , code is shared in the supplementary.480 ( d ) Did you discuss whether and how consent was obtained from people whose data you ’ re481 using/curating ? [ Yes ] 482 ( e ) Did you discuss whether the data you are using/curating contains personally identifiable483 information or offensive content ? [ N/A ] 484 5 . If you used crowdsourcing or conducted research with human subjects ... 485 ( a ) Did you include the full text of instructions given to participants and screenshots , if486 applicable ? [ N/A ] 487 ( b ) Did you describe any potential participant risks , with links to Institutional Review488 Board ( IRB ) approvals , if applicable ? [ N/A ] 489 ( c ) Did you include the estimated hourly wage paid to participants and the total amount490 spent on participant compensation ? [ N/A ] 491

Summary: In this paper, the authors proposed a new variant of GW distance with linear computational complexity. The proposed LinearGW-LR leverages the low-rank property of the distance matrices and the low-rank assumption imposed on the target optimal transport matrix, which combines the advantage of the quadratic entropic GW and that of the low-rank Sinkhorn factorization in [32]. In the experiments, the authors show the feasibility of the proposed method in various tasks, especially in the cases of low-dimensional data.