Fast unsupervised ground metric learning with tree-Wasserstein distance
Authors: Kira Michaela Düsterwald
Samo Hromadka
Makoto Yamada
Venue: ICLR 2025 Poster (2025)
Keywords: unsupervised learning
optimal transport
distance-based learning
clustering
trees
wasserstein distance
Primary area: learning on graphs and other geometries & topologies
TL;DR: Unsupervised ground metric learning with tree-based optimal transport is computationally efficient yet still more accurate than alternative approximation appraoches, and scales favourably on genomics datasets

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 interesting datasets are unlabelled, unsupervised ground metric learning approaches have been introduced. One promising option employs Wasserstein singular vectors (WSVs), which emerge when computing optimal transport distances between features and samples simultaneously. WSVs are effective, but can be prohibitively computationally expensive in some applications: $\mathcal{O}(n^2m^2(n \log(n) + m \log(m))$ for $n$ samples and $m$ features. 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 standard WSV approach than the best known alternatives, and does so with $\mathcal{O}(n^3+m^3+mn)$ 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 and 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 application.

Published as a conference paper at ICLR 2025
FAST UNSUPERVISED GROUND METRIC LEARNING WITH TREE- WASSERSTEIN DISTANCE
Kira M. Düsterwald Samo Hromadka', and Makoto Yamada²
I Gatsby Computational Neuroscience Unit, University College London, United Kingdom
2 Machine Learning and Data Science Unit, Okinawa Institute of Science and Technology, Japan
kira. wra.ddsttewwaad 21216... aacl.ac.u..
ABSTRACT
The performance unsupervised methods such as clustering depends on the choice of distance metric between features, ground metric. Commonly, ground metrics are decided with heuristics or learned via supervised algorithms. How- ever, since many interesting datasets unlabelled, unsupervised ground met- ric learning approaches have been introduced. One promising option employs Wasserstein singular vectors (WSVs), which emerge when computing optimal transport distances between features and samples simultaneously. WSVs are ef- fective, but can be prohibitively computationally expensive in some applications: ( log(n) + log(m)) for samples and m features. In this work, we propose to augment the WSV method by embedding samples and features on trees, which we compute the tree-War FreeWWasersstein distance (TWD). We demonstrate the- oretically and empirically that the algorithm converges to a better approximation of the standard WSV approach than the best known alternatives, and does so with O(n + m + mn) complexity. In addition, we prove that the initial tree struc- ture 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 and 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 several single-cell RNA sequencing genomics datasets, demon- strating 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 application.
INTRODUCTION
The geometry of data structures can help solve unsupervised learning tasks. For example, cluster- ing methods such as k-means compare groupings based on a distance between samples informed by distance features. Inherently, the heuristic that one chooses for the distance matrix function determines the outcome of the algorithm. So, how should one choose or learn this distance?
k-means gives equal weighting to the distance between each pair of features. From an optimal transport theory perspective, it could be argued that a more holistic sample distance is given by the Wasserstein distance: the minimum cost using an optimal 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. Hence the problem of learning sample distances is determined by the choice of 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. (2015) and gene2vec in bioinformatics (for identifying cell type from gene expression; Zou et al. (2019) Du et al. (2019). However, embeddings are not always available, or may be difficult to learn. This arises, for instance, when doing single-cell RNA sequencing on new tissue for which marker genes are not known.
We are interested in unsupervised ways to learn the ground metric. We do so by augmenting an iterative method by (Huizing et al. (0022) By harnessing the geometrical properties inherent in

---

Published as a conference paper at ICLR 2025
optimising distances between samples. which expressed vectors of features, Huizing et al. (2022) learn ground metrics via positive singular vectors, entirely unsupervised. Their algorithm uses power iterations 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 close, we learn 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. (2022)s 's algorithm expensive to run as requires multiple pairwise distance computations per iteration. Inspired by Yamada et al. (2022), we propose a method that reduces the algorithmic complexity of each iteration by at least a quadratic factor by representing the data on tree and learning weights tree edges, rather than the full pairwise distance matrix. Additionally, we prove a lemma that suggests a powerful, fast recursive algorithm to compute the set of basis vectors for the matrix of pairwise leaf paths on any tree. We provide source code at: com/ ki darradust/ttree-t
The unsupervised wessesssenn method for finding ground metrics offers geometric low-rank approximation interpretation of the full unsupervised ground metric learning problem. We show that it performs well 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- free-Wasserstein distances, inverse optimal transport and unsupervised ground metric learning.
Optimal transport: Optimal transport (OT) can be thought of as "natural geometry" for probabil- ity measures (Peyré& & Cuturi, 2019). The classical OT problem attributed to Monge (188I) asks how to find an optimal transport plan between two probability distributions subject to some cost. This is not always guaranteed to have a solution, and has non-linear constraints, which makes find- ing solutions difficult. For this reason, the modern OT usage hinges definition by Kantorovich [940). allowing probabilistic mappings.
Consider the empiric discrete OT problem. Let μ probability distributions ("histograms", normalised bin counts), with dm, the Dirac delta function at position Tg. Then the OT problem is to find optimal P satisfying:
Σ. a, and Eybb be discrete
Wi,o(u,u) = min (P,C)= PEREX
EIPiO minn
(1)
PER X
3
where a, 1, and C E R is some cost function. (1, v) is known as the (1)- (a--asssrsteiin distance. In this paper, we shall drop the I and assume that Wo is induced by the pairwise distance matrix C' its ground metric.
For example, in document similarity, xt 2y are words (or embedding vectors) and a (resp. b) tells us the frequency of words in a given document μ (resp. v). Usually, C is assumed to be the Euclidean distance between word embeddings, and then We-u,, v) is the word mover's 's distance (Kusner et al., 2015). 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 Wc via linear programming has complexity (n n) (Peyré & Cuturi 2019). Cuturi (003)) suggested using the celebrated Sinkhorn's algorithm, which speeds up the calculation through entropic-regularised OT to (nn²).
distance (SWD). in which We is computed via projection to one-dimensional subspace, improves on this complexity to On log (Kolouri et al. 2016). SWD can be expressed special case of geometric embedding, tree- Wasserstein distance (TWD), when the tree struc- ture comprises a root connected to leaves directly (Indyk & Thaper
2003 Le et al. 2019). TWD
has linear complexity, and so is faster even than SWD (Le et al. 2019).
2

---

Published as a conference paper at ICLR 2025
Tree- Wasserstein distance (TWD): The tree- rree-asserstein distance represents samples tree. Consider a tree T with tree metric dr (the unique and shortest path between any two nodes on T). Let w E - be the vector of weights between nodes and Z E {0, 1 1)x the tree parameter (where N is the number of nodes in the tree and NNear the number of leaves), defined by Zr, = 1 if the edge from a non-root node is from a parent of leaf j or i = j. and 0 otherwise. Given tree metric dy on leaves x, y and transport plan between measures μ4, v, the TWD can be written as (T(,,,) v) = in.xeU(H.)) Jx dr (2e, y dT (at, y). where Up, v) { E M2 (X' x )) : P py π = y is the set of joint probability distributions with marginals μ on X' and on y (Leetal 2019 Yamada et al. (022).
Furthermore, TWD takes on closed analytical form and therefore can be computed in O(dim( w)) Takezawa et al. 2021):
WT v) ||diag( () ((x y)1l
TWD is good empirical approximation of the usual 1- I-Wasserstein distance (Yamada et al. 2022
The inverse OT problem and ground metrie learning: Leveraging the geometric nature of the Wasserstein/OT distance for unsupervised learning tasks relies having a good idea of the ground metric between features that is then "lifted" to the OT distance hetween samples. Heuristics for ground metrics can be useful, especially with embeddings (Kusner et al. 2015 Du al. 2019). but we are interested in broader, principled ways to find ground metrics.
Inverse optimal transport provides one solution: given 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 codder-onstrrined, solutions exist given sufficient constraints on C' (Paty Wolfram 2019). However, inverse OT requires access to the full transport plan. In practice and for high-dimensional features, such as gene-cell data, this is not feasible.
Cuturi 2020 Li et al., 2019; [Stuart &
An alternative is to use partial information about distance similarity to learn the ground metric, for example with supervised or semi-supervised learning (Cuturi & Avis 2014). A related but distinct concept is the idea of supervised learning of sample Wasserstein distances (Huang et al. (0166) and supervised TWD (Takezawa et al.) 2021); here, the (feature) ground metric is still assumed to be Euclidean, but information about distances is used to learn metrics between samples.
Unsupervised ground metric learning with Wasserstein singular vectors (WSVs): The fully un- supervised approach differs in that neither feature nor sample ground metric is assumed. Unsuper- vised techniques harness the relationship between the geometry of features (embedded in samples) and the geometry of samples (embedded in features) (Paty & Cuturi 2020 Huizing et al.! (2022).
Consider some data matrix X E IR Let ag E R be a sample (normalised row such that it sums 1) of X, and by RT be feature (normalised column). These just histograms "embedding" a sample as a probability distribution over features, and vice versa. For example, the a; (of which there are n) could be histogram of words in a document, and the by (m) documents containing given word, As another example, might represent cells (samples) as histograms over gene expression, where by are the genes (features) histograms over cells that contain each gene.
Huizing et al. (2022) learn ground metric matrices A E ERR"" BE R mm satisfying the fixed point equations:
A.j = XA WDlai, ag),
BR.1 = Wa bt),
(3)
AB
vk, E (1, (1,.,,m) m). i,j (2022), we use A to denote the ground metric on the a; and B the ground metric on the by- Power iterations are used to compute the ground metric matrices via repeated application of mapping &b that "lifts ground metrics to pairwise distances":
(1, [I....,n) and some (AA, AB) E TR2 Note that as opposed to Huizing et al.
L (A)a). = W (b by) + +T|||| R(bb - bi),
where R is a norm regulariser and > 0. symmetric definition holds for B. Then the equivalent problem is to find Wasserstein singular vectors such that there exist (Aa,, E (2 satisfying T (B) = AAA, & (A) = Ag B (equal to (3) when = -). Maximal singu-
lar vectors can be extracted with power iterations: A
BnlB,
B [
(A(A,)
3

---

Published as a conference paper at ICLR 2025
The complexity of single power iteration is O(n2m² (n log(n) m log(m))), since it requires computing m² Wasserstein distances where each distance is O log(n)), then n² Wasserstein distances where each distance is O(m² log (oo)m)).. Huizing et al. (022) propose to reduce this com- plexity via stochastic optimisation (which suffers from requiring far more iterations to converge) and entropic regularisation, "Sinkhorn singular vectors" (SSV), which improves the complexity to best (n""m²) There are no other known reductions in complexity of this unsupervised method.
Tong et al. 2022) propose an efficient Lememdddinng of the Wasserstein distance on graphs in the unbalanced OT regime; however, the underlying ground metric is still assumed geodesic.
While our approach improve complexity through embedding samples and features as leaves on trees was inspired from the TWD literature, several authors have explored the relationship between samples (rows) and features (columns) as distributions over the other in general (Ankenman 2014
Gavish & Coifman 2012). including graphs (Shahid et al. 2016 and in tree-embeddings (Anken- man] [014] Mishne et al. 2018 Yair et al. 2017). These methods do not learn the ground metric in the same way our proposal, but Mishne et al. (2018), Ankenman (2014) and Yair et al. (2017)
describe iterative metric-learning of the tree metric and free construction, which is related.
1.2 CONTRIBUTIONS
We propose to improve on the WSV computational bottleneck by embedding the dataset in trees and approximating Wc with the TWD. We name method Tree-WSV. Our findings position Tree-WSV as fast, low-rank approximation of the standard WSV approach. We show that:
One can learn a complete set of strictly positive tree edge weights by solving a singulr system of linear equations (Algorithm D. and hence solve for the entire Wasserstein dis- tance matrix on features and samples, respectively.
This system is determined by a pairwise leaf-to-leaf paths matrix whose rank is guaranteed be equal to the number of edges in any tree with root degree 2 3. This path matrix's basis set can be found via a recursive algorithm, allowing fast computation of edge weights.
The Tree-WSV approach is a quadratic order more computationally efficient than WSV and outperforms Sinkhorn entropy regularisation (SSV) in terms of accuracy on toy data.
Tree- WSV, combined with meta-iterations, scales single-cell RNA sequencing datasets of size over 5000 x 5000, on which it is also faster and at least as accurate as SSV.
2 UNSUPERVISED TREE- WASSERSTEIN GROUND METRIC LEARNING
2.1 TWD SINGULAR VECTORS APPROXIMATE STANDARD WASSERSTEIN SVs
Let A = {ai), B (ba), "I....,,,, k E (1....,m) m) be the set of normalised rows (samples)
and set of normalised columns (features) respectively of the data matrix X E <R We embed A
and in respective trees Ta and TB. such that (ai...., an ) are the leaves of TA and {b). .... bm
the leaves of TB. as illustrated in Fig. Let (4) {0, (N ( E {0, 10 MM- be the tree parameters of each tree, where N and M are the number of nodes in Ta and TB respectively.
Let A I wB E TR. M-1 be the vectors of edge weights. We can port the WSV equations (3) to this tree setting:
Proposition 2.1.
The WSV fixed point equations 3 can be expressed on the trees TA, TB as:
1 WTa (ai, az), dTg (br. bi) WTA (br. bt).
1
dr (r,las, ag)
XA
A
where the singular vector update is to find w (and symmetrically wa) such that Vi,
NAWA T (2((4) + +00++2+0 (A) + 22 (A) o 2 x(a)) = |/iag (w (z ag)l
(4)
where: ((4)) the ith column of (A),, denotes element-wia product and (AA, E (R;)².
Proof: Appendix A
4

---

Published as a conference paper at ICLR 2025
r A
root B
Q22
b1
b2
b3
b4 by
Figure 1: Tree embeddings for samples A as leaves in Ta (left) and features B as leaves in TB (right). The tree metric d (a2; a3) shown as the shortest path between these leaves in orange on the left. Equivalently, we can use the relative expression of the features {b1. b2. b3. ba, bş } that represent 02 ag respectively (as shown in different hues of green on the right) to compute TWD, WTa (a2, 03), in TB. We assume d (22. az) ) is equal to WTa (a2, a3) to learn good embedding.
For intuition about Proposition 2.1 choose two samples ag, aj that are leaves in one of the trees; Ta (Fig. D. The tree metric d the shortest path between these leaves with edges in TA a4 aj
by- can also be expressed as histograms over features which are leaves on TB. i.e. 0₃ = > k==1 cb
(a.)
where Cb. 2 0 is the relative expression of each by in the histogram, so that Σ bb
[a,] = 1. The
1- I-Wasserstein distance WB Va,a az) between the two samples uses some unknown cost matrix C' so that Wn(a; aj) minp Σx E, PR.CCR. where P1, a; and P_ 1, aj. From Yamada al. 2022). Wa (ai ag) is approximated by the TWD between ai aj on TB. which we can express through summing proportions of dT between the leaves in TB: dr, (a2, az) WTB (a2, 03).
We can utilise Proposition 2.1 to learn low-rank approximations of WT (a;, aj) and WT (by, b1). Let Vi,j z4 (A) + zj (4)
(A) o zj (4) XOR{={4),=)) x= {0, 1 and concatenate the Vi,j into tensor (44) E f0, (JN 1)xnxx (i.e. all distinct nodes in the paths of all sets of pairwise leaves). Let Y' E {0, 11 be the matrix created when the tensor Y is unravelled along x n leaf-pairs paitspais...ooo Equation 4 for all i.j pairs then rewrites as a system of n² linear equations:
AwwA TY'
WTa where WTB = {
(diag(ws) ww Z (B) (a; ag) || vi. E " (5)
symmetric system of equations holds with the roles of w A. wB reversed.
2.2 NON-ZERO AND UNIQUE SOLUTIONS TO THE TREE WSV PROBLEM EXIST
Since Equation 5 is still singular vector problem, it can be solved with power iterations. As in Huizing et al. (022), we can show that solution to the singular vector problem exists:
Lemma 2.2. The singular value problem has a solution (w., B)
Proof: Appendix B
While Lemma 2.2 states that a solution exists, convergence to that solution is not guaranteed. How- ever, in practice we always observed convergence. with dynamics that were consistent with the linear convergence rate reported in Huizing et al. (2022). We show empirical convergence for power iterations on different dataset sizes in Appendix Figure
Each iteration of the singular value problem involves solving non-nege non-negative least squares (NNLS) problem. Theorem 2.4 contingent on Lemma 2.3] we show that as long we restrict ourselves to trees whose root node has degree greater than 2 (which we can do, since ClusterTree's 's hyperpa- rameters allow setting the number of children), NNLS finds unique, non-zero solutions to (5).
Lemma 2.3. Let the number of nodes in a tree be N, including the root and n leaves. Then there can be at most one non-leaf node of degree In addition, such node exists, is the root and Y node
has rank N- 2. Otherwise, every degree of the tree is at least 3, and Y has rank N - 1.
Proof: Appendix C
The proof of Lemma 2.3 suggests an efficient recursive algorithm that can also be used to find the basis set U for Y', and detailed in Appendix [
5

---

Published as a conference paper at ICLR 2025
Algorithm Unsupervised ground metric learning with trees (ree-WWVV)
Input: dataset X, size n x m
do once
A. B < (a: X (6)
row Σ x row (bx = x col Σ x (k)
Z A) wwwsteerrrreeeeee B), Clusterrrree
if n and m small (less than 500) then
((A), Y(B) tensor {Vi.j = (2 (A) _ (4) tensor {yk,1 (0R : (
U44), U(B) t basis sparse SVD upper triangular Y []]. upper triangular Y (B)
Zun An,e), (Baz) Z (bx bx) v(ba, bi} Bu, ( a3) V(ai, aj} Au
else
compute Us and Zais with recursive basis set algorithm (Appendix D end if
initialise
A, <- random vectors of length N 1, M 1, i.e. the number of edges in each tree
repeat
for t = 1 to num iterations do
(A)), WT (B) + {Il diag ( Zain (A || Vk, (iagg/////a/ee (waggww))))e Zan BBss) 11 vi,j
Wr norm WT noem f
W - NNLS WA (A) (T norm
w NNLS w (B) WT (A) norm
end for until convergence: max WA w A (/aiioou wB (previous)ll//m) < 10-6
an) the rows of data matrix X E REXm such that the root node of TA has degree 3 or more, and a tree TB with leaves given by B (bi,... bm} the columns X, there exists unique non-negative solution for wA (5).
Theorem 2.4. Given any tree TA with leaves A = {a1.....
Moreover, we can assume without loss of generality that the solution is strictly positive.
Proof:
As the matrix Y' has rank N the solution w W to AAW TY WTn equivalent to the solution of A Y (WTg B ).. where U is basis set of linearly independent columns of Y'. Y is an invertible (N - 1) x [N- - 1) matrix. Therefore, the NNLS problem
arg min IAAAA T Yú -WWr)ull S.t. WA A
is a minimisation of a strictly convex function over convex set, and has unique solution (existence also follows directly from the Rouché Rouché-Capellt Theorem). Further, as in Yamada et al. (2022). if w contains any zeros, these be removed, which corresponds merging nodes the tree TA-
L
2.3 COMPUTATIONAL COMPLEXITY AND SPEED-UPS
Note that the sparse LASSO- algorithm for learning w ("cTWD"") proposed in Yamada et al. (2022) could also solve Equation 5 However, cTWD cannot be applied directly in our setting due to its computational complexity: solving each TWD-LASSO problem is O + m m which worse than the cubic complexity for NNLS (Efron al., 2004), and cTWD requires computing all n² or m2 TWDs at each iteration. Further, memory consumption from storing the full Y' matrix quickly prohibitive.
Our Tree-WSS method suggests efficient way of computing tree- rree-Wassesttenn singular vectors in Algorithm [ Each inner power iteration computes tree- Wasserstein distances in time complexity do not
(n)) rather than (n² log(n)) ) for the 1-Wasserstein distance. Further, in each iteration need to compute all m² distances, but are capped to M < 2m the length of the weight vector (assuming redundant nodes). There 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 & set as l > 2 (Le al. (019)). Note that here, ClusterTree is modification of QuadTree (Samet (1984) and extended to higher dimensions via a grid construction in Indyk & Thaper 2003)). Our method follows the implementation in Le al. 2019) and

---

Published as a conference paper at ICLR 2025
1985
ClusterTree implemented in this way has complexity ((K) where k is the number of
eaff clftcrusters (in our case set automatically by the hyperparameter controlling the depth of the tree)
Gonzalez 1985). We compute (AA), ZBB) once for each of A, B, and Z ((4) (bą. bi) Ak, I and
ZiB) aa; az) j just once.
Similarly, we only learn the Y tensors once with the tree structure. In fact, we do not need to learn the whole tensor: we just need a full-rank sub-matrix of Y' of rank - 1, from Theorem 2.4 which call U the 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 U. Since Y' is sparse, for small dataset sizes, one can employ s sci ccipy sparse methods to find the basis set. Sparse QR decomposition sparse singular value decomposition (SVD) can be used; in practice sparse SVD performed slightly faster.
Because this operation requires computing large tensor Y reshaped into long rectangular matrix of size (N - 1) x n², the sparse method does not scale well with large n, m from a memory- consumption point of view. In this case, instead created recursive algorithm as suggested from the constructive proof of Lemma 2.3 This method directly computes basis set from just the (N-1) X (nodes by leaves) tree parameter matrix Z. Recursions are done on at most the number of nodes less the number of leaves (N < n 1 recursive loops); therefore, it is very fast and resists stack overflow. Stochasticity can be added by allowing randomness path assignments, but since the result is still basis set, this has minimal effect on finding the solution to the linear system of equations. The full recursive function detailed in Appendix D In practice, both approaches produced well-conditioned U basis set matrices (Section 3.2)
2.3.2 ITERATIVELY UPDATING THE WEIGHT VECTORS
At each iteration, we learn the vector wa satisfying equation coefficient matrix Y
5 using the basis set (A) of the
AAWA T (A)) WTB where WTg {|| diag (w (m)))))))) (a; ( az) Vai, aj E Av)
Here, (4) E (N (IN--))(N-I) is the basis vector set for Y' which corresponds with pairs of vectors {a,, ay} E Au. and z is the sub-matrix of the Z matrix these pairs {a az} E Au. We then symmetrically learn wB, and iterate.
From Theorem 2.4 solution exists, and it can be found with linear systems solver such non- negative least squares (NNLS). NNLS is efficient, with cubic complexity on the size of w A (i.e. O(N3) and ((3)) for WA, wB respectively). Each compute of pairwise TWD is linear in the length of the vectors (m or n for ag, by respectively), and there are N 1 or M - of these compute. This gives overall complexity per power iteration: O (N' 3 + m.N + M nM) < O (n' + m3 + mm). using N < 2n 1, M < 2m 1. This at least quadratic order faster than On?2m²(n log (n) + m log(m))) the complexity for a WSV power iteration, and if ss m, an order faster than O (2n?m²) the complexity for SSV power iteration.
Overall, we achieve a significant theoretical gain in time complexity. Computationally efficient unsupervised ground metric learning with TWD (Tree-WSV) is summarised in Algorithm !
EXPERIMENTAL RESULTS ON TOY DATASET
We demonstrate the improved speed and accuracy of our method as compared to SSV on synthetic dataset X E R= n = 80, m 60, modified from prior work (Huizing et al. 2022). On various synthetic dataset sizes and ratios, we observed empiric convergence of power iterations (Appendix and confirmed that both the SVD and recursive basis set approaches are well-conditioned.
3.1 DATASET CONSTRUCTION AND EXPERIMENTAL METHODS
The dataset uses various translations of unimodal periodic function on a 1D torus (with boundary conditions), such that Xst exp (- #x
4)² /)²/02) On X, the underlying (and learned) ground

---

Published as a conference paper at ICLR 2025
A
Time complexity comparison
B
Hilbert distance after 100 iteraaionns
Error compared to full WSV
100
Time for 100 iterations (s)
80
40
20
Wasserstein SV
Entropic WSV (SSV)
Tree-WSW: 10 children
Tree WSV: 3 children
Tree-WSV: children
Entropic WSV (SSV)
TreeWSV: 10 children
Tree-WSV: 3 children
Tree WSV: 2 children
Algorithm
Algorithm
Figure 2: Comparison of (A) mean time complexity and (B) mean Hilbert distance between learned and WSV ground metric matrices recovered by different algorithms (lower is better). Tree-WSV algorithms are specified by their hyperparameter the minimum number of children per node in the tree; all trees have maximum depth 5. Standard error of the mean is shown with black bars.
metries look like |sin( E | (and symmetrically for m). We ran experiments times for 100 iterations each on CPU with different methods, using a translation of the periodic function, to which we added 50% translation of the torus with half the magnitude. We set = 0.01. Thus
m 0.5)² /o²). In order 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 Hilbert metric between the learned cost and that recovered by the standard WSV algorithm (Fig. 2
X ik exp (- =) /)² + 0.5 +
3.2 RESULTS
All tree-based methods were faster than both standard WSV and entropic-regularised WSV by 2-3 orders of magnitude (3-children trees took ~1% of the time of the standard WSV, for example), as predicted by our theoretical time complexity calculations (Fig. AA) Interestingly, the tree-based methods had lower Hilbert distance were noted to converge a similar amount of time (less than 10 iterations) to the standard WSV.
when compared to entropic WSV (Fig. 2 B). Tree methods
As expected from Lemma 2.3 the trees initialised with 10 children had lower U-ranks than both 2- and 3-trees, and 3-trees had lower ranks than 2-trees. In general, trees with fewer children per node (higher rank) computed distance matrices slower, but produced more accurate results (Fig. 2 One exception follows directly from Lemma 2.3) 2-trees have rank N - while 3-trees are rank N - 1, so the learned edge weight vector w is not guaranteed be uniquc/optimal. As result, the 2-trees performed worse in terms of accuracy than the 3-trees (Fig. B).
We found that the basis set matrices U proaches. For the SVD approach, the condition numbers were within 1 + 10 15 dataset sizes 80x60. 100x200 and 500x500. For the recursive algorithm, the condition numbers were higher, 1000-3000 for random sample sizes between 1000 and 2000. Based on numerical precision, these higher condition numbers should not affect the accuracy of to within 2 orders of magnitude.
well-conditioned for both the SVD and recursive ap-
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 (Tagg al. 2009). scRNA-seq data 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 (Kzzak et al. 2019). Traditional
8

---

Published as a conference paper at ICLR 2025
A
B
CDST
Drysuee
15
Hane
Strama
DC
60-
-25
-50
C
D
UPAN1
GAkk-egic bouce
Dolamategg
Peob
Dipadendrayl Dipodeediin diccts
09577
Figure 3: UMAP embeddings based on the learned tree-singular vectors for cells, and coloured by provided annotation: (A) PBMCs (Wolf et al. 2018); (B) Lung tissue (Sikkema al. (2023); (C) Mouse VI neural tissue, by broad types (Tasic et al. 2018); and (D) GABAergic neurons in the same mouse VI neural tissue set C. by subtype (Tasic al. 2018).
methods assume a Euclidean metric on PCA embeddings, or use neural network model embeddings (Stuart et al.) 2019 Wolf al. 2018 Gayoso et al., (2022).
More recently, the gene mover distance has been used (Bellazzi et al. 2021 Du et al. 2019). Huizing jet al. (2022) showed that Sinkhorn singular vectors can achieve state-of-the oh-the-rrt clustering scores on small scRNA-seq dataset. We illustrate the potential utility of Tree-WSV by using it on the same dataset and achieving similar performance, as well as employing on other datasets with improved performance. Finally, showcase Tree-WSV on 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. 2018), which 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 (V1) cells in mice by Tasic et al. (2018), consisting of 1468 cells with 7 broad types and 1000 genes, and subset of the "Human Lung Cell Atlas (HLCA)" scRNA-seq consortium dataset of human respiratory tissue (Sikkema et al.) 2023),, consisting of 7000 cells with 4 broad types and 3923 genes. Details preprocessing, annotation and experimental set-up are identical to Huizing et al. (2022) for PBMC and similar in the other cases, and provided in Appendix F Experiments
NVIDIA
A 100 Tensor Core GPU with JAX (Bradbury et al. 2018).
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 with being better.
4.2 RESULTS USING TREE-WSV META-ITERATIONS
We noted that silhouette scores for iteration of the Tree- WSV algorithm while quick 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 Fuclidean distances the input to cluster nodes based on the

---

Published as a conference paper at ICLR 2025
Table 1: Comparison of ASW for metries computed and runtime for scRNA-seq datasets for WSV as compared to tree-based WSV. Euclidean metric for baseline. 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
> 12 hours
Tree-Wo WSV (ours- 4 iters)
0.299
8 min
0.436
min
0.104
27 min
(ours best)
0.313
72 min
0.602
32 min
0.457
71 min
incremental furthest search algorithm, resulting in low-depth initial trees. To account for this, we re- ran the Tree- TWD unsupervised 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 iterations and best (max 15 iterations) are summarised in Table []
The ASW metrics comparable in performance to Huizing al. 2002) Sinkhorn singular vec- tors approach and better than other alternatives for the PBMC dataset, as presented in Table and baseline. However, the Tree-WSV method vastly out-
Appendix G with Euclidean distance performs SSV on a similar task for a dataset of labelled neural cells (Tasic et al.) (018)) and in the HLCA data (Sikkema al. 2023). Further, our method has considerably lower runtime. Indeed, on the large HLCA dataset (7000 cells) SSV had not converged after 12 hours. Tree-WSY 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 The clusters for the PBMC dataset are qualitatively similar those of Huizing et al. (022 (Fig. 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. (2018) neuron dataset (Fig. 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 ict al. 2018). These clusters reveal themselves coarsely when we label by them in Figure Since Free-WSVine is agnostic to the number of clusters in the data, it is not surprising that more detail can be revealed. This advantagcous deeper level of complexity should be investigated further.
Finally, while we did not find that ASW scores had converged after 15 meta-iterations, vector solutions in the inner power-iteration loop tended to always converge within 10 iterations. Stochas- ticity inherent from choosing the basis set U, well as the distance used to decide 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.
5 CONCLUSIONS AND PERSPECTIVES
We present Tree-WSV, a new, computationally efficient approach to learning ground metrics in an unsupervised manner by harnessing the tree-Wase wrs-essserttein 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 new geometric underpinning of the restrictions needed for tree structures to flexibly represent data in low-rank approximations. Our method suggests a fast and recursive algorithm to construct basis sets for pairwise paths in trees which could have wide-spread applicability. While promising, there is need for further large- scale testing and development of the Tree- Tree-WSV algorithm for other use cases, including neuronal activity data. Lastly, understanding theoretical underpinnings of convergence, the role of metrics in designing the initial cluster tree, and stochasticity in the basis set construction could be fruitful avenues for future work.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
K.M.D. and S.H. are supported by the Gatsby Charitable Foundation, K.M.D. was hosted by the Okinawa Institute of Science and Technology (OIST) as a Visiting Research Scholar during this work, and the authors thank OIST for their hospitality. M.Y. is supported by MEXT KAKENHI, grant number 21 21H04874, and partly supported by MEXT KAKENHI, grant number 24K03004.
The authors thank Geert-Jan Huizing and Aaditya Singh for insightful discussions about the work and early comments. Geert-Jan Huizing also made available the PBMC3k pre-pro pre-processed dataset from Huizing et al. (2022), which facilitated direct reproduction of and comparison to existing sults.
REFERENCES
Shibla Abdulla, Brian Aevermann, Pedro Assis, Seve Badajoz, Sidney M. Bell, Emanuele Bezzi, Batuhan Cakir, Jim Chaffer, Signe Chambers, J. Michael Cherry, et al. CZ CELLxGENE Dis- cover: A single-cell data platform for scalable exploration, analysis and modeling of aggregated data. bioRxiv, 2023. doi:
Jerrod Isaac Ankenman. Geometry and analysis of dual networks on questionnaires. PhD thesis, Yale University, 2014. URL https::/wwit http::/:wwwwpppp...tt.... com http:///eeeeeee//111111
Dissertaat.ons$$00457777722
Riccardo Bellazzi, Andrea Codegoni, Stefano Gualandi, Giovanna Nicora, and Eleonora Vercesi.
The gene mover's distance: Single-cell similarity via optimal transport. arXiv, 2021. URL
http://arxiv. 22000 211./22222211111
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milen and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
//github. htt.hhh.ccoomm/cm jax 1alaammmmjaa jax
Emily Clough, Tanya Barrett, Stephen E. Wilhite, Pierre Ledoux, Carlos Evangelista, Irene F. Kim, Maxim Tomashevsky, Kimberly A. Marshall, Katherine H. Phillippy, Patti M. Sherman, et al. NCBI GEO: archive for gene expression and epigenomics data sets: 23-year update. Nucleic Acids Research, 52:D138-D144 2024. ISSN 13624962. doi: 10. 10.1093/na//kkad965.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In C.J. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K.Q. Weinberger (eds.), Ad- vances Neural Information Processing Systems, volume 26. Curran Associates, Inc., 2013. URL httpps://proceedddnss. neurips www.wneerii.... fil Fil4a422.0cssshhee////2222555333 wwtppppr.epppppppppppppe fi ftl.ess/ppppe//0033/
Marco Cuturi and David Avis. Ground metric learning. Journal of Machine Learning Research, 15 (17):533-564, 2014. [tl.:http///mmmrrorrr//pprrrrrrp
14a. 4..h.mm......mmmmmmmm...mm.mmm..
Jingeheng Du, Peilin Jia, Yulin Dai, Cui Tao, Zhongming Zhao and Degui Zhi. Gene2vec: Dis- tributed representation of genes based on co-expression. BMC Genomics, 20:7-15, 2 2019. ISSN 14712164. doi: TAt.://SSSS8000000-----000----7-
URL https: http::/s///mmmggoooeee /ttp////mmincc....////////./ 018 5370-x
bi piomedcentral. com www.w//wtttt.........tt.t.
Bradley Efron, Trevor Hastie, lain Johnstone, and Robert Tibshirani. Least angle regression. The
Annals of Statistics, 32(2):407 499, 2004. doi: 10.114400005366400000000000 URL https:
7/doi org./oorg//a /044:0000000
Matan Gavish and Ronald R. Coifman. Sampling. denoising and compression of matrices by coherent matrix organization. Applied and Computational Harmonic Analysis, 33(3):354-
369, 2012. ISSN 1063-5203. doi: http:://diiaggllo..ooo/s.....oop URL https: 7/www. ici enceii rect. c htt.t//ccccccneeeeaaeeeaaaaaaaa.
88

---

Published as a conference paper at ICLR 2025
Adam Gayoso, Romain Lopez, Galen Xing, Pierre Boyeau, Valeh Valiollah Pour Amiri, Justin Hong, Katherine Wu, Michael Jayasuriya, Edouard Mehlman, Maxime Langevin, et al. A Python library for probabilistic analysis of single-cell omics data. Nature Biotechnology, 40:163- 2022. ISSN 396.669. doi: 10. 1008//4115500 1a000441157--222126-----
Teofilo Gonzalez. Clustering to minimize the maximum intercluster distance. Theoretical Com- puter Science, 38:293-306, 1985. ISSN 0304-3975, doi: 10.101603304-995555)0022222
Yuhan Hao, Stephanie Hao, Erica Andersen-Nissee William M. Mauck, Shiwei Zheng. Andrew Butler, Maddie Lee, Aaron J. Wilk, Charlotte Darby, Michael Zager, et al. Integrated analysis of multimodal single-cell data. Cell, 184:3573, 2021. ISSN 10974172. doi:
00.0160..CELL..221....
URL
Bco:233s9//ppprrraasaaatttttottp axticles/PMC823640
nlm gov/pm////
Gao Huang, Chuan Guo, Matt J Kusner, Yu Sun, Fei Sha, and Kilian Q Weinberger. Su- pervised word mover's distance. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 29, Curran Associates, Inc., 2016. URL htttps::/proceediingg www.::ss..seuiiiisoommm ccppaper files/ Pra.aaaaee..rrteeeeeeee.eeeeeeer
Geert-Jan Huizing, Laura Cantini, and Gabriel Peyré. Unsupervised ground metric learning us- ing Wasserstein Singular Vectors. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Confer- ence Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 9429 9443. PMLR, 17-23 Jul 2022. huizing22a, html
URL htt.:/ss//ppcceccnngn...... www.:/ss.w....e..sseeese.eeeseee
Piotr Indyk and Nitin Thaper. Fast image retrieval via embeddings. In ICCV International Workshop on Statistical and Computational Theories of Vision, volume 2, 2003, URL/https: /www. cs.
princeton edu/ www.seee/accciee/eeeee/eee /ow.ueee//rrcceeeee///gg///////a
emd. pdf
L.V. Kantorovich. A new method of solving some classes of extremal problems. Doklady Akademii Nauk SSSR, 28:21 1-244, 1940.
S Kolouri, Yang Zou, and Gustavo K Rohde. Sliced Wasserstein kernels for probability distributions. In 2016 IEEE Conference on Computer Vision and Patlern Recognition, 2016. doi: 10.1 109/ CVPR.2016.568, URL https: http:://axxvv.o.oee. abs 1111. 03198
/
Monika Krzak, Yordan Raykov, Alexis Boukouvalas, Luisa Cutillo, and Claudia Angelini. Bench- mark and parameter sensitivity analysis of single-cell RNA sequencing clustering methods. Fron- fiers in Genetics, 10:486077, 12 2019. ISSN 16648021. doi: 10.389IFGEEN:22199.2555
BIBTEX. URL www. www.tfrrriiisioorrrrm
Matt Kusner, Yu Sun, Nicholas Kolkin, and Kilian Weinberger. From word embeddings to docu- ment distances. In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 957-9667666 Lille, France, 07-09 Jul 2015. PMLR. URL https: httpss://proceeeinggg press bosass/37/ee http:///////mmmmbb5.tttt.
Tam Le, Makoto Yamada, Kenji Fukumizu, and Marco Cuturi. Tree-sliced variants of Wasser- stein distances. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Gamett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https: http::////cceeecceeeeeee neur wwtppprrrcppppppppcppppppp.....e files/ Pap.rr..dappp.p....
Ed S. Lein, Michael J. Hawrylycz, Nancy Ao. Mikael Ayres, Amy Bensinger, Amy Bernard, An- drew F. Boe, Mark S. Boguski, Kevin S. Brockway, Emi J. Byrnes, et al. Genome-wide atlas of gene expression in the adult mouse brain. Nature, 445:168--76 2007. doi: 10.1038/aauue004553. URL htttp:://do.. orrg 1ttp:/00...000//tttro0o0
12

---

Published as a conference paper at ICLR 2025
Ruilin Li, Haomin Zhou, Hongyuan Zha, Xiaojing Ye, and Zha Li. Learning match via inverse optimal transport. Journal of Machine Learning Research, 20:1-37, 2019. doi: 10.48550/arXiv. 1802.03644.
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, 4(3):451-466, 2018. doi: 10.1109/TSIPN 201762743561.
Monge. Mémoire la théorie des déblais et des remblais. Histoire de l'Académie Royale des Sciences de Paris, avec les Mémoires de Mathématique et de Physique pour la même année, pp- 666-704, 1781.
Prançois-Pierre Paty and Marco Cuturi. Regularized optimal transport is ground cost adversarial. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 7532-7542. mlr wtw....ppre.ss/1///aaaaaaaa.m
PMLR, 13-18 Jul 2020. URL https:
html
Gabriel Peyré and Marco Cuturi. Computational optimal transport. Foundations and Trends in Ma- chine Learning, 1(5-6))355-607, 2019. URL https::///// /ttp///axvvvvvee///e wwwgg/ba//... wwtpaaaaa...../ 11.//1803.000 00567
Hanan Samet. The quadtree and related hierarchical data structures. ACM Comput. Surv, 16(2): 187-260, June 1984. ISSN 0360-0300. org/10. 100...1155335688 3566599322366300
doi: 10.114553569244355933 URL https::/doi..
Nauman Shahid, Nathanael Perraudin, Vassilis Kalofolias, Gilles Puy, and Pieme Vandergheynst. Fast robust PCA on graphs. IEEE Journal of Selected Topics in Signal Processing, 10(4):740- 756, 2016. doi: 10.1 109/1STSP2201225555299
Lisa Sikkema, Ciro Ramírez-Suústegui. Daniel C. Strobl, Tessa E. Gillett, Luke Zappia, Elo Madis- soon, Nikolay S. Markov, Laure Emmanuelle Zaragosi, Yuge Ji, Meshal Ansari, al. An inte- grated cell atlas of the lung in health and disease. Nature Medicine, 29:1563-1577, 2023. ISSN 1546170X. doi: 10.1038/s41591 19141199---232277----1
Andrew M. Stuart and Marie Therese Wolfram. Inverse optimal transport. SIAM Journal on Applied Mathematics, 80:399-619, 5 2019. ISSN 00361399. doi: 10. 1 137/19M1261 122. URL https:
Tim Stuart, Andrew Butler, Paul Hoffman, Christoph Hafemeister, Efthymia Papalexi, William M. Mauck, Yuhan Hao, Marlon Stoeckius, Peter Smibert, and Rahul Satija. Comprehensive in- tegration of single-cell data. Cell, 197: 2019. ISSN 10974172. doi:
(0.10166j.eel..211900.00
Yuki Takezawa, Ryoma Sato, and Makoto Yamada. Supervised tree- Sre-Wasserstein distance. In Ma- rina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning. volume 139 of Proceedings of Machine Learning Research, pp. 10086-10095. PMLR, 18-24 Jul 2021. html
URL https: htEps://ppooeeddnngg
Fuchou Tang, Catalin Barbacioru, Yangzhou Wang. Ellen Nordman, Clarence Lee, Nanlan Xu, Xi- aohui Wang, John Bodeau, Brian B Tuch, Asim Siddiqui, et al. mRNA-seq whole-transcriptome analysis of single cell. Nature Methods, 6:377-382, 5 2009. ISSN 1548-7091. doi: 10.1038//meth.133 1315.
Bosiljka Tasic Zizhen Yao, Lucas T. Graybuck, Kimberly A. Smith, Thuc Nghi Nguyen, Darren Bertagnolli, Jeff Goldy, Emma Garren, Michael Economo, Sarada Viswanathan, et Shared and distinct transcriptomic cell types across neocortical areas. Nature, 563:72-78, 11 2018. ISSN 14764687. doi: 10.10038/441886-088--
Alexander Tong. Guillaume Huguet, Dennis Shung, Amine Natik, Manik Kuchroo, Guillaume La- joie. Guy Wolf, and Smita Krishnaswamy. Embedding signals on graphs with unbalanced dif- fusion earth mover's distance. ICASSP, IEEE International Conference on Acoustics, Speech and Signal Processing Proceedings, volume 2022-May, pp. 5647-5651. Institute of Electrical
13

---

Published as a conference paper at ICLR 2025
and Electronics Engineers Inc., 2022. ISBN 9781665405409. doi: 10.1009/CAASSPP3922.222 9746556.
F. Alexander Wolf, Philipp Angerer, and Fabian J. Theis. SCANPY: Large-scale single-cell gene expression data analysis. Genome Biology, 19, 2 2018. ISSN 1474760X. doi: 10.1186/ si 13309-017-1382-0.
Or Yair, Ronen Talmon, Ronald R. Coifman, and loannis G. Kevrekidis. Reconstruction of nor- mal forms by learning informed observation geometries from data. Proceedings of the Na- lional Academy of Sciences, (14(38) 2017. doi: 10. 10.1003/ppas. 1etis411200411 14. URL https: /
/tttps::/www... www.wwwww.................w..... /oi/ 3 httppppa............ 1620045114
Makoto Yamada, Yuki Takezawa, Ryoma Sato, Han Bao, Zornitsa Kozareva, and Sujith Ravi. Ap- proximating I-Wasserstein distance with trees. Transactions Machine Learning Research, 2022. ISSN 2835-8856.
Quan Zou, Pengwei Xing, Leyi Wei, and Bin Liu. Gene2vec: gene subsequence embedding for prediction of mammalian N "--methyllddeooine sites from mRNA. RNA, 25:205-218. 2019. doi: 10. /mmaaa URL BtLhttp::// www. rna journal ott/c////oo//10 1261/rna.
14

---

Published as a conference paper at ICLR 2025
A PROOF OF PROPOSITION 2.1
From Yamada et al. (202) and Le al. (2019), given a tree and two probability measures sup- ported on 7, the TWD computed using 7 is the I-Wasserstein distance the tree metric dr. For any two measures, there also exists a tree such that can approximate the I-Wasserstein distance between the measures; however, since this is only proved for two measures, combining many mea- sures (samples) as we do in this work is not guaranteed to represent the Wasserstein distance (or indeed any distance) without distortion. In this proof, seek to show that we construct TWD-version of the Wasserstein singular vector problem, but we do not comment on how well this approximates the 1-Wasserstein distance theoretically.
(A)
2022), define a tree parameter-base based vector for every a, aj as Vitj Following Yamada et al.
((A) E Z(4,, that E {0, 1 (equivalently we will write this as (=R (4) (A) }), which is known from the tree structure. The tree metric (distance
2z;
(A)
with (A) ز
3
between any two leaves on Ta) is then given by:
(4) dr (a: aj) = WA Vi.j
(6)
Note that any vector by is just a distribution over leaves (a,) in Ta and vice versa. Using the distribution over the a, with the ground metric on the matrix/tree TA, distances between any by. by Since TWD admits the closed form in equation Yamada et al. WTA (b,, by) = []dia( Z (4) (()) b1)ll1-
derive tree Wasserstein
2022
Now, let assume that instead of learning the ground metric for the full distance matrix, want to learn a tree metric dTA (a,: aj) (6) The equivalent claim as in Huizing et al. (2022) is that we achieve the fixed points in Proposition 2.1 More formally, as before, we would like:
AgwB Y = (w with Aa, λB.T Rt, where & ( WT (a,, ay) + 7 A| tensor composed of the n² (or m²) Vi,j vectors. In this case, the ts map ground costs to Wasserstein distance matrices.
AAWA (Y()) = BB(B))
ay), and Y is the
Let A (AA (bn. b1). Since W E RNN-1 is positive, and letting absolute value,
denote element-wise
W TA (bk, b1)
||diag( W a (.) 1/2
(A)
M
(A) = Σ (wa) ( k,l
(((()
w A
So from 2 and Proposition
2.1
AAWA Vi.j
(4)
wB
[2(B]
A
yk WA T (12(A)
(B)
Vk, I (1, m), nI.
In full, the singular vector update becomes to find wA such that Vi,j AAWA (4) + 2 2²(2) 2(1))
diag w (B) (a: (an-aa))])
ag)l
(8)
(and symmetrically for the wB iteration).
D
B PROOF OF LEMMA 2.2
The proof proceeds identically to the proof of Theorem 2.3 in Huizing al. (022). As in Algorithm the singular value updates satisfy
WTB
w (A)
W
WTA
(B)
wB
| 1l
15

---

Published as a conference paper at ICLR 2025
The left-hand sides of these equations are the ground metrics, corresponding to A and B in Huizing et al. (2022). The matrices U(A), U(B) are invertible by construction, and so the existence of fixed point ( UUA),, wB (r))) to the left-hand sides equations above implies the existence of a corresponding fixed point for the weight vectors (WA, WB)
As the right-hand sides are -o normalised and non-negative, the ground metrics (i.e. left-hand sides) must have all components in the closed interval [o, 1]. Hence, reshaping the left-hand sides from n² (resp. m² m²)-dimensional vectors to (resp. x mm) distance matrices, the mapping that updates the ground metrics is a continuous mapping that maps the compact convex set
{(C, D) E ((0, 1jnx* x [0, 1]mxx) : C.2+ = 0 Vi, DR.E Vk}
to itself. By Brouwer's fixed point theorem, this mapping must have a fixed point.
C PROOF OF LEMMA 2.3
As preliminaries, assert that N > Neat m, and Z E 10,1)(N-1)xE has rank m. Note that Y E (0.1)((--)xxxx is strictly tensor, but in practice we only care for the reshaped matrix Y' {0, (NN-)xx.))
Assume that trees cannot have redundant nodes (i.e. nodes with exactly 1 child, not including themselves), and let every tree have root node and leaf nodes I..., e, which we also use to index the associated positions along vectors and y.
Note that 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 2 3.
Some bounds for N in terms of can be derived from the observation that rank is upper-bounded by the number of pairs that two distinct leaves can make, noting Vi.j = Vj,i and yi 0. For E2 4,
cll ((e-1) 1) 28 20 124. L.
Replacing N for N - we get bounds for 2 2. 28 2 - in 9 is derived from the upper bound on the number of nodes in the maximal spanning tree, binary tree.
We prove by induction on the number of tree nodes that the rank of Y' is N for any tree structure in which all nodes have degree 2 3, and N otherwise, for any & > 2.
N base case, N = 3: The tree can only be a root connected to two leaves. There is only one y-vector, between the two leaves. So the rank is N - 2 = 1.
N base case, N = 4: r connects to each of the three leaves. In general for N =+1, 3, we can show the rank is N - = f. Without loss of generality, choose leaf i and consider yi.j; Vj A Since every Zj is 0 everywhere except for positions indexed by node i and root r, where is 1, Vi.j i and 0 otherwise. Therefore the of é 1 vectors S, (yr.g. i * j, Vj} independent. Additionally, we can choose any two distinct nodes j. k where j * k * i, and S = {yjj& U ..) must also be linearly independent set. This follows because Yj,k is 1 at nodes j. k and 0 otherwise; and no linear combination of Yi,k: Vs.j (with = 1 for both vectors, which should be cancelled) can reproduce Yj.k-
linearly
Can we add any more vectors from Y to S such that the set is still linearly independent? No. To show this, assume that there exists some yg & SSan(S). Trivially, y # h 7 and {g. h) + [i, k). so yg.h is 0 everywhere except at the g, h positions. But note that we can write Vg.h using vectors in
Ng = Vi.g + yi,h - (yi.s + Yi.k yj.k
(10)
at otherwisc
So our assumption was incorrect and Y E (S) Therefore rank = -
Induction hypothesis, N < m: Assume that the rank of any tree structure with parameter N < m, m E N, has rank N 1 = m if all inner nodes are degree 3 or more, and N = m 2 otherwise. We call this m-tree.
16

---

Published as a conference paper at ICLR 2025
Inductive step, N = m + 4: We define a set of "adjacent" leaf nodes A to be the set of all leaves connected to the same parent node p where has no non-leaf children.
Note first that we 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 p (if depth is less than 2, we get the extended base case). These leaves have common ancestry, and their 2 a-parameters differ only at the indices of the leaves themselves. Since n 1, there must exist least one node p that is neither nor leaf. Collapse the p-sub-tree by removing and reconnecting its leaves to the immediate parent node of P. Then the new tree is an m-tree (we have lost one node), as claimed.
Every tree must have at least one such set A, and it must have at least elements. We consider the smallest A-sub-tree and the L/ A-u-b-rreee consisting of the rest of the tree, two cases.
Case 1, |A| > 3: From the base case, the rank of the A-sub-tree up to the parent node p is |A]. Now consider the sub-tree consisting of all the other leaves in the tree, L/ up until p. where the two sub-trees connect. Imagine p is leaf (ignore its children A). Then the L/ A-Auubb-rre has tree parameter NL/A = m + A, S m, so by the inductive hypothesis it has rank N - N --2 2 - A, depending on the degree of.. r.
or
How do we join the sub-trees? Let links including & now extend to some (the same) new leaf in {A), say ag The union of the Vi.j making up each sub-tree's 's basis set is linearly independent. To see this, note that the only Vi, that have at position o were those ending at p the leaf on {L/A}. So to express any ylu,a a link via is needed, that is the rank is lower-bounded by N - 1 or N -2 respectively.
We cannot add any more ys: since the sub- trees saturated by assumption, any addition y should also be a link, but any link can be expressed via an existing link through 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 - or N 2 based on the sub-tree, as required.
Case 2, |A| = 2: In this case, the L/ A-sub-tree has rank N - 3 or N - 4 and the A-sub-tree has internal rank since there is just one pair of leaves make yar. Without loss of generality, assume the links extend to a1. with az } for some (not necessarily single) i, necessarily linearly independent of aa The total rank then N
N 3.
In the 2-case, we can also create another yly.az, where j can either be the same as or different.
How do we know that 22 is not in the span of the two sub-tree's y vectors? Assume it is in the span. az has Is at p. The only other ys which shares this property are from Vl. So, if a1,,a could be written using already chosen ys. its decomposition must include yl, But since aj is known to be zero in yaz,ly but is non-zero in the path containing yl, [a,] = (where [ay] = refers to the component ya,,a, corresponding to aj) should "removed" through linear combinations with other ys; however, there is just one involving a1, which is yay,az subtract a2 resulting in a value of -1 at [a2]. But we require [a2] to be I: there is no linear combination that allows both [p] and [a2] to be 1.
So we must
So we must be able to add one link and so here, too, the rank is N or N 2 total (trivially any further ys must be links and would be in the span).
17

---

Published as a conference paper at ICLR 2025
D RECURSIVE ALGORITHM FOR TREE BASIS CALCULATION
Algorithm 2 Tree basis set recursion
basis tree
INPUTS: path_matrix, U basis initially parameter/parent matrix, rank initially 0
N, L f shape of (nniimmattixxndee leaves] (U has shape N x N)
if' N = 3 (base case: leaves and root) then
U [rank] < path of the 2 leaves (perform XOR on 2 path matrix columns), rank <- -
else
smallest-parent arg minlow index Brows UIL :), size < (smalllesst-ppaa smallest-parent s-leaves f- non-zero indices smallest-parent row if size = 2 (2 adjacent leaves) then
Urank] +- path amallest-parent s-lavvss (perform XOR on 2 path_matrix columns)
U 1rank+1] < path of I caaallsstppaeen seaf & a leaf & (smallest-paree s-leaves)
rank =
else
U Urrank] + path of smallest-parent's-eer s-leaves 0,size
rank+
for i 0 to size 2 do
Urank + i] f path of smallest taaeeestt-prettsss s-ravess-l end for rank+ =
end if
for in smallest- pattttt-peentss seavesss [ :] do
path matrix column at smallest-parent' s-leaves index i 0 (i.e. discount all leaves but one for the next recursion can be made stochastic) end for
U rank <- basis (rttsttteepaahhh matrix,, rankk end if
return rank
E EMPIRICAL CONVERGENCE
A
B
C
66
-0
-54
-35
-30
403
Figure 4: Empirical convergence (Hilbert metric) observed for different toy torus dataset sizes across power iterations. The smaller dimension is shown blue: (A) 60 x 80 (with SVD basis set calcula- tion); (B) 100 x 200 (with SVD basis set calculation); (C) 1500 x 300 (with Algorithm D basis set calculation).
F DETAILS ON GENOMICS DATA PREPROCESSING AND EXPERIMENTS
DATA PREPROCESSING
All single-cell RNA sequencing data was processed using standard protocols, including CPM- normalisation, loglp-tte ip-transformmtionn and selection of genes based on highest variability. Canonical markers were added according to source.
18

---

Published as a conference paper at ICLR 2025
PBMC3K data was recovered using Scanpy (Wolf al. 2018). Preprocessing procedure/code was provided on request from Huizing et al. (2022) and included cell type annotation using Azimuth marker genes (Hao et al., 2021), sampling only annotations for which confidence over 90%.
The results of SSV on this dataset were identical to that reported previously (Huizing et al. 2022).
Neural cell data and labels from Tasic al. (2018) were recovered from the NCBI GEO repository,
Series GSE71588 (Clough et al. 2024) and annotated using the Allen Brain Map (Lein et al.) 2007)
and CZ CELLXGENE (Abdulla al. 2023).
The Human Lung Cell Atlas (HLCA) was recovered, and annotated, from CZ CELLxGENE (Ab- dulla et al., 2023), and originally published by (Sikkema al. (2023). Because taking subset might disrupt the original annotation, low-gene-coum cells and cells with labelling entropy score less than 0.1 were removed.
F.2 ALGORITHM PARAMETERS
Genomics experiments using Tree- WSV were run for 20 iterations of the internal linear system of cquations singular vector loop (finding w) and 15 meta-itcrations with constructing trees based the previous weight matrices, through computing 2 new TWD-based distance matrices). A JAX random seed of 0 was used. The ASW and total runtime after the 4th meta-iteratio and the best score overall (usually the 12-15th meta-iteration) were reported.
the entire algorithm (starting
Genomics experiments using SSV were run for 15 iterations of the singular vectors loop with T 0.001, e 0.1; these follow the settings and replicate the results in Huizing al. (2022).
F.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 GB requested memory.
G COMPARISON TO OTHER METHODS FOR IMPROVING CELL-TYPE CLUSTERING
We reproduce the table from Huizing et al. (2022) showing other methods that can be used to cluster cell-type data from scRNA-seq experiments as compared to ours on PBMC dataset.
Table Comparison of average silhouette width (ASW) distance metrics computed and runtime in the PBMC dataset.
Dataset (size)
PBMC (2k)
Metric/method
ASW
Runtime
PCA / 1
0.238
Kernel PCA 8
0.241
scVI embedding / L
0.168
Sinkhorn
0.003
Gene Mover Distance
0.066
Euclidean
0.073
Entropic WSV (SSV)
0.348
110 min
Tree-WSV (ours iters)
0.299
8 min
WSV (ours best)
0.313
72 min
19


Reviews:
[{"review_id": "DjkdOR00mW", "reviewer": "Reviewer_1BGM", "summary": "The paper introduces Tree-WSV, integrating TWD with WSV (Huizing et al., 2022).", "questions": "- In line 194, what does “$a_i$ defined as for WSV” mean?\n- What is the size of $\\mathbf{U}$ in line 259?\n- 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?\n- How to decide whether the algorithm reaches convergence in Algorithm 1?\n- How fast does the proposed algorithm converge? How does the performance change across the iteration?\n- What are the hyperparameters of the proposed method? How sensitive are they in the experiments?\n- What is the Euclidean metric baseline in Table 1?\n- Why is WSV not considered as a baseline in Table 1?\n- Why are other TWD methods not considered baselines in Table 1?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "The proposed method demonstrates efficiency compared to WSV and SSV presented in 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.\n- 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.\n- The real-world application is restricted to single-cell RNA sequencing data, despite the introduction citing various data types as motivation.\n- The experiment section includes only two competing methods without considering the baselines used in WSV or other distance metric learning approaches.\n- $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.\n- 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.\n- 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.\n- 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.\n- 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.\n- 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?\n- It’s unclear what is $R$ in line 151.\n- 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.\n- It’s unclear what \"full WSV\" represents. Is it different from WSV?\n- 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.\n- 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.\n- 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}}$.\n- 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}}$\n- More details and explanations are needed for how Theorem 2.2 supports unique and non-zero solutions in Proposition 2.1.\n- 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)\n- 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.\"\n- 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\n\n## Minor\n- Missing “-” for tree-Wasserstein distances in line 077 and line 102\n- The acronym \"OT\" is used without being defined first\n- Missing punctuations in equations\n- It’s unclear what is $TB$ in line 214\n\n[1] Ankenman, J.I., 2014. Geometry and analysis of dual networks on questionnaires. Yale University.\n\n[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.\n\n[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.\n\n[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.\n\n[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.", "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": "BqzMt0mwQ3", "reviewer": "Reviewer_HU6t", "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.", "questions": "List of typos. I suggest that the authors give the manuscript a thorough proof-reading.\n\n* l.39 confusing, why is the optimal sample distance the Wasserstein distance?\n* l.106 in what sense is SWD a “geometric embedding”?\n* l.111 definition of pi?\n* l.116 what does “good approximation” mean here?\n* l.140 precise the meaning of normalized.\n* l.142 distribution -> “probability distribution”\n* l.150 R and tau are undefined\n* 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.\n* 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).\n* 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?\n* l.201 W_{T_B} instead of W_B ?\n* l.205 z_i^(A) undefined\n* l.205 Z^{(B)} or Z_B?\n* l.214 W_{TB} -> W_{T_B}", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The idea is sound, and there is a significant element of originality, especially in the development of the algorithm in Appendix C.", "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.\n\n1. l.238 – 241 these statements require a proof. Especially the convergence, since it was already somewhat delicate in Huizing et al. (2022).\n2. l.263 “Wasserstein” -> “Tree Wasserstein” or else requires a proof.\n3. 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?\n4. l.354 The previous work used n=100, m=80 and not n=80, m=60 as this paper states.\n5. 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?", "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": "gLXI5umixB", "reviewer": "Reviewer_pG9e", "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.\n\nThis 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.", "questions": "Refer to weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "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. \n2. The empirical results are presented clearly, with comparative metrics that directly illustrate the computational runtime saving and clustering performance.\n3. The paper provides a solid background review on optimal transport theory and the tree-Wasserstein distance.", "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)?\n2. 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?\n3. 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?\n4. Figure 3 could be better organized. The paper does not provide a comparison of how other metrics perform on these datasets.\n5. Line 557, The URL is invalid.", "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": "3UU2dTJR5N", "reviewer": "Reviewer_SCu1", "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.", "questions": "**Major**:\n\n- How is the condition number of eq. 5? Does the system become ill-conditioned as matrix becomes large?\n- Section 2.3, what is the convergence rate of the proposed iterative method? How it is affected by the data set size?\n- What is the complexity of Cluster Tree used here?\n- Please add details of memory consumption for your proposed iterative algorithm and approximation to SVD.\n- 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$.\n\n**Minor**:\n\n- In abstract, \"...a fast and recursive algorithm...\" and not \"a fast, recursive algorithm…”\n- In abstract, what does \"low-compute application\" mean?\n- P2, l82: \"... to have a solution\" and not \"to have solution\"\n- P3, l108: Is there a guarantee that the shortest path between any two nodes on a tree is unique?\n- Figure 1: Caption refers to b5, but I don’t see b5 in the graph.\n- In the references, make sure to use capital letters when needed, e.g. use “Wasserstein” instead of “wasserstein”", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 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)$.", "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.", "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)