Fast unsupervised ground metric learning with tree-Wasserstein distance

We were encouraged to read that the reviewer believes the paper can be accepted after a major revision, and we hope the revised manuscript will meet these criteria, particularly addressing the topics of convergence rate, well-posedness and memory consumption. We thank the reviewer for their clear revision asks.

Major:

• The condition number on eq. 5 (now rewritten as new eq. 5, the complete linear system of equations, to which I assume the reviewer is referring) is good (near 1) for the SVD approach, and we believe acceptable for the recursive approach (but we would appreciate any further insights or reinterpretation, as this is not our field of expertise!). For the SVD approach, the condition numbers on $\boldsymbol{U}$ are all very close to 1 with error on the order of 1e-15 (these were computed for (torus) dataset sizes 80 x 60, 100 x 200 and 500 x 500 using numpy.linalg's implementation of the condition number and the standard 2-norm). For the recursive algorithm, although we get full-rank matrices, the condition number is higher, on the order of 1000-3000 for sample sizes of 1000 and 2000 respectively. While this is much larger than 1, it should not affect accuracy within a margin of 1e-5 for the $\boldsymbol{w}$ vectors (based on a 1e-8 numerical precision of jax.float32) -- and these are generally around 1e-3 for large n. As the method scales larger, this could of course be a problem, and one option could be to compute multiple recursive basis sets and either take one with lowest condition number or iterate Tree-WSV with several and average. Thanks for the suggestion to confirm these! Would the reviewer recommend these numbers be included in the paper as a comment on the least squares solver's convergence?

• Regarding convergence rate, we now include empirical convergence (of power iterations) for various dataset sizes 80 x 60, 100 x 200 and 1500 x 300 in Appendix E. In practice, we always observed fast (within 10 iterations) convergence of power iterations, and built-in scipy and jax.numpy least squares / NNLS solvers are presumed to have converged with high accuracy based on the condition numbers. Convergence appears stable with larger dataset sizes.

• ClusterTree implemented in this way has complexity $\mathcal{O}(n\kappa)$ where $\kappa$ is the number of leaf-clusters (in our case set automatically by the hyperparameter controlling the depth of the tree) [1] (Gonzalez, 1985). We have added this to the paper.

• Regarding memory consumption of sparse SVD and the iterative algorithm, we have added the comment: "Because this operation requires computing a large tensor $\boldsymbol{\mathsf{Y}}$ reshaped into a long rectangular matrix of size $(N-1) \times n^2$, the sparse method does not scale well with large $n,m$ from a memory-consumption point of view." We elaborate on this in our response to reviewer 3. In summary, (sparse) SVD has complexity cubic in $n$ (resp. $m$), whereas our iterative algorithm has worst-case complexity quadratic in $n$ (resp. $m$). We observe that for small $n,m$, sparse SVD is faster, whereas for large $n,m$, our iterative algorithm is faster. We have not explained in this detail in-text besides the comment on memory consumption, but can do so if you feel it would help.

• In terms of showing complexity, we have been more cautious in explaining the calculation in section 2.3.2 ("This gives overall complexity per power iteration: $\mathcal{O}\left(N^3+mN+M^3+nM\right) < \mathcal{O}\left(n^3 + m^3 + mn\right)$, using $N<2n-1, M<2m-1$.") We could certainly also compare runtime for a variety of different dataset sizes if the reviewer feels it would help our paper a lot, but we hoped that including the runtime after a number of meta-iterations for each dataset in Table 2 (each of which ran for a similar number of inner power iterations) would provide some estimate for the interested reader. Exact time complexity curve fits will be limited to the exact number of edges ($N$) in any tree construction, too. Does the reviewer feel that including a time complexity fit in the appendix on the number of edges would improve the paper greatly? If so, we can produce one.

[1] Teofilo F. Gonzalez. Clustering to minimize the maximum intercluster distance. Theoretical Computer Science, 38:293–306, 1 1985. ISSN 0304-3975. doi: 10.1016/0304-3975(85)90224-5

(Part 2)

• The cost matrix on l. 211 should really read $C$, that is a typo (we have updated that section to aid explanation more generally). This matches the general definition for Wasserstein distance in section 1.1, and is illustrative in this paragraph.

• Figure 3 has been reorganised to be neater, which we hope offers an improvement.

• Regarding other metrics, do you mean other clustering metrics like Adjusted Rand Index, Davies-Bouldin? Or other methods (metric learning approaches) like those in the appendix table, but extended to the other datasets? If you have suggestions on particular preferred comparison metrics, we would be happy to include with this clarifying point.

• All references have been carefully checked and the URL previously on l.557 (Le et al., 2019) now points to the official, valid URL.

Overall, we thank the reviewer in particular for noting the need to review and explain ClusterTree and the SVD versus recursive algorithm approaches, since these are indeed core to the Tree-WSV approach. We look forward to discussing any ongoing concerns and suggestions further.

(Part 1)

We thank the reviewer for the response and for recognising the paper's strengths (including providing rigorous theoretical support for our TWD method, clear empirical results and a review of the related literature). We also agree that the paper is indeed an improved version on the workshop paper mentioned; however, since the workshop was non-archival, we would ask the reviewer to evaluate entirely and independently for this submission to ICLR (and, indeed, our read of the review supports that the reviewer has done this!).

In terms of the weaknesses:

• We apologise for not providing adequate background regarding ClusterTree. The details are indeed equivalent to the method in Le et al. (2019), although they do not yet refer to the method as ClusterTree; this was adopted by later authors who reference their work. It is not quite the same as Algorithm 1 in that paper but rather the extension described in the next paragraph (using Alg. 2). Further background on clustering to form trees (and the Farthest-Point Clustering Algorithm) are provided in their main-text and appendices. We have added a few sentences about this and referenced earlier authors, but would like to stress in reply that the exact clustering tree algorithm does not make much difference in our hands, since we learn edge weights. We have compared with using QuadTree for initialisation, although we prefer ClusterTree as there is more flexibility in the number of leaves per cluster. Note that Indyk and Thaper extend QuadTree to higher dimensions, but we agree the reference in that point in text was unclear -- we now also reference the original QuadTree paper and Indyk and Thaper as a follow-on. Here is how we have re-introduced ClusterTree: "Note that here, ClusterTree is a 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 et al. (2019) and Gonzalez (1985). ClusterTree implemented in this way has complexity $\mathcal{O}(n\kappa)$ where $\kappa$ is the number of leaf-clusters (in our case set automatically by the hyperparameter controlling the depth of the tree) (Gonzalez, 1985)." Also see [1]. We based our implementation of ClusterTree on the following publicly available repository and would be happy to include it in the paper if you feel it would strengthen the background on tree construction: https://github.com/oist/treeOT

• We agree that Algorithm 1 was scanty on ``bounds'' for large versus small datasets and we have now included some. This is really to do with memory issues when using SVD on large datasets. Because the sparse SVD method requires as input the entire $\boldsymbol{Y^\prime} \in \mathbb{R}^{(N-1) \times n^2}$ matrix, when $n$ is very large that matrix becomes prohibitive to store; in addition SVD itself scales poorly in terms of memory allocation (see https://fa.bianp.net/blog/2012/singular-value-decomposition-in-scipy/). Standard SVD with an input matrix of size $p \times q$ with $p>q$ has complexity $\mathcal{O}(pq^2)$, and sparse SVD should be similar, although we cannot find exact complexity references. Hence, as $n < N < 2n-1$, SVD on the matrix $\boldsymbol{Y^\prime}$ has complexity cubic in $n$. By our estimation, the recursive algorithm is fast: it recurses until at most depth $n$ (the number of leaves), and each recursion can include at most $n$ operations, so the total worst-case complexity is quadratic in $n$ (and symmetrically for $m$). In terms of the latter question of whether one or the other of $n,m$ is larger: separate methods could be used to compute each of the two basis sets, since there is no interaction between the trees. Towards the main points of our work, either method is fine and perhaps it would be simpler to present all work using the recursive algorithm. However, since it performs slower for smaller $n,m$ empirically, and since we felt that this was easily addressed with a more standard approach than the newly introduced algorithm, we wanted to include both. If one of $n,m$ is small and the other is large, we therefore suggest to use SVD on the small one and our recursive algorithm on the large one. There also may be some reasons to prefer SVD for its condition number (see comment by reviewer 4). We are interested to know what you think given this extra background -- please let us know!

• We have made the notation consistent throughout, using ICLR guidelines to bold italicise vectors and matrices respectively. We believe this is now all consistent and apologise for the poor notation in the submission rush before! We have also restated Theorem 2.2 as a lemma and a theorem which we believe will help to make sense of notation, and expanded on notation definitions throughout. Notation in section 1.1. is consistent with 2.1.

[1] Teofilo F. Gonzalez. Clustering to minimize the maximum intercluster distance. Theoretical Com- puter Science, 38:293–306, 1 1985.

Questions:

• We have addressed the typos. Thank you for your careful read!

• l. 39 previously called Wasserstein distance an optimal sample distance. Of course, this is a heuristic used by us to introduce OT, and the reality is more subtle (a metric that takes into account both mappings (weightings of features) and distances between features). We have tried to rephrase this as: "$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: ..."

• l. 106: SWD is a geometric embedding insofar as it is a projection. We have rephrased this part and incorporated it in the general computational complexity background section: "Sliced-Wasserstein distance (SWD), 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., 2016). SWD can be expressed as a special case of a geometric embedding, tree-Wasserstein distance (TWD), when the tree structure comprises a root connected to leaves directly (Indyk & Thaper, 2003; Le et al., 2019)."

• l. 111: pi is the transport plan. We now define it (as well as U, the set of joint probability distributions): "Given a tree metric $d_\mathcal{T}$ on leaves $x,y$ and transport plan $\pi$ between measures $\mu,\nu$, the TWD can be written as (previous expression, does not render here), where $U (\mu, \nu ) =$ (full expression inserted in-text, does not render here) is the set of joint probability distributions with marginals $\mu$ on $\mathcal{X}$ and $\nu$ on $\mathcal{Y}$ ."

• l. 116. Goodness of approximation of TWD to 1-WD: This is a good question. There are theoretical results that guarantee there exist trees that can fully express the 1-WD, and empirical results that show that learning edge weights provides good approximation -- see Yamada et al. 2022. We have updated to include ``TWD is a good empirical approximation ...'' An avenue for future work could include seeing whether our theory (specifically Lemma 2.3 / Theorem 2.4) could help to provide bounds on the approximation, via the constructive proof in Yamada et al. (2022).

• l. 140: normalized = sums to 1 (added).

• l. 142: changed.

• l. 150: $R, \tau$ are now defined, and $\Phi_{\boldsymbol{A}}$ described.

• l. 151: equivalence now states that $\tau = 0$, thank you.

• l. 157: We believe the remark in parentheses is correct, but have clarified it to read "(since it requires computing $m^2$ Wasserstein distances where each distance is $\mathcal{O}(n^3\log(n))$ then $n^2$ Wasserstein distances where each distance is $\mathcal{O}(m^3\log(m))$)". To explain: Huizing et al. (2022) state that a single power iteration is $\mathcal{O}(n^2m^2(n \log(n) + m \log(m)))$ for $n$ samples and $m$ features. Our comment in parentheses was to explain that calculation: each Wasserstein distance on $\mathbb{R}^n_+$ is computed in $n^3\log(n)$ and there are $m^2$ of these to compute in the distance matrix, whereas on $\mathbb{R}^m_+$ there are $n^2$ distances to compute and $m^3\log(m)$ for each. Summing and factorising you get the power iteration complexity. Note that we have changed the abstract to use this exact quantity rather than a (loose) approximation.

• l. 201: We agree and have replaced $\mathcal{W}_{T_B}$ instead of $\mathcal{W}_B$.

• l. 205: We now define $\boldsymbol{z_i^{(A)}}$ in the theorem: "where $\boldsymbol{z_i^{(A)}}$ is the $i$th column of $\boldsymbol{Z^{(A)}}$."

• l.205: Should read $\boldsymbol{Z^{(B)}}$, thanks!

• l.214: Typo fixed, thanks!

To reiterate, the idea to restructure the previous theorem as a lemma and theorem and to carefully review Huizing et al. helped to formalise the theory and improve the manuscript greatly -- thank you for this! We look forward to any further engagement or ideas you might have.

(Part 1: Weaknesses)

We thank the reviewer for recognising the originality and for their interest in the algorithm in the Appendix, about which we were particularly happy. We apologise for the rushed appearance and assure the reviewer that we have carefully proof-read, corrected typos, and rewritten the manuscript accordingly.

Please note the general updates across reviewers, which include:

• A careful re-read and incorporation of the theory in Huizing et al., including a new Lemma about existence of solutions for power iterations, and more caution regarding convergence guarantees.
• A full attempt to standardise notation and remove typos.
• Rephrasing Theorem 2.2 (thanks so much for this suggestion!)

In particular regarding weaknesses:

• Prev. l. 238-241: We now address more rigourously one of the guarantees from Huizing et al. in Lemma 2.2 (showing that the SV problem / power iterations has a solution). We believe that we could extend the argument from Huizing et al. using a large $\tau$ to support uniqueness and convergence of the SV solution in our case, too. For now, we show empirical convergence with the same linear dynamics as in Huizing et al. for different dataset sizes (implicitly for $\tau=0$) in Appendix E. In practice, we always observed quick convergence of power iterations.

• In prev. l. 263 "Wasserstein" is now replaced with `"tree-Wasserstein" (we agree with the reviewer on this point completely: in particular, the paper by Yamada et al. (2022) [1] showing TWD approximates the 1-WD does so mostly via empirical results, and while the equivalent empiric results are strong, we do not know a convincing proof of the bound of the approximation).

• We have rephrased Theorem 2.2 as a Lemma and the next paragraph as a Theorem (with some more rigour in the proof), as suggested -- we think that this was an excellent suggestion that helps to both formalise and clarify the text. Thank you! The lemma is identical to previous Theorem 2.2. The Theorem reads (sic): "Given any tree $T_A$ with leaves the rows of a data matrix $X$ such that the root node of $T_A$ has degree 3 or more, and a tree $T_B$ with leaves given by $B$ the columns of $X$, there exists a unique, non-zero solution for $\boldsymbol{w_A}$ in the system of linear equations."

• We agree that the sizes of the experimental datasets are slightly different. Note that we get similar results on the same dataset size ($n=100,m=80$) and would be happy to replace the figure with one using these dimensions if the reviewer feels it would strengthen the submission (we used the smaller just because it is faster to run the full WSV method multiple times). However, we also included row/column permutations of the matrix, so the dataset is already considerably different from the original. To clarify that this is the case, we replaced "as employed in previous work" to "modified from prior work" in the first paragraph of section 3.

• Regarding the different metric: we compare the final distance matrices between methods, as compared to the distance matrices found with WSV. We thought that the Frobenius norm is appropriate since it penalises each matrix element, whereas the Hilbert metric considers upper and lower bound logarithmic differences between the two matrices (nonetheless useful as these bounds get closer, i.e. for convergence). Is there another reason to prefer the Hilbert metric? Note that our results are broadly similar when using the Hilbert metric $d_H$: for example, on the same $60 \times 80$ dataset, Tree-WSV scores using $d_H$ are 1.38-1.41 while SSV scores are 21.2-21.7. If the reviewer would like us to replace the accuracy scores in Fig. 2 with this metric, we are happy to do so.

[1] M. Yamada, Y. Takezawa, R. Sato, H. Bao, Z. Kozareva, and S. Ravi. Approximating 1-Wasserstein distance with trees. Transactions on Machine Learning Research, 2022. ISSN 2835-8856. URL https://openreview.net/forum?id=Ig82l87ZVU.

We thank the reviewer for the comment and improved rating in response.

The reasoning to include the Hilbert metric to aid comparison / convergence proofs makes sense -- we will add that in for Figure 2 (we already adopted it for the convergence figure in the appendix). We are also happy to spell out without loss of generality (indeed, we ought not assume this a standard abbreviation!).

(Part 3)

• We tried to simplify the notation in Algorithm 1, and more importantly now introduce this notation in the text (under section 2.3.2 largely), which we hope will help the reader. We found a typo -- $A_{leaf}$ should be replaced with the number of edges -- and we now define convergence. Thanks for noticing this!

• References have been updates to include publisher, links to official papers, and capital letters where needed. We apologise for the rushed impression of the referencing section at the time of initial submission.

• Notation has all been updated to meet ICLR standards and to be consistent between section 1.1. and 2.1.

• Regarding minor points, all four have been corrected. We thank the reviewer for taking their time to find these!

Questions:

• l. 194: We mean that we use the same notation i.e. $\boldsymbol{a_i }$ is a normalised row in a data matrix. We now simply rewrite this to avoid reader fatigue paging back and forth.

• We have added the size of $\boldsymbol{U}$ in the manuscript: in text, it is number of edges x number of edges.

• We do not believe choice of tree construction influences the algorithm, up to the number of edges in the tree and whether the root has 3 or 2 children. We investigate these properties in Figure 2. We have also tried methods like QuadTree, and get similar results, except of course for differences in the time it takes to construct the tree. Alternative clustering methods to construct trees could be an avenue for future exploration as this could impact the meta-iterations, but learning the tree weights is identical on any tree structure.

• Regarding convergence, we have added our empiric stop-point / threshold to the algorithm for when weight vectors converge, and make a few more comments on convergence in Appendix E, where we also show empirical convergence for different dataset sizes, with dynamics matching those in Huizing et al.

• For hyperparameters, we set the number of children in each ClusterTree and the convergence threshold. The number of children / tree structure is briefly discussed and compared in Figure 2 and surrounds. The convergence threshold does not seem to matter as empirically we observe convergence to small values quickly.

• The Euclidean metric baseline uses the Euclidean metric (i.e. L2 norm) as the distance between features (rather than say our learned metric); we then compute the ASW based on that metric (a bit like k-means distances between samples compared on labelled clusters).

• We do not use WSV as a baseline in Table 1 because it is simply too computationally expensive to even run! It would take an estimated over 30 hours just to run the smallest genomics dataset in our empiric section. Note that even in the original paper by Huizing et al. (2022) they only used SSV on their genomics dataset, which is identical to our first dataset (PBMC) -- we did exactly the same, and got the same ASW score for SSV.

• Regarding the use of other TWD methods as baselines: we are not sure what you mean here. Can you clarify which other TWD methods you are referring to? We are not aware of any that learn ground metrics except our own.

Thanks again for these detailed comments and we look forward to engaging more!

(Part 2)

• We have updated the description of the sliced-Wasserstein distance. To answer the reviewer's question: SWD is a special case of TWD when the tree structure has just a root connected to leaves (which we now write). For this reason, we did not implement SWD as a comparison in this work: there is a loss of expressivity since there are fewer edge weights, so it is a worse low-rank approximation than a larger tree.

• We have updated all notation to match ICLR standards: $\boldsymbol{w}, \boldsymbol{Z}$ are a vector and a matrix respectively (we now state "Let $\boldsymbol{w}$ be the vector of weights"). We also rewrote l. 111 to explain the notation in more detail: "Given a tree metric $d_\mathcal{T}$ on leaves $x,y$ and transport plan $\pi$ between measures $\mu,\nu$, the TWD can be written as (previous expression, does not render here), where $U (\mu, \nu ) =$ (full expression inserted in-text, does not render here) is the set of joint probability distributions with marginals $\mu$ on $\mathcal{X}$ and $\nu$ on $\mathcal{Y}$ ."

• In terms of the notation of Wasserstein distance in line 92 and the notation of TWD in line 111: we have used former authors' notation, but we agree it is confusing. We hope that the subscript $\mathcal{T}$ makes clear when we refer to TWD rather than the full 1-WD and have now corrected so that this applies where appropriate; is that the case? If not, do you know of alternative notation that would help make this clear?

• Eq. (2): Thanks for pointing this out, as it was quite confusing. We now write $\boldsymbol{a,b}$ as $\boldsymbol{x,y}$, consistent with l.111 above. And size($\boldsymbol{w}$) is now rewritten as dim($\boldsymbol{w}$).

• Regarding the goodness of approximation of TWD to 1-WD: the paper by Yamada et al. (2022) [3] showing TWD approximates the 1-WD does so mostly with empirics, and while the equivalent empiric results are strong, we do not know a convincing proof of the bound of the approximation. There is a constructive proof in that paper which suggests that our Theorem 2.3 might actually help to show this, but we believe it should stand as future work as it would take significant effort. We have tried to clarify in the rewrite of Proposition 2.1 and onwards that our method takes inspiration from Huizing et al. (2022), who use the full 1-WD, but we use TWD. How well TWD approximates 1-WD is well-established empirically but should be proven (out of scope for our work).

• l. 151: We now define $R$ as a norm regulariser (identical to Huizing's description).

• We recognise the frequent and confusing use of the word "basis." These should all mean the same thing, and we have tried to systematise this in the new lemma / proof format for what was previously Theorem 2.2, as well as the first paragraph in section 2.3.2, where we now define $\boldsymbol{U}$ to be this basis set for $\boldsymbol{Y^\prime}$.

• ``Full WSV'' was used to refer to the full WSV approach, versus the speed-ups like Sinkhorn's SV mentioned in Huizing et al. We have replaced all mentions of "full WSV" with "standard WSV", as we agree that this may read more clearly.

• Regarding clarity on the size of edge weights vs number of nodes in the tree: Each node that is not a root has exactly one edge as parent of that node, and this counts all the edges in the tree. The root's edges are all counted by child-nodes of the root, so that means the number of edge weights is one less the number of nodes in the tree.

• We agree that Proposition 2.1 had a typo as noted by another reviewer, too -- we have amended so that $W_{B}$ is now $W_{\mathcal{T}_{B}}$ (and the same for the As), thanks!

• We apologise for the lack of definition of some notation. We have now made sure that $\circ, \lambda_A$ are defined in Prop 2.1 and $\boldsymbol{z_i^{(A)},Z^{(B)}}$ in the paragraph preceding. We have also made the notation in Prop 2.1's proof in the appendix consistent, and defined $\Phi$ in Section 1.1 when we address Huizing et al.'s iterative SV approach, and repeated this in the proof.

• We have rephrased Theorem 2.2. into a lemma and a theorem, a suggestion by another reviewer which we think significantly clears how the theory supports unique, non-zero solutions to prop 2.1. The lemma is identical to previous Theorem 2.2. The Theorem reads (sic): " Given any tree $T_A$ with leaves the rows of a data matrix $X$ such that the root node of $T_A$ has degree 3 or more, and a tree $T_B$ with leaves given by $B$ the columns of $X$, there exists a unique, non-zero solution for $\boldsymbol{w_A}$ in the system of linear equations."

• We also have a new Lemma and proof showing that power iterations also have a solution, which deepens the theoretical underpinnings.

(Part 1) We thank the reviewer for a thorough read and suggestions. We appreciated the comment that the "proposed method demonstrates efficiency compared to WSV and SSV presented in Huizing et al. (2022)" and we have addressed the following major points in our reply to all reviewers above:

• Notation
• Proving there is a solution to the power iterations problem
• Convergence (empirically)

In terms of the specific weaknesses (we answer in the same order to help the reviewer to cross-reference):

• We were glad to read the related works mentioned, and apologise for not being aware of these -- [2] (Mishne et al.) is particularly relevant for us (and could inspire possible extensions from our work!). We also thought [1] (Ankenman)'s construction of partition trees using a random walk algorithm was intriguing. We have now included a paragraph about these works starting at line 170. Specifically, we write:

While our approach to 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 on graphs (Shahid et al., 2016) and in tree-embeddings (Anken- man, 2014; Mishne et al., 2018; Yair et al., 2017). These methods do not learn the ground metric in the same way as our proposal, but Mishne et al. (2018), Ankenman (2014) and Yair et al. (2017) describe iterative metric-learning of the tree metric and tree construction, which is related.

• Regarding lack of clarity on the explanation of "how the Wasserstein distance can serve as a tree distance in Proposition 2.1", we have now rephrased the explanation starting at line 231 to explain the intuition in the embedding, and been more precise about the nature of the approximation of TWD to 1-WD.

• On ``It’s also not evident whether there exists a tree for which the tree distance would correspond to a Wasserstein distance'' please see Appendix A where Proposition 2.1 is proved and refer to Yamada et al., 2022 [3]: Proposition 1 in Yamada et al. ensures that for any distance metric on the support of two discrete measures, there exists a tree such that TWD equals to the 1-Wasserstein distance with that metric (using the shortest path on the tree for the tree metric).

• We agree that the real-world application was restricted to single-cell RNA sequencing data, while the introduction cites various data types as motivation. We have run the algorithm on orientation tuning in calcium-imaged neurons in V1 with interpretable results (we get sharper ring-like structures as compared to Euclidean metric, but these are still visible albeit with more noise using just a Euclidean metric, as might be expected for an orientation dataset). We could include these results as an appendix, but we wonder whether it would be best to remove the references to neural data in the introduction and leave these as options for future work, given the space limitations (which is what we have done for now). Do you have any preference as a reader to either of these two options?

• Regarding only comparing the two competing methods: we include a full table of comparisons for one of the datasets in the Appendix, which showcases other methods' poorer performance. If it would assure the reviewer, we can run a subset of these on the other datasets, too. On the second point, since our main aim was to show that the Tree-WSV algorithm improves on WSV in terms of compute time and outperforms SSV in terms of accuracy, we did not include other distance metric learning approaches. If the reviewer has any in particular to recommend, we are happy to compare these.

• In terms of complexity on $n$: We realise this was unclear in the abstract vs. later in the paper, and thank you for pointing it out. We assumed that the original $n,m$ in the Huizing paper were of a similar order and dropped the log term to get to an approximate quintic complexity. However, this is indeed unclear and not explained, so we will instead use the full complexity: $\mathcal{O}(n^2m^2(n log(n) + m log(m))$ for $n$ samples and $m$ features (l.018 in updated manuscript).

(See next part)

[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] M. Yamada, Y. Takezawa, R. Sato, H. Bao, Z. Kozareva, and S. Ravi. Approximating 1-Wasserstein distance with trees. Transactions on Machine Learning Research, 2022. ISSN 2835-8856. URL https://openreview.net/forum?id=Ig82l87ZVU.

We thank the reviewer for engaging and providing clarification on their point about alternative TWD methods (the first part of the follow-up). The question is a good one: indeed, the method by Yamada et al. (2022) really did inspire this work and is aligned since TWD is calculated in the same way and we also learn edge weights – so much so that when we explored this idea originally, we did it in the way suggested by the reviewer (using cTWD to compute cTWD between samples and subsequently using this as a ground metric between features, then iterating with the tree construction as per Yamada et al. (2022))! However, we found empirically that the algorithm did not seem to converge over several SV iterations, and reasoned that it could not scale to reasonable $n$ and $m$ greater than 500. This is because of two problems which motivated many of the extensions in our paper: first, cTWD as it is stated required computing all $n^2$ or $m^2$ tree-Wasserstein distances at each iteration. Second, solving the LASSO problem in the way suggested is computationally expensive and memory-inefficient, because the full $n^2$ or $m^2$ distances data is ideally used as input.

To expand on time complexity: each LASSO has complexity $\mathcal{O}(p^3+kp^2)$ where $p$ are the number of features and $k$ the number of samples (Efron et al., 2004 [1]). In the case of cTWD, $p=N, k=n^2$ (N is number of edges and n number of leaves), so the total complexity is at least $\mathcal{O}(n^4)$, since $n<N<n^2$. In addition, calculating all $n^2, m^2$ TWDs still has $mn^2$ and $nm^2$ (resp.) complexity. So each iteration is at least quartic in the bigger of $n,m$ in complexity, which is worse than the complexity of SSV. Our solution – to reduce the problem to solving for a subset of the full leaves (with related Lemma 2.3, Theorem 2.4) – moves this to cubic complexity by formalising the set-up in Yamada et al. to a linear system of equations problem. Indeed, one could apply our new theory to the Yamada et al. 2022 tree weight calculation (with non-negative LASSO or NNLS as the solver) and presumably speed up the tree weight calculation there!

To expand on memory consumption: for even relatively small $m$ or $n$, we have to calculate the same $Y^\prime$ matrix for each iteration (shape $(N-1)\times n^2$) and this is prohibitive quickly (even for $n=500$, this becomes a problem!). We solve this by using SVD or the recursive algorithm in the appendix to make $U$, which has size $(N-1)\times(N-1)$.

Why were these not a problem for the Yamada et al. paper which used a single weight-learning iteration to calculate a single tree? First, their use-case is on word-embeddings with a tree of ~5000 leaves (words), where each word has an embedding dimension of 300. Computing the Euclidean distances between words when each is 300-dimensional is fast and can be done in parallel (we also use parallel processing, albeit via JAX). In addition, while it is relatively quick to calculate a subset of TWDs (hundreds – as they used to compare to full Wasserstein distance), this becomes much worse computationally when computing all $5000^2$ TWDs for all document pairs – our method would require only computing at most 9999 at each iteration. Second, and importantly, Yamada et al. (2022) used a trick with their hyperparameters in the algorithm that sped things up: they set a max number of pairwise distances to sample in the weight-finding algorithm. This effectively reduces the size of the $Y^\prime$ matrix and the least squares (LASSO) problem, although it does so randomly. They account for some of this randomness using a sliced version and averaging across tree constructs. We think this could be an interesting (stochastic version) algorithm to explore, but the randomness seemed to induce errors when we tried to run a modified version of Yamada et al. with singular value iterations over edge-weight learning (both previously and in the last few days). In particular, the algorithm did not appear to converge in a reasonable number of iterations on the toy dataset (using $m,n=200,300$ and hyper-parameter of number of random sub-samples 100 or 1000). From Huizing et al., 2002 [2] Figure 3, we assumed that stochastic iterations in general may have low yield / take longer to converge, but notably we have not tested this thoroughly in the tree setting – we agree that this would be good to follow up.

We hope this explanation as to why we did not include cTWD as a comparative method helps! For these reasons, we chose to focus on the theory towards reducing the size of the least squares problem. However, it would be useful to try to compare stochastic-type approaches with averages across trees (or similar) as an alternative in the future, and we would be happy to add this as a suggestion.

[1] Efron et al., Annals of Statistics, 2004 (https://arxiv.org/pdf/math/0406456) [2] Huizing et al., ICML, 2022 (https://proceedings.mlr.press/v162/huizing22a/huizing22a.pdf)