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 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+γ+2γ) approximation for kernel costs, where γ∈[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=σ(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.
Pσ⋆←n⋅argminP∈Π(un,un)⟨C,P⟩F
Algorithm 1 Transport Clustering ( TC )
Figure 2: The relative cost of the rank K∈{50,75,…,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 ↓
AMI (A/B) ↑
ARI (A/B) ↑
CTA ↑
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 → 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
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.
Table 3: Low-rank OT costs with fixed ∣X∣=1024 and varying dataset size ∣Y∣ for task of ( Tong et al., 2023 ) (Kantorovich registration).
ε
LR-OT Cost
10−5
5.050
10−4
5.620
10−3
8.167
10−2
8.919
10−1
9.415
100
9.576
Appendix
Table 4: Low-rank OT cost of Transport Clustering as a function of the Sinkhorn regularization ε 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 , σ=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 ↓
AMI (A/B) ↑
ARI (A/B) ↑
CTA ↑
Runtime (s)
E8.5 → 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 → 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).
Figure 5: Comparison of low-rank OT methods on the stochastic block model dataset. (Left) Relative cost of the rank K∈{10,…,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 .
Figure 6: Comparison of low-rank OT methods on the shifted Gaussians dataset. (Top) Relative cost of the rank K∈{50,75,…,250} transport plan inferred by LOT , FRLC , and LatentOT compared to the cost of the transport plan inferred by TC across noise levels σ2∈{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 σ2∈{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∈{50,75,…,250} transport plan inferred by LOT , FRLC , and LatentOT compared to the cost of the transport plan inferred by TC across noise levels σ2∈{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∈{50,75,…,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 C (top), and the transport conjugated cost M† (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 ϵ . Lower ϵ is closer to an optimal full-rank solution.
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.
Pratik Jawanpuria, Bamdev Mishra
Centre for Machine Intelligence and Data Science, IIT Bombay · Microsoft India
We consider the fundamental problem of balanced k-means clustering. In particular, we introduce an optimal transport approach to alternating minimization called BalLOT, and we show that it delivers a fast and effective solution to this problem. We establish this with several theoretical guarantees and a variety of numerical experiments. On the theory front, we first prove that for generic data, BalLOT produces integral couplings at each step. Next, we perform a landscape analysis to provide theoretical guarantees for both exact and partial recoveries of planted clusters under the stochastic ball model. We also propose initialization schemes that achieve one-step recovery of planted clusters. To conclude, we present numerical experiments that corroborate our theoretical results.
Wenyan Luo, Dustin G. Mixon
Department of Mathematics, The Ohio State University, Columbus, Ohio, USA · Translational Data Analytics Institute, The Ohio State University, Columbus, Ohio, USA
Optimal transport (OT) provides a principled framework for mapping between probability distributions. Despite extensive progress, applying OT to large-scale data remains computationally demanding, and the resulting pointwise transport plans are often difficult to interpret. We introduce Optimal Mixture Transport (OMT), a scalable framework that shifts the transport paradigm from individual samples to mixtures of subpopulations, reformulating the transport problem as a strictly biconvex optimization with a unique global minimizer. We further establish theoretical guarantees on the stability of the OMT map, showing that bounded perturbations of the underlying distributions lead to bounded changes in the transport plan. By formulating subpopulations as exponential-family distributions, OMT decouples computational complexity from the sample size, scaling solely with the number of mixture components. We demonstrate the effectiveness and practicality of OMT on a wide range of synthetic benchmarks and real-world datasets, including image data and large-scale single-cell RNA sequencing measurements.