# Learning Ultrametric Trees for Optimal Transport Regression

arXiv: https://arxiv.org/abs/2210.12288 · HTML: https://mlsanity.com/paper/2210.12288/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 12 of its 42 references that are in the map.
- First posted on arXiv 21 Oct 2022.

## Builds on (most cited first, 12 of 12)

- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Fast Computation of Wasserstein Barycenters](https://mlsanity.com/paper/1310.4375/) (2013)
- [Generalized Sliced Wasserstein Distances](https://mlsanity.com/paper/1902.00434/) (2019)
- [Learning with a Wasserstein Loss](https://mlsanity.com/paper/1506.05439/) (2015)
- [Stochastic Optimization for Large-scale Optimal Transport](https://mlsanity.com/paper/1605.08527/) (2016)
- [Tree-Sliced Variants of Wasserstein Distances](https://mlsanity.com/paper/1902.00342/) (2019)
- [Scalable Nearest Neighbor Search for Optimal Transport](https://mlsanity.com/paper/1910.04126/) (2019)
- [A Sparse Multi-Scale Algorithm for Dense Optimal Transport](https://mlsanity.com/paper/1510.05466/) (2015)
- [Dynamical Optimal Transport on Discrete Surfaces](https://mlsanity.com/paper/1809.07083/) (2018)
- [Fixed Support Tree-Sliced Wasserstein Barycenter](https://mlsanity.com/paper/2109.03431/) (2021)
- [Approximating 1-Wasserstein Distance with Trees](https://mlsanity.com/paper/2206.12116/) (2022)
- [Tree-Wasserstein Barycenter for Large-Scale Multilevel Clustering and Scalable Bayes](https://mlsanity.com/paper/1910.04483/) (2019)

## Built on by (most cited first, 3 of 3)

- [Tree-Wasserstein Distance for High Dimensional Data with a Latent Feature Hierarchy](https://mlsanity.com/paper/2410.21107/) (2024)
- [An Empirical Study of Self-supervised Learning with Wasserstein Distance](https://mlsanity.com/paper/2310.10143/) (2023)
- [On Hamming-Lipschitz Type Stability of the Subdominant (Minmax) Ultrametric: Theory and Simple Proofs](https://mlsanity.com/paper/2608.04014/) (2026)

## Abstract (arXiv)

Optimal transport provides a metric which quantifies the dissimilarity between probability measures. For measures supported in discrete metric spaces, finding the optimal transport distance has cubic time complexity in the size of the space. However, measures supported on trees admit a closed-form optimal transport that can be computed in linear time. In this paper, we aim to find an optimal tree structure for a given discrete metric space so that the tree-Wasserstein distance approximates the optimal transport distance in the original space. One of our key ideas is to cast the problem in ultrametric spaces. This helps us optimize over the space of ultrametric trees -- a mixed-discrete and continuous optimization problem -- via projected gradient decent over the space of ultrametric matrices. During optimization, we project the parameters to the ultrametric space via a hierarchical minimum spanning tree algorithm, equivalent to the closest projection to ultrametrics under the supremum norm. Experimental results on real datasets show that our approach outperforms previous approaches (e.g. Flowtree, Quadtree) in approximating optimal transport distances. Finally, experiments on synthetic data generated on ground truth trees show that our algorithm can accurately uncover the underlying trees.