UltraTWD: Optimizing Ultrametric Trees for Tree-Wasserstein Distance
Yu, Fangchen Chen, Yanzhen Wei, Jiaxing Mao, Jianfeng Li, Wenye Sun, Qiang
Yu, Fangchen
Chen, Yanzhen
Wei, Jiaxing
Mao, Jianfeng
Li, Wenye
Sun, Qiang
Supervisor
Department
Statistics and Data Science
Embargo End Date
Type
Conference proceeding
Date
2025
License
Language
English
Collections
Research Projects
Organizational Units
Journal Issue
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.
Citation
F. Yu, Y. Chen, J. Wei, J. Mao, W. Li, and Q. Sun, “UltraTWD: Optimizing Ultrametric Trees for Tree-Wasserstein Distance,” Oct. 06, 2025, PMLR. [Online]. Available: https://proceedings.mlr.press/v267/yu25a.html
Source
Proceedings of Machine Learning Research
Conference
42nd International Conference on Machine Learning, ICML 2025
Keywords
Subjects
Source
42nd International Conference on Machine Learning, ICML 2025
Publisher
ML Research Press