Unbalanced Low-rank Optimal Transport Solvers
Authors: Meyer Scetbon
Michal Klein
Giovanni Palla
marco cuturi
Venue: NeurIPS 2023 poster (2023)
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
Meyer Scetbon" Microsoft Research t-ms www.sceeoommiccoooo..cooo
Michael Klein" Apple mi michalk@apple.com
Giovanni Palla Helmholtz Center Munich
Marco Cuturi Apple cuturi@apple.com com
giovanni. polteeallloeeeeeeoooooo-------ee mwwtmcceeeeenn...mmm de
Abstract
Two salient limitations have long hindered the relevance of optimal transport methods machine learning. First, the O(n²) 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 influ- enced 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 (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 tile. We propose custom algorithms to implement these extensions for the linear OT problem and its fused Gromov-Wasseee their practical relevance to challenging spatial transcriptomics matching problems. These algorithms are implemented the ott-jas jax toolbox [Cuturi et al., 2022].
solvers that are both scalable and versa-
generalization, and demonstrate
Introduction
Recent machine learning (ML) works have witnessed 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 distance (the Wasserstein and Kromov-Wassersee 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 al., 2023], attention mechanisms [Tay al., 2020, Sander et al., 2022], self-supervised learnin[[Caron et al.. 2020, Oquab et al., 2023], and learning graphs [Vincent- et al., 2023].
On the challenges of using OT. Despite their long history in MI. [Ruhner et al., 2000), OT methods have long suffered from various limitations, that arise from their statistical, computational, and the curse-of-dimensionaliy
modelling aspects. The statistical argument is commonly referred 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 pair of samples involves solving a (generalized) matching problem, with a price of (n²)) and above [Kuhn, 1955, Ahuja et al., 1993]. Finally, the original model for OT rests on
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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 more recent generation of solvers addressed these three crucial issues.
The Entropic Success Story. The winning approach, far, carry out that agenda has been entropic regularization methods [Cuturi, 2013J. The computational virtues of the Sinkhorn algorithm when solving OT [Altschuler al., 2017, Peyré et al., 2016, Solomon al., 2016, Le et al., 2021] come with statistical efficiency [Genevay et al., 2019, Mena and Niles- 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 18, Séjouméé et al., 2022, Fatras al., 2021, Pham 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 better statistical outlook [Scetbon and Cuturi, 2022] and extend to GW problems [Scetbon al., 2022]. In stark contrast entropic solvers, however, LR solvers benefit from linear complexity ((nrd)) w.r.t sample size n (using rank r and cost dimension d) that can scale ambitious tasks where entropic solvers fail [Klein 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 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 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é al., 2022], 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', dx) and (V, dy), as well as a cost function c: X' x + [0, +oo The simplex A; holds all positive n-vectors summing to For n, m 2 1,a E A and b E A given points Z1...... E X' and //..... ym ), we define two discrete probability measures and v as μ Σ 0,82,,
Σド= bydg, where 82 the Dirae mass 2.
Cost matrices. For 9 1, consider first two square pairwise geometries of points within and v, and a rectangular matrix that studies that across their support:
matrices, each encoding the
A = (x,,,, (r B == (d5(2,, By Misis'm C := [c(z,, lision. 375m
The Kantorovich Formulation of OT defined as the following parameterized linear program:
OT(H,)) = min (C,P), where 1e.bb {P R s.t. P1, aa PT1, =6 (I)
Pell,,
The Low-Rank Formulation of OT is best understood variant of (1) that rests on a low-rank property for cost matrix C, and low-rank constraints for couplings P. More precisely, Scetbon al. [2021] propose constraint the set of admissible couplings to those, within II, that have non-negative rank of 2 1. That set can be equivalently reparameterized as
(r, = (P E R m P = Q diag(1/g) Q E Ba.g, R Ib.9. and E A#).
The low-rank optimal transport (LOT) problem simply uses that restriction in (1) define LOT, v) == min
(C,P) REll
min
(C, Q diag(1/g)
(2)
rell.,(r)

---

Scetbon et al. [2021] propose and prove the convergence of derrent 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 mird, where d is the rank of C.
The Unbalanced Formulation of OT starts from (1) as well, but proposes to do without aT, and its marginal constraints [Frogner et al., 2015, Chizat et al., 2018], and rely instead on two regularizers:
UOT(H, = min (C,P) ( + (P',(1" PER
where TI, T2 > and KL(pl4) := E. log(p:/q) /:) + 9i 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.
This formulation is solved using
Gromov-Wase Wasserstein (GW) Considerations. The GW problem [Mémoli, 2011] is generalization of (1) where the energy QA,B is a quadratic function of P defined through inner cost matrices A, B:
(P.:(P):=)) E(Air (Aw 633
P" 0 PI, 1 1 2(APB, P)
(4)
where @ is the Hadamard product. To minimize (4), the default approach rests on entropic regulariza- tion [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 QA.B over low-rank matrices P. achieving linear-time complexity when A and B are themselves low-rank. Independently, [Séjourné et al., 2021] proposed an unbalanced generalization that also applies GW and which can be implemented practically using entropic regularization. Finally, the minimization of composite objective involving the sum of QA with (C,-) is known as the fusee GW problem [Vayer et al., 2018].
3 Unbalanced Low-Rank Transport
3.1 Unbalanced Low-rank Linear Optimal Transport
We incorporate unbalancedness low-rank solvers [Scetbon et al., 2021, 2022], moving gradu- ally 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:
UL,T,((, = min P: (P))
(C, KL + T2KLL(P" 1, lb),
(5)
where rk, (P) denotes the non-negative rank of P. Therefore by denoting II, == (,,
Rexr x T x [R: 1 R_1,, 9), and using the reparameterization of low-rank couplings, we obtain the following equivalent formulation of ULOT:
(L, min (Q. Le(Q.R.9)
(C.Qdiag(1/)))) + TI (a) + (21,
(6)
(Q.
We introduce the more compact notation (u, == Fria(Q1,) Fr. R1,), where (s) == =KL(s|22 for r 0 and 2 0 coordinate-we wise.. To solve (6), and using this split, we move away from mirror-descent and apply instead proximal gradient-descent for the KL diver- gence. At each iteration, we consider a linear approximation of Le where KL penalization is added to the objective (as the classical mirror descent scheme). However, we leave Ga.b intact at each iteration. Borrowing notations from [Scetbon et al., 2021], we must solve at each iteration the convex optimization problem:
(Qk+1; Rk+119k41) = argmin 1 _KI(S,E) + (TKI.(11,(a) + (0 (LL,
CET.. 7
where (Qo, Ro, 90) II, is the initialization, and the triplet En := (1) costs matrices that are re-computed at each iteration k:
5k c
holds synthetic
Qx ① e CR ding(1/sx), _e(() R& e C Qu ding(1/su))
(3)
gb C e N. /
Sk
3

---

with [ol] == [0] CR E R,, and (7k)k20 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 The convex optimization problem defined in (7) admits the following dual:
1 - - -1,56)
sup Du(fi, f2, h2) == - -F (-11) f1.21.2..22
7/k
(8)
(2) - -F (-2)
- (E7u(f20Ma) -,,
1
- (e (e-mm(()+ba)) - 1,, (3))
1
7k
7/k
where h1, E R,, f1 E R", f2 E R", (.). == sup, of Fr.z. addition strong duality holds and the primal problem admits a unique minimizer.
(y. =u ()) is the convex conjugate
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 and Fry. as considered by Chizat et al. [2018].
We use an alternating maximization scheme to solve (8). Starting from 00 h1 0,,, we apply
for & 2 0 the following updates (dropping iteration number k in (7) for simplicity):
((+1) = arg sup D(z, h1 :22 ((t)
),) pl(t1) == sug D(1
((())))
.11 2, 122 )
(++1) mett)) == arg sup D( (e+1) (hg 112
21, 12 (e+1) x2)
31
These maximizations can all be obtained in closed form, to result in the closed-fome updates:
exp())|(t1)) (()) (xp(>1(())
exp(n122t1))
b ((2) exp(>sh)))
ge+1 = (5(3) o (z( exp((x)(x)) + (1) expl> ((t())
1/3
exp(7n
ge+1
exp(T11 ((+1))
ge+1
(g()) exp(>> (
(g(2)) exp(1s24t1)
When using "scaling" representations for these dual variables, 2 0, u (0) = exp(7 (00) and
u = exp(>1))) for i {1,2], we obtain a simple update, provided the appendix (Alg. 5).
Initialization and Termination. We use the stopping criterion proposed in [Scetbon et al., 2021] to terminate the algorithm, AC,, = + (KI(C, ?) + KI. ()). Finding an efficient initialization for that problem is challenging, and various choices have been implemented for instance in [Cuturi 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 Ti at each iteration i 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, implemented in Alg. 5 (see ap- pendix), follows from [Bauschke and Combettes, 2008I). Scetbon et al. [2021] show the convergence of their scheme towards stationary point, w.r.t to the criterion A(; 7)), for fixed y. The stationary convergence of our proposed algorithm can be directly derived from their result.
Complexity. Given E, solving Eq. (7) requires time and memory complexity of O((n + m)r). However computing & requires in general O(nn² + )r)r) time and O(n² + m²) memory. Scetbon et al. [2021] propose consider low-rank factorizations of the cost matrix C' of the form C' C2 Cf where C E IR and C2 E Rdd In that case computing can be done in O(nn + m)rd) time and O( + m)(r + d)) memory. Such factorizations are either known explicitly (e.g. when using Squared-Fuclidess
distances) or can be obtained as approximations using the algorithm in [Indyk et al., 2019], which guarantees that for any distance matrix C E =m and a > 0 it can output matrices C1 E x C2 E Rmxd x in (( + n))oly( ())) algebraic operations such that with probability at least 0.99. ||C - C, S ||C - Call? + a a|I where Ca denotes the best rank-d approximation to C.

---

Algorithm 1 ULOTTC, 70, TI, T2,
Inputs: C, a. r, Jo. 71, T2,
g Initialization proposed in [Scetbon and Cuturi, 2022]
repeat
= Q R g -= 9,
Q CR (1/g). VR - =R= = diag(L/g),
+ CR), V -w
7 7 70/ max(| V E 1
((1) + Q @
g( t R O exp( (xp(-770 Vn), 5( g00
Q,R.g + (.
(Alg.
until A(QQR, (Q. 1) <6;
Result: Q, R.g
3.2 Improvements on the Unbalanced Dykstra Algorithm
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
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 our problem, to propose the following variant of (8):
(niLi F2, sup AL.AGER
D(fi + h1, λ1, f2 X2, h2 X2)
sup J. J2.
It is clear from the reparameterization that both problems (8) and (9) have the same value and also that
(f1, hu, f2, h2) is solution of (9) if and only if (fi + X1 hr - Xi + X₂, h2 >2) is solution of (8) where (Ai Xz) solves (nifi, hi, f2, h2). To solve (9), show that the variational formulation of the translation invariant dual objective targeted inside (9) can be obtained in closed form.
Proposition 2. Let fi E R", f2 E R and hy, hi IR", then the inner problem defined (9) by Drifi, b1 f2r h2) admits unique solution (Ai, X2) and we have that
7172 (1/7 (1)11/ T2)
T1/7
T1/7
72
A = 1
C2
(10)
1/7 +1/
1/7 +
1/7
(1-
1
7172
-1
T2/7
n1/7
72
A/ =
#1/7+7)(1/++))
1/7 +
.2
11/7 ti
1/7 +
C1
(11)
where
(exp(-fi/7i),00 (exp(->(h) (1,)),,(0)) g(2))
(exp(-f2/T2) T2) a) (exp(-7(h) h2)), {(3)))
G1 = :olog
and C2 = log
Using Proposition 2. we perform an alternate maximization scheme on the TI formulation of the dual DII. Indeed using Danskin's 's theorem (under the assumption that >j Xे do not diverge), one obtains a variant of Algorithm 5. That TI approach is summarized below in Algorithm 3, using Algorithm as a subroutine. We show in the experiments section (Exp. 1) that the TI approach has better computational performance simple task.
Algorithm 2 compute-lambbaas(..be 6, (3)) u1, U1. U2, 22, 7, T1, 72)
Inputs: a, ((3), u1, U1, U2 U2 T1,
di - u, 1/1/n ūz + un 1/2/7a
< (ūg((ūn, a)) log((g(() oi + ,) 02 - log( b) log((5(3), 11
Result: >j. >g as (10), (11)
5

---

Algorithm 3 UIRRTI-Dyyssraau, 6, 7: 7,,
Inputs: b. Ç (5(1, ( _ 7. T1, T2
V1 = U2 = 1,, 01 n, U2 = 1m
repeat
01 U1, 822 = U2 د₁ U1 و2 ug
A1 A2 t compute-lambdas a, b, (() U1, u2, 02, 71, T2)
(Alg. 2)
01 = c
exp
, u2 (grio
b
exp(- - 10-22/22) 1/9412,
A1 t compute-lambdars b, (00) u1 u u2 02. 7: 71,
(Alg. 2)
9 exp(n(A) + x2)) (5(0) @ (5 un @ (((2))T222)
1i
(g
2/2
until , max(|] log(u./ (ou/ū.)ll / ( ūz lo))) <;;
Result: diag(u))53 )5( diag(vi), diag(u2){5 diag(v2), g
3.3 Unbalanced Low-rank
The low-rank
(LGGW) problem [Scetbon et al., 2022] between the two discrete metric measure spaces (Hi dx) and (v, dy), written for compactness using (a, AA) and (b, B). reads
LGW,((a, A), (b, = min CA,B(P). PeTle,s(r)
(12)
Building upon $ 3.1 and leveraging the TI variant presented in 3.2, we introduce the unbalanced low-rank Cromovv-Waseersee appears when introducing unbalanced regularizers in (12): When P is constrained to be in ll the first two terms of the RHS in (12) simplify to af A + bz Bo%b. Hence, they are constant and discarded when optimizing. In an unbalanced setting, these terms vary and must be accounted for:
(ULGW) problem. There is, however, significant challenge that
ULGW, A), = min (A Q1, Q1,) + (8 21,,R1,
EQ.
(13)
- (iag(1ia(11/))))) B, B, + KLL +
(L(,
To solve the problem, we apply the same scheme as proposed for ULOT, that is proximal gradi- descent where we linearize QA.BB and add a KL penalization while leaving the soft marginal constraints unchanged. Therefore the algorithm to solve ULGW is the same that solving ULOT,, however, the kernels & now take into account the quadratic terms of the original LGW problem. More formally, at each iteration & of the outer loop, we propose to solve
(Q RE+1, 9k-11) = argmin 1 KL((CE&) + m (L((1,)a)) + (2K
(14)
CEll. 7x
where (Qo, Ro. go) E II, the initialization, (TR)R20 a sequence of positive step sizes. Using
notation Pa Qi diag(1/9n) Rf. the synthetic cost matrices & := (5t) . {{(3) are updated as:
(2)
(1) Q+ 2 expppp12 - 2 Qx1,1,) @ Axp(-4r.APPBB diag (1/9))))
AQ, diag(1/ga) l/gx))) (3) = 9h O wxpplxxxx/////oooooxxxxxxxpa g with [wx]: = AP ER.
6 = R& @ expl - 27, 02 R4 1, @
Note that (14) is the exact same optimization problem as (7), where only &a has changed and therefore can be solved using Algorithm 3. Algorithm 4 summarizes our strategy to solve (13).

---

Algorithm 4 ULGW b, 70, TI, T2,
Inputs: A, B, a, b, r. 7o. Tı, T2
g Initialization proposed in [Scetbon and Cuturi, 2022]
repeat
Q = Q R 9=9:
Q 4AQ
VR (aQdiig(1///) Rl BR diag (/g) + 2A 2 Q1,
diag(g/g)//) AQ diag (1/g) 2 R1,,1
w < D(Q" (Qdiagl//)))) BR), V - w/g"
7 - 70/ max(|/ /R (Vllss)
c(1) +- Q exp(
g(2) + o exp(->VR), 5(3) + g0epp((7aa
Q, 4 ULR-TI-Dykstraaa, b, TI, T2, (Alg. 3)
until A((QRR,9) (Q, R) 7) < <;;
Result: Q, R.g
Convergence and Complexity. Similarly to linear ULOT, the unbalanced Dykstra algorithm is guaranteed 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(n2² + m²)r) time and O(n m²) memory complexities. However, as in ULOT, we can take advantage of low-rank approximations of both costs matrices and B reach linear complexity. Indeed, assuming
A and B = Bj B2 where A1, A2 E Rxxx and B1, B2 Rmxdy, then the total time and memory complexities become respectively O(mr(r + dyy) + nr(r + dx)) and O((n + m)(r + dx + dy)). Again, when A and B are distance matrices, we use the algorithms from [Indyk et al., 2019].
3.4 Unbalanced Low-rank Fused- Fused-Gromor- Send-Grrooo-aassesttone
We finally focus on the increasingly impactful [Klein et al., 2023] BeaddGGromov-wasserscce problem, which merges linear and quadratic objectives [Vayer et al., 2018]:
FGW v) == min +
(15)
Pell..b
where a E [0, 1] and & := 1-c allows interpolating between the GW and linear OT geometries. This problem remains a GW problem, where one replaces the 4-way cost Mi,, i j.j'' (A,'-BJ,'') )) appearing in (4) by composite interpolated cost between the OT and GW geometries, redefined as Mii j''' = aCt,, + @( a(A,,² )² Our proposed unbalanced and low-rank version of the FGW problem includes PP| := ||P|| the mass of P. to homogenize linear and quadratic terms,
ULF
min PCC P\ + (QA. kL ma) 12KL((" 1, (b) P: (P))r
(16)
which is expanded through the explicit factorization of P, noticing that |P| [g| == 119
ULFGW, v) (Q.R.g)en, alg) R, + a R, Gn, R.
min
(17)
Then by linearizing again H. R.9) alg) R, aQA.B (...((4 (Q, R. g) with an added KL penalty and leaving Gu,b unchanged, we obtain at each iteration, the same optimization problem as (14) where the kernels &u are now defined as
& = (E( Sk
(2)
(3)
< gb O exp(->x H = algx algx|C'R diag L/9x) + (2A 1, + 4AP, diag(1/ l/th))
(1) = Qx O expl - TO HHH 4h = RR 2 R
V NH CCT Qs diag (19x))))))))))))))
(2B 1, 4BP AQu diag (1/ /9x))
V, = a (IC, Px) w / - 4aw quad /
[ @ CR
(1 Q AP. ViE (1,.....
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.
7

---

Negn measured
Nign predicted
Ste24a2 meseured
Ske24a2 predicted
4
24
(a) Gene Nrgn: measured vs. predicted expression
(b) Gene Slc24a2: measured vs. predicted expression
Figure Exp. Spatial visualization of two mouse brain sections, contrasting observed vs.
predicted (using ULFGW) spatial distributions of expression levels, for two different genes.
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 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 (LRR) counterparts on spatial transcriptomics task, and follow in Exp. 3 by comparing ULR solvers to entropic (E) counterparts on smaller task, to accommodate entropic solvers' quadratic complexity. We conclude Exp. by comparing ULR solvers to [Thual et al., 2022], which can learn sparse transport coupling, in the unbalanced FGW setting.
Datasets. We consider two real- world datasets, described in B.1, and two synthetic datasets, that are large enough to showcase our solvers. The real-world datasets consist of both shared feature space, used to compute the costs matrices for the linear term in the well
Tissue regian ground truth
Tissue region predicted
CTX
CTX_2
Hb: Vi
HPFimo Megs
OT and FGW settings, as geometries that are specific to each source and target mea- sures, and which are used to com- pute the costs matrices for the quadratic term in the GW and FGW settings. In Exp. 1, sim- ply consider two isotropic Gaussians in d = 30 to evaluate the performance of the TI variant on a liner problem. We the mouse brain STARmap spatial transcriptomies data from [Shi al., 2022] for Exp. 2 and Exp. 3. We use data from the Individual Brain Charting dataset [Pinho et al., 2018]. replicate the settings [Thual et al., 2022], in Exp.
STA
TH
Figure 2: Visualization of measured and predicted tissue regions
in the mouse brain in Exp. 2
Following Klein al. [2023], we evaluate maps by focusing on the two following metries: (i) pearson correlation computed between the (ground truth) source feature matrix F E Ruxx
Metries.
and the barycentric projection of the target to the source scaled by the target marginals D'. Writing P as the transport matrix from source to target, this can be computed as Pdiag( ) when assessing class transfer (among vector of labels (", taken in (1,... 111"", and the inferred labels for the same points, predicted for each by taking the argmax, Br,, where B matrix of n. X 11 row probabilities, each the barycentric projection of the target one-hot encoded labels L' e {0, m == Pdiag ( LA
Ft; (ii) F1 score possible classes), computed between the original source
Benchmarking The Translation Invariant Variant. We evaluate the effect 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 ott- implementation, and single NVIDIA GeForee RTX 2080
Experiment 1:
8

---

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 convergence threshold and 1000 maximal number of iterations for Dykstra, in 64-bit precision. We observe that the use of our proposed objective is consistently beneficial when solving ULR problems. See also Appendix B.3 for additional experiments.
rank=10, 10, T2 T2
rank=10, T2
method Sintham
16
methes
Sinkears
50
(TO
14
Siieher
[f!]
40
time (s)
time (s)
10
30
10 25
50
100
250
500
10 25 50
100
250 500
#samples (thousands)
#samples (thousands)
Figure Execution time of unbalanced LR Sinkhorn, with (Alg. 3) or without (Alg. 5) the TI variant. We fix the rank to 10; n points (displayed in thousands) are sampled from two Gaussian distributions in = 30 of means respectively 112 and 1.3, and standard deviations and 0.2. (left) displays large (close to balanced), (right) is smaller (more unbalanced). We use the same convergence threshold for the outer loop, for all sample sizes. As n gets bigger, this results in a relatively looser threshold, explaining why timings slightly decrease w.r.t. n. What matters is, therefore, the comparative performance TI vs non- for a fixed n, not the behaviour w.r.t. n.
Experiment 2: ULOT vs. LOT on Gene Expression / Cell Type Anno- tation. We evaluate the accuracy of ULOT solvers for large-scale spatial transcriptomics task, using gene ex- pression mapping and cell type anno- tation. We compare it to the balanced LR alternative using the Pearson cor- relation p as described in the metrics section. We leverage two coronal sec- tions of the mouse brain profiled by STARmap spatial transcriptomics by [Shi et al., 2022]. They consist of n Each cell is described by 1000 gene features, in addition to 2D spatial coordinates. As a result A, B are z 40k x 40k, and the fused term C' is -uuured-Euclidan distance matrix on 30D PCA space computed on the gene expression space. We selected 10 marker genes for the validation and test from the HPF__C cluster. We run extensive grid search 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 the other genes in Table as well 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.
solver
mass %
val p
test
Fl mac.
Fl mic.
Fl weig.
LOT
1.000
0.282
0,386
0.210
0.411
0.360
ULOT
0.889
0,301
0.409
0.200
0,425
0.363
LGW
1.000
0.227
0,288
0.487
0,716
0.692
ULGW
1.001
0,222
0,287
0.463
0,701
0.665
LFGW
1.000
0.365
0.443
0.576
0.720
0.714
ULFGW
0.443
0,379
0.463
0,582
0.733
0.724
Table I: Exp.2, Results for spatial transcriptomicss dataset (brain coronal section from Shi al. [2022]).
40, 000 cells in both the source and target brain section.
Experiment 3: ULOT vs. UEOT. We compare the performance of ULOT solvers to their unbal- anced entropic alternatives (UEOT). We use the same datasets as in Exp. 2, but must pick smaller subset (Olfactory bulb), to avoid OOM errors for entropic UGW solvers, since they cannot handle the 40k sizes considered in Exp. 2 (see B.I). This results in 20, 000 source and z 15, 000 target

---

source
ground-truth
ulfgw (ours), P = 0.8180
FUGW, 0.6324
24
9.3
4.6
-7.7
-4.6
-7.
-15
-24
-9.3
-14
Figure 4: Visualization of measured and predicted right auditory click contrast map in Exp.4.
cells, and 1000 genes. Similar to Exp. 2, the fused term C' is squared-Faclliden 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_ I. We run an extensive grid search, as in Exp. and B.2. Table shows that ULFGW out- performs entropic solvers w.r.t. Pear- correlation P, but worse when considering FI scores. On the other hand, ULFGW confirms its superior- ity compared to the balanced alterna- tive LFGW. Taken together, these re- sults 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.
solver
mass %
val /
test
Fl mac.
Ff mic.
weig.
UEOT
1.012
0.368
0.479
0,511
0.763
0.751
LOT
1.000
0.335
0.440
0,511
0.760
0.751
ULOT
0.998
0.356
0.461
0.518
0,770
0.762
UEFGW
1.015
0.343
0.475
0,564
0.839
0.831
LFGW
1.000
0.348
0.453
0,512
0.762
ULFGW
0.339
0.368
0.491
0.556
0.826
0.818
Table 2: Exp. 3: Results for spatial transcriptomics dataset (Olfactory bulb section from Shi et al. [2022]).
Experiment 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 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, draw inspiration from Pinho et al.'s smaller scale notebook implementations: We embed the n s 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 nt, i.e. 16k points. We evaluate methods by comparing the performance of the best hyperparameter combination, based on the average correlation between the barycentric projection and ground- truth value of features, across test set of 5 contrast maps. In Table 3, we observe that ULFGW and LFGW outperform FUGW-sparse In this setting, there is 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.
solver
val
test
FUGW-spart sparse
0.999
0.492
0.472
LFGW ULEGW
1.000
0.513
0.663
0.981
0.533
0.643
Table 3: Results the brain anatomy with functional signal data from Pinho et al. [2018] in Exp.4.
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 this paper to merge these two strains, and demonstrated the practical relevance of these unbalanced low-rank solvers on various challenging alignment tasks.
Conclusion.
10

---

References
Ravindra K Ahuja, Thomas L Magnanti, and James B Orlin. Network Flows: Theory, Algorithmss, 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 (rXiv:/705.096344 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 computer vision and pattern recognition, pages 1534-1543, 2016.
Heinz H Bauschke and Patrick Combettes, A Dykstra-like algorithm for two monotone operators.
Pacific Journal of Optimization, 4(3):383-39) 2008.
Heinz H Bauschke and Adrian S Lewis. Dykstras algorithm with bregman projections: A convergence proof. Optimization, 48(4):409-42) 2000.
Andrew] J Blumberg, Mathieu Carriere, Michael A Mandell, Raul Rabadan, and Soledad Villar. Mrec: fast and versatile framework for aligning and matching point clouds with applications single cell molecular data. arXiv preprini urXiv:2001.01660 2020.
Emmanuel Boissard and Thibaut Le Gouic. On the mean speed of convergence of empirical and occupation measures in wasserstein distance. In Annales de I'THP Probabilités et statistiques, volume 50, pages 339-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.11,,4777.7 URL https: //www. //ww..boorxiv org/ content/ conttnt/eearly 10021/22/1//02/// 12. 15.115.4775 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 Almet, Adam Stabell, Raul Ramos, Maksim V Plikus, Scott X Atwood, and Qing Nie. Screening cell-cell communication 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, 12:9912-9924, 2020.
Lénaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François- Seaner Vialard. Unbalanced optimal transport: Dynamic and kantorovich formulations. Journal of Functional Analysis, 274(11): 3090-3123, 2018
Lenaic Chizat, Pierre Roussillon, Flavien Léger, Trançois-Xaviee 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:2/04.02013 2021.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in neural information processing systems, pages 2292-2300. 2013.
Marco Cuturi, Lactitia Panaxaaapss,ssss, Yingtao Tian, Charlotte Bunne, Geoff Davis. and Olivier Teboul. Optimal transport tools (ott): A jax toolbox for all things wasserstein. arXiv preprint arXiv:2201. /2324, 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-wasserstemn optimal transport to align single-cell multi-omics data, bioRxiv, 2020. doi: 10.1 10.1101/2020004.8..66677
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 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 htttps://proceediiiiiie mlr. prress/v139/fatrrs2111 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-MIC 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 Alaya, Aurélie Boisbunon, Stanislas Chambon, Lactitia 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((8):1-8) 2021. URL attp://jmll org/ pttppppappp//////////ssppp 2/t///2222222111112110 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 Infelligence 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 PenngPaaaxantths, Michael Sterr, Aimee Bastidas- Poncee, Marta Tarquis-Meding Heiko Lickert, Mostafa Bakhti, Mor Nitzan, Marco Cuturi, and Fabian J. Theis. Mapping cells through time and space with moscot. bioRxiv, 2023. doi: 10. 1 101/2023.05 1,540374. URL hattps://wwt biorxiv. wwwp:orrgg/coaoooe hetg//ooteen//errrr 220///222222//2 00/00////1/////1//// 05. 11. 11.440374.
Harold W Kuhn. The hungarian method for the assignment problem. Naval research logistics
quarterly, 2(1-2):83 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 Learning Representations, 2023.
Facundo Mémoli.
distances and the metric approach to object matching. Foundations of Computational Mathematics, 11(4):417-487, 2011.
12

---

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, Mare Szafraniec, Vasil Khalidov, Pierre Fermandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, et al. Dinov2: Learning robust visual features without supervision. arXiv preprint arxivv2304.071193, 2023.
Gabriel Peyré, Marco Cuturi, and Justin Solomon. Gromov-wasserseinn averaging of kernel and distance matrices. In International Conference Machine Learning, pages 2674-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-Desmidr, Séverine Roger, Laurence Laurier, Véronique Joly- Testault, 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. hal- 0a8-718.
doi: 10. 1001008/sdata.2118..uk
URL https: //hal. science/
R Tyrelll Rockafellar. Convex analysis. Princeton university press, 1970.
Yossi Rubner, Carlo Tomasi, and Leonidas J. Guibas. The earth mover's distance metric for image retrieval. International Journal of Computer Vision, 40(2):99-12) 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 heterogeneous spaces. arXiv preprint arXiv:2002.016655 2020.
Meyer Scetbon and Marco Cuturi. Linear time sinkhorn divergences using positive features. Advances in Neural Information Processing Systems,
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-66 6814 Curran As- neurips cc/ cc/paper... llappe_rflllllllelpplpplllpppple fiwes/ pap.e///ppp///22//////// sociates, Inc., 2022. URL http:: Paper Conference Paperroofeerreeeeeeeeeeee pdf. tt10222668711022777661111811882
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. Lincar-time Soomov-Waseerser distances using low rank couplings and costs. ICML, 2022.
Geoffrey Schicbinger, Jian Shu, Marein Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. Optimal-transper analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell, 176(4):928-94)) 943, 2019.
Thibault Séjourné, François-Xavier Vialard, and Gabriel Peyré. The unbalanced gromov wasserstein distance: Conie formulation and relaxation. Advances in Neural Information Processing Systems, 14:8766-8779 2021.
13

---

Thibault Séjoumé, François Xavier Vialard, and Gabriel Peyré. Faster unbalanced optimal transport: Translation invariant sinkhorn and 1-d frank-wolfe. 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. 20..2999066..491144 URL https: //www. biorxiv wrg.cootent/eerlly////aa//// 06/22/2022 06. 496914.
doi: 10. 10.1
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: //tp::///rree http///proeeee....//o///// mlr.
www.mml....mm...eee prwss/vessvi paaess//vvvitttaa tatty/0aa...../ html.
Alexis Thual, Huy Tran, Tatiana Zemskova, Nicolas Courty, Rémi Flamary, Stanislas Dehaene, and Bertrand Thirion. Aligning individual brains with fused unbalanced Gromov-Wassertein. arXiv, 2022.
Titouan Vayer, Laetita Chapel, Rémi Flamary, Romain Tavenard, and Nicolas Courty. Fused gromov- wasserstein distance for structured objects: theoretical foundations and mathematical properties. arXiv preprint 1811 122834, 2018.
Cédric Vincent-Cuaz, Rémi Flamary, Marco Corneli, Titouan Vayer, and Nicolas Courty. Semi- relaxed gromov-wasserse divergence and applications on graphs 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 wromov-wasserstiii 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. Cur- Associates, Inc., 2019. URL U neursddddinnnnnnnnns.eee..nu wttpssscccppppepeeeeeeeeee saaper_fflles wiles/paper. 0001/111666ca222277711111////222 pdf.
14


Reviews:
[{"review_id": "LHD9cRodo6", "reviewer": "Reviewer_fjNH", "summary": "In 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.", "questions": "- 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$", "limitations": "not applicable", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "OT 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.", "weaknesses": "- 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Many thanks for your time and for your comments.\n\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\nYes, the formulation is indeed non-convex, as is also the case for all original unbalanced solvers we build upon.\n\nWhile this removes several theoretical guarantees, practice suggests a “mild” form of non-convexity, comparable to e.g. k-means or non-negative matrix factorisation (NMF). \n\nConvergence to a critical point is ensured through Dykstra, but the solution depends on the initialization. We use the standard approaches considered for the balanced case (see e.g. https://github.com/ott-jax/ott/blob/main/src/ott/initializers/linear/initializers_lr.py )\n\n> **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\nSure, in [Scetbon et al. , 2021]  \"Low-rank sinkhorn factorization\", the authors obtain in Appendix A a generic approach to show the stationary convergence of the mirror-descent scheme as soon as the objective function to minimize is relatively smooth w.r.t the negative entropy. Also the authors show the constrained version of the objective defined in (5) is relatively smooth w.r.t the negative entropy. Therefore the stationary convergence of our algorithm is a corollary of the result obtained in [Scetbon et al. , 2021]. We will clarify this point in the paragraph on \"Convergence and Complexity\" written page 4.\n\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\nFor convex OT problems (e.g. Sinkhorn), the usual approach is to compute the gradient norm (marginal violations in balanced case) to monitor convergence. For nonconvex OT problems (unbalanced solvers but also Gromov-Wasserstein and Wasserstein barycenter problems with free support), it is more usual to monitor changes in the solution (as in, e.g., k-means). \n\nNote also that our iterates are all bounded (they are transport plans $Q,R$ and simplex vector $g$), there is therefore no \"divergence\" we might miss (as in your harmonic series example), and we have never seen such behaviour in practice.\n\n> **In the last cost of eq(2) I believe it should be diag(1/g) not diag(g) [...]**\n\nIndeed, we apologize for this unfortunate typo early in the paper. Thanks for pointing it out.\n\n> **reparamaterized, repamatrization**\n\nFixed! we will use reparam**e**terized and reparameterization (https://en.wiktionary.org/wiki/reparameterize#English ).\n\n> **mass conversation for mass conservation**\n\nWe apologize for this typo, likely due to an unfortunate auto-correct.\n\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\nTo clarify: we use `\\citep` to refer to the *paper*, and `\\citet` to refer to the authors of the work. We meant therefore: “the framework of [these papers]” and “notations from [this paper]”. When crediting a group of authors we use `\\citet`, as in line 132 (`Séjourné et al. [2022] have proposed`). We have found a few mistakes disagreeing with this convention, we have fixed them.\n\n> **towards a stationary points, such that with probability at least 0.99 that: extra \"that\", Algo 1: (check me), paper uses \\star and * [...]**\n\nMany thanks for pointing out these typos and unfortunate \"left over\" comments.\n\n"}, {"review_id": "sLjMk0lLpM", "reviewer": "Reviewer_ckny", "summary": "The 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.", "questions": "In Algorithm 2: typo in 1 / gamma / tau1\n", "limitations": "The authors do not discuss the limitations of their work.", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The 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.", "weaknesses": "It 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Many thanks for the wealth of comments you have provided.\n\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\n⧐ Thanks. We agree with your suggestion and have removed many technical details from the main section to the appendix to free up space for experiments.\n\n> **While adding translation-invariant improves the convergence of the usual unbalanced OT, it is unclear about the real improvement in the experiments.**\n\n⧐ We agree that this was missing. We have now added such comparisons, showing a clear-cut benefit when using our TI variant. Please look at our Pdf, Figure X **[TODO Michal]**\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⧐ As you mention, we felt it was instructive to start with GW. While Alg. 4 is a particular case of Alg. 5, we believe the most important conceptual gap lies in going from 3 to 4. We prefer following your recommendation to free up space by moving Alg. 5 (ULFGW) to the appendix.\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⧐ Thanks for your suggestions, here are our answers:\n\n- **It is not clear to us why and how mini-batch OT could be a competitor to handle the large scale unbalanced GW problems we solve**. Mini-batch OT was introduced to *minimize* a _linear_ OT _loss_ (e.g. [FZFGC] Fatras, K., Zine, Y., Flamary, R., Gribonval, R., & Courty, N. (2019). *Learning with minibatch Wasserstein: asymptotic and gradient properties*), **to train typically a generative model**. There, the goal is to minimise a fitting term $\\mathcal{L}(\\theta)=W(p_\\theta, q),$ and both (or either) $p_\\theta$ and $q$ are continuous densities. The argument in favor of mini-batch is that sampling many times from $p_\\theta,q$ and/or $q$ to average multiple gradients for $\\theta$ is better than using the gradient computed from a single bigger batch.\n\n   All our experiments consider instead two **fixed**, large datasets of point-clouds, in a _quadratic_ GW setting. Here a mini-batch approach would result in an arbitrary partitioning of each measure into sub-measures, and then match arbitrarily these smaller sub-measures. This would be very naive, suboptimal, and we are not aware of papers advocating this except when doing more elaborate hierarchical approaches (e.g. MREC, https://arxiv.org/abs/2001.01666). MREC was compared to LRGW (see [Scetbon+22]), but does not have an unbalanced formulation.\n\n- **Quantized GW is not an unbalanced method**, therefore we feel it would be out of place in this study. Yet, we have considered one of their experiments (the scene dataset) in our new experiments (see pdf in global response). The authors of Quantized GW provide very little details on their experimental setup (hyperparameter tuning in particular) hence we did not re-run their solver. We only used their reported accuracy and timing in the text (which we assume is the best they could get). With the limited time of this rebuttal, we get roughly the same accuracy as shown in our pdf using ~3 min of compute.\n\n- Note that all experiments in 4.3 compare two brain meshes, i.e. two graphs endowed with shortest path metrics.\n\nTo summarize our thoughts, we believe that the natural competitor for our setting remains entropic unbalanced solvers, because:\n- they are by far the most studied alternative, having been used in dozens of papers;\n- they can be used in an unbalanced setting in a fairly comparable way to ours;\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\nYou are right. Because we were short of space at the time of submission, we mostly reported results and not experimental details (the benchmark is not ours). However, this can be easily corrected as we can expand on the material in our supplementary and bring it back to the main body. Thanks to the space freed by Algorithm 5, we can now rebalance this section, and move back relevant background in 4.3.\n\n\n"}, {"review_id": "wZNSsFnhgZ", "reviewer": "Reviewer_1XVN", "summary": "In 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. ", "questions": "1. 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.", "limitations": "The limitations are not discussed in this paper.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- 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", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Many thanks for providing so much feedback, and giving 3 **good** grades to the *soundness* / *presentation* / *contribution* of our paper. We hope we can convince you to raise your final score with the answers below.\n\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⧐ This is an excellent suggestion, and we will emphasise this more clearly around lines 91.\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).**\n\n⧐ We apologize for this lack of clarity. We have incorporated all the changes you have requested.\n\n> **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⧐ Finding an efficient initialization for Alg.1 is something of a research topic in itself. For instance, the ott-jax toolbox (https://github.com/ott-jax/ott/blob/main/src/ott/initializers/linear/initializers_lr.py) implements 4 different ways to do this. We will mention this line of work.\n\n> **In line 95, what is the definition of nonnegative rank of ? Please add this definition to the revision of this paper.**\n\n⧐ Yes, the definition is.that in [Scetbon21,22], we will provide it again.\n\n> **In equation (2), the term [...]**\n\n⧐ Yes, we apologize for these two bad typos\n\n> **In equation (7), are there any differences between KL? If any, they should be defined explicitly in the paper.**\n\n⧐ We're sorry about this typo. These 3 terms all refer to the generalised KL computed either for vectors of matrices with positive entries, $ \\textrm{KL}((p,q) = \\sum_i p_i \\log \\frac{p_i}{q_i} -p_i +q_i$. We will switch to $\\textrm{KL}(\\cdot|\\cdot)$ for all.\n\n> **Are there any theoretical guarantees that the tuples approximate the solution of the optimization problem in equation (6)?**\n\n⧐ Such guarantees cannot be achieved, to our knowledge, because the problem is at least as non-convex as the balanced case.\n\n> **In the proposed algorithms for solving ULOT, ULGW and ULFGW, which values of the rank hyperparameter should we choose? Would it be as low as possible?**\n\n⧐ In all experiments, we have used cross-validation on train sets to set rank $r$. Lower $r$ means faster computations, but they might end up in factorized transports that are too coarse to perform well. We believe $r$ should be set depending on the problem size vs. compute ability, much like k-means (or, to some extent, entropic regularization).\n\n> **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\n⧐ Thanks for these two very relevant references, and we apologize for having missed them, we have added them to our draft.\n\n> **In the definition of cost matrix in Section 2, the index notation [...]**\n\n⧐ Sure, we happy to change this.\n\n> **Notations: When introducing the unbalanced OT in equation (3), the authors should explain the notations KL in equation (4) should be defined explicitly.**\n\n⧐ We updated our paper, and followed your advice to clarify these notations from the start (KL and Hadamard exponent)\n\n> **Typo: in line 96, repamatrization --> reparametrization.**\n\n⧐ We had a typo indeed, thanks! It seems reparamet**e**rization is more common, so we will use it.\n\n\n"}, {"review_id": "bZ56g4Llnh", "reviewer": "Reviewer_GuCe", "summary": "The 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.", "questions": "See my question above.", "limitations": "They adressed the societal impact of their work.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- 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", "weaknesses": "- 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your encouraging review.\n\n> **The low-rank approximation is not sufficiently discussed in this paper. I mentioned that it is considered to accelerate the computations. However, 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. To my mind, using a low-rank plan seems contradictory with the property that the optimal plan is a (full rank) permutation in some setting. Could 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⧐ Your intuition is indeed correct:\n\nIn the general case, a coupling matrix $P^\\star$ minimizing the linear objective in Eq. 1 (e.g. that returned by an LP solver) will be (i) sparse (up to $n+m-1$ non-zeros values for a $n\\times m$ matrix) and (ii) full rank.\n\nThe idea of low-rank constraints and/or entropic regularization is therefore to move *away* from these properties. \n- For the low rank approach, this is akeen to requiring a \"cluster\" structure when computing transport (i.e. all of the mass must go through $r$ virtual points)\n- For entropic regularization, this forces all sources and targets to transfer mass between each other.\n\nBoth can be seen as **inductive biases** when trying to match points in large dimensions, that result in substantial computational and statistical gains. Both result in a maximally diffuse / rank 1 transport coupling $ab^T$ in the limit, when $r=1$ or when entropic regularization is infinite.\n\nThese aspects are already well covered in the references provided in lines 45-50, which is why we did not use to much space in the introduction for these reminders. Following your comment (and some space gained my moving some of the more technical aspects to the appendix) we will add the sentences above (abridged).\n\n> **To my mind, the authors do provide complexity per iteration of their algorithms,**\n\n⧐ We have improved this aspect, and have added compute complexities alongside iteration costs to be more explicit, as was done e.g. in [Scetbon'22]\n\n> **but the experimental aspect of measuring the time performance of their algorithms is omitted. In particular, have you checked experimentally that this 'translation-invariance' variant in Section (3.2) does accelerate the computations in your low-rank setting ?**\n\n⧐ This is an excellent suggestion. We have now added various experiments that illustrate such speed-ups (see Figures 1 and 2 in attached Pdf). We agree this strengthens our submission and will continue exploring to what extent this TI modification is always a \"win\" (we have seen it become slightly less effective for very \"loosened\", i.e. low $\\tau$ values). For now however, these figures reflect what we saw in practice in our experiments.\n\n> **Could 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 ?**\n\n⧐ Because Sinkhorn and Low-rank approaches solve slightly different problems, it is not easy to compare them in terms of timing. Note, however, that for the **balanced** such comparisons were proposed for the Linear OT problem in [Scetbon et al'20, Scetbon/Cuturi 22] whereas [Scetbon et al'22] compares them in a quadratic setting (GW) problem. \n\nIn the unbalanced case, this is even harder, because there is not a 1-to-1 correspondence between their regularizers, convergence criteria etc...  In other words, below 50k points, it's very easy to find a setup where unbalanced Sinkhorn and unbalanced LR Sinkhorn run roughly in the same time, by tweaking thresholds, regularizers etc... **What we can confidently report is that Unbalanced Sinkhorn basically blows up above 50k points, whereas Unbalanced LR Sinkhorn can be easily run \"off the shelf\", with 500k points and ranks of $r=10$ and aboveas shown in Fig. 1.\n\n> **Also, the use of unbalanced OT involves two extra parameters  which requires cross-validation. Could the authors give plots of the performance of their biology task as a function of…  This would provide interesting insights on the interpretation of these parameters, and their impact on learning tasks.**\n\n⧐ All of the results provided in the paper do indeed deal with these two parameters using cross-validation. We did not have the time (nor space left in the pdf) to provide accuracy (surface) plots as a function of these two, but we will prepare them for the camera ready. Since our grid for $\\tau$ was quite small, we will expand it to provide a more detailed picture as well. This might indeed be useful for practitioners, to see if there are \"jumps\" (we did not notice them so far, rather smaller variations)."}]

Author rebuttal:
**We thank the reviewers, AC and SAC assigned to this paper for their time and work looking into our submission**. 

We thank them in advance for reading our rebuttal and interacting with us for a few more days during the discussion period.

We were happy to see that the paper was overall well received by all 4 reviewers:

**fjNH** : *The proposed formulation is sound, the treatment is relatively easy to follow and the derivations are made clear.*

**ckny**: *The extensions [...] are not trivial and require some calculation effort [...] the writing is very instructive and all technical details are clearly presented and on point.*

**1XVN**: *the results in this paper are important to some extent[...]*, *This work is a novel combination*

**GuCe**: *The contributions might look incremental at first (combining unbalanced OT with low-rank OT), but are actually very thorough.*

The most important weaknesses highlighted by reviewers point to:
- some sloppiness in our presentation: we proposed several tightly related, yet distinct algorithms. This led to some redundancy, as discussed below in our answers. Because of a lack of space, some of the experimental sections were left to their minimum, with only (good!) results presented. As mentioned by reviewer **ckny**, this is true in particular for Section 4.3.

➡ *We have incorporated comments by reviewers in our draft, and have removed an algorithm from the main body to the appendix, and added back experimental details.*

- some clarifications in experiments, e.g. importance of using the translation-invariant adjustment, and comparing to quantized GW (which is not, however, an unbalanced method). 

➡ *We have run novel experiments following their remarks. They all reinforce the message of our submission, and will be added either to appendix or main body, depending on space.*

- a few "bad" typos (e.g. $D(g)$ instead of $D(1/g)$ and $R$ instead of $R^T$).

➡ *We apologize for these unfortunate typos, we have corrected them.*

We believe we have addressed all of the points raised by reviewers. In addition, we will release code for *all* of the tools presented in the paper in one of the major OT python toolboxes. Our implementation can run on CPU but also natively on GPU. We will share the code before the end of August.

At the moment, all reviewers have scored our paper as **good (3)** in *both* **soundness** and **presentation**. We have received an average 2.75 grade (2,2,3,4) in **contribution**, and we hope our rebuttal alleviates these concerns.

We believe that the very supportive words found in all reviews are not reflected in the current distribution of (fairly low) scores of 7,4,4,3 . If reviewers agree with our assessment, we humbly ask them to reconsider their score. 


Meta-review: The paper proposes a solution to the problem of efficiently compute discrete unbalanced Optimal Transport (OT) problems.
By assuming that the transport plan is a low-rank matrix, they leverage on this property to reduce
complexity of the global algorithm. The authors extend this to unbalanced GW and unbalanced Fused Growov

It is a sound and well built paper.

From my reading of the paper and discussions with the reviewer and SAC, it seems that the ideas proposed in the paper are
indeed incremental in several of their aspects albeit well executed, 

* Section 3.1 borrows heavily on Scetbon et al. (2021, 2022)
* Section 3.2 is an extension of Sejourné 2022 Proposition 2
* Section 3.3 is an extension of Scetbon et al (2022) 
* Section 3.4 is a mix of 3.1 and 3.3 and the algorithm subsumes the one in Section 3.3

However the methodology can indeed be useful for practictioners.

Given this very borderline situation, after many discussions, we have decided on acceptance, given the
full batch of papers of the SAC and the resulting calibration.
Decision: Accept (poster)