Title: Unbalanced Low-rank Optimal Transport Solvers
Venue: NeurIPS 2023
Accepted: yes
Keywords: Optimal Transport

Unbalanced
Abstract: The relevance of optimal transport methods to machine learning has long been hindered by two salient limitations.
First, the $O(n^3)$ computational cost of standard sample-based solvers (when used on batches of $n$ samples) is prohibitive.
Second, the mass conservation constraint makes OT solvers too rigid in practice: because they must match \textit{all} points from both measures, their output can be heavily influenced by outliers.
A flurry of recent works in OT has addressed these computational and modelling limitations, but has resulted in two separate strains of methods:
While the computational outlook was much improved by entropic regularization, more recent $O(n)$ linear-time \textit{low-rank} solvers hold the promise to scale up OT further.
On the other hand, modelling rigidities have been eased owing to unbalanced variants of OT, that rely on penalization terms to promote, rather than impose, mass conservation.
The goal of this paper is to merge these two strains, to achieve the promise of \textit{both} versatile/scalable unbalanced/low-rank OT solvers. 
We propose custom algorithms to implement these extensions for the linear OT problem and its Fused-Gromov-Wasserstein generalization, and demonstrate their practical relevance to challenging spatial transcriptomics matching problems.

Unbalanced Low-rank Optimal Transport Solvers

Abstract

Two salient limitations have long hindered the relevance of optimal transport methods to machine learning. First, the O ( n 3 ) computational cost of standard sample-based solvers (when used on batches of n samples) is prohibitive. Second, the mass conservation constraint makes OT solvers too rigid in practice: because they must match all points from both measures, their output can be heavily influenced by outliers. A flurry of recent works has addressed these computational and modeling limitations, but has resulted in two separate strains of methods: While the computational outlook was much improved by entropic regularization, more recent O ( n ) linear-time low-rank solvers hold the promise to scale up OT further. In terms of modeling flexibility, the rigidity of mass conservation has been eased for entropic regularized OT, thanks to unbalanced variants of OT that can penalize couplings whose marginals deviate from those specified by the source and target distributions. The goal of this paper is to merge these two strains, low-rank and unbalanced, to achieve the promise of solvers that are both scalable and versatile. We propose custom algorithms to implement these extensions for the linear OT problem and its fused-Gromov-Wasserstein generalization, and demonstrate their practical relevance to challenging spatial transcriptomics matching problems. These algorithms are implemented in the ott-jax toolbox [Cuturi et al., 2022].

1 Introduction

Recent machine learning (ML) works have witnessed a flurry of activity around optimal transport (OT) methods. The OT toolbox provides convenient, intuitive and versatile ways to quantify the difference between two probability measures, either to quantify a distance (the Wasserstein and Gromov-Wasserstein distances), or, in more elaborate scenarios, by computing a push-forward map that can transform one measure into the other [Peyré and Cuturi, 2019]. Recent examples include, e.g., single-cell omics [Bunne et al., 2021, 2022, Demetci et al., 2020, Nitzan et al., 2019, Cang et al., 2023, Klein et al., 2023], attention mechanisms [Tay et al., 2020, Sander et al., 2022], self-supervised learning[Caron et al., 2020, Oquab et al., 2023], and learning on graphs [Vincent-Cuaz et al., 2023].

On the challenges of using OT. Despite their long history in ML [Rubner et al., 2000], OT methods have long suffered from various limitations, that arise from their statistical, computational, and modelling aspects. The statistical argument is commonly referred to as the curse-of-dimensionality of OT estimators: the Wasserstein distance between two probability densities, and its associated optimal Monge map, is poorly approximated using samples as the dimension d of observation grows [Dudley et al., 1966, Boissard and Le Gouic, 2014]. On the computational side, computing OT between a pair of n samples involves solving a (generalized) matching problem, with a price of O ( n 3 ) and above [Kuhn, 1955, Ahuja et al., 1993]. Finally, the original model for OT rests on a

Michael Klein ∗

Marco Cuturi

mass conservation constraint: all observations from either samples must be accounted for, including outliers that are prevalent in machine learning datasets. Combined, these weaknesses have long hindered the use of OT, until a more recent generation of solvers addressed these three crucial issues.

The Entropic Success Story. The winning approach, so far, to carry out that agenda has been entropic regularization methods [Cuturi, 2013]. The computational virtues of the Sinkhorn algorithm when solving OT [Altschuler et al., 2017, Peyré et al., 2016, Solomon et al., 2016, Le et al., 2021] come with statistical efficiency [Genevay et al., 2019, Mena and Niles-Weed, 2019, Chizat et al., 2020], and can also be seamlessly combined with unbalanced formulations by penalizing - rather than constraint - mass conservation, both for the linear [Frogner et al., 2015, Chizat et al., 2018, Séjourné et al., 2022, Fatras et al., 2021, Pham et al., 2020] and quadratic [Séjourné et al., 2021] problems. These developments have all been implemented in popular OT packages [Feydy et al., 2019, Flamary et al., 2021, Cuturi et al., 2022].

The Low-Rank Alternative. A recent strain of solvers relies instead on low-rank (LR) properties of cost and coupling matrices [Forrow et al., 2018, Scetbon and Cuturi, 2020, Scetbon et al., 2021]. Much like entropic solvers, these LR solvers have a better statistical outlook [Scetbon and Cuturi, 2022] and extend to GW problems [Scetbon et al., 2022]. In stark contrast to entropic solvers, however, LR solvers benefit from linear complexity O ( nrd ) w.r.t sample size n (using rank r and cost dimension d ) that can scale to ambitious tasks where entropic solvers fail [Klein et al., 2023].

The Need for Unbalanced Low-Rank Solvers. LR solvers do suffer, however, from a major practical limitation: their inability to handle unbalanced problems. Yet, unbalancedness is a crucial ingredient for OT to be practically relevant. This is exemplified by the fact that unbalancedness played a crucial role in the seminal reference [Schiebinger et al., 2019], where it is used to model cell birth and death.

Our Contributions We propose in this work to lift this last limitation for LR solvers to:

Incorporate unbalanced regularizers to define a LR linear solver (§ 3.1);

Provide accelerated algorithms, inspired by some of the recent corrections proposed by [Séjourné et al., 2022], to isolate translation terms that appear in dual subroutines (§ 3.2);

Carry over and adapt these approaches to the GW (§ 3.3) and Fused-GW problems (§ 3.4);

Carry out an exhaustive hyperparameter selection procedure within large scale OT tasks (spatial transcriptomics, brain imaging), and demonstrate the benefits of our approach (§ 4).

2 Reminders on Low-Rank Transport and Unbalanced Transport

We consider two metric spaces ( X , d X ) and ( Y , d Y ) , as well as a cost function c : X × Y → [0 , + ∞ [ . The simplex ∆ + n holds all positive n -vectors summing to 1 . For n, m ≥ 1 , a ∈ ∆ + n , and b ∈ ∆ + m , given points x 1 , . . . , x n ∈ X and y 1 , . . . , y m ∈ Y , we define two discrete probability measures µ and ν as µ := ∑ n i =1 a i δ x i , ν := ∑ m j =1 b j δ y j where δ z is the Dirac mass at z .

Cost matrices. For q ≥ 1 , consider first two square pairwise cost matrices, each encoding the geometries of points within µ and ν , and a rectangular matrix that studies that across their support:

The Kantorovich Formulation of OT is defined as the following parameterized linear program:

The Low-Rank Formulation of OT is best understood as a variant of (1) that rests on a low-rank property for cost matrix C , and low-rank constraints for couplings P . More precisely, Scetbon et al. [2021] propose to constraint the set of admissible couplings to those, within Π a,b , that have a non-negative rank of r ≥ 1 . That set can be equivalently reparameterized as

The low-rank optimal transport (LOT) problem simply uses that restriction in (1) to define :

Scetbon et al. [2021] propose and prove the convergence of a mirror-descent scheme to solve (2), and obtain linear time and memory complexities with respect to the number of samples, where each iteration in that descent scales as ( n + m ) rd , where d is the rank of C .

The Unbalanced Formulation of OT starts from (1) as well, but proposes to do without Π a,b and its marginal constraints [Frogner et al., 2015, Chizat et al., 2018], and rely instead on two regularizers:

where τ 1 , τ 2 > 0 and KL ( p | q ) := ∑ i p i log( p i /q i ) + q i -p i . This formulation is solved using entropic regularization, with modified Sinkhorn updates [Frogner et al., 2015]. Proposing an efficient algorithm able to merge (2) with (3) is the first goal of this paper.

Gromov-Wasserstein (GW) Considerations. The GW problem [Mémoli, 2011] is a generalization of (1) where the energy Q A,B is a quadratic function of P defined through inner cost matrices A , B :

where /circledot is the Hadamard product. To minimize (4), the default approach rests on entropic regularization [Solomon et al., 2016, Peyré et al., 2016] and variants [Sato et al., 2020, Blumberg et al., 2020, Xu et al., 2019, Li et al., 2023]. Scetbon et al. [2022] adapted the low-rank framework to minimize Q A,B over low-rank matrices P , achieving a linear-time complexity when A and B are themselves low-rank. Independently, [Séjourné et al., 2021] proposed an unbalanced generalization that also applies to GW and which can be implemented practically using entropic regularization. Finally, the minimization of a composite objective involving the sum of Q A,B with 〈 C, · 〉 is known as the fused GWproblem [Vayer et al., 2018].

3 Unbalanced Low-Rank Transport

3.1 Unbalanced Low-rank Linear Optimal Transport

We incorporate unbalancedness to low-rank solvers [Scetbon et al., 2021, 2022], moving gradually from the linear problem to the more involved GW and FGW problem. Using the framework of [Frogner et al., 2015, Chizat et al., 2018], we extend first the definition of LOT, introduced in (2), to the unbalanced case by considering the following optimization problem:

where rk + ( P ) denotes the non-negative rank of P . Therefore by denoting Π r := { ( Q,R,g ) ∈ R n × r + × R m × r + × R r + : Q T 1 n = R T 1 m = g } , and using the reparameterization of low-rank couplings, we obtain the following equivalent formulation of ULOT:

We introduce the more compact notation G a,b ( Q,R,g ) := F τ 1 ,a ( Q 1 r ) + F τ 2 ,b ( R 1 r ) , where F τ ,z ( s ) := τ KL ( s | z ) for τ > 0 and z ≥ 0 coordinate-wise. To solve (6), and using this split, we move away from mirror-descent and apply instead proximal gradient-descent for the KL divergence. At each iteration, we consider a linear approximation of L C where a KL penalization is added to the objective (as in the classical mirror descent scheme). However, we leave G a,b intact at each iteration. Borrowing notations from [Scetbon et al., 2021], we must solve at each iteration the convex optimization problem:

where ( Q 0 , R 0 , g 0 ) ∈ Π r is the initialization, and the triplet ξ k := ( ξ (1) k , ξ (2) k , ξ (3) k ) holds synthetic costs matrices that are re-computed at each iteration k :

with [ ω k ] i := [ Q T k CR k ] i,i ∈ R r , and ( γ k ) k ≥ 0 is a sequence of positive step sizes.

Reformulation using Duality. To solve (7), we apply Dykstra's algorithm [1983], whose iterations correspond to an alternating maximization on the dual formulation of (7):

Proposition 1. The convex optimization problem defined in (7) admits the following dual:

where h 1 , h 2 ∈ R r , f 1 ∈ R n , f 2 ∈ R m , F /star τ ,z ( · ) := sup y { 〈 y, · 〉 -F τ ,z ( y ) } is the convex conjugate of F τ ,z . In addition strong duality holds and the primal problem admits a unique minimizer.

Remark 1. While we stick to KL regularizers in this work for simplicity, it is worth noting that this can be extended to more generic regularizers F τ 1 ,a and F τ 2 ,b , as considered by Chizat et al. [2018].

We use an alternating maximization scheme to solve (8). Starting from h (0) 1 = h (0) 2 = 0 r , we apply for /lscript ≥ 0 the following updates (dropping iteration number k in (7) for simplicity):

These maximizations can all be obtained in closed form, to result in the closed-form updates:

When using 'scaling' representations for these dual variables, /lscript ≥ 0 , u ( /lscript ) i := exp( γ f ( /lscript ) i ) and v ( /lscript ) i := exp( γ h ( /lscript ) i ) for i ∈ { 1 , 2 } , we obtain a simple update, provided in the appendix (Alg. 5).

Initialization and Termination. We use the stopping criterion proposed in [Scetbon et al., 2021] to terminate the algorithm, ∆ ( ζ , ˜ ζ , γ ) := 1 γ 2 (KL( ζ , ˜ ζ ) + KL( ˜ ζ , ζ )) . Finding an efficient initialization for that problem is challenging, and various choices have been implemented for instance in [Cuturi et al., 2022]. We adopt the practical choices proposed in [Scetbon and Cuturi, 2022], using either random subcoupling matrices or a k -means approach, and also follow them in adapting the choice of γ k at each iteration k of the outer loop. We summarize our proposal in Algorithm 1, which can be seen as an extension of [Scetbon et al., 2021, Alg.2].

Convergence. The convergence proof for Dykstra's algorithm, as implemented in Alg. 5 (see appendix), follows from [Bauschke and Combettes, 2008]). Scetbon et al. [2021] show the convergence of their scheme towards a stationary point, w.r.t to the criterion ∆ ( · , · , γ ) , for fixed γ . The stationary convergence of our proposed algorithm can be directly derived from their result.

Complexity. Given ξ , solving Eq. (7) requires a time and memory complexity of O (( n + m ) r ) . However computing ξ requires in general O (( n 2 + m 2 ) r ) time and O ( n 2 + m 2 ) memory. Scetbon et al. [2021] propose to consider low-rank factorizations of the cost matrix C of the form C /similarequal C 1 C T 2 where C 1 ∈ R n × d and C 2 ∈ R m × d . In that case computing ξ can be done in O (( n + m ) rd ) time and O (( n + m )( r + d )) memory. Such factorizations are either known explicitly (e.g. when using squared-Euclidean distances) or can be obtained as approximations using the algorithm in [Indyk et al., 2019], which guarantees that for any distance matrix C ∈ R n × m and α > 0 it can output matrices C 1 ∈ R n × d , C 2 ∈ R m × d in O (( m + n ) poly ( d α )) algebraic operations such that with probability at least 0 . 99 , ‖ C -C 1 C T 2 ‖ 2 F ≤ ‖ C -C d ‖ 2 F + α ‖ C ‖ 2 F , where C d denotes the best rankd approximation to C .

Algorithm 1 ULOT ( C, a, b, r, γ 0 , τ 1 , τ 2 , δ )

Inputs: C, a, b, r, γ , τ , τ , δ

3.2 Improvements on the Unbalanced Dykstra Algorithm

A well documented source of instability of unbalanced formulations of OT lies in the fact that the total mass of the optimal unbalanced coupling is not known beforehand. Séjourné et al. [2022] have proposed a technique to address this issue, with the benefit of reduced computational costs. They propose first a dual objective that is translation invariant (TI). We take inspiration from this strategy and adapt it to our problem, to propose the following variant of (8):

It is clear from the reparameterization that both problems (8) and (9) have the same value and also that ( ˜ f 1 , ˜ h 1 , ˜ f 2 , ˜ h 2 ) is solution of (9) if and only if ( ˜ f 1 + λ /star 1 , ˜ h 1 -λ /star 1 , ˜ f 2 + λ /star 2 , ˜ h 2 -λ /star 2 ) is solution of (8) where ( λ /star 1 , λ /star 2 ) solves D TI ( ˜ f 1 , ˜ h 1 , ˜ f 2 , ˜ h 2 ) . To solve (9), we show that the variational formulation of the translation invariant dual objective targeted inside (9) can be obtained in closed form.

Proposition 2. Let ˜ f 1 ∈ R n , ˜ f 2 ∈ R m and ˜ h 1 , ˜ h 2 ∈ R r , then the inner problem defined in (9) by D TI ( ˜ f 1 , ˜ h 1 , ˜ f 2 , ˜ h 2 ) admits a unique solution ( λ /star 1 , λ /star 2 ) and we have that

where

Using Proposition 2, we perform an alternate maximization scheme on the TI formulation of the dual D TI. Indeed using Danskin's theorem (under the assumption that λ /star 1 , λ /star 2 do not diverge), one obtains a variant of Algorithm 5. That TI approach is summarized below in Algorithm 3, using Algorithm 2 as a subroutine. We show in the experiments section ( Exp. 1 ) that the TI approach has better computational performance on a simple task.

Algorithm 2 compute-lambdas ( a, b, ξ (3) , u 1 , v 1 , u 2 , v 2 , γ , τ 1 , τ 2 ) Inputs: a, b, ξ (3) , u 1 , v 1 , u 2 , v 2 , γ , τ 1 , τ 2 ˜ u 1 ← u -1 / γ / τ 1 1 , ˜ u 2 ← u -1 / γ / τ 2 2 c 1 ← log( 〈 ˜ u 1 , a 〉 ) -log( 〈 ξ (3) , v -1 1 /circledot v -1 2 〉 ) , c 2 ← log( 〈 ˜ u 2 , b 〉 ) -log( 〈 ξ (3) , v -1 1 /circledot v -1 2 〉 ) Result: λ /star 1 , λ /star 2 as in (10), (11)

Algorithm 3 ULR-TI-Dykstra ( a, b, ξ , γ , τ 1 , τ 2 , δ )

Inputs: a, b, ξ = ( ξ (1) , ξ (2) , ξ (3) ) , γ , τ 1 , τ 2 , δ v 1 = v 2 = 1 r , u 1 = 1 n , u 2 = 1 m repeat ˜ v 1 = v 1 , ˜ v 2 = v 2 , ˜ u 1 = u 1 , ˜ u 2 = u 2 λ 1 , λ 2 ← compute-lambdas ( a, b, ξ (3) , u 1 , v 1 , u 2 , v 2 , γ , τ 1 , τ 2 ) ( Alg. 2) u 1 = ( a ξ (1) v 1 ) τ 1 τ 1 +1 / γ exp( -λ 1 / τ 1 ) τ 1 1 / γ + τ 1 , u 2 = ( b ξ (2) v 2 ) τ 2 τ 2 +1 / γ exp( -λ 2 / τ 2 ) τ 2 1 / γ + τ 2 , λ 1 , λ 2 ← compute-lambdas ( a, b, ξ (3) , u 1 , v 1 , u 2 , v 2 , γ , τ 1 , τ 2 ) ( Alg. 2) g = exp( γ ( λ 1 + λ 2 )) 1 / 3 ( ξ (3) /circledot ( ξ (1) ) T u 1 /circledot ( ξ (2) ) T u 2 ) 1 / 3 , v 1 = g ( ξ (1) ) T u 1 , v 2 = g ( ξ (2) ) T u 2 until 1 γ max( ‖ log( u i / ˜ u i ) ‖ ∞ , ‖ log( v i / ˜ v i ) ‖ ∞ ) < δ ; Result: diag( u 1 ) ξ (1) k diag( v 1 ) , diag( u 2 ) ξ (2) k diag( v 2 ) , g

3.3 Unbalanced Low-rank Gromov-Wasserstein

The low-rank Gromov-Wasssertein (LGW) problem [Scetbon et al., 2022] between the two discrete metric measure spaces ( µ, d X ) and ( ν , d Y ) , written for compactness using ( a, A ) and ( b, B ) , reads

Building upon § 3.1 and leveraging the TI variant presented in § 3.2, we introduce the unbalanced low-rank Gromov-Wasserstein (ULGW) problem. There is, however, a significant challenge that appears when introducing unbalanced regularizers in (12): When P is constrained to be in Π a,b , the first two terms of the RHS in (12) simplify to a T A /circledot 2 a + b T B /circledot 2 b . Hence, they are constant and discarded when optimizing. In an unbalanced setting, these terms vary and must be accounted for:

To solve the problem, we apply the same scheme as proposed for ULOT, that is a proximal gradient descent where we linearize Q A,B and add a KL penalization while leaving the soft marginal constraints unchanged. Therefore the algorithm to solve ULGW is the same as that solving ULOT, however, the kernels ξ k now take into account the quadratic terms of the original LGW problem. More formally, at each iteration k of the outer loop, we propose to solve

where ( Q 0 , R 0 , g 0 ) ∈ Π r is the initialization, ( γ k ) k ≥ 0 a sequence of positive step sizes. Using notation P k = Q k diag(1 /g k ) R T k , the synthetic cost matrices ξ k := ( ξ (1) k , ξ (2) k , ξ (3) k ) are updated as:

Note that (14) is the exact same optimization problem as (7), where only ξ k has changed and therefore can be solved using Algorithm 3. Algorithm 4 summarizes our strategy to solve (13).

Algorithm 4 ULGW ( A,B,a,b,r, γ 0 , τ 1 , τ 2 , δ )

Convergence and Complexity. Similarly to linear ULOT, the unbalanced Dykstra algorithm is guaranteed to converge [Bauschke and Lewis, 2000]. Because we use Algorithm 5, we retain exactly the same complexity, both in terms of time of memory, to solve these inner problems. The slight variation in kernel ξ compared to ULOT still retains the same O (( n 2 + m 2 ) r ) time and O ( n 2 + m 2 ) memory complexities. However, as in ULOT, we can take advantage of low-rank approximations of both costs matrices A and B to reach linear complexity. Indeed, assuming A /similarequal A 1 A T 2 and B /similarequal B 1 B 2 where A 1 , A 2 ∈ R n × d X and B 1 , B 2 ∈ R m × d Y , then the total time and memory complexities become respectively O ( mr ( r + d Y ) + nr ( r + d X )) and O (( n + m )( r + d X + d Y )) . Again, when A and B are distance matrices, we use the algorithms from [Indyk et al., 2019].

3.4 Unbalanced Low-rank Fused-Gromov-Wasserstein

Wefinally focus on the increasingly impactful [Klein et al., 2023] fused-Gromov-Wasserstein problem, which merges linear and quadratic objectives [Vayer et al., 2018]:

where α ∈ [0 , 1] and ¯ α := 1 -α allows interpolating between the GW and linear OT geometries. This problem remains a GW problem, where one replaces the 4-way cost M [ i, i ′ , j, j ′ ] := ( A i,i ′ -B j,j ′ ) 2 appearing in (4) by a composite interpolated cost between the OT and GW geometries, redefined as M [ i, i ′ , j, j ′ ] = α C i,j + ¯ α ( A i,i ′ -B j,j ′ ) 2 . Our proposed unbalanced and low-rank version of the FGW problem includes | P | := ‖ P ‖ 1 the mass of P , to homogenize linear and quadratic terms,

which is expanded through the explicit factorization of P , noticing that | P | = | g | := ‖ g ‖ 1 :

Then by linearizing again H : ( Q,R,g ) → α | g |L C ( Q,R,g ) + ¯ α Q A,B ( Q,R,g ) with an added KL penalty and leaving G a,b unchanged, we obtain at each iteration, the same optimization problem as in (14) where the kernels ξ k are now defined as

These steps are summarized in Alg. 6 (see appendix). These steps result in a quadratic complexity, both in time and memory, with respect to the number of points n and m . However, these complexities become linear when square matrices A,B and rectangular C all admit a low-rank factorization.

Nrgnmeasured

Nrgnpredicted

Slc24a2measured

Slc24a2predicted

1.0

1.0

0.8

0.8

0.6

0.6

0.4

0.4

0.2

0.2

0.0

0.0

(a) Gene Nrgn : measured vs. predicted expression

(b) Gene

Slc24a2

: measured vs. predicted expression

4 Experiments

We focus first in Exp. 1 on demonstrating the empirical benefits of the TI variant of our algorithm to solve linear ULOT, as implemented in Alg. 3 vs. Alg. 5; that algorithm is subsequently used as an inner routine to solve all quadratic ULR problems. We compare in Exp. 2 unbalanced low-rank (ULR) solvers to balanced low-rank (LR) counterparts on a spatial transcriptomics task, and follow in Exp. 3 by comparing ULR solvers to entropic (E) counterparts on a smaller task, to accommodate entropic solvers' quadratic complexity. We conclude in Exp. 4 by comparing ULR solvers to [Thual et al., 2022], which can learn a sparse transport coupling, in the unbalanced FGW setting.

Datasets. We consider two realworld datasets, described in B.1, and two synthetic datasets, that are large enough to showcase our solvers. The real-world datasets consist of both a shared feature space, used to compute the costs matrices for the linear term in the OT and FGW settings, as well as geometries that are specific to each source s and target t measures, and which are used to compute the costs matrices for the quadratic term in the GW and FGW settings. In Exp. 1 , we sim-

Tissueregiongroundtruth

Tissueregionpredicted

CTX_1

CTX_2

DG

FbTrt

HPF_CA

HblI_VS

L1_HPFmo_Mngs

LSX_HY_MB_HB

OB_1

STR

TH

ply consider two isotropic Gaussians in d = 30 to evaluate the performance of the TI variant on a liner problem. We use the mouse brain STARmap spatial transcriptomics data from [Shi et al., 2022] for Exp. 2 and Exp. 3 . We use data from the Individual Brain Charting dataset [Pinho et al., 2018], to replicate the settings of [Thual et al., 2022], in Exp. 4 .

Metrics. Following Klein et al. [2023], we evaluate maps by focusing on the two following metrics: (i) pearson correlation ρ computed between the (ground truth) source s feature matrix F s ∈ R n × d , and the barycentric projection of the target t to the source scaled by the target marginals b t . Writing P as the transport matrix from source to target, this can be computed as P diag ( 1 b t ) F t ; (ii) F1 score when assessing class transfer (among 11 possible classes), computed between the original source vector of labels l s , taken in { 1 , · · · , 11 } n , and the inferred labels for the same points, predicted for each i by taking the argmax j B i,j , where B is a matrix of n × 11 row probabilities, each the barycentric projection of the target t one-hot encoded labels L t ∈ { 0 , 1 } m × 11 , B := P diag ( 1 b t ) L t .

Experiment 1: Benchmarking The Translation Invariant Variant. We evaluate the effect of the proposed TI procedure on the computational cost of ULR solvers: We compare the time taken when solving unbalanced LR problems, with or without using the TI objective. In Figure 3, we compare the execution time (using our ott-jax implementation, and a single NVIDIA GeForce RTX 2080

Ti card) of unbalanced LR Sinkhorn on large and high dimensional Gaussian distributions. The results presented are averaged over 10 random seeds with error bars. We use a δ = 10 -9 convergence threshold and 1000 maximal number of iterations for Dykstra, in 64-bit precision. We observe that the use of our proposed TI objective is consistently beneficial when solving ULR problems. See also Appendix B.3 for additional experiments.

Experiment 2: ULOT vs. LOT on Gene Expression / Cell Type Annotation. We evaluate the accuracy of ULOT solvers for a large-scale spatial transcriptomics task, using gene expression mapping and cell type annotation. We compare it to the balanced LR alternative using the Pearson correlation ρ as described in the metrics section. We leverage two coronal sections of the mouse brain profiled by STARmap spatial transcriptomics by

[Shi et al., 2022]. They consist of n ≈ 40 , 000 cells in both the source and target brain section. Each cell is described by 1000 gene features, in addition to 2D spatial coordinates. As a result A,B are ≈ 40 k × 40 k , and the fused term C is a squared-Euclidean distance matrix on 30D PCA space computed on the gene expression space. We selected 10 marker genes for the validation and test sets from the HPF_CA cluster. We run an extensive grid search as reported in B.2, we pick the best hyperparameters combination using performance on the 10 validation genes as a criterion, and we report that metric on the other genes in Table 1, as well as qualitative results in Figure 1 and Figure 2. Clearly, ULFGW is the best performing solver across all metrics. Interestingly, the ULOT does not consistently outperforms its balanced version, and unbalancedness seems to hurt performance for the LGW solvers. Nevertheless, both solvers display inconsistent performance across metrics, whereas the ULFGW and LFGW are consistently superior to the rest of the solvers. These results highlight how the flexibility given by the FGW formulation to leverage common and disparate geometries, paired with the unbalancedness relaxation, can provide state of the art algorithms for matching problems in large-scale, real world biological problems.

Experiment 3: ULOT vs. UEOT. We compare the performance of ULOT solvers to their unbalanced entropic alternatives (UEOT). We use the same datasets as in Exp. 2 , but must pick a smaller subset (Olfactory bulb), to avoid OOM errors for entropic UGW solvers, since they cannot handle the 40 k sizes considered in Exp. 2 (see B.1). This results in n ≈ 20 , 000 source and ≈ 15 , 000 target

source

ground-truth

ulfgw (ours), p=0.8180

FUGW,p=0.6324

15

24

9.3

14

7.7

12

4.6

- 7.1

0

0

0

0

-7.7

-12

-4.6

-7.1

-15

-24

-9.3

-14

cells, and 1000 genes. Similar to Exp. 2 , the fused term C is a squared-Euclidean distance matrix on 30-D PCA space, computed on gene expressions. As done in Exp. 2 , we select 10 marker genes for the validation and 10 genes for the test set, from cluster OB_1 . We run an extensive grid search, as in Exp. 2 and B.2. Table 2 shows that ULFGW outperforms entropic solvers w.r.t. Pearson correlation ρ , but is worse when considering F1 scores. On the other hand, ULFGW confirms its superiority compared to the balanced alternative LFGW. Taken together, these results suggest that while unbalanced

LR solvers are on par with unbalanced entropic solvers in terms of performance, in small data regimes, they remain much faster and can unlock the applications of unbalanced OT to larger scales.

Experiment 4: ULOT to align brain meshes. In this experiment, we compare the performance of our ULFGW solver to FUGW-sparse [Thual et al., 2022], an alternative approach to solve unbalanced FGW problems, using a two-scale (corse/fine grid) approach to handle large sample sizes. This method was demonstrated to be effective in aligning brain anatomies, encompassing both mesh structures and functional signals associated with each vertex. For their empirical analysis, they use the individual brain charting dataset [Pinho et al., 2018].

In the absence of other information in the original paper, we draw inspiration from Pinho et al.'s smaller scale notebook implementations: We embed the n ≈ 160 , 000 vertices of the fsaverage7 mesh, into a 30-d space, using an approximation of the geodesic distances with landmark multi-dimensional scaling [De Silva and Tenenbaum, 2004] where 2048 points were used as anchors. Each vertex has an associated functional signal that entails 22 features. For both the quadratic and linear terms, we compute the costs based on the squared Euclidean distance. The coarse grid for FUGW-sparse is built using one-

tenth of n , i.e. ≈ 16 k points. We evaluate all methods by comparing the performance of the best hyperparameter combination, based on the average correlation between the barycentric projection and ground-truth value of 5 features, across a test set of 5 contrast maps. In Table 3, we observe that ULFGW and LFGW outperform FUGW-sparse. In this setting, there is no clear evidence that the unbalanced version performs better than its balanced counterpart for low-rank methods. See also Appendix B.2 for additional experimental details and results.

Conclusion. The practical success of OT methods to natural sciences demonstrates the relevance of OT to their analysis pipelines. Practitioners must, however, often deal with the poor scalability of OT algorithms, as well as their rigid assumptions w.r.t. mass conservation. While Low-rank OT approaches hold the promise of scaling OT methods to large sizes, unbalanced formulations have proved useful to relax mass conservation for entropic OT solvers. We have proposed in this paper to merge these two strains, and demonstrated the practical relevance of these unbalanced low-rank solvers on various challenging alignment tasks.

References

Ravindra K Ahuja, Thomas L Magnanti, and James B Orlin. Network Flows: Theory, Algorithms, and Applications . Prentice hall, 1993.

Jason Altschuler, Jonathan Weed, and Philippe Rigollet. Near-linear time approximation algorithms for optimal transport via sinkhorn iteration. arXiv preprint arXiv:1705.09634 , 2017.

Iro Armeni, Ozan Sener, Amir R Zamir, Helen Jiang, Ioannis Brilakis, Martin Fischer, and Silvio Savarese. 3d semantic parsing of large-scale indoor spaces. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 1534-1543, 2016.

Heinz H Bauschke and Patrick L Combettes. A Dykstra-like algorithm for two monotone operators. Pacific Journal of Optimization , 4(3):383-391, 2008.

Heinz H Bauschke and Adrian S Lewis. Dykstras algorithm with bregman projections: A convergence proof. Optimization , 48(4):409-427, 2000.

Andrew J Blumberg, Mathieu Carriere, Michael A Mandell, Raul Rabadan, and Soledad Villar. Mrec: a fast and versatile framework for aligning and matching point clouds with applications to single cell molecular data. arXiv preprint arXiv:2001.01666 , 2020.

Emmanuel Boissard and Thibaut Le Gouic. On the mean speed of convergence of empirical and occupation measures in wasserstein distance. In Annales de l'IHP Probabilités et statistiques , volume 50, pages 539-563, 2014.

Charlotte Bunne, Stefan G. Stark, Gabriele Gut, Jacobo Sarabia del Castillo, Kjong-Van Lehmann, Lucas Pelkmans, Andreas Krause, and Gunnar Rätsch. Learning single-cell perturbation responses using neural optimal transport. bioRxiv , 2021. doi: 10.1101/2021.12.15.472775. URL https: //www.biorxiv.org/content/early/2021/12/15/2021.12.15.472775 .

Charlotte Bunne, Andreas Krause, and Marco Cuturi. Supervised Training of Conditional Monge Maps. In Advances in Neural Information Processing Systems (NeurIPS) , 2022.

Zixuan Cang, Yanxiang Zhao, Axel A Almet, Adam Stabell, Raul Ramos, Maksim V Plikus, Scott X Atwood, and Qing Nie. Screening cell-cell communication in spatial transcriptomics via collective optimal transport. Nature Methods , 20(2):218-228, 2023.

Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. Advances in neural information processing systems , 33:9912-9924, 2020.

Lénaïc Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xavier Vialard. Unbalanced optimal transport: Dynamic and kantorovich formulations. Journal of Functional Analysis , 274(11): 3090-3123, 2018.

Lenaic Chizat, Pierre Roussillon, Flavien Léger, François-Xavier Vialard, and Gabriel Peyré. Faster wasserstein distance estimation with the sinkhorn divergence. Advances in Neural Information Processing Systems , 33, 2020.

Samir Chowdhury, David Miller, and Tom Needham. Quantized gromov-wasserstein. arXiv preprint arXiv:2104.02013 , 2021.

Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in neural information processing systems , pages 2292-2300, 2013.

Marco Cuturi, Laetitia Meng-Papaxanthos, Yingtao Tian, Charlotte Bunne, Geoff Davis, and Olivier Teboul. Optimal transport tools (ott): A jax toolbox for all things wasserstein. arXiv preprint arXiv:2201.12324 , 2022.

Vin De Silva and Joshua B Tenenbaum. Sparse multidimensional scaling using landmark points. Technical report, technical report, Stanford University, 2004.

Pinar Demetci, Rebecca Santorella, Björn Sandstede, William Stafford Noble, and Ritambhara Singh. Gromov-wasserstein optimal transport to align single-cell multi-omics data. bioRxiv , 2020. doi: 10.1101/2020.04.28.066787.

Richard Mansfield Dudley et al. Weak convergence of probabilities on nonseparable metric spaces and empirical measures on euclidean spaces. Illinois Journal of Mathematics , 10(1):109-126, 1966.

Richard L Dykstra. An algorithm for restricted least squares regression. Journal of the American Statistical Association , 78(384):837-842, 1983.

Kilian Fatras, Thibault Sejourne, Rémi Flamary, and Nicolas Courty. Unbalanced minibatch optimal transport; applications to domain adaptation. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pages 3186-3197. PMLR, 18-24 Jul 2021. URL https://proceedings. mlr.press/v139/fatras21a.html .

Jean Feydy, Pierre Roussillon, Alain Trouvé, and Pietro Gori. Fast and scalable optimal transport for brain tractograms. In Medical Image Computing and Computer Assisted Intervention-MICCAI 2019: 22nd International Conference, Shenzhen, China, October 13-17, 2019, Proceedings, Part III 22 , pages 636-644. Springer, 2019.

Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, Ievgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research , 22(78):1-8, 2021. URL http://jmlr.org/papers/v22/20-451.html .

Aden Forrow, Jan-Christian Hütter, Mor Nitzan, Philippe Rigollet, Geoffrey Schiebinger, and Jonathan Weed. Statistical optimal transport via factored couplings, 2018.

Charlie Frogner, Chiyuan Zhang, Hossein Mobahi, Mauricio Araya, and Tomaso A Poggio. Learning with a Wasserstein loss. In Advances in Neural Information Processing Systems , pages 2053-2061, 2015.

Aude Genevay, Lénaic Chizat, Francis Bach, Marco Cuturi, and Gabriel Peyré. Sample Complexity of Sinkhorn Divergences. In International Conference on Artificial Intelligence and Statistics (AISTATS) , volume 22, 2019.

Piotr Indyk, Ali Vakilian, Tal Wagner, and David Woodruff. Sample-optimal low-rank approximation of distance matrices, 2019.

Dominik Klein, Giovanni Palla, Marius Lange, Michal Klein, Zoe Piran, Manuel Gander, Laetitia Meng-Papaxanthos, Michael Sterr, Aimee Bastidas-Ponce, Marta Tarquis-Medina, Heiko Lickert, Mostafa Bakhti, Mor Nitzan, Marco Cuturi, and Fabian J. Theis. Mapping cells through time and space with moscot. bioRxiv , 2023. doi: 10.1101/2023.05.11.540374. URL https://www. biorxiv.org/content/early/2023/05/11/2023.05.11.540374 .

Harold W Kuhn. The hungarian method for the assignment problem. Naval research logistics quarterly , 2(1-2):83-97, 1955.

Khang Le, Huy Nguyen, Quang M Nguyen, Tung Pham, Hung Bui, and Nhat Ho. On robust optimal transport: Computational complexity and barycenter computation. Advances in Neural Information Processing Systems , 34:21947-21959, 2021.

Jiajin Li, Jianheng Tang, Lemin Kong, Huikang Liu, Jia Li, Anthony Man-Cho So, and Jose Blanchet. A convergent single-loop algorithm for relaxation of gromov-wasserstein in graph data. In The Eleventh International Conference on Learning Representations , 2023.

Facundo Mémoli. Gromov-Wasserstein distances and the metric approach to object matching. Foundations of Computational Mathematics , 11(4):417-487, 2011.

Gonzalo Mena and Jonathan Niles-Weed. Statistical bounds for entropic optimal transport: sample complexity and the central limit theorem. Advances in Neural Information Processing Systems , 32, 2019.

Mor Nitzan, Nikos Karaiskos, Nir Friedman, and Nikolaus Rajewsky. Gene expression cartography. Nature , 576(7785), 2019.

Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Marc Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, et al. Dinov2: Learning robust visual features without supervision. arXiv preprint arXiv:2304.07193 , 2023.

Gabriel Peyré, Marco Cuturi, and Justin Solomon. Gromov-wasserstein averaging of kernel and distance matrices. In International Conference on Machine Learning , pages 2664-2672, 2016.

Gabriel Peyré and Marco Cuturi. Computational optimal transport. Foundations and Trends in Machine Learning , 11(5-6), 2019. ISSN 1935-8245.

Khiem Pham, Khang Le, Nhat Ho, Tung Pham, and Hung Bui. On unbalanced optimal transport: An analysis of sinkhorn algorithm. In International Conference on Machine Learning , pages 7673-7682. PMLR, 2020.

Ana Luísa Pinho, Alexis Amadon, Torsten Ruest, Murielle Fabre, Elvis Dohmatob, Isabelle Denghien, Chantal Ginisty, Séverine Becuwe-Desmidt, Séverine Roger, Laurence Laurier, Véronique JolyTestault, Gaëlle Médiouni-Cloarec, Christine Doublé, Bernadette Martins, Philippe Pinel, Evelyn Eger, Gael Varoquaux, Christophe Pallier, Stanislas Dehaene, Lucie Hertz-Pannier, and Bertrand Thirion. Individual Brain Charting, a high-resolution fMRI dataset for cognitive mapping. Scientific Data , 5:180105, June 2018. doi: 10.1038/sdata.2018.105. URL https://hal.science/ hal-01817528 .

R Tyrrell Rockafellar. Convex analysis . Princeton university press, 1970.

Yossi Rubner, Carlo Tomasi, and Leonidas J. Guibas. The earth mover's distance as a metric for image retrieval. International Journal of Computer Vision , 40(2):99-121, November 2000.

Michael E Sander, Pierre Ablin, Mathieu Blondel, and Gabriel Peyré. Sinkformers: Transformers with doubly stochastic attention. In International Conference on Artificial Intelligence and Statistics , pages 3515-3530. PMLR, 2022.

Ryoma Sato, Marco Cuturi, Makoto Yamada, and Hisashi Kashima. Fast and robust comparison of probability measures in heterogeneous spaces. arXiv preprint arXiv:2002.01615 , 2020.

Meyer Scetbon and Marco Cuturi. Linear time sinkhorn divergences using positive features. Advances in Neural Information Processing Systems , 33:13468-13480, 2020.

Meyer Scetbon and Marco Cuturi. Low-rank optimal transport: Approximation, statistics and debiasing. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems , volume 35, pages 6802-6814. Curran Associates, Inc., 2022. URL https://proceedings.neurips.cc/paper_files/paper/2022/ file/2d69e771d9f274f7c624198ea74f5b98-Paper-Conference.pdf .

Meyer Scetbon, Marco Cuturi, and Gabriel Peyré. Low-rank sinkhorn factorization. In International Conference on Machine Learning , pages 9344-9354. PMLR, 2021.

Meyer Scetbon, Gabriel Peyré, and Marco Cuturi. Linear-time Gromov-Wasserstein distances using low rank couplings and costs. ICML , 2022.

Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell , 176(4):928-943, 2019.

Thibault Séjourné, François-Xavier Vialard, and Gabriel Peyré. The unbalanced gromov wasserstein distance: Conic formulation and relaxation. Advances in Neural Information Processing Systems , 34:8766-8779, 2021.

Thibault Séjourné, François-Xavier Vialard, and Gabriel Peyré. Faster unbalanced optimal transport: Translation invariant sinkhorn and 1-d frank-wolfe. In International Conference on Artificial Intelligence and Statistics , pages 4995-5021. PMLR, 2022.

Hailing Shi, Yichun He, Yiming Zhou, Jiahao Huang, Brandon Wang, Zefang Tang, Peng Tan, Morgan Wu, Zuwan Lin, Jingyi Ren, Yaman Thapa, Xin Tang, Albert Liu, Jia Liu, and Xiao Wang. Spatial atlas of the mouse central nervous system at molecular resolution. bioRxiv , 2022. doi: 10.1101/2022.06.20.496914. URL https://www.biorxiv.org/content/early/2022/ 06/22/2022.06.20.496914 .

Justin Solomon, Gabriel Peyré, Vladimir G Kim, and Suvrit Sra. Entropic metric alignment for correspondence problems. ACM Transactions on Graphics (TOG) , 35(4):1-13, 2016.

Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da-Cheng Juan. Sparse Sinkhorn attention. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning , volume 119 of Proceedings of Machine Learning Research , pages 9438-9447. PMLR, 13-18 Jul 2020. URL https://proceedings.mlr.press/v119/tay20a.html .

Alexis Thual, Huy Tran, Tatiana Zemskova, Nicolas Courty, Rémi Flamary, Stanislas Dehaene, and Bertrand Thirion. Aligning individual brains with fused unbalanced Gromov-Wasserstein. arXiv , 2022.

Titouan Vayer, Laetita Chapel, Rémi Flamary, Romain Tavenard, and Nicolas Courty. Fused gromovwasserstein distance for structured objects: theoretical foundations and mathematical properties. arXiv preprint arXiv:1811.02834 , 2018.

Cédric Vincent-Cuaz, Rémi Flamary, Marco Corneli, Titouan Vayer, and Nicolas Courty. Semirelaxed gromov-wasserstein divergence and applications on graphs. In International Conference on Learning Representations , 2023.

F Alexander Wolf, Philipp Angerer, and Fabian J Theis. SCANPY: large-scale single-cell gene expression data analysis. Genome Biology , 19(1), 2018.

Hongteng Xu, Dixin Luo, and Lawrence Carin. Scalable gromov-wasserstein learning for graph partitioning and matching. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019. URL https://proceedings.neurips.cc/paper_files/paper/ 2019/file/6e62a992c676f611616097dbea8ea030-Paper.pdf .

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors focus on the problem of efficiently compute discrete Optimal Transport (OT) problems, which are known to have a cubic complexity w.r.t. the number of input samples.\nThey propose to approximate it by assuming that the transport plan is a low-rank matrix and propagate this property in the computations, such that the complexity of matrix products is reduced.\nThis idea has already been proposed to approximate Balanced OT and Gromov-Wasserstein (GW) problems.\nThe authors extend this to variants called Unbalanced OT, as well as Unbalanced GW and Fused-Unbalanced GW.\nThe authors introduce these unbalanced OT variants, they derive in each setting the associated low-rank optimization algorithms using mirror descent.\nThen they perform experiments on brain and cell biology data, where unbalanced OT was extremely powerful, to assess the performance of these algorithms in an applied setting.\n\nSTRENGTHS:\n- The contributions might look incremental at first (combining unbalanced OT with low-rank OT), but are actually very thorough.\nThey do not restrict to this one combination, but consider all recent variants of unbalanced OT (Unbalanced OT, GW and Fused GW) which have been developped in the literature in the last years.\nThey also leverage one recent optimization idea of 'translation-invariance' which improves the computational efficiency of computing Unbalanced OT problems.\nAll in all, the authors aggregate several works altogether to propose a set of interesting algorithms and tools for practitioners, especially for the biology field.\n\nWEAKNESSES:\n- The low-rank approximation is not sufficiently discussed in this paper.\nI mentioned that it is considered to accelerate the computations.\nHowever, one might use this variant because it makes sense to have a prior of low-rank transport plan in their applications, because the data has some structural property which could be leveraged.\nTo my mind, using a low-rank plan seems contradictory with the property that the optimal plan is a (full rank) permutation in some setting.\nCould the authors discuss this in detail ? Why does it make sense to have a low-rank plan ? Which kind of data is relevant with such assumption ? This is unclear to me.\n\n- Some experimental illustrations could be provided on the methodological side of their algorithms, for completeness.\nThe authors propose some optimization problems and algorithms to compute them.\nTo my mind, the authors do provide complexity per iteration of their algorithms, but the experimental aspect of measuring the time performance of their algorithms is omitted.\nIn particular, have you checked experimentally that this 'translation-invariance' variant in Section (3.2) does accelerate the computations in your low-rank setting ?\nCould you provide a plot of the time to compute OT problems (standard or Sinkhorn vs. low-rank) as a function of the number of samples ?\nAlso, the use of unbalanced OT involves two extra parameters $(\\tau_1, \\tau_2)$ which requires cross-validation. \nCould the authors give plots of the performance of their biology task as a function of $(\\tau_1, \\tau_2)$.\nThis would provide interesting insights on the interpretation of these parameters, and their impact on learning tasks.\n\nQUESTIONS:\nSee my question above.\n\nLIMITATIONS:\nThey adressed the societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors combine two variants of optimal transport (OT) known as unbalanced OT and low-rank OT. While the former relaxes the marginal constraints to ease the modelling rigidities and discard possible outliers from input measures, the latter helps reduce the computational cost of $\\mathcal{O}(n^3)$, and therefore, makes OT scalable. In addition to standard OT which quantifies the discrepancy between two measures in the same dimensional space, they also apply this combination to the scenarios when two input measures belong to distinct dimensional spaces, which are Gromov-Wasserstein and Fused-Gromov-Wasserstein. Finally, the authors choose the spatial transcriptomic matching problems to justify the practical usage of their proposed methods.\n\nSTRENGTHS:\n- Originality: This work is a novel combination of two versatile and scalable variants of optimal transport (OT), which are unbalanced OT and low-rank OT. \n\n- Quality: The derivations of algorithm proposed in this work are associated with theoretical proof. The empirical performance of the proposed method is demonstrated via the spatial transcriptomic matching problems. \n\n- Significance: the results in this paper are important to some extent as it allows practitioners to find a low-rank solvers for the problem of quantifying the discrepancy between two arbitrary (not necessarily probability) measures.\n\nWEAKNESSES:\n- Originality: Although the combination of unbalanced OT and low-rank OT is novel, key tools and techniques (e.g. reparametrization of low-rank couplings, Dykstra algorithm) used in this paper have been already introduced in [1] and [2]. Thus, I think this work is incremental to some extent. To address this concern, I suggest the authors should highlight the main challenges of solving unbalanced low-rank OT compared to its balanced counterpart more clearly.\n\n- Clarity: The presentation of this paper is not good as there are a lot of notations which are either not carefully defined or define inaccurately (see Requested Changes). Additionally, there are some important parts, namely the initialization of Algorithm 1, which should be presented in this paper rather than merely refer to another paper.\n\n**References**\n\n[1] Meyer Scetbon, Marco Cuturi, and Gabriel Peyré. Low-rank sinkhorn factorization. In International Conference on Machine Learning, pages 9344–9354. PMLR, 2021.\n\n[2] Meyer Scetbon, Gabriel Peyré, and Marco Cuturi. Linear-time Gromov-Wasserstein distances using low rank couplings and costs. ICML, 2022.\n\nQUESTIONS:\n1. In line 95, what is the definition of nonnegative rank of $P$? Please add this definition to the revision of this paper.\n\n2. In equation (2), the term $Q\\text{diag}(g)R$ seems to be incorrect. Should it be $Q\\text{diag}(1/g)R^{\\top}$?\n\n3. In equation (7), are there any differences between $\\mathrm{KL}(\\cdot,\\cdot)$ and $\\mathrm{KL}(\\cdot | \\cdot)$? If any, they should be defined explicitly in the paper.\n\n4. Are there any theoretical guarantees that the tuples $(Q_{k+1},R_{k+1},g_{k+1})$ approximate the solution of the optimization problem in equation (6)? \n\n5. In the proposed algorithms for solving ULOT, ULGW and ULFGW, which values of the rank hyperparameter $r$ should we choose? Would it be as low as possible?\n\n**Requested Changes**:\n\n1. References: In lines 38 and 39, when mentioning the usage of Sinkhorn algorithm in solving OT, the authors should cite more relevant papers, namely [1]. Similarly, in line 78, regarding solving unbalanced OT using entropic regularization, the authors should cite the paper [2].\n\n2. In the definition of cost matrix $C$ in Section 2, the index notation $1\\leq i,j\\leq n,m$ is inacurrate. It should be changed to $1\\leq i\\leq n$ and $1\\leq j\\leq m$.\n\n3. Notations: When introducing the unbalanced OT in equation (3), the authors should explain the notations $\\mathrm{KL}$, $\\tau_1$ and $\\tau_2$ rather than assume that readers implicitly understand. Analogously, the notation $A^{\\odot2}$ in equation (4) should be defined explicitly.\n\n4. Typo: in line 96, repamatrization --> reparametrization.\n\n**References**\n\n[1] Khang Le, Huy Nguyen, Quang M Nguyen, Tung Pham, Hung Bui, and Nhat Ho. On robust optimal transport: Computational complexity and barycenter computation. Advances in Neural Information Processing Systems, 2021.\n\n[2] K. Pham, K. Le, N. Ho, T. Pham, and H. Bui. On unbalanced optimal transport: An analysis of sinkhorn algorithm. In ICML, 2020.\n\nLIMITATIONS:\nThe limitations are not discussed in this paper.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors proposed unbalanced low-rank OT (ULOT) solver, which is an extension of the balanced counterpart. They show how to adapt this solver to the other unbalanced low-rank settings, namely translation-invariant and GW. The experiments compare the performances of multiple low-rank solvers, and of the Unbalanced Fused GW.\n\nSTRENGTHS:\nThe extensions from the balanced to unbalanced LOT, as well as from ULOT to ULFGW are not trivial and require some calculation effort. I also find that the writing is very instructive and all technical details are clearly presented and on point.\n\nWEAKNESSES:\nIt seems to me that\n- The content in the main paper and appendix is not adequately partitioned, namely the section 3 is somewhat too long that there is few space left for the experiment section, and some experiment details in the Appendix should be put in the main. \n- The main contribution of the paper is quite incremental.\n\nMore precisely,\n\n+ While adding translation-invariant improves the convergence of the usual unbalanced OT, it is unclear about the real improvement in the experiments. IMHO, it is more or less an add-on of the ULOT solver and should be moved to Appendix because it would dilute the main message of the paper (which is about ULOT). It is enough that Algo 4 and 5 use the ULR-Dykstra solver (Algo 6), instead of Algo 3.\n\n+ The Algo 4 is merely a special case of Algo 5, where alpha = 0, thus should be removed. While it is more instructive to start with the GW setting before moving to the fused GW one, all the details can be moved to the Appendix and section 3.3 and 3.4 should be restructured or/and merged.\n\n+ It seems that the experiment section is not sufficiently diverse because the experiments and competing methods are restricted to just a family of the low-rank based methods (except FUGW). IMHO, since the contribution on the methodogoly (i.e. section 3) is not very significant, I would love to see more comparison with other methods, like (unbalanced) mini-batch OT, or quantized GW, and maybe on more experiments (e.g. on graphs).\n\n+ I find it weird that the section 4.3 is unreasonably brief and not sufficiently discussed, while most of its details are moved to the Appendix. This makes the section 4.3 look incomplete in the main paper.\n\nQUESTIONS:\nIn Algorithm 2: typo in 1 / gamma / tau1\n\nLIMITATIONS:\nThe authors do not discuss the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn its original exact formulation, optimal transport (OT) suffers from an $\\mathcal{O}(n^3)$ cost when applied to clouds of $n$ points. In addition, the exact constraint on the marginals makes it sensitive to outliers.\nOne solution for each of these problems exist:\n- unbalanced OT [Schiebinger et al., 2019] is less sensitive to outliers than regular OT, as it relaxes the marginal constraints into a KL-penalized form\n- entropic OT allows for the use of the celebrated Sinkhorn algorithm but still requires large matrix vector multiplications. Low rank OT  [Forrow et al., 2018] is a promising direction to further improve upon entropic OT, by using a low rank cost matrix and low rank constraints of the plan.\n\nBuilding on a these two seminal papers and a line of work by Scetbon, Cuturi and coauthors, the paper combines low rank OT with unbalanced OT (Formulation in Eq 5).\nA dedicated algorithm is proposed, interpreted as proximal gradient descent in the KL geometry. The latter is improved based on the works of Séjourné et al [2022].\n\nThe approach is extended to the Gromov-Wasserstein and Fused Gromov-Wasserstein OT problems.\n\nSTRENGTHS:\nOT has found many applications in practice in the last decade; alleviating its cost extends its range of applications. The proposed formulation is sound, the treatment is relatively easy to follow and the derivations are made clear.\n\nWEAKNESSES:\n- Although the paper heavily sells low rank optimal transport, isn't the formulation in eqs 5/6 non convex? If so, doesn't it make the exact solution intractable? How to ensure convergence towards a good critical point?\n- Convergence of the algorithm\n    - Known results about Dykstra's algorithm are used to compute the solution of Equation 7 via its dual. Can the authors provide a reference for the convergence of the whole algorithm, namely the iterations defined by 7? I may have missed it, but I could not find reference of convergence towards a critical point;\n    - convergence seems to be measured in terms of difference between two successive iterates going to zero. It is well-known that this does not guarantee convergence of the iterates (e.g. take the harmonic series). To clarify, do the author have a classical convergence result?\n\nQUESTIONS:\n- In the last cost of eq(2) I believe it should be diag(1/g) not diag(g), and $R^T$ not $R$\n\n\nTypos:\n- reparamaterized, repamatrization\n- mass conversation for mass conservation\n- misuse of \\cite or \\citep vs \\citet , e.g. \"of [Frogner et al., 2015, Chizat et al., 2018]\", and several other places e.g. \"notations from [Scetbon et al., 2021]\"\n-  towards a stationary points\n- such that with probability at least 0.99 that: extra \"that\"\n- Algo 1: (check me)\n- paper uses \\star and * for $\\lambda^\\star$\n\nLIMITATIONS:\nnot applicable", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}