Transport Clustering: Solving Low-Rank Optimal Transport via Clustering

This paper introduces "transport clustering," a polynomial-time algorithm that reduces the NP-hard low-rank optimal transport problem to a clustering task via transport registration, achieving constant-factor approximation guarantees and outperforming existing solvers on both synthetic and large-scale datasets.

The Gist

Imagine you have two massive, messy piles of laundry. One pile is from your bedroom (Dataset A), and the other is from your living room (Dataset B). Your goal is to match every sock in the bedroom pile to its perfect partner in the living room pile, but you don't know which sock goes with which.

Optimal Transport (OT) is the mathematical way of solving this. It asks: "What is the cheapest, most efficient way to move every sock from pile A to pile B?" Usually, this results in a giant, chaotic map where every single sock is matched to a specific sock in the other pile. It's accurate, but it's computationally exhausting and creates a "flat" map where every match looks equally important.

Low-Rank Optimal Transport tries to fix this by saying: "Wait, these socks probably belong to a few specific types of pairs." Maybe there are just 5 types of socks (e.g., black dress socks, white athletic socks, colorful argyles). Instead of matching sock-to-sock, we want to match sock-to-type. This reveals the hidden structure (the "latent factors") and makes the problem much more robust to noise (like a sock that got lost or is slightly dirty).

The Problem: Finding these hidden types is incredibly hard. It's like trying to solve a 3D puzzle where the pieces keep changing shape. The math is "non-convex," meaning if you start looking in the wrong spot, you get stuck in a local trap and never find the best solution. Existing methods are slow, sensitive to how you start, and often give different answers depending on the day.

The Solution: "Transport Clustering"

The authors of this paper introduce a clever new method called Transport Clustering. They realized that instead of trying to solve the hard 3D puzzle directly, you can break it down into two simpler steps.

Here is the analogy:

Step 1: The "Monge Registration" (The Perfect Matchmaker)

First, imagine a super-smart matchmaker who ignores the "types" for a moment and just finds the single best, one-to-one match for every sock, regardless of what it is.

• In math terms, this is solving the Full-Rank Optimal Transport problem. It's a standard, well-understood problem that computers can solve quickly and reliably.
• Let's say this matchmaker creates a "translation guide." It tells us: "Sock #1 in the bedroom corresponds to Sock #42 in the living room."

Step 2: The "Clustering" (Finding the Groups)

Now, here is the magic trick. Instead of trying to find the groups while matching, we use the matchmaker's guide to rearrange the living room pile.

• We take the living room socks and shuffle them around so that Sock #42 is now sitting right next to Sock #1.
• Once they are aligned, the problem changes. We no longer need to worry about matching two different piles. We just need to look at this single, rearranged pile and ask: "Which socks naturally belong together in a group?"
• This is now just a standard K-Means Clustering problem (the same algorithm used to group customers by shopping habits or photos by face). It's easy, fast, and has guaranteed good results.

Why is this a big deal?

1. It's a Shortcut: The authors proved that by doing Step 1 (the easy match) and then Step 2 (the easy grouping), you get a solution that is mathematically guaranteed to be very close to the perfect answer. You don't need to guess; the math says, "This will work within a small, predictable error margin."
2. It's Stable: Old methods were like trying to balance a house of cards; a tiny change in the starting point would make the whole thing collapse into a different, worse solution. Transport Clustering is like building with LEGOs; it's stable and reliable.
3. It's Fast: Because they reduced the hard problem to a standard clustering problem, they can use existing, super-fast computer algorithms to solve it.

Real-World Impact

The paper tested this on some huge, messy datasets:

• Images: They matched 60,000 images from the CIFAR-10 dataset. Transport Clustering found better groupings of similar images than previous methods.
• Biology: They analyzed millions of cells from mouse embryos to see how they change over time. This is like trying to trace the family tree of a cell. Transport Clustering was able to link cells across different time points more accurately than before, helping scientists understand how life develops.

The Takeaway

Think of Transport Clustering as a "divide and conquer" strategy for data matching.

• Old Way: Try to solve the matching and grouping simultaneously. It's a nightmare of complexity.
• New Way: First, force a perfect one-to-one match (Registration). Then, simply group the aligned data (Clustering).

By turning a difficult, NP-hard optimization problem into a simple clustering problem, the authors have given scientists and data analysts a powerful, reliable, and fast tool to find hidden structures in complex data. It's like realizing that to organize a messy library, you don't need to sort every book by hand; you just need to put them in the right order first, and then the groups will naturally fall into place.

Technical Summary

1. Problem Definition

Optimal Transport (OT) seeks a least-cost mapping between two probability distributions. While standard OT infers unstructured pointwise mappings (often resulting in full-rank permutation matrices), Low-Rank Optimal Transport (LR-OT) explicitly constrains the rank of the transport plan to a value $K \ll n$.

• Motivation: LR-OT reveals latent structure, improves statistical stability, yields sharper parametric rates for Wasserstein distance estimation, and generalizes K-means to co-clustering multiple datasets.
• Challenge: LR-OT is a non-convex, NP-hard optimization problem. Existing solvers rely on local optimization (e.g., mirror descent or Lloyd-type algorithms) over three or more variables, are sensitive to initialization, and lack provable approximation guarantees beyond convergence to stationary points.

2. Methodology: Transport Clustering (TC)

The authors propose Transport Clustering, an algorithm that reduces the LR-OT problem to a standard clustering problem via a "transport registration" step.

Core Algorithm (Algorithm 1)

The method proceeds in two distinct steps:

1. Transport Registration (Step 1): Compute the optimal full-rank transport plan (Monge map) $P_{\sigma^*}$ between the source ($X$) and target ($Y$) datasets. This is a convex problem solvable in polynomial time (e.g., via the Hungarian algorithm or Sinkhorn iterations).
2. Clustering (Step 2): "Register" the cost matrix $C$ using the Monge map to create a new cost matrix $\tilde{C} = C P_{\sigma^*}^\top$. The LR-OT problem is then reduced to solving a Generalized K-Means problem on this registered cost $\tilde{C}$ to find a single assignment matrix $Q$. The second factor $R$ is automatically derived as $R = P_{\sigma^*}^\top Q$.

Theoretical Reduction

The paper proves that the LR-OT problem (a co-clustering problem over two datasets) can be approximated by a clustering problem over a single set of correspondences.

• Hard Assignment: The method focuses on "hard" transport plans (permutation-like structures) rather than soft couplings, aligning with the discrete nature of K-means.
• Approximation Guarantees: The reduction yields constant-factor approximation algorithms:
  • Negative-type metrics: $(1 + \gamma)$ approximation.
  • Kernel costs (e.g., Squared Euclidean): $(1 + \gamma + \sqrt{2\gamma})$ approximation.
  • General metrics: $(1 + \gamma + \rho)$ approximation.
  • Here, $\gamma \in [0, 1]$ is the ratio of the optimal full-rank cost to the optimal rank-$K$ cost. Since $\gamma$ is typically small, these bounds are close to 1.

Solvers for Generalized K-Means

To solve the resulting Generalized K-Means problem, the authors propose:

1. GKMS (Mirror Descent): A local optimization algorithm using exponentiated gradient updates with Sinkhorn projections.
2. SDP Approach: A semidefinite programming relaxation (based on Burer-Monteiro factorization) for higher-quality solutions.
3. Initialization: The algorithm uses a "Transport Registered Initialization," running K-means/K-medians on the source and target separately and selecting the best initialization, which guarantees the constant-factor approximation is maintained.

3. Key Contributions

1. Algorithmic Reduction: Demonstrates that LR-OT can be reduced to a Generalized K-Means problem via Monge registration, eliminating the need for complex multi-variable optimization.
2. Provable Guarantees: Provides the first constant-factor approximation guarantees for LR-OT, a significant theoretical advance over existing heuristic solvers.
3. Practical Efficiency: The algorithm is simple to implement, inherits the stability of modern K-means solvers, and scales to large datasets.
4. Unified Framework: Shows that K-means is a strict subset of LR-OT, and TC generalizes K-means to the co-clustering of two distinct datasets.

4. Experimental Results

The authors evaluated TC against state-of-the-art LR-OT solvers (LOT, FRLC, LatentOT) on synthetic and real-world datasets.

• Synthetic Benchmarks:

  • Datasets: 2-Moons to 8-Gaussians, Shifted Gaussians, and Stochastic Block Models (SBM).
  • Performance: TC consistently achieved the lowest transport cost across all noise levels and ranks. On the Shifted Gaussians dataset, TC showed a 23% relative improvement over the next best method.
  • Clustering Accuracy: TC achieved high Adjusted Rand Index (ARI) and Adjusted Mutual Information (AMI), often outperforming competitors in recovering ground-truth clusters.

• Large-Scale Real-World Applications:

  • CIFAR-10 (60k images): TC achieved the lowest OT cost (231.20 vs. 234.73 for LOT) and superior cross-domain label transfer accuracy (CTA: 0.412 vs. 0.358).
  • Single-Cell Transcriptomics (Mouse Embryogenesis): On massive datasets (up to 131,040 cells), TC successfully aligned timepoints where other methods (like LOT) failed to compute solutions due to scalability issues. TC demonstrated higher AMI/ARI and better cell-type mapping (CTA: 0.722 vs. 0.611).

• Wasserstein Distance Estimation:

  • TC provided the most accurate robust estimators of the squared Wasserstein distance on the fragmented hypercube benchmark, outperforming full-rank OT (which suffers from the curse of dimensionality) and other low-rank methods.

5. Significance

• Theoretical Breakthrough: This work bridges the gap between the theoretical hardness of LR-OT and practical solvability by establishing a reduction to clustering with rigorous approximation bounds.
• Scalability: By leveraging efficient K-means solvers and polynomial-time OT registration, TC enables LR-OT on datasets with tens of thousands of points, a regime previously inaccessible to many LR-OT methods.
• Robustness: The method is less sensitive to initialization than existing non-convex solvers and provides statistically robust estimates of Wasserstein distances, making it highly valuable for applications in biology (cell trajectory inference), computer vision, and generative modeling.

In summary, Transport Clustering transforms a difficult, non-convex optimization problem into a tractable clustering task, offering a theoretically grounded, scalable, and empirically superior solution for low-rank optimal transport.