Title: Low-Rank Optimal Transport through Factor Relaxation with Latent Coupling
Venue: NeurIPS 2024
Accepted: yes
Keywords: Optimal Transport

Sinkhorn

Low-Rank

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

Low-Rank Optimal Transport through Factor Relaxation with Latent Coupling

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.

1 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 versatility different ground costs result in different ways to compare data - OT has found many applications in machine learning and beyond: from self-attention Tay et al. (2020); Sander et al. (2022); Geshkovski 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, 2022, 2023) derived algorithms to compute low-rank optimal transport plans for the primal OT problem with general costs, extending low-rank OT to GW and unbalanced problems using factored couplings.

Interestingly, a different factorization of P was proposed by Lin et al. (2021) in the context of k -Wasserstein barycenters. We call their factorization a latent coupling (LC) factorization, given by P = Q diag(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, 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: µ = ∑ n i =1 a i δ x i and ν = ∑ m j =1 b j δ y j . Let

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 :

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, 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:

The Fused Gromov-Wasserstein (FGW) objective function Vayer et al. (2020) is a convex combination of the W and GW objectives, given as α ⟨ C , P ⟩ F +(1 -α ) Q A , B ( P ) , for hyperparameter α ∈ (0 , 1) .

Relaxed marginal constraints. Balanced OT (1) constrains P to lie in Π a , b . Unbalanced OT relaxes constraints P 1 m = a and P T 1 n = b , replacing them with penalties in the form of KL divergences (or other divergences, see Chizat et al. (2018)):

where τ L , τ R > 0 control the strength of each penalty. Semi-relaxed optimal transport relaxes exactly one of the hard constraints P 1 m = a and P T 1 n = b in the same manner. The semi-relaxed version of (1) obtained by relaxing only the 'right' marginal constraint on b is:

while it's 'left' marginal counterpart SR L -W ( µ, ν ) is defined analogously over P ∈ Π · , b , using penalty τ KL( P 1 m ∥ a ) . Likewise, one can form semi-relaxed or unbalanced GW and FGW problems.

Entropy regularization. The seminal work Cuturi (2013b) introduced the Sinkhorn algorithm to solve an entropy regularized version of (1), W ϵ ( µ, ν ) := min P ∈ Π a , b ⟨ C , P ⟩ F -ϵH ( P ) , massively improving the O ( n 3 log n ) time complexity of classical techniques Orlin (1997); Tarjan (1997). Above, H is the entropy, H ( P ) = -∑ ij P ij (log P ij -1) , and ϵ > 0 is the regularization strength.

Low-rank regularization. The nonnegative rank rk + ( M ) of matrix M is the least number of nonnegative rank-one matrices summing to M . For r ≥ 1 , define

and let Π a , b ( r ) = Π a , · ( r ) ∩ Π · , b ( r ) . To estimate Wasserstein distances with greater stability and accuracy under sampling noise, Forrow et al. (2019) proposed a low-rank regularization on the coupling matrix, factoring the transport through a small set of anchor points. More explicitly, Scetbon et al. (2021) parameterized the set as Π a , b ( r ) through the set FC a , b ( r ) of factored couplings ,

The set FC a , b ( r ) parameterizes Π a , b ( r ) through ( Q , R , g ) ↦→ Q diag(1 / g ) R T , as shown by Cohen &Rothblum (1993).

Scetbon et al. (2021) apply this factorization to solve the Wasserstein problem subject to P ∈ Π a , b ( r ) for general cost matrices:

GW, unbalanced and semi-relaxed low-rank OT problems are defined as in (2), (3) and (4), replacing R n × m + , Π a , · , or Π · , b with rank-constrained counterparts (5). Scetbon & Cuturi (2022); Scetbon et al. (2022, 2023) developed a robust framework for solving all of these problems.

3 Factor Relaxation with Latent Coupling (FRLC) algorithm

3.1 Latent Coupling Factorization

We parameterize low-rank coupling matrices P ∈ Π a , b ( r ) using a factorization introduced in Lin et al. (2021), which we call the latent coupling (LC) factorization (Fig. 1). The key property of

!

"

!

!

!

"

#

$

%

#

$

&

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 = Q diag(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 call the factors Q , R , T in an LC factorization sub-couplings . Let R + := R n × r + × R m × r + × R r × r + . Given probability vectors a ∈ ∆ n , b ∈ ∆ m and a positive integer rank r , let

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 ) :

Since 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 ) ↦→ Q diag(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 that Q ∈ 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 .

3.2 The Balanced FRLC Algorithm

We introduce an algorithm Factor Relaxation with Latent Coupling (FRLC), to compute a LC factorization of minimum cost. We first describe the FRLC algorithm for the balanced Wasserstein problem. Extensions to other and marginal constraints are discussed later. The FRLC objective function, for low-rank, balanced Wasserstein OT, is

where 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:

The convex sets constraining the latent coupling matrix T are

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, 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:

The Sinkhorn kernels for the semi-relaxed OT problems arising from the factor relaxation step are:

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 Qk +1 , g R k +1 , after the factor relaxation step, we next take a coordinate MD step on the latent coupling T :

This is equivalent to applying Sinkhorn (Algorithm 5) to T given g Q and g R with the kernel:

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 ) remain in 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.

3.3 Initialization, convergence, and FRLC extensions

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

Convergence 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 ˜ O ( 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 { x i k } p i =0 where x k := x 0 k for k ∈ [ N ] outer iterations. Then one has:

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:

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 x k = ( Q k , R k , T k ) has non-asymptotic stationary convergence in the criterion ∆( · , · ) with:

(a)

(b)

(c)

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.

The 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 PB ⊙ 2 P T 1 n -2 ⟨ APB , P ⟩ , where ⊙ denotes the Hadamard (entrywise) product. Then the GW-specific Sinkhorn kernels are

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

4 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, 2023, 2022) and Latent OT refers to Lin et al. (2021).

(a)

(d)

(c)

(b)

FRLC (Rank 5)

FRLC (Rank 10)

LOT (Rank 5)

LOT (Rank 10)

(e)

Ground Truth (Full rank)

4.1 Evaluation of Low-rank Approximations for Balanced OT on Synthetic Data

We first compare the balanced OT version of FRLC with the the low-rank balanced OT algorithm LOT of Scetbon et al. (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).

4.2 Interpretation of the Latent Coupling and LC-Projection

We demonstrate the intepretability of the latent coupling T in the LC factorization. In both the LC factorization and factored couplings, the sub-couplings Q and R each have associated barycentric projection operators which coarse-grain input datasets Z (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 r 1 × r 2 + . 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

(a)

Objective

Algorithm

Spearman

ρ

ARI

AMI

W

LOT-U

0.394

0.332

0.397

FRLC-U

0.465

0.466

0.484

LOT-SR

0.391

0.319

0.396

FRLC-SR

0.467

0.475

0.492

GW

LOT-U

0.005

0.0

0.0

FRLC-U

0.266

0.299

0.364

LOT-SR

0.004

0.0

0.0

FRLC-SR

0.275

0.318

0.381

FGW

LOT-U

0.391

0.330

0.400

FRLC-U

0.368

0.391

0.420

LOT-SR

0.396

0.337

0.401

FRLC-SR

0.465

0.469

0.499

Tubb2bExpression

Tubb2bPredicted

1.0

0.8

0.6

0.4

0.2

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

4.3 Evaluation on Spatial Transcriptomics Alignment

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

Table 1: Comparing aspects of low-rank OT methods. Factorization indicates the structure of the inner matrix.

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

4.4 Additional Experiments

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

5 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).

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

Acknowledgments and Disclosure of Funding

This work is supported by NCI grant U24CA248453 to B.J.R. J.G. gratefully acknowledges support from the Schmidt DataX Fund at Princeton University made possible through a major gift from the Schmidt Futures Foundation.

References

Bakshi, A. and Woodruff, D. Sublinear Time Low-Rank Approximation of Distance Matrices. Advances in Neural Information Processing Systems , 31, 2018.

Banerjee, A., Chandrasekhar, A. G., Duflo, E., and Jackson, M. O. The Diffusion of Microfinance. Science , 341(6144):1236498, 2013.

Bauschke, H. H. and Lewis, A. S. Dykstras algorithm with Bregman projections: A convergence proof. Optimization , 48(4):409-427, January 2000. ISSN 1029-4945. doi: 10.1080/02331930008844513. URL http://dx.doi.org/10.1080/02331930008844513 .

Beck, A. and Tetruashvili, L. On the Convergence of Block Coordinate Descent Type Methods. SIAM J. Optim. , 23:2037-2060, 2013. URL https://api.semanticscholar.org/CorpusID: 6866704 .

Benamou, J.-D., Carlier, G., Cuturi, M., Nenna, L., and Peyré, G. Iterative Bregman Projections for Regularized Transportation Problems. SIAM Journal on Scientific Computing , 37 (2):A1111-A1138, January 2015. ISSN 1095-7197. doi: 10.1137/141000439. URL http: //dx.doi.org/10.1137/141000439 .

Bregman, L. M. The relaxation method of finding the common point of convex sets and its application to the solution of problems in convex programming. USSR computational mathematics and mathematical physics , 7(3):200-217, 1967.

Bunne, C., Stark, S. G., Gut, G., del Castillo, J. S., Levesque, M., Lehmann, K.-V., Pelkmans, L., Krause, A., and Rätsch, G. Learning single-cell perturbation responses using neural optimal transport. Nature Methods , 20(11):1759-1768, September 2023. ISSN 1548-7105. doi: 10.1038/ s41592-023-01969-x. URL http://dx.doi.org/10.1038/s41592-023-01969-x .

Charikar, M., Chen, B., Ré, C., and Waingarten, E. Fast Algorithms for a New Relaxation of Optimal Transport. In Neu, G. and Rosasco, L. (eds.), Proceedings of Thirty Sixth Conference on Learning Theory , volume 195 of Proceedings of Machine Learning Research , pp. 4831-4862. PMLR, 12-15 Jul 2023. URL https://proceedings.mlr.press/v195/charikar23a.html .

Chen, A., Liao, S., Cheng, M., Ma, K., Wu, L., Lai, Y., Qiu, X., Yang, J., Xu, J., Hao, S., et al. Spatiotemporal transcriptomic atlas of mouse organogenesis using DNA nanoball-patterned arrays. Cell , 185(10):1777-1792, 2022.

Chen, X. and Price, E. Condition number-free query and active learning of linear families. 2017.

Chizat, L., Peyré, G., Schmitzer, B., and Vialard, F.-X. Unbalanced Optimal Transport: Dynamic and Kantorovich Formulations. Journal of Functional Analysis , 274(11):3090-3123, June 2018. ISSN 0022-1236. doi: 10.1016/j.jfa.2018.03.008. URL http://dx.doi.org/10.1016/j.jfa.2018. 03.008 .

Chowdhury, S. and Needham, T. Generalized Spectral Clustering via Gromov-Wasserstein Learning. In International Conference on Artificial Intelligence and Statistics , pp. 712-720. PMLR, 2021.

Chung, F. Laplacians and the Cheeger Inequality for Directed Graphs. Annals of Combinatorics , 9: 1-19, 2005.

Cohen, J. E. and Rothblum, U. G. Nonnegative Ranks, Decompositions, and Factorizations of Nonnegative Matrices. Linear Algebra and its Applications , 190:149-168, 1993.

Courty, N., Flamary, R., and Tuia, D. Domain adaptation with regularized optimal transport. In Machine Learning and Knowledge Discovery in Databases: European Conference, ECML PKDD 2014, Nancy, France, September 15-19, 2014. Proceedings, Part I 14 , pp. 274-289. Springer, 2014.

Cuturi, M. Sinkhorn Distances: Lightspeed Computation of Optimal Transport. Advances in neural information processing systems , 26, 2013a.

Cuturi, M. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in Neural Information Processing Systems , pp. 2292-2300, 2013b. URL https://proceedings.neurips. cc/paper/2013/hash/af21d0c97db2e27e13572cbf59eb343d-Abstract.html .

Cuturi, M., Meng-Papaxanthos, L., Tian, Y., Bunne, C., Davis, G., and Teboul, O. Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein. arXiv preprint arXiv:2201.12324 , 2022.

Dang, C. D. and Lan, G. Stochastic Block Mirror Descent Methods for Nonsmooth and Stochastic Optimization. SIAM J. Optim. , 25(2):856-881, January 2015.

Dong, S., Pan, Z., Fu, Y., Xu, D., Shi, K., Yang, Q., Shi, Y ., and Zhuo, C. Partial Unbalanced Feature Transport for Cross-Modality Cardiac Image Segmentation. IEEE Transactions on Medical Imaging , 2023.

Dykstra, R. L. An Algorithm for Restricted Least Squares Regression. Journal of the American Statistical Association , 78(384):837-842, December 1983. ISSN 1537-274X. doi: 10.1080/ 01621459.1983.10477029. URL http://dx.doi.org/10.1080/01621459.1983.10477029 .

Forrow, A., Hütter, J.-C., Nitzan, M., Rigollet, P., Schiebinger, G., and Weed, J. Statistical Optimal Transport via Factored Couplings. In Chaudhuri, K. and Sugiyama, M. (eds.), Proceedings of the Twenty-Second International Conference on Artificial Intelligence and Statistics , volume 89 of Proceedings of Machine Learning Research , pp. 2454-2465. PMLR, 16-18 Apr 2019. URL https://proceedings.mlr.press/v89/forrow19a.html .

Frieze, A., Kannan, R., and Vempala, S. Fast Monte-Carlo Algorithms for Finding Low-rank Approximations. J. ACM , 51(6):1025-1041, nov 2004. ISSN 0004-5411. doi: 10.1145/1039488. 1039494. URL https://doi.org/10.1145/1039488.1039494 .

Frogner, C., Zhang, C., Mobahi, H., Araya, M., and Poggio, T. A. Learning with a Wasserstein Loss. Advances in neural information processing systems , 28, 2015.

Geshkovski, B., Letrouit, C., Polyanskiy, Y., and Rigollet, P. A mathematical perspective on Transformers. arXiv preprint arXiv:2312.10794 , 2023.

Ghadimi, S., Lan, G., and Zhang, H. Mini-batch stochastic approximation methods for nonconvex stochastic composite optimization. Mathematical Programming , 155(1-2):267-305, December 2014. ISSN 1436-4646. doi: 10.1007/s10107-014-0846-1. URL http://dx.doi.org/10. 1007/s10107-014-0846-1 .

Halmos, P., Liu, X., Gold, J., Chen, F., Ding, L., and Raphael, B. J. DeST-OT: Alignment of Spatiotemporal Transcriptomics Data. In International Conference on Research in Computational Molecular Biology , pp. 434-437. Springer, 2024.

Indyk, P., Vakilian, A., Wagner, T., and Woodruff, D. P. Sample-optimal low-rank approximation of distance matrices. In Beygelzimer, A. and Hsu, D. (eds.), Proceedings of the ThirtySecond Conference on Learning Theory , volume 99 of Proceedings of Machine Learning Research , pp. 1723-1751. PMLR, 25-28 Jun 2019. URL https://proceedings.mlr.press/ v99/indyk19a.html .

Kantorovich, L. On the Translocation of Masses: Doklady akademii nauk ussr. 1942.

Klein, D., Palla, G., Lange, M., Klein, M., Piran, Z., Gander, M., Meng-Papaxanthos, L., Sterr, M., Bastidas-Ponce, A., Tarquis-Medina, M., et al. Mapping cells through time and space with moscot. bioRxiv , pp. 2023-05, 2023.

Lin, C.-H., Azabou, M., and Dyer, E. L. Making transport more robust and interpretable by moving data through a small number of anchor points. Proceedings of machine learning research , 139: 6631, 2021.

Liu, W., Zhang, C., Zheng, N., and Qian, H. Approximating optimal transport via low-rank and sparse factorization. CoRR , abs/2111.06546, 2021. URL https://arxiv.org/abs/2111.06546 .

Liu, X., Zeira, R., and Raphael, B. J. Partial alignment of multislice spatially resolved transcriptomics data. Genome Research , 33(7):1124-1132, 2023.

Mémoli, F. On the use of Gromov-Hausdorff Distances for Shape Comparison. 2007.

Mémoli, F. Gromov-Wasserstein Distances and the Metric Approach to Object Matching. Foundations of computational mathematics , 11:417-487, 2011.

Nesterov, Y. Efficiency of Coordinate Descent Methods on Huge-Scale Optimization Problems. SIAM Journal on Optimization , 22(2):341-362, 2012.

Orlin, J. B. A polynomial time primal network simplex algorithm for minimum cost flows. Mathematical Programming , 78(2):109-129, Aug 1997. ISSN 1436-4646. doi: 10.1007/BF02614365. URL https://link.springer.com/content/pdf/10.1007/BF02614365.pdf .

Pham, K., Le, K., Ho, N., Pham, T., and Bui, H. H. On Unbalanced Optimal Transport: An Analysis of Sinkhorn Algorithm. In International Conference on Machine Learning , 2020. URL https://api.semanticscholar.org/CorpusID:211068892 .

Sander, M. E., Ablin, P., Blondel, M., and Peyré, G. Sinkformers: Transformers with Doubly Stochastic Attention. In Camps-Valls, G., Ruiz, F. J. R., and Valera, I. (eds.), Proceedings of The 25th International Conference on Artificial Intelligence and Statistics , volume 151 of Proceedings of Machine Learning Research , pp. 3515-3530. PMLR, 28-30 Mar 2022. URL https://proceedings.mlr.press/v151/sander22a.html .

Scetbon, M. and Cuturi, M. Low-rank Optimal Transport: Approximation, Statistics and Debiasing. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), Advances in Neural Information Processing Systems , 2022. URL https://openreview.net/forum?id=4btNeXKFAQ .

Scetbon, M., Cuturi, M., and Peyré, G. Low-Rank Sinkhorn Factorization. In International Conference on Machine Learning , 2021. URL https://api.semanticscholar.org/CorpusID: 232147563 .

Scetbon, M., Peyré, G., and Cuturi, M. Linear-time Gromov Wasserstein Distances using Low Rank Couplings and Costs. In International Conference on Machine Learning , pp. 19347-19365. PMLR, 2022.

Scetbon, M., Klein, M., Palla, G., and Cuturi, M. Unbalanced Low-rank Optimal Transport Solvers, 2023.

Schiebinger, G., Shu, J., Tabaka, M., Cleary, B., Subramanian, V., Solomon, A., Gould, J., Liu, S., Lin, S., Berube, P., et al. Optimal-Transport Analysis of Single-Cell Gene Expression Identifies Developmental Trajectories in Reprogramming. Cell , 176(4):928-943, 2019.

Solomon, J., De Goes, F., Peyré, G., Cuturi, M., Butscher, A., Nguyen, A., Du, T., and Guibas, L. Convolutional Wasserstein Distances: Efficient Optimal Transportation on Geometric Domains. ACM Transactions on Graphics (ToG) , 34(4):1-11, 2015.

Ståhl, P. L., Salmén, F., Vickovic, S., Lundmark, A., Navarro, J. F., Magnusson, J., Giacomello, S., Asp, M., Westholm, J. O., Huss, M., et al. Visualization and analysis of gene expression in tissue sections by spatial transcriptomics. Science , 353(6294):78-82, 2016.

Tarjan, R. E. Dynamic trees as search trees via Euler tours, applied to the network simplex algorithm. Mathematical Programming , 78(2):169-177, Aug 1997. ISSN 1436-4646. doi: 10.1007/BF02614369. URL https://link.springer.com/content/pdf/10.1007/ BF02614369.pdf .

Tay, Y., Bahri, D., Yang, L., Metzler, D., and Juan, D.-C. Sparse Sinkhorn Attention. In III, H. D. and Singh, A. (eds.), Proceedings of the 37th International Conference on Machine Learning , volume 119 of Proceedings of Machine Learning Research , pp. 9438-9447. PMLR, 13-18 Jul 2020. URL https://proceedings.mlr.press/v119/tay20a.html .

Tong, A., Malkin, N., Huguet, G., Zhang, Y., Rector-Brooks, J., Fatras, K., Wolf, G., and Bengio, Y. Improving and generalizing flow-based generative models with minibatch optimal transport. In ICML Workshop on New Frontiers in Learning, Control, and Dynamical Systems , 2023.

Vayer, T., Chapel, L., Flamary, R., Tavenard, R., and Courty, N. Fused Gromov-Wasserstein distance for structured objects. Algorithms , 13(9):212, August 2020. ISSN 1999-4893. doi: 10.3390/a13090212. URL http://dx.doi.org/10.3390/a13090212 .

Vincent-Cuaz, C., Flamary, R., Corneli, M., Vayer, T., and Courty, N. Semi-relaxed GromovWasserstein divergence and applications on graphs. In International Conference on Learning Representations , 2022. URL https://openreview.net/forum?id=RShaMexjc-x .

Wolf, F. A., Angerer, P., and Theis, F. J. SCANPY: large-scale single-cell gene expression data analysis. Genome biology , 19:1-5, 2018.

Xu, H., Luo, D., Zha, H., and Duke, L. C. Gromov-Wasserstein Learning for Graph Matching and Node Embedding. In International conference on machine learning , pp. 6932-6941. PMLR, 2019.

Yang, J. and Leskovec, J. Defining and Evaluating Network Communities based on Ground-truth. In Proceedings of the ACM SIGKDD Workshop on Mining Data Semantics , pp. 1-8, 2012.

Yang, K. D., Damodaran, K., Venkatachalapathy, S., Soylemezoglu, A. C., Shivashankar, G., and Uhler, C. Predicting cell lineages using autoencoders and optimal transport. PLoS computational biology , 16(4):e1007828, 2020.

Yin, H., Benson, A. R., Leskovec, J., and Gleich, D. F. Local Higher-Order Graph Clustering. In Proceedings of the 23rd ACM SIGKDD international conference on knowledge discovery and data mining , pp. 555-564, 2017.

Zeira, R., Land, M., Strzalkowski, A., and Raphael, B. J. Alignment and integration of spatial transcriptomics data. Nature Methods , 19(5):567-575, 2022.

A Low-rank optimal transport

A.1 Low-rank factorizations

The 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 :

Let a ∈ ∆ n , b ∈ ∆ m be probability vectors, and let Π a , b ( r ) denote the set of rankr coupling matrices with marginals a and b :

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, 2023) is

Cohen & Rothblum (1993) show that any P ∈ Π a , b ( r ) may be decomposed as P = Q diag(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:

Latent coupling factorization. The latent coupling parameterization of Π a , b ( r ) introduced in Lin et al. (2021), and used in the present work is

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, taking Q ′ ← Q diag(1 / g Q ) T one can refactor

or alternatively taking R ′ = R diag(1 / g R ) T T may refactor as

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.

A.2 Balanced low-rank optimal transport

The FRLC optimization problem Our optimization problem is over the variables ( Q , R , T ) and defined as follows:

where our objective function L LC is

Given ( Q , R , T ) ∈ LC a , b ( r ) , sub-couplings Q and R are constrained by:

while the convex sets constraining the latent coupling matrix T are

C 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. Under these definitions, LC a , b ( r ) = C 1 ∩ C 2 , where

To 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:

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.

By construction, the triple ( Q k +1 , R k +1 , T k +1 ) ∈ C 2 . However, because the set C 1 only constrains 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:

where the objective function L LOT is and where ˜ C = ˜ C 1 ∩ ˜ C 2 with:

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

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

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:

We 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:

After updating Q and R , the update on T then follows a similar form:

leading to balanced constraints on T of the form T 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 2 allows 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, 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 form H ( Q , R , T , g Q , g R ) = ⟨ Q diag(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 , Q diag(1 / g Q ) T diag(1 / g R ) R T ⟩ F and let L LC ( Q , R , T ) be as in (13) . One has the following equivalence:

Thus the semi-relaxed projections yield locally optimal inner marginals.

Proof. To see why this is true, notice that, so long as the outer marginals are tightly satisfied Q 1 r = a and R 1 r = b for Q ≥ 0 n × r , R ≥ 0 m × r , one has

and

Therefore, the inner marginals g ∗ 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 as Q T 1 n = g Q , T 1 r = g Q , and R T 1 m = g R , T T 1 r = g R .

B Block-Coordinate steps for the OT sub-problems

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:

For Q , R in either the appropriate set of couplings or a relaxation thereof (which we will describe shortly). X is related to T by:

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:

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:

A 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:

Without 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:

Thus, without the linearization assumption one may use product rule on X as an implicit function of Q (resp. R ) to take the total derivative:

Likewise for R ,

and so

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 :

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:

with kernels K ( k ) ζ j , for j = 1 , 2 , 3 given by

Or, dropping the linearization assumption on ( Q , R ) the updates are: K ( k ) ζ j , for j = 1 , 2 given by

The first projection is onto the set that satisfies the marginal constraint R 1 r = b . In particular, following the discussion above, one has the coordinate-MD step:

As before, there is no difference from the previous case, where one takes the unconstrained projection R = 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:

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:

To generalize this solution, we again consider adding a soft-regularization on the inner marginal of Q , where we consider the alternate problem:

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 rankr 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 1 n , g R = R T 1 m 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, 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:

This formulation amounts to solving a balanced Sinkhorn problem on T with respect to the proximal step kernel matrix.

Algorithm 2 SR R -projection (semi-relaxed OT, right marginal relaxed)

Algorithm 3 U -projection (unbalanced OT)

C The Latent Coupling Matrix

To solve the balanced form (and generalize the principle to the relaxed problems), we consider an alternative parametrization of the inner matrix. In particular, previous works Scetbon et al. (2021, 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 Q 1 r = a and R 1 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:

And

Thus, 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:

And:

We 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 X

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 :

And

This 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:

Noting the derivative D ( X -1 ) ◦ V = -X -1 V X -1 , we have from the first-order condition that:

This implies the kernel matrix update:

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 takes X = diag(1 / g Q ) T diag(1 / g R ) , then

and likewise

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:

Which yields the kernel matrix:

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

Initialize

2

g

Q

Q

0

0

,

,

a

g

T

R

,

,

if

r

=

r

X

0

2

0

else

X

then

0

end if while

∇

1

2

, τ, τ

, γ, δ, ε

,

b

R

1

r

2

2

r

r

r

=

1

,

1

Initialize-Couplings

T

(

1

-

0

∆((

Invertible case

#

diag(1

/

Q

Q

∇

R

γ

k

K

(

k

k

,

R

Q

k

CR

,

0

T

T

k

0

n

T

diag(1

1

)

k

X

C

Q

if

R

/

-

)

,

-

k

T

Q

X

T

γ/

)

,

1

(

Q

k

,

1

n

-

k

Q

{∥∇

k

R

k

diag

1

∥

m

a

b

g

,

,

R

-

-

1

0

Q

T

1

,

,

1

T

((

diag

∞

,

⊙

Q

m

k

g

)

-

1

))

CR

k

X

-

∥∇

1

R

γ

k

T

(diag(1

∥

)

/

∞

k

}

#

∇

-

(

Q

)

,

g

ℓ

R

∞

k

R

)

#

> ε

General case do

T

k

Q

R

diag(1

T

/

k

Q

g

X

))

k

R

C

Q

)

)

T

T

T

-normalization of Scetbon & Cuturi (2022)

⊙

,

, γ

k

, τ,

b

k

exp(

γ

k

R

)

exp(

K

k

(

k

)

R

k

else if

R

k

T

-

1

R

-

1

m

∇

, δ

)

#

Semi-relaxed OT

SR

-projection

K

else if

R

Semi-Relaxed Left

k

else if then

-projection

K

SR

R

(

Semi-Relaxed Right

)

(

k

R

then

k

k

R

end if

k

,

b

R

R

T

-

1

(

k

)

R

end while

Return P r = QXR T

Algorithm 5 Sinkhorn Algorithm (Cuturi (2013b), balanced OT)

D Gromov-Wasserstein (GW)

As defined in, the general Gromov-Wasserstein problem concerns a minimization of the energy:

T

-

1

k

m

U

-

projection

(

K

, γ

k

, γ

, τ,

, τ,

b

,

1

m

1

, δ

, δ

)

)

#

Semi-relaxed OT

Unbalanced OT

#

Where the minimization is over the set of all couplings Π a , b :

We consider extending the semi-relaxed framework to the GW problem under the low-rank restriction on P . This extension has thus far been considered for the balanced and unbalanced case in two previous works Scetbon et al. (2023, 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 ( ζ ) heretofore taken to be the Wasserstein loss ⟨ P ( ζ ) , C ⟩ F where the coupling P = QXR T is interpreted as a function of the low-rank sub-factor variables ζ ∈ { Q , R , X } . Taking L := Q A , B ( P ( ζ )) one can simply extend the gradient through the GW-loss and directly use it in place of the Wasserstein gradient in the update K ← ζ ⊙ exp( -γ k ∇ ζ L ( ζ )) for L ( ζ ) of each algorithm. The matrix-form of the GW-cost is expressed as:

Which, using the constraints of C 2 reduces the cost as a function of Q , R , X to:

Which is proportional to ∇ Q Q A , B ( Q , R , X ) ∝ -4 AQXR T BRX T for the balanced and rightmarginal semi-relaxed case. And:

Which likewise can be reduced in proportionality to ∇ R Q A , B ( Q , R , g ) ∝ -4 BRX T Q T AQX in 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 takes X ( 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:

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:

And 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:

So that the gradient with respect to the coupling matrix T is given as:

E Convergence Analysis and Other Proofs

E.1 Convergence and Smoothness of the Objective

We show in Proposition 3.3 that directly applying the block-descent lemma 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 n i . 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 × n i 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 ) :

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 n i , one has and 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 exists h i ∈ R n i such that v -u = U i h i . Then,

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 form

where X 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 .

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.

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

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

The prox-function (or Bregman divergence) associated with ω is then

and from this prox-function, γ > 0 , and some g ∈ R n , we define the generalized projection

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 :

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 via

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, where

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:

̸

Here, we have assumed that V can be written as a composite function of the block variables,

as is the case for the KL divergence. We also have assumed that h has this composite structure (as with entropy):

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 definition

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 , where

For Φ = f + h as above ( f has global smoothness constant L ), let Φ ∗ = arg min x ∈X Φ( x ) , and define

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 γ ( x -x + ) . Then for all x ∈ R n , all g ∈ R n , and γ > 0 , one has:

and where g X ,k is

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:

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:

Noting the definition g X ,k,i = 1 γ k,i ( x i -1 k -x i k ) , one has

Lemma 1 of Ghadimi et al. (2014) (stated as Lemma E.2 above) applies identically through blockwise optimality on X i because g X ,k,i = P X i ( x i -1 k , ∇ i f ( x i -1 k ) , γ k,i ) . Thus for any value ∇ i f ( x i -1 k ) takes,

and thus,

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 has

Noting the end-point condition Φ( x p k ) = Φ( x 0 k +1 ) , one may cancel all intermediate terms:

Thus one finds the upper bound in terms of Φ( x 0 ) and the minimum value Φ ∗ = min x ∈X Φ( x ) of Φ :

Taking γ k,i = α/L as in Ghadimi et al. (2014), the bracketed term directly above becomes α 2 / 2 L , and one has:

where D is as in (36), and where we used (37) to obtain the last line. Thus,

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:

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 n -g ( k -1) Q ∥ 2 < ϵ = 1 N ( 1 r -δ ) for g (0) Q = 1 r 1 r . Then, the FRLC objective

is 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, δ ) .

Proof. The addition of a pair of regularizations on the inner marginals τ KL( Q T 1 n ∥ Q T k 1 n ) and τ KL( R T 1 m ∥ R T k 1 m ) in 21 and 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 = ˜ O ( m 2 /ϵ ) iterations (hiding logarithmic and τ -factors) for each sub-problem solving for R k in Algorithm 4 and analogously N = ˜ O ( n 2 /ϵ ) for Q . Under the uniform initialization of g R , and after N iterations, we have:

With sufficiently large τ and N = ˜ O ( m 2 /ϵ ) sub-iterations, one can guarantee that:

This implies, for all iterations of the algorithm and all indices i , that

Thus, 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:

Where

Starting with the first term on the right side of (39), one has:

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,

To further bound the right-hand side above, consider:

While can easily be upper-bounded by an application of Jensen's inequality as

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

Next, we focus on the ∇ ( B ) Q term. Observe that

Thus:

Adding and subtracting terms in the norm and applying triangle inequality:

Invoking the lower-bound on the marginal, continuing from the above display,

Let us consider ∥∇ ( A ) Q k ∥ F ≤ ∥ X k ∥ F ∥ C ∥ F ∥ R k ∥ F ≤ ∥ X k ∥ F ∥ C ∥ F . We have that:

Where we always have that T ij ≤ g Q i and T ij ≤ g R j by definition of T as a coupling. As such:

Thus ∥∇ ( A ) Q k ∥ F ≤ √ mr δ ∥ C ∥ F , and the bound above reduces to

Further bounding the last term in the norm

Thus, the final bound on the ∇ ( B ) Q term is:

The total component-wise smoothness bound on Q is then

Identical reasoning applies for ∇ R L LC , where we similarly have the gradient split into two terms:

As before, we may first show smoothness in ∇ ( A ) R using the same steps as for Q

For ∇ ( B ) R , one may use the same reasoning as before to find:

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

holds trivially 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:

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:

Proof. 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 rankr 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:

We note that this bound also applies for the standard balanced optimal transport case, giving:

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:

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 ˜ q k ˜ r T k ≽ 0 :

Where we write this sum in terms of normalized vectors q k = ˜ q k / ∥ ˜ q k ∥ 1 , r k = ˜ r k / ∥ ˜ r k ∥ 1 , and λ k = ∥ ˜ r k ∥ 1 ∥ ˜ q k ∥ 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, as 1 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 is chosen 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 z -r +1 vectors of the decomposition to be:

And take the rankr approximation using the optimal r -1 vectors of OPT and this weighted average of the bottom to be:

Where, by the assumption that P ∗ ∈ Π b is feasible:

Thus ˜ P r ∈ Π b ,r is a feasible rankr solution by feasibility of P ∗ . One can verify that if P ∗ ∈ Π a , b , then ˜ P r 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 rankr solution:

Noting 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:

Concluding the proof. As discussed, this directly applies to the balanced case (for c b = 1 ).

F Initialization

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:

As C .,i , C .,j ≽ 0 , without loss of generality one may assume c 1 > 0 and c 2 < 0 . Then:

So clearly one has that C .,j -c C .,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 ( -c 1 c 2 ) 1 = C .,j -C .,i . Without loss of generality, assume 0 < -c 2 ≤ c 1 and C .,j > C .,i > 0 , so that δ = log ( -c 1 c 2 ) ≥ 0 . Then, considering constants q 1 , q 2 :

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 } i m =1 , the set { e C .,i } i m =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

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

One has the closed-form minimizer of g ∗ defined entrywise as:

For ω = 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:

Let us consider a rewriting of the inner product term:

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:

Thus, the first order condition tells us that the value of the coupling weight g j is related to λ as:

And by relying on the summation condition on the probability density g , yields the Langrange multiplier as

Plugging our Lagrange-multiplier into the above expression yields:

As the Hessian ∇ 2 g L = diag( ω g 3 ) ≽ 0 , we conclude that this value of g indeed minimizes the loss over ∆ r .

G Alternating updates on the dual variables

For the problem:

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 f 1 ) K Q diag( e γ k h 1 ) , and consider the entry-wise first-order condition required for the sub-coupling Q :

Thus:

And in matrix-form, this yields:

so that one finds

And expanding the Q 1 r term explicitly, noting that X diag( v ) 1 = Xv , we have:

Thus, we identify e γ k h 1 = e -γ k λ 1 as the right dual vector, and identify the following relationship in terms of the left dual vector:

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:

Implying an update for e γ k h 1 in the form:

Analogous reasoning applies for a relaxation of the other marginal, yielding the SR R -projection and SR L -projection (i.e. semi-relaxed OT).

H Discussion of Complexity

For ( Q , R , T ) ∈ R n × r 1 + × R m × r 2 + × R r 1 × r 2 + , the space complexity O ( nr 1 + r 1 r 2 + mr 2 ) is linear if the ranks r 1 , r 2 = o (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).

I Review of Background Material

I.1 Low-Rank Approximation of Pairwise Distance Matrices

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, 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 } j m =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 ˜ O ( 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:

For C r the optimal rankr 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 matrix C ∈ 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 C

Input point sets { x i } n i =1 , { y j } j M =1 in metric space X and metric d for

i

= 1

Pick indices i ∗ ∈ [ n ] , j ∗ ∈ [ m ] uniformly at random

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

to

n

do

Compute V using Chen & Price (2017)

return V , U

This 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:

With 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 note d ( x, y ) 2 ≤ ( d ( x, z ) + d ( z, y )) 2 = d ( x, z ) 2 + 2 d ( 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 a i n . Rewriting the norm of row i we have:

As we have the re-normalization p i ← p i ∑ i p i before sampling, we need to consider the value of the expectation E [ ∑ i p i ] to conclude.

By an application of Markov's inequality we have that:

Thus with probability 1 -δ we have:

This 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

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

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

J Connection of Optimal Transport to Projection Problems

Before discussing works which have address the problem of finding low-rank decompositions of the coupling matrix P , we discuss a few relevant preliminaries. The Bregman-divergence of some function F ( x ) , defined by the first-order Taylor expansion of F :

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)

Suppose C = ∩ 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:

Where the projection of ξ onto the set C is denoted by

Supposing that each C l forms a quotient space C l = V/U for a subspace U ⊂ V , defined as V/U = { v + U | v ∈ V } 2 , the iterative Bregman projection algorithm alternates projections onto each set C n as 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 space 3 , 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 as

One may initialize the residuals q -( L -1) = ... = q 0 = 0 and apply the algorithm:

Where C n := C n mod L as before and P C n 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 C n := 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:

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 .

J.1 Dykstra's algorithm with cyclic Bregman projections

As such, we can see the updates for F being the negative entropy and D F the KL-divergence. Without proof, the conjugate of the negative entropy is simply given as F ∗ ( ζ ) = exp { ζ -1 } . Thus, for the minimization problem:

Letting log q -( L -1) = ... = log q 0 = 0 , one may combine the two algorithms above to solve this minimization using generalized Dykstra's iterations. In particular, we have:

And:

So that:

Where division is interpreted to be element-wise, ⊙ refers to the Hadamard product, and the logarithm is applied elementwise.

J.2 Connection to Sinkhorn distances

Interestingly, the Sinkhorn algorithm described in Algorithm 5 can be alternatively derived in the context of Bregman iterations Benamou et al. (2015) as a minimization of the form:

Where ξ is the kernel ξ = e -C /ϵ and Π( µ, ν ) = C 1 ∩ C 2 for the convex constraint sets C 1 = { P : P 1 m = a } and C 2 = { P : P T 1 n = b } . To cast this into the Bregman-projection framework, one alternates between the two updates:

Where for the first projection one has the following first-order KKT condition:

With the constraint of C 1 that P 1 m = a , this implies diag( a / P ( l -1) 1 m ) = diag( e -λ /ϵ ) and recovers the first update of Sinkhorn P ( l ) ← diag( a / P ( l -1) 1 m ) P ( l -1) . An analogous argument gives the second, where all iterates satisfy P ( l ) = diag( u ( l ) ) e -C /ϵ diag( v ( l ) ) for u , v as defined in Algorithm 5.

0

20

40

60

80

100

Iterations

0

.

25

0

.

30

0

.

35

0

.

40

0

.

45

〈

C

,

P

〉

F

Rank-2

Full-rank

K Additional Simulations

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

L Graph Partitioning

L.1 Evaluation on a Graph Partitioning Task

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

60

80

100

120

140

160

180

200

Rank

0

.

165

0

.

170

0

.

175

0

.

180

0

.

185

0

.

190

0

.

195

〈

C

,

P

〉

F

FRLC (rank-2)

FRLC (full-rank random initialization)

LOT (rank-2)

LOT (default initialization)

LOT (k-means)

±

1

σ

(Full-rank, 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.

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.

L.2 Problem Statement

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 :

The 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 GWproblem 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 :

The learned h from semi-relaxed GW estimates the posterior proportion of nodes in G in each of the m clusters.

L.3 Dataset and Preprocessing

We run our semi-relaxed FRLC algorithm on four real-world graph datasets: 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.

L.4 Experiment Settings

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

20

40

60

80

100

Rank

0

.

44

0

.

46

0

.

48

0

.

50

0

.

52

0

.

54

0

.

56

〈

C

,

P

〉

F

FRLC (rank-2)

FRLC (full-rank random initialization)

LOT (rank-2)

LOT (default initialization)

η = 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.

M Spatial Transcriptomics Alignment

M.1 Problem Statement

In this problem, we use FRLC to find a low-rank alignment matrix between cells from two spatial transcriptomics (ST) Ståhl et al. (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 problem

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 problem

where 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 problem

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

M.2 Dataset and Preprocessing

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 .

Fig. 9 visualizes the two slices in this dataset, with each cell annotated with a cell type from the original publication.

Blood vessel

Kidney

Brain

Liver

Branchial arch

Lung primordium

Cavity

Meninges

Cartilage primordium

Mesentery

Connective tissue

Mesothelium

Choroid plexus

Mucosal epithelium

Dermomyotome

Muscle

Dorsal root ganglion

Ovary

Epidermis

Pancreas

GI tract

Sclerotome

Head mesenchyme

Spinal cord

Heart

Surface ectoderm

Inner ear

Sympathetic nerve

Jaw and tooth

Urogenital ridge

E11.5

E12.5

M.3 Experiment Settings

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.

M.4 Runtime

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.

N n th -roots of unity

N.1 Problem Statement

To test if FRLC can effectively coarse-grain transport between two datasets with obvious cluster structure, we generate a pair of two-dimensional datasets Z (1) and Z (2) , with ten and five clusters respectively. We run FRLC with T ∈ R 10 × 5 + to see if the barycentric projections for the LC factorization (Definition 4.1) can recover the cluster structure. These projections induce 10 and 5 barycenters for the first and second dataset, and are defined by

We examine, visually, whether these barycenters are good representatives of the clusters in each dataset, and whether the latent coupling depicts a reasonable transfer of mass between barycenters. We also run FRLC with T ∈ R 10 × 10 + and plot the barycenters from the resulting factorization for comparison. As discussed in § 4.2, the barycentric projections defined above, and in Definition 4.1 can be applied to factored couplings Forrow et al. (2019); Scetbon et al. (2021), yielding projections of the form:

Thus, we also ran the method of Scetbon et al. (2021) on this data, called LOT throughout the experiment, with g ∈ ( R ∗ + ) 5 and g ∈ ( R ∗ + ) 10 , plotting its barycenters in each case along with the diagonal latent coupling diag( g ) .

N.2 Dataset and Preprocessing

We instantiate a pair of two-dimensional datasets Z (1) and Z (2) as follows. Let U n denote the n th -roots of unity:

These complex numbers can be equivalently expressed as ordered pairs on the unit circle:

We consider n uniformly weighted mixtures of Gaussians, where for each sample X , we first sample a root of unity uniformly,

and conditionally on k , we sample X from an isotropic Gaussian centered at this root of unity

using a standard deviation σ = 0 . 1 . Samples are generated with the make_blobs function from sklearn.datasets . We generate two datasets in this way, Z (1) for n = 10 , and Z (2) for n = 5 . To generate dataset Z (1) , we first homogeneously scale the roots of unity to lie on a circle of radius 3 before sampling, so that the two datasets do not overlap but still use the same standard deviation. We do not scale the centers used for Z (2) , so they all lie on the unit circle. We generate Z (1) = { z (1) 1 , . . . , z (1) 1000 } and Z (2) = { z (2) 1 , . . . , z (2) 1000 } using 1000 samples each and form the empirical measures

supported on Z (1) and Z (2) , corresponding to uniform probability vectors a , b ∈ ∆ 1000 .

N.3 Experiment Settings

We used default hyperparameter settings for both methods. In particular, FRLC sets τ = 75 and γ = 90 , with a maximum of 200 iterations and a minimum of 25 iterations, subject to the stopping criterion ∆ given in (10). The LOT default settings are γ = 10 (there is no hyperparameter analogous to τ in LOT). Both methods use a random initialization and were run on CPU.

O Discussion of differences between FRLC and existing low-rank optimal transport algorithms

We 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, 2023, 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:

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

P Limitations

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.

Q Broader impacts

FRLC is general enough to be used modularly within any ML algorithm using OT as a subroutine to help with scalability. We also note that the LC factorization is similar to a PCA in the context of OT, yielding an optimal low-rank coupling with an interpretable latent coupling factor.

Blood vessel

Kidney

Brain

Liver

Branchial arch

Lung primordium

Cavity

Meninges

Cartilage primordium

Mesentery

Connective tissue

Mesothelium

Choroid plexus

Mucosal epithelium

Dermomyotome

Muscle

Dorsal root ganglion

Ovary

Epidermis

Pancreas

GI tract

Sclerotome

Head mesenchyme

Spinal cord

Heart

Surface ectoderm

Inner ear

Sympathetic nerve

Jaw and tooth

Urogenital ridge

Ground truth

Predicted

(c)

(a)

(b)

(d)

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The abstract and introduction exactly describe the background preceding this paper and placing it into context, and exactly describe the contribution and scope of the paper to the field.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We describe the limitations of our method in Section P. These include an additional hyperparameter that FRLC introduces relative to previous work, and the substantial room for code optimization.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: All theoretical results state their assumptions clearly and all propositions are followed by thorough line-by-line proofs with careful justifications made for each step. They have checked over by all of the authors. All proofs are provided in the supplement with detail and are referenced in the main body where appropriate.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide all code used for generating the results of our paper. This not only includes the source code for the optimization, but the experimental code used for benchmarking. The algorithm is implemented exactly as described in the paper, and reviewers can freely consult the code we provide to them.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We will upload all code, and all code required to generate the synthetic data experiments used. Any real data experiments using especially large-scale data have publicly available and easily accessible datasets which we provide references for. We provide comprehensive descriptions of how the data was pre-processed, and provide experimental code which can be followed easily.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We introduce all hyper-parameters with rigorous justification of their utility. The values of the default hyper-parameters are accessible in the code we upload and any experiment which does not use the default hyperparameter (e.g. for a validation search) has the full table provided and the code used for the experiment.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

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.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We indicate when experiments are run on CPU versus GPU in the experimental details.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: This research was conducted ethnically and conforms to the guidelines of the code of ethics. We do not anticipate any major negative societal impact 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).

11. 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?

Answer: [Yes]

Justification: All data that we use and all work that we build on are cited and credited heavily.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: We provide no new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: This paper involves no research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nAs the paper clearly outlines in the introduction, Optimal Transportation (OT) is widely used in various fields of machine learning, however, the cost for computing OT is computationaly expensive (quadratic scaling), even after Sinkhorn algorithm employing entropy regularization significantly alleviated the cost. The purpose of Low-rank Optimal Transport (LOT) is to reduce the complexity by constraining the transportation plan to possess a certain low-rank structure.\n\nThe key contribution of the paper is (1) new parametrization of LOT using a Latent Coupling (LC) factorization, (2) The eponymous \"Balance FRLC\" algorithm for resolving the optimization procedure for computing LOT with respect to various objectives associated with (LOT), (3) Theoretical guarantees for the smoothness of the regularized objectives, as well as upper bounds for the convergence criteria.  The algorithm is applied to illustrative toy examples and then to a more realistic example (spatial transcriptomics alignment) and the results indicate improvements over previous LOT.\n\nSTRENGTHS:\nThe contributions of the paper is significant. The writing is excellent. The LC factorization, adapted to the LOT framework from a previous work, is well-motivated. The Balanced FRLC algorithm appears a sound and non-trivial extension of previous approaches for LOT (Sinkhorn, semi-relaxed OT). The algorithm is particularly simple and shown to be effective. The experimental results are clear and illustrative, improvements are convincing.\n\nWEAKNESSES:\nWhat this reviewer found missing is discussion of the case the optimal transport map is itself not low-rank. It is well-known that the optimal transportation plan is not necessarily low-rank (as indicated e.g. by their connection to differential equations with non-regular solutions), so LOT may converge very slowly to the full-rank optimal transportation plan. In section 4.2, the authors provide a simple illustrative example, but this is when the transportation plan is known to be low-rank - in this case the transportation plan seems to stabilize after some threshold rank - however, what happens when the full-rank plan does not have an inherently low-rank structure?\n\nQUESTIONS:\n- line 103: Would it be good to say the what complexity \"massively improved\" to?\n- Can the authors comment on how to pick the rank for the LOT? (This of course ties in to the point above the rank of the transportation plan.)\n- Perhaps a more discussion of of how different the LC factorization is to, say the previous factored coupling with (the diagonal version $P = Q \\text{diag}(1/g) R^T$). This relation seems important when considering \"more\" LC factorizations: If this LC factorization improves the result while simultaneously allowing a simple optimization algorithm, would it then be natural to discuss further factorizations that are nested this way, with additional LC factorizations for $T$ and so on?\n\nLIMITATIONS:\nThis 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?", "TITLE": "Official Review", "RECOMMENDATION": 9, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nI found that the presented FRLC method is a quite elegant approach to low rank OT. It decouples optimization into entropy regularized semi relaxed OT problems for $Q$ and $R$, and $T$ is updated solving entropy regularized  OT of size $r$. This decoupling is enabled by the form of the inner coupling which is no longer assumed to be diagonal. This construction alleviates the need to use the Dykstra machinery as in (Scetbon 2021).\nWith the proposed approach, $Q$ and $R$ are not required to have the same inner marginal. From what I understand, the optimization over inner marginals is handled by the semi relaxed steps and does not require a separate step as in previous approaches. \n\nAnother strength of the approach is that it generalizes smoothly to the unbalanced setting. \n\nOther strengths include:\n* I carefully looked at the presented method and the appendix and the algorithmic approach made sense. \n* non asymptoptic convergence bounds are derived for the method. \n* the appendix includes a complete review of the necessary technical background presented in a clear way, easy to follow.\n* limitations are clearly stated (section N in appendix).\n* a new initialization scheme for LC decomposition is proposed. \n* there is a gain in interpretability compared to (Scetbon 2021) as the inner coupling captures the coupling between clusters in the data.\n\nWEAKNESSES:\n* from section 3.1 to 3.2, the new aspects of the approach compared to https://arxiv.org/pdf/2012.11589 should be stated more clearly. The contributions of the paper should be separated more clearly from previous works. I had to read the other paper to better understand each contribution. \n* In algorithm 1, it would be interesting to have the computational complexity of each step (presented in the algorithm directly).\n\nQUESTIONS:\n* How to choose $r$ in practical scenarios ? \n* Would it make sense to remove the entropic penalties (at the expense of smoothness) to recover results about hard clustering of the points via the decomposition of the OT coupling ? Put differently, what is the effect of entropic smoothing on the interpretability of the coupling in terms of clustering ?\n\nLIMITATIONS:\n/", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-15"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe introduction of the LC factorization and its integration into the FRLC algorithm represents an innovation in reducing the complexity of optimal transport problems.\n\nThe FRLC framework is versatile, handling multiple OT objectives and marginal constraints, making it applicable to a wide range of practical problems.\n\nThe latent coupling provides a interpretable description of the transport plan, which is beneficial for understanding and visualizing the results.\n\nWEAKNESSES:\nWhile the empirical results are good, the paper could benefit from more extensive comparisons with additional baseline methods to solidify the claims of superiority.\n\nThe theoretical results are limited in that, even though $\\Delta_k(x_k, x_{k+1})$ is small, it cannot be guaranteed that the iteration points converge to the optimal solution or even a stationary point.\n\nQUESTIONS:\nI suggest that Propositions E.3, E.5, E.6 should be stated in the main text to make the theoretical part easier to follow. \n\nTable 1 should be put in the main text to show the efficiency of FRLC in terms of runtime. It is hard to say that the method shows significant improvement because the runtime of LOT is already very short (less than one second). I suggest the authors conduct experiments on more challenging and diverse examples.\n\nLIMITATIONS:\nN.A.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-16"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- Generality: handles different OT objective costs and relaxations of the marginal constraints.\n- Significance: the proposed method, and improving OT using factorizations, has the potential to address the scalability issue of OT with large datasets, making it highly relevant for applications in machine learning and data science.\n\nWEAKNESSES:\n**Presentation and logical organization of the ideas**\n\nI found it difficult to discern the paper's particular contribution, especially concerning previous art like Latent Optimal Transport (Lin et al., 2021), Forrow et al. (2019), and Scetbon et al. (2021). For example, the main contribution is said to be \"compute a minimum cost low-rank transport plan using the LC factorization,\" where LC factorization is precisely the factorizations proposed by Lin et al. (2021), with the further contribution being the computation of this factorization using mirror-descent. I wondered how much of the description also applies to Lin et al. (2021) and which part is particular to the novel method. \n\nFurthermore, according to Lin et al. (2021), there is an algorithm for computing low-rank factorized plans based on factorized costs using a projection method and an algorithm to factorize the costs and select anchor points without the need to have a similar number of anchor points at source and target latent space. How does your approach compare with the approach? \n\n**Comparison and empirical validation**\nThe first empirical comparison analyzes LOT and FRLC, capturing how both methods behave as the factorization rank changes. This is a good way of understanding the effect of the rank on each case. Nevertheless, in the second setup, the comparison with LOT is not present anymore, which begs the question of how LOT performs in the second empirical evaluation. Furthermore, an extensive empirical study with the structured approach of Lin et al. (2021) would help clarify the particular contributions of the type of factorization and factoring algorithm used. \n\nI found some other issues lacking in a more in-depth discussion of the computational complexity, possible implementation difficulties, and overall limitations in the proposal.\n\nQUESTIONS:\n- Is the main contribution here actually the mirror descent formulation of the OT problem with the factorization with a different set of latent source and target points proposed by Lin et al. (2021)?\n- Furthermore, is the main difference between the fact that Lin et al. (2021) use a two-stage process and the novel mirror-descent method solves the OT problem in a more straightforward unified process? \n- What are the theoretical complexity of the algorithms, and what practical considerations are relevant when implementing the method?\n\nOne presentation improvement that would change my evaluation is a summary table locating the differences in the type of factorization, whether the cost or transport is being factorized, the algorithm used for computing the factorization, etc.\n\nLIMITATIONS:\nI found it necessary to discuss in more depth the limitations and the tradeoffs present in the proposal compared with previous art.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-26"}