Tree-Wasserstein Distance for High Dimensional Data with a Latent Feature Hierarchy
Authors: Ya-Wei Eileen Lin
Ronald R. Coifman
Gal Mishne
Ronen Talmon
Venue: ICLR 2025 Poster (2025)
Keywords: diffusion geometry
hyperbolic geometry
tree-Wasserstein distance
high-dimensional hyperbolic tree decoding
Primary area: learning on graphs and other geometries & topologies
TL;DR: We present a new Tree-Wasserstein distance using diffusion geometry and hyperbolic geometry for high-dimensional data with latent feature hierarchy.

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.

Published as a conference paper at ICLR 2025
FHEE-WASSESSTER DISTANCE FOR HIGH DIMEN- SIONAL DATA WITH A LATENT FEATURE HIERARCHY
Ya- Wei Eileen Lin' Ronald R. Coifmann Gal Mishne
Ronen Talmon'
1 Viterbi Faculty of Electrical and Computer Engineering, Technion
2 Department of Mathematics, Yale University 3 Halicioglu Data Science Institute, University of California San Diego
ABSTRACT
Finding meaningful distances between high-dimensional data samples is impor- tant scientific task. To this end, we propose a new tree- wre--Wssnnnttin distance (TWD) for high-dimensional data with two key aspects. First, our TWD is specifically de- signed for data with a latent feature hierarchy, i.e., the features lie in 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 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 word-document and single-cell RNA-sequencing datasets, demonstrating its advantages over existing TWDs and methods based on pre-trained models.
INTRODUCTION
High-dimensional datasets with latent feature hierarchy (Fefferman et al., 2016) are prevalent across various applications, involving data, e.g., genomic sequences (Tanay & Regev, 2017), text and images (Krishna et al., 2017), and movie user ratings (Bennett et al., 2007). In this setting, the assumption is that there is an unknown underlying hierarchical structure to the data features, which is not directly observable. For example, in word-document data, where samples are documents and features are words, the hierarchy underlying features (words) is often assumed because words are naturally organized as nodes in a tree, each connected to others through various hierarchical linguistic & Arenas, 2010). A key component of meaningful analyses of such
relationships (Borge-Holthodee data is finding distances between data samples that take into account the underlying hierarchical struc- ture of features. As opposed to relying on generic metrics, e.g., the Fuclidean or correlation distances, such meaningful distances serve as a fundamental building block in a broad range of downstream tasks, e.g., classification and clustering, often leading to significant performance improvement.
Recently, hyperbolic geometry (Ratcliffe et al., 1994) has gained prominence in hierarchical repre- sentation learning (Chamberlain et al., 2017; Nickel & Kiela, 2017) because the lengths of geodesic paths in hyperbolic spaces grow exponentially with the radius (Sarkar, 201 1). a property that natu- rally mirrors the exponential growth of the number of nodes in hierarchical structures as the depth increases. 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., 2018). 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 & Thaper, 2003) for this purpose, where we model samples distributions supported on a latent hierarchical structure. We propose a two-step approach. In the first step. we embed features into continuous hyperbolic spaces (Bowditch, 2007) utilizing diffusion geometry (Coifman & Lafon, 2006) to approximate the hidden

---

Published as a conference paper at ICLR 2025
hierarchical metric underlying features. In the second step, we construct tree that represents the latent feature hierarchy in a bottom-up manner from the hyperbolic embedding. The tree construction is based novel notion of 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 new TWD for high-dimensional data with latent feature hierarchy, based on diffusion geometry and hyperbolic geometry. We remark that in Lin al. (2023), integrating hyperbolic and diffusion geometries was shown be more effective in recovering latent hierarchies than existing hyperbolic embedding methods (Nickel & Kiela, 2017) 7) (see App. E.3.4). However, the work in Lin al. (2023) considers implicitly latent hierarchies of the samples, whereas here, we explicitly reveal the latent hierarchical structure of the features through a tree, which then utilized for finding a meaningful TWD between the samples (see App. F.4). Using TWD entails two useful attributes. First, the Wasserstein distance (Monge, 1781; Kantorovich, 1942) enables us to integrate feature relations into a meaningful distance between samples. Second, we use tree to represent the latent feature hierarchy, learn this tree based on the data, and then, naturally incorporate 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 based on the Euclidean ground metric.
From a theoretical standpoint, we show that our TWD provably recovers the Wasserstein distance based on the ground metric induced by the 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-sequeniing datasets. Specifically, we show that our TWD leads to superior classification performance compared to existing "WD-based methods (Indyk & Thaper, 2003; Backurs et al., 2020; Le et al., 2019; Yamada et al., 2022) and two application-spectne methods based on pre-trained models (WMD (Kusner et al., 2015) and GMD (Bellazzi et al., 2021)).
Our main contributions are as follows. (i) We present a new TWD for data with a hidden feature hierarchy, which provably recovers the Wasserstein distance 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, 2009) is 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, 1781; Kantorovich, 1942) with linear programming (Peyré et al., 2019) in cubic complexity. Various approaches, such as adding entropic regularization (Cuturi, 2013), using wavelet functions (Shird- honkar & Jacobs, 2008; Gavish al., 2010), applying Sirttroovic--uubistte 2016; Arjovsky et al., 2017), and computing the average Wasserstein distance across one- one-dimensional projections (Rabin et al., 2012), have been developed to reduce the computational complexity.
dual (Leeb & Coifman,
Tree- Free-Wasersstenn Distance Another popular alternative to reduce the computational complexity of the Wasserstein distance is the tree- wassesssrrnnnin distance (TWD) (Indyk & Thaper, 2003), which can be computed in linear time, enabling efficient comparisons of large datasets. Existing TWD methods (Leeb, 2018; Sato al., 2020; Takezawa al., 2021: Yamada et al., 2022) attempt to approximate the Wasserstein distance based on the Euclidean ground metric with tree metric by constructing a tree using the standard metric in the Euclidean ambient space (see App. B for more details). Our approach has a different goal. We assume that the features lie in a latent hierarchical space, infer a tree that represents this space, and compute a TWD based on the inferred tree to define meaningful distance between the samples, incorporating the latent feature hierarchy.
Hyperbolic Representation Learning. While effective in computing embeddings and distances for hierarchical data samples, current hyperbolic representation learning methods (Chamberlain et al., 2017; Nickel & Kiela, 2017; 2018) are usually given a hierarchical graph, e.g., a tree, and do not include decoding tree from observational data. One exception is the work in Chami et al. (2020), which introduces a method for decoding a tree in a two-dimensional Poincaré disk for hierarchical clustering (Murtagh& & Contreras, 2012), by minimizing the continuous Dasgupta's 's cost (Dasgupta,
2

---

Published as a conference paper at ICLR 2025
2016). However, this approach does not approximate the TWD between samples with latent feature hierarchy because it does not recover the tree underlying the features.
PRELIMINARIES
Notation. For N, we denote [m] [l...., "......,..... my. Let X E Raxm be data matrix with rows (samples) and m columns (features). We denote by X,, the i-th row (sample) and X. the j-th column (feature). Let d = d' denote the bilipschitz equivalence between metries d and d' defined on the set V. That is, there exist constants C, C 0 s.t. cd(y/1 V2) I d' (y1, U2) Σ Cd(y), y2) Vyn 32 E V.
Tree- Wassersssnntnn Distance. Consider tree T = (V, E, A) with leaves rooted node 1 w.l.ogg. where V is the vertex set, E' is the edge set and A E | contains edge weights. The tree metric dr between two nodes is the length of the shortest path on T. Let (v)) be the set of nodes in the subtree of T rooted at V. For any E V, there exists unique node v, which is the parent of u. We denote by w., = dr u) the length of the edge between any node u and its parent v (see App. A.4 for additional background on trees). The TWD (Indyk & Thaper, 2003; Evans & Matsen, 2012; Le et al., 2019) between two distributions P1, 442 E Ik" supported on the tree T' is defined by
TW( 12, (2,
Σ We VEV
Σ NETr(e)
(ur(4)-Hu(00) - -22()))
(1)
Diffusion Geometry. Consider set of points (a, E RR" 1j= assumed to lie on manifold embedded in R" Let Q Rm be an affinity matrix, given by Q3 = exp (-d'(j,j')/c)) where d(j. is some suitable distance between a, and ay,, and e > is a scale parameter, often chosen as the median pairwise distance among data points scaled by a constant factor (Ding & Wu, 2020) (see App. for hyperparameter tuning). Following Coifman & Lafon (2006), we consider the
affinity matrix Q D -DD to address the effects of non-uniform data sampling, where D is diagonal with entries Djj = Σj O3y', and the diffusion operator P, defined by
P=QD QD-1, where D1₁ = =E, Qif-
The matrix Q can be viewed edge weights of a graph G ([m]. E, Q). Since the diffusion operator P is chain on the graph G. Therefore, P can propagate densities between nodes. Specifically, the vector P = P'e, E R is the propagated density after diffusion time t E R of density concentrated at node where ej I is the indicator vector of the j-th node, p)
it can be viewed as a transition probability matrix of Markov
and ||D|||, =1
Poincaré Half-Space Model of Hyperbolic Geometry. Hyperbolic geometry is Riemannian geometry with constant negative curvature -1 (Beardon, 2012). We utilize the m-dimensional Poincaré half-space model H defined by H fx E IT' "='' (x(m >> 0} with the Riemannian metric tensor d (dx² (1) + dx² +... + dx² (x2 //m 2 (m). The Riemannian distance for any two hyperbolic point x, y E ! is defined by d (x. y) sinh " yll2 (2))(m))(m))).
Rooted Binary Tree. rooted and balanced binary tree with m leaves is a tree whose root node has degree 2, leaves have degree 1, and m internal nodes have degree 3. For any two leaves j. let j V j' and I F V j') denote their lowest common ancestor (LCA) and the set of leaves in the subtree of B rooted j V respectively. Given any three leaves J1. Ja. the LCA relation [jz jaa ~ n holds if I ((22 V ja) C Bii V j. 12 V ja) (Wang & Wang, 2018), 8), as shown in Fig. 1(a). For binary trees, only one of [i j2 ja]n, (i V ja f2]n, n, or (ja j2 fi) a can hold.
4 PROPOSED METHOD
4.1 PROBLEM FORMULATION
Consider a data matrix X E Raxm, whose n rows consist of samples and m columns consist of features. Suppose the features (X. 1X........ .......,..ll m C H. C R" lie underlying manifold H, which a complete and simply connected Riemannian manifold with negative curvature embedded in a high- dimensional ambient space R2' with geodesic distance dz. We view (H, dx) as hidden hierarchical
3

---

Published as a conference paper at ICLR 2025
metric space and consider weighted tree T E,, as its discrete approximation in the following sense. The tree node j E [m]] is associated with the j-th feature, E is the edge set, and A E Im is the weight matrix of the tree edges, defined such that the tree distance dr (i. between two nodes j and j', i.e., the length of the shortest path on 7, coincides with the geodesic distance (i)())") between the features j and j'. Therefore, we refer the hidden tree distance dr as the ground truth distance. Note that the assumption of an underlying manifold embedded in an ambient space I" is widely used in studies of data manifolds and manifold learning (Shnitzer et al., 2022; Katz et al., 2020; Tong et al., 2021; Huguet et al., 2022; Kapušniak et al., 2024), where the manifold typically underlies samples and graphs are viewed as discretizations of the underlying Riemannian manifold. In contrast, we consider the manifold underlying, features. Specifically, we focus on specification of the manifold to a hierarchical space H of features, and accordingly, a specification of the graph to a tree T. Our aim is to construct tree B from the data X, such that the TWD between samples using our learned tree B is (bilipschitzz) equivalent to the Wasserstein distance using the ground truth latent tree distance dr. To allow this construction of B, we assume that local affinities between (node) features { {X:. 1 that are close in I." embody information on the hidden hierarchical structure (see Sec. 4.4 for formal description of this assumption).
Following a large body of work (Villani, 2009; Cuturi, 2013; Courty et al., 2017; Solomon et al., 2015; Yurochkin al., 2019; Alvarez-Melis & Jaakkola, 2018), we focus on samples X;,- R with positive entries that can be normalized into discrete histograms x; = natural when the distribution over features is more important than the total mass. However, we note that our method is also applicable to data normalized into a histogram using other techniques.
which is
4.2 HYPERBOLIC DIFFUSION EMBEDDING
In our approach, we first construct hyperbolic representation following Lin et al. (2023). We briefly describe the steps of this construction below. The feature diffusion operator E mm is built according to Eq. (2), where d(j.j'") is an initial distance between the features j and j' embedded in the ambient space R In Sec. 5, we compute the initial distance using the cosine similarity in the X. I where is the dot product. We refer to App. D for further discussion on the initial distance metric, the scale parameter e, and the kernel type. Then, we construct the multi-scale diffusion densities Φ = =p² e; E R for each feature j E [m]] considering series of diffusion time steps on dyadic grid i. = 2- for k E z 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. + 2 = {(b))) 22/22-2TT Hm+1 based on the exponential growth of the Poincaré
ambient space, i.c., d'j = XX
half-space, where 4bt = Vo E R is the element-wise square root of the diffusion density Φy- The multi-scale embedding of the corresponding diffusion operators in hyperbolic space (Lin et al., 2023) results in distances that are exponentially scaled. This scaling aligns with the structure of a tree distance, effectively establishing natural hierarchical relation between different diffusion timescales. While it is well-established that diffusion geometry effectively recovers the underlying manifold (Coifman & Lafon, 2006), recent research by Lin et al. (2023) 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. (2017); Nickel & Kiela (2017)) could potentially be considered. However, these methods usually assume (partially) known graph T instead of high-dimensional observational data X, and therefore, we do not opt to use them (see App. E.3.4 for more details).
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 H and paths on trees. We define the following hyperbolic lowest common ancestor (LCA).
Definition 1 (Hyperbolic LCA). The hyperbolic LCA of any two points z zt. E Hm++ at the k-th scale is defined by their Fréchet mean (Fréchet, 1948) as 2 v z/z
arg min (iZ.m 12 (2, (2) + de 12 22)
aEHm+1
4

---

Published as a conference paper at ICLR 2025
R50
R50
H
H
% 10 projot vighT
% = prooj0 10 v "b)'' Yy
V
of. pTT.pr vitiT
V.
% - 10 (0',20 v valii'
J
% 10', proojot 201
Rs = 10 proj O8) 2012 A
R
(a)
(b)
(c)
Figure 1; (a) The LCA relation of J2. on tree B. (b) The hyperbolic LCA relation [i2 V ja ~ Jilg-+t for z 2 z E 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 p-oj(.) signifies the parent-child relation. (c) The HD-LCA hj defined as the Riemannian mean Ke of the orthogonal projections {o km- incorporating the multi-scale hyperbolic LCAs.
in Def. has a closed-form solution, given by z VzJ. Proposition 1. The hyperbolic LCA 2) Vzz z
1 (ab) + 2bt)) (z) v 2 where (roj v zz..) [ ( 2b.)) T 2
The proof of Prop. is in App. C. Note that proj V zz))) is the orthogonal projection of the hyperbolic LCA onto the (m-+1) 1-axis R in H=*+1 which is the "radius" of the geodesic connecting 25 and 2/., as depicted in Fig. (b). We assert that this orthogonal projection can be used identify the hierarchical relation at the k-th scale. This argument made formal in the following.
Definition 2 (Hyperbolic LCA Relation). Given any three points 2),, 2 2/ E [HHm+1 at the k-th scale, the relation [12 V js ~ -i holds if proj V Vzf,,) = min(proj V Vzlz), projj(z V 2 ). proj (z), V )
Proposition 2. For any h S k2. 2 (b2-k1) ~ dgm+i (z) to V
ko
// // (z) z zz
The proof of Prop. 2 is in App. C. Prop. 2 implies that the hyperbolic LCAs at different scales exhibit exponential growths, which 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=== where K2 E Z
Definition 3 (HD-LCA), The HD-LCA fe [m] with Ke scales is defined by the Fréchet mean hg.5' == arg min Ek= dgm=+i (h, y. where of. = =OT,, (z)) V vz )T E H hEAm+1
Proposition 3. The HD-LCA his in Def. 3 has a closed-form solution, given by has' [O.....0,, I T, where 03.5" = (Kel1) Vproj(z)) v 24.)) ... ...oj(z V vz3e).
Ke).
The proof of Prop. 3 is in App. C. Note that the HD-LCCA in Def. involves two applications of the geodesic paths in Hm+1; (i) to establish the hyperbolic LCA at each scale k by connecting two points on the same plane in 12++1 and (ii) to merge these LCAs across multiple levels along the (m + 1)-axis R>0 in H as in Fig. I(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. (2023), where it was shown that the distance between the embedding j + [<]]] ..... (z) K2 TIT
M = H x ... x m1 in product manifold of Poincaré half-spaces, given by dm(lij') M sinh -1 ((2 (2-22/22 k/2+1 |z 2/. [[)). recovers d''j.j"). However, (Lin et al., 2023) considers a latent hierarchical structure of the samples. is fundamentally different than the problem we consider in Sec. 4, 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.
5

---

Published as a conference paper at ICLR 2025
Note that the closed-form value "ر,ز' in Prop. 3 can be interpreted as the Riemannian average of the multi-scale hyperbolic LCA relations in Def. 1, resulting in the following HD-LCA relation. Definition 4 (HD-LCA Relation). Given any three features J2. E [m], the relation [52 V js min (ajndz jas Jj2.ss
jilM holds if jj.da
Algorithm HD Binary Tree Decoding
Algorithm 2 TWD with a latent feature hierarchy
Input: Embedding {(z)...., fzth } T-k- Output: B = (V, F, A) with m leaf nodes function HD_BTTT (00_BT((2 ((
Input: Data matrix X E R feature diffu- sion operator P. and maximal scale Ke Output: TWD W E R***
B <- leaves (j) j E [m])
function WWW__laeett_aee (ie K.)
for j.j' E [m] do
UAV -iig (P)
for k E (0, 11....... ..... Ko} do proj(z] 2
for k E {0, 1, 1,...... Ke) do
5. - 10
P + A
h55 t arg min Σke pir (h, ) heHm+1
for j E [m] do
x [V(UpL )e,, 24/2-22 20/22
S = {(ij")) sorted by hy.y (m + 1)} for (j.j") ) E do
- HD (T()
(2m 1ike
if and j' are not in the same subtree I, f internal node for node j I f internal node for node j' add an internal node for I, and Zy assign the geodesic edge weight
for i, [n| do
W W
Σ Oh
Σ (x:(u) - (x))
VEV
mET
return
return W
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 binary tree (Bowditch, 2007; Cormen et al., 2022), we suggest to construct 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, 2011) (see App. F.2 for more details). This is not only efficient but also aligns with established TWD (Indyk & Thaper, 2003; Backurs et al., 2020; Le et al., 2019). Then, pairs of features are merged iteratively by their HD-LCA (Def. 3). Specifically, the pair of features that exhibits the highest similarity in the HD-LCA (i.e., the smallest value of ajzz" in Prop. 3) is merged. Note that only one of the HD-LCA relations [i2 V ja- ~ Ji]M, [i V ja- ~ J2] or [ir V V.j2 ~ ja]M holds in B. Once the nodes are merged, the weight of the edge assigned with the induced f distance between the respective embeddings. This iterative process constructs the binary tree B in a bottom-up manner, similar to single linkage clustering (Gower & Ross, 1969), with the measure of similarity determined by their HD-LCCA 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 al., 2023) could be used in existing graph-based algorithms (Alon et al., 1995) or linkage methods (Murtagh & Contreras, 2012) 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 geometrically meaningful tree when given a tree metric (Sonthalia & Gilbert, 2020). because they do not exploit the hyperbolic geometry of the embedded space. Indeed, empirically shown in App. E, constructing 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. 3 is 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. Let r be the root node of B obtained in Alg. and z, be the corresponding embedding in M. The HDD-LCA depth is equivalent the Gromov product (Gromov, 1987), i.e., d (z( hj.) ~ (i- J") where (ir (dn(f,r) + dn(j .) d d')))

---

Published as a conference paper at ICLR 2025
The proof of Prop. 4 is in App. C. Prop. shows that the HD-LCA depth in the continuous embedding space is the tree depth in the decoded tree B.
Theorem 1. For sufficiently large Ke and m, there exist some constants C1. C2 > O such that (idr(ji:32) I (j) I C2dr (j1, f2) for J1. E. [mm] where dg is the decoded tree metric dg in Alg. 1, T the ground truth latent tree assumed from Sec, 4.1, and dr is the hidden tree metric between features defined the length of the shortest path T (i.., the geodesic distances).
The proof of Thm. 1 is in App. C. The validity of Thm. 1 relies the assumption that the local Gaussian affinities between features allow us to reveal the hidden hierarchy. Formally, in the proof of Thm. 1, assume that the discrete diffusion operator P constructed based on the local Gaussian affinities between features in the ambient space R", in the limit of m + oo and e - O, converges pointwise to the heat kernel of the underlying hierarchical metric space (H, dz). While T' is hidden, the diffusion operator P is accessible and used to construct another tree, B, such that, the tree distances of the inferred tree dg and the hidden tree dr are equivalent. This result is stated in Thm. 1. In addition, our HD-LCA is 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. (2020), where a continuous LCA in a single 2D) Poincaré disk was proposed for hierarchical clustering, and a tree was constructed by minimizing the Dasgupta's cost (Dasgupta, 2016). However, their method does not recover the latent feature hierarchy, as shown in Sec. 5,
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 digh-dimensional samples with latent feature hierarchy as follows.
Definition 5 (TWD for High- -igeeniimnii Data with Latent Feature Hierarchy.). high-dimensional samples with a latent feature hierarchy is defined as
The TWD for
TW(x, Xv, B)=
Σ ae
(xx(u) Xx xy( (E)
Σ
vEV
where B is the decoded tree in Alg. a,, = d v²), and node v' is the parent of node v.
Theorem 2. For sufficiently large Kg and m, there exist some constant Č1, Ö2 > O such that T,TW(x;, Xv, T) & <W(x;, Xp, , B) S xx,, T) for all samples i, E [m].
The proof of Thm. 2 is in App. C. Thm. 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 unsupervised Wasserstein metric learning between data samples with the latent feature hierarchy. In addition, our intermediate result in Thm. noteworthy, as it represents form of unsupervised ground metric learning (Cuturi& & Avis, 2014) 4) derived from the latent tree underlying the features.
We summarize the computation of our method. Notably, the input the algorithm is the observational data matrix X alone (i.e., the tree node attributes). and no prior knowledge about the latent feature tree graph 7' is given. First, we build the feature diffusion operator by Eq. (2). Then, for each scale k. we embed the features into -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. 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 kemee limits the applicability to small number of features. However, using (Shen & Wu, 2022), the proposed TWD can be computed in (m.-22 compared to the naïve implementation that requires Omm" m 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, 1781; Kantorovich, 1942) and the Wasserstein and TWDs (Villani, 2009; Cuturi, 2013; Indyk & Thaper, 2003; Backurs et al., 2020; Le et al., 2019; Takezawa al., 2021; Yamada et al., 2022), restricts its applicability. However, a large body of work (Villami, 2009; Cuturi, 2013; Courty et al., 2017; Solomon et al., 2015;
Yurochkin et al., 2019, 19;; -varez-Melis & Jaakkola, 2018) illustrates that this setting is relevant in
7

---

Published as a conference paper at ICLR 2025
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 R", which is in line with existing tree construction methods (Alon et al., 1995; Murtagh & Contreras, 2012). Note that the problem of finding distances between samples with feature hierarchy and possibly with negative entries has been addressed in Gavish al. (2010); Ankenman (2014); Mishne al. (2016; 2017) 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., 2019). Computing the TSWD consists of constructing multiple random trees of different depths from X. Then, the TSWD is given by the average of the TWDs corresponding to these trees. Here, considering only a single scale in Alg. (as detailed App. F.1) yields the decoded tree B* Similarly to the TSWD, we show that our TWD can be recast as the sum of the TWDS corresponding to B multiple scales k O,..., Ke.
Proposition 5. The TWD in Eq. (3) is bilipschitz equivalent to the sum of the multi-scale hyperbolic TWDs Xx
x,, xv, Bb) for
0, 0,,,.... Ke, Le., TW(X,, Xv, B) = Ek= TW(X;, Xx, B²).
The proof of Prop. 5 is in App. C. Prop. 5 implies that our TWD captures the (tree-sliced) information in the multi-scale trees B efficiently as it circumvents their explicit construction.
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 Xy, T)) from the data matrix X (i.e., the node attributes from the hidden T), which we achieve by the proposed TWD TW(x; Xy, B) in Eq. (3) as stated in Thm. 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 given ground tree metric between the features. While the complexity of this computation is typically O(m log m) (Pele& & Werman, 2009; Bonneel et al., 2011), we can approximate by the proposed TWD in linear time, as we demonstrate in the runtime analysis in App. E. Third, our method applies to general applications with or without latent hierarchical structure (see App. E.4). 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.
EXPERIMENTAL RESULTS
We apply the proposed method to a synthetic example and to word-document and single-cell RNA- sequencing (scRNA-seq) data. We use Word2vec (Mikolov et al., 2013) word embedding vectors for the word-document data and Gene2vec (Du et al., 2019) as gene embedding vectors in scRNA-seq experiments. The implementation details are in App. D Additional experiments are in App. E.
5.1 SYNTHETIC DERD--DT DATA EXAMPLE
We generate a synthetic dataset with latent feature hierarchy follows. The generated dataset is designed to resemble a word-document data matrix 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 the latent feature hicrarchy. 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 & Thaper, 2003) that constructs a random tree by a recursive division of hypercubes, (ii) Flowtree (Backurs et al., 2020) that computes the optimal flow of Quadtree, (iii) tree-sliced Wasserstein distance (TSWD) (Le et al., 2019) 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., 2024) that constructs an ultrametric tree (Johnson, 1967) by minimizing OT regression cost, (v) weighted cluster TWD (WCTWD) (Yamada et al., 2022), (vi) weighted Quadtree TWD (WQTWD) (Yamada et al., 2022), and their sliced variants (vii) SWCTWD and (viii SWQTWD (Yamada et al., 2022). We also compare to TWDs where the trees constructed by (i) MST (Prim, 1957), (ii) Tree Representation (TR) (Sonthalia & Gilbert, 2020) through a divide-and-conquee approach, (iii) gradicnt- based hierarchical clustering (HC) in hyperbolic space (gHHC) (Monath al., 2019), (iv) gradient-based Ultrametric Fitting (UttraFit)
8

---

Published as a conference paper at ICLR 2025
Ours
Quadtree
TSWD
WCTWD
UltmTree
Ar
Sitrr
MST-IWD
gHHC-IWD
UltaFit-TW
TR-TWD
HHC-TWD
Apple
Orangee
Feram
Kule
Lemoe
Caror
[0.1]
o
Kak
Kake
Spead
Can
to
Qut
Anke
Cam
Spead
tase
Cone
kike
beane
Orra
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 red, and the internal nodes are green.
(Chierchia & Perret, 2019), and (v) HC by hyperbolic Dasgupta's cost (HHC) (Chami et al., 2020). We refer to App. B for a detailed description of baselines and highlight the novel aspects of our TWD.
Fig. 2(b) presents the trees constructed using different methods. Our TWD constructs 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 (INNN) where the documents are labeled by the presence of fruits (with probability 0.5, creating two balanced classes), presented in Tab. 1. The depicted classification accuracy 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. see that our TWD demonstrates significant improvement over the competing TWDs.
5.2 REAI--WORLD
DATA ANALYSIS
We demonstrate the power of the proposed TWD in real-world document classification tasks. The latent hierarchical structure words (features) can be viewed as a node branching into related terms (Gollapalli & Caragea, 2014; Yurochkin et al., 2019; Chang & Blei, 2010), allowing for meaningful comparison between documents OT distance (Kusner al., 2015), as measures how words and their associated meanings are transported from one document to another.
We test four word-document benchmarks in Kusner et (2015), 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 Sec. 5.1. In addition, we include the Word Mover's Distance (WMD) (Kusner al., 2015) that computes the OT distance between documents using Word2Vec (Mikolov et al., 2013). A dissimilarity classification based on our TWD and the baselines using the ANN classifier is applied. We cross- cross-validane validatiit 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., 2015; Huang al., 2016).
Tab. presents the document classification accuracy. The results of Quadtree, Flowtree, TSWD, and WMD are taken from Huang et al. (2016); Takezawa et al. (2021). Our method outperforms all the competing baselines on BBCSPORT, TWITTER, and AMAZON datasets, including the pre-trained WMD that specifically designed for word-document data. On CLASSIC dataset, our proposed TWD yields the second-best classification accuracy and is competitive with WMD.
5.3 SINGLE-CELL GENE EXPRESSION DATA ANALYSIS
We further demonstrate the advantages of our TWD in cell classification tasks on scRNA-seg data. The hierarchical structure of genes in ScRNA-seq data (Tanay & Regev, 2017) are essential for

---

Published as a conference paper at ICLR 2025
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. (2016); Takezawa et al. (2021). The empirical comparisons to TWD methods demonstrate the advantages of our proposed approach, which effectively incorporates the latent feature hierarchy, leading to improved performance.
Real- World Word-Document
Real- World scRNA-scq
Synthetic
BBCSPORT
TWITTER
CLASSIC
AMAZON
Zeisel
CBMC
Quadtree
98.2+12
95.5=a5
69.6208*
95,9=04"
89.320.
8
80.6=06
Flowtree
97.3:109
9.
70.2:199
94,4106'
90.1 20.1
81.710.
81.8.09
WCTWD
91.4+12
92.6=21
69.1
93,7=29
88.2-144
81.3+49
78.4=33
WQTWD
92.8.11
94.3
69.4112
94,6132
87.4-18
80.9
79.1
UltraTree
93.2+3.1
93.1=115
68.1
92.3+19
86.2=311
83.9
82.3=22
TSWD-1
90,9113
8..6
69.8113'
94,5105*
85,5 1..5
79.6.1.8
72.6118
TSWD-5
92.4+11
88.1=13*
70.5
95,9=04"
90.8
81.3+1
74.9=11
TSWD-10
94,2=09
88.6=09*
70.7+13"
95,9=06"
91.1
83,2+08
76.5+07
SWCTWD
94.9+23
92.8=112
70.2+1.2
94.1#18
90.2=112
81.9+3.1
78.3=117
SWQTWD
95.1118
94,5111
70.611.9
95,4130
89.8
80.7.2.
79.8
MST-TWD
88.1
88.42:19
68.21:19
90.0.331
86.4=12
80.1
76,2:25
TR-TWD
93.9=97
89.2=a9
70.2+0.7
92,9+ux
88.7=11
80.3+0.7
78.4=12
HHC-TWD
94.1=17
85.3=18
70.4004
93.4+0a
88.5=07
82.3+007
77.3=11
gHHC-TWD
93,5
83,2=24
69.9
90.3+22
86,9=20
79 19.4
73.6=16
UlDDFit-TW
94,3+11
84.9=14
69.5+1.2
91,6a09
87.4=16
81.9#3.
77.8=12
WMD
95.4=07*
71.3+06"
97.2+a1*
92.6203*
GMD
84.2100
81.4207
Ours
99.810.1
96,1
73.4
96.9102
93.1
89.1
84.3
analyzing complex cellular processes (Ding & Regev, 2021), especially for computing OT distance
(Bellazzi et al., 2021: Tran et al., 2021: Bunne et al., 2023). For instance, in developmental biology, genes associated with stem cell differentiation follow a hierarchical pattern (Kang et al., 201 2). 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 Dumitrascu et al. (2021), 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 WMD, we include the Gene Mover's Distance (GMD) (Bellazzi et al., 2021). gene-based OT distance using Gene Vec (Du et al., 2019). The cell classification tasks are performed in the same way as in Sec. 5.1 and Sec. 5.2. Tab. shows the mean and the standard deviation of the cell classification accuracy. Our TWD outperforms the competing methods by 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 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 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 binary tree based on hyperbolic embeddings, relying on the analogy between the exponential growth of the metric in nigh-dineesiinal 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, present an unconventional use of the TWD. In addition to speeding up the computation of the Wasserstein distance, also use the tree represent the feature hierarchy. Second, latent feature hierarchy by a novel hyperbolic tree decoding algorithm, which differs from existing hierarchical (continuous) representation learning methods that consider (known) sample hierarchy.
consider and recover the
10

---

Published as a conference paper at ICLR 2025
ETHICS STATEMENT
This paper presents new machine learning techniques using publicly available and commonly used benchmarking datasets. These datasets are fully anonymized and do not contain any personal or sensitive information. We respect the privacy and integrity of all data, and we do not foresee any negative societal impact from our methods. We have no conflicts of interest to disclose.
REPRODUCIBILITY STATEMENT
We provide a detailed overview of our experimental setup, including the data splits, hyper- parameter configurations, and the criteria for their selection. Additionally, describe the type of classifier used and other relevant training details. For a complete explanation, please see App. D. Our code is available at https: /gittbb. com/y wwe-wwei-e wei-eileen-lin ewDri.thaateettaatuueeeeeee
ACKNOWLEDGMENTS
We thank the anonymous reviewers for their insightful feedback. We thank Gal Maman for proofread- ing an earlier version of the paper. 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-DIFFOO The work of RC was supported by the US Air Force Office of Scientific Research (AFOSR MURI FA9550-21-1-0317) The work of GM was supported by NSF CCF-2217058.
REFERENCES
Takuya Akiba, Shotaro Sano, Toshihiko Yanase, Takeru Ohta, and Masanori Koyama. Optuna: A next-generation hyperparameter optimization framework. In Proceedings of the 25th ACM SIGKDD international conference on knowledge discovery data mining, pp. 2623-2631, 2019.
Noga Alon, Richard M Karp, David Peleg, and Douglas West. A graphetheoretic game and its application to the k-server problem. SIAM Journal on Computing, 24(1): 1995.
David Alvarez-Melis and Tommi S Jaakkola. Wassoo--ssserttein alignment of word embedding spaces. arXiv preprint 1.0Kv:1800.0001... 2018.
David Alvarez-Melis, Youssef Mroueh, and Tommi Jaakkola. Unsupervised hierarchy matching with optimal transport over hyperbolic spaces. In International Conference on Artificial Intelligence and Statistics, pp. 1606-1617. PMLR, 2020.
Jerrod Isaac Ankenman. Geometry and analysis of dual networks on questionnaires. Yale University, 2014.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International Conference on Machine Learning, pp.
PMLR, 2017.
Arturs Backurs, Yihe Dong, Piotr Indyk, Ilya Razenshteyn, and Tal Wagner. Scalable nearest neighbor search for optimal transport. In International Conference Machine Learning, pp. 497-506. PMLR, 2020.
Alan F Beardon. The geometry of discrete groups, volume 91. Springer Science & Business Media, 2012.
Mikhail Belkin and Partha Niyogi. Towards theoretical foundation for Laplacian- manifold methods. Journal of Computer and System Sciences, 74(8): 2008
Riccardo Bellazzi, Andrea Codegoni, Stefano Gualandi, Giovanna Nicora, and Eleonora Vereesi. The gene mover's 's distance: Single-cell similarity via optimal transport. arXiv preprint 102.022022.01211 2021.
11

---

Published as a conference paper at ICLR 2025
James Bennett, Stan Lanning, et al. The netflix prize. In Proceedings of KDD cup and workshop. volume 2007, pp. 35. New York, 2007.
Clément Bonet, Laetitia Chapel, Lucas Drumetz, and Nicolas Courty. Hyperbolic sliced- Sered-Waseee Wasserstein via geodesic and horospherical projections. In Topological, Algebraic and Geometric Learning Workshops 2023, pp. 334-370. PMLR, 2023.
Nicolas Bonneel, Michiel Van De Panne, Sylvain Paris, and Wolfgang Heidrich. Displacement interpolation using Lagrangian mass transport. In Proceedings of the 2011 SIGGRAPH Asia conference, pp. 1-12, 2011.
Javier Holthoffe and Alex Arenas. Semantic networks: Structure and dynamics. Entropy, 12 (5):1264-1302, 2010.
Brian Hayward Bowditch. A course on geometric group theory. mathematical society of japan, tokyo, 2006. MR2243589 (2007e: 20085), 16, 2007.
Alexander M Bronstein, Michael M Bronstein, Ron Kimmel, Mona Mahmoudi, and Guillermo Sapiro. A Gromov-Hausdorff framework with diffusion geometry for topologically- non-rigid shape matching. International Journal of Computer Vision, (9(2))266-286 2010.
Daniel Brooks, Olivier Schwander, Frederic Barbaresco, Jean- Yves Schneider, and Matthieu Cord. Riemannian batch normalization for SPD neural networks. Advances in Neural Information Processing Systems, 32, 2019.
Charlotte Bunne, Stefan G Stark, Gabriele Gut, Jacobo Sarabia Del Castillo, Mitch Levesque, Kjong- Van Lehmann, Lucas Pelkmans, Andreas Krause, and Gunnar Rätsch. Learning single-cell perturbation responses using neural optimal transport. Nature methods, pp. 1-10, 2023.
Benjamin Paul Chamberlain, James Clough, and Marc Peter Deisenroth. Neural embeddings of graphs in hyperbolic space. arXiv preprint arXivv:70s. 10359, 2017.
Ines Chami, Zhitao Ying, Christopher Ré, and Jure Leskovec. Hyperbolic graph convolutional neural networks. Advances in Neural Information Processing Systems, 32, 2019.
Ines Chami, Albert Gu, Vaggos Chatziafratis, and Christopher Ré. From trees to continuous embed- dings and back: Hyperbolic hierarchical clustering. Advances in Neural Information Processing Systems, 00111565515000 2020.
Jonathan Chang and David M Blci. Hicrarchical relational models for document networks. 2010.
Samantha Chen, Puoya Tabaghi, and Yusu Wang. Learning ultrametric trees for optimal transport regression. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 20657-20665, 2024.
Victor Chepoi, Feodor Dragan, Bertrand Estellon, Michel Habib, and Yann Vaxès. Diameters, centers, and approximating trees of delta-hyperbolic geodesic spaces and graphs. In Proceedings of the !wenty-fourth annual symposium Computational geometry, pp. 59-68. 2008
Giovanni Chierchia and Benjamin Perret. Ultrametric fitting by gradient descent. Advances in neural information processing systems, 32, 2019.
Lenaic Chizat, Pierre Roussillon, Flavien Léger, François-Xave Xavier Vialard, and Gabriel Peyré. Faster Wasserstein distance estimation with the Sinkhorn divergence. Advances in Neural Information Processing Systems, 33:2257-2266 2020.
Ronald R Coifman and Matthew Hirn. Diffusion maps for changing data. Applied and computational harmonic analysis, 36(1):79-107, 2014.
Ronald R Coifman and Stéphane Lafon. Diffusion maps. Applied and computational harmonic analysis, 21(1):5-30, 2006.
Ronald Coifman and Mauro Maggioni. Diffusion wavelets. Applied and computational harmonic analysis, 21(1):53-94, 2006.
12

---

Published as a conference paper at ICLR 2025
Thomas H Cormen, Charles E Leiserson, Ronald L Rivest, and Clifford Stein. Introduction to algorithms. MIT press, 2022.
Nicolas Courty, Rémi Flamary, Amaury Habrard, and Alain Rakotomamoniy. Joint distribution optimal transportation for domain adaptation. Advances in neural information processing systems, 30, 2017.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in Neural Information Processing Systems, 26, 2013.
Marco Cuturi and David Avis. Ground metric learning. The Journal of Machine Learning Research, 15(1):533-564, 2014.
Sanjoy Dasgupta. A cost function for similarity-bar based hierarchical clustering. In Proceedings of the firtt-eighth annual ACM symposium Theory of Computing, pp. 118-127, 2016.
Jiarui Ding and Aviv Regev. Deep generative model embedding of single-cell RNA-Seq profiles hyperspheres and hyperbolic spaces. Nature communications, 12(1):1-17,2021.
Xiucai Ding and Hau-Tieng Wu. Impact of signal-to-noise ratio and bandwidth on graph Laplacian spectrum from high-dimensional noisy point cloud. arXiv preprint arXiv:2011.10725, 2020.
Jingcheng Du, Peilin Jia, Yulin Dai, Cui Tao, Zhongming Zhao, and Degui Zhi. Gene2vec: distributed representation of genes based on co-expression. BMC genomics, 20:7-15, 2019.
Bianca Dumitrascu, Soledad Villar, Dustin G Mixon, and Barbara E Engelhardt. Optimal marker gene selection for cell type discrimination in single cell analyses. Nature communications, 12(1): 1-8, 2021.
Steven N Evans and Frederick A Matsen. The phylogenetic Kantorovich-Rubinserin metric for environmental sequence samples. Journal of the Royal Statistical Society Series B: Statistical Methodology, 74(3):569-592, 2012.
Jittat Fakcharoenphol, Satish Rao, and Kunal Talwar. A tight bound approximating arbitrary metrics by tree metrics. In Proceedings of the thirty-fifth annual ACM symposium on Theory of computing, pp. 448-455, 2003.
Charles Fefferman, Sanjoy Mitter, and Hariharan Narayanan. Testing the manifold hypothesis.
Journal of the American Mathematical Society, 1944),9331 2016.
Maurice Fréchet. Les éléments aléatoires de nature quelconque dans un espace distancié. In Annales de I' l'institut Henri Poincaré, volume 10, pp. 2p.2111110. 1948.
Matan Gavish, Boaz Nadler, and Ronald R Coifman. Multiscale wavelets on trees, graphs and high dimensional data: Theory and applications semi supervised learning. In ICML,, 2010.
Sujatha Das Gollapalli and Cornelia Caragea. Extracting keyphrases from research papers using citation networks. In Proceedings of the AAAI conference artificial intelligence, volume 28, 2014.
John C Gower and Gavin IS Ross. Minimum spanning trees and single linkage cluster analysis. Journal of the Royal Statistical Society: Series (Applied Statistics). 18(1):54-64, 1969.
Mikhael Gromov. Hyperbolic groups. In Essays in group theory, pp. 75-263. Springer, 1987.
Andrés Hoyos-Idrobo. Aligning hyperbolic representations: an optimal transport-based, approach. arXiv preprint arxiv:20/2.01089, 2020.
Gao Huang, Chuan Guo, Matt J Kusner, Yu Sun, Fei Sha, and Kilian Q Weinberger. Supervised word mover's distance. Advances in Neural Information Processing Systems, 29, 2016.
Guillaume Huguet. Daniel Sumner Magruder, Alexander Tong, Oluwadamilola Fasina. Manik Kuchroo, Guy Wolf, and Smita Krishnaswamy. Manifold interpolating optimal-transport flows for trajectory inference. Advances in neural information processing systems, 35:29705-297113 2022.
13

---

Published as a conference paper at ICLR 2025
P. Indyk and N. Thaper. Fast image retrieval via embeddings. In 3rd international workshop statistical and computational theories of vision, 2, 2003.
Pablo A Jaskowiak, Ricardo JGB Campello, and Ivan Costa. On the selection of appropriate distances for gene expression data clustering. In BMC bioinformatics, volume 15, pp- 1-17. Springer, 2014.
Stephen Johnson. Hierarchical clustering schemes. Psychometrika, 32(3):241-254, 1967.
Peng Kang, Hyun Kyoung Lee, Stacey Glasgow, Meggie Finley, Tataka Donti, Zachary Gaber, Brett H Graham, Aaron E Foster, Bennett G Novitch, Richard M Gronostajski, al. Sox9 and NFIA coordinate a transcriptional regulatory cascade during the initiation of gliogenesis. Neuron, 74(1):79-94, 2012.
Leonid V Kantorovich. On the translocation of masses. In Dokl. Akad. Nauk. USSR (NS), volume 37, pp. 199,201, 1942.
Kacper Kapušniak, Peter Potaptchik, Teodora Reu, Leo Zhang, Alexander Tong, Michael Bronstein, Avishek Joey Bose, and Francesco Di Giovanni. Metric flow matching for smooth interpolations on the data manifold. arXiv preprint 2rXiv:2405.1478) 2024.
Ori Katz, Roy Lederman, and Ronen Talmon. Spectral flow on the manifold of SPD matrices for multimodal data processing. arXiv preprint 2xXiv:000.088622 2020.
Felix P Kemeth, Tom Bertalan, Nikolaos Evangelou, Tianqi Cui, Saurabh Malani, and Ioannis G Kevrekidis. Initializing LSTM internal states via manifold learning. Chaos: An Interdisciplinary Journal of Nonlinear Science, (9), 2021.
Tom Kenter and Maarten De Rijke. Short text similarity with word embeddings. In Proceedings of the 24th ACM international on conference on information and knowledge management, pp- 1411-1420, 2015.
Soheil Kolouri, Phillip E Pope, Charles E Martin, and Gustavo K Rohde. Sliced Wasserstein auto-encoders. In International Conference on Learning Representations, 2018.
Ranjay Krishna, Yuke Zhu, Oliver Groth, Justin Johnson, Kenji Hata, Joshua Kravitz, Stephanic Chen, Yannis Kalantidis, Li-Jia Li, David A Shamma, al. Visual genome: Connecting language and vision using crowdsourced dense image annotations. International journal of computer vision, 123:32-73, 2017.
Matt Kusner, Yu Sun, Nicholas Kolkin, and Kilian Weinberger. From word embeddings to document distances. In International Conference on Machine Learning, pp. 957-966. PMLR, 2015.
Hugo Lavenant, Sebastian Claici, Edward Chien, and Justin Solomon. Dynamical optimal transport discrete surfaces. ACM Transactions Graphics (TOG), 37(6): 2018.
Tam Le, Makoto Yamada, Kenji Fukumizu, and Marco Cuturi. Tree-sliced variants of Wasserstein distances. Advances in Neural Information Processing Systems, 32, 2019.
John M Lee. Riemannian manifolds: an introduction to curvature, volume 176. Springer Science & Business Media, 2006.
William Leeb. The mixed Lipschitz space and its dual for tree metries. Applied and Computational Harmonic Analysis, 44(3):584-61 610, 2018
William Leeb and Ronald Coifmann Hölder Lipschit norms and their duals on spaces with semi- groups, with applications to Earth mover's distance. Journal of Fourier Analysis and Applications, 22(4):910-95) 953, 2016.
YaiWei Eileen Lin, Yuval Kluger, and Ronen Talmon. Hyperbolic Procrustes analysis using Rieman- nian geometry. Advances in Neural Information Processing Systems, 34:5959-5977 2021.
Ya- Wei Eileen Lin, Ronald R. Coifman, Gal Mishne, and Ronen Talmon. Hyperbolic diffusion embedding and distance for hicrarchical representation learning. International Conference on Machine Learning, pp. 21003-21022 PMLR, 2023.
14

---

Published as a conference paper at ICLR 2025
Ya- Wei Eileen Lin, Yuval Kluger, and Ronen Talmon. Hyperbolic diffusion Procrustes analysis for intrinsic representation of hierarchical data sets. In ICASSP 2024-2024 IEEE International Conference Acoustics, Speech and Signal Processing (ICASSP), pp. 1325-6329, IEEE, 2024.
Ya- Wei Eileen Lin, Ronald R Coifman, Gal Mishne, and Ronen Talmon. Coupled hierarchical structure learning using tree- wasserstein distance. arXiv preprint arXlv:2501.03627, 2025.
Ofir Lindenbaum, Moshe Salhov, Arie Yeredor, and Amir Averbuch. Gaussian bandwidth selection for manifold learning and classification. Data mining and knowledge discovery, 34:1676-1712, 2020a.
Ofir Lindenbaum, Arie Yeredor, Moshe Salhov, and Amir Averbuch. Multi-view diffusion maps. Information Fusion, 15:127,, 2020b.
Aaron Lou, Isay Katsman, Qingxuan Jiang, Serge Belongie, Ser-Nam Lim, and Christopher De Sa. Differentiating through the Fréchet mean. In International Conference on Machine Learning, pp. 6393-6403 PMLR, 2020.
Kristin McLeod, Tommaso Mansi, Maxime Sermesant, Giacomo Pongiglione, and Xavier Pennec. Statistical shape analysis of surfaces medical images applied the tetralogy of fallot heart. In Modeling in Computational Biology and Biomedicine: A Multidisciplinary Endeavor, pp. 165-191. Springer, 2012.
Harel Mendelman and Ronen Talmon. Supervised and semi supisvueee diffusion maps with label- driven diffusion. In The Thirteenth International Conference on Learning Representations, 2025.
Tomas Mikolov, Ilya Sutskever, Kai Chen, Greg S Corrado, and Jeff Dean. Distributed representations of words and phrases and their compositionality. Advances in Neural Information Processing Systems, 26, 2013.
Gal Mishne, Ronen Talmon, Ron Meir, Jackie Schiller, Maria Lavzin, Uri Dubin, and Ronald R Coifman. Hierarchical coupled- geometry analysis for neuronal structure and activity pattern discovery. IEEE Journal of Selected Topics Signal Processing, (077):238-1253: 2016.
Gal Mishne, Ronen Talmon, Israel Cohen, Ronald R Coifman, and Yuval Kluger. Data-driven tree transforms and metrics. IEEE transactions on signal and information processing over networks, (3):451-466, 2017.
Gal Mishne, Uri Shaham, Alexander Cloninger, and Israel Cohen. Diffusion nets. Applied and Computational Harmonic Analysis, 47(2):259-285, 2019.
Nicholas Monath, Manzil Zaheer, Daniel Silva, Andrew McCallum, and Amr Ahmed. Gradient- based hierarchical clustering using continuous representations of trees in hyperbolic space. In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 714-722, 2019.
Gaspard Monge, Mémoire sur théorie des déblais et des remblais. Mem. Math. Phys. Acad. Royale Sci., pp. 666-704, 1781.
Fionn Murtagh and Pedro Contreras. Algorithms for hierarchical clustering: an overview. Wiley Interdisciplinary Reviews: Data Mining and Knowledge Discovery, 2(1):86-97, 97, 2012.
Boaz Nadler, Stephane Lafon, Ioannis Kevrekidis, and Ronald Coifman. Diffusion maps, spectral clustering and eigenfunctions of Bokker-Plande operators. Advances in Neural Information Processing Systems, 18, 2005.
Yoshihiro Nagano, Shoichiro Yamaguchi, Yasuhiro Fujita, and Masanori Koyama. A wrapped normal distribution on hyperbolic space for gradient-bart based learning. In International Conference on Machine Learning, pp. 4693 PMLR, 2019.
Maximillian Nickel and Douwe Kiela. Poincaré embeddings for learning hierarchical representations. Advances in Neural Information Processing Systems, 30, 2017.
15

---

Published as a conference paper at ICLR 2025
Maximillian Nickel and Douwe Kiela. Learning continuous hierarchies in the Lorentz model of hyperbolic geometry. International Conference on Machine Learning, pp. 3779-3788. PMLR, 2018.
Ofir Pele and Michael Werman. Fast and robust Earth mover's distances. In 2009 IEEE 12th international conference computer vision, pp- 460-467. IEEE, 2009.
Gabriel Peyré, Marco Cuturi, et Computational optimal transport: With applications data science. Foundations and Trends@ in Machine Learning, 11(1g..111111011 11(5-6):355-077, 2019.
Robert Clay Prim. Shortest connection networks and some generalizations. The Bell System Technical Journal, 36(6): 1389-111,,,,,, 1957.
Julien Rabin, Gabriel Peyré, Julie Delon, and Mare Bernot. Wasserstein barycenter and its application to texture mixing. In Scale Space and Variational Methods in Computer Vision: Third International Conference, SSVM 2011, Ein-Gedi, Israel, May 29-June 2, 2011, Revised Selected Papers 3, pp. 435-446. Springer, 2012.
John G Ratcliffe, Axler, and KA Ribet. Foundations hyperbolic manifolds, volume 149. Springer, 1994.
Havard Rue and Leonhard Held. Gaussian Markov random fields: theory and applications. CRC press, 2005.
Frederic Sala, Chris De Sa, Albert Gu, and Christopher Ré. Representation tradeoff's for hyperbolic embeddings. International Conference on Machine Learning, pp. 1460-4463 PMLR, 2018.
Rik Sarkar. Low distortion Delaunay embedding of trees in hyperbolic plane. In International Symposium on Graph Drawing, pp. 355-366. Springer, 2011.
Ryoma Sato, Makoto Yamada, and Hisashi Kashima. Fast unbalanced optimal transport on a tree.
dvances in Neural Information Processing Systems, 18:19999-1901,, 2020.
Chao Shen and Hau-Tieng Wu. Scalability and robustness of spectral embedding: Landmark diffusion is all you need. Information and Inference: A Journal of the IMA, 11(4):527711555, 2022.
Jie Shi, Wen Zhang, and Yalin Wang. Shape analysis with hyperbolic Wasserstein distance. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 5051-5061, 2016.
Sameer Shirdhonkar and David W Jacobs. Approximate Earth mover's distance in linear time. In 2008 IEEE Conference on Computer Vision and Pattern Recognition, pp. 1-8. IEEE, 2008
Tal Shnitzer, Ronen Talmon, and Jean-Jacques Slotine. Diffusion maps Kalman filter for a class of systems with gradient flows. IEEE Transactions Signal Processing, 18:2739-2753, 2020.
Tal Shnitzer, Mikhail Yurochkin, Kristjan Greenewald, and Justin Solomon. Log-Eute -uclldeld signatures for intrinsic distances between unaligned datasets. In International Conference Machine Learning, pp. 20106-20124. PMLR, 2022.
Tal Shnitzer, Hau-Tieng Wu, and Ronen Talmon. Spatiotemporal analysis using Riemannian com- position of diffusion operators. Applied and Computational Harmonic Analysis, 68:101583, 2024.
Amit Singer. From graph to manifold Laplacian: The convergence rate. Applied and Computational Harmonic Analysis, 21 2006.
Justin Solomon, Fernando De Goes, Gabriel Peyré, Marco Cuturi, Adrian Butscher, Andy Nguyen, Tao Du, and Leonidas Guibas. Convolutional Wasserstein distances: Efficient optimal transporta- tion on geometric domains. ACM Transactions on Graphics (ToG), 34(4):1-11, 2015.
Rishi Sonthalia and Anna Gilbert. Tree! am no Tree! am a low dimensional hyperbolic embedding. Advances in Neural Information Processing Systems, 33:845-856, 2020.
16

---

Published as a conference paper at ICLR 2025
Saul Stahl. The Poincaré half-plane: A gateway modern geometry. Jones & Bartlett Learning, 1993.
Marlon Stoeckius, Christoph Hafemeister, William Stephenson, Brian Houck-Loomis, Pratip K Chattopadhyay, Harold Swerdlow, Rahul Satjja, and Peter Smibert. Simultaneous epitope and transcriptome measurement in single cells. Nature methods, 14(9):855-868, 2017.
Yuki Takezawa, Ryoma Sato, and Makoto Yamada. Supervised tree-W Wasserstein distance. In International Conference on Machine Learning, pp. 10086-10095, PMLR, 2021.
Amos Tanay and Aviv Regev. Scaling single-cell genomics from phenomenology to mechanism. Nature, 541(7637:3331-333) 2017.
Alexander Y Tong, Guillaume Huguet, Amine Natik, Kincaid MacDonald, Manik Kuchroo, Ronald Coifman, Guy Wolf, and Smita Krishnaswamy, Diffusion Earth mover's distance and distribution embeddings. International Conference Machine Learning, pp. 10336-10346. PMLR, 2021.
Duc Tran, Hung Nguyen, Bang Tran, Carlo Vecchia, Hung Luu, and Tin Nguyen. Fast and precise single-cell data analysis using hierarchical autoencoder. Nature communications, 12(1): 1029, 2021.
Constantin Udriste. Convex
C functions and optimization methods on Riemannian manifolds, volume 297. Springer Science & Business Media, 2013.
Cédric Villani. Optimal transport: old and new, volume 338. Springer, 2009.
Dingkang Wang and Yusu Wang. An improved cost function for hierarchical cluster trees. arXiv preprint arXiv:18/2.022755 2018.
Simon Weber, Bar Zöngür, Nikita Araslanov, and Daniel Cremers. Flattening the parent bias: Hierar- chical semantic segmentation in the Poincaré ball. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 03223-28232, 2024.
Or Yair, Ronen Talmon, Ronald R Coifman, and Ioannis Kevrekidis. Reconstruction of normal forms by learning informed observation geometries from data. Proceedings of the National Academy of Sciences. 114(38):77655-778 -77744 2017.
Or Yair, Mirela Ben-Chen, and Ronen Talmon. Parallel transport on cone manifold of SPD matrices for domain adaptation. IEEE Transactions Signal Processing, 67(7):1797-1811, 2019.
Makoto Yamada, Yuki Takezawa, Ryoma Sato, Han Bao, Zomitsa Kozareva, and Sujith Ravi. Approximating 1-Wasserstein distance with trees. Transactions on Machine Learning Research, 2022.
Mikhail Yurochkin, Sebastian Claici, Edward Chien, Farzanch Mirzazadeh, and Justin M Solomon. Hierarchical optimal transport for document representation. Advances in Neural Information Processing Systems, 32, 2019.
Amit Zeisel, Ana Manchazaanchado, Simone Codeluppi, Peter Lönnerberg, Gioele La Manno, Anna Juréus, Sueli Marques, Hermany Munguba, Liqun He, Christer Betsholtz, et al. Cell types in the mouse cortex and hippocampus revealed by single-cell RNA-seq. Science, 342(6226) 1138 1138-1142,, 1142, 2015.
17

---

Published as a conference paper at ICLR 2025
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., 2017; Kolouri et al., 2018), aligning and comparing color distributions computer graphics (Solomon al., 2015; Lavenant al., 2018), and serving as a robust similarity metric in data analysis (Bellazzi et al., 2021; Kusner et al., 2015; Yair et al., 2017).
Let μr and #2 be two probability distributions on a measurable space Q with a ground metric I1(p1, (22 be the set of couplings on the product space S x Q such that π (B, x (2) = μr (B1) and T (S) x B2) = uzz B2)) for any sets B1, B2 C Ω. The Wasserstein distance (Villani, 2009) between
Hi and P2 is defined by
inf [ d(z, y)d((( v)-
(T( xEll(nn Jaxx
When replacing the inf in Eq. (4) with the arg min, the obtained solution is referred to as the Optimal Transport (OT) plan (Monge, 1781; Kantorovich, 1942). When computing the Wasserstein distance in Eq. (4) between two discrete probabilities with points, given by
11m m M1.
OT(H1, U42 C) = min
( $ER"X
C) s.t.
+
1m
μ2:
where C E mm is the ground pairwise distance matrix between the coordinates of the discrete probabilities, it can be computed in Om33 log (Peyré et al., 2019), preventing its use of OT for large-scale datasets (Bonneel et al., 2011). Therefore, existing works have been proposed to reduce this complexity. Among them, a notable work is the Sinkhorn algorithm (Cuturi, 2013; Chizat et al., 2020), which approximates the OT in quadratic time.
A.2 DIFFUSION GEOMETRY
Diffusion geometry (Coifman & Lafon, 2006) is mathematical framework designed to analyze high-dimensional data points {a, E _ J==1 by effectively capturing their underlying geometric structures (Belkin & Niyogi, 2008). It focuses on analyzing the similarity between data points through the process of diffusion propagation. Specifically, the constructed diffusion operator P E Rm derived from the observational data is associated with the heat kernel and the Laplacian of the underlying manifold (Belkin & Niyogi, 2008). Let Q(j."') exp (-d (i.j")/e) be the (j.j")th element of the pairwise affinity matrix between the data points. In the limit m -> oc and 6 - 0, the diffusion operator P converges to the heat kernel exp( where A is the Laplace-Beltrant operator on the manifold (Coifman & Lafon, 2006). By applying eigendecomposition to the diffusion operator P. obtain a sequence of eigenpairs ( with positive decreasing eigenvalues 1 1 2 122 2 ... 2 V' > The diffusion maps (DM) m' << m dimensions at the diffusion time embeds point a, into Tp.m by
m' :ay H> [uA522, usspa...... vt m1 1
It is shown that the Fuclidean distances defined in DM between any two points L m (y,) L m (11) approximate the following diffusion distances
dDM..(j.s') = P'ey - P'ey112
at the diffusion time t.
The diffusion operator has been shown to exhibit favorable convergence properties (Coifman & Lafon, 2006). In the limits, as the number of features increases toward infinity, the scaling parameter approaches zero, and the operator p'/c converges pointwise to the Neumann heat kernl H, = cxp( (tA), associated with the underlying manifold, where A is the Laplace- Laplace-Beltrami operator on the manifold. That is, the diffusion operator serves discrete approximation of the continuous heat kernel, thereby capturing the geometric structure of the manifold within a finite-dimensional framework (Coifman & Lafon, 2006; Singer, 2006; Belkin & Niyogi, 2008).
18

---

Published as a conference paper at ICLR 2025
Note that the Censity-normalizee affinity matrix was considered to mitigate the effect of non-uniform data sampling, which was proposed and shown theoretically in Coifman Lafon (2006). In Section 3.1 in Coifman & Lafon (2006), the construction of a family of diffusion operators is introduced with a parameter a E IR and the density q. In Section 3.4 in Coifman & Lafon (2006), it is shown that when a = 1 (in our case Q D QD !) the diffusion operator approximates the heat kernel, and this approximation is independent of the density 9. This technique has been further explored and utilized in (Mishne et al., 2019; Katz et al., 2020; Shnitzer et al., 2022; 2024; Lin et al., 2025).
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., 2005; Coifman & Maggioni, 2006). Its utility has been demonstrated across various applications spanning various fields, including nonlinear stochastic dynamical systems (Shnitzer et al., 2020), multi-view dimensionality reduction (Lindenbaum et al., 2020b), shape recognition (Bronstein et al., 2010), wavelet analysis (Coifman & Maggioni, 2006), internal state values of long short-term memory (Kemeth et al., 2021), and supervised and semi-supervised learning tools (Mendelman & Talmon, 2025).
A.3 HYPERBOLIC GEOMETRY
Hyperbolic geometry is Riemannian manifold of a constant negative curvature (Lee, 2006). 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 (Rateliffe et al., 1994). 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 natural representation of the exponential propagation of the diffusion densities on a dyadic grid. Formally,
the n-dimensional Poincaré half-space model H" with constant negative curvature -1 is defined by H" = {x E R" x(n) > 0} with the Riemannian metric tensor ds² = (dx² (1) + dx² + dx? (m))//2(m)) The Riemannian distance for any x, E H is defined by dg (x = 2 sinh -1 ([|x yllp (12)(()))))))) The distance function das (t,-) in the hyperbolic space H gives rise to two distinct types of geodesic paths (Rateliffe et al., 1994): straight lines that are parallel to the nith axis, and semi-circles that are perpendicular to the boundary of the plane Ik=-1
The Gromov product (Gromov, 1987) is used to define 8-hyperbolic spaces in hyperbolic geometry.
Definition A.1 (6-Hyperbolic Space). A metric space (X, dx) is 8-hyperbolic space (Gromov, 1987) if there exists 8 2 0 such that for any four points 21, I2, 33, T4 E X, we have
dx (xA, v1) + (x22 S max max (2) + dx (x3, na), dx(21, (3) + dx((x2, + 25.
Definition A.2 (Gromov Product (Gromov, 1987)). Let (X, dx) be metric space. The Gromov product (Gromov, 1987) of any two points æ1, 22 E X w.rl a reference point xa E X is defined by
(21, 2 (dx((11,5)) + dx(22, - dx(x1,
(9)
A.4 PRELIMINARIES ON TREES
Definition A.3. Consider undirected weighted graph G = (V,E,A) where V = 21, m} is the node set, E is the edge set, and A E Rmmx is the non-negative edge weight matrix. The shortest path metric de(j.j") of any two nodes j.j V is defined by the length of the shortest path on the graph G from node j to node j.
Lemma A.1. The tree metric is O-hyperbolic.
Definition A.4. A metric d: V X + B is a tree metric if there exists a weighted tree T = (V,E,A) such that for every j.3 V, the metric (") is equal the shortest path metric dz-(j-j').
Definition A.5. A balanced binary tree is a tree where each node has a degree of either 1, making it leaf node, or 3, making it internal node. This tree structure can be either empty 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 1 internal nodes.
19

---

Published as a conference paper at ICLR 2025
Definition A.6. A rooted and balanced binary tree is specific rype 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 degree of
Definition A.7 (Ultrametrie). A metric du V R is ultrametric if the triangle inequality for any three points J1, Js E V is strengthened by the ultrametric inequality du (2) S max fax(d,, (dr Ja), ; du Js):
Remark A.1. If a tree is ultrametric, the distances of all the leaves the root identical.
Definition A.8 (LCCA Relation (Wang & Wang, 2018).). Given tree T, we say that the L.CA relation (12 V js ~ jilr holds in T, the LCA of j2 and js, denoted as j2 V js, is proper descendant of the LCA of j2, 33 and jı, denoted as j2 V ja V j1-
Remark A.2. While very small or specific types (e.g., star-shaped) of trees may be isometrically embedded in Euclidean space, general trees or hierarchical structures cannot, and embedding them into Euclidean spaces results larger distortion. For larger or more complex trees, hyperbolic space often serves as more natural embedding space, preserving their hierarchical structure and distances.
B ADDITIONAL RELATED WORK
Here, we provide more details about related methods and objectives.
Wasserstein Distance on Hyperbolic Manifold. Wasserstein distance learning in hyperbolic spaces builds upon adapting the Wasserstein distance (Villani, 2009) to the structure of hyperbolic geometry (Lee, 2006). This adaptation involves redefining the cost function 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., 2016) and hierarchy matching (Alvarez- Melis et al., 2020; Hoyos-Idrobo, 2020). Recently, the work in Bonet al. (2023) proposed extending the sliced-Waas Wasserstein distance (Rabin et al., 2012) 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 hierarchy rather than probability distributions defined on hyperbolic spaces. Additionally, our focus 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 on the Euclidean metric. This approach approximates the original Euclidean metric using a tree metric, as described in several studies (Indyk & Thaper, 2003;
Le et al., 2019; Evans & Matsen, 2012: Leeb, 2018; Takezawa et al., 2021; Yamada et al., 2022). The resulting methods provide a coarser approximation of OT in an Euclidean space, facilitating more efficient computations from cubic complexity linear complexity. One notable method is the Quadtree (Indyk& & Thaper, 2003), which creates a random tree by recursively dividing hypercubes. This tree is then used to compute the TWD. Flowtree (Backurs al., 2020) 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 al., 2019) represents another TWD technique. It computes the average of TWDs by randomly sampling trees. The number of sampling trees is a hyperparameter, commonly set one, five, ten, resulting in the variants TSWD-1, TSWD- 5, and TSWD-10. Further advancements are presented in Yamada et al. (2022), which introduce WQTWD and WCTWD. These methods employ Quadtree or clustering based tree to approximate the Wasserstein distances by optimizing the weights on the tree. Lastly, UltraTree (Chen et al., 2024) constructs an ultrametric tree (Johnson, 1967) 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 using tree, 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 Fuclidean space.
20

---

Published as a conference paper at ICLR 2025
Tree Construction by Graph-Based and HC Methods. Another line of related work involves constructing trees using existing graph-baser algorithmss, such as Minimum Spanning Tree (MST) (Alon et al., 1995), 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 tree metrie. In the context of hierarchical clustering, various methods have been proposed. UltraFit (Chierchia & Perret, 2019) addresses an "ultrametric fitting" problem using Euclidean embeddings. gHHC (Monath al., 2019) operates under the assumption that partial information about the optimal clustering, specifically the hyperbolic embeddings of the leaves, is known. HHC (Chami et al., 2020) decodes a tree from a -wo-dimesioona hyperbolic embedding for hierarchical clustering by optimizing the continuous Dasgupta's 's cost (Dasgupta, 2016). 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 hicrarchy, ensuring a more accurate representation of the latent hierarchical structure underlying high- diteeniimninal features.
Fréchet Mean on Manifolds. The Fréchet mean (Fréchet, 1948) is a powerful tool for data analysis on Riemannian manifolds (Lee, 2006), including hyperbolic, spherical, and various matrix manifolds. Its application extends across a diverse range of fields, including batch normalization (Brooks et al., 2019; Lou et al., 2020), batch effect removal (Lin et al., 2021; 2024), domain adaptation (Yair et al., 2019), shape analysis and object recognition (Mcleed et al., 2012), 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 product manifold of hyperbolic spaces (see Def. and Def. 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).
C THEORETICAL ANALYSIS AND PROOFS
We note that the numbering of the statements corresponds 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
Proposition 1. The hyperbolic LCA 2} in Def. has a closed-fomm solution, given by z//22 [ (vb vb/)) T 2h/2 2
[ (v) + b/)² T proj(2)) 2tt]] where (z) V 2
Proof. Let mat (1) be the unique geodesic path from 2² to 2i such that Taf maa (0) = 2 and Tul (1) zÿ. Note that the geodesic path Yaf-mat is a semi-cirele, and the last coordinate of the points 2 and zi; are the same, i.e., 2)(m n 1) 2/. (m + 1) 222222 Therefore, the center point of the semi-circle projected on the plane Th is C2f (15) - +t)² .o] Moreover, the radius of the geodesic path is given by proj(z) V 2tt))
x/ z 2
2 ||2
C2t Cet
24/222
The Riemannian mean defined using the Fréchet mean (Fréchet, 1948) has a closed-form expression, and is located at the midpoint of the geodesic curve
v # =
(bb)
wt,) proj vali)
2
21

---

Published as a conference paper at ICLR 2025
Lemma C.1. The radius pz) V 2)))) is lower-bounded by the diffusson distance the diffusion time 2 k, given by
dDM;2 (j.5') S proj(z) v Vzli).
(10)
C.2 PROOF OF PROPOSITION 2
Proposition 2. For any k S 122 2 (12--(11) ~ dgmi (z) V 2333 2//0)/ /dgii (z)) V zz/, z²))
Proof. By Prop. 1, we have 2 V z
( + x proj(z)) v z and z4 Vzy
1/2 15 + } proj(z, v v
respectively. At the ky-th level, the Riemannian distance
between the 1i hyperbolic ICA 2j ki
V z and 2 at the ki-th level is given by
dgn.. (z) VzJi 2 sinh
(1 V 25 11
2 /
zz (m + 1)2 (m +++11)) +11 11))
Note that 2 (m + 1) = 22 and by Prop. I, we have 2 Ao V 2 (m + 1) = proj (z)) v Vz We
have the same relation at the be level. Therefore, for any ki S k2, we have
d (z) V z, x/)))////)))) / (2)2 V 2/7,2/1
sinh
2 1/3 v 2 23
1,/ (2 ug V 2 (m 1)2 (m + 1)
( v 2 23/1 / 2 (m + 1)z,' (m + 1)
sinh
-sinh"
-1 12-22/2+1
| - 2//000 [/-l[)
sinh -1 (2 / | - 3// 1,)
sinh
2 12-k2/241
1++1
|2b0 3/ 11) sinh (2 /2+1
|bb 3 l1)
-2 -kk
where the transition (1) is due to 1555 have
2 11, S ||1000 - vpppy 11 for any 1 S lgg Similarly, we
dgm+1
(z)
k2 V
ko 2gy
2.
kzz /d
(z/5) V 3,, x2)1
sinh
8/ V 2/ - 2//3 / 2
1/ 2 (m 1) 2 (m + 1)
-1 (1/2/)) V カ - -z3² /l/(2//2 (2v/2) v Vz (m + 1)2z"" (m + 11
sinh
-1 (2 (2-2//2 |:00 25/ 1b))
sinh- (2-k1/2 /2+1 | 3/ 1.)
sinh
(1)
2 2 2-.2-k11 R1/2+1 2-22211/+1 /3 2 21/
where the transition (1) because there lower bound v) of any Hellinger distance for j # j'
Therefore, we show that 2 2 (12-k1) dgm+i (z k2 2/²²) /d (z 2 2.)).
C.3 PROOF OF PROPOSITION 3
Proposition 3. The HD-LCA hy in Def. 3 has closed-form solution, given by h
0,,,5] where Ujj'' (Ku+1) pro v 29.) ... proj v vz
Proof. The orthogonal projections { I lie on the unique geodesic path 7o that is on the (m 1)-axis in H Note that the geodesic distance joining any two points o and o/ is
22

---

Published as a conference paper at ICLR 2025
given by In proj(a vx Ke
(Stahl, 1993; Udriste, 2013). The Fréchet mean (Fréchet, 1948) of fo jor /mmo 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 [O,..., .,....,,, O, agst where ajj' E Rso. We want find az that minimizes
, arg min
Pa (h, of
he
Ke m (h, of
arg min h= 0,0 eE
k=0
Ke
proj(zž V2t.
arg min 0,,] TE
In
k=0
Therefore, the geometric mean (K++1) (2 v 24.)) ... proj(z v x
is
obtained by
setting j
C.4 PROOF OF PROPOSITION
Proposition 4. Let r be the root node of B obtained in Alg. and zr be the corresponding embedding in M. The HD-LCA depth is equivalent to the Gromov product (Gromov, 1987), i.e. dm hy,5') (i,s'),, where (i.J") = (dd(i.r) + dn(s'r)) - dn(j.j')).
Proof. Let 2j and Zy be two points in M. denote Mayma, (1) the unique geodesic path from Zj to Zy such that -i, (0) = Zj and ii (1) = Zy. Following Bowditch (2007), for any point Zj" E M, we have dm 2) 2 (i.j")z", where z E Maymay (1). Note that if zy" E Te,- (t), then (j.j') = O. In addition, let 72,-=,,, (t) and 12,-- (t) be the unique geodesic path from Zj to Zy" and the unique geodesic path from Z,' to Zjn, respectively. Let Ajj'"" (Te,- (1), T2,=2,, (1), --,"" (L)) be the triangle connected by these geodesic paths. As M is O-hyperbolic, the triangle Aor"" is O-centre and by triangle inequality, have dm(2;, z) + dM( z) < dm 2z")
dm(zy, 2) + dm(zjr, < 2zr) dm(zz, 2) dz = d 2z)) Therefore, dzj 2) S (i. By the direct application of the aforementioned two inequality, we have (z, h),5') 2 (")
C.5 PROOF OF THEOREM
Theorem 1. For sufficiently large Ke and m, there exist some constants C1, C2 such that Cidr(ir.j2) < d & (zdr(ii.z)) for jı, E [m]) where dB is the decoded tree metric dp in Alg. T is the ground truth latent tree assumed from Sec. 4.1, and dr is the hidden tree metric between featurrss defined as the length of the shortest path on T (i.e., the geodesic distances).
Proof. We adopt the equivalence proof from Lin et al. (2023), which stems from Coifman & Lafon (2006) and (Leeb & Coifman, 2016). Note that diffusion operator P is typically used for manifold learning, where it be viewed as discrete proxy of the underlying manifold. It is shown that the
diffusion operator recovers the underlying manifold (Coifman & Lafon, 2006), i.e., P m-+xc--- H = exp(-A), where H is the Neumann heat kernel of the underlying manifold and A is the Beltrami operator on the manifold. In addition, the diffusion operator is based on a finite of points, which discrete approximation of H.
Laplace-Beltrami
23

---

Published as a conference paper at ICLR 2025
For a family of the operator hA', the multi-scale metric using the inverse hyperbolic sine function of the scaled Hellinger measure is defined by
(x,
22sinh
2 ++
k20
where (x)(x) = fx (x,) (xx")d" Our goal is to show that dn(i:j') = dH.M(i J'), and conclude with Thm. C.1 (Lin et al., 2023) which states the equivalence dH. = di
Note that for small enough
O, we have dp. dp.3"( dH.m( Specifically, following Coifman & Lafon (2006), we have |P -HIl + 0 and |P H H (MM) - as 0. We denote d'm dr, dp.'m).)") = (i.") for abbreviation.
For any three features J1, js E [m]] assume the HD-LCA relation [12 V js film holds. Let ja V js and j1 V js be LCAs of (32 and (5i, on B obtained in Alg. 1, respectively. We denote Zi V Zij and jjz" V Kis the corresponding embeddings. By Prop. I, we have
1 1
proj(z}, v val.a)
2 (wt + =ba) proj(z Ko v Vzja""
2//
+
2
and
1
1 (Rbs + vntt))
T proj(z v 2 17
( + +b)) proj(2], v vz
Zn V 2/aa"
Therefore, the tree distance between ji and js is given by (i. (l))=du(zzn dnl Bn Vzja (d
Zja Zj, V VZj,) + dz, V jis ). By triangle inequality, we have da (J1,3s) S dm (Z), Zj.) We denote dM((1::22) = dz(zz,, 2// for simplicity. By Prop. 2 and Prop. 4, we have e
S dn()i, ja). Moreover, since the HD-LCA relation [i2 V js jilM holds, the tree distance between j2. J3 E [m] on B can be directly obtained by d ja) = d2. As only one of the HD-LCA relations Li2 Vjs jilMM [i V js- j2l. or [ir V j2 ~ jalm holds in B, the proof can be applied to any three points
Ja E [m]] of which HD-LCA relation holds. Using Monte-Carlo integration, the convergence is with probability one as K₂ - dc. For finite set, we have high probability bound on the convergence. Thereeore, dg = dm = dr.
Theorem C.1 (Theorem 1 in Lin et al. (2023)). For a + 2 and sufficiently large Ke, there exists some constants Č₁, C2 > 0 such that de S dd S Özd2.
C.6 PROOF OF THEOREM 2
Theorem 2. For sufficiently large Ke and m, there exist some constant C1, Ō2 > 0 such that ÜiTW(x;, T) < TW(x;, B) < C2TW(x;, Xy, for all samples i, E [m]].
Proof. Let μu and pa be two probability distributions supported on T with a ground metric dr, Mpa, ((p))e) be the set of couplings on the product space 7 x T such that =(B, x T) (4(B)) and (T x B2) = py( ( (B2) ) for any sets B1, B2 C T. The Wasserstein distance (Villani, 2009) between N4 and py defined by
OT(Hss, pe, dr) = inf TEll(A.))) Irze d(x, y), v)-
(11)
Let Far and I1-1leay denote the Lipschitz functions w.r.t. dr and Lipschitz norm w.r.t. dr, respectively. The Kantorovich-Rubinsee dual (Villani, 2009) of Eq. (11) is given by
OT(1s, Py, dr) sup S. [(x)p,(d))) [ (()))edd)) (()))))ddd))))))))
By Evans & Matsen (2012), the witness function ff can be represented by Borel function T [-1,
i... i.c.
J (()))((())) = / (()))(d2))((T((())
24

---

Published as a conference paper at ICLR 2025
where À is the unique Borel measure on T. Therefore, following Le al. (2019), the Wasserstein distance DT( Pe, dr) can be written as
OT(H4 My dr) sup
= I, (u((TT(2))- - - ((dz)
Σ We Σ
(mi(u) - pr(())
eVV
NETr(e)
(W pe; T')
Note that based on Thm. 1, the decoded tree metric d=(:;)) and the underlying tree metric dy (+: are bilipschitz equivalent. Additionally, the TWD using the tree 7 is the Wasserstein distance with ground
pairwise distance dr (Evans & Matsen, 2012), i.e., TW( 122 T") = (T(M,, (2, dr). Therefore, for any two discrete histograms x; and X supported on T, we have TW Xy, B) XW, T). O
C.7 PROOF OF PROPOSITION
Proposition 5. The TWD in Eq. (3) is bilipschitz equivalent to the sum of the multi-scale hyperbolic TWDs TW(x,, xv, BB) for k = 0, l,.,.,, Ke i.e., TW(x;, Xv, B) = Ex TW(x;, X2, B*).
Proof. By Alg. 3, the k-th rooted binary tree B* = (V& hyperbolic LCA {2ī V 2 Js.',, and the length of the edge is determined by the hyperbolic distance d1 By Prop. 2, for all {0, 1,..., KeJ. the tree metric of the k-th rooted binary tree between the features j,j E [m] is bounded by
c Fe. A!) is decoded based on the k-th scale
Bg (i.d") s d (j.'') S B (ggg(i.j)))
where 8e, 3/ > and (:.) the tree metric of BO. Since the embedding distance induced by the E, distance on the product of Ka + 1 hyperbolic spaces, the tree metric d is the & distance on the product of K2 + 1 rooted binary tree {B* 1==0
[B i.c., for any two features j. E [m]] we have
K.
dn(j.j') Σ d (i.j')
k=0
Incorporating the above equation with Prop. 2, the tree metric of the decoded tree dg bounded by the tree metric dğa. given by
d
Σ 8 <dndini))
d
Σ
3x.
k=
==
Under the condition of uniform cost scaling in OT, we have
K.
K.
OT(x;, xv) S OT Xx, d dg
0T X d
k=0
k=0
where transition (1) is based on that OT(X,, Xy, dg+) is at most By OT X,, d d The lower bound can be obtained in similar way. Therefore, as TW( (H,1/42 OT f22 dr). the proposed TWD bilipschitz equivalent to the hyperbolic tree-sliced Wasserstein distance
K. TW(x;, Xv, B) TW Xv, B*).
25

---

Published as a conference paper at ICLR 2025
Table 2: Statistics on real-world word-document and scRNA-seq benchmarks.
Datasct
# 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 8617 cells
47
4000 genes
CBMC
56
500 genes
D ADDITIONAL DETAILS ON EXPERIMENTAL STUDY
We describe the setups and additional details of our experiments Sec. 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., 2015; Huang et al., 2016). This split was applied to both learning of the underlying feature tree and to the kNN model. Therefore, there no train data leakage to the test data. The feature tree constructed from the training data be reused process classifier with varying & E 7, 9, 11, 13, 15, 17, 19}. The evaluation is repeated five times, and we report the best result, averaged over these five runs.
samples. For each evaluation, we employ ANN
Datasets. In real-world word-document datasets, we test four word-document benchmarks in Kusner al. (2015): (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, (ii) 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 Word2 embedding (Mikolov et al., 2013) is used as the word embedding vectors that are trained on the Google News' dataset. 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. (2021), where the datasets are available at the link?. The first dataset, the Zeisel dataset, focuses on the mouse cortex and hippocampus (Zeisel et al., 2015). It comprises 4000 gene markers and 3005 single cells. The second dataset, the CBMC dataset, is derived from cord blood mononuclear cell study (Stoeckius et al., 2017). It comprises 500 gene markers and 8617 7 single cells. We apply the divisive biclustering method from Zeisel et al. (2015) to obtain 47 subclasses for the Zeisel dataset and 56 subclasses for the CBMC dataset. The cell subclasses are used as labels in cell classification tasks. The Gene Vec embeddings (Du et al., 2019) are used the gene embedding vectors following the study of Gene Mover's Distance (Bellazzi et al., 2021).
We report the statistics of these benchmarks in Tab.
Baselines. We compare the proposed TWD with existing TWDss including (i) Quadtree (Indyk & Thaper, 2003) that constructs random tree by recursive division of hypercubes, (ii) Flowtree (Backurs et al., 2020) that computes the optimal flow of Quadtree, (ii) tree-sliced Wasserstein distance (TSWD) (Le et al., 2019) that computes an average of TWDs by random sampling trees, where the numbers of sampling are set to 1, and 10, (iv) UltraTree (Chen et al., 2024) that constructs an ultrametric tree (Johnson, 1967) by minimizing OT regression cost, (v) weighted cluster TWD (WCTWD) (Yamada et al., 2022). (vi) weighted Quadtree TWD (WQTWD) (Yamada et al., 2022), and their sliced variants (vii) SWCTWD and (vii)) SWQTWD (Yamada et al., 2022). We also compare to TWDs where the trees constructed by (i) MST (Prim, 1957), (ii) Tree Representation (TR)
http:///////eeeeeteet/eeeee fes
26

---

Published as a conference paper at ICLR 2025
(Sonthalia & Gilbert, 2020) through divide-and -nn-uue approach, (iii) gradient-based hierarchical clustering (HC) in hyperbolic space (gHHC) (Monath et al., 2019), (iv) gradient-basee Ultrametric Fitting (UltraFit) (Chierchia & Perret, 2019), and (v) HC by hyperbolic Dasgupta's cost (HHC) (Chami et al., 2020). We refer to App. B for an extended description of these baselines. For non-TWD) baselines, in addition include the Word Mover's Distance (WMD) (Kusner al., 2015) for word-document data and we include the Gene Mover's Distance (GMD) (Bellazzi et al., 2021) for scRNA-seq data. In the experiments in Sec. 5, the Euclidean distance d is used as the ground metric for WMD, GMD, and the TWD baselines, which aligns with existing studies (Kusner al., 2015; Bellazzi et al., 2021; Indyk& & Thaper, 2003; Backurs et al., 2020; Le et al., 2019; Chen et al., 2024; Yamada et al., 2022). For the ablation study that uses cosine distance as the ground metric and directly performs classification tasks in the feature space using Euclidean distance and cosine distance, we refer to App. E.
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, 5, 10) x o, where is the median of the pairwise distances.
The diffusion operator's scale parameter e controls information propagation across data points. Smaller values preserve local structures, beneficial for distinct clusters but may cause overfitting and noise sensitivity. Larger c values capture broader relationships, useful for overlapping clusters but can blur distinctions. To balance these effects, set within standard range that is commonly-use in kernel and manifold learning.
The maximal scale K₂ determines the range of scales over which the hyperbolic manifold operates. We follow (Lin et al., 2023) and suggest setting K, E 1..... 19). Our empirical study, consistent with (Lin al., 2023), 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 broad range of K₂ values.
As these parameters are task-specific and do not have a one-size-fitssee solution, we use Optuna (Akiba al., 2019) to efficiently explore the parameter space and identify optimal settings, where (c. K.) are set (10. 12), (5, (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-fitse solution, continue to be of wide interest in ongoing research (Lindenbaum et al., 2020a).
In our study, employed the Gaussian kernel with distance measure based cosine similarity following prior research (Jaskowiak et al., 2014; Kenter & De Rijke, 2015), which is standard and commonly used in classification tasks and manifold learning. In empirical analysis, observed that employing the Fuclidean distance in the ambient space tends to be highly sensitive and underperforms compared to using 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, 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-clement vectors x; E {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 Bectroot 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 {x, jom we create data matrix with the hidden hierarchy of features (words).
27

---

Published as a conference paper at ICLR 2025
D.3 ACCOMMODATING LARGE DATASETS
In our work, note that constructing the diffusion kernels and their eigenvectors by applying the eigendecomposition to them the most computationally intensive step. However, recent methods in diffusion geometry have proposed various techniques to significantly reduce the time and space complexity of diffusion operators. Here, we consider the diffusion landmark (Shen & Wu, 2022) for scalability improvements which reduces the complexity O(n from the much larger O(n"), where r < is the proportion of the landmark set. We briefly describe the diffusion landmark approach below.
{a,} C R", let A' = (a S R be subset of AA, where 1> T = log,,, (m'). A landmark-set affinity matrix is constructed by K(j.'") = exp(-d (j. (e)), where d(j.j') denotes a suitable distance between the data points aj A and a/ E A'. Let D be diagonal matrix, where the diagonal elements are D e KK T
Consider a set of points
1m. The eigen- structure of an SPD matrix D D-111/ 10-1/2QD11 1/2 that is similar to the diffusion operator (Coifman & Him, 2014; Katz et al., 2020) is recovered by applying SVD to the matrix D 1 = ÜAV, where Q KK mm** is the landmark-affinity matrix. The technique of constructing the diffusion on A by going through the subset A' 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.
E ADDITIONAL EXPERIMENTAL RESULTS
E. TOY PROBLEM
Consider four-level balanced binary tree T AV rooted at node V1, where V (Vi...., ] is the vertex set organized from the root to the leaves of the tree, E is the edge set, and A R15x15 is the non-negative edge weights. Consider set of vectors (x; E R15) j==1 such that x; is non-negative realization from normal distribution N(u, L') (Rue & Held, 2005), where μ = 115 E I5 and is the graph Laplacian of T
The set of vectors {x were generated five times, and each time with different number of samples n E {3, 10, 32, 10 316,, 103, 3162, 10*}. Fig. illustrates the normalized Frobenius norm of the discrepancy between the pairwise TWD matrix obtained by Alg. and the pairwise TWD matrix TW(:: 7)) We see that the proposed TWD provides an accurate approximation of the TWD with the tree T, even based on limited sample size. In addition, the number of samples increases, the proposed TWD approaches closer to the true TWD, highlighting its efficacy in capturing the hidden feature hierarchy.
I[TW(. -, .) TWL: ., Till
10
TWW. TIII
10-1
10
10-3
10a
102
103
10a
# of samples (n)
Figure 3: The normalized Frobenius norm of the difference between the proposed TWD and ground truth TWD with different number of samples n.
Thhe mean is large compared to the covariance matrix, and the realizations are non-negative with high prohability: Realizations with negative elements rejected.
28

---

Published as a conference paper at ICLR 2025
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. outperforms deep-based methods like HHC-TWD and gHHC-TWD, 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-TWD, and MST-TWD. The work (Le et al., 2019) 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-TWD, and MST-TWD, our TWD demonstrates significant advantages in terms of classification accuracy. as shown in Tab. 1. We also report the runtime of the OT distance using the embedding distance between features as ground distance, denoted as OT-ED. We note that this computational expense Omm?3 + +m3 log m). In contrast, our TWD offers substantially reduced runtime of (mm thereby demonstrating a significant advantage when handling large datasets. Note that the scalability of the proposed TWD is enhanced by utilizing diffusion landmarks (Shen & Wu, 2022), detailed App. D.3. This enhancement allows our TWD to achieve computational performance comparable to TSWD-1, TR-TWD, and MST-TWD, particularly in datasets with a large number of features, such as the word-document datasets shown in Fig. 4.
10
eurs
Quadtree
Rowtree
Time Consumption
TSWD-
T5WD5-1
10
10²
OT:ED
WCTWD
WOTWD
SMCTIO
SMOTWW
10a
MSTTW
UHHCC-WW
BBCSPORT
TWITTER
CLASSIC
AMAZON
Zeisel
CBMC
Ultrai-TWW
Figure 4: Run time of the proposed TWD and competing TWD and OT distances on scRNA-seq datasets and word-document datasets.
E.3 ABLATION STUDIES
We note that method consists of several non-trivial steps. We provide theoretical justification for our method in Sec. 4. Specifically, in Thm. 2, we show that the TWD TW(..;,") based on the hidden feature tree T can be approximated by TWD using the decoded tree B, where the proposed TWD and are obtained solely from the data matrix, without access to ground truth T. In addition, in Thm. 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 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
input Alg. instead of using the distance based cosine similarity. The Euclidean distance leads to inferior classification compared to the cosine similarity. Indeed, as remarked in App. the selection of the distance metric influences the performance of method. In future work, we plan to comprehensively investigate the use of different metrics.
E.3.2 TWD BY MST AND AT
Fig. 6 shows the classification performance the scRNA-seg and word-document datasets using the TWD with the tree constructed using Minimum Spanning Trees (MST) (Prim, 1957) and Approximat-
29

---

Published as a conference paper at ICLR 2025
cosine Euclidean
97.5
92.5
Classifiattion accuracy
87.5
85
62.5
80
77.5
75
72.5
67.5
Zeisel
CBMC
BBCSPORT TWITTER CLASSIC AMAZON
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.
ing Tree (AT) (Chepoi et al., 2008). To ensure that the tree nodes are contained in a subtree, which required when computing the TWD in Eq. (I), 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 less effective than TWD, implying that MST and AT do not capture well the underlying tree structure between features compared to our binary tree construction.
ours
MST-TW
TV
97.5
95
92.5
Classification accurac
87.5
82.5
80
77.5
75
72.5
70
67.5
Zeisel
CBMC
BBCSPORT TWITTER CLASSIC
AMAZON
Figure 6: The classification accuracy of our TWD and TWDs computed with MST and AT for seRNA-seq and word-document datasets.
E.3.3 TREE CONSTRUCTION EMBEDDING DISTANCE
We note that the embedding distance (Lin et al., 2023) could be simply used for tree construction in existing TWDs (Indyk & Thaper, 2003; Le al., 2019; Yamada al., 2022), graph-based algorithms (Alon et al., 1995), or linkage methods (Murtagh & Contreras, 2012) 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 & Gilbert, 2020), because they not exploit the hyperbolic geometry the embedded space do not have thcoretical guarantees to recover the hidden feature hierarchy.
Tab. 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, important for revealing the hidden feature hierarchy, and therefore, our TWD is more effective and meaningful.
30

---

Published as a conference paper at ICLR 2025
Table 3: Document and single-cell classification accuracy of the proposed TWD and TWDs using embedding distance (ED).
Woal-World Word-Dacument
Weal-World scRNA-seg
BBCSPORT
TWITTER
CLASSIC
AMAZON
Zeisel
CBMC
Quadtree-HB ED
94.7+03
70.3=100
94.7209
88,2207
81.4
80.9
Flowtree-ED
95,0
71.2=112
93.2:109
90,6#111
81.0+12
81,
WCTWD-ED
93.1
68.7=22
92.4+3.1
87.9
81.9452
78.9
WQTWD-ED
94.0+2.1
67.5#29
93.5
86.5#2.1
80.0+3.7
79.7=3.
UltaTrre-ED
94,0+22
69.3
91.4+25
87.9
83.1
80.6#22
TSWD-1-ED
86.4+2.1
67.3
92.9+14
86.2+1.
79.4
74.8+22
TSWD-5-ED
87.3+19
69.2
93.7
89.8
80.8
75.4
TSWD-I0-ED
88.0+13
70.0=111
94.2:109
90.72033
81.9#11
76.0:11.
SWCTWD-ED
91,9+15
70.7=110
94.0#113
91.0=1.5
81.4+2255
76.2:1.5
SWQTWD-ED
93,8+14
71.2
94.2+22
90.3#1.5
82.6+2.7
80.4+29
MST-TWD-ED
89.2.444
69.1
89,2+27
87.8+2.1
81.3+22.
77.4+22
TR-TWD-ED
90.6+13
71.8=1.
94.5
90.6#1.
84.2
80.5
HHC-TWW-ED
84.1+113
71.2=22
91.8133
89.6115
80.7+112
79.5.09
TEI-C--WD-ED
82.5+23
70.8=1.4
91.0=18
87.4+1.
80.1
74.6
UlraFit-TWD-E
84.5+12
70.6=1.5
93.0#1.6
89.3#27
80.6429
78.1
Ours
96.1
73.4=002
96.9202
93.1
89.1+00
84.3
E.3.4 EXPLORING OTHER HYPERBOLIC EMBEDDING AND THE IMPACT OF CURVATURE
Hyperbolic embedding methods (Chamberlain et al., 2017; Nickel & Kiela, 2017; 2018) typically assume a fully known or partially known graph of T. which in turn is embedded into hyperbolic space. In Sonthalia & Gilbert (2020); Lin et al. (2023); Sala et al. (2018); 8);; Weber et al. (2024), these methods were shown to inadequately capture the intrinsic hierarchical relation underlying high-dimensional observational data X (i.e., the node attributes). Therefore, we opted to use hyperbolic embedding (Lin et al., 2023), a recent hyperbolic embedding method with theoretical guarantees. This embedding and its accompanying theory enabled to introduce theoretical results in Sec. 4.
is important to note that we could indeed apply traditional hyperbolic embedding methods directly to the affinity matrix Q. From graph perspective, this is the affinity of 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 graph induced by 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., 2023) is designed such that the recovered hyperbolic metric is close to that of the hidden tree T', guaranteeing will be bilipschitz equivalent 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 Q. Instead, Lin (Lin et al., 2023) designed a non-trivial process that neither iterative nor based 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.
best as possible the shortest path distance this non-tree
We here give an intuition as to the reason for the success of the approach from Lin et al. (2023). After normalizing the affinity matrix Q, the diffusion operator P is computed. This operator was proven to reveal the underlying manifold. This property is what allows embedding approach based on the diffusion operator to approximately recover the underlying tree metric dr rather than fit the metric of non-tree graph constructed from the affinity matrix Q. The diffusion operator reveals the underlying manifold the following sense (Coifman & Lafon, 2006). In the limit of the infinite number of features m + oo and small scale E - O, the operator pt/r pointwise converges to the Neumann heat kernel of the underlying manifold H, exp(-tA) where A is the Laplace- -elltrami operator the manifold. In other words, the diffusion operator is a discrete approximation of the heat kernel on the manifold. Intuitively, P is designed to reveal the local connectivity diffusion time scale
31

---

Published as a conference paper at ICLR 2025
Considering multiple timescales dyadie grid (2-*) "=== allows for the association of different dif- fusion timescales with neighborhoods of varying sizes. The multi-scale embedding of the associated diffusion operators in hyperbolic space generates distances that 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 latent hierarchical metric. From theoretical standpoint, it was shown that the & distance in this multi-scale embedding induced by the metric of the product of hyperbolic spaces can estimate the distance dr in the hidden hierarchical space that underlies the observations even when dr is not given or when we do not have access to the explicit T (Lin al., 2023). It is stated under the assumption that (P*) (P*) is a pointwise approximation of the heat kernel. Such an approximation, mainly in the limit mn + co and + O, was shown and studied in Coifman & Lafon (2006); Singer (2006); Belkin & Niyogi (2008). 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., 2023).
In our work, we use the hyperbolic diffusion embedding from Lin et al. (2023) to present a novel data-driven tree decoding method that operates in nigh-dimensional hyperbolic spaces. Our method is based on a new definition of the counterpart of the LCA in hyperbolic spaces. Our Thm. 1 is built on the multi-scale hyperbolic embedding from Lin et al. (2023), where we use the theoretical property (Lin et al., 2023) that the & distance in this embedding can estimate the hidden tree metric dy. The difference between Thm. 1 and the work from Lin et al. (2023) is that we construct an explicit tree B whose tree metric approximates the hidden tree metric dr. In Lin et al. (2023), no tree is recovered, and only the hidden tree metric is approximated by the hyperbolic embedding distances. Our Thm. is established based on Thm. 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., 2023), and (ii) there exists a closed-fore solution of the hyperbolic LCA and the HD-LCA, which are the key components in tree decoding algorithm in Alg. 1. Since the four standard models of hyperbolic space (Bowditch, 2007): the Poincaré disk, Lorentz, Poincaré half-space, and Beltrami- Klein, are equivalent and isometric, it is possible to 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 & Kicla, 2017; 2018; Sala et al., 2018; Chami al., 2020). Any negative constant curvature can indeed be used and considered as a hyperparameter, e.g., (Chami et al., 2019). The modifications that it will require from our algorithm are minimal, and most of the steps can be applied seamlessly in such 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/K where k >> O. Fig. 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 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/k-- -1 is a reasonable and commonly accepted practice.
We note that other TWD baselines could use other two- vro-dimensional Poincaré representations (e.g.. Poincaré embedding (PE) (Nickel & Kiela, 2017) or hyperbolic VAE (Nagano et al., 2019)) and decode the hyperbolic embedding into a tree by HHC (Chami et al., 2020). We denoted these baselines respectively, and present their results in Tab, 4. We see in the table that HHC (Chami et al., 2020) applied to a PE (Nickel & Kicla, 2017) or hyperbolic VAE (Nagano et al., 2019) is less effective than our TWD by a large margin. This discrepancy attributed to the sensitivity of the recovery of the hidden feature hierarchy using the competing methods, finding that is consistent with previous works (Sala et al., 2018; Sonthalia & Gilbert, 2020; Lin et al., 2023).
by PE-HHC-TWD and
32

---

Published as a conference paper at ICLR 2025
Zelsel
CBMC
Classif cation Accurscy
$6.0
60.00
65.5
Ciarsif
89.75
89.50
89.25
0.0
0.5
1.0
1.5
0.0
0.5
1.0
1.5
2.0
-logte]
-l9g14)
Figure 7: The classification accuracy for the Zeisel and CBMC datasets under different curvatures
-lKK. Decreasing curvature could slightly improve the cell classification accuracy.
Table 4: Document and single-cell classification accuracy of our TWD and TWDs constructed by other hyperbolic embeddings.
Real-Worla Word-Document
Real-Worke scRNA-seq
BBCSPORT
TWITTER
CLASSIC
AMAZON
Zeisel
CBMC
PE TWD
87.2+12
69.7.09
93.2+11
86. 86.2109
81.6109
74.9=1.
HVAE-HHC-TWD
88.4125
64.3
90.8+1.5
87.3+17
80.2+09
75.1
Ours
96.1
73.4
96,9
93.1104
89.1
84.3
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 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 provide reliable recovery guarantees even when the ground truth hicrarchical features are obscured by noise.
Zeisel
Cassificatoon Acturaey
80
85
80
Towtree
Duedrow
194D-1
75
20
assification
70
HHC TW
TR/7W
60
MCCTVO
-00
Ultraooo
MGF
PRICTV
Noise Varfance
Noise Variance
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 the cell classification performance, our method remains more robust than the competing baselines by large margin.
E.3.6 COMPARISON TO NON-TWI BASELINES
For non- TWD baselines, also compare method to Word Mover's Distance (WMD) (Kusner et al., 2015) for word-document datasets and Gene Mover's Distance (GMD) (Bellazzi et al., 2021) for ScRNA-seq datasets in Sec. 5. Note that WMD (Kusner al., 2015) computes the Wasserstein
33

---

Published as a conference paper at ICLR 2025
Table Document and single-cell classification accuracy of WMD and GMD using cosine distance as the ground metric.
Real-World World Word-Docunin
Real-
Real- World scRNA-seq
BBCSPORT
TWITTER CLASSIC
AMAZON
Zcisel
CBMC
WMD
91.2+04
68.4
93.5
89.4
GMD
86.9
82.7:0.9
Ours
96.1104
73.4:0.2
96,92022
93.1=00
89.1
84.3
Table Document and cell classification accuracy in the feature space using the Euclidean distance and cosine distance.
Real- World Cord-Dooumeet
Real RorldWorld scRNA -se
BBCSPORT
TWITTER
CLASSIC
AMAZON
Zeisel
CBMC
Euclidean
78.2=1.5
$7.3-06
64.0=a5
71.3204
67.241.
62.5.09
Cosine
91.S=09
68.1=06
90.3=a8
86.4
73.5
72.7+1.
Ours
96.1=04
73.4=u2
96,9=a2
93.1
89.1
84.3
distance using the Euclidean distance d between precomputed Word2Vec embeddings (Mikolov et al., 2013) as the ground cost, while GMD (Bellazzi et al., 2021) is defined similarly, using the Euclidean distance d between precomputed Gene2 embeddings as the ground cost. In the experiments in Sec. 5, the Euclidean distance d is used as the ground metric for WMD, GMD, and the TWD baselines, which aligns with existing studies (Kusner et al., 2015; Huang et al., 2016; Le et al., 2019; Takezawa et al., 2021; Yamada et al., 2022; Chen et al., 2024; Bellazzi et al., 2021).
We conduct additional experiments evaluating WMD on word-document datasets and GMD on scRNA-seq datasets, both using cosine distance as the ground metric. Tab. 5 presents the document and cell classification accuracy of WMD and GMD using cosine distance as the ground metric. The results show that WMD performs worse with cosine distance the ground metric compared to Euclidean distance, while GMD performs better with cosine distance than Euclidean distance. However, regardless of the ground metric (cosine or Euclidean), WMD and GMD exhibit less effective performance compared to our method. These findings emphasize the value of inferring latent feature hierarchies, as our approach outperforms the baselines in these tasks.
is important to note that using the ground metric à (e.g., Euclidean) in WMD, GMD, or existing TWD methods results in the (approximated) Wasserstein distance being based on the predefined metric d. In contrast, our method employs initial distance d, computed using the cosine similarity in the experiments, to compute an affinity matrix within the diffusion operator, which is central to kernel and manifold learning methods. This process leads to Wasserstein distance based on latent hierarchical metric induced by d. As noted in App. D.I, the kernel type, scale parameter, and initial distance metric selection in our approach significantly influence its effectiveness, These choices are task-specific and are critical in achieving optimal performance (Lindenbaum et al., 2020a).
In addition, we conduct additional tests evaluating classification performance directly the feature space, without using Wasserstein distance or TWD. Tab. 6 presents the document and cell classification accuracy in the feature space using the Fuclidean distance and cosine distance. We see that the classification performance relying on generic metrics (Euclidean or cosine) is much less effective than WMD and GMD, which aligns with the findings reported WMD (Kusner et al., 2015) and GMD (Bellazzi et al., 2021). This highlights that incorporating the feature relationship into distance between samples improves the effectiveness of distance metric learning. Furthermore, we see that our TWD method performs better than the distance relying generic metrics, WMD and GMD. This suggests that inferring explicitly the feature hierarchies is important and leads to the effectiveness of distance metric learning for data with hierarchical structure.
34

---

Published as a conference paper at ICLR 2025
E.4 EVALUATING OUR TWD ON MON-HIERARCHIAA DATA
We focus on the problem of data with a hidden feature hierarchy as it is of broad interest and prevalent across 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 in the paper. The digit classification was performed in the same way as Sec. 5. Fig. 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.
Classification Accuracy
92
90
88
86-
84
82
80-
78-
76
ours
Quadtree
Flowtree
UltraTree
TSWD-1
TSWD-5
TSWD-10
TR-TW
HHC-TW
WCTWD
WQTWD
SWCTWD
SWQTWD
MST-TW
gHHC-TW
Ultrafit-TW
Figure 9: Digit classification accuracy on MNIST dataset. In contrast to the word-document and scRNA-seq data, this case, there are no definitive latent hicrarchical structures. Yet, our method still demonstrates good results and outperforms the competing TWD methods.
We remark that the state-of-the-eee 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-base methods) to the state of the art is not entirely fair.
F ADDITIONAL REMARKS
F.1 SINGLE-SCALL HYPERBOLIC DIFFUSION BINARY TREE DECODING
In Alg. I, the tree decoding is based on the multi-scale embeddings and the corresponding HD-LCAs. In Alg. 3 present similar algorithm based the single scale hyperbolic embedding z and the corresponding je-th hyperbolic LCAs 2 V zz. Def. 1, giving rise (single-scale) binary tree B* Specifically, we merge the pair of features exhibiting the highest similarity at the k-th hyperbolic LCA, determined by the smallest value of pzoj V z in Prop. 1. Following Def. 2, only one of the hyperbolic LCA relations [jz ja 51 1 m Lnr js fel or [jr j2 ja) holds in 8 The merged nodes are then connected by edges whose weights are assigned by dgm=1 (...) between the corresponding points. It is worth noting that considering (Ka + 1) trees { (B*) ) k= generated by Alg. 3 gives rise to tree-sliced Wasserstein distance as defined in Prop. 5. In Prop. 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 & Kiela, 2017; 2018; Nagano al., 2019), extending its applicability and utility. While this extension technically feasible, does not carry the theoretical guarantees those established in Thm. 1 and Thm. 2.
Ke
35

---

Published as a conference paper at ICLR 2025
Algorithm 3 Single-Scale Hyperbolic Binary Tree Decoding
Input: single scale Poincaré half-space representation at the k-th level ((xf), (2/) ..., (2
Output: The k-th rooted binary tree B = (V*, D. A with m leaf nodes
function H BinaryTreel (f [[(2]] (
frm }1)
B [elvees (5}
{1, m m})
for j.j' E (1, , do Σ
2 Vzj. f arg min
L m (z, 2f()
end for S = [') pairs sorted by z4 V z (m + 1)}
for (j.j') S do
if nodes and j' are not in the same subtree in B* I, f internal node consisting node j I internal node consisting node add new internal node consisting Zj and Ij/ assign the edge weight by their geodesic distance end if end for return B*
F.2 PUTTING THE FEATURES ON THE LEAVES OF THE TREE
The ground truth latent tree does not assume that features are positioned at the leaves. In addition, note that the hyperbolic embedding (Lin al., 2023), which is consistent with other common hyperbolic embedding methods (Chamberlain et al., 2017; Nickel & Kiela, 2017; 2018; Sala et al., 2018), 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. I), which is in line with other tree decoding algorithms (Sarkar, 2011; Chami et al., 2020; Indyk & Thaper, 2003). Note that our tree decoding algorithm generates weighted binary tree. We assert that this does not pose a restriction because any tree can be effectively transformed into a binary tree (Bowditch, 2007: Cormen et al., 2022). For instance, consider the flexible tree in Fig. 10 (a): three children: "basketball", "soccer", and "football"). This tree can be accurately modeled as (weighted) binary tree as in Fig. 10 (b): le[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).
(where "sports" is parent node with
Sport
b
Soccer
Football
Basketball
Soccer
Football
Basketball
Sport
(a)
Figure 10: (a) In a flexible tree, "sports" is parent node with three children: "basketball", "soccer", and "football". (b) In binary tree, "sports", "basketball". "soccer". and "football" are all leaves, a, b, e are internal nodes.
36

---

Published as a conference paper at ICLR 2025
F.3 DIFFERENTIABILITY OF OUR TWD
The proposed TWD can be expressed in matrix form (Takezawa et al., 2021) as
TW Xy, B)
Σ OF Σ (xi(w) - xx(u)) = ||diag(o) (x; xy)ll.
(12)
VE
(ET
where a E R is a weights vector with ax ae. R E {0. 1) Mx is a parent-child relation matrix of the decoded tree B, where ( = 1 if node v' is the ancestor of the leaf node j. and
M = 2m + is the total number of the nodes B.
Note that the hyperbolic embedding in Lin al. (2023) can be recast function mapping from the feature diffusion operator P to the product manifold of the hyperbolic spaces M, denoted by fi P + .M. The weight a in the TWD then be expressed by
M x M
and the weights vector a E RR in Eq. (12) can be written as 92 fi the function 9i M times. Let : B be the hyperbolic diffusion binary tree decoding function in Alg. The parent-child relation matrix R. can be obtained by applying the function
can
by applying
h2 : B + {0, 1].11 xm
h2 : h1 h1(M) + {0, 10 [0,1)Mxm Mxm
h2 (h1 fP) + {0, 1]
where denotes function composition. Therefore, the gradient the our TWD w.r.t. j can be written
V ([diag((i(P) ((() o
This implies that the gradient of our TWD depends on the feature diffusion operator P. Notably, since the gradient of the feature diffusion operator hinges the gradient of the Gaussian kernel Qe (as detailed in Sec. 3), the gradient of the proposed TWD can be expressed only in terms of the gradient of Qe. Such formulation could be highly beneficial 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., 2023) is a hierarchical distance between the samples. Namely, in our TWD, the features are assumed to have hidden hierarchies, whereas in Lin al. (2023), the samples are assumed to have hidden hierarchies. The work in Lin al. (2023) does not involve TWD or OT.
More concretely, in comparison the work in Lin al. (2023), our work presents a new algorithm for tree decoding (Alg. I) and corresponding new TWD (Alg. 2). Our work also includes theoretical analysis of these two algorithms:
Our TWD recovers the TWD associated with the true hidden feature tree, as stated in Thm.
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. 1.
3. Our TWD captures the tree-sliced information encoded in the multi-scale trees efficiently, stated in Prop. 5.
4. Our TWD can be computed in linear time, enhancing computational efficiency in OT framework, as detailed in Sec. and App. E.
5. Our TWD is differentiable (details in App. F.3), which is important property that is lacking in existing TWD baselines.
In addition, we present favorable empirical results in Sec. 5 compared to several TWDs and Wasser- stein distances.
37

---

Published as a conference paper at ICLR 2025
Table Document and single-cell classification accuracy of our TWD, embedding distance (ED), and OT distance using embedding distance (OT-ED).
Real- World Word- Word-Document
scRNA-seq
BBCSPORT
TWITTER
CLASSIC
AMAZON
Zeisel
CBMC
ED
87.2
69.7.009
93.2
86.2=
84.5
81,4
OT-ED
95.910.3
73.5=003
96.720.1
93.2=002
89.0=03
84, 100
Ours
96.1+04
73.4-02
96.9202
93.1
89.1
84.3+03
Seemingly, from metric learning viewpoint, both our TWD and the embedding distance (Lin et al., 2023) 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--eq datasets of these two baselines in Tab. 7. We see clear empirical advantage of our TWD classifying documents and scRNA-seq data compared to embedding distance (ED), our TWD learns the latent hierarchy of features explicitly by constructing 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. 2). the computational expense is O(mn + m log m). In contrast, proposed TWD offers substantially reduced runtime of Om...2), thereby demonstrating significant advantage when handling large datasets. This computational advantage empirically validated and illustrated in Fig.
F.5 INTUITION BEHIND EMBEDDINGS
TREE CONSTRUCTION HYPERBOLIC
The single-linkage tree aligns well with hyperbolic embeddings because inherently emphasizes hierarchical relationships by connecting points based the measure of similarity determined by their HD-LCAs, where there is the "treelike" nature of a hyperbolic space (Bowditch, 2007). Specifically, Prop. 4 shows that the LCA depth closely approximates the Gromov product, which corresponds to the tree depth in O-hyperbolic metrics.
In the context of tree- Wasserstein distance (TWD), many tree construction methods aim to approxi- mate a given ground metric d (often Euclidean) using tree metries, enabling efficient computation of optimal transport distances. The quality of TWD methods is typically assessed by the mean relative error between the proposed TWD and the Wasserstein distance with the ground metric d. For instance, an ultrametric incurs a log m distortion with respect to Euclidean distance (Fakcharoenphol et al., 2003), utilized in UltraTree (Chen et al., 2024). Hence, the ultrametric Wasserstein distance (Chen et al., 2024) incurs a log m distortion compared to the original Euclidean Wasserstein distance.
Our approach differs fundamentally from existing TWD methods. Instead of approximating the Wasserstein distance with given ground metric d (e.g., Euclidean), we aim approximate the Wasserstein distance with a latent underlying hierarchical distance via TWD, where we use the tree to represent the latent feature hierarchy.
38


Reviews:
[{"review_id": "EQxOZa6mjs", "reviewer": "Reviewer_89iu", "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.", "questions": "I put the suggestions in the weakness.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "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.", "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\nSection 3 in Leeb & Coifman (2016) and Appendix C in Lin et al. (2023), the bilipschitz constants $c$\nand $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.\n\nFor 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$.\n\nI 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$\".\n\nAnother 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xIH45Bo7xk", "reviewer": "Reviewer_EDMh", "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.\n\nFor 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). \n\nThe 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.", "questions": "* l. 113 V is not defined, should this be [m]?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 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).\n\nThe 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.\n\nThe problem of computing distances between histograms based on a latent tree is well-motivated by practical applications to single-cell sequencing data.", "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.\n\nAt 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.\n\nI 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6QEYthC8wN", "reviewer": "Reviewer_7P2E", "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.", "questions": "MAJOR COMMENTS\n\n* 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?).\n\n* 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.\n\n* 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.\n\n\nMINOR COMMENTS\n\n* line 39: \"due to its exponential growth\" -- unclear formulation, whose growth?\n\n* line 113: What is V? This notation wasn't introduced.\n\n* 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?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "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.", "weaknesses": "The experimental evaluation is missing naive baselines, and in general lacks details.\n\nOverall, I am giving a borderline score which I would be willing to revise upwards (or downwards) based on authors' clarifications regarding the experiments.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kb5IRNeQwt", "reviewer": "Reviewer_sSsQ", "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.", "questions": "1. Can you clarify the construction of the diffusion operator as referred to above?\n2. Does the method require assuming the feature tree can be isometrically embedded in Euclidean space?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 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. \n\nQuality: The method appears to beat existing approaches as measured by tree reconstruction quality on synthetic data and nearest neighbor classification accuracy on real data. \n\nClarity: The paper is well-written and clear and uses consistent notation.\n\nSignificance: 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.", "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.\n\n2. 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?\n\n3. 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.\n\n4. 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?\n\n5. 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)