UltraTWD: Optimizing Ultrametric Trees for Tree-Wasserstein Distance

Thank you for your valuable questions. We will revise the paper accordingly and add more discussion on related work.

---

Question 1. Could you provide a clear definition of the tree structure used in your method? Specifically, is the method limited to binary trees, or does it accommodate more flexible tree structures?

Answer 1. Our method is based on ultrametric trees, which are rooted but not necessarily binary. While we empirically followed [1] and used the hierarchical minimum spanning tree procedure to construct a binary ultrametric tree, the method itself is not limited to binary structures. Non-binary ultrametric trees can also be constructed using algorithms like UPGMA. We plan to explore more flexible tree structures in future work.

[1] Chen, S., Tabaghi, P., and Wang, Y. Learning ultrametric trees for optimal transport regression. AAAI, 2024.

---

Question 2. How is the tree structure updated during the iterative projection process? Does the update involve regrouping leaves or modifying the overall tree topology?

Answer 2. Great question. During the iterative projection process, we update the distance matrix $D_T \in \mathbb{R}^{n \times n}$ to satisfy the strong triangle inequalities. This projection operates directly on the matrix entries and does not explicitly modify the tree structure or regroup leaves. The tree topology is constructed only once, in the final step of Algorithm 2, using a minimum spanning tree algorithm applied to the projected matrix.

---

Question 3. Could you provide a visualization of the constructed tree from the proposed methods?

Answer 3. Please see the figure. We will include it in the revision.

---

Question 4. Can you elaborate on the practical benefits of the $O(n^3)$ complexity for UltraTWD-IP, especially when compared to the $O(n^3 \log n)$ complexity of traditional optimal transport methods?

Answer 4. We would like to clarify that the $O(n^3)$ complexity of UltraTWD-IP refers to a one-time cost for learning the ultrametric tree. After this step, each pairwise distance $W_T(\mu, \nu)$ can be computed in just $O(n)$ time-much faster than the $O(n^3 \log n)$ cost of computing exact $W_1(\mu, \nu)$. As discussed in Section 3.6, the total time complexity for computing $W \in \mathbb{R}^{N \times N}$ is:

• $O(N^2 \cdot n + n^3)$ for UltraTWD-IP
• $O(N^2 \cdot n + n^2)$ for UltraTWD-GD
• $O(N^2 \cdot n^3 \log n)$ for traditional $W_1$.

For example, on the BBCSport dataset with 3,000 valid words per distribution, computing $W_T \in \mathbb{R}^{100 \times 100}$ using UltraTWD-IP takes only 0.9 hours (including 0.3 hours for tree learning), compared to 2.3 hours with traditional optimal transport.

---

Question 5. When reporting the computation time in Figures 7 and 8, which algorithm was referred to in UltraTWD?

Answer 5. Apologies for the confusion. Since all three UltraTWD variants use the same type of ultrametric tree, their computation times are nearly identical. We reported the average time across three algorithms and will clarify this in the revision.

---

Question 6. Could you include empirical comparisons with supervised tree-Wasserstein distance methods and standard WMD, given that your method claims to be an efficient approximation?

Answer 6. As suggested, we will include the following comparisons:

• Supervised methods: We compare with STW [1] and UltraTree [2] (UltraTree is discussed in the main text, Line 367, page 7). STW learns a new distance via contrastive loss rather than approximating $W_1$, resulting in large errors and poor precision. STW training is also time-consuming, taking 3.4 hours on BBCSport, whereas UltraTWD-GD requires only 34 seconds to train. As shown in Table R1, our methods consistently outperform both STW and UltraTree in accuracy and efficiency.

  Table R1. Performance comparison on the BBCSport dataset.

  Metric	RE-W$\downarrow$	Precision$\uparrow$	Total Time (hour)
  STW	0.643	0.335	3.5
  UltraTree	0.022	0.842	0.6
  UltraTWD-GD	0.016	0.868	0.1
  UltraTWD-IP	0.014	0.885	0.5

  [1] Takezawa, Y., Sato, R., & Yamada, M. Supervised tree-wasserstein distance. ICML, 2021.

  [2] Chen, S., Tabaghi, P., and Wang, Y. Learning ultrametric trees for optimal transport regression. AAAI, 2024.

• WMD: Standard WMD is used as the ground-truth $W_1$ in the main text. Figures 7 and 8 (page 16) show our methods are significantly faster, even when including tree learning time (see also Answer 4). Table R2 compares SVM classification accuracy, where WMD yields the highest accuracy due to exact optimization, while our methods remain competitive with much lower cost.

  Table R2. Comparison of SVM classification accuracy.

  Dataset	BBCSport	Reuters
  WMD	0.874	0.919
  UltraTWD-GD	0.838	0.900
  UltraTWD-IP	0.839	0.905

We truly appreciate your positive feedback. Your valuable suggestions will help us improve the work.

---

Comment 1. Comparing the running time of $W_T$ with the exact computation of $W_1$ is unfair. It is not clear how many iterations does the GD algorithm require.

Response 1. To ensure a fair comparison, we consider the total time complexity required to compute $W \in \mathbb{R}^{N \times N}$.

• For exact $W_1$: $O(N^2 \cdot n^3 \log n)$.
• For UltraTWD: (including the tree learning time)
  • $O(N^2 \cdot n + n^3)$ for UltraTWD-IP
  • $O(N^2 \cdot n + n^2)$ for UltraTWD-GD

As described in Algorithms 2 and 3, we use 1 iteration for UltraTWD-IP and 10 iterations for UltraTWD-GD. On the BBCSport dataset with 3,000 valid words:

• UltraTWD-GD learns the tree in 34s and computes each $W_T(\mu, \nu)$ in 0.4s using $O(n)$ time,
• Each $W_1(\mu, \nu)$ takes 1.7s with $O(n^3 \log n)$ time.

To compute $W \in \mathbb{R}^{100 \times 100}$, the total time is 0.9h (IP) and 0.6h (GD), both significantly faster than 2.3h for exact $W_1$. We will revise this comparison to be more rigorous.

---

Comment 2. For the Sinkhorn algorithm, the authors set $\lambda=1$ and the maximum number of iterations to 100. Can the authors explain how they chose these parameters? I also would like to know how the authors have chosen the regularization parameter $\lambda=0.001$ for the weight-optimized methods and the number 3 of trees in the sliced methods.

Response 2. We would like to clarify:

• For Sinkhorn, we followed [1], using $\lambda=1$ and 100 iterations to balance accuracy and runtime. Smaller $\lambda$ and more iterations may improve accuracy but greatly increase computation time.
• For weight-optimized and sliced methods, we followed [2], using 3 trees. We selected $\lambda=0.001$ based on the highest Pearson correlation with true Wasserstein distance.

As suggested, we test additional $\lambda$ values. As shown in Table R2, larger $\lambda$ improves RE-W but often lowers retrieval performance. UltraTWD methods consistently outperform these baselines across metrics. We will include these results in the revision.

Table R2. Performance comparison on the BBCSport dataset.

[1] Chen, S., Tabaghi, P., and Wang, Y. Learning ultrametric trees for optimal transport regression. AAAI, 2024.

[2] Yamada, M., Takezawa, Y., Sato, R., Bao, H., Kozareva, Z., & Ravi, S. Approximating 1-wasserstein distance with trees. TMLR, 2022.

---

Comment 3. I believe the following papers can be cited in this work.

Response 3. We will cite these relevant literature in the revision.

---

Comment 4. There are no theoretical guarantees on the convergence rate of their proposed IP and GD algorithms, and hence, it is hard to judge whether the proposed methods have a comparable time as the previous methods or not.

Response 4. We acknowledge the concern. While formal convergence rate analysis is challenging due to the non-convexity of the problem, our algorithms are designed to be lightweight and stable in practice. As shown in Figure 4 (Section 4.4), both IP and GD exhibit fast empirical convergence within several iterations. Moreover, Table 4 demonstrates that UltraTWD-GD achieves favorable runtime–performance trade-offs, learning trees in just 18–91 seconds across datasets while outperforming baselines.

---

Comment 5. The experimental results only compare the proposed methods with the existing TWDs and not other approximations of the Wasserstein distance, such as sliced Wasserstein distance (the SWD and not sliced TWD).

Response 5. As noted in Section 1, our work focuses on advancing the TWD framework. While sliced Wasserstein and its variants require equal-sized point clouds and are not directly applicable to our bag-of-words distributions, we reformulate the text data into weighted point clouds by treating each word as a point. We then compare with two sliced OT methods that support unequal masses—SOPT [1] and SPOT [2]. As shown in Table R3, our methods achieve significantly better performance across all metrics.

Table R3. Performance comparison on the BBCSport dataset.

[1] Bai, Y., Schmitzer, B., Thorpe, M., & Kolouri, S. Sliced optimal partial transport. CVPR, 2023.

[2] Bonneel, N., & Coeurjolly, D. Spot: sliced partial optimal transport. ACM TOG, 2019.

Thank you for your thoughtful review, and we sincerely appreciate your recognition of the strengths of our work, including the clear motivation, innovative algorithmic design, and strong application. We will revise the paper accordingly.

---

Comment 1. Is the closer your distance is to the Wasserstein distance, the better the performance of your new distance will be. If so, it stands to reason that the traditional distance effect is better than your method. If not, what is the significance of this approximation? Your experiments have verified that your distance is better than the traditional Wasserstein distance. How to explain this?

Response 1. Thank you for raising this important point. we apologize for any confusion and would like to clarify:

In general, the closer a distance is to the true Wasserstein distance, the better it preserves the underlying geometric structure, which is beneficial for downstream tasks. Our method aims to approximate $W_1$ more accurately than existing tree-based approaches, and it outperforms other approximations. However, it still performs slightly worse than the true Wasserstein distance. In the main text, we did not include comparisons with the true $W_1$ due to computational cost, but we have now conducted these experiments. As shown in Table R1, the true Wasserstein distance achieves the highest classification accuracy, while our method ranks second. We will include these results in the revised version.

Table R1. Comparison of document classification accuracy.

---

Comment 2. Although the authors provide some algorithms to solve the proposed problem, the paper does not provide strong theoretical guarantees on the convergence or optimality of the proposed algorithms.

Response 2. Theoretical convergence guarantees are extremely challenging to establish due to the non-convex and NP-hard nature of the problem. Similar to our work, prior studies [1, 2] employing gradient-based approaches also acknowledge the absence of such guarantees. For instance, [1] states: "While we provide no theoretical guarantee to find the global optimum..."

Despite this, our methods demonstrate empirical convergence, as shown in Figure 4. Moreover, they are grounded in well-established principles: projection theory for UltraTWD-IP and gradient descent for UltraTWD-GD, both of which exhibit robust performance across multiple datasets.

[1] Chierchia, G. and Perret, B. Ultrametric fitting by gradient descent. NeurIPS, 2019.

[2] Chen, S., Tabaghi, P., and Wang, Y. Learning ultrametric trees for optimal transport regression. AAAI, 2024.

---

Comment 3. How close the proposed distance is to the traditional distance has not been analyzed theoretically.

Response 3. Thank you for pointing this out. Our analysis focuses on empirical approximation quality, reported via relative error metrics RE-D and RE-W (Table 3). Theoretically, we now clarify that the approximation error of $|W_T(\mu,\nu)-W_1(\mu,\nu)|$ can be bounded by the difference in the cost matrices. Specifically, we have:

$

|W_T(\mu, \nu) - W_1(\mu, \nu)| \le 2||D||_\infty + ||D_T - D||_\infty \le 2||D||_\infty + ||D_T - D||_F

$

This result shows that the closer our learned ultrametric $D_T$ is to the cost matrix $D$, the closer the corresponding tree-Wasserstein distance is to the 1-Wasserstein distance. Since our optimization directly minimizes $||D_T-D||_\infty$ or $||D_T - D||_F^2$, the bound ensures that $W_T$ serves as a meaningful and controlled approximation to $W_1$. We will include this analysis in the revised version.

---

Comment 4. The generalization analysis of the proposed new method is also lacking.

Response 4. As an unsupervised method, UltraTWD does not rely on labeled training data and is inherently designed to generalize across domains without retraining. We provide an empirical generalization analysis in Appendix C.4 (Table 7, page 16). In this analysis, we build the tree using the BBCSport vocabulary and evaluate its generalization by testing on 100 randomly generated distribution pairs with varying sparsity levels, where each distribution contains $n$ valid words (sparsity = 1 - # valid words / # total words).

As shown in Table R2, both UltraTWD-GD and UltraTWD-IP consistently achieve low approximation errors (measured by RE-W), demonstrating strong generalization performance across different levels of sparsity and values of $n$.

Table R2. Approximation error (RE-W) under different test sparsity.

Thanks for your valuable comments. We will modify it accordingly.

---

Comment 1. No guarantee for the convergence of the proposed algorithms.

Response 1. Theoretical convergence is extremely hard to guarantee due to the non-convex and NP-hard nature of the problem. Prior work [1, 2] using similar gradient-based methods also lacks convergence guarantees. As [1] states: "While we provide no theoretical guarantee to find the global optimum...". Despite this, our methods show empirical convergence (Figure 4), and are built on solid principles: projection theory for UltraTWD-IP and gradient descent for UltraTWD-GD, both demonstrating strong performance across datasets.

[1] Chierchia, G. and Perret, B. Ultrametric fitting by gradient descent. NeurIPS, 2019.

[2] Chen, S., Tabaghi, P., and Wang, Y. Learning ultrametric trees for optimal transport regression. AAAI, 2024.

---

Comment 2. Algorithm 2 would be costly because the number of projections is too large.

Response 2. Although Algorithm 2 involves $O(n^3)$ projections, each is a closed-form, constant-time operation, and we use only one iteration in practice (Section 3.4). This keeps it practical even for moderately large datasets. For better scalability, UltraTWD-GD (Algorithm 3) has lower $O(n^2)$ complexity per iteration and offers a strong trade-off between speed and accuracy (Table 4, Section 4.4).

---

Comment 3. For the algorithm 3, theoretically is the $D_T$ determined uniquely by $H_T$? Since the number of entries reduce significantly?

Response 3. Yes. Given the rooted tree structure $T$, the node height vector $H_T \in \mathbb{R}^{2n-1}$ uniquely determines $D_T$, where $d_T(i, j) = h(\text{LCA}(l_i, l_j))$ is the height of the least common ancestor (LCA) of leaves $l_i$ and $l_j$ (Equation 8). Since the LCA structure is fixed by $T$, $H_T$ fully specifies $D_T$. This reduces the parameter space from $O(n^2)$ in $D_T$ to $O(n)$ in $H_T$, and allows efficient optimization over $H_T$ in Algorithm 3.

---

Comment 4. According to Table 2, the UltraTWD-IP and GD are competitive to each other, no dominates other in metrics RE-W and Precision. Is there any explanation for that, since UltraTWD-IP appears to be more natural and take longer time to optimize.

Response 4. UltraTWD-IP and GD solve the same problem but use different strategies: UltraTWD-IP uses local projections to enforce ultrametric constraints, while UltraTWD-GD applies gradient descent on node heights. Thus, their results may vary slightly. As a result, UltraTWD-IP tends to yield slightly better downstream performance due to precise constraint enforcement, while UltraTWD-GD runs significantly faster. This trade-off highlights the flexibility of our framework—both methods consistently outperform baselines.

---

Comment 5. UltraTree also produce quite competitive results, does this mean that the proposed method not really improve the overall performance significantly enough?

Response 5. While UltraTree is a strong supervised baseline, our UltraTWD methods show clear advantages in performance, generalizability, and efficiency.

• Performance gap: UltraTWD-IP and GD consistently outperform UltraTree across tasks. In document retrieval (Precision), UltraTWD-IP significantly improves over UltraTree by 4%–5% on all datasets (see Table R1).

  Table R1. Precision comparison of document retrieval.

  Dataset	BBCSport	Reuters	Ohsumed	Recipe
  Best Unsupervised	0.863	0.849	0.742	0.831
  UltraTree	0.842	0.834	0.749	0.830
  UltraTWD-GD	0.868	0.860	0.776	0.848
  UltraTWD-IP	0.885	0.876	0.788	0.866
  Improvement $\uparrow$	+5.1%	+5.0%	+5.2%	+4.3%

• Limitations of UltraTree: UltraTree requires training data with precomputed Wasserstein distances, which is costly and dataset-specific. Its performance also drops when training data sparsity differs from the test set (Table 3, Section 4.2).

• Advantages of UltraTWD: Our methods are fully unsupervised, robust to sparsity, and optimize $||D_T - D||_F^2$ directly. UltraTWD-GD achieves better performance while being up to 36$\times$ faster than UltraTree in tree learning time (Table 4, Section 4.4), making it much more scalable.

In summary, UltraTWD methods are more accurate, efficient, and broadly applicable than UltraTree.

---

Comment 6. Why in Algorithm 2: the weights are set to be $\frac{1}{t+1}$ and $\frac{t}{t+1}$ in which $t$ is the $t$-th iteration?

Response 6. The weight update in Algorithm 2 follows the HLWB projection scheme (Theorem 6, Section 3.4), where $\sigma_t = \frac{1}{t+1}$ forms a steering sequence. This gradually reduces the influence of the initial matrix $D$ while stabilizing updates ($\sigma_t \to 0$ as $t \to \infty$). While simple, this setting ensures convergence in convex settings and shows stable behavior empirically in our non-convex case (Figure 4, Section 4.4).