Research map / Learning theory and algorithms / Optimal transport
Fast unsupervised ground metric learning with tree-Wasserstein distance
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Tree-WSV, which embeds samples and features on trees and learns edge weights so tree-Wasserstein distances approximate Wasserstein singular vectors, cutting complexity to O(n^3+m^3+mn) and scaling unsupervised ground-metric learning to scRNA-seq clustering.
Problem
Unsupervised ground metric learning via Wasserstein singular vectors (WSVs) is accurate but prohibitively expensive, with complexity O(n^2 m^2 (n log n + m log m)) per power iteration, limiting use on large unlabelled datasets such as single-cell RNA-seq.
Approach
Replace the full pairwise Wasserstein distance matrices in WSV power iterations with tree-Wasserstein distances on trees built over samples and features, learn positive tree edge weights by solving a (non-negative) linear system using a recursively computed basis of leaf-to-leaf path matrices, and iterate with meta-iterations that rebuild trees from learned metrics.
Key results
- On PBMC 3k, Tree-WSV (best) reached ASW 0.313 in 72 min versus 0.348 for entropic WSV (SSV) in 110 min; on Neurons V1, Tree-WSV best reached ASW 0.602 vs 0.256 for SSV.
- On the 7k-cell Lung dataset SSV did not converge after 12 hours, while Tree-WSV (best) reached ASW 0.457 in 71 min.
- On the toy torus dataset, tree methods ran 2-3 orders of magnitude faster than standard WSV, with 3-children trees taking about 1% of standard WSV time, and had lower Frobenius error than entropic WSV.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Tree-WSV (best, max 15 iterations) instead of Entropic WSV (SSV) | PBMC 3k clustering · PBMC 3k (2k cells) | ASW: 0.348 -> 0.313 | Hurt |
| Use Tree-WSV (best, max 15 iterations) instead of Entropic WSV (SSV) | Neurons V1 clustering · Neurons V1 (1k cells) | ASW: 0.256 -> 0.602 | Helped |
| Use Tree-WSV (best, max 15 iterations) instead of Entropic WSV (SSV) | Human Lung Cell Atlas clustering · Human Lung Cell Atlas (7k cells) | ASW | Helped |
| Use Tree-WSV (4 iterations) instead of Entropic WSV (SSV) | PBMC 3k clustering · PBMC 3k (2k cells) | Runtime: 110 min -> 8 min | Helped |
| Use Tree-WSV (4 iterations) instead of Entropic WSV (SSV) | Neurons V1 clustering · Neurons V1 (1k cells) | Runtime: 70 min -> 7 min | Helped |
| Use Tree-WSV (4 iterations) instead of Entropic WSV (SSV) | Human Lung Cell Atlas clustering · Human Lung Cell Atlas (7k cells) | Runtime: > 12 hours -> 27 min | Helped |
| Use 3-child trees instead of 2-child trees (ablation) | Tree-WSV · Synthetic torus dataset | Frobenius norm error | Helped |
| Use trees with fewer children per node instead of trees with more children per node (ablation) | Tree-WSV · Synthetic torus dataset | Frobenius norm error | Helped |
| Use trees with fewer children per node instead of trees with more children per node (ablation) | Tree-WSV · Synthetic torus dataset | Runtime | Hurt |
| Run meta-iterations rather than only one Tree-WSV iteration (ablation) | Tree-WSV · PBMC 3k | ASW: 0.07 -> 0.299 | Helped |
| Run meta-iterations rather than only one Tree-WSV iteration (ablation) | Tree-WSV · Neurons V1 | ASW: 0.15 -> 0.436 | Helped |
| Run meta-iterations rather than only one Tree-WSV iteration (ablation) | Tree-WSV · Human Lung Cell Atlas | ASW: 0.05 -> 0.104 | Helped |
Surprising findings
- Tree-WSV with a single iteration gave low silhouette scores (0.07, 0.15, 0.05), which improved substantially only after meta-iterations that rebuild the tree from learned metrics.
- Entropic WSV (SSV) achieved the best ASW on PBMC (0.348), slightly above Tree-WSV (0.313).
Limitations
- Tree construction via ClusterTree depends on Euclidean distances at initialisation, and performance depends on meta-iterations; ASW had not converged after 15 meta-iterations.
- Theoretical convergence of the power iterations is not guaranteed; only empirically observed.
- Root degree must be at least 3 for uniqueness of solutions; 2-trees perform worse.
Experimental setup
Datasets: PBMC 3k (10x Genomics) Neurons V1 mouse scRNA-seq (GSE71585) Human Lung Cell Atlas subset synthetic torus dataset
Benchmarks: PBMC 3k Neurons V1 Human Lung Cell Atlas subset
Metrics: average silhouette width runtime Frobenius norm error
Baselines: Standard WSV Entropic WSV (Sinkhorn singular vectors, SSV) Euclidean distance PCA / l2 Kernel PCA / l2 scVI embedding / l2 Sinkhorn Gene Mover Distance
Builds on: Wasserstein singular vectors Sinkhorn singular vectors tree-Wasserstein distance ClusterTree
Closest papers by meaning
- Unsupervised Ground Metric Learning Using Wasserstein Singular Vectors (ICML 2022 · similarity 0.94)
- Tree-Wasserstein Distance for High Dimensional Data with a Latent Feature Hierarchy (ICLR 2025 · similarity 0.89)
- Joint Hierarchical Representation Learning of Samples and Features via Informed Tree-Wasserstein Distance (NeurIPS 2025 · similarity 0.88)
- UltraTWD: Optimizing Ultrametric Trees for Tree-Wasserstein Distance (ICML 2025 · similarity 0.87)
- A New Robust Partial p-Wasserstein-Based Metric for Comparing Distributions (ICML 2024 · similarity 0.85)
- Tree-Sliced Wasserstein Distance with Nonlinear Projection (ICML 2025 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.