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

This is comment 2 out of 2.

Intuition Behind Single-Linkage Tree Construction in Hyperbolic Embeddings

The single-linkage tree aligns well with hyperbolic embeddings because it inherently emphasizes hierarchical relationships by connecting points based on the measure of similarity determined by their HD-LCAs, where there is the “treelike” nature of a hyperbolic space [7]. Specifically, Proposition 4 shows that the LCA depth closely approximates the Gromov product, which corresponds to the tree depth in 0-hyperbolic metrics.

In the context of tree-Wasserstein distance (TWD), many tree construction methods aim to approximate a given ground metric $\widetilde{d}$ (often Euclidean) using tree metrics, 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 $\widetilde{d}$. For instance, an ultrametric incurs a $\log m$ distortion with respect to Euclidean distance [8], as utilized in UltraTree [9]. Hence, the ultrametric Wasserstein distance [9] 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 a given ground metric $\widetilde{d}$ (e.g., Euclidean), we aim to approximate the Wasserstein distance with a latent underlying hierarchical distance via TWD, where we use the tree to represent the latent feature hierarchy.

To clarify why our single-linkage tree construction approach performs better than other tree constructions given the hyperbolic embeddings, we have included the above discussion in Appendix F.5.

Thank you again for the valuable comment.

[7] Bowditch (2007). A course on geometric group theory.

[8] Fakcharoenphol et al. (2003). A tight bound on approximating arbitrary metrics by tree metrics.

[9] Chen et al. (2024). Learning ultrametric trees for optimal transport regression.

This is comment 1 out of 2.

We thank the reviewer for the insightful comments and questions. In addition, we thank the reviewer for appreciating the idea of learning an explicit hierarchy of features and its relevance to distance metric learning and TWD for data represented as distributions over features.

Assumption on Features Being Leaves in the Ground Truth Feature Hierarchy

We would like to clarify that we do not assume features are leaves in $T = ([m], E, \mathbf{A}, \mathbf{X}^\top)$. In this context, the vertex set $V = [m] = \\{1, \ldots, m\\}$ represents the nodes on the tree $T$.

In our decoded tree $B$, however, we assign all $m$ features to be leaf nodes in $B$, where the tree $B$ consists of $2m-1$ nodes.

To further emphasize that we do not assume features are leaves in $T$, we have extended the discussion in Appendix F.2, where we state that neither the underlying tree $T$ nor the hyperbolic embedding requires all features to be assigned to leaves. The assignment of features to leaf nodes is performed only during the tree decoding process.

Clarification of Feature Diffusion Operator Construction

To clarify the construction of the diffusion operator, we have incorporated the reviewer's suggestion and specified in the revised version in Section 4.2 that we computed the initial distance using the cosine similarity in the ambient space, defined as $d(j,j') = 1- \frac{\mathbf{X}\_{:, j} \cdot \mathbf{X}\_{:, j'}}{\left\lVert\mathbf{X}\_{:, j}\right\rVert_2 \left\lVert\mathbf{X}\_{:, j'}\right\rVert_2 }$, where $\cdot$ is the dot product. We also referred to Appendix D for further discussion on the initial distance metric, the scale parameter $\epsilon$, and the kernel type.

In the appendix, we highlighted that these elements, which do not have a one-size-fits-all solution, continue to be of wide interest in ongoing research [1]. Our primary objective is to highlight the unique aspects of our method rather than kernel selection, which, while crucial, is a common consideration across all kernel and manifold learning methods.

[1] Lindenbaum et al., 2020. Gaussian bandwidth selection for manifold learning and classification.

Definition of LCA Relation

In the revised version, we have incorporated the reviewer’s suggestion by including the definition of the LCA relation from [2] in Appendix A.4.

[2] Wang & Wang (2018). An improved cost function for hierarchical cluster trees.

Discussion on Non-Isometric Euclidean Embedding of Feature Trees

We would like to clarify that we do not assume the feature tree can be isometrically embedded in Euclidean space. While very small or specific types of trees (e.g., star-shaped trees) may allow for isometric embeddings in Euclidean space, general trees or hierarchical structures typically cannot. Embedding such structures into Euclidean space often introduces distortion [3]. For larger or more complex trees, hyperbolic space often serves as a more natural embedding space, preserving their hierarchical structure and distances [3-6].

To clarify that we do not require feature trees to be isometrically embedded in Euclidean space, we have included this remark in Appendix A.4.

[3] Nickel & Kiela (2017). Poincaré embeddings for learning hierarchical representations

[4] Nickel & Kiela (2018). Learning continuous hierarchies in the Lorentz model of hyperbolic geometry

[5] Chami et al. (2019). Hyperbolic graph convolutional neural networks.

[6] Monath el al. (2019). Gradient-based hierarchical clustering using continuous representations of trees in hyperbolic space.

We sincerely thank the reviewer for taking the time to carefully evaluate our responses and for their insightful questions.

• Regarding the assignment of features to leaves: Yes, your interpretation is correct.

• Regarding embedding trees in Euclidean space and the affinity matrix:

  Thank you for clarifying your concern.

  We agree with the reviewer that global distances, e.g., Euclidean or cosine distances, do not capture the hierarchical information or the underlying manifold structure of the features. However, we would like to point out that the affinity matrix is inherently local due to the use of the Gaussian kernel. It is not because local structures are computed with Euclidean components that the global geometry has to be Euclidean or that we are embedding in a Euclidean space. For instance, curved Riemannian manifolds are locally Euclidean, and the geodesic distance between nearby points can be approximated by the Euclidean distance, yet the geodesic distance between distant points is far from Euclidean.

  In our method, the Gaussian kernel enables the affinity matrix to capture local relationships. This approach, related to heat diffusion on a curved non-Euclidean manifold, does not mean that we are embedding the features in a Euclidean space, as global distances are not Euclidean. This kernel enables the affinity matrix to capture local relationships, where affinities between features that are close in the ambient space embody information on the hidden hierarchical structure. Specifically, the discrete diffusion operator, constructed from these local Gaussian affinities in the ambient space, asymptotically converges pointwise to the heat kernel of the underlying hierarchical metric space in the limit of $m\rightarrow \infty$ and $\epsilon \rightarrow 0$. This is a standard practice in the studies of data manifolds and manifold learning to approximate and represent data on non-Euclidean manifolds [1, 2].

  [1] Coifman & Lafon (2006). Diffusion maps.

  [2] Belkin et al. (2003). Laplacian eigenmaps for dimensionality reduction and data representation

This is comment 4 out of 4.

Clarification of "Exponential Growth"

We have revised the text for clarity as follows: "Hyperbolic geometry has gained prominence in hierarchical representation learning because the lengths of geodesic paths in hyperbolic spaces grow exponentially with the radius, a property that naturally mirrors the exponential growth of the number of nodes in hierarchical structures as the depth increases".

Typo $V$

Thank you for pointing this out. We have updated $V$ to $[m]$ in the revised version.

Clarification of the Construction of the Diffusion Operator

We would like to clarify that the density-normalized affinity matrix was considered to mitigate the effect of non-uniform data sampling, which was proposed and shown theoretically in Coifman & Lafon (2006) [13]. In Section 3.1 of Coifman & Lafon (2006) [13], the construction of a family of diffusion operators is introduced with a parameter $\alpha \in \mathbb{R}$ and the density $q$. In Section 3.4 of Coifman & Lafon (2006) [13], it is shown that when $\alpha = 1$ (in our case $\widehat{\mathbf{Q}} = \mathbf{D}^{-1}\mathbf{Q}\mathbf{D}^{-1}$), the diffusion operator approximates the heat kernel, and this approximation is independent of the density $q$. This technique has been further explored and utilized in [14-17].

In addition, we consider the column-stochastic version of the diffusion operator due to the propagation of the distributions on the graph. That is, given a probability density $\mathbf{w}$, the operation $\mathbf{Pw}$ keeps the resulting distribution normalized.

To clarify the affinity matrix normalization and the diffusion operator construction, we have extended the discussion on diffusion geometry in Appendix A.2.

---

Reference

[1] Kusner et al., (2015). From word embeddings to document distances.

[2] Mikolov et al., (2013). Distributed representations of words and phrases and their compositionally.

[3] Huang et al. (2016). Supervised word mover's distance.

[4] Le et al. (2019). Tree-sliced variants of Wasserstein distances.

[5] Takezawa et al. (2021). Supervised tree-Wasserstein distance.

[6] Yamada et al. (2022). Approximating 1-Wasserstein distance with trees.

[7] Chen et al. (2024). Learning ultrametric trees for optimal transport regression.

[8] Dumitrascu et al. (2021). Optimal marker gene selection for cell type discrimination in single cell analyses.

[9] Du et al. (2019). Gene2vec: distributed representation of genes based on co-expression.

[10] Bellazzi et al., (2021). The gene mover’s distance: Single-cell similarity via optimal transport.

[11] Fakcharoenphol et al. (2003). A tight bound on approximating arbitrary metrics by tree metrics.

[12] Lindenbaum et al., 2020. Gaussian bandwidth selection for manifold learning and classification.

[13] Coifman & Lafon (2006). Diffusion maps.

[14] Shnitzer et al. (2022). Log-Euclidean Signatures for Intrinsic Distances Between Unaligned Datasets.

[15] Shnitzer et al. (2024). Spatiotemporal Analysis Using Riemannian Composition of Diffusion Operators.

[16] Katz et al. (2020). Spectral Flow on the Manifold of SPD Matrices for Multimodal Data Processing.

[17] Mishne et al. (2019). Diffusion nets.

[18] Indyk & Thaper, (2003). Fast image retrieval via embeddings.

from scipy import io

f = io.loadmat("zeisel_data.mat")
X = f['zeisel_data'].T

f = io.loadmat("zeisel_labels1.mat")
y = f['zeisel_labels1'].ravel()

from sklearn.neighbors import KNeighborsClassifier
from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=1)

KNeighborsClassifier().fit(X_train, y_train).score(X_test, y_test)

Z = PCA(n_components=50).fit_transform(X)
X_train, X_test, y_train, y_test = train_test_split(Z, y, test_size=0.3, random_state=1)
KNeighborsClassifier(n_neighbors=10).fit(X_train, y_train).score(X_test, y_test)

This is comment 3 out of 4.

Comparison to Non-TWD Baselines

Thank you for your comment.

• New Naive Baselines in Feature Spaces. Following the reviewer's request, we have conducted additional tests evaluating classification performance directly in the feature space. In the revised manuscript, the new Table 6 in Appendix E presents the document and cell classification accuracy in the feature space using the Euclidean distance and cosine distance. We see that the classification performance relying on generic metrics (Euclidean or cosine) is less effective than our TWD method. This suggests that inferring explicitly the feature hierarchies is important and leads to the effectiveness of distance metric learning for data with a hierarchical structure.

• Additional Non-TWD baselines. We would like to clarify that in addition to comparisons with TWD methods, we also evaluated our method with non-TWD methods that are specifically designed for each type of data, which are considered strong baselines in their respective domains.

  We included Word Mover's Distance (WMD) [1] for word-document datasets and Gene Mover's Distance (GMD) [10] for scRNA-seq datasets. WMD computes the Wasserstein distance using the Euclidean distance $\widetilde{d}$ between precomputed Word2Vec embeddings as the ground metric, while GMD is defined similarly, using the Euclidean distance $\widetilde{d}$ between precomputed Gene2Vec embeddings as the ground metric.
  Each of these methods was designed specifically for a certain type of data (word-document, gene expression, etc.), while our approach is general and infers the latent hierarchy. We present favorable empirical results, demonstrating that our method, while generic and not tailored to specific data types, is highly effective for datasets with hidden hierarchies underlying the features since we infer and utilize the hierarchical structure. We have emphasized this point more clearly in Table 1, Section 5.2, and Appendix D.1 in our revised version.

  Motivated by the reviewer's comment, we have expanded our experiments to include additional evaluations of WMD on word-document datasets and GMD on scRNA-seq datasets, both using cosine distance as the ground metric (the new Table 5 in Appendix E in the revised manuscript). The results show that WMD performs worse with cosine distance as the ground metric compared to Euclidean ground metric, while GMD performs better with ground metric using cosine distance than Euclidean distance. However, regardless of the ground metric (cosine or Euclidean), our method outperforms WMD and GMD. Specifically, when using cosine distance as the ground metric in WMD and GMD, our method achieves higher classification accuracy, with improvements of 4.9% for BBCSPORT, 5% for TWITTER, 3.4% for CLASSIC, 3.7% for AMAZON, 2.2% for Zeisel, and 1.6% for CBMC. These findings emphasize the value of inferring latent feature hierarchies in our method.

  In response to the reviewer's comment on the non-TWD comparison, we have included the above discussion and the additional experiments in the revised version (Table 1, Section 5.2, Appendix D.1, Appendix E.3.6, and Appendix F.5). Thank you again for the valuable comment.

This is comment 2 out of 4.

Clarification of the Novelty of Our Approach Compared to Existing TWD Methods

Thank you for your insightful comment.

In Appendix B, we discussed further details of existing Tree-Wasserstein Distance (TWD) methods and related work and their objectives.

You are correct that, as outlined in the introduction, existing Tree-Wasserstein Distance (TWD) methods do not explicitly estimate the latent feature tree. Instead, these methods aim to approximate a given ground metric $\widetilde{d}$ (often Euclidean) using tree metrics. This allows for efficient computation of optimal transport distances (Wasserstein distances), as the computational complexity of TWD is linear, whereas that of standard optimal transport distances is cubic. For example, the Quadtree [18] is a widely used tree approximation method, that is used to define TWD in, e.g. [18, 5, 6, 7]. To construct the tree metric, the Quadtree method starts by defining a randomly shifted hypercube that contains all feature embedding vectors, and then recursively divides the hypercube into smaller hypercubes with half the side length, until each hypercube contains only a single feature embedding vector. Each hypercube corresponds to a node in the tree. 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 $\widetilde{d}$. For instance, an ultrametric incurs a $\log m$ distortion with respect to Euclidean distance [11], as utilized in UltraTree [7]. Hence, the ultrametric Wasserstein distance [7] 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 a given ground metric $\widetilde{d}$ (e.g., Euclidean), we aim to approximate the Wasserstein distance with a latent underlying hierarchical distance via TWD, where we use the tree to represent the latent feature hierarchy. The key difference in our method lies in how we construct the latent feature tree. Our approach explicitly incorporates the hierarchical structure into the tree estimation process, tailored for the data's underlying geometry.

Following the reviewer's comment, we have updated the text in Section 2 to emphasize that "We assume 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 a meaningful distance between the samples, incorporating the latent feature hierarchy". In addition, we have added the reference in the Related Work section (Section 2) and highlighted it in Section 5.1, pointing to Appendix B, where we introduce the existing methods used for comparison and highlight the novel aspects of our approach.

This is comment 1 out of 4.

We thank the reviewer for the insightful comments and questions. In addition, we thank the reviewer for their appreciation of the technical quality, clarity of the mathematical exposition, and high-quality figures in our paper.

Clarification of Data Representation for TWD Benchmarking

Regarding the data representation, we follow common methodologies in the Wasserstein-related fields. For word-document datasets, we use the benchmarks in [1], which uses normalized bag-of-words and the Word2Vec embedding [2] trained on the Google News dataset as word embedding vectors. These benchmark datasets are studied in the context of Word Mover’s Distance [1, 3] and Tree-Wasserstein Distance [4–7]. For the scRNA-seq application, we use the datasets in [8], where the datasets are provided at the link, and there the CBMC are sparse and normalized by $\log_2(1+X)$. Here, we use the Gene2Vec embeddings [9] as the gene embedding vector following the study in Gene Mover's Distance [10].

To clarify the data representations used in our experiments, we have updated the beginning of Section 5 and included the detailed information outlined above in Appendix D.1.

Rationale Behind the Empirical Comparison to TWD Methods in Table 1

To address "Why do all these methods of tree estimation work so much worse than yours?": for data with a hidden hierarchy, explicitly inferring the geometry of the data, i.e., the hierarchy underlying the features, is crucial and contributes to the effectiveness of the distance learning between samples in downstream tasks on datasets with a hierarchical feature structure, e.g., word-document and scRNA-seq data. Our method constructs the tree to represent the underlying hierarchical structure and incorporates it in the TWD framework. In contrast, existing TWD methods also use a tree, but they do not use it to estimate the latent feature hierarchy. Instead, they only use the tree to approximate a given ground metric $\widetilde{d}$ (often Euclidean) and approximate the Wasserstein distance on the ground metric.

The results in Table 1 demonstrate the importance of accounting for the feature hierarchy as our method yields consistently superior results to the many existing TWD-based methods, none of which accounts for that hierarchy. We have expanded the caption in Table 1 to highlight that the empirical comparisons to TWD methods demonstrate the advantages of our proposed approach.

It is important to note that using the ground metric $\widetilde{d}$ (e.g., Euclidean) in WMD, GMD, or existing TWD methods, results in the (approximated) Wasserstein distance being based on the predefined metric $\widetilde{d}$. In contrast, our method employs an initial distance $d$ to calculate local affinities, computed using the cosine similarity in the experiments. This initial distance is used to compute an affinity matrix within the diffusion operator, which is central to kernel and manifold learning methods. This process leads to a Wasserstein distance based on a latent hierarchical metric induced by $d$. As noted in Appendix D.1, 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 [12]. We have emphasized this point more clearly in Appendix F.5 in our revised version.

We thank the reviewer for the insightful comments and questions, and their appreciation of the motivation behind our method, the tree decoding algorithm in hyperbolic spaces, and the clarity of our exposition.

Definition of the Ground Truth Latent Tree

Thank you for this important comment.

In response to the reviewer’s comment, we have extended the description in Section 4.1 to further clarify the definition of the ground truth latent tree. Specifically, we assume the features $\\{\mathbf{X}\_{:,1}, \ldots, \mathbf{X}\_{:,m}\\} \subseteq \mathcal{H} \subset \mathbb{R}^n$ lie on an underlying manifold $\mathcal{H}$, which is a complete and simply connected Riemannian manifold with negative curvature embedded in a high-dimensional ambient space $\mathbb{R}^n$ with geodesic distance $d_\mathcal{H}$. We view $(\mathcal{H}, d_\mathcal{H})$ as a hidden hierarchical metric space and consider a weighted tree $T = ([m], E, \mathbf{A}, \mathbf{X}^\top)$ as its discrete approximation in the following sense. The tree node $j \in [m]$ is associated with the $j$-th feature, $E$ is the edge set, and $\mathbf{A}\in \mathbb{R}^{m \times m}$ is the weight matrix of the tree edges, defined such that the tree distance $d_T(j, j')$ between two nodes $j$ and $j'$, i.e., the length of the shortest path on $T$, coincides with the geodesic distance $d_\mathcal{H}(j, j')$ between the features $j$ and $j'$. Therefore, we refer to the hidden tree distance $d_T$ as the ground truth distance. Note that the assumption of an underlying manifold embedded in an ambient space $\mathbb{R}^n$ is widely used in studies of data manifolds and manifold learning [1-5], 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 a specification of the manifold to a hierarchical space $\mathcal{H}$ of features, and accordingly, a specification of the graph to a tree $T$. Our aim is to construct a tree $B$ from the data $\mathbf{X}$, such that the TWD between samples using our learned tree $B$ is (bilipschitz) equivalent to the Wasserstein distance using the ground truth latent tree distance $d_T$. To allow this construction of $B$, we assume that local affinities between (node) features $\\{\mathbf{X}\_{:,j}\\}\_{j=1}^m$ that are close in $\mathbb{R}^n$ embody information on the hidden hierarchical structure. More specifically, we assume that the discrete diffusion operator $\mathbf{P}$ constructed from the local Gaussian affinities computed on the features in the ambient space, in the limit of $m\rightarrow \infty$ and $\epsilon \rightarrow 0$, converges pointwise to the heat kernel of the underlying hierarchical metric space $(\mathcal{H},d_{\mathcal{H}})$.

[1] Tong et al. (2021). Diffusion Earth Mover’s Distance and Distribution Embeddings.

[2] Shnitzer et al. (2022). Log-Euclidean Signatures for Intrinsic Distances Between Unaligned Datasets.

[3] Huguet et al. (2022). Manifold Interpolating Optimal-Transport Flows for Trajectory Inference.

[4] Kapusniak et al. (2024). Metric Flow Matching for Smooth Interpolations on the Data Manifold.

[5] Kruiff et al. (2024). Pullback Flow Matching on Data Manifolds.

Improving Clarity on Features, Feature Hierarchy, and Method Setting

To improve the clarity of the problem setup, we have incorporated the reviewer's suggestions into the Introduction section. Specifically, to improve the clarity of "feature" and "feature hierarchy", we have included word-document data in the following illustrative example of high-dimensional datasets with a latent feature hierarchy. In word-document data, samples correspond to documents, and features correspond to words. The hierarchy underlying features (i.e., words) is often assumed because words are naturally organized as nodes in a tree, each connected to others through various hierarchical linguistic relationships. In addition, to improve the clarity of "the measures to which the Wasserstein distance is applied", we have highlighted that we model the data samples as distributions supported on a latent hierarchical structure.

Typo $V$

Thank you for pointing this out. We have updated $V$ to $[m]$ in the revised version.

We thank the reviewer for the insightful and thorough review, and their appreciation of our hyperbolic tree decoding algorithm and its integration with TWD.

Bilipschitz Constants in Theorem 1 and 2

We thank the reviewer for the detailed review and insightful comments.

In our proof, to show the equivalence between the ground truth distance $d_T$ and the inferred distance $d_B$, we used an auxiliary distance on the manifold $d_\mathcal{M}$, and show the equivalence in two steps, leading to $d_B \simeq d_\mathcal{M} \simeq d_T$.

In our initial submission, due to the haste in preparing the manuscript, we mistakenly took the constant from the intermediate equivalence $d_B$ and $d_\mathcal{M}$, while the constants between $d_\mathcal{M}$ and $d_T$ are more involved [1,2].

In the revised version, we have incorporated the reviewer’s suggestion and corrected the statement in Theorem 1: for sufficiently large $K_c$ and $m$, there exist constants $C_1, C_2 > 0$ such that $C_1 d_T \leq d_B\leq C_2 d_T$, and we have revised Theorem 2 accordingly.

Thank you for this important comment.

[1] Leeb & Coifman (2016). Hölder--Lipschitz norms and their duals on spaces with semigroups, with applications to earth mover's distance.

[2] Lin et al. (2023). Hyperbolic diffusion embedding and distance for hierarchical representation learning.

Hyperparameter Tuning

In the revised version, we have incorporated the reviewer’s comment in Section 2 regarding the scale parameter, noting that it is often chosen as the median pairwise distance among data points scaled by a constant factor. We also refer to Appendix D for additional details on hyperparameter tuning.