# Supervised Tree-Wasserstein Distance

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

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 17 of its 30 references that are in the map.
- First posted on arXiv 27 Jan 2021.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Distributed Representations of Words and Phrases and their Compositionality](https://mlsanity.com/paper/1310.4546/) (2013)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [SuperGlue: Learning Feature Matching with Graph Neural Networks](https://mlsanity.com/paper/1911.11763/) (2019)
- [Poincar'e Embeddings for Learning Hierarchical Representations](https://mlsanity.com/paper/1705.08039/) (2017)
- [Large Batch Training of Convolutional Networks](https://mlsanity.com/paper/1708.03888/) (2017)
- [From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label Classification](https://mlsanity.com/paper/1602.02068/) (2016)
- [Generalized Sliced Wasserstein Distances](https://mlsanity.com/paper/1902.00434/) (2019)
- [Hyperbolic Entailment Cones for Learning Hierarchical Embeddings](https://mlsanity.com/paper/1804.01882/) (2018)
- [Max-Sliced Wasserstein Distance and its use for GANs](https://mlsanity.com/paper/1904.05877/) (2019)
- [From Trees to Continuous Embeddings and Back: Hyperbolic Hierarchical Clustering](https://mlsanity.com/paper/2010.00402/) (2020)
- [A cost function for similarity-based hierarchical clustering](https://mlsanity.com/paper/1510.05043/) (2015)
- [Sliced Wasserstein Distance for Learning Gaussian Mixture Models](https://mlsanity.com/paper/1711.05376/) (2017)
- [Tree-Sliced Variants of Wasserstein Distances](https://mlsanity.com/paper/1902.00342/) (2019)
- [Hierarchical Optimal Transport for Document Representation](https://mlsanity.com/paper/1906.10827/) (2019)
- [Scalable Nearest Neighbor Search for Optimal Transport](https://mlsanity.com/paper/1910.04126/) (2019)
- [Fast Unbalanced Optimal Transport on a Tree](https://mlsanity.com/paper/2006.02703/) (2020)

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

- [Scalable Optimal Transport Methods in Machine Learning: A Contemporary Survey](https://mlsanity.com/paper/2305.05080/) (2023)
- [Re-evaluating Word Mover's Distance](https://mlsanity.com/paper/2105.14403/) (2021)
- [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)
- [SEAL: Simultaneous Label Hierarchy Exploration And Learning](https://mlsanity.com/paper/2304.13374/) (2023)
- [Tree-Wasserstein Distance for High Dimensional Data with a Latent Feature Hierarchy](https://mlsanity.com/paper/2410.21107/) (2024)
- [Learning Structured Representations by Embedding Class Hierarchy with Fast Optimal Transport](https://mlsanity.com/paper/2410.03052/) (2024)
- [Fast unsupervised ground metric learning with tree-Wasserstein distance](https://mlsanity.com/paper/2411.07432/) (2024)
- [An Empirical Study of Self-supervised Learning with Wasserstein Distance](https://mlsanity.com/paper/2310.10143/) (2023)
- [Label tree semantic losses for rich multi-class medical image segmentation](https://mlsanity.com/paper/2507.15777/) (2025)

## Abstract (arXiv)

To measure the similarity of documents, the Wasserstein distance is a powerful tool, but it requires a high computational cost. Recently, for fast computation of the Wasserstein distance, methods for approximating the Wasserstein distance using a tree metric have been proposed. These tree-based methods allow fast comparisons of a large number of documents; however, they are unsupervised and do not learn task-specific distances. In this work, we propose the Supervised Tree-Wasserstein (STW) distance, a fast, supervised metric learning method based on the tree metric. Specifically, we rewrite the Wasserstein distance on the tree metric by the parent-child relationships of a tree and formulate it as a continuous optimization problem using a contrastive loss. Experimentally, we show that the STW distance can be computed fast, and improves the accuracy of document classification tasks. Furthermore, the STW distance is formulated by matrix multiplications, runs on a GPU, and is suitable for batch processing. Therefore, we show that the STW distance is extremely efficient when comparing a large number of documents.