Transport Clustering: Solving Low-Rank Optimal Transport via Clustering

Henri Schmidt; Peter Halmos; Ben Raphael

cs.LGMar 3, 2026

Transport Clustering: Solving Low-Rank Optimal Transport via Clustering

Authors: Henri Schmidt, Peter Halmos, Ben Raphael

Organizations: Department of Computer Science, Princeton University, NJ, USA.

Abstract

Optimal transport (OT) finds a least cost transport plan between two probability distributions using a cost matrix defined on pairs of points. Unlike standard OT, which infers unstructured pointwise mappings, low-rank optimal transport explicitly constrains the rank of the transport plan to infer latent structure. This improves statistical stability and robustness, yields sharper parametric rates for estimating Wasserstein distances adaptive to the intrinsic rank, and generalizes $K$ -means to co-clustering. These advantages, however, come at the cost of a non-convex and NP-hard optimization problem. We introduce transport clustering, an algorithm to compute a low-rank OT plan that reduces low-rank OT to a clustering problem on correspondences obtained from a full-rank $\textit{transport registration}$ step. We prove that this reduction yields polynomial-time, constant-factor approximation algorithms for low-rank OT: specifically, a $(1+γ)$ approximation for negative-type metrics and a $(1+γ+\sqrt{2γ}\,)$ approximation for kernel costs, where $γ\in [0,1]$ denotes the approximation ratio of the optimal full-rank solution relative to the low-rank optimal. Empirically, transport clustering outperforms existing low-rank OT solvers on synthetic benchmarks and large-scale, high-dimensional datasets.

Figures & tables

Figure 1: TC on (a) a synthetic 2-Moons ( $X$ ) and 8-Gaussians ( $Y$ ) dataset ( $n=m=1024$ ) from ( Tong et al., 2023 ) with the (b) Monge map alignment of $X$ and $Y=\sigma(X)$ using ( Halmos et al., 2025b ) . TC reduces low-rank OT (co-clustering) to (c) clustering a single set of Monge registered correspondences using generalized $K$ -means.

\mathbf{P}_{\sigma^{\star}}\leftarrow n\cdot\argmin_{\mathbf{P}\in\Pi(\bm{u}_{n},\bm{u}_{n})}\langle\mathbf{C},\mathbf{P}\rangle_{F}

Algorithm 1 Transport Clustering ( TC )

Figure 2: The relative cost of the rank

K\in\{50,75,\ldots,250\}

transport plan inferred LOT , FRLC , and LatentOT compared to the cost of the transport plan inferred by TC across

315

synthetic instances (lower is better). Each dataset contains

n=m=5000

data points. LatentOT is excluded from the stochastic block model evaluation as it takes as input a squared Euclidean cost matrix.

Dataset	Method	Rank	OT Cost $\downarrow$	AMI (A/B) $\uparrow$	ARI (A/B) $\uparrow$	CTA $\uparrow$
CIFAR-10 ( $60{,}000$ )	TC	10	231.200	0.478 / 0.476	0.358 / 0.356	0.412
	FRLC	10	235.950	0.411 / 0.407	0.281 / 0.277	0.351
	LOT	10	234.733	0.430 / 0.427	0.306 / 0.303	0.358
Mouse embryo E8.5 $\to$ E8.75 ( $18{,}819$ )	TC	43	0.506	0.639 / 0.617	0.329 / 0.307	0.722
	FRLC	43	0.553	0.556 / 0.531	0.217 / 0.199	0.525
	LOT	43	0.520	0.605 / 0.592	0.283 / 0.272	0.611

Table 1: Comparison of low-rank OT methods across three datasets: CIFAR-10 (

n=60{,}000

), smallest mouse embryo split (

n=18{,}819

), and largest mouse embryo split (

n=131{,}040

Figure 3: Estimation of squared Wasserstein distance on the fractured hypercube of ( Forrow et al., 2019 ) . Convergence shown for fixed

d=30

K=10

, and averaged over

10

runs.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 4 — Figure 4: Geometric constructions providing lower bounds for Theorem 4.1 in the case of (left) Euclidean cost ( $k$ = 3) and (right) squared Euclidean cost ( $k$ = 2). Points in ${X}$ are colored black and points in ${Y}$ are colored white. Points connected by a line segment have identical coordinates and are separated for ease of visualization.

N	29	36	44	54	66	80	98	119
Full-Rank OT	14.520	14.121	13.887	13.485	13.213	12.863	12.523	12.249
K-Means	9.616	8.540	7.887	7.087	6.729	6.196	5.828	5.579
FactoredOT	5.560	4.493	3.742	2.774	2.173	1.899	1.291	0.827
FRLC	2.449	2.448	2.188	1.745	0.920	0.995	0.598	0.526
TC (Ours)	3.009	2.344	1.385	0.784	0.502	0.357	0.342	0.242

Appendix

Table 2: Estimation error

|\hat{W}_{2}^{2}-W_{2}^{2}|

for Varying Sample Size (d=30, k=10).

$\|Y\|$	TC	FRLC	LOT
64	8.983	10.508	10.108
128	9.407	11.211	10.607
256	8.831	10.289	10.153
512	8.931	10.595	10.168
1024	8.794	10.235	10.460

Appendix

Table 3: Low-rank OT costs with fixed

|X|=1024

and varying dataset size

|Y|

for task of ( Tong et al., 2023 ) (Kantorovich registration).

$\varepsilon$	LR-OT Cost
$10^{-5}$	5.050
$10^{-4}$	5.620
$10^{-3}$	8.167
$10^{-2}$	8.919
$10^{-1}$	9.415
$10^{0}$	9.576

Appendix

Table 4: Low-rank OT cost of Transport Clustering as a function of the Sinkhorn regularization

\varepsilon

used in the registration step.

$r$	TC (ours)	FRLC (rand)	FRLC (TC-init)
50	8.0022	8.4583	8.0088
100	7.7327	8.4360	7.7713
150	7.5021	8.4550	7.5679
200	7.2859	8.4490	7.3745
250	7.0762	8.4448	7.2210

Appendix

Table 5: Low-rank OT costs on Planted Gaussians (

k=250

\sigma=0.1

n=2500

) for varying rank

r

. TC denotes Transport Clustering; FRLC (rand) uses a standard random initialization for low-rank OT, and FRLC (TC-init) uses TC-derived initialization.

Timepoints	Method	Rank	OT Cost $\downarrow$	AMI (A/B) $\uparrow$	ARI (A/B) $\uparrow$	CTA $\uparrow$	Runtime (s)
E8.5 $\to$ E8.75 (18,819 cells)	TC	43	0.506	0.639 / 0.617	0.329 / 0.307	0.722	63.38
	FRLC	43	0.553	0.556 / 0.531	0.217 / 0.199	0.525	16.45
	LOT	43	0.520	0.605 / 0.592	0.283 / 0.272	0.611	8.77
E8.75 $\to$ E9.0 (30,240 cells)	TC	53	0.384	0.597 / 0.598	0.231 / 0.230	0.545	177.12
	FRLC	53	0.405	0.534 / 0.541	0.174 / 0.178	0.492	16.92
	LOT	53	0.390	0.559 / 0.567	0.193 / 0.197	0.487	10.88

Appendix

Table 6: Single-cell transcriptomics alignment on consecutive mouse embryo timepoints. We report OT cost (lower is better), AMI/ARI for each split (A/B), and class-transfer accuracy (CTA; higher is better).

Panel of Figure 5 — Figure 5: Comparison of low-rank OT methods on the stochastic block model dataset. (Left) Relative cost of the rank $K\in\{10,\ldots,100\}$ transport plan inferred by LOT and FRLC compared to the cost of the transport plan inferred by TC . (Right) Co-clustering accuracy (AMI/ARI) of TC , LOT , and FRLC at rank $K=100$ . The stochastic block model dataset consists of $100$ clusters of size $50$ .

Panel of Figure 6 — Figure 6: Comparison of low-rank OT methods on the shifted Gaussians dataset. (Top) Relative cost of the rank $K\in\{50,75,\ldots,250\}$ transport plan inferred by LOT , FRLC , and LatentOT compared to the cost of the transport plan inferred by TC across noise levels $\sigma^{2}\in\{0.1,0.2,0.3\}$ . (Bottom) Co-clustering accuracy (AMI/ARI) of TC , LOT , FRLC , and LatentOT at rank $K=250$ across noise levels $\sigma^{2}\in\{0.1,0.2,0.3\}$ . The shifted Gaussians dataset consists of $250$ clusters of unequal size.

Figure 7: Relative cost of the rank $K\in\{50,75,\ldots,250\}$ transport plan inferred by LOT , FRLC , and LatentOT compared to the cost of the transport plan inferred by TC across noise levels $\sigma^{2}\in\{0.1,0.2,0.3\}$ for the 2-Moons and 8-Gaussians ( Tong et al., 2023 ) dataset.

Figure 8: Runtime of TC , LOT , FRLC , and LatentOT versus the rank

K\in\{50,75,\ldots,250\}

for the 2-Moons and 8-Gaussians ( Tong et al., 2023 ) dataset and the Shifted Gaussians dataset across all noise levels.

Figure 9: (Left) An example co-clustering of the two-moons 8-Gaussians dataset ( Tong et al., 2023 ) with Algorithm 3 . (Right) A comparision between the raw cost matrix $\mathbf{C}$ (top), and the transport conjugated cost $\mathbf{M}^{\dagger}$ (bottom).

Figure 10: (Sensitivity of TC to error in the coupling) Low-Rank OT Cost of Transport Clustering as a function of entropy-regularization scale

\epsilon

. Lower

\epsilon

is closer to an optimal full-rank solution.

Optimal Transport

Optimal transport compares and transforms probability distributions by finding low-cost ways to move mass between them.

See all 289 papers

A Riemannian Approach to Low-Rank Optimal Transport
Jun 10, 2026Pratik Jawanpuria, Bamdev MishraRiemannian Optimization Second-Order Optimization
Centre for Machine Intelligence and Data Science, IIT Bombay · Microsoft India
BalLOT: Balanced $k$ -means clustering with optimal transport
Dec 5, 2025Wenyan Luo, Dustin G. MixonClustering Nonconvex Optimization
Department of Mathematics, The Ohio State University, Columbus, Ohio, USA · Translational Data Analytics Institute, The Ohio State University, Columbus, Ohio, USA
A Biconvex Formulation for Stable Transport of Mixture Models with a Unique Solution
Jun 1, 2026Yeganeh Marghi, Kelly Jin, Uygar SümbülOptimal Transport
Allen Institute Seattle, WA, USA

Jun 10, 2026J/K move · Enter open · S save

A Riemannian Approach to Low-Rank Optimal Transport

Pratik Jawanpuria, Bamdev Mishra

Centre for Machine Intelligence and Data Science, IIT Bombay · Microsoft India

Riemannian Optimization Second-Order Optimization Optimal Transport

Open paper AskPDF arXiv

Low-rank optimal transport (OT) mitigates the quadratic scaling of classical solvers, yet existing approaches rely heavily on first-order mirror-descent updates that require careful hyperparameter tuning and ignore the optimization landscape's curvature. To address these limitations, we propose a unified Riemannian geometric framework for low-rank OT, modeling balanced and unbalanced rank-

r

positive factored couplings as novel smooth embedded submanifolds of the positive orthant. By equipping these manifolds with the Fisher-Rao product metric, we derive tractable formulations for Riemannian projectors, retractions, and Hessian-vector products. Our cost-agnostic framework seamlessly extends to linear OT, Gromov-Wasserstein (GW), fused GW, and their unbalanced counterparts. For balanced OT, our geometric ingredients are computed via efficient conjugate-gradient and iterative Bregman updates. For the unbalanced OT, our operations elegantly reduce to closed-form scalings, completely eliminating inner iterative loops. In both regimes, per-iteration complexity scales linearly with dataset size, and we provide a rank-sufficiency certificate for global optimality verification. Extensive experiments across a range of problem sizes demonstrate that our regularization-free first- and second-order solvers achieve faster convergence and superior performance over existing state-of-the-art low-rank OT solvers.

Abstract

Optimal Transport

Explore similar work

A Riemannian Approach to Low-Rank Optimal Transport

BalLOT: Balanced kkk-means clustering with optimal transport

A Biconvex Formulation for Stable Transport of Mixture Models with a Unique Solution

BalLOT: Balanced $k$ -means clustering with optimal transport