Low-Rank Optimal Transport through Factor Relaxation with Latent Coupling

TL;DR

A new algorithm, FRLC, which handles multiple OT objectives, and marginal constraints (balanced, unbalanced, and semi-relaxed) with linear space complexity, and provides theoretical results on FRLC, and demonstrates superior performance on diverse applications -- including graph clustering and spatial transcriptomics -- while demonstrating its interpretability.

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 $\textit{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 $\textit{Factor Relaxation with Latent Coupling}$ (FRLC), which uses $\textit{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.

Figures (12)

• Figure 1: (Left) The LC factorization $\bm{P} = \bm{Q} \mathrm{diag}(1/\bm{g}_Q)\bm{T} \mathrm{diag}(1/\bm{g}_R) \bm{R}^\mathrm{T}$ of coupling matrix $\bm{P}$ with outer marginals $\bm{a},\bm{b}$, inner marginals $\bm{g}_Q,\bm{g}_R$, factors $\bm{Q}, \bm{R}$, and latent coupling $\bm{T}$. (Right) Full-rank coupling matrix $\bm{P}$.
• Figure 2: (a) Simulated dataset containing points from two moons (orange) and eight Gaussians (blue). (b) Transport cost $\langle \bm{C}, \bm{P} \rangle_F$ achieved by FRLC and LOT Scetbon2021LowRankSF for the balanced Wasserstein problem on the dataset in (a) for different ranks and initializations. FLRC full rank (blue curve) is average over 10 random initializations. (c) Results on the 10D mixture of Gaussians dataset.
• Figure 3: LC-projections of couplings of Gaussians centered on the 5th-roots of unity (green) and 10th roots of unity (yellow). (a) Ground-truth full-rank coupling. (b) Non-square rank-5 latent-coupling of FRLC (c) LC-projection barycenters aligned with rank-5 diagonal coupling of LOT Scetbon2021LowRankSF. (d) Square rank-10 latent coupling of FRLC. (e) Rank-10 diagonal coupling of LOT .
• Figure 4: (a) Brain marker gene Tubb2b expression and FRLC prediction. (b) Comparison of the low-rank unbalanced (LOT-U) algorithm of scetbon2023unbalanced and FRLC on aligning spatial transcriptomics data. Bold indicates top performing method for each metric on each objective.
• Figure 5: Transport cost $\langle \bm{C}, \bm{P} \rangle_F$ against number of iterations for FRLC with rank 200 on the synthetic dataset of two moons and eight Gaussians. Smooth convergence is observed for both rank-2 and full-rank random initialization.

Theorems & Definitions (29)

• Definition 3.1: Inner marginals
• Definition 3.2: LC factorization
• Proposition 3.3
• Proposition 3.4
• Definition 4.1: LC-Projection
• Lemma A.1
• proof
• Lemma E.1: Block descent lemma, Beck2013OnTC, Lemma 3.2.
• Lemma E.2: Ghadimi2014, Lemma 1
• Proposition E.3: Proposition \ref{['prop:convergence']}