UltraTWD: Optimizing Ultrametric Trees for Tree-Wasserstein Distance

Fangchen Yu, Yanzhen Chen, Jiaxing Wei, Jianfeng Mao, Wenye Li, Qiang Sun
Proceedings of the 42nd International Conference on Machine Learning, PMLR 267:72837-72852, 2025.

Abstract

The Wasserstein distance is a widely used metric for measuring differences between distributions, but its super-cubic time complexity introduces substantial computational burdens. To mitigate this, the tree-Wasserstein distance (TWD) offers a linear-time approximation by leveraging a tree structure; however, existing TWD methods often compromise accuracy due to suboptimal tree structures and edge weights. To address it, we introduce UltraTWD, a novel unsupervised framework that simultaneously optimizes both ultrametric tree structures and edge weights to more faithfully approximate the cost matrix. Specifically, we develop algorithms based on minimum spanning trees, iterative projection, and gradient descent to efficiently learn high-quality ultrametric trees. Empirical results across document retrieval, ranking, and classification tasks demonstrate that UltraTWD achieves superior approximation accuracy and competitive downstream performance. Code is available at: https://github.com/NeXAIS/UltraTWD.

Cite this Paper

BibTeX
@InProceedings{pmlr-v267-yu25a, title = {{U}ltra{TWD}: Optimizing Ultrametric Trees for Tree-{W}asserstein Distance}, author = {Yu, Fangchen and Chen, Yanzhen and Wei, Jiaxing and Mao, Jianfeng and Li, Wenye and Sun, Qiang}, booktitle = {Proceedings of the 42nd International Conference on Machine Learning}, pages = {72837--72852}, year = {2025}, editor = {Singh, Aarti and Fazel, Maryam and Hsu, Daniel and Lacoste-Julien, Simon and Berkenkamp, Felix and Maharaj, Tegan and Wagstaff, Kiri and Zhu, Jerry}, volume = {267}, series = {Proceedings of Machine Learning Research}, month = {13--19 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v267/main/assets/yu25a/yu25a.pdf}, url = {https://proceedings.mlr.press/v267/yu25a.html}, abstract = {The Wasserstein distance is a widely used metric for measuring differences between distributions, but its super-cubic time complexity introduces substantial computational burdens. To mitigate this, the tree-Wasserstein distance (TWD) offers a linear-time approximation by leveraging a tree structure; however, existing TWD methods often compromise accuracy due to suboptimal tree structures and edge weights. To address it, we introduce UltraTWD, a novel unsupervised framework that simultaneously optimizes both ultrametric tree structures and edge weights to more faithfully approximate the cost matrix. Specifically, we develop algorithms based on minimum spanning trees, iterative projection, and gradient descent to efficiently learn high-quality ultrametric trees. Empirical results across document retrieval, ranking, and classification tasks demonstrate that UltraTWD achieves superior approximation accuracy and competitive downstream performance. Code is available at: https://github.com/NeXAIS/UltraTWD.} }
Endnote
%0 Conference Paper %T UltraTWD: Optimizing Ultrametric Trees for Tree-Wasserstein Distance %A Fangchen Yu %A Yanzhen Chen %A Jiaxing Wei %A Jianfeng Mao %A Wenye Li %A Qiang Sun %B Proceedings of the 42nd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2025 %E Aarti Singh %E Maryam Fazel %E Daniel Hsu %E Simon Lacoste-Julien %E Felix Berkenkamp %E Tegan Maharaj %E Kiri Wagstaff %E Jerry Zhu %F pmlr-v267-yu25a %I PMLR %P 72837--72852 %U https://proceedings.mlr.press/v267/yu25a.html %V 267 %X The Wasserstein distance is a widely used metric for measuring differences between distributions, but its super-cubic time complexity introduces substantial computational burdens. To mitigate this, the tree-Wasserstein distance (TWD) offers a linear-time approximation by leveraging a tree structure; however, existing TWD methods often compromise accuracy due to suboptimal tree structures and edge weights. To address it, we introduce UltraTWD, a novel unsupervised framework that simultaneously optimizes both ultrametric tree structures and edge weights to more faithfully approximate the cost matrix. Specifically, we develop algorithms based on minimum spanning trees, iterative projection, and gradient descent to efficiently learn high-quality ultrametric trees. Empirical results across document retrieval, ranking, and classification tasks demonstrate that UltraTWD achieves superior approximation accuracy and competitive downstream performance. Code is available at: https://github.com/NeXAIS/UltraTWD.
APA
Yu, F., Chen, Y., Wei, J., Mao, J., Li, W. & Sun, Q.. (2025). UltraTWD: Optimizing Ultrametric Trees for Tree-Wasserstein Distance. Proceedings of the 42nd International Conference on Machine Learning, in Proceedings of Machine Learning Research 267:72837-72852 Available from https://proceedings.mlr.press/v267/yu25a.html.

Related Material