Research map / Learning theory and algorithms / Optimal transport
Learning Ultrametric Trees for Optimal Transport Regression
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.
Learns an ultrametric tree (topology and weights) for a discrete metric space by projected gradient descent on an optimal-transport regression cost, so tree-Wasserstein distances approximate true OT distances.
Problem
Approximating optimal transport (Wasserstein-1) distances on discrete metric spaces with linear-time tree-Wasserstein distances, where existing methods either fix the tree topology or build trees that approximate the metric rather than the Wasserstein distances.
Approach
Cast tree-metric learning as optimization over ultrametric matrices, using projected gradient descent where the projection is a hierarchical minimum spanning tree (closest ultrametric under sup norm), so both tree weights and topology change across iterations to minimize squared OT regression error.
Key results
- On the synthetic uniform dataset, Ult. Tree achieves mean relative error 0.104 at dim=2 and 0.011 at dim=17, versus 0.522 and 0.348 for Flowtree and 0.563 and 0.219 for cTWD.
- On graph datasets (USCA312, USAir97, Belfast), Ult. Tree outperforms the compared methods, e.g. 0.282 vs 0.602 for cTWD on USCA312 and 0.043 vs 0.122 for cTWD on USAir97.
- On the word datasets Twitter and BBCSport, Ult. Tree is worse than weight-optimized cTWD/qTWD variants, which the authors attribute to sparse supports; runtime is higher than cTWD (e.g. 0.42 s vs 0.0004 s on BBCSport).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| learn an ultrametric tree instead of using sliced-cTWD | 1-Wasserstein approximation · BBCSport | mean relative error: 0.020 $\pm$ 0.002 vs 0.023 $\pm$ 0.010 | Hurt |
| learn an ultrametric tree instead of using qTWD | 1-Wasserstein approximation · Twitter | mean relative error: 0.045 $\pm$ 0.012 vs 0.062 $\pm$ 0.034 | Hurt |
| learn an ultrametric tree instead of using cTWD | 1-Wasserstein approximation · USCA312, USAir97, Belfast | mean relative error: 0.602 $\pm$ 0.142 vs 0.282 $\pm$ 0.138; 0.122 $\pm$ 0.008 vs 0.043 $\pm$ 0.018; 1.089 $\pm$ 0.117 vs 0.368 $\pm$ 0.064 | Helped |
| learn an ultrametric tree instead of using Sinkhorn | 1-Wasserstein approximation · Uniform synthetic, dim=2 | mean relative error: 0.587 $\pm$ 0.257 vs 0.104 $\pm$ 0.074 | Helped |
| learn an ultrametric tree instead of using Sinkhorn | 1-Wasserstein approximation · Uniform synthetic, dim=5 | mean relative error: 0.071 $\pm$ 0.053 vs 0.029 $\pm$ 0.023 | Helped |
| learn an ultrametric tree instead of using Sinkhorn | 1-Wasserstein approximation · Uniform synthetic, dim=8 | mean relative error: 0.063 $\pm$ 0.046 vs 0.016 $\pm$ 0.013 | Helped |
| recompute MST projections throughout training rather than skipping them until the final iteration (ablation) | Ultrametric tree optimization · Noisy random tree, σ=5 | Wasserstein error: 0.279 $\pm$ 0.102 -> 0.262 $\pm$ 0.102 | Helped |
| recompute MST projections throughout training rather than skipping them until the final iteration (ablation) | Ultrametric tree optimization · Noisy random tree, σ=4 | Wasserstein error: 0.286 $\pm$ 0.109 -> 0.275 $\pm$ 0.257 | Helped |
| recompute MST projections throughout training rather than skipping them until the final iteration (ablation) | Ultrametric tree optimization · Noisy random tree, σ=3 | Wasserstein error: 0.281 $\pm$ 0.098 -> 0.265 $\pm$ 0.092 | Helped |
| recompute MST projections throughout training rather than skipping them until the final iteration (ablation) | Ultrametric tree optimization · Noisy random tree, σ=2 | Wasserstein error: 0.283 $\pm$ 0.108 -> 0.268 $\pm$ 0.095 | Helped |
| update tree topology during training rather than keeping the initialized topology fixed | Ultrametric tree optimization versus cTWD · Adversarially initialized tree example | mean relative error: 0.193 \pm 0.101 -> zero error | Helped |
| update tree topology during training rather than keeping the initialized topology fixed | Ultrametric tree optimization versus cTWD · 100 random incorrect tree initializations | mean relative error: 0.178 \pm 0.103 -> zero error | Helped |
Surprising findings
- On sparse word-frequency datasets (Twitter, BBCSport), the proposed topology-learning method underperforms fixed-topology weight-optimized methods such as cTWD and qTWD.
- Starting from an adversarial initial tree, cTWD with fixed topology has about 0.18-0.19 relative error, while the ultrametric method recovers the true tree exactly in all 100 random trials.
Limitations
- Projection to ultrametrics via hierarchical MST costs O(N^2) per iteration, a computational bottleneck.
- Underperforms weight-optimized fixed-topology methods on sparse real-world word datasets.
- Nonconvex optimization with no guarantee of achieving the log n distortion after training.
Experimental setup
Datasets: Twitter BBCSport USCA312 USAir97 Belfast public transit graph RNAseq (Allen Institute) synthetic uniform point sets in [-10,10]^d synthetic Gaussian point sets
Benchmarks: Twitter BBCSport USCA312 USAir97 Belfast RNAseq synthetic uniform dataset dim 2-17 synthetic Gaussian dataset
Metrics: mean relative error of 1-Wasserstein distance runtime tree distance dist(T_hat, T)
Baselines: Sinkhorn Flowtree Quadtree qTWD cTWD Sliced-qTWD Sliced-cTWD skip-MST
Builds on: Quadtree Wasserstein distance Flowtree cluster tree Wasserstein distance (cTWD) Sinkhorn distance
Closest papers by meaning
- Tree-Sliced Wasserstein Distance: A Geometric Perspective (ICML 2025 · similarity 0.88)
- Geometry-Aware Optimal Transport: Fast Intrinsic Dimension and Wasserstein Distance Estimation (NeurIPS 2026 · similarity 0.87)
- An Efficient Orlicz-Sobolev Approach for Transporting Unbalanced Measures on a Graph (NeurIPS 2025 · similarity 0.87)
- Spherical Tree-Sliced Wasserstein Distance (ICLR 2025 · similarity 0.87)
- Generalized Sobolev Transport for Probability Measures on a Graph (ICML 2024 · similarity 0.87)
- A Combinatorial Algorithm for the Semi-Discrete Optimal Transport Problem (NeurIPS 2024 · similarity 0.87)
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.