Tree-Wasserstein Distance for High Dimensional Data with a Latent Feature Hierarchy

## Abstract

Abstract Finding meaningful distances between high-dimensional data samples is an important scientific task. To this end, we propose a new tree-Wasserstein distance (TWD) for high-dimensional data with two key aspects. First, our TWD is specifically designed for data with a latent feature hierarchy, i.e., the features lie in a hierarchical space, in contrast to the usual focus on embedding samples in hyperbolic space. Second, while the conventional use of TWD is to speed up the computation of the Wasserstein distance, we use its inherent tree as a means to learn the latent feature hierarchy. The key idea of our method is to embed the features into a multi-scale hyperbolic space using diffusion geometry and then present a new tree decoding method by establishing analogies between the hyperbolic embedding and trees. We show that our TWD computed based on data observations provably recovers the TWD defined with the latent feature hierarchy and that its computation is efficient and scalable. We showcase the usefulness of the proposed TWD in applications to word-document and single-cell RNA-sequencing datasets, demonstrating its advantages over existing TWDs and methods based on pre-trained models.

## 1 Introduction

High-dimensional datasets with a latent feature hierarchy [Fefferman et al., ] are prevalent across a broad range of applications, involving data such as genomic sequences [Tanay and Regev, ], text and images [Krishna et al., ], and movie user ratings [Bennett et al., ]. In this setting, the assumption is that there is an unknown underlying hierarchical structure to the data features, which is not directly observable. A key component of meaningful analyses of such data is finding distances between the data samples that take into account the underlying hierarchical structure of the features.
As opposed to relying on generic metrics, e.g., the Euclidean or correlation distances, such meaningful distances serve as a fundamental building block in a broad range of downstream tasks, such as classification and clustering, often leading to significant performance improvement.

Recently, hyperbolic geometry [Ratcliffe et al., ] has gained prominence in hierarchical representation learning [Chamberlain et al., , Nickel and Kiela, ] due to its exponential growth [Sarkar, ].
Methods using hyperbolic geometry typically focus on finding a hyperbolic embedding of the samples, relying on a (partially) known graph, whose nodes represent the samples [Sala et al., ].
However, considering such a *known* hierarchical structure of the *samples* is fundamentally different than the problem we consider here, where we aim to find meaningful distances between data samples that incorporate the
*latent* hierarchical structure of the *features*.

In this paper, we introduce a new tree-Wasserstein distance (TWD) [Indyk and Thaper, ] for this purpose.
We propose a two-step approach.
In the first step, we embed the features into continuous hyperbolic spaces [Bowditch, ] utilizing diffusion geometry [Coifman and Lafon, ] to approximate the hidden hierarchical metric underlying the features.
In the second step, we construct a tree that represents the latent feature hierarchy in a bottom-up manner from the hyperbolic embedding of the features.
The tree construction is based on a novel notion of a continuous analog of the lowest common ancestor (LCA) in trees defined in hyperbolic spaces by associating shortest paths on trees with hyperbolic geodesic distances.
Utilizing the decoded feature tree, we introduce a new TWD for high-dimensional data with a latent feature hierarchy, based on diffusion geometry and hyperbolic geometry.
We remark that in Lin et al. [], integrating hyperbolic and diffusion geometries was shown to be more effective in recovering latent hierarchies than existing hyperbolic embedding methods [Nickel and Kiela, ] (see App. E.3.4 for more details).
However, the work in Lin et al. [] considers implicitly latent hierarchies of the samples, whereas here, we explicitly reveal the latent hierarchical structure of the features through a tree, which is then utilized for finding a meaningful TWD between the samples (see App. F.4).
Using a TWD entails two useful attributes.
First, the Wasserstein distance [Monge, , Kantorovich, ] enables us to inherently integrate mutual relations between the features into a geometrically meaningful distance between the samples.
Second, we use a tree to represent the latent feature hierarchy, learn this tree based on the data, and then, naturally incorporate it in the proposed TWD. Notably, this use of a tree significantly differs from the conventional use of trees in TWD for speeding up the computation of Wasserstein distance.

From a theoretical standpoint, we show that our proposed TWD provably recovers the TWD between samples with the true latent feature hierarchy.
From a practical standpoint, we show that the proposed TWD is computationally efficient and scalable, making it suitable for large datasets.
We showcase the efficacy of our method on an illustrative synthetic example, word-document datasets, and single-cell RNA-sequencing datasets. Specifically, we show that our
TWD leads to superior classification performance compared to existing TWD-based methods [Indyk and Thaper, , Backurs et al., , Le et al., , Yamada et al., ] and two application-specific methods based on pre-trained models (WMD [Kusner et al., ] and GMD [Bellazzi et al., ]).

Our main contributions are as follows.
(i) We present a new TWD for data with a hidden feature hierarchy, which provably recovers the TWD with the true latent feature tree and can be efficiently computed.
(ii) We present a novel data-driven tree decoding method based on a new definition of the LCA counterpart in high-dimensional hyperbolic spaces.
(iii) We demonstrate the empirical advantages of our TWD over existing TWDs and pre-trained Wasserstein distance methods.

## 2 Related Work

Wasserstein Distance.
The Wasserstein distance [Villani, ] is a powerful tool for transforming the relations between features into a distance between probability distributions (samples from the probability simplex).
It can be computed by solving the optimal transport (OT) problem [Monge, , Kantorovich, ] with linear programming [Peyré et al., ] in cubic complexity.
Various approaches, such as adding entropic regularization [Cuturi, ], using wavelet functions [Shirdhonkar and Jacobs, , Gavish et al., ], applying Kantorovich-Rubinstein dual [Leeb and Coifman, , Arjovsky et al., ],
and computing the average Wasserstein distance across one-dimensional projections [Rabin et al., ], have been developed to reduce the computational complexity.

Tree-Wasserstein Distance.
Another popular alternative to reduce the computational complexity of the Wasserstein distance is the tree-Wasserstein distance (TWD) [Indyk and Thaper, ], which can be computed in linear time, enabling efficient comparisons of large datasets.
Existing TWD methods [Leeb, , Sato et al., , Takezawa et al., , Yamada et al., ] and graph-based algorithms (e.g., MST) [Alon et al., ] attempt to approximate the original feature metric with a tree metric by constructing the tree using the standard metric in the Euclidean ambient space.
Our approach has a different goal.
In addition to accelerating the computation, we use the tree to represent the *latent feature hierarchy*.
Consequently, the resulting TWD between the distributions uses the latent feature hierarchy to define a meaningful distance between the samples.

Hyperbolic Representation Learning.
While effective in computing embeddings and distances for hierarchical data samples, current hyperbolic representation learning methods [Chamberlain et al., , Nickel and Kiela, , ] are usually given a hierarchical graph, e.g., a tree, and do not include decoding a tree from observational data.
One exception is the work in Chami et al. [], which introduces a method for decoding a tree in a two-dimensional Poincaré disk for hierarchical clustering [Murtagh and Contreras, ], by minimizing the continuous Dasgupta’s cost [Dasgupta, ].
However, this approach does not approximate the TWD between samples with a latent feature hierarchy because it does not recover the tree underlying the features.

## 3 Preliminaries

Notation.
For $m\in\mathbb{N}$, we denote $[m]=\{1,\ldots,m\}$.
Let $\mathbf{X}\in\mathbb{R}^{n\times m}$ denote a data matrix with $n$ rows (samples) and $m$ columns (features).
We denote by $\mathbf{X}_{i\cdot}$ the $i$-th row (sample).
Let $d\simeq d^{\prime}$ denote the bilipschitz equivalence between metrics $d$ and $d^{\prime}$ defined on the same set $\mathcal{Y}$.
That is, there exist constants $c,C>0$ such that $cd(y_{1},y_{2})\leq d^{\prime}(y_{1},y_{2})\leq Cd(y_{1},y_{2})$ for all $y_{1},y_{2}\in\mathcal{Y}$.

Tree-Wasserstein Distance.
Consider a tree $T=([m],E,\mathbf{A})$ rooted at node 1 w.l.o.g., where $E$ is the edge set and $\mathbf{A}\in\mathbb{R}^{m\times m}$ contains edge weights.
The tree metric $d_{T}$ between two nodes is the length of the shortest path on $T$.
Let $\Gamma_{T}(v)$ be the set of nodes in the subtree of $T$ rooted at $v\in V$.
For any $u\in V$, there exists a unique node $v$, which is the parent of $u$.
We denote by $\omega_{u}=d_{T}(u,v)$ the length of the edge between any node $u$ and its parent $v$ (see App. A.4 for additional background on trees).
The tree-Wasserstein distance (TWD) [Indyk and Thaper, , Evans and Matsen, , Le et al., ] between two probability distributions $\bm{\mu}_{1},\bm{\mu}_{2}\in\mathbb{R}^{m}$ supported on the tree $T$ is defined by

$$ $\mathtt{TW}(\bm{\mu}_{1},\bm{\mu}_{2},T)=\sum_{v\in V}\omega_{v}\left\lvert \sum_{u\in\Gamma_{T}(v)}\left(\bm{\mu}_{1}(u)-\bm{\mu}_{2}(u)\right)\right\rvert.$ (1) $$

Diffusion Geometry.
Consider a set of points $\{\mathbf{a}_{j}\in\mathbb{R}^{n}\}_{j=1}^{m}$ assumed to lie on a manifold embedded in $\mathbb{R}^{n}$.
Let $\mathbf{Q}\in\mathbb{R}^{m\times m}$ be an affinity matrix, whose elements are given by $\mathbf{Q}_{jj^{\prime}}=\exp\left(-d^{2}(j,j^{\prime})/\epsilon\right)$, where $d(j,j^{\prime})$ is some suitable distance between $\mathbf{a}_{j}$ and $\mathbf{a}_{j^{\prime}}$, and $\epsilon>0$ is a scale parameter.
Following Coifman and Lafon [], we consider the density-normalized affinity matrix $\widehat{\mathbf{Q}}=\mathbf{D}^{-1}\mathbf{Q}\mathbf{D}^{-1}$, where $\mathbf{D}$ is diagonal with entries $\mathbf{D}_{jj}=\textstyle\sum_{j^{\prime}}\mathbf{Q}_{jj^{\prime}}$, and the diffusion operator $\mathbf{P}$, defined by

$$ $\mathbf{P}=\widehat{\mathbf{Q}}\widehat{\mathbf{D}}^{-1},\text{ where } \widehat{\mathbf{D}}_{jj}=\textstyle\sum_{j^{\prime}}\widehat{\mathbf{Q}}_{jj^ {\prime}}.$ (2) $$

The matrix $\widehat{\mathbf{Q}}$ can be viewed as edge weights of a graph $G=([m],E,\widehat{\mathbf{Q}})$.
Since the diffusion operator $\mathbf{P}$ is column-stochastic, it can be viewed as a transition probability matrix of a Markov chain on the graph $G$.
Therefore, $\mathbf{P}$ can propagate densities between nodes.
Specifically, the vector $\mathbf{p}^{t}_{j}=\mathbf{P}^{t}\mathbf{e}_{j}\in\mathbb{R}^{m}$ is the propagated density after diffusion time $t\in\mathbb{R}$ of a density concentrated at node $j$, where $\mathbf{e}_{j}\in\mathbb{R}^{m}$ is the indicator vector of the $j$-th node, $\mathbf{p}^{t}_{j}(i)\geq 0$ and $\left\lVert\mathbf{p}^{t}_{j}\right\rVert_{1}=1$.

The diffusion operator has been shown to exhibit favorable convergence properties [Coifman and Lafon, ]. In the limits, as the number of features $m$ increases toward infinity, the scaling parameter $\epsilon$ approaches zero, and the operator $\mathbf{P}^{t/\epsilon}$ converges pointwise to the Neumann heat kernel $\mathbf{H}_{t}=\exp(-t\Delta)$, associated with the underlying manifold, where $\Delta$ is the Laplace–Beltrami operator on the manifold. That is, the diffusion operator serves as a discrete approximation of the continuous heat kernel, thereby capturing the geometric structure of the manifold within a finite-dimensional framework [Coifman and Lafon, , Singer, , Belkin and Niyogi, ].

Poincaré Half-Space Model of Hyperbolic Geometry.
Hyperbolic geometry is a Riemannian geometry with constant negative curvature $-1$ [Beardon, ].
We utilize the $m$-dimensional Poincaré half-space model $\mathbb{H}^{m}$, defined by $\mathbb{H}^{m}=\{\mathbf{x}\in\mathbb{R}^{m}|\mathbf{x}(m)>0\}$ with the Riemannian metric tensor $ds^{2}=(d\mathbf{x}^{2}(1)+d\mathbf{x}^{2}(2)+\ldots+d\mathbf{x}^{2}(m))/
\mathbf{x}^{2}(m)$.
The Riemannian distance for any two hyperbolic point $\mathbf{x},\mathbf{y}\in\mathbb{H}^{m}$ is defined by $d_{\mathbb{H}^{m}}(\mathbf{x},\mathbf{y})=2\sinh^{-1}(\left\lVert\mathbf{x}-
\mathbf{y}\right\rVert_{2}/(2\sqrt{\mathbf{x}(m)\mathbf{y}(m)}))$.

Rooted Binary Tree.
A rooted and balanced binary tree $B$ with $m$ leaves is a tree whose root node has degree 2, $m$ leaves have degree 1, and $m-2$ internal nodes have degree 3.
For any two leaves $j,j^{\prime}$, let $j\vee j^{\prime}$ and $\Gamma_{B}(j\vee j^{\prime})$ denote their lowest common ancestor (LCA) and the set of leaves in the subtree of $B$ rooted at $j\vee j^{\prime}$, respectively.
Given any three leaves $j_{1},j_{2},j_{3}$, the LCA relation $[j_{2}\vee j_{3}\sim j_{1}]_{B}$ holds if $\Gamma_{B}(j_{2}\vee j_{3})\subset\Gamma_{B}(j_{1}\vee j_{2}\vee j_{3})$, as shown in Fig. 1(a).
For binary trees, only one of $[j_{1}\vee j_{2}\sim j_{3}]_{B}$, $[j_{1}\vee j_{3}\sim j_{2}]_{B}$, or $[j_{3}\vee j_{2}\sim j_{1}]_{B}$ can hold [Wang and Wang, ].

## 4 Proposed Method

### 4.1 Problem Formulation

Consider a data matrix $\mathbf{X}\in\mathbb{R}^{n\times m}$, whose $n$ rows consist of data samples and $m$ columns consist of features.
Suppose the features exhibit a latent hierarchical structure.
Specifically, the feature hierarchy is modeled by a hidden rooted weighted tree $T=([m],E,\mathbf{A},\mathbf{X}^{\top})$ on the features, where the $j$-th feature is associated with tree node $j$, $E$ is the edge set, $\mathbf{A}\in\mathbb{R}^{m\times m}$ represents the weight matrix, and the data matrix $\mathbf{X}$ is the collection of the $m$ node attributes (i.e., $m$ data features) in $n$-dimensions (i.e., $n$ data samples). We assume that the data matrix $\mathbf{X}$ can be viewed as a function of the latent feature tree $T$ and that the tree is implicitly conveyed by the mutual relations between the data features.
One typical example is a word-document dataset, where $n$ is the number of documents (samples), $m$ is the number of words (features), and the element $\mathbf{X}_{ij}\geq 0$ quantifies the occurrence of the $j$-th word in the $i$-th document.
The hierarchy underlying the features (words) is often assumed because words are naturally organized as nodes in a tree, each connected to others through various hierarchical linguistic relationships [Borge-Holthoefer and Arenas, ].

We aim to devise an informative distance between samples (rows) in the data matrix $\mathbf{X}$ while incorporating the latent hierarchical structure $T$ underlying the features (columns).
Following a large body of work [Villani, , Cuturi, , Courty et al., , Solomon et al., , Yurochkin et al., , Alvarez-Melis and Jaakkola, ], we focus on samples with positive entries that can be normalized into discrete histograms, which is natural when the distribution over features is more important than the total mass.
In particular, we study word-document datasets [Kusner et al., , Huang et al., ] and single-cell RNA-sequencing (scRNA-seq) datasets [Bellazzi et al., , Luecken and Theis, ], which admit this characteristic.
To this end, we focus on the TWD between the samples $\mathtt{TW}(\mathbf{x}_{i},\mathbf{x}_{i^{\prime}},T)$, where the samples are represented as discrete histograms $\{\mathbf{x}_{i}\}_{i=1}^{n}$ by normalization to unit $L_{1}$ norm.
However, often in practice, the tree $T$ is not explicitly known, and we are given solely the observational data matrix $\mathbf{X}$ (i.e., the node attributes from the hidden $T$), where the features (columns) are presumed to have a latent hierarchy.
Therefore, our goal is to approximate the TWD $\mathtt{TW}(\mathbf{x}_{i},\mathbf{x}_{i^{\prime}},T)$ without access to the latent tree $T$ underlying the high-dimensional features.

### 4.2 Hyperbolic Diffusion Embedding

In our approach, we first construct a hyperbolic representation following Lin et al. []. We briefly describe the steps of this construction below.
The feature diffusion operator $\mathbf{P}\in\mathbb{R}^{m\times m}$ is built according to Eq. (2), where $d(j,j^{\prime})$ is an initial distance between the features embedded in the ambient space $\mathbb{R}^{n}$.
Then, we construct the multi-scale diffusion densities $\bm{\phi}^{k}_{j}=\mathbf{P}^{2^{-k}}\mathbf{e}_{j}\in\mathbb{R}^{m}$ for each feature $j\in[m]$, considering a series of diffusion time steps on a dyadic grid $t=2^{-k}$ for $k\in\mathbb{Z}_{0}^{+}$.
These densities give rise to a multi-scale view of the features on the dyadic grid.
We then embed each feature $j$ at the $k$-th scale by $(j,k)\mapsto\mathbf{z}_{j}^{k}=[(\bm{\psi}_{j}^{k})^{\top},2^{k/2-2}]^{\top}
\in\mathbb{H}^{m+1}$ based on the exponential growth of the Poincaré half-space, where $\bm{\psi}^{k}_{j}=\displaystyle\textstyle\sqrt{\bm{\phi}^{k}_{j}}\in\mathbb{R}
^{m}$ is the element-wise square root of the diffusion density $\bm{\phi}^{k}_{j}$.
The multi-scale embedding of the corresponding diffusion operators in hyperbolic space [Lin et al., ] results in distances that are exponentially scaled. This scaling aligns with the structure of a tree distance, effectively establishing a natural hierarchical relation between different diffusion timescales. While it is well-established that diffusion geometry effectively recovers the underlying manifold [Coifman and Lafon, ], recent research by Lin et al. [] builds on this by showing that propagated densities, when applied with diffusion times on dyadic grids, can reveal hidden hierarchical metric. This method effectively captures local information from exponentially expanding neighborhoods around each point.
We note that other hyperbolic embeddings (e.g., [Chamberlain et al., , Nickel and Kiela, ]) could potentially be considered.
However, these methods usually assume a (partially) known graph $T$ instead of high-dimensional observational data $\mathbf{X}$, and therefore, we do not opt to use them (see App. E.3.4 for more details).

Figure: Figure 1: (a) The LCA relation of $j_{1},j_{2},j_{3}$ on tree $B$. (b) The hyperbolic LCA relation $[j_{2}\vee j_{3}\sim j_{1}]_{\mathbb{H}^{m+1}}^{k}$ for $\mathbf{z}_{j_{1}}^{k},\mathbf{z}_{j_{2}}^{k},\mathbf{z}_{j_{3}}^{k}\in\mathbb {H}^{m+1}$ along with their geodesic paths (red semi-circles). Blue points indicate the hyperbolic LCAs, and yellow points represent their orthogonal projections. The value $\mathtt{proj}(\cdot)$ signifies the parent-child relation. (c) The HD-LCA $\overline{\mathbf{h}}_{j_{1},j_{2}}$ defined as the Riemannian mean of the orthogonal projections $\{\mathbf{o}_{j_{1},j_{2}}^{k}\}_{k=0}^{K_{c}}$, incorporating the multi-scale hyperbolic LCAs.
Refer to caption: x1.png

### 4.3 Hyperbolic Diffusion LCA

To reveal the latent hierarchy of the features based on the continuous hyperbolic embedding, we first establish the analogy between geodesic paths in hyperbolic space $\mathbb{H}^{m+1}$ and paths on trees. We define the following hyperbolic lowest common ancestor (LCA).

###### Definition 4.1 (Hyperbolic LCA) .

The hyperbolic LCA of any two points $\mathbf{z}_{j}^{k},\mathbf{z}_{j^{\prime}}^{k}\in\mathbb{H}^{m+1}$ at the $k$-th scale is defined by their Fréchet mean [Fréchet, ] as $\mathbf{z}_{j}^{k}\vee\mathbf{z}_{j^{\prime}}^{k}\coloneqq\underset{\mathbf{z}
\in\mathbb{H}^{m+1}}{\arg\min}\left(d_{\mathbb{H}^{m+1}}^{2}(\mathbf{z},
\mathbf{z}_{j}^{k})+d_{\mathbb{H}^{m+1}}^{2}(\mathbf{z},\mathbf{z}_{j^{\prime}
}^{k})\right)$.

###### Proposition 4.1 .

The hyperbolic LCA $\mathbf{z}_{j}^{k}\vee\mathbf{z}_{j^{\prime}}^{k}$ in Def. 4.1 has a closed-form solution, given by

$$ $\mathbf{z}_{j}^{k}\vee\mathbf{z}_{j^{\prime}}^{k}=\left[\frac{1}{2}\left(\bm{ \psi}_{j}^{k}+\bm{\psi}_{j^{\prime}}^{k}\right)^{\top},\mathtt{proj}(\mathbf{z }_{j}^{k}\vee\mathbf{z}_{j^{\prime}}^{k})\right]^{\top},$ $$

where $\mathtt{proj}(\mathbf{z}_{j}^{k}\vee\mathbf{z}_{j^{\prime}}^{k})=\left\lVert
\left[\frac{1}{2}\left(\bm{\psi}_{j}^{k}-\bm{\psi}_{j^{\prime}}^{k}\right)^{
\top},2^{k/2-2}\right]^{\top}\right\rVert_{2}$.

The proof of Prop. 4.1 is in App. C.
Note that $\mathtt{proj}(\mathbf{z}_{j}^{k}\vee\mathbf{z}_{j^{\prime}}^{k})$ is the orthogonal projection of the hyperbolic LCA onto the $(m+1)$-axis $\mathbb{R}_{>0}$ in $\mathbb{H}^{m+1}$, which is the “radius” of the geodesic connecting $\mathbf{z}_{j}^{k}$ and $\mathbf{z}_{j^{\prime}}^{k}$, as depicted in Fig. 1(b).
We assert that this orthogonal projection can be used to identify the hierarchical relation at the $k$-th scale.
This argument is made formal in the following.

###### Definition 4.2 (Hyperbolic LCA Relation) .

Given any three points $\mathbf{z}_{j_{1}}^{k},\mathbf{z}_{j_{2}}^{k},\mathbf{z}_{j_{3}}^{k}\in\mathbb
{H}^{m+1}$ at the $k$-th scale, the relation $[j_{2}\vee j_{3}\sim j_{1}]_{\mathbb{H}^{m+1}}^{k}$ holds if
$\mathtt{proj}(\mathbf{z}_{j_{2}}^{k}\vee\mathbf{z}_{j_{3}}^{k})=\min\{\mathtt{
proj}(\mathbf{z}_{j_{1}}^{k}\vee\mathbf{z}_{j_{2}}^{k}),\mathtt{proj}(\mathbf{
z}_{j_{1}}^{k}\vee\mathbf{z}_{j_{3}}^{k}),\mathtt{proj}(\mathbf{z}_{j_{2}}^{k}
\vee\mathbf{z}_{j_{3}}^{k})\}$.

###### Proposition 4.2 .

For any $k_{1}\leq k_{2}$, $2^{-(k_{2}-k_{1})}\simeq d_{\mathbb{H}^{m+1}}(\mathbf{z}_{j}^{k_{2}}\vee
\mathbf{z}_{j^{\prime}}^{k_{2}},\mathbf{z}_{j}^{k_{2}})/d_{\mathbb{H}^{m+1}}(
\mathbf{z}_{j}^{k_{1}}\vee\mathbf{z}_{j^{\prime}}^{k_{1}},\mathbf{z}_{j}^{k_{1
}})$.

The proof of Prop. 4.2 is in App. C.
Prop. 4.2 implies that the hyperbolic LCAs at different scales exhibit exponential growths, which is analogous to the exponential growth of the distance between parent and child nodes in trees.

We now define the Hyperbolic Diffusion LCA (HD-LCA) that incorporates the multi-scale hyperbolic LCAs by jointly considering the dyadic diffusion times $\{2^{-k}\}_{k=0}^{K_{c}}$, where $K_{c}\in\mathbb{Z}_{0}^{+}$.

###### Definition 4.3 (HD-LCA) .

The HD-LCA of features $j,j^{\prime}\in[m]$ with $K_{c}$ scales is defined by the Fréchet mean $\overline{\mathbf{h}}_{j,j^{\prime}}\coloneqq\underset{\mathbf{h}\in\mathbb{H}
^{m+1}}{\arg\min}\;\sum_{k=0}^{K_{c}}\;d_{\mathbb{H}^{m+1}}^{2}(\mathbf{h},
\mathbf{o}^{k}_{j,j^{\prime}})$, where $\mathbf{o}^{k}_{j,j^{\prime}}=[\mathbf{0}^{\top},\mathtt{proj}(\mathbf{z}_{j}^
{k}\vee\mathbf{z}_{j^{\prime}}^{k})]^{\top}\in\mathbb{H}^{m+1}$.

###### Proposition 4.3 .

The HD-LCA $\overline{\mathbf{h}}_{j,j^{\prime}}$ in Def. 4.3 has a closed-form solution, given by $\overline{\mathbf{h}}_{j,j^{\prime}}=\left[0,\ldots,0,a_{j,j^{\prime}}\right]^
{\top}$, where $a_{j,j^{\prime}}=\sqrt[(K_{c}+1)]{\mathtt{proj}(\mathbf{z}_{j}^{0}\vee\mathbf{
z}_{j^{\prime}}^{0})\cdots\mathtt{proj}(\mathbf{z}_{j}^{K_{c}}\vee\mathbf{z}_{
j^{\prime}}^{K_{c}})}$.

The proof of Prop. 4.3 is in App. C.
Note that the HD-LCA in Def. 4.3 involves two applications of the geodesic paths in $\mathbb{H}^{m+1}$:
(i) to establish the hyperbolic LCA at each scale $k$ by connecting two points on the same plane in $\mathbb{R}^{m+1}$, and (ii) to merge these LCAs across multiple levels along the $(m+1)$-axis $\mathbb{R}_{>0}$ in $\mathbb{H}^{m+1}$, as in Fig. 1(c).

Incorporating the multi-scale of the hyperbolic LCAs arranged on a dyadic grid via the HD-LCA facilitates the recovery of the underlying tree relations.
A similar approach was presented in Lin et al. [], where it was shown that the distance between the embedding $j\mapsto[(\mathbf{z}_{j}^{0})^{\top},\ldots,(\mathbf{z}_{j}^{K_{c}})^{\top}]^{
\top}\in\mathcal{M}=\mathbb{H}^{m+1}\times\ldots\times\mathbb{H}^{m+1}$ in a product manifold of Poincaré half-spaces, given by $d_{\mathcal{M}}(j,j^{\prime})=\sum_{k=0}^{K_{c}}~{}2\sinh^{-1}(2^{-k/2+1}\left
\lVert\mathbf{z}^{k}_{j}-\mathbf{z}^{k}_{j^{\prime}}\right\rVert_{2})$, recovers $d_{T}(j,j^{\prime})$.
However, [Lin et al., ] considers a latent hierarchical structure of the samples.
It is fundamentally different than the problem we consider in Sec. 4.1, as it does not learn the latent hierarchy of the features explicitly by constructing a tree, and then incorporate it in a meaningful distance between the samples as we do here (see App. F.4 for more details).
As shown in App. F.4, this distance is less effective for document and cell classifications.

Note that the closed-form value $a_{j,j^{\prime}}$ in Prop. 4.3 can be interpreted as the Riemannian average of the multi-scale hyperbolic LCA relations in Def. 4.1, resulting in the following HD-LCA relation.

###### Definition 4.4 (HD-LCA Relation) .

Given any three features $j_{1},j_{2},j_{3}\in[m]$, the relation $[j_{2}\vee j_{3}\sim j_{1}]_{\mathcal{M}}$ holds if $a_{j_{2},j_{3}}=\min\{a_{j_{1},j_{2}},a_{j_{1},j_{3}},a_{j_{2},j_{3}}\}$.

Figure: Algorithm 1 HD Binary Tree Decoding

Figure: Algorithm 2 TWD with a latent feature hierarchy

###### Definition 4.1 (Hyperbolic LCA) .

###### Proposition 4.1 .

###### Definition 4.2 (Hyperbolic LCA Relation) .

###### Proposition 4.2 .

###### Definition 4.3 (HD-LCA) .

###### Proposition 4.3 .

###### Definition 4.4 (HD-LCA Relation) .

### 4.4 Binary Tree Construction

Equipped with the definitions and properties above, we now present the proposed tree construction in hyperbolic spaces.
As any tree can be transformed into a binary tree [Bowditch, , Cormen et al., ], we suggest to construct a rooted binary tree $B$ through an iterative merging process.
We begin by assigning the features to the leaf nodes because it was shown that the leaves alone contain adequate information to fully reconstruct the tree [Sarkar, ] (see App. F.2 for more details).
This is not only efficient but also aligns with established TWD [Indyk and Thaper, , Backurs et al., , Le et al., ].
Then, pairs of features are merged iteratively by their HD-LCA (Def. 4.3).
Specifically, the pair of features that exhibits the highest similarity in the HD-LCA (i.e., the smallest value of $a_{j,j^{\prime}}$ in Prop. 4.3) is merged.
Note that only one of the HD-LCA relations $[j_{2}\vee j_{3}\sim j_{1}]_{\mathcal{M}}$, $[j_{1}\vee j_{3}\sim j_{2}]_{\mathcal{M}}$, or $[j_{1}\vee j_{2}\sim j_{3}]_{\mathcal{M}}$ holds in $B$.
Once the nodes are merged, the weight of the edge is assigned with the induced $\ell_{1}$ distance between the respective embeddings.
This iterative process constructs the binary tree $B$ in a bottom-up manner, similar to single linkage clustering [Gower and Ross, ], with the measure of similarity determined by their HD-LCAs and edge weights assigned with the geodesic distance.
We summarize the proposed HD binary tree decoding in Alg. 1.

We note that the distance in the embedding [Lin et al., ] could be used in existing graph-based algorithms [Alon et al., ] or linkage methods [Murtagh and Contreras, ] instead of the HD-LCA in the binary tree decoding (Alg. 1).
However, these generic algorithms, while providing upper bounds on the distortion of the constructed tree distances, do not build a geometrically meaningful tree when given a tree metric [Sonthalia and Gilbert, ], because they do not exploit the hyperbolic geometry of the embedded space.
Indeed, as empirically shown in App. E, constructing a feature tree using existing TWDs, graph-based, or linkage approaches with embedding distance is less effective.
This may suggest that the information encoded in the multi-scale hyperbolic LCAs is critical for revealing the latent feature tree.
In addition, we argue that having a closed-form solution to the basic element of the tree construction HD-LCA in Prop. 4.3 is an important advantage because it directly enables us to present an effective and interpretable algorithm (Alg. 1) for recovering the underlying tree relationships from the learned hyperbolic embedding.

###### Proposition 4.4 .

Let $r$ be the root node of $B$ obtained in Alg. 1 and $\mathbf{z}_{r}$ be the corresponding embedding in $\mathcal{M}$. The HD-LCA depth is equivalent to the Gromov product [Gromov, ], i.e., $d_{\mathcal{M}}(\mathbf{z}_{r},\overline{\mathbf{h}}_{j,j^{\prime}})\simeq
\langle j,j^{\prime}\rangle_{r}$, where $\langle j,j^{\prime}\rangle_{r}=\frac{1}{2}(d_{B}(j,r)+d_{B}(j^{\prime},r)-d_{
B}(j,j^{\prime}))$.

The proof of Prop. 4.4 is in App. C.
Prop. 4.4 shows that the HD-LCA depth in the continuous embedding space is the tree depth in the decoded tree $B$.

###### Theorem 4.1 .

For sufficiently large $K_{c}$ and $m$, the decoded tree metric $d_{B}$ in Alg. 1 is bilipschitz equivalent to the underlying tree metric $d_{T}$, where $T$ is the ground latent tree assumed from Sec. 4.1 and $d_{T}$ is the hidden tree metric between features defined as the length of the shortest path on $T$. Specifically, the bilipschitz constants can be determined by $\frac{1}{\sqrt{2}}d_{B}(j_{1},j_{2})\leq d_{T}(j_{1},j_{2})\leq d_{B}(j_{1},j_
{2})$ for all features $j_{1},j_{2}\in[m]$.

The proof of Thm. 4.1 is in App. C.
It is important to recognize that the validity of Thm. 4.1 relies on the construction of the diffusion operator $\mathbf{P}$ from the observations $\mathbf{X}$, which is tightly connected to the heat kernel and the Laplacian of the underlying hierarchy $T$ (see Sec. 3).
While $T$ is hidden, the diffusion operator $\mathbf{P}$ is used in our method to construct another tree, $B$, such that, the tree metrics of the data-driven tree $d_{B}$ and the hidden tree $d_{T}$ are equivalent. This result is stated in Thm. 4.1.
In addition, our $\mathsf{HD}$ LCA for decoded tree $B$ is a continuous LCA *in the product manifold of $(m+1)$-dimensional Poincaré half-spaces* that recovers hidden tree relations.
It is closely related to the hyperbolic LCA in Chami et al. [], where a continuous LCA in a *single* two-dimensional Poincaré disk was proposed for hierarchical clustering, and a tree was constructed by minimizing the Dasgupta’s cost [Dasgupta, ].
However, their method is not applicable to approximate $\mathtt{TW}(\mathbf{x}_{i},\mathbf{x}_{i^{\prime}},T)$ since it does not recover the latent feature hierarchy, as we empirically show in Sec. 5.

###### Proposition 4.4 .

###### Theorem 4.1 .

### 4.5 Tree-Wasserstein Distance for High-Dimensional Data with a Latent Feature Hierarchy

Finally, with the decoded tree $B$ at hand, we define a new TWD between high-dimensional samples with a latent feature hierarchy as follows.

###### Definition 4.5 (TWD for High-Dimensional Data with a Latent Feature Hierarchy.) .

The TWD for high-dimensional samples with a latent feature hierarchy is defined as

$$ $\mathtt{TW}(\mathbf{x}_{i},\mathbf{x}_{i^{\prime}},B)\coloneqq\sum_{v\in \widetilde{V}}\alpha_{v}\left\lvert\sum_{u\in\Gamma_{B}(v)}\left(\mathbf{x}_{i }(u)-\mathbf{x}_{i^{\prime}}(u)\right)\right\rvert,$ (3) $$

where $\alpha_{v}=d_{B}(v,v^{\prime})$, and node $v^{\prime}$ is the parent of node $v$.

###### Theorem 4.2 .

For sufficiently large $K_{c}$ and $m$, the TWD defined in Eq. (3) is bilipschitz equivalent to the TWD with the ground latent tree $T$. Specifically, the bilipschitz constants can be determined by $\frac{1}{\sqrt{2}}\mathtt{TW}(\mathbf{x}_{i},\mathbf{x}_{i^{\prime}},B)\leq
\mathtt{TW}(\mathbf{x}_{i},\mathbf{x}_{i^{\prime}},T)\leq\mathtt{TW}(\mathbf{x
}_{i},\mathbf{x}_{i^{\prime}},B)$ for all samples $i,i^{\prime}\in[n]$.

The proof of Thm. 4.2 is in App. C.
Thm. 4.2 shows that our TWD in Eq. (3) recovers the TWD associated with the latent feature tree $T$ without access to $T$.
This could be viewed as an unsupervised Wasserstein metric learning between data samples with the latent feature hierarchy.
In addition, our intermediate result in Thm. 4.1 is noteworthy, as it represents a form of unsupervised ground metric learning [Cuturi and Avis, ] derived from the latent tree underlying the features.

We summarize the computation of our method.
Notably, the input to the algorithm is the observational data matrix $\mathbf{X}$ alone (i.e., the tree node attributes), and no prior knowledge about the latent feature tree graph $T$ is given.
First, we build the feature diffusion operator $\mathbf{P}$ by Eq. (2).
Then, for each scale $k$, we embed the features into $\mathbb{H}^{m+1}$ using the $k$-th diffusion densities.
Collecting the embedding of all the features from all the scales, we decode a binary tree by Alg. 1, based on the HD-LCA in Def. 4.3.
Finally, we compute the proposed TWD as in Eq. (3).
We present the pseudo-code in Alg. 2.

Seemingly, the need to build a distance matrix and apply the eigendecomposition to the resulting kernel limits the applicability to a small number of features.
However, using [Shen and Wu, ], the proposed TWD can be computed in $O(m^{1.2})$ compared to the naïve implementation that requires $O(mn^{3}+m^{3}\log m)$ (see App. D).
Indeed, in Sec. 5, we show applications with thousands of features.

One limitation of our TWD in Eq. (3) is the restriction to samples from the probability simplex.
This requirement, which stems from the definitions of OT framework [Monge, , Kantorovich, ] and the Wasserstein and TWDs [Villani, , Cuturi, , Indyk and Thaper, , Backurs et al., , Le et al., , Takezawa et al., , Yamada et al., ], restricts its applicability.
However, a large body of work [Villani, , Cuturi, , Courty et al., , Solomon et al., , Yurochkin et al., , Alvarez-Melis and Jaakkola, ] illustrates that this setting is relevant in a broad range of applications.
In addition, this limitation does not affect the tree decoding in Alg. 1, as it can recover the latent feature hierarchy from $\mathbb{R}^{n}$, which is in line with existing tree construction methods [Alon et al., , Murtagh and Contreras, ].
Note that the problem of finding distances between samples with feature hierarchy and possibly with negative entries has been addressed in Gavish et al. [], Ankenman [], Mishne et al. [, ] using different approaches.

A notable property of our TWD in Eq. (3) is that it can be interpreted using the tree-sliced Wasserstein distance (TSWD) [Le et al., ].
Computing the TSWD consists of constructing multiple random trees of different depths from $\mathbf{X}$.
Then, the TSWD is given by the average of the TWDs corresponding to these trees.
Here, considering only a single scale $k$ in Alg. 1 (as detailed in App. F.1) yields the decoded tree $B^{k}$.
Similarly to the TSWD, we show that our TWD can be recast as the sum of the TWDs corresponding to $B^{k}$ at multiple scales $k=0,\ldots,K_{c}$.

###### Proposition 4.5 .

The TWD in Eq. (3) is bilipschitz equivalent to the sum of the multi-scale hyperbolic TWDs $\mathtt{TW}(\mathbf{x}_{i},\mathbf{x}_{i^{\prime}},B^{k})$ for $k=0,1,\ldots,K_{c}$, i.e., $\mathtt{TW}(\mathbf{x}_{i},\mathbf{x}_{i^{\prime}},B)\simeq\sum_{k=0}^{K_{c}}
\mathtt{TW}(\mathbf{x}_{i},\mathbf{x}_{i^{\prime}},B^{k})$.

The proof of Prop. 4.5 is in App. C.
Prop. 4.5 implies that our TWD captures the (tree-sliced) information in the multi-scale trees $B^{k}$ efficiently as it circumvents their explicit construction.
In addition, we show in Sec. 5 that the our TWD outperforms the TSWD [Le et al., ].

We conclude with a few remarks.
First, it is important to note that our primary objective is not to recover the exact hidden $T$, but to effectively approximate the TWD $\mathtt{TW}(\mathbf{x}_{i},\mathbf{x}_{i^{\prime}},T)$ from the data matrix $\mathbf{X}$ (i.e., the node attributes from the hidden $T$), which we achieve by the proposed TWD $\mathtt{TW}(\mathbf{x}_{i},\mathbf{x}_{i^{\prime}},B)$ in Eq. (3) as stated in Thm. 4.2.
Second, our TWD can also be used to approximate the OT distance in linear time.
For example, suppose we are interested in computing the OT distance based on a given ground tree metric between the features.
While the complexity of this computation is typically $O(m^{3}\log m)$ [Pele and Werman, , Bonneel et al., ], we can approximate it by the proposed TWD in linear time, as we demonstrate in the runtime analysis in App. E.
Finally, our TWD is differentiable, as detailed in App. F.3, in contrast to existing TWD methods.
This differentiability equips the proposed TWD with a continuous gradient, significantly enhancing its applicability in gradient-based optimization, a feature that is not typically found in TWD baselines.

###### Definition 4.5 (TWD for High-Dimensional Data with a Latent Feature Hierarchy.) .

###### Theorem 4.2 .

###### Proposition 4.5 .

## 5 Experimental Results

We apply the proposed method to a synthetic example and to word-document and single-cell gene expression datasets.
Additional implementation details are in App. D, including the hyperparameters and scaling to large datasets.
Additional experiments, including runtime analysis, ablation studies, the impact of curvature, and robustness analysis under Gaussian noise, are in App. E.

### 5.1 Synthetic Word-Document Data Example

Figure: Figure 2: (a) Illustration of the probabilistic hierarchical model for generating synthetic samples consisting of 8 binary elements. The orange nodes represent (sub)categories, and the black nodes present produce items. The edge weights represent the probabilities. (b) Feature trees constructed by our TWD and competing baselines. Nodes corresponding to fruits are colored in blue, those representing vegetables are in red, and the internal nodes are in green.
Refer to caption: x2.png

We generate a synthetic dataset with a latent feature hierarchy as follows.
The generated dataset is designed to resemble a word-document data matrix $\mathbf{X}$, whose rows (samples) represent documents and columns (features) represent word presence.
Our dataset consists of 100 documents and eight words of produce items: apple, orange, banana, carrot, beetroot, kale, spinach, and lettuce.
These items are divided into categories and sub-categories as depicted in Fig. 2(a), giving rise to the latent feature hierarchy.
The 100 documents are realizations of a random binary vector created according to the probabilistic hierarchical model in Fig. 2(a).
For details of the data generation, see App. D.

We compare the proposed TWD with existing TWDs, including (i) Quadtree [Indyk and Thaper, ] that constructs a random tree by a recursive division of hypercubes, (ii) Flowtree [Backurs et al., ] that computes the optimal flow of Quadtree, (iii) tree-sliced Wasserstein distance (TSWD) [Le et al., ] that computes an average of TWDs by random sampling trees, where the numbers of sampling are set to 1, 5, and 10, (iv) UltraTree [Chen et al., ] that constructs an ultrametric tree [Johnson, ] by minimizing OT regression cost, (v) weighted cluster TWD (WCTWD) [Yamada et al., ], (vi) weighted Quadtree TWD (WQTWD) [Yamada et al., ], and their sliced variants (vii) SWCTWD and (viii) SWQTWD [Yamada et al., ].
We also compare to TWDs where the trees are constructed by (i) MST [Prim, ], (ii) Tree Representation (TR) [Sonthalia and Gilbert, ] through a divide-and-conquer approach, (iii) gradient-based hierarchical clustering (HC) in hyperbolic space (gHHC) [Monath et al., ], (iv) gradient-based Ultrametric Fitting (UltraFit) [Chierchia and Perret, ], and (v) HC by hyperbolic Dasgupta’s cost (HHC) [Chami et al., ].
We refer to App. B for an extended description of these baselines.

Fig. 2(b) presents the trees constructed using different methods.
Our TWD constructs a binary tree that accurately represents the latent feature hierarchy, demonstrating a clear advantage over the trees obtained by the baselines.
Further evaluation through a binary document classification task using $k$-nearest neighbors ($k$NN), where the documents are labeled by the presence of fruits (with probability $0.5$, creating two balanced classes), is presented in Tab. LABEL:tab:classification.
The depicted classification accuracy is obtained by averaging over five runs, where in each run, the dataset is randomly split into $70\%$ training and $30\%$ testing sets.
The same random splits are used across all comparisons.
We see that our TWD demonstrates a significant improvement over the competing TWDs.

### 5.2 Real-World Word-Document Data Analysis

We demonstrate the power of the proposed TWD in real-world document classification tasks.
The latent hierarchical structure of words (features) can be viewed as a node branching into related terms [Gollapalli and Caragea, , Yurochkin et al., , Chang and Blei, ], allowing for meaningful comparison between documents in OT distance [Kusner et al., ], as it measures how words and their associated meanings are transported from one document to another.

We test four word-document benchmarks in Kusner et al. [], including BBCSPORT, TWITTER, CLASSIC, and AMAZON.
Detailed descriptions of these datasets are reported in App. D.
We compare the proposed TWD to the same competing methods as in Sec. 5.1.
In addition, we include the Word Mover’s Distance (WMD) [Kusner et al., ] that computes the OT distance between documents using Word2Vec [Mikolov et al., ].
A dissimilarity classification based on our TWD and the baselines using the $k$NN classifier is applied.
We use cross-validation with five trials, where the dataset is randomly divided into $70\%$ training set and $30\%$ testing set, following previous work in word-document data analysis [Kusner et al., , Huang et al., ].

Tab. LABEL:tab:classification presents the document classification accuracy.
The results of Quadtree, Flowtree, TSWD, and WMD are taken from Huang et al. [], Takezawa et al. [].
Our method outperforms all the competing baselines on BBCSPORT, TWITTER, and AMAZON datasets, including the pre-trained WMD.
On CLASSIC dataset, our proposed TWD yields the second-best classification accuracy and is competitive with WMD.

**Table 1: Document and cell classification accuracy (the highest in bold and the second highest underlined). Results marked with ^∗ are taken from Huang et al. [], Takezawa et al. [].**
|  |  | Synthetic |  | Real-World Word-Document |  | Real-World scRNA-seq |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  |  |  | BBCSPORT | TWITTER | CLASSIC | AMAZON |  | Zeisel | CBMC |  |
| Quadtree |  | 98.2$\pm$1.2 |  | 95.5$\pm$0.5^∗ | 69.6$\pm$0.8^∗ | 95.9$\pm$0.4^∗ | 89.3$\pm$0.3^∗ |  | 80.1$\pm$1.2 | 80.6$\pm$0.6 |
| Flowtree |  | 97.3$\pm$0.9 |  | 95.3$\pm$1.1^∗ | 70.2$\pm$0.9^∗ | 94.4$\pm$0.6^∗ | 90.1$\pm$0.3^∗ |  | 81.7$\pm$0.9 | 81.8$\pm$0.9 |
| WCTWD |  | 91.4$\pm$1.2 |  | 92.6$\pm$2.1 | 69.1$\pm$2.6 | 93.7$\pm$2.9 | 88.2$\pm$1.4 |  | 81.3$\pm$4.9 | 78.4$\pm$3.3 |
| WQTWD |  | 92.8$\pm$1.8 |  | 94.3$\pm$1.7 | 69.4$\pm$2.4 | 94.6$\pm$3.2 | 87.4$\pm$1.8 |  | 80.9$\pm$3.5 | 79.1$\pm$3.0 |
| UltraTree |  | 93.2$\pm$3.1 |  | 93.1$\pm$1.5 | 68.1$\pm$3.2 | 92.3$\pm$1.9 | 86.2$\pm$3.1 |  | 83.9$\pm$1.6 | 82.3$\pm$2.6 |
| TSWD-1 |  | 90.9$\pm$1.3 |  | 87.6$\pm$1.9^∗ | 69.8$\pm$1.3^∗ | 94.5$\pm$0.5^∗ | 85.5$\pm$0.6^∗ |  | 79.6$\pm$1.8 | 72.6$\pm$1.8 |
| TSWD-5 |  | 92.4$\pm$1.1 |  | 88.1$\pm$1.3^∗ | 70.5$\pm$1.1^∗ | 95.9$\pm$0.4^∗ | 90.8$\pm$0.1^∗ |  | 81.3$\pm$1.4 | 74.9$\pm$1.1 |
| TSWD-10 |  | 94.2$\pm$0.9 |  | 88.6$\pm$0.9^∗ | 70.7$\pm$1.3^∗ | 95.9$\pm$0.6^∗ | 91.1$\pm$0.5^∗ |  | 83.2$\pm$0.8 | 76.5$\pm$0.7 |
| SWCTWD |  | 94.9$\pm$2.3 |  | 92.8$\pm$1.2 | 70.2$\pm$1.2 | 94.1$\pm$1.8 | 90.2$\pm$1.2 |  | 81.9$\pm$3.1 | 78.3$\pm$1.7 |
| SWQTWD |  | 95.1$\pm$1.8 |  | 94.5$\pm$1.0 | 70.6$\pm$1.9 | 95.4$\pm$2.0 | 89.8$\pm$1.1 |  | 80.7$\pm$2.5 | 79.8$\pm$2.5 |
| MST-TW |  | 88.1$\pm$2.7 |  | 88.4$\pm$1.9 | 68.2$\pm$1.9 | 90.0$\pm$3.1 | 86.4$\pm$1.2 |  | 80.1$\pm$3.1 | 76.2$\pm$2.5 |
| TR-TW |  | 93.9$\pm$0.7 |  | 89.2$\pm$0.9 | 70.2$\pm$0.7 | 92.9$\pm$0.8 | 88.7$\pm$1.1 |  | 80.3$\pm$0.7 | 78.4$\pm$1.2 |
| HHC-TW |  | 94.1$\pm$1.7 |  | 85.3$\pm$1.8 | 70.4$\pm$0.4 | 93.4$\pm$0.8 | 88.5$\pm$0.7 |  | 82.3$\pm$0.7 | 77.3$\pm$1.1 |
| gHHC-TW |  | 93.5$\pm$1.9 |  | 83.2$\pm$2.4 | 69.9$\pm$1.8 | 90.3$\pm$2.2 | 86.9$\pm$2.0 |  | 79.4$\pm$1.9 | 73.6$\pm$1.6 |
| UltraFit-TW |  | 94.3$\pm$1.2 |  | 84.9$\pm$1.4 | 69.5$\pm$1.2 | 91.6$\pm$0.9 | 87.4$\pm$1.6 |  | 81.9$\pm$3.3 | 77.8$\pm$1.2 |
| WMD |  | - |  | 95.4$\pm$0.7^∗ | 71.3$\pm$0.6^∗ | 97.2$\pm$0.1^∗ | 92.6$\pm$0.3^∗ |  | - | - |
| GMD |  | - |  | - | - | - | - |  | 84.2$\pm$0.7 | 81.4$\pm$0.7 |
| Ours |  | 99.8$\pm$0.1 |  | 96.1$\pm$0.4 | 73.4$\pm$0.2 | 96.9$\pm$0.2 | 93.1$\pm$0.4 |  | 89.1$\pm$0.4 | 84.3$\pm$0.3 |

### 5.3 Single-Cell Gene Expression Data Analysis

We further demonstrate the advantages of our TWD in cell classification tasks on scRNA-seq data.
The hierarchical structure of genes in scRNA-seq data [Tanay and Regev, ] are essential for analyzing complex cellular processes [Ding and Regev, ], especially for computing OT distance [Bellazzi et al., , Tran et al., , Bunne et al., ].
For instance, in developmental biology, genes associated with stem cell differentiation follow a hierarchical pattern [Kang et al., ].
Revealing the latent gene hierarchy is essential for accurately computing the TWD between cells, as it captures the underlying biological relationships and functional similarities.

We test two scRNA-seq datasets in Dumitrascu et al. [], including the Zeisel and CBMC.
Detailed descriptions of the datasets are reported in App. D.
We compare the proposed TWD to the same baselines as in Sec. 5.2.
Instead of WMD, we include the Gene Mover’s Distance (GMD) [Bellazzi et al., ], a gene-based OT distance using Gene2Vec [Du et al., ].
The cell classification tasks are performed in the same way as in Sec. 5.1 and Sec. 5.2.
Tab. LABEL:tab:classification shows the mean and the standard deviation of the cell classification accuracy.
Our TWD outperforms the competing methods by a large margin.
This suggests that the proposed TWD effectively transforms the latent tree-like relationships in gene space, improving the cell classification accuracy compared to random sampling trees and trees constructed by graph-based and HC methods.
Notably, our TWD differentiates the cell types better than pre-trained GMD, which uses gene proximity, while ours is geometrically intrinsic.

## 6 Conclusions

We presented a novel distance between samples with a latent feature hierarchy.
The computation of the distance is efficient and can accommodate comparisons of large datasets.
We show experimental results on various applications, where our method outperforms competing baselines.
Theoretically, the proposed TWD provably recovers the TWD for data samples with a latent feature hierarchy without prior knowledge about the hidden feature hierarchy and based solely on the observational data matrix.
We note that the construction of the proposed distance consists of a new algorithm for decoding a binary tree based on hyperbolic embeddings, relying on the analogy between the exponential growth of the metric in high-dimensional hyperbolic spaces and the exponential expansion of trees.
This tree decoding is general and can be used for decoding trees in broader contexts.

Our work is conceptually different from existing TWD works in two main aspects.
First, we present an unconventional use of the TWD. In addition to speeding up the computation of the Wasserstein distance, we also use the tree to represent the feature hierarchy.
Second, we consider and recover the latent feature hierarchy by a novel hyperbolic tree decoding algorithm, which differs from existing hierarchical (continuous) representation learning methods that consider (known) sample hierarchy.

## Acknowledgments

The work of YEL and RT was supported by the European Union’s Horizon 2020 research and innovation programme under grant agreement No. 802735-ERC-DIFFOP.

## Appendix A Additional Background

While we have covered the most relevant background in the main paper, here we present additional background on the Wasserstein distance, diffusion geometry, hyperbolic geometry, and trees.

### A.1 Wasserstein Distance

The application of the Wasserstein distance spans various fields, including enhancing generative models in machine learning [Arjovsky et al., , Kolouri et al., ], aligning and comparing color distributions in computer graphics [Solomon et al., , Lavenant et al., ], and serving as a robust similarity metric in data analysis [Bellazzi et al., , Kusner et al., , Yair et al., ].

Let $\mu_{1}$ and $\mu_{2}$ be two probability distributions on a measurable space $\Omega$ with a ground metric $d$, $\Pi(\mu_{1},\mu_{2})$ be the set of couplings $\pi$ on the product space $\Omega\times\Omega$ such that $\pi(B_{1}\times\Omega)=\mu_{1}(B_{1})$ and $\pi(\Omega\times B_{2})=\mu_{2}(B_{2})$ for any sets $B_{1},B_{2}\subset\Omega$.
The Wasserstein distance [Villani, ] between $\mu_{1}$ and $\mu_{2}$ is defined by

$$ $\mathtt{OT}(\mu_{1},\mu_{2},d)\coloneqq\underset{\pi\in\Pi(\mu_{1},\mu_{2})}{ \inf}\int_{\Omega\times\Omega}d(x,y)\text{d}\pi(x,y).$ (4) $$

When replacing the $\inf$ in Eq. (4) with the $\arg\min$, the obtained solution $\pi$ is referred to as the Optimal Transport (OT) plan [Monge, , Kantorovich, ].
When computing the Wasserstein distance in Eq. (4) between two discrete probabilities with $m$ points, given by

$$ $\mathtt{OT}(\bm{\mu}_{1},\bm{\mu}_{2},\mathbf{C})\coloneqq\underset{\mathbf{ \Phi}\in\mathbb{R}^{m\times m}}{\min}\langle\mathbf{\Phi},\mathbf{C}\rangle \text{ s.t.}\begin{cases}\mathbf{\Phi 1}_{m}=\bm{\mu}_{1},\\ \mathbf{\Phi}^{\top}\mathbf{1}_{m}=\bm{\mu}_{2},\end{cases}$ (5) $$

where $\mathbf{C}\in\mathbb{R}^{m\times m}$ is the ground pairwise distance matrix between the coordinates of the discrete probabilities, it can be computed in $O(m^{3}\log m)$ [Peyré et al., ], preventing its use of OT for large-scale datasets [Bonneel et al., ].
Therefore, existing works have been proposed to reduce this complexity.
Among them, a notable work is the Sinkhorn algorithm [Cuturi, , Chizat et al., ], which approximates the OT in quadratic time.

### A.2 Diffusion Geometry

Diffusion geometry [Coifman and Lafon, ] is a mathematical framework designed to analyze high-dimensional data points $\{\mathbf{a}_{j}\in\mathbb{R}^{n}\}_{j=1}^{m}$ by effectively capturing their underlying geometric structures [Belkin and Niyogi, ].
It focuses on analyzing the similarity between data points through the process of diffusion propagation.
Specifically, the constructed diffusion operator $\mathbf{P}\in\mathbb{R}^{m\times m}$ derived from the observational data is associated with the heat kernel and the Laplacian of the underlying manifold [Belkin and Niyogi, ].
Let $\mathbf{Q}(j,j^{\prime})=\exp\left(-d^{2}(j,j^{\prime})/\epsilon\right)$ be the $(j,j^{\prime})$th element of the pairwise affinity matrix between the data points. In the limit $m\rightarrow\infty$ and $\epsilon\rightarrow 0$, the diffusion operator $\mathbf{P}$ converges to the heat kernel $\exp(-\Delta)$, where $\Delta$ is the Laplace–Beltrami operator on the manifold [Coifman and Lafon, ].
By applying eigendecomposition to the diffusion operator $\mathbf{P}$, we obtain a sequence of eigenpairs $\{\nu_{j},\bm{\varphi}_{j}\}_{j=1}^{m}$ with positive decreasing eigenvalues $1=\nu_{1}\geq\nu_{2}\geq\ldots\geq\nu_{m}>0$.
The diffusion maps (DM) in $m^{\prime}<<m$ dimensions at the diffusion time $t$ embeds a point $\mathbf{a}_{j}$ into $\mathbb{R}^{m^{\prime}}$ by

$$ $\bm{\Psi}_{m^{\prime}}:\mathbf{a}_{j}\mapsto[\nu_{2}^{t}\bm{\varphi}_{2},\nu_{ 3}^{t}\bm{\varphi}_{3},\ldots,\nu_{m^{\prime}+1}^{t}\bm{\varphi}_{m^{\prime}+1 }]^{\top}.$ (6) $$

It is shown that the Euclidean distances defined in DM between any two points $\left\lVert\bm{\Psi}_{m^{\prime}}(\mathbf{y}_{j})-\bm{\Psi}_{m^{\prime}}(
\mathbf{y}_{j^{\prime}})\right\rVert_{2}$ approximate the following diffusion distances

$$ $d_{\text{DM},t}(j,j^{\prime})=\left\lVert\mathbf{P}^{t}\mathbf{e}_{j}-\mathbf{ P}^{t}\mathbf{e}_{j^{\prime}}\right\rVert_{2}$ (7) $$

at the diffusion time $t$.

Diffusion geometry has mainly been utilized in manifold learning, leading to the development of multi-scale, low-dimensional representations and informative distances [Nadler et al., , Coifman and Maggioni, ].
Its utility has been demonstrated across various applications spanning various fields, including nonlinear stochastic dynamical systems [Shnitzer et al., ], multi-view dimensionality reduction [Lindenbaum et al., ], shape recognition [Bronstein et al., ], wavelet analysis [Coifman and Maggioni, ], and internal state values of long short-term memory [Kemeth et al., ].

### A.3 Hyperbolic Geometry

Hyperbolic geometry is a Riemannian manifold of a constant negative curvature [Lee, ].
There are four commonly used models for hyperbolic space: the Poincaré disk model, the Lorentz model, the Poincaré half-space model, and the Beltrami-Klein model.
These models are equivalent and isometric, representing the same geometry, and distances can be translated without distortion [Ratcliffe et al., ].
In this work, we used the Poincaré half-space model to represent the continuous tree structure underlying the data because the exponential growth of the metric serves as a natural representation of the exponential propagation of the diffusion densities on a dyadic grid.
Formally, the $n$-dimensional Poincaré half-space model $\mathbb{H}^{n}$ with constant negative curvature $-1$ is defined by $\mathbb{H}^{n}=\{\mathbf{x}\in\mathbb{R}^{n}\big{|}\mathbf{x}(n)>0\}$ with the Riemannian metric tensor $ds^{2}=(d\mathbf{x}^{2}(1)+d\mathbf{x}^{2}(2)+\ldots+d\mathbf{x}^{2}(n))/
\mathbf{x}^{2}(n)$.
The Riemannian distance for any $\bm{x},\bm{y}\in\mathbb{H}^{n}$ is defined by $d_{\mathbb{H}^{n}}(\mathbf{x},\mathbf{y})=2\sinh^{-1}(\left\lVert\mathbf{x}-
\mathbf{y}\right\rVert_{2}/(2\sqrt{\mathbf{x}(n)\mathbf{y}(n)}))$.
The distance function $d_{\mathbb{H}^{n}}(\cdot,\cdot)$ in the hyperbolic space $\mathbb{H}^{n}$ gives rise to two distinct types of geodesic paths [Ratcliffe et al., ]: straight lines that are parallel to the $n$-th axis, and semi-circles that are perpendicular to the boundary of the plane $\mathbb{R}^{n-1}$.

The Gromov product [Gromov, ] is used to define $\delta$-hyperbolic spaces in hyperbolic geometry.

###### Definition A.1 ( $\delta$ -Hyperbolic Space) .

A metric space $(\mathcal{X},d_{\mathcal{X}})$ is $\delta$-hyperbolic space [Gromov, ] if there exists $\delta\geq 0$ such that for any four points $x_{1},x_{2},x_{3},x_{4}\in\mathcal{X}$, we have

$$ $\displaystyle d_{\mathcal{X}}(x_{4},x_{1})+d_{\mathcal{X}}(x_{2},x_{3})\leq \max\{d_{\mathcal{X}}(x_{1},x_{2})+d_{\mathcal{X}}(x_{3},x_{4}),d_{\mathcal{X} }(x_{1},x_{3})+d_{\mathcal{X}}(x_{2},x_{4})\}+2\delta.$ (8) $$

###### Definition A.2 (Gromov Product [Gromov, 1987 ] ) .

Let $(\mathcal{X},d_{\mathcal{X}})$ be a metric space.
The Gromov product [Gromov, ] of any two points $x_{1},x_{2}\in\mathcal{X}$ w.r.t a reference point $x_{3}\in\mathcal{X}$ is defined by

$$ $\langle x_{1},x_{2}\rangle_{x_{3}}=\frac{1}{2}\left(d_{\mathcal{X}}(x_{1},x_{3 })+d_{\mathcal{X}}(x_{2},x_{3})-d_{\mathcal{X}}(x_{1},x_{2})\right).$ (9) $$

###### Definition A.1 ( δ 𝛿 \delta italic_δ -Hyperbolic Space) .

###### Definition A.2 (Gromov Product [Gromov, 1987 ] ) .

### A.4 Preliminaries on Trees

###### Definition A.3 .

Consider an undirected weighted graph $G=(V,E,\mathbf{A})$, where $V=\{1,2,\ldots,m\}$ is the node set, $E$ is the edge set, and $\mathbf{A}\in\mathbb{R}^{m\times m}$ is the non-negative edge weight matrix.
The shortest path metric $d_{G}(j,j^{\prime})$ of any two nodes $j,j^{\prime}\in V$ is defined by the length of the shortest path on the graph $G$ from node $j$ to node $j^{\prime}$.

###### Lemma A.1 .

The tree metric is 0 0-hyperbolic.

###### Definition A.4 .

A metric $d:V\times V\rightarrow\mathbb{R}$ is a tree metric if there exists a weighted tree $T=(V,E,\mathbf{A})$ such that for every $j,j^{\prime}\in V$, the metric $d(j,j^{\prime})$ is equal to the shortest path metric $d_{T}(j,j^{\prime})$.

###### Definition A.5 .

A balanced binary tree is a tree where each node has a degree of either 1, making it a leaf node, or 3, making it an internal node.
This tree structure can be either empty or consist of a root node linked to two disjoint binary trees, the left and right subtrees.
In a binary tree with $m$ leaves, there are exactly $m-1$ internal nodes.

###### Definition A.6 .

A rooted and balanced binary tree is a specific type of binary tree in which one internal node, designated as the root, has a degree of exactly 2, while the other internal nodes have a degree of 3. The remaining nodes are leaf nodes with a degree of 1.

###### Definition A.7 (Ultrametric) .

A metric $d_{u}:V\times V\rightarrow\mathbb{R}$ is ultrametric if the triangle inequality for any three points $j_{1},j_{2},j_{3}\in V$ is strengthened by the ultrametric inequality $d_{u}(j_{1},j_{2})\leq\max\{d_{u}(j_{1},j_{3}),d_{u}(j_{2},j_{3})\}$.

###### Remark .

If a tree is ultrametric, the distances of all the leaves to the root are identical.

###### Definition A.3 .

###### Lemma A.1 .

###### Definition A.4 .

###### Definition A.5 .

###### Definition A.6 .

###### Definition A.7 (Ultrametric) .

###### Remark .

## Appendix B Additional Related Work

Here, we provide more details about related methods and objectives.

Wasserstein Distance on Hyperbolic Manifold.
Wasserstein metric learning in hyperbolic spaces builds upon adapting the Wasserstein distance [Villani, ] to the structure of hyperbolic geometry [Lee, ].
This adaptation involves redefining the cost function in the optimization problem to incorporate the geodesic distance within the hyperbolic manifold.
The resulting Wasserstein metric maintains its intuitive interpretation as a measure of the “effort” required to transform one distribution into another while under the manifold constraints of hyperbolic structure.
It has proven valuable in various fields, for example, in shape analysis [Shi et al., ] and hierarchy matching [Alvarez-Melis et al., , Hoyos-Idrobo, ].
Recently, the work in Bonet et al. [] proposed extending the sliced-Wasserstein distance [Rabin et al., ] to hyperbolic representation, enabling efficient computation of the Wasserstein distance.
Our work, however, diverges fundamentally from these approaches.
We focus on observational data that inherently possess a hierarchy rather than probability distributions defined on hyperbolic spaces.
Additionally, our focus is on the TWD, aiming for efficient computation while addressing the specific challenges posed by the hidden feature hierarchy.

Existing TWDs.
The TWD is a popular method to reduce the computational complexity associated with calculating the Wasserstein distance.
This approach approximates the original metric using a tree metric, as described in several studies [Indyk and Thaper, , Le et al., , Evans and Matsen, , Leeb, , Takezawa et al., , Yamada et al., ].
The resulting methods provide a coarser approximation of OT in a Euclidean space, facilitating more efficient computations.
One notable method is the Quadtree [Indyk and Thaper, ], which creates a random tree by recursively dividing hypercubes.
This tree is then used to compute the TWD.
Flowtree [Backurs et al., ] modifies the Quadtree method by focusing on the optimal flow and evaluating the cost of this flow within the ground metric.
The Tree-Sliced Wasserstein Distance (TSWD) [Le et al., ] represents another TWD technique.
It computes the average of TWDs by randomly sampling trees. The number of sampling trees is a hyperparameter, commonly set to one, five, or ten, resulting in the variants TSWD-1, TSWD-5, and TSWD-10.
Further advancements are presented in Yamada et al. [], which introduce WQTWD and WCTWD. These methods employ a Quadtree or clustering-based tree to approximate the Wasserstein distances by optimizing the weights on the tree.
Lastly, UltraTree [Chen et al., ] constructs an ultrametric tree [Johnson, ] by minimizing the OT regression cost, aiming to approximate the Wasserstein distance within the original metric space.
Our approach, however, departs significantly from these existing methods.
While we also aim to accelerate the computation of the Wasserstein distance, we leverage the tree structure to introduce and recover the latent feature hierarchy. Specifically, our novel hyperbolic tree decoding algorithm is designed to recover hidden feature hierarchies.
Our TWD approximates the Wasserstein distance incorporating the latent feature hierarchy, rather than relying on the original metric in Euclidean space.

Tree Construction by Graph-Based and HC Methods.
Another line of related work involves constructing trees using existing graph-based algorithms, such as Minimum Spanning Tree (MST) [Alon et al., ], or hierarchical clustering (HC) methods, particularly those that leverage hyperbolic geometry. Generic graph-based algorithms, like those used in existing TWDs, aim to approximate the original metric with a tree metric.
In the context of hierarchical clustering, various methods have been proposed.
UltraFit [Chierchia and Perret, ] addresses an “ultrametric fitting” problem using Euclidean embeddings.
gHHC [Monath et al., ] operates under the assumption that partial information about the optimal clustering, specifically the hyperbolic embeddings of the leaves, is known.
HHC [Chami et al., ] decodes a tree from a two-dimensional hyperbolic embedding for hierarchical clustering by optimizing the continuous Dasgupta’s cost [Dasgupta, ].
While these methods may offer theoretical guarantees regarding clustering quality, they do not provide a theoretical guarantee for recovering the latent hierarchy underlying high-dimensional features, as our method does. Our approach uniquely focuses on revealing the latent feature hierarchy, ensuring a more accurate representation of the latent hierarchical structure underlying high-dimensional features.

Fréchet Mean on Manifolds.
The Fréchet mean [Fréchet, ] is a powerful tool for data analysis on Riemannian manifolds [Lee, ], including hyperbolic, spherical, and various matrix manifolds.
Its application extends across a diverse range of fields, including batch normalization [Brooks et al., , Lou et al., ], batch effect removal [Lin et al., ], domain adaptation [Yair et al., ], shape analysis and object recognition [McLeod et al., ], to name but a few.
In our work, we follow this line of work, using the Fréchet mean for constructing the HD-LCA in a product manifold of hyperbolic spaces (see Def. 4.1 and Def. 4.3).
Then, this HD-LCA serves as a primary building block in our binary tree decoding (Alg. 1), which in turn, is used in the computation of the proposed TWD (Alg. 2).

## Appendix C Theoretical Analysis and Proofs

We note that the numbering of the statements corresponds to the numbering used in the paper, and we have also included several separately numbered propositions and lemmas that are used in supporting the proofs presented.
We restate the claim of each statement for convenience.

### C.1 Proof of Proposition 4.1

##### Proposition 4.1 .

The hyperbolic LCA $\mathbf{z}_{j}^{k}\vee\mathbf{z}_{j^{\prime}}^{k}$ in Def. 4.1 has a closed-form solution, given by

$$ $\mathbf{z}_{j}^{k}\vee\mathbf{z}_{j^{\prime}}^{k}=\left[\frac{1}{2}\left(\bm{ \psi}_{j}^{k}+\bm{\psi}_{j^{\prime}}^{k}\right)^{\top},\mathtt{proj}(\mathbf{z }_{j}^{k}\vee\mathbf{z}_{j^{\prime}}^{k})\right]^{\top},$ $$

where $\mathtt{proj}(\mathbf{z}_{j}^{k}\vee\mathbf{z}_{j^{\prime}}^{k})=\left\lVert
\left[\frac{1}{2}\left(\bm{\psi}_{j}^{k}-\bm{\psi}_{j^{\prime}}^{k}\right)^{
\top},2^{k/2-2}\right]^{\top}\right\rVert_{2}$.

###### Proof.

Let $\gamma_{\mathbf{z}_{j}^{k}\leadsto\mathbf{z}_{j^{\prime}}^{k}}(t)$ be the unique geodesic path from $\mathbf{z}_{j}^{k}$ to $\mathbf{z}_{j^{\prime}}^{k}$ such that $\gamma_{\mathbf{z}_{j}^{k}\leadsto\mathbf{z}_{j^{\prime}}^{k}}(0)=\mathbf{z}_{
j}^{k}$ and $\gamma_{\mathbf{z}_{j}^{k}\leadsto\mathbf{z}_{j^{\prime}}^{k}}(1)=\mathbf{z}_{
j^{\prime}}^{k}$.
Note that the geodesic path $\gamma_{\mathbf{z}_{j}^{k}\leadsto\mathbf{z}_{j^{\prime}}^{k}}$ is a semi-circle, and the last coordinate of the points $\mathbf{z}_{j}^{k}$ and $\mathbf{z}_{j^{\prime}}^{k}$ are the same, i.e., $\mathbf{z}_{j}^{k}(m+1)=\mathbf{z}_{j^{\prime}}^{k}(m+1)=2^{k/2-2}$.
Therefore, the center point of the semi-circle projected on the plane $\mathbb{R}^{m}$ is $\mathbf{c}_{\mathbf{z}_{j}^{k}\leadsto\mathbf{z}_{j^{\prime}}^{k}}=\left[\frac
{1}{2}\left(\bm{\psi}_{j}^{k}+\bm{\psi}_{j^{\prime}}^{k}\right)^{\top},0\right
]^{\top}$.
Moreover, the radius of the geodesic path is given by

$$ $\displaystyle\mathtt{proj}(\mathbf{z}_{j}^{k}\vee\mathbf{z}_{j^{\prime}}^{k})$ $\displaystyle=\left\rVert\mathbf{c}_{\mathbf{z}_{j}^{k}\leadsto\mathbf{z}_{j^{ \prime}}^{k}}-\mathbf{z}_{j}^{k}\right\rVert_{2}$ $\displaystyle=\left\rVert\mathbf{c}_{\mathbf{z}_{j}^{k}\leadsto\mathbf{z}_{j^{ \prime}}^{k}}-\mathbf{z}_{j^{\prime}}^{k}\right\rVert_{2}$ $\displaystyle=\left\lVert\left[\frac{1}{2}\left(\bm{\psi}_{j}^{k}-\bm{\psi}_{j ^{\prime}}^{k}\right)^{\top},2^{k/2-2}\right]^{\top}\right\rVert_{2}.$ $$

The Riemannian mean defined using the Fréchet mean [Fréchet, ] has a closed-form expression, and is located at the midpoint of the geodesic curve

$$ $\mathbf{z}_{j}^{k}\vee\mathbf{z}_{j^{\prime}}^{k}=\left[\frac{1}{2}\left(\bm{ \psi}_{j}^{k}+\bm{\psi}_{j^{\prime}}^{k}\right)^{\top},\mathtt{proj}(\mathbf{z }_{j}^{k}\vee\mathbf{z}_{j^{\prime}}^{k})\right]^{\top}.$ $$

∎

###### Lemma C.1 .

The radius $\mathtt{proj}(\mathbf{z}_{j}^{k}\vee\mathbf{z}_{j^{\prime}}^{k})$ is lower-bounded by the diffusion distance at the diffusion time $2^{-k}$, given by

$$ $d_{\text{DM},2^{-k}}(j,j^{\prime})\leq\mathtt{proj}(\mathbf{z}_{j}^{k}\vee \mathbf{z}_{j^{\prime}}^{k}).$ (10) $$

###### Proof.

###### Lemma C.1 .

### C.2 Proof of Proposition 4.2

##### Proposition 4.2 .

For any $k_{1}\leq k_{2}$, $2^{-(k_{2}-k_{1})}\simeq d_{\mathbb{H}^{m+1}}(\mathbf{z}_{j}^{k_{2}}\vee
\mathbf{z}_{j^{\prime}}^{k_{2}},\mathbf{z}_{j}^{k_{2}})/d_{\mathbb{H}^{m+1}}(
\mathbf{z}_{j}^{k_{1}}\vee\mathbf{z}_{j^{\prime}}^{k_{1}},\mathbf{z}_{j}^{k_{1
}})$.

###### Proof.

By Prop. 4.1, we have

$$ $\mathbf{z}_{j}^{k_{1}}\vee\mathbf{z}_{j^{\prime}}^{k_{1}}=\left[\frac{1}{2} \left(\bm{\psi}_{j}^{k_{1}}+\bm{\psi}_{j^{\prime}}^{k_{1}}\right)^{\top}, \mathtt{proj}(\mathbf{z}_{j}^{k_{2}}\vee\mathbf{z}_{j^{\prime}}^{k_{1}})\right ]^{\top}$ $$

and

$$ $\mathbf{z}_{j}^{k_{2}}\vee\mathbf{z}_{j^{\prime}}^{k_{2}}=\left[\frac{1}{2} \left(\bm{\psi}_{j}^{k_{2}}+\bm{\psi}_{j^{\prime}}^{k_{2}}\right)^{\top}, \mathtt{proj}(\mathbf{z}_{j}^{k_{2}}\vee\mathbf{z}_{j^{\prime}}^{k_{2}})\right ]^{\top},$ $$

respectively.
At the $k_{1}$-th level, the Riemannian distance between the $k_{1}$-th hyperbolic LCA $\mathbf{z}_{j}^{k_{1}}\vee\mathbf{z}_{j^{\prime}}^{k_{1}}$ and $\mathbf{z}_{j}^{k_{1}}$ at the $k_{1}$-th level is given by

$$ $d_{\mathbb{H}^{m+1}}(\mathbf{z}_{j}^{k_{1}}\vee\mathbf{z}_{j^{\prime}}^{k_{1}} ,\mathbf{z}_{j}^{k_{1}})=2\sinh^{-1}\left(\left\lVert\mathbf{z}_{j}^{k_{1}} \vee\mathbf{z}_{j^{\prime}}^{k_{1}}-\mathbf{z}_{j}^{k_{1}}\right\rVert_{2}\Big {/}\left(2\sqrt{\mathbf{z}_{j}^{k_{1}}\vee\mathbf{z}_{j^{\prime}}^{k_{1}}(m+1) \mathbf{z}_{j}^{k_{1}}(m+1)}\right)\right).$ $$

Note that $\mathbf{z}_{j}^{k_{1}}(m+1)=2^{k_{1}/2-2}$ and by Prop. 4.1, we have $\mathbf{z}_{j}^{k_{1}}\vee\mathbf{z}_{j^{\prime}}^{k_{1}}(m+1)=\mathtt{proj}(
\mathbf{z}_{j}^{k_{1}}\vee\mathbf{z}_{j^{\prime}}^{k_{1}})$.
We have the same relation at the $k_{2}$ level. Therefore, for any $k_{1}\leq k_{2}$, we have

$$ $\displaystyle d_{\mathbb{H}^{m+1}}(\mathbf{z}_{j}^{k_{2}}\vee\mathbf{z}_{j^{ \prime}}^{k_{2}},\mathbf{z}_{j}^{k_{2}})/d_{\mathbb{H}^{m+1}}(\mathbf{z}_{j}^{ k_{1}}\vee\mathbf{z}_{j^{\prime}}^{k_{1}},\mathbf{z}_{j}^{k_{1}})$ $\displaystyle=$ $\displaystyle\frac{\sinh^{-1}\left(\left\lVert\mathbf{z}_{j}^{k_{2}}\vee \mathbf{z}_{j^{\prime}}^{k_{2}}-\mathbf{z}_{j}^{k_{2}}\right\rVert_{2}\Big{/} \left(2\sqrt{\mathbf{z}_{j}^{k_{2}}\vee\mathbf{z}_{j^{\prime}}^{k_{2}}(m+1) \mathbf{z}_{j}^{k_{2}}(m+1)}\right)\right)}{\sinh^{-1}\left(\left\lVert\mathbf {z}_{j}^{k_{1}}\vee\mathbf{z}_{j^{\prime}}^{k_{1}}-\mathbf{z}_{j}^{k_{1}} \right\rVert_{2}\Big{/}\left(2\sqrt{\mathbf{z}_{j}^{k_{1}}\vee\mathbf{z}_{j^{ \prime}}^{k_{1}}(m+1)\mathbf{z}_{j}^{k_{1}}(m+1)}\right)\right)}$ $\displaystyle=$ $\displaystyle\sinh^{-1}\left(2^{-k_{2}/2+1}\left\lVert\bm{\psi}_{j}^{k_{2}}- \bm{\psi}_{j^{\prime}}^{k_{2}}\right\rVert_{2}\right)\Big{/}\sinh^{-1}\left(2^ {-k_{1}/2+1}\left\lVert\bm{\psi}_{j}^{k_{1}}-\bm{\psi}_{j^{\prime}}^{k_{1}} \right\rVert_{2}\right)$ $\displaystyle\overset{(1)}{\geq}$ $\displaystyle\sinh^{-1}\left(2^{-k_{2}/2+1}\left\lVert\bm{\psi}_{j}^{k_{2}}- \bm{\psi}_{j^{\prime}}^{k_{2}}\right\rVert_{2}\right)\Big{/}\sinh^{-1}\left(2^ {-k_{1}/2+1}\left\lVert\bm{\psi}_{j}^{k_{2}}-\bm{\psi}_{j^{\prime}}^{k_{2}} \right\rVert_{2}\right)$ $\displaystyle=$ $\displaystyle 2^{-(k_{2}-k_{1})},$ $$

where the transition (1) is due to $\left\lVert\bm{\psi}_{j}^{k_{1}}-\bm{\psi}_{j^{\prime}}^{k_{1}}\right\rVert_{2
}\leq\left\lVert\bm{\psi}_{j}^{k_{2}}-\bm{\psi}_{j^{\prime}}^{k_{2}}\right
\rVert_{2}$ for any $k_{1}\leq k_{2}$.
Similarly, we have

$$ $\displaystyle d_{\mathbb{H}^{m+1}}(\mathbf{z}_{j}^{k_{2}}\vee\mathbf{z}_{j^{ \prime}}^{k_{2}},\mathbf{z}_{j}^{k_{2}})/d_{\mathbb{H}^{m+1}}(\mathbf{z}_{j}^{ k_{1}}\vee\mathbf{z}_{j^{\prime}}^{k_{1}},\mathbf{z}_{j}^{k_{1}})$ $\displaystyle=$ $\displaystyle\frac{\sinh^{-1}\left(\left\lVert\mathbf{z}_{j}^{k_{2}}\vee \mathbf{z}_{j^{\prime}}^{k_{2}}-\mathbf{z}_{j}^{k_{2}}\right\rVert_{2}\Big{/} \left(2\sqrt{\mathbf{z}_{j}^{k_{2}}\vee\mathbf{z}_{j^{\prime}}^{k_{2}}(m+1) \mathbf{z}_{j}^{k_{2}}(m+1)}\right)\right)}{\sinh^{-1}\left(\left\lVert\mathbf {z}_{j}^{k_{1}}\vee\mathbf{z}_{j^{\prime}}^{k_{1}}-\mathbf{z}_{j}^{k_{1}} \right\rVert_{2}\Big{/}\left(2\sqrt{\mathbf{z}_{j}^{k_{1}}\vee\mathbf{z}_{j^{ \prime}}^{k_{1}}(m+1)\mathbf{z}_{j}^{k_{1}}(m+1)}\right)\right)}$ $\displaystyle=$ $\displaystyle\sinh^{-1}\left(2^{-k_{2}/2+1}\left\lVert\bm{\psi}_{j}^{k_{2}}- \bm{\psi}_{j^{\prime}}^{k_{2}}\right\rVert_{2}\right)\Big{/}\sinh^{-1}\left(2^ {-k_{1}/2+1}\left\lVert\bm{\psi}_{j}^{k_{1}}-\bm{\psi}_{j^{\prime}}^{k_{1}} \right\rVert_{2}\right)$ $\displaystyle\overset{(1)}{\leq}$ $\displaystyle 2\cdot 2^{-k_{1}/2+1}\Big{/}\frac{\eta}{2}\cdot 2^{-k_{1}/2+1}$ $$

where the transition (1) is because there is a lower bound $\eta$ of any Hellinger distance for $j\neq j^{\prime}$.
Therefore, we show that $2^{-(k_{2}-k_{1})}\simeq d_{\mathbb{H}^{m+1}}(\mathbf{z}_{j}^{k_{2}}\vee
\mathbf{z}_{j^{\prime}}^{k_{2}},\mathbf{z}_{j}^{k_{2}})/d_{\mathbb{H}^{m+1}}(
\mathbf{z}_{j}^{k_{1}}\vee\mathbf{z}_{j^{\prime}}^{k_{1}},\mathbf{z}_{j}^{k_{1
}})$.
∎

###### Proof.

### C.3 Proof of Proposition 4.3

##### Proposition 4.3 .

The HD-LCA $\overline{\mathbf{h}}_{j,j^{\prime}}$ in Def. 4.3 has a closed-form solution, given by $\overline{\mathbf{h}}_{j,j^{\prime}}=\left[0,\ldots,0,a_{j,j^{\prime}}\right]^
{\top}$, where $a_{j,j^{\prime}}=\sqrt[(K_{c}+1)]{\mathtt{proj}(\mathbf{z}_{j}^{0}\vee\mathbf{
z}_{j^{\prime}}^{0})\cdots\mathtt{proj}(\mathbf{z}_{j}^{K_{c}}\vee\mathbf{z}_{
j^{\prime}}^{K_{c}})}$.

###### Proof.

The orthogonal projections $\{\mathbf{o}^{k}_{j,j^{\prime}}\}_{k=0}^{K_{c}}$ lie on the unique geodesic path $\gamma_{0}$ that is on the $(m+1)$-axis in $\mathbb{H}^{m+1}$.
Note that the geodesic distance joining any two points $\mathbf{o}^{k_{1}}_{j,j^{\prime}}$ and $\mathbf{o}^{k_{2}}_{j,j^{\prime}}$ is given by $\left\lvert\ln\left(\frac{\mathtt{proj}(\mathbf{z}_{j}^{k_{1}}\vee\mathbf{z}_{
j^{\prime}}^{k_{1}})}{\mathtt{proj}(\mathbf{z}_{j}^{k_{2}}\vee\mathbf{z}_{j^{
\prime}}^{k_{2}})}\right)\right\rvert$ [Stahl, , Udriste, ].
The Fréchet mean [Fréchet, ] of $\{\mathbf{o}^{k}_{j,j^{\prime}}\}_{k=0}^{K_{c}}$ is the point that minimizes the sum of squared distances to all points in the set.
For our set of points, the Fréchet mean mean will have the form $\left[0,\ldots,0,a_{j,j^{\prime}}\right]^{\top}$, where $a_{j,j^{\prime}}\in\mathbb{R}_{>0}$.
We want to find $a_{j,j^{\prime}}$ that minimizes

$$ $\displaystyle\overline{\mathbf{h}}_{j,j^{\prime}}$ $\displaystyle\coloneqq\underset{\mathbf{h}\in\mathbb{H}^{m+1}}{\arg\min}\;\sum _{k=0}^{K_{c}}\;d_{\mathbb{H}^{m+1}}^{2}(\mathbf{h},\mathbf{o}^{k}_{j,j^{ \prime}})$ $\displaystyle=\underset{\mathbf{h}=[0,\ldots,0,a]^{\top}\in\mathbb{H}^{m+1}}{ \arg\min}\sum_{k=0}^{K_{c}}\;d_{\mathbb{H}^{m+1}}^{2}(\mathbf{h},\mathbf{o}^{k }_{j,j^{\prime}})$ $\displaystyle=\underset{\mathbf{h}=[0,\ldots,0,a]^{\top}\in\mathbb{H}^{m+1}}{ \arg\min}\sum_{k=0}^{K_{c}}\;\left\lvert\ln\left(\frac{\mathtt{proj}(\mathbf{z }_{j}^{k}\vee\mathbf{z}_{j^{\prime}}^{k})}{a}\right)\right\rvert^{2}.$ $$

Therefore, the geometric mean $\overline{\mathbf{h}}_{j,j^{\prime}}$ is obtained by setting $a_{j,j^{\prime}}=\sqrt[(K_{c}+1)]{\mathtt{proj}(\mathbf{z}_{j}^{0}\vee\mathbf{
z}_{j^{\prime}}^{0})\cdots\mathtt{proj}(\mathbf{z}_{j}^{K_{c}}\vee\mathbf{z}_{
j^{\prime}}^{K_{c}})}$.

∎

###### Proof.

### C.4 Proof of Proposition 4.4

##### Proposition 4.4 .

Let $r$ be the root node of $B$ obtained in Alg. 1 and $\mathbf{z}_{r}$ be the corresponding embedding in $\mathcal{M}$. The HD-LCA depth is equivalent to the Gromov product [Gromov, ], i.e., $d_{\mathcal{M}}(\mathbf{z}_{r},\overline{\mathbf{h}}_{j,j^{\prime}})\simeq
\langle j,j^{\prime}\rangle_{r}$, where $\langle j,j^{\prime}\rangle_{r}=\frac{1}{2}(d_{B}(j,r)+d_{B}(j^{\prime},r)-d_{
B}(j,j^{\prime}))$.

###### Proof.

Let $\mathbf{z}_{j}$ and $\mathbf{z}_{j^{\prime}}$ be two points in $\mathcal{M}$.
We denote $\gamma_{\mathbf{z}_{j}\leadsto\mathbf{z}_{j^{\prime}}}(t)$ the unique geodesic path from $\mathbf{z}_{j}$ to $\mathbf{z}_{j^{\prime}}$ such that $\gamma_{\mathbf{z}_{j}\leadsto\mathbf{z}_{j^{\prime}}}(0)=\mathbf{z}_{j}$ and $\gamma_{\mathbf{z}_{j}\leadsto\mathbf{z}_{j^{\prime}}}(1)=\mathbf{z}_{j^{
\prime}}$.
Following Bowditch [], for any point $\mathbf{z}_{j^{\prime\prime}}\in\mathcal{M}$, we have $d_{\mathcal{M}}(\mathbf{z}_{j^{\prime\prime}},\mathbf{z})\geq\langle j,j^{
\prime}\rangle_{j^{\prime\prime}}$, where $\mathbf{z}\in\gamma_{\mathbf{z}_{j}\leadsto\mathbf{z}_{j^{\prime}}}(t)$.
Note that if $\mathbf{z}_{j^{\prime\prime}}\in\gamma_{\mathbf{z}_{j}\leadsto\mathbf{z}_{j^{
\prime}}}(t)$, then $\langle j,j^{\prime}\rangle_{j^{\prime\prime}}=0$.
In addition, let $\gamma_{\mathbf{z}_{j}\leadsto\mathbf{z}_{j^{\prime\prime}}}(t)$ and $\gamma_{\mathbf{z}_{j^{\prime}}\leadsto\mathbf{z}_{j^{\prime\prime}}}(t)$ be the unique geodesic path from $\mathbf{z}_{j}$ to $\mathbf{z}_{j^{\prime\prime}}$ and the unique geodesic path from $\mathbf{z}_{j^{\prime}}$ to $\mathbf{z}_{j^{\prime\prime}}$, respectively.
Let $\triangle_{j,j^{\prime},j^{\prime\prime}}=(\gamma_{\mathbf{z}_{j}\leadsto
\mathbf{z}_{j^{\prime}}}(t),\gamma_{\mathbf{z}_{j}\leadsto\mathbf{z}_{j^{
\prime\prime}}}(t),\gamma_{\mathbf{z}_{j^{\prime}}\leadsto\mathbf{z}_{j^{
\prime\prime}}}(t))$ be the triangle connected by these geodesic paths.
As $\mathcal{M}$ is 0 0-hyperbolic, the triangle $\triangle_{j,j^{\prime},j^{\prime\prime}}$ is 0-centre and by triangle inequality, we have

$$ $\displaystyle d_{\mathcal{M}}(\mathbf{z}_{j},\mathbf{z})+d_{\mathcal{M}}( \mathbf{z}_{j^{\prime}},\mathbf{z})\leq d_{\mathcal{M}}(\mathbf{z}_{j},\mathbf {z}_{j^{\prime\prime}})$ $\displaystyle d_{\mathcal{M}}(\mathbf{z}_{j^{\prime}},\mathbf{z})+d_{\mathcal{ M}}(\mathbf{z}_{j^{\prime\prime}},\mathbf{z})\leq d_{\mathcal{M}}(\mathbf{z}_{ j}^{\prime},\mathbf{z}_{j^{\prime\prime}})$ $\displaystyle d_{\mathcal{M}}(\mathbf{z}_{j},\mathbf{z})+d_{\mathcal{M}}( \mathbf{z}_{j^{\prime}},\mathbf{z})=d_{\mathcal{M}}(\mathbf{z}_{j},\mathbf{z}_ {j^{\prime}}).$ $$

Therefore, $d_{\mathcal{M}}(\mathbf{z}_{j^{\prime\prime}},\mathbf{z})\leq\langle j,j^{
\prime}\rangle_{j^{\prime\prime}}$.
By the direct application of the aforementioned two inequality, we have $d_{\mathcal{M}}(\mathbf{z}_{r},\overline{\mathbf{h}}_{j,j^{\prime}})\simeq
\langle j,j^{\prime}\rangle_{r}$.
∎

###### Proof.

### C.5 Proof of Theorem 4.1

##### Theorem 4.1 .

For sufficiently large $K_{c}$ and $m$, the decoded tree metric $d_{B}$ in Alg. 1 is bilipschitz equivalent to the underlying tree metric $d_{T}$, where $T$ is the ground latent tree assumed from Sec. 4.1 and $d_{T}$ is the hidden tree metric between features defined as the length of the shortest path on $T$. Specifically, the bilipschitz constants can be determined by $\frac{1}{\sqrt{2}}d_{B}(j_{1},j_{2})\leq d_{T}(j_{1},j_{2})\leq d_{B}(j_{1},j_
{2})$ for all features $j_{1},j_{2}\in[m]$.

###### Proof.

We adopt the equivalence proof from Lin et al. [], which stems from Coifman and Lafon [] and [Leeb and Coifman, ].
Note that diffusion operator $\mathbf{P}$ is typically used for manifold learning, where it be viewed as a discrete proxy of the underlying manifold.
It is shown that the diffusion operator recovers the underlying manifold [Coifman and Lafon, ], i.e., $\mathbf{P}\overset{m\rightarrow\infty,\epsilon\rightarrow 0}{\longrightarrow}
\mathbf{H}=\exp(-\Delta)$, where
$\mathbf{H}$ is the Neumann heat kernel of the underlying manifold and $\Delta$ is the Laplace–Beltrami operator on the manifold.
In addition, the diffusion operator $\mathbf{P}$ is based on a finite set of points, which is a discrete approximation of $\mathbf{H}$.

For a family of the operator $\{A^{t}\}_{t\in\mathbb{R}}$, the multi-scale metric using the inverse hyperbolic sine function of the scaled Hellinger measure is defined by

$$ $d_{A,\mathcal{M}}(x,x^{\prime}):=\sum_{k\geq 0}2\sinh^{-1}\left(2^{-\frac{k}{2 }+1}\left\lVert\sqrt{A^{2^{-k}}(x,\cdot)}-\sqrt{A^{2^{-k}}(x^{\prime},\cdot)} \right\rVert_{2}\right),$ $$

where $A^{t}f(x)=\int_{\mathcal{X}}A^{t}(x,x^{\prime})f(x^{\prime})dx^{\prime}$.
Our goal is to show that $d_{B}(j,j^{\prime})\simeq d_{\mathbf{H},\mathcal{M}}(j,j^{\prime})$, and conclude with Thm. C.1 [Lin et al., ] which states the equivalence $d_{\mathbf{H},\mathcal{M}}(j,j^{\prime})\simeq d_{T}(j,j^{\prime})$.

Note that for small enough $\epsilon>0$, we have $d_{\mathbf{P},\mathcal{M}}(j,j^{\prime})\simeq d_{\mathbf{H},\mathcal{M}}(j,j^
{\prime})$.
Specifically, following Coifman and Lafon [], we have $\left\lVert\mathbf{P}-\mathbf{H}\right\rVert_{L^{2}(\mathcal{M})}\rightarrow 0$ and $\left\lVert\mathbf{P}-\mathbf{H}\right\rVert_{L^{1}(\mathcal{M})}\rightarrow 0$ as $\epsilon\rightarrow 0$.
We denote $d_{\mathcal{M}}(j,j^{\prime})=d_{\mathbf{P},\mathcal{M}}(j,j^{\prime})\simeq d
_{\mathbf{H},\mathcal{M}}(j,j^{\prime})$ for abbreviation.

For any three features $j_{1},j_{2},j_{3}\in[m]$, assume the HD-LCA relation $[j_{2}\vee j_{3}\sim j_{1}]_{\mathcal{M}}$ holds.
Let $j_{2}\vee j_{3}$ and $j_{1}\vee j_{3}$ be LCAs of $(j_{2},j_{3})$ and $(j_{1},j_{3})$ on $B$ obtained in Alg. 1, respectively.
We denote $\mathbf{z}_{j_{2}}\vee\mathbf{z}_{j_{3}}$ and $\mathbf{z}_{j_{2}}\vee\mathbf{z}_{j_{3}}$ the corresponding embeddings.
By Prop. 4.1, we have

$$ $\displaystyle\mathbf{z}_{j_{2}}\vee\mathbf{z}_{j_{3}}=\left[\left[\frac{1}{2} \left(\bm{\psi}_{j_{2}}^{0}+\bm{\psi}_{j_{3}}^{0}\right)^{\top},\mathtt{proj}( \mathbf{z}_{j_{2}}^{0}\vee\mathbf{z}_{j_{3}}^{0})\right]^{\top},\ldots,\left[ \frac{1}{2}\left(\bm{\psi}_{j_{2}}^{k}+\bm{\psi}_{j_{3}}^{k}\right)^{\top}, \mathtt{proj}(\mathbf{z}_{j_{2}}^{K_{c}}\vee\mathbf{z}_{j_{3}}^{K_{c}})\right] ^{\top}\right]^{\top}$ and $\displaystyle\mathbf{z}_{j_{1}}\vee\mathbf{z}_{j_{3}}=\left[\left[\frac{1}{2} \left(\bm{\psi}_{j_{1}}^{0}+\bm{\psi}_{j_{3}}^{0}\right)^{\top},\mathtt{proj}( \mathbf{z}_{j_{1}}^{0}\vee\mathbf{z}_{j_{3}}^{0})\right]^{\top},\ldots,\left[ \frac{1}{2}\left(\bm{\psi}_{j_{1}}^{k}+\bm{\psi}_{j_{3}}^{k}\right)^{\top}, \mathtt{proj}(\mathbf{z}_{j_{1}}^{K_{c}}\vee\mathbf{z}_{j_{3}}^{K_{c}})\right] ^{\top}\right]^{\top}.$ $$

Therefore, the tree distance between $j_{1}$ and $j_{3}$ is given by $d_{B}(j_{1},j_{3})=d_{\mathcal{M}}(\mathbf{z}_{j_{1}},\mathbf{z}_{j_{1}}\vee
\mathbf{z}_{j_{3}})+d_{\mathcal{M}}(\mathbf{z}_{j_{1}}\vee\mathbf{z}_{j_{3}},
\mathbf{z}_{j_{2}}\vee\mathbf{z}_{j_{3}})+d_{\mathcal{M}}(\mathbf{z}_{j_{2}}
\vee\mathbf{z}_{j_{3}},\mathbf{z}_{j_{3}})$.
By triangle inequality, we have $d_{B}(j_{1},j_{3})\leq d_{\mathcal{M}}(\mathbf{z}_{j_{1}},\mathbf{z}_{j_{3}})$.
We denote $d_{\mathcal{M}}(j_{1},j_{2})=d_{\mathcal{M}}(\mathbf{z}_{j_{1}},\mathbf{z}_{j_
{3}})$ for simplicity.
By Prop. 4.2 and Prop. 4.4, we have $c\cdot d_{\mathcal{M}}({j_{1}},{j_{3}})\leq d_{B}(j_{1},j_{3})$.
Moreover, since the HD-LCA relation $[j_{2}\vee j_{3}\sim j_{1}]_{\mathcal{M}}$ holds, the tree distance between $j_{2},j_{3}\in[m]$ on $B$ can be directly obtained by $d_{B}(j_{2},j_{3})=d_{\mathcal{M}}(j_{2},j_{3})$.
As only one of the HD-LCA relations $[j_{2}\vee j_{3}\sim j_{1}]_{\mathcal{M}}$, $[j_{1}\vee j_{3}\sim j_{2}]_{\mathcal{M}}$, or $[j_{1}\vee j_{2}\sim j_{3}]_{\mathcal{M}}$ holds in $B$, the proof can be applied to any three points $j_{1},j_{2},j_{3}\in[m]$ of which HD-LCA relation holds.
Using Monte-Carlo integration, the convergence is with probability one as $K_{c}\rightarrow\infty$.
For a finite set, we have a high probability bound on the convergence.
Therefore, $d_{B}\simeq d_{\mathcal{M}}\simeq d_{T}$.
Specifically, by taking $\alpha\rightarrow\frac{1}{2}$ and $\beta=2$ in Section 3 in Leeb and Coifman [] and Appendix C in Lin et al. [], the bilipschitz constants $c$ and $C$ can be determined by $\frac{1}{\sqrt{2}}$ and $1$ in $d_{B}\simeq d_{T}$, respectively.
∎

###### Theorem C.1 (Theorem 1 in Lin et al. [ 2023 ] ) .

For $\alpha\rightarrow\frac{1}{2}$ and sufficiently large $K_{c}$, we have $d_{\mathcal{M}}\simeq d_{T}$.

###### Proof.

###### Theorem C.1 (Theorem 1 in Lin et al. [ 2023 ] ) .

### C.6 Proof of Theorem 4.2

##### Theorem 4.2 .

For sufficiently large $K_{c}$ and $m$, the TWD defined in Eq. (3) is bilipschitz equivalent to the TWD with the ground latent tree $T$. Specifically, the bilipschitz constants can be determined by $\frac{1}{\sqrt{2}}\mathtt{TW}(\mathbf{x}_{i},\mathbf{x}_{i^{\prime}},B)\leq
\mathtt{TW}(\mathbf{x}_{i},\mathbf{x}_{i^{\prime}},T)\leq\mathtt{TW}(\mathbf{x
}_{i},\mathbf{x}_{i^{\prime}},B)$ for all samples $i,i^{\prime}\in[n]$.

###### Proof.

Let $\mu_{i}$ and $\mu_{i^{\prime}}$ be two probability distributions supported on $T$ with a ground metric $d_{T}$, $\Pi(\mu_{i},\mu_{i^{\prime}})$ be the set of couplings $\pi$ on the product space $T\times T$ such that $\pi(B_{1}\times T)=\mu_{i}(B_{1})$ and $\pi(T\times B_{2})=\mu_{i^{\prime}}(B_{2})$ for any sets $B_{1},B_{2}\subset T$.
The Wasserstein distance [Villani, ] between $\mu_{i}$ and $\mu_{i^{\prime}}$ is defined by

$$ $\mathtt{OT}(\mu_{i},\mu_{i^{\prime}},d_{T})\coloneqq\underset{\pi\in\Pi(\mu_{i },\mu_{i^{\prime}})}{\inf}\int_{T\times T}d(x,y)\text{d}\pi(x,y).$ (11) $$

Let $\mathcal{F}_{d_{T}}$ and $\left\lVert\cdot\right\rVert_{L_{d_{T}}}$
denote the Lipschitz functions w.r.t. $d_{T}$ and Lipschitz norm w.r.t. $d_{T}$, respectively.
The Kantorovich-Rubinstein dual [Villani, ] of Eq. (11) is given by

$$ $\mathtt{OT}(\mu_{i},\mu_{i^{\prime}},d_{T})=\underset{\left\lVert f\right \rVert_{L_{d_{T}}}\leq 1}{\sup}\int_{T}f(x)\mu_{i}(dx)-\int_{T}f(y)\mu_{i^{ \prime}}(dy).$ $$

By Evans and Matsen [], the witness function $f$ can be represented by a Borel function $g:T\rightarrow[-1,1]$, i.e.,

$$ $\int_{T}f(x)\mu_{i}(dx)=\int_{T}g(z)\lambda(dz)\mu_{i}(\Gamma_{T}(z)),$ $$

where $\lambda$ is the unique Borel measure on $T$.
Therefore, following Le et al. [], the Wasserstein distance $\mathtt{OT}(\mu_{i},\mu_{i^{\prime}},d_{T})$ can be written as

$$ $\displaystyle\mathtt{OT}(\mu_{i},\mu_{i^{\prime}},d_{T})$ $\displaystyle=\sup\int_{T}\left(\mu_{i}(\Gamma_{T}(z))-\mu_{i^{\prime}}(\Gamma _{T}(z))\right)g(z)\lambda(dz)$ $\displaystyle=\int_{T}\left\lvert\mu_{i}(\Gamma_{T}(z))-\mu_{i^{\prime}}( \Gamma_{T}(z))\right\rvert\lambda(dz)$ $\displaystyle=\sum_{v\in V}\omega_{v}\left\lvert\sum_{u\in\Gamma_{T}(v)}\left( \mu_{i}(u)-\mu_{i^{\prime}}(u)\right)\right\rvert$ $\displaystyle=\mathtt{TW}(\mu_{i},\mu_{i^{\prime}},T).$ $$

Note that based on Thm. 4.1, the decoded tree metric $d_{B}(\cdot,\cdot)$ and the underlying tree metric $d_{T}(\cdot,\cdot)$ are bilipschitz equivalent.
Additionally, the TWD using the tree $T$ is the Wasserstein distance with ground pairwise distance $d_{T}$ [Evans and Matsen, ], i.e., $\mathtt{TW}(\mu_{1},\mu_{2},T)=\mathtt{OT}(\mu_{1},\mu_{2},d_{T})$.
Therefore, for any two discrete histograms $\mathbf{x}_{i}$ and $\mathbf{x}_{i^{\prime}}$ supported on $T$, we have $\mathtt{TW}(\mathbf{x}_{i},\mathbf{x}_{i^{\prime}},B)\simeq\mathtt{TW}(\mathbf
{x}_{i},\mathbf{x}_{i^{\prime}},T)$, and the bilipschitz constants are $\frac{1}{\sqrt{2}}$ and $1$, respectively.
∎

###### Proof.

### C.7 Proof of Proposition 4.5

##### Proposition 4.5 .

The TWD in Eq. (3) is bilipschitz equivalent to the sum of the multi-scale hyperbolic TWDs $\mathtt{TW}(\mathbf{x}_{i},\mathbf{x}_{i^{\prime}},B^{k})$ for $k=0,1,\ldots,K_{c}$, i.e., $\mathtt{TW}(\mathbf{x}_{i},\mathbf{x}_{i^{\prime}},B)\simeq\sum_{k=0}^{K_{c}}
\mathtt{TW}(\mathbf{x}_{i},\mathbf{x}_{i^{\prime}},B^{k})$.

###### Proof.

By Alg. 3, the $k$-th rooted binary tree $B^{k}=(\widehat{V}_{c}^{k},\widehat{E}_{c}^{k},\widehat{\mathbf{A}}_{c}^{k})$ is decoded based on the $k$-th scale hyperbolic LCA $\{\mathbf{z}_{j}^{k}\vee\mathbf{z}_{j^{\prime}}^{k}\}_{j,j^{\prime}}$, and the length of the edge is determined by the hyperbolic distance $d_{\mathbb{H}^{m+1}}$.
By Prop. 4.2, for all $k\in\{0,1,\ldots,K_{c}\}$, the tree metric of the $k$-th rooted binary tree between the features $j,j^{\prime}\in[m]$ is bounded by

$$ $\beta^{\prime}_{k}d_{\widehat{B}_{c}^{0}}(j,j^{\prime})\leq d_{B^{k}}(j,j^{ \prime})\leq\beta_{k}d_{\widehat{B}_{c}^{0}}(j,j^{\prime}),$ $$

where $\beta_{k},\beta^{\prime}_{k}>0$, and $d_{\widehat{B}_{c}^{0}}(\cdot,\cdot)$ is the tree metric of $\widehat{B}_{c}^{0}$.
Since the embedding distance is induced by the $\ell_{1}$ distance on the product of $K_{c}+1$ hyperbolic spaces, the tree metric $d_{B}$ is the $\ell_{1}$ distance on the product of $K_{c}+1$ rooted binary tree $\{B^{k}\}_{k=0}^{K_{c}}$, i.e., for any two features $j,j^{\prime}\in[m]$, we have

$$ $d_{B}(j,j^{\prime})=\sum_{k=0}^{K_{c}}d_{B^{k}}(j,j^{\prime}).$ $$

Incorporating the above equation with Prop. 4.2, the tree metric of the decoded tree $d_{B}$ is bounded by the tree metric $d_{\widehat{B}_{c}^{0}}$, given by

$$ $d_{\widehat{B}_{c}^{0}}(j,j^{\prime})\sum_{k=0}^{K_{c}}\;\beta^{\prime}_{k} \leq d_{B}(j,j^{\prime})\leq d_{\widehat{B}_{c}^{0}}(j,j^{\prime})\sum_{k=0}^{ K_{c}}\;\beta_{k}.$ $$

Under the condition of uniform cost scaling in OT, we have

$$ $\mathtt{OT}(\mathbf{x}_{i},\mathbf{x}_{i^{\prime}},d_{B})\leq\mathtt{OT}( \mathbf{x}_{i},\mathbf{x}_{i^{\prime}},d_{\widehat{B}_{c}^{0}})\sum_{k=0}^{K_{ c}}\;\beta_{k}\overset{(1)}{\leq}\sum_{k=0}^{K_{c}}\mathtt{OT}(\mathbf{x}_{i}, \mathbf{x}_{i^{\prime}},d_{B^{k}}),$ $$

where transition (1) is based on that $\mathtt{OT}(\mathbf{x}_{i},\mathbf{x}_{i^{\prime}},d_{B^{k}})$ is at most $\beta_{k}\cdot\mathtt{OT}(\mathbf{x}_{i},\mathbf{x}_{i^{\prime}},d_{\widehat{B
}_{c}^{0}})$.
The lower bound can be obtained in a similar way.
Therefore, as $\mathtt{TW}(\mu_{1},\mu_{2},T)=\mathtt{OT}(\mu_{1},\mu_{2},d_{T})$, the proposed TWD is bilipschitz equivalent to the hyperbolic tree-sliced Wasserstein distance

$$ $\mathtt{TW}(\mathbf{x}_{i},\mathbf{x}_{i^{\prime}},B)\simeq\sum_{k=0}^{K_{c}} \mathtt{TW}(\mathbf{x}_{i},\mathbf{x}_{i^{\prime}},B^{k}).$ $$

∎

###### Proof.

## Appendix D Additional Details on Experimental Study

We describe the setups and additional details of our experiments in Sec. 5.
The experiments are performed on NVIDIA DGX A100.

### D.1 Implementation Details

Setup.
In our experiments, the datasets are divided into a $70\%$ training set and a $30\%$ testing set.
The split of the word-document datasets in Sec. 5.2 follows previous work [Kusner et al., , Huang et al., ].
This split was applied to both learning of the underlying feature tree and to the kNN model.
Therefore, there is no train data leakage to the test data. The feature tree constructed from the training data can be reused to process new samples.
For each evaluation, we employ a $k$NN classifier with varying $k\in\{1,3,5,7,9,11,13,15,17,19\}$.
The evaluation is repeated five times, and we report the best result, averaged over these five runs.

**Table 2: Statistics on real-world word-document and scRNA-seq benchmarks.**
| Dataset | $\#$ Samples ($n$) | $\#$ Classes | $\#$ Features ($m$) |
| --- | --- | --- | --- |
| BBCSPORT | 517 documents | 5 | 13243 BOW |
| TWITTER | 2176 documents | 3 | 6344 BOW |
| CLASSIC | 4965 documents | 4 | 24277 BOW |
| AMAZON | 5600 documents | 4 | 42063 BOW |
| Zeisel | 3005 cells | 7 | 4000 genes |
| CBMC | 8617 cells | 13 | 500 genes |

Datasets.
In real-world word-document datasets, we test four word-document benchmarks in Kusner et al. []: (i) the BBCSPORT consisting of 13243 bags of words (BOW) and 517 articles categorized into five sports types, (ii) the TWITTER comprising 6344 BOW and 2176 tweets in three types of sentiment, (iii) the CLASSIC, including 24277 BOW and 4965 academic papers from four publishers, and (iv) the AMAZON containing 42063 BOW and 5600 reviews of four products.
The types of documents are used as labels in classification tasks.
In scRNA-seq data analysis, we test two scRNA-seq datasets in Dumitrascu et al. [].
The first dataset, the Zeisel dataset, focuses on the mouse cortex and hippocampus [Zeisel et al., ].
It comprises 4000 gene markers and 3005 single cells, categorized into seven distinct cell types.
The second dataset, the CBMC dataset, is derived from a cord blood mononuclear cell study [Stoeckius et al., ].
It comprises 500 gene markers and 8617 single cells, which belong to 13 different cell types.
The cell types are used as labels in cell classification tasks.
We report the statistics of these benchmarks in Tab. LABEL:tab:dataset_statistics.

Kernel Type, Scale, and Initial Distance Metric.
When constructing the Gaussian kernel in Sec. 3, the initial distance is derived from the cosine similarity calculated in the ambient space, and the kernel scale is set to be $\{0.1,1,2,5,10\}\times\sigma$, where $\sigma$ is the median of the pairwise distances.

The diffusion operator’s scale parameter $\epsilon$ controls information propagation across data points. Smaller $\epsilon$ values preserve local structures, beneficial for distinct clusters but may cause overfitting and noise sensitivity. Larger $\epsilon$ values capture broader relationships, useful for overlapping clusters but can blur distinctions. To balance these effects, we set $\epsilon$ within a standard range that is commonly-used in kernel and manifold learning.

The maximal scale $K_{c}$ determines the range of scales over which the hyperbolic manifold operates. We follow [Lin et al., ] and suggest setting $K_{c}\in\{0,1,\ldots,19\}$. Our empirical study, consistent with [Lin et al., ], shows that increasing the maximal scale does not significantly impact classification accuracy. This robustness simplifies the parameter selection process by indicating that the method performs reliably across a broad range of $K_{c}$ values.

As these parameters are task-specific and do not have a one-size-fits-all solution, we use Optuna [Akiba et al., ] to efficiently explore the parameter space and identify optimal settings, where $(\epsilon,K_{c})$ are set $(10,12)$, $(5,7)$, $(1,15)$, $(10,14)$, $(2,7)$, and $(5,9)$ for BBCSPORT, TWITTER, CLASSIC, AMAZON, Zeisel, and CBMC, respectively.
We note that the selection of kernel type, scale, and the initial distance metric in the kernel influences the effectiveness of our approach.
Such choices are not only crucial but also task-specific, which is in line with the common practice in kernel and manifold learning methods.
These elements, which do not have a one-size-fits-all solution, continue to be of wide interest in ongoing research [Lindenbaum et al., ].

In our study, we employed the Gaussian kernel with a distance measure based on cosine similarity following prior research [Jaskowiak et al., , Kenter and De Rijke, ], which is standard and commonly used in classification tasks and manifold learning.
In our empirical analysis, we observed that employing the Euclidean distance in the ambient space tends to be highly sensitive and underperforms compared to using a distance based on the cosine similarity in the tested datasets, as detailed in the App. E.
Our primary objective is to highlight the unique aspects of our method rather than kernel selection, which, while crucial, is a common consideration across all kernel and manifold learning methods.

### D.2 Probabilistic Hierarchical Model for Synthetic Data

To generate the synthetic dataset in Sec. 5.1, we employ a probabilistic model (see Fig. 2(a)) designed for generating 8-element vectors $\mathbf{x}_{i}\in\{0,1\}^{8}$, where each element indicates the presence of one produce item.
The generation of each vector is as follows.
First, each of the two, fruits and vegetables, is selected with an independent $50\%$ probability.
If fruits are selected, then each fruit (Apple, Orange, and Banana) is selected with a $50\%$ probability independently. If vegetables are selected, then each of the subcategories, green leaf, and root vegetable, is selected with a $50\%$ probability independently. If root vegetable is selected, then the vegetables Carrot and Beetroot are selected with a $50\%$ probability, independently. If the green leaf is selected, then the vegetables Kale, Spinach, and Lettuce are selected with a $50\%$ probability, independently.
By generating a set of 100 such vectors $\{\mathbf{x}_{i}\}_{i=1}^{100}$, we create a data matrix with the hidden hierarchy of features (words).

### D.3 Accommodating Large Datasets

In our work, we note that constructing the diffusion kernels and their eigenvectors by applying the eigendecomposition to them is the most computationally intensive step.
However, recent methods in diffusion geometry have proposed various techniques to significantly reduce the run time and space complexity of diffusion operators.
Here, we consider the diffusion landmark [Shen and Wu, ] for scalability improvements which reduces the complexity to $O(n^{1+2\tau})$ from the much larger $O(n^{3})$, where $\tau<1$ is the proportion of the landmark set.
We briefly describe the diffusion landmark approach below.

Consider a set of points $\mathcal{A}=\{\mathbf{a}_{j}\}_{j=1}^{m}\subseteq\mathbb{R}^{n}$, let $\mathcal{A}^{\prime}=\{\mathbf{a}_{j}^{\prime}\}_{j=1}^{m^{\prime}}\subseteq
\mathbb{R}^{n}$ be a subset of $\mathcal{A}$, where $1>\tau=\log_{m}(m^{\prime})$.
A landmark-set affinity matrix is constructed by $\mathbf{K}(j,j^{\prime})=\exp(-d^{2}(j,j^{\prime})/\epsilon)$, where $d(j,j^{\prime})$ denotes a suitable distance between the data points $\mathbf{a}_{j}\in\mathcal{A}$ and $\mathbf{a}_{j}^{\prime}\in\mathcal{A}^{\prime}$.
Let $\widetilde{\mathbf{D}}$ be a diagonal matrix, where the diagonal elements are $\widetilde{\mathbf{D}}(j,j)=\mathbf{e}_{j}^{\top}\mathbf{KK}^{\top}\mathbf{1}_
{m}$.
The eigen-structure of an SPD matrix $\mathbf{D}^{-1/2}\widetilde{\mathbf{Q}}\mathbf{D}^{-1/2}$ that is similar to the diffusion operator [Coifman and Hirn, , Katz et al., ] is recovered by applying SVD to the matrix $\widetilde{\mathbf{D}}^{-1/2}\mathbf{K}=\widetilde{\mathbf{U}}\widetilde{
\mathbf{\Lambda}}\widetilde{\mathbf{V}}$, where $\widetilde{\mathbf{Q}}=\mathbf{KK}^{\top}\in\mathbb{R}^{m\times m}$ is the landmark-affinity matrix.
The technique of constructing the diffusion on $\mathcal{A}$ by going through the subset $\mathcal{A}^{\prime}$ and their eigenvectors can be seamlessly integrated into our TWD, enabling the analysis of datasets larger than ten thousand data points, e.g., the word-document data in Sec. 5.2.

## Appendix E Additional Experimental Results

### E.1 Toy Problem

Consider a four-level balanced binary tree $T=(V,E,\mathbf{A})$ rooted at node $v_{1}$, where $V=\{v_{1},\ldots,v_{15}\}$ is the vertex set organized from the root to the leaves of the tree, $E$ is the edge set, and $\mathbf{A}\in\mathbb{R}^{15\times 15}$ is the non-negative edge weights. Consider a set of vectors $\{\mathbf{x}_{i}\in\mathbb{R}^{15}\}_{i=1}^{n}$ such that $\mathbf{x}_{i}$ is a non-negative realization from a normal distribution $\mathcal{N}(\bm{\mu},\mathbf{L}^{\dagger})$ [Rue and Held, ], where $\bm{\mu}=5\cdot\mathbf{1}_{15}\in\mathbb{R}^{15}$ and $\mathbf{L}$ is the graph Laplacian of $T$(^1^11The mean is large compared to the covariance matrix, and the realizations are non-negative with high probability. Realizations with negative elements are rejected.).

The set of vectors $\{\mathbf{x}_{i}\}_{i=1}^{n}$ were generated five times, and each time with a different number of samples $n\in\{3,10,32,10^{2},316,10^{3},3162,10^{4}\}$.
Fig. 3 illustrates the normalized Frobenius norm of the discrepancy between the pairwise TWD matrix obtained by Alg. 2 and the pairwise TWD matrix $\mathtt{TW}(\cdot,\cdot,T)$.
We see that the proposed TWD provides an accurate approximation of the TWD with the tree $T$, even based on a limited sample size.
In addition, as the number of samples increases, the proposed TWD approaches closer to the true TWD, highlighting its efficacy in capturing the hidden feature hierarchy.

Figure: Figure 3: The normalized Frobenius norm of the difference between the proposed TWD and ground truth TWD with different number of samples $n$.
Refer to caption: x3.png

### E.2 Runtime Analysis

In Fig. 4, we present the runtime performance of our TWD in comparison to other TWDs and Wasserstein distances when applied to scRNA-seq and word-document datasets. Although our TWD is not the fastest method, it consistently ranks as the fourth or fifth quickest. It outperforms deep-based methods like HHC-TW and gHHC-TW, as well as pre-trained OT baselines such as GMD for single-cell RNA sequencing datasets and WMD for word-document datasets.
The fastest methods identified are TSWD-1, TR-TW, and MST-TW.
The work in [Le et al., ] reports that the runtime of TSWD is influenced by the number of tree slices used in sampling. Similarly, slice variants SWCTWD and SWQTWD are also affected by the number of slices.
Despite being slightly slower compared to TSWD-1, TR-TW, and MST-TW, our TWD demonstrates significant advantages in terms of classification accuracy, as shown in Tab. LABEL:tab:classification.
We also report the runtime of the OT distance using the embedding distance between features as a ground distance, denoted as OT-ED.
We note that this computational expense is $O(mn^{3}+m^{3}\log m)$.
In contrast, our TWD offers a substantially reduced runtime of $O(m^{1.2})$, thereby demonstrating a significant advantage when handling large datasets.
Note that the scalability of the proposed TWD is enhanced by utilizing diffusion landmarks [Shen and Wu, ], as detailed in App. D.3.
This enhancement allows our TWD to achieve computational performance comparable to TSWD-1, TR-TW, and MST-TW, particularly in datasets with a large number of features, such as the word-document datasets shown in Fig. 4.

Figure: Figure 4: Run time of the proposed TWD and competing TWD and OT distances on scRNA-seq datasets and word-document datasets.
Refer to caption: x4.png

### E.3 Ablation Studies

We note that our method consists of several non-trivial steps.
We provide theoretical justification for our method in Sec. 4.
Specifically, in Thm. 4.2, we show that the TWD $\mathtt{TW}(\cdot,\cdot,T)$ based on the hidden feature tree $T$ can be approximated by our TWD using the decoded tree $B$, where the proposed TWD and $B$ are obtained solely from the data matrix, without access to ground truth $T$.
In addition, in Thm. 4.1, we show that the tree metric of the decoded tree $B$ is bilipschitz equivalent to the tree metric of the hidden tree $T$. These results stem directly from the particular steps of the proposed method and might not exist otherwise.
In the following, we performed a comprehensive empirical ablation study, demonstrating the importance of each component in our method.

#### E.3.1 Euclidean Initial Metric

Fig. 5 presents the classification performance on the scRNA-seq and word-document datasets using the Euclidean distance as an input to Alg. 2 instead of using the distance based on cosine similarity.
The Euclidean distance leads to inferior classification compared to the cosine similarity.
Indeed, as remarked in App. D, the selection of the distance metric influences the performance of our method.
In future work, we plan to comprehensively investigate the use of different metrics.

Figure: Figure 5: The classification accuracy of the proposed TWD using the distance based on cosine similarity and Euclidean distance for scRNA-seq and word-document datasets.
Refer to caption: x5.png

#### E.3.2 TWD by MST and AT

Fig. 6 shows the classification performance of the scRNA-seq and word-document datasets using the TWD with the tree constructed using Minimum Spanning Trees (MST) [Prim, ] and Approximating Tree (AT) [Chepoi et al., ].
To ensure that the tree nodes are contained in a subtree, which is required when computing the TWD in Eq. (1), we set the minimal leaf number as the number of features in MST for computing the TWD between samples.
We see that the TWD approximated by either the MST or the AT is less effective than our TWD, implying that MST and AT do not capture well the underlying tree structure between features compared to our binary tree construction.

Figure: Figure 6: The classification accuracy of our TWD and TWDs computed with MST and AT for scRNA-seq and word-document datasets.
Refer to caption: x6.png

#### E.3.3 Tree Construction by Embedding Distance

**Table 3: Document and single-cell classification accuracy of the proposed TWD and TWDs using embedding distance (ED).**
|  | Real-World Word-Document |  | Real-World scRNA-seq |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
|  | BBCSPORT | TWITTER | CLASSIC | AMAZON |  | Zeisel | CBMC |
| Quadtree-ED | 94.7$\pm$0.3 | 70.3$\pm$1.0 | 94.7$\pm$0.9 | 88.2$\pm$0.7 |  | 81.4$\pm$1.7 | 80.9$\pm$1.3 |
| Flowtree-ED | 95.0$\pm$1.6 | 71.2$\pm$1.2 | 93.2$\pm$0.9 | 90.6$\pm$1.1 |  | 81.0$\pm$1.2 | 81.0$\pm$1.3 |
| WCTWD-ED | 93.1$\pm$2.7 | 68.7$\pm$2.4 | 92.4$\pm$3.1 | 87.9$\pm$1.7 |  | 81.9$\pm$5.2 | 78.9$\pm$2.7 |
| WQTWD-ED | 94.0$\pm$2.1 | 67.5$\pm$2.9 | 93.5$\pm$3.9 | 86.5$\pm$2.1 |  | 80.0$\pm$3.7 | 79.7$\pm$3.6 |
| UltraTree-ED | 94.0$\pm$2.2 | 69.3$\pm$3.0 | 91.4$\pm$2.5 | 87.9$\pm$3.4 |  | 83.1$\pm$2.3 | 80.6$\pm$2.2 |
| TSWD-1-ED | 86.4$\pm$2.1 | 67.3$\pm$1.8 | 92.9$\pm$1.4 | 86.2$\pm$1.3 |  | 79.4$\pm$1.2 | 74.8$\pm$2.2 |
| TSWD-5-ED | 87.3$\pm$1.9 | 69.2$\pm$1.4 | 93.7$\pm$1.4 | 89.8$\pm$0.7 |  | 80.8$\pm$1.2 | 75.4$\pm$1.9 |
| TSWD-10-ED | 88.0$\pm$1.3 | 70.0$\pm$1.1 | 94.2$\pm$0.9 | 90.7$\pm$0.3 |  | 81.9$\pm$1.3 | 76.0$\pm$1.3 |
| SWCTWD-ED | 91.9$\pm$1.5 | 70.7$\pm$1.0 | 94.0$\pm$1.3 | 91.0$\pm$1.5 |  | 81.4$\pm$2.5 | 76.2$\pm$1.5 |
| SWQTWD-ED | 93.8$\pm$1.4 | 71.2$\pm$1.4 | 94.2$\pm$2.6 | 90.3$\pm$1.5 |  | 82.6$\pm$2.7 | 80.4$\pm$2.9 |
| MST-TW-ED | 89.2$\pm$4.4 | 69.1$\pm$1.4 | 89.2$\pm$2.7 | 87.8$\pm$2.1 |  | 81.3$\pm$2.2 | 77.4$\pm$2.7 |
| TR-TW-ED | 90.6$\pm$1.3 | 71.8$\pm$1.3 | 94.5$\pm$1.5 | 90.6$\pm$1.5 |  | 84.2$\pm$1.9 | 80.5$\pm$1.0 |
| HHC-TW-ED | 84.1$\pm$1.3 | 71.2$\pm$2.2 | 91.8$\pm$3.3 | 89.6$\pm$1.5 |  | 80.7$\pm$1.2 | 79.5$\pm$0.9 |
| gHHC-TW-ED | 82.5$\pm$2.3 | 70.8$\pm$1.4 | 91.0$\pm$1.8 | 87.4$\pm$3.1 |  | 80.1$\pm$0.8 | 74.6$\pm$1.7 |
| UltraFit-TW-ED | 84.5$\pm$1.2 | 70.6$\pm$1.5 | 93.0$\pm$1.6 | 89.3$\pm$2.7 |  | 80.6$\pm$2.9 | 78.1$\pm$1.3 |
| Ours | 96.1$\pm$0.4 | 73.4$\pm$0.2 | 96.9$\pm$0.2 | 93.1$\pm$0.4 |  | 89.1$\pm$0.4 | 84.3$\pm$0.3 |

We note that the embedding distance [Lin et al., ] could be simply used for tree construction in existing TWDs [Indyk and Thaper, , Le et al., , Yamada et al., ], graph-based algorithms [Alon et al., ], or linkage methods [Murtagh and Contreras, ] instead of the HD-LCA in the tree decoding (Alg. 1).
However, these generic algorithms, while providing upper bounds on the distortion of the constructed tree distances, do not build a geometrically meaningful tree when given a tree metric [Sonthalia and Gilbert, ], because they do not exploit the hyperbolic geometry of the embedded space or do not have theoretical guarantees to recover the hidden feature hierarchy.

Tab. LABEL:tab:tree_constructed_by_HDD shows the classification accuracy on the word-document and scRNA-seq datasets obtained by the proposed TWD and TWDs using embedding distance as input to construct the feature tree.
We see that our TWD leads to the best empirical performance among all the tested methods.
This result indicates that using HD-LCA, where the information is encoded in the multi-scale hyperbolic LCAs, is important for revealing the hidden feature hierarchy, and therefore, our TWD is more effective and meaningful.

#### E.3.4 Exploring Other Hyperbolic Embedding and the Impact of Curvature

Hyperbolic embedding methods [Chamberlain et al., , Nickel and Kiela, , ] typically assume a fully known or partially known graph of $T$, which in turn is embedded into hyperbolic space.
In Sonthalia and Gilbert [], Lin et al. [], Sala et al. [], Weber et al. [], these methods were shown to inadequately capture the intrinsic hierarchical relation underlying high-dimensional observational data $\mathbf{X}$ (i.e., the node attributes).
Therefore, we opted to use hyperbolic embedding [Lin et al., ], a recent hyperbolic embedding method with theoretical guarantees. This embedding and its accompanying theory enabled us to introduce our theoretical results in Sec. 4.

It is important to note that we could indeed apply traditional hyperbolic embedding methods directly to the affinity matrix $\mathbf{Q}$. From a graph perspective, this is the affinity of a graph, but this graph is not usually a tree. These traditional methods, due to the quantity they are minimizing, will estimate a hyperbolic distance that approximates as best as possible the shortest path distance on this non-tree graph induced by $\mathbf{Q}$. As such, they do not, in general, recover a hyperbolic distance related to the hyperbolic distance of the hidden tree $T$. In contrast, the hyperbolic diffusion embedding approach [Lin et al., ] is designed such that the recovered hyperbolic metric is close to that of the hidden tree $T$, guaranteeing it will be bilipschitz equivalent to it. One of the reasons for the success of this method is that the obtained embedding is not provided by the minimization of an objective function controlled by $\mathbf{Q}$. Instead, Lin [Lin et al., ] designed a non-trivial process that is neither iterative nor based on optimization that computes the embedding directly with theoretical guarantees. We build on the theoretical guarantees of this embedding in our own theoretical contribution in Thm. 4.1.

We here give an intuition as to the reason for the success of the approach from Lin et al. []. After normalizing the affinity matrix $\mathbf{Q}$, the diffusion operator $\mathbf{P}$ is computed. This operator was proven to reveal the underlying manifold. This property is what allows an embedding approach based on the diffusion operator to approximately recover the underlying tree metric $d_{T}$ rather than fit the metric of a non-tree graph constructed from the affinity matrix $\mathbf{Q}$. The diffusion operator reveals the underlying manifold in the following sense [Coifman and Lafon, ].
In the limit of the infinite number of features $m\rightarrow\infty$ and small scale $\epsilon\rightarrow 0$, the operator $\mathbf{P}^{t/\epsilon}$ pointwise converges to the Neumann heat kernel of the underlying manifold $\mathbf{H}_{t}=\exp(-t\Delta)$, where $\Delta$ is the Laplace–Beltrami operator on the manifold. In other words, the diffusion operator is a discrete approximation of the heat kernel on the manifold.
Intuitively, $\mathbf{P}^{t}$ is designed to reveal the local connectivity at diffusion time scale $t$.

Considering multiple timescales on a dyadic grid $\{2^{-k}\}_{k=0}^{K_{c}}$ allows for the association of different diffusion timescales with neighborhoods of varying sizes. The multi-scale embedding of the associated diffusion operators in hyperbolic space generates distances that are exponentially scaled, aligning with the structure of a tree distance. This scaling naturally introduces a hierarchical relationship between the diffusion timescales, enabling the estimation of a latent hierarchical metric. From a theoretical standpoint, it was shown that the $\ell_{1}$ distance in this multi-scale embedding induced by the metric of the product of hyperbolic spaces can estimate the distance $d_{T}$ in the hidden hierarchical space that underlies the observations even when $d_{T}$ is not given or when we do not have access to the explicit $T$ [Lin et al., ]. It is stated under the assumption that $(\mathbf{P}^{t})_{t\in(0,1]}$ is a pointwise approximation of the heat kernel. Such an approximation, mainly in the limit $m\rightarrow\infty$ and $\epsilon\rightarrow 0$, was shown and studied in Coifman and Lafon [], Singer [], Belkin and Niyogi []. In addition, three strong regularity conditions are required: an upper bound on the operator elements, a lower bound on the operator elements, and a Hölder continuity condition. Importantly, the heat kernel satisfies these conditions [Lin et al., ].

In our work, we use the hyperbolic diffusion embedding from Lin et al. [] to present a novel data-driven tree decoding method that operates in high-dimensional hyperbolic spaces. Our method is based on a new definition of the counterpart of the LCA in hyperbolic spaces. Our Thm. 4.1 is built on the multi-scale hyperbolic embedding from Lin et al. [], where we use the theoretical property [Lin et al., ] that the $\ell_{1}$ distance in this embedding can estimate the hidden tree metric $d_{T}$. The difference between Thm. 4.1 and the work from Lin et al. [] is that we construct an explicit tree $B$ whose tree metric approximates the hidden tree metric $d_{T}$. In Lin et al. [], no tree is recovered, and only the hidden tree metric is approximated by the hyperbolic embedding distances.
Our Thm. 4.2 is established based on Thm. 4.1, where the proposed TWD can estimate the TWD associated with the latent feature tree $T$ without access to $T$.
The uniqueness of our approach lies in its ability to work without the need for a pre-known hierarchical graph.

Among the existing hyperbolic models, we chose to work with the Poincaré half-space model because (i) it is natural to represent diffusion times on a dyadic grid in this model, following previous work [Lin et al., ], and (ii) there exists a closed-form solution of the hyperbolic LCA and the HD-LCA, which are the key components in our tree decoding algorithm in Alg. 1.
Since the four standard models of hyperbolic space [Bowditch, ]: the Poincaré disk, Lorentz, Poincaré half-space, and Beltrami-Klein, are equivalent and isometric, it is possible to use the closed-form solution of LCA in the alternative models, e.g., the Lorentz model for numerical stability.
However, we did not encounter any numerical instabilities in our empirical studies using our implementation in the Poincaré half-space model.

We consider hyperbolic space with negative constant curvature of $-1$ because it aligns with the common practice in hyperbolic representation [Nickel and Kiela, , , Sala et al., , Chami et al., ].
Any negative constant curvature can indeed be used and considered as a hyperparameter, e.g., [Chami et al., ].
The modifications that it will require from our algorithm are minimal, and most of the steps can be applied seamlessly in such a space.
We conducted experiments on scRNA-seq data to investigate the role of curvature in our method. Specifically, we vary the curvature and explore the effects of constant negative curvature $-1/\kappa$ where $\kappa>0$. Fig. 7 presents the classification accuracy of the Zeisel and CBMC datasets as a function of the curvature. We see that decreasing the curvature slightly improves the cell classification accuracy. This suggests that while the choice of curvature does have a large impact, our method remains robust across a range of curvatures. In addition, fine-tuning the curvature can enhance performance, and starting with the default choice of $-1/\kappa=-1$ is a reasonable and commonly accepted practice.

Figure: Figure 7: The classification accuracy for the Zeisel and CBMC datasets under different curvatures $-1/\kappa$. Decreasing curvature could slightly improve the cell classification accuracy.
Refer to caption: x7.png

We note that other TW baselines could use other two-dimensional Poincaré representations (e.g., Poincaré embedding (PE) [Nickel and Kiela, ] or hyperbolic VAE [Nagano et al., ]) and decode the hyperbolic embedding into a tree by HHC [Chami et al., ].
We denoted these baselines by PE-HHC-TW and HVAE-HHC-TW, respectively, and present their results in Tab. E.3.5. We see in the table that HHC [Chami et al., ] applied to a PE [Nickel and Kiela, ] or hyperbolic VAE [Nagano et al., ] is less effective than our TWD by a large margin.
This discrepancy is attributed to the sensitivity of the recovery of the hidden feature hierarchy using the competing methods, a finding that is consistent with previous works [Sala et al., , Sonthalia and Gilbert, , Lin et al., ].

#### E.3.5 Robustness Analysis Under Gaussian Noise

To demonstrate the robustness of our method, we conducted experiments on scRNA-seq data to evaluate the resilience of the proposed TWD to the presence of noise. Specifically, we added Gaussian noise to the measurements and examined the effect of varying the noise variance on classification performance.

Fig. 8 presents the cell classification accuracy on Zeisel and CBMC datasets under Gaussian noise perturbations. We see that while adding noise to the measurements hinders the extraction of hidden feature hierarchy and is directly reflected in degraded cell classification performance, our method remains more robust than the competing methods by a large margin. This demonstrates that our approach can provide reliable recovery guarantees even when the ground truth hierarchical features are obscured by noise.

Figure: Figure 8: Assessing the robustness of cell classification under Gaussian noise perturbations. While adding noise to the measurements hinders the extraction of hidden feature hierarchy and is directly reflected in the cell classification performance, our method remains more robust than the competing baselines by a large margin.
Refer to caption: x8.png

**Table 4: Document and single-cell classification accuracy of our TWD and TWDs constructed by other hyperbolic embeddings.**
|  | Real-World Word-Document |  | Real-World scRNA-seq |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
|  | BBCSPORT | TWITTER | CLASSIC | AMAZON |  | Zeisel | CBMC |
| PE-HHC-TW | 87.2$\pm$1.2 | 69.7$\pm$0.9 | 93.2$\pm$1.1 | 86.2$\pm$0.9 |  | 81.6$\pm$0.9 | 74.9$\pm$1.0 |
| HVAE-HHC-TW | 88.4$\pm$2.5 | 64.3$\pm$1.7 | 90.8$\pm$1.5 | 87.3$\pm$1.7 |  | 80.2$\pm$0.9 | 75.1$\pm$1.0 |
| Ours | 96.1$\pm$0.4 | 73.4$\pm$0.2 | 96.9$\pm$0.2 | 93.1$\pm$0.4 |  | 89.1$\pm$0.4 | 84.3$\pm$0.3 |

### E.4 Evaluating Our TWD on Non-Hierarchical Data

We focus on the problem of data with a hidden feature hierarchy as it is of broad interest and prevalent across a wide range of applications. Therefore, the existence of a hidden feature hierarchy is an important assumption in our method.

We tested our TWD on MNIST datasets, where the images are column-stacked. We compared the proposed TWD with the same competing methods as in the paper. The digit classification was performed in the same way as in Sec. 5. Fig. 9 presents the resulting digit classification accuracy. As expected, and in contrast to the word-document and scRNA-seq data, the classification results do not match the state of the art since there are no definitive latent hierarchical structures in this case. Yet, our method still demonstrates good results and outperforms all the other TWD-based methods.

Figure: Figure 9: Digit classification accuracy on MNIST dataset. In contrast to the word-document and scRNA-seq data, in this case, there are no definitive latent hierarchical structures. Yet, our method still demonstrates good results and outperforms the competing TWD methods.
Refer to caption: x9.png

We remark that the state-of-the-art methods of MNIST do not tackle the problem of finding meaningful distances between data samples that incorporate the latent hierarchical structure of the features. Moreover, they usually train on labeled data, whereas the distance we compute is unsupervised. Therefore, comparing the results obtained by our method (and the results obtained by the other TWD-based methods) to the state of the art is not entirely fair.

## Appendix F Additional Remarks

### F.1 Single-scale Hyperbolic Diffusion Binary Tree Decoding

Figure: Algorithm 3 Single-Scale Hyperbolic Binary Tree Decoding

In Alg. 1, the tree decoding is based on the multi-scale embeddings and the corresponding HD-LCAs.
In Alg. 3 we present a similar algorithm based on the single scale hyperbolic embedding $\mathbf{z}_{j}^{k}$ and the corresponding $k$-th hyperbolic LCAs ${\mathbf{z}_{j}^{k}\vee\mathbf{z}_{j^{\prime}}^{k}}$ in Def. 4.1, giving rise to a (single-scale) binary tree $B^{k}$.
Specifically, we merge the pair of features exhibiting the highest similarity at the $k$-th hyperbolic LCA, determined by the smallest value of $\mathtt{proj}(\mathbf{z}_{j}^{k}\vee\mathbf{z}_{j^{\prime}}^{k})$ in Prop. 4.1.
Following Def. 4.2, only one of the hyperbolic LCA relations $[j_{2}\vee j_{3}\sim j_{1}]_{\mathbb{H}^{m+1}}^{k}$, $[j_{1}\vee j_{3}\sim j_{2}]_{\mathbb{H}^{m+1}}^{k}$, or $[j_{1}\vee j_{2}\sim j_{3}]_{\mathbb{H}^{m+1}}^{k}$, holds in $B^{k}$.
The merged nodes are then connected by edges whose weights are assigned by $d_{\mathbb{H}^{m+1}}(\cdot,\cdot)$ between the corresponding points.
It is worth noting that considering $(K_{c}+1)$ trees $\{B^{k}\}_{k=0}^{K_{c}}$ generated by Alg. 3 gives rise to a tree-sliced Wasserstein distance as defined in Prop. 4.5.
In Prop. 4.5, we show that our TWD efficiently captures the tree-sliced information encoded within these multi-scale trees, circumventing the need for their explicit construction.
In addition, we posit that our high-dimensional tree decoding in Alg. 3 for single-scale hyperbolic embedding can be readily applied to other hyperbolic embeddings, e.g., [Nickel and Kiela, , , Nagano et al., ], extending its applicability and utility.
While this extension is technically feasible, it does not carry the theoretical guarantees as those established in Thm. 4.1 and Thm. 4.2.

### F.2 Putting the Features on the Leaves of the Tree

Note that the hyperbolic embedding [Lin et al., ], which is consistent with other common hyperbolic embedding methods [Chamberlain et al., , Nickel and Kiela, , , Sala et al., ], does not restrict all the features to the leaves.
Assigning the features to the leaf nodes is only carried out in the tree decoding (Alg. 1), which is in line with other tree decoding algorithms [Sarkar, , Chami et al., , Indyk and Thaper, ].
Note that our tree decoding algorithm generates a weighted binary tree.
We assert that this does not pose a restriction because any tree can be effectively transformed into a binary tree [Bowditch, , Cormen et al., ].
For instance, consider the flexible tree in Fig. 10 (a):
[sport[soccer,football,basketball]] (where “sports” is a parent node with three children: “basketball”, “soccer”, and “football”). This tree can be accurately modeled as a (weighted)
binary tree as in Fig. 10 (b):
[$c$[$b$[$a$[soccer, football], basketball]], sport] (where
“sports”, “basketball”, “soccer”, and “football” are all leaves, $a,b,c$ are internal nodes, and the edges encode the same relationships as the flexible tree above).

Figure: Figure 10: (a) In a flexible tree, “sports” is a parent node with three children: “basketball”, “soccer”, and “football”. (b) In a binary tree, “sports”, “basketball”, “soccer”, and “football” are all leaves, $a,b,c$ are internal nodes.
Refer to caption: x10.png

### F.3 Differentiability of Our TWD

The proposed TWD can be expressed in a matrix form [Takezawa et al., ] as

$$ $\displaystyle\begin{split}\mathtt{TW}(\mathbf{x}_{i},\mathbf{x}_{i^{\prime}},B )=\sum_{v\in\widetilde{V}}\alpha_{v}\left\lvert\sum_{u\in\Gamma_{B}(v)}\left( \mathbf{x}_{i}(u)-\mathbf{x}_{i^{\prime}}(u)\right)\right\rvert=\left\lVert \mathtt{diag}(\bm{\alpha})\mathbf{R}_{B}\left(\mathbf{x}_{i}-\mathbf{x}_{i^{ \prime}}\right)\right\rVert_{1},\end{split}$ (12) $$

where $\bm{\alpha}\in\mathbb{R}^{M}$ is a weights vector with $\bm{\alpha}(v)=\alpha_{v}$, $\mathbf{R}_{B}\in\{0,1\}^{M\times m}$ is a parent-child relation matrix of the decoded tree $B$, where $\mathbf{R}_{B}(v^{\prime},j)=1$ if node $v^{\prime}$ is the ancestor of the leaf node $j$, and $M=2m+1$ is the total number of the nodes in $B$.

Note that the hyperbolic embedding in Lin et al. [] can be recast as a function mapping from the feature diffusion operator $\mathbf{P}$ to the product manifold of the hyperbolic spaces $\mathcal{M}$, denoted by $f_{1}:\mathbf{P}\rightarrow\mathcal{M}$.
The weight $\alpha_{v}$ in the our TWD can then be expressed by

$$ $\displaystyle\quad g_{1}:\mathcal{M}\times\mathcal{M}\rightarrow\mathbb{R},$ $$

and the weights vector $\bm{\alpha}\in\mathbb{R}^{M}$ in Eq. (12) can be written as $g_{2}:f_{1}(\mathbf{P})\times f_{1}(\mathbf{P})\rightarrow\mathbb{R}^{M}$ by applying the function $g_{1}$ $M$ times.
Let $h_{1}:\mathcal{M}\rightarrow B$ be the hyperbolic diffusion binary tree decoding function in Alg. 1.
The parent-child relation matrix $\mathbf{R}_{B}$ can be obtained by applying the function

$$ $\displaystyle\quad h_{2}:B\rightarrow\{0,1\}^{M\times m}$ $\displaystyle\Rightarrow$ $\displaystyle\quad h_{2}:h_{1}(\mathcal{M})\rightarrow\{0,1\}^{M\times m}$ $\displaystyle\Rightarrow$ $\displaystyle\quad h_{2}:(h_{1}\circ f_{1})(\mathbf{P})\rightarrow\{0,1\}^{M \times m},$ $$

where $\circ$ denotes function composition.
Therefore, the gradient of the our TWD w.r.t. $j$ can be written as

$$ $\nabla\left\lVert\mathtt{diag}(f_{1}(\mathbf{P})\times f_{1}(\mathbf{P}))((h_{ 1}\circ f_{1})(\mathbf{P}))\left(\mathbf{x}_{i}-\mathbf{x}_{i^{\prime}}\right) \right\rVert_{1}.$ $$

This implies that the gradient of our TWD depends on the feature diffusion operator $\mathbf{P}$.
Notably, since the gradient of the feature diffusion operator hinges on the gradient of the Gaussian kernel $\mathbf{Q}_{c}$ (as detailed in Sec. 3), the gradient of the proposed TWD can be expressed only in terms of the gradient of $\mathbf{Q}_{c}$.
Such formulation could be highly beneficial in future applications that require rapid computation of the gradients of the proposed TWD.

### F.4 The Differences between Our TWD and the Embedding Distance [Lin et al., 2023 ]

Our TWD is a TWD between *samples* that incorporates the hidden hierarchical structure of the *features*, while the embedding distance [Lin et al., ] is a hierarchical distance between the samples.
Namely, in our TWD, the features are assumed to have hidden hierarchies, whereas in Lin et al. [], the samples are assumed to have hidden hierarchies.
The work in Lin et al. [] does not involve TWD or OT.

More concretely, in comparison to the work in Lin et al. [], our work presents a new algorithm for tree decoding (Alg. 1) and a corresponding new TWD (Alg. 2). Our work also includes theoretical analysis of these two algorithms:

- 1.
Our TWD recovers the TWD associated with the true hidden feature tree, as stated in Thm. 4.2.
- 2.
Our decoded tree metric represents a form of unsupervised ground metric learning derived from the latent tree underlying the features, as presented in Thm. 4.1.
- 3.
Our TWD captures the tree-sliced information encoded in the multi-scale trees efficiently, as stated in Prop. 4.5.
- 4.
Our TWD can be computed in linear time, enhancing computational efficiency in OT framework, as detailed in Sec. 4 and App. E.
- 5.
Our TWD is differentiable (details in App. F.3), which is an important property that is lacking in existing TW baselines.

In addition, we present favorable empirical results in Sec. 5 compared to several TW and Wasserstein distances.

**Table 5: Document and single-cell classification accuracy of our TWD, embedding distance (ED), and OT distance using embedding distance (OT-ED).**
|  | Real-World Word-Document |  | Real-World scRNA-seq |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
|  | BBCSPORT | TWITTER | CLASSIC | AMAZON |  | Zeisel | CBMC |
| ED | 87.2$\pm$1.2 | 69.7$\pm$0.9 | 93.2$\pm$1.1 | 86.2$\pm$0.9 |  | 84.5$\pm$1.7 | 81.4$\pm$2.8 |
| OT-ED | 95.9$\pm$0.3 | 73.5$\pm$0.3 | 96.7$\pm$0.4 | 93.2$\pm$0.2 |  | 89.0$\pm$0.3 | 84.1$\pm$0.5 |
| Ours | 96.1$\pm$0.4 | 73.4$\pm$0.2 | 96.9$\pm$0.2 | 93.1$\pm$0.4 |  | 89.1$\pm$0.4 | 84.3$\pm$0.3 |

Seemingly, from a metric learning viewpoint, both our TWD and the embedding distance [Lin et al., ] generate distances between the samples.
Since the two are fundamentally different, the comparison can be carried out in two ways: (i) directly compute the embedding distance between the samples (denoted as ED), and (ii) use the embedding distance between the features as the ground distance for the computation of the Wasserstein distance between the samples (denoted as OT-ED).
We present the classification accuracy of the word-document and scRNA-seq datasets of these two baselines in Tab. F.4.
We see a clear empirical advantage of our TWD in classifying documents and scRNA-seq data compared to embedding distance (ED), as our TWD learns the latent hierarchy of the features explicitly by constructing a tree, and then, incorporating it in a meaningful distance between the samples.
We note that while OT-ED performs similarly to our TWD in terms of accuracy (empirically corroborating Thm. 4.2), the computational expense is $O(mn^{3}+m^{3}\log m)$.
In contrast, our proposed TWD offers a substantially reduced runtime of $O(m^{1.2})$, thereby demonstrating a significant advantage when handling large datasets.
This computational advantage is empirically validated and illustrated in Fig. 4.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: I put the suggestions in the weakness.
soundness: 2
strengths: I think this paper is a nice combination of three ideas, 1) from finite data construct a hyperbolic embedding, 2) from points in hyperbolic space construct a tree 3) compute tree wasserstein distance on that tree. The part 1) is borrowed from Lin et al. (2023) and part 3) is the usual tree wasserstein distance, but in 2) algorithms and theoretical works seems to be their own and nice combinations from hyperbolic geometry. I felt this originality and signifcance should be appreciated. Although, my research is not in tree wasserstein distance or hyperbolic geometry with machine learning, so I cannot correctly evaluate originality or significance. I think the paper is is clearly written. I think the paper is in general in good quality, though some theoretical results are a bit skeptical: see weaknesses below.
summary: The paper presents a new Tree-Wasserstein Distance that is specifically designed for a hierarchical data and presents a novel flow matching method. This method can be applied to high dimensional data measured at non-equidistant time points without reducing dimensionality, so that the dynamics of the data are not oversimplified. The method is validated on several datasets and showed the improvement compared to existing methods.
weaknesses: In Theorem 1, the latent tree distance $d_{T}$ and the decided tree metric $d_{B}$ are claimed to be bilipschitz with constants $1/\sqrt{2}$ and $1$. However, I am not sure if the bilipschitz constants are correct. The authors just said that "Specifically, by taking $\alpha\rightarrow 1/2$ and $\beta =2$ in
Section 3 in Leeb & Coifman (2016) and Appendix C in Lin et al. (2023), the bilipschitz constants $c$
and $C$ can be determined by $1/\sqrt{2}$ and $1$ in $d_{H} \simeq d_{T}$, respectively." and did not provide detailed proof for why the constant can be $1/\sqrt{2}$ and $1$. However, proofs of Section 3 in Leeb & Coifman (2016) and Appendix C in Lin et al. (2023) are "all up to constants" and did not trace exact constants.

For example, for the Lipschitz constant $1$, i.e., $d_{B} \geq d_{T}$, it is somewhat necessary to argue that $d_{\mathcal{M}} \leq d_{T}$, where $d_{\mathcal{M}}$ is the distance in Lin et al. (2023). In Lin et al. (2023), the corresponding proof is Proposition C.2, and here the constant $1$ contains several constants as $A$ and $\epsilon$ (different from $epsilon$ in the authors' paper). In particular, the $\epsilon$ is from Lemma 3 in Leeb & Coifman (2016), which is $C_{1}\Psi(1) - C_{2}\Phi(A-1)$ in the corresponding proof. It does not seem possible to simplify everything and conclude that we can take the constant $1$.

I think authors should provide proof for how the bilipschitz constants $1/\sqrt{2}$ and $1$, or they should change the statement of Theorem 1 should be weakened to "there exists some constants $C_1$ and $C_2$ so that $C_1 d_B \leq d_T \leq C_2 d_B$".

Another thing is about the hyperparameter tuning. The suggested algorithm is based on the diffusion operator P in Equation (2) to provide similarity between data points. To compute Equation (2), the user needs to determine the scale parameter $epsilon>0$. Hence, I think providing a guideline for choosing $epsilon$ should be helpful to readers, for e.g., Lin et al. (2023) suggested the median of distances between data points multiplied by a constant.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: * l. 113 V is not defined, should this be [m]?
soundness: 3
strengths: The proposed method is well-motivated and combines an innovative application of previous results (on hyperbolic diffusion embeddings) with novel contributions (tree construction algorithm based on hyperbolic LCAs).

The paper clearly highlights the contributions. The related work and preliminaries sections are clearly written and contain necessary background material together with the appropriate references. The exposition of the method is at an appropriate level of detail and easy to follow once the preliminaries have been digested. The appendix contains further material related to the efficient implementation of the method.

The problem of computing distances between histograms based on a latent tree is well-motivated by practical applications to single-cell sequencing data.
summary: The authors propose an algorithm for constructing a binary tree from data points in the Poincaré half-space. They furthermore propose a method for learning a distance between histograms, assuming a latent hierarchical relation between the bins of the histograms. The method starts with embedding the columns of the data matrix (where each sample/row is a histogram) into Poincaré half-spaces using a combination of diffusion embeddings at multiple scales. In a second step, the tree construction algorithm is applied to these embeddings, resulting in a tree on the bins of the histograms. Finally, they plug this tree into the definition of the Tree Wasserstein distance, resulting in the desired distance between histograms.

For the tree construction algorithm the authors show that, asymptotically in the number of dimensions and scales, it recovers a tree defined by the hyperbolic lowest common ancestor metric on the Poincaré half-space. For the complete method, it is shown that (asymptotically in the number of dimensions and scales) the resulting distance is equivalent within a factor 1/sqrt(2) to TWD using a “ground latent tree” (which is unfortunately left undefined). 

The authors validate the theoretical results using experiments with a synthetic dataset (where “ground latent tree” is a probabilistic graphical model), a text dataset as well as single-cell gene expression dataset.
weaknesses: The paper never gives a definition of the "ground latent tree" used in the statements of Theorems 1 and 2. Theorem 1 references Section 4.1, which however does not seem to contain any formal definitions. The proof of Theorem 1 (via Theorem C.1) references Lin et al. (2023) Theorem 1 for the convergence of the manifold metric to the "hidden tree metric". Without a definition of the "hidden tree" or "ground latent tree", the Theorems and the paper as a whole are incomplete.

At first reading, it was unclear to me what is meant by "features" and "feature hierarchy" and which are the measures to which the Wasserstein distance is eventually applied. The paper states for the first time that it "focuses" on data points that are essentially histograms on page 4, l.172 and only confirms that the method is limited to histograms on page 7, line 372. It would greatly improve the presentation of the paper if the setting was made clear from the start.

I set my rating to weak reject because the paper is missing the definition of the "ground latent tree" used in Theorems 1 and 2. If that definition is added (and consistent with the experiments), I am willing to raise my rating to accept.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: MAJOR COMMENTS

* All comparisons in Table 1 are with other TWD distances. However, I would like to see a comparison to kNN classification accuracy directly in the feature space (R^M), as a baseline. The main point of the article is that inferring latent feature hierarchy is beneficial, so I'd like to see a comparison to a non-TWD approach. Please use the same distance function as elsewhere (I think you used cosine distance?).

* There are A LOT of TWD methods used for comparison in Table 1. Here I got a bit confused (because I am not familiar with this literature). From the introduction to this paper, my impression was that there are no existing TWD methods that estimate the latent feature tree. So what do all these TWD methods in Table 1 do instead? The description in section 5.1 suggests that you estimate the tree in many different ways and then use TWD with that. Why do all these methods of tree estimation work so much worse than yours? Maybe you need a part in the Introduction / Related work that would introduce these methods and explain what you are doing differently.

* The performance of kNN classification on the X matrix directly for RNAseq or text data will likely strongly depend on the normalization/transformation of the data. In RNAseq literature, row-normalized counts are usually log1p-transformed. In the NLP literature, words counts are usually transformed into TF-IDF with log-scaling. Operating on raw row-normalized counts without any transformation should perform much worse. It wasn't clear to me what data representation you use as input to your TWD benchmarking.


MINOR COMMENTS

* line 39: "due to its exponential growth" -- unclear formulation, whose growth?

* line 113: What is V? This notation wasn't introduced.

* line 126: Usually normalized Laplacian is defined with D^{-1/2} on both sides instead of D^{-1}. So I would expect to see Q hat defined as D^{-1/2} Q D^{-1/2}. Is this a typo? If not, why did you use D^{-1}? You say "following Coifman & Lafon 2006", but I am pretty sure they used D^{-1/2}. Furthermore, I would expect P to be defined directly as D^{-1} Q, instead of \hat D^{-1} \hat Q. Can you clarify?
soundness: 2
strengths: The paper considers an interesting problem setting, and suggests what seems to be a sensible algorithm. Technically the paper is well done, with high-quality figures and reasonably clear mathematical exposition.
summary: In their paper "Tree-Wasserstein distance for high dimensional data with a latent feature hierarchy", the authors consider the problem of constructing a NxN pairwise distance matrix between samples, given a data matrix X of NxM size and under assumption that the M features show some unknown hierarchical organization. The authors suggest a three-stage approach: (1) embed the M features into a hyperbolic space using diffusion eigenvectors; (2) construct a binary tree of the M features from the embedding; (3) use tree-Wasserstein distance between samples given the feature tree. In experiments, the authors argue that resulting distance yields superior kNN classification accuracy in bag-of-words and scRNA-seq datasets.
weaknesses: The experimental evaluation is missing naive baselines, and in general lacks details.

Overall, I am giving a borderline score which I would be willing to revise upwards (or downwards) based on authors' clarifications regarding the experiments.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: 1. Can you clarify the construction of the diffusion operator as referred to above?
2. Does the method require assuming the feature tree can be isometrically embedded in Euclidean space?
soundness: 3
strengths: Originality: The idea of learning an explicit hierarchy of features (instead of samples) is a clever one, and it lends itself to applying the TWD when data points are represented as distributions over features. 

Quality: The method appears to beat existing approaches as measured by tree reconstruction quality on synthetic data and nearest neighbor classification accuracy on real data. 

Clarity: The paper is well-written and clear and uses consistent notation.

Significance: Metric learning for data where features exhibit hierarchical structure is important, as this is a common phenomenon. So any effort to (a) recover the feature hierarchy, and (b) extract a metric that reflects that hierarchy is significant.
summary: This paper presents an approach for learning distances between samples when (a) features are known to exhibit hierarchical structure, and (b) data points are represented as distributions over features. The paper leverages previous work on multi-scale diffusion embeddings to obtain hyperbolic representations of features (not samples) and uses single-linkage clustering to build a tree of said features. The tree-Wasserstein distance is used to compute distances between samples, which leverages (a) the fact that data is represented as a distribution over features and (b) the learned feature hierarchy.
weaknesses: 1. The method assumes features are necessarily leaves in the ground-truth feature hierarchy. I can imagine simple situations (in language data) where this may not hold - for example, if 'fruit' was included as a word in the synthetic example provided in the paper, this assumption would not hold, it seems.

2. The construction of the feature diffusion operator could be explained in more detail. Primarily, the computation of the distance between features d(i,j) (which I assume to be the Euclidean distance between X[:, i] and X[:, j]) was not explicitly stated from what I can see. Can the authors add an explicit description of how these distances are computed?

3. On line 150 the LCA relaxation was referenced but never directly defined. Please add a restatement of the definition from Wang et. al in the body of the paper or add a description to the appendix.

4. The method doesn't address cases where the underlying feature tree cannot be isometrically embedded in Euclidean space. Can discussion be added to address this? For example, how would this affect the performance of the proposed approach?

5. No intuition is given for why a single-linkage tree construction approach performs better than other tree constructions given the hyperbolic embeddings. Could some intuition or theoretical justification be added?

Meta-review: The authors focus on estimating the hierarchical structure (i.e., tree structure) given observations from latent hierarchical features, and leverage the tree-Wasserstein for a fast computation for measures with hierarchical structure on supports. The authors propose to embed features into a multi-scale hyperbolic space and algorithmic approach to estimate the tree structure from embedded features. The authors show advantages of estimating tree structure from observed data (with a latent hierarchical structure) for tree Wasserstein on document and single-cell RNA sequencing datasets. All the Reviewers agree that it is a good submission for ICLR'2025. However, it is better to clarify whether the proposed approach can be applied for general applications which may or may not have latent hierarchical structure as in standard tree-sliced-Wasserstein setting (Le et al., 2019) (i.e., whether the proposed approach can be applied for applications without latent hierarchical structure? and/or how to know whether a given dataset in practice has prerequisite latent hierarchical structure or not?). We urge the Authors to incorporate the Reviewers' comments and discussions in the rebuttal into the updated version.

Decision: Accept (Poster)