Low-Rank Optimal Transport through Factor Relaxation with Latent Coupling

Authors: peter  halmos
xinhao  liu
julian  gold
benjamin j raphael
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 21:05:46

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.

Main paper: 
 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 to a ground cost C. The minimum cost yields a distance between a and b, while the optimal transport plan reveals key structural similarities between the distributions. Owing to its versatilitydifferent 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 et al. (2023) and domain adaptation Courty et al. (2014) ; Solomon et al. (2015) to 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-Wasserstein (GW) OT Mémoli (2011) compares distributions supported on distinct geometries through the expected metric distortion of P . Fused Gromov-Wasserstein (FGW) Vayer et al. (2020) OT is suited to structured data, taking a convex combination of the former two objectives. Independently, one can relax constraints on the marginals of P : in computational applications, P is a matrix whose row-sum P 1 m and column-sum P T 1 n are called its left and right marginals. Balanced OT requires P 1 m = a and P T 1 n = b. Unbalanced OT Frogner et al. (2015) replaces these constraints with penalties in the 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. (2024) .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) introduced another variant of OT, optimizing a 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. (2021) introduce the factorization P = Q diag (1/g)R T comprised of sub-coupling matrices Q and R sharing an inner marginal g, meaning Q T 1 n = R T 1 m = g. Building on this, Scetbon et al. (2021 Scetbon et al. ( , 2022 Scetbon et al. ( , 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 Lin et al. (2021) in the context of k-Wasserstein barycenters. We call their factorization a latent coupling (LC) factorization, given by P = Qdiag(1/g Q )T diag(1/g R )R T , with two inner marginals g Q = Q T 1 n and g R = R T 1 n and a general coupling T . Lin et al. (2021) constrain the transport between a 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 of Forrow et al. (2019) ; Scetbon et al. (2021) , who seek a minimal rank coupling with respect to a single, fixed cost C. We observe that factored couplings of Forrow et al. (2019) correspond to LC factorizations with diagonal T , suggesting the LC factorization of Lin et al. (2021) may provide an alternative parameterization of transport plans for the low-rank OT problem considered in Forrow et al. (2019) ; Scetbon et al. (2021) . 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 a 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 to 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 Scetbon 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 of Forrow et al. (2019) . FRLC computes the LC factorization using a novel coordinate mirror descent scheme, alternating descent steps on variables (Q, R) and T , inspired by the mirror descent approach of Scetbon et al. (2021) . 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, Gromov-Wasserstein, Fused Gromov-Wasserstein), and marginal constraints (balanced, unbalanced, and semi-relaxed). We show FRLC performs better than existing state-of-the-art low-rank methods on a range of synthetic and real datasets, retaining the interpretability of Lin et al. (2021) , and inheriting the broad applicability of Scetbon et al. (2021) ; Scetbon & Cuturi (2022) ; Scetbon et al. (2022 Scetbon et al. ( , 2023)) .2 Background Wasserstein OT. Let {x 1 , . . . , x n } and {y 1 , . . . y m } be datasets in a metric space X , and let ∆ d be the probability simplex of size d. Through probability vectors a ∈ ∆ n and b ∈ ∆ m , each dataset is encoded as a probability measure: µ = Thus, Π a,b is the set of transport plans (probabilistic coupling matrices) with marginals a and b. Given a cost function c : X × X → R + , define the cost matrix C ∈ R n×m + via C ij = c(x i , y j ). The Kantorovich formulation Kantorovich (1942) of discrete OT, also called the Wasserstein problem, seeks a transport plan P of minimal cost : W(µ, ν) := min P ∈Π a,b ⟨C, P ⟩ F .(1)Gromov-Wasserstein OT. In many applications, one wishes to compare datasets {x 1 , . . . , x n } ⊂ X and {y 1 , . . . , y m } ⊂ Y across distinct metric spaces X and Y. The Gromov-Wasserstein (GW) objective Mémoli (2007 Mémoli ( , 2011) ) addresses the absence of a common metric or coordinate system through intra-domain cost functions c 1 : X × X → R + and c 2 : Y × Y → R + , leading to intra-domain cost matrices A ik = c 1 (x i , x k ) and B jl = c 2 (y j , y l ). The GW objective function Q A,B (P ) := i,j,k,l (A ik -B jl ) 2 P ij P kl quantifies the expected metric distortion under P , leading to the optimization problem: GW(µ, ν) := min this factorization is the presence of a coupling matrix T linking two distinct inner marginals. For simplicity we describe this factorization using an r-dimensional latent space, but we also extend to non-square matrices linking two latent spaces of different dimensions, as demonstrated in the results. Definition 3.1 (Inner marginals). Given a factorization P = QXR T of a coupling matrix P ∈ Π a,b (r), the inner marginals of Q and R are g Q := Q T 1 n and g R := R T 1 m , respectively, where g Q , g R ∈ ∆ r .To distinguish the different marginals, we refer to a and b as outer marginals. Definition 3.2 (LC factorization). Given a coupling matrix P ∈ Π a,b (r), a latent coupling (LC) factorization of P is P = Qdiag(1/g Q )T diag(1/g R )R T , where g Q and g R are the inner marginals of Q and R, Q ∈ Π a,• , R ∈ Π b,• , and T ∈ Π g Q ,g R .We {(Q, R, T ) ∈ R + : Q ∈ Π a,• , R ∈ Π b,• , T ∈ Π g Q ,g R },be the set of admissible sub-couplings for the LC factorization. Definition 3.2 gives the following map from LC a,b (r) to Π a,b (r):EQUATIONSince this map is surjective, the set LC a,b (r) parameterizes Π a,b (r). Surjectivity follows from the fact that FC a,b (r) maps injectively into LC a,b (r), through (Q, R, g) → (Q, R, diag(g)), and FC a,b (r) maps surjectively onto Π a,b (r) via (Q, R, g) → Qdiag(1/g)R T . Definition 3.1 and Definition 3.2 are readily extended in two directions: the case when the outer marginal constraints are relaxed such thatQ ∈ R n×m + or R ∈ R n×m +, while maintaining the constraint that T ∈ Π g Q ,g R ; as well as the case of non-square T . 
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 to other and marginal constraints are discussed later. The FRLC objective function, for low-rank, balanced Wasserstein OT, isEQUATIONwhere P (Q,R,T ) is defined by (7) . Since LC a,b (r) parameterizes Π a,b (r), problem ( 8) is equivalent to 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 ) ∈ LC a,b (r), sub-couplings Q and R are constrained by:C 1 (a) := {(Q, R, T ) ∈ R + : Q1 r = a}, C 1 (b) := {(Q, R, T ) ∈ R + : R1 r = b}.The convex sets constraining the latent coupling matrix T areC 2 (g Q ) := {(Q, R, T ) ∈ R + : T 1 r = g Q }, C 2 (g R ) := {(Q, R, T ) ∈ R + : T T 1 r = g R },where g Q = Q T 1 n and g R = R T 1 m as per Definition 3.1. Writing C 1 = C 1 (a) ∩ C 1 (b) andC 2 = C 2 (g Q ) ∩ C 2 (g R ), one has LC a,b (r) = C 1 ∩ C 2 .We use coordinate mirror descent to optimize (8), building on the mirror descent (MD) approach of Scetbon et al. (2021) ; Scetbon & Cuturi (2022) ; Scetbon et al. (2022 Scetbon et al. ( , 2023) ) for the low-rank problem.First we take a 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 g Q and g R . 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 1. Let (γ k ) N k=1 be a sequence of step sizes. As in Scetbon & Cuturi (2022) , we choose ℓ ∞ -normalization for the step-sizes. Our coordinate mirror descent in the factor relaxation step is:(Q k+1 , R k+1 ) ← arg min (Q,R) : (Q,R,T k )∈C1 ⟨(Q, R), ∇ (Q,R) L LC ⟩ + 1 γ k KL((Q, R)∥(Q k , R k )) + τ KL((Q T 1 n , R T 1 m )∥(Q T k 1 n , R T k 1 m ))The Sinkhorn kernels for the semi-relaxed OT problems arising from the factor relaxation step are:K (k) Q := Q k ⊙ exp( -γ k (CR k X T k -1 n diag -1 ((CR k X T k ) T Q k diag(1/g Q k )) T ) ) K (k) R := R k ⊙ exp( -γ k (C T Q k X k -1 m diag -1 (diag(1/g R k )R T k C T Q k X k ) T ) ), introducing the shorthand X = diag (1/g Q ) T diag (1/g R )and where diag -1 (•) : R r×r → R r denotes the matrix-to-vector extraction of the diagonal. This τ -dependent regularization also allows us to show smoothness of the objective in Proposition E.5, from which the convergence guarantee Proposition 3.3 follows. We derive the semi-relaxed projection Algorithm 2 of the sub-couplings Q and R in Appendix G for completeness. We also show in Lemma A.1 that g Q and g R induced by the semi-relaxed projection are both feasible and locally optimal, not requiring separate optimization.As (Q k+1 , R k+1 , T ) ∈ C 2 if and only if T ∈ Π g Q k+1 ,g R k+1, after the factor relaxation step, we next take a coordinate MD step on the latent coupling T :EQUATIONThis is equivalent to applying Sinkhorn (Algorithm 5) to T given g Q and g R with the kernel:K (k) T := T k ⊙ exp( -γ k diag(1/g Q k+1 )Q T k+1 CR k+1 diag(1/g R k+1 ) ). After the final iteration of the coordinate-MD scheme, X = diag (1/g Q ) T diag (1/g R ) satisfies Xg R = 1 r and X T g Q = 1 r as T is a coupling between g Q and g R . Thus P r = QXR T ∈ Π a,b and the iterates (Q k , T k , R k ) remainin the intersection of the constraint sets. Thus, in contrast to other approaches Scetbon et al. (2021) ; Lin et al. (2021) ; 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-square latent coupling T , providing greater interpretability in problemspecific applications. Above, we presented FRLC in the simplest case that T is square.Full-rank random initializations of the sub-coupling matrices. We propose a new initialization of the sub-couplings (Q, R, T ) for the LC-factorization in Algorithm 6.This generates a full-rank initialization (Proposition F.1) in the set of rank-r couplings Π a,b (r) and is accomplished by applying Sinkhorn to random matrices. Our approach differs from Scetbon et al. (2021) ; Scetbon & Cuturi (2022) who use initializations for the diagonal factorization of Forrow et al. (2019) , and are not applicable to a latent coupling that is non-diagonal, non-square, or with two distinct inner marginals.Algorithm 1 Balanced FRLC Input C, r, a, b, τ, γ, δ, ε Initialize g Q , g R = 1 r 1 r Q 0 , R 0 , T 0 ← Initialize-Couplings(a, b, g Q , g R ) # Alg. 6 X 0 ← diag(1/Q T 0 1 n )T 0 diag(1/R T 0 1 m ) while ∆((Q k , R k , T k ), (Q k-1 , R k-1 , T k-1 )) > ε do # ∆ as in (10) ∇ Q ← CR k X T k -1 n diag -1 ((CR k X T k ) T Q k diag(1/g Q )) T ∇ R ← C T Q k X k -1 m diag -1 (diag(1/g R )R T k C T Q k X k ) T γ k ← γ/ max{∥∇ Q ∥ ∞ , ∥∇ R ∥ ∞ } # ℓ ∞ -normalization of Scetbon & Cuturi (2022) K (k) Q , K (k) R ← Q k ⊙ exp( -γ k ∇ Q ), R k ⊙ exp( -γ k ∇ R ) Q k ← SR R -projection(K (k) Q , γ k , τ, a, Q T k-1 1 n , δ) # Semi-relaxed OT, Alg. 2 R k ← SR R -projection(K (k) R , γ k , τ, b, R T k-1 1 m , δ) # Semi-relaxed OT g Q , g R ← Q T k 1 n , R T k 1 m ∇ T = diag(1/g Q )Q T k CR k diag(1/g R ) γ T = γ/∥∇ T ∥ ∞ # ℓ ∞ -normalization K (k) T ← T k ⊙ exp( -γ T ∇ T ) T k ← Sinkhorn(K (k) T , g R , g Q , δ) # Balanced OT, Alg. 5 X k ← diag(1/g Q )T diag(1/g R ) end while Return P r = QXR TConvergence analysis of FRLC. As objective ( 8) is non-convex, it is important to have convergence guarantees. Our convergence criterion ∆(•, •) is defined in (10). To prove convergence we require a lower bound on the entries of g Q and g R . Previous works introduce a lower-bound vector α ≤ g enforced element-wise for stability and smoothness Scetbon et al. (2021) . In FRLC the use of semi-relaxed projections naturally enforces a lower-bound. In Appendix E.5, we show that for any δ ∈ (0, 1 r ), the FRLC algorithm's τ -weighted regularization on the inner marginals can guarantee a uniform lower-bound of δ on the entries: for sufficiently large τ and Õ(m 2 /ϵ) iterations for the sub-coupling Pham et al. (2020) , one guarantees a lower bound of δ on g R and g Q . This allows us to show objective smoothness in Proposition E.5. Previous work on low-rank optimal transport Scetbon et al. (2021) 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 (2013) . Proposition 3.3. Suppose one has f ∈ C 1 (X , R) with block-coordinate Lipschitz gradient and block smoothness constants (L i ) p i=1 , and a function h ∈ C(X , R) which is α-strongly convex. For Φ = f + h, suppose one performs coordinate mirror descent on Φ minimized over a product of closed convex sets X = p i=1 X i . Let the sub-iterates with respect to the i-th block update be {xi k } p i=0 where x k := x 0 k for k ∈ [N ]outer iterations. Then one has:min k ∆(x k , x k-1 ) ≤ D 2 L N (α 2 /2L) = 2D 2 L 2 N α 2 ,where D is (36), L is the global smoothness constant, stepsizes γ k,i := α/L, and convergence criterion ∆(x k , x k-1 ) is given in (35).Specialized to the LC-parametrization, the criterion ∆ k (x k , x k+1 ) is:∆ k (x k , x k+1 ) := 1 γ 2 k ∥Q k+1 -Q k ∥ 2 F + ∥R k+1 -R k ∥ 2 F + ∥T k+1 -T k ∥ 2 F (10) for x k = (Q k , R k , T k ).We show through Propositions 3.3, E.5 the following result: Proposition 3.4. The FRLC algorithm with step-sizes γ k = α/L and iterates Where N is the number of iterations, D the optimality-gap as in (36), and L = max i∈{1,..,3} (L i ) the global smoothness for L i = poly(∥C∥ F , n, m, r, δ) the block-wise smoothness constants.x k = (Q k , R k , T k ) has non-asymptotic stationary convergence in the criterion ∆(•, •) with: min k∈1,..,N -1 ∆ k (x k , x k+1 ) ≤ 2D 2 L 2 /N α 2The proof of Proposition 3.4 follows directly from our extension of the non-asymptotic criterion with the block-descent result Proposition 3.3 and the proof that this lemma holds in FRLC Proposition E.5. We also mention two improvements to other low rank approximation results in literature.In Proposition F.2 we show that one can analytically solve for the block-optimal g for the factorization of Scetbon et al. (2021) , and we improve the bound on 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 4.We also extend the FRLC algorithm to the Gromov-Wasserstein problem. This consists of computing a 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 is 1 T m P T A ⊙2 P 1 m + 1 T n P B ⊙2 P T 1 n -2⟨AP B, P ⟩, where ⊙ denotes the Hadamard (entrywise) product. Then the GW-specific Sinkhorn kernels areK (k) Q ← exp 2γ k (2AQXR T BRX T -A ⊙2 Q1 r 1 T r ) , K (k) R ← exp 2γ k (2BRX T Q T AQX -B ⊙2 R1 r 1 T r ) , K (k) T ← exp(4γ k diag(g -1 Q )Q T AQXR T BR diag(g -1 R )).In Algorithm 4, one can solve the GW-problem by using the kernels above. Here, we present the kernels omitting a rank-1 perturbation, which is given in Appendix D. From the Wasserstein and GW gradients, the FGW gradient is easily taken as a convex combination of the two. In this work, we primarily focus on the LC-factorization for the rank r Wasserstein problem (6). 
experimental results: 
 We compare FRLC to existing low-rank and full-rank optimal transport algorithms on several datasets: simulated datasets previously used in Tong et al. (2023) and Scetbon et al. (2021) ; a massive spatialtranscriptomics dataset Chen et al. (2022) ; and a graph partitioning task Chowdhury & Needham (2021) . Further details of each experiment (e.g. pre-processing, validation) are in Appendices K, L, and M. In the section below, LOT refers to the works of Scetbon et al. (2021 Scetbon et al. ( , 2023 Scetbon et al. ( , 2022) ) and Latent OT refers to Lin et al. (2021) . We first compare the balanced OT version of FRLC with the the low-rank balanced OT algorithm LOT of Scetbon et al. (2021) on a 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 (1) 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. (2021) and the random initialization of Alg. 6. 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 ⟩ F with increasing rank (Fig. 2b ) and consistently 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. 6 (Fig. 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. 2c , Fig. 7 ), with FRLC 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 = m = 1000 points from all three datasets with rank r = 100 (Table 2 ). At the same time FRLC achieves lower primal cost ⟨C, P ⟩ F with tighter marginals ∥P 1 n -a∥ 2 and ∥P T 1 m -b∥ 2 . Lin et al. (2021) only solves a proxy for the rank-constrained Wasserstein problem, and thus is not the focus of our comparisons. Nevertheless, we verify that on all synthetic experiments that FRLC achieves significantly lower primal OT cost than Latent OT (Table 5 ).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 (1) , Z (2) . In particular, the LC projection is defined from the LC factorization as follows.Definition 4.1 (LC-Projection). Let Q diag(1/g Q )T diag(1/g R )R T be an LC factorization of of a coupling matrix P ∈ Π a,b (r) computed from datasets Z (1) ∈ R n×d , Z (2) ∈ R m×d , with T ∈ R r1×r2 + . The LC-projections Y (1) and Y (2) of Z (1) and Z (2) are Y (1) := diag(1/g Q )Q T Z (1) , and Y (2) := diag(1/g R )R T Z (2) .By interpreting any factored coupling (Q, R, g) as an LC factorization (Q, R, diag(g)), Definition 4.1 describes the barycentric projections for both factorizations. We compare the projections of the coupling computed by FRLC to those of LOT Scetbon et al. (2021) on a dataset containing 1000 samples from 2D-Gaussians centered at the 5 th -roots of unity and 1000 samples from 2D Gaussians centered at the 10 th -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 r = 10 to match the number of target and source clusters. In the rank-5 case, FRLC uses a non-square latent coupling T ∈ R 10×5 + which correctly captures the coupling between clusters (Fig. 3(b )), while the LOT rank-5 projection computes barycenters that are outside of the clusters (Fig. 3c ). A similar result is observed for square rank-10 latent couplings computed by FRLC (Fig. 3d ) and LOT (Fig. 3e ) demonstrating that the LOT barycenters in Fig. 3b ) are not an artifact of using the lowest rank. We observe similar results on other simulated datasets (Fig. 11 ).We compare FRLC and the algorithm (LOT-U) of Scetbon 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 E11.5 stage mouse embryo and an E12.5 stage mouse embryo Chen et al. (2022) . Optimal transport is a popular approach to align single-cell Schiebinger et al. (2019) and spatial trancriptomics datasets Zeira et al. (2022) ; Liu et al. (2023) ; Klein et al. (2023) . In 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 be 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 a gene in E12.5 from expression of the gene in E11.5 using the alignment; (2) a cell type prediction task where we predict the cell types of E12.5 from the cell type clustering of E11.5 (Appendix M). We evaluate the accuracy of the gene expression prediction task through the Spearman correlation ρ 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 a 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 a 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. Scetbon 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 state-of-the-art unbalanced low-rank method on all three objectives (Table 4 ). We also solve a semi-relaxed version of each problem motivated by the observation that all cells from E12.5 have an ancestor, but not all cells from E11.5 have the same number of descendants due to cell growth and death. Thus the former marginal is tight, and the latter relaxed Halmos et al. (2024) . We run both semi-relaxed FRLC (FRLC-SR) and a setting of LOT-U that recovers the semi-relaxed problem (LOT-SR). Semi-relaxed FRLC achieves the best results on all three metrics by a large margin (Table 4 ). As one example, the expression of Tubb2b, a mouse brain marker gene, agreeing with the expression predicted from the semi-relaxed alignment of FRLC (Fig. 4a ).We evaluate FRLC on an unsupervised graph partitioning problem Chowdhury & Needham (2021) on four real-world graph datasets Yang & Leskovec (2012) ; Yin et al. (2017) ; Banerjee et al. (2013) . We benchmark the performance of the semi-relaxed and GW settings of FRLC against (1) GWL Xu et al. (2019) , solving a balanced GW problem; (2) SpecGWL Chowdhury & Needham (2021) using the heat kernel on the graph Laplacian as 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). 
discussion: 
 We provide comparison of existing low-rank solvers in Table 1 . The FRLC algorithm has a number of advantages, including (1) coarsening a full-rank plan P to non-diagonal latent coupling T ; (2) minimizing the primal OT problem for general cost C rather than a barycenteric problem; (3) optimizing only sub-couplings; and (4) using Sinkhorn alone as the sub-routine for low-rank OT. While we 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 our 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 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 Liu et al. (2021) . 
conclusion: 
 We introduce FRLC, an 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 LC factorization provides an interpretable coarse-graining of the full transport plan and its marginals through the mapping (P , a, b) → (T , g Q , g R ). We demonstrate the superior performance of FRLC compared to state-of-the-art low-rank methods on real and synthetic datasets. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: We do not generate any new datasets other than a toy dataset of Gaussians centered at n 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 we encourage authors to take this into account and make a 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?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 an equivalent approval/review based on the requirements of your country or institution) were obtained? Answer: [NA] Justification: The paper does not involve crowdsourcing 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 we 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 as the institution conducting the review. 
 
 Appendix: A.1 Low-rank factorizationsThe set of low-rank couplings. Given M ∈ R n×m + , the nonnegative rank of M is the least number of nonnegative, rank-1 matrices that sum to M :rk + (M ) = min r≥1 M = r i=1 M i , such that rk(M i ) = 1 and M i ≥ 0 for all i .Let a ∈ ∆ n , b ∈ ∆ m be probability vectors, and let Π a,b (r) denote the set of rank-r coupling matrices with marginals a and b:Π a,b (r) = {P ∈ R n×m + : P T 1 m = a, P 1 n = b, rk + (P ) ≤ r}.To optimize any cost over Π a,b (r), one requires a parameterization of this set.Factored couplings. The factored coupling parameterization of Π a,b (r) introduced in Forrow et al. (2019) , and used by Scetbon et al. (2021) ; Scetbon & Cuturi (2022) ; Scetbon et al. (2022 Scetbon et al. ( , 2023) ) is Cohen & Rothblum (1993) show that any P ∈ Π a,b (r) may be decomposed as P = Qdiag(1/g)R T for some triple (Q, R, g) ∈ FC. Thus, for cost matrix C ∈ R n×m , the general low-rank optimal transport problem is equivalent to an optimization over factored couplings:FC a,b (r) := {(Q, R, g) ∈ R n×r + × R m×r + × (R * + ) r : Q ∈ Π a,g , R ∈ Π b,g }.EQUATIONLatent coupling factorization. The latent coupling parameterization of Π a,b (r) introduced in Lin et al. ( 2021), and used in the present work isLC a,b (r) := {(Q, R, T ) ∈ R n×r + × R m×r + × R r×r + : Q ∈ Π a,• , R ∈ Π b,• , T ∈ Π g Q ,g R },where g Q , g R are the inner marginals of Q and R.Latent coupling diagonalization. The LC-factorization recovers the factorization of Forrow et al. (2019) as a sub-case. While the diagonal factorization of previous works cannot be directly converted to the LC-factorization, the LC-factorization can easily recover the diagonal factorization. In particular, takingQ ′ ← Q diag(1/g Q )T one can refactor P r = Q diag(1/g Q )T diag(1/g R )R T = Q ′ diag(1/g R )R T or alternatively taking R ′ = R diag(1/g R )TT may refactor asP r = Q diag(1/g Q )T diag(1/g R )R T = Q diag(1/g Q )(R ′ ) T So that instead of returning (Q, R, T ) one may alternatively return (Q, R, T ) → (Q ′ , R, diag(g R )) or (Q, R, T ) → (Q, R ′ , diag(g Q ))to recover the Forrow et al. (2019) factorization. An example of this diagonal-conversion is offered in Figure 12 .The FRLC optimization problem Our optimization problem is over the variables (Q, R, T ) and defined as follows:EQUATIONwhere our objective function L LC isL LC (Q, R, T ) = ⟨C, Q(diag(1/Q T 1 n ))T (diag(1/R T 1 m ))R T ⟩ (13) Given (Q, R, T ) ∈ LC a,b (r), sub-couplings Q and R are constrained by: C 1 (a) := {(Q, R, T ) ∈ R + : Q1 r = a}, C 1 (b) := {(Q, R, T ) ∈ R + : R1 r = b},while the convex sets constraining the latent coupling matrix T areC 2 (g Q ) := {(Q, R, T ) ∈ R + : T 1 r = g Q }, C 2 (g R ) := {(Q, R, T ) ∈ R + : T T 1 r = g R },whereg Q = Q T 1 n and g R = R T 1 m as per Definition 3.1. Under these definitions, LC a,b (r) = C 1 ∩ C 2 ,whereEQUATIONTo solve (12), we separate the variables into two "blocks" of variables, (Q, R) and T , and perform two block updates per iteration, as follows. Let (γ k ) k≥0 be a positive sequence of stepsizes. Suppose we have(Q k , R k , T k ) ∈ C.We update the first variable block (Q, R) by taking a locally optimal (mirror descent) update step, while the second variable block T is held fixed:(Q k+1 , R k+1 ) ← arg min (Q,R):(Q,R,T k )∈C1 ⟨(Q, R), ∇ (Q,R) L LC ⟩ + 1 γ k KL((Q, R)∥(Q k , R k )) (15)Here, we slightly abused the notation by putting (Q, R) inside an inner product. The triple (Q k+1 , R k+1 , T k ) produced by this update lies in C 1 . Next, we update T , the second variable block, by taking another locally optimal (mirror descent) step, while the first variable block is held fixed.EQUATION)By construction, the triple(Q k+1 , R k+1 , T k+1 ) ∈ C 2 . However, because the set C 1 only con- strains the first variable block Q k+1 , R k+1 , we have that (Q k+1 , R k+1 , T k+1 ) ∈ C 2 , and hence (Q k+1 , R k+1 , T k+1 ) ∈ C. Thus, each iteration produces a feasible triple (Q k+1 , R k+1 , T k+1 )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:EQUATIONwhere the objective function L LOT isEQUATIONand where C = C 1 ∩ C 2 with:C 1 := {(Q, R, g) ∈ R n×r + × R m×r + × (R * + ) r : Q1 r = a, R1 r = b} C 2 := {(Q, R, g) ∈ R n×r + × R m×r + × R r + : Q T 1 n = g = R T 1 m }.Here, there are also three optimization variables (Q, R, g), but they are updated together in each iteration of LOT. That is, given feasible(Q k , R k , g k ) ∈ C, an iteration of LOT updates this triple via (Q k+1 , R k+1 , g k+1 ) := arg min (Q,R,g)∈ C ⟨(Q, R, g), ∇ (Q,R,g) L LOT ⟩ + 1 γ k KL((Q, R, g)∥(Q k , R k , g k )),where (γ k ) k≥0 is again a positive sequence of stepsizes. Scetbon et al. (2021) then compute the unconstrained argmin across all variables to yield a set of unconstrained kernels (K Q , K R , k g ), using Dykstra to jointly project the unconstrained update onto the intersection C of the constraint sets.OT subroutine in FRLC To see why we do not need Dykstra in the FRLC scheme, observe that the update (15) of variables (Q, R) can be equivalently expressed as(Q k+1 , R k+1 ) ← arg min Q,R : (Q,R,T k )∈C1 ⟨(Q, R), ∇ (Q,R) L LC ⟩ + 1 γ k KL((Q, R)∥(Q k , R k )),Thus, even though the pair (Q, R) is being updated at this step, solving for (Q k+1 , R k+1 ) above is equivalent to updating each individually because Q and R do not share an inner marginal:EQUATIONWe choose to add the regularization τ KL((Q T 1 n , R T 1 m )∥(Q T k 1 n , R T k 1 m ))to turn each update here into an entropy-regularized semi-relaxed optimal transport problem, and to ensure β-smoothness:EQUATIONAfter updating Q and R, the update on T then follows a similar form:EQUATIONleading to balanced constraints on T of the formT T 1 r = Q T k+1 1 n and T 1 r = R T k+1 1 m, allowing the problem to be solved by Sinkhorn.Importantly, there are no constraints in C 1 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 are not required to match the inner marginals exactly, the OT problems associated to Q and R are semi-relaxed by construction.The separation of our block updates into a step where(Q, R, T k ) ∈ C 1 and (Q k+1 , R k+1 , T ) ∈ C 2allows us to entirely remove the optimization over inner marginals g Q and g R , as done in all previous works on low-rank optimal transport which optimize g explicitly as both a variable and a constraint of the optimization Scetbon et al. (2021 Scetbon et al. ( , 2023)) ; Scetbon & Cuturi (2022) . If one were to introduce an extended loss in the style of previous works which adds g Q and g R as variables in the formH(Q, R, T , g Q , g R ) = ⟨Qdiag(1/g Q )T diag(1/g R )R T , C⟩ F , one observes an equivalence to simply taking a semi-relaxed projection. Lemma A.1. Define the function H(Q, R, T , g Q , g R ) = ⟨C, Qdiag(1/g Q )T diag(1/g R )R T ⟩ Fand let L LC (Q, R, T ) be as in (13). One has the following equivalence:EQUATIONThus 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 r = a andR1 r = b for Q ≥ 0 n×r , R ≥ 0 m×r , one has i g R,i = i ⟨R T i,. , 1 m ⟩ = i ⟨R .,i , 1 m ⟩ = ℓ i R ℓ,i = ℓ b ℓ = 1 and i g Q,i = i ⟨Q T i,. , 1 n ⟩ = i ⟨Q .,i , 1 m ⟩ = ℓ i Q ℓ,i = ℓ a ℓ = 1Therefore, the inner marginalsg * Q = (Q * ) T 1 n and g * R = (R * )T 1 m induced by the optimal Q * , R * of the optimization problem on the right hand side satisfy the constraints g * R ∈ ∆ r , g * Q ∈ ∆ r on the left hand side, so the two minimums coincide. This implies that an extra optimization for g Q and g R is unnecessary in a coordinate update which alternates (Q, R) and T . If we did not perform a block-update, in the form of standard MD in the objective described above, we would encounter some difficulty. In particular g Q and g R would optimized with the constraint g Q ∈ ∆ r and g R ∈ ∆ r , and would concurrently constrain all of the other variables asQ T 1 n = g Q , T 1 r = g Q , and R T 1 m = g R , T T 1 r = g R .We use a latent non-diagonal coupling instead of an inner diagonal coupling diag(g) of the form of Forrow et al. (2019) . This allows us to loosen the constraint that the inner marginals have to be joined by a common coupling Q T 1 n = R T 1 m = g. The fundamental advantage of this choice is that we can decouple the convex-optimization problem for (Q, R, T ) entirely. One can simply solve for the optimal Q and R independently, yield the associated inner marginals for each Q T 1 n = g Q and R T 1 m = g R , and then find the optimal T which links the two. This link is provided by the aforementioned form of the problem, where:P = QXR TFor Q, R in either the appropriate set of couplings or a relaxation thereof (which we will describe shortly). X is related to T by:X = diag(1/g Q )T diag(1/g R )And, T ∈ Π g Q ,g R consistently for all cases. As the semi-relaxed case is intermediate between fully-relaxed and balanced, it has ideas which generalize to both directly. As such, we use it as the leading example again. As in Scetbon et al. (2021) , we take proximal-steps of the form:min ζ ⟨∇L(ζ) | ζ k , ζ⟩ F + 1 γ k KL(ζ∥K (k) )Where these steps are now in block-wise fashion on (Q, R) and T , rather than joint. One may identify for each sub-factor in (Q, R) and T a linearized gradient as before, which yields a set of objectives which each solve an independent optimal-transport for the sub-factors. In particular, we have that:⟨QXR T , C⟩ F = Tr QXR T C T = ⟨CRX T , Q⟩ F ⟨QXR T , C⟩ F = Tr QXR T C T = ⟨C T QX, R⟩ F ⟨QXR T , C⟩ F = Tr R T C T QX = ⟨Q T CR, X⟩ FA linearization in the left-slot of the inner product as ⟨Q, CRX(Q k ) T ⟩ := ⟨Q, CRX T ⟩ or ⟨C T QX(R k ), R⟩ F := ⟨C T QX, R⟩ F is common practice for quadratic problems. In this case, the directional derivative of Q and R in the matrix-direction V are respectively:D⟨CRX T , Q⟩ F • (V ) = ⟨CRX T , V ⟩ F =⇒ ∇ Q L = CRX T D⟨C T QX, R⟩ F • (V ) = ⟨C T QX, V ⟩ F =⇒ ∇ R L = C T QXWithout this linearization assumption on X, the full gradient may be evaluated as well. In particular, we note that for diag -1 (•) the matrix-to-vector extraction of the diagonal, the directional derivative on Q is:D⟨C, QXR T ⟩ F • V = ⟨CRX T , V ⟩ F + ⟨C, QDX T • (V )R T ⟩ F = ⟨CRX T -1 n diag -1 ((CRX T ) T Qdiag(1/g Q )) T , V ⟩ FThus, without the linearization assumption one may use product rule on X as an implicit function of Q (resp. R) to take the total derivative:∇ Q L FRLC = CRX T -1 n diag -1 ((CRX T ) T Qdiag(1/g Q )) T Likewise for R, D⟨C, QXR T ⟩ F • V = ⟨C T QX, V ⟩ F + ⟨C, QDX • (V )R T ⟩ F = ⟨C T QX -1 m diag -1 (diag(1/g R )R T C T QX) T , V ⟩ F ,and so∇ R L FRLC = C T QX -1 m diag -1 (diag(1/g R )R T C T QX) T .Both the W and GW problem have these rank-one perturbations of the gradient from the derivative with respect to X.Lastly, owing to the block-coordinate updates which fix (Q, R) preceding the update on T , the derivative on T follows directly by chain rule on X:D⟨Q T CR, X⟩ F • (V ) = ⟨Q T CR, V ⟩ F =⇒ ∇ X L = Q T CR ∇ T L = diag(1/g Q )Q T CR diag(1/g R )Let (γ k ) be a 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:K (k) Q ← min Q ⟨CRX T , Q⟩ F + 1 γ k KL(Q k ∥K (k) Q ) K (k) R ← min R ⟨C T QX, R⟩ F + 1 γ k KL(R k ∥K (k) R ) K (k) T ← min T ⟨Q T CR, X⟩ F + 1 γ k KL(T k ∥K (k) T ), with kernels K (k)ζj , for j = 1, 2, 3 given byK (k) Q := Q k ⊙ exp( -γ k CR k X T k ) K (k) R := R k ⊙ exp( -γ k C T Q k X k ) K (k) T := T k ⊙ exp( -γ k diag(g -1 Q )Q T CRdiag(g -1 R ) )Or, dropping the linearization assumption on (Q, R) the updates are:K (k) ζj , for j = 1, 2 given by K (k) Q := Q k ⊙ exp( -γ k (CR k X T k -1 n diag -1 ((CR k X T k ) T Q k diag(1/g Q k )) T ) ) K (k) R := R k ⊙ exp( -γ k (C T Q k X k -1 m diag -1 (diag(1/g R )R T k C T Q k X k ) T ) )The first projection is onto the set that satisfies the marginal constraint R1 r = b. In particular, following the discussion above, one has the coordinate-MD step:min (Q,R,T ) 1 γ k KL ((Q, R, T ) ∥ (K Q , K R , K T )) + τ KL(Q1 r ∥a) s.t. R1 r = b (25)As before, there is no difference from the previous case, where one takes the unconstrained projectionR = diag(b/K R 1 r )K R .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 particular, we consider the problem:min (Q,R,T ) 1 γ k KL((Q, R, T )∥(K Q , K R , K T )) + τ KL(Q1 r ∥a) + τ KL(R T 1 m ∥g (k-1) R ≡ R T k-1 1 m ) s.t. R1 r = b (26) Which yields the relaxed solution of R = SR R -projection(K R , γ k , τ, b, g (k-1) R) which generalizes the original projection and is equivalent to it for τ = 0. This regularization is essential, as it ensures β-smoothness of the objective. For Q, as the constraint on g is fully relaxed, the Lagrange multiplier λ 1 = 0 entirely, such that the problem 40 now becomes fully-unconstrained:EQUATIONTo generalize this solution, we again consider adding a soft-regularization on the inner marginal of Q, where we consider the alternate problem:inf Q 1 γ k KL(Q∥K Q ) + τ KL(Q1 r ∥a) + τ KL(Q T 1 n ∥g (k-1) Q ≡ Q T k-1 1 n ) (28)Which trivially recovers the original for τ = 0. This form has a solution given by an unbalanced optimal transport with kernel K Q . We see these two, convex problems in sequence give independent optimal solutions for Q and R as the two matrices are not required to share an inner marginal. The last step is to link them via T , corresponding to the projection of(Q k , R k , T ) onto the set of valid rank-r couplings min (Q k ,R k ,T )∈C L LC (Q k , R k , T ) := min T ∈Π(g Q ,g R ) L LC (Q k , R k , T ). We verify in C that if T ∈ Π g Q =Q T 1n,g R =R T 1m then P ∈ Π a,b. As such, one does not require any projection onto the intersection of convex sets, as done in Scetbon et al. (2021 Scetbon et al. ( , 2023) ) via the Dykstra projection algorithm Dykstra (1983) . Alternating a coordinate-MD step in (Q, R) and a step on T , 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:EQUATIONThis 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, γ, τ, a, b, δ u ← 1 n v ← 1 r repeat ũ ← u ṽ ← v u ← (a/Kv) v ← b/K T u τ /(τ +γ -1 ) until γ -1 max{∥log ũ/u∥ ∞ , ∥log ṽ/v∥ ∞ } < δ return diag(u)K diag(v)Algorithm 3 U-projection (unbalanced OT)Input K, γ, τ, a, b, δ u ← 1 n v ← 1 r repeat ũ ← u ṽ ← v u ← (a/Kv) τ /(τ +γ -1 ) v ← b/K T u τ /(τ +γ -1 ) until γ -1 max{∥log ũ/u∥ ∞ , ∥log ṽ/v∥ ∞ } < δ return diag(u)K diag(v)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 et al. (2021 Scetbon et al. ( , 2023) ) consider diag(1/g) to be the inner matrix, with marginals Q T 1 n = R T 1 m = g to ensure that the outer conditions Q1 r = a and R1 r = b hold. We instead relax the constraint that Q T 1 n and R T 1 m are equal, by allowing Q T 1 n = g Q and R T 1 m = g R to vary arbitrarily, and considering a non-diagonal inner matrix X ∈ R r×r in the place of diag(1/g) where we have the conditions:Xg R = XR T 1 m = 1 r And X T g Q = XQ T 1 n = 1 rThus, if one considers the coupling formed as P r = QXR T , one maintains that P r ∈ Π a,b from the condition of C 1 , defined in the balanced case as in ( 14). We clearly have that:P r 1 m = QXR T 1 m = Q1 r = aAnd:P T r 1 n = RX T Q1 n = R1 r = bWe first consider two approaches to optimizing for such a X, given that it is not a coupling. First, we consider the appropriate proximal step for Xmin ζ ⟨∇L(ζ) | ξ k , ζ⟩ F + 1 γ k KL(ζ∥K (k) )We first note that the gradient of our loss with respect to X is given as Q T CR. If one supposes that X is invertible with X -1 = T , we have that for any such T :X -1 1 r = T 1 r = g R And X -T 1 r = T T 1 r = g QThis implies that this inverse matrix T is a coupling such that T ∈ Π g R ,g Q , which also suggests one might be able to update it using Sinkhorn. In fact, being a density in R r×r + it represents a transition matrix between the latent r-dimensional variables. In particular, writing the proximal step in full, we have: minT ⟨Q T CR, T -1 ⟩ F + 1 γ k KL(T k ∥K (k) T ) Noting the derivative D(X -1 ) • V = -X -1 V X -1, we have from the first-order condition that:-T -T Q T CRT -T + 1 γ k log T k K (k) T = 0This implies the kernel matrix update:K (k) T = T k ⊙ exp{+γ k T -T Q T CRT -T }Where one then takes the Sinkhorn projection J onto the setΠ g R ,g Q as P Πg R ,g Q (K (k)T ) using the Sinkhorn algorithm Cuturi (2013b) . However, a more stable and inversion-free update exists which ensures X remains positive by a diagonal re-scaling in the form introduced by Lin et al. (2021) . In particular, if one takesX = diag(1/g Q )T diag(1/g R ), then Xg R = diag(1/g Q )T diag(1/g R )g R = diag(1/g Q )T 1 r = 1 r and likewise X T g Q = diag(1/g R )T T diag(1/g Q )g Q = diag(1/g R )T T 1 r = 1 r so that X necessarily satisfies Xg R = 1 r and X T g Q = 1 r . Thus T 1 r = g Q and T T 1 r = g R and T ∈ Π g Q ,g R .With analogous reasoning to before, one has a step for the coupling T in the form:min T ⟨Q T CR, diag(1/g Q )T diag(1/g R )⟩ F + 1 γ k KL(T k ∥K (k) T )Which yields the kernel matrix:K (k) T = T k ⊙ exp{-γ k diag(g Q ) -1 Q T CR diag(g R ) -1 } Which is likewise projected onto Π g Q ,g R using the Sinkhorn algorithm. From this T ∈ Π g Q ,g R , one takes X = diag(1/g Q )T diag(1/g R )as the inner matrix that corresponds the unequal marginals g Q and g R which ensuring P r ∈ Π a,b .Algorithm 4 FRLC (General marginal constraint low-rank optimal transport)Input C, r, r 2 , a, b, τ, τ 2 , γ, δ, ε Initialize g Q , g R = 1 r 1 r , 1 r2 1 r2 Q 0 , R 0 , T 0 ← Initialize-Couplings(a, b, g Q , g R ) if r = r 2 then X 0 ← T -1 0 # Invertible case else X 0 ← diag(1/Q T 0 1 n )T 0 diag(1/R T 0 1 m ) # General case end if while ∆((Q k , R k , T k ), (Q k-1 , R k-1 , T k-1 )) > ε do ∇ Q ← CR k X T k -1 n diag -1 ((CR k X T k ) T Q k diag(1/g Q )) T ∇ R ← C T Q k X k -1 m diag -1 (diag(1/g R )R T k C T Q k X k ) T γ k ← γ/ max{∥∇ Q ∥ ∞ , ∥∇ R ∥ ∞ } # ℓ ∞ -normalization of Scetbon & Cuturi (2022) K (k) Q , K (k) R ← Q k ⊙ exp( -γ k ∇ Q ), R k ⊙ exp( -γ k ∇ R ) if Balanced then Q k ← SR R -projection(K (k) Q , γ k , τ, a, Q T k-1 1 n , δ) # Semi-relaxed OT R k ← SR R -projection(K (k) R , γ k , τ, b, R T k-1 1 m , δ) # Semi-relaxed OT else if Unbalanced then Q k ← U-projection(K (k) Q , γ k , τ, a, Q T k-1 1 n , δ) # Unbalanced OT R k ← U-projection(K (k) R , γ k , τ 2 , b, R T k-1 1 m , δ) # Unbalanced OT else if Semi-Relaxed Left then Q k ← U-projection(K (k) Q , γ k , τ, a, Q T k-1 1 n , δ) # Unbalanced OT R k ← SR R -projection(K (k) R , γ k , τ, b, R T k-1 1 m , δ) # Semi-relaxed OT else if Semi-Relaxed Right then Q k ← SR R -projection(K (k) Q , γ k , τ, a, Q T k-1 1 n , δ) # Semi-relaxed OT R k ← U-projection(K (k) R , γ k , τ, b, R T k-1 1 m , δ) # Unbalanced OT end if g Q , g R ← Q T k 1 n , R T k 1 m ∇ T = diag(g Q ) -1 Q T k CR k diag(g R ) -1 γ T = γ/∥∇ T ∥ ∞ # ℓ ∞ -normalization K (k) T ← T k ⊙ exp{-γ T ∇ T } T k ← Sinkhorn(K (k) T , g R , g Q , δ) # Balanced OT X k ← diag(1/g Q )T diag(1/g R ) end while Return P r = QXR T Algorithm 5 Sinkhorn Algorithm (Cuturi (2013b), balanced OT) Input K, a, b, δ u ← 1 n v ← 1 m while ∥diag(u)Kv -a∥ 1 + ∥diag(v)K T u -b∥ 1 > δ do u (l+1) ← a/Kv (l) v (l+1) ← b/K T u (l+1) end while Return diag(u)K diag(v) D Gromov-Wasserstein (GW)As defined in, the general Gromov-Wasserstein problem concerns a minimization of the energy:Q A,B (P ) = i,j,k,l (A ik -B jl ) 2 P ij P kl (30)Where the minimization is over the set of all couplings Π a,b :EQUATIONWe 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 et al. (2023 Scetbon et al. ( , 2022)) . In 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 (-γ k ∇ ζ L(ζ)) for L(ζ)Q A,B (P ) = 1 T m P T A ⊙2 P 1 m + 1 T n P B ⊙2 P T 1 n -2⟨AP B, P ⟩ Which, using the constraints of C 2 reduces the cost as a function of Q, R, X to: Q A,B (Q, R, X) = 1 T r Q T A ⊙2 Q1 r + 1 T r R T B ⊙2 R1 r -2⟨QXR T , AQXR T B⟩ F ∇ Q Q A,B (Q, R, X) = 2A ⊙2 Q1 r 1 T r -4AQXR T BRX T Which is proportional to ∇ Q Q A,B (Q, R, X) ∝ -4AQXR T BRX Tfor the balanced and rightmarginal semi-relaxed case. And:∇ R Q A,B (Q, R, X) = 2B ⊙2 R1 r 1 T r -4BRX T Q T AQX Which likewise can be reduced in proportionality to ∇ R Q A,B (Q, R, g) ∝ -4BRX T Q T AQXin the balanced and left-marginal semi-relaxed case. The gradients, as presented above, assume a linearization in X ← X k . If one does not make this assumption and takesX(Q, R) = diag(1/Q T 1 n )T diag(1/R T 1 m ), a rank-one perturbation must be added to the Q and R gradient of the form: ∇(2)Q = 41 n diag -1 (XR T B(QXR T ) T AQ diag(1/g Q )) T ∇ (2) R = 41 m diag -1 (X T Q T AQXR T BR diag(1/g R ))T Analogously, for the gradient on T one can simply take the gradient with respect to X, and subsequently T as X(T ) = diag(g Q ) -1 T diag(g R ) -1 . The gradient with respect to X is given as:∇ X Q A,B (Q, R, X) = -4Q T AQXR T BRAnd thus, the directional derivative with respect to X in the direction V X = diag(g Q ) -1 V T diag(g R ) -1 and thus by the chain rule V T is:DQ A,B (Q, R, X) • (V X ) = -4⟨Q T AQXR T BR, V X ⟩ F = -4⟨Q T AQXR T BR, diag(1/g Q )V T diag(1/g R )⟩ FSo that the gradient with respect to the coupling matrix T is given as:∇ T Q A,B (Q, R, T ) = -4 diag(1/g Q )Q T AQXR T BR diag(1/g R )We show in Proposition 3.3 that directly applying the block-descent lemma of Beck & Tetruashvili (2013) to the template of Ghadimi et al.'s proof Ghadimi et al. (2014) is sufficient to show the nonasymptotic 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 E.1 and Lemma E.2 below. For completeness, we define all notation used, describe the coordinate mirror descent algorithm in general, and discuss a few relevant preliminaries.Suppose that the vector of n variables x ∈ R n is partitioned into p blocks, x = (x(1), . . . , x(p)), where x(i) ∈ R ni . Here, n 1 , . . . , n p are positive integers summing to n. Following the notation of Beck & Tetruashvili (2013) ; Nesterov (2012), we define matrices U i ∈ R n×ni such that x(i) = U T i x for all i = 1, . . . , p. This also implies x = p i=1 U i x(i). This allows us to define the vector of partial derivatives corresponding to each block of variables x(i):∇ i f (x) := U T i ∇f (x).In Beck & Tetruashvili (2013) , the gradient of f is assumed to be block-coordinate-wise Lipschitz, with L i the smoothness constant associated to the i-th block of variables: for all h i ∈ R ni , one hasEQUATIONand for such functions we denote by L := max i L i the (global) smoothness constant of ∇f . To be clear, a smoothness constant associated to f is a Lipschitz constant of its gradient.Lemma E.1 (Block descent lemma, Beck & Tetruashvili (2013), Lemma 3.2.). Suppose f ∈ C 1 (R n , R) is a continuously differentiable function over R n whose gradient is block-coordinatewise Lipschitz (32) for L i the smoothness constant associated to the i-th block of variables x(i). Let u, v be two vectors differing only in the i-th block: there existsEQUATION)Lemma E.1 is central to adapting the proof of Theorem 1 of Ghadimi et al. (2014) to our case. Their Theorem 1 concerns the non-asymptotic convergence of mirror descent for objectives of the formmin x∈X f (x) + h(x),whereX is a closed, convex subset of R n , f ∈ C 1 (X , R) is a possibly non-convex objective, and where h is an α-strongly convex function. Using notation similar to Ghadimi et al. (2014) , we write Φ(x) = f (x) + h(x). Additionally, Ghadimi et al. (2014) assume ∇f is L-Lipschitz for some L > 0.To unify our assumptions on f , we suppose that ∇f ∈ C 1 (X , R) is block-coordinate Lipschitz (32) with block Lipschitz constants (L i ) p i=1 , and that X itself decomposes as a product X = p i=1 X i , where each X i is a closed convex set constraining the block variables x(i).The proof of Ghadimi relies on β-smoothness of the objective in all variables. We show that component-wise 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 ω : X → R is a distance generating function with modulus α > 0, with respect to the Euclidean norm ∥ • ∥, if ω is continuously differentiable and strongly convex, so that ⟨xz, ∇ω(x) -∇ω(z)⟩ ≥ α∥x -z∥ 2 , for all x, z ∈ X .The prox-function (or Bregman divergence) associated with ω is thenV (x, z) = ω(x) -ω(z) -⟨∇ω(z), x -z⟩,and from this prox-function, γ > 0, and some g ∈ R n , we define the generalized projectionEQUATION)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:P X (x, g, γ) := 1 γ (x -x + ).The mirror descent (MD) algorithm is as follows: given initial point x 0 ∈ X , a total number of iterations N , and positive stepsizes (γ k ) N k=1 , at step k, the (k + 1)-st iterate is computed viax k+1 ← arg min u∈X ⟨∇f (x k ), u⟩ + 1 γ k V (u, x k ) + h(u) .Among all iterates x k , the MD algorithm outputs the one at which the generalized projected gradient is of least norm. Concretely, x R is the output of the MD algorithm, whereR := arg min k=0,...,N ∥g X ,k ∥ 2 ,and where g X ,k isg X ,k := P X (x k , ∇f (x k ), γ k ).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, we suppose that in a given iteration k, the block variables are updated sequentially from i = 1, . . . , p. This leads to doubly-indexed iterates (x i k ) with k = 1, . . . , N indexing each full iteration through all variables, and i = 0, 1, . . . , p indexing the sub-iterations which update one block of variables at a time.The coordinate mirror descent (CMD) algorithm takes as input an initial point x 0 ∈ X , a number of iterations N , and a sequence of positive stepsizes (γ k,i ) N,p k=1,i=1 . We set x 0 0 = x 0 , and for k = 0, . . . , N -1 and i = 1, . . . , p, we compute x i k from x i-1 k as follows:x i k (i) ← arg min ui∈Xi ⟨∇ i f (x i-1 k ), u i ⟩ + 1 γ k,i V i (u i , U T i x i-1 k ) + h i (u i )x i k (j) ← x i-1 k (j) for j ̸ = i Here, we have assumed that V can be written as a composite function of the block variables,V (x, z) = p i=1 V i (x(i), z(i)),as is the case for the KL divergence. We also have assumed that h has this composite structure (as with entropy):h(x) = p i=1 h i (x(i)).Lastly, for k = 0, . . . , N -1, we set x 0 k+1 := x p k . We define g X ,k := (g X ,k,1 , . . . , g X ,k,p ) to be the collection of block-wise differences, where by definitiong X ,k,i = P Xi (x i-1 k , ∇ i f (x i-1 k ), γ k,i ) = 1 γ k,i x i-1 k -x i k = U T i g X ,k .The convergence criterion ∆ 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 x R , whereEQUATION)For Φ = f + h as above (f has global smoothness constant L), let Φ * = arg min x∈X Φ(x), and defineEQUATION)The other lemma used in the proof of Proposition 3.3 is as follows.Lemma E.2 (Ghadimi et al. (2014) , Lemma 1). Let x + = arg min u∈X {⟨g, u⟩ + 1 γ V (u, x) + h(u)} and P X (x, g, γ) = 1 γ (xx + ). Then for all x ∈ R n , all g ∈ R n , and γ > 0, one has:⟨g, P X (x, g, γ)⟩ ≥ α∥P X (x, g, γ)∥ 2 + 1 γ (h(x + ) -h(x)).Proposition E.3 (Proposition 3.3). Suppose one has f ∈ C 1 (X , R) whose gradient is blockcoordinate Lipschitz, with block smoothness constants (L i ) p i=1 , and a function h ∈ C(X , R) which is α-strongly convex. For Φ = f + h, suppose one performs a coordinate mirror descent on Φ minimized over a product of closed convex sets X = p i=1 X i . Let the sub-iterates with respect to the i-th block update be {x i k } p i=0 where x k := x 0 k for k ∈ [N ] outer iterations. Then one has:min k ∆(x k , x k-1 ) ≤ D 2 L N (α 2 /2L) = 2D 2 L 2 N α 2 ,where D is (36), L is the global smoothness constant of f , and convergence criterion ∆(x k , x k-1 ) is given in (35). Above, the stepsizes γ k,i in the coordinate mirror descent are γ k,i := α/L.Proof. As f satisfies the hypotheses of the block descent lemma, Lemma E.1, we apply (33) to obtain:f (x i k ) ≤ f (x i-1 k ) + ⟨∇ i f (x i-1 k ), x i k -x i-1 k ⟩ + L i 2 ∥x i k -x i-1 k ∥ 2 . Noting the definition g X ,k,i = 1 γ k,i x i-1 k -x i k , one has f (x i k ) ≤ f (x i-1 k ) -γ k,i ⟨∇ i f (x i-1 k ), g X ,k,i ⟩ + L i 2 γ 2 k,i ∥g X ,k,i ∥ 2 .Lemma 1 of Ghadimi et al. (2014) (stated as Lemma E.2 above) applies identically through blockwise optimality onX i because g X ,k,i = P Xi (x i-1 k , ∇ i f (x i-1 k ), γ k,i ). Thus for any value ∇ i f (x i-1 k ) takes, f (x i k ) ≤ f (x i-1 k ) -αγ k,i ∥g X ,k,i ∥ 2 + h(x i k ) -h(x i-1 k ) + L i 2 γ 2 k,i ∥g X ,k,i ∥ 2 ,and thus,f (x i k ) + h(x i k ) ≤ f (x i-1 k ) + h(x i-1 k ) -αγ k,i - L i 2 γ 2 k,i ∥g X ,k,i ∥ 2 .The right-hand side above only becomes larger, taking L i = L to be the global smoothness constant.Introducing a sum over sub-iterates and total iterates, one hasN,p k,i f (x i k ) + h(x i k ) ≤ N,p k,i f (x i-1 k ) + h(x i-1 k ) - N,p k,i αγ k,i - L 2 γ 2 k,i ∥g X ,k,i ∥ 2 , N,p k,i Φ(x i k ) ≤ N,p k,i Φ(x i-1 k ) - N,p k,i αγ k,i - L 2 γ 2 k,i ∥g X ,k,i ∥ 2 .Noting the end-point condition Φ(x p k ) = Φ(x 0 k+1 ), one may cancel all intermediate terms:Φ * ≤ Φ(x N ) ≤ Φ(x 0 ) - N,p k,i αγ k,i - L 2 γ 2 k,i ∥g X ,k,i ∥ 2Thus one finds the upper bound in terms of Φ(x 0 ) and the minimum valueEQUATION)Taking γ k,i = α/L as in Ghadimi et al. (2014) , the bracketed term directly above becomes α 2 /2L, and one has:N k α 2 2L min k ∆(x k , x k-1 ) = N k α 2 2L min k p i ∥g X ,k,i ∥ 2 ≤ N k α 2 2L p i ∥g X ,k,i ∥ 2 ≤ Φ(x 0 ) -Φ * = D 2 L,where D is as in ( 36), and where we used (37) to obtain the last line. Thus,min k ∆(x k , x k-1 ) ≤ D 2 L N (α 2 /2L) = 2D 2 L 2 N α 2 ,completing the proof.Definition E.4 (Relative smoothness). Let β > 0 and let g ∈ C 1 (R n , R) be continuously differentiable. Additionally, let ω be a distance generating function with associated prox-function V . The function g is β-smooth relative to ω if the following holds:g(y) ≤ g(x) + ⟨∇ω(x), x -y⟩ + βV (y, x)Proposition E.5. Let ϵ > 0 be a predefined error tolerance, r > 0 a small rank parameter, δ ∈ (0, 1 r ) a lower-bound parameter, and N the number of inner iterations required for the semi-relaxed projection to converge for each iteration k as ∥Q T 1 ng(k-1) Q ∥ 2 < ϵ = 1 N 1 r -δ for g (0) Q = 1 r 1 r . Then, the FRLC objective L LC (Q, R, T ) = ⟨Q diag(1/Q T 1 n )T diag(1/R T 1 m )R T , C⟩ Fis component-wise smooth with respect to the variables Q, R, T with smoothness constants {β i } 3 i=1 where β i = poly(∥C∥ F , n, m, r, δ).τ KL(Q T 1 n ∥Q T k 1 n ) and τ KL(R T 1 m ∥R T k 1 m ) in 21and 22 ensures that we may use standard results on relaxed optimaltransport which bound how far the marginal Q T k 1 n deviates across iterations. In particular, for all ϵ > 0 there exists τ and N (number of iterations) sufficiently large, so that∥R T 1 m -g R ∥ 2 2 < ϵ and ∥Q T 1 n -g Q ∥ 2 2 < ϵ.In particular, Pham et al. (2020) shows that one can attain convergence to any ϵ for the unbalanced problem in N = Õ(m 2 /ϵ) iterations (hiding logarithmic and τ -factors) for each sub-problem solving for R k in Algorithm 4 and analogously N = Õ(n 2 /ϵ) for Q. Under the uniform initialization of g R , and after N iterations, we have:∥g (0) R -g (N ) R ∥ 2 = 1 r 1 r -g (N ) R 2 ≤ N k=1 ∥g (k) R -g (k-1) R ∥ 2 .With sufficiently large τ and N = Õ(m 2 /ϵ) sub-iterations, one can guarantee that:∥g (k) R -g (k-1) R ∥ 2 < ϵ = 1 N 1 r -δ .This implies, for all iterations of the algorithm and all indices i, thatEQUATIONThus, by adding the regularization on the inner marginal, one may guarantee a lower-bound on the entries of g R and g Q . 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:∇ Q L LC (Q, R, T ) = ∇ (A) Q L LC + ∇ (B) Q L LC = CRX T -1 n diag -1 ((CRX T ) T Qdiag(1/g Q )) T = ∇ (A) Q L LC -1 n diag -1 ((∇ (A) Q L LC ) T Qdiag(1/g Q )) T Where ∥∇ Q L LC (Q k+1 , R k , T k ) -∇ Q L LC (Q k , R k , T k )∥ F ≤ ∥∇ (A) Q L LC (Q k+1 , R k , T k ) -∇ (A) Q L LC (Q k , R k , T k )∥ F + ∥∇ (B) Q L LC (Q k+1 , R k , T k ) -∇ (B) Q L LC (Q k , R k , T k )∥ F . (39)Starting with the first term on the right side of (39), one has:∥∇ (A) Q L LC (Q k+1 , R k , T k ) -∇ (A) Q L LC (Q k , R k , T k )∥ F = ∥CR k (diag(1/g Q k )T k diag(1/g R k )) T -CR k (diag(1/g Q k-1 )T k diag(1/g R k )) T ∥ F ≤ ∥diag(1/g R k )∥ F ∥C∥ F ∥R k ∥ F ∥T k ∥ F ∥diag(1/g Q k ) -diag(1/g Q k-1 )∥ F . Note that ∥R k ∥ 2 F = i,j (R k ) 2 i,j < i,j (R k ) i,j = 1, as R k has marginals which sum to one. The same bound holds for ∥T k ∥ 2 F , which is also a coupling. Invoking the lower-bound (38) of δ on the entries of the inner marginals, and continuing from the above display,∥∇ (A) Q L LC (Q k+1 , R k , T k ) -∇ (A) Q L LC (Q k , R k , T k )∥ F ≤ ∥C∥ F δ ∥diag(1/g Q k ) -diag(1/g Q k-1 )∥ F = ∥C∥ F δ ∥diag(1/g Q k-1 )diag(1/g Q k )(diag(g Q k-1 ) -diag(g Q k ))∥ F ≤ ∥C∥ F δ 3 ∥diag(g Q k-1 ) -diag(g Q k )∥ F .To further bound the right-hand side above, consider:∥diag(g Q k ) -diag(g Q k-1 )∥ 2 F = ∥Q T k 1 n -Q T k-1 1 n ∥ 2 2 = r i=1   r j=1 (Q k ) i,j -(Q k-1 ) i,j   2While can easily be upper-bounded by an application of Jensen's inequality as= r i=1 r 2   r j=1 1 r (Q k ) i,j -(Q k-1 ) i,j   2 ≤ r i=1 r 2   r j=1 1 r (Q k ) i,j -(Q k-1 ) i,j 2   = r r i=1 r j=1 (Q k ) i,j -(Q k-1 ) i,j 2 = r∥Q k -Q k-1 ∥ 2 F . Likewise, we have that ∥diag(g R k ) -diag(g R k-1 )∥ 2 F ≤ r∥R k -R k-1 ∥ 2 F . Thus, it holds that ∥C∥ F δ 3 ∥g Q k -g Q k-1 ∥ 2 ≤ ∥C∥ F √ r δ 3 ∥Q k -Q k-1 ∥ F .Next, we focus on the ∇(B) Q term. Observe that ∥1 n diag -1 X∥ 2 F = Tr(1 n diag -1 X) T (1 n diag -1 X) = n∥diag -1 X∥ 2 2 ≤ n∥X∥ 2 F . Thus: ∥∇ (B) Q k+1 -∇ (B) Q k ∥ F ≤ √ n∥(∇ (A) Q k+1 ) T Q k+1 diag(1/g Q k+1 ) -(∇ (A) Q k ) T Q k diag(1/g Q k )∥ F Adding and subtracting terms in the norm and applying triangle inequality:√ n∥(∇ (A) Q k+1 ) T Q k+1 diag(1/g Q k+1 ) -(∇ (A) Q k ) T Q k+1 diag(1/g Q k+1 ) + (∇ (A) Q k ) T Q k+1 diag(1/g Q k+1 ) -(∇ (A) Q k ) T Q k diag(1/g Q k )∥ F ≤ √ n∥∇ (A) Q k+1 -∇ (A) Q k ∥ F ∥Q k+1 ∥ F ∥diag(1/g Q k+1 )∥ F + √ n∥∇ (A) Q k ∥ F ∥(Q k+1 diag(1/g Q k+1 ) -Q k diag(1/g Q k )∥ FInvoking the lower-bound on the marginal, continuing from the above display,≤ √ n δ ∥∇ (A) Q k+1 -∇ (A) Q k ∥ F + √ n∥∇ (A) Q k ∥ F ∥(Q k+1 diag(1/g Q k+1 ) -Q k diag(1/g Q k )∥ F Let us consider ∥∇ (A) Q k ∥ F ≤ ∥X k ∥ F ∥C∥ F ∥R k ∥ F ≤ ∥X k ∥ F ∥C∥ F .We have that:∥X k ∥ 2 F = i,j 1 g 2 (Q k )i T 2 ij 1 g 2 (R k )jWhere we always have that T ij ≤ g Qi and T ij ≤ g Rj by definition of T as a coupling. As such:≤ ij 1 g 2 Qi (g Qi g Rj ) 1 g 2 Rj = ij 1 g Qi 1 g Rj = g -1 Q g -T R , 1 m 1 T r F ≤ mr δ 2Thus ∥∇ (A)Q k ∥ F ≤ √ mrδ ∥C∥ F , and the bound above reduces to≤ √ n δ ∥∇ (A) Q k+1 -∇ (A) Q k ∥ F + √ nmr δ ∥C∥ F ∥(Q k+1 diag(1/g Q k+1 ) -Q k diag(1/g Q k )∥ F Further bounding the last term in the norm √ nmr δ ∥C∥ F ∥(Q k+1 diag(1/g Q k+1 ) -Q k diag(1/g Q k+1 ) + Q k diag(1/g Q k+1 ) -Q k diag(1/g Q k )∥ F ≤ √ nmr δ ∥C∥ F (∥diag(1/g Q k+1 )∥ F ∥Q k+1 -Q k ∥ F + ∥Q k ∥ F ∥diag(1/g Q k+1 ) -diag(1/g Q k )∥ F ) ≤ √ nmr δ ∥C∥ F 1 δ + √ r δ 2 ∥Q k+1 -Q k ∥ FThus, the final bound on the ∇ (B)Q term is:≤ √ n δ ∥∇ (A) Q k+1 -∇ (A) Q k ∥ F + √ nmr δ ∥C∥ F ∥(Q k+1 diag(1/g Q k+1 ) -Q k diag(1/g Q k )∥ F ≤ ∥C∥ F √ nr δ 4 + √ nmr δ 2 ∥C∥ F 1 + √ r δ ∥Q k+1 -Q k ∥ FThe total component-wise smoothness bound on Q is then∥∇ Q L LC (Q k+1 , R k , T k ) -∇ Q L LC (Q k , R k , T k )∥ F ≤ ∥C∥ F δ 2 √ nr δ 2 + √ nmr 1 + √ r δ + √ r δ ∥Q k+1 -Q k ∥ FIdentical reasoning applies for ∇ R L LC , where we similarly have the gradient split into two terms:∇ R L LC (Q, R, T ) = ∇ (A) R L LC ∇ (B) R L LC = C T QX -1 m diag -1 (diag(1/g R )R T C T QX) T = ∇ (A) R L LC -1 m diag -1 (diag(1/g R )R T ∇ (A) R L LC ) TAs before, we may first show smoothness in ∇ (A) R using the same steps as for Q ∥∇(A) R L LC (Q k+1 , R k+1 , T k ) -∇ (A) R L LC (Q k+1 , R k , T k )∥ F = ∥C T Q k+1 diag(1/g Q k+1 )T k diag(1/g R k+1 ) -C T Q k+1 diag(1/g Q k+1 )T k diag(1/g R k )∥ F ≤ ∥C∥ F ∥diag(1/g Q k+1 )∥ F ∥Q k+1 ∥ F ∥T k ∥ F ∥diag(1/g R k+1 ) -diag(1/g R k )∥ F ≤ ∥C∥ F δ ∥diag(1/g R k+1 ) -diag(1/g R k )∥ F ≤ ∥C∥ F δ 3 ∥g R k+1 -g R k ∥ 2 ≤ ∥C∥ F √ r δ 3 ∥R k+1 -R k ∥ F . For ∇ (B)R , one may use the same reasoning as before to find:∥∇ (B) R L LC (Q k+1 , R k+1 , T k ) -∇ (B) R L LC (Q k+1 , R k , T k )∥ F ≤ ∥1 m diag -1 (diag(1/g R k+1 )R T k+1 ∇ (A) R k+1 L LC ) -diag -1 (diag(1/g R k )R T k ∇ (A) R k L LC ) T ∥ F ≤ √ m∥diag(1/g R k+1 )R T k+1 ∇ (A) R k+1 L LC -diag(1/g R k )R T k ∇ (A)R k L LC ∥ F As before, one may apply three rounds of triangle inequality inside the norm to bound this directly in terms of ∥∇(A) R k+1 -∇ (A) R k ∥ F , ∥diag(1/R k+1 ) -diag(1/R k )∥ F , and ∥R k+1 -R k ∥ F .Each of these terms is smooth in R by the lower-bound argument, so that smoothness in R holds analogously to Q. The remainder of the proof for smoothness in R thus follows identically to that of Q above.For T , the component-wise bound of∥∇ T k+1 L LC -∇ T k L LC ∥ F = ∥Q k+1 CR T k+1 -Q k+1 CR T k+1 ∥ 2 F ≤ L T ∥T k+1 -T k ∥ F holdstrivially for any L T > 0 as the gradient is uniquely determined by Q and R alone. Thus there exist L Q , L R , L T > 0 as component-wise smoothness constants for L LC (Q, R, T ).We next prove Proposition 3.4, restated just below for convenience. Proposition E.6 (Proposition 3.4). Consider the FRLC objective (8). The FRLC algorithm, Algorithm 4, yields β-smooth iterates for β = poly(∥C∥ F , m, r, δ), where δ denotes the lower-bound on the entries of g Q , g R . Consider the convergence metric of 3.3 adapted from Ghadimi et al. (2014) , given as:∆ k (x k , x k+1 ) = p i=1 ∥g X ,k,i ∥ 2 = 1 γ 2 k ∥Q k -Q k-1 ∥ 2 F + ∥R k -R k-1 ∥ 2 F + ∥T k -T k-1 ∥ 2 F for x k = (Q k , T k , R k ).Define the gap to the optimal solution D as in (36), and let L = sup i (L i ) to be the global smoothness constant across all components. Then for γ k = α/L as defined in 3.3 the FRLC algorithm has the non-asymptotic stationary convergence guarantee that:min k∈1,..,N -1 ∆ k ≤ 2D 2 L 2 N α 2Proof. The proof of the non-asymptotic stationary convergence of mirror descent of Ghadimi et al. (2014) , adapted for coordinate mirror descent using the block-descent lemma in 3.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(L Q , L R , L T ) = poly(∥C∥ F , m, r, δ).Proposition E.7. Low-rank Approximation Error. Let SR-W ⋆ r denote the optimal rank-r approximation for the semi-relaxed low-rank optimal transport problem, and let SR-W ⋆ denote the optimal solution for the full-rank semi-relaxed optimal transport problem.Additionally, suppose c b = m j=1 b j denotes the sum of the entries of the second marginal (c b = 1 if a probability measure). Then we have the following upper-bound on the objective error:|SR-W ⋆ r (µ b ) -SR-W ⋆ (µ b )| ≤ c b max p,q {C pq } -min p,q{C pq } ln (min{n, m}/(r -1))We note that this bound also applies for the standard balanced optimal transport case, giving:|W ⋆ r (µ a , µ b ) -W ⋆ (µ a , µ b )| ≤ max p,q {C pq } -min p,q{C pq } ln (min{n, m}/(r -1))and improves the previous bound of |W ⋆ r (µ a , µ b ) -W ⋆ (µ a , µ b )| ≤ ∥C∥ ∞ ln (min{n, m}/(r -1)) as the distance matrix C contains only non-negative entries.Proof. We adapt the proof from Scetbon & Cuturi (2022) for the balanced case, which previously gave the bound:|W ⋆ r (µ a , µ b ) -W ⋆ (µ a , µ b )| ≤ ∥C∥ ∞ ln (min{n, m}/(r -1))In particular, for z = min{m, n}, there exists an optimal rank + (P * ) ≤ z where one may express the optimal solution for the non-negative coupling matrix P as a sum of z rank-one, non-negative outer products qk rT k ≽ 0:P * = z k=1 qk rT k = z k=1 λ k q k r T kWhere we write this sum in terms of normalized vectors q k = qk /∥ qk ∥ 1 , r k = rk /∥r k ∥ 1 , andλ k = ∥r k ∥ 1 ∥ qk ∥ 1 .Without any loss of generality, (λ k ) z k=1 is ordered in terms of decreasing value such that λ 1 ≥ λ 2 ≥ ... ≥ λ z . Note that we have a fixed constraint for the sum of the entries of P for the semi-relaxed case, assuming b is a general positive measure, as1 T n P 1 m = P T 1 n T 1 m = b T 1 m = m j=1 b j := c b (where c b = 1 if b ischosen to be a probability measure, i.e. in the balanced case). Moreover, it is simple to observe for these ordered values that λ k ≤ (c b /k). As in Scetbon & Cuturi (2022) , define the weighted average of the bottom zr + 1 vectors of the decomposition to be:α r = z i=r λ i q i z i=r λ i β r = z i=r λ i r i z i=r λ iAnd take the rank-r approximation using the optimal r -1 vectors of OPT and this weighted average of the bottom to be:Pr = r-1 i=1 λ i q i r T i + z i=r λ i α r β T rWhere, by the assumption that P * ∈ Π b is feasible:P T r 1 n = r-1 i=1 λ i r i q T i 1 n + z i=r λ i β r α T r 1 n = r-1 i=1 λ i r i + z i=r λ i β r = z i=1 λ i r i = bThus Pr ∈ Π b,r is a feasible rank-r solution by feasibility of P * . One can verify that if P * ∈ Π a,b , then Pr 1 m = a and the solution is again feasible. From this, we observe that the difference between this solution and P * is an upper-bound to the difference between P * and the optimal rank-r solution:|SR-W ⋆ r (µ a , µ b ) -SR-W ⋆ (µ a , µ b )| = |⟨P * r , C⟩ F -⟨P * , C⟩ F | ≤ ⟨ Pr , C⟩ F -⟨P * , C⟩ F = r-1 i=1 λ i q i r T i + z i=r λ i α r β T r - z i=1 λ i q i r T i , C F = z i=r λ i α r β T r - z i=r λ i q i r T i , C FNoting that α r , β r and q i , r i are unit normalized positive vectors, the sum of the entries of the outer product 1 T n q i r T i 1 m = 1, and likewise for α r β T r . Thus, continuing from the above display:= ( z i=r λ i )⟨α r β T r , C⟩ F - z i=r λ i ⟨q i r T i , C⟩ F ≤ ( z i=r λ i )⟨α r β T r , C⟩ F -( z i=r λ i ) min p,q {C pq } ≤ max p,q {C pq } -min p,q {C pq } z i=r λ i ≤ c b max p,q {C pq } -min p,q {C pq } z i=r 1 i ≤ c b max p,q {C pq } -min p,q{C pq } ln (z/(r -1))Concluding the proof. As discussed, this directly applies to the balanced case (for c b = 1).We propose a new initialization of the sub-couplings Q, R, T for the LC-factorization. Algorithm 6 generates a random full-rank initial condition in the set of couplings Π a,b which still satisfies the marginal constraints. It accomplishes this by sampling random matrices which are full-rank and applying the Sinkhorn algorithm to each of them. Scetbon 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, is 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 in the optimization. This initialization is given in Algorithm 6.Algorithm 6 Initialize-Couplings Input a ∈ ∆ n , b ∈ ∆ m , g Q ∈ ∆ r , g R ∈ ∆ r C Q ∼ [0, 1] n×r , C R ∼ [0, 1] m×r , C T ∼ [0, 1] r×r K Q ← e C Q , K R ← e C R , K T ← e C T Q ← Sinkhorn(K Q , a, g Q ) R ← Sinkhorn(K R , b, g R ) T ← Sinkhorn(K T , g Q = Q T 1 n , g R = R T 1 m ) Return (Q, R, T )Proposition F.1. Suppose one samples an initial condition on the optimal transport coupling using Algorithm 6, where we assume C ij ∼ Unif(0, 1) such that P(rank(C) < min{n, m}) = 0.Additionally, suppose that a, b > 0 holds elementwise for both marginals a ∈ ∆ n , b ∈ ∆ r . Then the elementwise exponential exp{C} (or exp{-C}) has full-rank and the return Sinkhorn(e -C , a, b) has full-rank.Proof. It is established that a random matrix C ∼ [0, 1] n×m has full-rank with probability one. For K = exp{C}, it holds that the matrix must be entry-wise positive with K ij ≥ 0. If columns C .,i ̸ = C .,j then clearly C ⊙k .,i ̸ = C ⊙k .,j , and if C .,i , C .,j ≽ 0 and are independent remain so under element-wise powers. One may easily show this by contrapositive. Suppose there exist constants c 1 , c 2 such that:c 1 C ⊙k .,i + c 2 C ⊙k .,j = 0As C .,i , C .,j ≽ 0, without loss of generality one may assume c 1 > 0 and c 2 < 0. Then:c 1 C ⊙k .,i = c 1 1 ⊙ C ⊙k .,i = -c 2 1 ⊙ C ⊙k .,j = -c 2 C ⊙k .,j =⇒ - c 1 c 2 1/k 1 = c1 = C .,jC .,i So clearly one has that C .,j -cC .,i = 0 for c > 0. This implies the columns C .,j and C .,i are dependent. Thus it is 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 c 1 e C.,i + c 2 e C.,j = 0 for c 1 > 0, c 2 < 0, one finds log -c1 c2 1 = C .,j -C .,i . Without loss of generality, assume 0 < -c 2 ≤ c 1 and C .,j > C .,i > 0, so that δ = log -c1 c2 ≥ 0. Then, considering constants q 1 , q 2 : q 2 C .,j + q 1 C .,i = (q 1 + q 2 )C .,i + δq 2 1 = 0 Assuming C .,i has greater than one unique entry, which we assume as the entries are sampled densely in R, the two vectors are dependent if and only if q 1 = -q 2 and δ = 0, implying C .,i = C .,j . Thus, for the set of independent column vectors of C, given as {C .,i } m i=1 , the set {e C.,i } m i=1 , is also linearly independent. This holds analogously for the row vectors. As C is full-rank and span({C .,i }) = R min{m,n} , we have that span(K) = R min{m,n} as K = e C = ∞ k=0 C ⊙k k!(analogously e -C ) and remains full-rank.Sinkhorn expresses each variable as X = diag(u)K diag(v) where rank(X) = rank(diag(u)K diag(v)). As Cuturi (2013b) updates the vectors u ← a/Kv and v ← b/K T u from u 0 = 1 n and v 0 = 1 m , if a, b > 0 holds element-wise, one has that u, v > 0 elementwise as well. Then, one has that null diag(v) = {0} and null diag(u) = {0}, implying that rank(diag(u)K diag(v)) = rank(K) = min{n, m}.Thus, our initialization returns a random coupling matrix X ∈ Π a,b of full-rank.In the next proposition, we show that one can analytically solve for the block-optimal weights g for the factorization of the coupling matrix P as P = Q diag(1/g)R T Forrow et al. (2019); Scetbon et al. (2021) . Proposition F.2. For the minimization problem expressed as ming∈∆r ⟨Q diag(1/g)R T , C⟩ FOne has the closed-form minimizer of g * defined entrywise as:g * i = √ ω i r j=1√ ω jFor ω = diag -1 (Q T CR), when ω ≥ 0 holds entrywise.Proof. As ω ≥ 0, we consider the simplex condition r j=1 g j = 1 alone. Writing out the Lagrangian associated to our objective, with λ ∈ R our equality-condition dual variable, we have:L(g, λ) = ⟨Q diag(1/g)R T , C⟩ F + λ(1 - j g j )Let us consider a rewriting of the inner product term:⟨Q diag(1/g)R T , C⟩ F = n i=1 m j=1 C ij r k=1 Q ik 1 g k R T kj = r k=1 1 g k n i=1 m j=1 Q T ki C ij R jk = r k=1 1 g k Q T CR k,k = r k=1 ω k g k = ω T (1/g)Where division is interpreted element-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 a vector dot-product as:L(g, λ) = ω T (1/g) -λ(1 - r j=1 g j )Thus, the first order condition tells us that the value of the coupling weight g j is related to λ as:∂ gj L(g, λ) = - ω j g 2 j + λ = 0 =⇒ g j = ω j λAnd by relying on the summation condition on the probability density g, yields the Langrange multiplier as r j=1g j = 1 = r j=1ω j λ so that one finds1 = 1 √ λ r j=1 √ ω j =⇒ λ =   r j=1 √ ω j   2Plugging our Lagrange-multiplier into the above expression yields:g j = ω j λ = ω j r i=1 √ ω i 2 = √ ω j r i=1√ ω iAs the Hessian ∇ 2 g L = diag( ω g 3 ) ≽ 0, we conclude that this value of g indeed minimizes the loss over ∆ r .For the problem:inf (Q,R k ,g k )∈C1∩C2 1 γ k KL(Q∥K Q ) + τ KL(Q1 r ∥a) -λ T 1 Q T 1 n (40)one can find a simple set of semi-relaxed updates for the coupling matrix. We note the primal-dual relationship of Sinkhorn, Q = diag(e γ k f1 )K Q diag(e γ k h1 ), and consider the entry-wise first-order condition required for the sub-coupling Q:EQUATIONThus:Q ij = a i Q T i,. 1 r τ γ k (K Q ) ij e -λ1j γ kAnd in matrix-form, this yields:Q = diag a Q1 r τ γ k K Q diag(e -γ k λ1 ) = diag(e γ k f1 )K Q diag(e γ k h1 )And expanding the Q1 r term explicitly, noting that X diag(v)1 = Xv, we have:diag a Q1 r τ γ k K Q diag(e -γ k λ1 ) = diag a diag(e γ k f1 )K Q diag(e γ k h1 )1 r τ γ k K Q diag(e -γ k λ1 ) = diag a e γ k f1 ⊙ K Q e γ k h1 τ γ k K Q diag(e -γ k λ1 )Thus, we identify e γ k h1 = e -γ k λ1 as the right dual vector, and identify the following relationship in terms of the left dual vector:a e γ k f1 ⊙ K Q e γ k h1 τ γ k = e γ k f1 =⇒ e γ k f1 = a K Q e γ k h1 τ τ +1/γ k From 40, the condition that (Q, R k , g k ) ∈ C 1 ∩ C 2 implies that Q T 1 m = g k := g.As such, we find that:Q T 1 m = diag(e γ k h1 )K T Q diag(e γ k f1 )1 m = diag(e γ k h1 )K T Q e γ k f1 = gImplying an update for e γ k h1 in the form:e γ k h1 = g K T Q e γ k f1Analogous reasoning applies for a relaxation of the other marginal, yielding the SR R -projection and SR L -projection (i.e. semi-relaxed OT).Algorithm 7 SR L -projection (semi-relaxed OT, left marginal relaxed) 1) are taken to be small constants. The time-complexity of Algorithm 4 is O(BLr 2 (n + m)) for B the number of inner Sinkhorn iterations, L the number of mirror-descent steps, n, m the number of samples in the first and second dataset, and r = max{r 1 , r 2 , d} for r 1 , r 2 the ranks of the latent coupling and d the rank of the factorized distance matrix C (generally chosen to be a constant near r 1 , r 2 ). Each matrix-multiplication is of max order (n + m)r 2 , which happens a constant number of times in the computation of each gradient ∇ i , and for the respective Sinkhorn matrix-vector multiplications Kv and K T u. The L outer steps follow from the mirrordescent convergence rate and the number of iterations B required for each projection follow from the convergence of Sinkhorn. In particular, for ε a fixed error tolerance and η the entropy constant, one finds a ±εD approximation for D the diameter of the data in B = poly(1/ηε) iterations using the Sinkhorn algorithm Charikar et al. (2023) ; Cuturi (2013a) .Input K, γ, τ, a, b, δ u ← 1 n v ← 1 r repeat ũ ← u ṽ ← v u ← (a/Kv) τ /(τ +γ -1 ) v ← b/K T u until γ -1 max{∥log ũ/u∥ ∞ , ∥log ṽ/v∥ ∞ } < δ return diag(u)K diag(v) H Discussion of Complexity For (Q, R, T ) ∈ R n×r1 + × R m×r2 + × R r1×r2 + , the space complexity O(nr 1 + r 1 r 2 + mr 2 ) is linear if the ranks r 1 , r 2 = o(As mentioned previously, works such as Charikar et al. (2023) have developed algorithms with linear O((n + m) 1+o( 1) poly(1/ϵ)) time-complexity and O((n + m)d) space-complexity for sketching the optimal transport cost value. Recent works on low-rank factorization of the optimal transport coupling matrix P (the matrix associated to the coupling γ ∈ Π(µ, ν)) Scetbon & Cuturi (2022) ; Scetbon et al. (2023 Scetbon et al. ( , 2021) ) have achieved per-iteration time-complexities of O(T (n + m)dr) for some constant non-negative rank r ≥ 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(log (nm)/ϵ 2 ) while preserving pairwise distances, however, currently no proofs exist which offer the number of iterations T until convergence to some tolerance ε. This is partially due to 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 is O((n + m)dr), 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 & Woodruff (2018) developed an algorithm which, given two point sets {x i } n i=1 and {y j } m j=1 in some metric space X , finds a rank r approximation in O((n + m) 1+γ poly(r, 1/ε)) for γ > 0 an 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)r/ε) entries of the input matrix with a run-time which removes dependence on γ that is merely O((n + m)poly(r, 1/ε)). This algorithm is used by all of the low-rank optimal transport works. These works, by finding a low-rank approximation to the coupling matrix P ≈ AB T ∈ R n×m + due to space limitations on the coupling, necessarily cannot store the full distance matrix C ∈ R n×m + of the same size in memory either. As such, it must also be approximated as C ≈ V U 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-graining the distance matrix to be low-rank. As such, we present some of the details and algorithm of Indyk et al. (2019) as an 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 a lower-bound on the number of entries which must be read. Theorem I.1. Indyk et al. (2019) There is a randomized algorithm that, given a distance matrix C ∈ R n×m , reads O((n + m)r/ε) entries of C, runs in time Õ(n + m) • poly(r, 1/ε)1 and computes low-rank factors V ∈ R n×r , U ∈ R r×m that with probability 0.99 satisfy:EQUATION)For C r the optimal rank-r approximation of C.This is a remarkable result, especially in light of the next theorem.Theorem I.2. Indyk et al. (2019) Let r ≤ m ≤ n and ε > 0 such that r/ε = O(min{m, n 1/3 }).Any randomized and possibly adaptive algorithm that given a distance matrixC ∈ R n×m computes V ∈ R n×r , U ∈ R r×m satisfying ∥C -V U ∥ 2 F ≤ ∥C -C r ∥ 2 F + ε∥C∥ 2 F must read Ω((n + m)r/ε) entries of C in expectation.This lower bound also holds for symmetric distance matrices C ∈ S n .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 I.1 for completeness.Algorithm 8 Low-Rank approximation for distance matrix CInput point sets {x i } n i=1 , {y j } M j=1 in metric space X and metric d Pick indices i * ∈ [n], j * ∈ [m] uniformly at random for i = 1 to n do Update sample probability p i = d(x i , y j * ) 2 + d(x i * , y j * ) 2 + 1 m m j=1 d(x i * , y j ) 2 end for Sample O(r/ε) rows C i,. ∼ Categorical pi i pi Compute U using Frieze et al. (2004) Compute V using Chen & Price (2017) return V , UThis algorithm relies on two previous works, whose main results we summarize here.Theorem I.3. Frieze et al. (2004) Let C ∈ R n×m . For a sample of O(r/ε) rows according to a distribution p ∈ ∆ n which satisfies p i ≥ Ω(1)∥C i,. ∥ 2 2 /∥C∥ 2 F for i ∈ [n].Then in O(mr/ϵ + poly(r, 1/ε)) time one may compute a matrix U ∈ R r×m from this sample which satisfies:∥C -CU T U ∥ 2 F ≤ ∥C -C k ∥ 2 F + ε∥C∥ 2 FWith probability 0.99.Thus, to compute the first low-rank factor U , we need to ensure the p i generated from the algorithm satisfies this requirement and offer the (short) proof below.Proof. First, it is helpful to noted(x, y) 2 ≤ (d(x, z) + d(z, y)) 2 = d(x, z) 2 + 2d(x, z)d(y, z) + d(y, z) 2 ≤ 2(d(x, z) 2 + d(y, z) 2 ).Where in the last step one uses AM-GM where i a 1/n i ≤ i ai n . Rewriting the norm of row i we have:∥C i,. ∥ 2 2 = m j=1 d(x i , y j ) 2 ≤ 2 m j=1 d(x i , y j * ) 2 + d(y j * , y j ) 2 = 2md(x i , y j * ) 2 + 2 m j=1 d(y j * , y j ) 2 ≤ 2md(x i , y j * ) 2 + 4 m j=1 d(y j * , x i * ) 2 + d(y j , x i * ) 2 = 2md(x i , y j * ) 2 + 4md(y j * , x i * ) 2 + 4 m j=1 d(y j , x i * ) 2 = 4m   1 2 d(x i , y j * ) 2 + d(y j * , x i * ) 2 + 1 m m j=1 d(y j , x i * ) 2   ≤ 4mp iAs we have the re-normalization p i ← pi i pi before sampling, we need to consider the value of the expectation E[ i p i ] to conclude.E n i=1 p i = n i=1 E   d(x i , y j * ) 2 + d(x i * , y j * ) 2 + 1 m m j=1 d(x i * , y j ) 2   = n i=1 E j * ∼[m] d(x i , y j * ) 2 + E i * ∼[n],j * ∼[m] d(x i * , y j * ) 2 + 1 m m j=1 E i * ∼[n] d(x i * , y j ) 2 = n i=1 1 m m j=1 d(x i , y j ) + n n i=1 m j=1 1 nm d(x i , y j ) + n m m j=1 1 n n i=1 d(x i , y j ) = 3 m n i=1 m j=1 d(x i , y j ) = 3 m ∥C∥ 2F By an application of Markov's inequality we have that:P   n i=1 p i -1 ≥ 3∥C∥ 2 F mδ -1   = P n i=1 p i ≤ 3∥C∥ 2 F mδ ≥ 1 - E[ n i=1 p i ] 3∥C∥ 2 F mδ = 1 -δThus with probability 1δ we have:p i n i ′ =1 p i ′ ≥ mp i δ 3∥C∥ 2 F = 4mp i δ 12∥C∥ 2 F ≥ ∥C i,. ∥ 2 2 δ 12∥C∥ 2 F = Ω(δ) ∥C i,. ∥ 2 2 ∥C∥ 2 FThis indicates the algorithm presented has probabilities with an appropriate bound for using the algorithm of Frieze et al. (2004) to sample the O(r/ε) rows of C and generate a rank-r factor U .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 U . Theorem I.4. Chen & Price (2017) There is a randomized algorithm A, given matrices C ∈ R n×m , U ∈ R r×m reads only O(r/ε) columns of C with time-complexity O(mr) + poly(r/ε) and returns V ∈ R n×r which satisfies∥C -V U ∥ 2 F ≤ (1 + ε) min Z∈R n×r ∥C -ZU ∥ 2F with probability 0.99.Thus, using the result of Chen and Price Chen & Price (2017), one may find a satisfying V easily for a fixed U , C. In particular, Indyk 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∥C -V U ∥ 2 F ≤ (1 + ε) min Z ∥C -ZU ∥ 2 F ≤ (1 + ε)∥C -CU T U ∥ 2 F ≤ (1 + ε)(∥C -C r ∥ 2 F + ε∥C∥ 2 F ) = ∥C -C r ∥ 2 F + ε∥C -C r ∥ 2 F + (1 + ε)ε∥C∥ 2 F ≤ ∥C -C r ∥ 2 F + (1 + 2ε)ε∥C∥ 2 F ≤ ∥C -C r ∥ 2 F + ε∥C∥ 2 FWhich achieves the bound up to a constant scaling of ε and shows the result of Indyk et al. (2019) . 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.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-divergence of some function F (x), defined by the first-order Taylor expansion of F : D F (x, y) = F (x) -F (y) + ∇F (y) T (xy) For the negative entropy function -H(ξ), this corresponds to the KL-divergence KL(ζ | ξ). We introduce the iterative-Bregman projection algorithm in the context of the KL-divergence owing to the direct connection with entropically-regularized optimal transport. Definition J.1. Iterative Bregman ProjectionsBregman (1967) SupposeC = ∩ L l=1 C l is an intersection of closed convex sets {C l } L l=1 .For n > L, let the indexing be L-periodic as C n := C n mod L . Suppose we want to find a minimizer ζ of the KL-divergence with some positive vector ξ ∈ R n×m + 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 ζ∈C KL(ζ∥ξ)Where the projection of ξ onto the set C is denoted byζ * = arg min ζ∈C KL(ζ∥ξ) := P KL C (ξ)Supposing that each C l forms a quotient space C l = V /U for a subspace U ⊂ V , defined asV /U = {v + U | v ∈ V } 2 , the iterative Bregman projection algorithm alternates projections onto each set C n as ζ (n) ← P KL Cn (ζ (n-1) ) (43) starting from ζ (0) = ξ.One may show Bregman (1967) the convergence of Bregman projections to the unique minimizer in C, ζ * , where we have the guarantee that ζ (n) → ζ * as n → ∞. 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 a few more notions are required. Definition J.2. Dykstra's AlgorithmDykstra (1983) Given a point x 0 ∈ E for E a Euclidean space3 , to find the unique point in C = ∩ L l=1 C l for closed, convex sets C l that minimize the distance to x 0 asx * = arg minx∈C ∥x -x 0 ∥ 2One may initialize the residuals q -(L-1) = ... = q 0 = 0 and apply the algorithm:x n = P Cn (x n-1 + q n-1 ) q n = (x n-1x n ) + q n-L Where C n := C n mod L as before and P Cn denotes the projection operator onto the convex set C n .One can note that Dykstra's algorithm for projections onto intersections of convex sets no longer relies on the assumption that the set is a quotient space, and applies in the case that it is merely closed under convex-combinations. To generalize Dykstra's for more general functions that than the ℓ 2 -norm, one may define the projection with respect to the Bregman divergence of a cost function F and define the projection by the minimization: P Cn := arg min x D F (x, y). It was proven in Bauschke & Lewis (2000) that the generalized form of Dykstra's iterations are given as:x n = P Cn (∇F * (∇F (x n-1 ) + q n-1 )) q n = (∇F (x n-1 ) -∇F (x n )) + q n-L For F * denoting the Fenchel-conjugate of F , which we define later in connection to the low-rank dual problem. Notably, Bauschke & Lewis (2000) also provided guarantees of convergence to the optimal solution which extend to the case that F is the negative entropy and D F the KL-divergence. These constitute Dykstra's algorithm with cyclic Bregman projections, and project a point to the closest point in the intersection of convex sets C = ∩ L l=1 C l for an arbitrary cost function F and its associated Bregman-divergence D F . We tested on two additional synthetic datasets, both used as benchmarking datasets in Scetbon et al. (2021) . We follow exactly the parameters provided in Scetbon et al. (2021) to simulate these datasets. For the first one, we simulated n = m = 10, 000 points from two Gaussian mixtures in 2D (Fig. 6 ). The first Gaussian mixture is a mixture of three Gaussian distributions with means (0, 0), (0, 1), (1, 1) respectively. The mixture proportion is 1 3 and the covariance is 0.05 × identity for each Gaussian. The 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 1 2 and the covariance is 0.05 × identity for each Gaussian. For the second dataset, we simulated n = m = 5, 000 points from two Gaussian mixtures in 10D. The first Gaussian mixture is a mixture of three Gaussian distributions with means (0, 0, 0, • • • , 0), (0, 1, 0, • • • , 0), (1, 1, 0, • • • , 0) respectively. The mixture proportions is 1 3 and the covariance is 0.05 × identity for each Gaussian. The second Gaussian mixture is a mixture of two Gaussian distributions with means (0.5, 0.5, 0, • • • , 0), (-0.5, 0.5, 0, • • • , 0) respectively. The mixture proportions is 1 2 and the covariance is 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 as cost to find a low-rank coupling matrix between the two Gaussian mixtures with rank between 50 and 200. We observe the same pattern as Fig. 2b . FRLC obtains lower transport cost with increasing rank, and achieves lower cost for each rank than LOT under all initializations (Fig. 23c , Fig. 7 ).We 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, we represent the graph as G = (A, h), where A ∈ R n×n encodes the intra-graph node relationship (e.g. adjacency matrix) and h ∈ ∆ n is a uniform measure. We cluster the nodes of G by estimating a GW coupling 2 : Runtime of FRLC and LOT on the synthetic datasets of 1000 samples, as well as cost value ⟨C, P ⟩ F and marginal tightness for context. This was done with the FRLC setting max_inneriters_balanced=1000, max_inneriters_relaxed=50, min_iter=7 and rank r = 100. This time excludes the extra time incurred for ott-jax problem setup and includes the setup time for FRLC.between G and a smaller graph G = (B, h), where B ∈ R m×m represents the relationship between each of the m clusters and h ∈ ∆ m is the proportion of nodes of G in each cluster. Without a priori knowledge, h is set to be uniform and B is set as 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. 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 on four real-world graphs: a Wikipedia hyperlink network with 15 webpage categories Yang & Leskovec (2012) ; an email interaction network within a European institute with 42 departments Yin et al. (2017) ; an Amazon product network with 12 product categories Yang & Leskovec (2012) ; and a network of interactions between 12 Indian villages Banerjee et al. (2013) . We also test on the symmetric and noisy versions of each graph provided by Chowdhury & Needham (2021) . We compare with two OT-based methods: (1) GWL Xu et al. (2019) , which solves a balanced GW problem between G and G with the adjacency matrix of G as the intra-domain cost matrix A; (2) SpecGWL Chowdhury & Needham (2021) 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-GW) and heat kernel (denoted SpecFRLC-SR-GW). 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 9 out of 12 datasets. When using the heat kernel, our semi-relaxed algorithm achieves better clustering result than SpecGWL 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.As discussed in Vincent-Cuaz et 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 ∈ R n×n encodes the intra-graph node relationship (e.g. adjacency matrix, Laplacian) and h ∈ ∆n is a uniform measure over the nodes. If we want to cluster the nodes of G into m clusters, we can define a new graph G = (B, h), where B ∈ R m×m is a diagonal matrix representing the cluster's connections and h is a distribution over clusters estimating the proportion of nodes in G in each cluster. Since we don't know the density of clusters a 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 in G, with intra-domain cost matrices A and B:EQUATIONThe cluster assignment of each node in G can then be recovered from P by finding the node in G mapped to it with the maximum weight.However, since the proportion of nodes in G in each cluster is not known a priori, solving a balanced GW problem fixing the marginal of P on G to be a uniform h significantly constrains the expressivity of the algorithm. Therefore, as proposed in Vincent-Cuaz et al. ( 2022), we can instead solve a semi-relaxed GW problem with the right marginal relaxed, fixing the marginal on G to be h but allowing the marginal on G to deviate from h:EQUATIONThe learned h from semi-relaxed GW estimates the posterior proportion of nodes in G in each of the m clusters.We run our semi-relaxed FRLC algorithm on four real-world graph datasets: a Wikipedia hyperlink network with 1998 nodes and 15 clusters Yang & Leskovec (2012) , a directed graph of email interactions in a European research institute with 1005 nodes and 42 clusters Yin et al. (2017) , an Amazon product network with 1501 nodes and 12 clusters Yang & Leskovec (2012) , and a 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 to 10% additional edges following Chowdhury & Needham (2021) , leading to a total of 12 different graphs to cluster.We compare our algorithm with two baseline methods, GWL and SpecGWL. GWL Xu et al. (2019) solves the entropy-regularized version of the balanced GW problem of ( 46) with the adjacency matrix of G as A. We set h such that the density of each node is proportional to its degree. We set h to be a distribution estimated by sorting the weights of h and sampling m values via linear interpolation, following Chowdhury & Needham (2021) . We set B = diag(h). We set the entropy regularization 2021) on a spatial-transcriptomics dataset of mouse embryonic development Chen et al. (2022) . FRLC demonstrates a more robust trend of improved performance with higher rank. η = 10 -6 . SpecGWL Chowdhury & Needham (2021) solves the same problem as GWL, but instead of using the adjacency matrix as A, it uses the heat kernel on the normalized graph Laplacians Chung (2005) . We set the heat parameter t = 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 as C and label the result of the two representations as FRLC-SR-GW and SpecFRLC-SR-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.In this problem, we use FRLC to find a low-rank alignment matrix between cells from two spatial transcriptomics (ST) Ståhl et al. (2016) slices, collected at two timepoints, then use the computed alignment matrix for two downstream prediction tasks. An ST experiment on a 2D tissue slice yields a pair (X, Z). X ∈ R n×p is the gene expression matrix, where n is the number of cells on the slice and p is the number of genes measured. X ij ∈ R is the gene expression level of gene j in cell i, where a higher number indicates stronger expression. Z ∈ R n×2 is the spatial coordinate matrix, where each row i stores the x-y coordinate of cell i on the slice. Therefore, each cell on the slice has a gene expression vector of length p, which encodes the feature of 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 a pair of ST slices S 0 = (X 0 , Z 0 ), S 1 = (X 1 , Z 1 ), with n and m cells, of the same tissue region. We assume S 0 is collected at timepoint t = 0 and S 1 is collected at timepoint t = 1, hence the transition from S 0 to S 1 reflects the biological development of the tissue during the time period. We would like to find the ancestor-descendant relationship between cells from S 0 and S 1 by computing an optimal transport coupling matrix between cells from S 0 and cells from S 1 . The state-of-the-art spatial transcriptomics alignment method moscot Klein et al. (2023) claims that unbalanced OT is the most appropriate setup for this problem. Specifically, given discrete uniform measure a and b over cells from S 0 and S 1 , we solve the unbalanced Wasserstein problemmin P ∈R n×m + ⟨C, P ⟩ + τ KL(P 1 m ∥a) + τ KL(P T 1 n ∥b) (48) C ∈ R n×m + has entries C ij = c(X 0 i , X 1 j ), where c is the Euclidean distance between the features of cell i from S 0 and cell j from S 1 .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 problemEQUATIONwhere A ∈ R n×n is the Euclidean distance matrix between the spatial location of cells within S 0 , and B ∈ R m×m is the Euclidean distance matrix between the spatial location of cells within S 1 .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 1 should have an ancestor from S 1 , but not all cells from S 1 need to have a descendent in S 0 because of cell death. Therefore, the most natural OT task for this problem is semi-relaxed OT with the left marginal (the marginal on the first/ancestor slice) relaxed. Specifically, we can also solve the semi-relaxed Wasserstein problemEQUATIONas well as semi-relaxed GW and FGW problem.Gene expression prediction task Given the alignment matrix P linking cells from S 0 to S 1 , we can predict properties of cells in S 1 from properties of cells in S 0 . Let the expression of a gene j in S 0 be a vector f j ∈ R n , such that f ji is the expression level of gene j in cell i, we can predict the expression of gene j in S 1 as f j = m × P T × f j ∈ R m . The accuracy of the prediction can be measured by the Spearman correlation between the predicted expression and the ground truth expression f j of gene j in S 1 : ρ( f j , f j ). In this work, we test the prediction accuracy on 10 test genes: Tubb2b, Pantr1, Actc1, Tnni1, 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 0 to predict the cell type labels of cells in S 1 . Specifically, for each cell j in S 1 , we can assign it the type of the cell argmax i P ij in S 0 . 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 1 and the ground truth clustering.We use the large-scale real-world dataset of mouse embryo development Chen et 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 E11.5 and E12.5 embryos, consisting of 30,124 cells and 51,365 cells, respectively. We preprocess the dataset using the standard SCANPY Wolf et 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 log-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. We compare with the unbalanced low-rank optimal transport algorithm of Scetbon et al. (2023) , the backbone of the spatial transcriptomics alignment method moscot Klein et al. (2023) , which was shown by Scetbon et al. (2023) ; Klein et al. (2023) to achieve state-of-the art performance on spatial transcriptomics alignment. We use Scetbon et 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 a 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, a full-rank solution is not feasible, hence we solve for low-rank solutions with the rank validated. We also solve semi-relaxed versions of Wasserstein (result denoted as FRLC-SR-W), GW (FRLC-SR-GW), FGW (FRLC-SR-FGW) problems using our FRLC algorithm, as well as using a particular setting of LOT-U that is equivalent to semi-relaxed solver (results denoted as LOT-SR-W, LOT-SR-GW, LOT-SR-FGW).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 . 7 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, Myl4, Tnnt2, Apoa2, Hba-x, Tubb3, Epha7, Ldha, Col1a2, which are marker genes of various cell types as well.We report the performance of the alignment computed by each method using the Spearman correlation on the gene expression prediction task for 10 test genes, as well as the ARI and AMI on the cell type prediction task. Fig. 10 visualizes the ground truth cell type classification versus the classification predicted by our method FRLC-SR-W.We report the runtime and OT cost of FRLC and LOT Scetbon et al. (2021) on this dataset (mouse embryo E11.5-12.5) as well as two other datasets (E9.5-10.5, E10.5-11.5) from Chen et al. (2022) in Table 4 . For all three datasets, FRLC achieves a better OT cost in a shorter time. O Discussion of differences between FRLC and existing low-rank optimal transport algorithmsWe provide an extensive comparison of FRLC against the parametrization and objective of Scetbon et al. (2021) in Appendix A. As Latent-OT of Lin et al. ( 2021) is an extension of the k-Wasserstein barycenter problem, it has a distinct objective and thus performs worse on primal OT cost (Table 5 ), making a direct experimental comparison on primal cost minimization only appropriate relative to the works of Scetbon et al. (2021 Scetbon et al. ( , 2023 Scetbon et al. ( , 2022)) . This stated, we 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 et al. (2021) 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 x , z y ) 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 et al. are:[C Q ] ik = ∥x i -z x k ∥ 2 2 , [C R ] jℓ = ∥y j -z y ℓ ∥ 2 2 , [C T ] kℓ = ∥z x k -z y ℓ ∥ 2 2, optionally using a Wasserstein distance for the entries of C T . (3) FRLC 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 min P ∈Πr(a,b) ⟨C, P ⟩ F and differ substantially from those of the proxy objective in Lin et al. (2021); Forrow et al. (2019) . ( 4) The different objectives and variables lead to very different algorithms: Lin et al. 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 . (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 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 x k ) between x i and latent anchor z x k . 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. (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-Wasserstein (GW), or Fused GW, which FRLC readily generalizes to. A summary of the existing low-rank OT algorithms and key distinctions between them is given in Table 1 .For completeness, we offer a compare against the work Latent OT Lin 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 on various simulated datasets in Table 5 . Our method introduces an additional hyperparameter τ relative to previous approaches Scetbon et al. (2021) , controlling the strength of the KL penalty on the inner marginals when updating Q and R.Empirically, we found FRLC to be robust to different choices of τ , but applying the method optimally requires this additional hyperparamter in any grid-search.We also note that the non-asymptotic criterion ∆(•, •) is weak relative to stronger notions of convergence, and that often users might prefer to simply run the method up to some number of maximal iterations by setting the parameter for whether ∆(•, •) is used to False. The W optimization empirically converges to minima smoothly, so for the most part there is not much of a need for ∆ 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 minimum number of iterations should be very low (around 5) for simple datasets and substantially higher for high-dimensional, structured ones.Although we demonstrate strong performance already, there is massive room for improvement as our implementation is preliminary and not at the level of a high-performance library like ott-jax. We use lightweight vanilla implementations of Sinkhorn as a sub-routine, not taking advantage of the momentum-based techniques which could accelerate it massively. Thus, one can imagine that the potential scalability and speed of this method could be much higher than reported in this document.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. Answer: [Yes] Justification: The optimization is deterministic, so generally there's no need for reporting statistical significance. One small exception is that one of our proposed initializations is randomized, and for experiments which use it we do include ±1σ error bars. Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence 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, or 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).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. 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 to result from this work on optimal transport. Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.• The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in 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 in Section Q: FRLC can be used for scaling an interpretability in any ML method using OT as a 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 not tied to particular applications, let alone deployments. However, if there is a 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, it 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 or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).Where Õ(•) hides poly-log factors.This is to say, x, y ∈ C l =⇒ c1x + c2y ∈ V /U for any c1, c2 ∈ R.e.g. R, R d , R n×m , etc. 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: **Presentation and logical organization of the ideas**

I 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. 

Furthermore, 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? 

**Comparison and empirical validation**
The 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. 

I found some other issues lacking in a more in-depth discussion of the computational complexity, possible implementation difficulties, and overall limitations in the proposal.
questions: 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)?
- 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? 
- What are the theoretical complexity of the algorithms, and what practical considerations are relevant when implementing the method?

One 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: limitations: I found it necessary to discuss in more depth the limitations and the tradeoffs present in the proposal compared with previous art.
paper_summary: paper_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.
strength_weakness: strength_weakness: While the empirical results are good, the paper could benefit from more extensive comparisons with additional baseline methods to solidify the claims of superiority.

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.
questions: 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. 

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.
limitations: limitations: N.A.
paper_summary: paper_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.
strength_weakness: strength_weakness: * 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. 
* In algorithm 1, it would be interesting to have the computational complexity of each step (presented in the algorithm directly).
questions: questions: * How to choose $r$ in practical scenarios ? 
* 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: limitations: /
paper_summary: paper_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.

The 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.
strength_weakness: strength_weakness: 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?
questions: questions: - line 103: Would it be good to say the what complexity "massively improved" to?
- 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.)
- 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: 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?

Comments: comment: I thank the authors for the detailed response. I will keep the score.
Thank for your replies. I will increase my score from 5 to 6.
Dear reviewer jx3E, we thank you for your comments and are glad to see that this response helped clarify a few of your questions. We agree that placing this work in context is essential, and will add the table to the main text (rather than the appendix) following your suggestion. If you have any further questions, we are more than happy to offer any clarifications during the remainder of the discussion period.
Thank you very much for your answer.
I would like to keep my score.
Thank you for carefully engaging with the main points of my review. I believe the shared concern with other reviewers regarding the presentation of the idea and how it differs from previous art is fundamental to improving the quality of this paper. I believe the table with the comparisons could be in the main text; those technical differences and how they map to papers and methods could be immediate for those working directly with these problems, but it can easily be lost for the larger community, although being relevant.  

I am satisfied with the complexity analysis and extended empirical discussions.