Learning Ultrametric Trees for Optimal Transport Regression
share
Optimal transport provides a metric which quantifies the dissimilarity between probability measures. For measures supported in discrete metric spaces, finding optimal transport distance has cubic time complexity in the size of the space. However, measures supported on trees admit a closed-form optimal transport which 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 can best approximate the optimal transport distance in the original space. One of our key ideas is to cast the problem in ultrametric spaces. This helps define different tree structures and allows us to optimize the tree structure via projected gradient decent over space of ultrametric matrices. During optimization, we project the parameters to the ultrametric space via a hierarchical minimum spanning tree algorithm. Experimental results on real datasets show that our approach outperforms previous approaches in approximating optimal transport distances. Finally, experiments on synthetic data generated on ground truth trees show that our algorithm can accurately uncover the underlying tree metrics.
READ FULL TEXT