Low-Rank Optimal Transport through Factor Relaxation with Latent Coupling
Authors: Peter Halmos
Xinhao Liu
Julian Gold
Benjamin Raphael
Venue: NeurIPS 2024 poster (2024)
Keywords: Optimal Transport
Sinkhorn
Low-Rank
Matrix Factorization
Primary area: other
TL;DR: A general framework for low rank optimal transport using a latent coupling matrix and relaxed projections.

Abstract: Optimal transport (OT) is a general framework for finding a minimum-cost transport plan, or coupling, between probability distributions, and has many applications in machine learning. A key challenge in applying OT to massive datasets is the quadratic scaling of the coupling matrix with the size of the dataset. [Forrow et al. 2019] introduced a factored coupling for the k-Wasserstein barycenter problem, which [Scetbon et al. 2021] adapted to solve the primal low-rank OT problem. We derive an alternative parameterization of the low-rank problem based on the _latent coupling_ (LC) factorization previously introduced by [Lin et al. 2021] generalizing [Forrow et al. 2019]. The LC factorization has multiple  advantages for low-rank OT including decoupling the problem into three OT problems and greater flexibility and interpretability. We leverage these advantages to derive a new algorithm _Factor Relaxation with Latent Coupling_ (FRLC), which uses _coordinate_ mirror descent to compute the LC factorization. FRLC handles multiple OT objectives (Wasserstein, Gromov-Wasserstein, Fused Gromov-Wasserstein), and marginal constraints (balanced, unbalanced, and semi-relaxed) with linear space complexity. We provide theoretical results on FRLC, and demonstrate superior performance on diverse applications -- including graph clustering and spatial transcriptomics --  while demonstrating its interpretability.

Low-Rank Optimal Transport through Factor Relaxation with Latent Coupling
Peter Halmos Xinhao Liu"., Julian Gold"", and Benjamin J. Raphael
Department of Computer Science, Princeton University Center for Statistics and Machine Learning, Princeton University
Abstract
Optimal transport (OT) is general framework for finding a minimum-cost transport plan, or coupling, between probability distributions, and has many applications machine learning. key challenge in applying OT to massive datasets is the quadratic scaling of the coupling matrix with the size of the dataset. Forrow et al. (2019) introduced a factored coupling for the k- W-sssersteinn barycenter problem, which Scetbon et al. (2021) adapted to solve the primal low-rank OT problem. We derive an alternative parameterization of the low-rank problem based on the latent coupling (LC) factorization previously introduced by Lin et al. (221) generalizing Forrow et al. (2019). The LC factorization has multiple advantages for low-rank OT including decoupling the problem into three OT problems and greater flexibility and interpretability. We leverage these advantages to derive new algorithm Factor Relaxation with Latent Coupling (FRLC), which coordinate mirror descent to compute the LC factorization. FRLC handles multiple OT objectives (Wasserstein, wannoo--ssserttenn Fused constraints (balanced, unbalanced, and semi-relaxed) with linear space complexity. We provide thcoretical results on FRLC, and demonstrate superior performance on diverse applications including graph clustering and spatial transcriptomics while demonstrating its interpretability.
and marginal
Introduction
Optimal transport (OT) is a powerful geometric framework for comparing probability distributions. OT problems seek a transport plan P efficiently transforming one distribution (a) into another (b), subject ground cost C. The minimum cost yields distance between and b, while the optimal transport plan reveals key structural similarities between the distributions. Owing to its versatility different ground costs result in different ways to compare data OT has found many applications in
machine learning and beyond: from self-attention Tay et al. (2020); Sander et al. (2022); Geshkovski al. (2023) and domain adaptation Courty al. (2014); Solomon al. (2015) computational biology Schiebinger et al. (2019); Yang et al. (2020): Bunne et al. (2023) Liu et al. (2023).
This versatility is compounded by several variants using different forms of the objective function and/or constraints on the transport plan P. Wasserstein (W) OT Kantorovich (1942) compares distributions over the same space through the expected work of P. while Gromov- Wasoov-stissrntiinn (GW) OT Mémoli (201T) compares distributions supported on distinct geometries through the expected metric distortion of P. Fused Gromov-Wasserster Wasserstein (FGW) Vayer al.) (2020) OT is suited to structured data, taking a convex combination of the former two objectives. Independently, one can relax constraints the marginals of P: in computational applications, P is matrix whose row-sum PI,m and column-sum PT 1,, are called its left and right marginals. Balanced OT requires P1, = a and 1 1,, b. Unbalanced OT Frogner et al. (2015) replaces these constraints with penalties in the
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

transport cost, and is more robust to outliers. Semi-relaxed OT can be used to understand how one dataset embeds into another by imposing one hard constraint on either the left or right marginal. used for feature transfer Dong et al. (2023), and alignment of spatiotemporal data Halmos et al. (0024).
An important consideration in applying OT' is the quadratic space of the transport plan. To address both the quadratic complexity and to provide robustness under sampling noise, Forrow et al. (2019) in- troduced another variant of OT, optimizing k-Wasserstein Barycenter proxy for the rank-constrained Wasserstein objective. Their approach factors the transport plan through a small set of anchor points called hubs. Generalizing this approach, Scetbon et al. (202T) introduce the factorization
P = Q diag (1/g)R* comprised of sub-coupling matrices Q and R sharing an inner marginal
9. meaning Q'1n R 1m = g. Building on this, Scetbon et al. [2021 2022 2023) derived
algorithms to compute low-rank optimal transport plans for the primal OT problem with general costs, extending low-rank OT to GW and unbalanced problems using factored couplings.
Interestingly, a different factorization of P was proposed by [in et al. (2021) in the context of k- G-Wasserstein barycenters. We call their factorization a latent coupling (LC) factorization, given by
P Qdiag(1/ l/go) (ddag(1//g) R R with two inner marginals gQ Q"1, and 9R R² and general coupling T. Lin et al. (202I) constrain the transport between and b through two sets of learned anchor points, where the factorization is defined by three transport plans computed from three cost matrices between the points and their anchors. This objective differs from that Forrow et al.] (2019); SSettbon et al. (2021), who seek minimal rank coupling with respect to single, fixed cost C. We observe that factored couplings Forrow et al. (2019) correspond to LC factorizations with diagonal T, suggesting the I.C factorization of Lin et al. (202T) may provide an alternative parameterization of transport plans for the low-rank OT problem considered in Forrow et al. (2019): Scetbon et al. (22T)) To our knowledge, this idea has not yet been explored.
Contributions. We present a new algorithm, Factor Relaxation with Latent Coupling (FRLC, with the informal mnemonic "frolic"), to compute a minimum cost low-rank transport plan using the LC factorization. Parameterizing low-rank transport plans with the LC factorization has number of advantages. First, optimization of the low-rank OT objective decouples into three OT sub-problems on the LC factors Q, R, T, leading a simpler optimization algorithm. Second, this decoupling provides straightfoward extensions of FRLC to low-rank unbalanced and semi-relaxed OT; similar extensions for factored couplings required additional work [Seetbon et al. (2023) beyond the balanced case. Third, the latent coupling T in the LC factorization provides additional flexibility to model transport between datasets with different numbers of clusters, and to model mass-splitting between these clusters, providing a high-level and interpretable description of P that differs from the factored couplings Forrow et al. (2019). FRLC computes the I.C factorization using a novel coordinate mirror descent scheme, alternating descent steps on variables (Q, R) and T. inspired by the mirror descent approach of /Seeteeoonn et al. (202)) We call the descent step on (Q, R) factor relaxation, as the factors Q and R have relaxed inner marginals, allowing FRLC to be solved by OT sub-problems. FRLC handles multiple OT objectives (Wasserstein,
Fused and marginal constraints (balanced, unbalanced, and semi-relaxed). We show FRLC performs better than existing saute-of-the-a low-rank methods on a range of synthetic and real datasets, retaining the interpretability Lin et al. (2021). and inheriting the broad applicability of Scetbon et al. (2021); Scetbon & Cuturi (002)) Scetbon et al. (2022 2023).
2 Background
Wasserstein OT. Let {ai...., vn) and {yı, ... ym be datasets in a metric space X, and let Ad be the probability simplex of size d. Through probability vectors a E An and b E Am, each dataset encoded as a probability measure: f Ei= a,oz, and v = Σn byggy. Let
Ia. : [PERT
P1m = a), 11.b (P R PX 1',, = b}, Hl. := Ha, II..b
Thus, Ha. b is the set of transport plans (probabilistic coupling matrices) with marginals and b. Given cost function e: X x X + R,, define the cost matrix C' e IRx* via Cij c(z,, U5). The Kantorovich formulation Kantorovich (1922) of discrete OT, also called the Wasserstein problem, seeks a transport plan P of minimal cost:
W(1, v) == min (C, Prr-
(1)
Pella,b

---

Wasserstein OT. In many applications, one wishes to compare datasets fæi...., Fn C
X' and {y1...... ym y across distinct metric spaces X and V. The waanoo-asssrrttiin (GW) objective Mémoli (2007 201 addresses the absence of a common metric or coordinate system through intra-domain cost functions 61 x X - R+ and C2 y x y ~ IR+, leading to intra-domain cost matrices Aik c1 F&) and Bj, Cz(Uj, v1). The GW objective function
(P. := Eijaa BB))PP;: Par quantifies the expected metric distortion under P, leading to the optimization problem:
GW( := min QA.((P). Pe
The Fused Gromov-Wassersie Wasserstein (FGW) objective function Vayer al. (020) is convex combination of the W and GW objectives, given as a P)r (1- QA.B P). for hyperparameter a E (0, 1).
Relaxed marginal constraints. Balanced OT ① constrains P to lie in la Unbalanced OT relaxes constraints P1, a and pT = b. replacing them with penalties in the form of KI. divergences (or other divergences, see Chizat al. (008)):
U-W(u, = min (C, P) + KL (PLmlla))) 1. PER
where TL: Ti 0 control the strength of each penalty. Semi-relaxed optimal transport relaxes exactly one of the hard constraints Pi, a and 1, = b in the same manner. The semi-relaxed version of (T) obtained by relaxing only the "right" marginal constraint on is:
SRR-W u) = min Pella,.
+
4)
while it's "left" marginal counterpart SR - v/) is defined analogously over P E II..b. using penalty rKL((1...1 Likewise, one can form semi-relaxed unbalanced GW and FGW problems.
Entropy regularization. The seminal work Cuturi (013b) introduced the Sinkhorn algorithm to solve entropy regularized version of (I) W (H, v) min minpee.. (C, P) ell PP) massively
improving the (n² log n) time complexity of classical techniques Orlin 1997); Tarjan (197)). Above, H is the entropy, (P)) = - Eiù P.j(lop P.j- - 1), and > is the regularization strength.
Low-rank regularization. The nonnegative rank rk4 (MM of matrix M is the least number of nonnegative rank-one matrices summing to M. For 1, define
(5) and let (r) NT. (r) IT.. To estimate Wasserstein distances with greater stability and accuracy under sampling noise, Forrow et al.] (2019) proposed a low-rank regularization on the coupling matrix, factoring the transport through a small set of anchor points. More explicitly, [Seetbon lt al. (2221) parameterized the set lI) through the set FCa,b(r) of factored couplings.
Ia, (r) = (P E Na, (P)
II,,(r) =
<r).
Q E Ia R &. FC), = (.. E R x Xm X (
The set (C).b(r) parameterizes (r)) through R. + Qdiag (1/gg(1//)))) RT, as shown by Cohen & Rothblum (1993).
Scetbon 2021 apply this factorization to solve the Wasserstein problem subject P (r) For general cost matrices:
W W v) = min (C,P)
(6)
p(ll
GW, unbalanced and semi-relaxed low-rank OT problems defined in 2
3 and 4 replacing
R x II.... or 11.,b bo with aunk-constrained counterparts (5). Scetbon & Cuturi (2022): Scetbon et al.
2022
2023) developed robust framework for solving all of these problems.
3 Factor Relaxation with Latent Coupling (FRLC) algorithm
3.1 Latent Coupling Factorization
We parameterize low-rank coupling matrices P E (r) using a factorization introduced in Lin
et al. (2021). which call the latent coupling (LC) factorization (Fig. []. The key property of
3

---

b
a
90
T
P
Figure I: (Left) The LC factorization P = Qdiag(1 (1/ / ( diagg(L/ggg R of coupling matrix P
with outer marginals a, inner marginals giz gn. factors Q, R, and latent coupling T. (Right) Full-rank coupling matrix P.
this factorization the presence of a coupling matrix T linking two distinct inner marginals. For simplicity describe this factorization using r-dimensional latent space, but also extend to non-square matrices linking two latent spaces of different dimensions, as demonstrated in the results.
Definition 3.1 (Inner marginals). Given factorization P QX R. of coupling matrix P TT. (r)) the inner marginals of Q and R are gQ := Q²1,, 1, and gn == RT1m m. respectively, where 9Q.90 9R Ar.
To distinguish the different marginals, refer to and b as outer marginals.
Definition 3.2 (LC factorization). Given a coupling matrix P E Il,,(r)), a latent coupling (LC) factorization of P is P = Qdiag(1/gg) ( diag (afaa((1///))) /9 where go and gR are the inner marginals of Q and R. Q Ua,,, ., R ll.,,, and T 1l
We call factors Q, R, T an LC factorization sub-couplings. Let R. == R4 Ruxr x R Ruxr B Given probability vectors a E A,, b E Am and positive integer rank let
(C), == [(Q, R, T) E R+ Q E IT,,, R lT.,, T T 1ggggR
be the set of admissible sub-couplings for the LC factorization. Definition 3.2 gives the following map from (C), to (r):
(Q, R,T)) Qdiag(1/9q, Tdiag(1/gn) P.Q.R.T).
Since this map is surjective, the set (r). parameterizes 1I, (r). Surjectivity follows from the fact that (r) maps injectively into (r), through (Q, R. g) - (Q. 0 R. diag(g) and (r) maps surjectively onto (r) via (Q, R, + (1 Definition 3. and Definition 3.2 are readily extended in two directions: the case when the outer marginal constraints are relaxed such that Q E IT or E m while maintaining the constraint that E Ilgo.gn: well the case of non-square T.
3.2 The Balanced FRLC Algorithm
We introduce an algorithm Factor Relaxation with Latent Coupling (FRLC), to compute a LC factorization of minimum cost. We first describe the FRLC algorithm for the balanced Wasserstein problem. Extensions other and marginal constraints are discussed later. The FRLC objective function, for low-rank, balanced Wasserstein OT, is
LLC( R,T) = (C, (C.PIQ.T.)).
where P(Q.R.TY is defined by (7). Since (Ca.b(r) parameterizes (T,,(r), problem (8) is equivalent low rank problem (6) The FRLC algorithm is built from projections onto convex sets, described by constraints on the outer marginals alone for (Q, R) and by the inner marginals alone for T. Given (Q, R,T)) T) (r), sub-couplings Q and R constrained by:
(a) == (, R4 Q1, =a}, (b) ((Q, R,T) E R4 R1, = b).

---

The convex sets constraining the latent coupling matrix T are
C2(go) = (Q,, E R+ T1, = g2), C2(gR) = ((Q.R.T) E R4 : = =gR),
where 9Q = 1 and gn R 1m as per Definition onc has LCa.b(r) C1 C2.
3.1 Writing Ci = Ci(a) (C) and
C2 = (
We use coordinate mirror descent to optimize (8). building on the mirror descent (MD) approach of Scetbon al. (202T) Scetbon Cuturi (022);; Scetbon al. 2022 2023) for the low-rank problem. First we take descent step in the variables (Q, R) for a fixed T, using KL penalties on their inner marginals. These "soft" constraints allow the joint optimization in (Q. R) to decouple into two semi-relaxed OT problems, one for each variable. We call this step factor relaxation as this allows (Q, R) to have relaxed inner marginals go and gR- Next we take a descent step in the latent coupling variable T. fixing the Q and R. equivalent to solving a balanced OT problem. Thus, solving both coordinate descent steps corresponds to solving three OT problems.
We now provide further details on these coordinate descent steps, with the full algorithm given in Algorithm | Let (7*))== be sequence of step sizes. As in Scetbonn & Cuturi (2022), choose e normalization for the step-sizes. Our coordinate mirror descent in the factor relaxation step is:
(),, V (Q 1 KL((Q, R) Rx))
(Q++1, Rx+1)
arg min (Q.R) (.TTT))C
7KK
+ +TKL((Q 1,, 1m)[m R 1m))
The Sinkhorn kernels for the semi-relaxed OT problems arising from the factor relaxation step are:
K(K Qk C exp( 7h (CR, d,,i "(CC,,((C,""' "(CiR..""""ddiiil///e.... Qxdiag(L/gq.)))
KL K(6) R O exp( 7k ( Qx 1, diag diaad""(diag((" (diag( RRT C Qx
introducing the shorthand X diag (1/gQ) T diag (1/gn) and where (() T - R" denotes the matrix-to- us to show smoothness of the objective Proposition E.5 from which the convergence guarantee Proposition 3.3 follows. We derive the semi-relaxed projection Algorithm of the sub-couplings Q and in Appendix G for completeness. We also show Lemma A. that go and gR induced by the semi-relaxed projection are both feasible and locally optimal, not requiring separate optimization.
extraction of the diagonal. This r-dependent regularization also allows
As (Q++1, Rg+1, T) E C2 if and only if T Iggoues 99 after the factor relaxation step, we next take a coordinate MD step on the latent coupling T:
1 KL
Tat1 +
arg min T: Re T)EC2
(T. Vr +
(9)
This is equivalent applying Sinkhorn (Algorithm 5
to T given go and gn with the kerne::
T Tx @ exp(->>diag(1//)) (1g Q C CRs+1diag(1 gg..)))
After the final iteration of the coordinate-Mi scheme, X = diag (1/9Q) T diag (1/gR) satisfies Xgn = 1, and X 9Q 1, as T is a coupling between gQ and gn- Thus P, = QXR. E Ia. and the iterates (Q,, Tú. R2,) remain in the intersection of the constraint sets. Thus, in contrast to other approaches Scetbon et al. (202T): Lin et al. (202 I): Forrow et al. 2019). we do not require Dykstra projections back into the intersection to maintain feasability. We note that our implementation of FRLC allows for a non-s latent coupling T. providing greater interpretability in problem- specific applications. Above, presented FRLC in the simplest that T is square.
3.3 Initialization, convergence, and FRLC extensions
Full-rank random initializations of the sub-coupling matrices. We propose new initialization of the sub-couplings (Q, R. R)) for the LC-factorization in Algorithm 6 This generates full-rank initialization (Proposition F. in the set rank-r couplings lr. and is accomplished by applying Sinkhorn to random matrices. Our approach differs from Scetbon et al. ); (Scetoon & Cuturi (2022) who use initializations for the diagonal factorization of Forrow et al. (2019), and are not applicable to latent coupling that is non-diagonal, non-square, or with two distinct inner marginals.
5

---

Algorithm Balanced FRLC
Input C, ; a, T, 6,6
Initialize 9Q.9R = 1,
Qo. Ro: To <- SinttaalizeCouplinggen a, b, D, 9R)
Alg.
Xo + (iag(1/Q2 1 diag(1 mm)
while A (Q, Ry: Ti),, Qk-QQ-Q1 1- Rx Th
> do A in 10
Va= - CR& X 1,,iiag
Eaadiag-"(CCn) ((CR X )FQudiag() (1/g2))
Vn - CTQuX,
Imdiag (diag(1 (1/g)) R Qt xx))
7 2 max 1-mx//xa V Vall>>) # 8 normalization of Scetbon & Cuturi 2022
KQ (()) + Qk @ exp( -7k V(V R O exp( -7A -7 RR)
Qk + SR projectiion K
7k T, a, Q 1n, 6)
Semi-relaxed OT, Alg.
R2 + SR -poojection K R
7k, T, b, 1
1m.0) # Semi-relaxed OT
9Q. 9R + 1 R 1m
VT diag(L/gq) Q CR diag l/gR) -00 normmalliizaiinn
TT 7/1l V
K T + Ta @ exp(-PrVx)
Th + Sinkhorm K .R. 9Q.0) # Balanced OT, Alg.
X f diag(1/90 / diag(1/gn
end while
Return P. QXRT
Convergence analysis of FRLC. As objective (8) is non-convex, it important have convergence guarantees. Our convergence criterion A(,,) defined in (1O). To prove convergence we require lower bound on the entries of go and gn- Previous works introduce lower- bound vector a S g enforced element-wise for stability and smoothness (Scethon et al. (202T). In FRLC the use of semi-relaxed projections naturally enforces lower-hound. Appendix E.S] we show that for any guarantee uniform lower-bound of 6 on the entries: for sufficiently large т and (mm²/e) iterations for the sub-coupling Pham et al. (2020), one guarantees lower bound of d on In and JQ- This allows us to show objective smoothness in Proposition E.35 Previous work on low-rank optimal transport [Scetbon
5 E (0, )). the FRLC algorithm's r-weighted regularization the inner marginals
et al.a)222000 (202T) use the non-asymptotic convergence criterion of Ghadimi et al. (2014). Following existing works Dang & Lan 2015) establishing convergence rates of coordinate mirror-descent for smooth objectives, we show in Proposition 3.3 this criterion may be extended to coordinate-MD by adapting the block-descent lemma of Beck Tetruashvili (013)).
Lipschitz gradient and block smoothness constants (Li) and function h E C(X, R) which is a-strongly convex. For suppose one performs coordinate mirror descent on minimized over a product closed convex sets X -I X.. Let the sub-iterates with respect to the i-th block update be {xk]== where X == x for E [N] outer iterations. Then one has:
Proposition 3.3. Suppose one has f E CX, R) with block-coording
f =
DPL 2D
min MiinA(X,: A(Xx;(x- 5 N(Q22 No
where D is (6). L is the global smoothness constant, stepsizes This == a/ and convergence criterion A ( 1)) is given in (35).
Specialized to the LC- L.C-parametrization, the criterion (g 2b+11 is:
A (x Db+11) = 72 [[QQ+1 Qx1 | R&41 Rx | Txll
(10)
for æt = (Qu, R&, TE). We show through Propositions
3.3
E.5 the following result:
Proposition 3.4. The FRLC algorithm with step-sizes 7 0/ and iterates 2n. = (Qı, Rv. Te)
has non-asymptotic stationary convergence in the criterion A A(.-)) with:
min
Ax(ok, xx+1) < 2D2 1 22
kel kEI...N-- N-1

---

(a)
(b)
(c)
337.4
LoT poant
132
8.11
F&C Jubros
12T Tubii
230
Figure 2: (a) Simulated dataset containing points from two moons (orange) and eight Gaussians (blue). (b) Transport cost (C, P) achieved FRLC and LOT Scetbon al. (2021 for the balanced Wasserstein problem on the dataset in (a) for different ranks and initializations. FLRC full rank (blue curve) is average over 10 random initializations. (c) Results on the 10D mixture of Gaussians dataset.
Where is the number of iterations, D the optimality- gap as in (36). and L = maxie (3....3) (Li) the global smoothness for L, poly(I n, m, the block-wise smoothness constants.
The proof of Proposition 3.4 follows directly from extension of the non-asymptotic criterion with the block-descent result Proposition 3.3 and the proof that this lemma holds in FRLC Proposi- tion E.5 We also mention two improvements other low rank approximation results in literature. In Proposition
F.2 we show that one can analytically solve for the block-optimal for the factor- ization of Scetbon al. (2021), and we improve the bound the low-rank approximation error in Proposition
E.7
FRLC for other marginal constraints and objectives. The balanced FRLC algorithm can be extended simply to other marginal constraints owing to the decoupling of the coordinate MD scheme. In particular, by using either the semi-relaxed projections (Algorithm 2 or fully-relaxed (unbalanced) projections (Algorithm 3) on sub-couplings Q and R. one can solve the balanced problem, the problem with the left or right marginal relaxed, or the unbalanced problem. As such, all variants of marginal constraints can be handled by a single algorithm, given in Algorithm
We also extend the FRLC algorithm to the Gromov-Wasserster Wasserstein problem. This consists of computing GW-specific gradient with the appropriate marginal constraints applied to simplify their form, and re-computing Sinkhorn kernels as exponentiations of these gradients. The matrix form of the quadratic GW objective 1 p' A @ Pi,, + 1, 2 1"',, 2(APB, P), where O denotes the Hadamard (entrywise) product. Then the GW-specific Sinkhorn kernels are
(R) <- exp (2 BRX A A Q02Q1,1
(k) K
R exp (2BRX Q² AQX 0 (o)11)
k-lk) <- exp(4<x diag(go R.T ARiiii"
T
In Algorithm 4 one can solve the GW-problem by using the kernels above. Here, we present the kernels omitting rank-1 perturbation, which is given in Appendix D From the Wasserstein and GW gradients, the FGW gradient is easily taken convex combination of the two. In this work, primarily focus on the LC-factorization for the rank Wasserstein problem (6)
Experimental Results
We compare FRLC existing low-rank and full-rank optimal transport algorithms on several datasets:
simulated datasets previously used in Tong et al. (2023) and (Seebbon et al. (2021); a massive spatial- transcriptomics dataset Chen al. [2022); ; and graph partitioning task Chowdhury Needham (202T). Further details of each experiment (e.g. pre-processing, validation) are in Appendices K
and M the section below, LOT refers to the works of Scetbon et (221 2023 2022) and Latent
OT refers to Lin et al. 2021).

---

(a)
(b) FRLC [Rank
(c)
(d)
(e) LOT (Rank 10)
Ground Truch rank)
FRLC [Rank 10)
Figure 3: L.C-rojections of couplings of Gaussians centered on the -ooos of unity (green) and 10th roots of unity (yellow). (a) Ground- -round-truth full-rank coupling. (b) Non-square rank-5 catent-c of FRLC (c) LC-projection barycenters aligned with rank-5 diagonal coupling of LOT [Scetbon et al.) (2021) (d) Square rank- 10 latent coupling of FRLC. (e) Rank- 10 diagonal coupling of LOT.
4.1
Evaluation of Low-rank Approximations for Balanced OT on Synthetic Data
We first compare the balanced OT version of FRLC with the the low-rank balanced OT algorithm LOT of Scetbon et al. (202I) synthetic dataset following Tong et al. (2023). The dataset consists
of m 1000 points from two moons and n = 1000 points sampled from eight 2D Gaussian densities
(Fig. 2a). We solve the Wasserstein problem ① with cost matrix C' computed using the Euclidean distance. The full-rank coupling matrix P has rank 1000, and we compute both FRLC and LOT solutions with rank between 20 and 200. For each rank, we initialize FRLC adapting the deterministic rank-2 initialization proposed in Scetbon et al. (202I) and the random initialization of Alg. We initialize LOT using the rank-2 initialization and two other options in ott jax Cuturi et al. (2022)
We find that FRLC obtains lower transport cost (C, P)r with increasing rank (Fig. 2p) and consis- tently achieves lower transport cost than LOT across all ranks and all initializations. Specifically, starting both methods at the same rank-2 initialization, FRLC consistently achieves a lower cost than LOT for all ranks. Additionally, we observe smooth convergence of FRLC for both rank-2 initialization and the full-rank random initialization of Alg. (Tig. 5).
We also evaluate FRLC and LOT on two datasets of Gaussian mixtures, one in 2-dimensions and one in 10-dimensions, each with n = m = 5, 000 points from two mixtures of Gaussians, following Scetbon et al. 2021). with further details in Appendix K We observe the same trend as the previous simulation for both datasets (Fig. 2 Fig. 77 with FRIC achieving lower transport costs than LOT across all ranks and all initializations. In addition FRLC has half the runtime of LOT (CPU) including the setup time of FRLC but excluding the setup time of LOT in ott-jax on datasets of n = 1000 points from all three datasets with rank r = 100 (Table 2). At the same time FRLC achieves lower primal cost (C, P) with tighter marginals PI, all and ||
bll2- Lin et al. (202T) only solves proxy for the rank-constrained Wasserstein problem, and thus is not the focus our comparisons. Nevertheless, we verify that on all synthetic experiments that FRLC achieves significantly lower primal OT cost than Latent OT (Table
4.2 Interpretation of the Latent Coupling and LC-Projection
We demonstrate the intepretability of the latent coupling T in the LC factorization. In both the LC factorization and factored couplings, the sub-couplings Q and R each have associated barycentric projection operators which coarse-grain input datasets Z ( 7 In particular, the LC projection is defined from the LCC factorization as follows.
Definition 4.1 (LCC-rojection).. Let Q diag(1/ (1/9) T diag(1 (1/9n) RT be LC factorization of of a coupling matrix P (r), computed from datasets (1) E Raaxd Z E R with T E R xra The L.Cprojectionss (1) and (2) of (1)) and (2)) are (l) = (ling(/go)@TT
and ( == diag(1/gn) RT
By interpreting any factored coupling (Q, R. g) as an LC factorization (Q. R, diag(g)), Defini- tion 4.1 describes the barycentric projections for both factorizations, We compare the projections of the coupling computed by FRIC those of LOT |Seethon et al. (202T) on dataset containing 1000 samples from 2D- 2D-Gaussians centered at the 5th -roots of unity and 1000 samples from 2D Gaussians
8

---

Objective
Algorithm
Spearman
ARI
AMI
W
LOT-U
0.394
0.332
0.397
Tubb2b Expression
Tubo2b Predicted
1.0
FRLC-U
0.465
0.466
0.484
LOT-SR
0.391
0.319
0.396
0.8
FRLC-SR
0,467
0.475
0,492
GW
LOT-U
0.005
0.0
0.0
0.6
FRLC-U
0.266
0.299
0.364
LOT-SR 0.004
0.0
0.0
00.4
FRLC-SR
0.275
0.318
0,381
FGW
LOT-U
0.391
0.330
0.400
FRLC-U
0.368
0.391
0.420
LOT-SR
0.396
0.337
0.401
(a)
FRLC-SR
0,465
0.469
0.499
Figure 4: (a) Brain marker gene Tubb2b expression and FRLC prediction. (b) Comparison of the low-rank unbalanced (LOT-U) algorithm of |Seetbon et al. (2023) and FRLC on aligning spatial transcriptomics data. Bold indicates top performing method for each metric on each objective.
centered at the 10t-roots of unity (the latter scaled by a factor of two, Fig. 3a. In both cases, the latent coupling T or diag (g) is visualized as a transport between barycenters. We run FRLC and LOT with ranks r = 5 and = 10 to match the number of target and source clusters. In the rank-5 case, FRI.C uses a non-square latent coupling T E IR which correctly captures the coupling between clusters (Fig. 3)), while the LOT rank-5 projection computes barycenters that outside of the clusters (Fig. 3 A similar result is observed for square rank-10 latent couplings computed by FRLC (Fig. (d) and LOT (Fig. (ee demonstrating that the LOT barycenters in Fig. 3)) not an artifact of using the lowest rank. We observe similar results on other simulated datasets (Fig.
(Fig.3bb)
4.3 Evaluation on Spatial Transcriptomics Alignment
We compare FRLC and the algorithm (LOT-U) of (Seetbon et al. 2023) (which solves unbalanced low-rank Wasserstein, GW, and FGW problems) on the task of computing an alignment between cells from different time points during mouse embryonic development. Specifically, we compute an alignment between a spatial transcriptomics (ST) dataset of an E1 1.5 stage mouse embryo and an E12.5 stage mouse embryo Chen al. (2022). Optimal transport is popular approach to align single-cell Scchibinngrr et al. (2019) and spatial trancriptomics datasets Zeira et al. (2022); : Liu et al. (2023): Klein et al. (2023).
single-cell transcriptomics, one measures a gene expression vector for each cell, and in spatial transcriptomics one additionally measures the 2D location of each cell. The cost matrix C' describes the difference between gene expression vectors and intra-domain cost matrices A and B are derived from the 2D coordinates within each slice. Therefore, OT problems of W, GW, and FGW objectives can solved and the coupling matrix represents the cell-cell alignment (Appendix M). However, computation of a full-rank OT solution is not feasible in our large-scale dataset: the E11.5 slice has about 30,000 cells while the E12.5 slice has about 50,000 cells.
We evaluate the alignments by assessing performance on two prediction tasks from Scetbon et al. (2023): (1) a gene expression prediction task where we predict the expression of gene in E12.5 from expression of the gene in E11.5 using the alignment; (2) cell type prediction task where we predict the cell types of E12.5 from the cell type clustering of E1 (Appendix MM. We evaluate the accuracy of the gene expression prediction task through the Spearman correlation p between the predicted expression and the ground truth expression of 10 test marker genes. We evaluate the accuracy of the cell type prediction task by computing the Adjusted Rand Index (ARI) and Adjusted Mutual Information (AMI) between the predicted cell types and the cell types derived in the original publication Chen et al. (2022). Being comparison between different objectives, this relies on downstream metrics. For completeness, we validate the efficacy of FRLC on directly minimizing the balanced Wasserstein cost (C, P)F against Scetbon et al. 2021 in Figure 8
For direct comparison, we use FRLC to solve the same unbalanced problems (denoted FRLC-U). We perform an extensive grid search (Appendix M.3) to pick the best hyperparameters (including rank & 30, 000) for all algorithms. [Scethon et al. (2023) previously showed that unbalanced FGW algorithm has the best performance on ST alignment. We find that unbalanced FRLC achieves comparable or better results than the previous of unbalanced low-rank method on all three objectives (Table 4). We also solve a semi-relaxed version of each problem motivated by the

---

Sub-routine for coupling
Method
Factorization
Cost
Variables
Algorithm
Factored Coupling
k- Wasserstein barycenter
Anchors & sub cub-couplings
Factored coupling
Lloyd-type
Dykstra's
Forrow al. 2019)
Extension of k- Wasserstein barycenter
Latent OT
Anchors & sub-couplings
Latent coupling
Lloyd-type
Dykstra's
Lin et al.
2021
LOT
Sub-couplings & inner marginal
Factored coupling
Primal OT cost
Mirror-descent
Dykstra's
Scetbon et al. 2021
Coordinate
FRLC (this work)
Latent coupling
Primal OT cost
Sub-couplings
OT
mirror-descent
Table 1: Comparing aspects of low-rank OT methods. inner matrix.
Factorization indicates the structure of the
observation that all cells from E12.5 have ancestor, but not all cells from E11.5 have the same number of descendants due to cell growth and death. Thus the former marginal tight, and the latter relaxed Halmos et al. 22244.. We both semi-relaxed FRLC (FRLC-SR) and setting of LOT-U that recovers the semi-relaxed problem (LOT-SR). Semi-relaxed FRLC achieves the best results on all three metrics by large margin (Table 4) As one example, the expression of Tubb2b, mouse brain marker gene, agreeing with the expression predicted from the semi-relaxed alignment of FRLC (Fig. 4ap.
4.4 Additional Experiments
We evaluate FRLC unsupervised graph partitioning problem Chowdhury & Needham (2021 on four real-world graph datasets Yang & Leskoved (2012); Yin et al. 20 Banerjee et al. (2013). We
benchmark the performance of the semi-relaxed and GW settings of FRLC against (1) G GWL et al.)
(201)), solving a balanced GW problem; (2) SpecGWL/(Chodhhne (howdhurr
Needham (02T) using the
heat kernel on the graph Laplacian the cost matrix. We find FRLC achieves the better clustering performance than GWL and SpecGWL on 9/12 and 11/12 of the datasets (Table 3 and Appendix L.
5 Discussion
We provide comparison of existing low-rank solvers in Table D The FRLC algorithm has number of advantages, including (1) coarsening a full-rank plan P non-diagonal latent coupling T; (2)
minimizing the primal OT problem for general cost C rather than barycenteric problem; (3)
optimizing only couplign,, and (4) using Sinkhorn alone as the sub-routine for low-rank OT.
While argue these are substantial advantages, FRLC has limitations which warrant follow-up work. In particular, three key limitations of our work, common to the existing low-rank OT algorithms, are: (1) selecting values of the latent coupling ranks; (2) strengthening the convergence criterion; (3) addressing sensitivity to the initialization from non-convexity of the objective. A limitation specific to work is the selection of the hyperparameter controlling the smoothness of the trajectory. These and other limitations are discussed in Section N of the Appendix. Another direction for further investigation is better understand what structure LC factorizations capture when the optimal plan known to have full rank, e.g. when the Monge map exists, as has been explored by Liu et al. 2021
6 Conclusion
We introduce FRLC, algorithm to compute low-rank optimal transport plan from the latent coupling (LC) factorization. FRLC handles different OT objective costs and relaxations of the marginal constraints. Moreover, the LCC factorization provides interpretable coarse-graining of the full transport plan and its marginals through the mapping (P, a. b) - (T, 9Q IR). We demonstrate the superior performance of FRLC compared state-of-the-art low-rank methods on real and synthetic datasets.
10

---

Acknowledgments and Disclosure of Funding
This work is supported by NCI grant U24CA248453 to B.I.R. J.G. gratefully acknowledges support from the Schmidt DataX Fund at Princeton University made possible through a major gift from the Schmidt Futures Foundation.
References
Bakshi, A. and Woodruff, D. Sublinear Time Low-Rank Approximation of Distance Matrices.
Advances in Neural Information Processing Systems, 31, 2018.
Banerjee, A., Chandrasekhar, A. G., Duflo, E., and Jackson, M. O. The Diffusion of Microfinance. Science, 341(6144):1236498 2013.
Bauschke, H. H. and Lewis, S. Dykstras algorithm with Bregman projections: A convergence proof.
Optimization, 48(4):409 427, January 2000. ISSN 1029-4945. doi: 10.10800223311900084444111 URL htttp:/ddx.dii.orr org/10..0.11000 wtt.:///....00000339900000030330
Beck, A. and Tetruashvili, L. On the Convergence of Block Coordinate Descent Type Methods.
SIAM J. Optim., 23:2037-2060, 2013. URL https://api semanticscholar, orw.oo////oooo CorpusID
6866704
Benamou, J.-D., Carlier, G., Cuturi, M., Nenna, L., and Peyré, G. Iterative Bregman Projec- tions for Regularized Transportation Problems. SIAM Journal on Scientific Computing, 37 11:11111-A1138 January 2015. ISSN 1095-7197. doi: 10.1137/141000439. URL http
/ ddxxdoi. org/ 10. 10111113111000 /14100443
Bregman, L. M. The relaxation method of finding the common point of convex sets and its application to the solution of problems in convex programming. USSR computational mathematics and mathematical physics, 7(3):200-21,, 1967.
Bunne, C., Stark, S. G., Gut, G., del Castillo, J. S., Levesque, M., Lehmann, K.-V., Pelkmans, L., Krause, A., and Rätsch, G. Learning single-cell perturbation responses using neural optimal transport. Nature Methods, 20(11):155--1760 September 2023. ISSN 1548-7105, doi: 10.1038/
251592-023-01 0293-09999x.. URL/tttp://as doi org/ 10. 1038/ /1w.1///////11--11 023 023-0999------9999
Charikar, M., Chen, B., Ré, C., and Waingarten, F. Fast Algorithms for a New Relaxation of Optimal Transport. In Neu, G. and Rosasco, L. ), Proceedings of Thirty Sixth Conference on Learning Theory, volume 195 of Proceedings of Machine Learning Research, pp. 4831-4862. PMLR, 12-15 sitess/v100 v195/ http//hariiaa22. html
Jul 2023. URL httpth::s:p/cceeddiiiiic
mlr. prw.s
Chen, A., Liao, S., Cheng, M., Ma, K., Wu, L., Lai, Y., Qiu, X., Yang, J., Xu, J., Hao, S., et al. Spatiotemporal transcriptomiic atlas of mouse organogenesis using DNA nanoball-patternel arrays. Cell, 187-1197 2022.
Chen, X. and Price, E. Condition number- free query and active learning of linear families. 2017.
Chizat, L., Peyré, G., Schmitzer, B., and Vialard, F.-X. Unbalanced Optimal Transport: Dynamic and Kantorovich Formulations. Journal of Functional Analysis, 274 (74(1))309003123 June 2018. ISSN 0022- 1236. doi: jfa.66/..faa2018.00 03.008..7 URL http: doi org. 101.00...6jj........ 2018
03.008
Chowdhury, S. and Needham, T. Generalized Spectral Clustering via Kromov-Wasserste Learning. In International Conference on Artificial Intelligence and Statistics, pp. 7 712-720. PMLR, 2021.
Chung, F. Laplacians and the Cheeger Inequality for Directed Graphs. Annals Combinatorics, 1-19, 2005.
Cohen, J. E. and Rothblum, U. G. Nonnegative Ranks, Decompositions, and Factorizations of Nonnegative Matrices. Linear Algebra and its Applications, 190:149-163. 1993.
Courty, N., Flamary, R., and Tuia, D. Domain adaptation with regularized optimal transport. In Machine Learning and Knowledge Discovery Databases: European Conference, RCML.PKD) 2014, Nancy, France, September 15-19, 2014. Proceedings, Part I 14, pp. 279-289, Springer, 2014.
11

---

Cuturi, M. Sinkhorn Distances: Lightspeed Computation of Optimal Transport. Advances in neural information processing systems, 26, 2013a.
Cuturi, M. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in Neural Information Processing Systems, pp. 2292-2300, 2013b. U neurips html
Cuturi, M.,
L., Tian, Y., Bunne, C., Davis, G., and Teboul, O. Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein. arXiv preprint vrXi::2201.12324. 2022.
Dang, C. D. and Lan, G. Stochastic Block Mirror Descent Methods for Nonsmooth and Stochastic Optimization. SIAM J. Optim,, 25(2):856-8811 January 2015.
Dong, S., Pan, Z., Fu, Y., Xu, D., Shi, K., Yang, Q., Shi, Y., and Zhuo, C. Partial Unbalanced Feature Transport for Cross-Modality Cardiac Image Segmentation. IEEE Transactions on Medical Imaging, 2023.
Dykstra, R. L. An Algorithm for Restricted Least Squares Regression. Journal of the American Statistical Association, 78(384):837 842. December 1983. ISSN 1537-274X. doi: 10.1080/ 0162 (6691459.199 198.
URL http: //dxx doi. org. 10. 1080/01621459 1983. 10477029
Forrow, A., Hütter, J.-C., Nitzan, M., Rigollet, P., Schiebinger, G., and Weed, J. Statistical Optimal Transport via Factored Couplings. In Chaudhuri, and Sugiyama, M. (eds.). Proceedings of the Twenty-Sccond International Conference Artificial Intelligence and Statistics, volume 89 of Proceedings of Machine Learning Research, pp. 2454-2465, PMLR, 16-18 Apr 2019. URL https://proceediee
mlr. pwwps/v89/ ha09/fforrww99 html
Frieze, A., Kannan, R., and Vempala, S. Fast Monte-Carlo Algorithms for Finding Low-rank Approximations. ACM, 51(6): 1045-1041, nov 2004. ISSN 0004-5411. doi: 10.1145/1033488. 1039494. URL https: : org/ 1145/1039488 1039494
Frogner, C., Zhang, C., Mobahi, H., Araya, M., and Poggio, T. A. Learning with Wasserstein Loss. Advances in neural information processing systems, 28, 2015.
Geshkovski, B., Letrouit, C., Polyanskiy, Y., and Rigollet, P. A mathematical perspective on Transformers. arXiv preprint arXiv:2312.1079: 2023.
Ghadimi, S., Lan, G., and Zhang, H. Mini-batch stochastic approximation methods for nonconvex stochastic composite optimization. Mathematical Programming, 155(1-2):267-305 December 2014. ISSN 1436-4646. doi: 10.1007/510107-0144-08888 1007/s10107-014-0644
URL http://dx doi doi..rg/110
Halmos, P., Liu, X., Gold, J., Chen, F., Ding, L., and Raphael, B. J. DeST-OT: Alignment of Spatiotemporal Transcriptomics Data. International Conference on Reseurch in Computational Molecular Biology, pp. 434-437. Springer, 2024.
Indyk, P., Vakilian, A., Wagner, T., and Woodruff, D. P. Sample-optimal low-rank approxima- tion of distance matrices. In Beygelzimer, A. and Hsu, D. (eds.), Proceedings of the Thirty- Second Conference on Learning Theory, volume 99 of Proceedings of Machine Learning Re- search, pp. 17511 PMLR, 25-28 Jun 2019. URL petps::/proceediee v99 2999indyy19a.. hhmll
mlr, ppress/
Kantorovich, L. On the Translocation of Masses: Doklady akademii nauk ussr. 1942.
Klein, D., Palla, G., Lange, M., Klein, M., Piran, Z., Gander, M., Bastidas-Ponce, A., Tarquis-Medina, M., et al. Mapping cells through time and space with moscot. bioRxiv, pp. 2023-05, 2023.
L., Sterr, M.,
Lin, C.-H., Azabou, M., and Dyer, E. L. Making transport more robust and interpretable by moving data through a small number of anchor points. Proceedings of machine learning research, 139: 6631, 2021.
Liu, W., Zhang, C., Zheng, N., and Qian, H. Approximating optimal transport via low-rank and sparse factorization. CoRR, abs/2 1.10211611666666 2021. URL httts:://axxiiii http://aaxxvv.../aabb/a/ orggabs/2111 http///111.616666
12

---

Liu, X., Zeira, R., and Raphael, B. J. Partial alignment of multislice spatially resolved transcriptomics data. Genome Research, 33(7):1 2023.
Mémoli, F. On the use of Gromov-Hausdorff Distances for Shape Comparison. 2007.
Mémoli, F. Gromov-Wassersee Distances and the Metrie Approach to Object Matching. Founda-
tions of computational mathematics, 11:47-4877, 2011.
Nesterov, Y. Efficiency of Coordinate Descent Methods on Huge-Scale Optimization Problems. SIAM Journal on Optimization, 22(2):341-362 2012.
Orlin, J. B. A polynomial time primal network simplex algorithm for minimum cost flows. Mathe- matical Programmiing, 78(2):199-129, Aug 1997. ISSN 1436-4646. doi: 10.1007/BF2264466 URL https://link. spr springer cww www.o//////////ooooooooooo 10. 1007/FF02614365 pdf
Pham, K., Le. K., Ho, N., Pham, T., and Bui, H. H. On Unbalanced Optimal Transport: An Analysis of Sinkhorn Algorithm. In International Conference on Machine Learning, 2020. URL
semanticscholas. https://apl semanti cscholar org/ /ttg//Crrus... : 21 1068892
Sander, M. E., Ablin, P., Blondel, M., and Peyré, G. Sinkformers: Transformers with Doubly Stochastic Attention. In Camps- Calps-Valls, G., Ruiz, F. J. R., and Valera, I. (eds.), Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, volume 151 of Proceedings of Machine Learning Research, pp. attps://procceednng: mlr wwwpp/ www.5///aande22a...mm html
PMLR, 28-30 Mar 2022. URL
Scetbon, M. and Cuturi, M. Low-rank Optimal Transport: Approximation, Statistics and Debiasing. In Oh, A. H., Agarwal, A., Belgrave. D., and Cho, K. (eds.), Advances in Neural Information Processing Systems, 2022. URL
Scetbon, M., Cuturi, M., and Peyré, G. Low-Rank Sinkhorn Factorization. In International Con- ference on Machine Learning, 2021. Utt.ttts:://apl http://pp...smmniiinccccccoo.rr. org/ Corpusri 232147563
Scethon, M., Peyré, G., and Cuturi, M. Linear-time Gromov Wasserstein Distances using Low Rank Couplings and Costs. In International Conference on Machine Learning, pp. 19347-19365. PMLR, 2022.
Scetbon, M., Klein, M., Palla, G., and Cuturi, M. Unbalanced Low-rank Optimal Transport Solvers, 2023.
Schiebinger, G., Shu, J., Tabaka, M., Cleary, B., Subramanian, V., Solomon, A., Gould, J., Liu, S., Lin, S.. Berube, P., et al. Opaimal-Transport Analysis of Single- Cenl Gene Expression Identifies Developmental Trajectories in Reprogramming. Cell, 176(4):928-9))) 943, 2019.
Solomon, J., De Goes, F., Peyré, G., Cuturi, M., Butscher, A., Nguyen, A., Du, T., and Guibas, Convolutional Wasserstein Distances: Efficient Optimal Transportation on Geometric Domains. ACM Transactions on Graphics (ToG), 34(4):1-11, 2015.
Ståhl, P. L., Salmén, F., Vickovic, S., Lundmark, A., Navarro, J. F., Magnusson, J., Giacomello, S., Asp, M., Westholm, J. O., Huss, M., et al. Visualization and analysis of gene expression in tissue sections by spatial transcriptomics. Science, 353(6294)):7 82, 2016.
Tarjan, R. E. Dynamic trees search trees via Euler tours, applied to the network sim- plex algorithm. Mathematical Programming, 78(2):169-17) Aug 1997. ISSN 1436-4646. doi: 10. 1007/BF026143699 URL https////innnk...iggerrrr wwweeeer.oooooooooaeeeoeee http://oneettt/////c/ 1007//
BFO2614369.10
Tay, Y., Bahri, D., Yang, L., Metzler, D., and Juan, D.-C. Sparse Sinkhorn Attention. In III, H. D. and Singh. A. (eds.), Proceedings of the 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 9438-9447, PMLR, 13-18 Jul 2020. URL https://proceedings mlr http://iwwnnn.....e/eeseeeaea..e wttpsss///1//taaaaa.aaamm
Tong, A., Malkin, N., Huguet, G., Zhang, Y., Rector-Brooks, J., Fatras, K., Wolf, G., and Bengio, Y. Improving and generalizing flow-based generative models with minibatch optimal transport. In ICML Workshop on New Frontiers in Learning, Control, and Dynamical Systems, 2023.
13

---

Vayer, T., Chapel, L., Flamary, R., Tavenard, R., and Courty, N. Fused Gromov- Wasserstein distance for structured objects. Algorithms, 13(9):212, August 2020. ISSN 1999-4893. doi: (43390/a13090211 URI http://dx doi orog 10. 3390/a13090212
Vincent-Cuaz, C., Flamary, R., Coreeli, M., Vayer, T., and Courty, N. Semi-relaxed Gromov- Wasserstein divergence and applications on graphs. In International Conference on Learning Representations, 2022. URL
Wolf, F. A., Angerer, P., and Theis, F. J. SCANPY: large-scale single-cell gene expression data
analysis. Genome biology, 19:1-5, 2018.
Xu, H., Luo, D., Zha, H., and Duke, L. C. Gromov-Wassee Wasserstein Learning for Graph Matching and Node Embedding. In International conference on machine learning, pp. 6932-6941. PMLR, 2019.
Yang, and Leskovec, Defining and Evaluating Network Communities based on Ground-truth In Proceedings of the ACM SIGKDD Workshop on Mining Data Semantics, pp. 1-8, 2012.
Yang, K. D., Damodaran, K., Venkatachalapatty, S., Soylemezoglu. A. C., Shivashankar, G., and Uhler, C. Predicting cell lincages using autoencoders and optimal transport. PLoS computational biology, 16(4):e 1007828. 2020.
Yin, H., Benson, A. R., Leskovec, and Gleich, D. F. Local Higher-Order Graph Clustering. In Proceedings of the 23rd ACM SIGKDD international conference on knowledge discovery and data mining, pp. 555-564, 2017.
Zeira, R., Land, M., Strzalkowski, A., and Raphael, B. Alignment and integration of spatial transcriptomics data. Nature Methods, 19(5):567-577 2022.
14

---

A Low-rank optimal transport
A.1 Low-rank factorizations
The set of low-rank couplings. Given M E R4*", the nonnegative rank of M is the least number of nonnegative, rank-l matrices that sum to M:
r (M)
min
M M,, r21
M,, such that rk(M,) 1 and M, > 0 for all i=1
Let a E A. b E Am be probability vectors, and let Ir.)()) denote the set of rank-r coupling matrices with marginals a and b:
(r), E 1'1m a, Pl, = rk4 (P) r).
To optimize any cost over II, (r),(r), one requires parameterization of this set.
Factored couplings. The factored coupling parameterization of (rr,(()) introduced in Forrow et al. (0119) and used by |Scetbon et al. 2021); Scetbon & Cuturi (022); Scetbon et al. (202 2023) is
(r), = (QQ, R, E R Rmxrr (R1) Q E Il..g:
lbb.).
Cohen Rothblum 1993 show that any Pe (r), may be decomposed as P = Qdiag( R
for some triple (Q, g E FC. Thus, for cost matrix C R=xm the general low-rank optimal transport problem is equivalent to an optimization over factored couplings:
min (ella.b(r)
(C,P)
min (Q (G....))FCaas(r)
(C, Qdiag (odiag(L/g)RR)
(11)
Latent coupling factorization. The latent coupling parameterization of lr))(r)) introduced in Lin et al. (2021). and used in the present work is
LCa,b(r) = ((Q, R. T) E TRXX x TRxX x BET" Q Ia.., R lb. II gg..gn):
where DQ: gR are the inner marginals of Q and R.
Latent coupling diagonalization. The LC-factorization recovers the factorization of Forrow et al. (2119) as sub-case. While the diagonal factorization of previous works cannot be directly converted the LC-factorization, the LC-factorization can easily recover the diagonal factorization. In particular, taking Q' + Q diaag(L/g))) one can refactor
P. Q (diia(1//00 99 T diag(1/gR) R Q' (i
or alternatively taking R' [idiagg///g) may refactor
P. Q (tiag(1/gg) T diag(1/gR)
diag(1/g))
So that instead of returning (Q, R, T) one may alternatively return (Q, R,T) (Q', R, diag(gn)) (Q, R. > (Q, R', diag(gq) ) to recover the Forrow et al. (2019) factorization. An example of this diagonal-cooversinn is offered in Figure 12
A.2 Balanced low-rank optimal transport
The FRLC optimization problem Our optimization problem is over the variables (Q, R, T) and defined as follows:
min (Q (.. R,),Q,R,T),
(12)
where our objective function LLC is
L,c(Q,R,T) = (C, Q(diig(l/"" D 1 RRT)))
(13)
Given (Q, (, (r), sub-couplings Q and are constrained by:
(a) == (QQ. R. R4 Q1, = =a} (b) ((Q, R,T) E R4 R1, = bb]
15

---

while the convex sets constraining the latent coupling matrix T are
C2(go) = ((Q) R) R+ TI, = 9Q). C2(gn) == ((Q, E R4 : gg],
where gQ Qu1n and gR RR11, as per Definition 3.1 Under these definitions, LCn,b(r) C1 n C2, where
C = Ci (Ci(b) and C2 -2 (2(g))
(14)
To solve (2), we separate the variables into "blocks" of variables, (Q, R) and T, and perform two block updates per iteration, as follows. Let (7k)e20 be positive sequence of stepsizes. Suppose we have (Qu, Pk: Tx) C. We update the first variable block (Q, RR by taking a locally optimal (mirror descent) update step, while the second variable block T is held fixed:
1 (LL Rx)) (15) 7k
Q Qu+1, , + arg min (Q..):Q....T.EE.
(Q, R), V (LC, +
Here, we slightly abused the notation by putting (Q, R) inside an inner product. The triple (Q++1, Rx41, Tx) produced by this update lies in Ct. Next, we update T, the second variable block, by taking another locally optimal (mirror descent) step, while the first variable block is held fixed.
(T, V + 1 -L(TTTTT 7k
T f arg min
(16)
T)C)22
By construction, the triple (Qk+1, RR+1, Th+1) C2. However, because the set C1 only con- strains the first variable block Qk+1: RE+1, we have that (Q++1, RK+1, Tk+1) E C2, and hence (Q++1, Rg+1, Tk+1) C. Thus, each iteration produces feasible triple (Qk+1, Re41,Te41 through a pair of locally optimal block updates.
The LOT optimization problem Scetbon et al. 2021 For comparison, recall the optimization problem that is solved in the LOT framework:
min (Q.
LLOT
(17)
where the objective function CLTT is
LOrT (C, (ddia((/))RR))
(18)
and where č E1 n n with:
C ((Q) R, E X REXT x (R1)* Q1, a, R1, b}
C₂ = (QQ, R.g) E R x RT x R QT -
Here, there are also three optimization variables (Q, R. g). but they are updated together in each iteration of LOT. That is, given feasible (Qu, Re, gk) C, iteration of LOT updates this triple via
1 KL((Q, R. g. Rs. gk)),
(Qu+1, Rt-1. 9k-1) == argmin (Q, V Viq.ng) (Q.R.))eč
7ik
where (7R)и20 is again positive sequence of stepsizes. Scetbon al. 2021 then compute the unconstrained argmin across all variables to yield set of unconstrained kernels (K KR. kg),
using Dykstra to jointly project the unconstrained update onto the intersection Č of the constraint sets.
OT subroutine in FRLC To see why do need Dykstra in the FRLC scheme, observe that the update (15) of variables (Q, R) can be equivalently expressed as
(Qu+1. 16-11R R
arg min Q.R. )CC
V(..) + 7kk
RI( Rt)),
Thus, even though the pair (Q, R) is being updated at this step, solving for (Qk+1. R above is equivalent to updating each individually because Q and do not share an inner marginal:
Q Qk+1 + arg min (Q, VoLLo) + (
(19)
Q:Q1, Q1
16

---

R&+1 t arg min (R. V VnLEc)+ + KL(R] (L(RR)))) R. Rb
(20)
We choose to add the regularization 7 ( In, R² 1 I,, R Im) to turn each update
here into entropy-regularizzd semi-relaxed optimal transport problem, and ensure 3-smoothness:
1 - =L(Q) + I,
Qu+1 f arg min (Q, VQL(c)+ + Q:Q1,
(21)
7k
R&+1 < arg min RR. V Vnfic) + RA) TTT.(RR 1m)
(22)
R: R1.=b
Yk
After updating Q and R. the update on T then follows similar form:
(T, VTL(c) + 1 -LT(T)TK)
Te+1
arg min
(23)
T: (Q Re T)C,
7
leading to balanced constraints on T of the form T Qh 1n and T1, R Im allowing the problem to be solved by Sinkhorn.
Importantly, there
constraints in C1 involving both Q and R, which is what allows the optimization to split in this way. If there were such constraints, we would have needed to use Dykstra to update the pair Q_ R). Because our update scheme is equivalent to the three updates of individual
variables given in 21 22 23 we can solve for each update using optimal transport. As Q and R required match the inner marginals exactly, the OT problems associated Q and R are semi-relaxed by construction.
The separation of our block updates into a step where RQ, Tx) E C1 and (Qx+1. Rk+1, T) C2 allows us to entirely remove the optimization over inner marginals go and gR. as done all previous works on low-rank optimal transport which optimize g explicilly as both variable and constraint of the optimization Scetbon et al. (2021, 2023); Scetbon & Cuturi (022). one were to introduce an extended loss in the style of previous works which adds go and gn variables in the form H(Q, R, T, 9Q. 9R) (Qdiag(1) (l/g)) ( (Tdig (1/g)) C) observes equivalence to simply taking a semi-relaxed projection.
Lemma A.1. Define the function H(Q, R. T.q. 9R) (C, (ediaa((/gg)Tdiinnne (diag(())))/))))))) and let (Q, R. T) be in 13). One has the following equivalence:
min gnEA, a4 E.. QElla.g Rellb.g
H(Q. R, Tk, 9Q: gR)
min (Q.R.T))cC
L,,(Q, (24)
Thus the semi-relaxed projections yield locally optimal inner marginals.
Proof. To see why this is true, notice that, so long as the outer marginals are tightly satisfied Q1, a and R1, b for Q 2 Oaxry R
2 O xr. has
Σ gR. Σ R 1m)
Σ
R Im) ΣΣ R Σb:
and
Ega, 9Q. Σ In) Σ QQ... 1, 1m) ΣΣ Ca==1
Therefore, the inner marginals 9 = (Q") and g TR")TI,, induced by the optimal Q", R of the optimization problem on the right hand side satisfy the constraints g E Ar,g E A, on the left hand side, so the two minimums coincide.
This implies that an extra optimization for gQ and gR is unnecessary in a coordinate update which alternates (Q, R) and T. If we did not perform block-update, in the form of standard MD in the objective described above, we would encounter some difficulty. In particular go and gR would optimized with the constraint 9Q A, and gR E Ar, and would concurrently constrain all of the
other variables TTI, 9Q. T1, 9Q. and R 1m gR. T 1, 9R.
17

---

B Block-Coordinate steps for the OT sub-problems
We use latent non-diagonal coupling instead of an inner diagonal coupling diag(g) of the form Forrow et al. (2019). This allows us to loosen the constraint that the inner marginals have to be joined by common coupling QT = R 1. = 9. The fundamental advantage of this choice is that we can decouple the convex-optimization problem for (Q, R. T') entirely. One can simply solve the optimal Q and R independently, yield the associated inner marginals for each Q 1n - !Q and R mm gn. and then find the optimal T which links the two. This link provided by the aforementioned form of the problem, where:
P= QXR
For Q, R. in either the appropriate set of couplings relaxation thereof (which we will describe shortly). X related to by:
X = diag(1/gq) diag(1/gR)
And, T E Il go.gn consistently for all cases. As the semi-relaxed case is intermediate between fully-relaxed and balanced, has ideas which generalize to both directly. As such, we use it as the leading example again. As in |Seetoon et al. (202I). take proximal-ster -oxpps of the form:
min(V Iex. F) + KL(SI| (+))
7/k
Where these steps are now wise fashion on (Q, R) and T, rather than joint. One may identify for each sub-factor in (Q, R) and T linearized gradient as before, which yields set of objectives which each solve an independent pptimpl-transport for the sub-factors. particular, we have that:
.Q (QXRT, CF - R)
(QXRT, C)F = Tr =
(QRRT, C)F=IT+ R (Q
A lincarization in the left-slot of the inner product as (Q, CQ = (Q, CRX") or (CTQX QX (Rx), R == (CT Qx, R is common practice for quadratic problems. In this case, the directional derivative of Q and R in the matrix dirriiion V are respectively:
D(CRXT F (V) (CRXT, V)F D(C R) (V) (CTQX
CRXT CRX
C
Without this lincarization assumption X, the full gradient may be evaluated as well. In particular, we note that for (-) the matrix-to-vector extraction of the diagonal, the directional derivative on Q is:
(CRX", V)F (C. QDXT (V)R*)) = =CRX" ((Ciia ([(xRX")(Qtigg((eee)) Qdiag(L//g))
V)
Thus, without the linearization assumption one may use product rule on X as an implicit function of Q (resp. R) to take the total derivative:
VoLFRIC CRX ((CRXX) Qdiag(1//))))
Likewise for R.
D(C, (CC
PoV (C.
=C (di V)r.
and so
V -
= - d,ag (dig
Both the W and GW problem have these rank-one perturbations of the gradient from the derivative with respect to X.
18

---

Lastly, owing to the derivative on T follows directly by chain rule on X:
updates which fix (Q, R) preceding the update on T. the
D(Q' (V) = (Q' CR. = VxL=
CRRRRRCCRRR Vz = diag(1/gq Q C didiggg//
Let (7ks be sequence of step sizes and consider the first-order conditions required for the proximal step as before. From these we have the updated proximal-step updates:
1 (k) + min (CRX Q) + KL(Qu K (k) Q
K
7
KJ (b) min (C T'QXX R) F
1
KL.
R K (k)
R
7
K T
min (Q X)
(T
K )
7
with kernels K. & for j 1,2.3 given by
K (k) Qs exp( C
K (k) R4 exp( 7/ C xx
K( = Ta exp( -n I C Rd
Or, dropping the linearization assumption (Q, Rt) the updates are: K for j 1, given by
K( (k) = Qk @ exp( N d,di '(CRR Q d (1/9Q.))
K K(()) (k) R : R2 exp( (C" 1, diag (diagg"(iiag1//nnn (x))))))))
The first projection is onto the set that satisfies the marginal constraint R1, b. In particular, following the discussion above, one has the coordinate-MD step:
min (Q.RTT)
KL R, (KQ Kn. Kr)) + (KI((1,
7 R1,
(25)
s.t.
As before, there is difference from the previous case, where takes the unconstrained projection R = diag(b) K 1,,, Kn. To generalize this, we also consider adding a soft-constraint on the inner marginal of R to be near that of the previous iteration. In partieular, consider the problem:
min (Q.R.T)
KL((Q, KQ. KT))
7
TKI + -K 1mm 1n (K (k-1)= 1m))
(26)
s.t. R1,
Which yields the relaxed solution of = SR prrjecction KR. 7h. 9
(k-1)) which generalizes
the original projection and is equivalent to it for 0. This regularization is essential, as it ensures 3-smoothness objective. For Q, as the constraint on g is fully relaxed, Lagrange multiplier A 0 entirely, such that the problem 40 now becomes uulllyynncosttrinedd
1 74
inf Q
KL(Q] + (L
(27)
To generalize this solution, again consider adding soft-regularization the inner marginal of Q, where we consider the alternate problem:
L KL Ko) + +KL(QL)))a) + TKL(QT1,,ll)) (1-1) QF-11,) 7k
inf
(28)
Q
19

---

Which trivially recovers the original for r = 0. This form has a solution given by an unbalanced optimal transport with kernel Q-. We see these two, convex problems in sequence give independent optimal solutions for Q and R the two matrices are not required to share an inner marginal. The last step to link them via T, corresponding the projection of (Qx; Rg. T) onto the set of valid rank-r couplings min Ru. L Rg, T) == ge 9a)) (Lc(Qi: k, R2. T). then P E la b. As such, one does not require any We verify in C that T E 11 11 "tti===01...=RT"""" projection onto the intersection of convex sets, done in Scetbon et al. (202] 2023) via the Dykstra projection algorithm Dykstra (188)) Alternating a coordinate-Mo one not only minimizes the objective in an alternating fashion but remains in the feasible set without the need for projection algorithms beyond Sinkhorn. As such, the final linking step is done via:
step in Q Rj and a step on T,
min T
[T](T] KT)
(29)
8.t.
TI, = -q: T 1, = 9R
This formulation amounts to solving a balanced Sinkhorn problem on T with respect to the proximal step kernel matrix.
Algorithm 2 SR R projection
(semi-relaxed OT, right marginal relaxed)
Input K 7; T, a. 5
< 1,
v +- 1,
repeat
ú +- u
+ (v/
+- (b/KTu) 1/(r+r
until 7-1 max |loo ū/ull 0/0ll// < return diag(u) diag(v) v)
Algorithm 3
U-projection (unbalanced OT)
Input K, 7, T, a. 6,8 u + 1 v + 1, repeat
+- u +- (a/Kv) (//KK)//(tt>>]
(b/KTu) +/(r+r11
until 7 max |log u oor 0/01)/ < return diag(u) u) ) diag(v)
C The Latent Coupling Matrix
To solve the balanced form (and generalize the principle to the relaxed problems), we consider an alternative parametrization of the inner matrix. In particular, previous works Scetbon al. (202T 2023) consider 1iag(1/9) to be the inner matrix, with marginals Q11, = RT i, to ensure that the outer conditions Q1, a and RI,- b hold. We instead relax the constraint that Q 1, and RI'1,, are equal, by allowing QT 1,, go and R R' 9R to vary arbitrarily. and considering non-diagonal inner matrix X gr in the place of diag(1/g) where we have the conditions:
Xga X 1
1,
And
X 9Q XQ 1, 1,
20

---

Thus, if one considers the coupling formed as P. QXR". one maintains that P. E Ia, from
the condition of C, defined in the balanced case in (14)
We clearly have that:
P,1m - QXR 1m Q1, =a
And:
1 = R R1,
We first consider two approaches to optimizing for such X, given that it not a coupling. First, we consider the appropriate proximal step for X
min( VL((S) lex F)F + - (K((S]I 7
We first note that the gradient of our loss with respect to X is given as Q If one supposes that X invertible with --1 we have that for any such
X- TI, gn
And
T - 9Q
This implies that this inverse matrix T is a coupling such that T ll.a.go which also suggests one might be able to update using Sinkhorn. In fact, being density R2"x" it represents transition matrix between the latent r-dimensional variables. In particular, writing the proximal step in full, we have:
min QQ' CR,T FF 1 K (1))
7k
Noting the derivative D(X 1) 1V -- we have from the first-order condition that:
1
T
T QRTTT--
log
0
7k
K
(b)
This implies the kernel matrix update:
K, Ta @
Where one then takes the Sinkhorn projection onto the set N9...9 as Pl (k()) using the Sinkhorn algorithm Cuturi (2013b). However, more stable and inversion- update exists which ensures X remains positive by diagonal re-scaling in the form introduced by Lin al. (02TT. In particular, if one takes X diag(1/ (aleg(/go)/ (T tiag(1//R), then
XgR
diag(1 / (gg)gg diag(1/go) TI, 1,
and likewise
X XTgo
gQ diag(iigg))(() diag( (l/9)/q diag(I/gg](
1,
that X necessarily satisfies XgR 1, and X gQ 1,, Thus T1, 96 and T gR and T E II go-gx- With analogous reasoning to before, one has a step for the coupling T in the form:
1 KL 7kk
min ( diag(1/g))))) ddagg(//aa)))))
Which yields the kernel matrix:
K Th & exp exp(->k diag(gq) ""
Which is likewise projected onto II aR using the Sinkhorn algorithm. From this T E 0g a2 one takes X = diag(1/g2 l/9q) T diag(1/ga) as the inner matrix that corresponds the unequal marginals 9Q and 9R which ensuring P, ll.bb
21

---

Algorithm 4 FRLC
(General marginal constraint low-rank optimal transport)
Input C,r,22, a, b, 7,
Initialize 9Q-9R 1,,
Qu: Ro, To <- Initialize- Couplings(a, -niitalize-Coullnns( b, 9Q, 9n)gn)
if r = then
x To # Invertible case
else
Xo + (iag(11/// 1, TT diag(1/ RR 1m)
General case
end if
while
(((Qk; Rg. Tx), ( (Q k- R1 > E do
VQ= - CR2 x I - 1, diag ((CR, ) didigg(L/gg))
VR - C²Q X Imdia"" - (diagg1//g) R C Xx)²
Th 7/ max V|V || [ # 8 normalizzaiin of Scetbon & Cuturi
2022
k(k) k( + Qx @ exp( ~ VQ). ). R O expl 7V RR)
if Balanced then
Qx SR prrojeccci K
(k)
Tke T, a, Q 1n: 8)
Semi-relaxed OT
Rs projectton K
(k)
Nk T, b, R. mm 8)
Semi-relaxed OT
else if Unbalanced then
Qx U-projection
K
(k) Yk T, a, 1 8)
Unbalanced OT
Rs + U-projection
K R
(k)
Yk T2, b, R 1m;00
Unbalanced OT
else if Semi-Relaxed Left
then
Qx + U-projection K
7k T, a, Q 1 6)
#
Unbalanced OT
Q
Rs < SR prrojection K 7k r. b, RT
(k)
R k-1 1m mm 8)
# Semi-relaxed OT
else if Rela-Rel Right then
Qx + SR " -prjjcciin
(k)
Ns T, a, Q In, 8)
Semi-relaxed
Q
R₂ - U-projection (K(2)
R Yk T, b, RT
1m, 6)
Unbalanced OT
end if
gQ. 9R < 1. R 1m
Vr = diag(go)
Q
CR diag(gR)
Tr / V xc
# 8o normaaliizaion
K - Tx @ exp(-22 Vz}
Th Sinkhorn K (k)
gR, 0) X + diag(1/90 (1/92) (iiag(1/g)) end while
# Balanced OT
Return P. = QXR
Algorithm Sinkhorn Algorithm (Cuturi (2013b), balanced OT)
Input K, a. b,8 a +- In
v + 1m
while [diag(u) Kv all1 + (diag(v)KKT 6/11 do u((+1) a K (vll)
(1++)) t b/ K KTu (1-1) end while
Return diag(u) (o)g(()... K diag(v)
Gromov-Wasserstterr (GW)
As defined in, the general Crromo--asseeseee problem concerns minimization of the energy:
CA.B(P)
Σ CAi B,,) P.j
(30)
k.j.k.l
22

---

Where the minimization is over the set of all couplings Iaa
GW(P, v) == min QA.B(P)
(31)
Pella.b
We consider extending the semi-relaxed framework to the GW problem under the low-rank restriction on P. This extension has thus far been considered for the balanced and unbalanced case in two previous works Scetbon al. 2023 2022). both of these works, the algorithms for the Wasserstein case extend trivially to the Gromov- Wasserstein (GW) problem. In particular, each kernel with variable Ç has an update of the form K + Ç & exp (~>x (cC) for L({) heretofore taken to interpreted as a function of the low-rank sub-factor variables Ç E {Q, R, X) Taking L == DA.n(P(S5 one can simply extend the gradient through the GW-loss and directly use it in place of the Wasserstein gradient the update K <-6 Ç ④ exp (~ 7/ VeL(()) for f(() of each algorithm. The matrix-fomi of the GW-cost is expressed as:
be the Wasserstein loss (P(C), F where the coupling P
aAA P) = 1 PT Pl. + I P
02 1 1n 2(APB, P)
Which, using the constraints of C2 reduces the cost as function of Q. R, X to:
QA.((.RR.))=1 -" 1 Q1, + B R1, 2(QXR" AQXR B)
VQQA.B R = 2 Q1,1, - BAXX""BXX'
Which is proportional to VQQA B (Q, X) -4AQXR BRX for the balanced and right-
marginal semi-relaxed case. And:
VRLAB R X)= = 2B 02 - AQX
Which likewise can be reduced in proportionality to VAQAAB(Q.
-" AQX
in the balanced and left-marginal semi-relaxed case. The gradients, as presented above, assume a linearization in X &- Xx. If one does not make this assumption and takes X (Q. R) = diag(1/O) 1, diag(1 (//R 1m) rank-one perturbation must be added to the and R gradient of the form:
41, diag - (X B(Q AQ diag(1 (1/9)))
Q
( R 41m diag -1 (XT Q' AQX BR diag(1/gR)) (1/g)))
Analogously, for the gradient on T one can simply take the gradient with respect to X, and sub- sequently T as X(T) = diag(g))-11) diag(gR) The gradient with respect to X is given as:
VXQAB(Q.R..X -4Q" -QQT "QQARQRRII BR
And thus, the directional derivative with respect to X in the direction Vx diag(oo)- -v diag(gR)- and thus by the chain rule VT is:
DQA.(QQ, R, (Vx) = -QXQQ"AQX""
= =4Q -4Q""AQQXI BR. diag(1/gQ) Vr diag(1 / //gn))
So that the gradient with respect to the coupling matrix T given as:
VA -4diag(1/g))4 AQXR" DOXRBBRaR BRdiag(1/ (/gR)
E
Convergence Analysis and Other Proofs
E.1
Convergence and Smoothness of the Objective
We show in Proposition 3.3 that directly applying the block-descent lemma Beck & Tetruashvili (2013) to the template of Ghadimi et al.'s proof Ghadimi et al. (2014) is sufficient to show the non- asymptotic stationary convergence of a coordinate mirror descent procedure in Ghadimi's criterion. The non-asymptotic guarantee of Ghadimi et al. (2014) follows directly in the case of coordinate mirror descent using Lemma and [aaannnemaalllll describe the coordinate mirror descent algorithm in general, and discuss few relevant preliminaries.
Lemma F.2 below. For completeness, we define all notation used,
23

---

Suppose that the vector of n. variables x E R" is partitioned into blocks, x= (x(1), .T....... x(p)). where x(i) E R Here, ........ mp are positive integers summing n. Following the notation of Beck & Tetruashvili 2013 Nesterov 2012). we define matrices U; _R*x" such that x(i) Uf'x x for all i = ....... This also implies x E == U,x(t). This allows us to define the vector of partial derivatives corresponding each block of variables x(i):
v.f(x)= == (xx).
In Beck & Tetruashvili (2013) the gradient of f is assumed to be block-coordinate-wis Lipschitz, with L; the smoothness constant associated to the i-th block of variables: for all h; E R" one has
Vv + U,h.) v.f(x) s L, [b,l),
(32)
and for such functions we denote by L max; L, the (global) smoothness constant of V.. To be clear, a smoothness constant associated to f is Lipschitz constant of its gradient.
Lemma E.1 (Block descent lemma, Beck & Tetruashvili (203)), Lemma 3.2.). Suppose f E C'(R" R) is continuously differentiable function over I whose gradient is clockcoordinateeise Lipschitz ) for L, the smoothness constant associated to the i-th block of variables x(t). Let u, v U,h, Then,
be two vectors differing only in the i-th block: there exists h; E IR such that
Li S(v) S f(u) (VS(u), u) -u-
(33)
Lemma E. is central to adapting the proof of Theorem Ghadimi al. (2014) our case. Their Theorem concerns the non-asymptotic convergence of mirror descent for objectives of the form
min [xx) + h(x),
xEX
where X' is closed, convex subset of TR. f E C IR) is possibly non-convex objective, and where h is an a-strongly convex function. Using notation similar to Ghadimi et al. (2014), we write (xx) f(x) h(x). Additionally, Ghadimi al. (2014) assume V is L-Lipschitz for some L > 0.
To unify our assumptions on f, we suppose that Vf E C* (X is block-coordinate Lipschitz (32 with block Lipschitz constants (L.)-= and that X' itself decomposes as a product X' = -II X. where each X₁ is a closed convex set constraining the block variables x(i).
The proof of Ghadimi relies on f-smoothness of the objective in all variables. We show that component-we smoothness in each block is sufficient to achieve an analogous convergence result for a coordinate mirror descent. To provide context for Proposition 3.3 we now describe in general (1) mirror descent, and (2) block-coordinate mirror descent.
We again follow the notation of Ghadimi et al. (2014). A function w : X > R a distance generating function with modulus 0, with respect the Euclidean norm ||- if continuously differentiable and strongly convex, so that
(x(x) Vw(z)) 2 al|x z||? for all x, X.
The prox-function (or Bregman divergence) associated with is then
V(x, 2) (xx w(z) (Z),)) z):
and from this prox-function, > 0, and some g E TR. we define the generalized projection
(g. u) - V h(u))
arg min ueX
(34)
To describe the projected gradient descent algorithm (which coincides with mirror descent in our case of interest), we first define the generalized projected gradient of at x:
1
Px(xxg,1) =
x+).
The mirror descent (MD) algorithu is as follows: given initial point Xy E X, a total number of iterations N, and positive stepsizes (7x) ) at step ki the (k+ + 1)-st iterate is computed via
1 X &- arg min (vf(x))) u} + V(u, xx) + h(u))
ueX
7
24

---

Among all iterates Xx, the MD algorithm outputs the one at which the generalized projected gradient of least norm. Concretely, x the output of the MD algorithm, where
R == arg min |x,,,ll
O=
and where gx is
gY,A Px(x2, VI(()),
Having described the MD algorithm, let us consider a block-coordinate variant; we assume x admits the block-coordinate structure described above. To simplify the presentation, suppose that in a given iteration k. the block variables are updated sequentially from = ,, This leads to doubly-indexed iterates (xj) with k 1..... N indexing each full iteration through all variables, and = 0, indexing the sub-iterations which update one block of variables at a time.
The coordinate mirror descent (CMD) algorithm takes input initial point Xo X, number of iterations N, and sequence of positive stepsizes (7,, N.p
We set xf Xo, and for k O,...,. N and = 1,... ,....,P. we compute x from x as follows:
(1) f arg min
(Vif(x) ) u,) + Vi Uf + h:(wu))
1 K,,
u.EX
x x
for j *
Here, we have assumed that V can be written composite function of the block variables,
V(x, 2) Σ vi(x(),,())),
as is the case for the KL divergence. We also have assumed that has this composite structure (as with entropy):
n(x) = Σ (x())).
Lastly, for k = O...., N we set x x We define gr, (g gx...p)) collection of block-wise differences, where by definition
to be the
gx..i Px, (x V.f(x) Th,)
(x)
x
U. gx.
The convergence criterion A we use in Proposition 3.3 is the one used by Ghadimi et al. (2014)
summed across blocks. In particular, the CMD algorithm returns iterate XR, where
:= == arg min A A(xx, xk-1), O-.....
(35)
A(xe, xx- = | ?.
For & f + fo as above (f has global smoothness constant L). let d' arg argminxea (xx), and define
= (((o)) - Φb 1/2
(36)
L
The other lemma used in the proof of Proposition 3.3 as follows.
Lemma E.2 Ghadimi et al. 2014), Lemma 1). Lerx* = arg min u) + [V(
and Px(x, g, 7) (x x). Then for all x E IR", all g E IR, and 7 > one has:
(g, Px 7)) 2 a M) (h(xt) - h(x)).
25

---

Suppose one has f E (X,R) whose gradient is block- coordinate Lipschitz, with block smoothness constants (Li) and a function C(X,B) R) which is a-strongly convex. For + + h, suppose one performs coordinate mirror descent Ab minimized over a product of closed convex seis X' = -II== X3 Let the sub atb-iterates with respect to the i-th block update be {x² where X = for k [N] outer iterations. Then one has:
Proposition E.3 (Proposition 3.3
D
2D
min A Xk-1) N(a² /2L)
Na
where D is (36). L is the global smoothness constant of f, and convergence criterion A (X X{-1) given (3). Above, the stepsizes "k,s in the coordinate mirror descent are "kk := a/ L.
Proof. As f satisfies the hypotheses of the block descent lemma, Lemma obtain:
E.1
we apply
33
to
Li |
f(x)) s(({ + (V+f(x) ) x x 2)
xx 11
Noting the definition gx Mh.i
1
(xx 1 - x1) one has
f(x)) S f(xt - The (V f Bx.a,)) + L,
2
Lemma Ghadimi et al. (0014) (stated as Lemma H.2 above) applies identically through block- 1x)--1),, 1), nk.)). Thus for any value V. (xt
wise optimality on X; because g. Px. (x} vxf takes,
Lj 2 Th
f(x²) S f(x} (onm,ill h(xx)) - h(x; "] -
2
and thus,
Lg 2
f(xx) + h(x²)====)))))))))) S f(x f(xx=))))))))) h(x
11
07k,
Tk.
2
The right-hand side above only becomes larger, taking L; = to be the global smoothness constant. Introducing a sum over sub-iterates and total iterates, one has
N.p
N.p
N.
L 2 Th,s
Σ
f(xx) + M(xx))
Σ
f(xt
1(x
Σ
||Bx.,.:ll?
anki
2
R
2
N.p
N.p
N.p
T 2
Σ (x
Σ ((x+
Σ
anks
Thit
2
R,5
k
Noting the end-point condition (xxx") $ 1(x++1), one may cancel all intermediate terms:
N.p
Φ- < ((xN) S (xxo)= -
Σ
anik:
2 nh. l|g.x,..111
R,
Thus one finds the upper bound in terms of (xxo) and the minimum value क. = minxex +(x) of Φ
N
aryii
Tk.
2
|| < (xxx) - ."
(37)
2
Taking 7h.i = a/ L as in Ghadimi et al. (2014), the bracketed term directly above becomes a² /2/ and one has:
N
[a27
a
Σ
min A (XX xx-1)
Σ
min
2L
2L
k
N
P
gg xil2
2L
< x((x)
xb' = =D²L.
26

---

where D is as in (3). and where we used 377 to obtain the last line. Thus,
D
22
min A(x,
N(a2/2L)
Na
completing the proof.
Definition E,4 (Relative smoothness). Let and let C' (R", R) be continuously differen- tiable. Additionally, let w be a distance generating function with associated prox-fune prox-function V. The function g is B-smooth relative if the following holds:
g S g(x)) + (Vw(æ), x y) + BV (y,
Proposition E.5. Lete be a predefined error tolerance, lower-bound parameter, and N the number of inner tterations required for the semi-relaxed
small rank parameter, (0,
projection converge for each ùeration i as 1, -90 Then, the FRLC objective
(k-1) || < <====>- ( 6) for 9Q 1,
(0)
(L R,T) = (Q diag(1/Q"" diag R 1m R ,
component-wise smooth with respect to the variables Q, R, T with smoothness constants { where B; = pool((lCC)))) r,
Proof. The addition of a pair of regularizations on the inner marginals 1, and rKL(R" 1m R 1m) in 21 and 22 ensures that we may use standard results on relaxed optimal- transport which bound how far the marginal Q 1n deviates across iterations.
In particular, for all e 0 there exists and N (number of iterations) sufficiently large, so that | 1m gnll < 6 and ||Q"1, gQ l1 < 6. particular, |ham al. (2020) shows that one can attain convergence to any for the unbalanced problem in N = O(" /c) iterations (hiding logarithmic and --factors) for each sub-problem solving for R2 in Algorithm 4 and analogously
N Ónn /e) for Q. Under the uniform initialization of gn. and after N iterations, have:
Iggn
(N) In
2
1
In
(N)
In
(k)
9
(k-1) 1) 1|2.
With sufficiently large and N = Ō(m² /c) sub- -bbiiterations, one can guarantee that:
-
(k)
In (k-1)
11-
|
This implies, for all iterations of the algorithm and all indices i, that
(gn.): 6, and analogously, (gQ.):
(38)
Thus, by adding the regularization on the inner marginal, one may guarantee lower-bound on the entries of ga and gQ. This is essential for demonstrating smoothness of the objective.
First, we consider smoothness in Q. We note that the gradient in Q splits into two terms:
VQLic(Q,R.1
v( LLC ( LLC
= -CRXT - d,ai '(CCR"") " "Qiia((/gq)))
= v/ LLC 1, ((V( Qdiai
Where
Vofro(Q RS, Tb) VoLLo (Qk, Rτ, Ta)lFF
| (A) Lio(Q (Q++1 Rr. Tτ) ((A) Lrc(Qu. Re, Tx)
(39)
| B L1o(Qu41 R4, Tx) v( (nc(Qr, Ra, T.ll
27

---

Starting with the first term on the right side of (39. one has:
|| (A) (Lc(Qu+1, RR. Tk) V LLC(Qu. Ru. Tk)lle
= =dCCR( 1/ go. (1
(d gq. (TTediag(1/gn.)) )
<
R& || ||] P [ - ( .)lF
Note that R = E., (R&) < Ei, (Re),, = 1, as R& has marginals which sum to one. The
same bound holds for ||T 11. which is also coupling. Invoking the lower-bound (38) of 8 on the entries of the inner marginals, and continuing from the above display.
(1) (lo(Q R2. Ts) (/2 LLC (Qx: R3, TallE
| ||diag diag(1/ (1)0
|
(daag//)) diag(gq.) ))ll
C F ||diag(ge.-1) - diag(g).))l1
53
To further bound the right-hand side above, consider:
[[diag(gq.) diag(gq..
|Q Qi
Σ
Σ (Qx)ss - (Qk-1)
While can easily be upper-bounded by an application of Jensen's inequality as
(Qk)s.j- - (Qu=1))s))
==
2 (Qx)ss -(Qu-)iss)
Σ
r
j=l
(Qk)ss (Qu-1)as) )²
Q+ Q Qx-1l1 1
Likewise, we have that [diig(gn,)) - diag Diagg(Rn..)) (g .) S r|| Rg R Thus, holds that
|| F
||C|| FVT
9Q.. - 9Q. .1l2 <
1Q, - Qx-1llr.
83
83
Next, we focus on the ((B)
q term. Observe that
1 Ea.dia--X- =
Tr(11 diag - xx)
n||diag X1|2 < n X
Thus:
|V BB)
v 11 VVll(V Qut1)
(A) TQ Q diag(1/gq.
( ""Adiag(1/pp) (1/9q))))))))) 1
Q++
Adding and subtracting terms in the norm and applying triangle inequality:
(1) )T ((()( Qu_1) 1iigg diag(1/9q.....) 1/9q.+1) (V
(11) I 1
Q=
(V 21 k++ diag
(1/gg//)) - B
A)
/nl (A)
Q++1 v(4) || /
+ V (1)
1 - -Qd)in(1//qa.)1
Qu
28

---

Invoking the lower-bous on the marginal, continuing from the above display,
(4)
5
Qx+1
Q.
+
12
Q
P
(xdiag(L/gqq)))
Let us consider
|
(1
q
|r
X
|
R S Xx ||]
We have that:
Xxll
i. g(Q)):
{(R,),
Where we always have that Ti S IQ: and Tij S gRj by definition of T as a coupling. As such:
(QQ:RR))
Σ
1 1
(g"' 9R 1m17) I mr
gg
gi
gQ. gR.
8²
f
fj
Thus
|
(A) Qu
|F
x
mr
C p. and the bound above reduces to
A) Q++1
v(di |
v nmr
|C (Q Q
Qudiag(1/go)))
Further bounding the last term in the norm
/nmr
|CI
Qudiag(1/gaa l/gqu-))
+Qudiag(1/gq...) Qudiag(l/gq.)ll
Vnmr
||||| [[[==(() [|][[[diagl/gq Qu+1 QuIIF
8
/dial/// diag(1/9q.)))
/F 82
Q x++ QxIIF
F
8
Thus, the final bound on the V (B)
term is:
Q
(1)
Il + Vnmr
|CIlell@Q++.dia(((/go. [4tel(00--uta/l/ooooo--/--t///// Qudiag(1/ge.))l
8
Qut2
8
FVr
Vamrr 82
|C|| 1+ ) || - F
+
84
The total component-wise smoothness bound on Q is then
VVLLo(Qu+1..e... Rr. Tx) - -VfLc(Qu (Qu, Rr, TT)llF
||C||
((yrr + Vnmr
/r )) + vr) 1Q
[+ /
Qule
82
Identical reasoning applies for V VLLce where we similarly have the gradient split into two terms:
V VnLic(Q. R, T) v LLC V (B) LLC
=CTQX 1, diag -1 "ddaa(( R. Q'
= =v(A) n LLc 1 diag (diag(1//9) R R (A) R LTC)T
29

---

As before, may first show smoothness in ((A) using the same steps as for Q
I (Lo(Qu+1, Rx +1 Tx) V LLo(Q +1 Rk: Tt)lFe
-IC* (//9Q.+1 (T Q' /tt.a///o//dddiii/////oo/aaaa/aa
diag(1/g000 F
IClEE I|diag (ddiag(1/gR))) (1/ - (/ggg))
ICllE
[|g gR.
83
Rx+1 REIFF.
03
For V (B)
one may use the same reasoning as before to find:
R
V (B) L R8-1. V LLC CLc(Qu+1) R2. Talle
Ret1 F(C) diag (diag( (1/98. R V Lio) lee < 11 diag (diag(1//R. R V(A)
S //mlddiag(////
R (A))
Ru+1 LLC (() (A)
R₂ CLClLE
As before, one may apply three rounds of triangle inequality inside the norm to bound this directly in terms of [[V these terms smooth in R by the lower-bound argument, that smoothness in R holds analogousty to Q. The remainder of the proof for smoothness in R thus follows identically to that of Q above.
and [Rx+1 R Each of Re+1 v(A) Irr, ||diag(1/ R diag(1/ / RR,
(A)
For T, the component-wis bound of
TV LLC VT,
D Qu+1 CRRIC S LT TellF
holds trivially for any Lr > 0 as the gradient is uniquely determined by Q and R alone. Thus there exist LQ, LR, LT > 0 component-wise smoothness constants for (Lc R,
We next prove Proposition 3.4 restated just below for convenience.
Proposition E.6 (Proposition 3.4). Consider the TRLC (8). The FRLC algorithm, Algorithm
yields 3-smooth iterates for 3 = (Ily[I|Ci|| F, m, r. 8), where denotes the lower-bound on the
entries of 9Q.9R- Consider the convergence metric of 3.3 adapted from Chadimi et al. (2014) given as:
A (xtk
Σ 1.
1 11 7
Q 11-
R2 RR-
11 + [T T 1
for wx = (Qk, Tr, Ru). Define the gap to the optimal solution as in (36). and let L sup, (i)) to be the global smoothness constant across all components. Thenfor "ik- = a/ L as defined in 3.3 the FRLC algorithm has the non-asymptotic stationary convergence guarantee that:
2D2
min kE
Ao
S
2
Proof. The proof of the -sympttiii stationary convergence of mirror descent of Ghadimi et al. (2014), adapted for coordinate mirror descent using the block- lemma B.3 only requires component-wise smoothness in (Q, R, T). The proof of this for FRLC is given in ??, and the guarantee follows directly for this value of L max(Lq, LR. LT) = poly([|C|| p,m, m, 5).)
O
Proposition E.7. Low-rank Approximation Error. Let SR-W; denote the optimal rank-r approxi-
mation for the seli-relaxed low-rank optimal transport problem, and let SR- denote the optimal solution for the full-rank semi-relaxed optimal transport problem.
Additionally, suppose Cb = =ET by denotes the sum of the entries of the second marginal (Cb = 1 if a probability measure). Then we have the following upper-bound on the objective error:
(R6) - (R-W"(Pb)) s cb
mmx [max(Cpp)
min (Cra)) (minfn, m)/(r 1))
(nin(Cra) P-q
30

---

We note that this bound also applies for the standard balanced optimal transport case, giving:
(Ha, Pb) W (Ha, S
(maxx( - -min(Crg) In (min m)/(r 1))
P-9
P-9
and improves the previous bound of (Pa (b) as the distance matrix C contains only non-negative entries.
W (Ha.
- |C In (min(n, m)/ 1))
Proof. We adapt the proof from Scetbon & Cuturi gave the bound:
2022) for the balanced case, which previously
W Pb) W [b)
<
C |x In (min(n, m)/(r 1))
In particular, for 2 = min(mn n). there exists an optimal rank, (P*) S z where one may express the optimal solution for the non-negative coupling matrix P as of rank-one, non-negative outer products q:: k > 0:
P
ARqURT
k=1
k=1
Where we write this sum in terms of normalized vectors qu qu U Ph Fx/ // and A = || 1: Without any loss of generality, (A is ordered in terms of decreasing value such that d1 X2 2 2 Az Note that we have fixed constraint for the sum of the entries of P for the semi-relaxed case, assuming b a general positive measure, as 1 Plm = (PT1n)) 1m = b-1m m = Σド= by == do (where 6 if is chosen to be probability measure, i.e. in the balanced case). Moreover, it is simple to observe for these ordered values that X I (co/k). As in Scetbon & Cuturi 2022), define the weighted average of the bottom r 1 vectors of the decomposition to be:
X
a,
A
Apr:
B,
t=r di
And take the rank-r approximation using the optimal r vectors of OPT and this weighted average of the bottom to be:
P. Σ Nq +
Σ A i=r
a
Where, by the assumption that P E II, is feasible:
Amria +
Σ 11 -
Am++
A
3,
Arr = bb
2m1
Thus P, E Tbyr is feasible rank-r solution by feasibility of P One can verify that if p E Na.b. then P.1m = a and the solution is again feasible. From this, we observe that the difference between this solution and pa is upper-bound to the difference between P and the optimal rank-r solution:
SRR (H 1b) SR- SR-W"(L (Pa: ub)1 = (P, C)F (P,C))
(Pr.)) IP,C))
Agrr+
Σn a, - AygnT..
F
Σ
do a.1 - Eigarto
ier
31

---

Noting that ar,, 3r and qı, Py are unit normalized positive vectors, the sum of the entries of the outer
product 1, q.r. 1m and likewise for a, Thus, continuing from the above display:
(a, 3), 0)F=
Aila F)F
(a) C)F
min (Crp) P-q
t=r
maxx ppmp m P-q
m
min
Σ A
P.1
S cb
max (pm) - P.q
min
Cpg) p.g
Σ 7
imr
S
Cb
max - - Cpg)
In (z/(r- - 1))
P-9
Concluding the proof. As discussed, this directly applies to the balanced case (for Cb = 1).
F
Initialization
We propose a new initialization of the sub-couplings Q, R, T for the LC-factorization. Algorithm generates random full-rank initial condition in the set of couplings Il.. which still satisfies the marginal constraints. accomplishes this by sampling random matrices which are full-rank and applying the Sinkhorn algorithm to each of them. SSeetbon et al. (2021) proposed an initialization which represents an improvement over the rank-1 product measure which is rank-2. Follow-up work proposed initialization using k-means Scetbon & Cuturi (2022). However, this assumes the previous diagonal factorization and is thus not application for generating a latent coupling which may be non-diagonal, non-square, and with two distinct inner marginals. Our initialization is tailored to the LC-factorization, effective, and has a full-rank guarantee. In particular, higher-rank initializations may exhibit better convergence properties by allowing the gradient to explore a larger set of directions immediately the optimization. This initialization is given in Algorithm 6
Algorithm 6
Sitiaaize-Couplingg
Input E An. b E Am. 9Q E A IR E Ar
Cq ~ [0. 1| CR ~ 0, 1] mxr Cr [0, rxr
K Q t eCa, KR
CR, KT t Cr
Q t Sinkhorn( Kq a, (.)
R Sinkhorn K R. b, gn.
T < Sinkhorn KT, 9Q = Q 1,, In = 11,,,)
Return (Q,R,1) T)
Proposition F.1. Suppose one samples an initial condition on the optimal transport coupling using Algorithm 6 where assume Cig Unif 1) such that P(rank(C) < min(n, m)) Additionally, suppose that a, b > O holds elementwise for both marginals a E An, b A,- Then the elementwise exponential cxp(C) (or -C)) has full-rank and the return Sinkhorn(e-C a. has full-rank.
Proof. It is established that a random matrix C [0. 1 has full-rank with probability one.
For K = cxp(C), it holds that the matrix must be entry-wise positive with K,; 2 0. If columns C. * C.,, then clearly C C and if C.,, C.., > 0 and are independent remain under element-wise powers. One may easily show this by contrapositive. Suppose there exist constants C1, C2 such that:
Oo C2 C2 ok 0
32

---

As C,s: C..y > O. without loss of generality one may assume > 0 and < 0. Then:
G c =
④ Co C ... (-a)
c1
1/k
1= C.j
C,
So clearly one has that C.., -C.,- O for c > O. This implies the columns C.,, and C. are dependent. Thus clear that elementwise powers of entrywise positive independent vectors preserve independence. The same principle extends trivially to exponentiation of the columns, where if one assumes by contradiction that cje + oze 0 for c > O, C2 < 0, one finds
log
(-2a
L 02
C..
C.,.. Without loss of generality, assume 0 < c2 C1 and C.. C. > 0,
so that log
0. Then, considering constants 92
C2C., + CIC = (q1 + C2) 6q21 = 0
Assuming C.,; has greater than one unique entry, which we assume as the entries are sampled densely in IR, the two vectors are dependent if and only 1/1 = -02 and = O, implying C. = . Thus, for the set of independent column vectors of C, given as {C.,) ==1 the set fec... - is also linearly independent. This holds analogously for the row vectors. As C is full-rank and span((C.,i) [m..)))..) we have that span (K) pminfm.a) as K e Ek= (analogously and remains full-rank.
C
R
Sinkhorn expresses each variable as
X = diag(u) K diag(v)
where rank X) rank(diag(u))) K diag(v)). As Cuturi (013b) updates the vectors u + a Kr and v t b/ K u from uy = 1, and vo = 1m if a, > 0 holds element-wise, one has that u, v 0 elementwise well. Then, one has that null diag(v) {0} and null diag(u) {0}. implying that rank (diag(u) KK diag(v) rank(K) = min n2, m)
Thus, our initialization returns a random coupling matrix X TTa of full-rank.
In the next proposition, we show that one can analytically solve for the block optimal weights q for the factorization of the coupling matrix P as P Q diag(1/ / R Forrow et al. (2019) Scetbon
ct al. (202I).
Proposition F.2. For the minimization problem expressed as min (Q (iag(1/9)) Rg) C)F gEa.
One has the closed-form minimizer of g" defined entrywise as: /Wi
g
Σj=1
For diag (QTCR). when w 2 holds entrywise.
Proof. As w 2 O, we consider the simplex condition Σラ 9j = 1 alone. Writing out the Lagrangian
associated to our objective, with X E R our equality-conditio dual variable, we have:
C(g. A) = (Q diag(1/g)Rtt ,C' + + X(1-
Let us consider a rewriting of the inner product term:
(Q diag(1/// R" C =
ΣΣ
C
Qik
R
9k
k=1
Q
C.
Rib
ks
9k
1=1 )=
(QTCR) k.
Σ
TT (1/9)
gk
R=1
33

---

Where division is interpreted element-wist wise in the last line. Thus, one can interpret the problem as minimizing the weighted sum of reciprocals of a density, As a result, we can simplify our Lagrangian's Froebenius inner product to vector dot-product as:
L(g,
= w (L/9) -(
X(1 Σg, 9j)
Thus, the first order condition tells us that the value of the coupling weight gj is related to X as:
((g) = Wj - 9j 9j
W
And by relying on the summation condition on the probability density 9. yields the Langrange multiplier
1 Σ
EV
j=l
so that one finds
Σ
w -
d
w
j=l
Plugging our Lagrange-
Lagrange-multiplier multiplier into the above expression yields:
Wj
9j
(Ein vwo))²
Ei= wi
As the Hessian V2 diag( 7 > 0, we conclude that this value g indeed minimizes the loss over Ap.
G Alternating updates on the dual variables
For the problem:
inf (g.
+ KT 7
KL(Q]|Kg)
+KKL(11,11)) 1T
(40)
one can find a simple set of semi-relaxed updates for the coupling matrix. We note the primal-dup relationship of Sinkhorn, Q diag(e (e fa ) diag(e">hi) and consider the entry-wise first-order condition required for the sub-coupling Q:
Qú (Ko)is)
(QQi.) 1r)
0 = 7/ log
+
A1,
ai
(41)
a + log (KQ)is -A117k
a;
log Qij T7x log
(Q.... 1r)
Thus:
Qú =
a; QT 1,
(KQ)ip e -1s7ih
And in natrix-form. this yields:
Q = diag
a
\>T
Kº diag(e -1)))
Q1,
diag(e fi ) diag(e
34

---

R -projection and And expanding the Q1, term explicitly, noting that X diag(u)1 X we have:
diag
a
Ko diag(e
Q1,
The
-diag
K diag(e )
diag(e Q diaglenhhi) The Ko diagle TeAi)
diag
EnuJi O Koewhi
Thus, we identify e = e meAi as the right dual vector, and identify the following relationship in terms of the left dual vector:
TT.
= -
e not
a Koeenh
enahi O Koent
From 40 the condition that R.. C1 nC2 implies that Qu1m 1m 91. 9. As such, we find that:
1'mm diag(e Thh KJ diag(e (efi) )m de ee (
Implying update for hh in the form:
e
K Koeun
Analogous reasoning applies for relaxation of the other marginal, yielding the SR SR -pojjeoctti (i.e. semi-relaxed OT).
Algorithm 7 SR projjectcn (semi-relaxed OT, left marginal relaxed)
Input K, 7, 6,8
u + 1 + 1, repeat
u
--
(a/ (v) (/)//(++>"))
- (b/KTu)
until 7 max I|og // u/ullau Illo o/ulloo) <
return diag(u) u) diag(v)
H Discussion of Complexity
For (Q, R, T) Raxxi Rmxx2 R the space complexity O(uri +ira mr2) is linear the ranks r1,72 = o(1) are taken to small constants. The time-complexity of Algorithm 4 is O(BLr + m)) for B the number inner Sinkhorn iterations. L the number of mirror-descent steps, m the number of samples in the first and second dataset, and = max(r1, (a. T2, d} for 72,72 the ranks the latent coupling and d the rank of the factorized distance matrix C' (generally chosen to be constant near r2). Each matrix matrix-multin multiplication is of max order (n + m)r² which happens constant number of times in the computation of each gradient V,, and for the respective Sinkhorn matrix-vecte vector multiplications K v and KTu. u. The L outer steps follow from the mirror- descent convergence rate and the number of iterations required for each projection follow from the convergence of Sinkhorn. In particular, for e fixed error tolerance and the entropy constant, one finds + approximation for the diameter of the data B poly(1 //ne) iterations using the Sinkhorn algorithm Charikar al. 2023): Cuturi (2013a).
35

---

I Review of Background Material
1.1 Low-Rank Approximation of Pairwise Distance Matrices
As mentioned previously, works such as Charikar et al. (2023) have developed algorithms with linear O(nn + + (oll:+/oOOl111111 the optimal transport cost value. Recent works low-rank factorization of the optimal transport coupling matrix P (the matrix associated to the coupling E [[,,, v1) Scetbon & Cuturi (2022):
time-complexity and ( + m]d) space-complexity for sketching
Scetbon et al. (2023 2021 have achieved per-iteration time-complexittes of O + m (t)dr) for some constant non-negative rank r 2 1, d the dimension of the metric space, and T the number of iterations. By the JL-lemma one can simply embed the points in dimension d O (nm)//²² while preserving pairwise distances, however, currently no proof' cxist which offer the number of iterations I until convergence to some tolerance E. This is partially due how recent these works are, and also to the non-convexity of the objective which is sensitive to initial conditions(Scetbon et al. (2021). However, the space complexity of the algorithm O(nn m))rr) which is noteworthy for being linear in the number of points and avoids storing the potentially intractable O(nm) coupling matrix P. To accomplish this, however, these works rely on a low-rank approximation of the pairwise distance matrix. A number of works by Indyk and Woodruff have concerned algorithms for finding low-rank approximations for such distance matrices. A seminal work [Bakshi & Woodrufi (2018) developed an algorithm which, given two point sets { and {u, f=1 in some metric space X. finds rank r approximation in O(nn + m))++ 1/e)) for y > 0 arbitrarily small constant and & >> 0 an error parameter. A more recent work Indyk et al. (2019) improves on this one, by reading a sample-optimal O((n + m))/s) entries of the input matrix with a run-time which removes dependence on y that is merely ((n + moll(y( 1/E)). This algorithm is used by all of the low-rank optimal transport works. These works, by finding a low-rank approximation the coupling matrix P ~ ABT E R4 due to space limitations on the coupling, necessarily cannot store the full distance matrix C' E R of the same size in memory either. As such, it must also be approximated as C s VU before input to the algorithm, where we necessarily require very effective approximations of U and V to tolerate the additional source of error from coarse- graiining the distance matrix to be low-rank. As such, we present some of the details and algorithm of Indyk essential component of the existing low-rank optimal transport solvers. We begin by summarizing the main theorems in Indyk et al. (2019), which provide an algorithm (upper-bound) on the low-rank distance-matrix approximation problem and lower-bound on the number of entries which must be read.
et al. (2019)
Theorem L.1. Indyk et al. (2019) There a randomized algorithm that, given a distance matrix C e R reads O((n + m) m)r/ee entries of C, runs in time Ö(n + mn)yyor(rr, 1/e and computes low-rank factors V E Raxx, U Br that with probability 0.99 satisfy:
IC - VUII S ||C C, C|
(42)
For C, the optimal rank-r approximation C. C.
This is remarkable result, especially in light of the next theorem.
Theorem 1.2. Indyk al. 2219 Let S m S and y >
such that r/s Ommnn(m, n n//3))
Any randomized and possibly adaptive algorithm that given a distance matrix C E Raxm computes V Raxr, U E xm satisfying ||C- VU S I|C-rrll CrIl?+++IIC prr must read (0(n+m))r// entries of C in expectation. This lower bound also holds for symmetric distance matrices C E Sn-
The lower-bound follows from a difficult argument which involves constructing a hard distribution over distance matrices, involving the use of random matrix theory. The upper-bound, however, follows relatively straightforwardly from a number of previous algorithms and their associated guarantees, along with the algorithm presented below. We introduce the algorithm and also offer a proof of [.] for completeness.
Where Ö(-) hides poly-log factors.
36

---

Algorithm 8
Low-Rank approximation for distance matrix C
Input point sets {ar; }2
} in metric space X' and metric d Pick indices 1" [n], [m]] uniformly at random
for = 1 to do
Update sample probability p: d( yy d( Uz )2 1 m Eド=1 d( end for
Sample (rr/s) rows C.,. Categorical
Σ. Pi
Compute U using Frieze et al. (2044)
Compute V using Chen & Price (2017
return V, U
This algorithm relies on two previous works, whose main results we summarize here.
Theorem 1.3. Frieze et al. (004 Let C E Raxm For sample of O(r/e) rows according to distribution E An which satisfies P: 2 11. for i E [n]. Then in O(mr poly(r, 1/E)) time one may compute a matrix U R from this sample which satisfies:
|[C IC- - ||C Cxlli +
With probability 0.99.
Thus, compute the first low-rank factor U, we need to ensure the yy generated from the algorithm satisfies this requirement and offer the (short) proof below.
Proof. First, it is helpful note d(x,y)² S (d(x, 2) d(z,y))² = d(x, 2)2 + 2d(x, 2)d(y, d(y, 2)2 < 2(d(x, z)a + d(y, 2)²). Where in the last step one uses AM-GM where II, a 1/n Rewriting the norm of row we have:
< E. 24
Σ j=1
d (1 ,ソ)² 2
d( y( + ?2))
22d( yj- )2 22 Σ d(y)), y
j=1 m
< 22nd(x;; yj- )2 + Σ
d(y m.. + d(y )xr)
m
2mdd(x,
y2 yy
Amd(yj- 2, +
d Iz.)) )2
j=1
4m
d( 1,, d(y)),
d y xz,)²
< 4mps
2
==1
37

---

As we have the re- re-normalization expectation E
t Eip. before sampling, we need to consider the value of the
E
P.] to conclude.
1
E
pi
E
d(wi, Uz )2 + yg-) 2+ + m j=1
Σ dl axi yr)
1=1
13
E [m]] [d(z;; yj y/- + Ei- (m--(m))-~-m)) [d(2x.-.)) yz-) )2
E [(n) [d(z
((()))y;)']]
172
m
d( yi) +1
d(ar, yj)
d( Uj)
m
j=1 11m
m
11
j=1
3
a YR
3 |C
da
m
m
1=1
By an application of Markov's 's inequality we have that:
P
Σ
Ps
mo
Pi
C|C
2
E
C1
Pd
1
m2
Thus with probability
we have:
mp;6
4mp; C
|C C
| |
26)
Dr Pr
C|C
This indicates the algorithm presented has probabilities with appropriate bound for using the algorithm of Frieze et al. (2004) to sample the (rr/s) rows of C and generate a rank-r factor U. N
To conclude the result of Indyk et al. 2019 requires reference to an additional work which solves a regression problem for V given C and
Theorem L.4. Chen & Price (201)) There is a randomized algorithu A, given matrices C E nxxm U E T reads only O(r/e) columns of C with time-complexity 0(mr) + poly(r/e) and returns V Raxr which satisfies
- -UI (1++))
min ||C - zU 1
ZeR
with probability 0.99.
Thus, using the result of Chen and Price Chen & Price (011)) one may find a satisfying V easily for fixed U, C. In particular, [ndyk et al. (2019) concludes by tying together the low-rank distance matrix algorithm and the guarantees of the algorithms from Frieze et al. (2004) Chen & Price (2017) as follows
CC - vU < (1 m)n+mICiiiC- min 2UI
< (1++) - CUT UIB
I (1+ C2 + +
= |C 1. +el - C.I. + (1 + +c
C
+ < ||C + [=+ <
C,
Which achieves the bound up to constant scaling of and shows the result of Indyk et al. (019). We next investigate low-rank optimal transport solvers, which assume the result and algorithm of Indyk et al. (2019) to tractably scale to massive datasets.
38

---

J Connection of Optimal Transport to Projection Problems
Before discussing works which have address the problem of finding low-rank decompositions of the coupling matrix P. we discuss a few relevant preliminaries. The Bregman-divergane of some function F(x),, defined by the first-order Taylor expansion of F:
Dr(ar y) F(x) F(v) + VF(y)* -a
For the negative entropy function -) this corresponds the KL-divergence KL(Ş {)- We introduce the iterative- Breegann projection algorithm in the context of the KL-divergence owing to the direct connection with ceropiccaly-regulriiae optimal transport.
Definition J.1. Iterative Bregman ProjectionsBregne
(1967)
Suppose C = - Ci is an intersection of closed convex sets {CI For n. > L, let the indexing be L-periodic C, := Cn mod L. Suppose we want to find minimizer Ç of the KL-divergence with some positive vector & I such that C. This is to say, we hope to solve the problem of minimizing a distance subject to the condition that one remains in this intersection of convex sets:
min SEC
()()llE)
Where the projection of & onto the set C is denoted by
5 = arg min (T((CI))) == (C
Supposing that each C₁ forms a quotient space C. = V/U for a subspace U C V, defined as V/U (U+U E V the iterative Bregman projection algorithm alternates projections onto each set C, as
(x)) + (C
(43)
starting from = c.
One may show Bregman (967) the convergence of Bregman projections to the unique minimizer in C, Ç", where we have the guarantee that ((x) - 5* as n + co. These iterative projections only have convergence guarantees when the constraint sets are quotient spaces-this is clearly not the case for the constraints of the optimal transport LP, and few more notions are required. Definition J.2. Dykstra's 's EggorithnDykstte (983)
Given a point æo E for Euclidean space to find the unique point in C = n for closed, convex sets C₁ that minimize the distance to do as
x' = arg min soll2 #EC
One may initialize the residuals q- (-(L1)
40 and apply the algorithm:
I, = Pc. [ + q
9r. (In-1 xx,) qn-L
Where C2 Cn mod L as before and Pc., denotes the projection operator onto the convex set C...
One can note that Dykstra's algorithm for projections onto intersections of convex sets no longer relies on the assumption that the set is quotient space, and applies in the case that it is merely closed under convex-cone combinatooss.. To generalize Dykstra's 's for more general functions that than the fg-norm, may define the projection with respect to the Bregman divergence of cost function F and define the projection by the minimization: Pc, := arg min DF (x; y). was proven in Bauschke & Lewis 2000) that the generalized form of Dykstra's iterations given as:
CG Pc.,
g)-1))
4.. (VF(x,,-) VF(x,)) + 4n-L
For F denoting the Fenchel-conjugate of F. which define later in connection the low-rank dual problem. Notably, Bauschke & Lewis (2000) also provided guarantees of convergence to the optimal solution which extend the case that
the negative entropy and Dr the KL-divergence. These constitute Dykstra's 's algorithm with cyclic Bregman projections, and project point to the closest point in the intersection of convex sets n C, for an arbitrary cost function F and its associated Bregman -ivereende Dp.
2This to say, æ, y E C₁ = ciæ + c2y E V/U for any 02 E R.
3 e.g. H, gdd xm etc.
39

---

J.1 Dykstra's 's algorithm with cyclic Bregman projections
As such, we can see the updates for being the negative entropy and D the KL-divergence, Without proof, the conjugate of the negative entropy is simply given as (C) = exp(C 1}. Thus, for the minimization problem:
arg min (T(()l)) == (
<
Letting log q- (T.-1) = log O, one may combine the two algorithms above to solve this
minimization using generalized Dykstra's iterations. In particular, we have:
((m) Pc.. (VF (VF((F((((>)) + log qn qn-1))
Pc.. exp (VF(((n-1)) + log qn-1 - 11)
Pe. (xxp (log ((n-1) + +1 log qn-1 1001--11) -1)
Pc.. (c(n-1) 3 qn-1)
And:
log qn = )
-FF()))) log ((n-1) + (log Ç (1) + 1) + log qn-L)
+ log qn-N
((1-1)
log I ( ((n)
So that:
((m) Pen (c(n-1) 3 0qn-1)
(44)
((n-1)
On 9n-L
((m)
(45)
Where division is interpreted to be wise, @ refers to the Hadamard product, and the logarithm applied elementwise.
Connection Sinkhorn distances
Interestingly, the Sinkhorn algorithm described in Algorithm 5 can be alternatively derived in the context of Bregman iterations Benamou et al. (2015) as minimization of the form:
We v) = min KL(PIEE) PεII(a.))
Where & is the kernel & e C-C C/e and IT(,,)) = C1 n C2 for the convex constraint sets C1 (P: P1,, a} and C2 [P: pT1, = b}. To cast this into the Bregman alternates between the two updates:
Bregman-projection projection framework, one
p( =Pc, (P!l-1)
pll+1] = Pccs (P!))
Where for the first projection one has the following first-order KKT condition:
((L((P)p((D)) + (PIm -a)) = log
P
+ A1 =
P(1-1)
P diag(e 2/c) p((-1)
With the constraint of C, that P1m a, this implies diag (a/ p((-1) = diag(e >/5 (M/) and recovers the first update of Sinkhorn p()) +- diag(a, a (((-1)] p(l-1), An analogous argument gives the second, where all iterates satisfy p(t) diag(u (1) (1) )e -C/= C/+ diag(n for u, as defined in
Algorithm 5
40

---

0.45
Rank-2
Full-rank
0.40
(C, Pr
0.35
0.30
0.25
0
20
40
60
80
100
Iterations
Figure 5: Transport cost (C, P)F against number of iterations for FRLC with rank 200 on synthetic dataset of two moons and eight Gaussians. Smooth convergence is observed for both rank-2 and full-rank random initialization.
K Additional Simulations
We tested on two additional synthetic datasets, both used benchmarking datasets in Scetbon et al. (202T). We follow exactly the parameters provided in SSeetoon el al. (202T) to simulate these datasets. For the first one, simulated
m 10, 000 points from Gaussian mixtures in 2D (Fig. 6
The first Gaussian mixture is a mixture of three Gaussian distributions with means (0, 0), (0, 1), (1. I)
respectively. The mixture proportion is 1 and the covariance is 0.05 x identity for each Gaussian. The 3 second Gaussian mixture is a mixture of two Gaussian distributions with means (0.5, 0.5). (-0.5, 0.5) respectively. The mixture proportions is and the covariance 0.05 x identity for each Gaussian.
For the second dataset, we simulated n = m 5, 000 points from two Gaussian mixtures in 1OD. The first Gaussian mixture is mixture of three Gaussian distributions with means ..)) ,0) respectively. The mixture proportions is 3 and the
(0, 0, 0, ... 0), (0, 1, 0, ... 0), (1, 1, 0, .........00 covariance is 0.05 x identity for each Gaussian. The second Gaussian mixture mixture of two Gaussian distributions with means (0.5, 0.5, 0, ... 0), (-0.5, U......... respectively. The mixture proportions is 1/2 and the covariance 0.05 identity for each Gaussian.
For each dataset, we repeat the same procedure as in 4.1 running FRLC and LOT with Euclidean distance to find a low-rank coupling matrix between the Gaussian mixtures with rank between 50 and 200. We observe the same pattern as Fig. Pp. FRLC obtains lower transport cost with increasing rank, and achieves lower cost for each rank than LOT under all initializations (Fig. 2ec, Fig.
L Graph Partitioning
L.1 Evaluation on a Graph Partitioning Task
next evaluate FRLC on an unsupervised graph partitioning (node clustering) problem described by Chowdhury & Needham 2021 ). Specifically, given a graph G = (V, E)) of n nodes, represent the
graph as = (A, where A E RRxa encodes the intra-graph node relationship (e.g. adjacency matrix) and h A, uniform measure. We cluster the nodes of G by estimating GW coupling
41

---

Figure 6: Plot of the two simulated mixtures of Gaussians in 2D, following the same parameters as Scetbon al. (2021)
0.195
FRLC (rank-2)
FRLC (full-rank random initialization)
LOT (rank-2)
0.190
LOT (default initialization)
LOT (k-means)
0.185
+ 1 (Full-rank, FRLC)
0.180
0.175
0.170
0.165
120
160
200
Rank
Figure Transport cost (C, F) achieved by LOT Scetbon al. 2021 and FRLC across different ranks and different initializations on the Wasserstein problem on the synthetic dataset of two mixtures of Gaussians in 2D.
42

---

Method
Dataset
Time in seconds
Cost Value P)r
Pln - aal2
| b112
(CPU)
FRLC
2-moons,
0.379
0.207
1,56E-5
6.1E-18
LOT
8- 8-Gaussians
0.751
0,210
1.90E-5
1.89E-5
FRLC
Gaussian mixture
0.354
0.178
1.05E-5
7.0E-18
LOT
(2D)
0.735
0.181
1.65E-5
1.70E-5
FRLC
Gaussian mixture
0.323
0.294
2.48E-6
8.9E-18
LOT
(10D)
0.677
0.307
1.39E-5
1.46E-5
Table 2: Runtime of FRLC and LOT on the synthetic datasets of 1000 samples, as well as cost value P) and marginal tightness for context. This was done with the FRLC set- ting max_ inneriters satee__blllaaaeeeeeeeeeeaeeeeeae balenced= 1000, max_ inneiinneessseee. Seriiterrrreexee500 min_ iter=7 and rank 100. This time excludes the extra time incurred for ott-jax problem setup and includes the setup time for FRICC
between G and smaller graph = (B, (), where B IR m represents the relationship between each of the clusters and h E Am is the proportion of nodes of G in each cluster. Without priori knowledge, h is set to be uniform and B is set the identity matrix. |Vincent-Cuaz et al. (2022) notes that instead of solving a balanced GW problem, semi-relaxed GW with the right marginal h relaxed learns h from data and leads to more accurate node clustering.
Dataset
GWL
FRLC SR-GW
SpecGWL
SpecFRLC SR-GW
Wikipedia (1998 nodes, 2700 edges)
sym, raw
0.314
0.387
0.372
0.444
sym, noisy
0.250
0.361
0.293
0.400
asym, raw
0.263
0.276
0.194
0.304
asym, noisy
0.208
0.201
0.141
0.177
EU-email (1005 nodes, 25571 edges)
sym, raw
0.434
0.464
0.009
0.040
sym, noisy
0.392
0.422
0.009
0.014
asym, raw
0.388
0.398
0.012
0.028
asym, noisy
0.385
0.348
0.008
0.012
Amazon
raw
0.322
0.338
0,505
0.479
(1501 nodes, 4626 edges)
noisy 0.274
0.257
0.438
0.453
Village
0.531
0.710
0.553
0.579
(1991 nodes, 8423 edges)
noisy
0.413
0.536
0.397
0.829
Table 3: Performance (measured using Adjusted Mutual Information (AMI)) in graph partitioning for full-rank OT algorithms GWL, SpecGWL and full-rank semi-relaxed FRLC. The top performing method for each dataset is highlighted in bold.
We benchmark FRLC for semi-relaxed GW four real-world graphs: Wikipedia hyperlink network with 15 webpage categories Yang & Leskoved (002)); an email interaction network within
European institute with 42 departments Yin et al. (2017); ; Amazon product network with 12 product categories Yang & Leskovee (2012); ; and network of interactions between 12 Indian villages Banerjee al. (201)). We also test the symmetric and noisy versions of each graph provided by Cyhhhwdhuryy & Needham (2021). We compare with two OT-baaed methods: (1) GWL. Xu et al.) (2019). which solves a balanced GW problem between and G with the adjacency matrix G as the intra-domain cost matrix A; (2) SpecGWL Chowdhury & Needham (202l) which uses the heat kernel on the graph Laplacian as A. We similarly run our FRLC algorithm using with both the adjacency matrix (denoted FRLC-SR-GVV and heat kernel (denoted ESpecFLC--R--G GR) Since the number of clusters in each dataset is not large, we compute the full-rank coupling matrix in each case.
Overall, FRLC achieves the best clustering performance on 9 out of 12 datasets (Table 3). When using the adjacency matrix, our semi-relaxed algorithm achieves better clustering result than GWL. on out of 12 datasets. When using the heat kernel, our semi-relaxed algorithm achieves better
43

---

clustering result than SpeeGWL on 11 out of 12 datasets. These results show the importance of semi-relaxed OT on real-world problems, as well as the accuracy of FRLC.
L.2 Problem Statement
As discussed in Vincent-Cuaz al. 2022); Chowdhury & Needham (2021) it is possible to achieve unsupervised graph partitioning (node clustering) through Gromov- Wasserstein (GW) OT. Given a graph (V, E) of n nodes, we encode it as G = (A, h)), where A E R encodes the intra-graph node relationship (e.g. adjacency matrix, Laplacian) and E An is uniform measure over the nodes. If we want cluster the nodes of into m clusters, we can define new graph G = (B, where E m a diagonal matrix representing the cluster's connections and a distribution over clusters estimating the proportion of nodes in G in each cluster. Since we don't 't know the density of clusters priori, we can set h. to be uniform. Usually B is set as the identity matrix.
To cluster the nodes in G, we can solve a GW problem matching nodes in G with nodes G, with intra-domain matrices A and B:
min
Σ -Ay Pay
G''kl'
(46)
s.t.
P R.5
The cluster assignment of each node in G can then be recovered from P by finding the node in G mapped with the maximum weight.
However, since the proportion of nodes in G in each cluster is not known priori, solving balanced GW problem fixing the marginal P on G to be a uniform h significantly constrains the expressivity of the algorithm. Therefore,
proposed in /Vincent-Cuus et al. (2022), semi-relaxed GW problem with the right marginal relaxed, fixing the marginal on G to be h but allowing the marginal on to deviate from
instead solve
min
(A BBr) P.y Par 1KLL(P"1,,
(47)
s.t.
P np.
The learned h from semi-relaxed GW estimates the posterior proportion of nodes in G in each of the m clusters.
L.3
Dataset and Preprocessing
We run our semi-relaxed FRLC algorithm on four real-world graph datasets: Wikipedia hyperlink network with 1998 nodes and 15 clusters Yang & Leskoved (2012), a directed graph email interactions in European research institute with 1005 nodes and 42 clusters Yin et al. (2017) 7), an Amazon product network with 1501 nodes and clusters Yang & Leskovec (2012), and network of interactions between Indian villages with 1991 nodes and 12 clusters Banerjee et al. (2013). The Wikipedia and EU-email graphs are directed, so we also use undirected versions of them. We also use noisy version of each graph by adding up 10% additional edges following Chowdhury & Needham leading to a total of 12 different graphs to cluster.
(202
L.4 Experiment Settings
We compare algorithm with baseline methods, GWL and SpecGWL. GWL Xu et al. (019) solves the entropy-regularizee version the balanced GW problem of (46) with the adjacency matrix of G as A. We set such that the density each node is proportional to its degree. We set h to be distribution estimated by sorting the weights of h and sampling values via linear interpolation, We set B (h). We set the entropy regularization
following Chowdhury & Needham 202
44

---

0.56
FRLC (rank-2)
FRLC (full-rank random initialization)
0.54
LOT (rank-2)
LOT (default initialization)
0.52
11
11
1
Cr
0.50
0.48
0.46
0.44
20
40
60
80
100
Rank
Figure 8: FRLC achieves lower primal cost (C, for P Ha. than Scetbon et al. spatial-raanseriitonies dataset of mouse embryonic development Chen et al. (2022). demonstrates a more robust trend improved performance with higher rank.
2021
FRLC
?) = 10-6, SpecGWL Chowdhury & Needham (2021) solves the same problem GWL, but instead of using the adjacency matrix as A, it uses the heat kernel on the normalized graph Laplacians CaacciinnChhungg (2005). We set the heat parameter
10.
For our method, we solve the semi-relaxed GW problem of (47) with full rank solution. We use both adjacency matrix and heat kernel C² and label the result the two representations FRLC-SR-GI and POLFFSS---GW We set = 0.01 as to minimize the conformation to the right marginal. Since our method depends on random initialization, we run our method 10 times on each dataset and report the mean performance. We evaluate the resulting clusterings of all methods by computing the Adjusted Mutual Information (AMI) between the computed clustering and the ground truth clustering. This experiment, and the experiments on mouse embryo spatial transcriptomics, were conducted on cluster GPUs
M Spatial Transcriptomies Alignment
M.1 Problem Statement
In this problem, we use FRLC to find a low-rank alignment matrix between cells from two spatial transcriptomics (ST) Ståhl et al. (0166) slices, collected at two timepoints, then use the computed aligmment matrix for two downstream prediction tasks. An ST experiment on a 2D tissue slice yields pair (X, Z). and p is the number of genes measured. X1, E R is the gene expression level of gene j in cell i, where a higher number indicates stronger expression.
R is the gene expression matrix, where is the number of cells the slice
E g is the spatial coordinate matrix, where each row i stores the x-y coordinate of cell i on the slice. Therefore, each cell the slice has a gene expression vector of length p. which encodes the feature the cell, as well as a coordinate vector of length two, which encodes the geometrical information of the cell on the slice.
Our input data is pair of ST slices So (X", "r,,2),,, = (X'_ , with n and m cells, of the and S₁ is collected at timepoint I t= hence the transition from So to S reflects the biological development the time period. We would like to find the ancestor-descendant relationship between cells from So and S by computing an optimal transport coupling matrix between cells from Sn and cells from S
same tissue region. We assume Sp is collected at timepoint
the tissue during
The state-of-the-ee spatial transcriptomics alignment method moscot Klein et al. (023) claims that
45

---

unbalanced OT is the most appropriate setup for this problem. Specifically, given discrete uniform measure a and b over cells from So and Sj. we solve the unbalanced Wasserstein problem
min PERTX
C
1,,
(48)
Ce 1R has entries Ci = c(x; X))), where e is the Euclidean distance between the features of cell from So and cell j from St.
The above formulation only considers the feature information of the two slices, but not the geometrical information. Therefore, we can also solve the unbalanced GW problem
Σ (A B Par + +TKT(Pl + 1I..(1"...1))) PERTXM K.
min
(49)
where E R the Euclidean distance matrix between the spatial location of cells within So. and B E Rm is the Euclidean distance matrix between the spatial location of cells within Si. Similarly, we can combine the above two formulations to solve the unbalanced FGW problem.
Notice the inherent asymmetry stemming from the temporal nature of the problem: all cells on S should have an ancestor from S1, but not all cells from S need to have a descendent in Sy because of cell death. Therefore, the most natural OT task for this problem is semi-relaxe OT with the left marginal (the marginal on the first/ancestor slice) relaxed. Specifically, we can also solve the semi-relaxed Wasserstein problem
min P) P
(50)
as well as semi-relaxed GW and FGW problem.
Gene expression prediction task Given the alignment matrix P linking cells from So to S1, we can predict properties of cells S₁ from properties of cells in Sg. Let the expression of gene in So be a vector fa E R", such that fy; is the expression level of gene j in cell i, we can predict the expression of gene j in Si f3 m pT f₃ Rm The accuracy of the prediction can be measured by the Spearman correlation between the predicted expression and the ground truth expression f3 of gene j in S₁: p(f;, f3 In this work, we test the prediction accuracy on 10 test genes: Tubb2b, Pantri, Actc1, Tnni1, 1, Afp, Hbb-bh1, Fez1, Crabp1, Crabp2, Col3a1, which are markers genes for various cell types in mouse embryo.
Cell type prediction task We can also use the cell type labels of cells in S, predict the cell type labels of cells in Sj. Specifically, for each cell in Sj. we can assign it the type of the cell argmax, P.s in So. We can measure the accuracy of the cell type prediction by computing the Adjusted Rand Index (ARI) and Adjusted Mutual Information (AMI) between the predicted clustering of cells in S and the ground truth clustering.
M.2 Dataset and Preprocessing
We use the large-scale real-world dataset of mouse embryo development(Chene et al. al. (2022), consisting of eight timepoints of ST slices during the whole process of mouse embryo development. In this work, we align the pair of adjacent timepoints of E15. and E12.5 embryos, consisting of 124 124 cells and 51,365 cells, respectively. We preprocess the dataset using the standard SCANPY Wolf ctt al. (2018) pipeline. We first filter the two slices to have the same set of genes, resulting in 26,436 genes for all cells from both slices. We then fog-normalize the gene expression of all cells from the two slices, and apply Principle Component Analysis (PCA) to reduce the dimensionality of gene expressions to 30. We take the Euclidean distance between the gene expression in the PCA space as the cost matrix C. We take the Euclidean distance between the 2D coordinate of each cell within each slice as the intra-domain cost matrices A and B.
Fig. g visualizes the two slices in this dataset, with each cell annotated with a cell type from the original publication.
46

---

Blood ressel
Kidwyy
Broin
Liver
Branchicl arch
Lung primor dram
Cavity
Mereiges
Carl lage primodium
Mastrbery
Connective fisue
Mesothelum
Chomie pleaus
Mucocal eptthellmm
Demmmyytoom
Muscle
pangllo
Ovary
Esiderms
Pancrean
tract
Sclerotame
mesenchyme
Spiaal
estooerm
freee
Symputh narve
booth
Unogenital ridge
E11.5
E12.5
Figure 9: Visualization of the E11.5 and E12.5 mouse embryos, with each cell colored by the cell type annotated by Chen et al. (2022).
M.3 Experiment Settings
backbone of the spatial transcriptomics alignment method moscot Klein al. (2023), which the shown by |ceehhnn et al. (2023); y Klein et al. (2023) to achieve state-of-the art performance on spatial transeriptomies alignment. We use Sceetoon al. (2023) to solve three unbalanced problems for spatial transcriptomics alignment: the unbalanced Wasserstein problem of (48) (result denoted as LOT-U-W), the unbalanced GW problem of (49) (LOT-U-GW), and the unbalanced FGW problem with convex combination of the previous two costs (LOT-U-FGW), We use FRLC to solve the same three unbalanced problems (results denoted as FRLC-U-W, FRLC-U-GW. FRLC-U-FGW), Since the two slices contain > 30,000 cells and > 50,000 cells respectively, full-rank solution is not feasible, hence we solve for low-rank solutions with the rank validated. We also solve semi-relaxed versions GW (FRLC-SR-GW. FGW (FRI.C-SR-FGW)
We compare with the unbalanced low-rank optimal transport algorithm of Scetbon et al. 2023
of Wasserstein (result denoted as FRIC-SR-W problems using our FRLC algorithm, as well as using a particular setting of LOT-U that equivalent to semi-relaxed solver (results denoted as LOT-SR-V W, LOT-SR-GW, LOT-SR-GGW).
We perform extensive grid search to find the best hyperparameter combinations for each method and each problem. The grid of hyperparameters searched for each method is reported in Table. 6 The best performing hyperparameter combination for each method is reported in Table. along with the performance on the validation genes. We pick the best hyperparameters using the Spearman correlation on the gene expression prediction task for 10 validation genes: Ckb, Fabp7, My14. Tnnt2, Apoa2, Hba-x, Tubb3, Epha7. Ldha, Col 1 a2, which are marker genes of various cell types well. We report the performance of the aligmment computed by each method using the Spearman correlation on the gene expression prediction task for 10 test genes, well as the ARI and AMI the cell type prediction task. Fig. 10 visualizes the ground truth cell type classification versus the classification predicted by our method FRLC-SR-W.
M.4 Runtime
We report the runtime and OT cost of FRLC and LOT Scethon et al. (2021) on this dataset (mouse embryo E1 1.5--2.55 well as other datasets (E9.5-10.5, E10.5-11.5 from Chen et al. (2022) in Table For all three datasets, FRLC achieves a better OT cost in a shorter time.
Dataset
LOT (seconds)
FRIC (seconds)
LOT (OT Cost)
FRLC (OT Cost)
Mouse embryo (E9.5-100 10.5)
2.545
1.112
0.440
0.385
Mouse embryo (E10.5-11.5)
4.209
1.190
0.371
0.344
Mouse embryo (E15.5-12.5
8.667
1.889
0.478
0.439
Table 4: Comparison of methods on Stereo-Seg mouse embryo spatial transcriptomics datasets using
GPU and default settings: min ittr=10, max_ 100 rank 50.
47

---

-rooss of unity
N.1 Problem Statement
To test FRLC can effectively coarse-grain transport between two datasets with obvious cluster structure, we generate a pair of two-dimensional datasets z and z with ten and five clusters respectively. We run FRLC with T E I see if the barycentric projections for the LC factorization (Definition 4.1) can recover the cluster structure. These projections induce 10 and barycenters for the first and second dataset, and are defined by
Y diag(1/gq)) Q" Y diag(1/gn) (*)
We examine, visually. whether these barycenters are good representatives of the clusters in each dataset, and whether the latent coupling depicts reasonable transfer of mass between barycenters. We also run FRLC with T ε 1R 10 and plot the barycenters from the resulting factorization for comparison. As discussed in 4.2 the barycentric projections defined above, and in Definition 4.1 can be applied to factored couplings Forrow et al. (2019); Scetbon et al. (2021), yielding projections of the form:
(1) == diag(1/g)
(2) == diag
Thus, we also ran the method of Scetboon et al. 2021) on this data, called LOT throughout the experiment, with g (IR2) and g (R4 yrm ploting its barycenters in each case along with the diagonal latent coupling diag(g).
N.2 Dataset and Preprocessing
We instantiate pair of two-dimensional datasets (1) and (2) as follows. Let Un denote the n th roots of unity:
Un fc 2m2h ls ........... n).
These complex numbers can equivalently expressed as ordered pairs on the unit circle:
(Re(e2tik/n) CA [Im(e2rik/m)
cos(22kk/m) sin(2
We consider n uniformly weighted mixtures of Gaussians, where for each sample X, we first sample a root of unity uniformly,
k Uniform(n)).
and conditionally on k, we sample X from an isotropic Gaussian centered at this root of unity
X ~ I²Id2),
using standard deviation a 0.1. Samples are generated with the make bake_bloba function from sklearm datasets. We generate two datasets in this way, (1) for 10, and (2) for 5. To generate dataset ( we first homogeneously scale the roots of unity to lie a circle of radius 3 before sampling, so that the two datasets do not overlap but still use the same standard deviation. We do not scale the centers used for (2)) so they all lie on the unit circle. We generate
(11) = {z] (1) ...., ....,, and (2)
{z}
(o00a using 1000 samples each and form the
(2)
empirical measures
1000
1000
Σ
6 :==
Σ
1000
1000
1=1
supported (1) and 22), corresponding uniform probability vectors a, b E A1000-
N.3 Experiment Settings
We used default hyperparameter settings for both methods. In particular, FRLC sets T = 75 and 90, with maximum of 200 iterations and minimum of 25 iterations, subject to the stopping criterion A given in (T).. The LOT default settings are 7 = 10 (there no hyperparameter analogous to in LOT). Both methods use random initialization and were run CPU.
48

---

Discussion of differences between FRLC and existing low-rank optimal transport algorithms
We provide extensive comparison of FRLC against the parametrization and objective of Scetbon ct al. (202)) in Appendix A As Latent- OT Lin et al. (2021) is an extension of the k- Wassersten barycenter problem, it has distinct objective and thus performs worse on primal OT cost (Table [
making a direct experimental comparison on primal cost minimization only appropriate relative to the
works Scetbon al. (202 2023 2022). This stated, still list a number of distinctions between
FRLC and Latent-OT noting that most differences between FRLC and Forrow et al. (2019) transfer from this discussion since Lin al. (202T) extends the k-Barycenter problem of Forrow et al. (2019).
(1) Lin et al. optimize two sets of variables: sub-couplings (Q. R. T') and anchor points (z", 2 on which the sub-couplings depend. FRLC only has (Q, R, T) as variables of the optimization. (2) Cost matrices used in Lin et al. (2021) are built from distances between each dataset and its representative anchor points for Q, R. or the distances between the two sets of anchor points for T. In contrast, ground costs used in FRLC to update (Q, R, T) are always derived from the distance matrix C' in the Wasserstein objective (C, P) F. Specifically, the cost matrices used by Lin al. are:
[Cqlik = ||= 2 12 [Cnlje = lll, - 21 [CTlke = 1| - 22112
optionally using a Wasserstein distance for the entries of CT. (3) FRIC costs are given in the exponents of the Gibbs kernels written above and below Equation 9 These are derived directly from the rank-r Wasserstein problem mon (CC P)x and differ substantially from those of the proxy objective in Lin et al. (202I); Forrow et al. (219). (4) The different objectives and variables lead to very different algorithms: Lin et al. alternate updates to the sub-couplings (Q, R) using Dykstra, with updates to the latent anchor points (z*, using first-order conditions. In contrast, FRLC alternates semi-relaxed OT to update (Q. R)) and balanced OT to update T. (5) Because FRLC does not require anchor points to define costs, FRLC can handle cost matrices which are not simple functions of distance. For example, if C. is the price of transporting good to warehouse j one may not be able to re-evaluate a price c( =f{ between Fi and latent anchor z In such situations, while finding low-rank plan may make sense (e.g. approximate an "anchor" may not have clear definition in the setting of general cost matrices. (6) The Lin et al. objective is only a proxy for a Wasserstein-type loss, and Lin et al. do not explore extensions to Gromov-Wass Wasserstein (GW), or Fused GW, which FRLC readily generalizes to. A summary of the existing low-rank OT algorithms and key distinctions between them given in Table [
assignment for massive dataset),
For completeness, we offer a compare against the work Latent OT OT et al. 2021), which solves a variation of the k- &-Wasserstein barycenter problem. As discussed, while their factorization is similar, their problem is distinct from FRLC as it does not solve the primal OT problem for general cost. We report the cost obtained by FRLC and by Latent OT various simulated datasets in Table 5
Dataset
OT-coost (FRI.C)
OT-cosst (Lin et al.)
5th and 1hth roots of unity (rank r1, F2 = 5, 10)
1.174
2.124
Two-moons and 8- 8-Gaussians (rank r 20)
2.716
4.291
2D Gaussian mixture (rank = 20)
0.552
0.922
10D Gaussian mixture
1.038
1.298
(rank r = 20)
Table 5: Comparison against Lin et al. (2021) in primal OT-cost (C, F)
P Limitations
Our method introduces an additional hyperparameter relative to previous approaches Scethon et al.) (202I), controlling the strength of the KL penalty the inner marginals when updating Q and R.
(Spetoon
49

---

Empirically, we found FRLC to be robust to different choices of т. but applying the method optimally requires this additional hyperparamter any grid-search.
We also note that the non-asymptotic criterion ((:)) is weak relative to stronger notions of conver- gence, and that often users might prefer simply the method up some number of maximal iterations by setting the parameter for whether ((:-) is used to False. The W optimization em- pirically converges to minima smoothly, so for the most part there not much of a need for A except for early stopping. We recommend that users plot the loss over iterations and use it to set the tolerance parameter tol, and the minimum and maximum iteration parameters for the time-being. The needs for these parameters might vary widely-the the minimum number of iterations should be very low (around 5) for simple datasets and substantially higher for nigh-dimensional, structured ones.
Although we demonstrate strong performance already. there massive room for improvement as our implementation is preliminary and not at the level of high-performance library like ott-jax. We use lightweight vanilla implementations of Sinkhorn as sub-routine, not taking advantage of the momentum-basee techniques which could accelerate massively. Thus, one can imagine that the potential scalability and speed of this method could be much higher than reported in this document.
Broader impacts
FRLC is general enough to be used modularly within any ML algorithm using OT as a subroutine to help with scalability. We also note that the LC factorization is similar to a PCA in the context of OT, yielding an optimal low-rank coupling with an interpretable latent coupling factor.
50

---

Hyperparameter
Values
rank (Both)
50, 100, 200
(Ours)
30, 50, 100
T (LOT-U)
0.99, 0.9, 0.7
(LOT-U)
0.001, 0.01, 0.1
Table
Hyperparameter grid considered hyperparameter search for validation. Scetbon al. 1- and their 7' are in the same range.
(023
Klein et al. (2023) scales 7'
Solver
Rank
(Ours)
(UL)
eUL) Spearman p (Validation)
FRLC-SR-W (Ours)
200
30
0.465
FRI.C-S--GT (Ours)
100
100
0.288
FRLC-SR-FGW (Ours)
200
0.465
FRLC-U-W (Ours)
200
100
0.471
FRLC-U-GW (Ours)
50
0.282
FRLC-U-FGW (Ours)
200
30
0.353
LOT-U-W
200
0.001
0.394
LOT-U-GW
100
0.99
0.01
0.001
LOT-U-FGW
200
0.7
0.001
0.393
LOT-SR-W
200
0.7
0.001
0.394
LOT-SR-GE
200
0.7
0.1
0.003
LOT-SR-FGW
200
0.99
0.001
0.399
Table The best performing hyperparameters for each solver and performance on the validation genes.
Bboo weal
Konwyyy
Brift
Liver
Lung primondumm
Chomii pleas
optnectum
Demmanytoon
Muede
Dencal root ganglion
Ovary
Ca Caddems
Pancreas
Gi tract
Schretome
Head menanchyma
Sonral card
Ground truth
Predicted
Figure 10: Ground truth and the predicted cell type classification of the E12.5 embryo.
51

---

(a)
(b)
Target
Latent source points
Latent Larget points
Source
Target
Latent source points
Latent target points
(c)
(d)
Source
larget
Latent source points
Latent target points
Source
Target
target points
Figure 11: (a) LC-projection barycenters aligned with FRLC latent latent-coupling T on (a) two moons and eight Gaussians (r = 30), (b) LC- -rojjejettion barycenters aligned with diag(g) from Scetbon ett al. (202)) (r 30), (c) the checkerboard dataset with FRLC latent coupling aligned barycenters
(r = 12), and (d) with diagonal alignment [Scethon et al. (021) (r = 12). We show in A. that the
output of FRLC can be diagonalized to the factorization of Forrow et al. (2019) (Figure T2h.
52

---

Source
Target
Latent source points
Latent target points
Figure 12: As discussed in A.1 one may recover the factorization of Forrow al. (2019) sub-case of the LC- LC-factorization. Shown is the factorization found by diagonalizing the output of FRLC from (Q, R. T) + (Q', R, diag(gR)).
53

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The abstract and introduction exactly describe the background preceding this paper and placing it into context, and exactly describe the contribution and scope of the paper to the field.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
We describe the limitations of our method in Section P These include an addi- tional hyperparameter that FRLC introduces relative to previous work, and the substantial room for code optimization.
Justification:
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). ). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-tex system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
54

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: All theoretical results state their assumptions clearly and all propositions are followed by thorough line-by-line proofs with careful justifications made for each step. They have checked over by all of the authors. All proofs are provided in the supplement with detail and are referenced in the main body where appropriate.
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide all code used for generating the results of our paper. This not only includes the source code for the optimization, but the experimental code used for bench- marking. The algorithm is implemented exactly deseribed in the paper, and reviewers freely consult the code we provide to them.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm.
(b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
55

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We will upload all code, and all code required to generate the synthetic data experiments used. Any real data experiments using especially large-scale data have publicly available and easily accessible datasets which provide references for. We provide comprehensive descriptions of how the data was pre-processed, and provide experimental code which can be followed easily.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cm cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results. See the NeurlPS code and data submission guidelines (https //nips.ccl ltliiss.cc/publicguuuueee
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question:
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We introduce all hhper--yramene parameters with rigorous justification of their utility. The values of the default hyper-paramen parameters accessible in the code upload and any experiment which does not use the default hyperparameter (e.g. for a validation search) has the full table provided and the code used for the experiment.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
56

---

Answer: [Yes]
Justification: The optimization is deterministic, so generally there's 's no need for reporting statistical significance. One small exception is that one of proposed initializations is randomized, and for experiments which use it we do include lo error bars.
Guidelines:
The answer NA means that the paper does include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We indicate when experiments are run on CPU versus GPU the experimental details.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics httpp::/neurrpp:: httppp/llc//..iiccciiiiiieeoeeee
Answer: [Yes]
Justification: This research was conducted ethnically and conforms to the guidelines of the code of ethics. We do not anticipate any major negative societal impact result from this work on optimal transport.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
57

---

The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: We address broader impacts Section Q FRLC can be used for scaling an interpretability in any ML method using OT as subroutine.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: We do generate any new datasets other than toy dataset of Gaussians centered at th roots of unity, which has no societal impact. We do not believe our optimal transport work has significant risk of misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
58

---

Answer: [Yes]
All data that we use and all work that we build on are cited and credited heavily.
Justification:
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, wwwppsppppeaaaaeeeee thcode com datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We provide no new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/mooel as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Question:
Answer: [NA]
Justification: This paper involves research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
59

---

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
60


Reviews:
[{"review_id": "NUzEi5CVvZ", "reviewer": "Reviewer_jx3E", "summary": "The paper presents an approach for low-rank optimal transport (OT) leveraging a latent coupling (LC) factorization and solving it with mirror descent. This approach offers a new parameterization of the low-rank OT problem, providing advantages such as decoupling the problem into three OT problems and enhancing flexibility and interpretability. The authors introduce the Factor Relaxation with Latent Coupling (FRLC) algorithm, which utilizes coordinate mirror descent to compute the LC factorization. FRLC accommodates various OT objectives and marginal constraints. The paper includes theoretical results and demonstrates the performance on selected tasks like graph clustering and spatial transcriptomics.", "questions": "- Is the main contribution here actually the mirror descent formulation of the OT problem with the factorization with a different set of latent source and target points proposed by Lin et al. (2021)?\n- Furthermore, is the main difference between the fact that Lin et al. (2021) use a two-stage process and the novel mirror-descent method solves the OT problem in a more straightforward unified process? \n- What are the theoretical complexity of the algorithms, and what practical considerations are relevant when implementing the method?\n\nOne presentation improvement that would change my evaluation is a summary table locating the differences in the type of factorization, whether the cost or transport is being factorized, the algorithm used for computing the factorization, etc.", "limitations": "I found it necessary to discuss in more depth the limitations and the tradeoffs present in the proposal compared with previous art.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Generality: handles different OT objective costs and relaxations of the marginal constraints.\n- Significance: the proposed method, and improving OT using factorizations, has the potential to address the scalability issue of OT with large datasets, making it highly relevant for applications in machine learning and data science.", "weaknesses": "**Presentation and logical organization of the ideas**\n\nI found it difficult to discern the paper's particular contribution, especially concerning previous art like Latent Optimal Transport (Lin et al., 2021), Forrow et al. (2019), and Scetbon et al. (2021). For example, the main contribution is said to be \"compute a minimum cost low-rank transport plan using the LC factorization,\" where LC factorization is precisely the factorizations proposed by Lin et al. (2021), with the further contribution being the computation of this factorization using mirror-descent. I wondered how much of the description also applies to Lin et al. (2021) and which part is particular to the novel method. \n\nFurthermore, according to Lin et al. (2021), there is an algorithm for computing low-rank factorized plans based on factorized costs using a projection method and an algorithm to factorize the costs and select anchor points without the need to have a similar number of anchor points at source and target latent space. How does your approach compare with the approach? \n\n**Comparison and empirical validation**\nThe first empirical comparison analyzes LOT and FRLC, capturing how both methods behave as the factorization rank changes. This is a good way of understanding the effect of the rank on each case. Nevertheless, in the second setup, the comparison with LOT is not present anymore, which begs the question of how LOT performs in the second empirical evaluation. Furthermore, an extensive empirical study with the structured approach of Lin et al. (2021) would help clarify the particular contributions of the type of factorization and factoring algorithm used. \n\nI found some other issues lacking in a more in-depth discussion of the computational complexity, possible implementation difficulties, and overall limitations in the proposal.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer jx3E for their careful reading of our work, and for their feedback. We have abbreviated your questions due to space limitations for our responses.\n\n> [I found ... novel method.]\n\nThank you for this comment indicating an area where we can strengthen our presentation. Our FRLC algorithm differs from [Lin et al. (2021)](https://proceedings.mlr.press/v139/lin21a.html) in several respects. We omit them here due to space, but ask you to refer to the bulleted list of differences in our general response. We will include these points with additional details in our Appendix to clarify the differences between the approaches.\n\n> [Furthermore ... the approach?]\n\nWe hope our general response addresses most of this question.\n\nAdditionally, we want to emphasize that Lin et al. (2021) explicitly compute anchor points in each iteration of their algorithm, whereas FRLC has no anchor points. In our evaluation in Section 4.2, we computed LC projections from the output of FRLC on each dataset in order to demonstrate the interpretability of the latent coupling. These LC projections may be interpreted as the analogs of the anchor points used in Lin et al. (2021), but importantly the two concepts are not identical. Even if we specialize to using Euclidean distance, we are unaware of a way to write FRLC cost matrices as functions of the anchor-like points obtained from LC projection. \n\n> [The first ... empirical evaluation.]\n\nFirst, to clarify, in our experiments section (Section 4), LOT/UL/SRL all refer to the low-rank methods of Scetbon et al. ([Scetbon et al. (2021)](http://proceedings.mlr.press/v139/scetbon21a.html), [Scetbon et al. (2022b)](https://proceedings.mlr.press/v162/scetbon22b), [Scetbon et al. (2023)](https://openreview.net/forum?id=d2WsCmoITF)), and not to the Lin et al. (2021) method which coincidentally is also named LOT. UL and SRL denote the unbalanced and semi-relaxed versions (from Scetbon et al. (2023)) of Scetbon's LOT method Scetbon et al. (2021), respectively. We understand how this terminology can cause confusion, and we will clearly distinguish the method names in our revision. To avoid further confusion in this response (as in our paper), we use LOT to reference the work of Scetbon et al., and write \"Lin et al.\" to refer to Lin et al.'s Latent OT.  \n\nWe assume that \"second setup\" refers to Section 4.3, where we compare against UL and SRL rather than LOT. As noted above, UL and SRL are the unbalanced and semi-relaxed variants of LOT, respectively. We note that we also included a comparison between the balanced versions of FRLC and LOT for this dataset (Figure 8 in our Appendix). Thus, all of our comparisons -- including Section 4.1 (Figure 2) and Section 4.2 (Figure 3) -- are directly head-to-head with LOT (with the exception of Appendix K where we compare against the full-rank solvers of `ott-jax`).\n\n> [Furthermore ... algorithm used.]\n\nWe ran the Lin et al. (2021) method on several synthetic datasets and report the OT-cost (low-rank Wasserstein cost) found by this method and FRLC in Table 1 of our general response.\n\n> [I found ... the proposal.]\n\nTo address the first point, we offer a brief complexity analysis in the general response above.\nRegarding implementation difficulties and limitations, we included a discussion on limitations at the end of our Appendix, and  will add the sentences indicated in our general response to our Conclusion. \n\n> [Is the main ...  Lin et al. (2021)?]\n\nNot quite -- as we noted above the FRLC algorithm does not have latent points (anchor points) at all, but rather optimizes directly over the sub-couplings $Q$, $R$, and $T$. This has several advantages as noted above, including that the steps in our algorithm are themselves OT sub-problems, unlike Forrow et al. (2019), Scetbon et al. (2021), and Lin et al. (2021) which use Dykstra projections. Moreover, FRLC is the first algorithm which solves this factorization for general costs $C$ and other objectives (GW, FGW). We refer to the bullet points above in the general response for more discussion on the differences and contributions.\n\n> [Furthermore ... unified process?]\n\nYes, we agree that FRLC solves the OT problem in a more straightforward and unified process: we only optimize the sub-couplings, while Lin et al. (2021) alternate updates on the anchor points and the sub-couplings. However, this is only one of multiple differences between the methods.\n\n> [What are ... method?]\n\nWe give a brief analysis of the time complexity of FRLC in our general response above, and will include this in our main text and Appendix. Regarding practical considerations in implementing FRLC, in practice taking the hyper-parameter $\\tau$ to be close to $0$ approaches a theoretically optimal step on the sub-couplings while large $\\tau$ regularizes their trajectory. Thus, it is an open question how to choose this hyperparameter, as discussed in the limitations component of our general response, and likewise with the choice of dimensions for the latent coupling. \n\n> [One presentation ... etc.]\n\nThis is a great suggestion -- we include a table here, and will add this to our Appendix.\n\n| Method | Factorization  | Cost | Variables | Algorithm  | Subroutine for sub-couplings |\n|-|-|-|-|-|-|\n| Factored Coupling  (Forrow et al. (2019)) | Diagonal | $k$-Wasserstein barycenter| Anchors \\& sub-couplings| Lloyd-type | Dykstra |\n| Latent OT (Lin et al. (2021)) | Non-diagonal | Extension of $k$-Wasserstein barycenter| Anchors \\& sub-couplings| Lloyd-type | Dykstra |\n| LOT (Scetbon et al. (2021)) | Diagonal | Primal OT cost | Sub-couplings | Mirror-descent | Dykstra                |\n| FRLC (our work) | Non-diagonal | Primal OT cost | Sub-couplings | Coordinate mirror-descent | OT                       |\n\n**Table:** Comparison of latent and low-rank OT methods.\n\n> [I found ... previous art.]\n\nWe included a discussion on limitations in Appendix N but will now add the sentence given above regarding limitations to the main text.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "52os6NT0X1", "reviewer": "Reviewer_pnTm", "summary": "The paper introduces a novel framework called Factor Relaxation with Latent Coupling (FRLC) which is based on coordinate mirror descent to compute the low-rank LC factorization.The algorithm decouples the optimization into three sub-problems, offering greater flexibility, interpretability, and linear space complexity. FRLC is applicable to various OT objectives (Wasserstein, Gromov-Wasserstein, Fused Gromov-Wasserstein) and marginal constraints (balanced, unbalanced, and semi-relaxed). Theoretical results support its effectiveness, and empirical tests demonstrate superior performance in applications like graph clustering and spatial transcriptomics.", "questions": "I suggest that Propositions E.3, E.5, E.6 should be stated in the main text to make the theoretical part easier to follow. \n\nTable 1 should be put in the main text to show the efficiency of FRLC in terms of runtime. It is hard to say that the method shows significant improvement because the runtime of LOT is already very short (less than one second). I suggest the authors conduct experiments on more challenging and diverse examples.", "limitations": "N.A.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The introduction of the LC factorization and its integration into the FRLC algorithm represents an innovation in reducing the complexity of optimal transport problems.\n\nThe FRLC framework is versatile, handling multiple OT objectives and marginal constraints, making it applicable to a wide range of practical problems.\n\nThe latent coupling provides a interpretable description of the transport plan, which is beneficial for understanding and visualizing the results.", "weaknesses": "While the empirical results are good, the paper could benefit from more extensive comparisons with additional baseline methods to solidify the claims of superiority.\n\nThe theoretical results are limited in that, even though $\\Delta_k(x_k, x_{k+1})$ is small, it cannot be guaranteed that the iteration points converge to the optimal solution or even a stationary point.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer pnTm for their careful reading of our work, and for their feedback.\n\n> While the empirical results are good, the paper could benefit from more extensive comparisons with additional baseline methods to solidify the claims of superiority.\n\nThe only OT methods solving for a low-rank plan that we are aware of are Factored Coupling [Forrow et al. (2019)](http://proceedings.mlr.press/v89/forrow19a), LOT ([Scetbon et al. (2021)](http://proceedings.mlr.press/v139/scetbon21a.html), [Scetbon et al. (2022b)](https://proceedings.mlr.press/v162/scetbon22b), [Scetbon et al. (2023)](https://openreview.net/forum?id=d2WsCmoITF)), Latent OT [Lin et al. (2021)](https://proceedings.mlr.press/v139/lin21a.html). We did not compare to Forrow et al. (2019) because a prior work of Scetbon et al. (2021) established LOT as the current state of the art, and we did not compare to Lin et al. because they do not minimize the primal OT cost directly. However, we now added a comparison to Lin on the balanced Wasserstein objective cost, and we will add Table 1 in our above general response to our Appendix.\n\nImportantly, Lin et al. solve what they term the latent OT problem, which is a different objective than the low-rank Wasserstein objective computed by FRLC.  Thus we did not feel it was appropriate to include this comparison with the others in the main text, but will include as part of a new addition to clarify the differences betweeen FRLC and Lin et al.\n\nIn Section 4.4 (additional experiments) we refer to Section K of our Appendix where we also evaluated against full-rank OT solvers (from `ott-jax`) on downstream metrics for 12 graph partitioning datasets.\n\n> The theoretical results are limited in that, even though $\\Delta _ k(x _ k, x _ {k-1})$ is small, it cannot be guaranteed that the iteration points converge to the optimal solution or even a stationary point.\n\nWe agree. In our convergence results, we adapted the criterion of [Ghadimi et al. (2014)](https://link.springer.com/article/10.1007/s10107-014-0846-1) to coordinate-mirror descent to show non-asymptotic stationary convergence, as in other works on low-rank optimal transport [Scetbon et al. (2021)](http://proceedings.mlr.press/v139/scetbon21a.html), [Scetbon et al. (2022b)](https://proceedings.mlr.press/v162/scetbon22b). Such a result is the current baseline for convergence in the literature, but we acknowledge that the convergence criterion and guarantee is limited. A more thorough theoretical analysis with stronger guarantees might be difficult as the objective is non-convex, but would be valuable to analyze and understand more comprehensively in future work.\nWe mentioned this limitation in the Appendix, but in the revision we will add the following sentence to the Conclusion (Section 5) to more clearly indicate the limitations: \n\n_\"Two key limitations of our work are: (1) selecting values of the latent coupling ranks and the hyperparameter $\\tau$ controlling the smoothness of the trajectory; and (2) strengthening the convergence criterion. These and other limitations are discussed in Section N of the Appendix.\"_\n\n> I suggest that Propositions E.3, E.5, E.6 should be stated in the main text to make the theoretical part easier to follow.\n\nAbsolutely, we had removed them from the main text due to the page limit, but will include them in the updated version of the manuscript.\n\n> Table 1 should be put in the main text to show the efficiency of FRLC in terms of runtime. It is hard to say that the method shows significant improvement because the runtime of LOT is already very short (less than one second). I suggest the authors conduct experiments on more challenging and diverse examples.\n\nWe will move this table to the main text. We also performed additional comparisons on the large spatial transcriptomics datasets of [Chen et al. (2022)](https://www.sciencedirect.com/science/article/pii/S0092867422003993), which we include in Table 2 of our general response above.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "C4bsPp4JuB", "reviewer": "Reviewer_CGKE", "summary": "This work introduces a new low rank formulation of optimal transport based on the latent coupling decomposition introduced in https://arxiv.org/pdf/2012.11589 . Compared to previous formulations of low rank OT, this formulation allows easier extensions to unbalanced and Gromov Wasserstein settings.", "questions": "* How to choose $r$ in practical scenarios ? \n* Would it make sense to remove the entropic penalties (at the expense of smoothness) to recover results about hard clustering of the points via the decomposition of the OT coupling ? Put differently, what is the effect of entropic smoothing on the interpretability of the coupling in terms of clustering ?", "limitations": "/", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I found that the presented FRLC method is a quite elegant approach to low rank OT. It decouples optimization into entropy regularized semi relaxed OT problems for $Q$ and $R$, and $T$ is updated solving entropy regularized  OT of size $r$. This decoupling is enabled by the form of the inner coupling which is no longer assumed to be diagonal. This construction alleviates the need to use the Dykstra machinery as in (Scetbon 2021).\nWith the proposed approach, $Q$ and $R$ are not required to have the same inner marginal. From what I understand, the optimization over inner marginals is handled by the semi relaxed steps and does not require a separate step as in previous approaches. \n\nAnother strength of the approach is that it generalizes smoothly to the unbalanced setting. \n\nOther strengths include:\n* I carefully looked at the presented method and the appendix and the algorithmic approach made sense. \n* non asymptoptic convergence bounds are derived for the method. \n* the appendix includes a complete review of the necessary technical background presented in a clear way, easy to follow.\n* limitations are clearly stated (section N in appendix).\n* a new initialization scheme for LC decomposition is proposed. \n* there is a gain in interpretability compared to (Scetbon 2021) as the inner coupling captures the coupling between clusters in the data.", "weaknesses": "* from section 3.1 to 3.2, the new aspects of the approach compared to https://arxiv.org/pdf/2012.11589 should be stated more clearly. The contributions of the paper should be separated more clearly from previous works. I had to read the other paper to better understand each contribution. \n* In algorithm 1, it would be interesting to have the computational complexity of each step (presented in the algorithm directly).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer CGKE for their careful reading of our work, and for their feedback.\n\n> from section 3.1 to 3.2, the new aspects of the approach compared to [https://arxiv.org/pdf/2012.11589](https://arxiv.org/pdf/2012.11589) should be stated more clearly. The contributions of the paper should be separated more clearly from previous works. I had to read the other paper to better understand each contribution.\n\nWe agree that we could have been more clear about how our method is distinct from the approach of Lin et al. (2021). We greatly appreciate your comment, and we refer you to the bulleted list in our general response highlighting the differences between our work and that of Lin et al. (2021). We will add a section in our Appendix contrasting the approaches in more detail.\n\n> In algorithm 1, it would be interesting to have the computational complexity of each step (presented in the algorithm directly).\n\nWe thank you for your comment, and will make sure to include a discussion of time-complexity in the appendix and in-line with our algorithm in the main text. We refer you to the time-complexity analysis of FRLC included in our general response above.\n\n> How to choose $r$ in practical scenarios ?\n\nThis is a great question. The best we can say regarding the rank is that the choice is problem-specific, and is much like choosing $k$ in k-means or choosing the rank $r$ for a compact SVD (or PCA). We added the following sentence to the Conclusion (section 5) to directly address this limitation: \n\n_\"Two key limitations of our work are: (1) selecting values of the latent coupling ranks and the hyperparameter $\\tau$ controlling the smoothness of the trajectory; and (2) strengthening the convergence criterion. These and other limitations are discussed in Section N of the Appendix.\"_\n\n> Would it make sense to remove the entropic penalties (at the expense of smoothness) to recover results about hard clustering of the points via the decomposition of the OT coupling? Put differently, what is the effect of entropic smoothing on the interpretability of the coupling in terms of clustering?\n\nThis is a fantastic question. \n\nScetbon et al. [Scetbon et al. (2022a)](https://proceedings.neurips.cc/paper_files/paper/2022/hash/2d69e771d9f274f7c624198ea74f5b98-Abstract-Conference.html) show in their Proposition 9 that removing the entropic penalties and setting the second marginal to equal the first generalizes the $k$-means objective. [Lin et al. (2021)](https://proceedings.mlr.press/v139/lin21a.html) show that setting the distance between the anchors ($C _ {z}$) to zero also recovers $k$-means. \nHowever, because FRLC does not use Scetbon et al.'s factorization, and also does not evaluate a distance $C _ {z}$ between anchors, we do not know if removing the entropic penalties in FRLC would recover a known hard clustering (or co-clustering) method. We will add the following sentence to our Conclusion section discussing future work: \n\n_\"Another interesting question is whether a relationship exists between the minimization of the Wasserstein cost over LC-factorizations without entropic regularization and calculation of a hard clustering or co-clustering between data points, analogous to Proposition 9, [Scetbon et al. (2022a)](https://proceedings.neurips.cc/paper_files/paper/2022/hash/2d69e771d9f274f7c624198ea74f5b98-Abstract-Conference.html).\"_", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OHtYnslxF3", "reviewer": "Reviewer_1xw8", "summary": "As the paper clearly outlines in the introduction, Optimal Transportation (OT) is widely used in various fields of machine learning, however, the cost for computing OT is computationaly expensive (quadratic scaling), even after Sinkhorn algorithm employing entropy regularization significantly alleviated the cost. The purpose of Low-rank Optimal Transport (LOT) is to reduce the complexity by constraining the transportation plan to possess a certain low-rank structure.\n\nThe key contribution of the paper is (1) new parametrization of LOT using a Latent Coupling (LC) factorization, (2) The eponymous \"Balance FRLC\" algorithm for resolving the optimization procedure for computing LOT with respect to various objectives associated with (LOT), (3) Theoretical guarantees for the smoothness of the regularized objectives, as well as upper bounds for the convergence criteria.  The algorithm is applied to illustrative toy examples and then to a more realistic example (spatial transcriptomics alignment) and the results indicate improvements over previous LOT.", "questions": "- line 103: Would it be good to say the what complexity \"massively improved\" to?\n- Can the authors comment on how to pick the rank for the LOT? (This of course ties in to the point above the rank of the transportation plan.)\n- Perhaps a more discussion of of how different the LC factorization is to, say the previous factored coupling with (the diagonal version $P = Q \\text{diag}(1/g) R^T$). This relation seems important when considering \"more\" LC factorizations: If this LC factorization improves the result while simultaneously allowing a simple optimization algorithm, would it then be natural to discuss further factorizations that are nested this way, with additional LC factorizations for $T$ and so on?", "limitations": "This work builds on the LOT framework - so perhaps the authors comment on some scenarios where the LC factorization can potentially yield no better results, when compared to previous LOT?", "rating": 9, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The contributions of the paper is significant. The writing is excellent. The LC factorization, adapted to the LOT framework from a previous work, is well-motivated. The Balanced FRLC algorithm appears a sound and non-trivial extension of previous approaches for LOT (Sinkhorn, semi-relaxed OT). The algorithm is particularly simple and shown to be effective. The experimental results are clear and illustrative, improvements are convincing.", "weaknesses": "What this reviewer found missing is discussion of the case the optimal transport map is itself not low-rank. It is well-known that the optimal transportation plan is not necessarily low-rank (as indicated e.g. by their connection to differential equations with non-regular solutions), so LOT may converge very slowly to the full-rank optimal transportation plan. In section 4.2, the authors provide a simple illustrative example, but this is when the transportation plan is known to be low-rank - in this case the transportation plan seems to stabilize after some threshold rank - however, what happens when the full-rank plan does not have an inherently low-rank structure?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer 1xw8 for their careful reading of our work, and for their feedback.\n\n> What this reviewer found missing is discussion of the case the optimal transport map is itself not low-rank. It is well-known that the optimal transportation plan is not necessarily low-rank (as indicated e.g. by their connection to differential equations with non-regular solutions), so LOT may converge very slowly to the full-rank optimal transportation plan. In section 4.2, the authors provide a simple illustrative example, but this is when the transportation plan is known to be low-rank - in this case the transportation plan seems to stabilize after some threshold rank - however, what happens when the full-rank plan does not have an inherently low-rank structure?\n\nThis is a great question. We are actively investigating this, but do not have a complete answer during this short rebuttal period. We will include this point in the Conclusion of our main text: \n\n_\"Another direction for further investigation is to better understand what structure LC factorizations capture when the optimal plan is known to have full rank, e.g. when the Monge map exists as has been explored by works such as [Liu et al. (2021)](https://arxiv.org/abs/2111.06546).\"_\n\n> line 103: Would it be good to say the what complexity \"massively improved\" to?\n\nYes, thank you for the suggestion. For a fixed error tolerance $\\varepsilon$ and entropy constant $\\eta$, FRLC finds a $\\pm \\varepsilon D$ approximation for $D$ the diameter of the dataset in $B = \\mathrm{poly}(1/(\\eta \\varepsilon))$ iterations which is a major improvement over the previous cubic solvers. \n\n> Can the authors comment on how to pick the rank for the LOT? (This of course ties in to the point above the rank of the transportation plan.)\n\nThank you for your question. The best we can say regarding the rank is that the choice is problem-specific, and is much like choosing $k$ in k-means or choosing the rank $r$ for a compact SVD (or PCA). We will add the following sentence to the Conclusion (Section 5) to directly address this limitation: \n\n_\"Two key limitations of our work are: (1) selecting values of the latent coupling ranks and the hyperparameter $\\tau$ controlling the smoothness of the trajectory; and (2) strengthening the convergence criterion. These and other limitations are discussed in Section N of the Appendix.\"_\n\n> Perhaps a more discussion of of how different the LC factorization is to, say the previous factored coupling with (the diagonal version $P=Q\\mathrm{diag}(1/g)R^\\mathrm{T}$). This relation seems important when considering \"more\" LC factorizations: If this LC factorization improves the result while simultaneously allowing a simple optimization algorithm, would it then be natural to discuss further factorizations that are nested this way, with additional LC factorizations for  and so on? \n\nWe agree that one strength of the LC factorization is that it can be applied again to decompose a latent coupling, leading to a multi-scale family of optimal transport problems similar to [Gerber et al. (2017)](https://dl.acm.org/doi/abs/10.5555/3122009.3176816). We leave a thorough exploration of this to future work, but will will highlight the possibility of nesting LC factorizations in the main text -- thank you.\n\n> This work builds on the LOT framework - so perhaps the authors comment on some scenarios where the LC factorization can potentially yield no better results, when compared to previous LOT?\n\nFRLC generally has better performance, especially on large and structured datasets such as spatial transcriptomics, but the difference in performance is insignificant for ultra low rank (e.g. $r _ 1=2$) on synthetic datasets. There are also cases where one might correctly assume that transitions are diagonal between clusters -- in these cases, it might be better to use the previous factorization, which implicitly clusters the two datasets with a common set of labels. We note that we can recover the previous LOT factorization: by taking ${Q}' \\gets {Q}\\mathrm{diag}(1/{g} _ {Q}){T}$ one may diagonalize the return as $({Q},{R},{T}) \\to ({Q}',{R}, \\mathrm{diag}({g} _ {R}))$, with no change to the OT cost.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Here we include information for all reviewers. We thank each reviewer for their helpful feedback and questions.

> Regarding time complexity raised by reviewers jx3E, CGKE:

The time complexity of FRLC $O(BLr^{2}(n+m))$, where $B$ is the number of Sinkhorn iterations, $L$ the number of mirror descent iterations, $n$ is the number of samples in the first dataset, $m$ is the number of samples in the second dataset, and $r$ is both the rank of the latent coupling and the rank of the distance matrix $C$. The time complexity is thus linear in $n+m$ per iteration. The number $L$ of outer iterations follows from the convergence rate of coordinate mirror descent. The number of iterations required for each projection follow from the convergence of Sinkhorn where, for $\varepsilon$ a fixed error tolerance and $\eta$ the entropy constant, one finds a $\pm \varepsilon D$ approximation for $D$ the diameter of the data in $B = \mathrm{poly}(1/(\eta \varepsilon))$ iterations. Each step of the projection is $O((n+m)r^{2})$ assuming a rank-$r$ factorization of $C$, allowing every matrix multiplication to have one dimension on the order of the constant rank $r$. Therefore, the runtime is $O(B L r^2 (n+m))$.

> Regarding comparisons to Lin et al. (2021) brought up by jx3E, CGKE, we benchmarked FRLC against their method on the Wasserstein objective across four synthetic datasets:

| Dataset | OT-cost (FRLC) | OT-cost (Lin et al.) |
|-|-|-|
| 5th and 10th roots of unity (rank $r _ {1},r _ {2}=5,10$) | 1.174 | 2.124 |
| Two-moons and 8-Gaussians (rank $r=20$) | 2.716 | 4.291 |
| 2D Gaussian mixture (rank $r=20$)  | 0.552 | 0.922 |
| 10D Gaussian mixture (rank $r=20$) | 1.038 | 1.298 |

**Table 1:** Comparison against Lin et al. (2021) in primal OT-cost $\langle {C}, {P} \rangle _ {F}$

The table shows that FRLC achieves a lower cost on all datasets. Importantly, Lin et al. solve what they term the latent OT problem, which is a different objective than the low-rank Wasserstein objective minimized by FRLC. Thus we did not feel it was appropriate to include this comparison in the main text, but will include as part of a new addition to our Appendix to clarify the differences between FRLC and Lin et al.

> We also report our runtime on more challenging datasets, large spatial transcriptomics datasets of [Chen et al. (2022)](https://www.sciencedirect.com/science/article/pii/S0092867422003993), as requested by reviewer pnTm:

| Dataset | LOT (seconds) | FRLC (seconds) | OT cost (LOT) | OT Cost (FRLC) |
|-|-|-|-|-|
| Mouse embryo (E9.5-10.5) | 2.545 | 1.112 | 0.4396 | 0.385 |
| Mouse embryo (E10.5-11.5) | 4.209 | 1.190 | 0.3714 | 0.344 |
| Mouse embryo (E11.5-12.5) | 8.667 | 1.889 | 0.478 | 0.439 |

**Table 2:** Comparison of methods on MOSTA Stereo-Seq mouse embryo spatial transcriptomics datasets (GPU). This is using default settings with `min_iter=10`, `max_iter=100`, rank $r=50$. LOT denotes the low-rank method of Scetbon et al. (2021), as a point of comparison.

> In response to several reviewer questions about hyperparameter selection and our convergence criterion, we will add the following sentence to our Conclusion section:

_"Two key limitations of our work are: (1) selecting values of the latent coupling ranks and the hyperparameter $\tau$ controlling the smoothness of the trajectory; and (2) strengthening the convergence criterion. These and other limitations are discussed in Section N of the Appendix."_

> Regarding questions raised by reviewers jx3E, CGKE on how FRLC differs from the method of Lin et al. (2021), here is a short summary of the differences:

* Lin et al. (2021) optimize two sets of variables: sub-couplings $(Q, R, T)$ and anchor points $(z^x, z^y)$. FRLC only has $(Q,R,T)$ as variables of the optimization.
* Cost matrices used in Lin et al. (2021) are built from distances between each dataset and its representative anchor points for $Q, R$, or the distances between the two sets of anchor points for $T$. In contrast, ground costs used in FRLC to update $(Q,R,T)$ are always derived from the distance matrix $C$ in the Wasserstein objective $\langle C, P\rangle _ F$.
    * Specifically, the cost matrices used by Lin et al. (2021) are:
    $$
    [C _ Q^{\mathrm{Lin}}]_{ik} = \Vert x _ i - z _ k^x \Vert _ 2^2, \quad [C _ R^{\mathrm{Lin}}] _ {j\ell} = \Vert y _ j - z _ \ell^y \Vert _ 2^2, \quad [C _ T^{\mathrm{Lin}}] _ {k\ell} = \Vert z _ k^x - z _ \ell^y \Vert _ 2^2,
    $$
    optionally using a Wasserstein distance for the entries of $C_T^{\mathrm{Lin}}$
    * FRLC costs are given in the exponents of the Gibbs kernels written above and below equation (9) in our paper. These are derived directly from the rank-$r$ Wasserstein problem $\min _ {P \in \Pi _ r(a,b)} \langle C, P \rangle _ {F}$ and differ substantially from those of the proxy objective used in Lin et al. (2021).
* The different objectives and variables lead to different algorithms: Lin et al. (2021) alternate updates to the sub-couplings $(Q,R,T)$ using Dykstra, with updates to the latent anchor points $(z^{x}, z^{y})$ using first-order conditions. In contrast, FRLC alternates semi-relaxed OT to update $(Q,R)$ and balanced OT to update $T$.
* Moreover, because FRLC does not require anchor points to define costs, FRLC can handle cost matrices which are not simple functions of distance. For example, if $C _ {ij}$ is the price of transporting good $i$ to warehouse $j$ one may not be able to re-evaluate a price $c(x _ {i}, z _ {k}^x)$ between $x _ i$ and latent anchor $z _ {k}^x$. In such situations, while finding a low-rank plan may make sense (e.g. to approximate an assignment for a massive dataset), an "anchor" may not have clear definition in the setting of general cost matrices.
* The Lin et al. (2021) objective is only a proxy for a Wasserstein-type loss, and Lin et al. do not explore extensions to Gromov-Wasserstein (GW), or Fused GW, which FRLC readily generalizes to.

Meta-review: The paper presents a new numerical approach for low-rank optimal transport, introducing a novel mirror descent algorithm. The numerical results demonstrate the method's ability to address scalability issues with large datasets in optimal transport. The authors made a significant effort in the rebuttal, providing additional comparisons to previous works. However, the exposition could be improved to more clearly highlight the key novelties compared to existing low-rank optimal transport methods, particularly those presented by Scetbon et al. (2021). Additionally, the authors suggest that the factorization using (Q, R, g) variables was introduced in Forrow et al. (2019), but this formulation was introduced in Scetbon et al. (2021) (the formulation in Forrow et al. relies on an intermediate barycenter measure of vectors, and works with squared Euclidean costs); this should be clarified in the final version. I recommend acceptance.
Decision: Accept (poster)