  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.
