Fast unsupervised ground metric learning with tree-Wasserstein distance

## Abstract

Abstract The performance of unsupervised methods such as clustering depends on the choice of distance metric between features, or ground metric. Commonly, ground metrics are decided with heuristics or learned via supervised algorithms. However, since many datasets are unlabelled, unsupervised ground metric learning approaches have been introduced. One recent, promising option uses Wasserstein singular vectors (WSV), which emerge when computing optimal transport distances between features and samples simultaneously. While WSV is effective, it has complexity 𝒪 ⁢ ( n 5 ) 𝒪 superscript 𝑛 5 \mathcal{O}(n^{5}) caligraphic_O ( italic_n start_POSTSUPERSCRIPT 5 end_POSTSUPERSCRIPT ) , which is prohibitively expensive in some applications. In this work, we propose to augment the WSV method by embedding samples and features on trees, on which we compute the tree-Wasserstein distance (TWD). We demonstrate theoretically and empirically that the algorithm converges to a better approximation of the full WSV approach than the best known alternatives, and does so with 𝒪 ⁢ ( n 3 ) 𝒪 superscript 𝑛 3 \mathcal{O}(n^{3}) caligraphic_O ( italic_n start_POSTSUPERSCRIPT 3 end_POSTSUPERSCRIPT ) complexity. In addition, we prove that the initial tree structure can be chosen flexibly, since tree geometry does not constrain the richness of the approximation up to the number of edge weights. This proof suggests a fast, recursive algorithm for computing the tree parameter basis set, which we find crucial to realising the efficiency gains at scale. Finally, we employ the tree-WSV algorithm to several single-cell RNA sequencing genomics datasets, demonstrating its scalability and utility for unsupervised cell-type clustering problems. These results poise unsupervised ground metric learning with TWD as a low-rank approximation of WSV with the potential for widespread low-compute application.

## 1 Introduction

The geometry of data structures can help to solve unsupervised learning tasks. For example, clustering methods such as $k$-means compare groupings based on a distance between samples. Inherently, the heuristic that one chooses for the distance matrix or function determines the outcome of the algorithm. So, how should one choose or learn this distance?

From an optimal transport theory perspective, the optimal sample distance is the Wasserstein distance: that is, the minimum cost from the best mapping between samples, weighted by the cost, or ground metric, between their features. For example, samples might be documents, and features words. Then the distance/cost between two documents is determined by the relative expression of words in each, mappings between them, and the ground metric between words. This turns the problem of learning sample distances into one of learning – or choosing – the (feature) ground metric.

State-of-the-art methods usually employ heuristics for the ground metric, commonly Euclidean, with embeddings for the features, e.g. word2vec for documents (Kusner et al., ) and gene2vec in bioinformatics (for identifying cell type from gene expression) (Zou et al., ; Du et al., ). However, embeddings are not always available, or may be difficult to learn. This arises, for instance, when doing single-cell RNA sequencing on a new tissue or animal for which marker genes are not known, or in neuroscience, when large-scale recordings of neurons and imaging of brain areas allow for the matching of neural activity to behaviour or representations (Stringer et al., ; Kriegeskorte et al., ), but discovering the metric on neurons through supervised learning appears intractable.

We are interested in unsupervised ways to learn the ground metric. We do so by augmenting an iterative method by Huizing et al. (). By harnessing the geometrical properties inherent in optimising distances between samples, which are expressed as vectors of features, Huizing et al. () learn ground metrics via positive singular vectors, entirely unsupervised. Their algorithm uses power iterations to first learn distances on samples (expressed in histograms over features) and in turn learn distances on features (histograms over samples).

To illustrate the idea, consider the words “Sentosa” and “Singapore”, which one might expect to be close in word-space. If the metric between documents in the corpus that contain both of these words is close, we learn a smaller distance in the word dictionary between “Sentosa” and “Singapore”. This in turn decreases the distance between documents containing these words. Such an approach could be applied to gene expression (features) in cells (samples), V1 neuronal activity (features) corresponding to representation of images (samples), and other unlabelled high-dimensional data.

While effective, Huizing et al. ()’s algorithm is expensive to run as it requires multiple pairwise distance computations per iteration. Inspired by Yamada et al. (), we propose a method that reduces the algorithmic complexity of each iteration by at least a quadratic factor by representing the data on a tree and learning weights on tree edges, rather than the full pairwise distance matrix. Additionally, we prove a theorem that suggests a powerful, fast recursive algorithm to compute the set of basis vectors for the matrix of pairwise leaf paths on any tree.

The unsupervised tree-Wasserstein method for finding ground metrics offers a geometric low-rank approximation interpretation of the full unsupervised ground metric learning problem. We show that it performs well on a toy dataset, and scales favourably to large 5000-10000 single-cell RNA-sequencing datasets for cell type classification.

### 1.1 Previous and related work

In this section we introduce optimal transport distances, tree-Wasserstein distances, inverse optimal transport and unsupervised ground metric learning.

Optimal transport and optimal transport distances: Optimal transport (OT) can be thought of as a “natural geometry” for probability measures (Peyré and Cuturi, ). The classical OT problem – attributed to Monge () – asks how to find an optimal transport plan between two probability distributions subject to some cost. This is not always guaranteed to have solution, and has non-linear constraints, which makes finding solutions difficult. For this reason, the modern OT usage hinges on a definition by Kantorovich (), allowing probabilistic mappings.

Consider the empiric discrete OT problem. Let $\mu=\sum_{i}^{n}a_{i}\delta_{x_{i}}$ and $\nu=\sum_{j}^{n}b_{j}\delta_{x_{j}}$ be discrete probability distributions (“histograms”, or normalised bin counts), with $\delta_{x_{i}}$ the Dirac delta function at position $x_{i}$. Then the OT problem is to find optimal $P$ satisfying:

$$ $\mathcal{W}_{1,C}(\mu,\nu)=\min_{P\in\mathbb{R}_{+}^{n\times n}}\langle P,C \rangle=\min_{P\in\mathbb{R}_{+}^{n\times n}}\sum_{i}\sum_{j}P_{ij}C_{ij}$ (1) $$

where $P\mathbf{1}_{n}=a,P^{\top}\mathbf{1}_{n}=b$ and $C$ is some cost function. $\mathcal{W}_{1,C}(\mu,\nu)$ is known as the (1)-Wasserstein distance. In this paper, we shall drop the 1 and assume that $\mathcal{W}_{C}$ is induced by the pairwise distance matrix $C$ as its ground metric.

For example, in document similarity, $x_{i},x_{j}$ are words (or embedding vectors) and $a,b$ tell us the frequency of words in a given document $\mu$ (or $\nu$). Usually, $C$ is assumed to be the Euclidean distance between word embeddings, and then $\mathcal{W}_{C}(\mu,\nu)$ is the word mover’s distance (Kusner et al., ). Here, we instead take the approach of learning $C$ in an efficient manner, rather than utilising pre-trained embeddings, for documents represented as bag-of-words.

Computationally efficient OT: Solving the OT problem to compute $\mathcal{W}_{C}$ via linear programming has complexity $\mathcal{O}(n^{3}\log{n})$. Cuturi () suggested using the celebrated Sinkhorn’s algorithm, which speeds up the calculation through entropic-regularised OT to $\mathcal{O}(n^{2})$. Sliced-Wasserstein distance, in which $\mathcal{W}_{C}$ is computed via projection to a one-dimensional subspace, improves on this complexity to $\mathcal{O}(n\log{n})$ (Kolouri et al., ). Sliced-Wasserstein distance is a special case of another geometric embedding, tree-Wasserstein distance (Indyk and Thaper, ; Le et al., ). Tree-Wasserstein distance has linear complexity, and so is faster even than the sliced-Wasserstein distance (Le et al., ).

Tree-Wasserstein distance (TWD): The tree-Wasserstein distance represents samples on a tree. Consider a tree $\mathcal{T}$ with $d_{\mathcal{T}}$ the tree metric (the unique and shortest path between any two nodes on $\mathcal{T}$). Let $\mathbf{w}$ be the set of weights between nodes and let $\mathbf{Z}$ be the tree parameter ($\mathbf{Z}_{ij}=1$ if node $i=Parent(j)$ or itself, and 0 otherwise). Given a tree metric $d_{\mathcal{T}}(\mathbf{x},\mathbf{y})$, the TWD can be written as (Le et al., ; Yamada et al., )

$$ $\mathcal{W}_{\mathcal{T}}(\mu,\nu)=\inf_{\pi}\int_{\mathcal{X}\times\mathcal{Y }}d_{\mathcal{T}}(\mathbf{x},\mathbf{y})d\pi(\mathbf{x},\mathbf{y}).$ $$

Furthermore, TWD takes on a closed analytical form Takezawa et al. ():

$$ $\mathcal{W}_{\mathcal{T}}(\mu,\nu)=||\text{diag}(\mathbf{w})\mathbf{Z}(\mathbf {a}-\mathbf{b})||_{1}$ (2) $$

From (2), TWD can be computed in $\mathcal{O}(\text{size}(\mathbf{w}))$ (Takezawa et al., ). TWD is a good approximation of the full 1-Wasserstein distance (Yamada et al., ).

The inverse OT problem and ground metric learning: Leveraging the geometric nature of the Wasserstein/OT distance for unsupervised learning tasks relies on having a good idea of the ground metric between features that is then “lifted” to the OT distance between samples. Heuristics for ground metrics can be useful, especially with embeddings (Kusner et al., ; Du et al., ), but we are interested in broader, principled ways to find ground metrics.

Inverse optimal transport provides one solution: given a transport plan $P$, find the distance matrix $C$. For matching problems based on recommendation systems, for example the marriage dataset where preferred pairings and the associated features (i.e. the transport plan) are known, inverse OT can find the underlying distance between features. While in general this problem is under-constrained, solutions exist given sufficient constraints on $C$ (Paty and Cuturi, ; Li et al., ; Stuart and Wolfram, ). However, inverse OT requires access to the full transport plan – in practice and for high-dimensional features, such as gene-cell data, that is not feasible.

An alternative is to use partial information about distance or similarity to learn the ground metric, for example with supervised or semi-supervised learning (Cuturi and Avis, ). A related but distinct concept is the idea of supervised learning of sample Wasserstein distances (Huang et al., ) and supervised TWD (Takezawa et al., ); here, the (feature) ground metric is still assumed to be Euclidean, but information about distances is used to learn metrics between samples.

The fully unsupervised approach differs in that neither a feature nor a sample ground metric needs to be assumed.

Unsupervised ground metric learning with Wasserstein singular vectors (WSV): Unsupervised techniques harness the relationship between the geometry of features (embedded in samples) and the geometry of samples (embedded in features) (Paty and Cuturi, ; Huizing et al., ).

Consider some data matrix $X\in\mathbb{R}^{n\times m}_{+}$. Let $a_{i}\in\mathbb{R}^{m}_{+}$ be a sample (normalised row) of $X$, and $b_{k}\in\mathbb{R}^{n}_{+}$ be a feature (normalised column). These are just histograms “embedding” a sample as a distribution over features, and vice versa. For example, the $a_{i}$ (of which there are $n$) could be a histogram of words in a document, and the $b_{k}$ (of which there are $m$) documents containing a given word. As another example, $a_{i}$ might represent cells (samples) as histograms over gene expression, where $b_{k}$ are the genes (features) as histograms over cells that contain each gene.

Huizing et al. () learn ground metric matrices $\mathbb{A}\in\mathbb{R}^{m\times m},\mathbb{B}\in\mathbb{R}^{n\times n}$ satisfying the fixed point equations:

$$ $\mathbb{A}_{k,l}=\frac{1}{\lambda}\mathcal{W}_{\mathbb{B}}(b_{k},b_{l}),\qquad \mathbb{B}_{i,j}=\frac{1}{\mu}\mathcal{W}_{\mathbb{A}}(a_{i},a_{j}),$ (3) $$

$\forall k,l\in\{1,..,m\},i,j\in\{1,...,n\}$. Power iterations are used to compute the ground metric matrices via repeated application of the mapping

$$ $\Phi_{A}(\mathbb{A})_{i,j}:=W_{\mathbb{A}}(a_{i},a_{j})+\tau||{A}||_{\infty}R( a_{i}-a_{j}),$ $$

and symmetrically for $\mathbb{B}$. Then the equivalent problem is to find Wasserstein singular vectors such that there exist $(\lambda,\mu)\in(\mathbb{R}^{*}_{+})^{2}\text{ satisfying }\Phi_{B}(\mathbb{B}
)=\lambda\mathbb{A},\Phi_{A}(\mathbb{A})=\mu\mathbb{B}$.

Maximal singular vectors can be extracted with power iterations:

$$ $\mathbb{A}_{t+1}=\frac{\Phi_{B}(\mathbb{B}_{t})}{||\Phi_{B}(\mathbb{B})_{t}||_ {\infty}},\qquad\mathbb{B}_{t+1}=\frac{\Phi_{A}(\mathbb{A}_{t})}{||\Phi_{A}( \mathbb{A})_{t}||_{\infty}}.$ $$

The complexity of a single power iteration is $\mathcal{O}(n^{2}m^{2}(n\log(n)+m\log(m)))$ (since it requires computing $m^{2},n^{2}$ Wasserstein distances), which Huizing et al. () propose to reduce via stochastic optimisation (which suffers from requiring more iterations to converge) and entropic regularisation, or “Sinkhorn singular vectors” (SSV). The latter improves the time complexity to at best $\mathcal{O}(n^{2}m^{2})$. There are no other known direct reductions in complexity of this fully unsupervised method. Tong et al. () propose an efficient $L_{1}$-embedding of the Wasserstein distance on graphs in the unbalanced OT regime, and show that these embeddings are robust to noise; however, the underlying ground metric is still assumed geodesic, rather than learned.

### 1.2 Contributions

We propose to improve on the WSV computational bottleneck by embedding the dataset in trees and approximating $\mathcal{W}_{d}$ with the TWD. Our findings position Tree-WSV as a low-rank approximation of the full WSV approach with geometric interpretation, advantages for scaling and low-resource efficiency gains. We show that:

- •
One can learn a complete set of non-zero tree edge weights by solving a non-singular system of linear equations (Algorithm 1), and hence solve for the entire Wasserstein distance matrix on features and samples, respectively.
- •
This system is determined by a pairwise leaf-to-leaf paths matrix on the tree whose rank is guaranteed to be equal to the number of edges in the tree or one less the number of edges. This path matrix’s basis set can be found via a fast recursive algorithm.
- •
The Tree-WSV approach is a quadratic order more computationally efficient and outperforms Sinkhorn entropy regularisation (SSV) in terms of accuracy on toy data, as compared to the full WSV.
- •
The approach, using meta-iterations, is scalable to large datasets of over 5000 samples/features, including single-cell RNA sequencing data, where it is also much faster and at least as accurate as SSV.

## 2 Unsupervised tree-Wasserstein distance ground metric learning

We propose to use the TWD (i.e. distance between leaves) to learn the tree ground metric (i.e. distance between features, embedded in the tree structure).

### 2.1 TWD singular vectors approximate full Wasserstein SVs

Let $A=\{a_{i}\},B=\{b_{k}\},i\in\{1,...,n\},k\in\{1,...,m\}$ be the set of rows (samples) and set of columns (features) respectively of the data matrix $X\in\mathbb{R}^{n\times m}_{+}$. Embed each in a respective tree ($\mathcal{T}_{A}$ and $\mathcal{T}_{B}$).

Specifically, let the vectors in $A=\{a_{1},...,a_{n}\}$ be the leaves of $\mathcal{T}_{A}$, $a_{i}$ defined as for WSV, and let $\mathcal{T}_{B}$ have leaves the vectors in $B=\{b_{1},...,b_{m}\}$, as illustrated in Fig. 1. Let $Z_{A},Z_{B}$, $Z\in\{0,1\}^{N\times N_{\text{leaf}}}$ be the tree parameters for each tree, where $[Z]_{i,j}=1$ when node $i$ is the parent of node $j$ or itself ($i=j$), and otherwise is 0. For $\mathcal{T}_{A}$, $N_{\text{leaf}}=n$, and for $\mathcal{T}_{B}$, $N_{\text{leaf}}=m$. Let $\mathbf{w}_{A}\in\mathbb{R}^{N-1}_{+},\mathbf{w}_{B}\in\mathbb{R}^{M-1}_{+}$ be the vectors of edge weights (of size one less than the number of nodes in each tree: $N$ and $M$, respectively).

###### Proposition 2.1 .

The WSV fixed point equations (3) can be expressed on the trees $\mathcal{T}_{A},\mathcal{T}_{B}$ as:

$$ $d_{\mathcal{T}_{A}}(\mathbf{a_{i}},\mathbf{a_{j}})=\frac{1}{\lambda}\mathcal{W }_{B}(\mathbf{a_{i}},\mathbf{a_{j}}),\qquad d_{\mathcal{T}_{B}}(\mathbf{b_{k}} ,\mathbf{b_{l}})=\frac{1}{\mu}\mathcal{W}_{A}(\mathbf{b_{k}},\mathbf{b_{l}}),$ $$

where the singular vector update is to find $\mathbf{w_{A}}$ (and symmetrically $\mathbf{w_{B}}$) such that $\forall i,j$,

$$ $\lambda_{A}\mathbf{w_{A}}^{\top}\Big{(}\mathbf{z_{i}^{(A)}}+\mathbf{z_{j}^{(A) }}-2\mathbf{z_{i}^{(A)}}\circ\mathbf{z_{j}^{(A)}}\Big{)}=\Big{|}\Big{|}\mathrm {diag}(\mathbf{w_{B}})\mathbf{Z^{(B)}}(a_{i}-a_{j})\Big{|}\Big{|}_{1}$ (4) $$

Proof: Appendix A ∎

Figure: Figure 1: Tree embeddings for samples $A$ as leaves in $\mathcal{T}_{A}$ (left) and features $B$ as leaves in $\mathcal{T}_{B}$ (right). The tree metric $d_{\mathcal{T}_{A}}(a_{2},a_{3})$ is shown as the shortest path between these leaves in orange on the left. Equivalently, we can use the relative expression of the features $\{b_{1},b_{2},b_{3},b_{4}\}$ that represent $a_{2},a_{3}$ respectively (as shown in different hues of green on the right) to compute a TWD, $\mathcal{W}_{\mathcal{T}_{B}}(a_{2},a_{3})$, in $\mathcal{T}_{B}$. Thus we assume $d_{\mathcal{T}_{A}}(a_{2},a_{3})$ is equal to $\mathcal{W}_{\mathcal{T}_{B}}(a_{2},a_{3})$.

For intuition about Proposition 2.1, choose two samples $a_{i},a_{j}$ that are leaves in one of the trees; WLOG $\mathcal{T}_{A}$ (Fig. 1). The tree metric $d_{\mathcal{T}}$ is the shortest path between these leaves with edges in $\mathcal{T}_{A}$. Each sample can also be expressed as some histogram over features, which are leaves on $\mathcal{T}_{B}$. The 1-Wasserstein distance between the two samples is some cost matrix $B$ applied to this expression, that is $\min_{P}\sum_{k}\sum_{l}P_{kl}C_{kl}$ where $P\mathbf{1}_{n}=a_{i}$ and $P^{\top}\mathbf{1}_{n}=a_{j}$. From Yamada et al. (), $\mathcal{W}_{B}(a_{i},a_{j})$ is equivalent to the TWD between $a_{i},a_{j}$ on $\mathcal{T}_{B}$, which we can express through summing proportions of $d_{\mathcal{T}_{B}}$ between the leaves in $\mathcal{T}_{B}$: $d_{\mathcal{T}_{A}}(a_{2},a_{3})=\mathcal{W}_{\mathcal{T}_{B}}(a_{2},a_{3})$.

We can utilise Proposition 2.1 to learn low-rank approximations of $\mathcal{W}_{B}(\mathbf{a_{i}},\mathbf{a_{j}})$ and $\mathcal{W}_{A}(\mathbf{b_{k}},\mathbf{b_{l}})$.

Let $y_{ij}=\mathbf{z_{i}^{(A)}}+\mathbf{z_{j}^{(A)}}-2\mathbf{z_{i}^{(A)}\circ z_{
j}^{(A)}}=\text{XOR}\{\mathbf{z_{i}^{(A)}},\mathbf{z_{j}^{(A)}}\}$ and $\mathbf{Y}^{(A)}$ be the tensor of all $y_{ij}$s (i.e. all distinct nodes in the paths of all sets of pairwise leaves). Equation 4 presents a system of $n^{2}$ linear equations, since each $y_{ij}$ is of dimension $N$ and $\mathbf{Y^{(A)}}$ is a tensor of dimension $n\times n\times N$.(^1^11Note that in practice $\mathbf{w}$ is size $N$ because it is equivalent to the tree parameter (number of nodes, i.e. number of edges+1), but the component equal to 0 is dropped (which corresponds to the root node).) The system is over-determined since $N<<n^{2}$, which follows from $\max(N)=2n-1<n^{2}$.

Since Proposition 2.1 is still a singular vector problem, it can be solved with power iterations. It also satisfies the convergence guarantees in Huizing et al. () (up to the goodness of approximation of the true 1-Wasserstein distance), provided non-zero and unique solutions exist. We show in the next section that indeed the system always has a unique solution.

###### Proposition 2.1 .

### 2.2 Solutions to the tree WSV problem exist

From the Rouché–Capelli theorem, for there to exist a unique solution for all components of $\mathbf{w}$, the coefficient matrix to $\mathbf{w}$ in the system of linear equations needs to have rank equal to that of $\mathbf{w}:N-1$ or $M-1$, respectively. The coefficient matrix is the tensor $\mathbf{Y}$ unravelled along $n\times n$ leaf pairs to create a matrix $Y^{\prime}$ of size $n^{2}\times N$. This matrix describes the paths between all pairs of $n$ leaves. Note that $Y^{\prime}$’s rank cannot be greater than $\mathbf{w}$ since the rank is at most $N$, the smallest dimension, but the root node position is 0 for all pairs of leaves, so the rank is at most $N-1$. However, if $Y^{\prime}$’s rank were to be less than $N-1$, there could be more than one solution. Theorem 2.2 proves that we can choose trees for which the rank of $Y^{\prime}$ is $N-1$ (or $M-1$, respectively), ensuring that the problem is well-posed.

###### Theorem 2.2 .

Let the number of nodes in a tree, including the root and leaves, be $N$. Then there can be at most one non-leaf node of degree 2. In addition, if such a node exists, it is the root and $Y^{\prime}$ has rank $N-2$. Otherwise, every degree of the tree is at least 3, and $Y^{\prime}$ has rank $N-1$.

Proof: Appendix B ∎

This theorem implies that as long as we restrict ourselves to trees whose root node has degree greater than 2, we can ensure unique, non-zero solutions to (4), since the rank of $Y^{\prime}$ is guaranteed to be $N-1$, i.e. the length of $\mathbf{w}$. The proof also suggests an efficient recursive algorithm that can also be used to find the basis set for $Y^{\prime}$ , which we call $\mathbf{U}$, and is detailed in Appendix C.

Figure: Algorithm 1 Unsupervised ground metric learning with trees

###### Theorem 2.2 .

### 2.3 Computational complexity and speed-ups

The method detailed above suggests a computationally efficient way of computing Wasserstein singular vectors in Algorithm 1. Each inner power iteration computes Wasserstein distances in time complexity $\mathcal{O}(n)$ rather than $\mathcal{O}(n^{2})$. Further, in each iteration we do not need to compute the whole $m^{2}$ set of distances, but are rather capped to $M<2m-1$, the length of the weight vector $\mathbf{w}$. There are a few additional speed-ups possible, detailed below.

First, note that we only need to construct trees once (without edge weights), using ClusterTree with the number of children $k$ set as $k>2$ (Le et al. (); Indyk and Thaper ()). We compute $Z_{A},Z_{B}$ once for each of $A,B$, and $\mathbf{Z^{(A)}}(b_{k}-b_{l})\,\,\forall k,l$ and $\mathbf{Z^{(B)}}(a_{i}-a_{j})\,\,\forall i,j$ just once.

Similarly, we only learn the $\mathbf{Y}$ tensors once with the tree structure. In fact, we do not need to learn the whole tensor: we just need a square sub-matrix of $Y^{\prime}$ of rank $N-1$, from Theorem 2.2. This is equivalent to learning a basis vector set for the matrix of all pairwise paths between leaves.

#### 2.3.1 Learning the basis vector set for all pairwise leaf-leaf paths

There are two options to learn the basis set. Since $\mathbf{Y}$ is sparse, for small dataset sizes, one can employ scipy sparse methods to find the basis set, say $\mathbf{U}$. Sparse QR decomposition or sparse singular value decomposition (SVD) can be used; in practice sparse SVD performed marginally faster.

Because this operation requires computing a large tensor $\mathbf{Y}$ reshaped into a long rectangular matrix of size $N\times n^{2}$, the sparse method does not scale well with large $n,m$. In this case, we instead created a recursive algorithm as suggested from the constructive proof of Theorem 2.2. This method directly computes a basis set from just the $N\times n$ (nodes by leaves) tree parameter matrix. Recursions are done on at most the number of nodes less the number of leaves ($N-n<n-1$ recursive loops); therefore, it is very fast and resists stack overflow. Stochasticity can be added by allowing randomness in path assignments, but since the result is still a basis set, this has minimal effect on finding the (unique) solution to the linear system of equations. The full recursive function is detailed in Appendix C.

#### 2.3.2 Iteratively updating the weight vectors

Thereafter, the goal is to iteratively learn the single vector $\mathbf{w_{A}}$ satisfying

$$ $\lambda_{A}\mathbf{w_{A}}^{\top}\underbrace{\Big{(}\mathbf{z_{i}^{(A)}}+ \mathbf{z_{j}^{(A)}}-2\mathbf{z_{i}^{(A)}}\circ\mathbf{z_{j}^{(A)}}\Big{)}}_{ \mathbf{y^{(B)}}_{ij}\in\mathbf{Y^{(B)}}\text{ -- learn as basis vectors }U^{( B)}}=\Big{|}\Big{|}\text{diag}(\mathbf{w_{B}})\underbrace{\mathbf{Z^{(B)}}(a_{ i}-a_{j})}_{\text{learn once}}\Big{|}\Big{|}_{1}\forall i,j,$ $$

and symmetrically for $\mathbf{w_{B}}$ in terms of all $k,l$. We use the same basis set indices to compute the TWD on the right-hand side as appears in $\mathbf{U}^{(\textbf{B})}$. A solution exists, and can be found with a linear systems solver such as least squares or non-negative least squares (NNLS), which are both efficient, with cubic complexity on $N$. Each compute of a pairwise TWD is linear. This gives overall complexity per power iteration: $\mathcal{O}(N^{3})<\mathcal{O}((2n-1)^{3})<<\mathcal{O}(n^{5})$, the greatest quantity being the complexity for WSV, for comparison to our method.

In total, we achieve a quadratic order gain in time complexity per iteration. We summarise the full, computationally efficient unsupervised ground metric learning with TWD algorithm (Tree-WSV) in Algorithm 1.

## 3 Experimental results on toy dataset

We demonstrate the improved speed and accuracy of our method as compared to SSV on a synthetic dataset $\mathbf{X}\in\mathbb{R}^{n\times m}$, $n=80,m=60$, as employed in prior work (Huizing et al., ).

### 3.1 Dataset construction and experimental methods

The dataset uses various translations of a unimodal periodic function on a 1D torus (with boundary conditions), such that $\mathbf{X}_{ik}=\exp\Big{(}{-\big{(}\frac{i}{n}-\frac{k}{m}\big{)}^{2}}/{
\sigma^{2}}\Big{)}$. On $\mathbf{X}$, the underlying (and learned) ground metrics look like $\big{|}\sin(\frac{i}{n}-\frac{j}{n})\big{|}$ (and symmetrically for $m$).

We ran experiments 3 times for 100 iterations each on a CPU with different methods, using a translation of the periodic function, to which we added a 50% translation of the torus with half the magnitude. We set $\sigma=0.01$. Thus $\mathbf{X}_{ik}\propto\exp\Big{(}{-\big{(}\frac{i}{n}-\frac{k}{m}\big{)}^{2}}/
{\sigma^{2}}\Big{)}+0.5\Big{(}{-\big{(}\frac{i}{n}-\frac{k}{m}+0.5\big{)}^{2}}
/{\sigma^{2}}\Big{)}$. In order to control against potential gains from tree construction based on preferred data ordering, we randomly permuted the dataset rows and columns before each experiment. We compare performance via computation time and the Frobenius norm of the difference between the learned cost and the metric uncovered by the relevant WSV algorithm (Fig. 2).

### 3.2 Results

Figure: Figure 2: Comparison of mean time complexity (left) and Frobenius norm error (right) for different ground metric learning algorithms. Tree algorithms are specified by restrictions on the minimum number of children per node in the tree. In both, lower is better. Standard error of the mean is shown with black bars. “Many children” refers to ClusterTree initialised with 10 as the minimum children parameter.
Refer to caption: extracted/5993254/time_acc.png

All tree-based methods were faster than both standard WSV and entropic-regularised WSV by 2-3 orders of magnitude (3-children trees took $\sim$1% of the time of the full WSV, for example), as predicted by our theoretical time complexity calculations. Interestingly, the tree-based methods had lower Frobenius norm error when compared to entropic WSV. Tree methods were noted to converge in a similar amount of time (less than 10 iterations) to the full WSV.

As expected from Theorem 2.2, the “many” children trees (at least 10 children) had lower ranks than both 2- and 3-trees, and 3-trees had lower ranks than 2-trees. In general, rank correlated with accuracy and time – that is, trees with fewer children per node computed distance matrices slower, but produced more accurate results (Fig. 2). One exception follows directly from Theorem 2.2: the 2-trees have rank $N-2$ while the 3-trees are rank $N-1$, so the full weight vector cannot be learned for all edges on 2-trees (or $\mathbf{w}$ is not guaranteed to be unique/optimum). As a result, the 2-trees performed worse in terms of accuracy than the 3-trees (Fig. 2 right).

## 4 Experimental results on large genomics datasets

Single-cell RNA sequencing (scRNA-seq) is a powerful biotechnology method that measures gene expression in individual cells (Tang et al., ). scRNA-seq data can be used to determine cell types through clustering, useful for many applications. Clustering reliability across methods is limited by the choice of parameters and uncertainty in the number of clusters (Krzak et al., ). Traditional methods assume a Euclidean metric on PCA embeddings, or use neural network model embeddings (Stuart et al., ; Wolf et al., ; Gayoso et al., ).

More recently, the gene mover distance has been used (Bellazzi et al., ; Du et al., ). Huizing et al. () showed that Sinkhorn (entropy-regularised) singular vectors can achieve state-of-the-art clustering scores on a small scRNA-seq dataset. We illustrate the potential utility of our tree-based method (Tree-WSV) by using it on the same dataset and achieving similar performance, as well as employing it on other datasets with improved performance. Finally, we showcase the Tree-WSV method on a large dataset prohibitive to the standard WSV approach.

### 4.1 Cell-type clustering using genomics data

We first compare performance on the “PBMC 3k” dataset from 10x Genomics (Wolf et al., ). It consists of 2043 peripheral blood mononuclear cells with 6 broad types and 1030 genes. We also ran Tree-WSV on “Neurons V1”, a dataset of scRNA-seq neuronal visual area 1 (V1) cells in mice by Tasic et al. (), consisting of 1468 cells with 7 broad types and 1000 genes, and a subset of the “Human Lung Cell Atlas (HLCA)” scRNA-seq consortium dataset of human respiratory tissue (Sikkema et al., ), consisting of 7000 cells with 4 broad types and 3923 genes. Details on preprocessing, annotation and experimental set-up are identical to Huizing et al. () for PBMC and similar in the other cases, and provided in Appendix D. Experiments were run on GPU (A100 node) with JAX.

The average silhouette width (ASW) metric was used to compare performance on the clustering task. ASW measures how close each data point is to its own cluster compared to other clusters, using an input distance metric. Scores range between -1 and 1, with 1 being better.

### 4.2 Results using Tree-WSV meta-iterations

We noted that silhouette scores for one iteration of the Tree-WSV algorithm – while quick to run – were low across datasets (0.07, 0.15 and 0.05 respectively). This is because the initial tree structure is calculated via ClusterTree, which uses Euclidean distances on the input to cluster nodes based on the incremental furthest search algorithm, resulting in low-depth initial trees. To account for this, we re-ran the full unsupervised TWD ground metric algorithm using the ground metric outputs from the previous run as the distance metric for the initial tree construction. This approach increased the depth of the trees constructed, resulting in improved silhouette scores, but is still entirely unsupervised. Results at 4 iterations and best (max 15 iterations) are summarised in Table 1.

**Table 1: Comparison of average silhouette width (ASW) for distance metrics computed and runtime across scRNA-seq datasets for WSV as compared to tree-based WSV. Euclidean metric for baseline. $\ast$ indicates method did not converge.**
| Dataset (size) | PBMC | (2k) | Neurons V1 | (1k) | Lung | (7k) |
| --- | --- | --- | --- | --- | --- | --- |
| Metric/method | ASW | Runtime | ASW | Runtime | ASW | Runtime |
| Euclidean | 0.073 | N/A | 0.200 | N/A | 0.046 | N/A |
| Entropic WSV (SSV) | 0.348 | 110 min | 0.256 | 70 min | $\ast$ | $>$ 12 hours |
| Tree-WSV (ours - 4 iters) | 0.299 | 8 min | 0.436 | 7 min | 0.104 | 27 min |
| Tree-WSV (ours - best) | 0.313 | 72 min | 0.602 | 32 min | 0.457 | 71 min |

The ASW metrics are comparable in performance to Huizing et al. ()’s Sinkhorn singular vectors approach and better than other alternatives for the PBMC dataset, as presented in Table 1 and Appendix E, with Euclidean distance as a baseline. However, the Tree-WSV method vastly outperforms SSV on a similar task for a dataset of labelled neural cells (Tasic et al., ) and in the HLCA data (Sikkema et al., ). Further, our method has considerably lower runtime. Indeed, on the large HLCA dataset (7000 cells) SSV had not converged after 12 hours. Tree-WSV however can be run on matrices of over 10000 samples/features.

We visualise and validate the cell-type clusters via a 2-component UMAP projection in Figure 3. The clusters for the PBMC dataset are qualitatively similar to those of Huizing et al. () (Fig. 3(a)). Importantly, clusters are visually separable from each other.

Intriguingly, in some cases additional sub-clusters are visible further than the initial annotation would suggest: for example, for GABAergic neurons in the Tasic et al. () neuron dataset (Fig. 3(c)). This likely indicates a subsequent level of detail in the data: the GABAergic sub-type can have further sub-divisions, notably parvalbumin, somatostatin, VIP and other interneurons (Tasic et al., ). These clusters reveal themselves coarsely when we label by them in Figure 3(d). Since Tree-WSV is agnostic to the number of clusters in the data, it is not surprising that more detail can be revealed. This deeper level of complexity should be tested with longer runs and analysis, as it would be advantageous if the algorithm is flexible towards uncovering a greater number of clusters, or on other multi-omics applications (Huizing et al., ).

Finally, while we did not find that ASW scores on meta-iterations had converged after 15 iterations, $\mathbf{w}$ vector solutions in the inner loop tended to always converge quickly (within 10 iterations). Stochasticity inherent from choosing the basis set for the leaf-to-leaf path matrix, as well as the distance used to decide on the initial tree-structure clustering, could both be factors in the meta-algorithm. These are interesting avenues of future theoretical understanding and development that might improve the method.

Figure: (a) PBMCs (Wolf et al., )
Refer to caption: extracted/5993254/umap.png

## 5 Conclusions and perspectives

We present here a new, computationally efficient approach to learning ground metrics in an unsupervised manner by harnessing the tree-Wasserstein distance as an approximation of the 1-Wasserstein distance. The results on toy and real-world data are fast and scale to larger datasets than previously possible, with favourable results. Furthermore, we provide a new geometric underpinning of the restrictions needed for tree structures to flexibly represent data in low-rank approximations. Our method suggests a fast recursive algorithm to construct basis sets for pairwise paths in trees which could have wide-spread applicability. While promising, there is a need for further large-scale testing and development of the Tree-WSV algorithm for other use cases, including neuronal activity data. Lastly, understanding the role of distance metrics in designing the initial cluster tree and stochasticity in the basis set construction could be fruitful avenues for future work.

## Appendix A Proof of Proposition 2.1

First, note that for all discrete measures and distance metrics $d$, we are guaranteed theoretically that there exists a tree for which $\mathcal{W}_{1,d}=\mathcal{W}_{\mathcal{T}}$, and we can construct it by setting $d=d_{\mathcal{T}}$ [Le et al., , Yamada et al., ].

Following Yamada et al. [], define a tree parameter-based vector for every $a_{i},a_{j}$ as $\mathbf{y^{(A)}_{ij}}:=\mathbf{z^{(A)}_{i}}+\mathbf{z^{(A)}_{j}}-2\mathbf{z^{(
A)}_{i}}\circ\mathbf{z^{(A)}_{j}}$ with $\mathbf{z^{(A)}_{i}},\mathbf{z^{(A)}_{j}}\in\mathbf{Z^{(A)}}$, that is $\in\{0,1\}^{N}$ (equivalently we will write this as $\text{XOR}(\mathbf{z^{(A)}_{i}},\mathbf{z^{(A)}_{j}})$), which is known from the tree structure. The tree metric (distance between any two leaves on $\mathcal{T}_{A}$) is then given by:

$$ $d_{\mathcal{T}_{A}}(a_{i},a_{j})=\mathbf{w_{A}}^{\top}\mathbf{y^{(A)}_{ij}}.$ (5) $$

Note that any vector $b_{k}$ is just a distribution over leaves ($a_{i}$) in $\mathcal{T}_{A}$ and vice versa. Using the distribution over the $a_{i}$ with the ground metric on the matrix/tree $\mathcal{T}_{A}$, we can derive Wasserstein distances between any $b_{k},b_{l}$. Since TWD can approximate the 1-Wasserstein distance and admits the closed form in equation 2 [Yamada et al., ], $\mathcal{W}_{\mathcal{T}_{A}}(b_{k},b_{l})=||\text{diag}(\mathbf{w_{(A)}})
\mathbf{Z^{(A)}}(b_{k}-b_{l})||_{1}$.

Now, let us assume that instead of learning the ground metric for the full distance matrix, we want to learn a tree metric $d_{\mathcal{T}_{A}}(a_{i},a_{j})$ (5). The equivalent claim as in Huizing et al. [] is that we achieve the fixed points in Proposition 2.1. More formally, as before, we would like:

$$ $\lambda_{B}\mathbf{w_{B}}^{\top}\mathbf{Y^{(B)}}=\phi_{A}(\mathbf{w_{A}}), \qquad\lambda_{A}\mathbf{w_{A}}^{\top}\mathbf{Y^{(A)}}=\phi_{B}(\mathbf{w_{B}})$ $$

with $\lambda_{A},\lambda_{B},\tau\in\mathbb{R}_{+}$, where $\phi_{A}(\mathbf{W_{A}})_{ij}=\mathcal{W}_{\mathcal{T}_{A}}(a_{i},a_{j})+\tau|
|A||_{\infty}R(a_{i}-a_{j})$, and $\mathbf{Y}$ is the tensor composed of the $n^{2}$ (or $m^{2}$) $\mathbf{y_{ij}}$ vectors. In this case, $\Phi$s map ground costs to Wasserstein distance matrices.

Let $\mathbf{Z^{(A)}}(b_{k}-b_{l})=z_{kl}^{A}$. Since $\mathbf{w_{A}}\in\mathbb{R}^{N}_{+}$ is positive, and letting $|.|$ denote element-wise absolute value,

$$ $\displaystyle W_{\mathcal{T}_{A}}(b_{k},b_{l})$ $\displaystyle=$ $\displaystyle||\text{diag}(\mathbf{w_{A}})\mathbf{z_{kl}^{(A)}}||_{1}$ $\displaystyle=$ $\displaystyle\sum^{M}_{s}\mathbf{w_{A}}^{(s)}|\mathbf{z_{kl}^{(A(s))}}|$ $\displaystyle=$ $\displaystyle\mathbf{w_{A}}^{\top}|\mathbf{z_{kl}^{(A)}}|$ $$

So from (2) and Proposition 2.1:

$$ $\lambda\mathbf{w_{B}}^{\top}\mathbf{y^{(B)}_{kl}}=\mathbf{w_{A}}^{\top}| \mathbf{z_{kl}^{(A)}}|,\qquad\mu\mathbf{w_{A}}^{\top}\mathbf{y^{(A)}_{ij}}= \mathbf{w_{B}}^{\top}|\mathbf{z_{ij}^{(B)}}|$ (6) $$

$\forall k,l\in\{1,..,m\},i,j\in\{1,..,n\}$.

In full, the singular vector update becomes to find $\mathbf{w_{A}}$ such that $\forall i,j$

$$ $\lambda_{A}\mathbf{w_{A}}^{\top}\Big{(}\mathbf{z_{i}^{(A)}}+\mathbf{z_{j}^{(A) }}-2\mathbf{z_{i}^{(A)}}\cdot\mathbf{z_{j}^{(A)}}\Big{)}=\Big{|}\Big{|}\text{ diag}(\mathbf{w_{B}})\mathbf{Z^{(B)}}(a_{i}-a_{j})\Big{|}\Big{|}_{1}$ (7) $$

(and symmetrically for the $\mathbf{w_{B}}$ iteration). ∎

## Appendix B Proof of Theorem 2.2

As preliminaries, we assert that $N>N_{\text{leaf}}=m$, and $\mathbf{Z}\in\{0,1\}^{N\times\ell}$ has rank $m$. Note that $\mathbf{Y}\in\{0,1\}^{N\times l\times\ell}$ is strictly a tensor, but in practice we only care for the reshaped matrix $\mathbf{Y}\in\{0,1\}^{N\times\ell^{2}}$.

Assume that trees cannot have redundant nodes (i.e. nodes with exactly 1 child, not including themselves), and let every tree have root node $r$ and leaf nodes $1,...,\ell$, which we also use to index the associated positions along vectors $z$ and $y$.

Note that at most one inner (non-leaf) node can have degree less than 3, and it must be the root. In this case, the root has degree 2. Otherwise all inner nodes are of degree $\geq 3$.

Some bounds for $N$ in terms of $\ell$ can be derived from the observation that rank is upper-bounded by the number of pairs that two distinct leaves can make, noting $\mathbf{y_{ij}}=\mathbf{y_{ji}}$ and $\mathbf{y_{ii}}=0$. For $\ell\geq 4$,

$$ $l^{2}>\frac{\ell(\ell-1)}{2}\geq 2\ell-2\geq N-1\geq\ell.$ (8) $$

Replacing $N-1$ for $N-2$ we get bounds for $\ell\geq 2$. $2\ell-2\geq N-1$ in (8) is derived from the upper bound on the number of nodes in the maximal spanning tree, a binary tree.

We prove by induction on the number of tree nodes that the rank of $\mathbf{Y}$ is $N-1$ for any tree structure in which all nodes have degree $\geq 3$, and $N-2$ otherwise, for any $\ell\geq 2$.

$N-2$ base case, $N=3$: The tree can only be a root connected to two leaves. There is only one $\mathbf{y}$-vector, between the two leaves. So the rank is $N-2=1$.

$N-1$ base case, $N=4$:
$r$ connects to each of the three 3 leaves. In general for $N=\ell+1,\ell\geq 3$, we can show the rank is $N-1=\ell$. WLOG, choose a leaf $i$ and consider $\mathbf{y_{i,j}},\forall j\neq i$. Since every $\mathbf{z_{j}}$ is 0 everywhere except for positions indexed by node $i$ and root $r$, where it is 1, $\mathbf{y_{i,j}}$ is 1 at $i,j$ and 0 otherwise. Therefore the set of $\ell-1$ vectors $\mathbf{S}_{i}=\{\mathbf{y_{i,j}},i\neq j,\forall j\}$ is linearly independent. Additionally, we can choose any two distinct nodes $j,k$ where $j\neq k\neq i$, and $\mathbf{S}=\{\mathbf{y_{j,k}}\cup\mathbf{S}_{i}\}$ must also be a linearly independent set. This follows because $\mathbf{y_{j,k}}$ is 1 at nodes $j,k$ and 0 otherwise; and no linear combination of $\mathbf{y_{i,k}},\mathbf{y_{i,j}}$ (with $i=1$ for both vectors, which should be cancelled) can reproduce $\mathbf{y_{j,k}}$.

Can we add any more vectors from $\mathbf{Y}$ to $\mathbf{S}$ such that the set is still linearly independent? No. To show this, assume that there exists some $\mathbf{y_{g,h}}\notin\text{span}(\mathbf{S})$. Trivially, $g\neq h\neq i$ and $\{g,h\}\neq\{j,k\}$, so $\mathbf{y_{g,h}}$ is 0 everywhere except at the $g,h$ positions. But note that we can write $\mathbf{y_{g,h}}$ using vectors in $\mathbf{S}$:

$$ $\mathbf{y_{g,h}}=\mathbf{y_{i,g}}+\mathbf{y_{i,h}}-\big{(}\underbrace{\mathbf{ y_{i,j}}+\mathbf{y_{i,k}}-\mathbf{y_{j,k}}}_{=\text{2 at $i$; 0 otherwise}} \big{)}$ (9) $$

So our assumption was incorrect and $\mathbf{Y}\in\text{span}(\mathbf{S})$. Therefore $\text{rank}(\mathbf{Y})=\ell=N-1$.

Induction hypothesis, $N\leq m$: Assume that the rank of any tree structure with parameter $N\leq m,m\in\mathbb{N}$, has rank $N-1=m-1$ if all inner nodes are degree 3 or more, and $N-2=m-2$ otherwise. We call this a $m$-tree.

Inductive step, $N=m+1\geq 4$: We define a set of “adjacent” leaf nodes $\mathbf{A}$ to be the set of all leaves connected to the same parent node $\rho$ where $\rho$ has no non-leaf children.

Note first that we can construct any valid [$m+1$]-tree structure from some previous valid $m$-tree: for any [$m+1$]-tree, consider a complete set of adjacent leaf nodes of depth at least 2 from $r$ – that is, all the leaves connected to the same non-root node $\rho$ (if depth is less than 2, we get the extended base case). These leaves have common ancestry, and their $\mathbf{z}$-parameters differ only at the indices of the leaves themselves. Since $n>1$, there must exist at least one node $\rho$ that is neither $r$ nor a leaf. Collapse the $\rho$-sub-tree by removing $\rho$ and reconnecting its leaves to the immediate parent node of $\rho$. Then the new tree is an $m$-tree (we have lost one node), as claimed.

Every tree must have at least one such set $\mathbf{A}$, and it must have at least 2 elements. We consider the smallest $\mathbf{A}$-sub-tree and the $\mathbf{L}/\mathbf{A}$-sub-tree consisting of the rest of the tree, in two cases.

Case 1, $|\mathbf{A}|\geq 3$: From the base case, the rank of the $\mathbf{A}$-sub-tree up to the parent node $\rho$ is $|\mathbf{A}|$. Now consider the sub-tree consisting of all the other leaves in the tree, $\mathbf{L}/\mathbf{A}$, up until $\rho$, where the two sub-trees connect. Imagine $\rho$ is a leaf (ignore its children $\mathbf{A}$). Then the $\mathbf{L}/\mathbf{A}$-sub-tree has tree parameter $N_{\mathbf{L}/\mathbf{A}}=m+1-\mathbf{A},\leq m$, so by the inductive hypothesis it has rank $N-1-\mathbf{A}$ or $N-2-\mathbf{A}$, depending on the degree of $r$.

How do we join the sub-trees? Let links including $\rho$ now extend to some (the same) new leaf in $\{\mathbf{A}\}$, say $a_{i}$. The union of the $\mathbf{y_{i,j}}$ making up each sub-tree’s basis set is linearly independent. To see this, note that the only $\mathbf{y_{i,j}}$ that have a 1 at position $\rho$ were those ending at $\rho$ the leaf on $\{\mathbf{L}/\mathbf{A}\}$. So to express any $\mathbf{y_{l_{k},a_{i}}}$, a link via $\rho$ is needed, that is the rank is lower-bounded by $N-1$ or $N-2$, respectively.

We cannot add any more $\mathbf{y}$s: since the sub-trees are saturated by assumption, any addition $\mathbf{y}$ should also be a link, but any link can be expressed via an existing link through $\rho$ and combinations of the sub-tree basis sets (following the argument for the extended base case $N>3$, since we know that each sub-tree has more than 3 nodes). Therefore rank([$m+1$]-tree) $=N-1$ or $N-2$ based on the sub-tree, as required.

Case 2, $|\mathbf{A}|=2$: In this case, the $\mathbf{L}/\mathbf{A}$-sub-tree has rank $N-3$ or $N-4$ and the $\mathbf{A}$-sub-tree has internal rank 1 since there is just one pair of leaves to make $\mathbf{y_{a_{1},a_{2}}}$. WLOG, assume the $\rho$ links extend to $a_{1}$, with $\{\mathbf{y_{\ell_{i},a_{1}}}\}$ for some (not necessarily single) $i$, necessarily linearly independent of $y_{a_{1},a_{2}}$. The total rank is then $N-2$ or $N-3$.

In the 2-case, we can also create another $\mathbf{y_{l_{j},a_{2}}}$, where $j$ can either be the same as $i$ or different.

How do we know that $\mathbf{y_{l_{j},a_{2}}}$ is not in the span of the two sub-tree’s $\mathbf{y}$ vectors? Assume it is in the span. $\mathbf{y_{l_{j},a_{2}}}$ has 1s at $\rho$. The only other $\mathbf{y}$s which shares this property are from $\mathbf{y_{l_{i},a_{1}}}$. So, if $\mathbf{y_{l_{j},a_{2}}}$ could be written using already chosen $\mathbf{y}$s, its decomposition must include $\mathbf{y_{l_{i},a_{1}}}$. But since $a_{1}$ is known to be zero in $\mathbf{y_{a_{2},l_{j}}}$, but is non-zero in the path containing $\mathbf{y_{l_{i},a_{1}}}$, $a_{1}=1$ should be “removed” through linear combinations with other $\mathbf{y}$s; however, there is just one $\mathbf{y}$ involving $a_{1}$, which is $\mathbf{y_{a_{1},a_{2}}}$. So we must subtract $\mathbf{y_{a_{1},a_{2}}}$, resulting in a value of $-1$ at the $a_{2}$ position. But we require $a_{2}$ to be 1: there is no linear combination that allows both the ancestral node $\rho$ and $a_{2}$ to be 1.

So we must be able to add one link, and so here, too, the rank is $N-1$ or $N-2$ in total (trivially any further $\mathbf{y}$s must be links and so would be in the span).

∎

## Appendix C Recursive algorithm for tree basis calculation

Figure: Algorithm 2 Tree basis set recursion

## Appendix D Details on genomics data preprocessing and experiments

### D.1 Data preprocessing

All single-cell RNA sequencing data was processed using standard protocols, including CPM-normalisation, log1p-transformation, and selection of genes based on highest variability. Canonical markers were added according to source.

PBMC3K data was recovered using Scanpy [Wolf et al., ]. Preprocessing procedure/code was provided on request from Huizing et al. [] and included cell type annotation using Azimuth marker genes [Hao et al., ], sampling only annotations for which confidence was over 90%. The results of SSV on this dataset were identical to that reported previously [Huizing et al., ].

Neural cell data and labels from Tasic et al. [] were recovered from the NCBI GEO repository, Series GSE71585 [Clough et al., ] and annotated using the Allen Brain Map [Lein et al., ] and CZ CELLxGENE [Abdulla et al., ].

The Human Lung Cell Atlas (HLCA) was recovered, and annotated, from CZ CELLxGENE [Abdulla et al., ], and originally published by Sikkema et al. []. Because taking a subset might disrupt the original annotation, low-gene-count cells and cells with labelling entropy score less than 0.1 were removed.

### D.2 Algorithm parameters

Genomics experiments using Tree-WSV were run for 20 iterations of the internal linear system of equations singular vector loop (finding $\mathbf{w}$) and 15 meta-iterations of the entire algorithm (starting with constructing new trees based on the 2 previous weight matrices, through to computing 2 new TWD-based distance matrices). A JAX random seed of 0 was used. The ASW and total runtime after the 4th meta-iteration and the best score overall (usually the 12-15th meta-iteration) were reported.

Genomics experiments using SSV were run for 15 iterations of the singular vectors loop with $\tau=0.001,\epsilon=0.1$, and follow the settings and replicates the results in Huizing et al. [].

### D.3 Computational resources

Computations on CPU were done using a Apple M2 MacBook Air M2 with 16 GB RAM. GPU computations were performed on on an NVIDIA A100 node with 16 GB requested memory.

## Appendix E Comparison to other methods for improving cell-type clustering

We reproduce the table from Huizing et al. [] showing other methods that can be used to cluster cell-type data from scRNA-seq experiments as compared to ours on the PBMC dataset.

**Table 2: Comparison of average silhouette width (ASW) for distance metrics computed and runtime in the PBMC dataset.**
| Dataset (size) | PBMC | (2k) |
| --- | --- | --- |
| Metric/method | ASW | Runtime |
| PCA / $\ell^{2}$ | 0.238 |  |
| Kernel PCA / $\ell^{2}$ | 0.241 |  |
| scVI embedding / $\ell^{2}$ | 0.168 |  |
| Sinkhorn | 0.003 |  |
| Gene Mover Distance | 0.066 |  |
| Euclidean | 0.073 |  |
| Entropic WSV (SSV) | 0.348 | 110 min |
| Tree-WSV (ours - 4 iters) | 0.299 | 8 min |
| Tree-WSV (ours - best) | 0.313 | 72 min |

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 1
questions: - In line 194, what does “$a_i$ defined as for WSV” mean?
- What is the size of $\mathbf{U}$ in line 259?
- How does the choice of tree construction affect the proposed method? If a different tree construction method were used instead of ClusterTree, would it impact the method's outcome?
- How to decide whether the algorithm reaches convergence in Algorithm 1?
- How fast does the proposed algorithm converge? How does the performance change across the iteration?
- What are the hyperparameters of the proposed method? How sensitive are they in the experiments?
- What is the Euclidean metric baseline in Table 1?
- Why is WSV not considered as a baseline in Table 1?
- Why are other TWD methods not considered baselines in Table 1?
soundness: 3
strengths: The proposed method demonstrates efficiency compared to WSV and SSV presented in Huizing et al. (2022).
summary: The paper introduces Tree-WSV, integrating TWD with WSV (Huizing et al., 2022).
weaknesses: - The authors missed important related works [1, 2, 3, 4, 5] that consider the relationships between the samples are informed by the relationships between the columns, and vice versa, for Wasserstein distance and specifically in tree-related settings [2,5]. Specifically, the setup of randomly permuted dataset rows and columns in the toy datasets was one of the important tasks in these works.
- The explanation of how the Wasserstein distance can serve as a tree distance in Proposition 2.1 is unclear. It’s also not evident whether there exists a tree for which the tree distance would correspond to a Wasserstein distance.
- The real-world application is restricted to single-cell RNA sequencing data, despite the introduction citing various data types as motivation.
- The experiment section includes only two competing methods without considering the baselines used in WSV or other distance metric learning approaches.
- $n$ is not yet defined in the Abstract. It’s unclear what does it represent here. In addition, the computational complexity for WSV reported by the author differs from that presented in the original WSV paper.
- In Section 1.1, the authors mention the sliced Wasserstein distance. However, it’s unclear what it is here and how SWD is a special case of TWD. Also, it’s unclear why SWD is not considered an alternative for efficient computation for Wasserstein distance in the WSV framework.
- It’s unclear $\mathbf{w}$ and $\mathbf{Z}$ in line 109 are vector or matrix. Additionally, it’s unclear what is the connection between $\mathbf{x}$, $\mathbf{y}$, $\mu$, and $\nu$ in line 111.
- The notation of Wasserstein distance $\\mathcal{W}\_C$ in line 92 and the notation of TWD $\\mathcal{W}\_\\mathcal{T}$ in line 111 are confusing. In the formal $C$ is a pairwise distance matrix, and in the later $\\mathcal{T}$ is a tree.
- It’s unclear what are $\mathbf{a}$ and $\mathbf{b}$ in Eq.(2). In addition, it’s unclear what size($\mathbf{w}$) represents in line 116.
- More details are needed for how “TWD is a good approximation of the full 1-Wasserstein distance”. What do the authors refer to as “approximation”? What is the relation between TWD and full 1-Wasserstein distance?
- It’s unclear what is $R$ in line 151.
- The authors keep using the term “basis” throughout the paper, e.g.,  tree parameter basis set, the set of basis vector, matrix’s basis set. However, it’s unclear what does it represent in these contexts.
- It’s unclear what "full WSV" represents. Is it different from WSV?
- It’s unclear why the size of the vectors of edge weights is less than the number of nodes in the tree in Section 2.1.
- It’s unclear what $\mathcal{W}_A$ and $\mathcal{W}_B$ represent in Proposition 2.1. Here, $A$ and $B$ are the sets, which are not pairwise distance matrix nor tree as in previous notations.
- It’s unclear what $\circ$ denote in Proposition 2.1. Also, it’s unclear what are $\\lambda\_A$, $\\mathbf{z\_i}\^{(\\mathbf{A})}$, and $\\mathbf{Z}\^{\\mathbf{B}}$.
- The proof for Proposition 2.1 in Appendix A is very hard to follow. The notations used are not consistent with those used in the main texts. The newly defined notation is very dense. Also, it’s unclear what are $\\mathbf{W}\_{\\mathbf{A}}$, $\\Phi\_{\\mathbf{S}}$
- More details and explanations are needed for how Theorem 2.2 supports unique and non-zero solutions in Proposition 2.1.
- Algorithm 1 is very hard to follow. For example, it’s unclear what the line “$\\mathbf{Z}\_{\\mathbf{diff}}$ …… “ represents. It’s unclear what are $\\mathbf{A}\_{leaf}$, $\\mathbf{w}\_{\\mathbf{B}}$(prev)
- The reference style is inconsistent: some entries lack publisher information, some links are not official paper links, and "Wasserstein" is sometimes written with a lowercase "w."
- The notation style is inconsistent: vectors and matrices are inconsistently represented, with a mix of boldface and regular type. The notation for the tree parameter in Section 1.1 is different than in Section 2.1

## Minor
- Missing “-” for tree-Wasserstein distances in line 077 and line 102
- The acronym "OT" is used without being defined first
- Missing punctuations in equations
- It’s unclear what is $TB$ in line 214

[1] Ankenman, J.I., 2014. Geometry and analysis of dual networks on questionnaires. Yale University.

[2] Mishne, G., Talmon, R., Cohen, I., Coifman, R. R., & Kluger, Y. (2017). Data-driven tree transforms and metrics. IEEE transactions on signal and information processing over networks, 4(3), 451-466.

[3] Gavish, M. and Coifman, R.R., 2012. Sampling, denoising and compression of matrices by coherent matrix organization. Applied and Computational Harmonic Analysis, 33(3), pp.354-369.

[4] Shahid, N., Perraudin, N., Kalofolias, V., Puy, G. and Vandergheynst, P., 2016. Fast robust PCA on graphs. IEEE Journal of Selected Topics in Signal Processing, 10(4), pp.740-756.

[5] Yair, O., Talmon, R., Coifman, R.R. and Kevrekidis, I.G., 2017. Reconstruction of normal forms by learning informed observation geometries from data. Proceedings of the National Academy of Sciences, 114(38), pp.E7865-E7874.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: List of typos. I suggest that the authors give the manuscript a thorough proof-reading.

* l.39 confusing, why is the optimal sample distance the Wasserstein distance?
* l.106 in what sense is SWD a “geometric embedding”?
* l.111 definition of pi?
* l.116 what does “good approximation” mean here?
* l.140 precise the meaning of normalized.
* l.142 distribution -> “probability distribution”
* l.150 R and tau are undefined
* l.150 not clear what \Phi_A is. Huizing et al. (2022) describe Phi_A as ”lifts a ground metric to a pairwise distance matrix”. The authors need to explain the definition of Phi_A.
* l.151 equivalence to (3) assumes tau = 0. The authors need to be a bit more careful in their recap of Huizing et al. (2022).
* l.157 not sure that the remark in parentheses is correct, according to Huizing et al. (2022) a single Wasserstein iteration is n^3 log n. Also how many distances do we compute when we compute “m^2, n^2 W. distances”? Is is m^2 + n^2, m^2 * n^2, something else?
* l.201 W_{T_B} instead of W_B ?
* l.205 z_i^(A) undefined
* l.205 Z^{(B)} or Z_B?
* l.214 W_{TB} -> W_{T_B}
soundness: 3
strengths: The idea is sound, and there is a significant element of originality, especially in the development of the algorithm in Appendix C.
summary: The authors apply the unsupervised metric learning algorithm from Huizing et al. (2022) (which used Wasserstein and entropy regularized Wasserstein) to learn ground metrics for spaces of histograms based on the Tree Wasserstein distance.
weaknesses: The paper seems rushed overall, there are a large number of typos and a number of results that should have been presented as Theorems are merely stated informally.

1. l.238 – 241 these statements require a proof. Especially the convergence, since it was already somewhat delicate in Huizing et al. (2022).
2. l.263 “Wasserstein” -> “Tree Wasserstein” or else requires a proof.
3. I find Theorem 2.2 hard to interpret. Can the authors rephrase the interpretation in the next paragraph (l. 256-l.260) as a Theorem and include the current Theorem 2.2 as a Lemma?
4. l.354 The previous work used n=100, m=80 and not n=80, m=60 as this paper states.
5. l.367 Why are you using a different metric (Frobenius norm) than the original work? How does your method compare when using the same metric d_h(B, B’) = ||log(B/B’)||_V?

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: Refer to weaknesses.
soundness: 3
strengths: 1. The paper offers rigorous theoretical support for the TWD method, with proofs on the uniqueness and existence of solutions within specific tree configurations. 
2. The empirical results are presented clearly, with comparative metrics that directly illustrate the computational runtime saving and clustering performance.
3. The paper provides a solid background review on optimal transport theory and the tree-Wasserstein distance.
summary: The paper presents a novel approach for unsupervised ground metric learning using Tree-Wasserstein Distance (TWD) as a low-rank approximation of the computationally intensive Wasserstein Singular Vector (WSV) method. The proposed method embeds samples and features in tree structures, reducing the computational complexity from O(n⁵) in traditional WSV to O(n³) by learning distances between data points as TWD on trees. Empirical results indicate that the method achieves similar or better clustering accuracy compared to Sinkhorn singular vectors (SSV) while maintaining much faster runtimes.

This paper is the improved version of the workshop paper “Unsupervised Ground Metric Learning with Tree Wasserstein Distance”. The primary innovation of this work is adding recursive basis set computation for tree-based WSV.
weaknesses: 1. Although the ClusterTree algorithm plays a significant role in the tree structure initialization, there is limited background provided on how it operates, what assumptions it makes, or its typical applications. I reviewed both references—Le et al. (2019) and Indyk & Thaper (2003)—but did not find any mention of a ClusterTree. Could the author be referring to the ‘Partition_Tree_Metric’ described in Le et al. (2019)?
2. The algorithm section mentions differences in handling ‘large’ and ‘small’ datasets but does not specify the boundary between the two. What happens if either m or n is very large while the other meets the ‘small’ criteria?
3. The paper’s notation is sometimes inconsistent, making it challenging to reference equations or terms precisely. For example, on line 201, a_i and a_j are bold, but on line 205, a is not bold. On line 211, what does the cost matrix B represent?
4. Figure 3 could be better organized. The paper does not provide a comparison of how other metrics perform on these datasets.
5. Line 557, The URL is invalid.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: **Major**:

- How is the condition number of eq. 5? Does the system become ill-conditioned as matrix becomes large?
- Section 2.3, what is the convergence rate of the proposed iterative method? How it is affected by the data set size?
- What is the complexity of Cluster Tree used here?
- Please add details of memory consumption for your proposed iterative algorithm and approximation to SVD.
- I think the authors should make more effort to show the claimed complexity $\mathcal{O}(N^3)$. Consider one of the data set, and test the method against benchmark for a range of $N$.

**Minor**:

- In abstract, "...a fast and recursive algorithm..." and not "a fast, recursive algorithm…”
- In abstract, what does "low-compute application" mean?
- P2, l82: "... to have a solution" and not "to have solution"
- P3, l108: Is there a guarantee that the shortest path between any two nodes on a tree is unique?
- Figure 1: Caption refers to b5, but I don’t see b5 in the graph.
- In the references, make sure to use capital letters when needed, e.g. use “Wasserstein” instead of “wasserstein”
soundness: 3
strengths: I think the paper is proposing a novel idea for estimating ground metric learning of unlabeled data set, that reduces the complexity of each iteration from $\mathcal{O}(N^5)$ to $\mathcal{O}(N^3)$.
summary: The paper proposes an interesting variation of Wasserstein singular vectors by embedding samples and features on a tree. By doing so, they claim to have achieve a cubic complexity as opposed to quintic complexity of the standard method. While the authors show interesting results, I believe the manuscript can be accepted after a major revision.
weaknesses: Some important details of the proposed method are missing in the paper. This includes convergence rate, well-posedness, and memory consumption of the proposed method. See my questions/comments below.

Meta-review: The paper introduces a computationally efficient method for unsupervised ground metric learning using the tree-Wasserstein distance as a low-rank approximation of Wasserstein singular vectors. The key points raised during the reviewer discussion included concerns about theoretical guarantees, computational complexity, and real-world applications. Specifically, reviewers identified unclear explanations of the relationship between tree distances and Wasserstein metrics, as well as requests for comparisons with additional baselines and further theoretical justification. The authors addressed these concerns by clarifying the notation, improving the theoretical analysis (e.g., existence and convergence of solutions), and adding empirical results for alternative scenarios. These revisions significantly strengthened the paper’s presentation and rigor, leading to the decision to recommend its acceptance as a poster.

Decision: Accept (Poster)