Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Learning Ultrametric Trees for Optimal Transport Regression

Samantha Chen, Puoya Tabaghi, Yusu Wang

TL;DR

The paper tackles the cubic-time bottleneck of computing the 1-Wasserstein distance by learning an ultrametric tree that approximates OT on a discrete space. It casts the problem as optimal transport regression on ultrametrics and solves it with projected gradient descent, using a hierarchical minimum spanning tree as the projection operator to ultrametrics. The authors introduce a gradient-based learning scheme parameterized by the ultrametric’s height variables (LCAs) and provide a runtime analysis showing $O(k N^2)$ training complexity while enabling $O(N)$ inference on the learned tree. Empirical results demonstrate improved Wasserstein approximations over Flowtree and Quadtree on real graphs and synthetic dense distributions, with the added ability to recover ground-truth tree topology on synthetic data, highlighting practical impact for fast, accurate OT in discrete spaces.

Abstract

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.

Learning Ultrametric Trees for Optimal Transport Regression

TL;DR

The paper tackles the cubic-time bottleneck of computing the 1-Wasserstein distance by learning an ultrametric tree that approximates OT on a discrete space. It casts the problem as optimal transport regression on ultrametrics and solves it with projected gradient descent, using a hierarchical minimum spanning tree as the projection operator to ultrametrics. The authors introduce a gradient-based learning scheme parameterized by the ultrametric’s height variables (LCAs) and provide a runtime analysis showing training complexity while enabling inference on the learned tree. Empirical results demonstrate improved Wasserstein approximations over Flowtree and Quadtree on real graphs and synthetic dense distributions, with the added ability to recover ground-truth tree topology on synthetic data, highlighting practical impact for fast, accurate OT in discrete spaces.

Abstract

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.
Paper Structure (25 sections, 4 theorems, 15 equations, 6 figures, 6 tables, 1 algorithm)

This paper contains 25 sections, 4 theorems, 15 equations, 6 figures, 6 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Notation.
  3. Preliminaries
  4. Wasserstein Distance.
  5. Tree Wasserstein Distance.
  6. Optimal Transport Regression in Ultrametric Spaces
  7. A Projection Operator to Ultrametric Spaces.
  8. Computing the Gradient.
  9. Runtime Analysis.
  10. Why Use Ultrametrics?
  11. Experimental Results
  12. 1-Wasserstein Approximations
  13. Real-world Datasets.
  14. Synthetic Datasets.
  15. Changing Tree Structure
  16. ...and 10 more sections

Key Result

Theorem 2.1

Given two measures $\mu, \rho$ supported on $T = (V,E)$ with metric $d_T$, we have where $w_e$ is the weight of edge $e \in E$, and $v_e$ is the node of $e \in E$ that is farther from the root le2019tree.

Figures (6)

  • Figure 1: An overall summary of the our projected gradient descent procedure. $(a)$ Measurements are points in a metric space, vertices of a weighted graph, or a distance matrix. $(b_0)$ Hierarchical minimum spanning tree builds an ultrametric tree given a semimetric matrix. $(b_1)$ We compute a distance matrix for the tree leaves. $(b_2)$ We update the distance matrix by applying gradient descent on the optimal transport regression cost.
  • Figure 2: A stylized example of how the tree structure vary with heights of least common ancestors. Notice that in the second step, the heights of nodes $f$ and $c$ have increased while the height of node $g$ has decreased. Due to these changes in height, the recomputed distance between nodes $c$ and $d$ is smaller than the recomputed distance between $a$ and $b$ so $c$ and $d$ are the first to be clustered in the final projection to the new ultrametric (instead of node $c$ clustering with $a$ and $b$ in the original ultrametric).
  • Figure 3: $(a):$ A randomly generated tree $T$ with unit edge weights and leaf distance matrix $D_T$. $(b):$ We initialize the training process with noise-contaminated distance matrix, $\widetilde{D_T}$, and determine the initial tree topology on the left. We then train on true optimal transport distances and correctly recover the original tree structure on the right.
  • Figure 4: Training losses from epoch to epoch for real world datasets.
  • Figure 5: Changes in the mean relative error as the sparsity of the train and test distributions changes. Plot (a) shows the the changes in average relative error for discrete metric spaces with varying cardinalities as the sparsity of the train and test distributions increase. Plots (b) through (d) compare the average relative error for our learned ultrametric against cTWD and qTWD for discrete metric spaces with varying cardinalities.
  • ...and 1 more figures

Theorems & Definitions (5)

  • Theorem 2.1
  • Corollary 2.2
  • Definition 3.1
  • Theorem 3.3
  • Theorem A.1